From 6ee5bb0a0b3e32dd1a6d9fddb61b50905ccdd6e0 Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Date: Thu, 10 Sep 2026 08:43:00 -0700 Subject: [PATCH 001/774] [DCP][Kernel][Perf] Fuse the empty-shard LSE mask into the A2A pack kernel (#54889) Signed-off-by: Rita Brugarolas Brufau Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_dcp_a2a_pack_mask.py | 151 +++++++++++++++++++ vllm/v1/attention/ops/dcp.py | 67 +++++++- 2 files changed, 213 insertions(+), 5 deletions(-) create mode 100644 tests/v1/attention/test_dcp_a2a_pack_mask.py diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py new file mode 100644 index 000000000000..36032e7a69a0 --- /dev/null +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The empty-shard mask fused into the DCP A2A pack kernel must be bit-exact +with the eager ``mask_dcp_empty_shards_`` pass it replaces.""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.ops.dcp import ( + _dcp_a2a_lse_pack_dim, + _dcp_a2a_pack_send, + mask_dcp_empty_shards_, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="needs a CUDA or ROCm device for the Triton pack kernel", +) + + +def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc): + lse_pack_dim = _dcp_a2a_lse_pack_dim(out.dtype) + send = torch.zeros( + (world_size, out.shape[0], h_per_rank, head_dim + lse_pack_dim), + device=out.device, + dtype=out.dtype, + ) + _dcp_a2a_pack_send( + out, + lse, + send, + world_size, + h_per_rank, + head_dim, + lse_pack_dim, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + return send + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)]) +@pytest.mark.parametrize("tokens_per_req", [1, 2, 3]) +@pytest.mark.parametrize("num_pad_rows", [0, 2]) +@pytest.mark.parametrize( + "seq_lens_list", + [ + [16, 16, 16, 16], # no empty shards (the long-context decode case) + [16, 0, 16, 0], # interior empty shards + [0, 0, 0, 0], # every shard empty + [16, 16, 0, 0], # trailing empty shards, i.e. cudagraph padding + [0, 16, 16, 16], # leading empty shard + [16] * 52, # the profiled conc-52 decode batch + [7], # single request + ], +) +def test_fused_mask_matches_eager( + dtype, world_size, h_per_rank, tokens_per_req, num_pad_rows, seq_lens_list +): + torch.manual_seed(0) + device = "cuda" + head_dim = 512 + num_seqs = len(seq_lens_list) + # tokens_per_req > 1 is the MTP / multi-token-verify shape. + num_rows = num_seqs * tokens_per_req + num_pad_rows + num_heads = world_size * h_per_rank + + out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype) + lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32) + seq_lens = torch.tensor(seq_lens_list, device=device, dtype=torch.int32) + query_start_loc = torch.arange( + 0, + (num_seqs + 1) * tokens_per_req, + tokens_per_req, + device=device, + dtype=torch.int32, + ) # rows past query_start_loc[-1] are cudagraph padding + + # Reference: mask eagerly, then pack with masking disabled. + ref_lse = lse.clone() + mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc) + expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None) + + # Under test: pack with the mask fused in. + actual = _pack( + out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc + ) + + # The packed LSE slots hold halves of an fp32, not meaningful floats, so + # compare the send buffers bitwise: -inf's high half is a NaN in fp16. + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), ( + "fused mask is not bit-exact with the eager mask" + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("seed", range(8)) +def test_fused_mask_matches_eager_ragged(dtype, seed): + """Non-uniform per-request query lengths. + + The uniform-stride cases above never exercise an unequal gap between + consecutive ``query_start_loc`` entries, which is what the in-kernel + boundary search has to get right. MTP with partial acceptance produces + exactly this: some requests contribute one row, others several, in the + same batch. + """ + torch.manual_seed(seed) + device = "cuda" + world_size, h_per_rank, head_dim = 8, 2, 512 + num_seqs = 17 + + query_lens = torch.randint(1, 5, (num_seqs,), dtype=torch.int32) + query_start_loc = torch.zeros(num_seqs + 1, dtype=torch.int32) + query_start_loc[1:] = torch.cumsum(query_lens, 0) + # Mix real, empty, and (via padding rows) out-of-range shards. + seq_lens = torch.randint(0, 2, (num_seqs,), dtype=torch.int32) * 16 + num_rows = int(query_start_loc[-1]) + 3 + + query_start_loc = query_start_loc.to(device) + seq_lens = seq_lens.to(device) + num_heads = world_size * h_per_rank + out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype) + lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32) + + ref_lse = lse.clone() + mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc) + expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None) + actual = _pack( + out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc + ) + + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), ( + f"ragged query_start_loc mismatch: query_lens={query_lens.tolist()}, " + f"seq_lens={seq_lens.tolist()}" + ) + + +def test_mask_disabled_is_unmasked(): + """Passing no seq_lens/query_start_loc must leave the LSE untouched.""" + torch.manual_seed(0) + device = "cuda" + world_size, h_per_rank, head_dim = 8, 2, 512 + out = torch.randn( + 4, world_size * h_per_rank, head_dim, device=device, dtype=torch.bfloat16 + ) + lse = torch.randn(4, world_size * h_per_rank, device=device, dtype=torch.float32) + + packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None) + assert torch.isfinite(packed.float()).all() diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 96b25870b8d0..66a6bace46e0 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -45,13 +45,18 @@ # LSE/output combine -def mask_dcp_empty_shards_( - lse: torch.Tensor, +def _validate_dcp_empty_shard_args( seq_lens: torch.Tensor | None, query_start_loc: torch.Tensor | None, -) -> None: +) -> bool: + """Validate the empty-shard mask inputs. + + Returns True when masking is requested, False when it is disabled. Raises on + an inconsistent pair. Performs no device work so it is safe to call from a + CUDA-graph-captured region. + """ if seq_lens is None and query_start_loc is None: - return + return False if seq_lens is None or query_start_loc is None: raise ValueError("seq_lens and query_start_loc must be provided together") if ( @@ -60,6 +65,17 @@ def mask_dcp_empty_shards_( or query_start_loc.shape[0] != seq_lens.shape[0] + 1 ): raise ValueError("query_start_loc must contain one boundary per sequence") + return True + + +def mask_dcp_empty_shards_( + lse: torch.Tensor, + seq_lens: torch.Tensor | None, + query_start_loc: torch.Tensor | None, +) -> None: + if not _validate_dcp_empty_shard_args(seq_lens, query_start_loc): + return + assert seq_lens is not None and query_start_loc is not None row_indices = torch.arange( lse.shape[0], device=lse.device, dtype=query_start_loc.dtype @@ -603,6 +619,9 @@ def _dcp_a2a_pack_send_kernel( out_ptr, lse_ptr, send_ptr, + seq_lens_ptr, + query_start_loc_ptr, + num_seqs, out_stride_B, out_stride_H, out_stride_D, @@ -616,11 +635,33 @@ def _dcp_a2a_pack_send_kernel( HEAD_DIM: tl.constexpr, H_PER_RANK: tl.constexpr, LSE_PACK_DIM: tl.constexpr, + HAS_MASK: tl.constexpr, + NUM_SEQS_POW2: tl.constexpr, ): batch_idx = tl.program_id(0).to(tl.int64) local_head_idx = tl.program_id(1).to(tl.int64) d_offsets = tl.arange(0, HEAD_DIM) + # Empty-shard masking, fused in to avoid a separate per-layer pass over the + # LSE. Rows that belong to a request with no local KV shard, and rows past + # the end of the batch (CUDA-graph padding), must not contribute to the + # combine, which the reduction achieves by giving them an LSE of -inf. + row_is_empty = False + if HAS_MASK: + # searchsorted(query_start_loc[1:], batch_idx, right=True), i.e. the + # number of request boundaries at or below this row. + seq_offsets = tl.arange(0, NUM_SEQS_POW2) + boundaries = tl.load( + query_start_loc_ptr + 1 + seq_offsets, + mask=seq_offsets < num_seqs, + other=2**31 - 1, + ) + seq_idx = tl.sum((boundaries <= batch_idx).to(tl.int32), axis=0) + seq_idx = tl.minimum(seq_idx, num_seqs - 1) + num_rows = tl.load(query_start_loc_ptr + num_seqs) + seq_len = tl.load(seq_lens_ptr + seq_idx) + row_is_empty = (batch_idx >= num_rows) | (seq_len == 0) + for rank_idx in tl.static_range(N): src_head_idx = rank_idx * H_PER_RANK + local_head_idx send_base = ( @@ -642,6 +683,8 @@ def _dcp_a2a_pack_send_kernel( lse_val = tl.load( lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H ).to(tl.float32) + if HAS_MASK: + lse_val = tl.where(row_is_empty, -float("inf"), lse_val) if LSE_PACK_DIM == 1: tl.store( send_ptr + send_base + HEAD_DIM * send_stride_D, @@ -797,12 +840,21 @@ def _dcp_a2a_pack_send( seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, ) -> None: - mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) + # The empty-shard mask is fused into the pack kernel below rather than + # applied as a separate pass: the mask is a pure function of seq_lens and + # query_start_loc, which are identical for every MLA layer in a step, while + # the eager form costs eight launch-bound kernels per layer. + has_mask = _validate_dcp_empty_shard_args(seq_lens, query_start_loc) + num_seqs = seq_lens.shape[0] if has_mask and seq_lens is not None else 0 + grid = (cp_attn_out.shape[0], h_per_rank, 1) _dcp_a2a_pack_send_kernel[grid]( cp_attn_out, cp_attn_lse, send_buffer, + seq_lens, + query_start_loc, + num_seqs, cp_attn_out.stride(0), cp_attn_out.stride(1), cp_attn_out.stride(2), @@ -816,6 +868,11 @@ def _dcp_a2a_pack_send( HEAD_DIM=head_dim, H_PER_RANK=h_per_rank, LSE_PACK_DIM=lse_pack_dim, + HAS_MASK=has_mask, + # Only the arange bound is a constexpr: num_seqs itself is a runtime + # argument so the kernel specializes on at most log2(max_num_seqs) + # variants instead of one per distinct batch size. + NUM_SEQS_POW2=triton.next_power_of_2(max(num_seqs, 1)), ) From 3078e7cbe1efb7dfcdeabceea9a6f890bc724976 Mon Sep 17 00:00:00 2001 From: Josiah Davis Date: Thu, 10 Sep 2026 11:13:09 -0500 Subject: [PATCH 002/774] [Docs]: quote variable-bearing wheel URLs (#56286) Signed-off-by: Josiah Davis --- docs/getting_started/installation/cpu.arm.inc.md | 4 ++-- docs/getting_started/installation/cpu.x86.inc.md | 4 ++-- docs/getting_started/installation/gpu.cuda.inc.md | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index e719d65dd95e..db49e96386e0 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -20,12 +20,12 @@ Pre-built vLLM wheels for Arm are available since version 0.11.2. These wheels c ```bash export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --torch-backend cpu +uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl" --torch-backend cpu ``` ??? console "pip" ```bash - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl" --extra-index-url https://download.pytorch.org/whl/cpu ``` The `uv` approach works for vLLM `v0.6.6` and later. A unique feature of `uv` is that packages in `--extra-index-url` have [higher priority than the default index](https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes). If the latest public release is `v0.6.6.post1`, `uv`'s behavior allows installing a commit before `v0.6.6.post1` by specifying the `--extra-index-url`. In contrast, `pip` combines packages from `--extra-index-url` and the default index, choosing only the latest version, which makes it difficult to install a development version prior to the released version. diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index b2ef384bd93b..a64d0e570333 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -24,13 +24,13 @@ Pre-built vLLM wheels for x86 with AVX512/AVX2 are available since version 0.17. export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') # use uv -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --torch-backend cpu +uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl" --torch-backend cpu ``` ??? console "pip" ```bash # use pip - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl" --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" Before use vLLM CPU installed via wheels, make Intel OpenMP is added to `LD_PRELOAD`: diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index cd6f2892c390..d110e377da8d 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -43,7 +43,7 @@ As of now, vLLM's binaries are compiled with CUDA 12.9 and public PyTorch releas export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') export CUDA_VERSION=130 # or other export CPU_ARCH=$(uname -m) # x86_64 or aarch64 -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION} +uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl" --extra-index-url "https://download.pytorch.org/whl/cu${CUDA_VERSION}" ``` #### Install the latest code From 7470082f57a0245703be7482544f4b0fd410e1f2 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:18:37 +0200 Subject: [PATCH 003/774] [ROCm][Perf] Add bpreshuffled blockscaled fp8 GEMM (#51692) Signed-off-by: simondanielsson --- vllm/_aiter_ops.py | 32 ++++- .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/scaled_mm/aiter.py | 128 ++++++++++++++++++ 3 files changed, 159 insertions(+), 5 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9bb018b3d33d..6b63588150ed 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -175,18 +175,26 @@ def is_aiter_found_and_supported_on_rdna4() -> bool: @functools.cache def _load_gemm_tuned_configs( - q_dtype_w: torch.dtype, csv_path: str -) -> set[tuple[int, int, int]]: + csv_path: str, + filters: tuple[tuple[str, object], ...], + key_cols: tuple[str, ...] = ("N", "K", "M"), +) -> set[tuple[int, ...]]: try: df = pd.read_csv(csv_path).drop_duplicates() - df = df[df["q_dtype_w"] == str(q_dtype_w)] - return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int))) + for col, val in filters: + if col not in df.columns: + continue + if isinstance(val, int): + df = df[df[col].astype(int) == val] + else: + df = df[df[col].astype(str) == str(val)] + return set(zip(*(df[c].astype(int) for c in key_cols))) except Exception: return set() def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool: - configs = _load_gemm_tuned_configs(q_dtype_w, csv_path) + configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),)) l_m = ( [1, 2, 4] + list(range(8, 513, 8)) @@ -3200,6 +3208,20 @@ def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE return _check_kernel_tuned(N, K, q_dtype_w, csv_path) + @staticmethod + def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool: + """Whether (N, K) has a tuned aiter blockscale bpreshuffle config.""" + if not current_platform.is_rocm(): + return False + import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops + + csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE + gfx = aiter_gemm_a8w8_ops.get_gfx() + cu_num = aiter_gemm_a8w8_ops.get_cu_num() + return (n, k) in _load_gemm_tuned_configs( + csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K") + ) + @staticmethod def shuffle_weight( tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 9c52e7b413f7..1df73273ca39 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -178,6 +178,7 @@ AiterHipbMMPerTokenFp8ScaledMMLinearKernel, AiterInt8ScaledMMLinearKernel, AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledFp8BlockScaledMMKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( @@ -316,6 +317,7 @@ def _get_linear_backend() -> str: "aiter": { AiterInt8ScaledMMLinearKernel, AiterFp8BlockScaledMMKernel, + AiterPreshuffledFp8BlockScaledMMKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, AiterMxfp4LinearKernel, @@ -456,6 +458,7 @@ def _resolve_backend_kernels( BlockWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ + AiterPreshuffledFp8BlockScaledMMKernel, AiterFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], @@ -1210,6 +1213,7 @@ def register_linear_kernel( "ScaledMMLinearLayerConfig", "AiterHipbMMPerTokenFp8ScaledMMLinearKernel", "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", + "AiterPreshuffledFp8BlockScaledMMKernel", "AiterPerTokenFp8ScaledMMLinearKernel", "NvFp4LinearKernel", "NvFp4LinearLayerConfig", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index b1518f21cc82..45f6d418b969 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -457,3 +457,131 @@ def apply_block_scaled_mm( return gemm_a8w8_blockscale_op( A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype ) + + +class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported( + compute_capability + ) + + @classmethod + def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): + can_implement_base, reason = super().can_implement(config) + if not can_implement_base: + return can_implement_base, reason + + act_quant_desc = config.activation_quant_key.scale + if act_quant_desc.group_shape != GroupShape(1, 128): + return ( + False, + ( + "Supports only dynamic per token group activation " + "quantization with group_shape=(1,128)." + ), + ) + + # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8 + # 2D weight with N and K divisible by 128. Require a tuned config as the + # fallback can trigger faults or numerical errors. + if not rocm_aiter_ops.is_linear_fp8_enabled(): + return ( + False, + ( + "requires setting `VLLM_ROCM_USE_AITER=1` " + "and `VLLM_ROCM_USE_AITER_LINEAR=1`." + ), + ) + + n, k = config.weight_shape + if not (n % 128 == 0 and k % 128 == 0): + return ( + False, + ( + f"requires N and K dimensions divisible by 128, received " + f"N={n} and K={k}." + ), + ) + + if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k): + return ( + False, + ( + f"requires a tuned aiter blockscale bpreshuffle config for " + f"N={n} and K={k}." + ), + ) + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + + params = FP8BlockParams.from_layer(layer) + if params.weight_scale_inv is not None: + ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV + else: + ws, attr = params.weight_scale, params.WEIGHT_SCALE + if ws is not None and ws.dtype == torch.float8_e8m0fnu: + replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) + + weight = params.weight + # runtime safety net + assert ( + weight.dim() == 2 + and weight.dtype == current_platform.fp8_dtype() + and weight.shape[0] % 128 == 0 + and weight.shape[1] % 128 == 0 + ), ( + "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight " + "with N and K divisible by 128." + ) + + shuffled_weight = rocm_aiter_ops.shuffle_weight( + weight.contiguous(), layout=(16, 16) + ) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(shuffled_weight.data, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + params = self._get_layer_params(layer) + Bs = ( + params.weight_scale + if params.weight_scale_inv is None + else params.weight_scale_inv + ) + + x_2d = x.view(-1, x.shape[-1]) + A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=True) + output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + A, params.weight, As, Bs, output_dtype=self.config.out_dtype + ) + if bias is not None: + output = output + bias + return output.view(*x.shape[:-1], params.weight.shape[0]) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "AiterPreshuffledFp8BlockScaledMMKernel overrides apply_weights and " + "does not use apply_block_scaled_mm." + ) From e6cb56337b49e606f55fde1870adbbb051e23f9f Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 10 Sep 2026 10:36:04 -0700 Subject: [PATCH 004/774] [Core] MRV2 support for fast-prefill (#56145) Signed-off-by: Nick Hill --- .buildkite/test_areas/engine.yaml | 2 +- .../general/test_kv_sharing_fast_prefill.py | 16 ++- tests/v1/worker/test_attn_utils.py | 64 ++++++++++ tests/v1/worker/test_gpu_model_runner_v2.py | 2 +- vllm/config/cache.py | 1 - vllm/config/vllm.py | 4 - vllm/v1/attention/backends/utils.py | 9 +- vllm/v1/worker/gpu/attn_utils.py | 116 +++++++++++++++++- vllm/v1/worker/gpu/input_batch.py | 5 + vllm/v1/worker/gpu/model_runner.py | 24 +++- vllm/v1/worker/gpu/model_states/default.py | 1 + vllm/v1/worker/gpu/ubatch_utils.py | 1 + 12 files changed, 233 insertions(+), 12 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 4384484c0954..a6b98bf9313b 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -77,7 +77,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) E2E Core" device: h200_35gb key: e2e-core-1-gpu - timeout_in_minutes: 40 + timeout_in_minutes: 50 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index 3155da82ea85..d92aa5c77572 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -52,12 +52,23 @@ def test_prompts(): @use_fork_for_test -@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True]) +@pytest.mark.parametrize( + "kv_sharing_fast_prefill,use_v2_model_runner", + [ + (False, False), + (True, False), + (True, True), + # (False, True) omitted: fast-prefill-off behavior is runner-generic, + # and skipping it saves ~9 min of mostly torch.compile time in CI. The + # no-silent-fallback assertion below is exercised by the V2 cases. + ], +) @pytest.mark.parametrize("enforce_eager", [True, False]) def test_kv_sharing_fast_prefill( monkeypatch: pytest.MonkeyPatch, kv_sharing_fast_prefill: bool, enforce_eager: bool, + use_v2_model_runner: bool, ): if not enforce_eager and current_platform.is_rocm(): # Relevant context: https://github.com/vllm-project/vllm/pull/29244 @@ -84,6 +95,7 @@ def test_kv_sharing_fast_prefill( m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") else: m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0") prompts, answer, indices = prep_prompts(batch_size) @@ -95,6 +107,8 @@ def test_kv_sharing_fast_prefill( kv_sharing_fast_prefill=kv_sharing_fast_prefill, attention_backend="TRITON_ATTN", ) + # Guard against a silent fallback to the other model runner. + assert llm.llm_engine.vllm_config.use_v2_model_runner == use_v2_model_runner responses = llm.generate(prompts, sampling_params) check_answers( indices, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9751865794b7..906525494f4b 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -7,6 +7,8 @@ never addressed by the logical view. """ +from types import SimpleNamespace + import pytest import torch @@ -22,6 +24,7 @@ MLAAttentionSpec, compute_layout_strides, ) +from vllm.v1.worker.gpu import attn_utils from vllm.v1.worker.gpu.attn_utils import ( get_attn_cg_support, get_query_lens_mismatch_unsupported_backend, @@ -118,6 +121,67 @@ def test_attention_checks_preserve_global_and_target_scoped_support(): ) +def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch): + """Fast prefill applies to the contiguous suffix of KV-sharing layers. + + Draft-model layers register after the target model's and may share KV, so + they must not extend (or break) the target's eligible suffix. + """ + + def check( + layer_names: list[str], + shared: dict[str, str], + draft_layer_names: set[str] | None = None, + ) -> set[str]: + monkeypatch.setattr( + attn_utils, + "get_layers_from_vllm_config", + lambda *a, **k: {name: None for name in layer_names}, + ) + monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda *a: shared) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(kv_sharing_fast_prefill=True) + ) + return attn_utils.get_kv_sharing_fast_prefill_eligible_layers( + vllm_config, draft_layer_names + ) + + # No KV sharing: nothing is eligible. + assert check(["t0", "t1"], {}) == set() + + # Trailing run of sharing layers (YOCO-style second half). + assert check(["t0", "t1", "t2", "t3"], {"t2": "t1", "t3": "t1"}) == {"t2", "t3"} + + # A non-sharing layer after a sharing one breaks the suffix. + assert check(["t0", "t1", "t2", "t3"], {"t1": "t0", "t3": "t0"}) == {"t3"} + + # KV-sharing draft layers at the end are collected without an exclusion... + assert check( + ["t0", "t1", "t2", "t3", "d0", "d1"], + {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"}, + ) == {"t2", "t3", "d0", "d1"} + + # ...so the runner excludes them: skipped, not collected, and they do not + # break the target's trailing run. + assert check( + ["t0", "t1", "t2", "t3", "d0", "d1"], + {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"}, + draft_layer_names={"d0", "d1"}, + ) == {"t2", "t3"} + + # Feature flag off: nothing is eligible even with sharing layers. + monkeypatch.setattr( + attn_utils, "get_layers_from_vllm_config", lambda *a, **k: {"t0": None} + ) + monkeypatch.setattr( + attn_utils, "get_shared_kv_cache_layers", lambda *a: {"t0": "t0"} + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(kv_sharing_fast_prefill=False) + ) + assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set() + + def test_reshape_padded_kv_cache_strides_by_padded_page(): num_blocks = 3 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index eb05b9858af4..1118ce589ca5 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -84,7 +84,7 @@ def narrow(self, *args): monkeypatch.setattr( model_runner_module, "init_attn_backend", - lambda *args: ([], attn_cg_support, [8, 262144]), + lambda *args, **kwargs: ([], attn_cg_support, [8, 262144]), ) monkeypatch.setattr( model_runner_module, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 4c48fc291ebb..93114ab1d150 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -234,7 +234,6 @@ class CacheConfig: some layers can skip tokens corresponding to prefill. This flag enables attention metadata for eligible layers to be overridden with metadata necessary for implementing this optimization in some models (e.g. Gemma3n) - NOTE: KV cache sharing is not supported for MRv2 (v2 model runner). """ kv_cache_memory_bytes: int | None = None diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 294f80b143e3..f5bee6fa2c37 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2674,10 +2674,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: ): unsupported.append("custom logits processors") - if self.cache_config.kv_sharing_fast_prefill: - # Will be added by https://github.com/vllm-project/vllm/pull/35045 - unsupported.append("KV sharing fast prefill") - if self.cache_config.mamba_cache_mode == "all": unsupported.append("mamba cache mode 'all'") diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 716f3689592f..99716ca911e6 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -626,8 +626,13 @@ def make_kv_sharing_fast_prefill_common_attn_metadata( # Skip computing fast prefill path return common_attn_metadata - assert common_attn_metadata.logits_indices_padded is not None - assert common_attn_metadata.num_logits_indices is not None + if ( + common_attn_metadata.logits_indices_padded is None + or common_attn_metadata.num_logits_indices is None + ): + # Fast prefill not armed for this step (e.g. cudagraph capture, or a + # pure-decode step): run the KV-sharing layers on the full batch. + return common_attn_metadata logits_indices_padded = common_attn_metadata.logits_indices_padded num_logits_indices = common_attn_metadata.num_logits_indices diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 05949c25ac21..4a859924256e 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -3,11 +3,13 @@ from collections.abc import Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast +import numpy as np import torch from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec @@ -15,6 +17,7 @@ AttentionCGSupport, CommonAttentionMetadata, ) +from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, @@ -31,6 +34,12 @@ prepare_kernel_block_sizes, ) +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, + ) + @dataclass(frozen=True) class AttentionCGSupportInfo: @@ -50,6 +59,66 @@ def narrow( return self +@dataclass(frozen=True) +class FastPrefillBatchMetadata: + """Per-step inputs for the KV-sharing fast prefill path.""" + + logits_indices_padded: torch.Tensor + num_logits_indices: int + max_logits_per_req: int + + +class FastPrefillHelper: + """Decides per step whether to arm the KV-sharing fast prefill path, and + stages the logits indices it runs on. + """ + + def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int): + self.max_num_tokens = max_num_tokens + self.cudagraph_manager = cudagraph_manager + self.logits_indices_buf = torch.zeros( + max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device + ) + + def prepare( + self, + logits_indices: torch.Tensor, + num_reqs: int, + cu_num_logits_np: np.ndarray, + has_prefill: bool, + batch_desc: "BatchExecutionDescriptor", + ) -> FastPrefillBatchMetadata | None: + if ( + not has_prefill + or batch_desc.cg_mode == CUDAGraphMode.FULL + or batch_desc.num_ubatches > 1 + ): + return None + buf = self.logits_indices_buf + num_logits = logits_indices.shape[0] + assert num_logits > 0 + buf[:num_logits].copy_(logits_indices) + # There might be leftover indices in buf[num_logits:] from previous + # iterations. Broadcast the scalar GPU-side to keep them valid. + buf[num_logits:] = logits_indices[-1] + # Pad so the model's KV-sharing layers run their reduced (logits-only) + # batch at a captured piecewise cudagraph size. Pre-capture, or with + # cudagraphs off, dispatch returns the unpadded count. + desc = self.cudagraph_manager.dispatch( + num_reqs=num_reqs, + num_tokens=num_logits, + uniform_token_count=None, + num_active_loras=0, + ) + num_logits_padded = min(desc.num_tokens, self.max_num_tokens) + return FastPrefillBatchMetadata( + logits_indices_padded=buf[:num_logits_padded], + num_logits_indices=num_logits, + # Largest per-request logits count, known on the host. + max_logits_per_req=int(np.diff(cu_num_logits_np).max()), + ) + + def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: kv_cache_spec: dict[str, KVCacheSpec] = {} layer_type = cast(type[Any], AttentionLayerBase) @@ -75,11 +144,42 @@ def get_shared_kv_cache_layers(vllm_config: VllmConfig): } +def get_kv_sharing_fast_prefill_eligible_layers( + vllm_config: VllmConfig, draft_layer_names: set[str] | None = None +) -> set[str]: + """Trailing run of KV-sharing layers, eligible for fast prefill. + + In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar + KV sharing setups, only the layers that generate KV caches are involved in + the prefill phase, enabling prefill to early exit. Layers are registered in + execution order, so the eligible layers are the contiguous suffix of + KV-sharing layers. + + Speculator draft layers register after the target model's layers (and may + themselves share KV), so they are excluded from the walk. + """ + if not vllm_config.cache_config.kv_sharing_fast_prefill: + return set() + shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config) + if not shared_kv_cache_layers: + return set() + eligible_layers: set[str] = set() + attn_layers = get_layers_from_vllm_config(vllm_config, Attention) + for layer_name in reversed(attn_layers): + if draft_layer_names is not None and layer_name in draft_layer_names: + continue + if layer_name not in shared_kv_cache_layers: + break + eligible_layers.add(layer_name) + return eligible_layers + + def init_attn_backend( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device, active_layer_names: set[str] | None = None, + draft_layer_names: set[str] | None = None, ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]: # Phase 1: discover attention groups for each kv cache group. attn_groups: list[list[AttentionGroup]] = [] @@ -89,6 +189,9 @@ def init_attn_backend( add_kv_sharing_layers_to_kv_cache_groups( get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups ) + fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers( + vllm_config, draft_layer_names + ) # Phase 1: discover attention groups for each kv cache group. for kv_cache_group_id, kv_cache_group_spec in enumerate( @@ -106,6 +209,10 @@ def init_attn_backend( for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() + if layer_name in fast_prefill_eligible_layers: + attn_backend = create_fast_prefill_custom_backend( + "FastPrefill", attn_backend + ) layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): @@ -278,6 +385,7 @@ def build_attn_metadata( causal: bool | torch.Tensor | Mapping[int, bool] = True, rswa_prefix_lens: torch.Tensor | None = None, ubatch_idx: int = 0, + fast_prefill: FastPrefillBatchMetadata | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -305,6 +413,12 @@ def build_attn_metadata( group_is_prefilling = common_attn_metadata_extra_kwargs.pop( "is_prefilling", is_prefilling ) + if fast_prefill is not None: + common_attn_metadata_extra_kwargs.update( + logits_indices_padded=fast_prefill.logits_indices_padded, + num_logits_indices=fast_prefill.num_logits_indices, + max_logits_per_req=fast_prefill.max_logits_per_req, + ) common_attn_metadata = CommonAttentionMetadata( query_start_loc=query_start_loc_gpu, query_start_loc_cpu=query_start_loc_cpu, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 2a38c3d71061..30b9b8089c04 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -11,6 +11,7 @@ from vllm.utils.math_utils import cdiv if TYPE_CHECKING: + from vllm.v1.worker.gpu.attn_utils import FastPrefillBatchMetadata from vllm.v1.worker.gpu.block_table import BlockTables @@ -109,6 +110,10 @@ class InputBatch: # stays valid for every replay the graph serves. max_query_len: int | None = None + # Arms the KV-sharing fast prefill path for this step. Absent for dummy + # (cudagraph capture) batches, which run the KV-sharing layers in full. + fast_prefill: "FastPrefillBatchMetadata | None" = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ec5ce6350377..533fcbe1e709 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -89,6 +89,7 @@ StepTimingCollector, ) from vllm.v1.worker.gpu.attn_utils import ( + FastPrefillHelper, build_slot_mappings_by_layer, get_kv_cache_spec, init_attn_backend, @@ -321,6 +322,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): max_num_tokens=self.max_num_tokens, device=self.device, ) + self.fast_prefill: FastPrefillHelper | None = None if self.use_pp: self.pp_handler = PPHandler( max_num_reqs=self.max_num_reqs, @@ -610,10 +612,14 @@ def initialize_kv_cache( for group in self.kv_cache_config.kv_cache_groups for layer_name in group.layer_names } - self.speculator.draft_attn_layer_names + draft_attn_layer_names = None + if isinstance(self.speculator, DraftModelSpeculator): + draft_attn_layer_names = self.speculator.draft_attn_layer_names self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device, + draft_layer_names=draft_attn_layer_names, ) additional_attn_cg_support = self.model_state.get_additional_cg_support() attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support) @@ -665,7 +671,7 @@ def initialize_kv_cache( initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config, - use_replayssm=self.vllm_config.cache_config.use_replayssm, + use_replayssm=self.cache_config.use_replayssm, ) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE @@ -687,6 +693,10 @@ def initialize_kv_cache( lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, ) + if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: + self.fast_prefill = FastPrefillHelper( + self.cudagraph_manager, self.max_num_tokens + ) check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) @@ -1339,6 +1349,16 @@ def prepare_inputs( self.model_state.num_new_sampled_tokens_per_step, ) + fast_prefill = None + if self.fast_prefill is not None: + fast_prefill = self.fast_prefill.prepare( + logits_indices, + num_reqs, + cu_num_logits_np, + batch_req_state.has_prefill, + batch_desc, + ) + # CPU upper bound on seq_lens; padded entries left at zero. num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np] seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32) @@ -1385,6 +1405,7 @@ def prepare_inputs( cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, prompt_lens=prompt_lens, + fast_prefill=fast_prefill, max_query_len=( int(num_scheduled_tokens_upper_bound.max()) if adaptive_verification is not None @@ -2127,6 +2148,7 @@ def shutdown(self) -> None: memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() self.cudagraph_manager = None + self.fast_prefill = None if hasattr(self, "kv_caches"): self.kv_caches.clear() if hasattr(self, "attn_groups"): diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7d42745d1064..6baceac9bc52 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -226,5 +226,6 @@ def prepare_attn( for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, ubatch_idx=ubatch_idx, + fast_prefill=input_batch.fast_prefill, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 012ce6c471a5..8007e2dbb44c 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -187,6 +187,7 @@ def _slice_input_batch( if input_batch.prompt_lens is None else input_batch.prompt_lens[req_start:req_stop] ), + fast_prefill=None, ) From 8359e15aae32dee9dc1f259a9b2574fb72b5507e Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:48:42 +0200 Subject: [PATCH 005/774] [ROCm][Feature] Support KV connectors with ROCM_AITER_UNIFIED_ATTN (#53695) Signed-off-by: simondanielsson Signed-off-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../test_rocm_attention_backends_selection.py | 112 ++++++++++++++++++ .../unit/test_offloading_connector.py | 16 ++- .../kv_connector/v1/offloading/config.py | 13 +- .../backends/rocm_aiter_unified_attn.py | 16 +++ 4 files changed, 154 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8668fa26bcc8..2342a279a397 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -455,3 +455,115 @@ def test_sparse_not_supported(mock_vllm_config): RocmPlatform.get_attn_backend_cls( selected_backend=None, attn_selector_config=attn_selector_config ) + + +def _kv_connector_selector_config() -> AttentionSelectorConfig: + return AttentionSelectorConfig( + head_size=128, + dtype=torch.float16, + kv_cache_dtype="auto", + block_size=16, + use_mla=False, + has_sink=False, + use_sparse=False, + use_kv_connector=True, + ) + + +def test_unified_attn_declares_kv_connector_support(): + """ROCM_AITER_UNIFIED_ATTN opts into KV connectors and ROCM_ATTN does not.""" + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend + + assert RocmAiterUnifiedAttentionBackend.supports_kv_connector() is True + assert RocmAttentionBackend.supports_kv_connector() is False + + +def test_unified_attn_supports_kv_connector(mock_vllm_config, mock_get_cdna_version): + """ROCM_AITER_UNIFIED_ATTN can be selected with KV connectors.""" + from vllm.platforms.rocm import RocmPlatform + + backend_path = RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, + attn_selector_config=_kv_connector_selector_config(), + ) + + assert backend_path == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path() + + +def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version): + """Selecting ROCM_ATTN with a KV connector is illegal.""" + from vllm.platforms.rocm import RocmPlatform + + attn_selector_config = _kv_connector_selector_config() + + with pytest.raises(ValueError, match="KV connector not supported"): + RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_ATTN, + attn_selector_config=attn_selector_config, + ) + + +@pytest.mark.parametrize( + "aiter_found, expected_backend", + [ + (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN), + (False, AttentionBackendEnum.TRITON_ATTN), + ], +) +def test_auto_selection_for_kv_connector( + aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version +): + """Auto-selection with a KV connector and AITER enabled resolves to unified attn, + and to triton attn if AITER not enabled.""" + from vllm.platforms.rocm import RocmPlatform + + with patch( + "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found + ): + backend_path = RocmPlatform.get_attn_backend_cls( + selected_backend=None, + attn_selector_config=_kv_connector_selector_config(), + ) + + assert backend_path == expected_backend.get_path() + + +def test_unified_attn_prefers_block_contiguous_layout(): + """Unified attn prefers a block-first KV layout, hence ok with kv connectors.""" + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend + + unified_preferred = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()[0] + rocm_attn_preferred = RocmAttentionBackend.supported_kv_cache_layouts()[0] + + assert unified_preferred.is_block_contiguous is True + assert rocm_attn_preferred.is_block_contiguous is False + + +def test_unified_attn_drops_lhbnc_with_kv_connector(): + """Connectors move a block as one contiguous byte range, which LHBNC breaks.""" + from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.kv_cache_interface import KVCacheLayout + + assert KVCacheLayout.LHBNC in ( + RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts() + ) + + config = VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="ExampleConnector", kv_role="kv_both" + ) + ) + with set_current_vllm_config(config): + layouts = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts() + + assert layouts + assert all(layout.is_block_compact for layout in layouts) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 08ef476992d5..cde672bcadba 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -15,12 +15,26 @@ from vllm.distributed.kv_events import BlockStored, KVEventBatch from vllm.platforms import current_platform -CPU_BLOCK_SIZES: int = 64 if current_platform.is_xpu() else 48 +# The offloaded block size must stay a whole number of GPU blocks. +CPU_BLOCK_SIZES: int +if current_platform.is_xpu(): + CPU_BLOCK_SIZES = 64 +elif current_platform.is_rocm(): + # ROCM_AITER_UNIFIED_ATTN prefers 64-token GPU blocks, so 192 covers both that + # and the 16-token default. + CPU_BLOCK_SIZES = 192 +else: + CPU_BLOCK_SIZES = 48 + _ATTN_BACKENDS: list[str] = [] if current_platform.is_cuda(): _ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"] elif current_platform.is_rocm(): + from vllm._aiter_ops import is_aiter_found_and_supported + _ATTN_BACKENDS = ["TRITON_ATTN"] + if is_aiter_found_and_supported(): + _ATTN_BACKENDS.append("ROCM_AITER_UNIFIED_ATTN") elif current_platform.is_xpu(): _ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index e5ed896078ca..bc78b7298822 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING +from vllm.utils.math_utils import round_up from vllm.v1.core.kv_cache_utils import ( resolve_dcp_kv_block_size, resolve_kv_cache_block_sizes, @@ -90,8 +91,16 @@ def build_offloading_config( ) tokens_per_block = unique_tokens_per_block.pop() - assert tokens_per_chunk_int % tokens_per_block == 0 - blocks_per_chunk = tokens_per_chunk_int // tokens_per_block + if tokens_per_chunk_int % tokens_per_block == 0: + blocks_per_chunk = tokens_per_chunk_int // tokens_per_block + else: + raise ValueError( + f"'block_size'={tokens_per_chunk_int} in kv_connector_extra_config " + f"must be a multiple of the GPU KV cache block size " + f"({tokens_per_block} tokens). Use " + f"{round_up(tokens_per_chunk_int, tokens_per_block)} instead, or set " + f"'blocks_per_chunk' to express the chunk size in blocks." + ) worker_kv_bytes_per_block = 0 if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors: diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index cbd33e7bc65a..e0e4b385509c 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -8,6 +8,7 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import get_current_vllm_config_or_none from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -32,6 +33,14 @@ logger = init_logger(__name__) +def _kv_connector_enabled() -> bool: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return False + kv_transfer_config = vllm_config.kv_transfer_config + return kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance + + class RocmAiterUnifiedAttentionMetadataBuilder(RocmAttentionMetadataBuilder): def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata @@ -84,6 +93,10 @@ def supports_sink(cls) -> bool: def supports_non_causal(cls) -> bool: return False + @classmethod + def supports_kv_connector(cls) -> bool: + return True + forward_includes_kv_cache_update: bool = False @staticmethod @@ -110,6 +123,9 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: # K and V come out of the content dim as transposed views rather than # copies, so the head dim may sit on either side of the block dim, but # the layer must stay outermost. + if _kv_connector_enabled(): + # Connectors like MoRI assume contiguous blocks + return (KVCacheLayout.LBHNC,) return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) @staticmethod From 983b7e28c9f87c3a150081640e0fdeafec5d553a Mon Sep 17 00:00:00 2001 From: Julien Denize <40604584+juliendenize@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:57:44 +0200 Subject: [PATCH 006/774] [Bugfix] Avoid MistralCommonBackend for HF tokenizers (#54192) Signed-off-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> --- tests/tokenizers_/test_registry.py | 84 +++++++++++++++ .../test_mistral_common_tokenizer.py | 102 ------------------ vllm/tokenizers/hf.py | 32 +----- vllm/tokenizers/registry.py | 45 +++++++- vllm/v1/structured_output/__init__.py | 5 +- vllm/v1/structured_output/backend_guidance.py | 5 - vllm/v1/structured_output/utils.py | 13 --- 7 files changed, 129 insertions(+), 157 deletions(-) delete mode 100644 tests/v1/structured_output/test_mistral_common_tokenizer.py diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index 0a47426dc974..b7201c1cfe4c 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -10,6 +10,7 @@ from transformers.models.auto.configuration_auto import CONFIG_MAPPING from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.hf import CachedHfTokenizer from vllm.tokenizers.registry import ( TokenizerRegistry, cached_get_tokenizer, @@ -68,6 +69,89 @@ def test_resolve_tokenizer_args_idempotent(runner_type): ) +@pytest.mark.parametrize( + ("tokenizer_mode", "input_kwargs"), + [ + ("hf", {}), + ("hf", {"mistral_format": False}), + ("slow", {}), + ("slow", {"mistral_format": False}), + ], +) +def test_resolve_tokenizer_args_forces_hf_mistral_format_false( + tokenizer_mode, input_kwargs +): + resolved_mode, _, _, kwargs = resolve_tokenizer_args( + "mistralai/Mistral-Nemo-Instruct-2407", + tokenizer_mode=tokenizer_mode, + **input_kwargs, + ) + + assert resolved_mode == "hf" + assert kwargs["mistral_format"] is False + + +@pytest.mark.parametrize("tokenizer_mode", ["hf", "slow"]) +def test_resolve_tokenizer_args_rejects_hf_mistral_format_true(tokenizer_mode): + with pytest.raises( + ValueError, + match="mistral_format=True is not supported with tokenizer_mode='hf'", + ): + resolve_tokenizer_args( + "mistralai/Mistral-Nemo-Instruct-2407", + tokenizer_mode=tokenizer_mode, + mistral_format=True, + ) + + +@pytest.mark.parametrize( + ("transformers_version", "model_type", "native_filename", "expect_error"), + [ + pytest.param("5.14.0", "mistral", "tekken.json", True), + pytest.param("5.15.0", "mistral", "tekken.json", False), + pytest.param("5.14.0", "mistral3", "tekken.json", True), + pytest.param("5.15.0", "mistral3", "tekken.json", False), + pytest.param("5.14.0", "mixtral", "tokenizer.model.v1", True), + pytest.param("5.15.0", "mixtral", "tokenizer.model.v1", False), + pytest.param("5.14.0", "mistral", None, False), + pytest.param("5.14.0", "qwen3_5_moe", "tekken.json", False), + ], +) +def test_get_tokenizer_handles_hf_mistral_transformers_compatibility( + tmp_path: Path, + transformers_version: str, + model_type: str, + native_filename: str | None, + expect_error: bool, +): + if native_filename is not None: + (tmp_path / native_filename).touch() + + with ( + patch( + "vllm.tokenizers.registry.get_config", + return_value=SimpleNamespace(model_type=model_type), + ), + patch( + "vllm.tokenizers.registry.transformers.__version__", + transformers_version, + ), + patch.object( + CachedHfTokenizer, + "from_pretrained", + return_value=SimpleNamespace(is_fast=True), + ) as from_pretrained, + ): + if expect_error: + with pytest.raises(ValueError, match="requires transformers>=5.15.0"): + get_tokenizer(str(tmp_path), tokenizer_mode="hf") + else: + tokenizer = get_tokenizer(str(tmp_path), tokenizer_mode="hf") + assert tokenizer.is_fast is True + + assert from_pretrained.called == (not expect_error) + + def test_customized_tokenizer(): TokenizerRegistry.register("test_tokenizer", __name__, TestTokenizer.__name__) diff --git a/tests/v1/structured_output/test_mistral_common_tokenizer.py b/tests/v1/structured_output/test_mistral_common_tokenizer.py deleted file mode 100644 index 74d7cb085cc5..000000000000 --- a/tests/v1/structured_output/test_mistral_common_tokenizer.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Structured output with a `MistralCommonBackend` tokenizer.""" - -import pytest -from transformers import AutoTokenizer, MistralCommonBackend - -from vllm.config import StructuredOutputsConfig, VllmConfig -from vllm.tokenizers import get_tokenizer -from vllm.tokenizers.mistral import MistralTokenizer -from vllm.v1.structured_output.backend_outlines import OutlinesBackend -from vllm.v1.structured_output.backend_types import StructuredOutputOptions -from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend -from vllm.v1.structured_output.utils import ( - _reduced_vocabulary, - maybe_wrap_mistral_common_tokenizer, -) - -TOKENIZER = "mistralai/Mistral-Nemo-Instruct-2407" -JSON_SCHEMA = ( - '{"type": "object", "properties": {"x": {"type": "integer"}}, ' - '"required": ["x"], "additionalProperties": false}' -) -VALID_DOCUMENT = '{"x": 1}' - -BACKENDS = [("xgrammar", XgrammarBackend), ("outlines", OutlinesBackend)] - - -@pytest.fixture(scope="module") -def loaded_tokenizer(): - tokenizer = get_tokenizer(tokenizer_name=TOKENIZER, tokenizer_mode="hf") - assert isinstance(tokenizer, MistralCommonBackend) - return tokenizer - - -@pytest.fixture(scope="module") -def wrapped_tokenizer(loaded_tokenizer): - return maybe_wrap_mistral_common_tokenizer(loaded_tokenizer) - - -def _encode(tokenizer, text): - special = set(tokenizer.all_special_ids) - return [tid for tid in tokenizer.encode(text) if tid not in special] - - -def _backend(cls, name, tokenizer): - vllm_config = VllmConfig( - structured_outputs_config=StructuredOutputsConfig(backend=name) - ) - return cls(vllm_config, tokenizer=tokenizer, vocab_size=tokenizer.vocab_size) - - -def test_wrapping_produces_a_mistral_tokenizer(loaded_tokenizer, wrapped_tokenizer): - assert isinstance(wrapped_tokenizer, MistralTokenizer) - assert len(wrapped_tokenizer.vocab) == loaded_tokenizer.vocab_size - - -def test_other_tokenizers_are_left_alone(): - tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") - assert maybe_wrap_mistral_common_tokenizer(tokenizer) is tokenizer - - -def test_wrapping_fixes_the_outlines_vocabulary(loaded_tokenizer, wrapped_tokenizer): - unwrapped_vocab = _reduced_vocabulary(loaded_tokenizer) - wrapped_vocab = _reduced_vocabulary(wrapped_tokenizer) - - assert unwrapped_vocab.keys() == wrapped_vocab.keys() - - mismapped = { - token for token, ids in unwrapped_vocab.items() if ids != wrapped_vocab[token] - } - assert mismapped - - assert len({tuple(wrapped_vocab[t]) for t in mismapped}) == len(mismapped) - assert len({tuple(unwrapped_vocab[t]) for t in mismapped}) < len(mismapped) - - -@pytest.mark.parametrize("name,cls", BACKENDS) -def test_backend_accepts_a_valid_document(wrapped_tokenizer, name, cls): - backend = _backend(cls, name, wrapped_tokenizer) - grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA) - - for token in _encode(wrapped_tokenizer, VALID_DOCUMENT): - assert grammar.accept_tokens("req", [token]) - - -@pytest.mark.parametrize("name,cls", BACKENDS) -def test_backend_rejects_an_invalid_token(wrapped_tokenizer, name, cls): - backend = _backend(cls, name, wrapped_tokenizer) - grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA) - - assert not grammar.accept_tokens("req", [_encode(wrapped_tokenizer, "]")[-1]]) - - -def test_xgrammar_vocab_size_matches_the_tokenizer(wrapped_tokenizer): - backend = _backend(XgrammarBackend, "xgrammar", wrapped_tokenizer) - assert backend.vocab_size == len(wrapped_tokenizer.vocab) - - -def test_xgrammar_rejects_the_unwrapped_tokenizer(loaded_tokenizer): - with pytest.raises(ValueError, match="Unsupported tokenizer type"): - _backend(XgrammarBackend, "xgrammar", loaded_tokenizer) diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py index bdc767acd66c..3d35952c5f4c 100644 --- a/vllm/tokenizers/hf.py +++ b/vllm/tokenizers/hf.py @@ -116,17 +116,10 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: tokenizer_all_special_tokens = tokenizer.all_special_tokens tokenizer_vocab = tokenizer.get_vocab() tokenizer_len = len(tokenizer) - # The underlying tokenizer class could be MistralCommonBackend, - # which does not implement is_fast in Transformers + # The underlying tokenizer class could be a specific backend, + # which does not always implement is_fast in Transformers tokenizer_is_fast = getattr(tokenizer, "is_fast", True) - # MistralCommonBackend is tekken-backed and needs byte-fallback-aware tokenization. - mistral_tekkenizer = None - if getattr(getattr(tokenizer, "tokenizer", None), "instruct_tokenizer", None): - from vllm.tokenizers.mistral import mistral_common_tekkenizer - - mistral_tekkenizer = mistral_common_tekkenizer(tokenizer) - max_token_id = max(tokenizer_vocab.values()) max_chars_per_token = max(len(tok) for tok in tokenizer_vocab) @@ -159,27 +152,6 @@ def max_chars_per_token(self) -> int: def is_fast(self) -> bool: return tokenizer_is_fast - def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False): - if mistral_tekkenizer is not None: - from vllm.tokenizers.mistral import tekken_convert_ids_to_tokens - - return tekken_convert_ids_to_tokens(mistral_tekkenizer, ids) - return super().convert_ids_to_tokens( - ids, skip_special_tokens=skip_special_tokens - ) - - def convert_tokens_to_string(self, tokens: list[str]) -> str: - if mistral_tekkenizer is not None: - from vllm.tokenizers.mistral import tekken_convert_tokens_to_string - - return tekken_convert_tokens_to_string(mistral_tekkenizer, tokens) - try: - return super().convert_tokens_to_string(tokens) - except NotImplementedError: - # The underlying tokenizer class could be MistralCommonBackend, - # which does not implement convert_tokens_to_string in Transformers - return "".join(tokens) - def get_vocab(self) -> dict[str, int]: return tokenizer_vocab diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index b9c9226675e1..46446ecd6df7 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -7,9 +7,15 @@ from typing import TYPE_CHECKING import huggingface_hub +import transformers +from packaging.version import Version +from transformers.models.auto.tokenization_auto import ( + TOKENIZER_MAPPING_NAMES as _HF_TOKENIZER_MAPPING_NAMES, +) +from transformers.utils.import_utils import is_mistral_common_available from typing_extensions import TypeVar, assert_never -import vllm.envs as envs +from vllm import envs from vllm.logger import init_logger from vllm.transformers_utils.config import _maybe_register_hf_config, get_config from vllm.transformers_utils.repo_utils import ( @@ -26,6 +32,8 @@ logger = init_logger(__name__) +_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL = Version("5.15.0") + # Model types whose hub tokenizer_class is incorrect and should be overridden with # TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a @@ -74,8 +82,6 @@ def register(self, tokenizer_mode: str, module: str, class_name: str) -> None: self.tokenizers[tokenizer_mode] = (module, class_name) - return None - def load_tokenizer_cls(self, tokenizer_mode: str) -> type[TokenizerLike]: if tokenizer_mode not in self.tokenizers: raise ValueError(f"No tokenizer registered for {tokenizer_mode=!r}.") @@ -150,12 +156,15 @@ def resolve_tokenizer_args( if ( tokenizer_mode == "auto" and is_mistral_model_repo( - model_name_or_path=str(tokenizer_name), revision=revision + model_name_or_path=str(tokenizer_name), + revision=revision, + token=kwargs.get("token"), ) and any_pattern_in_repo_files( model_name_or_path=str(tokenizer_name), allow_patterns=["tekken.json", "tokenizer.model.v*"], revision=revision, + token=kwargs.get("token"), ) ): tokenizer_mode = "mistral" @@ -164,6 +173,13 @@ def resolve_tokenizer_args( if tokenizer_mode == "auto": tokenizer_mode = "hf" + if tokenizer_mode == "hf": + if kwargs.pop("mistral_format", False): + raise ValueError( + "mistral_format=True is not supported with tokenizer_mode='hf'" + ) + kwargs["mistral_format"] = False + return tokenizer_mode, tokenizer_name, args, kwargs @@ -231,6 +247,27 @@ def get_tokenizer( trust_remote_code=trust_remote_code, revision=revision, config_format=config_format, + token=kwargs.get("token"), + ) + + # TODO: delete when Transformers version dependency is bumped >= 5.15.0 + if ( + tokenizer_mode == "hf" + and Version(transformers.__version__) < _MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL + and is_mistral_common_available() + and _HF_TOKENIZER_MAPPING_NAMES.get(getattr(config, "model_type", None)) + == "MistralCommonBackend" + and any_pattern_in_repo_files( + model_name_or_path=str(tokenizer_name), + allow_patterns=["tekken.json", "tokenizer.model.v*"], + revision=revision, + token=kwargs.get("token"), + ) + ): + raise ValueError( + "Loading Mistral models with tokenizer_mode='hf' requires " + f"transformers>={_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL}. " + "Please upgrade transformers or delete Mistral tokenizer files." ) # Some models have an incorrect tokenizer_class on the hub. diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index fd48c99d1cf9..ca6adc736d06 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -18,7 +18,6 @@ StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend -from vllm.v1.structured_output.utils import maybe_wrap_mistral_common_tokenizer if TYPE_CHECKING: import numpy as np @@ -77,8 +76,8 @@ def __init__(self, vllm_config: VllmConfig): # of CPUs. max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2) self.executor = ThreadPoolExecutor(max_workers=max_workers) - self.tokenizer = maybe_wrap_mistral_common_tokenizer( - cached_tokenizer_from_config(model_config=self.vllm_config.model_config) + self.tokenizer = cached_tokenizer_from_config( + model_config=self.vllm_config.model_config ) reasoning_parser_plugin = ( self.vllm_config.structured_outputs_config.reasoning_parser_plugin diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 16e1f54be6b8..0d8b7377c557 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any import torch -from transformers import MistralCommonBackend from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger @@ -98,10 +97,6 @@ def __post_init__(self): if is_mistral_tokenizer(self.tokenizer): self.ll_tokenizer = self.tokenizer.llg_tokenizer - elif isinstance(self.tokenizer, MistralCommonBackend): - from mistral_common.guidance.tokenizer import from_mistral_tokenizer - - self.ll_tokenizer = from_mistral_tokenizer(self.tokenizer.tokenizer) else: self.ll_tokenizer = llguidance_hf.from_tokenizer( self.tokenizer, max(self.vocab_size, len(self.tokenizer)) diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index aa5d59c9c5b4..92122f85d0fc 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -14,7 +14,6 @@ import regex as re import torch from cachetools import LRUCache -from transformers import MistralCommonBackend import vllm.envs as envs from vllm.logger import init_logger @@ -320,18 +319,6 @@ def get_outlines_cache(): re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$") -def maybe_wrap_mistral_common_tokenizer(tokenizer: TokenizerLike) -> TokenizerLike: - """The grammar backends cannot consume a `MistralCommonBackend` directly.""" - if not isinstance(tokenizer, MistralCommonBackend): - return tokenizer - - # Deferred: `vllm.tokenizers.mistral` pulls in a large dependency tree, and - # this module is imported by `vllm.v1.worker.gpu_model_runner`. - from vllm.tokenizers.mistral import MistralTokenizer - - return MistralTokenizer(tokenizer) - - def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]: """Create a map from vocabulary tokens to lists of equivalent token ids. From 7cdd9304ae2e46572f220741bf86e0b3c2da569c Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:12:15 -0400 Subject: [PATCH 007/774] [KV Connector] Support symmetric DCP disagg for hybrid mamba models (#55531) Signed-off-by: Wei Zhao Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 Co-authored-by: Wei Zhao Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/config/parallel.py | 7 +++++++ vllm/config/vllm.py | 7 ++----- .../kv_connector/v1/nixl/base_worker.py | 8 +++++--- .../kv_connector/v1/nixl/pull_worker.py | 5 ++++- vllm/engine/arg_utils.py | 16 +++++++++++++--- 5 files changed, 31 insertions(+), 12 deletions(-) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index bb1d2095de7b..bd2532b1ccfb 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -393,8 +393,15 @@ class is dynamically inherited by the worker class. This is used to inject in (rank i+1, block j) only after (rank i, block j) is fully occupied. Block_size should be greater than or equal to cp_kv_cache_interleave_size. Block_size should be divisible by cp_kv_cache_interleave_size. + + When --cp-kv-cache-interleave-size is omitted (None), the interleave size + is resolved automatically based on NIXL transfer requirements. + Explicit settings take priority. """ + _allow_auto_resolve_cp_interleave_size: bool = True + """Whether NIXL may select the interleave size automatically.""" + data_parallel_index: int = Field(init=False) """Equal to the data parallel rank but not used for torch process groups and not overridden for dense models.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f5bee6fa2c37..893acb00cdb3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1237,11 +1237,6 @@ def __post_init__(self): "PD with decode_context_parallel_size > 1 is only " "supported for MLA models." ) - assert not (self.model_config.is_hybrid and dcp_size > 1), ( - "PD with decode_context_parallel_size > 1 is not " - "supported for hybrid Mamba/SSM models." - ) - if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) @@ -2883,6 +2878,8 @@ def adjust_dcp_kv_cache_interleave_size( "NixlConnector" ): return + if not self.parallel_config._allow_auto_resolve_cp_interleave_size: + return # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid # scaling by dcp_size (we need the local block_size here). diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 5cb76932d44b..ceef10e9fd6e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -530,9 +530,6 @@ def __init__( # DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully # replicated or fully sharded. A DCP rank is always derivable this way. self.dcp_rank = self.tp_rank % self.dcp_size - if self._has_mamba and self.dcp_size > 1: - # Prefix-cache-aware DCP slicing isn't implemented for the Mamba group. - raise ValueError("DCP is not supported for hybrid MLA+Mamba models.") self.num_blocks = kv_cache_config.num_blocks self.enable_permute_local_kv = False @@ -2177,6 +2174,11 @@ def _validate_remote_agent_handshake( f"DCP sizes must divide one another: local={self.dcp_size}, " f"remote={remote_dcp_size} (engine {remote_engine_id})." ) + if self._has_mamba and self.dcp_size != remote_dcp_size: + raise RuntimeError( + "Hybrid MLA+Mamba NIXL transfers require matching DCP sizes, " + f"got local={self.dcp_size}, remote={remote_dcp_size}." + ) tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size) block_size_ratio = self.transfer_topo.block_size_ratio( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 758f5c9fab4b..d39a9fe287c1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -17,6 +17,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( ReadSpec, + _is_attention_spec, ) from vllm.logger import init_logger @@ -213,7 +214,9 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]: local_ids = group_ids(meta.local_block_ids, rank) remote_ids = group_ids(remote_logical_block_ids, rank) for g in range(num_groups): - if not local_ids[g]: + if not local_ids[g] or not _is_attention_spec( + self._group_spec_types[g] + ): continue local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching( local_ids[g], diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index bf5cc62d367d..db398438d0dd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -487,7 +487,7 @@ class EngineArgs: dcp_comm_backend: DCPCommBackend | None = ParallelConfig.dcp_comm_backend dcp_q_replicate: bool | None = ParallelConfig.dcp_q_replicate dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size - cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size + cp_kv_cache_interleave_size: int | None = None data_parallel_size: int = ParallelConfig.data_parallel_size data_parallel_rank: int | None = None data_parallel_start_rank: int | None = None @@ -1104,7 +1104,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: ) parallel_group.add_argument( "--cp-kv-cache-interleave-size", - **parallel_kwargs["cp_kv_cache_interleave_size"], + **{ + **parallel_kwargs["cp_kv_cache_interleave_size"], + "default": None, + }, ) parallel_group.add_argument( "--prefill-context-parallel-size", @@ -2363,7 +2366,14 @@ def create_engine_config( dcp_comm_backend=self.dcp_comm_backend, dcp_q_replicate=self.dcp_q_replicate, dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size, - cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + cp_kv_cache_interleave_size=( + self.cp_kv_cache_interleave_size + if self.cp_kv_cache_interleave_size is not None + else ParallelConfig.cp_kv_cache_interleave_size + ), + _allow_auto_resolve_cp_interleave_size=( + self.cp_kv_cache_interleave_size is None + ), _api_process_count=self._api_process_count, _api_process_rank=self._api_process_rank, assigned_physical_gpu_ids=self._resolve_device_ids(), From 48cb12c184e9b753c5de1c8ecab7953ded127b87 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:39:33 -0700 Subject: [PATCH 008/774] [ROCm][Bugfix] Fix profiler in TheRock image (#56190) Signed-off-by: Matthew Wong --- vllm/env_override.py | 85 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 74 insertions(+), 11 deletions(-) diff --git a/vllm/env_override.py b/vllm/env_override.py index a931b2305fa8..a6b3c373face 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -5,21 +5,28 @@ import os -def _get_torch_cuda_version(): - """Peripheral function to _maybe_set_cuda_compatibility_path(). +def _get_torch_root(): + """Locate the installed torch package without importing it.""" + spec = importlib.util.find_spec("torch") + if not spec: + return None + if spec.origin: + return os.path.dirname(spec.origin) + if spec.submodule_search_locations: + return spec.submodule_search_locations[0] + return None + + +def _get_torch_version_attr(attr): + """Read an attribute of torch.version without importing torch. + PyTorch version must not be determined by importing directly because it will trigger the CUDA initialization, losing the chance to set the LD_LIBRARY_PATH beforehand. """ try: - spec = importlib.util.find_spec("torch") - if not spec: - return None - if spec.origin: - torch_root = os.path.dirname(spec.origin) - elif spec.submodule_search_locations: - torch_root = spec.submodule_search_locations[0] - else: + torch_root = _get_torch_root() + if not torch_root: return None version_path = os.path.join(torch_root, "version.py") if not os.path.exists(version_path): @@ -31,11 +38,16 @@ def _get_torch_cuda_version(): module = importlib.util.module_from_spec(ver_spec) # Avoid registering in sys.modules to not confuse future imports ver_spec.loader.exec_module(module) - return getattr(module, "cuda", None) + return getattr(module, attr, None) except Exception: return None +def _get_torch_cuda_version(): + """Peripheral function to _maybe_set_cuda_compatibility_path().""" + return _get_torch_version_attr("cuda") + + def _maybe_set_cuda_compatibility_path(): """Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled. @@ -82,7 +94,58 @@ def _maybe_set_cuda_compatibility_path(): os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_paths) +def _maybe_promote_torch_symbols_for_rocm(): + """Put libtorch_cpu.so in the global symbol scope on ROCm, for GPU profiling. + + Must run before 'import torch'. libkineto advertises itself to + rocprofiler-sdk by exporting rocprofiler_configure from libtorch_cpu.so, + and rocprofiler-sdk looks for that symbol exactly once, from a lazy + initializer that runs *during* 'import torch'. CPython loads extension + modules RTLD_LOCAL, so unless something has already placed the symbol in + the global scope the lookup misses and no client is ever registered. + + The failure is silent and total: no queue interception, so roctracer + yields no GPU records, while torch.profiler still writes a complete + CPU-only trace and reports success. + + Best effort. If the library cannot be loaded early we leave the process + exactly as it would have been, losing only GPU tracing. + """ + if _get_torch_version_attr("hip") is None: + return + try: + import ctypes + import glob + + torch_root = _get_torch_root() + if not torch_root: + return + lib = os.path.join(torch_root, "lib", "libtorch_cpu.so") + if not os.path.exists(lib): + return + try: + ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) + except OSError: + # Some ROCm wheel layouts leave a stale RUNPATH in libtorch_cpu.so, + # so a dependency that resolves once torch has set things up does + # not resolve for a standalone load. Satisfy it by search, then + # retry; give up quietly if it is not where we expect. + site_root = os.path.dirname(torch_root) + deps = glob.glob( + os.path.join( + site_root, "_rocm_sdk_*", "lib", "*", "lib", "librocm-openblas.so.*" + ) + ) + if not deps: + return + ctypes.CDLL(deps[0], mode=ctypes.RTLD_LOCAL) + ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) + except Exception: + return + + _maybe_set_cuda_compatibility_path() +_maybe_promote_torch_symbols_for_rocm() import torch From 7de70fa7ae9fdda688257eed13639b56cc677fde Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Thu, 10 Sep 2026 13:54:52 -0500 Subject: [PATCH 009/774] [Bugfix][ROCm] Create linear layer biases with `requires_grad=False` (#56161) Signed-off-by: Micah Williamson --- .buildkite/test-amd.yaml | 2 +- vllm/model_executor/layers/linear.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 9561fec5e04e..cca01e45ba80 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4022,7 +4022,7 @@ steps: #------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# - label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index cd1c0e69d2ae..2e77fac06c34 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -380,7 +380,8 @@ def __init__( if bias: self.bias = Parameter( - torch.empty(self.output_size, dtype=self.params_dtype) + torch.empty(self.output_size, dtype=self.params_dtype), + requires_grad=False, ) set_weight_attrs( self.bias, @@ -525,7 +526,8 @@ def __init__( if bias: self.bias = Parameter( - torch.empty(self.output_size_per_partition, dtype=params_dtype) + torch.empty(self.output_size_per_partition, dtype=params_dtype), + requires_grad=False, ) set_weight_attrs( self.bias, @@ -1708,7 +1710,10 @@ def __init__( ) if bias: - self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype)) + self.bias = Parameter( + torch.empty(self.output_size, dtype=params_dtype), + requires_grad=False, + ) set_weight_attrs( self.bias, { From a36dfc93cb1e4ca11425cfde48b5ed2264027aa2 Mon Sep 17 00:00:00 2001 From: Jiatai Wang Date: Thu, 10 Sep 2026 12:20:03 -0700 Subject: [PATCH 010/774] [Bugfix][Multimodal] Restore cached audio inputs with UUIDs (#56310) Signed-off-by: JiataiWang Co-authored-by: OpenAI Codex --- .../entrypoints/unit_tests/test_chat_utils.py | 8 ++++++-- tests/multimodal/test_parse.py | 18 ++++++++++++++++++ vllm/entrypoints/chat_utils.py | 2 +- vllm/multimodal/parse.py | 9 +++++++-- 4 files changed, 32 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py index 575bb904d664..630f641df59d 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/entrypoints/unit_tests/test_chat_utils.py @@ -2384,8 +2384,10 @@ def test_parse_chat_messages_include_thinking_chunk(mistral_model_config): assert conversation_with_thinking == expected_conversation +@pytest.mark.parametrize("input_audio", [None, {}, {"data": "", "format": "wav"}]) def test_parse_chat_messages_single_empty_audio_with_uuid( qwen2_audio_model_config, + input_audio, ): audio_uuid = "abcd" conversation, mm_data, mm_uuids = parse_chat_messages( @@ -2395,7 +2397,7 @@ def test_parse_chat_messages_single_empty_audio_with_uuid( "content": [ { "type": "input_audio", - "input_audio": {}, + "input_audio": input_audio, "uuid": audio_uuid, }, {"type": "text", "text": "What does the audio say?"}, @@ -2418,8 +2420,10 @@ def test_parse_chat_messages_single_empty_audio_with_uuid( @pytest.mark.asyncio +@pytest.mark.parametrize("input_audio", [None, {}, {"data": "", "format": "wav"}]) async def test_parse_chat_messages_single_empty_audio_with_uuid_async( qwen2_audio_model_config, + input_audio, ): audio_uuid = "abcd" conversation, mm_data, mm_uuids = await parse_chat_messages_async( @@ -2429,7 +2433,7 @@ async def test_parse_chat_messages_single_empty_audio_with_uuid_async( "content": [ { "type": "input_audio", - "input_audio": {}, + "input_audio": input_audio, "uuid": audio_uuid, }, {"type": "text", "text": "What does the audio say?"}, diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index 685e615f7832..c34c8eec8b8f 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -6,6 +6,7 @@ from PIL import Image from vllm.multimodal.parse import ( + AudioProcessorItems, ImageProcessorItems, MultiModalDataParser, VideoProcessorItems, @@ -58,6 +59,7 @@ def test_frame_size_hwc_chw(frame): @pytest.mark.parametrize( "modality,processor_cls", [ + ("audio", AudioProcessorItems), ("image", ImageProcessorItems), ("video", VideoProcessorItems), ], @@ -68,3 +70,19 @@ def test_parse_mm_data_accepts_none_cached_item(modality, processor_cls): assert isinstance(items, processor_cls) assert len(items) == 1 assert items.get(0) is None + + +def test_cached_audio_items_preserve_positions_during_resampling(): + waveform = np.arange(16, dtype=np.float32) + parser = MultiModalDataParser( + target_sr=16000, target_channels=1, audio_resample_method="scipy" + ) + items = parser.parse_mm_data( + {"audio": [None, (waveform, 8000), None, (waveform, 16000)]} + )["audio"] + + assert len(items) == 4 + assert items.get(0) is None + assert items.get(2) is None + assert len(items.get(1)) == 32 + np.testing.assert_array_equal(items.get(3), waveform) diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 327f9c4b24a4..addcc5f76df8 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -1773,7 +1773,7 @@ def _parse_chat_message_content_mm_part( # with url as a dict of {"url": url} audio_url = audio_url.get("url", None) return "audio_url", audio_url - if part.get("input_audio") is not None: + if "input_audio" in part: input_audio_params = _InputAudioParser(part).get("input_audio", None) return "input_audio", input_audio_params if "video_url" in part: diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index d96f19f39b95..fefaeef2b019 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -678,7 +678,7 @@ def _parse_audio_data( if self.is_embeddings(data): return AudioEmbeddingItems(data, self.expected_hidden_size) - data_items: list[AudioItem] + data_items: list[AudioItem | None] if ( (is_list_of(data, float) and len(data) > 0) or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 1) @@ -690,8 +690,13 @@ def _parse_audio_data( else: data_items = data # type: ignore[assignment] - new_audios = list[np.ndarray]() + new_audios = list[np.ndarray | None]() for data_item in data_items: + # Requests can omit audio samples when reusing a cached UUID. + if data_item is None: + new_audios.append(None) + continue + audio, orig_sr = self._get_audio_with_sr(data_item) if orig_sr is None: new_audio = audio From ae71862c51f2e069998e0dd5dc30a03a3b475ab9 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Thu, 10 Sep 2026 12:23:34 -0700 Subject: [PATCH 011/774] [Bugfix][Frontend] Check EC requirements for each metadata item (#56070) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- .../token_in_token_out/test_mm_serde.py | 45 +++++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 17 ++++--- 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index 8bd1bb9091bf..764d684445d8 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -5,6 +5,7 @@ token_in_token_out generate endpoint. """ +import pytest import torch from pydantic import ValidationError @@ -241,6 +242,50 @@ def test_metadata_only_generate_requires_ec_transfer_params(): assert set(item) == {"image_grid_thw"} +@pytest.mark.parametrize("modality", ["image", "video"]) +@pytest.mark.parametrize( + ("has_metadata", "has_ec"), + [(True, False), (True, True), (False, False)], + ids=["metadata-without-ec", "metadata-with-ec", "cache-hit-without-ec"], +) +def test_mixed_items_require_ec_only_for_metadata_only( + modality: str, has_metadata: bool, has_ec: bool +): + """A full item must not bypass the EC requirement for another item.""" + engine_input, _, image_grid_thw = _image_engine_input() + rendered = extract_mm_features(engine_input) + assert rendered is not None + features = rendered.model_dump() + metadata = ( + encode_mm_kwargs_item( + MultiModalKwargsItem({f"{modality}_grid_thw": image_grid_thw}) + ) + if has_metadata + else None + ) + second_item = { + "mm_hashes": "second-item", + "mm_placeholders": {"offset": 2, "length": 1}, + "kwargs_data": None, + "mm_metadata": metadata, + } + for key, value in second_item.items(): + features[key].setdefault(modality, []).append(value) + payload = { + "token_ids": [1, 2, 3], + "sampling_params": {"max_tokens": 1}, + "features": features, + } + if has_ec: + payload["ec_transfer_params"] = {"second-item": {"peer_host": "10.0.0.1"}} + + if has_metadata and not has_ec: + with pytest.raises(ValidationError, match="ec_transfer_params"): + GenerateRequest.model_validate(payload) + else: + GenerateRequest.model_validate(payload) + + def test_kwargs_and_metadata_generate_does_not_require_ec(): engine_input, pixel_values, image_grid_thw = _image_engine_input() rendered = extract_mm_features(engine_input) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 65643b9e9927..aec8431b4780 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -192,14 +192,19 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest": return self if not _has_serialized_mm_items(features.mm_metadata): return self - if _has_serialized_mm_items(features.kwargs_data): - return self if self.ec_transfer_params: return self - raise ValueError( - "features.mm_metadata without features.kwargs_data requires " - "ec_transfer_params so embeddings can be loaded by the EC connector" - ) + kwargs_data = features.kwargs_data or {} + for modality, metadata_items in (features.mm_metadata or {}).items(): + kwargs_items = kwargs_data.get(modality, [None] * len(metadata_items)) + if any( + metadata is not None and kwargs is None + for metadata, kwargs in zip(metadata_items, kwargs_items, strict=True) + ): + raise ValueError( + "metadata-only multimodal items require ec_transfer_params" + ) + return self sampling_params: SamplingParams """The sampling parameters for the model.""" From 9163190dda009a310d8c175d63860ac7c671ca90 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:14:22 -0700 Subject: [PATCH 012/774] [ROCm][Bugfix][Perf] Tune multi-stream shared experts use; wvSplitKrc fixes (#56098) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5 (1M context) --- csrc/rocm/skinny_gemms.cu | 248 +++++++++++++++--- .../quantization/test_rocm_skinny_gemms.py | 97 +++++++ vllm/distributed/parallel_state.py | 6 + .../layers/fused_moe/runner/moe_runner.py | 13 - .../layers/fused_moe/runner/shared_experts.py | 13 - vllm/model_executor/layers/utils.py | 69 +++-- 6 files changed, 369 insertions(+), 77 deletions(-) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 615cdabed585..79af0814550d 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -8,12 +8,20 @@ #include #include +#include +#include +#include #include "../cuda_compat.h" #include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" #include "core/batch_invariant.hpp" +// Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot. +static constexpr int64_t kWvSlots = 8; +// Slots held back for capturing streams, which cannot allocate one themselves. +static constexpr int64_t kWvCaptureSlots = 8; + // TODO(rasmith): The kernels in this file are susceptible to integer overflow // issues, do not take strides, and are unable to handle PyTorch tensors that // return is_contiguous() as False (the tensors may actually be contiguous @@ -1364,9 +1372,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) }; using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16; - __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB]; + // The main-loop staging areas are separated from the split-K readback by a + // __syncthreads(), so the readback may reuse all of LDS. + __shared__ union { + struct { + scalar_t stg[WvPrGrp * WVLDS / GrpsShrB]; + scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB]; + } mainloop; + scalar_t all[max_lds_len]; + } lds; + auto& stg = lds.mainloop.stg; + auto& s = lds.mainloop.s; unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]); - __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB]; #ifndef WVSPLITKRC_1KPASS constexpr int TUC_ = (THRDS * UNRL * A_CHUNK); @@ -1682,33 +1699,69 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if (my_cntr + 1 == k_rnd) { cntr[adr_] = 0; // clear for next round if constexpr (DTRMNSTC) { + constexpr uint32_t NT = N / NTILE / GrpsShrB; + // The readback stages k-shard partials through LDS, which bounds how + // many shards may be in flight at once. + constexpr uint32_t LDS_F4 = + max_lds_len * sizeof(scalar_t) / sizeof(float4); + constexpr uint32_t KSBMAX = LDS_F4 / (THRDS * 4 * NT); + static_assert(KSBMAX >= 1, "LDS cannot stage even one k-shard"); + const uint32_t KSB = (k_rnd < KSBMAX) ? k_rnd : KSBMAX; + + // First batch is peeled; it is the only batch for most shapes. #pragma unroll - for (int ks = 0; ks < k_rnd; ks++) { - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t kb = 0; kb < KSB; kb++) { + for (uint32_t nt = 0; nt < NT; nt++) { int g_nindx = (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4; int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4; __builtin_amdgcn_global_load_lds( - (float4*)(&glbl[g_adr + M * N * ks]), - &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 + - nt * THRDS * 4 * k_rnd]), + (float4*)(&glbl[g_adr + M * N * kb]), + &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 + + nt * THRDS * 4 * KSB]), 16, 0, 0); } } - if (BIAS) - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + if (BIAS) // overlaps the DMAs above + for (uint32_t nt = 0; nt < NT; nt++) { for (uint32_t j = 0; j < 4; j++) { int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx]; } } - asm volatile("s_waitcnt 0"); - for (int ks = 0; ks < k_rnd; ks++) { - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { - float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) + - ks * THRDS * 4 + nt * THRDS * 4 * k_rnd]; - vals[nt].f4 += eval; + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + for (uint32_t kb = 0; kb < KSB; kb++) { + for (uint32_t nt = 0; nt < NT; nt++) { + vals[nt].f4 += + ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) + + kb * THRDS * 4 + nt * THRDS * 4 * KSB]; + } + } + + for (uint32_t ks0 = KSB; ks0 < k_rnd; ks0 += KSB) { + const uint32_t ksn = (k_rnd - ks0 < KSB) ? (k_rnd - ks0) : KSB; + // prior reads must retire before the staging is overwritten + asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory"); + for (uint32_t kb = 0; kb < ksn; kb++) { + for (uint32_t nt = 0; nt < NT; nt++) { + int g_nindx = + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4; + int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4; + __builtin_amdgcn_global_load_lds( + (float4*)(&glbl[g_adr + M * N * (ks0 + kb)]), + &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 + + nt * THRDS * 4 * KSB]), + 16, 0, 0); + } + } + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + for (uint32_t kb = 0; kb < ksn; kb++) { + for (uint32_t nt = 0; nt < NT; nt++) { + vals[nt].f4 += + ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) + + kb * THRDS * 4 + nt * THRDS * 4 * KSB]; + } } } } else { @@ -1771,7 +1824,7 @@ __global__ void wvSplitKrc_(const int actlN, const int K, const int Kap, torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b, const std::optional& in_bias, const int64_t CuCount) { - int _DTRMNSTC = 1; // vllm::vllm_is_batch_invariant(); + static constexpr bool _DTRMNSTC = true; auto M_in = in_b.size(0); auto N_in = in_a.size(0); @@ -1819,26 +1872,157 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b, if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size"); - // Can we increase SplitK by shrinking the K-shared to 256? - int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1; - - static torch::Tensor axl_glbl = - torch::zeros( - 128 * 1024 * (_DTRMNSTC ? 12 : 1), - torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device())) - .detach(); - static torch::Tensor axl_cntr = - torch::zeros( - 128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4, - torch::TensorOptions().dtype(torch::kInt).device(in_a.device())) - .detach(); - auto glbl = axl_glbl.data_ptr(); - auto cntr = axl_cntr.data_ptr(); + // Can we increase SplitK by shrinking the K-shard to 256? That doubles both + // the CUs needed and the k-shards to read back, and the readback is serial in + // the shard count, so it pays only up to a crossover. The crossover is not a + // constant -- it falls as M grows -- so one cap cannot sit on it for every + // shape; this value minimizes total time over the admitted shapes. Retune it + // if the readback cost changes, and note the directions are not symmetric: + // capping too high costs about twice what capping too low does. + // Mirrored by rocm_unquantized_gemm_impl() in + // vllm/model_executor/layers/utils.py; both must pick the same chunkk or the + // host-side fit check bounds the wrong k_rnd. + constexpr int64_t CHUNKK2_MAX_SHARDS = 11; + const int64_t shards_chunkk2 = (K_in + 256 - 1) / 256; // 256 == 512 / 2 + const int chunkk = (N_p2 != 16 && CuNeeded * 2 <= CuCount && + shards_chunkk2 <= CHUNKK2_MAX_SHARDS) + ? 2 + : 1; + + // One fp32 partial per (M, N, k-shard) must fit the split-K workspace. + const int64_t k_rnd = (K_in + 512 / chunkk - 1) / (512 / chunkk); + constexpr int64_t glbl_cap = 128 * 1024 * (_DTRMNSTC ? 12 : 1); + const int64_t glbl_needed = + _DTRMNSTC ? (int64_t)N_p2 * M_in * k_rnd : (int64_t)N_p2 * M_in; + TORCH_CHECK(glbl_needed <= glbl_cap, + "wvSplitKrc: split-K workspace too small", " (need ", glbl_needed, + " > ", glbl_cap, " floats) for N=", N_in, " K=", K_in, + " M=", M_in); + + // Split-K workspaces: one pre-partitioned pool per device. + // These must be zeroed before first use, and the kernel re-zeros them after + // each use. + // To avoid unnecessary overhead, we do not allocate them zeroed + // per-invocation. However, a simple static allocation does not work with + // multi-streams, because concurrent streams can alias into the same + // workspace. Solution: pre-allocate kWvSlots = 8 slots per device, plus + // kWvCaptureSlots reserved for capturing streams, and zero the whole pool + // once at creation. On first use, a stream takes one of the slots if any are + // still available. On subsequent uses, it re-uses the slot. If no slots are + // available, fall back to allocating a zeroed one per call, which is safe but + // adds overhead. + // A capturing stream must neither allocate nor enqueue anything -- either + // would be captured into the graph -- so it is only handed a reserved slot. + struct WvSplitKrcPool { + torch::Tensor glbl, cntr; + int64_t glbl_stride = 0, cntr_stride = 0; + int next_slot = 0; + int next_capture_slot = 0; + std::map slot_of; + std::map capture_slot_of; + }; + // mutex not strictly necessary if caller is always single-threaded (Python) + // but this preserves Torch thread_local semantics for the C++ API + static std::mutex wv_pool_mu; + static std::map wv_pools; + + float* glbl; + int* cntr; + // Fallback space + torch::Tensor spill_glbl, spill_cntr; + { + std::lock_guard wv_lk(wv_pool_mu); + const int wv_dev = static_cast(in_a.device().index()); + WvSplitKrcPool& P = wv_pools[wv_dev]; + constexpr int64_t wv_gf = 128 * 1024 * (_DTRMNSTC ? 12 : 1); + constexpr int64_t wv_ci = wv_gf / 4; + constexpr int64_t wv_all_slots = kWvSlots + kWvCaptureSlots; + + // A stream whose capture was already invalidated reports an error; treat + // that as not capturing, since the capture is lost either way. + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + unsigned long long capture_id = 0; + if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) != + cudaSuccess) { + capture_status = cudaStreamCaptureStatusNone; + } + const bool wv_capturing = capture_status != cudaStreamCaptureStatusNone; + + if (!P.glbl.defined()) { + TORCH_CHECK(!wv_capturing, + "wvSplitKrc: the split-K workspace pool cannot be allocated " + "during graph capture. Call " + "warmup_rocm_skinny_gemm_workspaces() " + "(vllm/model_executor/layers/utils.py) eagerly first."); + P.glbl_stride = wv_gf; + P.cntr_stride = wv_ci; + P.glbl = torch::zeros(wv_gf * wv_all_slots, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(in_a.device())) + .detach(); + P.cntr = + torch::zeros( + wv_ci * wv_all_slots, + torch::TensorOptions().dtype(torch::kInt).device(in_a.device())) + .detach(); + // The mutex orders host access to the map, not device execution, so it + // alone would let another stream be handed a slot the fill above is + // still clearing. Synchronizing here, inside the lock, is what makes + // every slot safe to use without any further per-slot zeroing. + const cudaError_t wv_err = cudaStreamSynchronize(stream); + TORCH_CHECK(wv_err == cudaSuccess, + "wvSplitKrc: failed to zero the split-K workspace pool: ", + cudaGetErrorString(wv_err)); + } + + int wv_slot = -1; + if (wv_capturing) { + // One workspace per capturing stream, so graphs captured on the same + // stream share it -- safe because at most one of them replays at a time + // (vLLM captures one graph per padded batch size, and runs one per step). + auto cap_it = P.capture_slot_of.find(stream); + if (cap_it == P.capture_slot_of.end()) { + TORCH_CHECK(P.next_capture_slot < kWvCaptureSlots, + "wvSplitKrc: no free capture workspace slot; more than ", + kWvCaptureSlots, + " distinct streams have captured a graph containing " + "wvSplitKrc. Raise kWvCaptureSlots in " + "csrc/rocm/skinny_gemms.cu."); + cap_it = P.capture_slot_of + .emplace(stream, static_cast(kWvSlots) + + P.next_capture_slot++) + .first; + } + wv_slot = cap_it->second; + } else { + auto slot_it = P.slot_of.find(stream); + if (slot_it != P.slot_of.end()) { + wv_slot = slot_it->second; + } else if (P.next_slot < kWvSlots) { + wv_slot = P.next_slot++; + P.slot_of.emplace(stream, wv_slot); + } + } + + if (wv_slot >= 0) { + glbl = P.glbl.data_ptr() + wv_slot * P.glbl_stride; + cntr = P.cntr.data_ptr() + wv_slot * P.cntr_stride; + } else { + spill_glbl = torch::zeros( + wv_gf, + torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device())); + spill_cntr = torch::zeros( + wv_ci, + torch::TensorOptions().dtype(torch::kInt).device(in_a.device())); + glbl = spill_glbl.data_ptr(); + cntr = spill_cntr.data_ptr(); + } + } #define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \ { \ dim3 block(64, 4); \ - if (_DTRMNSTC) \ + if constexpr (_DTRMNSTC) \ wvSplitKrc_ \ <<>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \ af4, bf4, biasf4, glbl, cntr, c, \ diff --git a/tests/kernels/quantization/test_rocm_skinny_gemms.py b/tests/kernels/quantization/test_rocm_skinny_gemms.py index f0de9c834512..7a7fe326bfa4 100644 --- a/tests/kernels/quantization/test_rocm_skinny_gemms.py +++ b/tests/kernels/quantization/test_rocm_skinny_gemms.py @@ -53,6 +53,17 @@ DTYPES = [torch.bfloat16, torch.float16] +# kWvSlots in csrc/rocm/skinny_gemms.cu; beyond it streams take overflow +# workspaces. +WVSPLITKRC_SLOTS = 8 + +# The CU count is the grid size. Production passes num_compute_units(); a small +# value manufactures better concurrency for these tests. +WVSPLITKRC_TEST_CU = 16 + +# Gate spin, ~125 ms on gfx950 -- must outlast the host enqueue loop it gates. +_GATE_CYCLES = 300_000_000 + # Specific (N, K, M) combinations for targeted testing NKM_FACTORS_LLMM1 = [ # Small, medium, large cases @@ -116,6 +127,15 @@ # M tiles are 64 rows, +16 for a partial tile M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16] +# (N, K, M) with more K-shards than the readback can stage in one LDS pass. +# Spans both CHUNKK values and both N-tile counts. +NKM_FACTORS_WVSPLITKRC_LARGE_K = [ + (128, 6144, 128), + (96, 8192, 128), + (128, 12288, 128), + (32, 12288, 128), +] + NKM_FACTORS_WVSPLITK_FP8 = [ # FP8-specific cases with K % 16 == 0 (1, 16, 16), @@ -222,6 +242,83 @@ def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, padded_a, bias_mode, xnorm, seed torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2) +@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC_LARGE_K) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") +@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950") +def test_rocm_wvsplitkrc_large_k(n, k, m, dtype, seed): + """K large enough that the split-K readback must stage LDS in batches.""" + torch.manual_seed(seed) + + xavier = math.sqrt(2 / k) + A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier + B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier + + ref_out = torch.nn.functional.linear(A, B, None) + out = ops.wvSplitKrc(A, B, num_compute_units(), None) + + torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("n_streams", [WVSPLITKRC_SLOTS, WVSPLITKRC_SLOTS + 4]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") +@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950") +def test_rocm_wvsplitkrc_multistream(n_streams, dtype, seed): + """Concurrent streams must not share a split-K workspace. + + The workspace is cleared by protocol rather than per invocation, so sharing + one corrupts the reduction and returns wrong numbers rather than failing. + n_streams > WVSPLITKRC_SLOTS additionally exercises the overflow path. + """ + torch.manual_seed(seed) + + n, k, m = 16, 2048, 256 + iters = 4 + xavier = math.sqrt(2 / k) + + B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier + As = [ + torch.randn(n, k, dtype=dtype, device="cuda") * xavier for _ in range(n_streams) + ] + refs = [torch.nn.functional.linear(A, B, None) for A in As] + torch.accelerator.synchronize() + + streams = [torch.Stream() for _ in range(n_streams)] + outs: list[list[torch.Tensor]] = [[] for _ in range(n_streams)] + _release_together(streams) + # Round-robin the enqueues so kernels from different streams are in flight + # at the same time. + for _ in range(iters): + for i, s in enumerate(streams): + with torch.cuda.stream(s): + outs[i].append(ops.wvSplitKrc(As[i], B, WVSPLITKRC_TEST_CU, None)) + for s in streams: + torch.accelerator.current_stream().wait_stream(s) + torch.accelerator.synchronize() + + for i in range(n_streams): + for out in outs[i]: + torch.testing.assert_close(out, refs[i], atol=1e-3, rtol=1e-2) + + +def _release_together(streams): + """Release every stream at once, once everything below is enqueued. + + An event recorded on an idle stream is already complete, so waiting on it + gates nothing -- the spin, not the event, is what holds the streams. + """ + blocker = torch.Stream() + with torch.cuda.stream(blocker): + torch.cuda._sleep(_GATE_CYCLES) + gate = torch.Event() + gate.record() + for s in streams: + s.wait_event(gate) + + @pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16]) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 92deff118398..fd068dff1cea 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -2357,6 +2357,12 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): from vllm.platforms import current_platform if not current_platform.is_cpu(): + from vllm.triton_utils import HAS_TRITON + + if HAS_TRITON: + from vllm.v1.sample.ops.topk_topp_triton import reset_buffer_cache + + reset_buffer_cache() torch.accelerator.empty_cache() try: torch.accelerator.empty_host_cache() diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index daaff3b79bf6..2868690ebeb1 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -283,24 +283,11 @@ def __init__( self._shared_experts: SharedExperts | None = None if shared_experts is not None: can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts - # When unquantized, shared expert inputs alias the hidden states, - # which can lead to race condition in multi-stream mode. Quantized - # routed experts copy the input into a fresh buffer first, breaking - # the alias, so overlap is safe. Only observed on ROCm. - routed_input_is_quantized = lambda: ( - self.routed_experts.quant_method.moe_quant_config is not None - and self.routed_experts.quant_method.moe_quant_config.quant_dtype - is not None - ) - is_multistream_safe = lambda: ( - not current_platform.is_rocm() or routed_input_is_quantized() - ) self._shared_experts = SharedExperts( shared_experts, moe_config=moe_config, enable_dbo=enable_dbo, mk_can_overlap_shared_experts=can_overlap, - is_multistream_safe=is_multistream_safe, ) # Needed for string -> MoERunner layer lookup in custom ops. diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 5b9587c8530b..4687cc5469ff 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -43,7 +43,6 @@ def __init__( moe_config: FusedMoEConfig, enable_dbo: bool, mk_can_overlap_shared_experts: Callable[[], bool], - is_multistream_safe: Callable[[], bool], ): super().__init__() @@ -58,10 +57,6 @@ def __init__( self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts - # Might not be safe to run multi-stream mode if routed and shared experts - # alias the same inputs - self._is_multistream_safe = is_multistream_safe - # Allow disabling of the separate shared experts stream for # debug purposes. # TODO: Remove this after more extensive testings with TP/DP @@ -112,19 +107,11 @@ def _determine_shared_experts_order( if self._mk_can_overlap_shared_experts(): return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED - # On ROCm, empirically only DP-only deployments benefit from the overlap. - overlap_is_beneficial = not current_platform.is_rocm() or ( - self._moe_config.moe_parallel_config.dp_size > 1 - and self._moe_config.moe_parallel_config.tp_size == 1 - ) - should_run_shared_in_aux_stream = ( current_platform.is_cuda_alike() and self._stream is not None and hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - and overlap_is_beneficial - and self._is_multistream_safe() ) if should_run_shared_in_aux_stream: diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 3e19c2b38976..f483a619b293 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -249,6 +249,45 @@ def use_aiter_triton_gemm(n, m, k, dtype): ) +def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, bool]: + """Pick the K-shard split for wvSplitKrc and say whether the shape fits. + + Mirrors wvSplitKrc() in csrc/rocm/skinny_gemms.cu, which is also where the + shard cap is explained. Both must pick the same chunkk or the workspace + check here bounds the wrong k_rnd. + + Returns: + The CHUNKK the kernel will dispatch with, and whether the CU budget and + split-K workspace admit the shape at all. + """ + # Next ^2 of n + N_p2 = 1 << (n - 1).bit_length() + # How many of 4 waves in a group can work on same 16 Ms at same time? + # This reduces the Ms each group works on, i.e. increasing the CUs needed. + GrpsShrB = min(N_p2 // 16, 4) + # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), and each + # working on a 512-shard of K, how many CUs would we need? + CuNeeded = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) * GrpsShrB + + CHUNKK2_MAX_SHARDS = 11 + shards_chunkk2 = (k + 256 - 1) // 256 # 256-wide shards + chunkk = ( + 2 + if ( + N_p2 != 16 + and CuNeeded * 2 <= cu_count + and shards_chunkk2 <= CHUNKK2_MAX_SHARDS + ) + else 1 + ) + + # Deterministic reduction stores one fp32 partial per (M, N, k-shard); all + # of them must fit the split-K workspace. + k_rnd = (k + 512 // chunkk - 1) // (512 // chunkk) + fits = N_p2 * m * k_rnd <= 128 * 1024 * 12 and CuNeeded <= cu_count + return chunkk, fits + + def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: @@ -260,21 +299,7 @@ def rocm_unquantized_gemm_impl( cu_count = num_compute_units() - # Next ^2 of n - N_p2 = 1 << (n - 1).bit_length() - # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), - # and each working on a 512-shard of K, how many CUs would we need? - rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) - # How many of 4 waves in a group can work on same 16 Ms at same time? - # This reduces the Ms each group works on, i.e. increasing the number of CUs needed. - GrpsShrB = min(N_p2 // 16, 4) - # Given the above, how many CUs would we need? - CuNeeded = rndup_cus * GrpsShrB - # Deterministic reduction stores one float workspace value per K shard. - fits_wvsplitkrc = ( - N_p2 * m * ((k + 512 - 1) // 512) - ) <= 128 * 1024 * 12 # deterministic - fits_wvsplitkrc &= CuNeeded <= cu_count + _, fits_wvsplitkrc = wvsplitkrc_dispatch(n, k, m, cu_count) skinny_operands_compatible = weight.is_contiguous() and ( bias is None or bias.is_contiguous() @@ -362,12 +387,18 @@ def rocm_unquantized_gemm( @functools.cache def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None: - """Eagerly allocate wvSplitKrc's process-lifetime static workspaces. + """Eagerly allocate wvSplitKrc's per-device split-K workspace pool. + + wvSplitKrc partitions one per-device allocation into ``kWvSlots`` slots + (csrc/rocm/skinny_gemms.cu) and hands each stream one on first use, so that + two streams never share the split-K partials and counters. - They are otherwise created lazily on the first qualifying GEMM + The pool is otherwise created lazily on the first qualifying GEMM (csrc/rocm/skinny_gemms.cu), which can be the first real request — after - the KV cache backing buffer exists. If one landed in that segment's - rounding tail, it would pin the entire segment at engine shutdown. + the KV cache backing buffer exists. If it landed in that segment's rounding + tail, it would pin the entire segment at engine shutdown; it could also land + inside a cudagraph capture, where it would be taken from the graph's private + pool and its zero-fill would become a replayed graph node. """ from vllm.platforms.rocm import on_gfx950 From 360f33f61ce02bda04bcb7ea876f1d2e5af8d06d Mon Sep 17 00:00:00 2001 From: Colin McNamara Date: Thu, 10 Sep 2026 17:21:21 -0500 Subject: [PATCH 013/774] [Bugfix] Set stop_sequence explicitly in streaming message_delta event (#55325) Signed-off-by: Colin McNamara --- .../test_anthropic_messages_conversion.py | 26 +++++++++++++++++++ vllm/entrypoints/anthropic/serving.py | 4 ++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 436db76a9acd..0c986f129139 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1610,6 +1610,32 @@ async def sse_input(): assert msg_deltas[0]["delta"]["stop_reason"] == "stop_sequence" assert msg_deltas[0]["delta"]["stop_sequence"] == "" + @pytest.mark.asyncio + async def test_streaming_no_stop_string_emits_explicit_null_stop_sequence(self): + """exclude_unset=True drops stop_sequence unless it is set explicitly.""" + + async def sse_input(): + yield _make_stream_chunk(delta=DeltaMessage(role="assistant")) + yield _make_stream_chunk(delta=DeltaMessage(content="hi")) + yield _make_stream_chunk(finish_reason="stop") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + + events = _parse_sse_events(output) + msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"] + assert len(msg_deltas) == 1 + assert msg_deltas[0]["delta"]["stop_reason"] == "end_turn" + assert "stop_sequence" in msg_deltas[0]["delta"] + assert msg_deltas[0]["delta"]["stop_sequence"] is None + # ====================================================================== # Client-caused errors are 4xx, not 500 (Issue #52088) diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index fe0d752d4ecc..f8784284b49c 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -842,7 +842,9 @@ def stop_and_flush() -> list[str]: stop_delta = AnthropicDelta( stop_reason=self.stop_reason_map.get( finish_reason or "stop" - ) + ), + # Set explicitly so exclude_unset=True keeps it. + stop_sequence=None, ) chunk = AnthropicStreamEvent( type="message_delta", From a89de950a43896cfd97a3b7f051a8f0984ea6be0 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Thu, 10 Sep 2026 19:26:11 -0400 Subject: [PATCH 014/774] [CI/Build] Pin HyperCLOVAX V2 test model revision (#56335) Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex --- tests/models/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/models/registry.py b/tests/models/registry.py index c3973c4d0b94..a474a5a0fb72 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -963,6 +963,7 @@ def check_available_online( "HCXVisionV2ForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B", trust_remote_code=True, + revision="a6cdfd3464d1b767259cad23e164eaf39d3e3960", ), "HunYuanVLForConditionalGeneration": _HfExamplesInfo( "tencent/HunyuanOCR", From 5b6cf93e8e229aa0901e8c152eba4511a6195d28 Mon Sep 17 00:00:00 2001 From: Marceli Fylcek Date: Fri, 11 Sep 2026 01:58:28 +0200 Subject: [PATCH 015/774] [XPU] Add forward_xpu to Mixer2RMSNormGated and FusedRMSNormGated (#54968) Signed-off-by: Marceli Fylcek Co-authored-by: Claude Opus 5 --- vllm/model_executor/layers/mamba/mamba_mixer2.py | 7 +++++++ .../flash_linear_attention/ops/kda.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 5bc890d0078b..1496515e7735 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -180,6 +180,13 @@ def forward_cuda( norm_before_gate=False, ) + def forward_xpu( + self, + x: torch.Tensor, + gate: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self.forward_cuda(x, gate) + def mamba_v2_sharded_weight_loader( shard_spec: list[tuple[int, int, float]], diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 1701cd47aca8..548f600812e2 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -529,6 +529,22 @@ def forward_cuda( residual_in_fp32=residual_in_fp32, ) + def forward_xpu( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return self.forward_cuda( + x, + g, + residual=residual, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @triton.autotune( From 84030bbe3d74d99bad477a3d2e37a973ccd8865c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=84=A6=E6=81=BA=E6=96=87?= <15529091557@163.com> Date: Fri, 11 Sep 2026 08:40:10 +0800 Subject: [PATCH 016/774] [Bugfix][Core] Stop zero-progress preemption cascades for deferred KV frees (#49675) Signed-off-by: LearningMachine621 <147347705+LearningMachine621@users.noreply.github.com> --- tests/v1/core/test_async_scheduler.py | 54 ++++++++++++----- tests/v1/core/test_deferred_block_free.py | 74 ++++++++++++++++++++++- tests/v1/core/utils.py | 3 + vllm/v1/core/sched/scheduler.py | 26 +++++--- 4 files changed, 130 insertions(+), 27 deletions(-) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7f256f025e6d..0a21e896b030 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -701,16 +701,29 @@ def test_requires_kv_delivery_defaults_to_producer_role(): assert scheduler.requires_kv_delivery is expected, role -@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"]) -def test_kv_pressure_preempt_mid_handoff(kv_role: str): - """P/D race: a request is KV-pressure preempted while the output of its - final prefill chunk -- the hand-off token that would finish it -- is in - flight. - - On a producer, that output must be dropped so the request recomputes; - delivering it would finish the request and hand off blocks the preemption - already freed, so the consumer pulls garbage. A consumer hands nothing off, - so it keeps the lossless deliver-stale path. +@pytest.mark.parametrize( + ("kv_role", "defer_free"), + [ + ("kv_producer", False), + ("kv_consumer", False), + ("kv_consumer", True), + ], +) +def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool): + """P/D race: KV pressure hits while the output of a request's final + prefill chunk -- the hand-off token that would finish it -- is in flight. + + When the victim's blocks free immediately, the request is preempted. On a + producer, that output must be dropped so the request recomputes; delivering + it would finish the request and hand off blocks the preemption already + freed, so the consumer pulls garbage. A consumer hands nothing off, so it + keeps the lossless deliver-stale path. + + A consumer with overlapping batches (async scheduling or PP) instead fences + the victim's free behind its in-flight output, so the allocation retry + stops instead of preempting; the request then finishes from that output + once it lands. The gate depends on the platform (async scheduling is + force-disabled on CPU), so force the flag to cover both paths everywhere. """ is_producer = kv_role == "kv_producer" scheduler = create_scheduler( @@ -722,10 +735,13 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str): max_num_batched_tokens=512, ) assert scheduler.requires_kv_delivery is is_producer + # The production gate requires overlapping batches, which async + # scheduling only provides off-CPU. + scheduler.defer_block_free = defer_free # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the - # next decode allocation preempts the tail of the running queue (the handoff - # request) while its prefill output is still in flight. + # next decode allocation targets the tail of the running queue (the + # handoff request) while its prefill output is still in flight. decoder = create_requests( num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"] )[0] @@ -739,9 +755,17 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str): assert handoff.num_output_placeholders == 1 scheduler.schedule() - assert handoff.status == RequestStatus.PREEMPTED - assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens - assert handoff.drop_stale_output is is_producer + if defer_free: + # The victim's blocks are fenced behind its in-flight output, so + # preempting it could not satisfy the allocation; the retry stops + # instead of preempting. + assert handoff.status == RequestStatus.RUNNING + assert handoff.num_stale_output_tokens == 0 + else: + # Blocks free immediately, so the handoff request is preempted. + assert handoff.status == RequestStatus.PREEMPTED + assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens + assert handoff.drop_stale_output is is_producer scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py index 13789a396015..ebbf492e208d 100644 --- a/tests/v1/core/test_deferred_block_free.py +++ b/tests/v1/core/test_deferred_block_free.py @@ -48,13 +48,17 @@ def _make_model_runner_output( ) -def _create_deferring_scheduler(): +def _create_deferring_scheduler(scheduling_policy="fcfs"): """Async scheduler with deferred block freeing forced on. The production gate additionally requires a PD KV-consumer connector; the mechanism itself is independent of it. """ - scheduler = create_scheduler(model=MODEL, async_scheduling=True) + scheduler = create_scheduler( + model=MODEL, + async_scheduling=True, + scheduling_policy=scheduling_policy, + ) scheduler.defer_block_free = True return scheduler @@ -79,6 +83,24 @@ def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5): return request, out0, out1 +def _fail_one_allocation(scheduler, call_number: int): + real_allocate = scheduler.kv_cache_manager.allocate_slots + calls = 0 + + def allocate(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == call_number: + return None + return real_allocate(*args, **kwargs) + + return patch.object( + scheduler.kv_cache_manager, + "allocate_slots", + side_effect=allocate, + ) + + def test_gate_enabled_for_async_consumer(): # Overlapping batches + consumer-side connector enables the gate. Async # scheduling (which would give >1 concurrent batches) is force-disabled on @@ -224,6 +246,54 @@ def test_preempt_defers_free_and_clears_bookkeeping(): assert pool.get_num_free_blocks() == num_free_initially +@pytest.mark.parametrize( + ("policy", "failed_call"), + [ + ("fcfs", 1), + ("priority", 2), + ], +) +def test_allocation_retry_waits_for_fence_then_succeeds(policy, failed_call): + scheduler = _create_deferring_scheduler(policy) + requests = create_requests( + num_requests=3, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=10, + stop_token_ids=[STOP_TOKEN_ID], + ) + for request in requests: + scheduler.add_request(request) + out0 = scheduler.schedule() + out1 = scheduler.schedule() + + worst, trigger, tail = requests + worst.priority, trigger.priority, tail.priority = 9, 0, 1 + + with _fail_one_allocation(scheduler, failed_call) as allocate: + blocked = scheduler.schedule() + assert allocate.call_count == failed_call + assert not blocked.preempted_req_ids + retry_trigger = requests[failed_call - 1] + assert retry_trigger.request_id not in blocked.num_scheduled_tokens + + for output in (out0, out1, blocked): + if output.total_num_scheduled_tokens: + scheduler.update_from_output( + output, + _make_model_runner_output(output), + ) + + expected_victim = worst if policy == "priority" else tail + with _fail_one_allocation(scheduler, failed_call) as allocate: + resumed = scheduler.schedule() + + assert allocate.call_count > failed_call + assert expected_victim.request_id in resumed.preempted_req_ids + assert retry_trigger.request_id in resumed.num_scheduled_tokens + if policy == "priority": + assert tail.request_id in resumed.num_scheduled_tokens + + def test_multiple_deferred_frees_drain_in_order(): scheduler = _create_deferring_scheduler() pool = scheduler.kv_cache_manager.block_pool diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index d3090e2ddace..3c4a8960a3fa 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -17,6 +17,7 @@ SpeculativeConfig, VllmConfig, ) +from vllm.config.scheduler import SchedulerPolicy from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, @@ -74,6 +75,7 @@ def create_scheduler( use_v2_model_runner: bool | None = None, kv_cache_spec: KVCacheSpec | None = None, per_request_spec_decode_metrics: str = "none", + scheduling_policy: SchedulerPolicy = "fcfs", ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -113,6 +115,7 @@ def create_scheduler( is_encoder_decoder=model_config.is_encoder_decoder, # Ensure admission/preemption mechanics are deterministic watermark=0.0, + policy=scheduling_policy, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 01423212999f..8b2a551b5579 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -745,13 +745,16 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.running, key=lambda r: (r.priority, r.arrival_time), ) - # Record the index of the preemption victim to - # maintain accurate loop state. + else: + preempted_req = self.running[-1] + + # A deferred free will not help with immediate allocation. + if not self._request_blocks_can_be_freed(preempted_req): + break + + if self.policy == SchedulingPolicy.PRIORITY: victim_index = self.running.index(preempted_req) del self.running[victim_index] - # Decrement the loop cursor if the removed request - # preceded the current iteration, preventing the - # silent omission of the subsequent request. if victim_index < req_index: req_index -= 1 @@ -2605,15 +2608,18 @@ def set_pause_state(self, pause_state: PauseState) -> None: logger.info("setting pause state to %s", pause_state.name) self._pause_state = pause_state + def _request_blocks_can_be_freed(self, request: Request) -> bool: + # We must defer freeing blocks if an async kv connector may + # write to them immediately (not ordered with GPU stream). + return not self.defer_block_free or ( + request.last_sched_seq <= self.processed_step_seq + ) + def _free_request_blocks(self, request: Request): """Free the request's KV blocks, deferring the return to the block pool when an in-flight GPU step may still write them. """ - if not self.defer_block_free or ( - # Last scheduled step already processed: no in-flight write remains - # (always the case for a normal finish), so free now. - request.last_sched_seq <= self.processed_step_seq - ): + if self._request_blocks_can_be_freed(request): self.kv_cache_manager.free(request) return blocks = self.kv_cache_manager.pop_blocks_for_free(request) From 828f4f19b4d8ea6a97a047409a90b563d166002f Mon Sep 17 00:00:00 2001 From: James E T Smith Date: Thu, 10 Sep 2026 21:13:14 -0400 Subject: [PATCH 017/774] [ROCm][Bugfix] Route GLM-5.3-Flash MTP through ragged sparse MLA (#55239) Signed-off-by: jamesETsmith Signed-off-by: Cursor Co-authored-by: Cursor Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_rocm_glm5next_sparse.py | 5 ++++- .../attention/backends/mla/rocm_aiter_mla_sparse.py | 11 +++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index 509d8f1faea4..ccc9e5c6178e 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -76,7 +76,9 @@ def test_fit_kpool_indices_rejects_narrow_input(): ("auto", 512, 0, 2, 2, 1, True), ("fp8", 512, 1, 0, 0, 32, False), ("auto", 576, 1, 0, 0, 32, False), - ("auto", 512, 0, 2, 4, 2, False), + ("auto", 512, 0, 2, 4, 2, True), + ("auto", 512, 0, 2, 12, 6, True), + ("auto", 512, 0, 0, 0, 0, False), ], ) def test_rocm_sparse_triton_route( @@ -88,6 +90,7 @@ def test_rocm_sparse_triton_route( max_query_len, expected, ): + """Validate Triton routing for prefill, decode, and MTP verification.""" assert ( _use_rocm_sparse_triton( kv_cache_dtype=kv_cache_dtype, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 3a9a4c92b350..45a413aab381 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -53,13 +53,16 @@ def _use_rocm_sparse_triton( num_decode_tokens: int, max_query_len: int, ) -> bool: - """Select the rope-free BF16 path not supported by AITER sparse MLA.""" - plain_decode = num_decode_tokens == num_decodes + """Select the rope-free BF16 path not supported by AITER sparse MLA. + + The ragged Triton kernel indexes metadata per query token, so multi-token + speculative verification rows have the same capability requirements as + plain decode rows. + """ return ( not kv_cache_dtype.startswith("fp8") and head_size == kv_lora_rank - and plain_decode - and (num_prefills > 0 or (num_decodes > 0 and max_query_len == 1)) + and (num_prefills > 0 or num_decodes > 0) ) From 07b75534651a0dd584c6f3be75d8c62fb8426397 Mon Sep 17 00:00:00 2001 From: YukioZzz Date: Fri, 11 Sep 2026 09:36:37 +0800 Subject: [PATCH 018/774] [Feature][SimpleCPU] Load fine-grained hybrid prefix hits (#54736) Signed-off-by: Yichao Zhu --- tests/v1/simple_kv_offload/test_kv_events.py | 49 +- tests/v1/simple_kv_offload/test_scheduler.py | 672 ++++++++++++++++++ .../v1/simple_cpu_offload_connector.py | 7 + vllm/v1/core/kv_cache_coordinator.py | 16 +- vllm/v1/core/single_type_kv_cache_manager.py | 16 + vllm/v1/simple_kv_offload/manager.py | 332 +++++++-- 6 files changed, 982 insertions(+), 110 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index 3bc6c32e7ce0..f21dae7e0bd8 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -714,17 +714,8 @@ def test_disk_mode_block_removed_relabeled_storage() -> None: assert ev.medium == MEDIUM_STORAGE -def test_mamba_dcp_metadata_guard() -> None: - """Mamba+DCP: token_ids guard emits [] when capture and emission sizes differ. - - The metadata capture at _prepare_eager_store_specs uses - g_block_size = spec.block_size * cp_world_size = 32 * 2 = 64 to slice - tokens. The event emission uses block_size = spec.block_size * 1 = 32 - for Mamba (SSM state is not CP-sharded). The guard in - _process_store_completion detects len(token_ids) != block_size and emits - token_ids=[] instead of misleading data. When #49962 fixes the capture - path, the guard stops triggering and correct token_ids flow through. - """ +def test_mamba_align_skips_positional_event_metadata() -> None: + """Mamba align emits events only from explicit boundary handoffs.""" vbs = BLOCK_SIZE * 2 kv_cfg = _make_mixed_kv_cache_config( num_blocks=16, with_mamba=True, mamba_block_size=32 @@ -739,8 +730,6 @@ def test_mamba_dcp_metadata_guard() -> None: sched = fix.scheduler gpu_pool = fix.gpu_block_pool - # Mamba g_block_size = 32 * 2 = 64, so 2 Mamba blocks need 128 computed - # tokens to be "ready"; FA g_block_size = 16 * 2 = 32 needs only 64. req = _make_cp_request(num_blocks=4, virtual_block_size=vbs) fa_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=0) mamba_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=1) @@ -759,32 +748,12 @@ def test_mamba_dcp_metadata_guard() -> None: events = list(sched.take_events()) stored = [e for e in events if isinstance(e, BlockStored)] - assert len(stored) == 4, ( - f"expected 4 BlockStored (2 FA + 2 Mamba), got {len(stored)}" - ) - - by_group: dict[int, list[BlockStored]] = {} - for ev in stored: - by_group.setdefault(ev.group_idx, []).append(ev) - - # FA group: capture size (64) == emission size (64) → token_ids present. - fa_events = by_group[0] - assert len(fa_events) == 2 - for ev in fa_events: + assert len(stored) == 2 + assert {ev.group_idx for ev in stored} == {0} + for idx, ev in enumerate(stored): assert ev.block_size == vbs - assert len(ev.token_ids) == vbs, ( - f"FA token_ids should be {vbs}, got {len(ev.token_ids)}" - ) - - # Mamba group: capture size (64) != emission size (32) → guard emits []. - mamba_events = by_group[1] - assert len(mamba_events) == 2 - for ev in mamba_events: - assert ev.block_size == 32 - assert ev.token_ids == [], ( - f"Mamba token_ids should be [] (guard), got {ev.token_ids}" - ) - assert ev.parent_block_hash is None, ( - "Mamba parent_block_hash should be None (guard)" + assert ev.token_ids == req.prompt_token_ids[idx * vbs : (idx + 1) * vbs] + expected_hash = maybe_convert_block_hash( + get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0)) ) - assert ev.extra_keys is None, "Mamba extra_keys should be None (guard)" + assert ev.block_hashes == [expected_hash] diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index db7e8b7aa51c..1952e3a4b0b3 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -4,6 +4,7 @@ from __future__ import annotations +import logging from dataclasses import dataclass import pytest @@ -18,6 +19,7 @@ SchedulerConfig, VllmConfig, ) +from vllm.config.kv_events import KVEventsConfig from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) @@ -31,6 +33,7 @@ ) from vllm.v1.core.sched.output import ( CachedRequestData, + KVConnectorBlockState, NewRequestData, SchedulerOutput, ) @@ -42,6 +45,7 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MambaSpec, SlidingWindowSpec, ) from vllm.v1.outputs import KVConnectorOutput @@ -2128,3 +2132,671 @@ def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None: f"cp_world_size={cp_world_size}: target_cp={target_cp} should be " f"less than target_base={target_base}" ) + + +def _make_hybrid_attention_mamba_scheduler( + *, + num_cpu_blocks: int = 8, + num_gpu_blocks: int = 16, + attention_block_size: int = BLOCK_SIZE, + block_size: int = 4 * BLOCK_SIZE, + scheduler_block_size: int | None = None, + hash_block_size: int | None = None, + dcp_world_size: int = 4, + lazy: bool = False, + mamba_cache_mode: str = "align", + enable_kv_cache_events: bool = False, +) -> SchedulerFixture: + """Build a scheduler for one attention group plus one Mamba group.""" + scheduler_block_size = scheduler_block_size or block_size + hash_block_size = hash_block_size or block_size + attention_spec = FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + mamba_spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode=mamba_cache_mode, + ) + groups = [ + KVCacheGroupSpec(["attention"], attention_spec), + KVCacheGroupSpec(["mamba"], mamba_spec), + ] + tensors = [ + KVCacheTensor( + size=spec.page_size_bytes * num_gpu_blocks, + layers=group.layer_names, + layer_stride=spec.page_size_bytes * num_gpu_blocks, + block_stride=spec.page_size_bytes, + ) + for group, spec in zip(groups, (attention_spec, mamba_spec)) + ] + kv_cache_config = KVCacheConfig( + num_blocks=num_gpu_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + vllm_config = _make_cp_vllm_config(dcp_world_size=dcp_world_size) + vllm_config.cache_config.prefix_cache_retention_interval = 0 + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + if enable_kv_cache_events: + vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True) + # _derive_cpu_config() scales the requested capacity against + # kv_cache_tensors[0].size, so express it as a multiple of that tensor's + # per-block size to get exactly num_cpu_blocks. + cpu_capacity_bytes = tensors[0].size // num_gpu_blocks * num_cpu_blocks + sched = SimpleCPUOffloadScheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=cpu_capacity_bytes, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + lazy_offload=lazy, + ) + gpu_block_pool = BlockPool( + num_gpu_blocks=num_gpu_blocks, + enable_caching=True, + hash_block_size=hash_block_size, + ) + sched.bind_gpu_block_pool(gpu_block_pool) + return SchedulerFixture( + scheduler=sched, + gpu_block_pool=gpu_block_pool, + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + ) + + +def test_hybrid_store_uses_resolved_group_block_sizes() -> None: + """Replicated groups must not be scaled by the DCP world size. + + ``dcp_world_size_for_kv_cache_spec`` gives full attention the process DCP + size and every other spec 1. Reconstructing a group's geometry as + ``spec.block_size * cp_world_size`` over-scales the replicated groups, so + the eager store scan believes far fewer of their blocks are ready and + silently offloads only a fraction of them. + """ + attention_block_size = BLOCK_SIZE + mamba_block_size = 4 * BLOCK_SIZE + dcp_world_size = 2 + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=32, + num_gpu_blocks=32, + attention_block_size=attention_block_size, + block_size=mamba_block_size, + # Every prefix-cacheable group's resolved block size must be a multiple + # of the hash block, so it cannot exceed the smallest of them. + hash_block_size=BLOCK_SIZE, + dcp_world_size=dcp_world_size, + mamba_cache_mode="all", + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + attention_size, mamba_size = sched.group_block_sizes + assert attention_size == attention_block_size * dcp_world_size + assert mamba_size == mamba_block_size + + confirmed = 2 * sched.block_size + req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) + attention_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 + ) + mamba_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 + ) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched.build_connector_meta( + make_scheduler_output( + {req.request_id: confirmed}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + ) + req.num_computed_tokens = confirmed + meta = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 1}, + cached_req_new_blocks={req.request_id: None}, + ) + ) + + # Both groups are positionally stable here, so every confirmed block of + # both must be offloaded. Over-scaling the replicated group halves its + # count. + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ + confirmed // attention_size, + confirmed // mamba_size, + ] + assert set(meta.store_gpu_blocks) == { + *(b.block_id for b in attention_blocks), + *(b.block_id for b in mamba_blocks), + } + + +def test_dcp_mixed_cache_loads_fine_grained_external_hit() -> None: + """Map a fine-grained external suffix after a nonzero local prefix.""" + hash_block_size = BLOCK_SIZE + mamba_block_size = 4 * BLOCK_SIZE + local_tokens = 4 * BLOCK_SIZE + # Span two Mamba blocks and end partially in the second one. This exercises + # both cdiv(external_tokens, mamba_block_size) > 1 and the partial tail. + external_tokens = 6 * BLOCK_SIZE + total_cached_tokens = local_tokens + external_tokens + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=32, + num_gpu_blocks=32, + block_size=mamba_block_size, + hash_block_size=hash_block_size, + dcp_world_size=2, + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + + assert external_tokens % mamba_block_size != 0 + assert sched.cpu_coordinator.enable_partial_hash_hits + + producer = _make_cp_request(num_blocks=10, virtual_block_size=hash_block_size) + attention_blocks = _allocate_cp_gpu_blocks( + gpu_pool, + producer, + num_blocks=5, + virtual_block_size=2 * BLOCK_SIZE, + group_id=0, + ) + mamba_blocks = gpu_pool.get_new_blocks(3) + gpu_pool.cache_full_blocks( + request=producer, + blocks=mamba_blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=mamba_block_size, + kv_cache_group_id=1, + ) + gpu_pool.cache_partial_block( + request=producer, + block=mamba_blocks[2], + num_tokens=total_cached_tokens, + kv_cache_group_id=1, + block_size=mamba_block_size, + ) + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) + + # Positional scanning stores append-only FA blocks. Mamba align blocks use + # exact handoffs for every retained boundary, including the partial tail. + store_output = make_scheduler_output( + {producer.request_id: total_cached_tokens}, + new_reqs={producer.request_id: producer_blocks.get_block_ids()}, + ) + store_output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids=lambda _: (), + boundary_state_offloads={ + producer.request_id: [ + (1, mamba_blocks[0].block_id, mamba_block_size), + (1, mamba_blocks[1].block_id, 2 * mamba_block_size), + (1, mamba_blocks[2].block_id, total_cached_tokens), + ] + }, + ) + store_meta = sched.build_connector_meta(store_output) + # The exact handoffs land in the step they arrive; the append-only FA scan + # picks its blocks up once the scheduler has committed the tokens. + assert set(store_meta.store_gpu_blocks) == { + block.block_id for block in mamba_blocks + } + simulate_store_completion(sched, store_meta.store_event) + producer.num_computed_tokens = total_cached_tokens + fa_meta = sched.build_connector_meta( + make_scheduler_output( + {producer.request_id: 1}, + cached_req_new_blocks={producer.request_id: None}, + ) + ) + assert set(fa_meta.store_gpu_blocks) == { + block.block_id for block in attention_blocks + } + simulate_store_completion(sched, fa_meta.store_event) + + consumer = Request( + request_id="req-dcp-mixed-fine-load", + prompt_token_ids=producer.prompt_token_ids, + sampling_params=producer.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=producer._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens( + consumer, num_computed_tokens=local_tokens + ) + assert hit_tokens == external_tokens + assert is_async is True + pending_blocks, pending_tokens, _ = sched._pending_cpu_hits[consumer.request_id] + assert pending_tokens == external_tokens + assert len(pending_blocks[1]) == 2 + assert pending_blocks[1][0].is_null + assert not pending_blocks[1][1].is_null + + local_attention = attention_blocks[:2] + local_mamba = mamba_blocks[:1] + load_blocks = KVCacheBlocks( + blocks=( + local_attention + gpu_pool.get_new_blocks(3), + local_mamba + gpu_pool.get_new_blocks(2), + ) + ) + sched.update_state_after_alloc( + consumer, + load_blocks, + num_external_tokens=hit_tokens, + ) + load_meta = sched.build_connector_meta( + make_scheduler_output( + {consumer.request_id: 1}, + new_reqs={consumer.request_id: load_blocks.get_block_ids()}, + ) + ) + assert load_meta.load_gpu_blocks == [ + *(block.block_id for block in load_blocks.blocks[0][2:5]), + load_blocks.blocks[1][2].block_id, + ] + assert len(load_meta.load_cpu_blocks) == 4 + + +def test_finished_eager_store_caches_fa_partial_tail_for_hybrid_hit() -> None: + """A finished non-LCM prompt needs the FA partial tail in CPU too.""" + hash_block_size = BLOCK_SIZE + mamba_block_size = BLOCK_SIZE + attention_block_size = 4 * BLOCK_SIZE + dcp_world_size = 2 + scheduler_block_size = attention_block_size * dcp_world_size + total_cached_tokens = 5 * hash_block_size + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=32, + num_gpu_blocks=32, + attention_block_size=attention_block_size, + block_size=mamba_block_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + dcp_world_size=dcp_world_size, + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + assert sched.fa_block_size == scheduler_block_size + assert total_cached_tokens % sched.fa_block_size != 0 + + producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size) + attention_blocks = gpu_pool.get_new_blocks(1) + gpu_pool.cache_partial_block( + request=producer, + block=attention_blocks[0], + num_tokens=total_cached_tokens, + kv_cache_group_id=0, + block_size=scheduler_block_size, + ) + mamba_blocks = gpu_pool.get_new_blocks(5) + gpu_pool.cache_full_blocks( + request=producer, + blocks=mamba_blocks, + num_cached_blocks=0, + num_full_blocks=5, + kv_cache_group_id=1, + block_size=mamba_block_size, + ) + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) + + store_output = make_scheduler_output( + {producer.request_id: total_cached_tokens}, + new_reqs={producer.request_id: producer_blocks.get_block_ids()}, + ) + store_output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids=lambda _: (), + boundary_state_offloads={ + producer.request_id: [ + (1, mamba_blocks[4].block_id, total_cached_tokens), + ] + }, + ) + store_meta = sched.build_connector_meta(store_output) + simulate_store_completion(sched, store_meta.store_event) + + producer.num_computed_tokens = total_cached_tokens + sched.request_finished_all_groups(producer, producer_blocks.get_block_ids()) + finish_meta = sched.build_connector_meta(make_scheduler_output({})) + assert attention_blocks[0].block_id in finish_meta.store_gpu_blocks + simulate_store_completion(sched, finish_meta.store_event) + + consumer = Request( + request_id="req-finished-partial-tail-load", + prompt_token_ids=producer.prompt_token_ids, + sampling_params=producer.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=producer._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens( + consumer, num_computed_tokens=0 + ) + assert hit_tokens == total_cached_tokens + assert is_async is True + + +def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> None: + """Decode can fill the boundary block; the scan then already covers it.""" + hash_block_size = BLOCK_SIZE + mamba_block_size = BLOCK_SIZE + attention_block_size = 4 * BLOCK_SIZE + dcp_world_size = 2 + scheduler_block_size = attention_block_size * dcp_world_size + prompt_tokens = 5 * hash_block_size + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=32, + num_gpu_blocks=32, + attention_block_size=attention_block_size, + block_size=mamba_block_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + dcp_world_size=dcp_world_size, + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + assert prompt_tokens % sched.fa_block_size != 0 + + producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size) + attention_blocks = gpu_pool.get_new_blocks(1) + gpu_pool.cache_partial_block( + request=producer, + block=attention_blocks[0], + num_tokens=prompt_tokens, + kv_cache_group_id=0, + block_size=scheduler_block_size, + ) + mamba_blocks = gpu_pool.get_new_blocks(5) + gpu_pool.cache_full_blocks( + request=producer, + blocks=mamba_blocks, + num_cached_blocks=0, + num_full_blocks=5, + kv_cache_group_id=1, + block_size=mamba_block_size, + ) + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) + + # Decode carries the request past the physical attention block boundary, so + # the positional scan now considers the same block the partial tail names. + producer.num_computed_tokens = sched.fa_block_size + sched.request_finished_all_groups(producer, producer_blocks.get_block_ids()) + finish_meta = sched.build_connector_meta(make_scheduler_output({})) + + stored = finish_meta.store_gpu_blocks + assert attention_blocks[0].block_id in stored + assert len(stored) == len(set(stored)) + assert len(finish_meta.store_cpu_blocks) == len(stored) + + +def test_external_lookup_rejects_misaligned_local_prefix(caplog_vllm) -> None: + scheduler_block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler(block_size=scheduler_block_size) + + request = _make_cp_request(num_blocks=2, virtual_block_size=scheduler_block_size) + # The ``vllm`` logger sets propagate=False, so plain ``caplog`` can observe + # nothing; ``caplog_vllm`` re-enables propagation for the duration. + with caplog_vllm.at_level(logging.WARNING): + hit_tokens, is_async = fix.scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=BLOCK_SIZE + ) + + assert (hit_tokens, is_async) == (0, False) + assert "requires scheduler-block-aligned local tokens" in caplog_vllm.text + + +def test_fine_grained_external_hits_require_eager_offload() -> None: + hash_block_size = BLOCK_SIZE + scheduler_block_size = 4 * BLOCK_SIZE + + eager = _make_hybrid_attention_mamba_scheduler( + block_size=scheduler_block_size, + hash_block_size=hash_block_size, + ) + lazy = _make_hybrid_attention_mamba_scheduler( + block_size=scheduler_block_size, + hash_block_size=hash_block_size, + lazy=True, + ) + + assert eager.scheduler.cpu_coordinator.enable_partial_hash_hits + assert not lazy.scheduler.cpu_coordinator.enable_partial_hash_hits + + +def test_eager_store_does_not_scan_mamba_blocks_positionally() -> None: + """Mamba stores must use exact handoffs, not historical block positions.""" + block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers + assert attention_manager.has_positionally_stable_blocks + assert not mamba_manager.has_positionally_stable_blocks + + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + attn_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 + ) + gpu_blocks = gpu_pool.get_new_blocks(2) + gpu_blocks[1].set_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 1)) + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, gpu_blocks)) + req.num_computed_tokens = 0 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + output = make_scheduler_output( + {req.request_id: 2 * block_size}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids=lambda _: (), + boundary_state_offloads={ + req.request_id: [(1, gpu_blocks[1].block_id, 2 * block_size)] + }, + ) + meta1 = sched.build_connector_meta(output) + assert meta1.store_event >= 0 + # Only the exact handoff lands this step. build_connector_meta runs before + # the scheduler commits this step's tokens, so the positional scan has + # nothing confirmed yet. + assert set(meta1.store_gpu_blocks) == {gpu_blocks[1].block_id} + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [0, 0] + simulate_store_completion(sched, meta1.store_event) + + # Next step the attention blocks are confirmed and scanned positionally. A + # later hash update still must not make the mamba align snapshot a valid + # source; the explicit boundary-handoff tests cover the supported path. + gpu_blocks[0].set_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 1)) + req.num_computed_tokens = 2 * block_size + meta2 = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 1}, + cached_req_new_blocks={req.request_id: None}, + ) + ) + assert set(meta2.store_gpu_blocks) == {block.block_id for block in attn_blocks} + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0] + simulate_store_completion(sched, meta2.store_event) + + # Nothing further: mamba align is never picked up positionally. + meta3 = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 1}, + cached_req_new_blocks={req.request_id: None}, + ) + ) + assert meta3.store_event == -1 + assert meta3.store_gpu_blocks == [] + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0] + + +def test_eager_store_scans_mamba_all_blocks_positionally() -> None: + """Mamba all-mode retains positional identity and needs no handoff.""" + block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=16, + num_gpu_blocks=24, + block_size=block_size, + mamba_cache_mode="all", + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers + assert attention_manager.has_positionally_stable_blocks + assert mamba_manager.has_positionally_stable_blocks + + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + attention_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 + ) + mamba_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, 2, virtual_block_size=block_size, group_id=1 + ) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + # build_connector_meta runs before the scheduler commits this step's + # tokens, so the positional scan picks the blocks up on the next step. + first = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 2 * block_size}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + ) + assert first.store_gpu_blocks == [] + req.num_computed_tokens = 2 * block_size + meta = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: 1}, + cached_req_new_blocks={req.request_id: None}, + ) + ) + + assert set(meta.store_gpu_blocks) == { + *(block.block_id for block in attention_blocks), + *(block.block_id for block in mamba_blocks), + } + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 2] + + +@pytest.mark.parametrize("gone_via", ["preempted", "finished", "unregistered"]) +def test_boundary_handoff_dropped_for_departing_request(gone_via: str) -> None: + """Handoffs for a request leaving this step must not be read. + + The scheduler drains boundary handoffs without filtering by request + liveness, so a preempted or finished request can still offer block IDs. + Those blocks are being released and may already back another request, so + copying them would publish unrelated KV under a valid hash. + """ + block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + attn_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 + ) + mamba_blocks = gpu_pool.get_new_blocks(2) + mamba_blocks[1].set_block_hash( + make_block_hash_with_group_id(req.block_hashes[1], 1) + ) + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks)) + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + output = make_scheduler_output( + {req.request_id: 2 * block_size}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids=lambda _: (), + boundary_state_offloads={ + req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)] + }, + ) + if gone_via == "preempted": + output.preempted_req_ids = {req.request_id} + elif gone_via == "finished": + output.finished_req_ids = {req.request_id} + else: + sched._reqs_to_store.pop(req.request_id) + + meta = sched.build_connector_meta(output) + + # The exact-handoff block is never a store source for a departing request. + assert mamba_blocks[1].block_id not in meta.store_gpu_blocks + stats = sched.get_boundary_store_stats() + assert stats.dropped_request_gone == 1 + assert stats.stored == 0 + + +def test_boundary_handoff_keeps_block_meta_index_parallel() -> None: + """Boundary stores must extend block_meta alongside the block ids. + + ``_process_store_completion`` indexes ``block_meta`` by position, so a + boundary handoff that appends only block ids shortens the list, trips its + length assertion and shifts every following entry's metadata. + """ + block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler( + num_cpu_blocks=16, + num_gpu_blocks=24, + block_size=block_size, + enable_kv_cache_events=True, + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + assert sched.enable_kv_cache_events + + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + attn_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 + ) + mamba_blocks = gpu_pool.get_new_blocks(2) + mamba_blocks[1].set_block_hash( + make_block_hash_with_group_id(req.block_hashes[1], 1) + ) + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks)) + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + # Make the FA blocks positionally eligible in the same event as the + # explicit Mamba boundary handoff. + req.num_computed_tokens = 2 * block_size + + output = make_scheduler_output( + {req.request_id: 2 * block_size}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + output.kv_connector_block_state = KVConnectorBlockState( + req_ids=set(), + resolve_block_ids=lambda _: (), + boundary_state_offloads={ + req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)] + }, + ) + + gpu_ids, cpu_ids, req_ids, block_meta = sched.prepare_store_specs(output) + assert mamba_blocks[1].block_id in gpu_ids + assert set(gpu_ids) == { + *(block.block_id for block in attn_blocks), + mamba_blocks[1].block_id, + } + assert req_ids == [req.request_id] + assert block_meta is not None + assert len(block_meta) == len(gpu_ids) == len(cpu_ids) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 0589c83232f6..64bf8146c9e5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -19,6 +19,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput from vllm.v1.simple_kv_offload.manager import ( + BoundaryStoreStats, SimpleCPUOffloadScheduler, ) from vllm.v1.simple_kv_offload.metadata import ( @@ -300,6 +301,12 @@ def has_pending_push_work(self) -> bool: return self.scheduler_manager.has_pending_stores() return False + def get_boundary_store_stats(self) -> BoundaryStoreStats | None: + """Return cumulative boundary-handoff store diagnostics.""" + if self.scheduler_manager is not None: + return self.scheduler_manager.get_boundary_store_stats() + return None + def take_events(self) -> Iterable[KVCacheEvent]: if self.scheduler_manager is not None: return self.scheduler_manager.take_events() diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index a495f997bf3b..5fa881292467 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -155,6 +155,9 @@ def __init__( for manager in self.single_type_managers: if isinstance(manager, MambaManager): manager.drop_eagle_checkpoint_block = True + self.group_block_sizes = tuple( + manager.block_size for manager in self.single_type_managers + ) # A positive retention interval must be a multiple of the base hit granularity # (``scheduler_block_size``) to land on real cache-hit boundaries. @@ -605,6 +608,7 @@ def __init__( hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, num_prefill_lookahead: int = 0, + allow_partial_hash_hits: bool = True, ): super().__init__( kv_cache_config, @@ -629,7 +633,7 @@ def __init__( # Only groups that participate in prefix caching must satisfy the # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool # tail, block_size=kpool) are scratch buffers and excluded. - group_block_sizes = [ + cacheable_block_sizes = [ manager.block_size for manager, group in zip( self.single_type_managers, kv_cache_config.kv_cache_groups @@ -637,10 +641,10 @@ def __init__( if group.kv_cache_spec.prefix_cacheable ] assert all( - block_size % hash_block_size == 0 for block_size in group_block_sizes + block_size % hash_block_size == 0 for block_size in cacheable_block_sizes ), ( "Each KV cache group's real block_size must be divisible by " - f"hash_block_size. block_sizes={group_block_sizes}, " + f"hash_block_size. block_sizes={cacheable_block_sizes}, " f"hash_block_size={hash_block_size}" ) assert pcp_world_size == 1, "PCP not support hybrid attn now." @@ -669,7 +673,9 @@ def __init__( ) for g in kv_cache_config.kv_cache_groups ) - self.enable_partial_hash_hits = has_partial_mamba_group + self.enable_partial_hash_hits = ( + allow_partial_hash_hits and has_partial_mamba_group + ) if self.enable_partial_hash_hits: unsupported_partial_hit_managers = { type(manager).__name__ @@ -987,6 +993,7 @@ def get_kv_cache_coordinator( hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, num_prefill_lookahead: int = 0, + allow_partial_hash_hits: bool = True, ) -> KVCacheCoordinator: if not enable_caching: return KVCacheCoordinatorNoPrefixCache( @@ -1030,4 +1037,5 @@ def get_kv_cache_coordinator( hash_block_size=hash_block_size, metrics_collector=metrics_collector, num_prefill_lookahead=num_prefill_lookahead, + allow_partial_hash_hits=allow_partial_hash_hits, ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 15b35d9f7b57..93475e505526 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -48,6 +48,16 @@ class SingleTypeKVCacheManager(ABC): supports_fine_grained_hash_lookup: ClassVar[bool] = False + @property + def has_positionally_stable_blocks(self) -> bool: + """Whether positional offload scans can follow this block table. + + True means already-scanned indices will not be reused for a different + token range. Managers may still null blocks that left their retention + window, as long as the cursor's positional history remains valid. + """ + return True + def __init__( self, kv_cache_spec: KVCacheSpec, @@ -1423,6 +1433,12 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: class MambaManager(SingleTypeKVCacheManager): supports_fine_grained_hash_lookup: ClassVar[bool] = True + @property + def has_positionally_stable_blocks(self) -> bool: + # Align-mode Mamba can null interior states and relocate speculative + # blocks in place. Other modes retain positional identity. + return self.mamba_cache_mode != "align" + def __init__( self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs ) -> None: diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index afb51fb1d32b..9cae2f714cf7 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -5,6 +5,7 @@ import contextlib from collections.abc import Iterable from dataclasses import dataclass, field, replace +from enum import Enum, auto from typing import TYPE_CHECKING, Any from vllm.config import VllmConfig @@ -26,6 +27,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHashWithGroupId, ExternalBlockHash, + dcp_world_size_for_kv_cache_spec, generate_block_hash_extra_keys, get_block_hash, get_group_id, @@ -99,6 +101,36 @@ class StoreRequestState: finished: bool = False +class _StoreAdmission(Enum): + READY = auto() + NULL_BLOCK = auto() + NOT_HASHED = auto() + IN_FLIGHT = auto() + ALREADY_CACHED = auto() + + +@dataclass +class BoundaryStoreStats: + """Cumulative counters for boundary hand-off stores. + + Every branch that declines a hand-off is silent by design: the boundary is + simply not offloaded and a later request misses. Without counters the only + symptom is a lower cache hit rate with nothing in the logs, so each decline + reason is counted and exposed through + ``SimpleCPUOffloadConnector.get_boundary_store_stats()``. Reset by + ``reset()``; not otherwise cleared. + """ + + published: int = 0 + stored: int = 0 + dropped_cpu_full: int = 0 + dropped_request_gone: int = 0 + skipped_already_cached: int = 0 + skipped_in_flight: int = 0 + dropped_null_block: int = 0 + dropped_not_hashed: int = 0 + + class SimpleCPUOffloadScheduler: """Scheduler-side manager for CPU offloading.""" @@ -141,16 +173,6 @@ def __init__( self.fa_gidx = g_idx break assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups) - # FA group's own block_size; divides scheduler_block_size (the LCM) - # but is NOT assumed to equal it. - self.fa_block_size: int = ( - self.cpu_kv_cache_config.kv_cache_groups[ - self.fa_gidx - ].kv_cache_spec.block_size - * self.cp_world_size - ) - assert self.block_size % self.fa_block_size == 0 - logger.info( "SimpleCPUOffloadScheduler: Allocating %d offload blocks " "(%.2f GB, mode=%s, backend=%s)", @@ -175,7 +197,13 @@ def __init__( pcp_world_size=1, scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, + allow_partial_hash_hits=not lazy_offload, ) + self.group_block_sizes = self.cpu_coordinator.group_block_sizes + # FA group's own resolved block_size; divides scheduler_block_size (the + # LCM) but is NOT assumed to equal it. + self.fa_block_size: int = self.group_block_sizes[self.fa_gidx] + assert self.block_size % self.fa_block_size == 0 self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool # GPU block pool reference - bound after scheduler builds kv_cache_manager self._gpu_block_pool: BlockPool | None = None @@ -186,10 +214,13 @@ def __init__( # the worker reports completions by event index, not request id. self._load_event_to_reqs: dict[int, list[str]] = {} - # Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit, - # kept pinned via touch() while awaiting update_state_after_alloc(). + # Pending (cpu_hit_blocks, hit_length, num_computed_tokens) tuples from + # find_longest_cache_hit, kept pinned via touch() while awaiting + # update_state_after_alloc(). ``num_computed_tokens`` is the local + # prefix the hit was resolved against; the load path needs it to place + # external blocks and must not re-derive it from block metadata. self._pending_cpu_hits: dict[ - str, tuple[tuple[list[KVCacheBlock], ...], int] + str, tuple[tuple[list[KVCacheBlock], ...], int, int] ] = {} # Store metadata @@ -216,6 +247,7 @@ def __init__( # Event counters self._load_event_counter: int = 0 self._store_event_counter: int = 0 + self.boundary_store_stats = BoundaryStoreStats() # For TP/PP: track partial store completions across steps. # Events must be reported by all world_size workers before considered complete. @@ -267,7 +299,11 @@ def _estimate_lazy_target_blocks( target = 0 for g in kv_cache_config.kv_cache_groups: spec = g.kv_cache_spec - block_size = spec.block_size * cp_world_size + # Only full attention is sharded across DCP ranks; replicated specs + # (mamba, sliding window, chunked-local) keep their own block size. + block_size = spec.block_size * dcp_world_size_for_kv_cache_spec( + spec, cp_world_size + ) if isinstance(spec, MambaSpec): target += 2 elif isinstance(spec, SlidingWindowSpec): @@ -296,6 +332,20 @@ def get_num_new_matched_tokens( if request.skip_reading_prefix_cache: return 0, False + if num_computed_tokens % self.block_size != 0: + # Transfers are whole-block copies, so an external suffix cannot + # start in the middle of a destination block. When the GPU-side + # coordinator also serves fine-grained hits, the local prefix can + # land on a hash boundary that is not a scheduler-block boundary; + # such a request keeps its local hit and skips the external one. + logger.warning_once( + "SimpleCPU external lookup requires scheduler-block-aligned " + "local tokens, got %d tokens with scheduler block size %d.", + num_computed_tokens, + self.block_size, + ) + return 0, False + num_skipped_hashes = num_computed_tokens // self.hash_block_size remaining_hashes = request.block_hashes[num_skipped_hashes:] @@ -318,6 +368,7 @@ def get_num_new_matched_tokens( self._pending_cpu_hits[request.request_id] = ( cpu_hit_blocks, hit_length, + num_computed_tokens, ) return hit_length, True return 0, False @@ -371,20 +422,19 @@ def update_state_after_alloc( ) return - cpu_hit_blocks_full, _ = pending + cpu_hit_blocks_full, _, num_computed_tokens = pending - # ``num_external_tokens`` is LCM-aligned (checked per-group below), - # so this counts whole scheduler-aligned chunks of incoming tokens. - num_blocks_to_load = num_external_tokens // self.block_size - assert num_blocks_to_load > 0 - num_cached_fa_blocks = sum( - blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx] + assert num_computed_tokens % self.block_size == 0, ( + "SimpleCPU external loads must start at a scheduler-block boundary" + ) + # Fine-grained hits are hash-block aligned, not group-block aligned, so + # the per-group block count below rounds up instead of dividing exactly. + # Hash alignment is still required: it is what makes the accepted + # external suffix start on a boundary every group can address. + assert num_external_tokens % self.hash_block_size == 0, ( + f"num_external_tokens={num_external_tokens} is not aligned to " + f"hash_block_size={self.hash_block_size}" ) - num_computed_tokens = num_cached_fa_blocks * self.fa_block_size - - # Build transfer pairs across all groups. - total_computed_tokens = num_computed_tokens + num_external_tokens - kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups # The scheduler may have accepted fewer blocks than # get_num_new_matched_tokens() reported. @@ -393,14 +443,8 @@ def update_state_after_alloc( # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): - g_block_size = ( - kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size - ) - assert num_external_tokens % g_block_size == 0, ( - f"num_external_tokens={num_external_tokens} not aligned to " - f"group {g} block_size={g_block_size}" - ) - n_take_g = num_external_tokens // g_block_size + g_block_size = self.group_block_sizes[g] + n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) gpu_block_ids: list[int] = [] @@ -413,14 +457,8 @@ def update_state_after_alloc( if n_ext_g == 0: continue - # Number of blocks in the computed range for this group. - g_block_size = ( - kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size - ) - n_computed_g = cdiv(total_computed_tokens, g_block_size) - - # Back-trace: ext blocks sit at the tail of the computed range. - gpu_ext_start = n_computed_g - n_ext_g + g_block_size = self.group_block_sizes[g] + gpu_ext_start = num_computed_tokens // g_block_size group_gpu_ids = block_ids_by_group[g] for i, cpu_blk in enumerate(cpu_blocks_g): @@ -614,6 +652,72 @@ def _prepare_eager_store_specs( num_groups = len(self.cpu_kv_cache_config.kv_cache_groups) # Dedup against blocks already scheduled. in_flight = self._in_flight_store_gpu_blocks + num_free = cpu_block_pool.get_num_free_blocks() + + block_state = scheduler_output.kv_connector_block_state + boundary_offloads = ( + block_state.boundary_state_offloads if block_state is not None else {} + ) + stats = self.boundary_store_stats + preempted_req_ids = scheduler_output.preempted_req_ids or () + for req_id, entries in boundary_offloads.items(): + # The scheduler drains handoffs without filtering by request + # liveness. A request that finished or was preempted in this step + # is giving up its blocks, which may already have been reallocated + # to another request; reading them would publish unrelated KV under + # a valid hash. Drop conservatively, as the mooncake store does. + store_state = self._reqs_to_store.get(req_id) + if ( + store_state is None + or store_state.finished + or req_id in preempted_req_ids + or req_id in scheduler_output.finished_req_ids + ): + stats.published += len(entries) + stats.dropped_request_gone += len(entries) + continue + scheduled_for_req = False + # ``boundary_tokens`` is not needed here: the cache key comes from + # the handed-off block itself, not from the offered boundary. + for _, gpu_block_id, _ in entries: + stats.published += 1 + gpu_block = gpu_block_pool.blocks[gpu_block_id] + admission = self._classify_store_candidate(gpu_block) + if admission is _StoreAdmission.NULL_BLOCK: + stats.dropped_null_block += 1 + continue + if admission is _StoreAdmission.NOT_HASHED: + stats.dropped_not_hashed += 1 + continue + if admission is _StoreAdmission.IN_FLIGHT: + stats.skipped_in_flight += 1 + continue + if admission is _StoreAdmission.ALREADY_CACHED: + stats.skipped_already_cached += 1 + continue + if num_free <= 0: + stats.dropped_cpu_full += 1 + logger.warning_once( + "SimpleCPU dropped a boundary-state handoff because " + "the CPU block pool is full; this boundary will not be retried." + ) + continue + cpu_block = cpu_block_pool.get_new_blocks(1)[0] + merged_gpu_block_ids.append(gpu_block_id) + merged_cpu_block_ids.append(cpu_block.block_id) + if merged_block_meta is not None: + # Keep the metadata list index-parallel with the block ids. + # The hand-off block's own hashes are resolved at completion, + # so there is nothing to capture here; emit the event without + # a payload, as the size-mismatch fallback already does. + merged_block_meta.append({}) + in_flight.add(gpu_block_id) + gpu_block_pool.touch([gpu_block]) + num_free -= 1 + stats.stored += 1 + scheduled_for_req = True + if scheduled_for_req: + req_ids.append(req_id) for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output): state = self._reqs_to_store.get(req_id) @@ -641,7 +745,7 @@ def _prepare_eager_store_specs( self._select_eager_blocks_to_store(state, block_ids_by_group) ) - # --- Phase 2: Batch allocate CPU blocks --- + # Batch allocate the CPU destinations. n_to_alloc = len(gpu_block_ids) if n_to_alloc > 0: cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc) @@ -674,6 +778,9 @@ def _prepare_eager_store_specs( for g in range(num_groups): state.num_stored_blocks[g] += advanced_per_group[g] + # A request can contribute both boundary and positional blocks to the + # same event. Completion tracking is set-based, so keep one request ID. + req_ids = list(dict.fromkeys(req_ids)) return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta def _select_eager_blocks_to_store( @@ -694,19 +801,31 @@ def _select_eager_blocks_to_store( ) request = state.request confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders - aligned_tokens = confirmed_tokens // self.block_size * self.block_size + # Truncate to the granularity a lookup can actually land on. With + # fine-grained hits that is hash_block_size; otherwise hits only land on + # the scheduler block (the group LCM). Using the LCM unconditionally + # would drop whole blocks that sit between the last LCM boundary and a + # reachable fine-grained boundary, so the other groups would hold that + # boundary and this one would not, and the joint hybrid lookup would + # reconcile to zero. + store_alignment = ( + self.hash_block_size + if self.cpu_coordinator.enable_partial_hash_hits + else self.block_size + ) + aligned_tokens = confirmed_tokens // store_alignment * store_alignment num_free = self.cpu_block_pool.get_num_free_blocks() for g, group_gpu_ids in enumerate(block_ids_by_group): if len(gpu_block_ids) >= num_free: break + group_manager = self.cpu_coordinator.single_type_managers[g] + if not group_manager.has_positionally_stable_blocks: + continue # FIXME (yifan): handle CPU cache eviction, where # num_stored_blocks can be stale and omit evicted blocks in # the middle of the request. - group_size = ( - self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size - * self.cp_world_size - ) + group_size = self.group_block_sizes[g] ready = min(len(group_gpu_ids), aligned_tokens // group_size) resolved_hashes = ( resolve_block_hashes( @@ -720,20 +839,16 @@ def _select_eager_blocks_to_store( start = state.num_stored_blocks[g] for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start): gpu_block = self._gpu_block_pool.blocks[gpu_block_id] - if gpu_block.is_null or gpu_block.block_hash is None: - advanced_per_group[g] += 1 - continue if ( - gpu_block_id in self._in_flight_store_gpu_blocks - or self.cpu_block_pool.cached_block_hash_to_block.get_one_block( - gpu_block.block_hash - ) - is not None + self._classify_store_candidate(gpu_block) + is not _StoreAdmission.READY ): advanced_per_group[g] += 1 continue if len(gpu_block_ids) >= num_free: break + primary_block_hash = gpu_block.block_hash + assert primary_block_hash is not None gpu_block_ids.append(gpu_block_id) advanced_per_group[g] += 1 if block_meta is not None: @@ -750,7 +865,7 @@ def _select_eager_blocks_to_store( request, token_start, token_end, curr_mm_idx ) meta_by_hash = { - gpu_block.block_hash: BlockStoreMeta( + primary_block_hash: BlockStoreMeta( token_ids=list( request.all_token_ids[token_start:token_end] ), @@ -766,7 +881,7 @@ def _select_eager_blocks_to_store( block_hash = make_block_hash_with_group_id( request.block_hashes[hash_idx], g ) - if block_hash == gpu_block.block_hash: + if block_hash is None or block_hash == primary_block_hash: continue hash_start = hash_idx * self.hash_block_size hash_end = hash_start + self.hash_block_size @@ -796,6 +911,25 @@ def _select_eager_blocks_to_store( return gpu_block_ids, advanced_per_group, block_meta + def _classify_store_candidate(self, gpu_block: "KVCacheBlock") -> _StoreAdmission: + if gpu_block.is_null: + return _StoreAdmission.NULL_BLOCK + if gpu_block.block_hash is None: + return _StoreAdmission.NOT_HASHED + if gpu_block.block_id in self._in_flight_store_gpu_blocks: + return _StoreAdmission.IN_FLIGHT + if ( + self.cpu_block_pool.cached_block_hash_to_block.get_one_block( + gpu_block.block_hash + ) + is not None + ): + return _StoreAdmission.ALREADY_CACHED + return _StoreAdmission.READY + + def get_boundary_store_stats(self) -> BoundaryStoreStats: + return replace(self.boundary_store_stats) + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: """Handle async transfer completions from worker. @@ -886,12 +1020,7 @@ def _process_store_completion( if self.enable_kv_cache_events: meta_by_hash = block_meta[i] if block_meta is not None else {} primary_group_idx = get_group_id(primary_hash) - primary_spec = self.cpu_kv_cache_config.kv_cache_groups[ - primary_group_idx - ].kv_cache_spec - primary_block_size = primary_spec.block_size * ( - self.cp_world_size if not isinstance(primary_spec, MambaSpec) else 1 - ) + primary_block_size = self.group_block_sizes[primary_group_idx] events_to_emit: list[ tuple[BlockHashWithGroupId, BlockStoreMeta | None, int] ] = [ @@ -1028,6 +1157,21 @@ def _queue_finished_eager_store( if state is None or gpu_pool is None: return gpu_ids, _, block_meta = self._select_eager_blocks_to_store(state, block_ids) + partial_tail = self._find_fa_partial_tail_source(request, block_ids) + if ( + partial_tail is not None + # Decode tokens can fill the boundary block, in which case the + # positional scan above already selected it. Appending again would + # allocate two CPU blocks for one GPU block. + and partial_tail not in gpu_ids + and len(gpu_ids) < self.cpu_block_pool.get_num_free_blocks() + ): + gpu_ids.append(partial_tail) + if block_meta is not None: + # The block's own hashes are registered at completion, so no + # capture is needed here; emit the event without a payload, as + # the size-mismatch fallback above already does. + block_meta.append({}) if not gpu_ids: return cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids)) @@ -1037,9 +1181,64 @@ def _queue_finished_eager_store( self._in_flight_store_gpu_blocks.update(gpu_ids) gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids]) - def _free_pending_cpu_hit(self, pending: tuple) -> None: + def _find_fa_partial_tail_source( + self, request: "Request", block_ids: tuple[list[int], ...] + ) -> int | None: + """Locate the full-attention prompt-tail block for a fine-grained hit. + + Mirrors ``FullAttentionManager._cache_partial_tail_block``: only the + final prompt hash boundary is eligible, and boundaries that land on a + physical block edge are already covered by the positional scan. + + Attention block tables are append-only, so the block is located + positionally, as the mooncake store and offloading connector also do; + no new core hand-off is required. The boundary key is registered at + completion along with the block's other hashes, so nothing has to be + captured here. + """ + if not self.cpu_coordinator.enable_partial_hash_hits: + return None + assert self._gpu_block_pool is not None + boundary_tokens = ( + request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + ) + if boundary_tokens == 0 or boundary_tokens > request.num_computed_tokens: + return None + if boundary_tokens % self.fa_block_size == 0: + return None + block_idx = boundary_tokens // self.fa_block_size + fa_block_ids = block_ids[self.fa_gidx] + if block_idx >= len(fa_block_ids): + return None + gpu_block_id = fa_block_ids[block_idx] + gpu_block = self._gpu_block_pool.blocks[gpu_block_id] + if gpu_block.is_null or gpu_block_id in self._in_flight_store_gpu_blocks: + return None + hash_idx = boundary_tokens // self.hash_block_size - 1 + if hash_idx >= len(request.block_hashes): + return None + block_hash = make_block_hash_with_group_id( + request.block_hashes[hash_idx], self.fa_gidx + ) + # Only worth copying when the boundary key is actually registered on + # the source block; completion replays the block's own hashes. + if ( + gpu_block.block_hash != block_hash + and block_hash + not in self._gpu_block_pool.cached_block_hashes_by_block.get( + gpu_block_id, () + ) + ): + return None + if self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash): + return None + return gpu_block_id + + def _free_pending_cpu_hit( + self, pending: tuple[tuple[list["KVCacheBlock"], ...], int, int] + ) -> None: """Release the temporary CPU block pin taken in get_num_new_matched_tokens().""" - cpu_hit_blocks, _ = pending + cpu_hit_blocks, _hit_length, _num_computed_tokens = pending blocks_to_free = [ blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null ] @@ -1140,6 +1339,7 @@ def reset(self) -> bool: if event_idx in self._abandoned_store_event_to_blocks } self._cursor = None + self.boundary_store_stats = BoundaryStoreStats() # NOTE: _load_event_counter / _store_event_counter are not # reset as they are monotonic and must stay ahead of the workers # high-water marks to avoid event index collisions From ce08bb5b3463fd423be90d5f22c2b41142834bf6 Mon Sep 17 00:00:00 2001 From: Yicong Wang Date: Fri, 11 Sep 2026 10:57:19 +0800 Subject: [PATCH 019/774] [Bugfix][Mooncake] Fix heterogeneous PP transfer completion (#56033) Signed-off-by: wangyicong Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 62 +++++++++++++++++++ .../v1/mooncake/mooncake_connector.py | 26 +++++++- 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index f989fb4e32a4..43de313821be 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -892,6 +892,68 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout(): assert send_meta.need_send == 2 +@pytest.mark.asyncio +async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers(): + """P4/D2 must retain KV until both D stages finish, including a no-op pull.""" + config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer") + with ( + set_current_vllm_config(config), + patch_worker_dependencies(), + patch.object(MooncakeConnectorWorker, "_sync_block_size_with_kernel"), + ): + worker = MooncakeConnector( + config, KVConnectorRole.WORKER, _make_test_kv_cache_config() + ).connector_worker + try: + worker.pp_size = 4 + worker.kv_caches_base_addr = [0x1000] + worker.block_len_per_layer = [256] + worker.kv_block_len_per_layer = [256] + worker.registered_layer_names = ["model.layers.0.self_attn"] + worker.registered_layer_indices = [0] + send_meta = SendBlockMeta( + p_req_id="p-req", + transfer_id="transfer", + local_block_ids=[[1]], + ready=asyncio.Event(), + ) + send_meta.ready.set() + worker.reqs_need_send["transfer"] = send_meta + sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock()) + with ( + patch.object(worker, "sender_loop", asyncio.get_running_loop()), + patch.object(worker, "_send_blocks", return_value=0) as send, + ): + for pp_rank in range(2): + metadata = MooncakeXferMetadata( + remote_hostname="consumer", + remote_port=1234 + pp_rank, + remote_tp_size=1, + remote_tp_rank=0, + remote_pp_size=2, + req_blocks={"d-req": ("transfer", [[2]])}, + kv_caches_base_addr=[0x2000], + block_lens=[256], + kv_block_lens=[256], + registered_layer_names=[f"model.layers.{pp_rank}.self_attn"], + registered_layer_indices=[pp_rank], + ) + await worker.send_kv_to_decode(b"consumer", sock, metadata) + response = worker._xfer_resp_decoder.decode( + sock.send_multipart.call_args.args[0][1] + ) + assert response.status == MooncakeXferResponseStatus.FINISH + assert response.ok_reqs == ["d-req"] + assert ("transfer" in worker.reqs_need_send) == (pp_rank == 0) + assert worker.finished_sending_reqs == ( + set() if pp_rank == 0 else {"p-req"} + ) + send.assert_called_once_with("consumer:1234", [0x1100], [0x2200], [256]) + finally: + worker.shutdown() + worker.is_kv_consumer = True + + @pytest.mark.asyncio @patch( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index dc574a6219b0..0a899b7b62d2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -307,6 +307,8 @@ def _validate_asymmetric_region_lengths( def _align_transfer_regions( local_regions: list[TransferRegion], remote_regions: list[TransferRegion], + *, + allow_partial_layers: bool = False, ) -> tuple[list[TransferRegion], list[TransferRegion], str | None]: """Align KV transfer regions by registered layer-name occurrence. @@ -335,6 +337,11 @@ def keyed_regions( for key, local_region in local_keyed: remote_region = remote_by_key.get(key) if remote_region is None: + if ( + allow_partial_layers + and (local_region.layer_name, 0) not in remote_by_key + ): + continue return ( [], [], @@ -393,6 +400,7 @@ class MooncakeXferMetadata( registered_layer_names: list[str] = msgspec.field(default_factory=list) registered_layer_indices: list[int] = msgspec.field(default_factory=list) registered_group_indices: list[int] = msgspec.field(default_factory=list) + remote_pp_size: int = 1 class MooncakeXferResponseStatus(IntEnum): @@ -1232,7 +1240,11 @@ async def send_kv_to_decode( meta.registered_group_indices, ) local_regions, remote_regions, align_err = _align_transfer_regions( - local_regions, remote_regions + local_regions, + remote_regions, + allow_partial_layers=( + meta.remote_pp_size > 1 and meta.remote_pp_size != self.pp_size + ), ) if align_err is not None: response = MooncakeXferResponse( @@ -1322,7 +1334,9 @@ async def wait_and_ret( # Mark it sending to avoid expiration. send_meta.sending += 1 if not send_meta.need_send: - self.resolve_need_send(send_meta, remote_tp_ranks) + self.resolve_need_send( + send_meta, remote_tp_ranks, meta.remote_pp_size + ) ready_reqs.append((d_req_id, send_meta)) else: # Otherwise (expired, very unlikely), just forget it. @@ -1398,11 +1412,16 @@ def resolve_need_send( self, send_meta: SendBlockMeta, remote_tp_ranks: list[int], + remote_pp_size: int = 1, ): # Prepare for heterogeneous TP (one P pairs to multiple D) send_meta.need_send = len(remote_tp_ranks) + if remote_pp_size > 1 and remote_pp_size != self.pp_size: + # Each consumer PP stage pulls every producer stage, including + # peers with no shared layers. + send_meta.need_send *= remote_pp_size logger.debug( - "Mooncake request %s will be served by %d consumer TP workers: TP ranks=%s", + "Mooncake request %s will be served by %d consumer workers: TP ranks=%s", send_meta.transfer_id, send_meta.need_send, remote_tp_ranks, @@ -1851,6 +1870,7 @@ async def receive_kv_from_single_worker( remote_port=self.rpc_port, remote_tp_size=self.tp_size, remote_tp_rank=self.tp_rank, + remote_pp_size=self.pp_size, req_blocks={ req_id: (pull_meta.transfer_id, pull_meta.local_block_ids) for req_id, pull_meta in pull_metas.items() From b5d4186300fda544ec5ec51113d569159c5c52bd Mon Sep 17 00:00:00 2001 From: Kushal <72650064+kushaldabbe@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:52:33 +0900 Subject: [PATCH 020/774] [Bugfix][LoRA] Use stored rsLoRA scaling factor in MoE expert packing (#55548) Signed-off-by: Kushal <72650064+kushaldabbe@users.noreply.github.com> --- tests/lora/test_lora_weights.py | 52 +++++++++++++++++++++++++++++++++ vllm/lora/lora_weights.py | 7 +++-- 2 files changed, 56 insertions(+), 3 deletions(-) create mode 100644 tests/lora/test_lora_weights.py diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py new file mode 100644 index 000000000000..ff132147933e --- /dev/null +++ b/tests/lora/test_lora_weights.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights + +RANK, ALPHA = 16, 32 +# PEFTHelper stores alpha/sqrt(r) for use_rslora adapters. +RSLORA_SCALING = ALPHA / RANK**0.5 + + +@pytest.fixture(autouse=True) +def should_do_global_cleanup_after_test(): + # Pure-python weight packing; skip GPU/ray cleanup. + return False + + +def moe_loras(scaling: float, stacked: bool = False) -> list[LoRALayerWeights]: + a_shape = (3, RANK, 4) if stacked else (RANK, 4) + b_shape = (3, 6, RANK) if stacked else (6, RANK) + return [ + LoRALayerWeights( + module_name=f"experts.w{i}", + rank=RANK, + lora_alpha=ALPHA, + lora_a=torch.randn(*a_shape), + lora_b=torch.randn(*b_shape), + scaling=scaling, + ) + for i in (1, 2, 3) + ] + + +@pytest.mark.parametrize( + "packer,stacked", + [ + (PackedLoRALayerWeights.pack_moe, False), + (PackedLoRALayerWeights.pack_moe_stacked, True), + ], +) +def test_moe_packing_uses_stored_scaling(packer, stacked: bool): + packed = packer(moe_loras(RSLORA_SCALING, stacked), "experts") + assert packed.scaling == pytest.approx([RSLORA_SCALING] * 3) + + +def test_non_gated_moe_keeps_w3_unscaled(): + packed = PackedLoRALayerWeights.pack_moe( + moe_loras(RSLORA_SCALING), "experts", is_non_gated_moe=True + ) + assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2) + assert packed.scaling[2] == 1.0 diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 38a5875525da..bb241db77e13 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -201,8 +201,9 @@ def pack_moe( w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank) w2_lora_b = torch.stack(w2_lora_b_lst, dim=0) - # All w1, w2, w3 have the same scaling factor. - scaling = lora_alpha / rank + # All w1, w2, w3 have the same scaling factor. Use the per-adapter + # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank. + scaling = first_lora.scaling last_scaling = scaling if is_non_gated_moe: @@ -250,7 +251,7 @@ def pack_moe_stacked( assert w1_lora is not None and w2_lora is not None and w3_lora is not None rank = w1_lora.rank lora_alpha = w1_lora.lora_alpha - scaling = lora_alpha / rank + scaling = w1_lora.scaling return cls( module_name, rank, From e7edf17cea217e52701f913cd8491fcacf2d9490 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 11 Sep 2026 12:53:37 +0800 Subject: [PATCH 021/774] [EC] Automatically enable embedding inputs on EC/KV consumers (#56379) Signed-off-by: Tianyu Guo --- vllm/config/vllm.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 893acb00cdb3..cc443a14dd47 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1716,7 +1716,7 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) - self._resolve_allow_missing_mm_embeddings() + self._resolve_mm_embedding_inputs() self._resolve_mm_processor_device() self._validate_mm_processor_device() @@ -2497,8 +2497,8 @@ def __str__(self): f"kernel_config={self.kernel_config!r}" ) - def _resolve_allow_missing_mm_embeddings(self) -> None: - """Allow `*_embeds` tensors to be omitted on disaggregated consumers. + def _resolve_mm_embedding_inputs(self) -> None: + """Accept embedding inputs, tensor optional, on disaggregated consumers. An EC consumer loads embeddings from its connector. A KV consumer receives the prompt KV produced from those embeddings, so it does not @@ -2519,11 +2519,21 @@ def _resolve_allow_missing_mm_embeddings(self) -> None: mm_config.allow_missing_mm_embeddings = ( ec_config is not None and ec_config.is_ec_consumer ) or (kv_config is not None and kv_config.is_kv_consumer) - if mm_config.allow_missing_mm_embeddings: + if not mm_config.allow_missing_mm_embeddings: + return + + if not mm_config.enable_mm_embeds: + # Allowing missing tensors still requires enabling embedding inputs + # for the frontend to accept metadata-only requests. + mm_config.enable_mm_embeds = True logger.info_once( - "EC/KV consumer: pre-computed-embedding inputs may " - "omit the embedding tensor." + "EC/KV consumer: accepting pre-computed-embedding inputs, " + "which this role is sent by definition." ) + logger.info_once( + "EC/KV consumer: pre-computed-embedding inputs may " + "omit the embedding tensor." + ) def _resolve_mm_encoder_only(self) -> None: """Enable encoder-only mode for a dedicated EC producer.""" From 980c16c8e4c66dce0bc4d355e4ec7fa2ae7b5141 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Fri, 11 Sep 2026 01:07:53 -0400 Subject: [PATCH 022/774] [PCP][Spec Decode] Adds PCP support for single-module MTP and replicated DSpark. (#56107) Signed-off-by: Summer Yang Signed-off-by: Lucas Wilkinson Co-authored-by: Summer Yang Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex --- .../test_eagle_draft_attn_metadata.py | 1 + .../test_gpu_autoregressive_speculator.py | 1 + tests/v1/worker/test_gpu_pcp_manager.py | 11 +- .../layers/sparse_attn_indexer.py | 5 +- vllm/v1/worker/cp_utils.py | 10 +- vllm/v1/worker/gpu/model_runner.py | 16 +- vllm/v1/worker/gpu/pcp_manager.py | 137 +++++++++++------- .../spec_decode/autoregressive/speculator.py | 34 ++++- .../gpu/spec_decode/dflash/speculator.py | 12 ++ .../worker/gpu/spec_decode/mtp/speculator.py | 3 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 8 +- 11 files changed, 159 insertions(+), 79 deletions(-) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index b3006081e579..95e263f9de8d 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -48,6 +48,7 @@ def _make_fake_speculator( ) return SimpleNamespace( arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"), + draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool), block_tables=fake_block_tables, input_buffers=fake_input_buffers, attn_groups=[], diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 6d8b8cbdbca5..e9c7a3b290b2 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -94,6 +94,7 @@ def _make_speculator( speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False + speculator.pcp_manager = None speculator.vllm_config = None speculator.input_buffers = SimpleNamespace( input_ids=torch.arange(4), diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index eea668aafb6c..2d8941f6ffb3 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -11,7 +11,6 @@ from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.pcp_manager import PCPManager -from vllm.v1.worker.gpu.states import RequestState def _copy_to_cpu(value, out=None, device=None): @@ -125,6 +124,7 @@ def _make_global_decode_batch( base = InputBatch.make_dummy(num_reqs, num_tokens, buffers) buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device) + buffers.positions[:num_reqs] = torch.tensor(num_computed_tokens, device=device) query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to( device @@ -166,19 +166,10 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): the field afterwards from the PCP-owned buffers. """ device = torch.device("cuda:0") - req_states = RequestState( - max_num_reqs=4, - max_model_len=64, - max_num_batched_tokens=8, - num_speculative_steps=1, - vocab_size=8, - device=device, - ) manager = PCPManager( pcp_world_size=2, pcp_rank=0, device=device, - req_states=req_states, max_num_reqs=4, max_num_tokens=8, dcp_world_size=2, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 7d656f06e85f..bb00a9cfe103 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -401,9 +401,10 @@ def sparse_attn_indexer( # out-of-bounds reads in the kernel. # Keep PCP padding so every rank contributes the same all-gather shape. num_tokens = slot_mapping.shape[0] - if use_pcp: - num_tokens //= get_pcp_group().world_size if k is not None: + # MTP draft decodes are replicated and use an unexpanded slot mapping. + if use_pcp and num_tokens > k.shape[0]: + num_tokens //= get_pcp_group().world_size k = k[:num_tokens] if not skip_k_cache_insert: diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 7e208da39327..79651759aad4 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -18,16 +18,20 @@ logger = init_logger(__name__) -def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: +def check_attention_cp_compatibility( + vllm_config: VllmConfig, + target_layer_names: set[str] | None = None, +) -> None: pcp_size = vllm_config.parallel_config.prefill_context_parallel_size dcp_size = vllm_config.parallel_config.decode_context_parallel_size interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size if pcp_size * dcp_size > 1: layer_type = cast(type[Any], AttentionLayerBase) layers = get_layers_from_vllm_config(vllm_config, layer_type) - for layer in layers.values(): + for layer_name, layer in layers.items(): + check_pcp = target_layer_names is None or layer_name in target_layer_names get_attn_backend = getattr(layer, "get_attn_backend", None) - if pcp_size > 1 and get_attn_backend is not None: + if pcp_size > 1 and check_pcp and get_attn_backend is not None: backend = get_attn_backend() assert backend.supports_pcp(), ( "PCP requires attention backend support, " diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 533fcbe1e709..6a3d2e34e9ac 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -656,7 +656,6 @@ def initialize_kv_cache( self.vllm_config, self.device, self.supports_mm_inputs, - self.req_states, self.block_tables, cls=self.pcp_manager_cls, ) @@ -668,6 +667,8 @@ def initialize_kv_cache( self.kv_cache_config, self.max_num_reqs, ) + if self.speculator is not None: + self.speculator.pcp_manager = self.pcp_manager initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config, @@ -697,7 +698,7 @@ def initialize_kv_cache( self.fast_prefill = FastPrefillHelper( self.cudagraph_manager, self.max_num_tokens ) - check_attention_cp_compatibility(self.vllm_config) + check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) self.speculator.set_attn( @@ -1964,9 +1965,16 @@ def sample_tokens( return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens + draft_hidden_states = hidden_states + assert draft_hidden_states is not None hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( self.pcp_manager, hidden_states, input_batch ) + if self.pcp_manager is not None and aux_hidden_states is not None: + aux_hidden_states = [ + self.pcp_manager.restore_hidden_states(states) + for states in aux_hidden_states + ] sampler_output, num_sampled, num_rejected = self.sample( hidden_states, input_batch, grammar_output @@ -2043,10 +2051,10 @@ def sample_tokens( # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states + spec_hidden_states = draft_hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] with use_workspace_lane(self._draft_workspace_lane): draft_tokens = self.speculator.propose( input_batch, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index abec427b3793..6b97254621ec 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -6,7 +6,7 @@ import numpy as np import torch -from vllm.config import CUDAGraphMode, VllmConfig +from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -15,10 +15,7 @@ from vllm.v1.worker.gpu.input_batch import ( InputBatch, InputBuffers, - combine_sampled_and_draft_tokens, - prepare_pos_seq_lens, ) -from vllm.v1.worker.gpu.states import RequestState logger = init_logger(__name__) @@ -48,7 +45,6 @@ def __init__( pcp_world_size: int, pcp_rank: int, device: torch.device, - req_states: RequestState | None = None, max_num_reqs: int | None = None, max_num_tokens: int | None = None, block_tables: BlockTables | None = None, @@ -64,7 +60,9 @@ def __init__( self.cp_interleave = cp_interleave self._global_batch: InputBatch | None = None - self._req_states = req_states + self._local_batch: InputBatch | None = None + self._local_gather_idx: torch.Tensor | None = None + self.draft_prefill_batch: InputBatch | None = None self._block_tables = block_tables self._hidden_restore_idx: torch.Tensor | None = None self._padded_gather_idx: torch.Tensor | None = None @@ -77,11 +75,6 @@ def __init__( if max_num_local_reqs is not None and max_num_tokens is not None else None ) - self._local_req_idx = ( - torch.arange(max_num_local_reqs, dtype=torch.int32, device=device) - if max_num_local_reqs is not None - else None - ) self._local_block_tables: tuple[torch.Tensor, ...] | None self._local_block_table_ptrs: torch.Tensor | None if block_tables is not None and max_num_local_reqs is not None: @@ -144,19 +137,23 @@ def validate_config( raise NotImplementedError("MRV2 PCP does not support MM inputs yet.") if vllm_config.lora_config is not None: raise NotImplementedError("MRV2 PCP does not support LoRA yet.") - if vllm_config.speculative_config is not None: - raise NotImplementedError( - "MRV2 PCP does not support speculative decoding yet." - ) - is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk") - if ( - is_sparse_mla - and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - ): - raise NotImplementedError( - "MRV2 sparse MLA PCP does not support CUDA graphs yet. " - "Set -cc.cudagraph_mode=NONE." - ) + speculative_config = vllm_config.speculative_config + if speculative_config is not None: + if speculative_config.use_dspark(): + dcp_size = parallel_config.decode_context_parallel_size + if dcp_size not in (1, pcp_size): + raise NotImplementedError( + "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got " + f"DCP={dcp_size}, PCP={pcp_size}." + ) + elif ( + speculative_config.method != "mtp" + or speculative_config.use_multi_module_mtp() + ): + raise NotImplementedError( + "MRV2 PCP only supports DSpark or single-module MTP " + "speculative decoding." + ) if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") @@ -358,12 +355,8 @@ def partition_batch( input_batch: InputBatch, padded_num_tokens: int | None = None, ) -> InputBatch: - assert self._req_states is not None assert self._input_buffers is not None - req_states = self._req_states input_buffers = self._input_buffers - if input_batch.num_draft_tokens > 0: - raise NotImplementedError("MRV2 PCP does not support spec decode yet.") global_batch = input_batch self._global_batch = global_batch @@ -458,12 +451,21 @@ def partition_batch( local_gather_idx = self._padded_gather_idx[ rank_token_start : rank_token_start + num_local_tokens_padded ] + self._local_gather_idx = local_gather_idx torch.index_select( global_batch.input_ids, 0, local_gather_idx, out=input_buffers.input_ids[:num_local_tokens_padded], ) + # Keep the GPU request-state cursor materialized by prepare_inputs(). + # The CPU cursor can lag after speculative rejection. + torch.index_select( + global_batch.positions, + 0, + local_gather_idx, + out=input_buffers.positions[:num_local_tokens_padded], + ) local_query_start_loc_np = np.empty( input_buffers.max_num_reqs + 1, dtype=np.int32 @@ -478,17 +480,16 @@ def partition_batch( local_to_global_req_idx = async_copy_to_gpu( local_to_global_req_idx_np, device=self.device ) - local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device) - - assert self._local_req_idx is not None - prepare_pos_seq_lens( - self._local_req_idx[:num_local_reqs], - local_query_start_loc, - local_start_pos, - input_buffers.positions, - input_buffers.seq_lens[:num_local_reqs], - ) seq_lens = input_buffers.seq_lens[:num_local_reqs] + if num_local_tokens > 0: + local_end_positions = torch.index_select( + input_buffers.positions, + 0, + local_query_start_loc[1:] - 1, + ) + seq_lens.copy_(local_end_positions + 1) + else: + seq_lens.zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -511,18 +512,9 @@ def partition_batch( cu_num_logits = torch.zeros( num_local_reqs + 1, device=self.device, dtype=torch.int32 ) - logits_indices = combine_sampled_and_draft_tokens( - input_buffers.input_ids, - local_to_global_req_idx, - req_states.last_sampled_tokens, - local_query_start_loc, - seq_lens, - req_states.prefill_len.gpu, - req_states.draft_tokens, - cu_num_logits, - total_num_logits, - 1, - ) + # Local logits are never sampled. The complete hidden-state tensor is + # restored first and sampled with the untouched global InputBatch. + logits_indices = local_query_start_loc[1:] - 1 local_prefill_len_np = global_batch.prefill_len_np[ local_to_global_batch_req_idx_np @@ -536,7 +528,7 @@ def partition_batch( seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens - return replace( + self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, @@ -570,6 +562,7 @@ def partition_batch( cu_num_logits_np=cu_num_logits_np, prompt_lens=None, ) + return self._local_batch def prepare_attn( self, input_batch: InputBatch @@ -640,6 +633,46 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: gathered = get_pcp_group().all_gather(hidden_states, dim=0) return gathered[self._hidden_restore_idx] + def get_draft_input_buffers( + self, input_buffers: InputBuffers + ) -> InputBatch | InputBuffers: + return self.draft_prefill_batch or input_buffers + + def prepare_draft_prefill( + self, + input_batch: InputBatch, + input_ids: torch.Tensor, + ) -> None: + if input_batch is not self._global_batch or self._local_batch is None: + return + local_batch = self._local_batch + assert self._local_gather_idx is not None + num_local_tokens = self._local_gather_idx.shape[0] + torch.index_select( + input_ids, + 0, + self._local_gather_idx, + out=local_batch.input_ids[:num_local_tokens], + ) + self.draft_prefill_batch = local_batch + + def restore_draft_prefill( + self, + last_hidden_states: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.draft_prefill_batch is None: + return last_hidden_states, hidden_states + local_last_hidden_states = last_hidden_states + last_hidden_states = self.restore_hidden_states(local_last_hidden_states) + hidden_states = ( + last_hidden_states + if local_last_hidden_states is hidden_states + else self.restore_hidden_states(hidden_states) + ) + self.draft_prefill_batch = None + return last_hidden_states, hidden_states + def restore_for_sampling( self, hidden_states: torch.Tensor, @@ -686,7 +719,6 @@ def maybe_build_pcp_manager( vllm_config: VllmConfig, device: torch.device, supports_mm_inputs: bool, - req_states: RequestState, block_tables: BlockTables, cls: type[PCPManager] = PCPManager, ) -> PCPManager | None: @@ -705,7 +737,6 @@ def maybe_build_pcp_manager( pcp_world_size=pcp_size, pcp_rank=pcp_rank, device=device, - req_states=req_states, max_num_reqs=vllm_config.scheduler_config.max_num_seqs, max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens, block_tables=block_tables, diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index bba4e2b4bd31..cab250ff4e68 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -250,8 +250,6 @@ def propose( ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens_padded].copy_(hidden_states) - self._copy_request_inputs( num_reqs, input_batch.idx_mapping, @@ -272,6 +270,17 @@ def propose( self.max_num_reqs, ) + if self.pcp_manager is not None: + self.pcp_manager.prepare_draft_prefill( + input_batch, + self.input_buffers.input_ids[:num_tokens_padded], + ) + prefill = self.pcp_manager.draft_prefill_batch + if prefill is not None: + num_tokens = prefill.num_tokens + num_tokens_padded = prefill.num_tokens_after_padding + self.hidden_states[:num_tokens_padded].copy_(hidden_states) + # When all requests are decoding (no true prefills), each has # num_speculative_steps + 1 tokens, enabling FULL graph replay. uniform_token_count = get_uniform_decode_token_count( @@ -324,6 +333,11 @@ def propose( # Early exit. return self.draft_tokens[:num_reqs, :1] + if self.pcp_manager is not None and not dummy_run: + self.block_tables.gather_block_tables( + input_batch.idx_mapping, num_reqs_padded=num_reqs + ) + # Prepare the inputs for the decode steps. prepare_decode_inputs( self.draft_tokens[:num_reqs, 0], @@ -387,6 +401,11 @@ def _run_model( cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + input_buffers: InputBatch | InputBuffers = self.input_buffers + is_padding = None + if self.pcp_manager is not None: + input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers) + is_padding = input_buffers.is_padding[:num_tokens] batch_descriptor = BatchDescriptor(num_tokens=num_tokens) with set_forward_context( attn_metadata, @@ -396,6 +415,7 @@ def _run_model( num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, + is_padding=is_padding, ): inputs_embeds = None if self.supports_mm_inputs: @@ -406,15 +426,15 @@ def _run_model( is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens ) self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids( - self.input_buffers.input_ids[:num_input_tokens], + input_buffers.input_ids[:num_input_tokens], multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed, ) inputs_embeds = self.inputs_embeds[:num_tokens] model_inputs = dict( - input_ids=self.input_buffers.input_ids[:num_tokens], - positions=self.input_buffers.positions[:num_tokens], + input_ids=input_buffers.input_ids[:num_tokens], + positions=input_buffers.positions[:num_tokens], hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) @@ -463,6 +483,10 @@ def _prefill( cudagraph_runtime_mode=cudagraph_runtime_mode, mm_inputs=mm_inputs, ) + if self.pcp_manager is not None: + last_hidden_states, hidden_states = self.pcp_manager.restore_draft_prefill( + last_hidden_states, hidden_states + ) sample_hidden_states = last_hidden_states[last_token_indices] self.draft_tokens[:num_reqs, 0] = self.sample_draft( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a3e8df79243a..845781fab227 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -38,6 +38,13 @@ class DFlashSpeculator(DraftModelSpeculator): _speculator_name = "DFlash" # For logging, so we can share methods with subclasses def __init__(self, vllm_config: VllmConfig, device: torch.device): + parallel_config = vllm_config.parallel_config + if parallel_config.prefill_context_parallel_size > 1: + vllm_config = copy.copy(vllm_config) + vllm_config.parallel_config = replace( + parallel_config, + prefill_context_parallel_size=1, + ) super().__init__(vllm_config, device) self.hidden_states = torch.zeros( @@ -388,6 +395,11 @@ def propose( ) return self.draft_tokens[:num_reqs] + if self.pcp_manager is not None and not dummy_run: + self.block_tables.gather_block_tables( + input_batch.idx_mapping, num_reqs_padded=num_reqs + ) + # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. assert self.draft_kv_cache_group_id >= 0 diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index 4889e9a750f8..d94702612d9b 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -28,7 +28,8 @@ def load_draft_model( # toggles skip_topk so step 0 computes MTP's own indices and # steps 1+ reuse them. self.share_mtp_topk_indices = ( - getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + self.vllm_config.parallel_config.prefill_context_parallel_size == 1 + and getattr(draft_hf_config, "index_share_for_mtp_iteration", False) and hasattr(draft_model.model, "set_skip_topk") and hasattr(draft_model.model, "compact_topk_indices") ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 3c510e88850c..2681677c7603 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import replace -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -29,6 +29,9 @@ from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.utils import AttentionGroup +if TYPE_CHECKING: + from vllm.v1.worker.gpu.pcp_manager import PCPManager + logger = init_logger(__name__) @@ -149,6 +152,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.arange = torch.arange( self.max_num_reqs + 1, dtype=torch.int32, device="cpu" ) + self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) self.draft_logits: torch.Tensor | None = None if self.speculative_config.draft_sample_method == "probabilistic": @@ -166,6 +170,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) self.supports_mm_inputs = False + self.pcp_manager: PCPManager | None = None @abstractmethod def load_draft_model( @@ -326,6 +331,7 @@ def _build_draft_attn_metadata( kv_cache_config=self.kv_cache_config, causal=causal, seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, + is_prefilling=self.draft_is_prefilling[:num_reqs], ) return attn_metadata From ae48466cf33b4d7cfbe0f32370c425a3aa9fde61 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Fri, 11 Sep 2026 07:19:57 +0200 Subject: [PATCH 023/774] [Perf][ROCm] Add AITER custom AG/RS (DP only) (#48247) Signed-off-by: simondanielsson --- .../kernels/benchmark_device_communicators.py | 255 ++++++++++++++---- tests/distributed/test_comm_ops.py | 15 ++ .../distributed/test_rocm_aiter_custom_ar.py | 172 +++++++++++- tests/utils.py | 62 ++++- .../aiter_custom_all_reduce.py | 14 + .../device_communicators/cuda_communicator.py | 87 +++++- vllm/distributed/parallel_state.py | 22 +- vllm/envs.py | 2 +- 8 files changed, 552 insertions(+), 77 deletions(-) diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 24e22023b91d..6c3a726a9753 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -29,6 +29,10 @@ import torch.distributed as dist from torch.distributed import ProcessGroup +from vllm._aiter_ops import rocm_aiter_ops +from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( + AiterCustomAllreduce, +) from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce from vllm.distributed.device_communicators.flashinfer_all_reduce import ( FlashInferAllReduce, @@ -42,6 +46,7 @@ ) from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser logger = init_logger(__name__) @@ -56,6 +61,30 @@ # CUDA graph settings CUDA_GRAPH_CAPTURE_CYCLES = 10 +# All-gather and reduce-scatter are distinct collectives from all-reduce, so +# results are compared within a collective (each against its own NCCL/RCCL +# baseline) rather than across collectives. +COMM_COLLECTIVE = { + "ca_1stage": "all-reduce", + "ca_2stage": "all-reduce", + "pynccl": "all-reduce", + "pynccl-symm": "all-reduce", + "symm_mem_multimem": "all-reduce", + "symm_mem_two_shot": "all-reduce", + "flashinfer_trtllm": "all-reduce", + "flashinfer_mnnvl": "all-reduce", + "aiter_ag": "all-gather", + "pynccl_ag": "all-gather", + "aiter_rs": "reduce-scatter", + "pynccl_rs": "reduce-scatter", +} +COLLECTIVE_ORDER = ["all-reduce", "all-gather", "reduce-scatter"] +COLLECTIVE_BASELINE = { + "all-reduce": "pynccl", + "all-gather": "pynccl_ag", + "reduce-scatter": "pynccl_rs", +} + class CommunicatorBenchmark: """Benchmark class for testing device communicators.""" @@ -77,6 +106,11 @@ def __init__( max_seq_len = max(sequence_lengths) max_tensor_elements = max_seq_len * HIDDEN_SIZE self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 + # AITER has effective max size of half the max size, so we double it, and + # account for AG output size (which is world_size larger) + self.aiter_max_size = ( + 2 * world_size * max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 + ) # Initialize communicators self.custom_allreduce = None @@ -85,6 +119,7 @@ def __init__( self.symm_mem_comm_multimem = None self.symm_mem_comm_two_shot = None self.fi_ar_comm = None + self.aiter_comm = None self._init_communicators() @@ -181,6 +216,27 @@ def _init_communicators(self): ) self.fi_ar_comm = None + if current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled(): + try: + # Also used for AG/RS for DP + self.aiter_comm = AiterCustomAllreduce( + group=self.cpu_group, + device=self.device, + max_size=self.aiter_max_size, + ) + if not self.aiter_comm.disabled: + logger.info("Rank %s: AITER custom AG/RS initialized", self.rank) + else: + logger.info("Rank %s: AITER custom AG/RS disabled", self.rank) + self.aiter_comm = None + except Exception as e: + logger.warning( + "Rank %s: Failed to initialize AITER custom AG/RS: %s", + self.rank, + e, + ) + self.aiter_comm = None + def benchmark_allreduce( self, sequence_length: int, num_warmup: int, num_trials: int ) -> dict[str, float]: @@ -385,29 +441,116 @@ def benchmark_allreduce_single( f"CUDA graph benchmark failed for communicator: {e}" ) from e + def benchmark_ag_rs( + self, sequence_length: int, num_warmup: int, num_trials: int + ) -> dict[str, float]: + """Benchmark all-gather and reduce-scatter: custom (AITER) vs NCCL/RCCL. + + Each collective is measured for both the AITER custom kernel and the + PyNccl baseline so they compare like-for-like (print_results groups + columns by collective). All-gather maps ``(seq, hidden)`` -> + ``(world_size * seq, hidden)``; reduce-scatter maps ``(seq, hidden)`` -> + ``(seq // world_size, hidden)`` and so only runs when ``seq`` divides + evenly across ranks. Both reuse ``benchmark_allreduce_single`` timed on a + ``(seq, hidden)`` input; baselines write into a pre-allocated buffer. + """ + results: dict[str, float] = {} + entries = [] + + # ---- all-gather: (seq, hidden) -> (world_size * seq, hidden) ---- + if self.aiter_comm is not None: + c = self.aiter_comm + entries.append( + ( + "aiter_ag", + lambda t, c=c: c.custom_all_gather(t, dim=0), + lambda t, c=c: c.should_custom_ag(t), + c.capture(), + ) + ) + if self.pynccl_comm is not None: + c = self.pynccl_comm + ag_out = torch.empty( + sequence_length * self.world_size, + HIDDEN_SIZE, + dtype=BENCHMARK_DTYPE, + device=self.device, + ) + entries.append( + ( + "pynccl_ag", + lambda t, c=c, o=ag_out: c.all_gather(o, t), + lambda t: True, + nullcontext(), + ) + ) + + # ---- reduce-scatter: (seq, hidden) -> (seq // world_size, hidden) ---- + if sequence_length % self.world_size == 0: + rs_shape = (sequence_length // self.world_size, HIDDEN_SIZE) + if self.aiter_comm is not None: + c = self.aiter_comm + rs_out = torch.empty( + rs_shape, dtype=BENCHMARK_DTYPE, device=self.device + ) + entries.append( + ( + "aiter_rs", + lambda t, c=c, o=rs_out: c.custom_reduce_scatter(t, o, dim=0), + lambda t, c=c: c.should_custom_rs(t, dim=0), + c.capture(), + ) + ) + if self.pynccl_comm is not None: + c = self.pynccl_comm + rs_out = torch.empty( + rs_shape, dtype=BENCHMARK_DTYPE, device=self.device + ) + entries.append( + ( + "pynccl_rs", + lambda t, c=c, o=rs_out: c.reduce_scatter(o, t), + lambda t: True, + nullcontext(), + ) + ) -def _calculate_speedup_info(comm_results: dict[str, float]) -> str: - """Calculate speedup information for a single tensor size.""" + for name, fn, should_use_fn, context in entries: + latency = self.benchmark_allreduce_single( + sequence_length, fn, should_use_fn, context, num_warmup, num_trials + ) + if latency is not None: + results[name] = latency + + return results + + +def _calculate_speedup_info(comm_results: dict[str, float], baseline: str) -> str: + """Fastest comm in a collective group and its speedup vs ``baseline``.""" if not comm_results: return "N/A" - # Find the fastest communicator fastest_comm = min(comm_results.keys(), key=lambda k: comm_results[k]) fastest_time = comm_results[fastest_comm] - # Calculate speedup vs PyNccl if available - if "pynccl" in comm_results: - pynccl_time = comm_results["pynccl"] - speedup = pynccl_time / fastest_time + if baseline in comm_results and fastest_time > 0: + speedup = comm_results[baseline] / fastest_time return f"{fastest_comm} ({speedup:.2f}x)" - else: - return f"{fastest_comm} (N/A)" + return f"{fastest_comm} (N/A)" def print_results( - results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int + results: dict[int, dict[str, float]], sequence_lengths: list[int], world_size: int ): - """Print benchmark results in a formatted table.""" + """Print benchmark results as one table per collective. + + All-reduce, all-gather and reduce-scatter are different collectives, so each + gets its own table and its speedup is computed against that collective's + NCCL/RCCL baseline rather than across collectives. + """ + present: set[str] = set() + for size_results in results.values(): + present.update(size_results.keys()) print(f"\n{'=' * 130}") print("Device Communicator Benchmark Results") @@ -415,53 +558,45 @@ def print_results( f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, " f"Hidden Size: {HIDDEN_SIZE}" ) + print("Tensor Shape is the per-rank input; times are ms per operation.") print(f"{'=' * 130}") - # Get all communicator names - all_comms = set() - for size_results in results.values(): - all_comms.update(size_results.keys()) - - all_comms = sorted(list(all_comms)) - - # Print header - header = f"{'Tensor Shape':<20}{'Tensor Size':<15}" - for comm in all_comms: - header += f"{comm:<20}" - header += f"{'Best (Speedup vs PyNccl)':<30}" - print(header) - print("-" * len(header)) - - # Print results for each sequence length - for seq_len in sequence_lengths: - if seq_len in results: - # Calculate tensor size in elements and bytes - tensor_elements = seq_len * HIDDEN_SIZE - tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize - - # Format tensor size (MB) - tensor_size_mb = tensor_bytes / (1024 * 1024) - tensor_size_str = f"{tensor_size_mb:.2f} MB" - - # Format tensor shape + for collective in COLLECTIVE_ORDER: + comms = sorted(c for c in present if COMM_COLLECTIVE.get(c) == collective) + if not comms: + continue + baseline = COLLECTIVE_BASELINE[collective] + + print(f"\n{collective} (speedup vs {baseline})") + header = f"{'Tensor Shape':<20}{'Tensor Size':<15}" + for comm in comms: + header += f"{comm:<20}" + header += f"{'Best':<30}" + print(header) + print("-" * len(header)) + + for seq_len in sequence_lengths: + if seq_len not in results: + continue + row_results = { + c: results[seq_len][c] for c in comms if c in results[seq_len] + } + + tensor_bytes = seq_len * HIDDEN_SIZE * BENCHMARK_DTYPE.itemsize + tensor_size_str = f"{tensor_bytes / (1024 * 1024):.2f} MB" tensor_shape = f"({seq_len}, {HIDDEN_SIZE})" row = f"{tensor_shape:<20}{tensor_size_str:<15}" - for comm in all_comms: - if comm in results[seq_len]: - row += f"{results[seq_len][comm]:<20.3f}" + for comm in comms: + if comm in row_results: + row += f"{row_results[comm]:<20.3f}" else: row += f"{'N/A':<20}" - - # Calculate speedup information - speedup_info = _calculate_speedup_info(results[seq_len]) - row += f"{speedup_info:<30}" - + row += f"{_calculate_speedup_info(row_results, baseline):<30}" print(row) - print(f"{'=' * 130}") - print("All times are in milliseconds (ms) per allreduce operation") - print("Speedup column shows: fastest_algorithm (speedup_vs_pynccl)") + print(f"\n{'=' * 130}") + print("All times are in milliseconds (ms) per operation") def main(): @@ -526,6 +661,13 @@ def main(): num_warmup=args.num_warmup, num_trials=args.num_trials, ) + results.update( + benchmark.benchmark_ag_rs( + sequence_length=seq_len, + num_warmup=args.num_warmup, + num_trials=args.num_trials, + ) + ) all_results[seq_len] = results @@ -538,12 +680,23 @@ def main(): # Save to JSON if requested if args.output_json: - # Add speedup information to results + # Add per-collective speedup information to results. enhanced_results = {} for seq_len, comm_results in all_results.items(): + speedup_info = {} + for collective in COLLECTIVE_ORDER: + group = { + k: v + for k, v in comm_results.items() + if COMM_COLLECTIVE.get(k) == collective + } + if group: + speedup_info[collective] = _calculate_speedup_info( + group, COLLECTIVE_BASELINE[collective] + ) enhanced_results[seq_len] = { "timings": comm_results, - "speedup_info": _calculate_speedup_info(comm_results), + "speedup_info": speedup_info, } output_data = { diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index b5e552af4974..f35ba14c1f6b 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -410,6 +410,21 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> nccl_selector.assert_not_called() +def test_aiter_all_gather_precedes_pynccl(monkeypatch: pytest.MonkeyPatch) -> None: + """Ensure can use aiter all_gather if available even if pynccl is not available.""" + output = torch.empty(2) + aiter_comm = Mock() + aiter_comm.should_custom_ag.return_value = True + aiter_comm.custom_all_gather.return_value = output + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.world_size = 2 + communicator.aiter_ar_comm = aiter_comm + communicator.pynccl_comm = Mock(disabled=True) + monkeypatch.setattr(communicator, "_can_use_aiter_ag_rs", Mock(return_value=True)) + + assert communicator.all_gatherv(torch.empty(1)) is output + + def test_isend_object_posts_size_then_object_and_releases_on_wait( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/distributed/test_rocm_aiter_custom_ar.py b/tests/distributed/test_rocm_aiter_custom_ar.py index 0b85f36410d7..65018bfc4e5d 100644 --- a/tests/distributed/test_rocm_aiter_custom_ar.py +++ b/tests/distributed/test_rocm_aiter_custom_ar.py @@ -8,7 +8,7 @@ from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa -from vllm.distributed.parallel_state import get_tp_group, graph_capture +from vllm.distributed.parallel_state import get_dp_group, get_tp_group, graph_capture from vllm.envs import disable_envs_cache from vllm.platforms import current_platform @@ -132,3 +132,173 @@ def test_rocm_aiter_custom_allreduce( test_target, ): multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) + + +def _get_aiter_ag_rs_comm(): + device_communicator = get_dp_group().device_communicator + assert device_communicator.use_aiter_ag_rs, ( + "AITER custom AG/RS was not enabled on the DP group." + ) + aiter_comm = device_communicator.aiter_ar_comm + assert aiter_comm is not None, "AITER custom AG/RS was not initialized." + assert not aiter_comm.disabled, "AITER custom AG/RS is disabled." + return aiter_comm + + +def _assert_aiter_handles_ag(aiter_comm, inp: torch.Tensor) -> None: + assert aiter_comm.should_custom_ag(inp), ( + f"AITER custom all-gather does not support input shape {inp.shape}." + ) + + +def _assert_aiter_handles_rs(aiter_comm, inp: torch.Tensor) -> None: + assert aiter_comm.should_custom_rs(inp, dim=0), ( + f"AITER custom reduce-scatter does not support input shape {inp.shape}." + ) + + +@ray.remote(num_gpus=1, max_calls=1) +def eager_ag_rs( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, + data_parallel_size, + data_parallel_master_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + data_parallel_size=data_parallel_size, + data_parallel_master_port=data_parallel_master_port, + ) + + dp_group = get_dp_group() + group = dp_group.device_group + dp_world = dp_group.world_size + aiter_comm = _get_aiter_ag_rs_comm() + + for shape, dtype in test_cases: + num_tokens, hidden = shape + + # all-gather: each rank contributes (num_tokens, hidden). + inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1) + _assert_aiter_handles_ag(aiter_comm, inp) + expected = torch.empty( + (num_tokens * dp_world, hidden), dtype=dtype, device=device + ) + dist.all_gather_into_tensor(expected, inp, group=group) + out = aiter_comm.custom_all_gather(inp, dim=0) + assert out is not None + torch.testing.assert_close(out, expected) + + # reduce-scatter: each rank contributes (num_tokens * dp, hidden). + rs_in = torch.ones( + (num_tokens * dp_world, hidden), dtype=dtype, device=device + ) * (rank + 1) + _assert_aiter_handles_rs(aiter_comm, rs_in) + rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device) + dist.reduce_scatter_tensor(rs_expected, rs_in, group=group) + rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device) + aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0) + torch.testing.assert_close(rs_out, rs_expected) + + +@ray.remote(num_gpus=1, max_calls=1) +def graph_ag_rs( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, + data_parallel_size, + data_parallel_master_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + data_parallel_size=data_parallel_size, + data_parallel_master_port=data_parallel_master_port, + ) + + dp_group = get_dp_group() + group = dp_group.device_group + dp_world = dp_group.world_size + aiter_comm = _get_aiter_ag_rs_comm() + + # Warmup so DP comms is initialized before graph capture + data = torch.zeros(1, device=device) + dist.all_reduce(data, group=group) + torch.accelerator.synchronize() + del data + + for shape, dtype in test_cases: + num_tokens, hidden = shape + + # all-gather under graph capture. + inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1) + _assert_aiter_handles_ag(aiter_comm, inp) + ag_expected = torch.empty( + (num_tokens * dp_world, hidden), dtype=dtype, device=device + ) + dist.all_gather_into_tensor(ag_expected, inp, group=group) + with graph_capture(device=device) as graph_capture_context: + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + ag_out = aiter_comm.custom_all_gather(inp, dim=0) + graph.replay() + torch.testing.assert_close(ag_out, ag_expected) + + # reduce-scatter under graph capture. + rs_in = torch.ones( + (num_tokens * dp_world, hidden), dtype=dtype, device=device + ) * (rank + 1) + _assert_aiter_handles_rs(aiter_comm, rs_in) + rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device) + dist.reduce_scatter_tensor(rs_expected, rs_in, group=group) + rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device) + with graph_capture(device=device) as graph_capture_context: + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0) + graph.replay() + torch.testing.assert_close(rs_out, rs_expected) + + +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("tp_size", [1]) +@pytest.mark.parametrize("pipeline_parallel_size", [1]) +@pytest.mark.parametrize("data_parallel_size", [2]) +@pytest.mark.parametrize("test_target", [eager_ag_rs, graph_ag_rs]) +def test_rocm_aiter_custom_ag_rs( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pipeline_parallel_size, + data_parallel_size, + test_target, +): + multi_process_parallel( + monkeypatch, + tp_size, + pipeline_parallel_size, + test_target, + data_parallel_size=data_parallel_size, + ) diff --git a/tests/utils.py b/tests/utils.py index 0c2cfeb6a0b9..e0166785ba5f 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1416,6 +1416,8 @@ def init_test_distributed_environment( rank: int, distributed_init_port: str, local_rank: int = -1, + data_parallel_size: int = 1, + data_parallel_master_port: int | None = None, ) -> None: # Note: This function is often called from Ray worker processes, so we # can't rely on pytest fixtures to set the config. We check if the config @@ -1428,6 +1430,29 @@ def init_test_distributed_environment( distributed_init_method = f"tcp://localhost:{distributed_init_port}" + if data_parallel_size > 1: + # For DP we need to set a common DP master port + from vllm.config.parallel import ParallelConfig + + assert data_parallel_master_port is not None, ( + "data_parallel_master_port is required when data_parallel_size > 1" + ) + tp_pp_world = tp_size * pp_size + parallel_config = ParallelConfig( + data_parallel_size=data_parallel_size, + data_parallel_rank=rank // tp_pp_world, + _data_parallel_master_port_list=[int(data_parallel_master_port)], + ) + with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)): + init_distributed_environment( + world_size=tp_pp_world, + rank=rank % tp_pp_world, + distributed_init_method=distributed_init_method, + local_rank=local_rank if local_rank >= 0 else rank, + ) + ensure_model_parallel_initialized(tp_size, pp_size) + return + if get_current_vllm_config_or_none() is not None: # Config already set, use it directly init_distributed_environment( @@ -1454,6 +1479,7 @@ def multi_process_parallel( tp_size: int, pp_size: int, test_target: Any, + data_parallel_size: int = 1, ) -> None: import ray @@ -1474,18 +1500,34 @@ def multi_process_parallel( ) distributed_init_port = get_open_port() + # Separate port for the DP master group; only used when data_parallel_size > 1. + data_parallel_master_port = get_open_port() if data_parallel_size > 1 else None + world_size = data_parallel_size * tp_size * pp_size try: refs = [] - for rank in range(tp_size * pp_size): - refs.append( - test_target.remote( - monkeypatch, - tp_size, - pp_size, - rank, - distributed_init_port, - ), - ) + for rank in range(world_size): + if data_parallel_size > 1: + refs.append( + test_target.remote( + monkeypatch, + tp_size, + pp_size, + rank, + distributed_init_port, + data_parallel_size, + data_parallel_master_port, + ), + ) + else: + refs.append( + test_target.remote( + monkeypatch, + tp_size, + pp_size, + rank, + distributed_init_port, + ), + ) ray.get(refs) finally: ray.shutdown() diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index be2c901281bd..3247b33262e2 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -56,6 +56,20 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool: def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None: return self._impl.custom_all_reduce(inp) + def should_custom_ag(self, inp: torch.Tensor) -> bool: + return self._impl.should_custom_ag(inp) + + def should_custom_rs(self, inp: torch.Tensor, dim: int) -> bool: + return self._impl.should_custom_rs(inp, dim) + + def custom_all_gather(self, inp: torch.Tensor, dim: int = 0) -> torch.Tensor | None: + return self._impl.custom_all_gather(inp, dim=dim) + + def custom_reduce_scatter( + self, inp: torch.Tensor, out: torch.Tensor, dim: int = 0 + ) -> torch.Tensor | None: + return self._impl.custom_reduce_scatter(inp, out, dim=dim) + def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool: """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel. diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index faa9fd0987da..4afc28f65214 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -109,6 +109,7 @@ def __init__( self.fi_ar_comm: FlashInferAllReduce | None = None self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None self.aiter_ar_comm: AiterCustomAllreduce | None = None + self.use_aiter_ag_rs: bool = False if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( @@ -149,6 +150,22 @@ def __init__( ), ) + # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine + if ( + "dp" in unique_name + and self.world_size in (2, 4, 8) + and current_platform.is_rocm() + and rocm_aiter_ops.is_custom_all_reduce_enabled() + ): + self.aiter_ar_comm = AiterCustomAllreduce( + group=self.cpu_group, + device=self.device, + ) + if self.aiter_ar_comm.disabled: + self.aiter_ar_comm = None + else: + self.use_aiter_ag_rs = True + if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm(): # Initialize a custom quick all-reduce implementation for AMD. # Quick reduce is designed as a complement to custom allreduce @@ -285,7 +302,11 @@ def _log_all_reduce_backend_selection(self) -> None: enabled_ar_backends.append("NCCL_SYMM_MEM") if self.qr_comm is not None and not self.qr_comm.disabled: enabled_ar_backends.append("QUICK_REDUCE") - if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled: + if ( + self.use_aiter_allreduce + and self.aiter_ar_comm is not None + and not self.aiter_ar_comm.disabled + ): enabled_ar_backends.append("AITER_CUSTOM") if self.ca_comm is not None and not self.ca_comm.disabled: enabled_ar_backends.append("CUSTOM") @@ -340,7 +361,8 @@ def all_reduce(self, input_): return out aiter_ar_comm = self.aiter_ar_comm if ( - aiter_ar_comm is not None + self.use_aiter_allreduce + and aiter_ar_comm is not None and not aiter_ar_comm.disabled and aiter_ar_comm.should_custom_ar(input_) ): @@ -462,6 +484,11 @@ def reduce_scatterv( # Convert negative dim to positive. dim += input_.dim() + # 'sizes' is not needed if all inputs in the same group have the same + # shape + if sizes is not None and all(s == sizes[0] for s in sizes): + sizes = None + # Note: This will produce an incorrect answer if we don't make # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? input_tensor = input_.movedim(0, dim).contiguous() @@ -475,6 +502,16 @@ def reduce_scatterv( chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] + if self._can_use_aiter_ag_rs(sizes): + aiter_comm = self.aiter_ar_comm + assert aiter_comm is not None + if aiter_comm.should_custom_rs(input_tensor, dim=0): + output = torch.empty( + output_shape, dtype=input_tensor.dtype, device=input_tensor.device + ) + aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0) + return output.movedim(0, dim).contiguous() + # Symmetric memory is only used when all ranks have uniform sizes. # ncclCommWindowRegister is collective: asymmetric pool allocations # from variable per-rank sizes cause deadlocks. @@ -632,6 +669,34 @@ def destroy(self): self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] + def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool: + """Whether the AITER custom AG/RS fast path may run for this collective. + + Requires: + - uniform batches + - FULL CUDAgraphs + """ + if ( + not self.use_aiter_ag_rs + or self.aiter_ar_comm is None + or self.aiter_ar_comm.disabled + ): + return False + if sizes is not None: + return False + + from vllm.config.compilation import CUDAGraphMode + from vllm.forward_context import get_forward_context + + try: + ctx = get_forward_context() + except AssertionError: + return False + if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL: + return False + bd = ctx.batch_descriptor + return bd is not None and bd.uniform + def suspend(self) -> None: if self.pynccl_comm is not None: self.pynccl_comm.suspend() @@ -665,14 +730,28 @@ def all_gatherv( if dim != 0: raise NotImplementedError("only dim 0 all-gatherv is supported") world_size = self.world_size - pynccl_comm = self.pynccl_comm - assert pynccl_comm is not None and not pynccl_comm.disabled # 'sizes' is not needed if all inputs in the same group have the same # shape if sizes is not None and all(s == sizes[0] for s in sizes): sizes = None + if self._can_use_aiter_ag_rs(sizes): + aiter_comm = self.aiter_ar_comm + assert aiter_comm is not None + if isinstance(input_, torch.Tensor): + if aiter_comm.should_custom_ag(input_): + out = aiter_comm.custom_all_gather(input_, dim=0) + if out is not None: + return out + elif all(aiter_comm.should_custom_ag(inp) for inp in input_): + outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_] + if all(o is not None for o in outs): + return outs + + pynccl_comm = self.pynccl_comm + assert pynccl_comm is not None and not pynccl_comm.disabled + # Symmetric memory is only used when all ranks have uniform sizes. # ncclCommWindowRegister is collective: asymmetric pool allocations # from variable per-rank sizes cause deadlocks. diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index fd068dff1cea..b24ff64d4d92 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -673,11 +673,11 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None else: stream = graph_capture_context.stream - # only cuda uses this function, + # only cuda/rocm uses this function, # so we don't abstract it into the base class maybe_ca_context = nullcontext() maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext() - maybe_aiter_context = nullcontext() + maybe_aiter_ar_context = nullcontext() from vllm.distributed.device_communicators.cuda_communicator import ( CudaCommunicator, ) @@ -698,12 +698,10 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None if fi_pcie_ipc_ar_comm is not None: maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture() - from vllm._aiter_ops import rocm_aiter_ops - - if rocm_aiter_ops.is_enabled(): - aiter_ar = rocm_aiter_ops.get_aiter_allreduce() - if aiter_ar is not None: - maybe_aiter_context = aiter_ar.capture() # type: ignore + # Capture each group's own comm. A global lookup would double-capture + aiter_ar_comm = getattr(self.device_communicator, "aiter_ar_comm", None) + if aiter_ar_comm is not None: + maybe_aiter_ar_context = aiter_ar_comm.capture() # type: ignore # ensure all initialization operations complete before attempting to # capture the graph on another stream @@ -715,7 +713,7 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None torch.cuda.stream(stream), maybe_ca_context, maybe_fi_pcie_ipc_context, - maybe_aiter_context, + maybe_aiter_ar_context, ): yield graph_capture_context @@ -1640,7 +1638,11 @@ def graph_capture( context = graph_capture_context or GraphCaptureContext( torch.cuda.Stream(device=device) ) - with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context): + with ( + get_tp_group().graph_capture(context), + get_pp_group().graph_capture(context), + get_dp_group().graph_capture(context), + ): yield context diff --git a/vllm/envs.py b/vllm/envs.py index d429686d9cd1..18e04e3dc6a2 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1261,7 +1261,7 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1") ), # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's - # CudaCommunicator on ROCm. + # CudaCommunicator on ROCm. Also enables AITER AG/RS for DP communication. "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1") ), From 5ff50f3996d5b0f59dd48c73c1ec0f7d912b6466 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hung=20Hoang=20=28=E9=BB=84=E5=85=B4=EF=BC=89?= Date: Fri, 11 Sep 2026 12:25:51 +0700 Subject: [PATCH 024/774] [Bugfix][MM] Fix swapped H/W in dummy video profiling inputs (#56385) Signed-off-by: hungh Co-authored-by: Cursor Agent --- vllm/multimodal/processing/dummy_inputs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/multimodal/processing/dummy_inputs.py b/vllm/multimodal/processing/dummy_inputs.py index 29f373ccc897..369d4f41a4c0 100644 --- a/vllm/multimodal/processing/dummy_inputs.py +++ b/vllm/multimodal/processing/dummy_inputs.py @@ -193,5 +193,5 @@ def _get_dummy_videos( height, ) height = min(height, overrides.height) - video = np.full((num_frames, width, height, 3), 255, dtype=np.uint8) + video = np.full((num_frames, height, width, 3), 255, dtype=np.uint8) return [video] * num_videos From fe4d4109c24063fb5e60ede9d48a1c0b3e9e4cea Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 11 Sep 2026 13:33:26 +0800 Subject: [PATCH 025/774] [Frontend][last/N] Move all non-OpenAI content out of the OpenAI folder. (#56369) Signed-off-by: wang.yuqi --- .../openai/test_async_tokenization.py | 82 ------------------- .../utils}/test_sse_keep_alive.py | 2 +- .../openai/chat_completion/api_router.py | 2 +- .../openai/completion/api_router.py | 2 +- .../{openai => serve/utils}/sse_keep_alive.py | 0 5 files changed, 3 insertions(+), 85 deletions(-) delete mode 100644 tests/entrypoints/openai/test_async_tokenization.py rename tests/entrypoints/{openai => serve/utils}/test_sse_keep_alive.py (99%) rename vllm/entrypoints/{openai => serve/utils}/sse_keep_alive.py (100%) diff --git a/tests/entrypoints/openai/test_async_tokenization.py b/tests/entrypoints/openai/test_async_tokenization.py deleted file mode 100644 index 1d3d110d3027..000000000000 --- a/tests/entrypoints/openai/test_async_tokenization.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import random -from collections.abc import Callable - -import openai -import pytest -import pytest_asyncio - -from tests.utils import RemoteOpenAIServer - -MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" - - -@pytest.fixture(scope="module") -def server(): - args = [ - # use half precision for speed and memory savings in CI environment - "--dtype", - "bfloat16", - "--max-model-len", - "8192", - "--enforce-eager", - "--max-num-seqs", - "128", - "--load-format", - "dummy", - ] - - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: - yield remote_server - - -@pytest_asyncio.fixture -async def client(server): - async with server.get_async_client() as async_client: - yield async_client - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ids=["completion", "chat"], - argnames=["create_func_gen", "content_body"], - argvalues=[ - (lambda x: x.completions.create, {"prompt": " ".join(["A"] * 10_000)}), - ( - lambda x: x.chat.completions.create, - {"messages": [{"role": "user", "content": " ".join(["A"] * 10_000)}]}, - ), - ], -) -async def test_with_and_without_truncate( - server: RemoteOpenAIServer, - client: openai.AsyncOpenAI, - create_func_gen: Callable, - content_body: dict, -): - create_func = create_func_gen(client) - body = {"model": MODEL_NAME, **content_body, "max_tokens": 10} - - num_requests = 10 - truncate_prompt_tokens = [1000] * (num_requests // 2) + [None] * ( - num_requests - num_requests // 2 - ) - random.shuffle(truncate_prompt_tokens) - - bodies = [ - {**body, "extra_body": {"truncate_prompt_tokens": t}} - for t in truncate_prompt_tokens - ] - - async def get_status_code(**kwargs): - try: - await create_func(**kwargs) - return 200 - except openai.APIStatusError as e: - return e.status_code - - responses = await asyncio.gather(*[get_status_code(**b) for b in bodies]) - assert 500 not in responses diff --git a/tests/entrypoints/openai/test_sse_keep_alive.py b/tests/entrypoints/serve/utils/test_sse_keep_alive.py similarity index 99% rename from tests/entrypoints/openai/test_sse_keep_alive.py rename to tests/entrypoints/serve/utils/test_sse_keep_alive.py index 74f5ee3ba234..519b9bb4a129 100644 --- a/tests/entrypoints/openai/test_sse_keep_alive.py +++ b/tests/entrypoints/serve/utils/test_sse_keep_alive.py @@ -6,7 +6,7 @@ import pytest from starlette.responses import StreamingResponse -from vllm.entrypoints.openai.sse_keep_alive import ( +from vllm.entrypoints.serve.utils.sse_keep_alive import ( SSE_KEEP_ALIVE_COMMENT, with_sse_keep_alive, ) diff --git a/vllm/entrypoints/openai/chat_completion/api_router.py b/vllm/entrypoints/openai/chat_completion/api_router.py index 291907a77a56..aecf5bba4483 100644 --- a/vllm/entrypoints/openai/chat_completion/api_router.py +++ b/vllm/entrypoints/openai/chat_completion/api_router.py @@ -14,7 +14,6 @@ ChatCompletionResponse, ) from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat -from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -22,6 +21,7 @@ with_cancellation, ) from vllm.entrypoints.serve.utils.orca_metrics import metrics_header +from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive from vllm.logger import init_logger logger = init_logger(__name__) diff --git a/vllm/entrypoints/openai/completion/api_router.py b/vllm/entrypoints/openai/completion/api_router.py index 7b33fcdb6cc0..fe909c7baee6 100644 --- a/vllm/entrypoints/openai/completion/api_router.py +++ b/vllm/entrypoints/openai/completion/api_router.py @@ -12,7 +12,6 @@ CompletionResponse, ) from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion -from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -20,6 +19,7 @@ with_cancellation, ) from vllm.entrypoints.serve.utils.orca_metrics import metrics_header +from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive from vllm.logger import init_logger logger = init_logger(__name__) diff --git a/vllm/entrypoints/openai/sse_keep_alive.py b/vllm/entrypoints/serve/utils/sse_keep_alive.py similarity index 100% rename from vllm/entrypoints/openai/sse_keep_alive.py rename to vllm/entrypoints/serve/utils/sse_keep_alive.py From fbf51c70269bb0d8333e14c5788e628e33c72f5d Mon Sep 17 00:00:00 2001 From: Zhicheng Liu <2084597944@qq.com> Date: Fri, 11 Sep 2026 14:07:52 +0800 Subject: [PATCH 026/774] [Bugfix] Fix FlashInfer KV sharing with omitted K/V (#55864) Signed-off-by: zhicheng Liu Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/models/gemma4_mtp.py | 12 ++---------- vllm/v1/attention/backends/flashinfer.py | 20 ++++++++++++-------- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 5e25a80d6bbd..639af0ea2951 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -246,16 +246,8 @@ def forward( q, _ = self.rotary_emb(positions, q, None) - # Attention reads K/V from the target's cache via KV sharing; - # these dummy tensors are never consumed but required by the API. - num_tokens = q.shape[0] - kv_dummy = torch.empty( - num_tokens, - self.num_kv_heads * self.head_dim, - dtype=q.dtype, - device=q.device, - ) - attn_output = self.attn(q, kv_dummy, kv_dummy) + # Attention reads K/V from the target's cache via KV sharing. + attn_output = self.attn(q, None, None) output, _ = self.o_proj(attn_output) return output diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 4b79135caba6..3547a45c9814 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1962,8 +1962,8 @@ def forward( self, layer: torch.nn.Module, query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, kv_cache: torch.Tensor, attn_metadata: FlashInferMetadata, output: torch.Tensor, @@ -1974,8 +1974,10 @@ def forward( Args: query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] + key: shape = [num_tokens, num_kv_heads, head_size], or None for a + KV-sharing decoder layer. + value: shape = [num_tokens, num_kv_heads, head_size], or None for a + KV-sharing decoder layer. kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size] attn_metadata: Metadata for attention. Returns: @@ -2066,8 +2068,6 @@ def forward( # Inputs and outputs may be padded for CUDA graphs query = query[:num_actual_tokens] - key = key[:num_actual_tokens] - value = value[:num_actual_tokens] output_padded = output output = output[:num_actual_tokens] @@ -2158,6 +2158,10 @@ def forward( prefill_wrapper = attn_metadata.prefill.wrapper assert prefill_wrapper is not None if use_dcp: + if key is None or value is None: + raise NotImplementedError( + "FlashInfer DCP prefill does not support KV-sharing layers" + ) assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper) assert prefill_wrapper._context._window_left == self.window_left assert prefill_wrapper._context._logits_soft_cap == ( @@ -2176,8 +2180,8 @@ def forward( layer, prefill_query, kv_cache_tuple, - key[num_decode_tokens:], - value[num_decode_tokens:], + key[num_decode_tokens:num_actual_tokens], + value[num_decode_tokens:num_actual_tokens], out=output[num_decode_tokens:], ) else: From 4b839c3378c6a40fa613328a9b97069ed83b181b Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 11 Sep 2026 08:19:55 +0200 Subject: [PATCH 027/774] [Attention] Extend XQA decode support on SM90 (#50439) Signed-off-by: Andrii Skliar Co-authored-by: Andrii Skliar Co-authored-by: OpenAI Codex --- tests/v1/attention/test_attention_backends.py | 160 ++++++++++++++-- .../test_flashinfer_dcp_spec_reorder.py | 5 + vllm/v1/attention/backends/flashinfer.py | 179 +++++++++--------- 3 files changed, 239 insertions(+), 105 deletions(-) diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 6609705c2202..6efdf2a31b0b 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -889,7 +889,7 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding(): device = torch.device("cpu") builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder) - builder.use_dedicated_xqa = True + builder.use_xqa = True qo_indptr = torch.tensor([0, 3, 9, 15, 15], dtype=torch.int32, device=device) q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens( @@ -905,6 +905,95 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding(): assert q_cu_seq_lens.tolist() == [0, 3, 9, 15, 15] +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypatch): + """Non-ragged XQA keeps one query/output row per padded request.""" + import unittest.mock + + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + impl = object.__new__(flashinfer_backend.FlashInferImpl) + impl.scale = 1.0 + impl.bmm1_scale = None + impl.bmm2_scale = None + impl.kv_cache_dtype = "auto" + impl.is_kvcache_nvfp4 = False + impl.head_size = 16 + impl.dcp_world_size = 1 + impl.o_sf_scale = None + impl.window_left = -1 + impl.sinks = None + impl.cache_config = unittest.mock.Mock() + impl.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC + + layer = unittest.mock.Mock( + _q_scale=torch.tensor(1.0), + _q_scale_float=1.0, + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + decode = flashinfer_backend.FlashInferTrtllmAPIDecode( + kernel=flashinfer_backend.FlashInferDecodeKernel.XQA, + block_tables=torch.zeros((4, 1), dtype=torch.int32), + seq_lens=torch.tensor([8, 8, 0, 0], dtype=torch.int32), + max_seq_len=8, + q_len_per_req=1, + ) + attn_metadata = flashinfer_backend.FlashInferMetadata( + num_actual_tokens=2, + slot_mapping=torch.empty(0, dtype=torch.int64), + q_data_type_prefill=torch.bfloat16, + q_data_type_decode=torch.bfloat16, + num_decodes=4, + num_decode_tokens=2, + num_prefills=0, + num_prefill_tokens=0, + causal=True, + prefill=None, + decode=decode, + use_cascade=False, + cascade_wrapper=None, + ) + + seen_shapes = {} + + def mock_xqa(**kwargs): + seen_shapes["query"] = kwargs["query"].shape + seen_shapes["out"] = kwargs["out"].shape + seen_shapes["block_tables"] = kwargs["block_tables"].shape + + monkeypatch.setattr( + flashinfer_backend, + "_get_trtllm_workspace_buffer", + lambda: torch.empty(1, dtype=torch.uint8), + ) + monkeypatch.setattr( + flashinfer_backend, "flashinfer_xqa_batch_decode_with_kv_cache", mock_xqa + ) + + query = torch.zeros((4, 1, 16), dtype=torch.bfloat16) + output = torch.empty_like(query) + result = impl.forward( + layer, + query, + torch.empty_like(query), + torch.empty_like(query), + torch.zeros((1, 1, 1, 32), dtype=torch.bfloat16), + attn_metadata, + output, + ) + + assert result is output + assert seen_shapes == { + "query": torch.Size([4, 1, 16]), + "out": torch.Size([4, 1, 16]), + "block_tables": torch.Size([4, 1]), + } + + @pytest.mark.skipif( AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", @@ -913,7 +1002,7 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product(): from vllm.v1.attention.backends import flashinfer as flashinfer_backend builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder) - builder.use_dedicated_xqa = True + builder.use_xqa = True qo_indptr = torch.tensor([0, 3, 3, 3], dtype=torch.int32) q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens( @@ -929,6 +1018,42 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product(): assert q_cu_seq_lens.tolist() == [0, 3, 3, 3] +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +@pytest.mark.parametrize( + ("qo_indptr_values", "num_decode_tokens", "expected_q_len"), + [ + ([0, 3, 3], 3, 3), + ([0, 2, 4, 4, 4], 4, 2), + ], +) +def test_flashinfer_trtllm_gen_padded_decode_uses_varlen_offsets( + qo_indptr_values: list[int], + num_decode_tokens: int, + expected_q_len: int, +): + """Padded speculative decode keeps its actual packed query width.""" + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder) + builder.use_xqa = False + qo_indptr = torch.tensor(qo_indptr_values, dtype=torch.int32) + + q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens( + qo_indptr, + qo_indptr, + num_decodes=len(qo_indptr_values) - 1, + num_decode_tokens=num_decode_tokens, + ) + + assert q_len == expected_q_len + assert q_lens is None + assert q_cu_seq_lens is not None + assert q_cu_seq_lens.tolist() == qo_indptr_values + + @pytest.mark.skipif( AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", @@ -961,10 +1086,12 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype): reason="FlashInfer is not available.", ) def test_flashinfer_native_prefill_with_sinks(default_vllm_config): - if not ( - current_platform.is_cuda() and current_platform.is_device_capability_family(120) - ): - pytest.skip("Native FlashInfer prefill with sinks requires SM12x.") + supported = current_platform.is_cuda() and ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ) + if not supported: + pytest.skip("Native FlashInfer prefill with sinks requires SM90 or SM12x.") from vllm.v1.attention.backends.flashinfer import FlashInferBackend @@ -995,7 +1122,7 @@ def causal_mask_mod( reason="FlashInfer is not available.", ) def test_flashinfer_xqa_decode_correctness(default_vllm_config): - """FlashInfer should route supported decode through XQA and match SDPA.""" + """FlashInfer should route SM90/SM12x decode through XQA and match SDPA.""" supported = current_platform.is_cuda() and ( current_platform.is_device_capability(90) or current_platform.is_device_capability_family(120) @@ -1066,16 +1193,23 @@ def causal_mask_mod( ) attn_metadata = builder.build(0, common_attn_metadata) - expected_cg_support = ( - AttentionCGSupport.UNIFORM_BATCH - if current_platform.is_device_capability_family(120) - else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - ) assert ( flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support( vllm_config, kv_cache_spec ) - == expected_cg_support + == AttentionCGSupport.UNIFORM_BATCH + ) + wide_head_kv_cache_spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=kv_cache_spec.num_kv_heads, + head_size=512, + dtype=vllm_config.model_config.dtype, + ) + assert ( + flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support( + vllm_config, wide_head_kv_cache_spec + ) + == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE ) assert isinstance( attn_metadata.decode, diff --git a/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py index 868d59247af5..94d9ecd9a0cc 100644 --- a/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py +++ b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py @@ -13,6 +13,7 @@ from tests.v1.attention.utils import create_vllm_config from vllm.config import SpeculativeConfig, set_current_vllm_config +from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends import flashinfer as flashinfer_backend from vllm.v1.attention.backends.flashinfer import ( FlashInferDecodeKernel, @@ -73,3 +74,7 @@ def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch): builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN ) assert builder.reorder_batch_threshold == 1 + assert ( + FlashInferMetadataBuilder.get_cudagraph_support(vllm_config, kv_cache_spec) + == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 3547a45c9814..35e9db46a570 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -519,7 +519,7 @@ def supports_compute_capability(cls, capability: DeviceCapability) -> bool: @classmethod def supports_sink(cls) -> bool: - """FlashInfer supports sinks on SM12x XQA and SM100 trtllm-gen.""" + """Whether FlashInfer can serve attention sinks on this platform.""" from vllm.utils.flashinfer import ( force_use_trtllm_attention, ) @@ -527,7 +527,9 @@ def supports_sink(cls) -> bool: if force_use_trtllm_attention() is False: return False - if current_platform.is_device_capability_family(120): + if current_platform.is_device_capability( + 90 + ) or current_platform.is_device_capability_family(120): return supports_trtllm_attention(is_prefill=False) if not current_platform.is_device_capability_family(100): @@ -570,6 +572,10 @@ class FlashInferDecodeKernel(Enum): TRTLLM_GEN = "trtllm-gen" +def _is_xqa_head_dim_supported(head_dim: int) -> bool: + return 16 <= head_dim <= 256 and head_dim % 16 == 0 + + @dataclass class TRTLLMPrefill: """Metadata for the TRTLLM prefill pathway.""" @@ -827,7 +833,7 @@ def __init__( # the wider groups fall back to native FlashInfer decode. if ( self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA - and not (16 <= self.head_dim <= 256 and self.head_dim % 16 == 0) + and not _is_xqa_head_dim_supported(self.head_dim) ): logger.warning_once( "FlashInfer XQA decode does not support head_dim=%d; " @@ -847,9 +853,8 @@ def __init__( ) self.use_trtllm_decode_attention = False self.flashinfer_trtllm_api_decode_kernel = None - self.use_dedicated_xqa = ( - current_platform.is_device_capability_family(120) - and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA + self.use_xqa = ( + self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA ) # Adaptive verification trims drafts on device, so decode query lengths # must come from the device qo_indptr; only trtllm-gen supports that @@ -861,14 +866,11 @@ def __init__( == FlashInferDecodeKernel.TRTLLM_GEN and not self.use_dcp ) - supports_spec_as_decode = ( - self.flashinfer_trtllm_api_decode_kernel - == FlashInferDecodeKernel.TRTLLM_GEN - or self.use_dedicated_xqa - ) self._init_reorder_batch_threshold( 1, - supports_spec_as_decode=supports_spec_as_decode, + supports_spec_as_decode=( + self.flashinfer_trtllm_api_decode_kernel is not None + ), # trtllm-gen decode receives no cp_rank/global-seq-len information, # so its end-aligned causal mask is wrong for q_len > 1 over the # DCP-interleaved local KV shard (spec token i misses up to @@ -886,16 +888,6 @@ def __init__( per_layer_parameters = get_per_layer_parameters( vllm_config, layer_names, FlashInferImpl ) - if current_platform.is_device_capability(90) and any( - params.window_left != -1 for params in per_layer_parameters.values() - ): - # FlashInfer SM90 sliding-window prefill is not reliable with FP8-Q: - # https://github.com/flashinfer-ai/flashinfer/issues/3578 - raise NotImplementedError( - "FlashInfer backend on SM90 currently crashes with " - "sliding-window attention layers. Use the default attention " - "backend." - ) self.global_hyperparameters = infer_global_hyperparameters(per_layer_parameters) self.sm_scale = self.global_hyperparameters.sm_scale self.window_left = self.global_hyperparameters.window_left @@ -950,7 +942,9 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype: # if cache_config requests a quantized dtype globally. cache_dtype = self.cache_dtype - # XQA decode requires BF16/FP16-Q even with FP8 KV cache. + # On SM90/SM12x, XQA decode requires BF16/FP16-Q even with FP8 KV cache. + # FI native prefill on SM90 still uses FP8-Q in that case; SM12x prefill + # is fa2-only and keeps the model dtype (handled below). if ( ( current_platform.is_device_capability(90) @@ -964,7 +958,9 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype: # Otherwise, match Q dtype to the KV cache dtype. if cache_dtype.startswith("fp8"): - # FP8-Q requires an fp8 tensor-core attention path. + # FP8-Q requires an fp8 tensor-core attention path. SM90 uses + # native FA3 for prefill; XQA decode on SM90/SM12x is handled + # above with model-dtype Q, while SM100 uses trtllm-gen. # Architectures with only fa2 (e.g. SM89, SM120) cannot # consume FP8 queries, so keep the model dtype for Q there. if current_platform.is_device_capability( @@ -983,24 +979,19 @@ def get_cudagraph_support( vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec, ) -> AttentionCGSupport: - """Get the cudagraph support level for FlashInfer attention. - - SM90 XQA supports only single-token decode. SM12x uses the dedicated - XQA API, which supports speculative and non-causal decode. - """ - if current_platform.is_device_capability(90): - return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - - is_sm12x = current_platform.is_device_capability_family(120) - # XQA does not return LSE and therefore does not support DCP. - if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1: + """Get the cudagraph support level for FlashInfer attention.""" + # XQA lacks LSE for DCP; DCP also cannot graph variable-length trtllm-gen. + if vllm_config.parallel_config.decode_context_parallel_size > 1: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE kv_specs = iter_layer_specs(kv_cache_spec) num_qo_heads = vllm_config.model_config.get_num_attention_heads( vllm_config.parallel_config ) - has_trtllm_support: bool = len(kv_specs) > 0 + is_xqa_arch = current_platform.is_device_capability( + 90 + ) or current_platform.is_device_capability_family(120) + has_uniform_batch_support: bool = len(kv_specs) > 0 for spec in kv_specs: if not isinstance(spec, AttentionSpec): # FlashInfer only applies to attention, so we don't consider other types @@ -1010,13 +1001,12 @@ def get_cudagraph_support( num_qo_heads=num_qo_heads, num_kv_heads=spec.num_kv_heads, is_prefill=False, - ): - has_trtllm_support = False + ) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)): + has_uniform_batch_support = False break - if has_trtllm_support and ( - is_sm12x or not vllm_config.attention_config.use_non_causal - ): + use_non_causal = vllm_config.attention_config.use_non_causal + if has_uniform_batch_support and (not use_non_causal or is_xqa_arch): return AttentionCGSupport.UNIFORM_BATCH else: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE @@ -1066,10 +1056,27 @@ def _compute_decode_query_lens( num_decode_tokens: int, ) -> tuple[int, torch.Tensor | None, list[int] | None]: """Return the query width, ragged offsets, and effective query lengths.""" - assert self.use_dedicated_xqa if num_decodes == 0 or num_decode_tokens == 0: return 1, None, None + if not self.use_xqa: + decode_q_lens = ( + qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes] + ) + q_len_per_req = num_decode_tokens // num_decodes + if num_decode_tokens % num_decodes == 0 and bool( + (decode_q_lens == q_len_per_req).all().item() + ): + return q_len_per_req, None, None + + # CUDA-graph padding can leave zero-length requests after a + # uniform speculative decode batch. Keep the packed query and + # describe it through the trtllm-gen varlen API. + max_q_len = int(decode_q_lens.max().item()) + if max_q_len > 1: + return max_q_len, qo_indptr[: num_decodes + 1], None + return 1, None, None + decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes] nonzero = decode_q_lens[decode_q_lens > 0] q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1 @@ -1128,7 +1135,10 @@ def _get_prefill_wrapper( "NVFP4 KV cache." ) if self._noncausal_prefill_wrapper is None: - if self.has_sinks and current_platform.is_device_capability_family(120): + if self.has_sinks and ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ): self._noncausal_prefill_wrapper = ( BatchAttentionWithAttentionSinkWrapper( self._get_workspace_buffer(), @@ -1159,7 +1169,10 @@ def _get_prefill_wrapper( dcp_a2a=self.dcp_a2a, ) else: - if self.has_sinks and current_platform.is_device_capability_family(120): + if self.has_sinks and ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ): assert not self.is_kvcache_nvfp4 self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper( self._get_workspace_buffer(), @@ -1296,13 +1309,13 @@ def build( num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens causal = common_attn_metadata.causal - route_decode = causal or self.use_dedicated_xqa + route_decode = causal or self.use_xqa if route_decode: num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=not self.use_dedicated_xqa, + require_uniform=not self.use_xqa, ) ) else: @@ -1342,7 +1355,7 @@ def build( has_spec=uses_spec_reorder, ) decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and ( - causal or self.use_dedicated_xqa + causal or self.use_xqa ) if not causal and self.use_dcp: @@ -1360,13 +1373,12 @@ def build( ) if not all_uses_trtllm: - if self.has_sinks and ( - not self.use_dedicated_xqa or self.use_dcp or use_cascade - ): + sinks_would_be_dropped = self.use_dcp or use_cascade + if self.has_sinks and sinks_would_be_dropped: raise NotImplementedError( - "FlashInfer backend currently does not support attention " - "sinks, please use trtllm on blackwell or flash attention " - "on earlier GPUs." + "FlashInfer BatchDCPPrefillWrapper and " + "MultiLevelCascadeAttentionWrapper do not support attention sinks. " + "Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs." ) if not self.global_hyperparameters.has_same_window_lefts: @@ -1663,29 +1675,19 @@ def build( if num_decodes > 0: if decode_with_flashinfer_trtllm_api: assert self.flashinfer_trtllm_api_decode_kernel is not None - if not self.use_dedicated_xqa and not self.use_trtllm_gen_varlen_decode: - assert num_decode_tokens % num_decodes == 0, ( - "XQA/trtllm-gen decode requires uniform query lengths " - f"per request. Got {num_decode_tokens=} and {num_decodes=}." - ) seq_lens_decode = seq_lens[:num_decodes] if self.use_dcp: assert common_attn_metadata.dcp_local_seq_lens is not None seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[ :num_decodes ] - q_len_per_req = 1 - q_cu_seq_lens = None - decode_mask = None - if self.use_dedicated_xqa: - q_len_per_req, q_cu_seq_lens, ragged_q_lens = ( - self._compute_decode_query_lens( - qo_indptr, - qo_indptr_cpu, - num_decodes, - num_decode_tokens, - ) + q_len_per_req, q_cu_seq_lens, ragged_q_lens = ( + self._compute_decode_query_lens( + qo_indptr, qo_indptr_cpu, num_decodes, num_decode_tokens ) + ) + decode_mask = None + if self.use_xqa: decode_mask = self._get_decode_mask( q_len_per_req, ragged_q_lens, @@ -2066,7 +2068,10 @@ def forward( if fp8_view_dtype is not None: kv_cache = kv_cache.view(fp8_view_dtype) - # Inputs and outputs may be padded for CUDA graphs + # Inputs and outputs may be padded for CUDA graphs. Keep the original + # query for single-token XQA, whose non-ragged API requires one row for + # every (including padded) decode request. + query_padded = query query = query[:num_actual_tokens] output_padded = output output = output[:num_actual_tokens] @@ -2131,10 +2136,7 @@ def forward( kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) use_dcp = self.dcp_world_size > 1 - decode_with_dedicated_xqa = ( - decode_with_xqa and current_platform.is_device_capability_family(120) - ) - if decode_with_dedicated_xqa: + if decode_with_xqa: assert not use_dcp assert not self.is_kvcache_nvfp4 assert self.o_sf_scale is None @@ -2351,8 +2353,13 @@ def forward( ].copy_(out[:num_prefill_tokens]) if num_decode_tokens > 0: - decode_query = query[:num_decode_tokens] - assert decode_query.shape[0] == num_decode_tokens + decode_query_tokens = num_decode_tokens + if decode_with_xqa: + assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode) + if attn_metadata.decode.q_len_per_req == 1: + decode_query_tokens = attn_metadata.num_decodes + decode_query = query_padded[:decode_query_tokens] + assert decode_query.shape[0] == decode_query_tokens # Convert query to the expected dtype for decode if needed. decode_query = self.maybe_quant_query( @@ -2467,7 +2474,7 @@ def forward( ) decode_query = canonicalize_singleton_dim_strides(decode_query) - if decode_with_dedicated_xqa: + if decode_with_xqa: bmm1_scale = self.get_xqa_bmm1_scale( layer, attn_metadata.q_data_type_decode ) @@ -2483,7 +2490,7 @@ def forward( bmm1_scale=bmm1_scale, bmm2_scale=self.bmm2_scale, window_left=self.window_left, - out=output[:num_decode_tokens], + out=output_padded[:decode_query_tokens], sinks=self.sinks, kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), q_len_per_req=q_len_per_req, @@ -2492,6 +2499,7 @@ def forward( ) return output_padded + assert decode_with_trtllm_gen if output.dtype == FP4_DTYPE: assert self.o_sf_scale is not None out = FP4Tensor( @@ -2523,19 +2531,6 @@ def forward( else: q_len_per_req = num_decode_tokens // attn_metadata.num_decodes - if decode_with_xqa and q_len_per_req is not None and q_len_per_req > 1: - raise NotImplementedError( - "FlashInfer XQA speculative decode is not wired in vLLM yet." - ) - - # XQA decode can use model-dtype Q with FP8 KV, so only include - # q_scale when the decode query is actually FP8. - bmm1_scale = ( - self.get_xqa_bmm1_scale(layer, attn_metadata.q_data_type_decode) - if decode_with_xqa - else self.bmm1_scale - ) - lse = None if use_dcp: out = torch.empty( @@ -2558,7 +2553,7 @@ def forward( block_tables=block_tables_decode, seq_lens=seq_lens_decode, max_seq_len=attn_metadata.decode.max_seq_len, - bmm1_scale=bmm1_scale, + bmm1_scale=self.bmm1_scale, bmm2_scale=self.bmm2_scale, window_left=self.window_left, sinks=self.sinks, From 5c642796e871e41ae75be0b402a8eaa44ee771b1 Mon Sep 17 00:00:00 2001 From: Alexander Date: Fri, 11 Sep 2026 07:58:49 +0100 Subject: [PATCH 028/774] [Docs] Add Nebius Serverless AI deployment guide (#56414) Signed-off-by: Alexander Salikov --- docs/deployment/frameworks/nebius.md | 148 +++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 docs/deployment/frameworks/nebius.md diff --git a/docs/deployment/frameworks/nebius.md b/docs/deployment/frameworks/nebius.md new file mode 100644 index 000000000000..94d66c1de8b7 --- /dev/null +++ b/docs/deployment/frameworks/nebius.md @@ -0,0 +1,148 @@ +# Nebius Serverless AI + +[Nebius Serverless AI](https://docs.nebius.com/serverless/overview) runs GPU containers as persistent HTTP endpoints or finite jobs. This guide deploys the vLLM OpenAI-compatible server as a single-GPU Endpoint using the Nebius CLI. + +The Endpoint keeps running until you stop or delete it. This recipe uses a regular GPU instance; it does not configure autoscaling or automatic scale-to-zero. + +## Prerequisites + +- A Nebius project with permission to create and delete Serverless Endpoints, and sufficient GPU, VM, disk, and networking quotas. +- The [Nebius CLI](https://docs.nebius.com/cli/install), installed and authenticated for that project. +- A subnet in `eu-north1` with outbound access to Docker Hub and Hugging Face. The example uses the `gpu-l40s-a` platform and `1gpu-8vcpu-32gb` preset; check availability in your project before starting. +- `curl`, `jq`, and `openssl` on your local machine. + +The example model, `Qwen/Qwen3-0.6B`, does not require a Hugging Face token. Other models may require accepting their license and supplying a token through Nebius [secret environment variables](https://docs.nebius.com/serverless/endpoints/manage). + +## Create the Endpoint + +The commands below were tested with Nebius CLI `0.12.265`. Set your project and subnet IDs explicitly. Use the same CLI profile throughout the guide. + +```bash +export PROJECT_ID="" +export SUBNET_ID="" +export ENDPOINT_NAME="vllm-qwen-$(openssl rand -hex 6)" +export AUTH_TOKEN="$(openssl rand -hex 32)" +export MODEL_ID="Qwen/Qwen3-0.6B" +export MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca" +# vllm/vllm-openai:v0.19.1, Linux/amd64 image digest. +export VLLM_IMAGE="vllm/vllm-openai@sha256:89c1d0629d377daa3f7f369cbea6167a7b48ea89aaacd12555e2b0b2f7f740d3" + +nebius ai endpoint create \ + --parent-id "$PROJECT_ID" \ + --subnet-id "$SUBNET_ID" \ + --name "$ENDPOINT_NAME" \ + --image "$VLLM_IMAGE" \ + --container-command vllm \ + --args "serve $MODEL_ID --revision $MODEL_REVISION --tokenizer-revision $MODEL_REVISION --host 0.0.0.0 --port 8000 --tensor-parallel-size 1 --max-model-len 4096 --gpu-memory-utilization 0.8" \ + --platform gpu-l40s-a \ + --preset 1gpu-8vcpu-32gb \ + --container-port 8000/http \ + --auth token --token "$AUTH_TOKEN" \ + --disk-size 250Gi --shm-size 16Gi \ + --public=false --preemptible=false \ + --retries 1 +``` + +The image and model revisions are pinned for reproducibility. When changing the image, check its CUDA/driver requirements against the selected Nebius platform. The model downloads into the container disk on startup; this example does not configure persistent model storage. + +The Endpoint token is separate from your Nebius CLI credentials. Keep it private. Nebius token authentication requires exactly one HTTP port. This example relies on authentication at the managed HTTPS URL; it does not set a second vLLM API key. Access from inside the subnet must be restricted to trusted clients. + +A public VM IP is not required to use the managed HTTPS URL. The subnet still needs outbound connectivity for image and model downloads. + +Save the ID of the Endpoint you created: + +```bash +export ENDPOINT_ID="$(nebius ai endpoint get-by-name \ + --parent-id "$PROJECT_ID" --name "$ENDPOINT_NAME" \ + --format jsonpath='{.metadata.id}')" +``` + +If creation times out or your terminal disconnects, use the same project and name to find the Endpoint before attempting another creation. A local timeout does not mean that provisioning was cancelled. + +## Check readiness + +Inspect the state and recent logs: + +```bash +nebius ai endpoint get "$ENDPOINT_ID" +nebius ai endpoint logs "$ENDPOINT_ID" --tail 100 --timestamps +``` + +Wait for `RUNNING`, then select the managed HTTPS URL. The URL can appear while the Endpoint is still `STARTING`, when ingress may return HTTP 404. The `jq` expression requires exactly one HTTPS URL and preserves its scheme: + +```bash +ENDPOINT_URL="$(nebius ai endpoint get "$ENDPOINT_ID" --format json \ + | jq -er '[.status.public_endpoints[]? | select(startswith("https://"))] + | if length == 1 then .[0] else error("Expected one HTTPS URL") end')" +export ENDPOINT_URL="${ENDPOINT_URL%/}" + +curl --fail-with-body --silent --show-error --max-time 10 \ + "$ENDPOINT_URL/health" -H "Authorization: Bearer $AUTH_TOKEN" +curl --fail-with-body --silent --show-error --max-time 10 \ + "$ENDPOINT_URL/v1/models" -H "Authorization: Bearer $AUTH_TOKEN" | jq +``` + +`RUNNING` alone does not establish that the model has finished loading. Wait for `/health` to succeed and `/v1/models` to list `Qwen/Qwen3-0.6B` before sending a chat request. While waiting, inspect logs instead of creating another Endpoint. If startup does not succeed within your chosen time or spending limit, delete the Endpoint and investigate the error. + +## Send a chat request + +```bash +curl --fail-with-body --silent --show-error --max-time 120 \ + "$ENDPOINT_URL/v1/chat/completions" \ + -H "Authorization: Bearer $AUTH_TOKEN" \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "Say hello in one short sentence."}], + "max_tokens": 128, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": false} + }' | jq +``` + +Expect a chat completion containing an assistant message. The Qwen-specific template option disables thinking for this short example; it is not a universal option for other models. + +For streaming, add `"stream": true` to the JSON and use `curl -N` without piping into `jq`. A successful stream contains completion chunks followed by `data: [DONE]`. For an OpenAI-compatible client, use `${ENDPOINT_URL}/v1` as the base URL and the Endpoint token as the API key. + +Verify that the same inference URL rejects requests without the token: + +```bash +curl --silent --show-error --max-time 10 --output /dev/null \ + --write-out '%{http_code}\n' "$ENDPOINT_URL/v1/models" +``` + +Expect HTTP `401` or `403`. Do not share an Endpoint that accepts unauthenticated inference unexpectedly. + +## Stop or delete the Endpoint + +To stop serving while preserving the Endpoint configuration: + +```bash +nebius ai endpoint stop "$ENDPOINT_ID" +nebius ai endpoint get "$ENDPOINT_ID" +``` + +Let the synchronous stop command complete and confirm `STOPPED` before starting again. To start it again, run `nebius ai endpoint start "$ENDPOINT_ID"`, retrieve its current URL, and repeat the readiness checks. Starting again can require fresh capacity and new image/model downloads; do not rely on the container disk as persistent storage. + +When finished, delete the Endpoint and confirm that getting the same ID returns `NotFound`: + +```bash +nebius ai endpoint delete "$ENDPOINT_ID" +nebius ai endpoint get "$ENDPOINT_ID" +unset AUTH_TOKEN +``` + +Nebius deletes the managed VM and container disk with the Endpoint. Separately mounted storage has its own lifecycle and billing. Closing your terminal or interrupting a request does not stop the Endpoint. If deletion fails, keep the ID and retry cleanup; a generic connection or authentication error is not confirmation of deletion. + +## Troubleshooting + +| Symptom | Check | +| --- | --- | +| `PROVISIONING` or `NotEnoughResources` | Inspect quota, region, platform and preset. Capacity availability can differ from quota. | +| `RUNNING`, but HTTP 502/503 | Check model-download and engine-start logs, subnet egress, host binding and port 8000. | +| HTTP 401/403 | Use the Endpoint token, not the Nebius control-plane token. | +| CUDA or out-of-memory error | Check image/driver compatibility, the model's memory requirements and maximum context length. | +| Model not found | Verify the requested model ID and pinned revision. | +| Interrupted stream | Inspect ingress and server logs. A partial response should not be automatically replayed as a new generation. | + +See the [Nebius vLLM cookbook example](https://github.com/nebius/serverless-ai-cookbook/tree/main/inference/vllm-endpoint) and [Endpoint management documentation](https://docs.nebius.com/serverless/endpoints/manage) for further details. From a9271c750f6eb07ad7e4c52b9dffa65ff2ee336a Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 11 Sep 2026 02:30:39 -0500 Subject: [PATCH 029/774] [ROCm][CI] Accept `base-v2-preview` images during content hash lookup for ROCm base images (#56356) Signed-off-by: Micah Williamson --- .buildkite/scripts/rocm/refresh-base-image.sh | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh index 23d10492a99c..0e10b98c9eb2 100755 --- a/.buildkite/scripts/rocm/refresh-base-image.sh +++ b/.buildkite/scripts/rocm/refresh-base-image.sh @@ -384,6 +384,7 @@ build_base_image() { local stable_tag="${BASE_REPO}:base" local trusted_content_tag="" local scoped_content_tag="" + local preview_content_tag="" local writable_content_tag="" local build_ref="" local immutable_ref="" @@ -431,6 +432,7 @@ build_base_image() { dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}" trusted_content_tag=$(trusted_base_content_ref "${base_hash}" "${metadata_version}") scoped_content_tag=$(scoped_base_content_ref "${base_hash}" "${metadata_version}") + preview_content_tag="${scoped_content_tag}" # Preview writes share an exact-content namespace so stacked PRs can reuse # identical images. They may import the trusted ref, but never overwrite it. if is_trusted_main_build; then @@ -471,7 +473,8 @@ build_base_image() { if [[ "${trusted_content_tag}" == "${scoped_content_tag}" ]]; then immutable_ref=$(find_matching_base_content_ref \ "${base_hash}" "${metadata_version}" \ - "${trusted_content_tag}" "${stable_tag}") || reuse_status=$? + "${trusted_content_tag}" "${preview_content_tag}" \ + "${stable_tag}") || reuse_status=$? else immutable_ref=$(find_matching_base_content_ref \ "${base_hash}" "${metadata_version}" \ @@ -490,8 +493,9 @@ build_base_image() { if [[ ${reuse_status} -eq 0 && -n "${immutable_ref}" ]]; then echo "Reusing ROCm base image with matching content: ${immutable_ref}" - if [[ "${immutable_ref%@*}" == "${stable_tag}" \ - && "${writable_content_tag}" != "${stable_tag}" ]]; then + if [[ "${immutable_ref%@*}" != "${writable_content_tag}" \ + && ( "${immutable_ref%@*}" == "${stable_tag}" \ + || "${immutable_ref%@*}" == "${preview_content_tag}" ) ]]; then echo "Backfilling ROCm base content tag: ${writable_content_tag}" docker buildx imagetools create --prefer-index=false \ -t "${writable_content_tag}" "${immutable_ref}" From fadfe1c7d4df4c29b3dbc159ed790238a885b084 Mon Sep 17 00:00:00 2001 From: lucamotz Date: Fri, 11 Sep 2026 09:50:47 +0200 Subject: [PATCH 030/774] [Bugfix][Core] Retire Mamba states across null gaps (#55450) Signed-off-by: Luca Motz Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung --- .../core/test_single_type_kv_cache_manager.py | 76 +++++++++++++++++++ .../v1/e2e/general/test_mamba_prefix_cache.py | 2 +- vllm/v1/core/single_type_kv_cache_manager.py | 23 ++++++ 3 files changed, 100 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 39113723591b..14e46d7ba07c 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -103,6 +103,82 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership(): manager._relocate_speculative_block([pinned_block], 0) +def test_mamba_retirement_crosses_null_gaps(): + spec = MambaSpec( + block_size=4, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + pool = BlockPool(num_gpu_blocks=8, enable_caching=False, hash_block_size=4) + manager = MambaManager( + spec, + block_pool=pool, + enable_caching=False, + kv_cache_group_id=0, + scheduler_block_size=4, + ) + old, committed, in_flight = pool.get_new_blocks(3) + manager.req_to_blocks["r"] = [old, pool.null_block, committed, in_flight] + + # The state at token 12 and the following in-flight state must survive. + for _ in range(2): + manager.remove_skipped_blocks("r", processed_computed_tokens=12) + assert manager.req_to_blocks["r"] == [ + pool.null_block, + pool.null_block, + committed, + in_flight, + ] + assert old.ref_cnt == 0 + assert committed.ref_cnt == in_flight.ref_cnt == 1 + assert pool.get_num_free_blocks() == 5 + + manager.free("r") + manager.req_to_blocks["r"] = pool.get_new_blocks(2) + manager.remove_skipped_blocks("r", processed_computed_tokens=5) + assert manager.req_to_blocks["r"][0].is_null + assert manager.req_to_blocks["r"][1].ref_cnt == 1 + + +@pytest.mark.parametrize("block_size", [3584, 4608]) +@pytest.mark.parametrize("in_flight_chunks", [0, 1, 2]) +def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks): + spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=5, + ) + pool = BlockPool(num_gpu_blocks=1000, enable_caching=True, hash_block_size=256) + manager = MambaManager( + spec, + block_pool=pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=block_size, + ) + initial_free = pool.get_num_free_blocks() + chunk = 8192 // block_size * block_size + peak_held = 0 + # Replay full prefill chunks with an async committed-token lag. + for target in range(chunk, 480031 + 1, chunk): + computed = target - chunk + committed = max(0, computed - in_flight_chunks * chunk) + manager.remove_skipped_blocks("r", committed) + manager.get_num_blocks_to_allocate( + "r", target + 5, [], computed, computed, target + ) + manager.allocate_new_blocks("r", target + 5, target) + peak_held = max(peak_held, initial_free - pool.get_num_free_blocks()) + + # Five speculative blocks, current/previous states, and bounded in-flight states. + assert peak_held == 7 + in_flight_chunks + manager.free("r") + assert pool.get_num_free_blocks() == initial_free + + def get_sliding_window_manager( sliding_window_spec, block_pool, diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 27331516cbd3..aa1f730afbad 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -837,7 +837,7 @@ def get_mamba_prefix_cache_step_configs( StepAction( 560 * 10, 4, - [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1] + [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], (9, 10), diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 93475e505526..2015967dd2df 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1460,6 +1460,7 @@ def __init__( # Mapping from request ID to the index of the block # allocated in the previous step self.last_state_block_idx: dict[str, int] = {} + self._num_retired_blocks: dict[str, int] = {} # The set of the requests that have been allocated blocks self._allocated_block_reqs: set[str] = set() # checkpoint position and reserved block index for the current @@ -1610,6 +1611,27 @@ def reachable_block_mask( return mask + def _remove_blocks_in_range( + self, request_id: str, first_block: int, last_block: int + ) -> None: + if self.mamba_cache_mode != "align": + return super()._remove_blocks_in_range(request_id, first_block, last_block) + blocks = self.req_to_blocks.get(request_id, []) + first_block = max(first_block, self._num_retired_blocks.get(request_id, 0)) + last_block = min(last_block, len(blocks)) + if first_block >= last_block: + return + freed: list[KVCacheBlock] = [] + # Mamba prefill leaves null gaps between states awaiting retirement. + for i in range(last_block - 1, first_block - 1, -1): + if blocks[i].is_null: + continue + freed.append(blocks[i]) + blocks[i] = self._null_block + if freed: + self.block_pool.free_blocks(freed) + self._num_retired_blocks[request_id] = last_block + def remove_skipped_blocks( self, request_id: str, @@ -1930,6 +1952,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) + self._num_retired_blocks.pop(request_id, None) self._checkpoints.pop(request_id, None) self._producer_partial_tail_reqs.pop(request_id, None) # An offer is only guaranteed to hold committed bytes until the end From 6376c601ede7c8ff25d2bd89dbaa46297a45db10 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Fri, 11 Sep 2026 09:59:13 +0200 Subject: [PATCH 031/774] [XPU][Tests] Enable test_per_token_group_quant_int8 on XPU (#55681) Signed-off-by: pmanczak --- tests/kernels/quantization/test_per_token_group_quant.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index a2788d058304..d88a39d08469 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -413,9 +413,12 @@ def test_per_token_group_quant_fp8_packed_large_mn(): @pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)]) @pytest.mark.parametrize("group_size", [64, 128]) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Only test on CUDA/ROCm/XPU.", +) def test_per_token_group_quant_int8(shape, group_size: int): - device = "cuda" + device = current_platform.device_type torch.manual_seed(42) num_tokens, hidden_dim = shape From e3f755b732316df4db8437dcb973b9fa28bfd656 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:21:14 +0100 Subject: [PATCH 032/774] [CPU] Speedup LM Head on Arm CPUs (#55352) Signed-off-by: Fadi Arafeh Co-authored-by: Li, Jiang --- csrc/cpu/dnnl_helper.cpp | 11 +++++++---- csrc/cpu/dnnl_helper.h | 3 +++ csrc/cpu/dnnl_kernels.cpp | 28 ++++++++++++++++++++++++++-- tests/kernels/test_onednn.py | 1 + 4 files changed, 37 insertions(+), 6 deletions(-) diff --git a/csrc/cpu/dnnl_helper.cpp b/csrc/cpu/dnnl_helper.cpp index 14c136dcbbf0..1ac0401f540d 100644 --- a/csrc/cpu/dnnl_helper.cpp +++ b/csrc/cpu/dnnl_helper.cpp @@ -172,7 +172,8 @@ struct hash { size_t operator()( const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const { return hash()(val.b_n_size) ^ hash()(val.b_k_size) ^ - hash()(static_cast(val.b_type)); + hash()(static_cast(val.b_type)) ^ + hash()(static_cast(val.c_type)); } }; @@ -202,7 +203,7 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l, bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l, const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) { return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size && - l.b_type == r.b_type; + l.b_type == r.b_type && l.c_type == r.c_type; } bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l, @@ -482,8 +483,10 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) { dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache( const MSizeCacheKey& key) { if (m_size_cache_.get() == nullptr) { - ClassMatmulCacheKey class_key = { - .b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_}; + ClassMatmulCacheKey class_key = {.b_n_size = b_n_size_, + .b_k_size = b_k_size_, + .b_type = b_type_, + .c_type = c_type_}; m_size_cache_ = get_matul_class_primitive_cache(class_key, primitive_cache_size_); } diff --git a/csrc/cpu/dnnl_helper.h b/csrc/cpu/dnnl_helper.h index b841ed73058e..05fb116ecd1e 100644 --- a/csrc/cpu/dnnl_helper.h +++ b/csrc/cpu/dnnl_helper.h @@ -176,6 +176,7 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler { dnnl_dim_t b_n_size; dnnl_dim_t b_k_size; dnnl::memory::data_type b_type; + dnnl::memory::data_type c_type; friend bool operator==(const ClassMatmulCacheKey& l, const ClassMatmulCacheKey& r); @@ -203,6 +204,8 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler { public: MatMulPrimitiveHandler(const Args& args); + dnnl::memory::data_type get_output_type() const { return c_type_; } + void execute(ExecArgs& args); private: diff --git a/csrc/cpu/dnnl_kernels.cpp b/csrc/cpu/dnnl_kernels.cpp index 6dda09296160..e81d365628c0 100644 --- a/csrc/cpu/dnnl_kernels.cpp +++ b/csrc/cpu/dnnl_kernels.cpp @@ -513,6 +513,17 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b, args.ab_type = get_dnnl_type(); }); +#ifdef VLLM_USE_ACL + // TODO(fadara01): this is a hot-fix, remove once problem is addressed in the + // AArch64 stack. hybrid kernels in ACL which are the fastest for LM head + // layers are only enabled for BF16 x BF16 -> FP32 problem desc as they can't + // convert FP32->BF16 internally. In this case do: oneDNN matmul with + // bf16:bf16:fp32 then cast output down to bf16. + if (b.scalar_type() == at::kBFloat16 && b.size(1) >= 64000) { + args.c_type = dnnl::memory::data_type::f32; + } +#endif + return reinterpret_cast(new MatMulPrimitiveHandler(args)); } @@ -527,12 +538,22 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major MatMulPrimitiveHandler* ptr = reinterpret_cast(handler_tensor.item()); + torch::Tensor dnnl_output = c; + // AArch64 specific case where we do a bf16 x bf16 -> fp32 matmul + const bool convert_output = + ptr->get_output_type() == dnnl::memory::data_type::f32 && + c.scalar_type() != at::kFloat; + if (convert_output) { + dnnl_output = torch::empty(c.sizes(), c.options().dtype(at::kFloat)); + } + // ACL matmuls expect contiguous source tensors #ifdef VLLM_USE_ACL torch::Tensor a_contig = a.contiguous(); #endif MatMulPrimitiveHandler::ExecArgs exec_args; + exec_args.c_ptr = dnnl_output.data_ptr(); #ifdef VLLM_USE_ACL exec_args.a_m_size = a_contig.size(0); @@ -548,7 +569,7 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major #ifdef VLLM_USE_ACL // ACL matmuls in oneDNN do not support a bias. // We handle a matmul with bias by doing: c = bias; c += matmul(a, b) - c.copy_(bias.value()); + dnnl_output.copy_(bias.value()); #else exec_args.bias_ptr = bias->data_ptr(); #endif @@ -563,8 +584,11 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major exec_args.a_ptr = a.data_ptr(); #endif - exec_args.c_ptr = c.data_ptr(); ptr->execute(exec_args); }); + + if (convert_output) { + c.copy_(dnnl_output); + } } diff --git a/tests/kernels/test_onednn.py b/tests/kernels/test_onednn.py index 4e681ca6ac83..4e41791a5408 100644 --- a/tests/kernels/test_onednn.py +++ b/tests/kernels/test_onednn.py @@ -15,6 +15,7 @@ (256, 128), (4096, 4096), (16384, 4096), + (64128, 4096), (1023, 491), (1001, 15), ] From 69db1c26b4fe4474ab4c9df1c9701efac8bedde1 Mon Sep 17 00:00:00 2001 From: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Date: Fri, 11 Sep 2026 01:28:52 -0700 Subject: [PATCH 033/774] [CI/Build][Rust Frontend] Publish vllm-proto on crates.io (#56365) Co-authored-by: OpenAI Codex Co-authored-by: Bugen Zhao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Julien Darve Signed-off-by: Bugen Zhao --- .github/workflows/proto-crate.yml | 70 +++++++++++++++++++++++++++++ .pre-commit-config.yaml | 4 +- rust/Cargo.lock | 17 +++++-- rust/Cargo.toml | 2 + rust/proto/Cargo.toml | 22 +++++++++ rust/proto/README.md | 50 +++++++++++++++++++++ rust/{src/server => proto}/build.rs | 2 +- rust/proto/src/lib.rs | 6 +++ rust/src/server/Cargo.toml | 7 +-- rust/src/server/src/grpc/mod.rs | 4 +- 10 files changed, 168 insertions(+), 16 deletions(-) create mode 100644 .github/workflows/proto-crate.yml create mode 100644 rust/proto/Cargo.toml rename rust/{src/server => proto}/build.rs (88%) create mode 100644 rust/proto/src/lib.rs diff --git a/.github/workflows/proto-crate.yml b/.github/workflows/proto-crate.yml new file mode 100644 index 000000000000..43be5c8c9214 --- /dev/null +++ b/.github/workflows/proto-crate.yml @@ -0,0 +1,70 @@ +name: Rust proto crate + +on: + pull_request: + paths: + - "rust/proto/**" + - "rust/Cargo.toml" + - "rust/Cargo.lock" + - "rust-toolchain.toml" + - ".github/workflows/proto-crate.yml" + push: + tags: + - "proto-v*" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + semver: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2 + with: + manifest-path: rust/proto/Cargo.toml + package: vllm-proto + + check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - run: rustup show active-toolchain + - name: Check release version + if: github.event_name == 'push' + run: | + version=$(cargo metadata --manifest-path rust/Cargo.toml --no-deps --format-version 1 | + jq -er '.packages[] | select(.name == "vllm-proto") | .version') + if [[ "$GITHUB_REF_NAME" != "proto-v$version" ]]; then + echo "Tag $GITHUB_REF_NAME does not match vllm-proto version $version" >&2 + exit 1 + fi + - name: Verify package + run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run + + publish: + needs: [check, semver] + if: github.repository == 'vllm-project/vllm' && github.event_name == 'push' + runs-on: ubuntu-latest + environment: crates-io + permissions: + contents: read + id-token: write + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: ${{ github.sha }} + - run: rustup show active-toolchain + - name: Verify package before authentication + run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run + - uses: rust-lang/crates-io-auth-action@c6f97d42243bad5fab37ca0427f495c86d5b1a18 # v1.0.5 + id: auth + - name: Publish + env: + CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }} + run: cargo publish --manifest-path rust/proto/Cargo.toml --locked diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0954c53cb290..d1d0e9e3d15b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -285,7 +285,7 @@ repos: pass_filenames: false require_serial: true stages: [pre-commit] # Only run locally as Buildkite will cover this - files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$ + files: ^rust/(Cargo\.toml|proto/Cargo\.toml|src/.*/Cargo\.toml)$ - id: rust-cargo-sort name: Rust - Sort Cargo manifest sections entry: tools/pre_commit/rust-check.sh sort --workspace @@ -293,7 +293,7 @@ repos: pass_filenames: false require_serial: true stages: [pre-commit] # Only run locally as Buildkite will cover this - files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$ + files: ^rust/(Cargo\.toml|proto/Cargo\.toml|src/.*/Cargo\.toml)$ - id: rust-cargo-fmt name: Rust - Format code entry: tools/pre_commit/rust-check.sh fmt diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 7cd52d6b1c0e..bf19fd3225ce 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5938,6 +5938,18 @@ dependencies = [ "xgrammar-structural-tag", ] +[[package]] +name = "vllm-proto" +version = "0.1.0" +dependencies = [ + "prost", + "prost-types", + "protox", + "tonic", + "tonic-prost", + "tonic-prost-build", +] + [[package]] name = "vllm-server" version = "0.1.0" @@ -5960,9 +5972,7 @@ dependencies = [ "libc", "llm-multimodal", "openssl", - "prost", "prost-types", - "protox", "rmp-serde", "rmpv", "serde", @@ -5982,8 +5992,6 @@ dependencies = [ "tokio-util", "tonic", "tonic-health", - "tonic-prost", - "tonic-prost-build", "tower", "tower-http", "tracing", @@ -5997,6 +6005,7 @@ dependencies = [ "vllm-engine-core-client", "vllm-llm", "vllm-metrics", + "vllm-proto", "vllm-text", "vllm-tokenizer", "zeromq", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index a9421e7c4aca..75104642c345 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ + "proto", "src/bench", "src/build-info", "src/chat", @@ -146,6 +147,7 @@ vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } +vllm-proto = { version = "0.1.0", path = "proto" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml new file mode 100644 index 000000000000..3084c8c5b364 --- /dev/null +++ b/rust/proto/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "vllm-proto" +version = "0.1.0" +edition.workspace = true +license.workspace = true +description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" +repository = "https://github.com/vllm-project/vllm" +readme = "README.md" +include = ["*.proto", "src/**", "build.rs", "Cargo.toml", "README.md"] + +[dependencies] +prost.workspace = true +prost-types.workspace = true +tonic.workspace = true +tonic-prost.workspace = true + +[build-dependencies] +protox.workspace = true +tonic-prost-build.workspace = true + +[lints] +workspace = true diff --git a/rust/proto/README.md b/rust/proto/README.md index a9960f3bb8ab..3c5dbe0d3a61 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -10,3 +10,53 @@ The schema is published to `buf.build/vllm-project/vllm`: - Buf commits and generated SDK versions are immutable and can be pinned by consumers. Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically. + +## Rust crate + +`vllm-proto` exposes the Prost message types and Tonic client and server modules for both APIs: + +```toml +[dependencies] +vllm-proto = "0.1" +``` + +For example, import `vllm_proto::inference_client::InferenceClient` or +`vllm_proto::control_client::ControlClient`. The vLLM frontend uses this same crate. +The package includes the canonical `.proto` files and generates bindings during +compilation using `protox` and `tonic-prost-build`. Consumers do not need Buf +credentials, the vLLM checkout, or a separately installed `protoc`. + +### Versioning and releases + +The crate has its own version in `Cargo.toml`, independent of the vLLM release +number and the other Rust workspace crates. Release it when changes to the +schema, generated Rust API, or dependencies require a new version. Review both +Rust API and protobuf wire compatibility when choosing the version bump; +a wire-compatible schema addition can still break Rust callers. + +On pull requests and releases, `cargo-semver-checks` compares the crate with its +latest published version. Include any required version bump in the protocol +change PR. This check becomes available after the first manual publication. + +1. Update the crate version and the `vllm-proto` workspace dependency together, + and update `rust/Cargo.lock`. +2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run` + and the frontend gRPC tests. Record the tested vLLM releases or revisions in + the release notes; matching crate versions alone do not establish runtime compatibility. +3. After the change merges, create a `proto-v` tag on that commit. + The `proto-crate.yml` workflow checks that the tag matches the crate version, + verifies the package, and publishes it to crates.io. + +Use `proto-v*` tags so they do not trigger Buf's `v*` release workflow. There is +no scheduled crate publication. Pull requests and manual workflow runs only +verify packaging, including on forks. Retry a failed release by rerunning its +tag workflow; crates.io versions cannot be overwritten. + +### Publisher setup + +A vLLM maintainer must publish the first version and establish the crate's owners. +For subsequent releases, configure [crates.io Trusted Publishing](https://crates.io/docs/trusted-publishing) +for repository `vllm-project/vllm`, workflow `proto-crate.yml`, and environment +`crates-io`. Configure that GitHub environment to allow `proto-v*` release tags. +Only tag pushes in the upstream repository can publish; forks do not obtain +publishing credentials. diff --git a/rust/src/server/build.rs b/rust/proto/build.rs similarity index 88% rename from rust/src/server/build.rs rename to rust/proto/build.rs index 1a2ba8a6b095..27cae507727b 100644 --- a/rust/src/server/build.rs +++ b/rust/proto/build.rs @@ -3,7 +3,7 @@ fn main() -> Result<(), Box> { let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); - let proto_dir = std::path::PathBuf::from(format!("{manifest_dir}/../../proto")); + let proto_dir = std::path::PathBuf::from(manifest_dir); let protos = [ proto_dir.join("control.proto"), proto_dir.join("inference.proto"), diff --git a/rust/proto/src/lib.rs b/rust/proto/src/lib.rs new file mode 100644 index 000000000000..d4d44ea099f5 --- /dev/null +++ b/rust/proto/src/lib.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Protobuf messages and gRPC clients and servers for vLLM's Inference and Control APIs. + +tonic::include_proto!("vllm"); diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 6f5e95e93f44..e4d12bc96b39 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -19,7 +19,6 @@ itertools.workspace = true libc.workspace = true llm-multimodal.workspace = true openssl.workspace = true -prost.workspace = true prost-types.workspace = true rmpv.workspace = true serde.workspace = true @@ -37,7 +36,6 @@ tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true tonic-health.workspace = true -tonic-prost.workspace = true tower.workspace = true tower-http.workspace = true tracing.workspace = true @@ -51,12 +49,9 @@ vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true vllm-metrics.workspace = true +vllm-proto.workspace = true vllm-text.workspace = true -[build-dependencies] -protox.workspace = true -tonic-prost-build.workspace = true - [dev-dependencies] anyhow.workspace = true async-openai = { workspace = true, features = ["full"] } diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 06c7c9259eb8..3c96afd9b811 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -9,9 +9,7 @@ mod health; mod inference; /// Generated protobuf/gRPC types for the `vllm` package. -pub mod pb { - tonic::include_proto!("vllm"); -} +pub use vllm_proto as pb; pub(crate) use control::ControlGrpcService; pub use control::ControlServiceImpl; From 9dcab802154d89d25de22332c458168673afab89 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 11 Sep 2026 10:53:04 +0200 Subject: [PATCH 034/774] [Bugfix][Bench] Fix bench mm-processor crash in shared request sampling (#56300) Signed-off-by: NickLucche --- tests/benchmarks/test_throughput_cli.py | 89 +++++++++++++++++++++++++ vllm/benchmarks/throughput.py | 31 ++++++--- 2 files changed, 110 insertions(+), 10 deletions(-) diff --git a/tests/benchmarks/test_throughput_cli.py b/tests/benchmarks/test_throughput_cli.py index 6af4aad4d908..c4a39f24d948 100644 --- a/tests/benchmarks/test_throughput_cli.py +++ b/tests/benchmarks/test_throughput_cli.py @@ -17,6 +17,7 @@ from vllm.benchmarks.datasets import SampleRequest from vllm.benchmarks.throughput import ( _run_vllm_chat_requests, + _to_serve_args, add_cli_args, assign_loras, get_requests, @@ -277,3 +278,91 @@ def sample(self, *, num_requests, tokenizer, **kwargs): # noqa: ARG002 requests = get_requests(args, hf_tokenizer) assert len(requests) == 3 assert all(isinstance(r, SampleRequest) for r in requests) + + +def test_to_serve_args_supports_mm_processor_flags() -> None: + """The adapter must serve callers without throughput's legacy flags. + + ``bench mm-processor`` exposes neither --input-len/--output-len/--prefix-len + nor --backend, and leaves --random-prefix-len at its default of 0. The + adapter must not read the legacy flags unconditionally, and must keep the + explicit 0 instead of coercing it to None (which later crashed + RandomMultiModalDataset.get_prefix on ``None <= 0``). + """ + args = SimpleNamespace( + random_input_len=32, + random_output_len=1, + random_prefix_len=0, + ) + serve_args = _to_serve_args(args) + + assert serve_args.random_input_len == 32 + assert serve_args.random_output_len == 1 + assert serve_args.random_prefix_len == 0 + assert serve_args.backend == "vllm-chat" + assert serve_args.enable_multimodal_chat + assert serve_args.no_oversample is False + + +def test_to_serve_args_keeps_legacy_flag_fallback() -> None: + """Callers that only set the legacy flags still get them mapped.""" + args = SimpleNamespace( + input_len=256, + output_len=8, + prefix_len=16, + backend="vllm", + ) + serve_args = _to_serve_args(args) + + assert serve_args.random_input_len == 256 + assert serve_args.random_output_len == 8 + assert serve_args.random_prefix_len == 16 + assert serve_args.backend == "vllm" + assert not serve_args.enable_multimodal_chat + + +@pytest.mark.benchmark +def test_get_requests_random_mm_with_mm_processor_args( + hf_tokenizer: PreTrainedTokenizerBase, +) -> None: + """Regression: bench mm-processor's random-mm sampling must not crash. + + get_requests feeds mm-processor's parsed args (no legacy length/backend + flags) through the shared get_samples dispatch. With the default + --random-prefix-len 0 this previously crashed inside + RandomMultiModalDataset.get_prefix with a TypeError. + """ + from vllm.benchmarks.mm_processor import add_cli_args as add_mm_cli_args + + parser = FlexibleArgumentParser() + add_mm_cli_args(parser) + args = parser.parse_args( + [ + "--dataset-name", + "random-mm", + "--num-prompts", + "3", + "--random-input-len", + "32", + "--random-output-len", + "1", + "--random-mm-base-items-per-request", + "1", + "--random-mm-num-mm-items-range-ratio", + "0", + "--random-mm-limit-mm-per-prompt", + '{"image": 1}', + "--random-mm-bucket-config", + "{(224, 224, 1): 1.0}", + ] + ) + + requests = get_requests(args, hf_tokenizer) + assert len(requests) == 3 + assert all(isinstance(r, SampleRequest) for r in requests) + # mm-processor drives llm.chat, so prompts must be chat-formatted with + # multimodal content attached. + for req in requests: + assert isinstance(req.prompt, list) + assert req.prompt[0]["role"] == "user" + assert any(item.get("type") == "image_url" for item in req.prompt[0]["content"]) diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 17b2f2426477..45fdcd9694f0 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -539,27 +539,38 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace: """ d = vars(args).copy() # random_*: prefer --random-* over legacy --input/output/prefix-len. - d["random_input_len"] = getattr(args, "random_input_len", None) or args.input_len - d["random_output_len"] = getattr(args, "random_output_len", None) or args.output_len - d["random_prefix_len"] = getattr(args, "random_prefix_len", None) or args.prefix_len + # Callers without the legacy flags (e.g. bench mm-processor) omit them, so + # read them optionally. Identity checks, not truthiness: an explicit 0 + # (the default of --random-prefix-len) must not fall through to them. + input_len = getattr(args, "input_len", None) + output_len = getattr(args, "output_len", None) + prefix_len = getattr(args, "prefix_len", None) + d["random_input_len"] = getattr(args, "random_input_len", input_len) + d["random_output_len"] = getattr(args, "random_output_len", output_len) + d["random_prefix_len"] = getattr(args, "random_prefix_len", prefix_len) # --output-len maps to the per-dataset output-len entry points get_samples # reads. None passes through (each dataset applies its own default), # matching prior throughput behaviour of omitting output_len when unset. - d["hf_output_len"] = args.output_len - d["sharegpt_output_len"] = args.output_len + d["hf_output_len"] = output_len + d["sharegpt_output_len"] = output_len # sonnet reads dedicated attrs; fall back to SonnetDataset's own defaults. - d["sonnet_input_len"] = args.input_len if args.input_len is not None else 550 - d["sonnet_output_len"] = args.output_len if args.output_len is not None else 150 - d["sonnet_prefix_len"] = args.prefix_len + d["sonnet_input_len"] = input_len if input_len is not None else 550 + d["sonnet_output_len"] = output_len if output_len is not None else 150 + d["sonnet_prefix_len"] = prefix_len # Explicit --enable-multimodal-chat wins; otherwise auto-enable for the - # multimodal chat backend (preserves today's vllm-chat handling). + # multimodal chat backend (preserves today's vllm-chat handling). Callers + # without a --backend flag (e.g. bench mm-processor) drive every request + # through llm.chat, so default to the vllm-chat behaviour. + backend = getattr(args, "backend", "vllm-chat") + d["backend"] = backend d["enable_multimodal_chat"] = bool( - getattr(args, "enable_multimodal_chat", False) or args.backend == "vllm-chat" + getattr(args, "enable_multimodal_chat", False) or backend == "vllm-chat" ) # serve-only attrs throughput never exposed; keep serve's defaults. d.setdefault("disable_shuffle", False) d.setdefault("skip_chat_template", False) d.setdefault("no_stream", False) + d.setdefault("no_oversample", False) d.setdefault("request_id_prefix", "") d.setdefault("chat_template_kwargs", None) return argparse.Namespace(**d) From d0dfe587d5b20609af4de8e8d968713116b80eb6 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 11 Sep 2026 03:56:05 -0500 Subject: [PATCH 035/774] [Bugfix] Fall back to full decode graphs for noncompiled models (#55095) Signed-off-by: Andreas Karatzas Signed-off-by: Lucas Wilkinson Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson --- tests/test_config.py | 122 ++++++++++++++++++ .../spec_decode/test_adaptive_verification.py | 22 ++++ vllm/config/compilation.py | 22 ++++ vllm/config/vllm.py | 50 ++++--- vllm/platforms/rocm.py | 17 +++ vllm/v1/worker/gpu/model_runner.py | 11 +- .../gpu/spec_decode/adaptive_verification.py | 11 ++ 7 files changed, 237 insertions(+), 18 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 96ca8f36e4ed..827dbcd4f8c8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -273,6 +273,44 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): default_breakable_cudagraph_architectures.cache_clear() +@pytest.mark.parametrize( + ("architecture", "use_v2", "mode", "expected"), + [ + ("DeepseekV4ForCausalLM", True, None, CUDAGraphMode.NONE), + ("DeepseekV4ForConditionalGeneration", True, None, CUDAGraphMode.NONE), + ("DeepseekV4ForCausalLM", False, None, None), + ("LlamaForCausalLM", True, None, None), + ( + "DeepseekV4ForCausalLM", + True, + CUDAGraphMode.FULL_DECODE_ONLY, + CUDAGraphMode.FULL_DECODE_ONLY, + ), + ], +) +def test_rocm_gfx950_deepseek_v4_cudagraph_default( + monkeypatch, architecture, use_v2, mode, expected +): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.platforms import rocm + + monkeypatch.setattr(rocm, "on_gfx950", lambda: True) + monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: False) + monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: False) + monkeypatch.setattr( + rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", lambda: False + ) + config = SimpleNamespace( + compilation_config=CompilationConfig(cudagraph_mode=mode), + model_config=SimpleNamespace(architecture=architecture), + use_v2_model_runner=use_v2, + ) + + rocm.RocmPlatform.apply_config_platform_defaults(config) + + assert config.compilation_config.cudagraph_mode == expected + + @pytest.mark.parametrize( ("model", "architecture"), [ @@ -463,6 +501,66 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes +@pytest.mark.parametrize( + ("mode", "piecewise_capture_available", "attention_support", "expected"), + [ + ("PIECEWISE", False, "ALWAYS", "NONE"), + ("FULL_AND_PIECEWISE", False, "ALWAYS", "FULL_DECODE_ONLY"), + ("FULL_DECODE_ONLY", False, "ALWAYS", "FULL_DECODE_ONLY"), + ("FULL_DECODE_ONLY", False, "NEVER", "NONE"), + ], +) +def test_resolve_cudagraph_mode_uses_loaded_piecewise_provider( + mode, piecewise_capture_available, attention_support, expected +): + compilation_config = CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + cudagraph_mode=CUDAGraphMode[mode], + use_inductor_graph_partition=True, + ) + + resolved = compilation_config.resolve_cudagraph_mode_and_sizes( + AttentionCGSupport[attention_support], + "FakeAttentionBackend", + piecewise_capture_available=piecewise_capture_available, + ) + + assert resolved.name == expected + assert compilation_config.cudagraph_mode == resolved + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Requires CUDA graph support" +) +@pytest.mark.parametrize( + "engine_kwargs", + [ + {"runner": "pooling", "convert": "embed"}, + pytest.param( + {"prefill_context_parallel_size": 2, "tensor_parallel_size": 2}, + marks=pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm PCP graph restriction" + ), + ), + ], +) +def test_late_piecewise_restrictions_without_compilation(monkeypatch, engine_kwargs): + """Late compatibility overrides must not restore unavailable piecewise graphs.""" + from vllm.engine.arg_utils import EngineArgs + + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + config = EngineArgs( + model="facebook/opt-125m", + compilation_config=CompilationConfig(mode=CompilationMode.NONE), + **engine_kwargs, + ).create_engine_config() + + assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE + assert config.compilation_config.cudagraph_capture_sizes == [] + assert config.compilation_config.max_cudagraph_capture_size == 0 + + def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): """Cudagraph memory profiling uses a minimal KV cache, so the Mamba block-count guard must only fire for the real cache sizing.""" @@ -498,6 +596,30 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE +@pytest.mark.parametrize( + ("graph_mode", "should_raise"), + [ + (CUDAGraphMode.FULL_DECODE_ONLY, False), + (CUDAGraphMode.FULL_AND_PIECEWISE, False), + (CUDAGraphMode.NONE, True), + (CUDAGraphMode.PIECEWISE, True), + ], +) +def test_adaptive_verification_requires_full_cudagraphs(graph_mode, should_raise): + config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=True), + lora_config=None, + compilation_config=CompilationConfig(cudagraph_mode=graph_mode), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + ) + + if should_raise: + with pytest.raises(ValueError, match="requires full CUDA graphs"): + VllmConfig._validate_adaptive_verification(config) + else: + VllmConfig._validate_adaptive_verification(config) + + @pytest.mark.parametrize( ("model_config", "expected"), [ diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index d92f78a71535..b14fc126bbe9 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -4,7 +4,9 @@ from types import SimpleNamespace import numpy as np +import pytest +from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.worker.gpu.async_utils import StepTimingSample from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo @@ -12,6 +14,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, maybe_create_adaptive_verification_manager, + resolve_adaptive_cudagraph_mode, ) from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping @@ -35,6 +38,25 @@ def make_manager( return manager +@pytest.mark.parametrize( + ("mode", "piecewise_capture_available", "expected"), + [ + ("FULL_DECODE_ONLY", True, "FULL_DECODE_ONLY"), + ("FULL", True, "FULL_AND_PIECEWISE"), + ("FULL", False, "FULL_DECODE_ONLY"), + ("FULL_AND_PIECEWISE", True, "FULL_AND_PIECEWISE"), + ("FULL_AND_PIECEWISE", False, "FULL_DECODE_ONLY"), + ], +) +def test_resolve_adaptive_cudagraph_mode(mode, piecewise_capture_available, expected): + assert ( + resolve_adaptive_cudagraph_mode( + CUDAGraphMode[mode], piecewise_capture_available=piecewise_capture_available + ) + == CUDAGraphMode[expected] + ) + + def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch): class Backend: @classmethod diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 1219607d3abf..2cd5fae77554 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -77,6 +77,13 @@ def has_mode(self, mode: "CUDAGraphMode") -> bool: def requires_piecewise_compilation(self) -> bool: return self.has_mode(CUDAGraphMode.PIECEWISE) + def without_piecewise(self) -> "CUDAGraphMode": + if self == CUDAGraphMode.PIECEWISE: + return CUDAGraphMode.NONE + if self == CUDAGraphMode.FULL_AND_PIECEWISE: + return CUDAGraphMode.FULL_DECODE_ONLY + return self + def max_cudagraph_mode(self) -> "CUDAGraphMode": return CUDAGraphMode(max(self.value)) if self.separate_routine() else self @@ -1382,6 +1389,7 @@ def resolve_cudagraph_mode_and_sizes( kv_cache_config: "KVCacheConfig | None" = None, max_num_reqs: int | None = None, is_profiling: bool = False, + piecewise_capture_available: bool = True, ) -> CUDAGraphMode: from vllm.v1.attention.backend import AttentionCGSupport @@ -1464,6 +1472,20 @@ def resolve_cudagraph_mode_and_sizes( cudagraph_mode = CUDAGraphMode.NONE logger.warning(msg) + if ( + not piecewise_capture_available + and cudagraph_mode.requires_piecewise_compilation() + ): + fallback_mode = cudagraph_mode.without_piecewise() + logger.warning_once( + "Cudagraph mode %s requires piecewise capture, but the loaded " + "model provides neither a compiled submodule nor breakable CUDA " + "graphs. Overriding to %s.", + cudagraph_mode, + fallback_mode, + ) + cudagraph_mode = fallback_mode + # double check that we can support full cudagraph if they are requested # even after automatic downgrades if ( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index cc443a14dd47..5bd6af20deca 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -991,6 +991,28 @@ def _maybe_disable_dynamic_sd_for_data_parallel(self) -> None: ) speculative_config.num_speculative_tokens_per_batch_size = None + def _normalize_piecewise_cudagraph_mode( + self, *, breakable_cudagraph_enabled: bool + ) -> None: + compilation_config = self.compilation_config + if ( + compilation_config.cudagraph_mode.requires_piecewise_compilation() + and compilation_config.mode != CompilationMode.VLLM_COMPILE + and not breakable_cudagraph_enabled + ): + fallback_mode = compilation_config.cudagraph_mode.without_piecewise() + logger.info_once( + "Cudagraph mode %s is not compatible with compilation mode %s. " + "Overriding to %s.", + compilation_config.cudagraph_mode, + compilation_config.mode, + fallback_mode, + ) + compilation_config.cudagraph_mode = fallback_mode + if fallback_mode == CUDAGraphMode.NONE: + compilation_config.max_cudagraph_capture_size = 0 + compilation_config.cudagraph_capture_sizes = [] + def _post_init_kv_transfer_config(self) -> None: """Update KVTransferConfig based on top-level configs in VllmConfig. @@ -1536,18 +1558,9 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() - if ( - self.compilation_config.cudagraph_mode.requires_piecewise_compilation() - and self.compilation_config.mode != CompilationMode.VLLM_COMPILE - and not envs.VLLM_USE_BREAKABLE_CUDAGRAPH - ): - logger.info_once( - "Cudagraph mode %s is not compatible with compilation mode %s." - "Overriding to NONE.", - self.compilation_config.cudagraph_mode, - self.compilation_config.mode, - ) - self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + self._normalize_piecewise_cudagraph_mode( + breakable_cudagraph_enabled=breakable_cudagraph_enabled + ) # async tp is built on top of sequence parallelism and requires it. pass_config = self.compilation_config.pass_config @@ -1716,6 +1729,10 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) + self._normalize_piecewise_cudagraph_mode( + breakable_cudagraph_enabled=breakable_cudagraph_enabled + ) + self._resolve_mm_embedding_inputs() self._resolve_mm_processor_device() self._validate_mm_processor_device() @@ -2728,12 +2745,11 @@ def _validate_adaptive_verification(self) -> None: "Adaptive verification is not currently compatible with LoRA" ) - if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: - # The draft budget divides by step costs profiled from captured - # cudagraphs; eager execution captures none. + if not self.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise ValueError( - "Adaptive verification is not currently compatible with " - "enforce_eager/cudagraph_mode=none" + "Adaptive verification requires full CUDA graphs. Use cudagraph " + "mode FULL, FULL_DECODE_ONLY, or FULL_AND_PIECEWISE, or disable " + "adaptive verification." ) if self.parallel_config.pipeline_parallel_size > 1: diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index bdaa86c5decc..297efc632f92 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -880,8 +880,25 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: from vllm._aiter_ops import rocm_aiter_ops + from vllm.config.compilation import CUDAGraphMode compilation_config = vllm_config.compilation_config + model_config = vllm_config.model_config + if ( + compilation_config.cudagraph_mode is None + and model_config is not None + and on_gfx950() + and vllm_config.use_v2_model_runner + and model_config.architecture + in { + "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", + } + ): + # Default to eager after reported gfx950/MRV2 accuracy regressions: + # https://github.com/vllm-project/vllm/issues/52644 + compilation_config.cudagraph_mode = CUDAGraphMode.NONE + use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled() use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled() use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6a3d2e34e9ac..f41a5a8d3ae1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -104,6 +104,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, ModelCudaGraphManager, + has_compiled_submodule, make_cudagraph_stats, ) from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -152,6 +153,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, maybe_create_adaptive_verification_manager, + resolve_adaptive_cudagraph_mode, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, @@ -674,8 +676,14 @@ def initialize_kv_cache( self.kv_cache_config, use_replayssm=self.cache_config.use_replayssm, ) + piecewise_capture_available = bool( + envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model) + ) if self.adaptive_verification is not None: - self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode( + self.compilation_config.cudagraph_mode, + piecewise_capture_available=piecewise_capture_available, + ) cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, @@ -685,6 +693,7 @@ def initialize_kv_cache( kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, + piecewise_capture_available=piecewise_capture_available, ) self.cudagraph_manager = ModelCudaGraphManager( self.vllm_config, diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index d3b3743d67fd..af58b8993bfa 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -10,6 +10,7 @@ import torch import vllm.envs as envs +from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -33,6 +34,16 @@ from vllm.v1.worker.utils import AttentionGroup +def resolve_adaptive_cudagraph_mode( + mode: CUDAGraphMode, *, piecewise_capture_available: bool +) -> CUDAGraphMode: + """Select a separate decode route for variable-length verification.""" + assert mode.has_full_cudagraphs() + if mode == CUDAGraphMode.FULL_DECODE_ONLY or not piecewise_capture_available: + return CUDAGraphMode.FULL_DECODE_ONLY + return CUDAGraphMode.FULL_AND_PIECEWISE + + def _assign_draft_token_budget( confidence_probs: torch.Tensor, idx_mapping: torch.Tensor, From 1cf6555214446cfb98c42ff5d07baab784925c1b Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 11 Sep 2026 01:57:47 -0700 Subject: [PATCH 036/774] [Bugfix] Pin EPLB and MLA host-to-device transfer buffers (#56138) Signed-off-by: khluu Signed-off-by: Kevin H. Luu Co-authored-by: OpenAI Codex --- tests/distributed/test_eplb_execute.py | 25 ++++++++++++ tests/v1/attention/test_mla_context_chunks.py | 40 +++++++++++++++++++ vllm/distributed/eplb/eplb_communicator.py | 5 ++- .../layers/attention/sparse_mla_attention.py | 16 ++++++-- 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 21fa057fd201..1ff41f5df269 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -7,6 +7,8 @@ import torch import torch.distributed +import vllm.distributed.eplb.eplb_communicator as eplb_comm +import vllm.utils.gpu_sync_debug as gsd from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.eplb.eplb_communicator import ( create_eplb_communicator, @@ -25,6 +27,29 @@ from .eplb_utils import distributed_run, set_env_vars_and_device +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch): + """Gloo receives must reach the GPU without an implicit pageable copy.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + gsd._install_copy_checkers() + monkeypatch.setattr(eplb_comm, "is_local_first_rank", lambda: False) + + monkeypatch.setattr(eplb_comm, "P2POp", lambda op, tensor, peer, group: tensor) + + def receive(tensors): + for tensor in tensors: + tensor.fill_(7) + return [] + + monkeypatch.setattr(eplb_comm, "batch_isend_irecv", receive) + communicator = eplb_comm.TorchDistGlooStagedEplbCommunicator(cpu_group=None) + dst = torch.empty(32, device="cuda") + communicator.add_recv([dst], src_rank=1, expert_id=0) + gsd.with_gpu_sync_check(communicator.execute)() + torch.testing.assert_close(dst.cpu(), torch.full((32,), 7.0)) + + def create_expert_indices_with_redundancy( num_layers: int, num_logical_experts: int, diff --git a/tests/v1/attention/test_mla_context_chunks.py b/tests/v1/attention/test_mla_context_chunks.py index 091079acc639..2ed73b2c987e 100644 --- a/tests/v1/attention/test_mla_context_chunks.py +++ b/tests/v1/attention/test_mla_context_chunks.py @@ -8,14 +8,20 @@ longer needs an empty-span masking pass). """ +from types import SimpleNamespace + import pytest import torch +import vllm.utils.gpu_sync_debug as gsd from vllm.model_executor.layers.attention.mla_attention import ( build_mla_chunked_context_metadata, init_mla_context_partial, reorg_kvcache, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SparseMLACommonMetadataBuilder, +) BLOCK_SIZE = 16 @@ -45,6 +51,40 @@ def build_chunked_context( ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch): + """Sparse MLA must pin computed context lengths before the H2D copy.""" + builder = SimpleNamespace( + chunked_prefill_workspace=torch.empty((2048, 1)), + chunked_prefill_workspace_size=1024, + kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE), + device=torch.device("cuda"), + dcp_world_size=1, + dcp_local_block_size=1, + dcp_virtual_block_size=1, + dcp_manager=None, + ) + common_metadata = SimpleNamespace( + seq_lens_cpu_upper_bound=torch.tensor([17, 2052, 324], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1, 5, 9], dtype=torch.int32), + ) + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + gsd._install_copy_checkers() + metadata = gsd.with_gpu_sync_check( + SparseMLACommonMetadataBuilder._build_chunked_context_fields + )( + builder, + common_metadata, + num_decodes=1, + num_prefills=2, + prefill_query_lens_cpu=torch.tensor([4, 4], dtype=torch.int32), + ) + assert metadata is not None + assert metadata.context_lens.device.type == "cuda" + assert metadata.context_lens.tolist() == [2048, 320] + + @pytest.mark.parametrize( "context_lens,workspace_size", [ diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 77362e593c2a..162c0f58430b 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -34,6 +34,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import PIN_MEMORY logger = init_logger(__name__) @@ -204,7 +205,9 @@ def build_ops() -> None: ) ) continue - cpu_tensor = torch.empty_like(tensor, device="cpu") + cpu_tensor = torch.empty_like( + tensor, device="cpu", pin_memory=PIN_MEMORY + ) p2p_ops.append( P2POp( torch.distributed.irecv, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index ab6946132f4b..8c2bb72bf4a0 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -28,7 +28,11 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import has_flashinfer -from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor +from vllm.utils.torch_utils import ( + PIN_MEMORY, + is_quantized_kv_cache, + np_to_pinned_tensor, +) from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.utils import split_decodes_and_prefills @@ -241,9 +245,13 @@ def _build_chunked_context_fields( seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None - context_lens_cpu = ( - seq_lens_cpu[num_decodes : num_decodes + num_prefills] - - prefill_query_lens_cpu + context_lens_cpu = torch.empty( + num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY + ) + torch.subtract( + seq_lens_cpu[num_decodes : num_decodes + num_prefills], + prefill_query_lens_cpu, + out=context_lens_cpu, ) qsl_cpu = common_attn_metadata.query_start_loc_cpu prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes] From e6821ceac91f4d61c7a8603f6de21d9c8935effb Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 11 Sep 2026 16:58:27 +0800 Subject: [PATCH 037/774] [XPU][CI] Add decord to test requirements (#56355) Signed-off-by: Tony Lin --- requirements/test/xpu.in | 1 + requirements/test/xpu.txt | 3 +++ 2 files changed, 4 insertions(+) diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 7120b3104302..6848dd98b7d5 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -41,4 +41,5 @@ pqdm # --- Vision & Multimodal --- timm albumentations +decord==0.6.0 mistral-common[image,audio] \ No newline at end of file diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index f3e8bb1bc678..99e569999390 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -115,6 +115,8 @@ datasets==5.0.1 # mteb decorator==5.3.1 # via librosa +decord==0.6.0 + # via -r requirements/test/xpu.in defusedxml==0.7.1 # via nltk depyf==0.20.0 @@ -442,6 +444,7 @@ numpy==2.3.5 # albumentations # bm25s # datasets + # decord # evaluate # librosa # lm-eval From e77daef89e18e08321ae7b8b24827eedd5fe8673 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 11 Sep 2026 02:11:19 -0700 Subject: [PATCH 038/774] [Model] Support DeepSeek-V4.1-Flash (#56214) Signed-off-by: Yongye Zhu Signed-off-by: Jee Jee Li Signed-off-by: andyluo7 Signed-off-by: zjy0516 Co-authored-by: Jee Jee Li Co-authored-by: andyluo7 Co-authored-by: OpenAI Codex Co-authored-by: zjy0516 Co-authored-by: Claude Opus 5 (1M context) --- .buildkite/test_areas/kernels.yaml | 3 + .../fusion/test_quant_activation_contract.py | 8 + .../attention/test_rocm_triton_attn_dsv4.py | 25 + tests/kernels/core/test_fused_q_kv_rmsnorm.py | 189 +++++ .../quantization/test_rocm_mxfp8_linear.py | 144 ++++ tests/kernels/test_compressor_kv_cache.py | 432 ++++++++++ tests/kernels/test_engram.py | 778 ++++++++++++++++++ .../test_fused_indexer_q_rope_quant.py | 368 ++++++++- .../kernels/test_fused_inv_rope_fp8_quant.py | 105 ++- tests/kernels/test_mhc_kernels.py | 286 ++++++- .../layers/test_mla_short_prefill_indexer.py | 110 +++ .../multimodal/generation/test_maverick.py | 2 + tests/models/registry.py | 9 + tests/models/test_deepseek_v4_mega_moe.py | 215 ++++- tests/models/test_dspark_mla.py | 76 ++ tests/models/test_registry.py | 1 + tests/quantization/test_fp8.py | 378 +++++++++ tests/quantization/test_modelopt.py | 17 + .../test_indexer_deepseek_v4_slot_mapping.py | 44 +- tests/v1/core/test_contiguous_kv_packing.py | 102 +++ .../v1/cudagraph/test_breakable_cudagraph.py | 51 ++ .../unit/test_mooncake_store_hma_e2e.py | 99 +++ .../unit/test_nixl_desc_geometry.py | 109 +++ .../unit/test_simple_cpu_offload_connector.py | 57 ++ tests/v1/simple_kv_offload/test_scheduler.py | 65 +- vllm/_aiter_ops.py | 10 +- vllm/compilation/breakable_cudagraph.py | 16 +- vllm/config/engram.py | 28 +- vllm/config/kernel.py | 1 + vllm/config/speculative.py | 48 +- vllm/config/vllm.py | 1 + .../v1/mooncake/store/coordinator.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 8 +- .../layers/fusion/quant_activation.py | 12 +- .../layers/quantization/__init__.py | 17 +- .../layers/quantization/utils/fp8_utils.py | 26 +- .../layers/quantization/utils/mxfp8_utils.py | 10 +- vllm/model_executor/models/config.py | 5 +- vllm/model_executor/models/registry.py | 8 + .../warmup/deepseek_v4_mhc_warmup.py | 5 +- .../warmup/flashinfer_sparse_mla_warmup.py | 7 +- vllm/models/deepseek_v4_1/attention.py | 4 +- vllm/v1/attention/backends/mla/indexer.py | 17 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 112 ++- vllm/v1/core/kv_cache_utils.py | 37 +- vllm/v1/simple_kv_offload/manager.py | 34 +- .../v1/worker/gpu/spec_decode/dspark/utils.py | 4 +- .../gpu/spec_decode/eagle/eagle3_utils.py | 9 +- vllm/v1/worker/gpu_model_runner.py | 41 +- 49 files changed, 3996 insertions(+), 138 deletions(-) create mode 100644 tests/kernels/quantization/test_rocm_mxfp8_linear.py create mode 100644 tests/kernels/test_engram.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 2b43938060e9..ed4ea67f933f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -102,6 +102,8 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/kernels/ commands: + # Isolate mHC from FlashInfer's import-time ptxas override. + - pytest -v -s kernels/test_mhc_kernels.py - pytest -v -s kernels/ --ignore=kernels/attention --ignore=kernels/core @@ -117,6 +119,7 @@ steps: --ignore=kernels/test_kda.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_mhc_kernels.py --ignore=kernels/test_shuffle_rows.py # BROKEN on main, pending kernel fixes (B200): # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index a794ff1e3a67..f698a272f5ea 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -9,8 +9,13 @@ _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_FP8_KERNELS, _POSSIBLE_INT8_KERNELS, + _POSSIBLE_MXFP8_KERNELS, _POSSIBLE_NVFP4_KERNELS, ) +from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.base import ( NvFp4LinearKernel, NvFp4LinearLayerConfig, @@ -50,6 +55,8 @@ FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, } @@ -60,6 +67,7 @@ def _all_kernel_classes() -> list[type]: _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_INT8_KERNELS, _POSSIBLE_NVFP4_KERNELS, + _POSSIBLE_MXFP8_KERNELS, ): for kernels in registry.values(): for cls in kernels: diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index fb696edd840f..1c389a111733 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1268,6 +1268,31 @@ def test_get_cached_wo_a_bf16_plain_caches() -> None: torch.testing.assert_close(out2, expected, atol=0, rtol=0) +@pytest.mark.parametrize("scale_attr", ["weight_scale", "weight_scale_inv"]) +@torch.inference_mode() +def test_get_cached_wo_a_bf16_dequantized_ignores_retained_scale( + scale_attr: str, +) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 + + device = torch.device("cuda") + torch.manual_seed(6) + n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8 + weight = torch.randn( + n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device + ) + retained_scale = torch.full( + (n_local_groups, 1), 0.25, dtype=torch.float32, device=device + ) + wo_a = _FakeWoA(weight) + setattr(wo_a, scale_attr, retained_scale) + + out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) + + expected = weight.view(n_local_groups, o_lora_rank, hidden_dim) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + @torch.inference_mode() def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 8debe7c3a35b..7e5e41048791 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -112,3 +112,192 @@ def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int): rtol=1e-2, atol=1e-2, ) + + +@pytest.mark.skipif( + not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell" +) +@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 129, 1024]) +@pytest.mark.parametrize("q_size", [384, 1280, 1344]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype): + """Preserve native MXFP8 rounding and all swizzled scale padding bytes.""" + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + fused_q_kv_rmsnorm_quant, + ) + + torch.manual_seed(0) + kv_size = 512 + x = torch.randn(num_tokens, q_size + kv_size, device="cuda", dtype=dtype) + x[:1] = 0 + x[1:2] *= 1e-30 + qr, kv = x.split([q_size, kv_size], -1) + qw = torch.randn(q_size, device="cuda", dtype=x.dtype) + kvw = torch.randn(kv_size, device="cuda", dtype=x.dtype) + eps = 1e-20 + result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps) + assert result[0].orig_dtype == dtype + assert result[0].orig_shape == qr.shape + if not num_tokens: + assert result[0].data.shape == qr.shape + assert result[0].scale.numel() == 0 + assert result[1].shape == kv.shape + return + qr_ref, kv_ref = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps) + q_ref, scale_ref = mxfp8_e4m3_quantize(qr_ref, is_sf_swizzled_layout=True) + for replay in (False, True): + if replay: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps) + graph.replay() + torch.testing.assert_close( + result[0].data.view(torch.uint8), q_ref.view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close(result[0].scale, scale_ref, rtol=0, atol=0) + torch.testing.assert_close(result[1], kv_ref, rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell" +) +@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass"]) +@pytest.mark.parametrize("num_tokens", [1, 17, 256]) +@pytest.mark.parametrize("projection", ["attention", "indexer"]) +@pytest.mark.parametrize("bias", [False, True]) +def test_shared_query_quant_preserves_projection( + num_tokens, backend, projection, bias, monkeypatch +): + """The native-checkpoint projection keeps identical GEMM outputs.""" + from dataclasses import replace + from types import SimpleNamespace + from typing import Any, cast + + from vllm.config import ( + CompilationConfig, + CompilationMode, + VllmConfig, + set_current_vllm_config, + ) + from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearLayerConfig + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ) + from vllm.model_executor.layers.fusion.quant_activation import ( + expose_input_quant_key, + ) + from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptLinearMethod, + ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic + from vllm.models.deepseek_v4_1.attention import ( + DeepseekV4Attention, + DeepseekV4Indexer, + ) + from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + can_fuse_query_quant, + fused_q_kv_rmsnorm_quant, + ) + from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + + cls: ( + type[FlashInferCutedslMxfp8LinearKernel] + | type[FlashInferCutlassMxfp8LinearKernel] + ) + if backend == "cute-dsl": + cls = FlashInferCutedslMxfp8LinearKernel + else: + cls = FlashInferCutlassMxfp8LinearKernel + supported, reason = cls.is_supported() + if not supported: + pytest.skip(reason) + torch.manual_seed(0) + with set_current_vllm_config( + VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE)) + ): + kernel = cls(Mxfp8LinearLayerConfig()) + # Match the checkpoint's TP4 attention and replicated indexer widths. + if projection == "attention": + linear = ColumnParallelLinear.__new__(ColumnParallelLinear) + output_size = 8192 + else: + linear = ReplicatedLinear.__new__(ReplicatedLinear) + output_size = 4096 + torch.nn.Module.__init__(linear) + linear.bias = ( + torch.nn.Parameter( + torch.randn(output_size, device="cuda", dtype=torch.bfloat16) + ) + if bias + else None + ) + linear.skip_bias_add = False + linear.return_bias = projection == "indexer" + linear.gather_output = False + w, ws = mxfp8_e4m3_quantize( + torch.randn(output_size, 1280, device="cuda", dtype=torch.bfloat16) + ) + linear.weight = torch.nn.Parameter(w, requires_grad=False) + linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False) + quant_config = DeepseekV4FP8Config.from_config( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [32, 32], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4", + } + ) + method = quant_config.get_quant_method(linear, "model.layers.2.attn.wq_b") + assert type(method) is ModelOptLinearMethod + method.out_dtype = torch.bfloat16 + method.kernel = kernel + linear.quant_method = method + kernel.process_weights_after_loading(linear) + expose_input_quant_key(linear, kernel) + assert can_fuse_query_quant([linear, linear]) + x = torch.randn(num_tokens, 1792, device="cuda", dtype=torch.bfloat16) + qr, kv = x.split([1280, 512], -1) + qw = torch.randn(1280, device="cuda", dtype=x.dtype) + kvw = torch.randn(512, device="cuda", dtype=x.dtype) + qr_ref, _ = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-20) + q, _ = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, 1e-20) + owner = cast(Any, SimpleNamespace(wq_b=linear)) + project = ( + DeepseekV4Attention._wq_b_proj + if projection == "attention" + else DeepseekV4Indexer._wq_b_proj + ) + expected = project(owner, qr_ref) + + def reject_requantization(*args, **kwargs): + raise AssertionError("pre-quantized input must skip quantization") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.mxfp8.flashinfer.mxfp8_e4m3_quantize", + reject_requantization, + ) + actual = project(owner, q) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = project(owner, q) + graph.replay() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert not can_fuse_query_quant( + [linear, torch.nn.Linear(1280, 256, bias=False)] + ) + reshaped = project( + owner, replace(q, orig_shape=torch.Size((1, num_tokens, 1280))) + ) + torch.testing.assert_close(reshaped, expected.unsqueeze(0), rtol=0, atol=0) + with pytest.raises(AssertionError, match="consumer kernel"): + linear(replace(q, quant_key=kNvfp4Dynamic)) diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py new file mode 100644 index 000000000000..1d5e62c6c4b1 --- /dev/null +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``). + +DeepSeek V4.1 routes every attention linear through this kernel. It used to +carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears, +one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token +counts). These tests pin the two properties those workarounds stood in for -- +finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at +TP=4 and the small token counts decode actually runs. +""" + +import pytest +import torch + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + _mxfp8_e4m3_quantize_triton, + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm() or not current_platform.supports_mx(), + reason="native MX linear requires CDNA4 (gfx95x)", +) + +# (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4: +# fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is +# row-parallel (so K, not N, is sharded). +V41_ATTN_SHAPES = [ + ("fused_wqa_wkv", 1280 + 512, 5120), + ("wq_b", 64 * 512 // 4, 1280), + ("wo_a", 8 * 1024 // 4, 4096), + ("wo_b", 5120, 8192 // 4), +] + + +def _make_layer(n: int, k: int, device: str) -> torch.nn.Module: + weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05 + weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(weight, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + return layer + + +@pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v)) +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024]) +@torch.inference_mode() +def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens): + """Output is finite and matches an exact MXFP8 GEMM of the same inputs. + + The reference quantizes the activation with the same quantizer and then + contracts in FP32, so the tolerance covers only kernel/accumulation error, + not the MXFP8 quantization error both sides share. + """ + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + + torch.manual_seed(num_tokens) + device = "cuda" + layer = _make_layer(n, k, device) + kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None)) + kernel.process_weights_after_loading(layer) + + x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5 + out = kernel.apply_weights(layer, x) + + assert out.shape == (num_tokens, n) + assert torch.isfinite(out).all(), f"{name}: non-finite output" + + x_q, x_scale = mxfp8_e4m3_quantize(x) + expected = ( + dequant_mxfp8_to_bf16(x_q, x_scale).float() + @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T + ) + rel = (out.float() - expected).norm() / expected.norm() + assert rel < 5e-3, f"{name}: relative error {rel:.4f}" + + +@torch.inference_mode() +def test_rocm_mxfp8_linear_handles_extreme_magnitudes(): + """Rows that are all-zero or near the E4M3 range must not produce NaNs. + + A degenerate block makes the dynamic activation quantizer pick an extreme + E8M0 scale; a zero row makes it pick the tiny-clamped one. Both feed + ``tl.dot_scaled`` directly. + """ + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + + torch.manual_seed(0) + device = "cuda" + n, k = 5120, 2048 # wo_b at TP=4 + layer = _make_layer(n, k, device) + kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None)) + kernel.process_weights_after_loading(layer) + + x = torch.randn(8, k, device=device, dtype=torch.bfloat16) + x[0] = 0.0 + x[1] = 1e-30 + x[2] = 400.0 + x[3, ::2] = 0.0 + + out = kernel.apply_weights(layer, x) + assert torch.isfinite(out).all() + assert (out[0] == 0).all(), "an all-zero activation row must give a zero row" + + +@pytest.mark.parametrize("num_rows", [1, 64, 65]) +@torch.inference_mode() +def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows): + """The fused ROCm quantizer must agree with the torch one, zeros included. + + An all-zero 32-element block clamps the E8M0 scale to 0, i.e. 2**-127 -- + subnormal in fp32 and flushed to zero on CDNA. Dividing by it turned the + block into 0/0 == NaN, which then poisoned every output column of the GEMM + that consumed it. + """ + torch.manual_seed(num_rows) + device = "cuda" + x = torch.randn(num_rows, 128, device=device, dtype=torch.bfloat16) + x[0] = 0.0 # all-zero row: every block clamps to scale byte 0 + x[-1, :32] = 0.0 # single zero block in an otherwise normal row + x[-1, 32:64] = 1e-38 # denormal-ish block + + q_triton, s_triton = _mxfp8_e4m3_quantize_triton(x) + q_torch, s_torch = _mxfp8_e4m3_quantize_torch(x) + + assert not torch.isnan(q_triton.float()).any() + torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) + torch.testing.assert_close( + q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0 + ) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1049eb65811..1f6601dc6cfc 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -42,6 +42,438 @@ from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="graph capture coverage" +) +@pytest.mark.parametrize("compress_ratio", [1, 2]) +@pytest.mark.parametrize("use_graph", [False, True]) +@pytest.mark.parametrize( + "lengths", + [(1, 5, 13), (1, 1, 1), (1, 2, 0), (7, 8, 8), (2, 3, 3)], + ids=["prefill", "decode", "empty", "wrap_threshold", "packed_boundary"], +) +def test_v41_fused_save_compress_and_insert( + compress_ratio: int, use_graph: bool, lengths: tuple[int, ...] +): + """Ring states: a group boundary whose previous token is outside the chunk + pools the ring row that chunk left behind, even when this chunk overwrites + that row; only the last ``capacity`` tokens of a chunk are stored. + + Exercise historical group starts, adjacent requests, a chunk longer than + the ring, single-token decode pairs, empty requests, wrap thresholds at + odd and even starts, padded graph rows, and padded physical pages against + a Torch reference with BF16 rounding before RoPE. Replays change both + inputs and history so stale scratch/cache values cannot satisfy the test. + """ + from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + fused_save_compress_norm, + rope_quant_insert, + ) + + torch.manual_seed(42) + device = "cuda" + capacity = 8 + starts = [7, 8, 13] + actual = sum(lengths) + num_tokens = actual + 5 + positions = torch.zeros(num_tokens, dtype=torch.int64, device=device) + req_ids = torch.zeros(actual + 2, dtype=torch.int32, device=device) + state_slots = torch.full((actual + 2,), -1, dtype=torch.int64, device=device) + cache_slots = torch.full_like(state_slots, -1) + query_start_loc = torch.tensor( + [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32 + ).to(device) + ring_blocks = torch.randperm(6, device=device)[:3] + raw = torch.randn(num_tokens, 512 * compress_ratio, device=device) + norm = torch.randn(512, dtype=torch.bfloat16, device=device) + latent = torch.empty(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(32, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + # capacity x 1024 floats rounded up to 576-byte alignment. + state_backing = torch.empty(6, 8208, device=device) + state = state_backing.as_strided((6, capacity, 1024), (8208, 1024, 1)) + cache_block = 256 // compress_ratio + cache_stride = math.ceil(cache_block * 584 / 576) * 576 + cache_backing = torch.empty(3, cache_stride, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((3, cache_block, 584), (cache_stride, 584, 1)) + cursor = 0 + for req, (start, length) in enumerate(zip(starts, lengths)): + pos = torch.arange(start, start + length, device=device) + rows = slice(cursor, cursor + length) + positions[rows] = pos + req_ids[rows] = req + state_slots[rows] = ring_blocks[req] * capacity + pos % capacity + cache_slots[rows] = (2 - req) * cache_block + pos // compress_ratio + cursor += length + + has_ring = compress_ratio == 2 + + def run(): + fused_save_compress_norm( + raw, + positions, + state if has_ring else None, + state_slots, + query_start_loc if has_ring else None, + req_ids if has_ring else None, + norm, + 1e-20, + compress_ratio, + latent, + ) + rope_quant_insert( + latent, positions, cos_sin, cache, cache_slots, compress_ratio + ) + + state_backing.normal_() + run() + graph = None + if use_graph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + for _ in range(3): + raw.normal_() + state_backing.normal_() + history = state.clone() + expected_state = state_backing.clone() + expected_rows = expected_state.as_strided((6, capacity, 1024), (8208, 1024, 1)) + cache_backing.fill_(165) + expected_cache = cache_backing.clone() + latent.fill_(float("nan")) + # Ratio 1 has no ring; ratio 2 stores each chunk's last `capacity` rows. + for t in range(actual if has_ring else 0): + req = req_ids[t].item() + if query_start_loc[req + 1].item() - t > capacity: + continue + slot = state_slots[t].item() + expected_rows[slot // capacity, slot % capacity, :512] = raw[t, :512] + expected_rows[slot // capacity, slot % capacity, 512:] = raw[t, 512:] + + if graph is None: + run() + else: + graph.replay() + + torch.testing.assert_close(state_backing, expected_state, rtol=0, atol=0) + for t in range(num_tokens): + pos = positions[t].item() + if t >= actual or (pos + 1) % compress_ratio: + assert torch.isnan(latent[t]).all() + continue + req = req_ids[t].item() + group_rows = [] + for k in range(compress_ratio - 1, -1, -1): + if t - k >= query_start_loc[req].item(): + kv_score = raw[t - k] + if compress_ratio == 1: + kv_score = torch.cat((kv_score, torch.zeros_like(kv_score))) + group_rows.append(kv_score) + else: + ring = ring_blocks[req].item() + group_rows.append(history[ring, (pos - k) % capacity]) + group = torch.stack(group_rows) + pooled = (group[:, :512] * group[:, 512:].softmax(0)).sum(0) + normed = pooled * torch.rsqrt(pooled.square().mean() + 1e-20) * norm + torch.testing.assert_close(latent[t].float(), normed, rtol=0.004, atol=1e-6) + + slot = cache_slots[t].item() + page, row = divmod(slot, cache_block) + values = expected_cache[page, row * 576 : (row + 1) * 576] + quantized, scales = _ue8m0_reference(latent[t, :448], 64, 448.0) + values[:448] = quantized.view(torch.uint8) + c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2) + rope_input = latent[t, 448:].float() + rotated = ( + torch.stack( + ( + rope_input[0::2] * c - rope_input[1::2] * s, + rope_input[1::2] * c + rope_input[0::2] * s, + ), + dim=-1, + ) + .flatten() + .to(torch.bfloat16) + ) + actual_rope = cache_backing[page, row * 576 + 448 : (row + 1) * 576].view( + torch.bfloat16 + ) + torch.testing.assert_close(actual_rope, rotated, rtol=0.008, atol=1e-6) + values[448:] = actual_rope.view(torch.uint8) + scale_offset = cache_block * 576 + row * 8 + expected_cache[page, scale_offset : scale_offset + 7] = ( + scales.log2() + 127 + ).to(torch.uint8) + expected_cache[page, scale_offset + 7] = 0 + # Includes unwritten rows, page padding, quantized values and scales. + torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +@pytest.mark.parametrize("store_fp8", [False, True]) +def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): + """The FlashInfer plain-row compressed cache gets [448 NoPE | 64 RoPE] rows. + + bf16 caches store the latent verbatim with the RoPE tail rotated; per-tensor + fp8 caches scale the bf16-rounded row by 1/fp8_scale and clamp to e4m3. + Non-boundary positions, negative slots and untouched rows must keep their + prior contents. + """ + from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(3) + device = "cuda" + num_tokens = 9 + cache_block = 4 + positions = torch.arange(5, 5 + num_tokens, dtype=torch.int64, device=device) + cache_slots = torch.randperm(3 * cache_block, device=device)[:num_tokens] + cache_slots[2] = -1 + latent = torch.randn(num_tokens, 512, device=device).to(torch.bfloat16) + angles = torch.randn(32, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + dtype = torch.float8_e4m3fn if store_fp8 else torch.bfloat16 + # Pad the page stride so the kernel must honour strides, not shape. + cache_backing = torch.full((3, (cache_block + 1) * 512), 3.0, device=device) + cache_backing = cache_backing.to(dtype) + cache = cache_backing.as_strided( + (3, cache_block, 512), ((cache_block + 1) * 512, 512, 1) + ) + fp8_scale = torch.tensor([0.5], dtype=torch.float32, device=device) + expected = cache_backing.clone().float() + + rope_quant_insert( + latent, + positions, + cos_sin, + cache, + cache_slots, + compress_ratio, + fp8_scale=fp8_scale if store_fp8 else None, + ) + + for t in range(num_tokens): + pos = positions[t].item() + slot = cache_slots[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2) + rope_input = latent[t, 448:].float() + rotated = torch.stack( + ( + rope_input[0::2] * c - rope_input[1::2] * s, + rope_input[1::2] * c + rope_input[0::2] * s, + ), + dim=-1, + ).flatten() + row = torch.cat((latent[t, :448], rotated.to(torch.bfloat16))).float() + if store_fp8: + row = (row * (1.0 / fp8_scale)).clamp(-448.0, 448.0) + row = row.to(torch.float8_e4m3fn).float() + page, idx = divmod(slot, cache_block) + expected[page, idx * 512 : (idx + 1) * 512] = row + + actual = cache_backing.float() + if store_fp8: + # One e4m3 ulp of slack for FMA-vs-separate rounding in the RoPE tail. + torch.testing.assert_close(actual, expected, rtol=0.13, atol=1e-2) + else: + nope = torch.arange(512, device=device) < 448 + nope = nope.repeat(cache_block + 1) + torch.testing.assert_close(actual[:, nope], expected[:, nope], rtol=0, atol=0) + torch.testing.assert_close( + actual[:, ~nope], expected[:, ~nope], rtol=0.008, atol=1e-6 + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel") +def test_v41_compressor_metadata_maps_tokens_to_their_ring(): + """The ring group's generic slot mapping is disabled (all PAD), so the + builder must map every real token to ``ring_block * capacity + pos % + capacity`` and keep padding tokens at PAD.""" + from unittest.mock import MagicMock + + from vllm.models.deepseek_v4_1.compressor import CompressorMetadataBuilder + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.kv_cache_interface import CircularBufferSpec + + capacity = 8 + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_batched_tokens = 16 + spec = CircularBufferSpec( + block_size=capacity, + num_kv_heads=1, + head_size=1024, + head_size_v=0, + dtype=torch.float32, + ) + device = torch.device("cuda") + builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device) + + # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2 + # tokens at positions 7..8 on ring block 2; three padding tokens. + query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device) + positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device) + block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device) + common = CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc.cpu(), + seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device), + num_reqs=2, + num_actual_tokens=5, + max_query_len=3, + max_seq_len=16, + block_table_tensor=block_table, + slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device), + positions=positions, + ) + metadata = builder.build(0, common) + + expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1] + assert metadata.slot_mapping.tolist() == expected + assert metadata.query_start_loc is query_start_loc + assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage") +@pytest.mark.parametrize( + "use_aux,use_graph", [(False, False), (True, False), (True, True)] +) +@torch.inference_mode() +def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph): + """Both reused-event joins must publish this forward's states and cache rows.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v4_1.attention import ( + DeepseekV4Attention, + DeepseekV4Indexer, + ) + from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor + + torch.manual_seed(43) + raw = torch.randn(19, 1024, device="cuda") + q = torch.zeros(19, 512, dtype=torch.bfloat16, device="cuda") + positions = torch.arange(7, 26, device="cuda") + state_slots = positions.clone() + state_slots[-2:] = -1 + cache_slots = torch.where(state_slots >= 0, positions // 2, -1) + state = torch.empty(4, 8, 1024, device="cuda") + main = torch.empty(1, 128, 584, dtype=torch.uint8, device="cuda") + index = torch.empty(1, 128, 132, dtype=torch.uint8, device="cuda") + caches = (state, main, index) + observed = tuple(torch.empty_like(cache) for cache in caches) + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda")) + metadata = { + "state": SimpleNamespace( + slot_mapping=state_slots, + query_start_loc=torch.tensor([0, 19], dtype=torch.int32, device="cuda"), + token_to_req_indices=torch.zeros(19, dtype=torch.int32, device="cuda"), + ), + "main": SimpleNamespace(slot_mapping=cache_slots), + "index": SimpleNamespace(slot_mapping=cache_slots), + } + context = ForwardContext({}, metadata, {}) + compressor = DeepseekCompressor.__new__(DeepseekCompressor) + torch.nn.Module.__init__(compressor) + compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = ( + 512, + 64, + 2, + ) + compressor.rms_norm_eps = 1e-20 + compressor.norm = SimpleNamespace( + weight=torch.ones(512, dtype=torch.bfloat16, device="cuda") + ) + compressor.state_cache = SimpleNamespace(prefix="state", kv_cache=state) + compressor.k_cache_prefix = "main" + compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=main)} + indexer_weight = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda") + indexer = SimpleNamespace( + owns_k=True, + wk=lambda latent: (torch.nn.functional.linear(latent, indexer_weight), None), + k_norm=SimpleNamespace( + weight=torch.ones(128, dtype=torch.bfloat16, device="cuda"), + variance_epsilon=1e-20, + ), + k_cache=SimpleNamespace(prefix="index", kv_cache=index), + compress_ratio=2, + use_fp4_kv=False, + ) + + def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale): + DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary) + return None, None, None + + def observe(*args): + for output, cache in zip(observed, caches): + output.copy_(cache) + + auxiliary = [torch.cuda.Stream()] if use_aux else None + attention = SimpleNamespace( + compressor=compressor, + indexer=prepare_indexer, + aux_stream_list=auxiliary, + ln_events=[torch.cuda.Event(), torch.cuda.Event()], + rotary_emb=rotary, + indexer_rotary_emb=rotary, + n_local_heads=1, + head_dim=512, + _wq_b_proj=lambda qr, scale: qr.clone(), + _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q, + _sparse_indexer_and_attn=observe, + ) + + def run(): + with override_forward_context(context): + DeepseekV4Attention._prepare_and_attn( + attention, q, q, q, None, raw, q, positions, q + ) + + def reset(): + state.zero_() + main.fill_(165) + index.fill_(165) + + reset() + run() + graph = None + if use_graph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + previous = None + for _ in range(3): + raw.normal_() + reset() + attention.aux_stream_list = None + run() + expected = tuple(output.clone() for output in observed) + reset() + attention.aux_stream_list = auxiliary + if graph is None: + run() + else: + graph.replay() + for actual, reference in zip(observed, expected): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + if previous is not None: + assert all(not torch.equal(a, b) for a, b in zip(expected, previous)) + previous = expected + + def _on_gfx950() -> bool: if not current_platform.is_rocm(): return False diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py new file mode 100644 index 000000000000..f7aa6dd03e79 --- /dev/null +++ b/tests/kernels/test_engram.py @@ -0,0 +1,778 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import bisect +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.deepseek_v4_1.common import engram as engram_ops +from vllm.models.deepseek_v4_1.common.engram import ( + Engram, + NgramHashState, + ParallelEngramEmbedding, +) +from vllm.platforms import current_platform + + +def _reference_engram_post_wkv( + hidden_states, + kv, + q_weight, + k_weight, + token_mask, + eps, + clamp_value, + kv_start, +): + num_tokens, hc_mult, dim = hidden_states.shape + source = torch.arange(num_tokens, device=kv.device) + kv_start + valid = source < kv.shape[0] + local_kv = kv.new_zeros(num_tokens, kv.shape[1]) + local_kv[valid] = kv[source[valid]] + + key = local_kv[:, : hc_mult * dim].float().view(num_tokens, hc_mult, dim) + value = local_kv[:, hc_mult * dim :].float() + h = hidden_states.float() + weight = q_weight.float() * k_weight.float() + rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt( + key.square().mean(-1) + eps + ) + dot = (h * weight * key).sum(-1) * rstd * dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot)) + if token_mask is not None: + local_mask = torch.zeros(num_tokens, dtype=torch.bool, device=kv.device) + local_mask[valid] = token_mask[source[valid]] + gate = gate.masked_fill(~local_mask.unsqueeze(-1), 0) + return (h + gate.unsqueeze(-1) * value.unsqueeze(-2)).to(hidden_states.dtype) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "num_tokens,num_kv_tokens,hc_mult,dim,tp_size,tp_rank,use_mask", + [ + (1, 1, 4, 5120, 1, 0, False), + (3, 5, 4, 96, 2, 0, True), + (3, 5, 2, 257, 2, 1, False), + (3, 5, 4, 96, 2, 1, True), + (1, 1, 4, 96, 4, 3, True), + (0, 0, 4, 96, 4, 0, False), + ], +) +def test_fused_engram_post_wkv_matches_reference( + num_tokens, + num_kv_tokens, + hc_mult, + dim, + tp_size, + tp_rank, + use_mask, + monkeypatch, +): + """Preserve gate math, WKV layout, masking, and SP padding semantics.""" + torch.manual_seed(0) + device = "cuda" + hidden_states = torch.randn( + num_tokens, hc_mult, dim, dtype=torch.bfloat16, device=device + ) + kv = torch.randn( + num_kv_tokens, + (hc_mult + 1) * dim, + dtype=torch.bfloat16, + device=device, + ) + q_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device) + k_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device) + token_mask = None + if use_mask: + token_mask = torch.arange(num_kv_tokens, device=device) % 2 == 0 + eps = 1e-20 + clamp_value = 1e-6 + use_sequence_parallel = tp_size > 1 + shard_size = (num_kv_tokens + tp_size - 1) // tp_size + kv_start = tp_rank * shard_size if use_sequence_parallel else 0 + + expected = _reference_engram_post_wkv( + hidden_states, + kv, + q_weight, + k_weight, + token_mask, + eps, + clamp_value, + kv_start, + ) + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) + + module = Engram.__new__(Engram) + torch.nn.Module.__init__(module) + module.dim = dim + module.hc_mult = hc_mult + module.eps = eps + module.clamp_value = clamp_value + module.use_sequence_parallel = use_sequence_parallel + module.embed_tokens = torch.nn.Identity() + # `forward` reads rows staged by `prepare_embeddings`, so inject kv there. + module.embed_tokens.tp_size = 1 + module.staged_rows = kv.unsqueeze(1) + if use_sequence_parallel: + padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size)) + module.staged_rows = padded[ + tp_rank * shard_size : (tp_rank + 1) * shard_size + ].unsqueeze(1) + module.wkv = torch.nn.Identity() + + def check_local_projection(_, args): + assert args[0].shape[0] == num_tokens + + module.wkv.register_forward_pre_hook(check_local_projection) + module.q_weight = torch.nn.Parameter(q_weight, requires_grad=False) + module.k_weight = torch.nn.Parameter(k_weight, requires_grad=False) + + actual = module(hidden_states, kv.unsqueeze(1), token_mask) + + assert actual.shape == hidden_states.shape + assert actual.dtype == hidden_states.dtype + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +def _hash_state(use_slot_cache: bool) -> NgramHashState: + state = NgramHashState.__new__(NgramHashState) + torch.nn.Module.__init__(state) + state.block_size = 64 + state.use_slot_cache = use_slot_cache + state._cache = None + state._kv_cache_ref = None + state.swa_cache_module = torch.nn.Module() + state.swa_cache_module.kv_cache = torch.empty(0) + return state + + +@pytest.mark.parametrize("num_blocks", [4, 100]) +def test_engram_cache_follows_kv_cache_binding(num_blocks): + """Discard profiling history when the real KV cache replaces its storage.""" + state = _hash_state(use_slot_cache=True) + assert not state.ensure_cache() + + state.swa_cache_module.kv_cache = torch.empty(4, 64, 1) + assert state.ensure_cache() + profiling_cache = state._cache + profiling_cache.fill_(7) + assert state.ensure_cache() + assert state._cache is profiling_cache + assert torch.all(state._cache == 7) + + state.swa_cache_module.kv_cache = torch.empty(num_blocks, 64, 1) + assert state.ensure_cache() + assert state._cache is not profiling_cache + assert state._cache.shape == (num_blocks * 64,) + assert torch.count_nonzero(state._cache) == 0 + + state.swa_cache_module.kv_cache = torch.empty(0) + assert not state.ensure_cache() + assert state._cache is None + + +def test_engram_without_slot_cache_only_gates_on_kv_binding(): + """The V2 runner supplies every lookback, so no slot cache is allocated.""" + state = _hash_state(use_slot_cache=False) + assert not state.ensure_cache() + state.swa_cache_module.kv_cache = torch.empty(4, 64, 1) + assert state.ensure_cache() + assert state._cache is None + + +def _reference_hashes( + ids, + positions, + slots, + starts, + tables, + dead, + cache, + token_map, + multipliers, + primes, + offsets, +): + """Scalar n-gram oracle with persistent, physically addressed token history.""" + for token, slot, is_dead in zip(ids, slots, dead): + if slot >= 0: + cache[slot] = -1 if is_dead else token_map[token] + result = [] + for i, position in enumerate(positions): + req = min(bisect.bisect_right(starts[1:], i), len(tables) - 1) + history, blocked = [], False + for shift in range(4): + lookback = position - shift + p = max(0, min(lookback, len(tables[req]) * 16 - 1)) + source = cache[tables[req][p // 16] * 16 + p % 16] + blocked |= lookback < 0 or source == -1 + history.append(0 if blocked else source) + layers = [] + for layer, row in enumerate(multipliers): + rolling, hashes = history[0] * row[0], [] + for j, head_primes in enumerate(primes[layer], start=1): + rolling ^= history[j] * row[j] + for head, prime in enumerate(head_primes): + col = (j - 1) * len(head_primes) + head + hashes.append(rolling % prime + offsets[layer][col]) + layers.append(hashes) + result.append(layers) + return torch.tensor(result, dtype=torch.int32) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("query_len", [1, 17, 257]) +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("strided_slots", [False, True]) +@pytest.mark.parametrize("large_hashes", [False, True]) +def test_engram_hash_cache_replay(query_len, capture, strided_slots, large_hashes): + """Preserve history, dead tokens and padding as inputs change across replay.""" + num_tokens = 2 * query_len + 3 + blocks = (3 * query_len + 15) // 16 + tables = [list(range(blocks, 2 * blocks)), list(range(blocks))] + starts = [0, query_len, 2 * query_len] + + def tensor(data, dtype=torch.int64): + return torch.tensor(data, dtype=dtype, device="cuda") + + input_ids = tensor([0] * num_tokens) + positions = torch.zeros_like(input_ids) + slots = tensor([-1] * (num_tokens * (2 if strided_slots else 1))) + if strided_slots: + slots = slots[::2] + dead_mask = tensor([False] * num_tokens, torch.bool) + cache = tensor([-1] * (2 * blocks * 16), torch.int32) + token_map = [i * 3091 if large_hashes else i % 7 for i in range(32)] + multipliers = ( + engram_ops.compute_hash_multipliers((1, 14), 4, 99092).tolist() + if large_hashes + else [[3, 5, 7, 9], [11, 13, 15, 17]] + ) + primes = ( + [ + [ + [16000057 + 2 * (layer * 24 + j * 8 + h) for h in range(8)] + for j in range(3) + ] + for layer in range(2) + ] + if large_hashes + else [[[19], [23], [29]]] * 2 + ) + offsets = [] + for layer in primes: + row, offset = [], 0 + for heads in layer: + for prime in heads: + row.append(offset) + offset += prime + offsets.append(row) + state = NgramHashState.__new__(NgramHashState) + torch.nn.Module.__init__(state) + state.token_map = tensor(token_map, torch.int32) + state.multipliers = tensor(multipliers) + state.primes = tensor(primes) + state.offsets = tensor(offsets) + state._cache = cache + state.use_slot_cache = True + state.pad_id = 0 + state.block_size = 16 + lookback = tensor([[-1] * 3] * 2, torch.int32) + args = ( + input_ids, + positions, + tensor(starts, torch.int32), + dead_mask, + lookback, + torch.zeros_like(lookback, dtype=torch.bool), + slots, + tensor(tables, torch.int32), + ) + op = state + graph = None + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + op(*args) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = op(*args) + torch.cuda.current_stream().wait_stream(stream) + cache.fill_(-1) + expected_cache = [-1] * cache.numel() + + for step in range(3): + ids = [(i + 3 * step) % 32 for i in range(num_tokens)] + pos = list(range(step * query_len, (step + 1) * query_len)) * 2 + physical_slots = [ + tables[i // query_len][p // 16] * 16 + p % 16 for i, p in enumerate(pos) + ] + pos += [0] * 3 + physical_slots += [-1] * 3 + # A fully padded replay must leave every cache slot untouched. + if step == 2: + physical_slots = [-1] * num_tokens + dead = [(i + step) % 5 == 0 for i in range(num_tokens)] + input_ids.copy_(tensor(ids)) + positions.copy_(tensor(pos)) + slots.copy_(tensor(physical_slots)) + dead_mask.copy_(tensor(dead, torch.bool)) + if graph is None: + output = op(*args) + else: + graph.replay() + expected = _reference_hashes( + ids, + pos, + physical_slots, + starts, + tables, + dead, + expected_cache, + token_map, + multipliers, + primes, + offsets, + ) + if step != 2: + real = 2 * query_len + torch.testing.assert_close( + output[:real].cpu(), expected[:real], rtol=0, atol=0 + ) + assert cache.cpu().tolist() == expected_cache + + empty_output = state( + input_ids[:0], positions[:0], args[2], dead_mask[:0], *args[4:] + ) + assert empty_output.shape == (0, 2, len(offsets[0])) + assert cache.cpu().tolist() == expected_cache + + +def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False): + """Run the op on a batch of (token_ids, start, window) requests; `window` + is the runner's [depth] lookback for that request (None = all unknown). + Token IDs 14 and 22 stand in for image tokens that break n-grams. + Returns one [len(token_ids), ...] hash tensor per request.""" + token_map, multipliers, primes, offsets, block_size = tables_meta + depth = multipliers.shape[1] - 1 + + def tensor(data, dtype=torch.int64): + return torch.tensor(data, dtype=dtype, device="cuda") + + input_ids, positions, slots, windows, starts = [], [], [], [], [0] + for req, (token_ids, start, window) in enumerate(requests): + pos = list(range(start, start + len(token_ids))) + input_ids += token_ids + positions += pos + slots += [ + int(block_table[req, p // block_size]) * block_size + p % block_size + for p in pos + ] + windows.append(window or [-1] * depth) + starts.append(starts[-1] + len(token_ids)) + window = tensor(windows, torch.int32) + state = _hash_state(use_slot_cache=cache is not None) + state.token_map = token_map + state.multipliers = multipliers + state.primes = primes + state.offsets = offsets + state.pad_id = 0 + state.block_size = block_size + state._cache = cache + input_tensor = tensor(input_ids, torch.int32) + args = ( + input_tensor, + tensor(positions), + tensor(starts, torch.int32), + (input_tensor == 14) | (input_tensor == 22), + window, + (window == 14) | (window == 22), + tensor(slots) if cache is not None else None, + block_table if cache is not None else None, + ) + out = state(*args) + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = state(*args) + graph.replay() + return [out[a:b] for a, b in zip(starts, starts[1:])] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("runner", ["v2_window_only", "v1_prompt_window_plus_cache"]) +@pytest.mark.parametrize("capture", [False, True]) +def test_lookback_window_reproduces_single_instance(runner, capture): + """Decoding on a fresh instance whose prompt KV came from elsewhere (P/D, + offload) must hash like the instance that processed the whole sequence, + with two requests batched and chunks of 1-2 tokens so in-batch, window + and (V1) cache lookbacks mix within one step. + + V2 supplies every lookback from its device token history and needs no + slot cache. V1 supplies prompt positions only and reads generated + positions from the slot cache it fills itself. Without any window the + first decode token hashes stale slots (negative control).""" + block_size = 4 + block_table = torch.tensor( + [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda" + ) + tables_meta = ( + torch.arange(50, dtype=torch.int32, device="cuda"), + torch.tensor([[3, 5, 7, 9]], device="cuda"), + torch.tensor([[[97, 89], [83, 79], [73, 71]]], device="cuda"), + torch.tensor([[0, 97, 186, 269, 348, 421]], device="cuda"), + block_size, + ) + depth = 3 + prompts = [[11, 12, 13, 14, 15, 16], [31, 32, 33, 34, 35]] + decodes = [[21, 22, 23, 24], [41, 42, 43, 44]] + chunk_sizes = [1, 2, 1] + v2 = runner.startswith("v2") + + def fresh_cache(): + return torch.zeros(8 * block_size, dtype=torch.int32, device="cuda") + + reference = _hash_ids( + [(p + d, 0, None) for p, d in zip(prompts, decodes)], + tables_meta, + block_table, + cache=fresh_cache(), + ) + + histories = [list(p) for p in prompts] + cache = None if v2 else fresh_cache() + consumed = 0 + for chunk in chunk_sizes: + batch = [] + for history, prompt, decode in zip(histories, prompts, decodes): + start = len(history) + tokens = decode[consumed : consumed + chunk] + window = [ + history[start - 1 - j] if v2 or start - 1 - j < len(prompt) else -1 + for j in range(depth) + ] + batch.append((tokens, start, window)) + got = _hash_ids(batch, tables_meta, block_table, cache=cache, capture=capture) + for req, (tokens, start, _) in enumerate(batch): + assert torch.equal(got[req], reference[req][start : start + chunk]), ( + f"{runner} request {req} chunk at {start}" + ) + histories[req] += tokens + consumed += chunk + + unseeded = _hash_ids( + [(decodes[0][:1], len(prompts[0]), None)], + tables_meta, + block_table, + cache=fresh_cache(), + ) + assert not torch.equal(unseeded[0][0], reference[0][len(prompts[0])]) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_v2_model_state_gathers_lookback_window(): + """The window is gathered on device from the runner's token history.""" + from vllm.models.deepseek_v4_1.nvidia.model_state import DeepseekV41ModelState + + depth, max_num_reqs, max_model_len = 3, 4, 16 + state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) + state.rope_state = None + state.supports_mm_inputs = False + state.prompt_embeds_state = None + state.lookback_token_ids = torch.full( + (max_num_reqs, depth), -1, dtype=torch.int32, device="cuda" + ) + + all_token_ids = torch.zeros(max_num_reqs, max_model_len, dtype=torch.int32) + all_token_ids[2, :8] = torch.tensor([11, 12, 13, 14, 15, 16, 21, 22]) + all_token_ids[0, :5] = torch.arange(1, 6) + req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace( + gpu=torch.tensor([0, 9, 8, 6], dtype=torch.int32, device="cuda") + ), + all_token_ids=SimpleNamespace(gpu=all_token_ids.cuda()), + ) + # Batch rows -> request state rows: a request two decode steps in, a fresh + # request, and one at the end of a 5-token prompt. + input_batch = SimpleNamespace( + idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda") + ) + window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"] + assert window.cpu().tolist() == [ + [22, 21, 16], + [-1, -1, -1], + [0, 0, 0], + [-1, -1, -1], + ] + # Graph capture runs on the dummy inputs and replays read the same buffer. + dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"] + assert dummy is window and torch.all(dummy == -1) + + +def _reference_lookup(weight, scale_inv, ids, start, end, block=32): + """The torch expression the fused kernel replaces.""" + mask = (ids < start) | (ids >= end) + local = (ids - start).masked_fill(mask, 0).long() + values = torch.nn.functional.embedding(local, weight) + scales = torch.nn.functional.embedding(local, scale_inv) + scales = (scales.to(torch.int32) << 23).view(torch.float32) + values = values.float().unflatten(-1, (-1, block)) + values = (values * scales.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) + return values.masked_fill(mask.unsqueeze(-1), 0) + + +def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): + layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding) + torch.nn.Module.__init__(layer) + layer.dim, layer.block_size, layer.tp_size = dim, block, 1 + layer.n_hash_cols = layer.part_n_hash_cols = 24 + layer.head_start = 0 + layer.cpu_offload = cpu_offload + layer.part_num_embeddings = rows + # A window strictly inside the table, so unowned rows are exercised too. + layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2 + layer._views = layer._view_src = None + layer._num_sms = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + kwargs = ( + {"device": "cpu", "pin_memory": True} if cpu_offload else {"device": "cuda"} + ) + owned = layer.vocab_end_idx - layer.vocab_start_idx + layer.weight = torch.empty(owned, dim, dtype=torch.float8_e4m3fn, **kwargs) + layer.weight_scale_inv = torch.empty( + owned, dim // block, dtype=torch.uint8, **kwargs + ) + torch.manual_seed(0) + layer.weight.copy_((torch.randn(owned, dim) * 4).to(torch.float8_e4m3fn)) + layer.weight_scale_inv.copy_( + torch.randint(120, 134, (owned, dim // block), dtype=torch.uint8) + ) + return layer + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("cpu_offload", [False, True]) +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch): + """Keep complete buckets and reconstruct head order, including TP padding.""" + head_sizes = (17, 19, 23, 29, 31, 37) + num_rows, dim = sum(head_sizes), 64 + torch.manual_seed(0) + weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn) + scales = torch.randint(120, 134, (num_rows + 7, dim // 32), dtype=torch.uint8) + hashes = torch.empty(7, 2, len(head_sizes), dtype=torch.int32, device="cuda") + start = 0 + for head, size in enumerate(head_sizes): + hashes[:, :, head].random_(start, start + size) + hashes[0, :, head] = start + hashes[-1, :, head] = start + size - 1 + start += size + ids = hashes[:, 1] + expected = _reference_lookup(weight.cuda(), scales.cuda(), ids, 0, num_rows) + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + shards = [] + layers = [] + for rank in range(tp_size): + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank + ) + with torch.device("cuda"): + layer = ParallelEngramEmbedding( + num_rows + 7, dim, head_sizes, cpu_offload=cpu_offload + ) + layer.weight.weight_loader(layer.weight, weight) + layer.weight_scale_inv.weight_loader( + layer.weight_scale_inv, scales.view(torch.float8_e8m0fnu) + ) + out = torch.empty( + len(ids), + layer.part_n_hash_cols, + dim, + device="cuda", + dtype=torch.bfloat16, + ) + layer.lookup(ids, out) + shards.append(out) + layers.append(layer) + gathered = torch.cat(shards, dim=1) + torch.testing.assert_close(gathered[:, : len(head_sizes)], expected, rtol=0, atol=0) + assert torch.count_nonzero(gathered[:, len(head_sizes) :]) == 0 + assert sum(layer.part_num_embeddings for layer in layers) == num_rows + + def gather(local, dim): + torch.testing.assert_close(local, shards[0], rtol=0, atol=0) + return torch.cat(shards, dim=dim) + + monkeypatch.setattr(engram_ops, "tensor_model_parallel_all_gather", gather) + torch.testing.assert_close(layers[0](ids), expected, rtol=0, atol=0) + module = Engram.__new__(Engram) + torch.nn.Module.__init__(module) + module.embed_tokens = layers[0] + module.use_sequence_parallel = False + module.staged_rows = torch.empty_like(shards[0]) + module.prepare_embeddings(ids) + torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0) + # Slicing before head reordering must preserve padded heads and empty owners. + module.use_sequence_parallel = True + chunk = (len(ids) + tp_size - 1) // tp_size + padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size)) + for rank in range(tp_size): + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank + ) + torch.testing.assert_close( + module.embed(ids), padded[rank * chunk : (rank + 1) * chunk], rtol=0, atol=0 + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("cpu_offload", [False, True]) +@pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 7, 256]) +def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): + """The fused gather must be bit-exact with the torch dequant path it + replaces, from HBM and from pinned host memory alike, and must contribute + zeros for rows another TP rank owns.""" + layer = _make_embedding(cpu_offload) + cols, rows = 24, layer.part_num_embeddings + ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda") + expected = _reference_lookup( + layer.weight.cuda(), + layer.weight_scale_inv.cuda(), + ids, + layer.vocab_start_idx, + layer.vocab_end_idx, + ) + out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda") + layer.lookup(ids, out, background=background) + assert torch.equal(out, expected) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("cpu_offload", [False, True]) +@pytest.mark.parametrize("capture", ["eager", "full", "breakable"]) +def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture): + """Consume early lookup results after a break, with fresh IDs each replay.""" + _run_engram_prepared_rows(cpu_offload, capture) + + +def _run_engram_prepared_rows( + cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False +): + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + layer = _make_embedding(cpu_offload) + cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64) + layer.n_hash_cols = cols + layer.tp_size = tp_size + layer.part_n_hash_cols = (cols + tp_size - 1) // tp_size + layer.head_start = rank * layer.part_n_hash_cols + engram = Engram.__new__(Engram) + torch.nn.Module.__init__(engram) + engram.embed_tokens = layer + engram.use_sequence_parallel = use_sequence_parallel + engram.staged_rows = torch.empty( + num_tokens, + layer.part_n_hash_cols, + layer.dim, + dtype=torch.bfloat16, + device="cuda", + ) + # Match the non-contiguous per-layer slice of the model hash tensor. + hashes = torch.randint( + 0, + layer.part_num_embeddings, + (num_tokens, 2, cols), + dtype=torch.int32, + device="cuda", + ) + src = hashes[:, 1] + local_tokens = ( + (num_tokens + tp_size - 1) // tp_size if use_sequence_parallel else num_tokens + ) + out = torch.empty( + local_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda" + ) + embed = engram.embed + if capture == "compiled": + embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True) + + def step(cap=None): + engram.prepare_embeddings(src) + if cap is not None: + cap.add_eager(lambda: None) + out.copy_(embed(src)) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + + graph = None + if capture == "full": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + elif capture == "breakable": + graph = BreakableCUDAGraphCapture() + with torch.cuda.stream(warmup), graph: + step(graph) + torch.cuda.current_stream().wait_stream(warmup) + assert graph.num_graphs == 2 + + for _ in range(3): + hashes.random_(0, layer.part_num_embeddings) + if graph is None: + step() + else: + graph.replay() + expected = _reference_lookup( + layer.weight.cuda(), + layer.weight_scale_inv.cuda(), + src, + layer.vocab_start_idx, + layer.vocab_end_idx, + ) + if use_sequence_parallel: + expected = torch.nn.functional.pad( + expected, (0, 0, 0, 0, 0, (-num_tokens) % tp_size) + )[rank * local_tokens : (rank + 1) * local_tokens] + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +def _engram_tp_worker(rank, tp_size, port): + from tests.utils import init_test_distributed_environment + from vllm.distributed import cleanup_dist_env_and_memory + + torch.accelerator.set_device_index(rank) + init_test_distributed_environment(tp_size, 1, rank, str(port), local_rank=rank) + try: + for cpu_offload in (False, True): + for sp in (False, True): + for capture in ("eager", "compiled", "full", "breakable"): + _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp) + finally: + cleanup_dist_env_and_memory() + + +@pytest.mark.distributed(num_gpus=2) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_engram_head_collectives_survive_graph_breaks(): + """All-gather preserves head order and local SP tokens across graph replay.""" + from vllm.utils.network_utils import get_open_port + + if torch.accelerator.device_count() < 2: + pytest.skip("Requires two GPUs") + torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 5b568c938eec..8fd19ffb0b43 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -14,6 +14,7 @@ """ import contextlib +import importlib from unittest import mock import pytest @@ -29,9 +30,10 @@ HEAD_DIM = 128 ROPE_DIM = 64 -N_HEAD = 64 MAX_POS = 4096 +CUTEDSL_MODULE = "vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl" + def quantize_to_mxfp4( x: torch.Tensor, @@ -90,6 +92,7 @@ def _reference( weights: torch.Tensor, softmax_scale: float, head_scale: float, + n_head: int, use_fp4: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: q_rot = q.clone() @@ -105,7 +108,7 @@ def _reference( ) if use_fp4: - q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM)) + q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, n_head, HEAD_DIM)) # Pack 4 ue8m0 bytes into 1 int32 q_scale = ue8m0.view(torch.int32).squeeze(-1) # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar) @@ -118,8 +121,8 @@ def _reference( HEAD_DIM, use_ue8m0=True, ) - q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM) - q_scale = q_scale.view(-1, N_HEAD) + q_fp8 = q_fp8.view(-1, n_head, HEAD_DIM) + q_scale = q_scale.view(-1, n_head) weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale return q_fp8, weights_out @@ -144,9 +147,10 @@ def _reference( ], ) @pytest.mark.parametrize("use_cutedsl", [False, True]) +@pytest.mark.parametrize("n_head", [32, 64]) @torch.inference_mode() def test_fused_indexer_q_rope_quant_matches_unfused( - num_tokens, cache_dtype, use_fp4, use_cutedsl + num_tokens, cache_dtype, use_fp4, use_cutedsl, n_head ): if use_cutedsl and not has_cutedsl(): pytest.skip("cutedsl (cutlass) not installed") @@ -154,17 +158,29 @@ def test_fused_indexer_q_rope_quant_matches_unfused( device = "cuda" torch.manual_seed(0) - q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device) + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) positions = torch.randint( 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device ) cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device) - weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device) + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) softmax_scale = HEAD_DIM**-0.5 - head_scale = N_HEAD**-0.5 + # head_scale must be an exact power of two for bit-exactness to be a fair + # expectation: the triton kernel multiplies by softmax_scale and head_scale + # separately (two roundings) while the cutedsl kernels fold them into one + # fp32 scalar first (one rounding). Those agree only when one factor is a + # pow2. n_head**-0.5 is a pow2 at 64 heads but not at 32. + head_scale = 0.125 q_quant_ref, weights_ref = _reference( - positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4 + positions, + q, + cos_sin_cache, + weights, + softmax_scale, + head_scale, + n_head, + use_fp4, ) # use_cutedsl=False: force the triton path even when cutedsl is installed # by patching the dispatcher's has_cutedsl() binding to return False. @@ -210,3 +226,337 @@ def test_fused_indexer_q_rope_quant_matches_unfused( f"weights mismatch: max abs diff " f"{(weights_ref - weights_fused).abs().max().item()}" ) + + +@pytest.mark.skipif(not has_cutedsl(), reason="cutedsl (cutlass) not installed") +@pytest.mark.parametrize( + "use_fp4", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", + ), + ), + ], +) +@pytest.mark.parametrize("n_head", [32, 64]) +@torch.inference_mode() +def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head): + """The cutedsl kernels must not store past ``num_tokens``. + + The launcher rounds ``num_tokens * threads_per_token`` up to whole 128-thread + blocks, so the trailing block is only partially in range unless + ``threads_per_token`` is itself a multiple of the block size. With 32 heads + and coarsen=4 (picked for >= 512 tokens) a token needs 64 threads, so an odd + token count leaves half of the last block addressing token ``num_tokens``. + Every output gets one extra sentinel row that the kernel must leave alone. + """ + mod = importlib.import_module(CUTEDSL_MODULE) + device = "cuda" + torch.manual_seed(0) + num_tokens = 935 + + sentinel_u8, sentinel_f32 = 0x55, -12345.0 + # positions/q are padded too, so an out-of-bounds *load* stays addressable + # and the assertions below are about stores only. + positions = torch.zeros(num_tokens + 1, dtype=torch.int64, device=device) + q = torch.randn( + num_tokens + 1, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) + weights = torch.randn(num_tokens + 1, n_head, dtype=torch.bfloat16, device=device) + weights_out = torch.full( + (num_tokens + 1, n_head), sentinel_f32, dtype=torch.float32, device=device + ) + outputs = {"weights_out": weights_out} + + if use_fp4: + outputs["q_packed"] = torch.full( + (num_tokens + 1, n_head, HEAD_DIM // 2), + sentinel_u8, + dtype=torch.uint8, + device=device, + ) + outputs["q_scale"] = torch.full( + (num_tokens + 1, n_head, HEAD_DIM // 32), + sentinel_u8, + dtype=torch.uint8, + device=device, + ) + mod.fused_indexer_q_rope_quant_mxfp4_cutedsl( + positions[:num_tokens], + q[:num_tokens], + cos_sin_cache, + weights[:num_tokens], + 1.0, + 1.0, + outputs["q_packed"][:num_tokens], + outputs["q_scale"][:num_tokens], + weights_out[:num_tokens], + ) + else: + outputs["q_fp8"] = torch.full( + (num_tokens + 1, n_head, HEAD_DIM), + sentinel_u8, + dtype=torch.uint8, + device=device, + ) + mod.fused_indexer_q_rope_quant_fp8_cutedsl( + positions[:num_tokens], + q[:num_tokens], + cos_sin_cache, + weights[:num_tokens], + 1.0, + 1.0, + outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn), + weights_out[:num_tokens], + ) + torch.accelerator.synchronize() + + clobbered = {} + for name, tensor in outputs.items(): + sentinel = sentinel_f32 if tensor.dtype == torch.float32 else sentinel_u8 + n_changed = (tensor[num_tokens] != sentinel).sum().item() + if n_changed: + clobbered[name] = n_changed + assert not clobbered, ( + f"kernel wrote past token {num_tokens - 1} into the sentinel row: {clobbered}" + ) + + +def _indexer_k_reference( + k_pre, positions, cos_sin_cache, rms_norm_weight, compress_ratio, use_fp4 +): + k = k_pre.float() + k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-20) + k = (k * rms_norm_weight.float()).to(torch.bfloat16).float() + group_positions = positions // compress_ratio * compress_ratio + cos, sin = cos_sin_cache[group_positions].float().chunk(2, dim=-1) + even, odd = k[:, 64::2], k[:, 65::2] + rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), dim=-1) + k = torch.cat((k[:, :64], rotated.flatten(1)), dim=-1).to(torch.bfloat16) + if use_fp4: + return quantize_to_mxfp4(k) + exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448 + scale = exponent.log2().ceil().exp2() + values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + return values.view(torch.uint8), scale.view(torch.uint8) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel") +@pytest.mark.parametrize("num_tokens", [1, 17, 257]) +@pytest.mark.parametrize("compress_ratio", [1, 2]) +@pytest.mark.parametrize("use_fp4", [False, True]) +@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16]) +@torch.inference_mode() +def test_indexer_k_inserts_only_valid_groups_into_padded_pages( + num_tokens, compress_ratio, use_fp4, cache_dtype +): + """Insert group ends and preserve skipped slots, graph rows, and page padding.""" + from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + indexer_k_norm_rope_store, + ) + + if use_fp4 and not current_platform.is_device_capability_family(100): + pytest.skip("MXFP4 indexer cache requires an SM100-family GPU") + torch.manual_seed(17) + device = "cuda" + k_pre = torch.randn(num_tokens + 3, 128, device=device, dtype=torch.bfloat16) + norm_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + positions = torch.arange(7, num_tokens + 10, device=device) + slots = torch.randperm(2 * num_tokens, device=device)[:num_tokens].clone() + slots[5::7] = -1 + valid = (slots >= 0) & ((positions[:num_tokens] + 1) % compress_ratio == 0) + k_pre[3::7] = 0 + k_pre[8::11] *= 1e-10 + k_pre[:num_tokens][~valid] = float("nan") + k_pre[num_tokens:] = float("nan") + positions[num_tokens:] = MAX_POS + 10 + angles = torch.randn(MAX_POS, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(cache_dtype) + block_size = 16 + value_bytes = 64 if use_fp4 else 128 + page_bytes = block_size * (value_bytes + 4) + num_blocks = (2 * num_tokens + block_size - 1) // block_size + 1 + storage = torch.full( + (num_blocks, page_bytes + 128), 0xA5, device=device, dtype=torch.uint8 + ) + cache = storage.as_strided( + (num_blocks, block_size, value_bytes + 4), + (storage.stride(0), value_bytes + 4, 1), + ) + expected = storage.clone() + values, scales = _indexer_k_reference( + k_pre[:num_tokens][valid], + positions[:num_tokens][valid], + cos_sin, + norm_weight, + compress_ratio, + use_fp4, + ) + block_ids = slots[valid] // block_size + block_offsets = slots[valid] % block_size + value_offsets = block_offsets[:, None] * value_bytes + torch.arange( + value_bytes, device=device + ) + scale_offsets = block_size * value_bytes + block_offsets[:, None] * 4 + scale_offsets = scale_offsets + torch.arange(4, device=device) + expected[block_ids[:, None], value_offsets] = values + expected[block_ids[:, None], scale_offsets] = scales + + indexer_k_norm_rope_store( + k_pre, + positions, + cos_sin, + norm_weight, + 1e-20, + cache, + slots, + compress_ratio, + use_fp4, + ) + # Fused RoPE can change zero signs without changing their numerical values. + actual_values = storage[block_ids[:, None], value_offsets] + if use_fp4: + for shift in (0, 4): + zero_mask = (actual_values >> shift) & 7 == 0 + actual_values &= ~(zero_mask.to(torch.uint8) << (shift + 3)) + zero_mask = (values >> shift) & 7 == 0 + values &= ~(zero_mask.to(torch.uint8) << (shift + 3)) + else: + actual_values[actual_values == 128] = 0 + values[values == 128] = 0 + storage[block_ids[:, None], value_offsets] = actual_values + expected[block_ids[:, None], value_offsets] = values + torch.testing.assert_close(storage, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel") +@pytest.mark.parametrize("use_fp4", [False, True]) +@torch.inference_mode() +def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4): + """Replay must consume updated keys and slot mapping through the captured API.""" + from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + indexer_k_norm_rope_store, + ) + + if use_fp4 and not current_platform.is_device_capability_family(100): + pytest.skip("MXFP4 indexer cache requires an SM100-family GPU") + torch.manual_seed(18) + k_pre = torch.randn(19, 128, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.ones(128, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(19, device="cuda") + cos_sin = torch.zeros(32, 64, device="cuda") + cos_sin[:, :32] = 1 + slots = torch.arange(17, device="cuda") + cache = torch.full( + (2, 16, (64 if use_fp4 else 128) + 4), + 0xA5, + device="cuda", + dtype=torch.uint8, + ) + + def run(): + indexer_k_norm_rope_store( + k_pre, + positions, + cos_sin, + norm_weight, + 1e-20, + cache, + slots, + 2, + use_fp4, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + k_pre.neg_() + slots.add_(4) + cache.fill_(0xA5) + graph.replay() + replayed = cache.clone() + cache.fill_(0xA5) + run() + torch.testing.assert_close(cache, replayed, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm cache layout") +@pytest.mark.parametrize("block_size", [16, 64, 128]) +@pytest.mark.parametrize("compress_ratio", [1, 2]) +@torch.inference_mode() +def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_ratio): + """ROCm reads the indexer K cache back with the 16x16-tiled value layout. + + ``cp_gather_indexer_k_quant_cache_triton`` (prefill) and aiter's + ``Preshuffle=True`` paged kernel (decode) both pick that layout whenever + ``block_size > 1``, so a row-major store would feed the indexer permuted + key bytes and randomize its top-k. Assert the write/read pair is exact. + """ + from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + indexer_k_norm_rope_store, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, + ) + + torch.manual_seed(0) + device = "cuda" + num_tokens = 3 * block_size + 5 + num_blocks = num_tokens // block_size + 2 + k_pre = torch.randn(num_tokens, HEAD_DIM, device=device, dtype=torch.bfloat16) + norm_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16) + positions = torch.arange(num_tokens, device=device) + angles = torch.randn(MAX_POS, ROPE_DIM // 2, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).float() + cache = torch.zeros( + num_blocks, block_size, HEAD_DIM + 4, device=device, dtype=torch.uint8 + ) + slots = torch.arange(num_tokens, device=device) + + indexer_k_norm_rope_store( + k_pre, + positions, + cos_sin, + norm_weight, + 1e-20, + cache, + slots, + compress_ratio, + False, + ) + + k_fp8 = torch.zeros(num_tokens, HEAD_DIM, device=device, dtype=torch.float8_e4m3fn) + k_scale = torch.zeros(num_tokens, 4, device=device, dtype=torch.uint8) + cp_gather_indexer_k_quant_cache_triton( + cache, + k_fp8, + k_scale, + torch.arange(num_blocks, device=device, dtype=torch.int32).view(1, -1), + torch.tensor([0, num_tokens], device=device, dtype=torch.int32), + torch.zeros(num_tokens, device=device, dtype=torch.int32), + ) + + expected_values, expected_scale = _indexer_k_reference( + k_pre, positions, cos_sin, norm_weight, compress_ratio, False + ) + emitted = (positions + 1) % compress_ratio == 0 + gathered = k_fp8.view(torch.uint8) + gathered = torch.where(gathered == 128, gathered.new_zeros(()), gathered) + expected_values = torch.where( + expected_values == 128, expected_values.new_zeros(()), expected_values + ) + torch.testing.assert_close( + gathered[emitted], expected_values[emitted], rtol=0, atol=0 + ) + torch.testing.assert_close( + k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 + ) diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index 5b71e59baa9e..d27144a38e9b 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -159,13 +159,12 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8 Args: - x_f32: [..., quant_group_size] float32 — one or more 128-element blocks. + x_f32: [..., quant_group_size] float32 quantization blocks. Returns: x_fp8: same shape, float8_e4m3fn scales: [...] float32, one scale per block """ - assert x_f32.shape[-1] == QUANT_GROUP_SIZE absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS) scale_raw = absmax * (1.0 / FP8_MAX) scale = torch.exp2(torch.ceil(torch.log2(scale_raw))) @@ -253,6 +252,56 @@ def reference_inv_rope_fp8_quant( # ========================================================================= +@pytest.mark.parametrize("num_tokens", [1, 7, 128, 4096]) +@pytest.mark.parametrize("quant_group_size", [32, 128]) +@torch.inference_mode() +def test_packed_quantization_preserves_rotary_dimensions(num_tokens, quant_group_size): + """RoPE may span multiple scale blocks; pack every scale and padding row.""" + from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant as quantize, + ) + + torch.manual_seed(42) + groups, heads_per_group = 2, 2 + x = torch.randn( + num_tokens, + groups * heads_per_group, + HEAD_DIM, + device="cuda", + dtype=torch.bfloat16, + ) + cache = make_cos_sin_cache(4096) + positions = torch.randint(4096, (num_tokens,), device="cuda") + q, packed_scales = quantize( + x, + positions, + cache, + groups, + heads_per_group, + quant_group_size=quant_group_size, + tma_aligned_scales=True, + ) + scale_bytes = torch.stack( + [(packed_scales >> (8 * i)) & 0xFF for i in range(4)], dim=-1 + ).flatten(-2) + scales = torch.exp2(scale_bytes.float() - 127) + rotated = reference_inv_rope(x.float(), positions, cache) + ref_q, ref_scales = _ref_ue8m0_quant_block( + rotated.reshape(num_tokens, groups, -1, quant_group_size) + ) + torch.testing.assert_close(scales, ref_scales, rtol=0, atol=0) + actual = q.float().unflatten(-1, (-1, quant_group_size)) * scales.unsqueeze(-1) + expected = ref_q.float() * ref_scales.unsqueeze(-1) + assert (actual - expected).norm() / expected.norm() < 1e-3 + + aligned_tokens = (num_tokens + 3) // 4 * 4 + full_scales = packed_scales.transpose(0, 1).as_strided( + (groups, aligned_tokens, packed_scales.shape[-1]), + (packed_scales.stride(1), 1, packed_scales.stride(2)), + ) + assert torch.count_nonzero(full_scales[:, num_tokens:]) == 0 + + @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128]) @pytest.mark.parametrize( "num_heads,n_groups", @@ -260,8 +309,9 @@ def reference_inv_rope_fp8_quant( ids=["H64_G8", "H32_G4", "H128_G8"], ) @pytest.mark.parametrize("seed", [0, 42]) +@pytest.mark.parametrize("quantize", [True, False], ids=["fp8", "bf16"]) @torch.inference_mode() -def test_correctness(num_tokens, num_heads, n_groups, seed): +def test_correctness(num_tokens, num_heads, n_groups, seed, quantize): """Compare fused kernel against reference for FP8 values and scales.""" torch.manual_seed(seed) @@ -280,6 +330,26 @@ def test_correctness(num_tokens, num_heads, n_groups, seed): max_pos, ROPE_DIM, dtype=torch.float32, device=device ) + if not quantize: + ref_output = reference_inv_rope(o, positions, cos_sin_cache).view( + num_tokens, n_groups, -1 + ) + fused_output, fused_scale = fused_inv_rope_fp8_quant( + o.clone(), + positions, + cos_sin_cache, + n_groups, + heads_per_group, + quantize=False, + ) + d = heads_per_group * HEAD_DIM + assert fused_output.shape == (num_tokens, n_groups, d) + assert fused_scale.numel() == 0 + assert fused_output.dtype == torch.bfloat16 + assert fused_output.stride() == (d, num_tokens * d, 1) + torch.testing.assert_close(fused_output, ref_output, rtol=1e-2, atol=1e-2) + return + # Reference ref_fp8, ref_scale = reference_inv_rope_fp8_quant( o.clone(), @@ -726,6 +796,7 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups): This catches stride/layout bugs that only manifest when the einsum kernel actually consumes the quantized activations. """ + from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj from vllm.utils.deep_gemm import ( fp8_einsum, is_deep_gemm_supported, @@ -793,23 +864,23 @@ def ceil_div(a: int, b: int) -> int: ) # -- FUSED path -- - fused_fp8, fused_scale = fused_inv_rope_fp8_quant( + wo_a = torch.nn.Module() + wo_a.register_buffer("weight", w_fp8) + wo_a.register_buffer("weight_scale", w_scale_t) + z_fused = deep_gemm_fp8_o_proj( o.clone(), positions, cos_sin_cache, - n_groups, - heads_per_group, - ) - z_fused = torch.empty( - num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16 - ) - fp8_einsum( - "bhr,hdr->bhd", - (fused_fp8, fused_scale), - (w_fp8, w_scale_t), - z_fused, - recipe=recipe, - ) + wo_a, + torch.nn.Identity(), + n_groups=n_groups, + heads_per_group=heads_per_group, + nope_dim=NOPE_DIM, + rope_dim=ROPE_DIM, + o_lora_rank=o_lora_rank, + einsum_recipe=recipe, + tma_aligned_scales=True, + ).view_as(z_ref) # -- Checks -- # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index a146c1a5c973..61e0449cae19 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -12,7 +12,10 @@ from vllm.model_executor.kernels.mhc.tilelang import ( _tilelang_hc_prenorm_gemm, _torch_hc_prenorm_gemm, + mhc_pre_delayed_tilelang, ) +from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch +from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton from vllm.model_executor.layers.mhc import ( HAS_AITER_MHC, HAS_AITER_MHC_FUSED, @@ -26,12 +29,293 @@ DeepseekV4DecoderLayer, DeepseekV4Model, ) +from vllm.models.deepseek_v4_1.nvidia.model import ( + DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed DEVICE = current_platform.device_type +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "num_tokens,hc_mult,hidden_size", + [ + (0, 4, 5120), + (1, 4, 5120), + (33, 4, 5137), + (128, 4, 5120), + (1024, 4, 5120), + (2, 1, 513), + (7, 8, 5137), + (65537, 4, 65), + ], +) +@pytest.mark.parametrize("strided", [False, True]) +def test_hc_collapse_preserves_weighted_residual_sum( + num_tokens, hc_mult, hidden_size, strided +): + """Handle tile tails and strided streams without changing FP32 mixing.""" + set_random_seed(0) + x = torch.randn( + num_tokens, + hc_mult, + hidden_size * (2 if strided else 1), + dtype=torch.bfloat16, + device=DEVICE, + ) + pre = torch.rand(num_tokens, hc_mult, device=DEVICE) + if strided: + x = x[..., ::2] + pre = pre.t().contiguous().t() + if num_tokens: + pre[0].zero_() + pre[0, -1] = 1 + if num_tokens > 1 and hc_mult == 4: + pre[1].fill_(1) + x[1, 1] = -x[1, 0] + x[1, 3] = -x[1, 2] + expected = (pre.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype) + actual = hc_collapse_triton(x, pre) + assert actual.is_contiguous() + # BF16 rounding only: ptxas <13.1 contracts the FP32 mix into FFMA despite + # enable_fp_fusion=False, so compare within a ULP and pin the exact cases below. + torch.testing.assert_close(actual, expected, atol=1.6e-2, rtol=1e-2) + if num_tokens: + torch.testing.assert_close(actual[0], x[0, -1], atol=0, rtol=0) + if num_tokens > 1 and hc_mult == 4: + torch.testing.assert_close( + actual[1], torch.zeros_like(actual[1]), atol=0, rtol=0 + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_hc_collapse_custom_op_supports_compile(): + x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + pre = torch.rand(2, 4, device=DEVICE) + torch.library.opcheck(torch.ops.vllm.hc_collapse_triton.default, (x, pre)) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("num_tokens", [1, 7]) +def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): + """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm.""" + from vllm.models.deepseek_v4_1.nvidia import dspark + + set_random_seed(0) + hidden_size, hc_mult = 5120, 4 + streams = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + mixes = [torch.rand(num_tokens, hc_mult, device=DEVICE) for _ in range(2)] + carried_mixes = [] + + def make_layer(mix): + def forward(hidden, positions, ids, pre, post, res, residual): + carried_mixes.append(pre) + return hidden, None, None, None, mix + + return forward + + monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) + draft = SimpleNamespace( + use_sequence_parallel=False, + hc_mult=hc_mult, + layers=[make_layer(mix) for mix in mixes], + ) + actual = dspark.DSparkDeepseekV4Model.forward( + draft, + input_ids=torch.zeros(num_tokens, dtype=torch.long, device=DEVICE), + positions=torch.arange(num_tokens, device=DEVICE), + inputs_embeds=torch.zeros( + num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE + ), + ) + expected = (streams.float() * mixes[-1].unsqueeze(-1)).sum(dim=1) + assert carried_mixes[0] is None + assert carried_mixes[1] is mixes[0] + torch.testing.assert_close( + actual, expected.to(streams.dtype), atol=1.6e-2, rtol=1e-2 + ) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize( + "num_tokens,sinkhorn_repeat,norm_eps", + [ + (0, 20, 1e-6), + (1, 1, 1e-6), + (1, 20, 1e-6), + (33, 20, 1e-20), + (128, 20, 1e-6), + (1024, 20, 1e-6), + ], +) +@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"]) +@pytest.mark.parametrize("use_deep_gemm", [False, True]) +@pytest.mark.parametrize("fused_norm", [False, True]) +def test_deepseek_v41_mhc_pre_delayed( + num_tokens, + sinkhorn_repeat, + norm_eps, + entry, + use_deep_gemm, + fused_norm, + monkeypatch, +): + """Collapse with the carried pre-mix and save the new one for the next sublayer.""" + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if use_deep_gemm and not is_deep_gemm_supported(): + pytest.skip("DeepGEMM not available") + monkeypatch.setattr( + "vllm.utils.deep_gemm.is_deep_gemm_supported", lambda: use_deep_gemm + ) + set_random_seed(0) + hc_mult, hidden_size = 4, 5120 + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + if num_tokens > 1: + residual[0].zero_() + fn = torch.randn(24, hc_mult, hidden_size, device=DEVICE) * 0.02 + x = None + pre_mix = None + if entry == "broadcast": + x = residual[:, 0].contiguous() + residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + fn = fn.sum(1) + else: + fn = fn.flatten(1) + if entry == "carried": + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + if num_tokens > 1: + # A constant collapse just above 1 rounds to BF16 1 before RMSNorm. + residual[1].fill_(1) + pre_mix[1].zero_() + pre_mix[1, 0] = 1.003 + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) + base = torch.randn(24, device=DEVICE) + args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, sinkhorn_repeat) + expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix, x=x) + norm_kwargs = {} + if fused_norm: + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE) + weight.uniform_(0.5, 1.5) + norm_kwargs = dict(norm_weight=weight, norm_eps=norm_eps) + # Compare to the original collapse + CUDA RMSNorm, including BF16 rounding. + unfused = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x) + normalized = torch.empty_like(unfused[2]) + if num_tokens: + torch.ops._C.rms_norm(normalized, unfused[2], weight, norm_eps) + expected = (*unfused[:2], normalized, unfused[3]) + actual = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x, **norm_kwargs) + if fused_norm and num_tokens: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = mhc_pre_delayed_tilelang( + *args, pre_mix=pre_mix, x=x, **norm_kwargs + ) + graph.replay() + for eager, replayed in zip(actual, captured, strict=True): + torch.testing.assert_close(eager, replayed, atol=0, rtol=0) + if fused_norm and entry == "carried" and num_tokens > 1: + torch.testing.assert_close(actual[2][1], expected[2][1], atol=0, rtol=0) + # DeepGEMM uses TF32 for the projection; the TileLang fallback uses FP32. + atol, rtol = (1e-3, 1e-3) if use_deep_gemm else (1e-5, 1e-4) + for i in (0, 1, 3): + torch.testing.assert_close(actual[i], expected[i], atol=atol, rtol=rtol) + if pre_mix is None and not fused_norm: + torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0) + else: + torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) +def test_deepseek_v41_decoder_mixes_match_torch( + entry, monkeypatch, default_vllm_config +): + """Preserve carried pre-mixes and Engram ordering at every decoder entry.""" + set_random_seed(0) + decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) + nn.Module.__init__(decoder) + decoder.hc_mult = 4 + decoder.hc_sinkhorn_iters = 20 + decoder.hc_eps = 1e-6 + decoder.rms_norm_eps = 1e-20 + decoder.hc_post_alpha = 2.0 + decoder.use_sequence_parallel = False + decoder.engram = None + from vllm.model_executor.layers.layernorm import RMSNorm + + decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to( + device=DEVICE, dtype=torch.bfloat16 + ) + decoder.attn = lambda positions, x, _: x * 0.5 + decoder.ffn = lambda x, input_ids: x * 0.25 + with torch.device(DEVICE): + decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_attn_fn_broadcast = decoder.hc_attn_fn.view(24, 4, 5120).sum(1) + decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) + decoder.hc_attn_base = torch.randn(24) + decoder.hc_ffn_base = torch.randn(24) + x = torch.randn(3, 5120, dtype=torch.bfloat16) + positions = torch.arange(3) + kwargs = {} + if entry != "broadcast": + kwargs["pre_mix"] = torch.rand(3, 4) + if entry == "pipeline": + x = torch.randn(3, 4, 5120, dtype=torch.bfloat16) + if entry in ("residual", "engram"): + kwargs.update( + residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.rand(3, 4, 1), + res_mix=torch.rand(3, 4, 4), + ) + if entry == "engram": + + class FakeEngram(nn.Module): + layer_hash_index = 0 + + def forward(self, residual, hashes, mask): + return residual + 0.125 + + decoder.engram = FakeEngram() + kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32) + + actual = decoder(x, positions, None, **kwargs) + + def reference(*args, norm_weight, norm_eps, **kwargs): + post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) + return post, res, decoder.attn_norm(collapsed), pre + + monkeypatch.setattr( + "vllm.models.deepseek_v4_1.nvidia.model.mhc_pre_delayed_tilelang", + reference, + ) + expected = decoder(x, positions, None, **kwargs) + for result, ref in zip(actual, expected, strict=True): + torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("carried", [False, True]) +def test_mhc_pre_delayed_custom_op_supports_compile(carried): + set_random_seed(0) + x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + fn = torch.randn(24, 20480, device=DEVICE) * 0.02 + pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None + scale = torch.ones(3, device=DEVICE) + base = torch.zeros(24, device=DEVICE) + torch.library.opcheck( + torch.ops.vllm.mhc_pre_delayed_tilelang.default, + (x, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20, pre_mix), + ) + + def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: x = x.softmax(-1) + eps x = x / (x.sum(-2, keepdim=True) + eps) @@ -203,7 +487,7 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size): _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult) torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4) - torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4) + torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6) @pytest.mark.skipif( diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 8bd97a3ab47f..de6518ebec79 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -290,3 +290,113 @@ def record_forward_impl(*args): expected_args[1:], ) ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_select_candidate_blocks_tolerates_empty_rows(): + """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest + block pin must not index -1 for them (device-side assert); they select + no candidate blocks and real rows still pin their newest block.""" + block_size, topk_blocks = 8, 3 + logits = torch.zeros(3, 64, device="cuda") + logits[0, 3] = 5.0 # block 0 scores highest for row 0 + logits[2, 9] = 5.0 # block 1 scores highest for row 2 + row_ks = torch.zeros(3, dtype=torch.int64, device="cuda") + row_ke = torch.tensor([40, 0, 17], device="cuda") + out = torch.empty(3, topk_blocks, dtype=torch.int32, device="cuda") + + sparse_indexer._select_candidate_blocks( + logits, row_ks, row_ke, topk_blocks, block_size, out + ) + + assert out[1].tolist() == [-1, -1, -1] + assert out[0, 0].item() == 4 and 0 in out[0].tolist() # newest block pinned + assert out[2, 0].item() == 2 and 1 in out[2].tolist() + assert (out[0] >= 0).all() and (out[2, :2] >= 0).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "width,block_size,k,decode", + [ + (73, 8, 16, False), + (97, 3, 7, False), + (32768, 8, 2048, False), + (32768, 8, 2048, True), + ], +) +def test_candidate_kernels_preserve_packed_bounds_and_padding( + width, block_size, k, decode +): + """Preserve top-k ties, newest blocks, empty rows and candidate clamping.""" + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + select_candidate_blocks, + ) + + torch.manual_seed(42) + rows = 6 + logits = torch.randn(rows, width * 2, device="cuda")[:, ::2] + logits[0, :16] = 0 + logits[2, 5] = float("nan") + starts = torch.tensor([0, 5, 3, 0, 7, 1], device="cuda", dtype=torch.int32) + ends = torch.tensor([0, width, width - 1, 1, 11, width], device="cuda") + repeat = 1 + if decode: + starts = None + ends = torch.tensor([0, 1, width], device="cuda", dtype=torch.int32) + repeat = 2 + ks = torch.zeros(rows, device="cuda", dtype=torch.int64) if decode else starts + ke = ends.repeat_interleave(repeat) + cols = torch.arange(width, device="cuda") + valid = (cols >= ks[:, None]) & (cols < ke[:, None]) + scores = logits.masked_fill(~valid, -torch.inf) + blocks = ((cols - ks[:, None]) // block_size).clamp(min=0).long() + nblocks = (width + block_size - 1) // block_size + reduced = logits.new_full((rows, nblocks), -torch.inf) + reduced.scatter_reduce_(1, blocks, scores, reduce="amax", include_self=True) + lengths = ke - ks + last = ((lengths - 1) // block_size).clamp(min=0).long() + reduced.scatter_( + 1, last[:, None], torch.where(lengths > 0, torch.inf, -torch.inf)[:, None] + ) + top = reduced.topk(min(k, nblocks), dim=-1) + expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32) + expected[:, : top.indices.shape[1]] = torch.where( + top.values > -torch.inf, top.indices, -1 + ).int() + actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2] + select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + candidates = actual.clone() + candidates[:, 0] = nblocks + 10 + candidates[:, 1] = 0 + candidates[:, 2] = 0 + positions = ks[:, None, None] + candidates.long()[:, :, None] * block_size + positions = positions + torch.arange(block_size, device="cuda") + keep = torch.zeros(rows, width, device="cuda", dtype=torch.int8) + keep.scatter_reduce_( + 1, + positions.clamp(0, width - 1).reshape(rows, -1), + (candidates >= 0)[:, :, None] + .expand(-1, -1, block_size) + .reshape(rows, -1) + .to(torch.int8), + reduce="amax", + include_self=True, + ) + reference = logits.masked_fill((keep == 0) | ~valid, -torch.inf) + apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat) + torch.testing.assert_close(logits, reference, rtol=0, atol=0, equal_nan=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat) + for block in (1, -1): + candidates.fill_(-1) + candidates[:, 0] = block + logits.fill_(3.0) + graph.replay() + keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) + reference = torch.where(keep, 3.0, -torch.inf) + torch.testing.assert_close(logits, reference, rtol=0, atol=0) diff --git a/tests/models/multimodal/generation/test_maverick.py b/tests/models/multimodal/generation/test_maverick.py index ff6e523e5b25..d9afd93faaf3 100644 --- a/tests/models/multimodal/generation/test_maverick.py +++ b/tests/models/multimodal/generation/test_maverick.py @@ -616,6 +616,8 @@ def test_dummy_maverick( ) -> None: # Disable multiprocessing allows us to access model executor from LLM engine monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + # The forked test process may inherit CUDA state from test collection. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") model_path = create_reduced_maverick_model( original_model_name=original_model_name, diff --git a/tests/models/registry.py b/tests/models/registry.py index a474a5a0fb72..874fa9c11ca5 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -858,6 +858,9 @@ def check_available_online( "DeepseekV4ForConditionalGeneration": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp", ), + "DeepseekV41ForCausalLM": _HfExamplesInfo( + "deepseek-ai/DeepSeek-V4.1-Flash", + ), "Dots3NoteForCausalLM": _HfExamplesInfo( "dots-studio/dots3-note-prev", is_available_online=False, @@ -1537,6 +1540,12 @@ def check_available_online( is_available_online=False, use_original_num_layers=True, # DSpark has >1 draft block ), + "DSparkV41DraftModel": _HfExamplesInfo( + "deepseek-ai/DeepSeek-V4.1-Flash", + speculative_model="deepseek-ai/DeepSeek-V4.1-Flash", # draft in mtp.* + is_available_online=False, + use_original_num_layers=True, # DSpark has >1 draft block + ), "K3DSparkModel": _HfExamplesInfo( "moonshotai/Kimi-K3", speculative_model="Inferact/Kimi-K3-DSpark", diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 4e95904f6b2d..58d5239285dd 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,7 +18,13 @@ ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.models.deepseek_v4_1.common.mm_preprocess import ( + IMAGE_PAD_ID, + IMAGE_SENTINEL_BASE_ID, +) +from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config pytestmark = pytest.mark.skipif( not current_platform.is_cuda(), @@ -26,6 +32,133 @@ ) +@pytest.fixture +def v41_moe_config(dist_init): + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=DeepseekV41Config( + text_config=dict( + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + dspark_n_routed_experts=4, + dspark_num_experts_per_tok=3, + n_shared_experts=1, + moe_intermediate_size=128, + hidden_act="silu", + swiglu_limit=10.0, + norm_topk_prob=True, + topk_method="noaux_tc", + routed_scaling_factor=1.5, + ), + ), + ), + quant_config=None, + kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"), + parallel_config=SimpleNamespace( + enable_expert_parallel=True, + enable_eplb=False, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + + +@pytest.mark.parametrize("vision", [False, True]) +@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)]) +def test_deepseek_v41_moe_routes_without_hash_table( + v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k +): + """Main and draft layers select experts by score and preserve image routing.""" + config = v41_moe_config.model_config.hf_config + config.vision_n_layers = int(vision) + with torch.device("cuda"): + moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") + hidden_states = torch.randn(4, config.hidden_size) + input_ids = ( + torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257]) + if vision + else None + ) + assert moe.gate.tid2eid is None + assert moe.gate.weight.shape == (num_experts, config.hidden_size) + assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) + assert moe.experts.w13_weight.shape == (num_experts, 256, 64) + assert moe.experts.top_k == top_k + assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2) + + with torch.no_grad(): + moe.gate.weight.normal_(std=0.01) + moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda")) + if vision: + moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias) + + scores = torch.nn.functional.softplus( + torch.nn.functional.linear(hidden_states, moe.gate.weight) + ).sqrt() + bias = moe.gate.e_score_correction_bias + if vision: + image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID) + bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias) + expected_ids = (scores + bias).topk(top_k, dim=-1).indices + expected_weights = scores.gather(1, expected_ids) + expected_weights *= config.routed_scaling_factor / expected_weights.sum( + dim=-1, keepdim=True + ) + + def check_routing(x, weights, ids, *, activation_clamp): + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close(weights, expected_weights) + assert activation_clamp == config.swiglu_limit + return x.clone() + + monkeypatch.setattr(moe.experts, "forward", check_routing) + monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) + torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states) + + +@pytest.mark.parametrize( + "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))] +) +def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults( + v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected +): + config = v41_moe_config.model_config.hf_config + config.dspark_n_routed_experts = draft_experts + config.dspark_num_experts_per_tok = draft_top_k + v41_moe_config.kernel_config.moe_backend = "auto" + captured = {} + + def make_experts(**kwargs): + captured.update(kwargs) + return torch.nn.Identity() + + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts + ) + moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn") + assert moe.gate.weight.shape == (expected[0], 128) + assert (captured["num_experts"], captured["top_k"]) == expected + assert captured["hash_indices_table"] is None + + +def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config): + config = v41_moe_config.model_config.hf_config + config.num_hash_layers = 1 + config.vocab_size = 32 + moe = DeepseekV4MoE( + v41_moe_config, + prefix="model.layers.0.ffn", + num_hash_layers=config.num_hash_layers, + ) + assert moe.gate.tid2eid.shape == (32, 2) + assert moe.gate.e_score_correction_bias is None + with pytest.raises(ValueError, match="hash MoE routing requires input_ids"): + moe(torch.zeros(1, 128)) + + def test_deepseek_v4_mega_moe_expert_mapping(): mapping = make_deepseek_v4_expert_params_mapping(2) @@ -169,15 +302,79 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): assert torch.count_nonzero(experts.w13_weight[1]) == 0 +@pytest.mark.parametrize("intermediate_size", [512, 2304]) +def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size): + """Pad gate/up separately and zero every added down-projection column.""" + experts = DeepseekV4MegaMoEExperts( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ), + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=intermediate_size, + ) + originals = [] + for param in ( + experts.w13_weight, + experts.w13_weight_scale, + experts.w2_weight, + experts.w2_weight_scale, + ): + param.data.random_(1, 128) + originals.append(param.data.clone()) + monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", + lambda: SimpleNamespace( + transform_sf_into_required_layout=lambda sf, *args: sf, + transform_weights_for_mega_moe=lambda l1, l2: (l1, l2), + ), + ) + + experts.finalize_weights() + padded_size = experts.intermediate_size + assert padded_size == (intermediate_size + 511) // 512 * 512 + for actual, original in zip(experts._transformed_l1_weights, originals[:2]): + expected = ( + original + if actual.dtype == torch.int8 + else (experts._ue8m0_uint8_to_float(original)) + ) + actual = actual.view(expected.dtype).unflatten(1, (2, padded_size)) + assert torch.equal( + actual[:, :, :intermediate_size], + expected.unflatten(1, (2, intermediate_size)), + ) + assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0 + for actual, original in zip(experts._transformed_l2_weights, originals[2:]): + expected = ( + original + if actual.dtype == torch.int8 + else (experts._ue8m0_uint8_to_float(original)) + ) + actual = actual.view(expected.dtype) + assert torch.equal(actual[..., : original.shape[-1]], expected) + assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0 + + transformed_l1 = experts._transformed_l1_weights + experts.finalize_weights() + assert experts._transformed_l1_weights is transformed_l1 + + @pytest.mark.parametrize( - "hidden_size,intermediate_size,block_size,fused", + "hidden_size,intermediate_size,block_size,fused,mxfp8", [ - pytest.param(128, 512, 128, True, id="aligned-block128"), - pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"), + pytest.param(128, 512, 128, True, False, id="aligned-block128"), + pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"), + pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"), ], ) def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( - monkeypatch, hidden_size, intermediate_size, block_size, fused + monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8 ): """V4.1-Flash's routed padding must preserve the serial shared-expert fallback.""" @@ -261,6 +458,16 @@ def scale_parameter(*shape, dtype=torch.int32): ), ), ) + if mxfp8: + for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): + linear.weight_block_size = (1, 32) + linear.weight_scale = torch.nn.Parameter( + linear.weight_scale_inv.view(torch.uint8) + .repeat_interleave(128, dim=0) + .repeat_interleave(4, dim=1), + requires_grad=False, + ) + del linear.weight_scale_inv monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 01ee0b0533e5..0d1507a03382 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -186,3 +186,79 @@ def test_context_kv_weights_are_loaded_as_merged_linear_shards(): assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1] assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr() + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32] +) +@pytest.mark.parametrize( + ("checkpoint_name", "runtime_module", "shard_id"), + [ + ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0), + ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1), + ("mtp.0.main_proj.scale", "model.main_proj", None), + ( + "mtp.0.ffn.shared_experts.w1.scale", + "model.layers.0.ffn.shared_experts.gate_up_proj", + 0, + ), + ( + "mtp.0.ffn.shared_experts.w2.scale", + "model.layers.0.ffn.shared_experts.down_proj", + None, + ), + ], +) +def test_v41_dspark_loads_linear_scales( + monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id +): + """Checkpoint ``.scale`` maps to the quant method's scale parameter and + loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static + loader (see tests/quantization/test_modelopt.py), not in load_weights.""" + from vllm.models.deepseek_v4_1.nvidia import dspark + + mxfp8 = scale_dtype != torch.float32 + scale_name = "weight_scale" if mxfp8 else "weight_scale_inv" + runtime_name = f"{runtime_module}.{scale_name}" + raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8) + checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float() + param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False) + shards = [] + + def load_scale(param, weight, *args): + shards.append(args) + assert weight.dtype == checkpoint_scale.dtype + param.copy_(weight) + + param.weight_loader = load_scale + draft = SimpleNamespace( + config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1), + quant_config=SimpleNamespace( + weight_block_size=[32, 32] if mxfp8 else [128, 128] + ), + linear_scale_name=scale_name, + pad_shared_expert=False, + model=SimpleNamespace( + layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + confidence_head=None, + ), + named_parameters=lambda: [(runtime_name, param)], + process_weights_after_loading=lambda: None, + ) + draft._remap_dspark_name = lambda name: ( + dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name) + ) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: [] + ) + + loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights( + draft, [(checkpoint_name, checkpoint_scale)] + ) + + assert loaded == {runtime_name} + assert shards == [() if shard_id is None else (shard_id,)] + torch.testing.assert_close(param, checkpoint_scale) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 30eb80a3f88f..14d85fc09d84 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -93,6 +93,7 @@ def test_registry_imports(model_arch): [ ("LlamaForCausalLM", False, False, "bi-encoder"), ("LlavaForConditionalGeneration", True, True, "bi-encoder"), + ("DeepseekV41ForCausalLM", True, False, "bi-encoder"), ("BertForSequenceClassification", False, False, "cross-encoder"), ("RobertaForSequenceClassification", False, False, "cross-encoder"), ("XLMRobertaForSequenceClassification", False, False, "cross-encoder"), diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 1c73573d8dc0..a91c99bf5378 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -37,6 +37,7 @@ Fp8MoEMethod, ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, ) @@ -63,6 +64,383 @@ ] +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA") +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) +def test_deepseek_v41_mxfp8_scale_loading( + dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype +): + """Expand linear scales before TP slicing, including padded shared experts.""" + from vllm.model_executor.layers import linear as linear_module + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, + ) + from vllm.models.deepseek_v4_1 import quant_config as quant_module + from vllm.models.deepseek_v4_1.nvidia import model as model_module + + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + quant_config = quant_module.DeepseekV4FP8Config( + is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32] + ) + quant_config._resolved_expert_dtype = "fp4" + for module in (model_module, linear_module): + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank) + tp_args = dict(quant_config=quant_config, bias=False) + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args) + block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args) + block.ffn = torch.nn.Module() + block.ffn.shared_experts = torch.nn.Module() + block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear( + 128, [128, 128], **tp_args + ) + block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args) + block.ffn.experts = torch.nn.Module() + block.ffn.experts.register_parameter( + "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False) + ) + block.engram = torch.nn.Module() + block.engram.register_parameter( + "weight_scale_inv", + torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False), + ) + + def load_expert_scale(param, weight, *args, **kwargs): + param.data.copy_(weight) + return True + + block.ffn.experts.weight_scale.weight_loader = load_expert_scale + model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model) + torch.nn.Module.__init__(model) + model.layers = torch.nn.ModuleList([block]) + model.config = SimpleNamespace(num_attention_heads=2) + model.quant_config = quant_config + model.use_sequence_parallel = False + model.get_expert_mapping = lambda: [ + ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1") + ] + + checkpoints = [ + ("attn.wq_b", "attn.wq_b", 128, 128, 0, None), + ("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)), + ("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)), + ( + "ffn.shared_experts.w1", + "ffn.shared_experts.gate_up_proj", + 96, + 128, + 0, + slice(0, 64), + ), + ( + "ffn.shared_experts.w3", + "ffn.shared_experts.gate_up_proj", + 96, + 128, + 0, + slice(64, 128), + ), + ( + "ffn.shared_experts.down_proj", + "ffn.shared_experts.down_proj", + 128, + 96, + 1, + None, + ), + ] + weights = [] + expected = [] + for source, target, n, k, axis, shard in checkpoints: + weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn) + scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8) + weights.extend( + [ + (f"layers.0.{source}.weight", weight), + (f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)), + ] + ) + dequant = ( + weight.float().reshape(n // 32, 32, k // 32, 32) + * torch.exp2(scale_bytes.float() - 127)[:, None, :, None] + ).reshape(n, k) + if "shared_experts" in source: + padded = torch.zeros(128, 128) + padded[:n, :k] = dequant + dequant = padded + expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank])) + + expert_scale = torch.full((3, 4), 125, dtype=torch.uint8) + weights.append( + ("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype)) + ) + weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype))) + mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale") + loaded = model.load_weights( + (name.removeprefix("model."), weight) for name, weight in mapper.apply(weights) + ) + assert "layers.0.ffn.experts.weight_scale" in loaded + assert torch.equal(block.ffn.experts.weight_scale, expert_scale) + assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale) + for target, shard, reference in expected: + linear = block.get_submodule(target) + assert isinstance(linear.quant_method, ModelOptLinearMethod) + assert f"layers.0.{target}.weight_scale" in loaded + weight = linear.weight if shard is None else linear.weight[shard] + scale = linear.weight_scale if shard is None else linear.weight_scale[shard] + actual = ( + weight.float().unflatten(-1, (-1, 32)) + * torch.exp2(scale.float() - 127).unsqueeze(-1) + ).flatten(-2) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA") +@pytest.mark.parametrize( + "weight_block_size,expert_dtype,scale_name", + [ + ([32, 32], "fp4", "weight_scale"), + ([128, 128], "fp4", "weight_scale_inv"), + ([32, 32], "fp8", "weight_scale_inv"), + ], +) +def test_deepseek_v41_vl_mapper_routes_linear_scales( + weight_block_size, expert_dtype, scale_name +): + """The VL wrapper must map ``.scale`` keys to the parameter the linear + quant method registers, as the text model does. A hardcoded + ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" + from vllm.models.deepseek_v4_1.nvidia import model as model_module + from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module + + vllm_config = SimpleNamespace( + quant_config=SimpleNamespace(weight_block_size=weight_block_size) + ) + resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype) + assert resolved == scale_name + + mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved) + weight = torch.empty(0) + mapped = [ + name + for name, _ in mapper.apply( + [("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)] + ) + ] + assert mapped == [ + f"language_model.model.layers.0.attn.wq_a.{scale_name}", + "language_model.model.layers.0.attn.wkv.weight", + ] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") +@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) +@pytest.mark.parametrize( + "weight_shape,is_bmm", + [ + ((129, 160), False), + ((387, 160), True), + ((3, 129, 160), False), + ((3, 128, 512), False), + ], +) +def test_deepgemm_mxfp8_preserves_weight_and_scale_values( + scale_dtype, weight_shape, is_bmm +): + """DeepGEMM layout conversion preserves native weights and E8M0 scales.""" + from vllm.model_executor.layers.quantization.utils import fp8_utils + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( + 100 + ): + pytest.skip("DeepGEMM MXFP8 requires Blackwell") + + weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn) + scales = ( + torch.randint( + 1, + 255, + (weight.numel() // 32 + 1,), + device="cuda", + dtype=torch.uint8, + )[1:] + .view(*weight_shape[:-1], weight_shape[-1] // 32) + .view(scale_dtype) + ) + original_weight = weight.view(torch.uint8).clone() + original_scales = scales.view(torch.uint8).clone() + processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block( + weight, + scales, + quant_block_shape=(1, 32), + use_e8m0=True, + is_bmm=is_bmm, + bmm_batch_size=3 if is_bmm else 0, + ) + assert packed.dtype == torch.int32 + unpacked = torch.stack( + [(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1 + ).flatten(-2) + unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8) + torch.testing.assert_close( + unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0 + ) + torch.testing.assert_close( + processed_weight.view(torch.uint8).reshape_as(original_weight), + original_weight, + rtol=0, + atol=0, + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") +@pytest.mark.parametrize("prequantized", [False, True]) +@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"]) +@pytest.mark.parametrize("num_tokens", [1, 7, 128]) +def test_mxfp8_bmm_loads_and_projects_grouped_weights( + dist_init, default_vllm_config, prequantized, config_source, num_tokens +): + """BMM metadata set after construction selects grouped weight processing.""" + from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import ( + DeepGemmMxfp8BmmLinearKernel, + ) + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant, + ) + from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( + compute_fp8_einsum_recipe, + deep_gemm_fp8_o_proj, + ) + from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model + from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( + 100 + ): + pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell") + + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + quant_config = DeepseekV4FP8Config( + is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32] + ) + quant_config._resolved_expert_dtype = "fp4" + if config_source == "mxfp8": + quant_config = get_quantization_config("mxfp8").from_config( + {"quant_method": "mxfp8"} + ) + with torch.device("cuda"): + linear = ColumnParallelLinear( + 512, + 256, + bias=False, + quant_config=quant_config, + return_bias=False, + ) + linear.is_bmm = True + linear.bmm_batch_size = 2 + weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn) + scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8) + checkpoint_scales = ( + scales.view(torch.float8_e8m0fnu) + if config_source == "deepseek" + else scales.repeat_interleave(32, dim=0) + ) + model = DeepseekV4Model.__new__(DeepseekV4Model) + torch.nn.Module.__init__(model) + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.wo_a = linear + model.layers = torch.nn.ModuleList([block]) + model.config = SimpleNamespace(num_attention_heads=2) + model.quant_config = quant_config + model.use_sequence_parallel = False + model.get_expert_mapping = lambda: [] + model.load_weights( + [ + ("layers.0.attn.wo_a.weight", weight), + ("layers.0.attn.wo_a.weight_scale", checkpoint_scales), + ] + ) + reference_weight = ( + weight.float().reshape(8, 32, 16, 32) + * torch.exp2(scales.float() - 127)[:, None, :, None] + ).reshape(2, 128, 512) + linear.quant_method.process_weights_after_loading(linear) + linear.quant_method.process_weights_after_loading(linear) + assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel) + assert linear.weight.shape == (2, 128, 512) + torch.testing.assert_close( + linear.weight.view(torch.uint8).flatten(), + weight.view(torch.uint8).flatten(), + rtol=0, + atol=0, + ) + unpacked_scales = torch.stack( + [(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1 + ).flatten(-2) + torch.testing.assert_close( + unpacked_scales.to(torch.uint8).reshape(256, 16), + scales.repeat_interleave(32, dim=0), + rtol=0, + atol=0, + ) + + x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16) + reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight) + inputs = x.transpose(0, 1) + if prequantized: + cache = torch.cat( + ( + torch.ones(num_tokens, 32, device="cuda"), + torch.zeros(num_tokens, 32, device="cuda"), + ), + dim=1, + ) + inputs = fused_inv_rope_fp8_quant( + x.permute(1, 0, 2), + torch.arange(num_tokens, device="cuda"), + cache, + n_groups=2, + heads_per_group=1, + nope_dim=448, + rope_dim=64, + quant_group_size=32, + tma_aligned_scales=current_platform.has_device_capability(100), + ) + with torch.no_grad(): + output = linear(inputs) + assert output.shape == reference.shape + assert (output.float() - reference).norm() / reference.norm() < 0.06 + if prequantized: + recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32) + projected = deep_gemm_fp8_o_proj( + x.permute(1, 0, 2), + torch.arange(num_tokens, device="cuda"), + cache, + linear, + torch.nn.Identity(), + n_groups=2, + heads_per_group=1, + nope_dim=448, + rope_dim=64, + o_lora_rank=128, + einsum_recipe=recipe, + tma_aligned_scales=tma_aligned_scales, + ) + torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0) + compiled = torch.compile(linear, backend="eager", fullgraph=True) + with torch.no_grad(): + torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0) + + def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch): monkeypatch.setattr( flashinfer_utils, diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 59a333bfe46b..60b207c1840f 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -127,6 +127,23 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head(): assert method.spec.activation is kNvfp4Dynamic +def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch): + """Standard MXFP8 checkpoints already have one scale row per weight row.""" + from vllm.model_executor.layers.linear import ReplicatedLinear + + kernel = Mock() + kernel.input_quant_key.return_value = None + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel", + lambda **kwargs: kernel, + ) + config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"}) + linear = ReplicatedLinear(64, 64, bias=False, quant_config=config) + scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2) + linear.weight_scale.weight_loader(linear.weight_scale, scales) + assert torch.equal(linear.weight_scale, scales) + + def test_modelopt_fp8_updates_weight_dims_after_transpose(): """Humming reads weight.input_dim/output_dim. Swapping the ModelWeightParameter for a plain Parameter drops them, so the per-tensor diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index c0ce30d2144a..51f7e70f5fb8 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -7,19 +7,61 @@ import torch from tests.v1.attention.utils import create_vllm_config +from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend +from vllm.models.deepseek_v4_1.sparse_mla import ( + DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, +) from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla.compressor_utils import ( CompressedSlotMappingKernel, ) from vllm.v1.attention.backends.mla.indexer import ( BuildPrefillChunkMetadataKernel, + DeepseekV4IndexerBackend, DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, ) from vllm.v1.attention.backends.mla.sparse_utils import ( ConvertReqIndexToGlobalIndexKernel, ) -from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + MLAAttentionSpec, + compute_layer_kv_cache_shape_bytes, +) from vllm.v1.worker.block_table import get_block_table_width +from vllm.v1.worker.utils import select_common_block_size + + +def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla(): + """Packed MLA/indexer groups must retain 64 compressed rows per page.""" + kernel_block_size = select_common_block_size( + 256, [DeepseekV4SparseMLABackend, DeepseekV4IndexerBackend] + ) + spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=4, + ) + assert compute_layer_kv_cache_shape_bytes(spec, 2, kernel_block_size) == ( + 2, + 1, + 64, + 132, + ) + + +def test_indexer_preserves_deepseek_v41_mla_block_size(): + """V4.1 retains its smaller pages independently of V4's C4 indexer.""" + block_size = DeepseekV41SparseMLABackend.get_supported_kernel_block_sizes()[0] + assert isinstance(block_size, int) + assert ( + select_common_block_size( + block_size, [DeepseekV41SparseMLABackend, DeepseekV41IndexerBackend] + ) + == block_size + ) def test_indexer_warmup_normalizes_zero_compress_ratios(): diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 5249de7f8c11..2bdb5e74a7b9 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -398,6 +398,56 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid assert len(gdn) == 10 +class TestSlidingWindowBucketCap: + def test_sliding_window_bucket_is_capped_at_the_main_page(self): + """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an + unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would + set the block width and the MLA group would fill under a third of + every block; capped, it splits into groups no wider than the MLA + page.""" + config = _shared_layout_config() + specs: dict[str, KVCacheSpec] = {} + for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)): + specs[f"layers.{layer}.attn"] = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=ratio, + alignment=576, + ) + specs[f"layers.{layer}.idx"] = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=ratio, + alignment=576, + ) + for layer in range(43): + specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec( + block_size=64, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=128, + alignment=576, + ) + groups = get_kv_cache_groups(config, specs) + pages = _pages(groups) + main = next(g for g in groups if "layers.2.attn" in g.layer_names) + main_bytes = sum(pages[n] for n in main.layer_names) + assert _get_kv_cache_bytes_per_block(groups) == main_bytes + swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")] + per_group = main_bytes // pages["layers.0.swa"] + assert len(swa_groups) == -(-43 // per_group) + for g in swa_groups: + assert sum(pages[n] for n in g.layer_names) <= main_bytes + assert sorted(n for g in swa_groups for n in g.layer_names) == sorted( + f"layers.{i}.swa" for i in range(43) + ) + + class TestDensePacking: def test_bytes_per_block_is_largest_group(self): groups, g1, g2 = _mixed_page_groups() @@ -548,3 +598,55 @@ def test_bound_views_round_trip(self, layout): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +class TestCompressorRingGroup: + def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): + """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity + 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring + claims one block per request, never hashes, and does not disturb the + scheduler/hash block sizes the paged groups imply.""" + config = _shared_layout_config() + config.cache_config.enable_prefix_caching = True + specs: dict[str, KVCacheSpec] = { + "layers.2.attn": MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=2, + alignment=576, + ), + "layers.2.attn.compressor.state_cache": CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=1024, + head_size_v=0, + dtype=torch.float32, + ), + } + for layer in range(4): + specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec( + block_size=64, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=128, + alignment=576, + ) + groups = get_kv_cache_groups(config, specs) + ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]] + assert len(ring_groups) == 1 + assert not ring_groups[0].kv_cache_spec.prefix_cacheable + kv_cache_config = get_kv_cache_config_from_groups( + config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + ) + assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([kv_cache_config]), + max_model_len=8192, + enable_caching=True, + hash_block_size=64, + scheduler_block_size=128, + ) + assert len(manager.coordinator.single_type_managers) == len(groups) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index dc9c1ffb3067..ded83d800de9 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -333,6 +333,57 @@ def attention_like(t: torch.Tensor) -> None: assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) +@pytest.mark.parametrize("keyword", [False, True]) +def test_quantized_activation_replay_does_not_retain_tensors( + cuda_capture_stream, keyword +): + """Release the producer tensors while preserving both buffers on replay.""" + import gc + import weakref + + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic + + @eager_break_during_capture + def consume(x, data_out, scale_out): + data_out.copy_(x.data) + scale_out.copy_(x.scale) + + source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16) + data_out = source.to(torch.float8_e4m3fn) + scale_out = source[:, ::32].to(torch.uint8) + output = torch.empty_like(source) + cap = BreakableCUDAGraphCapture() + with cap: + data = source.to(torch.float8_e4m3fn) + scale = source[:, ::32].to(torch.uint8) + activation = QuantizedActivation( + data, scale, source.dtype, source.shape, kMxfp8Dynamic + ) + data_ref, scale_ref = weakref.ref(data), weakref.ref(scale) + if keyword: + consume(x=activation, data_out=data_out, scale_out=scale_out) + else: + consume(activation, data_out, scale_out) + output.copy_(data_out) + + del data, scale, activation + gc.collect() + assert data_ref() is None + assert scale_ref() is None + + for value in (1, 2, 3): + source.fill_(value) + cap.replay() + cuda_capture_stream.synchronize() + torch.testing.assert_close(output, source, rtol=0, atol=0) + torch.testing.assert_close(scale_out, source[:, ::32].to(torch.uint8)) + + def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): """Handle an eager attention break inside a multi-stream overlap region.""" from vllm.compilation.breakable_cudagraph import ( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index d6b0becc1d2c..e6b730242981 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -29,6 +29,7 @@ ) from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KpoolTailSpec, KVCacheConfig, @@ -798,3 +799,101 @@ def test_worker_setup_tolerates_finer_scratch_group(): # The scratch group's namespace never enters the store. scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0] assert not any(key.startswith(scratch_prefix) for key in store._data) + + +def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): + """A per-request ring group (CircularBufferSpec, capacity as block_size) + sits beside the paged group: worker setup tolerates its block size (the + DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged + blocks only, and lookup hits on the paged group alone.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + ring = CircularBufferSpec( + block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8 + ) + cfg = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=4 * full.page_size_bytes, + layers=["L0"], + layer_stride=4 * full.page_size_bytes, + block_stride=full.page_size_bytes, + ), + KVCacheTensor( + size=4 * ring.page_size_bytes, + layers=["L1"], + layer_stride=4 * ring.page_size_bytes, + block_stride=ring.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], ring), + ], + ) + vllm_config = _minimal_vllm_config(cache_block_size=16) + store = _DictStore() + + worker = _build_worker_with_dict_store(vllm_config, cfg, store) + worker.tp_size = 1 + worker.pp_size = 1 + worker.num_kv_head = 8 + assert worker.token_dbs[1].hash_block_size == 8 + + raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8) + raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8) + kv_caches = { + "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0], + "L1": dense_kv_cache_views(raw_ring, ring, 4, 1, KVCacheLayout.LBNHC)[0], + } + + def _fake_thread_init(*args, **kwargs): + for v in list(args) + list(kwargs.values()): + if isinstance(v, threading.Event): + v.set() + m = MagicMock() + m.start = lambda: None + return m + + with ( + patch.object( + mooncake_store_worker, + "KVCacheStoreSendingThread", + side_effect=_fake_thread_init, + ), + patch.object( + mooncake_store_worker, + "KVCacheStoreRecvingThread", + side_effect=_fake_thread_init, + ), + ): + worker.register_kv_caches(kv_caches) + + send_thread = KVCacheStoreSendingThread( + store=store, + token_databases=worker.token_dbs, + block_size=worker.block_size, + coord=worker.coord, + tp_rank=worker.tp_rank, + group_put_steps=worker._group_tp_replication_factors, + kv_role=worker.kv_role, + ready_event=threading.Event(), + enable_kv_event=False, + group_participates=[True, False], + ) + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)] + save_req = ReqMeta( + req_id="r0", + token_len_chunk=64, + # The ring group holds one block for the request's lifetime. + block_ids=([1, 2, 3, 4], [1]), + block_hashes=hs, + can_save=True, + store_job_id=1, + ) + send_thread.add_request(save_req) + send_thread._handle_request(send_thread.request_queue.get()) + worker.store = store + + assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64 + assert not any("@group:1" in key for key in store._data) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index a9f800739881..32908f78adf6 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -1251,3 +1251,112 @@ def test_csa_linear_remote_ple_is_copied_whole(): tp_ratio=-2, transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1), ) + + +def _make_ring_worker(): + """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a + model whose sliding-window KV lives in per-request rings.""" + from unittest.mock import MagicMock + + from vllm.config import set_current_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + paged_specs = { + f"paged.{index}": MLAAttentionSpec( + block_size=4, num_kv_heads=1, head_size=64, dtype=torch.float16 + ) + for index in range(2) + } + ring_specs = { + f"ring.{index}": CircularBufferSpec( + block_size=6, + num_kv_heads=1, + head_size=32, + head_size_v=0, + dtype=torch.float16, + ) + for index in range(2) + } + + def group(specs): + uniform = UniformTypeKVCacheSpecs.from_specs(specs) + assert uniform is not None + return KVCacheGroupSpec(list(specs), uniform) + + # Both rings overlay the second paged region, so the scratch group must + # address only that region while the paged group spans both. + tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1")) + region_size, page_size = 512, 256 + kv_cache_config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[ + KVCacheTensor( + size=len(tensor_regions) * region_size, + layers=[layer_name], + layer_stride=region_size, + block_stride=page_size, + offset=region_index * region_size, + ) + for region_index, layer_names in enumerate(tensor_regions) + for layer_name in layer_names + ], + kv_cache_groups=[group(paged_specs), group(ring_specs)], + ) + + vllm_config = create_vllm_config(block_size=4) + vllm_config.cache_config.enable_prefix_caching = False + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [4] + fake_backend.get_name.return_value = "TEST_ATTN" + fake_backend.full_cls_name.return_value = "test.AttentionBackend" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(2)] + worker.register_kv_caches( + { + layer_name: tensors[region_index] + for region_index, layer_names in enumerate(tensor_regions) + for layer_name in layer_names + } + ) + return worker + + +@pytest.mark.cpu_test +def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions(): + worker = _make_ring_worker() + + assert worker._ple_group_index is None + assert worker._scratch_region_indices == [1] + + desc_ids = worker._compute_desc_ids( + block_ids=([1], [0]), + dst_num_blocks=2, + block_size_ratio=None, + physical_blocks_per_logical=1, + ) + # Paged block 1 in both regions; ring block 0 only in its scratch region. + assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0]) diff --git a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py index 69f9c2da415b..af4c7858fcad 100644 --- a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py +++ b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py @@ -19,6 +19,7 @@ SimpleCPUOffloadConnector, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -90,3 +91,59 @@ def test_disk_mode_rejects_non_positive_capacity() -> None: "disk_capacity_bytes": 0, } ) + + +def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig: + """A paged full-attention group beside a per-request ring group.""" + full = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring = CircularBufferSpec( + block_size=6, num_kv_heads=1, head_size=32, head_size_v=0, dtype=DTYPE + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=_BYTES_PER_BLOCK * num_blocks, + layers=["attn"], + layer_stride=_BYTES_PER_BLOCK * num_blocks, + block_stride=_BYTES_PER_BLOCK, + ), + KVCacheTensor( + size=ring.page_size_bytes * num_blocks, + layers=["ring"], + layer_stride=ring.page_size_bytes * num_blocks, + block_stride=ring.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["attn"], full), + KVCacheGroupSpec(["ring"], ring), + ], + ) + + +def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None: + """The ring group's capacity is not a token granularity, it holds no + hashed blocks, and it never participates in CPU hits.""" + vllm_config = create_vllm_config( + kv_connector="SimpleCPUOffloadConnector", + kv_role="kv_both", + kv_connector_extra_config={"cpu_bytes_to_use": _BYTES_PER_BLOCK * 8}, + ) + connector = SimpleCPUOffloadConnector( + vllm_config=vllm_config, + role=KVConnectorRole.SCHEDULER, + kv_cache_config=_make_ring_kv_cache_config(), + ) + scheduler = connector.scheduler_manager + assert scheduler.fa_gidx == 0 + cpu_groups = scheduler.cpu_kv_cache_config.kv_cache_groups + assert [g.kv_cache_spec.prefix_cacheable for g in cpu_groups] == [True, False] + # Only the paged group takes part in CPU-side hit lookup. + attention_groups = scheduler.cpu_coordinator.attention_groups + assert [g.group_ids for g in attention_groups] == [[0]] diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 1952e3a4b0b3..901e076e0950 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1319,12 +1319,10 @@ def test_inflight_finish_deferred_cleanup() -> None: # Test 8: Null GPU blocks are skipped in store and load transfer pairs # --------------------------------------------------------------------------- def test_multi_group_null_blocks_skipped() -> None: - """Null GPU blocks (no block_hash) must not appear in store or load pairs. + """The null block id never appears in store or load pairs. - In eager store mode, _prepare_eager_store_specs skips blocks whose - block_hash is None (null blocks have no hash). We verify this by mixing - real hashed blocks with unhashed (null-like) blocks in a single group and - checking that only real blocks appear in the store list. + A null block-table slot whose hash is still cached on the GPU is stored + from the cached block instead; the null block itself is never copied. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False) sched = fix.scheduler @@ -1356,9 +1354,11 @@ def test_multi_group_null_blocks_skipped() -> None: f"Null block id {null_block_id} should not appear in store transfer pairs" ) - # Only real block should be scheduled for store - assert len(meta.store_gpu_blocks) == 1 + # The nulled slot's hash is still cached on the GPU, so that block is + # recovered from the pool and stored alongside the real one. + assert len(meta.store_gpu_blocks) == 2 assert gpu_blocks[0].block_id in meta.store_gpu_blocks + assert gpu_blocks[1].block_id in meta.store_gpu_blocks # Complete the store assert meta.store_event >= 0 @@ -1374,13 +1374,12 @@ def test_multi_group_null_blocks_skipped() -> None: block_hasher=req._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) - # Only 1 block was stored (the real one) - assert hit_tokens == BLOCK_SIZE + assert hit_tokens == 2 * BLOCK_SIZE assert is_async is True # Allocate new GPU blocks for the load - gpu_blocks2 = gpu_pool.get_new_blocks(1) - kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],)) + gpu_blocks2 = gpu_pool.get_new_blocks(2) + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) sched_out2 = make_scheduler_output({req2.request_id: 1}) @@ -2800,3 +2799,47 @@ def test_boundary_handoff_keeps_block_meta_index_parallel() -> None: assert req_ids == [req.request_id] assert block_meta is not None assert len(block_meta) == len(gpu_ids) == len(cpu_ids) + + +def test_finished_eager_store_recovers_nulled_window_tail() -> None: + """A sliding-window group nulls pages that left the window before the + connector sees the finished request's block table, while the GPU keeps + them hashed in its free queue. The eager store must copy those pages, or + the CPU lookup (which needs the whole window) misses every time.""" + fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + null_block = gpu_pool.null_block + + num_blocks = 4 # == sliding window of the test SWA group + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1) + # remove_skipped_blocks() ran: the two oldest window pages are freed (still + # hashed) and their slots hold the null block. + gpu_pool.free_blocks(swa_blocks[:2]) + swa_table = [null_block, null_block, *swa_blocks[2:]] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + block_ids = kv_blocks.get_block_ids() + sched.request_finished_all_groups(req, block_ids) + meta = sched.build_connector_meta(make_scheduler_output({})) + + assert null_block.block_id not in meta.store_gpu_blocks + assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)} + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-window-tail-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 6b63588150ed..2e1e8fc6194c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2143,13 +2143,15 @@ def register_ops_once() -> None: if not current_platform.is_rocm(): return - from vllm.platforms.rocm import on_gfx11 + from vllm.platforms.rocm import on_gfx11, on_gfx950 - if on_gfx11() and not _OPS_REGISTERED: + # This op has self-contained Triton/C++ implementations on gfx11 + # and gfx950. Only its optional top-k fast path comes from aiter. + if (on_gfx11() or on_gfx950()) and not _OPS_REGISTERED: direct_register_custom_op( op_name="rocm_aiter_sparse_attn_indexer", op_func=rocm_aiter_sparse_attn_indexer, - mutates_args=["topk_indices_buffer"], + mutates_args=["topk_indices_buffer", "candidate_blocks"], fake_impl=rocm_aiter_sparse_attn_indexer_fake, dispatch_key=current_platform.dispatch_key, ) @@ -2327,7 +2329,7 @@ def register_ops_once() -> None: direct_register_custom_op( op_name="rocm_aiter_sparse_attn_indexer", op_func=rocm_aiter_sparse_attn_indexer, - mutates_args=["topk_indices_buffer"], + mutates_args=["topk_indices_buffer", "candidate_blocks"], fake_impl=rocm_aiter_sparse_attn_indexer_fake, dispatch_key=current_platform.dispatch_key, ) diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 343b9ef398bd..8e25d05c280f 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -42,6 +42,7 @@ is_forward_context_available, ) from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors @@ -56,6 +57,12 @@ def is_breakable_cudagraph_enabled() -> bool: F = TypeVar("F", bound=Callable[..., Any]) +def _weak_ref_capture_arg(arg: Any) -> Any: + if isinstance(arg, QuantizedActivation): + return arg.weak_ref() + return weak_ref_tensor(arg) + + def eager_break_during_capture(fn: F) -> F: """Decorator that turns a custom-op Python kernel into a "break point" for the breakable cudagraph capture. @@ -105,13 +112,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: # Weak-ref args: strong refs in the replay lambda pin cudagraph-pool # slots across batch descriptors. cudagraph owns the slot, so the # weak_ref is safe to deref on replay. - weak_args = tuple( - weak_ref_tensor(a) if isinstance(a, torch.Tensor) else a for a in args - ) - weak_kwargs = { - k: weak_ref_tensor(v) if isinstance(v, torch.Tensor) else v - for k, v in kwargs.items() - } + weak_args = tuple(_weak_ref_capture_arg(a) for a in args) + weak_kwargs = {k: _weak_ref_capture_arg(v) for k, v in kwargs.items()} return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs)) return wrapper # type: ignore[return-value] diff --git a/vllm/config/engram.py b/vllm/config/engram.py index b46311315f37..2aede1868d55 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -15,6 +15,14 @@ logger = init_logger(__name__) +# Architecture -> the hf_text_config field naming its n-gram layers. A model is +# only configurable here if it actually has such layers to store. +_NGRAM_LAYER_FIELDS = { + "DeepseekV41ForCausalLM": "engram_layer_ids", + "Qwen4ExpForCausalLM": "ple_layer_ids", + "Qwen4ExpForConditionalGeneration": "ple_layer_ids", +} + def _default_cpu_offload() -> bool: """Honor the legacy environment variable only when the field is omitted.""" @@ -40,23 +48,25 @@ class EngramConfig: Otherwise, each DP rank has a separate TP-sharded embedding replica.""" def verify_model_config(self, model_config: "ModelConfig | None") -> None: - """Reject Engram configuration for models without supported embeddings.""" + """Reject Engram configuration for models without n-gram embeddings.""" from vllm.platforms import current_platform - supported_architectures = { - "Qwen4ExpForCausalLM", - "Qwen4ExpForConditionalGeneration", - } + field = ( + _NGRAM_LAYER_FIELDS.get(model_config.architecture) + if model_config is not None + else None + ) if ( model_config is None - or model_config.architecture not in supported_architectures + or field is None or not current_platform.is_cuda() - or not getattr(model_config.hf_text_config, "ple_layer_ids", None) + or not getattr(model_config.hf_text_config, field, None) ): raise ValueError( "EngramConfig requires a model with supported Engram " - "embeddings. Currently only the CUDA Qwen4Exp implementation " - "with non-empty ple_layer_ids is supported." + "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek " + "V4.1 implementations with non-empty n-gram layer ids are " + "supported." ) def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 1312cbb24157..a6dddc317875 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -165,6 +165,7 @@ def with_default( { "DeepseekV4ForCausalLM", "DeepSeekV4MTPModel", + "DeepseekV41ForCausalLM", } ) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index d4085d7d52f8..ae56e5880327 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -679,12 +679,18 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: ], } ) - if hf_config.model_type == "deepseek_v4": + if hf_config.model_type in ("deepseek_v4", "deepseek_v41"): + # V4.1 has no classic-MTP draft: its checkpoints ship DSpark stages + # under ``mtp.*``, so only V4 gets an MTP architecture here. The + # DSpark path rewrites ``architectures`` itself and needs only + # ``n_predict``; ``method="mtp"`` on V4.1 is rejected below. + is_v41 = hf_config.model_type == "deepseek_v41" hf_config.model_type = "deepseek_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", None) - hf_config.update( - {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]} - ) + overrides = {"n_predict": n_predict} + if not is_v41: + overrides["architectures"] = ["DeepSeekV4MTPModel"] + hf_config.update(overrides) if hf_config.model_type in ("pangu_ultra_moe"): hf_config.model_type = "pangu_ultra_moe_mtp" if hf_config.model_type == "pangu_ultra_moe_mtp": @@ -1335,6 +1341,18 @@ def __post_init__(self): MTPModelTypes ): self.method = "mtp" + if ( + self.target_model_config is not None + and self.target_model_config.hf_config.model_type + == "deepseek_v41" + ): + raise ValueError( + "DeepSeek V4.1 has no classic-MTP draft: its " + "checkpoints ship DSpark stages under mtp.* " + "(main_proj/markov_head/confidence_head) and carry " + "no e_proj/h_proj/enorm/hnorm/hc_head weights. Use " + "speculative method 'dspark' instead of 'mtp'." + ) if ( self.num_speculative_tokens > 1 and self.draft_model_config.hf_config.model_type @@ -1398,12 +1416,26 @@ def __post_init__(self): and "Gemma4DSparkModel" not in self.draft_model_config.architectures and "K3DSparkModel" not in self.draft_model_config.architectures ): - # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config + # DeepSeek-V4(.1) DSpark reuses the full target config # and its weights ship in the target checkpoint. - self.draft_model_config.hf_config.model_type = "deepseek_v4" - self.draft_model_config.hf_config.architectures = [ - "DSparkDraftModel" + is_v41 = ( + self.target_model_config.hf_config.model_type == "deepseek_v41" + ) + draft_hf_config = self.draft_model_config.hf_config + draft_hf_config.model_type = ( + "deepseek_v41" if is_v41 else "deepseek_v4" + ) + draft_hf_config.architectures = [ + "DSparkV41DraftModel" if is_v41 else "DSparkDraftModel" ] + if is_v41: + # hf_config_override set n_predict to the number of + # MTP stages (3), but one DSpark round drafts + # dspark_block_size tokens; num_speculative_tokens + # divisibility is checked against n_predict below. + draft_hf_config.n_predict = getattr( + draft_hf_config, "dspark_block_size", None + ) or getattr(draft_hf_config, "n_predict", None) self.draft_model_config.quantization = ( self.target_model_config.quantization ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5bd6af20deca..39c2cb4d52a0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -81,6 +81,7 @@ "DeepseekV4ForCausalLM", "DeepseekV4ForConditionalGeneration", "DeepSeekV4MTPModel", + "DeepseekV41ForCausalLM", "Dots3NoteForCausalLM", "Dots3NoteMTPModel", "Glm5NextForCausalLM", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 3f7008201329..67a935d7eedc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -91,6 +91,7 @@ def __init__( assert all( scheduler_block_size % g.kv_cache_spec.block_size == 0 for g in kv_cache_groups + if g.kv_cache_spec.prefix_cacheable ), "scheduler_block_size must be a multiple of each group's block_size" self.kv_cache_groups = kv_cache_groups self.mamba_group_ids = { diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index ceef10e9fd6e..dd2e97c96578 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -460,11 +460,13 @@ def __init__( for spec in iter_layer_specs(group.kv_cache_spec) ) ] - if len(ple_groups) != 1 or len(ple_groups[0][1].layer_names) != 1: + if len(ple_groups) > 1 or any( + len(group.layer_names) != 1 for _, group in ple_groups + ): raise ValueError( - "CSA-linear NIXL requires exactly one PLE cache owner." + "CSA-linear NIXL requires at most one PLE cache owner." ) - self._ple_group_index = ple_groups[0][0] + self._ple_group_index = ple_groups[0][0] if ple_groups else None if self._has_mamba: assert self._is_hma_required diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index 4be2f4f9ffe1..bd6cab668e8b 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -8,7 +8,7 @@ as_quantized_activation. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace import torch @@ -34,6 +34,16 @@ class QuantizedActivation: orig_shape: torch.Size quant_key: QuantKey + def weak_ref(self) -> "QuantizedActivation": + """Return a copy with non-owning tensor references for CUDA graph replay.""" + from vllm.utils.torch_utils import weak_ref_tensor + + return replace( + self, + data=weak_ref_tensor(self.data), + scale=weak_ref_tensor(self.scale), + ) + def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: """Advertise the kernel's pre-quantized input key on the layer, if any. diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index a2c65738e202..fd1d0470745f 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -115,7 +115,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # lazy import to avoid triggering `torch.compile` too early from vllm.config.quantization import _ONLINE_SHORTHANDS from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig - from vllm.models.deepseek_v4 import DeepseekV4FP8Config + + if current_platform.is_cuda() or current_platform.is_rocm(): + # The v4.1 class is the v4 one extended to accept model_type + # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has + # platform-specific implementations for both CUDA and ROCm. + from vllm.models.deepseek_v4_1 import ( + DeepseekV4FP8Config as DeepseekV41FP8Config, + ) + + deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config + else: + from vllm.models.deepseek_v4 import DeepseekV4FP8Config + + deepseek_config = DeepseekV4FP8Config from .auto_awq import AutoAWQConfig from .auto_gptq import AutoGPTQConfig @@ -161,7 +174,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "inc": INCConfig, "mxfp4": Mxfp4Config, "gpt_oss_mxfp4": GptOssMxfp4Config, - "deepseek_v4_fp8": DeepseekV4FP8Config, + "deepseek_v4_fp8": deepseek_config, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 199d856e3b78..e42f43b39c6c 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -937,14 +937,13 @@ def w8a8_triton_block_scaled_mm( assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] - # Triton cannot currently bind E8M0 scale tensors directly. On ROCm, - # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format, - # so decode them to fp32 before launching the kernel. - if current_platform.is_rocm() or current_platform.is_xpu(): - if As.dtype == torch.float8_e8m0fnu: - As = _upcast_e8m0_to_fp32(As).contiguous() - if Bs.dtype == torch.float8_e8m0fnu: - Bs = _upcast_e8m0_to_fp32(Bs).contiguous() + # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints + # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32 + # before launching the kernel. + if As.dtype == torch.float8_e8m0fnu: + As = _upcast_e8m0_to_fp32(As).contiguous() + if Bs.dtype == torch.float8_e8m0fnu: + Bs = _upcast_e8m0_to_fp32(Bs).contiguous() assert A.shape[-1] == B.shape[-1] assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() @@ -1089,8 +1088,6 @@ def deepgemm_post_process_weight_scale_block( is_sfa: bool = False, ) -> torch.Tensor: if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): - # Scales already in E8M0 from checkpoint; upcast to fp32 and let - # DeepGEMM pack the layout expected by the target architecture. ws = _upcast_e8m0_to_fp32(ws) else: assert ws.dtype == torch.float32, ( @@ -1121,12 +1118,7 @@ def deepgemm_post_process_fp8_weight_block( f"to be torch.float8_e4m3fn, got {wq.dtype} instead." ) - if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): - # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0 - # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization - # (weights already have power-of-two scales). - ws = _upcast_e8m0_to_fp32(ws) - else: + if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8): assert ws.dtype == torch.float32, ( f"Expected tensor scales dtype to be torch.float32 or " f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead" @@ -1137,7 +1129,7 @@ def deepgemm_post_process_fp8_weight_block( if is_bmm: # Reshape 2D weight/scale to 3D for grouped BMM (einsum): # wq: (g*r, d) -> (g, r, d) - # ws: (g*r/128, d/128) -> (g, r/128, d/128) + # Scale dimensions follow quant_block_shape. g = bmm_batch_size assert wq.ndim == 2 and ws.ndim == 2 d = wq.size(1) diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index 0c9324a5490a..bfcc69e4b1de 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -126,8 +126,14 @@ def _kernel( amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY) # [BLOCK_M] sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) - descale = tl.exp2(sb - 127.0) - xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty) + # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``: + # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127, + # which is subnormal in fp32 and flushes to zero on CDNA, turning the + # block's zeros into 0/0 == NaN. The reciprocal is normal for every + # reachable sb, and both forms are exact powers of two, so nothing + # else changes. + rescale = tl.exp2(127.0 - sb) + xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty) tl.store( xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, xq, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 485d15354deb..bbbe98bc8de0 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -330,7 +330,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: quant_config = getattr(model_config.hf_config, "quantization_config", None) if quant_config is not None and quant_config.get("quant_method") == "fp8": model_type = getattr(model_config.hf_config, "model_type", None) - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v41"): model_config.hf_config.quantization_config["quant_method"] = ( "deepseek_v4_fp8" ) @@ -343,7 +343,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: and hf_text_quant_config.get("quant_method") == "fp8" ): model_type = getattr(model_config.hf_text_config, "model_type", None) - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v41"): model_config.hf_text_config.quantization_config["quant_method"] = ( "deepseek_v4_fp8" ) @@ -1003,6 +1003,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "ColQwen3_5": ColQwen3_5Config, "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig, + "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index a9fe7c2ec2ad..69dd88f34918 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -376,6 +376,10 @@ "vllm.models.deepseek_v4", "DeepseekV4ForConditionalGeneration", ), + "DeepseekV41ForCausalLM": ( + "vllm.models.deepseek_v4_1", + "DeepseekV41ForCausalLM", + ), "Dots3NoteForCausalLM": ( "vllm.models.dots3_note", "Dots3NoteForCausalLM", @@ -640,6 +644,10 @@ "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), + "DSparkV41DraftModel": ( + "vllm.models.deepseek_v4_1", + "DSparkDeepseekV4ForCausalLM", + ), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "K3DSparkModel": ( diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index 6ca8c94fff0a..fc7731509d6f 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -173,7 +173,10 @@ def deepseek_v4_mhc_warmup( # large checkpoints; bail out for any model that is not DeepSeek V4. config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None - if model_type is not None and model_type != "deepseek_v4": + if model_type is not None and model_type not in ( + "deepseek_v4", + "deepseek_v41", + ): return layer = _find_first_mhc_layer(model) diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 80dfd5199614..5209b944345b 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -28,15 +28,20 @@ "FLASHMLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV4", "ROCM_FLASHMLA_SPARSE_DSV4", + "FLASHMLA_SPARSE_DSV41", + "FLASHINFER_MLA_SPARSE_DSV41", "DEEPSEEK_SPARSE_SWA", } ) _FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"}) -_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"}) +_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset( + {"FLASHINFER_MLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV41"} +) _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = { "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2", "FLASHINFER_MLA_SPARSE_DSV4": "DSv4", + "FLASHINFER_MLA_SPARSE_DSV41": "DSv4.1", } _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 77e98a376e95..9bb53526c99b 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -58,7 +58,7 @@ ) from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV4IndexerBackend, + DeepseekV41IndexerBackend, dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) @@ -1005,7 +1005,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: - return DeepseekV4IndexerBackend + return DeepseekV41IndexerBackend class DeepseekV4Indexer(nn.Module): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 5ec07ad211bb..c20eaab0e53c 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -245,9 +245,20 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows. return [256] +class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_INDEXER" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [64 if current_platform.is_device_capability_family(90) else 128] + + @dataclass class DeepseekV32IndexerPrefillChunkMetadata: block_table: torch.Tensor @@ -1239,9 +1250,9 @@ def build( ) block_table = self.indexer_decode_block_table_buffer[:rows, :cols] - seq_lens_is_buffer_view = (use_native and next_n > 1) or ( - not use_native and max_decode_len > 1 - ) + # Flattening always returns a buffer view, including single-token + # batches. Keep its address stable across varlen graph replays. + seq_lens_is_buffer_view = not use_native or next_n > 1 # DCP: localize the now-expanded per-token global bounds to this # rank's owned KV. Done here (after expansion) so each token's global diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 45a381d4a3c6..dd2392d6c239 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -11,7 +11,7 @@ import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CUDAGraphMode +from vllm.config import CUDAGraphMode, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.platforms import current_platform @@ -819,6 +819,32 @@ def rocm_fp8_mqa_logits( return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) +def _max_decode_logits_rows(num_batched_tokens: int) -> int: + """Upper bound on decode rows the paged-MQA logits buffer can ever hold. + + ``rocm_fp8_paged_mqa_logits`` sizes its workspace as + ``(batch_size * next_n, max_model_len)``. ``batch_size`` is bounded by + ``max_num_seqs`` and ``next_n`` by ``1 + num_speculative_tokens``, which is + far tighter than ``max_num_batched_tokens`` -- 192 vs 16384 for a typical + 32-seq DSpark-5 deployment. The loose bound is harmless at short contexts + but scales with ``max_model_len``, so at the model's full context it asks + for tens of TiB and the engine cannot start. Take whichever valid bound is + smaller; the workspace is locked after profiling, so it must not be under- + estimated. + """ + try: + vllm_config = get_current_vllm_config() + except Exception: + return num_batched_tokens + scheduler_config = getattr(vllm_config, "scheduler_config", None) + max_num_seqs = getattr(scheduler_config, "max_num_seqs", None) + if not max_num_seqs: + return num_batched_tokens + speculative_config = getattr(vllm_config, "speculative_config", None) + num_spec = getattr(speculative_config, "num_speculative_tokens", 0) or 0 + return min(num_batched_tokens, max_num_seqs * (1 + num_spec)) + + def rocm_aiter_sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: LayerNameType, @@ -835,6 +861,9 @@ def rocm_aiter_sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool = False, compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, ) -> torch.Tensor: return topk_indices_buffer @@ -856,6 +885,9 @@ def rocm_aiter_sparse_attn_indexer( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool = False, compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -881,15 +913,15 @@ def rocm_aiter_sparse_attn_indexer( ) # Decode logits buffer, used by rocm_fp8_paged_mqa_logits. - # batch_size * next_n <= hidden_states.shape[0] == max_num_batched_tokens + decode_rows = _max_decode_logits_rows(hidden_states.shape[0]) if _ON_GFX942 or _ON_GFX950: workspace_manager.get_simultaneous( - ((hidden_states.shape[0], max_model_len), torch.float32), + ((decode_rows, max_model_len), torch.float32), ) else: workspace_manager.get_simultaneous( ( - (q_fp8.shape[1], hidden_states.shape[0], max_model_len), + (q_fp8.shape[1], decode_rows, max_model_len), torch.float32, ), ) @@ -919,6 +951,9 @@ def rocm_aiter_sparse_attn_indexer( topk_indices_buffer, skip_k_cache_insert, compress_ratio, + candidate_blocks, + candidate_block_size, + candidate_write, ) layer_attn_metadata = attn_metadata[k_cache_prefix] assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata) @@ -974,6 +1009,30 @@ def rocm_aiter_sparse_attn_indexer( chunk.cu_seqlen_ks, chunk.cu_seqlen_ke, ) + if candidate_blocks is not None: + from vllm.model_executor.layers.sparse_attn_indexer import ( + _apply_candidate_mask, + _select_candidate_blocks, + ) + + chunk_candidates = candidate_blocks[chunk.token_start : chunk.token_end] + if candidate_write: + _select_candidate_blocks( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + chunk_candidates.shape[1], + candidate_block_size, + chunk_candidates, + ) + else: + _apply_candidate_mask( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + chunk_candidates, + candidate_block_size, + ) topk_indices = topk_indices_buffer[ chunk.token_start : chunk.token_end, :topk_tokens ] @@ -1041,6 +1100,37 @@ def rocm_aiter_sparse_attn_indexer( max_model_len=max_model_len, ) + if candidate_blocks is not None: + from vllm.model_executor.layers.sparse_attn_indexer import ( + _apply_candidate_mask, + _select_candidate_blocks, + ) + + num_rows = logits.shape[0] + visible = decode_metadata.seq_lens.reshape(-1) + if visible.numel() != num_rows: + visible = visible.repeat_interleave(next_n) + visible = visible[:num_rows].to(torch.int64) + row_starts = torch.zeros_like(visible) + decode_candidates = candidate_blocks[:num_rows] + if candidate_write: + _select_candidate_blocks( + logits, + row_starts, + visible, + decode_candidates.shape[1], + candidate_block_size, + decode_candidates, + ) + else: + _apply_candidate_mask( + logits, + row_starts, + visible, + decode_candidates, + candidate_block_size, + ) + topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] num_rows = logits.shape[0] @@ -1097,6 +1187,10 @@ def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor: ) return _upcast_e8m0_to_fp32(scale).contiguous() + if scale.dtype == torch.uint8: + # MXFP8 parameters preserve E8M0 scales as their raw exponent bytes. + # They are biased exponents, not numeric uint8 scale values. + return torch.exp2(scale.to(torch.int16).to(torch.float32) - 127.0) return scale.to(torch.float32) @@ -1226,7 +1320,15 @@ def _get_cached_wo_a_bf16( ) wo_a_scale_param = get_fp8_block_weight_scale(wo_a) - if wo_a_scale_param is not None: + if wo_a_scale_param is None: + # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the + # historical ``_inv`` suffix. + wo_a_scale_param = getattr(wo_a, "weight_scale", None) + # Emulated MXFP8 kernels can replace the original one-byte weight with an + # already-dequantized BF16 tensor while retaining the scale attribute for + # metadata. Applying that retained scale again would double-dequantize the + # weight. Block scaling is only valid while the one-byte FP8 storage remains. + if wo_a_scale_param is not None and wo_a.weight.element_size() == 1: wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to( torch.float32 ) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 38bae039095d..e2e38ec5d595 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KpoolTailSpec, @@ -2013,15 +2014,22 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: cdiv(len(names), n) * page for page, names in page_size_layers.items() ) - # Bytes a block must hold however the mamba buckets end up split: a mamba - # bucket can go down to one state per group, every other bucket's split is - # already fixed by the repeat pattern. + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + anchor_bytes = max( ( widest_group_bytes( page_size_layers, len(spec.kv_cache_specs) - if isinstance(spec.first_spec, MambaSpec) + if is_state_bucket(spec) else num_groups_for(spec, balanced), ) for spec, page_size_layers, balanced in bucketed @@ -2032,10 +2040,9 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: groups = [] for spec, page_size_layers, balanced in bucketed: num_groups = num_groups_for(spec, balanced) - # `_align_hybrid_block_size` pads a mamba state up to one attention - # page, so cap a mamba group at the states a block already fits rather - # than let it widen the block. - if anchor_bytes and isinstance(spec.first_spec, MambaSpec): + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) num_groups = max( num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) @@ -2071,8 +2078,9 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: if spec_config is None or not spec_config.use_eagle(): return False model_config = vllm_config.model_config - return ( - model_config is not None and model_config.hf_config.model_type == "deepseek_v4" + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", ) @@ -2093,9 +2101,9 @@ def _annotate_eagle_groups( spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses - the target's own decoder layer and so carries no spec marker. Its draft - attention layer is always the last registered layer, so flag whichever + 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block + reuses the target's own decoder layer and so carries no spec marker. Its + draft attention layer is always the last registered layer, so flag whichever group holds it. This rule is only valid where the groups partition exactly the layers of ``kv_cache_spec``, which is true on the packed grouping path and not in general; other callers must leave @@ -2108,7 +2116,8 @@ def _annotate_eagle_groups( kv_cache_spec: The kv cache spec of each attention layer, in layer registration order. Only read by rule 2. kv_cache_groups: Groups to annotate in place. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group. + use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed + group. """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle_block_drop(): diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 9cae2f714cf7..bc7e6c3a13b4 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -51,7 +51,7 @@ if TYPE_CHECKING: from vllm.v1.core.kv_cache_manager import KVCacheBlocks - from vllm.v1.core.kv_cache_utils import KVCacheBlock + from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -783,6 +783,18 @@ def _prepare_eager_store_specs( req_ids = list(dict.fromkeys(req_ids)) return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta + def _cached_gpu_block( + self, resolved_hashes: "BlockHashList", block_idx: int, group_id: int + ) -> "KVCacheBlock | None": + """Return the GPU block still cached under this group's block hash.""" + if block_idx >= len(resolved_hashes): + return None + assert self._gpu_block_pool is not None + blocks = self._gpu_block_pool.get_cached_block( + resolved_hashes[block_idx], [group_id] + ) + return blocks[0] if blocks else None + def _select_eager_blocks_to_store( self, state: StoreRequestState, @@ -827,18 +839,25 @@ def _select_eager_blocks_to_store( # the middle of the request. group_size = self.group_block_sizes[g] ready = min(len(group_gpu_ids), aligned_tokens // group_size) - resolved_hashes = ( - resolve_block_hashes( - request.block_hashes, self.hash_block_size, group_size - ) - if block_meta is not None - else None + resolved_hashes = resolve_block_hashes( + request.block_hashes, self.hash_block_size, group_size ) curr_mm_idx = 0 secondary_mm_idx = 0 start = state.num_stored_blocks[g] for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start): gpu_block = self._gpu_block_pool.blocks[gpu_block_id] + if gpu_block.is_null: + # Sliding-window groups null pages that left the window + # before the connector sees the block table, but the + # retained prefix-cache tail stays hashed in the GPU free + # queue; store it from there. + recovered = self._cached_gpu_block(resolved_hashes, i, g) + if recovered is None: + advanced_per_group[g] += 1 + continue + gpu_block = recovered + gpu_block_id = gpu_block.block_id if ( self._classify_store_candidate(gpu_block) is not _StoreAdmission.READY @@ -854,7 +873,6 @@ def _select_eager_blocks_to_store( if block_meta is not None: token_start = i * group_size token_end = token_start + group_size - assert resolved_hashes is not None parent_hash = ( None if i == 0 diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index f252a35a51f8..5b13530ef73f 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -17,9 +17,9 @@ def _resolve_dspark_attention_backend( ) -> AttentionBackendEnum | None: if draft_backend is not None: return draft_backend - # DeepSeek-V4 draft layers share the target's KV-cache layout. Other + # DeepSeek-V4(.1) draft layers share the target's KV-cache layout. Other # DSpark architectures may use a different attention kind. - if draft_model_config.hf_config.model_type == "deepseek_v4": + if draft_model_config.hf_config.model_type in ("deepseek_v4", "deepseek_v41"): if target_backend is not None: logger.info_once( "Using the target model's %s attention backend for the " diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index 0dca4588c332..6b0c5fe51d7a 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -112,7 +112,14 @@ def get_eagle3_aux_layers_from_config( if not layer_ids: dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None) if dspark_layer_ids: - layer_ids = [i + 1 for i in dspark_layer_ids] + if getattr(hf_config, "model_type", None) == "deepseek_v41": + # v4.1 reads the attention *inputs* of its target layers, and + # the target model captures the entry stream of layer L when + # idx+1 == L, so the ids are used as-is. (v4's ids are in + # capture-after semantics and keep the +1.) + layer_ids = list(dspark_layer_ids) + else: + layer_ids = [i + 1 for i in dspark_layer_ids] if not layer_ids: # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics. target_layer_ids = getattr(hf_config, "target_layer_ids", None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d6b8ed025c82..2730416f7698 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -813,6 +813,9 @@ def __init__( # Persistent buffers for CUDA graphs. self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) + # [max_num_reqs, depth] prompt token ids preceding each request's chunk + # start, for models exposing `token_lookback_depth`. Sized after load. + self.lookback_token_ids: CpuGpuBuffer | None = None self.positions = torch.zeros( self.max_num_tokens, dtype=torch.int64, device=self.device ) @@ -1065,9 +1068,36 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: ) return self._mamba_bufs - def _init_model_kwargs(self): + def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor: + """Gather, per request, the `depth` prompt token ids preceding its + first scheduled token (column j is position start - 1 - j); -1 where + the position is before the prompt or already past it. Generated + positions are left to the model: under async scheduling the CPU token + table holds placeholders for them.""" + buf = self.lookback_token_ids + assert buf is not None + buf.np.fill(-1) + if num_reqs > 0: + depth = buf.np.shape[1] + starts = self.input_batch.num_computed_tokens_cpu[:num_reqs, None] + pos = starts - np.arange(1, depth + 1) + num_prompt = self.input_batch.num_prompt_tokens[:num_reqs, None] + valid = (pos >= 0) & (pos < num_prompt) + rows = np.arange(num_reqs)[:, None] + ids = self.input_batch.token_ids_cpu[rows, np.clip(pos, 0, None)] + buf.np[:num_reqs] = np.where(valid, ids, -1) + return buf.copy_to_gpu() + + def _init_model_kwargs(self, num_reqs: int | None = None): model_kwargs = dict[str, Any]() + if self.lookback_token_ids is not None: + if num_reqs is None: + num_reqs = self.input_batch.num_reqs + model_kwargs["lookback_token_ids"] = self._prepare_lookback_token_ids( + num_reqs + ) + if not self.is_pooling_model: return model_kwargs @@ -5314,6 +5344,11 @@ def load_model(self, load_dummy_weights: bool = False) -> None: self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.model_config ) + lookback_depth = getattr(self.model, "token_lookback_depth", 0) + if lookback_depth > 0: + self.lookback_token_ids = self._make_buffer( + self.max_num_reqs, lookback_depth, dtype=torch.int32 + ) if self.lora_config: self.model = self.load_lora_model( self.model, self.vllm_config, self.device @@ -6076,7 +6111,7 @@ def _dummy_run( ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - model_kwargs = self._init_model_kwargs() + model_kwargs = self._init_model_kwargs(num_reqs=0) if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) @@ -6087,7 +6122,7 @@ def _dummy_run( elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - model_kwargs = self._init_model_kwargs() + model_kwargs = self._init_model_kwargs(num_reqs=0) else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None From eb7e89443287cc8ec95b8e6fb993bacbd310f9c3 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 11 Sep 2026 04:24:28 -0500 Subject: [PATCH 039/774] [ROCm][CI] Add HY-V4 generation coverage (#55667) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 8 + .buildkite/test_areas/models_basic.yaml | 5 + .buildkite/test_areas/models_distributed.yaml | 12 +- tests/models/test_hyv4_rocm.py | 137 ++++++++++++++++++ 4 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_hyv4_rocm.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index cca01e45ba80..8d2ba218be11 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -516,6 +516,11 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/models/hy_v4/ + - vllm/config/ + - vllm/v1/worker/gpu/ + - vllm/v1/attention/ops/ + - vllm/v1/attention/backend.py - vllm/model_executor/model_loader/sharded_state_loader.py - vllm/model_executor/models/ - vllm/model_executor/layers/ @@ -535,6 +540,7 @@ steps: - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' + - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' #-------------------------------------------------------- mi300 · benchmarks ---------------------------------------------------------# @@ -1978,9 +1984,11 @@ steps: - tests/models/test_terratorch.py - tests/models/transformers/test_backend.py - tests/models/test_registry.py + - tests/models/test_hyv4_rocm.py - tests/models/test_vision.py commands: - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge #----------------------------------------------------- mi300 · models / language -----------------------------------------------------# diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 93d6962868a4..717461261135 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -58,6 +58,11 @@ steps: timeout_in_minutes: 50 depends_on: - image-build-amd + source_file_dependencies: + - tests/models/test_hyv4_rocm.py + commands: + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" key: inkling-unit-tests-b200 diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index cf7623ce3ec8..f735a9b01871 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -42,6 +42,13 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/models/hy_v4/ + - vllm/config/ + - vllm/v1/worker/ + - vllm/v1/attention/ops/ + - vllm/v1/attention/backend.py + - tests/conftest.py + - tests/utils.py - vllm/model_executor/model_loader/sharded_state_loader.py - vllm/model_executor/models/ - vllm/model_executor/layers/ @@ -60,8 +67,9 @@ steps: HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; 1) - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' + pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && + pytest models/language -v -s -m 'distributed(num_gpus=2)' && + VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; 2) pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py diff --git a/tests/models/test_hyv4_rocm.py b/tests/models/test_hyv4_rocm.py new file mode 100644 index 000000000000..b6dc8d70423f --- /dev/null +++ b/tests/models/test_hyv4_rocm.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import pytest + +from tests.utils import multi_gpu_marks +from vllm import SamplingParams +from vllm.config import CompilationConfig +from vllm.config.compilation import CompilationMode, CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.metrics.reader import Counter + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +def _small_hyv4_config(config): + # Apply to both target and MTP draft while retaining MLA/indexer dimensions. + config.update( + { + "hidden_size": 512, + "num_hidden_layers": 2, + "num_attention_heads": 64, + "q_lora_rank": 128, + "intermediate_size": 512, + "moe_intermediate_size": 256, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "enable_ihc": True, + "index_n_heads": 32, + "index_topk": 2048, + "layer_types": ["deepseek_sparse_attention"] * 2, + "indexer_types": ["full", "shared"], + "num_nextn_predict_layers": 1, + "vocab_size": 256, + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + } + ) + return config + + +def _num_captured_graphs(_worker): + from vllm.compilation.counter import compilation_counter + + return compilation_counter.num_cudagraph_captured + + +@pytest.mark.parametrize( + "tp,enforce_eager,mtp,dtype,kv_cache_dtype", + [ + pytest.param(1, True, False, "bfloat16", "auto", id="eager-tp1"), + pytest.param(1, True, False, "float16", "auto", id="eager-fp16"), + pytest.param(1, True, False, "bfloat16", "fp8", id="eager-fp8-kv"), + pytest.param( + 2, + False, + False, + "bfloat16", + "auto", + id="graphs-tp2", + marks=multi_gpu_marks(num_gpus=2), + ), + pytest.param(1, True, True, "bfloat16", "auto", id="eager-mtp"), + ], +) +def test_hyv4_dummy_generation( + vllm_runner, + enable_pickle, + tp: int, + enforce_eager: bool, + mtp: bool, + dtype: str, + kv_cache_dtype: str, +) -> None: + """Exercise prefill, decode, and batch consistency with dummy HY weights.""" + with vllm_runner( + "tencent/Hy4-preview", + trust_remote_code=True, + skip_tokenizer_init=True, + load_format="dummy", + disable_log_stats=not mtp, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + hf_overrides=_small_hyv4_config, + max_model_len=256, + max_num_batched_tokens=128, + max_num_seqs=8, + enable_chunked_prefill=True, + enable_prefix_caching=False, + kv_cache_memory_bytes=128 * 1024 * 1024, + enforce_eager=enforce_eager, + tensor_parallel_size=tp, + speculative_config={"method": "mtp", "num_speculative_tokens": 2} + if mtp + else None, + seed=17, + # Exercise supported decode graphs independently of default graph policy. + compilation_config=CompilationConfig( + mode=CompilationMode.NONE, + cudagraph_mode=CUDAGraphMode.NONE + if enforce_eager + else CUDAGraphMode.FULL_DECODE_ONLY, + cudagraph_capture_sizes=[] if enforce_eager else [1, 2, 4], + ), + ) as model: + if not enforce_eager: + assert model.collective_rpc(_num_captured_graphs) == [3] * tp + prompts = [ + {"prompt_token_ids": [11, 12, 13, 14]}, + {"prompt_token_ids": list(range(20, 84))}, + ] + params = SamplingParams( + temperature=0, max_tokens=4, ignore_eos=True, logprobs=1 + ) + outputs = model.llm.generate(prompts, params) + assert len(outputs) == len(prompts) + for prompt, output in zip(prompts, outputs): + completion = output.outputs[0] + assert len(completion.token_ids) == 4 + assert completion.cumulative_logprob is not None + assert math.isfinite(completion.cumulative_logprob) + single = model.llm.generate([prompt], params)[0].outputs[0] + assert single.token_ids == completion.token_ids + + if mtp: + assert any( + isinstance(metric, Counter) + and metric.name == "vllm:spec_decode_num_drafts" + and metric.value > 0 + for metric in model.llm.get_metrics() + ), "MTP drafter did not run" From 8c1d1c2974ee42757ee2e93cc898932edfd9d265 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 11 Sep 2026 11:35:56 +0200 Subject: [PATCH 040/774] [Misc] Log FlashInfer allreduce workspace init failure as error (#55127) Signed-off-by: NickLucche Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../flashinfer_all_reduce.py | 4 ++-- vllm/logger.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index dd46f782b810..904bf2dae0c1 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -218,7 +218,7 @@ def _get_or_create(be: str): f"with backend={backend}" ) else: - logger.warning_once( + logger.error_once( "Failed to initialize FlashInfer Allreduce norm fusion workspace " f"with backend={backend}" ) @@ -292,7 +292,7 @@ def get_fi_ar_quant_workspace( f"fusion workspace with backend={backend}" ) else: - logger.warning_once( + logger.error_once( "Failed to initialize FlashInfer Allreduce norm quantization " f"fusion workspace with backend={backend}" ) diff --git a/vllm/logger.py b/vllm/logger.py index fde95662f172..2d053f608bb6 100644 --- a/vllm/logger.py +++ b/vllm/logger.py @@ -90,6 +90,12 @@ def _print_warning_once(logger: Logger, msg: str, *args: Hashable) -> None: logger.warning(msg, *args, stacklevel=3) +@lru_cache +def _print_error_once(logger: Logger, msg: str, *args: Hashable) -> None: + # Set the stacklevel to 3 to print the original caller's line info + logger.error(msg, *args, stacklevel=3) + + LogScope = Literal["process", "global", "local"] @@ -144,12 +150,22 @@ def warning_once( return _print_warning_once(self, msg, *args) + def error_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None: + """ + As [`error`][logging.Logger.error], but subsequent calls with + the same message are silently dropped. + """ + if not _should_log_with_scope(scope): + return + _print_error_once(self, msg, *args) + # Pre-defined methods mapping to avoid repeated dictionary creation _METHODS_TO_PATCH = { "debug_once": _VllmLogger.debug_once, "info_once": _VllmLogger.info_once, "warning_once": _VllmLogger.warning_once, + "error_once": _VllmLogger.error_once, } From 79c137c6ada6bf9384e14feb260e10d91bf58c0b Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Fri, 11 Sep 2026 18:19:55 +0800 Subject: [PATCH 041/774] [Bugfix][Rust Frontend][Renderer] Align DeepSeek tool-call arguments with deepseek-recipe (#56260) Co-authored-by: Claude Fable 5.1 Co-authored-by: Bugen Zhao Co-authored-by: Codex Signed-off-by: JinYan Su <751080330@qq.com> Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/deepseek.rs | 23 +++++----- .../chat/src/renderer/deepseek_v4/tests.rs | 46 +++++++++++++++++++ 2 files changed, 57 insertions(+), 12 deletions(-) diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index ab435028ebf3..5f900ab88b59 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -11,7 +11,7 @@ use std::collections::HashMap; use std::fmt::Write as _; use serde::Serialize; -use serde_json::Value; +use serde_json::{Map, Value}; use serde_json_fmt::JsonFormat; use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER; @@ -685,19 +685,18 @@ fn encode_arguments_to_dsml( dialect: DsDialect, ) -> Result<()> { let parameter_tag = dialect.parameter_tag(); - let arguments: Value = serde_json::from_str(&tool_call.arguments).map_err(|error| { - Error::ChatTemplate(format!( - "assistant tool call has invalid JSON arguments for DeepSeek: {error}" - )) - })?; - let Some(arguments) = arguments.as_object() else { - return Err(Error::ChatTemplate( - "assistant tool call arguments for DeepSeek must be a JSON object".to_string(), - )); - }; + // Match deepseek-recipe's render_tool_arguments for both V4 and V4.1: + // https://github.com/deepseek-ai/deepseek-recipe/blob/8cadfede7063c896b944e7bae05daa3549ae97ea/deepseek-recipe-encoding/src/v4/mod.rs#L48-L58 + let arguments = serde_json::from_str::>(&tool_call.arguments) + .unwrap_or_else(|_| { + Map::from_iter([( + "arguments".to_owned(), + Value::String(tool_call.arguments.clone()), + )]) + }); let mut wrote_parameter = false; - for (key, value) in arguments { + for (key, value) in &arguments { if wrote_parameter { out.push('\n'); } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index a0ded7bd992c..172b67463c42 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -9,6 +9,7 @@ use serde_json::Value; use super::DeepSeekV4ChatRenderer; use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::renderer::deepseek_v41::DeepSeekV41ChatRenderer; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ChatMessage, ChatRequest, ChatTool, GenerationPromptMode, ReasoningEffort}; @@ -399,6 +400,51 @@ fn reasoning_effort_template_kwarg_is_ignored() { assert!(rendered.ends_with("<｜Assistant｜>")); } +#[test] +fn tool_call_arguments_preserve_non_objects_like_recipe() { + let renderers: [(&dyn ChatRenderer, &str); 2] = [ + (&DeepSeekV4ChatRenderer::new(), "parameter"), + (&DeepSeekV41ChatRenderer::new(), " parameter"), + ]; + for arguments in [ + "not json", + "[1, 2]", + "null", + "42", + "true", + "", + r#""text""#, + r#""{\"query\":\"double\"}""#, + ] { + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall( + AssistantToolCall { + id: "call".to_string(), + name: "search".to_string(), + arguments: arguments.to_string(), + }, + )]), + ChatMessage::tool_response("result", "call"), + ], + ..ChatRequest::for_test() + }; + + for (renderer, tag) in renderers { + let prompt = renderer.render(&request).unwrap().prompt.into_text().unwrap(); + let parameters: Vec<_> = + prompt.lines().filter(|line| line.contains("parameter name=")).collect(); + assert_eq!( + parameters, + [format!( + "<｜DSML｜{tag} name=\"arguments\" string=\"true\">{arguments}" + )], + "arguments: {arguments:?}, tag: {tag:?}", + ); + } + } +} + #[test] fn tool_results_are_sorted_by_previous_assistant_tool_call_order() { let mut request = ChatRequest { From a2bc2ffb2c8bf46759c4bbeb88e2f4dc32968762 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Fri, 11 Sep 2026 03:26:44 -0700 Subject: [PATCH 042/774] [Bugfix][Pooling] Restore token limits for offline Jina scoring (#56415) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- .../test_jina_ranking_io_processor_unit.py | 87 ++++++++++++++++++- .../pooling/scoring/io_processor.py | 7 ++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py index c2fe62f609df..cb6452469987 100644 --- a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py @@ -1,15 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for JinaRankingIOProcessor online request building.""" +"""Unit tests for JinaRankingIOProcessor request building.""" +from types import SimpleNamespace from unittest.mock import MagicMock import pytest +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast +from vllm import PoolingParams from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor from vllm.entrypoints.pooling.scoring.protocol import RerankRequest from vllm.entrypoints.pooling.scoring.typing import ScoringData +from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext +from vllm.renderers import TokenizeParams pytestmark = pytest.mark.skip_global_cleanup @@ -57,3 +63,82 @@ def _spy_base(self, ctx): assert captured["truncation_side"] == "left" # The real request is restored after delegating. assert ctx.request is request + + +@pytest.fixture +def offline_processor_and_context(): + backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit() + proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + proc.tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) + proc.model_config = SimpleNamespace(max_model_len=1024, is_encoder_decoder=False) + proc.renderer = SimpleNamespace( + default_cmpl_tok_params=TokenizeParams(max_total_tokens=1024) + ) + ctx = OfflineScoringInputsContext( + pooling_task="token_embed", + scoring_data=ScoringData( + data_1=["query alpha beta", "query gamma delta"], + data_2=["document one two", "document three four"], + ), + pooling_params=PoolingParams( + extra_kwargs={ + "chat_template_kwargs": {"instruction": "Keep this instruction"} + } + ), + tokenization_kwargs=None, + chat_template=None, + lora_request=None, + priorities=None, + ) + return proc, ctx + + +@pytest.mark.parametrize( + ("query_limit", "doc_limit", "n_queries"), + [(None, None, 1), (0, 0, 2), (1, None, 1), (None, 2, 1), (1, 2, 1), (1, 2, 2)], + ids=["unset", "zero", "query", "doc", "both-1-to-n", "both-n-to-n"], +) +def test_offline_truncates_text_before_formatting( + offline_processor_and_context, query_limit, doc_limit, n_queries +): + proc, ctx = offline_processor_and_context + queries = ctx.scoring_data.data_1[:n_queries] + docs = ctx.scoring_data.data_2 + ctx.scoring_data = ScoringData(data_1=queries, data_2=docs) + ctx.pooling_params.extra_kwargs.update( + { + name: limit + for name, limit in ( + ("max_tokens_per_query", query_limit), + ("max_tokens_per_doc", doc_limit), + ) + if limit is not None + } + ) + + factory, num_requests = proc.get_request_factory_offline(ctx) + requests = list(factory()) + assert len(requests) == num_requests == n_queries + for i, request in enumerate(requests): + prompt = request["prompts"]["prompt"] + query = " ".join(queries[i].split()[: query_limit or None]) + assert f"\n{query}<|rerank_token|>\n" in prompt + prompt_docs = docs if n_queries == 1 else [docs[i]] + for j, doc in enumerate(prompt_docs): + doc = " ".join(doc.split()[: doc_limit or None]) + assert f'\n{doc}<|embed_token|>\n' in prompt + assert prompt.count("<|embed_token|>") == len(prompt_docs) + assert prompt.count("<|rerank_token|>") == 1 + assert "\nKeep this instruction\n" in prompt + + +@pytest.mark.parametrize("name", ["max_tokens_per_query", "max_tokens_per_doc"]) +@pytest.mark.parametrize("limit", [-1, 1024]) +def test_offline_rejects_invalid_token_limits( + offline_processor_and_context, name, limit +): + proc, ctx = offline_processor_and_context + ctx.pooling_params.extra_kwargs[name] = limit + with pytest.raises(ValueError, match=name): + proc.get_request_factory_offline(ctx) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 4feb40eadb6a..f8cdda20d5d1 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -845,7 +845,14 @@ def get_request_factory_offline( ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineScoringInputsContext) + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + pooling_params=ctx.pooling_params + ) scoring_data = ctx.scoring_data + if max_tokens_per_query > 0 or max_tokens_per_doc > 0: + scoring_data = self._truncate_scoring_data( + scoring_data, max_tokens_per_query, max_tokens_per_doc + ) prompt_extras = ctx.pooling_params.extra_kwargs queries = self.ensure_str(scoring_data.data_1) From b4da4d17ae0c1fb44f7358585f04fdba7e4e1c91 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Fri, 11 Sep 2026 04:44:57 -0700 Subject: [PATCH 043/774] [ROCm][Bugfix] Fix AITER preshuffled FP8 block-scale kernel (#56433) Signed-off-by: Matthew Wong --- .../kernels/linear/scaled_mm/aiter.py | 56 +++++++++++++-- .../layers/attention/mla_attention.py | 7 +- vllm/model_executor/layers/linear.py | 6 +- vllm/models/deepseek_v4/amd/model.py | 17 +++-- vllm/models/deepseek_v4/amd/rocm.py | 70 ++++++++++++++----- 5 files changed, 120 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 45f6d418b969..c9046e89f923 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import ClassVar + import torch from vllm import _custom_ops as ops @@ -462,6 +464,12 @@ def apply_block_scaled_mm( class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight.""" + # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major. + wants_transposed_act_scale: ClassVar[bool] = True + + # process_weights_after_loading shuffles layer.weight to layout (16, 16). + preshuffles_weight: ClassVar[bool] = True + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -519,9 +527,25 @@ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): return True, None + @staticmethod + def _reads_weight_directly(layer: torch.nn.Module) -> bool: + """True when something other than apply_weights consumes layer.weight. + + Such a weight must stay in the plain layout: ``is_bmm`` marks a stack + of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj. + Both are stamped after construction, so can_implement cannot see them. + """ + return bool( + getattr(layer, "is_bmm", False) + or getattr(layer, "skip_weight_relayout", False) + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: super().process_weights_after_loading(layer) + if self._reads_weight_directly(layer): + return + params = FP8BlockParams.from_layer(layer) if params.weight_scale_inv is not None: ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV @@ -565,11 +589,25 @@ def apply_weights( else params.weight_scale_inv ) + # Left unshuffled above; kv_b_proj still reaches here from MLA's + # prefill-context path. + plain = self._reads_weight_directly(layer) + x_2d = x.view(-1, x.shape[-1]) - A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=True) - output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, params.weight, As, Bs, output_dtype=self.config.out_dtype - ) + A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain) + if plain: + output = rocm_aiter_ops.gemm_a8w8_blockscale( + A, + params.weight, + As, + Bs, + list(self.weight_group_shape), + output_dtype=self.config.out_dtype, + ) + else: + output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + A, params.weight, As, Bs, output_dtype=self.config.out_dtype + ) if bias is not None: output = output + bias return output.view(*x.shape[:-1], params.weight.shape[0]) @@ -581,7 +619,11 @@ def apply_block_scaled_mm( As: torch.Tensor, Bs: torch.Tensor, ) -> torch.Tensor: - raise NotImplementedError( - "AiterPreshuffledFp8BlockScaledMMKernel overrides apply_weights and " - "does not use apply_block_scaled_mm." + """Block-scaled GEMM for callers that pre-quantize their activations. + + ``As`` must be column-major; a row-major one will not raise, it just + returns wrong numbers for M > 1. + """ + return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + A, B, As, Bs, output_dtype=self.config.out_dtype ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 34df61fc82b8..25eef7462e69 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -564,10 +564,9 @@ def __init__( ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) - # AMX reads kv_b_proj's weight directly and never calls it live; the - # reference CPU MLA backend calls it but isn't perf-critical. Skip - # the packed-kernel dispatch either way. - kv_b_proj._cpu_skip_gemm_dispatch = True + # MLA reads this weight directly to build W_UK/W_UV, so a backend must + # not relayout it at load time. + kv_b_proj.skip_weight_relayout = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 2e77fac06c34..212dd9537583 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -204,9 +204,9 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): - # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not - # perf-critical, so skip packing and use a plain fallback. - if getattr(layer, "_cpu_skip_gemm_dispatch", False): + # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical, + # so skip packing and use a plain fallback. + if getattr(layer, "skip_weight_relayout", False): layer.cpu_linear = torch.nn.functional.linear return diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 903635ffe751..e9ff6cd80f71 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -71,7 +71,10 @@ make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention +from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLAAttention, + weight_already_preshuffled, +) from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors @@ -147,11 +150,13 @@ def prepare_gateup_preshuffle(self) -> None: return if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - replace_parameter( - self.gate_up_proj, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Skip if the linear's kernel already shuffled it. + if not weight_already_preshuffled(self.gate_up_proj): + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) self._gateup_scale = ws def forward(self, x): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3b1df97dc737..a798053a4974 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -61,6 +61,16 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: return indptr +def weight_already_preshuffled(linear: torch.nn.Module) -> bool: + """True when the linear's kernel already B-preshuffled ``weight``. + + The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped + for those, since shuffle_weight is a permutation rather than an involution. + """ + kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) + return bool(getattr(kernel, "preshuffles_weight", False)) + + def apply_pre_quantized_block_scaled_mm( linear: torch.nn.Module, x_fp8: torch.Tensor, @@ -551,12 +561,13 @@ def _prep(linear) -> torch.Tensor | None: return None if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Shuffle the weight in place (single weight, no unshuffled copy). - replace_parameter( - linear, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Skip if the linear's kernel already shuffled it. + if not weight_already_preshuffled(linear): + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) return ws self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) @@ -654,16 +665,26 @@ def _run_parallel_input_projections( return qr_kv, kv_score, indexer_kv_score, indexer_weights @functools.cached_property - def _wq_b_uses_aiter_block_scaled(self) -> bool: - """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. + def _wq_b_act_scale_transpose(self) -> bool | None: + """Activation-scale byte order every wq_b consumer agrees on, else None. + + None means "do not take the fused norm+quant path". Otherwise the value + is the ``transpose_scale`` the producer must pass so that the scale it + writes matches what the selected GEMM reads. Cached: the linear kernels and the aiter env gates are fixed once the model is built, so this is evaluated at the first forward only. - The fused norm+quant path is only valid if the quant and GEMM it - replaces are exactly the aiter ones; otherwise fall back to the - shared path. + Two conditions, both necessary: + + * The fused norm+quant path is only valid if the quant and GEMM it + replaces are exactly the aiter ones; otherwise fall back to the + shared path. + * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have + different (N, K) and so can resolve to kernels wanting opposite byte + orders. A single producer cannot serve both, so refuse the fast path + when they disagree rather than guessing. """ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.kernels.linear.scaled_mm import ( @@ -671,16 +692,31 @@ def _wq_b_uses_aiter_block_scaled(self) -> bool: ) if not rocm_aiter_ops.is_linear_fp8_enabled(): - return False + return None linears = [self.wq_b] if self.indexer is not None: linears.append(self.indexer.wq_b) + layouts: set[bool] = set() for linear in linears: kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) if not isinstance(kernel, Fp8BlockScaledMMLinearKernel): - return False - return True + return None + layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False))) + if len(layouts) != 1: + logger.warning_once( + "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; " + "disabling the fused q/kv norm+quant path.", + scope="global", + ) + return None + transpose_scale = layouts.pop() + logger.debug_once( + "DeepSeek-V4 wq_b: emitting %s-major activation scales", + "column" if transpose_scale else "row", + scope="global", + ) + return transpose_scale def _split_qkv_and_norm( self, qr_kv: torch.Tensor @@ -697,11 +733,12 @@ def _split_qkv_and_norm( when the aiter linear path is not active. """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + transpose_scale = self._wq_b_act_scale_transpose if not ( qr.dim() == 2 and qr.shape[0] > 0 and self.q_lora_rank % 128 == 0 - and self._wq_b_uses_aiter_block_scaled + and transpose_scale is not None ): return super()._split_qkv_and_norm(qr_kv) @@ -715,7 +752,8 @@ def _split_qkv_and_norm( kv_weight=self.kv_norm.weight.data, kv_epsilon=self.eps, group_size=128, - transpose_scale=False, + # Emit the byte order the wq_b GEMMs read. + transpose_scale=transpose_scale, ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: From 7ef4d9bfed6311e3b78a40abb4a8bb6a2fc741b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicole=20LiHui=20=F0=9F=A5=9C?= Date: Fri, 11 Sep 2026 20:07:04 +0800 Subject: [PATCH 044/774] [Bugfix][Responses] Fix browser.find action type (#55305) Signed-off-by: nicole-lihui Co-authored-by: Chauncey --- .../openai/responses/test_harmony_utils.py | 17 ++++++++++++ .../openai/responses/test_streaming_events.py | 27 +++++++++++++++++++ vllm/entrypoints/openai/responses/harmony.py | 2 +- .../openai/responses/streaming_events.py | 2 +- 4 files changed, 46 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_harmony_utils.py b/tests/entrypoints/openai/responses/test_harmony_utils.py index 14e174dec0bb..54c034de4420 100644 --- a/tests/entrypoints/openai/responses/test_harmony_utils.py +++ b/tests/entrypoints/openai/responses/test_harmony_utils.py @@ -273,6 +273,23 @@ def test_browser_search_recipient_respects_incomplete(self, incomplete): assert output_items[0].action.type == "search" assert output_items[0].action.query == "cursor:weather in San Francisco" + def test_browser_find_recipient_uses_responses_action_type(self): + """browser.find must use the Responses API find_in_page action type.""" + message = ( + Message.from_role_and_content( + Role.ASSISTANT, '{"pattern": "vLLM", "cursor": 42}' + ) + .with_channel("analysis") + .with_recipient("browser.find") + ) + + output_items = harmony_to_response_output(message, frozenset()) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseFunctionWebSearch) + assert output_items[0].action.type == "find_in_page" + assert output_items[0].action.pattern == "vLLM" + def test_commentary_with_empty_content_and_no_recipient(self): """Test edge case: empty commentary with recipient=None.""" message = Message.from_role_and_content(Role.ASSISTANT, "") diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index 6b649aadad8b..e1a8bf46f640 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from openai.types.responses import ResponseFunctionWebSearch +from openai_harmony import Message, Role + from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -8,11 +11,35 @@ ) from vllm.entrypoints.openai.responses.streaming_events import ( SimpleStreamingEventProcessor, + StreamingState, _StateType, + emit_browser_tool_events, split_delta, ) +def test_browser_find_uses_responses_action_type(): + """Both streaming output items must use the Responses find action type.""" + message = ( + Message.from_role_and_content( + Role.ASSISTANT, '{"pattern": "vLLM", "cursor": 42}' + ) + .with_channel("analysis") + .with_recipient("browser.find") + ) + + events = emit_browser_tool_events(message, StreamingState()) + + added, done = events[0], events[-1] + assert added.type == "response.output_item.added" + assert done.type == "response.output_item.done" + for event in (added, done): + assert isinstance(event.item, ResponseFunctionWebSearch) + assert event.item.action.type == "find_in_page" + assert event.item.action.pattern == "vLLM" + assert event.item.action.url == "cursor:42" + + def _make_tool_call( index: int, name: str | None = None, arguments: str | None = None ) -> DeltaToolCall: diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py index 4e62828fd165..461ce1f3fa35 100644 --- a/vllm/entrypoints/openai/responses/harmony.py +++ b/vllm/entrypoints/openai/responses/harmony.py @@ -303,7 +303,7 @@ def _parse_browser_tool_call(message: Message, recipient: str) -> ResponseOutput action = ActionFind( pattern=browser_call.get("pattern", ""), url=f"cursor:{browser_call.get('url', '')}", - type="find", + type="find_in_page", ) else: raise ValueError(f"Unknown browser action: {recipient}") diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index b45264f8c3b4..40abc45b4261 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -667,7 +667,7 @@ def emit_browser_tool_events( ) elif function_name == "find": action = response_function_web_search.ActionFind( - type="find", + type="find_in_page", pattern=parsed_args["pattern"], # TODO: translate to url url=f"cursor:{parsed_args.get('cursor', '')}", From 295ac4e52e8a35772b2a63c028f6510e221a65cf Mon Sep 17 00:00:00 2001 From: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Date: Fri, 11 Sep 2026 20:16:05 +0800 Subject: [PATCH 045/774] [Bugfix][Multimodal] Parse decoded video frame lists as a single video (#55326) Signed-off-by: waizuichougou <2082431897@qq.com> Signed-off-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/multimodal/test_parse.py | 17 +++++++++++++++++ vllm/multimodal/parse.py | 10 ++++++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index c34c8eec8b8f..b4d37b0fe805 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -56,6 +56,23 @@ def test_frame_size_hwc_chw(frame): assert items.get_frame_size(0) == (W, H) +@pytest.mark.parametrize( + "frames", + [ + [np.zeros((H, W, 3), dtype=np.uint8) for _ in range(2)], + [torch.zeros((H, W, 3), dtype=torch.uint8) for _ in range(2)], + ], +) +def test_parse_video_frame_list_as_single_video(frames): + """A list of decoded frames must represent one video item.""" + items = MultiModalDataParser().parse_mm_data({"video": frames})["video"] + + assert items.get_count() == 1 + video = items.get(0) + assert isinstance(video, np.ndarray) + np.testing.assert_array_equal(video, np.stack([np.asarray(f) for f in frames])) + + @pytest.mark.parametrize( "modality,processor_cls", [ diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index fefaeef2b019..124a8845cccc 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -751,8 +751,14 @@ def _parse_video_data( return VideoEmbeddingItems(data, self.expected_hidden_size) data_items: list[VideoItem] - if (is_list_of(data, PILImage.Image) and len(data) > 0) or ( - isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4 + if ( + (is_list_of(data, PILImage.Image) and len(data) > 0) + or ( + is_list_of(data, (np.ndarray, torch.Tensor), check="all") + and len(data) > 0 + and all(item.ndim == 3 for item in data) + ) + or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4) ): data_items = [data] elif isinstance(data, (np.ndarray, torch.Tensor)): From 89dbb2644552d6e473a7e97da0ce8f0aa8e32c9d Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang Date: Fri, 11 Sep 2026 20:31:10 +0800 Subject: [PATCH 046/774] [Bugfix] Fix GLM-OCR MTP position masking during CUDA graph capture (#56447) Signed-off-by: zRzRzRzRzRzRzR --- vllm/model_executor/models/glm_ocr_mtp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py index 6c8e208c4b5c..05c3ac793578 100644 --- a/vllm/model_executor/models/glm_ocr_mtp.py +++ b/vllm/model_executor/models/glm_ocr_mtp.py @@ -86,7 +86,8 @@ def forward( ) -> torch.Tensor: assert inputs_embeds is not None # masking inputs at position 0, as not needed by MTP - inputs_embeds[positions[0] == 0] = 0 + token_positions = positions[0] if positions.ndim == 2 else positions + inputs_embeds.masked_fill_((token_positions == 0).unsqueeze(-1), 0) inputs_embeds = self.enorm(inputs_embeds) previous_hidden_states = self.hnorm(previous_hidden_states) From 912dfb37581b98c0b6eaeca5758d1e1462b7feac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EA=B9=80=EC=9E=AC=EC=96=B5?= Date: Fri, 11 Sep 2026 21:52:45 +0900 Subject: [PATCH 047/774] [Bugfix][Scoring] Warn when serving original Qwen3 reranker without chat template (#56017) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 김재억 Co-authored-by: Claude Opus 5 (1M context) --- .../entrypoints/pooling/test_io_processor.py | 147 ++++++++++++++++++ .../pooling/scoring/io_processor.py | 21 +++ 2 files changed, 168 insertions(+) diff --git a/tests/entrypoints/pooling/test_io_processor.py b/tests/entrypoints/pooling/test_io_processor.py index 6c64966e588e..9d2207fdf15d 100644 --- a/tests/entrypoints/pooling/test_io_processor.py +++ b/tests/entrypoints/pooling/test_io_processor.py @@ -74,3 +74,150 @@ def test_rejects_conflicting_pooling_task(processor: PoolingIOProcessor): ) assert exc_info.value.parameter is None assert exc_info.value.value is None + + +def test_qwen3_reranker_warns_without_chat_template(monkeypatch): + from unittest.mock import MagicMock + + from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor + + monkeypatch.setattr( + "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock() + ) + monkeypatch.setattr( + "vllm.model_executor.models.interfaces.supports_score_template", + lambda *_: False, + ) + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer", + lambda *_: False, + ) + + mock_logger = MagicMock() + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.logger", + mock_logger, + ) + + model_config = MagicMock() + model_config.architecture = "Qwen3ForSequenceClassification" + model_config.model = "Qwen/Qwen3-Reranker-0.6B" + model_config.hf_config.is_original_qwen3_reranker = True + model_config.use_sep_token = False + model_config.is_multimodal_model = False + model_config.max_model_len = 8192 + + vllm_config = MagicMock(model_config=model_config) + tokenizer = MagicMock() + tokenizer.pad_token_id = 0 + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + chat_template_config = MagicMock(chat_template=None) + + CrossEncoderIOProcessor( + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + mock_logger.warning.assert_called_once() + warning_call = mock_logger.warning.call_args + assert "Qwen/Qwen3-Reranker-0.6B" in warning_call[0][1] + assert "qwen3_reranker.jinja" in warning_call[0][2] + + +def test_qwen3_vl_reranker_warns_with_vl_template(monkeypatch): + from unittest.mock import MagicMock + + from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor + + monkeypatch.setattr( + "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock() + ) + monkeypatch.setattr( + "vllm.model_executor.models.interfaces.supports_score_template", + lambda *_: False, + ) + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer", + lambda *_: False, + ) + + mock_logger = MagicMock() + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.logger", + mock_logger, + ) + + model_config = MagicMock() + model_config.architecture = "Qwen3VLForSequenceClassification" + model_config.model = "Qwen/Qwen3-VL-Reranker-2B" + model_config.hf_config.is_original_qwen3_reranker = True + model_config.use_sep_token = False + model_config.is_multimodal_model = True + model_config.max_model_len = 8192 + + vllm_config = MagicMock(model_config=model_config) + tokenizer = MagicMock() + tokenizer.pad_token_id = 0 + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + chat_template_config = MagicMock(chat_template=None) + + CrossEncoderIOProcessor( + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + mock_logger.warning.assert_called_once() + warning_call = mock_logger.warning.call_args + assert "Qwen/Qwen3-VL-Reranker-2B" in warning_call[0][1] + assert "qwen3_vl_reranker.jinja" in warning_call[0][2] + + +def test_qwen3_reranker_no_warning_when_template_provided(monkeypatch): + from unittest.mock import MagicMock + + from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor + + monkeypatch.setattr( + "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock() + ) + monkeypatch.setattr( + "vllm.model_executor.models.interfaces.supports_score_template", + lambda *_: False, + ) + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer", + lambda *_: False, + ) + + mock_logger = MagicMock() + monkeypatch.setattr( + "vllm.entrypoints.pooling.scoring.io_processor.logger", + mock_logger, + ) + + model_config = MagicMock() + model_config.architecture = "Qwen3ForSequenceClassification" + model_config.model = "Qwen/Qwen3-Reranker-0.6B" + model_config.hf_config.is_original_qwen3_reranker = True + model_config.use_sep_token = False + model_config.is_multimodal_model = False + model_config.max_model_len = 8192 + + vllm_config = MagicMock(model_config=model_config) + tokenizer = MagicMock() + tokenizer.pad_token_id = 0 + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + chat_template_config = MagicMock(chat_template="custom template") + + CrossEncoderIOProcessor( + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + mock_logger.warning.assert_not_called() diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index f8cdda20d5d1..5e6104eedd22 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -7,6 +7,7 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt +from vllm.logger import init_logger from vllm.renderers import TokenizeParams from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( @@ -47,6 +48,8 @@ validate_score_input, ) +logger = init_logger(__name__) + ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest] @@ -425,6 +428,24 @@ def __init__(self, *args, **kwargs): self.model = model if self.supports_score_template else None self.use_sep_token = self.model_config.use_sep_token + if ( + getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False) + and self.chat_template is None + ): + suggested_template = ( + "examples/pooling/score/template/qwen3_vl_reranker.jinja" + if self.model_config.is_multimodal_model + else "examples/pooling/score/template/qwen3_reranker.jinja" + ) + logger.warning( + "Serving an original Qwen3 reranker (%s) without a " + "--chat-template. The model was trained with a specific prompt " + "template; running without it may produce inaccurate relevance " + "scores. Consider specifying `--chat-template %s`.", + self.model_config.model, + suggested_template, + ) + ####################################### # online APIs From a5b714eee4ef09c286368c53df28443b7ac4cf4a Mon Sep 17 00:00:00 2001 From: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> Date: Fri, 11 Sep 2026 20:58:15 +0800 Subject: [PATCH 048/774] [Bugfix] Fix Qwen3-VL and Cosmos3-Edge text architectures for CPU and pipeline parallelism (#53699) Signed-off-by: xiaoming <1259730330@qq.com> --- vllm/model_executor/models/cosmos3_edge.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py index 8c5714248b69..38e55405b895 100644 --- a/vllm/model_executor/models/cosmos3_edge.py +++ b/vllm/model_executor/models/cosmos3_edge.py @@ -575,7 +575,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): with self._mark_language_model(vllm_config): self.language_model = Cosmos3EdgeForCausalLM( - vllm_config=vllm_config, + vllm_config=vllm_config.with_hf_config( + config.text_config, architectures=["NemotronHForCausalLM"] + ), prefix=maybe_prefix(prefix, "language_model"), ) From c191787a6861868069bc4f6ed6f842af541de23a Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:01:38 +0100 Subject: [PATCH 049/774] [Bugfix] Align vLLM YaRN with Transformers and stop re-scaling max_model_len (#56446) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/test_config.py | 48 ++++++++++++++++ tests/transformers_utils/test_config.py | 26 +++++++++ vllm/config/model.py | 15 +++-- .../layers/rotary_embedding/__init__.py | 40 +------------ .../layers/rotary_embedding/common.py | 4 +- .../rotary_embedding/deepseek_scaling_rope.py | 12 +--- .../layers/rotary_embedding/mrope.py | 17 ++++-- .../telechat3_scaling_rope.py | 57 ------------------- .../rotary_embedding/yarn_scaling_rope.py | 29 +++++----- vllm/transformers_utils/config.py | 9 +++ vllm/transformers_utils/configs/mistral.py | 10 ++-- 11 files changed, 133 insertions(+), 134 deletions(-) delete mode 100644 vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py diff --git a/tests/test_config.py b/tests/test_config.py index 827dbcd4f8c8..3982507b6b63 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1853,6 +1853,54 @@ def test_get_and_verify_max_len_with_nope_layers( assert hf_config.rope_parameters["full_attention"] is None +@pytest.mark.parametrize( + ("rope_type", "factor", "expected_max_len"), + [ + # TeleChat3-36B-Thinking: 32768 already scaled from 8192 by 4 + ("yarn", 4.0, 32768), + # sarvam-105b: declares factor 40 but only serves 131072 of it + ("deepseek_yarn", 40.0, 32768), + ("deepseek_llama_scaling", 40.0, 32768), + # Non-YaRN scaling still multiplies + ("linear", 4.0, 131072), + ], +) +def test_get_and_verify_max_len_yarn_is_already_scaled( + rope_type, factor, expected_max_len +): + """YaRN variants must not re-apply `factor` to max_position_embeddings. + + Transformers treats max_position_embeddings as the final context length + for every YaRN variant, so scaling it again overstates the limit and lets + requests past the end of the cos/sin cache. + """ + from transformers import PretrainedConfig + + from vllm.config.model import _get_and_verify_max_len + from vllm.transformers_utils.model_arch_config_convertor import ( + ModelArchConfigConvertorBase, + ) + + hf_config = PretrainedConfig(max_position_embeddings=32768) + hf_config.rope_parameters = { + "rope_type": rope_type, + "factor": factor, + "original_max_position_embeddings": 8192, + } + model_arch_config = ModelArchConfigConvertorBase(hf_config, hf_config).convert() + + actual_max_len = _get_and_verify_max_len( + hf_config=hf_config, + model_arch_config=model_arch_config, + tokenizer_config=None, + max_model_len=None, + disable_sliding_window=False, + sliding_window=None, + ) + + assert actual_max_len == expected_max_len + + class MockConfig: """Simple mock object for testing maybe_pull_model_tokenizer_for_runai""" diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index ba1f0ef8ddd8..23a0f44a5dc4 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -6,6 +6,7 @@ `BaseRenderer.get_eos_token_id`. """ +import math from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock, patch @@ -52,6 +53,31 @@ def test_patch_legacy_rope_type_preserves_nope_layers(): } +def test_patch_legacy_rope_type_normalizes_telechat3_yarn(): + """TeleChat3's RoPE is YaRN with 0.07 in place of the usual 0.1. + + Encoding that as a precomputed attention_factor keeps the config + plain YaRN, which Transformers and every "yarn" guard understand. + `mscale` cannot express it: Transformers only applies mscale when + mscale_all_dim is also truthy. + """ + rope_parameters = { + "type": "telechat3-yarn", + "rope_type": "telechat3-yarn", + "factor": 4.0, + "original_max_position_embeddings": 8192, + } + + patch_legacy_rope_type(rope_parameters) + + assert rope_parameters == { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 8192, + "attention_factor": pytest.approx(0.07 * math.log(4.0) + 1.0), + } + + def test_glm5_next_accepts_deepseek_sparse_attention_layers(): layer_types = ["linear_attention", "deepseek_sparse_attention"] diff --git a/vllm/config/model.py b/vllm/config/model.py index e47f3eee96f0..89978c35d051 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2465,13 +2465,20 @@ def _get_and_verify_max_len( # loading HF config rope_type = rp["rope_type"] - if rope_type not in ("su", "longrope", "llama3"): + # YaRN variants leave max_position_embeddings already scaled, as + # Transformers' _compute_yarn_parameters assumes, so `factor` must + # not be applied to it again. + if rope_type not in ( + "su", + "longrope", + "llama3", + "yarn", + "deepseek_yarn", + "deepseek_llama_scaling", + ): # NOTE: rope_type == "default" does not define factor https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py # NOTE: This assumes all layer types have the same scaling factor. scaling_factor = rp.get("factor", scaling_factor) - - if rope_type == "yarn": - derived_max_model_len = rp["original_max_position_embeddings"] if scaling_factor is None: # Fallback the factor to 1.0 if a user assigned `null` logger.warning_once( diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index ab06448957fb..e23a5128b8a1 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -24,7 +24,6 @@ from .mrope_interleaved import MRotaryEmbeddingInterleaved from .ntk_scaling_rope import NTKScalingRotaryEmbedding from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding -from .telechat3_scaling_rope import TeleChat3RoPEScaledRotaryEmbedding from .yarn_scaling_rope import YaRNScalingRotaryEmbedding _ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {} @@ -246,16 +245,15 @@ def get_rope( for k, v in rope_parameters.items() if k in ( - "extrapolation_factor", - "attn_factor", "beta_fast", "beta_slow", - "apply_yarn_scaling", + "mscale", + "mscale_all_dim", + "attention_factor", "truncate", ) } if "mrope_section" in rope_parameters: - extra_kwargs.pop("apply_yarn_scaling", None) rotary_emb = MRotaryEmbedding( head_size, rotary_dim, @@ -288,8 +286,6 @@ def get_rope( for k, v in rope_parameters.items() if k in ( - "extrapolation_factor", - "attn_factor", "beta_fast", "beta_slow", "mscale", @@ -346,36 +342,6 @@ def get_rope( ) else: raise ValueError("Pangu mrope lacks necessary parameters.") - elif scaling_type == "telechat3-yarn": - scaling_factor = rope_parameters["factor"] - if "original_max_position_embeddings" in rope_parameters: - original_max_position = rope_parameters["original_max_position_embeddings"] - scaling_factor = max_position / original_max_position - else: - original_max_position = max_position - extra_kwargs = { - k: v - for k, v in rope_parameters.items() - if k - in ( - "extrapolation_factor", - "attn_factor", - "beta_fast", - "beta_slow", - "mscale", - "mscale_all_dim", - ) - } - rotary_emb = TeleChat3RoPEScaledRotaryEmbedding( - head_size, - rotary_dim, - original_max_position, - base, - is_neox_style, - scaling_factor, - dtype, - **extra_kwargs, - ) else: raise ValueError(f"Unknown RoPE scaling type {scaling_type}") _ROPE_DICT[key] = rotary_emb diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py index 0a392812b6bc..006dba17ec48 100644 --- a/vllm/model_executor/layers/rotary_embedding/common.py +++ b/vllm/model_executor/layers/rotary_embedding/common.py @@ -70,10 +70,10 @@ def yarn_linear_ramp_mask( return ramp_func -def yarn_get_mscale(scale: float = 1) -> float: +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: if scale <= 1: return 1.0 - return 0.1 * math.log(scale) + 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 def _flashinfer_rotary_embedding( diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py index bfaf81ad0077..f5e6472e235e 100644 --- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py @@ -39,8 +39,6 @@ def __init__( scaling_factor: float, dtype: torch.dtype, *, - extrapolation_factor: float = 1, - attn_factor: float = 1, beta_fast: int = 32, beta_slow: int = 1, mscale: float = 1, @@ -48,15 +46,12 @@ def __init__( init_cache: bool = True, ) -> None: self.scaling_factor = scaling_factor - self.extrapolation_factor = extrapolation_factor - self.attn_factor = attn_factor self.beta_fast = beta_fast self.beta_slow = beta_slow # Get n-d magnitude scaling corrected for interpolation. self.mscale = float( yarn_get_mscale(self.scaling_factor, float(mscale)) / yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) - * attn_factor ) self.use_flashinfer = ( self.enabled() @@ -96,10 +91,9 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: self.max_position_embeddings, ) # Get n-d rotational scaling corrected for extrapolation - inv_freq_mask = ( - 1 - - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float) - ) * self.extrapolation_factor + inv_freq_mask = 1 - yarn_linear_ramp_mask( + low, high, self.rotary_dim // 2, dtype=torch.float + ) inv_freq = ( inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 2cbffab8ddce..a7fe51c46598 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -252,21 +252,28 @@ def __init__( # YaRN parameters. *, scaling_factor: float | None = None, - extrapolation_factor: float = 1, - attn_factor: float = 1, beta_fast: int = 32, beta_slow: int = 1, + mscale: float | None = None, + mscale_all_dim: float | None = None, + attention_factor: float | None = None, truncate: bool = True, ) -> None: self.scaling_factor = scaling_factor - self.extrapolation_factor = extrapolation_factor - self.attn_factor = attn_factor self.beta_fast = beta_fast self.beta_slow = beta_slow self.truncate = truncate if self.scaling_factor is not None: # Get n-d magnitude scaling corrected for interpolation - self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor) + if attention_factor is not None: + self.mscale = float(attention_factor) + elif mscale and mscale_all_dim: + self.mscale = float( + yarn_get_mscale(self.scaling_factor, mscale) + / yarn_get_mscale(self.scaling_factor, mscale_all_dim) + ) + else: + self.mscale = float(yarn_get_mscale(self.scaling_factor)) else: self.mscale = 1.0 diff --git a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py deleted file mode 100644 index dd2fb9c320b6..000000000000 --- a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import math - -import torch - -from .base import RotaryEmbedding -from .yarn_scaling_rope import YaRNScalingRotaryEmbedding - - -class TeleChat3RoPEScaledRotaryEmbedding(YaRNScalingRotaryEmbedding): - """TeleChat3 uses a variant of YaRN method. - - To achieve code reuse as much as possible, we have rewritten the - `get_mscale` method in the initialization function - """ - - def __init__( - self, - head_size: int, - rotary_dim: int, - max_position_embeddings: int, - base: int, - is_neox_style: bool, - scaling_factor: float, - dtype: torch.dtype, - *, - extrapolation_factor: float = 1, - attn_factor: float = 1, - beta_fast: int = 32, - beta_slow: int = 1, - truncate: bool = True, - ) -> None: - self.scaling_factor = scaling_factor - self.extrapolation_factor = extrapolation_factor - self.attn_factor = attn_factor - self.beta_fast = beta_fast - self.beta_slow = beta_slow - self.truncate = truncate - - def get_mscale(scale, mscale=1): - if scale <= 1: - return 1.0 - return 0.07 * mscale * math.log(scale) + 1.0 - - self.mscale = float(get_mscale(self.scaling_factor) * attn_factor) - # Initialization must be performed after mscale, otherwise mscale is useless - RotaryEmbedding.__init__( - self, - head_size, - rotary_dim, - max_position_embeddings, - base, - is_neox_style, - dtype, - ) diff --git a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py index f01ca1e23121..01fbf0b2c5c2 100644 --- a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py @@ -23,25 +23,27 @@ def __init__( scaling_factor: float, dtype: torch.dtype, *, - extrapolation_factor: float = 1, - attn_factor: float = 1, beta_fast: int = 32, beta_slow: int = 1, - apply_yarn_scaling: bool = True, + mscale: float | None = None, + mscale_all_dim: float | None = None, + attention_factor: float | None = None, truncate: bool = True, ) -> None: self.scaling_factor = scaling_factor - self.extrapolation_factor = extrapolation_factor - self.attn_factor = attn_factor self.beta_fast = beta_fast self.beta_slow = beta_slow self.truncate = truncate # Get n-d magnitude scaling corrected for interpolation - self.mscale = ( - float(yarn_get_mscale(self.scaling_factor) * attn_factor) - if apply_yarn_scaling - else float(attn_factor) - ) + if attention_factor is not None: + self.mscale = float(attention_factor) + elif mscale and mscale_all_dim: + self.mscale = float( + yarn_get_mscale(self.scaling_factor, mscale) + / yarn_get_mscale(self.scaling_factor, mscale_all_dim) + ) + else: + self.mscale = float(yarn_get_mscale(self.scaling_factor)) super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype ) @@ -62,10 +64,9 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: self.truncate, ) # Get n-d rotational scaling corrected for extrapolation - inv_freq_mask = ( - 1 - - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float) - ) * self.extrapolation_factor + inv_freq_mask = 1 - yarn_linear_ramp_mask( + low, high, self.rotary_dim // 2, dtype=torch.float + ) inv_freq = ( inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 93e3ad4de88d..d34e391a2626 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math import os from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager @@ -539,6 +540,14 @@ def _patch_legacy_rope_type(rope_parameters: dict[str, Any]) -> None: ) rope_parameters["rope_type"] = "default" logger.warning("Replacing legacy rope_type 'mrope' with 'default'") + elif rope_parameters["rope_type"] == "telechat3-yarn": + # TeleChat3 is YaRN with 0.07 in place of YaRN's 0.1 attention + # scaling coefficient. Precompute it so the config is plain YaRN. + factor = rope_parameters["factor"] + rope_parameters["rope_type"] = "yarn" + rope_parameters["attention_factor"] = 0.07 * math.log(factor) + 1.0 + rope_parameters.pop("type", None) + logger.warning("Replacing rope_type 'telechat3-yarn' with 'yarn'") # Handle nested rope_parameters in interleaved sliding attention models if is_rope_parameters_nested(rope_parameters): diff --git a/vllm/transformers_utils/configs/mistral.py b/vllm/transformers_utils/configs/mistral.py index 2b0796691472..3ec94eb0bec8 100644 --- a/vllm/transformers_utils/configs/mistral.py +++ b/vllm/transformers_utils/configs/mistral.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any -from packaging.version import Version from transformers import PretrainedConfig, WhisperConfig -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.logger import init_logger @@ -119,7 +117,6 @@ def _remap_mistral_yarn_args(config: dict) -> dict: "original_max_position_embeddings": ("original_max_position_embeddings", int), "beta": ("beta_fast", float), "alpha": ("beta_slow", float), - "apply_scale": ("apply_yarn_scaling", bool), } yarn_config = config.get("yarn") or {} @@ -136,9 +133,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict: # Cast to remove Transformers > v5 type warnings config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name)) - # Ignore apply_yarn_scaling in Transformers > v5 RoPE validation to remove warnings - if Version(TRANSFORMERS_VERSION) >= Version("5.3.0.dev0"): - config["ignore_keys_at_rope_validation"] = {"apply_yarn_scaling"} + # `apply_scale: false` means no magnitude correction, which Transformers + # spells as an explicit attention_factor of 1. + if not yarn_config.pop("apply_scale", True): + config["rope_parameters"]["attention_factor"] = 1.0 assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}" From d5a9d0f59bbbbb55ac347e6126d01b1456dbd94f Mon Sep 17 00:00:00 2001 From: Guy Stone Date: Fri, 11 Sep 2026 15:12:36 +0100 Subject: [PATCH 050/774] [Metrics] Consolidate Prometheus histogram bucket defaults into a single module (#48866) Signed-off-by: Guy Stone Co-authored-by: Claude Fable 5 --- tests/v1/metrics/test_histogram_buckets.py | 273 +++++++++++++++++++++ vllm/v1/metrics/buckets.py | 228 +++++++++++++++++ vllm/v1/metrics/loggers.py | 170 ++----------- 3 files changed, 524 insertions(+), 147 deletions(-) create mode 100644 tests/v1/metrics/test_histogram_buckets.py create mode 100644 vllm/v1/metrics/buckets.py diff --git a/tests/v1/metrics/test_histogram_buckets.py b/tests/v1/metrics/test_histogram_buckets.py new file mode 100644 index 000000000000..54cc31da94fa --- /dev/null +++ b/tests/v1/metrics/test_histogram_buckets.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Snapshot tests locking the default Prometheus histogram buckets. + +The expected lists are deliberately hard-coded literals rather than imports: +any change to a default bucket boundary must show up as a failing test here. +""" + +import prometheus_client +import pytest + +from vllm.config import ModelConfig, VllmConfig +from vllm.config.observability import ObservabilityConfig +from vllm.v1.metrics.buckets import ( + BUCKET_FAMILY_KEYS, + build_1_2_5_buckets, + histogram_buckets, +) +from vllm.v1.metrics.loggers import PrometheusStatLogger +from vllm.v1.metrics.prometheus import unregister_vllm_metrics + +pytestmark = pytest.mark.cpu_test + +DEFAULT_BUCKET_SNAPSHOTS: dict[str, list[float]] = { + "request_latency": [ + 0.3, + 0.5, + 0.8, + 1.0, + 1.5, + 2.0, + 2.5, + 5.0, + 10.0, + 15.0, + 20.0, + 30.0, + 40.0, + 50.0, + 60.0, + 120.0, + 240.0, + 480.0, + 960.0, + 1920.0, + 7680.0, + ], + "time_to_first_token": [ + 0.001, + 0.005, + 0.01, + 0.02, + 0.04, + 0.06, + 0.08, + 0.1, + 0.25, + 0.5, + 0.75, + 1.0, + 2.5, + 5.0, + 7.5, + 10.0, + 20.0, + 40.0, + 80.0, + 160.0, + 640.0, + 2560.0, + ], + "inter_token_latency": [ + 0.01, + 0.025, + 0.05, + 0.075, + 0.1, + 0.15, + 0.2, + 0.3, + 0.4, + 0.5, + 0.75, + 1.0, + 2.5, + 5.0, + 7.5, + 10.0, + 20.0, + 40.0, + 80.0, + ], + "iteration_tokens": [ + 1, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + ], + "request_params_n": [1, 2, 5, 10, 20], + "request_num_preemptions": [1, 2, 3, 4, 5, 10, 20], + "kv_cache_residency": [ + 0.001, + 0.002, + 0.005, + 0.01, + 0.02, + 0.05, + 0.1, + 0.2, + 0.5, + 1, + 2, + 5, + 10, + 20, + 30, + 60, + 120, + 300, + 600, + 1200, + 1800, + ], +} + +METRIC_FAMILIES: dict[str, str] = { + "vllm:e2e_request_latency_seconds": "request_latency", + "vllm:request_queue_time_seconds": "request_latency", + "vllm:request_inference_time_seconds": "request_latency", + "vllm:request_prefill_time_seconds": "request_latency", + "vllm:request_decode_time_seconds": "request_latency", + "vllm:time_to_first_token_seconds": "time_to_first_token", + "vllm:inter_token_latency_seconds": "inter_token_latency", + "vllm:request_time_per_output_token_seconds": "inter_token_latency", + "vllm:iteration_tokens_total": "iteration_tokens", + "vllm:request_params_n": "request_params_n", + "vllm:request_num_preemptions": "request_num_preemptions", + "vllm:request_prompt_tokens": "request_tokens", + "vllm:request_generation_tokens": "request_tokens", + "vllm:request_max_num_generation_tokens": "request_tokens", + "vllm:request_params_max_tokens": "request_tokens", + "vllm:request_prefill_kv_computed_tokens": "request_tokens", + "vllm:kv_block_lifetime_seconds": "kv_cache_residency", + "vllm:kv_block_idle_before_evict_seconds": "kv_cache_residency", + "vllm:kv_block_reuse_gap_seconds": "kv_cache_residency", +} + +TEST_MODEL = "distilbert/distilgpt2" +TEST_MODEL_MAX_LEN = 1024 +REQUEST_TOKENS_SNAPSHOT_FOR_TEST_MODEL: list[float] = [ + 1, + 2, + 5, + 10, + 20, + 50, + 100, + 200, + 500, + 1000, +] + + +@pytest.mark.parametrize("family", sorted(DEFAULT_BUCKET_SNAPSHOTS)) +def test_default_bucket_snapshots(family): + """Static families must expose exactly today's default boundaries.""" + assert histogram_buckets(family) == DEFAULT_BUCKET_SNAPSHOTS[family] + + +def test_request_tokens_bucket_snapshot(): + """The 1-2-5 token-count series must be stable for a given model len.""" + assert histogram_buckets("request_tokens", max_model_len=32768) == [ + 1, + 2, + 5, + 10, + 20, + 50, + 100, + 200, + 500, + 1000, + 2000, + 5000, + 10000, + 20000, + ] + assert build_1_2_5_buckets(100) == [1, 2, 5, 10, 20, 50, 100] + + +def test_request_tokens_requires_max_model_len(): + with pytest.raises(ValueError, match="max_model_len is required"): + histogram_buckets("request_tokens") + + +def test_buckets_return_fresh_copies(): + """Mutating a returned list must not corrupt the shared defaults.""" + first = histogram_buckets("request_latency") + first.append(1e9) + second = histogram_buckets("request_latency") + assert second == DEFAULT_BUCKET_SNAPSHOTS["request_latency"] + assert second is not first + + +def test_bucket_family_keys(): + """The declared families are exactly these, and each one feeds a histogram.""" + assert { + "request_latency", + "time_to_first_token", + "inter_token_latency", + "iteration_tokens", + "request_params_n", + "request_num_preemptions", + "request_tokens", + "kv_cache_residency", + } == BUCKET_FAMILY_KEYS + assert set(METRIC_FAMILIES.values()) == BUCKET_FAMILY_KEYS + + +def collect_histogram_buckets() -> dict[str, list[float]]: + """Read back each vllm:* histogram's finite bucket bounds by metric name.""" + found: dict[str, list[float]] = {} + for metric in prometheus_client.REGISTRY.collect(): + if metric.name.startswith("vllm:") and metric.type == "histogram": + found[metric.name] = sorted( + { + float(sample.labels["le"]) + for sample in metric.samples + if sample.name.endswith("_bucket") and sample.labels["le"] != "+Inf" + } + ) + return found + + +def build_logger_config( + observability_config: ObservabilityConfig, +) -> VllmConfig: + return VllmConfig( + model_config=ModelConfig(model=TEST_MODEL), + observability_config=observability_config, + ) + + +def test_prometheus_logger_default_buckets(): + """Every core engine histogram must use its family's defaults. + + The config deliberately sets no kv_transfer_config: a KV connector would + register its own histograms, which own their boundaries and are not + covered by METRIC_FAMILIES. + """ + config = build_logger_config(ObservabilityConfig(kv_cache_metrics=True)) + assert config.model_config.max_model_len == TEST_MODEL_MAX_LEN + try: + PrometheusStatLogger(config) + found = collect_histogram_buckets() + assert set(found) == set(METRIC_FAMILIES) + for metric_name, family in METRIC_FAMILIES.items(): + if family == "request_tokens": + expected = REQUEST_TOKENS_SNAPSHOT_FOR_TEST_MODEL + else: + expected = DEFAULT_BUCKET_SNAPSHOTS[family] + assert found[metric_name] == [float(b) for b in expected], metric_name + finally: + unregister_vllm_metrics() diff --git a/vllm/v1/metrics/buckets.py b/vllm/v1/metrics/buckets.py new file mode 100644 index 000000000000..375d73072924 --- /dev/null +++ b/vllm/v1/metrics/buckets.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Default bucket boundaries for the Prometheus histograms emitted by vLLM. + +Every core engine histogram created by the Prometheus stat logger draws its +bucket boundaries from exactly one of the families defined here, so each +default list has a single source of truth. Histograms owned by other +subsystems, such as the KV connector and offloading metrics, keep their own +boundaries. +""" + +from typing import Literal, get_args + +BucketFamilyKey = Literal[ + "request_latency", + "time_to_first_token", + "inter_token_latency", + "iteration_tokens", + "request_params_n", + "request_num_preemptions", + "request_tokens", + "kv_cache_residency", +] +"""Canonical keys for the histogram bucket families.""" + +BUCKET_FAMILY_KEYS: frozenset[str] = frozenset(get_args(BucketFamilyKey)) +"""All bucket family keys, for membership checks.""" + +REQUEST_LATENCY_BUCKETS: tuple[float, ...] = ( + 0.3, + 0.5, + 0.8, + 1.0, + 1.5, + 2.0, + 2.5, + 5.0, + 10.0, + 15.0, + 20.0, + 30.0, + 40.0, + 50.0, + 60.0, + 120.0, + 240.0, + 480.0, + 960.0, + 1920.0, + 7680.0, +) +"""Sub-second scheduling delays through multi-hour batch requests; shared by +the accumulated request phase-timing histograms (e2e, queue, inference, +prefill, decode).""" + +TIME_TO_FIRST_TOKEN_BUCKETS: tuple[float, ...] = ( + 0.001, + 0.005, + 0.01, + 0.02, + 0.04, + 0.06, + 0.08, + 0.1, + 0.25, + 0.5, + 0.75, + 1.0, + 2.5, + 5.0, + 7.5, + 10.0, + 20.0, + 40.0, + 80.0, + 160.0, + 640.0, + 2560.0, +) +"""Millisecond-scale prefill for tiny prompts up to ~40-minute worst cases, +with the densest resolution around interactive (sub-second) latencies.""" + +INTER_TOKEN_LATENCY_BUCKETS: tuple[float, ...] = ( + 0.01, + 0.025, + 0.05, + 0.075, + 0.1, + 0.15, + 0.2, + 0.3, + 0.4, + 0.5, + 0.75, + 1.0, + 2.5, + 5.0, + 7.5, + 10.0, + 20.0, + 40.0, + 80.0, +) +"""Per-decode-step latencies: 10 ms fast decode up to multi-second stalls.""" + +ITERATION_TOKENS_BUCKETS: tuple[float, ...] = ( + 1, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, +) +"""Tokens processed per engine step: powers of two up to a typical +max-num-batched-tokens budget.""" + +REQUEST_PARAMS_N_BUCKETS: tuple[float, ...] = (1, 2, 5, 10, 20) +"""Small integer counts for the ``n`` sampling parameter.""" + +REQUEST_NUM_PREEMPTIONS_BUCKETS: tuple[float, ...] = (1, 2, 3, 4, 5, 10, 20) +"""Per-request preemption counts, recorded for every finished request: the +first bucket therefore holds never-preempted requests together with +singly-preempted ones, 2 through 5 are separated individually, and the tail +coarsens to 20. Distinct from `REQUEST_PARAMS_N_BUCKETS` because these are the +boundaries the metric shipped with, not a harmonized choice.""" + +KV_CACHE_RESIDENCY_BUCKETS: tuple[float, ...] = ( + 0.001, + 0.002, + 0.005, + 0.01, + 0.02, + 0.05, + 0.1, + 0.2, + 0.5, + 1, + 2, + 5, + 10, + 20, + 30, + 60, + 120, + 300, + 600, + 1200, + 1800, +) +"""KV cache block residency times: millisecond-scale reuse gaps up to +30-minute block lifetimes.""" + +_STATIC_FAMILY_DEFAULTS: dict[str, tuple[float, ...]] = { + "request_latency": REQUEST_LATENCY_BUCKETS, + "time_to_first_token": TIME_TO_FIRST_TOKEN_BUCKETS, + "inter_token_latency": INTER_TOKEN_LATENCY_BUCKETS, + "iteration_tokens": ITERATION_TOKENS_BUCKETS, + "request_params_n": REQUEST_PARAMS_N_BUCKETS, + "request_num_preemptions": REQUEST_NUM_PREEMPTIONS_BUCKETS, + "kv_cache_residency": KV_CACHE_RESIDENCY_BUCKETS, +} + + +def build_buckets(mantissa_lst: list[int], max_value: int) -> list[float]: + """Build buckets from mantissas scaled by increasing powers of 10. + + Args: + mantissa_lst: Mantissa values multiplied by each power of 10. + max_value: Largest bucket value to include. + + Returns: + Bucket values in increasing order, capped at `max_value`. + """ + exponent = 0 + buckets: list[float] = [] + while True: + for m in mantissa_lst: + value = m * 10**exponent + if value <= max_value: + buckets.append(value) + else: + return buckets + exponent += 1 + + +def build_1_2_5_buckets(max_value: int) -> list[float]: + """Build a 1-2-5 series of buckets capped at `max_value`. + + Example: + >>> build_1_2_5_buckets(100) + [1, 2, 5, 10, 20, 50, 100] + """ + return build_buckets([1, 2, 5], max_value) + + +def histogram_buckets( + family: BucketFamilyKey, + max_model_len: int | None = None, +) -> list[float]: + """Return the default bucket boundaries for a histogram family. + + Args: + family: Canonical bucket family key. + max_model_len: Cap for the token-count series of the + `request_tokens` family; required for that family and ignored + otherwise. + + Returns: + A fresh list of bucket upper bounds; callers may mutate it freely. + + Raises: + ValueError: If `family` is `request_tokens` and `max_model_len` + is None. + """ + if family == "request_tokens": + if max_model_len is None: + raise ValueError( + "max_model_len is required for the 'request_tokens' bucket family" + ) + return build_1_2_5_buckets(max_model_len) + return list(_STATIC_FAMILY_DEFAULTS[family]) diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 4d0b77add7f0..9e3f9c63a628 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -18,6 +18,7 @@ from vllm.logger import init_logger from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group from vllm.v1.engine import FinishReason +from vllm.v1.metrics.buckets import histogram_buckets from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import ( @@ -724,13 +725,23 @@ def __init__( for idx in engine_indexes } + request_tokens_buckets = histogram_buckets( + "request_tokens", max_model_len=max_model_len + ) + iteration_tokens_buckets = histogram_buckets("iteration_tokens") + request_params_n_buckets = histogram_buckets("request_params_n") + time_to_first_token_buckets = histogram_buckets("time_to_first_token") + inter_token_latency_buckets = histogram_buckets("inter_token_latency") + request_latency_buckets = histogram_buckets("request_latency") + request_num_preemptions_buckets = histogram_buckets("request_num_preemptions") + # # Histograms of counts # histogram_num_prompt_tokens_request = self._histogram_cls( name="vllm:request_prompt_tokens", documentation="Number of prefill tokens processed.", - buckets=build_1_2_5_buckets(max_model_len), + buckets=request_tokens_buckets, labelnames=labelnames, ) self.histogram_num_prompt_tokens_request = create_metric_per_engine( @@ -740,7 +751,7 @@ def __init__( histogram_num_generation_tokens_request = self._histogram_cls( name="vllm:request_generation_tokens", documentation="Number of generation tokens processed.", - buckets=build_1_2_5_buckets(max_model_len), + buckets=request_tokens_buckets, labelnames=labelnames, ) self.histogram_num_generation_tokens_request = create_metric_per_engine( @@ -753,7 +764,7 @@ def __init__( histogram_iteration_tokens = self._histogram_cls( name="vllm:iteration_tokens_total", documentation="Histogram of number of tokens per engine_step.", - buckets=[1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384], + buckets=iteration_tokens_buckets, labelnames=labelnames, ) self.histogram_iteration_tokens = create_metric_per_engine( @@ -763,7 +774,7 @@ def __init__( histogram_max_num_generation_tokens_request = self._histogram_cls( name="vllm:request_max_num_generation_tokens", documentation="Histogram of maximum number of requested generation tokens.", - buckets=build_1_2_5_buckets(max_model_len), + buckets=request_tokens_buckets, labelnames=labelnames, ) self.histogram_max_num_generation_tokens_request = create_metric_per_engine( @@ -773,7 +784,7 @@ def __init__( histogram_n_request = self._histogram_cls( name="vllm:request_params_n", documentation="Histogram of the n request parameter.", - buckets=[1, 2, 5, 10, 20], + buckets=request_params_n_buckets, labelnames=labelnames, ) self.histogram_n_request = create_metric_per_engine( @@ -783,7 +794,7 @@ def __init__( histogram_max_tokens_request = self._histogram_cls( name="vllm:request_params_max_tokens", documentation="Histogram of the max_tokens request parameter.", - buckets=build_1_2_5_buckets(max_model_len), + buckets=request_tokens_buckets, labelnames=labelnames, ) self.histogram_max_tokens_request = create_metric_per_engine( @@ -796,30 +807,7 @@ def __init__( histogram_time_to_first_token = self._histogram_cls( name="vllm:time_to_first_token_seconds", documentation="Histogram of time to first token in seconds.", - buckets=[ - 0.001, - 0.005, - 0.01, - 0.02, - 0.04, - 0.06, - 0.08, - 0.1, - 0.25, - 0.5, - 0.75, - 1.0, - 2.5, - 5.0, - 7.5, - 10.0, - 20.0, - 40.0, - 80.0, - 160.0, - 640.0, - 2560.0, - ], + buckets=time_to_first_token_buckets, labelnames=labelnames, ) self.histogram_time_to_first_token = create_metric_per_engine( @@ -829,27 +817,7 @@ def __init__( histogram_inter_token_latency = self._histogram_cls( name="vllm:inter_token_latency_seconds", documentation="Histogram of inter-token latency in seconds.", - buckets=[ - 0.01, - 0.025, - 0.05, - 0.075, - 0.1, - 0.15, - 0.2, - 0.3, - 0.4, - 0.5, - 0.75, - 1.0, - 2.5, - 5.0, - 7.5, - 10.0, - 20.0, - 40.0, - 80.0, - ], + buckets=inter_token_latency_buckets, labelnames=labelnames, ) self.histogram_inter_token_latency = create_metric_per_engine( @@ -859,56 +827,13 @@ def __init__( histogram_request_time_per_output_token = self._histogram_cls( name="vllm:request_time_per_output_token_seconds", documentation="Histogram of time_per_output_token_seconds per request.", - buckets=[ - 0.01, - 0.025, - 0.05, - 0.075, - 0.1, - 0.15, - 0.2, - 0.3, - 0.4, - 0.5, - 0.75, - 1.0, - 2.5, - 5.0, - 7.5, - 10.0, - 20.0, - 40.0, - 80.0, - ], + buckets=inter_token_latency_buckets, labelnames=labelnames, ) self.histogram_request_time_per_output_token = create_metric_per_engine( histogram_request_time_per_output_token, per_engine_labelvalues ) - request_latency_buckets = [ - 0.3, - 0.5, - 0.8, - 1.0, - 1.5, - 2.0, - 2.5, - 5.0, - 10.0, - 15.0, - 20.0, - 30.0, - 40.0, - 50.0, - 60.0, - 120.0, - 240.0, - 480.0, - 960.0, - 1920.0, - 7680.0, - ] histogram_e2e_time_request = self._histogram_cls( name="vllm:e2e_request_latency_seconds", documentation="Histogram of e2e request latency in seconds.", @@ -962,7 +887,7 @@ def __init__( histogram_request_num_preemptions = self._histogram_cls( name="vllm:request_num_preemptions", documentation="Histogram of the number of times a request was preempted.", - buckets=[1, 2, 3, 4, 5, 10, 20], + buckets=request_num_preemptions_buckets, labelnames=labelnames, ) self.histogram_request_num_preemptions = create_metric_per_engine( @@ -975,7 +900,7 @@ def __init__( "Histogram of new KV tokens computed during prefill " "(excluding cached tokens)." ), - buckets=build_1_2_5_buckets(max_model_len), + buckets=request_tokens_buckets, labelnames=labelnames, ) self.histogram_prefill_kv_computed_request = create_metric_per_engine( @@ -986,29 +911,7 @@ def __init__( # KV Cache residency metrics # if self.kv_cache_metrics_enabled: - kv_cache_residency_buckets = [ - 0.001, - 0.002, - 0.005, - 0.01, - 0.02, - 0.05, - 0.1, - 0.2, - 0.5, - 1, - 2, - 5, - 10, - 20, - 30, - 60, - 120, - 300, - 600, - 1200, - 1800, - ] + kv_cache_residency_buckets = histogram_buckets("kv_cache_residency") histogram_kv_block_lifetime = self._histogram_cls( name="vllm:kv_block_lifetime_seconds", @@ -1294,33 +1197,6 @@ def log_engine_initialized(self): self.log_metrics_info("cache_config", self.vllm_config.cache_config) -def build_buckets(mantissa_lst: list[int], max_value: int) -> list[int]: - """ - Builds a list of buckets with increasing powers of 10 multiplied by - mantissa values until the value exceeds the specified maximum. - - """ - exponent = 0 - buckets: list[int] = [] - while True: - for m in mantissa_lst: - value = m * 10**exponent - if value <= max_value: - buckets.append(value) - else: - return buckets - exponent += 1 - - -def build_1_2_5_buckets(max_value: int) -> list[int]: - """ - Example: - >>> build_1_2_5_buckets(100) - [1, 2, 5, 10, 20, 50, 100] - """ - return build_buckets([1, 2, 5], max_value) - - class StatLoggerManager: """ StatLoggerManager: From 5392fbca2a45d0224f5653e01885d271518d2dfd Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 11 Sep 2026 16:14:16 +0200 Subject: [PATCH 051/774] [ROCm][Docker] Pin AINIC apt repo to snapshot 1.117.5-a-77 (#56459) Signed-off-by: Djordje Ramic --- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.rocm_gfx1250 | 2 +- docs/features/moriio_connector_usage.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 8d5f4f8f408a..483aa9bb97e3 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -19,7 +19,7 @@ ARG NIC_BACKEND=all # Users can specify a custom version compatible with their host drivers. # The default version has been tested with ionic-dkms=26.03.3.001 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. -ARG AINIC_VERSION=1.117.5 +ARG AINIC_VERSION=1.117.5-a-77 ARG UBUNTU_CODENAME=jammy # Sccache configuration. Release builds use this today; CI can opt in when a diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index b5c9c3dfe88f..c68fdfebfe54 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -15,7 +15,7 @@ ARG NIC_BACKEND=all # Users can specify a custom version compatible with their host drivers. # The default version has been tested with ionic-dkms=26.03.3.001 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. -ARG AINIC_VERSION=1.117.5 +ARG AINIC_VERSION=1.117.5-a-77 ARG UBUNTU_CODENAME=jammy # Sccache configuration. Release builds use this today; CI can opt in when a diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index 0a5bda79a552..6c906e7db18a 100644 --- a/docs/features/moriio_connector_usage.md +++ b/docs/features/moriio_connector_usage.md @@ -253,7 +253,7 @@ To run MoRI with RDMA, your environment must have the necessary RDMA userspace l The official image `vllm/vllm-openai-rocm:nightly` comes pre-installed with userspace libraries for the following NICs and kernel module versions: -- AINIC (AMD Pensando Pollara): version `1.117.5`, ships `libionic1=54.0-187-1`, tested with `ionic-dkms=26.03.3.001` +- AINIC (AMD Pensando Pollara): version `1.117.5-a-77`, ships `libionic1=54.0-187-1`, tested with `ionic-dkms=26.03.3.001` - Thor2 (Broadcom): version `235.2.86.0`, tested with `bnxt-en-dkms=1.10.3.235.2.86.0`, `bnxt-re-dkms=235.2.86.0` Refer to [Dockerfile.rocm](../../docker/Dockerfile.rocm) for more details. For users with NICs, kernel modules, and/or FW other than those stated above we refer to From 22258a26bc090bccf5473cf681bbe9bac41bd035 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 11 Sep 2026 22:20:11 +0800 Subject: [PATCH 052/774] [Multimodal] Use GPU NVDEC for EPD encoder-only instance video media IO (#53675) Signed-off-by: Isotr0py --- tests/config/test_multimodal_config.py | 82 ++++++++++++++++++++ tests/multimodal/media/test_video.py | 44 +++++++++++ tests/multimodal/test_hasher.py | 27 +++++++ tests/multimodal/test_parse.py | 22 ++++++ tests/multimodal/test_video.py | 39 +++++++++- vllm/config/vllm.py | 56 +++++++++++++ vllm/multimodal/hasher.py | 5 +- vllm/multimodal/media/video.py | 27 +++++-- vllm/multimodal/parse.py | 16 ++-- vllm/multimodal/video.py | 25 ++++-- vllm/multimodal/video_decoders/__init__.py | 13 ++++ vllm/multimodal/video_decoders/torchcodec.py | 41 +++++++++- 12 files changed, 373 insertions(+), 24 deletions(-) diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 589eabf91d67..84cf092c2159 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -375,3 +375,85 @@ def test_vllm_config_runs_the_mm_processor_device_check(): pytest.raises(ValueError, match="also runs the language model"), ): VllmConfig._validate_mm_processor_device(vllm_config) + + +def _resolve_mm_video_decode_device( + *, + ec_role: ECRole | None, + mm_tensor_ipc: str = "torch_shm", + device: str | None = None, + video_kwargs: dict | None = None, + torchcodec_available: bool = True, +) -> dict: + """Run processor-device then video-decode-device resolution and report + the resulting video media IO kwargs.""" + mm_config = MultiModalConfig( + mm_processor_kwargs={} if device is None else {"device": device}, + mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type] + media_io_kwargs={} if video_kwargs is None else {"video": dict(video_kwargs)}, + ) + model_config = MagicMock(spec=ModelConfig) + model_config.multimodal_config = mm_config + vllm_config = MagicMock(spec=VllmConfig) + vllm_config.model_config = model_config + vllm_config.ec_transfer_config = ( + None + if ec_role is None + else ECTransferConfig(ec_connector="ECExampleConnector", ec_role=ec_role) + ) + + with ( + patch("vllm.platforms.current_platform.device_type", "cuda"), + patch( + "vllm.utils.import_utils.check_torchcodec_available", + side_effect=None if torchcodec_available else ImportError("torchcodec"), + ), + ): + VllmConfig._resolve_mm_processor_device(vllm_config) + VllmConfig._resolve_mm_video_decode_device(vllm_config) + return mm_config.media_io_kwargs.get("video", {}) + + +def test_auto_video_decode_uses_nvdec_on_encoder_instance(): + """The processor runs on the accelerator there, so decoded frames should + stay on-device too.""" + assert _resolve_mm_video_decode_device(ec_role="ec_producer") == { + "backend": "torchcodec", + "device": "cuda", + } + + +@pytest.mark.parametrize("ec_role", [None, "ec_consumer", "ec_both"]) +def test_auto_video_decode_stays_on_cpu_off_encoder_instance( + ec_role: ECRole | None, +): + """The processor stays on CPU off encode-only instances, so video + decoding should too.""" + assert _resolve_mm_video_decode_device(ec_role=ec_role) == {} + + +def test_auto_video_decode_follows_explicit_processor_device(): + assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cpu") == {} + assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cuda") == { + "backend": "torchcodec", + "device": "cuda", + } + + +def test_auto_video_decode_respects_explicit_media_io_kwargs(): + assert _resolve_mm_video_decode_device( + ec_role="ec_producer", video_kwargs={"backend": "opencv"} + ) == {"backend": "opencv"} + assert _resolve_mm_video_decode_device( + ec_role="ec_producer", video_kwargs={"device": "cpu"} + ) == {"device": "cpu"} + + +def test_auto_video_decode_skipped_without_torchcodec(): + assert ( + _resolve_mm_video_decode_device( + ec_role="ec_producer", + torchcodec_available=False, + ) + == {} + ) diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index a3ea0890a20b..79070ac50d5d 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -532,6 +532,50 @@ def test_strips_pool_size_from_runtime(self): def test_unknown_backend_not_treated_as_gpu(self): assert not VIDEO_LOADER_REGISTRY.backend_requires_gpu("totally_unknown") + def test_strips_request_level_device(self): + """The decode device is a startup-only knob: a request must not move + decoding onto the GPU when the startup config did not opt in, nor off + it when it did.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec"}, + runtime_kwargs={"device": "cuda"}, + ) + assert "device" not in result + + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec", "device": "cuda"}, + runtime_kwargs={"device": "cpu", "num_frames": 8}, + ) + assert result["device"] == "cuda" + assert result["num_frames"] == 8 + + +@pytest.mark.parametrize( + "default_kwargs", + [ + {"backend": "torchcodec", "device": "cuda", "seek_mode": "approximate"}, + {"backend": "pynvvideocodec", "hw_decoders": 2}, + ], +) +def test_switching_backend_drops_stale_codec_options(default_kwargs): + """Codec-specific options from the static config must not leak into a + different codec backend selected per-request, where they would fail the + new backend's option validation.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={**default_kwargs, "num_frames": 8}, + runtime_kwargs={"backend": "opencv"}, + ) + assert result == {"backend": "opencv", "num_frames": 8} + + +def test_same_backend_keeps_codec_options(): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec", "device": "cuda"}, + runtime_kwargs={"backend": "torchcodec", "num_frames": 8}, + ) + assert result["device"] == "cuda" + assert result["num_frames"] == 8 + @pytest.mark.parametrize("layout", ["nhwc", "nchw"]) def test_pynvvc_frames_normalized_to_nhwc(layout: str): diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 4810081ad40c..d0be316790d2 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -137,6 +137,33 @@ def item_for_hash(num_frames: int): ) +def test_hash_video_tensor_frames(): + """Videos holding tensor frames (e.g. NVDEC-decoded) hash like + array-framed ones, from the original bytes without a D2H copy.""" + source = b"x" * 100 + + def item_for_hash(frames): + metadata = { + "total_num_frames": 2, + "fps": 2.0, + "duration": 1.0, + "video_backend": "torchcodec", + "frames_indices": [0, 1], + "do_sample_frames": False, + } + video = MediaWithBytes((frames, metadata), source) + items = MultiModalDataParser()._parse_video_data([video]) + return items.get_all_items_for_hash()[0] + + np_frames = np.zeros((2, 8, 8, 3), dtype=np.uint8) + torch_frames = torch.zeros((2, 8, 8, 3), dtype=torch.uint8) + + hasher = MultiModalHasher + assert hasher.hash_kwargs("blake3", video=item_for_hash(np_frames)) == ( + hasher.hash_kwargs("blake3", video=item_for_hash(torch_frames)) + ) + + def test_hash_non_contiguous_array(): arr = np.arange(24).reshape(4, 6).T assert not arr.flags.c_contiguous diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index b4d37b0fe805..2649f80777c1 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -56,6 +56,28 @@ def test_frame_size_hwc_chw(frame): assert items.get_frame_size(0) == (W, H) +def test_video_with_metadata_tensor_passthrough(): + """Tensor frames pass through unchanged regardless of device: HF video + processors accept tensors, and device-resident frames (e.g. NVDEC-decoded) + must not be copied back to host.""" + frames = torch.zeros((4, H, W, 3), dtype=torch.uint8) + video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) + + assert video is frames + assert metadata is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_video_with_metadata_keeps_device_tensor(): + """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors, + so a device-side HF processor can consume them without a D2H copy.""" + frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda") + video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) + + assert video is frames + assert metadata is None + + @pytest.mark.parametrize( "frames", [ diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 9c24b5f96932..0bc3b21b42b4 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -11,6 +11,7 @@ import numpy as np import numpy.typing as npt import pytest +import torch from transformers import AutoVideoProcessor from transformers.video_utils import VideoMetadata @@ -173,9 +174,14 @@ def fake_import_module(name: str, package: str): [ ( "torchcodec", - {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + { + "min_frames": 4, + "num_ffmpeg_threads": 2, + "seek_mode": "approximate", + "device": "cuda", + }, {"min_frames": 4}, - {"num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + {"num_ffmpeg_threads": 2, "seek_mode": "approximate", "device": "cuda"}, ), ( "deepstream", @@ -205,6 +211,11 @@ def test_video_backend_rejects_options_for_another_decoder(): ): resolve_video_backend_kwargs("opencv", {"num_ffmpeg_threads": 2}) + with pytest.raises( + ValueError, match="device is not supported by the 'opencv' backend" + ): + resolve_video_backend_kwargs("opencv", {"device": "cuda"}) + @pytest.mark.parametrize( ("backend", "error"), @@ -1205,6 +1216,30 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path): + """With device="cuda", torchcodec decodes via NVDEC and keeps the frames + on the GPU instead of returning a host-side numpy array.""" + pytest.importorskip("torchcodec") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_data, num_frames=8, backend="torchcodec", device="cuda" + ) + + assert isinstance(frames, torch.Tensor) + assert frames.device.type == "cuda" + assert frames.dtype == torch.uint8 + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] == 8 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec" + + def test_torchcodec_backend_returns_target_frames_not_keyframes(): """Regression test: torchcodec must return the requested frames, not the GOP keyframe they seek back to. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 39c2cb4d52a0..05377a273513 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1736,6 +1736,7 @@ def has_blocked_weights(): self._resolve_mm_embedding_inputs() self._resolve_mm_processor_device() + self._resolve_mm_video_decode_device() self._validate_mm_processor_device() if self.use_v2_model_runner: @@ -2623,6 +2624,61 @@ def _resolve_mm_processor_device(self) -> None: device_type, ) + def _resolve_mm_video_decode_device(self) -> None: + """Default video decoding to torchcodec GPU backend for EPD encoder-only + instance if the mm processor runs on CUDA. + + The processor consumes the decoded frames on-device in that case, so + keeping the frames on the GPU skips the host round-trip through the + CPU media path. An explicit codec/backend choice in + `--media-io-kwargs` is left alone, and the default is skipped where + torchcodec (or its FFmpeg runtime) is unavailable. + """ + if self.model_config is None or self.model_config.multimodal_config is None: + return + mm_config = self.model_config.multimodal_config + + ec_config = self.ec_transfer_config + # An EC producer that is not also a consumer runs no forward pass and + # allocates no KV cache, so frontend accelerator work has the device to + # itself. + if ec_config is None or not ec_config.is_encode_only: + return + + from vllm.platforms import current_platform + + device_type = current_platform.device_type + if ( + device_type != "cuda" + or mm_config.get_mm_processor_device_type() != device_type + ): + return + + # User set video backend or device explicitly + video_kwargs = mm_config.media_io_kwargs.setdefault("video", {}) + if "backend" in video_kwargs or "device" in video_kwargs: + return + + from vllm.utils.import_utils import check_torchcodec_available + + try: + check_torchcodec_available() + except (ImportError, RuntimeError): + # torchcodec is not installed, or is installed without a usable + # FFmpeg runtime (it raises rather than returning False). + logger.info_once( + "EPD encoder instance: keeping CPU video decoding because " + "torchcodec is not available (needs a CUDA build with FFmpeg)." + ) + return + + video_kwargs["backend"] = "torchcodec" + video_kwargs["device"] = device_type + logger.info_once( + "EPD encoder instance: decoding video with NVDEC (torchcodec device=%s).", + device_type, + ) + def _validate_mm_processor_device(self) -> None: """Hand the EC config to `MultiModalConfig`, which owns the rule.""" model_config = self.model_config diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 3d86b33d5138..a432afdc9bbd 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -89,8 +89,11 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: ) return cls.iter_item_to_bytes("image", obj.original_bytes) - if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray): + if isinstance(obj, MediaWithBytes) and isinstance( + obj.media, (np.ndarray, torch.Tensor) + ): frames = obj.media + # Both np.ndarray and torch.Tensor expose .nbytes. if frames.nbytes < len(obj.original_bytes): return cls.iter_item_to_bytes("video", frames) return cls.iter_item_to_bytes("video", obj.original_bytes) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index e5c073f5f382..f1f5783b0e8a 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -19,17 +19,23 @@ safe_to_dense, ) -from ..video import VIDEO_LOADER_REGISTRY +from ..video import VIDEO_LOADER_REGISTRY, DecodedFrames +from ..video_decoders import incompatible_backend_options from .base import MediaIO, MediaWithBytes from .image import MAGIC_NUMPY_PREFIX, ImageMediaIO logger = init_logger(__name__) -class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): +class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]): """Configuration values can be user-provided either by --media-io-kwargs or by the runtime API field "media_io_kwargs". Ensure proper validation and error handling. + + Use --media-io-kwargs '{"video": {"backend": "torchcodec", + "device": "cuda"}}' to decode with NVDEC and keep frames on the GPU, + which pairs with a device-side multi-modal processor + (--mm-processor-device, EPD encode-only instances). """ @classmethod @@ -43,6 +49,9 @@ def merge_kwargs( runtime_kwargs = dict(runtime_kwargs) runtime_kwargs.pop("hw_decoders", None) runtime_kwargs.pop("pool_size", None) + # The decode device determines which device the decoded frames + # (and thus the processor pipeline) live on; startup-only. + runtime_kwargs.pop("device", None) # Block request-level selection of GPU video backends that # were not configured (and VRAM-reserved) at startup. @@ -70,6 +79,14 @@ def merge_kwargs( merged.pop("fps", None) elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs: merged.pop("num_frames", None) + + requested_backend = runtime_kwargs.get("backend") + static_backend = (default_kwargs or {}).get("backend") + if requested_backend and requested_backend != static_backend: + for name in incompatible_backend_options( + static_backend, requested_backend + ): + merged.pop(name, None) return merged def __init__( @@ -101,7 +118,7 @@ def __init__( def load_bytes( self, data: bytes - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) @@ -109,7 +126,7 @@ def load_bytes( def load_base64( self, media_type: str, data: str - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: if media_type.lower() == "video/jpeg": load_frame = partial( self.image_io.load_base64, @@ -177,7 +194,7 @@ def load_base64( def load_file( self, filepath: Path - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: with filepath.open("rb") as f: data = f.read() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 124a8845cccc..b2f8664246ae 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -36,6 +36,7 @@ VideoItem, ) from .media import MediaWithBytes +from .video import DecodedFrames _T = TypeVar("_T") _I = TypeVar("_I") @@ -653,7 +654,7 @@ def _get_audio_with_sr( def _get_video_with_metadata( self, video: VideoItem, - ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]: + ) -> tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any] | None]: if isinstance(video, MediaWithBytes): new_video, metadata = self._get_video_with_metadata(video.media) return MediaWithBytes(new_video, video.original_bytes), metadata @@ -661,10 +662,11 @@ def _get_video_with_metadata( return video if isinstance(video, list): return np.array(video), None - if isinstance(video, np.ndarray): + if isinstance(video, (np.ndarray, torch.Tensor)): + # Tensors pass through untouched: HF video processors accept them + # directly, and device tensors (e.g. NVDEC-decoded frames) must + # stay on-device to avoid a D2H round-trip. return video, None - if isinstance(video, torch.Tensor): - return video.numpy(), None assert_never(video) @@ -769,9 +771,9 @@ def _parse_video_data( data_items = data # type: ignore[assignment] new_videos = list[ - np.ndarray - | MediaWithBytes[np.ndarray] - | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]] + DecodedFrames + | MediaWithBytes[DecodedFrames] + | tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any]] | None ]() metadata_lst: list[dict[str, Any] | None] = [] diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 854b5572b3ef..4142febc3f92 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from abc import abstractmethod -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, TypeAlias, cast import numpy as np import numpy.typing as npt @@ -28,6 +28,10 @@ logger = init_logger(__name__) +DecodedFrames: TypeAlias = npt.NDArray | torch.Tensor +"""Decoded video frames: a host ``np.ndarray``, or a device ``torch.Tensor`` +when a GPU decoding codec is used (e.g. torchcodec with ``device="cuda"``).""" + class VideoLoaderRegistry(ExtensionManager): def __init__(self) -> None: @@ -164,8 +168,13 @@ def load_bytes( cls, data: bytes, **kwargs, - ) -> tuple[npt.NDArray, dict[str, Any]]: - """Load video frames from bytes and return (frames_array, metadata_dict).""" + ) -> tuple[DecodedFrames, dict[str, Any]]: + """Load video frames from bytes and return (frames, metadata_dict). + + ``frames`` is a CPU ``np.ndarray`` unless a GPU decoding codec is + used (e.g. torchcodec with ``device="cuda"``), in which case it is + a ``torch.Tensor`` living on that device. + """ raise NotImplementedError @classmethod @@ -239,7 +248,7 @@ def load_bytes( *, backend: VideoDecoderBackend = "opencv", **kwargs: Any, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> tuple[DecodedFrames, dict[str, Any]]: """Load sampled frames from raw video bytes. Args: @@ -265,6 +274,11 @@ def load_bytes( creation at the cost of relying on the file's metadata. See https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html for details. + - ``device`` (TorchCodec): ``"cpu"`` (default) decodes on the + host and returns a ``np.ndarray``; ``"cuda"`` decodes with + NVDEC and returns the frames as a CUDA ``torch.Tensor``, + which a device-side HF video processor can consume without + a host round-trip. - ``hw_decoders`` (PyNvVideoCodec): maximum number of concurrent decoder slots. Defaults to 2 and must be a positive integer. @@ -272,7 +286,8 @@ def load_bytes( size and pool acquisition timeout in seconds. Returns: - Tuple of ``(frames_array, metadata_dict)``. + Tuple of ``(frames, metadata_dict)``, where ``frames`` is a + CPU ``np.ndarray`` unless TorchCodec decodes on ``device="cuda"``. """ target = VideoTargetMetadata( num_frames=num_frames, fps=fps, max_duration=max_duration diff --git a/vllm/multimodal/video_decoders/__init__.py b/vllm/multimodal/video_decoders/__init__.py index 1e6106e19937..9931f60f871d 100644 --- a/vllm/multimodal/video_decoders/__init__.py +++ b/vllm/multimodal/video_decoders/__init__.py @@ -19,6 +19,7 @@ "torchcodec": { "num_ffmpeg_threads": 0, "seek_mode": "exact", + "device": "cpu", }, PYNVVIDEOCODEC_VIDEO_BACKEND: { "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, @@ -40,6 +41,17 @@ def _get_backend_option_defaults(backend: str) -> dict[str, Any]: ) from None +# TODO(Isotr0py): Make kwargs structured. +def incompatible_backend_options( + old_backend: str | None, + new_backend: str, +) -> set[str]: + """Option names accepted by ``old_backend`` but not by ``new_backend``.""" + old = _BACKEND_OPTION_DEFAULTS.get(old_backend or "", {}) + new = _BACKEND_OPTION_DEFAULTS.get(new_backend, {}) + return old.keys() - new.keys() + + def resolve_video_backend_kwargs( backend: str, kwargs: dict[str, Any], @@ -103,5 +115,6 @@ def decode_video( "VideoTargetMetadata", "check_frame_pixel_limit", "decode_video", + "incompatible_backend_options", "resolve_video_backend_kwargs", ] diff --git a/vllm/multimodal/video_decoders/torchcodec.py b/vllm/multimodal/video_decoders/torchcodec.py index f1f3290629a2..346f7b645b5b 100644 --- a/vllm/multimodal/video_decoders/torchcodec.py +++ b/vllm/multimodal/video_decoders/torchcodec.py @@ -5,7 +5,10 @@ import numpy as np import numpy.typing as npt +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available from .base import ( @@ -21,6 +24,8 @@ "decoders.VideoDecoder" ) +logger = init_logger(__name__) + def decode_torchcodec( loader_cls, @@ -30,12 +35,14 @@ def decode_torchcodec( *, num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", -) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + device: str = "cpu", +) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]: check_torchcodec_available() decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder( data, num_ffmpeg_threads=num_ffmpeg_threads, seek_mode=seek_mode, + device=device, ) check_frame_pixel_limit( decoder.metadata.width or 0, @@ -48,7 +55,7 @@ def decode_torchcodec( source=source, target=target, **sampling_kwargs ) frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames( - decoder, frame_idx + decoder, frame_idx, device=device ) return frames, source, frame_idx, valid @@ -67,7 +74,19 @@ def make_torchcodec_decoder( *, num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", + device: str = "cpu", ) -> "VideoDecoder": + torch_device = torch.device(device) + if torch_device.type == "cuda" and not current_platform.is_cuda(): + raise ValueError( + f"torchcodec video decoding on device {device!r} requires " + "a CUDA-capable platform." + ) + elif torch_device.type not in ("cpu", "cuda"): + raise ValueError( + f"torchcodec video decoding only supports 'cpu' and 'cuda' " + f"devices, got {device!r}." + ) # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest # of the pipeline expects, avoiding a transpose. return VideoDecoder( @@ -75,6 +94,7 @@ def make_torchcodec_decoder( dimension_order="NHWC", num_ffmpeg_threads=num_ffmpeg_threads, seek_mode=seek_mode, + device=device, ) @staticmethod @@ -91,10 +111,23 @@ def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: def decode_torchcodec_frames( decoder: "VideoDecoder", frame_indices: list[int], - ) -> tuple[npt.NDArray, list[int]]: + *, + device: str = "cpu", + ) -> tuple[npt.NDArray | torch.Tensor, list[int]]: """Decode the requested indices in one batched, index-exact call.""" if not frame_indices: return np.empty((0,), dtype=np.uint8), [] # Note: torchcodec releases the GIL for the entire call batch = decoder.get_frames_at(frame_indices) - return batch.data.numpy(), list(frame_indices) + frames = batch.data + if ( + torch.device(device).type != "cpu" + and (status := getattr(decoder, "cpu_fallback", None)) is not None + and bool(status) + ): + logger.warning_once( + "torchcodec could not use NVDEC for this video and " + "decoded on CPU instead; check codec support and " + "libnvcuvid." + ) + return frames, list(frame_indices) From 1e1060f9988fa188fd243c093610889594ba18fd Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:16:57 -0400 Subject: [PATCH 053/774] [Kimi Perf] Group fp8 mla cahche insertion, 4~6x kernel level performance improvement for small batch (#55356) Signed-off-by: yewentao256 --- csrc/libtorch_stable/cache_kernels.cu | 117 ++++++++++++++++++----- csrc/libtorch_stable/ops.h | 12 +-- csrc/libtorch_stable/torch_bindings.cpp | 8 +- tests/kernels/attention/test_cache.py | 88 +++++++++++++++++ vllm/_custom_ops.py | 4 + vllm/models/kimi_k3/nvidia/dspark_mla.py | 26 ++++- 6 files changed, 216 insertions(+), 39 deletions(-) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 79aa63a5b83a..f42d61a4f440 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -456,14 +456,14 @@ __global__ void concat_and_cache_mla_kernel( // Grouped variant of concat_and_cache_mla: inserts the context K/V for every // draft layer in a single launch. Grid is (num_tokens, num_layers); each layer // reads its own cache base pointer from kv_cache_ptrs (same pointer-array -// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no -// scaling or quantization; scalar_t is uint16_t for portability. -template +// pattern as copy_blocks_kernel). +template __global__ void concat_and_cache_mla_grouped_kernel( const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens, // kv_lora_rank] const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim] const int64_t* __restrict__ kv_cache_ptrs, // [num_layers] + const float* __restrict__ kv_scales, // [num_layers] or nullptr const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens] const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride, const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride, @@ -481,12 +481,16 @@ __global__ void concat_and_cache_mla_grouped_kernel( const int64_t block_idx = slot_idx / block_size; const int64_t block_offset = slot_idx % block_size; - scalar_t* __restrict__ kv_cache = - reinterpret_cast(kv_cache_ptrs[layer_idx]); + cache_t* __restrict__ kv_cache = + reinterpret_cast(kv_cache_ptrs[layer_idx]); const scalar_t* __restrict__ kv_c_layer = kv_c + layer_idx * kv_c_layer_stride; const scalar_t* __restrict__ k_pe_layer = k_pe + layer_idx * k_pe_layer_stride; + float scale = 0.0f; + if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) { + scale = kv_scales[layer_idx]; + } auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride, int size, int offset) { @@ -494,7 +498,12 @@ __global__ void concat_and_cache_mla_grouped_kernel( const int64_t src_idx = token_idx * src_token_stride + i; const int64_t dst_idx = block_idx * block_stride + block_offset * entry_stride + i + offset; - kv_cache[dst_idx] = src[src_idx]; + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + kv_cache[dst_idx] = src[src_idx]; + } else { + kv_cache[dst_idx] = + fp8::scaled_convert(src[src_idx], scale); + } } }; @@ -1001,20 +1010,55 @@ void concat_and_cache_mla_grouped( torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim] torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64 - int64_t block_size, int64_t block_stride, int64_t entry_stride) { - int num_layers = kv_c.size(0); - int num_tokens = kv_c.size(1); - int kv_lora_rank = kv_c.size(2); - int pe_dim = k_pe.size(2); + int64_t block_size, int64_t block_stride, int64_t entry_stride, + std::optional kv_scales, // [num_layers] or None + const std::string& kv_cache_dtype) { + const bool use_fp8 = kv_cache_dtype == "fp8" || + kv_cache_dtype == "fp8_e4m3" || + kv_cache_dtype == "fp8_e5m2"; +#ifdef USE_ROCM + STD_TORCH_CHECK(kv_cache_dtype != "fp8_e5m2", + "concat_and_cache_mla_grouped does not support fp8_e5m2 " + "KV cache on ROCm"); +#endif + STD_TORCH_CHECK( + use_fp8 || kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16", + "concat_and_cache_mla_grouped only supports BF16 and plain " + "FP8 KV cache; got ", + kv_cache_dtype); STD_TORCH_CHECK( kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 && k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=", + "concat_and_cache_mla_grouped requires BF16 inputs; got kv_c=", kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type()); STD_TORCH_CHECK( - kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, - "kv_cache_ptrs must be int64"); + kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long && + slot_mapping.scalar_type() == torch::headeronly::ScalarType::Long, + "cache pointers and slot mapping must be int64"); + + const int num_layers = kv_c.size(0); + const int num_tokens = kv_c.size(1); + const int kv_lora_rank = kv_c.size(2); + const int pe_dim = k_pe.size(2); + const float* kv_scales_ptr = nullptr; + if (use_fp8) { + STD_TORCH_CHECK(kv_scales.has_value(), + "FP8 grouped cache insert requires kv_scales"); + STD_TORCH_CHECK( + kv_scales->scalar_type() == torch::headeronly::ScalarType::Float, + "kv_scales must be float32"); + STD_TORCH_CHECK(kv_scales->numel() == num_layers, + "kv_scales must contain one scale per layer"); + STD_TORCH_CHECK(kv_scales->is_cuda() && kv_scales->get_device_index() == + kv_c.get_device_index(), + "kv_scales must be on the same CUDA device as kv_c"); + STD_TORCH_CHECK(kv_scales->is_contiguous(), "kv_scales must be contiguous"); + kv_scales_ptr = kv_scales->const_data_ptr(); + } else { + STD_TORCH_CHECK(!kv_scales.has_value(), + "BF16 grouped cache insert does not use kv_scales"); + } if (num_tokens == 0 || num_layers == 0) { return; @@ -1030,17 +1074,40 @@ void concat_and_cache_mla_grouped( kv_c.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); - dim3 grid(num_tokens, num_layers); - dim3 block(std::min(kv_lora_rank, 512)); - vllm::concat_and_cache_mla_grouped_kernel - <<>>( - reinterpret_cast(kv_c.data_ptr()), - reinterpret_cast(k_pe.data_ptr()), - kv_cache_ptrs.const_data_ptr(), - slot_mapping.const_data_ptr(), kv_c_layer_stride, - kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, - slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, - block_size); + const dim3 grid(num_tokens, num_layers); + const dim3 block(std::min(kv_lora_rank, 512)); + + if (!use_fp8) { + vllm::concat_and_cache_mla_grouped_kernel + <<>>( + reinterpret_cast(kv_c.data_ptr()), + reinterpret_cast(k_pe.data_ptr()), + kv_cache_ptrs.const_data_ptr(), nullptr, + slot_mapping.const_data_ptr(), kv_c_layer_stride, + kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, + slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, + block_size); + return; + } + +#define LAUNCH_GROUPED_FP8(KV_DTYPE) \ + vllm::concat_and_cache_mla_grouped_kernel<__nv_bfloat16, uint8_t, KV_DTYPE> \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + kv_cache_ptrs.const_data_ptr(), kv_scales_ptr, \ + slot_mapping.const_data_ptr(), kv_c_layer_stride, \ + kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, \ + slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, \ + block_size) + + if (kv_cache_dtype == "fp8_e5m2") { + LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E5M2); + } else { + LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E4M3); + } +#undef LAUNCH_GROUPED_FP8 } namespace vllm { diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index a8beef5cda3d..13429a54e694 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -654,12 +654,12 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, const std::string& kv_cache_dtype, torch::stable::Tensor& scale); -void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c, - torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache_ptrs, - torch::stable::Tensor& slot_mapping, - int64_t block_size, int64_t block_stride, - int64_t entry_stride); +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, int64_t entry_stride, + std::optional kv_scales, + const std::string& kv_cache_dtype); // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 3781eeaebf26..336c495113b3 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -976,14 +976,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " str kv_cache_dtype," " Tensor scale) -> ()"); - // Grouped concat_and_cache_mla across all layers (bf16 only). Each - // layer's cache base pointer is read from kv_cache_ptrs. + // Grouped concat_and_cache_mla across all layers. Each layer's cache base + // pointer and optional plain-FP8 scale are read from device tensors. ops.def( "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," " Tensor kv_cache_ptrs," " Tensor slot_mapping," " int block_size, int block_stride," - " int entry_stride) -> ()"); + " int entry_stride," + " Tensor? kv_scales=None," + " str kv_cache_dtype='auto') -> ()"); // Rotate Q and K, then write to kv cache for MLA ops.def( diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 6fabcaa789b3..0d5066d2f112 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -805,6 +805,94 @@ def test_concat_and_cache_mla( torch.testing.assert_close(kv_cache, ref_kv_cache) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("shared_slot_mapping", [False, True]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e5m2"]) +@torch.inference_mode() +def test_concat_and_cache_mla_grouped( + device: str, + shared_slot_mapping: bool, + kv_cache_dtype: str, +) -> None: + if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm(): + pytest.skip("fp8_e5m2 KV cache is not supported on ROCm/HIP") + + set_random_seed(0) + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + num_layers = 5 + num_tokens = 42 + num_blocks = 8 + block_size = 16 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + entry_size = kv_lora_rank + qk_rope_head_dim + total_slots = num_blocks * block_size + + kv_c = torch.randn(num_layers, num_tokens, kv_lora_rank, dtype=torch.bfloat16) + k_pe = torch.randn(num_layers, num_tokens, qk_rope_head_dim, dtype=torch.bfloat16) + kv_caches = torch.zeros( + num_layers, + num_blocks, + block_size, + entry_size, + dtype=torch.bfloat16 if kv_cache_dtype == "auto" else torch.uint8, + ) + reference = torch.zeros_like(kv_caches) + scales = torch.linspace(0.05, 0.25, num_layers, dtype=torch.float32) + + slot_mapping = torch.stack( + [torch.randperm(total_slots)[:num_tokens] for _ in range(num_layers)] + ) + slot_mapping[:, -1] = -1 + if shared_slot_mapping: + slot_mapping = slot_mapping[:1].expand(num_layers, -1) + + for layer_idx in range(num_layers): + ops.concat_and_cache_mla( + kv_c[layer_idx], + k_pe[layer_idx], + reference[layer_idx], + slot_mapping[layer_idx], + kv_cache_dtype, + scales[layer_idx], + ) + + cache_ptrs = torch.tensor( + [kv_caches[layer_idx].data_ptr() for layer_idx in range(num_layers)], + dtype=torch.int64, + ) + ref_cache = kv_caches[0] + + def run_grouped(kv_scales: torch.Tensor | None) -> None: + ops.concat_and_cache_mla_grouped( + kv_c, + k_pe, + cache_ptrs, + slot_mapping, + ref_cache.size(1), + ref_cache.stride(0), + ref_cache.stride(1), + kv_scales, + kv_cache_dtype, + ) + + run_grouped(None if kv_cache_dtype == "auto" else scales) + + torch.testing.assert_close(kv_caches, reference, rtol=0, atol=0) + + if kv_cache_dtype == "fp8" and not shared_slot_mapping: + noncontiguous_scales = torch.ones(num_layers, 2)[:, 0] + assert not noncontiguous_scales.is_contiguous() + for invalid_scales, error in ( + (scales.cpu(), "same CUDA device"), + (noncontiguous_scales, "must be contiguous"), + ): + with pytest.raises(RuntimeError, match=error): + run_grouped(invalid_scales) + + @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index db20adb9946c..c5fbb76d985d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2828,6 +2828,8 @@ def concat_and_cache_mla_grouped( block_size: int, block_stride: int, entry_stride: int, + kv_scales: torch.Tensor | None = None, + kv_cache_dtype: str = "auto", ) -> None: torch.ops._C_cache_ops.concat_and_cache_mla_grouped( kv_c, @@ -2837,6 +2839,8 @@ def concat_and_cache_mla_grouped( block_size, block_stride, entry_stride, + kv_scales, + kv_cache_dtype, ) diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 97febda0c86e..87f9e94d9b66 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -28,6 +28,10 @@ from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.worker.workspace import current_workspace_manager +_GROUPED_KV_CACHE_DTYPES = frozenset( + {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"} +) + def _duplicate_context_kv_weights( weights: Iterable[tuple[str, torch.Tensor]], num_layers: int @@ -233,6 +237,13 @@ def _build_fused_context_kv_metadata(self) -> None: self._context_kv_lora_rank = attn0.kv_lora_rank self._context_rope_dim = attn0.qk_rope_head_dim self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._context_kv_scales: torch.Tensor | None = None + if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache( + attn0.kv_cache_dtype + ): + self._context_kv_scales = torch.stack( + [attn._k_scale.reshape(()) for attn in attentions] + ) def _precompute_fused_context_kv( self, @@ -286,16 +297,18 @@ def _precompute_fused_context_kv( return cache_layers = [layer.self_attn for layer in self.layers] + cache_dtype = cache_layers[0].kv_cache_dtype if ( - not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype) + cache_dtype in _GROUPED_KV_CACHE_DTYPES + and all(cl.kv_cache_dtype == cache_dtype for cl in cache_layers) and self._has_uniform_block_layout(cache_layers) and ( isinstance(context_slot_mapping, torch.Tensor) or all(s is not None for s in context_slot_mapping) ) ): - # Grouped context KV insert only supports unquantized (bf16) KV cache - # and assumes that all layers share the same block layout. + # Grouped context KV insert assumes that all layers share the same + # cache dtype and block layout. if isinstance(context_slot_mapping, (list, tuple)): per_layer_slot_mappings = [ @@ -321,6 +334,8 @@ def _precompute_fused_context_kv( ref_cache.size(1), ref_cache.stride(0), ref_cache.stride(1), + self._context_kv_scales, + cache_dtype, ) return @@ -483,6 +498,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) # read: 1. all weights. 2. context kv weights weights = _duplicate_context_kv_weights(weights, len(self.model.layers)) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: self.model._build_fused_context_kv_metadata() - return loaded_weights From 5fe77aecfc7687c1f3cc49862022ce34d45d1784 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:17:02 -0400 Subject: [PATCH 054/774] [Deprecation] Deprecate items scheduled for 0.29 (#55353) Signed-off-by: yewentao256 --- benchmarks/attention_benchmarks/mla_runner.py | 6 -- tests/engine/test_arg_utils.py | 18 ----- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 4 +- tests/kernels/test_compressor_kv_cache.py | 8 +- .../attention/test_chunked_local_attention.py | 3 +- tests/v1/attention/utils.py | 9 --- .../test_partial_prefix_cache_hits.py | 4 +- .../e2e/spec_decode/draft_model/test_async.py | 76 ++----------------- vllm/config/attention.py | 25 +----- vllm/config/cache.py | 15 +--- vllm/config/multimodal.py | 18 +---- .../kv_connector/v1/offloading/scheduler.py | 4 +- vllm/envs.py | 18 ----- .../layers/attention/cross_attention.py | 2 +- vllm/model_executor/models/whisper_causal.py | 6 +- vllm/models/deepseek_v4/compressor.py | 12 ++- vllm/platforms/rocm.py | 10 --- vllm/v1/attention/backend.py | 42 +--------- vllm/v1/attention/backends/flashinfer.py | 2 +- vllm/v1/attention/backends/mamba_attn.py | 22 +++--- vllm/v1/attention/backends/mla/indexer.py | 4 +- .../backends/mla/rocm_aiter_mla_sparse.py | 7 +- vllm/v1/attention/backends/utils.py | 10 +-- vllm/v1/spec_decode/dflash.py | 2 - vllm/v1/spec_decode/llm_base_proposer.py | 11 --- vllm/v1/spec_decode/step3p5.py | 2 - vllm/v1/worker/gpu_model_runner.py | 10 +-- vllm/v1/worker/ubatch_utils.py | 17 ----- 28 files changed, 62 insertions(+), 305 deletions(-) diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 790b2f89442b..8a6ae3f87d4a 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -409,10 +409,6 @@ def _build_attention_metadata( seq_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32) seq_lens_gpu = seq_lens_cpu.to(device) - # Build num_computed_tokens (context length for each request) - context_lens = [kv_len - q_len for q_len, kv_len in zip(q_lens, kv_lens)] - num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32) - # Build block table num_blocks_per_req = [(kv + block_size - 1) // block_size for kv in kv_lens] max_num_blocks = max(num_blocks_per_req) @@ -454,8 +450,6 @@ def _build_attention_metadata( query_start_loc_cpu=q_start_cpu, seq_lens=seq_lens_gpu, seq_lens_cpu_upper_bound=seq_lens_cpu, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=num_computed_tokens_cpu, slot_mapping=slot_mapping, block_table_tensor=block_table_gpu, dcp_local_seq_lens=None, diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index f8948610226c..71b9dbbab851 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -561,24 +561,6 @@ def test_prefix_cache_default(): assert engine_args.prefix_cache_retention_interval == 64 -def test_prefix_cache_retention_interval_from_deprecated_env( - monkeypatch, caplog, disable_log_dedup -): - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") - - engine_args = EngineArgs() - - assert engine_args.prefix_cache_retention_interval == 64 - assert "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in caplog.text - assert "deprecated" in caplog.text - assert "prefix_cache_retention_interval" in caplog.text - - parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) - args = parser.parse_args(["--prefix-cache-retention-interval", "32"]) - engine_args = EngineArgs.from_cli_args(args) - assert engine_args.prefix_cache_retention_interval == 32 - - @pytest.mark.parametrize( ("arg", "expected", "option"), [ diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index 99da0af7d57b..aabce97fac9e 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -81,7 +81,7 @@ def _make_common_metadata(): query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc, seq_lens=seq_lens, - _seq_lens_cpu=seq_lens, + seq_lens_cpu_upper_bound=seq_lens, num_reqs=2, num_actual_tokens=2, max_query_len=1, @@ -99,7 +99,7 @@ def _make_mixed_common_metadata(): query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc, seq_lens=seq_lens, - _seq_lens_cpu=seq_lens, + seq_lens_cpu_upper_bound=seq_lens, num_reqs=2, num_actual_tokens=5, max_query_len=4, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 1f6601dc6cfc..082a60b7a5a6 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -705,9 +705,13 @@ def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None: ], ) def test_get_c128_boundary(starts, query_start_loc, expected): + query_start_loc_tensor = torch.tensor(query_start_loc) + query_lens = query_start_loc_tensor[1:] - query_start_loc_tensor[:-1] metadata = SimpleNamespace( - _num_computed_tokens_cpu=None if starts is None else torch.tensor(starts), - query_start_loc_cpu=torch.tensor(query_start_loc), + seq_lens_cpu_upper_bound=( + None if starts is None else torch.tensor(starts) + query_lens + ), + query_start_loc_cpu=query_start_loc_tensor, ) assert _get_c128_boundary(metadata) is expected diff --git a/tests/v1/attention/test_chunked_local_attention.py b/tests/v1/attention/test_chunked_local_attention.py index c2798c8f2fa4..0826b7073080 100644 --- a/tests/v1/attention/test_chunked_local_attention.py +++ b/tests/v1/attention/test_chunked_local_attention.py @@ -181,7 +181,8 @@ def test_local_attention_virtual_batches(test_data: LocalAttentionTestData): # Convert to numpy for easier comparison actual_q_seqlens = np.diff(result.query_start_loc_cpu.numpy()) - actual_k_seqlens = result.seq_lens_cpu.numpy() + assert result.seq_lens_cpu_upper_bound is not None + actual_k_seqlens = result.seq_lens_cpu_upper_bound.numpy() # Check that all query lengths are less than or equal to attn_chunk_size assert all(q_len <= attn_chunk_size for q_len in actual_q_seqlens) diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 590c9efb8ad6..48e618565c72 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -85,13 +85,6 @@ def create_common_attn_metadata( seq_lens_cpu = seq_lens.cpu() max_seq_len = int(seq_lens_cpu.max()) - # Create computed tokens (context length for each sequence) - context_lens = [ - batch_spec.seq_lens[i] - batch_spec.query_lens[i] - for i in range(batch_spec.batch_size) - ] - num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32) - # Create block table and slot mapping max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size if arange_block_indices: @@ -122,8 +115,6 @@ def create_common_attn_metadata( query_start_loc_cpu=query_start_loc_cpu, seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=num_computed_tokens_cpu, num_reqs=batch_spec.batch_size, num_actual_tokens=num_tokens, max_query_len=max_query_len, diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index c020a5db5483..97b846219fb4 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -2386,8 +2386,8 @@ def test_retention_snapshots_handed_off_with_exact_block_ids(): def test_snapshot_handoff_dense_default_retention(): """With dense (default) retention, every materialized mamba boundary state block is handed off — regular mamba-align + prefix-match-unit - deployments get store-able boundary snapshots without setting - VLLM_PREFIX_CACHE_RETENTION_INTERVAL.""" + deployments get store-able boundary snapshots without setting a + prefix-cache retention interval.""" hash_block_size = 2 block_size = 4 manager = make_kv_cache_manager( diff --git a/tests/v1/e2e/spec_decode/draft_model/test_async.py b/tests/v1/e2e/spec_decode/draft_model/test_async.py index d7f932b0c256..ebb471ea711b 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_async.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_async.py @@ -1,70 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that verifies no implicit GPU-CPU synchronization occurs during -speculative decoding generation under expected conditions. -""" - -import multiprocessing -import sys -import traceback +"""Tests for sync-free speculative decoding with async scheduling.""" import pytest -import torch - - -@pytest.fixture -def sync_tracker(): - """ - Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect - lazy init syncs. Prints stack traces immediately when syncs occur. - """ - from vllm.v1.attention.backend import CommonAttentionMetadata - - # Shared counter for cross-process communication (inherited by fork) - sync_count = multiprocessing.Value("i", 0) - - # Save original property - original_prop = CommonAttentionMetadata.seq_lens_cpu - original_fget = original_prop.fget - - # Create tracking wrapper - def tracking_seq_lens_cpu(self): - if self._seq_lens_cpu is None: - # Increment counter - with sync_count.get_lock(): - sync_count.value += 1 - count = sync_count.value - # Print stack trace immediately (shows in subprocess output) - print(f"\n{'=' * 60}", file=sys.stderr) - print(f"SYNC #{count}: seq_lens_cpu lazy init triggered!", file=sys.stderr) - print(f"{'=' * 60}", file=sys.stderr) - traceback.print_stack(file=sys.stderr) - print(f"{'=' * 60}\n", file=sys.stderr) - sys.stderr.flush() - return original_fget(self) - - # Apply patch - CommonAttentionMetadata.seq_lens_cpu = property(tracking_seq_lens_cpu) - - class SyncTracker: - @property - def count(self) -> int: - return sync_count.value - - def assert_no_sync(self, msg: str = ""): - count = sync_count.value - assert count == 0, ( - f"Unexpected GPU-CPU sync: seq_lens_cpu lazy init triggered " - f"{count} times. See stack traces above. {msg}" - ) - - yield SyncTracker() - - # Restore original property - CommonAttentionMetadata.seq_lens_cpu = original_prop - torch._dynamo.reset() - # Test configurations: (model, spec_model, method, num_spec_tokens, backend_env) SPEC_DECODE_CONFIGS = [ @@ -97,18 +35,16 @@ def assert_no_sync(self, msg: str = ""): SPEC_DECODE_CONFIGS, ) def test_no_sync_with_spec_decode( - sync_tracker, model: str, spec_model: str, method: str, num_spec_tokens: int, + monkeypatch, vllm_runner, ): - """ - Test that no implicit GPU-CPU sync occurs during speculative decoding - generation. - """ - # Import vLLM AFTER sync_tracker fixture has applied the patch + """Test generation while the worker rejects unintended GPU-CPU syncs.""" + monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", "error") + from vllm import SamplingParams from vllm.config import CompilationConfig @@ -158,5 +94,3 @@ def test_no_sync_with_spec_decode( assert outputs[0].outputs[0].text, ( f"{method} target={model}: generated completion text is empty" ) - - sync_tracker.assert_no_sync() diff --git a/vllm/config/attention.py b/vllm/config/attention.py index d5a150c92a31..799614a69934 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -7,12 +7,9 @@ from pydantic import field_validator from vllm.config.utils import config -from vllm.logger import init_logger from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum -logger = init_logger(__name__) - IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] @@ -64,10 +61,6 @@ class AttentionConfig: use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" - use_fp4_indexer_cache: bool | None = None - """Deprecated alias for `indexer_kv_dtype`; use that instead. True maps to - `mxfp4`, False is a no-op (it selected the model default already).""" - indexer_kv_dtype: IndexerKVDType = "auto" """Data type for the sparse-attention indexer K cache. "auto" picks the model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse @@ -116,20 +109,6 @@ def __post_init__(self) -> None: # layers still use the platform's normal automatic backend. self.backend = None - if self.use_fp4_indexer_cache is not None: - logger.warning( - "use_fp4_indexer_cache is deprecated and will be removed in " - "v0.29. Use indexer_kv_dtype instead (True -> 'mxfp4')." - ) - if self.use_fp4_indexer_cache: - if self.indexer_kv_dtype not in ("auto", "mxfp4"): - raise ValueError( - "use_fp4_indexer_cache=True conflicts with " - f"indexer_kv_dtype={self.indexer_kv_dtype!r}. Set only " - "indexer_kv_dtype." - ) - self.indexer_kv_dtype = "mxfp4" - def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType: """Resolve `indexer_kv_dtype`, substituting `default` for "auto".""" if self.indexer_kv_dtype == "auto": @@ -146,9 +125,7 @@ def compute_hash(self) -> str: """ from vllm.config.utils import get_hash_factors, hash_factors - # Folded into indexer_kv_dtype by __post_init__. - ignored_factors: set[str] = {"use_fp4_indexer_cache"} - factors = get_hash_factors(self, ignored_factors) + factors = get_hash_factors(self, set()) return hash_factors(factors) @field_validator("backend", mode="before") diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 93114ab1d150..9bd898731eb1 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -8,7 +8,7 @@ from pydantic import Field, field_validator, model_validator -from vllm.config.utils import config, get_from_deprecated_env_if_set +from vllm.config.utils import config from vllm.logger import init_logger from vllm.utils.torch_utils import ( is_quantized_kv_cache, @@ -58,15 +58,6 @@ def _layout_from_name(layout_name: str) -> KVCacheLayout: ] -def _get_prefix_cache_retention_interval() -> int | None: - env_value = get_from_deprecated_env_if_set( - "VLLM_PREFIX_CACHE_RETENTION_INTERVAL", - "v0.29", - "prefix_cache_retention_interval", - ) - return 0 if env_value is None else int(env_value) - - MambaDType = Literal["auto", "float32", "float16", "bfloat16"] MambaCacheMode = Literal["all", "align", "none"] PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"] @@ -154,9 +145,7 @@ class CacheConfig: security risk tolerance against the performance benefits before turning this on. - "xxhash_cbor" combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional ``xxhash`` package.""" - prefix_cache_retention_interval: int | None = Field( - default_factory=_get_prefix_cache_retention_interval, ge=0 - ) + prefix_cache_retention_interval: int | None = Field(default=0, ge=0) """Token interval between retained sliding-window and Mamba prefix-cache checkpoints. ``0`` retains only semantic checkpoints, including the latest replay boundary and shared-prefix junctions. Positive values additionally diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index ce2217d9fefd..0d0dc22a273b 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from pathlib import Path -from typing import Any, Literal, TypeAlias, TypedDict, cast, final +from typing import Any, Literal, TypeAlias, TypedDict, final import torch from pydantic import ConfigDict, Field, field_validator, model_validator @@ -11,7 +11,7 @@ import vllm.envs as envs from vllm.config.ec_transfer import ECTransferConfig -from vllm.config.utils import config, get_from_deprecated_env_if_set +from vllm.config.utils import config from vllm.logger import init_logger from vllm.utils.hashing import safe_hash from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -75,16 +75,6 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False): `"xpu"` on XPU). Validated against that set by the CLI.""" -def _get_mm_hasher_algorithm() -> MMHasherAlgorithm: - env_value = get_from_deprecated_env_if_set( - "VLLM_MM_HASHER_ALGORITHM", - "v0.27", - "mm_hasher_algorithm", - ) - env_value = "blake3" if env_value is None else env_value - return cast(MMHasherAlgorithm, env_value.lower()) - - MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions] """ A dictionary containing an entry for each modality type of dummy data. @@ -164,9 +154,7 @@ class MultiModalConfig: mm_processor_cache_type: MMCacheType = "lru" """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`, use shared memory FIFO cache. If `lru`, use mirrored LRU cache.""" - mm_hasher_algorithm: MMHasherAlgorithm = Field( - default_factory=_get_mm_hasher_algorithm - ) + mm_hasher_algorithm: MMHasherAlgorithm = "blake3" """Hash algorithm to use for multi-modal input caching. Use `"sha256"` or `"sha512"` for FIPS-compliant deployments.""" mm_shm_cache_max_object_size_mb: int = Field(default=128, ge=0) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a3cf1a5a1e68..a3edb7875030 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -305,7 +305,7 @@ def from_spec( if retention_interval is not None: if retention_interval < 0: raise ValueError( - f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL " + f"prefix_cache_retention_interval " f"({retention_interval}) must be non-negative." ) for config in kv_group_configs: @@ -314,7 +314,7 @@ def from_spec( and retention_interval % config.tokens_per_chunk != 0 ): raise ValueError( - f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL " + f"prefix_cache_retention_interval " f"({retention_interval}) must be a multiple of " f"tokens_per_chunk ({config.tokens_per_chunk})." ) diff --git a/vllm/envs.py b/vllm/envs.py index 18e04e3dc6a2..1c2246639822 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -86,7 +86,6 @@ VLLM_MAX_IMAGE_PIXELS: int = 178_956_970 VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" - VLLM_MM_HASHER_ALGORITHM: str = "blake3" VLLM_TARGET_DEVICE: str = "cuda" VLLM_MAIN_CUDA_VERSION: str = "13.0" VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest" @@ -325,7 +324,6 @@ VLLM_LORA_ENABLE_DUAL_STREAM: bool = False VLLM_GPU_NIC_PCIE_MAPPING: str = "" VLLM_NIC_SELECTION_VARS: str = "" - VLLM_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None VLLM_ENABLE_HPC_OPS: bool = False @@ -1070,17 +1068,6 @@ def _resolve_rust_cli_path() -> str | None: # imported at runtime. # If a non-existing backend is used, an AssertionError will be thrown. "VLLM_MEDIA_CONNECTOR": lambda: os.getenv("VLLM_MEDIA_CONNECTOR", "http"), - # Hash algorithm for multimodal content hashing. - # - "blake3": Default, fast cryptographic hash (not FIPS 140-3 compliant) - # - "sha256": FIPS 140-3 compliant, widely supported - # - "sha512": FIPS 140-3 compliant, faster on 64-bit systems - # Use sha256 or sha512 for FIPS compliance in government/enterprise deployments - "VLLM_MM_HASHER_ALGORITHM": env_with_choices( - "VLLM_MM_HASHER_ALGORITHM", - "blake3", - ["blake3", "sha256", "sha512"], - case_sensitive=False, - ), # Path to the XLA persistent cache directory. # Only used for XLA devices such as TPUs. "VLLM_XLA_CACHE_PATH": lambda: os.path.expanduser( @@ -1172,11 +1159,6 @@ def _resolve_rust_cli_path() -> str | None: if "VLLM_PLUGINS" not in os.environ else os.environ["VLLM_PLUGINS"].split(",") ), - "VLLM_PREFIX_CACHE_RETENTION_INTERVAL": lambda: ( - int(os.environ["VLLM_PREFIX_CACHE_RETENTION_INTERVAL"]) - if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ - else None - ), # a local directory to look in for unrecognized LoRA adapters. # only works if plugins are enabled and # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled. diff --git a/vllm/model_executor/layers/attention/cross_attention.py b/vllm/model_executor/layers/attention/cross_attention.py index 091f0a1856d4..9336bc17c241 100644 --- a/vllm/model_executor/layers/attention/cross_attention.py +++ b/vllm/model_executor/layers/attention/cross_attention.py @@ -115,7 +115,7 @@ def build( # needed here to know how many tokens to attend to from the cached # cross-attention KV cache. new_metadata.seq_lens = common_attn_metadata.encoder_seq_lens - new_metadata._seq_lens_cpu = torch.from_numpy( + new_metadata.seq_lens_cpu_upper_bound = torch.from_numpy( common_attn_metadata.encoder_seq_lens_cpu ) diff --git a/vllm/model_executor/models/whisper_causal.py b/vllm/model_executor/models/whisper_causal.py index d76e9b0d4a69..7004bba2838e 100644 --- a/vllm/model_executor/models/whisper_causal.py +++ b/vllm/model_executor/models/whisper_causal.py @@ -161,10 +161,8 @@ def build( new_common_attn_metadata.query_start_loc *= block_pool_size new_common_attn_metadata.query_start_loc_cpu *= block_pool_size new_common_attn_metadata.seq_lens *= block_pool_size - if new_common_attn_metadata._seq_lens_cpu is not None: - new_common_attn_metadata._seq_lens_cpu *= block_pool_size - if new_common_attn_metadata._num_computed_tokens_cpu is not None: - new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size + if new_common_attn_metadata.seq_lens_cpu_upper_bound is not None: + new_common_attn_metadata.seq_lens_cpu_upper_bound *= block_pool_size new_common_attn_metadata.num_actual_tokens *= block_pool_size new_common_attn_metadata.max_query_len *= block_pool_size new_common_attn_metadata.max_seq_len *= block_pool_size diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 3bc636945a71..4da343b2c98d 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -43,10 +43,12 @@ def _prefer_two_stage_compressor() -> bool: def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None: - starts = metadata._num_computed_tokens_cpu - if starts is None: + seq_lens_cpu = metadata.seq_lens_cpu_upper_bound + if seq_lens_cpu is None: return None + query_lens = metadata.query_start_loc_cpu[1:] - metadata.query_start_loc_cpu[:-1] + starts = seq_lens_cpu - query_lens starts_list = starts.tolist() query_start_loc = metadata.query_start_loc_cpu.tolist() return any( @@ -116,6 +118,10 @@ def build( _, _, num_decode_tokens, _ = split_decodes_and_prefills( common_attn_metadata, decode_threshold=1 ) + async_spec_decode = ( + self.vllm_config.scheduler_config.async_scheduling + and self.vllm_config.speculative_config is not None + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, @@ -124,7 +130,7 @@ def build( num_decode_tokens=num_decode_tokens, c128_boundary=( _get_c128_boundary(common_attn_metadata) - if self.block_size == 8 + if self.block_size == 8 and not async_spec_decode else None ), ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 297efc632f92..9459399441f9 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -130,14 +130,6 @@ def _sync_hip_cuda_env_vars(): hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None - if cuda_val is not None: - logger.warning_once( - "Using CUDA_VISIBLE_DEVICES on ROCm is deprecated and support " - "will be removed in vLLM v0.26.0. Please use HIP_VISIBLE_DEVICES " - "instead.", - scope="process", - ) - if hip_val is not None and cuda_val is not None: if hip_val != cuda_val: raise ValueError( @@ -148,8 +140,6 @@ def _sync_hip_cuda_env_vars(): ) elif hip_val is not None: os.environ["CUDA_VISIBLE_DEVICES"] = hip_val - elif cuda_val is not None: - os.environ["HIP_VISIBLE_DEVICES"] = cuda_val # Sync at import time - catches misconfigurations from process start. diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index fcc937fe8a67..d75c79fe0f70 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -8,7 +8,6 @@ import numpy as np import torch -from typing_extensions import deprecated from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic64Sym, @@ -452,10 +451,6 @@ class CommonAttentionMetadata: at the current decode run's last full-state write. write_pos counts from here, so a preemption-resumed request re-anchors past the prompt boundary.""" - # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) - _seq_lens_cpu: torch.Tensor | None = None - _num_computed_tokens_cpu: torch.Tensor | None = None - _num_computed_tokens_cache: torch.Tensor | None = None _token_to_req_indices_cache: torch.Tensor | None = None @@ -469,36 +464,6 @@ def naive_query_lens(self) -> torch.Tensor: def replace(self, **kwargs) -> "CommonAttentionMetadata": return replace(self, **kwargs) - @property - @deprecated( - """ - Prefer using device seq_lens directly to avoid implicit H<>D sync. - If a CPU copy is needed, use `seq_lens.cpu()` instead. - Will be removed in a future release, please migrate as soon as possible. - """ - ) - def seq_lens_cpu(self) -> torch.Tensor: - if self._seq_lens_cpu is None: - self._seq_lens_cpu = self.seq_lens.to("cpu") - return self._seq_lens_cpu - - @property - @deprecated( - """ - Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full - async scheduling. If a CPU copy is needed, it can be derived from - query_start_loc_cpu and seq_lens. - Will be removed in a future release, please migrate as soon as possible. - """ - ) - def num_computed_tokens_cpu(self) -> torch.Tensor: - if self._num_computed_tokens_cpu is None: - query_seq_lens = ( - self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1] - ) - self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens - return self._num_computed_tokens_cpu - def compute_num_computed_tokens(self) -> torch.Tensor: """Compute num_computed_tokens on device (seq_lens - query_lens).""" if self._num_computed_tokens_cache is None: @@ -542,12 +507,6 @@ def unpadded( query_start_loc=self.query_start_loc[: num_actual_reqs + 1], query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1], seq_lens=self.seq_lens[:num_actual_reqs], - _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs] - if self._seq_lens_cpu is not None - else None, - _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs] - if self._num_computed_tokens_cpu is not None - else None, num_reqs=num_actual_reqs, num_actual_tokens=num_actual_tokens, max_query_len=self.max_query_len, @@ -564,6 +523,7 @@ def unpadded( encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu), dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), + seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound), is_prefilling=maybe_slice_reqs(self.is_prefilling), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 35e9db46a570..0da68e3e3d93 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1421,7 +1421,7 @@ def build( needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm if needs_seq_lens_cpu: with gpu_sync_allowed(): - seq_lens_cpu = common_attn_metadata.seq_lens_cpu + seq_lens_cpu = common_attn_metadata.seq_lens.cpu() else: seq_lens_cpu = None diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 1b9ee263db5b..e0b5fab465a1 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -633,13 +633,21 @@ def _compute_common_metadata( if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0: decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu - num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu - if decode_base_cpu is None or num_computed_tokens_cpu is None: + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + async_spec_decode = ( + self.vllm_config.scheduler_config.async_scheduling + and self.vllm_config.speculative_config is not None + ) + if decode_base_cpu is None or seq_lens_cpu is None or async_spec_decode: raise ValueError( - "--use-replayssm requires CPU decode-base and " - "computed-token counts to derive decode write positions" + "--use-replayssm requires exact CPU sequence lengths and " + "decode-base counts to derive decode write positions" ) - num_computed_d = num_computed_tokens_cpu[:num_decodes] + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1] + - common_attn_metadata.query_start_loc_cpu[:num_decodes] + ) + num_computed_d = seq_lens_cpu[:num_decodes] - query_lens_cpu decode_base_d = decode_base_cpu[:num_decodes] align_mode = self.vllm_config.cache_config.mamba_cache_mode == "align" block_size = self.kv_cache_spec.block_size @@ -656,10 +664,6 @@ def _compute_common_metadata( # write_pos counts decode steps since the ring's last full-state # write (the anchor), so a resumed request re-anchors correctly. decode_steps_cpu = num_computed_d - effective_base - query_lens_cpu = ( - common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1] - - common_attn_metadata.query_start_loc_cpu[:num_decodes] - ) valid_decode_rows = query_lens_cpu > 0 # A single-token prefill row replayed as decode (query_len==1 with # prior state) has decode_steps < 0; force it to a one-token flush diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index c20eaab0e53c..833e0b21d8b2 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -695,7 +695,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: if self.vllm_config.speculative_config else 0 ) - self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config) + self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config) next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n @@ -708,7 +708,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.use_flattening, self.supports_varlen, next_n, - self.use_fp4_indexer_cache, + self.indexer_uses_fp4, ) sm_count = num_compute_units(self.device.index) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 45a413aab381..5fba8f42378c 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -17,6 +17,7 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, @@ -638,12 +639,14 @@ def build( reduce_partial_map = None if self._use_persistent_metadata and not use_triton_sparse: num_reqs = common_attn_metadata.num_reqs + with gpu_sync_allowed(): + seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].cpu().numpy() clamped_seq_lens = np.minimum( - common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(), + seq_lens_cpu, self.topk_tokens, ) clamped_context_lens = np.minimum( - common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths, + seq_lens_cpu - seg_lengths, self.topk_tokens, ) metadata_key = ( diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 99716ca911e6..bc127ccb8b08 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -478,7 +478,7 @@ def make_local_attention_virtual_batches( query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy() with gpu_sync_allowed(): # TODO see https://github.com/vllm-project/vllm/pull/31852 - seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy() + seq_lens_np = common_attn_metadata.seq_lens.cpu().numpy() block_table = common_attn_metadata.block_table_tensor device = common_attn_metadata.query_start_loc.device @@ -541,8 +541,6 @@ def make_local_attention_virtual_batches( # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1] seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32) seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block - num_computed_tokens_local = seqlens_k_local - seqlens_q_local - k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - ( rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks) ) @@ -612,9 +610,7 @@ def make_local_attention_virtual_batches( block_table_tensor=block_table_local, slot_mapping=common_attn_metadata.slot_mapping, causal=True, - seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local), + seq_lens_cpu_upper_bound=seq_lens_cpu, ), make_block_table @@ -692,8 +688,6 @@ def make_kv_sharing_fast_prefill_common_attn_metadata( slot_mapping=common_attn_metadata.slot_mapping, causal=True, seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, - _seq_lens_cpu=common_attn_metadata._seq_lens_cpu, - _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, ) return common_attn_metadata diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 626dd36dea9b..42d8ef035c66 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -192,8 +192,6 @@ def set_inputs_first_pass( torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone() * num_query_per_req ), - _seq_lens_cpu=None, - _num_computed_tokens_cpu=None, seq_lens_cpu_upper_bound=new_seq_lens_cpu_upper_bound, num_reqs=cad.num_reqs, num_actual_tokens=num_query_total, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 4b078c2aea4a..2edbd27d6ead 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -669,9 +669,6 @@ def propose( # (i.e., not the first proposal). if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None: common_attn_metadata.seq_lens -= num_rejected_tokens_gpu - # Invalidate the CPU-side shadows to avoid H<>D sync. - common_attn_metadata._seq_lens_cpu = None - common_attn_metadata._num_computed_tokens_cpu = None block_size = self.block_size assert block_size > 0, "block_size has not been initialized." @@ -798,10 +795,6 @@ def _update_positions_dependent_metadata( self.max_model_len, ) - if common_attn_metadata._seq_lens_cpu is not None: - common_attn_metadata._seq_lens_cpu += 1 - if common_attn_metadata._num_computed_tokens_cpu is not None: - common_attn_metadata._num_computed_tokens_cpu += 1 if common_attn_metadata.seq_lens_cpu_upper_bound is not None: common_attn_metadata.seq_lens_cpu_upper_bound += 1 @@ -1134,8 +1127,6 @@ def prepare_inputs_padded( query_start_loc=common_attn_metadata.query_start_loc, seq_lens=common_attn_metadata.seq_lens, query_start_loc_cpu=query_start_loc_cpu, - _seq_lens_cpu=common_attn_metadata._seq_lens_cpu, - _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, num_reqs=common_attn_metadata.num_reqs, num_actual_tokens=total_num_tokens, @@ -1245,8 +1236,6 @@ def prepare_inputs( query_start_loc=async_tensor_h2d(new_query_start_loc_cpu, device=device), seq_lens=async_tensor_h2d(new_seq_lens_cpu, device=device), query_start_loc_cpu=new_query_start_loc_cpu, - _seq_lens_cpu=new_seq_lens_cpu, - _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, seq_lens_cpu_upper_bound=new_seq_lens_cpu, num_reqs=common_attn_metadata.num_reqs, num_actual_tokens=total_num_tokens, diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 47821cc9b2d5..109d326356a1 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -386,8 +386,6 @@ def propose( if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None: common_attn_metadata.seq_lens -= num_rejected_tokens_gpu - common_attn_metadata._seq_lens_cpu = None - common_attn_metadata._num_computed_tokens_cpu = None block_size = self.block_size assert block_size > 0, "block_size has not been initialized." diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 2730416f7698..a14d00bd58b1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2377,8 +2377,7 @@ def _get_block_table(kv_cache_gid: int): num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[ :num_reqs_padded ] - seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded] - seq_lens_cpu_upper_bound = seq_lens_cpu + seq_lens_cpu_upper_bound = self.optimistic_seq_lens_cpu[:num_reqs_padded] # is_prefilling: True if request is still in prefill phase. # Used by mamba backends to distinguish actual decodes from @@ -2388,11 +2387,6 @@ def _get_block_table(kv_cache_gid: int): # misclassify padding as prefill in CUDA graph mode. is_prefilling[num_reqs:] = False - if self.use_async_spec_decode: - # GPU tensors are authoritative in async mode. - seq_lens_cpu = None - num_computed_tokens_cpu = None - # Compute mm_prefix bidirectional ranges before building # attention metadata so builders handle them during build(). # By default, ranges exceeding sliding_window are skipped to prevent @@ -2466,8 +2460,6 @@ def _get_block_table(kv_cache_gid: int): query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], seq_lens=self.seq_lens[:num_reqs_padded], - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=num_computed_tokens_cpu, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, replayssm_decode_base_cpu=replayssm_decode_base_cpu, num_reqs=num_reqs_padded, diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index f9612b1e7f82..9743322c4165 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -280,23 +280,11 @@ def _make_metadata_with_slice( query_start_loc[1:] -= tokens_skipped query_start_loc_cpu[1:] -= tokens_skipped seq_lens = attn_metadata.seq_lens[request_slice] - # Read raw fields to avoid triggering the deprecated D2H-syncing properties. - seq_lens_cpu = ( - attn_metadata._seq_lens_cpu[request_slice] - if attn_metadata._seq_lens_cpu is not None - else None - ) seq_lens_cpu_upper_bound = ( attn_metadata.seq_lens_cpu_upper_bound[request_slice] if attn_metadata.seq_lens_cpu_upper_bound is not None else None ) - num_computed_tokens_cpu = ( - attn_metadata._num_computed_tokens_cpu[request_slice] - if attn_metadata._num_computed_tokens_cpu is not None - else None - ) - if splits_last_request: # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate # the tokens skipped because query_start_loc_cpu might have been modified @@ -309,9 +297,6 @@ def _make_metadata_with_slice( # (not cudagraph compatible) seq_lens = seq_lens.clone() seq_lens[-1] -= tokens_skipped - if seq_lens_cpu is not None: - seq_lens_cpu = seq_lens_cpu.clone() - seq_lens_cpu[-1] -= tokens_skipped if seq_lens_cpu_upper_bound is not None: seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound.clone() seq_lens_cpu_upper_bound[-1] -= tokens_skipped @@ -346,8 +331,6 @@ def _make_metadata_with_slice( block_table_tensor=block_table_tensor, slot_mapping=slot_mapping, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=num_computed_tokens_cpu, ) From b87339888d29329c42c42573e34cc2beebdcc48b Mon Sep 17 00:00:00 2001 From: Lucain Date: Fri, 11 Sep 2026 17:29:31 +0200 Subject: [PATCH 055/774] [httpx migration] Import httpx from huggingface_hub (#56460) Signed-off-by: Lucain Pouget Co-authored-by: Claude Opus 5 (1M context) --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.txt | 2 +- .../kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py | 2 +- vllm/logger.py | 2 ++ 7 files changed, 8 insertions(+), 6 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 132979d8e678..acfc986c7533 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -8,7 +8,7 @@ tqdm blake3 py-cpuinfo transformers >= 5.10.4 -huggingface_hub >= 1.28.0 +huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export tokenizers >= 0.21.1 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 81eb1b15cfc4..2b1b8d372d6e 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -342,7 +342,7 @@ httpx==0.27.2 # schemathesis httpx2==2.12.0 # via mcp -huggingface-hub==1.28.0 +huggingface-hub==1.31.0 # via # -r requirements/test/../common.txt # accelerate diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index aee0e7003d1b..2910f01e9f95 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -361,7 +361,7 @@ httpx==0.27.2 # schemathesis httpx2==2.12.0 # via mcp -huggingface-hub==1.28.0 +huggingface-hub==1.31.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index aa0835cd1b87..ed47efded1db 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -354,7 +354,7 @@ httpx==0.27.2 # schemathesis httpx2==2.12.0 # via mcp -huggingface-hub==1.28.0 +huggingface-hub==1.31.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 99e569999390..76453f50f232 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -232,7 +232,7 @@ httpx2==2.10.0 # via # mcp # openai -huggingface-hub==1.28.0 +huggingface-hub==1.31.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 0a899b7b62d2..8fa9a8836a76 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -10,12 +10,12 @@ from enum import IntEnum from typing import TYPE_CHECKING, Any -import httpx import msgspec import numpy as np import torch import zmq import zmq.asyncio +from huggingface_hub.utils import httpx from vllm import envs from vllm.config import VllmConfig diff --git a/vllm/logger.py b/vllm/logger.py index 2d053f608bb6..ec67e504bb01 100644 --- a/vllm/logger.py +++ b/vllm/logger.py @@ -258,8 +258,10 @@ def current_formatter_type(logger: Logger) -> Literal["color", "newline", None]: # Transformers uses httpx to access the Hugging Face Hub. httpx is quite verbose, # so we set its logging level to WARNING when vLLM's logging level is INFO. +# httpx2 is the successor huggingface_hub switches to in its 2.x releases. if envs.VLLM_LOGGING_LEVEL == "INFO": logging.getLogger("httpx").setLevel(logging.WARNING) + logging.getLogger("httpx2").setLevel(logging.WARNING) logger = init_logger(__name__) From 127143e27f2ba15e3b0ee6f4d73eaeaa60097bc2 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Sat, 12 Sep 2026 00:17:26 +0800 Subject: [PATCH 056/774] Revert "[Rocm][Kimi-k3] Fix pipeline_parallel support for the kimik3 DCP mode (#53664)" (#56429) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- tests/models/test_registry.py | 2 -- vllm/model_executor/layers/attention/mla_attention.py | 10 ---------- 2 files changed, 12 deletions(-) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 14d85fc09d84..b3ba172822be 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -129,8 +129,6 @@ def test_registry_model_property(model_arch, is_mm, init_cuda, score_type): # in V1. # ("MLPSpeculatorPreTrainedModel", False, False), ("DeepseekV2ForCausalLM", True, False), - ("KimiLinearForCausalLM", True, False), - ("KimiK3ForConditionalGeneration", True, False), ("Qwen2VLForConditionalGeneration", True, True), ], ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 25eef7462e69..f9b86bd810cd 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2543,16 +2543,6 @@ def build( if num_decodes > 0: dcp_tot_seq_lens_device = None if self.dcp_world_size > 1: - assert seq_lens is not None, ( - "MLA DCP decode requires seq_lens on CommonAttentionMetadata" - ) - if dcp_local_seq_lens is None: - dcp_local_seq_lens = get_dcp_local_seq_lens( - seq_lens, - dcp_size=self.dcp_world_size, - dcp_rank=get_dcp_group().rank_in_group, - cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, - ) dcp_tot_seq_lens_device = seq_lens[:num_decodes] seq_lens = dcp_local_seq_lens From 9dd969da096e37256ee37e24f6a4689d860f39ce Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 11 Sep 2026 11:17:30 -0500 Subject: [PATCH 057/774] [Model][ROCm] Enable DeepSeek V4 Vision (#55107) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/models_basic.yaml | 6 +- .../attention/test_rocm_triton_attn_dsv4.py | 95 ++++ .../processing/test_tensor_schema.py | 7 +- tests/models/test_deepseek_v4_vl_rocm.py | 458 ++++++++++++++++++ tests/models/test_initialization.py | 4 +- tests/models/test_registry.py | 7 +- tests/test_config.py | 39 ++ .../spec_decode/test_adaptive_verification.py | 1 + vllm/models/deepseek_v4/__init__.py | 8 +- vllm/models/deepseek_v4/amd/model.py | 24 +- vllm/models/deepseek_v4/amd/mtp.py | 2 +- vllm/models/deepseek_v4/amd/rocm.py | 68 ++- vllm/models/deepseek_v4/common/vl_model.py | 331 +++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 1 + vllm/models/deepseek_v4/nvidia/vl_model.py | 332 +------------ vllm/models/deepseek_v4/vl_stub.py | 4 +- vllm/platforms/rocm.py | 15 + 17 files changed, 1055 insertions(+), 347 deletions(-) create mode 100644 tests/models/test_deepseek_v4_vl_rocm.py create mode 100644 vllm/models/deepseek_v4/common/vl_model.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 717461261135..855ec8c056cd 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -48,8 +48,9 @@ steps: - tests/models/test_terratorch.py - tests/models/transformers/test_backend.py - tests/models/test_registry.py + - tests/models/test_deepseek_v4_vl_rocm.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py mirror: amd: label: ":amd: (MI300) Basic Models (Other)" @@ -59,9 +60,10 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 1c389a111733..f0ac4d838511 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -336,6 +336,101 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +@torch.inference_mode() +def test_combine_topk_swa_indices_adds_image_visibility() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + num_tokens = 8 + topk_indices = torch.full((num_tokens, 1), -1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_lens = seq_lens.clone() + left_visible = torch.tensor( + [0, 0, 0, 1, 2, 3, 4, 0], dtype=torch.int32, device=device + ) + right_visible = torch.tensor( + [0, 0, 4, 3, 2, 1, 0, 0], dtype=torch.int32, device=device + ) + + indices, lens = combine_topk_swa_indices( + topk_indices, + query_start_loc, + seq_lens, + gather_lens, + window_size=4, + compress_ratio=1, + topk=0, + M=16, + N=0, + max_image_tokens=5, + left_visible=left_visible, + right_visible=right_visible, + ) + + expected_rows = [ + [0], + [0, 1], + [0, 1, 2, 3, 4, 5, 6], + [0, 1, 2, 3, 4, 5, 6], + [1, 2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [4, 5, 6, 7], + ] + for token_idx, expected in enumerate(expected_rows): + actual = indices[token_idx, : lens[token_idx]].cpu().tolist() + assert actual == expected + + +@torch.inference_mode() +def test_combine_topk_swa_indices_apc_hit_inside_image() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((2, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 2], dtype=torch.int32, device=device), + torch.tensor([10], dtype=torch.int32, device=device), + # Only positions [5, 10) exist in the gathered SWA workspace. + torch.tensor([5], dtype=torch.int32, device=device), + window_size=4, + compress_ratio=1, + topk=0, + M=10, + N=0, + max_image_tokens=10, + left_visible=torch.tensor([8, 9], dtype=torch.int32, device=device), + # Deliberately extends beyond seq_len to exercise the upper clamp too. + right_visible=torch.tensor([5, 5], dtype=torch.int32, device=device), + ) + + assert lens.cpu().tolist() == [5, 5] + assert indices[:, :5].cpu().tolist() == [list(range(5)), list(range(5))] + + +@torch.inference_mode() +def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((1, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + window_size=120, + compress_ratio=1, + topk=0, + M=256, + N=0, + max_image_tokens=16, + ) + + assert indices.shape == (1, 256) + assert lens.item() == 1 + + def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: from vllm.models.deepseek_v4.amd import rocm as mod diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 0ba18e9f0315..e1fb0cef7439 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -164,10 +164,11 @@ def test_model_tensor_schema(model_id: str): "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." ) - if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not ( - current_platform.is_cuda() + if ( + model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" + and not current_platform.is_cuda_alike() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") if model_id == "zai-org/GLM-5.3-Flash" and (current_platform.is_xpu()): pytest.skip("GLM-5.3-Flash is not supported on XPU") diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py new file mode 100644 index 000000000000..750927972ef3 --- /dev/null +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.utils import WeightsMapper +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +def _preshuffled_fp8_linear( + holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128) +) -> nn.Module: + pytest.importorskip("aiter") + from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterPreshuffledFp8BlockScaledMMKernel, + ) + from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + ) + + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8Static128BlockSym, + activation_quant_key=kFp8Dynamic128Sym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=weight_shape, + ) + layer = nn.Module() + layer.weight = nn.Parameter( + torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()), + requires_grad=False, + ) + layer.register_parameter( + "weight_scale" if holder == "scheme" else "weight_scale_inv", + nn.Parameter( + torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"), + requires_grad=False, + ), + ) + setattr( + layer, + holder, + SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)), + ) + return layer + + +@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) +def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 + + layer = _preshuffled_fp8_linear(holder) + layer.is_bmm = True + original = layer.weight.detach().clone() + kernel = getattr(layer, holder).fp8_linear + kernel.process_weights_after_loading(layer) + + actual = _get_cached_wo_a_bf16(layer, 2, 128, 128) + torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128)) + + +@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) +def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend( + holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config +) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.deepseek_v4.amd import model as rocm_model + + gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder) + original_gate_up = gate_up.weight.detach().clone() + original_down = down.weight.detach().clone() + monkeypatch.setattr( + rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up + ) + monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + model = rocm_model.DeepseekV4MLP(128, 128, "silu") + for linear in (gate_up, down): + getattr(linear, holder).fp8_linear.process_weights_after_loading(linear) + model.prepare_gateup_preshuffle() + + for linear, original in ((gate_up, original_gate_up), (down, original_down)): + expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16)) + torch.testing.assert_close( + linear.weight.view(torch.uint8), expected.view(torch.uint8) + ) + + +def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout( + monkeypatch: pytest.MonkeyPatch, default_vllm_config +) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLAAttention, + apply_pre_quantized_block_scaled_mm, + ) + + if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256): + pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM") + linear = _preshuffled_fp8_linear(weight_shape=(1024, 256)) + original = linear.weight.detach().clone() + linear.quant_method.fp8_linear.process_weights_after_loading(linear) + attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention) + nn.Module.__init__(attention) + attention.wq_b = linear + attention.indexer = None + attention.q_lora_rank, attention.head_dim = 256, 128 + attention.eps = 1e-5 + attention.q_norm = SimpleNamespace( + weight=torch.ones(256, dtype=torch.bfloat16, device="cuda") + ) + attention.kv_norm = SimpleNamespace( + weight=torch.ones(128, dtype=torch.bfloat16, device="cuda") + ) + monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True) + # Multiple rows and groups with different ranges expose byte-order errors. + inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda") + inputs[:, :128] *= 4 + quantized, scales, _ = attention._split_qkv_and_norm(inputs) + assert scales is not None + output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales) + q = inputs[:, :256].float() + normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps) + reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1) + reference_scales /= torch.finfo(current_platform.fp8_dtype()).max + dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1) + expected = dequantized @ original.float().T + torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01) + + +def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig() + + resolved_dtype, torch_dtype = _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) + + assert resolved_dtype == "fp8_ds_mla" + assert torch_dtype is torch.uint8 + assert cache_config.cache_dtype == "fp8_ds_mla" + + +def test_rocm_packed_kv_cache_rejects_unquantized_dtype() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig(cache_dtype="bfloat16") + + with pytest.raises(ValueError, match="only supports fp8 kv-cache"): + _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) + + +def test_vl_mapper_preserves_rocm_weight_mapping() -> None: + from vllm.models.deepseek_v4.amd.model import _make_deepseek_v4_weights_mapper + from vllm.models.deepseek_v4.common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper, + ) + + text_mapper = _make_deepseek_v4_weights_mapper("fp4", fuse_shared_experts=True) + mapper = _make_deepseek_v4_vl_weights_mapper(text_mapper, image_enabled=True) + + assert mapper._map_name("layers.3.attn.wq_a.input_scale") == ( + "language_model.model.layers.3.attn.wq_a.input_scale_2" + ) + assert mapper._map_name("layers.3.ffn.shared_experts.w2.weight") == ( + "language_model.model.layers.3.ffn.shared_experts.w2.weight" + ) + assert mapper._map_name("head.weight") == "language_model.lm_head.weight" + + +def test_rocm_moe_wires_vision_routing_on_hash_and_regular_layers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import model as rocm_model + from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + + captured: list[dict] = [] + + class FakeGate(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() + + def fake_factory(**kwargs): + captured.append(kwargs) + return nn.Identity() + + monkeypatch.setattr(rocm_model, "GateLinear", FakeGate) + monkeypatch.setattr(rocm_model, "FusedMoEFactory", fake_factory) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0) + + config = SimpleNamespace( + hidden_size=16, + n_routed_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + swiglu_limit=None, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + num_hash_layers=1, + vocab_size=32, + topk_method="noaux_tc", + vision_n_layers=1, + n_shared_experts=None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), quant_config=None + ) + + hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn") + regular_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.1.ffn") + + assert hash_moe.gate.tid2eid is not None + assert regular_moe.gate.tid2eid is None + for moe, factory_kwargs in zip((hash_moe, regular_moe), captured, strict=True): + assert moe.gate.e_score_correction_bias is not None + assert moe.gate.bias_vl is not None + assert factory_kwargs["bias_vl"] is moe.gate.bias_vl + assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID + + +def test_rocm_mtp_forwards_input_ids_for_vision_routing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import mtp as rocm_mtp + + hidden_size = 4 + hc_mult = 2 + + class FakeNorm(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = 1e-6 + + class FakeMTPBlock(nn.Module): + use_fused_mhc = False + + def __init__(self) -> None: + super().__init__() + self.seen_input_ids: torch.Tensor | None = None + + def forward( + self, + *, + positions: torch.Tensor, + x: torch.Tensor, + input_ids: torch.Tensor | None, + ): + self.seen_input_ids = input_ids + return x, None, None, None + + def passthrough_mtp_input( + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + *args, + ) -> tuple[torch.Tensor, torch.Tensor]: + return inputs_embeds, previous_hidden_states + + monkeypatch.setattr( + rocm_mtp, "_FUSED_MTP_INPUT_RMSNORM_KERNEL", passthrough_mtp_input + ) + + layer = object.__new__(rocm_mtp.DeepSeekV4MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.config = SimpleNamespace(hidden_size=hidden_size) + layer.hc_mult = hc_mult + layer.enorm = FakeNorm() + layer.hnorm = FakeNorm() + layer.e_proj = nn.Identity() + layer.h_proj = nn.Identity() + layer.mtp_block = FakeMTPBlock() + + input_ids = torch.tensor([11, 12]) + positions = torch.tensor([3, 4]) + inputs_embeds = torch.arange(8, dtype=torch.float32).view(2, hidden_size) + previous_hidden_states = torch.arange(16, dtype=torch.float32).view(2, -1) + + output = layer( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + ) + + assert layer.mtp_block.seen_input_ids is input_ids + expected = previous_hidden_states.view(2, hc_mult, hidden_size) + expected = expected + inputs_embeds.unsqueeze(-2) + torch.testing.assert_close(output, expected.flatten(1)) + + +def test_rocm_compute_logits_local_skips_gather() -> None: + from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM + + calls: list[tuple[nn.Module, torch.Tensor, bool]] = [] + + def logits_processor( + lm_head: nn.Module, hidden_states: torch.Tensor, *, skip_gather: bool = False + ) -> torch.Tensor: + calls.append((lm_head, hidden_states, skip_gather)) + return hidden_states + 1 + + model = object.__new__(DeepseekV4ForCausalLM) + nn.Module.__init__(model) + model.lm_head = nn.Identity() + model.logits_processor = logits_processor + hidden_states = torch.tensor([4.0]) + + result = model.compute_logits_local(hidden_states) + + assert torch.equal(result, torch.tensor([5.0])) + assert calls == [(model.lm_head, hidden_states, True)] + + +class _FakeLanguageModel(nn.Module): + finalizes_weights_during_load = False + + def __init__(self) -> None: + super().__init__() + self.tensor_a = nn.Parameter(torch.zeros(1)) + self.tensor_c = nn.Parameter(torch.zeros(1)) + self.finalized_values: list[tuple[float, float]] = [] + + def process_weights_after_loading(self) -> None: + self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + 1 + + +def test_vl_wrapper_streams_then_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + def interleaved_weights(): + yield "language_model.tensor_a", torch.tensor([1.0]) + assert model.language_model.tensor_a.item() == 1.0 + yield "vision.tensor_b", torch.tensor([2.0]) + assert model.vision.tensor_b.item() == 2.0 + yield "language_model.tensor_c", torch.tensor([3.0]) + + loaded = model.load_weights(interleaved_weights()) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.finalized_values == [] + + model.process_weights_after_loading() + + assert model.language_model.finalized_values == [(1.0, 3.0)] + assert torch.equal( + model.compute_logits_local(torch.tensor([4.0])), torch.tensor([5.0]) + ) + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +class _FakeFinalizingLanguageModel(_FakeLanguageModel): + finalizes_weights_during_load = True + + def __init__(self) -> None: + super().__init__() + self.load_calls = 0 + + def load_weights(self, weights) -> set[str]: + self.load_calls += 1 + loaded = set() + for name, value in weights: + getattr(self, name).data.copy_(value) + loaded.add(name) + self.process_weights_after_loading() + return loaded + + +def test_vl_wrapper_groups_child_that_finalizes_during_load() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + loaded = model.load_weights( + iter( + ( + ("language_model.tensor_a", torch.tensor([1.0])), + ("vision.tensor_b", torch.tensor([2.0])), + ("language_model.tensor_c", torch.tensor([3.0])), + ) + ) + ) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.load_calls == 1 + assert model.language_model.finalized_values == [(1.0, 3.0)] + + # The framework's later model-level hook must not double-finalize a child + # which already completed this work in load_weights. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +def test_vl_wrapper_dummy_load_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + + # DummyModelLoader bypasses model.load_weights(), so no finalized marker + # exists and the framework-level hook must still delegate to the child. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 3b68f424788f..cc982f699686 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -135,8 +135,8 @@ def _initialize_kv_caches_v1(self, vllm_config): if model_arch == "DeepseekV4ForConditionalGeneration": from vllm.platforms import current_platform - if not current_platform.is_cuda(): - pytest.skip("Deepseek V4 is only supported on CUDA") + if not (current_platform.is_cuda() or current_platform.is_rocm()): + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") with ( patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1), diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index b3ba172822be..b7b13397866d 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -64,10 +64,11 @@ def test_registry_imports(model_arch): ): pytest.skip("HY V4 is only supported on CUDA and ROCm") - if model_arch == "DeepseekV4ForConditionalGeneration" and not ( - current_platform.is_cuda() + if ( + model_arch == "DeepseekV4ForConditionalGeneration" + and not current_platform.is_cuda_alike() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") # Ensure all model classes can be imported successfully model_cls = ModelRegistry._try_load_model_cls(model_arch) diff --git a/tests/test_config.py b/tests/test_config.py index 3982507b6b63..2524968cce0d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -57,6 +57,45 @@ def _write_json(path: Path, value: object) -> None: path.write_text(json.dumps(value), encoding="utf-8") +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +@pytest.mark.parametrize( + ("is_mm_prefix_lm", "is_multimodal_model", "expected"), + [ + pytest.param(True, True, True, id="multimodal-prefix-lm"), + pytest.param(False, True, False, id="multimodal-causal"), + pytest.param(True, False, False, id="text-prefix-lm"), + pytest.param(None, True, False, id="missing-model-config"), + ], +) +def test_rocm_mm_prefix_lm_disables_chunked_mm_input( + is_mm_prefix_lm: bool | None, + is_multimodal_model: bool, + expected: bool, +) -> None: + from vllm.platforms.rocm import RocmPlatform + + config = SimpleNamespace( + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + worker_cls="test-worker", + ), + model_config=( + None + if is_mm_prefix_lm is None + else SimpleNamespace(is_mm_prefix_lm=is_mm_prefix_lm) + ), + scheduler_config=SimpleNamespace( + is_multimodal_model=is_multimodal_model, + disable_chunked_mm_input=False, + ), + ) + + RocmPlatform.check_and_update_config(config) + + assert config.scheduler_config.disable_chunked_mm_input is expected + + def test_kda_recoverssm_derivation_is_revalidated(): config = SimpleNamespace( cache_config=SimpleNamespace( diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index b14fc126bbe9..116f4a7c3612 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -42,6 +42,7 @@ def make_manager( ("mode", "piecewise_capture_available", "expected"), [ ("FULL_DECODE_ONLY", True, "FULL_DECODE_ONLY"), + ("FULL_DECODE_ONLY", False, "FULL_DECODE_ONLY"), ("FULL", True, "FULL_AND_PIECEWISE"), ("FULL", False, "FULL_DECODE_ONLY"), ("FULL_AND_PIECEWISE", True, "FULL_AND_PIECEWISE"), diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index 4170a39042ed..6bbd6e0f29b5 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -20,7 +20,7 @@ ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP - from .vl_stub import ( # type: ignore[assignment] + from .common.vl_model import ( # type: ignore[assignment] DeepseekV4ForConditionalGeneration, ) elif current_platform.is_xpu(): @@ -38,14 +38,14 @@ DeepseekV4ForConditionalGeneration, ) else: + from .common.vl_model import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, ) from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment] - from .nvidia.vl_model import ( # type: ignore[assignment] - DeepseekV4ForConditionalGeneration, - ) __all__ = [ "DSparkDeepseekV4ForCausalLM", diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index e9ff6cd80f71..349ff5267981 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -79,6 +79,8 @@ from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + logger = init_logger(__name__) @@ -549,6 +551,10 @@ def __init__( self.gate.e_score_correction_bias = None self.gate.tid2eid = None + self.gate.bias_vl = None + self.image_sentinel_lo = ( + IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0 + ) is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers self.hash_indices_dtype = torch.int32 if is_hash_moe: @@ -564,12 +570,20 @@ def __init__( ), requires_grad=False, ) - elif getattr(config, "topk_method", None) == "noaux_tc": + if getattr(config, "topk_method", None) == "noaux_tc" and ( + not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0 + ): self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32), requires_grad=False, ) + if getattr(config, "vision_n_layers", 0) > 0: + self.gate.bias_vl = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + self.n_shared_experts = config.n_shared_experts # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` @@ -638,6 +652,8 @@ def __init__( routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, hash_indices_table=self.gate.tid2eid, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=self.image_sentinel_lo, swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, routed_experts_cls=( @@ -660,6 +676,8 @@ def forward( ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + if self.gate.bias_vl is not None and input_ids is None: + raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") org_shape = hidden_states.shape final_hidden_states = self.experts( @@ -1351,6 +1369,7 @@ def _make_deepseek_v4_weights_mapper( class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model + finalizes_weights_during_load = False # Default mapper assumes the original FP4-expert checkpoint layout. # Overridden per-instance in __init__ when expert_dtype != "fp4". @@ -1397,6 +1416,9 @@ def compute_logits( logits = self.logits_processor(self.lm_head, hidden_states) return logits + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True) + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index cafb6bb0ed24..d94ff5c2709e 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -160,7 +160,7 @@ def forward( inputs_embeds ).unsqueeze(-2) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None + positions=positions, x=hidden_states, input_ids=input_ids ) if self.mtp_block.use_fused_mhc: hidden_states = self.mtp_block.hc_post( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index a798053a4974..b39ad953e75d 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -67,8 +67,13 @@ def weight_already_preshuffled(linear: torch.nn.Module) -> bool: The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped for those, since shuffle_weight is a permutation rather than an involution. """ - kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) - return bool(getattr(kernel, "preshuffles_weight", False)) + return any( + getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False) + for method in ( + getattr(linear, "quant_method", None), + getattr(linear, "scheme", None), + ) + ) def apply_pre_quantized_block_scaled_mm( @@ -116,13 +121,18 @@ def _combine_topk_swa_indices_kernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, + left_visible_ptr, + right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, + SWA_WIDTH: tl.constexpr, TOPK_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, + PADDED_SWA_WIDTH: tl.constexpr, + HAS_IMAGE: tl.constexpr, ): batch_idx = tl.program_id(0) worker_id = tl.program_id(1) @@ -141,7 +151,20 @@ def _combine_topk_swa_indices_kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + if HAS_IMAGE: + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) + # Prefix caching can resume inside an image span. Do not generate + # indices outside the SWA rows present in the gathered workspace. + swa_start = tl.maximum( + tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0), gather_start + ) + swa_end = tl.minimum(pos + right + 1, seq_len) + swa_len = tl.maximum(swa_end - swa_start, 0) topk_offset = tl.arange(0, PADDED_TOP_K) topk_mask = topk_offset < topk_len @@ -159,14 +182,14 @@ def _combine_topk_swa_indices_kernel( mask=topk_mask, ) - swa_offset = tl.arange(0, WINDOW_SIZE) + swa_offset = tl.arange(0, PADDED_SWA_WIDTH) tl.store( combined_indices_ptr + token_idx * combined_indices_stride + topk_len + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, + M * batch_idx + N + swa_offset + swa_start - gather_start, + mask=(swa_offset < swa_len) & (swa_offset < SWA_WIDTH), ) tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) @@ -182,12 +205,21 @@ def combine_topk_swa_indices( topk: int, M: int, N: int, + max_image_tokens: int = 0, + left_visible: torch.Tensor | None = None, + right_visible: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + if (left_visible is None) != (right_visible is None): + raise ValueError("left_visible and right_visible must be provided together") topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() num_tokens = topk_indices.shape[0] num_reqs = seq_lens.shape[0] + has_image = left_visible is not None + # Keep the row shape fixed for a vision model even when a particular batch + # has no image. + swa_width = window_size + max_image_tokens combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -211,13 +243,18 @@ def combine_topk_swa_indices( query_start_loc, seq_lens, gather_lens, + left_visible if left_visible is not None else topk_indices, + right_visible if right_visible is not None else topk_indices, M, N, TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, + SWA_WIDTH=swa_width, TOPK_WIDTH=topk_indices.shape[-1], PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + PADDED_SWA_WIDTH=triton.next_power_of_2(swa_width), + HAS_IMAGE=has_image, ) return combined_indices, combined_lens @@ -948,6 +985,12 @@ def _forward_prefill( assert query_start_loc_cpu is not None assert query_start_loc is not None prefill_token_base = query_start_loc_cpu[num_decodes] + left_visible = swa_metadata.prefill_left_visible + right_visible = swa_metadata.prefill_right_visible + if left_visible is not None: + left_visible = left_visible[num_decode_tokens:] + assert right_visible is not None + right_visible = right_visible[num_decode_tokens:] if not swa_only: if self.compress_ratio == 4: @@ -1026,6 +1069,17 @@ def _forward_prefill( top_k, M, N, + max_image_tokens=self.max_image_tokens, + left_visible=( + left_visible[query_start:query_end] + if left_visible is not None + else None + ), + right_visible=( + right_visible[query_start:query_end] + if right_visible is not None + else None + ), ) rocm_sparse_attn_prefill( q=q[query_start:query_end], diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py new file mode 100644 index 000000000000..4cccb85a8302 --- /dev/null +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). + +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: + +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / + ``image_newline`` / ``image_end``) fill the remaining sentinel positions. +- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` + (see ``common/vision.py`` for the tower itself). +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before + its hyper-connection stream expansion. Raw ``input_ids`` still flow into + the model so the MoE router can apply ``bias_vl`` to image tokens + (``requires_raw_input_tokens``). +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from .vision import DeepseekV4Aligner, DeepseekV4ViT + + +def _make_deepseek_v4_vl_weights_mapper( + text_mapper: WeightsMapper, image_enabled: bool +) -> WeightsMapper: + """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" + orig_to_new_prefix = { + src: None if dst is None else f"language_model.{dst}" + for src, dst in text_mapper.orig_to_new_prefix.items() + } + if not image_enabled: + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) + return WeightsMapper( + orig_to_new_renaming=text_mapper.orig_to_new_renaming, + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_regex=text_mapper.orig_to_new_regex, + orig_to_new_stacked=text_mapper.orig_to_new_stacked, + orig_to_new_suffix={ + **text_mapper.orig_to_new_suffix, + "head.weight": "language_model.lm_head.weight", + }, + orig_to_new_substr={ + **text_mapper.orig_to_new_substr, + # Draft models load the checkpoint's MTP weights separately from + # the target model. + "mtp.": None, + }, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class DeepseekV4ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. + + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) + delegates through ``language_model`` via the protocol defaults. + """ + + # The MoE router needs raw token ids to detect image sentinel tokens + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return IMAGE_PLACEHOLDER + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config, prefix: str = "") -> None: + super().__init__() + model_config = vllm_config.model_config + config = model_config.hf_config + self.config = config + self.multimodal_config = model_config.multimodal_config + assert self.multimodal_config is not None + + image_enabled = ( + config.vision_n_layers > 0 + and self.multimodal_config.get_limit_per_prompt("image") > 0 + ) + with self._mark_tower_model(vllm_config, {"image"}): + self.vision: DeepseekV4ViT | None = None + self.aligner: DeepseekV4Aligner | None = None + self.image_start: nn.Parameter | None = None + self.image_end: nn.Parameter | None = None + self.image_newline: nn.Parameter | None = None + self.image_pad: nn.Parameter | None = None + if image_enabled: + self.vision = DeepseekV4ViT(config) + self.aligner = DeepseekV4Aligner(config) + for name in ( + "image_start", + "image_end", + "image_newline", + "image_pad", + ): + setattr( + self, + name, + nn.Parameter( + torch.empty(config.hidden_size, dtype=torch.float32) + ), + ) + self.vision.to(dtype=model_config.dtype) + self.aligner.to(dtype=model_config.dtype) + + with self._mark_language_model(vllm_config): + # The arch convertor routes any config with a vision tower to + # this wrapper class; mark the copy handed to the text backbone + # so it resolves to DeepseekV4ForCausalLM instead of recursing + # (with_hf_config deepcopies the config, the marker survives). + config._dsv4_vl_inner = True # type: ignore[attr-defined] + try: + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["DeepseekV4ForCausalLM"], + ) + finally: + del config._dsv4_vl_inner # type: ignore[attr-defined] + # The outer mapper (see load_weights) fully resolves HF names into + # this wrapper's namespace before AutoWeightsLoader strips the + # "language_model." prefix and delegates to the child's load_weights, + # so the child's own mapper must be a no-op. Its suffix rules are not + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would + # re-fire "head.weight" -> "lm_head.weight"). + text_mapper = self.language_model.hf_to_vllm_mapper + self.language_model.hf_to_vllm_mapper = WeightsMapper() + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( + text_mapper, image_enabled + ) + self._weights_finalized = False + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + patches = kwargs.pop("patches", None) + if patches is None: + return None + vit_grid = kwargs.pop("vit_grid", None) + llm_grid = kwargs.pop("llm_grid", None) + perm = kwargs.pop("perm", None) + assert vit_grid is not None and llm_grid is not None and perm is not None + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + perm: torch.Tensor, + ) -> torch.Tensor: + assert self.vision is not None and self.aligner is not None + image_embeds = self.aligner( + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w + ) + # Reorder into the N-layout block order used in the prompt. + return image_embeds[perm.to(image_embeds.device)] + + def _process_image_input( + self, + patches: torch.Tensor, + vit_grid: torch.Tensor, + llm_grid: torch.Tensor, + perm: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + assert self.vision is not None and self.aligner is not None + patches = patches.to(self.aligner.w1.weight.dtype) + + embeds: list[torch.Tensor] = [] + vit_offset = 0 + llm_offset = 0 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( + vit_grid.tolist(), llm_grid.tolist(), strict=True + ): + n_vit = n_vit_h * n_vit_w + n_llm = n_llm_h * n_llm_w + embeds.append( + self._encode_image( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + perm[llm_offset : llm_offset + n_llm], + ) + ) + vit_offset += n_vit + llm_offset += n_llm + return tuple(embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None or self.vision is None: + return [] + return self._process_image_input( + image_input["patches"], + image_input["vit_grid"], + image_input["llm_grid"], + image_input["perm"], + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # All ids are in-vocab here: image-block sentinels are borrowed + # reserved tokens (their embedding rows are always overwritten below). + inputs_embeds = self.language_model.embed_input_ids(input_ids) + + if self.image_start is not None: + # Branch-free sentinel overwrite: safe inside compiled/captured + # regions (no data-dependent control flow). + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + # IMAGE positions get vision embeddings via the merge below. + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) + table = torch.stack( + [ + self.image_start, + self.image_pad, + self.image_pad, + self.image_newline, + self.image_end, + ] + ).to(inputs_embeds.dtype) + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) + inputs_embeds = torch.where( + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" + return self.language_model.get_mtp_target_hidden_states() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + child_finalizes = getattr( + self.language_model, "finalizes_weights_during_load", True + ) + mapped = self.hf_to_vllm_mapper.apply(weights) + if child_finalizes: + # A child which finalizes inside load_weights must see all of its + # weights in one contiguous delegation from AutoWeightsLoader. + mapped = iter(sorted(mapped, key=lambda x: x[0])) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(mapped) + self._weights_finalized = child_finalizes + return loaded_params + + def process_weights_after_loading(self) -> None: + # Backbones such as the ROCm implementation require this to run only + # after the loader's generic per-layer quantization finalization. + if getattr(self, "_weights_finalized", False): + return + self.language_model.process_weights_after_loading() + self._weights_finalized = True diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 021001b9ce01..97e7e617fe9c 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1790,6 +1790,7 @@ class DeepseekV4ForCausalLM( SupportsLoRA, DeepseekV4MixtureOfExperts, ): + finalizes_weights_during_load = True model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py index bd40b7bce9e7..a66217f2fe65 100644 --- a/vllm/models/deepseek_v4/nvidia/vl_model.py +++ b/vllm/models/deepseek_v4/nvidia/vl_model.py @@ -1,333 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). +"""Compatibility imports for the platform-neutral DeepSeek-V4 vision model.""" -Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: - -- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE - sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / - ``image_newline`` / ``image_end``) fill the remaining sentinel positions. -- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor - in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved - in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` - (see ``common/vision.py`` for the tower itself). -- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before - its hyper-connection stream expansion. Raw ``input_ids`` still flow into - the model so the MoE router can apply ``bias_vl`` to image tokens - (``requires_raw_input_tokens``). -""" - -from collections.abc import Iterable - -import torch -from torch import nn - -from vllm.model_executor.models.interfaces import ( - MultiModalEmbeddings, - SupportsEagle3, - SupportsMultiModal, - SupportsPP, +from ..common.vl_model import ( + DeepseekV4ForConditionalGeneration, ) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - WeightsMapper, - init_vllm_registered_model, - maybe_prefix, +from ..common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper as _make_common_vl_weights_mapper, ) -from vllm.multimodal import MULTIMODAL_REGISTRY - -from ..common.mm_preprocess import ( - IMAGE_PLACEHOLDER, - IMAGE_SENTINEL_BASE_ID, - DeepseekV4VLDummyInputsBuilder, - DeepseekV4VLMultiModalProcessor, - DeepseekV4VLProcessingInfo, - image_sentinel_mask, -) -from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT from .model import _make_deepseek_v4_weights_mapper -def _make_deepseek_v4_vl_weights_mapper( - expert_dtype: str, image_enabled: bool -) -> WeightsMapper: - """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" - base = _make_deepseek_v4_weights_mapper(expert_dtype) - orig_to_new_prefix: dict[str, str | None] = { - "layers.": "language_model.model.layers.", - "embed.": "language_model.model.embed.", - "norm.": "language_model.model.norm.", - "hc_head": "language_model.model.hc_head", - "mtp.": "language_model.model.mtp.", - } - if not image_enabled: - orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) - return WeightsMapper( - orig_to_new_prefix=orig_to_new_prefix, - orig_to_new_regex=base.orig_to_new_regex, - orig_to_new_suffix={ - "head.weight": "language_model.lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", - }, - orig_to_new_substr={ - ".shared_experts.w2": ".shared_experts.down_proj", - # The MTP/DSpark draft heads are not supported for the vision - # variant; drop their weights. - "mtp.": None, - }, +def _make_deepseek_v4_vl_weights_mapper(expert_dtype: str, image_enabled: bool): + """Retain the original NVIDIA helper signature for downstream imports.""" + return _make_common_vl_weights_mapper( + _make_deepseek_v4_weights_mapper(expert_dtype), image_enabled ) -@MULTIMODAL_REGISTRY.register_processor( - DeepseekV4VLMultiModalProcessor, - info=DeepseekV4VLProcessingInfo, - dummy_inputs=DeepseekV4VLDummyInputsBuilder, -) -class DeepseekV4ForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 -): - """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. - - ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) - delegates through ``language_model`` via the protocol defaults. - """ - - # The MoE router needs raw token ids to detect image sentinel tokens - # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. - requires_raw_input_tokens = True - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - if modality == "image": - return IMAGE_PLACEHOLDER - raise ValueError(f"Unsupported modality: {modality!r}") - - def __init__(self, *, vllm_config, prefix: str = "") -> None: - super().__init__() - model_config = vllm_config.model_config - config = model_config.hf_config - self.config = config - self.multimodal_config = model_config.multimodal_config - assert self.multimodal_config is not None - - image_enabled = ( - config.vision_n_layers > 0 - and self.multimodal_config.get_limit_per_prompt("image") > 0 - ) - with self._mark_tower_model(vllm_config, {"image"}): - self.vision: DeepseekV4ViT | None = None - self.aligner: DeepseekV4Aligner | None = None - self.image_start: nn.Parameter | None = None - self.image_end: nn.Parameter | None = None - self.image_newline: nn.Parameter | None = None - self.image_pad: nn.Parameter | None = None - if image_enabled: - self.vision = DeepseekV4ViT(config) - self.aligner = DeepseekV4Aligner(config) - for name in ( - "image_start", - "image_end", - "image_newline", - "image_pad", - ): - setattr( - self, - name, - nn.Parameter( - torch.empty(config.hidden_size, dtype=torch.float32) - ), - ) - self.vision.to(dtype=model_config.dtype) - self.aligner.to(dtype=model_config.dtype) - - with self._mark_language_model(vllm_config): - # The arch convertor routes any config with a vision tower to - # this wrapper class; mark the copy handed to the text backbone - # so it resolves to DeepseekV4ForCausalLM instead of recursing - # (with_hf_config deepcopies the config, the marker survives). - config._dsv4_vl_inner = True # type: ignore[attr-defined] - try: - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config, - prefix=maybe_prefix(prefix, "language_model"), - architectures=["DeepseekV4ForCausalLM"], - ) - finally: - del config._dsv4_vl_inner # type: ignore[attr-defined] - # The outer mapper (see load_weights) fully resolves HF names into - # this wrapper's namespace before AutoWeightsLoader strips the - # "language_model." prefix and delegates to the child's load_weights, - # so the child's own mapper must be a no-op. Its suffix rules are not - # idempotent (e.g. "lm_head.weight".endswith("head.weight") would - # re-fire "head.weight" -> "lm_head.weight"). - self.language_model.hf_to_vllm_mapper = WeightsMapper() - self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] - self.language_model.make_empty_intermediate_tensors - ) - - expert_dtype = getattr(config, "expert_dtype", "fp4") - self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( - expert_dtype, image_enabled - ) - - def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: - patches = kwargs.pop("patches", None) - if patches is None: - return None - vit_grid = kwargs.pop("vit_grid", None) - llm_grid = kwargs.pop("llm_grid", None) - perm = kwargs.pop("perm", None) - assert vit_grid is not None and llm_grid is not None and perm is not None - return { - "patches": patches, - "vit_grid": vit_grid, - "llm_grid": llm_grid, - "perm": perm, - } - - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - perm: torch.Tensor, - ) -> torch.Tensor: - assert self.vision is not None and self.aligner is not None - image_embeds = self.aligner( - self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w - ) - # Reorder into the N-layout block order used in the prompt. - return image_embeds[perm.to(image_embeds.device)] - - def _process_image_input( - self, - patches: torch.Tensor, - vit_grid: torch.Tensor, - llm_grid: torch.Tensor, - perm: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - assert self.vision is not None and self.aligner is not None - patches = patches.to(self.aligner.w1.weight.dtype) - - embeds: list[torch.Tensor] = [] - vit_offset = 0 - llm_offset = 0 - for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( - vit_grid.tolist(), llm_grid.tolist(), strict=True - ): - n_vit = n_vit_h * n_vit_w - n_llm = n_llm_h * n_llm_w - embeds.append( - self._encode_image( - patches[vit_offset : vit_offset + n_vit], - n_vit_h, - n_vit_w, - perm[llm_offset : llm_offset + n_llm], - ) - ) - vit_offset += n_vit - llm_offset += n_llm - return tuple(embeds) - - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - image_input = self._parse_and_validate_image_input(**kwargs) - if image_input is None or self.vision is None: - return [] - return self._process_image_input( - image_input["patches"], - image_input["vit_grid"], - image_input["llm_grid"], - image_input["perm"], - ) - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings: MultiModalEmbeddings | None = None, - *, - is_multimodal: torch.Tensor | None = None, - ) -> torch.Tensor: - from vllm.model_executor.models.utils import _merge_multimodal_embeddings - - # All ids are in-vocab here: image-block sentinels are borrowed - # reserved tokens (their embedding rows are always overwritten below). - inputs_embeds = self.language_model.embed_input_ids(input_ids) - - if self.image_start is not None: - # Branch-free sentinel overwrite: safe inside compiled/captured - # regions (no data-dependent control flow). - sentinel_mask = image_sentinel_mask(input_ids) - if is_multimodal is not None: - # IMAGE positions get vision embeddings via the merge below. - sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) - table = torch.stack( - [ - self.image_start, - self.image_pad, - self.image_pad, - self.image_newline, - self.image_end, - ] - ).to(inputs_embeds.dtype) - idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) - inputs_embeds = torch.where( - sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds - ) - - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds - - assert is_multimodal is not None - return _merge_multimodal_embeddings( - inputs_embeds=inputs_embeds, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors=None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - return self.language_model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - return self.language_model.compute_logits(hidden_states) - - def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.language_model.compute_logits_local(hidden_states) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.language_model.get_expert_mapping() - - def get_mtp_target_hidden_states(self) -> torch.Tensor | None: - """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" - return self.language_model.get_mtp_target_hidden_states() - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Map HF names into this wrapper's namespace up front and sort, so - # the "language_model." group reaches the child loader as one - # contiguous block (AutoWeightsLoader delegates per contiguous group, - # and the child's load_weights finalizes fused expert weights, which - # must not run on a partially loaded model). - mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) - loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(mapped) - # The child's load_weights already ran its post-load finalization. - self._weights_finalized = True - return loaded_params - - def process_weights_after_loading(self) -> None: - # Model-level post-load hook (called by the loader after any load - # format). Under DummyModelLoader the child's load_weights — and - # hence its finalize step — is bypassed, so run it here instead. - if getattr(self, "_weights_finalized", False): - return - self.language_model.process_weights_after_loading() +__all__ = ["DeepseekV4ForConditionalGeneration"] diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py index 0f15a82f0190..4ade7b17d387 100644 --- a/vllm/models/deepseek_v4/vl_stub.py +++ b/vllm/models/deepseek_v4/vl_stub.py @@ -9,6 +9,6 @@ class DeepseekV4ForConditionalGeneration(nn.Module): def __init__(self, *, vllm_config, prefix: str = ""): super().__init__() raise NotImplementedError( - "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only " - "supported on NVIDIA GPUs for now." + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is not " + "supported on the current platform; use CUDA or ROCm." ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9459399441f9..a9cf534716f3 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -943,6 +943,21 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: if parallel_config.worker_cls == "auto": parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" + model_config = vllm_config.model_config + scheduler_config = vllm_config.scheduler_config + # Note: model_config may be None during testing + if ( + model_config is not None + and model_config.is_mm_prefix_lm + and scheduler_config.is_multimodal_model + and not scheduler_config.disable_chunked_mm_input + ): + logger.warning_once( + "Forcing --disable_chunked_mm_input for models " + "with multimodal-bidirectional attention." + ) + scheduler_config.disable_chunked_mm_input = True + @classmethod def verify_model_arch(cls, model_arch: str) -> None: if model_arch in _ROCM_UNSUPPORTED_MODELS: From 9a35c081e80a94828af6f611525102bb70e3c67f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Dr=C3=A9ano?= <71752785+SorenDreano@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:54:35 +0200 Subject: [PATCH 058/774] [Docs] Correct API key authentication scope (/inference does NOT bypass the API key) (#56269) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sören Dréano --- docs/usage/security.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index ba730bcb2056..559f58702690 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -147,13 +147,13 @@ firewall configuration instructions. ### Overview -The `--api-key` flag (or `VLLM_API_KEY` environment variable) provides authentication for vLLM's HTTP server, but **only for OpenAI-compatible API endpoints under the `/v1` path prefix**, and other similar `/v2`, `/inference` path prefix**. Many other sensitive endpoints are exposed on the same HTTP server without any authentication enforcement. +The `--api-key` flag (or `VLLM_API_KEY` environment variable) provides authentication for vLLM's HTTP server, but **only for endpoints under the `/v1`, `/v2`, `/inference`, and `/cohere` path prefixes**. Many other sensitive endpoints are exposed on the same HTTP server without any authentication enforcement. **Important:** Do not rely exclusively on `--api-key` for securing access to vLLM. Additional security measures are required for production deployments. ### Protected Endpoints (Require API Key) -When `--api-key` is configured, the following `/v1` endpoints require Bearer token authentication: +When `--api-key` is configured, the following endpoints require Bearer token authentication: - `/v1/models` - List available models - `/v1/chat/completions` - Chat completions @@ -249,7 +249,7 @@ These endpoints are only available when profiling is enabled and should only be An attacker who can reach the vLLM HTTP server can: -1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials +1. **Bypass authentication** by using endpoints outside the protected path prefixes, such as `/invocations`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank`, to run arbitrary inference without credentials 2. **Cause denial of service** by calling `/pause`, `/scale_elastic_ep`, or `/abort_requests` without a token 3. **Access operational controls** to manipulate server state (e.g., pausing generation, updating model weights via `/update_weights`) 4. **If `--enable-tokenizer-info-endpoint` is set:** Access sensitive tokenizer configuration including chat templates, which may reveal prompt engineering strategies or other implementation details @@ -346,7 +346,7 @@ vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` en 1. **Only allowlist plugins you trust.** Set `VLLM_PLUGINS` to the exact plugin names you intend to run and never wildcard or copy an allowlist between deployments without reviewing what each named plugin does. 2. **Audit routes before deploying.** A plugin's `attach_router` can add routes under any path, including ones that duplicate existing `/v1/*` paths. There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)), so a malicious or buggy plugin can **shadow a core route** and silently replace its behavior. Prefer plugins that namespace their routes under a distinct prefix (e.g. `/plugins//...`) instead of reusing `/v1/...` and review `app.routes` after startup if you need certainty about what is actually being served. -3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, and `/inference` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally. +3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, `/inference`, and `/cohere` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally. 4. **Remember the `vllm.general_plugins` pairing.** A plugin that also needs new engine side behavior ships that half separately via `vllm.general_plugins` which loads in every worker process under the default (load all unless restricted) posture. Allowlisting the endpoint plugin does not by itself restrict its paired engine side plugin. Need to review both. ## gRPC Interface From 6fe67cbbf3e43da89bebf6ab0eeaca4ba6c75663 Mon Sep 17 00:00:00 2001 From: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:55:21 -0400 Subject: [PATCH 059/774] [Spec][V2] Support MTP speculative decoding under pipeline parallelism (#46994) Signed-off-by: Chris Eastwood Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/attention/test_sparse_mla_backends.py | 45 +++++++++++++++++++ .../e2e/spec_decode/test_mtp_parallel_load.py | 9 +--- .../layers/attention/sparse_mla_attention.py | 44 ++++++++++++++---- vllm/model_executor/models/deepseek_mtp.py | 18 +++++++- vllm/model_executor/models/qwen3_5_mtp.py | 12 +++-- vllm/models/deepseek_v32/amd/mtp.py | 18 +++++++- vllm/models/deepseek_v32/nvidia/model.py | 5 ++- vllm/models/deepseek_v32/nvidia/mtp.py | 18 +++++++- .../mla/flashinfer_mla_sparse_sm120.py | 13 +++--- .../backends/mla/rocm_aiter_mla_sparse.py | 14 +++--- .../attention/backends/mla/xpu_mla_sparse.py | 12 +++-- vllm/v1/worker/gpu/model_runner.py | 8 +++- 12 files changed, 167 insertions(+), 49 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 002aa5a4b72c..60f9fc83f985 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, + SparseMLACommonImpl, _masked_mha_workspace_fits, _topk_mask_shape, ) @@ -44,6 +45,9 @@ from vllm.model_executor.layers.attention.mla_attention import ( _canonicalize_sparse_mla_kv_cache_dtype, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, @@ -1864,3 +1868,44 @@ def run_kernel(**kwargs): assert out.is_contiguous() assert not out.isnan().any() assert not lse.isnan().any() + + +def test_sparse_impl_observes_repointed_indexer_buffer(): + """The MTP proposer repoints the draft's indexer at the target model's buffer + after the backend impl is built, so the impl must resolve the buffer per read. + Snapshotting it in __init__ leaves the layer reading indices nothing writes.""" + impl = object.__new__(FlashInferMLASparseImpl) + own = torch.zeros(4, 8, dtype=torch.int32) + target = torch.ones(4, 8, dtype=torch.int32) + indexer = SimpleNamespace(topk_indices_buffer=own) + impl.init_topk_indices_buffer(indexer, None) + + assert impl.topk_indices_buffer is own + + indexer.topk_indices_buffer = target + + assert impl.topk_indices_buffer is target + + +def test_explicit_topk_buffer_supersedes_indexer(): + """Backbone skip-topk layers have no indexer, and the proposer also assigns the + shared buffer directly onto draft submodules.""" + impl = object.__new__(FlashInferMLASparseImpl) + indexer = SimpleNamespace( + topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32), + ) + impl.init_topk_indices_buffer(indexer, None) + explicit = torch.ones(2, 2, dtype=torch.int32) + + impl.topk_indices_buffer = explicit + + assert impl.topk_indices_buffer is explicit + + +def test_sparse_mla_common_impl_resolves_buffer_lazily(): + """Guards the whole SparseMLACommonImpl family at once: a backend that + snapshots the buffer instead silently loses MTP buffer sharing.""" + assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer) + assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( + "topk_indices_buffer must stay a lazily-resolved property" + ) diff --git a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py index 236997261675..e7b2c36c4fcd 100644 --- a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py +++ b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py @@ -56,14 +56,7 @@ class InlineConfig: enable_expert_parallel=True, enable_eplb=True, ), - InlineConfig( - id="pp2", - pp_size=2, - skip_reason=( - "DeepSeek MTP pipeline-parallel support is in flight upstream; " - "see https://github.com/vllm-project/vllm/pull/44698" - ), - ), + InlineConfig(id="pp2", pp_size=2), ] diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 8c2bb72bf4a0..4de2c02d4f51 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -501,7 +501,40 @@ def _build_topk_mask( return out[:batch_size, :max_q_len] -class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]): +class SharedTopkIndicesBuffer: + """Resolves the shared top-k index buffer for sparse MLA implementations. + + The indexer owns the buffer, but `LLMBaseProposer.load_model` repoints the + draft's indexers at the target's buffer after the impls are constructed, so + it must be resolved per read rather than snapshotted. Backbone skip-topk + layers have no indexer and pass the buffer explicitly. + """ + + _indexer: object | None = None + _topk_indices_buffer: torch.Tensor | None = None + + def init_topk_indices_buffer( + self, + indexer: object | None, + topk_indices_buffer: torch.Tensor | None, + ) -> None: + self._indexer = indexer + self._topk_indices_buffer = topk_indices_buffer + + @property + def topk_indices_buffer(self) -> torch.Tensor | None: + if self._indexer is not None: + return self._indexer.topk_indices_buffer # type: ignore[attr-defined] + return self._topk_indices_buffer + + @topk_indices_buffer.setter + def topk_indices_buffer(self, buffer: torch.Tensor | None) -> None: + # An explicit assignment supersedes the indexer. + self._indexer = None + self._topk_indices_buffer = buffer + + +class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic[T]): """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True @@ -543,14 +576,7 @@ def __init__( kv_b_proj, ) - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer # type: ignore[attr-defined] - if indexer is not None - else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) self._use_flashinfer_concat_mla_k = ( has_flashinfer() and which("ninja") is not None diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 8ac43e0bb2a1..09f4a847d064 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -39,9 +39,11 @@ DeepseekV2MoE, _try_load_fp8_indexer_wk, ) +from .interfaces import SupportsPP from .utils import ( get_pp_missing_layer_names, get_spec_layer_idx_from_weight_name, + make_empty_intermediate_tensors_factory, maybe_prefix, ) @@ -230,7 +232,7 @@ def compute_logits( @support_torch_compile -class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -240,6 +242,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) # Set MoE hyperparameters self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -266,7 +271,7 @@ def set_moe_parameters(self): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -328,6 +333,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue is_fusion_moe_shared_experts_layer = ( self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 4c064bead747..7fb699b183b7 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -38,6 +38,7 @@ from .interfaces import ( MultiModalEmbeddings, SupportsMultiModal, + SupportsPP, _require_is_multimodal, ) from .utils import ( @@ -152,7 +153,9 @@ def forward( inputs_embeds: torch.Tensor | None = None, spec_step_idx: int = 0, ) -> torch.Tensor: - if get_pp_group().is_first_rank: + # Branch on the inputs, not the rank: the drafter is built entirely on + # the last PP stage, where `is_first_rank` is False. + if intermediate_tensors is None: if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) assert hidden_states.shape[-1] == inputs_embeds.shape[-1] @@ -162,7 +165,6 @@ def forward( hidden_states = self.fc(hidden_states) residual = None else: - assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] @@ -212,7 +214,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: "hidden_states": 0, } ) -class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): +class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP): packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -254,6 +256,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + def embed_input_ids( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 9ea86ee26b0f..46c35dec98fa 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -32,8 +32,10 @@ _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) +from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( get_pp_missing_layer_names, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.models.deepseek_v32.common.kernels import fused_eh_norm @@ -160,7 +162,7 @@ def compute_logits( return self.logits_processor(mtp_layer.shared_head.head, hidden_states) -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -169,6 +171,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( self.model.layers.values(), DeepseekV2MoE, @@ -192,7 +197,7 @@ def set_moe_parameters(self): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -268,6 +273,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue is_fusion_moe_shared_experts_layer = ( self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 76aa6ca8e161..92b74e536fc6 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -182,7 +182,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # DSA is always sparse (has index_topk); allocate the shared top-k # buffer the indexer writes and the sparse MLA backend reads. self.is_v32 = True - topk_indices_buffer = torch.empty( + # On the model, not a local: the MTP proposer shares it with the draft. + self.topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, config.index_topk, dtype=torch.int32, @@ -207,7 +208,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): lambda prefix: DeepseekV32DecoderLayer( vllm_config=vllm_config, prefix=prefix, - topk_indices_buffer=topk_indices_buffer, + topk_indices_buffer=self.topk_indices_buffer, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 2ea01dbdcb9c..a5348ed2204a 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -33,8 +33,10 @@ _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) +from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( get_pp_missing_layer_names, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm @@ -260,7 +262,7 @@ def get_top_tokens( ) -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -271,6 +273,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if self.config.model_type == "glm_moe_dsa": enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype) self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -289,7 +294,7 @@ def set_moe_parameters(self): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -368,6 +373,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 4de2c2f53f3b..b7459203e07f 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -6,6 +6,9 @@ import torch +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.v1.attention.backend import ( AttentionLayer, AttentionType, @@ -29,7 +32,9 @@ def _kv_scale_format_for_model(model_type: str | None) -> str: return "pow2_fp32" -class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]): +class FlashInferMLASparseSM120Impl( + MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer +): """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True @@ -86,11 +91,7 @@ def __init__( # Skip-topk layers are built with indexer=None and get the shared # buffer via mla_args instead (cf. FLASHMLA_SPARSE). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer - if indexer is not None - else mla_args.get("topk_indices_buffer") - ) + self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer")) from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 5fba8f42378c..b110f95d57cf 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -15,6 +15,9 @@ from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -750,7 +753,9 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor: return (result, lse) -class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]): +class ROCMAiterMLASparseImpl( + MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer +): is_sparse = True supports_dense_mha_prefill = False supports_dcp = False @@ -795,12 +800,7 @@ def __init__( self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 8aaa4749d506..158fbfc4155b 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, @@ -183,7 +186,7 @@ def build( return metadata -class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]): +class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata], SharedTopkIndicesBuffer): is_sparse = True def __init__( @@ -210,12 +213,7 @@ def __init__( self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) def _forward_bf16_kv( self, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f41a5a8d3ae1..fcadf9e3032b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2030,7 +2030,13 @@ def sample_tokens( ) mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None - if self.speculator is not None and self.speculator.supports_mm_inputs: + # The encoder runner exists only on the first PP rank, so later ranks + # have no cached embeddings to gather. + if ( + self.speculator is not None + and self.speculator.supports_mm_inputs + and self.model_state.supports_mm_inputs + ): # Get cached multimodal embeddings for draft forward. # NOTE: This is done here because postprocess updates # num_computed_prefill_tokens. From dc07f1638f73814b95776832b85df1cc92850416 Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Fri, 11 Sep 2026 19:02:13 +0100 Subject: [PATCH 060/774] [Bugfix][MLA] Read sparse model settings from text config (#56160) Signed-off-by: Andy Lo --- tests/v1/attention/test_dspark_noncausal_sparse_mla.py | 4 +++- .../attention/test_indexer_deepseek_v4_slot_mapping.py | 8 +++++--- tests/v1/attention/test_sparse_mla_backends.py | 10 ++++++++-- .../kernels/attention/dsa/dcp_indexer_cutedsl.py | 4 ++-- .../layers/attention/sparse_mla_attention.py | 4 ++-- vllm/v1/attention/backends/mla/compressor_utils.py | 6 +++--- vllm/v1/attention/backends/mla/flashattn_mla_sparse.py | 4 ++-- vllm/v1/attention/backends/mla/indexer.py | 8 ++++---- .../v1/attention/backends/mla/rocm_aiter_mla_sparse.py | 2 +- vllm/v1/attention/backends/mla/sparse_utils.py | 2 +- vllm/v1/attention/backends/mla/xpu_mla_sparse.py | 2 +- 11 files changed, 32 insertions(+), 22 deletions(-) diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 8fd74314d1da..d5c794131474 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -168,6 +168,7 @@ def _run_sparse_backend_vs_sdpa( ) model_config = vllm_config.model_config model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, q_lora_rank=None, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, @@ -175,6 +176,7 @@ def _run_sparse_backend_vs_sdpa( v_head_dim=v_head_dim, model_type="deepseek_v2", ) + del model_config.hf_config.index_topk # Composite configs only nest this field. model_config.dtype = dtype model_config.get_num_attention_heads = MethodType( lambda self, parallel_config: num_heads, model_config @@ -287,7 +289,7 @@ def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa): causal_reference = torch.cat(causal_reference_outputs, dim=0) vllm_config.cache_config.cache_dtype = kv_cache_dtype - vllm_config.model_config.hf_config.index_topk = topk_tokens + vllm_config.model_config.hf_text_config.index_topk = topk_tokens common_attn_metadata = create_common_attn_metadata( batch_spec, block_size, device, arange_block_indices=True diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 51f7e70f5fb8..800b0557dead 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -68,7 +68,9 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): config = SimpleNamespace( scheduler_config=SimpleNamespace(max_num_batched_tokens=8), model_config=SimpleNamespace( - hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32) + hf_text_config=SimpleNamespace( + compress_ratios=[0, 0, 4, 128, 0], index_kpool=32 + ) ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, @@ -89,7 +91,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): def test_compressed_slot_mapping_warmup_includes_index_kpool(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=256), - model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)), + model_config=SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=32)), ) keys = CompressedSlotMappingKernel().get_warmup_keys(config) @@ -101,7 +103,7 @@ def test_index_conversion_warmup_uses_physical_block_stride(): cache_config=SimpleNamespace(block_size=64), model_config=SimpleNamespace( max_model_len=1024, - hf_config=SimpleNamespace(index_topk=2048), + hf_text_config=SimpleNamespace(index_topk=2048), ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 60f9fc83f985..425db95a8a7f 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -423,6 +423,7 @@ def test_sparse_backend_decode_correctness( ) model_config = vllm_config.model_config model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, q_lora_rank=None, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, @@ -430,6 +431,7 @@ def test_sparse_backend_decode_correctness( v_head_dim=v_head_dim, model_type="deepseek_v2", ) + del model_config.hf_config.index_topk # Composite configs only nest this field. model_config.dtype = dtype model_config.get_num_attention_heads = MethodType( lambda self, parallel_config: num_heads, @@ -591,7 +593,7 @@ def test_sparse_backend_decode_correctness( sdpa_reference = torch.cat(reference_outputs, dim=0) vllm_config.cache_config.cache_dtype = kv_cache_dtype - vllm_config.model_config.hf_config.index_topk = topk_tokens + vllm_config.model_config.hf_text_config.index_topk = topk_tokens common_attn_metadata = create_common_attn_metadata( batch_spec, @@ -1149,6 +1151,7 @@ def test_sparse_backend_prefill_correctness( ) model_config = vllm_config.model_config model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, q_lora_rank=None, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, @@ -1156,6 +1159,7 @@ def test_sparse_backend_prefill_correctness( v_head_dim=v_head_dim, model_type="deepseek_v2", ) + del model_config.hf_config.index_topk # Composite configs only nest this field. model_config.dtype = dtype model_config.model_arch_config.total_num_attention_heads = num_heads model_config.get_num_attention_heads = MethodType( @@ -1255,7 +1259,7 @@ def test_sparse_backend_prefill_correctness( ref_output = torch.cat(reference_outputs, dim=0) vllm_config.cache_config.cache_dtype = kv_cache_dtype - vllm_config.model_config.hf_config.index_topk = topk_tokens + vllm_config.model_config.hf_text_config.index_topk = topk_tokens common_attn_metadata = create_common_attn_metadata( batch_spec, @@ -1747,6 +1751,7 @@ def _build_sparse_dcp_vllm_config( model_config = vllm_config.model_config model_config.dtype = torch.bfloat16 model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, q_lora_rank=None, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, @@ -1754,6 +1759,7 @@ def _build_sparse_dcp_vllm_config( v_head_dim=v_head_dim, model_type="deepseek_v2", ) + del model_config.hf_config.index_topk # Composite configs only nest this field. model_config.get_num_attention_heads = MethodType( lambda self, parallel_config: local_heads, model_config ) diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py index 276e1c465f4c..29cdb2c589e3 100644 --- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py +++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py @@ -173,7 +173,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: if dcp_world_size <= 1: return [] cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size - topk = vllm_config.model_config.hf_config.index_topk + topk = vllm_config.model_config.hf_text_config.index_topk if topk <= 0: return [] @@ -546,7 +546,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size if dcp_world_size <= 1: return [] - topk = vllm_config.model_config.hf_config.index_topk + topk = vllm_config.model_config.hf_text_config.index_topk if topk <= 0: return [] return self._trace_dispatch(self.dispatch)( diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 4de2c02d4f51..fd29debbb759 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -132,7 +132,7 @@ def __init__( self.device = device self.model_config = vllm_config.model_config self.mla_dims = get_mla_dims(self.model_config) - self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk + self.topk_tokens: int = vllm_config.model_config.hf_text_config.index_topk self.req_id_per_token_buffer = torch.empty( (vllm_config.scheduler_config.max_num_batched_tokens,), dtype=torch.int32, @@ -202,7 +202,7 @@ def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int: scheduler_config = vllm_config.scheduler_config cache_config = vllm_config.cache_config model_config = vllm_config.model_config - topk_tokens = model_config.hf_config.index_topk + topk_tokens = model_config.hf_text_config.index_topk workspace_size = min( max( diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index a7d801b0406c..78c30c0f5f2c 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -97,10 +97,10 @@ def dispatch( # type: ignore[override] ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: - hf_config = vllm_config.model_config.hf_config + hf_text_config = vllm_config.model_config.hf_text_config configured_ratios = ( - *(getattr(hf_config, "compress_ratios", None) or ()), - getattr(hf_config, "index_kpool", 1) or 1, + *(getattr(hf_text_config, "compress_ratios", None) or ()), + getattr(hf_text_config, "index_kpool", 1) or 1, ) compress_ratios = tuple( dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 4a78a700cfaa..2882a81b1152 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -99,8 +99,8 @@ def supports_combination( if vllm_config.parallel_config.decode_context_parallel_size > 1: return "FlashAttention MLA Sparse does not support DCP for now" - hf_config = vllm_config.model_config.hf_config - if not hasattr(hf_config, "index_topk"): + hf_text_config = vllm_config.model_config.hf_text_config + if not hasattr(hf_text_config, "index_topk"): return "FlashAttention MLA Sparse requires model with index_topk" return None diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 833e0b21d8b2..0eb642bd7632 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -396,7 +396,7 @@ def dispatch( # type: ignore[override] def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8)) - hf_config = vllm_config.model_config.hf_config + hf_text_config = vllm_config.model_config.hf_text_config parallel_config = vllm_config.parallel_config dcp_world = parallel_config.decode_context_parallel_size dcp_interleave = parallel_config.cp_kv_cache_interleave_size @@ -405,12 +405,12 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: dict.fromkeys( max(1, int(ratio)) for ratio in ( - *(getattr(hf_config, "compress_ratios", None) or (1,)), - getattr(hf_config, "index_kpool", 1) or 1, + *(getattr(hf_text_config, "compress_ratios", None) or (1,)), + getattr(hf_text_config, "index_kpool", 1) or 1, ) ) ) - index_kpool = getattr(hf_config, "index_kpool", None) + index_kpool = getattr(hf_text_config, "index_kpool", None) if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios: compress_ratios = compress_ratios + (index_kpool,) return self._trace_dispatch(self.dispatch)( diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index b110f95d57cf..86b77db79a81 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -435,7 +435,7 @@ def __init__( self.num_heads = self.model_config.get_num_attention_heads(parallel_config) self.mla_dims = get_mla_dims(self.model_config) - self.topk_tokens = vllm_config.model_config.hf_config.index_topk + self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk attention_context = vllm_config.compilation_config.static_forward_context # Sink decode must use AITER's nonpersistent path. In particular, # gfx942 has no persistent+LSE kernel, and its metadata heuristic diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 83e0e780e4c0..329d48be2141 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -232,7 +232,7 @@ def get_warmup_keys( dcp_size = vllm_config.parallel_config.decode_context_parallel_size dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0 - num_topk_tokens = vllm_config.model_config.hf_config.index_topk + num_topk_tokens = vllm_config.model_config.hf_text_config.index_topk max_num_blocks = cdiv( vllm_config.model_config.max_model_len, block_size * dcp_size, diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 158fbfc4155b..036dfe01730b 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -128,7 +128,7 @@ def __init__( self.num_heads = self.model_config.get_num_attention_heads(parallel_config) self.mla_dims = get_mla_dims(self.model_config) - self.topk_tokens = vllm_config.model_config.hf_config.index_topk + self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk self.topk_tokens_tensor = torch.tensor( [self.topk_tokens], device=device, dtype=torch.int32 ) From 8a7f98c8c3ae409df6130995e364393ead5dbc24 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 11 Sep 2026 15:05:22 -0400 Subject: [PATCH 061/774] [Kernel][MoE] Optimize batched_moe_align_block_size with cooperative writes (#53280) Signed-off-by: mgoin --- .../kernels/benchmark_moe_align_block_size.py | 83 ++++++++++++++++++- .../moe/moe_align_sum_kernels.cu | 78 +++++++++++++---- .../kernels/moe/test_moe_align_block_size.py | 20 +++-- 3 files changed, 154 insertions(+), 27 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe_align_block_size.py b/benchmarks/kernels/benchmark_moe_align_block_size.py index a340500379a0..7ddf5b3ecf7c 100644 --- a/benchmarks/kernels/benchmark_moe_align_block_size.py +++ b/benchmarks/kernels/benchmark_moe_align_block_size.py @@ -5,6 +5,7 @@ import torch +import vllm._custom_ops as ops from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) @@ -68,8 +69,87 @@ def benchmark(num_tokens, num_experts, topk, ep_size, provider): return 1000 * ms, 1000 * max_ms, 1000 * min_ms +batched_configs = list( + itertools.product( + [8, 32, 64, 256], + [16, 256, 512, 2048], + [8, 32, 48, 64], + ["uniform", "skewed", "mostly_empty"], + ) +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=[ + "num_batches", + "max_tokens_per_batch", + "block_size", + "distribution", + ], + x_vals=batched_configs, + line_arg="provider", + line_vals=["vllm"], + line_names=["vLLM"], + plot_name="batched-moe-align-block-size-performance", + args={}, + ) +) +def benchmark_batched( + num_batches, + max_tokens_per_batch, + block_size, + distribution, + provider, +): + set_random_seed(0) + if distribution == "uniform": + expert_num_tokens = torch.full( + (num_batches,), + max_tokens_per_batch, + dtype=torch.int32, + device="cuda", + ) + elif distribution == "skewed": + expert_num_tokens = torch.ones((num_batches,), dtype=torch.int32, device="cuda") + expert_num_tokens[0] = max_tokens_per_batch + else: + expert_num_tokens = torch.zeros( + (num_batches,), dtype=torch.int32, device="cuda" + ) + expert_num_tokens[::8] = max_tokens_per_batch + + max_num_tokens_padded = ( + num_batches + * ((max_tokens_per_batch + block_size - 1) // block_size) + * block_size + ) + sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda") + expert_ids = torch.empty( + (max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda" + ) + num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda") + + quantiles = [0.5, 0.2, 0.8] + if provider == "vllm": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: ops.batched_moe_align_block_size( + max_tokens_per_batch, + block_size, + expert_num_tokens, + sorted_ids, + expert_ids, + num_tokens_post_pad, + ), + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument("--batched", action="store_true") parser.add_argument( "--num_experts", type=int, @@ -85,4 +165,5 @@ def benchmark(num_tokens, num_experts, topk, ep_size, provider): ) args = parser.parse_args() - benchmark.run(print_data=True, show_plots=True) + selected_benchmark = benchmark_batched if args.batched else benchmark + selected_benchmark.run(print_data=True, show_plots=True) diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index 98a04de783d4..d66b78fe838a 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -23,20 +23,18 @@ namespace batched_moe_align_block_size { // Note num_threads needs to be 1024 for BlockScan Reduction in the kernel. static constexpr int32_t num_threads = 1024; static constexpr int32_t num_blocks = 1; +template __global__ void batched_moe_align_block_size_kernel( int32_t const num_batches, int32_t const max_tokens_per_batch, int32_t const block_size, int32_t const* __restrict__ batch_num_tokens, int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids, int32_t* __restrict__ num_tokens_post_pad) { - // TODO(varun): This is a naive implementation. Could be optimized. - - size_t const batch_id = threadIdx.x; size_t const stride = blockDim.x * gridDim.x; int32_t const num_blocks_per_batch = CEILDIV(max_tokens_per_batch, block_size); - int32_t const sorted_ids_size = - num_blocks_per_batch * num_batches * block_size; - int32_t const block_ids_size = sorted_ids_size / block_size; + size_t const sorted_ids_size = + static_cast(num_blocks_per_batch) * num_batches * block_size; + size_t const block_ids_size = sorted_ids_size / block_size; int32_t const SENTINEL = num_batches * max_tokens_per_batch; // To denote invalid entries. // Initialize sorted_ids @@ -49,8 +47,8 @@ __global__ void batched_moe_align_block_size_kernel( } int32_t b_num_tokens = 0; - if (batch_id < num_batches) { - b_num_tokens = batch_num_tokens[batch_id]; + if (threadIdx.x < num_batches) { + b_num_tokens = batch_num_tokens[threadIdx.x]; } int32_t const ceil_b_num_tokens = CEILDIV(b_num_tokens, block_size) * block_size; @@ -62,20 +60,56 @@ __global__ void batched_moe_align_block_size_kernel( BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val); __syncthreads(); - bool const is_last_batch = batch_id == (num_batches - 1); - if (is_last_batch) { + if (threadIdx.x == num_batches - 1) { *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens; } - if (batch_id < num_batches) { - int32_t const batch_offset = batch_id * max_tokens_per_batch; + if constexpr (cooperative_writes) { + __shared__ int32_t batch_cumsum[num_threads]; + __shared__ int32_t valid_tokens[num_threads]; + __shared__ int32_t batch_num_blocks[num_threads]; + if (threadIdx.x < num_batches) { + batch_cumsum[threadIdx.x] = cumsum_val; + valid_tokens[threadIdx.x] = b_num_tokens; + batch_num_blocks[threadIdx.x] = ceil_b_num_tokens / block_size; + } + __syncthreads(); + + int32_t const max_num_groups = blockDim.x / WARP_SIZE; + int32_t num_groups = 1; + while (num_groups < num_batches && num_groups < max_num_groups) { + num_groups *= 2; + } + int32_t const threads_per_batch = blockDim.x / num_groups; + int32_t const group_id = threadIdx.x / threads_per_batch; + int32_t const group_offset = threadIdx.x % threads_per_batch; + + // Assign at least one warp to each batch when possible. + for (int32_t batch_id = group_id; batch_id < num_batches; + batch_id += num_groups) { + size_t const batch_offset = + static_cast(batch_id) * max_tokens_per_batch; + size_t const cumsum = batch_cumsum[batch_id]; + for (size_t i = group_offset; i < valid_tokens[batch_id]; + i += threads_per_batch) { + sorted_ids[cumsum + i] = static_cast(batch_offset + i); + } + + size_t const block_start = cumsum / block_size; + for (size_t i = group_offset; i < batch_num_blocks[batch_id]; + i += threads_per_batch) { + block_ids[block_start + i] = batch_id; + } + } + } else if (threadIdx.x < num_batches) { + size_t const batch_id = threadIdx.x; + size_t const batch_offset = batch_id * max_tokens_per_batch; for (size_t i = 0; i < b_num_tokens; ++i) { - sorted_ids[cumsum_val + i] = batch_offset + i; + sorted_ids[cumsum_val + i] = static_cast(batch_offset + i); } - int32_t const block_start = cumsum_val / block_size; - int32_t const num_blocks = ceil_b_num_tokens / block_size; - for (size_t i = 0; i < num_blocks; ++i) { + size_t const block_start = cumsum_val / block_size; + for (size_t i = 0; i < ceil_b_num_tokens / block_size; ++i) { block_ids[block_start + i] = batch_id; } } @@ -747,8 +781,16 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch, STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1); STD_TORCH_CHECK(B <= batched_kernel::num_threads); - batched_kernel::batched_moe_align_block_size_kernel<<< - batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>( + // Avoid coordination overhead for small capacities or many batches. + int64_t const cooperative_threshold = std::max(256, 8 * B); + bool const use_cooperative_writes = + max_tokens_per_batch >= cooperative_threshold; + auto kernel = + use_cooperative_writes + ? batched_kernel::batched_moe_align_block_size_kernel + : batched_kernel::batched_moe_align_block_size_kernel; + kernel<<>>( B, max_tokens_per_batch, block_size, reinterpret_cast(batch_num_tokens.const_data_ptr()), reinterpret_cast(sorted_ids.mutable_data_ptr()), diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index 7d07af4080c4..8c16690785c3 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -329,13 +329,15 @@ def test_moe_align_block_size_deterministic(): @pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512]) @pytest.mark.parametrize("num_experts", [8, 16, 32, 64]) -@pytest.mark.parametrize("block_size", [8, 16, 32, 64]) -@pytest.mark.parametrize("simulate_empty_batches", [False, True]) +@pytest.mark.parametrize("block_size", [8, 16, 32, 48, 64]) +@pytest.mark.parametrize( + "token_distribution", ["random", "alternating_empty", "skewed"] +) def test_batched_moe_align_block_size( max_tokens_per_batch: int, num_experts: int, block_size: int, - simulate_empty_batches: bool, + token_distribution: str, ): def ref_outputs( expert_num_tokens: torch.Tensor, @@ -384,15 +386,17 @@ def ref_outputs( # Compute expert_num_tokens expert_num_tokens = torch.randint( low=0, - high=max_tokens_per_batch, + high=max_tokens_per_batch + 1, size=(num_experts,), device="cpu", dtype=torch.int32, ) - if simulate_empty_batches: - # mark half the batches to have 0 tokens - zero_batches = torch.randperm(num_experts)[: num_experts // 2] - expert_num_tokens[zero_batches] = 0 + if token_distribution == "alternating_empty": + expert_num_tokens[::2] = 0 + elif token_distribution == "skewed": + expert_num_tokens.zero_() + expert_num_tokens[0] = max_tokens_per_batch + expert_num_tokens[-1] = 1 # ref outputs ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs( From 00e5cda926ffdf35fb96027979e6692d77ea27bd Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 11 Sep 2026 15:05:41 -0400 Subject: [PATCH 062/774] [Agents] Link Triton skill to JIT kernel warmup guide (#56499) Signed-off-by: mgoin --- .agents/skills/triton-kernel-writing/SKILL.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.agents/skills/triton-kernel-writing/SKILL.md b/.agents/skills/triton-kernel-writing/SKILL.md index 6dd6a8520be0..5a524c8a718a 100644 --- a/.agents/skills/triton-kernel-writing/SKILL.md +++ b/.agents/skills/triton-kernel-writing/SKILL.md @@ -26,6 +26,10 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo unimportant runtime integer scalars in `do_not_specialize`, especially those that may alternate between values such as 0 and 1, which can produce different specialization keys. +- For new warmable kernels, follow the + [JIT Kernel Warmup guide](../../../docs/contributing/jit_kernel_warmup.md) + to define the warmup wrapper and call its `register_warmup(...)` where the + runtime implementation is selected during runner setup. - The Triton compiler does not guarantee safe ordering when a kernel writes to a pointer and subsequently reads from the same pointer. This pattern must have a `tl.debug_barrier()` between the write and read. The barrier From 0c1e89ceb92b95e02b5a7aa07781340aacfcf011 Mon Sep 17 00:00:00 2001 From: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:07:30 -0700 Subject: [PATCH 063/774] [Bugfix][Kimi-K3] Fix KDA projection overlap on Hopper (#55426) Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Tyler Michael Smith --- tests/kernels/test_bf16_skinny_gemm.py | 102 ++++++++++++++++- .../models/kimi_k3/nvidia/low_latency_gemm.py | 75 ++++++++----- .../nvidia/ops/cute_dsl/kda_skinny_gemm.py | 105 +++++++++++++++--- 3 files changed, 239 insertions(+), 43 deletions(-) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 98a49c753c27..c16b4d44e8e8 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -269,6 +269,70 @@ def test_kda_overlap_configs_match_measured_table() -> None: } +@pytest.mark.parametrize( + "compute_capability,supported", + [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)], +) +def test_kda_mixed_precision_bf16_fma_capability( + monkeypatch: pytest.MonkeyPatch, + compute_capability: tuple[int, int], + supported: bool, +) -> None: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl import kda_skinny_gemm + + capability = compute_capability[0] * 10 + compute_capability[1] + monkeypatch.setattr( + kda_skinny_gemm.current_platform, + "has_device_capability", + lambda target, device_id=0: capability >= target, + ) + + assert kda_skinny_gemm._has_mixed_precision_bf16_fma() is supported + + +@pytest.mark.parametrize("num_tokens", [2, 8, 14]) +def test_kda_skinny_gemms_sm90_cuda_graph(num_tokens: int) -> None: + """Hopper must compile and capture the pre-SM100 FMA specialization.""" + _require_capability_and_cute((9, 0)) + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import ( + KdaSkinnyGemm, + ) + + torch.manual_seed(43 + num_tokens) + hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + f_ab_weight = torch.randn(144, 7168, dtype=torch.bfloat16, device="cuda") + f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") + gemm = KdaSkinnyGemm() + + gemm.run_k(gemm.run_n(hidden_states, f_ab_weight), f_b_weight) + torch.accelerator.synchronize() + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + projected_fab = gemm.run_n(hidden_states, f_ab_weight) + projected_fb = gemm.run_k(projected_fab, f_b_weight) + torch.cuda.current_stream().wait_stream(capture_stream) + graph.replay() + torch.accelerator.synchronize() + + expected_fab = torch.nn.functional.linear( + hidden_states.float(), f_ab_weight.float() + ) + expected_fb = torch.nn.functional.linear( + expected_fab[:, :128].to(torch.bfloat16).float(), f_b_weight.float() + ) + for actual, expected in ( + (projected_fab, expected_fab), + (projected_fb, expected_fb), + ): + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + @pytest.mark.parametrize("tp_size", [1, 2, 4, 16]) def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None: from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention @@ -282,8 +346,34 @@ def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None: assert kda._projection_overlap_max_tokens == 0 +@pytest.mark.parametrize( + ("capability", "expected"), + [ + ((9, 0), None), + ((10, 0), "cute-dsl"), + ((10, 3), "cute-dsl"), + ((10, 7), None), + ((12, 0), None), + ], +) +def test_kda_qkvg_flashinfer_backend( + monkeypatch: pytest.MonkeyPatch, + capability: tuple[int, int], + expected: str | None, +) -> None: + monkeypatch.setattr( + current_platform, + "is_device_capability", + lambda candidate: candidate == capability, + ) + + assert k3_gemm._kda_qkvg_flashinfer_backend() == expected + + +@pytest.mark.parametrize("backend", ["cute-dsl", None]) def test_kda_qkvg_autotune_enables_full_overlap( monkeypatch: pytest.MonkeyPatch, + backend: str | None, ) -> None: flashinfer_gemm = pytest.importorskip("flashinfer.gemm") @@ -304,10 +394,20 @@ def fake_mm_bf16(a, b, *, pdl, backend): return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype) monkeypatch.setattr(flashinfer_gemm, "mm_bf16", fake_mm_bf16) + monkeypatch.setattr( + k3_gemm, + "_kda_qkvg_flashinfer_backend", + lambda: backend, + ) k3_gemm.autotune_kda_qkvg(kda) - assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")] + expected_calls = [] + if backend is not None: + expected_calls.append( + (torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl") + ) + assert calls == expected_calls assert ( kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS ) diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 02be89aa311a..e439ec5b6693 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -65,6 +65,20 @@ _KDA_TP_SIZE = 8 +def _kda_qkvg_flashinfer_backend() -> Literal["cute-dsl"] | None: + """Return the tested FlashInfer backend for the low-M QKVG branch. + + FlashInfer's ``cute-dsl`` BF16 GEMM supports SM100 and SM103, but rejects + SM90 during backend validation. Other capabilities use ``torch.mm`` for + QKVG while retaining the concurrent F_A/beta and F_B branch. + """ + if current_platform.is_device_capability( + (10, 0) + ) or current_platform.is_device_capability((10, 3)): + return "cute-dsl" + return None + + @dataclass(frozen=True, slots=True) class ProjectionSpec: n: int @@ -749,14 +763,16 @@ def run_qkvg() -> torch.Tensor: config, None, ) - if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS: + if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS and ( + backend := _kda_qkvg_flashinfer_backend() + ): from flashinfer.gemm import mm_bf16 return mm_bf16( hidden_states, qkvg_weight.t(), pdl=True, - backend="cute-dsl", + backend=backend, ) return torch.mm(hidden_states, qkvg_weight.t()) @@ -807,11 +823,10 @@ def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]: def autotune_kda_qkvg(model: nn.Module) -> None: - """Autotune the FlashInfer QKVG GEMM before CUDA graph capture.""" - from flashinfer.gemm import mm_bf16 - + """Autotune the supported QKVG GEMM before CUDA graph capture.""" from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + backend = _kda_qkvg_flashinfer_backend() children: list[KimiK3DeltaAttention] = [] weights_by_shape: dict[ tuple[int, int, int, torch.dtype, torch.device], torch.Tensor @@ -822,29 +837,33 @@ def autotune_kda_qkvg(model: nn.Module) -> None: if child._projection_overlap_max_tokens <= 0: continue children.append(child) - qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE] - shape = ( - KDA_PROJECTION_OVERLAP_MAX_TOKENS, - qkvg_weight.shape[0], - qkvg_weight.shape[1], - qkvg_weight.dtype, - qkvg_weight.device, - ) - weights_by_shape.setdefault(shape, qkvg_weight) - - for shape, qkvg_weight in weights_by_shape.items(): - num_tokens = shape[0] - hidden_states = torch.empty( - (num_tokens, qkvg_weight.shape[1]), - dtype=qkvg_weight.dtype, - device=qkvg_weight.device, - ) - mm_bf16( - hidden_states, - qkvg_weight.t(), - pdl=True, - backend="cute-dsl", - ) + if backend is not None: + qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE] + shape = ( + KDA_PROJECTION_OVERLAP_MAX_TOKENS, + qkvg_weight.shape[0], + qkvg_weight.shape[1], + qkvg_weight.dtype, + qkvg_weight.device, + ) + weights_by_shape.setdefault(shape, qkvg_weight) + + if backend is not None: + from flashinfer.gemm import mm_bf16 + + for shape, qkvg_weight in weights_by_shape.items(): + num_tokens = shape[0] + hidden_states = torch.empty( + (num_tokens, qkvg_weight.shape[1]), + dtype=qkvg_weight.dtype, + device=qkvg_weight.device, + ) + mm_bf16( + hidden_states, + qkvg_weight.t(), + pdl=True, + backend=backend, + ) for child in children: child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py index da69e2631bd9..9e8f93330018 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py @@ -18,6 +18,8 @@ from cutlass.cute.runtime import make_fake_stream, make_fake_tensor from cutlass.cutlass_dsl import T, dsl_user_op +from vllm.platforms import current_platform + _FAB_N = 144 _FAB_K = 7168 _FAB_VECTOR_WIDTH = 16 @@ -91,14 +93,61 @@ def _fma_f32_bf16( return Float32(result) +@dsl_user_op +def _fma_f32_bf16_portable( + a: BFloat16, + b: BFloat16, + acc: Float32, + *, + loc=None, + ip=None, +) -> Float32: + """BF16 multiply with FP32 accumulation for pre-SM100 GPUs. + + PTX ``fma.f32.bf16`` requires SM100 or newer. Converting the operands to + FP32 first keeps this kernel available on Hopper, where an unconditional + mixed-precision FMA fails libNVVM compilation for ``sm_90a``. + """ + a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.f32(), + [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .f32 a_f32, b_f32;\n\t" + "cvt.f32.bf16 a_f32, $1;\n\t" + "cvt.f32.bf16 b_f32, $2;\n\t" + "fma.rn.f32 $0, a_f32, b_f32, $3;\n\t" + "}", + "=f,h,h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float32(result) + + +def _has_mixed_precision_bf16_fma() -> bool: + """Whether PTX ``fma.f32.bf16`` is supported by the current GPU.""" + return current_platform.has_device_capability(100) + + class _KdaSkinnyNGemm: - def __init__(self, num_rows: int, split_n: int) -> None: + def __init__( + self, + num_rows: int, + split_n: int, + use_mixed_precision_fma: bool, + ) -> None: if _FAB_N % split_n: raise ValueError("split_n must divide the FAB output size") if num_rows * split_n > _FAB_BLOCK_SIZE: raise ValueError("the FAB output tile must fit in one block") self.num_rows = num_rows self.split_n = split_n + self.use_mixed_precision_fma = use_mixed_precision_fma @cute.jit def __call__( @@ -191,11 +240,18 @@ def kernel( for mi in cutlass.range_constexpr(num_rows): for ni in cutlass.range_constexpr(split_n): for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH): - acc[mi, ni] = _fma_f32_bf16( - a_regs[mi, vi], - weight_regs[ni, vi], - acc[mi, ni], - ) + if cutlass.const_expr(self.use_mixed_precision_fma): + acc[mi, ni] = _fma_f32_bf16( + a_regs[mi, vi], + weight_regs[ni, vi], + acc[mi, ni], + ) + else: + acc[mi, ni] = _fma_f32_bf16_portable( + a_regs[mi, vi], + weight_regs[ni, vi], + acc[mi, ni], + ) acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) partials_layout = cute.make_layout( @@ -232,13 +288,19 @@ def kernel( class _KdaSkinnyKGemm: - def __init__(self, num_rows: int, split_n: int) -> None: + def __init__( + self, + num_rows: int, + split_n: int, + use_mixed_precision_fma: bool, + ) -> None: if _FB_N % split_n: raise ValueError("split_n must divide the FB output size") if split_n * _FB_K_LANES > 1024: raise ValueError("the FB output tile exceeds the block-size limit") self.num_rows = num_rows self.split_n = split_n + self.use_mixed_precision_fma = use_mixed_precision_fma @cute.jit def __call__( @@ -314,11 +376,18 @@ def kernel( for mi in cutlass.range_constexpr(num_rows): acc = Float32(0.0) for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH): - acc = _fma_f32_bf16( - input_regs[mi, vi], - weight_regs[vi], - acc, - ) + if cutlass.const_expr(self.use_mixed_precision_fma): + acc = _fma_f32_bf16( + input_regs[mi, vi], + weight_regs[vi], + acc, + ) + else: + acc = _fma_f32_bf16_portable( + input_regs[mi, vi], + weight_regs[vi], + acc, + ) for offset in (8, 4, 2, 1): acc += cute.arch.shuffle_sync_bfly(acc, offset) output[mi, n_idx] = acc.to(BFloat16) @@ -341,7 +410,11 @@ def _compile_n(self, num_rows: int) -> None: return split_n = SKINNY_N_SPLIT_N[num_rows] self._compiled_n[key] = cute.compile( - _KdaSkinnyNGemm(num_rows, split_n), + _KdaSkinnyNGemm( + num_rows, + split_n, + use_mixed_precision_fma=_has_mixed_precision_bf16_fma(), + ), make_fake_tensor( BFloat16, (num_rows, _FAB_K), @@ -371,7 +444,11 @@ def _compile_k(self, num_rows: int) -> None: return split_n = SKINNY_K_SPLIT_N[num_rows] self._compiled_k[key] = cute.compile( - _KdaSkinnyKGemm(num_rows, split_n), + _KdaSkinnyKGemm( + num_rows, + split_n, + use_mixed_precision_fma=_has_mixed_precision_bf16_fma(), + ), make_fake_tensor( BFloat16, (num_rows, _FB_K), From dffb863c97c16f392207625d6bccf7e12bfe10ee Mon Sep 17 00:00:00 2001 From: Lyle Lai Date: Sat, 12 Sep 2026 03:32:47 +0800 Subject: [PATCH 064/774] Upgrade tpu-inference to v0.29.0 (#56388) Signed-off-by: Lyle Lai --- requirements/tpu.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/tpu.txt b/requirements/tpu.txt index 8eaf8850d877..15e0e9ceabf8 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.28.0 +tpu-inference==0.29.0 From cc5dd0a857cde0ec224c63fee40276269e32f278 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 11 Sep 2026 12:55:34 -0700 Subject: [PATCH 065/774] [MRV1] Scope breakable cudagraphs to the piecewise path only (#56312) Signed-off-by: Nick Hill Co-authored-by: Andreas Karatzas --- .../v1/cudagraph/test_breakable_cudagraph.py | 132 +++++++++++++++++- vllm/compilation/breakable_cudagraph.py | 31 ++-- vllm/compilation/cuda_graph.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 15 +- 4 files changed, 165 insertions(+), 17 deletions(-) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index ded83d800de9..0e10e4a51f1a 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -8,7 +8,7 @@ import threading from contextlib import nullcontext -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -536,3 +536,133 @@ def outer_op(t: torch.Tensor) -> None: # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113 assert torch.equal(x, torch.full((4,), 113.0, device="cuda")) assert inner_calls == 2 # replay invokes the outer's lambda again + + +# --------------------------------------------------------------------------- +# BreakableCUDAGraphWrapper: runtime_mode dispatch gating +# --------------------------------------------------------------------------- + + +def _mock_vllm_config(): + config = MagicMock() + config.parallel_config.data_parallel_size = 1 + config.parallel_config.use_sequence_parallel_moe = False + config.parallel_config.is_moe_model = False + return config + + +def test_wrapper_falls_through_on_runtime_mode_mismatch(): + """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing + nothing, so an outer CUDAGraphWrapper can own them.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.config import CUDAGraphMode + from vllm.forward_context import BatchDescriptor, set_forward_context + + calls = [] + + def model(x): + calls.append(x) + return x + 1 + + config = _mock_vllm_config() + wrapper = BreakableCUDAGraphWrapper( + model, config, runtime_mode=CUDAGraphMode.PIECEWISE + ) + desc = BatchDescriptor(num_tokens=4) + + with set_forward_context( + None, + config, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + batch_descriptor=desc, + ): + assert wrapper(1) == 2 + + assert calls == [1] + assert not wrapper.entries + + +def test_wrapper_captures_on_runtime_mode_match(): + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.config import CUDAGraphMode + from vllm.forward_context import BatchDescriptor, set_forward_context + + calls = [] + + def model(x): + calls.append(x) + return x + 1 + + config = _mock_vllm_config() + wrapper = BreakableCUDAGraphWrapper( + model, config, runtime_mode=CUDAGraphMode.PIECEWISE + ) + desc = BatchDescriptor(num_tokens=4) + + sentinel = object() + with ( + patch.object( + BreakableCUDAGraphWrapper, "_capture", return_value=sentinel + ) as mock_capture, + set_forward_context( + None, + config, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + batch_descriptor=desc, + ), + ): + assert wrapper(1) is sentinel + + mock_capture.assert_called_once() + assert calls == [] + assert desc in wrapper.entries + + +def test_wrapper_default_matches_any_non_none_mode(): + """runtime_mode=None matches any non-NONE mode (MRV2 relies on this).""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.config import CUDAGraphMode + from vllm.forward_context import BatchDescriptor, set_forward_context + + def model(x): + return x + 1 + + config = _mock_vllm_config() + wrapper = BreakableCUDAGraphWrapper(model, config) + desc = BatchDescriptor(num_tokens=4) + + sentinel = object() + with ( + patch.object( + BreakableCUDAGraphWrapper, "_capture", return_value=sentinel + ) as mock_capture, + set_forward_context( + None, + config, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + batch_descriptor=desc, + ), + ): + assert wrapper(1) is sentinel + + mock_capture.assert_called_once() + assert desc in wrapper.entries + + +def test_unwrap_recurses_through_nested_wrappers(): + """unwrap() reaches the raw model through nested wrappers.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.compilation.cuda_graph import CUDAGraphWrapper + from vllm.config import CUDAGraphMode + + def model(x): + return x + 1 + + config = _mock_vllm_config() + inner = BreakableCUDAGraphWrapper( + model, config, runtime_mode=CUDAGraphMode.PIECEWISE + ) + outer = CUDAGraphWrapper(inner, config, runtime_mode=CUDAGraphMode.FULL) + + assert inner.unwrap() is model + assert outer.unwrap() is model diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 8e25d05c280f..ca6aa1713774 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -256,6 +256,10 @@ class BreakableCUDAGraphWrapper: * If runtime mode mismatch / NONE, run eagerly. * Otherwise, lazily capture per ``batch_descriptor`` and replay on subsequent invocations with the same descriptor. + + When ``runtime_mode`` is given, only that mode is captured/replayed; + other non-NONE modes fall through, so the wrapper can be nested under + a ``CUDAGraphWrapper`` for another mode. """ _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = ( @@ -271,15 +275,14 @@ def __init__( self, runnable: Callable[..., Any], vllm_config: VllmConfig, + runtime_mode: CUDAGraphMode | None = None, ) -> None: - # Unlike the original CUDAGraphWrapper which strictly matches a - # single runtime_mode, this wrapper captures whatever the - # dispatcher emits (any non-NONE runtime_mode) -- breakable's - # capture is identical for prefill and decode, so there's nothing - # to dispatch on at the runtime_mode level. Entries are keyed by - # BatchDescriptor which already encodes batch shape / uniformity. + # runtime_mode=None matches any non-NONE dispatch mode; entries + # are keyed by BatchDescriptor, which already encodes batch + # shape / uniformity. self.runnable = runnable self.vllm_config = vllm_config + self.runtime_mode = runtime_mode self.compilation_config = vllm_config.compilation_config self.graph_pool = current_platform.get_global_graph_pool() self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG" @@ -298,7 +301,9 @@ def __getattr__(self, key: str) -> Any: raise AttributeError(key) def unwrap(self) -> Callable[..., Any]: - return self.runnable + # Recurse through nested wrappers to reach the original runnable. + runnable = self.runnable + return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable @property def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper: @@ -317,11 +322,13 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: batch_descriptor = forward_context.batch_descriptor cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode - # Capture whenever the dispatcher says "some cudagraph mode" -- - # breakable produces the same artifact regardless of PIECEWISE - # vs FULL, so we match either. Entries are keyed by batch - # descriptor, which already encodes prefill/decode distinctions. - if cudagraph_runtime_mode == CUDAGraphMode.NONE: + # By default (runtime_mode=None) match any non-NONE mode; a set + # runtime_mode restricts capture so other modes can be handled by + # an outer wrapper (e.g. CUDAGraphWrapper for FULL). + if cudagraph_runtime_mode == CUDAGraphMode.NONE or ( + self.runtime_mode is not None + and cudagraph_runtime_mode != self.runtime_mode + ): return self.runnable(*args, **kwargs) assert batch_descriptor is not None diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index b63d86199720..55771608837d 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -221,7 +221,9 @@ def __getattr__(self, key: str) -> Any: def unwrap(self) -> Callable[..., Any]: # in case we need to access the original runnable. - return self.runnable + runnable = self.runnable + # Recurse through nested wrappers. + return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable @property def cudagraph_wrapper(self) -> "CUDAGraphWrapper": diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a14d00bd58b1..4ff50f82b515 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5461,13 +5461,22 @@ def load_model(self, load_dummy_weights: bool = False) -> None: and cudagraph_mode != CUDAGraphMode.NONE and not self.parallel_config.use_ubatching ): - self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config) + # Scoped to PIECEWISE dispatch; FULL cudagraphs (below) are + # unaffected. PIECEWISE dispatch can also arise after wrapping + # (drafters under a FULL target mode, or a later FULL -> + # FULL_AND_PIECEWISE upgrade in _check_and_update_cudagraph_mode). + self.model = BreakableCUDAGraphWrapper( + self.model, self.vllm_config, runtime_mode=CUDAGraphMode.PIECEWISE + ) drafter = getattr(self, "drafter", None) if drafter is not None and hasattr(drafter, "model"): drafter.model = BreakableCUDAGraphWrapper( - drafter.model, self.vllm_config + drafter.model, + self.vllm_config, + runtime_mode=CUDAGraphMode.PIECEWISE, ) - elif ( + + if ( cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.use_ubatching ): From c1b69aa0d49b88016b5b2156ae61649c2c2df282 Mon Sep 17 00:00:00 2001 From: mevince <26759093+mevince@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:12:59 -0700 Subject: [PATCH 066/774] [Bugfix][KV Connector] Only enforce disk block alignment for O_DIRECT (#55424) Signed-off-by: Vincent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/simple_kv_offload/test_worker.py | 57 +++++++++++++++++++++++ vllm/v1/simple_kv_offload/disk_backend.py | 9 ++-- 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 66582f03bbd3..cbc7f2061330 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -43,6 +43,7 @@ build_params, pin_tensor, ) +from vllm.v1.simple_kv_offload.disk_backend import DiskBackend from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker from vllm.v1.worker.utils import allocate_kv_cache @@ -519,3 +520,59 @@ def test_mixed_page_byte_placement_is_dcp_invariant(): ) assert placements[0] == placements[1] + + +@pytest.mark.parametrize( + ("block_bytes", "use_page_cache", "expect_ok"), + [ + (4096, True, True), # aligned control + (3 * 512, True, True), # 512-aligned only: fine for buffered I/O + (3 * 512, False, False), # O_DIRECT keeps enforcing 4096 alignment + ], +) +def test_disk_backend_alignment_only_required_for_direct_io( + tmp_path, monkeypatch, block_bytes: int, use_page_cache: bool, expect_ok: bool +): + """The 4096 stride assert belongs to O_DIRECT; page-cache I/O must not hit it. + + Hybrid models (e.g. 101 per-rank cache tensors) sum to a block stride that + is 512- but not 4096-aligned, which previously made disk offload unbootable + even with use_page_cache=True. + """ + # Host pinning is a separate lifecycle concern with its own tests; the + # alignment gate must not depend on registration state, so stub both ends. + monkeypatch.setattr( + "vllm.v1.simple_kv_offload.disk_backend.pin_tensor", lambda t: None + ) + monkeypatch.setattr( + "vllm.v1.simple_kv_offload.disk_backend.unpin_tensor", + lambda t: None, + raising=False, # harmless on branches without teardown-side unpins + ) + num_blocks = 4 + gpu = {"k": torch.zeros((num_blocks, block_bytes), dtype=torch.int8, device="cuda")} + backend = DiskBackend() + disk_path = tmp_path / "kv-offload.bin" + + def _init() -> None: + backend.init( + gpu, + gpu["k"].device, + torch.cuda.Stream(), + torch.cuda.Stream(), + str(disk_path), + num_blocks, + block_bytes, + use_page_cache=use_page_cache, + ) + + if not expect_ok: + with pytest.raises(AssertionError, match="not aligned"): + _init() + return + + _init() + try: + assert disk_path.stat().st_size == num_blocks * block_bytes + finally: + backend.shutdown() diff --git a/vllm/v1/simple_kv_offload/disk_backend.py b/vllm/v1/simple_kv_offload/disk_backend.py index f39be7d05f56..dfa717cc9ac1 100644 --- a/vllm/v1/simple_kv_offload/disk_backend.py +++ b/vllm/v1/simple_kv_offload/disk_backend.py @@ -99,9 +99,12 @@ def init( t.stride(0) * t.element_size() for t in gpu_caches.values() ] - assert total_block_bytes % _ALIGNMENT == 0, ( - f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}" - ) + # The stride alignment is an O_DIRECT kernel-DMA requirement; buffered + # I/O through the page cache has no alignment constraint. + if not use_page_cache: + assert total_block_bytes % _ALIGNMENT == 0, ( + f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}" + ) # Separate buffer pools for store and load threads self._store_buffer_caches = {} From 9dcf6bf344caa7793bae0b45a7896d3f8e03a01a Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:33:53 -0700 Subject: [PATCH 067/774] [BugFix] Fix DP token padding in dflash attention metadata (#56181) Signed-off-by: Giancarlo Delfin --- .../test_eagle_draft_attn_metadata.py | 64 ++++++++++------ .../spec_decode/autoregressive/speculator.py | 14 ++-- .../gpu/spec_decode/dflash/speculator.py | 36 +-------- .../multi_module_mtp/speculator.py | 5 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 75 +++++++++++-------- 5 files changed, 97 insertions(+), 97 deletions(-) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index 95e263f9de8d..0cc78f3371c2 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the EAGLE speculator's draft attention metadata builder. -These tests guard the regression where ``_build_draft_attn_metadata`` did +These tests guard the regression where ``_build_uniform_attn_metadata`` did not populate ``seq_lens_cpu_upper_bound`` on the per-step ``CommonAttentionMetadata``. Several downstream attention backends and helpers (``split_decodes_prefills_and_extends``, the MLA indexer, @@ -13,11 +13,14 @@ AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None """ -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import patch +import numpy as np import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator @@ -30,7 +33,7 @@ def _make_fake_speculator( draft_max_seq_len: int = 1024, ) -> SimpleNamespace: """Build a fake EagleSpeculator with just the attributes used by - ``_build_draft_attn_metadata``. We deliberately avoid constructing a + ``_build_uniform_attn_metadata``. We deliberately avoid constructing a real ``EagleSpeculator`` because that requires a full ``VllmConfig`` and a draft model. """ @@ -46,8 +49,8 @@ def _make_fake_speculator( cp_rank=0, cp_interleave=1, ) - return SimpleNamespace( - arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"), + fake = SimpleNamespace( + arange_np=np.arange(max_num_reqs + 1, dtype=np.int32), draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool), block_tables=fake_block_tables, input_buffers=fake_input_buffers, @@ -56,21 +59,42 @@ def _make_fake_speculator( max_model_len=max_model_len, draft_max_seq_len=draft_max_seq_len, ) + # The uniform wrapper delegates through self; bind the real implementation. + fake._build_attn_metadata = MethodType(EagleSpeculator._build_attn_metadata, fake) + return fake -def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step): +def _run_build( + fake, + *, + num_reqs, + num_reqs_padded, + base, + step, + num_query_per_req=1, + cg_mode=CUDAGraphMode.FULL, +): captured: dict[str, object] = {} def fake_build_attn_metadata(**kwargs): captured.update(kwargs) return {} + # Request padding only occurs under FULL graphs, where the captured batch + # is a uniform decode of num_reqs_padded * num_query_per_req tokens. + batch_desc = BatchExecutionDescriptor( + cg_mode=cg_mode, + num_tokens=num_reqs_padded * num_query_per_req, + num_reqs=num_reqs_padded, + uniform_token_count=num_query_per_req, + ) + with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata): - EagleSpeculator._build_draft_attn_metadata( + EagleSpeculator._build_uniform_attn_metadata( fake, # type: ignore[arg-type] + batch_desc=batch_desc, num_reqs=num_reqs, - num_reqs_padded=num_reqs_padded, - num_tokens_padded=num_tokens_padded, + num_query_per_req=num_query_per_req, seq_lens_cpu_upper_bound=base, step=step, ) @@ -85,9 +109,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): fake = _make_fake_speculator() base = torch.tensor([100, 200, 300, 0], dtype=torch.int32) - captured = _run_build( - fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2 - ) + captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2) bound = captured["seq_lens_cpu_upper_bound"] assert isinstance(bound, torch.Tensor), ( @@ -102,20 +124,19 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32)) -def test_build_draft_attn_metadata_handles_zero_unpadded_reqs(): - """Edge case: when ``num_reqs == 0`` the upper-bound tensor must - still be a valid all-zero tensor of length ``num_reqs_padded``.""" +def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): + """The padded tail of the upper-bound tensor is zeroed, so it stays a + valid tensor of length ``num_reqs_padded`` regardless of padding.""" fake = _make_fake_speculator() base = torch.zeros(2, dtype=torch.int32) - captured = _run_build( - fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1 - ) + captured = _run_build(fake, num_reqs=1, num_reqs_padded=2, base=base, step=1) bound = captured["seq_lens_cpu_upper_bound"] assert isinstance(bound, torch.Tensor) assert bound.shape == (2,) - assert torch.equal(bound, torch.zeros(2, dtype=torch.int32)) + # base is all zeros, so the real entry is 0 + step and the pad entry is 0. + assert torch.equal(bound, torch.tensor([1, 0], dtype=torch.int32)) def test_build_draft_attn_metadata_clamps_to_max_model_len(): @@ -124,9 +145,7 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len(): fake = _make_fake_speculator(max_model_len=1024) base = torch.tensor([1023, 500], dtype=torch.int32) - captured = _run_build( - fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3 - ) + captured = _run_build(fake, num_reqs=2, num_reqs_padded=2, base=base, step=3) bound = captured["seq_lens_cpu_upper_bound"] # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected. @@ -154,7 +173,6 @@ def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave): fake, num_reqs=3, num_reqs_padded=4, - num_tokens_padded=4, base=torch.tensor([5, 9, 16]), step=0, ) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index cab250ff4e68..3d243b6ee5da 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -532,11 +532,10 @@ def _multi_step_decode( slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, self.kv_cache_config ) - attn_metadata = self._build_draft_attn_metadata( + attn_metadata = self._build_uniform_attn_metadata( num_reqs=num_reqs, - num_reqs_padded=batch_desc.num_reqs or num_reqs, - # One query per request; exclude DP-only model padding. - num_tokens_padded=batch_desc.num_reqs or num_reqs, + batch_desc=batch_desc, + num_query_per_req=1, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, step=step, ) @@ -581,11 +580,10 @@ def _fused_multi_step_decode( slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, self.kv_cache_config ) - attn_metadata = self._build_draft_attn_metadata( + attn_metadata = self._build_uniform_attn_metadata( num_reqs=num_reqs, - num_reqs_padded=batch_desc.num_reqs or num_reqs, - # One query per request; exclude DP-only model padding. - num_tokens_padded=batch_desc.num_reqs or num_reqs, + batch_desc=batch_desc, + num_query_per_req=1, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, step=1, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 845781fab227..a5710829e0ff 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -from collections.abc import Mapping from typing import Any -import numpy as np import torch import torch.nn as nn @@ -299,33 +297,6 @@ def _generate_draft( num_reqs, self.num_speculative_steps ) - def _build_draft_attn_metadata( - self, - num_reqs: int, - num_reqs_padded: int, - num_tokens_padded: int, - seq_lens_cpu_upper_bound: torch.Tensor, - step: int, - num_query_per_req: int | None = None, - causal: bool | Mapping[int, bool] = False, - query_start_loc_np: np.ndarray | None = None, - dcp_local_seq_lens: torch.Tensor | None = None, - ) -> dict[str, Any] | None: - if not self.draft_attn_layer_names: - return None - assert num_query_per_req is None # Omitted for DFlash, read from self instead - return super()._build_draft_attn_metadata( - num_reqs, - num_reqs_padded, - num_tokens_padded, - seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, - step=step, - num_query_per_req=self.num_query_per_req, - causal=causal, - query_start_loc_np=query_start_loc_np, - dcp_local_seq_lens=dcp_local_seq_lens, - ) - @torch.inference_mode() def propose( self, @@ -471,7 +442,6 @@ def propose( need_eager=is_profile, dp_sync=batch_sync, ) - num_reqs_padded = batch_desc.num_reqs or num_reqs num_tokens_padded = batch_desc.num_tokens num_tokens_across_dp = ( batch_sync.num_tokens_across_dp if batch_sync is not None else None @@ -479,10 +449,10 @@ def propose( # Rebuild the draft attention metadata even when replaying the FULL # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_draft_attn_metadata( + draft_attn_metadata = self._build_uniform_attn_metadata( num_reqs=num_reqs, - num_reqs_padded=num_reqs_padded, - num_tokens_padded=num_tokens_padded, + batch_desc=batch_desc, + num_query_per_req=self.num_query_per_req, seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, step=self.num_query_per_req, causal=self._group_causal, diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 90572c03ccd5..1c6ae315cf46 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -226,10 +226,9 @@ def propose( slot_mappings = build_slot_mappings_by_layer( slot_mappings_tensor, self.kv_cache_config ) - draft_attn_metadata = self._build_draft_attn_metadata( + draft_attn_metadata = self._build_attn_metadata( num_reqs=num_reqs, - num_reqs_padded=batch_desc.num_reqs or num_reqs, - num_tokens_padded=batch_desc.num_tokens, + batch_desc=batch_desc, query_start_loc_np=input_batch.query_start_loc_np, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, step=0, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 2681677c7603..4f20637d7020 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -23,6 +23,7 @@ ) from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -149,9 +150,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): dtype=torch.int64, device=device, ) - self.arange = torch.arange( - self.max_num_reqs + 1, dtype=torch.int32, device="cpu" - ) + self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32) self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) self.draft_logits: torch.Tensor | None = None @@ -254,42 +253,37 @@ def set_attn( self.target_input_buffers = target_input_buffers self.target_attn_groups = target_attn_groups - def _build_draft_attn_metadata( + def _build_attn_metadata( self, num_reqs: int, - num_reqs_padded: int, - num_tokens_padded: int, + batch_desc: BatchExecutionDescriptor, + query_start_loc_np: np.ndarray, seq_lens_cpu_upper_bound: torch.Tensor, step: int, - num_query_per_req: int = 1, causal: bool | Mapping[int, bool] = True, - query_start_loc_np: np.ndarray | None = None, dcp_local_seq_lens: torch.Tensor | None = None, ) -> dict[str, Any] | None: - if query_start_loc_np is not None: - # Non-uniform query layout (e.g. multi-module MTP's mixed - # prefill/decode queries); num_query_per_req is ignored. - query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) - query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( - query_start_loc_np[: num_reqs + 1] - ) - query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] - max_query_len = int( - (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max() - ) - else: - # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req. - # Clamp keeps the series non-decreasing past num_reqs, which some - # attention backends require. - query_start_loc_cpu = ( - torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs) - * num_query_per_req - ) - max_query_len = num_query_per_req + num_reqs_padded = batch_desc.num_reqs or num_reqs + # A FULL graph replays a captured shape whose padded requests each hold + # a full query width, so attention must see the padded token count. + # PIECEWISE/eager needs the actual token count, because batch_desc may + # carry graph or DP padding that no request owns, which would desync it + # from query_start_loc. + num_tokens = ( + batch_desc.num_tokens + if batch_desc.cg_mode == CUDAGraphMode.FULL + else int(query_start_loc_np[-1]) + ) + query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) + query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( + query_start_loc_np[: num_reqs + 1] + ) + query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] + max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()) block_tables = [ x[:num_reqs_padded] for x in self.block_tables.input_block_tables ] - slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded] + slot_mappings = self.block_tables.slot_mappings[:, :num_tokens] draft_seq_lens_cpu_upper_bound = torch.zeros( num_reqs_padded, dtype=torch.int32, device="cpu" ) @@ -313,7 +307,7 @@ def _build_draft_attn_metadata( attn_metadata = build_attn_metadata( attn_groups=self.attn_groups, num_reqs=num_reqs_padded, - num_tokens=num_tokens_padded, + num_tokens=num_tokens, query_start_loc_gpu=self.input_buffers.query_start_loc[ : num_reqs_padded + 1 ], @@ -434,3 +428,24 @@ def _build_uniform_batch_dp_sync( uniform_token_count=num_query_per_req, eager=False, ), num_batch_tokens + + def _build_uniform_attn_metadata( + self, + batch_desc: BatchExecutionDescriptor, + num_reqs: int, + num_query_per_req: int, + seq_lens_cpu_upper_bound: torch.Tensor, + step: int, + causal: bool | Mapping[int, bool] = True, + dcp_local_seq_lens: torch.Tensor | None = None, + ) -> dict[str, Any] | None: + query_start_loc_np = self.arange_np[: num_reqs + 1] * num_query_per_req + return self._build_attn_metadata( + num_reqs=num_reqs, + batch_desc=batch_desc, + query_start_loc_np=query_start_loc_np, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + step=step, + causal=causal, + dcp_local_seq_lens=dcp_local_seq_lens, + ) From 2d75e586fcaf88231f7a75f482dc8bfb5ad9da10 Mon Sep 17 00:00:00 2001 From: amd-xavierwang Date: Fri, 11 Sep 2026 17:34:17 -0400 Subject: [PATCH 068/774] [ROCm][Kernel][DSV4] Remove tl.constexpr to avoid cold-compile churn in indexer gather kernel (#56153) Signed-off-by: xawang --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index dd2392d6c239..598baee0b0b0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -292,16 +292,16 @@ def _cp_gather_indexer_quant_cache_kernel( HEAD_DIM: tl.constexpr, BLOCK_TILE_SIZE: tl.constexpr, HEAD_TILE_SIZE: tl.constexpr, - NUM_TOKENS: tl.constexpr, - NUM_BATCHES: tl.constexpr, - BLOCK_TABLE_WIDTH: tl.constexpr, - NUM_BLOCKS: tl.constexpr, + num_tokens, + num_batches, + block_table_width, + num_blocks, ): tid = tl.program_id(0) offset = tl.arange(0, HEAD_DIM) - valid_tid = tid < NUM_TOKENS + valid_tid = tid < num_tokens batch_id = tl.load(token_to_seq_ptr + tid, mask=valid_tid, other=-1) - valid_batch = (batch_id >= 0) & (batch_id < NUM_BATCHES) + valid_batch = (batch_id >= 0) & (batch_id < num_batches) safe_batch_id = tl.where(valid_batch, batch_id, 0) batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0) batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0) @@ -314,7 +314,7 @@ def _cp_gather_indexer_quant_cache_kernel( valid_block_table = ( valid_token & (block_table_id >= 0) - & (block_table_id < BLOCK_TABLE_WIDTH) + & (block_table_id < block_table_width) & (block_offset >= 0) & (block_offset < block_size) ) @@ -323,7 +323,7 @@ def _cp_gather_indexer_quant_cache_kernel( block_id = tl.load( block_table_ptr + block_table_offset, mask=valid_block_table, other=-1 ) - valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS) + valid_block = valid_block_table & (block_id >= 0) & (block_id < num_blocks) # The packed KV layout makes per-block strides large # enough that block_id * stride can exceed 32-bit range. safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64) From bbaa1b93402dbc3c9741b2d81f8f2ebac74ba339 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 11 Sep 2026 17:11:59 -0500 Subject: [PATCH 069/774] [ROCm][CI] Extend timeout for `Basic Models (other)` (#56522) Signed-off-by: Micah Williamson --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/models_basic.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 8d2ba218be11..efe0a0f8bce3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1973,7 +1973,7 @@ steps: - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - label: ":amd: (MI300) Basic Models (Other)" # TBD - timeout_in_minutes: 55 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 855ec8c056cd..24447cb50171 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -56,7 +56,7 @@ steps: label: ":amd: (MI300) Basic Models (Other)" dind: false device: mi300_1 - timeout_in_minutes: 50 + timeout_in_minutes: 65 depends_on: - image-build-amd source_file_dependencies: From 988d9b6777d077f843cd2164a222ac8535d36ed2 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:31:55 -0700 Subject: [PATCH 070/774] [Rust Frontend][gRPC] Surface engine generation errors (#56405) Signed-off-by: Alec Flowers Co-authored-by: Codex --- rust/src/server/src/grpc/convert.rs | 35 ++++++++-------- rust/src/server/src/grpc/inference.rs | 31 +++++++------- rust/src/server/src/grpc/tests.rs | 59 +++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 32 deletions(-) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 2a3f41ffa438..ca44f59c93ea 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -391,13 +391,14 @@ pub fn to_sequence_output( logprobs: Option<&DecodedLogprobs>, finished: Option<&Finished>, opts: &ResponseOpts, -) -> pb::SequenceOutput { +) -> Result { + let finish_info = finished.map(|f| to_finish_info(f, token_ids)).transpose()?; let (lp_values, rank_values, candidates) = match logprobs { Some(lp) if opts.output_logprobs => output_logprobs_to_proto(lp), _ => (vec![], vec![], vec![]), }; - pb::SequenceOutput { + Ok(pb::SequenceOutput { index: 0, // TODO: multi-sequence (n > 1) not supported text: if opts.output_text { delta.to_string() @@ -413,11 +414,11 @@ pub fn to_sequence_output( logprobs: lp_values, ranks: rank_values, candidate_tokens: candidates, - finish_info: finished.map(|f| to_finish_info(f, token_ids)), - } + finish_info, + }) } -fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo { +fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result { use pb::finish_info::FinishReason as PbFinishReason; let (finish_reason, stop_reason) = match &finished.finish_reason { @@ -438,18 +439,17 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo { (PbFinishReason::Stop as i32, sr) } FinishReason::Length => (PbFinishReason::Length as i32, None), - FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => { - (PbFinishReason::Aborted as i32, None) - } + FinishReason::Error => return Err(Status::internal("engine failed during generation")), + FinishReason::Abort | FinishReason::Repetition(_) => (PbFinishReason::Aborted as i32, None), }; - pb::FinishInfo { + Ok(pb::FinishInfo { num_output_tokens: finished.usage.output_token_count as u32, finish_reason, stop_reason, kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct), ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct), - } + }) } // ======================================================================================== @@ -709,7 +709,7 @@ mod tests { let fin = finished(FinishReason::Stop(None)); let token_ids = [1_u32, 2, 3, 151643]; - let info = to_finish_info(&fin, &token_ids); + let info = to_finish_info(&fin, &token_ids).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Stop as i32); assert_eq!(info.stop_reason, Some(PbStopReason::EosTokenId(151643))); @@ -719,7 +719,7 @@ mod tests { fn eos_stop_with_empty_token_ids_leaves_stop_reason_unset() { let fin = finished(FinishReason::Stop(None)); - let info = to_finish_info(&fin, &[]); + let info = to_finish_info(&fin, &[]).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Stop as i32); assert_eq!(info.stop_reason, None); @@ -730,7 +730,7 @@ mod tests { let fin = finished(FinishReason::Stop(Some(StopReason::TokenId(42)))); // Terminal token list should be ignored when an explicit stop reason is // present. - let info = to_finish_info(&fin, &[7, 42]); + let info = to_finish_info(&fin, &[7, 42]).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Stop as i32); assert_eq!(info.stop_reason, Some(PbStopReason::StopTokenId(42))); @@ -740,7 +740,7 @@ mod tests { fn explicit_stop_string_is_preserved() { let fin = finished(FinishReason::Stop(Some(StopReason::Text("".into())))); - let info = to_finish_info(&fin, &[1, 2, 3]); + let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Stop as i32); assert_eq!( @@ -753,7 +753,7 @@ mod tests { fn length_finish_has_no_stop_reason() { let fin = finished(FinishReason::Length); - let info = to_finish_info(&fin, &[1, 2, 3]); + let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Length as i32); assert_eq!(info.stop_reason, None); @@ -763,7 +763,7 @@ mod tests { fn abort_finish_is_mapped_to_aborted() { let fin = finished(FinishReason::Abort); - let info = to_finish_info(&fin, &[]); + let info = to_finish_info(&fin, &[]).expect("finish info"); assert_eq!(info.finish_reason, PbFinishReason::Aborted as i32); assert_eq!(info.stop_reason, None); @@ -778,7 +778,8 @@ mod tests { ..Default::default() }; - let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts); + let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts) + .expect("sequence output"); let finish = out.finish_info.expect("finish_info should be present"); assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32); diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index f864cbde5959..7051558cd882 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -241,12 +241,6 @@ impl pb::inference_server::Inference for InferenceServiceImpl { .instrument(request_span.clone()) .await .map_err(|error| log_text_error(&request_span, started_at, "collection", error))?; - info!( - parent: &request_span, - elapsed_ms = started_at.elapsed().as_millis() as u64, - "gRPC inference request completed" - ); - // Build the single aggregated response. let prompt_info = convert::to_prompt_info( &collected.prompt_token_ids, @@ -267,6 +261,11 @@ impl pb::inference_server::Inference for InferenceServiceImpl { collected.logprobs.as_ref(), Some(&finish_info), &response_opts, + )?; + info!( + parent: &request_span, + elapsed_ms = started_at.elapsed().as_millis() as u64, + "gRPC inference request completed" ); Ok(Response::new(pb::GenerateResponse { @@ -331,19 +330,21 @@ impl pb::inference_server::Inference for InferenceServiceImpl { logprobs, }, finished, - }) => Ok(pb::GenerateResponse { + }) => convert::to_sequence_output( + &decoded.text, + &token_ids, + logprobs.as_ref(), + finished.as_deref(), + &response_opts, + ) + .map(|outputs| pb::GenerateResponse { prompt_info: None, - outputs: Some(convert::to_sequence_output( - &decoded.text, - &token_ids, - logprobs.as_ref(), - finished.as_deref(), - &response_opts, - )), + outputs: Some(outputs), }), }; - if tx.send(response).await.is_err() { + let failed = response.is_err(); + if tx.send(response).await.is_err() || failed { break; } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index ea4364480cde..59403277ee94 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -680,6 +680,65 @@ async fn unary_generate_returns_collected_text() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn engine_error_finish_returns_internal_for_unary_and_streaming() { + for streaming in [false, true] { + for prefix in [vec![], vec![b'h' as u32]] { + let mut output_specs = Vec::new(); + if !prefix.is_empty() { + output_specs.push((prefix.clone(), None)); + } + output_specs.push((vec![], Some(EngineCoreFinishReason::Error))); + let (mut client, server_task, engine_task) = + grpc_test_server(b"engine-grpc-error", output_specs).await; + let request = pb::GenerateRequest { + request_id: "test-engine-error".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + response: Some(pb::ResponseOptions { + output_token_ids: true, + ..Default::default() + }), + ..Default::default() + }; + + let status = if streaming { + let mut stream = + client.generate_stream(request).await.expect("start stream").into_inner(); + let mut received_tokens = Vec::new(); + let status = loop { + match stream.message().await { + Ok(Some(response)) => { + if let Some(output) = response.outputs { + assert!( + output.finish_info.is_none(), + "error became a finish event" + ); + received_tokens.extend(output.token_ids); + } + } + Ok(None) => panic!("engine error ended as a successful stream"), + Err(status) => break status, + } + }; + assert_eq!(received_tokens, prefix); + status + } else { + client.generate(request).await.expect_err("engine error must fail the RPC") + }; + assert_eq!(status.code(), tonic::Code::Internal); + + engine_task.await.expect("mock engine task"); + server_task.abort(); + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_with_token_ids_prompt() { From 9d88ceb02694c6e3df182ec3c87201284f50e2c6 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:53:24 -0400 Subject: [PATCH 071/774] [KDA] Update flashKDA to support bf16 checkpoint state (#56485) Signed-off-by: wzhao18 --- cmake/external_projects/flashkda.cmake | 2 +- tests/models/kimi_k3/test_kda.py | 21 ++++++++++++++------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 6ffb46e6a793..5ac4f4b415ea 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG 3b225bf26bb8e218928a1fe14751cb48cf31d11b + GIT_TAG b59532f1f464fbd536272780e30df5bf6a2ccc02 GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 1fd4982e767e..84a5e7eed409 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -1421,14 +1421,21 @@ def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2)) +@pytest.mark.parametrize( + ("state_dtype", "tolerance"), + [ + pytest.param(torch.bfloat16, 0.03, id="bf16"), + pytest.param(torch.float32, 0.01, id="fp32"), + ], +) @torch.inference_mode() -def test_flashkda_checkpoint_correctness(): +def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: float): lower_bound = -3.0 - _require_kda_prefill_backend("flashkda", torch.float32, lower_bound) + _require_kda_prefill_backend("flashkda", state_dtype, lower_bound) import vllm._flashkda_C # noqa: F401 - inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound) + inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound) q, k, v = inputs.q, inputs.k, inputs.v raw_g, raw_beta = inputs.raw_g, inputs.raw_beta A_log, dt_bias = inputs.A_log, inputs.dt_bias @@ -1480,12 +1487,12 @@ def test_flashkda_checkpoint_correctness(): checkpoint_offsets=checkpoint_offsets, ) - assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01) - assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01) - assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01) + assert_close("checkpoint_o", expected_out, checkpoint_out, tolerance) + assert_close("checkpoint_ht", expected_state, checkpoint_final_state, tolerance) + assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], tolerance) conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE) - recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE) + recurrent_storage = torch.zeros(2, H * D * D + 8, dtype=state_dtype, device=DEVICE) recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D) conv_input = q[0].flatten(1) checkpoint_state_indices = torch.tensor( From 46d2b23ac5047a813ebb082122166e4ae09b5f39 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Fri, 11 Sep 2026 20:35:48 -0400 Subject: [PATCH 072/774] [ROCm][DSV4.1][Perf] Use AITER mHC for the delayed pre block (#56503) Signed-off-by: Fangzhou Ai Co-authored-by: Fangzhou Ai --- tests/kernels/test_mhc_kernels.py | 95 ++++++++++++ vllm/_aiter_ops.py | 131 +++++++++++++++++ vllm/model_executor/kernels/mhc/__init__.py | 2 + vllm/model_executor/kernels/mhc/aiter.py | 90 ++++++++++++ vllm/model_executor/kernels/mhc/triton.py | 131 +++++++++++++++++ vllm/model_executor/layers/mhc.py | 155 ++++++++++++++++++++ vllm/models/deepseek_v4_1/amd/model.py | 35 ++--- 7 files changed, 620 insertions(+), 19 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 61e0449cae19..ab5efb6578c6 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -23,6 +23,7 @@ HAS_AITER_MHC_PRE_NORM, HAS_TILELANG_MHC, MHCFusedPostPreOp, + MHCPreDelayedOp, MHCPreOp, ) from vllm.models.deepseek_v4.nvidia.model import ( @@ -697,6 +698,100 @@ def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch): torch.testing.assert_close(out[3], expected_layer_input) +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024]) +@pytest.mark.parametrize("carried", [False, True]) +def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried): + """AITER must reproduce the delayed reference on both seam variants. + + ``num_tokens`` spans the split-k choices AITER makes for the projection, + since the pre-mix is recovered from that unreduced output. + """ + set_random_seed(0) + hc_mult, hidden_size = 4, 5120 + residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs( + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult + ) + pre_mix = ( + torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 + if carried + else None + ) + rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6 + args = ( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + 1.0, + 20, + ) + + expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix) + actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix) + + for i in (0, 1, 3): + torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) + # The collapse is the same FP32 multiply-and-sum in both paths. + torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0) + + +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch): + """The broadcast seam and a fused norm must not take the AITER path. + + Neither is expressible with AITER's pre kernels: the broadcast projects a + narrower ``x``, and ``mhc_pre_gemm_sqrsum`` folds no RMSNorm. What is + asserted here is the routing decision, which is what this gate owns; the + numerics of whichever fallback it lands on are covered by + ``test_deepseek_v41_mhc_pre_delayed``. + """ + set_random_seed(0) + hc_mult, hidden_size = 4, 5120 + residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs( + num_tokens=4, hidden_size=hidden_size, hc_mult=hc_mult + ) + args = (residual, fn, hc_scale, hc_base, 1e-6, 1e-6, 1e-6, 1.0, 20) + op = object.__new__(MHCPreDelayedOp) + + aiter_op = torch.ops.vllm.mhc_pre_delayed_aiter + took_aiter = False + + def spy(*spy_args, **spy_kwargs): + nonlocal took_aiter + took_aiter = True + return aiter_op(*spy_args, **spy_kwargs) + + monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy) + + x = residual[:, 0].contiguous() + broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1) + broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x) + actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x) + assert not took_aiter, "the broadcast seam must not reach AITER" + for i in range(4): + torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) + + op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6) + assert not took_aiter, "a fused norm must not reach AITER" + + # Positive control: the same inputs without a norm do take the AITER path, + # so the two declines above are the gate discriminating rather than the + # path being unavailable in this environment. + op.forward_hip(*args) + assert took_aiter + + @pytest.mark.skipif( not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM), reason="AITER mHC with fused RMSNorm required", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 2e1e8fc6194c..454aba175e80 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -3599,6 +3599,137 @@ def mhc_pre( layer_input.view(*outer_shape, hidden_size), ) + @staticmethod + def mhc_pre_delayed( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """mHC pre using the pre-mix carried from the previous sublayer. + + Same gates as :meth:`mhc_pre`, but the stream collapse applies the + caller's ``pre_mix`` instead of the one computed here, and the one + computed here is returned for the next sublayer seam. AITER has no + single kernel for that shape, so drive its two stages directly: + ``mhc_pre_big_fuse`` still produces the post and comb gates (including + every sinkhorn iteration), while the pre gate is recovered from the + same split-k GEMM output and the collapse is done by the Triton + kernel. ``mhc_pre_big_fuse`` also writes a collapse we do not use; + that redundant store is the price of not having a native delayed + kernel, and is small next to the ~140 launches it replaces. + + Returns: + post_mix: shape (..., hc_mult, 1), dtype torch.float32 + comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 + layer_input: shape (..., hidden_size), dtype torch.bfloat16 + next_pre_mix: shape (..., hc_mult), dtype torch.float32 + """ + from aiter.ops.mhc import ( + get_mhc_pre_splitk, + mhc_pre_big_fuse, + mhc_pre_gemm_sqrsum, + ) + + assert residual.dtype == torch.bfloat16 + assert fn.dtype == torch.float32 + assert hc_scale.dtype == torch.float32 + assert hc_base.dtype == torch.float32 + + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult + hc_hidden_size = hc_mult * hidden_size + + assert fn.shape == (hc_mult3, hc_hidden_size) + assert hc_scale.shape == (3,) + assert hc_base.shape == (hc_mult3,) + + outer_shape = residual.shape[:-2] + residual_flat = residual.view(-1, hc_mult, hidden_size) + num_tokens = residual_flat.shape[0] + device = residual_flat.device + + if num_tokens == 0: + return ( + torch.empty(0, hc_mult, 1, dtype=torch.float32, device=device), + torch.empty(0, hc_mult, hc_mult, dtype=torch.float32, device=device), + torch.empty(0, hidden_size, dtype=torch.bfloat16, device=device), + torch.empty(0, hc_mult, dtype=torch.float32, device=device), + ) + + pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult) + + # AITER's Python wrappers allocate without explicit device arguments. + with torch.device(device): + splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size) + # AITER pads the GEMM output to a multiple of 32 columns. + gemm_pad = torch.empty( + splitk, + num_tokens, + (hc_mult3 + 31) // 32 * 32, + dtype=torch.float32, + device=device, + ) + gemm_out = gemm_pad[:, :, :hc_mult3] + sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device) + mhc_pre_gemm_sqrsum(gemm_out, sqrsum, residual_flat, fn, tile_k, 0) + + post_mix = torch.empty( + num_tokens, hc_mult, 1, dtype=torch.float32, device=device + ) + comb_mix = torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device + ) + unused_collapse = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=device + ) + mhc_pre_big_fuse( + post_mix, + comb_mix, + unused_collapse, + gemm_out, + sqrsum, + hc_scale, + hc_base, + residual_flat, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + + next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton( + gemm_out, + sqrsum, + hc_scale, + hc_base, + hc_mult, + hc_hidden_size, + rms_eps, + hc_pre_eps, + ) + + if pre_mix_flat is None: + # Model entry selects residual stream zero. + layer_input = residual_flat[:, 0] + else: + layer_input = torch.ops.vllm.hc_collapse_triton(residual_flat, pre_mix_flat) + + return ( + post_mix.view(*outer_shape, hc_mult, 1), + comb_mix.view(*outer_shape, hc_mult, hc_mult), + layer_input.view(*outer_shape, hidden_size), + next_pre_mix.view(*outer_shape, hc_mult), + ) + @staticmethod def hc_head( hs_flat: torch.Tensor, diff --git a/vllm/model_executor/kernels/mhc/__init__.py b/vllm/model_executor/kernels/mhc/__init__.py index 324e3cec6400..f9df7f6f2c33 100644 --- a/vllm/model_executor/kernels/mhc/__init__.py +++ b/vllm/model_executor/kernels/mhc/__init__.py @@ -12,6 +12,7 @@ "mhc_fused_post_pre_cuda", "hc_head_fused_kernel_cuda", "mhc_pre_aiter", + "mhc_pre_delayed_aiter", "mhc_post_aiter", "mhc_fused_post_pre_aiter", "hc_head_fused_aiter", @@ -26,6 +27,7 @@ "mhc_post_torch", "mhc_fused_post_pre_torch", "hc_head_fused_torch", + "mhc_pre_mix_triton", "mhc_pre_triton", "mhc_post_triton", "mhc_fused_post_pre_triton", diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index ce7690ec1d21..e04ce3322bb7 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -104,6 +104,90 @@ def _mhc_pre_aiter_fake( return post_mix, comb_mix, layer_input +def mhc_pre_delayed_aiter( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """mHC pre with the pre-mix carried in from the previous sublayer. + + Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix* + rather than the gate computed here, and that gate is returned as the + pre-mix for the next sublayer seam. + + Args: + residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 + fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32 + hc_scale: shape (3,), dtype torch.float32 + hc_base: shape (hc_mult3,), dtype torch.float32 + rms_eps: RMS normalization epsilon + hc_pre_eps: pre-mix epsilon + hc_sinkhorn_eps: sinkhorn epsilon + hc_post_mult_value: post-mix multiplier value + sinkhorn_repeat: number of sinkhorn iterations + pre_mix: shape (..., hc_mult) from the previous sublayer, or None at + model entry to select residual stream zero. + + Returns: + post_mix: shape (..., hc_mult, 1), dtype torch.float32 + comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 + layer_input: shape (..., hidden_size), dtype torch.bfloat16 + next_pre_mix: shape (..., hc_mult), dtype torch.float32 + """ + hidden_size = residual.shape[-1] + assert hidden_size % 256 == 0 + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.mhc_pre_delayed( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix, + ) + + +def _mhc_pre_delayed_aiter_fake( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + outer_shape = residual.shape[:-2] + return ( + torch.empty( + *outer_shape, hc_mult, 1, dtype=torch.float32, device=residual.device + ), + torch.empty( + *outer_shape, hc_mult, hc_mult, dtype=torch.float32, device=residual.device + ), + torch.empty( + *outer_shape, hidden_size, dtype=torch.bfloat16, device=residual.device + ), + torch.empty(*outer_shape, hc_mult, dtype=torch.float32, device=residual.device), + ) + + def mhc_post_aiter( x: torch.Tensor, residual: torch.Tensor, @@ -226,6 +310,12 @@ def _mhc_fused_post_pre_aiter_fake( mutates_args=[], fake_impl=_mhc_pre_aiter_fake, ) +direct_register_custom_op( + op_name="mhc_pre_delayed_aiter", + op_func=mhc_pre_delayed_aiter, + mutates_args=[], + fake_impl=_mhc_pre_delayed_aiter_fake, +) direct_register_custom_op( op_name="mhc_post_aiter", op_func=mhc_post_aiter, diff --git a/vllm/model_executor/kernels/mhc/triton.py b/vllm/model_executor/kernels/mhc/triton.py index d94d71a6ea1c..8724113779c8 100644 --- a/vllm/model_executor/kernels/mhc/triton.py +++ b/vllm/model_executor/kernels/mhc/triton.py @@ -141,6 +141,129 @@ def _hc_collapse_triton_fake(x: Tensor, pre_mix: Tensor) -> Tensor: return torch.empty(x.shape[0], x.shape[2], dtype=x.dtype, device=x.device) +@triton.jit +def _mhc_pre_mix_kernel( + gemm_ptr, + sqrsum_ptr, + hc_scale_ptr, + hc_base_ptr, + out_ptr, + splitk, + hc_mult: tl.constexpr, + gemm_stride_k, + gemm_stride_t, + gemm_stride_j, + sqrsum_stride_k, + sqrsum_stride_t, + out_stride_t, + out_stride_j, + inv_hc_hidden, + rms_eps, + hc_pre_eps, + SPLITK_BLOCK: tl.constexpr, + HC_BLOCK: tl.constexpr, +): + """Recover the pre-mix gate from a split-k mHC pre GEMM output.""" + token_idx = tl.program_id(0).to(tl.int64) + ks = tl.arange(0, SPLITK_BLOCK) + js = tl.arange(0, HC_BLOCK) + kmask = ks < splitk + jmask = js < hc_mult + + # Only the first hc_mult GEMM columns feed the pre gate. + gemm = tl.load( + gemm_ptr + + ks[:, None] * gemm_stride_k + + token_idx * gemm_stride_t + + js[None, :] * gemm_stride_j, + mask=kmask[:, None] & jmask[None, :], + other=0.0, + ).to(tl.float32) + mixes = tl.sum(gemm, 0) + + sqrsum = tl.load( + sqrsum_ptr + ks * sqrsum_stride_k + token_idx * sqrsum_stride_t, + mask=kmask, + other=0.0, + ).to(tl.float32) + rstd = tl.rsqrt(tl.sum(sqrsum, 0) * inv_hc_hidden + rms_eps) + + scale = tl.load(hc_scale_ptr).to(tl.float32) + base = tl.load(hc_base_ptr + js, mask=jmask, other=0.0).to(tl.float32) + + pre = tl.sigmoid(mixes * rstd * scale + base) + hc_pre_eps + tl.store(out_ptr + token_idx * out_stride_t + js * out_stride_j, pre, mask=jmask) + + +def mhc_pre_mix_triton( + gemm_out: Tensor, + sqrsum: Tensor, + hc_scale: Tensor, + hc_base: Tensor, + hc_mult: int, + hc_hidden_size: int, + rms_eps: float, + hc_pre_eps: float, +) -> Tensor: + """Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output. + + AITER's ``mhc_pre_big_fuse`` consumes the unreduced ``[splitk, tokens, + hc_mult3]`` GEMM output and the matching row square-sums, but only returns + the post and comb gates. The delayed formulation also needs the pre gate, + to carry into the next sublayer seam. It is the same slice of the same + numbers, so recover it here rather than repeating the projection. + """ + assert gemm_out.ndim == 3 and gemm_out.dtype == torch.float32 + assert sqrsum.ndim == 2 and sqrsum.dtype == torch.float32 + splitk, num_tokens = gemm_out.shape[0], gemm_out.shape[1] + assert sqrsum.shape == (splitk, num_tokens) + + out = torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=gemm_out.device) + if num_tokens == 0: + return out + + _mhc_pre_mix_kernel[(num_tokens,)]( + gemm_out, + sqrsum, + hc_scale, + hc_base, + out, + splitk, + hc_mult, + gemm_out.stride(0), + gemm_out.stride(1), + gemm_out.stride(2), + sqrsum.stride(0), + sqrsum.stride(1), + out.stride(0), + out.stride(1), + 1.0 / hc_hidden_size, + rms_eps, + hc_pre_eps, + SPLITK_BLOCK=triton.next_power_of_2(splitk), + HC_BLOCK=triton.next_power_of_2(hc_mult), + num_warps=1, + # Match the separate FP32 multiply and add in the Torch reference. + enable_fp_fusion=False, + ) + return out + + +def _mhc_pre_mix_triton_fake( + gemm_out: Tensor, + sqrsum: Tensor, + hc_scale: Tensor, + hc_base: Tensor, + hc_mult: int, + hc_hidden_size: int, + rms_eps: float, + hc_pre_eps: float, +) -> Tensor: + return torch.empty( + gemm_out.shape[1], hc_mult, dtype=torch.float32, device=gemm_out.device + ) + + def hc_head_reduce_triton_kernel( x: torch.Tensor, hc_fn: torch.Tensor, @@ -216,3 +339,11 @@ def _hc_head_triton( mutates_args=[], fake_impl=_hc_collapse_triton_fake, ) + + +direct_register_custom_op( + op_name="mhc_pre_mix_triton", + op_func=mhc_pre_mix_triton, + mutates_args=[], + fake_impl=_mhc_pre_mix_triton_fake, +) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 6ddf4ff0660f..14f1018597d9 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -292,6 +292,161 @@ def forward_cpu( ) +# --8<-- [start:mhc_pre_delayed] +@CustomOp.register("mhc_pre_delayed") +class MHCPreDelayedOp(CustomOp): + """MHC pre block using the pre-mix carried from the previous sublayer. + + Same gates as :class:`MHCPreOp`, but the stream collapse applies the + caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next + sublayer seam. Returns post_mix, comb_mix, layer_input, next_pre_mix. + """ + + # --8<-- [end:mhc_pre_delayed] + @classmethod + def enabled(cls) -> bool: + return True + + def forward_cuda( + self, + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + x: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops.vllm.mhc_pre_delayed_tilelang( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix, + x, + norm_weight, + norm_eps, + ) + + def forward_hip( + self, + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + x: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual` + # and folds no RMSNorm, so it cannot serve the model-entry broadcast + # (which projects a narrower `x`) or a fused norm. Both are handled by + # TileLang, or by the reference when TileLang is unavailable. + if ( + x is None + and norm_weight is None + and _aiter_mhc_supported(residual, None, supports_norm=False) + ): + return torch.ops.vllm.mhc_pre_delayed_aiter( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix, + ) + if HAS_TILELANG_MHC: + return torch.ops.vllm.mhc_pre_delayed_tilelang( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix, + x, + norm_weight, + norm_eps, + ) + return self.forward_native( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix, + x, + norm_weight, + norm_eps, + ) + + def forward_native( + self, + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + x: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + post_mix, comb_mix, layer_input, next_pre_mix = ( + mhc_kernels.mhc_pre_delayed_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + x=x, + ) + ) + return ( + post_mix, + comb_mix, + _apply_mhc_norm(layer_input, norm_weight, norm_eps), + next_pre_mix, + ) + + # --8<-- [start:mhc_post] @CustomOp.register("mhc_post") class MHCPostOp(CustomOp): diff --git a/vllm/models/deepseek_v4_1/amd/model.py b/vllm/models/deepseek_v4_1/amd/model.py index 069ff5f931a6..898d3ab7a025 100644 --- a/vllm/models/deepseek_v4_1/amd/model.py +++ b/vllm/models/deepseek_v4_1/amd/model.py @@ -17,15 +17,12 @@ ) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger -from vllm.model_executor.kernels.mhc.torch import ( - mhc_post_torch, - mhc_pre_delayed_torch, -) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import MHCPostOp, MHCPreDelayedOp from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -248,11 +245,13 @@ def __init__( ), requires_grad=False, ) + self.mhc_pre_delayed = MHCPreDelayedOp() + self.mhc_post = MHCPostOp() @staticmethod def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: - """Reference hyper-connection collapse used by DSpark on ROCm.""" - return (pre_mix.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype) + """Hyper-connection collapse used by DSpark on ROCm.""" + return torch.ops.vllm.hc_collapse_triton(x, pre_mix) def forward( self, @@ -275,7 +274,7 @@ def forward( # copies and the identity pre-mix selects copy 0. assert self.hc_attn_fn_broadcast is not None residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous() - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch( + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn_broadcast, self.hc_attn_scale, @@ -289,7 +288,7 @@ def forward( ) else: residual = x - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch( + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -302,7 +301,7 @@ def forward( pre_mix=pre_mix, ) else: - residual = mhc_post_torch(x, residual, post_mix, res_mix) + residual = self.mhc_post(x, residual, post_mix, res_mix) if self.engram is not None and engram_hashes is not None: # Engram injection happens between the previous sublayer's # post and this block's pre, on the full hc stream, so the @@ -312,7 +311,7 @@ def forward( engram_hashes[:, self.engram.layer_hash_index], engram_mask, ) - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch( + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -333,8 +332,8 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual = mhc_post_torch(x, residual, post_mix, res_mix) - post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_torch( + residual = self.mhc_post(x, residual, post_mix, res_mix) + post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, self.hc_ffn_fn, self.hc_ffn_scale, @@ -453,6 +452,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): else: self.norm = PPMissingLayer() + self.mhc_post = MHCPostOp() + spec_config = vllm_config.speculative_config needs_mtp_hidden_states = spec_config is not None and ( spec_config.use_eagle() or spec_config.uses_draft_model() @@ -603,7 +604,7 @@ def forward( ) if idx + 1 in self.aux_hidden_state_layers: # Reconstruct the aux hidden state for draft models - aux_recon = mhc_post_torch(hidden_states, residual, post_mix, res_mix) + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) aux_hidden_state = aux_recon.mean(dim=1) if self.use_sequence_parallel: aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] @@ -614,7 +615,7 @@ def forward( if self.end_layer in self.aux_hidden_state_layers: hidden_states = final_aux_recon else: - hidden_states = mhc_post_torch( + hidden_states = self.mhc_post( hidden_states, residual, post_mix, res_mix ) @@ -635,11 +636,7 @@ def forward( # mixes — the mix the reference applies via # ``last_layer.hc_pre(h, pre_mix)`` (v4.1 has no learned hc_head). assert pre_mix is not None - hidden_states = ( - (pre_mix.unsqueeze(-1) * hidden_states.float()) - .sum(dim=1) - .to(hidden_states.dtype) - ) + hidden_states = torch.ops.vllm.hc_collapse_triton(hidden_states, pre_mix) hidden_states = self.norm(hidden_states) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states From 56d72701380bc68955cc0aa9367636df8356e3c8 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 11 Sep 2026 18:06:09 -0700 Subject: [PATCH 073/774] [Bugfix][NIXL] Don't evict a remote engine a transfer is still reading from (#54689) Signed-off-by: Yizheng Jiao --- .../kv_connector/unit/test_nixl_connector.py | 42 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 21 +++++++++- 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 760a21c4c85d..282458d1a954 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2444,6 +2444,48 @@ def test_engine_ttl_eviction(default_vllm_config, dist_init): mock_rem.assert_any_call("agent_1") +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) +def test_engine_with_inflight_transfer_is_not_evicted(default_vllm_config, dist_init): + """A transfer that outlives the TTL must keep its engine registered. + + _engine_last_active is stamped when a read is issued and not refreshed + while it runs, so a stalled transfer -- a peer that has lost its NIC holds + one indefinitely -- leaves its engine looking idle. Evicting it releases + the dlist handle and remote agent the transfer is still reading through. + """ + worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0) + nixl_wrapper = worker.nixl_wrapper + + request_id = "req-still-reading" + worker._recving_transfers[request_id] = [MagicMock()] + worker._recving_metadata[request_id] = MagicMock( + remote=MagicMock(engine_id=engine_id) + ) + + with ( + patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel, + patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem, + ): + worker._engine_last_active[engine_id] = time.perf_counter() - 20.0 + + worker._evict_stale_engines() + + assert engine_id in worker._remote_agents + assert engine_id in worker.dst_xfer_side_handles + mock_rel.assert_not_called() + mock_rem.assert_not_called() + + # Once the transfer is done the engine is stale like any other. + del worker._recving_transfers[request_id] + worker._evict_stale_engines() + + assert engine_id not in worker._remote_agents + assert mock_rem.call_count == 2 + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index dd2e97c96578..e4b46cb64a96 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -3143,7 +3143,9 @@ def _evict_stale_engines(self) -> None: to be "optimal" until a new handshake is performed. Engines with active transfers or pending handshakes cannot be stale: - - Active transfers touch _engine_last_active in start_load_kv. + - Reads stamp _engine_last_active when they are issued, and engines a + transfer is still reading from are held back explicitly, since the + stamp is not refreshed while the read runs. - Pending handshakes don't have an _engine_last_active entry yet """ # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number @@ -3154,10 +3156,25 @@ def _evict_stale_engines(self) -> None: return now = time.perf_counter() + busy = self._engines_with_inflight_transfers() for eid, last_active in list(self._engine_last_active.items()): - if now - last_active > self._engine_ttl: + if now - last_active > self._engine_ttl and eid not in busy: self._cleanup_remote_engine(eid) + def _engines_with_inflight_transfers(self) -> set[EngineId]: + """Remote engines a transfer is still reading from. + + The timestamp is stamped when a read is issued and not refreshed while + it runs, so a transfer that outlives the TTL leaves its engine looking + idle. A peer that has lost its NIC holds one indefinitely. + """ + return { + meta.remote.engine_id + for req_id in self._recving_transfers + if (meta := self._recving_metadata.get(req_id)) is not None + and meta.remote is not None + } + def _cleanup_remote_engine( self, engine_id: EngineId, *, log_eviction: bool = True ) -> None: From 80650454dcf1a0b758b31ce59c004f49897963db Mon Sep 17 00:00:00 2001 From: Tahsin Tunan Date: Sat, 12 Sep 2026 07:08:05 +0600 Subject: [PATCH 074/774] [CI] Give every Buildkite step an explicit key, and a hook to enforce it (#56329) Signed-off-by: Tahsin Tunan --- .../intel_jobs/basic_correctness_intel.yaml | 1 + .buildkite/intel_jobs/engine_intel.yaml | 2 + .buildkite/intel_jobs/kernels_intel.yaml | 1 + .buildkite/intel_jobs/lora_intel.yaml | 6 ++ .buildkite/intel_jobs/misc_intel.yaml | 8 +- .../intel_jobs/model_runner_v2_intel.yaml | 4 + .buildkite/intel_jobs/test-intel.yaml | 7 ++ .buildkite/test_areas/docker.yaml | 1 + .buildkite/test_areas/lm_eval.yaml | 1 + .buildkite/test_areas/rust_frontend.yaml | 5 + .pre-commit-config.yaml | 7 ++ tools/pre_commit/check_buildkite_step_keys.py | 97 +++++++++++++++++++ 12 files changed, 139 insertions(+), 1 deletion(-) create mode 100644 tools/pre_commit/check_buildkite_step_keys.py diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 292920c47ea1..173baf40f62c 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-xpu steps: - label: XPU Sleep Mode + key: xpu-sleep-mode timeout_in_minutes: 45 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index d0aabd6d3299..2619e1c58830 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -41,6 +41,7 @@ steps: pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py' - label: Engine (1 GPU) + key: engine-1-gpu timeout_in_minutes: 30 device: intel_gpu agent_tags: @@ -65,6 +66,7 @@ steps: VLLM_XPU_ENABLE_XPU_GRAPH=1 pytest -v -s test_config.py' - label: V1 e2e (2 GPUs) + key: v1-e2e-2-gpus timeout_in_minutes: 30 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index e914c10d8bfe..da6086c1dc43 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-xpu steps: - label: vLLM IR Tests + key: vllm-ir-tests timeout_in_minutes: 30 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index cab0cb66aa45..2cecbf18f773 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-xpu steps: - label: LoRA Runtime + Utils + key: lora-runtime-utils timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -36,6 +37,7 @@ steps: pytest -v -s lora/test_worker.py' - label: LoRA Fused/MoE Kernels + key: lora-fused-moe-kernels timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -60,6 +62,7 @@ steps: pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' - label: LoRA Punica Kernels + key: lora-punica-kernels timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -86,6 +89,7 @@ steps: pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit' - label: LoRA Punica FP8/XPU Ops + key: lora-punica-fp8-xpu-ops timeout_in_minutes: 60 device: intel_gpu agent_tags: @@ -110,6 +114,7 @@ steps: pytest -v -s lora/test_punica_xpu_ops.py' - label: LoRA Models + key: lora-models timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -137,6 +142,7 @@ steps: pytest -s -v lora/test_minicpmv_tp.py' - label: LoRA Multimodal + key: lora-multimodal timeout_in_minutes: 45 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 752870abe36b..3818339662ed 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-xpu steps: - label: V1 Core + KV + Metrics + key: v1-core-kv-metrics-intel timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -32,9 +33,10 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && pytest -v -s v1/executor && - VLLM_BATCH_INVARIANT=1 pytest -v -s -m 'not cpu_test' v1/ec_connector/unit' + VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit' - label: V1 Sample + Logits + key: v1-sample-logits-intel timeout_in_minutes: 90 device: intel_gpu agent_tags: @@ -80,6 +82,7 @@ steps: parallelism: 4 - label: Basic Models Tests (Initialization) + key: basic-models-tests-initialization timeout_in_minutes: 60 device: intel_gpu agent_tags: @@ -104,6 +107,7 @@ steps: pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' - label: XPU CPU Offload + key: xpu-cpu-offload timeout_in_minutes: 60 device: intel_gpu agent_tags: @@ -131,6 +135,7 @@ steps: pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' - label: NixlConnector PD accuracy (4 GPUs) + key: nixlconnector-pd-accuracy-4-gpus timeout_in_minutes: 60 num_devices: 4 device: intel_gpu @@ -269,6 +274,7 @@ steps: pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' - label: Fusion Unit Tests + key: fusion-unit-tests timeout_in_minutes: 30 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml index 1f68333e976f..7f6767eba2f2 100644 --- a/.buildkite/intel_jobs/model_runner_v2_intel.yaml +++ b/.buildkite/intel_jobs/model_runner_v2_intel.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-xpu steps: - label: Model Runner V2 Core Tests (Intel) + key: model-runner-v2-core-tests-intel timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -34,6 +35,7 @@ steps: pytest -v -s v1/e2e/general/test_min_tokens.py' - label: Model Runner V2 Examples (Intel) + key: model-runner-v2-examples-intel timeout_in_minutes: 45 device: intel_gpu agent_tags: @@ -64,6 +66,7 @@ steps: python3 features/automatic_prefix_caching/prefix_caching_offline.py' - label: Model Runner V2 Distributed (2 GPUs) + key: model-runner-v2-distributed-2-gpus timeout_in_minutes: 50 device: intel_gpu agent_tags: @@ -90,6 +93,7 @@ steps: TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"' - label: Model Runner V2 Spec Decode + key: model-runner-v2-spec-decode timeout_in_minutes: 50 device: intel_gpu agent_tags: diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 0bc25261a2b3..00f172ca885d 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -15,6 +15,7 @@ steps: - exit_status: -10 # Agent was lost limit: 2 - label: "XPU example Test" + key: xpu-example-test depends_on: - image-build-xpu timeout_in_minutes: 50 @@ -50,6 +51,7 @@ steps: python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192 ' - label: "XPU W8A8 FP8 Linear Examples" + key: xpu-w8a8-fp8-linear-examples depends_on: - image-build-xpu timeout_in_minutes: 60 @@ -77,6 +79,7 @@ steps: python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096 ' - label: "XPU V1 test" + key: xpu-v1-test depends_on: - image-build-xpu timeout_in_minutes: 70 @@ -107,6 +110,7 @@ steps: pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py' parallelism: 2 - label: "XPU server test" + key: xpu-server-test depends_on: - image-build-xpu timeout_in_minutes: 45 @@ -130,6 +134,7 @@ steps: pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py && pytest -v -s benchmarks/test_serve_cli.py' - label: "XPU quantization test" + key: xpu-quantization-test depends_on: - image-build-xpu timeout_in_minutes: 30 @@ -158,6 +163,7 @@ steps: pytest -v -s quantization/test_auto_round.py && pytest -v -s quantization/test_online.py' - label: "XPU compressed tensors FP8 test" + key: xpu-compressed-tensors-fp8-test depends_on: - image-build-xpu timeout_in_minutes: 60 @@ -182,6 +188,7 @@ steps: pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8 && pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8_block_enabled' - label: "XPU Graph" + key: xpu-graph depends_on: - image-build-xpu timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index e88f19ad9793..ef50062a7f1c 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -3,6 +3,7 @@ depends_on: - image-build-cpu steps: - label: ":computer: (CPU) Docker Build Metadata" + key: docker-build-metadata timeout_in_minutes: 20 device: cpu-small source_file_dependencies: diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 0863dfd2767e..dc90dcb8529c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -470,6 +470,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" - label: ":nvidia: (H200 MIG 35GB) MRCR Eval Small Models" + key: mrcr-eval-small-models device: h200_35gb timeout_in_minutes: 25 source_file_dependencies: diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 8ccc0c50e926..97078512a1d7 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -3,6 +3,7 @@ depends_on: - image-build steps: - label: ":nvidia: (H200 MIG 18GB) Rust Frontend OpenAI Coverage" + key: rust-frontend-openai-coverage timeout_in_minutes: 30 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -67,6 +68,7 @@ steps: - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Serve/Admin Coverage" + key: rust-frontend-serve-admin-coverage timeout_in_minutes: 25 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -116,6 +118,7 @@ steps: - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Core Correctness" + key: rust-frontend-core-correctness timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -130,6 +133,7 @@ steps: - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use" + key: rust-frontend-tool-use device: h200_35gb timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -145,6 +149,7 @@ steps: - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" - label: ":nvidia: (L4) Rust Frontend Distributed" + key: rust-frontend-distributed timeout_in_minutes: 25 device: l4 num_devices: 4 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d1d0e9e3d15b..b12c8ce508ea 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -177,6 +177,13 @@ repos: files: ^\.github/mergify\.yml$ pass_filenames: false additional_dependencies: [PyYAML, regex] + - id: check-buildkite-step-keys + name: Check every Buildkite step declares an explicit key + language: python + entry: python tools/pre_commit/check_buildkite_step_keys.py + files: ^\.buildkite/.*\.ya?ml$ + pass_filenames: false + additional_dependencies: [PyYAML] - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward name: Run mypy for Python 3.10 entry: python tools/pre_commit/mypy.py "3.10" diff --git a/tools/pre_commit/check_buildkite_step_keys.py b/tools/pre_commit/check_buildkite_step_keys.py new file mode 100644 index 000000000000..41d65b4bc6db --- /dev/null +++ b/tools/pre_commit/check_buildkite_step_keys.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check that every Buildkite step the CI generator reads declares a key. + +A step with no `key:` is published under one the generator derives from its +label, which makes the job's address a function of its title: retitling the +step silently renames it and breaks retries, `depends_on:` and build history. +The generator only reads the directories a pipeline config names in +`job_dirs:`, so this walks the same files and fails on any step without a key. + +Usage: + python tools/pre_commit/check_buildkite_step_keys.py +""" + +import sys +from pathlib import Path + +import yaml + +BUILDKITE = Path(".buildkite") +CONFIG_GLOB = "ci_config*.yaml" +JOB_GLOB = "*.yaml" + + +def job_dirs() -> list[Path]: + """Every directory the pipeline generator reads, per the pipeline configs.""" + configs = sorted(BUILDKITE.glob(CONFIG_GLOB)) + if not configs: + raise SystemExit(f"no pipeline config matched {BUILDKITE / CONFIG_GLOB}") + dirs: set[Path] = set() + for config in configs: + with open(config, encoding="utf-8") as f: + doc = yaml.safe_load(f) + declared = doc.get("job_dirs") if isinstance(doc, dict) else None + if not isinstance(declared, list) or not declared: + raise SystemExit(f"{config}: `job_dirs` is missing or not a list") + for name in declared: + job_dir = Path(name) + if not job_dir.is_dir(): + raise SystemExit(f"{config}: job_dir `{name}` does not exist") + dirs.add(job_dir) + return sorted(dirs) + + +def scan(steps: object, source: Path) -> list[str]: + """Label of every step in a `steps:` list that declares no key.""" + if not isinstance(steps, list): + raise SystemExit(f"{source}: `steps:` is not a list") + keyless: list[str] = [] + for i, step in enumerate(steps): + # A scalar shorthand such as `- wait` has nowhere to put a key. Every + # other step does, so there are no other exemptions. A `mirror:` block + # is not a step here: the generator publishes it as `-`. + if not isinstance(step, dict): + continue + label = step.get("label") or step.get("group") or f"step {i}" + if "steps" in step: + raise SystemExit( + f"{source}: `{label}` nests its own `steps:`. The generator has " + "no group step and drops the children without an error, so give " + "each one its own entry." + ) + key = step.get("key") + if not (isinstance(key, str) and key.strip()): + keyless.append(label) + return keyless + + +def keyless_steps(path: Path) -> list[str]: + with open(path, encoding="utf-8") as f: + doc = yaml.safe_load(f) + if not isinstance(doc, dict) or "steps" not in doc: + return [] + return scan(doc["steps"], path) + + +def main() -> int: + files = sorted({p for d in job_dirs() for p in d.rglob(JOB_GLOB)}) + if not files: + raise SystemExit("no pipeline files under the declared job_dirs") + keyless = [(path, label) for path in files for label in keyless_steps(path)] + if not keyless: + return 0 + print(f"{len(keyless)} Buildkite step(s) have no explicit key:\n", file=sys.stderr) + for path, label in keyless: + print(f" {path}: {label}", file=sys.stderr) + print( + "\nAdd a `key:` naming the test each step runs. Without one the pipeline" + "\ngenerator derives the key from the label, so retitling the step renames" + "\nits job and breaks retries and build history.", + file=sys.stderr, + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 06e57f622cf53ed1c3f8d8295cdd48ea2c133bd1 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:17:39 -0400 Subject: [PATCH 075/774] [ROCm][Kimi-K3] Fix non-contiguous state_indices crash and GPU-sync assert in fused KDA/MLA prefill (#56526) Signed-off-by: Divakar Verma --- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 4 ++-- vllm/v1/attention/backends/mla/rocm_aiter_mla.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index efcdb1bc75d1..748581771b82 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -286,9 +286,9 @@ def fused_kda_chunk( None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32), None if checkpoint_state_indices is None - else checkpoint_state_indices.to(torch.int32), + else checkpoint_state_indices.to(torch.int32).contiguous(), state_cache, - None if state_indices is None else state_indices.to(torch.int32), + None if state_indices is None else state_indices.to(torch.int32).contiguous(), has_initial_state, ) return out, final_state diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index c341bfc3486e..b850deb07091 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -22,6 +22,7 @@ QueryLenSupport, ) from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor from vllm.v1.attention.backend import ( AttentionCGSupport, @@ -796,7 +797,8 @@ def _build_fp8_prefill_ps_metadata( # build) keeps it off the per-layer forward path where a sync would # break CUDA Graph capture. Using the device-side reduce_indptr is # acceptable since build is allowed to incur an occasional sync. - num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item()) + with gpu_sync_allowed(): + num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item()) # Attach PS metadata to the metadata object so forward_mha can read it. metadata.fp8_prefill_qo_indptr = qo_indptr From a0844fa6c6d6e1143917a033ad6984e40c212862 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 12 Sep 2026 09:24:30 +0800 Subject: [PATCH 076/774] [XPU][CI] skip DeepSeek-V4.1-Flash in `test_tensor_schema.py` (#56463) Signed-off-by: zhenwei-intel --- .buildkite/intel_jobs/models_multimodal_intel.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index d99a7debb0b5..44efa8e0f130 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -125,6 +125,7 @@ steps: pip install open-clip-torch --no-deps && cd tests && pytest -v -s models/multimodal/processing/test_tensor_schema.py + --deselect="tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[deepseek-ai/DeepSeek-V4.1-Flash]" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' parallelism: 4 From 31759ccb6d560e150849273b47e15680a0120169 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 11 Sep 2026 19:29:49 -0700 Subject: [PATCH 077/774] [Rust Frontend][Multimodal] Accept preprocessed multimodal gRPC features (#55047) Signed-off-by: Biswa Panda Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Codex --- rust/Cargo.lock | 2 +- rust/Cargo.toml | 2 +- rust/proto/Cargo.toml | 2 +- rust/proto/README.md | 2 +- rust/proto/inference.proto | 13 + rust/src/chat/src/error.rs | 3 + rust/src/chat/src/lib.rs | 19 +- rust/src/chat/src/multimodal.rs | 42 ++- rust/src/chat/src/multimodal/input.rs | 29 ++ rust/src/chat/src/multimodal/preprocessed.rs | 274 +++++++++++++++ rust/src/server/Cargo.toml | 2 +- rust/src/server/src/grpc/convert.rs | 91 ----- rust/src/server/src/grpc/inference.rs | 18 +- rust/src/server/src/grpc/media.rs | 341 +++++++++++++++++++ rust/src/server/src/grpc/mod.rs | 1 + rust/src/server/src/grpc/tests.rs | 258 +++++++++++++- 16 files changed, 986 insertions(+), 113 deletions(-) create mode 100644 rust/src/chat/src/multimodal/input.rs create mode 100644 rust/src/chat/src/multimodal/preprocessed.rs create mode 100644 rust/src/server/src/grpc/media.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index bf19fd3225ce..80edfe4ebc91 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5940,7 +5940,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.1.0" +version = "0.2.0" dependencies = [ "prost", "prost-types", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 75104642c345..bd9d99de623c 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -147,7 +147,7 @@ vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } -vllm-proto = { version = "0.1.0", path = "proto" } +vllm-proto = { version = "0.2.0", path = "proto" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index 3084c8c5b364..130fd69a5231 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.1.0" +version = "0.2.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/README.md b/rust/proto/README.md index 3c5dbe0d3a61..62b134280bd9 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -17,7 +17,7 @@ Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to ```toml [dependencies] -vllm-proto = "0.1" +vllm-proto = "0.2" ``` For example, import `vllm_proto::inference_client::InferenceClient` or diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index c854a9abdf1f..858d8a7a7e01 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -228,7 +228,20 @@ message MediaItem { string url = 2; // http:// or https:// string data_uri = 3; // data: bytes raw_bytes = 4; + PreprocessedMediaFeatures features = 7; } string mime_type = 5; string uuid = 6; } + +// Engine-ready multimodal features produced by a version-compatible frontend. +// `kwargs` is the MessagePack wire representation of one MmKwargsItem. +message PreprocessedMediaFeatures { + // Required inline payload. Cache-only references are not accepted. + optional bytes kwargs = 1; + string identifier = 2; + uint64 offset = 3; + uint64 length = 4; + optional string mm_hash = 5; + repeated bool is_embed = 6; +} diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index e3299dfdd4e1..a8304f17aa18 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -27,6 +27,8 @@ pub enum Error { UnsupportedModality { modality: String }, #[error("At most {limit} {modality}(s) may be provided in one prompt.")] MmLimitExceeded { modality: String, limit: usize }, + #[error("invalid inline multimodal features: {message}")] + InvalidPreprocessedMultimodal { message: String }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -106,6 +108,7 @@ impl Error { Self::UnsupportedMultimodalRenderer | Self::UnsupportedMultimodalContent(_) | Self::UnsupportedModality { .. } + | Self::InvalidPreprocessedMultimodal { .. } | Self::MmLimitExceeded { .. } => true, _ => false, diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 6250330e4868..1b5d5aea6f92 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -130,7 +130,7 @@ impl ChatRequestProcessor { /// Prepare media for an already-tokenized request. async fn prepare_media( &self, - media: Vec, + media: multimodal::MultimodalInput, token_ids: &mut Vec, ) -> Result> { if media.is_empty() { @@ -140,8 +140,15 @@ impl ChatRequestProcessor { .backend .multimodal_model_info() .ok_or(Error::UnsupportedMultimodalRenderer)?; - let model_dtype = self.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?; - let features = info.prepare_multimodal(media, token_ids, model_dtype).await?; + let features = match media { + multimodal::MultimodalInput::Raw(parts) => { + let model_dtype = self.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?; + info.prepare_multimodal(parts, token_ids, model_dtype).await? + } + multimodal::MultimodalInput::Preprocessed(features) => { + info.prepare_preprocessed(features, token_ids.len())? + } + }; Ok(Some(features)) } @@ -286,12 +293,14 @@ impl ChatLlm { } /// Prepare media for an already-tokenized request. + /// Raw content is preprocessed; inline features are checked against model + /// capabilities, modality limits, and the final prompt positions. pub async fn prepare_media( &self, - media: Vec, + media: impl Into, token_ids: &mut Vec, ) -> Result> { - self.processor.prepare_media(media, token_ids).await + self.processor.prepare_media(media.into(), token_ids).await } /// Effective tool-call parser name for this model, if parsing is enabled. diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index a0743b8769b6..674979e6ce63 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -39,11 +39,14 @@ use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; mod audio; mod expand; mod image; +mod input; mod item; +mod preprocessed; mod tensor; mod video; use self::expand::expand_prompt_token_ids; +pub use self::input::MultimodalInput; /// Resolved multimodal support for one loaded model. #[derive(Clone)] @@ -654,11 +657,16 @@ impl MultimodalModelInfo { /// /// Modalities without a configured count are unlimited. fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> { + self.validate_modality_limits(media_parts.iter().filter_map(media_part_limit_modality)) + } + + fn validate_modality_limits( + &self, + modalities: impl IntoIterator, + ) -> Result<()> { let mut counts: HashMap = HashMap::new(); - for part in media_parts { - if let Some(modality) = media_part_limit_modality(part) { - *counts.entry(modality).or_default() += 1; - } + for modality in modalities { + *counts.entry(modality).or_default() += 1; } for (modality, count) in counts { @@ -677,6 +685,32 @@ impl MultimodalModelInfo { Ok(()) } + /// Validate inline storage, batching, and placeholder ranges, then check + /// this model's supported modalities and item-count limits. + /// + /// `prompt_len` must include all expanded multimodal placeholders. + pub(crate) fn prepare_preprocessed( + &self, + mut features: MmFeatures, + prompt_len: usize, + ) -> Result { + preprocessed::validate_features(&mut features, prompt_len)?; + for feature in &features { + let supported = match feature.modality { + MmModality::Image => self.image.is_some(), + MmModality::Video => self.video.is_some(), + MmModality::Audio => self.audio.is_some(), + }; + if !supported { + return Err(Error::UnsupportedModality { + modality: feature.modality.as_str().to_owned(), + }); + } + } + self.validate_modality_limits(features.iter().map(|feature| feature.modality))?; + Ok(features) + } + /// Run media fetch, per-modality preprocessing, prompt expansion, and /// feature build. /// diff --git a/rust/src/chat/src/multimodal/input.rs b/rust/src/chat/src/multimodal/input.rs new file mode 100644 index 000000000000..53905f0a2654 --- /dev/null +++ b/rust/src/chat/src/multimodal/input.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use llm_multimodal::MediaContentPart; +use vllm_engine_core_client::protocol::multimodal::MmFeatures; + +/// Media preparation stage supplied by a frontend handler. +#[derive(Debug)] +pub enum MultimodalInput { + /// Raw media to fetch and preprocess, expanding prompt placeholders as needed. + Raw(Vec), + /// Preprocessed features to validate against storage, batching, model, and prompt constraints. + Preprocessed(MmFeatures), +} + +impl MultimodalInput { + pub fn is_empty(&self) -> bool { + match self { + Self::Raw(parts) => parts.is_empty(), + Self::Preprocessed(features) => features.is_empty(), + } + } +} + +impl From> for MultimodalInput { + fn from(parts: Vec) -> Self { + Self::Raw(parts) + } +} diff --git a/rust/src/chat/src/multimodal/preprocessed.rs b/rust/src/chat/src/multimodal/preprocessed.rs new file mode 100644 index 000000000000..a2ff4687b134 --- /dev/null +++ b/rust/src/chat/src/multimodal/preprocessed.rs @@ -0,0 +1,274 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; + +use vllm_engine_core_client::protocol::dtype::TensorDtype; +use vllm_engine_core_client::protocol::multimodal::{ + MmFeatureSpec, MmField, MmKwargValue, MmKwargsItem, +}; + +use crate::error::{ + Error, Result, bail_invalid_preprocessed_multimodal, invalid_preprocessed_multimodal, +}; + +const MAX_DEPTH: usize = 32; + +/// Validate inline storage, placeholder ranges, and batching compatibility, +/// then stably sort features by their prompt offsets. +/// Encoder cache identifiers are supplied by the producer, which must bind +/// them to the feature content. +pub(super) fn validate_features(features: &mut [MmFeatureSpec], prompt_len: usize) -> Result<()> { + let mut fields = HashMap::new(); + for feature in features.iter() { + if feature.identifier.is_empty() { + bail_invalid_preprocessed_multimodal!("feature identifier must be nonempty"); + } + let data = feature.data.as_ref().ok_or_else(|| { + invalid_preprocessed_multimodal!( + "inline data is required; cache-only features are unsupported" + ) + })?; + validate_kwargs(data)?; + for (key, elem) in data { + if let Some(previous) = fields.insert((feature.modality, key), &elem.field) + && !same_batching(previous, &elem.field) + { + bail_invalid_preprocessed_multimodal!( + "incompatible batching for {} field {key:?}", + feature.modality.as_str() + ); + } + } + let position = &feature.mm_position; + if position.length == 0 || position.offset.checked_add(position.length).is_none() { + bail_invalid_preprocessed_multimodal!( + "placeholder range must have a positive length and fit in usize" + ); + } + if position.offset + position.length > prompt_len { + bail_invalid_preprocessed_multimodal!( + "{} placeholder range exceeds prompt token IDs", + feature.modality.as_str() + ); + } + if let Some(mask) = &position.is_embed { + if mask.dtype != TensorDtype::Bool || mask.shape != [position.length] { + bail_invalid_preprocessed_multimodal!( + "is_embed must be a boolean tensor of shape [placeholder length]" + ); + } + mask.validate_inline() + .map_err(|message| Error::InvalidPreprocessedMultimodal { message })?; + } + } + // Mirror the Python frontend (`argsort_mm_positions`): features are + // ordered by their placeholder position in the prompt. + features.sort_by_key(|feature| feature.mm_position.offset); + Ok(()) +} + +fn validate_kwargs(kwargs: &MmKwargsItem) -> Result<()> { + if kwargs.is_empty() { + bail_invalid_preprocessed_multimodal!("kwargs must contain at least one field"); + } + for (key, elem) in kwargs { + if key.is_empty() { + bail_invalid_preprocessed_multimodal!("kwargs field name must be nonempty"); + } + let data = elem.data.as_ref().ok_or_else(|| { + invalid_preprocessed_multimodal!("field {key:?} requires inline data") + })?; + validate_value(data, 0)?; + if let MmField::Flat(field) = &elem.field { + match data { + MmKwargValue::Tensor(tensor) => { + let rank = tensor.shape.len() as i32; + if rank == 0 || field.dim < -rank || field.dim >= rank { + bail_invalid_preprocessed_multimodal!( + "flat field {key:?} dim is outside the tensor rank" + ); + } + } + MmKwargValue::List(_) if field.dim == 0 => {} + _ => bail_invalid_preprocessed_multimodal!( + "flat field {key:?} requires a tensor or a list with dim=0" + ), + } + } + } + Ok(()) +} + +fn validate_value(value: &MmKwargValue, depth: usize) -> Result<()> { + if depth > MAX_DEPTH { + bail_invalid_preprocessed_multimodal!("kwargs nesting exceeds {MAX_DEPTH} levels"); + } + match value { + MmKwargValue::Tensor(tensor) => tensor + .validate_inline() + .map_err(|message| Error::InvalidPreprocessedMultimodal { message }), + MmKwargValue::List(values) => { + for value in values { + validate_value(value, depth + 1)?; + } + Ok(()) + } + MmKwargValue::Int(_) | MmKwargValue::Float(_) => Ok(()), + } +} + +fn same_batching(left: &MmField, right: &MmField) -> bool { + // Slices and batch_size describe the producer's original batch. Python + // reduce_data combines the already-split item data without using them. + match (left, right) { + (MmField::Batched(a), MmField::Batched(b)) => a == b, + (MmField::Flat(a), MmField::Flat(b)) => a.dim == b.dim && a.keep_on_cpu == b.keep_on_cpu, + (MmField::Shared(a), MmField::Shared(b)) => a.keep_on_cpu == b.keep_on_cpu, + _ => false, + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmFieldElem, MmFlatField, MmModality, MmSharedField, MmSlice, + PlaceholderRange, SliceSpec, + }; + use vllm_engine_core_client::protocol::tensor::{WireArrayData, WireTensor}; + + use super::*; + use crate::multimodal::tests::qwen3_vl_info; + + fn feature(offset: usize, field: MmField) -> MmFeatureSpec { + MmFeatureSpec { + data: Some(BTreeMap::from([( + "pixels".to_owned(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![2], vec![1., 2.]).unwrap(), + )), + field, + }, + )])), + modality: MmModality::Image, + identifier: format!("image-{offset}"), + mm_position: PlaceholderRange { + offset, + length: 2, + is_embed: None, + }, + mm_hash: Some("processor-hash".to_owned()), + } + } + + fn batched() -> MmField { + MmField::Batched(MmBatchedField { keep_on_cpu: false }) + } + + #[test] + fn inline_features_preserve_overlaps_and_stably_sort_prompt_positions() { + let info = qwen3_vl_info(); + let mut a = feature(2, batched()); + a.mm_position.is_embed = Some(WireTensor::from_bool(vec![2], vec![true, false]).unwrap()); + let b = feature(1, batched()); + let mut c = feature(2, batched()); + c.modality = MmModality::Video; + let expected = vec![b.clone(), a.clone(), c.clone()]; + let features = info.prepare_preprocessed(vec![a, b, c], 4).unwrap(); + assert_eq!(features, expected); + } + + #[test] + fn independently_preprocessed_items_keep_their_original_batch_metadata() { + let info = qwen3_vl_info(); + let flat = |length| { + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(length), + step: None, + })], + dim: 0, + keep_on_cpu: false, + }) + }; + let shared = |batch_size| { + MmField::Shared(MmSharedField { + batch_size, + keep_on_cpu: true, + }) + }; + for (a, b) in [(flat(2), flat(7)), (shared(1), shared(3))] { + let features = vec![feature(0, a), feature(2, b)]; + let result = info.prepare_preprocessed(features.clone(), 4).unwrap(); + assert_eq!(result, features); + } + } + + #[test] + fn inline_features_reject_cache_references_and_invalid_tensor_storage() { + let info = qwen3_vl_info(); + let original = feature(0, batched()); + let mut cached = original.clone(); + cached.data = None; + assert!(info.prepare_preprocessed(vec![cached], 4).is_err()); + let mut field_cached = original.clone(); + field_cached.data.as_mut().unwrap().get_mut("pixels").unwrap().data = None; + assert!(info.prepare_preprocessed(vec![field_cached], 4).is_err()); + for tensor in [ + WireTensor { + dtype: TensorDtype::F32, + shape: vec![2], + data: WireArrayData::AuxIndex(1), + }, + WireTensor::from_raw(TensorDtype::F32, vec![2], vec![0; 4]), + WireTensor::from_raw(TensorDtype::F32, vec![usize::MAX, 2], vec![]), + ] { + let mut item = original.clone(); + item.data.as_mut().unwrap().get_mut("pixels").unwrap().data = + Some(MmKwargValue::Tensor(tensor)); + assert!(info.prepare_preprocessed(vec![item], 4).is_err()); + } + } + + #[test] + fn inline_features_validate_placeholder_and_mask_bounds() { + let info = qwen3_vl_info(); + for (offset, length) in [(usize::MAX, 2), (0, 0)] { + let mut item = feature(offset, batched()); + item.mm_position.length = length; + assert!(info.prepare_preprocessed(vec![item], 4).is_err()); + } + assert!(info.prepare_preprocessed(vec![feature(2, batched())], 3).is_err()); + for mask in [ + WireTensor::from_bool(vec![1], vec![true]).unwrap(), + WireTensor::from_u32(vec![2], vec![1, 0]).unwrap(), + ] { + let mut item = feature(0, batched()); + item.mm_position.is_embed = Some(mask); + assert!(info.prepare_preprocessed(vec![item], 4).is_err()); + } + } + + #[test] + fn inline_features_reject_incompatible_batching_and_flat_axes() { + let info = qwen3_vl_info(); + let flat = MmField::Flat(MmFlatField { + slices: vec![], + dim: 2, + keep_on_cpu: false, + }); + assert!(info.prepare_preprocessed(vec![feature(0, flat)], 4).is_err()); + let shared = MmField::Shared(MmSharedField { + batch_size: 1, + keep_on_cpu: false, + }); + assert!( + info.prepare_preprocessed(vec![feature(0, batched()), feature(2, shared)], 4) + .is_err() + ); + } +} diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index e4d12bc96b39..b9abe19c63b4 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -20,6 +20,7 @@ libc.workspace = true llm-multimodal.workspace = true openssl.workspace = true prost-types.workspace = true +rmp-serde.workspace = true rmpv.workspace = true serde.workspace = true serde_json.workspace = true @@ -58,7 +59,6 @@ async-openai = { workspace = true, features = ["full"] } bytes.workspace = true clap.workspace = true expect-test.workspace = true -rmp-serde.workspace = true serial_test.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["test-util"] } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index ca44f59c93ea..cf0045f71470 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -5,9 +5,7 @@ //! request/response types. use tonic::Status; -use url::Url; use uuid::Uuid; -use vllm_chat::MediaContentPart; use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_text::{ @@ -17,95 +15,6 @@ use vllm_text::{ use super::pb; -pub fn media_parts_from_request( - media: Vec, -) -> Result, Status> { - let mut parts = Vec::with_capacity(media.len()); - for (index, item) in media.into_iter().enumerate() { - let modality = item.modality(); - if modality == pb::Modality::Unspecified { - return Err(Status::invalid_argument(format!( - "media[{index}].modality is required" - ))); - } - let uuid = (!item.uuid.is_empty()).then_some(item.uuid); - let mime_type = (!item.mime_type.is_empty()).then_some(item.mime_type); - let source = item.source.ok_or_else(|| { - Status::invalid_argument(format!("media[{index}].source is required")) - })?; - match &source { - pb::media_item::Source::Url(url) => { - validate_media_uri(index, "url", url, &["http", "https"])?; - } - pb::media_item::Source::DataUri(uri) => { - validate_media_uri(index, "data_uri", uri, &["data"])?; - } - pb::media_item::Source::RawBytes(_) => {} - } - let part = match (modality, source) { - ( - pb::Modality::Image, - pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), - ) => MediaContentPart::ImageUrl { - url, - detail: None, - uuid, - }, - (pb::Modality::Image, pb::media_item::Source::RawBytes(data)) => { - MediaContentPart::ImageData { - data, - mime_type, - uuid, - detail: None, - } - } - ( - pb::Modality::Video, - pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), - ) => MediaContentPart::VideoUrl { url, uuid }, - (pb::Modality::Video, pb::media_item::Source::RawBytes(data)) => { - MediaContentPart::VideoData { - data, - mime_type, - uuid, - } - } - ( - pb::Modality::Audio, - pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), - ) => MediaContentPart::AudioUrl { url, uuid }, - (pb::Modality::Audio, pb::media_item::Source::RawBytes(data)) => { - MediaContentPart::AudioData { - data, - mime_type, - uuid, - } - } - (pb::Modality::Unspecified, _) => unreachable!("modality validated above"), - }; - parts.push(part); - } - Ok(parts) -} - -fn validate_media_uri( - index: usize, - field: &str, - value: &str, - allowed_schemes: &[&str], -) -> Result<(), Status> { - let uri = Url::parse(value).map_err(|_| { - Status::invalid_argument(format!("media[{index}].{field} is not a valid URI")) - })?; - if !allowed_schemes.contains(&uri.scheme()) { - return Err(Status::invalid_argument(format!( - "media[{index}].{field} must use the {} scheme", - allowed_schemes.join(" or ") - ))); - } - Ok(()) -} - // ======================================================================================== // Request conversion // ======================================================================================== diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 7051558cd882..72943b6af0bb 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -149,19 +149,23 @@ impl InferenceServiceImpl { })?); } - let media = convert::media_parts_from_request(media)?; + let media = super::media::from_proto(media)?; if !media.is_empty() { let Prompt::TokenIds(mut token_ids) = text_request.prompt else { return Err(Status::invalid_argument( "multimodal gRPC requests must provide token_ids input", )); }; - let mm_features = self - .state - .chat - .prepare_media(media, &mut token_ids) - .await - .map_err(|error| Status::internal(error.to_report_string()))?; + let mm_features = + self.state.chat.prepare_media(media, &mut token_ids).await.map_err( + |error| { + if error.is_request_validation_error() { + Status::invalid_argument(error.to_report_string()) + } else { + Status::internal(error.to_report_string()) + } + }, + )?; text_request.prompt = Prompt::TokenIds(token_ids); text_request.mm_features = mm_features; } diff --git a/rust/src/server/src/grpc/media.rs b/rust/src/server/src/grpc/media.rs new file mode 100644 index 000000000000..ff746b191e04 --- /dev/null +++ b/rust/src/server/src/grpc/media.rs @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Protobuf media adaptation into the shared multimodal preparation boundary. + +use std::io::Cursor; + +use serde::Deserialize as _; +use thiserror_ext::AsReport as _; +use tonic::Status; +use url::Url; +use vllm_chat::MediaContentPart; +use vllm_chat::multimodal::MultimodalInput; +use vllm_engine_core_client::protocol::multimodal::{ + MmFeatureSpec, MmKwargsItem, MmModality, PlaceholderRange, +}; +use vllm_engine_core_client::protocol::tensor::WireTensor; + +use super::pb; + +/// Maximum total serialized kwargs size accepted in one gRPC request. +const MAX_INLINE_MM_BYTES: usize = 16 * 1024 * 1024; +const MAX_MSGPACK_DEPTH: usize = 32; + +/// Adapt a request containing either raw media or preprocessed features. +/// +/// Require a single input kind and bound the total kwargs size before decoding. +/// Shared media preparation validates features, model support, item limits, and prompt bounds. +pub(super) fn from_proto(media: Vec) -> Result { + if !media + .iter() + .any(|item| matches!(item.source, Some(pb::media_item::Source::Features(_)))) + { + return raw_parts_from_proto(media).map(MultimodalInput::Raw); + } + let mut encoded_bytes = 0usize; + for item in &media { + let Some(pb::media_item::Source::Features(feature)) = &item.source else { + return Err(mixed_media_error()); + }; + encoded_bytes = encoded_bytes.saturating_add(feature.kwargs.as_ref().map_or(0, Vec::len)); + } + if encoded_bytes > MAX_INLINE_MM_BYTES { + return Err(Status::resource_exhausted(format!( + "inline multimodal payload exceeds {MAX_INLINE_MM_BYTES} bytes" + ))); + } + let features = media + .into_iter() + .enumerate() + .map(|(index, item)| { + feature_from_proto(item).map_err(|error| { + Status::new(error.code(), format!("media[{index}]: {}", error.message())) + }) + }) + .collect::, _>>()?; + Ok(MultimodalInput::Preprocessed(features)) +} + +fn feature_from_proto(item: pb::MediaItem) -> Result { + let modality = match item.modality() { + pb::Modality::Image => MmModality::Image, + pb::Modality::Video => MmModality::Video, + pb::Modality::Audio => MmModality::Audio, + pb::Modality::Unspecified => return Err(Status::invalid_argument("modality is required")), + }; + let Some(pb::media_item::Source::Features(feature)) = item.source else { + return Err(mixed_media_error()); + }; + let bytes = feature.kwargs.ok_or_else(|| { + Status::invalid_argument("features.kwargs is required; cache-only features are unsupported") + })?; + let offset = usize::try_from(feature.offset) + .map_err(|_| Status::invalid_argument("features.offset exceeds platform limits"))?; + let length = usize::try_from(feature.length) + .map_err(|_| Status::invalid_argument("features.length exceeds platform limits"))?; + let is_embed = if feature.is_embed.is_empty() { + None + } else { + Some( + WireTensor::from_bool(vec![length], feature.is_embed) + .map_err(Status::invalid_argument)?, + ) + }; + Ok(MmFeatureSpec { + modality, + data: Some(decode_inline_mm_kwargs(&bytes)?), + identifier: feature.identifier, + mm_hash: feature.mm_hash, + mm_position: PlaceholderRange { + offset, + length, + is_embed, + }, + }) +} + +/// Decode one inline item using the engine protocol's typed serde schema. +fn decode_inline_mm_kwargs(bytes: &[u8]) -> Result { + let mut decoder = rmp_serde::Deserializer::new(Cursor::new(bytes)); + decoder.set_max_depth(MAX_MSGPACK_DEPTH); + let kwargs = MmKwargsItem::deserialize(&mut decoder).map_err(|error| { + Status::invalid_argument(format!( + "invalid multimodal kwargs MessagePack: {}", + error.as_report() + )) + })?; + if decoder.position() != bytes.len() as u64 { + return Err(Status::invalid_argument( + "kwargs contains trailing MessagePack data", + )); + } + Ok(kwargs) +} + +fn mixed_media_error() -> Status { + Status::invalid_argument("raw media and preprocessed media features cannot be mixed") +} + +fn raw_parts_from_proto(media: Vec) -> Result, Status> { + let mut parts = Vec::with_capacity(media.len()); + for (index, item) in media.into_iter().enumerate() { + let modality = item.modality(); + if modality == pb::Modality::Unspecified { + return Err(Status::invalid_argument(format!( + "media[{index}].modality is required" + ))); + } + let uuid = (!item.uuid.is_empty()).then_some(item.uuid); + let mime_type = (!item.mime_type.is_empty()).then_some(item.mime_type); + let source = item.source.ok_or_else(|| { + Status::invalid_argument(format!("media[{index}].source is required")) + })?; + match &source { + pb::media_item::Source::Url(url) => { + validate_media_uri(index, "url", url, &["http", "https"])?; + } + pb::media_item::Source::DataUri(uri) => { + validate_media_uri(index, "data_uri", uri, &["data"])?; + } + pb::media_item::Source::RawBytes(_) => {} + pb::media_item::Source::Features(_) => return Err(mixed_media_error()), + } + let part = match (modality, source) { + ( + pb::Modality::Image, + pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), + ) => MediaContentPart::ImageUrl { + url, + detail: None, + uuid, + }, + (pb::Modality::Image, pb::media_item::Source::RawBytes(data)) => { + MediaContentPart::ImageData { + data, + mime_type, + uuid, + detail: None, + } + } + ( + pb::Modality::Video, + pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), + ) => MediaContentPart::VideoUrl { url, uuid }, + (pb::Modality::Video, pb::media_item::Source::RawBytes(data)) => { + MediaContentPart::VideoData { + data, + mime_type, + uuid, + } + } + ( + pb::Modality::Audio, + pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url), + ) => MediaContentPart::AudioUrl { url, uuid }, + (pb::Modality::Audio, pb::media_item::Source::RawBytes(data)) => { + MediaContentPart::AudioData { + data, + mime_type, + uuid, + } + } + (_, pb::media_item::Source::Features(_)) => return Err(mixed_media_error()), + (pb::Modality::Unspecified, _) => unreachable!("modality validated above"), + }; + parts.push(part); + } + Ok(parts) +} + +fn validate_media_uri( + index: usize, + field: &str, + value: &str, + allowed_schemes: &[&str], +) -> Result<(), Status> { + let uri = Url::parse(value).map_err(|_| { + Status::invalid_argument(format!("media[{index}].{field} is not a valid URI")) + })?; + if !allowed_schemes.contains(&uri.scheme()) { + return Err(Status::invalid_argument(format!( + "media[{index}].{field} must use the {} scheme", + allowed_schemes.join(" or ") + ))); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use super::*; + use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, SliceSpec, + }; + + fn preprocessed() -> pb::MediaItem { + let kwargs = BTreeMap::from([( + "x", + MmFieldElem { + data: Some(MmKwargValue::Int(1)), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + )]); + pb::MediaItem { + modality: pb::Modality::Image as i32, + source: Some(pb::media_item::Source::Features( + pb::PreprocessedMediaFeatures { + kwargs: Some(rmp_serde::to_vec_named(&kwargs).unwrap()), + identifier: "image-id".to_owned(), + length: 1, + ..Default::default() + }, + )), + ..Default::default() + } + } + + #[test] + fn mixed_media_is_rejected_in_both_orders() { + let raw = pb::MediaItem { + modality: pb::Modality::Image as i32, + source: Some(pb::media_item::Source::RawBytes(vec![1])), + ..Default::default() + }; + for media in [vec![raw.clone(), preprocessed()], vec![preprocessed(), raw]] { + assert_eq!( + from_proto(media).unwrap_err().code(), + tonic::Code::InvalidArgument + ); + } + } + + #[test] + fn protobuf_features_require_modality_inline_kwargs_and_a_matching_mask() { + let mut missing_modality = preprocessed(); + missing_modality.modality = 99; + let mut missing_kwargs = preprocessed(); + let Some(pb::media_item::Source::Features(feature)) = &mut missing_kwargs.source else { + unreachable!() + }; + feature.kwargs = None; + let mut invalid_mask = preprocessed(); + let Some(pb::media_item::Source::Features(feature)) = &mut invalid_mask.source else { + unreachable!() + }; + feature.is_embed = vec![true, false]; + for media in [missing_modality, missing_kwargs, invalid_mask] { + let error = from_proto(vec![media]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + assert!(error.message().starts_with("media[0]:")); + } + } + + #[test] + fn aggregate_payload_limit_is_checked_before_decoding_items() { + let mut item = preprocessed(); + let Some(pb::media_item::Source::Features(feature)) = &mut item.source else { + unreachable!() + }; + feature.kwargs = Some(vec![0; MAX_INLINE_MM_BYTES / 2 + 1]); + let error = from_proto(vec![item.clone(), item]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + + #[test] + fn kwargs_decoder_rejects_trailing_truncated_and_deeply_nested_input() { + let kwargs = BTreeMap::from([( + "pixels".to_owned(), + MmFieldElem { + data: Some(MmKwargValue::Int(1)), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + )]); + let valid = rmp_serde::to_vec_named(&kwargs).unwrap(); + assert_eq!(decode_inline_mm_kwargs(&valid).unwrap(), kwargs); + let mut trailing = valid.clone(); + trailing.push(0); + let mut nested_kwargs = kwargs.clone(); + let mut value = MmKwargValue::Int(1); + for _ in 0..64 { + value = MmKwargValue::List(vec![value]); + } + nested_kwargs.get_mut("pixels").unwrap().data = Some(value); + let nested = rmp_serde::to_vec_named(&nested_kwargs).unwrap(); + // A declared array length must not drive a proportional allocation. + let huge_array = vec![ + 0x81, 0xa1, b'x', 0x81, 0xa4, b'd', b'a', b't', b'a', 0xdd, 0xff, 0xff, 0xff, 0xff, + ]; + for bytes in [&valid[..valid.len() - 1], &trailing, &nested, &huge_array] { + assert!(decode_inline_mm_kwargs(bytes).is_err()); + } + } + + #[test] + fn kwargs_decode_python_encoder_tensor_and_factory_tuple() { + // MsgpackEncoder(size_threshold=1 << 30), vLLM 6cbb3c154e: + // pixel_values = tensor([1., 2.]), MultiModalFlatField(slices=[slice(0, 2)]). + let bytes = b"\x81\xacpixel_values\x82\xa4data\x93\xa7float32\x91\x02\xd7\x03\ + \x00\x00\x80\x3f\x00\x00\x00\x40\ + \xa5field\x92\xa4flat\x83\xabkeep_on_cpu\xc2\ + \xa6slices\x91\x93\x00\x02\xc0\xa3dim\x00"; + let kwargs = decode_inline_mm_kwargs(bytes).unwrap(); + let expected = MmFieldElem { + data: Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![2], vec![1., 2.]).unwrap(), + )), + field: MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(2), + step: None, + })], + dim: 0, + keep_on_cpu: false, + }), + }; + assert_eq!(kwargs["pixel_values"], expected); + } +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 3c96afd9b811..c61c1e5e1419 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -7,6 +7,7 @@ mod control; mod convert; mod health; mod inference; +mod media; /// Generated protobuf/gRPC types for the `vllm` package. pub use vllm_proto as pb; diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 59403277ee94..f13dc11b277d 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::BTreeMap; use std::fs; use std::future::Future; use std::io; @@ -23,6 +24,7 @@ use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatu use tonic_health::pb::health_client::HealthClient; use tonic_health::server::health_reporter; use tower::service_fn; +use vllm_chat::multimodal::{MmLimitPerPrompt, MmLimitSpec, MmModality}; use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt, @@ -32,12 +34,17 @@ use vllm_engine_core_client::mock_engine::{ default_ready_response, }; use vllm_engine_core_client::protocol::decode_value; +use vllm_engine_core_client::protocol::dtype::TensorDtype; use vllm_engine_core_client::protocol::handshake::{EngineCoreReadyResponse, KvEventsConfig}; +use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmKwargValue, +}; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, UtilityCallOutput, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::tensor::WireTensor; use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task_with_ready}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode}; @@ -292,6 +299,12 @@ const QWEN_IMAGE_TOKEN_ID: u32 = 151655; const TINY_PNG_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="; fn multimodal_backend() -> Arc { + multimodal_backend_with_limits(Default::default()) +} + +fn multimodal_backend_with_limits( + limit_mm_per_prompt: MmLimitPerPrompt, +) -> Arc { let config_path = std::env::temp_dir().join(format!( "vllm-grpc-qwen-config-{}.json", uuid::Uuid::new_v4() @@ -309,7 +322,7 @@ fn multimodal_backend() -> Arc { ..Default::default() }, Arc::new(TestTokenizer::new().with_regular_token("<|image_pad|>", QWEN_IMAGE_TOKEN_ID)), - Default::default(), + limit_mm_per_prompt, ) .expect("load multimodal info") .expect("qwen multimodal info is registered"); @@ -877,6 +890,249 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn generate_forwards_preprocessed_multimodal_features() { + for stream in [false, true] { + let kwargs = BTreeMap::from([( + "pixel_values".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(WireTensor::from_raw( + TensorDtype::U8, + vec![3], + vec![1, 2, 3], + ))), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + )]); + let encoded_kwargs = rmp_serde::to_vec_named(&kwargs).expect("encode multimodal kwargs"); + let expected_kwargs = kwargs.clone(); + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_backend( + b"engine-grpc-preprocessed-multimodal", + default_stream_output_specs(), + multimodal_backend(), + move |request| { + assert_eq!( + request.prompt_token_ids.as_deref(), + Some(&[11, 99, 99, 12][..]) + ); + let features = request.mm_features.as_ref().expect("multimodal features"); + assert_eq!(features.len(), 1); + let feature = &features[0]; + assert_eq!(feature.modality.as_str(), "image"); + assert_eq!(feature.identifier, "image-hash-a"); + assert_eq!(feature.mm_hash.as_deref(), Some("image-hash-a")); + assert_eq!(feature.mm_position.offset, 1); + assert_eq!(feature.mm_position.length, 2); + assert_eq!(feature.data.as_ref(), Some(&expected_kwargs)); + assert_eq!( + feature.mm_position.is_embed, + Some(WireTensor::from_bool(vec![2], vec![true, false]).unwrap()) + ); + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + + let request = pb::GenerateRequest { + request_id: "test-preprocessed-multimodal".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![11, 99, 99, 12], + })), + media: vec![pb::MediaItem { + modality: pb::Modality::Image as i32, + source: Some(pb::media_item::Source::Features( + pb::PreprocessedMediaFeatures { + kwargs: Some(encoded_kwargs), + identifier: "image-hash-a".to_string(), + offset: 1, + length: 2, + mm_hash: Some("image-hash-a".to_string()), + is_embed: vec![true, false], + }, + )), + mime_type: String::new(), + uuid: String::new(), + }], + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + }; + if stream { + let mut response = + client.generate_stream(request).await.expect("preprocessed stream").into_inner(); + while response.message().await.expect("stream message").is_some() {} + } else { + client.generate(request).await.expect("preprocessed generate"); + } + + engine_task.await.expect("mock engine task"); + server_task.abort(); + } +} + +fn preprocessed_grpc_media( + modality: pb::Modality, + identifier: &str, + offset: u64, + length: u64, +) -> pb::MediaItem { + let kwargs = BTreeMap::from([( + "pixel_values".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(WireTensor::from_raw( + TensorDtype::U8, + vec![1], + vec![1], + ))), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + )]); + pb::MediaItem { + modality: modality as i32, + source: Some(pb::media_item::Source::Features( + pb::PreprocessedMediaFeatures { + kwargs: Some(rmp_serde::to_vec_named(&kwargs).expect("encode kwargs")), + identifier: identifier.to_string(), + offset, + length, + mm_hash: Some(identifier.to_string()), + is_embed: Vec::new(), + }, + )), + mime_type: String::new(), + uuid: String::new(), + } +} + +fn preprocessed_grpc_request(media: Vec) -> pb::GenerateRequest { + pb::GenerateRequest { + request_id: "test-preprocessed-validation".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![11, 99, 99, 12], + })), + media, + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_rejects_preprocessed_features_for_text_only_model() { + let (inference_service, control_service, engine_health, _engine_task) = + setup_grpc_service_with_backend( + b"engine-grpc-preprocessed-text-only", + default_stream_output_specs(), + Arc::new(FakeTextBackend), + |_| {}, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + + let status = client + .generate(preprocessed_grpc_request(vec![preprocessed_grpc_media( + pb::Modality::Image, + "image-a", + 1, + 2, + )])) + .await + .expect_err("text-only models must reject preprocessed features"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn generate_validates_preprocessed_model_limits_and_prompt() { + let limits = MmLimitPerPrompt::from([(MmModality::Image, MmLimitSpec::Count(1))]); + let (inference_service, control_service, engine_health, _engine_task) = + setup_grpc_service_with_backend( + b"engine-grpc-preprocessed-limit", + default_stream_output_specs(), + multimodal_backend_with_limits(limits), + |_| {}, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + + let over_limit = preprocessed_grpc_request(vec![ + preprocessed_grpc_media(pb::Modality::Image, "image-a", 0, 1), + preprocessed_grpc_media(pb::Modality::Image, "image-b", 2, 1), + ]); + let bad_range = preprocessed_grpc_request(vec![preprocessed_grpc_media( + pb::Modality::Image, + "image", + 3, + 2, + )]); + let unsupported = preprocessed_grpc_request(vec![preprocessed_grpc_media( + pb::Modality::Audio, + "audio", + 0, + 1, + )]); + let mut text = preprocessed_grpc_request(vec![preprocessed_grpc_media( + pb::Modality::Image, + "image", + 0, + 1, + )]); + let mut truncated = text.clone(); + text.prompt = Some(pb::generate_request::Prompt::Text("hello".to_owned())); + truncated.truncate_prompt_tokens = 1; + let empty_identifier = + preprocessed_grpc_request(vec![preprocessed_grpc_media(pb::Modality::Image, "", 0, 1)]); + for request in [ + over_limit, + bad_range, + unsupported, + text, + truncated, + empty_identifier, + ] { + let status = client.generate(request.clone()).await.expect_err("invalid multimodal input"); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + let status = client + .generate_stream(request) + .await + .expect_err("invalid streaming multimodal input"); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + } + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_rejects_text_prompt_with_media() { From ec6b0494fef6bf665edc60ffbe315fcbab27bbd6 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Sat, 12 Sep 2026 11:11:29 +0800 Subject: [PATCH 078/774] [CI] Bump CUTLASS DSL to 4.7 (#54927) Signed-off-by: zjy0516 Co-authored-by: Codex --- cmake/external_projects/tml_fa4.cmake | 2 +- requirements/cuda.txt | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cmake/external_projects/tml_fa4.cmake b/cmake/external_projects/tml_fa4.cmake index a965295e2bbd..c6ddff82c339 100644 --- a/cmake/external_projects/tml_fa4.cmake +++ b/cmake/external_projects/tml_fa4.cmake @@ -14,7 +14,7 @@ else() FetchContent_Declare( tml_fa4 GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git - GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d + GIT_TAG 75765e76a9c2c012c1f6ecd64577eb646eb4d303 GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "") diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 427f86ddbeda..00a56ac3bd6e 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -26,8 +26,8 @@ fastsafetensors >= 0.3.3 instanttensor >= 0.1.9 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.6.2 -quack-kernels==0.6.4 # Required for CUTLASS DSL 4.6 by MSA +nvidia-cutlass-dsl[cu13]==4.7.1 +quack-kernels==0.6.5 # Required by MSA # Tokenspeed_MLA for faster mla with spec decode tokenspeed-mla==0.1.8; platform_system == "Linux" From 5f3e4b74474d4f2ad3e17334364a86d298ae2382 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 11 Sep 2026 20:15:02 -0700 Subject: [PATCH 079/774] Revert "[Agents] Link Triton skill to JIT kernel warmup guide" (#56546) --- .agents/skills/triton-kernel-writing/SKILL.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.agents/skills/triton-kernel-writing/SKILL.md b/.agents/skills/triton-kernel-writing/SKILL.md index 5a524c8a718a..6dd6a8520be0 100644 --- a/.agents/skills/triton-kernel-writing/SKILL.md +++ b/.agents/skills/triton-kernel-writing/SKILL.md @@ -26,10 +26,6 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo unimportant runtime integer scalars in `do_not_specialize`, especially those that may alternate between values such as 0 and 1, which can produce different specialization keys. -- For new warmable kernels, follow the - [JIT Kernel Warmup guide](../../../docs/contributing/jit_kernel_warmup.md) - to define the warmup wrapper and call its `register_warmup(...)` where the - runtime implementation is selected during runner setup. - The Triton compiler does not guarantee safe ordering when a kernel writes to a pointer and subsequently reads from the same pointer. This pattern must have a `tl.debug_barrier()` between the write and read. The barrier From 485421b1c3572597a4cfaece04836a843537435a Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 12 Sep 2026 11:39:17 +0800 Subject: [PATCH 080/774] [Rust Frontend] Honor HF revisions, offline mode, and cache directory (#56386) Signed-off-by: Bugen Zhao Co-authored-by: Codex --- rust/src/chat/src/backend/hf.rs | 3 +- rust/src/chat/src/backend/mod.rs | 2 + rust/src/cmd/src/cli.rs | 12 ++ rust/src/cmd/src/cli/tests.rs | 27 ++- rust/src/managed-engine/src/cli.rs | 5 + .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 2 + rust/src/server/src/lib.rs | 1 + rust/src/server/src/render.rs | 3 + rust/src/text/src/backend/hf/model_files.rs | 170 +++++++++++++++--- rust/src/text/src/lower.rs | 3 +- 11 files changed, 201 insertions(+), 28 deletions(-) diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 717bde289836..37212a96505c 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,7 +129,7 @@ pub(super) async fn load_model_backends( model_id: &str, options: LoadModelBackendsOptions, ) -> Result { - let files = ResolvedModelFiles::new(model_id).await?; + let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; let text_backend = HfTextBackend::from_resolved_model_files( files.clone(), model_id.to_string(), @@ -231,6 +231,7 @@ mod tests { resolved_files(config_json, tokenizer_config_json), "test-model".to_string(), LoadModelBackendsOptions { + revision: None, generation_config: Default::default(), renderer, language_model_only: false, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 77b9114a5852..028c0c40c8ec 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -61,6 +61,8 @@ pub type DynChatTextBackend = Arc; /// Frontend-side chat backend loading options. #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct LoadModelBackendsOptions { + /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). + pub revision: Option, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index c0b4c6d29550..f7fcf30022b9 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -110,6 +110,9 @@ pub enum BenchCommand { pub struct RenderArgs { /// Model identifier or local model directory containing tokenizer files. model: String, + /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). + #[arg(long)] + revision: Option, /// HTTP bind host. #[arg(long, default_value = "127.0.0.1")] host: String, @@ -157,6 +160,7 @@ impl RenderArgs { RenderConfig { model: self.model, + revision: self.revision, served_model_name: self.served_model_name, host: self.host, port: self.port, @@ -195,6 +199,11 @@ pub struct SharedRuntimeArgs { /// public model ID. pub model: String, + /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). + #[arg(long)] + #[serde(default)] + pub revision: Option, + /// The source of generation-config sampling defaults. `"auto"` loads the /// model's defaults, while `"vllm"` uses vLLM's neutral defaults. #[arg(long, default_value_t)] @@ -481,6 +490,7 @@ impl SharedRuntimeArgs { None => CoordinatorMode::None, }, model: self.model, + revision: self.revision, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, @@ -536,6 +546,7 @@ impl SharedRuntimeArgs { }, coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, + revision: self.revision, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode, @@ -743,6 +754,7 @@ impl ServeArgs { self.managed_engine.clone().into_config( self.runtime.model.clone(), + self.runtime.revision.clone(), self.runtime.max_logprobs, profiler_config, reasoning_parser.as_deref(), diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index e9411891e010..884da135d719 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -36,6 +36,8 @@ fn render_args_build_config_without_tls() { "vllm-rs", "render", "Qwen/Qwen2.5-0.5B-Instruct", + "--revision", + "release", "--port", "8080", "--max-model-len", @@ -55,6 +57,7 @@ fn render_args_build_config_without_tls() { let config = args.into_config(); assert_eq!(config.model, "Qwen/Qwen2.5-0.5B-Instruct"); + assert_eq!(config.revision.as_deref(), Some("release")); assert_eq!(config.host, "127.0.0.1"); assert_eq!(config.port, 8080); assert_eq!(config.max_model_len, Some(32768)); @@ -173,6 +176,7 @@ fn serve_args_forward_python_flags_with_separator() { uds: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1007,6 +1011,7 @@ fn frontend_args_accept_json() { data_parallel_size: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: None, @@ -1125,7 +1130,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "--output-address", "ipc:///tmp/output.sock", "--args-json", - r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, + r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) .unwrap(); @@ -1133,6 +1138,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { panic!("expected frontend args"); }; assert_eq!(args.runtime.engine_ready_timeout_secs, 42); + assert_eq!(args.runtime.revision.as_deref(), Some("release")); assert_eq!(args.runtime.generation_config, GenerationConfigMode::Vllm); assert_eq!( args.runtime.tool_call_parser, @@ -1604,6 +1610,7 @@ fn serve_args_accept_handshake_aliases() { uds: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1752,6 +1759,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { }, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1839,6 +1847,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { }, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1948,6 +1957,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present address: "tcp://127.0.0.1:7000", }, model: "Qwen/Qwen3-0.6B", + revision: None, generation_config: Auto, served_model_name: [], listener_mode: InheritedFd { @@ -2087,3 +2097,18 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() { let config = args.into_config(); assert_eq!(config.profiler, None); } + +#[test] +fn serve_revision_reaches_frontend_and_managed_engine() { + let cli = + Cli::try_parse_from(["vllm-rs", "serve", "test/model", "--revision", "release"]).unwrap(); + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + assert_eq!( + args.to_frontend_config("tcp://localhost:1234".into()).revision.as_deref(), + Some("release") + ); + let engine = args.to_managed_engine_config(1234); + assert!(engine.python_args.windows(2).any(|args| args == ["--revision", "release"])); +} diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index eadefced236e..f0cd7bf1715d 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -84,6 +84,7 @@ impl ManagedEngineArgs { pub fn into_config( self, model: String, + revision: Option, max_logprobs: Option, profiler_config: Option, reasoning_parser: Option<&str>, @@ -95,6 +96,10 @@ impl ManagedEngineArgs { ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. + if let Some(revision) = revision { + python_args.push("--revision".to_string()); + python_args.push(revision); + } if let Some(max_model_len) = self.max_model_len { python_args.push("--max-model-len".to_string()); python_args.push(max_model_len); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index d07ba242aed5..fc9914e3e2f2 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -60,6 +60,7 @@ async fn main() -> Result<()> { }, coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, + revision: None, generation_config: Default::default(), served_model_name: vec![], listener_mode: HttpListenerMode::BindTcp { diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index fde53183f5b3..4928cddf31e9 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -210,6 +210,8 @@ pub struct Config { pub coordinator_mode: CoordinatorMode, /// Backend model identifier used for engine-core loading. pub model: String, + /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). + pub revision: Option, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 098c80380557..95ae359f35b3 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -98,6 +98,7 @@ async fn build_state(config: &Config) -> Result> { let loaded = load_model_backends( &config.model, LoadModelBackendsOptions { + revision: config.revision.clone(), generation_config: config.generation_config, renderer: config.renderer, language_model_only: config.language_model_only, diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index 6942694afbf3..fe7033fa8cad 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -24,6 +24,7 @@ use crate::{ #[derive(Debug)] pub struct RenderConfig { pub model: String, + pub revision: Option, pub served_model_name: Vec, pub host: String, pub port: u16, @@ -67,6 +68,7 @@ async fn build_state(config: &RenderConfig) -> Result> { let loaded = load_model_backends( &config.model, LoadModelBackendsOptions { + revision: config.revision.clone(), generation_config: Default::default(), renderer: config.renderer, language_model_only: true, @@ -150,6 +152,7 @@ mod tests { let error = serve_render( RenderConfig { model: "test-model".to_string(), + revision: None, served_model_name: Vec::new(), host: "127.0.0.1".to_string(), port: 8000, diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index d0b0caecce4c..de802fb4390c 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -3,8 +3,8 @@ use std::path::{Path, PathBuf}; -use hf_hub::Cache; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; +use hf_hub::{Cache, Repo, RepoType}; use thiserror_ext::AsReport as _; use super::config::{HfTokenizerConfig, load_tokenizer_config}; @@ -56,14 +56,31 @@ impl ResolvedModelFiles { /// Resolve tokenizer/config files from a local model directory first when /// `model_id` points to one, otherwise consult the local HF cache and /// finally the Hub. - pub async fn new(model_id: &str) -> Result { + pub async fn new(model_id: &str, revision: Option<&str>) -> Result { if Path::new(model_id).is_dir() { return resolve_local_model_files(Path::new(model_id)); } - if let Some(files) = resolve_cached_model_files(model_id)? { + let repo = Repo::with_revision( + model_id.to_string(), + RepoType::Model, + revision.unwrap_or("main").to_string(), + ); + // TODO: Let hf-hub read HF_HUB_CACHE after upgrading to 1.0. + let cache = std::env::var_os("HF_HUB_CACHE") + .map(|path| Cache::new(path.into())) + .unwrap_or_else(Cache::from_env); + if let Some(files) = resolve_cached_model_files(&cache, &repo)? { return Ok(files); } - resolve_remote_model_files(model_id).await + if is_hf_hub_offline() { + return Err(Error::Tokenizer(format!( + "model '{model_id}' at revision '{}' has no supported tokenizer in cache '{}' \ + and HF_HUB_OFFLINE is enabled", + repo.revision(), + cache.path().display(), + ))); + } + resolve_remote_model_files(model_id, repo, &cache).await } } @@ -87,9 +104,13 @@ fn resolve_local_model_files(model_dir: &Path) -> Result { }) } -async fn resolve_remote_model_files(model_id: &str) -> Result { - let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?; - let repo = api.model(model_id.to_string()); +async fn resolve_remote_model_files( + model_id: &str, + model_repo: Repo, + cache: &Cache, +) -> Result { + let api = build_api(cache).map_err(|error| Error::Tokenizer(error.to_report_string()))?; + let repo = api.repo(model_repo); let info = repo.info().await.map_err(|error| { Error::Tokenizer(format!( "failed to fetch model '{model_id}': {}", @@ -146,12 +167,20 @@ async fn resolve_remote_model_files(model_id: &str) -> Result Result> { - let cache_repo = Cache::from_env().model(model_id.to_string()); +fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result> { + let cache_repo = cache.repo(repo.clone()); + // hf-hub 0.5 looks up refs even for commit SHAs; Python caches have no such ref. + let snapshot_dir = (repo.revision().len() == 40 + && repo.revision().bytes().all(|byte| byte.is_ascii_hexdigit())) + .then(|| cache.path().join(repo.folder_name()).join("snapshots").join(repo.revision())); + let get_file = |name: &str| match &snapshot_dir { + Some(dir) => local_file_if_exists(dir, name), + None => cache_repo.get(name), + }; - let tokenizer_config_path = cache_repo.get("tokenizer_config.json"); + let tokenizer_config_path = get_file("tokenizer_config.json"); let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?; - let tokenizer = match resolve_cached_tokenizer_source(&cache_repo, &tokenizer_config)? { + let tokenizer = match resolve_cached_tokenizer_source(&get_file, &tokenizer_config)? { Some(tokenizer) => tokenizer, None => return Ok(None), }; @@ -159,12 +188,12 @@ fn resolve_cached_model_files(model_id: &str) -> Result Option, tokenizer_config: &HfTokenizerConfig, ) -> Result> { - let tekken_path = cache_repo.get("tekken.json"); + let tekken_path = get_file("tekken.json"); if let Some(tekken_path) = tekken_path { return Ok(Some(TokenizerSource::Tekken(tekken_path))); } - let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| { + let Some(tokenizer_path) = get_file("tokenizer.json").or_else(|| { // tiktoken.model is the most common name, try it first. - cache_repo.get("tiktoken.model").or_else(|| { + get_file("tiktoken.model").or_else(|| { // Scan for any *.tiktoken file in the cache snapshot directory. - let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf(); + let snapshot_dir = get_file("config.json")?.parent()?.to_path_buf(); discover_tiktoken_in_dir(&snapshot_dir) }) }) else { @@ -315,8 +344,18 @@ async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> }) } -fn build_api() -> anyhow::Result { - let mut builder = ApiBuilder::from_env().with_progress(true); +fn is_hf_hub_offline() -> bool { + std::env::var("HF_HUB_OFFLINE").is_ok_and(|value| { + matches!( + value.to_ascii_uppercase().as_str(), + "1" | "ON" | "YES" | "TRUE" + ) + }) +} + +fn build_api(cache: &Cache) -> anyhow::Result { + let mut builder = + ApiBuilder::from_env().with_cache_dir(cache.path().clone()).with_progress(true); if let Ok(token) = std::env::var(HF_TOKEN_ENV) && !token.is_empty() { @@ -388,6 +427,7 @@ fn discover_chat_template_in_dir(dir: &std::path::Path) -> Option { #[cfg(test)] mod tests { use std::fs; + use std::path::PathBuf; use tempfile::tempdir; use vllm_tokenizer::{TiktokenTokenizer, Tokenizer}; @@ -405,7 +445,7 @@ mod tests { .expect("write tokenizer config"); fs::write(dir.path().join("config.json"), "{}").expect("write config"); - let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path")) + let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"), None) .await .expect("resolve local model files"); @@ -422,10 +462,90 @@ mod tests { ); } + #[test] + fn cached_model_files_respect_branch_tag_and_commit_revisions() { + let dir = tempdir().unwrap(); + let cache = hf_hub::Cache::new(dir.path().to_path_buf()); + let model = "test/model"; + let main_sha = "a".repeat(40); + let pinned_sha = "b".repeat(40); + for (revision, sha) in [("main", &main_sha), ("release/v1", &pinned_sha)] { + let repo = + hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into()); + cache.repo(repo.clone()).create_ref(sha).unwrap(); + let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha); + fs::create_dir_all(&snapshot).unwrap(); + fs::write(snapshot.join("tokenizer.json"), "{}").unwrap(); + fs::write(snapshot.join("config.json"), "{}").unwrap(); + } + for (revision, expected_sha) in [ + ("main", Some(&main_sha)), + ("release/v1", Some(&pinned_sha)), + (pinned_sha.as_str(), Some(&pinned_sha)), + ("missing", None), + ] { + let repo = + hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into()); + let files = super::resolve_cached_model_files(&cache, &repo).unwrap(); + match expected_sha { + Some(sha) => { + let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha); + let files = files.unwrap(); + assert_eq!(files.tokenizer.path(), snapshot.join("tokenizer.json")); + assert_eq!(files.config_path, Some(snapshot.join("config.json"))); + } + None => assert!(files.is_none()), + } + } + } + + #[tokio::test] + async fn offline_resolution_uses_hub_cache_and_reports_cache_misses() { + const CHILD: &str = "VLLM_TEST_HF_OFFLINE_CHILD"; + if std::env::var_os(CHILD).is_some() { + let files = ResolvedModelFiles::new("test/model", Some("release")).await.unwrap(); + let cache = PathBuf::from(std::env::var_os("HF_HUB_CACHE").unwrap()); + assert_eq!( + files.tokenizer.path(), + cache.join("models--test--model/snapshots/commit/tokenizer.json") + ); + let error = ResolvedModelFiles::new("test/model", Some("missing")).await.unwrap_err(); + let message = thiserror_ext::AsReport::as_report(&error).to_string(); + assert!(message.contains("HF_HUB_OFFLINE is enabled"), "{message}"); + assert!(message.contains("revision 'missing'"), "{message}"); + return; + } + + // Isolate environment variables from other tests and Tokio worker threads. + let dir = tempdir().unwrap(); + let cache = dir.path().join("hub-cache"); + let repo = cache.join("models--test--model"); + fs::create_dir_all(repo.join("refs")).unwrap(); + fs::create_dir_all(repo.join("snapshots/commit")).unwrap(); + fs::write(repo.join("refs/release"), "commit").unwrap(); + fs::write(repo.join("snapshots/commit/tokenizer.json"), "{}").unwrap(); + for value in ["1", "ON", "yes", "True"] { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "backend::hf::model_files::tests::offline_resolution_uses_hub_cache_and_reports_cache_misses", "--nocapture"]) + .env(CHILD, "1") + .env("HF_HOME", dir.path().join("hf-home")) + .env("HF_HUB_CACHE", &cache) + .env("HF_HUB_OFFLINE", value) + .env("HF_ENDPOINT", "http://127.0.0.1:1") + .output().unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } + } + #[tokio::test] #[ignore = "too slow for CI and requires network access to Hugging Face"] async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() { - let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5") + let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5", None) .await .expect("resolve real Kimi K2.5 model files"); diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 0b18fd423821..bff4a6d48bb1 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -802,7 +802,8 @@ mod tests { #[file_serial(hf_qwen3)] async fn lower_text_request_uses_real_qwen_generation_defaults() { let model_id = "Qwen/Qwen3-0.6B"; - let files = ResolvedModelFiles::new(model_id).await.expect("resolve qwen model files"); + let files = + ResolvedModelFiles::new(model_id, None).await.expect("resolve qwen model files"); let backend = HfTextBackend::from_resolved_model_files( files, model_id.to_string(), From 120ec4ebd2803c3cc63729e18fcfd42b5ba47447 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sat, 12 Sep 2026 06:05:59 +0200 Subject: [PATCH 081/774] [5/N][warmup][DSv4] Migrate NVIDIA CuTeDSL attention kernels (#53566) Signed-off-by: LopezCastroRoberto Co-authored-by: Codex --- .../test_jit_warmup_cutedsl_launcher.py | 41 +- .../test_jit_warmup_triton_launcher.py | 8 +- vllm/cute_utils/__init__.py | 8 + .../attention/dsa/dcp_indexer_cutedsl.py | 18 +- vllm/model_executor/warmup/jit_warmup.py | 23 + .../warmup/jit_warmup_cutedsl_helper.py | 51 +- .../warmup/jit_warmup_triton_helper.py | 58 +- vllm/models/common/ops/fused_qk_rmsnorm.py | 2 +- vllm/models/deepseek_v4/attention.py | 26 +- .../deepseek_v4/common/ops/cache_utils.py | 14 +- .../common/ops/fused_compress_quant_cache.py | 2 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 44 +- .../common/ops/fused_inv_rope_fp8_quant.py | 2 +- .../common/ops/fused_mtp_input_rmsnorm.py | 2 +- .../common/ops/save_partial_states.py | 2 +- vllm/models/deepseek_v4/compressor.py | 36 +- .../nvidia/ops/dequant_gather_k_cutedsl.py | 625 +-- .../nvidia/ops/fused_indexer_q_cutedsl.py | 1101 ++--- .../ops/sparse_attn_compress_cutedsl.py | 3551 +++++++++-------- vllm/models/deepseek_v4/sparse_mla.py | 2 +- .../deepseek_v4_1/common/ops/cache_utils.py | 12 +- .../backends/mla/compressor_utils.py | 3 +- vllm/v1/attention/backends/mla/indexer.py | 2 +- vllm/v1/attention/backends/mla/sparse_swa.py | 2 +- .../v1/attention/backends/mla/sparse_utils.py | 2 +- vllm/v1/attention/ops/common.py | 2 +- vllm/v1/attention/ops/dcp.py | 2 +- vllm/v1/worker/block_table.py | 2 +- 28 files changed, 3108 insertions(+), 2535 deletions(-) diff --git a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py index 2524c60d2e37..76ec4d1200b1 100644 --- a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py +++ b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py @@ -5,10 +5,10 @@ from typing import Any from vllm.model_executor.warmup import jit_warmup_cutedsl_helper +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( CuTeDSLLaunchSpec, VllmCuTeDSLJitKernel, - cutedsl_kernel_launcher, ) @@ -30,7 +30,7 @@ def get_warmup_keys(self) -> list[CompileKey]: def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: return (f"fake-{compile_key.variant}",) - @cutedsl_kernel_launcher + @kernel_launcher def __call__( self, payload: str, @@ -38,7 +38,7 @@ def __call__( variant: int, ) -> CuTeDSLLaunchSpec[CompileKey]: compile_key = self.dispatch(variant=variant) - return compile_key, (payload,), {"variant": variant} + return compile_key, (payload,) def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None: @@ -74,3 +74,38 @@ def executor(*runtime_args: Any) -> tuple[int, tuple[Any, ...]]: (("entry", 2), ("fake-2",)), ] assert launch_calls == [(1, ("runtime",)), (2, ("other",))] + + +def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None: + events: list[str] = [] + + def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any: + def executor(*runtime_args: Any) -> None: + events.append("launch") + + return executor + + class FinalizedKernel(_TestCuTeDSLKernel): + @kernel_launcher + def __call__( + self, + payload: str, + *, + variant: int, + ) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]: + compile_key = self.dispatch(variant=variant) + + def finalize() -> str: + events.append("finalize") + return "output" + + return compile_key, (payload,), "unused", finalize + + monkeypatch.setattr( + jit_warmup_cutedsl_helper, + "compile_cutedsl", + compile_cutedsl, + ) + + assert FinalizedKernel()("runtime", variant=1) == "output" + assert events == ["launch", "finalize"] diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index 2e624cc7f614..e76041532ff0 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -4,10 +4,10 @@ from dataclasses import dataclass from typing import Any +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, VllmTritonJitKernel, - kernel_launcher, ) @@ -107,13 +107,13 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None: ) launches: list[tuple[Any, ...]] = [] - def launch(grid: Any, inputs: Any, **kwargs: Any) -> None: - launches.append((grid, inputs, kwargs)) + def launch(launch_spec: Any, inputs: Any) -> None: + launches.append((launch_spec, inputs)) monkeypatch.setattr(owner, "launch", launch) owner.compile(compile_key) - grid, inputs, kwargs = launches[0] + (grid, kwargs), inputs = launches[0] assert grid == (2,) assert inputs["num_tokens"] == 2 assert inputs["block_table_stride"] == 128 diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index 23721aac440c..033533a5c9a0 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -22,6 +22,14 @@ torch.float32: Float32, } + +def torch_to_cute_dtype(dtype: torch.dtype): + try: + return _TORCH_TO_CUTE_DTYPE[dtype] + except KeyError as exc: + raise TypeError(f"Unsupported CuTe dtype {dtype}") from exc + + _CUTE_TO_PTX_DTYPE = { BFloat16: "bf16", Float16: "f16", diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py index 29cdb2c589e3..34ce32856ad0 100644 --- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py +++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py @@ -12,16 +12,15 @@ from quack.compile_utils import make_fake_tensor from vllm.cute_utils import recast_val +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( CuTeDSLLaunchSpec, VllmCuTeDSLJitKernel, - cutedsl_kernel_launcher, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton @@ -569,7 +568,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: ) return gathered, out - @cutedsl_kernel_launcher + @kernel_launcher def __call__( self, gathered: torch.Tensor, @@ -577,16 +576,11 @@ def __call__( *, topk: int, ) -> CuTeDSLLaunchSpec[CompileKey]: - compile_key = self.dispatch(topk=topk, num_candidates=gathered.shape[1]) - return ( - compile_key, - (gathered, out), - { - "gathered_shape": tuple(gathered.shape), - "out_shape": tuple(out.shape), - "topk": topk, - }, + compile_key = self.dispatch( + topk=topk, + num_candidates=gathered.shape[1], ) + return compile_key, (gathered, out) _PACK_DCP_TOPK_CANDIDATES_KERNEL = PackDCPTopkCandidatesKernel() diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 0be4bddea6b3..fc9d6fcdf895 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -15,6 +15,7 @@ from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, fields +from functools import wraps from typing import Any, Generic, TypeVar, cast __all__ = [ @@ -23,6 +24,7 @@ "WarmupIntRange", "get_ast_full_name", "get_function_source_node", + "kernel_launcher", "zip_inputs", ] @@ -30,6 +32,26 @@ CompileKeyT = TypeVar("CompileKeyT") +def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]: + """Delegate a declarative launch specification to its backend owner.""" + signature = inspect.signature(call_fn) + + @wraps(call_fn) + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + launch_spec = call_fn(self, *args, **kwargs) + if not self.bind_launch_inputs: + return self.launch(launch_spec, {}) + + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + inputs = { + name: value for name, value in bound.arguments.items() if name != "self" + } + return self.launch(launch_spec, inputs) + + return wrapper + + @dataclass(frozen=True) class WarmupIntRange: """Expand integers with range semantics or a custom monotonic progression.""" @@ -549,6 +571,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" CompileKey: type[CompileKeyT] + bind_launch_inputs = True def __init__(self) -> None: self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch) diff --git a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py index 7311cc18a419..879c8c6998d1 100644 --- a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py @@ -3,19 +3,25 @@ from abc import abstractmethod from collections.abc import Callable, Mapping -from functools import wraps from typing import Any, ClassVar, Generic, TypeAlias, TypeVar -from vllm.model_executor.warmup.jit_warmup import VllmJitKernel +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, +) DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi" CompileKeyT = TypeVar("CompileKeyT") -CuTeDSLLaunchSpec: TypeAlias = tuple[ - CompileKeyT, - tuple[Any, ...], - Mapping[str, Any] | None, -] +CuTeDSLLaunchSpec: TypeAlias = ( + tuple[CompileKeyT, tuple[Any, ...]] + | tuple[CompileKeyT, tuple[Any, ...], Any] + | tuple[ + CompileKeyT, + tuple[Any, ...], + Any, + Callable[[], Any], + ] +) def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any: @@ -44,6 +50,7 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): """CuTeDSL owner whose compiled executor is shared by warmup and runtime.""" kernel: ClassVar[Any] + bind_launch_inputs = False @abstractmethod def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]: @@ -58,23 +65,15 @@ def compile(self, compile_key: CompileKeyT) -> None: *self.warmup_inputs(compile_key), ) - -def cutedsl_kernel_launcher( - call_fn: Callable[..., CuTeDSLLaunchSpec[CompileKeyT]], -) -> Callable[..., Any]: - """Invoke a cached CuTeDSL executor from a declarative ``__call__``.""" - - @wraps(call_fn) - def wrapper( - self: VllmCuTeDSLJitKernel[CompileKeyT], - *args: Any, - **kwargs: Any, + def launch( + self, + launch_spec: CuTeDSLLaunchSpec[CompileKeyT], + _inputs: Mapping[str, Any], ) -> Any: - compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs) - executor = self._get_or_compile( - compile_key, - runtime_context=runtime_context, - ) - return executor(*launch_args) - - return wrapper + # (compile_key, args), optionally followed by output and epilogue. + compile_key, launch_args = launch_spec[:2] + executor = self._get_or_compile(compile_key) + result = executor(*launch_args) + if len(launch_spec) == 4: + return launch_spec[3]() + return launch_spec[2] if len(launch_spec) == 3 else result diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 7fffa12ec901..b019404cc9d1 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -5,7 +5,7 @@ from abc import abstractmethod from collections.abc import Callable, Mapping from dataclasses import dataclass -from functools import cached_property, wraps +from functools import cached_property from typing import Any, ClassVar, Generic, TypeVar, cast from vllm.model_executor.warmup.jit_warmup import ( @@ -128,14 +128,16 @@ def _kernel_arg_names(self) -> tuple[str, ...]: def launch( self, - grid: tuple[int, ...] | None, + launch_spec: LaunchSpec, inputs: Mapping[str, Any], - /, - *, - kernel: Any = None, - **kwargs: Any, ) -> Any: - kernel = kernel if kernel is not None else self.kernel + if len(launch_spec) == 3: + grid, launch_kwargs, outputs = launch_spec + else: + grid, launch_kwargs = launch_spec + outputs = None + kwargs = dict(launch_kwargs) + kernel = kwargs.pop("kernel", None) or self.kernel runtime_launcher = kwargs.pop("_runtime_launcher", None) runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0) for name, value in inputs.items(): @@ -151,52 +153,20 @@ def launch( kwargs["launch_pdl"] = self._warming_compile_key.launch_pdl warmup = getattr(kernel, "warmup", None) assert warmup is not None - return warmup(grid=(1,), **kwargs) + warmup(grid=(1,), **kwargs) + return outputs if grid is None: - return None + return outputs if runtime_launcher is not None: regular_args = [ kwargs.pop(name) for name in self._kernel_arg_names[:runtime_launcher_arg_count] ] - return runtime_launcher(kernel, grid, *regular_args, **kwargs) - return kernel[grid](**kwargs) - - -def kernel_launcher( - call_fn: Callable[..., LaunchSpec], -) -> Callable[..., Any]: - """Launch a Triton kernel from a declarative ``__call__`` specification. - - ``call_fn`` returns either ``(grid, launch_kwargs)`` or, when it allocates - its own outputs, ``(grid, launch_kwargs, outputs)``. The declared outputs are - returned to the caller. A ``grid`` of ``None`` skips the launch (e.g. an - empty-token batch) while still returning the outputs. - """ - signature = inspect.signature(call_fn) - - @wraps(call_fn) - def wrapper( - self: VllmTritonJitKernel[Any], - *args: Any, - **kwargs: Any, - ) -> Any: - spec = call_fn(self, *args, **kwargs) - if len(spec) == 3: - grid, launch_kwargs, outputs = spec + runtime_launcher(kernel, grid, *regular_args, **kwargs) else: - grid, launch_kwargs = spec - outputs = None - bound = signature.bind(self, *args, **kwargs) - bound.apply_defaults() - inputs = { - name: value for name, value in bound.arguments.items() if name != "self" - } - self.launch(grid, inputs, **launch_kwargs) + kernel[grid](**kwargs) return outputs - return wrapper - @dataclass(frozen=True) class TritonPointerInputVariant: diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py index 2ee6cc2854d2..e64e81821062 100644 --- a/vllm/models/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -5,11 +5,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 12205bc0431e..d32a48cbe43f 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -50,6 +50,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -420,7 +421,6 @@ def __init__( _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup() - from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl _FUSED_Q_KV_RMSNORM_KERNEL.register_warmup() @@ -447,7 +447,13 @@ def __init__( ) _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup() - if not has_cutedsl(): + if has_cutedsl(): + from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( # noqa: E501 + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, + ) + + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL.register_warmup() + else: _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL.register_warmup() elif backend_name == "FLASHINFER_MLA_SPARSE_DSV4": from vllm.models.deepseek_v4.common.ops.cache_utils import ( @@ -934,6 +940,21 @@ def __init__( "Using %s indexer cache for Lightning Indexer.", "MXFP4" if self.use_fp4_kv else "FP8", ) + if vllm_config.kernel_config.enable_jit_warmup: + from vllm.utils.import_utils import has_cutedsl + + if current_platform.is_cuda() and has_cutedsl(): + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( # noqa: E501 + _INDEXER_Q_FP8_KERNEL, + _INDEXER_Q_MXFP4_KERNEL, + ) + + indexer_q_kernel = ( + _INDEXER_Q_MXFP4_KERNEL + if self.use_fp4_kv + else _INDEXER_Q_FP8_KERNEL + ) + indexer_q_kernel.register_warmup() # no tensor parallel, just replicated self.wq_b = ReplicatedLinear( @@ -1016,7 +1037,6 @@ def __init__( ] if vllm_config.kernel_config.enable_jit_warmup: - from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl if not has_cutedsl() and not current_platform.is_xpu(): diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 1bd0706f0fe3..d98d6a297a43 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -24,6 +24,7 @@ ) from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( @@ -31,7 +32,6 @@ TritonPointerInputVariant, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.platforms import current_platform @@ -521,11 +521,17 @@ def dequantize_and_gather_k_cache( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( - dequantize_and_gather_k_cache_cutedsl, + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) - dequantize_and_gather_k_cache_cutedsl( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( + out=out, + k_cache=k_cache, + seq_lens=seq_lens, + gather_lens=gather_lens, + block_table=block_table, + block_size=block_size, + offset=offset, ) return diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index e0bc2a042831..326a94f4f43c 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -26,11 +26,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index ae014750bf90..cf6e7235f9ef 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -6,11 +6,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -598,19 +598,19 @@ def fused_indexer_q_rope_quant( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_mxfp4_cutedsl, + _INDEXER_Q_MXFP4_KERNEL, ) - fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_packed, - index_q_scale, - index_weights_out, + _INDEXER_Q_MXFP4_KERNEL( + positions=positions, + q=index_q, + cos_sin_cache=index_q_cos_sin_cache, + weights=index_weights, + weights_softmax_scale=index_weights_softmax_scale, + weights_head_scale=index_weights_head_scale, + q_packed=index_q_packed, + q_scale=index_q_scale, + weights_out=index_weights_out, ) elif current_platform.is_xpu(): torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( @@ -654,18 +654,18 @@ def fused_indexer_q_rope_quant( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_fp8_cutedsl, + _INDEXER_Q_FP8_KERNEL, ) - fused_indexer_q_rope_quant_fp8_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_fp8, - index_weights_out, + _INDEXER_Q_FP8_KERNEL( + positions=positions, + q=index_q, + cos_sin_cache=index_q_cos_sin_cache, + weights=index_weights, + weights_softmax_scale=index_weights_softmax_scale, + weights_head_scale=index_weights_head_scale, + q_fp8=index_q_fp8.view(torch.uint8), + weights_out=index_weights_out, ) elif current_platform.is_xpu(): torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8( diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py index 948d84be5c4a..77b1aa2236b5 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -12,11 +12,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py index c71043f1ae09..59ac4d27c648 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py +++ b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py @@ -22,11 +22,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 diff --git a/vllm/models/deepseek_v4/common/ops/save_partial_states.py b/vllm/models/deepseek_v4/common/ops/save_partial_states.py index 0537feab2956..6a554b9b4e02 100644 --- a/vllm/models/deepseek_v4/common/ops/save_partial_states.py +++ b/vllm/models/deepseek_v4/common/ops/save_partial_states.py @@ -6,11 +6,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 4da343b2c98d..41af0c76605a 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -203,8 +203,8 @@ class DeepseekCompressor(nn.Module): prologue (kv/score split, save_partial_states launch). The compress → norm → RoPE → store step is dispatched to a triton kernel (``compress_norm_rope_store_triton``) by default, except for the NVIDIA - head_dim=128 indexer path which uses the cutedsl kernel - (``compress_norm_rope_store_cutedsl``) for better performance. + head_dim=512 path which uses the CuTeDSL compressor kernels for better + performance. """ def __init__( @@ -317,7 +317,33 @@ def __init__( head_dim=self.head_dim, compress_ratio=self.compress_ratio, ) - if self.head_dim != 512: + if current_platform.is_cuda() and self.head_dim == 512: + from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import ( # noqa: E501 + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL, + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL, + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL, + ) + + store_full_kv = vllm_config.cache_config.cache_dtype != "fp8_ds_mla" + if self.compress_ratio == 4: + ( + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL + if store_full_kv + else _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL + ).register_warmup() + else: + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL.register_warmup() + if store_full_kv: + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL.register_warmup() + else: + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL.register_warmup( + vllm_config, + k_cache_prefix=self.k_cache_prefix, + compress_ratio=self.compress_ratio, + ) + else: from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( # noqa: E501 _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL, ) @@ -419,13 +445,13 @@ def forward( compress_norm_rope_store_fn: Any if current_platform.is_cuda() and self.head_dim == 512: from .nvidia.ops.sparse_attn_compress_cutedsl import ( - compress_norm_rope_store_cutedsl, + _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL, ) # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla # layout and the plain full-cache layout. The full-cache flags # are consumed only here. - compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl + compress_norm_rope_store_fn = _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL extra_kwargs: dict[str, Any] = dict( store_full_kv=store_full_kv, store_full_fp8=store_full_fp8, diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index a4514293c942..fcb12740f27f 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache +from dataclasses import dataclass +from typing import Any import cutlass import cutlass.cute as cute @@ -12,302 +13,340 @@ from quack.compile_utils import make_fake_tensor from vllm.cute_utils import _bf16x2_mul, cvt +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, +) -def dequantize_and_gather_k_cache_cutedsl( - out: torch.Tensor, - k_cache: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor | None, - block_table: torch.Tensor, - block_size: int, - offset: int, -) -> None: - DequantGatherKCacheKernel.compile( - block_size=block_size, - has_gather_lens=gather_lens is not None, - )(out, k_cache, seq_lens, gather_lens, block_table, offset) - - -class DequantGatherKCacheKernel: +class DequantGatherKCacheKernel( + VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] +): # Hard-coded for DSv4. head_dim = 512 + fp8_dim = 448 group_size = 64 # 1 scale per 64 elems + num_warps = 4 + tb_size = num_warps * 32 + num_stages = 4 - def __init__(self, fp8_dim: int = 448, block_size: int = 64): - self.fp8_dim = fp8_dim - self.bf16_dim = self.head_dim - fp8_dim - self.data_dim = fp8_dim + self.bf16_dim * 2 - self.block_size = block_size - - self.num_warps = 4 - self.tb_size = self.num_warps * 32 - self.num_stages = 4 - - @cute.jit - def __call__( - self, - out: cute.Tensor, - k_cache: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - stream: CUstream, - ): - # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] - # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. - k_data = cute.make_tensor( - k_cache.iterator, - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, self.data_dim), - stride=(k_cache.stride[0], self.data_dim, 1), - ), - ) - k_scale = cute.make_tensor( - k_cache.iterator + (self.block_size * self.data_dim), - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, 8), - stride=(k_cache.stride[0], 8, 1), - ), - ) - - num_reqs = out.shape[0] - max_rows = out.shape[1] - offset - num_workers = cutlass.max( - 1, cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)) - ) - grid = (num_reqs, num_workers, 1) - self.kernel( - out, - k_data, - k_scale, - seq_lens, - gather_lens, - block_table, - offset, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + @dataclass(frozen=True) + class CompileKey: + block_size: int + has_gather_lens: bool - @cute.jit - def load_g2s( - self, - k_data_slice: cute.Tensor, - k_scale: cute.Tensor, - block_table: cute.Tensor, - s_kdata_slice: cute.Tensor, - s_kscale: cute.Tensor, - req_id, - pos, - lane_id, - stage_id, - ): - # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] - # s_kdata_slice: [(4, data_dim/16), num_stages] - - op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) - cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64) - page_id = block_table[req_id, pos // self.block_size] - block_offset = pos % self.block_size - - # Load the first 512 bytes (32x16B). - idx = lane_id - src = k_data_slice[page_id, block_offset, (None, idx)] - cute.copy( - cp16_atom, - cute.recast_tensor(src, Uint32), - s_kdata_slice[(None, idx), stage_id], - ) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = DequantGatherKCacheKernel.head_dim + fp8_dim = DequantGatherKCacheKernel.fp8_dim + group_size = DequantGatherKCacheKernel.group_size + num_warps = DequantGatherKCacheKernel.num_warps + tb_size = DequantGatherKCacheKernel.tb_size + num_stages = DequantGatherKCacheKernel.num_stages + bf16_dim = head_dim - fp8_dim + data_dim = fp8_dim + bf16_dim * 2 + block_size = compile_key.block_size + + @cute.jit + def load_g2s( + k_data_slice: cute.Tensor, + k_scale: cute.Tensor, + block_table: cute.Tensor, + s_kdata_slice: cute.Tensor, + s_kscale: cute.Tensor, + req_id, + pos, + lane_id, + stage_id, + ): + op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) + cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) + cp8_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + ) + page_id = block_table[req_id, pos // block_size] + block_offset = pos % block_size - # Load the tail 64 bytes. - idx += 32 - if idx < cutlass.const_expr(self.data_dim // 16): + # Load the first 512 bytes (32x16B). + idx = lane_id src = k_data_slice[page_id, block_offset, (None, idx)] cute.copy( cp16_atom, cute.recast_tensor(src, Uint32), s_kdata_slice[(None, idx), stage_id], ) - elif idx == cutlass.const_expr(self.data_dim // 16): - cute.copy( - cp8_atom, - k_scale[page_id, block_offset, None], - s_kscale[None, stage_id], - ) - @cute.kernel - def kernel( - self, - out: cute.Tensor, - k_data: cute.Tensor, - k_scale: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - ): - req_id, worker_id, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - - _, num_workers, _ = cute.arch.grid_dim() - - # Prepare smem. - smem = cutlass.utils.SmemAllocator() - s_kdata = smem.allocate_tensor( - Uint32, - cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)), - byte_alignment=16, - )[None, warp_id, None] - s_kscale = smem.allocate_tensor( - Uint8, - cute.make_layout((8, self.num_warps, self.num_stages)), - byte_alignment=8, - )[None, warp_id, None] - - # Prepare for 16B cp.async, also for BF16 smem loads later. - k_data_slice = cute.logical_divide(k_data, (None, None, 16)) - s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) - - # Load FP8 elems in 8B units, so once dequantized, they are 16B units. - s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) - - # 16B st.global. - out_slice = cute.logical_divide(out, (None, None, 8)) - - cp_op = cute.nvgpu.CopyUniversalOp() - cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - - seq_len = seq_lens[req_id] - gather_len = seq_len - if cutlass.const_expr(gather_lens is not None): - gather_len = gather_lens[req_id] # type: ignore[index] - start_pos = seq_len - gather_len - - # Start prefetch. - for i in cutlass.range_constexpr(self.num_stages - 1): - next_pos = ( - start_pos - + worker_id * self.num_warps - + warp_id - + i * num_workers * self.num_warps - ) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - i, + # Load the tail 64 bytes. + idx += 32 + if idx < cutlass.const_expr(data_dim // 16): + src = k_data_slice[page_id, block_offset, (None, idx)] + cute.copy( + cp16_atom, + cute.recast_tensor(src, Uint32), + s_kdata_slice[(None, idx), stage_id], ) - cute.arch.cp_async_commit_group() - prefetch_stage = self.num_stages - 1 - compute_stage = 0 - - # Main loop. - for i in range( - worker_id * self.num_warps + warp_id, - gather_len, - num_workers * self.num_warps, - ): - pos = start_pos + i - - # Prefetch next stage. - next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - prefetch_stage, + elif idx == cutlass.const_expr(data_dim // 16): + cute.copy( + cp8_atom, + k_scale[page_id, block_offset, None], + s_kscale[None, stage_id], ) - prefetch_stage = (prefetch_stage + 1) % self.num_stages - cute.arch.cp_async_commit_group() - - # Wait for gmem->smem to finish. - cute.arch.cp_async_wait_group(self.num_stages - 1) - cute.arch.sync_warp() - - # There are 512 elems per token. As a warp, data0 holds the first - # 256 elems and data1 holds the second 256 elems, i.e. each thread - # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as - # contiguous 16B global stores. On Blackwell, this might not be - # necessary as we have 32B global stores, but doing it this way - # does not seem to be slower. - data0 = cute.make_rmem_tensor((2,), Uint32) - data1 = cute.make_rmem_tensor((2,), Uint32) - cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0) - cute.copy( - cp8_atom, - s_kdata_8B_slice[(None, lane_id + 32), compute_stage], - data1, - ) - # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 - # elements. An 8-element chunk advances the scale index by - # chunk_id * 8 // group_size. - scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage]) - scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) - scale1_u32 = Uint32( - s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage] - ) - scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) - - # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 - # (CUDA 13.2). - dequant0 = cute.make_rmem_tensor(4, Uint32) - dequant1 = cute.make_rmem_tensor(4, Uint32) - for j in cutlass.range_constexpr(2): - tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) - tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) - - # BF16 multiply is safe because the scales are exact powers of 2. - dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) - dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) - dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) - dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) - - # Last 64 elems are BF16 tail, corresponds to dequant1 of last - # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= self.fp8_dim // 8: - idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8 + @cute.kernel + def device_kernel( + out: cute.Tensor, + k_data: cute.Tensor, + k_scale: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + ): + req_id, worker_id, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + + _, num_workers, _ = cute.arch.grid_dim() + + # Prepare smem. + smem = cutlass.utils.SmemAllocator() + s_kdata = smem.allocate_tensor( + Uint32, + cute.make_layout((data_dim // 4, num_warps, num_stages)), + byte_alignment=16, + )[None, warp_id, None] + s_kscale = smem.allocate_tensor( + Uint8, + cute.make_layout((8, num_warps, num_stages)), + byte_alignment=8, + )[None, warp_id, None] + + # Prepare for 16B cp.async, also for BF16 smem loads later. + k_data_slice = cute.logical_divide(k_data, (None, None, 16)) + s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) + + # Load FP8 elems in 8B units, so once dequantized, they are 16B units. + s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) + + # 16B st.global. + out_slice = cute.logical_divide(out, (None, None, 8)) + + cp_op = cute.nvgpu.CopyUniversalOp() + cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + seq_len = seq_lens[req_id] + gather_len = seq_len + if cutlass.const_expr(gather_lens is not None): + gather_len = gather_lens[req_id] # type: ignore[index] + start_pos = seq_len - gather_len + + # Start prefetch. + for i in cutlass.range_constexpr(num_stages - 1): + next_pos = ( + start_pos + + worker_id * num_warps + + warp_id + + i * num_workers * num_warps + ) + if next_pos < seq_len: + load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + i, + ) + cute.arch.cp_async_commit_group() + prefetch_stage = num_stages - 1 + compute_stage = 0 + + # Main loop. + for i in range( + worker_id * num_warps + warp_id, + gather_len, + num_workers * num_warps, + ): + pos = start_pos + i + + # Prefetch next stage. + next_pos = pos + num_workers * num_warps * (num_stages - 1) + if next_pos < seq_len: + load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + prefetch_stage, + ) + prefetch_stage = (prefetch_stage + 1) % num_stages + cute.arch.cp_async_commit_group() + + # Wait for gmem->smem to finish. + cute.arch.cp_async_wait_group(num_stages - 1) + cute.arch.sync_warp() + + # There are 512 elems per token. As a warp, data0 holds the first + # 256 elems and data1 holds the second 256 elems, i.e. each thread + # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as + # contiguous 16B global stores. On Blackwell, this might not be + # necessary as we have 32B global stores, but doing it this way + # does not seem to be slower. + data0 = cute.make_rmem_tensor((2,), Uint32) + data1 = cute.make_rmem_tensor((2,), Uint32) cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, + cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0 + ) + cute.copy( + cp8_atom, + s_kdata_8B_slice[(None, lane_id + 32), compute_stage], + data1, ) - # Store two 16B BF16 chunks per lane: first half, then second half. - dst = out_slice[req_id, offset + i, (None, lane_id)] - cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) + # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 + # elements. An 8-element chunk advances the scale index by + # chunk_id * 8 // group_size. + scale0_u32 = Uint32(s_kscale[lane_id * 8 // group_size, compute_stage]) + scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) + scale1_u32 = Uint32( + s_kscale[(lane_id + 32) * 8 // group_size, compute_stage] + ) + scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) + + # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 + # (CUDA 13.2). + dequant0 = cute.make_rmem_tensor(4, Uint32) + dequant1 = cute.make_rmem_tensor(4, Uint32) + for j in cutlass.range_constexpr(2): + tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) + tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) + + # BF16 multiply is safe because the scales are exact powers of 2. + dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) + dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) + dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) + dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) + + # Last 64 elems are BF16 tail, corresponds to dequant1 of last + # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + + # 8x 16B for BF16. + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) + + # Store two 16B BF16 chunks per lane: first half, then second half. + dst = out_slice[req_id, offset + i, (None, lane_id)] + cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) + + dst = out_slice[req_id, offset + i, (None, lane_id + 32)] + cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) + + compute_stage = (compute_stage + 1) % num_stages + + @cute.jit + def host_entrypoint( + out: cute.Tensor, + k_cache: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + stream: CUstream, + ): + # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] + # block is actually a concat of + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + k_data = cute.make_tensor( + k_cache.iterator, + layout=cute.make_layout( + (k_cache.shape[0], block_size, data_dim), + stride=(k_cache.stride[0], data_dim, 1), + ), + ) + # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] + # s_kdata_slice: [(4, data_dim/16), num_stages] + k_scale = cute.make_tensor( + k_cache.iterator + (block_size * data_dim), + layout=cute.make_layout( + (k_cache.shape[0], block_size, 8), + stride=(k_cache.stride[0], 8, 1), + ), + ) - dst = out_slice[req_id, offset + i, (None, lane_id + 32)] - cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) + num_reqs = out.shape[0] + max_rows = out.shape[1] - offset + num_workers = cutlass.max( + 1, + cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)), + ) + grid = (num_reqs, num_workers, 1) + device_kernel( + out, + k_data, + k_scale, + seq_lens, + gather_lens, + block_table, + offset, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + block_size: int, + has_gather_lens: bool, + ) -> CompileKey: + return self.CompileKey( + block_size=block_size, + has_gather_lens=has_gather_lens, + ) - compute_stage = (compute_stage + 1) % self.num_stages + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + block_size = vllm_config.cache_config.block_size + compress_ratios = vllm_config.model_config.hf_config.compress_ratios + return self._trace_dispatch(self.dispatch)( + zip_inputs( + dict( + block_size=block_size, + has_gather_lens=True, + enabled=True, + ), + dict( + block_size=max(1, block_size // 4), + has_gather_lens=False, + enabled=4 in compress_ratios, + ), + dict( + block_size=max(1, block_size // 128), + has_gather_lens=False, + enabled=128 in compress_ratios, + ), + ), + valid_block_size=block_size > 0, + _when=lambda *, enabled, valid_block_size: enabled and valid_block_size, + ) - @cache - @staticmethod - def compile( - fp8_dim: int = 448, - block_size: int = 64, - has_gather_lens: bool = True, - ): + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() - head_dim = DequantGatherKCacheKernel.head_dim + head_dim = self.head_dim + fp8_dim = self.fp8_dim + block_size = compile_key.block_size head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) @@ -318,19 +357,47 @@ def compile( assumed_align=32, ) seq_lens = make_fake_tensor(Int32, (num_reqs,)) - gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None + gather_lens = ( + make_fake_tensor(Int32, (num_reqs,)) + if compile_key.has_gather_lens + else None + ) block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int())) - kernel = DequantGatherKCacheKernel(fp8_dim, block_size) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( out, k_cache, seq_lens, gather_lens, block_table, Int32(0), - stream, - options="--enable-tvm-ffi", ) + + @kernel_launcher + def __call__( + self, + *, + out: torch.Tensor, + k_cache: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor | None, + block_table: torch.Tensor, + block_size: int, + offset: int, + ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: + compile_key = self.dispatch( + block_size=block_size, + has_gather_lens=gather_lens is not None, + ) + launch_args = ( + out, + k_cache, + seq_lens, + gather_lens, + block_table, + offset, + ) + return compile_key, launch_args + + +_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL = DequantGatherKCacheKernel() diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 16e98abf23e1..6684e37f37e5 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache +# ruff: noqa: E501 + +from dataclasses import dataclass +from typing import Any import cutlass import cutlass.cute as cute @@ -14,594 +17,638 @@ _bf16x2_max, cvt, recast_val, + torch_to_cute_dtype, +) +from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, ) from vllm.vllm_flash_attn.cute import utils as cute_utils # MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. MXFP4_BLOCK_SIZE = 32 -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} - - -def fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_packed: torch.Tensor, - index_q_scale: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - # compile all variants at first invocation - for coarsen in (1, 4): - IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - # heuristic - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQMxFp4Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_packed, - index_q_scale, - index_weights_out, - scale, - ) - -def fused_indexer_q_rope_quant_fp8_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_fp8: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - for coarsen in (1, 4): - IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQFp8Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - # The cute kernel treats the FP8 buffer as raw bytes (Uint8). - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_fp8.view(torch.uint8), - index_weights_out, - scale, +@cute.jit +def _load_q_and_rope( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + num_heads: cutlass.Constexpr, + coarsen: cutlass.Constexpr, + tb_size: cutlass.Constexpr, + subwarp_size: cutlass.Constexpr, + threads_per_token: cutlass.Constexpr, + nope_dim: cutlass.Constexpr, + rope_dim: cutlass.Constexpr, + cos_sin_dtype: cutlass.Constexpr, +): + """Compute thread indices, load Q (BF16), and apply interleaved RoPE. + + Returns a tuple + (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, + head_start, in_bounds, num_token_heads) + where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed + bf16x2 pairs covering the 16 BF16 lanes owned by this thread for + each of ``coarsen`` heads. RoPE is applied in place to the + trailing ``rope_dim`` lanes; the leading nope lanes pass through. + """ + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + + num_tokens = q.shape[0] + num_token_heads = num_tokens * num_heads + global_tid = block_id * tb_size + tid + + global_subwarp_id = global_tid // subwarp_size + sublane = tid % subwarp_size + + token_id = global_subwarp_id // (num_heads // coarsen) + head_tile_id = global_subwarp_id % (num_heads // coarsen) + head_start = head_tile_id * coarsen + + # NOTE: token_id may exceed bounds, hence we need to add load/store guards + # we can't do early exit because CuteDSL doesn't support it. and we also need + # all threads in a warp to be active since we utilize warp shuffle later. + # must_in_bounds is constexpr, True when a token spans a whole number of + # threadblocks. Then num_tokens * threads_per_token is always a multiple of + # tb_size, so the trailing block never straddles num_tokens and the compiler + # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads + # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half + # of the last block addressing token num_tokens. + must_in_bounds = cutlass.const_expr(threads_per_token % tb_size == 0) + in_bounds = must_in_bounds or (token_id < num_tokens) + + cp_op = cute.nvgpu.CopyUniversalOp() + + _layout = cute.make_layout((coarsen, 8), stride=(8, 1)) + q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) + + if in_bounds: + # we can't do cute.copy() on the whole 2D tile directly because + # cute.copy() wants the 1st mode to be covered by the copy atom, + # and other modes as for loop. there is no fast way to + # "transpose" the tensor view. + q_tile = cute.local_tile( + q[token_id, None, None], + tiler=(coarsen, 16), + coord=(head_tile_id, sublane), + ) + cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) + for i in cutlass.range_constexpr(coarsen): + src = cute.recast_tensor(q_tile[i, None], Uint32) + cute.copy(cp_u32x8, src, q_bf16x2[i, None]) + + # RoPE applies only to the trailing rope_dim values. We keep the rounded + # BF16 result in q_bits so the later amax and quantization see BF16. + # cos_sin_cache layout: [max_pos, rope_dim] + if in_bounds and sublane * 16 >= nope_dim: + cos_vals = cute.make_rmem_tensor((8,), Float32) + sin_vals = cute.make_rmem_tensor((8,), Float32) + + pos = positions[token_id] + + # select 8 elems from cos and sin + cos_id = sublane - nope_dim // 16 + sin_id = cos_id + rope_dim // 16 + cos_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,)) + sin_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,)) + + cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + if const_expr(cos_sin_dtype is Float32): + cute.copy(cp_f32x8, cos_src, cos_vals) + cute.copy(cp_f32x8, sin_src, sin_vals) + else: + cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) + cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) + + for i in cutlass.range_constexpr(4): + cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) + sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) + cos_vals[i * 2] = cos0 + cos_vals[i * 2 + 1] = cos1 + sin_vals[i * 2] = sin0 + sin_vals[i * 2 + 1] = sin1 + + for i in cutlass.range_constexpr(coarsen): + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] + rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] + # convert back to BF16 to match numerics + q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) + + return ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, ) -class IndexerQRopeQuantKernel: - """Shared infrastructure for indexer-Q RoPE+quant fused kernels. +class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]): + tb_size = 128 - Subclasses implement ``kernel`` for a particular Q quantization scheme - (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the - common preamble: thread/token addressing, the BF16 Q load, and the - interleaved-RoPE pass over the trailing ``rope_dim`` lanes. - """ - - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - self.head_dim = head_dim - self.rope_dim = rope_dim - self.nope_dim = head_dim - rope_dim - self.num_heads = num_heads - self.cos_sin_dtype = cos_sin_dtype + @dataclass(frozen=True) + class CompileKey: + head_dim: int + rope_dim: int + num_heads: int + cos_sin_dtype: type[cutlass.Numeric] + coarsen: int + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_dim + rope_dim = compile_key.rope_dim + nope_dim = head_dim - rope_dim + num_heads = compile_key.num_heads + cos_sin_dtype = compile_key.cos_sin_dtype # process multiple heads at the same time to armotize RoPE load costs + coarsen = compile_key.coarsen + tb_size = IndexerQMxFp4Kernel.tb_size assert num_heads % coarsen == 0 - self.coarsen = coarsen # later we will use 32B load = 16 BF16 elems # thus, head_dim=128 requires 8 threads to handle. # let's call subwarp = 8 threads. - self.subwarp_size = head_dim // 16 - self.tb_size = 128 - self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size - - @cute.jit - def _load_q_and_rope( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - ): - """Compute thread indices, load Q (BF16), and apply interleaved RoPE. - - Returns a tuple - (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, - head_start, in_bounds, num_token_heads) - where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed - bf16x2 pairs covering the 16 BF16 lanes owned by this thread for - each of ``coarsen`` heads. RoPE is applied in place to the - trailing ``rope_dim`` lanes; the leading nope lanes pass through. - """ - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - - num_tokens = q.shape[0] - num_token_heads = num_tokens * self.num_heads - global_tid = block_id * self.tb_size + tid - - global_subwarp_id = global_tid // self.subwarp_size - sublane = tid % self.subwarp_size - - token_id = global_subwarp_id // (self.num_heads // self.coarsen) - head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen) - head_start = head_tile_id * self.coarsen - - # NOTE: token_id may exceed bounds, hence we need to add load/store guards - # we can't do early exit because CuteDSL doesn't support it. and we also need - # all threads in a warp to be active since we utilize warp shuffle later. - # must_in_bounds is constexpr, True when a token spans a whole number of - # threadblocks. Then num_tokens * threads_per_token is always a multiple of - # tb_size, so the trailing block never straddles num_tokens and the compiler - # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads - # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half - # of the last block addressing token num_tokens. - must_in_bounds = cutlass.const_expr(self.threads_per_token % self.tb_size == 0) - in_bounds = must_in_bounds or (token_id < num_tokens) - - cp_op = cute.nvgpu.CopyUniversalOp() - - _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1)) - q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) - - if in_bounds: - # we can't do cute.copy() on the whole 2D tile directly because - # cute.copy() wants the 1st mode to be covered by the copy atom, - # and other modes as for loop. there is no fast way to - # "transpose" the tensor view. - q_tile = cute.local_tile( - q[token_id, None, None], - tiler=(self.coarsen, 16), - coord=(head_tile_id, sublane), - ) - cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) - for i in cutlass.range_constexpr(self.coarsen): - src = cute.recast_tensor(q_tile[i, None], Uint32) - cute.copy(cp_u32x8, src, q_bf16x2[i, None]) - - # RoPE applies only to the trailing rope_dim values. We keep the rounded - # BF16 result in q_bits so the later amax and quantization see BF16. - # cos_sin_cache layout: [max_pos, rope_dim] - if in_bounds and sublane * 16 >= self.nope_dim: - cos_vals = cute.make_rmem_tensor((8,), Float32) - sin_vals = cute.make_rmem_tensor((8,), Float32) - - pos = positions[token_id] - - # select 8 elems from cos and sin - cos_id = sublane - self.nope_dim // 16 - sin_id = cos_id + self.rope_dim // 16 - cos_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,) - ) - sin_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,) + subwarp_size = head_dim // 16 + threads_per_token = (num_heads // coarsen) * subwarp_size + + @cute.kernel + def device_kernel( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, + ) = _load_q_and_rope( + positions, + q, + cos_sin_cache, + num_heads, + coarsen, + tb_size, + subwarp_size, + threads_per_token, + nope_dim, + rope_dim, + cos_sin_dtype, ) - cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - - if const_expr(self.cos_sin_dtype is Float32): - cute.copy(cp_f32x8, cos_src, cos_vals) - cute.copy(cp_f32x8, sin_src, sin_vals) - else: - cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) - cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) - - for i in cutlass.range_constexpr(4): - cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) - sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) - cos_vals[i * 2] = cos0 - cos_vals[i * 2 + 1] = cos1 - sin_vals[i * 2] = sin0 - sin_vals[i * 2 + 1] = sin1 - - for i in cutlass.range_constexpr(self.coarsen): - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] - rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] - # convert back to BF16 to match numerics - q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) - - return ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, - ) + cp_op = cute.nvgpu.CopyUniversalOp() + # layout: [coarsen, 8] + q_fp4_tile = cute.local_tile( + q_quant[token_id, None, None], + tiler=(coarsen, 8), + coord=(head_tile_id, sublane), + ) -class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row.""" - - @cute.jit - def __call__( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( - positions, - q, - cos_sin_cache, - weights, - q_quant, - q_scale, - weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + for i in cutlass.range_constexpr(coarsen): + # compute amax in packed bf16x2 to save instructions + # Each thread holds 16 elems. Two adjacent threads form one 32-elem + # MXFP4 block, so a width-2 shuffle gives the block amax. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=MXFP4_BLOCK_SIZE // 16, + ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) + + if in_bounds: + # compute block scale with bit manipulation + # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask + # increments the exponent whenever fp4_scale is not exactly a power of 2 + eps = cutlass.const_expr(float.fromhex("0x6p-126")) + fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) + bits = recast_val(fp4_scale, Uint32) + ue8m0 = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) + + # Only one of the two threads in an MXFP4 block writes the shared scale. + if tid % 2 == 0: + mx_block = sublane // 2 + q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) + + # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent + # -A + 127 = 254 - ue8m0. + inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) + inv_fp4_scale = recast_val(inv_scale_bits, Float32) + + vals = cute.make_rmem_tensor(16, Float32) + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + vals[j * 2] = q0 * inv_fp4_scale + vals[j * 2 + 1] = q1 * inv_fp4_scale + + # pack to FP4 + packed = cute.make_rmem_tensor((2,), Uint32) + packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) + packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) + + dst = q_fp4_tile[i, None] + cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) + + # Weight scaling is independent of the Q subwarp work. The first + # num_tokens * num_heads logical threads cover one weight each. + if global_tid < num_token_heads: + weight_token_id = global_tid // num_heads + weight_head_id = global_tid % num_heads + weights_out[weight_token_id, weight_head_id] = ( + weights[weight_token_id, weight_head_id].to(Float32) * scale + ) - @cute.kernel - def kernel( + @cute.jit + def host_entrypoint( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * threads_per_token + grid = (cute.ceil_div(total_threads, tb_size), 1, 1) + device_kernel( + positions, + q, + cos_sin_cache, + weights, + q_quant, + q_scale, + weights_out, + scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - - cp_op = cute.nvgpu.CopyUniversalOp() - - # layout: [coarsen, 8] - q_fp4_tile = cute.local_tile( - q_quant[token_id, None, None], - tiler=(self.coarsen, 8), - coord=(head_tile_id, sublane), + *, + cos_sin_dtype: type[cutlass.Numeric], + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + cos_sin_dtype=cos_sin_dtype, ) - for i in cutlass.range_constexpr(self.coarsen): - # compute amax in packed bf16x2 to save instructions - # Each thread holds 16 elems. Two adjacent threads form one 32-elem - # MXFP4 block, so a width-2 shuffle gives the block amax. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=MXFP4_BLOCK_SIZE // 16, - ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - - if in_bounds: - # compute block scale with bit manipulation - # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask - # increments the exponent whenever fp4_scale is not exactly a power of 2 - eps = cutlass.const_expr(float.fromhex("0x6p-126")) - fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) - bits = recast_val(fp4_scale, Uint32) - ue8m0 = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) - - # Only one of the two threads in an MXFP4 block writes the shared scale. - if tid % 2 == 0: - mx_block = sublane // 2 - q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) - - # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent - # -A + 127 = 254 - ue8m0. - inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) - inv_fp4_scale = recast_val(inv_scale_bits, Float32) - - vals = cute.make_rmem_tensor(16, Float32) - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - vals[j * 2] = q0 * inv_fp4_scale - vals[j * 2 + 1] = q1 * inv_fp4_scale - - # pack to FP4 - packed = cute.make_rmem_tensor((2,), Uint32) - packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) - packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) - - dst = q_fp4_tile[i, None] - cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) - - # Weight scaling is independent of the Q subwarp work. The first - # num_tokens * num_heads logical threads cover one weight each. - if global_tid < num_token_heads: - weight_token_id = global_tid // self.num_heads - weight_head_id = global_tid % self.num_heads - weights_out[weight_token_id, weight_head_id] = ( - weights[weight_token_id, weight_head_id].to(Float32) * scale - ) + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4 + + hf_config = vllm_config.model_config.hf_config + if hf_config is None or not dsa_indexer_uses_fp4(vllm_config): + return [] + + num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0) + head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0) + rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0) + if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: + return [] + + return self._trace_dispatch(self.dispatch)( + head_dim=head_dim, + rope_dim=rope_dim, + num_heads=num_heads, + cos_sin_dtype=Float32, + coarsen=(1, 4), + ) - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_tokens = cute.sym_int() max_pos = cute.sym_int() - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + BFloat16, + (num_tokens, compile_key.num_heads, compile_key.head_dim), + divisibility=16, ) positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), + compile_key.cos_sin_dtype, + (max_pos, compile_key.rope_dim), divisibility=8, ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) - q_fp4 = make_fake_tensor( + weights = make_fake_tensor( + BFloat16, (num_tokens, compile_key.num_heads), divisibility=8 + ) + q_packed = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim // 2), + (num_tokens, compile_key.num_heads, compile_key.head_dim // 2), divisibility=16, ) q_scale = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE), + ( + num_tokens, + compile_key.num_heads, + compile_key.head_dim // MXFP4_BLOCK_SIZE, + ), divisibility=4, ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQMxFp4Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + weights_out = make_fake_tensor( + Float32, (num_tokens, compile_key.num_heads), divisibility=4 ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( positions, q, cos_sin_cache, weights, - q_fp4, + q_packed, q_scale, weights_out, Float32(0.0), - stream, - options="--enable-tvm-ffi", ) - -class IndexerQFp8Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row and emit - float8 e4m3fn with a single per-(token, head) scalar scale folded - into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``). - """ - - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen) - # Each subwarp owns `coarsen` heads; we use the first `coarsen` - # threads of the subwarp to write the per-head weights using the - # fp8 scale computed in the matching loop iteration. - assert self.coarsen <= self.subwarp_size, ( - f"FP8 kernel requires coarsen ({self.coarsen}) <= " - f"subwarp_size ({self.subwarp_size}) for the weight-fold step" - ) - - @cute.jit + @kernel_launcher def __call__( self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( + *, + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + weights_softmax_scale: float, + weights_head_scale: float, + q_packed: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + ) -> CuTeDSLLaunchSpec["IndexerQMxFp4Kernel.CompileKey"]: + num_tokens, num_heads, head_dim = q.shape + compile_key = self.dispatch( + head_dim=head_dim, + rope_dim=cos_sin_cache.shape[-1], + num_heads=num_heads, + cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + coarsen=1 if num_tokens < 512 else 4, + ) + launch_args = ( positions, q, cos_sin_cache, weights, - q_fp8, + q_packed, + q_scale, weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + float(weights_softmax_scale * weights_head_scale), + ) + return compile_key, launch_args - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - _tid, - _global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - _num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - - cp_op = cute.nvgpu.CopyUniversalOp() - - # layout: [coarsen, 16] bytes (one e4m3fn per element). - q_fp8_tile = cute.local_tile( - q_fp8[token_id, None, None], - tiler=(self.coarsen, 16), - coord=(head_tile_id, sublane), + +class IndexerQFp8Kernel(VllmCuTeDSLJitKernel["IndexerQFp8Kernel.CompileKey"]): + tb_size = 128 + + @dataclass(frozen=True) + class CompileKey: + head_dim: int + rope_dim: int + num_heads: int + cos_sin_dtype: type[cutlass.Numeric] + coarsen: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_dim + rope_dim = compile_key.rope_dim + nope_dim = head_dim - rope_dim + num_heads = compile_key.num_heads + cos_sin_dtype = compile_key.cos_sin_dtype + # process multiple heads at the same time to armotize RoPE load costs + coarsen = compile_key.coarsen + tb_size = IndexerQFp8Kernel.tb_size + assert num_heads % coarsen == 0 + + # later we will use 32B load = 16 BF16 elems + # thus, head_dim=128 requires 8 threads to handle. + # let's call subwarp = 8 threads. + subwarp_size = head_dim // 16 + threads_per_token = (num_heads // coarsen) * subwarp_size + # Each subwarp owns `coarsen` heads; we use the first `coarsen` + # threads of the subwarp to write the per-head weights using the + # fp8 scale computed in the matching loop iteration. + assert coarsen <= subwarp_size, ( + f"FP8 kernel requires coarsen ({coarsen}) <= " + f"subwarp_size ({subwarp_size}) for the weight-fold step" ) - for i in cutlass.range_constexpr(self.coarsen): - # Reduce amax across the full head_dim: each thread already holds - # the max over its 16 lanes; a width=subwarp_size warp shuffle - # spreads the head-wide max to every lane in the subwarp. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=self.subwarp_size, + @cute.kernel + def device_kernel( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + _tid, + _global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + _num_token_heads, + ) = _load_q_and_rope( + positions, + q, + cos_sin_cache, + num_heads, + coarsen, + tb_size, + subwarp_size, + threads_per_token, + nope_dim, + rope_dim, + cos_sin_dtype, + ) + + cp_op = cute.nvgpu.CopyUniversalOp() + + # layout: [coarsen, 16] bytes (one e4m3fn per element). + q_fp8_tile = cute.local_tile( + q_fp8[token_id, None, None], + tiler=(coarsen, 16), + coord=(head_tile_id, sublane), ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - - # scale = max(amax, eps) / fp8_max, then rounded UP to the next - # power of two. Adding the mantissa mask before shifting out the - # mantissa bumps the exponent whenever s isn't a pure pow2. - fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) - bits = recast_val(fp32_scale, Uint32) - scale_exp = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) - - # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 - fp8_scale_bits = scale_exp << Uint32(23) - fp8_scale = recast_val(fp8_scale_bits, Float32) - # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 - inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) - inv_fp8_scale = recast_val(inv_scale_bits, Float32) - - # Weight fold: weights_out = weights * q_scale * scale_combined. - # All threads in the subwarp share the same fp8_scale after the - # warp_reduce above, so we let thread `sublane == i` write the - # weight for head `head_start + i`. - if in_bounds and sublane == i: - head_id = head_start + i - weights_out[token_id, head_id] = ( - weights[token_id, head_id].to(Float32) * scale * fp8_scale + + for i in cutlass.range_constexpr(coarsen): + # Reduce amax across the full head_dim: each thread already holds + # the max over its 16 lanes; a width=subwarp_size warp shuffle + # spreads the head-wide max to every lane in the subwarp. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=subwarp_size, ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) + + # scale = max(amax, eps) / fp8_max, then rounded UP to the next + # power of two. Adding the mantissa mask before shifting out the + # mantissa bumps the exponent whenever s isn't a pure pow2. + fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) + bits = recast_val(fp32_scale, Uint32) + scale_exp = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) - if in_bounds: - # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s - # (one cp.async-shaped 128-bit store per row). - packed = cute.make_rmem_tensor((4,), Uint32) - for j in cutlass.range_constexpr(4): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) - q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) - packed[j] = cvt.fp32x4_to_fp8x4( - q0 * inv_fp8_scale, - q1 * inv_fp8_scale, - q2 * inv_fp8_scale, - q3 * inv_fp8_scale, + # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 + fp8_scale_bits = scale_exp << Uint32(23) + fp8_scale = recast_val(fp8_scale_bits, Float32) + # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 + inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) + inv_fp8_scale = recast_val(inv_scale_bits, Float32) + + # Weight fold: weights_out = weights * q_scale * scale_combined. + # All threads in the subwarp share the same fp8_scale after the + # warp_reduce above, so we let thread `sublane == i` write the + # weight for head `head_start + i`. + if in_bounds and sublane == i: + head_id = head_start + i + weights_out[token_id, head_id] = ( + weights[token_id, head_id].to(Float32) * scale * fp8_scale ) - dst = q_fp8_tile[i, None] - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) + if in_bounds: + # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s + # (one cp.async-shaped 128-bit store per row). + packed = cute.make_rmem_tensor((4,), Uint32) + for j in cutlass.range_constexpr(4): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) + q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) + packed[j] = cvt.fp32x4_to_fp8x4( + q0 * inv_fp8_scale, + q1 * inv_fp8_scale, + q2 * inv_fp8_scale, + q3 * inv_fp8_scale, + ) + + dst = q_fp8_tile[i, None] + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) + + @cute.jit + def host_entrypoint( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * threads_per_token + grid = (cute.ceil_div(total_threads, tb_size), 1, 1) + device_kernel( + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + cos_sin_dtype: type[cutlass.Numeric], + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + cos_sin_dtype=cos_sin_dtype, + ) - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4 + + hf_config = vllm_config.model_config.hf_config + if hf_config is None or dsa_indexer_uses_fp4(vllm_config): + return [] + + num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0) + head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0) + rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0) + if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: + return [] + + return self._trace_dispatch(self.dispatch)( + head_dim=head_dim, + rope_dim=rope_dim, + num_heads=num_heads, + cos_sin_dtype=Float32, + coarsen=(1, 4), + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_tokens = cute.sym_int() max_pos = cute.sym_int() - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + BFloat16, + (num_tokens, compile_key.num_heads, compile_key.head_dim), + divisibility=16, ) positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), + compile_key.cos_sin_dtype, + (max_pos, compile_key.rope_dim), divisibility=8, ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) + weights = make_fake_tensor( + BFloat16, (num_tokens, compile_key.num_heads), divisibility=8 + ) q_fp8 = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim), + (num_tokens, compile_key.num_heads, compile_key.head_dim), divisibility=16, ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQFp8Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + weights_out = make_fake_tensor( + Float32, (num_tokens, compile_key.num_heads), divisibility=4 ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( positions, q, cos_sin_cache, @@ -609,6 +656,40 @@ def compile( q_fp8, weights_out, Float32(0.0), - stream, - options="--enable-tvm-ffi", ) + + @kernel_launcher + def __call__( + self, + *, + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + weights_softmax_scale: float, + weights_head_scale: float, + q_fp8: torch.Tensor, + weights_out: torch.Tensor, + ) -> CuTeDSLLaunchSpec["IndexerQFp8Kernel.CompileKey"]: + num_tokens, num_heads, head_dim = q.shape + compile_key = self.dispatch( + head_dim=head_dim, + rope_dim=cos_sin_cache.shape[-1], + num_heads=num_heads, + cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + coarsen=1 if num_tokens < 512 else 4, + ) + launch_args = ( + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + float(weights_softmax_scale * weights_head_scale), + ) + return compile_key, launch_args + + +_INDEXER_Q_MXFP4_KERNEL = IndexerQMxFp4Kernel() +_INDEXER_Q_FP8_KERNEL = IndexerQFp8Kernel() diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index 4ff4b232d10f..74866b096089 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -7,22 +7,25 @@ from __future__ import annotations -from functools import cache +from dataclasses import dataclass from typing import Any import cutlass import cutlass.cute as cute import torch from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr +from cutlass import Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op from quack.compile_utils import make_fake_tensor -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} +from vllm.cute_utils import torch_to_cute_dtype +from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, +) +from vllm.utils.math_utils import round_up @dsl_user_op @@ -72,361 +75,406 @@ def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16 return Uint16(out) -class SparseAttnCompressNormRopeStoreC4Kernel: +class SparseAttnCompressNormRopeStoreC4Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreC4Kernel.CompileKey"] +): min_scale = 1.0e-4 rcp_ln2 = 1.4426950408889634 - def __init__( - self, - head_size: int, - state_width: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - overlap: bool, - ): - self.head_dim = head_size - self.state_width = state_width - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.elems_per_lane = 8 - self.copy_elems = 4 - self.copy_chunks = self.elems_per_lane // self.copy_elems - self.lanes_per_group = quant_block // self.elems_per_lane - self.groups_per_warp = 32 // self.lanes_per_group - self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1 - self.scale_reduce_offset = self.lanes_per_group // 2 - self.num_warps = (head_size // quant_block) // self.groups_per_warp - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = self.num_warps * 32 - self.compress_ratio = compress_ratio - self.overlap = overlap - self.window = (1 + int(overlap)) * compress_ratio - - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - group_lane = lane_id % self.lanes_per_group - group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group - elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = ( - slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) - ) + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int + rope_head_dim: int + fp8_max: float + quant_block: int + token_stride: int + scale_dim: int + compress_ratio: int + overlap: bool + norm_weight_dtype: type[cutlass.Numeric] - if active: - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + min_scale = SparseAttnCompressNormRopeStoreC4Kernel.min_scale + rcp_ln2 = SparseAttnCompressNormRopeStoreC4Kernel.rcp_ln2 + head_dim = compile_key.head_size + state_width = compile_key.state_width + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + elems_per_lane = 8 + copy_elems = 4 + copy_chunks = elems_per_lane // copy_elems + lanes_per_group = quant_block // elems_per_lane + groups_per_warp = 32 // lanes_per_group + scale_reduce_steps = lanes_per_group.bit_length() - 1 + scale_reduce_offset = lanes_per_group // 2 + num_warps = (head_dim // quant_block) // groups_per_warp + nope_blocks = nope_dim // quant_block + tb_size = num_warps * 32 + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + window = (1 + int(overlap)) * compress_ratio + + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + group_lane = lane_id % lanes_per_group + group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group + elem_base = group_idx * quant_block + group_lane * elems_per_lane - smem = cutlass.utils.SmemAllocator() - s_block_numbers = smem.allocate_tensor( - Int32, cute.make_layout((self.window,)), byte_alignment=4 - ) - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 + has_req_idx = token_idx < token_to_req_indices.shape[0] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = ( + slot_id >= Int64(0) + and has_req_idx + and boundary + and kv_slot_idx >= Int64(0) ) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if tid == row: - block_number_i32 = Int32(0) - if pos >= Int64(0): - block_index = pos // block_size - block_number_i32 = block_table[req_idx, block_index] - s_block_numbers[row] = block_number_i32 - cute.arch.sync_threads() + if active: + req_idx = token_to_req_indices[token_idx] + start = position - Int64(window - 1) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + smem = cutlass.utils.SmemAllocator() + s_block_numbers = smem.allocate_tensor( + Int32, cute.make_layout((window,)), byte_alignment=4 + ) + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if tid == row: + block_number_i32 = Int32(0) + if pos >= Int64(0): + block_index = pos // block_size + block_number_i32 = block_table[req_idx, block_index] + s_block_numbers[row] = block_number_i32 + cute.arch.sync_threads() + + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (copy_chunks, copy_elems), + stride=(copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 - ) - copy_layout = cute.make_layout( - (self.copy_chunks, self.copy_elems), - stride=(self.copy_elems, 1), - ) - kv_vals = cute.make_rmem_tensor(copy_layout, Float32) - score_vals = cute.make_rmem_tensor(copy_layout, Float32) - - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if pos >= Int64(0): - block_index = pos // block_size - block_offset = pos - block_index * block_size - block_number = s_block_numbers[row].to(Int64) - head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_tensor = state_cache[block_number, block_offset, None] - for chunk in cutlass.range_constexpr(self.copy_chunks): - copy_elem = const_expr(chunk * self.copy_elems) - col_tile = ( - head_offset + (elem_base + Int32(copy_elem)).to(Int64) - ) // Int64(self.copy_elems) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=(col_tile,), - ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=( - col_tile + Int64(self.state_width // self.copy_elems), - ), - ) - cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) - cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - chunk = const_expr(e // self.copy_elems) - copy_elem = const_expr(e % self.copy_elems) - score = score_vals[chunk, copy_elem] - kv = kv_vals[chunk, copy_elem] - new_max = cute.arch.fmax(local_max[e], score) - old_scale = cute.math.exp2( - (local_max[e] - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - new_scale = cute.math.exp2( - (score - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - local_sum[e] = local_sum[e] * old_scale + new_scale - local_product[e] = local_product[e] * old_scale + kv * new_scale - local_max[e] = new_max + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if pos >= Int64(0): + block_index = pos // block_size + block_offset = pos - block_index * block_size + block_number = s_block_numbers[row].to(Int64) + head_offset = Int64((row // compress_ratio) * head_dim) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(copy_chunks): + copy_elem = const_expr(chunk * copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile + Int64(state_width // copy_elems),), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) + + for e in cutlass.range_constexpr(elems_per_lane): + chunk = const_expr(e // copy_elems) + copy_elem = const_expr(e % copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = ( + local_product[e] * old_scale + kv * new_scale + ) + local_max[e] = new_max + + x = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] + + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sumsq = Float32(0.0) - for e in cutlass.range_constexpr(self.elems_per_lane): - x[e] = local_product[e] / local_sum[e] - local_sumsq += x[e] * x[e] + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True + ) + cute.arch.sync_threads() - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + rrms = rrms_shared[0] + for e in cutlass.range_constexpr(elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride) + scale_base = ( + page * k_cache.stride[0] + + kv_cache_block_size * Int64(token_stride) + + kv_offset * Int64(scale_dim) ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - for e in cutlass.range_constexpr(self.elems_per_lane): - elem = elem_base + e - x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) - - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + kv_cache_block_size * Int64(self.token_stride) - + kv_offset * Int64(self.scale_dim) - ) - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = ( - value_base - + Int64(self.nope_dim) - + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64) + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio ) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_absmax = Float32(0.0) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - q0, q1 = _bf16x2_to_fp32(q_packed) - q[elem] = q0 - q[elem + 1] = q1 - local_absmax = cute.arch.fmax( - local_absmax, - cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = ( + value_base + + Int64(nope_dim) + + ((elem_base - nope_dim + Int32(elem)) * 2).to(Int64) + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_absmax = Float32(0.0) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + q0, q1 = _bf16x2_to_fp32(q_packed) + q[elem] = q0 + q[elem + 1] = q1 + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), + ) + absmax = local_absmax + group_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - lanes_per_group) << 8 + | (cute.arch.WARP_SIZE - 1) ) - absmax = local_absmax - group_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.lanes_per_group) << 8 - | (cute.arch.WARP_SIZE - 1) - ) - for step in cutlass.range_constexpr(self.scale_reduce_steps): - offset = const_expr(self.scale_reduce_offset >> step) - absmax = cute.arch.fmax( - absmax, - cute.arch.shuffle_sync_bfly( + for step in cutlass.range_constexpr(scale_reduce_steps): + offset = const_expr(scale_reduce_offset >> step) + absmax = cute.arch.fmax( absmax, - offset=offset, - mask_and_clamp=group_mask_and_clamp, - ), - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - y0 = cutlass.min( - cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), + cute.arch.shuffle_sync_bfly( + absmax, + offset=offset, + mask_and_clamp=group_mask_and_clamp, + ), + ) + scale_raw = cute.arch.fmax( + Float32(min_scale), + absmax / Float32(fp8_max), ) - y1 = cutlass.min( - cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val( + (Uint32(254) - ue8m0) << Uint32(23), Float32 ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if group_lane == 0: - k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8) - if group_idx == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + y0 = cutlass.min( + cute.arch.fmax(q[elem] * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if group_lane == 0: + k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to( + Uint8 ) + if group_idx == 0: + k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0) + + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + ) -> CompileKey: + overlap = compress_ratio == 4 + state_width = head_size * (2 if overlap else 1) + return self.CompileKey( + head_size=head_size, + state_width=state_width, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + token_stride=576, + scale_dim=8, + compress_ratio=compress_ratio, + overlap=overlap, + norm_weight_dtype=norm_weight_dtype, + ) - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 4, - overlap: bool = True, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=4, + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 4 in hf_config.compress_ratios + and cache_config.cache_dtype == "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + norm_weight_dtype = compile_key.norm_weight_dtype if compress_ratio != 4 or not overlap: raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") if head_size != 512: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires head_size=512." + "CuTe DSL C4 fused sparse-attn currently requires head_size=512.", ) if state_width != 2 * head_size: raise ValueError( - "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.", ) if quant_block != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.", ) if token_stride < head_size + rope_head_dim: raise ValueError("token_stride is too small for the packed FP8/BF16 row.") expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 if scale_dim < expected_scale_dim: raise ValueError("scale_dim is too small for the UE8M0 scale row.") - num_positions = cute.sym_int() num_slots = cute.sym_int() num_req_indices = cute.sym_int() @@ -437,7 +485,6 @@ def compile( block_table_width = cute.sym_int() max_pos = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_state_blocks, state_cache_block_size, state_cache_width), @@ -449,15 +496,21 @@ def compile( assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -468,29 +521,12 @@ def compile( k_cache = cute.runtime.make_fake_tensor( Uint8, (num_kv_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - kernel = SparseAttnCompressNormRopeStoreC4Kernel( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, @@ -503,60 +539,34 @@ def compile( k_cache, kv_slot_mapping, Int64(0), - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnCompressNormRopeStoreFullC4Kernel( - SparseAttnCompressNormRopeStoreC4Kernel -): - def __init__( + @kernel_launcher + def __call__( self, - head_size: int, - state_width: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, + *, + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, compress_ratio: int, - overlap: bool, - store_full_fp8: bool = False, - ): - super().__init__( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreC4Kernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, ) - self.store_full_fp8 = store_full_fp8 - - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( + launch_args = ( state_cache, token_to_req_indices, positions, @@ -566,283 +576,395 @@ def __call__( rms_norm_weight, rms_norm_eps, cos_sin_cache, - k_cache, + kv_cache, kv_slot_mapping, - kv_cache_block_size, - fp8_scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - group_lane = lane_id % self.lanes_per_group - group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group - elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = ( - slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) + kv_cache.shape[1], ) + return compile_key, launch_args - if active: - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) - - smem = cutlass.utils.SmemAllocator() - s_block_numbers = smem.allocate_tensor( - Int32, cute.make_layout((self.window,)), byte_alignment=4 - ) - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 - ) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if tid == row: - block_number_i32 = Int32(0) - if pos >= Int64(0): - block_index = pos // block_size - block_number_i32 = block_table[req_idx, block_index] - s_block_numbers[row] = block_number_i32 - cute.arch.sync_threads() +class SparseAttnCompressNormRopeStoreFullC4Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey"] +): + rcp_ln2 = 1.4426950408889634 - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int + rope_head_dim: int + fp8_max: float + quant_block: int + compress_ratio: int + overlap: bool + store_full_fp8: bool + norm_weight_dtype: type[cutlass.Numeric] - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + rcp_ln2 = SparseAttnCompressNormRopeStoreFullC4Kernel.rcp_ln2 + head_dim = compile_key.head_size + state_width = compile_key.state_width + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + elems_per_lane = 8 + copy_elems = 4 + copy_chunks = elems_per_lane // copy_elems + lanes_per_group = quant_block // elems_per_lane + groups_per_warp = 32 // lanes_per_group + num_warps = (head_dim // quant_block) // groups_per_warp + nope_blocks = nope_dim // quant_block + tb_size = num_warps * 32 + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + window = (1 + int(overlap)) * compress_ratio + store_full_fp8 = compile_key.store_full_fp8 + + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + group_lane = lane_id % lanes_per_group + group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group + elem_base = group_idx * quant_block + group_lane * elems_per_lane - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - copy_layout = cute.make_layout( - (self.copy_chunks, self.copy_elems), - stride=(self.copy_elems, 1), + has_req_idx = token_idx < token_to_req_indices.shape[0] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = ( + slot_id >= Int64(0) + and has_req_idx + and boundary + and kv_slot_idx >= Int64(0) ) - kv_vals = cute.make_rmem_tensor(copy_layout, Float32) - score_vals = cute.make_rmem_tensor(copy_layout, Float32) - - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if pos >= Int64(0): - block_index = pos // block_size - block_offset = pos - block_index * block_size - block_number = s_block_numbers[row].to(Int64) - head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_tensor = state_cache[block_number, block_offset, None] - for chunk in cutlass.range_constexpr(self.copy_chunks): - copy_elem = const_expr(chunk * self.copy_elems) - col_tile = ( - head_offset + (elem_base + Int32(copy_elem)).to(Int64) - ) // Int64(self.copy_elems) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=(col_tile,), - ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=( - col_tile + Int64(self.state_width // self.copy_elems), - ), - ) - cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) - cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - chunk = const_expr(e // self.copy_elems) - copy_elem = const_expr(e % self.copy_elems) - score = score_vals[chunk, copy_elem] - kv = kv_vals[chunk, copy_elem] - new_max = cute.arch.fmax(local_max[e], score) - old_scale = cute.math.exp2( - (local_max[e] - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - new_scale = cute.math.exp2( - (score - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - local_sum[e] = local_sum[e] * old_scale + new_scale - local_product[e] = local_product[e] * old_scale + kv * new_scale - local_max[e] = new_max - x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sumsq = Float32(0.0) - for e in cutlass.range_constexpr(self.elems_per_lane): - x[e] = local_product[e] / local_sum[e] - local_sumsq += x[e] * x[e] + if active: + req_idx = token_to_req_indices[token_idx] + start = position - Int64(window - 1) - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True + smem = cutlass.utils.SmemAllocator() + s_block_numbers = smem.allocate_tensor( + Int32, cute.make_layout((window,)), byte_alignment=4 + ) + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 ) - cute.arch.sync_threads() - rrms = rrms_shared[0] - for e in cutlass.range_constexpr(self.elems_per_lane): - elem = elem_base + e - x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if tid == row: + block_number_i32 = Int32(0) + if pos >= Int64(0): + block_index = pos // block_size + block_number_i32 = block_table[req_idx, block_index] + s_block_numbers[row] = block_number_i32 + cute.arch.sync_threads() + + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (copy_chunks, copy_elems), + stride=(copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if pos >= Int64(0): + block_index = pos // block_size + block_offset = pos - block_index * block_size + block_number = s_block_numbers[row].to(Int64) + head_offset = Int64((row // compress_ratio) * head_dim) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(copy_chunks): + copy_elem = const_expr(chunk * copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile + Int64(state_width // copy_elems),), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) + + for e in cutlass.range_constexpr(elems_per_lane): + chunk = const_expr(e // copy_elems) + copy_elem = const_expr(e % copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = ( + local_product[e] * old_scale + kv * new_scale + ) + local_max[e] = new_max + + x = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] + + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - if const_expr(self.store_full_fp8): - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - inv_fp8 = Float32(1.0) / fp8_scale[0] - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed_bf16 = _fp32x2_to_bf16x2(real, imag) - b0, b1 = _bf16x2_to_fp32(packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - b0, b1 = _bf16x2_to_fp32(packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed_bf16 = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + ((elem_base + Int32(elem)) * 2).to( - Int64 + cute.arch.sync_threads() + + rrms = rrms_shared[0] + for e in cutlass.range_constexpr(elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) + + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + + if const_expr(store_full_fp8): + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + inv_fp8 = Float32(1.0) / fp8_scale[0] + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio ) - k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed_bf16 = _fp32x2_to_bf16x2(real, imag) + b0, b1 = _bf16x2_to_fp32(packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + else: + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + b0, b1 = _bf16x2_to_fp32(packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 else: - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - out_base = value_base + ((elem_base + Int32(elem)) * 2).to( - Int64 + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio ) - k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed_bf16 = _fp32x2_to_bf16x2(real, imag) + out_base = value_base + ((elem_base + Int32(elem)) * 2).to( + Int64 + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + else: + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + out_base = value_base + ((elem_base + Int32(elem)) * 2).to( + Int64 + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + fp8_scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + store_full_fp8: bool, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + ) -> CompileKey: + overlap = compress_ratio == 4 + return self.CompileKey( + head_size=head_size, + state_width=head_size * (2 if overlap else 1), + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + compress_ratio=compress_ratio, + overlap=overlap, + store_full_fp8=store_full_fp8, + norm_weight_dtype=norm_weight_dtype, + ) - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 4, - overlap: bool = True, - store_full_fp8: bool = False, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=4, + store_full_fp8=( + cache_config.cache_dtype.startswith("fp8") + and cache_config.cache_dtype != "fp8_ds_mla" + ), + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 4 in hf_config.compress_ratios + and cache_config.cache_dtype != "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + norm_weight_dtype = compile_key.norm_weight_dtype if compress_ratio != 4 or not overlap: raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") if head_size != 512: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires head_size=512." + "CuTe DSL C4 fused sparse-attn currently requires head_size=512.", ) if state_width != 2 * head_size: raise ValueError( - "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.", ) if quant_block != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.", ) num_positions = cute.sym_int() num_slots = cute.sym_int() @@ -854,7 +976,6 @@ def compile( block_table_width = cute.sym_int() max_pos = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_state_blocks, state_cache_block_size, state_cache_width), @@ -866,15 +987,21 @@ def compile( assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -885,31 +1012,12 @@ def compile( k_cache = cute.runtime.make_fake_tensor( Uint8, (num_kv_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1) - - kernel = SparseAttnCompressNormRopeStoreFullC4Kernel( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - store_full_fp8, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, @@ -923,272 +1031,348 @@ def compile( kv_slot_mapping, Int64(0), fp8_scale, - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnCompressC128Block8Kernel: - head_tile = 64 - rows_per_warp = 16 - elems_per_lane = 2 - lanes_per_row = head_tile // elems_per_lane - num_warps = 8 - stats_lane_stride = lanes_per_row + 1 - final_reduce_steps = 3 - final_reduce_initial_offset = 4 - tb_size = num_warps * 32 - compress_ratio = 128 - state_block_size = 8 - rcp_ln2 = 1.4426950408889634 - - def __init__( + @kernel_launcher + def __call__( self, - head_size: int, - state_width: int, - ): - self.head_dim = head_size - self.num_splits = head_size // self.head_tile - self.state_width = state_width - - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - compressed_kv: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0] * self.num_splits, 1, 1) - self.kernel( + *, + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + store_full_fp8: bool, + fp8_scale: torch.Tensor | None, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, + ) + if fp8_scale is None: + fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device) + launch_args = ( state_cache, token_to_req_indices, positions, slot_mapping, block_table, - compressed_kv, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache.view(torch.uint8), + kv_slot_mapping, + kv_cache.shape[1], + fp8_scale, + ) + return compile_key, launch_args - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - compressed_kv: cute.Tensor, - ): - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - col_group = lane_id % self.lanes_per_row - - token_idx = block_id // self.num_splits - split_idx = block_id - token_idx * self.num_splits - col_base = split_idx * self.head_tile + col_group * self.elems_per_lane - - position = Int64(0) - req_idx = Int32(0) - slot_id = Int64(-1) - has_position = token_idx < positions.shape[0] - has_req_idx = token_idx < token_to_req_indices.shape[0] - if lane_id == 0: - slot_id = slot_mapping[token_idx] - if lane_id == 0 and has_position: - position = positions[token_idx] - if lane_id == 0 and has_req_idx: - req_idx = token_to_req_indices[token_idx] - slot_id = cute.arch.shuffle_sync(slot_id, offset=0) - position = cute.arch.shuffle_sync(position, offset=0) - req_idx = cute.arch.shuffle_sync(req_idx, offset=0) - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + +class SparseAttnCompressC128Block8Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressC128Block8Kernel.CompileKey"] +): + head_tile = 64 + rows_per_warp = 16 + elems_per_lane = 2 + lanes_per_row = head_tile // elems_per_lane + num_warps = 8 + stats_lane_stride = lanes_per_row + 1 + final_reduce_steps = 3 + final_reduce_initial_offset = 4 + tb_size = num_warps * 32 + compress_ratio = 128 + state_block_size = 8 + rcp_ln2 = 1.4426950408889634 + + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_tile = SparseAttnCompressC128Block8Kernel.head_tile + rows_per_warp = SparseAttnCompressC128Block8Kernel.rows_per_warp + elems_per_lane = SparseAttnCompressC128Block8Kernel.elems_per_lane + lanes_per_row = SparseAttnCompressC128Block8Kernel.lanes_per_row + num_warps = SparseAttnCompressC128Block8Kernel.num_warps + stats_lane_stride = SparseAttnCompressC128Block8Kernel.stats_lane_stride + final_reduce_steps = SparseAttnCompressC128Block8Kernel.final_reduce_steps + final_reduce_initial_offset = ( + SparseAttnCompressC128Block8Kernel.final_reduce_initial_offset ) - start = position - Int64(self.compress_ratio - 1) - active = slot_id >= Int64(0) and has_req_idx and boundary - - if active: - smem = cutlass.utils.SmemAllocator() - s_max = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, - ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, - ), - ), - byte_alignment=4, + tb_size = SparseAttnCompressC128Block8Kernel.tb_size + compress_ratio = SparseAttnCompressC128Block8Kernel.compress_ratio + state_block_size = SparseAttnCompressC128Block8Kernel.state_block_size + rcp_ln2 = SparseAttnCompressC128Block8Kernel.rcp_ln2 + head_dim = compile_key.head_size + num_splits = head_dim // head_tile + state_width = compile_key.state_width + + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + compressed_kv: cute.Tensor, + ): + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + col_group = lane_id % lanes_per_row + + token_idx = block_id // num_splits + split_idx = block_id - token_idx * num_splits + col_base = split_idx * head_tile + col_group * elems_per_lane + + position = Int64(0) + req_idx = Int32(0) + slot_id = Int64(-1) + has_position = token_idx < positions.shape[0] + has_req_idx = token_idx < token_to_req_indices.shape[0] + if lane_id == 0: + slot_id = slot_mapping[token_idx] + if lane_id == 0 and has_position: + position = positions[token_idx] + if lane_id == 0 and has_req_idx: + req_idx = token_to_req_indices[token_idx] + slot_id = cute.arch.shuffle_sync(slot_id, offset=0) + position = cute.arch.shuffle_sync(position, offset=0) + req_idx = cute.arch.shuffle_sync(req_idx, offset=0) + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - s_sum = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, - ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, + start = position - Int64(compress_ratio - 1) + active = slot_id >= Int64(0) and has_req_idx and boundary + + if active: + smem = cutlass.utils.SmemAllocator() + s_max = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), ), - ), - byte_alignment=4, - ) - s_product = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, + byte_alignment=4, + ) + s_sum = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, + byte_alignment=4, + ) + s_product = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), ), - ), - byte_alignment=4, - ) + byte_alignment=4, + ) - row_layout = cute.make_layout( - (self.rows_per_warp, self.elems_per_lane), - stride=(self.elems_per_lane, 1), - ) - kv_vals = cute.make_rmem_tensor(row_layout, Float32) - score_vals = cute.make_rmem_tensor(row_layout, Float32) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) - - first_block_index = start // Int64(self.state_block_size) - warp_block_index = first_block_index + (warp_id * 2).to(Int64) - block0_i32 = Int32(0) - block1_i32 = Int32(0) - if lane_id == 0: - block0_i32 = block_table[req_idx, warp_block_index] - block1_i32 = block_table[req_idx, warp_block_index + Int64(1)] - block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0) - block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0) + row_layout = cute.make_layout( + (rows_per_warp, elems_per_lane), + stride=(elems_per_lane, 1), + ) + kv_vals = cute.make_rmem_tensor(row_layout, Float32) + score_vals = cute.make_rmem_tensor(row_layout, Float32) + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + first_block_index = start // Int64(state_block_size) + warp_block_index = first_block_index + (warp_id * 2).to(Int64) + block0_i32 = Int32(0) + block1_i32 = Int32(0) + if lane_id == 0: + block0_i32 = block_table[req_idx, warp_block_index] + block1_i32 = block_table[req_idx, warp_block_index + Int64(1)] + block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0) + block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0) - cp_f32x2 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64 - ) - final_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1) - ) - col_tile = col_base.to(Int64) // Int64(self.elems_per_lane) - score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane) - - for i in cutlass.range_constexpr(self.rows_per_warp): - block_number_i32 = block0_i32 - block_offset = Int64(i) - if const_expr(i >= self.state_block_size): - block_number_i32 = block1_i32 - block_offset = Int64(i - self.state_block_size) - row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None] - kv_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=(col_tile,), + cp_f32x2 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64 ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=(score_col_tile,), + final_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1) ) - cute.copy(cp_f32x2, kv_src, kv_vals[i, None]) - cute.copy(cp_f32x2, score_src, score_vals[i, None]) + col_tile = col_base.to(Int64) // Int64(elems_per_lane) + score_col_tile = col_tile + Int64(state_width // elems_per_lane) + + for i in cutlass.range_constexpr(rows_per_warp): + block_number_i32 = block0_i32 + block_offset = Int64(i) + if const_expr(i >= state_block_size): + block_number_i32 = block1_i32 + block_offset = Int64(i - state_block_size) + row_tensor = state_cache[ + block_number_i32.to(Int64), block_offset, None + ] + kv_src = cute.local_tile( + row_tensor, + tiler=(elems_per_lane,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(elems_per_lane,), + coord=(score_col_tile,), + ) + cute.copy(cp_f32x2, kv_src, kv_vals[i, None]) + cute.copy(cp_f32x2, score_src, score_vals[i, None]) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) + for e in cutlass.range_constexpr(elems_per_lane): + for i in cutlass.range_constexpr(rows_per_warp): + exp_score = cute.math.exp2( + (score_vals[i, e] - local_max[e]) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] += exp_score + local_product[e] += kv_vals[i, e] * exp_score + + for e in cutlass.range_constexpr(elems_per_lane): + s_max[warp_id, col_group, e] = local_max[e] + s_sum[warp_id, col_group, e] = local_sum[e] + s_product[warp_id, col_group, e] = local_product[e] + cute.arch.sync_threads() + + out_group = tid // num_warps + final_lane = tid % num_warps + final_groups_per_pass = const_expr(tb_size // num_warps) + for pass_idx in cutlass.range_constexpr( + head_tile // final_groups_per_pass + ): + out_idx = pass_idx * final_groups_per_pass + out_group + out_lane = out_idx // elems_per_lane + out_elem = out_idx % elems_per_lane + + local_warp_max = s_max[final_lane, out_lane, out_elem] + global_max = local_warp_max + for step in cutlass.range_constexpr(final_reduce_steps): + offset = const_expr(final_reduce_initial_offset >> step) + global_max = cute.arch.fmax( + global_max, + cute.arch.shuffle_sync_bfly( + global_max, + offset=offset, + mask_and_clamp=final_mask_and_clamp, + ), + ) - for e in cutlass.range_constexpr(self.elems_per_lane): - for i in cutlass.range_constexpr(self.rows_per_warp): - exp_score = cute.math.exp2( - (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2), + scale = cute.math.exp2( + (local_warp_max - global_max) * Float32(rcp_ln2), fastmath=True, ) - local_sum[e] += exp_score - local_product[e] += kv_vals[i, e] * exp_score - - for e in cutlass.range_constexpr(self.elems_per_lane): - s_max[warp_id, col_group, e] = local_max[e] - s_sum[warp_id, col_group, e] = local_sum[e] - s_product[warp_id, col_group, e] = local_product[e] - cute.arch.sync_threads() - - out_group = tid // self.num_warps - final_lane = tid % self.num_warps - final_groups_per_pass = const_expr(self.tb_size // self.num_warps) - for pass_idx in cutlass.range_constexpr( - self.head_tile // final_groups_per_pass - ): - out_idx = pass_idx * final_groups_per_pass + out_group - out_lane = out_idx // self.elems_per_lane - out_elem = out_idx % self.elems_per_lane - - local_warp_max = s_max[final_lane, out_lane, out_elem] - global_max = local_warp_max - for step in cutlass.range_constexpr(self.final_reduce_steps): - offset = const_expr(self.final_reduce_initial_offset >> step) - global_max = cute.arch.fmax( - global_max, - cute.arch.shuffle_sync_bfly( - global_max, + global_sum = s_sum[final_lane, out_lane, out_elem] * scale + global_product = s_product[final_lane, out_lane, out_elem] * scale + for step in cutlass.range_constexpr(final_reduce_steps): + offset = const_expr(final_reduce_initial_offset >> step) + global_sum += cute.arch.shuffle_sync_bfly( + global_sum, offset=offset, mask_and_clamp=final_mask_and_clamp, - ), - ) - - scale = cute.math.exp2( - (local_warp_max - global_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - global_sum = s_sum[final_lane, out_lane, out_elem] * scale - global_product = s_product[final_lane, out_lane, out_elem] * scale - for step in cutlass.range_constexpr(self.final_reduce_steps): - offset = const_expr(self.final_reduce_initial_offset >> step) - global_sum += cute.arch.shuffle_sync_bfly( - global_sum, - offset=offset, - mask_and_clamp=final_mask_and_clamp, - ) - global_product += cute.arch.shuffle_sync_bfly( - global_product, - offset=offset, - mask_and_clamp=final_mask_and_clamp, - ) + ) + global_product += cute.arch.shuffle_sync_bfly( + global_product, + offset=offset, + mask_and_clamp=final_mask_and_clamp, + ) - if final_lane == 0: - compressed_kv.iterator[ - token_idx.to(Int64) * compressed_kv.stride[0] - + (split_idx * self.head_tile + out_idx).to(Int64) - ] = global_product / global_sum + if final_lane == 0: + compressed_kv.iterator[ + token_idx.to(Int64) * compressed_kv.stride[0] + + (split_idx * head_tile + out_idx).to(Int64) + ] = global_product / global_sum + + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + compressed_kv: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0] * num_splits, 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + compressed_kv, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, *, head_size: int, state_width: int + ) -> CompileKey: + return self.CompileKey(head_size=head_size, state_width=state_width) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + return self._trace_dispatch(self.dispatch)( + head_size=hf_config.head_dim, + state_width=hf_config.head_dim, + enabled=128 in hf_config.compress_ratios, + _when=lambda *, enabled: enabled, + ) - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 512, - ): + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0: raise ValueError("head_size must be divisible by the 64-wide head tile.") num_positions = cute.sym_int() @@ -1198,7 +1382,6 @@ def compile( state_cache_block_size = cute.sym_int() block_table_width = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_blocks, state_cache_block_size, state_cache_width), @@ -1210,12 +1393,16 @@ def compile( assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) compressed_kv = cute.runtime.make_fake_tensor( Float32, @@ -1223,238 +1410,332 @@ def compile( stride=(head_size, 1), assumed_align=4, ) - - kernel = SparseAttnCompressC128Block8Kernel( - head_size, - state_width, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, slot_mapping, block_table, compressed_kv, - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnNormRopeStoreKernel: - min_scale = 1.0e-4 - - def __init__( - self, - head_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - static_kv_cache_block_size: int, - ): - self.head_dim = head_size - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.num_warps = head_size // quant_block - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 - self.compress_ratio = compress_ratio - self.static_kv_cache_block_size = static_kv_cache_block_size - - @cute.jit + @kernel_launcher def __call__( self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - compressed_kv, + *, + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + head_dim: int, + compressed_kv: torch.Tensor | None = None, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressC128Block8Kernel.CompileKey]: + compile_key = self.dispatch(head_size=head_dim, state_width=head_dim) + if compressed_kv is None: + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) + launch_args = ( + state_cache, + token_to_req_indices, positions, slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - elem0 = tid * 2 - - position = Int64(0) - kv_slot_idx = Int64(-1) - has_position = token_idx < positions.shape[0] - slot_id = Int64(-1) - if lane_id == 0: - slot_id = slot_mapping[token_idx] - if lane_id == 0 and has_position: - position = positions[token_idx] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - if lane_id == 0 and has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - slot_id = cute.arch.shuffle_sync(slot_id, offset=0) - position = cute.arch.shuffle_sync(position, offset=0) - kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0) - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + block_table, + compressed_kv, ) - active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) - - if active: - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - static_block_size = Int64(self.static_kv_cache_block_size) - page = kv_slot_idx // static_block_size - kv_offset = kv_slot_idx - page * static_block_size - scale_row_offset = static_block_size * Int64(self.token_stride) - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + scale_row_offset - + kv_offset * Int64(self.scale_dim) - ) - weight0 = rms_norm_weight[elem0].to(Float32) - weight1 = rms_norm_weight[elem0 + 1].to(Float32) - - base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) - x0 = compressed_kv.iterator[base] - x1 = compressed_kv.iterator[base + Int64(1)] - - local_sumsq = x0 * x0 + x1 * x1 - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - smem = cutlass.utils.SmemAllocator() - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) + return compile_key, launch_args, compressed_kv - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - total = partial_sums[lane_id % self.num_warps] - sum_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1) +class SparseAttnNormRopeStoreKernel( + VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreKernel.CompileKey"] +): + min_scale = 1.0e-4 + + @dataclass(frozen=True) + class CompileKey: + head_size: int + rope_head_dim: int + fp8_max: float + quant_block: int + token_stride: int + scale_dim: int + kv_block_stride: int + compress_ratio: int + norm_weight_dtype: type[cutlass.Numeric] + kv_cache_block_size: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + min_scale = SparseAttnNormRopeStoreKernel.min_scale + head_dim = compile_key.head_size + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + num_warps = head_dim // quant_block + nope_blocks = nope_dim // quant_block + tb_size = head_dim // 2 + compress_ratio = compile_key.compress_ratio + static_kv_cache_block_size = compile_key.kv_cache_block_size + + @cute.kernel + def device_kernel( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + elem0 = tid * 2 + + position = Int64(0) + kv_slot_idx = Int64(-1) + has_position = token_idx < positions.shape[0] + slot_id = Int64(-1) + if lane_id == 0: + slot_id = slot_mapping[token_idx] + if lane_id == 0 and has_position: + position = positions[token_idx] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + if lane_id == 0 and has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + slot_id = cute.arch.shuffle_sync(slot_id, offset=0) + position = cute.arch.shuffle_sync(position, offset=0) + kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0) + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - for step in cutlass.range_constexpr(3): - offset = const_expr(4 >> step) - total += cute.arch.shuffle_sync_bfly( - total, - offset, - mask_and_clamp=sum_mask_and_clamp, + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + + if active: + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + static_block_size = Int64(static_kv_cache_block_size) + page = kv_slot_idx // static_block_size + kv_offset = kv_slot_idx - page * static_block_size + scale_row_offset = static_block_size * Int64(token_stride) + value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride) + scale_base = ( + page * k_cache.stride[0] + + scale_row_offset + + kv_offset * Int64(scale_dim) ) + weight0 = rms_norm_weight[elem0].to(Float32) + weight1 = rms_norm_weight[elem0 + 1].to(Float32) - rrms = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - x0 = x0 * rrms * weight0 - x1 = x1 * rrms * weight1 - - if warp_id == self.nope_blocks: - pair_idx = lane_id - compressed_pos = position - Int64(self.compress_ratio - 1) - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - real = x0 * cos_v - x1 * sin_v - imag = x0 * sin_v + x1 * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q_packed = _fp32x2_to_bf16x2(x0, x1) - q0, q1 = _bf16x2_to_fp32(q_packed) - abs0 = cute.math.absf(q0) - abs1 = cute.math.absf(q1) - local_absmax = cute.arch.fmax(abs0, abs1) - absmax = local_absmax + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) + x0 = compressed_kv.iterator[base] + x1 = compressed_kv.iterator[base + Int64(1)] + + local_sumsq = x0 * x0 + x1 * x1 + warp_sum = local_sumsq for step in cutlass.range_constexpr(5): offset = const_expr(16 >> step) - absmax = cute.arch.fmax( - absmax, cute.arch.shuffle_sync_bfly(absmax, offset) - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - y0 = cutlass.min( - cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + + smem = cutlass.utils.SmemAllocator() + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 ) - y1 = cutlass.min( - cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + + total = partial_sums[lane_id % num_warps] + sum_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1) ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to( - Int64 + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) + total += cute.arch.shuffle_sync_bfly( + total, + offset, + mask_and_clamp=sum_mask_and_clamp, + ) + + rrms = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True ) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if lane_id == 0: - k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8) - if warp_id == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 + x0 = x0 * rrms * weight0 + x1 = x1 * rrms * weight1 + + if warp_id == nope_blocks: + pair_idx = lane_id + compressed_pos = position - Int64(compress_ratio - 1) + cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( + Int64 + ) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + real = x0 * cos_v - x1 * sin_v + imag = x0 * sin_v + x1 * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = value_base + Int64(nope_dim) + (lane_id * 4).to(Int64) + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q_packed = _fp32x2_to_bf16x2(x0, x1) + q0, q1 = _bf16x2_to_fp32(q_packed) + abs0 = cute.math.absf(q0) + abs1 = cute.math.absf(q1) + local_absmax = cute.arch.fmax(abs0, abs1) + absmax = local_absmax + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + absmax = cute.arch.fmax( + absmax, cute.arch.shuffle_sync_bfly(absmax, offset) ) + scale_raw = cute.arch.fmax( + Float32(min_scale), + absmax / Float32(fp8_max), + ) + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val( + (Uint32(254) - ue8m0) << Uint32(23), Float32 + ) + y0 = cutlass.min( + cute.arch.fmax(q0 * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q1 * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (warp_id * quant_block + lane_id * 2).to( + Int64 + ) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if lane_id == 0: + k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to( + Uint8 + ) + if warp_id == 0: + k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0) + + @cute.jit + def host_entrypoint( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + cache_block_size: int, + cache_alignment: int, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + runtime_kv_block_stride: int | None = None, + ) -> CompileKey: + raw_kv_cache_block_size = cache_block_size // compress_ratio + kv_cache_block_size = ( + raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1 + ) + token_stride = 576 + scale_dim = 8 + return self.CompileKey( + head_size=head_size, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + token_stride=token_stride, + scale_dim=scale_dim, + kv_block_stride=( + runtime_kv_block_stride + if runtime_kv_block_stride is not None + else round_up( + kv_cache_block_size * (token_stride + scale_dim), + cache_alignment, + ) + ), + compress_ratio=compress_ratio, + norm_weight_dtype=norm_weight_dtype, + kv_cache_block_size=kv_cache_block_size, + ) - @cache - @staticmethod - def compile( - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 128, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - static_kv_cache_block_size: int = 0, - ): + def get_warmup_keys( + self, + vllm_config: Any, + *, + k_cache_prefix: str, + compress_ratio: int, + ) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + k_cache_layer = vllm_config.compilation_config.static_forward_context[ + k_cache_prefix + ] + kv_cache = k_cache_layer.kv_cache + return self._trace_dispatch(self.dispatch)( + compress_ratio=compress_ratio, + cache_block_size=cache_config.block_size, + cache_alignment=576, + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + runtime_kv_block_stride=kv_cache.stride(0), + enabled=( + cache_config.block_size > 0 + and compress_ratio == 128 + and kv_cache.ndim == 3 + and cache_config.cache_dtype == "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + kv_block_stride = compile_key.kv_block_stride + norm_weight_dtype = compile_key.norm_weight_dtype + static_kv_cache_block_size = compile_key.kv_cache_block_size if quant_block != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires quant_block=64." + "CuTe DSL sparse-attn store currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires rope_head_dim=64." + "CuTe DSL sparse-attn store currently requires rope_head_dim=64.", ) if head_size % quant_block != 0: raise ValueError("head_size must be divisible by quant_block.") @@ -1466,14 +1747,13 @@ def compile( if static_kv_cache_block_size <= 0: raise ValueError( "CuTe DSL sparse-attn store requires a positive static " - "kv_cache_block_size." + "kv_cache_block_size.", ) num_positions = cute.sym_int() num_slots = cute.sym_int() num_kv_slots = cute.sym_int() max_pos = cute.sym_int() num_blocks = cute.sym_int() - compressed_kv = cute.runtime.make_fake_tensor( Float32, (num_slots, head_size), @@ -1483,7 +1763,9 @@ def compile( positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -1494,28 +1776,11 @@ def compile( k_cache = cute.runtime.make_fake_tensor( Uint8, (num_blocks, cute.sym_int(), cute.sym_int()), - stride=( - kv_block_stride, - cute.sym_int64(divisibility=8), - 1, - ), + stride=(kv_block_stride, cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - - kernel = SparseAttnNormRopeStoreKernel( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - static_kv_cache_block_size, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( compressed_kv, positions, slot_mapping, @@ -1524,211 +1789,275 @@ def compile( cos_sin_cache, k_cache, kv_slot_mapping, - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnNormRopeStoreFullKernel: - def __init__( - self, - head_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - store_full_fp8: bool = False, - ): - # Standalone (not inheriting the #44230-restructured legacy kernel): - # set attrs directly so the full-cache C128 path is decoupled. - self.head_dim = head_size - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.num_warps = head_size // quant_block - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 - self.compress_ratio = compress_ratio - self.store_full_fp8 = store_full_fp8 - - @cute.jit + @kernel_launcher def __call__( self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( + *, + compressed_kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreKernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + cache_block_size=kv_cache.shape[1] * compress_ratio, + cache_alignment=576, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, + runtime_kv_block_stride=kv_cache.stride(0), + ) + launch_args = ( compressed_kv, positions, slot_mapping, rms_norm_weight, rms_norm_eps, cos_sin_cache, - k_cache, + kv_cache, kv_slot_mapping, - kv_cache_block_size, - fp8_scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - elem0 = tid * 2 - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) ) - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) - - if active: - base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) - x0 = compressed_kv.iterator[base] - x1 = compressed_kv.iterator[base + Int64(1)] - - local_sumsq = x0 * x0 + x1 * x1 - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - smem = cutlass.utils.SmemAllocator() - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 - ) + return compile_key, launch_args - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - rrms = rrms_shared[0] - x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) - x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) +class SparseAttnNormRopeStoreFullKernel( + VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreFullKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + head_size: int + rope_head_dim: int + fp8_max: float + quant_block: int + compress_ratio: int + store_full_fp8: bool + norm_weight_dtype: type[cutlass.Numeric] - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_size + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + num_warps = head_dim // quant_block + nope_blocks = nope_dim // quant_block + tb_size = head_dim // 2 + compress_ratio = compile_key.compress_ratio + store_full_fp8 = compile_key.store_full_fp8 + + @cute.kernel + def device_kernel( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + elem0 = tid * 2 - if const_expr(self.store_full_fp8): - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - inv_fp8 = Float32(1.0) / fp8_scale[0] - fp8_v0 = x0 - fp8_v1 = x1 - if warp_id == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - pair_idx = lane_id - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( - Int64 - ) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - fp8_v0 = x0 * cos_v - x1 * sin_v - fp8_v1 = x0 * sin_v + x1 * cos_v - fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1) - b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) + ) + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + + if active: + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) + x0 = compressed_kv.iterator[base] + x1 = compressed_kv.iterator[base + Int64(1)] + + local_sumsq = x0 * x0 + x1 * x1 + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + + smem = cutlass.utils.SmemAllocator() + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + elem0.to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - bf16_v0 = x0 - bf16_v1 = x1 - if warp_id == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True ) - pair_idx = lane_id - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( - Int64 + cute.arch.sync_threads() + + rrms = rrms_shared[0] + x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) + x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) + + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + + if const_expr(store_full_fp8): + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + inv_fp8 = Float32(1.0) / fp8_scale[0] + fp8_v0 = x0 + fp8_v1 = x1 + if warp_id == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + pair_idx = lane_id + cs_base = compressed_pos * cos_sin_cache.stride[ + 0 + ] + pair_idx.to(Int64) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + fp8_v0 = x0 * cos_v - x1 * sin_v + fp8_v1 = x0 * sin_v + x1 * cos_v + fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1) + b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), ) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - bf16_v0 = x0 * cos_v - x1 * sin_v - bf16_v1 = x0 * sin_v + x1 * cos_v - bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1) - out_base = value_base + (elem0 * 2).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed - - @cache - @staticmethod - def compile( - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 128, - store_full_fp8: bool = False, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + elem0.to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + else: + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + bf16_v0 = x0 + bf16_v1 = x1 + if warp_id == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + pair_idx = lane_id + cs_base = compressed_pos * cos_sin_cache.stride[ + 0 + ] + pair_idx.to(Int64) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + bf16_v0 = x0 * cos_v - x1 * sin_v + bf16_v1 = x0 * sin_v + x1 * cos_v + bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1) + out_base = value_base + (elem0 * 2).to(Int64) + k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed + + @cute.jit + def host_entrypoint( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + fp8_scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + store_full_fp8: bool, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + compress_ratio: int, + ) -> CompileKey: + return self.CompileKey( + head_size=head_size, + rope_head_dim=rope_head_dim, + compress_ratio=compress_ratio, + fp8_max=448.0, + quant_block=64, + store_full_fp8=store_full_fp8, + norm_weight_dtype=norm_weight_dtype, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=128, + store_full_fp8=( + cache_config.cache_dtype.startswith("fp8") + and cache_config.cache_dtype != "fp8_ds_mla" + ), + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 128 in hf_config.compress_ratios + and cache_config.cache_dtype != "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + norm_weight_dtype = compile_key.norm_weight_dtype if quant_block != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires quant_block=64." + "CuTe DSL sparse-attn store currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires rope_head_dim=64." + "CuTe DSL sparse-attn store currently requires rope_head_dim=64.", ) if head_size % quant_block != 0: raise ValueError("head_size must be divisible by quant_block.") @@ -1737,7 +2066,6 @@ def compile( num_kv_slots = cute.sym_int() max_pos = cute.sym_int() num_blocks = cute.sym_int() - compressed_kv = cute.runtime.make_fake_tensor( Float32, (num_slots, head_size), @@ -1747,7 +2075,9 @@ def compile( positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -1758,29 +2088,12 @@ def compile( k_cache = cute.runtime.make_fake_tensor( Uint8, (num_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1) - - kernel = SparseAttnNormRopeStoreFullKernel( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - store_full_fp8, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( compressed_kv, positions, slot_mapping, @@ -1791,83 +2104,57 @@ def compile( kv_slot_mapping, Int64(0), fp8_scale, - stream, - options="--enable-tvm-ffi", ) - -def compile_split_sparse_attn_cutedsl( - head_size: int, - state_width: int, - block_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - kv_cache_block_size: int, - kv_block_stride: int, - compress_ratio: int, - overlap: bool, - rms_norm_weight_dtype: torch.dtype, - store_full_kv: bool = False, - store_full_fp8: bool = False, -): - if not ( - head_size == 512 - and state_width == head_size - and compress_ratio == 128 - and not overlap - and block_size == 8 - ): - raise ValueError( - "CuTe DSL split sparse-attn wrapper only supports the real " - "DeepSeek V4 C128 layout: head_size=512, state_width=512, " - "compress_ratio=128, overlap=False, block_size=8." - ) - compress = SparseAttnCompressC128Block8Kernel.compile( - head_size=head_size, - state_width=state_width, - ) - norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype] - if store_full_kv: - # FlashInfer contiguous bf16/fp8 cache: standalone full-cache store. - store = SparseAttnNormRopeStoreFullKernel.compile( - head_size=head_size, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + @kernel_launcher + def __call__( + self, + *, + compressed_kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + store_full_fp8: bool, + fp8_scale: torch.Tensor | None, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreFullKernel.CompileKey]: + compile_key = self.dispatch( compress_ratio=compress_ratio, store_full_fp8=store_full_fp8, - norm_weight_dtype=norm_weight_dtype, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, ) - else: - store = SparseAttnNormRopeStoreKernel.compile( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - kv_block_stride, - compress_ratio, - norm_weight_dtype, - kv_cache_block_size, + if fp8_scale is None: + fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device) + launch_args = ( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache.view(torch.uint8), + kv_slot_mapping, + kv_cache.shape[1], + fp8_scale, ) - return compress, store + return compile_key, launch_args -def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( +def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( state_cache: torch.Tensor, token_to_req_indices: torch.Tensor, positions: torch.Tensor, slot_mapping: torch.Tensor, block_table: torch.Tensor, block_size: int, - compressed_kv: torch.Tensor, rms_norm_weight: torch.Tensor, rms_norm_eps: float, cos_sin_cache: torch.Tensor, @@ -1876,102 +2163,92 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( kv_cache_block_size: int, kv_block_stride: int, head_size: int = 512, - state_width: int = 512, + state_width: int = 1024, rope_head_dim: int = 64, fp8_max: float = 448.0, quant_block: int = 64, token_stride: int = 576, scale_dim: int = 8, - compress_ratio: int = 128, - overlap: bool = False, + compress_ratio: int = 4, + overlap: bool = True, store_full_kv: bool = False, store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, ) -> None: + del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride + if not ( + head_size == 512 + and state_width == 2 * head_size + and compress_ratio == 4 + and overlap + ): + raise ValueError( + "CuTe DSL fused sparse-attn wrapper only supports the real " + "DeepSeek V4 C4 layout: head_size=512, state_width=1024, " + "compress_ratio=4, overlap=True." + ) if k_cache.ndim != 3: raise ValueError( - "CuTe DSL sparse-attn store expects the real DeepSeek V4 " + "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." ) - if not store_full_kv and kv_cache_block_size != k_cache.shape[1]: + if kv_cache_block_size != k_cache.shape[1]: raise ValueError( - "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size " + "CuTe DSL fused sparse-attn wrapper expected kv_cache_block_size " f"to match k_cache.shape[1], got {kv_cache_block_size} and " f"{k_cache.shape[1]}." ) if positions.numel() == 0: return - if rms_norm_weight.dtype not in _TORCH_TO_CUTE: - raise ValueError( - "CuTe DSL sparse-attn store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." - ) if store_full_fp8 and not store_full_kv: raise ValueError("store_full_fp8 requires store_full_kv.") - compress, store = compile_split_sparse_attn_cutedsl( - head_size, - state_width, - block_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - kv_cache_block_size, - kv_block_stride, - compress_ratio, - overlap, - rms_norm_weight.dtype, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - ) - compress( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - compressed_kv, - ) if store_full_kv: - # Byte-addressed contiguous cache; block size + per-tensor scale are - # passed at call time (not baked into compile). - if fp8_scale is None: - fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device) - store( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache.view(torch.uint8), - kv_slot_mapping, - kv_cache_block_size, - fp8_scale, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) return - - store( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, + compress_ratio=compress_ratio, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) -def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( +def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( state_cache: torch.Tensor, token_to_req_indices: torch.Tensor, positions: torch.Tensor, slot_mapping: torch.Tensor, block_table: torch.Tensor, block_size: int, + compressed_kv: torch.Tensor, rms_norm_weight: torch.Tensor, rms_norm_eps: float, cos_sin_cache: torch.Tensor, @@ -1980,185 +2257,247 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( kv_cache_block_size: int, kv_block_stride: int, head_size: int = 512, - state_width: int = 1024, + state_width: int = 512, rope_head_dim: int = 64, fp8_max: float = 448.0, quant_block: int = 64, token_stride: int = 576, scale_dim: int = 8, - compress_ratio: int = 4, - overlap: bool = True, + compress_ratio: int = 128, + overlap: bool = False, store_full_kv: bool = False, store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, ) -> None: - if positions.numel() == 0: - return - norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) - if norm_weight_dtype is None: + del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride + if not ( + head_size == 512 + and state_width == head_size + and compress_ratio == 128 + and not overlap + and block_size == 8 + ): raise ValueError( - "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." + "CuTe DSL split sparse-attn wrapper only supports the real " + "DeepSeek V4 C128 layout: head_size=512, state_width=512, " + "compress_ratio=128, overlap=False, block_size=8." ) if k_cache.ndim != 3: raise ValueError( - "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " + "CuTe DSL sparse-attn store expects the real DeepSeek V4 " f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." ) + if not store_full_kv and kv_cache_block_size != k_cache.shape[1]: + raise ValueError( + "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size " + f"to match k_cache.shape[1], got {kv_cache_block_size} and " + f"{k_cache.shape[1]}." + ) + if positions.numel() == 0: + return if store_full_fp8 and not store_full_kv: raise ValueError("store_full_fp8 requires store_full_kv.") + + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL( + state_cache=state_cache, + num_actual=compressed_kv.shape[0], + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + head_dim=head_size, + compressed_kv=compressed_kv, + ) if store_full_kv: - # FlashInfer contiguous bf16/fp8 cache: byte-addressed full-cache C4 store. - if fp8_scale is None: - fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device) - compiled = SparseAttnCompressNormRopeStoreFullC4Kernel.compile( - head_size=head_size, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, store_full_fp8=store_full_fp8, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache.view(torch.uint8), - kv_slot_mapping, - kv_cache_block_size, - fp8_scale, + fp8_scale=fp8_scale, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) return - - compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile( - head_size=head_size, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) -def compress_norm_rope_store_cutedsl( - state_cache: torch.Tensor, - num_actual: int, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - state_width: int, - cos_sin_cache: torch.Tensor, - kv_cache: torch.Tensor, - k_cache_metadata: Any, - pdl_kwargs: dict, - head_dim: int, - rope_head_dim: int, - compress_ratio: int, - overlap: bool, - use_fp4_cache: bool, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - quant_block: int, - token_stride: int, - scale_dim: int, - store_full_kv: bool = False, - store_full_fp8: bool = False, - fp8_scale: torch.Tensor | None = None, -) -> None: - if compress_ratio == 4: - # For C4A, the single fused kernel is faster than the two-kernel version. - fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - compress_ratio=compress_ratio, - overlap=overlap, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - fp8_scale=fp8_scale, - ) - else: - # For C128, the two-kernel version is faster than the single fused kernel. - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, +class SparseAttnCompressorCuteDSL: + def __init__( + self, + *, + c4_store: SparseAttnCompressNormRopeStoreC4Kernel, + c4_store_full: SparseAttnCompressNormRopeStoreFullC4Kernel, + c128_compress: SparseAttnCompressC128Block8Kernel, + c128_store: SparseAttnNormRopeStoreKernel, + c128_store_full: SparseAttnNormRopeStoreFullKernel, + ) -> None: + self.c4_store = c4_store + self.c4_store_full = c4_store_full + self.c128_compress = c128_compress + self.c128_store = c128_store + self.c128_store_full = c128_store_full + + def __call__( + self, + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + store_full_kv: bool = False, + store_full_fp8: bool = False, + fp8_scale: torch.Tensor | None = None, + ) -> None: + del pdl_kwargs, state_width, overlap, use_fp4_cache + del quant_block, token_stride, scale_dim + if positions.numel() == 0: + return + if kv_cache.ndim != 3: + raise ValueError( + "CuTe DSL sparse-attn compressor expects the real DeepSeek V4 " + f"3D kv_cache layout, got ndim={kv_cache.ndim}.", + ) + if store_full_fp8 and (not store_full_kv): + raise ValueError("store_full_fp8 requires store_full_kv.") + if compress_ratio == 4: + if store_full_kv: + self.c4_store_full( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + else: + self.c4_store( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + return + if compress_ratio != 128: + raise ValueError( + "CuTe DSL sparse-attn compressor supports compress_ratio 4 or " + f"128, got {compress_ratio}.", + ) + if block_size != self.c128_compress.state_block_size: + raise ValueError( + "CuTe DSL C128 sparse-attn compressor requires state-cache " + f"block_size={self.c128_compress.state_block_size}, got " + f"{block_size}." + ) + compressed_kv = self.c128_compress( + state_cache=state_cache, + num_actual=num_actual, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + head_dim=head_dim, ) - split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - compressed_kv, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, + if store_full_kv: + self.c128_store_full( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + return + self.c128_store( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - fp8_scale=fp8_scale, + head_dim=head_dim, + rope_head_dim=rope_head_dim, ) + + +_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL = ( + SparseAttnCompressNormRopeStoreC4Kernel() +) +_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL = ( + SparseAttnCompressNormRopeStoreFullC4Kernel() +) +_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL = SparseAttnCompressC128Block8Kernel() +_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL = SparseAttnNormRopeStoreKernel() +_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL = SparseAttnNormRopeStoreFullKernel() +_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL = SparseAttnCompressorCuteDSL( + c4_store=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL, + c4_store_full=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL, + c128_compress=_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL, + c128_store=_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL, + c128_store_full=_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL, +) diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 81659f342a14..4953bdc9256f 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -9,11 +9,11 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability diff --git a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py b/vllm/models/deepseek_v4_1/common/ops/cache_utils.py index e2da6f5e2bed..64462eb73b4a 100644 --- a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4_1/common/ops/cache_utils.py @@ -416,11 +416,17 @@ def dequantize_and_gather_k_cache( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( - dequantize_and_gather_k_cache_cutedsl, + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) - dequantize_and_gather_k_cache_cutedsl( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( + out=out, + k_cache=k_cache, + seq_lens=seq_lens, + gather_lens=gather_lens, + block_table=block_table, + block_size=block_size, + offset=offset, ) return diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 78c30c0f5f2c..b32af7234b0c 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -5,12 +5,11 @@ import torch -from vllm.model_executor.warmup.jit_warmup import zip_inputs +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 0eb642bd7632..59fb109daccb 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -9,12 +9,12 @@ from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonPointerInputVariant, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.platforms import current_platform diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 02e6336a3835..f47554d81f0c 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -9,13 +9,13 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 329d48be2141..8d196f963064 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -10,13 +10,13 @@ from vllm.config import VllmConfig from vllm.distributed import get_dcp_group from vllm.model_executor.warmup.jit_warmup import ( + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 88ce37122626..b5618170556b 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -5,11 +5,11 @@ import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 66a6bace46e0..bd20c9020d3e 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -18,12 +18,12 @@ from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index e8e6f808534f..33b52652a154 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -11,11 +11,11 @@ from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton From d43bb2f37f87a63d3a9a299c97050443cf6a2520 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Sat, 12 Sep 2026 00:36:47 -0400 Subject: [PATCH 082/774] [3/N] HiSparse: host-resident sparse-MLA decode hot-buffering (#53781) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Signed-off-by: Nicolò Lucchesi Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: Nicolò Lucchesi Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Lucas Wilkinson --- CMakeLists.txt | 1 + csrc/libtorch_stable/cache_kernels.cu | 75 +- csrc/libtorch_stable/hisparse_kernels.cu | 1226 ++++++++++ csrc/libtorch_stable/ops.h | 53 +- csrc/libtorch_stable/torch_bindings.cpp | 63 +- docs/design/hisparse.md | 244 ++ tests/evals/gsm8k/gsm8k_eval.py | 7 + ...test_mla_cross_layer_kernel_equivalence.py | 12 +- tests/kernels/test_cp_gather_fp8.py | 42 + .../test_fused_deepseek_v32_norm_rope.py | 70 + .../layers/test_mla_short_prefill_indexer.py | 50 + .../test_flashinfer_autotune_warmup.py | 15 +- tests/test_config.py | 68 +- tests/v1/attention/test_backend_per_kind.py | 56 +- tests/v1/attention/test_mla_backends.py | 1 + tests/v1/attention/test_mla_context_chunks.py | 23 + .../v1/attention/test_sparse_mla_backends.py | 2107 +++++++++++++++-- .../test_partial_prefix_cache_hits.py | 73 +- tests/v1/core/test_async_scheduler.py | 28 +- tests/v1/core/test_contiguous_kv_packing.py | 1 + tests/v1/core/test_kv_cache_utils.py | 144 ++ tests/v1/core/test_prefix_caching.py | 1097 ++++++++- tests/v1/core/test_scheduler.py | 69 + tests/v1/e2e/general/test_hisparse.py | 225 ++ .../unit/offloading_connector/test_config.py | 32 +- .../offloading_connector/test_scheduler.py | 47 +- .../unit/test_hisparse_connector.py | 116 + .../unit/test_mooncake_store_worker.py | 85 + .../kv_connector/unit/test_multi_connector.py | 12 +- .../kv_connector/unit/test_nixl_connector.py | 178 +- .../unit/test_nixl_connector_hma.py | 216 +- .../unit/test_nixl_desc_geometry.py | 27 +- .../unit/test_nixl_push_connector.py | 76 +- .../unit/test_output_aggregator.py | 12 +- tests/v1/kv_connector/unit/utils.py | 1 + tests/v1/kv_offload/test_factory.py | 6 +- tests/v1/kv_offload/test_file_mapper.py | 5 +- tests/v1/simple_kv_offload/test_worker.py | 6 + tests/v1/test_kv_cache_spec_registry.py | 15 + tests/v1/worker/test_attn_utils.py | 127 +- .../test_dsv4_packed_zeroer_geometry.py | 1 + .../test_gpu_autoregressive_speculator.py | 2 + tests/v1/worker/test_gpu_kv_connector.py | 26 +- tests/v1/worker/test_gpu_model_runner.py | 4 + tests/v1/worker/test_gpu_model_runner_v2.py | 1 + tests/v1/worker/test_gpu_warmup_blocks.py | 1 + .../worker/test_kv_cache_allocation_scope.py | 34 +- tests/v1/worker/test_mixed_warmup_gate.py | 8 +- tests/v1/worker/test_utils.py | 894 ++++++- vllm/_custom_ops.py | 16 +- .../fusion/mla_rope_kvcache_cat_fusion.py | 5 +- vllm/config/__init__.py | 3 +- vllm/config/attention.py | 25 +- vllm/config/kv_transfer.py | 34 + vllm/config/speculative.py | 7 + vllm/config/vllm.py | 51 + .../kv_transfer/kv_connector/factory.py | 6 + .../kv_transfer/kv_connector/utils.py | 6 + .../kv_transfer/kv_connector/v1/base.py | 55 +- .../kv_connector/v1/hisparse/__init__.py | 2 + .../kv_connector/v1/hisparse/connector.py | 329 +++ .../kv_connector/v1/hisparse/worker.py | 832 +++++++ .../kv_connector/v1/lmcache_mp_connector.py | 6 +- .../v1/mooncake/mooncake_connector.py | 32 +- .../v1/mooncake/store/coordinator.py | 17 +- .../kv_connector/v1/mooncake/store/worker.py | 151 +- .../kv_connector/v1/multi_connector.py | 66 +- .../kv_connector/v1/nixl/base_scheduler.py | 5 +- .../kv_connector/v1/nixl/base_worker.py | 352 +-- .../kv_connector/v1/nixl/connector.py | 13 + .../kv_connector/v1/nixl/metadata.py | 4 +- .../kv_connector/v1/nixl/pull_worker.py | 28 +- .../kv_connector/v1/nixl/push_worker.py | 23 +- .../kv_connector/v1/offloading/config.py | 31 +- .../kv_connector/v1/offloading/scheduler.py | 103 +- .../kv_connector/v1/offloading/worker.py | 24 +- .../kv_connector/v1/offloading_connector.py | 47 +- .../kv_transfer/kv_transfer_state.py | 8 +- .../layers/attention/mla_attention.py | 131 +- .../layers/attention/sparse_mla_attention.py | 191 +- vllm/model_executor/layers/mla.py | 10 + .../layers/sparse_attn_indexer.py | 2 +- vllm/model_executor/models/deepseek_v2.py | 25 +- vllm/model_executor/models/glm4_moe_lite.py | 20 + .../models/glm4_moe_lite_mtp.py | 13 + .../warmup/flashinfer_sparse_mla_warmup.py | 29 + vllm/model_executor/warmup/kernel_warmup.py | 16 +- vllm/models/deepseek_v32/attention.py | 39 +- vllm/models/deepseek_v32/common/kernels.py | 19 +- vllm/models/deepseek_v32/nvidia/model.py | 11 + vllm/models/deepseek_v32/nvidia/mtp.py | 9 + vllm/models/minimax_m3/common/indexer.py | 2 + vllm/v1/attention/backend.py | 5 + .../backends/mla/flashattn_mla_sparse.py | 127 +- .../backends/mla/flashinfer_mla_sparse.py | 195 +- .../mla/flashinfer_mla_sparse_sm120.py | 131 +- .../attention/backends/mla/flashmla_sparse.py | 316 ++- vllm/v1/attention/backends/mla/index_group.py | 503 ++++ .../v1/attention/backends/mla/sparse_utils.py | 34 +- vllm/v1/core/block_pool.py | 56 +- vllm/v1/core/kv_cache_coordinator.py | 47 +- vllm/v1/core/kv_cache_manager.py | 33 +- vllm/v1/core/kv_cache_utils.py | 84 +- vllm/v1/core/sched/output.py | 3 + vllm/v1/core/sched/scheduler.py | 40 +- vllm/v1/core/single_type_kv_cache_manager.py | 408 +++- vllm/v1/hisparse/__init__.py | 2 + vllm/v1/hisparse/binding.py | 277 +++ vllm/v1/hisparse/block_pool.py | 27 + vllm/v1/hisparse/coordinator.py | 752 ++++++ vllm/v1/hisparse/layout.py | 345 +++ vllm/v1/hisparse/runtime.py | 1100 +++++++++ vllm/v1/hisparse/types.py | 30 + vllm/v1/kv_cache_interface.py | 157 +- vllm/v1/kv_cache_spec_registry.py | 17 +- vllm/v1/kv_offload/config.py | 2 + vllm/v1/outputs.py | 4 + vllm/v1/worker/gpu/attn_utils.py | 47 +- vllm/v1/worker/gpu/cudagraph_utils.py | 2 + vllm/v1/worker/gpu/kv_connector.py | 46 +- vllm/v1/worker/gpu/model_runner.py | 42 +- vllm/v1/worker/gpu/warmup.py | 10 + .../worker/kv_connector_model_runner_mixin.py | 7 +- vllm/v1/worker/utils.py | 61 +- 124 files changed, 14373 insertions(+), 1017 deletions(-) create mode 100644 csrc/libtorch_stable/hisparse_kernels.cu create mode 100644 docs/design/hisparse.md create mode 100644 tests/v1/e2e/general/test_hisparse.py create mode 100644 tests/v1/kv_connector/unit/test_hisparse_connector.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py create mode 100644 vllm/v1/attention/backends/mla/index_group.py create mode 100644 vllm/v1/hisparse/__init__.py create mode 100644 vllm/v1/hisparse/binding.py create mode 100644 vllm/v1/hisparse/block_pool.py create mode 100644 vllm/v1/hisparse/coordinator.py create mode 100644 vllm/v1/hisparse/layout.py create mode 100644 vllm/v1/hisparse/runtime.py create mode 100644 vllm/v1/hisparse/types.py diff --git a/CMakeLists.txt b/CMakeLists.txt index a7d40ae58603..018138f54cb4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -476,6 +476,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/hisparse_kernels.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index f42d61a4f440..887c83f3af87 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1458,7 +1458,11 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page( const int32_t block_size, const int32_t total_tokens, const int64_t block_table_stride, const int64_t cache_block_stride, const int64_t cache_entry_stride, const int64_t dst_entry_stride, - const int32_t* __restrict__ seq_starts) { + const int32_t* __restrict__ seq_starts, + const uint8_t* __restrict__ host_cache, + const int32_t* __restrict__ host_row_ids, + const int32_t* __restrict__ device_row_ids, + const int64_t host_entry_stride) { constexpr int32_t warps_per_cta = 16; __shared__ GatherPageTask page; __shared__ int32_t physical_block; @@ -1484,8 +1488,22 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page( page_token < page.page_token_end; page_token += warps_per_cta) { const int32_t output_token = page.output_token_begin + page_token - page.page_token_begin; - const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + - page_token * cache_entry_stride; + const int32_t source_row = physical_block * block_size + page_token; + const uint8_t* token_ptr; + if (device_row_ids == nullptr) { + token_ptr = src_cache + physical_block * cache_block_stride + + page_token * cache_entry_stride; + } else { + const int32_t device_row = device_row_ids[source_row]; + if (device_row >= 0) { + token_ptr = src_cache + (device_row / block_size) * cache_block_stride + + (device_row % block_size) * cache_entry_stride; + } else { + token_ptr = + host_cache + + static_cast(host_row_ids[source_row]) * host_entry_stride; + } + } __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride; gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id); } @@ -1694,7 +1712,10 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size, - std::optional seq_starts = std::nullopt) { + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt) { torch::stable::accelerator::DeviceGuard device_guard( src_cache.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); @@ -1714,6 +1735,12 @@ void cp_gather_and_upconvert_fp8_kv_cache( "seq_starts must be int32"); } + const bool has_host_rows = host_cache.has_value(); + STD_TORCH_CHECK( + has_host_rows == host_row_ids.has_value() && + has_host_rows == device_row_ids.has_value(), + "host_cache, host_row_ids, and device_row_ids must be provided together"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), "src_cache and dst must be on the same device"); STD_TORCH_CHECK(src_cache.device() == block_table.device(), @@ -1724,6 +1751,43 @@ void cp_gather_and_upconvert_fp8_kv_cache( STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), "src_cache and seq_starts must be on the same device"); } + const uint8_t* host_cache_ptr = nullptr; + const int32_t* host_row_ids_ptr = nullptr; + const int32_t* device_row_ids_ptr = nullptr; + int64_t host_entry_stride = 0; + if (has_host_rows) { + auto const& host = host_cache.value(); + auto const& host_rows = host_row_ids.value(); + auto const& device_rows = device_row_ids.value(); + STD_TORCH_CHECK( + host.device().is_cpu() && + host.scalar_type() == torch::headeronly::ScalarType::Byte && + host.dim() == 2 && host.is_contiguous() && + host.size(1) == src_cache.size(2), + "host_cache must be a contiguous uint8 CPU row matrix matching " + "src_cache"); + cudaPointerAttributes attributes{}; + const cudaError_t pointer_status = + cudaPointerGetAttributes(&attributes, host.const_data_ptr()); + if (pointer_status != cudaSuccess) { + cudaGetLastError(); + } + STD_TORCH_CHECK( + pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost, + "host_cache must be CUDA-accessible pinned memory"); + STD_TORCH_CHECK( + host_rows.is_cuda() && device_rows.is_cuda() && + host_rows.scalar_type() == torch::headeronly::ScalarType::Int && + device_rows.scalar_type() == torch::headeronly::ScalarType::Int && + host_rows.is_contiguous() && device_rows.is_contiguous() && + host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 && + device_rows.dim() <= 2, + "host/device row maps must be matching contiguous CUDA int32 tensors"); + host_cache_ptr = reinterpret_cast(host.const_data_ptr()); + host_row_ids_ptr = host_rows.const_data_ptr(); + device_row_ids_ptr = device_rows.const_data_ptr(); + host_entry_stride = host.stride(0); + } auto dtype = src_cache.scalar_type(); STD_TORCH_CHECK( dtype == torch::headeronly::ScalarType::Byte || // uint8 @@ -1766,7 +1830,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( workspace_starts.const_data_ptr(), static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, - dst_entry_stride, seq_starts_ptr); + dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr, + device_row_ids_ptr, host_entry_stride); } void cp_gather_and_upconvert_nvfp4_kv_cache( diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu new file mode 100644 index 000000000000..63d84576272b --- /dev/null +++ b/csrc/libtorch_stable/hisparse_kernels.cu @@ -0,0 +1,1226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// HiSparse decode hot-buffer kernels for sparse MLA (GLM-5 / DeepSeek-V3.2). +// +// The swap-in algorithm is a port of SGLang's hisparse +// load_cache_to_device_buffer kernel (sgl jit_kernel/csrc/hisparse.cuh), +// adapted to vLLM addressing: +// - tokens are keyed by their global KV slot id (block_table-converted +// indexer output) instead of in-request positions, so no per-request +// host-location table is needed: host pool row i is global slot i. +// - each batch row owns a fixed region of `region_stride` hot rows; +// slots [0, hot_size) are LRU-managed, slot `hot_size` holds the row's +// newest token (written directly by the KV-cache update). +// +// The kernels only move bytes; they are dtype agnostic. + +#include "torch_utils.h" +#include "ops.h" +#include "../cuda_utils.h" + +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kWarpSize = 32; +// Sentinel in the shared top-k scratch: entry already resolved (hit / +// newest / invalid), no miss handling needed. +constexpr int32_t kTokenDone = -1; +constexpr int32_t kHashEmpty = -1; + +bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) { + cudaPointerAttributes attributes{}; + const auto status = + cudaPointerGetAttributes(&attributes, tensor.const_data_ptr()); + if (status != cudaSuccess) { + cudaGetLastError(); + return false; + } + return attributes.type == cudaMemoryTypeHost; +} + +__device__ __forceinline__ int32_t hash_slot(int32_t key, int32_t hash_size) { + // Knuth multiplicative hash for the open-addressing table. + return static_cast((static_cast(key) * 2654435761u) % + static_cast(hash_size)); +} + +// Copy one row of `row_bytes` bytes with a single warp. MLA rows take the +// 16-byte vectorized path; packed indexer rows use the scalar tail-safe path. +__device__ __forceinline__ void copy_row_warp(int lane_id, const char* src, + char* dst, int64_t row_bytes) { + const auto alignment = reinterpret_cast(src) | + reinterpret_cast(dst) | + static_cast(row_bytes); + if ((alignment & 15) == 0) { + const int64_t num_vec = row_bytes / 16; + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_vec; j += kWarpSize) { + __stcg(dst4 + j, __ldcg(src4 + j)); + } + return; + } + + if ((alignment & 3) == 0) { + const int64_t num_words = row_bytes / 4; + const unsigned int* src_words = reinterpret_cast(src); + unsigned int* dst_words = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_words; j += kWarpSize) { + __stcg(dst_words + j, __ldcg(src_words + j)); + } + return; + } + + for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) { + __stcg(dst + j, __ldcg(src + j)); + } +} + +// Zero one row with a single warp (16B vectorized, L2-only stores), same +// addressing contract as copy_row_warp. +__device__ __forceinline__ void zero_row_warp(int lane_id, char* dst, + int64_t row_bytes) { + const auto alignment = + reinterpret_cast(dst) | static_cast(row_bytes); + if ((alignment & 15) == 0) { + const int64_t num_vec = row_bytes / 16; + uint64_t* dst8 = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_vec; j += kWarpSize) { + uint64_t* d = dst8 + j * 2; + asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(d), "l"(0ULL), + "l"(0ULL) + : "memory"); + } + return; + } + + if ((alignment & 3) == 0) { + const int64_t num_words = row_bytes / 4; + unsigned int* dst_words = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_words; j += kWarpSize) { + __stcg(dst_words + j, 0u); + } + return; + } + + for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) { + __stcg(dst + j, static_cast(0)); + } +} + +__device__ __forceinline__ char* cache_row_ptr(char* cache, int64_t row, + int32_t block_size, + int64_t block_stride, + int64_t row_bytes) { + return cache + (row / block_size) * block_stride + + (row % block_size) * row_bytes; +} + +__device__ __forceinline__ void zero_cache_row_warp(int lane_id, char* cache, + int64_t row, + int32_t block_size, + int64_t block_stride, + int64_t row_bytes) { + zero_row_warp(lane_id, + cache_row_ptr(cache, row, block_size, block_stride, row_bytes), + row_bytes); +} + +// In-place inclusive scan over s_data[offset, count) performed by warp 0, +// carrying `accumulator` across calls. Returns the running total. +__device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id, + int offset, int count, + int accumulator) { + int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; +#pragma unroll + for (int i = 1; i < 32; i *= 2) { + int n = __shfl_up_sync(0xffffffff, val, i); + if (lane_id >= i) val += n; + } + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + return __shfl_sync(0xffffffff, val, 31); +} + +__device__ __forceinline__ int64_t +get_physical_hot_row(const int32_t* hot_block_table, int32_t row, + int64_t table_stride, int32_t block_size, int32_t slot) { + const int32_t block = + hot_block_table[static_cast(row) * table_stride + + slot / block_size]; + return static_cast(block) * block_size + slot % block_size; +} + +__device__ __forceinline__ void store_hot_index( + int32_t* hot_indices, int32_t* attention_indices, int32_t index, + int32_t physical_row, int32_t hot_block_size, + int64_t attention_block_stride) { + hot_indices[index] = physical_row; + if (attention_indices != nullptr) { + attention_indices[index] = + physical_row < 0 + ? -1 + : (physical_row / hot_block_size) * attention_block_stride + + physical_row % hot_block_size; + } +} + +// One block per batch row. +// +// Shared memory layout (int32 region followed by int16 region): +// s_topk[top_k] top-k global ids; reused as miss scratch +// s_chunk_off[nbc + 1] prefix sums for hit (then miss) compaction +// s_evict_off[nbc + 1] prefix sums for evictable compaction +// s_hash_keys[hash_size] open addressing: global id -> top-k index +// s_counters[3] hits, phase-1 resolved, valid count +// s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd] +// s_hash_vals[hash_size] int16 hash values (top-k index) +// Valid global ids must be unique within each row. +__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( + const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks] + const int32_t* __restrict__ global_indices, // global or request-relative + const int32_t* __restrict__ request_ids, // [num_rows] or nullptr + const int32_t* __restrict__ source_block_table, // [num_reqs, max_blocks] + const int32_t* __restrict__ resident_block_table, + int32_t* __restrict__ resolved_global_indices, // [num_rows, top_k] + int32_t* __restrict__ valid_counts, // [num_rows] or nullptr + int32_t* __restrict__ swap_host_physical_rows, // [num_rows, top_k] + int32_t* __restrict__ swap_device_physical_rows, // [num_rows, top_k] + int32_t* __restrict__ swap_counts, // [num_rows] + int32_t* __restrict__ hot_indices, // [num_rows, top_k] + int32_t* __restrict__ attention_indices, // [num_rows, top_k] + int32_t* __restrict__ miss_mask, // [num_rows, top_k] or nullptr + int32_t* __restrict__ device_global_indices, // [max_rows, region_stride] + int16_t* __restrict__ lru_slots, // [max_rows, hot_size] + const int32_t* __restrict__ request_state_indices, // [num_requests] or + // nullptr + const int32_t request_state_count, const int64_t host_rows, + const int64_t hot_table_stride, const int32_t hot_block_size, + const int32_t top_k, const int32_t hot_size, const int32_t hash_size, + const int64_t region_stride, const int64_t attention_block_stride, + const int64_t source_bt_stride, const int32_t source_num_reqs, + const int32_t source_num_blocks, const int32_t source_block_size, + const int64_t resident_bt_stride, const int32_t resident_num_reqs, + const int32_t resident_num_blocks, const int32_t resident_block_size, + const int32_t resident_null_block, const int64_t input_row_stride, + const int64_t attention_row_stride, const int64_t valid_count_stride) { + const int NUM_WARPS = blockDim.x / kWarpSize; + const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; + const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize; + + const int batch_row = blockIdx.x; + const int request_row = + request_ids != nullptr ? request_ids[batch_row] : batch_row; + const int state_row = + request_state_indices != nullptr && request_row >= 0 && + request_row < request_state_count + ? request_state_indices[request_row] + : (request_state_indices == nullptr ? request_row : -1); + // V2 publishes -1 for CUDA-graph padding rows. + if (state_row < 0) { + for (int i = threadIdx.x; i < top_k; i += blockDim.x) { + const int64_t index = static_cast(batch_row) * top_k + i; + hot_indices[index] = -1; + if (attention_indices != nullptr) { + attention_indices[static_cast(batch_row) * + attention_row_stride + + i] = -1; + } + if (resolved_global_indices != nullptr) { + resolved_global_indices[index] = -1; + } + } + if (valid_counts != nullptr && threadIdx.x == 0) { + valid_counts[static_cast(batch_row) * valid_count_stride] = 0; + } + if (swap_counts != nullptr && threadIdx.x == 0) { + swap_counts[batch_row] = 0; + } + return; + } + const int tid = threadIdx.x; + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1; + + const int32_t* row_topk = + global_indices + static_cast(batch_row) * input_row_stride; + int32_t* row_out = hot_indices + static_cast(batch_row) * top_k; + int32_t* row_attention = + attention_indices != nullptr + ? attention_indices + + static_cast(batch_row) * attention_row_stride + : nullptr; + int32_t* row_miss = (miss_mask != nullptr) + ? miss_mask + static_cast(batch_row) * top_k + : nullptr; + int32_t* row_dgi = + device_global_indices + static_cast(state_row) * region_stride; + int16_t* row_lru = lru_slots + static_cast(state_row) * hot_size; + + extern __shared__ char smem_raw[]; + int32_t* s_topk = reinterpret_cast(smem_raw); + int32_t* s_chunk_off = s_topk + top_k; + int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1); + int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1); + int32_t* s_counters = s_hash_keys + hash_size; + int16_t* s_lru_out = reinterpret_cast(s_counters + 3); + int16_t* s_hash_vals = s_lru_out + hot_size; + + if (tid < 3) { + s_counters[tid] = 0; + } + __syncthreads(); + + // Phase 1: translate request-relative positions and resolve resident rows. + for (int i = tid; i < top_k; i += blockDim.x) { + const int32_t token_index = row_topk[i]; + int32_t g = token_index; + int32_t resident_row = -1; + if (source_block_table != nullptr) { + const int32_t request_id = request_ids[batch_row]; + const int32_t source_block = + token_index >= 0 ? token_index / source_block_size : -1; + if (request_id >= 0 && request_id < source_num_reqs && + source_block >= 0 && source_block < source_num_blocks) { + const int32_t physical_block = + source_block_table[static_cast(request_id) * + source_bt_stride + + source_block]; + g = physical_block > 0 ? physical_block * source_block_size + + token_index % source_block_size + : -1; + } else { + g = -1; + } + if (resident_block_table != nullptr) { + const int32_t resident_block = + token_index >= 0 ? token_index / resident_block_size : -1; + if (request_id >= 0 && request_id < resident_num_reqs && + resident_block >= 0 && resident_block < resident_num_blocks) { + const int32_t physical_block = + resident_block_table[static_cast(request_id) * + resident_bt_stride + + resident_block]; + if (physical_block != resident_null_block && physical_block >= 0) { + resident_row = physical_block * resident_block_size + + token_index % resident_block_size; + } + } + } + } + if (g >= host_rows) { + g = -1; + } + if (resolved_global_indices != nullptr) { + resolved_global_indices[static_cast(batch_row) * top_k + i] = g; + } + if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1); + if (row_miss != nullptr) row_miss[i] = 0; + if (resident_row >= 0) { + store_hot_index(row_out, row_attention, i, resident_row, hot_block_size, + attention_block_stride); + s_topk[i] = kTokenDone; + atomicAdd(&s_counters[1], 1); + } else if (g < 0) { + store_hot_index(row_out, row_attention, i, -1, hot_block_size, + attention_block_stride); + s_topk[i] = kTokenDone; + atomicAdd(&s_counters[1], 1); + } else { + s_topk[i] = g; + } + } + __syncthreads(); + if (valid_counts != nullptr && tid == 0) { + valid_counts[static_cast(batch_row) * valid_count_stride] = + s_counters[2]; + } + // Fully resident rows need only request-relative page translation. Avoid + // scanning or rewriting the hot LRU when no selected row can consult it. + if (s_counters[1] == top_k) { + if (tid == 0 && swap_counts != nullptr) { + swap_counts[batch_row] = 0; + } + return; + } + + for (int i = tid; i < hash_size; i += blockDim.x) { + s_hash_keys[i] = kHashEmpty; + } + for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) { + s_chunk_off[i] = 0; + s_evict_off[i] = 0; + } + __syncthreads(); + + for (int i = tid; i < top_k; i += blockDim.x) { + const int32_t g = s_topk[i]; + if (g == kTokenDone) { + continue; + } + int slot = hash_slot(g, hash_size); + while (true) { + const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g); + if (old == kHashEmpty || old == g) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = (slot + 1) % hash_size; + } + } + __syncthreads(); + + // Phase 2: walk hot slots in LRU order, classify hit / evictable, and + // compact them (hits forward, evictables backward) into s_lru_out. + const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS; + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < iters_buffer; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < num_buffer_chunks; + + const int pos = chunk_idx * kWarpSize + lane_id; + const bool has_valid_pos = has_valid_chunk && (pos < hot_size); + const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1); + // Corruption tripwire: lru/dgi are long-lived device state; if some + // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot + // out of [0, region_stride), treat it as no-hit so it degrades to a re-miss + // instead of an unbounded dgi read (phases 3/5 bound the write side). + const int32_t cached_g = + (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1; + + int found_topk_idx = -1; + if (cached_g >= 0) { + int h = hash_slot(cached_g, hash_size); + while (true) { + const int32_t k = s_hash_keys[h]; + if (k == cached_g) { + found_topk_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == kHashEmpty) break; + h = (h + 1) % hash_size; + } + } + const bool is_hit = found_topk_idx >= 0; + const bool is_evictable = has_valid_pos && !is_hit; + + if (is_hit) { + s_topk[found_topk_idx] = kTokenDone; + store_hot_index(row_out, row_attention, found_topk_idx, + static_cast(get_physical_hot_row( + hot_block_table, request_row, hot_table_stride, + hot_block_size, slot)), + hot_block_size, attention_block_stride); + } + + int local_hit_off = 0; + int local_evict_off = 0; + if (has_valid_chunk) { + const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit); + const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable); + local_hit_off = __popc(hit_mask & lanes_before); + local_evict_off = __popc(evict_mask & lanes_before); + if (lane_id == 0) { + s_chunk_off[chunk_idx + 1] = __popc(hit_mask); + s_evict_off[chunk_idx + 1] = __popc(evict_mask); + } + } + __syncthreads(); + + if (warp_id == 0) { + total_hit_count = + warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1, + num_buffer_chunks + 1, total_hit_count); + total_evict_count = + warp_inclusive_scan(s_evict_off, lane_id, chunk_idx + 1, + num_buffer_chunks + 1, total_evict_count); + if (tid == 0) { + s_counters[0] = total_hit_count; + } + } + __syncthreads(); + + if (is_hit) { + const int off = s_chunk_off[chunk_idx] + local_hit_off; + s_lru_out[off] = slot; + } + if (is_evictable) { + const int off = s_evict_off[chunk_idx] + local_evict_off; + s_lru_out[hot_size - 1 - off] = slot; + } + } + __syncthreads(); + + // Reset prefix sums for the miss compaction (token chunks <= buffer + // chunks because hot_size >= top_k). + for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) { + s_chunk_off[i] = 0; + } + __syncthreads(); + + // Phase 3: compact misses, assign them eviction slots (oldest first) and + // record the new ownership in device_global_indices. + const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS; + int miss_running_total = 0; + for (int iter = 0; iter < iters_token; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < num_token_chunks; + + const int i = chunk_idx * kWarpSize + lane_id; + const bool has_valid_token = has_valid_chunk && (i < top_k); + + int32_t g = 0; + bool is_miss = false; + if (has_valid_token) { + is_miss = s_topk[i] != kTokenDone; + if (is_miss) { + g = s_topk[i]; + } + } + + int local_miss_off = 0; + if (has_valid_chunk) { + const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss); + local_miss_off = __popc(miss_mask & lanes_before); + if (lane_id == 0) { + s_chunk_off[chunk_idx + 1] = __popc(miss_mask); + } + } + __syncthreads(); + + if (warp_id == 0) { + miss_running_total = + warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1, + num_token_chunks + 1, miss_running_total); + } + __syncthreads(); + + if (is_miss) { + const int m = s_chunk_off[chunk_idx] + local_miss_off; + const int16_t evict_slot = s_lru_out[hot_size - 1 - m]; + if (evict_slot < 0 || evict_slot >= region_stride) { + // Corruption tripwire: an out-of-range slot (corrupted lru state, + // see phase 2) must not become a hot-cache/dgi write. Resolve the + // entry as invalid (-1, masked by attention, not in the miss set); + // it re-misses on a later step. Phase 5 re-checks the same + // s_lru_out value, so its copy is skipped consistently. + store_hot_index(row_out, row_attention, i, -1, hot_block_size, + attention_block_stride); + if (swap_host_physical_rows != nullptr) { + const int64_t compact_index = + static_cast(batch_row) * top_k + m; + swap_host_physical_rows[compact_index] = g; + swap_device_physical_rows[compact_index] = -1; + } + } else { + // Reuse s_topk as compacted miss scratch: m < i always holds (done + // entries are skipped), so writes never overrun pending reads. + s_topk[m] = g; + const int32_t physical_row = static_cast( + get_physical_hot_row(hot_block_table, request_row, hot_table_stride, + hot_block_size, evict_slot)); + store_hot_index(row_out, row_attention, i, physical_row, hot_block_size, + attention_block_stride); + if (swap_host_physical_rows != nullptr) { + const int64_t compact_index = + static_cast(batch_row) * top_k + m; + swap_host_physical_rows[compact_index] = g; + swap_device_physical_rows[compact_index] = physical_row; + } + if (row_miss != nullptr) row_miss[i] = 1; + row_dgi[evict_slot] = g; + } + } + } + __syncthreads(); + + const int total_hits = s_counters[0]; + const int total_misses = top_k - total_hits - s_counters[1]; + if (swap_counts != nullptr && tid == 0) { + swap_counts[batch_row] = total_misses; + } + // Phase 4: write back the LRU order: stale evictables at the front, + // freshly loaded misses next, then hits at MRU. + const int total_evictable = hot_size - total_hits; + const int remaining_evictable = total_evictable - total_misses; + for (int i = tid; i < hot_size; i += blockDim.x) { + if (i < remaining_evictable) { + row_lru[i] = s_lru_out[hot_size - 1 - total_misses - i]; + } else if (i < remaining_evictable + total_misses) { + row_lru[i] = s_lru_out[hot_size - 1 - (i - remaining_evictable)]; + } else { + row_lru[i] = s_lru_out[i - remaining_evictable - total_misses]; + } + } +} + +// Shared-layer plan replay. Given a plan (hot_indices + miss_mask) already +// computed by the group's index-producing layer, gather THIS +// layer's own missed KV rows into the planned hot slots. No LRU resolution: +// index-sharing shared layers see identical global_indices/hot_indices, so the +// slot assignment is identical -- only the per-layer bytes differ. Fixed shape +// (num_rows x top_k), so it is CUDA-graph-capture safe. +__global__ void hisparse_gather_plan_kernel( + const char* __restrict__ host_cache, // [host_rows, row_bytes] + char* __restrict__ hot_cache, // [hot_rows, row_bytes] + const int32_t* __restrict__ global_indices, // [num_rows, top_k] + const int32_t* __restrict__ hot_indices, // [num_rows, top_k] abs hot rows + const int32_t* __restrict__ miss_mask, // [num_rows, top_k] + int32_t* __restrict__ attention_indices, // [num_rows, top_k] + const int32_t* __restrict__ request_state_indices, // [num_rows] or nullptr + const int64_t host_rows, const int64_t hot_rows, const int64_t row_bytes, + const int64_t hot_block_stride, const int32_t hot_block_size, + const int32_t top_k, const int64_t attention_block_stride) { + const int NUM_WARPS = blockDim.x / kWarpSize; + // Columns are interleaved across gridDim.y blocks per row so few-row + // launches (local-prefill staging's single-row layout, small decode + // batches) still fill the device with outstanding host reads instead of + // starving on one block per row. + const int row = blockIdx.x; + const bool is_padding = + request_state_indices != nullptr && request_state_indices[row] < 0; + const int warp_id = threadIdx.x / kWarpSize; + const int lane_id = threadIdx.x % kWarpSize; + const int col_start = blockIdx.y * NUM_WARPS + warp_id; + const int col_stride = gridDim.y * NUM_WARPS; + const int64_t base = static_cast(row) * top_k; + for (int col = col_start; col < top_k; col += col_stride) { + const int32_t dst = hot_indices[base + col]; + if (lane_id == 0 && attention_indices != nullptr) { + attention_indices[base + col] = + is_padding + ? -1 + : (dst < 0 ? -1 + : (dst / hot_block_size) * attention_block_stride + + dst % hot_block_size); + } + if (is_padding) { + continue; + } + if (miss_mask[base + col] == 0) { + continue; + } + const int32_t g = global_indices[base + col]; + if (g < 0 || dst < 0 || dst >= hot_rows) { + continue; + } + if (g < host_rows) { + copy_row_warp(lane_id, host_cache + static_cast(g) * row_bytes, + cache_row_ptr(hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes), + row_bytes); + } else { + // No source row for g: zero the planned slot rather than serving + // whatever bytes it held (see the swap-in kernel's phase 5). + zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes); + } + } +} + +// Copy the compact swap rows produced by the residency resolver. +__global__ void hisparse_gather_compact_kernel( + const char* __restrict__ host_cache, char* __restrict__ hot_cache, + const int32_t* __restrict__ miss_global_indices, + const int32_t* __restrict__ miss_hot_indices, + const int32_t* __restrict__ miss_counts, const int64_t host_rows, + const int64_t hot_rows, const int64_t row_bytes, + const int64_t hot_block_stride, const int32_t hot_block_size, + const int32_t top_k) { + const int NUM_WARPS = blockDim.x / kWarpSize; + const int row = blockIdx.x; + const int warp_id = threadIdx.x / kWarpSize; + const int lane_id = threadIdx.x % kWarpSize; + const int miss_count = min(max(miss_counts[row], 0), top_k); + const int col_start = blockIdx.y * NUM_WARPS + warp_id; + const int col_stride = gridDim.y * NUM_WARPS; + const int64_t base = static_cast(row) * top_k; + for (int col = col_start; col < miss_count; col += col_stride) { + const int32_t g = miss_global_indices[base + col]; + const int32_t dst = miss_hot_indices[base + col]; + if (g < 0 || dst < 0 || dst >= hot_rows) { + continue; + } + if (g < host_rows) { + copy_row_warp(lane_id, host_cache + static_cast(g) * row_bytes, + cache_row_ptr(hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes), + row_bytes); + } else { + zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes); + } + } +} + +__global__ void hisparse_invalidate_written_slots_kernel( + int32_t* __restrict__ device_global_indices, + const int32_t* __restrict__ request_state_indices, + const int32_t* __restrict__ req_id_per_token, + const int64_t* __restrict__ written_slots, const int64_t num_tokens, + const int64_t num_request_ids, const int64_t num_state_rows, + const int64_t region_stride) { + const int64_t token_idx = blockIdx.x; + if (token_idx >= num_tokens) { + return; + } + const int32_t req_idx = req_id_per_token[token_idx]; + if (req_idx < 0 || req_idx >= num_request_ids) { + return; + } + const int32_t state_idx = request_state_indices[req_idx]; + const int64_t written_slot = written_slots[token_idx]; + if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0) { + return; + } + int32_t* row = device_global_indices + state_idx * region_stride; + for (int64_t offset = threadIdx.x; offset < region_stride; + offset += blockDim.x) { + if (row[offset] == written_slot) { + row[offset] = -1; + } + } +} + +int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name, + int64_t row_bytes) { + STD_TORCH_CHECK(t.dim() == 2, name, " must be 2D"); + STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + STD_TORCH_CHECK(t.size(1) * t.element_size() == row_bytes, name, + " row width mismatch"); + return t.size(0); +} + +} // namespace + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots) { + STD_TORCH_CHECK(device_global_indices.is_cuda() && + request_state_indices.is_cuda() && + req_id_per_token.is_cuda() && written_slots.is_cuda(), + "HiSparse invalidation tensors must be on CUDA"); + STD_TORCH_CHECK( + device_global_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + device_global_indices.dim() == 2 && + device_global_indices.is_contiguous(), + "device_global_indices must be a contiguous 2D int32 CUDA tensor"); + STD_TORCH_CHECK( + request_state_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + request_state_indices.dim() == 1 && + request_state_indices.is_contiguous(), + "request_state_indices must be a contiguous 1D int32 CUDA tensor"); + STD_TORCH_CHECK( + req_id_per_token.scalar_type() == torch::headeronly::ScalarType::Int && + req_id_per_token.dim() == 1 && req_id_per_token.is_contiguous(), + "req_id_per_token must be a contiguous 1D int32 CUDA tensor"); + STD_TORCH_CHECK( + written_slots.scalar_type() == torch::headeronly::ScalarType::Long && + written_slots.dim() == 1 && written_slots.is_contiguous(), + "written_slots must be a contiguous 1D int64 CUDA tensor"); + STD_TORCH_CHECK(req_id_per_token.numel() == written_slots.numel(), + "req_id_per_token and written_slots must have equal length"); + const int device_index = device_global_indices.get_device_index(); + STD_TORCH_CHECK( + request_state_indices.get_device_index() == device_index && + req_id_per_token.get_device_index() == device_index && + written_slots.get_device_index() == device_index, + "HiSparse invalidation tensors must be on the same CUDA device"); + + const int64_t num_tokens = written_slots.numel(); + const int64_t num_state_rows = device_global_indices.size(0); + const int64_t region_stride = device_global_indices.size(1); + if (num_tokens == 0 || num_state_rows == 0 || region_stride == 0) { + return; + } + STD_TORCH_CHECK(num_tokens <= INT32_MAX, + "HiSparse invalidation grid exceeds CUDA limits"); + + constexpr int kBlockSize = 256; + const torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_invalidate_written_slots_kernel<<(num_tokens), + kBlockSize, 0, stream>>>( + device_global_indices.mutable_data_ptr(), + request_state_indices.const_data_ptr(), + req_id_per_token.const_data_ptr(), + written_slots.const_data_ptr(), num_tokens, + request_state_indices.numel(), num_state_rows, region_stride); + const cudaError_t launch_error = cudaGetLastError(); + STD_TORCH_CHECK(launch_error == cudaSuccess, + "HiSparse invalidation kernel launch failed: ", + cudaGetErrorString(launch_error)); +} + +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, + std::optional const& miss_mask, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + hot_block_table.is_cuda() && + hot_block_table.scalar_type() == torch::headeronly::ScalarType::Int && + hot_block_table.dim() == 2, + "hot_block_table must be a 2D int32 CUDA tensor"); + STD_TORCH_CHECK(global_indices.is_cuda() && hot_indices.is_cuda() && + device_global_indices.is_cuda() && lru_slots.is_cuda(), + "index tensors must be on CUDA"); + STD_TORCH_CHECK( + global_indices.scalar_type() == torch::headeronly::ScalarType::Int && + hot_indices.scalar_type() == torch::headeronly::ScalarType::Int, + "global_indices/hot_indices must be int32"); + STD_TORCH_CHECK( + device_global_indices.scalar_type() == torch::headeronly::ScalarType::Int, + "device_global_indices must be int32"); + STD_TORCH_CHECK( + lru_slots.scalar_type() == torch::headeronly::ScalarType::Short, + "lru_slots must be int16"); + STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.stride(1) == 1, + "global_indices must be row-major 2D"); + const int64_t num_rows = global_indices.size(0); + STD_TORCH_CHECK(num_rows >= 0 && num_rows <= INT32_MAX, + "num_rows must fit the CUDA grid"); + STD_TORCH_CHECK(hot_indices.size(0) == num_rows && + hot_indices.size(1) == global_indices.size(1) && + hot_indices.is_contiguous(), + "hot_indices must have one contiguous row per launch row"); + STD_TORCH_CHECK( + device_global_indices.dim() == 2 && device_global_indices.is_contiguous(), + "device_global_indices must be contiguous 2D"); + STD_TORCH_CHECK( + lru_slots.dim() == 2 && + lru_slots.size(0) == device_global_indices.size(0) && + lru_slots.is_contiguous(), + "lru_slots must be contiguous with one row per request state"); + STD_TORCH_CHECK(hot_cache.dim() == 3, + "hot_cache must be [num_blocks, block_size, row_width]"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes, + "hot-cache rows must be contiguous"); + const int64_t hot_block_size = hot_cache.size(1); + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + const auto launch_rows = static_cast(num_rows); + const auto top_k = static_cast(global_indices.size(1)); + const auto hot_size = static_cast(lru_slots.size(1)); + STD_TORCH_CHECK(hot_size >= top_k, "hot buffer size must be >= top_k"); + STD_TORCH_CHECK(hot_size <= 32768, "hot buffer size must fit int16 slots"); + STD_TORCH_CHECK(region_stride == hot_size, + "region_stride must match the LRU size"); + STD_TORCH_CHECK(device_global_indices.size(1) == region_stride, + "device_global_indices must cover the full hot region"); + STD_TORCH_CHECK(hot_block_table.size(1) >= + (region_stride + hot_block_size - 1) / hot_block_size, + "hot_block_table has too few columns"); + STD_TORCH_CHECK(hot_rows < INT32_MAX, "hot indices must fit int32"); + + const int32_t* request_ids_ptr = nullptr; + const int32_t* source_block_table_ptr = nullptr; + int64_t source_bt_stride = 0; + int32_t source_num_reqs = 0; + int32_t source_num_blocks = 0; + STD_TORCH_CHECK( + request_ids.has_value() == source_block_table.has_value(), + "request_ids and source_block_table must be provided together"); + if (source_block_table.has_value()) { + auto const& req = request_ids.value(); + auto const& table = source_block_table.value(); + STD_TORCH_CHECK( + req.is_cuda() && + req.scalar_type() == torch::headeronly::ScalarType::Int && + req.numel() >= launch_rows && req.is_contiguous(), + "request_ids must be contiguous int32 on CUDA with one entry per row"); + STD_TORCH_CHECK( + table.is_cuda() && + table.scalar_type() == torch::headeronly::ScalarType::Int && + table.dim() == 2 && table.stride(1) == 1, + "source_block_table must be a row-major 2D int32 CUDA tensor"); + STD_TORCH_CHECK(source_block_size > 0, + "source_block_size must be positive"); + request_ids_ptr = req.const_data_ptr(); + source_block_table_ptr = table.const_data_ptr(); + source_bt_stride = table.stride(0); + source_num_reqs = static_cast(table.size(0)); + source_num_blocks = static_cast(table.size(1)); + } + const int32_t required_hot_rows = + request_ids_ptr != nullptr ? source_num_reqs : launch_rows; + STD_TORCH_CHECK(hot_block_table.size(0) >= required_hot_rows, + "hot_block_table has too few rows"); + + const int32_t* resident_block_table_ptr = nullptr; + int64_t resident_bt_stride = 0; + int32_t resident_num_reqs = 0; + int32_t resident_num_blocks = 0; + if (resident_block_table.has_value()) { + auto const& table = resident_block_table.value(); + STD_TORCH_CHECK( + source_block_table.has_value() && request_ids.has_value(), + "resident lookup requires request_ids and source_block_table"); + STD_TORCH_CHECK( + table.is_cuda() && + table.scalar_type() == torch::headeronly::ScalarType::Int && + table.dim() == 2 && table.stride(1) == 1, + "resident_block_table must be a row-major 2D int32 CUDA tensor"); + STD_TORCH_CHECK(resident_block_size == hot_block_size, + "resident and hot block sizes must match"); + STD_TORCH_CHECK(resident_null_block >= 0, + "resident null block must be non-negative"); + resident_block_table_ptr = table.const_data_ptr(); + resident_bt_stride = table.stride(0); + resident_num_reqs = static_cast(table.size(0)); + resident_num_blocks = static_cast(table.size(1)); + } + + int32_t* resolved_global_indices_ptr = nullptr; + if (resolved_global_indices.has_value()) { + auto const& resolved = resolved_global_indices.value(); + STD_TORCH_CHECK( + resolved.is_cuda() && + resolved.scalar_type() == torch::headeronly::ScalarType::Int && + resolved.dim() == 2 && resolved.size(0) == launch_rows && + resolved.size(1) == global_indices.size(1) && + resolved.is_contiguous(), + "resolved_global_indices must be contiguous int32 matching indices"); + resolved_global_indices_ptr = resolved.mutable_data_ptr(); + } + + int32_t* valid_counts_ptr = nullptr; + if (valid_counts.has_value()) { + auto const& counts = valid_counts.value(); + STD_TORCH_CHECK( + counts.is_cuda() && + counts.scalar_type() == torch::headeronly::ScalarType::Int && + counts.dim() == 1 && counts.size(0) == launch_rows, + "valid_counts must be int32 on CUDA with one entry per row"); + valid_counts_ptr = counts.mutable_data_ptr(); + } + + int32_t* swap_host_physical_rows_ptr = nullptr; + int32_t* swap_device_physical_rows_ptr = nullptr; + int32_t* swap_counts_ptr = nullptr; + const bool has_compact_swaps = swap_host_physical_rows.has_value(); + STD_TORCH_CHECK(has_compact_swaps == swap_device_physical_rows.has_value() && + has_compact_swaps == swap_counts.has_value(), + "compact swap tensors must be provided together"); + if (has_compact_swaps) { + auto const& globals = swap_host_physical_rows.value(); + auto const& hots = swap_device_physical_rows.value(); + auto const& counts = swap_counts.value(); + STD_TORCH_CHECK( + globals.is_cuda() && hots.is_cuda() && counts.is_cuda() && + globals.scalar_type() == torch::headeronly::ScalarType::Int && + hots.scalar_type() == torch::headeronly::ScalarType::Int && + counts.scalar_type() == torch::headeronly::ScalarType::Int && + globals.dim() == 2 && hots.dim() == 2 && + globals.size(0) == launch_rows && + globals.size(1) == global_indices.size(1) && + hots.size(0) == launch_rows && + hots.size(1) == global_indices.size(1) && + counts.numel() >= launch_rows && globals.is_contiguous() && + hots.is_contiguous() && counts.is_contiguous(), + "compact swaps must be contiguous int32 matching indices"); + swap_host_physical_rows_ptr = globals.mutable_data_ptr(); + swap_device_physical_rows_ptr = hots.mutable_data_ptr(); + swap_counts_ptr = counts.mutable_data_ptr(); + } + + const int32_t* request_state_ptr = nullptr; + int32_t request_state_count = 0; + if (request_state_indices.has_value()) { + auto const& state_indices = request_state_indices.value(); + STD_TORCH_CHECK( + state_indices.is_cuda() && + state_indices.scalar_type() == torch::headeronly::ScalarType::Int && + state_indices.dim() == 1 && state_indices.is_contiguous() && + state_indices.numel() <= INT32_MAX, + "request_state_indices must be contiguous 1D int32 on CUDA"); + request_state_ptr = state_indices.const_data_ptr(); + request_state_count = static_cast(state_indices.numel()); + } else { + STD_TORCH_CHECK(device_global_indices.size(0) >= launch_rows, + "device_global_indices has too few rows"); + } + + // Optional output: 1 at columns requiring a host-to-device swap, 0 elsewhere. + int32_t* miss_mask_ptr = nullptr; + if (miss_mask.has_value()) { + auto const& mm = miss_mask.value(); + STD_TORCH_CHECK( + mm.is_cuda() && + mm.scalar_type() == torch::headeronly::ScalarType::Int && + mm.dim() == 2 && mm.is_contiguous() && mm.size(0) == launch_rows && + mm.size(1) == global_indices.size(1), + "miss_mask must be a contiguous int32 CUDA tensor matching " + "global_indices"); + miss_mask_ptr = mm.mutable_data_ptr(); + } + + int32_t* attention_indices_ptr = nullptr; + if (attention_indices.has_value()) { + auto const& indices = attention_indices.value(); + STD_TORCH_CHECK( + indices.is_cuda() && + indices.scalar_type() == torch::headeronly::ScalarType::Int && + indices.dim() == 2 && indices.size(0) == launch_rows && + indices.stride(1) == 1 && indices.size(1) == global_indices.size(1), + "attention_indices must be row-major int32 matching global_indices"); + STD_TORCH_CHECK(attention_block_stride >= hot_block_size, + "attention block stride must cover one hot block"); + attention_indices_ptr = indices.mutable_data_ptr(); + } + + if (launch_rows == 0 || top_k == 0) { + return; + } + + constexpr int kBlockSize = 1024; + const int hash_size = 2 * top_k; + const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; + const size_t smem_bytes = + sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) + + sizeof(int16_t) * (hot_size + hash_size); + + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + const int64_t attention_row_stride = + attention_indices.has_value() ? attention_indices.value().stride(0) : 0; + const int64_t valid_count_stride = + valid_counts.has_value() ? valid_counts.value().stride(0) : 0; + auto kernel = hisparse_resolve_residency_kernel; + if (smem_bytes > 48 * 1024) { + const cudaError_t attribute_error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + STD_TORCH_CHECK(attribute_error == cudaSuccess, + "failed to configure HiSparse swap-in shared memory: ", + cudaGetErrorString(attribute_error)); + } + kernel<<>>( + hot_block_table.const_data_ptr(), + global_indices.const_data_ptr(), request_ids_ptr, + source_block_table_ptr, resident_block_table_ptr, + resolved_global_indices_ptr, valid_counts_ptr, + swap_host_physical_rows_ptr, swap_device_physical_rows_ptr, + swap_counts_ptr, hot_indices.mutable_data_ptr(), + attention_indices_ptr, miss_mask_ptr, + device_global_indices.mutable_data_ptr(), + lru_slots.mutable_data_ptr(), request_state_ptr, + request_state_count, host_rows, hot_block_table.stride(0), hot_block_size, + top_k, hot_size, hash_size, region_stride, attention_block_stride, + source_bt_stride, source_num_reqs, source_num_blocks, + static_cast(source_block_size), resident_bt_stride, + resident_num_reqs, resident_num_blocks, + static_cast(resident_block_size), + static_cast(resident_null_block), global_indices.stride(0), + attention_row_stride, valid_count_stride); + const cudaError_t launch_error = cudaGetLastError(); + STD_TORCH_CHECK(launch_error == cudaSuccess, + "HiSparse residency kernel launch failed: ", + cudaGetErrorString(launch_error)); +} + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + global_indices.is_cuda() && hot_indices.is_cuda() && miss_mask.is_cuda(), + "plan tensors must be on CUDA"); + STD_TORCH_CHECK( + global_indices.scalar_type() == torch::headeronly::ScalarType::Int && + hot_indices.scalar_type() == torch::headeronly::ScalarType::Int && + miss_mask.scalar_type() == torch::headeronly::ScalarType::Int, + "plan tensors must be int32"); + STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.is_contiguous(), + "global_indices must be contiguous 2D"); + STD_TORCH_CHECK( + hot_indices.size(0) == global_indices.size(0) && + hot_indices.size(1) == global_indices.size(1) && + miss_mask.size(0) == global_indices.size(0) && + miss_mask.size(1) == global_indices.size(1) && + hot_indices.is_contiguous() && miss_mask.is_contiguous(), + "hot_indices/miss_mask must match contiguous 2D global_indices"); + + STD_TORCH_CHECK(hot_cache.dim() == 2 || hot_cache.dim() == 3, + "hot_cache must be a 2D staging buffer or paged 3D cache"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK( + hot_cache.stride(hot_cache.dim() - 2) * hot_cache.element_size() == + row_bytes, + "hot-cache rows must be contiguous"); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + + const auto num_rows = static_cast(global_indices.size(0)); + const auto top_k = static_cast(global_indices.size(1)); + + const int32_t* request_state_ptr = nullptr; + if (request_state_indices.has_value()) { + auto const& state_indices = request_state_indices.value(); + STD_TORCH_CHECK( + state_indices.is_cuda() && + state_indices.scalar_type() == torch::headeronly::ScalarType::Int && + state_indices.numel() >= num_rows, + "request_state_indices must be int32 on CUDA with one entry per row"); + request_state_ptr = state_indices.const_data_ptr(); + } + + int32_t* attention_indices_ptr = nullptr; + if (attention_indices.has_value()) { + auto const& indices = attention_indices.value(); + STD_TORCH_CHECK( + indices.is_cuda() && + indices.scalar_type() == torch::headeronly::ScalarType::Int && + indices.dim() == 2 && indices.is_contiguous() && + indices.size(0) == global_indices.size(0) && + indices.size(1) == global_indices.size(1), + "attention_indices must be a contiguous int32 CUDA tensor matching " + "global_indices"); + attention_indices_ptr = indices.mutable_data_ptr(); + } + + if (num_rows == 0 || top_k == 0) { + return; + } + + // Match the swap-in kernel's block size: the gather serves 3 of every 4 + // layers' misses (index-sharing replay), so per-row copy parallelism is + // the throughput limiter on cold rows. + constexpr int kBlockSize = 1024; + constexpr int kNumWarps = kBlockSize / kWarpSize; + // Interleave columns over enough blocks per row to cover the device even + // for few-row launches (local-prefill staging's single-row layout, small + // decode batches), keeping >= 1 column per warp. + constexpr int kTargetBlocks = 256; + const int max_chunks = std::max(1, (top_k + kNumWarps - 1) / kNumWarps); + const int num_chunks = + std::min(max_chunks, std::max(1, kTargetBlocks / num_rows)); + const dim3 grid(num_rows, num_chunks); + const int32_t hot_block_size = + hot_cache.dim() == 3 ? static_cast(hot_cache.size(1)) : 1; + if (attention_indices.has_value()) { + STD_TORCH_CHECK(attention_block_stride >= hot_block_size, + "attention block stride must cover one hot block"); + } + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_gather_plan_kernel<<>>( + static_cast(host_cache.const_data_ptr()), + static_cast(hot_cache.mutable_data_ptr()), + global_indices.const_data_ptr(), + hot_indices.const_data_ptr(), + miss_mask.const_data_ptr(), attention_indices_ptr, + request_state_ptr, host_rows, hot_rows, row_bytes, hot_block_stride, + hot_block_size, top_k, attention_block_stride); +} + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + miss_global_indices.is_cuda() && miss_hot_indices.is_cuda() && + miss_counts.is_cuda() && + miss_global_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + miss_hot_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + miss_counts.scalar_type() == torch::headeronly::ScalarType::Int && + miss_global_indices.dim() == 2 && miss_hot_indices.dim() == 2 && + miss_global_indices.size(0) == miss_hot_indices.size(0) && + miss_global_indices.size(1) == miss_hot_indices.size(1) && + miss_counts.numel() >= miss_global_indices.size(0) && + miss_global_indices.is_contiguous() && + miss_hot_indices.is_contiguous() && miss_counts.is_contiguous(), + "compact miss plan must be matching contiguous int32 CUDA tensors"); + STD_TORCH_CHECK(hot_cache.dim() == 3, + "hot_cache must be [num_blocks, block_size, row_width]"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes, + "hot-cache rows must be contiguous"); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + const auto num_rows = static_cast(miss_global_indices.size(0)); + const auto top_k = static_cast(miss_global_indices.size(1)); + if (num_rows == 0 || top_k == 0) { + return; + } + + constexpr int kBlockSize = 512; + constexpr int kTargetBlocks = 64; + const int num_chunks = std::min(8, std::max(1, kTargetBlocks / num_rows)); + const dim3 grid(num_rows, num_chunks); + const int32_t hot_block_size = static_cast(hot_cache.size(1)); + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_gather_compact_kernel<<>>( + static_cast(host_cache.const_data_ptr()), + static_cast(hot_cache.mutable_data_ptr()), + miss_global_indices.const_data_ptr(), + miss_hot_indices.const_data_ptr(), + miss_counts.const_data_ptr(), host_rows, hot_rows, row_bytes, + hot_block_stride, hot_block_size, top_k); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 13429a54e694..6893c6913d9e 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -661,6 +661,54 @@ void concat_and_cache_mla_grouped( std::optional kv_scales, const std::string& kv_cache_dtype); +#ifndef USE_ROCM +// HiSparse kernels use raw PTX in the row copy; CUDA-only. +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, + std::optional const& miss_mask, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block); + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots); + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride); + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts); + +#endif // !USE_ROCM + // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, @@ -704,7 +752,10 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size, - std::optional seq_starts = std::nullopt); + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt); // Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace void cp_gather_and_upconvert_nvfp4_kv_cache( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 336c495113b3..cb6f46339683 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -987,6 +987,57 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor? kv_scales=None," " str kv_cache_dtype='auto') -> ()"); +#ifndef USE_ROCM + ops.def( + "hisparse_resolve_residency(Tensor host_cache," + " Tensor! hot_cache," + " Tensor hot_block_table," + " Tensor global_indices," + " Tensor! hot_indices," + " Tensor! device_global_indices," + " Tensor! lru_slots," + " Tensor? request_state_indices," + " int region_stride," + " Tensor(a!)? miss_mask=None," + " Tensor(b!)? attention_indices=None," + " int attention_block_stride=0," + " Tensor? request_ids=None," + " Tensor? source_block_table=None," + " int source_block_size=0," + " Tensor(c!)? resolved_global_indices=None," + " Tensor(d!)? valid_counts=None," + " Tensor(e!)? swap_host_physical_rows=None," + " Tensor(f!)? swap_device_physical_rows=None," + " Tensor(g!)? swap_counts=None," + " Tensor? resident_block_table=None," + " int resident_block_size=0," + " int resident_null_block=0) -> ()"); + + ops.def( + "hisparse_invalidate_written_slots(Tensor! device_global_indices," + " Tensor request_state_indices," + " Tensor req_id_per_token," + " Tensor written_slots) -> ()"); + + ops.def( + "hisparse_gather_plan(Tensor host_cache," + " Tensor! hot_cache," + " Tensor global_indices," + " Tensor hot_indices," + " Tensor miss_mask," + " Tensor? request_state_indices," + " Tensor(a!)? attention_indices=None," + " int attention_block_stride=0) -> ()"); + + ops.def( + "hisparse_gather_compact(Tensor host_cache," + " Tensor! hot_cache," + " Tensor miss_global_indices," + " Tensor miss_hot_indices," + " Tensor miss_counts) -> ()"); + +#endif // !USE_ROCM + // Rotate Q and K, then write to kv cache for MLA ops.def( "concat_and_cache_mla_rope_fused(" @@ -1023,7 +1074,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { ops.def( "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts) -> ()"); + "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " + "device_row_ids=None) -> ()"); ops.def( "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " @@ -1091,6 +1143,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); ops.impl("concat_and_cache_mla_grouped", TORCH_BOX(&concat_and_cache_mla_grouped)); + +#ifndef USE_ROCM + ops.impl("hisparse_resolve_residency", + TORCH_BOX(&hisparse_resolve_residency)); + ops.impl("hisparse_invalidate_written_slots", + TORCH_BOX(&hisparse_invalidate_written_slots)); + ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); + ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); +#endif // !USE_ROCM ops.impl("concat_and_cache_mla_rope_fused", TORCH_BOX(&concat_and_cache_mla_rope_fused)); ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md new file mode 100644 index 000000000000..965ab81d8da1 --- /dev/null +++ b/docs/design/hisparse.md @@ -0,0 +1,244 @@ +# HiSparse local KV offload architecture + +Status: experimental + +## The short version + +There are three different jobs: + +1. The normal KV cache system manages GPU block pools and tables. +2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and + host/GPU residency transitions. +3. `HiSparseConnector` carries residency work between scheduler and worker; + `HiSparseWorker` coordinates transfers, while per-cache + `HiSparseRuntime` objects own host/hot views and GPU replacement state. + +Neither worker-side object allocates or frees logical blocks. HMA provides the +GPU allocation shared by the resident and hot groups; it does not manage CPU +memory or KV identity. + +```text +request ────► HiSparseCoordinator ── source blocks + residency policy + │ + ├──► KV cache manager ── resident/hot GPU leases (HMA) + │ + └──► HiSparseConnector ── host bytes + copies + GPU LRU +``` + +This is a local KV connector. When P/D or another offload connector is also +configured, `MultiConnector` composes it with `HiSparseConnector`. + +`host_pool_gib` is configured on `HiSparseConnector` and is the usable +host-cache capacity per data-parallel replica, not a node-wide memory budget. +Tensor-parallel ranks hold replicated views of that logical cache. Those views +may use private per-rank backing or one shared +physical allocation without changing the configured capacity. Physical host +memory consumption is therefore topology- and implementation-dependent. The +realized capacity may be slightly smaller because the budget is rounded down +to complete host blocks. + +## Ownership + +| Thing | Owner | What “owner” means | +| --- | --- | --- | +| HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks | +| Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks | +| Resident block tables | normal KV cache manager | tells attention where resident pages are | +| Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions | +| Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle | +| Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown | +| Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases | +| Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU | +| Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention | +| Sparse attention | attention backend | consumes a device cache and physical row IDs | +| HMA | allocator | provides GPU capacity; owns no KV meaning | + +The key distinction is logical allocation versus contents. `HiSparseCoordinator` +owns host block IDs and their request/prefix associations. `HiSparseWorker` and +its per-cache runtimes own the corresponding bytes. The normal cache manager +sees only device pools. +The source group has `block_pool_id=None`; device-pool consumers must narrow it +before indexing, so host ownership cannot masquerade as a numeric GPU pool. + +## Code boundary + +```text +scheduler process worker process + +HiSparseConnector HiSparseConnector + └─ HiSparseCoordinator └─ HiSparseWorker + │ │ + │ connector metadata ├─ host bytes + │ - page transfers ├─ copy scheduling + │ - block-table replacements └─ per-layer hot state + └───────────────────────────────────────────────►│ + ◄──────── connector worker metadata ─────────────┘ + enqueued and completed transfer IDs +``` + +The command travels in `kv_connector_metadata`; transfer updates return in +`KVConnectorOutput.kv_connector_worker_meta`. The model runner does not +interpret page transfers. Enqueue acknowledgements let the scheduler release +source leases in stream order; completion acknowledgements publish the copied +host pages. + +## Resident device pages + +Resident pages are intentionally outside `HiSparseRuntime`. + +KV-cache initialization binds cache-manager allocations to the attention-facing +`HiSparseCacheHandle` before constructing `HiSparseWorker`. That same +handle's runtime retains the resident source index needed by a transfer plan. +There is no second resident object or registration wrapper. + +```text +KV cache setup + │ + ├─ bind resident allocation ──► HiSparseCacheHandle + │ cache + block table + slot mapping + │ + ├─ bind host/hot allocation ──► HiSparseRuntime + │ host + hot + GPU LRU + │ + └─ register cache handles ────► HiSparseWorker + step-level transfers + +HiSparseWorker registers the same HiSparseCacheHandle objects directly +``` + +Attention construction links each layer to the most recent layer that actually +owns an indexer. This releases a follower's duplicate LRU tensors before GPU +memory profiling. Cache binding only attaches storage; it does not infer +semantic groups from the physical packed-tensor order. The construction cursor +is discarded with the worker's pinned state. + +Every HiSparse decode batch uses the same fused resolver. It checks resident +pages first, then hot rows, then pinned host memory. A resident hit exits inside +the kernel before hot-LRU lookup or host copying; there is no framework-level +residency route or separate CUDA graph. No CPU decision is added to the decode +path. The resolver consumes the existing graph-stable request mapping from +attention metadata; neither the worker nor individual cache handles keep a +duplicate mapping. + +Speculative decoding resolves and consumes each verification step in order. +Each step receives distinct replayable plan rows while sharing the request's +hot-cache state, so a later step cannot reuse a hot row before an earlier step +has consumed it. + +## P/D import target + +The decoder chooses the landing target once per request from the normal cache +admission calculation. If the complete imported prefix fits the device pools, +NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed +host-backed GPU footprint and host source blocks fit, the request imports into +the host tier. There is no context-length threshold or other heuristic, and a +request waiting for capacity retains its choice across admission retries. + +A host import reads through a bounded decoder-GPU staging pool before copying +into registered host memory. Pages needed immediately are mirrored into their +resident destinations during that copy. Both landing targets then use the same +fused decode resolver described above. + +## Indexer KV offloading + +HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a +normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured +with HiSparse, it stores and restores that group through the generic KV +offloading path; HiSparse continues to own only the sparse MLA host tier. + +The two prefix sources can have different hit lengths. When the HiSparse host +prefix extends beyond the GPU-resident indexer prefix, the scheduler asks +`OffloadingConnector` to restore only the missing indexer suffix, capped at the +host prefix boundary. If that suffix is unavailable, all groups fall back to +the shorter prefix they share. NIXL P/D transfers continue to place indexer KV +directly in its GPU group. + +## Spill transaction + +A resident block cannot be reused until its contents have been handed to the +worker. + +```text +HiSparseCoordinator HiSparseWorker + │ │ + │ pin source and destination leases │ + │── SparseKVPageTransfer ─────────────►│ + │ │ enqueue GPU-to-host copy + │◄── enqueued transfer ID ────────────│ + │ replace resident table entry │ + │ release resident lease to HMA │ + │ │ copy reaches its event + │◄── completed transfer ID ───────────│ + │ mark host page valid │ + │ release destination host lease │ +``` + +“Enqueued” means the copy has entered the worker stream. Stream ordering makes +it safe to reuse the resident GPU block for later work, but the host page is +not yet published. “Completed” means the worker has observed the copy's event; +only then does the coordinator publish the host page for prefix reuse and +release its destination lease. A host-write event separately protects direct +CPU readers from writes already queued on the accelerator. + +The worker transfer contains only its transfer ID and physical copy +coordinates. Request identity and logical page state remain in the scheduler. + +## Hot lookup and LRU + +The NVIDIA CUDA path keeps replacement entirely on the accelerator: + +```text +top-K logical positions + │ + ▼ +resident page? ── yes ──► resident physical row + │ no + ▼ +hot row? ──────── yes ──► existing hot physical row + update GPU LRU + │ no + ▼ +choose GPU LRU victim ──► copy pinned host row ──► hot physical row +``` + +ROCm is not currently supported because the fused HiSparse cache operations are +implemented only by CUDA kernels. A future platform-specific worker may provide +the same command, output, and cache-resolution boundaries. + +## Main classes + +| Class | Inherits / implements | Responsibility | +| --- | --- | --- | +| `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine | +| `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary | +| `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes | +| `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding | +| `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle | +| `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution | +| `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution | +| `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work | + +## Performance invariants + +- Resident hits bypass hot-LRU lookup and host copies inside the fused resolver. +- Hot lookup, victim selection, and LRU updates stay on the GPU. +- A hot miss still copies directly from registered pinned host memory. +- Top-K resolution stays inside the attention invocation and remains graph + capturable. +- Compatible layers still share one miss plan. +- Index-sharing followers release their private LRU state before memory sizing. +- Indexer KV is untouched by HiSparse unless a generic KV offloader is configured. +- Resident and hot leases can still share one packed HMA allocation. +- No device scalar readback or CPU/device metadata round trip is added. +- The abstraction wraps the fused kernel; it does not add another kernel + launch. +- When HiSparse is disabled, the scheduler does not construct an offload + command or empty update table. + +## What remains platform-specific + +The command/result and attention-layer boundaries can be shared. The host +allocator, copy implementation, hot layout, and replacement policy should stay +platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot +kernel. ROCm is not currently supported; AMD or other accelerator backends can +implement their own worker without forcing NVIDIA's policy into the shared +boundary. diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index b5a1e6a9cc29..83e54afd9e7f 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -364,6 +364,12 @@ def main() -> None: type=int, help="Maximum number of concurrent requests", ) + parser.add_argument( + "--request-timeout-seconds", + type=float, + default=600, + help="Timeout for each request, including time waiting for a connection", + ) parser.add_argument("--save-results", type=str, help="Save results to JSON file") args = parser.parse_args() @@ -377,6 +383,7 @@ def main() -> None: temperature=args.temperature, seed=args.seed, max_concurrency=args.max_concurrency, + request_timeout_seconds=args.request_timeout_seconds, ) # Print results to terminal diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index ac4607355999..ffdb76e176f8 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -217,7 +217,8 @@ def run(kv): assert (out_ref - out_view).abs().max().item() == 0.0 -def test_flashmla_fp8_sparse_decode_unified_slot_view(): +@pytest.mark.parametrize("next_n", [1, 2, 4]) +def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a unified-slot block-major view bit-identically to a contiguous fp8_ds_mla cache, with finite nonzero output.""" @@ -240,7 +241,10 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(): n_layers = 3 layer = 1 - q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1 + q = ( + torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16) + * 0.1 + ) # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B # bf16 rope (random bytes corrupt the scale region and yield NaNs). @@ -273,13 +277,13 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(): # Sparse indices: each batch uses its own disjoint blocks. blocks_per_batch = num_blocks // batch - idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32) + idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32) for b in range(batch): slots: list[int] = [] for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch): slots.extend(blk * page + off for off in range(page)) slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32) - idx[b, 0, : slots_t.numel()] = slots_t + idx[b, :, : slots_t.numel()] = slots_t def run(kc): meta, num_splits = fm.get_mla_metadata() diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index 35afe33808cc..fe539abb4f72 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -316,6 +316,48 @@ def test_cp_gather_fp8_shuffled_blocks(): assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) +def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows(): + seq_lens = [7, 5] + block_size = 4 + ( + cache, + block_table, + workspace_starts, + num_reqs, + total_tokens, + expected, + ) = _build_test_case(seq_lens, block_size) + + device_cache = cache.clone() + host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES) + device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda") + use_device = device_rows % 2 == 0 + host_rows = device_rows.clone() + device_rows.masked_fill_(~use_device, -1) + + host_cache[use_device.cpu()] = 0 + device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0 + dst = torch.empty( + total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" + ) + + ops.cp_gather_and_upconvert_fp8_kv_cache( + device_cache, + dst, + block_table, + workspace_starts, + num_reqs, + host_cache=host_cache, + host_row_ids=host_rows, + device_row_ids=device_rows, + ) + + torch.testing.assert_close( + dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 + ) + assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) + + @pytest.mark.parametrize( "gather_seq_lens,seq_starts", [ diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index d0562f6d2d25..fbe22af2f897 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -234,6 +234,76 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str assert (topk == -1).all(), "topk buffer not cleared on indexer layer" +def test_fused_norm_rope_packed_indexer_block_stride(): + """Indexer writes use their HMA block stride and independent slot map.""" + torch.manual_seed(7) + dev = "cuda" + num_tokens = block_size = 2 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + cos_sin = make_cos_sin(32, ROPE_DIM, dev) + + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 + packed_block_stride = block_size * idx_row + 64 + backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8) + idx_cache = torch.as_strided( + backing, + (2, block_size, idx_row), + (packed_block_stride, idx_row, 1), + ) + mla_cache = torch.zeros( + 1, + block_size, + KV_LORA + ROPE_DIM, + device=dev, + dtype=torch.bfloat16, + ) + topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32) + + K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + cos_sin, + ik, + ikw, + ikb, + EPS, + cos_sin, + topk, + slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64), + indexer_slot_mapping=torch.arange( + block_size, block_size + num_tokens, device=dev, dtype=torch.int64 + ), + indexer_k_cache=idx_cache, + mla_kv_cache=mla_cache, + has_indexer=True, + ) + + ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False) + q_ref, s_ref = ue8m0_quant(ik_ref) + packed = idx_cache[1].reshape(-1) + values = ( + packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM) + ) + scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32) + assert_fp8(values, q_ref, "packed indexer-K fp8") + torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) + assert (backing[block_size * idx_row : packed_block_stride] == 0).all() + + def test_fused_norm_rope_materializes_qk_without_local_cache_slots(): """DCP non-owner ranks still need valid query shards for query AllGather.""" torch.manual_seed(7) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index de6518ebec79..f8e7a89ccad2 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -7,6 +7,7 @@ import torch import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.models.deepseek_v32.attention as deepseek_attention from vllm.config import CUDAGraphMode from vllm.models.deepseek_v32 import attention as deepseek_v32_attention from vllm.models.deepseek_v32.attention import DeepseekV32Attention @@ -16,6 +17,51 @@ MLA_LAYER = "model.layers.0.self_attn.attn" +def test_sparse_attention_refreshes_batch_state_inside_eager_segment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + metadata = SimpleNamespace(num_actual_tokens=1) + + class BatchStateRefreshed(Exception): + pass + + def refresh_batch_state(attn_metadata): + assert attn_metadata is metadata + raise BatchStateRefreshed + + layer = SimpleNamespace( + indexer=None, + skip_topk=True, + layer_name=MLA_LAYER, + impl=SimpleNamespace( + prepare_for_batch=refresh_batch_state, + record_logical_topk_ready=lambda: None, + ), + ) + monkeypatch.setattr( + deepseek_attention, + "get_attention_context", + lambda layer_name: (metadata, None, torch.empty(1), None), + ) + + with pytest.raises(BatchStateRefreshed): + DeepseekV32Attention._sparse_indexer_and_attn( + layer, + torch.empty(1, dtype=torch.long), + torch.empty(1, 1), + torch.empty(1, 1, 1), + torch.empty(1, 1, 1), + None, + None, + None, + None, + None, + torch.empty(1, 1, 1), + torch.empty(1, 1, 1), + torch.empty(1, 1), + ) + + def make_indexer_metadata( *, num_decodes: int = 0, @@ -246,6 +292,10 @@ def record_forward_impl(*args): _fp8_query=fp8_query, _use_sparse_mha=lambda _: True, rotary_emb=lambda _positions, q: (q + 1, None), + impl=SimpleNamespace( + record_logical_topk_ready=lambda: None, + prepare_for_batch=lambda _: None, + ), forward_impl=record_forward_impl, ) q_nope = torch.randn(2, 1, 2) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index a8ecd1276e16..dd162629b8f4 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from inspect import signature from types import SimpleNamespace from typing import Any from unittest.mock import Mock, call, patch @@ -76,8 +77,16 @@ def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): assert token_counts == (32,) -def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): +@pytest.mark.parametrize("skip_attn", [False, True]) +def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner + runner = _make_runner([]) + dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run) + runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind( + None, **kwargs + ) max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192) deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128) @@ -93,7 +102,7 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): ) as get_buckets, patch("vllm.utils.flashinfer.autotune") as autotune, ): - _run_flashinfer_autotune_dummy_runs(runner) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) assert get_buckets.call_args_list == [call(8192), call(128)] assert autotune.call_args_list == [ @@ -106,11 +115,13 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): skip_eplb=True, is_profile=True, randomize_inputs=True, + **({"skip_attn": True} if skip_attn else {}), ), call( num_tokens=128, skip_eplb=True, is_profile=True, randomize_inputs=True, + **({"skip_attn": True} if skip_attn else {}), ), ] diff --git a/tests/test_config.py b/tests/test_config.py index 2524968cce0d..7b6027313af0 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -19,10 +19,12 @@ import vllm.envs as envs from vllm.compilation.backends import VllmBackend from vllm.config import ( + AttentionConfig, CacheConfig, CompilationConfig, DeviceConfig, EngramConfig, + HiSparseConfig, KernelConfig, KVTransferConfig, ModelConfig, @@ -279,6 +281,62 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected +def test_hisparse_requires_v2_model_runner(): + config = object.__new__(VllmConfig) + config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig()) + + with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None): + assert config.use_v2_model_runner + with ( + patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", False), + pytest.raises(ValueError, match="requires Model Runner V2"), + ): + _ = config.use_v2_model_runner + + +def test_hisparse_rejects_decode_context_parallelism(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + with pytest.raises(ValueError, match="decode context parallelism"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + decode_context_parallel_size=2, + ), + ) + + +def test_hisparse_rejects_pipeline_parallelism(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + with pytest.raises(ValueError, match="pipeline parallelism"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + parallel_config=ParallelConfig(pipeline_parallel_size=2), + ) + + +def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + with pytest.raises(ValueError, match="requires the hybrid KV cache manager"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=True, + ), + ) + + +def test_hisparse_rejects_non_cuda(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + with pytest.raises(ValueError, match="requires NVIDIA CUDA"): + VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" @@ -305,7 +363,8 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): # (warning_once args must be hashable for its lru_cache). monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) config = SimpleNamespace( - model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]) + model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]), + attention_config=AttentionConfig(), ) assert VllmConfig.use_v2_model_runner.fget(config) is False finally: @@ -385,6 +444,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( ) config = SimpleNamespace( model_config=model_config, + attention_config=AttentionConfig(), speculative_config=SimpleNamespace(method="mtp") if with_mtp else None, parallel_config=SimpleNamespace(prefill_context_parallel_size=1), compilation_config=CompilationConfig( @@ -897,7 +957,10 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch): monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) monkeypatch.setattr(current_platform, "is_rocm", lambda: False) - config = SimpleNamespace(model_config=model_config) + config = SimpleNamespace( + model_config=model_config, + attention_config=AttentionConfig(), + ) config._get_v2_model_runner_unsupported_features = lambda: [] assert VllmConfig.use_v2_model_runner.fget(config) is expected @@ -2787,6 +2850,7 @@ def test_draft_sample_method_gumbel_is_rejected(): def _watermarked_vllm_config() -> VllmConfig: config = object.__new__(VllmConfig) config.watermark_config = WatermarkConfig(key=42) + config.attention_config = AttentionConfig() config.speculative_config = None return config diff --git a/tests/v1/attention/test_backend_per_kind.py b/tests/v1/attention/test_backend_per_kind.py index 9d180d3aee53..ab77d588253e 100644 --- a/tests/v1/attention/test_backend_per_kind.py +++ b/tests/v1/attention/test_backend_per_kind.py @@ -1,16 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for per-KV-group attention backend selection (backend_per_kind).""" +"""Tests for attention configuration and per-KV-group backend selection.""" + +from types import SimpleNamespace import pytest -from vllm.config.attention import AttentionConfig +from vllm.config.attention import AttentionConfig, HiSparseConfig from vllm.model_executor.layers.attention.attention import ( _largest_kernel_block_within, ) from vllm.v1.attention.backend import AttentionType, MultipleOf from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import get_attn_spec_kind +from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig from vllm.v1.kv_cache_interface import KVCacheSpecKind @@ -86,3 +89,52 @@ def test_backend_per_kind_rejects_unknown_kind(): def test_backend_per_kind_defaults_empty(): assert AttentionConfig().backend_per_kind == {} + + +def test_hisparse_device_buffer_size_boundaries(): + vllm_config = SimpleNamespace( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + speculative_config=None, + ) + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 256 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=127 + ) + with pytest.raises(ValueError, match="expected at least 128"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=32768 + ) + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 32768 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=32769 + ) + with pytest.raises(ValueError, match="int16 slot-index limit"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + + +def test_hisparse_device_buffer_covers_speculative_window(): + vllm_config = SimpleNamespace( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + speculative_config=SimpleNamespace( + num_speculative_tokens=3, + parallel_drafting=False, + ), + ) + + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 5 * 128 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=4 * 128 - 1 + ) + with pytest.raises(ValueError, match="expected at least 512"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index edc6cbaded32..ee5737a99573 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -200,6 +200,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( layer = SimpleNamespace( kv_cache_dtype=cache_dtype, head_size=576, + indexer=None, non_causal_multi_token_decode=False, sliding_window=None, ) diff --git a/tests/v1/attention/test_mla_context_chunks.py b/tests/v1/attention/test_mla_context_chunks.py index 2ed73b2c987e..2d63b165a055 100644 --- a/tests/v1/attention/test_mla_context_chunks.py +++ b/tests/v1/attention/test_mla_context_chunks.py @@ -55,6 +55,7 @@ def build_chunked_context( def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch): """Sparse MLA must pin computed context lengths before the H2D copy.""" builder = SimpleNamespace( + _context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True), chunked_prefill_workspace=torch.empty((2048, 1)), chunked_prefill_workspace_size=1024, kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE), @@ -113,6 +114,7 @@ def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_si query_lens = [4] * len(context_lens) metadata = build_chunked_context(context_lens, query_lens, workspace_size) assert metadata is not None + assert metadata.context_lens.tolist() == context_lens gathered: dict[int, list[tuple[int, int]]] = {} previous_request_start = -1 @@ -253,6 +255,7 @@ def test_dcp_chunks_fit_the_per_rank_row_budget(): dcp_local_block_size=interleave, ) assert metadata is not None + assert metadata.context_lens.tolist() == context_lens virtual_block_size = interleave * dcp_world_size local_cursor: dict[int, int] = {} @@ -330,3 +333,23 @@ def test_dcp_reorg_uses_each_chunks_local_starts(): ) torch.testing.assert_close(reorganized, torch.cat(expected)) + + +def test_chunked_context_metadata_owns_context_lengths(): + """Reusing CPU scratch must not overwrite an earlier batch's context lengths.""" + context_lens = torch.tensor([16, 32], dtype=torch.int32) + metadata = build_mla_chunked_context_metadata( + context_lens_cpu=context_lens, + prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), + chunked_prefill_workspace=torch.empty((64, 1)), + chunked_prefill_workspace_size=64, + block_size=BLOCK_SIZE, + align_chunk_to_block=True, + device=torch.device("cpu"), + dcp_world_size=1, + dcp_local_block_size=1, + dcp_virtual_block_size=1, + ) + context_lens.zero_() + assert metadata is not None + assert metadata.context_lens.tolist() == [16, 32] diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 425db95a8a7f..e7239320c23a 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3,7 +3,9 @@ """Unit tests for the sparse MLA backends and utilities.""" import math +from collections import deque from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -20,13 +22,19 @@ create_vllm_config, ) from vllm import _custom_ops as ops -from vllm.config import set_current_vllm_config +from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.model_executor.layers.attention import mla_attention from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, SparseMLACommonImpl, + SparseMLAPrefillMetadata, _masked_mha_workspace_fits, _topk_mask_shape, + _use_dense_mha_prefill, ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform @@ -41,7 +49,6 @@ allow_module_level=True, ) -import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse_mod from vllm.model_executor.layers.attention.mla_attention import ( _canonicalize_sparse_mla_kv_cache_dtype, ) @@ -49,10 +56,19 @@ SharedTopkIndicesBuffer, ) from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import current_stream +from vllm.v1.attention.backends.mla import index_group as index_group_module +from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseImpl, +) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, + FlashInferMLASparseMetadataBuilder, FlashInferMLASparseTRTLLMBackend, ) +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( + FlashInferMLASparseSM120Impl, +) from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseBackend, FlashMLASparseImpl, @@ -60,12 +76,29 @@ FlashMLASparseMetadataBuilder, triton_convert_req_index_to_global_index, ) -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, + SparseMLAIndexGroupBuilder, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, +) +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, split_prefill_chunks, ) from vllm.v1.attention.ops import flashmla +from vllm.v1.hisparse import runtime as hisparse_runtime +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + HiSparseRuntime, + ResolvedHiSparseConfig, + _has_hisparse_ops, + build_hisparse_prefill_staging_plan, + hisparse_prefill_staging_remap, +) +from vllm.v1.hisparse.types import SparseKVRowMirror SPARSE_BACKEND_BATCH_SPECS = { name: BATCH_SPECS[name] @@ -113,10 +146,11 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): impl.bmm2_scale = None impl.is_nope_mla = True impl.need_to_return_lse_for_decode = False - monkeypatch.setattr( - flashinfer_sparse_mod, - "triton_convert_req_index_to_global_index", - lambda *args, **kwargs: (topk, torch.ones(1, dtype=torch.int32)), + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.index_group.convert_logical_to_physical_topk.return_value = ( + topk, + torch.ones(1, dtype=torch.int32), ) import flashinfer.decode @@ -142,6 +176,41 @@ def fake_flashinfer(**kwargs): assert recorded_scale != kv_lora_rank**-0.5 +def test_hisparse_routes_prefill_to_sparse_mqa(): + layer = SimpleNamespace(hisparse_cache=object()) + + assert not mla_attention.MLAAttention._use_sparse_mha(layer, SimpleNamespace()) + + +def test_hisparse_metadata_keeps_short_prefill_indexer_enabled(): + config = SimpleNamespace( + attention_config=SimpleNamespace( + hisparse_config=object(), sparse_mla_force_mqa=False + ) + ) + + assert not _use_dense_mha_prefill(config, prefill_max_seq_len=512, topk_tokens=2048) + + +def test_unified_mla_prepares_batch_without_slot_mapping(monkeypatch): + layer = SimpleNamespace( + prepare_kv_cache_update=MagicMock(), + update_kv_cache=MagicMock(), + ) + metadata = object() + monkeypatch.setattr( + mla_attention, + "get_attention_context", + lambda _: (metadata, layer, torch.empty(0), None), + ) + tensor = torch.empty(0, device=DEVICE_TYPE) + + mla_attention.unified_mla_kv_cache_update(tensor, tensor, "layer", "auto", tensor) + + layer.prepare_kv_cache_update.assert_called_once_with(metadata) + layer.update_kv_cache.assert_not_called() + + def _float_to_e8m0_truncate(f: float) -> float: """Simulate SM100's float -> e8m0 -> bf16 scale conversion. e8m0 format only stores the exponent (power of 2). @@ -618,8 +687,6 @@ def test_sparse_backend_decode_correctness( # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. - from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend - prefill_backend = get_mla_prefill_backend(vllm_config)( num_heads=num_heads, scale=scale, @@ -972,6 +1039,8 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): ), block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10), block_size=block_size, + num_decode_tokens=num_mqa_tokens, + num_actual_tokens=num_mqa_tokens, ) assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens @@ -990,11 +1059,30 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): def _stub_kernel(q, kv, indices, lengths, actual_num_heads): captured["indices"] = indices captured["actual_num_heads"] = actual_num_heads - return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device) + return ( + torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device), + None, + ) + + def _convert_topk(indices, metadata, **kwargs): + return triton_convert_req_index_to_global_index( + metadata.req_id_per_token[: indices.shape[0]], + metadata.block_table, + indices, + BLOCK_SIZE=metadata.block_size, + BLOCK_STRIDE_ROWS=kwargs["block_stride_rows"], + NUM_TOPK_TOKENS=indices.shape[1], + return_valid_counts=kwargs["return_valid_counts"], + ) - stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel) + stub_impl = SimpleNamespace( + _bf16_flash_mla_kernel=_stub_kernel, + _convert_logical_to_physical_topk=_convert_topk, + index_group=None, + index_group_index=0, + ) - out = FlashMLASparseImpl._forward_bf16_kv( + out, _ = FlashMLASparseImpl._forward_bf16_kv( stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1] ) @@ -1283,8 +1371,6 @@ def test_sparse_backend_prefill_correctness( # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. - from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend - prefill_backend = get_mla_prefill_backend(vllm_config)( num_heads=num_heads, scale=scale, @@ -1524,189 +1610,1775 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): assert torch.all(result[row, num_valid:] == -1) -def test_flashmla_cache_dtype_aliases_use_ds_layout(): - # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer - # stores kv_cache_dtype, so they cannot bypass the gate. - for alias in ("fp8", "fp8_e4m3"): - assert ( - _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias) - == "fp8_ds_mla" - ) +def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch): + device = torch.device(DEVICE_TYPE) + logical = torch.full((2, 128), -1, dtype=torch.int32, device=device) + logical[:, :4] = torch.tensor( + [[0, 2, 5, -1], [1, 3, 7, -1]], dtype=torch.int32, device=device + ) + builder = SparseMLAIndexGroupBuilder(logical) + group, leader_index = builder.register_layer(True) + follower_group, follower_index = builder.register_layer(False) + assert follower_group is group + + block_table = torch.tensor([[2, 3], [4, 5]], dtype=torch.int32, device=device) + metadata = SimpleNamespace( + block_table=block_table, + block_size=4, + req_id_per_token=torch.tensor([0, 1], dtype=torch.int32, device=device), + num_decode_tokens=2, + num_actual_tokens=2, + ) + calls = [] + convert = index_group_module.triton_convert_req_index_to_global_index + def convert_spy(*args, **kwargs): + calls.append((args, kwargs)) + return convert(*args, **kwargs) -def test_flashmla_fp8_metadata_reuses_common_batch_split(): - builder = SimpleNamespace( - device=torch.device(DEVICE_TYPE), - vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)), + monkeypatch.setattr( + index_group_module, + "triton_convert_req_index_to_global_index", + convert_spy, ) - common_metadata = SimpleNamespace( - num_actual_tokens=1, - seq_lens_cpu_upper_bound=torch.tensor([1]), - query_start_loc_cpu=torch.tensor([0, 1]), - block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + group.set_logical_topk_ready(leader_index) + leader_result = group.convert_logical_to_physical_topk( + leader_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=True, ) - metadata = FlashMLASparseMetadata( - num_reqs=1, - max_query_len=1, - max_seq_len=1, - num_actual_tokens=1, - query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE), - slot_mapping=torch.tensor([0], device=DEVICE_TYPE), - block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), - req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE), - num_decodes=0, - num_prefills=1, - num_decode_tokens=0, + follower_result = group.convert_logical_to_physical_topk( + follower_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=True, ) + torch.accelerator.synchronize() - fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( - builder, common_metadata, metadata + assert len(calls) == 1 + leader_indices, leader_counts = leader_result + follower_indices, follower_counts = follower_result + torch.testing.assert_close(follower_indices, leader_indices) + torch.testing.assert_close(follower_counts, leader_counts) + + +def test_sparse_mla_index_group_falls_back_for_prefill_sized_batch(monkeypatch): + device = torch.device(DEVICE_TYPE) + logical = torch.full((4, 128), -1, dtype=torch.int32, device=device) + logical[:, 0] = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device) + builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2) + group, leader_index = builder.register_layer(True) + follower_group, follower_index = builder.register_layer(False) + assert follower_group is group + + metadata = SimpleNamespace( + block_table=torch.tensor( + [[2], [3], [4], [5]], dtype=torch.int32, device=device + ), + block_size=4, + req_id_per_token=torch.arange(4, dtype=torch.int32, device=device), + ) + calls = 0 + convert = index_group_module.triton_convert_req_index_to_global_index + + def convert_spy(*args, **kwargs): + nonlocal calls + calls += 1 + return convert(*args, **kwargs) + + monkeypatch.setattr( + index_group_module, + "triton_convert_req_index_to_global_index", + convert_spy, + ) + leader_result = group.convert_logical_to_physical_topk( + leader_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + follower_result = group.convert_logical_to_physical_topk( + follower_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=False, ) + torch.accelerator.synchronize() - assert fp8_metadata.num_decodes == 0 - assert fp8_metadata.num_prefills == 1 - assert fp8_metadata.num_decode_tokens == 0 - assert fp8_metadata.num_prefill_tokens == 1 + assert calls == 2 + torch.testing.assert_close(follower_result, leader_result) -def test_flashmla_common_metadata_requires_uniform_decodes(): - common_metadata = SimpleNamespace( - max_query_len=3, - num_reqs=3, - num_actual_tokens=6, - query_start_loc_cpu=torch.tensor([0, 1, 3, 6]), - is_prefilling=None, +def test_sparse_mla_index_groups_own_distinct_physical_buffers(): + logical = torch.empty((8, 128), dtype=torch.int32, device=DEVICE_TYPE) + builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2) + + first, _ = builder.register_layer(True) + follower, _ = builder.register_layer(False) + second, _ = builder.register_layer(True) + + assert follower is first + assert second is not first + assert first.physical_topk_indices[:2].shape == (2, 128) + assert second.physical_topk_indices[:2].shape == (2, 128) + assert first.physical_topk_indices.data_ptr() != ( + second.physical_topk_indices.data_ptr() ) - split = split_decodes_and_prefills( - common_metadata, - decode_threshold=128, - require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes, + +# HiSparse is host-resident-only and kernel-only: runtime construction +# raises without the compiled CUDA ops. +requires_hisparse_ops = pytest.mark.skipif( + not _has_hisparse_ops(), + reason="HiSparse CUDA ops not compiled", +) + + +def _make_hisparse_runtime( + *, + top_k: int = 4, + device_buffer_size: int = 5, + max_num_reqs: int = 2, + row_width: int = 8, + block_size: int = 64, + max_swap_rows: int | None = None, + index_group: hisparse_runtime.HiSparseIndexGroup | None = None, +) -> HiSparseRuntime: + runtime = HiSparseRuntime( + config=ResolvedHiSparseConfig( + top_k=top_k, + device_buffer_size=device_buffer_size, + ), + max_num_reqs=max_num_reqs, + row_width=row_width, + kv_dtype=torch.float32, + device=DEVICE_TYPE, + max_swap_rows=max_swap_rows, + index_group=index_group, + ) + blocks_per_request = cdiv(runtime.region_stride, block_size) + # Leave one extra physical block for resident-tier tests. HMA group layouts + # alias this backing, but distinct group allocations use distinct block IDs. + num_blocks = max_num_reqs * blocks_per_request + 1 + raw = torch.zeros( + num_blocks * block_size * row_width, + dtype=torch.float32, + device=DEVICE_TYPE, + ).view(torch.int8) + block_table = torch.arange( + max_num_reqs * blocks_per_request, + dtype=torch.int32, + device=DEVICE_TYPE, + ).view(max_num_reqs, blocks_per_request) + runtime.bind_hot_cache( + raw, + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=num_blocks, + block_size=block_size, + block_table=block_table, + ) + runtime.request_state_indices = torch.arange( + max_num_reqs, dtype=torch.int32, device=runtime.device + ) + return runtime + + +def _make_hisparse_cache_handle( + *, + top_k: int = 4, + device_buffer_size: int = 5, + max_num_reqs: int = 2, + row_width: int = 8, + block_size: int = 64, + max_swap_rows: int | None = None, +) -> HiSparseCacheHandle: + runtime = _make_hisparse_runtime( + top_k=top_k, + device_buffer_size=device_buffer_size, + max_num_reqs=max_num_reqs, + row_width=row_width, + block_size=block_size, + max_swap_rows=max_swap_rows, ) + return HiSparseCacheHandle(runtime) - assert split == (1, 2, 1, 5) +@requires_hisparse_ops +def test_hisparse_uses_graph_stable_request_state_mapping(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + runtime = _make_hisparse_runtime( + top_k=1, + device_buffer_size=2, + max_num_reqs=2, + row_width=row_width, + block_size=block_size, + ) + runtime.bind_source_cache( + torch.arange(2 * block_size * row_width, dtype=torch.float32) + .view(2, block_size, row_width) + .pin_memory() + ) + runtime.request_state_indices = torch.tensor([1], dtype=torch.int32, device=device) -def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch): - monkeypatch.setattr( - "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata", - lambda: (object(), None), + cache_handle = HiSparseCacheHandle(runtime) + cache_handle.all_context_pages_resident = False + runtime.begin_forward() + cache_handle.swap_in( + req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device), + block_table=torch.tensor([[1]], dtype=torch.int32, device=device), + logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device), + block_size=block_size, ) - builder = SimpleNamespace( - device=torch.device(DEVICE_TYPE), - dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE), - max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE), + torch.accelerator.synchronize() + + assert (runtime.index_group.device_global_indices[0] == -1).all() + assert (runtime.index_group.device_global_indices[1] == block_size).any() + + +@requires_hisparse_ops +def test_hisparse_maps_speculative_rows_through_request_state(): + """Multiple verification rows for one request share its persistent state.""" + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + runtime = _make_hisparse_runtime( + top_k=1, + device_buffer_size=2, + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + max_swap_rows=4, ) - query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660]) - common_metadata = SimpleNamespace( - num_actual_tokens=660, - query_start_loc_cpu=query_start_loc_cpu, - seq_lens=torch.arange(7, device=DEVICE_TYPE), + runtime.bind_source_cache( + torch.zeros(1, block_size, row_width, dtype=torch.float32).pin_memory() ) - metadata = FlashMLASparseMetadata( - num_reqs=7, - max_query_len=110, - max_seq_len=110, - num_actual_tokens=660, - query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE), - slot_mapping=torch.arange(660, device=DEVICE_TYPE), - block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE), - req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE), - num_decodes=7, - num_prefills=0, - num_decode_tokens=660, + + cache_handle = HiSparseCacheHandle(runtime) + resident_table = torch.tensor([[1]], dtype=torch.int32, device=device) + cache_handle.bind_cache( + runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=resident_table, + slot_mapping=torch.arange(4, dtype=torch.int64, device=device), + ) + cache_handle.all_context_pages_resident = True + runtime.begin_forward() + physical = cache_handle.swap_in( + req_id_per_token=torch.zeros(4, dtype=torch.int32, device=device), + block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), + logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device), + block_size=block_size, ) + torch.accelerator.synchronize() - fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( - builder, common_metadata, metadata + assert physical.tolist() == [[block_size]] * 4 + + +@requires_hisparse_ops +def test_hisparse_resident_rows_bypass_hot_lru(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + cache_handle = _make_hisparse_cache_handle( + top_k=4, + device_buffer_size=5, + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + ) + raw = cache_handle.runtime.hot.cache.view(torch.int8) + resident_table = torch.tensor([[1]], dtype=torch.int32, device=device) + resident_slots = torch.tensor([block_size], dtype=torch.int64, device=device) + cache_handle.bind_cache( + raw, + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=resident_table, + slot_mapping=resident_slots, + ) + + host = torch.randn(1, block_size, row_width).pin_memory() + cache_handle.runtime.bind_source_cache(host) + resident_rows = host[0].to(device).add(7) + cache_handle.runtime.hot.cache[1].copy_(resident_rows) + source_table = torch.tensor([[5]], dtype=torch.int32, device=device) + request_ids = torch.zeros(1, dtype=torch.int32, device=device) + topk = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=device) + lru_before = cache_handle.runtime.index_group.lru_slots.clone() + + cache_handle.runtime.begin_forward() + indices, valid_counts = cache_handle.swap_in( + req_id_per_token=request_ids, + block_table=source_table, + logical_topk_indices=topk, + block_size=block_size, + return_valid_counts=True, ) + torch.accelerator.synchronize() - assert fp8_metadata.num_decodes == 6 - assert fp8_metadata.num_decode_tokens == 660 - assert fp8_metadata.decode is not None - assert fp8_metadata.decode.decode_query_len == 110 + expected_indices = topk + block_size + torch.testing.assert_close(indices, expected_indices) torch.testing.assert_close( - fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE) + valid_counts, torch.tensor([topk.shape[1]], dtype=torch.int32, device=device) + ) + torch.testing.assert_close(cache_handle.runtime.index_group.lru_slots, lru_before) + torch.testing.assert_close( + cache_handle.runtime.hot.attention_cache.view(-1, row_width)[ + indices.to(torch.long) + ], + resident_rows[topk.to(torch.long)], ) -@pytest.mark.parametrize("use_mixed_batch", [False, True]) -def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool): - num_decode_tokens = 2 - num_batch_tokens = 5 - q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE) - topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE) - kernel_q_shapes = [] - - def convert_indices(*args, **kwargs): # noqa: ARG001 - assert not kwargs.get("HAS_PREFILL_WORKSPACE", False) - if not kwargs.get("return_valid_counts", False): - return topk_indices - valid_counts = torch.full( - (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE +@requires_hisparse_ops +def test_hisparse_swap_in_preserves_rows_across_eviction(): + device = torch.device(DEVICE_TYPE) + block_size = 64 + row_width = 64 + num_blocks = 64 + top_k = 128 + buf = 256 + num_reqs = 4 + + kv_pool = torch.randn( + (num_blocks + 1, block_size, row_width), dtype=torch.float32 + ).pin_memory() + flat_pool = kv_pool.reshape(-1, row_width) + + runtime = _make_hisparse_runtime( + top_k=top_k, + device_buffer_size=buf, + max_num_reqs=num_reqs, + row_width=row_width, + ) + runtime.bind_source_cache(kv_pool) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + cache.runtime.begin_forward() + + blocks_per_req = num_blocks // num_reqs + block_table = torch.arange( + 1, num_blocks + 1, dtype=torch.int32, device=device + ).view(num_reqs, blocks_per_req) + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + seq_len = blocks_per_req * block_size + base = torch.arange(top_k, dtype=torch.int32, device=device) + for step in range(3): + cache.runtime.begin_forward() + topk = torch.stack( + [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)] ) - return topk_indices, valid_counts + topk[:, -1] = -1 + + hot_indices, valid_counts = cache.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk.clone(), + block_size=block_size, + return_valid_counts=True, + ) + torch.accelerator.synchronize() - monkeypatch.setattr( - "vllm.v1.attention.backends.mla.flashmla_sparse." - "triton_convert_req_index_to_global_index", - convert_indices, - ) + global_ref = _triton_convert_reference_impl( + req_ids, block_table, topk, block_size, top_k + ) + valid = hot_indices >= 0 + torch.testing.assert_close( + valid_counts, + (global_ref >= 0).sum(dim=1, dtype=torch.int32), + rtol=0, + atol=0, + ) + gathered = cache.runtime.hot.attention_cache.reshape(-1, row_width)[ + hot_indices[valid].to(torch.long) + ].cpu() + expected = flat_pool[global_ref[valid].cpu().to(torch.long)] + torch.testing.assert_close(gathered, expected) - def run_kernel(**kwargs): - kernel_q_shapes.append(kwargs["q"].shape) - return kwargs["q"][..., :1], None - if use_mixed_batch: - fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata( - scheduler_metadata=object(), # type: ignore[arg-type] - dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), - cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), +@requires_hisparse_ops +def test_hisparse_multi_step_swaps_match_independent(): + """The shared LRU must retain every speculative step until followers run.""" + device = torch.device(DEVICE_TYPE) + torch.manual_seed(0) + block_size = 64 + row_width = 64 + num_blocks = 64 + top_k = 128 + buf = 3 * top_k + num_reqs = 4 + + kv_pool = torch.randn( + (num_blocks, block_size, row_width), dtype=torch.float32 + ).pin_memory() + + def make( + index_group: hisparse_runtime.HiSparseIndexGroup | None = None, + ) -> HiSparseCacheHandle: + runtime = _make_hisparse_runtime( + top_k=top_k, + max_num_reqs=num_reqs, + device_buffer_size=buf, + row_width=row_width, + block_size=block_size, + max_swap_rows=2 * num_reqs, + index_group=index_group, ) - else: - FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode - fp8_metadata = FP8Meta( - num_decodes=1, - num_prefills=1, - num_decode_tokens=num_decode_tokens, - num_prefill_tokens=num_batch_tokens - num_decode_tokens, - decode=FP8Meta.Decode( - seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), - kernel_metadata=object(), # type: ignore[arg-type] - decode_query_len=num_decode_tokens, - ), - prefill=FP8Meta.Prefill( - request_ids=torch.empty( - num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE - ), - workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), - chunks=[], - ), + runtime.bind_source_cache(kv_pool) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + return cache + + blocks_per_req = num_blocks // num_reqs + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view( + num_reqs, blocks_per_req + ) + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + seq_len = blocks_per_req * block_size + producer = make() + shared = make(producer.runtime.index_group) + indep = make() + base = torch.arange(top_k, dtype=torch.int32, device=device) + for iteration in range(2): + producer.runtime.begin_forward() + shared.runtime.begin_forward() + indep.runtime.begin_forward() + topks = [ + torch.stack( + [ + (base + (iteration * 2 + step) * top_k + row * 17) % seq_len + for row in range(num_reqs) + ] + ) + for step in range(2) + ] + kw = dict( + req_id_per_token=req_ids, + block_table=block_table, + block_size=block_size, ) - metadata = SimpleNamespace( - fp8_extra_metadata=fp8_metadata, - fp8_use_mixed_batch=use_mixed_batch, - num_actual_tokens=num_batch_tokens, - req_id_per_token=torch.empty( - num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE - ), - block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), - block_size=64, + producer_indices = [ + producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + independent_indices = [ + indep.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + follower_indices = [ + shared.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + torch.accelerator.synchronize() + + for idx_full, idx_indep, idx_shared in zip( + producer_indices, independent_indices, follower_indices + ): + torch.testing.assert_close(idx_shared, idx_full, rtol=0, atol=0) + torch.testing.assert_close(idx_indep, idx_full, rtol=0, atol=0) + torch.testing.assert_close(shared.runtime.hot.cache, indep.runtime.hot.cache) + + shared_rows = shared.runtime.hot.attention_cache.reshape(-1, row_width) + host_rows = kv_pool.reshape(-1, row_width) + for physical_rows, logical_rows in zip(follower_indices, topks): + global_rows = _triton_convert_reference_impl( + req_ids, block_table, logical_rows, block_size, top_k + ) + valid = physical_rows >= 0 + torch.testing.assert_close( + shared_rows[physical_rows[valid].to(torch.long)].cpu(), + host_rows[global_rows[valid].cpu().to(torch.long)], + ) + + +@requires_hisparse_ops +def test_hisparse_multi_step_writes_request_major_output(): + device = torch.device(DEVICE_TYPE) + block_size = 64 + row_width = 8 + num_reqs = 2 + query_len = 3 + top_k = 4 + cache = _make_hisparse_cache_handle( + top_k=top_k, + device_buffer_size=query_len * top_k, + max_num_reqs=num_reqs, + row_width=row_width, + block_size=block_size, + max_swap_rows=num_reqs * query_len, + ) + host = torch.arange( + (num_reqs + 1) * block_size * row_width, dtype=torch.float32 + ).view(num_reqs + 1, block_size, row_width) + cache.runtime.bind_source_cache(host.pin_memory()) + block_table = torch.arange(1, num_reqs + 1, dtype=torch.int32, device=device).view( + -1, 1 + ) + request_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + logical = torch.tensor( + [ + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], + [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], + ], + dtype=torch.int32, + device=device, ) - impl = SimpleNamespace( - kv_cache_dtype="fp8_ds_mla", - topk_indices_buffer=topk_indices, - num_heads=2, - kv_lora_rank=1, - dcp_world_size=1, - need_to_return_lse_for_decode=False, - _fp8_flash_mla_kernel=run_kernel, + physical = torch.full_like(logical, -1) + valid_counts = torch.zeros((num_reqs, query_len), dtype=torch.int32, device=device) + + cache.runtime.begin_forward() + for step in range(query_len): + cache.swap_in( + request_ids, + block_table, + logical[:, step], + block_size=block_size, + return_valid_counts=True, + attention_indices_out=physical[:, step], + valid_counts_out=valid_counts[:, step], + ) + torch.accelerator.synchronize() + + torch.testing.assert_close( + valid_counts, + torch.full_like(valid_counts, top_k), ) - impl._forward_fp8_kv_mixed_batch = MethodType( - FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl + global_rows = _triton_convert_reference_impl( + torch.arange(num_reqs, device=device, dtype=torch.int32).repeat_interleave( + query_len + ), + block_table, + logical.view(num_reqs * query_len, top_k), + block_size, + top_k, ) - impl._forward_fp8_kv_separate_prefill_decode = MethodType( - FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl + hot_rows = cache.runtime.hot.attention_cache.reshape(-1, row_width) + host_rows = host.reshape(-1, row_width) + torch.testing.assert_close( + hot_rows[physical.view(num_reqs * query_len, top_k).to(torch.long)].cpu(), + host_rows[global_rows.cpu().to(torch.long)], + ) + + +@requires_hisparse_ops +def test_hisparse_kv_update_writes_resident_and_staging_caches(): + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + cache_handle = _make_hisparse_cache_handle( + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + ) + resident_slots = torch.tensor([4, 5, -1], dtype=torch.int64, device=device) + cache_handle.bind_cache( + cache_handle.runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=cache_handle.runtime.hot.cache.shape[0], + block_size=block_size, + block_table=torch.tensor([[1]], dtype=torch.int32, device=device), + slot_mapping=resident_slots, + ) + cache_handle.mirror_staging_cache = torch.empty( + (1, block_size, row_width), dtype=torch.float32, device=device + ) + cache_handle.mirror_staging_slots = torch.arange( + block_size, dtype=torch.int64, device=device + ) + slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device) + kv_c = torch.randn(8, row_width - 2, device=device) + k_pe = torch.randn(8, 1, 2, device=device) + cache_handle.num_actual_tokens = slots.numel() + cache_handle.decode_batch = False + cache_handle.host_mirror_required = True + source_cache = torch.zeros_like(cache_handle.view.cache) + impl = object.__new__(FlashMLASparseImpl) + layer = SimpleNamespace( + hisparse_cache=cache_handle, + use_pcp=False, + impl=impl, + ) + mla_attention.MLAAttention.update_kv_cache( + layer, + kv_c, + k_pe, + source_cache, + slots, + SimpleNamespace(num_decode_tokens=0), + "auto", + torch.tensor(1.0, device=device), + ) + torch.accelerator.synchronize() + + expected = torch.cat([kv_c[:2], k_pe[:2, 0]], dim=-1).cpu() + torch.testing.assert_close( + cache_handle.view.cache.view(-1, row_width)[resident_slots[:2]], + expected.to(device), + ) + torch.testing.assert_close(source_cache, torch.zeros_like(source_cache)) + staged = cache_handle.mirror_staging_cache.view(-1, row_width) + staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1) + torch.testing.assert_close(staged[:3], staged_expected) + + +@requires_hisparse_ops +def test_hisparse_row_dma_copies_discontiguous_spans_across_layers(): + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + num_rows = 16 + resident_caches = tuple( + ( + torch.arange(num_rows * row_width, dtype=torch.uint8, device=device) + .add_(layer_index * 17) + .view(num_rows // block_size, block_size, row_width) + ) + for layer_index in range(2) + ) + host_caches = tuple( + torch.zeros((num_rows, row_width), dtype=torch.uint8, pin_memory=True) + for _ in resident_caches + ) + worker = object.__new__(HiSparseConnectorWorker) + worker.is_host_writer = True + worker.kernel_block_size = block_size + worker.resident_caches = resident_caches + worker.host_caches = host_caches + worker.cache_handles = [ + SimpleNamespace( + decode_batch=True, + runtime=SimpleNamespace(resident_source_index=0), + ) + for _ in resident_caches + ] + worker.dma_stream = torch.cuda.Stream(device=device) + worker.host_write_event = torch.Event() + worker._dma_free_descriptors = [] + worker._pending_dma_descriptors = deque() + worker._pending_transfer_events = deque() + worker._enqueued_transfer_ids = [] + worker._dma_submitted = False + worker._set_row_mirrors( + ( + SparseKVRowMirror((1,), 5, 2), + SparseKVRowMirror((8,), 10, 4), + ) + ) + + worker._enqueue_row_dma(range(2)) + current_stream().wait_event(worker.host_write_event) + torch.accelerator.synchronize() + + for resident, host in zip(resident_caches, host_caches): + flat_resident = resident.view(num_rows, row_width).cpu() + torch.testing.assert_close(host[5:7], flat_resident[1:3]) + torch.testing.assert_close(host[10:14], flat_resident[8:12]) + + +@requires_hisparse_ops +def test_hisparse_remaps_strided_hma_rows_for_attention(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 4, 8 + runtime = _make_hisparse_runtime( + max_num_reqs=1, + block_size=block_size, + row_width=row_width, + ) + page_elements = block_size * row_width + stride_elements = 3 * page_elements + byte_offset = page_elements * torch.float32.itemsize + raw = torch.zeros(2 * stride_elements, dtype=torch.float32, device=device).view( + torch.int8 + ) + runtime.bind_hot_cache( + raw, + byte_offset=byte_offset, + block_stride=stride_elements * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=torch.tensor([[0, 1]], dtype=torch.int32, device=device), + ) + + source = ( + torch.arange(36 * row_width, dtype=torch.float32) + .view(9, block_size, row_width) + .pin_memory() + ) + runtime.bind_source_cache(source) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + runtime.begin_forward() + attention_indices = cache.swap_in( + req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device), + block_table=torch.arange(1, 9, dtype=torch.int32, device=device).view(1, 8), + logical_topk_indices=torch.tensor( + [[0, 1, 2, 3]], dtype=torch.int32, device=device + ), + block_size=block_size, + ) + + assert runtime.hot.attention_cache is not None + assert runtime.hot.attention_cache.is_contiguous() + physical_indices = runtime.index_group.shared_topk.device_topk_rows[0] + expected = ( + physical_indices // block_size * (stride_elements // row_width) + + physical_indices % block_size + ) + torch.testing.assert_close(attention_indices[0], expected) + + +@requires_hisparse_ops +def test_hisparse_gather_zeroes_unaligned_destination(): + """Invalid host rows are safely zeroed even at a byte-unaligned address.""" + device = torch.device(DEVICE_TYPE) + row_bytes = 16 + host_cache = torch.ones(1, row_bytes, dtype=torch.uint8).pin_memory() + storage = torch.full((2 * row_bytes + 1,), 7, dtype=torch.uint8, device=device) + hot_cache = torch.as_strided( + storage, + size=(2, 1, row_bytes), + stride=(row_bytes + 1, row_bytes, 1), + ) + global_indices = torch.tensor([[1]], dtype=torch.int32, device=device) + hot_indices = torch.tensor([[1]], dtype=torch.int32, device=device) + miss_mask = torch.ones_like(global_indices) + + torch.ops._C_cache_ops.hisparse_gather_plan( + host_cache, + hot_cache, + global_indices, + hot_indices, + miss_mask, + None, + ) + torch.accelerator.synchronize() + + torch.testing.assert_close(hot_cache[1, 0], torch.zeros_like(hot_cache[1, 0])) + assert storage[row_bytes].item() == 7 + + +@requires_hisparse_ops +def test_hisparse_gather_rejects_short_attention_block_stride(): + """Reject a stride that would let the gather kernel cross block bounds.""" + device = torch.device(DEVICE_TYPE) + host_cache = torch.ones(1, 16, dtype=torch.uint8).pin_memory() + hot_cache = torch.empty((1, 2, 16), dtype=torch.uint8, device=device) + global_indices = torch.zeros((1, 1), dtype=torch.int32, device=device) + hot_indices = torch.zeros_like(global_indices) + miss_mask = torch.ones_like(global_indices) + attention_indices = torch.empty_like(global_indices) + + with pytest.raises(RuntimeError, match="stride must cover one hot block"): + torch.ops._C_cache_ops.hisparse_gather_plan( + host_cache, + hot_cache, + global_indices, + hot_indices, + miss_mask, + None, + attention_indices, + 1, + ) + + +@requires_hisparse_ops +def test_hisparse_newest_write_and_recycled_slot_invalidation(): + """Newest writes clamp padding and invalidated rows are loaded again.""" + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + num_blocks = 8 + + kv_pool = ( + torch.arange(num_blocks * block_size * row_width, dtype=torch.float32) + .view(num_blocks, block_size, row_width) + .pin_memory() + ) + flat_pool = kv_pool.reshape(-1, row_width) + + cache_handle = _make_hisparse_cache_handle(block_size=block_size) + cache_handle.runtime.eager_host_mirror = True + + block_table = torch.tensor([[2, 0, 4]], dtype=torch.int32, device=device) + resident_slots = torch.tensor([4, -1], dtype=torch.int64, device=device) + cache_handle.bind_cache( + cache_handle.runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=cache_handle.runtime.hot.cache.shape[0], + block_size=block_size, + block_table=torch.tensor([[0, 0, 1]], dtype=torch.int32, device=device), + slot_mapping=resident_slots, + ) + req_ids = torch.tensor([0], dtype=torch.int32, device=device) + newest_global = 4 * block_size + padded_global = 1 + padded_row = flat_pool[padded_global].clone() + slot_mapping = torch.tensor( + [newest_global, padded_global], dtype=torch.int64, device=device + ) + request_state_indices = torch.tensor([0, -1], dtype=torch.int32, device=device) + cache_handle.runtime.request_state_indices = request_state_indices + + kv_c = torch.randn(3, row_width - 2, device=device) + k_pe = torch.randn(3, 1, 2, device=device) + cache_handle.runtime.bind_source_cache(kv_pool) + cache_handle.num_actual_tokens = slot_mapping.numel() + ops.concat_and_cache_mla( + kv_c[:1], + k_pe[:1, 0], + cache_handle.view.cache, + resident_slots[:1], + kv_cache_dtype="auto", + scale=torch.tensor(1.0, device=device), + ) + torch.accelerator.synchronize() + expected_row = torch.cat([kv_c[0], k_pe[0, 0]]).cpu() + flat_pool[newest_global].copy_(expected_row) + torch.testing.assert_close(flat_pool[newest_global], expected_row) + torch.testing.assert_close(flat_pool[padded_global], padded_row) + flat_hot = cache_handle.runtime.hot.cache.view(-1, row_width) + torch.testing.assert_close(flat_hot[resident_slots[0]].cpu(), expected_row) + + topk = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32, device=device) + cache_handle.all_context_pages_resident = False + cache_handle.runtime.begin_forward() + hot_indices = cache_handle.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk, + block_size=block_size, + ) + torch.accelerator.synchronize() + stale_hot_slot = hot_indices[0, 0].item() + stale_row = flat_hot[stale_hot_slot].clone() + + flat_pool[8] += 1000 + cache_handle.runtime.invalidate_written_slots( + torch.tensor([8], dtype=torch.int64, device=device), req_ids + ) + cache_handle.runtime.begin_forward() + hot_indices = cache_handle.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk, + block_size=block_size, + ) + torch.accelerator.synchronize() + idx = hot_indices.cpu().tolist()[0][0] + assert not torch.equal(flat_hot[idx], stale_row) + torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8]) + + +@requires_hisparse_ops +def test_hisparse_mixed_batch_bf16_row_split( + default_vllm_config, dist_init, workspace_init +): + """Host-resident mixed batch on the bf16 path is row-split. + + Two long-context speculative-decode requests + one short local-prefill + chunk: every decode step must be served from the bounded hot buffer before + it is reused, while only the prefill rows' blocks are staged host->GPU. + """ + ok, reason = flashmla.is_flashmla_sparse_supported() + if not ok: + pytest.skip(reason) + + device = torch.device(DEVICE_TYPE) + dtype = torch.bfloat16 + torch.manual_seed(0) + + num_heads = 64 + kv_lora_rank = 512 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + v_head_dim = 128 + head_size = kv_lora_rank + qk_rope_head_dim + topk_tokens = 128 + block_size = 64 + + # Long decode contexts + a short prefill chunk (router shortcut shape). + batch_spec = BatchSpec(seq_lens=[2048, 2048, 192], query_lens=[2, 2, 64]) + max_seqlen = max(batch_spec.seq_lens) + total_cache_tokens = sum(batch_spec.seq_lens) + total_tokens = batch_spec.compute_num_tokens() + + vllm_config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat", + tensor_parallel_size=1, + max_model_len=max_seqlen, + num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1), + block_size=block_size, + hf_config_override={ + "index_topk": topk_tokens, + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}], + }, + ) + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=2 * topk_tokens, + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=1 + ) + model_config = vllm_config.model_config + model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + model_type="deepseek_v2", + ) + model_config.dtype = dtype + model_config.get_num_attention_heads = MethodType( + lambda self, parallel_config: num_heads, model_config + ) + model_config.get_num_kv_heads = MethodType( + lambda self, parallel_config: 1, model_config + ) + model_config.get_head_size = MethodType(lambda self: head_size, model_config) + model_config.get_sliding_window = MethodType(lambda self: None, model_config) + # create_hisparse_cache_handle sizes hot buffers per layer. + model_config.get_num_layers = MethodType( + lambda self, parallel_config: 1, model_config + ) + + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + common_attn_metadata = create_common_attn_metadata( + batch_spec, + vllm_config.cache_config.block_size, + device, + arange_block_indices=True, + ) + # Prepopulate every position of every sequence so the forward needs no + # KV-cache update (the row split itself is what is under test). + kv_c_contexts = [ + torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device) + for s_len in batch_spec.seq_lens + ] + k_pe_contexts = [ + torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device) + for s_len in batch_spec.seq_lens + ] + kv_cache = create_and_prepopulate_kv_cache( + kv_c_contexts=kv_c_contexts, + k_pe_contexts=k_pe_contexts, + block_size=vllm_config.cache_config.block_size, + head_size=head_size, + dtype=dtype, + device=device, + num_blocks=vllm_config.cache_config.num_gpu_blocks, + common_attn_metadata=common_attn_metadata, + randomize_blocks=False, + kv_cache_dtype="auto", + ) + + prefill_backend = get_mla_prefill_backend(vllm_config)( + num_heads=num_heads, + scale=1.0 / math.sqrt(head_size), + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + ) + vllm_config.compilation_config.static_forward_context["placeholder"] = ( + SimpleNamespace(prefill_backend=prefill_backend) + ) + + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + metadata = builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + assert isinstance(metadata.prefill, SparseMLAPrefillMetadata) + + # Per-token sparse indices bounded by each token's position, with unique + # offsets and -1 padding (same construction as the decode parity test). + positions: list[int] = [] + for i in range(batch_spec.batch_size): + ctx_len = batch_spec.seq_lens[i] - batch_spec.query_lens[i] + positions.extend(ctx_len + q_idx for q_idx in range(batch_spec.query_lens[i])) + sparse_indices = torch.empty( + total_tokens, topk_tokens, dtype=torch.int32, device=device + ) + for tok_idx in range(total_tokens): + max_valid_idx = positions[tok_idx] + offset = tok_idx * 7 # Prime number for varied offsets + num_valid = min(topk_tokens // 2, max_valid_idx + 1) + valid_range = torch.arange(num_valid, device=device, dtype=torch.int32) + sparse_indices[tok_idx] = torch.cat( + [ + (valid_range + offset) % (max_valid_idx + 1), + torch.full( + (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32 + ), + ] + ) + + q = torch.rand(total_tokens, num_heads, head_size, dtype=dtype, device=device) + mock_indexer = SimpleNamespace( + topk_indices_buffer=sparse_indices, topk_tokens=topk_tokens + ) + + impl_cls = FlashMLASparseBackend.get_impl_cls() + with set_current_vllm_config(vllm_config): + impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + scale=1.0 / math.sqrt(head_size), + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_b_proj=None, + indexer=mock_indexer, + ) + assert isinstance(impl.index_group, HiSparseMLAIndexGroup) + cache_handle = impl.index_group.cache(impl.index_group_index) + blocks_per_request = cdiv(cache_handle.runtime.region_stride, block_size) + num_hot_blocks = cache_handle.runtime.max_num_reqs * blocks_per_request + hot_cache = torch.zeros( + num_hot_blocks * block_size * head_size, + dtype=dtype, + device=device, + ).view(torch.int8) + cache_handle.runtime.bind_hot_cache( + hot_cache, + byte_offset=0, + block_stride=block_size * head_size * dtype.itemsize, + num_blocks=num_hot_blocks, + block_size=block_size, + block_table=torch.arange(num_hot_blocks, dtype=torch.int32, device=device).view( + cache_handle.runtime.max_num_reqs, blocks_per_request + ), + ) + cache_handle.runtime.request_state_indices = torch.arange( + cache_handle.runtime.max_num_reqs, dtype=torch.int32, device=device + ) + impl.prepare_for_batch(metadata) + + # Device-resident reference: the whole batch converted against the full + # block table and run as one kernel call over the GPU cache. + ref_topk, ref_topk_length = triton_convert_req_index_to_global_index( + metadata.req_id_per_token, + metadata.block_table, + sparse_indices, + BLOCK_SIZE=metadata.block_size, + NUM_TOPK_TOKENS=topk_tokens, + return_valid_counts=True, + ) + reference, _ = impl._bf16_flash_mla_kernel(q, kv_cache, ref_topk, ref_topk_length) + + # Host-resident pool with identical contents. + kv_pool = kv_cache.squeeze(1).cpu().pin_memory() + cache_handle.runtime.bind_source_cache(kv_pool) + cache_handle.source_block_table = metadata.block_table + + staging_calls = [] + original_gather = cache_handle.runtime.gather_prefill_cache + + def spy_gather(self, kv, plan, resident_cache=None): + staged = original_gather(kv, plan, resident_cache) + staging_calls.append((plan, staged.shape)) + return staged + + cache_handle.runtime.gather_prefill_cache = MethodType( + spy_gather, cache_handle.runtime + ) + + backend_output, _ = impl._forward_bf16_kv( + q, kv_pool, sparse_indices, metadata, q.shape[1] + ) + torch.accelerator.synchronize() + + # Only the prefill rows' blocks were staged: the decode rows' 2048-token + # contexts (32 blocks each) must stay off the staging gather. + assert len(staging_calls) == 1 + plan, staged_shape = staging_calls[0] + assert plan is metadata.prefill.host_staging_plan + prefill_blocks = cdiv(batch_spec.seq_lens[-1], block_size) + assert staged_shape[0] <= prefill_blocks + 1 # +1: block-0 tail padding + + torch.testing.assert_close(backend_output, reference, rtol=0.01, atol=0.01) + + +def test_hisparse_prefill_staging_remap(): + """Compacted staging references the same host rows as direct indexing.""" + block_size = 4 + block_table = torch.tensor( + [[5, 2, -1, -1], [2, 7, 3, -1], [9, -1, -1, -1]], dtype=torch.int32 + ) + + new_bt, row_ids = hisparse_prefill_staging_remap(block_table, block_size) + + assert new_bt.shape == block_table.shape + assert row_ids.dtype == torch.int32 and new_bt.dtype == torch.int32 + n_unique = len({0, 2, 3, 5, 7, 9}) # -1 clamps to block 0 + flat_rows = row_ids.flatten() + assert (flat_rows >= 0).sum() == n_unique * block_size + for i in range(block_table.shape[0]): + for j in range(block_table.shape[1]): + orig = max(int(block_table[i, j]), 0) + staged = int(new_bt[i, j]) + for k in range(block_size): + assert int(flat_rows[staged * block_size + k]) == orig * block_size + k + + +def test_hisparse_prefill_staging_plan_masks_unused_blocks(): + block_table = torch.tensor([[5, 2, 7], [9, 3, 4]], dtype=torch.int32) + plan = build_hisparse_prefill_staging_plan( + block_table, + seq_lens=torch.tensor([5, 8], dtype=torch.int32), + block_size=4, + staging_block_capacity=4, + ) + + assert (plan.block_table[:, 2] == 0).all() + valid_rows = plan.row_ids[plan.row_ids >= 0] + assert set((valid_rows[::4] // 4).tolist()) == {0, 2, 3, 5, 9} + for row, used in enumerate((2, 2)): + for column in range(used): + staged_block = int(plan.block_table[row, column]) + host_block = int(block_table[row, column]) + assert int(valid_rows[staged_block * 4]) == host_block * 4 + + +def test_hisparse_prefill_staging_plan_resolves_resident_sources(): + """Per-page resident hits become device sources; misses stay host DMAs.""" + block_size = 4 + resident_block_size = 2 + block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32) + seq_lens = torch.tensor([5, 8], dtype=torch.int32) + plan = build_hisparse_prefill_staging_plan( + block_table, seq_lens, block_size, staging_block_capacity=4 + ) + # Two resident pages per host block; 0 entries are null (not resident). + resident_table = torch.tensor( + [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32 + ) + + plan.ensure_gpu_sources(resident_table, resident_block_size) + + assert plan.gpu_row_ids is not None + unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() + gpu_rows = plan.gpu_row_ids[0].view(-1, block_size) + reps = {5: (0, 0), 2: (0, 1), 9: (1, 0), 3: (1, 1)} + for u, host_id in enumerate(unique_hosts): + for t in range(block_size): + if host_id < 0: + assert int(gpu_rows[u, t]) == -1 + continue + if host_id == 0: + assert int(gpu_rows[u, t]) == -1 + continue + row, col = reps[host_id] + page = t // resident_block_size + res_block = int(resident_table[row, col * 2 + page]) + expected = ( + res_block * resident_block_size + t % resident_block_size + if res_block > 0 + else -1 + ) + assert int(gpu_rows[u, t]) == expected + valid_rows = plan.row_ids >= 0 + torch.testing.assert_close( + plan.miss_mask, + ((plan.gpu_row_ids < 0) & valid_rows).int(), + check_dtype=False, + ) + + plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size) + + assert plan.gpu_row_ids is not None + assert (plan.gpu_row_ids == -1).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_hisparse_gather_prefill_cache_prefers_resident_rows(): + """Staged rows come from the resident cache when a shadow page exists.""" + if not _has_hisparse_ops(): + pytest.skip("hisparse CUDA ops unavailable") + device = torch.device("cuda") + block_size, resident_block_size, row_width = 4, 2, 16 + block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device) + seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device) + plan = build_hisparse_prefill_staging_plan( + block_table, seq_lens, block_size, staging_block_capacity=4 + ) + resident_table = torch.tensor( + [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], + dtype=torch.int32, + device=device, + ) + plan.ensure_gpu_sources(resident_table, resident_block_size) + + num_host_blocks, num_res_blocks = 10, 24 + host_cache = ( + torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32) + .view(num_host_blocks, block_size, row_width) + .pin_memory() + ) + resident_cache = ( + ( + -torch.arange( + num_res_blocks * resident_block_size * row_width, dtype=torch.float32 + ) + - 1.0 + ) + .view(num_res_blocks, resident_block_size, row_width) + .to(device) + ) + + staged = HiSparseRuntime.gather_prefill_cache( + None, host_cache, plan, resident_cache=resident_cache + ) + + staged_flat = staged.view(-1, row_width).cpu() + host_flat = host_cache.view(-1, row_width) + resident_flat = resident_cache.reshape(-1, row_width).cpu() + gpu_rows = plan.gpu_row_ids[0].cpu() + host_rows = plan.row_ids[0].cpu() + for i in range(staged_flat.shape[0]): + if int(host_rows[i]) < 0: + assert int(plan.miss_mask[0, i]) == 0 + continue + expected = ( + resident_flat[int(gpu_rows[i])] + if int(gpu_rows[i]) >= 0 + else host_flat[int(host_rows[i])] + ) + torch.testing.assert_close(staged_flat[i], expected) + + +def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch): + """FP8 decode must use active dimensions when common metadata is padded.""" + device = torch.device("cpu") + num_decodes = 2 + query_len = 3 + num_tokens = num_decodes * query_len + q = torch.randn(num_tokens, 2, 4, device=device) + topk = ( + torch.arange(num_tokens, dtype=torch.int32, device=device) + .view(-1, 1) + .expand(-1, 4) + ) + steps: list[int] = [] + kernel_shapes: list[torch.Size] = [] + + def swap_in(req_ids, *, logical_topk_indices, **kwargs): # noqa: ARG001 + steps.append(int(logical_topk_indices[0, 0])) + output = kwargs["attention_indices_out"] + if output is not None: + output.copy_(logical_topk_indices + 10) + return topk[:num_decodes] + + def run_kernel(self, *, q, **kwargs): # noqa: ARG001 + kernel_shapes.append(q.shape) + return q[..., :1], None + + runtime = SimpleNamespace( + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + ) + leader_cache = SimpleNamespace( + runtime=runtime, + source_block_table=torch.empty( + num_decodes, 1, dtype=torch.int32, device=device + ), + swap_in=MagicMock(side_effect=swap_in), + ) + follower_cache = SimpleNamespace( + runtime=runtime, + source_block_table=torch.empty( + num_decodes, 1, dtype=torch.int32, device=device + ), + swap_in=MagicMock(side_effect=swap_in), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [leader_cache, follower_cache] + index_group.physical_topk_indices = torch.empty( + (num_tokens + 1, 4), dtype=torch.int32, device=device + ) + index_group.valid_topk_counts = torch.empty( + num_tokens + 1, dtype=torch.int32, device=device + ) + index_group.request_ids = torch.arange( + num_decodes, dtype=torch.int32, device=device + ) + impl = SimpleNamespace( + kv_lora_rank=1, + index_group=index_group, + index_group_index=0, + ) + impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl) + metadata = SimpleNamespace( + num_decodes=num_tokens, + num_decode_tokens=num_tokens, + decode_max_query_len=1, + query_start_loc=torch.arange( + 0, + num_tokens + 1, + query_len, + dtype=torch.int32, + device=device, + ), + block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device), + block_size=64, + ) + + output = FlashMLASparseImpl._host_backed_fp8_decode( + impl, + q, + topk, + metadata, + SimpleNamespace(), + num_decodes, + query_len, + ) + + assert steps == [0, 1, 2] + assert kernel_shapes == [(num_decodes, query_len, 2, 4)] + assert output.shape == (num_tokens, 2, 1) + torch.testing.assert_close( + index_group.physical_topk_indices[:num_tokens], topk + 10 + ) + + follower_indices = index_group.convert_decode_logical_to_physical_topk( + 1, + topk, + metadata, + return_valid_counts=False, + num_decodes=num_decodes, + decode_query_len=query_len, + ) + assert steps == [0, 1, 2, 0, 1, 2] + torch.testing.assert_close(follower_indices, topk + 10) + + +def test_hisparse_varlen_decode_uses_canonical_request_rows(): + """Graph metadata's query upper bound must preserve packed token rows.""" + num_decodes = 2 + topk = torch.arange(18, dtype=torch.int32).view(9, 2) + + def swap_in(_request_ids, *, logical_topk_indices, **_kwargs): + return logical_topk_indices + 10 + + cache = SimpleNamespace( + source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32), + swap_in=MagicMock(side_effect=swap_in), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32) + index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32) + index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32) + metadata = SimpleNamespace( + num_decodes=num_decodes, + num_decode_tokens=topk.shape[0], + decode_max_query_len=15, + query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32), + block_size=64, + ) + + physical_topk = index_group.convert_decode_logical_to_physical_topk( + 0, + topk, + metadata, + return_valid_counts=False, + ) + + torch.testing.assert_close(physical_topk, topk + 10) + request_rows = cache.swap_in.call_args_list[0].args[0] + torch.testing.assert_close(request_rows, index_group.request_ids) + + +def test_flashinfer_hisparse_decode_runs_batched_attention(): + device = torch.device("cpu") + num_tokens = 6 + q = torch.randn(num_tokens, 2, 4, device=device) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device) + physical_topk = topk + 10 + valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) + kernel_shapes: list[torch.Size] = [] + + def convert_decode(self, *args, **kwargs): # noqa: ARG001 + return physical_topk, valid_counts + + def prepare_kernel(self, *args, **kwargs): # noqa: ARG001 + pass + + def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 + kernel_shapes.append(q.shape) + return q[..., :1], None + + cache_handle = SimpleNamespace( + runtime=SimpleNamespace( + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + ), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache_handle] + index_group.convert_decode_logical_to_physical_topk = MethodType( + convert_decode, index_group + ) + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl) + impl._run_mqa_kernel = MethodType(run_kernel, impl) + metadata = SimpleNamespace(num_decode_tokens=num_tokens) + + output, lse = FlashInferMLASparseImpl.forward_mqa( + impl, + q, + torch.empty(1, device=device), + metadata, + SimpleNamespace(), + ) + + assert kernel_shapes == [q.shape] + assert output.shape == (num_tokens, 2, 1) + assert lse is None + + +def test_flashattn_hisparse_decode_uses_index_group(): + num_tokens = 4 + q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) + q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) + physical = topk.clone() + counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.convert_decode_logical_to_physical_topk = MagicMock( + return_value=(physical, counts) + ) + index_group.physical_kv_cache = MagicMock( + return_value=torch.empty(1, device=DEVICE_TYPE) + ) + impl = object.__new__(FlashAttnMLASparseImpl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) + metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) + + output, lse = FlashAttnMLASparseImpl.forward_mqa( + impl, + (q_nope, q_rope), + torch.empty(1, device=DEVICE_TYPE), + metadata, + SimpleNamespace(), + ) + + assert output.shape == (num_tokens, 2, 1) + assert lse is None + index_group.convert_decode_logical_to_physical_topk.assert_called_once() + impl._run_mqa_kernel.assert_called_once() + + +def test_flashinfer_sm120_hisparse_decode_uses_index_group(): + num_tokens = 4 + q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) + physical = topk.clone() + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.convert_decode_logical_to_physical_topk = MagicMock( + return_value=physical + ) + index_group.physical_kv_cache = MagicMock( + return_value=torch.empty(1, device=DEVICE_TYPE) + ) + impl = object.__new__(FlashInferMLASparseSM120Impl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._run_mqa_kernel = MagicMock(return_value=q[..., :1]) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + topk_tokens=topk.shape[1], + ) + + output, lse = FlashInferMLASparseSM120Impl.forward_mqa( + impl, + q, + torch.empty(1, device=DEVICE_TYPE), + metadata, + SimpleNamespace(), + ) + + assert output.shape == (num_tokens, 2, 1) + assert lse is None + index_group.convert_decode_logical_to_physical_topk.assert_called_once() + impl._run_mqa_kernel.assert_called_once() + + +def test_hisparse_resident_prefill_uses_attention_block_stride(): + expected = torch.tensor([[19]], dtype=torch.int32) + cache_handle = SimpleNamespace( + all_context_pages_resident=True, + view=SimpleNamespace(block_size=64, attention_block_stride=832), + block_table=torch.tensor([[3]], dtype=torch.int32), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache_handle] + index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32) + index_group._convert_once = MagicMock(return_value=expected) + topk = torch.zeros((2, 1), dtype=torch.int32) + metadata = SimpleNamespace( + req_id_per_token=torch.zeros(2, dtype=torch.int32), + ) + + result = index_group.convert_logical_to_physical_topk( + 0, + topk, + metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + + assert result is expected + assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832 + + +def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill(): + """Default sparse backends must route multi-token chunks to prefill.""" + builder = object.__new__(FlashInferMLASparseMetadataBuilder) + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object()), + speculative_config=SimpleNamespace( + num_speculative_tokens=4, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + builder._init_reorder_batch_threshold( + 1024, + supports_spec_as_decode=True, + supports_dcp_with_varlen=True, + ) + + assert builder.reorder_batch_threshold == 1 + + +def test_hisparse_flashmla_reorders_full_speculative_window_as_decode(): + builder = object.__new__(FlashMLASparseMetadataBuilder) + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object()), + speculative_config=SimpleNamespace( + num_speculative_tokens=3, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + builder._init_reorder_batch_threshold( + 256, + supports_spec_as_decode=True, + supports_dcp_with_varlen=True, + ) + + assert builder.reorder_batch_threshold == 4 + + +def test_hisparse_mixed_single_token_batch_is_not_decode_only(): + cache = SimpleNamespace( + runtime=SimpleNamespace( + begin_forward=lambda: None, + eager_host_mirror=True, + ), + ) + metadata = SimpleNamespace( + num_decode_tokens=1, + num_actual_tokens=2, + max_query_len=1, + num_reqs=2, + req_id_per_token=torch.arange(2), + ) + + HiSparseCacheHandle._prepare_for_batch(cache, metadata) + + assert not cache.decode_batch + + +def test_flashmla_cache_dtype_aliases_use_ds_layout(): + # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer + # stores kv_cache_dtype, so they cannot bypass the gate. + for alias in ("fp8", "fp8_e4m3"): + assert ( + _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias) + == "fp8_ds_mla" + ) + + +def test_flashmla_common_metadata_requires_uniform_decodes(): + common_metadata = SimpleNamespace( + max_query_len=3, + num_reqs=3, + num_actual_tokens=6, + query_start_loc_cpu=torch.tensor([0, 1, 3, 6]), + is_prefilling=None, + ) + + split = split_decodes_and_prefills( + common_metadata, + decode_threshold=128, + require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes, + ) + + assert split == (1, 2, 1, 5) + + +def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata", + lambda: (object(), None), + ) + builder = SimpleNamespace( + device=torch.device(DEVICE_TYPE), + dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE), + max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE), + ) + query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660]) + common_metadata = SimpleNamespace( + num_actual_tokens=660, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=torch.arange(7, device=DEVICE_TYPE), + ) + metadata = FlashMLASparseMetadata( + num_reqs=7, + max_query_len=110, + max_seq_len=110, + num_actual_tokens=660, + query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE), + slot_mapping=torch.arange(660, device=DEVICE_TYPE), + block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE), + req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE), + num_decodes=7, + num_prefills=0, + num_decode_tokens=660, + ) + + fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( + builder, common_metadata, metadata + ) + + assert fp8_metadata.num_decodes == 6 + assert fp8_metadata.num_decode_tokens == 660 + assert fp8_metadata.decode is not None + assert fp8_metadata.decode.decode_query_len == 110 + torch.testing.assert_close( + fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE) + ) + + +@pytest.mark.parametrize("use_mixed_batch", [False, True]) +def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool): + num_decode_tokens = 2 + num_batch_tokens = 5 + q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE) + topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE) + kernel_q_shapes = [] + + def convert_indices(*args, **kwargs): # noqa: ARG001 + assert not kwargs.get("HAS_PREFILL_WORKSPACE", False) + if not kwargs.get("return_valid_counts", False): + return topk_indices + valid_counts = torch.full( + (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE + ) + return topk_indices, valid_counts + + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse." + "triton_convert_req_index_to_global_index", + convert_indices, + ) + + def run_kernel(**kwargs): + kernel_q_shapes.append(kwargs["q"].shape) + return kwargs["q"][..., :1], None + + if use_mixed_batch: + fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=object(), # type: ignore[arg-type] + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + ) + else: + FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode + fp8_metadata = FP8Meta( + num_decodes=1, + num_prefills=1, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_batch_tokens - num_decode_tokens, + decode=FP8Meta.Decode( + seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + kernel_metadata=object(), # type: ignore[arg-type] + decode_query_len=num_decode_tokens, + ), + prefill=FP8Meta.Prefill( + request_ids=torch.empty( + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE + ), + workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + chunks=[], + ), + ) + metadata = SimpleNamespace( + fp8_extra_metadata=fp8_metadata, + fp8_use_mixed_batch=use_mixed_batch, + num_actual_tokens=num_batch_tokens, + num_decode_tokens=num_decode_tokens, + req_id_per_token=torch.empty( + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE + ), + block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + block_size=64, + ) + impl = SimpleNamespace( + kv_cache_dtype="fp8_ds_mla", + topk_indices_buffer=topk_indices, + num_heads=2, + kv_lora_rank=1, + index_group=None, + index_group_index=0, + dcp_world_size=1, + need_to_return_lse_for_decode=False, + _fp8_flash_mla_kernel=run_kernel, + _convert_logical_to_physical_topk=( + lambda indices, metadata, **kwargs: convert_indices( + metadata.req_id_per_token[: indices.shape[0]], + metadata.block_table, + indices, + return_valid_counts=kwargs["return_valid_counts"], + ) + ), + ) + impl._forward_fp8_kv_mixed_batch = MethodType( + FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl + ) + impl._forward_fp8_kv_separate_prefill_decode = MethodType( + FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl ) output, lse = FlashMLASparseImpl.forward_mqa( @@ -1856,6 +3528,8 @@ def run_kernel(**kwargs): cp_kv_cache_interleave_size=1, ) impl = SimpleNamespace( + index_group=None, + index_group_index=0, dcp_world_size=2, dcp_rank=0, need_to_return_lse_for_decode=True, @@ -1876,6 +3550,117 @@ def run_kernel(**kwargs): assert not lse.isnan().any() +def test_hisparse_prefill_reuses_builder_staging_plan(): + """Every layer must reuse the batch plan instead of synchronizing to dedupe.""" + plan = SimpleNamespace( + block_table=torch.tensor([[0]], dtype=torch.int32), + ensure_gpu_sources=MagicMock(), + ) + staged = torch.empty((1, 1, 8)) + calls = [] + + def gather(kv_cache, staging_plan, resident_cache=None): + calls.append((kv_cache, staging_plan, resident_cache)) + return staged + + resident_cache = torch.empty((1, 1, 8)) + resident_block_table = torch.tensor([[1]], dtype=torch.int32) + cache = SimpleNamespace( + runtime=SimpleNamespace( + gather_prefill_cache=gather, + ), + view=SimpleNamespace(cache=resident_cache, block_size=1), + block_table=resident_block_table, + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + source = torch.empty((1, 1, 8)) + metadata = SimpleNamespace( + num_decodes=0, + num_decode_tokens=0, + seq_lens=torch.tensor([1], dtype=torch.int32), + prefill=SimpleNamespace(host_staging_plan=plan), + req_id_per_token=torch.tensor([0], dtype=torch.int32), + ) + + result, block_table, request_ids = index_group.stage_prefill_rows( + 0, source, metadata + ) + + assert result is staged + assert block_table is plan.block_table + torch.testing.assert_close(request_ids, metadata.req_id_per_token) + plan.ensure_gpu_sources.assert_called_once() + args = plan.ensure_gpu_sources.call_args.args + torch.testing.assert_close(args[0], resident_block_table) + assert args[1] == 1 + assert calls == [(source, plan, resident_cache)] + + +def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): + compute_stream = object() + prefill_stream = MagicMock() + gather = MagicMock() + monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr( + index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather + ) + + resident_cache = torch.empty((2, 4, 656), dtype=torch.uint8) + ready = MagicMock() + cache = SimpleNamespace( + view=SimpleNamespace(cache=resident_cache, block_size=4), + block_table=torch.tensor([[2, 3]], dtype=torch.int32), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + index_group.prefill_stream = prefill_stream + index_group.prefill_ready_events = [ready] + source = torch.empty((3, 4, 656), dtype=torch.uint8) + dst = torch.empty((8, 576), dtype=torch.bfloat16) + block_table = torch.tensor([[0, 1]], dtype=torch.int32) + workspace_starts = torch.tensor([0], dtype=torch.int32) + plan = SimpleNamespace( + block_table=torch.tensor([[0, 1]], dtype=torch.int32), + row_ids=torch.arange(8, dtype=torch.int32).view(1, -1), + gpu_row_ids=torch.full((1, 8), -1, dtype=torch.int32), + ensure_gpu_sources=MagicMock(), + ) + metadata = SimpleNamespace( + prefill=SimpleNamespace(host_staging_plan=plan), + num_decodes=0, + ) + + result = index_group.gather_fp8_prefill( + 0, + source, + dst, + block_table, + workspace_starts, + 1, + metadata, + 0, + ) + + prefill_stream.wait_stream.assert_called_once_with(compute_stream) + ready.record.assert_called_once_with(prefill_stream) + assert result is ready + gather.assert_called_once() + args = gather.call_args + assert args.args[0] is resident_cache + assert args.args[1] is dst + torch.testing.assert_close(args.args[2], plan.block_table) + assert args.args[3] is workspace_starts + assert args.args[4] == 1 + assert args.kwargs["host_cache"].data_ptr() == source.data_ptr() + assert args.kwargs["host_row_ids"] is plan.row_ids + assert args.kwargs["device_row_ids"] is plan.gpu_row_ids + plan.ensure_gpu_sources.assert_called_once() + ensure_args = plan.ensure_gpu_sources.call_args.args + torch.testing.assert_close(ensure_args[0], cache.block_table) + assert ensure_args[1] == 4 + + def test_sparse_impl_observes_repointed_indexer_buffer(): """The MTP proposer repoints the draft's indexer at the target model's buffer after the backend impl is built, so the impl must resolve the buffer per read. diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 97b846219fb4..8012c82f7841 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -41,7 +41,9 @@ def test_connector_without_divergent_hit_support_uses_common_lookup(): manager = MagicMock() manager.get_computed_blocks.return_value = (common_blocks, 0, 0) scheduler = SimpleNamespace( - connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False), + connector=SimpleNamespace( + supports_divergent_local_hybrid_hits=False, + ), kv_cache_manager=manager, ) @@ -61,7 +63,9 @@ def test_capable_connector_uses_divergent_partial_hit_lookup(): True, ) scheduler = SimpleNamespace( - connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True), + connector=SimpleNamespace( + supports_divergent_local_hybrid_hits=True, + ), kv_cache_manager=manager, ) @@ -922,71 +926,6 @@ def request_finished_all_groups(self, request, block_ids): scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once() -def test_block_pool_touch_pins_released_cow_target(): - """The connector can rescue an offered CoW target after its step-scoped - retention is released by using the bound BlockPool's touch method.""" - hash_block_size = 2 - block_size = 2 * hash_block_size - kv_cache_config = KVCacheConfig( - num_blocks=24, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full"], - FullAttentionSpec( - block_size=hash_block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mamba"], - MambaSpec( - block_size=block_size, - shapes=(1, 1), - dtypes=(torch.float32,), - mamba_cache_mode="align", - ), - ), - ], - ) - manager = make_kv_cache_manager( - kv_cache_config=kv_cache_config, - max_model_len=8192, - enable_caching=True, - hash_block_size=hash_block_size, - ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) - computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) - assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None - req0.num_computed_tokens = 6 - req0.append_output_token_ids([3]) - assert manager.allocate_slots(req0, 1) is not None - - # Retention released before the drain (defer_block_free=False ordering). - _copies, retained = manager.take_kv_cache_block_copies() - manager.block_pool.free_blocks(retained) - - offloads = drain_boundary_state_offloads(manager) - ((_group_id, block_id, boundary_tokens),) = offloads["0"] - assert boundary_tokens == 6 - cow_block = manager.block_pool.blocks[block_id] - assert cow_block.ref_cnt == 0 - assert block_id in _free_block_ids(manager) - - manager.block_pool.touch([cow_block]) - assert cow_block.ref_cnt == 1 - assert block_id not in _free_block_ids(manager) - - # The connector-pinned block is out of the free queue: draining every free block - # neither trips the allocator's ref_cnt assert nor hands it out. - new_blocks = manager.block_pool.get_new_blocks( - manager.block_pool.get_num_free_blocks() - ) - assert block_id not in {b.block_id for b in new_blocks} - - def test_boundary_state_offload_dropped_when_request_freed_before_drain(): """A hand-off recorded in the same scheduling pass as the request's death must not be drained: its release hook has already run, so draining would diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 0a21e896b030..7ff63d019c85 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -13,7 +13,7 @@ from vllm.v1.structured_output import StructuredOutputGrammar from vllm.v1.utils import ConstantList -from .utils import create_requests, create_scheduler +from .utils import create_requests, create_scheduler, mock_kv pytestmark = pytest.mark.cpu_test @@ -145,6 +145,32 @@ def abort_request(): assert req.num_output_tokens == abort_order_copy.index(i) +def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch): + """Mirroring must see earlier unresolved outputs, excluding the current step.""" + scheduler = create_scheduler( + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + ) + (request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4) + scheduler.add_request(request) + + assert scheduler.connector is not None + build_connector_meta = scheduler.connector.build_connector_meta + observed_placeholders = [] + + def record_placeholders(scheduler_output): + observed_placeholders.append(request.num_output_placeholders) + return build_connector_meta(scheduler_output) + + monkeypatch.setattr( + scheduler.connector, "build_connector_meta", record_placeholders + ) + scheduler.schedule() + scheduler.schedule() + + assert observed_placeholders == [0, 1] + + def test_preempt(): scheduler = create_scheduler(async_scheduling=True) requests = create_requests(num_requests=10, max_tokens=20) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 2bdb5e74a7b9..2fcf70c6efa7 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -75,6 +75,7 @@ def _mock_vllm_config(layout: str | None): config.cache_config = CacheConfig() config.cache_config.num_gpu_blocks_override = None config.cache_config.kv_cache_layout = layout + config.attention_config.hisparse_config = None return config diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index f0bd296b7183..e7eee03e7bff 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -22,6 +22,7 @@ SchedulerConfig, VllmConfig, ) +from vllm.config.attention import HiSparseConfig from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( @@ -52,10 +53,16 @@ make_block_hash_with_group_id, tensor_data, ) +from vllm.v1.hisparse.layout import ( + create_hisparse_layout, + get_hisparse_gpu_memory_usage, +) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, @@ -68,18 +75,139 @@ SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, + SparseCacheRole, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, is_full_attention_spec, iter_layer_specs, ) +from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats from vllm.v1.request import Request pytestmark = pytest.mark.cpu_test +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + group_spec = UniformTypeKVCacheSpecs.from_specs(specs) + assert group_spec is not None + group = KVCacheGroupSpec(list(specs), group_spec) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + indexer_spec = specs["model.layers.0.self_attn.indexer"] + assert get_hisparse_gpu_memory_usage(config, [group]) == ( + indexer_spec.max_memory_usage_bytes(config) + ) + + monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) + cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + config, [group], available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_config, + SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +def test_hisparse_rejects_deepseek_v4(): + full_specs = { + "model.layers.0.attn": MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + model_version="deepseek_v4", + ) + } + full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs) + assert full_uniform is not None + group = KVCacheGroupSpec(list(full_specs), full_uniform) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)), + cache_config=SimpleNamespace(num_gpu_blocks_override=7), + ) + + with pytest.raises(ValueError, match="does not support DeepSeek V4"): + create_hisparse_layout( + config, + [group], + host_budget=2**30, + ) + + @pytest.fixture(autouse=True) def _auto_init_hash_fn(request): hash_fn: Callable @@ -3848,6 +3976,22 @@ def test_iter_layer_specs_returns_group_members(): assert list(iter_layer_specs(wrapped)) == [full, mla] +def test_wrapped_mamba_group_requires_block_zeroing(): + mamba = MambaSpec( + block_size=4, + shapes=((4, 1),), + dtypes=(torch.float32,), + ) + wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba}) + config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)], + ) + + assert config.needs_kv_cache_zeroing + + def _spec_decode_grouping_config(method="dspark", model_type=None): """Grouping config with an EAGLE-family speculative method enabled.""" return SimpleNamespace( diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 1760ca14d0d6..d91260e30ae0 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -7,6 +7,7 @@ from dataclasses import replace from math import lcm from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -14,11 +15,21 @@ import vllm.v1.core.kv_cache_manager as kv_cache_manager import vllm.v1.core.kv_cache_utils as kv_cache_utils from vllm.distributed.kv_events import ( + MEDIUM_CPU, MEDIUM_GPU, AllBlocksCleared, BlockRemoved, BlockStored, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorScheduler, + HiSparseConnectorWorkerMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import ( + OffloadingConnector, +) from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -40,16 +51,23 @@ init_none_hash, make_block_hash_with_group_id, ) +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KVCacheConfig, + KVCacheGroupRole, KVCacheGroupSpec, KVCacheSpecKind, MambaSpec, MLAAttentionSpec, SlidingWindowSpec, ) +from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.request import RequestStatus pytestmark = pytest.mark.cpu_test @@ -118,7 +136,10 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa kv_cache_config, prefix_cache_retention_interval=kwargs.pop("retention_interval"), ) - return KVCacheManager(kv_cache_config, **kwargs) + manager = KVCacheManager(kv_cache_config, **kwargs) + if kv_cache_config.hisparse_host_num_blocks is not None: + get_hisparse_coordinator(manager) + return manager def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: @@ -139,6 +160,999 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: ) +HISPARSE_BLOCK_SIZE = 16 + + +def make_hisparse_kv_cache_config( + num_blocks: int, + host_num_blocks: int, + *, + transfer_device_cache: bool = False, +) -> KVCacheConfig: + source_spec = FullAttentionSpec( + block_size=HISPARSE_BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) + groups = [ + KVCacheGroupSpec( + ["source"], + source_spec, + role=KVCacheGroupRole.HISPARSE_SOURCE, + host_resident=True, + ), + KVCacheGroupSpec( + ["indexer"], + source_spec, + enable_kv_transfer=transfer_device_cache, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ), + ] + groups.append( + KVCacheGroupSpec( + ["resident"], + HiSparseResidentSpec( + block_size=HISPARSE_BLOCK_SIZE, + page_size=HISPARSE_BLOCK_SIZE * 4, + ), + enable_kv_transfer=transfer_device_cache, + ) + ) + groups.append( + KVCacheGroupSpec( + ["hot"], + HiSparseHotSpec( + block_size=HISPARSE_BLOCK_SIZE, + page_size=HISPARSE_BLOCK_SIZE * 4, + blocks_per_request=2, + ), + enable_kv_transfer=False, + ) + ) + return KVCacheConfig( + num_blocks=num_blocks, + hisparse_host_num_blocks=host_num_blocks, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + + +def make_hisparse_kv_cache_manager( + num_blocks: int, + host_num_blocks: int, + *, + max_model_len: int = 128, + max_in_flight_tokens: int | None = None, + enable_caching: bool = False, + **config_kwargs, +) -> KVCacheManager: + config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs) + return make_kv_cache_manager( + config, + max_model_len=max_model_len, + max_in_flight_tokens=max_in_flight_tokens, + enable_caching=enable_caching, + hash_block_size=HISPARSE_BLOCK_SIZE, + ) + + +def test_hisparse_builds_dma_row_mirrors_across_pages(): + manager = make_hisparse_kv_cache_manager(32, 16) + request = make_request( + "request", + list(range(2 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + assert manager.allocate_slots(request, num_new_tokens=32) is not None + coordinator = get_hisparse_coordinator(manager) + resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id] + host_blocks = coordinator.host_manager.req_to_blocks[request.request_id] + + mirrors = coordinator.build_row_mirrors([(request.request_id, 14, 4)]) + + assert [mirror.source_starts for mirror in mirrors] == [ + (resident_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,), + (resident_blocks[1].block_id * HISPARSE_BLOCK_SIZE,), + ] + assert [mirror.destination_start for mirror in mirrors] == [ + host_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14, + host_blocks[1].block_id * HISPARSE_BLOCK_SIZE, + ] + assert [mirror.num_rows for mirror in mirrors] == [2, 2] + + +def test_hisparse_async_speculation_mirrors_uncertain_position_range(): + """Unresolved drafts must not leave gaps in the eager host mirror.""" + coordinator = MagicMock(host_group_id=0) + coordinator.take_block_table_updates.return_value = {} + coordinator.build_offload_command.return_value = None + coordinator.build_row_mirrors.return_value = () + scheduler = HiSparseConnectorScheduler( + async_speculative=True, + draft_kv_lookahead=4, + ) + scheduler.bind_coordinator(coordinator) + request = SimpleNamespace(request_id="request", num_output_placeholders=3) + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = scheduler + connector.update_state_after_alloc(request, None, 0) + scheduler_output = SimpleNamespace( + block_table_updates=None, + kv_cache_block_copies=None, + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + new_block_ids=[], + req_ids=["request"], + num_computed_tokens=[103], + ), + num_scheduled_tokens={"request": 4}, + ) + + scheduler.build_connector_meta(scheduler_output) + + coordinator.build_row_mirrors.assert_called_once_with((("request", 100, 11),)) + + # Acceptance arrives between steps; the connector must read the live count. + request.num_output_placeholders = 1 + scheduler.build_connector_meta(scheduler_output) + coordinator.build_row_mirrors.assert_called_with((("request", 102, 9),)) + + connector.request_finished_all_groups(request, ()) + assert not scheduler.requests + + +def test_hisparse_reports_when_context_is_fully_resident(): + """A page the pool reuses under a host-reading request stops being resident.""" + manager = make_hisparse_kv_cache_manager(32, 16) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + coordinator = get_hisparse_coordinator(manager) + scheduled = ((request.request_id, len(tokens) - 1, 1),) + assert coordinator.all_context_pages_resident(scheduled) + + for hot_manager in coordinator.hot_managers: + hot_manager.require_hot(request.request_id) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + pool = manager.block_pool + pool.get_new_blocks(pool.get_num_free_blocks()) + + assert not coordinator.all_context_pages_resident(scheduled) + assert coordinator.take_block_table_updates().keys() == {request.request_id} + + +def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): + """Keep indexer-only imports host-backed in the resident block table.""" + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + spill_counts = {spill.transfer_id: 1 for spill in spills} + get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts) + _, indexer_blocks, _, _ = manager.get_blocks(original.request_id).blocks + evicted_indexer_id = indexer_blocks[2].block_id + manager.free(original) + manager.block_pool.evict_blocks({evicted_indexer_id}) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + blocks, num_local, _ = manager.get_computed_blocks(resumed) + max_completion = HISPARSE_BLOCK_SIZE + assert num_local == 2 * HISPARSE_BLOCK_SIZE + assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0] + + allocated = manager.allocate_slots( + resumed, + num_new_tokens=1, + num_new_computed_tokens=num_local, + new_computed_blocks=blocks, + num_external_computed_tokens=max_completion, + ) + + assert allocated is not None + source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks + assert len(source) == len(indexer) == len(resident) == 4 + assert len(hot) == 2 + # The local prefix is adopted from shadow pages (GPU-resident), while the + # externally imported page stays host-backed until its tail allocation. + assert not any(block.is_null for block in resident[:2]) + assert resident[2].is_null + assert not resident[3].is_null + + +def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix(): + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + spill_counts = {spill.transfer_id: 1 for spill in spills} + get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts) + host_blocks, _, _, _ = manager.get_blocks(original.request_id).blocks + evicted_host_id = host_blocks[0].block_id + manager.free(original) + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {evicted_host_id} + ) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + _, num_local, _ = manager.get_computed_blocks(resumed) + connector = object.__new__(OffloadingConnector) + connector._kv_cache_config = manager.kv_cache_config + connector.bind_kv_cache_manager(manager) + assert num_local == 0 + assert connector._max_loadable_tokens(resumed, num_local) == 0 + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize( + "available_tokens,imported_tokens,host_blocks", + [(0, 0, 2), (16, 16, 2), (64, 32, 3)], +) +def test_connector_completes_partial_prefix_without_importing_missing_host_kv( + nested, available_tokens, imported_tokens, host_blocks +): + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(64)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=64) is not None + _publish_hisparse_pages(manager) + host, indexer, _, _ = manager.get_blocks(original.request_id).blocks + missing_indexer = indexer[1].block_id + missing_host = host[0].block_id + manager.free(original) + manager.block_pool.evict_blocks({missing_indexer}) + + defer = True + + def lookup(request, num_computed_tokens, max_num_new_tokens=None): + nonlocal defer + if defer: + defer = False + return None, False + available = available_tokens + if max_num_new_tokens is not None: + available = min(available, max_num_new_tokens) + return available, available > 0 + + connector = object.__new__(OffloadingConnector) + connector._kv_cache_config = manager.kv_cache_config + connector.connector_scheduler = SimpleNamespace(get_num_new_matched_tokens=lookup) + if nested: + for _ in range(2): + outer = object.__new__(MultiConnector) + outer._connectors = (connector,) + outer._requests_to_connector = {} + connector = outer + connector.bind_kv_cache_manager(manager) + scheduler = SimpleNamespace(connector=connector, kv_cache_manager=manager) + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + + blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit( + scheduler, resumed + ) + assert local == 16 and not diverged + assert connector.get_num_new_matched_tokens(resumed, local) == (None, False) + + blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit( + scheduler, resumed + ) + external, asynchronous = connector.get_num_new_matched_tokens(resumed, local) + assert external == imported_tokens + assert asynchronous == (external > 0) + assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0] + assert ( + manager.allocate_slots( + resumed, + int(external == 0), + delay_cache_blocks=True, + num_new_computed_tokens=local, + new_computed_blocks=blocks, + num_external_computed_tokens=external, + ) + is not None + ) + allocated = manager.get_blocks(resumed.request_id).blocks + assert len(allocated[0]) == host_blocks + assert allocated[2][1].is_null == (external > 0) + manager.free(resumed) + + # A surviving indexer offload must not make absent host KV look computed. + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {missing_host} + ) + retry = make_request("retry", tokens, HISPARSE_BLOCK_SIZE, sha256) + _, local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, retry) + assert local == 0 + assert connector.get_num_new_matched_tokens(retry, local) == (0, False) + + +def allocate_external_prefix( + manager: KVCacheManager, request: Request, num_tokens: int +) -> KVCacheBlocks | None: + return manager.allocate_slots( + request, + num_new_tokens=0, + num_external_computed_tokens=num_tokens, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + + +@pytest.mark.parametrize("enable_caching", [False, True]) +def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching): + """Later admissions must not remove a scheduled request's only readable tier.""" + manager = make_hisparse_kv_cache_manager( + 32, 32, max_model_len=256, enable_caching=enable_caching + ) + first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(first, num_new_tokens=128) is not None + _publish_hisparse_pages(manager) + first.num_computed_tokens = 128 + pool = manager.block_pool + held = pool.get_new_blocks(pool.get_num_free_blocks() - 3) + + assert manager.allocate_slots(first, num_new_tokens=1) is not None + second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) + manager.allocate_slots(second, num_new_tokens=16) + coordinator = get_hisparse_coordinator(manager) + scheduled = (("first", 128, 1),) + assert coordinator.all_context_pages_resident(scheduled) + assert all(not m.has_hot("first") for m in coordinator.hot_managers) + + # Once another allocation finishes, A can acquire its hot buffer and + # safely release its clean pages for other requests to reuse. + pool.free_blocks(held) + assert manager.allocate_slots(first, num_new_tokens=1) is not None + assert all(m.has_hot("first") for m in coordinator.hot_managers) + pool.get_new_blocks(pool.get_num_free_blocks()) + assert not coordinator.all_context_pages_resident(scheduled) + assert "first" in coordinator.take_block_table_updates() + + +def test_hisparse_full_pool_keeps_pages_pinned_until_preemption(): + """Without room for a hot buffer, admission must defer instead of losing KV.""" + manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160) + first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(first, num_new_tokens=128) is not None + _publish_hisparse_pages(manager) + assert manager.block_pool.get_num_free_blocks() == 1 + first.num_computed_tokens = 128 + assert manager.allocate_slots(first, num_new_tokens=1) is None + assert get_hisparse_coordinator(manager).all_context_pages_resident( + (("first", 127, 1),) + ) + + manager.free(first) + second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(second, num_new_tokens=16) is not None + + +def test_hisparse_materializes_prefix_without_allocating_hot_blocks(): + """A host prefix becomes visible only when every page is durable. + + Completing a later page first must not expose a prefix with a hole. + """ + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + blocks = manager.get_block_ids(request.request_id) + assert len(blocks[2]) == 2 + assert blocks[3] == [] + + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + assert len(spills) == 2 + assert all(spill.after_forward for spill in spills) + resident_blocks = manager.get_blocks(request.request_id).blocks[2] + assert [block.ref_cnt for block in resident_blocks] == [2, 2] + spill_counts = {spill.transfer_id: 1 for spill in spills} + duplicate = make_request( + "duplicate", + list(range(3 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills(spill_counts, {}) + assert [block.ref_cnt for block in resident_blocks] == [2, 2] + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills({}, {spills[1].transfer_id: 1}) + assert [block.ref_cnt for block in resident_blocks] == [2, 1] + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills({}, {spills[0].transfer_id: 1}) + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 2 * HISPARSE_BLOCK_SIZE + blocks = manager.get_block_ids(request.request_id) + assert len(blocks[2]) == 2 + assert blocks[3] == [] + + +def test_hisparse_materialization_respects_per_step_spill_budget(): + """Prefix publication must not bypass the configured spill batch limit.""" + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + coordinator = get_hisparse_coordinator(manager) + coordinator.max_spill_pages = 1 + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256) + + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + first = coordinator.build_offload_command().page_transfers + assert len(first) == 1 + + coordinator.plan_prefix_materialization(request.request_id, len(tokens)) + second = coordinator.build_offload_command().page_transfers + assert len(second) == 1 + assert first[0].transfer_id != second[0].transfer_id + + +def test_hisparse_deferred_free_retains_host_blocks_until_fence(): + """Detaching a request must not make in-flight host pages reusable.""" + manager = make_hisparse_kv_cache_manager(16, 16) + host_pool = get_hisparse_coordinator(manager).get_host_block_pool() + assert host_pool is not None + host_free = host_pool.get_num_free_blocks() + device_free = manager.block_pool.get_num_free_blocks() + request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None + blocks = manager.pop_blocks_for_free(request) + assert host_pool.get_num_free_blocks() < host_free + assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks) + manager.block_pool.free_blocks(reversed(blocks)) + assert host_pool.get_num_free_blocks() == host_free + assert manager.block_pool.get_num_free_blocks() == device_free + + +def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool(): + """Host-only copy-on-write work must reach the worker copy queue.""" + manager = make_hisparse_kv_cache_manager(16, 16) + source_manager = manager.coordinator.single_type_managers[0] + source_block = source_manager.block_pool.get_new_blocks(1)[0] + source_manager.req_to_blocks["cow"] = [source_block] + source_manager._partial_hit_reqs["cow"] = (0, source_block) + + new_blocks = source_manager.allocate_new_blocks( + "cow", HISPARSE_BLOCK_SIZE, HISPARSE_BLOCK_SIZE + ) + new_block_ids = manager.take_new_block_ids() + coordinator = get_hisparse_coordinator(manager) + assert manager.take_kv_cache_block_copies() == ([], []) + connector = HiSparseConnectorScheduler(async_speculative=False) + connector.bind_coordinator(coordinator) + output = SchedulerOutput.make_empty() + copies = connector.build_connector_meta(output).host_block_copies + assert output.has_sync_kv_loads + + assert new_blocks and new_block_ids == [] + assert len(copies) == 1 + assert copies[0].src_block_id == source_block.block_id + assert copies[0].dst_block_id == new_blocks[0].block_id + source_manager.free("cow") + assert source_block.ref_cnt > 0 and new_blocks[0].ref_cnt > 0 + assert coordinator.has_pending_work() + connector.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=HiSparseConnectorWorkerMetadata( + {}, {}, completed_host_copy_dst_ids=(copies[0].dst_block_id,) + ) + ) + ) + assert source_block.ref_cnt == new_blocks[0].ref_cnt == 0 + assert not coordinator.has_pending_work() + + +def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages(): + """An in-flight host import must reserve its unwritten resident pages.""" + manager = make_hisparse_kv_cache_manager( + 11, + 16, + transfer_device_cache=True, + ) + request = make_request( + "partial-import", + list(range(6 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + imported_tokens = 4 * HISPARSE_BLOCK_SIZE + + assert allocate_external_prefix(manager, request, imported_tokens) is not None + required = manager.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=request.num_tokens, + new_computed_blocks=manager.empty_kv_cache_blocks.blocks, + num_encoder_tokens=0, + total_computed_tokens=imported_tokens, + num_local_computed_tokens=0, + num_tokens_main_model=request.num_tokens, + apply_admission_cap=True, + ) + + assert required == 4 + + +@pytest.mark.parametrize( + "full_sequence_must_fit,host_num_blocks,admitted", + [(True, 2, False), (True, 3, True), (False, 2, False), (False, 3, True)], +) +def test_hisparse_async_admission_requires_only_import_destinations( + full_sequence_must_fit, host_num_blocks, admitted, tmp_path +): + """Imports need host destinations, but do not reserve future prefill pages.""" + from .utils import create_scheduler, mock_kv + + (tmp_path / "config.json").write_text( + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}' + ) + scheduler = create_scheduler( + model=str(tmp_path), + skip_tokenizer_init=True, + max_model_len=128, + use_kv_connector=mock_kv(matched_tokens=HISPARSE_BLOCK_SIZE, is_async=True), + ) + manager = make_hisparse_kv_cache_manager(32, host_num_blocks) + scheduler.kv_cache_manager = manager + scheduler.kv_cache_config = manager.kv_cache_config + scheduler.scheduler_reserve_full_isl = full_sequence_must_fit + inflight = make_request( + "inflight", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + scheduler.add_request(inflight) + scheduler.schedule() + assert inflight in scheduler._inflight_prefills + + request = make_request( + "waiting", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + scheduler.add_request(request) + scheduler.schedule() + + assert (request in scheduler._inflight_prefills) == admitted + assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted + + +@pytest.mark.parametrize( + "evictable,local_tokens,free_blocks,admitted", + [ + (False, HISPARSE_BLOCK_SIZE, 1, True), + (True, HISPARSE_BLOCK_SIZE, 1, False), + (False, HISPARSE_BLOCK_SIZE - 1, 1, False), + (False, HISPARSE_BLOCK_SIZE - 1, 2, True), + ], +) +def test_hisparse_import_capacity_includes_hits_and_cow( + evictable, local_tokens, free_blocks, admitted +): + """Touching an evictable hit and copying a partial hit consume capacity.""" + manager = make_hisparse_kv_cache_manager(32, free_blocks + 2 - int(evictable)) + coordinator = get_hisparse_coordinator(manager) + host_pool = coordinator.get_host_block_pool() + assert host_pool is not None + hit = host_pool.get_new_blocks(1)[0] + if evictable: + host_pool.free_blocks([hit]) + assert host_pool.get_num_free_blocks() == free_blocks + assert ( + coordinator.host_manager.get_num_blocks_to_allocate( + "waiting", 64, [hit], 32, local_tokens, 64 + ) + <= host_pool.get_num_free_blocks() + ) == admitted + + +@pytest.mark.parametrize("enable_caching", [False, True]) +def test_hisparse_host_exhaustion_keeps_gpu_pages_readable(enable_caching): + """Pages without a host destination must survive GPU cache reclamation.""" + manager = make_hisparse_kv_cache_manager(32, 2, enable_caching=enable_caching) + request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + coordinator = get_hisparse_coordinator(manager) + host = coordinator.host_manager + assert host is not None + assert [b.is_null for b in host.req_to_blocks[request.request_id]] == [ + False, + True, + True, + True, + ] + _publish_hisparse_pages(manager) + for hot in coordinator.hot_managers: + hot.require_hot(request.request_id) + hot.allocate_new_blocks(request.request_id, 64, 64) + coordinator.update_residency(request.request_id) + for resident in coordinator.resident_managers: + for block in resident.req_to_blocks[request.request_id][1:]: + assert not block.is_null and block.ref_cnt > 0 + assert coordinator.request_states[request.request_id].valid_pages == {0} + pool = coordinator.get_host_block_pool() + assert pool is not None + manager.free(request) + assert pool.get_num_free_blocks() == 1 + + +def test_hisparse_host_cow_takes_priority_over_new_pages(): + """The last host block must hold the private tail, not a fresh page.""" + manager = make_hisparse_kv_cache_manager(32, 3) + coordinator = get_hisparse_coordinator(manager) + host = coordinator.host_manager + assert host is not None + source = host.block_pool.get_new_blocks(1)[0] + host.req_to_blocks["cow"] = [source] + host._partial_hit_reqs["cow"] = (0, source) + + host.allocate_new_blocks("cow", 32, 32) + + private, missing = host.req_to_blocks["cow"] + assert not private.is_null and private is not source + assert missing.is_null + assert host.take_host_cow_copies() == [(source, private)] + + +@pytest.mark.parametrize("failed", [False, True]) +def test_host_receive_completion_without_spill_metadata(failed): + """Only successful external receives make host pages readable.""" + manager = make_hisparse_kv_cache_manager(32, 16) + request = make_request("import", list(range(32)), HISPARSE_BLOCK_SIZE, sha256) + assert allocate_external_prefix(manager, request, 32) is not None + request.num_computed_tokens = 32 + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(get_hisparse_coordinator(manager)) + scheduler.requests[request.request_id] = request + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = scheduler + composite = object.__new__(MultiConnector) + composite._connectors = (connector,) + composite.update_connector_output( + KVConnectorOutput( + finished_recving={request.request_id}, + failed_recving={request.request_id} if failed else set(), + ) + ) + state = get_hisparse_coordinator(manager).request_states.get(request.request_id) + if failed: + assert state is None + else: + assert state is not None and state.valid_pages == {0, 1} + + +def test_hisparse_recomputation_does_not_complete_failed_host_import(): + """Allocating recomputation is not evidence that imported host KV is valid.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + request = make_request( + "partial-failure", list(range(48)), HISPARSE_BLOCK_SIZE, sha256 + ) + assert allocate_external_prefix(manager, request, 32) is not None + # A partial load failure leaves only the first page computed. The next + # allocation reaches the original import boundary before the worker runs. + request.num_computed_tokens = 16 + manager.cache_blocks(request, 16) + assert manager.allocate_slots(request, 16) is not None + state = get_hisparse_coordinator(manager).request_states.get(request.request_id) + assert state is None or 1 not in state.valid_pages + + +def test_hisparse_host_import_ignores_unsealed_tail(): + """A partial imported page must not become readable from stale host data.""" + manager = make_hisparse_kv_cache_manager(16, 16) + coordinator = get_hisparse_coordinator(manager) + + coordinator.complete_host_import("partial", HISPARSE_BLOCK_SIZE + 1) + + state = coordinator.request_states["partial"] + assert state.valid_pages == {0} + assert state.ready_prefix_pages == 1 + + +def test_hisparse_resident_request_can_grow_without_hot_capacity(): + """Running resident history must not be mistaken for a new host-prefix hit.""" + manager = make_hisparse_kv_cache_manager(8, 16) + request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=32) is not None + request.num_computed_tokens = 32 + + # Three free blocks fit the next indexer/resident pages, but not a hot region. + assert manager.block_pool.get_num_free_blocks() == 3 + assert manager.allocate_slots(request, num_new_tokens=1) is not None + assert manager.get_block_ids(request.request_id)[3] == [] + + +def test_hisparse_capacity_query_does_not_require_hot_blocks(): + """A read-only capacity query must not mutate hot-block requirements.""" + manager = make_hisparse_kv_cache_manager(16, 16) + coordinator = get_hisparse_coordinator(manager) + host_pool = coordinator.get_host_block_pool() + assert host_pool is not None + computed: tuple[list[KVCacheBlock], ...] = ( + [host_pool.blocks[1]], + [], + [], + [], + ) + + manager.coordinator.get_num_blocks_to_allocate( + request_id="query-only", + num_tokens=HISPARSE_BLOCK_SIZE, + new_computed_blocks=computed, + num_encoder_tokens=0, + total_computed_tokens=HISPARSE_BLOCK_SIZE, + num_local_computed_tokens=HISPARSE_BLOCK_SIZE, + num_tokens_main_model=HISPARSE_BLOCK_SIZE, + ) + + assert all( + "query-only" not in hot_manager.hot_required + for hot_manager in coordinator.hot_managers + ) + + +def _publish_hisparse_pages(manager: KVCacheManager) -> None: + """Ack all planned spills so host pages publish to the prefix cache.""" + command = get_hisparse_coordinator(manager).build_offload_command() + counts = {spill.transfer_id: 1 for spill in command.page_transfers} + get_hisparse_coordinator(manager).update_spills(counts, counts) + + +def test_hisparse_events_report_host_and_device_placement(): + """Host publication and eviction must not advertise GPU-resident KV.""" + manager = make_kv_cache_manager( + make_hisparse_kv_cache_config(32, 16), + max_model_len=128, + enable_caching=True, + hash_block_size=HISPARSE_BLOCK_SIZE, + enable_kv_cache_events=True, + ) + request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + _publish_hisparse_pages(manager) + + stored = [ + event for event in manager.take_events() if isinstance(event, BlockStored) + ] + assert {(event.group_idx, event.medium) for event in stored} == { + (0, MEDIUM_CPU), + (1, MEDIUM_GPU), + } + + host_block = manager.get_blocks(request.request_id).blocks[0][0] + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {host_block.block_id} + ) + [removed] = manager.take_events() + assert isinstance(removed, BlockRemoved) + assert removed.medium == MEDIUM_CPU + assert removed.group_idx == 0 + + +def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks(): + """Draining unused connector offers must not retain old recurrent states.""" + manager = make_kv_cache_manager( + KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ) + ], + ), + max_model_len=128, + enable_caching=True, + hash_block_size=16, + ) + request = make_request("request", list(range(64)), 16, sha256) + old_blocks = [] + for end in (16, 32, 48): + assert manager.allocate_slots(request, 16) is not None + request.num_computed_tokens = end + offers = manager.take_boundary_state_offloads() + [(_, block_id, boundary)] = offers[request.request_id] + assert boundary == end + old_blocks.append(manager.block_pool.blocks[block_id]) + manager.remove_skipped_blocks(request.request_id, end, 64) + manager.new_step_starts() + + assert all(block.ref_cnt == 0 for block in old_blocks[:-1]) + + +def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): + """A host prefix hit must come back GPU-resident while shadows survive.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + original_resident_ids = [ + block.block_id for block in manager.get_blocks("original").blocks[2] + ] + manager.free(original) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + resident_blocks = manager.get_blocks("resumed").blocks[2] + assert [block.block_id for block in resident_blocks[:3]] == ( + original_resident_ids[:3] + ) + assert not any(block.is_null for block in resident_blocks[:3]) + assert get_hisparse_coordinator(manager).all_context_pages_resident( + ((resumed.request_id, num_computed, len(tokens) - num_computed),) + ) + + +def test_hisparse_host_backed_request_accepts_local_prefix_hit(): + """Local group-completion hits need no external host import allocation.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(original) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + num_external_computed_tokens=0, + ) + + +def test_hisparse_finished_request_leaves_free_copies(): + """Clean pages of a finished request are free-queue blocks, not pins.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(original) + + coordinator = get_hisparse_coordinator(manager) + assert coordinator.copies + pool = manager.block_pool + assert pool.get_num_free_blocks() == pool.num_gpu_blocks - 1 + pool.get_new_blocks(pool.get_num_free_blocks()) + assert not coordinator.copies + assert not coordinator.spills_to_send + + +def test_hisparse_reset_prefix_cache_drops_copies(): + """Reset must forget GPU copies along with the host hashes they mirror.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(request) + + coordinator = get_hisparse_coordinator(manager) + assert coordinator.copies + assert manager.reset_prefix_cache() + assert not coordinator.copies + + +def test_hisparse_reused_copy_is_not_adopted(): + """A host prefix hit must not adopt a GPU copy the pool has handed out.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + indexer_blocks = list(manager.get_blocks("original").blocks[1]) + manager.free(original) + pool = manager.block_pool + pool.touch(indexer_blocks) + reused = pool.get_new_blocks(pool.get_num_free_blocks()) + pool.free_blocks(reused) + pool.free_blocks(indexer_blocks) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + resident_blocks = manager.get_blocks("resumed").blocks[2] + assert all(block.is_null for block in resident_blocks[:3]) + + +def test_hisparse_external_import_uses_hard_gpu_footprint(): + """An external prefix larger than resident capacity must remain admissible. + + This guards the admission boundary where indexer + one writable resident + page + the fixed hot allocation fit, but indexer + full resident history do + not. Requiring one more resident page would silently strand long P/D inputs. + """ + num_prompt_blocks = 4 + num_prompt_tokens = num_prompt_blocks * HISPARSE_BLOCK_SIZE + manager = make_hisparse_kv_cache_manager( + 8, + 9, + transfer_device_cache=True, + ) + request = make_request( + "host-import", + list(range(num_prompt_tokens)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + + allocated = allocate_external_prefix(manager, request, num_prompt_tokens) + + assert allocated is not None + source, indexer, resident, hot = manager.get_blocks(request.request_id).blocks + assert len(source) == num_prompt_blocks + assert len(indexer) == num_prompt_blocks + assert len(resident) == num_prompt_blocks + assert all(block.is_null for block in resident) + assert len(hot) == 2 + assert manager.block_pool.get_num_free_blocks() == 1 + + +def test_hisparse_external_import_survives_capacity_retry(): + manager = make_hisparse_kv_cache_manager( + 9, + 7, + transfer_device_cache=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + first = make_request("first", tokens, HISPARSE_BLOCK_SIZE, sha256) + second = make_request("second", tokens, HISPARSE_BLOCK_SIZE, sha256) + + assert allocate_external_prefix(manager, first, len(tokens)) is not None + + assert allocate_external_prefix(manager, second, len(tokens)) is None + + manager.free(first) + assert allocate_external_prefix(manager, second, len(tokens)) is not None + _, _, resident, hot = manager.get_blocks(second.request_id).blocks + assert all(block.is_null for block in resident) + assert len(hot) == 2 + + def make_kv_cache_config_hybrid_model( block_size: int, num_blocks: int, @@ -3874,6 +4888,42 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): ) +@pytest.mark.parametrize(("reserved_blocks", "fits"), [(4, True), (10, False)]) +def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk( + reserved_blocks, fits +): + """Reservations constrain the current allocation, not the entire prompt.""" + block_size = 16 + config = KVCacheConfig( + num_blocks=11, # Ten usable blocks after the null block. + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + manager = make_kv_cache_manager( + config, max_model_len=128, enable_caching=True, hash_block_size=block_size + ) + request = make_request("chunk", list(range(128)), block_size, sha256) + + blocks = manager.allocate_slots( + request, + block_size, + full_sequence_must_fit=True, + reserved_blocks=reserved_blocks, + ) + + assert (blocks is not None) is fits + + def test_can_fit_full_sequence_full_attention_still_gates_oversized(): """The cap only loosens the SWA group; a prompt that exceeds the full-attention pool capacity must still be rejected.""" @@ -3926,6 +4976,30 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None +def test_can_fit_full_sequence_hisparse_caps_resident_pages(): + manager = make_hisparse_kv_cache_manager( + num_blocks=11, + host_num_blocks=9, + max_model_len=8 * HISPARSE_BLOCK_SIZE, + max_in_flight_tokens=2 * HISPARSE_BLOCK_SIZE, + ) + request = make_request( + "long-hisparse", + list(range(8 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + + assert ( + manager.allocate_slots( + request, + num_new_tokens=HISPARSE_BLOCK_SIZE, + full_sequence_must_fit=True, + ) + is not None + ) + + def test_cache_hit_local_and_external(): # Regression test for #33775: when a request hits the local prefix cache # in one KV cache group and needs external (connector) blocks in another, @@ -4771,3 +5845,24 @@ def distinct(v): assert last_req_hit(retention=0, pin=False) == 0 # retention=0 with the pin keeps the junction window -> reuse restored. assert last_req_hit(retention=0, pin=True) == 4 * block_size + + +def test_device_eviction_keeps_host_prefix(): + """Device eviction must not target a same-numbered host block.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + _publish_hisparse_pages(manager) + host, indexer, _, _ = manager.get_blocks(request.request_id).blocks + device_block = indexer[0] + host_block = host[0] + assert device_block.block_id == host_block.block_id + manager.free(request) + assert host_block.block_hash is not None and device_block.block_hash is not None + manager.evict_blocks({device_block.block_id}) + assert device_block.block_hash is None, ( + "Device eviction left the indexer prefix cached" + ) + assert host_block.block_hash is not None, ( + "Device eviction invalidated unrelated host KV" + ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 1c56eaa6a500..3a972e682e4e 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -19,7 +19,12 @@ SpeculativeConfig, VllmConfig, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorScheduler, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, @@ -34,6 +39,7 @@ from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.engine import FinishReason +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -148,6 +154,30 @@ def test_get_num_unfinished_requests(): assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1 +def _bind_hisparse_connector(scheduler): + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) + from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager + + coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16)) + connector.connector_scheduler.bind_coordinator(coordinator) + composite = object.__new__(MultiConnector) + composite._connectors = (connector,) + scheduler.connector = composite + return coordinator + + +def test_pending_hisparse_spill_keeps_scheduler_alive(): + """A final host spill must complete after the last request finishes.""" + scheduler = create_scheduler() + coordinator = _bind_hisparse_connector(scheduler) + pending = coordinator.pending_spills + pending[0] = Mock() + assert scheduler.has_requests() + pending.clear() + assert not scheduler.has_requests() + + @pytest.mark.parametrize( "enable_prefix_caching, prompt_logprobs", [ @@ -1181,6 +1211,45 @@ def test_preempt_during_execution(): assert requests[1].output_token_ids[0] == 42 +def test_pending_hisparse_reclamation_defers_preemption(monkeypatch): + """Pending reclamation must stall rather than preempt its owning request.""" + scheduler = create_scheduler(enable_prefix_caching=False) + request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0] + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + monkeypatch.setattr( + scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None) + ) + coordinator = _bind_hisparse_connector(scheduler) + monkeypatch.setattr( + scheduler.connector, "build_connector_meta", Mock(return_value=None) + ) + monkeypatch.setattr( + coordinator, + "has_pending_reclamation", + Mock(return_value=True), + ) + + deferred_output = scheduler.schedule() + + assert not deferred_output.num_scheduled_tokens + assert scheduler.running == [request] + assert request.status == RequestStatus.RUNNING + assert request.num_preemptions == 0 + + def test_prefix_cache_query_not_inflated_by_connector_defer(): """The GPU prefix-cache query is recorded at admission, so a request the connector defers several times is counted once, not once per retry.""" diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py new file mode 100644 index 000000000000..dd4c5886570d --- /dev/null +++ b/tests/v1/e2e/general/test_hisparse.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import prometheus_client +import pytest +import torch + +from tests.conftest import VllmRunner +from tests.utils import fork_new_process_for_each_test +from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( + MultiConnector, +) +from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator + +MODEL = "deepseek-ai/DeepSeek-V3.2" + +HF_OVERRIDES = { + "num_hidden_layers": 8, + "hidden_size": 256, + "intermediate_size": 512, + "num_attention_heads": 8, + "num_key_value_heads": 1, + "n_routed_experts": 8, + "num_experts_per_tok": 2, + "index_topk": 128, +} + + +def _shrink_config(config): + overrides = HF_OVERRIDES.copy() + if any("MTP" in architecture for architecture in config.architectures): + overrides.pop("num_hidden_layers") + config.update(overrides) + return config + + +def _num_hisparse_spills(runner: VllmRunner) -> int: + client = runner.llm.llm_engine.engine_core + coordinator = get_hisparse_coordinator( + client.engine_core.scheduler.kv_cache_manager + ) + return coordinator.next_spill_id + + +def _offload_load_bytes() -> float: + return sum( + sample.value + for metric in prometheus_client.REGISTRY.collect() + if metric.name == "vllm:kv_offload_load_bytes" + for sample in metric.samples + if sample.name == "vllm:kv_offload_load_bytes_total" + ) + + +def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: + engine_core = runner.llm.llm_engine.engine_core.engine_core + model_runner = engine_core.model_executor.driver_worker.worker.model_runner + connector = model_runner.kv_connector.kv_connector + if isinstance(connector, MultiConnector): + connector = next( + child + for child in connector.sub_connectors + if isinstance(child, HiSparseConnector) + ) + assert isinstance(connector, HiSparseConnector) + assert connector.connector_worker is not None + return connector.connector_worker + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" +) +@pytest.mark.parametrize( + "with_offloading", [False, True], ids=["standalone", "offload"] +) +@fork_new_process_for_each_test +def test_hisparse_spill_and_prefix_restore( + monkeypatch: pytest.MonkeyPatch, + vllm_runner: type[VllmRunner], + with_offloading: bool, +): + """Spilled prefixes restore and FULL-graph decode writes reach host KV. + + A regression omitted attention metadata before FULL graph replay, so decode + completed normally while its newly written KV rows were never copied to host. + """ + capability = current_platform.get_device_capability() + if capability is None or capability.major < 9: + pytest.skip("Sparse MLA requires Hopper or newer") + + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") + + target = [1000 + i % 64 for i in range(257)] + pressure = [ + [2000 + request_idx * 128 + i % 64 for i in range(257)] + for request_idx in range(4) + ] + hisparse_connector = { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 1}, + } + kv_transfer_config = KVTransferConfig(**hisparse_connector) + if with_offloading: + kv_transfer_config = KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + hisparse_connector, + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, + }, + ] + }, + ) + + with vllm_runner( + MODEL, + load_format="dummy", + hf_overrides=_shrink_config, + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig( + device_buffer_size=512, + ) + ), + kv_transfer_config=kv_transfer_config, + block_size=64, + max_model_len=320, + max_num_batched_tokens=1024, + max_num_seqs=4, + num_gpu_blocks_override=128, + disable_log_stats=False, + enable_chunked_prefill=True, + enable_prefix_caching=True, + speculative_config={"method": "mtp", "num_speculative_tokens": 3}, + compilation_config={ + "cudagraph_mode": "FULL_AND_PIECEWISE", + "cudagraph_capture_sizes": [1, 4], + }, + ) as runner: + worker = _get_hisparse_worker(runner) + engine_core = runner.llm.llm_engine.engine_core.engine_core + model_runner = engine_core.model_executor.driver_worker.worker.model_runner + full_graph_calls = 0 + full_replay_pending = False + original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph + + def record_fullgraph(batch_desc): + nonlocal full_graph_calls, full_replay_pending + full_graph_calls += 1 + full_replay_pending = True + return original_run_fullgraph(batch_desc) + + monkeypatch.setattr( + model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph + ) + + original_finish_forward = worker.finish_forward + verified_full_replay_rows = 0 + copied_nonzero_kv = False + + def finish_forward(): + nonlocal full_replay_pending, verified_full_replay_rows + nonlocal copied_nonzero_kv + try: + original_finish_forward() + if not full_replay_pending: + return + # Synchronize only the test's inspection, not finish_forward itself. + with gpu_sync_allowed(): + torch.accelerator.synchronize() + assert worker._row_mirrors + for layer_index, cache in enumerate(worker.cache_handles): + source_index = cache.runtime.resident_source_index + for mirror in worker._row_mirrors: + source_row = mirror.source_starts[source_index] + destination_row = mirror.destination_start + num_rows = mirror.num_rows + source_block, source_offset = divmod( + source_row, worker.kernel_block_size + ) + with gpu_sync_allowed(): + gpu_rows = worker.resident_caches[layer_index][ + source_block, source_offset : source_offset + num_rows + ].cpu() + host_rows = worker.host_caches[layer_index][ + destination_row : destination_row + num_rows + ] + assert torch.equal(host_rows, gpu_rows) + copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item()) + verified_full_replay_rows += num_rows + finally: + full_replay_pending = False + + monkeypatch.setattr(worker, "finish_forward", finish_forward) + + expected = runner.generate_greedy([target], max_tokens=8) + + assert full_graph_calls > 0 + assert verified_full_replay_rows > 0 + assert copied_nonzero_kv + + runner.generate_greedy(pressure, max_tokens=8) + + assert _num_hisparse_spills(runner) > 0 + load_bytes = _offload_load_bytes() + actual = runner.generate_greedy([target], max_tokens=8) + runner.generate_greedy([[42]], max_tokens=1) + + assert actual == expected + if with_offloading: + assert _offload_load_bytes() > load_bytes diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index cd5356031e82..10464d950dde 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -21,6 +21,7 @@ FullAttentionSpec, HiddenStateCacheSpec, KVCacheConfig, + KVCacheGroupRole, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor, @@ -336,6 +337,36 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): assert offloading_config.cache.blocks_per_chunk == 2 +def test_hisparse_offloads_only_indexer_group(): + source = KVCacheGroupSpec( + ["source"], + _full_attention_spec(), + host_resident=True, + ) + indexer = KVCacheGroupSpec( + ["indexer"], + _full_attention_spec(), + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[source, indexer], + hisparse_host_num_blocks=4, + ) + config = _make_vllm_config() + + offloading_config = build_offloading_config(config, kv_cache_config) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), config, kv_cache_config + ) + + assert [ + (group.group_id, group.layer_names) for group in offloading_config.groups + ] == [(1, ("indexer",))] + assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1] + + def test_zero_blocks_skips_tensor_layout_validation(): kv_cache_config = _make_sizing_kv_cache_config(packed=False) kv_cache_config.num_blocks = 0 @@ -792,7 +823,6 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): False, id="uniform-uncertifiable-inner", ), - pytest.param([], False, id="attention-free"), ], ) def test_canonical_layout_gate(kv_cache_groups, certified): diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 47fb7103d742..26664e3ac550 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -98,6 +98,14 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): ) manager = CPUOffloadingManager(num_chunks=100) spec = SimpleNamespace( + config=SimpleNamespace( + groups=tuple( + SimpleNamespace( + group_id=i, tokens_per_block=group.kv_cache_spec.block_size + ) + for i, group in enumerate(groups) + ) + ), tokens_per_block=(256, 64, 8), tokens_per_hash=8, blocks_per_chunk=4, @@ -303,7 +311,24 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail(): assert src_spec.block_indices == [0, 0] -def test_aligned_boundary_store_flushes_before_cow_destination_reuse(): +def test_aligned_boundary_store_maps_sparse_group_id_to_dense_transfer_slot(): + scheduler = _make_partial_tail_scheduler() + indexer_config = scheduler.config.kv_group_configs[1] + scheduler.config = scheduler.config._replace(kv_group_configs=(indexer_config,)) + _make_partial_tail_request(scheduler) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 98, 16)]}) + + [job] = jobs.values() + assert job.src_spec.group_sizes == [1] + assert job.src_spec.block_indices == [0] + + +@pytest.mark.parametrize("cow_reuse", [False, True]) +def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse): scheduler = _make_partial_tail_scheduler() _make_partial_tail_request(scheduler) scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -320,7 +345,11 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse(): [job_id] = meta.store_jobs output = SchedulerOutput.make_empty() - output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)] + if cow_reuse: + output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)] + else: + output.scheduled_cached_reqs.req_ids = ["req"] + output.scheduled_cached_reqs.new_block_ids = [([], [99])] meta = scheduler.build_connector_meta(output) assert meta.jobs_to_flush == {job_id} @@ -384,6 +413,20 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys(): assert req_status.partial_tail_boundary is None +def test_lookup_cap_stops_at_authoritative_prefix_boundary(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + request.skip_reading_prefix_cache = False + scheduler.manager.lookup.return_value = LookupResult.HIT + + tokens, load_async = scheduler.get_num_new_matched_tokens( + request, 0, max_num_new_tokens=20 + ) + + assert (tokens, load_async) == (20, True) + assert scheduler._req_status["req"].partial_tail_boundary == 20 + + def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): """A recurrent group may hold unhashed blocks inside the computed prefix. diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py new file mode 100644 index 000000000000..6c4e79efbe08 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole +from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( + ExampleConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorMetadata, + HiSparseConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator +from vllm.v1.hisparse.runtime import HiSparseCacheHandle +from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + +@pytest.mark.parametrize("nested", [False, True]) +def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested): + """Composite binding must reach HiSparse and existing pool-only connectors.""" + hisparse = object.__new__(HiSparseConnector) + hisparse._role = KVConnectorRole.SCHEDULER + hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) + legacy = object.__new__(ExampleConnector) + legacy.bind_gpu_block_pool = MagicMock() + connector = object.__new__(MultiConnector) + connector._connectors = [hisparse, legacy] + if nested: + parent = object.__new__(MultiConnector) + parent._connectors = [connector] + connector = parent + from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager + + manager = make_hisparse_kv_cache_manager(16, 16) + + connector.bind_kv_cache_manager(manager) + + assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager) + legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool) + + +def test_hisparse_requires_block_outermost_device_layout(): + assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC" + + +def test_no_forward_enqueues_deferred_hisparse_transfers(): + """A zero-token step must still enqueue deferred post-forward transfers.""" + connector = object.__new__(ActiveKVConnector) + connector._disabled = False + connector.pre_forward = MagicMock() + connector.finish_forward = MagicMock() + connector.post_forward = MagicMock(return_value=None) + + scheduler_output = SimpleNamespace(finished_req_ids=set()) + connector.no_forward(scheduler_output) + + connector.pre_forward.assert_called_once_with(scheduler_output) + connector.finish_forward.assert_called_once_with() + + +def test_full_graph_step_prepares_host_mirror_outside_model(): + """Graph replay must restore host-mirror state cleared at step start.""" + runtime = SimpleNamespace( + is_group_leader=True, + eager_host_mirror=True, + begin_forward=MagicMock(), + invalidate_written_slots=MagicMock(), + ) + handle = HiSparseCacheHandle(runtime) + handle.mirror_slot_mapping = torch.tensor([4, 5]) + worker = object.__new__(HiSparseConnectorWorker) + worker.cache_layer_names = ["layer"] + worker.cache_handles = [handle] + worker._group_leaders = (("layer", handle),) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker.is_host_writer = True + worker._enqueue_row_dma = MagicMock() + worker.start_step = MagicMock( + side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state() + ) + + connector = object.__new__(HiSparseConnector) + connector.connector_worker = worker + connector._get_connector_metadata = MagicMock( + return_value=HiSparseConnectorMetadata(None, (), (), {}, True) + ) + req_id_per_token = torch.tensor([0, 1], dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_actual_tokens=2, + num_decode_tokens=2, + num_reqs=2, + max_query_len=1, + req_id_per_token=req_id_per_token, + ) + + connector.start_load_kv( + SimpleNamespace(), + request_state_indices=None, + request_ids=[], + attn_metadata={"layer": attn_metadata}, + ) + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None) + runtime.invalidate_written_slots.assert_called_once() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 52df7531e715..f613ba720e2e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -3297,7 +3297,9 @@ def _make_bare_worker( # path; everything flows through the coordinator). from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, + KVCacheTensor, ) worker.disk_offload_buffer_budget_bytes = None @@ -3311,6 +3313,18 @@ def _make_bare_worker( block_size=block_size, num_kv_heads=8, head_size=64, dtype=None ) group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec) + worker._kv_cache_config = KVCacheConfig( + num_blocks=num_gpu_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=0, + layers=["layer0", "__cross_layer__"], + layer_stride=0, + block_stride=0, + ) + ], + kv_cache_groups=[group], + ) worker._kv_cache_groups = [group] worker.pcp_size = 1 worker.dcp_size = 1 @@ -4059,6 +4073,77 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout): worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) +def test_register_kv_caches_uses_transfer_group_memory_domain(): + """Derived GPU caches must not change host-source transfer addresses.""" + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupRole, + KVCacheGroupSpec, + KVCacheTensor, + ) + + host_num_blocks = 4 + gpu_num_blocks = 2 + page_size = 32 + source_spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16 + ) + indexer_spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16 + ) + source_group = KVCacheGroupSpec( + ["source"], + source_spec, + host_resident=True, + ) + indexer_group = KVCacheGroupSpec( + ["indexer"], + indexer_spec, + enable_kv_transfer=False, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + config = KVCacheConfig( + num_blocks=gpu_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=host_num_blocks * page_size, + layers=["source"], + layer_stride=host_num_blocks * page_size, + block_stride=page_size, + host_resident=True, + ), + KVCacheTensor( + size=gpu_num_blocks * page_size, + layers=["indexer"], + layer_stride=gpu_num_blocks * page_size, + block_stride=page_size, + ), + ], + kv_cache_groups=[source_group, indexer_group], + hisparse_host_num_blocks=host_num_blocks, + ) + worker = _make_bare_worker(num_gpu_blocks=gpu_num_blocks) + worker._kv_cache_config = config + worker._kv_cache_groups = list(config.transfer_groups) + worker.token_dbs = [ + ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16) + ] + source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8) + indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8) + + _register_with_mocked_threads( + worker, + {"source": source, "indexer": indexer}, + ) + + db = worker.token_dbs[0] + assert db.kv_caches_base_addr == [source.data_ptr()] + assert db.block_len == [page_size] + worker.store.register_buffer.assert_called_once_with( + source.data_ptr(), source.nbytes + ) + + def test_register_kv_caches_separate_head_groups(): # LHBNC gives each head group its own region; the K/V split doubles heads. layout = KVCacheLayout.LHBNC diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index bac0b0373807..737d03ac6900 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -304,10 +304,10 @@ def test_multi_example_connector_consistency(): events = get_connector_events() storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) - # Initial events bind the block pool, query completion counts, and exchange + # Initial events bind the cache manager, query completion counts, and exchange # handshake metadata before the request is enqueued. assert storage1_scheduler_events[:7] == [ - "bind_gpu_block_pool", + "bind_kv_cache_manager", "get_finished_count", "set_xfer_handshake_metadata_pp_aware", "on_new_request", @@ -315,8 +315,8 @@ def test_multi_example_connector_consistency(): "update_state_after_alloc num_blocks=[7] 0", "build_connector_meta", ] - # First three events are from initialization. During generate(), layer hooks - # run before the deferred load is started after the forward pass. + # First three events are from initialization. Layer hooks run before the + # deferred load starts after the forward pass. expected_worker_prefix = [ "register_kv_caches", "set_host_xfer_buffer_ops", @@ -328,12 +328,12 @@ def test_multi_example_connector_consistency(): ] for connector_name in ("storage1-WORKER", "storage2-WORKER"): worker_events = events[connector_name] - assert worker_events[:7] == expected_worker_prefix + assert worker_events[: len(expected_worker_prefix)] == expected_worker_prefix assert worker_events.index("start_load_kv") > worker_events.index( "save_kv_layer" ) assert storage2_scheduler_events[:7] == [ - "bind_gpu_block_pool", + "bind_kv_cache_manager", "get_finished_count", "set_xfer_handshake_metadata_pp_aware", "on_new_request", diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 282458d1a954..3fcbda4c6255 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -29,7 +29,10 @@ get_current_attn_backend, ) from vllm.distributed.kv_transfer.kv_connector.v1 import nixl -from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorRole, + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiKVConnectorStats, @@ -518,6 +521,7 @@ def _nixl_handshake( slot_size_bytes = 4096 self.slot_size_per_layer = [slot_size_bytes] self.block_len_per_layer = [slot_size_bytes * self.block_size] + self.num_blocks = self.kv_cache_config.num_blocks self.num_regions = 1 self.block_stride_per_layer = list(self.block_len_per_layer) self.region_num_blocks = [self.num_blocks] @@ -635,6 +639,12 @@ def test_pcp_producer_uses_canonical_replica( if pcp_rank > 0: assert connector.get_finished(set()) == (set(), set()) + worker.get_transfer_results = MagicMock( + return_value=KVConnectorTransferResults(finished_sending={"sent"}) + ) + results = connector.get_transfer_results(set()) + assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set()) + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, @@ -1195,6 +1205,7 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices( ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, physical_blocks_per_logical_kv_block=1, + region_num_blocks=None, ) assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ @@ -1435,9 +1446,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init): metadata = NixlConnectorMetadata() metadata.add_new_req_to_recv( request_id=request_id, - local_block_ids=([1, 2, 3],), + local_block_ids=([0],), kv_transfer_params={ - "remote_block_ids": ([4, 5, 6],), + "remote_block_ids": ([0],), "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, "remote_request_id": f"prefill-{request_id}", "remote_host": "localhost", @@ -1482,6 +1493,10 @@ def test_kv_connector_stats(default_vllm_config, dist_init): assert stats_after_reset is None +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init): """reqs_to_send deadlines are stamped with the scheduler process's perf_counter, whose epoch differs across processes and (by boot-time @@ -1906,9 +1921,10 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert handle == 22 assert worker._dram_src_handles_by_block_size[worker.block_size] == 11 assert [block[2] for block in blocks] == [0, 0, 3, 3] - assert [ + memory_types = [ call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list - ] == ["DRAM", "VRAM"] + ] + assert memory_types == ["DRAM", "VRAM"] def test_mixed_memory_read_notifies_after_both_transfers_finish(): @@ -1919,6 +1935,7 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish(): worker._recving_metadata = {"request": MagicMock()} worker._recving_transfers = defaultdict(list) worker._pending_recv_notifs = {} + worker._failed_inflight_recvs = set() worker.xfer_stats = MagicMock() worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102] @@ -1959,6 +1976,7 @@ def test_mixed_memory_read_failure_does_not_notify_producer(): worker._recving_metadata = {"request": MagicMock()} worker._recving_transfers = {"request": [101, 102]} worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} + worker._failed_inflight_recvs = set() worker._is_hma_required = True worker._failed_recv_reqs = MagicMock() worker._log_failure = MagicMock() # type: ignore[method-assign] @@ -2871,9 +2889,11 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) -def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init): - """Test that transfer setup failures mark blocks invalid - and return via get_finished.""" +@pytest.mark.parametrize("is_hma", [False, True]) +def test_transfer_setup_failure_returns_finished( + default_vllm_config, dist_init, is_hma +): + """Setup failures report the request; only non-HMA reports block IDs.""" vllm_config = create_vllm_config() connector = NixlConnector( @@ -2883,6 +2903,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) vllm_config, connector.engine_id, hand_shake_latency=0 ) connector.connector_worker.nixl_wrapper.fail_transfer_setup = True + connector.connector_worker._is_hma_required = is_hma request_id = "test_transfer_fail" metadata = NixlConnectorMetadata() @@ -2912,13 +2933,106 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) time.sleep(0.1) connector.start_load_kv(dummy_ctx) - # check that blocks were marked invalid + results = connector.get_transfer_results(finished_req_ids=set()) + assert request_id in results.finished_recving + assert results.failed_recving == {request_id} invalid_blocks = connector.get_block_ids_with_load_errors() - assert invalid_blocks == {7, 8, 9} + assert invalid_blocks == (set() if is_hma else {7, 8, 9}) + + +class _ScriptedXferWrapper(FakeNixlWrapper): + """Scripts per-handle xfer states; forbids releasing in-flight handles.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # handle -> successive states; the last one repeats. + self.xfer_states: dict[int, list[str]] = {} + self.released: list[int] = [] + + def check_xfer_state(self, handle: int) -> str: + states = self.xfer_states[handle] + return states.pop(0) if len(states) > 1 else states[0] + + def release_xfer_handle(self, handle: int) -> None: + # A posted-but-unfinished transfer cannot be aborted: releasing it + # leaves the RDMA READ armed. Production code must never do this. + assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", ( + f"released in-flight handle {handle}" + ) + self.released.append(handle) + - # ensure request appears in get_finished +def _make_split_read_connector(vllm_config, request_id, states): + """Seed a request with scripted in-flight xfer handles + recv metadata.""" + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + wrapper = _ScriptedXferWrapper("agent") + worker.nixl_wrapper = wrapper + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id=request_id, + local_block_ids=([7, 8, 9],), + kv_transfer_params={ + "remote_block_ids": ([10, 11, 12],), + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": f"prefill-{request_id}", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + }, + ) + worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id] + wrapper.xfer_states = dict(states) + worker._recving_transfers[request_id] = list(states) + return connector, worker, wrapper + + +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) +def test_split_read_failure_defers_report_until_last_handle( + default_vllm_config, dist_init +): + """One half of a split (mixed DRAM/VRAM) read failing must not report the + request — nor invalidate its blocks — while the sibling xfer is still in + flight: a posted READ cannot be aborted and would DMA into blocks the + scheduler could free and reuse. The report happens exactly once, when the + last handle is terminal.""" + request_id = "split_read_partial_failure" + err_handle, live_handle = 11, 22 + connector, worker, wrapper = _make_split_read_connector( + create_vllm_config(), + request_id, + {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]}, + ) + + # Poll 1: one half fails; the sibling is in flight -> nothing reported. _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id in done_recving + assert done_recving == set() + assert connector.get_block_ids_with_load_errors() == set() + assert request_id in worker._recving_metadata + assert wrapper.released == [err_handle] + + # Poll 2: sibling still in flight. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() + + # Poll 3: sibling terminal -> reported exactly once, blocks invalidated. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == {request_id} + assert connector.get_block_ids_with_load_errors() == {7, 8, 9} + assert request_id not in worker._recving_metadata + assert wrapper.released == [err_handle, live_handle] + + # Poll 4: nothing left; no double report. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() @patch( @@ -2938,7 +3052,7 @@ def test_failed_request_skips_kv_postprocessing( default_vllm_config, dist_init, failure_mode ): """Test that failed requests skip KV sync and post-processing in - get_finished(). + get_transfer_results(). This is the core safety behavior: when a KV transfer fails at any stage, the request must still appear in done_recving (so the scheduler can apply @@ -3023,11 +3137,10 @@ def test_failed_request_skips_kv_postprocessing( patch.object(worker, "sync_recved_kv_to_device") as mock_sync, patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess, ): - _, done_recving = connector.get_finished(finished_req_ids=set()) + results = connector.get_transfer_results(finished_req_ids=set()) - # The failed request must appear in done_recving so the scheduler - # can handle it (e.g., trigger recompute via kv_load_failure_policy). - assert request_id in done_recving + assert request_id in results.finished_recving + assert results.failed_recving == {request_id} # Critical: KV sync and post-processing must NOT have been called # since no valid KV data was received for the failed request. @@ -3049,19 +3162,7 @@ def test_failed_request_skips_kv_postprocessing( def test_handles_failing_in_separate_polls_do_not_kill_the_engine( default_vllm_config, dist_init ): - """A request whose handles fail in different polls must not crash the engine. - - One transfer handle is created per remote rank, and when a peer goes away - they do not all fail in the same poll: one errors while another is still - PROC. The first failure reports the request via _failed_recv_reqs and - get_finished() pops its metadata; when the remaining handles fail in a - later poll, the request must be cleaned up without being reported again. - - Reporting twice kills the EngineCore: the scheduler's assert in - _update_from_kv_xfer_finished only expects a finished recv for a request - still waiting for KVs, having moved this one out of - WAITING_FOR_REMOTE_KVS to recompute locally after the first report. - """ + """Failures across polls report once, after the last writer is terminal.""" vllm_config = create_vllm_config() connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) @@ -3107,20 +3208,21 @@ def first_failed(handle): ): _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id in done_recving - assert request_id not in worker._recving_metadata + assert request_id not in done_recving + assert request_id in worker._recving_metadata assert worker._recving_transfers[request_id] == [second] + assert connector.get_block_ids_with_load_errors() == set() - # Poll 2: the second handle fails too. The request was already reported, - # so it must not be reported a second time: the scheduler has since moved - # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a - # finished recv for a request in that state. Its remaining handles are - # still released and the transfer entry removed. + # Only the last terminal handle permits the scheduler to reuse the blocks. with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"): _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id not in done_recving + assert done_recving == {request_id} assert request_id not in worker._recving_transfers + assert request_id not in worker._recving_metadata + assert connector.get_block_ids_with_load_errors() == {1, 2, 3} + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() def _set_test_speculative_config( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 08506d82f246..259e9536c1ca 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -5,16 +5,31 @@ import gc from unittest.mock import MagicMock, patch +import msgspec import pytest import torch from tests.v1.attention.utils import MockMambaBuilder from vllm import LLM, SamplingParams -from vllm.config import KVTransferConfig +from vllm.config import KVTransferConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, +) from vllm.v1.core.single_type_kv_cache_manager import ( FullAttentionManager, SlidingWindowManager, ) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MLAAttentionSpec, +) from .utils import ( create_request, @@ -272,26 +287,24 @@ def test_logical_to_kernel_block_ids_with_hma(): @pytest.mark.cpu_test @pytest.mark.parametrize( - "is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs", + "is_rocm,has_mamba,use_host_buffer,done_recving,expected_syncs", [ - (True, True, False, {"req"}, set(), 1), - (False, True, False, {"req"}, set(), 0), - (True, False, False, {"req"}, set(), 0), - (True, True, True, {"req"}, set(), 0), - (True, True, False, set(), set(), 0), - (True, True, False, {"req"}, {"req"}, 0), + (True, True, False, {"req"}, 1), + (False, False, False, {"req"}, 0), + (False, True, False, {"req"}, 0), + (True, True, True, {"req"}, 0), + (True, True, False, set(), 0), ], ) -def test_sync_device_after_mamba_recv_gates( +def test_sync_device_after_direct_recv_gates( monkeypatch, is_rocm, has_mamba, use_host_buffer, done_recving, - failed_recving, expected_syncs, ): - """Only direct-GPU Mamba receives on ROCm need a device fence.""" + """Only backends requiring visibility fences synchronize direct receives.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, @@ -309,7 +322,7 @@ def test_sync_device_after_mamba_recv_gates( lambda: sync_calls.append(True), ) - worker._sync_device_after_mamba_recv(done_recving, failed_recving) + worker._sync_device_after_direct_recv(done_recving) assert len(sync_calls) == expected_syncs @@ -399,6 +412,8 @@ def test_read_blocks_for_req_expands_remote_ids( ) remote_engine_id = "remote-engine" + worker.region_group_ids = list(range(len(resolved_types))) + worker.dst_region_group_ids = {remote_engine_id: list(worker.region_group_ids)} worker.transfer_topo = MagicMock() # tp_ratio not exercised (all_source_ranks is empty so no reads run), @@ -888,6 +903,45 @@ def _make_mock_worker_for_desc_ids( return worker +@pytest.mark.cpu_test +def test_shared_nixl_region_broadcasts_every_group_block_table(): + """A shared region must receive every group's block table exactly once.""" + worker = _make_mock_worker_for_desc_ids( + num_regions=3, + has_mamba=False, + group_spec_types=(FullAttentionSpec,) * 3, + block_len_per_layer=[100] * 3, + ) + worker.region_group_ids = [0, bw._SHARED_REGION_GROUP_ID, 2] + worker._uses_region_group_mapping = True + + result = worker._compute_desc_ids( + block_ids=([1], [2], [3]), + dst_num_blocks=10, + block_size_ratio=None, + physical_blocks_per_logical=1, + region_num_blocks=[10, 10, 10], + ) + + assert result.tolist() == [1, 11, 12, 13, 23] + + +@pytest.mark.cpu_test +def test_nixl_region_sort_preserves_pipeline_layer_order(): + """Numeric layer order prevents PP peers from pairing different regions.""" + names = [ + "model.layers.10.self_attn", + "model.layers.2.self_attn", + "model.layers.1.self_attn", + ] + + assert sorted(names, key=bw._region_sort_key) == [ + "model.layers.1.self_attn", + "model.layers.2.self_attn", + "model.layers.10.self_attn", + ] + + @pytest.mark.cpu_test def test_get_block_descs_ids_hybrid_ssm(): """Test _compute_desc_ids uses per-group strides for hybrid @@ -1097,10 +1151,7 @@ def test_post_process_zeroes_untransferred_tail(): worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} fa_group = MagicMock(layer_names=["attn.0"]) ssm_group = MagicMock(layer_names=["mamba.0"]) - worker.kv_cache_config = MagicMock( - kv_cache_groups=[fa_group, ssm_group], - transfer_groups=[fa_group, ssm_group], - ) + worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group]) # The cached property filters mamba layers out of the permuted caches. attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) assert len(attn_caches) == 1 and attn_caches[0] is attn_cache @@ -1188,6 +1239,11 @@ def _make_fake_kv_cache_manager(): manager = object.__new__(KVCacheManager) manager.coordinator = MagicMock() + manager.kv_cache_config = MagicMock() + manager.kv_cache_config.kv_cache_groups = ( + MagicMock(host_resident=False), + MagicMock(host_resident=False), + ) manager.coordinator.single_type_managers = ( _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba @@ -1723,16 +1779,6 @@ def test_exchange_clipped_blocks_ssm_positional_states(): assert clipped == ([1, 2, 3], [0, 5, 6, 7]) -@pytest.mark.cpu_test -def test_exchange_clipped_blocks_excludes_nontransfer_groups(): - """Scheduler block tables must match the worker's registered regions.""" - sched = make_nixl_scheduler() - sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True) - sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False - - assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],) - - # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ───────────────────── @@ -1786,6 +1832,95 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): ) +@pytest.mark.cpu_test +@pytest.mark.parametrize("kernel_block_size", [16, 8]) +def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size): + host_num_blocks = 4 + gpu_num_blocks = 9 + spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + ) + kv_cache_config = KVCacheConfig( + num_blocks=gpu_num_blocks, + hisparse_host_num_blocks=host_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=host_num_blocks * spec.page_size_bytes, + layers=["mla.host"], + layer_stride=host_num_blocks * spec.page_size_bytes, + block_stride=spec.page_size_bytes, + host_resident=True, + ), + KVCacheTensor( + size=gpu_num_blocks * spec.page_size_bytes, + layers=["mla.device"], + layer_stride=gpu_num_blocks * spec.page_size_bytes, + block_stride=spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["mla.host"], + spec, + host_resident=True, + ), + KVCacheGroupSpec(["mla.device"], spec), + ], + ) + vllm_config = create_vllm_config(block_size=16) + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_backend.get_kv_cache_shape.return_value = (1, kernel_block_size, 1, 1) + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + physical_per_logical = spec.block_size // kernel_block_size + physical_page_bytes = spec.page_size_bytes // physical_per_logical + host_cache = torch.zeros( + host_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8 + ) + device_cache = torch.zeros( + gpu_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8 + ) + + with ( + patch.object(bw, "NixlWrapper"), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config) + worker.use_mla = True + worker.nixl_wrapper.get_agent_metadata.return_value = b"metadata" + worker.register_kv_caches({"mla.host": host_cache, "mla.device": device_cache}) + + assert worker.num_blocks == gpu_num_blocks * physical_per_logical + assert dict(zip(worker.region_group_ids, worker.region_num_blocks)) == { + 0: host_num_blocks * physical_per_logical, + 1: gpu_num_blocks * physical_per_logical, + } + caches = {"mla.host": host_cache, "mla.device": device_cache} + expected_addrs = [ + caches[name].data_ptr() + block * stride + for name, stride, count in zip( + worker.region_names, + worker.block_stride_per_layer, + worker.region_num_blocks, + strict=True, + ) + for block in range(count) + ] + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs + + @pytest.mark.cpu_test def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): """Hybrid MLA+KDA registration: HMA tensors shared by both layer types @@ -1803,6 +1938,9 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): kv_cache_config = _make_hybrid_mla_kv_cache_config() unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes vllm_config = create_vllm_config(block_size=12) + # The engine rewrites this global value to the smallest scheduler-group + # block size. NIXL must retain the transferable groups' 12-token geometry. + vllm_config.cache_config.block_size = 4 # kv_buffer_device defaults to the *real* platform's device type, which on # a CPU-only test host would make this a host-buffer worker: host xfer # buffers are per-layer, so the HMA shared tensors would not be @@ -1833,7 +1971,13 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): worker.use_mla = True # opt-125m test config is not MLA; force the flag worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata" - tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)] + kernel_page = unified_page // 3 + block_stride = 2 * kernel_page + backing = torch.zeros(12, block_stride, dtype=torch.uint8) + tensors = [ + backing[:, :kernel_page], + backing[:, kernel_page:], + ] # KDA layer first per tensor: exercises the dual-purpose flag merge. worker.register_kv_caches( { @@ -1849,6 +1993,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): # 12-token logical blocks over the 4-token MLA kernel block. assert worker._physical_blocks_per_logical_kv_block == 3 assert worker.block_size == 4 and worker.num_blocks == 12 + assert worker.region_group_ids == [bw._SHARED_REGION_GROUP_ID] * 2 + assert worker.block_stride_per_layer == [block_stride, block_stride] # Both shared tensors are dual-purpose: their FA view is MLA even though # a KDA layer registered them first. assert worker._region_is_mla == [True, True] @@ -1861,8 +2007,20 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks. assert worker.src_blocks_data.shape == (24 + 32, 3) fa_descs = worker.src_blocks_data[:24] - assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3 + expected_addrs = [ + tensor.data_ptr() + block * block_stride + for tensor in tensors + for block in range(12) + ] + assert fa_descs[:, 0].tolist() == expected_addrs assert all(size == unified_page // 3 for size in fa_descs[:, 1]) + worker.nixl_wrapper.register_memory.assert_called_once() + metadata = msgspec.msgpack.decode( + worker.xfer_handshake_metadata.agent_metadata_bytes, + type=NixlAgentMetadata, + ) + assert metadata.block_strides == [block_stride, block_stride] + assert metadata.region_num_blocks == [12, 12] @pytest.mark.cpu_test @@ -1909,6 +2067,8 @@ def test_push_write_hybrid_mla_replicates_attention(): worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}} worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]} worker.src_xfer_handles_by_block_size = {4: 300} + worker.region_group_ids = [0, 1] + worker.dst_region_group_ids = {engine_id: [0, 1]} worker._sending_transfers = defaultdict(list) worker._sending_transfers_lock = threading.Lock() worker.kv_cache_config = _make_hybrid_mla_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 32908f78adf6..607e4822b974 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -186,11 +186,16 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker.host_buffer_kv_cache_layout = "NHD" worker._physical_blocks_per_logical_kv_block = 1 worker._logical_num_blocks = num_blocks + worker.region_mem_types = [] worker.region_group_ids = [] worker.region_mem_types = [] worker._mixed_mem_types = False worker.region_names = [] worker.region_num_blocks = [] + worker._mixed_mem_types = False + worker._desc_is_dram_by_block_size = {} + worker._desc_pos_by_block_size = {} + worker._dram_src_handles_by_block_size = {} worker._region_is_mla = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] @@ -304,8 +309,6 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo # would not be deduplicated. Pin it to the faked device type. vllm_config.kv_transfer_config.kv_buffer_device = "cuda" - from unittest.mock import MagicMock - fake_backend = MagicMock() fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] fake_backend.get_name.return_value = "FLASHMLA" @@ -552,6 +555,16 @@ def _make_remote_meta( ) +def _register_remote_agents(worker, metadata, tp_size): + """Mirror the async handshake callback that publishes prepared agents.""" + worker._remote_agents[metadata.engine_id] = { + (0, rank): worker.add_remote_agent( + metadata, remote_tp_rank=rank, remote_tp_size=tp_size + ) + for rank in range(tp_size) + } + + def _owned_byte_ranges(worker, group_logical_ids): """Byte ranges owned by a request: for each HMA region tensor, every logical block id of every group maps to one unified page.""" @@ -609,8 +622,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): remote_num_logical=12, remote_ssm_sizes=(24, 32), ) - for rank in (0, 1): - worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + _register_remote_agents(worker, meta_r, 2) # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. @@ -626,7 +638,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): "remote_block_ids": remote_ids, "remote_engine_id": "remote-engine", "remote_request_id": "prefill-req-b", - "remote_host": "localhost", + "remote_host": "remote-host", "remote_port": 1234, "tp_size": 2, }, @@ -709,8 +721,7 @@ def _run_hetero_case( remote_num_logical=max(2 * n_remote + 4, 8), remote_ssm_sizes=(48 // tp_size, 64 // tp_size), ) - for rank in range(tp_size): - worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + _register_remote_agents(worker, meta_r, tp_size) # Sparse ids so neighbors exist between the request's blocks. local_attn = [2 * i + 1 for i in range(n_local)] @@ -726,7 +737,7 @@ def _run_hetero_case( "remote_block_ids": remote_ids, "remote_engine_id": "remote-engine", "remote_request_id": "prefill-req-b", - "remote_host": "localhost", + "remote_host": "remote-host", "remote_port": 1234, "tp_size": tp_size, }, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 5b7760e7e33c..ce233bbb603b 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -15,7 +15,7 @@ * The worker matches D registrations against P finished blocks (both scenario directions) and forwards non-PUSH_REG NIXL notifs to the main thread's ``_get_new_notifs``. -* ``get_finished`` enqueues evictions for the writer. +* ``get_transfer_results`` enqueues evictions for the writer. """ from __future__ import annotations @@ -32,6 +32,9 @@ import msgspec import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( PUSH_REG_NOTIF_PREFIX, NixlAgentMetadata, @@ -347,6 +350,8 @@ def fresh(cls) -> _StubWriterWorker: # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} + w._pending_recv_notifs = {} + w._failed_inflight_recvs = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() w._reqs_to_send = {} @@ -742,21 +747,19 @@ def test_get_new_notifs_processes_forwarded_completion_notif(self): assert notified == set() assert request_id in w._recving_transfers - def test_get_finished_evicts_completed_state(self): - """``get_finished`` should enqueue evictions and wake the writer.""" + def test_get_transfer_results_evicts_completed_state(self): + """Transfer completion should enqueue evictions and wake the writer.""" w = _StubWriterWorker.fresh() - # Stub the base ``get_finished`` to return one done_sending entry. - # Patch via the MRO's parent class. with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=({"req-done"}, set()), + "get_transfer_results", + return_value=KVConnectorTransferResults(finished_sending={"req-done"}), ): - done_sending, done_recving = w.get_finished() + results = w.get_transfer_results() - assert "req-done" in done_sending - assert done_recving == set() + assert results.finished_sending == {"req-done"} + assert results.finished_recving == set() # Eviction enqueued for the writer. evicted = [] while True: @@ -928,19 +931,21 @@ def test_handle_push_reg_idempotent_for_same_request_id(self): assert len(w._pending_d_registrations) == 1 assert w.start_push_calls == [] - def test_get_finished_enqueues_eviction_for_each_done_request(self): - """``get_finished`` must enqueue an eviction for every request + def test_get_transfer_results_enqueues_each_completed_request(self): + """``get_transfer_results`` must enqueue every completed request in ``done_sending`` so the writer can drop stale matching state. Unlike the happy-path test, this verifies the *cardinality*: N completed requests -> N evictions, in order.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=({"req-1", "req-2", "req-3"}, set()), + "get_transfer_results", + return_value=KVConnectorTransferResults( + finished_sending={"req-1", "req-2", "req-3"} + ), ): - done_sending, _ = w.get_finished() - assert done_sending == {"req-1", "req-2", "req-3"} + results = w.get_transfer_results() + assert results.finished_sending == {"req-1", "req-2", "req-3"} evicted: list[str] = [] while True: @@ -950,23 +955,47 @@ def test_get_finished_enqueues_eviction_for_each_done_request(self): break assert sorted(evicted) == ["req-1", "req-2", "req-3"] - def test_get_finished_with_no_completions_does_not_enqueue_eviction(self): + def test_empty_transfer_results_do_not_enqueue_eviction(self): """If there's nothing newly done, no eviction should be enqueued. The wake event IS still set because ``get_finished`` always wakes the writer to drain notifs.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=(set(), set()), + "get_transfer_results", + return_value=KVConnectorTransferResults(), ): - done_sending, done_recving = w.get_finished() - assert done_sending == set() - assert done_recving == set() + results = w.get_transfer_results() + assert results.finished_sending == set() + assert results.finished_recving == set() assert w._evict_finished_inbox.qsize() == 0 # Wake set so the writer drains NIXL notifs even when idle. assert w._push_writer_wake.is_set() + def test_failed_push_transfer_finishes_sending(self): + w = _StubWriterWorker.fresh() + w.nixl_wrapper = MagicMock() + w.nixl_wrapper.check_xfer_state.return_value = "ERR" + w.xfer_stats = MagicMock() + w._log_failure = MagicMock() + w._sending_transfers["req-failed"] = [17] + w._reqs_to_send["req-failed"] = time.perf_counter() + 30 + w._reqs_to_process.add("req-failed") + + with patch.object( + NixlPushConnectorWorker.__mro__[1], + "get_transfer_results", + return_value=KVConnectorTransferResults(), + ): + results = w.get_transfer_results() + + assert results.finished_sending == {"req-failed"} + assert results.finished_recving == set() + assert "req-failed" not in w._sending_transfers + assert "req-failed" not in w._reqs_to_send + assert "req-failed" not in w._reqs_to_process + w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17) + def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): """A completion notif for a request the worker doesn't know about should be logged but not crash.""" @@ -1160,6 +1189,7 @@ def _mla_worker_writing_to(d_ranks): ) } w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids + w.dst_region_group_ids[engine_id] = [0] w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}} w.src_xfer_handles_by_block_size = {16: 2000} w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}} @@ -1240,6 +1270,8 @@ def test_csa_linear_writes_every_group_to_each_target(self): ) } w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids + w.region_group_ids = [0, 1, 2, 3] + w.dst_region_group_ids[engine_id] = [0, 1, 2, 3] w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}} w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]} diff --git a/tests/v1/kv_connector/unit/test_output_aggregator.py b/tests/v1/kv_connector/unit/test_output_aggregator.py index b083ccef9819..e757a658425f 100644 --- a/tests/v1/kv_connector/unit/test_output_aggregator.py +++ b/tests/v1/kv_connector/unit/test_output_aggregator.py @@ -15,12 +15,14 @@ def __init__( finished_sending: set[str] | None = None, finished_recving: set[str] | None = None, invalid_block_ids: set[int] | None = None, + failed_recving: set[str] | None = None, expected_finished_count: int = 0, ): self.kv_connector_output = KVConnectorOutput( finished_sending=finished_sending, finished_recving=finished_recving, invalid_block_ids=invalid_block_ids or set(), + failed_recving=failed_recving or set(), expected_finished_count=expected_finished_count, ) @@ -73,7 +75,9 @@ def test_aggregate_workers_output(): output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4}) output2 = DummyModelRunnerOutput( - finished_recving={"req2"}, invalid_block_ids={4, 5} + finished_recving={"req2"}, + invalid_block_ids={4, 5}, + failed_recving={"req3"}, ) aggregated = aggregator.aggregate([output1, output2]) @@ -83,6 +87,12 @@ def test_aggregate_workers_output(): assert aggregated.finished_sending is None assert aggregated.finished_recving == {"req2"} assert aggregated.invalid_block_ids == {3, 4, 5} + assert not aggregated.failed_recving + + output1 = DummyModelRunnerOutput(finished_recving={"req3"}) + output2 = DummyModelRunnerOutput(finished_recving={"req3"}) + aggregated = aggregator.aggregate([output1, output2]) + assert aggregated.kv_connector_output.failed_recving == {"req3"} def test_aggregate_workers_output_with_expected_finished_count(): diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 80cab94b3484..7b70065d7f95 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -539,6 +539,7 @@ def make_nixl_scheduler( sched._heartbeat_interval = kv_lease_duration // 6 # Fields touched by build_connector_meta / request_finished: sched._reqs_need_recv = {} + sched._hisparse_host_blocks_to_recv = {} sched._reqs_need_send = {} sched._reqs_in_batch = set() sched._reqs_not_processed = set() diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 251cef325143..16539677c490 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -62,7 +62,7 @@ def _make_offloading_config( normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use if groups is None: - groups = (OffloadingGroupConfig(16, ("layer",)),) + groups = (OffloadingGroupConfig(16, ("layer",), 0),) return OffloadingConfig( groups=groups, @@ -493,8 +493,8 @@ def test_offloading_spec_has_replicated_layout_default(): def test_offloading_spec_uses_normalized_chunk_geometry(): groups = ( - OffloadingGroupConfig(12, ("full_layer",)), - OffloadingGroupConfig(16, ("mla_layer",)), + OffloadingGroupConfig(12, ("full_layer",), 0), + OffloadingGroupConfig(16, ("mla_layer",), 1), ) spec = _create_spec( groups=groups, diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 26da4079b470..8cd1ae9b10a1 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -26,8 +26,11 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: OffloadingGroupConfig( tokens_per_block=tokens_per_block, layer_names=(layer_name,), + group_id=group_id, + ) + for group_id, (tokens_per_block, layer_name) in enumerate( + kwargs.get("groups", ()) ) - for tokens_per_block, layer_name in kwargs.get("groups", ()) ), worker_kv_bytes_per_block=0, enable_kv_cache_events=False, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index cbc7f2061330..a62dacc31028 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -400,6 +400,8 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.kv_cache_layout = layout.name vllm_config.cache_config.num_gpu_blocks_override = None + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None pages = [spec.page_size_bytes for spec in specs.values()] num_blocks = 4 @@ -458,6 +460,8 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.kv_cache_layout = layout.name vllm_config.cache_config.num_gpu_blocks_override = None + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None pages = [spec.page_size_bytes for spec in specs.values()] kv_cache_config = get_kv_cache_config_from_groups( @@ -507,6 +511,8 @@ def test_mixed_page_byte_placement_is_dcp_invariant(): vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name vllm_config.cache_config.num_gpu_blocks_override = None vllm_config.parallel_config.decode_context_parallel_size = dcp_size + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None page_bytes = sum(spec.page_size_bytes for spec in specs.values()) config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4) diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py index 58c23ce285b1..6da14b4e426b 100644 --- a/tests/v1/test_kv_cache_spec_registry.py +++ b/tests/v1/test_kv_cache_spec_registry.py @@ -16,6 +16,7 @@ ChunkedLocalAttentionManager, CrossAttentionManager, FullAttentionManager, + HiSparseSourceManager, MambaManager, SingleTypeKVCacheManager, SinkFullAttentionManager, @@ -27,6 +28,7 @@ CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + KVCacheGroupRole, KVCacheSpec, KVCacheSpecKind, MambaSpec, @@ -176,6 +178,19 @@ def test_builtin_kvcache_specs_registered(self): is spec_uniform_base_map[spec_cls] ) + @pytest.mark.parametrize("role", list(KVCacheGroupRole)) + def test_mla_manager_selection_by_role(self, role): + """Only the source role overrides ordinary MLA manager selection.""" + expected = ( + HiSparseSourceManager + if role == KVCacheGroupRole.HISPARSE_SOURCE + else FullAttentionManager + ) + assert ( + KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role) + is expected + ) + @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) def test_custom_spec_register(self, spec_cls): """A decorated custom spec resolves to the declared manager.""" diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 906525494f4b..dc75ba4ea19b 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -13,15 +13,18 @@ import torch from tests.v1.attention.utils import dense_kv_cache_views -from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + HiSparseResidentSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, KVCacheTensor, MLAAttentionSpec, + SparseCacheRole, compute_layout_strides, ) from vllm.v1.worker.gpu import attn_utils @@ -36,6 +39,60 @@ ) +@pytest.mark.parametrize( + ("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"), + [ + (True, 256, [64], [64], 64), + (True, 64, [32, 64], [16, 32], 32), + (True, 64, [MultipleOf(16)], [32], 32), + (True, 64, [64], [32], None), + (False, 256, [64], [64], 256), + ], +) +def test_get_kv_cache_spec_resolves_hisparse_block_size( + monkeypatch, enabled, block_size, main_sizes, indexer_sizes, expected +): + """Resolve shared MLA geometry before planning; leave other specs alone.""" + specs = { + "main": MLAAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16 + ), + "indexer": MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + "dense": FullAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ), + } + layers = {} + for name, sizes in zip(specs, [main_sizes, indexer_sizes, [block_size]]): + backend = SimpleNamespace( + customize_spec=AttentionBackend.customize_spec, + get_supported_kernel_block_sizes=lambda sizes=sizes: sizes, + ) + layers[name] = SimpleNamespace( + get_kv_cache_spec=lambda _, spec=specs[name]: spec, + get_attn_backend=lambda backend=backend: backend, + ) + monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object() if enabled else None) + ) + if expected is None: + with pytest.raises(ValueError, match="supported by every sparse"): + attn_utils.get_kv_cache_spec(config) + return + + resolved = attn_utils.get_kv_cache_spec(config) + assert resolved["main"].block_size == resolved["indexer"].block_size == expected + assert resolved["dense"] is specs["dense"] + assert all(spec.block_size == block_size for spec in specs.values()) + + class _FakeMetadataBuilder: def __init__(self, support: AttentionCGSupport): self.support = support @@ -390,3 +447,71 @@ def test_copy_kv_cache_blocks_with_virtual_block_splitting( torch.testing.assert_close( cache[dst_start + physical_idx], expected[layer_idx][physical_idx] ) + + +def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs(): + """Host tensors get their own backing; view-less specs keep the raw one.""" + spec = FullAttentionSpec( + block_size=2, num_kv_heads=1, head_size=4, dtype=torch.float32 + ) + page = spec.page_size_bytes + resident_spec = HiSparseResidentSpec(block_size=2, page_size=page) + device_size = 4 * page + config = KVCacheConfig( + num_blocks=4, + hisparse_host_num_blocks=3, + kv_cache_tensors=[ + KVCacheTensor( + size=3 * page, + layers=["source"], + layer_stride=3 * page, + block_stride=page, + host_resident=True, + ), + KVCacheTensor( + size=device_size, + layers=["indexer"], + layer_stride=device_size, + block_stride=page, + ), + KVCacheTensor( + size=device_size, + layers=["resident"], + layer_stride=device_size, + block_stride=page, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["source"], spec, host_resident=True), + KVCacheGroupSpec(["indexer"], spec), + KVCacheGroupSpec(["resident"], resident_spec), + ], + ) + host_buffers: list[torch.Tensor] = [] + + def host_allocator(size: int) -> torch.Tensor: + host_buffers.append(torch.zeros(size, dtype=torch.int8)) + return host_buffers[-1] + + caches = allocate_hisparse_kv_caches( + config, + torch.device("cpu"), + KVCacheLayout.LBHNC, + [2, 2, 2], + SimpleNamespace(allocate=host_allocator), + ) + assert len(config.kv_cache_tensors) == 3 + + assert [buf.numel() for buf in host_buffers] == [3 * page] + assert caches["source"].shape[0] == 3 + assert ( + caches["source"].untyped_storage().data_ptr() + == host_buffers[0].untyped_storage().data_ptr() + ) + assert caches["indexer"].shape[0] == 4 + backing = caches["resident"] + assert backing.dtype == torch.int8 and backing.numel() >= device_size + assert ( + backing.untyped_storage().data_ptr() + == caches["indexer"].untyped_storage().data_ptr() + ) diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index f84c91872faa..4cc8f9ffe2fd 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -134,6 +134,7 @@ def make_spec(head_size): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.num_gpu_blocks_override = None vllm_config.cache_config.kv_cache_layout = "BLHNC" + vllm_config.attention_config.hisparse_config = None config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024) views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None) buf_ptr = views["g1.big"].data_ptr() diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index e9c7a3b290b2..920c9a25f062 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -351,6 +351,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected( speculator = object.__new__(_TestSpeculator) speculator.num_speculative_steps = 4 speculator.current_draft_step = torch.tensor(0) + speculator.slot_mapping_observer = None + speculator.host_mirror_forward_observer = None speculator.input_buffers = SimpleNamespace( positions=torch.arange(2), query_start_loc=torch.arange(3), diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index 7fb40fff3711..03c483dce165 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -5,9 +5,11 @@ from unittest.mock import Mock import pytest +import torch import vllm.v1.worker.gpu.kv_connector as kv_connector_module from vllm.config import KVTransferConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector @@ -18,9 +20,9 @@ def _make_connector( backend = Mock() backend.handle_preemptions.side_effect = lambda _: events.append("handle") backend.bind_connector_metadata.side_effect = lambda _: events.append("bind") - backend.start_load_kv.side_effect = lambda _: events.append("start") + backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start") backend.wait_for_save.side_effect = lambda: events.append("wait") - backend.get_finished.side_effect = lambda _: (set(), set()) + backend.get_transfer_results.return_value = KVConnectorTransferResults() backend.get_block_ids_with_load_errors.return_value = set() backend.get_kv_connector_stats.return_value = None backend.get_kv_connector_kv_cache_events.return_value = None @@ -61,7 +63,15 @@ def test_load_start_phase( connector = _make_connector(monkeypatch, events) output = _scheduler_output(has_sync_kv_loads) - connector.pre_forward(output) # type: ignore[arg-type] + request_indices = torch.tensor([3, 1]) + request_ids = ["first", "second"] + attn_metadata = {"layer": object()} + connector.pre_forward( # type: ignore[arg-type] + output, + request_state_indices=request_indices, + request_ids=request_ids, + attn_metadata=attn_metadata, + ) assert events == ( ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"] ) @@ -69,6 +79,16 @@ def test_load_start_phase( connector.post_forward(set()) assert events == ["handle", "bind", "start", "wait", "clear"] + kwargs = connector.kv_connector.start_load_kv.call_args.kwargs + assert kwargs["request_state_indices"] is request_indices + assert kwargs["request_ids"] is request_ids + assert kwargs["attn_metadata"] is attn_metadata + + # A subsequent step without a forward must not reuse the prior batch. + connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type] + assert connector.kv_connector.start_load_kv.call_count == 2 + assert connector.kv_connector.start_load_kv.call_args.kwargs == {} + def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch): events: list[str] = [] diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 99331e7b354e..bb77f90ced2a 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -312,6 +312,10 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_select_common_block_size_without_active_backends_uses_manager_size(): + assert select_common_block_size(256, []) == 256 + + def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch): monkeypatch.setattr(current_platform, "is_rocm", lambda: True) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 1118ce589ca5..c9ce409f5cb9 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -237,6 +237,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner: runner.attn_groups = None runner.kv_cache_config = None runner.use_aux_hidden_state_outputs = False + runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR return runner diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index fcace4ae4a20..067f31e9350c 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -89,6 +89,7 @@ def _make_runner( return SimpleNamespace( num_speculative_steps=num_spec_steps, adaptive_verification=None, + rejection_sampler=None, decode_query_len=num_spec_steps + 1, is_pooling_model=False, is_encoder_decoder=False, diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index a9579f1af69b..bacd1cad8482 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -5,8 +5,10 @@ from types import SimpleNamespace from typing import Any, cast +import pytest import torch +import vllm.v1.hisparse.binding as hisparse_binding import vllm.v1.worker.gpu.attn_utils as attn_utils import vllm.v1.worker.gpu_model_runner as gpu_model_runner from vllm.v1.worker.gpu_worker import Worker @@ -26,7 +28,8 @@ def __exit__(self, *args: Any) -> None: self.active = False -def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: +@pytest.mark.parametrize("hisparse", [False, True]) +def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> None: scope = _AllocationScope() kv_caches = {"layer": torch.empty(0)} @@ -38,23 +41,48 @@ def bind(*args, **kwargs): assert not scope.active monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate) - monkeypatch.setattr(attn_utils, "bind_kv_cache", bind) + monkeypatch.setattr(attn_utils, "bind_kv_cache_to_layers", bind) monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {}) + hisparse_bindings = [] + if hisparse: + host_pool = SimpleNamespace() + + def bind_hisparse(**kwargs): + hisparse_bindings.append(kwargs) + assert scope.active + assert kwargs["kv_caches"] is kv_caches + assert kwargs["host_pool"] is host_pool + return [ + SimpleNamespace( + view=SimpleNamespace(cache=torch.empty(1, 1, 1), block_size=1), + runtime=SimpleNamespace(), + ) + ] + + monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool) + monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) + monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse) + config = SimpleNamespace( + attention_config=SimpleNamespace( + hisparse_config=object() if hisparse else None + ), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None), model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")), ) result = attn_utils.init_kv_cache( - [], {}, SimpleNamespace(kv_cache_groups=[]), torch.device("cpu"), [], config, kv_cache_allocation_context=scope, + block_tables=SimpleNamespace(), ) + assert len(hisparse_bindings) == int(hisparse) assert result is kv_caches assert not scope.active diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index df339135eb5a..04f456a7c9cd 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -35,10 +35,15 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup): """Startup must warm fixed drafts before calibration and retain its manager.""" manager = SimpleNamespace(cost_tables=None) - runner = SimpleNamespace(adaptive_verification=manager) + rejection_sampler = SimpleNamespace(enable_adaptive_verification=True) + runner = SimpleNamespace( + adaptive_verification=manager, + rejection_sampler=rejection_sampler, + ) def run_steps(model_runner, execute, sample): assert model_runner.adaptive_verification is None + assert not model_runner.rejection_sampler.enable_adaptive_verification if fail_warmup: raise RuntimeError("warmup failed") @@ -50,3 +55,4 @@ def run_steps(model_runner, execute, sample): warmup.warmup_kernels(runner, _fail, _fail) assert runner.adaptive_verification is manager assert manager.cost_tables is None + assert rejection_sampler.enable_adaptive_verification diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a4474c8a5843..86c8688906e9 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,13 +1,897 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections import deque +from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import MagicMock, call +import numpy as np +import pytest import torch +import vllm.v1.attention.backends.mla.index_group as index_group_module +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from vllm.config import CUDAGraphMode from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import ( + worker as hisparse_worker_module, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, + _flatten_row_mirrors, + _select_written_row_mirrors, + _SlotMappingStaging, +) from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 -from vllm.v1.worker.utils import bind_kv_cache +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror +from vllm.v1.worker.utils import ( + bind_kv_cache, + bind_kv_cache_to_layers, + copy_kv_cache_blocks_inplace, +) + + +def _make_hisparse_worker() -> HiSparseConnectorWorker: + worker = object.__new__(HiSparseConnectorWorker) + worker.refines_row_mirrors = False + worker._slot_mapping_staging = None + worker._row_mirror_num_rows = 0 + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._pending_dma_descriptors = deque() + worker._dma_free_descriptors = [] + worker.host_write_events = (MagicMock(), MagicMock()) + worker.host_write_event = worker.host_write_events[1] + worker._next_host_write_event = 0 + worker.dma_stream = None + worker.shared_host_region = None + worker._completed_host_copy_dst_ids = [] + worker._slot_mapping_staging = None + return worker + + +def test_hisparse_row_mirrors_follow_runner_request_order(): + first = SparseKVRowMirror((1,), 10, 1) + second = SparseKVRowMirror((2,), 20, 1) + row_mirrors = {"first": (first,), "second": (second,)} + + assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == ( + second, + first, + ) + + +def test_hisparse_gpu_slots_select_written_rows_from_async_envelope(): + """Rejected speculative rows must never become durable host KV.""" + candidates = ( + SparseKVRowMirror((100, 500), 1000, 4), + SparseKVRowMirror((20, 300), 2000, 4), + ) + source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64) + + assert _select_written_row_mirrors(candidates, source_slots, 0) == ( + SparseKVRowMirror((102, 502), 1002, 2), + SparseKVRowMirror((21, 301), 2001, 3), + ) + + +def test_hisparse_written_rows_preserve_scheduler_page_boundaries(): + candidates = ( + SparseKVRowMirror((0,), 100, 4), + SparseKVRowMirror((4,), 104, 4), + ) + + assert _select_written_row_mirrors( + candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0 + ) == ( + SparseKVRowMirror((2,), 102, 2), + SparseKVRowMirror((4,), 104, 2), + ) + + +def test_hisparse_written_rows_without_host_destinations_are_skipped(): + assert ( + _select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == () + ) + + +def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch): + """Separate context and query writes must share one mirror phase.""" + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.refines_row_mirrors = True + state = _SlotMappingStaging( + MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64) + ) + worker._slot_mapping_staging = state + handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1)) + worker.cache_layer_names = ["layer"] + worker.cache_handles = [handle] + worker._layer_mirror_callbacks = (MagicMock(),) + worker._row_mirror_num_rows = 1 + worker._submitted_mirror_layers = set() + main_stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + + handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64) + worker._stage_row_mirror_mapping(2) + handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64) + worker._stage_row_mirror_mapping(2) + + torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10])) + assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2 + assert state.event.record.call_args_list == [call(state.stream)] * 2 + assert state.num_tokens == 4 + assert state.source_index == 1 + + +@pytest.mark.parametrize( + ("mode", "expected_calls"), + [(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)], +) +def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary( + monkeypatch, mode, expected_calls +): + """Piecewise replay must pipeline DMA; FULL replay uses the batch fallback.""" + handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle) + handle.dummy_batch = False + handle.decode_batch = False + handle.submit_layer_mirror = MagicMock() + monkeypatch.setattr( + hisparse_runtime_module, + "get_forward_context", + lambda: SimpleNamespace(cudagraph_runtime_mode=mode), + ) + + handle.finish_kv_update() + + assert handle.submit_layer_mirror.call_count == expected_calls + + +def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): + backing = torch.full((10, 2, 3), -1, dtype=torch.float32) + cache = backing[1:9] + cache[2:4] = 7 + cache[6:8] = 11 + + copy_kv_cache_blocks_inplace( + [cache], + num_blocks=4, + kv_cache_block_copies=[ + KVCacheBlockCopy(1, 0), + KVCacheBlockCopy(3, 2), + ], + ) + + torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7)) + torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11)) + assert (backing[0] == -1).all() + assert (backing[9] == -1).all() + + +@pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)]) +@pytest.mark.parametrize("has_copies", [False, True]) +def test_hisparse_host_copy_waits_for_writes_on_writer_only( + monkeypatch, shared, rank, has_copies +): + """Host CoW must read completed writes; shared non-writers only barrier.""" + worker = _make_hisparse_worker() + cache = torch.zeros((2, 1, 1)) + worker.host_caches = [cache] + worker.host_num_blocks = 2 + worker.shared_host_region = object() if shared else None + event = MagicMock() + event.synchronize.side_effect = lambda: cache[0].fill_(7) + barrier = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank + ) + monkeypatch.setattr( + hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier) + ) + copies = [KVCacheBlockCopy(0, 1)] if has_copies else [] + + worker._copy_host_blocks(copies, event) + + writes = has_copies and (not shared or rank == 0) + assert cache[1].item() == (7 if writes else 0) + assert event.synchronize.call_count == int(writes) + assert barrier.call_count == int(shared and has_copies) + assert worker._completed_host_copy_dst_ids == ([1] if has_copies else []) + + +def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch): + worker = _make_hisparse_worker() + worker.request_state_indices = torch.arange(4, dtype=torch.int32) + worker._pending_invalid_block_ids = [5] + invalidations = [] + worker.invalidate_blocks = lambda blocks, states: invalidations.append( + (blocks.copy(), states.clone()) + ) + original_ptr = worker.request_state_indices.data_ptr() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32)) + + assert worker.request_state_indices.data_ptr() == original_ptr + assert worker.request_state_indices.tolist() == [3, 1, -1, -1] + assert len(invalidations) == 1 + assert invalidations[0][0] == [5] + torch.testing.assert_close( + invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32) + ) + assert worker._pending_invalid_block_ids == [] + + +def test_hisparse_pre_forward_transfer_builds_page_descriptors(): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.kernel_block_size = 2 + source = torch.empty((3, 2, 4), dtype=torch.uint8) + destination = torch.empty((6, 4), dtype=torch.uint8) + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0)) + ] + worker._dma_free_descriptors = [] + worker._submit_dma_descriptors = MagicMock() + + worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)]) + + descriptors, count = worker._submit_dma_descriptors.call_args.args + assert count == 1 + assert descriptors.src[:count].tolist() == [ + source.data_ptr() + source.stride(0) * source.element_size() + ] + assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4] + assert descriptors.sizes[:count].tolist() == [2 * 4] + assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,) + + +def test_hisparse_eager_mirror_records_transfer_without_page_copy(): + worker = _make_hisparse_worker() + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True)) + ] + worker._record_transfer_completion = MagicMock() + worker._enqueue_transfers = MagicMock() + transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] + + worker._submit_transfers(transfers) + + worker._record_transfer_completion.assert_called_once_with(transfers) + worker._enqueue_transfers.assert_not_called() + + +def test_hisparse_lazy_mirror_copies_transfer_pages(): + worker = _make_hisparse_worker() + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False)) + ] + worker._record_transfer_completion = MagicMock() + worker._enqueue_transfers = MagicMock() + transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] + + worker._submit_transfers(transfers) + + worker._enqueue_transfers.assert_called_once_with(transfers) + worker._record_transfer_completion.assert_not_called() + + +def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.kernel_block_size = 2 + source = torch.empty((3, 2, 4), dtype=torch.uint8) + destination = torch.empty((6, 4), dtype=torch.uint8) + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(resident_source_index=0), decode_batch=True + ) + ] + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker.dma_stream = MagicMock() + worker.host_write_event = MagicMock() + worker._dma_free_descriptors = [] + worker._pending_dma_descriptors = deque() + worker._pending_transfer_events = deque() + worker._enqueued_transfer_ids = [] + worker._dma_submitted = False + current_stream = MagicMock() + event = MagicMock() + swap_blocks_batch = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch, "Event", lambda: event) + monkeypatch.setattr( + hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch + ) + + worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),)) + worker._enqueue_row_dma(range(1)) + + src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args + assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4] + assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4] + assert sizes.tolist() == [4] + worker.dma_stream.wait_stream.assert_called_once_with(current_stream) + worker.host_write_event.record.assert_called_once_with(worker.dma_stream) + event.record.assert_called_once_with(worker.dma_stream) + assert worker._dma_submitted + + +def test_hisparse_row_dma_uses_resident_spans(): + worker = _make_hisparse_worker() + source = torch.empty((2, 2, 4), dtype=torch.uint8) + destination = torch.empty((12, 4), dtype=torch.uint8) + worker.is_host_writer = True + worker.kernel_block_size = 2 + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(resident_source_index=0), decode_batch=True + ) + ] + worker._set_row_mirrors( + ( + SparseKVRowMirror((0,), 4, 2), + SparseKVRowMirror((2,), 9, 1), + ) + ) + worker._dma_free_descriptors = [] + worker._submit_dma_descriptors = MagicMock() + + worker._enqueue_row_dma(range(1)) + + descriptors, count = worker._submit_dma_descriptors.call_args.args + assert count == 2 + assert descriptors.src[:count].tolist() == [ + source.data_ptr(), + source.data_ptr() + 2 * 4, + ] + assert descriptors.dst[:count].tolist() == [ + destination.data_ptr() + 4 * 4, + destination.data_ptr() + 9 * 4, + ] + assert descriptors.sizes[:count].tolist() == [2 * 4, 4] + + +def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch): + dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64) + req_ids = torch.tensor([0, 1, 2], dtype=torch.int32) + leader = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=True, + invalidate_written_slots=MagicMock(), + ) + follower = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=3, + num_decode_tokens=2, + req_id_per_token=req_ids, + mirror_slot_mapping=dst_slots, + ) + for runtime in (leader, follower) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),)) + worker._dma_submitted = False + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._post_forward_transfers = [] + worker._enqueue_row_dma = MagicMock() + worker.host_write_event = MagicMock() + worker._forward_ready_event = MagicMock() + current_stream = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + + worker.finish_forward() + + worker._forward_ready_event.record.assert_called_once_with() + worker._enqueue_row_dma.assert_called_once_with( + (0, 1), ready_event=worker._forward_ready_event + ) + leader.invalidate_written_slots.assert_called_once() + torch.testing.assert_close( + leader.invalidate_written_slots.call_args.args[0], dst_slots + ) + follower.invalidate_written_slots.assert_not_called() + worker.host_write_event.record.assert_called_once_with(current_stream) + + +def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors(): + slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for _ in range(2) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) + worker._per_layer_mirrored = {0, 1} + worker._submitted_mirror_layers = {0, 1} + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_not_called() + + +def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch): + runtime = SimpleNamespace(eager_host_mirror=False) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)] + transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True) + worker._post_forward_transfers = [transfer] + worker._forward_ready_event = MagicMock() + worker._enqueue_host_mirror = MagicMock() + worker._submit_transfers = MagicMock() + worker._dma_submitted = False + worker._submitted_mirror_layers = set() + worker._finish_mirror_phase = MagicMock() + worker._release_completed_dma_descriptors = MagicMock() + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.finish_forward() + + worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event) + worker._submit_transfers.assert_called_once_with([transfer]) + worker.host_write_event.record.assert_called_once_with(stream) + + +def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group(): + slots = torch.tensor([7, 8], dtype=torch.int64) + source_indices = [0, 0, 1, 1, 1, 1, 2] + handles = [ + SimpleNamespace( + runtime=SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + resident_source_index=source_index, + ), + decode_batch=False, + host_mirror_required=layer_index < 5, + num_actual_tokens=2 if layer_index < 5 else 0, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for layer_index, source_index in enumerate(source_indices) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),)) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_ready_events = tuple(MagicMock() for _ in handles) + worker._enqueue_row_dma = MagicMock() + + for layer_index in range(5): + worker._enqueue_layer_mirror(layer_index) + + worker._enqueue_row_dma.assert_called_once_with( + (0, 1), ready_event=worker._layer_ready_events[1] + ) + worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1]) + assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),) + assert worker._enqueue_row_dma.call_args_list[1].kwargs == { + "ready_event": worker._layer_ready_events[-1] + } + + +def test_hisparse_finish_forward_rejects_partial_per_layer_mirror(): + slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for _ in range(2) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) + worker._per_layer_mirrored = {0} + worker._submitted_mirror_layers = set() + worker._enqueue_row_dma = MagicMock() + + with pytest.raises(RuntimeError, match="did not mirror every active layer"): + worker._enqueue_host_mirror() + + +def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch): + current_stream = MagicMock() + worker = _make_hisparse_worker() + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker.is_host_writer = True + worker._dma_submitted = True + worker.host_write_event = MagicMock() + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0 + ) + ] + worker._post_forward_transfers = [] + worker._pending_dma_descriptors = deque() + worker._dma_free_descriptors = [] + worker._forward_ready_event = MagicMock() + worker._finish_mirror_phase = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + + worker.finish_forward() + + current_stream.wait_event.assert_called_once_with(worker.host_write_event) + assert not worker._dma_submitted + + +def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch): + dst_slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + inactive = SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=False, + num_actual_tokens=0, + num_decode_tokens=0, + req_id_per_token=None, + mirror_slot_mapping=dst_slots, + ) + mtp = SimpleNamespace( + runtime=runtime, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=2, + req_id_per_token=torch.tensor([0, 1], dtype=torch.int32), + mirror_slot_mapping=dst_slots, + ) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = [inactive, mtp] + worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None) + + +def test_hisparse_shared_host_reader_skips_mirror(monkeypatch): + """A non-writer TP rank must not mirror rows into the shared host pool.""" + dst_slots = torch.tensor([7, 8], dtype=torch.int64) + leader = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=True, + invalidate_written_slots=MagicMock(), + ) + handle = SimpleNamespace( + runtime=leader, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=1, + req_id_per_token=torch.tensor([0], dtype=torch.int32), + mirror_slot_mapping=dst_slots, + ) + worker = _make_hisparse_worker() + worker.is_host_writer = False + worker.cache_handles = [handle] + worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) + worker._per_layer_mirrored = set() + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_not_called() + leader.invalidate_written_slots.assert_called_once() + + +def test_hisparse_shared_host_reader_skips_transfer_completion(): + """A non-writer TP rank must not acknowledge host writes.""" + worker = _make_hisparse_worker() + worker.is_host_writer = False + worker.kernel_block_size = 1 + worker._enqueued_transfer_ids = [] + worker._pending_transfer_events = [] + + worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)]) + + assert worker._enqueued_transfer_ids == [] + assert worker._pending_transfer_events == [] + + +def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.dma_stream = MagicMock() + worker._enqueued_transfer_ids = [] + worker._pending_transfer_events = deque() + event = MagicMock() + monkeypatch.setattr(torch, "Event", lambda: event) + + worker._record_transfer_completion( + [ + SparseKVPageTransfer(3, 2, (1,), False), + SparseKVPageTransfer(7, 4, (5,), True), + ] + ) + + event.record.assert_called_once_with(worker.dma_stream) + assert worker._enqueued_transfer_ids == [3, 7] + assert worker._pending_transfer_events == deque([(event, (3, 7))]) + + +@pytest.mark.parametrize("is_host_writer", [False, True]) +def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer): + worker = _make_hisparse_worker() + worker.is_host_writer = is_host_writer + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1")) + host_write_events = (MagicMock(), MagicMock()) + worker.host_write_events = host_write_events + worker.host_write_event = host_write_events[1] + worker._next_host_write_event = 0 + worker.host_caches = () + worker.host_num_blocks = 1 + worker._post_forward_transfers = [] + worker._pending_invalid_block_ids = [] + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), + decode_batch=False, + host_mirror_required=False, + num_actual_tokens=0, + num_decode_tokens=0, + req_id_per_token=None, + ) + ] + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_mirror_callbacks = () + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + + assert stream.wait_event.call_args_list == [ + call(host_write_events[1]), + call(host_write_events[0]), + ] + assert worker.host_write_event is host_write_events[1] + + +@pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)]) +def test_hisparse_shared_host_block_copy_has_one_writer( + monkeypatch, tp_rank, expected_copies +): + worker = _make_hisparse_worker() + worker.shared_host_region = object() + worker.host_caches = (torch.empty(4, 1),) + worker.host_num_blocks = 4 + previous_event = MagicMock() + copies = (KVCacheBlockCopy(0, 1),) + copy_blocks = MagicMock() + tp_group = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group) + monkeypatch.setattr( + hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks + ) + + worker._copy_host_blocks(copies, previous_event) + + assert copy_blocks.call_count == expected_copies + tp_group.barrier.assert_called_once_with() + assert worker.take_completed_host_copies() == [1] + assert worker.take_completed_host_copies() == [] + + +def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): + handle = SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), + decode_batch=True, + num_actual_tokens=2, + num_decode_tokens=2, + req_id_per_token=torch.tensor([0, 1]), + mirror_slot_mapping=torch.tensor([4, 5]), + ) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.hot_backing = SimpleNamespace(device=torch.device("cpu")) + worker.host_write_event = MagicMock() + worker._forward_ready_event = MagicMock() + worker.host_caches = () + worker.host_num_blocks = 1 + worker.cache_handles = [handle] + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_mirror_callbacks = (MagicMock(),) + worker._post_forward_transfers = [] + worker._pending_invalid_block_ids = [] + worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror) + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + worker.finish_forward() + + worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event) + assert handle.num_actual_tokens == 0 + torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5])) + + +def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch): + """Followers must not allocate duplicate runtime state before profiling.""" + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_seqs=2, + max_num_batched_tokens=2, + async_scheduling=False, + ), + speculative_config=None, + kv_transfer_config=None, + ) + resolved = hisparse_runtime_module.ResolvedHiSparseConfig( + top_k=4, + device_buffer_size=8, + ) + monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True) + monkeypatch.setattr( + hisparse_runtime_module.ResolvedHiSparseConfig, + "from_vllm_config", + classmethod(lambda cls, vllm_config, model_top_k: resolved), + ) + shared_states: list[object] = [] + streams: list[object] = [] + + def create_shared_state(_device, _max_rows, _top_k): + shared_states.append(object()) + return shared_states[-1] + + def create_stream(_device): + streams.append(object()) + return streams[-1] + + monkeypatch.setattr( + hisparse_runtime_module, "_create_shared_topk_state", create_shared_state + ) + monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream) + monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object()) + monkeypatch.setattr(index_group_module, "_create_event", lambda: object()) + index_group_builder = index_group_module.SparseMLAIndexGroupBuilder( + torch.empty((2, 4), dtype=torch.int32) + ) + + def make_cache_handle(is_leader: bool): + index_group, _ = index_group_builder.register_layer(is_leader) + cache_handle = hisparse_runtime_module.create_hisparse_cache_handle( + config, + model_top_k=4, + is_index_group_leader=is_leader, + row_width=8, + kv_dtype=torch.float32, + index_group=index_group, + device="cpu", + ) + assert cache_handle is not None + return cache_handle + + first_leader = make_cache_handle(True) + first_index_group = index_group_builder.current_group + first_follower = make_cache_handle(False) + second_leader = make_cache_handle(True) + second_index_group = index_group_builder.current_group + second_follower = make_cache_handle(False) + + assert first_index_group is not None + assert second_index_group is not None + assert first_follower.runtime.index_group is first_leader.runtime.index_group + assert second_follower.runtime.index_group is second_leader.runtime.index_group + assert first_leader.runtime.index_group is not second_leader.runtime.index_group + assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream + assert ( + first_leader.runtime.index_group.logical_topk_ready + is first_index_group.logical_topk_ready + ) + assert ( + second_leader.runtime.index_group.copy_stream is second_index_group.side_stream + ) + assert ( + second_leader.runtime.index_group.logical_topk_ready + is second_index_group.logical_topk_ready + ) + assert len(shared_states) == 2 + assert streams == [] class _TestReplaySSMMixer(MambaMixer2): @@ -30,7 +914,8 @@ def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8) -def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): +@pytest.mark.parametrize("layers_only", [False, True]) +def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only): mixers = [_TestReplaySSMMixer() for _ in range(3)] layer_names = [f"layers.{i}.mixer" for i in range(3)] ctx = dict(zip(layer_names, mixers)) @@ -45,7 +930,10 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): SimpleNamespace(layer_names=[layer_names[1]]), ] - bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) + if layers_only: + bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups) + else: + bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) assert ( mixers[0]._replayssm_ring_start.data_ptr() diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index c5fbb76d985d..3311215016a7 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3008,6 +3008,9 @@ def cp_gather_and_upconvert_fp8_kv_cache( workspace_starts: torch.Tensor, batch_size: int, seq_starts: torch.Tensor | None = None, + host_cache: torch.Tensor | None = None, + host_row_ids: torch.Tensor | None = None, + device_row_ids: torch.Tensor | None = None, ) -> None: """Gather and upconvert FP8 KV cache to BF16 workspace. @@ -3018,9 +3021,20 @@ def cp_gather_and_upconvert_fp8_kv_cache( workspace_starts: Workspace start offsets [num_reqs] batch_size: Number of requests seq_starts: Optional source sequence offsets [num_reqs] + host_cache: Optional pinned host rows used for non-resident entries + host_row_ids: Host source row for each remapped cache row + device_row_ids: Device source row, or -1, for each remapped cache row """ torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( - src_cache, dst, block_table, workspace_starts, batch_size, seq_starts + src_cache, + dst, + block_table, + workspace_starts, + batch_size, + seq_starts, + host_cache, + host_row_ids, + device_row_ids, ) diff --git a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py index 5998a0313840..e89864fe4315 100644 --- a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py +++ b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py @@ -35,7 +35,10 @@ def fused_rope_unified_mla_kv_cache_update_impl( layer_name: LayerNameType, ) -> torch.Tensor: layer_name = _resolve_layer_name(layer_name) - attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name) + attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context( + layer_name + ) + attn_layer.prepare_kv_cache_update(attn_metadata) if layer_slot_mapping is not None: ops.concat_and_cache_mla_rope_fused( positions, diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 1cfa372d60f7..ff2ba0493d8d 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.config.attention import AttentionConfig +from vllm.config.attention import AttentionConfig, HiSparseConfig from vllm.config.cache import CacheConfig from vllm.config.compilation import ( CompilationConfig, @@ -68,6 +68,7 @@ __all__ = [ # From vllm.config.attention "AttentionConfig", + "HiSparseConfig", # From vllm.config.cache "CacheConfig", # From vllm.config.compilation diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 799614a69934..05da048bdbd6 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -4,7 +4,7 @@ from dataclasses import field from typing import Any, Literal -from pydantic import field_validator +from pydantic import Field, field_validator from vllm.config.utils import config from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -14,6 +14,25 @@ MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] +@config +class HiSparseConfig: + """Configuration for HiSparse sparse-MLA KV offloading.""" + + device_buffer_size: int | None = Field(default=None, gt=0) + """Total per-request GPU hot-buffer rows, including the newest-token slot. + + Defaults to one top-k per decode query plus one top-k of LRU slack. The + physical allocation is rounded up to the GPU cache block size selected + from the active backends. + """ + + eager_host_mirror: bool = True + """Mirror decode-written KV rows to the host pool during the forward so + page spills complete without moving data. When disabled, decode rows stay + resident-only and evicted pages are copied to host at spill time. Prefill + rows are always mirrored during the forward.""" + + @config class AttentionConfig: """Configuration for attention mechanisms in vLLM.""" @@ -67,6 +86,10 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + hisparse_config: HiSparseConfig | None = None + """HiSparse host-resident KV configuration. Setting this enables experimental + Model Runner V2-only HiSparse sparse-MLA decode hot-buffering.""" + use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index d1344a5211cd..b299aae6f0a7 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -13,6 +13,40 @@ KVRole = Literal[KVProducer, KVConsumer] +def hisparse_host_pool_gib( + kv_transfer_config: "KVTransferConfig | None", +) -> float | None: + if kv_transfer_config is None: + return None + if kv_transfer_config.kv_connector == "MultiConnector": + connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", []) + else: + connectors = [ + { + "kv_connector": kv_transfer_config.kv_connector, + "kv_connector_extra_config": ( + kv_transfer_config.kv_connector_extra_config + ), + } + ] + entries = [ + connector.get("kv_connector_extra_config", {}) + for connector in connectors + if connector.get("kv_connector") == "HiSparseConnector" + ] + if len(entries) > 1: + raise ValueError("Only one HiSparseConnector may be configured") + if not entries: + return None + host_pool_gib = entries[0].get("host_pool_gib") + if host_pool_gib is None: + raise ValueError("HiSparseConnector requires host_pool_gib") + host_pool_gib = float(host_pool_gib) + if host_pool_gib <= 0: + raise ValueError("HiSparseConnector host_pool_gib must be positive") + return host_pool_gib + + def kv_buffer_device_default_factory() -> str: from vllm.platforms import current_platform diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ae56e5880327..cfecff72b618 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1919,6 +1919,13 @@ def uses_dynamic_speculative_decoding(self) -> bool: def uses_draft_model(self) -> bool: return self.method == "draft_model" + def uses_draft_kv_cache(self) -> bool: + return ( + self.use_eagle() + or self.uses_draft_model() + or self.uses_extract_hidden_states() + ) + def uses_extract_hidden_states(self) -> bool: return self.method == "extract_hidden_states" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 05377a273513..3c197df0daff 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -666,6 +666,14 @@ def uniform_decode_query_len(self) -> int: @property def use_v2_model_runner(self) -> bool: + if self.attention_config.hisparse_config is not None: + if envs.VLLM_USE_V2_MODEL_RUNNER is False: + raise ValueError( + "HiSparse requires Model Runner V2; remove " + "VLLM_USE_V2_MODEL_RUNNER=0." + ) + return True + if getattr(self, "watermark_config", None) is not None: if envs.VLLM_USE_V2_MODEL_RUNNER is False: logger.info_once( @@ -1559,6 +1567,39 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if self.attention_config.hisparse_config is not None: + if not current_platform.is_cuda(): + raise ValueError("HiSparse currently requires NVIDIA CUDA.") + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError("HiSparse does not support pipeline parallelism.") + if self.parallel_config.decode_context_parallel_size > 1: + raise ValueError( + "HiSparse does not support decode context parallelism." + ) + if self.model_config is not None and not hasattr( + self.model_config.hf_config, "index_topk" + ): + raise ValueError( + "HiSparse is only supported for DSA models with index_topk." + ) + if self.kv_transfer_config is not None and ( + self.kv_transfer_config.kv_connector + not in ( + None, + "NixlConnector", + "MooncakeStoreConnector", + "MultiConnector", + ) + ): + logger.warning( + "HiSparse host-resident KV is configured with connector " + "%s. NixlConnector (GPU-staged host imports) and " + "MooncakeStoreConnector (shared-store offload) are the " + "validated paths; other connectors are treated as " + "debug/fallback paths.", + self.kv_transfer_config.kv_connector, + ) + self._normalize_piecewise_cudagraph_mode( breakable_cudagraph_enabled=breakable_cudagraph_enabled ) @@ -1937,6 +1978,16 @@ def has_blocked_weights(): # Default to enable HMA if not explicitly disabled by user or logic above. self.scheduler_config.disable_hybrid_kv_cache_manager = False + if ( + self.attention_config.hisparse_config is not None + and self.scheduler_config.disable_hybrid_kv_cache_manager + ): + raise ValueError( + "HiSparse requires the hybrid KV cache manager; remove " + "--disable-hybrid-kv-cache-manager or use connectors that " + "support HMA." + ) + if self.compilation_config.debug_dump_path: self.compilation_config.debug_dump_path = ( self.compilation_config.debug_dump_path.absolute().expanduser() diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index aad7999d08a2..c08064bd8456 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -197,6 +197,12 @@ def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool: "MultiConnector", ) +KVConnectorFactory.register_connector( + "HiSparseConnector", + "vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector", + "HiSparseConnector", +) + KVConnectorFactory.register_connector( "MoRIIOConnector", "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector", diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 39f8d8a1b927..93552d51b521 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -60,6 +60,7 @@ def __init__(self, expected_finished_count: int): # [req_id -> n_remaining_workers] self._recv_remaining_count = dict[str, int]() self._send_remaining_count = dict[str, int]() + self._failed_recving_pending = set[str]() self._expected_finished_count = expected_finished_count @classmethod @@ -156,6 +157,10 @@ def update_finished_set( combined_kv_cache_events.increment_workers(1) invalid_block_ids |= kv_output.invalid_block_ids + self._failed_recving_pending |= kv_output.failed_recving + + failed_recving = self._failed_recving_pending & finished_recving + self._failed_recving_pending -= failed_recving # select output of the worker specified by output_rank output = outputs[output_rank] @@ -168,6 +173,7 @@ def update_finished_set( kv_cache_events=combined_kv_cache_events or None, kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None, invalid_block_ids=invalid_block_ids, + failed_recving=failed_recving, expected_finished_count=self._expected_finished_count, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index fa764571acf2..c97374832ffe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -34,8 +34,8 @@ save_kv_layer() - starts saving KV for layer i (maybe async) wait_for_save() - blocks until all saves are done - get_finished() - called with ids of finished requests, returns - ids of requests that have completed async sending/recving. + get_transfer_results() - returns async send/receive completions and + receive failures in one snapshot. build_connector_worker_meta() - builds metadata to be sent back to the scheduler-side connector """ @@ -43,6 +43,7 @@ import enum from abc import ABC, abstractmethod from collections.abc import Callable, Iterable +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal import torch @@ -63,7 +64,7 @@ ) from vllm.forward_context import ForwardContext from vllm.v1.core.block_pool import BlockPool - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -82,6 +83,19 @@ logger = init_logger(__name__) +@dataclass +class KVConnectorTransferResults: + """Asynchronous transfer completions from one worker snapshot. + + Failed receives also appear in ``finished_recving`` so the scheduler can + release the request from its transfer wait state. + """ + + finished_sending: set[str] = field(default_factory=set) + finished_recving: set[str] = field(default_factory=set) + failed_recving: set[str] = field(default_factory=set) + + class SupportsHMA(ABC): """ The class that indicates the corresponding connector supports hybrid memory @@ -211,6 +225,7 @@ def __init__( else: raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1") self._kv_cache_config = kv_cache_config + self._kv_cache_manager: KVCacheManager | None = None self._role = role @property @@ -278,6 +293,15 @@ def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): """ return + # TODO(NickLucche): group model-runner lifecycle hooks in the interface. + def finish_forward(self) -> None: + """Notify the connector that the model no longer reads this step's KV.""" + return + + def reset_capture_state(self) -> None: + """Reset worker state mutated while capturing CUDA graphs.""" + return + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata): """ Handle preempted requests or evicted blocks BEFORE they are overwritten. @@ -368,6 +392,16 @@ def get_finished( """ return None, None + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + """Return completed sends, receives, and receive failures together.""" + finished_sending, finished_recving = self.get_finished(finished_req_ids) + return KVConnectorTransferResults( + finished_sending=set(finished_sending or ()), + finished_recving=set(finished_recving or ()), + ) + def get_block_ids_with_load_errors(self) -> set[int]: """ Get the set of block IDs that failed to load. @@ -380,9 +414,9 @@ def get_block_ids_with_load_errors(self) -> set[int]: - Applies to both sync- and async-loading requests. - Async loading: failed blocks may be reported in any forward pass up to and including the pass where the request ID is returned by - `get_finished()`. Even if failures occur, the request must still - be reported via `get_finished()`, and the failed block IDs must - appear here no later than that same pass. + `get_transfer_results()`. Even if failures occur, the request + must still be reported as finished receiving, and the failed + block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. """ @@ -436,6 +470,11 @@ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: # Scheduler-side methods # ============================== + def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: + """Bind the scheduler's cache manager after it has been constructed.""" + self._kv_cache_manager = kv_cache_manager + self.bind_gpu_block_pool(kv_cache_manager.block_pool) + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: """ Bind the GPU block pool to the connector for per-GPU block status tracking. @@ -583,6 +622,10 @@ def take_events(self) -> Iterable["KVCacheEvent"]: """ return () + def has_pending_block_frees(self) -> bool: + """Whether pending transfers can release blocks instead of preemption.""" + return False + def has_pending_push_work(self) -> bool: """Return True if the connector has push-mode work that requires the engine main loop to keep stepping (e.g. a P-side request whose diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py new file mode 100644 index 000000000000..338fd5cc6e33 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Connector boundary for HiSparse scheduler and worker state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + KVConnectorRole, + KVConnectorWorkerMetadata, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator +from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror +from vllm.v1.outputs import KVConnectorOutput + +if TYPE_CHECKING: + from vllm.forward_context import ForwardContext + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager + from vllm.v1.hisparse.coordinator import HiSparseCoordinator + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + + +@dataclass +class HiSparseConnectorMetadata(KVConnectorMetadata): + command: SparseKVOffloadCommand | None + host_block_copies: tuple[KVCacheBlockCopy, ...] + source_block_ids: tuple[int, ...] + row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]] + all_context_pages_resident: bool + + +@dataclass +class HiSparseConnectorWorkerMetadata(KVConnectorWorkerMetadata): + enqueued_transfer_counts: dict[int, int] + completed_transfer_counts: dict[int, int] + completed_host_copy_dst_ids: tuple[int, ...] = () + + def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata: + assert isinstance(other, HiSparseConnectorWorkerMetadata) + + def add_counts(first: dict[int, int], second: dict[int, int]) -> dict[int, int]: + combined = first.copy() + for transfer_id, count in second.items(): + combined[transfer_id] = combined.get(transfer_id, 0) + count + return combined + + return HiSparseConnectorWorkerMetadata( + enqueued_transfer_counts=add_counts( + self.enqueued_transfer_counts, other.enqueued_transfer_counts + ), + completed_transfer_counts=add_counts( + self.completed_transfer_counts, other.completed_transfer_counts + ), + # Every worker runs the same copies, so a union is the ack. + completed_host_copy_dst_ids=tuple( + {*self.completed_host_copy_dst_ids, *other.completed_host_copy_dst_ids} + ), + ) + + +class HiSparseConnectorScheduler: + def __init__( + self, + *, + async_speculative: bool, + draft_kv_lookahead: int = 0, + ) -> None: + self.coordinator: HiSparseCoordinator | None = None + self.async_speculative = async_speculative + self.draft_kv_lookahead = draft_kv_lookahead + self.requests: dict[str, Request] = {} + + def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None: + assert self.coordinator is None + self.coordinator = coordinator + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> HiSparseConnectorMetadata: + assert self.coordinator is not None + # HiSparse rebinds worker state every step, so the load must start + # before the forward rather than being deferred to post-forward. + scheduler_output.has_sync_kv_loads = True + scheduler_output.block_table_updates = ( + self.coordinator.take_block_table_updates() or None + ) + command = self.coordinator.build_offload_command() + host_block_copies = self.coordinator.take_host_block_copies() + source_group_id = self.coordinator.host_group_id + assert source_group_id is not None + source_block_ids = [ + block_id + for request in scheduler_output.scheduled_new_reqs + for block_id in request.block_ids[source_group_id] + ] + for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids: + if new_block_ids is not None: + source_block_ids.extend(new_block_ids[source_group_id]) + num_computed_tokens = { + request.req_id: request.num_computed_tokens + for request in scheduler_output.scheduled_new_reqs + } + num_computed_tokens.update( + zip( + scheduler_output.scheduled_cached_reqs.req_ids, + scheduler_output.scheduled_cached_reqs.num_computed_tokens, + ) + ) + scheduled_requests = tuple( + ( + request_id, + num_computed_tokens[request_id], + scheduled_count, + ) + for request_id, scheduled_count in ( + scheduler_output.num_scheduled_tokens.items() + ) + ) + row_mirrors = {} + for request_id, scheduled_start, scheduled_count in scheduled_requests: + mirror_start = scheduled_start + if self.async_speculative: + mirror_start = max( + 0, + scheduled_start - self.requests[request_id].num_output_placeholders, + ) + row_mirrors[request_id] = self.coordinator.build_row_mirrors( + ( + ( + request_id, + mirror_start, + scheduled_count + + scheduled_start + - mirror_start + + self.draft_kv_lookahead, + ), + ), + ) + return HiSparseConnectorMetadata( + command, + host_block_copies, + tuple(source_block_ids), + row_mirrors, + self.coordinator.all_context_pages_resident(scheduled_requests), + ) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + assert self.coordinator is not None + for request_id in connector_output.finished_recving or (): + self.coordinator.finish_host_import( + request_id, failed=request_id in connector_output.failed_recving + ) + metadata = connector_output.kv_connector_worker_meta + if metadata is None: + return + assert isinstance(metadata, HiSparseConnectorWorkerMetadata) + self.coordinator.release_completed_host_copies( + metadata.completed_host_copy_dst_ids + ) + self.coordinator.update_spills( + metadata.enqueued_transfer_counts, + metadata.completed_transfer_counts, + ) + + +class HiSparseConnector(KVConnectorBase_V1, SupportsHMA): + """Join the scheduler coordinator to the worker's transfer engine.""" + + @classmethod + def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str: + return "BLHNC" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: KVCacheConfig, + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + if kv_cache_config.hisparse_host_num_blocks is None: + raise ValueError("HiSparseConnector requires a HiSparse host pool") + self.connector_scheduler: HiSparseConnectorScheduler | None = None + self.connector_worker: HiSparseConnectorWorker | None = None + if role == KVConnectorRole.SCHEDULER: + speculative_config = vllm_config.speculative_config + self.connector_scheduler = HiSparseConnectorScheduler( + async_speculative=bool( + vllm_config.scheduler_config.async_scheduling + and speculative_config is not None + ), + draft_kv_lookahead=vllm_config.num_lookahead_tokens, + ) + elif role == KVConnectorRole.WORKER: + self.connector_worker = HiSparseConnectorWorker( + vllm_config, kv_cache_config + ) + else: + raise ValueError(f"Unsupported KV connector role: {role}") + + def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.bind_coordinator( + get_hisparse_coordinator(kv_cache_manager) + ) + + @property + def requires_kv_delivery(self) -> bool: + return False + + def has_pending_push_work(self) -> bool: + assert self.connector_scheduler is not None + assert self.connector_scheduler.coordinator is not None + return self.connector_scheduler.coordinator.has_pending_work() + + def has_pending_block_frees(self) -> bool: + assert self.connector_scheduler is not None + assert self.connector_scheduler.coordinator is not None + return self.connector_scheduler.coordinator.has_pending_reclamation() + + def finish_forward(self) -> None: + assert self.connector_worker is not None + self.connector_worker.finish_forward() + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def reset_capture_state(self) -> None: + assert self.connector_worker is not None + self.connector_worker.reset_hot_state() + + def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: + assert self.connector_worker is not None + metadata = self._get_connector_metadata() + assert isinstance(metadata, HiSparseConnectorMetadata) + request_state_indices = kwargs.get("request_state_indices") + assert request_state_indices is None or isinstance( + request_state_indices, torch.Tensor + ) + request_ids = kwargs.get("request_ids") + assert request_ids is None or isinstance(request_ids, list) + attn_metadata = kwargs.get("attn_metadata") + assert attn_metadata is None or isinstance(attn_metadata, dict) + self.connector_worker.start_step( + metadata, + request_state_indices, + request_ids, + num_tokens=int(kwargs.get("num_tokens") or 0), + ) + self.connector_worker.prepare_forward(attn_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + return + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs: Any, + ) -> None: + return + + def wait_for_save(self) -> None: + return + + def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: + assert self.connector_worker is not None + enqueued, completed = self.connector_worker.take_transfer_updates() + host_copies = self.connector_worker.take_completed_host_copies() + if not enqueued and not completed and not host_copies: + return None + return HiSparseConnectorWorkerMetadata( + enqueued_transfer_counts={transfer_id: 1 for transfer_id in enqueued}, + completed_transfer_counts={transfer_id: 1 for transfer_id in completed}, + completed_host_copy_dst_ids=tuple(host_copies), + ) + + def shutdown(self) -> None: + if self.connector_worker is not None: + self.connector_worker.shutdown() + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int | None, bool]: + return 0, False + + def update_state_after_alloc( + self, + request: Request, + blocks: KVCacheBlocks, + num_external_tokens: int, + ) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.requests[request.request_id] = request + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + + def request_finished_all_groups( + self, + request: Request, + block_ids: tuple[list[int], ...], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + self.connector_scheduler.requests.pop(request.request_id, None) + return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py new file mode 100644 index 000000000000..1f723e53443c --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -0,0 +1,832 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HiSparse worker-side host/hot state and data movement.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Any, NamedTuple + +import numpy as np +import torch + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tp_group, +) +from vllm.utils.torch_utils import current_stream +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX +from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state +from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + KVCacheConfig, +) +from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace + +if TYPE_CHECKING: + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnectorMetadata, + ) + from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion + + +class _DMADescriptors(NamedTuple): + src: torch.Tensor + dst: torch.Tensor + sizes: torch.Tensor + src_np: np.ndarray + dst_np: np.ndarray + sizes_np: np.ndarray + + +@dataclass +class _SlotMappingStaging: + stream: torch.Stream + event: torch.Event + slots: torch.Tensor + candidates: tuple[SparseKVRowMirror, ...] = () + num_tokens: int = 0 + source_index: int = 0 + + +def _allocate_dma_descriptors(size: int) -> _DMADescriptors: + src, dst, sizes = (torch.empty(size, dtype=torch.int64) for _ in range(3)) + return _DMADescriptors( + src, + dst, + sizes, + src.numpy(), + dst.numpy(), + sizes.numpy(), + ) + + +def _get_hisparse_cache( + forward_context: dict[str, Any], layer_name: str +) -> HiSparseCacheHandle: + attention_layer = forward_context[layer_name] + hisparse_cache = attention_layer.hisparse_cache + assert hisparse_cache is not None + return hisparse_cache + + +def _flatten_row_mirrors( + row_mirrors: Mapping[str, tuple[SparseKVRowMirror, ...]], + request_ids: Sequence[str] | None, +) -> tuple[SparseKVRowMirror, ...]: + ordered_ids = row_mirrors if request_ids is None else request_ids + return tuple( + mirror + for request_id in ordered_ids + for mirror in row_mirrors.get(request_id, ()) + ) + + +def _select_written_row_mirrors( + candidates: tuple[SparseKVRowMirror, ...], + source_slots: np.ndarray, + source_index: int, +) -> tuple[SparseKVRowMirror, ...]: + """Select and coalesce GPU-written rows from a scheduler-owned envelope.""" + source_slots = source_slots[source_slots >= 0] + if source_slots.size == 0 or not candidates: + return () + starts = np.asarray([mirror.source_starts for mirror in candidates], dtype=np.int64) + counts = np.fromiter((mirror.num_rows for mirror in candidates), dtype=np.int64) + destinations = np.fromiter( + [mirror.destination_start for mirror in candidates], dtype=np.int64 + ) + if source_index >= starts.shape[1]: + raise RuntimeError("HiSparse row-mirror source index is out of range.") + source_starts = starts[:, source_index] + order = np.argsort(source_starts) + matches = np.searchsorted(source_starts[order], source_slots, side="right") - 1 + candidate_ids = order[np.maximum(matches, 0)] + offsets = source_slots - source_starts[candidate_ids] + valid = (matches >= 0) & (offsets < counts[candidate_ids]) + if not np.all(valid): + # A page the scheduler could not map has no host destination, so its + # rows are simply not mirrorable this step. + source_slots = source_slots[valid] + candidate_ids = candidate_ids[valid] + offsets = offsets[valid] + if source_slots.size == 0: + return () + sources = starts[candidate_ids] + offsets[:, None] + destinations = destinations[candidate_ids] + offsets + boundaries = ( + np.flatnonzero( + (candidate_ids[1:] != candidate_ids[:-1]) + | (np.diff(destinations) != 1) + | np.any(np.diff(sources, axis=0) != 1, axis=1) + ) + + 1 + ) + run_starts = np.concatenate(([0], boundaries)) + run_ends = np.concatenate((boundaries, [destinations.size])) + return tuple( + SparseKVRowMirror( + source_starts=tuple(int(value) for value in sources[start]), + destination_start=int(destinations[start]), + num_rows=int(end - start), + ) + for start, end in zip(run_starts, run_ends, strict=True) + ) + + +def _create_hisparse_host_events( + shared_host_region: SharedOffloadRegion | None, + is_host_writer: bool, + device: torch.device, +) -> tuple[torch.Event, torch.Event]: + if shared_host_region is None: + return torch.Event(), torch.Event() + + events: tuple[torch.Event, torch.Event] | None = None + ipc_handles = None + if is_host_writer: + events = ( + torch.cuda.Event(interprocess=True), + torch.cuda.Event(interprocess=True), + ) + stream = current_stream() + for event in events: + event.record(stream) + ipc_handles = tuple(event.ipc_handle() for event in events) + ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0) + if events is None: + events = ( + torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]), + torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]), + ) + return events + + +class HiSparseConnectorWorker: + """Own HiSparse host/hot state and execute its worker-side transfers.""" + + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: + self.vllm_config = vllm_config + self.kv_cache_config = kv_cache_config + self._initialized = False + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + forward_context = self.vllm_config.compilation_config.static_forward_context + cache_handles: list[HiSparseCacheHandle] = [] + cache_layer_names: list[str] = [] + for group in self.kv_cache_config.kv_cache_groups: + if not isinstance(group.kv_cache_spec, HiSparseHotSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_HOT_SUFFIX) + layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] + cache_handles.append(_get_hisparse_cache(forward_context, layer_name)) + cache_layer_names.append(layer_name) + + if not cache_handles: + raise RuntimeError("HiSparse connector found no hot-cache handles.") + hot_backings: dict[int, torch.Tensor] = {} + registered_host_pools: dict[int, torch.Tensor] = {} + for cache in cache_handles: + hot_backing = cache.runtime.hot_backing + hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing + registered_pool = cache.runtime.registered_host_pool + registered_host_pools[registered_pool.data_ptr()] = registered_pool + if len(hot_backings) != 1: + raise RuntimeError("HiSparse hot tensors must share one GPU backing.") + hot_backing = next(iter(hot_backings.values())) + pinned_host_pools = list(registered_host_pools.values()) + + resident = cache_handles[0].view + assert resident is not None + host_num_blocks = self.kv_cache_config.hisparse_host_num_blocks + assert host_num_blocks is not None + try: + self.initialize( + cache_handles, + cache_layer_names, + hot_backing, + self.vllm_config.scheduler_config.max_num_seqs, + host_num_blocks, + hot_backing.device, + pinned_host_pools, + ) + except Exception: + release_pinned_state( + [cache.runtime for cache in cache_handles], + pinned_host_pools, + ) + raise + + def initialize( + self, + cache_handles: list[HiSparseCacheHandle], + cache_layer_names: list[str], + hot_backing: torch.Tensor, + max_num_reqs: int, + host_num_blocks: int, + device: torch.device, + pinned_host_pools: list[torch.Tensor], + *, + shared_host_region: SharedOffloadRegion | None = None, + is_host_writer: bool = True, + ) -> None: + if self._initialized: + raise RuntimeError("HiSparse connector worker is already initialized.") + resident = cache_handles[0].view + assert resident is not None + self.kernel_block_size = resident.block_size + self.host_num_blocks = host_num_blocks + self.pinned_host_pools = pinned_host_pools + self.shared_host_region = shared_host_region + self.is_host_writer = is_host_writer + self.dma_stream = ( + torch.cuda.Stream(device=device) if self.is_host_writer else None + ) + self._slot_mapping_staging = None + if self.is_host_writer: + max_mirror_rows = ( + self.vllm_config.scheduler_config.max_num_batched_tokens + + max_num_reqs * (self.vllm_config.num_lookahead_tokens + 1) + ) + self._slot_mapping_staging = _SlotMappingStaging( + stream=torch.Stream(device=device), + event=torch.Event(), + slots=torch.empty( + max_mirror_rows, + dtype=torch.int64, + pin_memory=True, + ), + ) + self._dma_free_descriptors: list[_DMADescriptors] = [] + self._pending_dma_descriptors: deque[tuple[torch.Event, _DMADescriptors]] = ( + deque() + ) + self._dma_submitted = False + self._per_layer_mirrored: set[int] = set() + self._submitted_mirror_layers: set[int] = set() + self._layer_ready_events = tuple(torch.Event() for _ in cache_handles) + self._forward_ready_event = torch.Event() + self._set_row_mirrors(()) + self.cache_layer_names = cache_layer_names + self._group_leaders = tuple( + (layer_name, cache) + for layer_name, cache in zip(cache_layer_names, cache_handles, strict=True) + if cache.runtime.is_group_leader + ) + self.host_write_events = _create_hisparse_host_events( + shared_host_region, is_host_writer, device + ) + self.host_write_event = self.host_write_events[1] + self._next_host_write_event = 0 + self.cache_handles = cache_handles + self.leader_runtimes = [cache.runtime for _, cache in self._group_leaders] + request_state_indices = { + indices.data_ptr(): indices + for cache in cache_handles + if (indices := cache.runtime.request_state_indices) is not None + } + if len(request_state_indices) != 1: + raise RuntimeError( + "HiSparse runtimes must share one request-state mapping." + ) + self.request_state_indices = next(iter(request_state_indices.values())) + if self.request_state_indices.numel() != max_num_reqs: + raise RuntimeError( + "HiSparse request-state mapping does not match max_num_seqs." + ) + self.hot_backing = hot_backing + self._pending_invalid_block_ids: list[int] = [] + # Destination block ids of host copies this worker has run. + self._completed_host_copy_dst_ids: list[int] = [] + self._post_forward_transfers: list[SparseKVPageTransfer] = [] + self._enqueued_transfer_ids: list[int] = [] + self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = ( + deque() + ) + self._init_dma() + if self.is_host_writer: + for layer_index, handle in enumerate(cache_handles): + handle.submit_layer_mirror = partial( + self._enqueue_layer_mirror, layer_index + ) + self._initialized = True + + def set_request_state_indices(self, indices: torch.Tensor) -> None: + if indices.numel() > self.request_state_indices.numel(): + raise ValueError( + "HiSparse request-state mapping exceeds max_num_seqs: " + f"{indices.numel()} > {self.request_state_indices.numel()}." + ) + if torch.cuda.is_current_stream_capturing(): + return + # Attention indexes persistent request state by input-batch row. Refresh + # that indirection after every batch compaction or reorder. + self.request_state_indices.fill_(-1) + self.request_state_indices[: indices.numel()].copy_(indices) + if self._pending_invalid_block_ids: + self.invalidate_blocks(self._pending_invalid_block_ids, indices) + self._pending_invalid_block_ids.clear() + + def _init_dma(self) -> None: + host_caches = tuple(cache.runtime.host_cache for cache in self.cache_handles) + resident_caches = [] + for cache in self.cache_handles: + assert cache.view is not None and cache.slot_mapping is not None + resident_caches.append(cache.view.cache) + self.host_caches = host_caches + self.resident_caches = tuple(resident_caches) + + for resident_cache, host_cache in zip(resident_caches, host_caches): + row_bytes = resident_cache.shape[-1] * resident_cache.element_size() + if ( + resident_cache.ndim != 3 + or resident_cache.shape[1] != self.kernel_block_size + or resident_cache.stride(1) * resident_cache.element_size() != row_bytes + ): + raise RuntimeError("HiSparse DMA requires contiguous resident rows.") + if ( + host_cache.ndim != 2 + or not host_cache.is_contiguous() + or host_cache.shape[1] * host_cache.element_size() != row_bytes + ): + raise RuntimeError("HiSparse DMA requires contiguous host rows.") + + def start_step( + self, + metadata: HiSparseConnectorMetadata, + request_state_indices: torch.Tensor | None, + request_ids: list[str] | None = None, + num_tokens: int = 0, + ) -> None: + self._stage_row_mirror_mapping(num_tokens) + previous_host_write_event = self.host_write_event + self.host_write_event = self.host_write_events[self._next_host_write_event] + self._next_host_write_event ^= 1 + current_stream().wait_event(previous_host_write_event) + self._release_completed_dma_descriptors() + mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids) + if self._slot_mapping_staging is not None: + self._slot_mapping_staging.candidates = mirrors + self._set_row_mirrors(mirrors) + self._dma_submitted = False + self._clear_forward_mirror_state() + for handle in self.cache_handles: + handle.all_context_pages_resident = metadata.all_context_pages_resident + handle.mirror_from_resident = True + self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event) + transfers = ( + metadata.command.page_transfers if metadata.command is not None else [] + ) + self._post_forward_transfers = [ + transfer for transfer in transfers if transfer.after_forward + ] + self._submit_transfers( + [transfer for transfer in transfers if not transfer.after_forward] + ) + self._pending_invalid_block_ids.extend(metadata.source_block_ids) + if request_state_indices is not None: + self.set_request_state_indices(request_state_indices) + + def _clear_forward_mirror_state(self) -> None: + self._per_layer_mirrored.clear() + self._submitted_mirror_layers.clear() + for handle in self.cache_handles: + handle.decode_batch = False + handle.host_mirror_required = False + handle.num_actual_tokens = 0 + handle.num_decode_tokens = 0 + handle.req_id_per_token = None + + def prepare_forward(self, attn_metadata: Mapping[str, Any] | None) -> None: + if attn_metadata is None: + return + for layer_name, handle in self._group_leaders: + metadata = attn_metadata.get(layer_name) + if metadata is not None: + handle.prepare_group_for_batch(metadata) + + def _stage_row_mirror_mapping(self, num_tokens: int) -> None: + """Snapshot the rows this forward will write, off the compute stream. + + The resident slot mapping is a persistent view bound at registration + time, so the rows are already staged by the time the forward launches. + """ + state = self._slot_mapping_staging + if state is None or not num_tokens: + return + handle = self.cache_handles[0] + slots = handle.slot_mapping + assert slots is not None + source_index = handle.runtime.resident_source_index + start = state.num_tokens + end = start + num_tokens + if end > state.slots.shape[0]: + raise ValueError( + "HiSparse row mapping exceeds staging capacity: " + f"{end} > {state.slots.shape[0]}." + ) + main_stream = current_stream() + state.stream.wait_stream(main_stream) + with torch.cuda.stream(state.stream): + state.slots[start:end].copy_( + slots[:num_tokens], + non_blocking=True, + ) + state.event.record(state.stream) + state.num_tokens = end + state.source_index = source_index + + def _resolve_row_mirrors(self, state: _SlotMappingStaging) -> None: + """Compute the final row mirrors from the staged GPU slot mapping.""" + self._set_row_mirrors( + _select_written_row_mirrors( + state.candidates, + state.slots[: state.num_tokens].numpy(), + state.source_index, + ) + ) + state.num_tokens = 0 + + def _copy_host_blocks( + self, + host_block_copies: Sequence[KVCacheBlockCopy], + previous_host_write_event: torch.Event, + ) -> None: + if not host_block_copies: + return + self._completed_host_copy_dst_ids.extend( + copy.dst_block_id for copy in host_block_copies + ) + if self.shared_host_region is None or get_tensor_model_parallel_rank() == 0: + if self.host_caches: + previous_host_write_event.synchronize() + copy_kv_cache_blocks_inplace( + self.host_caches, + self.host_num_blocks, + host_block_copies, + ) + if self.shared_host_region is not None: + get_tp_group().barrier() + + def invalidate_blocks( + self, block_ids: list[int], request_state_indices: torch.Tensor + ) -> None: + """Invalidate recycled host slots in this worker's leader runtimes.""" + if not block_ids: + return + device = self.cache_handles[0].runtime.device + staging = torch.tensor(block_ids, dtype=torch.int32, pin_memory=True) + blocks = staging.to(device, dtype=torch.int32, non_blocking=True) + offsets = torch.arange(self.kernel_block_size, dtype=torch.int32, device=device) + slots = (blocks[:, None] * self.kernel_block_size + offsets[None, :]).flatten() + sorted_slots = torch.sort(slots).values + state_indices = request_state_indices.to(device=device, dtype=torch.long) + for runtime in self.leader_runtimes: + runtime.invalidate_sorted_slots(sorted_slots, state_indices) + + def reset_hot_state(self) -> None: + for runtime in self.leader_runtimes: + runtime.reset_hot_state() + + def _release_completed_dma_descriptors(self) -> None: + pending = self._pending_dma_descriptors + while pending and pending[0][0].query(): + _, descriptors = pending.popleft() + self._dma_free_descriptors.append(descriptors) + + def _acquire_dma_descriptors(self, size: int) -> _DMADescriptors: + for index, descriptors in enumerate(self._dma_free_descriptors): + if descriptors.src.numel() >= size: + return self._dma_free_descriptors.pop(index) + return _allocate_dma_descriptors(size) + + def _submit_dma_descriptors( + self, + descriptors: _DMADescriptors, + descriptor_count: int, + transfer_ids: tuple[int, ...] = (), + ready_event: torch.Event | None = None, + ) -> None: + stream = self.dma_stream + assert stream is not None + if ready_event is None: + stream.wait_stream(current_stream()) + else: + stream.wait_event(ready_event) + completion_event = torch.Event() + with torch.cuda.stream(stream): + ops.swap_blocks_batch( + descriptors.src[:descriptor_count], + descriptors.dst[:descriptor_count], + descriptors.sizes[:descriptor_count], + ) + self.host_write_event.record(stream) + completion_event.record(stream) + self._pending_dma_descriptors.append((completion_event, descriptors)) + self._dma_submitted = True + if transfer_ids: + self._pending_transfer_events.append((completion_event, transfer_ids)) + self._enqueued_transfer_ids.extend(transfer_ids) + + def _set_row_mirrors(self, mirrors: tuple[SparseKVRowMirror, ...]) -> None: + self._row_mirrors = mirrors + self._row_mirror_destination_starts = np.fromiter( + (mirror.destination_start for mirror in mirrors), + dtype=np.int64, + count=len(mirrors), + ) + self._row_mirror_counts = np.fromiter( + (mirror.num_rows for mirror in mirrors), + dtype=np.int64, + count=len(mirrors), + ) + if mirrors: + self._row_mirror_source_starts = np.asarray( + [mirror.source_starts for mirror in mirrors], dtype=np.int64 + ) + if self._row_mirror_source_starts.ndim != 2: + raise RuntimeError("HiSparse DMA source mappings must be rectangular.") + else: + self._row_mirror_source_starts = np.empty((0, 0), dtype=np.int64) + self._row_mirror_num_rows = int(self._row_mirror_counts.sum()) + + def _enqueue_row_dma( + self, layer_indices: Sequence[int], ready_event: torch.Event | None = None + ) -> None: + if ( + not layer_indices + or not self._row_mirror_num_rows + or not self.is_host_writer + ): + return + mirrors = self._row_mirrors + num_layers = len(layer_indices) + descriptor_count = len(mirrors) * num_layers + descriptors = self._acquire_dma_descriptors(descriptor_count) + destination_starts = self._row_mirror_destination_starts + row_counts = self._row_mirror_counts + for descriptor_offset, layer_index in enumerate(layer_indices): + cache = self.cache_handles[layer_index] + source_index = cache.runtime.resident_source_index + if source_index >= self._row_mirror_source_starts.shape[1]: + raise RuntimeError("HiSparse row DMA source index is out of range.") + source_rows = self._row_mirror_source_starts[:, source_index] + source = self.resident_caches[layer_index] + destination = self.host_caches[layer_index] + row_bytes = source.shape[-1] * source.element_size() + if ( + source.stride(1) * source.element_size() != row_bytes + or destination.shape[1] * destination.element_size() != row_bytes + ): + raise RuntimeError("HiSparse row DMA requires contiguous rows.") + source_blocks, source_row_offsets = np.divmod( + source_rows, self.kernel_block_size + ) + source_out_of_range = ( + np.any(source_blocks < 0) + or np.any(source_blocks >= source.shape[0]) + or np.any(source_row_offsets + row_counts > self.kernel_block_size) + ) + if source_out_of_range: + raise RuntimeError("HiSparse row DMA source is out of range.") + if np.any(destination_starts < 0) or np.any( + destination_starts + row_counts > destination.shape[0] + ): + raise RuntimeError("HiSparse row DMA index is out of range.") + descriptor_slice = slice(descriptor_offset, descriptor_count, num_layers) + descriptors.src_np[descriptor_slice] = ( + source.data_ptr() + + source_blocks * source.stride(0) * source.element_size() + + source_row_offsets * row_bytes + ) + descriptors.dst_np[descriptor_slice] = ( + destination.data_ptr() + destination_starts * row_bytes + ) + descriptors.sizes_np[descriptor_slice] = row_counts * row_bytes + + self._submit_dma_descriptors( + descriptors, descriptor_count, ready_event=ready_event + ) + + def _enqueue_layer_mirror(self, layer_index: int) -> None: + handle = self.cache_handles[layer_index] + if not handle.host_mirror_required: + return + if layer_index in self._per_layer_mirrored: + raise RuntimeError(f"HiSparse layer {layer_index} mirrored twice.") + state = self._slot_mapping_staging + if state is not None and state.num_tokens: + # The staging stream waited on the compute stream before the + # forward launched, so this only drains pre-forward work that is + # already queued ahead of the layer kernels: no GPU bubble. + state.event.synchronize() + self._resolve_row_mirrors(state) + self._per_layer_mirrored.add(layer_index) + next_layer = layer_index + 1 + if ( + next_layer < len(self.cache_handles) + and self.cache_handles[next_layer].runtime.resident_source_index + == handle.runtime.resident_source_index + ): + return + ready_event = self._layer_ready_events[layer_index] + ready_event.record() + pending_layers = tuple( + sorted(self._per_layer_mirrored - self._submitted_mirror_layers) + ) + self._enqueue_row_dma( + pending_layers, + ready_event=ready_event, + ) + self._submitted_mirror_layers.update(pending_layers) + + def _record_transfer_completion( + self, transfers: list[SparseKVPageTransfer] + ) -> None: + if not transfers or not self.is_host_writer: + return + stream = self.dma_stream + assert stream is not None + completion_event = torch.Event() + completion_event.record(stream) + transfer_ids = tuple(transfer.transfer_id for transfer in transfers) + self._pending_transfer_events.append((completion_event, transfer_ids)) + self._enqueued_transfer_ids.extend(transfer_ids) + + def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: + if self.cache_handles[0].runtime.eager_host_mirror: + self._record_transfer_completion(transfers) + else: + self._enqueue_transfers(transfers) + + def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: + if not transfers or not self.is_host_writer: + return + num_layers = len(self.cache_handles) + descriptor_count = len(transfers) * num_layers + descriptors = self._acquire_dma_descriptors(descriptor_count) + destination_rows = np.fromiter( + ( + transfer.destination_block_id * self.kernel_block_size + for transfer in transfers + ), + dtype=np.int64, + count=len(transfers), + ) + source_blocks_by_transfer = np.asarray( + [transfer.source_block_ids for transfer in transfers], dtype=np.int64 + ) + if source_blocks_by_transfer.ndim != 2: + raise RuntimeError( + "HiSparse spill DMA source mappings must be rectangular." + ) + for layer_index, cache in enumerate(self.cache_handles): + source_index = cache.runtime.resident_source_index + if source_index >= source_blocks_by_transfer.shape[1]: + raise RuntimeError("HiSparse spill DMA source index is out of range.") + source_blocks = source_blocks_by_transfer[:, source_index] + source = self.resident_caches[layer_index] + destination = self.host_caches[layer_index] + if np.any(source_blocks < 0) or np.any(source_blocks >= source.shape[0]): + raise RuntimeError("HiSparse spill DMA source is out of range.") + if np.any(destination_rows < 0) or np.any( + destination_rows + self.kernel_block_size > destination.shape[0] + ): + raise RuntimeError("HiSparse spill DMA destination is out of range.") + row_bytes = source.shape[-1] * source.element_size() + descriptor_slice = slice(layer_index, descriptor_count, num_layers) + descriptors.src_np[descriptor_slice] = ( + source.data_ptr() + + source_blocks * source.stride(0) * source.element_size() + ) + descriptors.dst_np[descriptor_slice] = ( + destination.data_ptr() + destination_rows * row_bytes + ) + descriptors.sizes_np[descriptor_slice] = self.kernel_block_size * row_bytes + self._submit_dma_descriptors( + descriptors, + descriptor_count, + transfer_ids=tuple(transfer.transfer_id for transfer in transfers), + ) + + def _enqueue_host_mirror( + self, + ready_event: torch.Event | None = None, + ) -> None: + active = [ + (index, handle) + for index, handle in enumerate(self.cache_handles) + if handle.num_actual_tokens != 0 + ] + if not active: + return + + def mirror_key(handle: HiSparseCacheHandle) -> tuple: + slots = handle.mirror_slot_mapping + return ( + handle.num_actual_tokens, + handle.num_decode_tokens, + handle.decode_batch, + handle.host_mirror_required, + handle.runtime.eager_host_mirror, + None if slots is None else slots.data_ptr(), + ) + + keys = {mirror_key(handle) for _, handle in active} + if len(keys) > 1: + raise RuntimeError( + f"HiSparse cache layers disagree on mirror metadata: {list(keys)}." + ) + cache = active[0][1] + if not cache.host_mirror_required: + return + dst_slots = cache.mirror_slot_mapping + if dst_slots is None: + raise RuntimeError("HiSparse host mirror has no source slot mapping.") + num_rows = min(cache.num_actual_tokens, dst_slots.numel()) + if num_rows == 0: + return + if self.is_host_writer: + expected_layers = {index for index, _ in active} + if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers: + raise RuntimeError( + "HiSparse per-layer DMA did not mirror every active layer: " + f"expected {sorted(expected_layers)}, got " + f"{sorted(self._per_layer_mirrored)}." + ) + pending_layers = tuple( + sorted(expected_layers - self._submitted_mirror_layers) + ) + if pending_layers: + self._enqueue_row_dma(pending_layers, ready_event=ready_event) + self._submitted_mirror_layers.update(pending_layers) + assert cache.req_id_per_token is not None + for _, handle in active: + if handle.runtime.is_group_leader: + handle.runtime.invalidate_written_slots( + dst_slots[:num_rows], + cache.req_id_per_token[:num_rows], + ) + + def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None: + state = self._slot_mapping_staging + if state is not None and state.num_tokens: + # Measured free: blocking here costs no throughput (26.7 vs 27.6 + # gen tok/s) and lets every step mirror exactly the written rows. + state.event.synchronize() + self._resolve_row_mirrors(state) + self._enqueue_host_mirror(ready_event) + self._clear_forward_mirror_state() + + def finish_forward(self) -> None: + compute_stream = current_stream() + self._forward_ready_event.record() + self._finish_mirror_phase(self._forward_ready_event) + transfers = self._post_forward_transfers + self._post_forward_transfers = [] + self._submit_transfers(transfers) + if self.is_host_writer: + if self._dma_submitted: + compute_stream.wait_event(self.host_write_event) + self._dma_submitted = False + else: + self.host_write_event.record(compute_stream) + self._release_completed_dma_descriptors() + + def take_completed_host_copies(self) -> list[int]: + """Drain host copies this worker has enqueued for this step.""" + completed = self._completed_host_copy_dst_ids + self._completed_host_copy_dst_ids = [] + return completed + + def take_transfer_updates(self) -> tuple[list[int], list[int]]: + enqueued = self._enqueued_transfer_ids + self._enqueued_transfer_ids = [] + completed: list[int] = [] + while self._pending_transfer_events: + event, transfer_ids = self._pending_transfer_events[0] + if not event.query(): + break + self._pending_transfer_events.popleft() + completed.extend(transfer_ids) + return enqueued, completed + + def shutdown(self) -> None: + if not self._initialized: + return + if self._slot_mapping_staging is not None: + self._slot_mapping_staging.stream.synchronize() + if self.dma_stream is not None: + self.dma_stream.synchronize() + release_pinned_state( + [cache.runtime for cache in self.cache_handles], self.pinned_host_pools + ) + self._initialized = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index f465fcad92a8..e28076c89b09 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -702,9 +702,9 @@ def get_block_ids_with_load_errors(self) -> set[int]: - Applies to both sync- and async-loading requests. - Async loading: failed blocks may be reported in any forward pass up to and including the pass where the request ID is returned by - `get_finished()`. Even if failures occur, the request must still - be reported via `get_finished()`, and the failed block IDs must - appear here no later than that same pass. + `get_transfer_results()`. Even if failures occur, the request + must still be reported as finished receiving, and the failed + block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 8fa9a8836a76..4009af7009ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -635,6 +635,7 @@ def __init__( ): self.vllm_config = vllm_config self.block_size = vllm_config.cache_config.block_size + self.kv_cache_config = kv_cache_config assert vllm_config.kv_transfer_config self.is_kv_producer: bool = ( @@ -649,7 +650,7 @@ def __init__( not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager and any( not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ) ) # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec @@ -670,7 +671,7 @@ def __init__( (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) if isinstance(g.kv_cache_spec, SlidingWindowSpec) else (0, self.block_size) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ] # cdiv(n_tokens, block_size) gives blocks/window; add 1 to # conservatively account for boundary overlap. @@ -684,11 +685,17 @@ def get_sw_clipped_blocks( block_ids: tuple[list[int], ...] | list[list[int]], ) -> list[list[int]]: """Clip per-group block IDs to sliding window size.""" - if len(block_ids) == 0 or not self._is_hma_required: - return list(block_ids) + if len(block_ids) == 0: + return [] + if len(block_ids) == len(self.kv_cache_config.transfer_group_ids): + selected = list(block_ids) + else: + selected = list(self.kv_cache_config.select_transfer_block_ids(block_ids)) + if len(selected) == 0 or not self._is_hma_required: + return selected return [ blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks - for i, blocks in enumerate(block_ids) + for i, blocks in enumerate(selected) ] def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: @@ -1046,18 +1053,13 @@ def __init__( self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} self._layer_specs: dict[str, KVCacheSpec] = {} - for group in kv_cache_config.kv_cache_groups: + for group in kv_cache_config.transfer_groups: group_spec = group.kv_cache_spec specs_by_layer = getattr(group_spec, "kv_cache_specs", {}) for layer_name in group.layer_names: self._layer_specs[layer_name] = specs_by_layer.get( layer_name, group_spec ) - self._layer_group_indices: dict[str, int] = { - layer: group_index - for group_index, group in enumerate(kv_cache_config.kv_cache_groups) - for layer in group.layer_names - } self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.tp_size, @@ -1441,7 +1443,7 @@ def _logical_to_kernel_block_ids( block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape( 1, -1 ) - group_specs = self.kv_cache_config.kv_cache_groups + group_specs = self.kv_cache_config.transfer_groups return [ BlockTable.map_to_kernel_blocks( np.array(group), @@ -1494,7 +1496,7 @@ async def _build_transfer_params( local_block_ids_by_group: list[list[int]] = [] remote_block_ids_by_group: list[list[int]] = [] has_block_error = False - group_specs = self.kv_cache_config.kv_cache_groups + group_specs = self.kv_cache_config.transfer_groups for group_index, (local_group, remote_group) in enumerate( zip(send_meta.local_block_ids, remote_block_ids_per_group) ): @@ -1750,7 +1752,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.registered_layer_names.append(layer_name) self.registered_layer_indices.append(layer_index) self.registered_group_indices.append( - self._layer_group_indices[layer_name] + self.kv_cache_config.transfer_group_index_by_layer[layer_name] ) storage = cache.untyped_storage() storage_addr = storage.data_ptr() @@ -2121,7 +2123,7 @@ def _get_transfer_regions( ) -> list[TransferRegion]: if not group_indices: group_indices = [ - self._layer_group_indices.get(layer_name, 0) + self.kv_cache_config.transfer_group_index_by_layer.get(layer_name, 0) for layer_name in layer_names ] return _expand_transfer_regions( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 67a935d7eedc..375730807d5d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -3,18 +3,20 @@ """External-store cache-hit coordinator for MooncakeStoreConnector.""" from collections.abc import Sequence -from typing import cast +from typing import NamedTuple, cast from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( chunk_hashes_for_block_size, ) from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool -from vllm.v1.core.kv_cache_coordinator import SpecGroup from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, ) +from vllm.v1.core.single_type_kv_cache_manager import ( + SingleTypeKVCacheManager, +) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -25,6 +27,13 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +class StoreSpecGroup(NamedTuple): + spec: KVCacheSpec + group_ids: list[int] + manager_cls: type[SingleTypeKVCacheManager] + use_eagle: bool + + class ExternalCachedBlockPool: """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set.""" @@ -121,7 +130,7 @@ def _verify_and_split_kv_cache_groups(self) -> None: """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but dispatches via spec_manager_map (we don't allocate managers). """ - attention_groups: list[SpecGroup] = [] + attention_groups: list[StoreSpecGroup] = [] for i, g in enumerate(self.kv_cache_groups): # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash # kpool tail): per-request scratch, never shareable, so they must @@ -143,7 +152,7 @@ def _verify_and_split_kv_cache_groups(self) -> None: break else: attention_groups.append( - SpecGroup(spec, [i], manager_cls, g.is_eagle_group) + StoreSpecGroup(spec, [i], manager_cls, g.is_eagle_group) ) # Full attention first (matches upstream convergence ordering). attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 870500e84068..ebcf5708c89d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -67,6 +67,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket +from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -1909,21 +1910,31 @@ def register_kv_caches( logger.warning("No KV caches to offload.") return + # Resolve each entry to a representative tensor for storage + # deduplication. For attention layers the value is already a tensor; + # for Mamba layers it is a list of tensors that all share the same + # underlying raw storage, so we take the first one. + def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: + assert isinstance(v, torch.Tensor | list) + return v if isinstance(v, torch.Tensor) else v[0] + assert self.cache_config.num_gpu_blocks is not None self.num_blocks = self.cache_config.num_gpu_blocks + use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None - seen_storage_ptrs: set[int] = set() - cache_tensors: list[torch.Tensor] = [] - - for cache in kv_caches.values(): - cache = group_kernel_blocks(cache, self.num_blocks) - cache_tensors.append(cache) - cache_storage = cache.untyped_storage() - base_addr = cache_storage.data_ptr() - region_len = cache_storage.nbytes() - - if base_addr not in seen_storage_ptrs: + if not use_group_regions: + seen_storage_ptrs: set[int] = set() + cache_tensors = [ + group_kernel_blocks(_repr_tensor(cache), self.num_blocks) + for cache in kv_caches.values() + ] + for cache in cache_tensors: + cache_storage = cache.untyped_storage() + base_addr = cache_storage.data_ptr() + if base_addr in seen_storage_ptrs: + continue seen_storage_ptrs.add(base_addr) + region_len = cache_storage.nbytes() ret = self.store.register_buffer(base_addr, region_len) if ret != 0: logger.error( @@ -1932,17 +1943,110 @@ def register_kv_caches( region_len, ret, ) + for db in self.token_dbs: + db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) + + registered_buffers: dict[int, int] = {} + num_host_segments = 0 + num_device_segments = 0 + + for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True): + if not use_group_regions: + continue + missing_layers = [ + layer_name + for layer_name in group.layer_names + if layer_name not in kv_caches + ] + if missing_layers: + raise ValueError( + f"Missing transferable KV caches for layers: {missing_layers}" + ) + group_caches = [ + (layer_name, kv_caches[layer_name]) for layer_name in group.layer_names + ] + + seen_group_ptrs: set[int] = set() + addrs: list[int] = [] + block_lens: list[int] = [] + for layer_name, value in group_caches: + cache = _repr_tensor(value) + is_host_resident = group.host_resident + if is_host_resident: + num_blocks = self._kv_cache_config.hisparse_host_num_blocks + assert num_blocks is not None + else: + num_blocks = self._kv_cache_config.num_blocks + cache = group_kernel_blocks(cache, num_blocks) + cache_storage = cache.untyped_storage() + if is_host_resident: + base_addr = cache.data_ptr() + region_len = cache.nbytes + else: + base_addr = cache_storage.data_ptr() + region_len = cache_storage.nbytes() + if region_len % num_blocks != 0: + raise ValueError( + f"KV cache region for {layer_name} has size {region_len}, " + f"which is not divisible by its {num_blocks} blocks." + ) + + registered_len = registered_buffers.get(base_addr) + if registered_len is None: + ret = self.store.register_buffer(base_addr, region_len) + if ret != 0: + mem_kind = "host" if is_host_resident else "device" + raise RuntimeError( + "Mooncake register_buffer failed for addr " + f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}" + ) + registered_buffers[base_addr] = region_len + elif registered_len != region_len: + raise ValueError( + f"KV cache views at {base_addr:#x} expose inconsistent " + f"region sizes: {registered_len} and {region_len}." + ) + + regions: list[tuple[int, int]] = [] + if not is_non_overlapping_and_dense(cache[0]): + for head_idx in range(cache.shape[1]): + head_cache = cache[:, head_idx] + assert is_non_overlapping_and_dense(head_cache[0]) + regions.append( + ( + head_cache.data_ptr(), + head_cache.stride(0) * head_cache.element_size(), + ) + ) + elif cache.stride(0) * cache.element_size() * num_blocks == region_len: + regions.append((base_addr, region_len // num_blocks)) + else: + regions.append( + (cache.data_ptr(), cache.stride(0) * cache.element_size()) + ) + + for region_addr, block_len in regions: + if region_addr in seen_group_ptrs: + continue + seen_group_ptrs.add(region_addr) + addrs.append(region_addr) + block_lens.append(block_len) + if is_host_resident: + num_host_segments += 1 + else: + num_device_segments += 1 + + db.set_kv_caches_base_addr(addrs) + db.set_block_len(block_lens) logger.info( - "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d", + "Registered KV caches: num_groups=%d, num_segments=%d (host=%d, device=%d)", len(self.token_dbs), - len(cache_tensors), - self.num_blocks, + num_host_segments + num_device_segments, + num_host_segments, + num_device_segments, ) - for db in self.token_dbs: - db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) - # Start transfer threads if self.can_put: ready_event_sending = threading.Event() @@ -2316,6 +2420,19 @@ def close(self) -> None: if store is None: return self.store = None + deadline = time.monotonic() + 5.0 + for thread in (self.kv_send_thread, *self.kv_recv_threads): + if thread is None: + continue + while thread.request_queue.unfinished_tasks and time.monotonic() < deadline: + time.sleep(0.05) + if thread.request_queue.unfinished_tasks: + logger.warning( + "Mooncake store %s still has %d in-flight requests at " + "close; closing anyway.", + thread.name, + thread.request_queue.unfinished_tasks, + ) try: store.close() except Exception as e: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 591ded1c6761..d291d3852804 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -17,6 +17,7 @@ KVConnectorHandshakeMetadata, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, KVConnectorWorkerMetadata, SupportsHMA, ) @@ -35,7 +36,7 @@ from vllm.distributed.kv_events import KVCacheEvent from vllm.forward_context import ForwardContext from vllm.v1.core.block_pool import BlockPool - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -208,6 +209,10 @@ def __init__( # Propagated from scheduler to worker side via the connector metadata. self._extra_async_saves: dict[str, int] = {} + @property + def sub_connectors(self) -> list[KVConnectorBase_V1]: + return list(self._connectors) + @property def supports_divergent_local_hybrid_hits(self) -> bool: return bool(self._connectors) and all( @@ -248,6 +253,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): for c in self._connectors: c.register_kv_caches(kv_caches) + def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: + for connector in self._connectors: + connector.bind_kv_cache_manager(kv_cache_manager) + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: for c in self._connectors: c.bind_gpu_block_pool(gpu_block_pool) @@ -292,6 +301,14 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: for c in self._connectors: c.start_load_kv(forward_context, **kwargs) + def finish_forward(self) -> None: + for c in self._connectors: + c.finish_forward() + + def reset_capture_state(self) -> None: + for c in self._connectors: + c.reset_capture_state() + def wait_for_layer_load(self, layer_name: str) -> None: for c in self._connectors: c.wait_for_layer_load(layer_name) @@ -313,29 +330,31 @@ def wait_for_save(self): def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - finished_sending: set[str] = set() - finished_recving: set[str] = set() - for c in self._connectors: - sending, recving = c.get_finished(finished_req_ids) - if not recving and not sending: - continue - # Aggregate finished recving request ids. - finished_recving.update(recving or ()) - # Aggregate finished sending request ids - only include - # once we've drained the "extra" count (for cases where - # more than one connector is async-saving the same request). - for req_id in sending or (): + results = self.get_transfer_results(finished_req_ids) + return ( + results.finished_sending or None, + results.finished_recving or None, + ) + + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + results = KVConnectorTransferResults() + for connector in self._connectors: + child_results = connector.get_transfer_results(finished_req_ids) + results.finished_recving.update(child_results.finished_recving) + results.failed_recving.update(child_results.failed_recving) + for req_id in child_results.finished_sending: extra_pending = self._extra_async_saves.get(req_id) if extra_pending is None: - finished_sending.add(req_id) - continue - assert extra_pending > 0 - if extra_pending == 1: - del self._extra_async_saves[req_id] + results.finished_sending.add(req_id) else: - self._extra_async_saves[req_id] = extra_pending - 1 - - return finished_sending or None, finished_recving or None + assert extra_pending > 0 + if extra_pending == 1: + del self._extra_async_saves[req_id] + else: + self._extra_async_saves[req_id] = extra_pending - 1 + return results def get_block_ids_with_load_errors(self) -> set[int]: agg_block_ids: set[int] = set() @@ -563,6 +582,9 @@ def take_events(self) -> Iterable["KVCacheEvent"]: for c in self._connectors: yield from c.take_events() + def has_pending_block_frees(self) -> bool: + return any(c.has_pending_block_frees() for c in self._connectors) + def has_pending_push_work(self) -> bool: return any(c.has_pending_push_work() for c in self._connectors) @@ -630,7 +652,7 @@ def build_kv_connector_stats( return MultiKVConnectorStats(data=reconstructed_data) - def get_kv_connector_stats(self) -> MultiKVConnectorStats | None: + def get_kv_connector_stats(self) -> KVConnectorStats | None: # Group connector stats by connector type. stats_by_connector: MultiKVConnectorStats | None = None for c in self._connectors: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index ddfcf35e033b..00579c40f032 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -93,10 +93,7 @@ def __init__( for g in kv_cache_config.transfer_groups ) ) - self._has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) - for g in kv_cache_config.transfer_groups - ) + self._has_mamba = kv_cache_config.has_mamba_layers logger.info("Initializing NIXL Scheduler %s", engine_id) if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index e4b46cb64a96..17660da451a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Base worker-side logic for the NIXL connector.""" +import contextlib import itertools import logging +import math import os import queue import threading @@ -31,7 +33,10 @@ kv_postprocess_blksize_on_receive, kv_postprocess_layout_on_receive, ) -from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + CopyBlocksOp, + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( GET_META_MSG, @@ -186,23 +191,28 @@ def _compute_desc_ids( num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) + num_regions = self.num_regions or len(self.block_len_per_layer) if region_num_blocks is None: - region_num_blocks = [num_blocks] * self.num_regions + region_num_blocks = [num_blocks] * num_regions elif block_size_ratio is not None: region_num_blocks = [ int(count * block_size_ratio) for count in region_num_blocks ] - assert len(region_num_blocks) == self.num_regions + if region_num_blocks: + num_regions = len(region_num_blocks) + assert len(region_num_blocks) == num_regions region_offsets = np.cumsum([0, *region_num_blocks[:-1]]) num_fa_descs = sum(region_num_blocks) - # All-attention fast path: single vectorized broadcast. + # All-attention fast path. if num_ssm_regions == 0: if region_group_ids is None: region_group_ids = self.region_group_ids if uses_region_group_mapping is None: uses_region_group_mapping = self._uses_region_group_mapping - assert len(region_group_ids) == self.num_regions + if not region_group_ids and num_regions == 1: + region_group_ids = [0] + assert len(region_group_ids) == num_regions if uses_region_group_mapping is None: uses_region_group_mapping = len(set(region_group_ids)) > 1 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32) @@ -397,8 +407,6 @@ def __init__( # Config. self.vllm_config = vllm_config - # mypy will complain on re-assignment otherwise. - self.block_size: int = cast(int, vllm_config.cache_config.block_size) if vllm_config.kv_transfer_config is None: raise ValueError("kv_transfer_config must be set for NixlConnector") @@ -420,6 +428,17 @@ def __init__( ) self.kv_cache_config = kv_cache_config + transfer_block_sizes = [ + group.kv_cache_spec.block_size + for group in kv_cache_config.transfer_groups + if get_representative_spec_type(group.kv_cache_spec) + not in (MambaSpec, CircularBufferSpec) + ] + self.block_size = ( + math.lcm(*transfer_block_sizes) + if transfer_block_sizes + else cast(int, vllm_config.cache_config.block_size) + ) # Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers. self._layer_specs: dict[str, KVCacheSpec] = {} for group in kv_cache_config.transfer_groups: @@ -599,12 +618,15 @@ def __init__( # Number of NIXL regions. Currently one region per cache # (so 1 per layer for MLA, otherwise 2 per layer) self.num_regions = 0 + self.region_mem_types: list[str] = [] self.region_group_ids: list[int] = [] self._uses_region_group_mapping = False self.region_names: list[str] = [] self.region_num_blocks: list[int] = [] - self.region_mem_types: list[str] = [] self._mixed_mem_types = False + self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} + self._desc_pos_by_block_size: dict[int, np.ndarray] = {} + self._dram_src_handles_by_block_size: dict[int, int] = {} # PP>1 (push mode): this worker holds a contiguous layer slice and # transfers into the matching sub-range of a PP=1 remote's regions. @@ -644,9 +666,6 @@ def __init__( self.dst_region_group_ids: dict[EngineId, list[int]] = {} self.dst_uses_region_group_mapping: dict[EngineId, bool] = {} self.dst_region_mem_types: dict[EngineId, list[str]] = {} - self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} - self._desc_pos_by_block_size: dict[int, np.ndarray] = {} - self._dram_src_handles_by_block_size: dict[int, int] = {} self._registered_descs: list[Any] = [] # In progress transfers. @@ -667,6 +686,7 @@ def __init__( # Uses Queue for thread-safe cross-thread coordination with the # background handshake thread, matching the _ready_requests pattern. self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() + self._failed_inflight_recvs: set[ReqId] = set() self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {} # Handshake metadata of this worker for NIXL transfers. @@ -690,7 +710,6 @@ def __init__( "engine_ttl", 3600.0 ) - self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config self.use_mla = self.model_config.use_mla @@ -1240,7 +1259,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.use_host_buffer, ) - registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {} + registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {} region_mem_types: list[str] = [] seen_base_addresses: list[int] = [] self._ssm_region_indices = [] @@ -1280,10 +1299,16 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). + # P and D may allocate equivalent transferable layers in different + # cache-group orders. Keep their region lists aligned without putting + # layers.10 before layers.2, which would break PP region slicing. layer_names = ( xfer_buffers if self._is_csa_linear - else sorted(xfer_buffers, key=_region_sort_key) + else sorted( + xfer_buffers, + key=_region_sort_key, + ) ) for layer_name in layer_names: cache = xfer_buffers[layer_name] @@ -1299,6 +1324,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): layer_name, ) continue + group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name] # `layer_spec.page_size_bytes` only accounts for logical page_size, that is # the page_size assuming constant `self._logical_num_blocks`. physical_page_size = ( @@ -1307,17 +1333,48 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): else layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) - group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name] + group = self.kv_cache_config.transfer_groups[group_index] + if group.host_resident: + logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks + assert logical_num_blocks is not None + else: + logical_num_blocks = self.kv_cache_config.num_blocks + group_id = group_index num_blocks = ( - self._logical_num_blocks + logical_num_blocks if isinstance(layer_spec, MambaSpec) - else self.num_blocks + else logical_num_blocks * self._physical_blocks_per_logical_kv_block + ) + base_addr = cache.data_ptr() + is_mla_region = isinstance( + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) ) logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) storage = cache.untyped_storage() storage_addr = storage.data_ptr() + is_host_resident = group.host_resident + if cache.device.type == "cpu": + mem_type = "DRAM" + region_device_id = 0 + else: + mem_type = self.nixl_memory_type + region_device_id = max(cache.get_device(), 0) + self.device_id = region_device_id + if is_host_resident: + registration_base = cache.data_ptr() + registration_len = cache.nbytes + else: + registration_base = storage_addr + registration_len = storage.nbytes() + storage_key = (registration_base, mem_type) + if storage_key not in registration_ranges: + registration_ranges[storage_key] = ( + registration_base, + registration_base + registration_len, + region_device_id, + ) if isinstance(layer_spec, KpoolTailSpec): compressed_owner = compressed_region_owners.get(cache.data_ptr()) @@ -1345,22 +1402,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) continue - if cache.device.type == "cpu": - mem_type = "DRAM" - region_device_id = 0 - else: - mem_type = self.nixl_memory_type - region_device_id = max(cache.get_device(), 0) - self.device_id = region_device_id - registration_ranges.setdefault( - (storage_addr, mem_type), - (storage_addr, storage.nbytes(), region_device_id, ""), - ) - - is_mla_region = isinstance( - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) - ) - if isinstance(layer_spec, MambaSpec): physical_ratio = self._physical_blocks_per_logical_kv_block block_len = physical_page_size // physical_ratio @@ -1381,14 +1422,13 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): block_stride = physical_page_size else: block_stride = cache.stride(0) * cache.element_size() - storage_is_block_major = num_blocks * block_stride == storage.nbytes() - # A layer whose [H, N, C] interior is dense addresses its own page - # as one chunk. Otherwise the block's whole row is the only - # contiguous transfer unit. - hnc_contiguous = ( - cache.ndim == 4 - and cache.stride(2) == cache.shape[3] - and cache.stride(1) == cache.shape[2] * cache.shape[3] + storage_is_block_major = num_blocks * block_stride == registration_len + # A view with a contiguous page can own a strided region even + # when sibling layers share the same backing allocation. + page_contiguous = ( + cache.ndim > 1 + and cache[0].is_contiguous() + and cache[0].nbytes == physical_page_size ) virtual_transfer_pages = _uses_dense_virtual_transfer_pages( layer_spec, cache, physical_page_size, num_blocks @@ -1399,13 +1439,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): (cache.data_ptr(), physical_page_size, physical_page_size) ] elif storage_is_block_major and ( - (packed_storage and is_mla_region) - or (not hnc_contiguous and not self._is_csa_linear) + not page_contiguous + and ((packed_storage and is_mla_region) or not self._is_csa_linear) ): - # Packed MLA layouts transfer the complete storage row. - storage_block_len = storage.nbytes() // num_blocks + # TODO(Lucas): handle TP slicing for packed_storage; for now + # restrict to MLA (DSv4) where kv is replicated. + storage_block_len = registration_len // num_blocks region_specs = [ - (storage_addr, storage_block_len, storage_block_len) + (registration_base, storage_block_len, storage_block_len) ] elif storage_is_block_major: region_specs = [ @@ -1440,6 +1481,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_index = seen_base_addresses.index(base_addr) assert region_mem_types[region_index] == mem_type self._region_is_mla[region_index] |= is_mla_region + if is_mla_region: + self.block_len_per_layer[region_index] = block_len + self.block_stride_per_layer[region_index] = block_stride + self.region_num_blocks[region_index] = num_blocks if self.region_group_ids[region_index] != group_id: self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID else: @@ -1450,8 +1495,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.region_group_ids.append(group_id) self.region_names.append(layer_name) self.region_num_blocks.append(num_blocks) - region_mem_types.append(mem_type) self._region_is_mla.append(is_mla_region) + region_mem_types.append(mem_type) if isinstance(layer_spec, MambaSpec): if layer_spec.tp_replicated: @@ -1496,7 +1541,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): == len(self.region_group_ids) == len(self.region_names) == len(self.region_num_blocks) - == len(region_mem_types) ) # Descriptor ids must be region-ordered, matching the remote side. self._scratch_region_indices.sort() @@ -1521,16 +1565,23 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # Total local FA descriptors (boundary between FA and mamba descs). self.num_descs = sum(self.region_num_blocks) + self._mixed_mem_types = len(set(region_mem_types)) > 1 + if self._mixed_mem_types: + assert self.use_mla and not self._has_mamba, ( + "Mixed-device KV registration is only supported for MLA " + "models without Mamba layers." + ) for mem_type in sorted(set(region_mem_types)): - ranges = [ - registration - for ( - _, - registration_mem_type, - ), registration in registration_ranges.items() - if registration_mem_type == mem_type + ranges_for_mem_type = [ + (start, end - start, device_id, "") + for (_, cache_mem_type), ( + start, + end, + device_id, + ) in registration_ranges.items() + if cache_mem_type == mem_type ] - descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type) + descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type) self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) self._registered_descs.append(descs) @@ -1854,7 +1905,11 @@ def register_local_xfer_handler( if self._mixed_mem_types: desc_is_dram = np.concatenate( [ - np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool) + np.full( + count * block_size_ratio, + mem_type == "DRAM", + dtype=bool, + ) for mem_type, count in zip( self.region_mem_types, self.region_num_blocks, @@ -1864,24 +1919,29 @@ def register_local_xfer_handler( ) assert len(desc_is_dram) == len(blocks_data) desc_pos = np.empty(len(desc_is_dram), dtype=np.int64) - dram_indices = np.flatnonzero(desc_is_dram) - device_indices = np.flatnonzero(~desc_is_dram) - desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64) - desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64) + dram_idx = np.where(desc_is_dram)[0] + vram_idx = np.where(~desc_is_dram)[0] + desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64) + desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64) self._desc_is_dram_by_block_size[block_size] = desc_is_dram self._desc_pos_by_block_size[block_size] = desc_pos - blocks_data[desc_is_dram, 2] = 0 - dram_descs = self.nixl_wrapper.get_xfer_descs( - blocks_data[desc_is_dram], "DRAM" - ) + + # DRAM descriptors are registered under CPU device 0. + blocks_data = [ + (addr, length, 0) if is_dram else (addr, length, dev) + for (addr, length, dev), is_dram in zip( + blocks_data, desc_is_dram, strict=True + ) + ] + dram_blocks = [blocks_data[i] for i in dram_idx] + vram_blocks = [blocks_data[i] for i in vram_idx] + dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") self._dram_src_handles_by_block_size[block_size] = ( self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) ) - device_descs = self.nixl_wrapper.get_xfer_descs( - blocks_data[~desc_is_dram], self.nixl_memory_type - ) + descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type) return ( - self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs), + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data, ) @@ -1968,10 +2028,6 @@ def add_remote_agent( ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] - if nixl_agent_meta.region_num_blocks is not None: - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ - start:end - ] if nixl_agent_meta.region_group_ids is not None: nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[ start:end @@ -1982,6 +2038,10 @@ def add_remote_agent( nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ + start:end + ] ### Register remote engine in TransferTopology (idempotent). assert self.transfer_topo is not None @@ -2026,11 +2086,12 @@ def add_remote_agent( nixl_agent_meta.region_num_blocks or [nixl_agent_meta.num_blocks] * num_remote_regions ) - self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or ( + remote_region_group_ids = nixl_agent_meta.region_group_ids or ( self.region_group_ids if len(self.region_group_ids) == num_remote_regions else [0] * num_remote_regions ) + self.dst_region_group_ids[engine_id] = remote_region_group_ids self.dst_uses_region_group_mapping[engine_id] = ( len(set(self.dst_region_group_ids[engine_id])) > 1 ) @@ -2138,8 +2199,6 @@ def add_remote_agent( mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info) blocks_data = np.concatenate([blocks_data, mamba]) - # A mixed local destination is split into multiple READs below. The - # producer remains a single descriptor list in this prerequisite. remote_mem_types = set(self.dst_region_mem_types[engine_id]) if len(remote_mem_types) != 1: raise NotImplementedError( @@ -2540,9 +2599,10 @@ def post_process_device_kv_on_receive_heterogeneous_attn( indices=indices, ) - def get_finished(self) -> tuple[set[str], set[str]]: + def get_transfer_results(self) -> KVConnectorTransferResults: """ - Get requests that are done sending or recving on this specific worker. + Get transfers that completed on this specific worker. + The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ @@ -2564,7 +2624,6 @@ def get_finished(self) -> tuple[set[str], set[str]]: # Add failed requests to done_recving for scheduler tracking # (blocks are already marked invalid, scheduler will handle recompute) done_recving.update(failed_recv_reqs) - if len(done_sending) > 0 or len(done_recving) > 0: logger.debug( "Rank %s, get_finished: %s requests done sending " @@ -2577,6 +2636,7 @@ def get_finished(self) -> tuple[set[str], set[str]]: block_ids_for_blocksize_post_process = defaultdict(list) block_ids_for_heterogeneous_attn_post_process = list[list[int]]() + direct_device_recving = set[str]() for req_id in done_recving: # clean up metadata for completed requests meta = self._recving_metadata.pop(req_id, None) @@ -2594,6 +2654,8 @@ def get_finished(self) -> tuple[set[str], set[str]]: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) + direct_device_recving.add(req_id) + # Post processing for heteroblocksize/layout, and for blocks the # transfer clipped. The latter happens either at remote-block # granularity (block_size_ratio > 1) or at kernel-block @@ -2642,7 +2704,7 @@ def get_finished(self) -> tuple[set[str], set[str]]: for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) - self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs) + self._sync_device_after_direct_recv(direct_device_recving) # Handle timeout to avoid stranding blocks on remote. now = time.perf_counter() @@ -2664,20 +2726,21 @@ def get_finished(self) -> tuple[set[str], set[str]]: del self._reqs_to_send[req_id] done_sending.add(req_id) - return done_sending, done_recving + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recv_reqs, + ) - def _sync_device_after_mamba_recv( - self, - done_recving: set[str], - failed_recv_reqs: set[str], - ) -> None: - """Synchronize ROCm direct-GPU Mamba receives before model execution.""" - if ( - not current_platform.is_rocm() - or not self._has_mamba - or self.use_host_buffer - or not (done_recving - failed_recv_reqs) - ): + def get_finished(self) -> tuple[set[str], set[str]]: + """Compatibility wrapper for the legacy completion API.""" + results = self.get_transfer_results() + return results.finished_sending, results.finished_recving + + def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None: + """Make direct NIXL writes visible before model execution.""" + requires_sync = current_platform.is_rocm() and self._has_mamba + if self.use_host_buffer or not done_recving or not requires_sync: return torch.accelerator.synchronize() @@ -2731,11 +2794,11 @@ def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: self.nixl_wrapper.release_xfer_handle(handle) elif xfer_state == "PROC": in_progress.append(handle) - continue else: self._log_failure( failure_type="transfer_failed", - msg="Marking blocks as invalid", + msg="Deferring request completion until its last " + "xfer is terminal", req_id=req_id, xfer_state=xfer_state, ) @@ -2743,21 +2806,20 @@ def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: except Exception as e: self._log_failure( failure_type="transfer_exception", - msg="Marking blocks as invalid", + msg="Handle is unpollable; treating it as terminal", req_id=req_id, error=e, ) self._handle_failed_transfer(req_id, handle) if not in_progress: - # Only report request as completed when all transfers are done. - # A request failed in an earlier poll was already reported via - # _failed_recv_reqs and its metadata popped by get_finished(); - # don't report it again, just drop the remaining handles. - if req_id in self._recving_metadata: - done_req_ids.add(req_id) - self._send_pending_recv_notifs(req_id) del transfers[req_id] + done_req_ids.add(req_id) + if req_id in self._failed_inflight_recvs: + self._failed_inflight_recvs.remove(req_id) + self._report_failed_recv(req_id) + else: + self._send_pending_recv_notifs(req_id) else: transfers[req_id] = in_progress return done_req_ids @@ -2771,22 +2833,25 @@ def _handle_failed_transfer(self, req_id: str, handle: int | None): req_id: The request ID. handle: The transfer handle. """ - # (multi-read) One handle is created per remote rank, and they do not - # all fail in the same _pop_done_transfers poll. The request is - # reported failed on the first one, which pops its metadata in - # get_finished(); on the later failures only the handle cleanup is left. - # TODO (NickLucche) handle failed transfer for HMA. - # A split READ's notification is sent only after every handle succeeds. + # A sibling READ may still be writing these blocks. Retain metadata + # and defer invalidation until every handle is terminal. self._pending_recv_notifs.pop(req_id, None) + if req_id in self._recving_transfers: + self._failed_inflight_recvs.add(req_id) + else: + self._report_failed_recv(req_id) + if handle is not None: + self.nixl_wrapper.release_xfer_handle(handle) + self.xfer_stats.record_failed_transfer() + + def _report_failed_recv(self, req_id: str) -> None: if (meta := self._recving_metadata.get(req_id)) is not None: if not self._is_hma_required: self._invalid_block_ids.put(set(meta.local_block_ids[0])) self._failed_recv_reqs.put(req_id) - if handle is not None: - self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() def _send_pending_recv_notifs(self, req_id: str) -> None: + """Send notifications deferred by split DRAM/VRAM reads.""" for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []): try: self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id) @@ -2909,10 +2974,10 @@ def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> Block return block_ids block_arange = np.arange(0, ratio).reshape(1, -1) # Mamba blocks have no logical<>physical discrepancy (block-size=1) - group_specs = self.kv_cache_config.transfer_groups physical_block_ids = [] for i, group in enumerate(block_ids): - if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)): + spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec + if _is_ssm_spec(get_representative_spec_type(spec)): physical_block_ids.append(group) else: physical_block_ids.append( @@ -3215,9 +3280,44 @@ def _cleanup_remote_engine( ) def __del__(self): - self.shutdown() + with contextlib.suppress(Exception): + self.shutdown() - def shutdown(self): + def _finish_shutdown(self) -> None: + self._recving_transfers.clear() + try: + for handle in self.src_xfer_handles_by_block_size.values(): + self.nixl_wrapper.release_dlist_handle(handle) + for handles in self.src_xfer_handles_by_tp_ratio.values(): + for handle in handles: + self.nixl_wrapper.release_dlist_handle(handle) + for handles in self._dram_src_handles_by_tp_ratio.values(): + for handle in handles: + self.nixl_wrapper.release_dlist_handle(handle) + for handle in self._dram_src_handles_by_block_size.values(): + self.nixl_wrapper.release_dlist_handle(handle) + except Exception: + logger.exception("NIXL dlist-handle release failed at shutdown.") + self.src_xfer_handles_by_block_size.clear() + self.src_xfer_handles_by_tp_ratio.clear() + self._dram_src_handles_by_tp_ratio.clear() + self._dram_src_handles_by_block_size.clear() + try: + for engine_id in list(self._remote_agents): + self._cleanup_remote_engine(engine_id, log_eviction=False) + except Exception: + logger.exception("NIXL remote-engine cleanup failed at shutdown.") + try: + for desc in self._registered_descs: + self.nixl_wrapper.deregister_memory(desc) + finally: + self._registered_descs.clear() + # Drop cache references before their owners release registered + # host memory; handshake futures may outlive model-runner shutdown. + self.device_kv_caches = {} + self.host_xfer_buffers = {} + + def shutdown(self) -> None: """Shutdown the connector worker.""" if not hasattr(self, "_handshake_initiation_executor"): # error happens during init, no need to shutdown @@ -3227,22 +3327,4 @@ def shutdown(self): for handle in handles: self.nixl_wrapper.release_xfer_handle(handle) self._recving_transfers.clear() - for handle in self.src_xfer_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_block_size.clear() - for handles in self.src_xfer_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_tp_ratio.clear() - for handles in self._dram_src_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - self._dram_src_handles_by_tp_ratio.clear() - for handle in self._dram_src_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) - self._dram_src_handles_by_block_size.clear() - for engine_id in list(self._remote_agents): - self._cleanup_remote_engine(engine_id, log_eviction=False) - for desc in self._registered_descs: - self.nixl_wrapper.deregister_memory(desc) - self._registered_descs.clear() + self._finish_shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 7914b7ab3a19..105617bb8dac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -26,6 +26,7 @@ KVConnectorHandshakeMetadata, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -234,6 +235,18 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + results = self.connector_worker.get_transfer_results() + if ( + self.kv_transfer_config.kv_role == "kv_producer" + and self.connector_worker.pcp_rank > 0 + ): + results.finished_sending.clear() + return results + def get_block_ids_with_load_errors(self) -> set[int]: """Get block IDs that failed to load via NIXL.""" assert self.connector_worker is not None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 526d9c688133..f61fdfa83512 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -305,7 +305,9 @@ def add_new_req_to_recv( local_num_computed_blocks: tuple[int, ...] = (), ): req = self._add_new_req( - local_block_ids, kv_transfer_params, local_num_computed_blocks + local_block_ids, + kv_transfer_params, + local_num_computed_blocks, ) req.remote = RemoteMeta( block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index d39a9fe287c1..ceaa4fce577f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -166,11 +166,10 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1 local_block_ids = meta.local_physical_block_ids - local_region_groups = getattr(self, "region_group_ids", []) - remote_region_groups = getattr(self, "dst_region_group_ids", {}).get( - engine_id, local_region_groups - ) - if local_region_groups != remote_region_groups: + remote_region_groups = self.dst_region_group_ids[engine_id] + local_region_groups = self.region_group_ids or remote_region_groups + groups_differ = local_region_groups != remote_region_groups + if groups_differ: if not self.use_mla or self._has_mamba: raise NotImplementedError( "Different NIXL cache-group layouts are only supported for " @@ -292,6 +291,9 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]: spec.remote_rank ] + # Once a read routes the request to failure reporting, the + # scheduler may free and reuse its blocks, so no sibling READs + # may be posted (and P must not be notified). if not self._read_blocks( read_spec=spec, request_id=req_id, @@ -329,6 +331,10 @@ def _read_blocks( """ Post a READ point-to-point xfer request from a single local worker to a single remote worker. + + Returns True when the read was posted (or was unnecessary), False + when the request was routed to failure reporting — the caller must + not post further transfers for it. """ assert self.transfer_topo is not None remote_rank = read_spec.remote_rank @@ -416,11 +422,11 @@ def _read_blocks( dst_num_blocks=self.dst_num_blocks[dst_engine_id], block_size_ratio=None, physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical, - region_num_blocks=self.dst_region_num_blocks[dst_engine_id], + region_num_blocks=(self.dst_region_num_blocks.get(dst_engine_id) or None), region_group_ids=( list(range(self.num_regions)) if read_spec.block_ids_by_region - else self.dst_region_group_ids[dst_engine_id] + else (self.dst_region_group_ids.get(dst_engine_id) or None) ), uses_region_group_mapping=( self.num_regions > 1 @@ -433,11 +439,11 @@ def _read_blocks( dst_num_blocks=self.dst_num_blocks[self.engine_id], block_size_ratio=block_size_ratio, physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block, - region_num_blocks=self.dst_region_num_blocks[self.engine_id], + region_num_blocks=(self.dst_region_num_blocks.get(self.engine_id) or None), region_group_ids=( list(range(self.num_regions)) if read_spec.block_ids_by_region - else self.region_group_ids + else (self.region_group_ids or None) ), uses_region_group_mapping=( self.num_regions > 1 @@ -537,11 +543,11 @@ def _read_blocks_mixed( self._pending_recv_notifs.setdefault(request_id, []).append( (notif_agent, notif_id) ) - for index, handle in enumerate(handles): + for i, handle in enumerate(handles): try: self.nixl_wrapper.transfer(handle) except Exception: - for unstarted in handles[index:]: + for unstarted in handles[i:]: self.nixl_wrapper.release_xfer_handle(unstarted) raise self._recving_transfers[request_id].append(handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 21bd956c67e4..c3bcd2c8f907 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -41,6 +41,9 @@ import msgspec from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -145,7 +148,7 @@ def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]): self._push_writer_thread.start() logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank) - def shutdown(self): + def shutdown(self) -> None: self._push_writer_stop.set() # Unblock the writer if it's waiting in the no-active-state branch. self._push_writer_wake.set() @@ -522,7 +525,14 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): ) remote_block_ids = meta.remote.block_ids local_block_ids = meta.local_physical_block_ids - num_groups = len(local_block_ids) + local_region_groups = self.region_group_ids + remote_region_groups = self.dst_region_group_ids[engine_id] + groups_differ = local_region_groups != remote_region_groups + if groups_differ: + raise NotImplementedError( + "NixlPushConnector does not support different producer and " + "consumer cache-group layouts" + ) # MLA latent is replicated across D's TP ranks: the tp-mapping # collapses it to one rank (fine for reads), but push must WRITE every @@ -537,6 +547,8 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): else: write_ranks = list(plan.all_source_ranks) + num_groups = len(local_block_ids) + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds: return [ list(block_ids[g]) @@ -777,13 +789,14 @@ def _get_new_notifs(self) -> set[str]: self._reqs_to_send.pop(req_id, None) return notified_req_ids - def get_finished(self) -> tuple[set[str], set[str]]: + def get_transfer_results(self) -> KVConnectorTransferResults: # Engine main thread asking for completions: also wake the writer # so it gets a chance to drain NIXL notifs (heartbeats, completion # notifs, late PUSH_REGs) even if it had been parked. self._push_writer_wake.set() - done_sending, done_recving = super().get_finished() + results = super().get_transfer_results() + done_sending = results.finished_sending # ``_pop_done_transfers`` mutates ``_sending_transfers``; the # writer thread also appends to it, so guard the pop. @@ -803,4 +816,4 @@ def get_finished(self) -> tuple[set[str], set[str]]: if done_sending: self._push_writer_wake.set() - return done_sending, done_recving + return results diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index bc78b7298822..42b9dd6aa392 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -12,6 +12,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, FullAttentionSpec, + KVCacheGroupRole, KVCacheSpec, MLAAttentionSpec, SlidingWindowMLASpec, @@ -31,6 +32,16 @@ from vllm.v1.kv_cache_interface import KVCacheConfig +def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: + if kv_cache_config.hisparse_host_num_blocks is None: + return tuple(range(len(kv_cache_config.kv_cache_groups))) + return tuple( + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if group.role is KVCacheGroupRole.HISPARSE_INDEXER + ) + + def build_offloading_config( vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig", @@ -43,15 +54,22 @@ def build_offloading_config( engine_id = kv_transfer_config.engine_id parallel_config = vllm_config.parallel_config + selected_groups = tuple( + (group_id, kv_cache_config.kv_cache_groups[group_id]) + for group_id in get_offloading_group_ids(kv_cache_config) + ) + if not selected_groups: + raise ValueError("KV offloading found no eligible cache groups.") groups = tuple( OffloadingGroupConfig( + group_id=group_id, tokens_per_block=resolve_dcp_kv_block_size( group.kv_cache_spec, parallel_config.decode_context_parallel_size, ), layer_names=tuple(group.layer_names), ) - for group in kv_cache_config.kv_cache_groups + for group_id, group in selected_groups ) _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config) @@ -103,11 +121,20 @@ def build_offloading_config( ) worker_kv_bytes_per_block = 0 - if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors: + all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups) + if ( + all_groups_selected + and kv_cache_config.num_blocks > 0 + and kv_cache_config.kv_cache_tensors + ): # Every KVCacheTensor describes placement within the same backing allocation, # so its size is the total, not a per-tensor share. total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks + elif kv_cache_config.num_blocks > 0: + worker_kv_bytes_per_block = sum( + group.kv_cache_spec.page_size_bytes for _, group in selected_groups + ) single_group_spec = ( kv_cache_config.kv_cache_groups[0].kv_cache_spec diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a3edb7875030..86e351093917 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -152,8 +152,9 @@ def resolve_mamba_align_size( such groups agree on the same value. """ mamba_align_size: int | None = None - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec + for group in spec.config.groups: + tokens_per_block = group.tokens_per_block + kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in ( "align", "all", @@ -188,8 +189,9 @@ def from_spec( # each segment can never serve a load hit. Relevant for hybrid # architectures like DeepSeek V4 (MLA + SWA groups). full_attn_tokens_per_chunk: set[int] = set() - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec + for group in spec.config.groups: + tokens_per_block = group.tokens_per_block + kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec sw = get_sliding_window_size_in_chunks( kv_spec, tokens_per_block * spec.blocks_per_chunk ) @@ -214,9 +216,9 @@ def from_spec( retention_interval = vllm_config.cache_config.prefix_cache_retention_interval eagle_groups = { - idx - for idx, g in enumerate(kv_cache_config.kv_cache_groups) - if g.is_eagle_group + group.group_id + for group in spec.config.groups + if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group } use_eagle_block_drop = ( @@ -249,8 +251,10 @@ def from_spec( ) kv_group_configs_list: list[GroupOffloadConfig] = [] - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_cache_group = kv_cache_config.kv_cache_groups[idx] + for group in spec.config.groups: + group_id = group.group_id + tokens_per_block = group.tokens_per_block + kv_cache_group = kv_cache_config.kv_cache_groups[group_id] kv_spec = kv_cache_group.kv_cache_spec sw = get_sliding_window_size_in_chunks( kv_spec, tokens_per_block * spec.blocks_per_chunk @@ -261,7 +265,7 @@ def from_spec( ) kv_group_configs_list.append( GroupOffloadConfig( - group_idx=idx, + group_idx=group_id, tokens_per_block=tokens_per_block, tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk, hashes_per_chunk=( @@ -272,7 +276,7 @@ def from_spec( kv_cache_spec=kv_spec, manager_cls=manager_cls, kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group), - is_eagle_group=idx in eagle_groups, + is_eagle_group=group_id in eagle_groups, requires_cow_source=( isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align" @@ -408,9 +412,10 @@ def update_block_id_groups( if new_block_id_groups is None: return - assert len(new_block_id_groups) == len(self.group_states) - for group_state, new_blocks in zip(self.group_states, new_block_id_groups): - group_state.block_ids.extend(new_blocks) + for group_config, group_state in zip( + self.config.kv_group_configs, self.group_states + ): + group_state.block_ids.extend(new_block_id_groups[group_config.group_idx]) def storable_chunks( self, @@ -541,11 +546,11 @@ def __init__( full_attention_groups: list[int] = [] sliding_window_groups: list[int] = [] - for group_config in self.config.kv_group_configs: + for config_idx, group_config in enumerate(self.config.kv_group_configs): if group_config.sliding_window_size_in_chunks is None: - full_attention_groups.append(group_config.group_idx) + full_attention_groups.append(config_idx) else: - sliding_window_groups.append(group_config.group_idx) + sliding_window_groups.append(config_idx) # sort sliding window groups by window size in decreasing order def _sliding_window_sort_key(i: int) -> int: @@ -737,7 +742,11 @@ def _touch(self, req_status: RequestOffloadState): req_status.req_context, ) - def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None: + def _lookup_complete_chunks( + self, + req_status: RequestOffloadState, + max_num_new_tokens: int | None = None, + ) -> int | None: """ Find how many tokens beyond num_locally_computed_tokens can be loaded. @@ -748,6 +757,10 @@ def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None """ num_computed_tokens = req_status.num_locally_computed_tokens max_hit_size_tokens: int = req_status.req.num_tokens + if max_num_new_tokens is not None: + max_hit_size_tokens = min( + max_hit_size_tokens, num_computed_tokens + max_num_new_tokens + ) if self._sliding_window_groups: # the last prompt token has to be recomputed to get the logprobs # for sliding window attention, we must reduce by 1 to make sure @@ -929,8 +942,12 @@ def _make_boundary_key( hash_idx = boundary_tokens // self.config.tokens_per_hash - 1 return make_offload_key(request.block_hashes[hash_idx], group_idx) - def _lookup(self, req_status: RequestOffloadState) -> int | None: - complete_hit = self._lookup_complete_chunks(req_status) + def _lookup( + self, + req_status: RequestOffloadState, + max_num_new_tokens: int | None = None, + ) -> int | None: + complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens) req_status.partial_tail_boundary = None if complete_hit is None or not self.config.supports_partial_tail: return complete_hit @@ -939,9 +956,10 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: complete_boundary = local_tokens + complete_hit tokens_per_hash = self.config.tokens_per_hash block_end = complete_boundary + self._partial_tail_block_size - max_boundary = round_down( - min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash - ) + max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) + if max_num_new_tokens is not None: + max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) + max_boundary = round_down(max_boundary, tokens_per_hash) if max_boundary <= complete_boundary: return complete_hit @@ -990,7 +1008,10 @@ def on_new_request(self, request: Request) -> None: self._req_status[request.request_id] = req_status def get_num_new_matched_tokens( - self, request: Request, num_computed_tokens: int + self, + request: Request, + num_computed_tokens: int, + max_num_new_tokens: int | None = None, ) -> tuple[int | None, bool]: """ Get number of new tokens that can be loaded beyond the @@ -1030,7 +1051,7 @@ def get_num_new_matched_tokens( num_hit_tokens = 0 else: lookup_start = time.monotonic() - num_hit_tokens = self._lookup(req_status) + num_hit_tokens = self._lookup(req_status, max_num_new_tokens) self._connector_stats.observe_histogram( _ConnectorMetricName.LOOKUP_SYNC_DELAY, time.monotonic() - lookup_start, @@ -1065,11 +1086,11 @@ def update_state_after_alloc( # per group group_sizes: list[int] = [] block_indices: list[int] = [] - for group_config, group_state, group_blocks in zip( + for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states, - blocks.blocks, ): + group_blocks = blocks.blocks[group_config.group_idx] self._current_batch_allocated_block_ids.update( block.block_id for block in group_blocks if block.block_id != 0 ) @@ -1184,7 +1205,8 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: for grp_idx in self._sliding_window_groups ) req_status.update_block_id_groups(new_block_id_groups) - for new_blocks in new_block_id_groups: + for group_config in self.config.kv_group_configs: + new_blocks = new_block_id_groups[group_config.group_idx] for bid in new_blocks: if bid != 0: self._current_batch_allocated_block_ids.add(bid) @@ -1218,6 +1240,10 @@ def _build_aligned_boundary_store_jobs( ) -> dict[int, TransferJob]: store_jobs: dict[int, TransferJob] = {} num_groups = len(self.config.kv_group_configs) + config_idx_by_group = { + config.group_idx: idx + for idx, config in enumerate(self.config.kv_group_configs) + } for req_id, entries in handoffs.items(): req_status = self._req_status.get(req_id) if req_status is None: @@ -1225,7 +1251,10 @@ def _build_aligned_boundary_store_jobs( req = req_status.req max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens) for group_idx, block_id, boundary in entries: - group_config = self.config.kv_group_configs[group_idx] + config_idx = config_idx_by_group.get(group_idx) + if config_idx is None: + continue + group_config = self.config.kv_group_configs[config_idx] if ( block_id == 0 or boundary > max_boundary @@ -1254,9 +1283,11 @@ def _build_aligned_boundary_store_jobs( fenced_block_ids=[block_id], ) group_sizes = [0] * num_groups - group_sizes[group_idx] = 1 + group_sizes[config_idx] = 1 block_indices = [0] * num_groups - block_indices[group_idx] = boundary // group_config.tokens_per_block - 1 + block_indices[config_idx] = ( + boundary // group_config.tokens_per_block - 1 + ) store_jobs[job_id] = TransferJob( req_id=req_id, src_spec=GPULoadStoreSpec( @@ -1300,6 +1331,12 @@ def _build_partial_tail_store_jobs( assert len(boundaries) == 1 boundary = boundaries.pop() req = req_status.req + group_states = { + group.group_idx: state + for group, state in zip( + self.config.kv_group_configs, req_status.group_states + ) + } max_boundary = min( req.num_prompt_tokens, req_status.max_offload_tokens or req.num_prompt_tokens, @@ -1315,7 +1352,7 @@ def _build_partial_tail_store_jobs( block_idx = boundary // self._partial_tail_block_size if any( group.group_idx not in self._cow_source_groups - and block_idx >= len(req_status.group_states[group.group_idx].block_ids) + and block_idx >= len(group_states[group.group_idx].block_ids) for group in self.config.kv_group_configs ): continue @@ -1326,7 +1363,7 @@ def _build_partial_tail_store_jobs( block_ids = [ cow_blocks[group.group_idx] if group.group_idx in self._cow_source_groups - else req_status.group_states[group.group_idx].block_ids[block_idx] + else group_states[group.group_idx].block_ids[block_idx] for group in self.config.kv_group_configs ] assert all(block_id != 0 for block_id in block_ids) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index ecdae9c4f2d0..14e826e5525c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -14,6 +14,9 @@ OffloadingWorkerMetadata, ReqId, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + get_offloading_group_ids, +) from vllm.logger import init_logger from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -65,7 +68,10 @@ def _init_worker(self, kv_caches: CanonicalKVCaches) -> None: def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): kv_cache_config = self.kv_cache_config - num_blocks = kv_cache_config.num_blocks + selected_group_ids = get_offloading_group_ids(kv_cache_config) + selected_groups = tuple( + kv_cache_config.kv_cache_groups[group_id] for group_id in selected_group_ids + ) mappings = derive_canonical_mappings( self.vllm_config, kv_cache_config, kv_caches ) @@ -76,7 +82,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} - for kv_cache_group in kv_cache_config.kv_cache_groups: + for kv_cache_group in selected_groups: + assert not kv_cache_group.host_resident + num_blocks = kv_cache_config.num_blocks group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs): @@ -131,8 +139,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ), None, ) - if packed_layer_name is not None: + if packed_layer_name is not None and len(selected_groups) == len( + kv_cache_config.kv_cache_groups + ): (tensor,) = tensors_per_block[packed_layer_name] + num_blocks = tensor.shape[0] block_stride = tensor.stride(0) packed_tensor = tensor.as_strided( (num_blocks, block_stride), @@ -142,10 +153,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._init_worker( CanonicalKVCaches( [CanonicalKVCacheTensor(packed_tensor, block_stride)], - [ - [CanonicalKVCacheRef(0, block_stride)] - for _ in kv_cache_config.kv_cache_groups - ], + [[CanonicalKVCacheRef(0, block_stride)] for _ in selected_groups], ) ) return @@ -192,7 +200,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) group_data_refs: list[list[CanonicalKVCacheRef]] = [] - for kv_cache_group in kv_cache_config.kv_cache_groups: + for kv_cache_group in selected_groups: group_refs: list[CanonicalKVCacheRef] = [] for layer_name in kv_cache_group.layer_names: group_refs += block_data_refs[layer_name] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index a64f97805ecc..73570881753d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from functools import cached_property from typing import Any import torch @@ -25,6 +26,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, + get_offloading_group_ids, ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, @@ -38,6 +40,7 @@ ) from vllm.forward_context import ForwardContext from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig @@ -47,6 +50,25 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): + @cached_property + def _bounding_group_ids(self) -> tuple[int, ...]: + """Prefix-cacheable groups this connector does not offload. + + Never offer tokens for a range some group cannot cover: past such a + group's own cached prefix its KV would be left unwritten. + """ + offloaded = set(get_offloading_group_ids(self._kv_cache_config)) + return tuple( + group_id + for group_id, group in enumerate(self._kv_cache_config.kv_cache_groups) + if group.kv_cache_spec.prefix_cacheable and group_id not in offloaded + ) + + @property + def scheduler(self) -> OffloadingConnectorScheduler: + assert self.connector_scheduler is not None + return self.connector_scheduler + @property def requires_kv_delivery(self) -> bool: # Runs as kv_both, but is a best-effort cache: a dropped save is just a @@ -138,8 +160,29 @@ def get_num_new_matched_tokens( ) -> tuple[int | None, bool]: assert self.connector_scheduler is not None return self.connector_scheduler.get_num_new_matched_tokens( - request, num_computed_tokens + request, + num_computed_tokens, + max_num_new_tokens=self._max_loadable_tokens(request, num_computed_tokens), + ) + + def _max_loadable_tokens( + self, request: "Request", num_computed_tokens: int + ) -> int | None: + """How far past ``num_computed_tokens`` a load may reach, if bounded. + + Bounded by the deepest prefix the groups this connector does not + offload already hold, since nothing refills them beyond it. + """ + if not self._bounding_group_ids: + return None + assert self._kv_cache_manager is not None + coordinator = self._kv_cache_manager.coordinator + assert isinstance(coordinator, HybridKVCacheCoordinator) + _, per_group_hits = coordinator.find_longest_cache_hit_per_group( + request.block_hashes, request.num_tokens - 1 ) + bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids) + return max(0, bound - num_computed_tokens) def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int @@ -185,6 +228,8 @@ def take_events(self) -> Iterable[KVCacheEvent]: @classmethod def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None: + if vllm_config.attention_config.hisparse_config is not None: + return "BLHNC" return "LBHNC" def reset_cache(self) -> bool | None: diff --git a/vllm/distributed/kv_transfer/kv_transfer_state.py b/vllm/distributed/kv_transfer/kv_transfer_state.py index 211d33fff215..690eb9fa23d9 100644 --- a/vllm/distributed/kv_transfer/kv_transfer_state.py +++ b/vllm/distributed/kv_transfer/kv_transfer_state.py @@ -78,13 +78,11 @@ def ensure_kv_transfer_initialized( global _KV_CONNECTOR_AGENT - if vllm_config.kv_transfer_config is None: + if _KV_CONNECTOR_AGENT is not None: return - if ( - vllm_config.kv_transfer_config.is_kv_transfer_instance - and _KV_CONNECTOR_AGENT is None - ): + kv_transfer_config = vllm_config.kv_transfer_config + if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance: # NIXL P/D requires an interleave_size equal to block_size. vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config) _sync_engine_id_across_tp(vllm_config) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f9b86bd810cd..21596eb04365 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -207,7 +207,7 @@ from dataclasses import dataclass, replace from enum import Enum from math import lcm -from typing import ClassVar, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast import numpy as np import torch @@ -284,6 +284,7 @@ CommonAttentionMetadata, MLAAttentionImpl, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.prefill import ( MLAPrefillBackend, get_mla_prefill_backend, @@ -309,6 +310,11 @@ get_kv_quant_mode, ) +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) + logger = init_logger(__name__) _FP8_DTYPE = current_platform.fp8_dtype() @@ -433,6 +439,7 @@ def __init__( use_sparse: bool = False, indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None, non_causal_multi_token_decode: bool = False, sliding_window: int | None = None, prefill_backend_cls: type[MLAPrefillBackend] | None = None, @@ -538,9 +545,11 @@ def __init__( # explicitly so backbone "skip" layers (indexer=None) still find it. if use_sparse: extra_impl_args["topk_indices_buffer"] = topk_indices_buffer + if index_group_builder is not None: + extra_impl_args["index_group_builder"] = index_group_builder impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) - self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass + impl = impl_cls( num_heads=self.num_heads, head_size=self.head_size, scale=self.scale, @@ -562,6 +571,13 @@ def __init__( indexer=indexer, **extra_impl_args, ) + self.impl = impl # type: ignore[assignment] + index_group = getattr(impl, "index_group", None) + self.hisparse_cache = ( + index_group.cache(cast(Any, impl).index_group_index) + if isinstance(index_group, HiSparseMLAIndexGroup) + else None + ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) # MLA reads this weight directly to build W_UK/W_UV, so a backend must @@ -723,6 +739,61 @@ def chunked_prefill_workspace_size(self) -> int: ) return self._chunked_prefill_workspace_size + def update_kv_cache( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor | None, + attn_metadata: "MLACommonMetadata | None", + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + cache = self.hisparse_cache + if slot_mapping is None or (cache is not None and cache.dummy_batch): + return + kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs( + kv_c_normed, + k_pe, + slot_mapping, + attn_metadata.num_decode_tokens if attn_metadata is not None else None, + self.use_pcp, + ) + assert slot_mapping is not None + if cache is not None: + kv_cache, slot_mapping, num_rows = cache.write_target( + kv_c_normed.shape[0], slot_mapping.numel() + ) + kv_c_normed = kv_c_normed[:num_rows] + k_pe = k_pe[:num_rows] + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_c_normed, + k_pe, + kv_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + ) + if cache is not None: + mirror_target = cache.mirror_write_target(kv_c_normed.shape[0]) + if mirror_target is not None: + mirror_cache, mirror_slots = mirror_target + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_c_normed, + k_pe, + mirror_cache, + mirror_slots, + kv_cache_dtype, + k_scale, + ) + + def prepare_kv_cache_update( + self, attn_metadata: "MLACommonMetadata | None" + ) -> None: + cache = self.hisparse_cache + if cache is not None and cache.runtime.is_group_leader: + cache.prepare_group_for_batch(attn_metadata) + def forward( self, q: torch.Tensor, @@ -750,25 +821,18 @@ def forward( f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. " ) layer_slot_mapping = slot_mapping.get(self.layer_name) - kv_for_cache, kpe_for_cache, layer_slot_mapping = ( - maybe_gather_mla_latent_cache_inputs( - kv_c_normed, - k_pe, - layer_slot_mapping, - attn_metadata.num_decode_tokens - if attn_metadata is not None - else None, - self.use_pcp, - ) - ) - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_for_cache, - kpe_for_cache, + self.prepare_kv_cache_update(attn_metadata) + self.update_kv_cache( + kv_c_normed, + k_pe, self_kv_cache, layer_slot_mapping, + attn_metadata, self.kv_cache_dtype, self._k_scale, ) + if self.hisparse_cache is not None: + self.hisparse_cache.finish_kv_update() output = torch.empty(output_shape, dtype=q.dtype, device=q.device) self.forward_impl( q, @@ -1113,6 +1177,8 @@ def forward_impl( return output_padded def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool: + if self.hisparse_cache is not None: + return False prefill = attn_metadata.prefill if prefill is None: return False @@ -1296,6 +1362,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: ) return MLAAttentionSpec( **common_kwargs, + is_index_group_leader=self.indexer is not None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) @@ -1349,19 +1416,14 @@ def unified_mla_kv_cache_update( attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context( layer_name ) + attn_layer.prepare_kv_cache_update(attn_metadata) if layer_slot_mapping is not None: - kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs( - kv_c_normed, - k_pe, - layer_slot_mapping, - attn_metadata.num_decode_tokens if attn_metadata is not None else None, - attn_layer.use_pcp, - ) - attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined] + attn_layer.update_kv_cache( kv_c_normed, k_pe, kv_cache, layer_slot_mapping, + attn_metadata, kv_cache_dtype, k_scale, ) @@ -1409,6 +1471,8 @@ def unified_mla_attention_with_output( del kv_cache_dummy_dep layer_name = _resolve_layer_name(layer_name) attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name) + if layer.hisparse_cache is not None: + layer.hisparse_cache.finish_kv_update() layer.forward_impl( q, kv_c_normed, @@ -1501,6 +1565,7 @@ class MLACommonBackend(AttentionBackend): def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack an inline fp32 scale pair after the latent data (single-sided: ``head_size_v == 0`` for MLA).""" + spec = super().customize_spec(spec) mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: return spec @@ -2090,7 +2155,11 @@ def padded_rows(rows: int) -> int: local_token_offset += num_local_tokens seq_lens_cpu = _flat_int32(seq_lens_flat) - starts = _flat_int32(starts_flat).to(device, non_blocking=True) + starts_and_context_lens = _flat_int32(starts_flat + context_lens).to( + device, non_blocking=True + ) + starts = starts_and_context_lens[: len(starts_flat)] + context_lens_gpu = starts_and_context_lens[len(starts_flat) :] cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True) cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True) token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to( @@ -2145,7 +2214,7 @@ def padded_rows(rows: int) -> int: chunks.append(chunk) return MLACommonPrefillMetadata.ChunkedContextMetadata( - context_lens=context_lens_cpu.to(device, non_blocking=True), + context_lens=context_lens_gpu, workspace=chunked_prefill_workspace, chunks=chunks, context_lens_list=context_lens, @@ -2835,12 +2904,12 @@ def _compute_prefill_context( output_lse = None for chunk in chunked_context.chunks: toks = chunk.num_context_tokens - block_table = prefill_metadata.block_table[chunk.request_slice] + chunk_block_table = prefill_metadata.block_table[chunk.request_slice] if self.kv_cache_dtype == "fp8_ds_mla": ops.cp_gather_and_upconvert_fp8_kv_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, workspace_starts=chunk.cu_seq_lens, batch_size=chunk.num_requests, seq_starts=chunk.starts, @@ -2853,7 +2922,7 @@ def _compute_prefill_context( ops.gather_mla_context_cache_cpu( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, starts=chunk.starts, cu_seq_lens=chunk.cu_seq_lens, ) @@ -2861,7 +2930,7 @@ def _compute_prefill_context( ops.gather_and_maybe_dequant_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace, - block_table=block_table, + block_table=chunk_block_table, cu_seq_lens=chunk.cu_seq_lens, token_to_seq=chunk.token_to_seq, num_tokens=toks, @@ -2874,7 +2943,7 @@ def _compute_prefill_context( ops.cp_gather_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, cu_seq_lens=chunk.cu_seq_lens, batch_size=chunk.num_requests, seq_starts=chunk.starts, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index fd29debbb759..5baddcd0a5cb 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -3,13 +3,15 @@ """Shared MHA implementation and metadata builder for sparse MLA backends.""" import math +from dataclasses import dataclass from shutil import which -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar import numpy as np import torch from vllm import _custom_ops as ops +from vllm.config import get_current_vllm_config from vllm.distributed import ( get_dcp_group, get_tensor_model_parallel_world_size, @@ -17,6 +19,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBaseImpl, + MLACommonDecodeMetadata, MLACommonMetadata, MLACommonPrefillMetadata, accumulate_mla_context_chunk, @@ -33,11 +36,19 @@ is_quantized_kv_cache, np_to_pinned_tensor, ) -from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder +from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroup, + SparseMLAIndexGroupBuilder, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.hisparse.runtime import ( + HiSparsePrefillStagingPlan, + build_hisparse_prefill_staging_plan, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -47,8 +58,6 @@ logger = init_logger(__name__) -T = TypeVar("T", bound=AttentionMetadata) - GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024 # 128 MiB @@ -116,9 +125,45 @@ def _is_masked_mha_available( return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype) +def _use_dense_mha_prefill( + vllm_config: "VllmConfig", prefill_max_seq_len: int, topk_tokens: int +) -> bool: + attention_config = vllm_config.attention_config + return ( + attention_config.hisparse_config is None + and not attention_config.sparse_mla_force_mqa + and prefill_max_seq_len <= topk_tokens + ) + + +@dataclass +class SparseMLAPrefillMetadata(MLACommonPrefillMetadata): + host_staging_plan: HiSparsePrefillStagingPlan | None = None + + +@dataclass(kw_only=True) +class SparseMLACommonMetadata(MLACommonMetadata[MLACommonDecodeMetadata]): + block_table: torch.Tensor + req_id_per_token: torch.Tensor + seq_lens: torch.Tensor | None = None + block_size: int = 64 + topk_tokens: int = 2048 + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + decode_max_query_len: int = 0 + prefill_max_seq_len: int = 0 + prefill: SparseMLAPrefillMetadata | None = None + cp_kv_cache_interleave_size: int = 1 + + +T = TypeVar("T", bound=SparseMLACommonMetadata) + + class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): metadata_cls: type[T] require_uniform_decodes: ClassVar[bool] = False + hisparse_supports_multi_token_decode: ClassVar[bool] = False def __init__( self, @@ -138,6 +183,12 @@ def __init__( dtype=torch.int32, device=device, ) + self._context_lens_cpu = torch.empty( + vllm_config.scheduler_config.max_num_seqs, + dtype=torch.int32, + device="cpu", + pin_memory=PIN_MEMORY, + ) parallel_config = vllm_config.parallel_config self.use_pcp = parallel_config.prefill_context_parallel_size > 1 try: @@ -197,6 +248,22 @@ def __init__( layer_prefill_backend.clone() if layer_prefill_backend is not None else None ) + def _init_reorder_batch_threshold( + self, + reorder_batch_threshold: int | None = 1, + supports_spec_as_decode: bool = False, + supports_dcp_with_varlen: bool = False, + ) -> None: + if self.vllm_config.attention_config.hisparse_config is not None: + reorder_batch_threshold = 1 + if not self.hisparse_supports_multi_token_decode: + supports_spec_as_decode = False + super()._init_reorder_batch_threshold( + reorder_batch_threshold, + supports_spec_as_decode, + supports_dcp_with_varlen, + ) + @staticmethod def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int: scheduler_config = vllm_config.scheduler_config @@ -245,9 +312,7 @@ def _build_chunked_context_fields( seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None - context_lens_cpu = torch.empty( - num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY - ) + context_lens_cpu = self._context_lens_cpu[:num_prefills] torch.subtract( seq_lens_cpu[num_decodes : num_decodes + num_prefills], prefill_query_lens_cpu, @@ -289,34 +354,66 @@ def build( prefill_max_query_len, prefill_query_lens_cpu, ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills) + decode_max_query_len = 0 + if num_decodes: + if common_attn_metadata.max_logits_per_req is not None: + decode_max_query_len = common_attn_metadata.max_logits_per_req + elif num_prefills == 0: + decode_max_query_len = common_attn_metadata.max_query_len + else: + decode_query_lens_cpu = torch.diff( + common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] + ) + decode_max_query_len = int(decode_query_lens_cpu.max().item()) prefill_max_seq_len = 0 - prefill: MLACommonPrefillMetadata | None = None + prefill: SparseMLAPrefillMetadata | None = None if num_prefills > 0 and self._prefill_backend is not None: seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None prefill_max_seq_len = int( seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item() ) - prefill = MLACommonPrefillMetadata( - block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...], + block_table = common_attn_metadata.block_table_tensor[num_decodes:, ...] + chunked_context = self._build_chunked_context_fields( + common_attn_metadata, + num_decodes, + num_prefills, + prefill_query_lens_cpu, + ) + staging_plan = None + if self.vllm_config.attention_config.hisparse_config is not None: + prefill_seq_lens_cpu = seq_lens_cpu[ + num_decodes : num_decodes + num_prefills + ] + staging_block_capacity = int( + ( + (prefill_seq_lens_cpu + self.kv_cache_spec.block_size - 1) + // self.kv_cache_spec.block_size + ).sum() + ) + staging_plan = build_hisparse_prefill_staging_plan( + block_table, + common_attn_metadata.seq_lens[num_decodes:], + self.kv_cache_spec.block_size, + staging_block_capacity, + ) + if chunked_context is not None: + block_table = staging_plan.block_table + prefill = SparseMLAPrefillMetadata( + block_table=block_table, query_start_loc=prefill_query_start_loc, max_query_len=prefill_max_query_len, query_lens_cpu=prefill_query_lens_cpu, - chunked_context=self._build_chunked_context_fields( - common_attn_metadata, - num_decodes, - num_prefills, - prefill_query_lens_cpu, - ), + chunked_context=chunked_context, q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, - use_dense_mha=( - prefill_max_seq_len <= self.topk_tokens - and not self.vllm_config.attention_config.sparse_mla_force_mqa + use_dense_mha=_use_dense_mha_prefill( + self.vllm_config, prefill_max_seq_len, self.topk_tokens ), topk_mask_workspace=self.topk_mask_workspace, + host_staging_plan=staging_plan, ) self._prefill_backend.prepare_metadata(prefill) @@ -335,6 +432,7 @@ def build( num_decodes=num_decodes, num_prefills=num_prefills, num_decode_tokens=num_decode_tokens, + decode_max_query_len=decode_max_query_len, prefill_max_seq_len=prefill_max_seq_len, prefill=prefill, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, @@ -560,6 +658,7 @@ def __init__( kv_b_proj: "ColumnParallelLinear", indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, q_pad_num_heads: int | None = None, ) -> None: super().__init__( @@ -577,6 +676,21 @@ def __init__( ) self.init_topk_indices_buffer(indexer, topk_indices_buffer) + self.index_group: SparseMLAIndexGroup | None = None + self.index_group_index = 0 + if index_group_builder is None and self.topk_indices_buffer is not None: + index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) + if index_group_builder is not None: + vllm_config = get_current_vllm_config() + self.index_group, self.index_group_index = ( + index_group_builder.register_layer( + indexer is not None, + vllm_config, + head_size=head_size, + kv_cache_dtype=kv_cache_dtype, + ) + ) + self._use_flashinfer_concat_mla_k = ( has_flashinfer() and which("ninja") is not None @@ -593,6 +707,33 @@ def __init__( kv_cache_dtype=kv_cache_dtype, ) + def record_logical_topk_ready(self) -> None: + if self.index_group is not None: + self.index_group.set_logical_topk_ready(self.index_group_index) + + def prepare_for_batch(self, attn_metadata: T | None) -> None: + if self.index_group is not None: + self.index_group.prepare_for_batch(self.index_group_index, attn_metadata) + + def _convert_logical_to_physical_topk( + self, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + group = self.index_group + assert group is not None + assert self.dcp_world_size == 1 + return group.convert_logical_to_physical_topk( + self.index_group_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=return_valid_counts, + ) + @staticmethod def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool: """Whether this prefill batch's top-k masks fit the workspace.""" @@ -853,8 +994,10 @@ def forward_mha( # type: ignore[override] output: torch.Tensor, output_scale: torch.Tensor | None = None, ) -> None: - prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined] + prefill_metadata = attn_metadata.prefill + assert prefill_metadata is not None + prefill_max_seq_len = attn_metadata.prefill_max_seq_len + topk_tokens = attn_metadata.topk_tokens force_dense = getattr(self, "_sparse_mla_force_dense_mha", False) force_masked = getattr(self, "_sparse_mla_force_masked_mha", False) if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked): @@ -863,7 +1006,7 @@ def forward_mha( # type: ignore[override] kv_c_normed, k_pe, kv_c_and_k_pe_cache, - cast(MLACommonMetadata, attn_metadata), + attn_metadata, k_scale, output, output_scale, @@ -871,13 +1014,11 @@ def forward_mha( # type: ignore[override] assert output_scale is None assert self.masked_mha_available - prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined] - assert prefill_metadata is not None assert prefill_metadata.query_lens_cpu is not None assert self.topk_indices_buffer is not None q_lens = prefill_metadata.query_lens_cpu.tolist() - num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined] + num_decode_tokens = attn_metadata.num_decode_tokens topk_all = self.topk_indices_buffer[ num_decode_tokens : num_decode_tokens + q.shape[0] ] diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 421181306236..9dd5ad67736c 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from typing import TYPE_CHECKING import torch @@ -11,6 +12,11 @@ from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) + @dataclass class MLAModules: @@ -30,6 +36,7 @@ class MLAModules: topk_indices_buffer: torch.Tensor | None indexer_rotary_emb: torch.nn.Module | None = None g_proj: torch.nn.Module | None = None + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None # --8<-- [start:multi_head_latent_attention] @@ -128,6 +135,7 @@ def __init__( use_sparse=self.is_sparse, indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, + index_group_builder=mla_modules.index_group_builder, non_causal_multi_token_decode=non_causal_multi_token_decode, ) indexer_op = getattr(self.indexer, "indexer_op", None) @@ -223,6 +231,8 @@ def forward( if self.indexer and self.is_sparse and not self.skip_topk: self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) + if self.is_sparse: + self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined] if llama_4_scaling is not None: q *= llama_4_scaling diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index bb00a9cfe103..817f8ad361fb 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -425,7 +425,6 @@ def sparse_attn_indexer( quant_block_size, scale_fmt, ) - # The indexer and main MLA may classify the same short extend differently # because they use independent decode thresholds. Only the main MLA route # can determine whether the top-k indices will be consumed. @@ -955,6 +954,7 @@ def forward_cuda( self.dcp_rank, self.dcp_world_size, self.cp_kv_cache_interleave_size, + False, candidate_blocks=self.candidate_blocks, candidate_block_size=self.candidate_block_size, candidate_write=self.candidate_write, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index bc99509c7f45..bdc634102e2f 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -103,10 +103,14 @@ from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionBackend, AttentionType +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerBackend, ) -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, SparseCacheRole from .interfaces import ( MixtureOfExperts, @@ -656,6 +660,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, + cache_role=SparseCacheRole.INDEXER, ) # Only has one vector instead of K + V def forward(self): ... @@ -1012,6 +1017,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, input_size: int | None = None, reduce_results: bool = True, non_causal_multi_token_decode: bool = False, @@ -1199,6 +1205,7 @@ def __init__( indexer_rotary_emb=self.indexer_rope_emb, is_sparse=self.is_v32, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) self.mla_attn = MultiHeadLatentAttentionWrapper( @@ -1242,6 +1249,7 @@ def __init__( prefix: str, config: DeepseekV2Config | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -1290,6 +1298,11 @@ def __init__( and parallel_config.pipeline_parallel_size == 1 and is_moe_layer ) + attn_kwargs = ( + {"index_group_builder": index_group_builder} + if attn_cls is DeepseekV2MLAAttention + else {} + ) self.self_attn = attn_cls( vllm_config=vllm_config, config=config, @@ -1306,6 +1319,7 @@ def __init__( prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, reduce_results=not self.use_sequence_parallel_moe, + **attn_kwargs, ) if is_moe_layer: @@ -1429,6 +1443,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -1445,6 +1467,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index d0119719af0e..019380540d08 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -67,6 +67,10 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -108,6 +112,7 @@ def __init__( prefix: str, config: "Glm4MoeLiteConfig | None" = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -136,6 +141,11 @@ def __init__( attn_cls = Glm4MoeLiteMLAAttention else: attn_cls = Glm4MoeLiteAttention + attn_kwargs: dict[str, typing.Any] = ( + {"index_group_builder": index_group_builder} + if attn_cls is Glm4MoeLiteMLAAttention + else {} + ) self.self_attn = attn_cls( vllm_config=vllm_config, @@ -152,6 +162,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + **attn_kwargs, ) if ( @@ -234,6 +245,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -252,6 +271,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config=config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index 1d30a0eaf6c6..2ef9407fb750 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -51,6 +51,10 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm4_moe_lite import ( Glm4MixtureOfExperts, @@ -108,6 +112,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config @@ -117,6 +129,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=prefix, config=self.config, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) def forward( diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 5209b944345b..91ae7c44c239 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -47,6 +47,35 @@ _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 +def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None: + """Autotune each HiSparse FlashInfer sparse-MLA configuration.""" + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + ) + + tuned: set[tuple[object, ...]] = set() + for layer in runner.vllm_config.compilation_config.static_forward_context.values(): + impl = getattr(layer, "impl", None) + if not isinstance(impl, FlashInferMLASparseImpl): + continue + if getattr(layer, "hisparse_cache", None) is None: + continue + if impl.topk_indices_buffer is None: + continue + key = ( + impl.kv_cache_dtype, + impl.num_heads, + impl.qk_nope_head_dim, + impl.qk_rope_head_dim, + impl.kv_lora_rank, + impl.topk_indices_buffer.shape[1], + ) + if key in tuned: + continue + impl.autotune_hisparse_decode(layer) + tuned.add(key) + + def _attention_backend_name(backend: object) -> str | None: get_name = getattr(backend, "get_name", None) if get_name is None: diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c9c740f3684c..2991dc1027a4 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -25,6 +25,7 @@ write_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( + autotune_hisparse_flashinfer_attention, deepseek_v4_sparse_mla_attention_warmup, flashinfer_sparse_mla_decode_autotune_warmup, ) @@ -353,9 +354,12 @@ def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, .. return tuple(dict.fromkeys(token_counts)) -def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None: +def _run_flashinfer_autotune_dummy_runs( + runner: "GPUModelRunner", *, skip_attn: bool = False +) -> None: import vllm.utils.flashinfer as fi_utils + dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} for num_tokens in _flashinfer_autotune_token_counts(runner): tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) logger.info( @@ -369,6 +373,7 @@ def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None: skip_eplb=True, is_profile=True, randomize_inputs=True, + **dummy_run_kwargs, ) @@ -431,7 +436,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: torch.inference_mode(), fi_utils.autotune(tune_mode=True, **autotune_kwargs), ): - _run_flashinfer_autotune_dummy_runs(runner) + hisparse_enabled = ( + runner.vllm_config.attention_config.hisparse_config is not None + ) + if hisparse_enabled: + # HiSparse hot-buffer attention is bounded by decode batch + # size, not the prefill-sized batch used for the full model. + autotune_hisparse_flashinfer_attention(runner) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 6a0d207ff568..b0fe554e7e10 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -42,6 +42,9 @@ if TYPE_CHECKING: from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) class DeepseekV32Indexer(nn.Module): @@ -126,6 +129,7 @@ def __init__( prefix: str, topk_indices_buffer: torch.Tensor | None = None, attn_backend: "type | None" = None, + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None, ) -> None: quant_config = vllm_config.quant_config cache_config = vllm_config.cache_config @@ -209,6 +213,7 @@ def __init__( use_sparse=True, indexer=indexer, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, attn_backend=attn_backend, ) @@ -304,6 +309,14 @@ def forward( # type: ignore[override] slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) + indexer_slot = ( + slot_mapping.get(self.indexer.k_cache.prefix) + if self.indexer is not None + else None + ) + hisparse_cache = self.hisparse_cache + layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name) + self.impl.prepare_for_batch(layer_attn_metadata) if self.indexer is not None and not self.skip_topk: has_indexer = True @@ -331,12 +344,17 @@ def forward( # type: ignore[override] mla_k_scale = None indexer_k_cache = None mla_slot = None + indexer_slot = None else: - mla_kv_cache = self.kv_cache + mla_kv_cache = None if hisparse_cache is not None else self.kv_cache mla_k_scale = self._k_scale - kv_c_out = torch.empty_like(kv_c) - k_pe_out = torch.empty_like(k_pe) + if self.use_pcp or hisparse_cache is None: + kv_c_out = torch.empty_like(kv_c) + k_pe_out = torch.empty_like(k_pe) + else: + kv_c_out = torch.empty_like(kv_c) + k_pe_out = torch.empty_like(k_pe) q_c = fused_norm_rope( positions, q_c, @@ -354,6 +372,7 @@ def forward( # type: ignore[override] indexer_k_rope_cos_sin_cache, self.topk_indices_buffer, slot_mapping=mla_slot, + indexer_slot_mapping=indexer_slot, indexer_k_cache=indexer_k_cache, mla_kv_cache=mla_kv_cache, mla_kv_cache_dtype=self.kv_cache_dtype, @@ -365,6 +384,18 @@ def forward( # type: ignore[override] index_k_out=index_k_out, ) + if hisparse_cache is not None and mla_slot is not None: + assert kv_c_out is not None and k_pe_out is not None + self.update_kv_cache( + kv_c_out, + k_pe_out, + self.kv_cache, + mla_slot, + layer_attn_metadata, + self.kv_cache_dtype, + self._k_scale, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) @@ -459,6 +490,7 @@ def _sparse_indexer_and_attn( ), skip_topk_buffer_clear=True, ) + self.impl.record_logical_topk_ready() # type: ignore[attr-defined] attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context( self.layer_name @@ -467,6 +499,7 @@ def _sparse_indexer_and_attn( output.zero_() return attn_metadata = cast("MLACommonMetadata", attn_metadata) + self.impl.prepare_for_batch(attn_metadata) if self.use_pcp: assert kv_c is not None and k_pe is not None diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 1c91d3015774..ef7fcde5d341 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -68,7 +68,7 @@ def _fp8_quant_and_cache_write( kv_cache_ptr, kv_cache_scale_ptr, cache_block_size, - cache_stride, + cache_block_stride, offsets, HEAD_DIM: tl.constexpr, ): @@ -76,7 +76,7 @@ def _fp8_quant_and_cache_write( block_idx = slot_idx // cache_block_size block_offset = slot_idx % cache_block_size - block_start = block_idx * cache_block_size * cache_stride + block_start = block_idx * cache_block_stride tl.store( kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets, @@ -129,8 +129,9 @@ def _fused_norm_rope_kernel( index_k_out_ptr, index_k_out_stride, INDEX_K_HALF_ROT_DIM: tl.constexpr, - # Cache params (shared by indexer K and MLA) + # Cache params slot_mapping_ptr, + indexer_slot_mapping_ptr, # Index K FP8 cache indexer_cache_ptr, indexer_cache_scale_ptr, @@ -400,8 +401,8 @@ def _fused_norm_rope_kernel( ) # PCP inserts index K after gathering; other paths write it directly. - if indexer_cache_ptr is not None and slot_mapping_ptr is not None: - slot_idx = tl.load(slot_mapping_ptr + tok_idx) + if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None: + slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx) _fp8_quant_and_cache_write( result, index_k_mask, @@ -431,8 +432,9 @@ def fused_norm_rope( index_k_layer_norm_eps: float, index_k_rope_cos_sin_cache: torch.Tensor | None, topk_indices_buffer: torch.Tensor, - # Cache params for fused writes (single slot_mapping for both caches) + # Cache params for fused writes slot_mapping: torch.Tensor | None = None, + indexer_slot_mapping: torch.Tensor | None = None, indexer_k_cache: torch.Tensor | None = None, mla_kv_cache: torch.Tensor | None = None, mla_kv_cache_dtype: str = "auto", @@ -476,9 +478,11 @@ def fused_norm_rope( # --- Indexer K cache setup --- if indexer_k_cache is not None: assert slot_mapping is not None + if indexer_slot_mapping is None: + indexer_slot_mapping = slot_mapping idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) idx_cache_block_size = indexer_k_cache.shape[1] - idx_cache_stride = indexer_k_cache.shape[2] + idx_cache_stride = indexer_k_cache.stride(0) if indexer_k_cache.dtype == torch.uint8: indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn) else: @@ -578,6 +582,7 @@ def fused_norm_rope( index_k_rope_cos_sin_cache.shape[-1] // 2, # Cache params slot_mapping, + indexer_slot_mapping, indexer_k_cache, idx_cache_scale_view, idx_cache_block_size, diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 92b74e536fc6..b42ba28c8df8 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -46,6 +46,10 @@ ) from vllm.models.deepseek_v32.attention import DeepseekV32Attention from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm @@ -57,6 +61,7 @@ def __init__( prefix: str, config=None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -80,6 +85,7 @@ def __init__( config=config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) if ( @@ -189,6 +195,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dtype=torch.int32, device=self.device, ) + index_group_builder = SparseMLAIndexGroupBuilder( + self.topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) if get_pp_group().is_first_rank: self.embed_tokens = make_input_embedding( @@ -209,6 +219,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=prefix, topk_indices_buffer=self.topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index a5348ed2204a..bdc6b95c15e4 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -48,6 +48,10 @@ from vllm.models.deepseek_v32.common.kernels import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm52_low_latency_gemm import ( build_glm52_plan, @@ -80,6 +84,10 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: dtype=torch.int32, device=current_platform.device_type, ) + index_group_builder = SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config ) @@ -88,6 +96,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: prefix, config=config, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) def forward( diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index f5a8db0c2443..ac02bca1fcc8 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -55,6 +55,7 @@ AttentionSpec, KVCacheSpec, MLAAttentionSpec, + SparseCacheRole, ) logger = init_logger(__name__) @@ -148,6 +149,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, + cache_role=SparseCacheRole.INDEXER, ) def forward(self) -> None: ... diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d75c79fe0f70..f4031a7ee01e 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms.interface import DeviceCapability + from vllm.v1.hisparse.runtime import HiSparseCacheHandle from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheLayout, @@ -838,6 +839,9 @@ def __new__(cls, *args, **kwargs): def process_weights_after_loading(self, act_dtype: torch.dtype): pass + def prepare_for_batch(self, attn_metadata: T | None) -> None: + """Prepare implementation-specific state for the current batch.""" + class AttentionImpl(AttentionImplBase[T], Generic[T]): """Standard attention implementation with forward method.""" @@ -956,6 +960,7 @@ def do_rope_and_kv_cache_update( class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): """MLA attention implementation with forward_mqa and forward_mha methods.""" + hisparse_cache: "HiSparseCacheHandle | None" = None supports_pcp: bool = True @abstractmethod diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 2882a81b1152..43720fb3ed02 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -8,9 +8,9 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) from vllm.platforms.interface import DeviceCapability @@ -18,11 +18,11 @@ AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, MLAAttentionImpl, MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -106,26 +106,8 @@ def supports_combination( @dataclass -class FlashAttnMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - seq_lens: torch.Tensor - block_size: int = 64 - topk_tokens: int = 2048 - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - cp_kv_cache_interleave_size: int = 1 +class FlashAttnMLASparseMetadata(SparseMLACommonMetadata): + pass class FlashAttnMLASparseMetadataBuilder( @@ -214,6 +196,76 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + outputs = [] + if num_decode_tokens: + physical_topk, valid_counts = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + outputs.append( + self._run_mqa_kernel( + q_nope[:num_decode_tokens], + q_rope[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + attn_metadata.block_size, + ) + ) + if num_decode_tokens < num_actual_toks: + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = ( + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + ) + prefill_cache = index_group.physical_kv_cache( + self.index_group_index + ).view(kv_c_and_k_pe_cache.dtype) + else: + prefill_cache, block_table, req_ids = ( + index_group.stage_prefill_rows( + self.index_group_index, + kv_c_and_k_pe_cache, + attn_metadata, + ) + ) + physical_topk, valid_counts = ( + triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) + ) + outputs.append( + self._run_mqa_kernel( + q_nope[num_decode_tokens:], + q_rope[num_decode_tokens:], + prefill_cache, + physical_topk, + valid_counts, + attn_metadata.block_size, + ) + ) + return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None + kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -226,9 +278,36 @@ def forward_mqa( NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) + return ( + self._run_mqa_kernel( + q_nope, + q_rope, + kv_rows, + topk_indices, + valid_counts, + attn_metadata.block_size, + cache_is_flat=True, + ), + None, + ) + + def _run_mqa_kernel( + self, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + block_size: int, + *, + cache_is_flat: bool = False, + ) -> torch.Tensor: + kv_rows = ( + kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] + ) cu_seqlens_q = torch.arange( - 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device + 0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device ) k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1) v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) @@ -247,4 +326,4 @@ def forward_mqa( causal=True, fa_version=3, ) - return out, None + return out diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index b34c98375261..ae2e51284015 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -11,22 +11,23 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, AttentionType, MLAAttentionImpl, MultipleOf, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -224,33 +225,9 @@ def supports_combination( @dataclass -class FlashInferMLASparseMetadata(AttentionMetadata): +class FlashInferMLASparseMetadata(SparseMLACommonMetadata): """Attention metadata for FlashInfer MLA Sparse backend.""" - num_reqs: int - max_query_len: int - max_seq_len: int - num_actual_tokens: int - - # Query start locations - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - block_table: torch.Tensor - req_id_per_token: torch.Tensor - - # Sequence lengths for all requests (context + query) - seq_lens: torch.Tensor - num_decodes: int - num_prefills: int - num_decode_tokens: int - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - - # Sparse-specific - block_size: int = 64 - topk_tokens: int = 2048 - cp_kv_cache_interleave_size: int = 1 - class FlashInferMLASparseMetadataBuilder( SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata] @@ -286,6 +263,7 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde """Metadata builder for the SM100 TRT-LLM sparse MLA kernel.""" _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + hisparse_supports_multi_token_decode: ClassVar[bool] = True def _build_req_id_per_token( self, @@ -402,6 +380,78 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + self._prepare_mqa_kernel(layer, q.device) + + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + decode_out: torch.Tensor | None = None + decode_lse: torch.Tensor | None = None + if num_decode_tokens > 0: + physical_topk, valid_counts = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + decode_out, decode_lse = self._run_mqa_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + ) + if num_decode_tokens == num_actual_toks: + return decode_out, decode_lse + + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = ( + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + ) + return self._run_mqa_kernel( + q, + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + ) + + prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( + self.index_group_index, kv_c_and_k_pe_cache, attn_metadata + ) + prefill_indices, prefill_lens = triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) + prefill_out, prefill_lse = self._run_mqa_kernel( + q[num_decode_tokens:], + prefill_cache, + prefill_indices, + prefill_lens, + ) + if decode_out is None: + return prefill_out, prefill_lse + output = torch.cat((decode_out, prefill_out)) + if decode_lse is None: + return output, None + assert prefill_lse is not None + return output, torch.cat((decode_lse, prefill_lse)) + _, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -420,18 +470,27 @@ def forward_mqa( return_valid_counts=True, ) else: - topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, + topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk( topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], + attn_metadata, + block_stride_rows=block_stride_rows, return_valid_counts=True, ) + return self._run_mqa_kernel( + q, + kv_c_and_k_pe_cache, + topk_indices_physical, + seq_lens, + ) + + def _prepare_mqa_kernel( + self, + layer: AttentionLayer, + device: torch.device, + ) -> None: if self._workspace_buffer is None: - self._workspace_buffer = _get_workspace_buffer(q.device) + self._workspace_buffer = _get_workspace_buffer(device) if self.bmm1_scale is None: self.bmm1_scale = self.scale @@ -442,15 +501,71 @@ def forward_mqa( if is_quantized_kv_cache(self.kv_cache_dtype): self.bmm2_scale *= layer._k_scale_float + def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: + """Autotune the largest legal HiSparse decode batch.""" + assert isinstance(self.index_group, HiSparseMLAIndexGroup) + cache = self.index_group.cache(self.index_group_index) + assert self.topk_indices_buffer is not None + + runtime = cache.runtime + kv_cache = runtime.hot.attention_cache + num_tokens = runtime.max_num_reqs + topk_tokens = self.topk_indices_buffer.shape[1] + self._prepare_mqa_kernel(layer, kv_cache.device) + + q_dtype = ( + current_platform.fp8_dtype() + if is_quantized_kv_cache(self.kv_cache_dtype) + else kv_cache.dtype + ) + q = torch.zeros( + ( + num_tokens, + self.num_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + ), + dtype=q_dtype, + device=kv_cache.device, + ) + topk_indices = ( + torch.arange( + topk_tokens, + dtype=torch.int32, + device=kv_cache.device, + ) + .expand(num_tokens, -1) + .contiguous() + ) + seq_lens = torch.full( + (num_tokens,), + topk_tokens, + dtype=torch.int32, + device=kv_cache.device, + ) + self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens) + + def _run_mqa_kernel( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + seq_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self._workspace_buffer is not None + assert self.bmm1_scale is not None + assert self.bmm2_scale is not None + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + kv_cache = kv_cache.view(q.dtype) + # Single-token sparse decode. trtllm-gen requires the q_len_per_request # dim, but the sparse attention mask is fully per-token (each query token # carries its own top-k index row), so unsqueeze is sufficient and # correct. The MTP/multi-token q_len grouping is a perf-only layout and is # deferred until MTP is validated end-to-end for this backend. query = q.unsqueeze(1) - block_tables = topk_indices_physical.unsqueeze(1) + block_tables = topk_indices.unsqueeze(1) seq_lens_arg = seq_lens # page_table width = topk buffer width, which kpool widens past @@ -460,7 +575,7 @@ def forward_mqa( # count), so the -1 padding slots past seq_lens are never attended to. # Use the actual buffer width instead of the fixed topk_tokens, which # mismatches the page_table when index_kpool > 1. - sparse_topk_capacity = topk_indices_physical.shape[1] + sparse_topk_capacity = topk_indices.shape[1] extra_kwargs: dict[str, torch.Tensor] = {} empty_rows: torch.Tensor | None = None @@ -473,14 +588,12 @@ def forward_mqa( # valid indices into a contiguous prefix, which is what the kernel # requires of the page table. empty_rows = seq_lens == 0 - topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill( - empty_rows, 0 - ) + topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0) extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1) kernel_out = trtllm_batch_decode_with_kv_cache_mla( query=query, - kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), + kv_cache=kv_cache.unsqueeze(1), workspace_buffer=self._workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, @@ -506,7 +619,7 @@ def forward_mqa( if lse is not None: lse = self._normalize_lse(lse, out.shape[0], out.shape[1]) if empty_rows is None and lse is not None: - empty_rows = (topk_indices_physical == -1).all(dim=-1) + empty_rows = (topk_indices == -1).all(dim=-1) if empty_rows is not None: out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) if lse is not None: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index b7459203e07f..e59c9ce46db9 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -7,17 +7,14 @@ import torch from vllm.model_executor.layers.attention.sparse_mla_attention import ( - SharedTopkIndicesBuffer, -) -from vllm.v1.attention.backend import ( - AttentionLayer, - AttentionType, - MLAAttentionImpl, + SparseMLACommonImpl, ) +from vllm.v1.attention.backend import AttentionLayer, AttentionType from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseMetadata, _get_workspace_buffer, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, ) @@ -32,9 +29,7 @@ def _kv_scale_format_for_model(model_type: str | None) -> str: return "pow2_fp32" -class FlashInferMLASparseSM120Impl( - MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer -): +class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True @@ -65,20 +60,28 @@ def __init__( "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention" ) - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads - self.kv_cache_dtype = kv_cache_dtype - if self.kv_cache_dtype != "fp8_ds_mla": + if kv_cache_dtype != "fp8_ds_mla": raise NotImplementedError( "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla " f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}." ) - self.kv_lora_rank: int = mla_args["kv_lora_rank"] - self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] - self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + topk_indices_buffer = mla_args.pop("topk_indices_buffer", None) + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + indexer=indexer, + topk_indices_buffer=topk_indices_buffer, + **mla_args, + ) from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() @@ -89,9 +92,6 @@ def __init__( ) self.kv_scale_format = _kv_scale_format_for_model(model_type) - # Skip-topk layers are built with indexer=None and get the shared - # buffer via mla_args instead (cf. FLASHMLA_SPARSE). - self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer")) from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): @@ -118,6 +118,70 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + outputs = [] + if num_decode_tokens: + topk_indices_physical = cast( + torch.Tensor, + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=False, + ), + ) + outputs.append( + self._run_mqa_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index), + topk_indices_physical, + ) + ) + if num_decode_tokens < num_actual_toks: + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + topk_indices_physical = cast( + torch.Tensor, + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ), + ) + prefill_cache = index_group.physical_kv_cache( + self.index_group_index + ) + else: + prefill_cache, block_table, req_ids = ( + index_group.stage_prefill_rows( + self.index_group_index, + kv_c_and_k_pe_cache, + attn_metadata, + ) + ) + topk_indices_physical = cast( + torch.Tensor, + triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ), + ) + outputs.append( + self._run_mqa_kernel( + q[num_decode_tokens:], + prefill_cache, + topk_indices_physical, + ) + ) + output = torch.cat(outputs) if len(outputs) > 1 else outputs[0] + return output, None topk_indices_physical = cast( torch.Tensor, @@ -129,6 +193,22 @@ def forward_mqa( NUM_TOPK_TOKENS=topk_indices.shape[1], ), ) + return ( + self._run_mqa_kernel( + q, + kv_c_and_k_pe_cache, + topk_indices_physical, + ), + None, + ) + + def _run_mqa_kernel( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices_physical: torch.Tensor, + ) -> torch.Tensor: + num_actual_toks = q.shape[0] output = q.new_empty( (num_actual_toks, self.num_heads, self.kv_lora_rank), @@ -142,20 +222,21 @@ def forward_mqa( flashinfer_trtllm_batch_decode_with_kv_cache_mla, ) + sparse_capacity = topk_indices_physical.shape[1] out = flashinfer_trtllm_batch_decode_with_kv_cache_mla( query=q.unsqueeze(1), - kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1), + kv_cache=kv_cache.view(torch.uint8).unsqueeze(1), workspace_buffer=self._workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=topk_indices_physical.unsqueeze(1), seq_lens=None, - max_seq_len=attn_metadata.topk_tokens, + max_seq_len=sparse_capacity, out=output.unsqueeze(1), bmm1_scale=self.scale, bmm2_scale=1.0, - sparse_mla_top_k=attn_metadata.topk_tokens, + sparse_mla_top_k=sparse_capacity, kv_scale_format=self.kv_scale_format, ) - return out.squeeze(1), None + return out.squeeze(1) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index b837cc366baf..4ecf5596832e 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -9,23 +9,23 @@ from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -172,28 +172,7 @@ def supports_combination( @dataclass -class FlashMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int # Number of tokens excluding padding. - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - block_size: int = 64 - topk_tokens: int = 2048 - - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - seq_lens: torch.Tensor | None = None - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - cp_kv_cache_interleave_size: int = 1 - +class FlashMLASparseMetadata(SparseMLACommonMetadata): @dataclass class FP8KernelMetadata: scheduler_metadata: FlashMLASchedMeta @@ -233,6 +212,7 @@ class Chunk: req_start_idx: int workspace_starts: torch.Tensor chunk_tot_seqlen: int + seq_lens: torch.Tensor | None = None chunks: list[Chunk] @@ -263,6 +243,7 @@ class FlashMLASparseMetadataBuilder( ): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH require_uniform_decodes: ClassVar[bool] = True + hisparse_supports_multi_token_decode: ClassVar[bool] = True metadata_cls = FlashMLASparseMetadata def __init__( @@ -283,6 +264,9 @@ def __init__( ) else: threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) + self.use_hisparse = vllm_config.attention_config.hisparse_config is not None + if self.use_hisparse: + threshold = 1 # Varlen decodes are safe under DCP: causality comes from the # indexer's top-k indices, not from the kernel metadata. self._init_reorder_batch_threshold( @@ -347,7 +331,9 @@ def __init__( device=device, ) - self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL + self.fp8_use_mixed_batch = ( + self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse + ) if parallel_config.decode_context_parallel_size > 1: if parallel_config.dcp_comm_backend != "ag_rs": @@ -385,37 +371,21 @@ def __init__( def _build_fp8_mixed_decode_prefill( self, - common_attn_metadata: CommonAttentionMetadata, - ) -> "FlashMLASparseMetadata.FP8KernelMetadata": + ) -> FlashMLASparseMetadata.FP8KernelMetadata: """Build FP8 metadata treating MQA tokens as one batch. The scheduler initializes lazily from the runtime query shape, which may be the full batch or only decodes when prefills use dense MHA. This avoids the BF16 prefill kernel's head-padding overhead at high TP. """ - num_tokens = common_attn_metadata.num_actual_tokens - - # Use padded head count since that's what the kernel will see - padded_heads = self.fp8_decode_padded_heads - - # Build metadata for all tokens as a single batch - scheduler_metadata, _ = get_mla_metadata( - cache_seqlens=self.topk_tokens_tensor[:1], # Single batch - num_q_tokens_per_head_k=num_tokens * padded_heads, - topk=self.topk_tokens, - num_heads_q=padded_heads, - num_heads_k=1, - is_fp8_kvcache=True, - ) - fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata, _ = get_mla_metadata() + return FlashMLASparseMetadata.FP8KernelMetadata( scheduler_metadata=scheduler_metadata, cache_lens=self.max_model_len_tensor[:1], dummy_block_table=self.dummy_block_table[:1], ) - return fp8_metadata - def _build_fp8_separate_prefill_decode( self, common_attn_metadata: CommonAttentionMetadata, @@ -520,6 +490,9 @@ def _build_fp8_separate_prefill_decode( chunk_block_table = common_attn_metadata.block_table_tensor[ num_decodes + chunk_start : num_decodes + chunk_end ] + chunk_seq_lens = common_attn_metadata.seq_lens[ + num_decodes + chunk_start : num_decodes + chunk_end + ] prefill_chunks.append( FP8Meta.Prefill.Chunk( @@ -528,6 +501,7 @@ def _build_fp8_separate_prefill_decode( req_start_idx=chunk_start, workspace_starts=chunk_workspace_starts, chunk_tot_seqlen=chunk_tot_seqlen, + seq_lens=chunk_seq_lens, ) ) @@ -569,9 +543,7 @@ def build( metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch if self.use_fp8_kv_cache: if self.fp8_use_mixed_batch: - metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill( - common_attn_metadata - ) + metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill() else: metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode( common_attn_metadata, metadata @@ -679,30 +651,83 @@ def _forward_bf16_kv( topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, actual_num_heads: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Convert per-request indices to global slots (decode) or workspace - # offsets (prefill). req_id_per_token covers the whole batch; slice it - # to the MQA tokens (q may exclude prefill tokens routed to dense MHA). + ) -> tuple[torch.Tensor, torch.Tensor | None]: + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + cache = index_group.cache(self.index_group_index) + else: + cache = None + block_table = attn_metadata.block_table + # req_id_per_token covers the whole batch; slice it to the MQA tokens + # (q may exclude prefill tokens routed to dense MHA). + req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] + decode_out: torch.Tensor | None = None + if cache is not None: + assert isinstance(index_group, HiSparseMLAIndexGroup) + num_decode_tokens = attn_metadata.num_decode_tokens + if num_decode_tokens > 0: + decode_topk, decode_lengths = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + decode_out, _ = self._bf16_flash_mla_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index), + decode_topk, + decode_lengths, + actual_num_heads, + ) + if num_decode_tokens == q.shape[0]: + return decode_out, None + q = q[num_decode_tokens:] + topk_indices = topk_indices[num_decode_tokens:] + kv_c_and_k_pe_cache, block_table, req_id_per_token = ( + index_group.stage_prefill_rows( + self.index_group_index, kv_c_and_k_pe_cache, attn_metadata + ) + ) + # Convert per-request indices to global slots (decode) or workspace offsets. kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) - topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, + decode_only = ( + attn_metadata.num_decode_tokens + == attn_metadata.num_actual_tokens + == topk_indices.shape[0] ) + uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) + if not uses_host_cache and decode_only: + topk_indices, topk_length = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + else: + topk_indices, topk_length = triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) - return self._bf16_flash_mla_kernel( + attn_out, lse = self._bf16_flash_mla_kernel( q, kv_rows, topk_indices, topk_length, actual_num_heads, ) + if decode_out is None: + return attn_out, lse + return torch.cat([decode_out, attn_out], dim=0), None def _forward_fp8_kv_separate_prefill_decode( self, @@ -721,6 +746,28 @@ def _forward_fp8_kv_separate_prefill_decode( "FP8 sparse MLA expects either the decode subset or the full batch" ) + decode_topk: torch.Tensor | None = None + index_group = self.index_group + uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) + if uses_host_cache and num_decode_tokens > 0: + decode_topk = topk_indices[:num_decode_tokens] + + prefill_ready = None + if num_prefill_tokens > 0 and uses_host_cache: + assert fp8_metadata.prefill is not None + first_chunk = fp8_metadata.prefill.chunks[0] + assert isinstance(index_group, HiSparseMLAIndexGroup) + prefill_ready = index_group.gather_fp8_prefill( + self.index_group_index, + kv_c_and_k_pe_cache, + self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen], + first_chunk.block_table, + first_chunk.workspace_starts, + len(first_chunk.block_table), + attn_metadata, + first_chunk.req_start_idx, + ) + prefill_request_ids = None prefill_workspace_starts = None has_prefill_workspace = False @@ -736,17 +783,26 @@ def _forward_fp8_kv_separate_prefill_decode( # For BF16 cache: always use global cache slots (no workspace) # prefill_workspace_starts has been adjusted in-place per chunk so # prefill indices automatically come out chunk-local - topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - HAS_PREFILL_WORKSPACE=has_prefill_workspace, - prefill_workspace_request_ids=prefill_request_ids, - prefill_workspace_starts=prefill_workspace_starts, - return_valid_counts=True, - ) + topk_length = None + if num_prefill_tokens == 0 and not uses_host_cache: + topk_indices, topk_length = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + elif num_prefill_tokens > 0: + topk_indices, topk_length = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[: topk_indices.shape[0]], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + HAS_PREFILL_WORKSPACE=has_prefill_workspace, + prefill_workspace_request_ids=prefill_request_ids, + prefill_workspace_starts=prefill_workspace_starts, + return_valid_counts=True, + ) fp8_metadata = attn_metadata.fp8_extra_metadata assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) @@ -755,6 +811,16 @@ def _fp8_decode( q: torch.Tensor, topk_indices: torch.Tensor, ) -> torch.Tensor: + assert fp8_metadata.decode is not None + if uses_host_cache: + return self._host_backed_fp8_decode( + q, + topk_indices, + attn_metadata, + fp8_metadata.decode.kernel_metadata, + num_decodes, + fp8_metadata.decode.decode_query_len, + ) # Reshape q: (num_decode_tokens, num_heads, head_dim) # -> (num_decodes, seq_len, num_heads, head_dim) q = reshape_query_for_spec_decode(q, num_decodes) @@ -762,7 +828,6 @@ def _fp8_decode( # Reshape topk_indices: (num_decode_tokens, topk) # -> (num_decodes, seq_len, topk) topk_indices = topk_indices.view(num_decodes, seq_len, -1) - assert fp8_metadata.decode is not None attn_out, _ = self._fp8_flash_mla_kernel( q=q, kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, @@ -776,7 +841,9 @@ def _fp8_decode( # Pure decode: direct call without allocation if num_decode_tokens > 0 and num_prefill_tokens == 0: assert fp8_metadata.decode is not None - attn_out = _fp8_decode(q, topk_indices) + attn_out = _fp8_decode( + q, decode_topk if decode_topk is not None else topk_indices + ) else: # Mixed or pure prefill: allocate output tensor attn_out = q.new_empty( @@ -788,13 +855,30 @@ def _fp8_decode( if num_decode_tokens > 0: attn_out[:num_decode_tokens] = _fp8_decode( q[:num_decode_tokens], - topk_indices[:num_decode_tokens], + decode_topk + if decode_topk is not None + else topk_indices[:num_decode_tokens], ) assert fp8_metadata.prefill is not None - for chunk in fp8_metadata.prefill.chunks: + for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks): chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen] - if self.kv_cache_dtype == "fp8_ds_mla": + if uses_host_cache and chunk_index > 0: + assert isinstance(index_group, HiSparseMLAIndexGroup) + prefill_ready = index_group.gather_fp8_prefill( + self.index_group_index, + kv_c_and_k_pe_cache, + chunk_workspace, + chunk.block_table, + chunk.workspace_starts, + len(chunk.block_table), + attn_metadata, + chunk.req_start_idx, + ) + if uses_host_cache: + assert prefill_ready is not None + current_stream().wait_event(prefill_ready) + elif self.kv_cache_dtype == "fp8_ds_mla": ops.cp_gather_and_upconvert_fp8_kv_cache( kv_c_and_k_pe_cache, chunk_workspace, @@ -813,6 +897,7 @@ def _fp8_decode( chunk_q = q[chunk.tokens_slice] chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] + assert topk_length is not None chunk_topk_length = topk_length[chunk.tokens_slice] attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( @@ -839,6 +924,17 @@ def _forward_fp8_kv_mixed_batch( The lse is only returned when DCP needs it, otherwise None. """ + assert attn_metadata.fp8_extra_metadata is not None + assert isinstance( + attn_metadata.fp8_extra_metadata, + FlashMLASparseMetadata.FP8KernelMetadata, + ) + fp8_metadata = attn_metadata.fp8_extra_metadata + + block_table = attn_metadata.block_table + # req_id_per_token covers the whole batch; slice it to the MQA tokens + # (q may exclude prefill tokens routed to dense MHA). + req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] if self.dcp_world_size > 1: # The indexer emits global token ids; keep this rank's shard and # convert to local slots. compact_valid_to_front=False keeps the @@ -846,8 +942,8 @@ def _forward_fp8_kv_mixed_batch( # empty-row neutralization below relies on. req_id is sliced to # topk_indices rows (the converter grids from req_id). topk_indices = triton_filter_and_convert_dcp_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, + req_id_per_token, + block_table, topk_indices, dcp_size=self.dcp_world_size, dcp_rank=self.dcp_rank, @@ -859,19 +955,26 @@ def _forward_fp8_kv_mixed_batch( else: # Convert per-request indices to global slots (decode) or workspace # offsets (prefill). - topk_indices = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], + decode_only = ( + attn_metadata.num_decode_tokens + == attn_metadata.num_actual_tokens + == topk_indices.shape[0] ) - - assert attn_metadata.fp8_extra_metadata is not None - assert isinstance( - attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata - ) - fp8_metadata = attn_metadata.fp8_extra_metadata + if decode_only: + topk_indices = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + else: + topk_indices = triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ) _attn_out, _lse = self._fp8_flash_mla_kernel( q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) @@ -937,6 +1040,37 @@ def _fp8_flash_mla_kernel( return out, lse + def _host_backed_fp8_decode( + self, + q: torch.Tensor, + topk_indices: torch.Tensor, + attn_metadata: FlashMLASparseMetadata, + kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata, + num_decodes: int, + decode_query_len: int, + ) -> torch.Tensor: + assert isinstance(self.index_group, HiSparseMLAIndexGroup) + physical_topk = self.index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + return_valid_counts=False, + num_decodes=num_decodes, + decode_query_len=decode_query_len, + ) + assert isinstance(physical_topk, torch.Tensor) + q = reshape_query_for_spec_decode(q, num_decodes) + physical_topk = physical_topk.view(num_decodes, q.shape[1], -1) + output, _ = self._fp8_flash_mla_kernel( + q=q, + kv_c_and_k_pe_cache=self.index_group.physical_kv_cache( + self.index_group_index + ), + topk_indices=physical_topk, + kernel_metadata=kernel_metadata, + ) + return reshape_attn_output_for_spec_decode(output) + def _bf16_flash_mla_kernel( self, q: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py new file mode 100644 index 000000000000..8ba0171835fd --- /dev/null +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -0,0 +1,503 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig +from vllm.utils.torch_utils import current_stream +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) +from vllm.v1.hisparse.runtime import ( + FP8_DS_MLA_ROW_BYTES, + HiSparseCacheHandle, + HiSparsePrefillStagingPlan, + create_hisparse_cache_handle, +) + + +def _create_side_stream(device: torch.device) -> torch.Stream: + return torch.Stream(device=device) + + +def _create_event() -> torch.Event: + return torch.Event() + + +@dataclass +class SparseMLAIndexGroup: + """Layers that consume one sparse-indexer result.""" + + logical_topk_indices: torch.Tensor + physical_topk_indices: torch.Tensor + valid_topk_counts: torch.Tensor + request_ids: torch.Tensor + side_stream: torch.Stream + logical_topk_ready: torch.Event + physical_topk_ready: torch.Event + has_indexer: bool + num_layers: int = 0 + + def register_layer( + self, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> int: + layer_index = self.num_layers + self.num_layers += 1 + return layer_index + + def set_logical_topk_ready(self, layer_index: int) -> None: + if layer_index == 0 and self.has_indexer: + self.logical_topk_ready.record(current_stream()) + + def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: + pass + + def gather_fp8_prefill( + self, + layer_index: int, + source_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, + attn_metadata: Any, + request_start: int, + ) -> torch.Event | None: + ops.cp_gather_and_upconvert_fp8_kv_cache( + source_cache, + dst, + block_table, + workspace_starts, + batch_size, + ) + return None + + def convert_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Convert one index result once for all layers in this group.""" + return self._convert_once( + layer_index, + logical_topk_indices, + attn_metadata.req_id_per_token[: logical_topk_indices.shape[0]], + attn_metadata.block_table, + attn_metadata.block_size, + block_stride_rows=block_stride_rows, + return_valid_counts=return_valid_counts, + ) + + def _convert_once( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + num_tokens = logical_topk_indices.shape[0] + if num_tokens > self.physical_topk_indices.shape[0]: + return triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=logical_topk_indices.shape[1], + return_valid_counts=return_valid_counts, + ) + + physical_topk_indices = self.physical_topk_indices[:num_tokens] + valid_topk_counts = self.valid_topk_counts[:num_tokens] + compute_stream = current_stream() + if layer_index == 0: + if self.has_indexer: + self.side_stream.wait_event(self.logical_topk_ready) + else: + self.side_stream.wait_stream(compute_stream) + with self.side_stream: + triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=logical_topk_indices.shape[1], + return_valid_counts=return_valid_counts, + out=physical_topk_indices, + valid_counts_out=( + valid_topk_counts if return_valid_counts else None + ), + ) + self.physical_topk_ready.record(self.side_stream) + compute_stream.wait_event(self.physical_topk_ready) + + if return_valid_counts: + return physical_topk_indices, valid_topk_counts + return physical_topk_indices + + +@dataclass +class HiSparseMLAIndexGroup(SparseMLAIndexGroup): + """Sparse index group with per-layer host-backed KV caches.""" + + caches: list[HiSparseCacheHandle] = field(default_factory=list) + hisparse_group: Any | None = None + prefill_stream: torch.Stream | None = None + prefill_ready_events: list[torch.Event] = field(default_factory=list) + + def register_layer( + self, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> int: + assert vllm_config is not None + assert head_size is not None + assert kv_cache_dtype is not None + layer_index = super().register_layer() + if self.prefill_stream is None: + self.prefill_stream = _create_side_stream(self.logical_topk_indices.device) + if kv_cache_dtype == "fp8_ds_mla": + row_width = FP8_DS_MLA_ROW_BYTES + kv_dtype = torch.uint8 + else: + from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype + + row_width = head_size + kv_dtype = kv_cache_dtype_str_to_dtype( + kv_cache_dtype, vllm_config.model_config + ) + cache = create_hisparse_cache_handle( + vllm_config, + self.logical_topk_indices.shape[1], + is_index_group_leader=layer_index == 0, + row_width=row_width, + kv_dtype=kv_dtype, + index_group=self, + ) + assert cache is not None + self.caches.append(cache) + cache.index_group_caches = self.caches + self.prefill_ready_events.append(_create_event()) + return layer_index + + def cache(self, layer_index: int) -> HiSparseCacheHandle: + return self.caches[layer_index] + + def physical_kv_cache(self, layer_index: int) -> torch.Tensor: + cache = self.cache(layer_index) + return cache.runtime.hot.attention_cache + + def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: + if layer_index == 0: + self.cache(0).prepare_group_for_batch(attn_metadata) + + def convert_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + req_id_per_token: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + cache = self.cache(layer_index) + num_tokens = logical_topk_indices.shape[0] + if req_id_per_token is None: + req_id_per_token = attn_metadata.req_id_per_token[:num_tokens] + if num_tokens > self.physical_topk_indices.shape[0]: + # Prefill-sized batches do not fit the decode residency workspace. + # Non-resident prefills are staged before reaching this path. + assert cache.all_context_pages_resident + leader = self.cache(0) + assert leader.view is not None and leader.block_table is not None + return self._convert_once( + layer_index, + logical_topk_indices, + req_id_per_token, + leader.block_table, + leader.view.block_size, + block_stride_rows=leader.view.attention_block_stride, + return_valid_counts=return_valid_counts, + ) + source_block_table = cache.source_block_table + assert source_block_table is not None + return cache.swap_in( + req_id_per_token, + block_table=source_block_table, + logical_topk_indices=logical_topk_indices, + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + ) + + def convert_decode_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + return_valid_counts: bool, + num_decodes: int | None = None, + decode_query_len: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + num_tokens = logical_topk_indices.shape[0] + if num_decodes is None: + num_decodes = attn_metadata.num_decodes + if decode_query_len is None: + decode_query_len = attn_metadata.decode_max_query_len + if decode_query_len == 1: + return self.convert_logical_to_physical_topk( + layer_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=return_valid_counts, + req_id_per_token=self.request_ids[:num_decodes], + ) + + if num_tokens != num_decodes * decode_query_len: + query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1] + physical_topk_indices = self.physical_topk_indices[: num_tokens + 1] + valid_topk_counts = self.valid_topk_counts[: num_tokens + 1] + request_ids = self.request_ids[:num_decodes] + cache = self.cache(layer_index) + source_block_table = cache.source_block_table + assert source_block_table is not None + for step in range(decode_query_len): + token_indices = query_start_loc[:-1] + step + active = (token_indices < query_start_loc[1:]) & ( + token_indices < num_tokens + ) + output_indices = torch.where( + active, token_indices, torch.full_like(token_indices, num_tokens) + ).long() + token_indices = token_indices.clamp(0, num_tokens - 1).long() + step_topk = logical_topk_indices.index_select( + 0, token_indices + ).masked_fill(~active.unsqueeze(1), -1) + step_result = cache.swap_in( + request_ids, + block_table=source_block_table, + logical_topk_indices=step_topk, + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + ) + if layer_index == 0: + if return_valid_counts: + step_indices, step_counts = step_result + valid_topk_counts.index_copy_(0, output_indices, step_counts) + else: + step_indices = step_result + physical_topk_indices.index_copy_(0, output_indices, step_indices) + physical_topk_indices = physical_topk_indices[:num_tokens] + if return_valid_counts: + return physical_topk_indices, valid_topk_counts[:num_tokens] + return physical_topk_indices + + assert num_tokens == num_decodes * decode_query_len + logical_topk_by_request = logical_topk_indices.view( + num_decodes, decode_query_len, -1 + ) + physical_topk_by_request = self.physical_topk_indices[:num_tokens].view( + num_decodes, decode_query_len, -1 + ) + valid_topk_by_request = self.valid_topk_counts[:num_tokens].view( + num_decodes, decode_query_len + ) + request_ids = self.request_ids[:num_decodes] + for step in range(decode_query_len): + cache = self.cache(layer_index) + source_block_table = cache.source_block_table + assert source_block_table is not None + cache.swap_in( + request_ids, + block_table=source_block_table, + logical_topk_indices=logical_topk_by_request[:, step], + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + attention_indices_out=( + physical_topk_by_request[:, step] if layer_index == 0 else None + ), + valid_counts_out=( + valid_topk_by_request[:, step] + if layer_index == 0 and return_valid_counts + else None + ), + ) + physical_topk_indices = physical_topk_by_request.view(num_tokens, -1) + if return_valid_counts: + return physical_topk_indices, valid_topk_by_request.view(num_tokens) + return physical_topk_indices + + def stage_prefill_rows( + self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cache = self.cache(layer_index) + prefill = attn_metadata.prefill + staging_plan = prefill.host_staging_plan if prefill is not None else None + assert staging_plan is not None + resident_cache = None + if cache.view is not None and cache.block_table is not None: + staging_plan.ensure_gpu_sources( + cache.block_table[attn_metadata.num_decodes :], + cache.view.block_size, + ) + resident_cache = cache.view.cache + staged_cache = cache.runtime.gather_prefill_cache( + kv_cache, + staging_plan, + resident_cache=resident_cache, + ) + req_ids = attn_metadata.req_id_per_token[attn_metadata.num_decode_tokens :] + if attn_metadata.num_decodes > 0: + req_ids = req_ids - attn_metadata.num_decodes + return staged_cache, staging_plan.block_table, req_ids + + def _prefill_gather_plan( + self, layer_index: int, attn_metadata: Any + ) -> HiSparsePrefillStagingPlan: + cache = self.cache(layer_index) + prefill = attn_metadata.prefill + plan = prefill.host_staging_plan if prefill is not None else None + assert plan is not None + assert cache.view is not None and cache.block_table is not None + plan.ensure_gpu_sources( + cache.block_table[attn_metadata.num_decodes :], + cache.view.block_size, + ) + assert plan.gpu_row_ids is not None + return plan + + def gather_fp8_prefill( + self, + layer_index: int, + source_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, + attn_metadata: Any, + request_start: int, + ) -> torch.Event | None: + cache = self.cache(layer_index) + plan = self._prefill_gather_plan(layer_index, attn_metadata) + assert cache.view is not None and plan.gpu_row_ids is not None + row_width = source_cache.shape[-1] + prefill_stream = self.prefill_stream + assert prefill_stream is not None + block_table = plan.block_table[request_start : request_start + batch_size] + prefill_stream.wait_stream(current_stream()) + with prefill_stream: + ops.cp_gather_and_upconvert_fp8_kv_cache( + cache.view.cache, + dst, + block_table, + workspace_starts, + batch_size, + host_cache=source_cache.view(-1, row_width), + host_row_ids=plan.row_ids, + device_row_ids=plan.gpu_row_ids, + ) + ready = self.prefill_ready_events[layer_index] + ready.record(prefill_stream) + return ready + + +class SparseMLAIndexGroupBuilder: + """Assign sparse MLA layers to their index-producing layer.""" + + def __init__( + self, logical_topk_indices: torch.Tensor, max_decode_rows: int | None = None + ) -> None: + self.logical_topk_indices = logical_topk_indices + self.max_decode_rows = ( + logical_topk_indices.shape[0] + if max_decode_rows is None + else max_decode_rows + ) + self.current_group: SparseMLAIndexGroup | None = None + + def register_layer( + self, + is_index_producing_layer: bool, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> tuple[SparseMLAIndexGroup, int]: + if is_index_producing_layer or self.current_group is None: + group_cls = ( + HiSparseMLAIndexGroup + if vllm_config is not None + and vllm_config.attention_config.hisparse_config is not None + else SparseMLAIndexGroup + ) + workspace_rows = self.max_decode_rows + physical_topk_indices = torch.empty( + (workspace_rows + 1, self.logical_topk_indices.shape[1]), + dtype=self.logical_topk_indices.dtype, + device=self.logical_topk_indices.device, + ) + self.current_group = group_cls( + logical_topk_indices=self.logical_topk_indices, + physical_topk_indices=physical_topk_indices, + valid_topk_counts=torch.empty( + workspace_rows + 1, + dtype=torch.int32, + device=self.logical_topk_indices.device, + ), + request_ids=torch.arange( + workspace_rows, + dtype=torch.int32, + device=self.logical_topk_indices.device, + ), + side_stream=_create_side_stream(self.logical_topk_indices.device), + logical_topk_ready=_create_event(), + physical_topk_ready=_create_event(), + has_indexer=is_index_producing_layer, + ) + group = self.current_group + group_index = group.register_layer( + vllm_config, + head_size=head_size, + kv_cache_dtype=kv_cache_dtype, + ) + return group, group_index + + +def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int: + max_query_len = 1 + speculative_config = vllm_config.speculative_config + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + max_query_len += speculative_config.num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + scheduler_config = vllm_config.scheduler_config + return min( + scheduler_config.max_num_batched_tokens, + scheduler_config.max_num_seqs * max_query_len, + ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 8d196f963064..d9b45892a671 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -407,6 +407,8 @@ def triton_convert_req_index_to_global_index( prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, return_valid_counts: bool = False, + out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ out[token_id, indice_id] = @@ -462,16 +464,36 @@ def triton_convert_req_index_to_global_index( # contiguous prefix [0, valid_count) and leaves the tail unwritten, so # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to # [:topk_length] == exactly the valid set (no dropped tokens). - if return_valid_counts: - out = torch.full_like(token_indices_c, -1) + if out is None: + out = ( + torch.full_like(token_indices_c, -1) + if return_valid_counts + else torch.empty_like(token_indices_c) + ) else: - out = torch.empty_like(token_indices_c) + assert out.dtype == token_indices_c.dtype + assert out.device == token_indices_c.device + assert out.shape == token_indices_c.shape + assert out.is_contiguous() + if return_valid_counts: + out.fill_(-1) valid_counts: torch.Tensor | None = None if return_valid_counts: - # Zero-init only matters for the atomic accumulation path. - alloc = torch.empty if single_tile else torch.zeros - valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) + if valid_counts_out is None: + # Zero-init only matters for the atomic accumulation path. + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc( + num_tokens, dtype=torch.int32, device=token_indices.device + ) + else: + assert valid_counts_out.dtype == torch.int32 + assert valid_counts_out.device == token_indices.device + assert valid_counts_out.shape == (num_tokens,) + assert valid_counts_out.is_contiguous() + valid_counts = valid_counts_out + if not single_tile: + valid_counts.zero_() # Prepare prefill pointers if HAS_PREFILL_WORKSPACE: diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 1efa0e0cea6a..e2a06e7e4e16 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from typing import Any from vllm.distributed.kv_events import ( @@ -157,6 +157,7 @@ class BlockPool: actual block size can be a multiple of hash_block_size. enable_kv_cache_events: Whether to enable kv cache events. metrics_collector: Optional metrics collector for tracking block residency. + medium: Storage medium reported in KV cache events. """ def __init__( @@ -166,14 +167,16 @@ def __init__( hash_block_size: int, enable_kv_cache_events: bool = False, metrics_collector: KVCacheMetricsCollector | None = None, + medium: str = MEDIUM_GPU, ): assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0 self.num_gpu_blocks = num_gpu_blocks + self.medium = medium self.enable_caching = enable_caching self.hash_block_size = hash_block_size # All kv-cache blocks. self.blocks: list[KVCacheBlock] = [ - KVCacheBlock(idx) for idx in range(num_gpu_blocks) + KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks) ] # Free block queue that constructs and manipulates a doubly linked # list of free blocks (including eviction candidates when caching is @@ -194,6 +197,9 @@ def __init__( self.kv_event_queue: list[KVCacheEvent] = [] self.metrics_collector = metrics_collector + # Callbacks for blocks released with ``unpin_blocks`` whose contents + # are still being read until the pool reuses them. + self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {} def get_cached_block( self, block_hash: BlockHash, kv_cache_group_ids: list[int] @@ -364,7 +370,7 @@ def _build_block_stored_event( token_ids=request.all_token_ids[start_token_idx:end_token_idx], block_size=block_size, lora_id=request.lora_request.adapter_id if request.lora_request else None, - medium=MEDIUM_GPU, + medium=self.medium, lora_name=request.lora_request.name if request.lora_request else None, extra_keys=extra_keys_list if extra_keys_list else None, group_idx=kv_cache_group_id, @@ -543,7 +549,7 @@ def cache_partial_block( lora_id=request.lora_request.adapter_id if request.lora_request else None, - medium=MEDIUM_GPU, + medium=self.medium, lora_name=request.lora_request.name if request.lora_request else None, @@ -610,7 +616,7 @@ def _emit_block_removed_events( self.kv_event_queue.append( BlockRemoved( block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))], - medium=MEDIUM_GPU, + medium=self.medium, group_idx=get_group_id(block_hash), ) ) @@ -671,6 +677,9 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks) + if self._reuse_watchers: + self._notify_reuse(ret) + # In order to only iterate the list once, we duplicated code a bit if self.enable_caching: for block in ret: @@ -687,6 +696,33 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: self.metrics_collector.on_block_allocated(block) return ret + def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None: + for block in blocks: + watcher = self._reuse_watchers.pop(block.block_id, None) + if watcher is not None: + watcher(block) + + def unpin_blocks( + self, + blocks: Iterable[KVCacheBlock], + on_reuse: Callable[[KVCacheBlock], None], + ) -> None: + """Release references to blocks that stay readable until reused. + + The blocks become last-resort eviction candidates regardless of prefix + caching: they join the tail of the free queue and count as free. + ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the + cache is reset), so the caller can stop reading it. + """ + released: list[KVCacheBlock] = [] + for block in blocks: + assert block.ref_cnt > 0 and not block.is_null + block.ref_cnt -= 1 + self._reuse_watchers[block.block_id] = on_reuse + if block.ref_cnt == 0: + released.append(block) + self.free_block_queue.append_n(released) + def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: """ If a block is cached in `cached_block_hash_to_block`, we reset its hash @@ -742,7 +778,11 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: # Identify blocks with hash (LRU cache) and without it (never match APC) blocks_to_evict_last = [] blocks_to_evict_first = [] + other_pools: dict[BlockPool, list[KVCacheBlock]] = {} for block in ordered_blocks: + if block.pool is not None and block.pool is not self: + other_pools.setdefault(block.pool, []).append(block) + continue block.ref_cnt -= 1 if block.ref_cnt == 0 and not block.is_null: if block.block_hash is None or not self.enable_caching: @@ -756,6 +796,8 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: self.free_block_queue.prepend_n(blocks_to_evict_first) # Blocks to reuse last are appended to the end of the free queue. self.free_block_queue.append_n(blocks_to_evict_last) + for pool, blocks in other_pools.items(): + pool.free_blocks(blocks) def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. @@ -797,6 +839,10 @@ def reset_prefix_cache(self) -> bool: # Remove all hashes so that no new blocks will hit. self.cached_block_hash_to_block = BlockHashToBlockMap() self.cached_block_hashes_by_block.clear() + if self._reuse_watchers: + self._notify_reuse( + [self.blocks[block_id] for block_id in list(self._reuse_watchers)] + ) # Remove all hashes from all blocks. for block in self.blocks: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 5fa881292467..afb03ac3eac9 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -139,6 +139,7 @@ def __init__( max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, block_pool=self.block_pool, + role=kv_cache_group.role, enable_caching=enable_caching, kv_cache_group_id=i, dcp_world_size=dcp_world_size_for_kv_cache_spec( @@ -179,7 +180,7 @@ def get_num_blocks_to_allocate( apply_admission_cap: bool = False, ) -> int: """ - Get the number of blocks needed to be allocated for the request. + Get the number of device blocks needed to be allocated for the request. Args: request_id: The request ID. @@ -343,6 +344,8 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """ boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: + if not manager.enable_caching: + continue # Only cache tokens with finalized KV. The last num_reprefillable_tokens # tokens can be re-prefilled during multi-module MTP. num_tokens_to_cache = max( @@ -355,6 +358,27 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: replay_boundaries=boundaries, ) + def emit_cached_block_events( + self, request: Request, computed_blocks: tuple[list[KVCacheBlock], ...] + ) -> None: + for group_idx, group_blocks in enumerate(computed_blocks): + if group_blocks: + manager = self.single_type_managers[group_idx] + group = self.kv_cache_config.kv_cache_groups[group_idx] + manager.block_pool.emit_cached_block_events( + request, + len(group_blocks), + group.kv_cache_spec.block_size, + group_idx, + ) + + def reset_prefix_cache(self) -> bool: + """Reset each manager's pool once.""" + pools = dict.fromkeys( + manager.block_pool for manager in self.single_type_managers + ) + return all([pool.reset_prefix_cache() for pool in pools]) + def free(self, request_id: str) -> None: """ Free the blocks for the request. @@ -728,10 +752,12 @@ def verify_and_split_kv_cache_groups(self) -> None: # Try to find an existing group with the same spec for idx, group in enumerate(self.attention_groups): - if group.spec == spec: - assert manager_cls is group.manager_cls, ( - "Expected same manager class for identical KV cache specs." - ) + if ( + group.spec == spec + and group.manager_cls is manager_cls + and self.single_type_managers[group.group_ids[0]].block_pool + is self.single_type_managers[i].block_pool + ): group.group_ids.append(i) if use_eagle and not group.use_eagle: self.attention_groups[idx] = group._replace(use_eagle=True) @@ -782,6 +808,8 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: + if not manager.enable_caching: + continue num_tokens_to_cache = cached_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop # it, so make that lookahead block eligible to be cached. @@ -895,7 +923,7 @@ def find_longest_cache_hit( block_hashes=block_hashes, max_length=_max_length, kv_cache_group_ids=group_ids, - block_pool=self.block_pool, + block_pool=self.single_type_managers[first_group_id].block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, alignment_tokens=self._cache_hit_alignment_tokens, @@ -929,7 +957,10 @@ def find_longest_cache_hit( for group in self.attention_groups: if not isinstance(group.spec, FullAttentionSpec): continue - group_block_size = self.single_type_managers[group.group_ids[0]].block_size + manager = self.single_type_managers[group.group_ids[0]] + if manager.retains_longer_hit: + continue + group_block_size = manager.block_size num_blocks = cdiv(hit_length, group_block_size) for group_id in group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: @@ -966,7 +997,7 @@ def find_longest_cache_hit_per_group( block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=group_ids, - block_pool=self.block_pool, + block_pool=manager.block_pool, kv_cache_spec=spec, drop_eagle_block=use_eagle, alignment_tokens=self._cache_hit_alignment_tokens, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 084ba24311c0..20e037f60ea1 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -183,6 +183,11 @@ def __init__( manager.fine_grained_prefix_cache = True self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool + self.retained_hit_group_ids = tuple( + manager.kv_cache_group_id + for manager in self.coordinator.single_type_managers + if manager.retains_longer_hit + ) self.kv_cache_config = kv_cache_config # Watermark: minimum number of KV cache blocks to keep free when @@ -287,17 +292,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int and self.enable_kv_cache_events and getattr(request, "kv_cache_report_mode", "incremental") == "full" ): - for group_idx, group_blocks in enumerate(computed_blocks): - num_blocks = len(group_blocks) - if num_blocks > 0: - group = self.kv_cache_config.kv_cache_groups[group_idx] - block_size = group.kv_cache_spec.block_size - self.block_pool.emit_cached_block_events( - request, - num_blocks, - block_size, - group_idx, - ) + self.coordinator.emit_cached_block_events(request, computed_blocks) # The junction to pin is where the lagging sparse-retention group stops # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e. @@ -463,6 +458,16 @@ def allocate_slots( if new_computed_blocks is not None: new_computed_block_list = new_computed_blocks.blocks + if self.retained_hit_group_ids: + # Only retain the prefix covered by the local and external hits. + reused = num_new_computed_tokens + num_external_computed_tokens + groups = list(new_computed_block_list) + for group_id in self.retained_hit_group_ids: + manager = self.coordinator.single_type_managers[group_id] + groups[group_id] = groups[group_id][ + : cdiv(reused, manager.block_size) + ] + new_computed_block_list = tuple(groups) else: new_computed_block_list = self.empty_kv_cache_blocks.blocks @@ -564,7 +569,7 @@ def allocate_slots( # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. - if not self.enable_caching or delay_cache_blocks: + if delay_cache_blocks: return self.create_kv_cache_blocks(new_blocks) # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens @@ -639,7 +644,7 @@ def reset_prefix_cache(self) -> bool: bool: True if the prefix cache is successfully reset, False otherwise. """ - if not self.block_pool.reset_prefix_cache(): + if not self.coordinator.reset_prefix_cache(): return False if self.log_stats: assert self.prefix_cache_stats is not None @@ -741,6 +746,8 @@ def estimate_cached_tokens(self, request: Request) -> int: self.kv_cache_config.kv_cache_groups, self.get_blocks(request.request_id).blocks, ): + if not group.kv_cache_spec.prefix_cacheable: + continue if isinstance( group.kv_cache_spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec), diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index e2e38ec5d595..996cce8c43f5 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -8,23 +8,30 @@ import os from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from functools import partial -from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.hisparse.layout import ( + get_hisparse_gpu_memory_usage, + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, @@ -44,6 +51,10 @@ from vllm.v1.request import Request from vllm.v1.utils import tensor_data +if TYPE_CHECKING: + from vllm.v1.core.block_pool import BlockPool + + # BlockHash represents the hash of a single KV-cache block used for # prefix caching. Treating it as a distinct type from `bytes` helps # catch accidental misuse when passing around raw byte strings. @@ -169,6 +180,10 @@ class KVCacheBlock: block_id: int # Reference count. ref_cnt: int = 0 + # Deferred frees may contain blocks from multiple pools. + pool: "BlockPool | None" = field( + default=None, repr=False, compare=False, kw_only=True + ) # The hash key (block hash + group id) of the block, only available # when the block is full and cached. _block_hash: BlockHashWithGroupId | None = None @@ -1057,21 +1072,32 @@ def get_max_concurrency_for_kv_cache_config( A request at max_model_len consumes whole blocks from each group's block table — cdiv(per-request bytes, page bytes) of the group's spec — and all - groups draw those block ids from one shared pool, so the per-request + device groups draw those block ids from one shared pool, so the per-request total is the sum over groups. The memory/page ratio is identical whether a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or a representative per-layer spec (scheduler config), so both capacity call sites agree. + + Host groups use a separate pool; the smaller concurrency limit applies. """ - num_blocks_per_request = sum( - cdiv( + num_blocks_per_request = 0 + host_blocks_per_request = 0 + for group in kv_cache_config.kv_cache_groups: + required = cdiv( group.kv_cache_spec.max_memory_usage_bytes(vllm_config), group.kv_cache_spec.page_size_bytes, ) - for group in kv_cache_config.kv_cache_groups - ) - max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request - return max_concurrency + if group.host_resident: + host_blocks_per_request += required + else: + num_blocks_per_request += required + limits = [kv_cache_config.num_blocks / num_blocks_per_request] + if host_blocks_per_request: + assert kv_cache_config.hisparse_host_num_blocks is not None + limits.append( + kv_cache_config.hisparse_host_num_blocks / host_blocks_per_request + ) + return min(limits) def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: @@ -1575,6 +1601,13 @@ def _get_kv_cache_bytes_per_block( for group in kv_cache_groups ) assert bytes_per_block > 0 + hot_page_sizes = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + if hot_page_sizes: + bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) return bytes_per_block @@ -1640,6 +1673,12 @@ def get_kv_cache_config_from_groups( ), ) + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( attn_group, @@ -2210,6 +2249,9 @@ def get_kv_cache_groups( # attention free models. return [] + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + if is_kv_cache_spec_uniform(kv_cache_spec): # KV cache of all layers are the same, which is true for # most models. Allocate the same amount of memory for @@ -2287,6 +2329,10 @@ def generate_scheduler_kv_cache_config( assert all( [cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs] ) + assert all( + cfg.hisparse_host_num_blocks == kv_cache_configs[0].hisparse_host_num_blocks + for cfg in kv_cache_configs + ) # All workers have the same kv_cache_config except layer names, so use # an arbitrary one to initialize the scheduler. cfg = copy.deepcopy(kv_cache_configs[0]) @@ -2347,6 +2393,9 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 + if vllm_config.attention_config.hisparse_config is not None: + return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( attn_group, @@ -2396,9 +2445,21 @@ def _estimate_max_model_len_from_groups( Returns 0 if even 1 token doesn't fit. """ original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) def fits(model_len: int) -> bool: vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + config = get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1 return ( _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) <= available_memory @@ -2629,6 +2690,9 @@ def get_kv_cache_configs( adjusted_memory.append(override * bytes_per_block) available_memory = adjusted_memory + if vllm_config.attention_config.hisparse_config is not None: + available_memory = [min(available_memory)] * len(available_memory) + # Reserve the null block BlockPool permanently holds back, so auto-fit and # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 93fc2d83c151..3a46fa40a71e 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -294,6 +294,9 @@ class SchedulerOutput: # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + # Complete block-table rows that replace incrementally appended block IDs. + block_table_updates: dict[str, tuple[list[int], ...]] | None = None + # Scheduler-local; always None by the time this reaches a worker. kv_connector_block_state: KVConnectorBlockState | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8b2a551b5579..82203738459b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -316,10 +316,9 @@ def __init__( self.cache_config.enable_mamba_fine_grained_prefix_cache ), ) - # Bind GPU block pool to the KV connector. This must happen after - # kv_cache_manager is constructed so block_pool is available. + # Bind after construction so connectors can access the cache manager. if self.connector is not None: - self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + self.connector.bind_kv_cache_manager(self.kv_cache_manager) self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = vllm_config.use_v2_model_runner @@ -738,6 +737,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # The request can be scheduled. break + if ( + self.connector is not None + and self.connector.has_pending_block_frees() + ): + break + # The request cannot be scheduled. # Preempt the lowest-priority request. if self.policy == SchedulingPolicy.PRIORITY: @@ -1895,14 +1900,20 @@ def update_from_output( outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - failed_kv_load_req_ids = None + failed_kv_load_req_ids: set[str] = set() + if kv_connector_output and kv_connector_output.failed_recving: + failed_kv_load_req_ids.update( + self._handle_failed_recving(kv_connector_output.failed_recving) + ) if kv_connector_output and kv_connector_output.invalid_block_ids: # These blocks contain externally computed tokens that failed to # load. Identify affected requests and adjust their computed token # count to trigger recomputation of the invalid blocks. - failed_kv_load_req_ids = self._handle_invalid_blocks( - kv_connector_output.invalid_block_ids, - num_scheduled_tokens, + failed_kv_load_req_ids.update( + self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) ) # Persist per-step routed experts into the scheduler-side slot @@ -3135,6 +3146,21 @@ def _update_requests_with_invalid_blocks( return affected_req_ids, total_affected_tokens, blocks_to_evict + def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]: + """Fail closed for layouts whose block IDs are not globally unique.""" + affected_req_ids: set[str] = set() + for req_id in failed_req_ids: + request = self.requests.get(req_id) + if ( + request is not None + and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ): + affected_req_ids.add(req_id) + if self.recompute_kv_load_failures: + request.num_computed_tokens = 0 + self.failed_recving_kv_req_ids.add(req_id) + return set() if self.recompute_kv_load_failures else affected_req_ids + def _handle_invalid_blocks( self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] ) -> set[str]: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2015967dd2df..d4e68d8adc86 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -4,8 +4,9 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Sequence -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar +from vllm.distributed.kv_events import MEDIUM_CPU from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool @@ -16,6 +17,7 @@ KVCacheBlock, resolve_block_hashes, ) +from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, @@ -23,7 +25,10 @@ CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KpoolTailSpec, + KVCacheGroupRole, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -37,6 +42,9 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request +if TYPE_CHECKING: + from vllm.v1.hisparse.coordinator import HiSparseCoordinator + logger = init_logger(__name__) @@ -48,6 +56,9 @@ class SingleTypeKVCacheManager(ABC): supports_fine_grained_hash_lookup: ClassVar[bool] = False + # Keep this group's longer prefix until external cache lookup completes. + retains_longer_hit: bool = False + @property def has_positionally_stable_blocks(self) -> bool: """Whether positional offload scans can follow this block table. @@ -2191,10 +2202,388 @@ def __init__( assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0 +class HiSparseSourceManager(FullAttentionManager): + """Host-tier manager with a private pool; publishes hashes once durable. + + Host capacity is best effort: a page that cannot get a host block keeps a + null host entry, so its GPU copy is never written back and stays pinned. + """ + + coordinator: "HiSparseCoordinator | None" = None + # The host tier keeps prefixes the device groups have already lost. + retains_longer_hit = True + + @property + def records_new_block_ids(self) -> bool: + return False + + def take_new_block_ids(self) -> list[int]: + self.new_block_ids = [] + return [] + + def bind_host_pool(self, num_blocks: int) -> None: + """Replace the device pool this group was built with by a host one.""" + device_pool = self.block_pool + self.block_pool = SharedEventQueueBlockPool( + num_gpu_blocks=num_blocks, + enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable, + hash_block_size=device_pool.hash_block_size, + enable_kv_cache_events=device_pool.enable_kv_cache_events, + metrics_collector=device_pool.metrics_collector, + medium=MEDIUM_CPU, + event_owner=device_pool, + ) + self._null_block = self.block_pool.null_block + + def take_pending_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Host copies never reach the worker's generic block-copy path.""" + return [] + + def take_host_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Drain host copies for the coordinator to hand to the connector.""" + copies = self._pending_cow_copies + self._pending_cow_copies = [] + return copies + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + if ( + total_computed_tokens > num_local_computed_tokens + or self._has_partial_local_hit( + new_computed_blocks, num_local_computed_tokens + ) + ): + # External loads need real destinations; future GPU-computed pages + # remain best effort. Use the same admission sentinel as Mamba. + required = super().get_num_blocks_to_allocate( + request_id, + total_computed_tokens, + new_computed_blocks, + total_computed_tokens, + num_local_computed_tokens, + total_computed_tokens, + ) + if required > self.block_pool.get_num_free_blocks(): + assert self.coordinator is not None + assert self.coordinator.gpu_pool is not None + return self.coordinator.gpu_pool.num_gpu_blocks + 1 + return 0 + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + num_required = cdiv(num_tokens, self.block_size) + num_free = self.block_pool.get_num_free_blocks() + if request_id in self._partial_hit_reqs: + num_free -= 1 + fit_blocks = min(num_required, len(req_blocks) + max(0, num_free)) + new_blocks = super().allocate_new_blocks( + request_id, + fit_blocks * self.block_size, + min(num_tokens_main_model, fit_blocks * self.block_size), + ) + req_blocks.extend([self._null_block] * (num_required - len(req_blocks))) + return new_blocks + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + if num_external_computed_tokens <= 0: + return + # The connector writes these pages; only a successful receive makes + # them readable, not advancing the request's computed-token count. + assert self.coordinator is not None + self.coordinator.record_pending_host_import( + request_id, num_local_computed_tokens + num_external_computed_tokens + ) + super().allocate_external_computed_blocks( + request_id, num_local_computed_tokens, num_external_computed_tokens + ) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + assert self.coordinator is not None + self.coordinator.publish_when_ready( + request, + num_tokens, + retention_interval, + replay_boundaries=replay_boundaries, + ) + + def publish_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + super().cache_blocks( + request, + num_tokens, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + assert self.coordinator is not None + self.coordinator.free(request_id) + return super().pop_blocks_for_free(request_id) + + +class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager): + """Base for ephemeral groups whose host source owns prefix caching.""" + + coordinator: "HiSparseCoordinator | None" = None + + def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None: + # Never prefix-cached, but the per-step ``cache_blocks`` hook is where + # residency work runs, so stay opted in regardless of prefix caching. + kwargs["enable_caching"] = True + super().__init__(kv_cache_spec, **kwargs) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + return None + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + return 0 + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + return tuple([] for _ in kv_cache_group_ids), 0 + + +class HiSparseHotManager(_HiSparseAuxiliaryManager): + """Allocate a hot region only after a request acquires CPU-only history.""" + + def __init__(self, kv_cache_spec: HiSparseHotSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.blocks_per_request = kv_cache_spec.blocks_per_request + self.hot_required: set[str] = set() + + def require_hot(self, request_id: str) -> None: + self.hot_required.add(request_id) + + def has_hot(self, request_id: str) -> bool: + return len(self.req_to_blocks.get(request_id, ())) == self.blocks_per_request + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + assert not new_computed_blocks + # A hot region is needed to read host-backed history: an external + # import, a new request resuming a host prefix, or one already asked + # to transition. Running requests keep their earlier answer. + host_import = total_computed_tokens > num_local_computed_tokens + resumes_host_prefix = ( + num_local_computed_tokens > 0 and request_id not in self.num_cached_block + ) + if host_import or resumes_host_prefix or request_id in self.hot_required: + return self.get_num_required_blocks(request_id) + return 0 + + def get_num_required_blocks(self, request_id: str) -> int: + return max( + self.blocks_per_request - len(self.req_to_blocks.get(request_id, ())), + 0, + ) + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + assert not new_computed_blocks + self.num_cached_block[request_id] = 0 + assert self.coordinator is not None + if num_local_computed_tokens > 0 or not self.coordinator.resident_managers: + self.require_hot(request_id) + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + self.require_hot(request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + # Cold admissions can bypass add_local_computed_blocks. + self.num_cached_block[request_id] = 0 + if request_id not in self.hot_required: + return [] + req_blocks = self.req_to_blocks[request_id] + num_new_blocks = self.blocks_per_request - len(req_blocks) + if num_new_blocks <= 0: + return [] + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + req_blocks.extend(new_blocks) + return new_blocks + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + self.hot_required.discard(request_id) + return super().pop_blocks_for_free(request_id) + + +class HiSparseResidentManager(_HiSparseAuxiliaryManager): + """Track GPU-resident pages for otherwise host-backed KV.""" + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + del num_tokens_main_model + assert not new_computed_blocks + if total_computed_tokens > num_local_computed_tokens: + if total_computed_tokens % self.block_size != 0: + raise ValueError( + "A host-only HiSparse import must end on a cache-block boundary." + ) + imported_pages = total_computed_tokens // self.block_size + return max(cdiv(num_tokens, self.block_size) - imported_pages, 0) + existing = len(self.req_to_blocks.get(request_id, ())) + host_pages = cdiv(num_local_computed_tokens, self.block_size) + required = cdiv(num_tokens, self.block_size) + if apply_admission_cap: + assert self._max_admission_blocks_per_request is not None + required = min(required, self._max_admission_blocks_per_request) + return max(required - max(existing, host_pages), 0) + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + """Represent imported host history with null resident pages.""" + assert num_external_computed_tokens > 0 + num_tokens = num_local_computed_tokens + num_external_computed_tokens + blocks = self.req_to_blocks[request_id] + tail_page = cdiv(num_tokens, self.block_size) - 1 + if len(blocks) <= tail_page: + blocks.extend([self._null_block] * (tail_page + 1 - len(blocks))) + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + assert not new_computed_blocks + req_blocks = self.req_to_blocks[request_id] + assert not req_blocks + num_host_pages = cdiv(num_local_computed_tokens, self.block_size) + req_blocks.extend([self._null_block] * num_host_pages) + self.num_cached_block[request_id] = 0 + assert self.coordinator is not None + self.coordinator.commit_computed_blocks(request_id, num_host_pages) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + assert self.coordinator is not None + self.coordinator.plan_prefix_materialization(request.request_id, num_tokens) + self.coordinator.update_residency(request.request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + del num_tokens_main_model + req_blocks = self.req_to_blocks[request_id] + num_new_blocks = cdiv(num_tokens, self.block_size) - len(req_blocks) + if num_new_blocks <= 0: + return [] + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + req_blocks.extend(new_blocks) + return new_blocks + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + assert self.coordinator is not None + self.coordinator.free(request_id) + return super().pop_blocks_for_free(request_id) + + def adopt_resident_page( + self, request_id: str, block_idx: int, block: KVCacheBlock + ) -> bool: + """Point a null prefix page at a pinned GPU copy of its contents.""" + blocks = self.req_to_blocks.get(request_id) + if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null: + return False + blocks[block_idx] = block + return True + + def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None: + blocks = self.req_to_blocks.get(request_id) + if blocks is None or block_idx >= len(blocks): + return None + block = blocks[block_idx] + return None if block.is_null else block + + def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, max_in_flight_tokens: int, max_model_len: int, + role: str | None = None, **kwargs, ) -> SingleTypeKVCacheManager: """ @@ -2212,7 +2601,7 @@ def get_manager_for_kv_cache_spec( Returns: An instance of the appropriate SingleTypeKVCacheManager subclass """ - manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec) + manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec, role) assert manager_class is not None, ( f"No manager registered for KVCacheSpec {type(kv_cache_spec)}" ) @@ -2224,7 +2613,7 @@ def get_manager_for_kv_cache_spec( # FullAttentionSpec sizing without a separate admission cap. if isinstance( kv_cache_spec, - (SlidingWindowSpec, ChunkedLocalAttentionSpec), + (SlidingWindowSpec, ChunkedLocalAttentionSpec, HiSparseResidentSpec), ): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( @@ -2238,11 +2627,24 @@ def get_manager_for_kv_cache_spec( def register_all_kvcache_specs(vllm_config): """Built-in spec registration""" + KVCacheSpecRegistry.register_role_manager( + KVCacheGroupRole.HISPARSE_SOURCE, HiSparseSourceManager + ) KVCacheSpecRegistry.register( FullAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec, ) + KVCacheSpecRegistry.register( + HiSparseHotSpec, + HiSparseHotManager, + uniform_type_base_spec=HiSparseHotSpec, + ) + KVCacheSpecRegistry.register( + HiSparseResidentSpec, + HiSparseResidentManager, + uniform_type_base_spec=HiSparseResidentSpec, + ) KVCacheSpecRegistry.register( SlidingWindowSpec, diff --git a/vllm/v1/hisparse/__init__.py b/vllm/v1/hisparse/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/v1/hisparse/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py new file mode 100644 index 000000000000..263e553d32c4 --- /dev/null +++ b/vllm/v1/hisparse/binding.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker-side wiring of HiSparse caches to attention layers.""" + +from collections.abc import Mapping +from copy import copy +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.v1.hisparse.layout import ( + HISPARSE_HOT_SUFFIX, + HISPARSE_RESIDENT_SUFFIX, +) +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + HiSparseHostPool, + initialize_hisparse_runtime_buffers, + release_pinned_state, +) +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheLayout, + KVCacheSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + create_kv_cache_views, +) +from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.v1.worker.gpu.block_table import BlockTables + + +def resolve_hisparse_block_size( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + attn_layers: Mapping[str, "AttentionLayerBase"], +) -> None: + """Resolve a common kernel block size in-place for HiSparse MLA specs.""" + if vllm_config.attention_config.hisparse_config is None: + return + mla_specs = { + name: spec + for name, spec in kv_cache_spec.items() + if isinstance(spec, MLAAttentionSpec) + } + if not mla_specs: + return + block_sizes = {spec.block_size for spec in mla_specs.values()} + if len(block_sizes) != 1: + raise ValueError("HiSparse requires one scheduler block size.") + backends = [attn_layers[name].get_attn_backend() for name in mla_specs] + try: + block_size = select_common_block_size(block_sizes.pop(), backends) + except ValueError as error: + raise ValueError( + "HiSparse requires a GPU block size supported by every sparse " + f"attention and indexer backend: {error}" + ) from error + kv_cache_spec.update( + (name, spec.copy_with_new_block_size(block_size)) + for name, spec in mla_specs.items() + ) + + +def allocate_hisparse_kv_caches( + kv_cache_config: KVCacheConfig, + device: torch.device, + layout: KVCacheLayout, + kernel_block_sizes: list[int], + host_pool: HiSparseHostPool, +) -> dict[str, torch.Tensor]: + """Allocate the host pool separately from the shared device backing.""" + device_config = copy(kv_cache_config) + device_config.kv_cache_tensors = [ + tensor + for tensor in kv_cache_config.kv_cache_tensors + if not tensor.host_resident + ] + kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes) + host_tensors = [ + tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident + ] + (host_size,) = {tensor.size for tensor in host_tensors} + backing = host_pool.allocate(host_size) + (host_group_id,) = kv_cache_config.host_group_ids + host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec + for tensor in host_tensors: + spec = ( + host_spec.kv_cache_specs[tensor.layers[0]] + if isinstance(host_spec, UniformTypeKVCacheSpecs) + else host_spec + ) + kernel_block_size = kernel_block_sizes[host_group_id] + if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: + kernel_block_size = spec.storage_block_size + views = create_kv_cache_views( + backing, + spec, + kv_cache_config.num_blocks_of(tensor), + layout, + tensor, + kernel_block_size=kernel_block_size, + ) + kv_caches.update(zip(tensor.layers, views)) + return kv_caches + + +def init_hisparse_kv_cache( + kv_cache_config: KVCacheConfig, + device: torch.device, + kernel_block_sizes: list[int], + vllm_config: VllmConfig, + forward_context: dict[str, Any], + block_tables: "BlockTables", +) -> dict[str, torch.Tensor]: + """Allocate and bind HiSparse caches within the caller's allocation context.""" + host_pool = HiSparseHostPool() + kv_caches = allocate_hisparse_kv_caches( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + host_pool, + ) + cache_handles = bind_hisparse_kv_caches( + forward_context=forward_context, + kv_cache_config=kv_cache_config, + kv_caches=kv_caches, + block_tables=block_tables, + host_pool=host_pool, + ) + initialize_hisparse_runtime_buffers( + cache_handles, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, + ) + return kv_caches + + +def _get_hisparse_cache( + forward_context: dict[str, Any], layer_name: str +) -> HiSparseCacheHandle: + attention_layer = forward_context[layer_name] + hisparse_cache = attention_layer.hisparse_cache + assert hisparse_cache is not None + return hisparse_cache + + +def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: + cache_handles = [ + cache + for layer in forward_context.values() + if (cache := getattr(layer, "hisparse_cache", None)) is not None + ] + runtimes = { + id(cache.runtime): cache.runtime + for cache in cache_handles + if hasattr(cache.runtime, "_host_cache") + } + if not runtimes: + return + + registered_pools = list( + { + runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool + for runtime in runtimes.values() + }.values() + ) + release_pinned_state(list(runtimes.values()), registered_pools) + for cache in cache_handles: + cache.mirror_staging_cache = None + cache.mirror_staging_slots = None + + +def bind_hisparse_kv_caches( + *, + forward_context: dict[str, Any], + kv_cache_config: KVCacheConfig, + kv_caches: dict[str, torch.Tensor], + block_tables: "BlockTables", + host_pool: HiSparseHostPool, +) -> list[HiSparseCacheHandle]: + """Bind existing cache storage and block tables; return the bound handles.""" + assert host_pool.registered is not None + tensor_configs = { + name: tensor_config + for tensor_config in kv_cache_config.kv_cache_tensors + for name in tensor_config.layers + } + resident_source_index = 0 + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) + layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] + tensor_config = tensor_configs[cache_name] + assert not tensor_config.host_resident + cache_handle = _get_hisparse_cache(forward_context, layer_name) + cache_handle.bind_cache( + kv_caches[cache_name], + byte_offset=tensor_config.offset, + block_stride=tensor_config.block_stride, + num_blocks=kv_cache_config.num_blocks, + block_size=group.kv_cache_spec.block_size, + block_table=block_tables.input_block_tables[group_id], + slot_mapping=block_tables.slot_mappings[group_id], + ) + assert cache_handle.view is not None + kv_caches[cache_name] = cache_handle.view.cache + cache_handle.runtime.resident_source_index = resident_source_index + resident_source_index += 1 + + hot_backing: torch.Tensor | None = None + cache_handles: list[HiSparseCacheHandle] = [] + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, HiSparseHotSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_HOT_SUFFIX) + raw_tensor = kv_caches[cache_name] + if hot_backing is None: + hot_backing = raw_tensor + elif hot_backing.untyped_storage().data_ptr() != ( + raw_tensor.untyped_storage().data_ptr() + ): + raise RuntimeError("HiSparse hot tensors must share one GPU backing.") + layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] + cache_handle = _get_hisparse_cache(forward_context, layer_name) + tensor_config = tensor_configs[cache_name] + assert not tensor_config.host_resident + cache_handle.runtime.bind_hot_cache( + raw_tensor, + byte_offset=tensor_config.offset, + block_stride=tensor_config.block_stride, + num_blocks=kv_cache_config.num_blocks, + block_size=group.kv_cache_spec.block_size, + block_table=block_tables.input_block_tables[group_id], + ) + resident = cache_handle.view + hot = cache_handle.runtime.hot + assert resident is not None + if ( + resident.cache.untyped_storage().data_ptr() + != hot.cache.untyped_storage().data_ptr() + or resident.cache.data_ptr() != hot.cache.data_ptr() + or resident.cache.stride() != hot.cache.stride() + ): + raise RuntimeError("HiSparse resident and hot layouts must match.") + source_cache = kv_caches[layer_name] + assert source_cache.untyped_storage().data_ptr() == ( + host_pool.registered.untyped_storage().data_ptr() + ) + cache_handle.runtime.bind_source_cache( + source_cache, + registered_host_pool=host_pool.registered, + ) + cache_handles.append(cache_handle) + # The hot slab is raw storage owned by the runtime, not a layer + # cache: nothing downstream binds or registers it. + del kv_caches[cache_name] + + if hot_backing is None or not cache_handles: + raise RuntimeError("HiSparse found no hot-cache handles.") + (source_group_id,) = kv_cache_config.host_group_ids + for cache_handle in cache_handles: + cache_handle.source_block_table = block_tables.input_block_tables[ + source_group_id + ] + cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id] + return cache_handles diff --git a/vllm/v1/hisparse/block_pool.py b/vllm/v1/hisparse/block_pool.py new file mode 100644 index 000000000000..cbb620d64769 --- /dev/null +++ b/vllm/v1/hisparse/block_pool.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.kv_events import KVCacheEvent +from vllm.v1.core.block_pool import BlockPool + + +class SharedEventQueueBlockPool(BlockPool): + """A pool that publishes into another pool's live KV event queue. + + The owner rebinds its queue on every drain, so this reads it through the + owner rather than holding a reference. + """ + + def __init__(self, *args, event_owner: BlockPool, **kwargs) -> None: + self._event_owner = event_owner + super().__init__(*args, **kwargs) + + @property + def kv_event_queue(self) -> list[KVCacheEvent]: + return self._event_owner.kv_event_queue + + @kv_event_queue.setter + def kv_event_queue(self, events: list[KVCacheEvent]) -> None: + # ``BlockPool.__init__`` seeds an empty queue and ``take_events`` + # swaps in a fresh one; both belong to the owner, which drains it. + assert not events diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py new file mode 100644 index 000000000000..ecbc0a270ac1 --- /dev/null +++ b/vllm/v1/hisparse/coordinator.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from vllm.utils.math_utils import cdiv +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_utils import ( + BlockHashWithGroupId, + KVCacheBlock, + KVCacheBlockCopy, +) +from vllm.v1.core.single_type_kv_cache_manager import ( + HiSparseHotManager, + HiSparseResidentManager, + HiSparseSourceManager, + SingleTypeKVCacheManager, +) +from vllm.v1.hisparse.types import ( + SparseKVOffloadCommand, + SparseKVPageTransfer, + SparseKVRowMirror, +) +from vllm.v1.kv_cache_interface import ( + HiSparseResidentSpec, + KVCacheConfig, +) +from vllm.v1.request import Request + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheManager + +# Sealed pages this many positions behind the block-table tail stay pinned so +# a page written by an in-flight step is never handed out under it. +_ACTIVE_TAIL_PAGES = 2 + + +@dataclass +class _PendingPublication: + request: Request + num_computed_tokens: int + num_pages: int + retention_interval: int | None + replay_boundaries: Sequence[int] + + +@dataclass +class _HiSparseRequestState: + """Residency bookkeeping for one request. + + A resident page moves through three states: dirty (only the GPU copy + exists), clean (``valid_pages``: the host copy is complete) and unpinned + (``unpinned_pages``: clean, and its allocation reference was released so + the pool may reuse it). ``pinned_clean`` tracks clean pages still holding + their reference, which only happens before the request can read from host. + """ + + valid_pages: set[int] = field(default_factory=set) + ready_prefix_pages: int = 0 + pending_pages: dict[int, int] = field(default_factory=dict) + publication: _PendingPublication | None = None + copies_recorded_blocks: int = 0 + pinned_clean: set[int] = field(default_factory=set) + unpinned_pages: set[int] = field(default_factory=set) + + +@dataclass +class _PendingSpill: + transfer_id: int + page: tuple[str, int] + request_state: _HiSparseRequestState + host_block: KVCacheBlock + resident_blocks: tuple[KVCacheBlock, ...] + expected_worker_completions: int = 0 + worker_completions: int = 0 + enqueue_applied: bool = False + + +class HiSparseCoordinator: + """Own HiSparse host allocation, host publication, spills, and GPU residency. + + GPU-resident pages are a write-back cache of the host tier. Once a page's + host copy is durable and its owner can read from host (it has, or has + asked for, a hot region), the page's allocation reference is released with + ``BlockPool.unpin_blocks``: the block stays in the block table and keeps + being read, but the pool counts it as free and may hand it out, at which + point the owner's page is nulled and its block table republished. + """ + + def __init__( + self, + kv_cache_config: KVCacheConfig, + managers: tuple[SingleTypeKVCacheManager, ...], + max_model_len: int, + ) -> None: + self.managers = managers + self.max_model_len = max_model_len + groups = kv_cache_config.kv_cache_groups + + resident_managers: list[HiSparseResidentManager] = [] + hot_managers: list[HiSparseHotManager] = [] + resident_group_ids: list[int] = [] + for group_id, (group, manager) in enumerate(zip(groups, managers)): + if isinstance(group.kv_cache_spec, HiSparseResidentSpec): + if not isinstance(manager, HiSparseResidentManager): + raise TypeError( + "HiSparse resident specs require resident cache managers." + ) + resident_managers.append(manager) + resident_group_ids.append(group_id) + assert not group.host_resident + if isinstance(manager, HiSparseHotManager): + hot_managers.append(manager) + assert not group.host_resident + if isinstance(manager, (HiSparseHotManager, HiSparseResidentManager)): + manager.coordinator = self + self.resident_managers = tuple(resident_managers) + self.hot_managers = tuple(hot_managers) + self.host_manager: HiSparseSourceManager | None = None + self.host_group_id: int | None = None + self.max_spill_pages = 0 + for group_id, manager in enumerate(managers): + if isinstance(manager, HiSparseSourceManager): + if self.host_manager is not None: + raise ValueError("Only one HiSparse host group is supported.") + num_host_blocks = kv_cache_config.hisparse_host_num_blocks + if num_host_blocks is None: + raise ValueError("HiSparse host group needs host capacity.") + self.host_group_id = group_id + self.host_manager = manager + manager.coordinator = self + manager.bind_host_pool(num_host_blocks) + self.has_host_cache = self.host_manager is not None + self.gpu_pool: BlockPool | None = None + self.transition_watermark = 0 + if self.resident_managers: + assert self.host_manager is not None + assert self.host_group_id is not None + if self.host_group_id > min(resident_group_ids): + raise ValueError("HiSparse host group must precede resident groups.") + resident_block_sizes = { + manager.block_size for manager in self.resident_managers + } + if len(resident_block_sizes) != 1: + raise ValueError( + "HiSparse resident cache groups must use one block size." + ) + resident_block_size = resident_block_sizes.pop() + if self.host_manager.block_size != resident_block_size: + raise ValueError("HiSparse host and resident block sizes must match.") + self.max_spill_pages = max_model_len // resident_block_size + self.gpu_pool = self.resident_managers[0].block_pool + # Requests start reading from host once the shared pool runs this + # low, so admissions find evictable pages rather than pinned ones. + hot_cost = sum(manager.blocks_per_request for manager in hot_managers) + self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10) + + # Host copies handed to the worker, by destination block id, until it + # reports having run them. + self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {} + # request -> prefix length an external load is still filling in. + self._pending_imports: dict[str, int] = {} + self.block_table_updates: set[str] = set() + self.spills_to_send: list[SparseKVPageTransfer] = [] + self.pending_spills: dict[int, _PendingSpill] = {} + self.request_states: dict[str, _HiSparseRequestState] = {} + self.next_spill_id = 0 + # Published host block hash -> GPU copies of that page, readable until + # the pool reuses them. Lets a host prefix hit come back GPU-resident. + self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {} + self._copy_by_block: dict[int, BlockHashWithGroupId] = {} + # Unpinned GPU block id -> (request, page) pairs still reading it. + self._owners: dict[int, set[tuple[str, int]]] = {} + + def _get_request_state(self, request_id: str) -> _HiSparseRequestState: + state = self.request_states.get(request_id) + if state is None: + state = _HiSparseRequestState() + self.request_states[request_id] = state + return state + + def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None: + """Account for a prefix hit on host pages the request now references.""" + if not self.resident_managers or self.host_manager is None: + return + host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) + num_host_pages = min(num_host_pages, len(host_blocks)) + if num_host_pages == 0: + return + state = self._get_request_state(request_id) + state.valid_pages.update(range(num_host_pages)) + state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages) + self._adopt_copies(request_id, state, host_blocks[:num_host_pages]) + state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages) + + def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]: + """Drain host copy-on-write work, retaining both endpoints. + + The worker runs the copies in the step that carries them, so the + endpoints are released when that step's output comes back. + """ + if self.host_manager is None: + return () + pairs = self.host_manager.take_host_cow_copies() + for source_block, cow_block in pairs: + self._retained_copies[cow_block.block_id] = (source_block, cow_block) + return tuple( + KVCacheBlockCopy( + src_block_id=source_block.block_id, + dst_block_id=cow_block.block_id, + ) + for source_block, cow_block in pairs + ) + + def release_completed_host_copies(self, dst_block_ids: Iterable[int]) -> None: + """Release the endpoints of host copies the worker reports having run.""" + blocks: list[KVCacheBlock] = [] + for dst_block_id in dst_block_ids: + pair = self._retained_copies.pop(dst_block_id, None) + if pair is not None: + blocks.extend(pair) + if not blocks: + return + assert self.host_manager is not None + self.host_manager.block_pool.free_blocks(reversed(blocks)) + + def get_host_block_pool(self) -> BlockPool | None: + manager = self.host_manager + return manager.block_pool if manager is not None else None + + # ------------------------------------------------------------------ + # GPU copies of published host pages + # ------------------------------------------------------------------ + + def _adopt_copies( + self, + request_id: str, + state: _HiSparseRequestState, + host_blocks: Sequence[KVCacheBlock], + ) -> None: + """Point null prefix pages at readable GPU copies of their contents. + + Runs after the hit length is fully reconciled, so copy hits and misses + can be scattered per page without affecting the prefix hit. + """ + if not self.copies: + return + for host_idx, host_block in enumerate(host_blocks): + if host_block.is_null or host_block.block_hash is None: + continue + blocks = self.copies.get(host_block.block_hash) + if blocks is None: + continue + for manager, block in zip(self.resident_managers, blocks): + if manager.adopt_resident_page(request_id, host_idx, block): + manager.block_pool.touch([block]) + state.pinned_clean.add(host_idx) + + def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: + """Index the GPU copies of just-published host blocks for later hits.""" + if not self.resident_managers: + return + assert self.host_manager is not None + host_blocks = self.host_manager.req_to_blocks.get(request_id) + if not host_blocks: + return + state = self._get_request_state(request_id) + num_host_blocks = min( + num_computed_tokens // self.host_manager.block_size, len(host_blocks) + ) + for host_idx in range(state.copies_recorded_blocks, num_host_blocks): + host_block = host_blocks[host_idx] + if host_block.is_null or host_block.block_hash is None: + continue + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, host_idx) + if block is None: + break + blocks.append(block) + if len(blocks) != len(self.resident_managers): + continue + self._drop_copy(host_block.block_hash) + self.copies[host_block.block_hash] = tuple(blocks) + for block in blocks: + self._copy_by_block[block.block_id] = host_block.block_hash + state.copies_recorded_blocks = max( + state.copies_recorded_blocks, num_host_blocks + ) + + def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None: + blocks = self.copies.pop(host_hash, None) + if blocks is None: + return + for block in blocks: + self._copy_by_block.pop(block.block_id, None) + + # ------------------------------------------------------------------ + # Residency: pinning, unpinning and losing pages + # ------------------------------------------------------------------ + + def _can_read_from_host(self, request_id: str) -> bool: + """Whether resident pages may be dropped under the request.""" + return bool(self.hot_managers) and all( + manager.has_hot(request_id) for manager in self.hot_managers + ) + + def _resident_page_blocks( + self, request_id: str, page_idx: int + ) -> list[KVCacheBlock] | None: + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + req_blocks = manager.req_to_blocks.get(request_id) + if ( + req_blocks is None + or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES + or req_blocks[page_idx].is_null + ): + return None + blocks.append(req_blocks[page_idx]) + return blocks + + def _unpin_page( + self, request_id: str, state: _HiSparseRequestState, page_idx: int + ) -> bool: + if page_idx in state.pending_pages or page_idx not in state.valid_pages: + return False + blocks = self._resident_page_blocks(request_id, page_idx) + if blocks is None: + return False + for manager, block in zip(self.resident_managers, blocks): + self._owners.setdefault(block.block_id, set()).add((request_id, page_idx)) + manager.block_pool.unpin_blocks([block], self._on_block_reused) + state.pinned_clean.discard(page_idx) + state.unpinned_pages.add(page_idx) + return True + + def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None: + for page_idx in sorted(state.pinned_clean): + self._unpin_page(request_id, state, page_idx) + + def _page_became_clean( + self, request_id: str, state: _HiSparseRequestState, page_idx: int + ) -> None: + state.pinned_clean.add(page_idx) + if self._can_read_from_host(request_id): + self._unpin_page(request_id, state, page_idx) + + def _on_block_reused(self, block: KVCacheBlock) -> None: + host_hash = self._copy_by_block.get(block.block_id) + if host_hash is not None: + self._drop_copy(host_hash) + for request_id, page_idx in self._owners.pop(block.block_id, ()): + self._lose_page(request_id, page_idx) + + def _lose_page(self, request_id: str, page_idx: int) -> None: + state = self.request_states.get(request_id) + if state is None: + return + for manager in self.resident_managers: + blocks = manager.req_to_blocks.get(request_id) + if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null: + continue + block = blocks[page_idx] + blocks[page_idx] = manager._null_block + owners = self._owners.get(block.block_id) + if owners is not None: + owners.discard((request_id, page_idx)) + state.unpinned_pages.discard(page_idx) + self.block_table_updates.add(request_id) + for hot_manager in self.hot_managers: + hot_manager.require_hot(request_id) + + def update_residency(self, request_id: str) -> None: + """Per-step residency policy for a scheduled request. + + A request that can read from host releases every clean sealed page to + the pool. One that cannot keeps its pages pinned until the shared pool + runs low, then asks for a hot region. Pages remain pinned until that region + is allocated on a subsequent scheduling pass. + """ + if not self.resident_managers: + return + state = self._get_request_state(request_id) + if not self._can_read_from_host(request_id): + assert self.gpu_pool is not None + if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark: + return + for manager in self.hot_managers: + manager.require_hot(request_id) + return + self._unpin_clean_pages(request_id, state) + + # ------------------------------------------------------------------ + # Host publication and spills + # ------------------------------------------------------------------ + + def plan_prefix_materialization( + self, request_id: str, num_computed_tokens: int + ) -> None: + """Eagerly write back every sealed page that has no host copy yet.""" + if not self.resident_managers: + return + assert self.host_manager is not None + host_block_size = self.host_manager.block_size + num_pages = num_computed_tokens // host_block_size + state = self._get_request_state(request_id) + budget = max(self.max_spill_pages - len(self.spills_to_send), 0) + for page_idx in range(num_pages): + if budget == 0: + break + if page_idx in state.pending_pages: + continue + if page_idx in state.valid_pages: + continue + if not all( + manager.get_resident_page(request_id, page_idx) is not None + for manager in self.resident_managers + ): + continue + if self._plan_spill(request_id, page_idx, after_forward=True): + budget -= 1 + + def publish_when_ready( + self, + request: Request, + num_computed_tokens: int, + retention_interval: int | None, + *, + replay_boundaries: Sequence[int], + ) -> None: + """Publish host-source hashes only after their pages are durable.""" + manager = self.host_manager + if manager is None: + return + num_pages = num_computed_tokens // manager.block_size + request_id = request.request_id + state = self._get_request_state(request_id) + if state.ready_prefix_pages >= num_pages: + manager.publish_blocks( + request, + num_computed_tokens, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + self._record_copies(request_id, num_computed_tokens) + state.publication = None + return + state.publication = _PendingPublication( + request=request, + num_computed_tokens=num_computed_tokens, + num_pages=num_pages, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + + def _publish_host_blocks_if_ready(self, request_id: str) -> None: + state = self.request_states.get(request_id) + if state is None: + return + publication = state.publication + if publication is None or state.ready_prefix_pages < publication.num_pages: + return + assert self.host_manager is not None + self.host_manager.publish_blocks( + publication.request, + publication.num_computed_tokens, + retention_interval=publication.retention_interval, + replay_boundaries=publication.replay_boundaries, + ) + self._record_copies(request_id, publication.num_computed_tokens) + state.publication = None + + def record_pending_host_import(self, request_id: str, num_tokens: int) -> None: + """Note a prefix an external load is populating in host pages.""" + self._pending_imports[request_id] = num_tokens + + def finish_host_import(self, request_id: str, *, failed: bool) -> None: + num_tokens = self._pending_imports.pop(request_id, None) + if num_tokens is not None and not failed: + self.complete_host_import(request_id, num_tokens) + + def complete_host_import(self, request_id: str, num_computed_tokens: int) -> None: + """Publish externally populated host pages after connector completion.""" + if not self.resident_managers: + return + block_size = self.resident_managers[0].block_size + num_pages = num_computed_tokens // block_size + state = self._get_request_state(request_id) + state.valid_pages = set(range(num_pages)) + state.ready_prefix_pages = num_pages + self._publish_host_blocks_if_ready(request_id) + + def _plan_spill( + self, + request_id: str, + page_idx: int, + *, + after_forward: bool, + ) -> bool: + assert self.host_manager is not None + state = self._get_request_state(request_id) + host_blocks = self.host_manager.req_to_blocks.get(request_id) + if host_blocks is None or page_idx >= len(host_blocks): + return False + host_block = host_blocks[page_idx] + if host_block.is_null: + return False + + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, page_idx) + if block is None: + return False + blocks.append(block) + + self.host_manager.block_pool.touch([host_block]) + for manager, block in zip(self.resident_managers, blocks): + manager.block_pool.touch([block]) + spill_id = self.next_spill_id + self.next_spill_id += 1 + plan = SparseKVPageTransfer( + transfer_id=spill_id, + destination_block_id=host_block.block_id, + source_block_ids=tuple(block.block_id for block in blocks), + after_forward=after_forward, + ) + self.pending_spills[spill_id] = _PendingSpill( + transfer_id=spill_id, + page=(request_id, page_idx), + request_state=state, + host_block=host_block, + resident_blocks=tuple(blocks), + ) + state.pending_pages[page_idx] = spill_id + self.spills_to_send.append(plan) + return True + + # ------------------------------------------------------------------ + # Worker-facing views + # ------------------------------------------------------------------ + + def build_row_mirrors( + self, + requests: Iterable[tuple[str, int, int]], + ) -> tuple[SparseKVRowMirror, ...]: + """Map scheduled rows to scheduler-owned resident and host blocks.""" + if not self.resident_managers or self.host_manager is None: + return () + block_size = self.resident_managers[0].block_size + mirrors: list[SparseKVRowMirror] = [] + for request_id, num_computed_tokens, num_scheduled_tokens in requests: + host_blocks = self.host_manager.req_to_blocks.get(request_id) + resident_blocks = [ + manager.req_to_blocks.get(request_id) + for manager in self.resident_managers + ] + if host_blocks is None or any(blocks is None for blocks in resident_blocks): + continue + token_position = num_computed_tokens + end_position = token_position + num_scheduled_tokens + while token_position < end_position: + page_idx, row_offset = divmod(token_position, block_size) + num_rows = min(block_size - row_offset, end_position - token_position) + # A page without a GPU copy has no mirror source, and a page + # without a host block has no destination. Skip just that page: + # its neighbours in the window are still mirrorable, and a + # window only ever moves forward, so dropping them here would + # leave their host rows permanently stale. + source_starts = [] + for blocks in resident_blocks: + assert blocks is not None + if page_idx >= len(blocks) or blocks[page_idx].is_null: + break + source_starts.append( + blocks[page_idx].block_id * block_size + row_offset + ) + if len(source_starts) != len(resident_blocks): + token_position += num_rows + continue + host_block_idx = page_idx + if host_block_idx >= len(host_blocks): + token_position += num_rows + continue + host_block = host_blocks[host_block_idx] + if host_block.is_null: + token_position += num_rows + continue + mirrors.append( + SparseKVRowMirror( + source_starts=tuple(source_starts), + destination_start=( + host_block.block_id * block_size + row_offset + ), + num_rows=num_rows, + ) + ) + token_position += num_rows + return tuple(mirrors) + + def all_context_pages_resident( + self, + requests: Iterable[tuple[str, int, int]], + ) -> bool: + """Return whether every scheduled request can read only resident KV.""" + if not self.resident_managers: + return False + block_size = self.resident_managers[0].block_size + for request_id, num_computed_tokens, num_scheduled_tokens in requests: + num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size) + for page_idx in range(num_pages): + if any( + manager.get_resident_page(request_id, page_idx) is None + for manager in self.resident_managers + ): + return False + return True + + def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: + updates = { + request_id: tuple( + [block.block_id for block in manager.req_to_blocks.get(request_id, [])] + for manager in self.managers + ) + for request_id in self.block_table_updates + } + self.block_table_updates.clear() + return updates + + def build_offload_command(self) -> SparseKVOffloadCommand | None: + """Package residency decisions for the worker connector.""" + if not self.resident_managers: + return None + command = SparseKVOffloadCommand(page_transfers=self.spills_to_send) + self.spills_to_send = [] + return command + + def has_pending_reclamation(self) -> bool: + """Whether an in-flight spill will free GPU blocks on completion.""" + return any( + pending.page[0] in self.request_states + and self._can_read_from_host(pending.page[0]) + for pending in self.pending_spills.values() + ) + + def has_pending_work(self) -> bool: + return bool(self.spills_to_send or self.pending_spills or self._retained_copies) + + def update_spills( + self, + enqueued_counts: Mapping[int, int], + completed_counts: Mapping[int, int], + ) -> None: + for spill_id, count in enqueued_counts.items(): + pending = self.pending_spills.get(spill_id) + if pending is not None and pending.expected_worker_completions == 0: + pending.expected_worker_completions = count + for spill_id, count in completed_counts.items(): + pending = self.pending_spills.get(spill_id) + if pending is not None: + pending.worker_completions += count + + self._apply_enqueued_spills() + self._complete_host_writes() + + def _apply_enqueued_spills(self) -> None: + """Drop the DMA pins once the worker has enqueued a spill.""" + for pending in self.pending_spills.values(): + if ( + pending.enqueue_applied + or not pending.expected_worker_completions + or pending.worker_completions < pending.expected_worker_completions + ): + continue + for manager, block in zip(self.resident_managers, pending.resident_blocks): + manager.block_pool.free_blocks([block]) + pending.enqueue_applied = True + + def _complete_host_writes(self) -> None: + completed = [ + pending + for pending in self.pending_spills.values() + if pending.enqueue_applied + and pending.expected_worker_completions + and pending.worker_completions >= pending.expected_worker_completions + ] + completed_request_ids: set[str] = set() + for pending in completed: + self.pending_spills.pop(pending.transfer_id, None) + request_id, page_idx = pending.page + pending.request_state.pending_pages.pop(page_idx, None) + if self.request_states.get(request_id) is pending.request_state: + pending.request_state.valid_pages.add(page_idx) + self._page_became_clean(request_id, pending.request_state, page_idx) + completed_request_ids.add(request_id) + assert self.host_manager is not None + self.host_manager.block_pool.free_blocks([pending.host_block]) + for request_id in completed_request_ids: + state = self.request_states[request_id] + while state.ready_prefix_pages in state.valid_pages: + state.ready_prefix_pages += 1 + self._publish_host_blocks_if_ready(request_id) + + def free(self, request_id: str) -> None: + """Detach the request; its clean pages stay readable copies in the pool.""" + self._pending_imports.pop(request_id, None) + state = self.request_states.pop(request_id, None) + if state is None: + return + for manager in self.resident_managers: + blocks = manager.req_to_blocks.get(request_id) + if not blocks: + continue + for page_idx, block in enumerate(blocks): + if block.is_null: + continue + if page_idx in state.unpinned_pages: + owners = self._owners.get(block.block_id) + if owners is not None: + owners.discard((request_id, page_idx)) + elif ( + page_idx in state.valid_pages + and page_idx not in state.pending_pages + ): + manager.block_pool.unpin_blocks([block], self._on_block_reused) + else: + continue + blocks[page_idx] = manager._null_block + + +def get_hisparse_coordinator( + kv_cache_manager: "KVCacheManager", +) -> HiSparseCoordinator: + """Return the coordinator shared by a KV cache manager's HiSparse groups. + + Built on first call and memoised on the HiSparse managers it binds itself + to. Must run before the first request is admitted. + """ + managers = tuple(kv_cache_manager.coordinator.single_type_managers) + for manager in managers: + coordinator = getattr(manager, "coordinator", None) + if coordinator is not None: + assert isinstance(coordinator, HiSparseCoordinator) + return coordinator + coordinator = HiSparseCoordinator( + kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len + ) + if coordinator.host_manager is None: + raise ValueError("No HiSparse cache group is configured.") + return coordinator diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py new file mode 100644 index 000000000000..d74b0e8cfd49 --- /dev/null +++ b/vllm/v1/hisparse/layout.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from dataclasses import dataclass +from typing import cast + +from vllm.config import VllmConfig +from vllm.config.kv_transfer import hisparse_host_pool_gib +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv +from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupRole, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MLAAttentionSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + compute_layout_strides, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +logger = init_logger(__name__) + +HISPARSE_HOT_SUFFIX = ".hisparse_hot" +HISPARSE_RESIDENT_SUFFIX = ".hisparse_resident" + + +@dataclass(frozen=True) +class HiSparseLayout: + source_group: KVCacheGroupSpec + device_groups: list[KVCacheGroupSpec] + host_num_blocks: int + + +def get_hisparse_kv_cache_groups( + vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec] +) -> list[KVCacheGroupSpec] | None: + attention_config = getattr(vllm_config, "attention_config", None) + if attention_config is None or attention_config.hisparse_config is None: + return None + + mla_specs: dict[str, KVCacheSpec] = { + name: spec + for name, spec in kv_cache_spec.items() + if isinstance(spec, MLAAttentionSpec) + } + other_specs = { + name: spec + for name, spec in kv_cache_spec.items() + if not isinstance(spec, MLAAttentionSpec) + } + if not mla_specs or not other_specs: + return None + + from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + + mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) + assert mla_group_spec is not None + mla_group = KVCacheGroupSpec(list(mla_specs), mla_group_spec) + return [mla_group, *get_kv_cache_groups(vllm_config, other_specs)] + + +def get_hisparse_host_pool_bytes(vllm_config: VllmConfig) -> int: + host_pool_gib = hisparse_host_pool_gib(vllm_config.kv_transfer_config) + if host_pool_gib is None: + raise ValueError("HiSparse requires HiSparseConnector with host_pool_gib") + return int(host_pool_gib * 2**30) + + +def _partition_hisparse_specs( + groups: list[KVCacheGroupSpec], +) -> tuple[dict[str, MLAAttentionSpec], dict[str, MLAAttentionSpec]]: + group_spec = groups[0].kv_cache_spec + if not isinstance(group_spec, UniformTypeKVCacheSpecs): + raise ValueError("HiSparse requires uniform sparse-MLA cache specs.") + + specs = group_spec.kv_cache_specs + if any( + isinstance(spec, MLAAttentionSpec) and spec.model_version == "deepseek_v4" + for spec in specs.values() + ): + raise ValueError("HiSparse does not support DeepSeek V4.") + if not all(isinstance(spec, MLAAttentionSpec) for spec in specs.values()): + raise ValueError("HiSparse requires its first cache group to contain MLA only.") + + source_specs = { + name: spec + for name, spec in specs.items() + if isinstance(spec, MLAAttentionSpec) + and spec.cache_role is SparseCacheRole.SPARSE + } + indexer_specs = { + name: spec + for name, spec in specs.items() + if isinstance(spec, MLAAttentionSpec) + and spec.cache_role is SparseCacheRole.INDEXER + } + if not source_specs or not indexer_specs: + raise ValueError("HiSparse requires sparse-MLA and indexer cache specs.") + return source_specs, indexer_specs + + +def get_hisparse_gpu_memory_usage( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + _, indexer_specs = _partition_hisparse_specs(kv_cache_groups) + return sum( + spec.max_memory_usage_bytes(vllm_config) for spec in indexer_specs.values() + ) + sum( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config) + for group in kv_cache_groups[1:] + ) + + +def create_hisparse_layout( + vllm_config: VllmConfig, + groups: list[KVCacheGroupSpec], + host_budget: int, +) -> HiSparseLayout: + source_specs, indexer_specs = _partition_hisparse_specs(groups) + block_sizes = { + spec.block_size for spec in (*source_specs.values(), *indexer_specs.values()) + } + if len(block_sizes) != 1: + raise ValueError("HiSparse requires one resolved GPU block size.") + gpu_block_size = block_sizes.pop() + + config = ResolvedHiSparseConfig.from_vllm_config( + vllm_config, + vllm_config.model_config.hf_config.index_topk, + gpu_block_size, + ) + assert config is not None + indexer_group_spec = UniformTypeKVCacheSpecs.from_specs( + cast(dict[str, KVCacheSpec], indexer_specs) + ) + assert indexer_group_spec is not None + indexer_group = KVCacheGroupSpec( + list(indexer_specs), + indexer_group_spec, + enable_kv_transfer=True, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + + indexer_page = sum(spec.page_size_bytes for spec in indexer_specs.values()) + hot_blocks_per_request = cdiv(config.device_buffer_size, gpu_block_size) + hot_units: list[list[tuple[str, MLAAttentionSpec]]] = [] + for layer_name, layer_spec in source_specs.items(): + if layer_spec.is_index_group_leader or not hot_units: + hot_units.append([]) + hot_units[-1].append((f"{layer_name}{HISPARSE_HOT_SUFFIX}", layer_spec)) + + resident_groups: list[KVCacheGroupSpec] = [] + hot_groups: list[KVCacheGroupSpec] = [] + + def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None: + page_sizes = {spec.page_size_bytes for _, spec in layers} + if len(page_sizes) != 1: + raise ValueError( + "HiSparse hot-cache groups require one page size, got " + f"{sorted(page_sizes)}." + ) + page_size = page_sizes.pop() + names = [name for name, _ in layers] + resident_groups.append( + KVCacheGroupSpec( + [ + name[: -len(HISPARSE_HOT_SUFFIX)] + HISPARSE_RESIDENT_SUFFIX + for name in names + ], + HiSparseResidentSpec( + block_size=gpu_block_size, + page_size=page_size, + ), + enable_kv_transfer=False, + ) + ) + hot_groups.append( + KVCacheGroupSpec( + names, + HiSparseHotSpec( + block_size=gpu_block_size, + page_size=page_size, + blocks_per_request=hot_blocks_per_request, + ), + enable_kv_transfer=False, + ) + ) + + current: list[tuple[str, MLAAttentionSpec]] = [] + current_page = 0 + for unit in hot_units: + unit_page = sum(spec.page_size_bytes for _, spec in unit) + if current and current_page + unit_page > indexer_page: + append_hot_group(current) + current = [] + current_page = 0 + current.extend(unit) + current_page += unit_page + if current: + append_hot_group(current) + + source_group_spec = UniformTypeKVCacheSpecs.from_specs( + cast(dict[str, KVCacheSpec], source_specs) + ) + assert source_group_spec is not None + source_group = KVCacheGroupSpec( + list(source_specs), + source_group_spec, + role=KVCacheGroupRole.HISPARSE_SOURCE, + host_resident=True, + enable_kv_transfer=True, + ) + regular_groups = groups[1:] + gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups] + + host_num_blocks = host_budget // sum( + spec.page_size_bytes for spec in source_specs.values() + ) + if host_num_blocks <= 0: + raise ValueError("HiSparse has no allocatable host blocks.") + + return HiSparseLayout( + source_group=source_group, + device_groups=gpu_groups, + host_num_blocks=host_num_blocks, + ) + + +def _build_hisparse_kv_cache_tensors( + kv_cache_groups: list[KVCacheGroupSpec], + num_blocks: int, + size: int, + layout: KVCacheLayout, + bytes_per_block: int, + *, + host_resident: bool = False, +) -> list[KVCacheTensor]: + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + tensors: list[KVCacheTensor] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + if isinstance(spec, (HiSparseHotSpec, HiSparseResidentSpec)): + if not layout.is_block_outermost: + raise ValueError("HiSparse requires a block-outermost KV layout.") + layer_stride = spec.page_size_bytes + block_stride = bytes_per_block + else: + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), + ) + offset = ( + byte_offset + * max(layer_stride, spec.page_size_bytes) + // spec.page_size_bytes + ) + tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, + host_resident=host_resident, + ) + ) + byte_offset += len(layer_names) * spec.page_size_bytes + return tensors + + +def get_hisparse_kv_cache_config( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, + host_budget: int, +) -> KVCacheConfig: + from vllm.v1.core.kv_cache_utils import ( + _get_kv_cache_bytes_per_block, + may_override_num_blocks, + validate_kv_cache_layout, + ) + + hisparse_layout = create_hisparse_layout(vllm_config, kv_cache_groups, host_budget) + device_groups = hisparse_layout.device_groups + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, device_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(device_groups) + num_blocks = may_override_num_blocks( + vllm_config, available_memory // bytes_per_block + ) + size = bytes_per_block * num_blocks + kv_cache_tensors = _build_hisparse_kv_cache_tensors( + device_groups, num_blocks, size, layout, bytes_per_block + ) + + host_groups = [hisparse_layout.source_group] + host_layout = KVCacheLayout.LBNHC + validate_kv_cache_layout(host_layout, host_groups) + host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups) + host_size = host_bytes_per_block * hisparse_layout.host_num_blocks + kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors( + host_groups, + hisparse_layout.host_num_blocks, + host_size, + host_layout, + host_bytes_per_block, + host_resident=True, + ) + logger.info_once( + "HiSparse HMA: %.1f GiB host source (%d blocks), %.1f GiB shared " + "GPU indexer/resident/hot pool (%d blocks).", + host_size / 2**30, + hisparse_layout.host_num_blocks, + size / 2**30, + num_blocks, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[*host_groups, *device_groups], + hisparse_host_num_blocks=hisparse_layout.host_num_blocks, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py new file mode 100644 index 000000000000..96235f1570ac --- /dev/null +++ b/vllm/v1/hisparse/runtime.py @@ -0,0 +1,1100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA HiSparse resident, host, and hot-cache data plane.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeAlias + +import psutil +import torch + +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import current_stream +from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup + +# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. +FP8_DS_MLA_ROW_BYTES = 656 +HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + + +def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool: + return ( + attn_metadata is not None + and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens + and attn_metadata.max_query_len == 1 + and attn_metadata.num_reqs == attn_metadata.num_actual_tokens + ) + + +def _get_max_decode_query_len(vllm_config: VllmConfig) -> int: + speculative_config = getattr(vllm_config, "speculative_config", None) + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + return 1 + speculative_config.num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + return 1 + + +def _get_max_swap_rows(vllm_config: VllmConfig) -> int: + max_query_len = _get_max_decode_query_len(vllm_config) + scheduler_config = vllm_config.scheduler_config + return min( + scheduler_config.max_num_batched_tokens, + scheduler_config.max_num_seqs * max_query_len, + ) + + +@dataclass(frozen=True) +class PagedCacheView: + cache: torch.Tensor + attention_cache: torch.Tensor + block_size: int + attention_block_stride: int + + @classmethod + def bind( + cls, + raw_tensor: torch.Tensor, + *, + dtype: torch.dtype, + row_width: int, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + ) -> PagedCacheView: + itemsize = dtype.itemsize + assert byte_offset % itemsize == 0 and block_stride % itemsize == 0 + cache = torch.as_strided( + raw_tensor.view(dtype), + size=(num_blocks, block_size, row_width), + stride=(block_stride // itemsize, row_width, 1), + storage_offset=byte_offset // itemsize, + ) + row_bytes = row_width * itemsize + if block_stride % (block_size * row_bytes): + return cls(cache, cache, block_size, block_size) + attention_cache = ( + raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width) + ) + return cls(cache, attention_cache, block_size, block_stride // row_bytes) + + +@dataclass(frozen=True) +class ResolvedHiSparseConfig: + top_k: int + device_buffer_size: int + eager_host_mirror: bool = True + + @classmethod + def from_vllm_config( + cls, + vllm_config: VllmConfig, + model_top_k: int, + block_size: int | None = None, + ) -> ResolvedHiSparseConfig | None: + config = vllm_config.attention_config.hisparse_config + if config is None: + return None + + max_query_len = _get_max_decode_query_len(vllm_config) + min_device_buffer_size = max_query_len * model_top_k + # Retain the worst-case union of every speculative step's top-k and + # leave one additional top-k of LRU slack. + configured_size = config.device_buffer_size + if configured_size is None: + device_buffer_size = (max_query_len + 1) * model_top_k + else: + device_buffer_size = configured_size + + if device_buffer_size < min_device_buffer_size: + raise ValueError( + "HiSparse device_buffer_size must cover every decode query's " + "index_topk rows. Got " + f"device_buffer_size={device_buffer_size}, " + f"max_decode_query_len={max_query_len}, " + f"index_topk={model_top_k}; expected at least " + f"{min_device_buffer_size}." + ) + max_device_buffer_size = torch.iinfo(torch.int16).max + 1 + if device_buffer_size > max_device_buffer_size: + raise ValueError( + "HiSparse device_buffer_size exceeds the int16 slot-index " + f"limit: got {device_buffer_size}, maximum is " + f"{max_device_buffer_size}." + ) + if configured_size is not None and block_size is not None: + padding = -device_buffer_size % block_size + if padding: + logger.warning( + "HiSparse device_buffer_size=%d is not aligned to the " + "%d-token kernel block size and allocates %d unused " + "padding rows per hot group. Use %d to use all allocated " + "rows.", + device_buffer_size, + block_size, + padding, + device_buffer_size + padding, + ) + return cls( + top_k=model_top_k, + device_buffer_size=device_buffer_size, + eager_host_mirror=config.eager_host_mirror, + ) + + +def check_hisparse_host_memory(pool_bytes: int) -> None: + """Reject a physical host-pool allocation that cannot fit in RAM.""" + mem = psutil.virtual_memory() + if pool_bytes > mem.available * 0.95: + raise ValueError( + f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB " + f"but only {mem.available / 2**30:.0f} GiB of RAM is available. " + "Lower the HiSparseConnector host_pool_gib or leave headroom " + "for co-tenants." + ) + + +def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate and deterministically register an exact-size host KV region.""" + page = 4096 + padded_size = round_up(size, page) + backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu") + aligned_offset = (-backing.data_ptr()) % page + registered = backing[aligned_offset : aligned_offset + padded_size] + pin_tensor(registered) + return registered[:size], registered + + +class HiSparseHostPool: + """Keep the single host backing and its exact CUDA registration range.""" + + def __init__(self) -> None: + self.backing: torch.Tensor | None = None + self.registered: torch.Tensor | None = None + + def allocate(self, size: int) -> torch.Tensor: + assert self.backing is None, "HiSparse host pool is already allocated." + check_hisparse_host_memory(size) + self.backing, self.registered = allocate_pinned_host_pool(size) + return self.backing + + +def release_pinned_state( + runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor] +) -> None: + """Synchronize and release registered host KV pools.""" + if pinned_host_pools: + try: + torch.accelerator.synchronize() + except RuntimeError as e: + logger.warning( + "HiSparse: CUDA context unusable at teardown (%s); leaving " + "%d host-pool tensors pinned for kernel exit reclaim.", + e, + len(pinned_host_pools), + ) + return + + cudart = torch.cuda.cudart() + release_start = time.perf_counter() + freed_bytes = 0 + while pinned_host_pools: + tensor = pinned_host_pools[-1] + err = cudart.cudaHostUnregister(tensor.data_ptr()) + if err.value != 0: + logger.warning( + "HiSparse: cudaHostUnregister failed (code=%d); leaving " + "%d host-pool tensors pinned for kernel exit reclaim.", + err.value, + len(pinned_host_pools), + ) + cudart.cudaGetLastError() + break + freed_bytes += tensor.nbytes + pinned_host_pools.pop() + if freed_bytes: + logger.info( + "HiSparse: unpinned %.1f GiB of host pool in %.1fs.", + freed_bytes / 2**30, + time.perf_counter() - release_start, + ) + + for runtime in runtimes: + del runtime._host_cache + del runtime.registered_host_pool + del runtime.hot_backing + + +def hisparse_prefill_staging_remap( + block_table: torch.Tensor, block_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Renumber blocks without a data-dependent CUDA output allocation.""" + block_ids = torch.cat( + ( + torch.zeros(1, dtype=torch.int32, device=block_table.device), + block_table.flatten().clamp(min=0).to(torch.int32), + ) + ) + sorted_ids, permutation = torch.sort(block_ids) + is_new = torch.cat( + ( + torch.ones(1, dtype=torch.bool, device=block_table.device), + sorted_ids[1:] != sorted_ids[:-1], + ) + ) + compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1 + inverse = torch.empty_like(compact_sorted) + inverse.scatter_(0, permutation, compact_sorted) + + unique_ids = torch.full_like(sorted_ids, -1) + unique_ids.scatter_(0, compact_sorted, sorted_ids) + valid = ( + torch.arange(sorted_ids.numel(), device=block_table.device) + <= compact_sorted[-1] + ) + unique_ids = torch.where(valid, unique_ids, -1) + new_bt = inverse[1:].view_as(block_table) + offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device) + row_ids = ( + torch.where( + valid[:, None], + unique_ids[:, None] * block_size + offsets, + -1, + ) + ).view(1, -1) + return new_bt, row_ids + + +@dataclass +class HiSparsePrefillStagingPlan: + block_table: torch.Tensor + row_ids: torch.Tensor + dst_rows: torch.Tensor + miss_mask: torch.Tensor + block_size: int + # Host rows with a GPU-resident copy (adopted shadow pages): the flat + # resident-cache row to read instead of DMAing from host, -1 for misses. + gpu_row_ids: torch.Tensor | None = None + gpu_source_key: tuple[int, int] | None = None + + def ensure_gpu_sources( + self, + resident_block_table: torch.Tensor, + resident_block_size: int, + ) -> None: + """Resolve which staged rows can be served from the resident cache. + + Computed once per plan (the resident block table is shared by every + layer in the group); non-null resident pages become miss_mask=0 rows + gathered device-to-device by ``gather_prefill_cache``. + """ + source_key = (resident_block_table.data_ptr(), resident_block_size) + if self.gpu_source_key == source_key: + return + block_size = self.block_size + if resident_block_size <= 0 or block_size % resident_block_size != 0: + return + device = self.row_ids.device + num_unique = self.row_ids.shape[1] // block_size + host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size + new_bt = self.block_table.to(torch.int64) + num_rows, num_cols = new_bt.shape + if num_rows == 0 or resident_block_table.shape[0] < num_rows: + return + # One representative (row, col) per unique host block: any request + # referencing the block holds an equivalent (refcounted) resident view. + flat_pos = torch.arange(num_rows * num_cols, device=device) + rep = torch.full( + (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64 + ) + rep.scatter_reduce_( + 0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True + ) + rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1) + rep_col = rep % num_cols + pages_per_block = block_size // resident_block_size + res_pos = rep_col[:, None] * pages_per_block + torch.arange( + pages_per_block, device=device + ) + res_cols = resident_block_table.shape[1] + res_blocks = torch.where( + res_pos < res_cols, + torch.gather( + resident_block_table.to(torch.int64)[rep_row], + 1, + res_pos.clamp(max=max(res_cols - 1, 0)), + ), + torch.zeros_like(res_pos), + ) + offsets = torch.arange(block_size, device=device) + per_off_block = res_blocks[:, offsets // resident_block_size] + gpu_rows = ( + per_off_block * resident_block_size + + (offsets % resident_block_size)[None, :] + ) + # Block id 0 is the null block in both id spaces. + valid_rows = self.row_ids.view(num_unique, block_size) >= 0 + hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows + self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32) + self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int() + self.gpu_source_key = source_key + + +def build_hisparse_prefill_staging_plan( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + staging_block_capacity: int, +) -> HiSparsePrefillStagingPlan: + """Build an asynchronous layer-independent host-cache staging remap.""" + device = block_table.device + used = (seq_lens.to(torch.int64) + block_size - 1) // block_size + packed_offsets = torch.arange(staging_block_capacity, device=device) + row_ends = torch.cumsum(used, dim=0) + rows = torch.searchsorted(row_ends, packed_offsets, right=True) + valid = rows < block_table.shape[0] + rows = rows.clamp(max=block_table.shape[0] - 1) + row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1])) + columns = packed_offsets - row_starts[rows] + valid &= columns < block_table.shape[1] + flat_positions = rows * block_table.shape[1] + columns.clamp( + max=block_table.shape[1] - 1 + ) + referenced = torch.where(valid, block_table.flatten()[flat_positions], 0) + referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size) + scratch_index = block_table.numel() + scatter_positions = torch.where(valid, flat_positions, scratch_index) + remapped = torch.where(valid, referenced_bt.flatten(), 0) + new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device) + new_bt_storage.scatter_(0, scatter_positions, remapped) + new_bt = new_bt_storage[:-1].view_as(block_table) + valid_rows = row_ids >= 0 + dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view( + 1, -1 + ) + return HiSparsePrefillStagingPlan( + block_table=new_bt, + row_ids=row_ids, + dst_rows=dst_rows, + miss_mask=valid_rows.to(torch.int32), + block_size=block_size, + ) + + +def _has_hisparse_ops() -> bool: + if not hasattr(torch.ops, "_C_cache_ops"): + return False + return ( + hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency") + and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots") + and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan") + and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact") + ) + + +class _SharedTopKState: + """CUDA-graph-safe residency result shared by index-sharing layers.""" + + def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None: + self.device_topk_rows = torch.full( + (max_rows, top_k), -1, dtype=torch.int32, device=device + ) + self.physical_topk_indices = torch.empty_like(self.device_topk_rows) + self.swap_host_physical_rows = torch.empty( + (max_rows, top_k), dtype=torch.int32, device=device + ) + self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows) + self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device) + self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device) + + +def _create_shared_topk_state( + device: torch.device, max_rows: int, top_k: int +) -> _SharedTopKState: + return _SharedTopKState(device, max_rows, top_k) + + +def _create_copy_stream(device: torch.device) -> torch.Stream: + return torch.Stream(device=device) + + +class HiSparseIndexGroup: + """GPU index-resolution state shared by one leader and its followers.""" + + leader: HiSparseRuntime + + def __init__( + self, + device: torch.device, + max_num_reqs: int, + region_stride: int, + max_swap_rows: int, + top_k: int, + copy_stream: torch.Stream | None = None, + logical_topk_ready: torch.Event | None = None, + ) -> None: + self.device_global_indices = torch.full( + (max_num_reqs, region_stride), + -1, + dtype=torch.int32, + device=device, + ) + self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device) + self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous() + self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k) + self.copy_stream = ( + copy_stream if copy_stream is not None else _create_copy_stream(device) + ) + self.logical_topk_ready = logical_topk_ready + self.followers: list[HiSparseRuntime] = [] + + +class HiSparseRuntime: + """Per-cache host/hot data plane and GPU replacement state.""" + + hot: PagedCacheView + hot_backing: torch.Tensor + hot_block_table: torch.Tensor + _host_cache: torch.Tensor + registered_host_pool: torch.Tensor + + def __init__( + self, + config: ResolvedHiSparseConfig, + max_num_reqs: int, + row_width: int, + kv_dtype: torch.dtype, + device: torch.device | str, + max_swap_rows: int | None = None, + index_group: HiSparseIndexGroup | None = None, + copy_stream: torch.Stream | None = None, + logical_topk_ready: torch.Event | None = None, + ) -> None: + if not _has_hisparse_ops(): + raise RuntimeError( + "HiSparse requires its compiled _C_cache_ops CUDA kernels " + "(host-resident decode has no Python fallback). Rebuild vLLM " + "from source so " + "csrc/libtorch_stable/hisparse_kernels.cu is included." + ) + self.max_num_reqs = max_num_reqs + self.row_width = row_width + self.kv_dtype = kv_dtype + self.device = torch.device(device) + # Logical slots per request. Physical rows come from its ephemeral HMA + # block table and need not be contiguous in the shared slab. + self.region_stride = config.device_buffer_size + + row_bytes = row_width * kv_dtype.itemsize + if row_bytes % 16 != 0: + raise ValueError( + f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B." + ) + + self._layer_ready_event: torch.Event | None = None + self._swap_staged = False + self._swap_step = 0 + self.is_group_leader = index_group is None + if index_group is None: + index_group = HiSparseIndexGroup( + self.device, + max_num_reqs, + self.region_stride, + max_swap_rows or max_num_reqs, + config.top_k, + copy_stream, + logical_topk_ready, + ) + index_group.leader = self + else: + index_group.followers.append(self) + self.index_group = index_group + + self.eager_host_mirror = config.eager_host_mirror + self.resident_source_index = -1 + self.request_state_indices: torch.Tensor | None = None + + @property + def host_cache(self) -> torch.Tensor: + return self._host_cache + + def bind_hot_cache( + self, + raw_tensor: torch.Tensor, + *, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + block_table: torch.Tensor, + ) -> None: + """Bind this runtime's strided view into the shared GPU HMA slab.""" + self.hot = PagedCacheView.bind( + raw_tensor, + dtype=self.kv_dtype, + row_width=self.row_width, + byte_offset=byte_offset, + block_stride=block_stride, + num_blocks=num_blocks, + block_size=block_size, + ) + self.hot_backing = raw_tensor + self.hot_block_table = block_table + + def bind_source_cache( + self, + kv_cache: torch.Tensor, + *, + registered_host_pool: torch.Tensor | None = None, + ) -> None: + if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width: + raise ValueError( + "HiSparse runtime bound to a KV cache with mismatched " + f"layout: expected ({self.row_width}, {self.kv_dtype}), got " + f"({kv_cache.shape[-1]}, {kv_cache.dtype})." + ) + # Host-resident pool: the cache itself is the only full-size store; + # every allocated slot is written before the indexer can select it. + if kv_cache.device.type != "cpu": + raise ValueError( + "HiSparse requires a host-resident KV pool; got a KV cache " + f"on {kv_cache.device}." + ) + # The pool is pinned via cudaHostRegister (exact-size, deterministic + # unpin at shutdown); torch's is_pinned() only recognizes its own + # caching-host-allocator memory, so also accept ranges the model + # allocator explicitly registered. + if not (kv_cache.is_pinned() or registered_host_pool is not None): + raise ValueError("HiSparse host-resident KV pool must be pinned memory.") + + self._host_cache = kv_cache.view(-1, kv_cache.shape[-1]) + self.registered_host_pool = ( + registered_host_pool if registered_host_pool is not None else kv_cache + ) + + def gather_prefill_cache( + self, + kv_cache: torch.Tensor, + plan: HiSparsePrefillStagingPlan, + resident_cache: torch.Tensor | None = None, + ) -> torch.Tensor: + """Gather this runtime's host cache using a shared staging plan.""" + if kv_cache.shape[1] != plan.block_size: + raise ValueError( + f"HiSparse staging block size {plan.block_size} does not match " + f"the KV cache block size {kv_cache.shape[1]}." + ) + row_width = kv_cache.shape[-1] + + staged = torch.empty( + ( + plan.row_ids.shape[1] // plan.block_size, + plan.block_size, + row_width, + ), + dtype=kv_cache.dtype, + device=plan.block_table.device, + ) + if plan.gpu_row_ids is not None and resident_cache is not None: + gpu_rows = plan.gpu_row_ids[0].to(torch.long) + src = gpu_rows.clamp_min(0) + resident_block_size = resident_cache.shape[1] + rows = resident_cache[src // resident_block_size, src % resident_block_size] + if rows.dtype != staged.dtype: + rows = rows.contiguous().view(staged.dtype) + staged.view(-1, row_width).copy_(rows) + torch.ops._C_cache_ops.hisparse_gather_plan( + kv_cache.view(-1, row_width), + staged, + plan.row_ids, + plan.dst_rows, + plan.miss_mask, + None, + None, + 0, + ) + return staged + + def reset_hot_state(self) -> None: + """Drop all hot-buffer bookkeeping (hits become misses).""" + group = self.index_group + compute_stream = current_stream() + compute_stream.wait_stream(group.copy_stream) + group.device_global_indices.fill_(-1) + group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots)) + group.copy_stream.wait_stream(compute_stream) + + def invalidate_sorted_slots( + self, + sorted_slots: torch.Tensor, + state_indices: torch.Tensor, + ) -> None: + """Drop hot copies using shared, preprocessed invalidation inputs.""" + device_global_indices = self.index_group.device_global_indices + active_indices = device_global_indices.index_select(0, state_indices) + positions = torch.searchsorted(sorted_slots, active_indices) + positions.clamp_(max=sorted_slots.numel() - 1) + active_indices.masked_fill_( + sorted_slots[positions] == active_indices, + -1, + ) + device_global_indices.index_copy_(0, state_indices, active_indices) + + def invalidate_written_slots( + self, + written_slots: torch.Tensor, + req_id_per_token: torch.Tensor, + ) -> None: + num_tokens = min(written_slots.numel(), req_id_per_token.numel()) + if num_tokens == 0: + return + assert self.request_state_indices is not None + torch.ops._C_cache_ops.hisparse_invalidate_written_slots( + self.index_group.device_global_indices, + self.request_state_indices, + req_id_per_token[:num_tokens], + written_slots[:num_tokens], + ) + + def begin_forward(self) -> None: + self._swap_step = 0 + + def _step_rows(self, num_tokens: int) -> slice: + start = self._swap_step * num_tokens + stop = start + num_tokens + if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]: + raise ValueError( + "HiSparse swap rows exceed the configured speculative decode " + f"capacity: stop={stop}, capacity=" + f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}." + ) + self._swap_step += 1 + return slice(start, stop) + + def _resolve_residency( + self, + *, + resident: HiSparseCacheHandle | None = None, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool = False, + shared_rows: slice, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> None: + """Resolve logical top-k positions and compact rows requiring swaps.""" + num_tokens = topk_indices.shape[0] + hot = self.hot + host_cache = self.host_cache + assert hot.block_size == block_size + request_state_indices = self.request_state_indices + assert request_state_indices is not None + group = self.index_group + shared = group.shared_topk + + device_topk_rows = shared.device_topk_rows[shared_rows] + valid_topk_counts = None + if return_valid_counts: + valid_topk_counts = ( + shared.valid_topk_counts[shared_rows] + if valid_counts_out is None + else valid_counts_out + ) + physical_topk_indices = ( + shared.physical_topk_indices[shared_rows] + if attention_indices_out is None + else attention_indices_out + ) + swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows] + swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows] + swap_counts = shared.swap_counts[shared_rows] + + # Padded rows are skipped by the kernel (request_state_indices) and must + # come out as -1 so the attention kernel masks them. + torch.ops._C_cache_ops.hisparse_resolve_residency( + host_cache, + hot.cache, + self.hot_block_table, + topk_indices, + device_topk_rows, + group.device_global_indices, + group.lru_slots, + request_state_indices, + self.region_stride, + None, + physical_topk_indices, + hot.attention_block_stride, + req_id_per_token[:num_tokens].contiguous(), + block_table, + block_size, + None, + valid_topk_counts, + swap_host_physical_rows, + swap_device_physical_rows, + swap_counts, + resident.block_table if resident is not None else None, + resident.view.block_size + if resident is not None and resident.view is not None + else 0, + 0, + ) + + def _swap_rows(self, shared_rows: slice) -> None: + hot = self.hot + shared = self.index_group.shared_topk + torch.ops._C_cache_ops.hisparse_gather_compact( + self.host_cache, + hot.cache, + shared.swap_host_physical_rows[shared_rows], + shared.swap_device_physical_rows[shared_rows], + shared.swap_counts[shared_rows], + ) + + def _resolve_and_stage_group( + self, + *, + resident: HiSparseCacheHandle, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool, + shared_rows: slice, + attention_indices_out: torch.Tensor | None, + valid_counts_out: torch.Tensor | None, + ) -> None: + group = self.index_group + compute_stream = current_stream() + if group.logical_topk_ready is not None: + group.copy_stream.wait_event(group.logical_topk_ready) + else: + group.copy_stream.wait_stream(compute_stream) + with group.copy_stream: + self._resolve_residency( + resident=resident, + req_id_per_token=req_id_per_token, + block_table=block_table, + topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + shared_rows=shared_rows, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + runtimes = [self] + if resident.decode_batch: + runtimes.extend(group.followers) + for runtime in runtimes: + runtime._swap_rows(shared_rows) + if runtime._layer_ready_event is None: + runtime._layer_ready_event = torch.Event() + runtime._layer_ready_event.record(group.copy_stream) + runtime._swap_staged = True + + def _stage_this_layer(self, shared_rows: slice) -> None: + group = self.index_group + compute_stream = current_stream() + group.copy_stream.wait_stream(compute_stream) + with group.copy_stream: + self._swap_rows(shared_rows) + if self._layer_ready_event is None: + self._layer_ready_event = torch.Event() + self._layer_ready_event.record(group.copy_stream) + self._swap_staged = True + + def wait_for_staged_swap(self) -> None: + if not self._swap_staged or self._layer_ready_event is None: + raise RuntimeError("HiSparse layer swap was not prefetched.") + current_stream().wait_event(self._layer_ready_event) + self._swap_staged = False + + def swap_in( + self, + *, + resident: HiSparseCacheHandle, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool = False, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> HiSparseTopKResult: + """Resolve once per group and ensure this layer's rows are available.""" + num_tokens = logical_topk_indices.shape[0] + shared_rows = self._step_rows(num_tokens) + group = self.index_group + hot = self.hot + assert hot.block_size == block_size + assert hot.attention_block_stride == group.leader.hot.attention_block_stride + + if self.is_group_leader: + self._resolve_and_stage_group( + resident=resident, + req_id_per_token=req_id_per_token, + block_table=block_table, + logical_topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + shared_rows=shared_rows, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + + if not self._swap_staged: + self._stage_this_layer(shared_rows) + self.wait_for_staged_swap() + + physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows] + if return_valid_counts: + return ( + physical_topk_indices, + group.shared_topk.valid_topk_counts[shared_rows], + ) + return physical_topk_indices + + +class HiSparseCacheHandle: + """Attention-facing handle for resident KV and sparse offload state.""" + + def __init__(self, runtime: HiSparseRuntime) -> None: + self.view: PagedCacheView | None = None + self.block_table: torch.Tensor | None = None + self.source_block_table: torch.Tensor | None = None + self.slot_mapping: torch.Tensor | None = None + self.runtime = runtime + self.dummy_batch = False + self.decode_batch = False + self.all_context_pages_resident = True + self.num_actual_tokens = 0 + self.num_decode_tokens = 0 + self.req_id_per_token: torch.Tensor | None = None + self.host_mirror_required = False + self.mirror_from_resident = False + self.mirror_slot_mapping: torch.Tensor | None = None + self.mirror_staging_cache: torch.Tensor | None = None + self.mirror_staging_slots: torch.Tensor | None = None + self.submit_layer_mirror: Callable[[], None] | None = None + self.index_group_caches: list[HiSparseCacheHandle] = [self] + + def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: + assert self.runtime.is_group_leader + for cache in self.index_group_caches: + cache._prepare_for_batch(attn_metadata) + + def _prepare_for_batch(self, attn_metadata: Any | None) -> None: + self.dummy_batch = attn_metadata is None + self.runtime.begin_forward() + self.num_actual_tokens = ( + attn_metadata.num_actual_tokens if attn_metadata is not None else 0 + ) + self.num_decode_tokens = ( + attn_metadata.num_decode_tokens if attn_metadata is not None else 0 + ) + self.req_id_per_token = ( + attn_metadata.req_id_per_token if attn_metadata is not None else None + ) + self.decode_batch = is_hisparse_decode_batch(attn_metadata) + self.host_mirror_required = attn_metadata is not None and ( + not self.decode_batch or self.runtime.eager_host_mirror + ) + + def write_target( + self, num_input_rows: int, num_slot_rows: int + ) -> tuple[torch.Tensor, torch.Tensor, int]: + assert self.view is not None and self.slot_mapping is not None + num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens) + if self.decode_batch: + num_rows = min(num_rows, self.runtime.max_num_reqs) + return self.view.cache, self.slot_mapping[:num_rows], num_rows + + def mirror_write_target( + self, num_rows: int + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if ( + not self.host_mirror_required + or self.decode_batch + or self.mirror_from_resident + ): + return None + cache = self.mirror_staging_cache + slots = self.mirror_staging_slots + if cache is None or slots is None: + raise RuntimeError("HiSparse prefill mirror staging is not bound.") + if num_rows > slots.numel(): + raise RuntimeError( + "HiSparse prefill mirror exceeds staging capacity: " + f"{num_rows} > {slots.numel()}." + ) + return cache, slots[:num_rows] + + def finish_kv_update(self) -> None: + if self.dummy_batch: + return + if ( + self.submit_layer_mirror is not None + and not self.decode_batch + and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL + ): + self.submit_layer_mirror() + + def bind_cache( + self, + raw_tensor: torch.Tensor, + *, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + self.view = PagedCacheView.bind( + raw_tensor, + dtype=self.runtime.kv_dtype, + row_width=self.runtime.row_width, + byte_offset=byte_offset, + block_stride=block_stride, + num_blocks=num_blocks, + block_size=block_size, + ) + self.block_table = block_table + self.slot_mapping = slot_mapping + + def swap_in( + self, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + *, + block_size: int, + return_valid_counts: bool = False, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> HiSparseTopKResult: + return self.runtime.swap_in( + resident=self, + req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]], + block_table=block_table, + logical_topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + + +def initialize_hisparse_runtime_buffers( + cache_handles: list[HiSparseCacheHandle], + *, + max_num_reqs: int, + max_num_batched_tokens: int, +) -> None: + """Allocate shared request state and per-layer prefill staging buffers.""" + assert cache_handles + resident = cache_handles[0].view + assert resident is not None + device = resident.cache.device + request_state_indices = torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ) + staging_blocks = ( + max_num_batched_tokens + resident.block_size - 1 + ) // resident.block_size + mirror_staging_caches = torch.empty( + ( + len(cache_handles), + staging_blocks, + resident.block_size, + resident.cache.shape[-1], + ), + dtype=resident.cache.dtype, + device=device, + ) + mirror_staging_slots = torch.arange( + max_num_batched_tokens, dtype=torch.int64, device=device + ) + for layer_index, cache_handle in enumerate(cache_handles): + cache_handle.runtime.request_state_indices = request_state_indices + cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index] + cache_handle.mirror_staging_slots = mirror_staging_slots + + +def create_hisparse_cache_handle( + vllm_config: VllmConfig, + model_top_k: int, + *, + is_index_group_leader: bool, + row_width: int, + kv_dtype: torch.dtype, + index_group: HiSparseMLAIndexGroup | None = None, + device: torch.device | str | None = None, +) -> HiSparseCacheHandle | None: + config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k) + if config is None: + return None + + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + max_swap_rows = _get_max_swap_rows(vllm_config) + if device is None: + device = torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ) + + hisparse_group = ( + None + if is_index_group_leader or index_group is None + else index_group.hisparse_group + ) + runtime = HiSparseRuntime( + config=config, + max_num_reqs=max_num_reqs, + max_swap_rows=max_swap_rows, + row_width=row_width, + kv_dtype=kv_dtype, + device=device, + index_group=hisparse_group, + copy_stream=index_group.side_stream if index_group is not None else None, + logical_topk_ready=( + index_group.logical_topk_ready + if index_group is not None and index_group.has_indexer + else None + ), + ) + if is_index_group_leader and index_group is not None: + index_group.hisparse_group = runtime.index_group + logger.info_once( + "Enabled experimental HiSparse HMA hot cache: top_k=%d, " + "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.", + config.top_k, + config.device_buffer_size, + config.device_buffer_size, + max_num_reqs, + ) + handle = HiSparseCacheHandle(runtime) + speculative_config = vllm_config.speculative_config + handle.mirror_from_resident = bool( + vllm_config.scheduler_config.async_scheduling + and speculative_config is not None + and speculative_config.uses_draft_kv_cache() + ) + return handle diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py new file mode 100644 index 000000000000..8618f3cb20a4 --- /dev/null +++ b/vllm/v1/hisparse/types.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SparseKVPageTransfer: + """Copy one logical KV page between cache-manager and worker-owned tiers.""" + + transfer_id: int + destination_block_id: int + source_block_ids: tuple[int, ...] + after_forward: bool + + +@dataclass(frozen=True) +class SparseKVRowMirror: + """Mirror one contiguous resident-row span into the host cache.""" + + source_starts: tuple[int, ...] + destination_start: int + num_rows: int + + +@dataclass +class SparseKVOffloadCommand: + """Opaque scheduler-to-worker command for one model step.""" + + page_transfers: list[SparseKVPageTransfer] diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 1888404a573a..fb839bcd6855 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -147,6 +147,11 @@ class KVCacheSpecKind(str, Enum): UNKNOWN = "unknown" +class SparseCacheRole(str, Enum): + SPARSE = "sparse" + INDEXER = "indexer" + + @dataclass(frozen=True) class KVCacheSpec: """ @@ -192,6 +197,10 @@ def get_num_kernel_states(self, kernel_block_size: int) -> int: return kernel_block_size // self.tokens_per_state return 1 + @property + def block_table_token_alignment(self) -> int | None: + return 128 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: """ The maximum possible memory usage of this KV cache in bytes. @@ -245,6 +254,25 @@ def is_uniform_with_collection( isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() ) + @property + def has_layer_views(self) -> bool: + """Whether generic allocation creates per-layer cache views. + + Specs without a per-layer shape keep the raw backing tensor and lay it + out themselves. They also have no attention module of their own, so + they take no part in attention-backend or kernel-block selection. + """ + return True + + @property + def uses_slot_mapping(self) -> bool: + """Whether the worker computes a per-token slot mapping for this spec. + + Specs that address their pages themselves (raw storage, ring buffers) + take no slot mapping row. + """ + return True + def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor: """View a kernel-block-granular layer cache with manager blocks as dim 0. @@ -383,6 +411,70 @@ def create_kv_cache_views( return views +@dataclass(frozen=True, kw_only=True) +class HiSparseHotSpec(KVCacheSpec): + """Ephemeral per-request HiSparse hot-cache allocation.""" + + page_size: int + blocks_per_request: int + + @property + def prefix_cacheable(self) -> bool: + return False + + @property + def page_size_bytes(self) -> int: + return self.page_size + + @property + def block_table_token_alignment(self) -> int | None: + return None + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + return self.blocks_per_request * self.page_size + + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + return self.blocks_per_request + + @property + def has_layer_views(self) -> bool: + return False + + @property + def uses_slot_mapping(self) -> bool: + return False + + +@dataclass(frozen=True, kw_only=True) +class HiSparseResidentSpec(KVCacheSpec): + """Reclaimable GPU-resident pages for host-backed HiSparse KV.""" + + page_size: int + + @property + def prefix_cacheable(self) -> bool: + return False + + @property + def page_size_bytes(self) -> int: + return self.page_size + + def max_admission_blocks_per_request( + self, max_in_flight_tokens: int, max_model_len: int + ) -> int: + num_tokens = min(max_in_flight_tokens, max_model_len) + return cdiv(num_tokens, self.block_size) + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + return cdiv(vllm_config.model_config.max_model_len, self.block_size) * ( + self.page_size + ) + + @property + def has_layer_views(self) -> bool: + return False + + @dataclass(frozen=True, kw_only=True) class AttentionSpec(KVCacheSpec): num_kv_heads: int @@ -557,6 +649,8 @@ class MLAAttentionSpec(FullAttentionSpec): # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. alignment: int | None = None # Default to None for no padding. model_version: str | None = None + cache_role: SparseCacheRole = SparseCacheRole.SPARSE + is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" # Group capability enabled when any member flattens a non-causal query block @@ -577,16 +671,20 @@ def merge(cls, specs: list[Self]) -> Self: cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) tokens_per_state_set = set(spec.tokens_per_state for spec in specs) model_version_set = set(spec.model_version for spec in specs) + cache_role_set = {spec.cache_role for spec in specs} + index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 + and len(cache_role_set) == 1 + and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, tokens per state, model version, and storage " - "block size." + "quantization method, tokens per state, model version, cache role, " + "index-sharing role, and storage block size." ) merged_spec = cls( block_size=specs[0].block_size, @@ -600,6 +698,8 @@ def merge(cls, specs: list[Self]) -> Self: cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), + cache_role=cache_role_set.pop(), + is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs @@ -772,6 +872,10 @@ class CircularBufferSpec(AttentionSpec): reads the open group's committed keys from the ring. """ + @property + def block_table_token_alignment(self) -> int | None: + return None + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: # The ring occupies one block per request for its whole lifetime. del vllm_config @@ -792,6 +896,10 @@ def is_uniform_with_collection( def prefix_cacheable(self) -> bool: return False + @property + def uses_slot_mapping(self) -> bool: + return False + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): @@ -919,6 +1027,10 @@ def page_size_bytes(self) -> int: return self.page_size_padded return page_size + @property + def block_table_token_alignment(self) -> int | None: + return None + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: if vllm_config.cache_config.mamba_cache_mode == "all": max_model_len = vllm_config.model_config.max_model_len @@ -1093,6 +1205,10 @@ def first_spec(self) -> KVCacheSpec: def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) + @property + def block_table_token_alignment(self) -> int | None: + return self.first_spec.block_table_token_alignment + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_num_pages = max( cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) @@ -1154,6 +1270,10 @@ def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int: for spec in self.kv_cache_specs.values() ) + @property + def uses_slot_mapping(self) -> bool: + return self.first_spec.uses_slot_mapping + def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]: """The per-layer specs a KV cache group spec covers. @@ -1259,6 +1379,13 @@ class KVCacheTensor: layer_stride: int block_stride: int offset: int = 0 # byte offset of layers[0]'s block 0 + host_resident: bool = False + + +class KVCacheGroupRole(str, Enum): + DEFAULT = "default" + HISPARSE_SOURCE = "hisparse_source" + HISPARSE_INDEXER = "hisparse_indexer" @dataclass @@ -1274,8 +1401,13 @@ class KVCacheGroupSpec: kv_cache_spec: KVCacheSpec # Whether this group contains EAGLE/MTP draft attention layers. is_eagle_group: bool = False + # Host groups use the dedicated HiSparse pool; others share the device pool. + host_resident: bool = False # Whether this group is part of the externally transferable KV state. + # Ephemeral accelerator-side replicas are rebuilt from their transferable + # source group after a connector load. enable_kv_transfer: bool = True + role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT @dataclass @@ -1300,6 +1432,8 @@ class KVCacheConfig: """Resolved retention policy for local prefix-cache checkpoints.""" kv_cache_layout: str | None = None """The KV cache layout resolved by the engine core, adopted by all workers.""" + hisparse_host_num_blocks: int | None = None + """Capacity of the dedicated HiSparse host-block manager, when enabled.""" @cached_property def transfer_group_ids(self) -> tuple[int, ...]: @@ -1337,6 +1471,21 @@ def select_transfer_block_ids( ) return tuple(block_ids[group_id] for group_id in self.transfer_group_ids) + @property + def host_group_ids(self) -> tuple[int, ...]: + return tuple( + group_id + for group_id, group in enumerate(self.kv_cache_groups) + if group.host_resident + ) + + def num_blocks_of(self, tensor: KVCacheTensor) -> int: + """Number of blocks addressable by the pool backing ``tensor``.""" + if not tensor.host_resident: + return self.num_blocks + assert self.hisparse_host_num_blocks is not None + return self.hisparse_host_num_blocks + @property def has_mamba_layers(self) -> bool: return any( @@ -1347,9 +1496,11 @@ def has_mamba_layers(self) -> bool: @property def has_mixed_precision_kv_cache(self) -> bool: - """Whether attention groups store their KV cache at more than one precision.""" + """Whether device attention caches use more than one precision.""" kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set() for group in self.kv_cache_groups: + if group.host_resident: + continue kv_cache_precisions.update( (spec.dtype, spec.kv_quant_mode) for spec in iter_layer_specs(group.kv_cache_spec) diff --git a/vllm/v1/kv_cache_spec_registry.py b/vllm/v1/kv_cache_spec_registry.py index 816a6862dae2..62cb291892a2 100644 --- a/vllm/v1/kv_cache_spec_registry.py +++ b/vllm/v1/kv_cache_spec_registry.py @@ -34,6 +34,7 @@ class KVCacheSpecMetadata: _REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {} +_REGISTRY_ROLE_MANAGERS: dict[str, type["SingleTypeKVCacheManager"]] = {} class KVCacheSpecRegistry: @@ -102,20 +103,34 @@ def register( uniform_type_base_spec=uniform_type_base_spec, ) + @classmethod + def register_role_manager( + cls, role: str, manager_class: type["SingleTypeKVCacheManager"] + ) -> None: + """Select ``manager_class`` for every group with the given role.""" + registered = _REGISTRY_ROLE_MANAGERS.get(role) + assert registered is None or registered is manager_class, ( + f"Conflicting manager registration for group role {role}" + ) + _REGISTRY_ROLE_MANAGERS[role] = manager_class + @classmethod def get_manager_class( - cls, kvcache_spec: "KVCacheSpec" + cls, kvcache_spec: "KVCacheSpec", role: str | None = None ) -> type["SingleTypeKVCacheManager"] | None: """ Get the single type kvcache manager class for a given kvcache spec instance. Args: kvcache_spec: A KVCacheSpec instance + role: The group's role; a role-registered manager takes precedence. Returns: The SingleTypeKVCacheManager class to use for this kvcache_spec """ cls._ensure_registered() + if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)): + return manager_class kvcache_spec_cls = type(kvcache_spec) # Walk up the MRO to find a registered base class diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index 637d3bb664dc..74b619dc0f8a 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -14,6 +14,8 @@ class OffloadingGroupConfig: tokens_per_block: int # Layer names belonging to this group. layer_names: tuple[str, ...] + # Original KVCacheConfig group index. + group_id: int @dataclass(frozen=True) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index e94380825f0e..d8924b273881 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -281,6 +281,9 @@ class KVConnectorOutput: # IDs of externally computed KV blocks that failed to load. # Requests referencing these blocks should be rescheduled to recompute them invalid_block_ids: set[int] = field(default_factory=set) + # Receive failures keyed by request identity. This remains unambiguous for + # hybrid/multi-pool cache layouts where numeric block IDs overlap. + failed_recving: set[str] = field(default_factory=set) # Configuration describing how many finished sending/receiving # notifications should be expected for each request. This allows # handshake-based connectors like Nixl to update the KVOutputAggregator. @@ -295,6 +298,7 @@ def is_empty(self): and not self.kv_connector_stats and not self.kv_cache_events and not self.invalid_block_ids + and not self.failed_recving and not self.kv_connector_worker_meta ) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 4a859924256e..7e7a212b5fc5 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -18,6 +18,10 @@ CommonAttentionMetadata, ) from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend +from vllm.v1.hisparse.binding import ( + init_hisparse_kv_cache, + resolve_hisparse_block_size, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, @@ -30,11 +34,12 @@ AttentionGroup, add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, - bind_kv_cache, + bind_kv_cache_to_layers, prepare_kernel_block_sizes, ) if TYPE_CHECKING: + from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, @@ -132,6 +137,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if isinstance(spec, AttentionSpec): spec = attn_module.get_attn_backend().customize_spec(spec) kv_cache_spec[layer_name] = spec + resolve_hisparse_block_size(vllm_config, kv_cache_spec, attn_layers) return kv_cache_spec @@ -198,6 +204,9 @@ def init_attn_backend( kv_cache_config.kv_cache_groups ): layer_names = kv_cache_group_spec.layer_names + if not kv_cache_group_spec.kv_cache_spec.has_layer_views: + attn_groups.append([]) + continue if active_layer_names is not None: layer_names = list(active_layer_names.intersection(layer_names)) @@ -316,22 +325,34 @@ def get_query_lens_mismatch_unsupported_backend( def init_kv_cache( - runner_kv_caches: list[torch.Tensor | list[torch.Tensor]], forward_context: dict[str, Any], kv_cache_config: KVCacheConfig, device: torch.device, kernel_block_sizes: list[int], vllm_config: VllmConfig, kv_cache_allocation_context: AbstractContextManager | None = None, + *, + block_tables: "BlockTables | None" = None, ) -> dict[str, Any]: allocation_context = kv_cache_allocation_context or nullcontext() with allocation_context: - kv_caches = allocate_kv_cache( - kv_cache_config, - device, - vllm_config.cache_config.get_resolved_kv_cache_layout(), - kernel_block_sizes, - ) + if vllm_config.attention_config.hisparse_config is not None: + assert block_tables is not None + kv_caches = init_hisparse_kv_cache( + kv_cache_config, + device, + kernel_block_sizes, + vllm_config, + forward_context, + block_tables, + ) + else: + kv_caches = allocate_kv_cache( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + ) for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): kv_caches[layer_name] = kv_caches[target] # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per @@ -342,10 +363,12 @@ def init_kv_cache( in ("longcat_flash", "longcat_flash_ngram") else 1 ) - bind_kv_cache( - kv_caches, + bindable_caches = { + name: cache for name, cache in kv_caches.items() if name in forward_context + } + bind_kv_cache_to_layers( + bindable_caches, forward_context, - runner_kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, ) @@ -396,6 +419,8 @@ def build_attn_metadata( attn_metadata: dict[str, Any] = {} num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): + if not attn_groups[i]: + continue block_table = block_tables[i] slot_mapping = slot_mappings[i] # Per-group causal for hybrid drafters (mixed SWA/full attention). diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 287e64ce300c..12354f3c229e 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -34,6 +34,7 @@ from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import current_stream +from vllm.v1.hisparse.binding import release_hisparse_profiling_cache from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -931,6 +932,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: if (model := getattr(runner, "model", None)) is not None: layers = itertools.chain(layers, model.modules()) clear_layer_kv_caches(layers) + release_hisparse_profiling_cache(runner.compilation_config.static_forward_context) runner.cache_config.num_gpu_blocks = None runner.maybe_remove_all_loras(runner.lora_config) gc.collect() diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index eaef0c291661..88b7b0cf9d29 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -29,7 +29,10 @@ class KVConnector: """KVConnector interface used by GPUModelRunner.""" - def pre_forward(self, scheduler_output: "SchedulerOutput") -> None: + def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + pass + + def finish_forward(self) -> None: pass def post_forward( @@ -43,6 +46,9 @@ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: def set_disabled(self, disabled: bool) -> None: pass + def reset_capture_state(self) -> None: + pass + class ActiveKVConnector(KVConnector): def __init__( @@ -54,10 +60,10 @@ def __init__( self.kv_connector.register_kv_caches(kv_caches_dict) self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks) - self._pending_load_start = False + self._pending_load_kwargs: dict[str, Any] | None = None self._disabled = False - def pre_forward(self, scheduler_output: "SchedulerOutput") -> None: + def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: if self._disabled: return @@ -65,23 +71,31 @@ def pre_forward(self, scheduler_output: "SchedulerOutput") -> None: assert kv_connector_metadata is not None self.kv_connector.handle_preemptions(kv_connector_metadata) self.kv_connector.bind_connector_metadata(kv_connector_metadata) + self._pending_load_kwargs = kwargs if scheduler_output.has_sync_kv_loads: # Sync loads need to run before this step's forward. self._start_load_kv() - else: - # Start any async loads in post-forward instead, keeping - # their host-side submission cost off the critical path. - self._pending_load_start = True + # Otherwise start the async load after forward to keep submission + # off the critical path. def _start_load_kv(self) -> None: - self._pending_load_start = False + load_kwargs = self._pending_load_kwargs + assert load_kwargs is not None + self._pending_load_kwargs = None # TODO: sort out KV Connectors' use of forward_context if is_forward_context_available(): - self.kv_connector.start_load_kv(get_forward_context()) + self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs) else: with set_forward_context(None, self.vllm_config): - self.kv_connector.start_load_kv(get_forward_context()) + self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs) + + def finish_forward(self) -> None: + if not self._disabled: + self.kv_connector.finish_forward() + + def reset_capture_state(self) -> None: + self.kv_connector.reset_capture_state() def post_forward( self, finished_req_ids: set[str], wait_for_save: bool = True @@ -89,15 +103,16 @@ def post_forward( if self._disabled: return None - if self._pending_load_start: + if self._pending_load_kwargs is not None: self._start_load_kv() output = KVConnectorOutput() if wait_for_save: self.kv_connector.wait_for_save() - output.finished_sending, output.finished_recving = ( - self.kv_connector.get_finished(finished_req_ids) - ) + transfer_results = self.kv_connector.get_transfer_results(finished_req_ids) + output.finished_sending = transfer_results.finished_sending or None + output.finished_recving = transfer_results.finished_recving or None + output.failed_recving = transfer_results.failed_recving output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors() output.kv_connector_stats = self.kv_connector.get_kv_connector_stats() output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events() @@ -112,6 +127,7 @@ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: return EMPTY_MODEL_RUNNER_OUTPUT self.pre_forward(scheduler_output) + self.finish_forward() finished_req_ids = scheduler_output.finished_req_ids kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False) return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fcadf9e3032b..14e3823c0410 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -67,7 +67,6 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( - CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -590,7 +589,7 @@ def initialize_kv_cache( layer_spec = ( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) - slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec)) + slot_mapping_enabled.append(layer_spec.uses_slot_mapping) # Let each cache type account for CP. Attention KV is DCP-sharded, # while Mamba/GDN recurrent state is replicated across DCP ranks. max_num_blocks = spec.max_num_blocks_per_req( @@ -598,12 +597,11 @@ def initialize_kv_cache( ) # Preserve each cache type's alignment requirements after applying # its topology-aware block-table width. - if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)): - max_num_blocks = get_block_table_width( - max_num_blocks, spec.block_size, token_alignment=None - ) - else: - max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size) + max_num_blocks = get_block_table_width( + max_num_blocks, + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) max_num_blocks_per_group.append(max_num_blocks) target_attn_layer_names = None @@ -722,18 +720,20 @@ def initialize_kv_cache( # to its own attention support. self.speculator.init_cudagraph_manager(cudagraph_mode) - self.kv_caches: list[torch.Tensor] = [] # Capture warmup providers that depend on allocated KV-cache strides. with self.jit_warmup_registry.activate(): kv_caches_dict = init_kv_cache( - self.kv_caches, self.compilation_config.static_forward_context, self.kv_cache_config, self.device, self.kernel_block_sizes, self.vllm_config, kv_cache_allocation_context=kv_cache_allocation_context, + block_tables=self.block_tables, ) + self.kv_caches = [ + cache for cache in kv_caches_dict.values() if cache.device == self.device + ] if is_profiling: self.kv_connector = NO_OP_KV_CONNECTOR else: @@ -1014,6 +1014,7 @@ def capture_model(self, *, profile_only: bool = False) -> int: ): self._dummy_run(**batch) self.adaptive_verification.set_initial_cost_curves(timings) + self.kv_connector.reset_capture_state() end_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -1141,13 +1142,17 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: def update_requests(self, scheduler_output: SchedulerOutput) -> None: # Add new blocks and update num_computed_tokens for the existing requests. reqs = scheduler_output.scheduled_cached_reqs + table_updates = scheduler_output.block_table_updates or {} + for req_id, block_ids in table_updates.items(): + if (idx := self.req_states.req_id_to_index.get(req_id)) is not None: + self.block_tables.append_block_ids(idx, block_ids, overwrite=True) num_computed_tokens_np = self.req_states.num_computed_tokens_np for req_id, num_computed_tokens, req_new_block_ids in zip( reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids ): req_index = self.req_states.req_id_to_index[req_id] num_computed_tokens_np[req_index] = num_computed_tokens - if req_new_block_ids is not None: + if req_new_block_ids is not None and req_id not in table_updates: self.block_tables.append_block_ids( req_index, req_new_block_ids, overwrite=False ) @@ -1841,13 +1846,22 @@ def execute_model( ) self.step_timing.forward_start() + connector_kwargs = dict( + scheduler_output=scheduler_output, + request_state_indices=input_batch.idx_mapping, + request_ids=input_batch.req_ids, + num_tokens=input_batch.num_tokens, + ) + # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. assert self.cudagraph_manager is not None - self.kv_connector.pre_forward(scheduler_output) + self.kv_connector.pre_forward( + **connector_kwargs, attn_metadata=attn_metadata + ) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: # For piecewise and eager mode, just call model(). @@ -1871,7 +1885,7 @@ def execute_model( skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): - self.kv_connector.pre_forward(scheduler_output) + self.kv_connector.pre_forward(**connector_kwargs) if ubatch_state is not None: assert self.ubatch_runner is not None model_output = self.ubatch_runner.run( @@ -1889,6 +1903,8 @@ def execute_model( # Eager (NONE): call the raw model directly. model_output = self.model(**model_inputs) + self.kv_connector.finish_forward() + if self.is_last_pp_rank: if self.use_aux_hidden_state_outputs: assert isinstance(model_output, tuple) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 7072210cb79e..56f7d438fc3f 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -224,10 +224,20 @@ def warmup_kernels( # fixed draft counts here, then restore the manager for capture and serving. adaptive_verification = model_runner.adaptive_verification model_runner.adaptive_verification = None + rejection_sampler = model_runner.rejection_sampler + adaptive_sampling = ( + rejection_sampler is not None and rejection_sampler.enable_adaptive_verification + ) + if adaptive_sampling: + assert rejection_sampler is not None + rejection_sampler.enable_adaptive_verification = False try: _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens) finally: model_runner.adaptive_verification = adaptive_verification + if adaptive_sampling: + assert rejection_sampler is not None + rejection_sampler.enable_adaptive_verification = True def _warmup_kernels( diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 5d6235bf3359..6054bb129aba 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -94,9 +94,12 @@ def _get_kv_connector_output( if wait_for_save and not defer_finalize: kv_connector.wait_for_save() - output.finished_sending, output.finished_recving = ( - kv_connector.get_finished(scheduler_output.finished_req_ids) + transfer_results = kv_connector.get_transfer_results( + scheduler_output.finished_req_ids ) + output.finished_sending = transfer_results.finished_sending + output.finished_recving = transfer_results.finished_recving + output.failed_recving = transfer_results.failed_recving output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors() output.kv_connector_stats = kv_connector.get_kv_connector_stats() diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 067a2b0d2267..df2ebc7d36ff 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -166,6 +166,8 @@ def __init__( kv = static_forward_context[layer_name].kv_cache if not isinstance(kv, torch.Tensor): continue + if kv.device.type != self.device.type: + continue dp = kv.data_ptr() assert kv.shape[0] % num_blocks == 0, ( @@ -366,32 +368,22 @@ def block_size_is_supported( return False return True - # Case 1: if the block_size of kv cache manager is supported by all backends, - # return it directly. if block_size_is_supported(backends, kv_manager_block_size): return kv_manager_block_size - # Case 2: otherwise, the block_size must be an `int`-format supported size of - # at least one backend. Iterate over all `int`-format supported sizes in - # descending order and return the first one that is supported by all backends. - # Simple proof: - # If the supported size b is in MultipleOf(x_i) format for all attention - # backends i, and b a factor of kv_manager_block_size, then - # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will - # return kv_manager_block_size in case 1. - all_int_supported_sizes = set( - supported_size + # MultipleOf constraints also accept the manager size if they accept a divisor. + # Any remaining candidate must therefore be an explicit size from a backend. + candidates = { + size for backend in backends - for supported_size in backend.get_supported_kernel_block_sizes() - if isinstance(supported_size, int) - ) + for size in backend.get_supported_kernel_block_sizes() + if isinstance(size, int) and kv_manager_block_size % size == 0 + } - for supported_size in sorted(all_int_supported_sizes, reverse=True): - if kv_manager_block_size % supported_size != 0: - continue - if block_size_is_supported(backends, supported_size): - return supported_size - raise ValueError(f"No common block size for {kv_manager_block_size}. ") + for size in sorted(candidates, reverse=True): + if block_size_is_supported(backends, size): + return size + raise ValueError(f"No common block size for {kv_manager_block_size}.") def allocate_kv_cache( @@ -440,7 +432,11 @@ def allocate_kv_cache( if isinstance(spec, UniformTypeKVCacheSpecs): spec = spec.kv_cache_specs[layer_name] - num_blocks = kv_cache_config.num_blocks + if not spec.has_layer_views: + kv_caches.update((name, buf) for name in tensor.layers) + continue + + num_blocks = kv_cache_config.num_blocks_of(tensor) kernel_block_size = None if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[group_id] @@ -485,7 +481,9 @@ def prepare_kernel_block_sizes( kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values())) if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): continue - if isinstance(kv_cache_spec, AttentionSpec): + if not kv_cache_spec.has_layer_views: + kernel_block_sizes.append(kv_cache_spec.block_size) + elif isinstance(kv_cache_spec, AttentionSpec): # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] @@ -621,7 +619,6 @@ def bind_kv_cache( for layer_name in kv_caches: index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name) - ordered_layer_names: list[str] = [] for layer_index in sorted(index2name.keys()): layer_names = index2name[layer_index] if len(layer_names) > 1: @@ -635,8 +632,19 @@ def bind_kv_cache( current_platform.check_runner_kv_caches_multi_layer() for layer_name in layer_names: runner_kv_caches.append(kv_caches[layer_name]) - ordered_layer_names.append(layer_name) + bind_kv_cache_to_layers( + kv_caches, forward_context, num_attn_module, kv_cache_groups + ) + + +def bind_kv_cache_to_layers( + kv_caches: dict[str, torch.Tensor], + forward_context: dict[str, Attention], + num_attn_module: int = 1, + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, +) -> None: + """Bind layer caches and share ReplaySSM trackers in model-layer order.""" # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks # its raw allocation into the per-layer view(s) it needs (e.g. Mamba # splits conv/ssm), so the kv_caches dict can hold a single tensor per @@ -644,6 +652,9 @@ def bind_kv_cache( for layer_name, kv_cache in kv_caches.items(): forward_context[layer_name].bind_kv_cache(kv_cache) + ordered_layer_names = sorted( + kv_caches, key=lambda name: extract_layer_index(name, num_attn_module) + ) share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) From eed1f3d0c6043bd494424a22443ee198dd56f657 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 12 Sep 2026 12:56:55 +0800 Subject: [PATCH 083/774] [Rust Frontend] Support `generation` blocks in HF chat template (#56378) Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/hf/generation.rs | 187 ++++++++++++++++++++ rust/src/chat/src/renderer/hf/mod.rs | 17 ++ rust/src/chat/src/renderer/hf/template.rs | 3 + 3 files changed, 207 insertions(+) create mode 100644 rust/src/chat/src/renderer/hf/generation.rs diff --git a/rust/src/chat/src/renderer/hf/generation.rs b/rust/src/chat/src/renderer/hf/generation.rs new file mode 100644 index 000000000000..39e61b5fb6e6 --- /dev/null +++ b/rust/src/chat/src/renderer/hf/generation.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Support Hugging Face's `{% generation %}` blocks in MiniJinja chat templates. +//! +//! Transformers uses these blocks to track assistant-token masks via scoped call +//! blocks. Inference only needs their rendered text. MiniJinja lacks this tag, so +//! we rewrite `{% generation %}...{% endgeneration %}` to +//! `{% call __hf_generation() %}...{% endcall %}` before compiling the template. +//! The registered `__hf_generation` callback renders and returns `caller()`, +//! preserving the block's local scope just as Transformers' AssistantTracker does. +//! +//! We use MiniJinja's lexer to identify statement tags and replace only their +//! keyword spans. This preserves delimiters and whitespace controls (`-` / `+`), +//! while keeping literals, comments, and raw blocks intact. Original `generation` +//! and `call` pairings are checked before rewriting both to call blocks. + +use minijinja::machinery::{Token, tokenize}; +use minijinja::value::{Kwargs, Value}; +use minijinja::{Error, ErrorKind, State}; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum CallBlockKind { + Call, + Generation, +} + +/// Rewrite Hugging Face's assistant-mask blocks as scoped MiniJinja call blocks. +pub(super) fn rewrite_generation_blocks(source: String) -> Result { + let tokens = tokenize(&source, false, Default::default(), Default::default()) + .collect::, _>>()?; + // Track native calls too: rewriting makes both block kinds use `endcall`, + // so validate the original pairing before that distinction is lost. + let mut blocks = Vec::new(); + let mut replacements = Vec::new(); + + for (index, pair) in tokens.windows(2).enumerate() { + let [(Token::BlockStart, _), (Token::Ident(name), span)] = pair else { + continue; + }; + match *name { + "call" => blocks.push(CallBlockKind::Call), + "endcall" if blocks.pop() == Some(CallBlockKind::Generation) => { + return Err(Error::new( + ErrorKind::SyntaxError, + "expected endgeneration, got endcall", + )); + } + "generation" | "endgeneration" + if matches!(tokens.get(index + 2), Some((Token::BlockEnd, _))) => + { + let replacement = if *name == "generation" { + blocks.push(CallBlockKind::Generation); + "call __hf_generation()" + } else { + if blocks.pop() != Some(CallBlockKind::Generation) { + return Err(Error::new( + ErrorKind::SyntaxError, + "unexpected endgeneration", + )); + } + "endcall" + }; + replacements.push(( + span.start_offset as usize..span.end_offset as usize, + replacement, + )); + } + _ => {} + } + } + if blocks.contains(&CallBlockKind::Generation) { + return Err(Error::new( + ErrorKind::SyntaxError, + "unclosed generation block", + )); + } + + // Replace only keyword spans, preserving delimiters and their whitespace controls. + let mut source = source; + for (span, replacement) in replacements.into_iter().rev() { + source.replace_range(span, replacement); + } + Ok(source) +} + +/// Render the caller with the same local scope as Transformers' AssistantTracker. +pub(super) fn render_generation(state: &State, kwargs: Kwargs) -> Result { + let caller: Value = kwargs.get("caller")?; + kwargs.assert_all_used()?; + caller.call(state, &[]) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use minijinja::Environment; + + use super::*; + + #[test] + fn generation_blocks_reject_invalid_syntax_and_mismatched_end_tags() { + let env = Environment::new(); + for template in [ + "{% generation %}unclosed", + "{% endgeneration %}", + "{% generation extra %}{% endgeneration %}", + "{% generation %}{% endgeneration extra %}", + "{% generation %}{% endcall %}", + "{% call f() %}{% endgeneration %}", + "{% generation %}{% endcall %}{% call f() %}{% endgeneration %}", + "{% generation %}{% if true %}{% endgeneration %}{% endif %}", + ] { + let result = rewrite_generation_blocks(template.to_string()) + .and_then(|source| env.template_from_str(&source).map(|_| ())); + assert!(result.is_err(), "accepted invalid template: {template}"); + } + } + + #[test] + fn rewrite_generation_blocks_preserves_hf_rendering_semantics() { + let cases = [ + ( + "whitespace", + "前 {%- generation -%} 中 {%- endgeneration -%} 后", + ), + ( + "keep_newlines", + "A{%+ generation +%}\nB{%+ endgeneration +%}\nC", + ), + ("literal", "{{ '{% generation %}' }}"), + ("raw", "{% raw %}{%- generation -%}{% endraw %}"), + ("comment", "A{# {% generation %} #}B"), + ( + "scope", + "{% set x = 'outer' %}{% generation %}{% set x = 'inner' %}{{ x }}{% endgeneration %}{{ x }}", + ), + ( + "nested_calls", + "{% macro wrap() %}[{{ caller() }}]{% endmacro %}{% generation %}A{% call wrap() %}{% generation %}B{% endgeneration %}{% endcall %}C{% endgeneration %}", + ), + ]; + let results = cases.map(|(name, source)| { + let template = rewrite_generation_blocks(source.to_string()).unwrap(); + let mut env = Environment::new(); + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + env.add_function("__hf_generation", render_generation); + let rendered = + env.template_from_str(&template).unwrap().render(Value::UNDEFINED).unwrap(); + (name, rendered) + }); + expect![[r#" + [ + ( + "whitespace", + "前中后", + ), + ( + "keep_newlines", + "A\nB\nC", + ), + ( + "literal", + "{% generation %}", + ), + ( + "raw", + "{%- generation -%}", + ), + ( + "comment", + "AB", + ), + ( + "scope", + "innerouter", + ), + ( + "nested_calls", + "A[B]C", + ), + ] + "#]] + .assert_debug_eq(&results); + } +} diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 7461acf2955a..5e7ba75e0a7f 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -25,6 +25,7 @@ use crate::{ mod error; mod format; +mod generation; mod template; mod tojson; @@ -642,6 +643,22 @@ mod tests { .map_err(|_| unreachable!("HF renderer should return text prompt")) } + #[test] + fn generation_blocks_allow_content_format_detection_and_request_overrides() { + let template = "{% for message in messages %}{% generation %}{% for part in message.content %}{{ part.text }}{% endfor %}{% endgeneration %}{% endfor %}"; + let mut request = sample_request(vec![ChatMessage::user("hello")]); + let default = render(Some(template), &request).unwrap(); + request.chat_options.chat_template = Some(template.to_string()); + let overridden = render(Some("unused"), &request).unwrap(); + expect![[r#" + ( + "hello", + "hello", + ) + "#]] + .assert_debug_eq(&(default, overridden)); + } + fn render_mm( template: &str, request: &ChatRequest, diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index b7de0172f36f..b8907837012b 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -20,6 +20,7 @@ use super::error::TemplateError; use super::format::{ ChatTemplateContentFormat, ChatTemplateContentFormatOption, detect_chat_template_content_format, }; +use super::generation::{render_generation, rewrite_generation_blocks}; use super::tojson::hf_tojson_filter; use crate::renderer::hf::{TemplateMessage, TemplateTool}; @@ -36,6 +37,7 @@ fn build_environment(template: String) -> Result> { env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); env.add_filter("tojson", hf_tojson_filter); + env.add_function("__hf_generation", render_generation); Ok(env) } @@ -107,6 +109,7 @@ pub(super) struct CompiledChatTemplate { impl CompiledChatTemplate { /// Compile the given chat template string into a [`CompiledChatTemplate`]. pub fn new(template: String, content_format: ChatTemplateContentFormatOption) -> Result { + let template = rewrite_generation_blocks(template)?; let content_format = match content_format { ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template), ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String, From dca96bf97b350204937a72c4d386b79b425f3989 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Sat, 12 Sep 2026 14:22:30 +0800 Subject: [PATCH 084/774] [Bugfix][Spec Decode] Avoid fastsafetensors deadlock for PP draft models (#54416) Signed-off-by: Canlin Co-authored-by: Codex --- .../test_draft_attention_backend_override.py | 9 ++++++++- .../test_draft_moe_backend_override.py | 9 ++++++++- vllm/v1/worker/gpu/spec_decode/dflash/utils.py | 2 ++ vllm/v1/worker/gpu/spec_decode/dspark/utils.py | 2 ++ vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 4 ++++ vllm/v1/worker/gpu/spec_decode/utils.py | 18 ++++++++++++++++++ 6 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_draft_attention_backend_override.py b/tests/v1/spec_decode/test_draft_attention_backend_override.py index 5984d337692b..410afe8d61d3 100644 --- a/tests/v1/spec_decode/test_draft_attention_backend_override.py +++ b/tests/v1/spec_decode/test_draft_attention_backend_override.py @@ -6,11 +6,13 @@ assert on the config ``load_eagle_model`` hands to ``get_model``. """ -from dataclasses import dataclass +from dataclasses import dataclass, field +from types import SimpleNamespace from unittest.mock import patch import pytest +from vllm.config import LoadConfig from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model @@ -43,6 +45,7 @@ class _VllmConfig: kernel_config: _KernelConfig cache_config: _CacheConfig speculative_config: _SpeculativeConfig + load_config: LoadConfig = field(default_factory=LoadConfig) def _config(target_backend: str, draft_backend: str | None) -> _VllmConfig: @@ -65,6 +68,10 @@ def _fake_get_model(*, vllm_config, model_config): with ( patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model), + patch( + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", + return_value=SimpleNamespace(world_size=1), + ), pytest.raises(_Captured) as exc, ): load_eagle_model(object(), cfg) diff --git a/tests/v1/spec_decode/test_draft_moe_backend_override.py b/tests/v1/spec_decode/test_draft_moe_backend_override.py index d4acd9e4be49..5a4522e5b301 100644 --- a/tests/v1/spec_decode/test_draft_moe_backend_override.py +++ b/tests/v1/spec_decode/test_draft_moe_backend_override.py @@ -8,11 +8,13 @@ server fails to start rather than falling back. """ -from dataclasses import dataclass +from dataclasses import dataclass, field +from types import SimpleNamespace from unittest.mock import patch import pytest +from vllm.config import LoadConfig from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model @@ -39,6 +41,7 @@ class _VllmConfig: kernel_config: _KernelConfig cache_config: _CacheConfig speculative_config: _SpeculativeConfig + load_config: LoadConfig = field(default_factory=LoadConfig) def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig: @@ -62,6 +65,10 @@ def _fake_get_model(*, vllm_config, model_config): with ( patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model), + patch( + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", + return_value=SimpleNamespace(world_size=1), + ), pytest.raises(_Captured) as exc, ): load_eagle_model(object(), cfg) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 6226dea1c2db..7323b29312c1 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -9,6 +9,7 @@ get_target_lm_head, maybe_share_target_embed, ) +from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: @@ -37,6 +38,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if speculative_config.kv_cache_dtype is not None else vllm_config.cache_config ), + load_config=get_pp_safe_draft_load_config(vllm_config.load_config), ) with set_model_tag("dflash_head"): dflash_model = get_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 5b13530ef73f..db1e17a9f23b 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -6,6 +6,7 @@ from vllm.config import ModelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config logger = init_logger(__name__) @@ -73,6 +74,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if speculative_config.kv_cache_dtype is not None else vllm_config.cache_config ), + load_config=get_pp_safe_draft_load_config(vllm_config.load_config), ) # VllmConfig post-init restores the target's quant config because the target # config is retained for DSpark's target-layer metadata, so we must override it. diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index a85f21a3cd66..b93a79e45513 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -8,6 +8,7 @@ from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.model_loader import get_model from vllm.model_executor.models.utils import PPMissingLayer +from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: @@ -103,6 +104,9 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod backend=speculative_config.attention_backend, ), ) + draft_load_config = get_pp_safe_draft_load_config(vllm_config.load_config) + if draft_load_config is not vllm_config.load_config: + vllm_config = replace(vllm_config, load_config=draft_load_config) with set_model_tag("eagle_head"): eagle_model = get_model( vllm_config=vllm_config, model_config=draft_model_config diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index e25672e953bd..f73acb2bdaed 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -3,10 +3,28 @@ import numpy as np import torch +from vllm.config import LoadConfig, replace +from vllm.distributed import get_pp_group +from vllm.logger import init_logger from vllm.v1.outputs import DraftTokenIds from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.input_batch import InputBatch +logger = init_logger(__name__) + + +def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig: + """Avoid collectives that include PP ranks without a draft model.""" + if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors": + logger.warning_once( + "fastsafetensors cannot load a draft model instantiated on only " + "one pipeline stage; falling back to the standard safetensors " + "loader for the draft model. The target model still uses " + "fastsafetensors." + ) + return replace(load_config, load_format="auto") + return load_config + class DraftTokensHandler: def __init__(self, device: torch.device | None = None): From 372477d2d9890381dbc794cf66ef86ff989f04e0 Mon Sep 17 00:00:00 2001 From: junuxyz <216036880+junuxyz@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:22:34 +0900 Subject: [PATCH 085/774] [Bugfix][Examples] Launch prefill and decode concurrently in NixlPushConnector demo (#55088) Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com> --- .../disagg_proxy_pushconnector_demo.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py index 672162f8ba34..a35a44fbba2d 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py @@ -31,6 +31,7 @@ """ import argparse +import asyncio import contextlib import ipaddress import itertools @@ -259,17 +260,24 @@ async def _push_completion(self, raw_request: Request, path: str): decode_instance = self.schedule(self.decode_cycler) headers = self._common_headers(request_id) - # Fire prefill; we don't read its body but must drain the - # connection so the upstream server can free its slot. - async for _ in self.forward_request( - f"http://{prefill_instance}{path}", prefill_request, headers - ): - continue + async def drain_prefill(): + async for _ in self.forward_request( + f"http://{prefill_instance}{path}", prefill_request, headers + ): + continue - generator = self.forward_request( - f"http://{decode_instance}{path}", decode_request, headers - ) - return StreamingResponse(generator, media_type="application/json") + prefill_task = asyncio.create_task(drain_prefill()) + + async def stream_decode(): + try: + async for chunk in self.forward_request( + f"http://{decode_instance}{path}", decode_request, headers + ): + yield chunk + finally: + await prefill_task + + return StreamingResponse(stream_decode(), media_type="application/json") async def create_completion(self, raw_request: Request): try: From 30118ba27d1d923bdd91f97d945528dcb4a862c1 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 12 Sep 2026 15:07:36 +0800 Subject: [PATCH 086/774] [DSV4.1] Remove compressor-aware image sentinel token padding (#56554) Signed-off-by: Isotr0py --- vllm/models/deepseek_v4_1/amd/vl_model.py | 8 -- .../deepseek_v4_1/common/mm_preprocess.py | 120 ++---------------- vllm/models/deepseek_v4_1/nvidia/vl_model.py | 8 -- .../configs/deepseek_v41.py | 6 - 4 files changed, 10 insertions(+), 132 deletions(-) diff --git a/vllm/models/deepseek_v4_1/amd/vl_model.py b/vllm/models/deepseek_v4_1/amd/vl_model.py index 91a13bd72506..3227b50a4e45 100644 --- a/vllm/models/deepseek_v4_1/amd/vl_model.py +++ b/vllm/models/deepseek_v4_1/amd/vl_model.py @@ -47,9 +47,7 @@ IMAGE, IMAGE_END, IMAGE_NEW_LINE, - IMAGE_PAD_ID, IMAGE_PLACEHOLDER, - IMAGE_SENTINEL_BASE_ID, IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, @@ -277,12 +275,6 @@ def embed_input_ids( ) -> torch.Tensor: from vllm.model_executor.models.utils import _merge_multimodal_embeddings - # Compressor-alignment pads borrow a reserved id; embed them as the - # plain image token (the checkpoint has no image_pad vector). - # Branch-free: safe inside compiled/captured regions. - input_ids = input_ids.masked_fill( - input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID - ) inputs_embeds = self.language_model.embed_input_ids(input_ids) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: diff --git a/vllm/models/deepseek_v4_1/common/mm_preprocess.py b/vllm/models/deepseek_v4_1/common/mm_preprocess.py index 8bd8d70caea3..4f87b816c9a5 100644 --- a/vllm/models/deepseek_v4_1/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4_1/common/mm_preprocess.py @@ -12,14 +12,9 @@ rows in reading order; the delimiters take the learned ``image_start`` / ``image_newline`` / ``image_end`` vectors. -One vLLM-side deviation from the reference token stream: a leading -compressor-alignment pad (``COMPRESS_PAD_TO - 1 - start % COMPRESS_PAD_TO`` -positions, so the span always starts at the same compressor phase) is -prepended when the block is spliced into the final prompt. Pad positions -borrow the reserved in-vocab token ``<|place_holder_mm_span_0436|>`` so -they stay distinguishable from real span positions; they embed as the -plain image token (v4.1 has no ``image_pad`` vector) and are routed and -engram-deadened like image tokens. +The token stream matches the reference exactly: no compressor-alignment +pad is inserted (the reference pools image tokens across compressor-group +boundaries freely). """ import math @@ -30,7 +25,6 @@ import torch from PIL import Image, ImageOps from transformers import BatchFeature -from typing_extensions import assert_never from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions from vllm.inputs import MultiModalDataDict @@ -44,42 +38,26 @@ PromptUpdate, PromptUpdateDetails, ) -from vllm.multimodal.processing.processor import ( - MultiModalPromptUpdates, - MultiModalPromptUpdatesApplyResult, - PlaceholderFeaturesInfo, - UpdateMode, - _plan_prompt_updates, -) from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4) -# Image spans are padded so they always start at the same compressor group -# phase; v4.1 uses ratio-2 compressors (v4.0 used 4). -COMPRESS_PAD_TO = 2 IMAGE_PLACEHOLDER = "<｜deepseek_image｜>" # The checkpoint tokenizer's id for IMAGE_PLACEHOLDER (the config's # ``image_token_id``). Every image-span position carries this id, exactly # like the reference; the MoE router detects image tokens as the range -# [IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID + 5), which also covers -# IMAGE_PAD_ID below. +# [IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID + 5). IMAGE_SENTINEL_BASE_ID = 129264 -# Leading compressor-alignment pads borrow this reserved in-vocab token -# (``<|place_holder_mm_span_0436|>``) so they can be told apart from real -# span positions when the per-position role is decided from token ids. -IMAGE_PAD_ID = 129265 -IMAGE_PAD_TOKEN_NAME = "<|place_holder_mm_span_0436|>" def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor: - """Boolean mask for image-span positions (span tokens and align pads).""" - return (token_ids == IMAGE_SENTINEL_BASE_ID) | (token_ids == IMAGE_PAD_ID) + """Boolean mask for image-span positions.""" + return token_ids == IMAGE_SENTINEL_BASE_ID def validate_image_sentinel_ids(tokenizer) -> None: - """Check the image/pad token ids against the tokenizer.""" + """Check the image token id against the tokenizer.""" image_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER) if image_id != IMAGE_SENTINEL_BASE_ID: raise ValueError( @@ -88,13 +66,6 @@ def validate_image_sentinel_ids(tokenizer) -> None: "image_token_id); the DeepSeek-V4.1 vision path keys image " "routing and engram masking off this id." ) - pad_id = tokenizer.convert_tokens_to_ids(IMAGE_PAD_TOKEN_NAME) - if pad_id != IMAGE_PAD_ID: - raise ValueError( - f"Image pad token {IMAGE_PAD_TOKEN_NAME!r} has id {pad_id}, " - f"expected {IMAGE_PAD_ID}; the DeepSeek-V4.1 vision path " - "borrows this reserved id for compressor-alignment pads." - ) def llm_grid(best_height, best_width, patch_size, downsample_ratio): @@ -134,8 +105,7 @@ def safe_resize( height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token ): """Shrink the pixel size until the image costs at most max_n_token LLM - tokens (minus the reservation for the compressor-alignment pad).""" - max_n_token -= COMPRESS_PAD_TO - 1 + tokens.""" n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio) if num_image_tokens(n_llm_h, n_llm_w) > max_n_token: best_height, best_width = solve_resize_ratio( @@ -293,13 +263,7 @@ def get_mm_max_tokens_per_item( seq_len: int, mm_counts: Mapping[str, int], ) -> Mapping[str, int]: - # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of - # vision_max_n_token for the compressor-alignment pad, so the full - # image span is bounded by vision_max_n_token; the margin is kept - # in case that reservation changes. - return { - "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1, - } + return {"image": self.get_hf_config().vision_max_n_token} def get_image_placeholder_token_id(self) -> int: token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER) @@ -314,7 +278,7 @@ def get_image_size_with_most_features(self) -> ImageSize: # A square maximizes the ViT patch count (area) within the token # budget; solve the budget-derived size directly to keep the dummy # image small. - budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1) + budget = hf_config.vision_max_n_token side = budget * patch_size * downsample_ratio best_h, best_w = solve_resize_ratio( side, side, patch_size, downsample_ratio, budget @@ -397,67 +361,3 @@ def get_image_replacement(item_idx: int) -> PromptUpdateDetails: replacement=get_image_replacement, ), ] - - def _apply_token_matches_with_placeholders( - self, - token_ids: list[int], - mm_prompt_updates: MultiModalPromptUpdates, - ) -> tuple[ - list[int], - MultiModalPromptUpdatesApplyResult, - Mapping[str, list[PlaceholderFeaturesInfo]], - ]: - # Same as the base implementation, except that each image block gets - # its compressor-alignment pad (``COMPRESS_PAD_TO - 1 - start % - # COMPRESS_PAD_TO`` IMAGE_PAD_ID tokens) prepended while splicing: - # the pad depends on the block's final position in the prompt, which - # is unknown when the (cacheable) prompt updates are built. Pad - # positions are not embed positions (their id differs from - # image_token_id), so ``is_embed`` stays aligned with the pad-free - # replacement content. - matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates) - placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { - modality: [] for modality in mm_prompt_updates - } - - new_token_ids = list[int]() - prev_end_idx = 0 - for matched_update in matched_updates: - update = matched_update.update - match = matched_update.match - - if update.mode == UpdateMode.INSERT: - end_idx_to_insert = match.end_idx - elif update.mode == UpdateMode.REPLACE: - end_idx_to_insert = match.start_idx - else: - assert_never(update.mode) - - new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert]) - start_idx = len(new_token_ids) - - tokens = list(update.content.full) - if tokens and update.modality == "image": - compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO - tokens = [IMAGE_PAD_ID] * compress_pad + tokens - - if tokens: - content_is_embed = update.content.is_embed - is_embed = ( - content_is_embed(tokens) if content_is_embed is not None else None - ) - placeholders[update.modality].append( - PlaceholderFeaturesInfo( - modality=update.modality, - item_idx=update.item_idx, - start_idx=start_idx, - tokens=tokens, - is_embed=is_embed, - ) - ) - new_token_ids.extend(tokens) - - prev_end_idx = match.end_idx - - new_token_ids.extend(token_ids[prev_end_idx:]) - return new_token_ids, result, placeholders diff --git a/vllm/models/deepseek_v4_1/nvidia/vl_model.py b/vllm/models/deepseek_v4_1/nvidia/vl_model.py index 111936b43115..bfbcaa89a1a1 100644 --- a/vllm/models/deepseek_v4_1/nvidia/vl_model.py +++ b/vllm/models/deepseek_v4_1/nvidia/vl_model.py @@ -47,9 +47,7 @@ IMAGE, IMAGE_END, IMAGE_NEW_LINE, - IMAGE_PAD_ID, IMAGE_PLACEHOLDER, - IMAGE_SENTINEL_BASE_ID, IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, @@ -277,12 +275,6 @@ def embed_input_ids( ) -> torch.Tensor: from vllm.model_executor.models.utils import _merge_multimodal_embeddings - # Compressor-alignment pads borrow a reserved id; embed them as the - # plain image token (the checkpoint has no image_pad vector). - # Branch-free: safe inside compiled/captured regions. - input_ids = input_ids.masked_fill( - input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID - ) inputs_embeds = self.language_model.embed_input_ids(input_ids) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py index 3a406d24bdc5..0be36d2480db 100644 --- a/vllm/transformers_utils/configs/deepseek_v41.py +++ b/vllm/transformers_utils/configs/deepseek_v41.py @@ -68,9 +68,3 @@ def __init__( # in-kernel, so mm-prefix ranges longer than sliding_window must be # kept (they are only consumed by those kernels). self.mm_prefix_clamp_sliding_window = vision_n_layers > 0 - # The visibility span covers the image block [IMAGE_START, - # IMAGE_END]; the mm placeholder additionally carries a leading - # compressor-alignment pad of ``COMPRESS_PAD_TO - 1 - offset % - # COMPRESS_PAD_TO`` tokens (see common/mm_preprocess.py), with - # COMPRESS_PAD_TO = 2 for v4.1's ratio-2 compressors. - self.mm_prefix_span_leading_pad_modulus = 2 if vision_n_layers > 0 else 0 From d406f09eb8d05eaf81d768bf74212bc6d24cd11d Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 12 Sep 2026 02:40:14 -0500 Subject: [PATCH 087/774] [ROCm][CI] Stage F gating (#55252) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/disaggregated.yaml | 35 ++++++++++++++ .buildkite/test_areas/distributed.yaml | 39 +++++++++++++++ .buildkite/test_areas/lm_eval.yaml | 42 +++++++++++++++++ .buildkite/test_areas/misc.yaml | 47 +++++++++++++++++++ .buildkite/test_areas/models_basic.yaml | 9 ++++ .buildkite/test_areas/models_multimodal.yaml | 9 ++++ .buildkite/test_areas/pytorch.yaml | 8 ++++ .buildkite/test_areas/rust_frontend.yaml | 29 ++++++++++++ .../test_areas/rust_frontend_cargo.yaml | 18 +++++++ 9 files changed, 236 insertions(+) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 0100b99993ef..5ab0e836ebd8 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -172,6 +172,23 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh + mirror: + amd: + label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache" + dind: false + device: mi355_2 + timeout_in_minutes: 25 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy" key: multiconnector-nixl-offloading-pd-accuracy-2-gpus @@ -234,6 +251,24 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + mirror: + amd: + label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases" + dind: false + device: mi355_2 + timeout_in_minutes: 25 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh # P TP 4 - D DPEP 4 test case for DSv4-Flash - label: ":nvidia: (H200) DSv4-Flash Disaggregated DP EP" diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 44e96ce63264..dfe900d827f3 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -105,6 +105,26 @@ steps: - export NCCL_CUMEM_HOST_ENABLE=0 - pytest -v -s entrypoints/llm/test_collective_rpc.py - pytest -v -s ./compile/test_wrapper.py + mirror: + amd: + label: ":amd: (MI355) Distributed Compile + RPC" + dind: false + device: mi355_2 + timeout_in_minutes: 65 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/compile/test_wrapper.py + - tests/entrypoints/llm/test_collective_rpc.py + - vllm/platforms/rocm.py - label: ":nvidia: (L4) Distributed Torchrun + Shutdown" key: distributed-torchrun-shutdown-tests-2-gpus @@ -335,6 +355,25 @@ steps: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py - pytest -v -s v1/distributed/test_pp_dp_v2.py + mirror: + amd: + label: ":amd: (MI250) Pipeline + Context Parallelism" + device: mi250_4 + timeout_in_minutes: 90 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/executor/ + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (L4) RayExecutorV2" key: rayexecutorv2-4-gpus diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index dc90dcb8529c..1c2e3bb54f18 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -79,6 +79,26 @@ steps: autorun_on_main: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt + mirror: + amd: + label: ":amd: (MI355) LM Eval Small Models Distributed" + dind: false + device: mi355_2 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 @@ -375,6 +395,28 @@ steps: commands: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt + mirror: + amd: + label: ":amd: (MI355) GPQA Eval (GPT-OSS)" + dind: false + device: mi355_2 + timeout_in_minutes: 40 + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - tests/evals/gpt_oss/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt - label: ":nvidia: (B200) Qwen3.8-Flash-Next-FP8 Accuracy Eval" key: accuracy-eval-qwen3-8-flash-next-fp8-b200 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 5ae3bcf73ecb..d24ec5b57f46 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -251,6 +251,35 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration + mirror: + amd: + label: ":amd: (MI355) Extract Hidden States Integration" + dind: false + device: mi355_2 + timeout_in_minutes: 20 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/config/speculative.py + - vllm/distributed/kv_transfer/kv_connector/ + - vllm/model_executor/layers/attention/ + - vllm/model_executor/layers/mamba/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/models/extract_hidden_states.py + - vllm/model_executor/models/llama.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/registry.py + - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/v1/kv_cache_interface.py + - vllm/v1/spec_decode/extract_hidden_states.py + - vllm/v1/worker/gpu_model_runner.py + - vllm/_aiter_ops.py + - tests/v1/kv_connector/extract_hidden_states_integration/ + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) Regression" key: regression @@ -537,6 +566,24 @@ steps: - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + mirror: + amd: + label: ":amd: (MI250) Batch Invariance" + device: mi250_1 + timeout_in_minutes: 30 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - label: ":nvidia: (B200) Batch Invariance" key: batch-invariance-b200 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 24447cb50171..bf453a12d8b3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -156,3 +156,12 @@ steps: device: cpu-small commands: - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ + mirror: + amd: + label: ":amd: (MI250) Basic Models Other" + device: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace/tests" + timeout_in_minutes: 25 + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 9893383bff1b..b09e8cc2d907 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -98,6 +98,15 @@ steps: commands: - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 4 + mirror: + amd: + label: ":amd: (MI250) Multimodal Processor Shard %N" + device: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace/tests" + timeout_in_minutes: 80 + depends_on: + - image-build-amd - label: ":nvidia: (H200 MIG 18GB) Multimodal Processor" key: multi-modal-processor diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index eb5b8407033c..f34ff50e672d 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -160,6 +160,14 @@ steps: # Use `find` to launch multiple instances of pytest so that # they do not suffer from https://github.com/vllm-project/vllm/issues/28965 - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" + mirror: + amd: + label: ":amd: (MI355) PyTorch Fullgraph" + dind: false + device: mi355_1 + timeout_in_minutes: 90 + depends_on: + - image-build-amd # Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model # instances and cannot fit a 35GB MIG slice. L4 retains the original coverage: diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 97078512a1d7..770d6b64df22 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -131,6 +131,20 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + mirror: + amd: + label: ":amd: (MI355) Rust Frontend Core Correctness" + dind: false + device: mi355_1 + timeout_in_minutes: 20 + depends_on: + - image-build-amd + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - tests/utils.py + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use" key: rust-frontend-tool-use @@ -147,6 +161,21 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" + mirror: + amd: + label: ":amd: (MI355) Rust Frontend Tool Use" + dind: false + device: mi355_1 + timeout_in_minutes: 25 + depends_on: + - image-build-amd + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - vllm/tool_parsers/ + - tests/utils.py + - tests/tool_use/ + - vllm/platforms/rocm.py - label: ":nvidia: (L4) Rust Frontend Distributed" key: rust-frontend-distributed diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 11bca6252667..0b49787d5d76 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -14,6 +14,15 @@ steps: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy + mirror: + amd: + label: ":amd: (MI250) Rust Frontend Cargo Style + Clippy" + device: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace" + timeout_in_minutes: 35 + depends_on: + - image-build-amd - label: ":computer: (CPU) Rust Frontend Cargo" key: rust-frontend-cargo-tests @@ -28,3 +37,12 @@ steps: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test + mirror: + amd: + label: ":amd: (MI250) Rust Frontend Cargo" + device: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace" + timeout_in_minutes: 35 + depends_on: + - image-build-amd From 756794a9a7f08900c00fbfaa6d8332631503f528 Mon Sep 17 00:00:00 2001 From: Walter Beller-Morales Date: Sat, 12 Sep 2026 04:04:46 -0400 Subject: [PATCH 088/774] [Frontend] create unified Cohere parser (#56392) Signed-off-by: walterbm Co-authored-by: Flora Feng <4florafeng@gmail.com> --- docs/features/reasoning_outputs.md | 1 + docs/features/tool_calling.md | 9 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/rocm.txt | 2 +- tests/parser/cohere/__init__.py | 2 + tests/parser/cohere/conftest.py | 18 + tests/parser/cohere/test_citations.py | 172 ++++ tests/parser/cohere/test_reasoning.py | 203 +++++ tests/parser/cohere/test_structural_tags.py | 362 ++++++++ tests/parser/cohere/test_tool_calls.py | 288 +++++++ tests/parser/cohere/utils.py | 131 +++ .../test_cohere_command_reasoning_parser.py | 796 ------------------ .../test_cohere_command_tool_parser.py | 510 ----------- .../entrypoints/cohere/cohere_chat_message.py | 11 +- vllm/entrypoints/cohere/serving.py | 23 +- vllm/parser/cohere_command.py | 687 +++++++++++++++ vllm/parser/parser_manager.py | 15 + .../cohere_command_reasoning_parser.py | 680 +-------------- vllm/renderers/cohere.py | 2 +- .../cohere_command_tool_parser.py | 138 +-- 21 files changed, 1917 insertions(+), 2137 deletions(-) create mode 100644 tests/parser/cohere/__init__.py create mode 100644 tests/parser/cohere/conftest.py create mode 100644 tests/parser/cohere/test_citations.py create mode 100644 tests/parser/cohere/test_reasoning.py create mode 100644 tests/parser/cohere/test_structural_tags.py create mode 100644 tests/parser/cohere/test_tool_calls.py create mode 100644 tests/parser/cohere/utils.py delete mode 100644 tests/reasoning/test_cohere_command_reasoning_parser.py delete mode 100644 tests/tool_parsers/test_cohere_command_tool_parser.py create mode 100644 vllm/parser/cohere_command.py diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index fbd23fcdd581..a58dc6033e13 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models: | Model Series | Parser Name | Structured Output Support | Tool Calling | | ------------ | ----------- | ---------------- | ----------- | | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ | +| [Cohere Command A Plus](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) | `cohere_command4` | `json`, `regex` | ✅ | | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ | | [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ | | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ | diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index ae65231919a3..9628bc431a1e 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -366,15 +366,16 @@ Flags: * For non-reasoning: `--tool-call-parser hunyuan_a13b` * For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b` -### Cohere Command A Reasoning (`cohere_command3`) +### Cohere Command (`cohere_command3` / `cohere_command4`) Supported models: -* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) +* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) (`cohere_command3`) +* [`CohereLabs/command-a-plus-05-2026`](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) and [`CohereLabs/North-Mini-Code-1.0`](https://huggingface.co/CohereLabs/North-Mini-Code-1.0) (`cohere_command4`) -Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3` +Flags: `--tool-call-parser cohere_command4 --reasoning-parser cohere_command4`. -Note: the Cohere tool parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package. +Note: the Cohere parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package. ### LongCat-Flash-Chat Models (`longcat`) diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 2b1b8d372d6e..052896c9f22b 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -127,7 +127,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 2910f01e9f95..db6e35918289 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -132,7 +132,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index ed47efded1db..a6b9271989af 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -134,7 +134,7 @@ cloudpickle==3.1.2 # tilelang cohere==7.0.8 # via -r requirements/test/rocm.in -cohere-melody==0.11.1 +cohere-melody==0.14.0 # via -r requirements/test/rocm.in colorama==0.4.6 # via diff --git a/tests/parser/cohere/__init__.py b/tests/parser/cohere/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/parser/cohere/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/parser/cohere/conftest.py b/tests/parser/cohere/conftest.py new file mode 100644 index 000000000000..4477cacf43a3 --- /dev/null +++ b/tests/parser/cohere/conftest.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + +from .utils import MockCohereTokenizer + + +@pytest.fixture(scope="package") +def tokenizer() -> MockCohereTokenizer: + return MockCohereTokenizer() + + +@pytest.fixture +def request_obj() -> ChatCompletionRequest: + return ChatCompletionRequest(messages=[], model="test-model") diff --git a/tests/parser/cohere/test_citations.py b/tests/parser/cohere/test_citations.py new file mode 100644 index 000000000000..14ee1612ead4 --- /dev/null +++ b/tests/parser/cohere/test_citations.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Citation source resolution and the streamed citation wire shape.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from vllm.parser.cohere_command import ( + _melody_citations_to_vllm, + _melody_sources_to_vllm, +) + +from .utils import make_parser, stream_parser + + +class TestMelodySourceResolution: + """Pin the parser-side source resolution (see + ``_melody_sources_to_vllm``). The parser receives melody's numeric + ``(tool_call_index, tool_result_indices)`` addressing and resolves + it against a ``position_to_source`` map handed in by + ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same + map that would otherwise live in the serving layer's resolver. + """ + + @staticmethod + def _fake_melody_source(bucket: int, indices: list[int]) -> Any: + return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) + + def test_multi_index_source_fans_out(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), + (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), + } + raw = [self._fake_melody_source(0, [0, 1])] + out = _melody_sources_to_vllm(raw, position_map) + assert [s.id for s in out] == ["d0", "d1"] + # Verify type / payload were plumbed through, not just the id. + assert out[0].type == "document" + assert out[0].document == {"id": "d0"} + + def test_unresolvable_position_skipped(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [self._fake_melody_source(9, [0])] + assert _melody_sources_to_vllm(raw, position_map) == [] + + def test_missing_position_map_drops_all_sources(self): + # A parser instance without a position map (parser wired + # outside of ``CohereServingChatV2``) can't attribute anything, + # so every source is dropped. Callers downstream will see the + # citation with empty ``sources`` and drop it entirely. + raw = [self._fake_melody_source(0, [0])] + assert _melody_sources_to_vllm(raw, None) == [] + + def test_citations_pass_through_is_thinking_tag(self): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + + position_map: dict[tuple[int, int], CitationSource] = { + (0, 0): CitationSource(type="document", id="d0"), + } + raw = [ + SimpleNamespace( + start_index=0, + end_index=5, + text="hello", + is_thinking=True, + sources=[self._fake_melody_source(0, [0])], + ) + ] + out = _melody_citations_to_vllm(raw, position_map) + assert out is not None + assert out[0].type == "THINKING_CONTENT" + assert out[0].sources[0].id == "d0" + + +class TestParserStreamingEndToEnd: + """End-to-end streaming shape: raw model output text is tokenized + with the mock byte-level tokenizer and fed one token at a time + through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`. + We collect the ordered sequence of ``DeltaMessage``s the parser + hands back -- i.e. what the client would receive on the wire -- + and pin the exact shape. + + Combines the three flavors that reach ``delta``: a thinking block + (populates ``reasoning``), a text block (populates ``content``), + and a resolved citation (populates ``citations``) referencing + ``(start, end)`` offsets into the already-streamed content. + """ + + _MODEL_OUTPUT = ( + "<|START_THINKING|>Let me check.<|END_THINKING|>" + "<|START_RESPONSE|>The capital is Paris.<|END_RESPONSE|>" + ) + + def test_full_wire_stream_shape(self, tokenizer, request_obj): + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource + from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY + + # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs`` + # installs on a request whose tool_call_index=0 tool result at + # index 1 is the "France" document. + position_to_source = { + (0, 1): CitationSource( + type="document", + id="doc-paris", + document={"id": "doc-paris", "title": "France"}, + ), + } + + parser = make_parser( + tokenizer, + "cohere_command3", + chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source}, + ) + deltas = stream_parser(parser, request_obj, tokenizer, self._MODEL_OUTPUT) + + # -- 1. Reasoning and content stream separately and reach + # the wire only as their real payload. Framing tokens + # (``<|START_THINKING|>`` etc.) are consumed by the parser's + # state machine and never produce a delta. + reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None] + content_stream = [d.content for d in deltas if d.content is not None] + citation_deltas = [d for d in deltas if getattr(d, "citations", None)] + + assert "".join(reasoning_stream) == "Let me check." + assert "".join(content_stream) == "The capital is Paris." + + # -- 2. Exactly one citation reaches the wire, resolved to the + # source from ``position_to_source`` (not the raw + # ``(bucket, idx)`` coordinates). + assert len(citation_deltas) == 1 + (cite_delta,) = citation_deltas + assert cite_delta.citations is not None + assert len(cite_delta.citations) == 1 + cite = cite_delta.citations[0] + assert cite.text == "Paris" + # Char offsets are into the accumulated ``content``: "The + # capital is " is 15 chars, "Paris" is 5. + assert cite.start == 15 + assert cite.end == 20 + assert len(cite.sources) == 1 + assert cite.sources[0].id == "doc-paris" + assert cite.sources[0].type == "document" + + # -- 3. The citation's ``(start, end)`` offsets point into + # bytes that were already emitted as content-deltas earlier + # in the stream -- i.e. the citation is anchored to real + # already-shipped content. + cite_delta_idx = deltas.index(cite_delta) + earlier_content = "".join( + d.content for d in deltas[:cite_delta_idx] if d.content is not None + ) + assert earlier_content[cite.start : cite.end] == "Paris" + + # Everything after the citation delta is the trailing period. + later_content = "".join( + d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None + ) + assert later_content == "." + + # -- 4. Reasoning and content are never mixed on the same + # delta: the parser flips modes on the ``<|END_THINKING|>`` + # boundary and never emits a delta with both fields set. + for d in deltas: + assert not (d.reasoning is not None and d.content is not None) diff --git a/tests/parser/cohere/test_reasoning.py b/tests/parser/cohere/test_reasoning.py new file mode 100644 index 000000000000..8c3311e8388c --- /dev/null +++ b/tests/parser/cohere/test_reasoning.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning extraction, reasoning-end gating and parser selection for the +unified Cohere Command parser.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser import ParserManager +from vllm.reasoning.cohere_command_reasoning_parser import ( + CohereCommand3ReasoningParser, + CohereCommand4ReasoningParser, +) + +from .utils import make_parser, stream_parser + + +@dataclass +class ReasoningCase: + parser_name: str + model_output: str + expected_reasoning: str | None + expected_content: str | None + + +REASONING_CASES = [ + pytest.param( + ReasoningCase( + parser_name="cohere_command3", + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content=None, + ), + id="cmd3-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command4", + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content=None, + ), + id="cmd4-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command3", + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd3-citations_with_emoji", + ), + pytest.param( + ReasoningCase( + parser_name="cohere_command4", + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd4-citations_with_emoji", + ), +] + + +@pytest.mark.parametrize("case", REASONING_CASES) +class TestExtractReasoning: + def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase): + parser = make_parser(tokenizer, case.parser_name) + reasoning, content, tool_calls = parser.parse( + case.model_output, request_obj, enable_auto_tools=True + ) + assert reasoning == case.expected_reasoning + assert content == case.expected_content + # No tools on the request, so any tool call is dropped, not surfaced. + assert not tool_calls + + def test_streaming(self, tokenizer, request_obj, case: ReasoningCase): + parser = make_parser(tokenizer, case.parser_name) + deltas = stream_parser(parser, request_obj, tokenizer, case.model_output) + + reasoning = "".join(d.reasoning for d in deltas if d.reasoning) or None + content = "".join(d.content for d in deltas if d.content) or None + assert reasoning == case.expected_reasoning + assert content == case.expected_content + assert all(not d.tool_calls for d in deltas) + + +class TestFramingTokensStripped: + @pytest.mark.parametrize( + ("parser_name", "content_tags"), + [ + pytest.param( + "cohere_command4", ("<|START_TEXT|>", "<|END_TEXT|>"), id="cmd4" + ), + pytest.param( + "cohere_command3", ("<|START_RESPONSE|>", "<|END_RESPONSE|>"), id="cmd3" + ), + ], + ) + @pytest.mark.parametrize("chunk_size", [1, 4], ids=["per_token", "batched"]) + @pytest.mark.parametrize("with_tools", [False, True], ids=["no_tools", "tools"]) + def test_content_framing_tokens_stripped( + self, tokenizer, parser_name, content_tags, chunk_size, with_tools + ): + parser = make_parser(tokenizer, parser_name) + tools = [{"type": "function", "function": {"name": "foo"}}] + request = ChatCompletionRequest( + messages=[], model="test-model", tools=tools if with_tools else None + ) + parser.adjust_request(request) + + start_tag, end_tag = content_tags + generation = ( + f"<|START_THINKING|>Think deeply. The user greets us.<|END_THINKING|>" + f"{start_tag}I'm doing well, thank you{end_tag}" + ) + deltas = stream_parser(parser, request, tokenizer, generation, chunk_size) + + assert "".join(d.reasoning or "" for d in deltas) == ( + "Think deeply. The user greets us." + ) + assert not any(d.tool_calls for d in deltas) + assert "".join(d.content or "" for d in deltas) == "I'm doing well, thank you" + + +class TestIsReasoningEnd: + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser], + ids=["cmd3", "cmd4"], + ) + def test_is_reasoning_end(self, tokenizer, parser_cls): + parser = parser_cls(tokenizer) + start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") + content_ids = [99, 100] + + # Generation-only tokens have no chatbot marker, so the whole sequence + # is considered. + assert parser.is_reasoning_end([end_id]) + assert parser.is_reasoning_end([start_id, *content_ids, end_id]) + assert not parser.is_reasoning_end([start_id, *content_ids]) + + # Full prompt/history tokens are scoped to the latest chatbot marker, + # so stray thinking tokens from the preamble or previous turns are ignored. + assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids]) + assert parser.is_reasoning_end( + [start_id, end_id, chatbot_id, start_id, *content_ids, end_id] + ) + + +@pytest.mark.parametrize("parser_name", ["cohere_command3", "cohere_command4"]) +def test_count_reasoning_tokens(tokenizer, parser_name): + parser = make_parser(tokenizer, parser_name) + start = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + end = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + + assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2 + assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3 + assert parser.count_reasoning_tokens([start, 1, 2]) == 2 + assert parser.count_reasoning_tokens([1, 2, end]) == 0 + + +class TestParserSelection: + @pytest.mark.parametrize("with_tool_parser", [True, False]) + def test_adjust_request_keeps_framing_tokens(self, tokenizer, with_tool_parser): + """Melody needs the special tokens, with or without a tool parser.""" + cls = ParserManager.get_parser( + "cohere_command4" if with_tool_parser else None, + "cohere_command4", + enable_auto_tools=True, + ) + request = ChatCompletionRequest(messages=[], model="m") + assert request.skip_special_tokens + cls(tokenizer).adjust_request(request) + assert request.skip_special_tokens is False + + def test_non_cohere_selection_does_not_import_melody_parser(self, monkeypatch): + """``cohere_melody`` is optional; other parsers must not depend on it.""" + import sys + + monkeypatch.setitem(sys.modules, "vllm.parser.cohere_command", None) + assert ParserManager.get_parser("hermes", "deepseek_r1", True) is not None diff --git a/tests/parser/cohere/test_structural_tags.py b/tests/parser/cohere/test_structural_tags.py new file mode 100644 index 000000000000..27440d52af81 --- /dev/null +++ b/tests/parser/cohere/test_structural_tags.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``adjust_request``: folding tools and JSON schemas into Cohere structural tags.""" + +from __future__ import annotations + +import json +from collections import UserDict +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.generate.base.protocol import ( + JsonSchemaResponseFormat, + ResponseFormat, + StructuralTagResponseFormat, +) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.cohere_command import ( + CohereCommandParser, + _has_effective_tools, + _response_format_type, + _schema_dict_from_structured_outputs, + convert_schema_to_structural_tags, +) +from vllm.sampling_params import StructuredOutputsParams + +from .utils import MockCohereTokenizer, make_parser + +SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}} +SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}} +GET_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, +} +VALID_STRUCTURAL_TAG = { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "tags": [ + { + "begin": "", + "content": {"type": "any_text"}, + "end": "", + } + ], + "triggers": [""], + }, +} + + +def _model_config(arch: str) -> SimpleNamespace: + return SimpleNamespace( + architecture=arch, + architectures=[arch], + hf_text_config=SimpleNamespace(architectures=[arch]), + ) + + +def _make_chat_request(**kwargs) -> ChatCompletionRequest: + data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} + data.update(kwargs) + return ChatCompletionRequest.model_validate(data) + + +def _first_json_schema(tag_json: str) -> dict | None: + outer = json.loads(tag_json) + for t in (outer.get("format") or {}).get("tags") or []: + c = t.get("content") or {} + if c.get("type") == "json_schema": + js = c.get("json_schema") + return js if isinstance(js, dict) else None + return None + + +def _content_types(tag_json: str) -> set[str]: + outer = json.loads(tag_json) + out: set[str] = set() + for t in (outer.get("format") or {}).get("tags") or []: + ty = (t.get("content") or {}).get("type") + if isinstance(ty, str): + out.add(ty) + return out + + +@pytest.fixture(scope="module") +def parser(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser configured with a supported Cohere architecture.""" + return make_parser( + tokenizer, "cohere_command4", model_config=_model_config("Cohere2ForCausalLM") + ) + + +@pytest.fixture(scope="module") +def parser_no_model_config(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser with no ``model_config`` (cannot resolve architecture).""" + return make_parser(tokenizer, "cohere_command4") + + +@pytest.fixture(scope="module") +def parser_unsupported_arch(tokenizer: MockCohereTokenizer) -> CohereCommandParser: + """Parser configured with an architecture that has no structural tag style.""" + return make_parser( + tokenizer, "cohere_command4", model_config=_model_config("LlamaForCausalLM") + ) + + +class TestAdjustRequestPassthrough: + def test_structured_outputs_structural_tag_not_modified(self, parser) -> None: + tag = json.dumps(VALID_STRUCTURAL_TAG) + r = _make_chat_request(structured_outputs={"structural_tag": tag}) + o = parser.adjust_request(r) + assert o.structured_outputs.structural_tag == tag + + def test_response_format_structural_tag_short_circuit(self, parser) -> None: + # ``ChatCompletionRequest`` validates ``response_format`` as a union; + # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model). + rf = StructuralTagResponseFormat( + type="structural_tag", + format=VALID_STRUCTURAL_TAG["format"], + ) + r = _make_chat_request(response_format=rf) + o = parser.adjust_request(r) + assert _response_format_type(o.response_format) == "structural_tag" + assert o.structured_outputs is None + + +class TestAdjustRequestNoOp: + def test_no_schema_no_tools(self, parser) -> None: + o = parser.adjust_request(_make_chat_request()) + assert o.structured_outputs is None + assert o.response_format is None + + def test_no_model_config(self, parser_no_model_config) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + o = parser_no_model_config.adjust_request(r) + assert o.response_format is not None + assert o.structured_outputs is None + + +class TestAdjustRequestUnsupportedArchitecture: + def test_json_schema_raises(self, parser_unsupported_arch) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + with pytest.raises(ValueError, match="does not support"): + parser_unsupported_arch.adjust_request(r) + + +class TestAdjustRequestFoldFromResponseFormat: + @pytest.mark.parametrize( + "response_format, expected_schema", + [ + pytest.param( + ResponseFormat( + type="json_schema", + json_schema=JsonSchemaResponseFormat( + name="n", json_schema=SCHEMA_A + ), + ), + SCHEMA_A, + id="json_schema_pydantic", + ), + pytest.param( + { + "type": "json_schema", + "json_schema": {"name": "n", "schema": SCHEMA_A}, + }, + SCHEMA_A, + id="json_schema_dict", + ), + pytest.param( + {"type": "json_object"}, + {"type": "object"}, + id="json_object", + ), + ], + ) + def test_response_format_cleared( + self, parser, response_format, expected_schema + ) -> None: + r = _make_chat_request(response_format=response_format) + o = parser.adjust_request(r) + assert o.response_format is None + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + +class TestHasEffectiveTools: + @pytest.mark.parametrize( + "tools, expected", + [ + pytest.param(None, False, id="none"), + pytest.param([], False, id="empty_list"), + pytest.param(" ", False, id="blank_str"), + pytest.param( + [{"type": "function", "function": {"name": "f"}}], + True, + id="non_empty_list", + ), + pytest.param('{"x": 1}', True, id="non_empty_str"), + ], + ) + def test_has_effective_tools(self, tools, expected) -> None: + assert _has_effective_tools(tools) is expected + + def test_convert_schema_json_only_with_empty_tools_list(self) -> None: + tag = convert_schema_to_structural_tags( + schema=SCHEMA_B, + tools=[], + model_architecture="Cohere2ForCausalLM", + ) + assert tag is not None + assert _first_json_schema(tag) == SCHEMA_B + + +class TestAdjustRequestFoldFromStructuredOutputs: + @pytest.mark.parametrize( + "structured_outputs, expected_schema", + [ + pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"), + pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"), + pytest.param( + {"json_object": True}, {"type": "object"}, id="json_object_flag" + ), + pytest.param( + StructuredOutputsParams(json=SCHEMA_B), + SCHEMA_B, + id="structured_outputs_dataclass", + ), + pytest.param( + {"json": {"name": "n", "schema": SCHEMA_A}}, + SCHEMA_A, + id="openai_wrapper_dict_unwrapped", + ), + ], + ) + def test_structured_outputs_folded( + self, parser, structured_outputs, expected_schema + ) -> None: + o = parser.adjust_request( + _make_chat_request(structured_outputs=structured_outputs), + ) + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + def test_responses_request_default_empty_tools(self, parser) -> None: + """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``.""" + r = ResponsesRequest.model_validate( + { + "input": "hi", + "model": "m", + "structured_outputs": {"json": SCHEMA_B}, + } + ) + assert r.tools == [] + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B + + def test_json_userdict_mapping_unwrapped(self) -> None: + inner = {"type": "object", "properties": {"u": {"type": "number"}}} + so = StructuredOutputsParams(json=UserDict(inner)) + assert _schema_dict_from_structured_outputs(so) == inner + + @pytest.mark.parametrize( + "json_value, match", + [ + pytest.param("{not json}", "valid JSON", id="invalid_json_string"), + pytest.param( + json.dumps(["a", "b"]), "JSON object", id="non_object_json_string" + ), + pytest.param(" ", "empty", id="empty_json_string"), + ], + ) + def test_structured_outputs_json_string_raises( + self, parser, json_value, match + ) -> None: + with pytest.raises(ValueError, match=match): + parser.adjust_request( + _make_chat_request(structured_outputs={"json": json_value}), + ) + + @pytest.mark.parametrize( + "construct", + [ + pytest.param( + lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}), + id="chat_completion_request", + ), + pytest.param( + lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type] + id="structured_outputs_params", + ), + ], + ) + def test_json_wrong_type_raises(self, construct) -> None: + """Non-str / non-dict ``json`` fails at Pydantic validation.""" + with pytest.raises(ValidationError): + construct() + + +class TestAdjustRequestPrecedence: + def test_response_format_over_structured_outputs_json(self, parser) -> None: + s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}} + s_so = {"type": "object", "properties": {"so": {"type": "number"}}} + inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + structured_outputs={"json": s_so}, + ) + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf + + +class TestAdjustRequestTextPlusStructuredOutputs: + def test_text_response_format_preserved(self, parser) -> None: + sch = {"type": "object", "properties": {"k": {"type": "string"}}} + r = _make_chat_request( + response_format=ResponseFormat(type="text"), + structured_outputs={"json": sch}, + ) + o = parser.adjust_request(r) + assert o.response_format is not None + assert o.response_format.type == "text" + assert _first_json_schema(o.structured_outputs.structural_tag) == sch + + +class TestAdjustRequestTools: + def test_tools_only_command_a_grammar(self, parser) -> None: + o = parser.adjust_request( + _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"), + ) + assert "grammar" in _content_types(o.structured_outputs.structural_tag) + + def test_tools_plus_json_schema_both_kinds(self, parser) -> None: + inner = JsonSchemaResponseFormat( + name="n", + json_schema={"type": "object", "properties": {"r": {"type": "string"}}}, + ) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + tools=[GET_WEATHER_TOOL], + tool_choice="auto", + ) + o = parser.adjust_request(r) + types = _content_types(o.structured_outputs.structural_tag) + assert "grammar" in types + assert "json_schema" in types diff --git a/tests/parser/cohere/test_tool_calls.py b/tests/parser/cohere/test_tool_calls.py new file mode 100644 index 000000000000..a9fe7edb8db6 --- /dev/null +++ b/tests/parser/cohere/test_tool_calls.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tool-call parsing for the unified Cohere Command parser. + +The ``cohere_command3`` / ``cohere_command4`` tool parsers are registry shims; +``CohereCommandParser`` (resolved through ``ParserManager``) does the parsing. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.abstract_parser import Parser + +from .utils import MockCohereTokenizer, drive_parser, make_parser, token_deltas + +SPECIAL_TOKEN_MARKERS = ( + "<|START_THINKING|>", + "<|END_THINKING|>", + "<|START_RESPONSE|>", + "<|END_RESPONSE|>", + "<|START_ACTION|>", + "<|END_ACTION|>", + "<|START_TEXT|>", + "<|END_TEXT|>", +) + +DUMMY_TOOLS = [ + {"type": "function", "function": {"name": name, "parameters": {"type": "object"}}} + for name in ("foo", "bar") +] + +# (registry name, answer framing tag) +VARIANTS = (("cohere_command3", "RESPONSE"), ("cohere_command4", "TEXT")) + + +@dataclass +class ExpectedToolCall: + id: str + name: str + arguments: dict + + +@dataclass +class ToolCallCase: + parser_name: str + model_output: str + expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) + expected_reasoning: str | None = None + expected_content: str | None = None + + +SINGLE_TOOL_CALL = """\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""" +MULTI_TOOL_CALL = """\ +<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, + {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} +] +<|END_ACTION|>""" +CITATIONS_NO_TOOL_CALL = ( + "<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|>\n" + "<|START_{tag}|>foo bar<|END_{tag}|>" +) + +# (case id, model output template, expectations) +_CASE_TABLE = ( + ( + "single_tool_call", + SINGLE_TOOL_CALL, + dict( + expected_tool_calls=[ExpectedToolCall("0", "foo", {"query": "query1"})], + expected_reasoning="i will call foo with query1", + ), + ), + ( + "citations_no_tool_calls", + CITATIONS_NO_TOOL_CALL, + dict( + expected_reasoning="This is a rainbow emoji: 🌈", expected_content="foo bar" + ), + ), + ( + "multiple_tool_calls", + MULTI_TOOL_CALL, + dict( + expected_tool_calls=[ + ExpectedToolCall("0", "foo", {"query": "query1"}), + ExpectedToolCall("1", "bar", {"x": 42}), + ], + expected_reasoning="first I think about foo", + ), + ), + ( + "reasoning_only", + "<|START_THINKING|>just think, no response<|END_THINKING|>", + dict(expected_reasoning="just think, no response"), + ), +) + +TOOL_CALL_CASES = [ + pytest.param( + ToolCallCase(name, output.replace("{tag}", tag), **expected), + id=f"{name}-{case_id}", + ) + for name, tag in VARIANTS + for case_id, output, expected in _CASE_TABLE +] + +PARSER_NAMES = pytest.mark.parametrize("parser_name", [name for name, _ in VARIANTS]) + + +def _tool_parser(tokenizer: MockCohereTokenizer, parser_name: str) -> Parser: + return make_parser(tokenizer, parser_name, DUMMY_TOOLS) + + +def _tool_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + messages=[], model="test-model", tools=DUMMY_TOOLS, tool_choice="auto" + ) + + +@dataclass +class StreamingResult: + tool_calls: dict[int, dict] + reasoning: str | None + content: str | None + + +def run_streaming( + tokenizer: MockCohereTokenizer, + parser_name: str, + model_output: str, + chunk_size: int = 1, +) -> StreamingResult: + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, model_output, chunk_size), + ) + accumulated: dict[int, dict] = {} + for delta in deltas: + for tc in delta.tool_calls or []: + slot = accumulated.setdefault( + tc.index, {"id": "", "name": "", "arguments": ""} + ) + if tc.id: + slot["id"] = tc.id + if tc.function and tc.function.name: + slot["name"] = tc.function.name + if tc.function and tc.function.arguments: + slot["arguments"] += tc.function.arguments + return StreamingResult( + tool_calls=accumulated, + reasoning="".join(d.reasoning for d in deltas if d.reasoning) or None, + content="".join(d.content for d in deltas if d.content) or None, + ) + + +def run_nonstreaming( + tokenizer: MockCohereTokenizer, parser_name: str, model_output: str +) -> tuple[str | None, str | None, list]: + reasoning, content, tool_calls = _tool_parser(tokenizer, parser_name).parse( + model_output, _tool_request(), enable_auto_tools=True + ) + return reasoning, content, tool_calls or [] + + +@pytest.mark.parametrize("case", TOOL_CALL_CASES) +class TestExtractToolCalls: + def test_streaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert len(streamed.tool_calls) == len(case.expected_tool_calls) + for i, expected in enumerate(case.expected_tool_calls): + tc = streamed.tool_calls[i] + assert tc["id"] == expected.id + assert tc["name"] == expected.name + assert json.loads(tc["arguments"]) == expected.arguments + + def test_streaming_reasoning( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert streamed.reasoning == case.expected_reasoning + + def test_streaming_content( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + assert streamed.content == case.expected_content + + def test_nonstreaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase): + reasoning, content, tool_calls = run_nonstreaming( + tokenizer, case.parser_name, case.model_output + ) + assert reasoning == case.expected_reasoning + assert content == case.expected_content + assert len(tool_calls) == len(case.expected_tool_calls) + for actual, expected in zip(tool_calls, case.expected_tool_calls): + assert actual.name == expected.name + assert json.loads(actual.arguments) == expected.arguments + + def test_streaming_nonstreaming_agree( + self, tokenizer: MockCohereTokenizer, case: ToolCallCase + ): + streamed = run_streaming(tokenizer, case.parser_name, case.model_output) + _, _, tool_calls = run_nonstreaming( + tokenizer, case.parser_name, case.model_output + ) + assert len(streamed.tool_calls) == len(tool_calls) + for i, actual in enumerate(tool_calls): + assert streamed.tool_calls[i]["name"] == actual.name + assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads( + actual.arguments + ) + + +class TestSpeculativeDecodingMultiTokenDelta: + @PARSER_NAMES + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_no_special_token_leak_in_streaming_deltas( + self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int + ): + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, SINGLE_TOOL_CALL, chunk_size), + ) + for delta in deltas: + fields: list[tuple[str, str | None]] = [ + ("reasoning", delta.reasoning), + ("content", delta.content), + ] + for tc in delta.tool_calls or []: + if tc.function: + fields.append(("tool_call.name", tc.function.name)) + fields.append(("tool_call.arguments", tc.function.arguments)) + for marker in SPECIAL_TOKEN_MARKERS: + for field_name, value in fields: + assert value is None or marker not in value, ( + f"special token {marker!r} leaked into {field_name} " + f"with chunk_size={chunk_size} delta={delta!r}" + ) + + @PARSER_NAMES + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_multi_token_chunks_still_produce_correct_tool_call( + self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int + ): + streamed = run_streaming(tokenizer, parser_name, SINGLE_TOOL_CALL, chunk_size) + assert len(streamed.tool_calls) == 1 + tc = streamed.tool_calls[0] + assert tc["id"] == "0" + assert tc["name"] == "foo" + assert json.loads(tc["arguments"]) == {"query": "query1"} + + +class TestStreamingDeltaShape: + @PARSER_NAMES + def test_reasoning_and_tool_calls_are_separate_deltas( + self, tokenizer: MockCohereTokenizer, parser_name: str + ): + deltas = drive_parser( + _tool_parser(tokenizer, parser_name), + _tool_request(), + token_deltas(tokenizer, SINGLE_TOOL_CALL), + ) + for delta in deltas: + populated = [ + delta.content is not None, + delta.reasoning is not None, + bool(delta.tool_calls), + ] + assert sum(populated) == 1, ( + "A single streaming delta must carry exactly one of " + f"content/reasoning/tool_calls, got {delta!r}" + ) + assert any(d.reasoning is not None for d in deltas) + assert any(d.tool_calls for d in deltas) diff --git a/tests/parser/cohere/utils.py b/tests/parser/cohere/utils.py new file mode 100644 index 000000000000..f26df5296c6d --- /dev/null +++ b/tests/parser/cohere/utils.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared helpers for the unified Cohere Command parser tests.""" + +from __future__ import annotations + +import itertools + +import regex as re + +from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser import ParserManager +from vllm.parser.abstract_parser import Parser + +_SPECIAL_TOKEN_RE = re.compile(r"(<\|[A-Z_]+\|>)") +REPLACEMENT_CHAR = "�" + + +class MockCohereTokenizer: + """Byte-level stand-in for the Cohere tokenizer. + + Text round-trips through UTF-8 bytes so a multi-byte character can split + across "tokens" (reproducing trailing U+FFFD buffering), while each + ``<|...|>`` marker is a single id so token-id reasoning-end gating behaves + as with the real tokenizer. + """ + + _SPECIAL_TOKEN_IDS = { + tok: 256 + i + for i, tok in enumerate( + ( + "<|START_THINKING|>", + "<|END_THINKING|>", + "<|CHATBOT_TOKEN|>", + "<|START_RESPONSE|>", + "<|END_RESPONSE|>", + "<|START_TEXT|>", + "<|END_TEXT|>", + "<|START_ACTION|>", + "<|END_ACTION|>", + ) + ) + } + _ID_TO_SPECIAL_TOKEN = {v: k for k, v in _SPECIAL_TOKEN_IDS.items()} + + def convert_tokens_to_ids(self, token: str) -> int: + return self._SPECIAL_TOKEN_IDS.get(token, 0) + + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + ids: list[int] = [] + for part in _SPECIAL_TOKEN_RE.split(text): + if part in self._SPECIAL_TOKEN_IDS: + ids.append(self._SPECIAL_TOKEN_IDS[part]) + else: + ids.extend(part.encode("utf-8")) + return ids + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + out: list[str] = [] + for special, run in itertools.groupby( + ids, self._ID_TO_SPECIAL_TOKEN.__contains__ + ): + if not special: + out.append(bytes(run).decode("utf-8", errors="replace")) + elif not skip_special_tokens: + out.extend(self._ID_TO_SPECIAL_TOKEN[i] for i in run) + return "".join(out) + + +def token_deltas( + tokenizer: MockCohereTokenizer, text: str, chunk_size: int = 1 +) -> list[tuple[str, list[int]]]: + """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of + ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing + U+FFFD) into the next step as real streaming does.""" + ids = tokenizer.encode(text) + deltas: list[tuple[str, list[int]]] = [] + prev = "" + pending: list[int] = [] + for start in range(0, len(ids), chunk_size): + end = start + chunk_size + pending.extend(ids[start:end]) + current = tokenizer.decode(ids[:end]) + if current.endswith(REPLACEMENT_CHAR) and end < len(ids): + continue + deltas.append((current[len(prev) :], pending)) + prev, pending = current, [] + return deltas + + +def make_parser( + tokenizer: MockCohereTokenizer, + name: str, + tools: list[dict] | None = None, + chat_template_kwargs: dict | None = None, + model_config=None, +) -> Parser: + """Resolve the unified Cohere parser through ``ParserManager`` by registry + name (``cohere_command3`` / ``cohere_command4``).""" + cls = ParserManager.get_parser(name, name, enable_auto_tools=True) + assert cls is not None + return cls( + tokenizer, + tools, + model_config=model_config, + chat_template_kwargs=chat_template_kwargs or {}, + ) + + +def drive_parser( + parser: Parser, + request: ChatCompletionRequest, + deltas: list[tuple[str, list[int]]], +) -> list[DeltaMessage]: + """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and + return the non-``None`` messages in order.""" + out: list[DeltaMessage] = [] + for i, (text, ids) in enumerate(deltas): + delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1) + if delta is not None: + out.append(delta) + return out + + +def stream_parser(parser, request, tokenizer, text: str, chunk_size: int = 1): + """Stream ``text`` through ``parser`` in ``chunk_size``-token deltas.""" + return drive_parser(parser, request, token_deltas(tokenizer, text, chunk_size)) diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py deleted file mode 100644 index ea9ab63aef72..000000000000 --- a/tests/reasoning/test_cohere_command_reasoning_parser.py +++ /dev/null @@ -1,796 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -from collections import UserDict -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any - -import pytest -from pydantic import ValidationError - -from vllm.entrypoints.generate.base.protocol import ( - JsonSchemaResponseFormat, - ResponseFormat, - StructuralTagResponseFormat, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.reasoning.cohere_command_reasoning_parser import ( - CohereCommand3ReasoningParser, - CohereCommand4ReasoningParser, - _has_effective_tools, - _melody_citations_to_vllm, - _melody_sources_to_vllm, - _response_format_type, - _schema_dict_from_structured_outputs, - convert_schema_to_structural_tags, -) -from vllm.sampling_params import StructuredOutputsParams - - -@dataclass -class ReasoningCase: - parser_cls: Any - model_output: str - expected_reasoning: str | None - expected_content: str | None - - -REASONING_CASES = [ - pytest.param( - ReasoningCase( - parser_cls=CohereCommand3ReasoningParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_reasoning="i will call foo with query1", - expected_content="""\ -<|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - ), - id="cmd3-single_tool_call", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand4ReasoningParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_reasoning="i will call foo with query1", - expected_content="""\ -<|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - ), - id="cmd4-single_tool_call", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand3ReasoningParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd3-citations_with_emoji", - ), - pytest.param( - ReasoningCase( - parser_cls=CohereCommand4ReasoningParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd4-citations_with_emoji", - ), -] - - -class MockCohereTokenizer: - """Minimal byte-level stand-in for the Cohere tokenizer. - - ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a - multi-byte character (e.g. an emoji) across "tokens" reproduces the - trailing U+FFFD buffering that real streaming exhibits. Cohere special - tokens map to distinct synthetic ids; everything else shares a default id. - ``adjust_request`` only needs the token ids, not real tokenization. - """ - - _SPECIAL_TOKEN_IDS = { - "<|START_THINKING|>": -1, - "<|END_THINKING|>": -2, - "<|CHATBOT_TOKEN|>": -3, - } - - def convert_tokens_to_ids(self, token: str) -> int: - return self._SPECIAL_TOKEN_IDS.get(token, 0) - - def get_vocab(self) -> dict[str, int]: - return {} - - def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: - return list(text.encode("utf-8")) - - def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: - return bytes(ids).decode("utf-8", errors="replace") - - -@pytest.fixture(scope="module") -def tokenizer() -> MockCohereTokenizer: - return MockCohereTokenizer() - - -@pytest.fixture -def request_obj(): - return ChatCompletionRequest(messages=[], model="test-model") - - -REPLACEMENT_CHAR = "\ufffd" - - -def _token_deltas(tokenizer, text: str) -> list[str]: - """Progressively decode the token sequence and return per-step string - deltas. Incomplete multi-byte sequences (trailing U+FFFD) are buffered - until the next token completes them, matching real streaming behaviour.""" - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - for i in range(1, len(ids) + 1): - current = tokenizer.decode(ids[:i], skip_special_tokens=False) - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -@pytest.mark.parametrize("case", REASONING_CASES) -class TestExtractReasoning: - def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase): - parser = case.parser_cls(tokenizer) - reasoning, content = parser.extract_reasoning(case.model_output, request_obj) - - assert reasoning == case.expected_reasoning - assert content == case.expected_content - - def test_streaming(self, tokenizer, case: ReasoningCase): - parser = case.parser_cls(tokenizer) - token_strings = _token_deltas(tokenizer, case.model_output) - - reasoning_parts: list[str] = [] - content_parts: list[str] = [] - tool_call_deltas: list[dict] = [] - - previous_text = "" - previous_token_ids: list[int] = [] - - for token_str in token_strings: - current_text = previous_text + token_str - current_token_ids = previous_token_ids + [0] - - delta = parser.extract_reasoning_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=[0], - ) - if delta is not None: - if delta.reasoning is not None: - reasoning_parts.append(delta.reasoning) - if delta.content is not None: - content_parts.append(delta.content) - for tc in delta.tool_calls: - tool_call_deltas.append( - { - "id": tc.id, - "index": tc.index, - "name": tc.function.name if tc.function else None, - "arguments": ( - tc.function.arguments if tc.function else None - ), - } - ) - - previous_text = current_text - previous_token_ids = current_token_ids - - reasoning = "".join(reasoning_parts) if reasoning_parts else None - assert reasoning == case.expected_reasoning - - content = "".join(content_parts) if content_parts else None - assert content == case.expected_content - assert tool_call_deltas == [] - - -class TestIsReasoningEnd: - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser], - ids=["cmd3", "cmd4"], - ) - def test_is_reasoning_end(self, tokenizer, parser_cls): - parser = parser_cls(tokenizer) - start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") - end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") - chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") - content_ids = [99, 100] - - # Generation-only tokens have no chatbot marker, so the whole sequence - # is considered. - assert parser.is_reasoning_end([end_id]) - assert parser.is_reasoning_end([start_id, *content_ids, end_id]) - assert not parser.is_reasoning_end([start_id, *content_ids]) - - # Full prompt/history tokens are scoped to the latest chatbot marker, - # so stray thinking tokens from the preamble or previous turns are ignored. - assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids]) - assert parser.is_reasoning_end( - [start_id, end_id, chatbot_id, start_id, *content_ids, end_id] - ) - - -@pytest.mark.parametrize( - "parser_cls", - [ - CohereCommand3ReasoningParser, - CohereCommand4ReasoningParser, - ], - ids=["cmd3", "cmd4"], -) -def test_count_reasoning_tokens(tokenizer, parser_cls): - parser = parser_cls(tokenizer) - start = parser.start_token_id - end = parser.end_token_id - - assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2 - assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3 - assert parser.count_reasoning_tokens([start, 1, 2]) == 2 - assert parser.count_reasoning_tokens([1, 2, end]) == 0 - - -SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}} -SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}} -GET_WEATHER_TOOL = { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - }, -} -VALID_STRUCTURAL_TAG = { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "tags": [ - { - "begin": "", - "content": {"type": "any_text"}, - "end": "", - } - ], - "triggers": [""], - }, -} - - -def _model_config(arch: str) -> SimpleNamespace: - return SimpleNamespace( - architecture=arch, - architectures=[arch], - hf_text_config=SimpleNamespace(architectures=[arch]), - ) - - -def _make_chat_request(**kwargs) -> ChatCompletionRequest: - data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} - data.update(kwargs) - return ChatCompletionRequest.model_validate(data) - - -def _first_json_schema(tag_json: str) -> dict | None: - outer = json.loads(tag_json) - for t in (outer.get("format") or {}).get("tags") or []: - c = t.get("content") or {} - if c.get("type") == "json_schema": - js = c.get("json_schema") - return js if isinstance(js, dict) else None - return None - - -def _content_types(tag_json: str) -> set[str]: - outer = json.loads(tag_json) - out: set[str] = set() - for t in (outer.get("format") or {}).get("tags") or []: - ty = (t.get("content") or {}).get("type") - if isinstance(ty, str): - out.add(ty) - return out - - -@pytest.fixture(scope="module") -def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser: - """Parser configured with a supported Cohere architecture.""" - return CohereCommand4ReasoningParser( - tokenizer, - model_config=_model_config("Cohere2ForCausalLM"), - ) - - -@pytest.fixture(scope="module") -def parser_no_model_config( - tokenizer: MockCohereTokenizer, -) -> CohereCommand4ReasoningParser: - """Parser with no ``model_config`` (cannot resolve architecture).""" - return CohereCommand4ReasoningParser(tokenizer, model_config=None) - - -@pytest.fixture(scope="module") -def parser_unsupported_arch( - tokenizer: MockCohereTokenizer, -) -> CohereCommand4ReasoningParser: - """Parser configured with an architecture that has no structural tag style.""" - return CohereCommand4ReasoningParser( - tokenizer, - model_config=_model_config("LlamaForCausalLM"), - ) - - -class TestAdjustRequestPassthrough: - def test_structured_outputs_structural_tag_not_modified(self, parser) -> None: - tag = json.dumps(VALID_STRUCTURAL_TAG) - r = _make_chat_request(structured_outputs={"structural_tag": tag}) - o = parser.adjust_request(r) - assert o.structured_outputs.structural_tag == tag - - def test_response_format_structural_tag_short_circuit(self, parser) -> None: - # ``ChatCompletionRequest`` validates ``response_format`` as a union; - # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model). - rf = StructuralTagResponseFormat( - type="structural_tag", - format=VALID_STRUCTURAL_TAG["format"], - ) - r = _make_chat_request(response_format=rf) - o = parser.adjust_request(r) - assert _response_format_type(o.response_format) == "structural_tag" - assert o.structured_outputs is None - - -class TestAdjustRequestNoOp: - def test_no_schema_no_tools(self, parser) -> None: - o = parser.adjust_request(_make_chat_request()) - assert o.structured_outputs is None - assert o.response_format is None - - def test_no_model_config(self, parser_no_model_config) -> None: - inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - ) - o = parser_no_model_config.adjust_request(r) - assert o.response_format is not None - assert o.structured_outputs is None - - -class TestAdjustRequestUnsupportedArchitecture: - def test_json_schema_raises(self, parser_unsupported_arch) -> None: - inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - ) - with pytest.raises(ValueError, match="does not support"): - parser_unsupported_arch.adjust_request(r) - - -class TestAdjustRequestFoldFromResponseFormat: - @pytest.mark.parametrize( - "response_format, expected_schema", - [ - pytest.param( - ResponseFormat( - type="json_schema", - json_schema=JsonSchemaResponseFormat( - name="n", json_schema=SCHEMA_A - ), - ), - SCHEMA_A, - id="json_schema_pydantic", - ), - pytest.param( - { - "type": "json_schema", - "json_schema": {"name": "n", "schema": SCHEMA_A}, - }, - SCHEMA_A, - id="json_schema_dict", - ), - pytest.param( - {"type": "json_object"}, - {"type": "object"}, - id="json_object", - ), - ], - ) - def test_response_format_cleared( - self, parser, response_format, expected_schema - ) -> None: - r = _make_chat_request(response_format=response_format) - o = parser.adjust_request(r) - assert o.response_format is None - assert ( - _first_json_schema(o.structured_outputs.structural_tag) == expected_schema - ) - - -class TestHasEffectiveTools: - @pytest.mark.parametrize( - "tools, expected", - [ - pytest.param(None, False, id="none"), - pytest.param([], False, id="empty_list"), - pytest.param(" ", False, id="blank_str"), - pytest.param( - [{"type": "function", "function": {"name": "f"}}], - True, - id="non_empty_list", - ), - pytest.param('{"x": 1}', True, id="non_empty_str"), - ], - ) - def test_has_effective_tools(self, tools, expected) -> None: - assert _has_effective_tools(tools) is expected - - def test_convert_schema_json_only_with_empty_tools_list(self) -> None: - tag = convert_schema_to_structural_tags( - schema=SCHEMA_B, - tools=[], - model_architecture="Cohere2ForCausalLM", - ) - assert tag is not None - assert _first_json_schema(tag) == SCHEMA_B - - -class TestAdjustRequestFoldFromStructuredOutputs: - @pytest.mark.parametrize( - "structured_outputs, expected_schema", - [ - pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"), - pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"), - pytest.param( - {"json_object": True}, {"type": "object"}, id="json_object_flag" - ), - pytest.param( - StructuredOutputsParams(json=SCHEMA_B), - SCHEMA_B, - id="structured_outputs_dataclass", - ), - pytest.param( - {"json": {"name": "n", "schema": SCHEMA_A}}, - SCHEMA_A, - id="openai_wrapper_dict_unwrapped", - ), - ], - ) - def test_structured_outputs_folded( - self, parser, structured_outputs, expected_schema - ) -> None: - o = parser.adjust_request( - _make_chat_request(structured_outputs=structured_outputs), - ) - assert ( - _first_json_schema(o.structured_outputs.structural_tag) == expected_schema - ) - - def test_responses_request_default_empty_tools(self, parser) -> None: - """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``.""" - r = ResponsesRequest.model_validate( - { - "input": "hi", - "model": "m", - "structured_outputs": {"json": SCHEMA_B}, - } - ) - assert r.tools == [] - o = parser.adjust_request(r) - assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B - - def test_json_userdict_mapping_unwrapped(self) -> None: - inner = {"type": "object", "properties": {"u": {"type": "number"}}} - so = StructuredOutputsParams(json=UserDict(inner)) - assert _schema_dict_from_structured_outputs(so) == inner - - @pytest.mark.parametrize( - "json_value, match", - [ - pytest.param("{not json}", "valid JSON", id="invalid_json_string"), - pytest.param( - json.dumps(["a", "b"]), "JSON object", id="non_object_json_string" - ), - pytest.param(" ", "empty", id="empty_json_string"), - ], - ) - def test_structured_outputs_json_string_raises( - self, parser, json_value, match - ) -> None: - with pytest.raises(ValueError, match=match): - parser.adjust_request( - _make_chat_request(structured_outputs={"json": json_value}), - ) - - @pytest.mark.parametrize( - "construct", - [ - pytest.param( - lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}), - id="chat_completion_request", - ), - pytest.param( - lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type] - id="structured_outputs_params", - ), - ], - ) - def test_json_wrong_type_raises(self, construct) -> None: - """Non-str / non-dict ``json`` fails at Pydantic validation.""" - with pytest.raises(ValidationError): - construct() - - -class TestAdjustRequestPrecedence: - def test_response_format_over_structured_outputs_json(self, parser) -> None: - s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}} - s_so = {"type": "object", "properties": {"so": {"type": "number"}}} - inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - structured_outputs={"json": s_so}, - ) - o = parser.adjust_request(r) - assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf - - -class TestAdjustRequestTextPlusStructuredOutputs: - def test_text_response_format_preserved(self, parser) -> None: - sch = {"type": "object", "properties": {"k": {"type": "string"}}} - r = _make_chat_request( - response_format=ResponseFormat(type="text"), - structured_outputs={"json": sch}, - ) - o = parser.adjust_request(r) - assert o.response_format is not None - assert o.response_format.type == "text" - assert _first_json_schema(o.structured_outputs.structural_tag) == sch - - -class TestAdjustRequestTools: - def test_tools_only_command_a_grammar(self, parser) -> None: - o = parser.adjust_request( - _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"), - ) - assert "grammar" in _content_types(o.structured_outputs.structural_tag) - - def test_tools_plus_json_schema_both_kinds(self, parser) -> None: - inner = JsonSchemaResponseFormat( - name="n", - json_schema={"type": "object", "properties": {"r": {"type": "string"}}}, - ) - r = _make_chat_request( - response_format=ResponseFormat(type="json_schema", json_schema=inner), - tools=[GET_WEATHER_TOOL], - tool_choice="auto", - ) - o = parser.adjust_request(r) - types = _content_types(o.structured_outputs.structural_tag) - assert "grammar" in types - assert "json_schema" in types - - -class TestMelodySourceResolution: - """Pin the parser-side source resolution (see - ``_melody_sources_to_vllm``). The parser receives melody's numeric - ``(tool_call_index, tool_result_indices)`` addressing and resolves - it against a ``position_to_source`` map handed in by - ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same - map that would otherwise live in the serving layer's resolver. - """ - - @staticmethod - def _fake_melody_source(bucket: int, indices: list[int]) -> Any: - return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) - - def test_multi_index_source_fans_out(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), - (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), - } - raw = [self._fake_melody_source(0, [0, 1])] - out = _melody_sources_to_vllm(raw, position_map) - assert [s.id for s in out] == ["d0", "d1"] - # Verify type / payload were plumbed through, not just the id. - assert out[0].type == "document" - assert out[0].document == {"id": "d0"} - - def test_unresolvable_position_skipped(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0"), - } - raw = [self._fake_melody_source(9, [0])] - assert _melody_sources_to_vllm(raw, position_map) == [] - - def test_missing_position_map_drops_all_sources(self): - # A parser instance without a position map (parser wired - # outside of ``CohereServingChatV2``) can't attribute anything, - # so every source is dropped. Callers downstream will see the - # citation with empty ``sources`` and drop it entirely. - raw = [self._fake_melody_source(0, [0])] - assert _melody_sources_to_vllm(raw, None) == [] - - def test_citations_pass_through_is_thinking_tag(self): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - - position_map: dict[tuple[int, int], CitationSource] = { - (0, 0): CitationSource(type="document", id="d0"), - } - raw = [ - SimpleNamespace( - start_index=0, - end_index=5, - text="hello", - is_thinking=True, - sources=[self._fake_melody_source(0, [0])], - ) - ] - out = _melody_citations_to_vllm(raw, position_map) - assert out is not None - assert out[0].type == "THINKING_CONTENT" - assert out[0].sources[0].id == "d0" - - -class TestParserStreamingEndToEnd: - """End-to-end streaming shape: raw model output text is tokenized - with the mock byte-level tokenizer and fed one token at a time - through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`. - We collect the ordered sequence of ``DeltaMessage``s the parser - hands back -- i.e. what the client would receive on the wire -- - and pin the exact shape. - - Combines the three flavors that reach ``delta``: a thinking block - (populates ``reasoning``), a text block (populates ``content``), - and a resolved citation (populates ``citations``) referencing - ``(start, end)`` offsets into the already-streamed content. - """ - - _MODEL_OUTPUT = ( - "<|START_THINKING|>Let me check.<|END_THINKING|>" - "<|START_RESPONSE|>The capital is Paris.<|END_RESPONSE|>" - ) - - @staticmethod - def _drive_parser( - parser: CohereCommand3ReasoningParser, text: str, tok: MockCohereTokenizer - ) -> list[Any]: - """Feed ``text`` through the parser one token at a time - (byte-level) and return the ordered list of non-``None`` - ``DeltaMessage``s. - """ - deltas: list[Any] = [] - previous_text = "" - previous_token_ids: list[int] = [] - for token_str in _token_deltas(tok, text): - current_text = previous_text + token_str - current_token_ids = previous_token_ids + [0] - delta = parser.extract_reasoning_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=[0], - ) - if delta is not None: - deltas.append(delta) - previous_text = current_text - previous_token_ids = current_token_ids - return deltas - - def test_full_wire_stream_shape(self, tokenizer): - from vllm.entrypoints.cohere.cohere_chat_message import CitationSource - from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY - - # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs`` - # installs on a request whose tool_call_index=0 tool result at - # index 1 is the "France" document. - position_to_source = { - (0, 1): CitationSource( - type="document", - id="doc-paris", - document={"id": "doc-paris", "title": "France"}, - ), - } - - parser = CohereCommand3ReasoningParser( - tokenizer, - chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source}, - ) - - deltas = self._drive_parser(parser, self._MODEL_OUTPUT, tokenizer) - - # -- 1. Reasoning and content stream separately and reach - # the wire only as their real payload. Framing tokens - # (``<|START_THINKING|>`` etc.) are consumed by the parser's - # state machine and never produce a delta. - reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None] - content_stream = [d.content for d in deltas if d.content is not None] - citation_deltas = [d for d in deltas if getattr(d, "citations", None)] - - assert "".join(reasoning_stream) == "Let me check." - assert "".join(content_stream) == "The capital is Paris." - - # -- 2. Exactly one citation reaches the wire, resolved to the - # source from ``position_to_source`` (not the raw - # ``(bucket, idx)`` coordinates). - assert len(citation_deltas) == 1 - (cite_delta,) = citation_deltas - assert cite_delta.citations is not None - assert len(cite_delta.citations) == 1 - cite = cite_delta.citations[0] - assert cite.text == "Paris" - # Char offsets are into the accumulated ``content``: "The - # capital is " is 15 chars, "Paris" is 5. - assert cite.start == 15 - assert cite.end == 20 - assert len(cite.sources) == 1 - assert cite.sources[0].id == "doc-paris" - assert cite.sources[0].type == "document" - - # -- 3. The citation's ``(start, end)`` offsets point into - # bytes that were already emitted as content-deltas earlier - # in the stream -- i.e. the citation is anchored to real - # already-shipped content. - cite_delta_idx = deltas.index(cite_delta) - earlier_content = "".join( - d.content for d in deltas[:cite_delta_idx] if d.content is not None - ) - assert earlier_content[cite.start : cite.end] == "Paris" - - # Everything after the citation delta is the trailing period. - later_content = "".join( - d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None - ) - assert later_content == "." - - # -- 4. Reasoning and content are never mixed on the same - # delta: the parser flips modes on the ``<|END_THINKING|>`` - # boundary and never emits a delta with both fields set. - for d in deltas: - assert not (d.reasoning is not None and d.content is not None) diff --git a/tests/tool_parsers/test_cohere_command_tool_parser.py b/tests/tool_parsers/test_cohere_command_tool_parser.py deleted file mode 100644 index f102474b6289..000000000000 --- a/tests/tool_parsers/test_cohere_command_tool_parser.py +++ /dev/null @@ -1,510 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from typing import Any - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.tool_parsers.cohere_command_tool_parser import ( - CohereCommand3ToolParser, - CohereCommand4ToolParser, -) - - -@dataclass -class ExpectedToolCall: - id: str - name: str - arguments: dict - - -@dataclass -class ToolCallCase: - parser_cls: Any - model_output: str - expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) - expected_reasoning: str | None = None - expected_content: str | None = None - - -TOOL_CALL_CASES = [ - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], - expected_reasoning="i will call foo with query1", - ), - id="cmd3-single_tool_call", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], - expected_reasoning="i will call foo with query1", - ), - id="cmd4-single_tool_call", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd3-citations_no_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> -<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", - expected_reasoning="This is a rainbow emoji: 🌈", - expected_content="foo bar", - ), - id="cmd4-citations_no_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, - {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), - ], - expected_reasoning="first I think about foo", - ), - id="cmd3-multiple_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> -[ - {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, - {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} -] -<|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), - ], - expected_reasoning="first I think about foo", - ), - id="cmd4-multiple_tool_calls", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand3ToolParser, - model_output="""\ -<|START_THINKING|>just think, no response<|END_THINKING|>""", - expected_reasoning="just think, no response", - ), - id="cmd3-reasoning_only", - ), - pytest.param( - ToolCallCase( - parser_cls=CohereCommand4ToolParser, - model_output="""\ -<|START_THINKING|>just think, no response<|END_THINKING|>""", - expected_reasoning="just think, no response", - ), - id="cmd4-reasoning_only", - ), -] - - -class MockCohereTokenizer: - """Minimal byte-level stand-in for the Cohere tokenizer.""" - - def get_vocab(self) -> dict[str, int]: - return {} - - def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: - return list(text.encode("utf-8")) - - def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: - return bytes(ids).decode("utf-8", errors="replace") - - -@pytest.fixture(scope="module") -def tokenizer() -> MockCohereTokenizer: - return MockCohereTokenizer() - - -@pytest.fixture -def request_obj() -> ChatCompletionRequest: - return ChatCompletionRequest(messages=[], model="test-model") - - -REPLACEMENT_CHAR = "\ufffd" - - -def _token_deltas(tokenizer: MockCohereTokenizer, text: str) -> list[str]: - """Decode per-token string deltas, buffering incomplete multi-byte chars.""" - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - for i in range(1, len(ids) + 1): - current = tokenizer.decode(ids[:i], skip_special_tokens=False) - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -@dataclass -class StreamingResult: - tool_calls: dict[int, dict] - reasoning: str | None - content: str | None - - -def _run_streaming_over_deltas( - parser, - deltas: list[str], - request_obj: ChatCompletionRequest, -) -> StreamingResult: - accumulated: dict[int, dict] = {} - reasoning_parts: list[str] = [] - content_parts: list[str] = [] - previous_text = "" - - for token_str in deltas: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - if delta is not None: - if delta.reasoning is not None: - reasoning_parts.append(delta.reasoning) - if delta.content is not None: - content_parts.append(delta.content) - for tc in delta.tool_calls: - idx = tc.index - if idx not in accumulated: - accumulated[idx] = {"id": "", "name": "", "arguments": ""} - if tc.id: - accumulated[idx]["id"] = tc.id - if tc.function and tc.function.name: - accumulated[idx]["name"] = tc.function.name - if tc.function and tc.function.arguments: - accumulated[idx]["arguments"] += tc.function.arguments - previous_text = current_text - - return StreamingResult( - tool_calls=accumulated, - reasoning="".join(reasoning_parts) if reasoning_parts else None, - content="".join(content_parts) if content_parts else None, - ) - - -def _run_streaming( - parser, - tokenizer: MockCohereTokenizer, - model_output: str, - request_obj: ChatCompletionRequest, -) -> StreamingResult: - return _run_streaming_over_deltas( - parser, - _token_deltas(tokenizer, model_output), - request_obj, - ) - - -@pytest.mark.parametrize("case", TOOL_CALL_CASES) -class TestExtractToolCalls: - def test_streaming( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert len(streamed.tool_calls) == len(case.expected_tool_calls) - - for i, expected_tc in enumerate(case.expected_tool_calls): - tc = streamed.tool_calls[i] - assert tc["id"] == expected_tc.id - assert tc["name"] == expected_tc.name - assert json.loads(tc["arguments"]) == expected_tc.arguments - - def test_streaming_reasoning( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert streamed.reasoning == case.expected_reasoning - - def test_streaming_content( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) - - assert streamed.content == case.expected_content - - def test_nonstreaming( - self, - request_obj: ChatCompletionRequest, - tokenizer: MockCohereTokenizer, - case: ToolCallCase, - ): - parser = case.parser_cls(tokenizer) - result = parser.extract_tool_calls(case.model_output, request_obj) - - assert result.tools_called == (len(case.expected_tool_calls) > 0) - assert len(result.tool_calls) == len(case.expected_tool_calls) - - for actual_tc, expected_tc in zip(result.tool_calls, case.expected_tool_calls): - assert actual_tc.type == "function" - assert actual_tc.function.name == expected_tc.name - assert json.loads(actual_tc.function.arguments) == expected_tc.arguments - - def test_streaming_nonstreaming_agree( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - case: ToolCallCase, - ): - parser_streaming = case.parser_cls(tokenizer) - parser_nonstreaming = case.parser_cls(tokenizer) - - streamed = _run_streaming( - parser_streaming, - tokenizer, - case.model_output, - request_obj, - ) - result = parser_nonstreaming.extract_tool_calls( - case.model_output, - request_obj, - ) - - assert len(streamed.tool_calls) == len(result.tool_calls) - - for i, actual_tc in enumerate(result.tool_calls): - assert streamed.tool_calls[i]["name"] == actual_tc.function.name - assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads( - actual_tc.function.arguments - ) - - -SPECIAL_TOKEN_MARKERS = ( - "<|START_THINKING|>", - "<|END_THINKING|>", - "<|START_RESPONSE|>", - "<|END_RESPONSE|>", - "<|START_ACTION|>", - "<|END_ACTION|>", - "<|START_TEXT|>", - "<|END_TEXT|>", -) - - -def _multi_token_deltas( - tokenizer: MockCohereTokenizer, - text: str, - chunk_size: int, -) -> list[str]: - ids = tokenizer.encode(text, add_special_tokens=False) - deltas: list[str] = [] - prev = "" - i = 0 - while i < len(ids): - end = min(len(ids), i + chunk_size) - current = tokenizer.decode(ids[:end], skip_special_tokens=False) - i = end - if current.endswith(REPLACEMENT_CHAR): - continue - delta = current[len(prev) :] - if delta: - deltas.append(delta) - prev = current - return deltas - - -class TestSpeculativeDecodingMultiTokenDelta: - MODEL_OUTPUT = ( - "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" - "<|START_ACTION|>\n" - '[\n {"tool_call_id": "0", "tool_name": "foo", ' - '"parameters": {"query": "query1"}}\n]\n' - "<|END_ACTION|>" - ) - - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) - def test_no_special_token_leak_in_streaming_deltas( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - chunk_size: int, - ): - parser = parser_cls(tokenizer) - chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) - - previous_text = "" - for token_str in chunked: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - previous_text = current_text - if delta is None: - continue - - fields: list[tuple[str, str | None]] = [ - ("reasoning", delta.reasoning), - ("content", delta.content), - ] - for tc in delta.tool_calls or []: - if tc.function: - fields.append(("tool_call.name", tc.function.name)) - fields.append(("tool_call.arguments", tc.function.arguments)) - - for marker in SPECIAL_TOKEN_MARKERS: - for field_name, value in fields: - assert value is None or marker not in value, ( - f"special token {marker!r} leaked into {field_name} " - f"with chunk_size={chunk_size} delta={delta!r}" - ) - - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) - def test_multi_token_chunks_still_produce_correct_tool_call( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - chunk_size: int, - ): - parser = parser_cls(tokenizer) - chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) - streamed = _run_streaming_over_deltas(parser, chunked, request_obj) - - assert len(streamed.tool_calls) == 1 - tc = streamed.tool_calls[0] - assert tc["id"] == "0" - assert tc["name"] == "foo" - assert json.loads(tc["arguments"]) == {"query": "query1"} - - -class TestStreamingDeltaShape: - @pytest.mark.parametrize( - "parser_cls", - [CohereCommand3ToolParser, CohereCommand4ToolParser], - ids=["cmd3", "cmd4"], - ) - def test_reasoning_and_tool_calls_are_separate_deltas( - self, - tokenizer: MockCohereTokenizer, - request_obj: ChatCompletionRequest, - parser_cls, - ): - parser = parser_cls(tokenizer) - model_output = ( - "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" - "<|START_ACTION|>\n" - '[\n {"tool_call_id": "0", "tool_name": "foo", ' - '"parameters": {"query": "query1"}}\n]\n' - "<|END_ACTION|>" - ) - - token_strings = _token_deltas(tokenizer, model_output) - previous_text = "" - saw_reasoning = False - saw_tool_call = False - - for token_str in token_strings: - current_text = previous_text + token_str - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=token_str, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request_obj, - ) - if delta is not None: - populated = [ - delta.content is not None, - delta.reasoning is not None, - bool(delta.tool_calls), - ] - assert sum(populated) == 1, ( - "A single streaming delta must carry exactly one of " - f"content/reasoning/tool_calls, got {delta!r}" - ) - if delta.reasoning is not None: - saw_reasoning = True - if delta.tool_calls: - saw_tool_call = True - previous_text = current_text - - assert saw_reasoning, "expected at least one reasoning delta" - assert saw_tool_call, "expected at least one tool-call delta" diff --git a/vllm/entrypoints/cohere/cohere_chat_message.py b/vllm/entrypoints/cohere/cohere_chat_message.py index 042856dea68c..85925d0f1878 100644 --- a/vllm/entrypoints/cohere/cohere_chat_message.py +++ b/vllm/entrypoints/cohere/cohere_chat_message.py @@ -7,8 +7,8 @@ It provides two things: * :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation - representation produced by the Cohere reasoning parser - (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by + representation produced by the Cohere parser + (:mod:`vllm.parser.cohere_command`) and consumed by :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere SDK shape -- that conversion happens in the serving layer. * :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage` @@ -42,9 +42,9 @@ class CitationSource(OpenAIBaseModel): document/tool-output identifier; ``document`` and ``tool_output`` carry the original payload that produced the citation. - Sources are fully resolved by the reasoning parser at emit time + Sources are fully resolved by the parser at emit time (see :func:`_melody_sources_to_vllm` in - :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the + :mod:`vllm.parser.cohere_command`) using the ``POSITION_TO_SOURCE_KEY`` map forwarded through ``chat_template_kwargs`` by :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`. @@ -121,8 +121,7 @@ def _serialize(self, handler): class CohereDeltaMessage(DeltaMessage): """:class:`DeltaMessage` extension carrying grounding citations for streaming. - Emitted by - :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser` + Emitted by :class:`vllm.parser.cohere_command.CohereCommandParser` on delta events whose payload includes citations. Non-Cohere parsers return plain :class:`DeltaMessage`, so their streamed shape is unchanged. diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 8ec4701907ed..788558645243 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -714,8 +714,8 @@ def _apply_cohere_template_kwargs( # emits, so citation deltas already carry SDK-shape sources by # the time they reach the streaming loop. See # ``_build_position_to_source`` for the numbering rule and - # ``_melody_sources_to_vllm`` in the reasoning parser for the - # consumer. + # ``_melody_sources_to_vllm`` in ``vllm.parser.cohere_command`` for + # the consumer. position_to_source = cls._build_position_to_source(request) if position_to_source: kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source) @@ -1109,12 +1109,11 @@ def _finalize_response_message( *, parser: Parser | None, ) -> ChatMessage: - """Copy grounding citations off the reasoning parser onto the message. + """Copy grounding citations off the unified parser onto the message. - The Cohere reasoning parser - (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the - citations produced by its most recent unary ``extract_reasoning`` - call on ``last_unary_citations``. We surface them here on + :class:`vllm.parser.cohere_command.CohereCommandParser` caches the + citations produced by its most recent unary ``parse`` call on + ``last_unary_citations``. We surface them here on :class:`CohereChatMessage` so downstream response conversion can pick them up without the base :class:`OpenAIServingChat` having to know about citations. @@ -1128,11 +1127,7 @@ def _finalize_response_message( # ``model_dump`` via the extras bucket, so the wire is correct # either way. message = cast(CohereChatMessage, message) - citations = getattr( - getattr(parser, "reasoning_parser", None), - "last_unary_citations", - None, - ) + citations = getattr(parser, "last_unary_citations", None) if citations: message.citations = citations return message @@ -1143,10 +1138,10 @@ def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None: :class:`CohereChatMessage` carries a ``citations: list[vllm...Citation] | None`` field populated by :meth:`_finalize_response_message` (which reads it off the - reasoning parser's ``last_unary_citations`` cache). Sources are + unified parser's ``last_unary_citations`` cache). Sources are already fully resolved by the parser (see :func:`_melody_sources_to_vllm` in - :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using + :mod:`vllm.parser.cohere_command`) using the position map forwarded via ``chat_template_kwargs``. This method's remaining job is: diff --git a/vllm/parser/cohere_command.py b/vllm/parser/cohere_command.py new file mode 100644 index 000000000000..cf8d79b10ec5 --- /dev/null +++ b/vllm/parser/cohere_command.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere Command parser. + +Uses ``cohere_melody`` as a stateful parser filter. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard + +import regex as re +import xgrammar as xgr + +try: + from cohere_melody import ( + AccumulatedToolCall, + FilterAggregatedResult, + PyFilter, + PyFilterOptions, + ) +except ImportError as e: + raise ImportError( + "The Cohere Command parser requires the `cohere_melody` package. " + "Install it with:\n pip install 'cohere-melody>=0.14.0'" + ) from e + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation, + CitationSource, + CohereDeltaMessage, +) +from vllm.entrypoints.generate.base.protocol import ( + AnyResponseFormat, + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionCall, +) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.cohere_command_reasoning_parser import ( + BaseCohereCommandReasoningParser, +) +from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY +from vllm.sampling_params import StructuredOutputsParams +from vllm.tool_parsers.cohere_command_tool_parser import BaseCohereCommandToolParser + +if TYPE_CHECKING: + # Declared in melody's stub only; not exported at runtime. + from cohere_melody import FilterCitation, Source + + from vllm.config import ModelConfig + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.utils import Tool + +PositionToSource = Mapping[tuple[int, int], CitationSource] + + +class CohereTagRegistry(NamedTuple): + """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``).""" + + trigger: str + end: str + + +class CohereTagStyle(NamedTuple): + """The structural tags style for a given model architecture. + + ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses + both response and text delimiters). ``tools`` is the tool-call wrapper. + """ + + json_tags: tuple[CohereTagRegistry, ...] + tools: CohereTagRegistry + + +class CohereNormalizedTool(TypedDict): + """A tool definition normalized to the shape ``collect_tool_schema`` expects. + + ``parameters`` is a JSON Schema object (possibly empty) describing the tool's + call signature. + """ + + name: str + parameters: dict[str, Any] + + +COMMAND_A_TOOLS_TAG = CohereTagRegistry( + trigger="<|START_ACTION|>", + end="<|END_ACTION|>", +) +COMMAND_A_JSON_TAG = CohereTagRegistry( + trigger="<|START_RESPONSE|>", + end="<|END_RESPONSE|>", +) +COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry( + trigger="<|START_TEXT|>", + end="<|END_TEXT|>", +) + +MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = { + "Cohere2ForCausalLM": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG,), + tools=COMMAND_A_TOOLS_TAG, + ), + "Cohere2VisionForConditionalGeneration": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), + tools=COMMAND_A_TOOLS_TAG, + ), + "Cohere2MoeForCausalLM": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), + tools=COMMAND_A_TOOLS_TAG, + ), +} + + +def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: + """Build an xgrammar EBNF grammar that matches a JSON array of tool calls. + + The grammar shape is architecture-independent; callers are responsible for + wrapping it in the correct structural tag (see ``CohereTagStyle.tools``). + """ + tool_dictionary: dict[str, str] = {} + for tool in tool_schema: + tool_name = tool["name"] + tool_parameters = json.dumps(tool["parameters"]) + json_schema = f"""{{ + "type": "object", + "properties": {{ + "tool_call_id": {{ + "type": "string", + "pattern": "^[0-9]+$" + }}, + "tool_name": {{ + "type": "string", + "const": "{tool_name}" + }}, + "parameters": {tool_parameters} + }} + }}""" + tool_grammar = str(xgr.Grammar.from_json_schema(json_schema)) + for match in re.findall(r"\b(\w+)\s*::=", tool_grammar): + tool_grammar = re.sub( + rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar + ) + tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}" + # Emitted grammar shape: + # root ::= tools + # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws + # ws ::= (" " | "\t" | "\n")* + # tool ::= | | ... (one alternative per input) + # ::= root (per-tool xgrammar rules) + # root ::= ... (from xgr.Grammar.from_json_schema) + tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys()) + tool_rules = "\n ".join(tool_dictionary.values()) + grammar = f"""root ::= tools + tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws + ws ::= (" " | "\\t" | "\\n")* + {tool_alternatives} + {tool_rules} + """ + return grammar + + +def _tool_definitions_to_schema_list( + tools: str | list[Any], +) -> list[CohereNormalizedTool]: + """ + Build the list of ``CohereNormalizedTool`` dicts expected by + ``collect_tool_schema``. + + Accepts: + - JSON string + - list of dicts with top-level ``name`` / ``parameters`` + - list of Chat Completions-style ``{"type": "function", "function": {...}}`` + - list of Pydantic models with ``model_dump()`` + """ + if isinstance(tools, str): + try: + parsed = json.loads(tools) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + else: + parsed = list(tools) + + out: list[CohereNormalizedTool] = [] + for raw in parsed: + t = raw.model_dump() if hasattr(raw, "model_dump") else raw + if not isinstance(t, dict): + continue + # Unwrap Chat Completions' ``{"type": "function", "function": {...}}`` + # shape; otherwise take the dict as-is. + if t.get("type") == "function" and isinstance(t.get("function"), dict): + t = t["function"] + name = t.get("name") + if not isinstance(name, str): + continue + params = t.get("parameters") + out.append( + CohereNormalizedTool( + name=name, + parameters=params if isinstance(params, dict) else {}, + ) + ) + return out + + +def _has_effective_tools( + tools: str | list[Any] | None, +) -> TypeGuard[str | list[Any]]: + """ + True when ``tools`` contains at least one tool definition to convert. + + ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest`` + uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are + treated as effective only when non-blank. + """ + if tools is None: + return False + if isinstance(tools, str): + return bool(tools.strip()) + return len(tools) > 0 + + +# Builder: produces vLLM response_format in xgrammar's canonical format. +# See xgrammar docs: type "structural_tag" with "format" = triggered_tags +# and tag content type = json_schema | grammar. +def convert_schema_to_structural_tags( + schema: dict | None = None, + tools: str | list[Any] | None = None, + model_architecture: str | None = None, +) -> str | None: + """ + Returns a response_format string accepted by xgrammar's structural tag format. + Uses the canonical shape: {"type": "structural_tag", "format": {...}} with + format.type "triggered_tags" and tag content type "json_schema" or "grammar". + + Callers that are not on an engine path (e.g. the reasoning parser) must pass + ``model_architecture`` explicitly. + """ + if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE: + return None + style = MODEL_TO_TAG_STYLE[model_architecture] + + tags: list[dict] = [] + triggers: list[str] = [] + + def _add_tag(tag: CohereTagRegistry, content: dict) -> None: + tags.append({"begin": tag.trigger, "content": content, "end": tag.end}) + triggers.append(tag.trigger) + + if schema is not None: + # One structural tag per JSON wrapper (e.g. MOE: response + text). + # Same for schema-only and "tools plus JSON mode" (North: schema when + # the model does not call tools). + for jt in style.json_tags: + _add_tag(jt, {"type": "json_schema", "json_schema": schema}) + + if _has_effective_tools(tools): + # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS) + # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts). + tool_schema_list = _tool_definitions_to_schema_list(tools) + if not tool_schema_list: + raise ValueError( + "No valid tool definitions could be parsed from the request for " + "structural tag conversion." + ) + tool_grammar = collect_tool_schema(tool_schema_list) + _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar}) + + if not tags: + return None + return json.dumps( + { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "triggers": triggers, + "tags": tags, + }, + } + ) + + +def _response_format_type( + response_format: AnyResponseFormat | dict | None, +) -> str | None: + if response_format is None: + return None + if isinstance(response_format, dict): + t = response_format.get("type") + return t if isinstance(t, str) else None + return response_format.type + + +def _maybe_parse_json_dict(value: Any) -> dict | None: + """If value is a JSON string, parse to dict; otherwise require dict.""" + if isinstance(value, dict): + return value + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, json.JSONDecodeError): + return None + return parsed if isinstance(parsed, dict) else None + return None + + +def _unwrap_nested_schema(candidate: Any) -> dict | None: + """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present. + + Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict. + """ + cand = _maybe_parse_json_dict(candidate) + if not isinstance(cand, dict): + return None + nested = cand.get("schema") + return nested if isinstance(nested, dict) else cand + + +def _schema_from_json_schema_field(js_wr: Any) -> dict | None: + """ + Extract the JSON Schema object from Chat Completions ``json_schema`` payload. + + Accepts: + - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field + - dict in OpenAI shape ``{"name": ..., "schema": {...}}`` + - dict with ``json_schema`` key holding either the schema or a nested wrapper + - dict that is already a JSON Schema document (some clients omit the wrapper) + - JSON strings for any of the above + """ + if js_wr is None: + return None + + parsed_wr = _maybe_parse_json_dict(js_wr) + if parsed_wr is not None: + js_wr = parsed_wr + + if hasattr(js_wr, "model_dump"): + for by_alias in (True, False): + try: + data = js_wr.model_dump(by_alias=by_alias, exclude_none=False) + except TypeError: + data = js_wr.model_dump(by_alias=by_alias) + out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema")) + if out is not None: + return out + inner_attr = getattr(js_wr, "json_schema", None) + return inner_attr if isinstance(inner_attr, dict) else None + + if isinstance(js_wr, dict): + for key in ("schema", "json_schema"): + out = _unwrap_nested_schema(js_wr.get(key)) + if out is not None: + return out + return js_wr + + return None + + +def _schema_dict_from_chat_response_format( + rf: AnyResponseFormat | dict | None, +) -> dict | None: + """JSON schema dict from Chat Completions ``request.response_format`` only.""" + if rf is None: + return None + rf_type = _response_format_type(rf) + if rf_type == "json_object": + return {"type": "object"} + if rf_type != "json_schema": + return None + js_wr = ( + rf.get("json_schema") + if isinstance(rf, dict) + else getattr(rf, "json_schema", None) + ) + return _schema_from_json_schema_field(js_wr) + + +def _schema_dict_from_structured_outputs( + so: StructuredOutputsParams | None, +) -> dict | None: + """Schema dict from ``structured_outputs`` (``json`` / ``json_object``). + + Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or + ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other + types raise ``ValueError`` only if a caller bypasses that validation. + """ + if so is None: + return None + if so.json_object: + return {"type": "object"} + raw: Any = so.json + if raw is None: + return None + + if hasattr(raw, "model_dump"): + out = _schema_from_json_schema_field(raw) + if out is None: + raise ValueError( + "structured_outputs.json model has no extractable JSON Schema." + ) + return out + + if isinstance(raw, str): + if not raw.strip(): + raise ValueError("structured_outputs.json cannot be empty.") + try: + raw = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError("structured_outputs.json must be valid JSON.") from e + if not isinstance(raw, dict): + raise ValueError("structured_outputs.json must decode to a JSON object.") + + if isinstance(raw, Mapping): + body = raw if isinstance(raw, dict) else dict(raw) + return _schema_from_json_schema_field(body) or body + + raise ValueError( + f"structured_outputs.json has unsupported type {type(raw).__name__}." + ) + + +def _melody_sources_to_vllm( + sources: list[Source], position_to_source: PositionToSource | None +) -> list[CitationSource]: + """Resolve melody ``Source`` addresses into :class:`CitationSource` objects. + + Melody reports each source as ``(tool_call_index, tool_result_indices)``. + ``position_to_source`` is built by + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source` + and forwarded through ``chat_template_kwargs``; each index fans out to one + resolved source, and unresolved positions are skipped. Without a map (a + request that did not come through the Cohere v2 handler) every source is + dropped and the serving layer discards the citation at wire-coercion time. + """ + if not position_to_source: + return [] + out: list[CitationSource] = [] + for source in sources: + for idx in source.tool_result_indices: + info = position_to_source.get((source.tool_call_index, idx)) + if info is not None: + out.append(info) + return out + + +def _melody_citations_to_vllm( + citations: list[FilterCitation], position_to_source: PositionToSource | None +) -> list[Citation] | None: + """Convert melody ``FilterCitation`` objects into :class:`Citation`. + + Citations whose sources all fail to resolve are still returned with + ``sources=[]``; the serving layer drops them so the fail-closed policy + lives in one place. + """ + if not citations: + return None + return [ + Citation( + start=c.start_index, + end=c.end_index, + text=c.text, + sources=_melody_sources_to_vllm(c.sources, position_to_source), + type="THINKING_CONTENT" if c.is_thinking else "TEXT_CONTENT", + ) + for c in citations + ] + + +class CohereCommandParser(DelegatingParser): + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + model_config: ModelConfig | None = None, + **kwargs, + ): + super().__init__(tokenizer, tools, *args, model_config=model_config, **kwargs) + self._model_config = model_config + self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + shim = self._reasoning_parser or self._tool_parser + assert isinstance( + shim, (BaseCohereCommandReasoningParser, BaseCohereCommandToolParser) + ) + options = getattr(PyFilterOptions(), shim.melody_preset)() + self._streaming = PyFilter(options) + self._unary = PyFilter(options) + # Citations from the most recent ``parse`` call, surfaced on + # ``CohereChatMessage`` by ``CohereServingChatV2``. + self.last_unary_citations: list[Citation] | None = None + # Melody may emit a tool-call id before the function name; hold it until + # the first name delta so clients receive both together. + self._pending_tool_ids: dict[int, str] = {} + ctk = kwargs.get("chat_template_kwargs") or {} + raw_map = ctk.get(POSITION_TO_SOURCE_KEY) + self._position_to_source: PositionToSource | None = ( + raw_map if isinstance(raw_map, Mapping) else None + ) + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + # keep special tokens to ensure melody can correctly identify boundaries. + request.skip_special_tokens = False + request = self._apply_structural_tags(request) + return super().adjust_request(request) + + def _apply_structural_tags( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Fold ``tools`` and any JSON schema into a Cohere ``structural_tag``.""" + so = request.structured_outputs + if so is not None and so.structural_tag: + return request + # Schema: prefer ``response_format`` (OpenAI Chat Completions), then + # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay + # on ``request.tools``. + rf = ( + request.response_format + if isinstance(request, ChatCompletionRequest) + else None + ) + if rf is not None and _response_format_type(rf) == "structural_tag": + return request + model_architecture = ( + self._model_config.architecture if self._model_config is not None else None + ) + tools = request.tools + # ``response_format`` wins if both it and ``structured_outputs`` supply JSON. + schema = _schema_dict_from_chat_response_format(rf) + if schema is None: + schema = _schema_dict_from_structured_outputs(so) + if schema is None and not _has_effective_tools(tools): + return request + if model_architecture is None: + return request + result = convert_schema_to_structural_tags( + schema=schema, + tools=tools, + model_architecture=model_architecture, + ) + if result is None: + # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``. + raise ValueError( + "Failed to build structural_tag guided decoding constraints from " + "this request's JSON schema and/or tools. The configured model " + f"architecture ({model_architecture!r}) does not support Cohere " + "command structural tags, or the schema cannot be expressed in " + "that format." + ) + request.structured_outputs = StructuredOutputsParams(structural_tag=result) + # Folded JSON constraints into ``structural_tag``; drop ``response_format`` + # when it was the source so ``to_sampling_params`` does not also set ``json`` / + # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``). + if isinstance(request, ChatCompletionRequest) and rf is not None: + rf_type = _response_format_type(rf) + if rf_type in ("json_schema", "json_object"): + request.response_format = None + return request + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + count = 0 + depth = 0 + for token_id in token_ids: + if token_id == self.start_token_id: + depth += 1 + continue + if token_id == self.end_token_id: + if depth > 0: + depth -= 1 + continue + if depth > 0: + count += 1 + return count + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + result = self._unary.process_full_text(model_output) + self.last_unary_citations = _melody_citations_to_vllm( + result.citations, self._position_to_source + ) + tool_calls: list[FunctionCall] | None = None + if self._tools_allowed(request): + tool_calls = [ + FunctionCall(id=tc.id, name=tc.name, arguments=tc.arguments) + for tc in result.tool_calls + ] or None + return result.reasoning, result.content, tool_calls + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + results = [self._streaming.write_decoded(delta_text)] + if finished: + results.append(self._streaming.flush_partials()) + return self._to_delta(results, request) + + def _tools_allowed(self, request: ChatCompletionRequest | ResponsesRequest) -> bool: + return self._tool_parser is not None and request.tool_choice != "none" + + def _to_delta( + self, + results: list[FilterAggregatedResult], + request: ChatCompletionRequest | ResponsesRequest, + ) -> CohereDeltaMessage | None: + content = self._join(r.content for r in results) + reasoning = ( + self._join(r.reasoning for r in results) + if request.include_reasoning + else None + ) + tool_calls: list[DeltaToolCall] = [] + citations: list[Citation] = [] + for result in results: + if result.tool_calls and self._tools_allowed(request): + tool_calls.extend(self._tool_call_deltas(result.tool_calls)) + citations.extend( + _melody_citations_to_vllm(result.citations, self._position_to_source) + or [] + ) + if not (content or reasoning or tool_calls or citations): + return None + # ``CohereDeltaMessage`` carries citations to the wire via + # ``SerializeAsAny[DeltaMessage]``; it drops the field when unset so the + # shape matches a plain ``DeltaMessage``. + message = CohereDeltaMessage() + if content is not None: + message.content = content + if reasoning is not None: + message.reasoning = reasoning + if tool_calls: + message.tool_calls = tool_calls + if citations: + message.citations = citations + return message + + @staticmethod + def _join(parts: Any) -> str | None: + present = [p for p in parts if p is not None] + return "".join(present) if present else None + + def _tool_call_deltas( + self, melody_tool_calls: list[AccumulatedToolCall] + ) -> list[DeltaToolCall]: + tool_calls: list[DeltaToolCall] = [] + for tc in melody_tool_calls: + name = tc.name or None + arguments = tc.arguments or None + if tc.id: + self._pending_tool_ids[tc.index] = tc.id + # Empty strings are melody streaming placeholders; skip fragments + # that carry neither a name nor argument text. + if name is None and arguments is None: + continue + function = DeltaFunctionCall() + if name is not None: + function.name = name + if arguments is not None: + function.arguments = arguments + delta = DeltaToolCall(index=tc.index, function=function) + if name is not None: + resolved_id = tc.id or self._pending_tool_ids.pop(tc.index, None) + if resolved_id is not None: + delta.id = resolved_id + delta.type = "function" + tool_calls.append(delta) + return tool_calls diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 420873a58da8..cc9dd3a6c659 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -126,6 +126,21 @@ class _KimiK3Parser(KimiK3Parser): return _KimiK3Parser + if {reasoning_parser_name, tool_parser_name} & { + "cohere_command3", + "cohere_command4", + }: + from vllm.parser.cohere_command import CohereCommandParser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _CohereCommandParser(CohereCommandParser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + + return _CohereCommandParser + from vllm.parser.abstract_parser import DelegatingParser r_cls = reasoning_parser_cls diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 4dada76ee7c2..d02091e6003e 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -3,531 +3,28 @@ from __future__ import annotations -import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar -import regex as re -import xgrammar as xgr - -try: - from cohere_melody import PyFilter, PyFilterOptions -except ImportError as e: - raise ImportError( - "The Cohere reasoning parser requires the `cohere_melody` " - "package, which is not installed. Install it with:\n" - " pip install 'cohere-melody>=0.11.1'" - ) from e - - -from vllm.entrypoints.cohere.cohere_chat_message import ( - Citation, - CitationSource, - CohereDeltaMessage, -) -from vllm.entrypoints.generate.base.protocol import ( - AnyResponseFormat, - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.reasoning import ReasoningParser -from vllm.sampling_params import StructuredOutputsParams -from vllm.tokenizers import TokenizerLike if TYPE_CHECKING: - from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - -REPLACEMENT_CHAR = "\ufffd" - - -class CohereTagRegistry(NamedTuple): - """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``).""" - - trigger: str - end: str - - -class CohereTagStyle(NamedTuple): - """The structural tags style for a given model architecture. - - ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses - both response and text delimiters). ``tools`` is the tool-call wrapper. - """ - - json_tags: tuple[CohereTagRegistry, ...] - tools: CohereTagRegistry - - -class CohereNormalizedTool(TypedDict): - """A tool definition normalized to the shape ``collect_tool_schema`` expects. - - ``parameters`` is a JSON Schema object (possibly empty) describing the tool's - call signature. - """ - - name: str - parameters: dict[str, Any] - - -COMMAND_A_TOOLS_TAG = CohereTagRegistry( - trigger="<|START_ACTION|>", - end="<|END_ACTION|>", -) -COMMAND_A_JSON_TAG = CohereTagRegistry( - trigger="<|START_RESPONSE|>", - end="<|END_RESPONSE|>", -) -COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry( - trigger="<|START_TEXT|>", - end="<|END_TEXT|>", -) - -MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = { - "Cohere2ForCausalLM": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG,), - tools=COMMAND_A_TOOLS_TAG, - ), - "Cohere2VisionForConditionalGeneration": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), - tools=COMMAND_A_TOOLS_TAG, - ), - "Cohere2MoeForCausalLM": CohereTagStyle( - json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), - tools=COMMAND_A_TOOLS_TAG, - ), -} - - -def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: - """Build an xgrammar EBNF grammar that matches a JSON array of tool calls. - - The grammar shape is architecture-independent; callers are responsible for - wrapping it in the correct structural tag (see ``CohereTagStyle.tools``). - """ - tool_dictionary: dict[str, str] = {} - for tool in tool_schema: - tool_name = tool["name"] - tool_parameters = json.dumps(tool["parameters"]) - json_schema = f"""{{ - "type": "object", - "properties": {{ - "tool_call_id": {{ - "type": "string", - "pattern": "^[0-9]+$" - }}, - "tool_name": {{ - "type": "string", - "const": "{tool_name}" - }}, - "parameters": {tool_parameters} - }} - }}""" - tool_grammar = str(xgr.Grammar.from_json_schema(json_schema)) - for match in re.findall(r"\b(\w+)\s*::=", tool_grammar): - tool_grammar = re.sub( - rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar - ) - tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}" - # Emitted grammar shape: - # root ::= tools - # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws - # ws ::= (" " | "\t" | "\n")* - # tool ::= | | ... (one alternative per input) - # ::= root (per-tool xgrammar rules) - # root ::= ... (from xgr.Grammar.from_json_schema) - tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys()) - tool_rules = "\n ".join(tool_dictionary.values()) - grammar = f"""root ::= tools - tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws - ws ::= (" " | "\\t" | "\\n")* - {tool_alternatives} - {tool_rules} - """ - return grammar - - -def _tool_definitions_to_schema_list( - tools: str | list[Any], -) -> list[CohereNormalizedTool]: - """ - Build the list of ``CohereNormalizedTool`` dicts expected by - ``collect_tool_schema``. - - Accepts: - - JSON string - - list of dicts with top-level ``name`` / ``parameters`` - - list of Chat Completions-style ``{"type": "function", "function": {...}}`` - - list of Pydantic models with ``model_dump()`` - """ - if isinstance(tools, str): - try: - parsed = json.loads(tools) - except json.JSONDecodeError: - return [] - if not isinstance(parsed, list): - return [] - else: - parsed = list(tools) - - out: list[CohereNormalizedTool] = [] - for raw in parsed: - t = raw.model_dump() if hasattr(raw, "model_dump") else raw - if not isinstance(t, dict): - continue - # Unwrap Chat Completions' ``{"type": "function", "function": {...}}`` - # shape; otherwise take the dict as-is. - if t.get("type") == "function" and isinstance(t.get("function"), dict): - t = t["function"] - name = t.get("name") - if not isinstance(name, str): - continue - params = t.get("parameters") - out.append( - CohereNormalizedTool( - name=name, - parameters=params if isinstance(params, dict) else {}, - ) - ) - return out - - -def _has_effective_tools( - tools: str | list[Any] | None, -) -> TypeGuard[str | list[Any]]: - """ - True when ``tools`` contains at least one tool definition to convert. - - ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest`` - uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are - treated as effective only when non-blank. - """ - if tools is None: - return False - if isinstance(tools, str): - return bool(tools.strip()) - return len(tools) > 0 - - -# Builder: produces vLLM response_format in xgrammar's canonical format. -# See xgrammar docs: type "structural_tag" with "format" = triggered_tags -# and tag content type = json_schema | grammar. -def convert_schema_to_structural_tags( - schema: dict | None = None, - tools: str | list[Any] | None = None, - model_architecture: str | None = None, -) -> str | None: - """ - Returns a response_format string accepted by xgrammar's structural tag format. - Uses the canonical shape: {"type": "structural_tag", "format": {...}} with - format.type "triggered_tags" and tag content type "json_schema" or "grammar". - - Callers that are not on an engine path (e.g. the reasoning parser) must pass - ``model_architecture`` explicitly. - """ - if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE: - return None - style = MODEL_TO_TAG_STYLE[model_architecture] - - tags: list[dict] = [] - triggers: list[str] = [] - - def _add_tag(tag: CohereTagRegistry, content: dict) -> None: - tags.append({"begin": tag.trigger, "content": content, "end": tag.end}) - triggers.append(tag.trigger) - - if schema is not None: - # One structural tag per JSON wrapper (e.g. MOE: response + text). - # Same for schema-only and "tools plus JSON mode" (North: schema when - # the model does not call tools). - for jt in style.json_tags: - _add_tag(jt, {"type": "json_schema", "json_schema": schema}) - - if _has_effective_tools(tools): - # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS) - # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts). - tool_schema_list = _tool_definitions_to_schema_list(tools) - if not tool_schema_list: - raise ValueError( - "No valid tool definitions could be parsed from the request for " - "structural tag conversion." - ) - tool_grammar = collect_tool_schema(tool_schema_list) - _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar}) + from vllm.tokenizers import TokenizerLike - if not tags: - return None - return json.dumps( - { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "triggers": triggers, - "tags": tags, - }, - } - ) - -def _response_format_type( - response_format: AnyResponseFormat | dict | None, -) -> str | None: - if response_format is None: - return None - if isinstance(response_format, dict): - t = response_format.get("type") - return t if isinstance(t, str) else None - return response_format.type - - -def _maybe_parse_json_dict(value: Any) -> dict | None: - """If value is a JSON string, parse to dict; otherwise require dict.""" - if isinstance(value, dict): - return value - if isinstance(value, str): - try: - parsed = json.loads(value) - except (TypeError, json.JSONDecodeError): - return None - return parsed if isinstance(parsed, dict) else None - return None - - -def _unwrap_nested_schema(candidate: Any) -> dict | None: - """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present. - - Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict. - """ - cand = _maybe_parse_json_dict(candidate) - if not isinstance(cand, dict): - return None - nested = cand.get("schema") - return nested if isinstance(nested, dict) else cand - - -def _schema_from_json_schema_field(js_wr: Any) -> dict | None: - """ - Extract the JSON Schema object from Chat Completions ``json_schema`` payload. - - Accepts: - - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field - - dict in OpenAI shape ``{"name": ..., "schema": {...}}`` - - dict with ``json_schema`` key holding either the schema or a nested wrapper - - dict that is already a JSON Schema document (some clients omit the wrapper) - - JSON strings for any of the above - """ - if js_wr is None: - return None - - parsed_wr = _maybe_parse_json_dict(js_wr) - if parsed_wr is not None: - js_wr = parsed_wr - - if hasattr(js_wr, "model_dump"): - for by_alias in (True, False): - try: - data = js_wr.model_dump(by_alias=by_alias, exclude_none=False) - except TypeError: - data = js_wr.model_dump(by_alias=by_alias) - out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema")) - if out is not None: - return out - inner_attr = getattr(js_wr, "json_schema", None) - return inner_attr if isinstance(inner_attr, dict) else None - - if isinstance(js_wr, dict): - for key in ("schema", "json_schema"): - out = _unwrap_nested_schema(js_wr.get(key)) - if out is not None: - return out - return js_wr - - return None - - -def _schema_dict_from_chat_response_format( - rf: AnyResponseFormat | dict | None, -) -> dict | None: - """JSON schema dict from Chat Completions ``request.response_format`` only.""" - if rf is None: - return None - rf_type = _response_format_type(rf) - if rf_type == "json_object": - return {"type": "object"} - if rf_type != "json_schema": - return None - js_wr = ( - rf.get("json_schema") - if isinstance(rf, dict) - else getattr(rf, "json_schema", None) - ) - return _schema_from_json_schema_field(js_wr) - - -def _schema_dict_from_structured_outputs( - so: StructuredOutputsParams | None, -) -> dict | None: - """Schema dict from ``structured_outputs`` (``json`` / ``json_object``). - - Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or - ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other - types raise ``ValueError`` only if a caller bypasses that validation. - """ - if so is None: - return None - if so.json_object: - return {"type": "object"} - raw: Any = so.json - if raw is None: - return None - - if hasattr(raw, "model_dump"): - out = _schema_from_json_schema_field(raw) - if out is None: - raise ValueError( - "structured_outputs.json model has no extractable JSON Schema." - ) - return out - - if isinstance(raw, str): - if not raw.strip(): - raise ValueError("structured_outputs.json cannot be empty.") - try: - raw = json.loads(raw) - except json.JSONDecodeError as e: - raise ValueError("structured_outputs.json must be valid JSON.") from e - if not isinstance(raw, dict): - raise ValueError("structured_outputs.json must decode to a JSON object.") - - if isinstance(raw, Mapping): - body = raw if isinstance(raw, dict) else dict(raw) - return _schema_from_json_schema_field(body) or body - - raise ValueError( - f"structured_outputs.json has unsupported type {type(raw).__name__}." - ) - - -def _melody_sources_to_vllm( - raw_sources: Any, - position_to_source: Mapping[tuple[int, int], CitationSource] | None, -) -> list[CitationSource]: - """Convert melody's ``Source`` objects into resolved :class:`CitationSource`. - - Melody's ``Source`` shape is - ``{tool_call_index, tool_result_indices, document_ids}``. Without - :meth:`PyFilterOptions.with_message_history` the ``document_ids`` - list is empty, so we resolve the numeric address ourselves using - ``position_to_source``, which is - :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source` - applied to the inbound request and forwarded through - ``chat_template_kwargs`` at parser-construction time. Each melody - source with ``tool_result_indices=[i, j, ...]`` fans out to one - :class:`CitationSource` per resolved position, populating the - wire-shape ``type`` / ``id`` / ``document`` / ``tool_output`` - fields; unresolved positions are silently skipped (the calling - citation is dropped at wire-coercion time if every source in it - ends up unresolved -- see ``_to_wire_citation`` in - :mod:`vllm.entrypoints.cohere.serving`). - - When ``position_to_source`` is ``None`` (parser wired outside of - :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion - request that happens to hit a cohere reasoning parser) sources - can't be resolved and get dropped; citations then never reach any - citation-aware handler downstream. - """ - if not position_to_source: - return [] - out: list[CitationSource] = [] - for s in raw_sources or []: - bucket = getattr(s, "tool_call_index", None) - indices = list(getattr(s, "tool_result_indices", None) or []) - if bucket is None or not indices: - continue - for idx in indices: - info = position_to_source.get((bucket, idx)) - if info is not None: - out.append(info) - return out - - -def _melody_citations_to_vllm( - raw_citations: Any, - position_to_source: Mapping[tuple[int, int], CitationSource] | None, -) -> list[Citation] | None: - """Convert melody's ``FilterCitation`` objects into :class:`Citation`. - - Resolves each source in-place via ``position_to_source`` (see - :func:`_melody_sources_to_vllm`). Citations whose sources all fail - to resolve are still emitted here with ``sources=[]``; the serving - layer drops them at wire-coercion time so the fail-closed policy - lives in exactly one place. +class BaseCohereCommandReasoningParser(ReasoningParser): + """Expose is_reasoning_end for the structured-output engine. + Everything else runs in + :class:`vllm.parser.cohere_command.CohereCommandParser`. """ - if not raw_citations: - return None - out: list[Citation] = [] - for c in raw_citations: - out.append( - Citation( - start=getattr(c, "start_index", None), - end=getattr(c, "end_index", None), - text=getattr(c, "text", None), - sources=_melody_sources_to_vllm( - getattr(c, "sources", None), position_to_source - ), - type=( - "THINKING_CONTENT" - if getattr(c, "is_thinking", False) - else "TEXT_CONTENT" - ), - ) - ) - return out + melody_preset: ClassVar[str] -class BaseCohereCommandReasoningParser(ReasoningParser): - def __init__( - self, - tokenizer: TokenizerLike, - *args, - streaming_opts: PyFilterOptions, - unary_opts: PyFilterOptions, - **kwargs, - ): + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__(tokenizer, *args, **kwargs) self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") - self.unary_opts = unary_opts - self.melody_unary = PyFilter(unary_opts) - self.melody_streaming = PyFilter(streaming_opts) - # Citations extracted by the most recent ``extract_reasoning`` call. - # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its - # ``_finalize_response_message`` hook) read this back from the - # parser instance (which is constructed per-request) and attach - # the result to :class:`CohereChatMessage.citations`. ``None`` when - # the last parse produced no citations. - self.last_unary_citations: list[Citation] | None = None - # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved - # CitationSource`` map, forwarded from - # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via - # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm` - # for how it's consumed. Absent for parser instances that - # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI - # chat completion request routed through a cohere reasoning - # parser); citations from those requests are dropped since we - # have no way to attribute their sources. - ctk = kwargs.get("chat_template_kwargs") or {} - # Lazy import: ``renderers.cohere`` transitively imports - # ``cohere_melody`` and other heavy deps; hoisting the constant - # to module scope would tie parser importability to that graph. - from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY - - raw_map = ctk.get(POSITION_TO_SOURCE_KEY) - self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = ( - raw_map if isinstance(raw_map, Mapping) else None - ) @property def reasoning_start_str(self) -> str | None: @@ -537,80 +34,6 @@ def reasoning_start_str(self) -> str | None: def reasoning_end_str(self) -> str | None: return "<|END_THINKING|>" - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - r = self.melody_streaming.write_decoded(delta_text) - citations = _melody_citations_to_vllm( - getattr(r, "citations", None), self._position_to_source - ) - if ( - r.content is None - and r.reasoning is None - and not r.tool_calls - and not citations - ): - return None - # Always emit CohereDeltaMessage so citations reach the wire via - # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope. - # When ``citations`` is unset, ``CohereDeltaMessage._serialize`` - # drops the field so the emitted shape matches plain - # :class:`DeltaMessage` - msg = CohereDeltaMessage() - if r.content is not None: - msg.content = r.content - if r.reasoning is not None: - msg.reasoning = r.reasoning - if r.tool_calls: - msg.tool_calls = [ - DeltaToolCall( - id=tc.id, - index=tc.index, - type="function", - function=DeltaFunctionCall(name=tc.name, arguments=tc.arguments), - ) - for tc in r.tool_calls - ] - if citations: - msg.citations = citations - return msg - - def extract_reasoning( - self, model_output: str, request: ChatCompletionRequest | ResponsesRequest - ) -> tuple[str | None, str | None]: - result = self.melody_unary.process_full_text(model_output) - # Cache citations so citation-aware handlers can surface them - # on :class:`CohereChatMessage.citations`. The base - # :meth:`ReasoningParser.extract_reasoning` contract only - # returns ``(reasoning, content)``, so citations are passed - # back via parser-instance state -- safe because the parser is - # constructed per-request. - self.last_unary_citations = _melody_citations_to_vllm( - getattr(result, "citations", None), self._position_to_source - ) - return result.reasoning, result.content - - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - token_buf: list[int] = [] - content_ids: list[int] = [] - content_filter = PyFilter(self.unary_opts) - for t in input_ids: - token_buf.append(t) - s = self.model_tokenizer.decode(token_buf, skip_special_tokens=False) - if s.endswith(REPLACEMENT_CHAR): - continue - r = content_filter.write_decoded(s) - if r.content is not None: - content_ids.extend(token_buf) - token_buf = [] - return content_ids - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: chatbot = self.chatbot_token_id start = self.start_token_id @@ -628,91 +51,10 @@ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: return has_end_token - def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - count = 0 - depth = 0 - for token_id in token_ids: - if token_id == self.start_token_id: - depth += 1 - continue - if token_id == self.end_token_id: - if depth > 0: - depth -= 1 - continue - if depth > 0: - count += 1 - return count - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - so = request.structured_outputs - if so is not None and so.structural_tag: - return request - # Schema: prefer ``response_format`` (OpenAI Chat Completions), then - # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay - # on ``request.tools``. - rf = ( - request.response_format - if isinstance(request, ChatCompletionRequest) - else None - ) - if rf is not None and _response_format_type(rf) == "structural_tag": - return request - model_architecture = ( - self._model_config.architecture if self._model_config is not None else None - ) - tools = request.tools - # ``response_format`` wins if both it and ``structured_outputs`` supply JSON. - schema = _schema_dict_from_chat_response_format(rf) - if schema is None: - schema = _schema_dict_from_structured_outputs(so) - if schema is None and not _has_effective_tools(tools): - return request - if model_architecture is None: - return request - result = convert_schema_to_structural_tags( - schema=schema, - tools=tools, - model_architecture=model_architecture, - ) - if result is None: - # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``. - raise ValueError( - "Failed to build structural_tag guided decoding constraints from " - "this request's JSON schema and/or tools. The configured model " - f"architecture ({model_architecture!r}) does not support Cohere " - "command structural tags, or the schema cannot be expressed in " - "that format." - ) - request.structured_outputs = StructuredOutputsParams(structural_tag=result) - # Folded JSON constraints into ``structural_tag``; drop ``response_format`` - # when it was the source so ``to_sampling_params`` does not also set ``json`` / - # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``). - if isinstance(request, ChatCompletionRequest) and rf is not None: - rf_type = _response_format_type(rf) - if rf_type in ("json_schema", "json_object"): - request.response_format = None - return request - class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): - def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): - super().__init__( - tokenizer, - *args, - streaming_opts=PyFilterOptions().cmd3().no_tools(), - unary_opts=PyFilterOptions().cmd3().no_tools(), - **kwargs, - ) + melody_preset = "cmd3" class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser): - def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): - super().__init__( - tokenizer, - *args, - streaming_opts=PyFilterOptions().cmd4().no_tools(), - unary_opts=PyFilterOptions().cmd4().no_tools(), - **kwargs, - ) + melody_preset = "cmd4" diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index 8e6c4a5ac875..e22ebf1ebf1b 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -244,7 +244,7 @@ def _try_import_melody(): raise ImportError( "The `cohere` tokenizer/renderer mode requires the " "`cohere_melody` package. Install it via " - "`pip install 'cohere-melody>=0.11.1'` or build from " + "`pip install 'cohere-melody>=0.14.0'` or build from " "https://github.com/cohere-ai/melody." ) from e diff --git a/vllm/tool_parsers/cohere_command_tool_parser.py b/vllm/tool_parsers/cohere_command_tool_parser.py index 3739b222e5ff..eca5c5496729 100644 --- a/vllm/tool_parsers/cohere_command_tool_parser.py +++ b/vllm/tool_parsers/cohere_command_tool_parser.py @@ -1,148 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Sequence +from typing import ClassVar -try: - from cohere_melody import PyFilter, PyFilterOptions -except ImportError as e: - raise ImportError( - "The Cohere tool parser requires the `cohere_melody` " - "package, which is not installed. Install it with:\n" - " pip install 'cohere-melody>=0.11.1'" - ) from e - -from vllm.entrypoints.generate.base.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ( - ResponsesRequest, -) -from vllm.tokenizers import TokenizerLike from vllm.tool_parsers import ToolParser -from vllm.tool_parsers.utils import Tool class BaseCohereCommandToolParser(ToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - streaming_opts: PyFilterOptions, - unary_opts: PyFilterOptions, - ): - super().__init__(tokenizer) - self.melody_streaming = PyFilter(streaming_opts) - self.melody_unary = PyFilter(unary_opts) - # Melody can emit the tool-call id before the function name. Keep it - # until the first real name delta so clients receive both together. - self._pending_streaming_tool_call_ids: dict[int, str] = {} - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - request = super().adjust_request(request) - request.skip_special_tokens = False - return request - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - r = self.melody_streaming.write_decoded(delta_text) - if r.content is not None: - return DeltaMessage(content=r.content) - if r.reasoning is not None: - return DeltaMessage(reasoning=r.reasoning) - if r.tool_calls: - tool_calls: list[DeltaToolCall] = [] - for tc in r.tool_calls: - if tc.id: - self._pending_streaming_tool_call_ids[tc.index] = tc.id - name = tc.name or None - arguments = tc.arguments or None - # Empty strings are placeholders in Melody's streaming output; - # omit them from OpenAI-compatible deltas instead of sending - # invalid tool names or empty argument fragments. - if name is None and arguments is None: - continue - - function_kwargs = {} - if name is not None: - function_kwargs["name"] = name - if arguments is not None: - function_kwargs["arguments"] = arguments - tool_call_kwargs = { - "index": tc.index, - "function": DeltaFunctionCall(**function_kwargs), - } - if name is not None: - tool_call_id = tc.id or self._pending_streaming_tool_call_ids.pop( - tc.index, None - ) - if tool_call_id is not None: - tool_call_kwargs["id"] = tool_call_id - tool_call_kwargs["type"] = "function" - tool_calls.append(DeltaToolCall(**tool_call_kwargs)) - if tool_calls: - return DeltaMessage(tool_calls=tool_calls) - return None - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - result = self.melody_unary.process_full_text(model_output) - tool_calls = [ - ToolCall( - id=tc.id, - type="function", - function=FunctionCall(name=tc.name, arguments=tc.arguments), - ) - for tc in result.tool_calls - ] - return ExtractedToolCallInformation( - tools_called=len(tool_calls) > 0, - tool_calls=tool_calls, - content=result.content, - ) + melody_preset: ClassVar[str] class CohereCommand3ToolParser(BaseCohereCommandToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - tools: list[Tool] | None = None, - ): - super().__init__( - tokenizer, - streaming_opts=PyFilterOptions().cmd3().start_in_answer(), - unary_opts=PyFilterOptions().cmd3().start_in_answer(), - ) + melody_preset = "cmd3" class CohereCommand4ToolParser(BaseCohereCommandToolParser): - def __init__( - self, - tokenizer: TokenizerLike, - tools: list[Tool] | None = None, - ): - super().__init__( - tokenizer, - streaming_opts=PyFilterOptions().cmd4().start_in_answer(), - unary_opts=PyFilterOptions().cmd4().start_in_answer(), - ) + melody_preset = "cmd4" From 33fa95a0cf9ed6bf384a0a51667a01a4f0bb02d5 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 12 Sep 2026 02:41:54 -0700 Subject: [PATCH 089/774] [CI] Sample GPU utilization and memory alongside test timelines (#56541) Signed-off-by: khluu Co-authored-by: Codex --- .buildkite/scripts/ci-otel/README.md | 43 ++++++ .buildkite/scripts/ci-otel/ci_gpu.py | 145 ++++++++++++++++++ .buildkite/scripts/ci-otel/ci_otel.py | 28 +++- .buildkite/scripts/ci-otel/ci_otel.sh | 19 +++ .../scripts/ci-otel/tests/test_ci_otel.py | 136 +++++++++++++++- 5 files changed, 366 insertions(+), 5 deletions(-) create mode 100644 .buildkite/scripts/ci-otel/README.md create mode 100644 .buildkite/scripts/ci-otel/ci_gpu.py diff --git a/.buildkite/scripts/ci-otel/README.md b/.buildkite/scripts/ci-otel/README.md new file mode 100644 index 000000000000..3ad552b5dbf0 --- /dev/null +++ b/.buildkite/scripts/ci-otel/README.md @@ -0,0 +1,43 @@ +# CI timing and GPU samples + +The pipeline generator activates these helpers for trusted, traced CI commands. +`ci_otel.sh` records command spans and `ci_pytest.sh` injects the pytest timing +plugin. `ci_gpu.py` runs once per command, outside pytest workers, when +`nvidia-smi` is available. It queries up to 16 container-visible NVIDIA devices +once per second without importing PyTorch or creating a CUDA context. +Set `CI_INFRA_GPU_SAMPLING=0` in the job environment to disable sampling. + +Samples contain device UUID, index, name, utilization percent, and used/total +device memory in bytes. They are `ci.gpu.sample` OTLP events inside +`ci.gpu.samples` spans, parented to the command span and authenticated using the +existing job-scoped Buildkite OIDC identity. Batches upload every 30 seconds; +the final batch is spooled for the shell's normal job-end upload. No host agent +or long-lived dashboard token is needed. The dashboard companion change adds +GPU controls alongside jobs, commands, and pytest rows in the Builds timeline. + +Device queries and periodic uploads each have a two-second timeout. Three +consecutive query failures stop sampling. Failed periodic uploads are dropped; +memory stays bounded, and the dashboard shows the missing intervals as gaps. +Uploads can delay the next poll. The shell stops and joins the sampler before +flushing traces, allows at most three seconds for shutdown, and preserves the +original command exit status. Samples already uploaded survive an interrupted +job; the last unflushed batch may be lost on forced termination. + +These are device readings during each test's time interval, not process-level +kernel attribution. Other processes and overlapping tests can contribute to +the readings. Subsecond tests may have no sample. Missing readings never become +zero utilization. MIG parent-device memory and utilization are omitted because +they cannot be attributed to the job's partition. The initial implementation +covers NVIDIA commands in the existing tracing rollout; CPU jobs and AMD +mirrors do not collect GPU samples. + +Run the helper tests without loading vLLM or requiring a GPU: + +```bash +PATH="$PWD/.venv/bin:$PATH" .venv/bin/python -m pytest \ + .buildkite/scripts/ci-otel/tests/test_ci_otel.py -q +``` + +The suite uses a fake `nvidia-smi` to exercise the shell lifecycle, device-query +failure handling, disable switch, event encoding, and bounded uploads. Real +hardware sampling overhead still needs a canary run before broad deployment. diff --git a/.buildkite/scripts/ci-otel/ci_gpu.py b/.buildkite/scripts/ci-otel/ci_gpu.py new file mode 100644 index 000000000000..304d738b886b --- /dev/null +++ b/.buildkite/scripts/ci-otel/ci_gpu.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Best-effort, container-visible NVIDIA device samples for command timelines.""" + +from __future__ import annotations + +import csv +import io +import math +import os +import secrets +import signal +import subprocess +import sys +import threading +import time + +from ci_otel import Span, export_spans, record_spans + +INTERVAL_SECONDS = 1.0 +BATCH_SECONDS = 30 +MAX_DEVICES = 16 +QUERY_TIMEOUT_SECONDS = 2 +QUERY = "index,uuid,name,utilization.gpu,memory.used,memory.total,mig.mode.current" + + +def _number(value: str, scale: int = 1, maximum: int | None = None) -> int | None: + try: + number = float(value) + if not math.isfinite(number) or number < 0: + return None + if maximum is not None and number > maximum: + return None + return round(number * scale) + except ValueError: + return None + + +def parse_samples(output: str, timestamp_ns: int) -> list[dict]: + events = [] + for row in csv.reader(io.StringIO(output)): + if len(row) != 7 or len(events) >= MAX_DEVICES: + continue + index, uuid, name, util, used, total, mig = (value.strip() for value in row) + if not index.isdigit() or not uuid.startswith("GPU-"): + continue + attributes: dict[str, str | int | bool] = { + "gpu.uuid": uuid, + "gpu.index": int(index), + "gpu.name": name, + } + # Parent-device memory on a MIG GPU is not the job's allocated memory. + if mig.lower() == "enabled": + attributes["gpu.status"] = "unsupported_mig" + else: + for key, value in ( + ("gpu.utilization", _number(util, maximum=100)), + ("gpu.memory.used", _number(used, 1024**2)), + ("gpu.memory.total", _number(total, 1024**2)), + ): + if value is not None: + attributes[key] = value + events.append( + {"time_ns": timestamp_ns, "name": "ci.gpu.sample", "attributes": attributes} + ) + return events + + +def sample_span(events: list[dict]) -> Span: + return Span( + trace_id=os.environ["CI_INFRA_TRACE_ID"], + span_id=secrets.token_hex(8), + parent_span_id=os.environ["CI_INFRA_COMMAND_SPAN_ID"], + name="ci.gpu.samples", + start_ns=events[0]["time_ns"], + end_ns=events[-1]["time_ns"], + attributes={ + "ci.span.kind": "gpu-samples", + "gpu.sample.interval_ms": int(INTERVAL_SECONDS * 1000), + "gpu.scope": "container-visible-devices", + }, + events=events, + ) + + +def collect(parent_pid: int, stop: threading.Event) -> None: + events: list[dict] = [] + batch_started = time.monotonic() + failures = 0 + try: + while not stop.is_set() and os.getppid() == parent_pid: + started = time.monotonic() + try: + result = subprocess.run( + [ + "nvidia-smi", + f"--query-gpu={QUERY}", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=QUERY_TIMEOUT_SECONDS, + check=True, + ) + if stop.is_set(): + break + samples = parse_samples(result.stdout, time.time_ns()) + if not samples: + break + events.extend(samples) + failures = 0 + except (OSError, subprocess.SubprocessError): + failures += 1 + if failures >= 3: + print( + "CI GPU sampling stopped: device query unavailable", + file=sys.stderr, + ) + break + if events and time.monotonic() - batch_started >= BATCH_SECONDS: + # Bounded memory and upload time even if the receiver is down. + # Dropped batches appear as gaps, never as zero utilization. + export_spans([sample_span(events)], timeout_seconds=2) + events = [] + batch_started = time.monotonic() + stop.wait(max(0, INTERVAL_SECONDS - (time.monotonic() - started))) + finally: + # The shell joins this process before the ordinary job-end spool flush. + if events: + record_spans([sample_span(events)]) + + +def main() -> None: + stop = threading.Event() + signal.signal(signal.SIGTERM, lambda *_: stop.set()) + signal.signal(signal.SIGINT, lambda *_: stop.set()) + try: + collect(int(sys.argv[1]), stop) + except Exception: + print("CI GPU sampling skipped", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/.buildkite/scripts/ci-otel/ci_otel.py b/.buildkite/scripts/ci-otel/ci_otel.py index 0d25594bec19..cec3f4130986 100644 --- a/.buildkite/scripts/ci-otel/ci_otel.py +++ b/.buildkite/scripts/ci-otel/ci_otel.py @@ -14,12 +14,13 @@ import time import urllib.request from collections.abc import Iterable -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path ENDPOINT = os.getenv("CI_INFRA_OTEL_ENDPOINT", "https://ci.vllm.ai/api/otel/v1/traces") AUDIENCE = os.getenv("CI_INFRA_OTEL_AUDIENCE", "https://ci.vllm.ai/api/otel") MAX_BATCH_SIZE = 2_000 +MAX_BATCH_BYTES = 1024 * 1024 TEST_SPOOL_BATCH_SIZE = 100 @@ -56,6 +57,7 @@ class Span: end_ns: int attributes: dict[str, str | int | bool] status_code: int = 1 + events: list[dict] = field(default_factory=list) def _varint(value: int) -> bytes: @@ -137,6 +139,12 @@ def _encode_span(span: Span) -> bytes: ) for name, value in attributes.items(): encoded += _bytes_field(9, _attribute(name, value)) + for event in span.events: + encoded_event = _fixed64_field(1, int(event["time_ns"])) + encoded_event += _string_field(2, event["name"]) + for name, value in event["attributes"].items(): + encoded_event += _bytes_field(3, _attribute(name, value)) + encoded += _bytes_field(11, encoded_event) return encoded + _bytes_field(15, _varint_field(3, span.status_code)) @@ -482,10 +490,24 @@ def export_spans(spans: list[Span], timeout_seconds: float = 3.0) -> bool: try: deadline = time.monotonic() + max(timeout_seconds, 0.1) token = _oidc_token(deadline) - for offset in range(0, len(spans), MAX_BATCH_SIZE): + batches: list[list[Span]] = [] + batch: list[Span] = [] + size = 0 + for span in spans: + span_size = len(_encode_span(span)) + if batch and ( + size + span_size > MAX_BATCH_BYTES or len(batch) >= MAX_BATCH_SIZE + ): + batches.append(batch) + batch, size = [], 0 + batch.append(span) + size += span_size + if batch: + batches.append(batch) + for batch in batches: request = urllib.request.Request( ENDPOINT, - data=encode_request(spans[offset : offset + MAX_BATCH_SIZE]), + data=encode_request(batch), method="POST", headers={ "Authorization": f"Bearer {token}", diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index feec16fa6528..78b18207ad27 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -82,6 +82,13 @@ ci_otel_start() { CI_INFRA_COMMAND_SPAN_ID="${_CI_INFRA_OTEL_COMMAND_SPAN_ID}" export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID _CI_INFRA_OTEL_ACTIVE=1 + _CI_INFRA_GPU_PID="" + if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && + [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && + command -v nvidia-smi >/dev/null 2>&1; then + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" /dev/null || : + _CI_INFRA_GPU_WAIT=0 + while kill -0 "${_CI_INFRA_GPU_PID}" 2>/dev/null && + [ "${_CI_INFRA_GPU_WAIT}" -lt 30 ]; do + sleep 0.1 + _CI_INFRA_GPU_WAIT=$((_CI_INFRA_GPU_WAIT + 1)) + done + kill -KILL "${_CI_INFRA_GPU_PID}" 2>/dev/null || : + wait "${_CI_INFRA_GPU_PID}" 2>/dev/null || : + _CI_INFRA_GPU_PID="" + fi CI_INFRA_TRACE_ID="" CI_INFRA_COMMAND_SPAN_ID="" export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ff279de953a5..f636e7ed1826 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -8,6 +8,7 @@ import shlex import subprocess import sys +import threading import time from pathlib import Path @@ -16,6 +17,7 @@ SCRIPTS_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(SCRIPTS_DIR)) +import ci_gpu # noqa: E402 import ci_otel # noqa: E402 from ci_otel import Span # noqa: E402 @@ -251,9 +253,15 @@ def test_export_failure_is_soft_and_bounded(monkeypatch): assert time.monotonic() - started < 0.5 -def test_export_batches_with_one_oidc_token(monkeypatch): +@pytest.mark.parametrize("byte_limited", [False, True]) +def test_export_batches_with_one_oidc_token(monkeypatch, byte_limited): monkeypatch.setenv("BUILDKITE", "true") - monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2) + if byte_limited: + monkeypatch.setattr( + ci_otel, "MAX_BATCH_BYTES", len(ci_otel._encode_span(_span())) * 2 + ) + else: + monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2) token_calls = [] requests = [] @@ -786,3 +794,127 @@ def test_pytest_plugin_failure_cannot_fail_pytest(tmp_path): assert result.returncode == 0, result.stderr assert "1 passed" in result.stdout + + +def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics(): + events = ci_gpu.parse_samples( + "0, GPU-a, H200, 0, 1024, 2048, Disabled\n" + "1, GPU-b, H200, [N/A], [N/A], 2048, Disabled\n" + "2, GPU-c, H200, 50, 1024, 2048, Enabled\n" + "3, GPU-d, H200, NaN, -1, inf, Disabled\n", + 1234, + ) + assert len(events) == 4 + assert events[0]["attributes"]["gpu.utilization"] == 0 + assert events[0]["attributes"]["gpu.memory.used"] == 1024**3 + for event in events[1:]: + assert "gpu.utilization" not in event["attributes"] + assert "gpu.memory.used" not in event["attributes"] + assert events[2]["attributes"]["gpu.status"] == "unsupported_mig" + assert "gpu.memory.total" not in events[2]["attributes"] + + +def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip( + monkeypatch, tmp_path +): + monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) + monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) + monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path)) + span = ci_gpu.sample_span( + ci_gpu.parse_samples("0, GPU-a, H200, 90, 1, 2, Disabled", 1234) + ) + assert span.parent_span_id == "02" * 8 + assert span.trace_id == "01" * 16 + assert ci_otel.record_spans([span]) + assert ci_otel.load_spans() == [span] + assert b"ci.gpu.sample" in ci_otel.encode_request([span]) + + +def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch): + calls = [] + + def query(*args, **kwargs): + calls.append(kwargs["timeout"]) + raise subprocess.TimeoutExpired("nvidia-smi", kwargs["timeout"]) + + monkeypatch.setattr(ci_gpu.subprocess, "run", query) + monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) + monkeypatch.setattr( + ci_gpu, "record_spans", lambda _: pytest.fail("fabricated sample") + ) + ci_gpu.collect(os.getppid(), threading.Event()) + assert calls == [2, 2, 2] + + +def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch): + monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) + monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) + monkeypatch.setattr(ci_gpu, "BATCH_SECONDS", 0) + monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) + monkeypatch.setattr( + ci_gpu.subprocess, + "run", + lambda *a, **k: subprocess.CompletedProcess( + a, 0, "0, GPU-a, H200, 50, 1, 2, Disabled" + ), + ) + stop = threading.Event() + batches = [] + + def upload(spans, timeout_seconds): + batches.append(spans) + assert timeout_seconds == 2 + if len(batches) == 3: + stop.set() + return False + + monkeypatch.setattr(ci_gpu, "export_spans", upload) + ci_gpu.collect(os.getppid(), stop) + assert [len(batch[0].events) for batch in batches] == [1, 1, 1] + + +@pytest.mark.parametrize("sampling", ["1", "0"]) +def test_gpu_sampler_stops_with_command_and_preserves_failure_status( + tmp_path, sampling +): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + smi = bin_dir / "nvidia-smi" + smi.write_text('#!/bin/sh\necho "0, GPU-a, H200, 70, 1024, 2048, Disabled"\n') + smi.chmod(0o755) + # Keep the spool for inspection, and avoid any network upload. + shell = ( + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"\n' + "ci_otel_run 1 workload sh -c 'sleep 1.3; exit 7'\n" + "exit $?" + ) + result = subprocess.run( + ["/bin/sh", "-c", shell], + capture_output=True, + text=True, + timeout=8, + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "BUILDKITE": "false", + "CI_INFRA_GPU_SAMPLING": sampling, + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), + "CI_INFRA_OTEL_RUNTIME_DIR": str(tmp_path / "runtime"), + "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), + }, + ) + assert result.returncode == 7, result.stderr + records = [ + json.loads(line) + for path in (tmp_path / "spans").glob("spans-*.jsonl") + for line in path.read_text().splitlines() + ] + command = next(record for record in records if record["name"] == "ci.command") + batches = [record for record in records if record["name"] == "ci.gpu.samples"] + assert bool(batches) == (sampling == "1") + if batches: + assert batches[0]["parent_span_id"] == command["span_id"] + assert all( + command["start_ns"] <= event["time_ns"] <= command["end_ns"] + for event in batches[0]["events"] + ) From 6b153463a8a6721fd4bb89ec9afa014409bcce36 Mon Sep 17 00:00:00 2001 From: Srijan Keshri Date: Sat, 12 Sep 2026 15:15:32 +0530 Subject: [PATCH 090/774] [Build] Define _USE_MATH_DEFINES for FlashMLA targets (#54007) Signed-off-by: arcusbuilds Co-authored-by: Shengqi Chen --- cmake/external_projects/flashmla.cmake | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 4cf1e421f0a5..382fd0289a2a 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -186,6 +186,15 @@ if(FLASH_MLA_ARCHS) if(VLLM_GPU_LANG STREQUAL "CUDA") target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA) endif() + + # FlashMLA's sources use M_LOG2E from . MSVC's UCRT only defines the + # M_* constants when _USE_MATH_DEFINES is set before is first + # included, so define it for both targets on MSVC. Scoped to MSVC so the + # compile flags (and build caches) on other platforms stay unchanged. + if(MSVC) + target_compile_definitions(_flashmla_C PRIVATE _USE_MATH_DEFINES) + target_compile_definitions(_flashmla_extension_C PRIVATE _USE_MATH_DEFINES) + endif() else() message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}") # Create empty targets for setup.py on unsupported systems From d86257e2833e0c3add2c694a380435c451c1b176 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 12 Sep 2026 03:29:22 -0700 Subject: [PATCH 091/774] [CI/Build] Give the Torch ABI audit time to start (#56596) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/torch_abi.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml index 9e8375587adb..201f9b6093a3 100644 --- a/.buildkite/test_areas/torch_abi.yaml +++ b/.buildkite/test_areas/torch_abi.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: ":nvidia: (L4) Torch Stable ABI Audit" key: torch-stable-abi-audit - timeout_in_minutes: 5 + timeout_in_minutes: 10 source_file_dependencies: - .buildkite/check-torch-abi.py - csrc/ From bde4feb06123d5bdefbf817b840e3489113548c5 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Sat, 12 Sep 2026 20:03:12 +0800 Subject: [PATCH 092/774] [XPU][CI] Remove pip install dependency in test yaml files (#55171) Signed-off-by: Kunshang Ji --- .buildkite/intel_jobs/lm_eval_intel.yaml | 3 +- .buildkite/intel_jobs/misc_intel.yaml | 10 +- .../intel_jobs/model_executor_intel.yaml | 2 - .../intel_jobs/models_multimodal_intel.yaml | 10 +- .buildkite/intel_jobs/test-intel.yaml | 6 +- requirements/test/xpu.in | 22 +- requirements/test/xpu.txt | 198 +++++++++++++++--- 7 files changed, 200 insertions(+), 51 deletions(-) diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 966ced8156d9..11d3edc90e10 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -23,6 +23,5 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install "gpt-oss[eval]==0.0.5" && - cd tests && + 'cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 3818339662ed..39b61b3f7d90 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -69,8 +69,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install lm_eval[api]>=0.4.12 && - export VLLM_WORKER_MULTIPROC_METHOD=spawn && + 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py && pytest -v -s v1/test_oracle.py v1/test_request.py v1/test_outputs.py && @@ -193,8 +192,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install modelscope\<1.38 && - cd tests && + 'cd tests && pytest -v -s test_regression.py' - label: Metrics, Tracing (2 GPUs) @@ -229,8 +227,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 && - cd tests && + 'cd tests && pytest -v -s v1/tracing' - label: Async Engine, Inputs, Utils, Worker @@ -268,7 +265,6 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pip install av && pytest -v -s detokenizer && pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 14a853544eca..d7c683f8cea3 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -25,8 +25,6 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'apt-get update && apt-get install -y curl libsodium23 && - pip3 install tensorizer==2.10.1 && - pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 && export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 44efa8e0f130..58ecb2d28209 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -22,8 +22,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av && - cd tests && + 'cd tests && pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" && pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model' @@ -96,8 +95,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av && - cd tests && + 'cd tests && pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' - label: Multi-Modal Processor # 44min @@ -121,9 +119,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av matplotlib ftfy && - pip install open-clip-torch --no-deps && - cd tests && + 'cd tests && pytest -v -s models/multimodal/processing/test_tensor_schema.py --deselect="tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[deepseek-ai/DeepSeek-V4.1-Flash]" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 00f172ca885d..8f4e83eee8ab 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -34,8 +34,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install tblib==3.1.0 && - python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager && + 'python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN && @@ -129,8 +128,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av && - cd tests && + 'cd tests && pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py && pytest -v -s benchmarks/test_serve_cli.py' - label: "XPU quantization test" diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 6848dd98b7d5..e93b418e9a65 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -1,7 +1,7 @@ -r ../common.txt # --- Test Infrastructure --- -tblib +tblib==3.1.0 pytest pytest_asyncio pytest-timeout @@ -17,9 +17,12 @@ accelerate arctic-inference lm_eval[api]>=0.4.12 modelscope<1.38 +runai-model-streamer[s3,gcs,azure]>=0.15.7 +tensorizer==2.10.1 transformers==5.16.1 # --- Audio Processing --- +av librosa audioread soxr @@ -29,7 +32,7 @@ soundfile # --- Tool Parsing & Evaluation --- blobfile rapidfuzz -gpt-oss +gpt-oss[eval]==0.0.5 schemathesis jiwer bm25s @@ -41,5 +44,18 @@ pqdm # --- Vision & Multimodal --- timm albumentations +ftfy +matplotlib decord==0.6.0 -mistral-common[image,audio] \ No newline at end of file +mistral-common[image,audio] + +open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py + +# --- Observability --- +opentelemetry-api>=1.26.0 +opentelemetry-exporter-otlp>=1.26.0 +opentelemetry-sdk>=1.26.0 +opentelemetry-semantic-conventions-ai>=0.4.1 + +cohere_melody>=0.11.1 # required for cohere command reasoning/tool parser tests +cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 76453f50f232..6c8bc1fcdd7e 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -54,14 +54,34 @@ attrs==26.1.0 # referencing audioread==3.1.0 # via -r requirements/test/xpu.in +av==18.1.0 + # via -r requirements/test/xpu.in +azure-core==1.41.0 + # via + # azure-identity + # azure-storage-blob +azure-identity==1.25.3 + # via runai-model-streamer-azure +azure-storage-blob==12.30.1 + # via runai-model-streamer-azure blake3==1.0.9 # via -r requirements/test/../common.txt blobfile==3.2.0 - # via -r requirements/test/xpu.in + # via + # -r requirements/test/xpu.in + # gpt-oss bm25s==0.3.10 # via # -r requirements/test/xpu.in # mteb +boto3==1.43.87 + # via + # runai-model-streamer-s3 + # tensorizer +botocore==1.43.87 + # via + # boto3 + # s3transfer bounded-pool-executor==0.0.3 # via pqdm cachetools==7.1.7 @@ -94,16 +114,29 @@ click==8.4.2 # uvicorn cloudpickle==3.1.2 # via -r requirements/test/../common.txt +cohere==7.1.1 + # via -r requirements/test/xpu.in +cohere-melody==0.13.4 + # via -r requirements/test/xpu.in colorama==0.4.6 # via sacrebleu compressed-tensors==0.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt +contourpy==1.3.3 + # via matplotlib coverage==7.15.4 # via pytest-cov cryptography==50.0.0 - # via pyjwt + # via + # azure-identity + # azure-storage-blob + # google-auth + # msal + # pyjwt +cycler==0.12.1 + # via matplotlib dataproperty==1.1.1 # via # pytablewriter @@ -174,6 +207,8 @@ fastar==0.11.0 # via # fastapi # fastapi-cloud-cli +fastavro==1.12.2 + # via cohere filelock==3.32.3 # via # -c requirements/common.txt @@ -183,6 +218,8 @@ filelock==3.32.3 # huggingface-hub # modelscope # torch +fonttools==4.64.0 + # via matplotlib frozenlist==1.8.0 # via # aiohttp @@ -193,11 +230,36 @@ fsspec==2026.6.0 # evaluate # huggingface-hub # torch +ftfy==6.3.1 + # via + # -r requirements/test/xpu.in + # open-clip-torch +google-api-core==2.34.0 + # via + # google-cloud-core + # google-cloud-storage +google-auth==2.57.0 + # via + # google-api-core + # google-cloud-core + # google-cloud-storage + # runai-model-streamer-gcs +google-cloud-core==2.7.0 + # via google-cloud-storage +google-cloud-storage==3.13.1 + # via runai-model-streamer-gcs +google-crc32c==1.8.0 + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.10.2 + # via google-cloud-storage googleapis-common-protos==1.75.1 # via + # google-api-core # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.9 +gpt-oss==0.0.5 # via -r requirements/test/xpu.in graphql-core==3.2.11 # via hypothesis-graphql @@ -212,6 +274,8 @@ harfile==0.5.0 # via schemathesis hf-xet==1.6.0 # via huggingface-hub +hiredis==3.4.1 + # via tensorizer html2text==2025.4.15 # via gpt-oss httpcore==1.0.9 @@ -223,6 +287,7 @@ httptools==0.8.0 httpx==0.28.1 # via # anthropic + # cohere # datasets # fastapi # fastapi-cloud-cli @@ -239,10 +304,13 @@ huggingface-hub==1.31.0 # accelerate # datasets # evaluate + # open-clip-torch # sentence-transformers # timm # tokenizers # transformers +humanize==4.16.0 + # via runai-model-streamer hypothesis==6.165.10 # via # hypothesis-graphql @@ -305,10 +373,13 @@ intel-sycl-rt==2026.0.0 # torch interegular==0.3.3 # via lm-format-enforcer +isodate==0.7.2 + # via azure-storage-blob jinja2==3.1.6 # via # -c requirements/xpu.txt # fastapi + # gpt-oss # lm-eval # torch jiter==0.16.0 @@ -318,7 +389,10 @@ jiter==0.16.0 jiwer==4.0.0 # via -r requirements/test/xpu.in jmespath==1.1.0 - # via model-hosting-container-standards + # via + # boto3 + # botocore + # model-hosting-container-standards joblib==1.5.3 # via # librosa @@ -335,16 +409,16 @@ jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema -jupyter-client==8.9.1 - # via gpt-oss -jupyter-core==5.9.1 - # via jupyter-client +kiwisolver==1.5.1 + # via matplotlib lark==1.2.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt lazy-loader==0.5 # via librosa +libnacl==2.1.0 + # via tensorizer librosa==1.0.0 # via -r requirements/test/xpu.in llguidance==1.7.6 @@ -372,6 +446,8 @@ markupsafe==3.0.3 # via # jinja2 # werkzeug +matplotlib==3.11.1 + # via -r requirements/test/xpu.in mbstrdecoder==1.1.5 # via # dataproperty @@ -408,6 +484,12 @@ more-itertools==11.1.0 # via lm-eval mpmath==1.3.0 # via sympy +msal==1.38.0 + # via + # azure-identity + # msal-extensions +msal-extensions==1.3.1 + # via azure-identity msgpack==1.2.1 # via librosa msgspec==0.21.1 @@ -436,18 +518,21 @@ numba==0.65.0 # via # -c requirements/xpu.txt # librosa -numpy==2.3.5 +numpy==2.2.6 # via # -r requirements/test/../common.txt # accelerate # albucore # albumentations # bm25s + # contourpy # datasets # decord # evaluate + # gpt-oss # librosa # lm-eval + # matplotlib # mistral-common # mteb # numba @@ -455,12 +540,14 @@ numpy==2.3.5 # pandas # pytrec-eval-terrier # rouge-score + # runai-model-streamer # sacrebleu # scikit-learn # scipy # sentence-transformers # soundfile # soxr + # tensorizer # torchvision # transformers # xgrammar @@ -487,10 +574,13 @@ onemkl-sycl-rng==2026.0.0 # via torch onemkl-sycl-sparse==2026.0.0 # via torch +open-clip-torch==2.32.0 + # via -r requirements/test/xpu.in openai==3.2.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # gpt-oss openai-harmony==0.0.8 # via # -c requirements/common.txt @@ -507,6 +597,7 @@ opentelemetry-api==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # -r requirements/test/xpu.in # mcp # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http @@ -516,6 +607,7 @@ opentelemetry-exporter-otlp==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # -r requirements/test/xpu.in opentelemetry-exporter-otlp-proto-common==1.44.0 # via # opentelemetry-exporter-otlp-proto-grpc @@ -533,6 +625,7 @@ opentelemetry-sdk==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # -r requirements/test/xpu.in # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # opentelemetry-semantic-conventions-ai @@ -544,6 +637,7 @@ opentelemetry-semantic-conventions-ai==0.5.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # -r requirements/test/xpu.in outlines-core==0.2.14 # via # -c requirements/common.txt @@ -557,6 +651,7 @@ packaging==26.3 # huggingface-hub # lazy-loader # lm-format-enforcer + # matplotlib # modelscope # pooch # pytest @@ -567,6 +662,7 @@ pandas==3.0.5 # via # datasets # evaluate + # gpt-oss partial-json-parser==0.2.1.1.post7 # via -r requirements/test/../common.txt pathvalidate==3.3.1 @@ -574,12 +670,11 @@ pathvalidate==3.3.1 pillow==12.3.0 # via # -r requirements/test/../common.txt + # matplotlib # mistral-common # torchvision platformdirs==4.11.3 - # via - # jupyter-core - # pooch + # via pooch pluggy==1.6.0 # via # pytest @@ -609,20 +704,30 @@ propcache==0.5.2 # via # aiohttp # yarl +proto-plus==1.28.4 + # via google-api-core protobuf==7.35.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # google-api-core # googleapis-common-protos # opentelemetry-proto + # proto-plus + # tensorizer psutil==7.2.2 # via # -r requirements/test/../common.txt # accelerate + # tensorizer py-cpuinfo==9.0.0 # via -r requirements/test/../common.txt pyarrow==25.0.1 # via datasets +pyasn1==0.6.4 + # via pyasn1-modules +pyasn1-modules==0.4.2 + # via google-auth pybase64==1.5.0 # via -r requirements/test/../common.txt pycountry==26.2.16 @@ -637,6 +742,7 @@ pydantic==2.13.4 # -r requirements/test/../common.txt # albumentations # anthropic + # cohere # compressed-tensors # fastapi # fastapi-cloud-cli @@ -653,7 +759,9 @@ pydantic==2.13.4 # pydantic-settings # xgrammar pydantic-core==2.46.4 - # via pydantic + # via + # cohere + # pydantic pydantic-extra-types==2.11.1 # via # fastapi @@ -667,7 +775,11 @@ pygments==2.21.0 # pytest # rich pyjwt==2.13.0 - # via mcp + # via + # mcp + # msal +pyparsing==3.3.2 + # via matplotlib pyrate-limiter==4.4.0 # via schemathesis pystemmer==3.1.0 @@ -700,7 +812,8 @@ pytest-timeout==2.4.0 # via -r requirements/test/xpu.in python-dateutil==2.9.0.post0 # via - # jupyter-client + # botocore + # matplotlib # pandas # typepy python-dotenv==1.2.3 @@ -735,11 +848,12 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # jupyter-client rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in # jiwer +redis==8.1.0 + # via tensorizer referencing==0.37.0 # via # jsonschema @@ -748,6 +862,7 @@ regex==2026.7.19 # via # -r requirements/test/../common.txt # nltk + # open-clip-torch # sacrebleu # tiktoken # transformers @@ -755,13 +870,18 @@ requests==2.34.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # azure-core + # cohere # datasets # docker # evaluate + # google-api-core + # google-cloud-storage # gpt-oss # lm-eval # mistral-common # modelscope + # msal # mteb # opentelemetry-exporter-otlp-proto-http # pooch @@ -786,6 +906,16 @@ rpds-py==2026.6.3 # via # jsonschema # referencing +runai-model-streamer==0.16.1 + # via -r requirements/test/xpu.in +runai-model-streamer-azure==0.16.1 + # via runai-model-streamer +runai-model-streamer-gcs==0.16.1 + # via runai-model-streamer +runai-model-streamer-s3==0.16.1 + # via runai-model-streamer +s3transfer==0.19.2 + # via boto3 sacrebleu==2.6.0 # via lm-eval safetensors==0.8.0 @@ -793,6 +923,7 @@ safetensors==0.8.0 # -c requirements/common.txt # -r requirements/test/../common.txt # accelerate + # open-clip-torch # timm # transformers schemathesis==4.24.3 @@ -887,7 +1018,7 @@ tbb==2023.0.0 # intel-opencl-rt # mkl # torch -tblib==3.2.2 +tblib==3.1.0 # via -r requirements/test/xpu.in tcmlib==1.5.0 # via @@ -900,6 +1031,8 @@ tenacity==9.1.4 # via # gpt-oss # lm-eval +tensorizer==2.10.1 + # via -r requirements/test/xpu.in termcolor==3.3.0 # via gpt-oss threadpoolctl==3.6.0 @@ -912,11 +1045,14 @@ tiktoken==0.14.0 # lm-eval # mistral-common timm==1.0.28 - # via -r requirements/test/xpu.in + # via + # -r requirements/test/xpu.in + # open-clip-torch tokenizers==0.23.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # cohere # sentence-transformers # transformers torch==2.13.0+xpu @@ -925,32 +1061,33 @@ torch==2.13.0+xpu # accelerate # compressed-tensors # mteb + # open-clip-torch + # runai-model-streamer # sentence-transformers + # tensorizer # timm # torchvision # xgrammar torchvision==0.28.0+xpu - # via timm -tornado==6.5.8 - # via jupyter-client + # via + # open-clip-torch + # timm tqdm==4.70.0 # via # -r requirements/test/../common.txt # datasets # evaluate + # gpt-oss # huggingface-hub # lm-eval # modelscope # mteb # nltk + # open-clip-torch # openai # pqdm # sentence-transformers # transformers -traitlets==5.16.1 - # via - # jupyter-client - # jupyter-core transformers==5.16.1 # via # -c requirements/common.txt @@ -982,6 +1119,8 @@ typer==0.27.1 # fastapi-cli # fastapi-cloud-cli # transformers +types-requests==2.33.0.20260712 + # via cohere typing-extensions==4.16.0 # via # -c requirements/common.txt @@ -991,12 +1130,15 @@ typing-extensions==4.16.0 # anthropic # anyio # apache-tvm-ffi + # azure-core + # azure-identity + # azure-storage-blob # chz + # cohere # fastapi # grpcio # httpx2 # huggingface-hub - # jupyter-client # lm-eval # mcp # mcp-types @@ -1035,10 +1177,12 @@ umf==1.1.0 urllib3==2.7.0 # via # blobfile + # botocore # docker # modelscope # requests # sentry-sdk + # types-requests uvicorn==0.52.3 # via # fastapi @@ -1052,6 +1196,8 @@ watchfiles==1.2.0 # via # -r requirements/test/../common.txt # uvicorn +wcwidth==0.8.3 + # via ftfy websockets==17.0.1 # via uvicorn werkzeug==3.1.8 From a0914ab7d07b9a788db15b499edf3b5a280e7a14 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sat, 12 Sep 2026 08:07:10 -0500 Subject: [PATCH 093/774] [Nano-Nemotron] Fix Nano-Nemotron precomputed multimodal embeddings (#56398) Signed-off-by: Taneem Ibrahim --- tests/models/multimodal/test_nano_nemotron_vl.py | 12 ++++++++++++ vllm/model_executor/models/nano_nemotron_vl.py | 15 +++++++++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 0aef83a8271b..f4ca19ed969c 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -5,6 +5,7 @@ import numpy as np import pytest +import torch from vllm import envs from vllm.model_executor.models.nano_nemotron_vl import ( @@ -14,6 +15,17 @@ from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems +@pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"]) +def test_precomputed_multimodal_embeddings(input_key: str): + model = object.__new__(NemotronH_Nano_VL_V2) + embeds = torch.randn(2, 4, 8) + + outputs = model.embed_multimodal(**{input_key: embeds}) + + assert len(outputs) == len(embeds) + assert all(torch.equal(output, embed) for output, embed in zip(outputs, embeds)) + + class _TextOnlyMultiModalConfig: def get_limit_per_prompt(self, modality: str) -> int: return 0 diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index ac6e254cb2df..2edf56871dbf 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1103,7 +1103,8 @@ def extract_feature( def _parse_and_validate_image_input( self, **kwargs: object ) -> NanoNemotronVLImageInputs | None: - if image_embeds := kwargs.pop("image_embeds", None): + image_embeds = kwargs.pop("image_embeds", None) + if image_embeds is not None: return NanoNemotronVLImageEmbeddingInputs( type="image_embeds", data=image_embeds, @@ -1341,7 +1342,7 @@ def _create_final_video_embeddings( def _parse_and_validate_video_input( self, **kwargs: object - ) -> NanoNemotronVLVideoPixelInputs | None: + ) -> NanoNemotronVLVideoInputs | None: pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None) video_num_patches = kwargs.pop("video_num_patches", None) video_embeds = kwargs.pop("video_embeds", None) @@ -1408,7 +1409,10 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: and "images" not in modalities ): modalities["images"] = self._parse_and_validate_image_input(**kwargs) - if input_key in ("pixel_values_flat_video",) and "videos" not in modalities: + if ( + input_key in ("pixel_values_flat_video", "video_embeds") + and "videos" not in modalities + ): modalities["videos"] = self._parse_and_validate_video_input(**kwargs) if ( input_key @@ -1450,7 +1454,10 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: multimodal_embeddings += tuple(image_embeddings) if modality == "videos": video_input = modalities["videos"] - video_embeddings = self._process_video_input(video_input) + if video_input["type"] == "video_embeds": + video_embeddings = video_input["data"] + else: + video_embeddings = self._process_video_input(video_input) multimodal_embeddings += tuple(video_embeddings) if modality == "audios": audio_input = modalities["audios"] From 1b29c508bcb25fbfd9ab4c3d61ae131d7e0700d5 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sat, 12 Sep 2026 08:07:20 -0500 Subject: [PATCH 094/774] [Bugfix] Initialize data parser in Nano-Nemotron audio test (#56401) Signed-off-by: Taneem Ibrahim --- tests/models/multimodal/test_nano_nemotron_vl.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index f4ca19ed969c..75900cde53a3 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -12,7 +12,11 @@ NanoNemotronVLMultiModalProcessor, NemotronH_Nano_VL_V2, ) -from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems +from vllm.multimodal.parse import ( + MultiModalDataItems, + MultiModalDataParser, + VideoProcessorItems, +) @pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"]) @@ -166,6 +170,9 @@ def test_extract_audio_from_videos_passes_max_duration(): mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64) processor = object.__new__(NanoNemotronVLMultiModalProcessor) + processor.data_parser = MultiModalDataParser( + target_sr=dummy_audio[1], target_channels=1 + ) target = "vllm.model_executor.models.nano_nemotron_vl.load_audio_pyav" with patch(target, return_value=dummy_audio) as mock_load: From 986e2f870abc8f87a1830af7d335961b3e4234b2 Mon Sep 17 00:00:00 2001 From: JartX Date: Sat, 12 Sep 2026 15:20:26 +0200 Subject: [PATCH 095/774] =?UTF-8?q?[Refactor][ROCm]=20Migrate=20the=20RDNA?= =?UTF-8?q?3=20W4A16=20MoE=20to=20the=20oracle/experts=20pa=E2=80=A6=20(#5?= =?UTF-8?q?5522)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: JartX --- .../quantization/test_rdna3_compile_guards.py | 127 +++++---- tests/quantization/test_moe_wna16.py | 34 +++ vllm/config/kernel.py | 2 + .../layers/fused_moe/experts/rdna3_moe.py | 235 ++++++++++++++++ .../layers/fused_moe/oracle/int_wna16.py | 110 ++++++++ .../layers/fused_moe/routed_experts.py | 1 - .../compressed_tensors_moe.py | 7 - .../compressed_tensors_moe_wna16.py | 6 +- .../compressed_tensors_moe_wna16_rdna3.py | 260 ------------------ .../compressed_tensors_moe/rocm_moe_rdna.py | 48 ---- 10 files changed, 453 insertions(+), 377 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py delete mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py delete mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py index 3c1e226d2c07..194441f0f909 100644 --- a/tests/kernels/quantization/test_rdna3_compile_guards.py +++ b/tests/kernels/quantization/test_rdna3_compile_guards.py @@ -34,6 +34,13 @@ if not current_platform.is_rocm(): pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, + kInt8Static, +) from vllm.platforms.rocm import on_gfx1100 # noqa: E402 gfx1100_only = pytest.mark.skipif( @@ -166,15 +173,13 @@ def test_op_absent_on_non_gfx1100(op_name): @not_gfx1100 def test_rocm_moe_not_supported_on_non_gfx1100(): - """rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware.""" - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 - rocm_moe_rdna, + """The RDNA3 MoE experts must not be selectable on non-gfx1100 hardware.""" + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, ) - wq = type("WQ", (), {"num_bits": 4})() - assert rocm_moe_rdna.is_supported(wq) is False, ( - "rocm_moe_rdna.is_supported() returned True on non-gfx1100 — " - "dispatch guard is broken" + assert Rdna3WNA16Experts._supports_current_device() is False, ( + "Rdna3WNA16Experts reported support on non-gfx1100 — dispatch guard is broken" ) @@ -359,54 +364,58 @@ def test_no_toplevel_rocm_c_import(self): # ============================================================================ -class _FakeWeightQuant: - """Minimal stand-in for a weight quantization config.""" - - def __init__(self, num_bits): - self.num_bits = num_bits - - class TestMoEDispatchMocked: """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable.""" - def test_is_supported_false_when_mocked_cdna(self): - """rocm_moe_rdna.is_supported() must return False when not on gfx1100.""" - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 - rocm_moe_rdna, + def test_kernel_unavailable_when_mocked_cdna(self): + """The device gate must reject when not on gfx1100.""" + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + rdna3_moe_kernel_available, ) with patch("vllm.platforms.rocm.on_gfx1100", return_value=False): - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False + assert rdna3_moe_kernel_available() is False - @pytest.mark.parametrize("num_bits", [2, 3, 8, 16]) - def test_is_supported_rejects_non_w4(self, num_bits): - """is_supported() rejects non-4-bit even before checking arch.""" - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 - rocm_moe_rdna, + @pytest.mark.parametrize( + "weight_key", + [kInt8Static, kInt4StaticAsym, kInt4Static32Asym, None], + ) + def test_quant_scheme_rejects_non_symmetric_int4(self, weight_key): + """Only symmetric int4 weight-only schemes reach the kernel.""" + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, ) - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False + assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is False - def test_is_supported_false_when_op_missing(self): - """is_supported() returns False when the C++ op doesn't exist.""" - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 - rocm_moe_rdna, + @pytest.mark.parametrize("weight_key", [kInt4Static, kInt4Static32]) + def test_quant_scheme_accepts_symmetric_int4(self, weight_key): + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, + ) + + assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is True + + def test_kernel_unavailable_when_op_missing(self): + """The device gate returns False when the C++ op doesn't exist.""" + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + rdna3_moe_kernel_available, ) fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})() with patch.object(torch, "ops", create=True) as mock_ops: mock_ops._rocm_C = fake_rocm_c - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False + assert rdna3_moe_kernel_available() is False - def test_is_supported_false_when_rocm_c_absent(self): - """is_supported() returns False when _rocm_C doesn't exist at all.""" - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 - rocm_moe_rdna, + def test_kernel_unavailable_when_rocm_c_absent(self): + """The device gate returns False when _rocm_C doesn't exist at all.""" + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + rdna3_moe_kernel_available, ) fake_ops = type("FakeOps", (), {})() with patch.object(torch, "ops", fake_ops): - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False + assert rdna3_moe_kernel_available() is False class TestDenseKernelSelectionMocked: @@ -471,30 +480,30 @@ def test_chooser_skips_rdna3_when_mocked_cdna(self): ) -class TestCompressedTensorsMoEDispatchGuard: - """Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm().""" +class TestWNA16OracleWiring: + """The RDNA3 backend must reach the kernel only through the oracle.""" - def test_rocm_guard_in_dispatch_source(self): - """The rocm_moe_rdna import and call must be inside an is_rocm() check.""" - src = _read_pkg_source_or_skip( - "model_executor", - "layers", - "quantization", - "compressed_tensors", - "compressed_tensors_moe", - "compressed_tensors_moe.py", + def test_backend_maps_to_rdna3_experts(self): + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, + ) + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + backend_to_kernel_cls, + map_wna16_backend, ) - lines = src.splitlines() - for i, line in enumerate(lines, 1): - stripped = line.strip() - if "rocm_moe" in stripped and not stripped.startswith("#"): - found_guard = False - for j in range(i - 1, max(0, i - 15), -1): - if "is_rocm()" in lines[j - 1]: - found_guard = True - break - assert found_guard, ( - f"L{i}: rocm_moe_rdna reference not protected by " - f"is_rocm() guard: {stripped}" - ) + assert backend_to_kernel_cls(WNA16MoEBackend.RDNA3) == [Rdna3WNA16Experts] + assert map_wna16_backend("rdna3") == WNA16MoEBackend.RDNA3 + + def test_backend_is_offered_before_the_triton_fallback(self): + """Priority order: the native kernel outranks Triton when supported.""" + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + _get_priority_backends, + ) + + backends = _get_priority_backends() + assert backends.index(WNA16MoEBackend.RDNA3) < backends.index( + WNA16MoEBackend.TRITON + ) diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 475ca84f6a17..0fd5a10120b3 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -93,6 +93,40 @@ def test_moe_wna16_accepts_channelwise_gptq_activation_order(): False, "MoeWNA16 checkpoint layout", ), + ( + WNA16MoEBackend.RDNA3, + AutoGPTQConfig(4, 128, False, True, False, {}, {}), + False, + False, + "compressed-tensors", + ), + ( + WNA16MoEBackend.RDNA3, + QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=False, + dynamic=False, + group_size=128, + ), + True, + False, + "asymmetric", + ), + ( + WNA16MoEBackend.RDNA3, + QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.CHANNEL, + symmetric=True, + dynamic=False, + ), + False, + False, + "group-wise scales", + ), ], ) def test_wna16_oracle_rejects_incompatible_quant_structures( diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index a6dddc317875..c16e4fc169c6 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -141,6 +141,7 @@ def with_default( "flydsl", "hpc", "emulation", + "rdna3", ] # Backends that run the mega-MoE model path through the flashinfer moe_ep @@ -262,6 +263,7 @@ class KernelConfig: - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16 (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250) - "flydsl": Use AMD FlyDSL kernels (ROCm only) + - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only) - "hpc": Use HPC kernels (FP8 and Hopper only) - "emulation": use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py new file mode 100644 index 000000000000..dbbce6b0b182 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MoE W4A16 experts on the RDNA3 (gfx1100) HIP kernel. + +``moe_gptq_gemm_rdna3`` is a single HIP kernel launch per GEMM that handles +expert routing + W4A16 dequant + dot product with atomic output accumulation. + +Weight format (per expert, same as the dense RDNA3 W4A16 kernel): + - Packed int32 ``[E, K/8, N]`` with exllama shuffle + - Scales ``[E, groups, N]`` in activation dtype + - Zero points ``[E, groups, N/8]`` packed int32 (synthesized, symmetric only) +""" + +import torch + +import vllm._custom_ops as ops +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation_supported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static, + kInt4Static32, +) +from vllm.platforms import current_platform + + +def rdna3_moe_kernel_available() -> bool: + """Whether the fused RDNA3 MoE HIP kernel is built into this binary.""" + if not current_platform.is_rocm(): + return False + + from vllm.platforms.rocm import on_gfx1100 + + return ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3") + ) + + +class Rdna3WNA16Experts(mk.FusedMoEExpertsModular): + """W4A16 experts backed by ``moe_gptq_gemm_rdna3`` (gfx1100). + + Both GEMMs accumulate atomically, so their destinations are zeroed first. + The second GEMM is given ``output_topk``, which makes it reduce over the + top-k dimension while it accumulates — the ``moe_sum`` launch and the + ``[M * top_k, K]`` intermediate are fused away, hence + ``TopKWeightAndReduceNoOP``. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + assert quant_config.use_int4_w4a16, "Supports only int4_w4a16" + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + self._empty_topk_weights: torch.Tensor | None = None + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + @staticmethod + def _supports_current_device() -> bool: + return rdna3_moe_kernel_available() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # Symmetric int4 weights only: the kernel consumes synthesized zero + # points and has no path for checkpoint zero points. + return activation_key is None and weight_key in (kInt4Static, kInt4Static32) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return apply_moe_activation_supported(activation) + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # Weights are packed K-first: w1 is [E, K // 8, w13_shards * N] and + # w2 is [E, N // 8, K], so neither N nor K can be read off a trailing + # dimension the way the base implementation does. + assert w1.dim() == 3 and w2.dim() == 3 + assert a1.dim() == 2 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + return w1.size(0), a1.size(0), w2.size(1) * 8, a1.size(-1), topk_ids.size(1) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # The modular kernel provisions the output buffer out of workspace13, + # so both GEMM-1 buffers have to live in workspace2 — the second GEMM + # reads the activation output while accumulating into the output. + gate_up = N if not activation.is_gated else 2 * N + return ((M, K), (M * topk * (gate_up + N),), (M, K)) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + assert self.w1_scale is not None and self.w2_scale is not None + assert self.w1_zp is not None and self.w2_zp is not None + + E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids) + if global_num_experts == -1: + global_num_experts = E + + gate_up = w1.size(2) + act_n = self.adjust_N_for_activation(gate_up, activation) + rows = M * top_k + + scratch = workspace2.view(-1) + gate_up_out = scratch[: rows * gate_up].view(rows, gate_up) + act_out = scratch[rows * gate_up : rows * (gate_up + act_n)].view(rows, act_n) + + # BLOCK_SIZE_M=1 for decode (no padding waste), 4 for prefill. + block_size_m = 1 if M <= 4 else 4 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + block_size_m, + global_num_experts, + expert_map, + ) + + if ( + self._empty_topk_weights is None + or self._empty_topk_weights.device != hidden_states.device + ): + self._empty_topk_weights = torch.empty( + 0, dtype=torch.float32, device=hidden_states.device + ) + topk_weights_f32 = topk_weights.reshape(-1).float() + no_topk_weights = self._empty_topk_weights + + gate_up_out.zero_() + ops.moe_gptq_gemm_rdna3( + hidden_states, + gate_up_out, + w1, + self.w1_scale, + self.w1_zp, + topk_weights_f32 if apply_router_weight_on_input else no_topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + apply_router_weight_on_input, + ) + + self.activation(activation, act_out, gate_up_out) + + # output_topk=top_k makes the kernel accumulate into out[token_id], + # fusing the top-k reduction into the atomic write-back. + output.zero_() + ops.moe_gptq_gemm_rdna3( + act_out, + output, + w2, + self.w2_scale, + self.w2_zp, + no_topk_weights if apply_router_weight_on_input else topk_weights_f32, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + 1, + block_size_m, + not apply_router_weight_on_input, + output_topk=top_k, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 899f9bc3d94a..d47ad9c97e31 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -7,6 +7,7 @@ import torch from compressed_tensors.quantization import ( QuantizationArgs, + QuantizationStrategy, ) import vllm._custom_ops as ops @@ -41,8 +42,10 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + pack_quantized_values_into_int32, ) from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types logger = init_logger(__name__) @@ -56,6 +59,7 @@ class WNA16MoEBackend(Enum): TRITON = "TRITON" XPU = "XPU" EMULATION = "EMULATION" + RDNA3 = "RDNA3" def backend_to_kernel_cls( @@ -100,6 +104,12 @@ def backend_to_kernel_cls( ) return [Int4EmulationTritonExperts] + elif backend == WNA16MoEBackend.RDNA3: + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, + ) + + return [Rdna3WNA16Experts] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -114,6 +124,8 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: return [WNA16MoEBackend.XPU] return [ + # Native HIP kernel, gated on gfx1100 by _supports_current_device(). + WNA16MoEBackend.RDNA3, WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, @@ -134,6 +146,18 @@ def _backend_incompatibility_reason( if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias): return "zero points and bias are not supported" + if backend == WNA16MoEBackend.RDNA3: + if not isinstance(quant_config, QuantizationArgs): + return "only compressed-tensors checkpoints are supported" + if may_have_zp: + return "asymmetric checkpoints are not supported" + if may_have_bias: + return "expert bias is not supported" + if quant_config.actorder == "group": + return "group activation ordering is not supported" + if quant_config.strategy != QuantizationStrategy.GROUP: + return "only group-wise scales are supported" + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config @@ -190,6 +214,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, "emulation": WNA16MoEBackend.EMULATION, + "rdna3": WNA16MoEBackend.RDNA3, } if backend := mapping.get(runner_backend): return backend @@ -372,6 +397,9 @@ def make_wna16_moe_kernel( from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( Int4EmulationTritonExperts, ) + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import ( + Rdna3WNA16Experts, + ) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, ) @@ -387,6 +415,7 @@ def make_wna16_moe_kernel( XPUExpertsWNA16, CPUExpertsInt4, Int4EmulationTritonExperts, + Rdna3WNA16Experts, ) if backend == WNA16MoEBackend.HUMMING: allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING)) @@ -811,6 +840,78 @@ def _process_awq_weights_marlin( ) +def _synthesize_rdna3_qzeros( + groups: int, out_features: int, device: torch.device +) -> torch.Tensor: + """Create the packed zero-point tensor for symmetric quantization. + + GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode + (bias - 1) = 7 for uint4b8 (bias=8). + """ + zeros = torch.full( + (groups, out_features), + scalar_types.uint4b8.bias - 1, + dtype=torch.int32, + device=device, + ) + return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1) + + +def _process_weights_rdna3( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + group_size: int, +) -> tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor | None, # w13_qzeros + torch.Tensor | None, # w2_qzeros + torch.Tensor | None, # w13_input_global_scale + torch.Tensor | None, # w2_input_global_scale + torch.Tensor | None, # w13_bias + torch.Tensor | None, # w2_bias +]: + """RDNA3 (gfx1100) W4A16 weight post-processing. + + Interleaves the packed nibbles per expert (the exllama shuffle the dense + RDNA3 kernel also uses) and synthesizes the symmetric zero points that + ``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout + ``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the + kernel wants, so neither is repacked. + """ + device = w13.device + num_experts = w13.size(0) + + for e in range(num_experts): + w13_e = w13[e].contiguous() + ops.gptq_shuffle(w13_e, 4) + w13[e] = w13_e + w2_e = w2[e].contiguous() + ops.gptq_shuffle(w2_e, 4) + w2[e] = w2_e + + def _qzeros(w: torch.Tensor) -> torch.Tensor: + qz = _synthesize_rdna3_qzeros((w.size(1) * 8) // group_size, w.size(2), device) + return qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous() + + return ( + w13, + w2, + w13_scale.contiguous(), + w2_scale.contiguous(), + _qzeros(w13), + _qzeros(w2), + None, # w13_input_global_scale + None, # w2_input_global_scale + None, # w13_bias + None, # w2_bias + ) + + def _process_weights_cpu( quant_config: QuantizationConfig | QuantizationArgs | None, w13: torch.Tensor, @@ -1432,6 +1533,15 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) + elif backend == WNA16MoEBackend.RDNA3: + assert isinstance(quant_config, QuantizationArgs) + return _process_weights_rdna3( + w13, + w2, + w13_scale, + w2_scale, + quant_config.group_size, + ) elif backend == WNA16MoEBackend.CPU: return _process_weights_cpu( quant_config, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 5d284bcc80f3..2dafb6c493d6 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -617,7 +617,6 @@ def weight_loader( # against known CompressionFormat enum values that have this quality if quant_method_name in ( "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsWNA16RDNA3MoEMethod", "CompressedTensorsW4A16FlydslMoEMethod", ): if is_transposed: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index bcc930471da7..f2d2a9c9452c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -88,14 +88,7 @@ def get_moe_method( f"and bits: {weight_quant.num_bits}" ) - # Native ROCm HIP kernels (RDNA3, etc.) if current_platform.is_rocm(): - from . import rocm_moe_rdna - - if rocm_moe_rdna.is_supported(weight_quant): - return rocm_moe_rdna.make_method( - weight_quant, input_quant, layer.moe_config - ) from vllm.platforms.rocm import on_gfx950 vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 6a3f73e5392d..7a5b2d949b9f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -443,10 +443,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) - # CPU fused_experts_cpu requires zero points even for symmetric quant. + # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points + # even for symmetric quant (the oracle synthesizes them). # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None. if ( - not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU + not self.symmetric + or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3) ) and self.wna16_backend != WNA16MoEBackend.EMULATION: assert w13_qzeros is not None and w2_qzeros is not None replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py deleted file mode 100644 index 05db6ec3ddd8..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ /dev/null @@ -1,260 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel. - -Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that -handles expert routing + W4A16 dequant + dot product with atomic output. - -Weight format (per expert, same as dense RDNA3 W4A16): - - Packed int32 ``[E, K/8, N]`` with exllama shuffle - - Scales ``[E, groups, N]`` in activation dtype - - Zero points ``[E, groups, N/8]`` packed int32 (synthesized) -""" - -import torch - -from vllm import _custom_ops as ops -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import ( - RoutedExperts, - SharedExperts, -) -from vllm.model_executor.layers.fused_moe.activation import ( - MoEActivation, - apply_moe_activation, -) -from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( - moe_align_block_size, -) -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 - CompressedTensorsWNA16MoEMethod, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - pack_quantized_values_into_int32, -) -from vllm.scalar_type import scalar_types - -logger = init_logger(__name__) - - -def _synthesize_qzeros( - groups: int, out_features: int, device: torch.device -) -> torch.Tensor: - """Create packed zero-point tensor for symmetric quant. - - GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode - (bias - 1) = 7 for uint4b8 (bias=8). - """ - zeros = torch.full( - (groups, out_features), - scalar_types.uint4b8.bias - 1, - dtype=torch.int32, - device=device, - ) - return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1) - - -class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod): - """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3). - - Weights are in RDNA3 format (shuffled int32 [E, K/8, N]), - NOT Triton format (transposed uint8). apply() dispatches through - the fused HIP kernel directly. - """ - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - device = layer.w13_weight_packed.device - num_experts = layer.w13_weight_packed.shape[0] - - # Shuffle weights in-place per expert (exllama nibble interleave) - for e in range(num_experts): - w13_e = layer.w13_weight_packed.data[e].contiguous() - ops.gptq_shuffle(w13_e, 4) - layer.w13_weight_packed.data[e] = w13_e - w2_e = layer.w2_weight_packed.data[e].contiguous() - ops.gptq_shuffle(w2_e, 4) - layer.w2_weight_packed.data[e] = w2_e - - # Keep scales as [E, groups, N] in activation dtype - act_dtype = layer.w13_weight_scale.dtype - layer.w13_weight_scale = torch.nn.Parameter( - layer.w13_weight_scale.to(dtype=act_dtype).contiguous(), - requires_grad=False, - ) - layer.w2_weight_scale = torch.nn.Parameter( - layer.w2_weight_scale.to(dtype=act_dtype).contiguous(), - requires_grad=False, - ) - - # Synthesize packed zero points: [E, groups, N/8] int32 - w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size - w13_N = layer.w13_weight_packed.shape[2] - w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size - w2_N = layer.w2_weight_packed.shape[2] - - w13_qz = _synthesize_qzeros(w13_groups, w13_N, device) - w2_qz = _synthesize_qzeros(w2_groups, w2_N, device) - layer.w13_qzeros = torch.nn.Parameter( - w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(), - requires_grad=False, - ) - layer.w2_qzeros = torch.nn.Parameter( - w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(), - requires_grad=False, - ) - - # Pre-allocate reusable buffers for decode (sizes based on top_k=8) - N_gate_up = w13_N - hidden_size = w2_N - intermediate = N_gate_up // 2 # gated activation - # Max tokens we expect in decode; prefill will re-allocate if needed - max_decode_tokens = 16 - top_k = 8 # conservative default - buf_size = max_decode_tokens * top_k - layer.rdna3_w1_buf = torch.zeros( - buf_size, N_gate_up, dtype=act_dtype, device=device - ) - layer.rdna3_act_buf = torch.empty( - buf_size, intermediate, dtype=act_dtype, device=device - ) - layer.rdna3_out_buf = torch.zeros( - max_decode_tokens, hidden_size, dtype=act_dtype, device=device - ) - layer.rdna3_empty_tw = torch.empty(0, device=device) - - def apply( - self, - layer: RoutedExperts, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts: SharedExperts | None, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - activation = ( - layer.activation - if isinstance(layer.activation, MoEActivation) - else MoEActivation.from_str(layer.activation) - ) - return _rdna3_fused_moe( - x, - topk_weights, - topk_ids, - layer=layer, - activation=activation, - apply_router_weight_on_input=(layer.apply_router_weight_on_input), - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - ) - - -def _rdna3_fused_moe( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - layer: RoutedExperts, - activation: MoEActivation, - apply_router_weight_on_input: bool, - global_num_experts: int, - expert_map: torch.Tensor | None, -) -> torch.Tensor: - """Fused MoE forward using the RDNA3 W4A16 HIP kernel. - - Optimizations vs naive dispatch: - - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path) - - Pre-allocated buffers (no torch.zeros per call) - - Inline token sorting for small M (skip moe_align_block_size) - - moe_sum fused into output accumulation - """ - num_tokens = hidden_states.shape[0] - top_k = topk_ids.shape[1] - total_tokens = num_tokens * top_k - N_gate_up = layer.w13_weight_packed.shape[2] - hidden_size = layer.w2_weight_packed.shape[2] - dtype = hidden_states.dtype - device = hidden_states.device - - intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up - - if global_num_experts <= 0: - global_num_experts = layer.w13_weight_packed.shape[0] - - # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill - block_size_m = 1 if num_tokens <= 4 else 4 - - # --- Token routing --- - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, - block_size_m, - global_num_experts, - expert_map, - ) - - # --- Reuse pre-allocated buffers when possible --- - if total_tokens <= layer.rdna3_w1_buf.shape[0]: - w1_out = layer.rdna3_w1_buf[:total_tokens] - w1_out.zero_() - act_out = layer.rdna3_act_buf[:total_tokens] - else: - w1_out = torch.zeros( - total_tokens, - N_gate_up, - dtype=dtype, - device=device, - ) - act_out = torch.empty( - total_tokens, - intermediate_size, - dtype=dtype, - device=device, - ) - - # --- topk weights (pre-cast to float32 for kernel) --- - topk_w_float = topk_weights.view(-1).float() - empty_tw = layer.rdna3_empty_tw - - # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] --- - ops.moe_gptq_gemm_rdna3( - hidden_states, - w1_out, - layer.w13_weight_packed, - layer.w13_weight_scale, - layer.w13_qzeros, - topk_w_float if apply_router_weight_on_input else empty_tw, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - top_k, - block_size_m, - apply_router_weight_on_input, - ) - - # --- Activation (silu_and_mul etc.) --- - apply_moe_activation(activation, act_out, w1_out) - - # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) --- - # output_topk=top_k: kernel writes to out[token_id / top_k] directly, - # fusing moe_sum into the atomic accumulation — saves one kernel launch - # and the w2_out intermediate buffer. - out = torch.zeros( - num_tokens, - hidden_size, - dtype=dtype, - device=device, - ) - ops.moe_gptq_gemm_rdna3( - act_out, - out, - layer.w2_weight_packed, - layer.w2_weight_scale, - layer.w2_qzeros, - topk_w_float if not apply_router_weight_on_input else empty_tw, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - 1, - block_size_m, - not apply_router_weight_on_input, - output_topk=top_k, - ) - return out diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py deleted file mode 100644 index e72caa0796a5..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm MoE kernel dispatcher. - -Selects architecture-specific native HIP MoE kernels in priority order. -Falls back to the Triton WNA16 path when no native kernel is available. -""" - -import torch - -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -def is_supported(weight_quant) -> bool: - """Check if a native ROCm MoE kernel is available for this config.""" - if weight_quant.num_bits != 4: - return False - - from vllm.platforms.rocm import on_gfx1100 - - # RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc. - return ( - on_gfx1100() - and hasattr(torch.ops, "_rocm_C") - and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3") - ) - - -def make_method(weight_quant, input_quant, moe_config): - """Create the native ROCm MoE method. Call only after is_supported().""" - from vllm.platforms.rocm import on_gfx1100 - - if on_gfx1100(): - from .compressed_tensors_moe_wna16_rdna3 import ( - CompressedTensorsWNA16RDNA3MoEMethod, - ) - - logger.info_once( - "Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)" - ) - return CompressedTensorsWNA16RDNA3MoEMethod( - weight_quant, input_quant, moe_config - ) - - # Future: RDNA4, CDNA, etc. - raise RuntimeError("is_supported() returned True but no kernel matched") From 2f59050eda23bccf530dbe924ce4877b25a031cc Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 12 Sep 2026 21:26:55 +0800 Subject: [PATCH 096/774] [XPU][CI] Skip ROCm test on non-ROCm platforms (#56555) Signed-off-by: zhenwei-intel --- tests/test_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_config.py b/tests/test_config.py index 7b6027313af0..282f2351221b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -371,6 +371,7 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): default_breakable_cudagraph_architectures.cache_clear() +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") @pytest.mark.parametrize( ("architecture", "use_v2", "mode", "expected"), [ From 13e221f8308416564a446909c2fd4805c78fbb8a Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 12 Sep 2026 07:20:04 -0700 Subject: [PATCH 097/774] [Perf] Fuse DSV4.1 input metadata preparation with Triton (#56562) --- .../test_indexer_deepseek_v4_slot_mapping.py | 105 +++++++++++++++ vllm/v1/attention/backend.py | 24 ++-- vllm/v1/attention/backends/mla/indexer.py | 121 ++++++++++-------- vllm/v1/attention/ops/metadata.py | 87 +++++++++++++ 4 files changed, 274 insertions(+), 63 deletions(-) create mode 100644 vllm/v1/attention/ops/metadata.py diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 800b0557dead..aabe88eeee5c 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -32,6 +32,111 @@ from vllm.v1.worker.utils import select_common_block_size +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("query_lens", [[1], [6] * 16, [0, 2, 6, 0, 4], [0, 0, 0]]) +@pytest.mark.parametrize("padding", [0, 17]) +def test_fused_indexer_decode_metadata(query_lens, padding): + """Flatten device query boundaries and clear graph padding on every replay.""" + from vllm.v1.attention.ops.metadata import _indexer_decode_metadata_kernel + + device = "cuda" + lengths = torch.tensor(query_lens, device=device, dtype=torch.int32) + qsl = torch.cat( + [torch.zeros(1, device=device, dtype=torch.int32), lengths.cumsum(0).int()] + ) + reqs = len(query_lens) + n = sum(query_lens) + tokens = n + padding + capacity = max(tokens, reqs) + 31 + seq = lengths + 4096 + # Exercise block-table row padding and masked columns. + bt = torch.arange(reqs * 74, device=device, dtype=torch.int32).view(reqs, 74)[ + :, :65 + ] + out_bt = torch.full((capacity, 69), -99, device=device, dtype=torch.int32) + outputs = [ + torch.full((capacity,), -99, device=device, dtype=torch.int32) for _ in range(4) + ] + out_seq, out_lens, indices, per_req = outputs + grid = max(reqs, tokens + (capacity - tokens + 255) // 256) + _indexer_decode_metadata_kernel[(grid,)]( + qsl, + seq, + bt, + out_seq, + out_bt, + out_lens, + indices, + per_req, + reqs, + tokens, + capacity, + bt.stride(0), + out_bt.stride(0), + BLOCK_COLS=bt.shape[1], + num_warps=4, + ) + expected_req = torch.repeat_interleave(torch.arange(reqs, device=device), lengths) + expected_seq = torch.cat( + [torch.arange(4097, 4097 + q, device=device) for q in query_lens] + ) + torch.testing.assert_close(out_seq[:n], expected_seq.int(), rtol=0, atol=0) + assert torch.count_nonzero(out_seq[n:]) == 0 + torch.testing.assert_close(out_bt[:n, :65], bt[expected_req], rtol=0, atol=0) + assert torch.count_nonzero(out_bt[n:tokens, :65]) == 0 + torch.testing.assert_close(indices[:n], expected_req.int(), rtol=0, atol=0) + torch.testing.assert_close( + indices[n:tokens], + torch.arange(reqs, reqs + padding, device=device, dtype=torch.int32), + ) + assert torch.all(out_lens[:tokens] == 1) + torch.testing.assert_close(per_req[:reqs], lengths) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0]]) +def test_device_token_request_mapping(query_lens): + """Graph replay follows device boundaries even when CPU lengths are stale.""" + lengths = torch.tensor(query_lens, device="cuda", dtype=torch.int32) + qsl = torch.cat( + [torch.zeros(1, device="cuda", dtype=torch.int32), lengths.cumsum(0).int()] + ) + n = sum(query_lens) + output = torch.full((n + 7,), -99, device="cuda", dtype=torch.int32) + common = CommonAttentionMetadata( + query_start_loc=qsl, + query_start_loc_cpu=qsl.cpu(), + seq_lens=lengths, + num_reqs=len(query_lens), + num_actual_tokens=output.numel(), + max_query_len=max(query_lens), + max_seq_len=max(query_lens), + block_table_tensor=torch.empty( + (len(query_lens), 1), device="cuda", dtype=torch.int32 + ), + slot_mapping=torch.full((n + 7,), -1, device="cuda", dtype=torch.int64), + ) + result = common.token_to_req_indices(output) + assert result.data_ptr() == output.data_ptr() + expected = torch.repeat_interleave( + torch.arange(len(query_lens), device="cuda", dtype=torch.int32), lengths + ) + torch.testing.assert_close(output[:n], expected) + assert torch.count_nonzero(output[n:]) == 0 + graph = torch.cuda.CUDAGraph() + common._token_to_req_indices_cache = None + with torch.cuda.graph(graph): + common.token_to_req_indices(output) + reversed_lens = lengths.flip(0) + qsl[1:].copy_(reversed_lens.cumsum(0)) + graph.replay() + expected = torch.repeat_interleave( + torch.arange(len(query_lens), device="cuda", dtype=torch.int32), reversed_lens + ) + torch.testing.assert_close(output[:n], expected) + assert torch.count_nonzero(output[n:]) == 0 + + def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla(): """Packed MLA/indexer groups must retain 64 compressed rows per page.""" kernel_block_size = select_common_block_size( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index f4031a7ee01e..addd68ec6ae0 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -484,19 +484,21 @@ def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: # Built from the device query_start_loc: adaptive verification decides the # per-request draft split on device, so the CPU copy carries the right total # but not the right per-request boundaries. Padding requests have a query - # length of zero and drop out of the repeat. + # length of zero and are skipped by the device boundary search. num_mapped_tokens = int(self.query_start_loc_cpu[-1]) - query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1] - assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens) - token_to_req_indices = torch.repeat_interleave( - torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device), - query_lens, - output_size=num_mapped_tokens, + from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel + + num_output_tokens = max(num_mapped_tokens, num_tokens) + assert buffer.shape[0] >= num_output_tokens + _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)]( + self.query_start_loc, + buffer, + self.query_start_loc.shape[0] - 1, + num_mapped_tokens, + num_output_tokens, + num_warps=4, ) - buffer[:num_mapped_tokens].copy_(token_to_req_indices) - if num_mapped_tokens < num_tokens: - buffer[num_mapped_tokens:num_tokens].zero_() - self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)] + self._token_to_req_indices_cache = buffer[:num_output_tokens] return self._token_to_req_indices_cache[:num_tokens] # TODO(lucas): remove once we have FULL-CG spec-decode support diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 59fb109daccb..7bc0b442c34b 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -960,28 +960,6 @@ def _prepare_global_decode_seq_lens( self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0 return self.global_decode_seq_lens_buffer[:num_decode_tokens] - def _build_varlen_decode_indices( - self, - decode_lens: torch.Tensor, - decode_lens_cpu: torch.Tensor, - num_decode_tokens: int, - ) -> torch.Tensor: - """Build request ids for flattened SM100 varlen rows.""" - indices = self.decode_indices_buffer[:num_decode_tokens] - actual_expanded = int(decode_lens_cpu.sum().item()) - num_decodes = decode_lens.shape[0] - indices[:actual_expanded] = torch.repeat_interleave( - self.arange_buffer[:num_decodes], - decode_lens, - output_size=actual_expanded, - ) - if actual_expanded < num_decode_tokens: - pad = num_decode_tokens - actual_expanded - indices[actual_expanded:num_decode_tokens] = ( - num_decodes + self.arange_buffer[:pad] - ) - return indices - @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1156,12 +1134,15 @@ def build( decode_metadata = None if num_decodes > 0: - torch.diff( - common_attn_metadata.query_start_loc[: num_decodes + 1], - out=self.decode_lens_buffer[:num_decodes], - ) + if not self.supports_varlen: + torch.diff( + common_attn_metadata.query_start_loc[: num_decodes + 1], + out=self.decode_lens_buffer[:num_decodes], + ) + self.per_req_decode_lens_buffer[:num_decodes].copy_( + self.decode_lens_buffer[:num_decodes] + ) decode_lens = self.decode_lens_buffer[:num_decodes] - self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens) decode_lens_cpu = torch.diff( common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] ) @@ -1196,38 +1177,74 @@ def build( and step_next_n_ok ) - global_seq_lens_for_decode = self._prepare_global_decode_seq_lens( - global_seq_lens=global_seq_lens_for_decode, - decode_lens=decode_lens, - decode_lens_cpu=decode_lens_cpu, - query_start_loc=common_attn_metadata.query_start_loc[:num_decodes], - num_decode_tokens=num_decode_tokens, - use_native=use_native, - max_decode_len=max_decode_len, - ) - - decode_indices = None - if self.supports_varlen: - decode_indices = self._build_varlen_decode_indices( - decode_lens=decode_lens, - decode_lens_cpu=decode_lens_cpu, - num_decode_tokens=num_decode_tokens, - ) - - seq_lens, block_table, decode_lens, batch_size, requires_padding = ( - self._prepare_decode_tensors( - seq_lens=seq_lens, - block_table=block_table, + if not self.supports_varlen: + global_seq_lens_for_decode = self._prepare_global_decode_seq_lens( + global_seq_lens=global_seq_lens_for_decode, decode_lens=decode_lens, decode_lens_cpu=decode_lens_cpu, query_start_loc=common_attn_metadata.query_start_loc[:num_decodes], - num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, use_native=use_native, - next_n=next_n, max_decode_len=max_decode_len, ) - ) + + decode_indices = None + if self.supports_varlen: + from vllm.v1.attention.ops.metadata import ( + _indexer_decode_metadata_kernel, + ) + + capacity = self.decode_seq_lens_buffer.numel() + grid = max( + num_decodes, + num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256), + ) + _indexer_decode_metadata_kernel[(grid,)]( + query_start_loc, + seq_lens, + block_table, + self.decode_seq_lens_buffer, + self.expanded_block_table_buffer, + self.decode_lens_buffer, + self.decode_indices_buffer, + self.per_req_decode_lens_buffer, + num_decodes, + num_decode_tokens, + capacity, + block_table.stride(0), + self.expanded_block_table_buffer.stride(0), + BLOCK_COLS=block_table.shape[1], + num_warps=4, + ) + seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] + block_table = self.expanded_block_table_buffer[:num_decode_tokens] + decode_lens = self.decode_lens_buffer[:num_decode_tokens] + decode_indices = self.decode_indices_buffer[:num_decode_tokens] + requires_padding = False + if global_seq_lens_for_decode is not None and max_decode_len > 1: + self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_( + seq_lens + ) + global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[ + :num_decode_tokens + ] + else: + seq_lens, block_table, decode_lens, batch_size, requires_padding = ( + self._prepare_decode_tensors( + seq_lens=seq_lens, + block_table=block_table, + decode_lens=decode_lens, + decode_lens_cpu=decode_lens_cpu, + query_start_loc=common_attn_metadata.query_start_loc[ + :num_decodes + ], + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + use_native=use_native, + next_n=next_n, + max_decode_len=max_decode_len, + ) + ) if self.compress_ratio > 1: kernel_block_size = self.kernel_block_size diff --git a/vllm/v1/attention/ops/metadata.py b/vllm/v1/attention/ops/metadata.py new file mode 100644 index 000000000000..f5baa51a9297 --- /dev/null +++ b/vllm/v1/attention/ops/metadata.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device-side request mapping and sparse indexer metadata.""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _token_request(query_start_loc, token, num_reqs): + lo = tl.full(token.shape, 0, tl.int32) + hi = tl.full(token.shape, num_reqs, tl.int32) + while tl.sum((lo < hi).to(tl.int32), 0) > 0: + mid = (lo + hi) // 2 + end = tl.load(query_start_loc + mid + 1, mid < num_reqs, other=0x7FFFFFFF) + right = (lo < hi) & (end <= token) + lo = tl.where(right, mid + 1, lo) + hi = tl.where((lo < hi) & ~right, mid, hi) + return tl.minimum(lo, num_reqs - 1) + + +@triton.jit(do_not_specialize=["num_reqs", "num_mapped", "num_tokens"]) +def _token_request_mapping_kernel( + query_start_loc, output, num_reqs, num_mapped, num_tokens +): + token = tl.program_id(0) * 256 + tl.arange(0, 256) + req = _token_request(query_start_loc, token, num_reqs) + tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens) + + +@triton.jit( + do_not_specialize=[ + "num_reqs", + "num_tokens", + "capacity", + "input_stride", + "output_stride", + ] +) +def _indexer_decode_metadata_kernel( + query_start_loc, + seq_lens, + block_table, + output_seq_lens, + output_block_table, + decode_lens, + indices, + per_req_lens, + num_reqs, + num_tokens, + capacity, + input_stride, + output_stride, + BLOCK_COLS: tl.constexpr, +): + token = tl.program_id(0) + cols = tl.arange(0, triton.next_power_of_2(BLOCK_COLS)) + actual_tokens = tl.load(query_start_loc + num_reqs) + if token < num_tokens: + req = _token_request(query_start_loc, tl.full((1,), token, tl.int32), num_reqs) + req = tl.sum(req, 0) + valid = token < actual_tokens + start = tl.load(query_start_loc + req) + end = tl.load(query_start_loc + req + 1) + seq = tl.load(seq_lens + req) + length = tl.where(valid, seq - (end - start) + token - start + 1, 0) + tl.store(output_seq_lens + token, length) + tl.store(decode_lens + token, 1) + tl.store( + indices + token, tl.where(valid, req, num_reqs + token - actual_tokens) + ) + blocks = tl.load( + block_table + req.to(tl.int64) * input_stride + cols, + valid & (cols < BLOCK_COLS), + other=0, + ) + tl.store( + output_block_table + token.to(tl.int64) * output_stride + cols, + blocks, + cols < BLOCK_COLS, + ) + else: + tail = num_tokens + (token - num_tokens) * 256 + tl.arange(0, 256) + tl.store(output_seq_lens + tail, 0, tail < capacity) + if token < num_reqs: + start = tl.load(query_start_loc + token) + end = tl.load(query_start_loc + token + 1) + tl.store(per_req_lens + token, end - start) From 658c8131c73194f812cd80d493e0409475b94235 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:26:39 +0300 Subject: [PATCH 098/774] [Elastic EP] Reuse CUDA graphs across reconfiguration (#54985) --- requirements/kv_connectors.txt | 2 +- tests/distributed/test_elastic_ep.py | 98 +++------ tests/distributed/test_eplb_utils.py | 15 +- vllm/config/parallel.py | 18 +- .../device_communicators/all2all.py | 64 ++++-- .../base_device_communicator.py | 12 ++ .../distributed/elastic_ep/elastic_execute.py | 193 +++++++----------- vllm/distributed/elastic_ep/elastic_state.py | 2 +- vllm/distributed/eplb/eplb_state.py | 78 +++++-- vllm/engine/arg_utils.py | 6 + vllm/envs.py | 5 - .../layers/fused_moe/all2all_utils.py | 33 ++- .../model_executor/layers/fused_moe/config.py | 1 + .../layers/fused_moe/eep_reconfigure.py | 4 +- vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/modular_kernel.py | 7 - .../fused_moe/prepare_finalize/nixl_ep.py | 17 +- vllm/model_executor/layers/fused_moe/utils.py | 3 +- .../warmup/flashinfer_autotune_cache.py | 39 ++++ vllm/v1/engine/core_client.py | 6 + vllm/v1/worker/dp_utils.py | 32 ++- vllm/v1/worker/gpu/dp_utils.py | 6 +- 22 files changed, 372 insertions(+), 270 deletions(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index 1760ed651734..9d9d1278818e 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,7 +2,7 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. cupy-cuda13x != 14.1.0 -nixl == 1.3.2 +nixl == 1.4.1 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. mooncake-transfer-engine >= 0.3.12 diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 01c254c2d039..8ccea8628389 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -27,7 +27,7 @@ def cleanup_ray_between_tests(): NUM_GSM8K_QUESTIONS = 256 EXPECTED_ACCURACY = 0.58 ACCURACY_TOL = 0.08 -MAX_NUM_SEQS = 32 +DEFAULT_ALL2ALL_BACKEND = "allgather_reducescatter" def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: @@ -161,7 +161,12 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: return accuracy -def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]: +def _base_serve_args( + dp_size: int = 2, + max_dp_size: int = 4, + enforce_eager: bool = False, + all2all_backend: str = DEFAULT_ALL2ALL_BACKEND, +) -> list[str]: args = [ "--trust-remote-code", "--tensor-parallel-size", @@ -170,12 +175,12 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str] "0.8", "--max-model-len", "4096", - "--max-num-seqs", - str(MAX_NUM_SEQS), "--enable-expert-parallel", "--all2all-backend", - "allgather_reducescatter", + all2all_backend, "--enable-elastic-ep", + "--elastic-ep-max-dp-size", + str(max_dp_size), "--enable-eplb", "--eplb-config.num_redundant_experts", "0", @@ -205,28 +210,43 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str] @pytest.mark.parametrize( - ("enforce_eager", "traffic_mode"), + ("enforce_eager", "traffic_mode", "all2all_backend"), [ - pytest.param(True, "none", id="enforce_eager_none"), - pytest.param(True, "light", id="enforce_eager_light"), - pytest.param(True, "heavy", id="enforce_eager_heavy"), - pytest.param(False, "heavy", id="cuda_graphs_heavy"), + pytest.param(True, "none", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_none"), + pytest.param(True, "light", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_light"), + pytest.param(True, "heavy", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_heavy"), + pytest.param( + False, + "heavy", + DEFAULT_ALL2ALL_BACKEND, + id="cuda_graphs_heavy", + ), + pytest.param(False, "heavy", "nixl_ep", id="cuda_graphs_heavy_nixl_ep"), ], ) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): +def test_elastic_ep_scaling( + enforce_eager: bool, traffic_mode: str, all2all_backend: str +): from vllm.distributed.eplb.eplb_communicator import has_nixl + from vllm.utils.import_utils import has_nixl_ep if not has_nixl(): pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if all2all_backend == "nixl_ep" and not has_nixl_ep(): + pytest.skip("NIXL EP is not installed") initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) assert target_dp_size > initial_dp_size - vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager) + vllm_serve_args = _base_serve_args( + initial_dp_size, target_dp_size, enforce_eager, all2all_backend + ) with RemoteOpenAIServer( - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 + MODEL_NAME, + vllm_serve_args, + max_wait_seconds=1200, ) as server: initial_accuracy = _run_gsm8k_eval(server, "Initial") @@ -255,55 +275,3 @@ def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})" ) print(f" Tolerance: {ACCURACY_TOL:.3f}") - - -@multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling_uneven(): - """Test scale up with uneven worker distribution. - - This tests the case where num_new_workers % old_dp_size != 0, - specifically 2 -> 3 where remainder = 1 % 2 = 1. - This exercises the remainder handling in sender-receiver pairing. - """ - from vllm.distributed.eplb.eplb_communicator import has_nixl - - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - - vllm_serve_args = _base_serve_args() - - with RemoteOpenAIServer( - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 - ) as server: - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)") - - # Scale 2 -> 3: This has remainder = 1 % 2 = 1 - # Tests uneven sender-receiver pairing - assert _send_scale_command(server, 3) - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)") - - assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( - f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" - ) - - # Scale back down to 2 - assert _send_scale_command(server, 2) - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") - - assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( - f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" - ) - - print("\nAccuracy Summary (Uneven Scaling):") - print(f" Initial: {initial_accuracy:.3f}") - print( - f" Scale up: {scale_up_accuracy:.3f} " - f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})" - ) - print( - f" Scale down: {scale_down_accuracy:.3f} " - f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})" - ) - print(f" Tolerance: {ACCURACY_TOL:.3f}") diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py index 53a4ce21af2e..a6cfb0d29d8e 100644 --- a/tests/distributed/test_eplb_utils.py +++ b/tests/distributed/test_eplb_utils.py @@ -16,10 +16,14 @@ def _make_model_state( phy2log: torch.Tensor, log2phy: torch.Tensor, logcnt: torch.Tensor, + phy2log_storage: torch.Tensor | None = None, ) -> MagicMock: - """Build a minimal EplbModelState mock with only the three map tensors.""" + """Build a minimal EplbModelState mock with its map tensors.""" state = MagicMock() state.physical_to_logical_map = phy2log + state.physical_to_logical_map_buffer = ( + phy2log if phy2log_storage is None else phy2log_storage + ) state.logical_to_physical_map = log2phy state.logical_replica_count = logcnt return state @@ -27,15 +31,17 @@ def _make_model_state( def test_commit_eplb_maps_shape_change(): """ - The normal path copies the physical_to_logical map in-place. When the number of - physical experts changes, the old map should be replaced entirely. + When the number of physical experts changes, resize the active map within its + preallocated storage. """ num_layers, num_logical, num_physical = 2, 4, 6 max_replicas = 3 # Build current state tensors + storage = torch.full((num_layers, num_physical + 2), -1, dtype=torch.long) model_state = _make_model_state( - phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long), + phy2log=storage[:, :num_physical], + phy2log_storage=storage, log2phy=torch.full( (num_layers, num_logical, max_replicas), -1, dtype=torch.long ), @@ -54,6 +60,7 @@ def test_commit_eplb_maps_shape_change(): # Check that the number of physical experts has been updated and that the values # match assert model_state.physical_to_logical_map.shape[1] == num_physical + 2 + assert model_state.physical_to_logical_map.data_ptr() == storage.data_ptr() assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index bd2532b1ccfb..0213e06283df 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -216,6 +216,8 @@ class ParallelConfig: enable_elastic_ep: bool = False """Enable elastic expert parallelism with stateless NCCL groups for DP/EP.""" + elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment] + """Maximum data parallel size supported by elastic expert parallelism.""" enable_dbo: bool = False """Enable dual batch overlap for the model executor.""" @@ -435,7 +437,9 @@ class is dynamically inherited by the worker class. This is used to inject ) """The configurations for fault tolerance.""" - @field_validator("disable_nccl_for_dp_synchronization", mode="wrap") + @field_validator( + "disable_nccl_for_dp_synchronization", "elastic_ep_max_dp_size", mode="wrap" + ) @classmethod def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: """Skip validation if the value is `None` when initialisation is delayed.""" @@ -941,6 +945,18 @@ def __post_init__(self) -> None: " for dense models." ) + max_dp_size = self.elastic_ep_max_dp_size + self.elastic_ep_max_dp_size = ( + max_dp_size + if self.enable_elastic_ep and max_dp_size is not None + else self.data_parallel_size + ) + if self.elastic_ep_max_dp_size < self.data_parallel_size: + raise ValueError( + "--elastic-ep-max-dp-size must be greater than or equal to " + f"the initial data_parallel_size ({self.data_parallel_size})." + ) + self.data_parallel_index = self.data_parallel_rank if self.distributed_executor_backend == "external_launcher": diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index eb187a99f4be..080e7feeff1f 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from typing import Any @@ -377,6 +379,7 @@ def clean_buffers(self) -> None: @dataclass class _NixlEPBufferState: buffer: Any + max_num_ep_ranks: int connected_ep_size: int active_ep_size: int @@ -402,7 +405,10 @@ def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) self.support_fault_tolerance = True - self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS + @property + def max_num_ep_ranks(self) -> int: + assert NixlEPAll2AllManager._buffer is not None + return NixlEPAll2AllManager._buffer.max_num_ep_ranks def _init_buffer( self, @@ -412,11 +418,19 @@ def _init_buffer( ) -> None: from nixl_ep import Buffer # type: ignore[import-not-found] - max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank + parallel_config = get_current_vllm_config().parallel_config + max_num_ep_ranks = ( + self.world_size + * parallel_config.elastic_ep_max_dp_size + // parallel_config.data_parallel_size + if parallel_config.enable_elastic_ep + else self.world_size + ) + max_num_global_experts = max_num_ep_ranks * num_experts_per_rank num_rdma_bytes = Buffer.get_rdma_size_hint( num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank, hidden=token_hidden_size, - num_ranks=self.max_num_ep_ranks, + num_ranks=max_num_ep_ranks, num_experts=max_num_global_experts, ) assert NixlEPAll2AllManager._buffer is None, ( @@ -427,7 +441,7 @@ def _init_buffer( tcp_store_group=self.tcp_store_group.store, ) buffer.update_memory_buffers( - num_ranks=self.max_num_ep_ranks, + num_ranks=max_num_ep_ranks, num_experts_per_rank=num_experts_per_rank, num_rdma_bytes=num_rdma_bytes, ) @@ -435,6 +449,7 @@ def _init_buffer( buffer.connect_ranks(ranks_to_connect) NixlEPAll2AllManager._buffer = _NixlEPBufferState( buffer=buffer, + max_num_ep_ranks=max_num_ep_ranks, connected_ep_size=self.world_size, active_ep_size=self.world_size, ) @@ -474,6 +489,7 @@ def _unmask_connected_ranks(self, target_ep_size: int) -> None: for rank in range(state.active_ep_size, target_ep_size): state.buffer.update_mask_buffer(rank, mask=False) + torch.accelerator.synchronize() state.active_ep_size = target_ep_size def _stage_ep_size(self) -> None: @@ -486,8 +502,12 @@ def _stage_ep_size(self) -> None: if target_ep_size > state.connected_ep_size: self._connect_to_ep_size(target_ep_size, make_active=False) - def commit_staged_state(self) -> None: - """Commit staged NIXL EP state to the active communication set.""" + def stage_ep_size(self) -> None: + with NixlEPAll2AllManager._lock: + self._stage_ep_size() + + def commit_ep_size(self) -> None: + """Commit the staged EP size to the active communication set.""" with NixlEPAll2AllManager._lock: assert NixlEPAll2AllManager._buffer is not None state = NixlEPAll2AllManager._buffer @@ -500,31 +520,31 @@ def commit_staged_state(self) -> None: self._unmask_connected_ranks(target_ep_size) - def _ensure_ep_size(self, *, stage: bool) -> None: - if stage: - self._stage_ep_size() - else: - self.commit_staged_state() + @contextmanager + def mask_remote_ranks(self) -> Iterator[None]: + peers = [rank for rank in range(self.world_size) if rank != self.rank] + assert NixlEPAll2AllManager._buffer is not None + buffer = NixlEPAll2AllManager._buffer.buffer + for rank in peers: + buffer.update_mask_buffer(rank, mask=True) + torch.accelerator.synchronize() + try: + yield + finally: + for rank in peers: + buffer.update_mask_buffer(rank, mask=False) + torch.accelerator.synchronize() def get_handle(self, kwargs): with NixlEPAll2AllManager._lock: - stage = bool(kwargs.get("stage", False)) state = NixlEPAll2AllManager._buffer if state is None: - assert not stage, ( - "NIXL EP staged initialization requires an existing buffer" - ) max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"] - num_experts_per_rank = ( - kwargs["num_global_experts"] // kwargs["num_ep_ranks"] - ) self._init_buffer( max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank, token_hidden_size=kwargs["token_hidden_size"], - num_experts_per_rank=num_experts_per_rank, + num_experts_per_rank=kwargs["num_local_experts"], ) - else: - self._ensure_ep_size(stage=stage) assert NixlEPAll2AllManager._buffer is not None handle = NixlEPAll2AllManager._buffer.buffer @@ -564,7 +584,7 @@ def query_active_mask(self) -> torch.Tensor: assert state is not None if NixlEPAll2AllManager._mask is None: NixlEPAll2AllManager._mask = torch.zeros( - self.max_num_ep_ranks, device="cuda", dtype=torch.int32 + state.max_num_ep_ranks, device="cuda", dtype=torch.int32 ) state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask) return NixlEPAll2AllManager._mask[: state.active_ep_size] diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 83ae25ff4437..4c1900590b99 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +from collections.abc import Iterator +from contextlib import contextmanager from weakref import WeakValueDictionary import torch @@ -157,6 +159,16 @@ def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = Fals def destroy(self): pass + def stage_ep_size(self) -> None: + pass + + def commit_ep_size(self) -> None: + pass + + @contextmanager + def mask_remote_ranks(self) -> Iterator[None]: + yield + class DeviceCommunicatorBase: """ diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 6124d22b44d8..343edd171fe3 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -2,34 +2,32 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import weakref -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Iterator, Sequence from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager, nullcontext from dataclasses import replace from typing import TYPE_CHECKING import torch import torch.nn as nn -import torch.nn.functional as F from torch.distributed import P2POp -from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.compilation.wrapper import reset_compile_wrapper -from vllm.config import ( - CompilationMode, - set_current_vllm_config, -) +from vllm.config import ParallelConfig, set_current_vllm_config from vllm.distributed import ( get_dp_group, get_ep_group, get_pcp_group, get_tp_group, + get_world_group, ) from vllm.distributed.elastic_ep.standby_state import ( create_standby_groups, get_standby_dp_group, get_standby_ep_group, get_standby_eplb_group, + get_standby_world_group, pop_standby_groups, ) from vllm.distributed.eplb.eplb_communicator import ( @@ -41,13 +39,20 @@ ) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + get_ep_all2all_manager, +) from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) +from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + sync_flashinfer_autotune_cache, +) from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType +from vllm.v1.worker.dp_utils import skip_dp_coordination from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.workspace import lock_workspace, unlock_workspace @@ -59,6 +64,15 @@ ) +def can_reuse_fused_moe_kernel(parallel_config: ParallelConfig) -> bool: + """Whether fused MoE kernels can be reused across EP rank changes. + + The EP backend must preserve its CUDA graph entries when ranks are connected, + disconnected, masked, or unmasked. + """ + return parallel_config.all2all_backend == "nixl_ep" + + def batch_transfer_weights( model: nn.Module, is_sender: bool, @@ -105,31 +119,20 @@ def batch_transfer_weights( def broadcast_expert_mapping( physical_to_logical: torch.Tensor | None, - num_local_physical_experts: int | None, - num_logical_experts: int | None, dp_group: StatelessGroupCoordinator, device: torch.device, src_rank: int = 0, -) -> tuple[torch.Tensor, int, int]: +) -> torch.Tensor: if dp_group.rank_in_group == src_rank: assert physical_to_logical is not None - assert num_local_physical_experts is not None - assert num_logical_experts is not None assert physical_to_logical.dtype == torch.int64 shape_tensor = torch.tensor( list(physical_to_logical.shape), dtype=torch.int64, device="cpu" ) - metadata_tensor = torch.tensor( - [num_local_physical_experts, num_logical_experts], - dtype=torch.int64, - device="cpu", - ) else: shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu") - metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu") shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank) - metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank) if dp_group.rank_in_group != src_rank: assert device is not None @@ -140,11 +143,7 @@ def broadcast_expert_mapping( ) assert physical_to_logical is not None - physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank) - num_local_physical_experts = int(metadata_tensor[0].item()) - num_logical_experts = int(metadata_tensor[1].item()) - - return physical_to_logical, num_local_physical_experts, num_logical_experts + return dp_group.broadcast(physical_to_logical, src_rank) class ElasticEPScalingExecutor: @@ -166,6 +165,21 @@ def worker(self): raise RuntimeError("Worker has been garbage collected") return worker + def _can_reuse_fused_moe_kernel(self) -> bool: + return can_reuse_fused_moe_kernel( + self.worker.vllm_config.parallel_config, + ) + + @contextmanager + def _disable_flashinfer_autotune(self) -> Iterator[None]: + kernel_config = self.worker.vllm_config.kernel_config + enabled = kernel_config.enable_flashinfer_autotune + kernel_config.enable_flashinfer_autotune = False + try: + yield + finally: + kernel_config.enable_flashinfer_autotune = enabled + def execute(self, execute_method: str, *args, **kwargs): method = getattr(self, execute_method, None) if method is None: @@ -260,11 +274,20 @@ def prepare_reconfiguration( use_all2all=use_all2all, enable_eplb=parallel_config.enable_eplb, ) - self.stage_standby_moe_quant_methods() + standby_ep_group = get_standby_ep_group() + assert standby_ep_group is not None + all2all_manager = get_ep_all2all_manager(standby_ep_group) + all2all_manager.stage_ep_size() + if not self._can_reuse_fused_moe_kernel(): + self.stage_standby_moe_quant_methods(all2all_manager) self._prepare_eplb_communicator(get_standby_eplb_group()) if new_dp_size > old_dp_size: self.transfer_weights(old_dp_size, new_dp_size) - self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group()) + self._warm_target_groups(get_standby_dp_group(), standby_ep_group) + if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel(): + target_world_group = get_standby_world_group() + assert target_world_group is not None + sync_flashinfer_autotune_cache(self.worker.model_runner, target_world_group) def _prepare_eplb_communicator(self, eplb_group) -> None: assert eplb_group is not None @@ -338,14 +361,9 @@ def broadcast_expert_mapping(self) -> None: assert eplb_state is not None eplb_state.drain_async() eplb_model_state = eplb_state.model_states[model_config.compute_hash()] - physical_to_logical = eplb_model_state.physical_to_logical_map - num_physical_experts = physical_to_logical.shape[1] - num_local_physical_experts = num_physical_experts // get_ep_group().world_size - num_logical_experts = eplb_model_state.logical_replica_count.shape[1] + physical_to_logical = eplb_model_state.physical_to_logical_map_buffer broadcast_expert_mapping( physical_to_logical=physical_to_logical, - num_local_physical_experts=num_local_physical_experts, - num_logical_experts=num_logical_experts, dp_group=standby_dp_group, src_rank=0, device=self.worker.device, @@ -368,7 +386,7 @@ def _make_eep_moe_config(self, module, dp_group, ep_group): moe_parallel_config=moe_parallel_config, ) - def stage_standby_moe_quant_methods(self) -> None: + def stage_standby_moe_quant_methods(self, all2all_manager) -> None: standby_dp_group = get_standby_dp_group() standby_ep_group = get_standby_ep_group() model = self.worker.model_runner.get_model() @@ -383,6 +401,7 @@ def stage_standby_moe_quant_methods(self) -> None: standby_dp_group, standby_ep_group, ), + all2all_manager=all2all_manager, ) if staged_quant_method is not None: self._staged_moe_quant_methods[module] = staged_quant_method @@ -396,7 +415,6 @@ def _commit_staged_moe_quant_methods(self) -> None: continue assert staged_quant_method.moe_kernel is not None module._replace_quant_method(staged_quant_method) - staged_quant_method.moe_kernel.prepare_finalize.on_commit() self._staged_moe_quant_methods.clear() def _release_cuda_graphs(self) -> None: @@ -405,7 +423,7 @@ def _release_cuda_graphs(self) -> None: wrapper.concrete_cudagraph_entries = {} elif isinstance(self.worker.model_runner.model, UBatchWrapper): - raise RuntimeError("DBO is not yet supported in elastic EP") + self.worker.model_runner.model.clear_graphs() torch.compiler.reset() with set_current_vllm_config(self.worker.vllm_config): @@ -427,10 +445,8 @@ def switch_and_remove(self) -> None: self._wait_for_group_cleanup() def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: - old_dp_size = get_dp_group().world_size - old_ep_size = get_ep_group().world_size - - self._release_cuda_graphs() + if not self._can_reuse_fused_moe_kernel(): + self._release_cuda_graphs() retired_groups = _replace_active_groups(**pop_standby_groups()) parallel_config = self.worker.vllm_config.parallel_config @@ -487,47 +503,13 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: parallel_config.eplb_config.num_redundant_experts = ( num_physical_experts - num_logical_experts ) - old_physical_to_logical = eplb_model_state.physical_to_logical_map - num_moe_layers = old_physical_to_logical.shape[0] - num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size - if new_dp_size > old_dp_size: - expanded_physical_to_logical = torch.full( - (num_moe_layers, num_local_experts * new_ep_size), - -1, - dtype=old_physical_to_logical.dtype, - device=old_physical_to_logical.device, - ) - expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = ( - old_physical_to_logical - ) - eplb_model_state.physical_to_logical_map = expanded_physical_to_logical - - old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1] - pad_size = num_physical_experts - old_num_physical_experts - if new_dp_size > old_dp_size: - assert pad_size > 0 - expanded_expert_load_pass = F.pad( - eplb_model_state.expert_load_pass, (0, pad_size), value=0 - ) - expanded_expert_load_window = F.pad( - eplb_model_state.expert_load_window, (0, pad_size), value=0 - ) - eplb_model_state.expert_load_pass = expanded_expert_load_pass - eplb_model_state.expert_load_window = expanded_expert_load_window - else: - assert pad_size < 0 - eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[ - :, :num_physical_experts - ] - eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[ - :, :, :num_physical_experts - ] + eplb_state.reconfigure_physical_expert_slots(model_config, num_physical_experts) model = self.worker.model_runner.get_model() model.expert_weights = [] with set_current_vllm_config(self.worker.vllm_config): model.set_eplb_state( - eplb_model_state.expert_load_pass, + eplb_model_state.expert_load_pass_buffer, eplb_model_state.logical_to_physical_map, eplb_model_state.logical_replica_count, ) @@ -539,6 +521,7 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: num_local_physical_experts=num_local_experts, ) self._commit_staged_moe_quant_methods() + get_ep_all2all_manager().commit_ep_size() # Legacy modular methods need to be recreated for the new EP size. for module in moe_modules: if getattr(module._quant_method, "wraps_legacy_quant_method", False): @@ -548,21 +531,6 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: eplb_state.update_communicator(model_config, self._prepared_eplb_communicator) self._prepared_eplb_communicator = None - if ( - self.worker.vllm_config.compilation_config.mode - == CompilationMode.STOCK_TORCH_COMPILE - ): - # NOTE(yongji): when using stock torch.compile, - # torch.compile is triggered during GPUModelRunner's load_model() - # TODO(yongji):check do we need to re-trigger torch.compile here? - # any changes to the tensor shapes in execution should already - # be handled internally by torch.compile. - backend = self.worker.vllm_config.compilation_config.init_backend( - self.worker.vllm_config - ) - compilation_counter.stock_torch_compile_count += 1 - self.worker.model_runner.model.compile(fullgraph=True, backend=backend) - return retired_groups def _perform_eplb_reshuffle( @@ -608,7 +576,8 @@ def commit_scale_up(self, is_existing_worker: bool) -> None: else: mapping = self.receive_expert_mapping() self.worker.model_runner.setup_eplb_from_mapping(mapping) - self.warm_and_capture() + if not self._can_reuse_fused_moe_kernel(): + self.warm_and_capture() self._perform_eplb_reshuffle(async_op=True) if is_existing_worker: self._start_group_cleanup(retired_groups) @@ -619,7 +588,8 @@ def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: self.switch_and_remove() else: retired_groups = self.switch_and_prepare() - self.warm_and_capture() + if not self._can_reuse_fused_moe_kernel(): + self.warm_and_capture() self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: @@ -677,42 +647,26 @@ def prepare_new_worker(self) -> None: ) torch.accelerator.synchronize() self._warm_target_groups(get_dp_group(), get_ep_group()) + if self._can_reuse_fused_moe_kernel(): + sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group()) def receive_expert_mapping(self) -> torch.Tensor: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) - physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping( + return broadcast_expert_mapping( physical_to_logical=None, - num_local_physical_experts=None, - num_logical_experts=None, dp_group=dp_group, src_rank=0, device=self.worker.device, ) - num_moe_layers = physical_to_logical.shape[0] - new_dp_size = get_dp_group().world_size - tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size - new_ep_size = new_dp_size * tp_size - expanded_physical_to_logical = torch.full( - (num_moe_layers, num_local_physical_experts * new_ep_size), - -1, - dtype=physical_to_logical.dtype, - device=physical_to_logical.device, - ) - expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = ( - physical_to_logical - ) - return expanded_physical_to_logical - def warmup_local_kernels(self) -> None: + def warmup_new_worker(self) -> None: with set_current_vllm_config(self.worker.vllm_config): kernel_warmup(self.worker, process_local_only=True) + if self._can_reuse_fused_moe_kernel(): + self.warm_and_capture() def warm_and_capture(self) -> None: - # Must run on every DP sibling in lockstep: _dummy_run calls - # coordinate_batch_across_dp whenever data_parallel_size > 1 - # (gpu_model_runner.py:3663), which deadlocks if any rank skips it. - # Save and clear block tables so the dummy MoE forward doesn't # write dummy slot mappings into real KV-cache blocks. multi_block_table = self.worker.model_runner.input_batch.block_table @@ -735,8 +689,15 @@ def warm_and_capture(self) -> None: # directly with skip_eplb=True so dummy routing doesn't pollute the # just-rebalanced EPLB stats. runner = self.worker.model_runner - runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) - self.worker.compile_or_warm_up_model() + all2all_manager = get_ep_all2all_manager() + reuse_kernel = self._can_reuse_fused_moe_kernel() + with ( + skip_dp_coordination() if reuse_kernel else nullcontext(), + all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(), + self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(), + ): + runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) + self.worker.compile_or_warm_up_model() lock_workspace() diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 6718cd8ae873..5ba1891b7aca 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -204,7 +204,7 @@ def _progress_new_engine(self) -> bool: return True elif state == ScaleUpNewEngineState.PREPARE: - self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",)) + self._collective_rpc("elastic_ep_execute", args=("warmup_new_worker",)) self._mark_ready_for_switch() tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") torch.distributed.all_reduce( diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 580477ba239e..bc3aa36bb2c5 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -121,6 +121,8 @@ class EplbModelState: [0, 2, 0, 1, 0, 3]] ``` """ + physical_to_logical_map_buffer: torch.Tensor + """Maximum-capacity buffer backing ``physical_to_logical_map``.""" logical_to_physical_map: torch.Tensor """ Mapping from logical experts to physical experts. @@ -168,6 +170,8 @@ class EplbModelState: Shape: (num_moe_layers, num_physical_experts) """ + expert_load_pass_buffer: torch.Tensor + """Maximum-capacity buffer backing ``expert_load_pass``.""" expert_load_window: torch.Tensor """ A sliding window of expert load. @@ -373,6 +377,19 @@ def add_model( physical_to_logical_map_list, device=self.device, ) + physical_expert_capacity = model.num_physical_experts + if self.parallel_config.enable_elastic_ep: + physical_expert_capacity = ( + model.num_local_physical_experts + * self.parallel_config.elastic_ep_max_dp_size + * get_ep_group().world_size + // self.parallel_config.data_parallel_size + ) + physical_to_logical_map = torch.nn.functional.pad( + physical_to_logical_map, + (0, physical_expert_capacity - model.num_physical_experts), + value=-1, + ) # Assuming 8 GPUs per node, this supports up to # (1023 + 1) / 8 = 128 nodes for now. # TODO(rui): make this configurable @@ -407,6 +424,10 @@ def add_model( ) .contiguous() ) + physical_to_logical_map_buffer = physical_to_logical_map + physical_to_logical_map = physical_to_logical_map_buffer[ + :, : model.num_physical_experts + ] logical_to_physical_map = ( logical_to_physical_map.unsqueeze(0) .expand( @@ -425,11 +446,12 @@ def add_model( .contiguous() ) - expert_load_pass = torch.zeros( - (model.num_moe_layers, model.num_physical_experts), + expert_load_pass_buffer = torch.zeros( + (model.num_moe_layers, physical_expert_capacity), dtype=torch.int32, device=self.device, ) + expert_load_pass = expert_load_pass_buffer[:, : model.num_physical_experts] self.expert_load_window_size = self.parallel_config.eplb_config.window_size expert_load_window = torch.zeros( ( @@ -460,7 +482,7 @@ def add_model( ] model.set_eplb_state( - expert_load_pass, + expert_load_pass_buffer, logical_to_physical_map, logical_replica_count, ) @@ -479,9 +501,11 @@ def add_model( model_state = EplbModelState( physical_to_logical_map=physical_to_logical_map, + physical_to_logical_map_buffer=physical_to_logical_map_buffer, logical_to_physical_map=logical_to_physical_map, logical_replica_count=logical_replica_count, expert_load_pass=expert_load_pass, + expert_load_pass_buffer=expert_load_pass_buffer, expert_load_window=expert_load_window, model_name=model_config.model, model=model, @@ -1075,10 +1099,14 @@ def update_mapping( expanded_physical_to_logical: torch.Tensor, ) -> None: eplb_model_state = self.model_states[model_config.compute_hash()] - eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical) + eplb_model_state.physical_to_logical_map_buffer.copy_( + expanded_physical_to_logical + ) + eplb_model_state.expert_load_pass.zero_() + physical_to_logical_map = eplb_model_state.physical_to_logical_map (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps( - expanded_physical_to_logical.cpu(), + physical_to_logical_map.cpu(), eplb_model_state.model.num_logical_experts, ) @@ -1097,6 +1125,33 @@ def update_mapping( eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map) eplb_model_state.logical_replica_count.copy_(logical_replica_count) + def reconfigure_physical_expert_slots( + self, + model_config: ModelConfig, + num_physical_experts: int, + ) -> None: + """Replace physical_to_logical_map and expert_load_pass with views + covering the new active size. + """ + model_state = self.model_states[model_config.compute_hash()] + old_num_physical_experts = model_state.model.num_physical_experts + first_slot, last_slot = sorted((old_num_physical_experts, num_physical_experts)) + physical_to_logical_map_buffer = model_state.physical_to_logical_map_buffer + expert_load_pass_buffer = model_state.expert_load_pass_buffer + assert last_slot <= physical_to_logical_map_buffer.shape[1] + expert_slots = slice(first_slot, last_slot) + physical_to_logical_map_buffer[:, expert_slots].fill_(-1) + expert_load_pass_buffer[:, expert_slots].zero_() + model_state.physical_to_logical_map = physical_to_logical_map_buffer[ + :, :num_physical_experts + ] + model_state.expert_load_pass = expert_load_pass_buffer[:, :num_physical_experts] + + pad_size = num_physical_experts - model_state.expert_load_window.shape[-1] + model_state.expert_load_window = torch.nn.functional.pad( + model_state.expert_load_window, (0, pad_size) + ) + def create_communicator( self, model_config: ModelConfig, group_coordinator: GroupCoordinator ) -> EplbCommunicator: @@ -1333,14 +1388,13 @@ def _commit_eplb_maps( if PIN_MEMORY and src.is_cpu: src = src.new_empty(src.shape, pin_memory=True).copy_(src) - # Rare Case: When the number of physical experts has changed, discard the old - # physical to logical expert map and use the new one. This only happens when the - # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the - # new map into the old one. + # When the number of physical experts changes, refresh the active map view to + # cover the new size before copying it. if src.shape[1] != dst.shape[1]: - model_state.physical_to_logical_map = src.to(dst.device, non_blocking=True) - else: - dst.copy_(src, non_blocking=True) + dst = model_state.physical_to_logical_map_buffer[:, : src.shape[1]] + assert dst.shape == src.shape + model_state.physical_to_logical_map = dst + dst.copy_(src, non_blocking=True) num_logical_experts = model_state.logical_to_physical_map.shape[1] new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index db398438d0dd..ae6f14b1df41 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -507,6 +507,7 @@ class EngineArgs: linear_backend: LinearBackend = KernelConfig.linear_backend all2all_backend: All2AllBackend = ParallelConfig.all2all_backend enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep + elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size enable_dbo: bool = ParallelConfig.enable_dbo ubatch_size: int = ParallelConfig.ubatch_size dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold @@ -1201,6 +1202,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: parallel_group.add_argument( "--enable-elastic-ep", **parallel_kwargs["enable_elastic_ep"] ) + parallel_group.add_argument( + "--elastic-ep-max-dp-size", + **parallel_kwargs["elastic_ep_max_dp_size"], + ) parallel_group.add_argument( "--dbo-decode-token-threshold", **parallel_kwargs["dbo_decode_token_threshold"], @@ -2345,6 +2350,7 @@ def create_engine_config( enable_ep_weight_filter=self.enable_ep_weight_filter, all2all_backend=self.all2all_backend, enable_elastic_ep=self.enable_elastic_ep, + elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, enable_dbo=self.enable_dbo, ubatch_size=self.ubatch_size, dbo_decode_token_threshold=self.dbo_decode_token_threshold, diff --git a/vllm/envs.py b/vllm/envs.py index 1c2246639822..aa32fc38bafb 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -316,7 +316,6 @@ VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True - VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False VLLM_XPU_FORCE_N_CONTIG_WEIGHT: bool = False VLLM_XPU_USE_SAMPLER_KERNEL: bool = True @@ -2129,10 +2128,6 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) ), - # NIXL EP environment variables - "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int( - os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32") - ), # Whether enable XPU graph on Intel GPU "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool( int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0")) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 7861345453a0..5019867d53ed 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -8,6 +8,7 @@ from vllm.config import get_current_vllm_config from vllm.distributed import ( + GroupCoordinator, get_ep_group, ) from vllm.logger import init_logger @@ -104,16 +105,12 @@ def flashinfer_one_sided_dispatch_layout( ) -def get_ep_all2all_manager(eep_stage: bool = False) -> Any: - if eep_stage: - from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group - - ep_group = get_standby_ep_group() - assert ep_group is not None - device_communicator = ep_group.device_communicator - else: - device_communicator = get_ep_group().device_communicator - +def get_ep_all2all_manager( + ep_group: GroupCoordinator | None = None, +) -> Any: + if ep_group is None: + ep_group = get_ep_group() + device_communicator = ep_group.device_communicator assert device_communicator is not None all2all_manager = device_communicator.all2all_manager assert all2all_manager is not None @@ -168,7 +165,7 @@ def maybe_make_prepare_finalize( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, use_monolithic: bool = False, - eep_stage: bool = False, + all2all_manager: Any | None = None, ) -> FusedMoEPrepareAndFinalize | None: if not moe.moe_parallel_config.use_all2all_kernels: if not allow_new_interface: @@ -190,7 +187,8 @@ def maybe_make_prepare_finalize( "Detected DP deployment with no --enable-expert-parallel. " "Falling back to AllGather+ReduceScatter dispatch/combine." ) - all2all_manager = get_ep_all2all_manager(eep_stage) + if all2all_manager is None: + all2all_manager = get_ep_all2all_manager() return make_moe_prepare_and_finalize_naive_dp_ep( is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=all2all_manager.world_size, @@ -199,7 +197,8 @@ def maybe_make_prepare_finalize( else: return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) - all2all_manager = get_ep_all2all_manager(eep_stage) + if all2all_manager is None: + all2all_manager = get_ep_all2all_manager() prepare_finalize: FusedMoEPrepareAndFinalize | None = None @@ -361,10 +360,7 @@ def maybe_make_prepare_finalize( all_to_all_args = dict( max_num_tokens_per_dp_rank=moe.max_num_tokens, token_hidden_size=moe.hidden_dim, - num_ep_ranks=all2all_manager.world_size, - num_global_experts=moe.num_experts, - num_local_experts=moe.num_experts // all2all_manager.world_size, - stage=eep_stage, + num_local_experts=moe.num_local_experts, ) handle = all2all_manager.get_handle(all_to_all_args) @@ -378,7 +374,8 @@ def maybe_make_prepare_finalize( prepare_finalize = NixlEPPrepareAndFinalize( handle, max_tokens_per_rank=moe.max_num_tokens, - num_dispatchers=all2all_manager.world_size, + num_dispatchers=all2all_manager.max_num_ep_ranks, + expert_capacity=(moe.num_local_experts * all2all_manager.max_num_ep_ranks), use_fp8_dispatch=use_fp8_dispatch, global_to_physical=global_to_physical, physical_to_global=physical_to_global, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 8736b1393d82..4c7760d00792 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1306,6 +1306,7 @@ class FusedMoEConfig: moe_backend: MoEBackend = "auto" max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP + elastic_ep_max_dp_size: int | None = None has_bias: bool = False is_lora_enabled: bool = False diff --git a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py index 2c6c3f6aa40a..9b6a1da0f9e3 100644 --- a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py +++ b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py @@ -60,6 +60,8 @@ def _make_eep_experts( def make_eep_staged_quant_method( module: "MoERunner", moe_config: FusedMoEConfig, + *, + all2all_manager: Any, ) -> FusedMoEMethodBase | None: quant_method = module._quant_method if not quant_method.supports_internal_mk: @@ -93,7 +95,7 @@ def make_eep_staged_quant_method( routing_tables=None, allow_new_interface=True, use_monolithic=quant_method.is_monolithic, - eep_stage=True, + all2all_manager=all2all_manager, ) assert prepare_finalize is not None assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 4b633dd99b7a..a9c1a66d87dc 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -348,6 +348,7 @@ def FusedMoEFactory( moe_backend=vllm_config.kernel_config.moe_backend, router_logits_dtype=router_logits_dtype, max_num_tokens=max_num_batched_tokens, + elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, has_bias=has_bias, is_lora_enabled=vllm_config.lora_config is not None, activation=moe_activation, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 778b95bf47ee..120c9e985753 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -252,13 +252,6 @@ def supports_async(self) -> bool: """ return False - def on_commit(self) -> None: - """ - Runs after this prepare/finalize has been committed to the active - MoE kernel. - """ - return - # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index fd7d45655dd6..b2c85ce0c0ae 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -6,8 +6,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.distributed import get_ep_group -from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( @@ -79,6 +77,7 @@ def __init__( buffer: nixl_ep.Buffer, max_tokens_per_rank: int, num_dispatchers: int, + expert_capacity: int, use_fp8_dispatch: bool = False, global_to_physical: torch.Tensor | None = None, physical_to_global: torch.Tensor | None = None, @@ -94,6 +93,7 @@ def __init__( # combine function. self.handles: list[tuple | None] = [None, None] self.num_dispatchers_ = num_dispatchers + self.expert_capacity = expert_capacity topk_indices_dtype = self.topk_indices_dtype() @@ -141,17 +141,6 @@ def activation_format(self) -> mk.FusedMoEActivationFormat: def max_num_tokens_per_rank(self) -> int | None: return self.max_tokens_per_rank - def on_commit(self) -> None: - device_communicator = get_ep_group().device_communicator - assert device_communicator is not None - all2all_manager = device_communicator.all2all_manager - assert isinstance(all2all_manager, NixlEPAll2AllManager) - # maybe_make_prepare_finalize(..., eep_stage=True) initializes self.buffer - # with get_handle(..., stage=True), which stages global NIXL state for the - # new config but leaves it inactive while the old config remains active. - # When EEP commit switches to this P/F, this P/F needs to commit that state. - all2all_manager.commit_staged_state() - def topk_indices_dtype(self) -> torch.dtype | None: return NIXL_EP_TOPK_INDICES_DTYPE @@ -267,7 +256,7 @@ def prepare_async( a1, dispatch_topk_ids, self.max_tokens_per_rank, - num_experts, + self.expert_capacity, use_fp8=self.use_fp8_dispatch, # round_scale needs to be set to dispatch in ue8m0 round_scale=self.use_ue8m0_dispatch, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 48e217b1a91a..cfc3be3e8e12 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -500,7 +500,8 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int: For a detailed explanation, see: `docs/serving/data_parallel_deployment.md` """ - return max(moe_config.max_num_tokens * moe_config.dp_size, 8192) + dp_size = moe_config.elastic_ep_max_dp_size or moe_config.dp_size + return max(moe_config.max_num_tokens * dp_size, 8192) @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index 37ea4812fc6c..afc3ef3ac014 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -12,6 +12,7 @@ import vllm.envs as envs if TYPE_CHECKING: + from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -40,6 +41,44 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: return output_dir / "autotune_configs.json" +def sync_flashinfer_autotune_cache( + runner: "GPUModelRunner", + group: "GroupCoordinator", +) -> None: + cache: bytes | str | None = None + if ( + group.rank_in_group == 0 + and runner.vllm_config.kernel_config.enable_flashinfer_autotune + ): + try: + from vllm.platforms import current_platform + from vllm.utils.flashinfer import has_flashinfer + + if has_flashinfer() and current_platform.has_device_capability(90): + from flashinfer.autotuner import AutoTuner + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "autotune_configs.json" + AutoTuner.get().save_configs(str(path)) + cache = path.read_bytes() + except Exception as exc: + cache = f"{type(exc).__name__}: {exc}" + + cache = group.broadcast_object(cache) + if isinstance(cache, str): + raise RuntimeError(f"Failed to serialize FlashInfer autotune state: {cache}") + if cache is None or group.rank_in_group == 0: + return + + from flashinfer.autotuner import AutoTuner + + with tempfile.NamedTemporaryFile() as f: + f.write(cache) + f.flush() + if not AutoTuner.get().load_configs(f.name): + raise RuntimeError("FlashInfer autotune cache is incompatible") + + def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None: cache_path.parent.mkdir(parents=True, exist_ok=True) fd, tmp_path = tempfile.mkstemp( diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index ffd10ab5572f..c354ba699c75 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1711,6 +1711,12 @@ async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: "Only ray DP backend supports scaling elastic EP" ) parallel_config = self.vllm_config.parallel_config + if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size: + raise ValueError( + f"Cannot scale to data_parallel_size {new_data_parallel_size}; " + "--elastic-ep-max-dp-size is " + f"{parallel_config.elastic_ep_max_dp_size}." + ) num_experts = self.vllm_config.model_config.get_num_experts() num_physical_experts = ( num_experts + parallel_config.eplb_config.num_redundant_experts diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index 07bda704833e..fe82ec1427f1 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar import torch import torch.distributed as dist @@ -18,6 +19,22 @@ logger = init_logger(__name__) +_skip_dp_coordination = ContextVar("skip_dp_coordination", default=False) + + +@contextmanager +def skip_dp_coordination(): + """Run without coordinating DP metadata with other DP ranks.""" + token = _skip_dp_coordination.set(True) + try: + yield + finally: + _skip_dp_coordination.reset(token) + + +def should_skip_dp_coordination() -> bool: + return _skip_dp_coordination.get() + def _get_device_and_group(parallel_config: ParallelConfig): # Use the actual device assigned to the DP group, not just the device type @@ -223,6 +240,19 @@ def coordinate_batch_across_dp( if num_tokens_padded is None: num_tokens_padded = num_tokens_unpadded + if should_skip_dp_coordination(): + should_ubatch = should_attempt_ubatching and not is_last_ubatch_empty( + num_tokens_unpadded, + num_tokens_padded, + parallel_config.num_ubatches, + ) + num_tokens_after_padding = torch.full( + (parallel_config.data_parallel_size,), + num_tokens_padded, + dtype=torch.int32, + ) + return should_ubatch, num_tokens_after_padding, cudagraph_mode + (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = ( _synchronize_dp_ranks( num_tokens_unpadded, diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index dddf630cc335..b7c73789de6f 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -10,6 +10,7 @@ from vllm.config import ParallelConfig from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import get_dp_group +from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, @@ -71,7 +72,10 @@ def sync_cudagraph_and_dp_padding( tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) tensor[4][dp_rank] = int(allow_ubatching) tensor[5][dp_rank] = num_reqs - dist.all_reduce(tensor, group=group) + if should_skip_dp_coordination(): + tensor[:] = tensor[:, dp_rank, None].clone() + else: + dist.all_reduce(tensor, group=group) num_tokens_across_dp = tensor[0] cg_mode_across_dp = tensor[1] From 22f6e4eccb674b534f62810c66838317169b6b98 Mon Sep 17 00:00:00 2001 From: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:12:25 -0400 Subject: [PATCH 099/774] [Frontend][Rust] Reject empty structured-output values (#54821) Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> --- .../src/protocol/structured_outputs.rs | 90 ++++++++++++++++++- rust/src/server/src/grpc/convert.rs | 18 ++++ .../routes/openai/chat_completions/convert.rs | 57 +++++++++++- .../routes/openai/utils/structured_outputs.rs | 21 +++-- 4 files changed, 175 insertions(+), 11 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs.rs b/rust/src/engine-core-client/src/protocol/structured_outputs.rs index 745ca3150b9b..99fa66ae2b02 100644 --- a/rust/src/engine-core-client/src/protocol/structured_outputs.rs +++ b/rust/src/engine-core-client/src/protocol/structured_outputs.rs @@ -104,6 +104,41 @@ impl StructuredOutputsParams { backend: StructuredOutputBackend::default(), } } + + pub fn validate(&self) -> Result<()> { + match &self.constraint { + StructuredOutputConstraint::Grammar(grammar) if grammar.trim().is_empty() => { + Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.grammar cannot be an empty string".to_string(), + }) + } + StructuredOutputConstraint::Json(Value::String(value)) => { + if value.trim().is_empty() { + return Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.json cannot be an empty string".to_string(), + }); + } + serde_json::from_str::(value).map_err(|_| { + Error::InvalidStructuredOutputsParams { + message: "Invalid JSON grammar specification.".to_string(), + } + })?; + Ok(()) + } + StructuredOutputConstraint::Choice(choice) if choice.is_empty() => { + Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.choice cannot be empty".to_string(), + }) + } + StructuredOutputConstraint::StructuralTag(tag) if tag.trim().is_empty() => { + Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.structural_tag cannot be an empty string" + .to_string(), + }) + } + _ => Ok(()), + } + } } /// Wire-compatible structured-output payload used by Python engine-core. @@ -176,7 +211,7 @@ impl TryFrom for StructuredOutputsParams { } insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag)); - Ok(Self { + let params = Self { constraint: constraint.map(|(_, c)| c).ok_or_else(|| { Error::InvalidStructuredOutputsParams { message: "missing structured output constraint".to_string(), @@ -188,7 +223,9 @@ impl TryFrom for StructuredOutputsParams { whitespace_pattern: raw.whitespace_pattern, }, backend: raw.backend, - }) + }; + params.validate()?; + Ok(params) } } @@ -289,6 +326,55 @@ mod tests { assert!(error.to_string().contains("json_object must be true")); } + #[test] + fn structured_outputs_rejects_empty_json_string() { + let error = serde_json::from_value::(serde_json::json!({ + "json": " ", + })) + .unwrap_err(); + + assert!(error.to_string().contains("json cannot be an empty string")); + } + + #[test] + fn validate_rejects_empty_grammar_from_constructor() { + let error = StructuredOutputsParams::grammar(" ").validate().unwrap_err(); + assert!(error.to_string().contains("grammar cannot be an empty string")); + } + + #[test] + fn validate_rejects_empty_json_string_from_constructor() { + let error = StructuredOutputsParams::json(Value::String(" ".to_string())) + .validate() + .unwrap_err(); + assert!(error.to_string().contains("json cannot be an empty string")); + } + + #[test] + fn validate_rejects_malformed_json_string() { + let error = StructuredOutputsParams::json(Value::String("not json".to_string())) + .validate() + .unwrap_err(); + assert!(error.to_string().contains("Invalid JSON grammar specification.")); + } + + #[test] + fn validate_rejects_empty_choice() { + let error = StructuredOutputsParams::choice(vec![]).validate().unwrap_err(); + assert!(error.to_string().contains("choice cannot be empty")); + } + + #[test] + fn validate_rejects_empty_structural_tag() { + let error = StructuredOutputsParams::structural_tag("\n\t").validate().unwrap_err(); + assert!(error.to_string().contains("structural_tag cannot be an empty string")); + } + + #[test] + fn validate_accepts_non_empty_grammar() { + assert!(StructuredOutputsParams::grammar("root ::= \"a\"").validate().is_ok()); + } + #[test] fn structured_outputs_serializes_through_raw_shape() { let params = StructuredOutputsParams { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index cf0045f71470..61eb21ef70ee 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -4,6 +4,7 @@ //! Conversion between gRPC protobuf types and internal `vllm-text` //! request/response types. +use thiserror_ext::AsReport as _; use tonic::Status; use uuid::Uuid; use vllm_engine_core_client::protocol::output::StopReason; @@ -259,6 +260,9 @@ fn convert_structured_output( StructuredOutputsParams::structural_tag(tag.clone()) } }; + params + .validate() + .map_err(|error| Status::invalid_argument(error.to_report_string()))?; Ok(Some(params)) } @@ -546,6 +550,20 @@ mod tests { assert_eq!(text.sampling_params.temperature, Some(0.0)); } + #[test] + fn grpc_rejects_empty_grammar_before_engine() { + use super::pb::decoding_parameters::StructuredOutput; + let req = pb::GenerateRequest { + decoding: Some(pb::DecodingParameters { + structured_output: Some(StructuredOutput::Grammar(" ".to_string())), + ..Default::default() + }), + ..base_request() + }; + let err = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); + assert!(err.message().contains("grammar cannot be an empty string")); + } + #[test] fn absent_seed_is_none() { let req = pb::GenerateRequest { diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 3aaba8df28dc..2c64ffbef48b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -437,7 +437,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use axum::http::HeaderMap; + use axum::http::{HeaderMap, StatusCode}; use expect_test::expect; use llm_multimodal::ImageDetail; use serde_json::json; @@ -583,6 +583,61 @@ mod tests { assert_eq!(prepared.chat_request.decode_options.min_tokens, 0); } + #[test] + fn prepare_chat_request_rejects_empty_json_schema() { + let mut request = base_request(); + request.response_format = Some(ResponseFormat::JsonSchema { + json_schema: JsonSchemaFormat { + name: "answer".to_string(), + description: None, + schema: json!(" "), + strict: None, + }, + }); + + let error = prepare_chat_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .unwrap_err(); + + assert_eq!(error.status_code(), StatusCode::BAD_REQUEST); + assert!( + error + .to_error_response() + .error + .message + .contains("json cannot be an empty string") + ); + } + + #[test] + fn prepare_chat_request_rejects_whitespace_only_structured_outputs_grammar() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hello"}], + "structured_outputs": {"grammar": " \t\n"}, + })) + .expect("parse structured_outputs"); + + let error = prepare_chat_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .unwrap_err(); + + assert_eq!(error.status_code(), StatusCode::BAD_REQUEST); + assert!( + error + .to_error_response() + .error + .message + .contains("grammar cannot be an empty string") + ); + } + #[test] fn prepare_chat_request_defaults_function_tool_fields() { let request: ChatCompletionRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/utils/structured_outputs.rs b/rust/src/server/src/routes/openai/utils/structured_outputs.rs index 849870f32aaa..864e00818beb 100644 --- a/rust/src/server/src/routes/openai/utils/structured_outputs.rs +++ b/rust/src/server/src/routes/openai/utils/structured_outputs.rs @@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; +use thiserror_ext::AsReport as _; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use crate::error::ApiError; @@ -80,12 +81,12 @@ pub fn convert_from_response_format( let Some(fmt) = response_format else { return Ok(None); }; - match fmt { - ResponseFormat::Text => Ok(None), - ResponseFormat::JsonObject => Ok(Some(StructuredOutputsParams::json_object())), - ResponseFormat::JsonSchema { json_schema } => Ok(Some(StructuredOutputsParams::json( - json_schema.schema.clone(), - ))), + let params = match fmt { + ResponseFormat::Text => return Ok(None), + ResponseFormat::JsonObject => StructuredOutputsParams::json_object(), + ResponseFormat::JsonSchema { json_schema } => { + StructuredOutputsParams::json(json_schema.schema.clone()) + } ResponseFormat::StructuralTag { .. } => { // The Python frontend dumps the entire response_format object (including the // `type` field) as a JSON string for the engine-core backend. @@ -95,9 +96,13 @@ pub fn convert_from_response_format( Some("response_format"), ) })?; - Ok(Some(StructuredOutputsParams::structural_tag(tag_json))) + StructuredOutputsParams::structural_tag(tag_json) } - } + }; + params.validate().map_err(|error| { + ApiError::invalid_request(error.to_report_string(), Some("response_format")) + })?; + Ok(Some(params)) } /// Convert raw `response_format` and/or `structured_outputs` JSON blobs into From c377114636db645e13594e66a6c4f4f0084adac2 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:12:58 +0200 Subject: [PATCH 100/774] [Bugfix][CI] Update CuTeDSL indexer Q sentinel test for migrated wrappers (#56594) Signed-off-by: LopezCastroRoberto --- .../test_fused_indexer_q_rope_quant.py | 38 +++++++++---------- .../nvidia/ops/fused_indexer_q_cutedsl.py | 2 +- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 8fd19ffb0b43..73b147584c51 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -289,16 +289,16 @@ def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head): dtype=torch.uint8, device=device, ) - mod.fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions[:num_tokens], - q[:num_tokens], - cos_sin_cache, - weights[:num_tokens], - 1.0, - 1.0, - outputs["q_packed"][:num_tokens], - outputs["q_scale"][:num_tokens], - weights_out[:num_tokens], + mod._INDEXER_Q_MXFP4_KERNEL( + positions=positions[:num_tokens], + q=q[:num_tokens], + cos_sin_cache=cos_sin_cache, + weights=weights[:num_tokens], + weights_softmax_scale=1.0, + weights_head_scale=1.0, + q_packed=outputs["q_packed"][:num_tokens], + q_scale=outputs["q_scale"][:num_tokens], + weights_out=weights_out[:num_tokens], ) else: outputs["q_fp8"] = torch.full( @@ -307,15 +307,15 @@ def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head): dtype=torch.uint8, device=device, ) - mod.fused_indexer_q_rope_quant_fp8_cutedsl( - positions[:num_tokens], - q[:num_tokens], - cos_sin_cache, - weights[:num_tokens], - 1.0, - 1.0, - outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn), - weights_out[:num_tokens], + mod._INDEXER_Q_FP8_KERNEL( + positions=positions[:num_tokens], + q=q[:num_tokens], + cos_sin_cache=cos_sin_cache, + weights=weights[:num_tokens], + weights_softmax_scale=1.0, + weights_head_scale=1.0, + q_fp8=outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn), + weights_out=weights_out[:num_tokens], ) torch.accelerator.synchronize() diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 6684e37f37e5..8052178a03d9 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -684,7 +684,7 @@ def __call__( q, cos_sin_cache, weights, - q_fp8, + q_fp8.view(torch.uint8), weights_out, float(weights_softmax_scale * weights_head_scale), ) From 9d3e991ece7f54c173c0cc9261ab5969b1dbeb4c Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Sat, 12 Sep 2026 19:32:12 +0300 Subject: [PATCH 101/774] Add @arpera to CODEOWNERS of Structured Output (#56501) Signed-off-by: Artem Perevedentsev --- .github/CODEOWNERS | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1f605db84a7e..10fd12576893 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -73,7 +73,7 @@ /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni -/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett +/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett @arpera /vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium @@ -108,7 +108,7 @@ /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm -/tests/v1/structured_output @mgoin @russellb @aarnphm +/tests/v1/structured_output @mgoin @russellb @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee From 3bb23e2def9527d941be57c5a7d47401a8dbaedc Mon Sep 17 00:00:00 2001 From: jcotant-inferact Date: Sat, 12 Sep 2026 09:50:16 -0700 Subject: [PATCH 102/774] [CI/Build] Add DSv4.1 auto-label rules and narrow DSv4 (#56332) Signed-off-by: Joe Cotant Co-authored-by: Claude --- .github/mergify.yml | 22 +++++++++++++++++++++- .github/workflows/issue_autolabel.yml | 7 +++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/.github/mergify.yml b/.github/mergify.yml index 04a623a499d1..ba3a69ee3575 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -330,12 +330,32 @@ pull_request_rules: - label != stale - or: - files~=^vllm/models/deepseek_v4/ - - title~=(?i)(?:\bDSv4\b|deepseek[-\s]?v4) + # Trailing (?![._]?1) keeps V4.1 out: DeepseekV41ForCausalLM is a + # distinct architecture with its own tree and label (DSv4.1). + - title~=(?i)(?:\bDSv4\b|deepseek[-\s_]?v4)(?![._]?1) actions: label: add: - DSv4 +- name: label-dsv4-1 + description: Automatically apply DSv4.1 label + conditions: + - label != stale + - or: + # The 4.1 sources spell the model three ways: deepseek_v4_1 (models), + # deepseek_v41 (most) and deepseekv41 (tool parsers). + - files~=^vllm/models/deepseek_v4_1/ + - files~=^vllm/(?:parser|tokenizers|reasoning|transformers_utils/configs)/deepseek_v41.*\.py + - files~=^vllm/tool_parsers/deepseekv41.*\.py + - files~=^rust/src/(?:chat|parser)/src/.*deepseek_v41 + - files~=^tests/.*deepseek_?v41 + - title~=(?i)(?:\bDSv4[._]?1\b|deepseek[-\s_]?v4[._]?1) + actions: + label: + add: + - DSv4.1 + - name: label-nvidia description: Automatically apply nvidia label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index b2392e87ce5c..fc6e4921ef83 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -197,9 +197,12 @@ jobs: deepseek: { keywords: [{ term: "DeepSeek", searchIn: "title" }], }, + // DSv4 and DSv4.1 are distinct architectures; these keep them exclusive. DSv4: { - keywords: [{ term: "DSv4", searchIn: "title" }], - substrings: [{ term: "deepseek-v4", searchIn: "title" }], + regexPatterns: [{ pattern: "(?:\\bDSv4\\b|deepseek[-\\s_]?v4)(?![._]?1)", flags: "gi", searchIn: "title" }], + }, + "DSv4.1": { + regexPatterns: [{ pattern: "(?:\\bDSv4[._]?1\\b|deepseek[-\\s_]?v4[._]?1)", flags: "gi", searchIn: "title" }], }, llama: { keywords: [{ term: "Llama", searchIn: "title" }], From dff76bc3e8d702901e6dda5971f9506a6c1bc00f Mon Sep 17 00:00:00 2001 From: Zeyu Yang <532183776@qq.com> Date: Sun, 13 Sep 2026 01:03:28 +0800 Subject: [PATCH 103/774] [Bugfix][CI] skip conftest for NPU compatibility test (#55799) Signed-off-by: yangzeyu <532183776@qq.com> --- .../scripts/hardware_ci/run-npu-test.sh | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-npu-test.sh b/.buildkite/scripts/hardware_ci/run-npu-test.sh index f7227095ac2e..d0c973f6daff 100644 --- a/.buildkite/scripts/hardware_ci/run-npu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-npu-test.sh @@ -76,21 +76,12 @@ FROM ${BASE_IMAGE_NAME} ENV DEBIAN_FRONTEND=noninteractive ENV SOC_VERSION="ascend910b1" -# Install uv and configure uv to use internal PyPI mirror -RUN curl -LsSf https://astral.sh/uv/install.sh | sh -ENV PATH="/root/.local/bin:$PATH" -ENV UV_SYSTEM_PYTHON=true -ENV UV_INDEX_URL="http://cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_PORT}/pypi/simple" -ENV UV_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" - -RUN apt-get update -y && \ - apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \ - rm -rf /var/cache/apt/* && \ - rm -rf /var/lib/apt/lists/* +RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_PORT}/pypi/simple && \ + pip config set global.trusted-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local # Install for pytest to make the docker build cache layer always valid -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install pytest>=6.0 'modelscope<1.38' --quiet +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install pytest>=6.0 WORKDIR /workspace/vllm @@ -167,5 +158,5 @@ docker run \ "${image_name}" \ bash -c ' set -e - pytest -v -s tests/e2e/vllm_interface/test_vllm_pr_interface_compatibility.py + pytest --noconftest -v -s tests/e2e/vllm_interface/test_vllm_pr_interface_compatibility.py ' From 29332cf93628a67c49252ef7c89a8aa3185def4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Sat, 12 Sep 2026 19:47:55 +0200 Subject: [PATCH 104/774] [4/N] Expose HiSparse cache metrics via KV connector stats (#56061) Signed-off-by: NickLucche Signed-off-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- csrc/libtorch_stable/hisparse_kernels.cu | 20 +++- csrc/libtorch_stable/ops.h | 1 + csrc/libtorch_stable/torch_bindings.cpp | 13 +-- .../kv_connector/unit/test_hisparse_stats.py | 62 ++++++++++++ tests/v1/worker/test_utils.py | 39 +++++++- .../kv_connector/v1/hisparse/connector.py | 37 +++++++ .../kv_connector/v1/hisparse/stats.py | 97 +++++++++++++++++++ .../kv_connector/v1/hisparse/worker.py | 39 ++++++++ vllm/v1/hisparse/runtime.py | 8 ++ 9 files changed, 308 insertions(+), 8 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_hisparse_stats.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu index 63d84576272b..97ce9246cc75 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/libtorch_stable/hisparse_kernels.cu @@ -203,6 +203,7 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( int32_t* __restrict__ miss_mask, // [num_rows, top_k] or nullptr int32_t* __restrict__ device_global_indices, // [max_rows, region_stride] int16_t* __restrict__ lru_slots, // [max_rows, hot_size] + unsigned long long* __restrict__ stats, // [2] hits,misses or nullptr const int32_t* __restrict__ request_state_indices, // [num_requests] or // nullptr const int32_t request_state_count, const int64_t host_rows, @@ -552,6 +553,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( if (swap_counts != nullptr && tid == 0) { swap_counts[batch_row] = total_misses; } + if (stats != nullptr && tid == 0) { + atomicAdd(&stats[0], static_cast(total_hits)); + atomicAdd(&stats[1], static_cast(total_misses)); + } + // Phase 4: write back the LRU order: stale evictables at the front, // freshly loaded misses next, then hits at MRU. const int total_evictable = hot_size - total_hits; @@ -779,6 +785,7 @@ void hisparse_resolve_residency( std::optional const& request_state_indices, int64_t region_stride, std::optional const& miss_mask, + std::optional const& stats, std::optional const& attention_indices, int64_t attention_block_stride, std::optional const& request_ids, @@ -997,6 +1004,17 @@ void hisparse_resolve_residency( miss_mask_ptr = mm.mutable_data_ptr(); } + unsigned long long* stats_ptr = nullptr; + if (stats.has_value()) { + auto const& st = stats.value(); + STD_TORCH_CHECK( + st.is_cuda() && st.is_contiguous() && st.dim() == 1 && + st.scalar_type() == torch::headeronly::ScalarType::UInt64 && + st.numel() == 2, + "stats must be a contiguous two-element uint64 CUDA tensor"); + stats_ptr = static_cast(st.mutable_data_ptr()); + } + int32_t* attention_indices_ptr = nullptr; if (attention_indices.has_value()) { auto const& indices = attention_indices.value(); @@ -1046,7 +1064,7 @@ void hisparse_resolve_residency( swap_counts_ptr, hot_indices.mutable_data_ptr(), attention_indices_ptr, miss_mask_ptr, device_global_indices.mutable_data_ptr(), - lru_slots.mutable_data_ptr(), request_state_ptr, + lru_slots.mutable_data_ptr(), stats_ptr, request_state_ptr, request_state_count, host_rows, hot_block_table.stride(0), hot_block_size, top_k, hot_size, hash_size, region_stride, attention_block_stride, source_bt_stride, source_num_reqs, source_num_blocks, diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 6893c6913d9e..68a24ca4340a 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -673,6 +673,7 @@ void hisparse_resolve_residency( std::optional const& request_state_indices, int64_t region_stride, std::optional const& miss_mask, + std::optional const& stats, std::optional const& attention_indices, int64_t attention_block_stride, std::optional const& request_ids, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index cb6f46339683..110f8595ee99 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -999,16 +999,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor? request_state_indices," " int region_stride," " Tensor(a!)? miss_mask=None," - " Tensor(b!)? attention_indices=None," + " Tensor(b!)? stats=None," + " Tensor(c!)? attention_indices=None," " int attention_block_stride=0," " Tensor? request_ids=None," " Tensor? source_block_table=None," " int source_block_size=0," - " Tensor(c!)? resolved_global_indices=None," - " Tensor(d!)? valid_counts=None," - " Tensor(e!)? swap_host_physical_rows=None," - " Tensor(f!)? swap_device_physical_rows=None," - " Tensor(g!)? swap_counts=None," + " Tensor(d!)? resolved_global_indices=None," + " Tensor(e!)? valid_counts=None," + " Tensor(f!)? swap_host_physical_rows=None," + " Tensor(g!)? swap_device_physical_rows=None," + " Tensor(h!)? swap_counts=None," " Tensor? resident_block_table=None," " int resident_block_size=0," " int resident_null_block=0) -> ()"); diff --git a/tests/v1/kv_connector/unit/test_hisparse_stats.py b/tests/v1/kv_connector/unit/test_hisparse_stats.py new file mode 100644 index 000000000000..09904d0a3caf --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hisparse_stats.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( + HiSparseKVConnectorStats, +) + + +def test_is_empty_on_fresh_stats(): + stats = HiSparseKVConnectorStats() + assert stats.is_empty() + + +def test_record_snapshot_and_reduce(): + stats = HiSparseKVConnectorStats() + stats.record_snapshot(hits=7, misses=3, host_to_device_bytes=48) + stats.record_snapshot(hits=5, misses=1, host_to_device_bytes=16) + assert not stats.is_empty() + + reduced = stats.reduce() + assert reduced["HiSparse hot-buffer hits"] == 12 + assert reduced["HiSparse hot-buffer misses"] == 4 + assert reduced["HiSparse host-to-device bytes"] == 64 + + +def test_aggregate_extends_snapshot_deltas(): + first = HiSparseKVConnectorStats() + first.record_snapshot(hits=7, misses=3, host_to_device_bytes=48) + second = HiSparseKVConnectorStats() + second.record_snapshot(hits=5, misses=1, host_to_device_bytes=16) + + first.aggregate(second) + + assert first.data == { + "cache_hits": [7, 5], + "cache_misses": [3, 1], + "host_to_device_bytes": [48, 16], + } + + +def test_aggregate_skips_empty_stats(): + stats = HiSparseKVConnectorStats() + stats.record_snapshot(hits=2, misses=1, host_to_device_bytes=16) + + stats.aggregate(HiSparseKVConnectorStats()) + + assert stats.data["cache_hits"] == [2] + + +def test_build_kv_connector_stats_round_trip(): + stats = HiSparseKVConnectorStats() + stats.record_snapshot(hits=12, misses=4, host_to_device_bytes=64) + payload = stats.to_dict() + + rebuilt = HiSparseConnector.build_kv_connector_stats(data=payload) + + assert rebuilt is not None + assert isinstance(rebuilt, HiSparseKVConnectorStats) + assert rebuilt.reduce() == stats.reduce() diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 86c8688906e9..8422a591c358 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -48,10 +48,47 @@ def _make_hisparse_worker() -> HiSparseConnectorWorker: worker.dma_stream = None worker.shared_host_region = None worker._completed_host_copy_dst_ids = [] - worker._slot_mapping_staging = None + worker._metrics_calls = 0 + worker._metrics_event = MagicMock() + worker._metrics_pending = False + worker.leader_runtimes = [] return worker +def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeypatch): + worker = _make_hisparse_worker() + worker._metrics_calls = hisparse_worker_module._METRICS_INTERVAL - 1 + worker._metrics_pending = False + worker._metrics_event = MagicMock() + worker._metrics_event.query.return_value = True + compute_stream = MagicMock() + group = SimpleNamespace( + swap_stats=torch.tensor([12, 4], dtype=torch.uint64), + swap_stats_host=torch.empty(2, dtype=torch.uint64), + stats_row_bytes=16, + copy_stream=MagicMock(), + ) + worker.leader_runtimes = [SimpleNamespace(index_group=group)] + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: compute_stream + ) + + assert worker.get_kv_connector_stats() is None + compute_stream.wait_stream.assert_called_once_with(group.copy_stream) + group.copy_stream.wait_stream.assert_called_once_with(compute_stream) + worker._metrics_event.record.assert_called_once_with() + assert group.swap_stats.tolist() == [0, 0] + + stats = worker.get_kv_connector_stats() + assert stats is not None + assert stats.data == { + "cache_hits": [12], + "cache_misses": [4], + "host_to_device_bytes": [64], + } + + def test_hisparse_row_mirrors_follow_runner_request_order(): first = SparseKVRowMirror((1,), 10, 1) second = SparseKVRowMirror((2,), 20, 1) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 338fd5cc6e33..5acb7e3bbd87 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -17,9 +17,19 @@ KVConnectorWorkerMetadata, SupportsHMA, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( + HiSparseKVConnectorStats, + HiSparsePromMetrics, +) from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( HiSparseConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.core.sched.output import SchedulerOutput @@ -244,6 +254,33 @@ def reset_capture_state(self) -> None: assert self.connector_worker is not None self.connector_worker.reset_hot_state() + def get_kv_connector_stats(self) -> KVConnectorStats | None: + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> KVConnectorStats | None: + return ( + HiSparseKVConnectorStats(data=data) + if data is not None + else HiSparseKVConnectorStats() + ) + + @classmethod + def build_prom_metrics( + cls, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> KVConnectorPromMetrics: + return HiSparsePromMetrics( + vllm_config, metric_types, labelnames, per_engine_labelvalues + ) + def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: assert self.connector_worker is not None metadata = self._get_connector_metadata() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py new file mode 100644 index 000000000000..f2abf51fc4e9 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stats and Prometheus metrics for the HiSparse connector.""" + +from dataclasses import dataclass +from typing import Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) +from vllm.v1.metrics.utils import create_metric_per_engine + +_HISPARSE_COUNTERS: tuple[tuple[str, str], ...] = ( + ("cache_hits", "Number of HiSparse device hot-buffer hits."), + ("cache_misses", "Number of HiSparse device hot-buffer misses."), + ( + "host_to_device_bytes", + "Bytes transferred from host KV storage to HiSparse hot buffers.", + ), +) + + +@dataclass +class HiSparseKVConnectorStats(KVConnectorStats): + """Container for HiSparse hot-buffer residency metrics. + + Each list entry is the delta recorded over one device counter snapshot + interval, so a list can hold multiple snapshots per logging interval. + """ + + def __post_init__(self): + if not self.data: + # Empty container init, no data is passed in. + self.reset() + + def reset(self): + # Must be serializable + self.data: dict[str, list[int]] = { + "cache_hits": [], + "cache_misses": [], + "host_to_device_bytes": [], + } + + def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): + self.data["cache_hits"].append(hits) + self.data["cache_misses"].append(misses) + self.data["host_to_device_bytes"].append(host_to_device_bytes) + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + if not other.is_empty(): + for k, v in other.data.items(): + accumulator = self.data[k] + assert isinstance(accumulator, list) + accumulator.extend(v) + return self + + def reduce(self) -> dict[str, int | float]: + # Compute compact representative stats suitable for CLI logging. + return { + "HiSparse hot-buffer hits": sum(self.data["cache_hits"]), + "HiSparse hot-buffer misses": sum(self.data["cache_misses"]), + "HiSparse host-to-device bytes": sum(self.data["host_to_device_bytes"]), + } + + def is_empty(self) -> bool: + return all(len(values) == 0 for values in self.data.values()) + + +class HiSparsePromMetrics(KVConnectorPromMetrics): + def __init__( + self, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ): + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) + + self.hisparse_counters: dict[str, Any] = {} + for name, documentation in _HISPARSE_COUNTERS: + counter = self._counter_cls( + name=f"vllm:hisparse_{name}", + documentation=documentation, + labelnames=labelnames, + ) + self.hisparse_counters[name] = create_metric_per_engine( + counter, self.per_engine_labelvalues + ) + + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): + for name, counter in self.hisparse_counters.items(): + for value in transfer_stats_data.get(name, []): + counter[engine_idx].inc(value) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index 1f723e53443c..69721803be79 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -15,6 +15,9 @@ from vllm import _custom_ops as ops from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( + HiSparseKVConnectorStats, +) from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tp_group, @@ -68,6 +71,9 @@ def _allocate_dma_descriptors(size: int) -> _DMADescriptors: ) +_METRICS_INTERVAL = 2000 + + def _get_hisparse_cache( forward_context: dict[str, Any], layer_name: str ) -> HiSparseCacheHandle: @@ -311,6 +317,9 @@ def initialize( self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = ( deque() ) + self._metrics_calls = 0 + self._metrics_event = torch.Event() + self._metrics_pending = False self._init_dma() if self.is_host_writer: for layer_index, handle in enumerate(cache_handles): @@ -496,6 +505,36 @@ def reset_hot_state(self) -> None: for runtime in self.leader_runtimes: runtime.reset_hot_state() + def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None: + stats = None + if self._metrics_pending and self._metrics_event.query(): + stats = HiSparseKVConnectorStats() + for runtime in self.leader_runtimes: + group = runtime.index_group + hits, misses = group.swap_stats_host.tolist() + if hits or misses: + stats.record_snapshot(hits, misses, misses * group.stats_row_bytes) + self._metrics_pending = False + if stats.is_empty(): + stats = None + + self._metrics_calls += 1 + if ( + self._metrics_calls % _METRICS_INTERVAL == 0 + and not self._metrics_pending + and not torch.cuda.is_current_stream_capturing() + ): + compute_stream = current_stream() + for runtime in self.leader_runtimes: + group = runtime.index_group + compute_stream.wait_stream(group.copy_stream) + group.swap_stats_host.copy_(group.swap_stats, non_blocking=True) + group.swap_stats.zero_() + group.copy_stream.wait_stream(compute_stream) + self._metrics_event.record() + self._metrics_pending = True + return stats + def _release_completed_dma_descriptors(self) -> None: pending = self._pending_dma_descriptors while pending and pending[0][0].query(): diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 96235f1570ac..3dd8c017f5b5 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -464,6 +464,11 @@ def __init__( ) self.logical_topk_ready = logical_topk_ready self.followers: list[HiSparseRuntime] = [] + self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device) + self.swap_stats_host = torch.empty( + 2, dtype=torch.uint64, device="cpu", pin_memory=True + ) + self.stats_row_bytes = 0 class HiSparseRuntime: @@ -526,6 +531,7 @@ def __init__( else: index_group.followers.append(self) self.index_group = index_group + index_group.stats_row_bytes += row_bytes self.eager_host_mirror = config.eager_host_mirror self.resident_source_index = -1 @@ -639,6 +645,7 @@ def reset_hot_state(self) -> None: compute_stream.wait_stream(group.copy_stream) group.device_global_indices.fill_(-1) group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots)) + group.swap_stats.zero_() group.copy_stream.wait_stream(compute_stream) def invalidate_sorted_slots( @@ -741,6 +748,7 @@ def _resolve_residency( request_state_indices, self.region_stride, None, + group.swap_stats, physical_topk_indices, hot.attention_block_stride, req_id_per_token[:num_tokens].contiguous(), From 7bcca16729b71294bd1b67ecde9aca584618234d Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:20:38 -0400 Subject: [PATCH 105/774] [Bugfix] Fix DeepGEMM FP8 warmup coverage (#56452) Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/model_executor/test_deep_gemm_warmup.py | 90 +++++++++++++++++++ .../scaled_mm/BlockScaledMMLinearKernel.py | 38 ++++---- .../kernels/linear/scaled_mm/deep_gemm.py | 25 +++--- .../model_executor/warmup/deep_gemm_warmup.py | 86 ++++-------------- 4 files changed, 140 insertions(+), 99 deletions(-) create mode 100644 tests/model_executor/test_deep_gemm_warmup.py diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py new file mode 100644 index 000000000000..6f2a31cf2787 --- /dev/null +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( + DeepGemmFp8BlockScaledMMKernel, +) +from vllm.model_executor.warmup import deep_gemm_warmup + + +def _block_fp8_layer(n: int, k: int, scale_name: str = "weight_scale"): + layer = SimpleNamespace( + weight=torch.empty((n, k), dtype=torch.float8_e4m3fn), + weight_block_size=[128, 128], + deep_gemm_warmup_provider=object.__new__(DeepGemmFp8BlockScaledMMKernel), + ) + setattr(layer, scale_name, torch.empty((n // 128, k // 128), dtype=torch.float32)) + return layer + + +def _run_warmup(monkeypatch, layers) -> list[dict]: + calls: list[dict] = [] + monkeypatch.setattr( + deep_gemm_warmup, + "get_mk_alignment_for_contiguous_layout", + lambda: [128, 128], + ) + monkeypatch.setattr( + deep_gemm_warmup, + "_deepgemm_fp8_gemm_nt_warmup", + lambda **kwargs: calls.append(kwargs), + ) + model = SimpleNamespace(modules=lambda: iter(layers)) + deep_gemm_warmup.deepgemm_fp8_gemm_nt_warmup(model, max_tokens=16) + return calls + + +@pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"]) +def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None: + """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the + quantization method that owns it or the name of its scale parameter.""" + layer = _block_fp8_layer(256, 128, scale_name) + + calls = _run_warmup(monkeypatch, [layer]) + + assert len(calls) == 1 + assert calls[0]["w"] is layer.weight + assert calls[0]["ws"] is getattr(layer, scale_name) + + +def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None: + """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576), + so warmup must not require N % 128 == 0.""" + assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1 + + +def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None: + mismatched = _block_fp8_layer(256, 128) + mismatched.weight_block_size = [1, 32] + layers = [SimpleNamespace(), mismatched, _block_fp8_layer(256, 128)] + + calls = _run_warmup(monkeypatch, layers) + + assert len(calls) == 1 + assert calls[0]["w"] is layers[-1].weight + + +@pytest.mark.parametrize("is_bmm", [False, True]) +def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: + """bmm layers bypass the kernel at runtime, so they must not be warmed.""" + kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel) + kernel.is_deep_gemm_supported = False + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.empty((256, 128), dtype=torch.float8_e4m3fn), requires_grad=False + ) + layer.weight_scale = torch.nn.Parameter( + torch.empty((2, 1), dtype=torch.float32), requires_grad=False + ) + layer.weight_block_size = [128, 128] + layer.is_bmm = is_bmm + + kernel.process_weights_after_loading(layer) + + provider = getattr(layer, "deep_gemm_warmup_provider", None) + assert provider is (None if is_bmm else kernel) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index c83885ca51c4..3a053412f3cf 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -38,6 +38,22 @@ def from_layer(cls, layer: torch.nn.Module) -> Self: input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None), ) + @property + def block_scale_attr(self) -> str: + """Fp8LinearMethod registers the block scale as ``weight_scale_inv``, + compressed-tensors as ``weight_scale``.""" + return ( + self.WEIGHT_SCALE + if self.weight_scale_inv is None + else self.WEIGHT_SCALE_INV + ) + + @property + def block_scale(self) -> torch.Tensor: + scale = getattr(self, self.block_scale_attr) + assert scale is not None + return scale + class Fp8BlockScaledMMLinearKernel( MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC @@ -74,25 +90,13 @@ def _get_layer_params(self, layer: torch.nn.Module, **kwargs) -> FP8BlockParams: def process_weights_after_loading(self, layer: torch.nn.Module): params = self._get_layer_params(layer) - # Fp8LinearMethod registered weight scale - # buffer as weight_scale_inv unlike compressed tensors. - weight_scale = ( - params.weight_scale - if params.weight_scale_inv is None - else params.weight_scale_inv - ) - scale_attr_name = ( - params.WEIGHT_SCALE - if params.weight_scale_inv is None - else params.WEIGHT_SCALE_INV - ) new_weight, new_weight_scale = process_fp8_weight_block_strategy( params.weight, - weight_scale, + params.block_scale, ) replace_parameter(layer, params.WEIGHT, new_weight.data) - replace_parameter(layer, scale_attr_name, new_weight_scale.data) + replace_parameter(layer, params.block_scale_attr, new_weight_scale.data) def apply_weights( self, @@ -104,11 +108,7 @@ def apply_weights( out_dtype = self.config.out_dtype params = self._get_layer_params(layer) weight = params.weight - weight_scale = ( - params.weight_scale - if params.weight_scale_inv is None - else params.weight_scale_inv - ) + weight_scale = params.block_scale input_scale = params.input_scale scale_up = params.input_scale_ub diff --git a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py index 79458f5cda85..19f9347e311e 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py @@ -86,25 +86,28 @@ def process_weights_after_loading(self, layer): params = self._get_layer_params(layer) assert layer.weight_block_size is not None + is_bmm = getattr(layer, "is_bmm", False) + if self.is_deep_gemm_supported: - weight_scale_invs = params.weight_scale_inv - scale_attr = ( - params.WEIGHT_SCALE_INV - if weight_scale_invs is not None - else params.WEIGHT_SCALE - ) dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block( wq=params.weight, - ws=weight_scale_invs - if weight_scale_invs is not None - else params.weight_scale, + ws=params.block_scale, quant_block_shape=tuple(layer.weight_block_size), use_e8m0=self.use_deep_gemm_e8m0, - is_bmm=getattr(layer, "is_bmm", False), + is_bmm=is_bmm, bmm_batch_size=getattr(layer, "bmm_batch_size", 0), ) replace_parameter(layer, params.WEIGHT, dg_weight) - replace_parameter(layer, scale_attr, dg_weight_scale) + replace_parameter(layer, params.block_scale_attr, dg_weight_scale) + # bmm layers bypass this kernel and run through fp8_einsum. + if not is_bmm: + layer.deep_gemm_warmup_provider = self + + def get_deep_gemm_warmup_weights( + self, layer: torch.nn.Module + ) -> tuple[torch.Tensor, torch.Tensor]: + params = self._get_layer_params(layer) + return params.weight, params.block_scale def apply_block_scaled_mm( self, diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 44f2121659ea..35fd6bb8f9a9 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -22,9 +22,6 @@ from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) -from vllm.model_executor.layers.linear import LinearBase -from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod -from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineLinearMethod from vllm.tracing import instrument from vllm.utils.deep_gemm import ( fp8_gemm_nt, @@ -82,28 +79,6 @@ def _generate_optimal_warmup_m_values( return sorted(m_values) -def _extract_data_from_linear_base_module( - m: torch.nn.Module, -) -> tuple[torch.Tensor, torch.Tensor, list[int]]: - """ - Extract weights, weight scales and quantization block sizes from the given - LinearBase module. - """ - assert isinstance(m, LinearBase) - assert isinstance(m.quant_method, Fp8LinearMethod) - assert m.quant_method.block_quant - assert m.quant_method.quant_config is not None - - w = m.weight - ws = m.weight_scale_inv if hasattr(m, "weight_scale_inv") else m.weight_scale - quant_block_size = m.quant_method.quant_config.weight_block_size - - assert isinstance(w, torch.Tensor) - assert isinstance(ws, torch.Tensor) - assert quant_block_size is not None - return (w, ws, quant_block_size) - - def _extract_data_from_fused_moe_module( m_: torch.nn.Module, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: @@ -133,45 +108,19 @@ def _extract_data_from_fused_moe_module( return w13, w13_s, w2, w2_s, num_topk -def _is_deep_gemm_backed_kernel(fp8_linear: object) -> bool: - """ - Return True if the selected linear kernel dispatches to DeepGEMM, either - directly or as the fallback branch of a dynamic wrapper. - """ - if isinstance(fp8_linear, DeepGemmFp8BlockScaledMMKernel): - return True - return isinstance( - getattr(fp8_linear, "fallback", None), DeepGemmFp8BlockScaledMMKernel - ) - - -def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: +def _deep_gemm_linear_data( + module: torch.nn.Module, +) -> tuple[torch.Tensor, torch.Tensor] | None: """ - Return True if the input module/layer could be processed with DeepGEMM. + Return (weight, weight_scale) for a layer whose kernel stamped itself as + ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None. """ - - if not ( - isinstance(module, LinearBase) - and isinstance(module.quant_method, Fp8LinearMethod) - and not isinstance(module.quant_method, Mxfp8OnlineLinearMethod) - and getattr(module.quant_method, "block_quant", False) - and not getattr(module.quant_method, "use_marlin", True) - ): - return False - - fp8_linear = getattr(module.quant_method, "fp8_linear", None) - if not _is_deep_gemm_backed_kernel(fp8_linear): - return False - - block_size = get_mk_alignment_for_contiguous_layout()[0] - - w, _, block_sizes = _extract_data_from_linear_base_module(module) - return ( - block_sizes == get_mk_alignment_for_contiguous_layout() - and w.ndim == 2 - and w.shape[0] % block_size == 0 - and w.shape[1] % block_size == 0 - ) + provider = getattr(module, "deep_gemm_warmup_provider", None) + if not isinstance(provider, DeepGemmFp8BlockScaledMMKernel): + return None + if module.weight_block_size != get_mk_alignment_for_contiguous_layout(): + return None + return provider.get_deep_gemm_warmup_weights(module) def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool: @@ -361,11 +310,10 @@ def _warmup(w: torch.Tensor, w_scale: torch.Tensor): def deepgemm_fp8_gemm_nt_warmup( model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None ): - dg_modules = [m for m in model.modules() if _fp8_linear_may_use_deep_gemm(m)] - - for dgm in dg_modules: - w, ws, _ = _extract_data_from_linear_base_module(dgm) - _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar) + for m in model.modules(): + if (data := _deep_gemm_linear_data(m)) is not None: + w, ws = data + _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar) def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( @@ -392,8 +340,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: total = 0 for m in model.modules(): - if _fp8_linear_may_use_deep_gemm(m): - w, _, _ = _extract_data_from_linear_base_module(m) + if (data := _deep_gemm_linear_data(m)) is not None: + w, _ = data if w.size() not in seen_fp8_sizes: total += len(_get_fp8_gemm_nt_m_values(w, max_tokens)) seen_fp8_sizes.add(w.size()) From fd7cd4b883e7720fbbfbc0ca6aa9404bf2f8e3f1 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 12 Sep 2026 11:29:58 -0700 Subject: [PATCH 106/774] [CI] Fix Qwen3 Omni DSpark load test config (#56600) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/model_executor/test_qwen3_omni.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index e507b82d9c5e..92d527e6252c 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -354,6 +354,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): parallel_config=SimpleNamespace(tensor_parallel_size=1), attention_config=SimpleNamespace(backend=None), cache_config=SimpleNamespace(), + load_config=SimpleNamespace(), model_config=SimpleNamespace(get_vocab_size=Mock(return_value=100)), ) @@ -368,6 +369,10 @@ def fake_replace(config, **changes): "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group", return_value=SimpleNamespace(world_size=1), ), + patch( + "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", + return_value=SimpleNamespace(world_size=1), + ), patch( "vllm.compilation.backends.set_model_tag", return_value=nullcontext(), From 1ee4be4dbc576f2158913a89598ccfa9dec9391e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 12 Sep 2026 11:31:13 -0700 Subject: [PATCH 107/774] [CI] Update DeepSeek V4.1 MegaMoE routing test (#56599) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/models/test_deepseek_v4_mega_moe.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 58d5239285dd..d3d9ccabad09 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,10 +18,7 @@ ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.models.deepseek_v4_1.common.mm_preprocess import ( - IMAGE_PAD_ID, - IMAGE_SENTINEL_BASE_ID, -) +from vllm.models.deepseek_v4_1.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config @@ -78,7 +75,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") hidden_states = torch.randn(4, config.hidden_size) input_ids = ( - torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257]) + torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]) if vision else None ) @@ -100,7 +97,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( ).sqrt() bias = moe.gate.e_score_correction_bias if vision: - image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID) + image_mask = input_ids == IMAGE_SENTINEL_BASE_ID bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias) expected_ids = (scores + bias).topk(top_k, dim=-1).indices expected_weights = scores.gather(1, expected_ids) From 7ee8a6dd013819838da8012ca549d724bee7c6c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Rialland?= <36076211+TQCB@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:42:01 +0200 Subject: [PATCH 108/774] [watermarking] Dual-key gumbel-max watermarking for speculative decoding support (#56122) Signed-off-by: Raphael Rialland Signed-off-by: Simon Veitner Signed-off-by: Tomas Ruiz Co-authored-by: OpenAI Codex Co-authored-by: Simon Veitner Co-authored-by: Tomas Ruiz --- .buildkite/test_areas/lm_eval.yaml | 22 + docs/features/watermarking.md | 120 ++- tests/engine/test_arg_utils.py | 25 +- .../Qwen3-4B-FP8-DSpark-watermark-TP2.yaml | 29 + .../configs/models-watermark-dspark-tp2.txt | 1 + .../test_spec_decode_rejection_warmup.py | 121 +++ .../test_watermark_sample_warmup.py | 300 ++++++ tests/test_config.py | 183 +++- .../test_rejection_sampler_utils.py | 279 ++++++ tests/watermarking/test_gumbel.py | 135 ++- tests/watermarking/test_watermarking.py | 907 +++++++++++++++++- vllm/config/vllm.py | 71 +- vllm/config/watermarking.py | 57 +- vllm/model_executor/warmup/kernel_warmup.py | 4 + .../warmup/spec_decode_rejection_warmup.py | 80 +- .../warmup/watermark_sample_warmup.py | 155 +++ vllm/v1/watermarking/__init__.py | 13 +- vllm/v1/watermarking/factory.py | 30 +- vllm/v1/watermarking/gpu_sampler.py | 131 ++- vllm/v1/watermarking/gumbel.py | 165 +++- vllm/v1/watermarking/spec_decode.py | 153 +++ vllm/v1/watermarking/watermarker.py | 73 +- vllm/v1/worker/gpu/model_runner.py | 31 +- vllm/v1/worker/gpu/sample/watermark.py | 369 ++++++- .../gpu/spec_decode/dspark/speculator.py | 10 +- .../gpu/spec_decode/rejection_sampler.py | 27 + .../spec_decode/rejection_sampler_utils.py | 163 +++- vllm/v1/worker/gpu/spec_decode/speculator.py | 33 +- 28 files changed, 3451 insertions(+), 236 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml create mode 100644 tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt create mode 100644 tests/model_executor/test_spec_decode_rejection_warmup.py create mode 100644 tests/model_executor/test_watermark_sample_warmup.py create mode 100644 vllm/model_executor/warmup/watermark_sample_warmup.py create mode 100644 vllm/v1/watermarking/spec_decode.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 1c2e3bb54f18..f33c5825548f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -122,6 +122,28 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt +- label: ":nvidia: (H100) LM Eval Watermark Feature Combination" + key: lm-eval-dspark-watermark-2xh100 + device: h100 + num_devices: 2 + # Local TP2 runs took 84–379 seconds, including startup. + timeout_in_minutes: 15 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/speculative.py + - vllm/config/vllm.py + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/warmup/ + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt + - label: ":nvidia: (B200) LM Eval Spec Decode" key: lm-eval-spec-decode-4xb200 timeout_in_minutes: 120 diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 7c601e13aaa2..f9ceda054773 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -34,11 +34,17 @@ Deployments that require watermarking must restrict this field to trusted callers, or strip and validate it at the ingress boundary, so untrusted clients cannot opt out. -`context_width` controls how many prior output tokens seed each watermark -decision and defaults to 4. Larger values make the watermark less robust to +`context_width` controls how many prior tokens seed each watermark decision +and defaults to 4. Larger values make the watermark less robust to edits because an insertion, deletion, or substitution changes more subsequent contexts. Values above 16 are allowed but emit a warning. +`allow_target_only_watermarking` defaults to false and only has an effect when +speculative decoding is enabled. It permits speculative decoding with a +watermarking algorithm that does not support it natively, at the cost of +weaker detectability. See +[Speculative decoding](#speculative-decoding). + ## Architecture `WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the @@ -52,21 +58,108 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers remain responsible for using the tokenizer and watermark profile that match generation. +## Speculative decoding + +Watermarking requires speculative decoding to use probabilistic draft sampling, +standard rejection sampling, and an autoregressive model-based method (`dspark`, +`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`. + +A watermarking algorithm without native speculative-decoding support is +rejected before model loading. Set +`"allow_target_only_watermarking": true` to allow it: accepted draft tokens are +not watermarked, while target-side rejection recovery and bonus sampling remain +watermarked. The watermark signal is diluted in proportion to the share of +output tokens supplied by accepted drafts; rejected drafts do not dilute it +because their recovery tokens are watermarked. + +Speculative-decoding token paths do not currently support generation-side +context deduplication. The configured `deduplicate_contexts` policy is not +applied to accepted drafts, rejection-recovery tokens, or bonus tokens. + +For `dual_key_gumbel`, `alpha` has no effect under speculative decoding. The +speculative protocol selects the key for each token instead. + ## Algorithms ### Gumbel-max -Gumbel-max derives a deterministic pseudorandom value from the key, prior -generated-token context, and every candidate token, then uses the resulting -Gumbel noise for categorical sampling. See +Gumbel-max derives a deterministic pseudorandom value from the key, prior token +context, and every candidate token, then uses the resulting Gumbel noise for +categorical sampling. See [Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf). Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`) bypass watermarking and emit a warning once per worker. -Single-key Gumbel-max watermarking may increase degenerate generations. In -particular, reusing the same keyed random vector when a context repeats can -reinforce repetition loops that continue until the maximum token limit. +When a token context is repeated, generation can use ordinary sampling for that +occurrence. Reusing the repeated context leads to a bias over the sequence, as +certain token choices would be correlated. Ordinary sampling at these positions +allows for single-sequence non-distortion (see section G.3 of the +[SynthID-Text supplementary materials](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf)). +The detector independently deduplicates contexts so repeated keyed random +vectors are not treated as independent evidence, meaning that context +deduplication at generation time does not reduce the watermarking signal, unless +a user tampers with the output and removes the first occurrence while leaving an +unwatermarked repetition. + +`deduplicate_contexts` controls which history is searched for a repeated +context: + +- `"none"` disables context deduplication. +- `"single_turn"`, the default, searches the tokens generated for the current + request. This gives single-turn non-distortion. +- `"all"` also searches the prompt, which extends non-distortion across the + turns of a conversation. The first `context_width` generated tokens use + ordinary sampling, and subsequent tokens whose context already occurs in the + prompt are not watermarked. Previous turns of a session, or a prompt that + already contains the structure of the answer such as a tool result the model + extends, may therefore leave little of the answer watermarked. Given the + impact of `"all"` on both performance (sampling steps must scan over a + potentially large amount of previous context) and watermarking detectability + (we may significantly reduce the amount of watermarked tokens), you should + only use `"all"` when non-distortion across multiple turns is required. + +`deduplicate_contexts_max_history` limits the search to the most recent +positions and defaults to 8,192. Each position is compared over the +`context_width` tokens before it, so the window of tokens read is that much +longer. A smaller value reduces scanning cost but only provides the guarantee +within that window. Set it to `null` to search back to the start of the +generation for `"single_turn"`, or of the request for `"all"`; an unbounded +search costs more as the sequence grows. The setting has no effect when +`deduplicate_contexts` is `"none"`. Values below 1,024 emit a warning because a +short window can miss repetition loops whose contexts recur farther apart. + +For example, this checks prompt and completion history within the default +8,192-position window: + +```bash +vllm serve MODEL \ + --watermark-config \ + '{"algorithm":"gumbel","key":42,"deduplicate_contexts":"all"}' +``` + +### Dual-key Gumbel-max + +Dual-key Gumbel-max derives independent keys A and B from one configured master +key. During ordinary generation, each token uses key A with probability +`1 - alpha` and key B with probability `alpha`. The algorithm-specific `alpha` +parameter defaults to 0.1. Detection scores every token against both keys. + +The same two key streams support speculative decoding without changing its +acceptance rate. In this mode, the speculative protocol selects the key instead +of `alpha`: draft tokens use key A, while rejection recovery and bonus tokens +use key B. The ordinary target-to-draft probability-ratio test remains unchanged, +implementing [SynthID-Text Supplementary Algorithm +6](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf). + +Select `dual_key_gumbel` together with probabilistic drafting: + +```bash +vllm serve MODEL \ + --speculative-config \ + '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \ + --watermark-config '{"algorithm":"dual_key_gumbel","key":42,"alpha":0.1}' +``` ### SynthID-Text @@ -106,7 +199,14 @@ result = GumbelWatermarkDetector(key=42, prf="philox").detect(token_ids) print(result.p_value, result.is_watermarked) ``` -The tokenizer, algorithm, PRF, key, and context width must match generation. +The detection configuration must match the generation configuration, including +the tokenizer, PRF, watermarking algorithm, algorithm-specific watermarking +configuration, and key. In practice, this information is often unavailable +when checking a piece of text. Deployments should therefore retain the set of +candidate configurations they have served, test the text against each +candidate, and correct for multiple testing, for example with a Bonferroni +correction to the resulting p-values. + Gumbel-max detection scores repeated contexts once by default so identical PRF random vectors are not treated as independent evidence. Keep `deduplicate_contexts=True` unless the detector's calibration has been adjusted @@ -140,7 +240,7 @@ watermarked output or to modify watermarked text so it is no longer detected. ## Limitations - Watermarking is currently available only with Model Runner V2. -- Gumbel-max cannot be configured with speculative decoding. +- Not all watermarking algorithms have native speculative-decoding support. - Beam search expands candidates from model log probabilities and does not apply Gumbel-max watermarking. - Models that replace the vLLM sampler with a custom sampler cannot use diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 71b9dbbab851..0262b42d11f3 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -53,18 +53,37 @@ def test_watermark_config_cli(): "--model", "dummy", "--watermark-config", - '{"algorithm":"gumbel","key":42,"prf":"philox"}', + '{"algorithm":"dual_key_gumbel","key":42,"prf":"philox","alpha":0.25,' + '"allow_target_only_watermarking":true}', ] ) config = EngineArgs.from_cli_args(args).create_watermark_config() assert config is not None - assert config.algorithm == "gumbel" + assert config.algorithm == "dual_key_gumbel" assert config.key == 42 + assert config.alpha == 0.25 assert config.context_width == 4 + assert config.deduplicate_contexts == "single_turn" + assert config.deduplicate_contexts_max_history == 8192 assert config.prf == "philox" - assert not config.supports_speculative_decoding + assert config.allow_target_only_watermarking + + args = parser.parse_args( + [ + "--model", + "dummy", + "--watermark-config", + '{"key":42,"deduplicate_contexts":"none",' + '"deduplicate_contexts_max_history":32}', + ] + ) + config = EngineArgs.from_cli_args(args).create_watermark_config() + + assert config is not None + assert config.deduplicate_contexts == "none" + assert config.deduplicate_contexts_max_history == 32 @pytest.mark.parametrize( diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml new file mode 100644 index 000000000000..c1298bebfa62 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml @@ -0,0 +1,29 @@ +model_name: "Qwen/Qwen3-4B-FP8" +num_questions: 1319 +num_fewshot: 5 +max_tokens: 256 +use_chat_completions: false +temperature: 1.0 +seed: 42 + +# Reference accuracy from the non-watermarked DSpark regression test. +accuracy_threshold: 0.801 +tolerance: 0.02 +# Measured TP1 acceptance length: 4.115; allow a 0.2 regression. +min_acceptance_length: 3.915 + +env: + VLLM_USE_V2_MODEL_RUNNER: "1" +server_args: >- + --tensor-parallel-size 2 + --max-model-len 4096 + --gpu-memory-utilization 0.92 + --no-enable-prefix-caching + --speculative-config.method dspark + --speculative-config.model deepseek-ai/dspark_qwen3_4b_block7 + --speculative-config.num_speculative_tokens 7 + --speculative-config.draft_sample_method probabilistic + --watermark-config.algorithm dual_key_gumbel + --watermark-config.key 42 + --watermark-config.context_width 4 + --watermark-config.deduplicate_contexts none diff --git a/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt new file mode 100644 index 000000000000..127b0b47509f --- /dev/null +++ b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt @@ -0,0 +1 @@ +Qwen3-4B-FP8-DSpark-watermark-TP2.yaml diff --git a/tests/model_executor/test_spec_decode_rejection_warmup.py b/tests/model_executor/test_spec_decode_rejection_warmup.py new file mode 100644 index 000000000000..6abfb1a4d1c1 --- /dev/null +++ b/tests/model_executor/test_spec_decode_rejection_warmup.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config.watermarking import WatermarkConfig, derive_watermark_key +from vllm.model_executor.warmup.spec_decode_rejection_warmup import ( + spec_decode_rejection_warmup, +) + +MASTER_KEY = 20250910 + + +def _worker( + watermark_config: WatermarkConfig | None, + draft_sample_method: str = "greedy", +) -> SimpleNamespace: + return SimpleNamespace( + device=torch.device("cpu"), + vllm_config=SimpleNamespace( + speculative_config=SimpleNamespace( + num_speculative_tokens=2, + rejection_sample_method="standard", + draft_sample_method=draft_sample_method, + ), + model_config=SimpleNamespace( + get_vocab_size=lambda: 128, dtype=torch.bfloat16 + ), + watermark_config=watermark_config, + ), + ) + + +def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + calls: list[dict] = [] + + def fake_rejection_sample(**kwargs): + calls.append(kwargs) + return None, None + + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample", + fake_rejection_sample, + ) + return calls + + +def test_warmup_token_id_dtypes_match_the_runtime_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + spec_decode_rejection_warmup(_worker(None)) + + assert calls + for kwargs in calls: + assert kwargs["draft_sampled"].dtype == torch.int32 + assert kwargs["pos"].dtype == torch.int64 + + +def test_warmup_index_mapping_dtypes_match_the_runtime_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + spec_decode_rejection_warmup(_worker(None)) + + assert calls + for kwargs in calls: + assert kwargs["idx_mapping"].dtype == torch.int64 + assert kwargs["expanded_idx_mapping"].dtype == torch.int64 + assert kwargs["expanded_local_pos"].dtype == torch.int32 + + +def test_warmup_watermark_args_match_the_runtime_kernel_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + watermark_config = WatermarkConfig( + key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3 + ) + calls = _capture_calls(monkeypatch) + spec_decode_rejection_warmup( + _worker(watermark_config, draft_sample_method="probabilistic") + ) + + assert calls + key_b = derive_watermark_key(MASTER_KEY, b"key_b") + for kwargs in calls: + contexts = kwargs["contexts"] + assert contexts.dtype == torch.int32 + assert contexts.shape == (3, 3) and contexts.stride() == (3, 1) + assert kwargs["watermarking"].dtype == torch.bool + assert kwargs["watermark_key"] == key_b + assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31 + + +def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + monkeypatch.setattr( + "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key", + lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")), + ) + spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) + + assert calls == [] + + +def test_warmup_draft_logits_match_the_runtime_draft_sample_method( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + spec_decode_rejection_warmup(_worker(None)) + assert calls + assert all(kwargs["draft_logits"] is None for kwargs in calls) + + calls = _capture_calls(monkeypatch) + spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic")) + assert calls + assert all(kwargs["draft_logits"] is not None for kwargs in calls) diff --git a/tests/model_executor/test_watermark_sample_warmup.py b/tests/model_executor/test_watermark_sample_warmup.py new file mode 100644 index 000000000000..73ecb95ee36a --- /dev/null +++ b/tests/model_executor/test_watermark_sample_warmup.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The warmup must cover every specialization the sampler will launch. + +``_philox_gumbel_kernel`` is keyed on the Philox key halves, on the logits +dtype (model dtype when the sampler passes logits through, fp32 when logits +processing copies them) and on whether a skip mask is passed, so a missing +combination leaves a request paying the compile this module exists to remove. +""" + +import logging +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config.watermarking import WatermarkConfig, derive_watermark_key +from vllm.model_executor.warmup.watermark_sample_warmup import watermark_sample_warmup +from vllm.platforms import current_platform +from vllm.v1.watermarking.factory import create_watermarker +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler +from vllm.v1.watermarking.spec_decode import create_speculative_target_watermarker + +MASTER_KEY = 20250910 +KEY_A = derive_watermark_key(MASTER_KEY, b"key_a") +KEY_B = derive_watermark_key(MASTER_KEY, b"key_b") + + +@pytest.fixture(autouse=True) +def cuda_alike_platform(monkeypatch: pytest.MonkeyPatch) -> None: + """The warmup returns early off CUDA, so pin the platform it reads. + + Everything here runs on the fake worker's cpu device; without this the + tests below would pass only on a CUDA-alike host and silently warm nothing + anywhere else. + """ + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) + + +def _sampler( + watermark_config: WatermarkConfig, + speculative: bool, + use_fp64_gumbel: bool, +) -> GPUWatermarkSampler: + """The sampler ``GPUModelRunner`` builds for this config. + + Mirrors ``vllm/v1/worker/gpu/model_runner.py:473-486``. ``__new__`` skips + ``Sampler.__init__``, which allocates UVA-backed device buffers; the warmup + only reads the three attributes set here. + """ + watermarker = create_watermarker(watermark_config) + if speculative: + watermarker = create_speculative_target_watermarker(watermarker) + sampler = GPUWatermarkSampler.__new__(GPUWatermarkSampler) + sampler.watermarker = watermarker + sampler.deduplicate_contexts = watermark_config.deduplicate_contexts + sampler.use_fp64_gumbel = use_fp64_gumbel + return sampler + + +def _worker( + watermark_config: WatermarkConfig | None, + speculative: bool = False, + use_fp64_gumbel: bool = False, + with_sampler: bool = True, +) -> SimpleNamespace: + sampler = ( + _sampler(watermark_config, speculative, use_fp64_gumbel) + if watermark_config is not None and with_sampler + else None + ) + return SimpleNamespace( + device=torch.device("cpu"), + model_runner=SimpleNamespace(sampler=sampler), + vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + get_vocab_size=lambda: 128, + dtype=torch.bfloat16, + ), + watermark_config=watermark_config, + ), + ) + + +def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + calls: list[dict] = [] + + def fake_philox_gumbel_sample(logits, contexts, key, **kwargs): + calls.append({"logits": logits, "contexts": contexts, "key": key, **kwargs}) + return torch.zeros(logits.shape[0], dtype=torch.int64) + + monkeypatch.setattr( + "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", + fake_philox_gumbel_sample, + ) + return calls + + +def _variants(calls: list[dict]) -> set[tuple[int, torch.dtype, bool]]: + return { + (call["key"], call["logits"].dtype, call.get("skip_mask") is not None) + for call in calls + } + + +def test_no_warmup_without_a_watermark_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + watermark_sample_warmup(_worker(None)) + + assert calls == [] + + +def test_no_warmup_without_a_watermarked_sampler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A rank that does not sample launches nothing to warm. + + Only the final pipeline stage of a generative model holds a + ``GPUWatermarkSampler``; on any other rank the kernel is never launched. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker(WatermarkConfig(key=MASTER_KEY), with_sampler=False) + ) + + assert calls == [] + + +def test_no_warmup_off_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + """``GumbelWatermarker`` launches the kernel only on a CUDA-alike device. + + Elsewhere -- an XPU worker inherits ``kernel_warmup`` from the CUDA worker + -- the sampler takes the torch path, so warming would compile a kernel + nothing launches. + """ + calls = _capture_calls(monkeypatch) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: False) + watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) + + assert calls == [] + + +def test_single_key_warms_both_logits_dtypes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _capture_calls(monkeypatch) + watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) + + assert _variants(calls) == { + (MASTER_KEY, torch.bfloat16, True), + (MASTER_KEY, torch.float32, True), + } + assert len(calls) == len(_variants(calls)) + + +def test_dual_key_warms_both_derived_keys(monkeypatch: pytest.MonkeyPatch) -> None: + """``DualKeyGumbelWatermarker.sample`` launches the kernel once per key. + + The halves of a derived key exceed 2**31, so each key is typed i64/i32 by + magnitude and a wrong key compiles a specialization the engine never uses. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker(WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel")) + ) + + assert _variants(calls) == { + (KEY_A, torch.bfloat16, True), + (KEY_A, torch.float32, True), + (KEY_B, torch.bfloat16, True), + (KEY_B, torch.float32, True), + } + + +@pytest.mark.parametrize("alpha,key", [(0.0, KEY_A), (1.0, KEY_B)]) +def test_dual_key_alpha_bounds_warm_one_key( + monkeypatch: pytest.MonkeyPatch, alpha: float, key: int +) -> None: + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker( + WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel", alpha=alpha) + ) + ) + + assert _variants(calls) == { + (key, torch.bfloat16, True), + (key, torch.float32, True), + } + + +def test_speculative_config_warms_the_target_role_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Under spec decode the sampler holds ``create_target_watermarker()``. + + Rows without draft tokens sample through it, so the sampler path launches + key B alone; the draft key runs in the speculator's own sampler. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker( + WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"), + speculative=True, + ) + ) + + assert _variants(calls) == { + (KEY_B, torch.bfloat16, True), + (KEY_B, torch.float32, True), + } + + +def test_dedup_none_also_warms_the_mask_free_specialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without deduplication an all-watermarked batch passes no skip mask. + + A later batch mixing in an opted-out request passes one, so both + specializations are reachable and both must be warmed. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker(WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none")) + ) + + assert _variants(calls) == { + (MASTER_KEY, torch.bfloat16, True), + (MASTER_KEY, torch.float32, True), + (MASTER_KEY, torch.bfloat16, False), + (MASTER_KEY, torch.float32, False), + } + for call in calls: + if call.get("skip_mask") is None: + assert "temperatures" not in call + + +def test_warmup_argument_dtypes_match_the_runtime_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Contexts are int32, the sampling state int64 indices and fp32 temperatures. + + ``GPUWatermarkSampler._get_contexts`` gathers from the int32 request-state + token ids, ``expanded_idx_mapping``/``positions``/``seeds`` are int64 and + ``SamplingStates.temperature`` is fp32; warming other dtypes compiles + specializations no engine launches. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY, context_width=3))) + + assert calls + for call in calls: + num_tokens = call["logits"].shape[0] + assert call["contexts"].dtype == torch.int32 + assert call["contexts"].shape == (num_tokens, 3) + assert call["skip_mask"].dtype == torch.bool + assert call["expanded_idx_mapping"].dtype == torch.int64 + assert call["positions"].dtype == torch.int64 + assert call["seeds"].dtype == torch.int64 + assert call["temperatures"].dtype == torch.float32 + + +def test_use_fp64_follows_the_path_that_forwards_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only the mixed path forwards the sampler's ``use_fp64_gumbel``. + + ``GumbelWatermarker._sample_watermarked`` launches the mask-free variant + with the ``use_fp64=False`` default, so warming it with fp64 would compile + a specialization the engine never launches and miss the one it does. + """ + calls = _capture_calls(monkeypatch) + watermark_sample_warmup( + _worker( + WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"), + use_fp64_gumbel=True, + ) + ) + + assert calls + for call in calls: + assert call["use_fp64"] is (call.get("skip_mask") is not None) + + +def test_warmup_failure_does_not_abort_startup( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + def failing_philox_gumbel_sample(*args, **kwargs): + raise RuntimeError("no triton") + + monkeypatch.setattr( + "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", + failing_philox_gumbel_sample, + ) + with caplog.at_level(logging.WARNING): + watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) + + assert "Skipping watermark sampler warmup." in caplog.text diff --git a/tests/test_config.py b/tests/test_config.py index 282f2351221b..e6d040caf5ba 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2856,14 +2856,183 @@ def _watermarked_vllm_config() -> VllmConfig: return config -def test_gumbel_watermark_rejects_speculative_decoding(): +def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log_dedup): config = _watermarked_vllm_config() - config.speculative_config = SpeculativeConfig( + config.watermark_config = WatermarkConfig( + key=42, allow_target_only_watermarking=True + ) + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with caplog_vllm.at_level(logging.WARNING): + config._check_watermarking_unsupported() + + assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text + assert "Context deduplication is not supported" in caplog_vllm.text + + +def test_speculative_watermarking_without_context_dedup_does_not_warn( + caplog_vllm, disable_log_dedup +): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig( + algorithm="dual_key_gumbel", key=42, deduplicate_contexts="none" + ) + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with caplog_vllm.at_level(logging.WARNING): + config._check_watermarking_unsupported() + + assert "Context deduplication is not supported" not in caplog_vllm.text + + +def test_gumbel_rejects_speculative_decoding_without_target_only(): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig( + algorithm="gumbel", key=42, allow_target_only_watermarking=False + ) + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"): + config._check_watermarking_unsupported() + + +def test_dual_key_gumbel_warns_that_configured_alpha_is_unused( + caplog_vllm, disable_log_dedup +): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig( + algorithm="dual_key_gumbel", key=42, alpha=0.25 + ) + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with caplog_vllm.at_level(logging.WARNING): + config._check_watermarking_unsupported() + + assert "The configured alpha=0.25 is not used" in caplog_vllm.text + + +@pytest.mark.parametrize( + ("alpha", "speculative"), + [(0.1, True), (0.25, False)], + ids=["default-alpha", "no-specdec"], +) +def test_dual_key_gumbel_alpha_warning_is_not_emitted( + caplog_vllm, disable_log_dedup, alpha, speculative +): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig( + algorithm="dual_key_gumbel", key=42, alpha=alpha + ) + if speculative: + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with caplog_vllm.at_level(logging.WARNING): + config._check_watermarking_unsupported() + + assert "is not used" not in caplog_vllm.text + + +def test_dual_key_gumbel_requires_probabilistic_drafting(): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + config.speculative_config = SimpleNamespace( + method="mtp", + draft_sample_method="greedy", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"): + config._check_watermarking_unsupported() + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"]) +def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + config.speculative_config = SimpleNamespace( + method=method, + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, + ) + + config._check_watermarking_unsupported() + + +def test_dual_key_gumbel_supports_dspark(): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + config.speculative_config = SimpleNamespace( + method="dspark", + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=True, + ) + + config._check_watermarking_unsupported() + + +def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + config.speculative_config = SimpleNamespace( method="ngram", - num_speculative_tokens=1, + draft_sample_method="probabilistic", + rejection_sample_method="standard", + parallel_drafting=False, ) - with pytest.raises(ValueError, match="does not support speculative decoding"): + with pytest.raises(ValueError, match="autoregressive model-based"): + config._check_watermarking_unsupported() + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"rejection_sample_method": "synthetic"}, "rejection_sample_method"), + ({"rejection_sample_method": "block"}, "rejection_sample_method"), + ({"parallel_drafting": True}, "Parallel speculative drafting"), + ], +) +def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match): + config = _watermarked_vllm_config() + config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + values = { + "method": "mtp", + "draft_sample_method": "probabilistic", + "rejection_sample_method": "standard", + "parallel_drafting": False, + **overrides, + } + config.speculative_config = SimpleNamespace(**values) + + with pytest.raises(ValueError, match=match): config._check_watermarking_unsupported() @@ -2882,6 +3051,12 @@ def test_watermark_key_must_fit_in_64_bits(): WatermarkConfig(key=2**64) +@pytest.mark.parametrize("alpha", [-0.1, 1.1]) +def test_watermark_alpha_must_be_a_probability(alpha): + with pytest.raises(ValidationError): + WatermarkConfig(key=42, algorithm="dual_key_gumbel", alpha=alpha) + + def test_unknown_watermark_prf_is_rejected(): with pytest.raises(ValidationError): pydantic.TypeAdapter(WatermarkConfig).validate_python( diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 86787b53a105..a64b0eb6e940 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -6,7 +6,10 @@ import pytest import torch +from vllm.v1.watermarking import GumbelWatermarker +from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( rejection_sample, ) @@ -139,6 +142,93 @@ def _assert_distribution_match( ) +def test_watermarked_bonus_uses_target_key(): + target_logits = torch.randn(16, device="cuda") + inputs = _build_rejection_sample_inputs( + target_logits, + target_logits, + num_speculative_steps=1, + temperature=1.0, + num_trials=1, + ) + draft_token = inputs["draft_sampled"][1] + contexts = torch.tensor( + [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda" + ) + + sampled, num_sampled = watermarked_rejection_sample( + **inputs, + num_speculative_steps=1, + contexts=contexts, + watermarking=torch.tensor([True], device="cuda"), + watermarker=GumbelWatermarker(key=42, context_width=2), + ) + expected = GumbelWatermarker(key=42, context_width=2).sample( + target_logits.unsqueeze(0), contexts[1:], lambda _: None + ) + + assert num_sampled.item() == 2 + assert sampled[0, 1] == expected.token_ids[0] + + +def test_watermarked_recovery_uses_target_key(): + target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda") + draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda") + inputs = _build_rejection_sample_inputs( + target_logits, + draft_logits, + num_speculative_steps=1, + temperature=1.0, + num_trials=1, + ) + inputs["draft_sampled"][1] = 0 + contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda") + + sampled, num_sampled = watermarked_rejection_sample( + **inputs, + num_speculative_steps=1, + contexts=contexts, + watermarking=torch.tensor([True], device="cuda"), + watermarker=GumbelWatermarker(key=42, context_width=2), + ) + expected = GumbelWatermarker(key=42, context_width=2).sample( + target_logits.unsqueeze(0), contexts[:1], lambda _: None + ) + + assert num_sampled.item() == 1 + assert sampled[0, 0] == expected.token_ids[0] + + +def test_watermarked_recovery_supports_smaller_draft_vocabulary(): + target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0, 100.0], device="cuda") + draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda") + inputs = _build_rejection_sample_inputs( + target_logits[:4], + draft_logits, + num_speculative_steps=1, + temperature=1.0, + num_trials=1, + ) + inputs["target_logits"] = target_logits.unsqueeze(0).expand(2, -1).contiguous() + inputs["draft_sampled"][1] = 0 + contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda") + watermarker = GumbelWatermarker(key=42, context_width=2) + + sampled, num_sampled = watermarked_rejection_sample( + **inputs, + num_speculative_steps=1, + contexts=contexts, + watermarking=torch.tensor([True], device="cuda"), + watermarker=watermarker, + ) + expected = watermarker.sample( + target_logits[:4].unsqueeze(0), contexts[:1], lambda _: None + ) + + assert num_sampled.item() == 1 + assert sampled[0, 0] == expected.token_ids[0] + + @pytest.mark.parametrize( "num_speculative_steps,temperature", [ @@ -747,3 +837,192 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool): steps = torch.arange(num_speculative_steps + 1, device=device) valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1) assert torch.equal(chunked_sampled[valid], sampled[valid]) + + +def _residual_recovery_logits( + inputs: dict, + num_sampled: torch.Tensor, + vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + target_logits = inputs["target_logits"] + draft_logits = inputs["draft_logits"] + cu_num_logits = inputs["cu_num_logits"] + resample_rows = cu_num_logits[:-1].to(torch.int64) + num_sampled.to(torch.int64) - 1 + req_state_indices = inputs["expanded_idx_mapping"][resample_rows].to(torch.int64) + temperatures = inputs["temperature"][req_state_indices] + + is_bonus = resample_rows == cu_num_logits[1:] - 1 + next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1) + needs_residual = ~is_bonus & (inputs["draft_sampled"][next_rows] >= 0) + + target_rows = target_logits[resample_rows][:, :vocab_size].float() + target_rows = torch.where(target_rows.isnan(), float("-inf"), target_rows) + if draft_logits is None: + rejected = inputs["draft_sampled"][next_rows].clamp_min(0) + one_hot = torch.arange(vocab_size, device=target_rows.device).unsqueeze( + 0 + ) == rejected.unsqueeze(-1) + residual_logits = torch.where(one_hot, float("-inf"), target_rows) + else: + draft_steps = (num_sampled.to(torch.int64) - 1).clamp_max( + draft_logits.shape[1] - 1 + ) + draft_rows = draft_logits[req_state_indices, draft_steps][ + :, :vocab_size + ].float() + draft_rows = torch.where(draft_rows.isnan(), float("-inf"), draft_rows) + draft_rows = draft_rows / temperatures.unsqueeze(-1) + target_log_probs = torch.log_softmax(target_rows, dim=-1) + draft_log_probs = torch.log_softmax(draft_rows, dim=-1) + ratio = torch.exp(draft_log_probs - target_log_probs) + residual_logits = torch.where( + ratio < 1, + target_log_probs + torch.log1p(-ratio.clamp_max(1)), + float("-inf"), + ) + recovery = torch.where(needs_residual.unsqueeze(-1), residual_logits, target_rows) + return recovery, resample_rows + + +def _watermark_inputs( + num_trials: int, + num_speculative_steps: int, + context_width: int, + vocab_size: int, + one_hot_draft: bool, + temperature: float = 1.0, +) -> tuple[dict, torch.Tensor]: + device = "cuda" + target_logits_1d = torch.randn(vocab_size, device=device) / temperature + draft_logits_1d = torch.randn(vocab_size, device=device) + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + num_speculative_steps, + temperature=temperature, + num_trials=num_trials, + ) + if one_hot_draft: + inputs["draft_logits"] = None + inputs["draft_sampled"].view(num_trials, num_speculative_steps + 1)[::5, 1] = -1 + num_logits = inputs["target_logits"].shape[0] + contexts = torch.randint( + 0, vocab_size, (num_logits, context_width), dtype=torch.int64, device=device + ) + contexts[::7, 0] = -1 + return inputs, contexts + + +@pytest.mark.parametrize("one_hot_draft", [False, True]) +# Cover one full block and a partial multi-block vocabulary. +@pytest.mark.parametrize("vocab_size", [1024, 3500]) +def test_watermarked_recovery_matches_philox_gumbel_sample( + one_hot_draft: bool, vocab_size: int +): + torch.manual_seed(7) + key = 0xDA39A3EE5E6B4B0D + context_width = 4 + inputs, contexts = _watermark_inputs( + num_trials=128, + num_speculative_steps=2, + context_width=context_width, + vocab_size=vocab_size, + one_hot_draft=one_hot_draft, + ) + watermarker = GumbelWatermarker(key=key, context_width=context_width) + + sampled, num_sampled = watermarked_rejection_sample( + **inputs, + num_speculative_steps=2, + contexts=contexts, + watermarking=torch.ones(128, dtype=torch.bool, device="cuda"), + watermarker=watermarker, + ) + + recovery, resample_rows = _residual_recovery_logits(inputs, num_sampled, vocab_size) + expected = philox_gumbel_sample(recovery, contexts[resample_rows], key) + recovered = sampled[torch.arange(128, device="cuda"), num_sampled.long() - 1] + torch.testing.assert_close(recovered, expected, rtol=0, atol=0) + + is_bonus = resample_rows == inputs["cu_num_logits"][1:] - 1 + next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1) + is_placeholder = ~is_bonus & (inputs["draft_sampled"][next_rows] < 0) + assert int((~is_bonus & ~is_placeholder).sum()) > 0 + assert int(is_placeholder.sum()) > 0 + if not one_hot_draft: + # One-hot drafts almost never reach the bonus row. + assert int(is_bonus.sum()) > 0 + + +def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw(): + torch.manual_seed(11) + num_trials = 96 + vocab_size = 512 + inputs, contexts = _watermark_inputs( + num_trials=num_trials, + num_speculative_steps=3, + context_width=2, + vocab_size=vocab_size, + one_hot_draft=False, + ) + watermarking = torch.zeros(num_trials, dtype=torch.bool, device="cuda") + watermarking[::2] = True + inputs["temperature"][:8] = 0.0 + + stock_sampled, stock_num_sampled = rejection_sample( + **inputs, num_speculative_steps=3 + ) + sampled, num_sampled = watermarked_rejection_sample( + **inputs, + num_speculative_steps=3, + contexts=contexts, + watermarking=watermarking, + watermarker=GumbelWatermarker(key=99, context_width=2), + ) + + torch.testing.assert_close(num_sampled, stock_num_sampled, rtol=0, atol=0) + rows = torch.arange(num_trials, device="cuda") + positions = num_sampled.long() - 1 + disabled = ~watermarking | (inputs["temperature"] == 0) + torch.testing.assert_close( + sampled[rows[disabled], positions[disabled]], + stock_sampled[rows[disabled], positions[disabled]], + rtol=0, + atol=0, + ) + for req in range(num_trials): + n = int(num_sampled[req]) - 1 + torch.testing.assert_close( + sampled[req, :n], stock_sampled[req, :n], rtol=0, atol=0 + ) + assert int(disabled.sum()) > 0 and int((~disabled).sum()) > 0 + + +def test_watermarked_recovery_is_unchanged_by_fp64_gumbel(): + torch.manual_seed(13) + inputs, contexts = _watermark_inputs( + num_trials=64, + num_speculative_steps=2, + context_width=4, + vocab_size=2048, + one_hot_draft=False, + ) + watermarker = GumbelWatermarker(key=2**63 + 12345, context_width=4) + kwargs = dict( + num_speculative_steps=2, + contexts=contexts, + watermarking=torch.ones(64, dtype=torch.bool, device="cuda"), + watermarker=watermarker, + ) + sampled32, num_sampled32 = watermarked_rejection_sample(**inputs, **kwargs) + sampled64, num_sampled64 = watermarked_rejection_sample( + **inputs, **kwargs, use_fp64=True + ) + rows = torch.arange(64, device="cuda") + torch.testing.assert_close(num_sampled64, num_sampled32, rtol=0, atol=0) + torch.testing.assert_close( + sampled64[rows, num_sampled64.long() - 1], + sampled32[rows, num_sampled32.long() - 1], + rtol=0, + atol=0, + ) diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index e3b4125172b0..dd32bda62793 100644 --- a/tests/watermarking/test_gumbel.py +++ b/tests/watermarking/test_gumbel.py @@ -5,8 +5,15 @@ import torch from vllm.platforms import current_platform -from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker +from vllm.v1.watermarking import ( + DualKeyGumbelWatermarkDetector, + GumbelWatermarkDetector, + GumbelWatermarker, + derive_watermark_key, +) from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample def test_gamma_survival_integer_shape(): @@ -34,6 +41,67 @@ def test_detector_deduplicates_context_even_when_target_differs(): assert detection.num_scored_tokens == 3 +def test_dual_key_detector_scores_each_token_against_both_keys(): + token_ids = [1, 2, 3, 4, 5] + key_a_detector = GumbelWatermarkDetector( + key=derive_watermark_key(42, b"key_a"), + context_width=1, + deduplicate_contexts=False, + ) + key_b_detector = GumbelWatermarkDetector( + key=derive_watermark_key(42, b"key_b"), + context_width=1, + deduplicate_contexts=False, + ) + detector = DualKeyGumbelWatermarkDetector( + key=42, + context_width=1, + deduplicate_contexts=False, + alpha=0.5, + ) + + key_a = key_a_detector.detect(token_ids) + key_b = key_b_detector.detect(token_ids) + dual = detector.detect(token_ids) + + assert dual.score == pytest.approx((key_a.score + key_b.score) / 2) + assert dual.p_value == pytest.approx( + _gamma_survival_integer_shape(dual.score * 2, dual.num_scored_tokens * 2) + ) + + +@pytest.mark.parametrize("alpha", [0.2, 0.7]) +def test_dual_key_detector_weights_and_recalibrates_scores(alpha: float): + token_ids = [1, 2, 3, 4, 5] + key_a = GumbelWatermarkDetector( + key=derive_watermark_key(42, b"key_a"), deduplicate_contexts=False + ).detect(token_ids) + key_b = GumbelWatermarkDetector( + key=derive_watermark_key(42, b"key_b"), deduplicate_contexts=False + ).detect(token_ids) + dual = DualKeyGumbelWatermarkDetector( + key=42, deduplicate_contexts=False, alpha=alpha + ).detect(token_ids) + + variance = (1 - alpha) ** 2 + alpha**2 + expected_p_value = torch.special.gammaincc( + torch.tensor(dual.num_scored_tokens / variance, dtype=torch.float64), + torch.tensor(dual.score / variance, dtype=torch.float64), + ).item() + + assert dual.score == pytest.approx((1 - alpha) * key_a.score + alpha * key_b.score) + assert dual.p_value == pytest.approx(expected_p_value) + + +def test_dual_key_detector_rejects_invalid_alpha(): + with pytest.raises(ValueError, match="alpha must be between 0 and 1"): + DualKeyGumbelWatermarkDetector(key=42, alpha=1.1) + + +def test_dual_key_detector_default_alpha(): + assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2 + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" ) @@ -45,10 +113,8 @@ def test_fused_watermarker_matches_cpu(key: int, context_width: int): logits = torch.randn(32, 8193) watermarker = GumbelWatermarker(key, context_width) - expected = watermarker.sample(logits, contexts, lambda values: None).token_ids - actual = watermarker.sample( - logits.cuda(), contexts.cuda(), lambda values: None - ).token_ids.cpu() + expected = watermarker.sample(logits, contexts).token_ids + actual = watermarker.sample(logits.cuda(), contexts.cuda()).token_ids.cpu() torch.testing.assert_close(actual, expected, rtol=0, atol=0) @@ -61,7 +127,7 @@ def test_fused_watermarker_handles_nan_logits(): logits = torch.full((2, 1025), float("nan"), device="cuda") watermarker = GumbelWatermarker(key=42, context_width=4) - token_ids = watermarker.sample(logits, contexts, lambda values: None).token_ids + token_ids = watermarker.sample(logits, contexts).token_ids assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1])) @@ -76,9 +142,58 @@ def test_fused_watermarker_handles_noncontiguous_inputs(): logits = torch.randn(8, 2050, device="cuda")[:, ::2] watermarker = GumbelWatermarker(key=42, context_width=4) - expected = watermarker.sample( - logits.contiguous(), contexts.contiguous(), lambda values: None - ).token_ids - actual = watermarker.sample(logits, contexts, lambda values: None).token_ids + expected = watermarker.sample(logits.contiguous(), contexts.contiguous()).token_ids + actual = watermarker.sample(logits, contexts).token_ids + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("use_fp64", [False, True]) +def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool): + torch.manual_seed(0) + logits = torch.randn(8, 8193, device="cuda") + context_storage = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda") + contexts = context_storage[:, ::2] + repeated_mask = torch.tensor( + [False, False, False, True, False, False, True, False], device="cuda" + ) + watermarking = torch.tensor( + [True, False, True, True, False, True, True, False], device="cuda" + ) + req_indices = torch.tensor([3, 1, 7, 0, 4, 2, 6, 5], device="cuda") + temperatures = torch.ones(8, dtype=torch.float32, device="cuda") + temperatures[2] = 0 + seeds = torch.arange(8, dtype=torch.int64, device="cuda") + 1000 + positions = torch.arange(8, dtype=torch.int64, device="cuda") + 100 + watermark_mask = ( + watermarking[req_indices] & (temperatures[req_indices] != 0) & ~repeated_mask + ) + + watermarked = philox_gumbel_sample(logits, contexts, 42) + ordinary = gumbel_sample( + logits, + req_indices, + temperatures, + seeds, + positions, + apply_temperature=False, + is_drafting=False, + use_fp64=use_fp64, + ) + expected = torch.where(watermark_mask, watermarked, ordinary) + actual = philox_gumbel_sample( + logits, + contexts, + 42, + skip_mask=~watermark_mask, + expanded_idx_mapping=req_indices, + temperatures=temperatures, + seeds=seeds, + positions=positions, + use_fp64=use_fp64, + ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 0452957ffcb3..65ea9e93ad40 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -9,44 +9,98 @@ from vllm import SamplingParams from vllm.config.watermarking import WatermarkConfig -from vllm.v1.watermarking import create_watermarker +from vllm.platforms import current_platform +from vllm.v1.watermarking import ( + DualKeyGumbelWatermarker, + SupportsSpeculativeDecoding, + create_watermarker, + derive_watermark_key, +) from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler -from vllm.v1.watermarking.watermarker import WatermarkSample +from vllm.v1.watermarking.gumbel import GumbelWatermarker +from vllm.v1.watermarking.spec_decode import ( + DraftWatermarker, + _resolve_watermark_key, + create_speculative_draft_watermarker, + create_speculative_target_watermarker, + speculative_target_watermark_key, +) +from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.sample.watermark import ( + philox_gumbel_sample, + repeated_context_mask, +) +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator -@pytest.mark.parametrize("algorithm", ["gumbel"]) +class StubWatermarker(Watermarker): + context_width = 1 + + def _sample_watermarked(self, logits, contexts): + return WatermarkSample(torch.tensor([7, 7]), logits + 10) + + +def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor: + return logits.argmax(dim=-1) + + +@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) def test_watermarker_contract(algorithm: str): watermarker = create_watermarker( WatermarkConfig(algorithm=algorithm, key=42, context_width=4) ) logits = torch.zeros(2, 128) contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) - random_sample = lambda sample_logits: sample_logits.argmax(dim=-1) - first = watermarker.sample(logits, contexts, random_sample) - second = watermarker.sample(logits, contexts, random_sample) + first = watermarker.sample(logits, contexts, _argmax_sampler) + second = watermarker.sample(logits, contexts, _argmax_sampler) assert first.token_ids.shape == (2,) assert first.logits.shape == logits.shape assert torch.equal(first.token_ids, second.token_ids) -def test_gumbel_config_warns_about_degenerate_generations(monkeypatch): +@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) +@pytest.mark.parametrize( + "config_overrides", + [ + {"deduplicate_contexts": "none"}, + {"deduplicate_contexts_max_history": 1023}, + ], +) +def test_gumbel_config_warns_when_context_deduplication_is_weak( + monkeypatch, algorithm, config_overrides +): messages: list[str] = [] monkeypatch.setattr( "vllm.config.watermarking.logger.warning_once", lambda message, *, scope: messages.append(message), ) - WatermarkConfig(key=42) + WatermarkConfig(key=42, algorithm=algorithm) + WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024) + WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None) + WatermarkConfig(key=42, algorithm=algorithm, **config_overrides) assert messages == [ - "Single-key Gumbel-max watermarking may increase the frequency of " - "degenerate generations, including repetition loops." + ( + "Gumbel-max watermarking with context deduplication disabled " + "or limited to fewer than 1024 positions may increase the frequency of " + "degenerate generations, including repetition loops. Use " + "deduplicate_contexts='single_turn' or 'all' with " + "deduplicate_contexts_max_history at least 1024 or null to mitigate this." + ) ] +def test_context_deduplication_history_can_be_unbounded(): + config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None) + + assert config.deduplicate_contexts_max_history is None + + def test_large_context_width_warns_but_is_allowed(): config = WatermarkConfig(key=42, context_width=17) @@ -56,6 +110,158 @@ def test_large_context_width_warns_but_is_allowed(): assert watermarker.context_width == 17 +def test_dual_key_watermarker_uses_domain_separated_keys(): + config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + + watermarker = create_watermarker(config) + assert isinstance(watermarker, SupportsSpeculativeDecoding) + draft = watermarker.draft_watermarker + target = watermarker.target_watermarker + + assert isinstance(watermarker, DualKeyGumbelWatermarker) + assert watermarker.alpha == 0.1 + assert draft.prf.key == derive_watermark_key(42, b"key_a") + assert target.prf.key == derive_watermark_key(42, b"key_b") + assert target.prf.key != draft.prf.key + + +def test_dual_key_watermarker_routes_tokens_with_alpha(): + watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25) + logits = torch.zeros(2, 16) + contexts = torch.tensor([[1, 2], [3, 4]]) + key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None) + key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None) + + def route(routing_logits): + torch.testing.assert_close( + routing_logits.softmax(dim=-1), + torch.tensor([[0.75, 0.25], [0.75, 0.25]]), + ) + return torch.tensor([0, 1]) + + sampled = watermarker.sample(logits, contexts, route) + + assert torch.equal( + sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]]) + ) + + +def test_dual_key_routing_logits_are_cached(): + alpha = 0.25 + watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha) + logits = torch.zeros(2, 16) + contexts = torch.tensor([[1, 2], [3, 4]]) + recorded: list[torch.Tensor] = [] + + def route(routing_logits): + recorded.append(routing_logits) + return torch.tensor([0, 1]) + + watermarker.sample(logits, contexts, route) + watermarker.sample(logits, contexts, route) + + expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log() + expected = expected.expand(2, -1) + assert len(recorded) == 2 + assert all(torch.equal(seen, expected) for seen in recorded) + assert len(watermarker._routing_logits_cache) == 1 + + +@pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0]) +def test_dual_key_watermarker_skips_masked_rows(alpha: float): + watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha) + torch.manual_seed(0) + logits = torch.randn(4, 32) + contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]]) + skip_mask = torch.tensor([True, False, True, False]) + + unmasked = watermarker.sample(logits, contexts, _argmax_sampler) + masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask) + + ordinary = _argmax_sampler(logits) + assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask]) + assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask]) + assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask]) + + +def test_dual_key_watermarker_routing_requires_a_random_sampler(): + watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25) + + with pytest.raises(ValueError, match="random sampler"): + watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]])) + + +def test_speculative_decoding_uses_fixed_dual_key_roles(): + watermarker = create_watermarker( + WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25) + ) + + target = create_speculative_target_watermarker(watermarker) + draft = create_speculative_draft_watermarker( + watermarker, + max_num_reqs=1, + device=torch.device("cpu"), + allow_target_only=False, + ) + + assert target.prf.key == derive_watermark_key(42, b"key_b") + assert draft is not None + assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a") + + +def test_recovery_key_rejects_the_unsplit_dual_key_watermarker(): + watermarker = create_watermarker( + WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25) + ) + + with pytest.raises(ValueError, match="keys the target role separately"): + _resolve_watermark_key(watermarker) + + target = create_speculative_target_watermarker(watermarker) + assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b") + + +@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) +def test_config_key_resolution_matches_the_model_runner(algorithm: str): + config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25) + runtime_key = _resolve_watermark_key( + create_speculative_target_watermarker(create_watermarker(config)) + ) + + assert speculative_target_watermark_key(config) == runtime_key + + +def test_config_key_resolution_returns_none_when_watermarking_is_disabled(): + assert speculative_target_watermark_key(None) is None + + +def test_target_only_speculative_watermarking_skips_draft_watermarker(): + watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42)) + + assert not isinstance(watermarker, SupportsSpeculativeDecoding) + with pytest.raises(ValueError, match="does not support speculative decoding"): + create_speculative_draft_watermarker( + watermarker, + max_num_reqs=1, + device=torch.device("cpu"), + allow_target_only=False, + ) + assert ( + create_speculative_draft_watermarker( + watermarker, + max_num_reqs=1, + device=torch.device("cpu"), + allow_target_only=True, + ) + is None + ) + + +def test_dual_key_derivation_is_stable(): + assert derive_watermark_key(32, b"key_a") == 16368605726115524094 + assert derive_watermark_key(32, b"key_b") == 4799302812959726346 + + def test_sampling_params_can_disable_watermarking(): assert SamplingParams().watermarking assert not SamplingParams.from_optional(watermarking=False).watermarking @@ -83,14 +289,9 @@ def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch): def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch): - class StubWatermarker: - context_width = 1 - - def sample(self, logits, contexts, random_sample): - return WatermarkSample(torch.tensor([7, 7]), logits + 10) - sampler = object.__new__(GPUWatermarkSampler) sampler.watermarker = StubWatermarker() + sampler.deduplicate_contexts = "single_turn" sampler.watermarking = SimpleNamespace( np=np.array([True, False]), gpu=torch.tensor([True, False]) ) @@ -102,8 +303,11 @@ def sample(self, logits, contexts, random_sample): sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( 2, 1, dtype=torch.int64 ) + sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros( + 2, dtype=torch.bool + ) monkeypatch.setattr( - "vllm.v1.watermarking.gpu_sampler.gumbel_sample", + "vllm.v1.watermarking.watermarker.gumbel_sample", lambda *args, **kwargs: torch.tensor([3, 4]), ) logits = torch.zeros(2, 8) @@ -123,11 +327,524 @@ def sample(self, logits, contexts, random_sample): assert torch.equal(output_logits[1], logits[1]) +def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = StubWatermarker() + sampler.deduplicate_contexts = "single_turn" + sampler.watermarking = SimpleNamespace( + np=np.array([True, True]), gpu=torch.tensor([True, True]) + ) + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), + seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), + ) + sampler.use_fp64_gumbel = False + sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( + 2, 1, dtype=torch.int64 + ) + sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: ( + torch.tensor([True, False]) + ) + monkeypatch.setattr( + "vllm.v1.watermarking.watermarker.gumbel_sample", + lambda *args, **kwargs: torch.tensor([3, 4]), + ) + logits = torch.zeros(2, 8) + + sampled, output_logits = sampler._sample_random( + logits, + torch.tensor([0, 1]), + np.array([0, 1]), + torch.zeros(2, dtype=torch.int64), + None, + None, + False, + ) + + assert torch.equal(sampled, torch.tensor([3, 7])) + assert torch.equal(output_logits[0], logits[0]) + assert torch.equal(output_logits[1], torch.full((8,), 10.0)) + + +def test_gpu_sampler_can_disable_context_deduplication(monkeypatch): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = StubWatermarker() + sampler.deduplicate_contexts = "none" + sampler.watermarking = SimpleNamespace( + np=np.array([True, True]), gpu=torch.tensor([True, True]) + ) + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), + seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), + ) + sampler.use_fp64_gumbel = False + sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( + 2, 1, dtype=torch.int64 + ) + sampler._get_repeated_contexts = lambda *args: pytest.fail( + "context deduplication should not run" + ) + logits = torch.zeros(2, 8) + + sampled, output_logits = sampler._sample_random( + logits, + torch.tensor([0, 1]), + np.array([0, 1]), + torch.zeros(2, dtype=torch.int64), + None, + None, + False, + ) + + assert torch.equal(sampled, torch.tensor([7, 7])) + assert torch.equal(output_logits, torch.full((2, 8), 10.0)) + + +def test_repeated_context_mask_ignores_prompt_tokens(): + all_token_ids = torch.tensor( + [ + [8, 9, 1, 2, 1, 2], + [3, 4, 1, 2, 3, 4], + [0, 0, 0, 0, 0, 0], + ], + dtype=torch.int32, + ) + req_indices = torch.tensor([0, 1, -1]) + prompt_lens = torch.tensor([2, 2, 0]) + total_lens = torch.tensor([6, 6, 0]) + contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]]) + + repeated = repeated_context_mask( + all_token_ids, + req_indices, + prompt_lens, + total_lens, + contexts, + ) + + assert torch.equal(repeated, torch.tensor([True, False, False])) + + +def test_repeated_context_mask_can_include_prompt_tokens(): + all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) + req_indices = torch.tensor([0]) + prompt_lens = torch.tensor([4]) + total_lens = torch.tensor([10]) + contexts = torch.tensor([[1, 2, 3, 4]]) + + single_turn = repeated_context_mask( + all_token_ids, req_indices, prompt_lens, total_lens, contexts + ) + all_history = repeated_context_mask( + all_token_ids, + req_indices, + prompt_lens, + total_lens, + contexts, + include_prompt=True, + ) + + assert not single_turn.item() + assert all_history.item() + + +def test_repeated_context_mask_can_skip_partial_contexts(): + all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32) + req_indices = torch.tensor([0]) + prompt_lens = torch.tensor([4]) + contexts = torch.tensor([[-1, -1, -1, 1]]) + + partial = repeated_context_mask( + all_token_ids, + req_indices, + prompt_lens, + torch.tensor([5]), + contexts, + include_prompt=True, + skip_partial_context=True, + ) + complete = repeated_context_mask( + all_token_ids, + req_indices, + prompt_lens, + torch.tensor([8]), + torch.tensor([[1, 2, 3, 4]]), + include_prompt=True, + skip_partial_context=True, + ) + + assert partial.item() + assert not complete.item() + + +def test_gpu_sampler_contexts_are_completion_local_for_all_scopes(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = SimpleNamespace(context_width=4) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])), + prompt_len=SimpleNamespace(gpu=torch.tensor([4])), + total_len=SimpleNamespace(gpu=torch.tensor([6])), + ) + request_indices = torch.tensor([0]) + + sampler.deduplicate_contexts = "single_turn" + single_turn = sampler._get_contexts(request_indices) + sampler.deduplicate_contexts = "all" + all_history = sampler._get_contexts(request_indices) + + assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]])) + assert torch.equal(all_history, single_turn) + + +def test_gpu_sampler_all_history_skips_partial_context(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = SimpleNamespace(context_width=4) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])), + prompt_len=SimpleNamespace(gpu=torch.tensor([4])), + total_len=SimpleNamespace(gpu=torch.tensor([6])), + ) + sampler.deduplicate_contexts_max_history = None + request_indices = torch.tensor([0]) + contexts = sampler._get_contexts(request_indices) + + sampler.deduplicate_contexts = "single_turn" + single_turn = sampler._get_repeated_contexts(request_indices, contexts) + sampler.deduplicate_contexts = "all" + all_history = sampler._get_repeated_contexts(request_indices, contexts) + + assert not single_turn.item() + assert all_history.item() + + +def test_repeated_context_mask_respects_max_history(): + all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) + req_indices = torch.tensor([0]) + prompt_lens = torch.tensor([0]) + total_lens = torch.tensor([10]) + contexts = torch.tensor([[1, 2, 3, 4]]) + + full_history = repeated_context_mask( + all_token_ids, req_indices, prompt_lens, total_lens, contexts + ) + last_six = repeated_context_mask( + all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6 + ) + last_five = repeated_context_mask( + all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5 + ) + + assert full_history.item() + assert last_six.item() + assert not last_five.item() + + +@pytest.mark.parametrize("max_history", [0, -1]) +def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int): + with pytest.raises(ValueError, match="max_history must be positive or None"): + repeated_context_mask( + torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([4], dtype=torch.int32), + torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), + max_history=max_history, + ) + + +def test_gpu_sampler_wires_request_history_scope(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace( + gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) + ), + prompt_len=SimpleNamespace(gpu=torch.tensor([4])), + total_len=SimpleNamespace(gpu=torch.tensor([10])), + ) + sampler.deduplicate_contexts_max_history = None + request_indices = torch.tensor([0]) + contexts = torch.tensor([[1, 2, 3, 4]]) + + sampler.deduplicate_contexts = "single_turn" + single_turn = sampler._get_repeated_contexts(request_indices, contexts) + sampler.deduplicate_contexts = "all" + all_history = sampler._get_repeated_contexts(request_indices, contexts) + + assert not single_turn.item() + assert all_history.item() + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("max_history", [5, 6]) +@pytest.mark.parametrize("include_prompt", [False, True]) +def test_repeated_context_mask_max_history_accelerator_parity( + max_history: int, include_prompt: bool +): + inputs = ( + torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]), + torch.tensor([0]), + torch.tensor([0]), + torch.tensor([10]), + torch.tensor([[1, 2, 3, 4]]), + ) + + expected = repeated_context_mask( + *inputs, max_history=max_history, include_prompt=include_prompt + ) + actual = repeated_context_mask( + *(value.cuda() for value in inputs), + max_history=max_history, + include_prompt=include_prompt, + ).cpu() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_repeated_context_mask_partial_context_accelerator_parity(): + inputs = ( + torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64), + ) + + expected = repeated_context_mask( + *inputs, include_prompt=True, skip_partial_context=True + ) + actual = repeated_context_mask( + *(value.cuda() for value in inputs), + include_prompt=True, + skip_partial_context=True, + ).cpu() + + assert torch.equal(actual, expected) + + +def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]: + """One row whose only repeated context ends at position 680. + + 680 is not a multiple of the kernel's 512-position block, so the scan window + selected by ``max_history`` starts in the middle of a block. + """ + history = list(range(1_000, 2_200)) + history[676:680] = [1, 2, 3, 4] + return ( + torch.tensor([history], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([len(history)], dtype=torch.int32), + torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize( + "max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)] +) +def test_repeated_context_mask_unaligned_max_history_accelerator_parity( + max_history: int, expected: bool +): + inputs = _unaligned_max_history_inputs() + + reference = repeated_context_mask(*inputs, max_history=max_history) + actual = repeated_context_mask( + *(value.cuda() for value in inputs), max_history=max_history + ).cpu() + + assert torch.equal(reference, torch.tensor([expected])) + assert torch.equal(actual, reference) + + +def _repeated_context_inputs( + context_width: int, device: str = "cpu" +) -> tuple[torch.Tensor, ...]: + """Build repeated, unique, and padded request rows for mask tests.""" + prompt = [101, 102, 103, 104] + repeated_output = [*range(1, context_width + 1)] * 2 + unique_output = [*range(1, context_width + 2)] + max_len = len(prompt) + len(repeated_output) + rows = [ + prompt + repeated_output, + prompt + unique_output, + [], + ] + all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device) + for row, token_ids in enumerate(rows): + all_token_ids[row, : len(token_ids)] = torch.tensor( + token_ids, dtype=torch.int32, device=device + ) + return ( + all_token_ids, + torch.tensor([0, 1, -1], dtype=torch.int32, device=device), + torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device), + torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device), + torch.tensor( + [ + repeated_output[-context_width:], + unique_output[-context_width:], + [-1] * context_width, + ], + dtype=torch.int64, + device=device, + ), + ) + + +@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17]) +def test_repeated_context_mask(context_width: int): + repeated = repeated_context_mask(*_repeated_context_inputs(context_width)) + + assert torch.equal(repeated, torch.tensor([True, False, False])) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17]) +def test_repeated_context_mask_accelerator_parity(context_width: int): + cpu_inputs = _repeated_context_inputs(context_width) + accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda")) + contexts = accelerator_inputs[-1] + storage = torch.empty( + contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda" + ) + storage[:, ::2] = contexts + accelerator_inputs[-1] = storage[:, ::2] + + expected = repeated_context_mask(*cpu_inputs) + actual = repeated_context_mask(*accelerator_inputs).cpu() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_repeated_context_mask_scans_multiple_blocks(): + context = [1, 2, 3, 4] + repeated_output = [*range(10_000, 11_030), *context, 99, *context] + unique_output = list(range(20_000, 21_039)) + prompt = [101, 102, 103, 104] + rows = [prompt + repeated_output, prompt + unique_output] + all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32) + for row, token_ids in enumerate(rows): + all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids) + inputs = ( + all_token_ids, + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([len(prompt), len(prompt)], dtype=torch.int32), + torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32), + torch.tensor([context, unique_output[-4:]], dtype=torch.int64), + ) + + expected = repeated_context_mask(*inputs) + actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu() + + assert torch.equal(expected, torch.tensor([True, False])) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_repeated_context_mask_last_request_at_capacity(): + max_model_len = 1024 + all_token_ids = torch.arange( + 2 * max_model_len, dtype=torch.int32, device="cuda" + ).reshape(2, max_model_len) + inputs = ( + all_token_ids, + torch.tensor([1], dtype=torch.int32, device="cuda"), + torch.tensor([0, 0], dtype=torch.int32, device="cuda"), + torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"), + all_token_ids[1, -4:].to(torch.int64).unsqueeze(0), + ) + + repeated = repeated_context_mask(*inputs) + + assert not repeated.item() + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts(): + device = torch.device("cuda") + logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device) + all_token_ids = torch.tensor( + [[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]], + dtype=torch.int32, + device=device, + ) + request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device) + request_indices_np = np.array([0, 1]) + temperatures = torch.ones(2, device=device) + seeds = torch.tensor([11, 22], dtype=torch.int64, device=device) + positions = torch.tensor([6, 6], dtype=torch.int64, device=device) + + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = GumbelWatermarker(key=42, context_width=1) + sampler.deduplicate_contexts = "single_turn" + sampler.deduplicate_contexts_max_history = None + sampler.watermarking = SimpleNamespace( + np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device) + ) + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures), + seeds=SimpleNamespace(gpu=seeds), + ) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=all_token_ids), + prompt_len=SimpleNamespace( + gpu=torch.tensor([2, 2], dtype=torch.int32, device=device) + ), + total_len=SimpleNamespace( + gpu=torch.tensor([6, 6], dtype=torch.int32, device=device) + ), + ) + sampler.use_fp64_gumbel = False + + contexts = sampler._get_contexts(request_indices) + repeated = sampler._get_repeated_contexts(request_indices, contexts) + expected = philox_gumbel_sample( + logits, + contexts, + 42, + skip_mask=repeated, + expanded_idx_mapping=request_indices, + temperatures=temperatures, + seeds=seeds, + positions=positions, + ) + + actual, output_logits = sampler._sample_random( + logits, + request_indices, + request_indices_np, + positions, + None, + None, + False, + ) + + assert torch.equal(repeated, torch.tensor([True, False], device=device)) + assert torch.equal(actual, expected) + assert output_logits is logits + + def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch): class StubWatermarker: context_width = 1 - def sample(self, logits, contexts, random_sample): + def sample(self, logits, contexts, random_sampler=None, skip_mask=None): raise AssertionError("watermarker should not run for greedy requests") sampler = object.__new__(GPUWatermarkSampler) @@ -155,3 +872,157 @@ def sample(self, logits, contexts, random_sample): ) assert actual is expected + + +def test_gpu_sampler_builds_speculative_contexts_from_drafts(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = SimpleNamespace(context_width=2) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])), + prompt_len=SimpleNamespace(gpu=torch.tensor([2])), + total_len=SimpleNamespace(gpu=torch.tensor([4])), + ) + + contexts = sampler._get_contexts( + torch.tensor([0, 0, 0]), + torch.tensor([0, 1, 2]), + torch.tensor([11, 20, 21]), + ) + + assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]])) + + +def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = SimpleNamespace(context_width=2) + sampler.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace( + gpu=torch.tensor( + [ + [100, 101, 10, 11, 0, 0], + [200, 201, 30, 31, 0, 0], + [300, 301, 50, 51, 0, 0], + ] + ) + ), + prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])), + total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])), + ) + + contexts = sampler._get_contexts( + torch.tensor([2, 2, 1, 1]), + torch.tensor([0, 1, 0, 1]), + torch.tensor([51, 60, 31, 40]), + ) + + assert torch.equal( + contexts, + torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]), + ) + + +def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): + class StubSpeculator(DraftModelSpeculator): + def capture(self): ... + + def init_cudagraph_manager(self, cudagraph_mode): ... + + def load_draft_model(self, target_model, target_attn_layer_names): ... + + def propose(self, *args, **kwargs): ... + + class StubModel: + @staticmethod + def compute_logits(hidden_states): + return torch.zeros(hidden_states.shape[0], 8) + + class StubWatermarker: + @staticmethod + def sample(logits, contexts, random_sampler=None): + assert random_sampler is None + return WatermarkSample(torch.tensor([7, 7]), logits) + + speculator = object.__new__(StubSpeculator) + speculator.model = StubModel() + speculator.use_fp64_gumbel = False + draft_watermarker = object.__new__(DraftWatermarker) + draft_watermarker.watermarker = StubWatermarker() + draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]]) + draft_watermarker.enabled = torch.tensor([True, False]) + speculator.draft_watermarker = draft_watermarker + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample", + lambda *args, **kwargs: torch.tensor([3, 4]), + ) + + sampled = speculator.sample_draft( + hidden_states=torch.zeros(2, 4), + sample_src_positions=torch.zeros(2, dtype=torch.int64), + idx_mapping=torch.tensor([0, 1]), + temperature=torch.ones(2), + seeds=torch.zeros(2, dtype=torch.int64), + draft_step=torch.tensor(0), + draft_logits=torch.zeros(2, 1, 8), + ) + + assert torch.equal(sampled, torch.tensor([7, 4])) + assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]])) + + +def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): + speculator = object.__new__(DSparkSpeculator) + speculator.draft_logits = torch.zeros(2, 1, 8) + speculator._d2t_scatter_index = torch.tensor([1, 5]) + speculator._draft_scatter_buf = torch.full((2, 8), float("-inf")) + speculator.temperature = torch.ones(2) + speculator.seeds = torch.zeros(2, dtype=torch.int64) + speculator._step_cols = torch.tensor([0]) + speculator.use_fp64_gumbel = False + watermark_logits: list[torch.Tensor] = [] + + def sample(logits, sampled, idx_map, temperature): + watermark_logits.append(logits.clone()) + return sampled + 1 + + speculator.draft_watermarker = SimpleNamespace(sample=sample) + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample", + lambda *args, **kwargs: torch.tensor([3, 4]), + ) + + sampled = speculator._sample_logits( + torch.tensor([[10.0, 20.0], [30.0, 40.0]]), + torch.tensor([0, 1]), + torch.tensor([1, 1]), + 0, + ) + + assert torch.equal(sampled, torch.tensor([4, 5])) + assert torch.equal( + watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]]) + ) + assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all() + + +def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch): + speculator = object.__new__(DSparkSpeculator) + speculator.draft_logits = torch.zeros(2, 1, 8) + speculator._d2t_scatter_index = None + speculator.temperature = torch.ones(2) + speculator.seeds = torch.zeros(2, dtype=torch.int64) + speculator._step_cols = torch.tensor([0]) + speculator.use_fp64_gumbel = False + speculator.draft_watermarker = None + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample", + lambda *args, **kwargs: torch.tensor([3, 4]), + ) + + sampled = speculator._sample_logits( + torch.zeros(2, 8), + torch.tensor([0, 1]), + torch.tensor([1, 1]), + 0, + ) + + assert torch.equal(sampled, torch.tensor([3, 4])) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3c197df0daff..3bdfe1f3edb5 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -50,7 +50,7 @@ from .scheduler import SchedulerConfig from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig from .structured_outputs import StructuredOutputsConfig -from .utils import SupportsHash, config, replace +from .utils import SupportsHash, config, get_field, replace from .watermarking import WatermarkConfig from .weight_transfer import WeightTransferConfig @@ -1142,14 +1142,67 @@ def _check_watermarking_unsupported( watermark_config = getattr(self, "watermark_config", None) if watermark_config is None: return - if ( - self.speculative_config is not None - and not watermark_config.supports_speculative_decoding - ): - raise ValueError( - f"The {watermark_config.algorithm} watermarking algorithm " - "does not support speculative decoding." - ) + if self.speculative_config is not None: + speculative_config = self.speculative_config + if speculative_config.draft_sample_method != "probabilistic": + raise ValueError( + "Speculative decoding with watermarking requires " + "draft_sample_method='probabilistic'." + ) + if speculative_config.rejection_sample_method != "standard": + raise ValueError( + "Speculative decoding with watermarking requires " + "rejection_sample_method='standard'." + ) + if ( + speculative_config.parallel_drafting + and speculative_config.method != "dspark" + ): + raise ValueError( + "Parallel speculative drafting is not supported with watermarking." + ) + if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"): + raise ValueError( + "Watermarking supports only autoregressive model-based " + "speculative decoding." + ) + if ( + not watermark_config.allow_target_only_watermarking + and not watermark_config.supports_speculative_decoding + ): + raise ValueError( + f"The '{watermark_config.algorithm}' watermarking algorithm " + "does not support speculative decoding. Set " + "allow_target_only_watermarking=true to leave draft tokens " + "unwatermarked." + ) + if ( + watermark_config.allow_target_only_watermarking + and not watermark_config.supports_speculative_decoding + ): + logger.warning_once( + "Target-only watermarking leaves accepted draft tokens " + "unwatermarked, weakening detectability in proportion to the " + "share of output tokens supplied by accepted drafts.", + scope="global", + ) + if watermark_config.deduplicate_contexts != "none": + logger.warning_once( + "Context deduplication is not supported with speculative " + "decoding and will not be applied to accepted drafts, " + "rejection-recovery tokens, or bonus tokens.", + scope="global", + ) + if watermark_config.algorithm == "dual_key_gumbel" and ( + watermark_config.alpha != get_field(WatermarkConfig, "alpha").default + ): + logger.warning_once( + "Speculative decoding selects the watermark key by token role: " + "draft tokens use key A, recovery and bonus tokens use key B. " + "The configured alpha=%s is not used.", + watermark_config.alpha, + scope="global", + ) if beam_search: raise ValueError("Beam search is not supported with watermarking.") if custom_sampler: diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py index 2247dd8248ea..8701c7b72e4d 100644 --- a/vllm/config/watermarking.py +++ b/vllm/config/watermarking.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib from typing import Literal from pydantic import Field, model_validator @@ -11,12 +12,16 @@ logger = init_logger(__name__) -WatermarkingAlgorithm = Literal["gumbel"] +WatermarkingAlgorithm = Literal["gumbel", "dual_key_gumbel"] WatermarkPRFName = Literal["philox"] +WatermarkContextScope = Literal["none", "single_turn", "all"] -_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = { - "gumbel": False, -} +_MIN_RECOMMENDED_DEDUP_HISTORY = 1024 + + +def derive_watermark_key(key: int, domain: bytes) -> int: + digest = hashlib.sha256(domain + key.to_bytes(8, "big")).digest() + return int.from_bytes(digest[:8], "big") @config @@ -27,23 +32,49 @@ class WatermarkConfig: """Secret key used to watermark generated text.""" algorithm: WatermarkingAlgorithm = "gumbel" """Algorithm used to watermark generated text.""" + alpha: float = Field(default=0.1, ge=0, le=1) + """Probability of selecting key B for dual-key watermarking.""" context_width: int = Field(default=4, ge=1) - """Number of prior output tokens used by the watermark PRF.""" + """Number of prior tokens used by the watermark PRF.""" + deduplicate_contexts: WatermarkContextScope = "single_turn" + """Which history is searched for a repeated context before a token is + sampled; a repeated context is sampled without the watermark. `none` + disables the search. `single_turn` (default) searches this request's + generated tokens. `all` also searches the prompt and samples the first + `context_width` generated tokens without watermarking.""" + deduplicate_contexts_max_history: int | None = Field(default=8192, ge=1) + """Number of most recent history positions searched (default 8192), or + `None` for the whole scope. Each position is compared over the + `context_width` tokens before it. Ignored when `deduplicate_contexts` is + `none`.""" prf: WatermarkPRFName = "philox" """Pseudorandom function used by the watermarking algorithm.""" + allow_target_only_watermarking: bool = False + """Allow speculative decoding without watermarking draft tokens.""" + + @property + def supports_speculative_decoding(self) -> bool: + return self.algorithm == "dual_key_gumbel" @model_validator(mode="after") - def validate_key(self) -> Self: + def validate_watermark_settings(self) -> Self: if self.key > 2**64 - 1: raise ValueError("philox keys must fit in 64 bits") - if self.algorithm == "gumbel": + history_is_too_short = ( + self.deduplicate_contexts_max_history is not None + and self.deduplicate_contexts_max_history < _MIN_RECOMMENDED_DEDUP_HISTORY + ) + if self.algorithm in ("gumbel", "dual_key_gumbel") and ( + self.deduplicate_contexts == "none" or history_is_too_short + ): logger.warning_once( - "Single-key Gumbel-max watermarking may increase the frequency of " - "degenerate generations, including repetition loops.", + "Gumbel-max watermarking with context deduplication " + "disabled or limited to fewer than " + f"{_MIN_RECOMMENDED_DEDUP_HISTORY} positions may increase the " + "frequency of degenerate generations, including repetition loops. " + "Use deduplicate_contexts='single_turn' or 'all' with " + "deduplicate_contexts_max_history at least " + f"{_MIN_RECOMMENDED_DEDUP_HISTORY} or null to mitigate this.", scope="global", ) return self - - @property - def supports_speculative_decoding(self) -> bool: - return _SPECULATIVE_DECODING_SUPPORT[self.algorithm] diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 2991dc1027a4..bbb0a66d538c 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -44,6 +44,9 @@ from vllm.model_executor.warmup.spec_decode_rejection_warmup import ( spec_decode_rejection_warmup, ) +from vllm.model_executor.warmup.watermark_sample_warmup import ( + watermark_sample_warmup, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -207,6 +210,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): if enable_jit_warmup: kimi_k3_triton_warmup(worker) spec_decode_rejection_warmup(worker) + watermark_sample_warmup(worker) qwen4_exp_qsa_triton_warmup(worker) if enable_jit_warmup and current_platform.is_device_capability_family(100): diff --git a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py index 8cedbfc4db42..6f6f08a2d58f 100644 --- a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py +++ b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py @@ -1,17 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up spec-decode rejection-sampler Triton kernels. - -The rejection sampler kernels (``_compute_local_logits_stats_kernel``, -``_rejection_kernel``, ``_resample_kernel``) are JIT-compiled by Triton on -first use. Without warmup, the first spec-decode request pays a multi-second -compilation cost. This pre-compiles them with dummy data matching the -server's vocab size and speculative config. -""" +"""Warm up speculative-decoding rejection kernels.""" from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -39,42 +32,76 @@ def spec_decode_rejection_warmup(worker: Worker) -> None: if num_spec <= 0 or vocab_size <= 0: return - # Mirror the constexpr-relevant flags the runtime uses. rejection_method = getattr(spec_config, "rejection_sample_method", None) use_block_verification = rejection_method == "block" use_synthetic = rejection_method == "synthetic" + # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime. + use_draft_logits = ( + getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic" + ) device = worker.device num_reqs = 1 tokens_per_req = num_spec + 1 num_logits = num_reqs * tokens_per_req - # Triton JIT-specializes on tensor dtypes. The target logits may be fp32 - # (apply_sampling_params copies to fp32 when processing is needed) or the - # model dtype (pass-through otherwise), while draft logits are always the - # model dtype. Warm every (target, draft) combination the runtime can hit. + watermark_config = getattr(worker.vllm_config, "watermark_config", None) + watermark_kwargs: dict[str, Any] = {} + if watermark_config is not None: + from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key + + try: + watermark_key = speculative_target_watermark_key(watermark_config) + except Exception: + logger.warning( + "Skipping spec-decode rejection sampler warmup: could not derive " + "the watermark recovery key.", + exc_info=True, + ) + return + # Tensor dtypes and integer widths affect Triton specialization. + watermark_kwargs = { + "contexts": torch.zeros( + (num_logits, watermark_config.context_width), + dtype=torch.int32, + device=device, + ), + "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device), + "watermark_key": watermark_key, + } + + # Sampling-parameter processing may promote either distribution to fp32. model_dtype = model_config.dtype - warmup_dtype_pairs = { - (model_dtype, model_dtype), - (torch.float32, torch.float32), - (torch.float32, model_dtype), - (model_dtype, torch.float32), - } + warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]] + if use_draft_logits: + warmup_dtype_pairs = { + (model_dtype, model_dtype), + (torch.float32, torch.float32), + (torch.float32, model_dtype), + (model_dtype, torch.float32), + } + else: + warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)} logger.info( "Warming up spec-decode rejection sampler kernels " - "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s).", + "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).", vocab_size, num_spec, [(str(t), str(d)) for t, d in warmup_dtype_pairs], use_block_verification, + watermark_config is not None, ) for tgt_dtype, draft_dtype in warmup_dtype_pairs: target_logits = torch.zeros( (num_logits, vocab_size), dtype=tgt_dtype, device=device ) - draft_logits = torch.zeros( - (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device + draft_logits = ( + torch.zeros( + (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device + ) + if draft_dtype is not None + else None ) synthetic_rates = ( torch.full((num_spec,), 0.5, dtype=torch.float32, device=device) @@ -85,14 +112,14 @@ def spec_decode_rejection_warmup(worker: Worker) -> None: rejection_sample( target_logits=target_logits, draft_logits=draft_logits, - draft_sampled=torch.zeros(num_logits, dtype=torch.int64, device=device), + draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device), cu_num_logits=torch.tensor( [0, num_logits], dtype=torch.int32, device=device ), pos=torch.zeros(num_logits, dtype=torch.int64, device=device), - idx_mapping=torch.zeros(num_reqs, dtype=torch.int32, device=device), + idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), expanded_idx_mapping=torch.zeros( - num_logits, dtype=torch.int32, device=device + num_logits, dtype=torch.int64, device=device ), expanded_local_pos=torch.arange( num_logits, dtype=torch.int32, device=device @@ -103,6 +130,7 @@ def spec_decode_rejection_warmup(worker: Worker) -> None: synthetic_conditional_rates=synthetic_rates, use_fp64=False, use_block_verification=use_block_verification, + **watermark_kwargs, ) except Exception: logger.warning( diff --git a/vllm/model_executor/warmup/watermark_sample_warmup.py b/vllm/model_executor/warmup/watermark_sample_warmup.py new file mode 100644 index 000000000000..d315f10310b7 --- /dev/null +++ b/vllm/model_executor/warmup/watermark_sample_warmup.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up the watermarked sampler's Triton kernel. + +``_philox_gumbel_kernel`` is JIT-compiled once per (key, logits dtype, +skip-mask, ``USE_FP64``) specialization. The generic sampler warmup samples with +``SamplingParams.for_sampler_warmup()``, whose feature-heavy logits processing +forces the fp32 copy in ``apply_sampling_params``, so the first ordinary +temperature-1.0 request is the first launch with model-dtype logits and pays +the compilation inside inference. This pre-compiles every specialization the +sampler path can launch for the configured watermark. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.v1.watermarking.watermarker import Watermarker + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_NUM_WARMUP_TOKENS = 1 + + +def _philox_key(watermarker: Watermarker) -> int | None: + """None when the watermarker takes the torch path and compiles no kernel.""" + from vllm.v1.watermarking.gumbel import GumbelWatermarker + from vllm.v1.watermarking.prfs import PhiloxPRF + + if not isinstance(watermarker, GumbelWatermarker): + return None + prf = watermarker.prf + return prf.key if type(prf) is PhiloxPRF else None + + +def _philox_sampler_keys(watermarker: Watermarker) -> list[int]: + """Philox keys the sampler launches the kernel with.""" + from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker + + watermarkers: list[Watermarker] = [watermarker] + if isinstance(watermarker, DualKeyGumbelWatermarker): + # ``DualKeyGumbelWatermarker.sample`` launches once per key, and + # collapses to a single key at the alpha bounds, where the other key's + # specialization is never launched. + if watermarker.alpha == 0: + watermarkers = [watermarker.draft_watermarker] + elif watermarker.alpha == 1: + watermarkers = [watermarker.target_watermarker] + else: + watermarkers = [ + watermarker.draft_watermarker, + watermarker.target_watermarker, + ] + return [key for key in map(_philox_key, watermarkers) if key is not None] + + +@torch.inference_mode() +def watermark_sample_warmup(worker: Worker) -> None: + if getattr(worker.vllm_config, "watermark_config", None) is None: + return + # ``GumbelWatermarker`` launches the kernel only on CUDA + # (vllm/v1/watermarking/gumbel.py:63,:82); elsewhere the sampler takes the + # torch path and warming here would compile a kernel nothing launches. + if not current_platform.is_cuda_alike(): + return + + from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler + from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample + + # Read the sampler the runner built instead of rebuilding it from config: + # its watermarker (the target role under a speculative config), its context + # deduplication setting and its fp64 flag are what the launches have to + # match, and cannot drift from the runner. A rank that does not sample -- a + # non-final pipeline stage, a pooling model -- holds no such sampler. + sampler = getattr(worker.model_runner, "sampler", None) + if not isinstance(sampler, GPUWatermarkSampler): + return + + model_config = worker.vllm_config.model_config + device = worker.device + try: + keys = _philox_sampler_keys(sampler.watermarker) + if not keys: + return + + # The sampler passes model-dtype logits through when no logits + # processing is needed and fp32 otherwise. ``ModelConfig.dtype`` is + # resolved to a torch dtype during validation. + dtypes = {cast("torch.dtype", model_config.dtype), torch.float32} + # ``GPUWatermarkSampler._sample_random`` builds a skip mask whenever + # context deduplication is on; with it off only a batch that mixes in an + # opted-out or greedy row does, and the mask-free specialization stays + # reachable. + with_skip_mask = ( + (True,) if sampler.deduplicate_contexts != "none" else (True, False) + ) + + vocab_size = model_config.get_vocab_size() + contexts = torch.zeros( + (_NUM_WARMUP_TOKENS, sampler.watermarker.context_width), + dtype=torch.int32, + device=device, + ) + # Dtypes mirror the runtime buffers the sampler passes and are part of + # the specialization key: int32 request-state token ids, int64 index and + # seed buffers, fp32 sampling temperatures. + sampling_state = { + "skip_mask": torch.zeros( + _NUM_WARMUP_TOKENS, dtype=torch.bool, device=device + ), + "expanded_idx_mapping": torch.zeros( + _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device + ), + "temperatures": torch.ones( + _NUM_WARMUP_TOKENS, dtype=torch.float32, device=device + ), + "seeds": torch.zeros(_NUM_WARMUP_TOKENS, dtype=torch.int64, device=device), + "positions": torch.zeros( + _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device + ), + } + + logger.info( + "Warming up watermark sampler kernel (vocab=%d, keys=%d, dtypes=%s, " + "skip_mask=%s).", + vocab_size, + len(keys), + [str(dtype) for dtype in dtypes], + list(with_skip_mask), + ) + for dtype in dtypes: + logits = torch.zeros( + (_NUM_WARMUP_TOKENS, vocab_size), dtype=dtype, device=device + ) + for key in keys: + for use_skip_mask in with_skip_mask: + philox_gumbel_sample( + logits, + contexts, + key, + # Only the mixed path forwards the sampler's setting; + # ``GumbelWatermarker._sample_watermarked`` launches + # the mask-free variant with the fp32 default. + use_fp64=use_skip_mask and sampler.use_fp64_gumbel, + **(sampling_state if use_skip_mask else {}), + ) + except Exception: + logger.warning("Skipping watermark sampler warmup.", exc_info=True) diff --git a/vllm/v1/watermarking/__init__.py b/vllm/v1/watermarking/__init__.py index bb53981461ce..23e9fd20fda6 100644 --- a/vllm/v1/watermarking/__init__.py +++ b/vllm/v1/watermarking/__init__.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.config.watermarking import derive_watermark_key from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector from vllm.v1.watermarking.factory import create_watermarker from vllm.v1.watermarking.gumbel import ( + DualKeyGumbelWatermarkDetector, + DualKeyGumbelWatermarker, GumbelWatermarkDetector, GumbelWatermarker, ) @@ -12,12 +15,19 @@ WatermarkPRF, create_prf, ) -from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample +from vllm.v1.watermarking.watermarker import ( + SupportsSpeculativeDecoding, + Watermarker, + WatermarkSample, +) __all__ = [ + "DualKeyGumbelWatermarkDetector", + "DualKeyGumbelWatermarker", "GumbelWatermarkDetector", "GumbelWatermarker", "PhiloxPRF", + "SupportsSpeculativeDecoding", "WatermarkDetection", "WatermarkDetector", "WatermarkPRF", @@ -25,4 +35,5 @@ "WatermarkSample", "create_watermarker", "create_prf", + "derive_watermark_key", ] diff --git a/vllm/v1/watermarking/factory.py b/vllm/v1/watermarking/factory.py index d9111e7a9e8c..2bc1919aef01 100644 --- a/vllm/v1/watermarking/factory.py +++ b/vllm/v1/watermarking/factory.py @@ -1,24 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - from vllm.config.watermarking import WatermarkConfig -from vllm.v1.watermarking.gumbel import GumbelWatermarker +from vllm.v1.watermarking.gumbel import ( + DualKeyGumbelWatermarker, + GumbelWatermarker, +) from vllm.v1.watermarking.watermarker import Watermarker -def _create_gumbel(config: WatermarkConfig) -> Watermarker: - return GumbelWatermarker(config.key, config.context_width, config.prf) - - -_WATERMARKERS: dict[str, Callable[[WatermarkConfig], Watermarker]] = { - "gumbel": _create_gumbel, -} - - def create_watermarker(config: WatermarkConfig) -> Watermarker: - watermarker_factory = _WATERMARKERS.get(config.algorithm) - if watermarker_factory is None: - raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}") - return watermarker_factory(config) + if config.algorithm == "gumbel": + return GumbelWatermarker(config.key, config.context_width, config.prf) + if config.algorithm == "dual_key_gumbel": + return DualKeyGumbelWatermarker( + config.key, + config.context_width, + config.prf, + config.alpha, + ) + raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}") diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index 7e943cc04035..afa95ee8781b 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -1,16 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Literal + import numpy as np import torch from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p -from vllm.v1.watermarking.watermarker import Watermarker +from vllm.v1.watermarking.watermarker import RandomSampler, Watermarker from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.sample.watermark import repeated_context_mask logger = init_logger(__name__) @@ -20,10 +22,14 @@ def __init__( self, watermarker: Watermarker, *args, + deduplicate_contexts: Literal["none", "single_turn", "all"] = "single_turn", + deduplicate_contexts_max_history: int | None = 8192, **kwargs, ) -> None: super().__init__(*args, **kwargs) self.watermarker = watermarker + self.deduplicate_contexts = deduplicate_contexts + self.deduplicate_contexts_max_history = deduplicate_contexts_max_history self.watermarking = UvaBackedTensor( self.sampling_states.max_num_reqs, dtype=torch.bool ) @@ -72,37 +78,38 @@ def _sample_random( processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) contexts = self._get_contexts(expanded_idx_mapping) + repeated_contexts = None + if self.deduplicate_contexts != "none": + repeated_contexts = self._get_repeated_contexts( + expanded_idx_mapping, contexts + ) - def random_sample(sample_logits: torch.Tensor) -> torch.Tensor: - return gumbel_sample( - sample_logits, - expanded_idx_mapping, - self.sampling_states.temperature.gpu, - self.sampling_states.seeds.gpu, - pos, - apply_temperature=False, - is_drafting=False, - use_fp64=self.use_fp64_gumbel, + temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping] + needs_mixed_sampling = repeated_contexts is not None or not np.all(enabled) + skip_mask = None + if needs_mixed_sampling: + watermarking = self.watermarking.gpu[expanded_idx_mapping] & ( + temperatures != 0 ) + if repeated_contexts is not None: + watermarking &= ~repeated_contexts + skip_mask = ~watermarking + random_sampler = RandomSampler( + expanded_idx_mapping=expanded_idx_mapping, + temperatures=self.sampling_states.temperature.gpu, + seeds=self.sampling_states.seeds.gpu, + positions=pos, + use_fp64=self.use_fp64_gumbel, + ) output = self.watermarker.sample( processed_logits, contexts, - random_sample, + random_sampler, + skip_mask=skip_mask, ) - temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping] - watermarking = self.watermarking.gpu[expanded_idx_mapping] & (temperatures != 0) - if not np.all(enabled): - unwatermarked = random_sample(processed_logits) - sampled = torch.where(watermarking, output.token_ids, unwatermarked) - output_logits = output.logits - if output.logits is not processed_logits: - output_logits = torch.where( - watermarking.unsqueeze(-1), output.logits, processed_logits - ) - else: - sampled = output.token_ids - output_logits = output.logits + sampled = output.token_ids + output_logits = output.logits sampled = torch.where( temperatures == 0, processed_logits.argmax(dim=-1), @@ -110,22 +117,78 @@ def random_sample(sample_logits: torch.Tensor) -> torch.Tensor: ) return sampled, output_logits - def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor: + def _get_repeated_contexts( + self, + expanded_idx_mapping: torch.Tensor, + contexts: torch.Tensor, + ) -> torch.Tensor: + return repeated_context_mask( + self.req_states.all_token_ids.gpu, + expanded_idx_mapping, + self.req_states.prompt_len.gpu, + self.req_states.total_len.gpu, + contexts, + self.deduplicate_contexts_max_history, + include_prompt=self.deduplicate_contexts == "all", + skip_partial_context=self.deduplicate_contexts == "all", + ) + + def _get_contexts( + self, + expanded_idx_mapping: torch.Tensor, + expanded_local_pos: torch.Tensor | None = None, + draft_sampled: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build contexts for flattened verification rows. + + With committed tokens ``[10, 11]``, drafts ``[20, 21]``, and width 2, + local positions ``[0, 1, 2]`` use contexts ``[10, 11]``, ``[11, 20]``, + and ``[20, 21]``. The first verification row precedes the first draft. + """ context_width = self.watermarker.context_width req_indices = expanded_idx_mapping.to(torch.int64) valid_reqs = req_indices >= 0 safe_req_indices = req_indices.clamp_min(0) total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64) prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64) + if expanded_local_pos is None: + expanded_local_pos = torch.zeros_like(req_indices) + else: + expanded_local_pos = expanded_local_pos.to(torch.int64) offsets = torch.arange( -context_width, 0, dtype=torch.int64, device=req_indices.device ) - positions = total_lens.unsqueeze(-1) + offsets - valid_positions = valid_reqs.unsqueeze(-1) & ( - positions >= prompt_lens.unsqueeze(-1) + positions = ( + total_lens.unsqueeze(-1) + expanded_local_pos.unsqueeze(-1) + offsets ) - positions = positions.clamp_min(0) - contexts = self.req_states.all_token_ids.gpu[ - safe_req_indices.unsqueeze(-1), positions + committed = positions < total_lens.unsqueeze(-1) + valid_committed = ( + valid_reqs.unsqueeze(-1) + & committed + & (positions >= prompt_lens.unsqueeze(-1)) + ) + committed_contexts = self.req_states.all_token_ids.gpu[ + safe_req_indices.unsqueeze(-1), + positions.clamp(0, self.req_states.all_token_ids.gpu.shape[1] - 1), ] - return torch.where(valid_positions, contexts, -1) + contexts = torch.where(valid_committed, committed_contexts, -1) + if draft_sampled is None: + return contexts + + row_indices = torch.arange( + len(req_indices), dtype=torch.int64, device=req_indices.device + ) + draft_offsets = positions - total_lens.unsqueeze(-1) + draft_indices = ( + row_indices.unsqueeze(-1) + - expanded_local_pos.unsqueeze(-1) + + draft_offsets + + 1 + ) + valid_drafts = ( + valid_reqs.unsqueeze(-1) + & ~committed + & (draft_offsets < expanded_local_pos.unsqueeze(-1)) + ) + draft_contexts = draft_sampled[draft_indices.clamp(0, len(draft_sampled) - 1)] + return torch.where(valid_drafts, draft_contexts, contexts) diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 856d697aeaf9..cacb620b5e4b 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/vllm/v1/watermarking/gumbel.py @@ -8,13 +8,14 @@ import torch -from vllm.config.watermarking import WatermarkPRFName +from vllm.config.watermarking import WatermarkPRFName, derive_watermark_key from vllm.v1.watermarking.detector import ( WatermarkDetector, ) from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf from vllm.v1.watermarking.watermarker import ( RandomSampler, + SupportsSpeculativeDecoding, Watermarker, WatermarkSample, ) @@ -54,11 +55,10 @@ def __init__( def context_width(self) -> int: return self._context_width - def sample( + def _sample_watermarked( self, logits: torch.Tensor, contexts: torch.Tensor, - _random_sample: RandomSampler, ) -> WatermarkSample: if type(self.prf) is PhiloxPRF and logits.device.type == "cuda": from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample @@ -72,6 +72,110 @@ def sample( noise = -torch.log(-torch.log(uniforms)) return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits) + def _try_sample_mixed( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + skip_mask: torch.Tensor, + random_sampler: RandomSampler, + ) -> WatermarkSample | None: + if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda": + return None + + from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample + + token_ids = philox_gumbel_sample( + logits, + contexts, + self.prf.key, + skip_mask=skip_mask, + expanded_idx_mapping=random_sampler.expanded_idx_mapping, + temperatures=random_sampler.temperatures, + seeds=random_sampler.seeds, + positions=random_sampler.positions, + use_fp64=random_sampler.use_fp64, + ) + return WatermarkSample(token_ids, logits) + + +class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding): + def __init__( + self, + key: int, + context_width: int = 4, + prf: WatermarkPRFName = "philox", + alpha: float = 0.1, + ) -> None: + if not 0 <= alpha <= 1: + raise ValueError("alpha must be between 0 and 1") + self.alpha = alpha + self.draft_watermarker = GumbelWatermarker( + derive_watermark_key(key, b"key_a"), context_width, prf + ) + self.target_watermarker = GumbelWatermarker( + derive_watermark_key(key, b"key_b"), context_width, prf + ) + self._routing_logits_cache: dict[torch.device, torch.Tensor] = {} + + def _routing_logits(self, device: torch.device) -> torch.Tensor: + # Building this tensor on every CUDA call stages a pageable H2D copy and + # synchronizes the stream. + cached = self._routing_logits_cache.get(device) + if cached is None: + cached = torch.tensor( + [1 - self.alpha, self.alpha], + dtype=torch.float32, + device=device, + ).log() + self._routing_logits_cache[device] = cached + return cached + + @property + def context_width(self) -> int: + return self.draft_watermarker.context_width + + def sample( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + random_sampler: RandomSampler | None = None, + skip_mask: torch.Tensor | None = None, + ) -> WatermarkSample: + if self.alpha == 0: + return self.draft_watermarker.sample( + logits, contexts, random_sampler, skip_mask + ) + if self.alpha == 1: + return self.target_watermarker.sample( + logits, contexts, random_sampler, skip_mask + ) + if random_sampler is None: + raise ValueError("dual-key Gumbel routing requires a random sampler") + + key_a_sample = self.draft_watermarker.sample( + logits, contexts, random_sampler, skip_mask + ) + key_b_sample = self.target_watermarker.sample( + logits, contexts, random_sampler, skip_mask + ) + routing_logits = self._routing_logits(logits.device).expand(logits.shape[0], -1) + use_key_a = random_sampler(routing_logits) == 0 + return WatermarkSample( + torch.where( + use_key_a, + key_a_sample.token_ids, + key_b_sample.token_ids, + ), + logits, + ) + + def _sample_watermarked( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + ) -> WatermarkSample: + raise ValueError("dual-key Gumbel routing requires a random sampler") + class GumbelWatermarkDetector(WatermarkDetector): def __init__( @@ -100,6 +204,61 @@ def _aggregate_scores(self, token_scores: torch.Tensor) -> float: return token_scores.sum().item() +class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector): + """Detect a dual-key watermark with weighted early fusion. + + ``alpha`` is the weight assigned to key B; key A receives ``1 - alpha``. + """ + + def __init__( + self, + key: int, + context_width: int = 4, + p_value_threshold: float = 0.01, + prf: WatermarkPRFName = "philox", + deduplicate_contexts: bool = True, + alpha: float = 0.2, + ) -> None: + if not 0 <= alpha <= 1: + raise ValueError("alpha must be between 0 and 1") + super().__init__( + derive_watermark_key(key, b"key_a"), + context_width, + p_value_threshold, + prf, + deduplicate_contexts, + ) + self.key_b_prf = create_prf(prf, derive_watermark_key(key, b"key_b")) + self.alpha = alpha + + def _score_tokens( + self, contexts: torch.Tensor, targets: torch.Tensor + ) -> torch.Tensor: + scores = torch.stack( + [ + -torch.log1p( + -prf.uniform(contexts, targets.unsqueeze(-1)) + .squeeze(-1) + .to(torch.float64) + ) + for prf in (self.prf, self.key_b_prf) + ], + dim=-1, + ) + weights = scores.new_tensor([1 - self.alpha, self.alpha]) + return scores @ weights + + def _get_p_value(self, score: float, num_scored_tokens: int) -> float: + variance = (1 - self.alpha) ** 2 + self.alpha**2 + if self.alpha in (0, 0.5, 1): + return _gamma_survival_integer_shape( + score / variance, round(num_scored_tokens / variance) + ) + shape = torch.tensor(num_scored_tokens / variance, dtype=torch.float64) + scaled_score = torch.tensor(score / variance, dtype=torch.float64) + return torch.special.gammaincc(shape, scaled_score).item() + + def _validate_context_width(context_width: int) -> None: if context_width < 1: raise ValueError("context_width must be positive") diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py new file mode 100644 index 000000000000..c3675a3ab953 --- /dev/null +++ b/vllm/v1/watermarking/spec_decode.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.config.watermarking import WatermarkConfig +from vllm.v1.watermarking.factory import create_watermarker +from vllm.v1.watermarking.gumbel import GumbelWatermarker +from vllm.v1.watermarking.prfs import PhiloxPRF +from vllm.v1.watermarking.watermarker import ( + SupportsSpeculativeDecoding, + Watermarker, +) +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample + + +class DraftWatermarker: + def __init__( + self, + watermarker: Watermarker, + max_num_reqs: int, + device: torch.device, + ) -> None: + self.watermarker = watermarker + self.contexts = torch.zeros( + max_num_reqs, + watermarker.context_width, + dtype=torch.int64, + device=device, + ) + self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + + def prepare(self, contexts: torch.Tensor, enabled: torch.Tensor) -> None: + num_reqs = contexts.shape[0] + self.contexts[:num_reqs].copy_(contexts) + self.enabled[:num_reqs].copy_(enabled) + + def sample( + self, + logits: torch.Tensor, + ordinary_sampled: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + ) -> torch.Tensor: + request_temperatures = temperature[idx_mapping] + processed_logits = logits / torch.where( + request_temperatures == 0, 1, request_temperatures + ).unsqueeze(-1) + watermarked = self.watermarker.sample( + processed_logits, + self.contexts[: logits.shape[0]], + ).token_ids + enabled = self.enabled[: logits.shape[0]] & (request_temperatures != 0) + sampled = torch.where(enabled, watermarked, ordinary_sampled) + contexts = self.contexts[: logits.shape[0]] + contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1)) + return sampled + + +def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker: + if isinstance(watermarker, SupportsSpeculativeDecoding): + return watermarker.target_watermarker + return watermarker + + +def create_speculative_draft_watermarker( + watermarker: Watermarker, + max_num_reqs: int, + device: torch.device, + allow_target_only: bool, +) -> DraftWatermarker | None: + if isinstance(watermarker, SupportsSpeculativeDecoding): + return DraftWatermarker(watermarker.draft_watermarker, max_num_reqs, device) + if allow_target_only: + return None + raise ValueError( + f"The {type(watermarker).__name__} watermarking algorithm does not support " + "speculative decoding. Set allow_target_only_watermarking=true to " + "leave draft tokens unwatermarked." + ) + + +def _philox_key(watermarker: Watermarker) -> int: + if not ( + isinstance(watermarker, GumbelWatermarker) + and type(watermarker.prf) is PhiloxPRF + ): + raise NotImplementedError( + "In-kernel watermarked recovery supports only the Philox PRF " + "GumbelWatermarker" + ) + return watermarker.prf.key + + +def _resolve_watermark_key(watermarker: Watermarker) -> int: + """Return the Philox recovery key.""" + if isinstance(watermarker, SupportsSpeculativeDecoding): + raise ValueError( + f"{type(watermarker).__name__} keys the target role separately from " + "the draft. Pass create_speculative_target_watermarker(watermarker) " + "so the recovery draw carries the target's key." + ) + return _philox_key(watermarker) + + +def speculative_target_watermark_key( + watermark_config: WatermarkConfig | None, +) -> int | None: + """Resolve the resample kernel's Philox key from configuration.""" + if watermark_config is None: + return None + + return _resolve_watermark_key( + create_speculative_target_watermarker(create_watermarker(watermark_config)) + ) + + +def watermarked_rejection_sample( + target_logits: torch.Tensor, + draft_logits: torch.Tensor | None, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + pos: torch.Tensor, + idx_mapping: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + expanded_local_pos: torch.Tensor, + temperature: torch.Tensor, + seed: torch.Tensor, + num_speculative_steps: int, + contexts: torch.Tensor, + watermarking: torch.Tensor, + watermarker: Watermarker, + use_fp64: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Use the target key for rejection recovery and bonus tokens.""" + assert contexts.shape == (target_logits.shape[0], watermarker.context_width) + return rejection_sample( + target_logits, + draft_logits, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seed, + num_speculative_steps, + use_fp64=use_fp64, + contexts=contexts, + watermarking=watermarking, + watermark_key=_resolve_watermark_key(watermarker), + ) diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py index f00659393efd..12570f4ae5ca 100644 --- a/vllm/v1/watermarking/watermarker.py +++ b/vllm/v1/watermarking/watermarker.py @@ -2,13 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod -from collections.abc import Callable from dataclasses import dataclass -from typing import TypeAlias +from typing import Protocol, runtime_checkable import torch -RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor] +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @dataclass(frozen=True) @@ -17,17 +16,81 @@ class WatermarkSample: logits: torch.Tensor +@dataclass(frozen=True) +class RandomSampler: + expanded_idx_mapping: torch.Tensor + temperatures: torch.Tensor + seeds: torch.Tensor + positions: torch.Tensor + use_fp64: bool = False + + def __call__(self, logits: torch.Tensor) -> torch.Tensor: + return gumbel_sample( + logits, + self.expanded_idx_mapping, + self.temperatures, + self.seeds, + self.positions, + apply_temperature=False, + is_drafting=False, + use_fp64=self.use_fp64, + ) + + class Watermarker(ABC): @property @abstractmethod def context_width(self) -> int: raise NotImplementedError - @abstractmethod def sample( self, logits: torch.Tensor, contexts: torch.Tensor, - random_sample: RandomSampler, + random_sampler: RandomSampler | None = None, + skip_mask: torch.Tensor | None = None, + ) -> WatermarkSample: + if skip_mask is not None: + if random_sampler is None: + raise ValueError("skip_mask requires a random sampler") + mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler) + if mixed is not None: + return mixed + + watermarked = self._sample_watermarked(logits, contexts) + if skip_mask is None: + return watermarked + + assert random_sampler is not None + token_ids = torch.where( + skip_mask, random_sampler(logits), watermarked.token_ids + ) + output_logits = watermarked.logits + if output_logits is not logits: + output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits) + return WatermarkSample(token_ids, output_logits) + + @abstractmethod + def _sample_watermarked( + self, + logits: torch.Tensor, + contexts: torch.Tensor, ) -> WatermarkSample: raise NotImplementedError + + def _try_sample_mixed( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + skip_mask: torch.Tensor, + random_sampler: RandomSampler, + ) -> WatermarkSample | None: + return None + + +@runtime_checkable +class SupportsSpeculativeDecoding(Protocol): + """A watermarker compatible with standard speculative rejection sampling.""" + + draft_watermarker: Watermarker + target_watermarker: Watermarker diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 14e3823c0410..0ce0288f310a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -79,6 +79,10 @@ ) from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler +from vllm.v1.watermarking.spec_decode import ( + create_speculative_target_watermarker, + speculative_target_watermark_key, +) from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu import pcp_manager as pcp @@ -472,8 +476,18 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: if self.vllm_config.watermark_config is None: self.sampler = Sampler(**sampler_kwargs) else: - watermarker = create_watermarker(self.vllm_config.watermark_config) - self.sampler = GPUWatermarkSampler(watermarker, **sampler_kwargs) + wm_config = self.vllm_config.watermark_config + watermarker = create_watermarker(wm_config) + if self.speculative_config is not None: + watermarker = create_speculative_target_watermarker(watermarker) + self.sampler = GPUWatermarkSampler( + watermarker, + deduplicate_contexts=wm_config.deduplicate_contexts, + deduplicate_contexts_max_history=( + wm_config.deduplicate_contexts_max_history + ), + **sampler_kwargs, + ) custom = self.model_state.custom_sampler(self.sampler) if custom: @@ -484,6 +498,9 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: self.sampler, self.speculative_config, self.device, + watermark_key=speculative_target_watermark_key( + self.vllm_config.watermark_config + ), ) self.prompt_logprobs_worker = PromptLogprobsWorker( self.max_num_reqs, @@ -860,6 +877,11 @@ def _dummy_run( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + if isinstance(self.sampler, GPUWatermarkSampler): + self.speculator.prepare_watermarking( + self.sampler._get_contexts(input_batch.idx_mapping), + self.sampler.watermarking.gpu[input_batch.idx_mapping], + ) with use_workspace_lane(self._draft_workspace_lane): self.speculator.propose( input_batch=input_batch, @@ -2086,6 +2108,11 @@ def sample_tokens( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] + if isinstance(self.sampler, GPUWatermarkSampler): + self.speculator.prepare_watermarking( + self.sampler._get_contexts(input_batch.idx_mapping), + self.sampler.watermarking.gpu[input_batch.idx_mapping], + ) with use_workspace_lane(self._draft_workspace_lane): draft_tokens = self.speculator.propose( input_batch, diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index a0105cc1b7f0..0a9c64cbf732 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/vllm/v1/worker/gpu/sample/watermark.py @@ -25,6 +25,7 @@ from vllm.v1.watermarking.prfs.philox import ( _UINT32_MASK as _UINT32_MASK_VALUE, ) +from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax if HAS_TRITON: from triton.language import math as tl_math @@ -42,6 +43,162 @@ _TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE +@triton.jit +def _repeated_context_mask_kernel( + output_ptr, + all_token_ids_ptr, + all_token_ids_stride, + req_indices_ptr, + prompt_lens_ptr, + total_lens_ptr, + contexts_ptr, + context_stride, + CONTEXT_WIDTH: tl.constexpr, + MAX_HISTORY: tl.constexpr, + INCLUDE_PROMPT: tl.constexpr, + SKIP_PARTIAL_CONTEXT: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + req_idx = tl.load(req_indices_ptr + row).to(tl.int64) + if req_idx < 0: + tl.store(output_ptr + row, 0) + return + prompt_len = tl.load(prompt_lens_ptr + req_idx) + total_len = tl.load(total_lens_ptr + req_idx) + if SKIP_PARTIAL_CONTEXT and total_len - prompt_len < CONTEXT_WIDTH: + tl.store(output_ptr + row, 1) + return + sequence_start = 0 if INCLUDE_PROMPT else prompt_len + history_len = total_len - sequence_start + + offsets = tl.arange(0, BLOCK) + repeated = tl.full((), 0, tl.int32) + scan_start = 0 + if MAX_HISTORY > 0: + scan_start = tl.maximum(history_len - MAX_HISTORY, 0) + # Aligned loop bounds let the per-lane offsets fold into immediates. + aligned_start = (scan_start // BLOCK) * BLOCK + for block_start in tl.range(aligned_start, history_len, BLOCK): + previous_pos = block_start + offsets + in_window = (previous_pos >= scan_start) & (previous_pos < history_len) + matches = in_window + for offset in range(CONTEXT_WIDTH): + context_token = tl.load(contexts_ptr + row * context_stride + offset) + historical_pos = previous_pos + offset - CONTEXT_WIDTH + historical_token = tl.load( + all_token_ids_ptr + + req_idx * all_token_ids_stride + + sequence_start + + tl.maximum(historical_pos, 0), + mask=in_window & (historical_pos >= 0), + other=-1, + ) + matches &= historical_token == context_token + repeated |= tl.max(matches.to(tl.int32), axis=0) + tl.store(output_ptr + row, repeated) + + +def _repeated_context_mask_cpu( + all_token_ids: torch.Tensor, + req_indices: torch.Tensor, + prompt_lens: torch.Tensor, + total_lens: torch.Tensor, + contexts: torch.Tensor, + max_history: int | None = None, + include_prompt: bool = False, + skip_partial_context: bool = False, +) -> torch.Tensor: + """Reference implementation; the parity tests check the Triton kernel against it.""" + repeated = torch.zeros(len(req_indices), dtype=torch.bool) + for row, req_idx_tensor in enumerate(req_indices): + req_idx = int(req_idx_tensor) + if req_idx < 0: + continue + prompt_len = int(prompt_lens[req_idx]) + total_len = int(total_lens[req_idx]) + if skip_partial_context and total_len - prompt_len < contexts.shape[-1]: + repeated[row] = True + continue + sequence_start = 0 if include_prompt else prompt_len + history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist() + prefix = [-1] * contexts.shape[-1] + current_context = tuple(contexts[row].tolist()) + history_start = ( + max(0, len(history_tokens) - max_history) if max_history is not None else 0 + ) + for history_pos, token_id in enumerate(history_tokens): + if history_pos >= history_start and ( + tuple(prefix[-contexts.shape[-1] :]) == current_context + ): + repeated[row] = True + break + prefix.append(token_id) + return repeated + + +def repeated_context_mask( + all_token_ids: torch.Tensor, + req_indices: torch.Tensor, + prompt_lens: torch.Tensor, + total_lens: torch.Tensor, + contexts: torch.Tensor, + max_history: int | None = None, + include_prompt: bool = False, + skip_partial_context: bool = False, +) -> torch.Tensor: + """Return, per row, whether the row's context already occurred in its history. + + Args: + all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request. + req_indices: Request slot per sampled row; -1 marks a padding row, which + is reported as not repeated. + prompt_lens: Prompt length per request slot. + total_lens: Prompt plus generated length per request slot. + contexts: `[num_rows, context_width]` context per row, padded with -1 + before the start of the scanned history. + max_history: Number of most recent history positions searched, or + ``None`` for all of them. The window compared at each position + reaches `context_width` tokens further back. + include_prompt: Search the prompt as well as the generated tokens. + skip_partial_context: Mark contexts containing start padding so they use + ordinary sampling. + """ + if max_history is not None and max_history < 1: + raise ValueError("max_history must be positive or None") + if all_token_ids.device.type == "cpu": + return _repeated_context_mask_cpu( + all_token_ids, + req_indices, + prompt_lens, + total_lens, + contexts, + max_history, + include_prompt, + skip_partial_context, + ) + + if contexts.stride(-1) != 1: + contexts = contexts.contiguous() + repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device) + _repeated_context_mask_kernel[(len(req_indices),)]( + repeated, + all_token_ids, + all_token_ids.stride(0), + req_indices, + prompt_lens, + total_lens, + contexts, + contexts.stride(0), + CONTEXT_WIDTH=contexts.shape[-1], + MAX_HISTORY=0 if max_history is None else max_history, + INCLUDE_PROMPT=include_prompt, + SKIP_PARTIAL_CONTEXT=skip_partial_context, + BLOCK=512, + ) + return repeated + + @triton.jit def _mulhilo32(multiplier: tl.constexpr, value): high = tl_math.umulhi(multiplier, value) @@ -74,13 +231,101 @@ def _uint32_to_uniform(value): @triton.jit -def _gumbel_value(logits_ptr, output, mask): - logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) +def _philox_context_state( + contexts_row_ptr, + key_0, + key_1, + CONTEXT_WIDTH: tl.constexpr, +): + state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32) + state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32) + state_2 = tl.full((), 0, tl.uint32) + state_3 = tl.full((), 0, tl.uint32) + for offset in range(0, CONTEXT_WIDTH, 4): + context_0 = tl.load(contexts_row_ptr + offset).to(tl.uint32) + if offset + 1 < CONTEXT_WIDTH: + context_1 = tl.load(contexts_row_ptr + offset + 1).to(tl.uint32) + else: + context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32) + if offset + 2 < CONTEXT_WIDTH: + context_2 = tl.load(contexts_row_ptr + offset + 2).to(tl.uint32) + else: + context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32) + if offset + 3 < CONTEXT_WIDTH: + context_3 = tl.load(contexts_row_ptr + offset + 3).to(tl.uint32) + else: + context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32) + state_0, state_1, state_2, state_3 = _philox4x32_10( + (state_0 ^ context_0) & _UINT32_MASK, + (state_1 ^ context_1) & _UINT32_MASK, + (state_2 ^ context_2) & _UINT32_MASK, + (state_3 ^ context_3) & _UINT32_MASK, + (key_0 ^ offset) & _UINT32_MASK, + key_1, + ) + return state_0, state_1, state_2, state_3 + + +@triton.jit +def _philox_candidate_words(groups, state_0, state_1, state_2, state_3, key_0, key_1): + candidate_words = groups.to(tl.uint32) + vector_zero = candidate_words * 0 + return _philox4x32_10( + candidate_words & _UINT32_MASK, + state_0 + vector_zero, + state_1 + vector_zero, + state_2 + vector_zero, + ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero, + ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero, + ) + + +@triton.jit +def _philox_gumbel_from_logits(logits, word): logits = tl.where(logits != logits, float("-inf"), logits) - uniform = _uint32_to_uniform(output) + uniform = _uint32_to_uniform(word) return logits - tl.log(-tl.log(uniform)) +@triton.jit +def _gumbel_value(logits_ptr, output, mask): + logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) + return _philox_gumbel_from_logits(logits, output) + + +@triton.jit +def philox_gumbel_block_argmax( + logits, + mask, + block_idx, + contexts_row_ptr, + key_0, + key_1, + CONTEXT_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Apply the detector-compatible PRF to one vocabulary block.""" + tl.static_assert(BLOCK_SIZE % 4 == 0) + state_0, state_1, state_2, state_3 = _philox_context_state( + contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH + ) + groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4) + output_0, output_1, output_2, output_3 = _philox_candidate_words( + groups, state_0, state_1, state_2, state_3, key_0, key_1 + ) + # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3. + words = tl.interleave( + tl.interleave(output_0, output_2), + tl.interleave(output_1, output_3), + ) + values = tl.where( + mask, + _philox_gumbel_from_logits(logits.to(tl.float32), words), + float("-inf"), + ) + return tl.max(values, axis=0, return_indices=True) + + @triton.jit def _select_max(best_value, best_token, candidate_value, candidate): take_candidate = candidate_value > best_value @@ -100,60 +345,63 @@ def _philox_gumbel_kernel( logits_stride, contexts_ptr, context_stride, + skip_mask_ptr, + expanded_idx_mapping_ptr, + seeds_ptr, + pos_ptr, + temp_ptr, key_0_value, key_1_value, vocab_size, CONTEXT_WIDTH: tl.constexpr, BLOCK_SIZE: tl.constexpr, + USE_FP64: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) block_index = tl.program_id(1) groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4) - key_0 = key_0_value.to(tl.uint32) - key_1 = key_1_value.to(tl.uint32) - state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32) - state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32) - state_2 = tl.full((), 0, tl.uint32) - state_3 = tl.full((), 0, tl.uint32) - for offset in range(0, CONTEXT_WIDTH, 4): - context_0 = tl.load(contexts_ptr + row * context_stride + offset).to(tl.uint32) - if offset + 1 < CONTEXT_WIDTH: - context_1 = tl.load(contexts_ptr + row * context_stride + offset + 1).to( - tl.uint32 + if skip_mask_ptr is not None: + req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64) + valid_req = req_state_idx >= 0 + if tl.load(skip_mask_ptr + row): + temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to( + tl.float32 ) - else: - context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32) - if offset + 2 < CONTEXT_WIDTH: - context_2 = tl.load(contexts_ptr + row * context_stride + offset + 2).to( - tl.uint32 + seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0) + pos = tl.load(pos_ptr + row) + candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = candidate < vocab_size + logits_row = logits_ptr + row * logits_stride + logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to( + tl.float32 ) - else: - context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32) - if offset + 3 < CONTEXT_WIDTH: - context_3 = tl.load(contexts_ptr + row * context_stride + offset + 3).to( - tl.uint32 + value, index = gumbel_noised_argmax( + logits, + candidate, + mask, + seed, + pos, + temp, + IS_DRAFTING=False, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, ) - else: - context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32) - state_0, state_1, state_2, state_3 = _philox4x32_10( - (state_0 ^ context_0) & _UINT32_MASK, - (state_1 ^ context_1) & _UINT32_MASK, - (state_2 ^ context_2) & _UINT32_MASK, - (state_3 ^ context_3) & _UINT32_MASK, - (key_0 ^ offset) & _UINT32_MASK, - key_1, - ) + token_id = block_index * BLOCK_SIZE + index + tl.store( + local_argmax_ptr + row * local_argmax_stride + block_index, + token_id, + ) + tl.store(local_max_ptr + row * local_max_stride + block_index, value) + return - candidate_words = groups.to(tl.uint32) - vector_zero = candidate_words * 0 - output_0, output_1, output_2, output_3 = _philox4x32_10( - candidate_words & _UINT32_MASK, - state_0 + vector_zero, - state_1 + vector_zero, - state_2 + vector_zero, - ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero, - ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero, + key_0 = key_0_value.to(tl.uint32) + key_1 = key_1_value.to(tl.uint32) + state_0, state_1, state_2, state_3 = _philox_context_state( + contexts_ptr + row * context_stride, key_0, key_1, CONTEXT_WIDTH + ) + output_0, output_1, output_2, output_3 = _philox_candidate_words( + groups, state_0, state_1, state_2, state_3, key_0, key_1 ) candidate_0 = groups * 4 candidate_1 = candidate_0 + 1 @@ -184,17 +432,42 @@ def _philox_gumbel_kernel( def philox_gumbel_sample( - logits: torch.Tensor, contexts: torch.Tensor, key: int + logits: torch.Tensor, + contexts: torch.Tensor, + key: int, + *, + skip_mask: torch.Tensor | None = None, + expanded_idx_mapping: torch.Tensor | None = None, + temperatures: torch.Tensor | None = None, + seeds: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + use_fp64: bool = False, ) -> torch.Tensor: + sampling_state = (expanded_idx_mapping, temperatures, seeds, positions) + if skip_mask is None and any(value is not None for value in sampling_state): + raise ValueError("sampling state requires skip_mask") + if skip_mask is not None and any(value is None for value in sampling_state): + raise ValueError("skip_mask requires complete sampling state") + if logits.stride(-1) != 1: logits = logits.contiguous() if contexts.stride(-1) != 1: contexts = contexts.contiguous() + if skip_mask is not None: + skip_mask = skip_mask.contiguous() + if expanded_idx_mapping is not None: + expanded_idx_mapping = expanded_idx_mapping.contiguous() + if positions is not None: + positions = positions.contiguous() num_tokens, vocab_size = logits.shape block_size = 1024 num_blocks = triton.cdiv(vocab_size, block_size) local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64) - local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32) + local_max = logits.new_empty( + num_tokens, + num_blocks, + dtype=torch.float64 if use_fp64 else torch.float32, + ) _philox_gumbel_kernel[(num_tokens, num_blocks)]( local_argmax, local_argmax.stride(0), @@ -204,11 +477,17 @@ def philox_gumbel_sample( logits.stride(0), contexts, contexts.stride(0), + skip_mask, + expanded_idx_mapping, + seeds, + positions, + temperatures, key & _UINT32_MASK_VALUE, key >> 32, vocab_size, CONTEXT_WIDTH=contexts.shape[-1], BLOCK_SIZE=block_size, + USE_FP64=use_fp64, ) max_block_index = local_max.argmax(dim=-1, keepdim=True) return local_argmax.gather(dim=-1, index=max_block_index).view(-1) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index e9eca550ca9a..036c3014f3cb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -137,7 +137,7 @@ def _sample_logits( # sample_pos is the predicted token's position P. Sampling keys a draw # by the position before the sampled token, P-1. - return gumbel_sample( + sampled = gumbel_sample( logits, idx_map, self.temperature, @@ -149,6 +149,14 @@ def _sample_logits( logits_cache_col=self._step_cols[step], use_fp64=self.use_fp64_gumbel, ) + if self.draft_watermarker is not None: + sampled = self.draft_watermarker.sample( + logits, + sampled, + idx_map, + self.temperature, + ) + return sampled def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: if self._draft_topk is not None: diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index a9027e0550f8..609924e87fd2 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Iterator +from typing import Any import numpy as np import torch @@ -10,6 +11,7 @@ from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.worker.gpu.input_batch import ( InputBatch, get_num_sampled_and_rejected, @@ -78,8 +80,13 @@ def __init__( sampler: Sampler, spec_config: SpeculativeConfig, device: torch.device, + *, + watermark_key: int | None = None, ): self.sampler = sampler + self.watermark_key = watermark_key + if watermark_key is not None: + assert isinstance(sampler, GPUWatermarkSampler) self.num_speculative_steps = spec_config.num_speculative_tokens self.enable_adaptive_verification = spec_config.enable_adaptive_verification rejection_sample_method = spec_config.rejection_sample_method @@ -141,6 +148,23 @@ def _get_logprobs_tensors( in ("raw_logits", "processed_logits"), ) + def _watermarking_kwargs( + self, + draft_sampled: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + expanded_local_pos: torch.Tensor, + ) -> dict[str, Any]: + if self.watermark_key is None: + return {} + assert isinstance(self.sampler, GPUWatermarkSampler) + return { + "contexts": self.sampler._get_contexts( + expanded_idx_mapping, expanded_local_pos, draft_sampled + ), + "watermarking": self.sampler.watermarking.gpu, + "watermark_key": self.watermark_key, + } + def _verify( self, logits: torch.Tensor, @@ -177,6 +201,9 @@ def _verify( self.synthetic_conditional_rates, use_fp64=self.sampler.use_fp64_gumbel, use_block_verification=self.use_block_verification, + **self._watermarking_kwargs( + draft_sampled, expanded_idx_mapping, expanded_local_pos + ), ) return processed_logits, sampled, num_sampled diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 29b20afea69a..9d64c1ecbcf4 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -4,6 +4,7 @@ from vllm.triton_utils import tl, tldevice, triton from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32 +from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax @triton.jit @@ -693,6 +694,39 @@ def _rejection_kernel( @triton.jit +def _seeded_resample_argmax( + residual_logits, + block, + mask, + resample_token_idx, + expanded_idx_mapping_ptr, + temp_ptr, + seed_ptr, + pos_ptr, + vocab_size, + USE_FP64: tl.constexpr, +): + return gumbel_block_argmax( + residual_logits, + block, + mask, + resample_token_idx, + expanded_idx_mapping_ptr, + temp_ptr, + seed_ptr, + pos_ptr, + None, # logits_cache_ptr + 0, # logits_cache_stride_0 + 0, # logits_cache_stride_1 + None, # logits_cache_col_ptr + vocab_size, + IS_DRAFTING=False, + APPLY_TEMPERATURE=False, + USE_FP64=USE_FP64, + ) + + +@triton.jit(do_not_specialize=["watermark_key_0", "watermark_key_1"]) def _resample_kernel( # [num_reqs, num_blocks] resampled_local_argmax_ptr, @@ -727,11 +761,20 @@ def _resample_kernel( pos_ptr, # [num_logits] cumulative_log_p_ptr, + # [num_logits, CONTEXT_WIDTH] + contexts_ptr, + contexts_stride, + # [max_num_reqs], uint8 view of a bool tensor + watermarking_ptr, + watermark_key_0, + watermark_key_1, vocab_size, BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, USE_FP64: tl.constexpr, USE_BLOCK_VERIFICATION: tl.constexpr, + CONTEXT_WIDTH: tl.constexpr, + WATERMARK: tl.constexpr, ): req_idx = tl.program_id(0) resample_idx = tl.load(rejected_step_ptr + req_idx) @@ -763,10 +806,7 @@ def _resample_kernel( other=float("-inf"), ).to(tl.float32) - # Compute the residual logits to resample the rejected token from. if is_bonus or not is_valid_rejected_draft: - # Bonus token (no rejections) or -1 placeholder token. In either case, - # directly use the target logits. residual_logits = target_logits elif HAS_DRAFT_LOGITS: # draft_logits is stored pre-temperature, so apply scale first. @@ -796,12 +836,7 @@ def _resample_kernel( ) target_log_probs += log_p_tau draft_log_probs = draft_logits - draft_lse - # Compute the residual: - # r(x) = max(p(x) - q(x), 0) - # Gumbel sampling needs logits, so we compute it in log space: - # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0)) - # The more numerically stable form is: - # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0)) + # Compute log(max(p - q, 0)) without subtracting probabilities. ratio = tl.exp(draft_log_probs - target_log_probs) residual_logits = tl.where( ratio < 1.0, @@ -809,38 +844,62 @@ def _resample_kernel( float("-inf"), ).to(tl.float32) else: - # One-hot draft. The residual is just the target distribution with - # the rejected draft token probability zeroed out. - # NOTE: During block verification, the residual becomes: - # 0 if x == rejected_draft_token - # p_tau * M_b(x) / Z otherwise - # Therefore p_tau is a constant that cancels under normalization, - # and does not need to be applied. + # The block-verification factor is constant and cancels on normalization. residual_logits = tl.where( block != rejected_draft_token, target_logits, float("-inf"), ).to(tl.float32) - # Resample the rejected/bonus token. - value, idx = gumbel_block_argmax( - residual_logits, - block, - mask, - resample_token_idx, - expanded_idx_mapping_ptr, - temp_ptr, - seed_ptr, - pos_ptr, - None, # logits_cache_ptr - 0, # logits_cache_stride_0 - 0, # logits_cache_stride_1 - None, # logits_cache_col_ptr - vocab_size, - IS_DRAFTING=False, - APPLY_TEMPERATURE=False, - USE_FP64=USE_FP64, - ) + if WATERMARK: + # Padded and greedy rows retain the stock draw. + is_watermarked = ( + tl.load( + watermarking_ptr + req_state_idx, + mask=req_state_idx >= 0, + other=0, + ) + != 0 + ) + if is_watermarked & (temp != 0.0): + watermark_value, idx = philox_gumbel_block_argmax( + residual_logits, + mask, + block_idx, + contexts_ptr + resample_token_idx * contexts_stride, + watermark_key_0.to(tl.uint32), + watermark_key_1.to(tl.uint32), + CONTEXT_WIDTH, + BLOCK_SIZE, + ) + # Detector compatibility fixes the keyed draw at fp32. + value = watermark_value.to(tl.float64) if USE_FP64 else watermark_value + else: + value, idx = _seeded_resample_argmax( + residual_logits, + block, + mask, + resample_token_idx, + expanded_idx_mapping_ptr, + temp_ptr, + seed_ptr, + pos_ptr, + vocab_size, + USE_FP64=USE_FP64, + ) + else: + value, idx = _seeded_resample_argmax( + residual_logits, + block, + mask, + resample_token_idx, + expanded_idx_mapping_ptr, + temp_ptr, + seed_ptr, + pos_ptr, + vocab_size, + USE_FP64=USE_FP64, + ) token_id = block_idx * BLOCK_SIZE + idx tl.store( resampled_local_argmax_ptr @@ -947,6 +1006,9 @@ def rejection_sample( synthetic_conditional_rates: torch.Tensor | None = None, use_fp64: bool = False, use_block_verification: bool = False, + contexts: torch.Tensor | None = None, + watermarking: torch.Tensor | None = None, + watermark_key: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: assert target_logits.ndim == 2 and target_logits.stride(-1) == 1 assert draft_logits is None or ( @@ -954,6 +1016,30 @@ def rejection_sample( ) num_reqs = cu_num_logits.shape[0] - 1 num_logits, vocab_size = target_logits.shape + + watermark = contexts is not None + assert watermark == (watermarking is not None) == (watermark_key is not None), ( + "contexts, watermarking and watermark_key must be set together." + ) + contexts_stride = 0 + context_width = 1 + watermarking_bytes: torch.Tensor | None = None + watermark_key_0 = 0 + watermark_key_1 = 0 + if contexts is not None: + assert watermarking is not None and watermark_key is not None + assert contexts.ndim == 2 and contexts.shape[0] == num_logits + # Context words are hashed as uint32, so int32 (the request-state token + # dtype) and int64 both land on the same PRF stream, -1 included. + assert contexts.dtype in (torch.int32, torch.int64) + if contexts.stride(-1) != 1: + contexts = contexts.contiguous() + assert watermarking.ndim == 1 and watermarking.dtype == torch.bool + contexts_stride = contexts.stride(0) + context_width = contexts.shape[-1] + watermarking_bytes = watermarking.view(torch.uint8) + watermark_key_0 = watermark_key & 0xFFFFFFFF + watermark_key_1 = watermark_key >> 32 draft_logits_stride_0 = 0 draft_logits_stride_1 = 0 if has_draft_logits := draft_logits is not None: @@ -1165,11 +1251,18 @@ def rejection_sample( seed, pos, cumulative_log_p, + contexts, + contexts_stride, + watermarking_bytes, + watermark_key_0, + watermark_key_1, vocab_size, BLOCK_SIZE=RESAMPLE_BLOCK_SIZE, HAS_DRAFT_LOGITS=has_draft_logits, USE_FP64=use_fp64, USE_BLOCK_VERIFICATION=use_block_verification, + CONTEXT_WIDTH=context_width, + WATERMARK=watermark, ) # Insert the resampled tokens into the output sampled. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 4f20637d7020..f6704412191a 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -17,6 +17,11 @@ from vllm.model_executor.models import supports_multimodal_embeddings from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.watermarking import create_watermarker +from vllm.v1.watermarking.spec_decode import ( + DraftWatermarker, + create_speculative_draft_watermarker, +) from vllm.v1.worker.gpu.attn_utils import ( build_attn_metadata, init_attn_backend, @@ -168,6 +173,17 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=device, ) + self.draft_watermarker: DraftWatermarker | None = None + watermark_config = getattr(vllm_config, "watermark_config", None) + if watermark_config is not None: + watermarker = create_watermarker(watermark_config) + self.draft_watermarker = create_speculative_draft_watermarker( + watermarker, + self.max_num_reqs, + device, + watermark_config.allow_target_only_watermarking, + ) + self.supports_mm_inputs = False self.pcp_manager: PCPManager | None = None @@ -373,7 +389,7 @@ def sample_draft( ) -> torch.Tensor: if draft_logits is not None: logits = self.model.compute_logits(hidden_states) - return gumbel_sample( + sampled = gumbel_sample( logits, idx_mapping, temperature, @@ -385,8 +401,23 @@ def sample_draft( logits_cache_col=draft_step, use_fp64=self.use_fp64_gumbel, ) + if self.draft_watermarker is not None: + sampled = self.draft_watermarker.sample( + logits, + sampled, + idx_mapping, + temperature, + ) + return sampled return self._greedy_sample_draft(hidden_states) + def prepare_watermarking( + self, contexts: torch.Tensor, watermarking: torch.Tensor + ) -> None: + if self.draft_watermarker is None: + return + self.draft_watermarker.prepare(contexts, watermarking) + def _copy_request_inputs( self, num_reqs: int, From 72d4d8315749c146cf19f047fecfec491911770d Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:09:19 -0700 Subject: [PATCH 109/774] [ROCm][Bugfix] Fix elastic EP scaling deadlock (#56610) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 4.8 (1M context) --- .../device_communicators/pynccl.py | 31 ++++++++++ .../distributed/elastic_ep/elastic_execute.py | 15 ++++- vllm/distributed/elastic_ep/standby_state.py | 37 ++++++----- vllm/distributed/parallel_state.py | 61 ++++++++++++------- 4 files changed, 106 insertions(+), 38 deletions(-) diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index 916d0308bf76..ece4db82dae2 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/vllm/distributed/device_communicators/pynccl.py @@ -3,6 +3,7 @@ # ===================== import region ===================== +import contextlib import threading import torch @@ -26,6 +27,33 @@ logger = init_logger(__name__) +_warmup_tls = threading.local() + + +def _warmup_deferred() -> bool: + return getattr(_warmup_tls, "deferred", False) + + +@contextlib.contextmanager +def defer_comm_warmup_on_rocm(): + """Build communicators on this thread without their warm-up on ROCm, where + the warm-up can deadlock if built next to a live engine whose existing + communicators still have collectives in flight; instead the warm-up is + deferred to when the engine is quiet (during commit). + """ + from vllm.platforms import current_platform + + if not current_platform.is_rocm(): + yield + return + prev = _warmup_deferred() + _warmup_tls.deferred = True + try: + yield + finally: + _warmup_tls.deferred = prev + + _NCCL_SYMM_OPS_REGISTERED = False _NCCL_SUSPEND_MEM = 0x01 @@ -153,6 +181,9 @@ def _init_comm(self, device: int | str | torch.device) -> None: self.world_size, self.unique_id, self.rank ) + if _warmup_deferred(): + return + stream = current_stream() # A small all_reduce for warmup. data = torch.zeros(1, device=device) diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 343edd171fe3..19fc15b27584 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -50,6 +50,7 @@ sync_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.kernel_warmup import kernel_warmup +from vllm.platforms import current_platform from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.worker.dp_utils import skip_dp_coordination @@ -283,7 +284,9 @@ def prepare_reconfiguration( self._prepare_eplb_communicator(get_standby_eplb_group()) if new_dp_size > old_dp_size: self.transfer_weights(old_dp_size, new_dp_size) - self._warm_target_groups(get_standby_dp_group(), standby_ep_group) + # On ROCm this deadlocks while serving; warm_and_capture warms these at commit. + if not current_platform.is_rocm(): + self._warm_target_groups(get_standby_dp_group(), standby_ep_group) if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel(): target_world_group = get_standby_world_group() assert target_world_group is not None @@ -577,7 +580,11 @@ def commit_scale_up(self, is_existing_worker: bool) -> None: mapping = self.receive_expert_mapping() self.worker.model_runner.setup_eplb_from_mapping(mapping) if not self._can_reuse_fused_moe_kernel(): + # warm_and_capture's dummy run warms the deferred dp/ep groups. self.warm_and_capture() + elif current_platform.is_rocm(): + # Reuse path skips warm_and_capture; warm the deferred groups here. + self._warm_target_groups(get_dp_group(), get_ep_group()) self._perform_eplb_reshuffle(async_op=True) if is_existing_worker: self._start_group_cleanup(retired_groups) @@ -590,6 +597,9 @@ def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: retired_groups = self.switch_and_prepare() if not self._can_reuse_fused_moe_kernel(): self.warm_and_capture() + elif current_platform.is_rocm(): + # Reuse path skips warm_and_capture; warm the deferred groups here. + self._warm_target_groups(get_dp_group(), get_ep_group()) self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: @@ -646,7 +656,8 @@ def prepare_new_worker(self) -> None: expert_weights=expert_weights, ) torch.accelerator.synchronize() - self._warm_target_groups(get_dp_group(), get_ep_group()) + if not current_platform.is_rocm(): + self._warm_target_groups(get_dp_group(), get_ep_group()) if self._can_reuse_fused_moe_kernel(): sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group()) diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py index 1892f3e79426..c9b12f447d9a 100644 --- a/vllm/distributed/elastic_ep/standby_state.py +++ b/vllm/distributed/elastic_ep/standby_state.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm from vllm.distributed.parallel_state import ( _init_stateless_group, _node_count, @@ -78,27 +79,35 @@ def create_standby_groups( ) standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0) standby_dp_ranks = [x.tolist() for x in standby_dp_ranks] - _STANDBY_DP = _init_stateless_group( - standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store - ) - standby_ep_ranks = ( - all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0) - ) - standby_ep_ranks = [x.tolist() for x in standby_ep_ranks] - _STANDBY_EP = _init_stateless_group( - standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all - ) + # Deferred to commit so the warm-up runs while the engine is paused. + with defer_comm_warmup_on_rocm(): + _STANDBY_DP = _init_stateless_group( + standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store + ) - if enable_eplb: - _STANDBY_EPLB = _init_stateless_group( + standby_ep_ranks = ( + all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0) + ) + standby_ep_ranks = [x.tolist() for x in standby_ep_ranks] + _STANDBY_EP = _init_stateless_group( standby_ep_ranks, - "eplb", + "ep", master_ip, backend, - coord_store=coord_store, + coord_store, + use_all2all=use_all2all, ) + if enable_eplb: + _STANDBY_EPLB = _init_stateless_group( + standby_ep_ranks, + "eplb", + master_ip, + backend, + coord_store=coord_store, + ) + def pop_standby_groups() -> dict: """Return all standby groups and clear the standby state.""" diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index b24ff64d4d92..0466779c3477 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1919,6 +1919,20 @@ def init_distributed_environment( _INNER_DP_WORLD = _WORLD +def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]: + """Defer communicator warm-up when joining a live elastic-EP cluster on ROCm. + + The existing ranks defer their matching warm-up in create_standby_groups(), + so a joining worker must too or the warm-up collective has no peers. Both + sides warm up at commit. + """ + from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm + + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: + return defer_comm_warmup_on_rocm() + return nullcontext() + + def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, @@ -2101,13 +2115,14 @@ def initialize_model_parallel( group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0) group_ranks = [x.tolist() for x in group_ranks] if enable_elastic_ep: - _DP = _init_stateless_group( - group_ranks, - "dp", - parallel_config.data_parallel_master_ip, - backend, - coord_store=coord_store, - ) + with _elastic_join_warmup_ctx(): + _DP = _init_stateless_group( + group_ranks, + "dp", + parallel_config.data_parallel_master_ip, + backend, + coord_store=coord_store, + ) else: _DP = init_model_parallel_group( group_ranks, get_world_group().local_rank, backend, group_name="dp" @@ -2130,14 +2145,15 @@ def initialize_model_parallel( group_ranks = [x.tolist() for x in group_ranks] use_all2all = parallel_config.use_all2all if enable_elastic_ep: - _EP = _init_stateless_group( - group_ranks, - "ep", - parallel_config.data_parallel_master_ip, - backend, - coord_store=coord_store, - use_all2all=use_all2all, - ) + with _elastic_join_warmup_ctx(): + _EP = _init_stateless_group( + group_ranks, + "ep", + parallel_config.data_parallel_master_ip, + backend, + coord_store=coord_store, + use_all2all=use_all2all, + ) else: _EP = init_model_parallel_group( group_ranks, @@ -2155,13 +2171,14 @@ def initialize_model_parallel( assert _EPLB is None, "EPLB group is already initialized" if config.parallel_config.enable_eplb: if enable_elastic_ep: - _EPLB = _init_stateless_group( - group_ranks, - "eplb", - parallel_config.data_parallel_master_ip, - backend, - coord_store=coord_store, - ) + with _elastic_join_warmup_ctx(): + _EPLB = _init_stateless_group( + group_ranks, + "eplb", + parallel_config.data_parallel_master_ip, + backend, + coord_store=coord_store, + ) else: _EPLB = init_model_parallel_group( group_ranks, From ebe1dec2da16ac559e40446a4c6e06a73c9efe00 Mon Sep 17 00:00:00 2001 From: PatrykSaffer Date: Sun, 13 Sep 2026 01:24:44 +0200 Subject: [PATCH 110/774] [PCP][DCP] Enable PCP+DCP on sparse-MLA models (#56157) Signed-off-by: patryk.saffer Signed-off-by: patryk.saffer Signed-off-by: PatrykSaffer Signed-off-by: Lucas Wilkinson Co-authored-by: patryk.saffer Co-authored-by: patryk.saffer Co-authored-by: Lucas Wilkinson Co-authored-by: Patryk Saffer Co-authored-by: Claude Fable 5.1 --- .buildkite/test_areas/lm_eval.yaml | 5 + .../GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml | 19 ++ tests/evals/gsm8k/configs/models-pcp.txt | 1 + .../v1/attention/test_indexer_dcp_localize.py | 35 +++ .../test_indexer_deepseek_v4_slot_mapping.py | 1 + .../v1/attention/test_sparse_mla_backends.py | 3 + tests/v1/worker/test_gpu_pcp_manager.py | 149 ++++++++++ .../layers/attention/mla_attention.py | 7 + .../layers/attention/sparse_mla_attention.py | 11 +- .../layers/sparse_attn_indexer.py | 94 +++++- .../layers/sparse_attn_indexer_kpool.py | 10 + vllm/models/deepseek_v32/attention.py | 31 +- vllm/v1/attention/backend.py | 17 +- .../attention/backends/mla/flashmla_sparse.py | 272 ++++++++++++++---- vllm/v1/attention/backends/mla/indexer.py | 230 +++++++++++++-- .../v1/attention/backends/mla/sparse_utils.py | 59 +++- vllm/v1/worker/gpu/attn_utils.py | 8 + vllm/v1/worker/gpu/input_batch.py | 3 + vllm/v1/worker/gpu/model_states/default.py | 2 + vllm/v1/worker/gpu/pcp_manager.py | 96 +++++-- vllm/v1/worker/gpu_model_runner.py | 2 +- 21 files changed, 917 insertions(+), 138 deletions(-) create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f33c5825548f..3c6561196610 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -110,12 +110,17 @@ steps: - csrc/ - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml + - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml - tests/evals/gsm8k/configs/models-pcp.txt - vllm/model_executor/layers/quantization - vllm/config/parallel.py - vllm/distributed/parallel_state.py - vllm/model_executor/layers/attention/mla_attention.py - vllm/model_executor/layers/attention/pcp.py + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/model_executor/layers/sparse_attn_indexer.py + - vllm/v1/attention/backends/mla/flashmla_sparse.py + - vllm/v1/attention/backends/mla/indexer.py - vllm/v1/worker/gpu/model_runner.py - vllm/v1/worker/gpu/pcp_manager.py autorun_on_main: true diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml new file mode 100644 index 000000000000..f57467ac7418 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -0,0 +1,19 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 100 +server_args: >- + --compilation-config '{"cudagraph_mode":"PIECEWISE"}' + --max-model-len 4096 + --max-num-batched-tokens 32768 + --safetensors-load-strategy prefetch + --moe-backend flashinfer_cutlass + --prefill-context-parallel-size 4 + --decode-context-parallel-size 4 + --enable-expert-parallel + --kv-cache-dtype fp8 +env: + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + VLLM_LOGGING_LEVEL: "DEBUG" + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/models-pcp.txt b/tests/evals/gsm8k/configs/models-pcp.txt index df645edad57a..e6d5bc9d39e9 100644 --- a/tests/evals/gsm8k/configs/models-pcp.txt +++ b/tests/evals/gsm8k/configs/models-pcp.txt @@ -1,2 +1,3 @@ GLM-5.2-NVFP4-TP2-PCP2-EP.yaml GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 7f7c20d53e96..ffcfcbfac1e5 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch @@ -962,3 +963,37 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp(): dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("dcp_world_size", [2, 4, 8]) +@pytest.mark.parametrize("req_lens", [[7], [8, 8], [1, 9], [256, 1, 2730]]) +def test_pcp_plan_deinterleave_restores_global_order(dcp_world_size, req_lens): + """The index gather must undo DCP sharding per request""" + from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan + + scheduled = np.array(req_lens, dtype=np.int64) + shard_rows = -(-scheduled // dcp_world_size) + rows = np.arange(len(scheduled)) + plan = build_pcp_global_chunk_plan( + rows, shard_rows, dcp_world_size, torch.device("cpu") + ) + + starts = np.concatenate([[0], np.cumsum(scheduled)]) + + # Build each rank's padded shard exactly as the cache gather would. + padded_cu = plan.padded_local_cu.tolist() + shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1) + for r in range(dcp_world_size): + for i, g in enumerate(req_lens): + for t in range(r, g, dcp_world_size): + shards[r, padded_cu[i] + t // dcp_world_size, 0] = starts[i] + t + + gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1) + restored = gathered[plan.deinterleave_idx] + # The layout is padded per request; only the real positions are read. + row_start = plan.row_start_cu.tolist() + for i, g in enumerate(req_lens): + torch.testing.assert_close( + restored[row_start[i] : row_start[i] + g, 0], + torch.arange(starts[i], starts[i] + g, dtype=torch.float32), + ) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index aabe88eeee5c..49462ea8fba6 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -212,6 +212,7 @@ def test_index_conversion_warmup_uses_physical_block_stride(): ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, + prefill_context_parallel_size=1, cp_kv_cache_interleave_size=1, ), ) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index e7239320c23a..0766584f5451 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1080,6 +1080,7 @@ def _convert_topk(indices, metadata, **kwargs): _convert_logical_to_physical_topk=_convert_topk, index_group=None, index_group_index=0, + pcp_dcp_kv_gather=False, ) out, _ = FlashMLASparseImpl._forward_bf16_kv( @@ -3257,6 +3258,7 @@ def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch): device=torch.device(DEVICE_TYPE), dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE), max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE), + pcp_dcp_kv_gather=False, ) query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660]) common_metadata = SimpleNamespace( @@ -3363,6 +3365,7 @@ def run_kernel(**kwargs): index_group=None, index_group_index=0, dcp_world_size=1, + pcp_dcp_kv_gather=False, need_to_return_lse_for_decode=False, _fp8_flash_mla_kernel=run_kernel, _convert_logical_to_physical_topk=( diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index 2d8941f6ffb3..dc935929d4c4 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -114,6 +114,155 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +def _rank_rows( + pcp_rank: int, + pcp_world_size: int, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + is_prefilling: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """One PCP+DCP rank's rows as (global request, extent).""" + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=pcp_rank, + device=torch.device("cpu"), + dcp_world_size=2, + ) + query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype( + np.int32 + ) + segments = manager._get_rank_segments( + pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np + ) + rows = np.array([segment.global_batch_req_idx for segment in segments]) + extents = (num_computed_tokens + num_scheduled_tokens)[rows] + return rows, extents + + +@pytest.mark.parametrize( + ("num_scheduled_tokens", "num_computed_tokens"), + [ + ([32, 32], [0, 0]), + ([32, 32, 32, 32], [0, 0, 0, 0]), + ([32, 32], [100, 0]), # one continued, one fresh + ([64, 32, 48], [0, 512, 0]), # ragged lengths and mixed contexts + # A replicated prefill BEHIND a split one. Sorting replicated rows first + # put request 1 ahead of request 0. + ([64, 3], [0, 0]), + ([64, 3, 64], [0, 0, 0]), + ], +) +@pytest.mark.parametrize("pcp_world_size", [2, 4]) +def test_published_row_order_is_identical_on_every_pcp_rank( + num_scheduled_tokens, num_computed_tokens, pcp_world_size +): + """Every rank must map its rows to the same global requests, in the same order.""" + num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32) + num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32) + is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_) + + orders = [ + _rank_rows( + rank, + pcp_world_size, + num_scheduled_tokens, + num_computed_tokens, + is_prefilling, + )[0] + for rank in range(pcp_world_size) + ] + for rank, order in enumerate(orders[1:], start=1): + assert np.array_equal(orders[0], order), ( + f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}" + ) + # Grouped by request and ascending: what the indexer plan indexes by. + assert np.all(np.diff(orders[0]) >= 0) + + +def test_published_row_order_puts_every_decode_before_every_prefill(): + """split_decodes_and_prefills takes the FIRST prefilling row as the boundary.""" + # req 0 is a continued prefill, req 1 is a decode. + req_idx, _ = _rank_rows( + pcp_rank=0, + pcp_world_size=2, + num_scheduled_tokens=np.array([32, 1], dtype=np.int32), + num_computed_tokens=np.array([100, 20], dtype=np.int32), + is_prefilling=np.array([True, False], dtype=np.bool_), + ) + # The decode (request 1) must be row 0. + assert req_idx[0] == 1 + + +def test_split_prefill_rows_repeat_the_request_and_its_full_extent(): + """Equal adjacent request indices are what let a backend share a KV region.""" + req_idx, extents = _rank_rows( + pcp_rank=1, + pcp_world_size=2, + num_scheduled_tokens=np.array([64, 3], dtype=np.int32), + num_computed_tokens=np.array([100, 0], dtype=np.int32), + is_prefilling=np.ones(2, dtype=np.bool_), + ) + assert req_idx.tolist() == [0, 0, 1] + assert extents.tolist() == [164, 164, 3] + + +@pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9]) +def test_dcp_replicates_prefills_too_short_to_split(query_len): + """A prefill that cannot fill 2*pcp chunks is replicated, not split.""" + pcp_world_size = 2 + num_scheduled_tokens = np.array([query_len], dtype=np.int32) + is_prefilling = np.ones(1, dtype=np.bool_) + + rows_per_rank = [] + for rank in range(pcp_world_size): + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=rank, + device=torch.device("cpu"), + dcp_world_size=2, + ) + assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0] + rows = list( + manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling) + ) + assert rows == [(0, 0, query_len)] + rows_per_rank.append(rows) + assert rows_per_rank[0] == rows_per_rank[1] + + +@pytest.mark.parametrize("pcp_world_size", [2, 4, 8]) +@pytest.mark.parametrize( + "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097] +) +def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len): + num_scheduled_tokens = np.array([query_len], dtype=np.int32) + is_prefilling = np.ones(1, dtype=np.bool_) + + charged_per_rank = [] + for rank in range(pcp_world_size): + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=rank, + device=torch.device("cpu"), + dcp_world_size=pcp_world_size, + ) + chunk_lens = [ + chunk_len + for _, _, chunk_len in manager._iter_rank_chunks( + rank, num_scheduled_tokens, is_prefilling + ) + ] + assert chunk_lens, f"rank {rank} got no rows for {query_len=}" + assert chunk_lens[0] == max(chunk_lens), ( + f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}" + ) + charged_per_rank.append(len(chunk_lens) * chunk_lens[0]) + + assert len(set(charged_per_rank)) == 1, ( + f"ranks would chunk differently for {query_len=}: {charged_per_rank}" + ) + + def _make_global_decode_batch( num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device ) -> InputBatch: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 21596eb04365..1c2e5b396263 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -420,6 +420,9 @@ class MLAAttention(nn.Module, AttentionLayerBase): """ supports_dense_mha_prefill: ClassVar[bool] = True + # Under PCP+DCP only the decode rows carry an LSE; the base forward + # merges a full-batch LSE, so subclasses opt in with their own forward. + supports_pcp_dcp: ClassVar[bool] = False def __init__( self, @@ -660,6 +663,10 @@ def __init__( _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup() + if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp: + raise NotImplementedError( + f"{type(self).__name__} does not support PCP+DCP." + ) self.dcp_manager: MLADCPManager | None = None if self.impl.dcp_world_size > 1: query_dtype = ( diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 5baddcd0a5cb..79ca61203255 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -190,7 +190,8 @@ def __init__( pin_memory=PIN_MEMORY, ) parallel_config = vllm_config.parallel_config - self.use_pcp = parallel_config.prefill_context_parallel_size > 1 + self.pcp_world_size = parallel_config.prefill_context_parallel_size + self.use_pcp = self.pcp_world_size > 1 try: self.dcp_world_size = get_dcp_group().world_size except AssertionError: @@ -409,8 +410,12 @@ def build( q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, - use_dense_mha=_use_dense_mha_prefill( - self.vllm_config, prefill_max_seq_len, self.topk_tokens + use_dense_mha=( + _use_dense_mha_prefill( + self.vllm_config, prefill_max_seq_len, self.topk_tokens + ) + # Dense MHA cannot be used under PCP+DCP. + and not (self.use_pcp and self.dcp_world_size > 1) ), topk_mask_workspace=self.topk_mask_workspace, host_staging_plan=staging_plan, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 817f8ad361fb..d663c54de367 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -3,6 +3,7 @@ """Custom Sparse Attention Indexer layers.""" import torch +import torch.distributed as dist import vllm.envs as envs from vllm import _custom_ops as ops @@ -130,6 +131,23 @@ def _merge_dcp_topk_global( ) +def dcp_gather_kv_rows( + local_padded: torch.Tensor, + deinterleave_idx: torch.Tensor, + gathered_buf: torch.Tensor, + out_buf: torch.Tensor, +) -> torch.Tensor: + """All-gather this rank's KV shard and return it in global token order.""" + + gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]] + dist.all_gather_into_tensor( + gathered, local_padded.contiguous(), group=get_dcp_group().device_group + ) + return torch.index_select( + gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]] + ) + + @triton.jit def _fused_indexer_q_rope_quant_kernel( positions, @@ -346,11 +364,19 @@ def sparse_attn_indexer( values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, fp8_dtype, use_fp4_cache ) - current_workspace_manager().get_simultaneous( + profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ values_spec, scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) + ] + if use_pcp and dcp_world_size > 1: + # The PCP+DCP path takes an all-gather destination and a + # de-interleaved result. + gather_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + profile_specs.extend(gather_spec * 2) + current_workspace_manager().get_simultaneous(*profile_specs) # Dummy allocation to simulate for peak logits tensor memory during inference. # FP8 elements so elements == bytes @@ -462,9 +488,27 @@ def sparse_attn_indexer( values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, fp8_dtype, use_fp4_cache ) - k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + # PCP + DCP needs two more pairs: the rank-major all-gather destination + # and the de-interleaved result. + pcp_chunks = [ + (c, c.pcp_deinterleave_idx) + for c in prefill_metadata.chunks + if c.pcp_deinterleave_idx is not None + ] + gather_specs: list[tuple[tuple[int, int], torch.dtype]] = [] + if pcp_chunks: + gathered_rows = max( + dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks + ) + deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks) + for rows in (gathered_rows, deinterleaved_rows): + gather_specs.extend( + _gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache) + ) + k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous( values_spec, scales_spec, + *gather_specs, ) for chunk in prefill_metadata.chunks: cu_seqlen_ks = chunk.cu_seqlen_ks @@ -481,6 +525,23 @@ def sparse_attn_indexer( chunk.local_cu_seq_lens, ) + # PCP + DCP KV all-gather. + deinterleave_idx = chunk.pcp_deinterleave_idx + if deinterleave_idx is not None: + # local_total_seq_lens is the PCP-padded extent here, so every + # rank contributes an identically shaped shard. + gathered_values, gathered_scales, out_values, out_scales = gather_bufs + shard = k_quant[: chunk.local_total_seq_lens] + k_quant = dcp_gather_kv_rows( + shard, deinterleave_idx, gathered_values, out_values + ) + k_scale = dcp_gather_kv_rows( + k_scale[: chunk.local_total_seq_lens], + deinterleave_idx, + gathered_scales, + out_scales, + ) + q_slice = q_quant[chunk.token_start : chunk.token_end] q_scale_slice = ( q_scale[chunk.token_start : chunk.token_end] @@ -491,7 +552,7 @@ def sparse_attn_indexer( chunk.token_start : chunk.token_end, :topk_tokens ] - if chunk.local_total_seq_lens == 0: + if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) else: @@ -561,15 +622,18 @@ def sparse_attn_indexer( topk_tokens, ) - _merge_dcp_topk_global( - logits, - topk_indices, - topk_tokens, - dcp_rank, - dcp_world_size, - cp_kv_cache_interleave_size, - row_starts=chunk.cu_seqlen_ks, - ) + if deinterleave_idx is None: + # Under the PCP path the top-k already ran over the whole + # context. + _merge_dcp_topk_global( + logits, + topk_indices, + topk_tokens, + dcp_rank, + dcp_world_size, + cp_kv_cache_interleave_size, + row_starts=chunk.cu_seqlen_ks, + ) if has_decode: decode_metadata = attn_metadata_narrowed.decode @@ -910,6 +974,10 @@ def forward_native( if current_platform.is_cuda() or current_platform.is_xpu(): return self.forward_cuda(hidden_states, q_quant, k, weights) elif current_platform.is_rocm(): + if self.use_pcp and self.dcp_world_size > 1: + raise NotImplementedError( + "The ROCm sparse-indexer path does not support PCP+DCP." + ) return self.forward_hip(hidden_states, q_quant, k, weights) elif current_platform.is_cpu(): return self.forward_cpu(hidden_states, q_quant, k, weights) diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py index ba7a5adff160..2a314042deef 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py +++ b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py @@ -919,6 +919,16 @@ def __init__( raise RuntimeError( "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." ) + _cfg = get_current_vllm_config_or_none() + _parallel = _cfg.parallel_config if _cfg is not None else None + if ( + _parallel is not None + and _parallel.prefill_context_parallel_size > 1 + and _parallel.decode_context_parallel_size > 1 + ): + raise NotImplementedError( + "SparseAttnIndexerKpool does not support PCP+DCP." + ) def forward_native( self, diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index b0fe554e7e10..00dd2e18a8d2 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, ClassVar, cast import torch import torch.nn as nn @@ -121,6 +121,7 @@ def __init__( class DeepseekV32Attention(MLAAttention): indexer: "DeepseekV32Indexer | None" indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer + supports_pcp_dcp: ClassVar[bool] = True def __init__( self, @@ -569,8 +570,9 @@ def _sparse_indexer_and_attn( seq_lens: torch.Tensor | None query_start_loc: torch.Tensor | None if self.use_pcp: - if attn_metadata.decode is not None: - seq_lens = attn_metadata.decode.seq_lens + decode_metadata = getattr(attn_metadata, "decode", None) + if decode_metadata is not None: + seq_lens = decode_metadata.seq_lens else: all_seq_lens = cast( torch.Tensor, @@ -585,12 +587,23 @@ def _sparse_indexer_and_attn( # PCP-only empty-shard metadata is needed. seq_lens = None query_start_loc = None - attn_out = self.dcp_manager.combine( - attn_out, - lse, - seq_lens=seq_lens, # type: ignore[arg-type] - query_start_loc=query_start_loc, # type: ignore[arg-type] - ) + # Under PCP+DCP the prefill rows attended over the gathered KV, so + # only the decode rows carry an LSE and take part in the merge. + num_merge_rows = lse.shape[0] + if num_merge_rows == attn_out.shape[0]: + attn_out = self.dcp_manager.combine( + attn_out, + lse, + seq_lens=seq_lens, # type: ignore[arg-type] + query_start_loc=query_start_loc, # type: ignore[arg-type] + ) + elif num_merge_rows > 0: + attn_out[:num_merge_rows] = self.dcp_manager.combine( + attn_out[:num_merge_rows], + lse, + seq_lens=seq_lens, # type: ignore[arg-type] + query_start_loc=query_start_loc, # type: ignore[arg-type] + ) if self.use_pcp: attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index addd68ec6ae0..c427d3a1de87 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -415,9 +415,14 @@ class CommonAttentionMetadata: encoder_seq_lens_cpu: np.ndarray | None = None dcp_local_seq_lens: torch.Tensor | None = None - dcp_local_seq_lens_cpu: torch.Tensor | None = None """Sequence lengths of the local rank in decode context parallelism world""" + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it + holds, on every row, the largest shard any DCP rank has of the row's + whole request, identical on every PCP rank, so the sparse backends pad + their KV gather to it.""" + positions: torch.Tensor | None = None """(num_actual_tokens,) token positions. Optional; set when the caller has positions available so that builders can pre-compute position-dependent @@ -434,6 +439,11 @@ class CommonAttentionMetadata: decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" + req_idx: np.ndarray | None = None + """(batch_size,) index of each row's request in the runner's request + table. Rows of one request are adjacent, so equal neighbours are PCP + chunks sharing one KV context.""" + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None """PrefixLM bidirectional ranges for multimodal tokens. Maps request index to list of (start, end) token position ranges @@ -525,9 +535,12 @@ def unpadded( encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu), dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), - dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), + dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs( + self.dcp_local_seq_lens_cpu_upper_bound + ), seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound), is_prefilling=maybe_slice_reqs(self.is_prefilling), + req_idx=maybe_slice_reqs(self.req_idx), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), ) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 4ecf5596832e..433cded5679d 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -3,11 +3,14 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar +import numpy as np import torch +import torch.distributed as dist from vllm import _custom_ops as ops from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType +from vllm.distributed.parallel_state import get_dcp_group from vllm.logger import init_logger from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, @@ -28,10 +31,12 @@ from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, + request_row_bounds, triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import ( + get_dcp_local_seq_lens, reshape_attn_output_for_spec_decode, reshape_query_for_spec_decode, split_prefill_chunks, @@ -43,6 +48,7 @@ get_mla_metadata, ) from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -171,6 +177,28 @@ def supports_combination( return None +def gathered_prefill_shards( + row_req_idx: np.ndarray, row_shard_rows: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Under PCP+DCP, group this rank's prefill rows by request and take each + request's slice of this rank's KV shard. + + PCP gives a rank two adjacent chunk rows of a split prefill; they share + one context, so the workspace holds it once. ``row_shard_rows`` is + ``CommonAttentionMetadata.dcp_local_seq_lens_cpu_upper_bound``: the + largest DCP shard of the whole request on every row, identical on every + PCP rank, so the all-gather is shaped the same everywhere. + + Returns the request row bounds and, per request, its shard rows. + """ + row_bounds = request_row_bounds(row_req_idx) + shard_rows = row_shard_rows[row_bounds[:-1]].astype(np.int64) + assert np.all(shard_rows > 0), ( + f"PCP+DCP prefill got an empty context: {shard_rows.tolist()}" + ) + return row_bounds, shard_rows + + @dataclass class FlashMLASparseMetadata(SparseMLACommonMetadata): @dataclass @@ -331,8 +359,13 @@ def __init__( device=device, ) + # PCP+DCP attends prefill rows over the DCP-gathered KV, which only the + # separate prefill/decode path can do. + self.pcp_dcp_kv_gather = self.use_pcp and self.dcp_world_size > 1 self.fp8_use_mixed_batch = ( - self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse + self.num_heads < MIN_HEADS_FOR_BF16_PREFILL + and not self.use_hisparse + and not self.pcp_dcp_kv_gather ) if parallel_config.decode_context_parallel_size > 1: @@ -342,7 +375,12 @@ def __init__( "default 'ag_rs' DCP comm backend; got " f"'{parallel_config.dcp_comm_backend}'" ) - if not self.fp8_use_mixed_batch: + if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": + raise NotImplementedError( + "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " + f"upconvert; got a {cache_config.cache_dtype} cache" + ) + if not self.fp8_use_mixed_batch and not self.pcp_dcp_kv_gather: raise NotImplementedError( "DCP for FlashMLA sparse is only supported on the " "mixed-batch fp8 path (num_heads < " @@ -353,9 +391,16 @@ def __init__( # Head padding (and the tile-scheduler metadata sized from it) is # computed from the local head count, but the kernel runs on the # DCP-gathered heads. - gathered_num_heads = ( - self.num_heads * parallel_config.decode_context_parallel_size - ) + if self.use_pcp: + gathered_num_heads = ( + self.num_heads * parallel_config.tensor_parallel_size + if self.dcp_world_size > self.pcp_world_size + else self.num_heads + ) + else: + gathered_num_heads = ( + self.num_heads * parallel_config.decode_context_parallel_size + ) gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads( gathered_num_heads ) @@ -433,6 +478,40 @@ def _build_fp8_separate_prefill_decode( query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] + # Chunk prefill requests to fit within workspace size + max_prefill_buffer_size = get_prefill_workspace_size( + self.vllm_config.model_config.max_model_len + ) + # Workspace rows per entry, and the prefill rows each entry covers. + workspace_rows = prefill_seq_lens_cpu + row_bounds = np.arange(num_prefills + 1) + prefill_block_table = common_attn_metadata.block_table_tensor[num_decodes:] + prefill_seq_lens = common_attn_metadata.seq_lens[num_decodes:] + if self.pcp_dcp_kv_gather: + # One entry per request, holding this rank's KV shard: the + # all-gather of every rank's shard is the whole context. + # A dummy batch bypasses the PCP manager: one row per request, + # so the row's own extent is the request's. + row_req_idx = common_attn_metadata.req_idx + if row_req_idx is None: + row_req_idx = np.arange(common_attn_metadata.num_reqs) + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound + if shard_rows_cpu is None: + shard_rows_cpu = get_dcp_local_seq_lens( + seq_lens_cpu, + self.dcp_world_size, + 0, + self.cp_kv_cache_interleave_size, + ) + row_bounds, rows_per_rank = gathered_prefill_shards( + row_req_idx[num_decodes:], shard_rows_cpu[num_decodes:].numpy() + ) + workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) + max_prefill_buffer_size //= self.dcp_world_size + entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device) + prefill_block_table = prefill_block_table.index_select(0, entry_rows) + prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) + num_entries = len(workspace_rows) # Build prefill_request_id: -1 for decode, request index for # prefill. This enables a single @@ -441,33 +520,28 @@ def _build_fp8_separate_prefill_decode( (num_tokens,), -1, dtype=torch.int32, device=self.device ) # Map prefill tokens to their request IDs (0, 1, 2, ...) - for req_idx in range(num_prefills): + for req_idx in range(num_entries): # Get query token range for this prefill request - global_req_idx = num_decodes + req_idx - req_query_start = query_start_loc_cpu[global_req_idx] - req_query_end = query_start_loc_cpu[global_req_idx + 1] + req_query_start = query_start_loc_cpu[ + num_decodes + int(row_bounds[req_idx]) + ] + req_query_end = query_start_loc_cpu[ + num_decodes + int(row_bounds[req_idx + 1]) + ] prefill_request_id[req_query_start:req_query_end] = req_idx # will be adjusted by chunk loop prefill_workspace_starts_cpu = torch.zeros( - num_prefills, dtype=torch.int32, pin_memory=True - ) - prefill_workspace_starts_cpu[1:] = torch.cumsum( - prefill_seq_lens_cpu[:-1], dim=0 + num_entries, dtype=torch.int32, pin_memory=True ) + prefill_workspace_starts_cpu[1:] = torch.cumsum(workspace_rows[:-1], dim=0) # populated by non-blocking copy after prefill_workspace_starts_cpu is # updated by each chunk prefill_workspace_starts = torch.empty( - num_prefills, dtype=torch.int32, device=self.device + num_entries, dtype=torch.int32, device=self.device ) - # Chunk prefill requests to fit within workspace size - max_prefill_buffer_size = get_prefill_workspace_size( - self.vllm_config.model_config.max_model_len - ) - chunk_bounds = split_prefill_chunks( - prefill_seq_lens_cpu, max_prefill_buffer_size - ) + chunk_bounds = split_prefill_chunks(workspace_rows, max_prefill_buffer_size) prefill_chunks = [] for chunk_start, chunk_end in chunk_bounds: @@ -480,19 +554,19 @@ def _build_fp8_separate_prefill_decode( offset = prefill_workspace_starts_cpu[chunk_start].item() prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset - chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum() - token_start = query_start_loc_cpu[num_decodes + chunk_start].item() - token_end = query_start_loc_cpu[num_decodes + chunk_end].item() + chunk_tot_seqlen = workspace_rows[chunk_start:chunk_end].sum() + token_start = query_start_loc_cpu[ + num_decodes + int(row_bounds[chunk_start]) + ].item() + token_end = query_start_loc_cpu[ + num_decodes + int(row_bounds[chunk_end]) + ].item() tokens_slice = slice(token_start, token_end) # Create chunk view of gpu tensor chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end] - chunk_block_table = common_attn_metadata.block_table_tensor[ - num_decodes + chunk_start : num_decodes + chunk_end - ] - chunk_seq_lens = common_attn_metadata.seq_lens[ - num_decodes + chunk_start : num_decodes + chunk_end - ] + chunk_block_table = prefill_block_table[chunk_start:chunk_end] + chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end] prefill_chunks.append( FP8Meta.Prefill.Chunk( @@ -516,14 +590,19 @@ def _build_fp8_separate_prefill_decode( ) if num_decodes > 0: - # Use padded head count since that's what the kernel will see - scheduler_metadata, _ = get_mla_metadata() + if self.pcp_dcp_kv_gather: + # Decode rows take the mixed-batch kernel call, which returns + # the LSE the DCP merge needs. + kernel_meta = self._build_fp8_mixed_decode_prefill() + else: + # Use padded head count since that's what the kernel will see + scheduler_metadata, _ = get_mla_metadata() - kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( - scheduler_metadata=scheduler_metadata, - dummy_block_table=self.dummy_block_table[:active_num_decodes], - cache_lens=self.max_model_len_tensor[:active_num_decodes], - ) + kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=scheduler_metadata, + dummy_block_table=self.dummy_block_table[:active_num_decodes], + cache_lens=self.max_model_len_tensor[:active_num_decodes], + ) fp8_metadata.decode = FP8Meta.Decode( seq_lens=common_attn_metadata.seq_lens[:active_num_decodes], kernel_metadata=kernel_meta, @@ -626,19 +705,35 @@ def __init__( "the bf16 sparse path is not supported under DCP." ) + self.pcp_dcp_kv_gather = False + self.gathered_kv_workspace: torch.Tensor | None = None if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None prefill_workspace_size = get_prefill_workspace_size( vllm_config.model_config.max_model_len ) - self.prefill_workspace_shape = (prefill_workspace_size, head_size) - self.q_concat_buffer, self.prefill_bf16_workspace = ( - current_workspace_manager().get_simultaneous( - (q_concat_shape, torch.bfloat16), - (self.prefill_workspace_shape, torch.bfloat16), - ) + parallel_config = vllm_config.parallel_config + self.pcp_dcp_kv_gather = ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 ) + shard_rows = prefill_workspace_size + if self.pcp_dcp_kv_gather: + # PCP+DCP upconverts this rank's KV shard, then all-gathers the + # shards into a workspace of the full prefill size. + shard_rows //= parallel_config.decode_context_parallel_size + self.prefill_workspace_shape = (shard_rows, head_size) + specs = [ + (q_concat_shape, torch.bfloat16), + (self.prefill_workspace_shape, torch.bfloat16), + ] + if self.pcp_dcp_kv_gather: + specs.append(((prefill_workspace_size, head_size), torch.bfloat16)) + buffers = current_workspace_manager().get_simultaneous(*specs) + self.q_concat_buffer, self.prefill_bf16_workspace = buffers[:2] + if self.pcp_dcp_kv_gather: + self.gathered_kv_workspace = buffers[2] else: (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( (q_concat_shape, torch.bfloat16), @@ -729,13 +824,49 @@ def _forward_bf16_kv( return attn_out, lse return torch.cat([decode_out, attn_out], dim=0), None + def _gather_prefill_chunk( + self, + chunk: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill.Chunk", + shard: torch.Tensor, + topk_indices: torch.Tensor, + attn_metadata: FlashMLASparseMetadata, + prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill", + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """All-gather this rank's upconverted KV shard so the chunk's rows attend + the whole context, and map their top-k onto the rank-major result.""" + shard_rows = int(chunk.chunk_tot_seqlen) + assert self.gathered_kv_workspace is not None + gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows] + dist.all_gather_into_tensor( + gathered_kv, shard, group=get_dcp_group().device_group + ) + + tokens = chunk.tokens_slice + topk_indices, topk_length = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[tokens], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + HAS_PREFILL_WORKSPACE=True, + prefill_workspace_request_ids=prefill_meta.request_ids[tokens], + prefill_workspace_starts=prefill_meta.workspace_starts, + prefill_workspace_rank_stride=shard_rows, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + return_valid_counts=True, + ) + return gathered_kv, topk_indices, topk_length + def _forward_fp8_kv_separate_prefill_decode( self, q: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """The lse covers the decode rows and is only returned under PCP+DCP.""" fp8_metadata = attn_metadata.fp8_extra_metadata assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) num_decodes = fp8_metadata.num_decodes @@ -783,8 +914,13 @@ def _forward_fp8_kv_separate_prefill_decode( # For BF16 cache: always use global cache slots (no workspace) # prefill_workspace_starts has been adjusted in-place per chunk so # prefill indices automatically come out chunk-local - topk_length = None - if num_prefill_tokens == 0 and not uses_host_cache: + # Under PCP+DCP the decode rows are converted by the mixed-batch call + # and the prefill rows per chunk, against the gathered workspace. + lse: torch.Tensor | None = None + topk_length: torch.Tensor | None = None + if self.pcp_dcp_kv_gather: + pass + elif num_prefill_tokens == 0 and not uses_host_cache: topk_indices, topk_length = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, @@ -812,6 +948,17 @@ def _fp8_decode( topk_indices: torch.Tensor, ) -> torch.Tensor: assert fp8_metadata.decode is not None + if self.pcp_dcp_kv_gather: + # Attend this rank's KV shard; the LSE feeds the DCP merge. + nonlocal lse + out, lse = self._forward_fp8_kv_mixed_batch( + q, + kv_c_and_k_pe_cache, + topk_indices, + attn_metadata, + kernel_metadata=fp8_metadata.decode.kernel_metadata, + ) + return out if uses_host_cache: return self._host_backed_fp8_decode( q, @@ -896,9 +1043,20 @@ def _fp8_decode( ) chunk_q = q[chunk.tokens_slice] - chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] - assert topk_length is not None - chunk_topk_length = topk_length[chunk.tokens_slice] + if self.pcp_dcp_kv_gather: + chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length = ( + self._gather_prefill_chunk( + chunk, + chunk_workspace, + topk_indices[chunk.tokens_slice], + attn_metadata, + fp8_metadata.prefill, + ) + ) + else: + chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] + assert topk_length is not None + chunk_topk_length = topk_length[chunk.tokens_slice] attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( chunk_q, @@ -907,7 +1065,10 @@ def _fp8_decode( chunk_topk_length, ) - return attn_out + if self.pcp_dcp_kv_gather and lse is None: + # No decode rows: the DCP merge still expects an LSE, of no rows. + lse = q.new_empty((0, q.shape[1]), dtype=torch.float32) + return attn_out, lse def _forward_fp8_kv_mixed_batch( self, @@ -915,6 +1076,7 @@ def _forward_fp8_kv_mixed_batch( kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, + kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata | None" = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Mixed batch FP8 forward path that treats all tokens as one batch. @@ -924,12 +1086,8 @@ def _forward_fp8_kv_mixed_batch( The lse is only returned when DCP needs it, otherwise None. """ - assert attn_metadata.fp8_extra_metadata is not None - assert isinstance( - attn_metadata.fp8_extra_metadata, - FlashMLASparseMetadata.FP8KernelMetadata, - ) - fp8_metadata = attn_metadata.fp8_extra_metadata + fp8_metadata = kernel_metadata or attn_metadata.fp8_extra_metadata + assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8KernelMetadata) block_table = attn_metadata.block_table # req_id_per_token covers the whole batch; slice it to the MQA tokens @@ -1160,7 +1318,7 @@ def forward_mqa( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) else: - attn_out = self._forward_fp8_kv_separate_prefill_decode( + attn_out, lse = self._forward_fp8_kv_separate_prefill_decode( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 7bc0b442c34b..164a3f89169f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import Any +import numpy as np import torch import vllm.envs as envs @@ -25,6 +26,7 @@ native_next_n_supported, ) from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -33,6 +35,7 @@ MultipleOf, ) from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping +from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, split_decodes_and_prefills, @@ -43,6 +46,7 @@ KVCacheSpec, MLAAttentionSpec, ) +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) @@ -259,6 +263,87 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [64 if current_platform.is_device_capability_family(90) else 128] +@dataclass(frozen=True) +class PCPGlobalChunkPlan: + """PCP packing for one indexer prefill chunk under PCP + DCP.""" + + # [num_reqs+1] cumsum of the padded per-request context (shard rows x W): + # the global layout. + row_start_cu: torch.Tensor + # [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its + # extent. + global_cu: torch.Tensor + # [num_reqs+1] cumsum of shard rows: this rank's padded layout. + padded_local_cu: torch.Tensor + padded_local_total: int + total: int + # [total] for each global row, its position in the rank-major gathered buffer. + deinterleave_idx: torch.Tensor + + +def build_pcp_global_chunk_plan( + row_req_idx: np.ndarray, + row_shard_rows: np.ndarray, + dcp_world_size: int, + device: torch.device, +) -> PCPGlobalChunkPlan: + """Plan the PCP packing for one chunk from its requests' KV shard rows. + + The global layout is padded to ``shard rows x W`` per request. Positions + past a request's true extent lie beyond every token's causal bound, so + the kernel never reads them. + """ + shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64) + assert row_req_idx.shape == shard_rows.shape + num_rows = len(shard_rows) + + row_bounds = request_row_bounds(row_req_idx) + region_first_row = row_bounds[:-1] + region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds)) + + region_padded = shard_rows[region_first_row] + assert np.all(region_padded > 0), ( + f"PCP+DCP prefill got an empty context: {region_padded.tolist()}" + ) + + # Pinned host buffers, filled in place so the copies below are async. + cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY) + row_start_cu, global_cu, padded_cu = cu + # Only a region's first row carries its extent. + first_row = torch.from_numpy(region_first_row) + row_padded = torch.zeros(num_rows, dtype=torch.int32) + row_padded[first_row] = torch.from_numpy(region_padded).int() + torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:]) + torch.cumsum(row_padded, 0, out=padded_cu[1:]) + row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)] + row_start_cu[num_rows] = global_cu[num_rows] + total = int(global_cu[num_rows]) + padded_total = int(padded_cu[num_rows]) + + idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY) + idx_np = idx.numpy() + region_start = global_cu[first_row].numpy() + region_padded_start = padded_cu[first_row].numpy() + for i in range(len(region_first_row)): + g = int(region_padded[i]) * dcp_world_size + t = np.arange(g, dtype=np.int64) + idx_np[region_start[i] : region_start[i] + g] = ( + (t % dcp_world_size) * padded_total + + region_padded_start[i] + + t // dcp_world_size + ) + + cu = async_copy_to_gpu(cu, device=device) + return PCPGlobalChunkPlan( + row_start_cu=cu[0], + global_cu=cu[1], + padded_local_cu=cu[2], + padded_local_total=padded_total, + total=total, + deinterleave_idx=async_copy_to_gpu(idx, device=device), + ) + + @dataclass class DeepseekV32IndexerPrefillChunkMetadata: block_table: torch.Tensor @@ -277,6 +362,8 @@ class DeepseekV32IndexerPrefillChunkMetadata: local_total_seq_lens: int = 0 max_local_total_seq_lens: int = 0 + pcp_deinterleave_idx: torch.Tensor | None = None + class BuildPrefillChunkMetadataKernel( VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"] @@ -960,6 +1047,48 @@ def _prepare_global_decode_seq_lens( self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0 return self.global_decode_seq_lens_buffer[:num_decode_tokens] + def _split_pcp_dcp_prefill_chunks( + self, + row_req_idx: np.ndarray, + row_shard_rows: np.ndarray, + row_query_lens_cpu: torch.Tensor, + max_logits_bytes: int, + request_offset: int, + ) -> list[tuple[slice, slice]]: + """Chunk by request rather than by row, so a split prefill's two rows + charge their shared context once, then widen each chunk back to rows. + + ``row_shard_rows`` holds the whole request's largest DCP shard on + every row, so the plan is identical on every PCP rank. + """ + row_bounds = request_row_bounds(row_req_idx) + first_rows = row_bounds[:-1] + # Each request's context, padded to whole DCP shards. + seq_lens = row_shard_rows[first_rows] * self.dcp_world_size + # Every rank holds a full-size chunk of a split request (shorter ones + # are replicated), so rows x longest row is the same on every rank, + # whichever row holds the short tail. + row_query_lens = row_query_lens_cpu.numpy() + query_lens = np.diff(row_bounds) * np.maximum.reduceat( + row_query_lens, first_rows + ) + chunk_specs = self._split_indexer_prefill_chunks( + torch.from_numpy(seq_lens.astype(np.int32)), + torch.from_numpy(query_lens.astype(np.int32)), + self.max_prefill_buffer_size, + max_logits_bytes, + ) + return [ + ( + slice( + request_offset + int(row_bounds[request_slice.start]), + request_offset + int(row_bounds[request_slice.stop]), + ), + query_slice, + ) + for request_slice, query_slice in chunk_specs + ] + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1094,16 +1223,52 @@ def build( # slice below). assert common_attn_metadata.seq_lens_cpu_upper_bound is not None seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound - chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], - prefill_query_lens_cpu, - self.max_prefill_buffer_size, - max_logits_bytes, - request_offset=num_decodes, - ) + req_idx = None + shard_rows = None + if self.use_pcp and self.dcp_world_size > 1: + # The gathered KV must be packed identically on every PCP rank: + # chunk by request from its DCP shard rows, which every rank + # holds. A dummy batch bypasses the PCP manager and has one + # row per request, so its own extent is the request's. + req_idx = common_attn_metadata.req_idx + if req_idx is None: + req_idx = np.arange(num_reqs) + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound + if shard_rows_cpu is None: + shard_rows_cpu = get_dcp_local_seq_lens( + seq_lens_cpu, + self.dcp_world_size, + 0, + self.cp_kv_cache_interleave_size, + ) + shard_rows = shard_rows_cpu.numpy() + chunk_specs = self._split_pcp_dcp_prefill_chunks( + req_idx[num_decodes:], + shard_rows[num_decodes:], + prefill_query_lens_cpu, + max_logits_bytes, + request_offset=num_decodes, + ) + else: + chunk_specs = self._split_indexer_prefill_chunks( + compressed_seq_lens_cpu[num_decodes:], + prefill_query_lens_cpu, + self.max_prefill_buffer_size, + max_logits_bytes, + request_offset=num_decodes, + ) chunks = [] for req_slice, query_slice in chunk_specs: + pcp_plan = None + if req_idx is not None: + assert shard_rows is not None + pcp_plan = build_pcp_global_chunk_plan( + req_idx[req_slice], + shard_rows[req_slice], + self.dcp_world_size, + self.device, + ) metadata = build_prefill_chunk_metadata( req_slice.start, req_slice.stop, @@ -1119,6 +1284,7 @@ def build( dcp_rank=self.dcp_rank, dcp_world_size=self.dcp_world_size, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + pcp_plan=pcp_plan, ) # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: @@ -1353,8 +1519,12 @@ def build_prefill_chunk_metadata( dcp_rank: int = 0, dcp_world_size: int = 1, cp_kv_cache_interleave_size: int = 1, + pcp_plan: PCPGlobalChunkPlan | None = None, ) -> DeepseekV32IndexerPrefillChunkMetadata | None: - total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item() + if pcp_plan is not None: + total_seq_lens = pcp_plan.total + else: + total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item() if total_seq_lens == 0: return None @@ -1362,15 +1532,22 @@ def build_prefill_chunk_metadata( device = block_table.device token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device) - cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) - # Assigning to slice avoids cpu sync. - cu_seq_lens[:1] = 0 - torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:]) + if pcp_plan is not None: + cu_seq_lens = pcp_plan.global_cu + else: + cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) + # Assigning to slice avoids cpu sync. + cu_seq_lens[:1] = 0 + torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:]) local_cu_seq_lens = cu_seq_lens local_total_seq_lens = total_seq_lens max_local_total_seq_lens = total_seq_lens - if dcp_world_size > 1: + if pcp_plan is not None: + local_cu_seq_lens = pcp_plan.padded_local_cu + local_total_seq_lens = pcp_plan.padded_local_total + max_local_total_seq_lens = pcp_plan.padded_local_total + elif dcp_world_size > 1: # Per-rank local KV length under interleave-aware DCP sharding, shape # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the # sharding matches the rest of the DCP pipeline (decode/prefill). @@ -1394,35 +1571,47 @@ def build_prefill_chunk_metadata( (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item() ) if query_slice is not None: - qs_start = query_slice.start - qs_stop = query_slice.stop + qs_start = min(query_slice.start, total_query_len) + qs_stop = min(query_slice.stop, total_query_len) else: qs_start = 0 qs_stop = total_query_len output_query_len = qs_stop - qs_start + assert 0 <= output_query_len <= total_query_len - qs_start, ( + f"query slice [{qs_start}, {qs_stop}) does not fit this rank's " + f"{total_query_len} tokens; the metadata tail would stay uninitialized" + ) cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device) cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device) - # Under DCP the kernel writes this rank's local row bounds into - # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. + if pcp_plan is not None: + row_start_cu = pcp_plan.row_start_cu + kernel_dcp_rank, kernel_dcp_world = 0, 1 + else: + # Under DCP the kernel writes this rank's local row bounds into + # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. + row_start_cu = local_cu_seq_lens + kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size _BUILD_PREFILL_CHUNK_METADATA_KERNEL( query_start_loc, uncompressed_seq_lens[start_idx:end_idx], cu_seq_lens, - local_cu_seq_lens, + row_start_cu, token_to_seq, cu_seq_len_ks, cu_seq_len_ke, qs_start, qs_stop, - dcp_rank, - dcp_world_size, + kernel_dcp_rank, + kernel_dcp_world, cp_kv_cache_interleave_size, num_reqs=num_reqs, COMPRESS_RATIO=compress_ratio, ) + pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None + token_start = query_start_loc_cpu[start_idx].item() if query_slice is not None: token_end = token_start + qs_stop @@ -1432,6 +1621,7 @@ def build_prefill_chunk_metadata( token_end = query_start_loc_cpu[end_idx].item() return DeepseekV32IndexerPrefillChunkMetadata( + pcp_deinterleave_idx=pcp_deinterleave_idx, cu_seqlen_ks=cu_seq_len_ks, cu_seqlen_ke=cu_seq_len_ke, cu_seq_lens=cu_seq_lens, diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index d9b45892a671..242c8f3b7843 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from typing import Any +import numpy as np import torch from vllm.config import VllmConfig @@ -24,6 +25,21 @@ from vllm.v1.worker.block_table import get_block_table_width +def request_row_bounds(req_idx: np.ndarray) -> np.ndarray: + """Bounds of the runs of adjacent rows that belong to one request: run + ``r`` is rows ``[bounds[r], bounds[r + 1])``. + + Under PCP a rank holds two adjacent chunk rows of a split prefill; the + sparse backends give such a run one KV region. + """ + assert req_idx.size > 0 + bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True]) + assert bounds.size - 1 == np.unique(req_idx).size, ( + "rows of one request must be adjacent" + ) + return bounds + + def flat_kv_row_view( kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] block_size: int, @@ -69,7 +85,7 @@ class CompileKey: block_table_stride: int @staticmethod - @triton.jit(do_not_specialize=["max_num_blocks_per_req"]) + @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"]) def kernel( req_id_ptr, # int32 [num_tokens] block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] @@ -78,6 +94,7 @@ def kernel( valid_count_ptr, # int32 [num_tokens] - output valid count per row prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr + workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace # shapes (compile-time where possible) max_num_blocks_per_req, BLOCK_SIZE: tl.constexpr, @@ -146,7 +163,7 @@ def kernel( # Guard block_table access valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 - is_invalid_tok |= ~valid_block | is_remote + is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill) base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0) out_val = base * BLOCK_STRIDE_ROWS + inblock_off @@ -155,7 +172,11 @@ def kernel( workspace_start = tl.load( workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 ) - prefill_out = workspace_start + tok + # Under DCP the prefill workspace is the all-gather of every rank's + # shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok. + prefill_out = ( + owning_rank * workspace_rank_stride + workspace_start + local_idx + ) out_val = tl.where(is_prefill, prefill_out, out_val) out_val = tl.where(is_invalid_tok, -1, out_val) @@ -238,6 +259,22 @@ def get_warmup_keys( block_size * dcp_size, ) max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size) + pcp_dcp_combos = [] + if ( + vllm_config.parallel_config.prefill_context_parallel_size > 1 + and dcp_size > 1 + ): + # PCP + DCP attends prefill rows over a DCP-gathered workspace. + pcp_dcp_combos.append( + dict( + HAS_PREFILL_WORKSPACE=True, + COUNT_VALID=True, + COMPACT_TO_FRONT=True, + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, + DCP_INTERLEAVE=dcp_interleave, + ) + ) return self._trace_dispatch(self.dispatch)( zip_inputs( dict( @@ -280,6 +317,7 @@ def get_warmup_keys( DCP_RANK=dcp_rank, DCP_INTERLEAVE=dcp_interleave, ), + *pcp_dcp_combos, ), BLOCK_SIZE=block_size, BLOCK_STRIDE_ROWS=block_stride_rows, @@ -316,6 +354,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: int32_ptr if compile_key.has_prefill_workspace else None ), max_num_blocks_per_req=1, + workspace_rank_stride=1, BLOCK_SIZE=compile_key.block_size, BLOCK_STRIDE_ROWS=compile_key.block_stride_rows, BLOCK_N=compile_key.block_n, @@ -339,6 +378,7 @@ def __call__( prefill_workspace_request_ids: torch.Tensor | None, prefill_workspace_starts: torch.Tensor | None, *, + workspace_rank_stride: int, max_num_blocks_per_req: int, BLOCK_SIZE: int, BLOCK_STRIDE_ROWS: int, @@ -406,6 +446,10 @@ def triton_convert_req_index_to_global_index( HAS_PREFILL_WORKSPACE: bool = False, prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, + prefill_workspace_rank_stride: int | None = None, + dcp_size: int = 1, + dcp_rank: int = 0, + cp_kv_cache_interleave_size: int = 1, return_valid_counts: bool = False, out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, @@ -511,6 +555,7 @@ def triton_convert_req_index_to_global_index( prefill_workspace_request_ids, prefill_workspace_starts, # shapes / constexprs + workspace_rank_stride=prefill_workspace_rank_stride or 0, max_num_blocks_per_req=max_num_blocks_per_req, BLOCK_SIZE=BLOCK_SIZE, BLOCK_STRIDE_ROWS=( @@ -521,10 +566,9 @@ def triton_convert_req_index_to_global_index( HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE, COUNT_VALID=return_valid_counts, COMPACT_TO_FRONT=return_valid_counts, - # DCP disabled (no-op de-interleave) - DCP_SIZE=1, - DCP_RANK=0, - DCP_INTERLEAVE=1, + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, + DCP_INTERLEAVE=cp_kv_cache_interleave_size, ) if return_valid_counts: @@ -618,6 +662,7 @@ def triton_filter_and_convert_dcp_index( # No prefill workspace on the DCP decode path. None, None, + workspace_rank_stride=0, max_num_blocks_per_req=max_num_blocks_per_req, BLOCK_SIZE=BLOCK_SIZE, BLOCK_STRIDE_ROWS=( diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 7e7a212b5fc5..a3f6541b0866 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -400,6 +400,7 @@ def build_attn_metadata( kv_cache_config: KVCacheConfig, seq_lens_cpu_upper_bound: torch.Tensor | None = None, dcp_local_seq_lens: torch.Tensor | None = None, + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None, positions: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None, @@ -409,10 +410,15 @@ def build_attn_metadata( rswa_prefix_lens: torch.Tensor | None = None, ubatch_idx: int = 0, fast_prefill: FastPrefillBatchMetadata | None = None, + req_idx: np.ndarray | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs] + if dcp_local_seq_lens_cpu_upper_bound is not None: + dcp_local_seq_lens_cpu_upper_bound = dcp_local_seq_lens_cpu_upper_bound[ + :num_reqs + ] if seq_lens_cpu_upper_bound is not None: seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] @@ -457,10 +463,12 @@ def build_attn_metadata( slot_mapping=slot_mapping, causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound, positions=positions, is_prefilling=group_is_prefilling, mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, + req_idx=req_idx, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 30b9b8089c04..5cbd642d2529 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -114,6 +114,9 @@ class InputBatch: # (cudagraph capture) batches, which run the KV-sharing layers in full. fast_prefill: "FastPrefillBatchMetadata | None" = None + # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata). + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 6baceac9bc52..27a9591e4be0 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -220,11 +220,13 @@ def prepare_attn( kv_cache_config=kv_cache_config, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + dcp_local_seq_lens_cpu_upper_bound=input_batch.dcp_local_seq_lens_cpu_upper_bound, positions=input_batch.positions, is_prefilling=torch.from_numpy(input_batch.is_prefilling_np), mm_req_doc_ranges=req_doc_ranges, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, + req_idx=input_batch.idx_mapping_np, ubatch_idx=ubatch_idx, fast_prefill=input_batch.fast_prefill, ) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 6b97254621ec..3a51532f57a1 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -6,10 +6,10 @@ import numpy as np import torch -from vllm.config import VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger -from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import ( @@ -154,30 +154,48 @@ def validate_config( "MRV2 PCP only supports DSpark or single-module MTP " "speculative decoding." ) - if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk") + if parallel_config.decode_context_parallel_size > 1 and not is_sparse_mla: + # Dense MLA prefill sizes its DCP KV gather from each rank's own + # chunk rows, so the ranks' collectives diverge (#53573). + raise NotImplementedError("MRV2 PCP + DCP supports sparse MLA models only.") + if ( + is_sparse_mla + and parallel_config.decode_context_parallel_size == 1 + and cudagraph_mode != CUDAGraphMode.NONE + ): + raise NotImplementedError( + "MRV2 sparse MLA PCP does not support CUDA graphs yet. " + "Set -cc.cudagraph_mode=NONE." + ) + if cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + parallel_config.decode_context_parallel_size > 1 + and parallel_config.dcp_comm_backend != "ag_rs" + ): + raise NotImplementedError( + "MRV2 PCP + DCP requires dcp_comm_backend='ag_rs'; got " + f"'{parallel_config.dcp_comm_backend}'." + ) @staticmethod def _reorder_segments( segments: list[RankSegment], - num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, - query_start_loc_np: np.ndarray, ) -> list[RankSegment]: - """Move pure prefills last to match the batch ordering expected by - attention backends like MLA and sparse MLA. - """ + """Order this rank's rows decodes-first, then prefills, canonically.""" - def is_pure_prefill(segment: RankSegment) -> bool: + def sort_key(segment: RankSegment) -> tuple[bool, int, int]: req_idx = segment.global_batch_req_idx - start_pos = ( - num_computed_tokens[req_idx] - + segment.global_batch_slice.start - - query_start_loc_np[req_idx] + return ( + bool(is_prefilling[req_idx]), + req_idx, + segment.global_batch_slice.start, ) - return is_prefilling[req_idx] and start_pos == 0 - segments.sort(key=is_pure_prefill) + segments.sort(key=sort_key) rank_offset = 0 for index, segment in enumerate(segments): segments[index] = replace( @@ -189,6 +207,21 @@ def is_pure_prefill(segment: RankSegment) -> bool: rank_offset += segment.num_tokens return segments + def replicated_requests( + self, + num_scheduled_tokens: np.ndarray, + is_prefilling: np.ndarray, + ) -> np.ndarray: + """Per global request, whether every PCP rank gets the whole query.""" + num_chunks = 2 * self.pcp_world_size + query_lens = np.asarray(num_scheduled_tokens, dtype=np.int64) + replicated = ~np.asarray(is_prefilling, dtype=np.bool_) + if self.dcp_world_size > 1: + chunk_sizes = (query_lens + num_chunks - 1) // num_chunks + drops_a_chunk = (num_chunks - 1) * chunk_sizes >= query_lens + replicated |= drops_a_chunk + return replicated + def _iter_rank_chunks( self, rank: int, @@ -204,17 +237,20 @@ def _iter_rank_chunks( rank 1: 1 6 rank 2: 2 5 rank 3: 3 4 + + Decodes, and prefills too short to fill all eight, are replicated instead. """ num_chunks = 2 * self.pcp_world_size + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling) for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens): query_len = int(num_tokens) if query_len == 0: continue chunk_indices: tuple[int, ...] - if bool(is_prefilling[global_batch_req_idx]): + if not replicated[global_batch_req_idx]: chunk_size = (query_len + num_chunks - 1) // num_chunks chunk_indices = (rank, num_chunks - 1 - rank) - else: # decodes are replicated + else: # decodes, and short prefills under DCP, are replicated chunk_size = query_len chunk_indices = (0,) @@ -229,7 +265,6 @@ def _get_rank_segments( self, rank: int, num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, query_start_loc_np: np.ndarray, ) -> list[RankSegment]: @@ -248,12 +283,7 @@ def _get_rank_segments( ) ) rank_offset += chunk_len - return self._reorder_segments( - rank_segments, - num_computed_tokens, - is_prefilling, - query_start_loc_np, - ) + return self._reorder_segments(rank_segments, is_prefilling) def _build_batch_layout( self, @@ -263,13 +293,13 @@ def _build_batch_layout( query_start_loc_np: np.ndarray, padded_num_tokens: int | None = None, ) -> tuple[list[list[RankSegment]], list[int]]: + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling) segments_by_rank = [] per_rank_num_tokens = [] for rank in range(self.pcp_world_size): segments = self._get_rank_segments( rank, num_scheduled_tokens, - num_computed_tokens, is_prefilling, query_start_loc_np, ) @@ -309,7 +339,7 @@ def _build_batch_layout( dtype=np.int64, ) # Cache insertion pairs one slot entry with each rank's local decode. - if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0: + if replicated[segment.global_batch_req_idx] and rank != 0: continue gathered_kv_write_mask[padded_gathered_slice] = True hidden_restore_idx[segment.global_batch_slice] = np.arange( @@ -527,6 +557,19 @@ def partition_batch( ) seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens + dcp_local_seq_lens_cpu_upper_bound = None + if self.dcp_world_size > 1: + # The largest DCP shard of each row's whole request, identical on + # every PCP rank: the sparse backends pad their KV gather to it. + request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ + local_to_global_batch_req_idx_np + ] + dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + torch.from_numpy(request_seq_lens.astype(np.int32)), + self.dcp_world_size, + 0, + self.cp_interleave, + ) self._local_batch = replace( input_batch, @@ -549,6 +592,7 @@ def partition_batch( seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound, num_computed_tokens_np=local_start_pos_np, prefill_len_np=local_prefill_len_np, num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4ff50f82b515..d6a27b1f8ba9 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2486,7 +2486,7 @@ def _get_block_table(kv_cache_gid: int): self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded) cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded] - cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[ + cm_base.dcp_local_seq_lens_cpu_upper_bound = self.dcp_local_seq_lens.cpu[ :num_reqs_padded ] From e19a3e172ecc1d1ece8ecce9dfd91bf811a814a0 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sat, 12 Sep 2026 19:48:25 -0400 Subject: [PATCH 111/774] [5/N] Share HiSparse host cache across TP ranks (reopens #52760) (#56629) Signed-off-by: Matthew Bonanni Signed-off-by: Chao Lei Signed-off-by: Lucas Wilkinson Signed-off-by: Lucas Wilkinson Co-authored-by: Chao Lei Co-authored-by: Claude Fable 5 Co-authored-by: Matthew Bonanni Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex --- docs/design/hisparse.md | 10 + tests/v1/core/test_kv_cache_utils.py | 26 +- .../unit/offloading_connector/test_config.py | 36 +++ .../cpu/test_shared_offload_region.py | 130 +++++++-- tests/v1/worker/test_attn_utils.py | 82 ++++++ .../worker/test_kv_cache_allocation_scope.py | 6 +- tests/v1/worker/test_utils.py | 253 ++++++++++++++++++ .../kv_connector/v1/hisparse/worker.py | 26 +- .../kv_connector/v1/offloading/config.py | 18 +- vllm/v1/hisparse/binding.py | 71 +++-- vllm/v1/hisparse/layout.py | 19 +- vllm/v1/hisparse/runtime.py | 195 +++++++++++++- vllm/v1/kv_cache_interface.py | 6 + .../kv_offload/cpu/shared_offload_region.py | 85 ++++-- 14 files changed, 878 insertions(+), 85 deletions(-) diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md index 965ab81d8da1..53a309371783 100644 --- a/docs/design/hisparse.md +++ b/docs/design/hisparse.md @@ -60,6 +60,16 @@ sees only device pools. The source group has `block_pool_id=None`; device-pool consumers must narrow it before indexing, so host ownership cannot masquerade as a numeric GPU pool. +For single-node MP tensor parallelism, every TP worker maps the same pinned host +pool and uses the same block and layer offsets. MLA source KV is replicated +across TP ranks, so this stores one physical copy instead of one copy per rank. +TP rank 0 writes the shared host pool; peers wait on its IPC events before +reading it. Other executor and parallel layouts retain private per-rank pools. + +The shared layout backs the per-replica logical capacity with one physical pool; +the private layout allocates one physical pool per rank. Physical pool size +includes block-stride alignment. + ## Code boundary ```text diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index e7eee03e7bff..3b29799d3f89 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,6 +3,7 @@ import copy import hashlib import importlib +import mmap import subprocess import sys from collections.abc import Callable @@ -15,6 +16,7 @@ import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module from vllm.config import ( CacheConfig, KVTransferConfig, @@ -90,7 +92,13 @@ @pytest.mark.parametrize("gpu_block_size", [32, 64]) -def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( block_size=gpu_block_size, @@ -116,7 +124,15 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): hf_config=SimpleNamespace(index_topk=128), max_model_len=gpu_block_size, ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), cache_config=SimpleNamespace( num_gpu_blocks_override=7, prefix_cache_retention_interval=None, @@ -137,6 +153,12 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): assert cache_config.hisparse_host_num_blocks > 7 host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride assert host_group.host_resident assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) host_layers = set(host_group.layer_names) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 10464d950dde..f597c8fa27c8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -367,6 +367,42 @@ def test_hisparse_offloads_only_indexer_group(): assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1] +def test_hisparse_partial_group_size_survives_scheduler_flattening(): + """Worker and scheduler representations must produce identical mmap rows.""" + layer_specs = {name: _full_attention_spec() for name in ("indexer.0", "indexer.1")} + wrapped = UniformTypeKVCacheSpecs.from_specs(layer_specs) + assert wrapped is not None + source = KVCacheGroupSpec( + ["source"], + _full_attention_spec(), + host_resident=True, + role=KVCacheGroupRole.HISPARSE_SOURCE, + ) + worker_indexer = KVCacheGroupSpec( + list(layer_specs), wrapped, role=KVCacheGroupRole.HISPARSE_INDEXER + ) + scheduler_indexer = KVCacheGroupSpec( + list(layer_specs), + next(iter(layer_specs.values())), + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + + def make_config(indexer: KVCacheGroupSpec) -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[source, indexer], + hisparse_host_num_blocks=4, + ) + + config = _make_vllm_config() + worker = build_offloading_config(config, make_config(worker_indexer)) + scheduler = build_offloading_config(config, make_config(scheduler_indexer)) + + assert worker.worker_kv_bytes_per_block == scheduler.worker_kv_bytes_per_block + assert worker.worker_kv_bytes_per_block == wrapped.page_size_bytes + + def test_zero_blocks_skips_tensor_layout_validation(): kv_cache_config = _make_sizing_kv_cache_config(packed=False) kv_cache_config.num_blocks = 0 diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index c7ddac4feffa..c448e76a9c51 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -10,11 +10,12 @@ import threading import time import uuid -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest from vllm.utils.system_utils import get_mp_context +from vllm.v1.kv_offload.cpu import shared_offload_region as region_module from vllm.v1.kv_offload.cpu.shared_offload_region import ( SharedOffloadRegion, _wait_for_file_size, @@ -192,28 +193,43 @@ def _mp_barrier_construct_and_hold( barrier, fill_value: int, done_queue, + replicated: bool = False, ) -> None: """Construct with a real cross-process barrier, write, then hold the mapping (no cleanup) until the parent SIGKILLs this process.""" try: + populate_calls = 0 + get_populate_write_fn = region_module._get_populate_write_fn + + def track_population(mmap_obj): + nonlocal populate_calls + populate_calls += 1 + return get_populate_write_fn(mmap_obj) + + region_module._get_populate_write_fn = track_population region = SharedOffloadRegion( engine_id=engine_id, num_chunks=2, - rank=rank, - kv_bytes_per_chunk=num_workers * PAGE_SIZE, + rank=0 if replicated else rank, + kv_bytes_per_chunk=PAGE_SIZE if replicated else num_workers * PAGE_SIZE, cpu_page_size=PAGE_SIZE, barrier=lambda: barrier.wait(30), + populate_only_on_creator=replicated, ) # The constructor's barrier precedes the creator's unlink. barrier.wait(30) t = region.create_next_worker_view(PAGE_SIZE) - t[:, :] = fill_value + if not replicated or rank == 0: + t[:, :] = fill_value + barrier.wait(30) done_queue.put( { "rank": rank, "error": None, "inode": os.fstat(region.fd).st_ino, "path_exists": os.path.exists(region.mmap_path), + "populate_calls": populate_calls, + "contents_match": bool((t == (1 if replicated else fill_value)).all()), } ) time.sleep(60) # hold the mapping until killed @@ -752,6 +768,26 @@ def test_cleanup_idempotent(iid): r.cleanup() # must be a no-op +def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): + """cleanup() must release every chunk from a split host registration.""" + r = _make_region(iid) + cudart = MagicMock() + cudart.cudaHostUnregister.return_value = MagicMock(value=0) + monkeypatch.setattr(region_module, "current_platform", MagicMock()) + region_module.current_platform.is_cuda_alike.return_value = True + monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart) + r.pinned_addresses = [0x100000, 0x200000, 0x300000] + r.is_pinned = True + + r.cleanup() + + assert cudart.cudaHostUnregister.call_args_list == [ + call(0x300000), + call(0x200000), + call(0x100000), + ] + + def test_cleanup_after_create_next_worker_view_releases_mmap(iid): """cleanup() must close the mmap even after create_next_worker_view was called. create_next_worker_view returns a view that shares storage with _base; both must be @@ -811,11 +847,53 @@ def test_wait_for_file_size_timeout(tmp_path): os.close(fd) +@pytest.mark.skip_global_cleanup +def test_wait_for_file_size_rejects_unlinked_file(tmp_path): + path = tmp_path / "unlinked.mmap" + fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600) + try: + os.unlink(path) + with pytest.raises(RuntimeError, match="creator failed"): + _wait_for_file_size(fd, PAGE_SIZE, timeout=1.0) + finally: + os.close(fd) + + # --------------------------------------------------------------------------- # Constructor — capacity validation # --------------------------------------------------------------------------- +@pytest.mark.skip_global_cleanup +@pytest.mark.skipif(not os.path.isdir("/dev/shm"), reason="requires /dev/shm") +def test_creator_memory_check_runs_only_for_creator(iid): + checked_sizes: list[int] = [] + creator = SharedOffloadRegion( + engine_id=iid, + num_chunks=4, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + creator_memory_check=checked_sizes.append, + ) + joiner: SharedOffloadRegion | None = None + try: + joiner = SharedOffloadRegion( + engine_id=iid, + num_chunks=4, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + creator_memory_check=checked_sizes.append, + ) + assert checked_sizes == [4 * PAGE_SIZE] + finally: + if joiner is not None: + joiner.cleanup() + creator.cleanup() + _cleanup_file(creator.mmap_path) + + def test_insufficient_space_raises_clear_error(monkeypatch): """A failed creator capacity check must clean up and give a clear error.""" import vllm.v1.kv_offload.cpu.shared_offload_region as region @@ -916,7 +994,8 @@ def test_barrier_failure_unlinks_creator_and_raises(iid): assert not os.path.exists(path) -def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): +@pytest.mark.parametrize("replicated", [False, True]) +def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): """With a real cross-process barrier every worker maps one shared inode, the name is gone while they run, and SIGKILL leaks nothing.""" num_workers = 2 @@ -927,7 +1006,7 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): procs = [ ctx.Process( target=_mp_barrier_construct_and_hold, - args=(iid, rank, num_workers, barrier, rank + 1, done_queue), + args=(iid, rank, num_workers, barrier, rank + 1, done_queue, replicated), ) for rank in range(num_workers) ] @@ -939,6 +1018,10 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): assert r["error"] is None, f"rank {r['rank']}: {r['error']}" assert len({r["inode"] for r in results}) == 1, "workers split onto two files" assert not any(r["path_exists"] for r in results) + assert all(r["contents_match"] for r in results) + assert sum(r["populate_calls"] for r in results) == ( + 1 if replicated else num_workers + ) assert not os.path.exists(path) finally: for p in procs: @@ -950,23 +1033,36 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): assert restarted._creator is True, "restart must be able to create anew" -def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch): +@pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"]) +def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase): """A worker that dies before the rendezvous must still arrive at the barrier, or its peers block in the collective until it times out.""" - import vllm.v1.kv_offload.cpu.shared_offload_region as region - - monkeypatch.setattr( - region, - "check_shm_free_space", - MagicMock(side_effect=RuntimeError("Insufficient space")), + failure = MagicMock(side_effect=ValueError("Insufficient space")) + memory_check = failure if failure_phase == "memory" else None + if failure_phase == "shm": + monkeypatch.setattr(region_module, "check_shm_free_space", failure) + elif failure_phase == "population": + monkeypatch.setattr(region_module, "_get_populate_write_fn", failure) + path = f"/dev/shm/vllm_offload_{iid}.mmap" + seen_at_barrier = [] + barrier = MagicMock( + side_effect=lambda: seen_at_barrier.append(os.path.exists(path)) ) - barrier = MagicMock() - with pytest.raises(RuntimeError, match="Insufficient space"): - _make_region(iid, barrier=barrier) + with pytest.raises(ValueError, match="Insufficient space"): + SharedOffloadRegion( + engine_id=iid, + num_chunks=1, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + barrier=barrier, + creator_memory_check=memory_check, + populate_only_on_creator=True, + ) barrier.assert_called_once_with() - assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap") + assert seen_at_barrier == [False] def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index dc75ba4ea19b..2dadc8c980e6 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -12,6 +12,7 @@ import pytest import torch +import vllm.v1.hisparse.binding as attn_utils_module from tests.v1.attention.utils import dense_kv_cache_views from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -239,6 +240,87 @@ def check( assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set() +class _FakeSharedHostRegion: + def __init__(self) -> None: + self.cleanup_calls = 0 + self.base_tensor = torch.empty(1, dtype=torch.int8) + + def cleanup(self) -> None: + self.cleanup_calls += 1 + + +def test_profiling_cleanup_releases_tp_shared_region_once(monkeypatch): + """TP-shared profiling pools must use region-aware chunk cleanup.""" + region = _FakeSharedHostRegion() + runtime = SimpleNamespace( + _host_cache=object(), + registered_host_pool=region.base_tensor, + hot_backing=object(), + shared_host_region=region, + ) + forward_context = { + "layer": SimpleNamespace( + hisparse_cache=SimpleNamespace(runtime=runtime), + ) + } + released = [] + + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): + released.append((runtimes, pinned_host_pools, shared_host_region)) + + monkeypatch.setattr( + attn_utils_module, + "release_pinned_state", + release_pinned_state, + ) + + attn_utils_module.release_hisparse_profiling_cache(forward_context) + + assert released == [([runtime], [], region)] + + +@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"]) +def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): + """A failure after mmap allocation must not leak the shared registration.""" + region = _FakeSharedHostRegion() + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC + ), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), + ) + + def allocate(*args): + args[-1].shared_region = region + if failure_phase == "allocation": + raise RuntimeError("initialization failed") + return {} + + def bind(**kwargs): + if failure_phase == "binding": + raise RuntimeError("initialization failed") + return [] + + def buffers(*args, **kwargs): + raise RuntimeError("initialization failed") + + monkeypatch.setattr(attn_utils_module, "allocate_hisparse_kv_caches", allocate) + monkeypatch.setattr(attn_utils_module, "bind_hisparse_kv_caches", bind) + monkeypatch.setattr( + attn_utils_module, "initialize_hisparse_runtime_buffers", buffers + ) + with pytest.raises(RuntimeError, match="initialization failed"): + attn_utils_module.init_hisparse_kv_cache( + SimpleNamespace(), + torch.device("cpu"), + [], + vllm_config, + {}, + SimpleNamespace(), + ) + assert region.cleanup_calls == 1 + + def test_reshape_padded_kv_cache_strides_by_padded_page(): num_blocks = 3 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index bacd1cad8482..7f4bea52c227 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -60,7 +60,11 @@ def bind_hisparse(**kwargs): ) ] - monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool) + monkeypatch.setattr( + hisparse_binding, + "HiSparseHostPool", + lambda vllm_config, kv_cache_config: host_pool, + ) monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 8422a591c358..55cab287fd26 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap from collections import deque from contextlib import nullcontext from types import SimpleNamespace @@ -188,6 +189,232 @@ def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary( assert handle.submit_layer_mirror.call_count == expected_calls +def _hisparse_parallel_config(**overrides): + values = { + "tensor_parallel_size": 2, + "pipeline_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "world_size": 2, + "distributed_executor_backend": "mp", + "nnodes_within_dp": 1, + "data_parallel_index": 3, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + config = SimpleNamespace(parallel_config=_hisparse_parallel_config()) + + assert hisparse_runtime_module.use_shared_hisparse_host_pool(config) + + unsupported = ( + {"tensor_parallel_size": 1, "world_size": 1}, + {"pipeline_parallel_size": 2, "world_size": 4}, + {"prefill_context_parallel_size": 2, "world_size": 4}, + {"decode_context_parallel_size": 2}, + {"world_size": 4}, + {"distributed_executor_backend": "ray"}, + {"nnodes_within_dp": 2}, + ) + for overrides in unsupported: + config.parallel_config = _hisparse_parallel_config(**overrides) + assert not hisparse_runtime_module.use_shared_hisparse_host_pool(config) + + +@pytest.mark.skip_global_cleanup +def test_hisparse_shared_host_pool_uses_one_replicated_mmap(monkeypatch): + page = mmap.PAGESIZE + tp_group = MagicMock() + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group) + + class FakeSharedOffloadRegion: + BLOCK_SIZE_ALIGNMENT = page + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.total_size_bytes = kwargs["num_chunks"] * kwargs["kv_bytes_per_chunk"] + self.base_tensor = torch.empty(self.total_size_bytes, dtype=torch.int8) + self.is_pinned = False + self.pinned_addresses = [] + self.view_sizes = [] + self.offset = 0 + + def create_next_canonical_view(self, size): + self.view_sizes.append(size) + view = torch.as_strided( + self.base_tensor, + size=(self.kwargs["num_chunks"], size), + stride=(self.kwargs["kv_bytes_per_chunk"], 1), + storage_offset=self.offset, + ) + self.offset += size + return view + + def cleanup(self): + pass + + monkeypatch.setattr( + hisparse_runtime_module, "SharedOffloadRegion", FakeSharedOffloadRegion + ) + registration_ranges = MagicMock(return_value=((0, page), (page, 4 * page))) + monkeypatch.setattr( + hisparse_runtime_module, + "_hisparse_registration_ranges", + registration_ranges, + ) + pinned: list[torch.Tensor] = [] + + def pin_tensor(tensor): + pinned.append(tensor) + + monkeypatch.setattr(hisparse_runtime_module, "pin_tensor", pin_tensor) + config = SimpleNamespace( + instance_id="instance", + parallel_config=_hisparse_parallel_config( + tensor_parallel_size=1, + world_size=1, + distributed_executor_backend="uni", + ), + ) + + pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools( + config, + [24, 40], + num_blocks=4, + host_block_stride=page, + use_shared_host_pool=True, + ) + + assert region is not None + assert private_pools == [] + assert region.kwargs == { + "engine_id": "hisparse_instance_dp3", + "num_chunks": 1, + "rank": 0, + "kv_bytes_per_chunk": 4 * page, + "cpu_page_size": 64, + "barrier": tp_group.barrier, + "creator_memory_check": hisparse_runtime_module.check_hisparse_host_memory, + "populate_only_on_creator": True, + } + assert region.view_sizes == [24, 40] + assert [pool.shape for pool in pools] == [(24,), (40,)] + registration_ranges.assert_called_once_with([24, 40], 4, page) + assert [ + (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes) + for tensor in pinned + ] == [(0, page), (page, 3 * page)] + assert region.is_pinned + + +@pytest.mark.parametrize("fail_registration", [None, 1, 2]) +def test_shared_host_pool_tracks_successful_registrations( + monkeypatch, fail_registration +): + """Only successful registrations may be unregistered on allocation failure.""" + page = mmap.PAGESIZE + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", MagicMock()) + backing = torch.empty(3 * page, dtype=torch.uint8) + region = MagicMock(base_tensor=backing, pinned_addresses=[], is_pinned=False) + monkeypatch.setattr( + hisparse_runtime_module, "SharedOffloadRegion", lambda **kw: region + ) + monkeypatch.setattr( + hisparse_runtime_module, + "_hisparse_registration_ranges", + lambda *args: ((0, page), (page, 3 * page)), + ) + cudart = MagicMock() + cudart.cudaHostRegister.side_effect = [ + SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2) + ] + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + config = SimpleNamespace( + instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0) + ) + context = ( + pytest.raises(RuntimeError, match="cudaHostRegister failed") + if fail_registration + else nullcontext() + ) + with context: + hisparse_runtime_module.allocate_hisparse_host_pools( + config, [3 * page], 3, page, use_shared_host_pool=True + ) + successful = 2 if fail_registration is None else fail_registration - 1 + assert region.pinned_addresses == [ + backing.data_ptr() + i * page for i in range(successful) + ] + assert region.is_pinned == bool(successful) + assert region.cleanup.call_count == int(fail_registration is not None) + attempts = 2 if fail_registration is None else fail_registration + assert ( + cudart.cudaHostRegister.call_args_list + == [ + call(backing.data_ptr(), page, 0), + call(backing.data_ptr() + page, 2 * page, 0), + ][:attempts] + ) + + +def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch): + """Aliased tensor configs must not count the shared backing once per view.""" + region = SimpleNamespace(base_tensor=torch.empty(64, dtype=torch.int8)) + allocate = MagicMock(return_value=([], [], region)) + monkeypatch.setattr( + hisparse_runtime_module, "allocate_hisparse_host_pools", allocate + ) + config = SimpleNamespace( + hisparse_shared_host_pool=True, + hisparse_host_num_blocks=4, + hisparse_host_block_stride=16, + kv_cache_tensors=[ + SimpleNamespace( + host_resident=True, + layers=["a", "b"], + offset=0, + layer_stride=12, + size=44, + ), + SimpleNamespace( + host_resident=True, layers=["c"], offset=24, layer_stride=20, size=44 + ), + ], + ) + vllm_config = SimpleNamespace() + pool = hisparse_runtime_module.HiSparseHostPool(vllm_config, config) + backing = pool.allocate(44) + allocate.assert_called_once_with( + vllm_config, [12, 12, 20], 4, 16, use_shared_host_pool=True + ) + assert backing.numel() == 44 + assert backing.data_ptr() == region.base_tensor.data_ptr() + assert pool.registered is region.base_tensor + + +def test_hisparse_registration_chunks_end_between_host_blocks(): + """Registration seams must not bisect a DMA-addressable host block.""" + page = mmap.PAGESIZE + ranges = hisparse_runtime_module._hisparse_registration_ranges( + tensor_sizes=[4 * 3 * page, 4 * 5 * page], + num_blocks=4, + host_block_stride=8 * page, + max_chunk_bytes=10 * page, + ) + + assert ranges == ( + (0, 9 * page), + (9 * page, 17 * page), + (17 * page, 27 * page), + (27 * page, 32 * page), + ) + + def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): backing = torch.full((10, 2, 3), -1, dtype=torch.float32) cache = backing[1:9] @@ -931,6 +1158,32 @@ def make_cache_handle(is_leader: bool): assert streams == [] +def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): + worker = object.__new__(HiSparseConnectorWorker) + worker._initialized = True + worker._slot_mapping_staging = None + worker.dma_stream = None + worker.cache_handles = [] + worker.pinned_host_pools = [] + worker.shared_host_region = object() + released = False + + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): + nonlocal released + assert runtimes == [] + assert pinned_host_pools == [] + assert shared_host_region is worker.shared_host_region + released = True + + monkeypatch.setattr( + hisparse_worker_module, "release_pinned_state", release_pinned_state + ) + + worker.shutdown() + + assert released + + class _TestReplaySSMMixer(MambaMixer2): def __init__(self): torch.nn.Module.__init__(self) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index 69721803be79..b7ffe9ec97f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -147,6 +147,12 @@ def _select_written_row_mirrors( ) +def _is_hisparse_host_writer( + shared_host_region: SharedOffloadRegion | None, +) -> bool: + return shared_host_region is None or get_tensor_model_parallel_rank() == 0 + + def _create_hisparse_host_events( shared_host_region: SharedOffloadRegion | None, is_host_writer: bool, @@ -200,15 +206,26 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: raise RuntimeError("HiSparse connector found no hot-cache handles.") hot_backings: dict[int, torch.Tensor] = {} registered_host_pools: dict[int, torch.Tensor] = {} + shared_host_regions: dict[int, SharedOffloadRegion] = {} for cache in cache_handles: hot_backing = cache.runtime.hot_backing hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing registered_pool = cache.runtime.registered_host_pool registered_host_pools[registered_pool.data_ptr()] = registered_pool + if (region := cache.runtime.shared_host_region) is not None: + shared_host_regions[id(region)] = region if len(hot_backings) != 1: raise RuntimeError("HiSparse hot tensors must share one GPU backing.") hot_backing = next(iter(hot_backings.values())) - pinned_host_pools = list(registered_host_pools.values()) + if len(shared_host_regions) > 1: + raise RuntimeError("HiSparse caches must share one host region.") + shared_host_region = next(iter(shared_host_regions.values()), None) + pinned_host_pools = ( + [] + if shared_host_region is not None + else list(registered_host_pools.values()) + ) + is_host_writer = _is_hisparse_host_writer(shared_host_region) resident = cache_handles[0].view assert resident is not None @@ -223,11 +240,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: host_num_blocks, hot_backing.device, pinned_host_pools, + shared_host_region=shared_host_region, + is_host_writer=is_host_writer, ) except Exception: release_pinned_state( [cache.runtime for cache in cache_handles], pinned_host_pools, + shared_host_region, ) raise @@ -866,6 +886,8 @@ def shutdown(self) -> None: if self.dma_stream is not None: self.dma_stream.synchronize() release_pinned_state( - [cache.runtime for cache in self.cache_handles], self.pinned_host_pools + [cache.runtime for cache in self.cache_handles], + self.pinned_host_pools, + self.shared_host_region, ) self._initialized = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 42b9dd6aa392..c4e6d0f2a2ed 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -17,6 +17,7 @@ MLAAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, + UniformTypeKVCacheSpecs, iter_layer_specs, ) from vllm.v1.kv_offload.config import ( @@ -29,7 +30,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: @@ -42,6 +43,19 @@ def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ... ) +def _group_kv_bytes_per_block(group: "KVCacheGroupSpec") -> int: + """Return the physical bytes occupied by one block of a cache group. + + Worker configs may retain ``UniformTypeKVCacheSpecs`` while scheduler + configs flatten that wrapper to one representative per-layer spec. Keep + the result invariant across those two representations. + """ + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.page_size_bytes + return spec.page_size_bytes * len(group.layer_names) + + def build_offloading_config( vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig", @@ -133,7 +147,7 @@ def build_offloading_config( worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks elif kv_cache_config.num_blocks > 0: worker_kv_bytes_per_block = sum( - group.kv_cache_spec.page_size_bytes for _, group in selected_groups + _group_kv_bytes_per_block(group) for _, group in selected_groups ) single_group_spec = ( diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 263e553d32c4..c23abc9cb4b8 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -120,27 +120,32 @@ def init_hisparse_kv_cache( block_tables: "BlockTables", ) -> dict[str, torch.Tensor]: """Allocate and bind HiSparse caches within the caller's allocation context.""" - host_pool = HiSparseHostPool() - kv_caches = allocate_hisparse_kv_caches( - kv_cache_config, - device, - vllm_config.cache_config.get_resolved_kv_cache_layout(), - kernel_block_sizes, - host_pool, - ) - cache_handles = bind_hisparse_kv_caches( - forward_context=forward_context, - kv_cache_config=kv_cache_config, - kv_caches=kv_caches, - block_tables=block_tables, - host_pool=host_pool, - ) - initialize_hisparse_runtime_buffers( - cache_handles, - max_num_reqs=vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, - ) - return kv_caches + host_pool = HiSparseHostPool(vllm_config, kv_cache_config) + try: + kv_caches = allocate_hisparse_kv_caches( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + host_pool, + ) + cache_handles = bind_hisparse_kv_caches( + forward_context=forward_context, + kv_cache_config=kv_cache_config, + kv_caches=kv_caches, + block_tables=block_tables, + host_pool=host_pool, + ) + initialize_hisparse_runtime_buffers( + cache_handles, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, + ) + return kv_caches + except Exception: + if host_pool.shared_region is not None: + host_pool.shared_region.cleanup() + raise def _get_hisparse_cache( @@ -166,13 +171,24 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: if not runtimes: return - registered_pools = list( - { - runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool - for runtime in runtimes.values() - }.values() + shared_regions = { + id(runtime.shared_host_region): runtime.shared_host_region + for runtime in runtimes.values() + if runtime.shared_host_region is not None + } + assert len(shared_regions) <= 1 + shared_region = next(iter(shared_regions.values()), None) + registered_pools = ( + [] + if shared_region is not None + else list( + { + runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool + for runtime in runtimes.values() + }.values() + ) ) - release_pinned_state(list(runtimes.values()), registered_pools) + release_pinned_state(list(runtimes.values()), registered_pools, shared_region) for cache in cache_handles: cache.mirror_staging_cache = None cache.mirror_staging_slots = None @@ -257,6 +273,7 @@ def bind_hisparse_kv_caches( assert source_cache.untyped_storage().data_ptr() == ( host_pool.registered.untyped_storage().data_ptr() ) + cache_handle.runtime.shared_host_region = host_pool.shared_region cache_handle.runtime.bind_source_cache( source_cache, registered_host_pool=host_pool.registered, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index d74b0e8cfd49..b3b45d3334a5 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -9,7 +9,11 @@ from vllm.config.kv_transfer import hisparse_host_pool_gib from vllm.logger import init_logger from vllm.utils.math_utils import cdiv -from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig +from vllm.v1.hisparse.runtime import ( + ResolvedHiSparseConfig, + get_hisparse_host_block_stride, + use_shared_hisparse_host_pool, +) from vllm.v1.kv_cache_interface import ( HiSparseHotSpec, HiSparseResidentSpec, @@ -36,6 +40,8 @@ class HiSparseLayout: source_group: KVCacheGroupSpec device_groups: list[KVCacheGroupSpec] host_num_blocks: int + host_block_stride: int + shared_host_pool: bool def get_hisparse_kv_cache_groups( @@ -221,9 +227,12 @@ def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None: regular_groups = groups[1:] gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups] - host_num_blocks = host_budget // sum( - spec.page_size_bytes for spec in source_specs.values() + shared_host_pool = use_shared_hisparse_host_pool(vllm_config) + host_block_stride = get_hisparse_host_block_stride( + sum(spec.page_size_bytes for spec in source_specs.values()), + use_shared_host_pool=shared_host_pool, ) + host_num_blocks = host_budget // host_block_stride if host_num_blocks <= 0: raise ValueError("HiSparse has no allocatable host blocks.") @@ -231,6 +240,8 @@ def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None: source_group=source_group, device_groups=gpu_groups, host_num_blocks=host_num_blocks, + host_block_stride=host_block_stride, + shared_host_pool=shared_host_pool, ) @@ -339,6 +350,8 @@ def get_hisparse_kv_cache_config( kv_cache_tensors=kv_cache_tensors, kv_cache_groups=[*host_groups, *device_groups], hisparse_host_num_blocks=hisparse_layout.host_num_blocks, + hisparse_host_block_stride=hisparse_layout.host_block_stride, + hisparse_shared_host_pool=hisparse_layout.shared_host_pool, prefix_cache_retention_interval=( vllm_config.cache_config.prefix_cache_retention_interval ), diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 3dd8c017f5b5..314d3ba86b7a 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -4,6 +4,8 @@ from __future__ import annotations +import math +import mmap import time from collections.abc import Callable from dataclasses import dataclass @@ -13,17 +15,22 @@ import torch from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed import get_tp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import current_stream +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor logger = init_logger(__name__) +HOST_REGISTER_CHUNK_BYTES = 256 * 2**30 + if TYPE_CHECKING: from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup + from vllm.v1.kv_cache_interface import KVCacheConfig # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. FP8_DS_MLA_ROW_BYTES = 656 @@ -185,22 +192,195 @@ def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: class HiSparseHostPool: """Keep the single host backing and its exact CUDA registration range.""" - def __init__(self) -> None: + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: + self.vllm_config = vllm_config + self.kv_cache_config = kv_cache_config self.backing: torch.Tensor | None = None self.registered: torch.Tensor | None = None + self.shared_region: SharedOffloadRegion | None = None def allocate(self, size: int) -> torch.Tensor: assert self.backing is None, "HiSparse host pool is already allocated." - check_hisparse_host_memory(size) - self.backing, self.registered = allocate_pinned_host_pool(size) + config = self.kv_cache_config + if config.hisparse_shared_host_pool: + tensor_sizes = [ + size + for _, size in sorted( + (tensor.offset + i * tensor.layer_stride, tensor.layer_stride) + for tensor in config.kv_cache_tensors + if tensor.host_resident + for i in range(len(tensor.layers)) + ) + ] + assert config.hisparse_host_num_blocks is not None + assert config.hisparse_host_block_stride is not None + _, _, self.shared_region = allocate_hisparse_host_pools( + self.vllm_config, + tensor_sizes, + config.hisparse_host_num_blocks, + config.hisparse_host_block_stride, + use_shared_host_pool=True, + ) + assert self.shared_region is not None + self.registered = self.shared_region.base_tensor + self.backing = self.registered[:size] + else: + check_hisparse_host_memory(size) + self.backing, self.registered = allocate_pinned_host_pool(size) return self.backing +def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool: + """Whether replicated MLA host KV can share one local mmap.""" + parallel = vllm_config.parallel_config + return ( + current_platform.is_cuda_alike() + and parallel.tensor_parallel_size > 1 + and parallel.pipeline_parallel_size == 1 + and parallel.prefill_context_parallel_size == 1 + and parallel.decode_context_parallel_size == 1 + and parallel.world_size == parallel.tensor_parallel_size + and parallel.distributed_executor_backend == "mp" + and parallel.nnodes_within_dp == 1 + ) + + +def get_hisparse_host_block_stride( + page_size_bytes: int, *, use_shared_host_pool: bool +) -> int: + if use_shared_host_pool: + return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT) + return page_size_bytes + + +def _hisparse_registration_ranges( + tensor_sizes: list[int], + num_blocks: int, + host_block_stride: int, + max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES, +) -> tuple[tuple[int, int], ...]: + """Split a shared pool without bisecting any tensor's host block. + + CUDA batch copies reject descriptors spanning adjacent registrations. + """ + total_size = num_blocks * host_block_stride + if total_size % mmap.PAGESIZE: + raise ValueError("HiSparse host pool must be OS-page aligned.") + if max_chunk_bytes < mmap.PAGESIZE: + raise ValueError("HiSparse registration chunks must span at least one page.") + + tensors: list[tuple[int, int, int]] = [] + tensor_start = 0 + for tensor_size in tensor_sizes: + if tensor_size % num_blocks: + raise ValueError( + f"HiSparse host tensor size {tensor_size} is not divisible by " + f"{num_blocks} blocks." + ) + tensor_end = tensor_start + tensor_size + tensors.append((tensor_start, tensor_end, tensor_size // num_blocks)) + tensor_start = tensor_end + if tensor_start > total_size: + raise ValueError("HiSparse host tensors exceed the shared host pool.") + + def last_safe_boundary(lower: int, upper: int) -> int | None: + best = None + for start, end, block_bytes in tensors: + first_block = max(0, (lower - start) // block_bytes + 1) + last_block = min(num_blocks, (min(upper, end) - start) // block_bytes) + if first_block > last_block: + continue + divisor = math.gcd(block_bytes, mmap.PAGESIZE) + if start % divisor: + continue + period = mmap.PAGESIZE // divisor + if period == 1: + residue = 0 + else: + # Find block boundaries that are also OS-page boundaries. + residue = ( + (-start // divisor) * pow(block_bytes // divisor, -1, period) + ) % period + block = last_block - (last_block - residue) % period + if block >= first_block: + candidate = start + block * block_bytes + best = candidate if best is None else max(best, candidate) + + if upper >= tensor_start: + candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE + if candidate >= tensor_start and candidate > lower: + best = candidate if best is None else max(best, candidate) + return best + + ranges: list[tuple[int, int]] = [] + start = 0 + while total_size - start > max_chunk_bytes: + end = last_safe_boundary(start, start + max_chunk_bytes) + if end is None: + raise ValueError( + "A HiSparse host block is too large for the registration chunk limit." + ) + ranges.append((start, end)) + start = end + ranges.append((start, total_size)) + return tuple(ranges) + + +def allocate_hisparse_host_pools( + vllm_config: VllmConfig, + tensor_sizes: list[int], + num_blocks: int, + host_block_stride: int, + *, + use_shared_host_pool: bool, +) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]: + """Allocate private host tensors or one mmap shared by local TP ranks.""" + if not use_shared_host_pool: + check_hisparse_host_memory(sum(tensor_sizes)) + private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes] + return ( + [pool for pool, _ in private_pools], + [registered for _, registered in private_pools], + None, + ) + + region = SharedOffloadRegion( + engine_id=( + f"hisparse_{vllm_config.instance_id}_" + f"dp{vllm_config.parallel_config.data_parallel_index}" + ), + num_chunks=1, + rank=0, + kv_bytes_per_chunk=num_blocks * host_block_stride, + cpu_page_size=sum(tensor_sizes), + barrier=get_tp_group().barrier, + creator_memory_check=check_hisparse_host_memory, + populate_only_on_creator=True, + ) + try: + for start, end in _hisparse_registration_ranges( + tensor_sizes, num_blocks, host_block_stride + ): + tensor = region.base_tensor[start:end] + pin_tensor(tensor) + region.pinned_addresses.append(tensor.data_ptr()) + region.is_pinned = True + pools = [ + region.create_next_canonical_view(size).view(-1) for size in tensor_sizes + ] + except Exception: + region.cleanup() + raise + return pools, [], region + + def release_pinned_state( - runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor] + runtimes: list[HiSparseRuntime], + pinned_host_pools: list[torch.Tensor], + shared_host_region: SharedOffloadRegion | None = None, ) -> None: """Synchronize and release registered host KV pools.""" - if pinned_host_pools: + if pinned_host_pools or shared_host_region is not None: try: torch.accelerator.synchronize() except RuntimeError as e: @@ -208,7 +388,7 @@ def release_pinned_state( "HiSparse: CUDA context unusable at teardown (%s); leaving " "%d host-pool tensors pinned for kernel exit reclaim.", e, - len(pinned_host_pools), + len(pinned_host_pools) + int(shared_host_region is not None), ) return @@ -240,6 +420,8 @@ def release_pinned_state( del runtime._host_cache del runtime.registered_host_pool del runtime.hot_backing + if shared_host_region is not None: + shared_host_region.cleanup() def hisparse_prefill_staging_remap( @@ -536,6 +718,7 @@ def __init__( self.eager_host_mirror = config.eager_host_mirror self.resident_source_index = -1 self.request_state_indices: torch.Tensor | None = None + self.shared_host_region: SharedOffloadRegion | None = None @property def host_cache(self) -> torch.Tensor: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index fb839bcd6855..e2869e047f54 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1435,6 +1435,12 @@ class KVCacheConfig: hisparse_host_num_blocks: int | None = None """Capacity of the dedicated HiSparse host-block manager, when enabled.""" + hisparse_host_block_stride: int | None = None + """Physical bytes between consecutive HiSparse host blocks.""" + + hisparse_shared_host_pool: bool = False + """Whether local TP ranks share one physical HiSparse host pool.""" + @cached_property def transfer_group_ids(self) -> tuple[int, ...]: """IDs of cache groups that participate in external KV transfer.""" diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 1a5dde21e62d..46a1e7ae6460 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import errno import mmap import os @@ -27,6 +28,10 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N while True: if os.fstat(fd).st_size >= expected_size: return + if os.fstat(fd).st_nlink == 0: + raise RuntimeError( + "Shared offload region creator failed during initialization." + ) if time.monotonic() > deadline: raise TimeoutError( f"Timed out waiting for mmap file to reach {expected_size} bytes" @@ -75,6 +80,9 @@ class SharedOffloadRegion: given, the path is unlinked once every worker has mapped the file, so the kernel reclaims the memory when the last worker exits, no matter how it exits; mappings taken before the unlink stay valid. + + Creator-only population pre-faults the entire region before the barrier + and requires that barrier to keep joiners from using unpopulated pages. """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE @@ -87,7 +95,12 @@ def __init__( kv_bytes_per_chunk: int, cpu_page_size: int, barrier: Callable[[], None] | None = None, + *, + creator_memory_check: Callable[[int], None] | None = None, + populate_only_on_creator: bool = False, ) -> None: + if populate_only_on_creator and barrier is None: + raise ValueError("Creator-only population requires a barrier.") self.page_size = mmap.PAGESIZE assert kv_bytes_per_chunk % self.page_size == 0 @@ -112,25 +125,18 @@ def __init__( # Joiner path — another worker won O_EXCL. Reopen and wait # for the file to reach expected size. self.fd = os.open(self.mmap_path, os.O_RDWR) - try: - _wait_for_file_size(self.fd, self.total_size_bytes) - except (TimeoutError, OSError): - os.close(self.fd) - raise + _wait_for_file_size(self.fd, self.total_size_bytes) logger.info("Opened existing mmap file %s", self.mmap_path) else: # Creator path. We won O_EXCL, so we own the file: any # failure here must clean up so concurrent joiners don't # land on a 0-byte stub and spin in _wait_for_file_size # for the full 30 s timeout. - try: - check_shm_free_space(self.total_size_bytes) - os.ftruncate(self.fd, self.total_size_bytes) - except (RuntimeError, OSError): - os.unlink(self.mmap_path) - os.close(self.fd) - raise self._creator = True + if creator_memory_check is not None: + creator_memory_check(self.total_size_bytes) + check_shm_free_space(self.total_size_bytes) + os.ftruncate(self.fd, self.total_size_bytes) logger.info( "Created mmap file %s (%.2f GB)", self.mmap_path, @@ -143,10 +149,21 @@ def __init__( flags=mmap.MAP_SHARED, prot=mmap.PROT_READ | mmap.PROT_WRITE, ) + + if populate_only_on_creator and self._creator: + populate_write_fn = _get_populate_write_fn(self.mmap_obj) + populate_write_fn(self.mmap_obj, 0, self.total_size_bytes) except Exception: if self._creator: - os.unlink(self.mmap_path) + with contextlib.suppress(FileNotFoundError): + os.unlink(self.mmap_path) self._creator = False + if hasattr(self, "mmap_obj") and self.mmap_obj is not None: + self.mmap_obj.close() + self.mmap_obj = None + if hasattr(self, "fd") and self.fd is not None: + os.close(self.fd) + self.fd = None # Peers block inside the barrier until the collective times out if # we die before reaching it. Arrive anyway so every worker calls # barrier() exactly once and they fail on their own errors instead @@ -169,16 +186,28 @@ def __init__( barrier() except Exception: if self._creator: - os.unlink(self.mmap_path) + with contextlib.suppress(FileNotFoundError): + os.unlink(self.mmap_path) self._creator = False self.mmap_obj.close() os.close(self.fd) + self.mmap_obj = None + self.fd = None raise if self._creator: os.unlink(self.mmap_path) self._creator = False logger.info("Unlinked mmap file %s", self.mmap_path) + self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) + self._views: list[torch.Tensor] = [] + self._canonical_offset = 0 + self.is_pinned: bool = False + self.pinned_addresses: list[int] = [] + + if populate_only_on_creator: + return + populate_write_fn = _get_populate_write_fn(self.mmap_obj) if rank is not None: @@ -205,10 +234,11 @@ def __init__( "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0 ) - self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) - self._views: list[torch.Tensor] = [] - self._canonical_offset = 0 - self.is_pinned: bool = False + @property + def base_tensor(self) -> torch.Tensor: + if self._base is None: + raise RuntimeError("Shared offload region has been released.") + return self._base def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: """Allocate a strided int8 view for this worker, one canonical tensor. @@ -311,13 +341,18 @@ def cleanup(self) -> None: if self.is_pinned and self._base is not None: if current_platform.is_cuda_alike(): base_ptr = self._base.data_ptr() - result = torch.cuda.cudart().cudaHostUnregister(base_ptr) - if result.value != 0: - logger.warning( - "cudaHostUnregister failed for rank=%d (code=%d)", - self.rank, - result, - ) + addresses = self.pinned_addresses or [base_ptr] + for address in reversed(addresses): + result = torch.cuda.cudart().cudaHostUnregister(address) + if result.value != 0: + logger.warning( + "cudaHostUnregister failed for rank=%d, " + "address=%#x (code=%d)", + self.rank, + address, + result.value, + ) + self.pinned_addresses.clear() self.is_pinned = False # Release views before _base: each view holds a _base reference and a # direct StorageImpl reference. Freeing views first lets both refcounts From aed894c190aa3db74c502d521b3bc3935ad85de1 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:24:24 +0200 Subject: [PATCH 112/774] [6/N][warmup][DSv4] Migrate sampling, and DFlash JIT kernels (#56323) Signed-off-by: Lucas Wilkinson Signed-off-by: LopezCastroRoberto Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .buildkite/test_areas/models_basic.yaml | 2 +- .buildkite/test_areas/samplers.yaml | 2 +- tests/model_executor/test_jit_warmup.py | 32 ++ .../test_jit_warmup_triton_launcher.py | 335 ++++++++++- .../test_spec_decode_rejection_warmup.py | 121 ---- tests/v1/sample/test_topk_topp_sampler.py | 4 + tests/v1/shutdown/test_delete.py | 16 +- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- .../test_gpu_autoregressive_speculator.py | 1 + .../layers/mamba/ops/scatter_states.py | 49 +- vllm/model_executor/warmup/jit_warmup.py | 339 +++++++++++- .../warmup/jit_warmup_triton_helper.py | 520 +++++++++++++++++- vllm/model_executor/warmup/kernel_warmup.py | 9 +- .../warmup/spec_decode_rejection_warmup.py | 139 ----- vllm/models/glm5next/nvidia/kda.py | 9 + vllm/v1/sample/ops/topk_topp_sampler.py | 21 +- vllm/v1/sample/ops/topk_topp_triton.py | 333 ++++++++--- vllm/v1/sample/rejection_sampler.py | 201 ++++++- vllm/v1/spec_decode/dflash.py | 18 +- vllm/v1/spec_decode/llm_base_proposer.py | 35 +- vllm/v1/spec_decode/utils.py | 258 ++++++++- vllm/v1/worker/cpu_model_runner.py | 7 +- vllm/v1/worker/gpu/model_runner.py | 97 ++-- vllm/v1/worker/gpu/sample/states.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 131 +++-- .../multi_module_mtp/speculator.py | 88 ++- vllm/v1/worker/gpu_model_runner.py | 67 ++- 27 files changed, 2286 insertions(+), 556 deletions(-) delete mode 100644 tests/model_executor/test_spec_decode_rejection_warmup.py delete mode 100644 vllm/model_executor/warmup/spec_decode_rejection_warmup.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index bf453a12d8b3..870cc9b393de 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -41,7 +41,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Basic Models (Other)" device: h200_35gb key: basic-models-tests-other - timeout_in_minutes: 35 + timeout_in_minutes: 38 source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index cfde3717c860..dae7e9646846 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -5,7 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Samplers" device: h200_35gb key: samplers-test - timeout_in_minutes: 40 + timeout_in_minutes: 60 source_file_dependencies: - vllm/model_executor/layers - vllm/sampling_metadata.py diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py index 62ecb05f6b78..7d38f4b40d42 100644 --- a/tests/model_executor/test_jit_warmup.py +++ b/tests/model_executor/test_jit_warmup.py @@ -171,6 +171,19 @@ def test_warmup_range_validates_custom_advancement() -> None: ) +def test_warmup_cases_support_lambda_advancement() -> None: + def warmup_inputs() -> dict[str, Any]: + value: Any = WarmupIntRange(1, 9, advance=lambda value: value * 2) + return dict(value=value) + + assert list(ToyKernel()._expand_warmup_cases(warmup_inputs)) == [ + {"value": 1}, + {"value": 2}, + {"value": 4}, + {"value": 8}, + ] + + def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None: cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True) @@ -577,6 +590,25 @@ def test_registry_records_only_inside_model_setup_context() -> None: assert kernel.compiled == [] +def test_registry_capture() -> None: + class Owner: + @JitWarmupRegistry.capture + def __init__( + self, + vllm_config: Any, + kernel: RecordingToyKernel, + ) -> None: + kernel.register_warmup(3, vllm_config) + + kernel = RecordingToyKernel() + config = _config() + owner = Owner(config, kernel) + + assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined] + kernel.register_warmup(5, config) + assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined] + + def test_registry_expands_requests_and_deduplicates_owner_keys() -> None: registry = JitWarmupRegistry(_config()) kernel = RecordingToyKernel() diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index e76041532ff0..b36361e637de 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -4,22 +4,68 @@ from dataclasses import dataclass from typing import Any -from vllm.model_executor.warmup.jit_warmup import kernel_launcher +import pytest +import torch + +from vllm.model_executor.warmup import jit_warmup_triton_helper +from vllm.model_executor.warmup.jit_warmup import ( + WarmupChoices, + WarmupIntRange, + kernel_launcher, +) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, + TritonJitKey, + TritonWarmupTensor, VllmTritonJitKernel, + triton_kernel_dispatcher_with_warmup, + triton_warmup_inputs, ) +from vllm.triton_utils import tl, triton + + +@triton.jit +def _binder_test_kernel(x, value, BLOCK: tl.constexpr): + pass + + +@triton.jit +def _second_binder_test_kernel(x, value, BLOCK: tl.constexpr): + pass + + +@triton.jit +def _pointer_group_test_kernel(x_ptr, value, BLOCK: tl.constexpr): + pass class _FakeTritonKernel: - arg_names = ("first", "second", "CONST") + arg_names: tuple[str, ...] = ("first", "second", "CONST") def __init__(self) -> None: self.warmup_calls: list[dict[str, Any]] = [] + self.runtime_calls: list[tuple[Any, tuple[Any, ...], dict[str, Any]]] = [] def warmup(self, **kwargs: Any) -> None: self.warmup_calls.append(kwargs) + def __getitem__(self, grid: Any) -> Any: + def launch(*args: Any, **kwargs: Any) -> None: + self.runtime_calls.append((grid, args, kwargs)) + + return launch + + +def _patch_key_deriver( + monkeypatch: pytest.MonkeyPatch, + derive: Any, +) -> None: + monkeypatch.setattr( + jit_warmup_triton_helper, + "_triton_key_deriver", + lambda kernel: lambda kwargs: derive(kernel, kwargs), + ) + class _TestTritonKernel(VllmTritonJitKernel["_TestTritonKernel.CompileKey"]): kernel = _FakeTritonKernel() @@ -91,6 +137,291 @@ class TestCpuKernel(_TestTritonKernel): assert calls == [("runtime", 2, 7)] +def test_triton_kernel_decorator_returns_launcher( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeTritonKernel() + kernel.arg_names = ( + "FIRST_PTR", + "FIRST_STRIDE0", + "aliased_ptr", + "aliased_stride_0", + "SECOND", + "CONST", + ) + + def warmup_inputs() -> dict[str, Any]: + return dict( + first=TritonWarmupTensor(torch.float32, shape=(2, 3), strides=(5, 1)), + second=WarmupChoices(1, 2), + ) + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def launch(first: str, second: int = 2, config: int = 7) -> LaunchSpec: + return (2,), dict(aliased_ptr=first, CONST=config) + + def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]: + return {TritonJitKey(id(kernel), "fake", 0, kwargs["SECOND"])} + + _patch_key_deriver(monkeypatch, fake_keys) + + keys = launch.get_warmup_keys() + keys_with_config = launch.get_warmup_keys(vllm_config=object()) + assert keys_with_config == keys + assert [dict(key.inputs)["second"] for key in keys] == [1, 2] + launch.compile(keys[0]) + assert kernel.warmup_calls == [ + { + "grid": (1,), + "FIRST_PTR": TritonWarmupTensor( + torch.float32, shape=(2, 3), strides=(5, 1) + ), + "FIRST_STRIDE0": 5, + "aliased_ptr": TritonWarmupTensor( + torch.float32, shape=(2, 3), strides=(5, 1) + ), + "aliased_stride_0": 5, + "SECOND": 1, + "CONST": 7, + } + ] + first = torch.empty_strided((2, 3), (5, 1)) + launch(first) + assert kernel.runtime_calls == [ + ( + (2,), + (), + { + "FIRST_PTR": first, + "FIRST_STRIDE0": 5, + "aliased_ptr": first, + "aliased_stride_0": 5, + "SECOND": 2, + "CONST": 7, + }, + ) + ] + assert launch.__name__ == "launch" + with pytest.raises(TypeError, match="unexpected keyword"): + launch(first, 1, 7, stale_constexpr=True) + + +def test_triton_kernel_decorator_exhausts_large_ranges_before_deduplication( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeTritonKernel() + dispatched: list[int] = [] + + def warmup_inputs() -> dict[str, Any]: + tokens: Any = WarmupIntRange(1, 8193) + return dict(first="warmup", second=tokens) + + def specialization(second: int) -> int: + return 7 if second <= 97 else 8 + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def dispatch(first: str, second: int) -> LaunchSpec: + dispatched.append(second) + return (second,), dict(CONST=specialization(second)) + + def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]: + return {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])} + + _patch_key_deriver(monkeypatch, fake_keys) + + keys = dispatch.get_warmup_keys() + assert dispatched == list(range(1, 8193)) + assert len(keys) == 2 + assert {specialization(dict(key.inputs)["second"]) for key in keys} == {7, 8} + + +def test_triton_kernel_decorator_propagates_dispatch_assertions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeTritonKernel() + + def warmup_inputs() -> dict[str, Any]: + return dict(first="warmup", second=WarmupChoices(1, 2)) + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def dispatch(first: str, second: int) -> LaunchSpec: + assert second != 2, "broken dispatch" + return (1,), dict(CONST=second) + + _patch_key_deriver( + monkeypatch, + lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])}, + ) + + with pytest.raises(AssertionError, match="broken dispatch"): + dispatch.get_warmup_keys() + + +def test_triton_kernel_dispatch_uses_cuda_fake_tensors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeTritonKernel() + + def warmup_inputs() -> dict[str, Any]: + return dict( + first=TritonWarmupTensor(torch.float32, shape=(2, 3)), + second=1, + ) + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def dispatch(first: torch.Tensor, second: int) -> LaunchSpec: + assert isinstance(first, torch.Tensor) + assert first.is_cuda + assert first[0].is_contiguous() + return (first.shape[0],), dict(CONST=first[0].numel()) + + _patch_key_deriver( + monkeypatch, + lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])}, + ) + + keys = dispatch.get_warmup_keys() + assert len(keys) == 1 + assert dict(keys[0].inputs)["first"].device.type == "cuda" + + +def test_triton_kernel_decorates_native_launchers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeTritonKernel() + + def warmup_inputs(kernel: Any = kernel) -> dict[str, Any]: + return triton_warmup_inputs( + kernel, + "warmup", + WarmupChoices(1, 2), + grid=(2,), + CONST=7, + ) + + launcher = triton_kernel_dispatcher_with_warmup(warmup_inputs=warmup_inputs)(kernel) + launcher[(3,)]("runtime", 4, CONST=8) + + assert kernel.runtime_calls == [((3,), ("runtime", 4), {"CONST": 8})] + + _patch_key_deriver( + monkeypatch, + lambda kernel, kwargs: { + TritonJitKey(id(kernel), "fake", 0, (kwargs["second"], kwargs["CONST"])) + }, + ) + + keys = launcher.get_warmup_keys() + assert len(keys) == 2 + launcher.compile(keys[0]) + assert kernel.warmup_calls == [ + {"grid": (1,), "first": "warmup", "second": 1, "CONST": 7} + ] + + +def test_triton_warmup_inputs_expands_explicit_pointer_dtypes() -> None: + inputs = triton_warmup_inputs( + _pointer_group_test_kernel, + grid=(1,), + pointer_dtypes={torch.float32: ("x_ptr",)}, + value=2, + BLOCK=16, + ) + + assert inputs["x_ptr"] == TritonWarmupTensor(torch.float32) + assert inputs["value"] == 2 + + with pytest.raises(ValueError, match="not a pointer: value"): + triton_warmup_inputs( + _pointer_group_test_kernel, + grid=(1,), + pointer_dtypes={torch.float32: ("value",)}, + x_ptr=TritonWarmupTensor(torch.float32), + BLOCK=16, + ) + with pytest.raises(ValueError, match="Missing Triton pointer inputs: x_ptr"): + triton_warmup_inputs( + _pointer_group_test_kernel, + grid=(1,), + pointer_dtypes={}, + value=2, + BLOCK=16, + ) + + +def test_triton_key_derivation_applies_wrappers_and_runtime_options( + monkeypatch: pytest.MonkeyPatch, +) -> None: + knobs = triton.knobs + driver = triton.runtime.driver + + device = "test-device" + monkeypatch.setattr(driver.active, "get_current_device", lambda: device) + calls: list[dict[str, Any]] = [] + + def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any: + calls.append(dict(BLOCK=BLOCK, **options)) + specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)] + return {}, specialization, options + + cache: dict[Any, Any] = {} + _binder_test_kernel.device_caches[device] = ({}, cache, None, None, binder) + heuristic_kernel = triton.heuristics( + {"BLOCK": lambda args: 16 if args["value"] <= 16 else 32} + )(_binder_test_kernel) + + keys = jit_warmup_triton_helper._triton_key_deriver(heuristic_kernel)( + {"x": TritonWarmupTensor(tl.float32), "value": 2} + ) + + assert len(keys) == 1 + expected_options = { + "debug": knobs.runtime.debug, + "instrumentation_mode": knobs.compilation.instrumentation_mode, + } + if hasattr(knobs.compilation, "fpsan_homomorphic_casts"): + expected_options["fpsan_homomorphic_casts"] = ( + knobs.compilation.fpsan_homomorphic_casts + ) + assert calls == [{"BLOCK": 16, **expected_options}] + + +def test_triton_key_derivation_covers_autotune_configs_and_jit_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + driver = triton.runtime.driver + + device = "test-device" + monkeypatch.setattr(driver.active, "get_current_device", lambda: device) + + def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any: + specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)] + return {}, specialization, options + + for kernel in (_binder_test_kernel, _second_binder_test_kernel): + kernel.device_caches[device] = ({}, {}, None, None, binder) + autotuned = triton.autotune( + configs=[ + triton.Config({"BLOCK": 16}, num_warps=2), + triton.Config({"BLOCK": 32}, num_warps=4), + ], + key=["value"], + )(_binder_test_kernel) + inputs = {"x": TritonWarmupTensor(tl.float32), "value": 2} + + autotune_keys = jit_warmup_triton_helper._triton_key_deriver(autotuned)(inputs) + other_kernel_keys = jit_warmup_triton_helper._triton_key_deriver( + _second_binder_test_kernel + )(inputs | {"BLOCK": 16}) + + assert len(autotune_keys) == 2 + assert {key.jit_function_key for key in autotune_keys} == { + _binder_test_kernel.cache_key + } + assert {key.jit_function_id for key in autotune_keys} == {id(_binder_test_kernel)} + assert autotune_keys.isdisjoint(other_kernel_keys) + + def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None: from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.block_table import ComputeSlotMappingKernel diff --git a/tests/model_executor/test_spec_decode_rejection_warmup.py b/tests/model_executor/test_spec_decode_rejection_warmup.py deleted file mode 100644 index 6abfb1a4d1c1..000000000000 --- a/tests/model_executor/test_spec_decode_rejection_warmup.py +++ /dev/null @@ -1,121 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace - -import pytest -import torch - -from vllm.config.watermarking import WatermarkConfig, derive_watermark_key -from vllm.model_executor.warmup.spec_decode_rejection_warmup import ( - spec_decode_rejection_warmup, -) - -MASTER_KEY = 20250910 - - -def _worker( - watermark_config: WatermarkConfig | None, - draft_sample_method: str = "greedy", -) -> SimpleNamespace: - return SimpleNamespace( - device=torch.device("cpu"), - vllm_config=SimpleNamespace( - speculative_config=SimpleNamespace( - num_speculative_tokens=2, - rejection_sample_method="standard", - draft_sample_method=draft_sample_method, - ), - model_config=SimpleNamespace( - get_vocab_size=lambda: 128, dtype=torch.bfloat16 - ), - watermark_config=watermark_config, - ), - ) - - -def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: - calls: list[dict] = [] - - def fake_rejection_sample(**kwargs): - calls.append(kwargs) - return None, None - - monkeypatch.setattr( - "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample", - fake_rejection_sample, - ) - return calls - - -def test_warmup_token_id_dtypes_match_the_runtime_buffers( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = _capture_calls(monkeypatch) - spec_decode_rejection_warmup(_worker(None)) - - assert calls - for kwargs in calls: - assert kwargs["draft_sampled"].dtype == torch.int32 - assert kwargs["pos"].dtype == torch.int64 - - -def test_warmup_index_mapping_dtypes_match_the_runtime_buffers( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = _capture_calls(monkeypatch) - spec_decode_rejection_warmup(_worker(None)) - - assert calls - for kwargs in calls: - assert kwargs["idx_mapping"].dtype == torch.int64 - assert kwargs["expanded_idx_mapping"].dtype == torch.int64 - assert kwargs["expanded_local_pos"].dtype == torch.int32 - - -def test_warmup_watermark_args_match_the_runtime_kernel_arguments( - monkeypatch: pytest.MonkeyPatch, -) -> None: - watermark_config = WatermarkConfig( - key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3 - ) - calls = _capture_calls(monkeypatch) - spec_decode_rejection_warmup( - _worker(watermark_config, draft_sample_method="probabilistic") - ) - - assert calls - key_b = derive_watermark_key(MASTER_KEY, b"key_b") - for kwargs in calls: - contexts = kwargs["contexts"] - assert contexts.dtype == torch.int32 - assert contexts.shape == (3, 3) and contexts.stride() == (3, 1) - assert kwargs["watermarking"].dtype == torch.bool - assert kwargs["watermark_key"] == key_b - assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31 - - -def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = _capture_calls(monkeypatch) - monkeypatch.setattr( - "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key", - lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")), - ) - spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) - - assert calls == [] - - -def test_warmup_draft_logits_match_the_runtime_draft_sample_method( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = _capture_calls(monkeypatch) - spec_decode_rejection_warmup(_worker(None)) - assert calls - assert all(kwargs["draft_logits"] is None for kwargs in calls) - - calls = _capture_calls(monkeypatch) - spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic")) - assert calls - assert all(kwargs["draft_logits"] is not None for kwargs in calls) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index fdf704169288..5a0ce5cc089f 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -79,6 +79,10 @@ class MockPlatform: def is_cuda(): return False + @staticmethod + def is_cuda_alike(): + return True + @staticmethod def is_cpu(): return False diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index 991ad6891856..3a513945eaaa 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/tests/v1/shutdown/test_delete.py @@ -11,6 +11,7 @@ SHUTDOWN_TEST_TIMEOUT_SEC, ) from vllm import LLM, SamplingParams +from vllm.config import KernelConfig from vllm.engine.arg_utils import AsyncEngineArgs from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind @@ -39,7 +40,10 @@ async def test_async_llm_delete( pytest.skip(reason="Not enough CUDA devices") engine_args = AsyncEngineArgs( - model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size + model=model, + enforce_eager=True, + tensor_parallel_size=tensor_parallel_size, + kernel_config=KernelConfig(enable_jit_warmup=False), ) # Instantiate AsyncLLM; make request to complete any deferred @@ -94,7 +98,10 @@ def test_llm_delete( # Instantiate LLM; make request to complete any deferred # initialization; then delete instance llm = LLM( - model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size + model=model, + enforce_eager=True, + tensor_parallel_size=tensor_parallel_size, + kernel_config=KernelConfig(enable_jit_warmup=False), ) if send_one_request: llm.generate( @@ -122,7 +129,10 @@ def test_llm_delete_inprocess( with monkeypatch.context() as m: m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - with VllmRunner(model) as vllm_model: + with VllmRunner( + model, + kernel_config=KernelConfig(enable_jit_warmup=False), + ) as vllm_model: if send_one_request: vllm_model.generate( ["Hello my name is"], diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 16d6d8e516df..ce47cb8ae625 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -43,7 +43,7 @@ def _run_prepare( num_scheduled_tokens=np.array([4], dtype=np.int32), positions=torch.tensor(target_positions, dtype=torch.int64, device=device), query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device), - idx_mapping=torch.tensor([2], dtype=torch.int32, device=device), + idx_mapping=torch.tensor([2], dtype=torch.int64, device=device), ) query_slot_mapping = torch.full( (max_num_tokens,), -2, dtype=torch.int64, device=device diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 920c9a25f062..a38265a6f7af 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -157,6 +157,7 @@ def init_base(speculator, vllm_config, device): speculator.device = device speculator.max_num_tokens = 4 speculator.max_num_reqs = 2 + speculator.num_speculative_steps = 1 speculator.hidden_size = 3 speculator.dtype = torch.float32 speculator.draft_model_config = draft_model_config diff --git a/vllm/model_executor/layers/mamba/ops/scatter_states.py b/vllm/model_executor/layers/mamba/ops/scatter_states.py index 0102b6cf12b4..04b5cb61c2b9 100644 --- a/vllm/model_executor/layers/mamba/ops/scatter_states.py +++ b/vllm/model_executor/layers/mamba/ops/scatter_states.py @@ -1,8 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch +from vllm.model_executor.warmup.jit_warmup import WarmupIntRange +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -33,11 +41,39 @@ def _scatter_states_kernel( tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask) +def _scatter_states_warmup_inputs( + *, + state_shape: tuple[int, ...], + state_dtype: torch.dtype, + indices_dtype: torch.dtype, + max_num_tokens: int, +) -> dict[str, object]: + num_tokens: Any = WarmupIntRange(1, max_num_tokens + 1) + return dict( + state=TritonWarmupTensor( + state_dtype, + shape=(1,) + state_shape, + ), + src=TritonWarmupTensor( + state_dtype, + shape=(num_tokens,) + state_shape, + ), + indices=TritonWarmupTensor( + indices_dtype, + shape=(num_tokens,), + ), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_scatter_states_kernel, + warmup_inputs=_scatter_states_warmup_inputs, +) def scatter_states( state: torch.Tensor, src: torch.Tensor, indices: torch.Tensor, -) -> None: +) -> DispatchSpec: """Scatter ``src`` rows into ``state`` at ``indices`` (in place). Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound, @@ -58,13 +94,10 @@ def scatter_states( assert src[0].is_contiguous() block_size = min(triton.next_power_of_2(row_size), 1024) grid = (triton.cdiv(row_size, block_size), indices.numel()) - _scatter_states_kernel[grid]( - state, - src, - indices, - state.stride(0), - src.stride(0), - indices.stride(0), + return grid, dict( + stride_state_batch=state.stride(0), + stride_src_batch=src.stride(0), + stride_indices=indices.stride(0), row_size=row_size, BLOCK_SIZE=block_size, num_warps=8, diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index fc9d6fcdf895..72626964ae67 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -21,6 +21,7 @@ __all__ = [ "JitWarmupRegistry", "VllmJitKernel", + "WarmupChoices", "WarmupIntRange", "get_ast_full_name", "get_function_source_node", @@ -62,6 +63,21 @@ class WarmupIntRange: advance: Callable[[int], int] | None = None +@dataclass(frozen=True) +class WarmupChoices: + """Expand an explicit finite set of values in a traced warmup method.""" + + values: tuple[Any, ...] + + def __init__(self, *values: Any) -> None: + object.__setattr__(self, "values", values) + + +def _when(condition: bool) -> None: + """Filter symbolic cases in an AST-traced warmup-cases method.""" + raise RuntimeError("_when() is only valid in an AST-traced warmup method") + + WarmupValues = Any CompileKeyDispatchFn = Callable[..., CompileKeyT] WarmupPredicateFn = Callable[..., bool] @@ -87,6 +103,8 @@ class _WarmupInputRows: def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]: + if isinstance(values, WarmupChoices): + return values.values if isinstance(values, WarmupIntRange): if values.advance is None: return tuple(range(values.start, values.stop, values.step)) @@ -278,6 +296,27 @@ def visit_Name(self, node: ast.Name) -> Any: def visit_Constant(self, node: ast.Constant) -> Any: return node.value + def visit_Lambda(self, node: ast.Lambda) -> Callable[..., Any]: + arguments = node.args + if ( + arguments.posonlyargs + or arguments.vararg is not None + or arguments.kwonlyargs + or arguments.kwarg is not None + or arguments.defaults + ): + raise _dispatch_expr_error(node, "Traced lambdas require positional args") + names = tuple(argument.arg for argument in arguments.args) + + def evaluate(*values: Any) -> Any: + if len(values) != len(names): + raise TypeError(f"Expected {len(names)} lambda arguments") + return type(self)( + dict(self.values) | dict(zip(names, values)), self.globals + ).eval(node.body) + + return evaluate + def visit_IfExp(self, node: ast.IfExp) -> Any: return self.visit(node.body if self.visit(node.test) else node.orelse) @@ -333,7 +372,12 @@ def visit_BinOp(self, node: ast.BinOp) -> Any: return op(self.visit(node.left), self.visit(node.right)) def visit_Call(self, node: ast.Call) -> Any: - args = [self.visit(arg) for arg in node.args] + args: list[Any] = [] + for arg in node.args: + if isinstance(arg, ast.Starred): + args.extend(self.visit(arg.value)) + else: + args.append(self.visit(arg)) fn = self.visit(node.func) call_kwargs: dict[str, Any] = {} for keyword in node.keywords: @@ -388,6 +432,45 @@ def _eval_dispatch_expr( return _DispatchExprEvaluator(kwargs, globals_).eval(node) +def _validate_dispatch_expr(node: ast.AST) -> None: + """Enforce the dispatch AST subset before compiling traced code.""" + allowed_nodes = ( + ast.Name, + ast.Constant, + ast.Lambda, + ast.arguments, + ast.arg, + ast.IfExp, + ast.Tuple, + ast.List, + ast.BoolOp, + ast.And, + ast.Or, + ast.Compare, + *tuple(_CMP_OPS), + ast.UnaryOp, + ast.Not, + ast.USub, + ast.BinOp, + *tuple(_BIN_OPS), + ast.Call, + ast.keyword, + ast.Starred, + ast.Attribute, + ast.Subscript, + ast.Load, + ) + for child in ast.walk(node): + if not isinstance(child, allowed_nodes): + raise _dispatch_expr_error(child, "Unsupported dispatch expression") + if isinstance(child, ast.Call) and any( + keyword.arg is None for keyword in child.keywords + ): + raise _dispatch_expr_error( + child, "Dispatch helper calls cannot use **kwargs" + ) + + def _collect_input_names( node: ast.AST, candidate_names: set[str], @@ -406,6 +489,23 @@ def _collect_input_names( } +def _named_assignment( + statement: ast.stmt, + subject: str, +) -> tuple[str, ast.AST] | None: + match statement: + case ast.Assign(targets=[ast.Name(id=name)], value=value): + return name, value + case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None: + return name, value + case ast.Assign() | ast.AnnAssign(): + raise _dispatch_expr_error( + statement, f"{subject} assignments require one name" + ) + case _: + return None + + def _collect_expression_body( fn: Callable[..., Any], function_def: ast.FunctionDef | ast.Lambda, @@ -422,38 +522,19 @@ def _collect_expression_body( ): continue - if isinstance(statement, ast.Assign): - if len(statement.targets) != 1 or not isinstance( - statement.targets[0], ast.Name - ): - raise _dispatch_expr_error( - statement, - "Traced warmup helper assignments must target one local name", - ) - local_exprs.append((statement.targets[0].id, statement.value)) + if ( + assignment := _named_assignment(statement, "Traced warmup helper") + ) is not None: + local_exprs.append(assignment) continue - if isinstance(statement, ast.AnnAssign): + if isinstance(statement, ast.Return): if statement.value is None: raise _dispatch_expr_error( - statement, "Traced warmup helper annotations must assign a value" + statement, "Traced warmup helper must return an expression" ) - if not isinstance(statement.target, ast.Name): - raise _dispatch_expr_error( - statement, - "Traced warmup helper assignments must target one local name", - ) - local_exprs.append((statement.target.id, statement.value)) - continue - - if isinstance(statement, ast.Return) and statement.value is not None: return local_exprs, statement.value - if isinstance(statement, ast.Return): - raise _dispatch_expr_error( - statement, "Traced warmup helper must return an expression" - ) - raise _dispatch_expr_error( statement, "Traced warmup helper may only contain local assignments before return", @@ -574,10 +655,17 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): bind_launch_inputs = True def __init__(self) -> None: - self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + dispatch = type(self).dispatch + self._dispatch_trace = ( + None + if dispatch is VllmJitKernel.dispatch + else _trace_compile_key_dispatch(self.dispatch) + ) self._compiled_cache: dict[Any, Any] = {} def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: + if self._dispatch_trace is None: + raise TypeError(f"{type(self).__name__} does not define dispatch()") return self._dispatch_trace.compile_key(self.CompileKey, kwargs) def _get_or_compile( @@ -683,7 +771,184 @@ def traced( return traced - @abstractmethod + def _expand_warmup_cases( + self, + cases_fn: Callable[..., Any], + *args: Any, + _value_expander: Callable[[WarmupValues], tuple[Any, ...]] = ( + _expand_warmup_values + ), + **kwargs: Any, + ) -> Iterator[Mapping[str, Any]]: + """Expand symbolic domains declared inside a warmup-cases method.""" + function_def = get_function_source_node(cases_fn) + if isinstance(function_def, ast.Lambda): + raise _dispatch_expr_error( + function_def, "Warmup cases must be a function definition" + ) + signature = inspect.signature(cases_fn) + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + static_values = dict(bound.arguments) + bound_self = getattr(cases_fn, "__self__", None) + if bound_self is not None: + static_values["self"] = bound_self + globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__ + + domains: list[tuple[str, tuple[Any, ...]]] = [] + local_exprs: list[tuple[str, ast.AST]] = [] + predicates: list[ast.AST] = [] + return_expr: ast.AST | None = None + for statement in function_def.body: + if ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str) + ): + continue + assignment = _named_assignment(statement, "Warmup case") + if assignment is not None: + name, value_expr = assignment + call_name = ( + (get_ast_full_name(value_expr.func) or "").split(".")[-1] + if isinstance(value_expr, ast.Call) + else "" + ) + if call_name in {"WarmupIntRange", "WarmupChoices"}: + value = _eval_dispatch_expr(value_expr, static_values, globals_) + domains.append((name, _value_expander(value))) + else: + local_exprs.append((name, value_expr)) + continue + if ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Call) + and (get_ast_full_name(statement.value.func) or "").endswith("_when") + ): + if len(statement.value.args) != 1 or statement.value.keywords: + raise _dispatch_expr_error( + statement, "_when requires exactly one positional expression" + ) + predicates.append(statement.value.args[0]) + continue + if isinstance(statement, ast.Return): + if return_expr is not None or statement.value is None: + raise _dispatch_expr_error( + statement, "Warmup cases require one final return expression" + ) + return_expr = statement.value + continue + raise _dispatch_expr_error( + statement, + "AST-traced warmup cases support assignments, _when, and return", + ) + + if return_expr is None or not isinstance(return_expr, ast.Call): + raise ValueError("AST-traced warmup cases must return one case(...) call") + if any(keyword.arg is None for keyword in return_expr.keywords): + raise _dispatch_expr_error( + return_expr, "Warmup case expressions do not support **kwargs" + ) + + class _InlineDomainRewriter(ast.NodeTransformer): + def visit_Call(self, node: ast.Call) -> ast.AST: + call_name = (get_ast_full_name(node.func) or "").split(".")[-1] + if call_name not in {"WarmupIntRange", "WarmupChoices"}: + return self.generic_visit(node) + name = f"__warmup_domain_{len(domains)}" + domain = _eval_dispatch_expr(node, static_values, globals_) + domains.append((name, _value_expander(domain))) + return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node) + + return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr)) + + domain_names = tuple(name for name, _ in domains) + dynamic_names = set(domain_names) + dynamic_local_exprs: list[tuple[str, ast.AST]] = [] + for name, expr in local_exprs: + referenced_names = { + node.id for node in ast.walk(expr) if isinstance(node, ast.Name) + } + if referenced_names & dynamic_names: + dynamic_local_exprs.append((name, expr)) + dynamic_names.add(name) + else: + static_values[name] = _eval_dispatch_expr(expr, static_values, globals_) + + for _, expr in dynamic_local_exprs: + _validate_dispatch_expr(expr) + for predicate in predicates: + _validate_dispatch_expr(predicate) + _validate_dispatch_expr(return_expr) + + case_body: list[ast.stmt] = [ + ast.Assign( + targets=[ast.Name(id=name, ctx=ast.Store())], + value=cast(ast.expr, expr), + ) + for name, expr in dynamic_local_exprs + ] + if predicates: + predicate = ( + predicates[0] + if len(predicates) == 1 + else ast.BoolOp( + op=ast.And(), + values=[cast(ast.expr, value) for value in predicates], + ) + ) + case_body.append( + ast.If( + test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)), + body=[ast.Return(value=ast.Constant(value=None))], + orelse=[], + ) + ) + case_body.append(ast.Return(value=cast(ast.expr, return_expr))) + # FunctionDef fields vary across the supported Python versions, so no + # single constructor overload matches every mypy target. + case_function = ast.FunctionDef( # type: ignore[call-overload] + name="__vllm_warmup_case", + args=ast.arguments( + posonlyargs=[], + args=[ast.arg(arg=name) for name in domain_names], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ), + body=case_body, + decorator_list=[], + returns=None, + type_comment=None, + ) + case_globals = dict(globals_) + case_globals.update(static_values) + exec( + compile( + ast.fix_missing_locations( + ast.Module(body=[case_function], type_ignores=[]) + ), + "", + "exec", + ), + case_globals, + ) + evaluate_case = cast(Callable[..., Any], case_globals[case_function.name]) + + domain_values = tuple(values for _, values in domains) + for values in itertools.product(*domain_values): + case = evaluate_case(*values) + if case is None: + continue + if not isinstance(case, Mapping): + raise TypeError("AST-traced warmup cases must return a mapping") + if any( + isinstance(value, WarmupChoices | WarmupIntRange) + for value in case.values() + ): + raise TypeError("Warmup domains must be direct expressions") + yield case + def dispatch(self, **kwargs: Any) -> CompileKeyT: """Build one compile key from one concrete dispatch point.""" raise NotImplementedError @@ -751,6 +1016,19 @@ def __init__(self, vllm_config: Any) -> None: list[tuple[tuple[Any, ...], dict[str, Any]]], ] = {} + @classmethod + def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]: + """Collect warmup registrations made while the decorated callable runs.""" + + @wraps(init_fn) + def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None: + registry = cls(vllm_config) + instance.jit_warmup_registry = registry + with registry.activate(): + init_fn(instance, vllm_config, *args, **kwargs) + + return wrapped + @contextmanager def activate(self) -> Iterator[None]: """Collect registrations made in this context.""" @@ -808,9 +1086,10 @@ def warmup(self) -> None: if ( not args and not kwargs - and inspect.signature(kernel.get_warmup_keys).parameters + and "vllm_config" + in inspect.signature(kernel.get_warmup_keys).parameters ): - args = (self.vllm_config,) + kwargs = {"vllm_config": self.vllm_config} for compile_key in kernel.get_warmup_keys(*args, **kwargs): compile_keys[compile_key] = None if compile_keys: diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index b019404cc9d1..dab055902aee 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -3,25 +3,130 @@ import ast import inspect from abc import abstractmethod -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from functools import cached_property -from typing import Any, ClassVar, Generic, TypeVar, cast +from collections.abc import Callable, Hashable, Iterable, Mapping +from contextlib import nullcontext, suppress +from dataclasses import dataclass, field +from functools import cache, cached_property, update_wrapper +from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload from vllm.model_executor.warmup.jit_warmup import ( VllmJitKernel, get_ast_full_name, get_function_source_node, ) +from vllm.platforms import current_platform +from vllm.triton_utils import triton CompileKeyT = TypeVar("CompileKeyT") -# ``(grid, launch_kwargs)`` or ``(grid, launch_kwargs, outputs)`` for -# self-allocating kernels. ``grid=None`` skips the launch (e.g. empty batch) -# but still returns the declared outputs. +P = ParamSpec("P") +# ``kernel_launcher`` also supports a ``None`` grid when no launch is needed. LaunchSpec = ( tuple[tuple[int, ...] | None, dict[str, Any]] | tuple[tuple[int, ...] | None, dict[str, Any], Any] ) +# Dispatchers select a concrete launch configuration for a Triton kernel. +DispatchSpec = ( + tuple[tuple[int, ...], dict[str, Any]] | tuple[tuple[int, ...], dict[str, Any], Any] +) + + +@dataclass(frozen=True) +class _LaunchBindingPlan: + input_targets: tuple[tuple[int, str], ...] + stride_targets: tuple[tuple[str, str, int], ...] + + +@cache +def _launch_binding_plan( + arg_names: tuple[str, ...], + input_names: tuple[str, ...], + launch_names: tuple[str, ...], +) -> _LaunchBindingPlan: + arg_name_set = set(arg_names) + launch_name_set = set(launch_names) + + def kernel_arg(input_name: str) -> str | None: + pointer_name = f"{input_name}_ptr" + candidates = ( + input_name, + pointer_name, + input_name.upper(), + pointer_name.upper(), + ) + return next( + (candidate for candidate in candidates if candidate in arg_name_set), + None, + ) + + input_targets = tuple( + (index, target) + for index, name in enumerate(input_names) + if (target := kernel_arg(name)) is not None and target not in launch_name_set + ) + bound_names = launch_name_set | {target for _, target in input_targets} + stride_targets = [] + for target in arg_names: + if target in bound_names: + continue + name, separator, dim = target.rpartition("_stride") + if not separator: + name, separator, dim = target.rpartition("_STRIDE") + dim = dim.removeprefix("_") + if separator and (not dim or dim.isdigit()): + pointer = f"{name}_PTR" if separator == "_STRIDE" else f"{name}_ptr" + source = name if name in bound_names else pointer + if source in bound_names: + stride_targets.append((target, source, int(dim or 0))) + return _LaunchBindingPlan(input_targets, tuple(stride_targets)) + + +def triton_warmup_inputs( + kernel: Any, + *args: Any, + grid: tuple[int, ...], + pointer_dtypes: ( + Mapping[Any, Iterable[str]] | Iterable[tuple[Any, Iterable[str]]] | None + ) = None, + **kwargs: Any, +) -> dict[str, Any]: + """Build launcher inputs from Triton's native positional argument order.""" + arg_names = tuple(kernel.arg_names) + if len(args) > len(arg_names): + raise ValueError( + f"Received {len(args)} positional inputs for {len(arg_names)} " + "Triton kernel arguments" + ) + inputs = dict(zip(arg_names, args)) + duplicate_names = inputs.keys() & kwargs.keys() + if duplicate_names: + raise ValueError( + f"Triton inputs passed twice: {', '.join(sorted(duplicate_names))}" + ) + inputs.update(kwargs) + if pointer_dtypes is not None: + function_def = get_function_source_node(kernel) + if not isinstance(function_def, ast.FunctionDef): + raise ValueError("Expected Triton kernel to be defined as a function") + pointer_names = _pointer_arg_names(function_def, arg_names) + groups = ( + pointer_dtypes.items() + if isinstance(pointer_dtypes, Mapping) + else pointer_dtypes + ) + for dtype, names in groups: + for name in names: + if name not in arg_names: + raise ValueError(f"Unknown Triton pointer argument: {name}") + if name not in pointer_names: + raise ValueError(f"Triton argument is not a pointer: {name}") + if name in inputs: + raise ValueError(f"Triton input passed twice: {name}") + inputs[name] = TritonWarmupTensor(dtype) + missing_pointers = pointer_names - inputs.keys() + if missing_pointers: + names = ", ".join(sorted(missing_pointers)) + raise ValueError(f"Missing Triton pointer inputs: {names}") + return {"grid": grid, **inputs} def triton_scalar_specialization_rep(value: int) -> int: @@ -72,6 +177,7 @@ class TritonWarmupTensor: aligned: bool = True shape: tuple[int, ...] = (1,) strides: tuple[int, ...] | None = None + init: Any = 0 def data_ptr(self) -> int: return 0 if self.aligned else 1 @@ -79,6 +185,12 @@ def data_ptr(self) -> int: def ptr_range(self) -> int: return 0 + @property + def device(self) -> Any: + import torch + + return torch.device(current_platform.device_type) + def stride(self, dim: int | None = None) -> int | tuple[int, ...]: if self.strides is None: strides: list[int] = [] @@ -95,9 +207,10 @@ def stride(self, dim: int | None = None) -> int | tuple[int, ...]: class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): """Triton owner whose runtime launch specification is reused for warmup.""" - kernel: ClassVar[Any] + kernel: Any _warming = False _warming_compile_key: CompileKeyT | None = None + _run_autotune = False @abstractmethod def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]: @@ -126,6 +239,21 @@ def _kernel_arg_names(self) -> tuple[str, ...]: f"Cannot inspect kernel parameters for {type(self.kernel).__name__}" ) + def _prepare_launch_kwargs( + self, + input_names: tuple[str, ...], + input_values: tuple[Any, ...], + launch_kwargs: dict[str, Any], + ) -> dict[str, Any]: + plan = _launch_binding_plan( + self._kernel_arg_names, input_names, tuple(launch_kwargs) + ) + for index, target in plan.input_targets: + launch_kwargs[target] = input_values[index] + for target, source, dim in plan.stride_targets: + launch_kwargs[target] = launch_kwargs[source].stride(dim) + return launch_kwargs + def launch( self, launch_spec: LaunchSpec, @@ -136,15 +264,20 @@ def launch( else: grid, launch_kwargs = launch_spec outputs = None - kwargs = dict(launch_kwargs) + kwargs = self._prepare_launch_kwargs( + tuple(inputs), tuple(inputs.values()), dict(launch_kwargs) + ) kernel = kwargs.pop("kernel", None) or self.kernel runtime_launcher = kwargs.pop("_runtime_launcher", None) runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0) - for name, value in inputs.items(): - target = name if name in self._kernel_arg_names else f"{name}_ptr" - if target in self._kernel_arg_names and target not in kwargs: - kwargs[target] = value if self._warming: + if self._run_autotune: + assert grid is not None + kernel[grid](**kwargs) + return outputs + kwargs = { + name: _triton_metadata_arg(value) for name, value in kwargs.items() + } if ( self._warming_compile_key is not None and "launch_pdl" in kwargs @@ -168,6 +301,367 @@ def launch( return outputs +@dataclass(frozen=True) +class TritonJitKey: + """Process-local identity of one Triton JIT specialization.""" + + jit_function_id: int + jit_function_key: Hashable + device: Hashable + cache_key: Hashable + + +@dataclass(frozen=True) +class TritonCompileKey: + """Triton-derived compile key with one non-comparing replay input.""" + + jit_keys: frozenset[TritonJitKey] + inputs: tuple[tuple[str, Any], ...] = field(compare=False, hash=False, repr=False) + + +def _triton_key_deriver( + kernel: Any, +) -> Callable[[Mapping[str, Any]], set[TritonJitKey]]: + """Prepare Triton's key derivation once for one kernel and device.""" + knobs = triton.knobs + Autotuner = triton.runtime.autotuner.Autotuner + Heuristics = triton.runtime.autotuner.Heuristics + driver = triton.runtime.driver + JITFunction = triton.runtime.jit.JITFunction + compute_cache_key = triton.runtime.jit.compute_cache_key + + if isinstance(kernel, Heuristics): + derive_inner = _triton_key_deriver(kernel.fn) + + def derive_heuristic(kwargs: Mapping[str, Any]) -> set[TritonJitKey]: + heuristic_kwargs = dict(kwargs) + for name, heuristic in kernel.values.items(): + heuristic_kwargs[name] = heuristic(heuristic_kwargs) + return derive_inner(heuristic_kwargs) + + return derive_heuristic + + if isinstance(kernel, Autotuner): + derive_inner = _triton_key_deriver(kernel.fn) + + def derive_autotuned(kwargs: Mapping[str, Any]) -> set[TritonJitKey]: + previous_nargs = getattr(kernel, "nargs", None) + kernel.nargs = {} + try: + configs = kernel.prune_configs(dict(kwargs)) + finally: + kernel.nargs = previous_nargs + keys: set[TritonJitKey] = set() + for config in configs: + conflicts = kwargs.keys() & config.kwargs.keys() + if conflicts: + names = ", ".join(sorted(conflicts)) + raise ValueError(f"Conflicting autotune parameters: {names}") + keys.update(derive_inner(dict(kwargs) | config.all_kwargs())) + return keys + + return derive_autotuned + + if not isinstance(kernel, JITFunction): + raise TypeError(f"Unsupported Triton kernel wrapper: {type(kernel).__name__}") + + device = cast(Hashable, driver.active.get_current_device()) + _, kernel_key_cache, _, _, binder = kernel.device_caches[device] + jit_function_id = id(kernel) + jit_function_key = cast(Hashable, kernel.cache_key) + debug = kernel.debug or knobs.runtime.debug + instrumentation_mode = knobs.compilation.instrumentation_mode + fpsan_homomorphic_casts = getattr( + knobs.compilation, "fpsan_homomorphic_casts", None + ) + inspection_hook = knobs.runtime.add_stages_inspection_hook + + def derive_jit(kwargs: Mapping[str, Any]) -> set[TritonJitKey]: + binder_kwargs = dict(kwargs) + binder_kwargs["debug"] = binder_kwargs.get("debug", debug) or debug + binder_kwargs["instrumentation_mode"] = instrumentation_mode + if fpsan_homomorphic_casts is not None: + binder_kwargs["fpsan_homomorphic_casts"] = fpsan_homomorphic_casts + _, specialization, options = binder(**binder_kwargs) + if inspection_hook is not None: + _, inspection_hash = inspection_hook() + specialization.append(f'("custom_pipeline", {inspection_hash})') + cache_key = cast( + Hashable, + compute_cache_key(kernel_key_cache, specialization, options), + ) + return { + TritonJitKey( + jit_function_id, + jit_function_key, + device, + cache_key, + ) + } + + return derive_jit + + +WarmupCases = Mapping[str, Any] | Iterable[Mapping[str, Any]] + + +class _AutomaticTritonJitKernel(VllmTritonJitKernel[TritonCompileKey]): + def warmup_inputs(self, compile_key: TritonCompileKey) -> dict[str, Any]: + return dict(compile_key.inputs) + + def _provider_cases( + self, + provider: Callable[..., WarmupCases], + *args: Any, + **kwargs: Any, + ) -> Iterable[Mapping[str, Any]]: + cases_node = get_function_source_node(provider) + if any( + isinstance(node, ast.Call) + and (get_ast_full_name(node.func) or "").split(".")[-1] + in {"WarmupIntRange", "WarmupChoices", "_when"} + for node in ast.walk(cases_node) + ): + return self._expand_warmup_cases(provider, *args, **kwargs) + cases = provider(*args, **kwargs) + return (cases,) if isinstance(cases, Mapping) else cases + + +def _is_autotuned(kernel: Any) -> bool: + Autotuner = triton.runtime.autotuner.Autotuner + + while kernel is not None: + if isinstance(kernel, Autotuner): + return True + kernel = getattr(kernel, "fn", None) + return False + + +def _triton_metadata_arg(value: Any) -> Any: + from torch._subclasses.fake_tensor import FakeTensor + + if not isinstance(value, FakeTensor): + return value + return TritonWarmupTensor( + value.dtype, + aligned=getattr(value, "_vllm_warmup_aligned", True), + shape=tuple(value.shape), + strides=tuple(value.stride()), + ) + + +def _materialize_warmup_case( + case: Mapping[str, Any], + *, + real: bool, + mode: Any, + cache: dict[TritonWarmupTensor, Any], +) -> dict[str, Any]: + import torch + + def materialize(value: Any) -> Any: + if not isinstance(value, TritonWarmupTensor): + return value + if not real: + try: + return cache[value] + except KeyError: + pass + except TypeError: + # Unhashable initializers are uncommon and safe to materialize directly. + pass + + strides = cast(tuple[int, ...], value.stride()) + with mode: + tensor = torch.empty_strided( + value.shape, + strides, + dtype=value.dtype, + device=value.device, + ) + if real: + if callable(value.init): + value.init(tensor) + else: + tensor.fill_(value.init) + else: + tensor._vllm_warmup_aligned = value.aligned + with suppress(TypeError): + cache[value] = tensor + return tensor + + return {name: materialize(value) for name, value in case.items()} + + +class _DecoratedTritonJitKernel(_AutomaticTritonJitKernel): + def __init__( + self, + kernel: Any, + warmup_inputs: Callable[..., WarmupCases], + dispatch: Callable[..., DispatchSpec] | None, + ) -> None: + self.kernel = kernel + self._run_autotune = _is_autotuned(kernel) + self._warmup_inputs_fn = warmup_inputs + self._dispatch_fn = dispatch + self._dispatch_arg_names: tuple[str, ...] = () + self._dispatch_defaults: dict[str, Any] = {} + if dispatch is not None: + parameters = tuple(inspect.signature(dispatch).parameters.values()) + if any( + parameter.kind + in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + for parameter in parameters + ): + raise TypeError("Triton dispatchers require explicit parameters") + self._dispatch_arg_names = tuple(parameter.name for parameter in parameters) + self._dispatch_defaults = { + parameter.name: parameter.default + for parameter in parameters + if parameter.default is not inspect.Parameter.empty + } + super().__init__() + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + if self._dispatch_fn is None: + grid = kwargs.pop("grid") if not args else args[0] + kernel_args = () if not args else args[1:] + inputs = triton_warmup_inputs( + self.kernel, *kernel_args, grid=grid, **kwargs + ) + inputs.pop("grid") + return self.launch((grid, inputs), {}) + spec = self._dispatch_fn(*args, **kwargs) + input_values = args + if len(args) != len(self._dispatch_arg_names): + input_values += tuple( + kwargs[name] if name in kwargs else self._dispatch_defaults[name] + for name in self._dispatch_arg_names[len(args) :] + ) + return self.launch( + spec, + dict(zip(self._dispatch_arg_names, input_values)), + ) + + def __getitem__(self, grid: tuple[int, ...]) -> Callable[..., Any]: + return cast(Callable[..., Any], self.kernel[grid]) + + def __getattr__(self, name: str) -> Any: + return getattr(self.kernel, name) + + def get_warmup_keys( + self, *, vllm_config: Any | None = None, **kwargs: Any + ) -> list[TritonCompileKey]: + keys: list[TritonCompileKey] = [] + seen_jit_keys: set[frozenset[TritonJitKey]] = set() + if "vllm_config" in inspect.signature(self._warmup_inputs_fn).parameters: + kwargs.setdefault("vllm_config", vllm_config) + cases = self._provider_cases(self._warmup_inputs_fn, **kwargs) + derive_keys = _triton_key_deriver(self.kernel) + materialize = self._run_autotune + mode: Any = nullcontext() if self._run_autotune else None + materialized: dict[TritonWarmupTensor, Any] = {} + for case in cases: + inputs = ( + _materialize_warmup_case( + case, + real=self._run_autotune, + mode=mode, + cache=materialized, + ) + if materialize + else case + ) + if self._dispatch_fn is None: + prepared = dict(inputs) + prepared.pop("grid") + else: + dispatch_inputs = { + name: inputs[name] + for name in self._dispatch_arg_names + if name in inputs + } + try: + _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2] + except (AssertionError, AttributeError, TypeError): + if materialize: + raise + from torch._subclasses.fake_tensor import FakeTensorMode + + materialize = True + mode = FakeTensorMode() + inputs = _materialize_warmup_case( + case, + real=False, + mode=mode, + cache=materialized, + ) + dispatch_inputs = { + name: inputs[name] + for name in self._dispatch_arg_names + if name in inputs + } + _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2] + prepared = self._prepare_launch_kwargs( + tuple(inputs), + tuple(inputs.values()), + dict(launch_kwargs), + ) + if materialize: + prepared = { + name: _triton_metadata_arg(value) + for name, value in prepared.items() + } + jit_keys = frozenset(derive_keys(prepared)) + if jit_keys and jit_keys not in seen_jit_keys: + seen_jit_keys.add(jit_keys) + keys.append(TritonCompileKey(jit_keys, tuple(inputs.items()))) + return keys + + +class TritonKernelDispatcher(Protocol[P]): + """Callable Triton launcher with compile-time warmup support.""" + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any: ... + + def register_warmup(self, *args: Any, **kwargs: Any) -> None: ... + + +@overload +def triton_kernel_dispatcher_with_warmup( + *, + warmup_inputs: Callable[..., WarmupCases], + kernel: None = None, +) -> Callable[[Any], Any]: ... + + +@overload +def triton_kernel_dispatcher_with_warmup( + *, + warmup_inputs: Callable[..., WarmupCases], + kernel: Any, +) -> Callable[[Callable[P, DispatchSpec]], TritonKernelDispatcher[P]]: ... + + +def triton_kernel_dispatcher_with_warmup( + *, + warmup_inputs: Callable[..., WarmupCases], + kernel: Any | None = None, +) -> Callable[[Any], Any]: + """Decorate a dispatch function or a native Triton kernel for warmup.""" + + def decorate(dispatch: Any) -> Any: + owner = _DecoratedTritonJitKernel( + dispatch if kernel is None else kernel, + warmup_inputs, + None if kernel is None else dispatch, + ) + update_wrapper(owner, dispatch, updated=()) + return owner + + return decorate + + @dataclass(frozen=True) class TritonPointerInputVariant: # Named pointer-alignment variant for compile-only Triton warmup. diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index bbb0a66d538c..e56ae50ebbca 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -41,9 +41,6 @@ from vllm.model_executor.warmup.replayssm_warmup import ( replayssm_autotune_warmup, ) -from vllm.model_executor.warmup.spec_decode_rejection_warmup import ( - spec_decode_rejection_warmup, -) from vllm.model_executor.warmup.watermark_sample_warmup import ( watermark_sample_warmup, ) @@ -178,7 +175,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): logger.info("JIT kernel warmup starting.") jit_warmup_start = time.perf_counter() try: - worker.model_runner.jit_warmup_registry.warmup() + registry = ( + worker.model_runner.jit_warmup_registry # type: ignore[attr-defined] + ) + registry.warmup() except Exception: logger.exception( "JIT kernel warmup failed after %.2fs.", @@ -209,7 +209,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): # Run next so input-prep kernels JIT against pristine runner state. if enable_jit_warmup: kimi_k3_triton_warmup(worker) - spec_decode_rejection_warmup(worker) watermark_sample_warmup(worker) qwen4_exp_qsa_triton_warmup(worker) diff --git a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py deleted file mode 100644 index 6f6f08a2d58f..000000000000 --- a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py +++ /dev/null @@ -1,139 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up speculative-decoding rejection kernels.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -import torch - -from vllm.logger import init_logger - -if TYPE_CHECKING: - from vllm.v1.worker.gpu_worker import Worker - -logger = init_logger(__name__) - - -@torch.inference_mode() -def spec_decode_rejection_warmup(worker: Worker) -> None: - spec_config = worker.vllm_config.speculative_config - if spec_config is None: - return - - from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( - rejection_sample, - ) - - model_config = worker.vllm_config.model_config - vocab_size = model_config.get_vocab_size() - num_spec = spec_config.num_speculative_tokens - if num_spec <= 0 or vocab_size <= 0: - return - - rejection_method = getattr(spec_config, "rejection_sample_method", None) - use_block_verification = rejection_method == "block" - use_synthetic = rejection_method == "synthetic" - # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime. - use_draft_logits = ( - getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic" - ) - - device = worker.device - num_reqs = 1 - tokens_per_req = num_spec + 1 - num_logits = num_reqs * tokens_per_req - - watermark_config = getattr(worker.vllm_config, "watermark_config", None) - watermark_kwargs: dict[str, Any] = {} - if watermark_config is not None: - from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key - - try: - watermark_key = speculative_target_watermark_key(watermark_config) - except Exception: - logger.warning( - "Skipping spec-decode rejection sampler warmup: could not derive " - "the watermark recovery key.", - exc_info=True, - ) - return - # Tensor dtypes and integer widths affect Triton specialization. - watermark_kwargs = { - "contexts": torch.zeros( - (num_logits, watermark_config.context_width), - dtype=torch.int32, - device=device, - ), - "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device), - "watermark_key": watermark_key, - } - - # Sampling-parameter processing may promote either distribution to fp32. - model_dtype = model_config.dtype - warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]] - if use_draft_logits: - warmup_dtype_pairs = { - (model_dtype, model_dtype), - (torch.float32, torch.float32), - (torch.float32, model_dtype), - (model_dtype, torch.float32), - } - else: - warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)} - - logger.info( - "Warming up spec-decode rejection sampler kernels " - "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).", - vocab_size, - num_spec, - [(str(t), str(d)) for t, d in warmup_dtype_pairs], - use_block_verification, - watermark_config is not None, - ) - for tgt_dtype, draft_dtype in warmup_dtype_pairs: - target_logits = torch.zeros( - (num_logits, vocab_size), dtype=tgt_dtype, device=device - ) - draft_logits = ( - torch.zeros( - (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device - ) - if draft_dtype is not None - else None - ) - synthetic_rates = ( - torch.full((num_spec,), 0.5, dtype=torch.float32, device=device) - if use_synthetic - else None - ) - try: - rejection_sample( - target_logits=target_logits, - draft_logits=draft_logits, - draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device), - cu_num_logits=torch.tensor( - [0, num_logits], dtype=torch.int32, device=device - ), - pos=torch.zeros(num_logits, dtype=torch.int64, device=device), - idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), - expanded_idx_mapping=torch.zeros( - num_logits, dtype=torch.int64, device=device - ), - expanded_local_pos=torch.arange( - num_logits, dtype=torch.int32, device=device - ), - temperature=torch.zeros(num_reqs, dtype=torch.float32, device=device), - seed=torch.full((num_reqs,), 42, dtype=torch.int64, device=device), - num_speculative_steps=num_spec, - synthetic_conditional_rates=synthetic_rates, - use_fp64=False, - use_block_verification=use_block_verification, - **watermark_kwargs, - ) - except Exception: - logger.warning( - "Skipping spec-decode rejection sampler warmup.", exc_info=True - ) - return diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index e6ba6ac5c659..192327cf4a5a 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -175,6 +175,15 @@ def __init__( projection_size = self.head_dim * self.num_heads self.local_projection_size = divide(projection_size, self.tp_size) + _, recurrent_state_shape = self.get_state_shape() + _, recurrent_state_dtype = self.get_state_dtype() + scatter_states.register_warmup( + state_shape=recurrent_state_shape, + state_dtype=recurrent_state_dtype, + indices_dtype=torch.int32, + max_num_tokens=vllm_config.scheduler_config.max_num_seqs, + ) + # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches). # Order matches checkpoint's fused_qkvbfg_a_proj convention. # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks. diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index d714976d85ee..aeb82dec4232 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -13,11 +13,27 @@ from vllm.triton_utils import HAS_TRITON if HAS_TRITON: - from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton + from vllm.v1.sample.ops.topk_topp_triton import ( + _topk_topp, + _topp_split_mask, + _topp_split_stats, + _topp_split_step, + apply_top_k_top_p_triton, + ) logger = init_logger(__name__) +def register_top_k_top_p_warmups() -> None: + """Register every native accelerator sampling kernel used at runtime.""" + if HAS_TRITON and not current_platform.is_cpu(): + _topk_topp.register_warmup() + if current_platform.is_cuda_alike(): + _topp_split_stats.register_warmup() + _topp_split_step.register_warmup() + _topp_split_mask.register_warmup() + + def _skip_aiter_sampler_on_gfx1250() -> bool: # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU. from vllm.platforms.rocm import on_gfx1250 @@ -128,6 +144,9 @@ def __init__( else: self.forward = self.forward_native + # Every accelerator backend can fall back to native sampling at runtime. + register_top_k_top_p_warmups() + def forward_native( self, logits: torch.Tensor, diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 7c1f080fcdf7..695ba624f77f 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -9,8 +9,21 @@ """ +from typing import Any + import torch +from vllm.model_executor.warmup.jit_warmup import ( + WarmupChoices, + WarmupIntRange, + _when, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import next_power_of_2 @@ -863,6 +876,90 @@ def _topk_topp_kernel( tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n) +def _max_sampler_batch_size(vllm_config: Any) -> int: + return vllm_config.scheduler_config.max_num_seqs * ( + vllm_config.num_speculative_tokens + 1 + ) + + +def _topk_topp_warmup_inputs(vllm_config: Any) -> dict[str, Any]: + vocab_size = vllm_config.model_config.get_vocab_size() + split_enabled = current_platform.is_cuda_alike() + mode: Any = WarmupChoices((True, False), (True, True), (False, True)) + logits_stride: Any = WarmupChoices(16, 2) + batch_size: Any = WarmupIntRange(1, _max_sampler_batch_size(vllm_config) + 1) + topk_enabled = mode[0] + topp_enabled = mode[1] + _when(topk_enabled or not (split_enabled and batch_size <= _SPLIT_MAX_BATCH)) + logits = TritonWarmupTensor( + torch.float32, + shape=(batch_size, vocab_size), + strides=(logits_stride, 1), + ) + return dict( + logits=logits, + buffer=TritonWarmupTensor(torch.float32), + percentile_to_std_table=TritonWarmupTensor(torch.float32), + normal_cdf_to_sigma_table=TritonWarmupTensor(torch.float32), + k=TritonWarmupTensor(torch.int32) if topk_enabled else None, + p=TritonWarmupTensor(torch.float32) if topp_enabled else None, + mask_value=float("-inf"), + num_sm=num_compute_units(), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_topk_topp_kernel, + warmup_inputs=_topk_topp_warmup_inputs, +) +def _topk_topp( + logits: torch.Tensor, + buffer: torch.Tensor, + percentile_to_std_table: torch.Tensor, + normal_cdf_to_sigma_table: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + mask_value: float, + num_sm: int, +) -> DispatchSpec: + batch_size, vocab_size = logits.shape + topk_enabled = k is not None + topp_enabled = p is not None + k_ptr = k if k is not None else logits + p_ptr = p if p is not None else logits + num_programs = min(num_sm, batch_size) + # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. + # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot + # approximation; use smaller tiles for accurate top-p results. + if logits.device.type == "cpu": + block_size, block_size_trunc, num_warps = 256, 128, 4 + elif logits.device.type == "xpu": + block_size, block_size_trunc, num_warps = 4096, 2048, 4 + else: + block_size, block_size_trunc, num_warps = 8192, 4096, 8 + # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so + # per-tile latency bounds kernel latency, and Triton's default of 4 + # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is + # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not. + split_covers_ponly = ( + current_platform.is_cuda_alike() + and topp_enabled + and batch_size <= _SPLIT_MAX_BATCH + ) + return (num_programs,), dict( + K=k_ptr, + P=p_ptr, + BATCH_SIZE=batch_size, + VOCAB_SIZE=vocab_size, + BLOCK_SIZE=block_size, + BLOCK_SIZE_TRUNC=block_size_trunc, + TOPK_ENABLED=topk_enabled, + TOPP_ENABLED=topp_enabled, + SPLIT_COVERS_PONLY=split_covers_ponly, + num_warps=num_warps, + ) + + # --------------------------------------------------------------------------- # Split-row pipeline for small batches. # @@ -892,6 +989,13 @@ def _topk_topp_kernel( _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {} +def _topp_split_count(batch_size: int, num_sm: int) -> int: + splits = 1 + while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm: + splits *= 2 + return splits + + @triton.jit def _topp_sb_stats_kernel( LOGITS, @@ -939,6 +1043,51 @@ def _topp_sb_stats_kernel( tl.store(STATS + base + 3, cnt) +def _topp_split_stats_warmup_inputs(vllm_config: Any) -> dict[str, Any]: + vocab_size = vllm_config.model_config.get_vocab_size() + num_sm = num_compute_units() + logits_stride: Any = WarmupChoices(16, 2) + batch_size: Any = WarmupIntRange( + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 + ) + has_k: Any = WarmupChoices(False, True) + logits = TritonWarmupTensor( + torch.float32, + shape=(batch_size, vocab_size), + strides=(logits_stride, 1), + ) + return dict( + logits=logits, + stats=TritonWarmupTensor(torch.float32), + k=TritonWarmupTensor(torch.int32) if has_k else None, + p=TritonWarmupTensor(torch.float32), + num_sm=num_sm, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_topp_sb_stats_kernel, + warmup_inputs=_topp_split_stats_warmup_inputs, +) +def _topp_split_stats( + logits: torch.Tensor, + stats: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor, + num_sm: int, +) -> DispatchSpec: + batch_size, vocab_size = logits.shape + splits = _topp_split_count(batch_size, num_sm) + return (logits.shape[0] * splits,), dict( + K=k if k is not None else logits, + HAS_K=k is not None, + VOCAB_SIZE=vocab_size, + S=splits, + BLOCK=8192, + num_warps=8, + ) + + @triton.jit def _topp_sb_row_stats(STATS, row, S: tl.constexpr): """Combine per-slice partials into row (max, sum_exp, min, n_finite).""" @@ -1183,6 +1332,58 @@ def _topp_sb_step_kernel( tl.store(pbase + fidx * 3 + 2, nmin) +def _topp_split_step_warmup_inputs(vllm_config: Any) -> dict[str, Any]: + vocab_size = vllm_config.model_config.get_vocab_size() + num_sm = num_compute_units() + logits_stride: Any = WarmupChoices(16, 2) + batch_size: Any = WarmupIntRange( + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 + ) + round: Any = WarmupIntRange(0, _SPLIT_ROUNDS) + has_k: Any = WarmupChoices(False, True) + logits = TritonWarmupTensor( + torch.float32, + shape=(batch_size, vocab_size), + strides=(logits_stride, 1), + ) + return dict( + logits=logits, + stats=TritonWarmupTensor(torch.float32), + parts=TritonWarmupTensor(torch.float32), + k=TritonWarmupTensor(torch.int32) if has_k else None, + p=TritonWarmupTensor(torch.float32), + round=round, + num_sm=num_sm, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_topp_sb_step_kernel, + warmup_inputs=_topp_split_step_warmup_inputs, +) +def _topp_split_step( + logits: torch.Tensor, + stats: torch.Tensor, + parts: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor, + round: int, + num_sm: int, +) -> DispatchSpec: + batch_size, vocab_size = logits.shape + splits = _topp_split_count(batch_size, num_sm) + return (logits.shape[0] * splits,), dict( + K=k if k is not None else logits, + HAS_K=k is not None, + S=splits, + F=_SPLIT_FANOUT, + NUM_ROUNDS=_SPLIT_ROUNDS, + VOCAB_SIZE=vocab_size, + BLOCK=2048, + num_warps=8, + ) + + @triton.jit def _topp_sb_mask_kernel( LOGITS, @@ -1313,6 +1514,57 @@ def _topp_sb_mask_kernel( tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n) +def _topp_split_mask_warmup_inputs(vllm_config: Any) -> dict[str, Any]: + vocab_size = vllm_config.model_config.get_vocab_size() + num_sm = num_compute_units() + logits_stride: Any = WarmupChoices(16, 2) + batch_size: Any = WarmupIntRange( + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1 + ) + has_k: Any = WarmupChoices(False, True) + logits = TritonWarmupTensor( + torch.float32, + shape=(batch_size, vocab_size), + strides=(logits_stride, 1), + ) + return dict( + logits=logits, + stats=TritonWarmupTensor(torch.float32), + parts=TritonWarmupTensor(torch.float32), + k=TritonWarmupTensor(torch.int32) if has_k else None, + p=TritonWarmupTensor(torch.float32), + mask_value=float("-inf"), + num_sm=num_sm, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_topp_sb_mask_kernel, + warmup_inputs=_topp_split_mask_warmup_inputs, +) +def _topp_split_mask( + logits: torch.Tensor, + stats: torch.Tensor, + parts: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor, + mask_value: float, + num_sm: int, +) -> DispatchSpec: + batch_size, vocab_size = logits.shape + splits = _topp_split_count(batch_size, num_sm) + return (logits.shape[0] * splits,), dict( + K=k if k is not None else logits, + HAS_K=k is not None, + S=splits, + F=_SPLIT_FANOUT, + NUM_ROUNDS=_SPLIT_ROUNDS, + VOCAB_SIZE=vocab_size, + BLOCK=8192, + num_warps=8, + ) + + def _apply_topp_split( logits: torch.Tensor, k: torch.Tensor | None, @@ -1321,11 +1573,6 @@ def _apply_topp_split( num_sm: int, ) -> None: """Split-row top-p pipeline for small batches; masks logits in place.""" - batch_size, vocab_size = logits.shape - splits = 1 - while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm: - splits *= 2 - ws = _TRITON_SPLIT_CACHE.get(logits.device) if ws is None: ws = { @@ -1336,53 +1583,31 @@ def _apply_topp_split( } _TRITON_SPLIT_CACHE[logits.device] = ws - k_ptr = k if k is not None else logits # dummy pointer when HAS_K=False - has_k = k is not None - - _topp_sb_stats_kernel[(batch_size * splits,)]( + _topp_split_stats( logits, - logits.stride(0), ws["stats"], - k_ptr, + k, p, - HAS_K=has_k, - VOCAB_SIZE=vocab_size, - S=splits, - BLOCK=8192, - num_warps=8, + num_sm, ) for round_i in range(_SPLIT_ROUNDS): - _topp_sb_step_kernel[(batch_size * splits,)]( + _topp_split_step( logits, - logits.stride(0), ws["stats"], ws["parts"], - k_ptr, + k, p, round_i, - HAS_K=has_k, - S=splits, - F=_SPLIT_FANOUT, - NUM_ROUNDS=_SPLIT_ROUNDS, - VOCAB_SIZE=vocab_size, - BLOCK=2048, - num_warps=8, + num_sm, ) - _topp_sb_mask_kernel[(batch_size * splits,)]( + _topp_split_mask( logits, - logits.stride(0), ws["stats"], ws["parts"], - k_ptr, + k, p, - HAS_K=has_k, - MASK_VALUE=mask_value, - S=splits, - F=_SPLIT_FANOUT, - NUM_ROUNDS=_SPLIT_ROUNDS, - VOCAB_SIZE=vocab_size, - BLOCK=8192, - num_warps=8, + mask_value, + num_sm, ) @@ -1477,39 +1702,15 @@ def apply_top_k_top_p_triton( else: normal_cdf_to_sigma_table, percentile_to_std_table = tables - # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. - # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot - # approximation; use smaller tiles for accurate top-p results. - launch_kwargs = {} - if logits.device.type == "cpu": - block_size, block_size_trunc = 256, 128 - elif logits.device.type == "xpu": - block_size, block_size_trunc = 4096, 2048 - else: - block_size, block_size_trunc = 8192, 4096 - # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so - # per-tile latency bounds kernel latency, and Triton's default of 4 - # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is - # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not. - launch_kwargs["num_warps"] = 8 - - _topk_topp_kernel[(NUM_PROGRAMS,)]( + _topk_topp( logits, - logits.stride(0), buffer, percentile_to_std_table, normal_cdf_to_sigma_table, - k_ptr, - p_ptr, - BATCH_SIZE=batch_size, - MASK_VALUE=mask_value, - VOCAB_SIZE=vocab_size, - BLOCK_SIZE=block_size, - BLOCK_SIZE_TRUNC=block_size_trunc, - TOPK_ENABLED=topk_enabled, - TOPP_ENABLED=topp_enabled, - SPLIT_COVERS_PONLY=use_split, - **launch_kwargs, + k_ptr if topk_enabled else None, + p_ptr if topp_enabled else None, + mask_value, + num_sm, ) if use_split: _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index a1baf0a80851..5aa83454739e 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -5,13 +5,19 @@ from collections.abc import Sequence from dataclasses import replace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import WarmupChoices +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor @@ -88,6 +94,14 @@ def __init__( device=device, ) self.synthetic_mode = self.synthetic_conditional_rates is not None + _rejection_greedy_sample.register_warmup(synthetic_mode=self.synthetic_mode) + _rejection_random_sample.register_warmup( + synthetic_mode=self.synthetic_mode, + ) + _expand.register_warmup() + _sample_recovered_tokens.register_warmup( + use_fp64_gumbel=self.use_fp64_gumbel, + ) def forward( self, @@ -453,7 +467,7 @@ def rejection_sample( if not sampling_metadata.all_random: # Rejection sampling for greedy sampling requests. target_argmax = target_logits.argmax(dim=-1) - rejection_greedy_sample_kernel[(batch_size,)]( + _rejection_greedy_sample( output_token_ids, cu_num_draft_tokens, draft_token_ids, @@ -463,7 +477,6 @@ def rejection_sample( max_spec_len, uniform_probs, synthetic_conditional_rates, - SYNTHETIC_MODE=synthetic_mode, ) if sampling_metadata.all_greedy: return output_token_ids @@ -488,7 +501,7 @@ def rejection_sample( # Rejection sampling for random sampling requests. assert uniform_probs is not None - rejection_random_sample_kernel[(batch_size,)]( + _rejection_random_sample( output_token_ids, cu_num_draft_tokens, draft_token_ids, @@ -499,10 +512,7 @@ def rejection_sample( uniform_probs, is_greedy, max_spec_len, - vocab_size, synthetic_conditional_rates, - NO_DRAFT_PROBS=draft_probs is None, - SYNTHETIC_MODE=synthetic_mode, ) return output_token_ids @@ -594,13 +604,12 @@ def expand_batch_to_tokens( batch_size = x.shape[0] assert cu_num_tokens.shape[0] == batch_size expanded_x = x.new_empty(num_tokens) - expand_kernel[(batch_size,)]( + _expand( expanded_x, x, cu_num_tokens, replace_from, replace_to, - MAX_NUM_TOKENS=MAX_SPEC_LEN, # To avoid recompilation. ) return expanded_x @@ -694,18 +703,14 @@ def sample_recovered_tokens( inv_q = q.reciprocal() recovered_token_ids = torch.empty_like(draft_token_ids) - BLOCK_SIZE = 8192 - sample_recovered_tokens_kernel[(batch_size, max_spec_len)]( + _sample_recovered_tokens( recovered_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, inv_q, - vocab_size, - BLOCK_SIZE, - NO_DRAFT_PROBS=draft_probs is None, - USE_FP64_GUMBEL=use_fp64_gumbel, + max_spec_len=max_spec_len, ) return recovered_token_ids @@ -769,6 +774,50 @@ def rejection_greedy_sample_kernel( ) +def _rejection_greedy_sample_warmup_inputs( + *, synthetic_mode: bool +) -> dict[str, object]: + all_greedy = WarmupChoices(False, True) + int32 = TritonWarmupTensor(torch.int32) + return dict( + output_token_ids=int32, + cu_num_draft_tokens=int32, + draft_token_ids=int32, + target_argmax=TritonWarmupTensor(torch.int64), + bonus_token_ids=TritonWarmupTensor(torch.int64), + is_greedy=None if all_greedy else TritonWarmupTensor(torch.bool), + max_spec_len=5, + uniform_probs=( + None + if all_greedy and not synthetic_mode + else TritonWarmupTensor(torch.float64) + ), + synthetic_conditional_rates=( + TritonWarmupTensor(torch.float32) if synthetic_mode else None + ), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=rejection_greedy_sample_kernel, + warmup_inputs=_rejection_greedy_sample_warmup_inputs, +) +def _rejection_greedy_sample( + output_token_ids: torch.Tensor, + cu_num_draft_tokens: torch.Tensor, + draft_token_ids: torch.Tensor, + target_argmax: torch.Tensor, + bonus_token_ids: torch.Tensor, + is_greedy: torch.Tensor | None, + max_spec_len: int, + uniform_probs: torch.Tensor | None, + synthetic_conditional_rates: torch.Tensor | None, +) -> DispatchSpec: + return (output_token_ids.shape[0],), dict( + SYNTHETIC_MODE=synthetic_conditional_rates is not None, + ) + + # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation. @triton.jit(do_not_specialize=["max_spec_len"]) def rejection_random_sample_kernel( @@ -845,6 +894,57 @@ def rejection_random_sample_kernel( ) +def _rejection_random_sample_warmup_inputs( + *, synthetic_mode: bool +) -> dict[str, object]: + vocab_size: Any = WarmupChoices(2, 16) + no_draft_probs = WarmupChoices(False, True) + int32 = TritonWarmupTensor(torch.int32) + probabilities = TritonWarmupTensor( + torch.float32, + shape=(1, vocab_size), + ) + return dict( + output_token_ids=int32, + cu_num_draft_tokens=int32, + draft_token_ids=int32, + draft_probs=None if no_draft_probs else probabilities, + target_probs=probabilities, + bonus_token_ids=TritonWarmupTensor(torch.int64), + recovered_token_ids=int32, + uniform_probs=TritonWarmupTensor(torch.float64), + is_greedy=TritonWarmupTensor(torch.bool), + max_spec_len=5, + synthetic_conditional_rates=( + TritonWarmupTensor(torch.float32) if synthetic_mode else None + ), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=rejection_random_sample_kernel, + warmup_inputs=_rejection_random_sample_warmup_inputs, +) +def _rejection_random_sample( + output_token_ids: torch.Tensor, + cu_num_draft_tokens: torch.Tensor, + draft_token_ids: torch.Tensor, + draft_probs: torch.Tensor | None, + target_probs: torch.Tensor, + bonus_token_ids: torch.Tensor, + recovered_token_ids: torch.Tensor, + uniform_probs: torch.Tensor, + is_greedy: torch.Tensor, + max_spec_len: int, + synthetic_conditional_rates: torch.Tensor | None, +) -> DispatchSpec: + return (output_token_ids.shape[0],), dict( + vocab_size=target_probs.shape[1], + NO_DRAFT_PROBS=draft_probs is None, + SYNTHETIC_MODE=synthetic_conditional_rates is not None, + ) + + # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation. @triton.jit(do_not_specialize=["replace_from", "replace_to"]) def expand_kernel( @@ -869,6 +969,31 @@ def expand_kernel( tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens) +def _expand_warmup_inputs() -> dict[str, object]: + dtype = WarmupChoices(torch.float32, torch.int32, torch.int64) + data = TritonWarmupTensor(dtype) + return dict( + output=data, + input=data, + cu_num_tokens=TritonWarmupTensor(torch.int32), + replace_from=0, + replace_to=0, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=expand_kernel, warmup_inputs=_expand_warmup_inputs +) +def _expand( + output: torch.Tensor, + input: torch.Tensor, + cu_num_tokens: torch.Tensor, + replace_from: int, + replace_to: int, +) -> DispatchSpec: + return (input.shape[0],), dict(MAX_NUM_TOKENS=MAX_SPEC_LEN) + + @triton.jit def sample_recovered_tokens_kernel( output_token_ids_ptr, # [num_tokens] @@ -951,3 +1076,49 @@ def sample_recovered_tokens_kernel( recovered_id = tl.minimum(recovered_id, vocab_size - 1) tl.store(output_token_ids_ptr + token_idx, recovered_id) + + +def _sample_recovered_tokens_warmup_inputs( + *, use_fp64_gumbel: bool +) -> dict[str, object]: + vocab_size: Any = WarmupChoices(2, 16) + no_draft_probs = WarmupChoices(False, True) + int32 = TritonWarmupTensor(torch.int32) + probabilities = TritonWarmupTensor( + torch.float32, + shape=(1, vocab_size), + ) + return dict( + output_token_ids=int32, + cu_num_draft_tokens=int32, + draft_token_ids=int32, + draft_probs=None if no_draft_probs else probabilities, + target_probs=probabilities, + inv_q=TritonWarmupTensor( + torch.float64 if use_fp64_gumbel else torch.float32, + shape=(1, vocab_size), + ), + max_spec_len=5, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=sample_recovered_tokens_kernel, + warmup_inputs=_sample_recovered_tokens_warmup_inputs, +) +def _sample_recovered_tokens( + output_token_ids: torch.Tensor, + cu_num_draft_tokens: torch.Tensor, + draft_token_ids: torch.Tensor, + draft_probs: torch.Tensor | None, + target_probs: torch.Tensor, + inv_q: torch.Tensor, + *, + max_spec_len: int, +) -> DispatchSpec: + return (cu_num_draft_tokens.shape[0], max_spec_len), dict( + vocab_size=target_probs.shape[1], + BLOCK_SIZE=8192, + NO_DRAFT_PROBS=draft_probs is None, + USE_FP64_GUMBEL=inv_q.dtype == torch.float64, + ) diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 42d8ef035c66..284db7dee71f 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -13,7 +13,7 @@ from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer from vllm.v1.spec_decode.utils import ( - copy_and_expand_dflash_inputs_kernel, + _copy_and_expand_dflash_inputs, next_power_of_2, ) @@ -79,6 +79,13 @@ def __init__( self.dflash_causal = not dflash_has_any_non_causal( self.draft_model_config.hf_config ) + _copy_and_expand_dflash_inputs.register_warmup( + block_table_stride=(self.max_model_len + self.block_size - 1) + // self.block_size, + parallel_drafting_token_id=self.parallel_drafting_token_id, + block_size=self.block_size, + num_speculative_tokens=self.num_speculative_tokens, + ) @override def _create_draft_vllm_config(self) -> VllmConfig: @@ -137,11 +144,9 @@ def set_inputs_first_pass( max_ctx_per_req = cad.max_query_len max_tokens_per_req = max_ctx_per_req + num_query_per_req BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req)) - num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE - grid = (batch_size, num_blocks) has_num_rejected = num_rejected_tokens_gpu is not None - copy_and_expand_dflash_inputs_kernel[grid]( + _copy_and_expand_dflash_inputs( # Inputs next_token_ids_ptr=next_token_ids, target_positions_ptr=target_positions, @@ -166,8 +171,9 @@ def set_inputs_first_pass( num_query_per_req=num_query_per_req, num_speculative_tokens=self.num_speculative_tokens, total_input_tokens=num_context, - BLOCK_SIZE=BLOCK_SIZE, - HAS_NUM_REJECTED=has_num_rejected, + max_tokens_per_req=max_tokens_per_req, + triton_block_size=BLOCK_SIZE, + has_num_rejected=has_num_rejected, ) query_slot_mapping = self._slot_mapping_buffer[:num_query_total] diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 2edbd27d6ead..c0e456521ca0 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -52,10 +52,11 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.spec_decode.utils import ( PADDING_SLOT_ID, + _copy_and_expand_eagle_inputs, + _eagle_prepare_inputs_padded, + _eagle_prepare_next_token_padded, + _eagle_step_slot_mapping_metadata, compute_new_slot_mapping, - copy_and_expand_eagle_inputs_kernel, - eagle_prepare_inputs_padded_kernel, - eagle_prepare_next_token_padded_kernel, eagle_step_update_slot_mapping_and_metadata, extend_all_queries_by_N, next_power_of_2, @@ -140,6 +141,23 @@ def __init__( self.max_batch_size = vllm_config.scheduler_config.max_num_seqs self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + _eagle_step_slot_mapping_metadata.register_warmup( + max_model_len=self.max_model_len, + max_batch_size=self.max_batch_size, + ) + _eagle_prepare_inputs_padded.register_warmup(max_batch_size=self.max_batch_size) + _eagle_prepare_next_token_padded.register_warmup( + vocab_size=vllm_config.model_config.get_vocab_size(), + num_sampled_tokens_per_req=self.num_speculative_tokens + 1, + max_batch_size=self.max_batch_size, + ) + if self.needs_extra_input_slots: + _copy_and_expand_eagle_inputs.register_warmup( + parallel_drafting_token_id=self.parallel_drafting_token_id, + num_padding_slots_per_request=self.extra_slots_per_request, + shift_input_ids=self.pass_hidden_states_to_model, + max_num_tokens=self.max_num_tokens, + ) self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32) # Can be specialized by methods like DFlash to reduce the limit @@ -872,13 +890,12 @@ def set_inputs_first_pass( ) # Kernel grid: one program per request (row) - grid = (batch_size, num_blocks) query_start_loc = cad.query_start_loc query_end_loc = cad.query_start_loc[1:] - 1 if num_rejected_tokens_gpu is not None: query_end_loc = query_end_loc - num_rejected_tokens_gpu - copy_and_expand_eagle_inputs_kernel[grid]( + _copy_and_expand_eagle_inputs( # (Padded) Inputs from the target model target_token_ids_ptr=target_token_ids, target_positions_ptr=target_positions, @@ -901,6 +918,8 @@ def set_inputs_first_pass( num_padding_slots_per_request=self.extra_slots_per_request, shift_input_ids=self.pass_hidden_states_to_model, BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS, + batch_size=batch_size, + num_blocks=num_blocks, ) if self.pass_hidden_states_to_model: assert self.parallel_drafting_hidden_state_tensor is not None @@ -1065,11 +1084,10 @@ def prepare_next_token_ids_padded( valid_sampled_tokens_count = next_token_ids.new_empty(batch_size) # Kernel grid: one program per request (row) - grid = (batch_size,) # Find the next power of 2 for block sizes BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens) - eagle_prepare_next_token_padded_kernel[grid]( + _eagle_prepare_next_token_padded( sampled_token_ids, discard_request_mask, backup_tokens_gpu, @@ -1108,8 +1126,7 @@ def prepare_inputs_padded( (num_reqs,), dtype=torch.int32, device=device ) - grid = (num_reqs,) - eagle_prepare_inputs_padded_kernel[grid]( + _eagle_prepare_inputs_padded( spec_decode_metadata.cu_num_draft_tokens, valid_sampled_tokens_count, common_attn_metadata.query_start_loc, diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 65b9408a8901..fbfe4438422a 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -1,7 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch +from vllm.model_executor.warmup.jit_warmup import ( + WarmupChoices, + WarmupIntRange, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import ( @@ -85,6 +96,49 @@ def eagle_step_slot_mapping_metadata_kernel( tl.store(seq_lens_ptr + req_idx, new_seq_len) +def _eagle_step_warmup_inputs(*, max_model_len: int, max_batch_size: int): + block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2) + n_blocks_per_req = triton.cdiv(max_model_len, block_size) + return dict( + positions=TritonWarmupTensor(torch.int64), + block_table=TritonWarmupTensor( + torch.int32, shape=(1, n_blocks_per_req), strides=(n_blocks_per_req, 1) + ), + block_table_stride=n_blocks_per_req, + seq_lens=TritonWarmupTensor(torch.int32), + out_clamped_positions=TritonWarmupTensor(torch.int64), + out_slot_mapping=TritonWarmupTensor(torch.int64), + block_size=block_size, + max_model_len=max_model_len, + n_blocks_per_req=n_blocks_per_req, + PAD_ID=PADDING_SLOT_ID, + batch_size=WarmupIntRange(1, max_batch_size + 1), + input_batch_size=1, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=eagle_step_slot_mapping_metadata_kernel, + warmup_inputs=_eagle_step_warmup_inputs, +) +def _eagle_step_slot_mapping_metadata( + positions: torch.Tensor, + block_table: torch.Tensor, + block_table_stride: int, + seq_lens: torch.Tensor, + out_clamped_positions: torch.Tensor, + out_slot_mapping: torch.Tensor, + *, + block_size: int, + max_model_len: int, + n_blocks_per_req: int, + PAD_ID: int, + batch_size: int, + input_batch_size: int, +) -> DispatchSpec: + return (input_batch_size,), {} + + def eagle_step_update_slot_mapping_and_metadata( positions_1d: torch.Tensor, block_table_tensor: torch.Tensor, @@ -118,7 +172,7 @@ def eagle_step_update_slot_mapping_and_metadata( input_batch_size = batch_size n_blocks_per_req = block_table_tensor.shape[1] - eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)]( + _eagle_step_slot_mapping_metadata( positions_1d, block_table_tensor, block_table_tensor.stride(0), @@ -130,6 +184,7 @@ def eagle_step_update_slot_mapping_and_metadata( n_blocks_per_req=n_blocks_per_req, PAD_ID=PADDING_SLOT_ID, batch_size=batch_size, + input_batch_size=input_batch_size, ) @@ -175,6 +230,33 @@ def eagle_prepare_inputs_padded_kernel( tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens) +def _eagle_prepare_inputs_warmup_inputs(*, max_batch_size: int): + int32 = TritonWarmupTensor(torch.int32) + return dict( + cu_num_draft_tokens=int32, + valid_sampled_tokens_count=int32, + query_start_loc_gpu=int32, + token_indices_to_sample=int32, + num_rejected_tokens_gpu=int32, + num_reqs=WarmupIntRange(1, max_batch_size + 1), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=eagle_prepare_inputs_padded_kernel, + warmup_inputs=_eagle_prepare_inputs_warmup_inputs, +) +def _eagle_prepare_inputs_padded( + cu_num_draft_tokens: torch.Tensor, + valid_sampled_tokens_count: torch.Tensor, + query_start_loc_gpu: torch.Tensor, + token_indices_to_sample: torch.Tensor, + num_rejected_tokens_gpu: torch.Tensor, + num_reqs: int, +) -> DispatchSpec: + return (num_reqs,), {} + + @triton.jit def eagle_prepare_next_token_padded_kernel( sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req] @@ -238,6 +320,49 @@ def eagle_prepare_next_token_padded_kernel( tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) +def _eagle_prepare_next_token_warmup_inputs( + *, vocab_size: int, num_sampled_tokens_per_req: int, max_batch_size: int +): + int32 = TritonWarmupTensor(torch.int32) + sampled_width: Any = WarmupIntRange(1, num_sampled_tokens_per_req + 1) + return dict( + sampled_token_ids=TritonWarmupTensor( + torch.int32, + shape=(1, sampled_width), + strides=(sampled_width, 1), + ), + discard_request_mask=TritonWarmupTensor(torch.bool), + backup_next_token_ids=int32, + next_token_ids=int32, + valid_sampled_tokens_count=int32, + vocab_size=vocab_size, + num_sampled_tokens_per_req=sampled_width, + num_reqs=WarmupIntRange(1, max_batch_size + 1), + stride_sampled_token_ids=sampled_width, + BLOCK_SIZE_TOKENS=next_power_of_2(sampled_width), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=eagle_prepare_next_token_padded_kernel, + warmup_inputs=_eagle_prepare_next_token_warmup_inputs, +) +def _eagle_prepare_next_token_padded( + sampled_token_ids: torch.Tensor, + discard_request_mask: torch.Tensor, + backup_next_token_ids: torch.Tensor, + next_token_ids: torch.Tensor, + valid_sampled_tokens_count: torch.Tensor, + vocab_size: int, + num_sampled_tokens_per_req: int, + num_reqs: int, + stride_sampled_token_ids: int, + *, + BLOCK_SIZE_TOKENS: int, +) -> DispatchSpec: + return (num_reqs,), {} + + def compute_new_slot_mapping( cad: CommonAttentionMetadata, new_positions: torch.Tensor, @@ -454,6 +579,68 @@ def copy_and_expand_eagle_inputs_kernel( ) +def _copy_and_expand_eagle_warmup_inputs( + *, + parallel_drafting_token_id: int, + num_padding_slots_per_request: int, + shift_input_ids: bool, + max_num_tokens: int, +): + int32 = TritonWarmupTensor(torch.int32) + int64 = TritonWarmupTensor(torch.int64) + boolean = TritonWarmupTensor(torch.bool) + return dict( + target_token_ids_ptr=int32, + target_positions_ptr=int64, + next_token_ids_ptr=int32, + out_input_ids_ptr=int32, + out_positions_ptr=int64, + out_is_rejected_token_mask_ptr=boolean, + out_is_masked_token_mask_ptr=boolean, + out_new_token_indices_ptr=int32, + out_hidden_state_mapping_ptr=int32, + query_start_loc_ptr=int32, + query_end_loc_ptr=int32, + padding_token_id=0, + parallel_drafting_token_id=parallel_drafting_token_id, + total_input_tokens=WarmupIntRange(1, max_num_tokens + 1), + num_padding_slots_per_request=num_padding_slots_per_request, + shift_input_ids=shift_input_ids, + BLOCK_SIZE_TOKENS=WarmupIntRange(1, 257, advance=lambda value: value * 2), + batch_size=1, + num_blocks=1, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=copy_and_expand_eagle_inputs_kernel, + warmup_inputs=_copy_and_expand_eagle_warmup_inputs, +) +def _copy_and_expand_eagle_inputs( + target_token_ids_ptr: torch.Tensor, + target_positions_ptr: torch.Tensor, + next_token_ids_ptr: torch.Tensor, + out_input_ids_ptr: torch.Tensor, + out_positions_ptr: torch.Tensor, + out_is_rejected_token_mask_ptr: torch.Tensor, + out_is_masked_token_mask_ptr: torch.Tensor, + out_new_token_indices_ptr: torch.Tensor, + out_hidden_state_mapping_ptr: torch.Tensor, + query_start_loc_ptr: torch.Tensor, + query_end_loc_ptr: torch.Tensor, + padding_token_id: int, + parallel_drafting_token_id: int, + total_input_tokens: int, + num_padding_slots_per_request: int, + shift_input_ids: bool, + *, + BLOCK_SIZE_TOKENS: int, + batch_size: int, + num_blocks: int, +) -> DispatchSpec: + return (batch_size, num_blocks), {} + + @triton.jit def copy_and_expand_dflash_inputs_kernel( # Inputs @@ -562,6 +749,75 @@ def copy_and_expand_dflash_inputs_kernel( ) +def _copy_and_expand_dflash_warmup_inputs( + *, + block_table_stride: int, + parallel_drafting_token_id: int, + block_size: int, + num_speculative_tokens: int, +): + int32 = TritonWarmupTensor(torch.int32) + int64 = TritonWarmupTensor(torch.int64) + triton_block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2) + has_num_rejected = WarmupChoices(False, True) + return dict( + next_token_ids_ptr=int32, + target_positions_ptr=int64, + out_input_ids_ptr=int32, + out_context_positions_ptr=int64, + out_query_positions_ptr=int64, + out_context_slot_mapping_ptr=int64, + out_query_slot_mapping_ptr=int64, + out_token_indices_ptr=int32, + block_table_ptr=int32, + query_start_loc_ptr=int32, + num_rejected_tokens_ptr=(int32 if has_num_rejected else 0), + block_table_stride=block_table_stride, + parallel_drafting_token_id=parallel_drafting_token_id, + block_size=block_size, + num_query_per_req=1 + num_speculative_tokens, + num_speculative_tokens=num_speculative_tokens, + total_input_tokens=WarmupChoices(1, 2, 16), + max_tokens_per_req=triton_block_size, + triton_block_size=triton_block_size, + has_num_rejected=has_num_rejected, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=copy_and_expand_dflash_inputs_kernel, + warmup_inputs=_copy_and_expand_dflash_warmup_inputs, +) +def _copy_and_expand_dflash_inputs( + *, + next_token_ids_ptr: torch.Tensor, + target_positions_ptr: torch.Tensor, + out_input_ids_ptr: torch.Tensor, + out_context_positions_ptr: torch.Tensor, + out_query_positions_ptr: torch.Tensor, + out_context_slot_mapping_ptr: torch.Tensor, + out_query_slot_mapping_ptr: torch.Tensor, + out_token_indices_ptr: torch.Tensor, + block_table_ptr: torch.Tensor, + query_start_loc_ptr: torch.Tensor, + num_rejected_tokens_ptr: torch.Tensor | int, + block_table_stride: int, + parallel_drafting_token_id: int, + block_size: int, + num_query_per_req: int, + num_speculative_tokens: int, + total_input_tokens: int, + max_tokens_per_req: int, + triton_block_size: int, + has_num_rejected: bool, +) -> DispatchSpec: + num_reqs = query_start_loc_ptr.shape[0] - 1 + num_blocks = triton.cdiv(max_tokens_per_req, triton_block_size) + return (num_reqs, num_blocks), dict( + BLOCK_SIZE=triton_block_size, + ) + + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def update_num_computed_tokens_for_batch_change( num_computed_tokens: torch.Tensor, diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 67ddf83baec5..04ca2b9194e4 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -80,16 +80,15 @@ def _postprocess_triton(self) -> None: # Speculative decoding fallbacks import vllm.v1.sample.rejection_sampler - import vllm.v1.spec_decode.llm_base_proposer import vllm.v1.spec_decode.utils as spec_decode_utils - vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = ( + spec_decode_utils._eagle_prepare_inputs_padded.kernel = ( cpu_tl.eagle_prepare_inputs_padded_kernel ) - vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_next_token_padded_kernel = ( + spec_decode_utils._eagle_prepare_next_token_padded.kernel = ( cpu_tl.eagle_prepare_next_token_padded_kernel ) - vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = ( + spec_decode_utils._copy_and_expand_eagle_inputs.kernel = ( cpu_tl.copy_and_expand_eagle_inputs_kernel ) spec_decode_utils.copy_and_expand_dflash_inputs_kernel = ( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0ce0288f310a..14e0ac46297e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -188,6 +188,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + @JitWarmupRegistry.capture def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -202,7 +203,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.speculative_config is not None and self.speculative_config.use_dspark() ) self.observability_config = vllm_config.observability_config - self.jit_warmup_registry = JitWarmupRegistry(vllm_config) self.device = device self.dtype = self.model_config.dtype @@ -399,7 +399,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: logger.info_once("Loading model from scratch...") # Capture warmup providers selected while constructing the model. - with self.jit_warmup_registry.activate(): + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.vllm_config.model_config, @@ -461,47 +461,50 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: - sampler_kwargs: dict[str, Any] = { - "max_num_reqs": self.max_num_reqs, - "vocab_size": self.vocab_size, - "device": self.device, - "req_states": self.req_states, - "logprobs_mode": self.model_config.logprobs_mode, - "num_speculative_tokens": self.decode_query_len, - "use_fp64_gumbel": self.model_config.use_fp64_gumbel, - "enable_trace_replay": self.model_config.enable_trace_replay, - "reasoning_config": self.vllm_config.reasoning_config, - "return_sampling_mask": self.model_config.return_sampling_mask, - } - if self.vllm_config.watermark_config is None: - self.sampler = Sampler(**sampler_kwargs) - else: - wm_config = self.vllm_config.watermark_config - watermarker = create_watermarker(wm_config) - if self.speculative_config is not None: - watermarker = create_speculative_target_watermarker(watermarker) - self.sampler = GPUWatermarkSampler( - watermarker, - deduplicate_contexts=wm_config.deduplicate_contexts, - deduplicate_contexts_max_history=( - wm_config.deduplicate_contexts_max_history - ), - **sampler_kwargs, - ) - custom = self.model_state.custom_sampler(self.sampler) - - if custom: - self.vllm_config._check_watermarking_unsupported(custom_sampler=True) - self.sampler, self.rejection_sampler = custom - elif self.speculative_config is not None: - self.rejection_sampler = RejectionSampler( - self.sampler, - self.speculative_config, - self.device, - watermark_key=speculative_target_watermark_key( - self.vllm_config.watermark_config - ), - ) + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] + sampler_kwargs: dict[str, Any] = { + "max_num_reqs": self.max_num_reqs, + "vocab_size": self.vocab_size, + "device": self.device, + "req_states": self.req_states, + "logprobs_mode": self.model_config.logprobs_mode, + "num_speculative_tokens": self.decode_query_len, + "use_fp64_gumbel": self.model_config.use_fp64_gumbel, + "enable_trace_replay": self.model_config.enable_trace_replay, + "reasoning_config": self.vllm_config.reasoning_config, + "return_sampling_mask": self.model_config.return_sampling_mask, + } + if self.vllm_config.watermark_config is None: + self.sampler = Sampler(**sampler_kwargs) + else: + wm_config = self.vllm_config.watermark_config + watermarker = create_watermarker(wm_config) + if self.speculative_config is not None: + watermarker = create_speculative_target_watermarker(watermarker) + self.sampler = GPUWatermarkSampler( + watermarker, + deduplicate_contexts=wm_config.deduplicate_contexts, + deduplicate_contexts_max_history=( + wm_config.deduplicate_contexts_max_history + ), + **sampler_kwargs, + ) + custom = self.model_state.custom_sampler(self.sampler) + + if custom: + self.vllm_config._check_watermarking_unsupported( + custom_sampler=True + ) + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + watermark_key=speculative_target_watermark_key( + self.vllm_config.watermark_config + ), + ) self.prompt_logprobs_worker = PromptLogprobsWorker( self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode, @@ -738,7 +741,7 @@ def initialize_kv_cache( self.speculator.init_cudagraph_manager(cudagraph_mode) # Capture warmup providers that depend on allocated KV-cache strides. - with self.jit_warmup_registry.activate(): + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] kv_caches_dict = init_kv_cache( self.compilation_config.static_forward_context, self.kv_cache_config, @@ -924,6 +927,12 @@ def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: # during actual execution. assert self.sampler is not None self.sampler(logits, dummy_input_batch) + if self.rejection_sampler is not None: + # Warm head-dtype logits; scheduler warmup covers processed FP32 logits. + assert self.speculator is not None + self.rejection_sampler( + logits, dummy_input_batch, self.speculator.draft_logits + ) @torch.inference_mode() def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index fe4dee6a6b10..08cae8212836 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -4,7 +4,10 @@ import torch from vllm.sampling_params import SamplingParams -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p +from vllm.v1.sample.ops.topk_topp_sampler import ( + apply_top_k_top_p, + register_top_k_top_p_warmups, +) from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.sample.gumbel import apply_temperature from vllm.v1.worker.gpu.sample.min_p import apply_min_p @@ -18,6 +21,7 @@ class SamplingStates: def __init__(self, max_num_reqs: int, vocab_size: int): self.max_num_reqs = max_num_reqs self.vocab_size = vocab_size + register_top_k_top_p_warmups() self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a5710829e0ff..10f4346b5420 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +from types import SimpleNamespace from typing import Any +import numpy as np import torch import torch.nn as nn @@ -10,6 +12,15 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import ( + WarmupChoices, + WarmupIntRange, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -108,6 +119,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 + prepare_dflash_inputs.register_warmup(speculator=self) @property def attn_vllm_config(self) -> VllmConfig: @@ -682,6 +694,64 @@ def _prepare_dflash_inputs_kernel( tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) +def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]: + int32 = TritonWarmupTensor(torch.int32) + int64 = TritonWarmupTensor(torch.int64) + float32 = TritonWarmupTensor(torch.float32) + input_buffers = SimpleNamespace( + input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32 + ) + gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids) + max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2) + block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0)) + block_table = TritonWarmupTensor( + torch.int32, + shape=(1, block_table_stride), + strides=(block_table_stride, 1), + ) + input_batch = SimpleNamespace( + num_reqs=1, + num_scheduled_tokens=np.array([max_target_query_len]), + positions=int64, + query_start_loc=int32, + idx_mapping=int64, + ) + return dict( + input_buffers=input_buffers, + query_slot_mapping=int64, + context_positions=int64, + context_slot_mapping=int64, + sample_indices=int64, + sample_pos=int64, + sample_idx_mapping=int32, + temperature=float32, + seeds=int64, + input_batch=input_batch, + num_sampled=int32, + num_rejected=int32, + last_sampled=int64, + next_prefill_tokens=int32, + input_temperature=float32, + input_seeds=int64, + block_table=block_table, + block_size=speculator.block_tables.kernel_block_sizes[gid], + cp_rank=speculator.block_tables.cp_rank, + cp_size=speculator.block_tables.cp_size, + cp_interleave=speculator.block_tables.cp_interleave, + parallel_drafting_token_id=speculator.parallel_drafting_token_id, + num_query_per_req=speculator.num_query_per_req, + num_speculative_steps=speculator.num_speculative_steps, + max_num_reqs=speculator.max_num_reqs, + max_num_tokens=speculator.max_num_tokens, + max_model_len=speculator.max_model_len, + sample_from_anchor=speculator.sample_from_anchor, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_prepare_dflash_inputs_kernel, + warmup_inputs=_prepare_dflash_warmup_inputs, +) def prepare_dflash_inputs( input_buffers: InputBuffers, query_slot_mapping: torch.Tensor, @@ -718,50 +788,33 @@ def prepare_dflash_inputs( max_num_tokens: int, max_model_len: int, sample_from_anchor: bool = False, -) -> None: +) -> DispatchSpec: num_reqs = input_batch.num_reqs assert num_reqs > 0 # Cover the longest possible per-request span (ctx + query). Use the max # per-request query length, not the total token count across the batch. max_target_query_len = int(input_batch.num_scheduled_tokens.max()) max_tokens_per_req = max_target_query_len + num_query_per_req - BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) - num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) - _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( - input_buffers.input_ids, - input_buffers.positions, - input_buffers.query_start_loc, - input_buffers.seq_lens, - query_slot_mapping, - context_positions, - context_slot_mapping, - sample_indices, - sample_pos, - sample_idx_mapping, - temperature, - seeds, - input_batch.positions, - input_batch.query_start_loc, - input_batch.idx_mapping, - last_sampled, - next_prefill_tokens, - num_sampled, - num_rejected, - input_temperature, - input_seeds, - block_table, - block_table.stride(0), - parallel_drafting_token_id, - block_size, - num_query_per_req, - num_speculative_steps, - max_num_reqs, - max_num_tokens, - max_model_len, - cp_rank, - SAMPLE_FROM_ANCHOR=sample_from_anchor, + block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) + num_blocks = triton.cdiv(max_tokens_per_req, block) + return (num_reqs, num_blocks), dict( + out_input_ids_ptr=input_buffers.input_ids, + out_query_positions_ptr=input_buffers.positions, + out_query_start_loc_ptr=input_buffers.query_start_loc, + out_seq_lens_ptr=input_buffers.seq_lens, + out_query_slot_mapping_ptr=query_slot_mapping, + out_context_positions_ptr=context_positions, + out_context_slot_mapping_ptr=context_slot_mapping, + out_sample_indices_ptr=sample_indices, + out_sample_pos_ptr=sample_pos, + out_sample_idx_mapping_ptr=sample_idx_mapping, + out_temperature_ptr=temperature, + out_seeds_ptr=seeds, + target_positions_ptr=input_batch.positions, + target_query_start_loc_ptr=input_batch.query_start_loc, + idx_mapping_ptr=input_batch.idx_mapping, + temperature_ptr=input_temperature, + seeds_ptr=input_seeds, PAD_SLOT_ID=PAD_SLOT_ID, - CP_SIZE=cp_size, - CP_INTERLEAVE=cp_interleave, - BLOCK_SIZE=BLOCK_SIZE, + BLOCK_SIZE=block, ) diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c6ae315cf46..1c37ccb6c2db 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -9,6 +9,11 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -69,6 +74,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.cudagraph_manager: SpeculatorCudaGraphManager | None = None + _shift_input_ids.register_warmup() + _shift_input_embeds.register_warmup(speculator=self) + def load_draft_model( self, target_model: nn.Module, @@ -1008,6 +1016,32 @@ def _shift_input_ids_kernel( tl.store(input_ids_ptr + last_token_index, draft_token) +def _shift_input_ids_warmup_inputs() -> dict[str, Any]: + return dict( + input_ids=TritonWarmupTensor(torch.int32), + idx_mapping=TritonWarmupTensor(torch.int32), + query_start_loc=TritonWarmupTensor(torch.int32), + last_token_indices=TritonWarmupTensor(torch.int64), + draft_tokens=TritonWarmupTensor(torch.int64), + num_reqs=1, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_shift_input_ids_kernel, + warmup_inputs=_shift_input_ids_warmup_inputs, +) +def _shift_input_ids( + input_ids: torch.Tensor, + idx_mapping: torch.Tensor, + query_start_loc: torch.Tensor, + last_token_indices: torch.Tensor, + draft_tokens: torch.Tensor, + num_reqs: int, +) -> DispatchSpec: + return (num_reqs,), dict(BLOCK_SIZE=1024) + + @triton.jit def _shift_input_embeds_kernel( input_embeds_ptr, @@ -1065,6 +1099,44 @@ def _shift_input_embeds_kernel( ) +def _shift_input_embeds_warmup_inputs( + *, speculator: MultiModuleMTPSpeculator +) -> dict[str, Any]: + hidden = TritonWarmupTensor( + speculator.dtype, + shape=(1, speculator.hidden_size), + strides=(speculator.hidden_size, 1), + ) + return dict( + input_embeds=hidden, + draft_embeds=hidden, + idx_mapping=TritonWarmupTensor(torch.int32), + query_start_loc=TritonWarmupTensor(torch.int32), + last_token_indices=TritonWarmupTensor(torch.int64), + num_reqs=1, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_shift_input_embeds_kernel, + warmup_inputs=_shift_input_embeds_warmup_inputs, +) +def _shift_input_embeds( + input_embeds: torch.Tensor, + draft_embeds: torch.Tensor, + idx_mapping: torch.Tensor, + query_start_loc: torch.Tensor, + last_token_indices: torch.Tensor, + num_reqs: int, +) -> DispatchSpec: + hidden_size = input_embeds.shape[-1] + return (num_reqs, triton.cdiv(hidden_size, 256)), dict( + hidden_size=hidden_size, + BLOCK_SIZE_Q=16, + BLOCK_SIZE_H=256, + ) + + def update_draft_inputs( draft_tokens: torch.Tensor, draft_embeds: torch.Tensor | None, @@ -1074,29 +1146,21 @@ def update_draft_inputs( idx_mapping: torch.Tensor, num_reqs: int, ) -> None: - _shift_input_ids_kernel[(num_reqs,)]( + _shift_input_ids( input_buffers.input_ids, idx_mapping, input_buffers.query_start_loc, last_token_indices, draft_tokens, - BLOCK_SIZE=1024, + num_reqs=num_reqs, ) if input_embeds is not None: assert draft_embeds is not None - hidden_size = input_embeds.shape[-1] - hidden_block_size = 256 - _shift_input_embeds_kernel[ - (num_reqs, triton.cdiv(hidden_size, hidden_block_size)) - ]( + _shift_input_embeds( input_embeds, - input_embeds.stride(0), draft_embeds, - draft_embeds.stride(0), idx_mapping, input_buffers.query_start_loc, last_token_indices, - hidden_size, - BLOCK_SIZE_Q=16, - BLOCK_SIZE_H=hidden_block_size, + num_reqs=num_reqs, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d6a27b1f8ba9..97a59bd273c3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -498,6 +498,7 @@ class ExecuteModelState(NamedTuple): class GPUModelRunner( LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin ): + @JitWarmupRegistry.capture def __init__( self, vllm_config: VllmConfig, @@ -514,7 +515,6 @@ def __init__( self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config - self.jit_warmup_registry = JitWarmupRegistry(vllm_config) model_config = self.model_config cache_config = self.cache_config @@ -750,36 +750,35 @@ def __init__( self.parallel_config.cp_kv_cache_interleave_size ) # Capture warmup providers registered by the initial placeholder InputBatch - with self.jit_warmup_registry.activate(): - self.input_batch = InputBatch( - max_num_reqs=self.max_num_reqs, - # We need to use the encoder length for encoder-decoder - # because of KV cache for cross-attention. - max_model_len=max(self.max_model_len, self.max_encoder_len), - max_num_batched_tokens=self.max_num_tokens, - device=self.device, - vocab_size=self.model_config.get_vocab_size(), - block_sizes=[placeholder_block_size], - kernel_block_sizes=[placeholder_block_size], - max_num_blocks_per_req=[placeholder_max_num_blocks], - num_spec_tokens=self.num_spec_tokens, - logitsprocs=build_logitsprocs( - self.vllm_config, - self.device, - PIN_MEMORY, - self.is_pooling_model, - custom_logitsprocs, - ), - # We currently don't know whether a particular custom logits processor - # uses output token ids so we set this conservatively. Thinking-budget - # tracking is requested dynamically when a budgeted request is in the - # batch. - logitsprocs_need_output_token_ids=bool(custom_logitsprocs), - is_pooling_model=self.is_pooling_model, - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, - reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, - ) + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + # We need to use the encoder length for encoder-decoder + # because of KV cache for cross-attention. + max_model_len=max(self.max_model_len, self.max_encoder_len), + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=[placeholder_block_size], + kernel_block_sizes=[placeholder_block_size], + max_num_blocks_per_req=[placeholder_max_num_blocks], + num_spec_tokens=self.num_spec_tokens, + logitsprocs=build_logitsprocs( + self.vllm_config, + self.device, + PIN_MEMORY, + self.is_pooling_model, + custom_logitsprocs, + ), + # We currently don't know whether a particular custom logits processor + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the + # batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + ) # Separate cuda stream for overlapping transfer of sampled token ids from # GPU to CPU when async scheduling is enabled. @@ -5332,7 +5331,7 @@ def load_model(self, load_dummy_weights: bool = False) -> None: self.load_config.load_format = "dummy" model_loader = get_model_loader(self.load_config) # Capture warmup providers selected while constructing the model. - with self.jit_warmup_registry.activate(): + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.model_config ) @@ -7304,7 +7303,7 @@ def may_reinitialize_input_batch( self.parallel_config.cp_kv_cache_interleave_size ) # Capture warmup providers registered after final KV-cache geometry is known - with self.jit_warmup_registry.activate(): + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] self.input_batch = InputBatch( max_num_reqs=self.max_num_reqs, max_model_len=max_model_len, @@ -7454,7 +7453,7 @@ def initialize_kv_cache( # Reinitialize need to after initialize_attn_backend self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) # Capture warmup providers that depend on allocated KV-cache strides. - with self.jit_warmup_registry.activate(): + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] kv_caches = self.initialize_kv_cache_tensors( kv_cache_config, kernel_block_sizes, From e52be1a62d3879b1202f4f355d3c3472b560c6f2 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:25:23 +0200 Subject: [PATCH 113/774] [9/N][warmup][DSv4] Migrate MHC TileLang kernels (#50178) Signed-off-by: LopezCastroRoberto Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Codex --- tests/kernels/test_mhc_jit_warmup.py | 296 +++++++ tests/kernels/test_mhc_kernels.py | 13 +- vllm/model_executor/kernels/mhc/tilelang.py | 381 +++------ .../kernels/mhc/tilelang_kernels.py | 755 ++++++++++++++++++ .../warmup/deepseek_v4_mhc_warmup.py | 211 ----- .../warmup/jit_warmup_tilelang_helper.py | 172 ++++ vllm/model_executor/warmup/kernel_warmup.py | 12 - vllm/models/deepseek_v4/nvidia/model.py | 59 ++ vllm/models/glm5next/nvidia/model.py | 44 + 9 files changed, 1454 insertions(+), 489 deletions(-) create mode 100644 tests/kernels/test_mhc_jit_warmup.py delete mode 100644 vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py create mode 100644 vllm/model_executor/warmup/jit_warmup_tilelang_helper.py diff --git a/tests/kernels/test_mhc_jit_warmup.py b/tests/kernels/test_mhc_jit_warmup.py new file mode 100644 index 000000000000..28e746c09295 --- /dev/null +++ b/tests/kernels/test_mhc_jit_warmup.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate JIT dispatch against pre-contract behavior.""" + +from typing import Any + +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_cuda_alike(): + pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True) + +from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + HcHeadFusedTileLangKernel, + HcPrenormGemmTileLangKernel, + MhcFusedTileLangKernel, + MhcPostTileLangKernel, + MhcPreBigFuseTileLangKernel, +) +from vllm.model_executor.warmup import jit_warmup_tilelang_helper + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + [ + ( + dict( + num_tokens=64, + hc_hidden_size=4096, + hidden_size=2048, + hc_mult=2, + n_out=128, + ), + (2048, 2, 128, 1024, 4, 1, False, 1), + ), + ( + dict( + num_tokens=1024, + hc_hidden_size=4096, + hidden_size=2048, + hc_mult=2, + n_out=128, + ), + (2048, 2, 128, 512, 12, 1, True, 2), + ), + ( + dict( + num_tokens=64, + hc_hidden_size=4096, + hidden_size=2048, + hc_mult=2, + n_out=128, + n_thr=256, + tile_n=8, + n_splits=4, + ), + (2048, 2, 128, 256, 8, 4, False, 1), + ), + ], +) +def test_hc_prenorm_gemm_dispatch_matches_legacy_runtime_config( + kwargs: dict[str, Any], + expected: tuple[int, int, int, int, int, int, bool, int], +) -> None: + kernel = HcPrenormGemmTileLangKernel() + + assert kernel.dispatch(**kwargs) == kernel.CompileKey(*expected) + + +@pytest.mark.parametrize( + ("is_broadcast", "use_norm_weight", "expected_use_norm", "expected_eps"), + [ + (False, False, False, 0.0), + (False, True, True, 1.0e-5), + (True, False, True, 2.0e-5), + ], +) +def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config( + is_broadcast: bool, + use_norm_weight: bool, + expected_use_norm: bool, + expected_eps: float, +) -> None: + kernel = MhcPreBigFuseTileLangKernel() + + assert kernel.dispatch( + hidden_size=4096, + hc_mult=4, + n_splits=2, + is_broadcast=is_broadcast, + use_norm_weight=use_norm_weight, + rms_eps=1.0e-6, + hc_pre_eps=2.0e-6, + hc_sinkhorn_eps=3.0e-6, + hc_post_mult_value=0.5, + sinkhorn_repeat=3, + norm_eps=1.0e-5, + broadcast_norm_eps=2.0e-5, + ) == kernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=2, + use_norm_weight=expected_use_norm, + is_broadcast=is_broadcast, + rms_eps=1.0e-6, + hc_pre_eps=2.0e-6, + hc_sinkhorn_eps=3.0e-6, + hc_post_mult_value=0.5, + sinkhorn_repeat=3, + norm_eps=expected_eps, + ) + + +@pytest.mark.parametrize( + ("num_tokens", "hidden_size", "expected_n_splits", "expected_tile_n"), + [(4, 4096, 8, 2), (4, 8192, 4, 2), (8, 4096, 4, 3)], +) +def test_mhc_fused_dispatch_matches_legacy_runtime_config( + num_tokens: int, + hidden_size: int, + expected_n_splits: int, + expected_tile_n: int, +) -> None: + kernel = MhcFusedTileLangKernel() + + assert kernel.dispatch( + num_tokens=num_tokens, + hidden_size=hidden_size, + hc_mult=4, + ) == kernel.CompileKey( + hidden_size=hidden_size, + hc_mult=4, + n_splits=expected_n_splits, + tile_n=expected_tile_n, + ) + + +@pytest.mark.parametrize( + ("kernel", "compile_key"), + [ + ( + HcPrenormGemmTileLangKernel(), + HcPrenormGemmTileLangKernel.CompileKey( + hidden_size=2048, + hc_mult=2, + n_out=128, + n_thr=1024, + tile_n=4, + n_splits=1, + use_block_m=False, + block_m=1, + ), + ), + ( + HcPrenormGemmTileLangKernel(), + HcPrenormGemmTileLangKernel.CompileKey( + hidden_size=2048, + hc_mult=2, + n_out=128, + n_thr=512, + tile_n=12, + n_splits=1, + use_block_m=False, + block_m=1, + ), + ), + ( + HcPrenormGemmTileLangKernel(), + HcPrenormGemmTileLangKernel.CompileKey( + hidden_size=2048, + hc_mult=2, + n_out=128, + n_thr=512, + tile_n=12, + n_splits=1, + use_block_m=True, + block_m=2, + ), + ), + ( + MhcPreBigFuseTileLangKernel(), + MhcPreBigFuseTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=2, + use_norm_weight=True, + is_broadcast=False, + rms_eps=1.0e-6, + hc_pre_eps=2.0e-6, + hc_sinkhorn_eps=3.0e-6, + hc_post_mult_value=0.5, + sinkhorn_repeat=3, + norm_eps=1.0e-5, + ), + ), + ( + MhcPreBigFuseTileLangKernel(), + MhcPreBigFuseTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=1, + use_norm_weight=False, + is_broadcast=False, + rms_eps=1.0e-6, + hc_pre_eps=2.0e-6, + hc_sinkhorn_eps=3.0e-6, + hc_post_mult_value=0.5, + sinkhorn_repeat=3, + norm_eps=0.0, + ), + ), + ( + MhcPreBigFuseTileLangKernel(), + MhcPreBigFuseTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=2, + use_norm_weight=True, + is_broadcast=True, + rms_eps=1.0e-6, + hc_pre_eps=2.0e-6, + hc_sinkhorn_eps=3.0e-6, + hc_post_mult_value=0.5, + sinkhorn_repeat=3, + norm_eps=1.0e-5, + ), + ), + ( + MhcPostTileLangKernel(), + MhcPostTileLangKernel.CompileKey(hidden_size=4096, hc_mult=4), + ), + ( + MhcFusedTileLangKernel(), + MhcFusedTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=4, + tile_n=3, + ), + ), + ( + MhcFusedTileLangKernel(), + MhcFusedTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + n_splits=8, + tile_n=2, + ), + ), + ( + MhcFusedTileLangKernel(), + MhcFusedTileLangKernel.CompileKey( + hidden_size=8192, + hc_mult=4, + n_splits=4, + tile_n=2, + ), + ), + ( + HcHeadFusedTileLangKernel(), + HcHeadFusedTileLangKernel.CompileKey( + hidden_size=4096, + hc_mult=4, + rms_eps=1.0e-6, + hc_eps=2.0e-6, + ), + ), + ], +) +def test_tilelang_warmup_inputs_reproduce_compile_key( + monkeypatch: pytest.MonkeyPatch, + kernel: Any, + compile_key: Any, +) -> None: + compiled: list[Any] = [] + single_kernel = isinstance( + kernel, + (MhcPostTileLangKernel, MhcFusedTileLangKernel, HcHeadFusedTileLangKernel), + ) + expected_kernel = kernel.kernel() if single_kernel else kernel.kernel(compile_key) + if single_kernel: + monkeypatch.setattr( + kernel, + "dispatch", + lambda **kwargs: pytest.fail("single-kernel launch must not dispatch"), + ) + monkeypatch.setattr( + jit_warmup_tilelang_helper, + "compile_tilelang", + lambda jit_impl, *args, **kwargs: compiled.append(jit_impl), + ) + + kernel.compile(compile_key) + + assert compiled == [expected_kernel] diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index ab5efb6578c6..7f0515d11ade 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -10,10 +10,12 @@ import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( - _tilelang_hc_prenorm_gemm, _torch_hc_prenorm_gemm, mhc_pre_delayed_tilelang, ) +from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, +) from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton from vllm.model_executor.layers.mhc import ( @@ -485,7 +487,14 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size): sqrsum = torch.empty_like(sqrsum_ref) _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref) - _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult) + _HC_PRENORM_GEMM_TILELANG_KERNEL( + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + ) torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4) torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index c4fe6dd16328..97ca23b5418c 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import direct_register_custom_op @@ -18,73 +19,56 @@ def _torch_hc_prenorm_gemm( sqrsum[0].copy_(x_float.square().sum(dim=-1)) -def _tilelang_hc_prenorm_gemm( +def _hc_prenorm_gemm_outputs( x: torch.Tensor, fn: torch.Tensor, - out: torch.Tensor, - sqrsum: torch.Tensor, + *, hidden_size: int, hc_mult: int, - tile_n: int = 12, - n_thr: int = 512, - n_splits: int = 1, -) -> None: + use_tilelang_fallback: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - hc_prenorm_gemm_block_m_tilelang, - hc_prenorm_gemm_tilelang, - ) - - assert out.shape[0] == n_splits - assert sqrsum.shape[0] == n_splits - assert x.shape[1] == hc_mult * hidden_size - assert x.shape[1] % n_splits == 0 - assert (x.shape[1] // n_splits) % n_thr == 0 - use_default_config = tile_n == 12 and n_thr == 512 - if n_splits == 1 and use_default_config and x.shape[0] >= 1024: - hc_prenorm_gemm_block_m_tilelang( - x, - fn, - out, - sqrsum, - hidden_size, - hc_mult, - fn.shape[0], - n_thr, - tile_n, - 2, + compute_num_split, + ) + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + tf32_hc_prenorm_gemm, + ) + + use_deep_gemm = is_deep_gemm_supported() or not use_tilelang_fallback + num_tokens = x.shape[0] + n_splits = ( + compute_num_split(64, x.shape[1], cdiv(num_tokens, 64)) if use_deep_gemm else 1 + ) + out = torch.empty( + n_splits, + num_tokens, + fn.shape[0], + dtype=torch.float32, + device=x.device, + ) + sqrsum = torch.empty( + n_splits, + num_tokens, + dtype=torch.float32, + device=x.device, + ) + if use_deep_gemm: + tf32_hc_prenorm_gemm(x, fn, out, sqrsum, n_splits) + else: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, ) - return - if ( - n_splits == 1 - and use_default_config - and x.shape[0] < 128 - and x.shape[1] % 1024 == 0 - ): - hc_prenorm_gemm_tilelang( + + _HC_PRENORM_GEMM_TILELANG_KERNEL( x, fn, out, sqrsum, hidden_size, hc_mult, - fn.shape[0], - 1024, - 4, - n_splits, ) - return - hc_prenorm_gemm_tilelang( - x, - fn, - out, - sqrsum, - hidden_size, - hc_mult, - fn.shape[0], - n_thr, - tile_n, - n_splits, - ) + return out, sqrsum def mhc_pre_delayed_tilelang( @@ -127,6 +111,7 @@ def mhc_pre_delayed_tilelang( (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult). """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, mhc_pre_big_fuse_tilelang, ) from vllm.model_executor.kernels.mhc.warmup import ( @@ -186,7 +171,14 @@ def mhc_pre_delayed_tilelang( if use_deep_gemm: tf32_hc_prenorm_gemm(x, fn, mixes, sqrsum, n_splits) else: - _tilelang_hc_prenorm_gemm(x, fn, mixes, sqrsum, input_size, 1) + _HC_PRENORM_GEMM_TILELANG_KERNEL( + x, + fn, + mixes, + sqrsum, + input_size, + 1, + ) if norm_weight is not None: assert norm_weight.shape == (hidden_size,) assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous() @@ -298,7 +290,8 @@ def mhc_pre_tilelang( hc_sinkhorn_eps: sinkhorn epsilon hc_post_mult_value: post-mix multiplier value sinkhorn_repeat: number of sinkhorn iterations - n_splits: split-k factor; + n_splits: retained for the shared MHC operator API; the active GEMM + backend selects its split factor internally. norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype torch.bfloat16. When provided, RMSNorm is fused into the layer_input write path of the big_fuse kernel. @@ -311,12 +304,8 @@ def mhc_pre_tilelang( layer_input: shape (..., hidden_size), dtype torch.bfloat16 """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, - mhc_pre_big_fuse_tilelang, - mhc_pre_big_fuse_with_norm_tilelang, + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, ) - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - from vllm.utils.math_utils import cdiv assert residual.dtype == torch.bfloat16 assert fn.dtype == torch.float32 @@ -346,17 +335,6 @@ def mhc_pre_tilelang( residual_flat = residual.view(-1, hc_mult, hidden_size) num_tokens = residual_flat.shape[0] - from vllm.utils.deep_gemm import is_deep_gemm_supported - - use_deep_gemm = is_deep_gemm_supported() - if use_deep_gemm: - # these numbers are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) - else: - n_splits = 1 - post_mix = torch.empty( num_tokens, hc_mult, dtype=torch.float32, device=residual.device ) @@ -367,78 +345,30 @@ def mhc_pre_tilelang( num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device ) - gemm_out_mul = torch.empty( - n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs( + residual_flat.view(num_tokens, hc_mult * hidden_size), + fn, + hidden_size=hidden_size, + hc_mult=hc_mult, ) - gemm_out_sqrsum = torch.empty( - n_splits, num_tokens, dtype=torch.float32, device=residual.device + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual_flat, + post_mix, + comb_mix, + layer_input, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight=norm_weight, + norm_eps=norm_eps, ) - residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size) - if use_deep_gemm: - tf32_hc_prenorm_gemm( - residual_2d, - fn, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, - ) - else: - _tilelang_hc_prenorm_gemm( - residual_2d, - fn, - gemm_out_mul, - gemm_out_sqrsum, - hidden_size, - hc_mult, - ) - - if norm_weight is None: - mhc_pre_big_fuse_tilelang( - gemm_out_mul, - gemm_out_sqrsum, - hc_scale, - hc_base, - residual_flat, - post_mix, - comb_mix, - layer_input, - # Pre-mix buffers are unused in V4 mode. - post_mix, - post_mix, - hidden_size, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - n_splits, - hc_mult, - ) - else: - mhc_pre_big_fuse_with_norm_tilelang( - gemm_out_mul, - gemm_out_sqrsum, - hc_scale, - hc_base, - residual_flat, - post_mix, - comb_mix, - layer_input, - norm_weight, - post_mix, - post_mix, - hidden_size, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - norm_eps, - n_splits, - hc_mult, - ) - return ( post_mix.view(*outer_shape, hc_mult, 1), comb_mix.view(*outer_shape, hc_mult, hc_mult), @@ -505,11 +435,10 @@ def mhc_pre_broadcast_tilelang( fn_broadcast: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """First-layer mHC pre for a residual broadcast from ``(T, H)``.""" + # n_splits is retained for the shared API; split selection is internal. from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, - mhc_pre_big_fuse_broadcast_with_norm_tilelang, + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, ) - from vllm.utils.math_utils import cdiv assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm" assert residual.dtype == torch.bfloat16 @@ -537,8 +466,6 @@ def mhc_pre_broadcast_tilelang( residual_flat = residual num_tokens = residual.shape[0] - n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64)) - residual_out = torch.empty( num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device ) @@ -551,42 +478,31 @@ def mhc_pre_broadcast_tilelang( layer_input = torch.empty( num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device ) - gemm_out_mul = torch.empty( - n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device - ) - gemm_out_sqrsum = torch.empty( - n_splits, num_tokens, dtype=torch.float32, device=residual.device - ) - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - - tf32_hc_prenorm_gemm( + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs( residual_flat, fn_broadcast, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, + hidden_size=hidden_size, + hc_mult=hc_mult, + use_tilelang_fallback=False, ) - mhc_pre_big_fuse_broadcast_with_norm_tilelang( + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual_flat, - residual_out, post_mix, comb_mix, layer_input, - norm_weight, - hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - norm_eps, - n_splits, - hc_mult, + residual_out=residual_out, + norm_weight=norm_weight, + norm_eps=norm_eps, ) return ( residual_out, @@ -603,11 +519,11 @@ def mhc_post_tilelang( comb_res_mix: torch.Tensor, ) -> torch.Tensor: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - mhc_post_tilelang as _mhc_post_kernel, + _MHC_POST_TILELANG_KERNEL, ) out = torch.empty_like(residual) - _mhc_post_kernel( + _MHC_POST_TILELANG_KERNEL( comb_res_mix, residual, post_layer_mix.squeeze(-1), @@ -644,6 +560,9 @@ def mhc_fused_post_pre_tilelang( RMSNorm'd activation (fused into the kernel); otherwise it is the raw pre-norm activation as before. + ``n_splits`` and ``tile_n`` are retained for the shared MHC operator API. + The TileLang path selects both values internally from the runtime shape. + Returns: residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size) post_mix_cur: shape (..., hc_mult, 1) @@ -652,13 +571,10 @@ def mhc_fused_post_pre_tilelang( """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( - compute_num_split, - mhc_fused_tilelang, - mhc_post_tilelang, - mhc_pre_big_fuse_tilelang, - mhc_pre_big_fuse_with_norm_tilelang, + _MHC_FUSED_TILELANG_KERNEL, + _MHC_POST_TILELANG_KERNEL, + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, ) - from vllm.utils.math_utils import cdiv assert residual.dtype == torch.bfloat16 assert x.dtype == torch.bfloat16 @@ -692,48 +608,14 @@ def mhc_fused_post_pre_tilelang( if not norm_weight.is_contiguous(): norm_weight = norm_weight.contiguous() - assert n_splits in (1, 2, 4, 8) - assert hidden_size % n_splits == 0 - residual_flat = residual.view(-1, hc_mult, hidden_size) num_tokens = residual_flat.shape[0] x_flat = x.view(num_tokens, hidden_size) post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult) comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) - from vllm.utils.deep_gemm import is_deep_gemm_supported - - use_deep_gemm = is_deep_gemm_supported() use_small_fma = num_tokens <= 16 - if use_small_fma: - # TODO(gnovack): investigate autotuning these heuristics - tile_n = 2 if num_tokens < 8 else 3 - n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 - else: - if use_deep_gemm: - # these number are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split( - block_k, hc_hidden_size, cdiv(num_tokens, block_m) - ) - else: - n_splits = 1 - - gemm_out_mul = torch.empty( - n_splits, - num_tokens, - hc_mult3, - dtype=torch.float32, - device=residual.device, - ) - gemm_out_sqrsum = torch.empty( - n_splits, - num_tokens, - dtype=torch.float32, - device=residual.device, - ) - residual_cur = torch.empty_like(residual_flat) + post_mix_cur = torch.empty( num_tokens, hc_mult, @@ -754,55 +636,17 @@ def mhc_fused_post_pre_tilelang( ) if use_small_fma: - mhc_fused_tilelang( + gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL( comb_res_mix_flat, residual_flat, post_layer_mix_flat, x_flat, fn.view(hc_mult3, hc_mult, hidden_size), - gemm_out_mul, - gemm_out_sqrsum, - residual_cur, hc_mult, hidden_size, hc_mult3, - tile_n=tile_n, - n_splits=n_splits, ) - else: - mhc_post_tilelang( - comb_res_mix_flat, - residual_flat, - post_layer_mix_flat, - x_flat, - residual_cur, - residual.shape[-2], - residual.shape[-1], - ) - - residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size) - if use_deep_gemm: - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - - tf32_hc_prenorm_gemm( - residual_cur_2d, - fn, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, - ) - else: - _tilelang_hc_prenorm_gemm( - residual_cur_2d, - fn, - gemm_out_mul, - gemm_out_sqrsum, - hidden_size, - hc_mult, - ) - - if norm_weight is None: - mhc_pre_big_fuse_tilelang( + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL( gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -811,20 +655,32 @@ def mhc_fused_post_pre_tilelang( post_mix_cur, comb_mix_cur, layer_input_cur, - # Pre-mix buffers are unused in V4 mode. - post_mix_cur, - post_mix_cur, - hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - n_splits, - hc_mult, + norm_weight=norm_weight, + norm_eps=norm_eps, ) else: - mhc_pre_big_fuse_with_norm_tilelang( + residual_cur = torch.empty_like(residual_flat) + _MHC_POST_TILELANG_KERNEL( + comb_res_mix_flat, + residual_flat, + post_layer_mix_flat, + x_flat, + residual_cur, + residual.shape[-2], + residual.shape[-1], + ) + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs( + residual_cur.view(num_tokens, hc_mult * hidden_size), + fn, + hidden_size=hidden_size, + hc_mult=hc_mult, + ) + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL( gemm_out_mul, gemm_out_sqrsum, hc_scale, @@ -833,18 +689,13 @@ def mhc_fused_post_pre_tilelang( post_mix_cur, comb_mix_cur, layer_input_cur, - norm_weight, - post_mix_cur, - post_mix_cur, - hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - norm_eps, - n_splits, - hc_mult, + norm_weight=norm_weight, + norm_eps=norm_eps, ) return ( @@ -926,9 +777,11 @@ def hc_head_fused_kernel_tilelang( ) if num_tokens == 0: return out - from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_HEAD_FUSED_TILELANG_KERNEL, + ) - hc_head_fuse_tilelang( + _HC_HEAD_FUSED_TILELANG_KERNEL( hs_flat, fn, hc_scale, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 1c23d152b916..b4d3190c0eff 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -4,10 +4,19 @@ from __future__ import annotations import math +from dataclasses import dataclass from functools import cache +from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs +from vllm.model_executor.warmup.jit_warmup_tilelang_helper import ( + TileLangLaunchSpec, + VllmTileLangJitKernel, + kernel_launcher, + make_tilelang_warmup_tensor, +) from vllm.platforms import current_platform from vllm.tilelang_utils import T, tilelang, tilelang_jit from vllm.utils.math_utils import cdiv @@ -28,6 +37,10 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k +def _mhc_fused_n_splits(num_tokens: int, hidden_size: int) -> int: + return 8 if num_tokens < 8 and hidden_size <= 4096 else 4 + + @tilelang_jit def mhc_pre_big_fuse_tilelang( gemm_out_mul, @@ -987,3 +1000,745 @@ def hc_head_fuse_tilelang( if ENABLE_PDL: T.pdl_trigger() + + +class HcPrenormGemmTileLangKernel( + VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + n_out: int + n_thr: int + tile_n: int + n_splits: int + use_block_m: bool + block_m: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + if compile_key.use_block_m: + return hc_prenorm_gemm_block_m_tilelang + return hc_prenorm_gemm_tilelang + + def dispatch( # type: ignore[override] + self, + *, + num_tokens: int, + hc_hidden_size: int, + n_thr: int = 512, + tile_n: int = 12, + n_splits: int = 1, + **compile_key_fields: int, + ) -> CompileKey: + use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512 + use_block_m = use_default_config and num_tokens >= 1024 + use_small_tile = ( + use_default_config + and not use_block_m + and num_tokens < 128 + and hc_hidden_size % 1024 == 0 + ) + effective_n_thr = 1024 if use_small_tile else n_thr + effective_tile_n = 4 if use_small_tile else tile_n + return self.CompileKey( + **compile_key_fields, + n_thr=effective_n_thr, + tile_n=effective_tile_n, + n_splits=n_splits, + use_block_m=use_block_m, + block_m=2 if use_block_m else 1, + ) + + def get_warmup_keys( + self, + vllm_config: Any, + *, + hidden_size: int, + hc_mult: int, + n_out: int, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_tokens <= 0: + return [] + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange( + 1, + max_tokens + 1, + advance=lambda value: ( + 128 if value < 128 else 1024 if value < 1024 else max_tokens + 1 + ), + ), + hc_hidden_size=hidden_size * hc_mult, + hidden_size=hidden_size, + hc_mult=hc_mult, + n_out=n_out, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + num_tokens = 128 + if compile_key.use_block_m: + num_tokens = 1024 + elif compile_key.n_thr == 1024: + num_tokens = 1 + hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult + return dict( + x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size), + fn=make_tilelang_warmup_tensor( + torch.float32, compile_key.n_out, hc_hidden_size + ), + out=make_tilelang_warmup_tensor( + torch.float32, + compile_key.n_splits, + num_tokens, + compile_key.n_out, + ), + sqrsum=make_tilelang_warmup_tensor( + torch.float32, compile_key.n_splits, num_tokens + ), + hidden_size=compile_key.hidden_size, + hc_mult=compile_key.hc_mult, + tile_n=compile_key.tile_n, + n_thr=compile_key.n_thr, + n_splits=compile_key.n_splits, + ) + + @kernel_launcher + def __call__( + self, + x: Any, + fn: Any, + out: Any, + sqrsum: Any, + hidden_size: int, + hc_mult: int, + tile_n: int = 12, + n_thr: int = 512, + n_splits: int = 1, + ) -> TileLangLaunchSpec: + assert out.shape[0] == n_splits + assert sqrsum.shape[0] == n_splits + assert x.shape[1] == hc_mult * hidden_size + assert x.shape[1] % n_splits == 0 + assert (x.shape[1] // n_splits) % n_thr == 0 + compile_key = self.dispatch( + num_tokens=x.shape[0], + hc_hidden_size=x.shape[1], + hidden_size=hidden_size, + hc_mult=hc_mult, + n_out=fn.shape[0], + n_thr=n_thr, + tile_n=tile_n, + n_splits=n_splits, + ) + kernel_arg = ( + compile_key.block_m if compile_key.use_block_m else compile_key.n_splits + ) + return (compile_key,), ( + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + fn.shape[0], + compile_key.n_thr, + compile_key.tile_n, + kernel_arg, + ) + + +class MhcPreBigFuseTileLangKernel( + VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + n_splits: int + use_norm_weight: bool + is_broadcast: bool + rms_eps: float + hc_pre_eps: float + hc_sinkhorn_eps: float + hc_post_mult_value: float + sinkhorn_repeat: int + norm_eps: float + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + if compile_key.is_broadcast: + return mhc_pre_big_fuse_broadcast_with_norm_tilelang + if compile_key.use_norm_weight: + return mhc_pre_big_fuse_with_norm_tilelang + return mhc_pre_big_fuse_tilelang + + @staticmethod + def _kernel_args( + compile_key: CompileKey, + *, + gemm_out_mul: Any, + gemm_out_sqrsum: Any, + hc_scale: Any, + hc_base: Any, + residual: Any, + residual_out: Any | None, + post_mix: Any, + comb_mix: Any, + layer_input: Any, + norm_weight: Any | None, + ) -> tuple[Any, ...]: + common_args = ( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + ) + output_args = (post_mix, comb_mix, layer_input) + eps_args = ( + compile_key.rms_eps, + compile_key.hc_pre_eps, + compile_key.hc_sinkhorn_eps, + compile_key.hc_post_mult_value, + compile_key.sinkhorn_repeat, + ) + middle_args: tuple[Any, ...] + if compile_key.is_broadcast: + assert residual_out is not None + assert norm_weight is not None + middle_args = (residual_out, *output_args, norm_weight) + elif compile_key.use_norm_weight: + assert norm_weight is not None + middle_args = (*output_args, norm_weight, post_mix, post_mix) + else: + middle_args = (*output_args, post_mix, post_mix) + norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else () + return ( + *common_args, + *middle_args, + compile_key.hidden_size, + *eps_args, + *norm_eps_args, + compile_key.n_splits, + compile_key.hc_mult, + ) + + def dispatch( # type: ignore[override] + self, + *, + hidden_size: int, + hc_mult: int, + n_splits: int, + is_broadcast: bool, + use_norm_weight: bool, + use_pre_gemm_splits: bool = False, + pre_gemm_k: int = 0, + sinkhorn_repeat: int, + norm_eps: float, + broadcast_norm_eps: float = 0.0, + num_tokens: int = 1, + use_fused_tilelang: bool = False, + **compile_key_fields: float, + ) -> CompileKey: + pre_gemm_n_splits = ( + compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64) + if use_pre_gemm_splits + else n_splits + ) + fused_n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + actual_n_splits = fused_n_splits if use_fused_tilelang else pre_gemm_n_splits + actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps + actual_use_norm_weight = use_norm_weight or is_broadcast + return self.CompileKey( + **compile_key_fields, + hidden_size=hidden_size, + hc_mult=hc_mult, + n_splits=actual_n_splits, + use_norm_weight=actual_use_norm_weight, + is_broadcast=is_broadcast, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0, + ) + + def get_warmup_keys( + self, + vllm_config: Any, + *, + hidden_size: int, + hc_mult: int, + use_norm_weight: bool, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + norm_eps: float | tuple[float, ...], + broadcast_norm_eps: float | tuple[float, ...] | None = None, + include_pre_gemm_splits: bool = False, + include_broadcast_splits: bool = False, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_tokens <= 0: + return [] + + warmup_cases = zip_inputs( + dict( + is_broadcast=False, + use_pre_gemm_splits=False, + pre_gemm_k=0, + use_fused_tilelang=False, + ), + dict( + is_broadcast=False, + use_pre_gemm_splits=include_pre_gemm_splits, + pre_gemm_k=hc_mult * hidden_size, + use_fused_tilelang=False, + ), + dict( + is_broadcast=include_broadcast_splits, + use_pre_gemm_splits=include_broadcast_splits, + pre_gemm_k=hidden_size, + use_fused_tilelang=False, + ), + dict( + is_broadcast=False, + use_pre_gemm_splits=False, + pre_gemm_k=0, + use_fused_tilelang=True, + ), + ) + return self._trace_dispatch(self.dispatch)( + warmup_cases, + num_tokens=WarmupIntRange( + 1, + max_tokens + 1, + advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1, + ), + n_splits=1, + hidden_size=hidden_size, + hc_mult=hc_mult, + use_norm_weight=use_norm_weight, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=norm_eps, + broadcast_norm_eps=broadcast_norm_eps + if broadcast_norm_eps is not None + else norm_eps, + _when=lambda *, use_fused_tilelang, num_tokens: ( + not use_fused_tilelang or num_tokens <= 16 + ), + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + hc_mult3 = hc_mult * (2 + hc_mult) + num_tokens = 1 + + gemm_out_mul = make_tilelang_warmup_tensor( + torch.float32, compile_key.n_splits, num_tokens, hc_mult3 + ) + gemm_out_sqrsum = make_tilelang_warmup_tensor( + torch.float32, compile_key.n_splits, num_tokens + ) + hc_scale = make_tilelang_warmup_tensor(torch.float32, 3) + hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3) + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) + comb_mix = make_tilelang_warmup_tensor( + torch.float32, num_tokens, hc_mult * hc_mult + ) + layer_input = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hidden_size + ) + residual_shape = ( + (num_tokens, hidden_size) + if compile_key.is_broadcast + else (num_tokens, hc_mult, hidden_size) + ) + residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape) + residual_out = ( + make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hc_mult, hidden_size + ) + if compile_key.is_broadcast + else None + ) + norm_weight = ( + make_tilelang_warmup_tensor(torch.bfloat16, hidden_size) + if compile_key.is_broadcast or compile_key.use_norm_weight + else None + ) + return dict( + gemm_out_mul=gemm_out_mul, + gemm_out_sqrsum=gemm_out_sqrsum, + hc_scale=hc_scale, + hc_base=hc_base, + residual=residual, + post_mix=post_mix, + comb_mix=comb_mix, + layer_input=layer_input, + rms_eps=compile_key.rms_eps, + hc_pre_eps=compile_key.hc_pre_eps, + hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps, + hc_post_mult_value=compile_key.hc_post_mult_value, + sinkhorn_repeat=compile_key.sinkhorn_repeat, + residual_out=residual_out, + norm_weight=norm_weight, + norm_eps=compile_key.norm_eps, + ) + + @kernel_launcher + def __call__( + self, + gemm_out_mul: Any, + gemm_out_sqrsum: Any, + hc_scale: Any, + hc_base: Any, + residual: Any, + post_mix: Any, + comb_mix: Any, + layer_input: Any, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + *, + residual_out: Any | None = None, + norm_weight: Any | None = None, + norm_eps: float = 0.0, + ) -> TileLangLaunchSpec: + is_broadcast = residual_out is not None + use_norm_weight = norm_weight is not None + n_splits = gemm_out_mul.shape[0] + if is_broadcast: + assert residual_out is not None + hidden_size = residual.shape[-1] + hc_mult = residual_out.shape[-2] + else: + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + compile_key = self.dispatch( + hidden_size=hidden_size, + hc_mult=hc_mult, + n_splits=n_splits, + is_broadcast=is_broadcast, + use_norm_weight=use_norm_weight, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=norm_eps, + broadcast_norm_eps=norm_eps, + ) + return ( + (compile_key,), + self._kernel_args( + compile_key, + gemm_out_mul=gemm_out_mul, + gemm_out_sqrsum=gemm_out_sqrsum, + hc_scale=hc_scale, + hc_base=hc_base, + residual=residual, + residual_out=residual_out, + post_mix=post_mix, + comb_mix=comb_mix, + layer_input=layer_input, + norm_weight=norm_weight, + ), + ) + + +class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + + @staticmethod + def kernel() -> Any: + return mhc_post_tilelang + + def dispatch( # type: ignore[override] + self, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey(**compile_key_fields) + + def get_warmup_keys( + self, + *, + hidden_size: int, + hc_mult: int, + ) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + hidden_size=hidden_size, + hc_mult=hc_mult, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + num_tokens = 1 + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + comb_mix = make_tilelang_warmup_tensor( + torch.float32, num_tokens, hc_mult, hc_mult + ) + residual = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hc_mult, hidden_size + ) + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) + layer_input = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hidden_size + ) + out = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hc_mult, hidden_size + ) + return dict( + comb_mix=comb_mix, + residual=residual, + post_mix=post_mix, + layer_input=layer_input, + out=out, + hc_mult=hc_mult, + hidden_size=hidden_size, + ) + + @kernel_launcher + def __call__( + self, + comb_mix: Any, + residual: Any, + post_mix: Any, + layer_input: Any, + out: Any, + hc_mult: int, + hidden_size: int, + ) -> TileLangLaunchSpec: + return (), ( + comb_mix, + residual, + post_mix, + layer_input, + out, + hc_mult, + hidden_size, + ) + + +class MhcFusedTileLangKernel( + VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + n_splits: int + tile_n: int + + @staticmethod + def kernel() -> Any: + return mhc_fused_tilelang + + def dispatch( # type: ignore[override] + self, + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + ) -> CompileKey: + # TODO(gnovack): investigate autotuning these heuristics + tile_n = 2 if num_tokens < 8 else 3 + n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + return self.CompileKey( + hidden_size=hidden_size, + hc_mult=hc_mult, + n_splits=n_splits, + tile_n=tile_n, + ) + + def get_warmup_keys( + self, + vllm_config: Any, + *, + hidden_size: int, + hc_mult: int, + ) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + hidden_size=hidden_size, + hc_mult=hc_mult, + _when=lambda *, num_tokens: num_tokens <= 16, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + num_tokens = 1 if compile_key.tile_n == 2 else 8 + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + hc_mult3 = hc_mult * (2 + hc_mult) + comb_mix = make_tilelang_warmup_tensor( + torch.float32, num_tokens, hc_mult, hc_mult + ) + residual_in = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hc_mult, hidden_size + ) + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) + x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) + weight_t = make_tilelang_warmup_tensor( + torch.float32, hc_mult3, hc_mult, hidden_size + ) + return dict( + comb_mix=comb_mix, + residual_in=residual_in, + post_mix=post_mix, + x_in=x_in, + weight_t=weight_t, + hc_mult=hc_mult, + hidden_size=hidden_size, + hc_mult3=hc_mult3, + ) + + @kernel_launcher + def __call__( + self, + comb_mix: Any, + residual_in: Any, + post_mix: Any, + x_in: Any, + weight_t: Any, + hc_mult: int, + hidden_size: int, + hc_mult3: int, + ) -> TileLangLaunchSpec: + num_tokens = residual_in.shape[0] + tile_n = 2 if num_tokens < 8 else 3 + n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + yp_out = residual_in.new_empty( + (n_splits, num_tokens, hc_mult3), dtype=torch.float32 + ) + rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32) + residual_out = residual_in.new_empty(residual_in.shape) + return ( + (), + ( + comb_mix, + residual_in, + post_mix, + x_in, + weight_t, + yp_out, + rp_out, + residual_out, + hc_mult, + hidden_size, + hc_mult3, + ), + dict(tile_n=tile_n, split_k=n_splits), + (yp_out, rp_out, residual_out), + ) + + +class HcHeadFusedTileLangKernel( + VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + hidden_size: int + hc_mult: int + rms_eps: float + hc_eps: float + + @staticmethod + def kernel() -> Any: + return hc_head_fuse_tilelang + + def dispatch( # type: ignore[override] + self, + *, + hidden_size: int, + hc_mult: int, + rms_eps: float, + hc_eps: float, + ) -> CompileKey: + return self.CompileKey( + hidden_size=hidden_size, + hc_mult=hc_mult, + rms_eps=rms_eps, + hc_eps=hc_eps, + ) + + def get_warmup_keys( + self, + *, + hidden_size: int, + hc_mult: int, + rms_eps: float, + hc_eps: float, + ) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + hidden_size=hidden_size, + hc_mult=hc_mult, + rms_eps=rms_eps, + hc_eps=hc_eps, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: + num_tokens = 1 + hidden_size = compile_key.hidden_size + hc_mult = compile_key.hc_mult + residual = make_tilelang_warmup_tensor( + torch.bfloat16, num_tokens, hc_mult, hidden_size + ) + fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size) + hc_scale = make_tilelang_warmup_tensor(torch.float32, 1) + hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult) + out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) + return dict( + residual=residual, + fn=fn, + hc_scale=hc_scale, + hc_base=hc_base, + out=out, + hidden_size=hidden_size, + rms_eps=compile_key.rms_eps, + hc_eps=compile_key.hc_eps, + hc_mult=hc_mult, + ) + + @kernel_launcher + def __call__( + self, + residual: Any, + fn: Any, + hc_scale: Any, + hc_base: Any, + out: Any, + hidden_size: int, + rms_eps: float, + hc_eps: float, + hc_mult: int, + ) -> TileLangLaunchSpec: + return (), ( + residual, + fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_eps, + hc_eps, + hc_mult, + ) + + +_HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel() +_MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel() +_MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel() +_MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel() +_HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel() diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py deleted file mode 100644 index fc7731509d6f..000000000000 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests. - -Ported from lucifer1004/vllm-jasl with the two env-var knobs removed -(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`). -Gating is intrinsic: non-DSv4 models and layers without hc_* attributes -return early, so the warmup is a no-op except where it's needed. -""" - -import time -from collections.abc import Iterable - -import torch - -from vllm.logger import init_logger -from vllm.tracing import instrument - -logger = init_logger(__name__) - -_AUTO_WARMUP_MAX_TOKENS = 16_384 -_DEFAULT_TOKEN_SIZE_CANDIDATES = ( - 1, - 2, - 4, - 8, - 16, - 32, - 64, - 128, - 256, - 512, - 1024, - 2048, - 4096, - 8192, - 16_384, -) - - -def _normalize_token_sizes( - token_sizes: Iterable[int], - *, - max_tokens: int, -) -> list[int]: - return sorted({size for size in token_sizes if 1 <= size <= max_tokens}) - - -def _select_mhc_warmup_token_sizes( - *, - max_tokens: int, - cudagraph_capture_sizes: list[int], -) -> list[int]: - if max_tokens <= 0: - return [] - - max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) - candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES) - candidates.extend(cudagraph_capture_sizes) - candidates.append(max_auto_tokens) - return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens) - - -def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: - for module in model.modules(): - if module.__class__.__name__ != "DeepseekV4DecoderLayer": - continue - if all( - hasattr(module, attr) - for attr in ( - "hc_pre", - "hc_post", - "hc_attn_fn", - "hc_attn_scale", - "hc_attn_base", - "hc_ffn_fn", - "hc_ffn_scale", - "hc_ffn_base", - ) - ): - return module - return None - - -def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None: - for module in model.modules(): - if module.__class__.__name__ != "DeepseekV4Model": - continue - if all( - hasattr(module, attr) - for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base") - ): - return module - return None - - -def _warmup_layer_mhc( - layer: torch.nn.Module, - token_sizes: list[int], -) -> None: - max_tokens = max(token_sizes) - hidden_size = int(layer.hidden_size) - hc_mult = int(layer.hc_mult) - device = layer.hc_attn_fn.device - residual = torch.zeros( - max_tokens, - hc_mult, - hidden_size, - dtype=torch.bfloat16, - device=device, - ) - - for size in token_sizes: - residual_slice = residual[:size] - for fn, scale, base in ( - (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base), - (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base), - ): - layer_input, post_mix, comb_mix = layer.hc_pre( - residual_slice, - fn, - scale, - base, - ) - layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) - - -def _warmup_hc_head( - model: torch.nn.Module, - token_sizes: list[int], -) -> None: - # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored - # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp - # instance attached to the model as ``hc_head_op``. We call through - # that instance so the warmup exercises the same dispatched - # implementation as the inference path. - hc_head_op = getattr(model, "hc_head_op", None) - if hc_head_op is None: - return - - max_tokens = max(token_sizes) - hidden_size = int(model.config.hidden_size) - hc_mult = int(model.hc_mult) - device = model.hc_head_fn.device - hidden_states = torch.zeros( - max_tokens, - hc_mult, - hidden_size, - dtype=torch.bfloat16, - device=device, - ) - - for size in token_sizes: - hc_head_op( - hidden_states[:size], - model.hc_head_fn, - model.hc_head_scale, - model.hc_head_base, - model.rms_norm_eps, - model.hc_eps, - ) - - -@instrument(span_name="DeepSeek V4 mHC warmup") -def deepseek_v4_mhc_warmup( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: list[int] | None = None, -) -> None: - # Cheap model-type gate before walking ``model.modules()``. The class - # walk below is O(num_layers) and shows up in startup time on very - # large checkpoints; bail out for any model that is not DeepSeek V4. - config = getattr(model, "config", None) - model_type = getattr(config, "model_type", None) if config is not None else None - if model_type is not None and model_type not in ( - "deepseek_v4", - "deepseek_v41", - ): - return - - layer = _find_first_mhc_layer(model) - if layer is None: - return - - device = layer.hc_attn_fn.device - if device.type != "cuda": - return - - deepseek_model = _find_deepseek_v4_model(model) - token_sizes = _select_mhc_warmup_token_sizes( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes or [], - ) - if not token_sizes: - return - - started = time.perf_counter() - logger.info( - "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s", - token_sizes, - ) - with torch.inference_mode(): - _warmup_layer_mhc(layer, token_sizes) - if deepseek_model is not None: - _warmup_hc_head(deepseek_model, token_sizes) - torch.accelerator.synchronize() - logger.info( - "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.", - time.perf_counter() - started, - ) diff --git a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py new file mode 100644 index 000000000000..056ca5180508 --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compile-only helpers for TileLang JIT warmup.""" + +from __future__ import annotations + +import logging +from abc import abstractmethod +from collections.abc import Callable, Iterator, Mapping, MutableMapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import wraps +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar, cast + +import torch + +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel + +CompileKeyT = TypeVar("CompileKeyT") +# (kernel_args, launch_args), optionally followed by keywords and a result. +TileLangLaunchSpec: TypeAlias = ( + tuple[tuple[Any, ...], tuple[Any, ...]] + | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None] + | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None, Any] +) + + +@contextmanager +def _quiet_tilelang_warmup_logs() -> Iterator[None]: + loggers = [ + logging.getLogger("tilelang.cache.kernel_cache"), + logging.getLogger("tilelang.jit.kernel"), + ] + previous_levels = [logger.level for logger in loggers] + try: + for logger in loggers: + logger.setLevel(logging.ERROR) + yield + finally: + for logger, level in zip(loggers, previous_levels): + logger.setLevel(level) + + +@dataclass(frozen=True) +class TileLangWarmupTensor: + """Minimal tensor-like object accepted by TileLang compile(). + + TileLang builds its cache key and TIR from tensor dtype, shape and stride. + This object deliberately has no storage, so compile-only warmup does not + allocate GPU memory and does not launch the kernel. + """ + + dtype: torch.dtype + shape: tuple[int, ...] = (1,) + strides: tuple[int, ...] | None = field(default=None) + + def stride(self) -> tuple[int, ...]: + if self.strides is not None: + return self.strides + + strides: list[int] = [] + stride = 1 + for size in reversed(self.shape): + strides.append(stride) + stride *= size + return tuple(reversed(strides)) + + def new_empty( + self, + shape: tuple[int, ...], + *, + dtype: torch.dtype | None = None, + device: Any = None, + ) -> TileLangWarmupTensor: + return TileLangWarmupTensor(dtype or self.dtype, shape) + + +def make_tilelang_warmup_tensor( + dtype: torch.dtype, + *shape: int, + strides: tuple[int, ...] | None = None, +) -> TileLangWarmupTensor: + return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape), strides=strides) + + +def _tilelang_call_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + call_kwargs = dict(kwargs) + tune_params = call_kwargs.pop("__tune_params", {}) + if isinstance(tune_params, dict): + call_kwargs.update(tune_params) + return call_kwargs + + +def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None: + """Compile one TileLang specialization and populate its call cache. + + TileLang's ``compile()`` materializes the kernel without launching it. + We also store the compiled kernel in ``_kernel_cache`` using the same + parsed key as ``__call__`` so runtime does not report a cache miss for an + already materialized specialization. + """ + + with _quiet_tilelang_warmup_logs(): + compiled = jit_impl.compile(*args, **kwargs) + func = getattr(jit_impl, "func", None) + parse_args = getattr(func, "parse_args", None) + cache = getattr(jit_impl, "_kernel_cache", None) + if not callable(parse_args) or not isinstance(cache, MutableMapping): + return + + key, _ = parse_args(*args, **_tilelang_call_kwargs(kwargs)) + cache[key] = compiled + + +class VllmTileLangJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): + """TileLang owner whose runtime launch specification is reused for warmup.""" + + kernel: ClassVar[Any] + _warming_key: CompileKeyT | None = None + + @abstractmethod + def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]: + """Return runtime-shaped inputs that reproduce one compile key.""" + raise NotImplementedError + + def compile(self, compile_key: CompileKeyT) -> None: + self._warming_key = compile_key + try: + cast(Callable[..., Any], self)(**self.warmup_inputs(compile_key)) + finally: + self._warming_key = None + + def launch( + self, + jit_impl: Any, + args: tuple[Any, ...], + kwargs: Mapping[str, Any] | None = None, + ) -> Any: + call_kwargs = dict(kwargs or {}) + if self._warming_key is not None: + return compile_tilelang(jit_impl, *args, **call_kwargs) + return jit_impl(*args, **call_kwargs) + + +def kernel_launcher( + call_fn: Callable[..., TileLangLaunchSpec], +) -> Callable[..., Any]: + """Launch TileLang from declarative kernel and runtime argument tuples.""" + + @wraps(call_fn) + def wrapper( + self: VllmTileLangJitKernel[Any], + *args: Any, + **kwargs: Any, + ) -> Any: + launch_spec = call_fn(self, *args, **kwargs) + kernel_args, launch_args = launch_spec[:2] + if self._warming_key is not None and kernel_args not in ( + (), + (self._warming_key,), + ): + raise RuntimeError( + "TileLang warmup inputs produced a different compile key: " + f"expected {self._warming_key!r}, got kernel arguments " + f"{kernel_args!r}" + ) + launch_kwargs = launch_spec[2] if len(launch_spec) > 2 else None + output = launch_spec[3] if len(launch_spec) > 3 else None + result = self.launch(self.kernel(*kernel_args), launch_args, launch_kwargs) + return output if len(launch_spec) > 3 else result + + return wrapper diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e56ae50ebbca..80b5160d26c8 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -17,9 +17,6 @@ from vllm.model_executor.warmup.b12x_warmup import b12x_warmup from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup -from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( - deepseek_v4_mhc_warmup, -) from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, write_flashinfer_autotune_cache, @@ -197,15 +194,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): compilation_config = worker.vllm_config.compilation_config cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or []) - # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder - # layer per token; warm them across token sizes first so the first real - # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). - deepseek_v4_mhc_warmup( - worker.get_model(), - max_tokens=worker.scheduler_config.max_num_batched_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - # Run next so input-prep kernels JIT against pristine runner state. if enable_jit_warmup: kimi_k3_triton_warmup(worker) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 97e7e617fe9c..b1379cc9eb45 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1210,6 +1210,53 @@ def __init__( requires_grad=False, ) + if vllm_config.kernel_config.enable_jit_warmup: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, + _MHC_FUSED_TILELANG_KERNEL, + _MHC_POST_TILELANG_KERNEL, + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, + ) + from vllm.utils.deep_gemm import is_deep_gemm_supported + + include_pre_gemm_splits = is_deep_gemm_supported() + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.hc_mult, + use_norm_weight=True, + include_pre_gemm_splits=include_pre_gemm_splits, + include_broadcast_splits=( + get_pp_group().is_first_rank and extract_layer_index(prefix) == 0 + ), + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + norm_eps=( + self.attn_norm.variance_epsilon, + self.ffn_norm.variance_epsilon, + ), + broadcast_norm_eps=self.attn_norm.variance_epsilon, + ) + if not include_pre_gemm_splits: + _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.hc_mult, + n_out=self.hc_mult * (2 + self.hc_mult), + ) + _MHC_POST_TILELANG_KERNEL.register_warmup( + hidden_size=self.hidden_size, + hc_mult=self.hc_mult, + ) + _MHC_FUSED_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.hc_mult, + ) + def forward( self, x: torch.Tensor, @@ -1402,6 +1449,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): else: self._mtp_hidden_buffer = None + if vllm_config.kernel_config.enable_jit_warmup and get_pp_group().is_last_rank: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_HEAD_FUSED_TILELANG_KERNEL, + ) + + _HC_HEAD_FUSED_TILELANG_KERNEL.register_warmup( + hidden_size=config.hidden_size, + hc_mult=self.hc_mult, + rms_eps=self.rms_norm_eps, + hc_eps=self.hc_eps, + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 8444c8141611..21918bdf6bed 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -403,6 +403,50 @@ def __init__( self.mhc_post_op = MHCPostOp() self.mhc_fused_post_pre_op = MHCFusedPostPreOp() + if vllm_config.kernel_config.enable_jit_warmup: + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, + _MHC_FUSED_TILELANG_KERNEL, + _MHC_POST_TILELANG_KERNEL, + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, + ) + from vllm.utils.deep_gemm import is_deep_gemm_supported + + include_pre_gemm_splits = is_deep_gemm_supported() + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.n, + use_norm_weight=True, + include_pre_gemm_splits=include_pre_gemm_splits, + include_broadcast_splits=False, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.mhc_post_mult_value, + sinkhorn_repeat=self.mhc_sinkhorn_iterations, + norm_eps=( + self.input_layernorm.variance_epsilon, + self.post_attention_layernorm.variance_epsilon, + ), + ) + if not include_pre_gemm_splits: + _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.n, + n_out=self.n * (2 + self.n), + ) + _MHC_POST_TILELANG_KERNEL.register_warmup( + hidden_size=self.hidden_size, + hc_mult=self.n, + ) + _MHC_FUSED_TILELANG_KERNEL.register_warmup( + vllm_config, + hidden_size=self.hidden_size, + hc_mult=self.n, + ) + def forward( self, positions: torch.Tensor, From 533dbe424b9bdde6a7ffe100d507a69a24dfee1a Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 13 Sep 2026 11:16:39 +0800 Subject: [PATCH 114/774] [Misc] Clean up fast loader daemon quant method verification (#55656) Signed-off-by: Isotr0py --- .../model_loader/base_loader.py | 18 ++- .../model_loader/weight_cache/__init__.py | 2 - .../model_loader/weight_cache/daemon.py | 38 +++++-- .../model_loader/weight_cache/ipc_loader.py | 28 ++--- .../model_loader/weight_cache/protocol.py | 106 ++++-------------- 5 files changed, 75 insertions(+), 117 deletions(-) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 618f874fd0d6..6f356464585a 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -39,6 +39,19 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: inplace weights loading for an already-initialized model""" raise NotImplementedError + def create_model( + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" + ) -> nn.Module: + """Create a model with the given configurations.""" + model = initialize_model( + vllm_config=vllm_config, + model_config=model_config, + prefix=prefix, + ) + log_online_quantization(vllm_config) + log_model_inspection(model) + return model + @instrument(span_name="Load model") def load_model( self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" @@ -52,15 +65,12 @@ def load_model( target_device = torch.device(load_device) with set_default_torch_dtype(model_config.dtype): with target_device: - model = initialize_model( + model = self.create_model( vllm_config=vllm_config, model_config=model_config, prefix=prefix, ) - log_online_quantization(vllm_config) - log_model_inspection(model) - logger.debug("Loading weights on %s ...", load_device) self.load_weights(model, model_config) diff --git a/vllm/model_executor/model_loader/weight_cache/__init__.py b/vllm/model_executor/model_loader/weight_cache/__init__.py index 0cea5cf72a21..f9dbe9bb32a1 100644 --- a/vllm/model_executor/model_loader/weight_cache/__init__.py +++ b/vllm/model_executor/model_loader/weight_cache/__init__.py @@ -11,7 +11,6 @@ WeightCacheUnavailableError, check_ipc_platform_support, check_ipc_quant_support, - is_ipc_quant_supported, ) __all__ = [ @@ -24,5 +23,4 @@ "WeightCacheUnavailableError", "check_ipc_platform_support", "check_ipc_quant_support", - "is_ipc_quant_supported", ] diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 083910abb4d9..31f59a6ac5ff 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -39,6 +39,8 @@ init_distributed_environment, ) from vllm.logger import init_logger +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.model_executor.model_loader.weight_cache.protocol import ( TensorEntry, WeightCacheKey, @@ -53,6 +55,7 @@ verify_peer_is_owner, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_default_torch_dtype logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") @@ -98,6 +101,30 @@ def _add(name: str, tensor: torch.Tensor, kind: str) -> None: return entries, aliases +def get_daemon_model(vllm_config: VllmConfig) -> torch.nn.Module: + """Load the daemon's model, composed from the configured loader. + + Runs the quantization check after model creation but before the slow + weight load, so an unsupported method fails fast. Online quantization + always fails the check, so load_model's finalize step for it is + unnecessary here. + """ + model_config = vllm_config.model_config + load_config = vllm_config.load_config + loader = get_model_loader(load_config) + device_config = vllm_config.device_config + target_device = torch.device( + device_config.device if load_config.device is None else load_config.device + ) + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = loader.create_model(vllm_config, model_config) + check_ipc_quant_support(model) + loader.load_weights(model, model_config) + process_weights_after_loading(model, model_config, target_device) + return model.eval() + + class WeightCacheDaemon: """Per-GPU process that loads one TP shard and serves CUDA IPC handles.""" @@ -124,8 +151,6 @@ def __init__( ) def load_model(self) -> None: - from vllm.model_executor.model_loader import get_model - tp_size = self.cache_config.tp_size torch.accelerator.set_device_index(self.tp_rank) init_distributed_environment( @@ -137,7 +162,7 @@ def load_model(self) -> None: ) with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(tp_size, 1) - self.model = get_model(vllm_config=self.vllm_config) + self.model = get_daemon_model(self.vllm_config) self._export_entries() logger.info( "Weight cache daemon rank %d cached %d tensors", @@ -331,10 +356,9 @@ def main() -> None: "The weight cache daemon itself must load from disk; use the " "default --load-format" ) - # Checked before loading anything: an unsupported quantization method would - # otherwise only surface in the engine, after a full load. - check_ipc_platform_support(where="daemon") - check_ipc_quant_support(vllm_config.model_config, where="daemon") + # Config-only so it can fail before any model loading; the quant method + # check needs the created model and runs in get_daemon_model. + check_ipc_platform_support() parallel_config = vllm_config.parallel_config _reject_unsupported_parallelism(parallel_config) tp_size = parallel_config.tensor_parallel_size diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 2132a55fbea5..be2e5ebae4cf 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -119,9 +119,9 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: def load_model( self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" ) -> nn.Module: - # Unsupported quantization is a permanent misconfiguration rather than a - # transient daemon outage, so it is raised even when fallback is on. - self._check_supported(vllm_config, model_config) + # An unsupported platform is a permanent misconfiguration rather than + # a transient daemon outage, so it is raised even when fallback is on. + check_ipc_platform_support() state_fetched = False try: entries, aliases = self._fetch_entries(model_config) @@ -135,6 +135,11 @@ def load_model( logger.warning( "Weight cache unusable (%s); falling back to disk loading", e ) + except UnsupportedQuantForIPCError: + # Unsupported quantization is a permanent misconfiguration rather + # than a transient daemon outage, so it is raised even when + # fallback is on. + raise except Exception: if not self.fallback: raise @@ -176,6 +181,7 @@ def _build_model( model_config=model_config, prefix=prefix, ) + check_ipc_quant_support(model) self._apply_entries(model, entries, aliases, device_index) # The daemon exports tensors that already went through # process_weights_after_loading; re-run it in pre-processed mode @@ -197,22 +203,6 @@ def _build_model( ) return model.eval() - @staticmethod - def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None: - check_ipc_platform_support(where="engine") - check_ipc_quant_support(model_config, where="engine") - cache_dtype = vllm_config.cache_config.cache_dtype - if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"): - # BaseKVCacheMethod.process_weights_after_loading turns the loaded - # k/v scale parameters into plain float attributes. For fp8 cache - # dtypes those are rebuilt from the exported scale buffers when - # process_weights_after_loading runs in pre-processed mode; other - # quantized cache dtypes are not verified. - raise UnsupportedQuantForIPCError( - f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not " - "supported by the weight cache; use --kv-cache-dtype auto." - ) - def _apply_entries( self, model: nn.Module, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 08e5e725bda7..fff1fcca2a0b 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -26,6 +26,7 @@ import vllm.version from vllm.config import ModelConfig +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash @@ -59,15 +60,12 @@ class UnsupportedPlatformForIPCError(Exception): """Raised when the current platform cannot share CUDA IPC handles.""" -def check_ipc_platform_support(*, where: str) -> None: +def check_ipc_platform_support() -> None: """Hard-error unless the current platform can share CUDA IPC handles. Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other platforms (e.g. XPU) would silently ship every tensor by value instead. - Args: - where: Short tag ("daemon"/"engine") used in the error message. - Raises: UnsupportedPlatformForIPCError: If the current platform is not CUDA/ROCm. @@ -75,95 +73,33 @@ def check_ipc_platform_support(*, where: str) -> None: if current_platform.is_cuda_alike(): return raise UnsupportedPlatformForIPCError( - f"[weight_cache:{where}] platform {current_platform.device_name!r} " - "does not support CUDA IPC weight sharing; only CUDA and ROCm are " - "supported. Use the default --load-format for this platform." + f"platform {current_platform.device_name!r} does not support CUDA IPC " + "weight sharing; only CUDA and ROCm are supported. Use the default " + "--load-format for this platform." ) -# The daemon exports tensor data only, so sharing is correct just for methods -# whose post-load effect is either fully captured by that data or rebuilt when -# process_weights_after_loading runs under weights_already_processed. Anything -# absent from this registry hard-errors: a method that silently repacks into -# shapes the client cannot reproduce or stamps unrebuildable Python-side state -# would serve wrong numerics. Extend it only after an end-to-end check against -# a disk-loaded baseline. -# -# This check is config-only so the daemon fails fast before a slow full load, -# rather than surfacing the incompatibility only after the engine has built the -# model (where process_weights_after_loading's supports_pre_processed_weights -# guard would eventually catch it). - - -def _quant_config_field(quant_config: Any, key: str) -> Any: - if quant_config is None: - return None - if isinstance(quant_config, dict): - return quant_config.get(key) - return getattr(quant_config, key, None) - - -def _fp8_round_trips_via_ipc(quant_config: Any) -> bool: - """Only block-wise FP8 is verified. - - Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes - ``layer.weight`` during post-processing -- a shape the client's - meta-initialized model cannot reproduce. - """ - return _quant_config_field(quant_config, "weight_block_size") is not None - - -# quantization name -> predicate(quant_config) -> True when verified safe. -IPC_QUANT_ALLOWLIST: dict[str | None, Any] = { - None: lambda _quant_config: True, # unquantized - "fp8": _fp8_round_trips_via_ipc, - "modelopt_fp4": lambda _quant_config: True, - # Kimi-K3 routed experts (mxfp4-pack): the MegaMoE experts keep their packed - # weights and scales as plain parameters and defer the DeepGEMM transform to - # forward(), so the exported tensors are complete. Any other mxfp4 method - # that repacks or drops parameters still fails closed at the engine's - # supports_pre_processed_weights guard. - "mxfp4": lambda _quant_config: True, -} - - -def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool: - predicate = IPC_QUANT_ALLOWLIST.get(quantization) - return False if predicate is None else bool(predicate(quant_config)) - - -def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None: - """Hard-error unless the model's quantization is verified for IPC sharing. +# The daemon transfers post processed weights directly +def check_ipc_quant_support(model: torch.nn.Module) -> None: + """Hard-error unless every quant method supports pre-processed weights. Args: - model_config: Model configuration to inspect. - where: Short tag ("daemon"/"engine") used in the error message. + model: The model to inspect (weights need not be loaded). Raises: - UnsupportedQuantForIPCError: If the quantization method is not on the - verified allowlist. + UnsupportedQuantForIPCError: If any quant method does not declare + ``supports_pre_processed_weights``. """ - quantization = model_config.quantization - # Prefer the canonical, nested-aware config (multimodal models keep it under - # text_config); fall back to the raw hf_config for older code paths. - quant_config = getattr( - getattr(model_config, "model_arch_config", None), "quantization_config", None - ) - if quant_config is None: - quant_config = getattr(model_config.hf_config, "quantization_config", None) - if is_ipc_quant_supported(quantization, quant_config): - return - verified = ", ".join( - "unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST - ) - raise UnsupportedQuantForIPCError( - f"[weight_cache:{where}] quantization {quantization!r} is not verified " - f"for CUDA IPC weight sharing: its post-load processing may repack " - f"weights into shapes the client cannot reproduce or stamp Python-side " - f"state that tensor export cannot carry, which would silently serve " - f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size " - f"set, i.e. block-wise). Use the default --load-format for this model." - ) + for name, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if ( + isinstance(quant_method, QuantizeMethodBase) + and not quant_method.supports_pre_processed_weights + ): + raise UnsupportedQuantForIPCError( + f"layer {name or ''}: {type(quant_method).__name__} " + "does not support loading from pre-processed weights." + ) def get_physical_device_id(device_index: int) -> int | None: From c2ea9f1a5f10caec36eac91c461fcec6564f5394 Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Sat, 12 Sep 2026 23:27:08 -0400 Subject: [PATCH 115/774] [Kernel][Perf][Quantization] Fix odd-row performance cliff in per-token-group quantization (#56478) Signed-off-by: Ayushman Singh Co-authored-by: Ayushman Singh --- .../w8a8/fp8/per_token_group_quant.cu | 146 +++++++++++------- .../test_per_token_group_quant.py | 49 +++++- 2 files changed, 131 insertions(+), 64 deletions(-) diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index 06fa29eaba29..cfd1fab1ad00 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -96,7 +96,8 @@ __device__ __forceinline__ void QuantizeGroup( } template + bool SCALE_UE8M0 = false, typename scale_packed_t = float, + bool CHECK_BOUNDS = false> __global__ void per_token_group_quant_8bit_kernel( const T* __restrict__ input, void* __restrict__ output_q, scale_packed_t* __restrict__ output_s, const int group_size, @@ -114,49 +115,54 @@ __global__ void per_token_group_quant_8bit_kernel( using scale_element_t = float; static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); - const T* group_input = input + block_group_offset; - DST_DTYPE* group_output = - static_cast(output_q) + block_group_offset; - scale_element_t* scale_output; - #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - if constexpr (IS_COLUMN_MAJOR) { - const int num_elems_per_pack = - static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); - const int scale_num_rows_element = scale_num_rows * num_elems_per_pack; - const int row_idx = global_group_id / scale_num_rows_element; - const int col_idx_raw = global_group_id % scale_num_rows_element; - const int col_idx = col_idx_raw / num_elems_per_pack; - const int pack_idx = col_idx_raw % num_elems_per_pack; - scale_output = reinterpret_cast(output_s) + - (col_idx * scale_stride * num_elems_per_pack + - row_idx * num_elems_per_pack + pack_idx); - } else { - scale_output = output_s + global_group_id; - } - // shared memory to cache each group's data to avoid double DRAM reads. extern __shared__ __align__(16) char smem_raw[]; T* smem = reinterpret_cast(smem_raw); T* smem_group = smem + local_group_id * group_size; + const bool valid_group = !CHECK_BOUNDS || global_group_id < num_groups; + float y_s = 1.0f; + if (valid_group) { + const T* group_input = input + block_group_offset; + scale_element_t* scale_output; + + if constexpr (IS_COLUMN_MAJOR) { + const int num_elems_per_pack = + static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); + const int scale_num_rows_element = scale_num_rows * num_elems_per_pack; + const int row_idx = global_group_id / scale_num_rows_element; + const int col_idx_raw = global_group_id % scale_num_rows_element; + const int col_idx = col_idx_raw / num_elems_per_pack; + const int pack_idx = col_idx_raw % num_elems_per_pack; + scale_output = reinterpret_cast(output_s) + + (col_idx * scale_stride * num_elems_per_pack + + row_idx * num_elems_per_pack + pack_idx); + } else { + scale_output = output_s + global_group_id; + } - const float y_s = ComputeGroupScale( - group_input, smem_group, group_size, lane_id, threads_per_group, eps, - max_8bit); + y_s = ComputeGroupScale(group_input, smem_group, group_size, + lane_id, threads_per_group, eps, + max_8bit); - scale_element_t y_s_quant = y_s; + scale_element_t y_s_quant = y_s; - if (lane_id == 0) { - *scale_output = y_s_quant; + if (lane_id == 0) { + *scale_output = y_s_quant; + } } __syncthreads(); - QuantizeGroup(smem_group, group_output, group_size, lane_id, - threads_per_group, y_s, min_8bit, max_8bit); + if (valid_group) { + DST_DTYPE* group_output = + static_cast(output_q) + block_group_offset; + QuantizeGroup(smem_group, group_output, group_size, lane_id, + threads_per_group, y_s, min_8bit, max_8bit); + } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -209,10 +215,20 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, constexpr int THREADS_PER_GROUP = 16; - const int groups_per_block = GetGroupsPerBlock(num_groups); + if (num_groups == 0) { + return; + } + STD_TORCH_CHECK(input.size(-1) % group_size == 0); + const int64_t groups_per_row = input.size(-1) / group_size; + const int64_t num_rows = num_groups / groups_per_row; + // Odd row counts inherit the next even row count's launch geometry. + const int64_t policy_num_groups = + (num_rows + (num_rows & 1)) * groups_per_row; + const int groups_per_block = GetGroupsPerBlock(policy_num_groups); + const bool needs_bounds_check = num_groups % groups_per_block != 0; auto dst_type = output_q.scalar_type(); - const int num_blocks = num_groups / groups_per_block; + const int num_blocks = (num_groups + groups_per_block - 1) / groups_per_block; const int num_threads = groups_per_block * THREADS_PER_GROUP; const bool is_column_major = output_s.stride(0) < output_s.stride(1); @@ -220,38 +236,51 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, const int scale_stride = output_s.stride(1); #ifndef USE_ROCM - #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ - do { \ - cudaLaunchConfig_t config = {}; \ - config.gridDim = dim3(num_blocks); \ - config.blockDim = dim3(num_threads); \ - config.dynamicSmemBytes = (SMEM_BYTES); \ - config.stream = stream; \ - cudaLaunchAttribute attrs[1]; \ - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ - attrs[0].val.programmaticStreamSerializationAllowed = 1; \ - config.numAttrs = 1; \ - config.attrs = attrs; \ - cudaLaunchKernelEx( \ - &config, \ - per_token_group_quant_8bit_kernel, \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, num_groups, \ - groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \ - scale_num_rows, scale_stride); \ + #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \ + SMEM_BYTES) \ + do { \ + cudaLaunchConfig_t config = {}; \ + config.gridDim = dim3(num_blocks); \ + config.blockDim = dim3(num_threads); \ + config.dynamicSmemBytes = (SMEM_BYTES); \ + config.stream = stream; \ + cudaLaunchAttribute attrs[1]; \ + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ + attrs[0].val.programmaticStreamSerializationAllowed = 1; \ + config.numAttrs = 1; \ + config.attrs = attrs; \ + cudaLaunchKernelEx( \ + &config, \ + per_token_group_quant_8bit_kernel, \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, num_groups, \ + groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \ + scale_num_rows, scale_stride); \ } while (0) #else - #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ - do { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit, scale_num_rows, scale_stride); \ + #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \ + SMEM_BYTES) \ + do { \ + per_token_group_quant_8bit_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, \ + num_groups, groups_per_block, (float)eps, (float)min_8bit, \ + (float)max_8bit, scale_num_rows, scale_stride); \ } while (0) #endif +#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ + do { \ + if (needs_bounds_check) { \ + LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, true, SMEM_BYTES); \ + } else { \ + LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, false, SMEM_BYTES); \ + } \ + } while (0) + #define LAUNCH_KERNEL(T, DST_DTYPE) \ do { \ size_t smem_bytes = \ @@ -284,6 +313,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, #undef LAUNCH_KERNEL #undef LAUNCH_KERNEL_INST +#undef LAUNCH_KERNEL_IMPL } // Register-resident fast path for group_size==128. diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index d88a39d08469..f2326bfd25fc 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -11,8 +11,24 @@ @pytest.mark.parametrize( - "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)] + "shape", + [ + (0, 128), + (1, 128), + (3, 384), + (5, 384), + (7, 384), + (15, 384), + (31, 128), + (32, 128), + (63, 256), + (64, 256), + (16, 512), + (3, 2688), + (3, 5376), + ], ) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("column_major", [False, True]) @pytest.mark.parametrize("tma_aligned", [False, True]) @pytest.mark.parametrize("scale_ue8m0", [False, True]) @@ -22,14 +38,20 @@ reason="Only test on CUDA/ROCm/XPU.", ) def test_per_token_group_quant_fp8( - shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int + shape, + dtype, + column_major: bool, + tma_aligned: bool, + scale_ue8m0: bool, + group_size: int, ): + """Tail groups preserve FP8 values and scales in every supported layout.""" device = current_platform.device_type torch.manual_seed(42) num_tokens, hidden_dim = shape - x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 + x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8 # native kernel path out_q, scale = fp8_utils.per_token_group_quant_fp8( @@ -411,19 +433,34 @@ def test_per_token_group_quant_fp8_packed_large_mn(): assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch" -@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)]) +@pytest.mark.parametrize( + "shape", + [ + (0, 128), + (1, 128), + (3, 384), + (5, 384), + (7, 384), + (15, 384), + (32, 128), + (64, 256), + (16, 512), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("group_size", [64, 128]) @pytest.mark.skipif( not (current_platform.is_cuda_alike() or current_platform.is_xpu()), reason="Only test on CUDA/ROCm/XPU.", ) -def test_per_token_group_quant_int8(shape, group_size: int): +def test_per_token_group_quant_int8(shape, dtype, group_size: int): + """The shared launch policy also handles INT8 tail groups.""" device = current_platform.device_type torch.manual_seed(42) num_tokens, hidden_dim = shape - x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 + x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8 # cuda path out_q, scale = int8_utils.per_token_group_quant_int8( From c711f740b7c7540c8c4d7f2c26ff6960b427a4cc Mon Sep 17 00:00:00 2001 From: Markus Hartikainen Date: Sat, 12 Sep 2026 20:47:07 -0700 Subject: [PATCH 116/774] [ROCm] Auto-enable breakable CUDA graphs for DeepseekV41ForCausalLM (#56349) Signed-off-by: Markus Hartikainen Co-authored-by: Jee Jee Li --- tests/test_config.py | 4 ++++ vllm/config/vllm.py | 13 ++++++++++--- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index e6d040caf5ba..218ec6144562 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -358,6 +358,10 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): assert "DeepseekV32ForCausalLM" not in breakable_architectures assert "DeepseekV32MTPModel" not in breakable_architectures assert "GlmMoeDsaForCausalLM" not in breakable_architectures + # V4.1 cannot torch.compile and the ROCm sparse SWA backend only + # supports uniform-batch CUDA graphs, so it must opt in or the + # default FULL_AND_PIECEWISE serve path cannot start. + assert "DeepseekV41ForCausalLM" in breakable_architectures # The carve-out takes effect via the runner-selection property # (warning_once args must be hashable for its lru_cache). diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3bdfe1f3edb5..6fab6f0c7177 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -110,10 +110,17 @@ def default_breakable_cudagraph_architectures() -> frozenset[str]: from vllm.platforms import current_platform if current_platform.is_rocm(): - # Breakable CUDA graphs currently regress performance on ROCm, so no - # architecture opts in by default here. Users can still force it with + # Breakable CUDA graphs currently regress performance on ROCm for + # models that can use torch.compile piecewise graphs instead. Do not + # opt those in by default. Users can still force them with # VLLM_USE_BREAKABLE_CUDAGRAPH=1. - return frozenset() + # + # DeepseekV41ForCausalLM cannot torch.compile, and the ROCm sparse + # SWA backend only reports AttentionCGSupport.UNIFORM_BATCH. Default + # FULL_AND_PIECEWISE then dies at capture unless breakable CUDA + # graphs are on. Enable this architecture so the published AMD + # recipe can start. + return frozenset({"DeepseekV41ForCausalLM"}) return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES From 2671fedfc7ae604761990603fc736c0c4f21de57 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 12 Sep 2026 21:07:56 -0700 Subject: [PATCH 117/774] [MRV2] Revert explicit Triton JIT warmup migration (#56654) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- .../test_gpu_autoregressive_speculator.py | 1 - vllm/v1/worker/gpu/model_runner.py | 97 ++++++------- vllm/v1/worker/gpu/sample/states.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 131 ++++++------------ .../multi_module_mtp/speculator.py | 88 ++---------- 6 files changed, 97 insertions(+), 228 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index ce47cb8ae625..16d6d8e516df 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -43,7 +43,7 @@ def _run_prepare( num_scheduled_tokens=np.array([4], dtype=np.int32), positions=torch.tensor(target_positions, dtype=torch.int64, device=device), query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device), - idx_mapping=torch.tensor([2], dtype=torch.int64, device=device), + idx_mapping=torch.tensor([2], dtype=torch.int32, device=device), ) query_slot_mapping = torch.full( (max_num_tokens,), -2, dtype=torch.int64, device=device diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index a38265a6f7af..920c9a25f062 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -157,7 +157,6 @@ def init_base(speculator, vllm_config, device): speculator.device = device speculator.max_num_tokens = 4 speculator.max_num_reqs = 2 - speculator.num_speculative_steps = 1 speculator.hidden_size = 3 speculator.dtype = torch.float32 speculator.draft_model_config = draft_model_config diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 14e0ac46297e..0ce0288f310a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -188,7 +188,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): - @JitWarmupRegistry.capture def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -203,6 +202,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.speculative_config is not None and self.speculative_config.use_dspark() ) self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) self.device = device self.dtype = self.model_config.dtype @@ -399,7 +399,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: logger.info_once("Loading model from scratch...") # Capture warmup providers selected while constructing the model. - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] + with self.jit_warmup_registry.activate(): self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.vllm_config.model_config, @@ -461,50 +461,47 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] - sampler_kwargs: dict[str, Any] = { - "max_num_reqs": self.max_num_reqs, - "vocab_size": self.vocab_size, - "device": self.device, - "req_states": self.req_states, - "logprobs_mode": self.model_config.logprobs_mode, - "num_speculative_tokens": self.decode_query_len, - "use_fp64_gumbel": self.model_config.use_fp64_gumbel, - "enable_trace_replay": self.model_config.enable_trace_replay, - "reasoning_config": self.vllm_config.reasoning_config, - "return_sampling_mask": self.model_config.return_sampling_mask, - } - if self.vllm_config.watermark_config is None: - self.sampler = Sampler(**sampler_kwargs) - else: - wm_config = self.vllm_config.watermark_config - watermarker = create_watermarker(wm_config) - if self.speculative_config is not None: - watermarker = create_speculative_target_watermarker(watermarker) - self.sampler = GPUWatermarkSampler( - watermarker, - deduplicate_contexts=wm_config.deduplicate_contexts, - deduplicate_contexts_max_history=( - wm_config.deduplicate_contexts_max_history - ), - **sampler_kwargs, - ) - custom = self.model_state.custom_sampler(self.sampler) - - if custom: - self.vllm_config._check_watermarking_unsupported( - custom_sampler=True - ) - self.sampler, self.rejection_sampler = custom - elif self.speculative_config is not None: - self.rejection_sampler = RejectionSampler( - self.sampler, - self.speculative_config, - self.device, - watermark_key=speculative_target_watermark_key( - self.vllm_config.watermark_config - ), - ) + sampler_kwargs: dict[str, Any] = { + "max_num_reqs": self.max_num_reqs, + "vocab_size": self.vocab_size, + "device": self.device, + "req_states": self.req_states, + "logprobs_mode": self.model_config.logprobs_mode, + "num_speculative_tokens": self.decode_query_len, + "use_fp64_gumbel": self.model_config.use_fp64_gumbel, + "enable_trace_replay": self.model_config.enable_trace_replay, + "reasoning_config": self.vllm_config.reasoning_config, + "return_sampling_mask": self.model_config.return_sampling_mask, + } + if self.vllm_config.watermark_config is None: + self.sampler = Sampler(**sampler_kwargs) + else: + wm_config = self.vllm_config.watermark_config + watermarker = create_watermarker(wm_config) + if self.speculative_config is not None: + watermarker = create_speculative_target_watermarker(watermarker) + self.sampler = GPUWatermarkSampler( + watermarker, + deduplicate_contexts=wm_config.deduplicate_contexts, + deduplicate_contexts_max_history=( + wm_config.deduplicate_contexts_max_history + ), + **sampler_kwargs, + ) + custom = self.model_state.custom_sampler(self.sampler) + + if custom: + self.vllm_config._check_watermarking_unsupported(custom_sampler=True) + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + watermark_key=speculative_target_watermark_key( + self.vllm_config.watermark_config + ), + ) self.prompt_logprobs_worker = PromptLogprobsWorker( self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode, @@ -741,7 +738,7 @@ def initialize_kv_cache( self.speculator.init_cudagraph_manager(cudagraph_mode) # Capture warmup providers that depend on allocated KV-cache strides. - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] + with self.jit_warmup_registry.activate(): kv_caches_dict = init_kv_cache( self.compilation_config.static_forward_context, self.kv_cache_config, @@ -927,12 +924,6 @@ def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: # during actual execution. assert self.sampler is not None self.sampler(logits, dummy_input_batch) - if self.rejection_sampler is not None: - # Warm head-dtype logits; scheduler warmup covers processed FP32 logits. - assert self.speculator is not None - self.rejection_sampler( - logits, dummy_input_batch, self.speculator.draft_logits - ) @torch.inference_mode() def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index 08cae8212836..fe4dee6a6b10 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -4,10 +4,7 @@ import torch from vllm.sampling_params import SamplingParams -from vllm.v1.sample.ops.topk_topp_sampler import ( - apply_top_k_top_p, - register_top_k_top_p_warmups, -) +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.sample.gumbel import apply_temperature from vllm.v1.worker.gpu.sample.min_p import apply_min_p @@ -21,7 +18,6 @@ class SamplingStates: def __init__(self, max_num_reqs: int, vocab_size: int): self.max_num_reqs = max_num_reqs self.vocab_size = vocab_size - register_top_k_top_p_warmups() self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 10f4346b5420..a5710829e0ff 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -from types import SimpleNamespace from typing import Any -import numpy as np import torch import torch.nn as nn @@ -12,15 +10,6 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup import ( - WarmupChoices, - WarmupIntRange, -) -from vllm.model_executor.warmup.jit_warmup_triton_helper import ( - DispatchSpec, - TritonWarmupTensor, - triton_kernel_dispatcher_with_warmup, -) from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -119,7 +108,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 - prepare_dflash_inputs.register_warmup(speculator=self) @property def attn_vllm_config(self) -> VllmConfig: @@ -694,64 +682,6 @@ def _prepare_dflash_inputs_kernel( tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) -def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]: - int32 = TritonWarmupTensor(torch.int32) - int64 = TritonWarmupTensor(torch.int64) - float32 = TritonWarmupTensor(torch.float32) - input_buffers = SimpleNamespace( - input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32 - ) - gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids) - max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2) - block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0)) - block_table = TritonWarmupTensor( - torch.int32, - shape=(1, block_table_stride), - strides=(block_table_stride, 1), - ) - input_batch = SimpleNamespace( - num_reqs=1, - num_scheduled_tokens=np.array([max_target_query_len]), - positions=int64, - query_start_loc=int32, - idx_mapping=int64, - ) - return dict( - input_buffers=input_buffers, - query_slot_mapping=int64, - context_positions=int64, - context_slot_mapping=int64, - sample_indices=int64, - sample_pos=int64, - sample_idx_mapping=int32, - temperature=float32, - seeds=int64, - input_batch=input_batch, - num_sampled=int32, - num_rejected=int32, - last_sampled=int64, - next_prefill_tokens=int32, - input_temperature=float32, - input_seeds=int64, - block_table=block_table, - block_size=speculator.block_tables.kernel_block_sizes[gid], - cp_rank=speculator.block_tables.cp_rank, - cp_size=speculator.block_tables.cp_size, - cp_interleave=speculator.block_tables.cp_interleave, - parallel_drafting_token_id=speculator.parallel_drafting_token_id, - num_query_per_req=speculator.num_query_per_req, - num_speculative_steps=speculator.num_speculative_steps, - max_num_reqs=speculator.max_num_reqs, - max_num_tokens=speculator.max_num_tokens, - max_model_len=speculator.max_model_len, - sample_from_anchor=speculator.sample_from_anchor, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_prepare_dflash_inputs_kernel, - warmup_inputs=_prepare_dflash_warmup_inputs, -) def prepare_dflash_inputs( input_buffers: InputBuffers, query_slot_mapping: torch.Tensor, @@ -788,33 +718,50 @@ def prepare_dflash_inputs( max_num_tokens: int, max_model_len: int, sample_from_anchor: bool = False, -) -> DispatchSpec: +) -> None: num_reqs = input_batch.num_reqs assert num_reqs > 0 # Cover the longest possible per-request span (ctx + query). Use the max # per-request query length, not the total token count across the batch. max_target_query_len = int(input_batch.num_scheduled_tokens.max()) max_tokens_per_req = max_target_query_len + num_query_per_req - block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) - num_blocks = triton.cdiv(max_tokens_per_req, block) - return (num_reqs, num_blocks), dict( - out_input_ids_ptr=input_buffers.input_ids, - out_query_positions_ptr=input_buffers.positions, - out_query_start_loc_ptr=input_buffers.query_start_loc, - out_seq_lens_ptr=input_buffers.seq_lens, - out_query_slot_mapping_ptr=query_slot_mapping, - out_context_positions_ptr=context_positions, - out_context_slot_mapping_ptr=context_slot_mapping, - out_sample_indices_ptr=sample_indices, - out_sample_pos_ptr=sample_pos, - out_sample_idx_mapping_ptr=sample_idx_mapping, - out_temperature_ptr=temperature, - out_seeds_ptr=seeds, - target_positions_ptr=input_batch.positions, - target_query_start_loc_ptr=input_batch.query_start_loc, - idx_mapping_ptr=input_batch.idx_mapping, - temperature_ptr=input_temperature, - seeds_ptr=input_seeds, + BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) + num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) + _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( + input_buffers.input_ids, + input_buffers.positions, + input_buffers.query_start_loc, + input_buffers.seq_lens, + query_slot_mapping, + context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch.positions, + input_batch.query_start_loc, + input_batch.idx_mapping, + last_sampled, + next_prefill_tokens, + num_sampled, + num_rejected, + input_temperature, + input_seeds, + block_table, + block_table.stride(0), + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + max_model_len, + cp_rank, + SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, - BLOCK_SIZE=block, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, + BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c37ccb6c2db..1c6ae315cf46 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -9,11 +9,6 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup_triton_helper import ( - DispatchSpec, - TritonWarmupTensor, - triton_kernel_dispatcher_with_warmup, -) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -74,9 +69,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.cudagraph_manager: SpeculatorCudaGraphManager | None = None - _shift_input_ids.register_warmup() - _shift_input_embeds.register_warmup(speculator=self) - def load_draft_model( self, target_model: nn.Module, @@ -1016,32 +1008,6 @@ def _shift_input_ids_kernel( tl.store(input_ids_ptr + last_token_index, draft_token) -def _shift_input_ids_warmup_inputs() -> dict[str, Any]: - return dict( - input_ids=TritonWarmupTensor(torch.int32), - idx_mapping=TritonWarmupTensor(torch.int32), - query_start_loc=TritonWarmupTensor(torch.int32), - last_token_indices=TritonWarmupTensor(torch.int64), - draft_tokens=TritonWarmupTensor(torch.int64), - num_reqs=1, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_shift_input_ids_kernel, - warmup_inputs=_shift_input_ids_warmup_inputs, -) -def _shift_input_ids( - input_ids: torch.Tensor, - idx_mapping: torch.Tensor, - query_start_loc: torch.Tensor, - last_token_indices: torch.Tensor, - draft_tokens: torch.Tensor, - num_reqs: int, -) -> DispatchSpec: - return (num_reqs,), dict(BLOCK_SIZE=1024) - - @triton.jit def _shift_input_embeds_kernel( input_embeds_ptr, @@ -1099,44 +1065,6 @@ def _shift_input_embeds_kernel( ) -def _shift_input_embeds_warmup_inputs( - *, speculator: MultiModuleMTPSpeculator -) -> dict[str, Any]: - hidden = TritonWarmupTensor( - speculator.dtype, - shape=(1, speculator.hidden_size), - strides=(speculator.hidden_size, 1), - ) - return dict( - input_embeds=hidden, - draft_embeds=hidden, - idx_mapping=TritonWarmupTensor(torch.int32), - query_start_loc=TritonWarmupTensor(torch.int32), - last_token_indices=TritonWarmupTensor(torch.int64), - num_reqs=1, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_shift_input_embeds_kernel, - warmup_inputs=_shift_input_embeds_warmup_inputs, -) -def _shift_input_embeds( - input_embeds: torch.Tensor, - draft_embeds: torch.Tensor, - idx_mapping: torch.Tensor, - query_start_loc: torch.Tensor, - last_token_indices: torch.Tensor, - num_reqs: int, -) -> DispatchSpec: - hidden_size = input_embeds.shape[-1] - return (num_reqs, triton.cdiv(hidden_size, 256)), dict( - hidden_size=hidden_size, - BLOCK_SIZE_Q=16, - BLOCK_SIZE_H=256, - ) - - def update_draft_inputs( draft_tokens: torch.Tensor, draft_embeds: torch.Tensor | None, @@ -1146,21 +1074,29 @@ def update_draft_inputs( idx_mapping: torch.Tensor, num_reqs: int, ) -> None: - _shift_input_ids( + _shift_input_ids_kernel[(num_reqs,)]( input_buffers.input_ids, idx_mapping, input_buffers.query_start_loc, last_token_indices, draft_tokens, - num_reqs=num_reqs, + BLOCK_SIZE=1024, ) if input_embeds is not None: assert draft_embeds is not None - _shift_input_embeds( + hidden_size = input_embeds.shape[-1] + hidden_block_size = 256 + _shift_input_embeds_kernel[ + (num_reqs, triton.cdiv(hidden_size, hidden_block_size)) + ]( input_embeds, + input_embeds.stride(0), draft_embeds, + draft_embeds.stride(0), idx_mapping, input_buffers.query_start_loc, last_token_indices, - num_reqs=num_reqs, + hidden_size, + BLOCK_SIZE_Q=16, + BLOCK_SIZE_H=hidden_block_size, ) From 93e051e290d8a87ff7dab5c5ab2e5a3d4a53e760 Mon Sep 17 00:00:00 2001 From: Suraj Date: Sun, 13 Sep 2026 01:20:54 -0400 Subject: [PATCH 118/774] [Bugfix][Benchmark] Make streaming TTFT/E2E latency accounting consistent across endpoints (#55508) Signed-off-by: surajm20061998 Co-authored-by: Claude Opus 5 --- benchmarks/backend_request_func.py | 50 +- .../test_endpoint_request_func_timing.py | 464 ++++++++++++++++++ vllm/benchmarks/lib/endpoint_request_func.py | 14 +- 3 files changed, 505 insertions(+), 23 deletions(-) create mode 100644 tests/benchmarks/test_endpoint_request_func_timing.py diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py index 6349095ad729..bae0149f106f 100644 --- a/benchmarks/backend_request_func.py +++ b/benchmarks/backend_request_func.py @@ -105,7 +105,7 @@ async def async_request_tgi( timestamp = time.perf_counter() # First token if ttft == 0.0: - ttft = time.perf_counter() - st + ttft = timestamp - st output.ttft = ttft # Decoding phase @@ -330,8 +330,7 @@ async def async_request_openai_completions( # First token if not first_chunk_received: first_chunk_received = True - ttft = time.perf_counter() - st - output.ttft = ttft + output.ttft = timestamp - st # Decoding phase else: @@ -416,7 +415,7 @@ async def async_request_openai_chat_completions( output.prompt_len = request_func_input.prompt_len generated_text = "" - ttft = 0.0 + first_chunk_received = False st = time.perf_counter() most_recent_timestamp = st try: @@ -443,22 +442,30 @@ async def async_request_openai_chat_completions( if choices := data.get("choices"): content = choices[0]["delta"].get("content") # First token - if ttft == 0.0: - ttft = timestamp - st - output.ttft = ttft + if not first_chunk_received: + first_chunk_received = True + output.ttft = timestamp - st # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) generated_text += content or "" + # Only token chunks advance the request end; + # the trailing usage chunk carries no token. + most_recent_timestamp = timestamp elif usage := data.get("usage"): output.output_tokens = usage.get("completion_tokens") - most_recent_timestamp = timestamp - output.generated_text = generated_text - output.success = True + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!" + ) output.latency = most_recent_timestamp - st else: output.error = response.reason or "" @@ -530,7 +537,7 @@ def to_bytes(y, sr): output.prompt_len = request_func_input.prompt_len generated_text = "" - ttft = 0.0 + first_chunk_received = False st = time.perf_counter() most_recent_timestamp = st try: @@ -551,9 +558,9 @@ def to_bytes(y, sr): if choices := data.get("choices"): content = choices[0]["delta"].get("content") # First token - if ttft == 0.0: - ttft = timestamp - st - output.ttft = ttft + if not first_chunk_received: + first_chunk_received = True + output.ttft = timestamp - st # Decoding phase else: @@ -562,15 +569,24 @@ def to_bytes(y, sr): ) generated_text += content or "" + # Only token chunks advance the request + # end; the trailing usage chunk carries no + # token. + most_recent_timestamp = timestamp elif usage := data.get("usage"): output.output_tokens = usage.get( "completion_tokens" ) - most_recent_timestamp = timestamp - output.generated_text = generated_text - output.success = True + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!" + ) output.latency = most_recent_timestamp - st else: output.error = response.reason or "" diff --git a/tests/benchmarks/test_endpoint_request_func_timing.py b/tests/benchmarks/test_endpoint_request_func_timing.py new file mode 100644 index 000000000000..34394e3b014a --- /dev/null +++ b/tests/benchmarks/test_endpoint_request_func_timing.py @@ -0,0 +1,464 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Timing-accounting invariants for the streaming benchmark request functions. + +serve.py derives TPOT as (latency - ttft) / (output_len - 1), which only means +"average inter-token latency" if latency - ttft == sum(itl). These tests pin +that identity, and pin that the trailing choice-less usage chunk does not +extend the measured end of the request. + +Driven against a fake session and a scripted clock: no GPU, model, or socket. +""" + +import asyncio +import importlib.util +import json +from pathlib import Path +from typing import Any + +import aiohttp +import numpy as np +import pytest + +import vllm.benchmarks.lib.endpoint_request_func as request_func_module +from vllm.benchmarks.lib.endpoint_request_func import ( + RequestFuncInput, + async_request_openai_audio, + async_request_openai_chat_completions, + async_request_openai_completions, +) + +pytestmark = pytest.mark.skip_global_cleanup + +AUDIO_SAMPLE_RATE = 16_000 + + +class _ScriptedClock: + """A deterministic stand-in for time.perf_counter. + + Each read advances the clock by exactly one tick, so what a code path + observes depends only on how many times it reads the clock. A stray or + misattributed read then shows up as a whole tick of drift rather than as + sub-microsecond jitter, which keeps these assertions free of flakiness. + """ + + def __init__(self, tick: float = 1.0, start: float = 1000.0) -> None: + self.tick = tick + self._now = start + + def __call__(self) -> float: + value = self._now + self._now += self.tick + return value + + +class _FakeTime: + """Namespace substituted for the module-level time import.""" + + def __init__(self, perf_counter) -> None: + self.perf_counter = perf_counter + + +def _sse(payload: dict[str, Any]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _token_chunk(index: int, *, chat: bool) -> bytes: + if chat: + choice = {"index": 0, "delta": {"content": f"tok{index} "}} + else: + choice = {"index": 0, "text": f"tok{index} "} + return _sse({"id": "cmpl-test", "object": "x", "choices": [choice]}) + + +def _usage_chunk(n_tokens: int) -> bytes: + return _sse( + { + "id": "cmpl-test", + "object": "x", + "choices": [], + "usage": { + "prompt_tokens": 7, + "completion_tokens": n_tokens, + "total_tokens": 7 + n_tokens, + }, + } + ) + + +def _build_stream( + n_tokens: int, + *, + chat: bool, + ttft_delay: float = 0.0, + itl_delay: float = 0.0, + usage_delay: float = 0.0, +) -> list[tuple[float, bytes]]: + """Build ``(delay_before_chunk, payload)`` pairs for one streamed response. + + ``vllm bench serve`` always requests ``stream_options.include_usage``, so a + choice-less usage chunk always trails the final token. + """ + chunks = [ + (ttft_delay if i == 0 else itl_delay, _token_chunk(i, chat=chat)) + for i in range(n_tokens) + ] + chunks.append((usage_delay, _usage_chunk(n_tokens))) + chunks.append((0.0, b"data: [DONE]\n\n")) + return chunks + + +class _FakeContent: + def __init__(self, chunks: list[tuple[float, bytes]]) -> None: + self._chunks = chunks + + async def iter_any(self): + for delay, payload in self._chunks: + if delay: + await asyncio.sleep(delay) + yield payload + + # The legacy script iterates ``response.content`` directly. + def __aiter__(self): + return self.iter_any() + + +class _FakeResponse: + def __init__(self, chunks: list[tuple[float, bytes]]) -> None: + self.status = 200 + self.reason = "OK" + self.content = _FakeContent(chunks) + + async def __aenter__(self) -> "_FakeResponse": + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + +class _FakeSession: + """Returns a scripted SSE stream for any request, without touching a socket.""" + + def __init__(self, chunks: list[tuple[float, bytes]]) -> None: + self._chunks = chunks + + def post(self, **kwargs: Any) -> _FakeResponse: + del kwargs + return _FakeResponse(self._chunks) + + async def __aenter__(self) -> "_FakeSession": + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + +# (id, request function, api url, is_chat_payload) +STREAMING_ENDPOINTS = [ + ( + "completions", + async_request_openai_completions, + "http://test/v1/completions", + False, + ), + ( + "chat", + async_request_openai_chat_completions, + "http://test/v1/chat/completions", + True, + ), + ( + "audio", + async_request_openai_audio, + "http://test/v1/audio/transcriptions", + True, + ), +] + + +def _make_input(api_url: str, n_tokens: int, endpoint_id: str) -> RequestFuncInput: + multi_modal_content = None + if endpoint_id == "audio": + # A short silent clip; only its duration is read by the request func. + multi_modal_content = { + "audio": ( + np.zeros(AUDIO_SAMPLE_RATE // 10, dtype=np.float32), + AUDIO_SAMPLE_RATE, + ) + } + return RequestFuncInput( + prompt="hello", + api_url=api_url, + prompt_len=7, + output_len=n_tokens, + model="test-model", + multi_modal_content=multi_modal_content, + ) + + +def _run( + request_func, + api_url: str, + endpoint_id: str, + chunks, + n_tokens: int, + *, + expect_success: bool = True, +): + output = asyncio.run( + request_func( + _make_input(api_url, n_tokens, endpoint_id), + _FakeSession(chunks), + ) + ) + if expect_success: + assert output.success, output.error + return output + + +@pytest.mark.parametrize( + "endpoint_id,request_func,api_url,chat", + STREAMING_ENDPOINTS, + ids=[e[0] for e in STREAMING_ENDPOINTS], +) +def test_decode_span_identity_is_exact( + monkeypatch: pytest.MonkeyPatch, + endpoint_id: str, + request_func, + api_url: str, + chat: bool, +) -> None: + """latency - ttft must equal sum(itl) exactly. + + On a scripted clock this can only drift if a code path reads the clock an + extra time or attributes a read to the wrong event. + """ + tick = 1.0 + monkeypatch.setattr( + request_func_module, "time", _FakeTime(_ScriptedClock(tick=tick)) + ) + + n_tokens = 5 + chunks = _build_stream(n_tokens, chat=chat) + output = _run(request_func, api_url, endpoint_id, chunks, n_tokens) + + assert len(output.itl) == n_tokens - 1 + assert output.output_tokens == n_tokens + + residual = (output.latency - output.ttft) - sum(output.itl) + assert residual == pytest.approx(0.0, abs=tick / 1000), ( + f"{endpoint_id}: (latency - ttft) - sum(itl) = {residual!r}; " + f"TPOT is derived from latency - ttft, so it must match the measured " + f"inter-token latencies" + ) + + +@pytest.mark.parametrize( + "endpoint_id,request_func,api_url,chat", + STREAMING_ENDPOINTS, + ids=[e[0] for e in STREAMING_ENDPOINTS], +) +def test_trailing_usage_chunk_does_not_extend_latency( + monkeypatch: pytest.MonkeyPatch, + endpoint_id: str, + request_func, + api_url: str, + chat: bool, +) -> None: + """E2E latency must end at the last token, not at the usage trailer. + + The usage chunk carries no token, so counting it stretches E2EL (and + therefore TPOT) by the delivery time of one extra SSE frame. The [DONE] + sentinel is already excluded on the same grounds. + """ + tick = 1.0 + clock = _ScriptedClock(tick=tick) + monkeypatch.setattr(request_func_module, "time", _FakeTime(clock)) + + n_tokens = 5 + chunks = _build_stream(n_tokens, chat=chat) + output = _run(request_func, api_url, endpoint_id, chunks, n_tokens) + + # Clock reads: 1 for the start time, then 1 per non-[DONE] chunk. The last + # token is therefore read at start + n_tokens ticks; the usage chunk one + # tick later. + expected_latency = n_tokens * tick + assert output.latency == pytest.approx(expected_latency, abs=tick / 1000), ( + f"{endpoint_id}: latency {output.latency} should stop at the final " + f"token ({expected_latency}), not run on to the usage chunk " + f"({expected_latency + tick})" + ) + + +@pytest.mark.parametrize( + "endpoint_id,request_func,api_url,chat", + STREAMING_ENDPOINTS, + ids=[e[0] for e in STREAMING_ENDPOINTS], +) +def test_decode_span_identity_holds_with_real_clock( + endpoint_id: str, + request_func, + api_url: str, + chat: bool, +) -> None: + """The identity also holds against the real clock, with real gaps. + + The residual differences out the same recorded timestamps, so wall clock + jitter cancels and only float rounding remains; a microsecond bound is + ample rather than tight. + """ + n_tokens = 4 + chunks = _build_stream( + n_tokens, chat=chat, ttft_delay=0.02, itl_delay=0.01, usage_delay=0.02 + ) + output = _run(request_func, api_url, endpoint_id, chunks, n_tokens) + + residual = (output.latency - output.ttft) - sum(output.itl) + assert abs(residual) < 1e-6, ( + f"{endpoint_id}: (latency - ttft) - sum(itl) drifted by {residual:.9f}s" + ) + + +@pytest.mark.parametrize( + "endpoint_id,request_func,api_url,chat", + STREAMING_ENDPOINTS, + ids=[e[0] for e in STREAMING_ENDPOINTS], +) +def test_usage_only_stream_is_not_reported_as_success( + endpoint_id: str, + request_func, + api_url: str, + chat: bool, +) -> None: + """A stream with no token chunk must fail rather than report 0 ms. + + Only token chunks move the end of the request, so a degenerate stream + carrying nothing but a usage trailer leaves latency at zero. Reporting that + as a success would feed a zero-duration request into the benchmark + aggregates. + """ + chunks = [(0.0, _usage_chunk(0)), (0.0, b"data: [DONE]\n\n")] + output = _run(request_func, api_url, endpoint_id, chunks, 0, expect_success=False) + + assert not output.success + assert output.itl == [] + assert "TTFT" in output.error + + +# ``benchmarks/backend_request_func.py`` is a standalone script rather than a +# package module, so it is loaded by path. It also builds its own ClientSession, +# so the scripted stream is injected by swapping the module's aiohttp reference. +_LEGACY_PATH = Path(__file__).parents[2] / "benchmarks" / "backend_request_func.py" + + +class _FakeAiohttp: + """Stands in for the legacy module's aiohttp import.""" + + def __init__(self, chunks: list[tuple[float, bytes]]) -> None: + self._chunks = chunks + + def __getattr__(self, name: str) -> Any: + return getattr(aiohttp, name) + + def ClientSession(self, **kwargs: Any) -> _FakeSession: # noqa: N802 + del kwargs + return _FakeSession(self._chunks) + + +def _load_legacy(): + spec = importlib.util.spec_from_file_location( + "legacy_backend_request_func", _LEGACY_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +LEGACY_ENDPOINTS = [ + ( + "completions", + "async_request_openai_completions", + "http://test/v1/completions", + False, + ), + ( + "chat", + "async_request_openai_chat_completions", + "http://test/v1/chat/completions", + True, + ), + ( + "audio", + "async_request_openai_audio", + "http://test/v1/audio/transcriptions", + True, + ), +] + + +def _run_legacy( + monkeypatch, func_name, api_url, endpoint_id, chunks, n_tokens, clock=None +): + module = _load_legacy() + monkeypatch.setattr(module, "aiohttp", _FakeAiohttp(chunks)) + if clock is not None: + monkeypatch.setattr(module, "time", _FakeTime(clock)) + request_func = getattr(module, func_name) + return asyncio.run(request_func(_make_input(api_url, n_tokens, endpoint_id))) + + +@pytest.mark.parametrize( + "endpoint_id,func_name,api_url,chat", + LEGACY_ENDPOINTS, + ids=[e[0] for e in LEGACY_ENDPOINTS], +) +def test_legacy_decode_span_identity_is_exact( + monkeypatch: pytest.MonkeyPatch, + endpoint_id: str, + func_name: str, + api_url: str, + chat: bool, +) -> None: + """benchmarks/backend_request_func.py must honour the same identity.""" + tick = 1.0 + n_tokens = 5 + chunks = _build_stream(n_tokens, chat=chat) + output = _run_legacy( + monkeypatch, + func_name, + api_url, + endpoint_id, + chunks, + n_tokens, + clock=_ScriptedClock(tick=tick), + ) + + assert output.success, output.error + assert len(output.itl) == n_tokens - 1 + + residual = (output.latency - output.ttft) - sum(output.itl) + assert residual == pytest.approx(0.0, abs=tick / 1000), ( + f"legacy {endpoint_id}: (latency - ttft) - sum(itl) = {residual!r}" + ) + assert output.latency == pytest.approx(n_tokens * tick, abs=tick / 1000) + + +@pytest.mark.parametrize( + "endpoint_id,func_name,api_url,chat", + LEGACY_ENDPOINTS, + ids=[e[0] for e in LEGACY_ENDPOINTS], +) +def test_legacy_usage_only_stream_is_not_reported_as_success( + monkeypatch: pytest.MonkeyPatch, + endpoint_id: str, + func_name: str, + api_url: str, + chat: bool, +) -> None: + """A usage-only stream must fail rather than report a zero-duration success.""" + chunks = [(0.0, _usage_chunk(0)), (0.0, b"data: [DONE]\n\n")] + output = _run_legacy(monkeypatch, func_name, api_url, endpoint_id, chunks, 0) + + assert not output.success + assert output.itl == [] diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index f426e78c744c..b55c5f4e164d 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -235,8 +235,7 @@ async def async_request_openai_completions( # First token if not first_chunk_received: first_chunk_received = True - ttft = time.perf_counter() - st - output.ttft = ttft + output.ttft = timestamp - st # Decoding phase else: @@ -415,13 +414,14 @@ async def async_request_openai_chat_completions( output.itl.append(timestamp - most_recent_timestamp) generated_text += content or "" + # Only token chunks advance the request end; + # the trailing usage chunk carries no token. + most_recent_timestamp = timestamp elif usage := data.get("usage"): output.output_tokens = usage.get("completion_tokens") if (pt := usage.get("prompt_tokens")) is not None: output.prompt_len = pt - most_recent_timestamp = timestamp - output.generated_text = generated_text if first_chunk_received: output.success = True @@ -540,13 +540,15 @@ async def send_audio_file( ) generated_text += content or "" + # Only token chunks advance the request + # end; the trailing usage chunk carries no + # token. + most_recent_timestamp = timestamp elif usage := data.get("usage"): output.output_tokens = usage.get( "completion_tokens" ) - most_recent_timestamp = timestamp - output.generated_text = generated_text if first_chunk_received: output.success = True From 0cf266a469bd6f06e90fa267d466bbfe5b11dce5 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 13 Sep 2026 00:49:45 -0500 Subject: [PATCH 119/774] [Pooling] Support prompt embeddings in MRV2 decoder pooling (#56639) Signed-off-by: Taneem Ibrahim --- vllm/v1/worker/gpu/model_runner.py | 1 - vllm/v1/worker/gpu/pool/pooling_runner.py | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0ce0288f310a..fe423d6e38a6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1128,7 +1128,6 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: if self.pooling_runner is not None: assert new_req_data.pooling_params is not None - assert new_req_data.prompt_token_ids is not None self.pooling_runner.add_request( req_id, req_index, diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index 839edc8b8cab..ce947648514d 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/vllm/v1/worker/gpu/pool/pooling_runner.py @@ -64,7 +64,7 @@ def add_request( req_id: str, req_index: int, pooling_params: PoolingParams, - prompt_token_ids: list[int], + prompt_token_ids: list[int] | None, ) -> None: task = pooling_params.task if task not in self.supported_tasks: @@ -77,6 +77,7 @@ def add_request( self.pooling_params[req_index] = pooling_params self.pooling_states[req_index] = PoolingStates() if pooling_params.requires_token_ids: + assert prompt_token_ids is not None self.prompt_token_ids[req_index] = torch.tensor( prompt_token_ids, dtype=torch.int64 ) From 410f6da5c4bb62010728502035bee1b5f0eab2ac Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 01:56:14 -0400 Subject: [PATCH 120/774] [Bugfix][KV Offload] Submit CPU stores on no-forward steps (#56621) Signed-off-by: Lucas Wilkinson --- tests/v1/simple_kv_offload/test_worker.py | 103 +++++++++++++++++++++- vllm/v1/simple_kv_offload/worker.py | 11 ++- 2 files changed, 110 insertions(+), 4 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index a62dacc31028..173929005151 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -10,6 +10,7 @@ from __future__ import annotations import time +from contextlib import nullcontext from unittest.mock import MagicMock import pytest @@ -21,7 +22,13 @@ pytest.skip("Requires CUDA or ROCm", allow_module_level=True) from tests.v1.attention.utils import dense_kv_cache_tensor, dense_kv_cache_views +from tests.v1.kv_connector.unit.test_kv_connector_lifecycle import ( + _make_empty_scheduler_output, +) from vllm.config import CacheConfig +from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( + SimpleCPUOffloadConnector, +) from vllm.v1.core.kv_cache_utils import ( get_kv_cache_config_from_groups, is_kv_cache_spec_uniform, @@ -41,11 +48,12 @@ CU_MEMCPY_SRC_ACCESS_ORDER_ANY, CU_MEMCPY_SRC_ACCESS_ORDER_STREAM, build_params, - pin_tensor, ) from vllm.v1.simple_kv_offload.disk_backend import DiskBackend from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker +from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector +from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin from vllm.v1.worker.utils import allocate_kv_cache NUM_BLOCKS = 64 @@ -58,8 +66,11 @@ def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]: gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")} - cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")} - pin_tensor(cpu["k"]) + cpu = { + "k": torch.zeros( + (NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu", pin_memory=True + ) + } low_pri, _ = torch.cuda.Stream.priority_range() backend = DmaCopyBackend() backend.init( @@ -72,6 +83,92 @@ def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]: return backend, gpu["k"], cpu["k"] +@pytest.mark.parametrize("use_v2", [False, True]) +def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2): + """A finished request's final store must drain without another model step.""" + backend, gpu, cpu = _make_backend() + worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0) + worker._backend = backend + connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector) + connector.worker_handler = worker + output = _make_empty_scheduler_output() + output.kv_connector_metadata = SimpleCPUOffloadMetadata( + store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0] + ) + module = "vllm.v1.worker.kv_connector_model_runner_mixin" + monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector) + monkeypatch.setattr(f"{module}.set_forward_context", lambda *a: nullcontext()) + monkeypatch.setattr(f"{module}.get_forward_context", lambda: None) + try: + gpu.fill_(91) + if use_v2: + active = ActiveKVConnector.__new__(ActiveKVConnector) + active.kv_connector = connector + active._disabled = False + active._pending_load_start = False + module = "vllm.v1.worker.gpu.kv_connector" + monkeypatch.setattr(f"{module}.is_forward_context_available", lambda: True) + monkeypatch.setattr(f"{module}.get_forward_context", lambda: None) + result = active.no_forward(output) + else: + result = KVConnectorModelRunnerMixin.kv_connector_no_forward(output, None) + completion = ( + result.kv_connector_output.kv_connector_worker_meta + if result.kv_connector_output is not None + else None + ) + deadline = time.monotonic() + 5 + while completion is None and time.monotonic() < deadline: + connector.get_finished(set()) + completion = connector.build_connector_worker_meta() + time.sleep(0.001) + assert completion is not None, "CPU store never completed on the empty step" + assert completion.completed_store_events == {0: 1} + assert torch.equal(cpu[0], gpu[0].cpu()) + finally: + backend.shutdown() + + +def test_deferred_store_waits_for_draft_forward(monkeypatch): + """Target completion must not submit the store before draft finalization.""" + backend, gpu, cpu = _make_backend() + worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0) + worker._backend = backend + connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector) + connector.worker_handler = worker + output = _make_empty_scheduler_output() + output.kv_connector_metadata = SimpleCPUOffloadMetadata( + store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0] + ) + module = "vllm.v1.worker.kv_connector_model_runner_mixin" + monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector) + monkeypatch.setattr(f"{module}.has_kv_transfer_group", lambda: True) + monkeypatch.setattr(f"{module}.get_forward_context", lambda: None) + launch = MagicMock(wraps=backend.launch_copy) + monkeypatch.setattr(backend, "launch_copy", launch) + try: + with KVConnectorModelRunnerMixin._get_kv_connector_output( + output, defer_finalize=True + ): + gpu.fill_(17) + launch.assert_not_called() + gpu.fill_(91) + KVConnectorModelRunnerMixin.finalize_kv_connector() + assert launch.call_count == 1 + completion = None + deadline = time.monotonic() + 5 + while completion is None and time.monotonic() < deadline: + connector.get_finished(set()) + completion = connector.build_connector_worker_meta() + time.sleep(0.001) + assert completion is not None + assert completion.completed_store_events == {0: 1} + assert torch.equal(cpu[0], gpu[0].cpu()) + assert launch.call_count == 1 + finally: + backend.shutdown() + + def _drive_store( backend: DmaCopyBackend, gpu: torch.Tensor, diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index 9efc544f2f5a..fc950af40fa6 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -67,6 +67,7 @@ def __init__( # Metadata for the current step self._connector_metadata: SimpleCPUOffloadMetadata | None = None + self._store_submitted = False # Compute-done event recorded before each store; reused across steps # (get_finished runs once per step, copy queue is FIFO). @@ -244,12 +245,15 @@ def _init_cpu_mode( def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None: self._connector_metadata = metadata + self._store_submitted = False if metadata.load_event >= 0: self._pending_load_event_indices.add(metadata.load_event) if metadata.store_event >= 0: self._pending_store_event_indices.add(metadata.store_event) def clear_connector_metadata(self) -> None: + # No-forward steps skip the normal wait_for_save hook. + self.wait_for_save() self._connector_metadata = None def start_load_kv(self) -> None: @@ -278,7 +282,11 @@ def wait_for_save(self) -> None: #45704 for the bug and #39306 for the srcAccessOrder rationale. """ metadata = self._connector_metadata - if metadata is not None and metadata.store_gpu_blocks: + if ( + metadata is not None + and metadata.store_gpu_blocks + and not self._store_submitted + ): backend = self._backend assert backend is not None if self._store_compute_done is None: @@ -292,6 +300,7 @@ def wait_for_save(self) -> None: events_list=self._store_events, wait_event=self._store_compute_done, ) + self._store_submitted = True def get_finished( self, finished_req_ids: set[str] From 1cfd9728166c73747acdb423e7a49b23f27cdc44 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 03:36:06 -0400 Subject: [PATCH 121/774] [CI] Fix NIXL transfer-rank geometry fixture (#56672) Signed-off-by: Lucas Wilkinson --- tests/v1/kv_connector/unit/test_nixl_desc_geometry.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 607e4822b974..bfec51da6901 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -156,6 +156,8 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker = object.__new__(NixlConnectorWorker) worker.tp_rank = 0 worker.world_size = 1 + worker.transfer_tp_rank = 0 + worker.transfer_tp_size = 1 worker.block_size = 4 worker.engine_id = "local-engine" worker.use_mla = True From fa008bdccf10f2f31f84553112a227cf1b8947d7 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Sun, 13 Sep 2026 15:44:37 +0800 Subject: [PATCH 122/774] [Perf][Kernel] Integrate DeepSelect TopK for the DSA sparse indexer (#56464) Signed-off-by: zjy0516 Signed-off-by: Roger Wang Co-authored-by: Kimi Code Co-authored-by: Roger Wang --- CMakeLists.txt | 1 + cmake/external_projects/deepselect.cmake | 97 +++++ setup.py | 4 + tests/kernels/test_top_k_per_row.py | 323 +++++++++++++++++ vllm/config/kernel.py | 32 ++ vllm/engine/arg_utils.py | 17 +- vllm/model_executor/layers/indexer_topk.py | 339 ++++++++++++++++++ .../layers/sparse_attn_indexer.py | 68 +--- 8 files changed, 829 insertions(+), 52 deletions(-) create mode 100644 cmake/external_projects/deepselect.cmake create mode 100644 vllm/model_executor/layers/indexer_topk.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 018138f54cb4..fbacfd9091ce 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1572,6 +1572,7 @@ endif() # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) + include(cmake/external_projects/deepselect.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) include(cmake/external_projects/flashkda.cmake) diff --git a/cmake/external_projects/deepselect.cmake b/cmake/external_projects/deepselect.cmake new file mode 100644 index 000000000000..e2a5592e6167 --- /dev/null +++ b/cmake/external_projects/deepselect.cmake @@ -0,0 +1,97 @@ +include(FetchContent) + +# If DEEPSELECT_SRC_DIR is set, DeepSelect is built from that directory +# instead of downloading. +# It can be set as an environment variable or passed as a cmake argument. +# The environment variable takes precedence. +if (DEFINED ENV{DEEPSELECT_SRC_DIR}) + set(DEEPSELECT_SRC_DIR $ENV{DEEPSELECT_SRC_DIR}) +endif() + +if(DEEPSELECT_SRC_DIR) + FetchContent_Declare( + deepselect + SOURCE_DIR ${DEEPSELECT_SRC_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +else() + FetchContent_Declare( + deepselect + GIT_REPOSITORY https://github.com/vllm-project/DeepSelect.git + GIT_TAG d96d33afe1fab0d6066da49cdc91e64c2bee65ea + GIT_PROGRESS TRUE + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +endif() + +FetchContent_MakeAvailable(deepselect) +message(STATUS "DeepSelect is available at ${deepselect_SOURCE_DIR}") + +# DeepSelect kernels require CUDA 12.9+; 10.0f (family) covers SM100 and SM103. +set(DEEPSELECT_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND DEEPSELECT_SUPPORT_ARCHS "10.0f") +endif() + +cuda_archs_loose_intersection(DEEPSELECT_ARCHS "${DEEPSELECT_SUPPORT_ARCHS}" "${CUDA_ARCHS}") +if(DEEPSELECT_ARCHS) + message(STATUS "DeepSelect CUDA architectures: ${DEEPSELECT_ARCHS}") + + # All explicit template instantiations must be compiled: api.cpp + # references every dispatch combination. + file(GLOB DeepSelect_SOURCES CONFIGURE_DEPENDS + ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3/instantiations/*.cu + ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_fp32/instantiations/*.cu + ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_cluster/instantiations/*.cu + ) + list(PREPEND DeepSelect_SOURCES ${deepselect_SOURCE_DIR}/csrc/api.cpp) + + set(DeepSelect_INCLUDES + ${deepselect_SOURCE_DIR}/csrc + ${deepselect_SOURCE_DIR}/csrc/3rdparty/cutlass/include + ${deepselect_SOURCE_DIR}/csrc/3rdparty/kerutils/include + ) + + set_gencode_flags_for_srcs( + SRCS "${DeepSelect_SOURCES}" + CUDA_ARCHS "${DEEPSELECT_ARCHS}") + + set(VLLM_DEEPSELECT_GPU_FLAGS ${VLLM_GPU_FLAGS}) + list(APPEND VLLM_DEEPSELECT_GPU_FLAGS + "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda" + "--use_fast_math" "--ftz=false" + "-U__CUDA_NO_HALF_OPERATORS__" "-U__CUDA_NO_HALF_CONVERSIONS__" + "-U__CUDA_NO_HALF2_OPERATORS__" "-U__CUDA_NO_BFLOAT16_CONVERSIONS__") + + define_extension_target( + _deepselect_C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${DeepSelect_SOURCES} + COMPILE_FLAGS ${VLLM_DEEPSELECT_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${DeepSelect_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + # Only use C-shim APIs available in PyTorch 2.10. + # _deepselect_C is abi compatible with PyTorch >= TORCH_TARGET_VERSION. + target_compile_definitions(_deepselect_C PRIVATE + TORCH_TARGET_VERSION=0x020A000000000000ULL) + + # Needed to use cuda APIs from C-shim + if(VLLM_GPU_LANG STREQUAL "CUDA") + target_compile_definitions(_deepselect_C PRIVATE USE_CUDA) + endif() + + # DeepSelect requires C++20 + target_compile_options(_deepselect_C PRIVATE + $<$:-std=c++20> + $<$:-std=c++20>) +else() + message(STATUS "DeepSelect will not compile: unsupported CUDA architecture ${CUDA_ARCHS}") + # Create an empty target for setup.py on unsupported systems + add_custom_target(_deepselect_C) +endif() diff --git a/setup.py b/setup.py index 81ff499be4d5..13692f3d20ca 100644 --- a/setup.py +++ b/setup.py @@ -1015,6 +1015,7 @@ def extract_precompiled_and_patch_package( "vllm/_flashmla_extension_C.abi3.so", "vllm/_flashkda_C.abi3.so", "vllm/_sparse_flashmla_C.abi3.so", + "vllm/_deepselect_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", "vllm/cumem_allocator.abi3.so", @@ -1381,6 +1382,9 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append( CMakeExtension(name="vllm._flashmla_extension_C", optional=True) ) + # DeepSelect requires CUDA 12.9 or later (SM100a/SM103a only) + # Optional since it won't build on unsupported architectures + ext_modules.append(CMakeExtension(name="vllm._deepselect_C", optional=True)) if USE_PRECOMPILED_EXTENSIONS or ( CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0") ): diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index a1e27d2e50a3..c629ae15dea9 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import numpy as np import pytest import torch @@ -21,6 +22,13 @@ def _has_device_capability(major: int) -> bool: return current_platform.is_cuda() and current_platform.has_device_capability(major) +# DeepSelect is compiled for sm_100a/sm_103a only. +requires_sm100 = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="DeepSelect requires SM100a/SM103a", +) + + COOPERATIVE_TOPK_BACKEND = pytest.param( "cooperative_topk", marks=pytest.mark.skipif( @@ -1315,3 +1323,318 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: f"Row {i}: {backend} with padded stride doesn't match. " f"seq_len={sl}, stride={padded_stride}" ) + + +@requires_sm100 +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("batch_size", [1, 8, 64]) +@pytest.mark.parametrize("vocab_size", [4096, 131072]) +@pytest.mark.parametrize("top_k", [512, 2048]) +@torch.inference_mode() +def test_deep_select_topk( + dtype: torch.dtype, + batch_size: int, + vocab_size: int, + top_k: int, +) -> None: + """DeepSelect wrapper vs torch.topk with variable per-row ends.""" + from vllm.model_executor.layers import indexer_topk + + set_random_seed(0) + torch.set_default_device("cuda:0") + + row_ends = torch.randint( + 1, vocab_size + 1, (batch_size,), dtype=torch.int32, device="cuda" + ) + # Force one short row to exercise the -1 out-of-bounds fill. + row_ends[0] = min(top_k - 1, vocab_size) + logits = torch.randn(batch_size, vocab_size, dtype=dtype, device="cuda") + # Values beyond each row's end must never be read by the kernel: fill + # them with NaN (production tails are uninitialized; DeepSelect traps on + # NaN with abort_when_nan_found=True, so a wrong read fails loudly). + col_idx = torch.arange(vocab_size, device="cuda") + logits[col_idx[None, :] >= row_ends[:, None]] = float("nan") + + indices = indexer_topk.deep_select_topk(logits, top_k, end=row_ends) + torch.accelerator.synchronize() + + assert indices.shape == (batch_size, top_k) + row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") + for i in range(batch_size): + row_end = int(row_ends[i]) + k_i = min(top_k, row_end) + idx = logits[i, :row_end].topk(k_i, dim=-1)[1] + torch_indices[i, :k_i] = idx + assert torch.all(indices[i, k_i:] == -1), ( + f"Row {i}: expected -1 fill after {k_i} valid indices" + ) + # All selected indices must lie inside [0, row_end). + assert torch.all(indices[i, :k_i] < row_end) + + assert compare_top_k_results( + logits, indices, torch_indices, row_starts, row_ends, top_k + ), "DeepSelect topk results don't match torch.topk" + + +@requires_sm100 +@torch.inference_mode() +def test_deep_select_topk_preallocated_output() -> None: + """DeepSelect writes into a preallocated aligned buffer (no `end`).""" + from vllm.model_executor.layers import indexer_topk + + set_random_seed(0) + torch.set_default_device("cuda:0") + + batch_size = 8 + vocab_size = 131072 + top_k = 2048 + logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda") + + # Mimic vLLM's topk_indices_buffer: slice of a wider contiguous buffer, + # whose stride(0) is 32B-aligned. + buffer = torch.full((batch_size, 4096), -2, dtype=torch.int32, device="cuda") + output_idx = buffer[:, :top_k] + result = indexer_topk.deep_select_topk(logits, top_k, output_idx=output_idx) + torch.accelerator.synchronize() + + assert result.data_ptr() == output_idx.data_ptr() + assert torch.all(buffer[:, top_k:] == -2), "kernel wrote out of bounds" + + row_ends = torch.full((batch_size,), vocab_size, dtype=torch.int32, device="cuda") + row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") + for i in range(batch_size): + torch_indices[i] = logits[i].topk(top_k, dim=-1)[1] + + assert compare_top_k_results( + logits, output_idx, torch_indices, row_starts, row_ends, top_k + ), "DeepSelect topk (preallocated output) results don't match torch.topk" + + +def _has_flashinfer_topk() -> bool: + if not current_platform.is_cuda(): + return False + # find_spec only: importing flashinfer initializes CUDA at import time. + import importlib.util + + return importlib.util.find_spec("flashinfer") is not None + + +def _has_cooperative_topk() -> bool: + return _has_device_capability( + 90 + ) and not current_platform.is_device_capability_family(120) + + +SPARSE_INDEXER_EXPLICIT_BACKENDS = [ + pytest.param( + "per_row", + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="requires CUDA" + ), + ), + "torch", + pytest.param( + "persistent", + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="requires CUDA" + ), + ), + pytest.param( + "cooperative", + marks=pytest.mark.skipif( + not _has_cooperative_topk(), + reason="cooperative_topk requires SM90+ (non-SM12x)", + ), + ), + pytest.param( + "flashinfer", + marks=pytest.mark.skipif( + not _has_flashinfer_topk(), reason="requires flashinfer top-k" + ), + ), + pytest.param("deep_select", marks=requires_sm100), +] + + +@pytest.mark.parametrize("next_n", [1, 4]) +@pytest.mark.parametrize("backend", SPARSE_INDEXER_EXPLICIT_BACKENDS) +@torch.inference_mode() +def test_sparse_indexer_decode_topk_explicit_backends( + backend: str, next_n: int, workspace_init +) -> None: + """Every explicit sparse-indexer decode top-k backend must restrict each + row to its seq_len (dirty data past the end must never be selected) and + match torch.topk on the valid region.""" + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk + + set_random_seed(0) + torch.set_default_device("cuda:0") + + top_k = 2048 + batch_size = 4 + num_rows = batch_size * next_n + + # Per-row effective lens in the (B, next_n) "native spec decode" form; + # cooperative/persistent_topk require lengths.numel() == num_rows. + row_ends = torch.randint(4000, 16000, (num_rows,), dtype=torch.int32, device="cuda") + seq_lens = row_ends.reshape(batch_size, next_n) + max_seq_len = int(row_ends.max()) + # Pad the stride: multiple of 4 for cooperative_topk (TMA) and + # 1024B-aligned (256 fp32) for DeepSelect. + stride = (max_seq_len + 255) // 256 * 256 + + row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") + + logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda") + col_idx = torch.arange(stride, device="cuda") + logits[col_idx[None, :] >= row_ends[:, None]] = 1e30 # dirty tail + + indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda") + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, max_seq_len) + torch.accelerator.synchronize() + + # k_i == top_k for every row here (row_ends >= 3997 > top_k). + assert torch.all(indices >= 0) + torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + for i in range(num_rows): + row_end = int(row_ends[i]) + torch_indices[i] = logits[i, :row_end].topk(top_k, dim=-1)[1] + assert torch.all(indices[i] < row_end), ( + f"{backend}: row {i} selected indices past its seq_len" + ) + + assert compare_top_k_results( + logits, indices, torch_indices, row_starts, row_ends, top_k + ), f"{backend} results don't match torch.topk" + + +@pytest.mark.parametrize( + "backend", + [ + "torch", + pytest.param( + "per_row", + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="requires CUDA" + ), + ), + pytest.param( + "flashinfer", + marks=pytest.mark.skipif( + not _has_flashinfer_topk(), reason="requires flashinfer top-k" + ), + ), + pytest.param("deep_select", marks=requires_sm100), + ], +) +@torch.inference_mode() +def test_sparse_indexer_decode_topk_short_seq_lens( + backend: str, workspace_init +) -> None: + """seq_len < next_n: derived per-row ends must clamp at 0 (the reference + kernels do max(0, ...)); negative ends are OOB for the ragged kernels.""" + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk + + set_random_seed(0) + next_n = 4 + top_k = 512 + vocab = 8192 # 1024B-aligned rows for DeepSelect + # 1D (B,) seq_lens: rows of a request share its len and each kernel + # derives per-row ends (with clamping) itself. + seq_lens = torch.tensor([2, 6], dtype=torch.int32, device="cuda") + num_rows = 2 * next_n + logits = torch.randn(num_rows, vocab, dtype=torch.float32, device="cuda") + indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda") + + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, vocab) + torch.accelerator.synchronize() + + # Derived ends: [2,6] - next_n + 1 + arange(next_n) -> [-1,0,1,2] and + # [3,4,5,6], clamped to [0,0,1,2] and [3,4,5,6]. + for row, end in enumerate([0, 0, 1, 2, 3, 4, 5, 6]): + k_i = min(top_k, end) + assert torch.all(indices[row, k_i:] == -1), f"{backend}: row {row}" + if k_i: + ref = logits[row, :end].topk(k_i).indices.sort().values + got = indices[row, :k_i].sort().values + assert torch.equal(got, ref), f"{backend}: row {row}" + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@torch.inference_mode() +def test_sparse_indexer_topk_backend_resolution() -> None: + """Auto heuristic chain and fail-fast validation of explicit backends.""" + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk + + # 1024B-aligned stride: 131072 * 4 bytes. + logits = torch.randn(64, 131072, dtype=torch.float32, device="cuda") + # 4-divisible (cooperative TMA ok) but not 1024B-aligned (DeepSelect no): + # 131076 * 4 = 524304 bytes. + unaligned_logits = torch.randn(128, 131076, dtype=torch.float32, device="cuda")[ + :, :131072 + ] + + def resolve( + backend: str, + t: torch.Tensor = logits, + k: int = 2048, + num_rows: int = 64, + ) -> str: + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + return SparseIndexerTopk().resolve_backend(t, k, num_rows) + + is_sm100 = current_platform.is_device_capability_family(100) + has_coop = _has_cooperative_topk() + has_fi = _has_flashinfer_topk() + + # "auto" is exactly the pre-existing chain: cooperative -> persistent -> + # per_row. It must never select the opt-in backends by itself. + if has_coop: + assert resolve("auto") == "cooperative" + assert resolve("auto", k=512, num_rows=8) == "cooperative" + # Past cooperative's 64-row limit -> persistent, even where + # DeepSelect would be applicable. + assert resolve("auto", k=512, num_rows=128) == "persistent" + # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: + # still cooperative at <= 64 rows... + assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative" + # ...and persistent past the row limit. + assert resolve("auto", unaligned_logits, num_rows=128) == "persistent" + # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels) + # -> per_row fallback. + assert resolve("auto", k=5000) == "per_row" + + # Explicit values are returned as-is when constraints are met. + assert resolve("per_row") == "per_row" + assert resolve("torch") == "torch" + assert resolve("persistent") == "persistent" + if has_coop: + assert resolve("cooperative") == "cooperative" + if has_fi: + assert resolve("flashinfer") == "flashinfer" + if is_sm100: + # Explicit deep_select has no row-count requirement. + assert resolve("deep_select", num_rows=1) == "deep_select" + with pytest.raises(RuntimeError, match="constraints"): + resolve("deep_select", unaligned_logits) + + # Fail-fast on unmet constraints. + with pytest.raises(RuntimeError, match="num_rows must be <= 64"): + resolve("cooperative", num_rows=128) + with pytest.raises(RuntimeError, match="topk_tokens must be in"): + resolve("persistent", k=3000) + with pytest.raises(RuntimeError, match="topk_tokens must be in"): + resolve("cooperative", k=3000) + if not is_sm100: + with pytest.raises(RuntimeError, match="SM100a/SM103a"): + resolve("deep_select") diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index c16e4fc169c6..d44e83bf67c5 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -160,6 +160,16 @@ def with_default( # moe_ep variants. MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS +SparseIndexerTopkBackend = Literal[ + "auto", + "deep_select", + "cooperative", + "persistent", + "per_row", + "flashinfer", + "torch", +] + # Architectures whose model code wires up the flashinfer moe_ep experts. MTP # and DSpark draft variants inherit the setting from these target models. FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( @@ -269,6 +279,21 @@ class KernelConfig: running QDQ on activations. """ + sparse_indexer_topk_backend: SparseIndexerTopkBackend = "auto" + """Backend for the DSA sparse indexer decode top-k kernel. Available options: + + - "auto": The pre-existing chain (cooperative -> persistent -> per_row); + the other backends are opt-in + - "deep_select": Use DeepSelect kernels (SM100a/SM103a only) + - "cooperative": Use vLLM's cooperative_topk kernel + - "persistent": Use vLLM's persistent_topk kernel + - "per_row": Use vLLM's top_k_per_row_decode kernel + - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel + - "torch": Use a plain torch.topk implementation (debug reference) + + Explicit values raise RuntimeError when their constraints are not met. + """ + linear_backend: LinearBackend = "auto" """Backend for linear layer GEMM kernels. Available options: @@ -312,6 +337,13 @@ def _normalize_linear_backend(cls, value: Any) -> Any: return value.lower().replace("-", "_") return value + @field_validator("sparse_indexer_topk_backend", mode="before") + @classmethod + def _normalize_sparse_indexer_topk_backend(cls, value: Any) -> Any: + if isinstance(value, str): + return value.lower().replace("-", "_") + return value + def compute_hash(self) -> str: """ Produces a hash unique to the pass configuration. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ae6f14b1df41..b420e84f542b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -76,7 +76,12 @@ PrefixCachingHashAlgo, ) from vllm.config.device import Device -from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend +from vllm.config.kernel import ( + IrOpPriorityConfig, + LinearBackend, + MoEBackend, + SparseIndexerTopkBackend, +) from vllm.config.load import SafetensorsLoadStrategy from vllm.config.lora import MaxLoRARanks from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm @@ -505,6 +510,9 @@ class EngineArgs: enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter moe_backend: MoEBackend = KernelConfig.moe_backend linear_backend: LinearBackend = KernelConfig.linear_backend + sparse_indexer_topk_backend: SparseIndexerTopkBackend = ( + KernelConfig.sparse_indexer_topk_backend + ) all2all_backend: All2AllBackend = ParallelConfig.all2all_backend enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size @@ -1670,6 +1678,11 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: linear_backend_kwargs = kernel_kwargs["linear_backend"] linear_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--linear-backend", **linear_backend_kwargs) + sparse_indexer_topk_kwargs = kernel_kwargs["sparse_indexer_topk_backend"] + sparse_indexer_topk_kwargs["type"] = lambda s: s.lower().replace("-", "_") + kernel_group.add_argument( + "--sparse-indexer-topk-backend", **sparse_indexer_topk_kwargs + ) # vLLM arguments vllm_kwargs = get_kwargs(VllmConfig) @@ -2548,6 +2561,8 @@ def create_engine_config( kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": kernel_config.linear_backend = self.linear_backend + if self.sparse_indexer_topk_backend != "auto": + kernel_config.sparse_indexer_topk_backend = self.sparse_indexer_topk_backend # Transfer top-level ir_op_priority into KernelConfig.ir_op_priority for op_name, op_priority in asdict(self.ir_op_priority).items(): diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py new file mode 100644 index 000000000000..a3b9da4e1ead --- /dev/null +++ b/vllm/model_executor/layers/indexer_topk.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Top-k kernels for the DSA sparse attention indexer.""" + +import functools + +import torch + +from vllm import _custom_ops as ops +from vllm.config import get_current_vllm_config +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer +from vllm.v1.worker.workspace import current_workspace_manager + +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + +# --------------------------------------------------------------------------- +# DeepSelect (vllm._deepselect_C) +# --------------------------------------------------------------------------- + +# Matches the -1 fill convention used for topk_indices_buffer elsewhere. +IDX_OOB_FILL_VALUE = -1 + +try: + import vllm._deepselect_C # noqa: F401 (registers torch.ops.deep_select) +except ImportError as e: + from vllm.logger import init_logger + + init_logger(__name__).warning( + "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e + ) + + +@functools.lru_cache(maxsize=1) +def get_deep_select_stride_requirement() -> tuple[int, int]: + """Stride alignment requirement (input, output) in bytes.""" + return torch.ops.deep_select.get_alignment_requirement() + + +def is_deep_select_supported(input: torch.Tensor, topk: int) -> bool: + """Whether the kernel accepts this input (dtype/stride/topk constraints).""" + return ( + topk <= 4096 + and input.shape[1] < 2**23 + and input.dtype in (torch.float32, torch.bfloat16) + and input.stride(1) == 1 + and input.stride(0) + * input.element_size() + % get_deep_select_stride_requirement()[0] + == 0 + ) + + +def _get_empty_and_aligned_tensor( + dim0: int, dim1: int, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Tensor with shape (dim0, dim1) whose stride(0) is 32B-aligned.""" + output_stride_requirement = get_deep_select_stride_requirement()[1] // ( + dtype.itemsize + ) + assert output_stride_requirement > 0 + dim1_rounded = ( + (dim1 + output_stride_requirement - 1) // output_stride_requirement + ) * output_stride_requirement + return torch.empty((dim0, dim1_rounded), device=device, dtype=dtype)[:, :dim1] + + +def deep_select_topk( + input: torch.Tensor, + topk: int, + end: torch.Tensor | None = None, + output_idx: torch.Tensor | None = None, + indices_dtype: torch.dtype = torch.int32, +) -> torch.Tensor: + """Select the top-k indices per row of `input` with DeepSelect. + + Args: + input: (num_rows, vocab_size), bf16 or fp32. stride(1) must be 1 and + stride(0) must be 1024B-aligned. + topk: Number of elements to select per row; must be <= 4096. + end: Optional (num_rows,) int32 tensor with the exclusive right + boundary of each row. Rows with `end[i] < topk` get their + remaining indices filled with -1. + output_idx: Optional preallocated (num_rows, topk) output tensor whose + stride(0) is 32B-aligned (e.g. a slice of a wider buffer). + indices_dtype: Output dtype when `output_idx` is not provided. + + Returns: + The (num_rows, topk) indices tensor. + """ + assert input.dim() == 2 and input.stride(1) == 1 + assert ( + input.stride(0) * input.element_size() % get_deep_select_stride_requirement()[0] + == 0 + ) + + num_rows = input.shape[0] + if output_idx is None: + output_idx = _get_empty_and_aligned_tensor( + num_rows, topk, input.device, indices_dtype + ) + else: + assert output_idx.dtype == indices_dtype + assert output_idx.shape[0] >= num_rows and output_idx.shape[1] >= topk + + torch.ops.deep_select.topk( + input, + topk, + None, # begin is not supported + end, + False, # sorted_value + False, # sorted_index + None, # output_value + output_idx, + None, # output_idx_offset + IDX_OOB_FILL_VALUE, + float("-inf"), # value_oob_fill_value + False, # return_value + False, # abort_when_nan_found: dummy/capture passes can feed NaN + # logits from uninitialized KV cache; the reference topk kernels do + # not trap on NaN, and trapping would abort CUDA graph capture. + ) + return output_idx + + +# --------------------------------------------------------------------------- +# Backend selection +# --------------------------------------------------------------------------- + + +@functools.cache +def get_indexer_topk(backend: str) -> "SparseIndexerTopk": + return SparseIndexerTopk(backend) + + +class SparseIndexerTopk(torch.nn.Module): + """The sparse indexer's decode top-k stage. + + Selects among the available top-k kernels (see + kernel_config.sparse_indexer_topk_backend) and runs the chosen one. + """ + + def __init__(self, backend: str | None = None) -> None: + super().__init__() + if backend is None: + backend = ( + get_current_vllm_config().kernel_config.sparse_indexer_topk_backend + ) + self._backend = backend + self._is_cuda = current_platform.is_cuda() + self._has_deep_select = self._is_cuda and ( + current_platform.is_device_capability_family(100) + ) + self._has_flashinfer_topk = has_flashinfer() + self._cooperative_capable = self._is_cuda and ( + current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + + def resolve_backend( + self, logits: torch.Tensor, topk_tokens: int, num_rows: int + ) -> str: + """Resolve the decode top-k implementation from the configured + backend ("auto" = the pre-existing chain, or a validated explicit + value).""" + if self._backend == "auto": + return self._resolve_auto(logits, topk_tokens, num_rows) + + failures: list[str] = [] + if self._backend == "cooperative": + failures = self._cooperative_constraints(logits, topk_tokens, num_rows) + elif self._backend == "persistent": + if not self._is_cuda: + failures.append("requires a CUDA platform") + if topk_tokens not in (512, 1024, 2048): + failures.append( + f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" + ) + elif self._backend == "deep_select": + if not self._is_cuda: + failures.append("requires a CUDA platform") + elif not self._has_deep_select: + failures.append("requires SM100a/SM103a (10.x device family)") + elif not is_deep_select_supported(logits, topk_tokens): + failures.append( + f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," + f" stride={logits.stride()}, topk={topk_tokens}" + ) + elif self._backend == "flashinfer": + if not self._is_cuda: + failures.append("requires a CUDA platform") + if not self._has_flashinfer_topk: + failures.append( + "flashinfer.topk.top_k_ragged_transform is not importable" + ) + if logits.dtype != torch.float32 or logits.stride(1) != 1: + failures.append( + f"requires fp32 logits with stride(1) == 1, got " + f"dtype={logits.dtype}, stride={logits.stride()}" + ) + if failures: + raise RuntimeError( + f"sparse_indexer_topk_backend='{self._backend}' was requested, but: " + + "; ".join(failures) + ) + return self._backend + + def _resolve_auto( + self, logits: torch.Tensor, topk_tokens: int, num_rows: int + ) -> str: + """The pre-existing priority chain: cooperative -> persistent -> + per_row. deep_select/flashinfer/torch are opt-in only.""" + if not self._cooperative_constraints(logits, topk_tokens, num_rows): + return "cooperative" + if self._is_cuda and topk_tokens in (512, 1024, 2048): + return "persistent" + return "per_row" + + def _cooperative_constraints( + self, logits: torch.Tensor, topk_tokens: int, num_rows: int + ) -> list[str]: + """Unmet constraints of cooperative_topk (empty when applicable).""" + failures = [] + if not self._is_cuda: + failures.append("requires a CUDA platform") + if topk_tokens not in (512, 1024, 2048): + failures.append( + f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" + ) + if num_rows > 64: + failures.append(f"num_rows must be <= 64, got {num_rows}") + if logits.stride(0) % 4 != 0: + failures.append( + f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}" + ) + if self._is_cuda and not self._cooperative_capable: + failures.append("requires SM90+ and is not supported on the SM12x family") + return failures + + @staticmethod + def _row_ends(seq_lens: torch.Tensor, next_n: int, num_rows: int) -> torch.Tensor: + """Per-row exclusive end offsets (int32, (num_rows,)) for top-k + kernels that take ragged lengths (DeepSelect, FlashInfer, torch + reference). + + seq_lens is (B, next_n) per-row effective lens for native spec decode + and (B, 1) otherwise, in which case per-row lens are derived the same + way as the other decode top-k kernels. + """ + if seq_lens.numel() == num_rows: + row_ends = seq_lens.reshape(-1) + else: + next_n_offsets = torch.arange( + next_n, dtype=torch.int32, device=seq_lens.device + ) + row_ends = ( + (seq_lens.reshape(-1, 1) - next_n + 1 + next_n_offsets) + .clamp_(min=0) + .reshape(-1) + ) + assert row_ends.dtype == torch.int32 + return row_ends + + def forward( + self, + logits: torch.Tensor, + seq_lens: torch.Tensor, + next_n: int, + topk_indices: torch.Tensor, + topk_tokens: int, + max_seq_len: int, + ) -> None: + """Run the resolved decode top-k implementation, writing into + topk_indices (int32, -1 fill for rows shorter than topk_tokens).""" + backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) + if backend == "deep_select": + row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) + deep_select_topk(logits, topk_tokens, end=row_ends, output_idx=topk_indices) + elif backend == "cooperative": + (topk_workspace,) = current_workspace_manager().get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.cooperative_topk( + logits, + seq_lens, + topk_indices, + topk_workspace, + topk_tokens, + max_seq_len, + ) + elif backend == "persistent": + (topk_workspace,) = current_workspace_manager().get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.persistent_topk( + logits, + seq_lens, + topk_indices, + topk_workspace, + topk_tokens, + logits.shape[1], + ) + elif backend == "flashinfer": + # Deferred: importing flashinfer initializes CUDA at import time. + from flashinfer.topk import top_k_ragged_transform + + row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) + offsets = torch.zeros( + logits.shape[0], dtype=torch.int32, device=logits.device + ) + # top_k_ragged_transform selects within [0, row_ends[i]) per + # row and -1-fills past the row length. + indices = top_k_ragged_transform(logits, offsets, row_ends, topk_tokens) + topk_indices.copy_(indices) + elif backend == "torch": + # Debug reference: mask everything past each row's end, then topk. + row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) + cols = torch.arange(logits.shape[1], device=logits.device) + masked = logits.masked_fill( + cols.unsqueeze(0) >= row_ends.unsqueeze(1), float("-inf") + ) + indices = masked.topk(topk_tokens, dim=-1).indices + in_range = torch.arange(topk_tokens, device=logits.device).unsqueeze( + 0 + ) < row_ends.unsqueeze(1) + indices = torch.where(in_range, indices, -1) + topk_indices.copy_(indices) + else: + assert backend == "per_row", f"unknown topk backend: {backend}" + ops.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_tokens, + ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d663c54de367..32902901c72b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -20,6 +20,10 @@ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( select_candidate_blocks as _select_candidate_blocks, ) +from vllm.model_executor.layers.indexer_topk import ( + RADIX_TOPK_WORKSPACE_SIZE, + get_indexer_topk, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -46,8 +50,6 @@ logger = init_logger(__name__) -RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 - # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. MXFP4_BLOCK_SIZE = 32 @@ -343,6 +345,7 @@ def sparse_attn_indexer( candidate_blocks: torch.Tensor | None = None, candidate_block_size: int = 0, candidate_write: bool = False, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -750,56 +753,16 @@ def sparse_attn_indexer( ) topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - use_cooperative_topk = ( - current_platform.is_cuda() - and topk_tokens in (512, 1024, 2048) - and num_rows <= 64 - and logits.stride(0) % 4 == 0 # TMA 16-byte alignment - and current_platform.has_device_capability(90) - and not current_platform.is_device_capability_family(120) - ) - use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( - 512, - 1024, - 2048, + # The backend comes from the layer (config is only readable at model + # construction); dispatchers are cached per backend. + get_indexer_topk(topk_backend)( + logits, + seq_lens, + next_n, + topk_indices, + topk_tokens, + attn_metadata_narrowed.max_seq_len, ) - if use_cooperative_topk: - workspace_manager = current_workspace_manager() - (topk_workspace,) = workspace_manager.get_simultaneous( - ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) - torch.ops._C.cooperative_topk( - logits, - seq_lens, - topk_indices, - topk_workspace, - topk_tokens, - attn_metadata_narrowed.max_seq_len, - ) - elif use_persistent_topk: - workspace_manager = current_workspace_manager() - (topk_workspace,) = workspace_manager.get_simultaneous( - ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) - torch.ops._C.persistent_topk( - logits, - seq_lens, - topk_indices, - topk_workspace, - topk_tokens, - logits.shape[1], - ) - else: - ops.top_k_per_row_decode( - logits, - next_n, - seq_lens, - topk_indices, - num_rows, - logits.stride(0), - logits.stride(1), - topk_tokens, - ) if decode_metadata.global_seq_lens is not None: _merge_dcp_topk_global( @@ -851,6 +814,7 @@ def sparse_attn_indexer_fake( candidate_blocks: torch.Tensor | None = None, candidate_block_size: int = 0, candidate_write: bool = False, + topk_backend: str = "auto", ) -> torch.Tensor: return topk_indices_buffer @@ -917,6 +881,7 @@ def __init__( # than threading them through per-step metadata. vllm_config = get_current_vllm_config() parallel_config = vllm_config.parallel_config + self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend self._parallel_config = parallel_config self.dcp_world_size = parallel_config.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 @@ -1026,6 +991,7 @@ def forward_cuda( candidate_blocks=self.candidate_blocks, candidate_block_size=self.candidate_block_size, candidate_write=self.candidate_write, + topk_backend=self.topk_backend, ) def forward_xpu( From a987777755c83863b5d563c15630dd324033d53f Mon Sep 17 00:00:00 2001 From: Ylean <61880605+Woolgathererer@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:07:12 +0800 Subject: [PATCH 123/774] [Bugfix][Gemma4] Keep image kwargs out of video preprocessing (#56652) Signed-off-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Co-authored-by: Codex --- .../multimodal/processing/test_gemma4.py | 68 +++++++++++++++++++ vllm/model_executor/models/gemma4_mm.py | 3 + 2 files changed, 71 insertions(+) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index f30afe47dde2..a28cc50be5fb 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -3,6 +3,7 @@ from collections.abc import Mapping +import numpy as np import pytest import torch from PIL import Image as PILImage @@ -13,6 +14,7 @@ Gemma4ImagePixelInputs, ) from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.inputs import MultiModalFieldConfig from vllm.utils.mem_constants import GiB_bytes @@ -199,6 +201,72 @@ def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str): assert replacement == expected +@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) +@pytest.mark.parametrize("kwargs_on_init", [False, True]) +@pytest.mark.parametrize( + "image_kwargs", [{"rescale_factor": 1 / 127.5}, {"max_soft_tokens": 560}] +) +@pytest.mark.parametrize("video_uuid", [None, "same-video"]) +def test_video_cache_is_independent_of_image_kwargs( + model_id: str, + kwargs_on_init: bool, + image_kwargs: dict[str, object], + video_uuid: str | None, +): + """Image overrides must not change video frames or depend on cache warmth.""" + kwargs = {"images_kwargs": image_kwargs} + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"image": 1, "video": 1}, + mm_processor_cache_gb=1, + ) + cache = MultiModalProcessorOnlyCache(ctx.model_config) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=cache) + hf_processor = processor.info.get_hf_processor() + image = PILImage.new("RGB", (48, 48), color=(128, 128, 128)) + frames = np.stack([np.asarray(image)] * 2) + metadata = {"fps": 2.0, "frames_indices": [0, 1]} + mm_items = processor.info.parse_mm_data( + {"image": image, "video": [(frames, metadata)]} + ) + + def process(mm_kwargs): + return processor( + hf_processor.image_token + hf_processor.video_token, + mm_items, + mm_uuid_items={"video": [video_uuid]}, + hf_processor_mm_kwargs=mm_kwargs, + ) + + baseline = process({}) + if kwargs_on_init: + ctx = build_model_context( + model_id, + mm_processor_kwargs=kwargs, + limit_mm_per_prompt={"image": 1, "video": 1}, + mm_processor_cache_gb=1, + ) + cache = MultiModalProcessorOnlyCache(ctx.model_config) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=cache) + request_kwargs = {} if kwargs_on_init else kwargs + process(request_kwargs) + cached = process(request_kwargs) + cache.clear_cache() + fresh = process(request_kwargs) + + def pixels(result, modality, field): + return result["mm_kwargs"][modality][0][field].data + + video_pixels = pixels(fresh, "video", "pixel_values_videos") + assert torch.equal(pixels(cached, "video", "pixel_values_videos"), video_pixels) + assert torch.equal(pixels(baseline, "video", "pixel_values_videos"), video_pixels) + assert baseline["mm_hashes"]["video"] == fresh["mm_hashes"]["video"] + assert not torch.equal( + pixels(baseline, "image", "pixel_values"), + pixels(fresh, "image", "pixel_values"), + ) + + @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) def test_limit_mm_per_prompt( image_assets: ImageTestAssets, diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 326d179cc237..2dddc72b218d 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -629,6 +629,9 @@ def _apply_hf_processor_main( # Process frames as images with max_soft_tokens=70 video_mm_kwargs = dict(hf_processor_mm_kwargs) + # Override configured image options too: these inputs are video + # frames, whose cache keys exclude images_kwargs. + video_mm_kwargs["images_kwargs"] = {} video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS dummy_prompt = ("\t" + processor.image_token) * len(frames) From e7a3963339654aa18b713f9ff233f129af9c7c8b Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 13 Sep 2026 02:23:20 -0700 Subject: [PATCH 124/774] [Bugfix] Avoid repeated dummy initialization and random CPU Engram fills (#56682) Signed-off-by: Woosuk Kwon --- vllm/model_executor/model_loader/dummy_loader.py | 8 ++++++-- vllm/model_executor/model_loader/weight_utils.py | 4 ++++ vllm/models/deepseek_v4_1/common/engram.py | 4 ++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/model_loader/dummy_loader.py b/vllm/model_executor/model_loader/dummy_loader.py index 98f82a5b7197..0f263fe11709 100644 --- a/vllm/model_executor/model_loader/dummy_loader.py +++ b/vllm/model_executor/model_loader/dummy_loader.py @@ -14,7 +14,6 @@ from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo from vllm.model_executor.model_loader.reload.utils import get_layer_tensors from vllm.model_executor.model_loader.weight_utils import ( - initialize_dummy_weights, initialize_single_dummy_weight, ) @@ -41,7 +40,12 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: else: # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. - initialize_dummy_weights(layer, model_config) + # Initialize only local tensors: state_dict() would recursively + # initialize descendants that model.modules() also visits. + # Exclude nonpersistent buffers, matching state_dict(). + for name, tensor in get_layer_tensors(layer).items(): + if name not in layer._non_persistent_buffers_set: + initialize_single_dummy_weight(tensor) def _process_online_quant_layer( self, diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 3085df85c6f8..4cb8249e48f6 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1341,6 +1341,10 @@ def initialize_single_dummy_weight( if param.device.type == "meta": return # deferred to finalize_layerwise_processing (e.g. online quant) + if (dummy_weight_value := getattr(param, "dummy_weight_value", None)) is not None: + param.fill_(dummy_weight_value) + return + if not torch.is_floating_point(param): if current_platform.is_rocm(): # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v4_1/common/engram.py index f30c3a5606d1..4cfba842b04f 100644 --- a/vllm/models/deepseek_v4_1/common/engram.py +++ b/vllm/models/deepseek_v4_1/common/engram.py @@ -687,6 +687,10 @@ def __init__( }, ) if cpu_offload: + # Constant dummy values avoid randomizing huge CPU lookup tables. + set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) + # The ue8m0 encoding of scale 1.0 is exponent byte 127. + set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) logger.info( "Engram table offloaded to pinned host memory: %d rows x %d, " "%.2f GiB per rank", From d2d649e674c75425d2d6975c87eb89fd4d55fff8 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Sun, 13 Sep 2026 02:44:23 -0700 Subject: [PATCH 125/774] [DS V4.1][Engram] Support async prefetch for offloaded engram lookups and engram DP sharding (#56512) Signed-off-by: Juntian Liu Signed-off-by: Woosuk Kwon Co-authored-by: Roger Wang Co-authored-by: ywang96 <> Co-authored-by: Woosuk Kwon Co-authored-by: Codex --- .buildkite/test_areas/distributed.yaml | 4 + tests/distributed/test_engram_dp_shard.py | 776 +++++++++++++++++++++ tests/engine/test_arg_utils.py | 30 +- tests/kernels/test_engram.py | 240 ++++++- tests/test_config.py | 158 +++-- vllm/config/engram.py | 60 +- vllm/config/vllm.py | 31 +- vllm/distributed/parallel_state.py | 80 +++ vllm/envs.py | 9 +- vllm/models/deepseek_v4_1/common/engram.py | 205 +++--- vllm/models/deepseek_v4_1/nvidia/engram.py | 403 +++++++++++ vllm/models/deepseek_v4_1/nvidia/model.py | 18 +- 12 files changed, 1820 insertions(+), 194 deletions(-) create mode 100644 tests/distributed/test_engram_dp_shard.py create mode 100644 vllm/models/deepseek_v4_1/nvidia/engram.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index dfe900d827f3..070664ac7ac3 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -206,6 +206,9 @@ steps: - tests/v1/distributed - tests/v1/engine/test_engine_core_client.py - tests/distributed/test_utils + - tests/distributed/test_engram_dp_shard.py + - vllm/config/engram.py + - vllm/models/deepseek_v4_1/ commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 @@ -217,6 +220,7 @@ steps: - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py + - pytest -v -s distributed/test_engram_dp_shard.py - label: ":nvidia: (L4) Distributed Compile + Comm" key: distributed-compile-comm-4-gpus diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py new file mode 100644 index 000000000000..ebb02542898c --- /dev/null +++ b/tests/distributed/test_engram_dp_shard.py @@ -0,0 +1,776 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engram tables shared by co-located DP replicas. + +Four GPUs on one node split the hash heads four ways, whether as DEP +(TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns +the right rows if the ids of every replica are gathered, each rank looks up +its own heads, and each replica trades those tokens back for the heads its +peers own. + +With shared host storage, DP replicas instead map the same TP slice and +prefetch only their own tokens without DP gathers. +""" + +from itertools import product +from pathlib import Path +from tempfile import TemporaryDirectory +from threading import Event +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + +from vllm.config import ( + EngramConfig, + VllmConfig, + get_current_vllm_config_or_none, + set_current_vllm_config, +) +from vllm.config.load import LoadConfig +from vllm.config.parallel import ParallelConfig +from vllm.config.scheduler import SchedulerConfig +from vllm.distributed import cleanup_dist_env_and_memory, parallel_state +from vllm.distributed.parallel_state import ( + get_engram_dp_group, + get_engram_dp_size, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.forward_context import set_forward_context +from vllm.models.deepseek_v4_1.common.engram import EngramLayout +from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops +from vllm.models.deepseek_v4_1.nvidia.engram import ( + Engram, + ParallelEngramEmbedding, + engram_head_shard_rank, + gather_engram_hashes, +) +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.system_utils import update_environment_variables + +WORLD_SIZE = 4 +DIM, BLOCK = 64, 32 +# Prime-sized buckets, as the model's layout builds them. +HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131) + + +def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]: + """Even loads, then uneven ones with an idle replica to pad around.""" + return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size) + + +def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]: + """The unsharded table, identical on every rank.""" + torch.manual_seed(0) + weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn) + scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8) + return weight, scale_inv + + +def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor: + """Ids of each head, drawn from that head's own bucket range.""" + torch.manual_seed(seed) + offset = 0 + columns = [] + for size in head_sizes: + columns.append(torch.randint(offset, offset + size, (num_tokens, 1))) + offset += size + return torch.cat(columns, dim=1).to(torch.int32).cuda() + + +def _reference(weight, scale_inv, ids) -> torch.Tensor: + """Dequantized lookup against the whole table.""" + valid = ids >= 0 + safe_ids = ids.clamp_min(0).long() + values = torch.nn.functional.embedding(safe_ids, weight).float() + scales = torch.nn.functional.embedding(safe_ids, scale_inv) + scales = (scales.to(torch.int32) << 23).view(torch.float32) + values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1) + return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0) + + +@pytest.mark.parametrize( + "architecture,options,edp_ranks,etp_ranks", + [ + ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]), + ("Qwen4ExpForCausalLM", {}, None, [0, 1]), + ( + "Qwen4ExpForCausalLM", + {"embedding_across_dp": True}, + None, + [0, 1, 2, 3], + ), + ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]), + ], +) +def test_engram_group_creation_honors_model_and_options( + monkeypatch, architecture, options, edp_ranks, etp_ranks +): + config = SimpleNamespace( + model_config=SimpleNamespace(architecture=architecture, is_moe=False), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=2, + distributed_executor_backend="mp", + ), + engram_config=EngramConfig(**options), + ) + for name in ( + "_TP", + "_ETP", + "_ENGRAM_DP", + "_DCP", + "_PCP", + "_PP", + "_DP", + "_EP", + "_EPLB", + ): + monkeypatch.setattr(parallel_state, name, None) + monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0)) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1) + + def make_group(group_ranks, *args, **kwargs): + ranks = next(ranks for ranks in group_ranks if 0 in ranks) + return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks)) + + monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group) + with set_current_vllm_config(config): + initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl") + group = get_engram_dp_group() + assert (group.ranks if group is not None else None) == edp_ranks + assert parallel_state.get_etp_group().ranks == etp_ranks + + +def _worker(rank: int, tp_size: int, port: int) -> None: + torch.accelerator.set_device_index(rank) + update_environment_variables( + { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(WORLD_SIZE), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(port), + } + ) + dp_size = WORLD_SIZE // tp_size + dp_rank, tp_rank = divmod(rank, tp_size) + vllm_config = VllmConfig( + parallel_config=ParallelConfig( + tensor_parallel_size=tp_size, + data_parallel_size=dp_size, + data_parallel_rank=dp_rank, + ), + scheduler_config=SchedulerConfig( + max_model_len=8, + is_encoder_decoder=False, + max_num_batched_tokens=8, + max_num_seqs=8, + ), + ) + # Supply the resolved Engram settings without constructing a full model. + vllm_config.engram_config = EngramConfig() + vllm_config.model_config = SimpleNamespace( + architecture="DeepseekV41ForCausalLM", is_moe=True + ) + try: + init_distributed_environment() + with set_current_vllm_config(vllm_config): + initialize_model_parallel(tensor_model_parallel_size=tp_size) + assert get_engram_dp_size() == dp_size + assert get_engram_dp_group().ranks == [ + tp_rank + replica * tp_size for replica in range(dp_size) + ] + assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank + + for failure in ("create", "open", "mmap", "register", "pinned", None): + _check_shared_storage_lifetime_and_failures(failure) + + # Reuse the processes and communication groups across storage modes. + for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product( + [(False, False), (True, False), (True, True)], + [8, 7], + [False, True] if tp_size > 1 else [False], + ): + _check_table( + vllm_config, + cpu_offload, + dp_shared_memory, + n_heads, + sequence_parallel, + ) + finally: + cleanup_dist_env_and_memory() + + +def _check_shared_storage_lifetime_and_failures(failure): + """One peer's failure must roll back every mapping/registration, without hangs.""" + import gc + import weakref + + group = get_engram_dp_group() + runtime = torch.cuda.cudart() + make_file = engram_ops.tempfile.NamedTemporaryFile + map_file = engram_ops.mmap.mmap + paths, mappings, registrations, unregistrations = [], [], [], [] + failed_rank = 0 if failure == "create" else 1 + + def fail(): + raise OSError(f"injected {failure} failure") + + def temporary_file(*args, **kwargs): + if failure == "create": + fail() + file = make_file(*args, **kwargs) + paths.append(Path(file.name)) + return file + + def open_file(*args, **kwargs): + if failure == "open" and group.rank_in_group == failed_rank: + fail() + return open(*args, **kwargs) + + def mapping(*args, **kwargs): + if failure == "mmap" and group.rank_in_group == failed_rank: + fail() + result = map_file(*args, **kwargs) + mappings.append(weakref.ref(result)) + return result + + def register(pointer, size, flags): + if failure == "register" and group.rank_in_group == failed_rank: + return SimpleNamespace(value=1) + result = runtime.cudaHostRegister(pointer, size, flags) + assert result.value == 0 + registrations.append(pointer) + return result + + def unregister(pointer): + unregistrations.append(pointer) + result = runtime.cudaHostUnregister(pointer) + assert result.value == 0 + return result + + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + engram_ops, + "tempfile", + SimpleNamespace(NamedTemporaryFile=temporary_file), + ) + patch.setattr(engram_ops, "open", open_file, raising=False) + patch.setattr( + engram_ops, + "mmap", + SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED), + ) + patch.setattr( + torch.cuda, + "cudart", + lambda: SimpleNamespace( + cudaHostRegister=register, + cudaHostUnregister=unregister, + ), + ) + if failure == "pinned" and group.rank_in_group == failed_rank: + patch.setattr(torch.Tensor, "is_pinned", lambda self: False) + if failure is not None: + stage = "cudaHostRegister" if failure in ("register", "pinned") else failure + with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"): + engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) + assert unregistrations == registrations + else: + storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) + assert all(not path.exists() for path in paths) + storage_ref = weakref.ref(storage) + parameter = torch.nn.Parameter(storage.weight, requires_grad=False) + views = storage.get_views(parameter, storage.weight_scale_inv) + del storage + gc.collect() + assert storage_ref() is None + assert not unregistrations + del parameter + gc.collect() + assert not unregistrations # UVA views still retain the CPU allocation. + del views + gc.collect() + assert unregistrations == registrations + for ref in mappings: + mapped = ref() + assert mapped is None or mapped.closed + assert all(not path.exists() for path in paths) + + +def _check_table( + vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel +): + parallel = vllm_config.parallel_config + dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size + dp_rank = parallel.data_parallel_rank + tp_rank = engram_ops.get_tensor_model_parallel_rank() + multithread = ( + dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1) + ) + load_config = ( + LoadConfig( + load_format="safetensors", + use_tqdm_on_load=False, + model_loader_extra_config={ + "enable_multithread_load": True, + "num_threads": 2, + }, + ) + if multithread + else vllm_config.load_config + ) + with pytest.MonkeyPatch.context() as m: + head_sizes = HEAD_SIZES[:n_heads] + rows = sum(head_sizes) + config = SimpleNamespace( + hidden_size=DIM, + hc_mult=2, + rms_norm_eps=1e-6, + engram_layer_ids=[0], + engram_num_embeddings=[rows], + engram_max_ngram_size=2, + engram_n_heads=n_heads, + engram_head_dim=DIM, + engram_compressed_vocab_size=32, + engram_pad_token_id=0, + engram_vocab_size=HEAD_SIZES[0], + ) + with set_current_vllm_config(vllm_config): + weight, scale_inv = _full_table(rows) + layout = EngramLayout(config) + # Stub configuration inputs, not Engram's initialization or execution. + component_config = SimpleNamespace( + engram_config=EngramConfig( + cpu_offload=cpu_offload, + dp_shared_memory=dp_shared_memory, + ), + scheduler_config=vllm_config.scheduler_config, + load_config=load_config, + ) + with m.context() as init_patch, torch.device("cuda"): + init_patch.setattr( + engram_ops, "get_current_vllm_config", lambda: component_config + ) + engram = Engram( + config, + quant_config=None, + layout=layout, + layer_hash_index=0, + use_sequence_parallel=sequence_parallel, + prefix="model.layers.0.engram", + ) + layer = engram.embed_tokens + if multithread: + _load_shared_multithread(layer, weight, scale_inv, load_config) + else: + # Only the leader has valid checkpoint payload in shared mode. + loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None + loader_scale = ( + scale_inv if not dp_shared_memory or dp_rank == 0 else None + ) + layer.weight.weight_loader(layer.weight, loader_weight) + layer.weight_scale_inv.weight_loader( + layer.weight_scale_inv, loader_scale + ) + if dp_shared_memory: + assert layer.dp_size == 1 + assert layer.head_start == tp_rank * layer.part_n_hash_cols + assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned() + + def unexpected_collective(*args, **kwargs): + raise AssertionError("shared lookup must not use DP collectives") + + m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective) + + token_counts = _token_counts(dp_size) + + # Production forward runs after the construction-time config has exited. + assert get_current_vllm_config_or_none() is None + for batch_idx, counts in enumerate(token_counts): + num_tokens = counts[dp_rank] + # Ids differ per replica but must agree across its TP ranks. + ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank) + if batch_idx == 2 and dp_rank == 1: + ids.fill_(engram_ops.DEAD_ID) + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32), + ): + gathered = gather_engram_hashes( + ids, dp_shared_memory=layer.dp_shared_memory + ) + expected_tokens = ( + num_tokens if dp_shared_memory else max(counts) * dp_size + ) + assert gathered.shape[0] == expected_tokens + engram.prepare_embeddings(gathered) + looked_up = engram.embed(ids) + standalone = layer(ids) + expected = _reference(weight.cuda(), scale_inv.cuda(), ids) + torch.testing.assert_close(standalone, expected, atol=0, rtol=0) + if sequence_parallel: + # SP hands back only this rank's token shard, zero-padded. + chunk = -(-num_tokens // tp_size) + expected = torch.nn.functional.pad( + expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens) + ) + expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk] + torch.testing.assert_close(looked_up, expected, atol=0, rtol=0) + + if ( + cpu_offload + and n_heads == 7 + and not sequence_parallel + and (not dp_shared_memory or tp_size == 1) + ): + _check_dummy_hash_model_forward( + vllm_config, + engram, + head_sizes, + weight, + scale_inv, + dp_shared_memory, + ) + + replay = n_heads == 7 and ( + (dp_shared_memory and sequence_parallel == (tp_size > 1)) + or ( + cpu_offload + and not dp_shared_memory + and sequence_parallel == (tp_size > 1) + ) + ) + if replay: + captures = ("full", "breakable") + for capture in captures: + _check_prefetch_replay( + engram, + head_sizes, + weight, + scale_inv, + dp_shared_memory, + capture, + ) + if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1): + with m.context() as storage_patch: + storage_patch.setattr(layer.weight, "data", layer.weight.data.clone()) + with pytest.raises(RuntimeError, match="storage must not be replaced"): + layer(ids) + + +def _check_dummy_hash_model_forward( + vllm_config, engram, head_sizes, weight, scales, dp_shared_memory +): + """A metadata-less replica joins DP sharding but skips a shared lookup.""" + from vllm.models.deepseek_v4_1.common.engram import NgramHashState + from vllm.models.deepseek_v4_1.nvidia import model as model_ops + + dp_rank = vllm_config.parallel_config.data_parallel_rank + if dp_shared_memory and dp_rank != 1: + return + tokens = 4 + ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank) + state = NgramHashState.__new__(NgramHashState) + torch.nn.Module.__init__(state) + state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda") + state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda") + state.ensure_cache = lambda: True + # Hash arithmetic is covered separately; keep the real model's branch and + # NgramHashState.dummy_hashes, plus real preparation/lookup collectives. + state.forward = lambda *args: ids.unsqueeze(1) + + class Decoder(SimpleNamespace): + def __call__(self, hidden, positions, input_ids, *args): + hashes, keep = args[-2:] + if dp_shared_memory and dp_rank == 1: + assert hashes is None and keep is None + output = hidden.new_zeros( + tokens, len(head_sizes), DIM, dtype=torch.bfloat16 + ) + else: + assert hashes is not None and keep is not None + assert bool(keep.all()) == (dp_rank != 1) + if dp_rank == 1: + assert torch.all(hashes == engram_ops.DEAD_ID) + output = engram.embed(hashes[:, 0]) + return output, None, None, None, None + + model = SimpleNamespace( + use_mega_moe=False, + use_sequence_parallel=False, + engram_hash=state, + engram_swa_prefix="swa", + engram_dp_shared_memory=dp_shared_memory, + layers=[Decoder(engram=engram)], + start_layer=0, + end_layer=1, + aux_hidden_state_layers=(), + ) + metadata = ( + None + if dp_rank == 1 + else { + "swa": SimpleNamespace( + query_start_loc=None, slot_mapping=None, block_table=None + ) + } + ) + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + model_ops, + "get_pp_group", + lambda: SimpleNamespace( + is_first_rank=True, + is_last_rank=False, + ), + ) + patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden) + with set_forward_context( + metadata, + vllm_config, + num_tokens=tokens, + num_tokens_across_dp=torch.full( + (vllm_config.parallel_config.data_parallel_size,), + tokens, + dtype=torch.int32, + ), + ): + output = model_ops.DeepseekV4Model.forward( + model, + torch.arange(tokens, device="cuda"), + torch.arange(tokens, device="cuda"), + intermediate_tensors=None, + inputs_embeds=torch.zeros(tokens, DIM, device="cuda"), + lookback_token_ids=torch.full((1, 1), -1, device="cuda"), + )["hidden_states"] + expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids + expected = _reference(weight.cuda(), scales.cuda(), expected_ids) + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +def _load_shared_multithread(layer, weight, scale_inv, load_config): + """Opposite shard completion orders must leave every shared parameter ready.""" + from safetensors.torch import save_file + + from vllm.model_executor.model_loader import weight_utils + from vllm.model_executor.model_loader.default_loader import DefaultModelLoader + + tensors = { + "weight": weight, + "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu), + } + order = list(tensors) + if get_engram_dp_group().rank_in_group % 2: + order.reverse() + first_loaded = Event() + original_load = weight_utils.load_file + + def load_file(path, **kwargs): + if Path(path).stem != order[0]: + assert first_loaded.wait(timeout=30), "First shared weight load stalled" + return original_load(path, **kwargs) + + with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m: + for name, tensor in tensors.items(): + save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors")) + m.setattr(weight_utils, "load_file", load_file) + loader = DefaultModelLoader(load_config) + source = DefaultModelLoader.Source(directory, revision=None) + loaded = [] + for name, tensor in loader._get_weights_iterator(source): + param = getattr(layer, name) + param.weight_loader(param, tensor) + loaded.append(name) + first_loaded.set() + assert loaded == order + + # Check before any further collective could mask incomplete shared writes. + for name, tensor in tensors.items(): + param = getattr(layer, name) + expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0]) + torch.testing.assert_close( + param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0 + ) + + +def _check_prefetch_replay( + engram, head_sizes, weight, scale_inv, dp_shared_memory, capture +): + """DP-sharded and shared prefetches must replay with new IDs.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + ids = _make_ids(head_sizes, 8, seed=501) + tp_rank = engram_ops.get_tensor_model_parallel_rank() + tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8 + output = torch.empty( + tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16 + ) + + def step(): + gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory) + engram.prepare_embeddings(gathered) + output.copy_(engram.embed(ids)) + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + if capture == "breakable": + graph = BreakableCUDAGraphCapture() + with torch.cuda.stream(warmup), graph: + step() + assert graph.num_eager_breaks == 2 + else: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + for iteration in range(2): + ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration)) + graph.replay() + expected = _reference(weight.cuda(), scale_inv.cuda(), ids) + if engram.use_sequence_parallel: + expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens] + torch.testing.assert_close(output, expected, atol=0, rtol=0) + torch.accelerator.synchronize() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch): + """Check all storage modes with one process launch per TP/DP topology.""" + if torch.accelerator.device_count() < WORLD_SIZE: + pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.") + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") + mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE) + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize( + "cpu_offload,dp_size,error", + [ + (False, 4, "requires cpu_offload=True"), + (True, 1, "effective Engram DP size is 1"), + ], +) +def test_engram_dp_shared_memory_runtime_requirements( + monkeypatch, cpu_offload, dp_size, error +): + """Report unmet storage/topology requirements before allocating CUDA memory.""" + monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size) + with pytest.raises(ValueError, match=error): + ParallelEngramEmbedding( + sum(HEAD_SIZES), + DIM, + HEAD_SIZES, + cpu_offload=cpu_offload, + dp_shared_memory=True, + ) + + +@pytest.mark.parametrize( + "node_ids,dp_size,replica_size,expected", + [ + ([0] * 4, 1, 4, 1), + ([0] * 4, 4, 1, 4), + ([0] * 4, 2, 2, 2), + ([0] * 8 + [1] * 8, 8, 2, 4), + ([0] * 6 + [1] * 6, 4, 1, 2), + ([0] * 3 + [1] * 4, 7, 1, 1), + ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1), + ([0, 1] * 4, 4, 2, 1), + ([0] * 3 + [1] * 5, 4, 2, 1), + ], +) +def test_engram_dp_shard_size_respects_node_boundaries( + node_ids, dp_size, replica_size, expected, monkeypatch +): + """Sharing must fall back when replicas or node rank ranges do not align.""" + monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids))) + monkeypatch.setattr( + parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None) + ) + monkeypatch.setattr( + parallel_state, + "in_the_same_node_as", + lambda group, src: [node == node_ids[src] for node in node_ids], + ) + assert ( + parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size) + == expected + ) + + +@pytest.mark.parametrize("num_tokens", [0, 2, 4, 5]) +def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch): + """Idle/padded slots must not query row zero; oversized batches are rejected.""" + monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4) + monkeypatch.setattr( + engram_ops, + "get_engram_dp_group", + lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)), + ) + ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32) + if num_tokens > 4: + with pytest.raises(ValueError, match="exceeds the DP token slot"): + gather_engram_hashes(ids) + return + gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3) + for replica in gathered: + torch.testing.assert_close(replica[:num_tokens], ids) + assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID) + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("uniform", [False, True]) +def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch): + """Two four-GPU nodes pad only to their own maximum, including TP x DP.""" + group_size = 4 // tp_size + counts = ( + [4] * (8 // tp_size) + if uniform + else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2) + ) + monkeypatch.setattr( + engram_ops, + "get_forward_context", + lambda: SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts)) + ), + ) + for dp_rank in range(8 // tp_size): + slot = 4 if uniform else (4096 if dp_rank < group_size else 3) + ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32) + monkeypatch.setattr( + engram_ops, + "get_dp_group", + lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank), + ) + monkeypatch.setattr( + engram_ops, + "get_engram_dp_group", + lambda rank=dp_rank: SimpleNamespace( + rank_in_group=rank % group_size, + world_size=group_size, + all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim), + ), + ) + gathered = gather_engram_hashes(ids) + expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID) + expected[: ids.shape[0]] = ids + torch.testing.assert_close( + gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}" + ) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 0262b42d11f3..d6ce5b36b514 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -89,7 +89,10 @@ def test_watermark_config_cli(): @pytest.mark.parametrize( "options", [ - ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'], + [ + "--engram-config", + '{"cpu_offload": false, "embedding_across_dp": true}', + ], [ "--engram-config.cpu_offload", "false", @@ -98,9 +101,8 @@ def test_watermark_config_cli(): ], ], ) -def test_engram_config_cli(options, monkeypatch): - """CLI settings take precedence over the legacy offload environment.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1") +def test_engram_config_cli(options): + """JSON and dotted CLI options independently control Engram settings.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = EngineArgs.from_cli_args(parser.parse_args(options)) assert args.engram_config is not None @@ -109,14 +111,26 @@ def test_engram_config_cli(options, monkeypatch): @pytest.mark.parametrize( - "options, provided", - [([], False), (["--engram-config", "{}"], True)], + "options,provided,dp_shared_memory", + [ + ([], False, False), + (["--engram-config", "{}"], True, False), + ( + ["--engram-config", '{"dp_shared_memory": true}'], + True, + True, + ), + (["--engram-config.dp_shared_memory", "true"], True, True), + ], ) -def test_engram_config_cli_optional(options, provided): - """An explicit empty config must remain distinct from an omitted config.""" +def test_engram_config_cli_optional(options, provided, dp_shared_memory): + """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = EngineArgs.from_cli_args(parser.parse_args(options)) assert (args.engram_config is not None) == provided + if provided: + assert args.engram_config.cpu_offload is True + assert args.engram_config.dp_shared_memory is dp_shared_memory @pytest.mark.parametrize( diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index f7aa6dd03e79..c7e08ced82ae 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import bisect +import inspect from types import SimpleNamespace import pytest @@ -8,10 +9,13 @@ from vllm.models.deepseek_v4_1.common import engram as engram_ops from vllm.models.deepseek_v4_1.common.engram import ( - Engram, + Engram as CommonEngram, +) +from vllm.models.deepseek_v4_1.common.engram import ( NgramHashState, - ParallelEngramEmbedding, ) +from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops +from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding from vllm.platforms import current_platform @@ -107,7 +111,7 @@ def test_fused_engram_post_wkv_matches_reference( ) monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) - module = Engram.__new__(Engram) + module = CommonEngram.__new__(CommonEngram) torch.nn.Module.__init__(module) module.dim = dim module.hc_mult = hc_mult @@ -116,7 +120,7 @@ def test_fused_engram_post_wkv_matches_reference( module.use_sequence_parallel = use_sequence_parallel module.embed_tokens = torch.nn.Identity() # `forward` reads rows staged by `prepare_embeddings`, so inject kv there. - module.embed_tokens.tp_size = 1 + module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1 module.staged_rows = kv.unsqueeze(1) if use_sequence_parallel: padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size)) @@ -151,6 +155,25 @@ def _hash_state(use_slot_cache: bool) -> NgramHashState: return state +@pytest.mark.parametrize("num_tokens", [0, 7]) +def test_engram_dummy_hashes_leave_history_untouched(num_tokens): + """Dummy forwards have no valid table IDs and do not alter cached history.""" + state = _hash_state(use_slot_cache=True) + state.multipliers = torch.empty(2, 4, dtype=torch.int64) + state.primes = torch.empty(2, 3, 8, dtype=torch.int64) + state._cache = torch.arange(16, dtype=torch.int32) + history = state._cache.clone() + + hashes, keep = state.dummy_hashes(torch.arange(num_tokens)) + + assert hashes.shape == (num_tokens, 2, 24) + assert hashes.dtype == torch.int32 + assert torch.all(hashes == engram_ops.DEAD_ID) + assert keep.shape == (num_tokens,) and keep.dtype == torch.bool + assert not keep.any() + torch.testing.assert_close(state._cache, history) + + @pytest.mark.parametrize("num_blocks", [4, 100]) def test_engram_cache_follows_kv_cache_binding(num_blocks): """Discard profiling history when the real KV cache replaces its storage.""" @@ -533,9 +556,11 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding) torch.nn.Module.__init__(layer) layer.dim, layer.block_size, layer.tp_size = dim, block, 1 + layer.dp_size = 1 layer.n_hash_cols = layer.part_n_hash_cols = 24 layer.head_start = 0 layer.cpu_offload = cpu_offload + layer.dp_shared_memory = False layer.part_num_embeddings = rows # A window strictly inside the table, so unowned rows are exercised too. layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2 @@ -559,12 +584,27 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): return layer +@pytest.mark.parametrize( + "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)] +) +def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch): + """Reject empty owners before allocating weights or accessing CUDA.""" + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + with pytest.raises(AssertionError, match="ranks without hash heads"): + ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch): """Keep complete buckets and reconstruct head order, including TP padding.""" - head_sizes = (17, 19, 23, 29, 31, 37) + head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73) num_rows, dim = sum(head_sizes), 64 torch.manual_seed(0) weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn) @@ -623,7 +663,7 @@ def gather(local, dim): module.staged_rows = torch.empty_like(shards[0]) module.prepare_embeddings(ids) torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0) - # Slicing before head reordering must preserve padded heads and empty owners. + # Slicing before head reordering must preserve padded heads and tokens. module.use_sequence_parallel = True chunk = (len(ids) + tp_size - 1) // tp_size padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size)) @@ -636,6 +676,31 @@ def gather(local, dim): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_engram_lookup_reuses_jit_across_token_shapes(): + """Runtime token counts and launch grids must share one JIT variant.""" + layer = _make_embedding(cpu_offload=False) + kernel = engram_ops._engram_lookup_kernel + kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0] + kernel_cache.clear() + + cache_sizes = [] + for background in (False, True): + for num_tokens in (1, 7, 256): + ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda") + out = torch.empty( + num_tokens, + 24, + layer.dim, + dtype=torch.bfloat16, + device="cuda", + ) + layer.lookup(ids, out, background=background) + cache_sizes.append(len(kernel_cache)) + + assert cache_sizes == [1] * len(cache_sizes) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) @@ -660,19 +725,147 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") -@pytest.mark.parametrize("cpu_offload", [False, True]) +@pytest.mark.parametrize( + "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)] +) +def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload): + """Default offload uses pinned storage and a side stream; False uses HBM.""" + from vllm.config import EngramConfig + + config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6) + layout = SimpleNamespace( + num_embeddings=(72,), + head_dim=32, + max_ngram_size=2, + n_heads=24, + primes=(((3,) * 24,),), + ) + vllm_config = SimpleNamespace( + engram_config=EngramConfig() + if cpu_offload is None + else EngramConfig(cpu_offload=cpu_offload), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + ) + offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload + if backend == "common": + vllm_config.engram_config = None + monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr( + nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config + ) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity() + ) + with torch.device("cuda"): + cls = Engram if backend == "nvidia" else CommonEngram + module = cls(config, None, layout, 0, False, "engram") + layer = module.embed_tokens + if offloaded: + assert layer.weight.device.type == "cpu" and layer.weight.is_pinned() + else: + assert layer.weight.device.type == "cuda" + layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn)) + layer.weight_scale_inv.data.fill_(127) + streams = [] + lookup = layer.lookup + + def track_lookup(ids, out, background=False): + streams.append(torch.cuda.current_stream()) + lookup(ids, out, background=background) + + layer.lookup = track_lookup + main = torch.cuda.current_stream() + ids = torch.randint( + layer.vocab_start_idx, + layer.vocab_end_idx, + (8, 24), + device="cuda", + dtype=torch.int32, + ) + module.prepare_embeddings(ids) + expected = _reference_lookup( + layer.weight.cuda(), + layer.weight_scale_inv.cuda(), + ids, + layer.vocab_start_idx, + layer.vocab_end_idx, + ) + torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0) + if offloaded: + assert streams == [module._prefetch_stream] + assert streams[0] != main + else: + assert streams == [main] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "cpu_offload,delay", + [(False, None), (True, None), (True, "producer"), (True, "lookup")], +) @pytest.mark.parametrize("capture", ["eager", "full", "breakable"]) -def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture): - """Consume early lookup results after a break, with fresh IDs each replay.""" - _run_engram_prepared_rows(cpu_offload, capture) +def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay): + """Temporary lookup IDs survive allocator reuse and graph replay.""" + _run_engram_prepared_rows(cpu_offload, capture, delay=delay) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"] +) +def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency): + """Negative controls must expose incorrect rows when a dependency is removed.""" + if ( + missing_dependency == "record_stream" + and torch.cuda.memory.get_allocator_backend() != "native" + ): + pytest.skip( + "The allocator reuse negative control requires the native allocator" + ) + + def start(self, hash_ids, rows, stream): + if missing_dependency != "producer_wait": + stream.wait_stream(torch.cuda.current_stream()) + if missing_dependency != "record_stream": + hash_ids.record_stream(stream) + with torch.cuda.stream(stream): + self.embed_tokens.lookup(hash_ids, rows, background=True) + + def finish(self, stream): + pass + + monkeypatch.setattr(Engram, "_start_prefetch", start) + if missing_dependency == "consumer_wait": + monkeypatch.setattr(Engram, "_finish_prefetch", finish) + delay = "producer" if missing_dependency == "producer_wait" else "lookup" + try: + with pytest.raises( + AssertionError, match=r"Tensor-likes are not (?:equal|close)!" + ): + _run_engram_prepared_rows(True, "eager", delay=delay) + finally: + torch.accelerator.synchronize() def _run_engram_prepared_rows( - cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False + cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None ): - from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) layer = _make_embedding(cpu_offload) + if delay == "lookup": + lookup = layer.lookup + + def delayed_lookup(indices, out, background=False): + torch.cuda._sleep(2_000_000) + lookup(indices, out, background=background) + + layer.lookup = delayed_lookup cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64) layer.n_hash_cols = cols layer.tp_size = tp_size @@ -681,6 +874,17 @@ def _run_engram_prepared_rows( engram = Engram.__new__(Engram) torch.nn.Module.__init__(engram) engram.embed_tokens = layer + engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None + # Exercise the production eager boundaries even when the test process + # imported Engram before breakable graphs were enabled. + if capture == "breakable": + import vllm.envs as envs + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True) + for name in ("_start_prefetch", "_finish_prefetch"): + fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name))) + setattr(engram, name, fn.__get__(engram, Engram)) engram.use_sequence_parallel = use_sequence_parallel engram.staged_rows = torch.empty( num_tokens, @@ -709,7 +913,12 @@ def _run_engram_prepared_rows( embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True) def step(cap=None): - engram.prepare_embeddings(src) + if delay == "producer": + torch.cuda._sleep(2_000_000) + # Drop the last reference to a non-contiguous input after launching lookup. + engram.prepare_embeddings(hashes.clone()[:, 1]) + # Exercise same-size allocator reuse before consuming the prefetched rows. + torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1) if cap is not None: cap.add_eager(lambda: None) out.copy_(embed(src)) @@ -730,9 +939,12 @@ def step(cap=None): with torch.cuda.stream(warmup), graph: step(graph) torch.cuda.current_stream().wait_stream(warmup) - assert graph.num_graphs == 2 + assert graph.num_eager_breaks == (3 if cpu_offload else 1) for _ in range(3): + # Eager prefill can overwrite staging rows between decode replays. + hashes.random_(0, layer.part_num_embeddings) + step() hashes.random_(0, layer.part_num_embeddings) if graph is None: step() diff --git a/tests/test_config.py b/tests/test_config.py index 218ec6144562..c2629a75ab0a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1158,31 +1158,68 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in assert config.get_parallel_size(parallel) == expected -def test_engram_rejects_elastic_cross_dp(): +@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"]) +def test_engram_rejects_elastic_cross_dp(option): parallel = ParallelConfig( tensor_parallel_size=4, data_parallel_size=2, distributed_executor_backend="mp", ) parallel.enable_elastic_ep = True - with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"): - EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel) + with pytest.raises(ValueError, match=f"{option}.*elastic EP"): + EngramConfig(**{option: True}).verify_parallel_config(parallel) -@pytest.mark.parametrize("legacy", [None, "0", "1"]) -def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy): - """Explicit settings must override the legacy environment fallback.""" - monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) - if legacy is not None: - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) - assert EngramConfig().cpu_offload == (legacy == "1") - assert EngramConfig(cpu_offload=False).cpu_offload is False - assert EngramConfig(cpu_offload=True).cpu_offload is True +def test_engram_dp_shared_memory_requires_cpu_offload(): + with pytest.raises(ValueError, match="requires cpu_offload"): + EngramConfig(cpu_offload=False, dp_shared_memory=True) + + +@pytest.mark.parametrize( + "dp_size,load_format,multithread,error", + [ + (1, "auto", False, "requires data_parallel_size > 1"), + (2, "dummy", False, "requires load_format"), + (2, "sharded_state", False, "requires load_format"), + (2, "auto", False, None), + (2, "safetensors", True, None), + (2, "pt", True, None), + ], +) +def test_engram_dp_shared_memory_config_validation( + monkeypatch, dp_size, load_format, multithread, error +): + """Reject invalid shared configs before distributed init; allow threaded loads.""" + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + config = cast( + VllmConfig, + SimpleNamespace( + model_config=SimpleNamespace( + architecture="DeepseekV41ForCausalLM", + hf_text_config=SimpleNamespace(engram_layer_ids=[1]), + ), + speculative_config=None, + engram_config=EngramConfig(cpu_offload=True, dp_shared_memory=True), + parallel_config=ParallelConfig(data_parallel_size=dp_size), + load_config=LoadConfig( + load_format=load_format, + model_loader_extra_config={"enable_multithread_load": multithread}, + ), + ), + ) + if error: + with pytest.raises(ValueError, match=error): + VllmConfig._resolve_and_verify_engram_config(config) + else: + VllmConfig._resolve_and_verify_engram_config(config) @pytest.mark.parametrize( "architecture, ple_layers, cuda, supported", [ + ("DeepseekV41ForCausalLM", [1], True, True), + ("DeepseekV41ForCausalLM", [], True, False), + ("DeepseekV41ForCausalLM", [1], False, False), ("Qwen4ExpForCausalLM", [1], True, True), ("Qwen4ExpForConditionalGeneration", [1], True, True), ("Qwen4ExpForCausalLM", [], True, False), @@ -1201,7 +1238,9 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo ModelConfig, SimpleNamespace( architecture=architecture, - hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers), + hf_text_config=SimpleNamespace( + ple_layer_ids=ple_layers, engram_layer_ids=ple_layers + ), ), ) if architecture is not None @@ -1214,55 +1253,73 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo with pytest.raises(ValueError, match="requires a model with supported Engram"): config.verify_model_config(model) + resolved = cast( + VllmConfig, + SimpleNamespace( + model_config=model, + speculative_config=None, + parallel_config=ParallelConfig(), + load_config=LoadConfig(load_format="dummy"), + engram_config=None, + ), + ) + VllmConfig._resolve_and_verify_engram_config(resolved) + assert (resolved.engram_config is not None) == supported + if supported: + assert resolved.engram_config.cpu_offload is True -def test_engram_config_defaults_to_none(monkeypatch): + +@pytest.mark.parametrize( + ("value", "expected"), [(None, True), ("0", False), ("1", True)] +) +def test_engram_cpu_offload_environment_default(monkeypatch, value, expected): monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) + if value is not None: + monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", value) + assert EngramConfig().cpu_offload is expected + assert EngramConfig(cpu_offload=False).cpu_offload is False + assert EngramConfig(cpu_offload=True).cpu_offload is True + + +def test_engram_config_defaults_to_none(): config = VllmConfig() assert config.engram_config is None assert config.compute_hash() -@pytest.mark.parametrize("legacy", ["0", "1"]) -def test_engram_none_resolves_legacy_offload(monkeypatch, legacy): - """Legacy enablement materializes a config before model validation.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("enable_dbo,ubatch_size", [(True, 0), (False, 2)]) +def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size): + """Reject overlapping Engram staging even without an explicit EngramConfig.""" config = cast( VllmConfig, SimpleNamespace( model_config=SimpleNamespace( - architecture="Qwen4ExpForCausalLM", - hf_text_config=SimpleNamespace(ple_layer_ids=[1]), + architecture="DeepseekV41ForCausalLM", + hf_text_config=SimpleNamespace(engram_layer_ids=[1]), ), speculative_config=None, - engram_config=None, - parallel_config=ParallelConfig(), + engram_config=EngramConfig(cpu_offload=False) if explicit else None, + parallel_config=ParallelConfig( + enable_dbo=enable_dbo, ubatch_size=ubatch_size + ), ), ) - VllmConfig._resolve_and_verify_engram_config(config) - if legacy == "1": - assert config.engram_config is not None - assert config.engram_config.cpu_offload is True - else: - assert config.engram_config is None + with pytest.raises( + ValueError, match="Engram does not support DBO or microbatching" + ): + VllmConfig._resolve_and_verify_engram_config(config) -@pytest.mark.parametrize("legacy", ["0", "1"]) -def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy): +def test_engram_explicit_config_requires_supported_model(): """Explicit all-false settings still opt into model validation.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) with pytest.raises(ValueError, match="requires a model with supported Engram"): VllmConfig(engram_config=EngramConfig(cpu_offload=False)) -def test_engram_legacy_offload_requires_supported_model(monkeypatch): - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1") - with pytest.raises(ValueError, match="requires a model with supported Engram"): - VllmConfig() - - @pytest.mark.parametrize("target_has_ple", [False, True]) -def test_engram_draft_config_validates_target(monkeypatch, target_has_ple): +@pytest.mark.parametrize("explicit", [False, True]) +def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit): """MTP may inherit cross-DP sharding without having its own PLE layers.""" monkeypatch.setattr(current_platform, "is_cuda", lambda: True) target = SimpleNamespace( @@ -1277,24 +1334,29 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple): speculative_config=SimpleNamespace( draft_model_config=draft, target_model_config=target ), - engram_config=EngramConfig(embedding_across_dp=True), + engram_config=EngramConfig(embedding_across_dp=True) if explicit else None, parallel_config=ParallelConfig(), + load_config=LoadConfig(load_format="dummy"), ), ) - if target_has_ple: - VllmConfig._resolve_and_verify_engram_config(config) - else: + if explicit and not target_has_ple: with pytest.raises(ValueError, match="requires a model with supported Engram"): VllmConfig._resolve_and_verify_engram_config(config) + else: + VllmConfig._resolve_and_verify_engram_config(config) + assert (config.engram_config is not None) == target_has_ple + if target_has_ple: + assert config.engram_config.cpu_offload is True -def test_engram_hash_tracks_storage_and_sharding(): +def test_engram_hash_tracks_execution_options(): configs = [ - EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp) - for offload in (False, True) - for across_dp in (False, True) + EngramConfig(cpu_offload=True), + EngramConfig(cpu_offload=False), + EngramConfig(cpu_offload=True, embedding_across_dp=True), + EngramConfig(cpu_offload=True, dp_shared_memory=True), ] - assert len({config.compute_hash() for config in configs}) == 4 + assert len({config.compute_hash() for config in configs}) == len(configs) @pytest.mark.parametrize("port", [1, 29550, 65535]) diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 2aede1868d55..5abb2bd34a5e 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -3,18 +3,17 @@ from typing import TYPE_CHECKING -from pydantic import Field +from pydantic import Field, model_validator +from typing_extensions import Self import vllm.envs as envs from vllm.config.utils import config, get_hash_factors, hash_factors -from vllm.logger import init_logger if TYPE_CHECKING: + from vllm.config.load import LoadConfig from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig -logger = init_logger(__name__) - # Architecture -> the hf_text_config field naming its n-gram layers. A model is # only configurable here if it actually has such layers to store. _NGRAM_LAYER_FIELDS = { @@ -25,28 +24,44 @@ def _default_cpu_offload() -> bool: - """Honor the legacy environment variable only when the field is omitted.""" - if envs.is_set("VLLM_PLE_CPU_OFFLOAD"): - logger.warning_once( - "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a " - "future release. Use --engram-config.cpu_offload instead." - ) return envs.VLLM_PLE_CPU_OFFLOAD +def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: + """Whether the model carries n-gram embedding layers.""" + if model_config is None: + return False + field = _NGRAM_LAYER_FIELDS.get(model_config.architecture) + if field is None: + return False + return bool(getattr(model_config.hf_text_config, field, None)) + + @config class EngramConfig: """Configuration for Engram embedding storage and sharding.""" cpu_offload: bool = Field(default_factory=_default_cpu_offload) """Store embedding weights in pinned CPU memory for UVA lookup. - Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility. - An explicit value takes precedence over the legacy environment variable.""" + Defaults to VLLM_PLE_CPU_OFFLOAD, which is enabled by default. An explicit + value takes precedence over the environment variable.""" embedding_across_dp: bool = False """Shard embeddings across TP and all DP ranks when enabled. Otherwise, each DP rank has a separate TP-sharded embedding replica.""" + dp_shared_memory: bool = False + """Share CPU-offloaded embedding weights between co-located + DP replicas. Each node stores one copy of every TP shard, reducing host + memory without per-step Engram DP collectives. Requires sufficient + /dev/shm capacity and a shared IPC namespace.""" + + @model_validator(mode="after") + def _validate_shared_memory(self) -> Self: + if self.dp_shared_memory and not self.cpu_offload: + raise ValueError("dp_shared_memory requires cpu_offload=True") + return self + def verify_model_config(self, model_config: "ModelConfig | None") -> None: """Reject Engram configuration for models without n-gram embeddings.""" from vllm.platforms import current_platform @@ -64,13 +79,16 @@ def verify_model_config(self, model_config: "ModelConfig | None") -> None: ): raise ValueError( "EngramConfig requires a model with supported Engram " - "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek " - "V4.1 implementations with non-empty n-gram layer ids are " - "supported." + "embeddings, non-empty n-gram layer ids, and CUDA." ) def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: """Reject unsupported embedding parallel topologies.""" + if self.dp_shared_memory: + if parallel_config.data_parallel_size <= 1: + raise ValueError("dp_shared_memory requires data_parallel_size > 1.") + if parallel_config.enable_elastic_ep: + raise ValueError("dp_shared_memory is not supported with elastic EP.") if ( self.embedding_across_dp and parallel_config.data_parallel_size > 1 @@ -80,6 +98,18 @@ def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: "Engram embedding_across_dp is not supported with elastic EP yet." ) + def verify_load_config(self, load_config: "LoadConfig") -> None: + """Shared tables require a loader that invokes parameter weight callbacks.""" + if self.dp_shared_memory and load_config.load_format not in ( + "auto", + "safetensors", + "pt", + ): + raise ValueError( + "dp_shared_memory requires load_format 'auto', " + f"'safetensors' or 'pt'; got {load_config.load_format!r}." + ) + def get_parallel_size(self, parallel_config: "ParallelConfig") -> int: """Derive the embedding group size from the parallel configuration.""" size = parallel_config.tensor_parallel_size diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6fab6f0c7177..7f0a3f0d8723 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -34,7 +34,7 @@ from .diffusion import DiffusionConfig from .ec_manager_config import EncoderCacheManagerConfig from .ec_transfer import ECTransferConfig -from .engram import EngramConfig +from .engram import EngramConfig, model_has_engram_layers from .kernel import KernelConfig from .kv_events import KVEventsConfig from .kv_transfer import KVTransferConfig @@ -381,7 +381,7 @@ class VllmConfig: attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" engram_config: EngramConfig | None = None - """Optional Engram configuration, only valid for supported PLE models.""" + """N-gram embedding storage and sharding settings.""" mamba_config: MambaConfig = Field(default_factory=MambaConfig) """Mamba configuration.""" kernel_config: KernelConfig = Field(default_factory=KernelConfig) @@ -1218,22 +1218,37 @@ def _check_watermarking_unsupported( ) def _resolve_and_verify_engram_config(self) -> None: - """Resolve legacy offload settings and validate model and parallel configs.""" - if self.engram_config is None: - if not envs.VLLM_PLE_CPU_OFFLOAD: - return - self.engram_config = EngramConfig() + """Resolve defaults and validate n-gram embedding settings.""" + from vllm.platforms import current_platform + model_config = self.model_config speculative_config = self.speculative_config # Draft configs inherit the target's communication groups and settings. - # Qwen4Exp MTP itself disables PLE, so validate its target instead. + # Validate the target because the draft may disable n-gram embeddings. if ( speculative_config is not None and model_config is speculative_config.draft_model_config ): model_config = speculative_config.target_model_config + if ( + model_config is not None + and model_config.architecture == "DeepseekV41ForCausalLM" + and getattr(model_config.hf_text_config, "engram_layer_ids", None) + and self.parallel_config.use_ubatching + ): + raise ValueError( + "DeepSeek V4.1 Engram does not support DBO or microbatching. " + "Disable --enable-dbo and set --ubatch-size to 0." + ) + if self.engram_config is None: + if not current_platform.is_cuda() or not model_has_engram_layers( + model_config + ): + return + self.engram_config = EngramConfig() self.engram_config.verify_model_config(model_config) self.engram_config.verify_parallel_config(self.parallel_config) + self.engram_config.verify_load_config(self.load_config) logger.info_once("Resolved Engram configuration: %s", str(self.engram_config)) def __post_init__(self): diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 0466779c3477..dc3022a16fec 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -32,6 +32,7 @@ from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from datetime import timedelta +from math import gcd from multiprocessing import shared_memory from typing import TYPE_CHECKING, Any, Protocol from unittest.mock import patch @@ -1558,6 +1559,23 @@ def get_etp_group() -> GroupCoordinator: return _ETP +_ENGRAM_DP: GroupCoordinator | None = None + + +def get_engram_dp_group() -> GroupCoordinator | None: + """Return the DP replicas that share one engram embedding table. + + None when every replica holds a full (TP-sharded) copy, which is the case + for models without engram layers and for replicas that span nodes. + """ + return _ENGRAM_DP + + +def get_engram_dp_size() -> int: + """Number of DP replicas one engram embedding table is sharded over.""" + return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1 + + _DCP: GroupCoordinator | None = None @@ -1933,6 +1951,29 @@ def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]: return nullcontext() +def _engram_dp_shard_size( + world_size: int, + data_parallel_size: int, + replica_size: int, +) -> int: + """Size of DP shards aligned with complete replicas and node boundaries.""" + node_count = get_node_count() + replicas_per_node, remainder = divmod(world_size, node_count * replica_size) + if remainder: + return 1 + # Shards must divide both the DP group and each node's replica count. + shard_size = gcd(data_parallel_size, replicas_per_node) + if shard_size == 1 or node_count == 1: + return shard_size + ranks_per_node = replicas_per_node * replica_size + for start in range(0, world_size, ranks_per_node): + same_node = in_the_same_node_as(get_world_group().cpu_group, start) + node_ranks = [rank for rank, local in enumerate(same_node) if local] + if node_ranks != list(range(start, start + ranks_per_node)): + return 1 + return shard_size + + def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, @@ -2055,6 +2096,40 @@ def initialize_model_parallel( group_name="etp", ) + # Build the node-local DP group used to shard or share an Engram table. + global _ENGRAM_DP + assert _ENGRAM_DP is None, "engram data parallel group is already initialized" + engram_dp_size = 1 + engram_config = config.engram_config + if ( + engram_config is not None + and config.model_config is not None + and config.model_config.architecture == "DeepseekV41ForCausalLM" + and not enable_elastic_ep + ): + engram_dp_size = _engram_dp_shard_size( + world_size, + data_parallel_size, + tensor_model_parallel_size + * pipeline_model_parallel_size + * prefill_context_model_parallel_size, + ) + if engram_dp_size > 1: + group_ranks = ( + all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0) + ) + _ENGRAM_DP = init_model_parallel_group( + [x.tolist() for x in group_ranks], + get_world_group().local_rank, + backend, + group_name="edp", + ) + logger.info_once( + "Engram node-local DP group size: %d (TP=%d)", + engram_dp_size, + tensor_model_parallel_size, + ) + # Build the DCP model-parallel groups. global _DCP assert _DCP is None, "decode context model parallel group is already initialized" @@ -2304,6 +2379,11 @@ def destroy_model_parallel(): _TP.destroy() _TP = None + global _ENGRAM_DP + if _ENGRAM_DP: + _ENGRAM_DP.destroy() + _ENGRAM_DP = None + global _DCP if _DCP: _DCP.destroy() diff --git a/vllm/envs.py b/vllm/envs.py index aa32fc38bafb..0269a96409fe 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -154,7 +154,7 @@ VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False VLLM_ENABLE_V1_MULTIPROCESSING: bool = True VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 - VLLM_PLE_CPU_OFFLOAD: bool = False + VLLM_PLE_CPU_OFFLOAD: bool = True VLLM_DISABLE_COMPILE_CACHE: bool = False VLLM_REPLICATE_EMBED: bool = False VLLM_USE_LAYERNAME: bool = True @@ -2059,11 +2059,8 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_LOG_MODEL_INSPECTION": lambda: bool( int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0")) ), - # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their - # rows through UVA on a dedicated CUDA stream. - # Legacy fallback for EngramConfig.cpu_offload, which takes precedence. - # This environment variable may be removed in a future release. - "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))), + # Store n-gram embedding tables in pinned CPU memory for UVA lookup. + "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "1"))), # Debug logging for --enable-mfu-metrics "VLLM_DEBUG_MFU_METRICS": lambda: bool( int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0")) diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v4_1/common/engram.py index 4cfba842b04f..c415bbaa6e42 100644 --- a/vllm/models/deepseek_v4_1/common/engram.py +++ b/vllm/models/deepseek_v4_1/common/engram.py @@ -51,8 +51,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.utils import set_weight_attrs from vllm.triton_utils import tl, triton -from vllm.utils.platform_utils import is_uva_available -from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor logger = init_logger(__name__) @@ -236,6 +234,7 @@ def _write_hash_cache_kernel( tl.store(cache + slot, value, valid) +# Keep request shapes out of the cache key so one warmed variant covers runtime batches. @triton.jit( do_not_specialize=[ "num_tokens", @@ -462,6 +461,21 @@ def ensure_cache(self) -> bool: self._kv_cache_ref = weakref.ref(kv_cache) return True + def dummy_hashes( + self, input_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Participate in DP lookups without valid rows or hash-cache updates.""" + num_tokens = input_ids.shape[0] + num_layers, max_ngram = self.multipliers.shape + num_heads = self.primes.shape[-1] + hashes = input_ids.new_full( + (num_tokens, num_layers, (max_ngram - 1) * num_heads), + DEAD_ID, + dtype=torch.int32, + ) + keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device) + return hashes, keep + def forward( self, input_ids: torch.Tensor, @@ -561,7 +575,16 @@ def _engram_head_shard_weight_loader( param.data.copy_(shard) -@triton.jit +@triton.jit( + do_not_specialize=[ + "vocab_start", + "vocab_end", + "num_rows", + "ids_stride_t", + "ids_stride_h", + "GRID", + ] +) def _engram_lookup_kernel( weight, scales, @@ -578,7 +601,7 @@ def _engram_lookup_kernel( DIM: tl.constexpr, QUANT_BLOCK: tl.constexpr, BLOCK_R: tl.constexpr, - GRID: tl.constexpr, + GRID, ): """Gather fp8 rows, apply their ue8m0 block scales, write bf16. @@ -620,12 +643,9 @@ def _engram_lookup_kernel( class ParallelEngramEmbedding(nn.Module): - """The n-gram hash table, sharded by complete hash heads over TP ranks. - Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup. + """TP-sharded hash heads with FP8 rows and per-block E8M0 scales.""" - With `cpu_offload` the shard lives in pinned host memory and is read over - UVA instead of HBM; the TP sharding is unchanged either way. - """ + _weight_loader = staticmethod(_engram_head_shard_weight_loader) def __init__( self, @@ -633,89 +653,55 @@ def __init__( dim: int, head_sizes: tuple[int, ...], block_size: int = 32, - cpu_offload: bool = False, - ): + ) -> None: super().__init__() - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() assert head_sizes and all(size > 0 for size in head_sizes) assert sum(head_sizes) <= num_embeddings - if cpu_offload and not is_uva_available(): - raise RuntimeError("Engram CPU offload requires UVA support") self.num_embeddings = num_embeddings self.dim = dim self.block_size = block_size self.n_hash_cols = len(head_sizes) - self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size) - self.head_start = tp_rank * self.part_n_hash_cols + self.tp_size = get_tensor_model_parallel_world_size() + num_shards, head_rank = self._get_shard_info() + self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards) + # TODO: Support row-wise sharding when there are too few hash heads. + assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, ( + f"Engram sharding leaves ranks without hash heads: " + f"{self.n_hash_cols} heads over {num_shards} shards" + ) + self.head_start = head_rank * self.part_n_hash_cols head_end = self.head_start + self.part_n_hash_cols self.vocab_start_idx = sum(head_sizes[: self.head_start]) self.vocab_end_idx = sum(head_sizes[:head_end]) self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx - self.tp_size = tp_size - self.cpu_offload = cpu_offload - self._views: tuple[torch.Tensor, torch.Tensor] | None = None - self._view_src: tuple[int, int] | None = None self._num_sms = torch.cuda.get_device_properties( torch.accelerator.current_device_index() ).multi_processor_count - - # Explicit device: model init runs under a `torch.device("cuda")` - # context, which would otherwise put the shard in HBM. - kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {} - self.weight = nn.Parameter( - torch.empty( - self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs - ), - requires_grad=False, - ) - self.weight_scale_inv = nn.Parameter( - torch.empty( - self.part_num_embeddings, - dim // block_size, - dtype=torch.uint8, - **kwargs, - ), - requires_grad=False, - ) + weight, scales = self._allocate_weights() + self.weight = nn.Parameter(weight, requires_grad=False) + self.weight_scale_inv = nn.Parameter(scales, requires_grad=False) for param in (self.weight, self.weight_scale_inv): set_weight_attrs( param, { - "weight_loader": _engram_head_shard_weight_loader, + "weight_loader": self._weight_loader, "engram_vocab_start": self.vocab_start_idx, }, ) - if cpu_offload: - # Constant dummy values avoid randomizing huge CPU lookup tables. - set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) - # The ue8m0 encoding of scale 1.0 is exponent byte 127. - set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) - logger.info( - "Engram table offloaded to pinned host memory: %d rows x %d, " - "%.2f GiB per rank", - self.part_num_embeddings, - dim, - self.part_num_embeddings * (dim + dim // block_size) / 1024**3, - ) - def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: - """Parameters when resident, else cached UVA views of the pinned shard. + def _get_shard_info(self) -> tuple[int, int]: + return self.tp_size, get_tensor_model_parallel_rank() - Rebuilt if anything swaps `.data`, so a stale device pointer cannot - survive silently. - """ - if not self.cpu_offload: - return self.weight.data, self.weight_scale_inv.data - src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) - if self._view_src != src: - self._views = ( - get_accelerator_view_from_cpu_tensor(self.weight.data), - get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), - ) - self._view_src = src - assert self._views is not None - return self._views + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn), + torch.empty( + self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8 + ), + ) + + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.weight.data, self.weight_scale_inv.data def lookup( self, indices: torch.Tensor, out: torch.Tensor, background: bool = False @@ -753,7 +739,7 @@ def lookup( def forward(self, indices: torch.Tensor) -> torch.Tensor: """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] - bf16, gathered from all TP shards.""" + bf16, gathered from all shards for this replica's tokens.""" out = torch.empty( (indices.shape[0], self.part_n_hash_cols, self.dim), dtype=torch.bfloat16, @@ -762,8 +748,7 @@ def forward(self, indices: torch.Tensor) -> torch.Tensor: self.lookup(indices, out) if self.tp_size > 1: out = tensor_model_parallel_all_gather(out, dim=1) - out = out[:, : self.n_hash_cols] - return out + return out[:, : self.n_hash_cols] @triton.jit(do_not_specialize=["num_kv_tokens"]) @@ -859,7 +844,7 @@ def _fused_engram_post_wkv_kernel( @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"]) -def _engram_sp_rows_kernel( +def _engram_select_rows_kernel( gathered, output, num_tokens, @@ -880,6 +865,32 @@ def _engram_sp_rows_kernel( tl.store(output + offsets, values, offsets < num_elements) +def _engram_select_rows( + gathered: torch.Tensor, + output: torch.Tensor, + source_tokens: int, + token_start: int, + local_width: int, +) -> None: + """Copy one token window out of a rank-major gathered buffer. + + Both gathers land rank-major ([rank][token][local width]); this walks the + window the rank keeps and lays its ranks out side by side as width. + """ + if output.numel() == 0: + return + _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)]( + gathered, + output, + source_tokens, + token_start, + output.numel(), + local_width, + output.shape[1] * output.shape[2], + BLOCK_SIZE=1024, + ) + + class Engram(nn.Module): """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. @@ -910,13 +921,7 @@ def __init__( # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight`` # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight`` # suffix rule. - engram_config = get_current_vllm_config().engram_config - self.embed_tokens = ParallelEngramEmbedding( - layout.num_embeddings[layer_hash_index], - layout.head_dim, - tuple(size for order in layout.primes[layer_hash_index] for size in order), - cpu_offload=engram_config.cpu_offload if engram_config else True, - ) + self.embed_tokens = self._create_embedding(layout, layer_hash_index) n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads self.wkv = ReplicatedLinear( n_hash_cols * layout.head_dim, @@ -936,40 +941,54 @@ def __init__( ) max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens - # Keep lookup results alive across breakable graph segments. + self._init_staging(max_tokens, layout.head_dim) + + def _create_embedding( + self, layout: EngramLayout, layer_hash_index: int + ) -> ParallelEngramEmbedding: + return ParallelEngramEmbedding( + layout.num_embeddings[layer_hash_index], + layout.head_dim, + tuple(size for order in layout.primes[layer_hash_index] for size in order), + ) + + def _init_staging(self, max_tokens: int, head_dim: int) -> None: + # Persistent storage keeps lookup addresses stable across graph replays. self.staged_rows = torch.empty( max_tokens, self.embed_tokens.part_n_hash_cols, - layout.head_dim, + head_dim, dtype=torch.bfloat16, ) def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: - """Gather this layer's rows on the main stream before decoder layers.""" - self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]]) + """Look up rows before the decoder layers consume them.""" + rows = self.staged_rows[: hash_ids.shape[0]] + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small" + self.embed_tokens.lookup(hash_ids, rows) + + def _ready_rows(self, num_tokens: int) -> torch.Tensor: + return self.staged_rows[:num_tokens] def embed(self, hash_ids: torch.Tensor) -> torch.Tensor: """Gather heads, returning only local tokens when SP is enabled.""" - rows = self.staged_rows[: hash_ids.shape[0]] + rows = self._ready_rows(hash_ids.shape[0]) if self.embed_tokens.tp_size == 1: - return rows + return rows[:, : self.embed_tokens.n_hash_cols] if self.use_sequence_parallel: tp_size = self.embed_tokens.tp_size num_tokens, local_heads, dim = rows.shape gathered = tensor_model_parallel_all_gather(rows, dim=0) chunk = (num_tokens + tp_size - 1) // tp_size - rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) - _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)]( + out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) + _engram_select_rows( gathered, - rows, + out, num_tokens, get_tensor_model_parallel_rank() * chunk, - rows.numel(), local_heads * dim, - self.embed_tokens.n_hash_cols * dim, - BLOCK_SIZE=1024, ) - return rows + return out rows = tensor_model_parallel_all_gather(rows, dim=1) return rows[:, : self.embed_tokens.n_hash_cols] diff --git a/vllm/models/deepseek_v4_1/nvidia/engram.py b/vllm/models/deepseek_v4_1/nvidia/engram.py new file mode 100644 index 000000000000..b141b6c34975 --- /dev/null +++ b/vllm/models/deepseek_v4_1/nvidia/engram.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch.""" + +import mmap +import tempfile +import weakref +from contextlib import ExitStack + +import numpy as np +import torch + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_engram_dp_group, + get_engram_dp_size, + get_tensor_model_parallel_rank, + tensor_model_parallel_all_gather, +) +from vllm.distributed.parallel_state import GroupCoordinator +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v4_1.common.engram import ( + DEAD_ID, + EngramLayout, + _engram_head_shard_weight_loader, + _engram_select_rows, +) +from vllm.models.deepseek_v4_1.common.engram import ( + Engram as BaseEngram, +) +from vllm.models.deepseek_v4_1.common.engram import ( + ParallelEngramEmbedding as BaseParallelEngramEmbedding, +) +from vllm.utils.platform_utils import is_uva_available +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +logger = init_logger(__name__) + + +def engram_head_shard_rank() -> int: + """This rank's slot among the hash-head shards of one engram table. + + TP-major, so the shards a DP gather brings in are contiguous heads and + the following TP gather completes the head order. + """ + dp_group = get_engram_dp_group() + dp_size = dp_group.world_size if dp_group is not None else 1 + dp_rank = dp_group.rank_in_group if dp_group is not None else 0 + return get_tensor_model_parallel_rank() * dp_size + dp_rank + + +def engram_gathered_num_tokens() -> int: + """Per-replica token slot for the node-local Engram DP group.""" + dp_metadata = get_forward_context().dp_metadata + if dp_metadata is None: + raise RuntimeError("a DP-shared engram table needs DP token metadata") + group = get_engram_dp_group() + assert group is not None + # Engram groups are contiguous slices of the full DP group. + start = get_dp_group().rank_in_group - group.rank_in_group + return int( + dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max() + ) + + +def gather_engram_hashes( + hash_ids: torch.Tensor, *, dp_shared_memory: bool = False +) -> torch.Tensor: + """Collect the n-gram ids of every DP replica sharing one table. + + Replicas are padded to a common token slot, so the gathered shape is + static under CUDA graph capture (where DP already pads alike). + """ + dp_group = get_engram_dp_group() + if dp_group is None or dp_shared_memory: + return hash_ids + slot = engram_gathered_num_tokens() + if hash_ids.shape[0] > slot: + raise ValueError("Engram token count exceeds the DP token slot") + if hash_ids.shape[0] < slot: + pad = hash_ids.new_full( + (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID + ) + hash_ids = torch.cat((hash_ids, pad)) + return dp_group.all_gather(hash_ids, dim=0) + + +class DPSharedEngramStorage: + """Registered host weights shared by a node-local DP group with one writer.""" + + def __init__( + self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator + ) -> None: + self.group = group + weight_bytes = num_rows * dim + storage = self._allocate(weight_bytes + weight_bytes // block_size) + self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim) + self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size) + self._views: tuple[torch.Tensor, torch.Tensor] | None = None + + def _allocate(self, num_bytes: int) -> torch.Tensor: + """Map and register one physical allocation across a node-local DP group.""" + from vllm.distributed.device_communicators.shm_broadcast import ( + check_shm_free_space, + ) + + group = self.group + with ExitStack() as stack: + path, error = None, None + if group.rank_in_group == 0: + try: + check_shm_free_space(num_bytes) + backing_file = stack.enter_context( + tempfile.NamedTemporaryFile( + prefix="vllm_engram_", dir="/dev/shm" + ) + ) + backing_file.truncate(num_bytes) + path = backing_file.name + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + path, error = group.broadcast_object((path, error)) + if error is not None: + raise RuntimeError( + "Engram shared-memory creation failed on EDP rank 0: " + error + ) + + mapping = owner = tensor = finalizer = None + stage = "open" + try: + try: + with open(path, "r+b") as file: + stage = "mmap" + mapping = mmap.mmap( + file.fileno(), num_bytes, flags=mmap.MAP_SHARED + ) + stage = "cudaHostRegister" + owner = np.frombuffer(mapping, dtype=np.uint8) + pointer = owner.ctypes.data + tensor = torch.from_numpy(owner) + result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0) + if result.value != 0: + raise RuntimeError(f"cudaHostRegister failed: {result}") + finalizer = weakref.finalize( + owner, self._unregister, mapping, pointer + ) + finalizer.atexit = False # type: ignore[misc] + # The UVA helper otherwise allocates a private pinned copy. + if not tensor.is_pinned(): + raise RuntimeError( + "CUDA did not recognize the shared Engram registration" + ) + except Exception as exc: + error = f"{stage}: {type(exc).__name__}: {exc}" + + errors: list[str | None] = [None] * group.world_size + # Also fences peer mappings before the leader unlinks the file. + torch.distributed.all_gather_object( + errors, error, group=group.cpu_group + ) + failures = "; ".join( + f"EDP rank {rank}: {error}" + for rank, error in enumerate(errors) + if error is not None + ) + if failures: + raise RuntimeError( + "Engram shared-memory initialization failed: " + failures + ) + assert tensor is not None + return tensor + except Exception: + if finalizer is not None: + finalizer() + tensor = owner = None + if mapping is not None: + mapping.close() + raise + + @staticmethod + def _unregister(mapping: mmap.mmap, pointer: int) -> None: + # Torch storage retains the numpy owner, including through cached UVA views. + # Keep its mmap alive until CUDA has released the registration. + result = torch.cuda.cudart().cudaHostUnregister(pointer) + if result.value != 0: + logger.warning("Engram cudaHostUnregister failed: %s", result) + + def load_weight( + self, param: torch.nn.Parameter, loaded_weight: torch.Tensor + ) -> None: + if self.group.rank_in_group == 0: + _engram_head_shard_weight_loader(param, loaded_weight) + # Read order may differ across ranks. Equal load counts ensure all shared + # weights are ready after the last weight-loader call returns. + torch.distributed.barrier(group=self.group.cpu_group) + + def get_views( + self, weight: torch.Tensor, scales: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if (weight.data_ptr(), scales.data_ptr()) != ( + self.weight.data_ptr(), + self.weight_scale_inv.data_ptr(), + ): + raise RuntimeError("Shared Engram parameter storage must not be replaced") + if self._views is None: + self._views = ( + get_accelerator_view_from_cpu_tensor(self.weight), + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv), + ) + return self._views + + +class ParallelEngramEmbedding(BaseParallelEngramEmbedding): + """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices.""" + + _shared_memory: DPSharedEngramStorage | None = None + + def __init__( + self, + num_embeddings: int, + dim: int, + head_sizes: tuple[int, ...], + block_size: int = 32, + cpu_offload: bool = False, + dp_shared_memory: bool = False, + ) -> None: + self.cpu_offload = cpu_offload + self.dp_shared_memory = dp_shared_memory + self.dp_size = get_engram_dp_size() + if dp_shared_memory: + if not cpu_offload: + raise ValueError("dp_shared_memory requires cpu_offload=True") + if self.dp_size <= 1: + raise ValueError( + "dp_shared_memory requires a node-local Engram DP " + f"group with size > 1; effective Engram DP size is {self.dp_size}. " + "Check that the node layout and rank placement allow complete " + "DP replicas to be co-located." + ) + self.dp_size = 1 + if cpu_offload and not is_uva_available(): + raise RuntimeError("Engram CPU offload requires UVA support") + self._views: tuple[torch.Tensor, torch.Tensor] | None = None + self._view_src: tuple[int, int] | None = None + super().__init__(num_embeddings, dim, head_sizes, block_size) + if cpu_offload: + # Constant dummy values avoid randomizing huge CPU lookup tables. + set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) + # The ue8m0 encoding of scale 1.0 is exponent byte 127. + set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) + logger.info( + "Engram table offloaded to pinned host memory: %d rows x %d, " + "%.2f GiB %s", + self.part_num_embeddings, + dim, + self.part_num_embeddings * (dim + dim // block_size) / 1024**3, + "shared across DP replicas" if dp_shared_memory else "per rank", + ) + + def _get_shard_info(self) -> tuple[int, int]: + if self.dp_size == 1: + return super()._get_shard_info() + return self.tp_size * self.dp_size, engram_head_shard_rank() + + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.dp_shared_memory: + group = get_engram_dp_group() + assert group is not None + storage = DPSharedEngramStorage( + self.part_num_embeddings, self.dim, self.block_size, group + ) + self._shared_memory = storage + self._weight_loader = storage.load_weight + return storage.weight, storage.weight_scale_inv + if not self.cpu_offload: + return super()._allocate_weights() + # Model initialization may be inside a CUDA device context. + return ( + torch.empty( + self.part_num_embeddings, + self.dim, + dtype=torch.float8_e4m3fn, + device="cpu", + pin_memory=True, + ), + torch.empty( + self.part_num_embeddings, + self.dim // self.block_size, + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ), + ) + + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: + if self._shared_memory is not None: + return self._shared_memory.get_views(self.weight, self.weight_scale_inv) + if not self.cpu_offload: + return super()._storage() + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) + if self._view_src != src: + self._views = ( + get_accelerator_view_from_cpu_tensor(self.weight.data), + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), + ) + self._view_src = src + assert self._views is not None + return self._views + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + if self.dp_size == 1: + return super().forward(indices) + num_tokens = indices.shape[0] + indices = gather_engram_hashes(indices) + out = torch.empty( + (indices.shape[0], self.part_n_hash_cols, self.dim), + dtype=torch.bfloat16, + device=indices.device, + ) + self.lookup(indices, out) + out = _gather_engram_rows(out, num_tokens) + if self.tp_size > 1: + out = tensor_model_parallel_all_gather(out, dim=1) + return out[:, : self.n_hash_cols] + + +def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor: + """Exchange DP tokens for heads, retaining only this replica's tokens.""" + dp_group = get_engram_dp_group() + assert dp_group is not None + slot, remainder = divmod(staged.shape[0], dp_group.world_size) + assert remainder == 0 and 0 <= num_tokens <= slot + gathered = dp_group.all_gather(staged, dim=0) + local_heads, dim = staged.shape[1:] + rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim)) + _engram_select_rows( + gathered, + rows, + staged.shape[0], + dp_group.rank_in_group * slot, + local_heads * dim, + ) + return rows + + +class Engram(BaseEngram): + """NVIDIA Engram with asynchronous offload and node-local DP lookup.""" + + _prefetch_stream: torch.cuda.Stream | None = None + + def _create_embedding( + self, layout: EngramLayout, layer_hash_index: int + ) -> ParallelEngramEmbedding: + engram_config = get_current_vllm_config().engram_config + assert engram_config is not None + return ParallelEngramEmbedding( + layout.num_embeddings[layer_hash_index], + layout.head_dim, + tuple(size for order in layout.primes[layer_hash_index] for size in order), + cpu_offload=engram_config.cpu_offload, + dp_shared_memory=engram_config.dp_shared_memory, + ) + + def _init_staging(self, max_tokens: int, head_dim: int) -> None: + super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim) + if self.embed_tokens.cpu_offload: + self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device) + + def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: + """Prefetch local shared rows or the DP group's gathered hash IDs.""" + if self._prefetch_stream is None: + return super().prepare_embeddings(hash_ids) + rows = self.staged_rows[: hash_ids.shape[0]] + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small" + self._start_prefetch(hash_ids, rows, self._prefetch_stream) + + @eager_break_during_capture + def _start_prefetch( + self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream + ) -> None: + # Eager boundaries let the lookup span piecewise graph segments. + stream.wait_stream(torch.cuda.current_stream()) + # Keep temporary hash storage alive until lookup finishes reading it. + hash_ids.record_stream(stream) + with torch.cuda.stream(stream): + self.embed_tokens.lookup(hash_ids, rows, background=True) + + @eager_break_during_capture + def _finish_prefetch(self, stream: torch.cuda.Stream) -> None: + torch.cuda.current_stream().wait_stream(stream) + + def _ready_rows(self, num_tokens: int) -> torch.Tensor: + if self._prefetch_stream is not None: + self._finish_prefetch(self._prefetch_stream) + if self.embed_tokens.dp_size > 1: + slot = engram_gathered_num_tokens() + staged = self.staged_rows[: slot * self.embed_tokens.dp_size] + return _gather_engram_rows(staged, num_tokens) + return super()._ready_rows(num_tokens) diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index 6d803736ec5c..0868089da64b 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -12,6 +12,7 @@ from vllm.config import VllmConfig from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( + get_engram_dp_size, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -74,8 +75,9 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id -from ..common.engram import Engram, EngramLayout, NgramHashState +from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask +from .engram import Engram, gather_engram_hashes if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -466,6 +468,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # layer's sliding-window KV cache. Only PP ranks owning an engram # layer need it. self.engram_hash: NgramHashState | None = None + self.engram_dp_shared_memory = bool( + vllm_config.engram_config and vllm_config.engram_config.dp_shared_memory + ) self.engram_swa_prefix: str | None = None if self.engram_layout is not None: local_engram = any( @@ -594,12 +599,21 @@ def forward( swa_metadata.slot_mapping, swa_metadata.block_table, ) + elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1: + # DP-sharded lookups are collective, so a replica skipping the + # hash still has to reach them. + engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids) + if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. + # One gather feeds every layer sharing the DP-split table. + gathered_hashes = gather_engram_hashes( + engram_hashes, dp_shared_memory=self.engram_dp_shared_memory + ) for layer in islice(self.layers, self.start_layer, self.end_layer): engram = getattr(layer, "engram", None) if engram is not None: engram.prepare_embeddings( - engram_hashes[:, engram.layer_hash_index] + gathered_hashes[:, engram.layer_hash_index] ) full_num_tokens = positions.shape[0] From fa1b3b1922ccb2d7a597925765c4bec7f358d748 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Sun, 13 Sep 2026 18:08:30 +0800 Subject: [PATCH 126/774] [Bugfix][EPD] Preserve explicit multimodal UUIDs with caches disabled (#56432) Signed-off-by: Tianyu Guo --- .../test_process_multi_modal_uuids.py | 44 ++++++++++++------- vllm/renderers/base.py | 20 +++++---- 2 files changed, 38 insertions(+), 26 deletions(-) diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py index 76b394897ee1..8987782d71bc 100644 --- a/tests/renderers/test_process_multi_modal_uuids.py +++ b/tests/renderers/test_process_multi_modal_uuids.py @@ -176,9 +176,32 @@ def test_multi_modal_uuids_accepts_empty( assert processed_mm_uuids == mm_uuids -def test_multi_modal_uuids_ignored_when_caching_disabled(): - # When both processor cache is 0 and prefix caching disabled, the - # processor builds overrides from request id instead of using user UUIDs. +@pytest.mark.parametrize( + "mm_uuids, expected", + [ + ( + {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}, + {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}, + ), + ( + {"image": [None, "hash_stop"], "video": None}, + {"image": ["req-42-image-0", "hash_stop"], "video": ["req-42-video-0"]}, + ), + ( + {"image": ["", None]}, + {"image": ["", "req-42-image-1"], "video": ["req-42-video-0"]}, + ), + ( + {}, + { + "image": ["req-42-image-0", "req-42-image-1"], + "video": ["req-42-video-0"], + }, + ), + ], +) +def test_multi_modal_uuids_preserved_when_caching_disabled(mm_uuids, expected): + """Only missing UUIDs get request-local IDs when both caches are disabled.""" renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False) request_id = "req-42" @@ -186,7 +209,6 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(): "image": [cherry_pil_image, stop_pil_image], "video": baby_reading_np_ndarrays, } - mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]} mm_processor = renderer.get_mm_processor() mm_data_items = mm_processor.info.parse_mm_data(mm_data) @@ -196,16 +218,4 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(): mm_data, mm_data_items, mm_uuid_items, request_id ) - # Expect request-id-based overrides are passed through - assert set(mm_uuids.keys()) == {"image", "video"} - assert len(mm_uuids["image"]) == 2 - assert len(mm_uuids["video"]) == 1 - assert processed_mm_uuids["image"][0].startswith( - f"{request_id}-image-" - ) and processed_mm_uuids["image"][0].endswith("-0") - assert processed_mm_uuids["image"][1].startswith( - f"{request_id}-image-" - ) and processed_mm_uuids["image"][1].endswith("-1") - assert processed_mm_uuids["video"][0].startswith( - f"{request_id}-video-" - ) and processed_mm_uuids["video"][0].endswith("-0") + assert processed_mm_uuids == expected diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 0fd1b7610bae..4aca17fa4334 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -824,20 +824,22 @@ def _process_mm_uuids( ) -> MultiModalUUIDItems: model_config = self.model_config - # NOTE: When users explicitly turn off BOTH prefix caching and input - # processing caching, no multimodal features or embeddings will be - # reused across requests, therefore identifying multimodal data items - # by their content is no longer necessary, and we create uuids with - # `--`, overriding even user-provided ones. + # Skip content hashing when both caches are disabled, but preserve + # explicit UUIDs used to identify embeddings across EC instances. if ( model_config.multimodal_config and model_config.multimodal_config.mm_processor_cache_gb == 0 and not self.config.cache_config.enable_prefix_caching ): - mm_uuid_items = { - modality: [f"{mm_req_id}-{modality}-{i}" for i in range(data_count)] - for modality, data_count in mm_data_items.get_all_counts().items() - } + mm_uuid_items = dict(mm_uuid_items) + for modality, data_count in mm_data_items.get_all_counts().items(): + uuids = mm_uuid_items.get(modality) + if uuids is None: + uuids = [None] * data_count + mm_uuid_items[modality] = [ + uuid if uuid is not None else f"{mm_req_id}-{modality}-{i}" + for i, uuid in enumerate(uuids) + ] self._validate_mm_uuids(mm_data, mm_data_items, mm_uuid_items) From c351fd3c649569d8a585eb22cfaf2ee728e5fc22 Mon Sep 17 00:00:00 2001 From: Kebe Date: Sun, 13 Sep 2026 22:03:05 +0900 Subject: [PATCH 127/774] [Core] Fix ValueError on KV load failure with a hybrid KV cache (#50388) Signed-off-by: Kebe --- .../unit/test_kv_load_failure_recovery.py | 342 +++++++++++++++++- tests/v1/kv_connector/unit/utils.py | 2 + vllm/v1/core/sched/scheduler.py | 40 +- 3 files changed, 380 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 4f35527b0e3f..74bac5559823 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -2,18 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock +from typing import Literal +from unittest.mock import Mock, patch import pytest +import torch +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler -from vllm.v1.request import Request, RequestStatus +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) +from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( create_model_runner_output, create_request, create_scheduler, create_vllm_config, + make_kv_cache_config, ) @@ -34,6 +44,18 @@ def scheduler(): return create_scheduler(vllm_config) +@pytest.fixture +def hybrid_scheduler(): + """Scheduler for a hybrid model: full attention plus a Mamba group.""" + vllm_config = create_vllm_config(kv_load_failure_policy="recompute") + kv_cache_config = make_kv_cache_config( + block_size=vllm_config.cache_config.block_size, + mamba_enabled=True, + num_blocks=10000, + ) + return create_scheduler(vllm_config, kv_cache_config=kv_cache_config) + + @pytest.mark.parametrize( "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs", [ @@ -337,3 +359,319 @@ def test_async_progressive_load_failure( assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS assert scheduler.failed_recving_kv_req_ids == {request.request_id} assert scheduler.connector.get_num_new_matched_tokens.call_count == 1 + + +def _schedule_hybrid_async_load( + scheduler: Scheduler, num_prompt_blocks: int, num_external_computed_blocks: int +) -> Request: + """Put one request into an async external load on a hybrid model.""" + request = create_request(num_tokens=num_prompt_blocks * scheduler.block_size) + scheduler.add_request(request=request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_computed_blocks * scheduler.block_size}, + async_load=True, + ) + ) + scheduler.connector.take_events.return_value = () + return request + + +def _report_invalid_blocks( + scheduler: Scheduler, scheduler_output, invalid_block_ids: set[int] +) -> None: + scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids=invalid_block_ids, + use_eos=True, + ), + ) + + +@pytest.mark.parametrize("failed_group_idx", [0, 1]) +def test_hybrid_load_failure_recomputes_whole_request( + hybrid_scheduler: Scheduler, failed_group_idx: int +): + """A hybrid model must recover, not crash. + + Its groups sit on different block grids, so there is no single truncation + point; the request restarts from scratch instead. + """ + num_prompt_blocks, num_external_computed_blocks = 10, 9 + request = _schedule_hybrid_async_load( + hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks + ) + scheduler_output = hybrid_scheduler.schedule() + + block_ids_per_group = hybrid_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert len(block_ids_per_group) == 2, "expected a hybrid (multi-group) layout" + null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id + # The Mamba group backs only its aligned snapshot with a real block, so + # pick a block that actually belongs to this request. + failed_block = next( + block_id + for block_id in block_ids_per_group[failed_group_idx] + if block_id != null_block_id + ) + + _report_invalid_blocks(hybrid_scheduler, scheduler_output, {failed_block}) + + assert request.num_computed_tokens == 0 + assert hybrid_scheduler.failed_recving_kv_req_ids == {request.request_id} + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_hybrid_load_failure_ignores_null_block(hybrid_scheduler: Scheduler): + """The null block is shared by every request and holds no KV, so reporting + it must not restart anyone.""" + num_prompt_blocks, num_external_computed_blocks = 10, 9 + request = _schedule_hybrid_async_load( + hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks + ) + scheduler_output = hybrid_scheduler.schedule() + + null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id + _, mamba_block_ids = hybrid_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + # The Mamba group parks its unaligned positions on the null block, so a + # request that never failed would otherwise be restarted by this report. + assert mamba_block_ids[0] == null_block_id + + num_computed_tokens = request.num_computed_tokens + _report_invalid_blocks(hybrid_scheduler, scheduler_output, {null_block_id}) + + assert request.num_computed_tokens == num_computed_tokens + assert not hybrid_scheduler.failed_recving_kv_req_ids + + +def _make_deepseek_v4_five_group_config(num_blocks: int) -> KVCacheConfig: + """Build the scheduler-side projection of V4-Flash's packed KV layout.""" + + def mla_spec() -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + tokens_per_state=4, + model_version="deepseek_v4", + ) + + def swa_spec( + *, block_size: int, sliding_window: int, head_size: int + ) -> SlidingWindowMLASpec: + return SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version="deepseek_v4", + ) + + # DeepSeek-V4 packs these into five physical groups. The two SWA groups + # have identical scheduling semantics but are split to align layer tuples. + groups = [ + KVCacheGroupSpec(["full_mla_and_indexer"], mla_spec()), + KVCacheGroupSpec( + ["c4_state"], + swa_spec(block_size=4, sliding_window=8, head_size=2048), + ), + KVCacheGroupSpec( + ["c128_state"], + swa_spec(block_size=8, sliding_window=128, head_size=2048), + ), + KVCacheGroupSpec( + ["swa.0"], + swa_spec(block_size=64, sliding_window=512, head_size=512), + ), + KVCacheGroupSpec( + ["swa.1"], + swa_spec(block_size=64, sliding_window=512, head_size=512), + ), + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + + +def _create_deepseek_v4_scheduler( + kv_load_failure_policy: Literal["recompute", "fail"] = "recompute", +) -> Scheduler: + num_blocks = 10000 + vllm_config = create_vllm_config( + block_size=256, + max_num_batched_tokens=4096, + max_model_len=8192, + # max_model_len only reaches SchedulerConfig; opt-125m still caps the + # model at 2048, which is below the 9-block requests used here. + hf_overrides={"max_position_embeddings": 8192}, + kv_load_failure_policy=kv_load_failure_policy, + ) + return create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=_make_deepseek_v4_five_group_config(num_blocks), + hash_block_size=4, + ) + + +@pytest.fixture +def deepseek_v4_scheduler() -> Scheduler: + return _create_deepseek_v4_scheduler() + + +def _schedule_deepseek_v4_async_load( + scheduler: Scheduler, +) -> tuple[Request, SchedulerOutput]: + num_external_tokens = 8 * scheduler.block_size + request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_tokens}, async_load=True + ) + ) + scheduler.connector.take_events.return_value = () + return request, scheduler.schedule() + + +@pytest.mark.parametrize("failed_group_idx", range(5)) +def test_deepseek_v4_load_failure_recomputes_whole_request( + deepseek_v4_scheduler: Scheduler, failed_group_idx: int +): + """A failed load in any V4-Flash KV group must not kill EngineCore.""" + request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) + + block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert len(block_ids_per_group) == 5 + null_block_id = ( + deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id + ) + failed_block_id = next( + block_id + for block_id in block_ids_per_group[failed_group_idx] + if block_id != null_block_id + ) + + deepseek_v4_scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids={failed_block_id}, + use_eos=True, + ), + ) + + assert request.num_computed_tokens == 0 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert deepseek_v4_scheduler.failed_recving_kv_req_ids == {request.request_id} + + +def test_deepseek_v4_load_failure_ignores_shared_null_block( + deepseek_v4_scheduler: Scheduler, +): + """The shared null block contains no request KV and cannot fail a load.""" + request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) + null_block_id = ( + deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id + ) + block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert any(null_block_id in group for group in block_ids_per_group[1:]) + num_computed_tokens = request.num_computed_tokens + + deepseek_v4_scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids={null_block_id}, + use_eos=True, + ), + ) + + assert request.num_computed_tokens == num_computed_tokens + assert not deepseek_v4_scheduler.failed_recving_kv_req_ids + + +def test_deepseek_v4_sync_load_failure_fails_and_evicts_real_blocks(): + """Fail policy must terminate cleanly and never evict the null block.""" + scheduler = _create_deepseek_v4_scheduler(kv_load_failure_policy="fail") + num_external_tokens = 8 * scheduler.block_size + request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_tokens}, async_load=False + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.take_events.return_value = () + scheduler_output = scheduler.schedule() + + block_ids_per_group = scheduler.kv_cache_manager.get_block_ids(request.request_id) + assert len(block_ids_per_group) == 5 + null_block_id = scheduler.kv_cache_manager.block_pool.null_block.block_id + real_block_ids = { + block_id + for group_block_ids in block_ids_per_group + for block_id in group_block_ids + if block_id != null_block_id + } + failed_block_id = next( + block_id for block_id in block_ids_per_group[2] if block_id != null_block_id + ) + + evicted_block_ids: list[set[int]] = [] + original_evict_blocks = scheduler.kv_cache_manager.evict_blocks + + def evict_blocks_spy(block_ids: set[int]) -> None: + evicted_block_ids.append(set(block_ids)) + original_evict_blocks(block_ids) + + with ( + patch.object( + scheduler.kv_cache_manager, + "evict_blocks", + side_effect=evict_blocks_spy, + ), + patch.object(scheduler, "_connector_finished", return_value=(False, None)), + ): + outputs = scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + [request], + invalid_block_ids={failed_block_id}, + use_eos=True, + ), + ) + + assert request.status == RequestStatus.FINISHED_ERROR + assert request.get_finished_reason() == FinishReason.ERROR + assert request.request_id not in scheduler.requests + assert evicted_block_ids == [real_block_ids] + assert null_block_id not in evicted_block_ids[0] + + engine_outputs = next(iter(outputs.values())) + assert len(engine_outputs.outputs) == 1 + assert engine_outputs.outputs[0].request_id == request.request_id + assert engine_outputs.outputs[0].finish_reason == FinishReason.ERROR diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 7b70065d7f95..07a40dc4ec00 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -155,6 +155,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -185,6 +186,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 82203738459b..1b17eb3b1518 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3058,6 +3058,11 @@ def _update_requests_with_invalid_blocks( For observability, it also accumulates the total number of tokens that will need to be recomputed across all affected requests. + A hybrid model spreads one token span over KV cache groups that disagree + on block size, so a block index maps to a different token position in + each group and there is no single longest valid prefix. Those requests + recompute in full instead. + Args: requests: The set of requests to scan for invalid blocks. invalid_block_ids: IDs of invalid blocks. @@ -3081,18 +3086,49 @@ def _update_requests_with_invalid_blocks( # these requests must be rescheduled, but only the first will recompute # it. This set tracks blocks already marked for recomputation. marked_invalid_block_ids: set[int] = set() + # Positions a group holds no KV for -- skipped sliding-window blocks, + # and every Mamba block but the aligned snapshot -- all point at this + # one block, shared by every request. It carries no request's data, so + # it can neither fail for a request nor be evicted on its behalf. + null_block_id = self.kv_cache_manager.block_pool.null_block.block_id for request in requests: is_affected = False marked_invalid_block = False req_id = request.request_id - # TODO (davidb): add support for hybrid memory allocator - (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) + # One block-id list per KV cache group. + req_block_ids_per_group = self.kv_cache_manager.get_block_ids(req_id) # We iterate only over blocks that may contain externally computed # tokens req_num_computed_tokens = ( request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) ) + if len(req_block_ids_per_group) > 1: + # No single truncation point exists across groups, so the whole + # request is recomputed. Blocks are still marked, so a request + # sharing one of them is rescheduled too. + request_invalid_block_ids = { + block_id + for req_block_ids in req_block_ids_per_group + for block_id in req_block_ids + if block_id != null_block_id and block_id in invalid_block_ids + } + if request_invalid_block_ids: + marked_invalid_block_ids |= request_invalid_block_ids + total_affected_tokens += req_num_computed_tokens + request.num_computed_tokens = 0 + if evict_blocks: + for req_block_ids in req_block_ids_per_group: + blocks_to_evict.update( + block_id + for block_id in req_block_ids + if block_id != null_block_id + ) + affected_req_ids.add(req_id) + continue + + (req_block_ids,) = req_block_ids_per_group + req_num_computed_blocks = ( req_num_computed_tokens + self.block_size - 1 ) // self.block_size From 2e9f7bb45475c7e461079284a73ec9c53761c6b5 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 13 Sep 2026 08:20:21 -0500 Subject: [PATCH 128/774] [Pooling] Report actual input token usage for scoring APIs (#56573) Signed-off-by: Taneem Ibrahim --- .../language/pooling/test_max_tokens_per_doc.py | 16 ++++++++-------- vllm/entrypoints/pooling/scoring/serving.py | 16 +++++++++------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py index b6d6ba145478..d99aa4990982 100644 --- a/tests/models/language/pooling/test_max_tokens_per_doc.py +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -80,10 +80,10 @@ class TestConfig: ], # This model has produced both prompt-token totals in CI/local cache; # keep truncation checks exact while tolerating the boundary delta. - without_truncated_prompt_tokens=(285, 286), - with_max_tokens_per_query_prompt_tokens=(155, 156), - with_max_tokens_per_doc_prompt_tokens=155, - with_max_tokens_per_query_and_doc_prompt_tokens=25, + without_truncated_prompt_tokens=(284, 285), + with_max_tokens_per_query_prompt_tokens=(154, 155), + with_max_tokens_per_doc_prompt_tokens=154, + with_max_tokens_per_query_and_doc_prompt_tokens=24, ), # 4. late-interaction TestConfig( @@ -94,10 +94,10 @@ class TestConfig: "512", "--trust-remote-code", ], - without_truncated_prompt_tokens=285, - with_max_tokens_per_query_prompt_tokens=155, - with_max_tokens_per_doc_prompt_tokens=155, - with_max_tokens_per_query_and_doc_prompt_tokens=25, + without_truncated_prompt_tokens=284, + with_max_tokens_per_query_prompt_tokens=154, + with_max_tokens_per_doc_prompt_tokens=154, + with_max_tokens_per_query_and_doc_prompt_tokens=24, ), # 5. jinaai/jina-reranker-v3 TestConfig( diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index b32a7a5730af..674522b7ebde 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -80,6 +80,11 @@ def _build_response( request_id = ctx.request_id created_time = ctx.created_time model_name = ctx.model_name + assert ctx.engine_inputs is not None + num_prompt_tokens = sum( + self._extract_prompt_len(engine_input["prompts"]) + for engine_input in ctx.engine_inputs + ) if isinstance(ctx.request, ScoreRequest): return self._request_output_to_score_response( @@ -87,6 +92,7 @@ def _build_response( request_id, created_time, model_name, + num_prompt_tokens, ) elif isinstance(ctx.request, RerankRequest): return self._request_output_to_rerank_response( @@ -95,6 +101,7 @@ def _build_response( model_name, ctx.request.documents, ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch), + num_prompt_tokens, ) else: raise ValueError(f"Invalid {self.request_id_prefix} request type") @@ -105,9 +112,9 @@ def _request_output_to_score_response( request_id: str, created_time: int, model_name: str, + num_prompt_tokens: int, ) -> JSONResponse: items: list[ScoreResponseData] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): classify_res = ScoringRequestOutput.from_base(final_res) @@ -116,10 +123,7 @@ def _request_output_to_score_response( index=idx, score=classify_res.outputs.score, ) - prompt_token_ids = final_res.prompt_token_ids - items.append(item) - num_prompt_tokens += len(prompt_token_ids) usage = UsageInfo( prompt_tokens=num_prompt_tokens, @@ -143,12 +147,12 @@ def _request_output_to_rerank_response( model_name: str, documents: ScoreInput | list[ScoreInput], top_n: int, + num_prompt_tokens: int, ) -> JSONResponse: if not isinstance(documents, list): documents = [documents] results: list[RerankResult] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): classify_res = ScoringRequestOutput.from_base(final_res) @@ -166,8 +170,6 @@ def _request_output_to_rerank_response( relevance_score=classify_res.outputs.score, ) results.append(result) - prompt_token_ids = final_res.prompt_token_ids - num_prompt_tokens += len(prompt_token_ids) # sort by relevance, then return the top n if set results.sort(key=lambda x: x.relevance_score, reverse=True) From de500290ffffeb7bd834872c0dea5331590dd40d Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sun, 13 Sep 2026 21:44:24 +0800 Subject: [PATCH 129/774] [XPU][CI] fix jit_warmup_triton_launcher (#56670) Signed-off-by: zhenwei-intel --- tests/model_executor/test_jit_warmup_triton_launcher.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index b36361e637de..2e81cb6000f8 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -21,6 +21,7 @@ triton_kernel_dispatcher_with_warmup, triton_warmup_inputs, ) +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -257,10 +258,11 @@ def dispatch(first: str, second: int) -> LaunchSpec: dispatch.get_warmup_keys() -def test_triton_kernel_dispatch_uses_cuda_fake_tensors( +def test_triton_kernel_dispatch_uses_device_fake_tensors( monkeypatch: pytest.MonkeyPatch, ) -> None: kernel = _FakeTritonKernel() + device_type = current_platform.device_type def warmup_inputs() -> dict[str, Any]: return dict( @@ -271,7 +273,7 @@ def warmup_inputs() -> dict[str, Any]: @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) def dispatch(first: torch.Tensor, second: int) -> LaunchSpec: assert isinstance(first, torch.Tensor) - assert first.is_cuda + assert first.device.type == device_type assert first[0].is_contiguous() return (first.shape[0],), dict(CONST=first[0].numel()) @@ -282,7 +284,7 @@ def dispatch(first: torch.Tensor, second: int) -> LaunchSpec: keys = dispatch.get_warmup_keys() assert len(keys) == 1 - assert dict(keys[0].inputs)["first"].device.type == "cuda" + assert dict(keys[0].inputs)["first"].device.type == device_type def test_triton_kernel_decorates_native_launchers( From dd4c8410707bdb833b155ff3258ce5271c0f476e Mon Sep 17 00:00:00 2001 From: linitra24 Date: Sun, 13 Sep 2026 22:10:54 +0800 Subject: [PATCH 130/774] [LoRA] Add LoRA support for DeepSeek-V4 Flash Vision (#55897) Signed-off-by: linitra24 --- docs/models/supported_models.md | 2 +- vllm/lora/ops/triton_ops/lora_shrink_op.py | 2 +- vllm/models/deepseek_v4/common/vl_model.py | 41 +++++++++++++++++++++- 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d06154fb5bd5..81f4fb7df1f6 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -547,7 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I+ | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ | | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I+ | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ | -| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I+ | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | | ✅︎ | +| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I+ | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | ✅︎ | ✅︎ | | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + IE+ | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ | | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ | diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py index 4144c16eadd7..d05c8f7c2a52 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py @@ -179,7 +179,7 @@ def _lora_shrink( assert weight.dtype in [torch.float16, torch.bfloat16] assert inputs.size(1) == lora_a_weights[0].size(-1) - assert inputs.is_contiguous() + inputs = inputs.contiguous() assert output_tensor.is_contiguous() # metadata sanity check diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py index 4cccb85a8302..97701cc2d507 100644 --- a/vllm/models/deepseek_v4/common/vl_model.py +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -25,9 +25,11 @@ from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsLoRA, SupportsMultiModal, SupportsPP, ) +from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -35,6 +37,7 @@ maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalKwargsItem from .mm_preprocess import ( IMAGE_PLACEHOLDER, @@ -81,7 +84,7 @@ def _make_deepseek_v4_vl_weights_mapper( dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) class DeepseekV4ForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3, SupportsLoRA ): """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. @@ -89,9 +92,19 @@ class DeepseekV4ForConditionalGeneration( delegates through ``language_model`` via the protocol defaults. """ + packed_modules_mapping = { + "gate_up_proj": ["w1", "w3"], + "fused_wqa_wkv": ["wq_a", "wkv"], + "fused_wkv_wgate": ["wkv", "wgate"], + # for visual encoder + "wqkv": ["wqkv"], + "w1": ["w1"], + } + # The MoE router needs raw token ids to detect image sentinel tokens # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. requires_raw_input_tokens = True + supports_tower_connector_lora = True @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: @@ -329,3 +342,29 @@ def process_weights_after_loading(self) -> None: return self.language_model.process_weights_after_loading() self._weights_finalized = True + + def get_mm_mapping(self) -> MultiModelKeys: + """Get the module prefixes in the multimodal model.""" + return MultiModelKeys.from_string_field( + language_model="language_model", + connector="aligner", + tower_model="vision.", + ) + + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + if modality != "image": + raise ValueError(f"Unsupported modality: {modality!r}") + + patches = mm_kwargs.get("patches") if mm_kwargs else None + if patches is not None and isinstance(patches.data, torch.Tensor): + tower_tokens = patches.data.shape[0] + else: + tower_tokens = num_mm_embeds * self.config.vision_downsample_ratio**2 + + return tower_tokens, num_mm_embeds From 8cf9de9080ea5cded1c93e3d9a5ac531da0c93b6 Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Sun, 13 Sep 2026 07:22:03 -0700 Subject: [PATCH 131/774] [Test][Determinism] Cover VLM batch invariance in default execution mode (#56528) Signed-off-by: Shenglei Fu Co-authored-by: Claude Sonnet 5 --- .../determinism/test_batch_invariance_vlm.py | 48 ++++++++++++++++--- 1 file changed, 41 insertions(+), 7 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance_vlm.py b/tests/v1/determinism/test_batch_invariance_vlm.py index 266768f67253..00f138d62c81 100644 --- a/tests/v1/determinism/test_batch_invariance_vlm.py +++ b/tests/v1/determinism/test_batch_invariance_vlm.py @@ -59,15 +59,22 @@ def _assert_batch_invariant(llm: LLM, inputs: list[dict], sampling) -> None: _extract_step_logprobs(llm.generate([inp], sampling, use_tqdm=False)[0]) for inp in inputs ] - bsN = [_extract_step_logprobs(out) for out in llm.generate(inputs, sampling)] + bsN_outputs = llm.generate(inputs, sampling, use_tqdm=False) + assert len(bs1) == len(inputs) + assert len(bsN_outputs) == len(inputs) + bsN = [_extract_step_logprobs(out) for out in bsN_outputs] failures = [] for i, ((lp1, t1), (lpN, tN)) in enumerate(zip(bs1, bsN)): if lp1 is None or lpN is None: - pytest.skip( - "Logits are not available on RequestOutput; " - "enable logprobs return to run this test." + failures.append(f"req {i}: selected-token logprobs are unavailable") + continue + if len(t1) != sampling.max_tokens or len(tN) != sampling.max_tokens: + failures.append( + f"req {i}: expected {sampling.max_tokens} tokens, " + f"got bs1={len(t1)} bsN={len(tN)}" ) + continue if t1 != tN: failures.append(f"req {i}: token mismatch bs1={t1} bsN={tN}") elif not torch.equal(lp1, lpN): @@ -84,16 +91,43 @@ def _assert_batch_invariant(llm: LLM, inputs: list[dict], sampling) -> None: @skip_if_not_cuda @pytest.mark.parametrize("input_type", ["image", "video"]) @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"]) -def test_vlm_batch_invariance_bs1_vs_bsN(input_type: str, mm_encoder_attn_backend: str): +def test_vlm_batch_invariance_bs1_vs_bsN( + input_type: str, + mm_encoder_attn_backend: str, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True) + + +@skip_if_not_cuda +def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False) + + +def _run_vlm_batch_invariance( + input_type: str, mm_encoder_attn_backend: str, enforce_eager: bool +) -> None: inputs = _make_inputs(input_type, num_reqs=4) - sampling = SamplingParams(temperature=0.0, max_tokens=8, seed=1234, logprobs=5) + sampling = SamplingParams( + temperature=0.0, + max_tokens=8, + seed=1234, + ignore_eos=True, + logprobs=5, + ) llm = LLM( model=VLM_TEST_MODEL, dtype="bfloat16", - enforce_eager=True, + enforce_eager=enforce_eager, max_num_seqs=16, + max_model_len=4096, + max_num_batched_tokens=2048, gpu_memory_utilization=0.85, + enable_prefix_caching=False, + mm_processor_cache_gb=0, mm_encoder_attn_backend=mm_encoder_attn_backend, ) try: From 82a85dc1d2d5b3ad4453eaa1bd7596de888a66b3 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Sun, 13 Sep 2026 10:29:41 -0400 Subject: [PATCH 132/774] [Attention] Add Triton/FlashInfer composite for multimodal prefix attention (#56305) Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Co-authored-by: Codex Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- docs/design/attention_backends.md | 49 ++ .../attention/test_attention_selector.py | 139 +++++ tests/v1/attention/test_mm_prefix.py | 246 +++++++++ tests/v1/spec_decode/test_dflash2.py | 1 + vllm/platforms/cuda.py | 8 + vllm/v1/attention/backend.py | 10 + vllm/v1/attention/backends/composite.py | 509 ++++++++++++++++++ vllm/v1/attention/backends/flash_attn.py | 4 + vllm/v1/attention/backends/flex_attention.py | 4 + vllm/v1/attention/backends/registry.py | 6 + vllm/v1/attention/backends/triton_attn.py | 4 + .../attention/backends/triton_flash_attn.py | 67 +++ .../attention/backends/triton_flashinfer.py | 22 + vllm/v1/attention/backends/utils.py | 21 +- vllm/v1/attention/selector.py | 8 +- 15 files changed, 1095 insertions(+), 3 deletions(-) create mode 100644 vllm/v1/attention/backends/composite.py create mode 100644 vllm/v1/attention/backends/triton_flash_attn.py create mode 100644 vllm/v1/attention/backends/triton_flashinfer.py diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 3fdc23ff8422..d175add40b84 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -56,6 +56,55 @@ llm = LLM( ## Backend Selection Behavior +### Triton/FlashAttention Composite + +On Hopper, `TRITON_FLASH_ATTN` is preferred for compatible multimodal-prefix +configurations. It uses Triton when the current queries need bidirectional +image attention and FlashAttention for causal text prefills and decode, with a +shared KV cache. The causal child must resolve to FA4. This happens +automatically for FA4-only shapes such as head size 512 and models whose +version policy promotes all layers to FA4, including Gemma 4. A standalone +head-size-256 configuration whose version policy selects FA3 falls back to +Triton for the whole layer. + +No attention override is needed for supported multimodal models: + +```bash +vllm serve google/gemma-4-31B-it +``` + +### Triton/FlashInfer Composite + +On Blackwell, `TRITON_FLASHINFER` is preferred for compatible +multimodal-prefix configurations, including Gemma 4 with BF16 or FP8 KV cache. It uses +Triton when a batch's current queries require bidirectional image attention, +and FlashInfer for causal text prefills and decode. Historical image tokens +alone do not select Triton. + +This instantiates `create_composite_attention_backend` from +`vllm/v1/attention/backends/composite.py` with Triton, +FlashInfer, and `MMPrefixAttentionRouting`. The reusable factory owns child +implementations, metadata dispatch, compatible cache requirements, and workspace +sharing. The routing policy selects the child and defines graph-capture safety. +Other combinations can reuse the same machinery. + +The backend supports head dimensions 256/512, FP16/BF16 and FP8 KV cache, +and 64-token kernel pages with a head-major cache layout. TRTLLM handles causal +attention at both head dimensions; its hdim512 kernels do not support 128-token +pages. +For Gemma 4, full CUDA graphs cover single-token batches; multi-token batches +use the non-full-graph execution path. Models whose image masks extend beyond +the sliding window cannot use full attention graphs with this composite. +Context parallelism, R-SWA, attention sinks, +and adaptive verification are not supported by this composite. + +It can also be selected explicitly: + +```bash +vllm serve google/gemma-4-31B-it \ + --attention-backend TRITON_FLASHINFER +``` + ### Manual Selection When you explicitly set a backend via `--attention-backend` or `AttentionConfig`: diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index d74a119e654d..53867a5edf93 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -808,3 +808,142 @@ def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected): ) assert impl.vllm_flash_attn_version == expected assert impl.fa4_hd256 == (expected == 4) + + +@pytest.fixture +def blackwell_selection(): + with ( + patch.object( + type(current_platform), + "get_device_capability", + return_value=DeviceCapability(10, 0), + ), + patch( + "vllm.v1.attention.backends.fa_utils.is_fa_version_supported", + return_value=True, + ), + ): + yield + + +@pytest.fixture +def hopper_selection(): + with ( + patch.object( + type(current_platform), + "get_device_capability", + return_value=DeviceCapability(9, 0), + ), + patch( + "vllm.v1.attention.backends.fa_utils.is_fa_version_supported", + return_value=True, + ), + ): + yield + + +@blackwell_only +@pytest.mark.parametrize("use_mm_prefix", [False, True]) +@pytest.mark.parametrize("flash_attn_version", [None, 3, 4]) +def test_hopper_mm_prefix_selects_triton_flash_attn( + use_mm_prefix, flash_attn_version, hopper_selection +): + """Hopper selects the composite only when its causal route resolves FA4.""" + from vllm.engine.arg_utils import EngineArgs + + config = EngineArgs( + model="google/gemma-4-31B-it", + dtype="bfloat16", + attention_config=AttentionConfig(flash_attn_version=flash_attn_version), + ).create_engine_config() + with set_current_vllm_config(config): + backend = get_attn_backend( + 256, torch.bfloat16, None, use_mm_prefix=use_mm_prefix + ) + fa4_resolved = config.attention_config.flash_attn_version == 4 + if use_mm_prefix: + expected = "TRITON_FLASH_ATTN" if fa4_resolved else "TRITON_ATTN" + else: + expected = "FLASH_ATTN" + assert backend.get_name() == expected + + +@blackwell_only +@pytest.mark.parametrize("use_mm_prefix", [False, True]) +@pytest.mark.parametrize("kv_cache_dtype", [None, "fp8_e4m3"]) +def test_mm_prefix_selects_composite_without_changing_causal_default( + use_mm_prefix, kv_cache_dtype, blackwell_selection +): + """The image-mask requirement must reach CUDA's automatic backend priority.""" + from vllm.engine.arg_utils import EngineArgs + + config = EngineArgs( + model="google/gemma-4-31B-it", dtype="bfloat16" + ).create_engine_config() + with set_current_vllm_config(config): + backend = get_attn_backend( + 256, torch.bfloat16, kv_cache_dtype, use_mm_prefix=use_mm_prefix + ) + expected = "TRITON_FLASHINFER" if use_mm_prefix else "FLASHINFER" + assert backend.get_name() == expected + + +@blackwell_only +@pytest.mark.parametrize( + "feature_kwargs,reason", + [ + ({"use_rswa": True}, "R-SWA not supported"), + ({"kv_cache_dtype": "int8_per_token_head"}, "kv_cache_dtype not supported"), + ({"block_size": 16}, "block_size not supported"), + ({"use_dcp": True}, "DCP not supported"), + ({"use_pcp": True}, "PCP not supported"), + ({"use_adaptive_verification": True}, "device-cpu query lens mismatch"), + ], +) +def test_composite_rejects_features_not_shared_by_both_routes(feature_kwargs, reason): + """Automatic selection must not admit features only one child can execute.""" + from vllm.engine.arg_utils import EngineArgs + from vllm.v1.attention.backends.triton_flashinfer import TritonFlashInferBackend + + config = EngineArgs( + model="google/gemma-4-31B-it", dtype="bfloat16" + ).create_engine_config() + args = dict( + head_size=256, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=128, + use_mla=False, + has_sink=False, + use_sparse=False, + use_mm_prefix=True, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(10, 0), + attn_type=AttentionType.DECODER, + ) + args.update(feature_kwargs) + with set_current_vllm_config(config): + reasons = TritonFlashInferBackend.validate_configuration(**args) + assert any(reason in message for message in reasons), reasons + + +@blackwell_only +def test_rswa_selection_does_not_reuse_causal_result(blackwell_selection): + """An R-SWA requirement must invalidate a previously cached causal selection.""" + from vllm.engine.arg_utils import EngineArgs + + config = EngineArgs( + model="google/gemma-4-31B-it", + dtype="bfloat16", + attention_config={"backend": "TRITON_FLASHINFER"}, + ).create_engine_config() + with set_current_vllm_config(config): + assert ( + get_attn_backend(256, torch.bfloat16, None).get_name() + == "TRITON_FLASHINFER" + ) + config.model_config.model_arch_config.rswa_window = 128 + with pytest.raises(ValueError, match="R-SWA"): + get_attn_backend(256, torch.bfloat16, None) + config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN + assert get_attn_backend(256, torch.bfloat16, None).get_name() == "TRITON_ATTN" diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index 6b8717389f45..f878ede4865f 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/tests/v1/attention/test_mm_prefix.py @@ -638,3 +638,249 @@ def test_mm_prefix_kv_cache_path(head_size: int): assert not torch.allclose(expected, causal_only, atol=2e-2, rtol=2e-2), ( "test batch does not actually exercise the mm_prefix branch" ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA backend imports") +@pytest.mark.parametrize( + "query_len,seq_len,spans,expected,unclamped_expected", + [ + (16, 32, [(20, 28)], True, True), + (16, 32, [(0, 16)], False, True), + (16, 32, [(0, 17)], True, True), + (16, 32, [(32, 40)], False, False), + (16, 32, [(31, 40)], False, True), + (1, 24, [(20, 28)], False, True), + (0, 24, [(20, 28)], False, False), + (16, 32, [(20, 20)], False, False), + (16, 64, [(20, 28)], False, False), + ], +) +def test_composite_routes_queries_that_need_image_masking( + query_len, seq_len, spans, expected, unclamped_expected +): + """Historical images and single-query steps must use the causal graph path.""" + from types import SimpleNamespace + + from vllm.v1.attention.backends.composite import requires_mm_prefix + + common = SimpleNamespace( + max_query_len=query_len, + mm_req_doc_ranges={0: spans}, + query_start_loc_cpu=torch.tensor([0, query_len]), + seq_lens_cpu_upper_bound=torch.tensor([seq_len]), + ) + assert requires_mm_prefix(common) is expected + assert requires_mm_prefix(common, unclamped_window=True) is unclamped_expected + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA attention kernels") +@pytest.mark.parametrize("head_size", [256, 512]) +@pytest.mark.parametrize( + "reverse_children,cache_dtype", + [(False, "auto"), (True, "auto"), (False, "fp8_e4m3")], +) +@pytest.mark.parametrize("causal_backend_name", ["FLASHINFER", "FLASH_ATTN"]) +def test_composite_shared_cache_across_image_and_causal_steps( + head_size, reverse_children, cache_dtype, causal_backend_name +): + """Changing routes must preserve KV writes, image masking and graph replay.""" + if causal_backend_name == "FLASHINFER": + if not current_platform.is_device_capability_family(100): + pytest.skip("The Triton/FlashInfer composite supports Blackwell") + else: + if not current_platform.is_device_capability(90): + pytest.skip("The Triton/FlashAttention composite supports Hopper") + if reverse_children or cache_dtype != "auto": + pytest.skip("Reverse-child and FP8 coverage is backend-independent") + + from vllm.config import set_current_vllm_config + from vllm.engine.arg_utils import EngineArgs + from vllm.model_executor.layers.attention import Attention + from vllm.utils.torch_utils import set_default_torch_dtype + from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata + from vllm.v1.attention.backends.composite import ( + MMPrefixAttentionRouting, + create_composite_attention_backend, + ) + from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionMetadata, + ) + from vllm.v1.attention.backends.flashinfer import ( + FlashInferBackend, + FlashInferMetadata, + FlashInferTrtllmAPIDecode, + TRTLLMPrefill, + ) + from vllm.v1.attention.backends.triton_attn import ( + TritonAttentionBackend, + TritonAttentionMetadata, + ) + from vllm.v1.attention.backends.triton_flash_attn import ( + TritonFlashAttentionBackend, + ) + from vllm.v1.attention.backends.triton_flashinfer import ( + TritonFlashInferBackend, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec + + if causal_backend_name == "FLASHINFER": + backend = TritonFlashInferBackend + causal_metadata_type = FlashInferMetadata + block_size = backend.get_supported_kernel_block_sizes()[0] + else: + backend = TritonFlashAttentionBackend + causal_metadata_type = FlashAttentionMetadata + block_size = backend.get_preferred_block_size(16) + num_blocks = (209 + block_size - 1) // block_size + torch.manual_seed(42) + attention_config = {"backend": backend.get_name()} + if causal_backend_name == "FLASH_ATTN": + attention_config["flash_attn_version"] = 4 + cfg = EngineArgs( + model="google/gemma-4-31B-it", + dtype="bfloat16", + max_model_len=512, + max_num_seqs=4, + max_num_batched_tokens=512, + block_size=block_size, + kv_cache_dtype=cache_dtype, + enforce_eager=True, + attention_config=attention_config, + ).create_engine_config() + cfg.cache_config.kv_cache_layout = "LBHNC" + window = 128 if head_size == 256 else None + storage_dtype = torch.uint8 if cache_dtype == "fp8_e4m3" else DTYPE + spec_kwargs = dict( + block_size=block_size, + num_kv_heads=16, + head_size=head_size, + dtype=storage_dtype, + kv_quant_mode=get_kv_quant_mode(cache_dtype), + ) + spec = ( + SlidingWindowSpec(**spec_kwargs, sliding_window=window) + if window + else FullAttentionSpec(**spec_kwargs) + ) + if reverse_children: + # Swapping the children and policy must preserve outputs and KV writes; + # dispatch cannot depend on a particular child's class or position. + class ReversedRouting(MMPrefixAttentionRouting): + capture_variant = 0 + + def select(self, metadata): + return 1 - super().select(metadata) + + @staticmethod + def variant_uses_mm_prefix(variant): + return variant == 1 + + backend = create_composite_attention_backend( + FlashInferBackend, + TritonAttentionBackend, + name="ReversedCompositeBackend", + backend_name="CUSTOM", + module=__name__, + routing_policy=ReversedRouting, + ) + with set_current_vllm_config(cfg), set_default_torch_dtype(DTYPE): + layer = Attention( + 32, + head_size, + head_size**-0.5, + num_kv_heads=16, + cache_config=cfg.cache_config, + prefix="composite_test", + per_layer_sliding_window=window, + attn_backend=backend, + ).to(DEVICE) + layer.mm_prefix_clamp_sliding_window = window is not None + assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) + if window: + layer.mm_prefix_clamp_sliding_window = False + assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == ( + AttentionCGSupport.NEVER + ) + layer.mm_prefix_clamp_sliding_window = True + builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE) + builder.set_kernel_block_size(block_size) + # Match production, where unwritten page tails are zero-initialized. + cache = torch.zeros( + num_blocks, + 16, + block_size, + 2 * head_size, + dtype=storage_dtype, + device=DEVICE, + ) + keys = torch.randn(209, 16, head_size, dtype=DTYPE, device=DEVICE) + values = torch.randn_like(keys) + if cache_dtype == "fp8_e4m3": + keys = keys.to(torch.float8_e4m3fn).to(DTYPE) + values = values.to(torch.float8_e4m3fn).to(DTYPE) + table = torch.arange(num_blocks, dtype=torch.int32, device=DEVICE).unsqueeze(0) + previous = 0 + for end, spans, metadata_type in [ + (128, [(32, 111)], TritonAttentionMetadata), + (129, [(32, 111)], causal_metadata_type), + (145, [(32, 111)], causal_metadata_type), + (209, [(32, 111), (160, 207)], TritonAttentionMetadata), + ]: + qlen = end - previous + starts = torch.tensor([0, qlen], dtype=torch.int32) + lengths = torch.tensor([end], dtype=torch.int32) + common = CommonAttentionMetadata( + query_start_loc=starts.to(DEVICE), + query_start_loc_cpu=starts, + seq_lens=lengths.to(DEVICE), + seq_lens_cpu_upper_bound=lengths, + num_reqs=1, + num_actual_tokens=qlen, + max_query_len=qlen, + max_seq_len=end, + block_table_tensor=table, + slot_mapping=torch.arange(previous, end, device=DEVICE), + mm_req_doc_ranges={0: spans}, + ) + metadata = builder.build(0, common) + assert isinstance(metadata, metadata_type) + if isinstance(metadata, FlashInferMetadata): + if qlen == 1: + assert isinstance(metadata.decode, FlashInferTrtllmAPIDecode) + else: + assert isinstance(metadata.prefill, TRTLLMPrefill) + query = torch.randn(qlen, 32, head_size, dtype=DTYPE, device=DEVICE) + if cache_dtype == "fp8_e4m3": + query = query.to(torch.float8_e4m3fn).to(DTYPE) + key, value = keys[previous:end], values[previous:end] + output = torch.empty_like(query) + layer.impl.do_kv_cache_update(layer, key, value, cache, common.slot_mapping) + layer.impl.forward(layer, query, key, value, cache, metadata, output) + expected = _dense_reference( + query, + [keys[:end]], + [values[:end]], + [qlen], + [end], + {0: spans}, + window, + head_size**-0.5, + mm_clamp_sw=window or 0, + ) + torch.testing.assert_close(output.float(), expected, atol=2e-2, rtol=2e-2) + if qlen == 1: + metadata = builder.build_for_cudagraph_capture(common) + assert isinstance(metadata, metadata_type) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + layer.impl.forward( + layer, query, key, value, cache, metadata, output + ) + output.zero_() + graph.replay() + torch.testing.assert_close( + output.float(), expected, atol=2e-2, rtol=2e-2 + ) + previous = end diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 0bff2d537eb0..e7ff2b6d5d0b 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -159,6 +159,7 @@ def test_dflash2_model_decoder_layer_cls(monkeypatch): model_config=SimpleNamespace( dtype=torch.float32, is_mm_prefix_lm=False, + rswa_window=None, ), kernel_config=SimpleNamespace( linear_backend="auto", diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 1e5a5a3ad02b..6d9444567a04 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -87,6 +87,7 @@ def _get_backend_priorities( kv_cache_dtype: CacheDType | None = None, use_non_causal: bool = False, head_size: int | None = None, + use_mm_prefix: bool = False, ) -> list[AttentionBackendEnum]: """Get backend priorities with lazy import to avoid circular dependency.""" from vllm.utils.torch_utils import is_quantized_kv_cache @@ -156,6 +157,7 @@ def _get_backend_priorities( # So prefer FlashAttention when non-causal on SM100f. if device_capability.major == 10 and not use_non_causal: return [ + *([AttentionBackendEnum.TRITON_FLASHINFER] if use_mm_prefix else []), AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN, @@ -164,6 +166,11 @@ def _get_backend_priorities( ] else: return [ + *( + [AttentionBackendEnum.TRITON_FLASH_ATTN] + if device_capability.major == 9 and use_mm_prefix + else [] + ), AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.TRITON_ATTN, @@ -394,6 +401,7 @@ def get_valid_backends( kv_cache_dtype=attn_selector_config.kv_cache_dtype, use_non_causal=attn_selector_config.use_non_causal, head_size=attn_selector_config.head_size, + use_mm_prefix=attn_selector_config.use_mm_prefix, ) for priority, backend in enumerate(backend_priorities): try: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index c427d3a1de87..01599e9e9c72 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -185,6 +185,10 @@ def supports_per_head_quant_scales(cls) -> bool: def supports_sliding_window(cls) -> bool: return False + @classmethod + def supports_rswa(cls) -> bool: + return False + @classmethod def supports_non_causal(cls) -> bool: """Check if backend supports non-causal (bidirectional) attention @@ -287,6 +291,7 @@ def validate_configuration( use_pcp: bool = False, use_adaptive_verification: bool = False, use_dcp: bool = False, + use_rswa: bool = False, ) -> list[str]: invalid_reasons = [] if not cls.supports_head_size(head_size): @@ -321,6 +326,8 @@ def validate_configuration( invalid_reasons.append(f"attention type {attn_type} not supported") if has_sliding_window and not cls.supports_sliding_window(): invalid_reasons.append("sliding window not supported") + if use_rswa and not cls.supports_rswa(): + invalid_reasons.append("R-SWA not supported") if use_non_causal and not cls.supports_non_causal(): invalid_reasons.append("non-causal attention not supported") if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp(): @@ -854,6 +861,9 @@ def __new__(cls, *args, **kwargs): def process_weights_after_loading(self, act_dtype: torch.dtype): pass + def get_impl_variants(self) -> tuple["AttentionImplBase", ...]: + return (self,) + def prepare_for_batch(self, attn_metadata: T | None) -> None: """Prepare implementation-specific state for the current batch.""" diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py new file mode 100644 index 000000000000..d091a29d8793 --- /dev/null +++ b/vllm/v1/attention/backends/composite.py @@ -0,0 +1,509 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reusable attention backend composition and routing policies.""" + +import functools +import math +from typing import ClassVar, Protocol, cast + +import torch + +from vllm.config import VllmConfig +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionImpl, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts +from vllm.v1.kv_cache_interface import KVCacheSpec + + +def requires_mm_prefix( + metadata: CommonAttentionMetadata, *, unclamped_window: bool = False +) -> bool: + ranges = metadata.mm_req_doc_ranges + # With one query per request there are no future keys to unmask. This + # also keeps single-token prefills on the captured FlashInfer path. + if not ranges or (metadata.max_query_len <= 1 and not unclamped_window): + return False + seq_lens = metadata.seq_lens_cpu_upper_bound + assert seq_lens is not None + starts = metadata.query_start_loc_cpu + for req_idx, spans in ranges.items(): + query_len = int(starts[req_idx + 1] - starts[req_idx]) + if query_len == 0 or (query_len == 1 and not unclamped_window): + continue + end = int(seq_lens[req_idx]) + begin = end - query_len + # An unclamped image can also expose keys behind the sliding window. + if unclamped_window and any( + start < stop and start < end and stop >= begin for start, stop in spans + ): + return True + # Prefill lengths are exact. Historical image ranges must not divert + # later text/decode queries to the image backend. + if any( + start < stop and start < end - 1 and stop > begin for start, stop in spans + ): + return True + return False + + +def _has_unclamped_window(layers) -> bool: + return any( + getattr(layer, "sliding_window", None) is not None + and not getattr(layer, "mm_prefix_clamp_sliding_window", False) + for layer in layers + ) + + +class CompositeAttentionRouting(Protocol): + """Select a child backend and declare routing constraints for composition.""" + + capture_variant: ClassVar[int] + """Child used during graph capture; replay must select the same child.""" + + def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ... + + def select(self, metadata: CommonAttentionMetadata) -> int: + """Return 0 for the first child or 1 for the second child.""" + ... + + @staticmethod + def get_cudagraph_support( + vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec + ) -> AttentionCGSupport: + """Limit graph support to batches whose routing stays invariant.""" + ... + + @staticmethod + def variant_uses_mm_prefix(variant: int) -> bool: + """Whether this child must support the model's image-mask requirement.""" + ... + + +class MMPrefixAttentionRouting(CompositeAttentionRouting): + """Route image masks to variant zero and causal queries to variant one.""" + + def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: + layers = vllm_config.compilation_config.static_forward_context + self.unclamped_window = _has_unclamped_window( + layers[name] for name in layer_names + ) + + def select(self, metadata: CommonAttentionMetadata) -> int: + return int( + not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window) + ) + + capture_variant: ClassVar[int] = 1 + + @staticmethod + def get_cudagraph_support( + vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec + ) -> AttentionCGSupport: + if _has_unclamped_window( + vllm_config.compilation_config.static_forward_context.values() + ): + return AttentionCGSupport.NEVER + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + + @staticmethod + def variant_uses_mm_prefix(variant: int) -> bool: + return variant == 0 + + +def _intersect_kernel_block_sizes( + first: type[AttentionBackend], second: type[AttentionBackend] +): + from vllm.v1.attention.backend import MultipleOf + + result: list[int | MultipleOf] = [] + seen: set[tuple[type, int]] = set() + for lhs in first.get_supported_kernel_block_sizes(): + for rhs in second.get_supported_kernel_block_sizes(): + candidate: int | MultipleOf | None = None + if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf): + candidate = MultipleOf(math.lcm(lhs.base, rhs.base)) + elif isinstance(lhs, MultipleOf) and isinstance(rhs, int): + candidate = rhs if rhs % lhs.base == 0 else None + elif isinstance(lhs, int) and isinstance(rhs, MultipleOf): + candidate = lhs if lhs % rhs.base == 0 else None + elif lhs == rhs: + candidate = lhs + if candidate is None: + continue + value = candidate.base if isinstance(candidate, MultipleOf) else candidate + key = (type(candidate), value) + if key not in seen: + result.append(candidate) + seen.add(key) + return result + + +@functools.cache +def create_composite_attention_backend( + general_backend: type[AttentionBackend], + causal_backend: type[AttentionBackend], + *, + name: str, + backend_name: str, + module: str, + routing_policy: type[CompositeAttentionRouting], + head_sizes: tuple[int, ...] = (), + kernel_block_sizes: tuple[int, ...] = (), + device_major: int | None = None, +) -> type[AttentionBackend]: + """Compose two backends with shared storage and an explicit routing policy.""" + if general_backend.full_cls_name() == causal_backend.full_cls_name(): + return general_backend + general_impl_cls = general_backend.get_impl_cls() + causal_impl_cls = causal_backend.get_impl_cls() + general_builder_cls = general_backend.get_builder_cls() + causal_builder_cls = causal_backend.get_builder_cls() + if ( + general_backend.forward_includes_kv_cache_update + or causal_backend.forward_includes_kv_cache_update + ): + raise ValueError("Composite backends require separate KV cache updates") + + class CompositeAttentionImpl(AttentionImpl): + def __init__(self, *args, **kwargs) -> None: + self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs)) + try: + self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs)) + except TypeError as e: + raise ValueError( + f"Causal attention backend {causal_backend.get_name()} " + f"does not accept the attention arguments required by " + f"the general backend {general_backend.get_name()}: {e}" + ) from e + for name in ( + "num_heads", + "num_kv_heads", + "head_size", + "scale", + "kv_cache_dtype", + ): + if hasattr(self.general_impl, name): + setattr(self, name, getattr(self.general_impl, name)) + self.supports_quant_query_input = ( + self.general_impl.supports_quant_query_input + and self.causal_impl.supports_quant_query_input + ) + self.supports_pcp = self.general_impl.supports_pcp + self.supports_dcp = self.causal_impl.supports_dcp + self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode + self.lse_base_on_e = self.causal_impl.lse_base_on_e + self.need_to_return_lse_for_decode = ( + self.causal_impl.need_to_return_lse_for_decode + ) + self.supports_mtp_with_cp_non_trivial_interleave_size = ( + self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size + ) + + def get_impl_for_metadata(self, attn_metadata): + return ( + self.causal_impl + if getattr(attn_metadata, "_attention_backend_variant", 0) + else self.general_impl + ) + + def get_impl_variants(self): + return self.general_impl, self.causal_impl + + def forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale=None, + output_block_scale=None, + ): + impl = self.get_impl_for_metadata(attn_metadata) + return impl.forward( + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, + ) + + def process_weights_after_loading(self, act_dtype: torch.dtype): + for impl in self.get_impl_variants(): + impl.process_weights_after_loading(act_dtype) + + def fused_output_quant_supported(self, quant_key): + return all( + impl.fused_output_quant_supported(quant_key) + for impl in self.get_impl_variants() + ) + + def fused_qk_norm_rope_kvcache_supported(self): + return self.general_impl.fused_qk_norm_rope_kvcache_supported() + + def do_qk_norm_rope_kvcache_update(self, *args, **kwargs): + method = self.general_impl.do_qk_norm_rope_kvcache_update + return method(*args, **kwargs) + + def fused_rope_kvcache_supported(self): + return self.general_impl.fused_rope_kvcache_supported() + + def do_rope_and_kv_cache_update(self, *args, **kwargs): + method = self.general_impl.do_rope_and_kv_cache_update + return method(*args, **kwargs) + + def do_kv_cache_update(self, *args, **kwargs): + method = self.general_impl.do_kv_cache_update # type: ignore[attr-defined] + return method(*args, **kwargs) + + class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder): + requires_block_table_width = ( + general_builder_cls.requires_block_table_width + or causal_builder_cls.requires_block_table_width + ) + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + args = (kv_cache_spec, layer_names, vllm_config, device) + # Only the variant that asked for the extra ctor kwargs receives them. + self.general_builder = general_builder_cls( + *args, + **(kwargs if general_builder_cls.requires_block_table_width else {}), + ) + self.causal_builder = causal_builder_cls( + *args, + **(kwargs if causal_builder_cls.requires_block_table_width else {}), + ) + self.reorder_batch_threshold = min( + ( + builder.reorder_batch_threshold + for builder in self._builders + if builder.reorder_batch_threshold is not None + ), + default=None, + ) + self.routing = routing_policy(layer_names, vllm_config) + + def set_kernel_block_size(self, kernel_block_size): + super().set_kernel_block_size(kernel_block_size) + for builder in self._builders: + builder.set_kernel_block_size(kernel_block_size) + + @classmethod + def get_cudagraph_support(cls, vllm_config, kv_cache_spec): + return min( + general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec), + causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec), + routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec), + key=lambda support: support.value, + ) + + def _builder(self, common_attn_metadata): + backend_variant = self.routing.select(common_attn_metadata) + builder = self.causal_builder if backend_variant else self.general_builder + return builder, backend_variant + + def _build(self, method, common_attn_metadata, **kwargs): + builder, backend_variant = self._builder(common_attn_metadata) + metadata = getattr(builder, method)( + common_attn_metadata=common_attn_metadata, **kwargs + ) + metadata._attention_backend_variant = backend_variant + return metadata + + def build(self, common_prefix_len, common_attn_metadata, fast_build=False): + return self._build( + "build", + common_attn_metadata, + common_prefix_len=common_prefix_len, + fast_build=fast_build, + ) + + def build_for_cudagraph_capture(self, common_attn_metadata): + variant = routing_policy.capture_variant + metadata = self._builders[variant].build_for_cudagraph_capture( + common_attn_metadata + ) + metadata._attention_backend_variant = variant + return metadata + + @property + def _builders(self): + return self.general_builder, self.causal_builder + + # Forward the optional workspace protocol so that a wrapped builder + # which allocates one (FlashInfer) still joins the runner's cross-group + # sharing instead of allocating a second buffer behind the composite. + def _get_workspace_buffer(self): + for builder in self._builders: + if hasattr(builder, "_get_workspace_buffer"): + workspace = builder._get_workspace_buffer() + # The runner only calls the getter for the group that + # provides the buffer, so hand it to the sibling here. + self.set_workspace_buffer(workspace) + return workspace + return None + + def set_workspace_buffer(self, workspace_buffer): + for builder in self._builders: + if hasattr(builder, "set_workspace_buffer"): + builder.set_workspace_buffer(workspace_buffer) + + class CompositeAttentionBackend(AttentionBackend): + forward_includes_kv_cache_update = False + general_backend_cls = general_backend + causal_backend_cls = causal_backend + + @staticmethod + def get_name(): + return backend_name + + @staticmethod + def get_impl_cls(): + return CompositeAttentionImpl + + @staticmethod + def get_builder_cls(): + return CompositeAttentionMetadataBuilder + + @staticmethod + def get_supported_kernel_block_sizes(): + if kernel_block_sizes: + return [ + size + for size in kernel_block_sizes + if all( + backend.supports_block_size(size) + for backend in (general_backend, causal_backend) + ) + ] + return _intersect_kernel_block_sizes(general_backend, causal_backend) + + @classmethod + def supports_block_size(cls, block_size): + return bool( + cls.get_supported_kernel_block_sizes() + ) and super().supports_block_size(block_size) + + @classmethod + def get_supported_head_sizes(cls): + return list(head_sizes) + + @classmethod + def supports_head_size(cls, head_size): + return (not head_sizes or head_size in head_sizes) and all( + backend.supports_head_size(head_size) + for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_compute_capability(cls, capability): + return (device_major is None or capability.major == device_major) and all( + backend.supports_compute_capability(capability) + for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_dtype(cls, dtype): + return all( + backend.supports_dtype(dtype) + for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_kv_cache_dtype(cls, dtype): + return all( + backend.supports_kv_cache_dtype(dtype) + for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_mm_prefix(cls): + return all( + not routing_policy.variant_uses_mm_prefix(variant) + or backend.supports_mm_prefix() + for variant, backend in enumerate((general_backend, causal_backend)) + ) + + @classmethod + def supports_sliding_window(cls): + return all( + backend.supports_sliding_window() + for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_rswa(cls): + return all( + backend.supports_rswa() for backend in (general_backend, causal_backend) + ) + + @classmethod + def supports_pcp(cls): + return False + + @classmethod + def supports_dcp(cls): + return False + + @classmethod + def supports_device_cpu_query_lens_mismatch(cls): + return False + + @classmethod + def supports_combination( + cls, + head_size, + dtype, + kv_cache_dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ): + for variant, backend in enumerate((general_backend, causal_backend)): + reason = backend.supports_combination( + head_size, + dtype, + kv_cache_dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant), + device_capability, + ) + if reason: + return f"{backend.get_name()}: {reason}" + return None + + @classmethod + def customize_spec(cls, spec): + general_spec = general_backend.customize_spec(spec) + if general_spec != causal_backend.customize_spec(spec): + raise ValueError("Composite backends require the same KV cache spec") + return general_spec + + @classmethod + def supported_kv_cache_layouts(cls): + return tuple( + get_supported_kv_cache_layouts((general_backend, causal_backend)) + ) + + CompositeAttentionBackend.__name__ = name + CompositeAttentionBackend.__qualname__ = name + CompositeAttentionBackend.__module__ = module + return CompositeAttentionBackend diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 9ff543024fde..9109a9c6d7ab 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -129,6 +129,10 @@ def get_preferred_block_size(cls, default_block_size: int) -> int: def get_name() -> str: return "FLASH_ATTN" + @classmethod + def supports_rswa(cls) -> bool: + return True + @classmethod def supports_sliding_window(cls) -> bool: return True diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 7529d8d20f8e..70eded1d093f 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -105,6 +105,10 @@ class FlexAttentionBackend(AttentionBackend): def get_name() -> str: return "FLEX_ATTENTION" + @classmethod + def supports_rswa(cls) -> bool: + return True + @classmethod def supports_sliding_window(cls) -> bool: return True diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index c1c63305767c..085f2358fabb 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -49,6 +49,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): TRITON_ATTN_DIFFKV = ( "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend" ) + TRITON_FLASH_ATTN = ( + "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend" + ) ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" ROCM_AITER_TRITON_MLA = ( @@ -63,6 +66,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend" TORCH_SDPA = "" # this tag is only used for ViT FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" + TRITON_FLASHINFER = ( + "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend" + ) FLASHINFER_MLA = ( "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend" ) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 48dd30d66635..b7e2a41d00e5 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -323,6 +323,10 @@ def supports_non_causal(cls) -> bool: def get_name() -> str: return "TRITON_ATTN" + @classmethod + def supports_rswa(cls) -> bool: + return True + @classmethod def supports_sliding_window(cls) -> bool: return True diff --git a/vllm/v1/attention/backends/triton_flash_attn.py b/vllm/v1/attention/backends/triton_flash_attn.py new file mode 100644 index 000000000000..0853b291cda2 --- /dev/null +++ b/vllm/v1/attention/backends/triton_flash_attn.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton image-mask and FlashAttention causal attention composite.""" + +import torch + +from vllm.config.cache import CacheDType +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.composite import ( + MMPrefixAttentionRouting, + create_composite_attention_backend, +) +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + + +class _FA4Backend(FlashAttentionBackend): + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + reason = super().supports_combination( + head_size, + dtype, + kv_cache_dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ) + if reason is not None: + return reason + if ( + get_flash_attn_version( + head_size=head_size, + has_sinks=has_sink, + kv_cache_block_size=block_size, + supports_fa4_hd256=True, + ) + != 4 + ): + return "causal route requires FlashAttention v4" + return None + + +TritonFlashAttentionBackend = create_composite_attention_backend( + TritonAttentionBackend, + _FA4Backend, + name="TritonFlashAttentionBackend", + backend_name="TRITON_FLASH_ATTN", + module=__name__, + routing_policy=MMPrefixAttentionRouting, + head_sizes=(256, 512), + device_major=9, +) diff --git a/vllm/v1/attention/backends/triton_flashinfer.py b/vllm/v1/attention/backends/triton_flashinfer.py new file mode 100644 index 000000000000..07912275913e --- /dev/null +++ b/vllm/v1/attention/backends/triton_flashinfer.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton image-mask and FlashInfer causal attention composite.""" + +from vllm.v1.attention.backends.composite import ( + MMPrefixAttentionRouting, + create_composite_attention_backend, +) +from vllm.v1.attention.backends.flashinfer import FlashInferBackend +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + +TritonFlashInferBackend = create_composite_attention_backend( + TritonAttentionBackend, + FlashInferBackend, + name="TritonFlashInferBackend", + backend_name="TRITON_FLASHINFER", + module=__name__, + routing_policy=MMPrefixAttentionRouting, + head_sizes=(256, 512), + kernel_block_sizes=(64,), + device_major=10, +) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index bc127ccb8b08..a3adf70d2588 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -8,6 +8,7 @@ TYPE_CHECKING, Any, Protocol, + TypeVar, cast, ) @@ -45,6 +46,19 @@ PAD_SLOT_ID = -1 NULL_BLOCK_ID = 0 + +ImplT = TypeVar("ImplT", bound=AttentionImpl) + + +def find_attention_impl_variant( + impl: AttentionImpl, impl_cls: type[ImplT] +) -> ImplT | None: + for variant in impl.get_impl_variants(): + if isinstance(variant, impl_cls): + return cast(ImplT, variant) + return None + + _LN_2 = math.log(2.0) @@ -342,8 +356,8 @@ def get_per_layer_parameters( per_layer_params: dict[str, PerLayerParameters] = {} for key, layer in layers.items(): - impl = layer.impl - assert isinstance(impl, cls_) + impl = find_attention_impl_variant(layer.impl, cls_) + assert impl is not None # Infer hyperparameters from the attention layer window_size = getattr(impl, "sliding_window", None) @@ -1013,6 +1027,9 @@ def __init__(self, metadata, common_attn_metadata): common_attn_metadata.logits_indices_padded ) self.num_logits_indices = common_attn_metadata.num_logits_indices + self._attention_backend_variant = getattr( + metadata, "_attention_backend_variant", 0 + ) return KVSharingFastPrefillAttentionMetadata(metadata, common_attn_metadata) diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index c9631fd45715..286fdcb8dfd2 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple): use_pcp: bool = False use_adaptive_verification: bool = False use_dcp: bool = False + use_rswa: bool = False def __repr__(self): return ( @@ -55,7 +56,8 @@ def __repr__(self): f"use_kv_connector={self.use_kv_connector}, " f"use_adaptive_verification={self.use_adaptive_verification}, " f"use_pcp={self.use_pcp}, " - f"use_dcp={self.use_dcp})" + f"use_dcp={self.use_dcp}, " + f"use_rswa={self.use_rswa})" ) @@ -168,6 +170,10 @@ def get_attn_backend( use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1, use_adaptive_verification=use_adaptive_verification, use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1, + use_rswa=( + vllm_config.model_config is not None + and vllm_config.model_config.rswa_window is not None + ), ) # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over From b7e0cdac5d11e3a1b32745079cf5591a685aa721 Mon Sep 17 00:00:00 2001 From: Pavel Zakharov Date: Sun, 13 Sep 2026 07:58:09 -0700 Subject: [PATCH 133/774] [Bugfix] Detect unloaded NVFP4 weight scales with a NaN sentinel (#52501) Signed-off-by: pavelzak Co-authored-by: Claude Sonnet 5 Co-authored-by: OpenAI Codex --- .../layers/quantization/modelopt.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 48482acd75bb..e3e8b373696f 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1942,7 +1942,9 @@ def create_weights(self, layer, role, ctx, shapes, wl) -> None: PerTensorScaleParameter, wl, ) - # Per-block (group_size) weight scale. + # Per-block (group_size) weight scale. Initialized to NaN (rather than + # left as torch.empty garbage) so an unloaded scale is unambiguously + # detectable in process() below, matching the sentinel pattern in #45320. self.register_params( layer, "weight_scale", @@ -1955,11 +1957,26 @@ def create_weights(self, layer, role, ctx, shapes, wl) -> None: wl, input_dim=1, output_dim=0, + init=float("nan"), ) def process(self, layer, role) -> None: if role is not WEIGHT: self.reject(role) + # Sanity-check: weight_scale must have been overwritten by the weight + # loader. A value still equal to the NaN sentinel it was created with + # means the FP4 weights were never actually loaded (e.g. the + # checkpoint stores this layer as BF16 and the weight loader silently + # skipped it). + if torch.isnan(layer.weight_scale.float()).any(): + raise RuntimeError( + f"NVFP4 weight_scale for layer " + f"{getattr(layer, 'name', repr(layer))!r} was never loaded " + "(still NaN). The checkpoint likely stores this layer as " + "BF16 (not FP4). Fix: pass quant_config=None when " + "constructing this layer, or add it to the quantization " + "ignore list." + ) if torch.unique(layer.weight_scale_2).numel() != 1: logger.warning_once( "In NVFP4 linear, the global weight scale differs across " From f09c52a587f205963d3a2b8f1c42820a831ee4b0 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 13 Sep 2026 12:56:02 -0400 Subject: [PATCH 134/774] [ROCm][Performance] Avoid blocking MiniMax M3 scalar upload (#56170) Signed-off-by: fai Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/minimax_m3/common/sparse_attention.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 50c03dd9bf0a..f385f82067fd 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -314,7 +314,10 @@ def build( prefill_cu_seqlens_k = torch.empty( num_prefills + 1, dtype=torch.int32, device=seq_lens.device ) - prefill_cu_seqlens_k[0] = 0 + if current_platform.is_rocm(): + prefill_cu_seqlens_k[:1].zero_() + else: + prefill_cu_seqlens_k[0] = 0 torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:]) prefill_cu_seqlens_q = ( query_start_loc[num_decodes:] - num_decode_tokens From dec0b5d63fc42dc1d1789caf572195c73dd86ea6 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 13 Sep 2026 10:05:29 -0700 Subject: [PATCH 135/774] [Bugfix] Skip Triton autotune inspection without Triton (#56676) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- .../test_jit_warmup_triton_launcher.py | 20 +++++++++++++++++++ .../warmup/jit_warmup_triton_helper.py | 4 +++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index 2e81cb6000f8..725a383833a6 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -207,6 +207,26 @@ def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]: launch(first, 1, 7, stale_constexpr=True) +def test_triton_kernel_decorator_without_triton( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(jit_warmup_triton_helper, "HAS_TRITON", False) + monkeypatch.setattr(jit_warmup_triton_helper, "triton", object()) + kernel = _FakeTritonKernel() + + def warmup_inputs() -> dict[str, Any]: + return dict(first="warmup", second=1) + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def launch(first: str, second: int) -> LaunchSpec: + return (2,), dict(CONST=7) + + launch("runtime", 2) + assert kernel.runtime_calls == [ + ((2,), (), {"first": "runtime", "second": 2, "CONST": 7}) + ] + + def test_triton_kernel_decorator_exhausts_large_ranges_before_deduplication( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index dab055902aee..198e6fdbe2f4 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -15,7 +15,7 @@ get_function_source_node, ) from vllm.platforms import current_platform -from vllm.triton_utils import triton +from vllm.triton_utils import HAS_TRITON, triton CompileKeyT = TypeVar("CompileKeyT") P = ParamSpec("P") @@ -428,6 +428,8 @@ def _provider_cases( def _is_autotuned(kernel: Any) -> bool: + if not HAS_TRITON: + return False Autotuner = triton.runtime.autotuner.Autotuner while kernel is not None: From b6e2aa748b2bef8d72f2003cd881a77ac47dcb57 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 13:34:09 -0400 Subject: [PATCH 136/774] [CI] Initialize warmup registry in V2 QSA runner fixture (#56707) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- tests/v1/worker/test_gpu_model_runner_v2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index c9ce409f5cb9..86ff1a074538 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -8,6 +8,7 @@ import torch import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, @@ -37,6 +38,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): parallel_config=parallel_config, cache_config=SimpleNamespace(mamba_cache_mode="none"), ) + runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config) runner.model_state = SimpleNamespace( get_additional_cg_support=lambda: (), num_new_sampled_tokens_per_step=1, From 71888f507ad5df5af2fae2ab39a6c1b0cc465eb1 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 13 Sep 2026 12:06:13 -0700 Subject: [PATCH 137/774] [Perf] Avoid redundant conversions in FP8 dummy initialization (#56688) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- vllm/model_executor/model_loader/weight_utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 4cb8249e48f6..abe24fa06eca 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1380,10 +1380,9 @@ def initialize_single_dummy_weight( generator.manual_seed(seed) if torch.finfo(param.data.dtype).bits < 16: # uniform_ doesn't support < 16-bit datatypes (FP8) - dtype = param.data.dtype - tmp_param = param.data.to(torch.float16) - tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype) - param.data.copy_(tmp_param) + tmp_param = torch.empty_like(param, dtype=torch.float16) + tmp_param.uniform_(low, high, generator=generator) + param.copy_(tmp_param) else: param.uniform_(low, high, generator=generator) From 7fe8fc803d904795fd10b785776e20181aee64c5 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 16:23:15 -0400 Subject: [PATCH 138/774] [NIXL][PCP][DCP] Expose PCP producer KV shards as transfer ranks (#56645) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../kv_connector/unit/test_nixl_connector.py | 73 +++++++++++++++---- .../unit/test_nixl_push_connector.py | 15 ++-- tests/v1/kv_connector/unit/test_tp_mapping.py | 1 + tests/v1/kv_connector/unit/utils.py | 2 + vllm/config/vllm.py | 12 +-- .../kv_connector/v1/nixl/base_scheduler.py | 6 ++ .../kv_connector/v1/nixl/base_worker.py | 28 ++++--- .../kv_connector/v1/nixl/connector.py | 13 +++- .../kv_connector/v1/nixl/pull_scheduler.py | 2 +- .../kv_connector/v1/nixl/pull_worker.py | 2 +- .../kv_connector/v1/nixl/push_scheduler.py | 3 +- .../kv_connector/v1/nixl/push_worker.py | 2 +- vllm/v1/worker/gpu_worker.py | 9 +++ 13 files changed, 124 insertions(+), 44 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 3fcbda4c6255..9b54f626f265 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -580,19 +580,38 @@ def _nixl_handshake( class TestNixlHandshake: - @pytest.mark.parametrize("pcp_rank", [0, 1]) + @pytest.mark.parametrize( + ("pcp_rank", "pcp_size", "dcp_size", "expected_tracked"), + [ + (0, 2, 1, True), + (1, 2, 1, False), + (0, 2, 2, True), + (1, 2, 2, True), + (0, 4, 4, True), + (1, 4, 4, True), + (2, 4, 4, True), + (3, 4, 4, True), + ], + ) @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) - def test_pcp_producer_uses_canonical_replica( - self, default_vllm_config, dist_init, pcp_rank + def test_pcp_producer_exposes_dcp_shards_or_canonical_replica( + self, + default_vllm_config, + dist_init, + pcp_rank, + pcp_size, + dcp_size, + expected_tracked, ): - """Only PCP rank zero publishes, but every rank reports completion.""" + """Replicated PCP is canonicalized; PCP-DCP publishes every shard.""" from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend vllm_config = create_vllm_config(kv_role="kv_producer") - vllm_config.parallel_config.prefill_context_parallel_size = 2 + vllm_config.parallel_config.prefill_context_parallel_size = pcp_size + vllm_config.parallel_config.decode_context_parallel_size = dcp_size with ( patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl." @@ -614,13 +633,15 @@ def test_pcp_producer_uses_canonical_replica( worker = connector.connector_worker assert worker is not None assert worker.pcp_rank == pcp_rank + assert worker.pcp_dcp_sharded is (dcp_size > 1) + assert worker.transfer_tp_rank == (pcp_rank if dcp_size > 1 else 0) + assert worker.transfer_tp_size == (pcp_size if dcp_size > 1 else 1) req_id = "req" metadata = NixlConnectorMetadata() metadata.reqs_in_batch.add(req_id) metadata.reqs_to_send[req_id] = time.perf_counter() + 10 worker.start_load_kv(metadata) - expected_tracked = pcp_rank == 0 assert (req_id in worker._reqs_to_process) == expected_tracked assert (req_id in worker._reqs_to_send) == expected_tracked @@ -628,22 +649,22 @@ def test_pcp_producer_uses_canonical_replica( worker.xfer_handshake_metadata = payload worker.transfer_topo = MagicMock() worker._get_new_notifs = MagicMock( - side_effect=lambda: {"sent"} if pcp_rank == 0 else set() + side_effect=lambda: {"sent"} if expected_tracked else set() ) - expected_payload = payload if pcp_rank == 0 else None + expected_payload = payload if expected_tracked else None assert connector.get_handshake_metadata() is expected_payload done_sending, done_recving = connector.get_finished(set()) - assert done_sending == ({"sent"} if pcp_rank == 0 else {req_id}) + assert done_sending == ({"sent"} if expected_tracked else {req_id}) assert done_recving == set() - if pcp_rank > 0: + if not expected_tracked: assert connector.get_finished(set()) == (set(), set()) worker.get_transfer_results = MagicMock( return_value=KVConnectorTransferResults(finished_sending={"sent"}) ) results = connector.get_transfer_results(set()) - assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set()) + assert results.finished_sending == ({"sent"} if expected_tracked else set()) @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", @@ -3354,9 +3375,15 @@ def test_transfer_mode_changes_compatibility_hash(): @pytest.mark.skip_global_cleanup -def test_scheduler_advertises_transfer_mode(): - # Each scheduler advertises its transfer mode in kv_transfer_params so an - # external router can route pull (READ) vs push (WRITE) producers. +@pytest.mark.parametrize("mode", ["pull", "push"]) +@pytest.mark.parametrize( + "tp_size,pcp_size,dcp_size,transfer_tp_size", + [(1, 1, 1, 1), (4, 1, 1, 4), (4, 1, 4, 4), (1, 4, 1, 1), (1, 4, 4, 4)], +) +def test_scheduler_advertises_transfer_topology( + mode, tp_size, pcp_size, dcp_size, transfer_tp_size +): + """The consumer must address every distinct producer KV shard.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( NixlPullConnectorScheduler, ) @@ -3364,8 +3391,22 @@ def test_scheduler_advertises_transfer_mode(): NixlPushConnectorScheduler, ) - assert NixlPullConnectorScheduler._TRANSFER_MODE == "pull" - assert NixlPushConnectorScheduler._TRANSFER_MODE == "push" + config = create_vllm_config(kv_role="kv_producer") + config.parallel_config.tensor_parallel_size = tp_size + config.parallel_config.prefill_context_parallel_size = pcp_size + config.parallel_config.decode_context_parallel_size = dcp_size + cls = NixlPullConnectorScheduler if mode == "pull" else NixlPushConnectorScheduler + scheduler = cls(config, "prefiller", make_kv_cache_config(block_size=16)) + request = create_request(request_id=1, num_tokens=32, do_remote_decode=True) + request.status = RequestStatus.FINISHED_LENGTH_CAPPED + try: + delay, params = scheduler.request_finished(request, ([0, 1],)) + assert delay + assert params["transfer_mode"] == mode + assert params["tp_size"] == transfer_tp_size + assert params["dcp_size"] == dcp_size + finally: + scheduler.shutdown() @pytest.mark.parametrize( diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index ce233bbb603b..1cdafcfe482f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -358,6 +358,7 @@ def fresh(cls) -> _StubWriterWorker: w.consumer_notification_counts_by_req = defaultdict(int) w.tp_rank = 0 w.pcp_rank = 0 + w.pcp_dcp_sharded = False w.world_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} @@ -545,9 +546,11 @@ def test_start_load_kv_enqueues_to_writer(self): assert w._push_writer_wake.is_set() assert w.start_push_calls == [] - def test_noncanonical_pcp_rank_skips_producer_work(self): + @pytest.mark.parametrize("sharded", [False, True]) + def test_noncanonical_pcp_rank_only_pushes_distinct_shards(self, sharded): w = _StubWriterWorker.fresh() w.pcp_rank = 1 + w.pcp_dcp_sharded = sharded w._send_heartbeats = MagicMock() meta = NixlConnectorMetadata() @@ -557,11 +560,11 @@ def test_noncanonical_pcp_rank_skips_producer_work(self): w.start_load_kv(meta) - assert w._finished_blocks_inbox.empty() - assert "req" not in w._reqs_to_process - assert "req" not in w._reqs_to_send - assert not w._push_writer_wake.is_set() - w._send_heartbeats.assert_not_called() + assert w._finished_blocks_inbox.empty() is (not sharded) + assert ("req" in w._reqs_to_process) is sharded + assert ("req" in w._reqs_to_send) is sharded + assert w._push_writer_wake.is_set() is sharded + assert w._send_heartbeats.called is sharded # The P→D handshake must run on the base worker's background executor, never diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 52067bc2fa6f..3ddc568a524d 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -75,6 +75,7 @@ def test_source_ranks_p_gt_d(self): @pytest.mark.parametrize( "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks", [ + (0, 1, 8, 1, 8, tuple(range(8))), (0, 4, 4, 1, 4, (0, 1, 2, 3)), (2, 4, 4, 4, 4, (2,)), (0, 2, 4, 2, 4, (0, 2)), diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 07a40dc4ec00..1ee1e817766a 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -548,6 +548,7 @@ def make_nixl_scheduler( sched._reqs_need_save = {} sched.use_host_buffer = False sched.engine_id = "test-engine" + sched.transfer_tp_size = 1 sched.side_channel_host = "localhost" sched.side_channel_port = 5555 sched.blocks_per_sw = [] @@ -587,6 +588,7 @@ def make_nixl_push_scheduler( sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl sched.use_host_buffer = False sched.engine_id = "decode-engine" + sched.transfer_tp_size = 1 sched.side_channel_host = "127.0.0.1" sched.side_channel_port = 5600 sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7f0a3f0d8723..fbe021385524 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1329,14 +1329,14 @@ def __post_init__(self): self.kv_transfer_config is not None and self.kv_transfer_config.has_connector("NixlConnector") ): - assert self.parallel_config.prefill_context_parallel_size == 1, ( - "NIXL does not support prefill context parallelism." - ) dcp_size = self.parallel_config.decode_context_parallel_size - tp_size = self.parallel_config.tensor_parallel_size - assert dcp_size in (1, tp_size), ( + transfer_tp_size = max( + self.parallel_config.tensor_parallel_size, + self.parallel_config.prefill_context_parallel_size, + ) + assert dcp_size in (1, transfer_tp_size), ( f"decode_context_parallel_size={dcp_size} must be 1 or equal " - f"to tensor_parallel_size={tp_size} when using NixlConnector." + f"to the NIXL transfer parallel size={transfer_tp_size}." ) if self.model_config is not None: assert self.model_config.use_mla or dcp_size == 1, ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 00579c40f032..5c01d996968e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -63,6 +63,12 @@ def __init__( kv_cache_config: "KVCacheConfig", ): self.vllm_config = vllm_config + parallel_config = vllm_config.parallel_config + # TP1 PCP+DCP exposes its DCP shards as transfer ranks. + self.transfer_tp_size = max( + parallel_config.tensor_parallel_size, + parallel_config.decode_context_parallel_size, + ) self.block_size = vllm_config.cache_config.block_size self.engine_id: EngineId = engine_id self.kv_cache_config = kv_cache_config diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 17660da451a1..4edfaefa2f5f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -548,9 +548,15 @@ def __init__( self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 - # DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully - # replicated or fully sharded. A DCP rank is always derivable this way. - self.dcp_rank = self.tp_rank % self.dcp_size + # TP1 PCP+DCP owns distinct KV shards; replicated PCP uses rank zero. + self.pcp_dcp_sharded = self.pcp_size > 1 and self.dcp_size > 1 + self.transfer_tp_size = ( + self.pcp_size if self.pcp_dcp_sharded else self.world_size + ) + self.transfer_tp_rank = self.pcp_rank if self.pcp_dcp_sharded else self.tp_rank + + # MLA is either fully replicated or fully sharded across transfer ranks. + self.dcp_rank = self.transfer_tp_rank % self.dcp_size self.num_blocks = kv_cache_config.num_blocks self.enable_permute_local_kv = False @@ -781,12 +787,16 @@ def _validate_remote_parallel_config( local_dcp_size = self.dcp_size remote_pcp_size = agent_metadata.pcp_size remote_dcp_size = agent_metadata.dcp_size - if (local_pcp_size > 1 and remote_dcp_size > 1) or ( - remote_pcp_size > 1 and local_dcp_size > 1 + if remote_pcp_size > 1 and remote_dcp_size not in (1, remote_pcp_size): + raise NotImplementedError( + "Remote NixlConnector PCP+DCP does not span the full PCP group. " + f"Remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}." + ) + if (local_pcp_size > 1 and local_dcp_size == 1 and remote_dcp_size > 1) or ( + remote_pcp_size > 1 and remote_dcp_size == 1 and local_dcp_size > 1 ): raise NotImplementedError( - "NixlConnector PCP requires decode_context_parallel_size=1 " - "on both instances. " + "Replicated PCP cannot be paired with a DCP-sharded NIXL peer. " f"Local PCP/DCP={local_pcp_size}/{local_dcp_size}; " f"remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}." ) @@ -1210,8 +1220,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" self.transfer_topo = TransferTopology( - tp_rank=self.tp_rank, - tp_size=self.world_size, + tp_rank=self.transfer_tp_rank, + tp_size=self.transfer_tp_size, block_size=self.block_size, engine_id=self.engine_id, is_mla=self.use_mla, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 105617bb8dac..99c84765e9db 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -108,10 +108,15 @@ def __init__( "Consumers and kv_both require " "prefill_context_parallel_size=1." ) - if pcp_size > 1 and parallel_config.decode_context_parallel_size > 1: + dcp_size = parallel_config.decode_context_parallel_size + if ( + pcp_size > 1 + and dcp_size > 1 + and (parallel_config.tensor_parallel_size != 1 or dcp_size != pcp_size) + ): raise NotImplementedError( - "NixlConnector PCP producers currently require " - "decode_context_parallel_size=1." + "NixlConnector PCP+DCP requires TP1 with DCP spanning " + "the full PCP group." ) # TODO: Support PCP with bidirectional KV transfer by tracking separate # send and receive completion counts. @@ -243,6 +248,7 @@ def get_transfer_results( if ( self.kv_transfer_config.kv_role == "kv_producer" and self.connector_worker.pcp_rank > 0 + and not self.connector_worker.pcp_dcp_sharded ): results.finished_sending.clear() return results @@ -324,6 +330,7 @@ def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: if ( self.kv_transfer_config.kv_role == "kv_producer" and self.connector_worker.pcp_rank > 0 + and not self.connector_worker.pcp_dcp_sharded ): return None return self.connector_worker.xfer_handshake_metadata diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 9740b518a22c..446126c328ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -282,7 +282,7 @@ def request_finished( remote_request_id=request.request_id, remote_host=self.side_channel_host, remote_port=self.side_channel_port, - tp_size=self.vllm_config.parallel_config.tensor_parallel_size, + tp_size=self.transfer_tp_size, dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size, pp_size=self.vllm_config.parallel_config.pipeline_parallel_size, remote_num_tokens=remote_num_tokens, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index ceaa4fce577f..59205b943a5e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -80,7 +80,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): while not self._ready_requests.empty(): self._read_blocks_for_req(*self._ready_requests.get_nowait()) - if self.pcp_rank > 0: + if self.pcp_rank > 0 and not self.pcp_dcp_sharded: # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank # has nothing to send. Report the requests as sent right away so # the scheduler-side aggregation (world_size workers, and any diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index 53394cf9c138..c11a63bb3414 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -286,7 +286,8 @@ def request_finished( remote_request_id=request.request_id, remote_host=self.side_channel_host, remote_port=self.side_channel_port, - tp_size=self.vllm_config.parallel_config.tensor_parallel_size, + tp_size=self.transfer_tp_size, + dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size, pp_size=self.vllm_config.parallel_config.pipeline_parallel_size, remote_num_tokens=remote_num_tokens, transfer_mode=self._TRANSFER_MODE, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index c3bcd2c8f907..e506aff1f425 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -166,7 +166,7 @@ def shutdown(self) -> None: def start_load_kv(self, metadata: NixlConnectorMetadata): """Pre-process metadata; defer NIXL ops to the writer thread.""" - if self.pcp_rank > 0: + if self.pcp_rank > 0 and not self.pcp_dcp_sharded: return # D-side: track reqs waiting for P to push. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 53bc0550a524..2e84fd511a22 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -45,6 +45,7 @@ Handle, checkpoint_prepare_distributed_state, checkpoint_restore_distributed_state, + get_pcp_group, get_pp_group, get_tp_group, resume_device_comms, @@ -703,6 +704,14 @@ def get_kv_connector_handshake_metadata( pp_rank = get_pp_group().rank_in_group tp_rank = get_tp_group().rank_in_group + parallel_config = self.vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 + ): + tp_rank += ( + get_pcp_group().rank_in_group * parallel_config.tensor_parallel_size + ) return {(pp_rank, tp_rank): metadata} def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: From 586f652f8d9e841901c40634c024a2ebbfe3402f Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 16:45:17 -0400 Subject: [PATCH 139/774] [Bugfix][PCP][DCP] Respect interleave in indexer KV gather mapping (#56715) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../v1/attention/test_indexer_dcp_localize.py | 43 ++++++++++++++----- vllm/v1/attention/backends/mla/indexer.py | 10 ++++- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index ffcfcbfac1e5..eab4759640a6 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -966,34 +966,55 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp(): @pytest.mark.parametrize("dcp_world_size", [2, 4, 8]) -@pytest.mark.parametrize("req_lens", [[7], [8, 8], [1, 9], [256, 1, 2730]]) -def test_pcp_plan_deinterleave_restores_global_order(dcp_world_size, req_lens): - """The index gather must undo DCP sharding per request""" +@pytest.mark.parametrize("interleave", [1, 64]) +@pytest.mark.parametrize("rows_per_req", [1, 2]) +@pytest.mark.parametrize( + "req_lens", [[7], [8, 8], [1, 9], [63, 64, 65], [255, 256, 257], [256, 1, 2730]] +) +def test_pcp_plan_deinterleave_restores_global_order( + monkeypatch, dcp_world_size, interleave, rows_per_req, req_lens +): + """Gathering block-interleaved shards must restore each request's token order.""" from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan + monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False) + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu", + lambda x, device: x.to(device), + ) scheduled = np.array(req_lens, dtype=np.int64) - shard_rows = -(-scheduled // dcp_world_size) - rows = np.arange(len(scheduled)) + shard_rows = np.array( + [ + max( + _local_count(g, r, dcp_world_size, interleave) + for r in range(dcp_world_size) + ) + for g in req_lens + ] + ) + rows = np.repeat(np.arange(len(scheduled)), rows_per_req) plan = build_pcp_global_chunk_plan( - rows, shard_rows, dcp_world_size, torch.device("cpu") + rows, shard_rows[rows], dcp_world_size, torch.device("cpu"), interleave ) starts = np.concatenate([[0], np.cumsum(scheduled)]) # Build each rank's padded shard exactly as the cache gather would. - padded_cu = plan.padded_local_cu.tolist() + padded_cu = plan.padded_local_cu[::rows_per_req].tolist() shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1) for r in range(dcp_world_size): for i, g in enumerate(req_lens): - for t in range(r, g, dcp_world_size): - shards[r, padded_cu[i] + t // dcp_world_size, 0] = starts[i] + t + owned = [t for t in range(g) if (t // interleave) % dcp_world_size == r] + for local, t in enumerate(owned): + shards[r, padded_cu[i] + local, 0] = starts[i] + t gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1) restored = gathered[plan.deinterleave_idx] # The layout is padded per request; only the real positions are read. row_start = plan.row_start_cu.tolist() - for i, g in enumerate(req_lens): + for row, i in enumerate(rows): + g = req_lens[i] torch.testing.assert_close( - restored[row_start[i] : row_start[i] + g, 0], + restored[row_start[row] : row_start[row] + g, 0], torch.arange(starts[i], starts[i] + g, dtype=torch.float32), ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 164a3f89169f..828f2c2ac49c 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -25,6 +25,7 @@ has_deep_gemm, native_next_n_supported, ) +from vllm.utils.math_utils import round_down from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.attention.backend import ( @@ -286,6 +287,7 @@ def build_pcp_global_chunk_plan( row_shard_rows: np.ndarray, dcp_world_size: int, device: torch.device, + interleave: int = 1, ) -> PCPGlobalChunkPlan: """Plan the PCP packing for one chunk from its requests' KV shard rows. @@ -327,10 +329,13 @@ def build_pcp_global_chunk_plan( for i in range(len(region_first_row)): g = int(region_padded[i]) * dcp_world_size t = np.arange(g, dtype=np.int64) + local = round_down(t // dcp_world_size, interleave) + t % interleave + # Padded positions are outside causal bounds but must remain in-bounds. + local = np.minimum(local, region_padded[i] - 1) idx_np[region_start[i] : region_start[i] + g] = ( - (t % dcp_world_size) * padded_total + ((t // interleave) % dcp_world_size) * padded_total + region_padded_start[i] - + t // dcp_world_size + + local ) cu = async_copy_to_gpu(cu, device=device) @@ -1268,6 +1273,7 @@ def build( shard_rows[req_slice], self.dcp_world_size, self.device, + self.cp_kv_cache_interleave_size, ) metadata = build_prefill_chunk_metadata( req_slice.start, From 186a1e62221c02cdf8ef0433eb4ad33d1b48c7f7 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sun, 13 Sep 2026 22:47:21 +0200 Subject: [PATCH 140/774] [Warmup] Gemma 4 de-JITification (#55768) Signed-off-by: LopezCastroRoberto Co-authored-by: Codex --- vllm/v1/attention/backends/flash_attn.py | 237 ++++++++++++++++++- vllm/v1/worker/gpu/attn_utils.py | 5 +- vllm/v1/worker/gpu/model_runner.py | 18 +- vllm/vllm_flash_attn/flash_attn_interface.py | 54 ++++- 4 files changed, 287 insertions(+), 27 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 9109a9c6d7ab..eb2268cc624c 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -5,12 +5,13 @@ import copy import functools from dataclasses import dataclass -from typing import ClassVar +from typing import Any, ClassVar import numpy as np import torch from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, WarmupChoices from vllm.platforms import current_platform from vllm.utils.torch_utils import ( PIN_MEMORY, @@ -46,6 +47,7 @@ if is_flash_attn_varlen_func_available(): from vllm.v1.attention.backends.fa_utils import ( + compile_flash_attn_varlen_func_from_specs, flash_attn_supports_sinks, flash_attn_varlen_func, get_scheduler_metadata, @@ -67,7 +69,13 @@ AttentionMetadataBuilder, CommonAttentionMetadata, ) -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, + KVCacheSpec, + KVQuantMode, + SlidingWindowSpec, +) from vllm.v1.worker.cp_utils import ( run_split_fa2_dcp_context_attention, should_skip_dcp_context_attention, @@ -77,6 +85,200 @@ logger = init_logger(__name__) +FA4_DENSE_FLOAT_DTYPES = (torch.bfloat16, torch.float16) +FA4_DENSE_Q_TILE = 128 +FA4_DENSE_NUM_BLOCKS = 256 +FA4_DENSE_MAX_SEQLEN_K = 8192 + + +class FA4DenseAttentionKernel(VllmJitKernel["FA4DenseAttentionKernel.CompileKey"]): + """Warm paged, causal FA4 on Blackwell and Hopper.""" + + @dataclass(frozen=True) + class CompileKey: + q_stage: int + dtype: torch.dtype + qhead_per_kvhead: int + head_dim: int + page_size: int + has_window_left: bool + has_window_right: bool + softcap: float + is_single_batch: bool + num_splits: int + + @staticmethod + def kernel(*args: Any, **kwargs: Any) -> Any: + assert flash_attn_varlen_func is not None + return flash_attn_varlen_func(*args, **kwargs) + + def dispatch( # type: ignore[override] + self, + *, + q_stage: int, + dtype: torch.dtype, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + window_size: tuple[int, int], + softcap: float, + is_single_batch: bool = True, + num_splits: int = 1, + ) -> CompileKey: + return self.CompileKey( + q_stage=q_stage, + dtype=dtype, + qhead_per_kvhead=num_qo_heads // num_kv_heads, + head_dim=head_dim, + page_size=page_size, + has_window_left=window_size[0] >= 0, + has_window_right=window_size[1] >= 0, + softcap=softcap, + is_single_batch=is_single_batch, + num_splits=num_splits, + ) + + def get_warmup_keys( + self, + *, + vllm_config: VllmConfig, + kv_cache_spec: AttentionSpec, + num_qo_heads: int, + ) -> list[CompileKey]: + capability = current_platform.get_device_capability() + major = capability.major if capability is not None else None + if ( + vllm_config.parallel_config.decode_context_parallel_size != 1 + or vllm_config.model_config.rswa_window is not None + or kv_cache_spec.kv_quant_mode != KVQuantMode.NONE + or kv_cache_spec.dtype not in FA4_DENSE_FLOAT_DTYPES + or not isinstance(kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec)) + ): + return [] + head_dim = kv_cache_spec.head_size + page_size = kv_cache_spec.block_size + sliding_window = kv_cache_spec.sliding_window + window_size = ( + (sliding_window - 1, 0) if sliding_window is not None else (-1, -1) + ) + softcap = float( + getattr( + vllm_config.model_config.hf_text_config, "attn_logit_softcapping", 0 + ) + or 0 + ) + if ( + major in (10, 11) + and uses_fa4_hd256_kernel(head_dim) + and page_size == FA4_HD256_PAGE_SIZE + ): + q_stages = WarmupChoices(1, 2) + single_batches = WarmupChoices(True, False) + split_counts = WarmupChoices(1) + elif ( + major == 9 and head_dim == 512 and window_size == (-1, -1) and softcap == 0 + ): + # SM90 forward has split/non-split variants; its transitive + # combine specializes at 32/64/128/256 splits. + q_stages = WarmupChoices(1) + single_batches = WarmupChoices(False) + split_counts = WarmupChoices(1, 32, 64, 128, 256) + else: + return [] + + return self._trace_dispatch(self.dispatch)( + q_stage=q_stages, + is_single_batch=single_batches, + num_splits=split_counts, + dtype=vllm_config.model_config.dtype, + num_qo_heads=num_qo_heads, + num_kv_heads=kv_cache_spec.num_kv_heads, + head_dim=head_dim, + page_size=page_size, + window_size=WarmupChoices(window_size), + softcap=softcap, + ) + + def compile(self, compile_key: CompileKey) -> None: + assert compile_flash_attn_varlen_func_from_specs is not None + max_seqlen_q = FA4_DENSE_Q_TILE + 1 if compile_key.q_stage == 2 else 1 + batch_size = 1 if compile_key.is_single_batch else 2 + kv_shape = ( + FA4_DENSE_NUM_BLOCKS, + compile_key.page_size, + 1, + compile_key.head_dim, + ) + kv_stride = ( + 2 * compile_key.page_size * compile_key.head_dim, + 2 * compile_key.head_dim, + 2 * compile_key.head_dim, + 1, + ) + compile_flash_attn_varlen_func_from_specs( + q_shape=( + batch_size * max_seqlen_q, + compile_key.qhead_per_kvhead, + compile_key.head_dim, + ), + k_shape=kv_shape, + v_shape=kv_shape, + k_stride=kv_stride, + v_stride=kv_stride, + q_dtype=compile_key.dtype, + cu_seqlens_q_shape=(batch_size + 1,), + seqused_k_shape=(batch_size,), + page_table_shape=( + batch_size, + FA4_DENSE_MAX_SEQLEN_K // compile_key.page_size, + ), + max_seqlen_q=max_seqlen_q, + max_seqlen_k=FA4_DENSE_MAX_SEQLEN_K, + causal=True, + window_size=[ + 1 if compile_key.has_window_left else -1, + 1 if compile_key.has_window_right else -1, + ], + softcap=compile_key.softcap, + num_splits=2 if compile_key.num_splits > 1 else 1, + fa_version=4, + return_softmax_lse=False, + ) + if compile_key.head_dim == 512 and compile_key.num_splits > 1: + from torch._subclasses.fake_tensor import FakeTensorMode + + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd_combine + + # FA4's forward compile-only entry returns before its transitive + # SplitKV combine. Call the native combine entry in FakeTensorMode: + # it populates the native cache without launching the kernel. + with FakeTensorMode(): + _flash_attn_fwd_combine( + torch.empty( + (compile_key.num_splits, 2, 1, 512), + dtype=torch.float32, + device="cuda", + ), + torch.empty( + (compile_key.num_splits, 2, 1), + dtype=torch.float32, + device="cuda", + ), + torch.empty((2, 1, 512), dtype=compile_key.dtype, device="cuda"), + cu_seqlens=torch.empty((3,), dtype=torch.int32, device="cuda"), + ) + + def __call__( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + **kwargs: Any, + ) -> Any: + return self.kernel(q=q, k=k, v=v, **kwargs) + class FlashAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -447,15 +649,27 @@ def __init__( self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 + head_size_v = getattr(kv_cache_spec, "head_size_v", None) + fa_version = get_flash_attn_version( + head_size=self.headdim, + head_size_v=head_size_v, + kv_cache_block_size=self.block_size, + supports_fa4_hd256=True, + ) + self.fa4_hd256 = fa_version == 4 and uses_fa4_hd256_kernel( + self.headdim, head_size_v + ) - self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and ( - get_flash_attn_version( - head_size=self.headdim, - kv_cache_block_size=self.block_size, - supports_fa4_hd256=True, + if ( + vllm_config.kernel_config.enable_jit_warmup + and self.model_config.hf_config.model_type in ("gemma4", "gemma4_unified") + and fa_version == 4 + ): + _FA4_DENSE_ATTENTION_KERNEL.register_warmup( + vllm_config=vllm_config, + kv_cache_spec=kv_cache_spec, + num_qo_heads=self.num_heads_q, ) - == 4 - ) try: from vllm.distributed.parallel_state import get_dcp_group @@ -1164,7 +1378,7 @@ def forward( block_table = block_table[:, :num_pages] num_splits = 1 - flash_attn_varlen_func( + _FA4_DENSE_ATTENTION_KERNEL( q=query[:num_actual_tokens], k=key_cache, v=value_cache, @@ -1837,3 +2051,6 @@ def cascade_attention( # Merge prefix and suffix outputs, and store the result in output. merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse) + + +_FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel() diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index a3f6541b0866..eccfdbabcdfd 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -8,7 +8,10 @@ import numpy as np import torch -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import ( + VllmConfig, + get_layers_from_vllm_config, +) from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fe423d6e38a6..9d7918126f62 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -632,12 +632,18 @@ def initialize_kv_cache( draft_attn_layer_names = None if isinstance(self.speculator, DraftModelSpeculator): draft_attn_layer_names = self.speculator.draft_attn_layer_names - self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( - self.kv_cache_config, - self.vllm_config, - self.device, - draft_layer_names=draft_attn_layer_names, - ) + # Metadata builders select attention kernels that need JIT warmup. + with self.jit_warmup_registry.activate(): + ( + self.attn_groups, + attn_cg_support, + self.kernel_block_sizes, + ) = init_attn_backend( + self.kv_cache_config, + self.vllm_config, + self.device, + draft_layer_names=draft_attn_layer_names, + ) additional_attn_cg_support = self.model_state.get_additional_cg_support() attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support) # The speculator clears the flag at load time when the checkpoint has diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 81768e12092e..4ba4692f6778 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -445,6 +445,7 @@ def compile_flash_attn_varlen_func_from_specs( v_shape: tuple[int, ...], q_dtype: torch.dtype, v_stride: tuple[int, ...] | None = None, + k_stride: tuple[int, ...] | None = None, cu_seqlens_q_shape: tuple[int, ...] | None = None, cu_seqlens_k_shape: tuple[int, ...] | None = None, max_seqlen_q: int | None = None, @@ -457,6 +458,9 @@ def compile_flash_attn_varlen_func_from_specs( return_softmax_lse=False, num_splits: int = 0, fa_version: int = DEFAULT_FA_VERSION, + seqused_k_shape: tuple[int, ...] | None = None, + page_table_shape: tuple[int, ...] | None = None, + softcap: float | None = None, ) -> None: if fa_version != 4: raise ValueError( @@ -467,7 +471,8 @@ def compile_flash_attn_varlen_func_from_specs( del deterministic from vllm.vllm_flash_attn.cute.interface import ( - compile_flash_attn_varlen_func_from_specs as _fa4_compile_flash_attn_varlen_func_from_specs, + _flash_attn_fwd, + _make_compile_only_tensor_spec, ) real_window_size: tuple[int, int] @@ -480,21 +485,50 @@ def compile_flash_attn_varlen_func_from_specs( if softmax_scale is None: softmax_scale = q_shape[-1] ** (-0.5) - return _fa4_compile_flash_attn_varlen_func_from_specs( - q_shape=q_shape, - k_shape=k_shape, - v_shape=v_shape, - q_dtype=q_dtype, - v_stride=v_stride, - cu_seqlens_q_shape=cu_seqlens_q_shape, - cu_seqlens_k_shape=cu_seqlens_k_shape, + q = _make_compile_only_tensor_spec(q_shape, q_dtype) + k = _make_compile_only_tensor_spec(k_shape, q_dtype, stride=k_stride) + v = _make_compile_only_tensor_spec(v_shape, q_dtype, stride=v_stride) + cu_seqlens_q = _make_compile_only_tensor_spec(cu_seqlens_q_shape, torch.int32, 4) + out = _make_compile_only_tensor_spec( + (*q_shape[:-1], v_shape[-1]), + q_dtype, + ) + lse = None + if return_softmax_lse: + assert q is not None + if cu_seqlens_q_shape is None: + lse_shape = (*q.shape[:-3], q.shape[-2], q.shape[-3]) + lse_stride = None + else: + lse_shape = (q.shape[-2], q.shape[0]) + lse_stride = (None, 1) + lse = _make_compile_only_tensor_spec( + lse_shape, + torch.float32, + 4, + stride=lse_stride, + ) + + _flash_attn_fwd( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=_make_compile_only_tensor_spec(cu_seqlens_k_shape, torch.int32, 4), + seqused_k=_make_compile_only_tensor_spec(seqused_k_shape, torch.int32, 4), + page_table=_make_compile_only_tensor_spec(page_table_shape, torch.int32, 4), max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, softmax_scale=softmax_scale, causal=causal, - window_size=real_window_size, + softcap=softcap, + window_size_left=real_window_size[0], + window_size_right=real_window_size[1], num_splits=num_splits, return_lse=return_softmax_lse, + out=out, + lse=lse, + compile_only=True, ) From b23433088bf29980d20dff1bd64c753dd8883506 Mon Sep 17 00:00:00 2001 From: stu-cao Date: Sun, 13 Sep 2026 13:56:28 -0700 Subject: [PATCH 141/774] [Attention] Remove DCP indexer interleave guard and test TP1 output parity (#55802) Signed-off-by: Stu Cao Signed-off-by: <> Signed-off-by: Summer Yang Signed-off-by: Lucas Wilkinson Co-authored-by: Stu Cao Co-authored-by: OpenAI Codex Co-authored-by: Summer Yang Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- vllm/v1/attention/backends/mla/indexer.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 828f2c2ac49c..2321cf8d13d4 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -771,15 +771,6 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.pcp_world_size = parallel_config.prefill_context_parallel_size self.use_pcp = self.pcp_world_size > 1 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size - # The DCP sparse-indexer code is parameterized by interleave size, but - # interleave > 1 is not yet validated end-to-end (gsm8k parity fails), - # so fail closed here rather than silently produce wrong output. - if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size > 1: - raise NotImplementedError( - "DCP sparse indexer currently supports only " - f"cp_kv_cache_interleave_size=1 (got " - f"{self.cp_kv_cache_interleave_size})." - ) # NOTE(Chen):an estimated max size of flattened_kv. Need to double check. self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config) self.num_speculative_tokens = ( From 319cc5ef19946d34c2e66cbbec5bda29d0bfa328 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Mon, 14 Sep 2026 06:07:12 +0800 Subject: [PATCH 142/774] [Performance][EPD] Reduce Python proxy serialization overhead (#56657) Signed-off-by: Tianyu Guo --- .../disaggregated_encoder/disagg_epd_proxy.py | 52 ++++--- .../ec_connector/unit/test_epd_proxy_retry.py | 133 +++++++++++++++++- 2 files changed, 160 insertions(+), 25 deletions(-) diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index ace44558287e..3912c1c67a4b 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -33,9 +33,10 @@ from typing import Any import aiohttp +import msgspec import uvicorn from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse ############################################################################### # FastAPI app & global state @@ -264,7 +265,7 @@ async def fanout_encoder_primer( for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)): # Derive a *child* request id: :: child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}" - headers = {"x-request-id": child_req_id} + headers = {"x-request-id": child_req_id, "Content-Type": "application/json"} # With --no-rewrite the decoder still receives the raw image and derives # the cache key by hashing it, so the encoder must do the same -- passing @@ -307,7 +308,7 @@ async def fanout_encoder_primer( tasks.append( encode_session.post( f"{target_url}/v1/chat/completions", - json=encoder_req, + data=msgspec.json.encode(encoder_req), headers=headers, ) ) @@ -347,7 +348,7 @@ async def fanout_encoder_primer( # The encoder reports each mm_hash's metadata (e.g. the grid) here, # keyed by the same uuid this proxy assigned above. try: - params = (await r.json()).get("ec_transfer_params") or {} + params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {} except Exception: logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx) params = {} @@ -407,7 +408,7 @@ async def maybe_prefill( prefill_response = await process_prefill_stage(req_data, p_url, req_id) # for nixl connector to facilitate kv transfer... - prefill_response_json = await prefill_response.json() + prefill_response_json = msgspec.json.decode(await prefill_response.read()) kv_transfer_params = prefill_response_json.get("kv_transfer_params", {}) if kv_transfer_params: return {**req_data, "kv_transfer_params": kv_transfer_params} @@ -439,10 +440,12 @@ async def process_prefill_stage( if "stream_options" in prefill_request: del prefill_request["stream_options"] - headers = {"x-request-id": req_id} + headers = {"x-request-id": req_id, "Content-Type": "application/json"} try: prefill_response = await prefill_session.post( - f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers + f"{p_url}/v1/chat/completions", + data=msgspec.json.encode(prefill_request), + headers=headers, ) prefill_response.raise_for_status() @@ -598,7 +601,7 @@ async def forward_non_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, -) -> dict: +) -> Response: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() @@ -608,12 +611,14 @@ async def forward_non_stream( _t2 = time.perf_counter() logger.info("[%s] Forwarding to decode: %s", req_id, d_url) - headers = {"x-request-id": req_id} + headers = {"x-request-id": req_id, "Content-Type": "application/json"} if dp_rank is not None: headers["X-data-parallel-rank"] = str(dp_rank) async with decode_session.post( - f"{d_url}/v1/chat/completions", json=prepared, headers=headers + f"{d_url}/v1/chat/completions", + data=msgspec.json.encode(prepared), + headers=headers, ) as resp: if resp.status >= 400: detail = await resp.text() @@ -637,7 +642,7 @@ async def forward_non_stream( detail, ) raise HTTPException(status_code=resp.status, detail=detail) - out = await resp.json() + out = await resp.read() _t3 = time.perf_counter() logger.info( "STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f " @@ -649,7 +654,15 @@ async def forward_non_stream( (_t3 - _t0) * 1e3, attempt, ) - return out + return Response( + content=out, + status_code=resp.status, + headers={ + "Content-Type": resp.headers.get( + "Content-Type", "application/json" + ) + }, + ) raise HTTPException(status_code=500, detail="Decode failed after re-encoding") except HTTPException: @@ -667,7 +680,7 @@ async def forward_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, -) -> AsyncIterator[str]: +) -> AsyncIterator[bytes]: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() @@ -677,14 +690,14 @@ async def forward_stream( _t2 = time.perf_counter() logger.info("[%s] Starting streaming from decode: %s", req_id, d_url) - headers = {"x-request-id": req_id} + headers = {"x-request-id": req_id, "Content-Type": "application/json"} if dp_rank is not None: headers["X-data-parallel-rank"] = str(dp_rank) _first = None async with decode_session.post( f"{d_url}/v1/chat/completions", - json=prepared, + data=msgspec.json.encode(prepared), headers=headers, ) as resp: # Retry only before the first chunk: once anything reached the @@ -701,11 +714,11 @@ async def forward_stream( ) continue resp.raise_for_status() - async for chunk in resp.content.iter_chunked(1024): + async for chunk in resp.content.iter_any(): if chunk: if _first is None: _first = time.perf_counter() - yield chunk.decode("utf-8", errors="ignore") + yield chunk _t3 = time.perf_counter() logger.info( @@ -739,7 +752,7 @@ async def forward_stream( @app.post("/v1/chat/completions") async def chat_completions(request: Request): try: - req_data = await request.json() + req_data = msgspec.json.decode(await request.body()) req_id = request.headers.get("x-request-id", str(uuid.uuid4())) e_urls = app.state.e_urls # we want the full list for fan-out @@ -764,10 +777,9 @@ async def chat_completions(request: Request): ), media_type="text/event-stream", ) - result = await forward_non_stream( + return await forward_non_stream( req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank ) - return JSONResponse(content=result) except HTTPException: raise diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index b026b5537a20..9a92b1fdbcca 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -10,7 +10,11 @@ import importlib.util from pathlib import Path +import httpx +import msgspec import pytest +from aiohttp import web +from aiohttp.test_utils import TestServer PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" @@ -29,8 +33,8 @@ class _Response: def __init__(self, params): self._params = params - async def json(self): - return {"kv_transfer_params": self._params} + async def read(self): + return msgspec.json.encode({"kv_transfer_params": self._params}) def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch): @@ -62,8 +66,8 @@ def __init__(self, params): self.status = 200 self._params = params - async def json(self): - return {"ec_transfer_params": self._params} + async def read(self): + return msgspec.json.encode({"ec_transfer_params": self._params}) async def text(self): return "" @@ -75,7 +79,7 @@ class _EncoderSession: def __init__(self, replies): self._replies = list(replies) - async def post(self, url, json=None, headers=None): + async def post(self, url, data=None, headers=None): return _EncoderResponse(self._replies.pop(0)) @@ -147,6 +151,125 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): assert params["ec_items"][0]["transfer_id"] +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("prefill", [False, True]) +async def test_http_roundtrip_preserves_payload_and_response_bytes( + proxy, monkeypatch, stream, prefill +): + """Exercise real HTTP hops, rewrite and a decode retry without model servers.""" + seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} + prefix = b'data: {"content":"' + # Split a multibyte character at the old 1024-byte forwarding boundary. + expected = ( + prefix + b"x" * (1023 - len(prefix)) + '图片🌍"}\n\ndata: [DONE]\n\n'.encode() + if stream + else '{ "choices": [{"message": {"content": "图片🌍"}}] }\n'.encode() + ) + + async def backend(request): + stage = request.match_info["stage"] + assert request.content_type == "application/json" + body = await request.json() + seen[stage].append(body) + if stage == "encode": + item = body["messages"][0]["content"][0] + return web.json_response( + { + "ec_transfer_params": { + item["uuid"]: { + "metadata": {"image_grid_thw": [[1, 2, 2]]}, + "peer_port": len(seen[stage]), + } + } + } + ) + if stage == "prefill": + assert body["max_tokens"] == 1 and body["stream"] is False + return web.json_response( + {"kv_transfer_params": {"remote_block_ids": [len(seen[stage])]}} + ) + if len(seen[stage]) == 1: + return web.Response(status=500, text="retry") + return web.Response( + body=expected, + content_type="text/event-stream" if stream else "application/json", + ) + + backend_app = web.Application() + backend_app.router.add_post("/{stage}/v1/chat/completions", backend) + async with TestServer(backend_app) as server: + for stage, field in [ + ("encode", "e_urls"), + ("prefill", "p_urls"), + ("decode", "d_urls"), + ]: + urls = [str(server.make_url(f"/{stage}"))] + if stage == "prefill" and not prefill: + urls = [] + monkeypatch.setattr(proxy.app.state, field, urls, raising=False) + monkeypatch.setattr(proxy.app.state, "d_ec_urls", [], raising=False) + monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False) + monkeypatch.setattr(proxy, "DECODE_RETRIES", 1) + monkeypatch.setattr(proxy, "NO_REWRITE", False) + item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}} + body = { + "model": "test", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "描述图片🌍"}, + item, + item, + ], + } + ], + "stream": stream, + "temperature": 0.01, + "max_tokens": 32, + "seed": 42, + "structured_outputs": {"choice": ["A", "B"]}, + } + await proxy.on_startup() + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy" + ) as client: + response = await client.post("/v1/chat/completions", json=body) + assert response.status_code == 200 + assert response.content == expected + assert response.headers["content-type"].startswith( + "text/event-stream" if stream else "application/json" + ) + finally: + await proxy.on_shutdown() + + assert len(seen["encode"]) == 4 + assert len(seen["decode"]) == 2 + final = seen["decode"][-1] + for key in [ + "model", + "stream", + "temperature", + "max_tokens", + "seed", + "structured_outputs", + ]: + assert final[key] == body[key] + content = final["messages"][0]["content"] + assert content[0] == body["messages"][0]["content"][0] + for reference in content[1:]: + assert reference == { + "type": "image_embeds", + "image_embeds": {"image_grid_thw": [1, 2, 2]}, + "uuid": proxy.content_uuid(item), + } + assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4) + if prefill: + assert final["kv_transfer_params"] == {"remote_block_ids": [2]} + + @pytest.mark.parametrize("server_keep_alive", ["0.1", "1", "5", "2", "30"]) def test_pooled_connections_are_retired_before_the_server_closes_them( proxy, monkeypatch, server_keep_alive From 9f03b510c33575fe40c320b2d266a289e8a4b83a Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sun, 13 Sep 2026 20:22:08 -0400 Subject: [PATCH 143/774] [DSv4 Bug] fix dsv4 start up error `NotImplementedError: DeepSeek V4 MegaMoE currently requires expert parallel` (#55914) Signed-off-by: yewentao256 --- vllm/config/speculative.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cfecff72b618..186126682093 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1746,6 +1746,7 @@ def create_draft_parallel_config( draft_parallel_config = ParallelConfig( pipeline_parallel_size=1, tensor_parallel_size=speculative_draft_tensor_parallel_size, + enable_expert_parallel=target_parallel_config.enable_expert_parallel, distributed_executor_backend=target_parallel_config.distributed_executor_backend, max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce, From 3b533197a9419302987c12bcea5e4e06c9d40858 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 13 Sep 2026 18:58:20 -0700 Subject: [PATCH 144/774] [Perf] Parallelize mHC pre-norm JIT warmup (#56683) Signed-off-by: Woosuk Kwon --- vllm/model_executor/kernels/mhc/warmup.py | 16 ++++++++++++++++ vllm/model_executor/warmup/jit_warmup.py | 11 +++++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 6a42ca622270..0abe1a2161c2 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable from dataclasses import asdict, dataclass from typing import Any @@ -61,6 +62,21 @@ def compile(self, compile_key: CompileKey) -> None: **asdict(compile_key) ) + def compile_many(self, compile_keys: Iterable[CompileKey]) -> None: + missing = list( + dict.fromkeys( + key for key in compile_keys if key not in self._compiled_cache + ) + ) + if len(missing) < 2: + return super().compile_many(missing) + + # TileLang elaborates serially before compiling independent TIR modules. + compiled = self.kernel.par_compile( + [asdict(key) for key in missing], num_workers=min(4, len(missing)) + ) + self._compiled_cache.update(zip(missing, compiled, strict=True)) + def __call__(self, *tensors, **fields): compile_key = self.dispatch( n_splits=tensors[0].shape[0], diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 72626964ae67..84f713aa6489 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -963,14 +963,18 @@ def compile(self, compile_key: CompileKeyT) -> None: """Compile one warmup key.""" raise NotImplementedError + def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None: + """Compile a batch of warmup keys, allowing backend-specific scheduling.""" + for compile_key in compile_keys: + self.compile(compile_key) + def register_warmup(self, *args: Any, **kwargs: Any) -> None: """Register this kernel with the active runner's warmup registry.""" JitWarmupRegistry.register(self, *args, **kwargs) def warmup(self, *args: Any, **kwargs: Any) -> None: """Compile this kernel's warmup keys.""" - for compile_key in self.get_warmup_keys(*args, **kwargs): - self.compile(compile_key) + self.compile_many(self.get_warmup_keys(*args, **kwargs)) def _same_value(left: Any, right: Any) -> bool: @@ -1111,5 +1115,4 @@ def warmup(self) -> None: f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", refresh=False, ) - for compile_key in compile_keys: - kernel.compile(compile_key) + kernel.compile_many(compile_keys) From fe49665eb3b28283eb95d6fb143e5c6ea26c0923 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 14 Sep 2026 10:01:50 +0800 Subject: [PATCH 145/774] [XPU][CI] update test requirements (#56704) Signed-off-by: zhenwei-intel --- requirements/test/xpu.in | 10 +++++----- requirements/test/xpu.txt | 28 ++++++++++++++++++---------- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index e93b418e9a65..3aab5bdb00e1 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -32,8 +32,8 @@ soundfile # --- Tool Parsing & Evaluation --- blobfile rapidfuzz -gpt-oss[eval]==0.0.5 -schemathesis +gpt-oss >= 0.0.7; python_version > '3.11' +schemathesis>=4.0.0 # Required for openai schema test. jiwer bm25s pystemmer @@ -42,12 +42,12 @@ num2words pqdm # --- Vision & Multimodal --- -timm +timm>=1.0.17 # required for internvl and gemma3n-mm test albumentations ftfy matplotlib -decord==0.6.0 -mistral-common[image,audio] +decord==0.6.0; platform_machine == "x86_64" +mistral-common[image,audio]>=1.11.6 # required for voxtral test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6c8bc1fcdd7e..b7cf49649b7e 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -67,9 +67,7 @@ azure-storage-blob==12.30.1 blake3==1.0.9 # via -r requirements/test/../common.txt blobfile==3.2.0 - # via - # -r requirements/test/xpu.in - # gpt-oss + # via -r requirements/test/xpu.in bm25s==0.3.10 # via # -r requirements/test/xpu.in @@ -259,7 +257,7 @@ googleapis-common-protos==1.75.1 # google-api-core # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.5 +gpt-oss==0.0.9 # via -r requirements/test/xpu.in graphql-core==3.2.11 # via hypothesis-graphql @@ -379,7 +377,6 @@ jinja2==3.1.6 # via # -c requirements/xpu.txt # fastapi - # gpt-oss # lm-eval # torch jiter==0.16.0 @@ -409,6 +406,10 @@ jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema +jupyter-client==8.10.0 + # via gpt-oss +jupyter-core==5.9.1 + # via jupyter-client kiwisolver==1.5.1 # via matplotlib lark==1.2.2 @@ -529,7 +530,6 @@ numpy==2.2.6 # datasets # decord # evaluate - # gpt-oss # librosa # lm-eval # matplotlib @@ -580,7 +580,6 @@ openai==3.2.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # gpt-oss openai-harmony==0.0.8 # via # -c requirements/common.txt @@ -662,7 +661,6 @@ pandas==3.0.5 # via # datasets # evaluate - # gpt-oss partial-json-parser==0.2.1.1.post7 # via -r requirements/test/../common.txt pathvalidate==3.3.1 @@ -674,7 +672,9 @@ pillow==12.3.0 # mistral-common # torchvision platformdirs==4.11.3 - # via pooch + # via + # jupyter-core + # pooch pluggy==1.6.0 # via # pytest @@ -813,6 +813,7 @@ pytest-timeout==2.4.0 python-dateutil==2.9.0.post0 # via # botocore + # jupyter-client # matplotlib # pandas # typepy @@ -848,6 +849,7 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # jupyter-client rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in @@ -1072,12 +1074,13 @@ torchvision==0.28.0+xpu # via # open-clip-torch # timm +tornado==6.5.8 + # via jupyter-client tqdm==4.70.0 # via # -r requirements/test/../common.txt # datasets # evaluate - # gpt-oss # huggingface-hub # lm-eval # modelscope @@ -1088,6 +1091,10 @@ tqdm==4.70.0 # pqdm # sentence-transformers # transformers +traitlets==5.16.1 + # via + # jupyter-client + # jupyter-core transformers==5.16.1 # via # -c requirements/common.txt @@ -1139,6 +1146,7 @@ typing-extensions==4.16.0 # grpcio # httpx2 # huggingface-hub + # jupyter-client # lm-eval # mcp # mcp-types From 58d45fd767ff81229afd92543400f5c8c93f4000 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Sun, 13 Sep 2026 19:18:16 -0700 Subject: [PATCH 146/774] [Bugfix][KV Connector] MooncakeStore: exclude non-prefix-cacheable (QSA ring) groups; fix align-mode check (#55027) Signed-off-by: Zhewen Li Co-authored-by: Zhewen Li Co-authored-by: OpenAI Codex Co-authored-by: Claude Fable 5.1 --- tests/v1/core/test_kv_cache_utils.py | 42 ++++++- .../unit/test_mooncake_store_connector.py | 48 ++++++++ .../unit/test_mooncake_store_hma_e2e.py | 23 ++-- .../unit/test_mooncake_store_scheduler.py | 110 ++++++++++++++++-- .../unit/test_mooncake_store_worker.py | 90 ++++++++++++++ .../v1/mooncake/store/connector.py | 22 ++-- .../v1/mooncake/store/coordinator.py | 2 +- .../v1/mooncake/store/scheduler.py | 47 +++++--- .../kv_connector/v1/mooncake/store/worker.py | 13 ++- vllm/v1/core/kv_cache_manager.py | 19 ++- vllm/v1/kv_cache_interface.py | 17 +++ 11 files changed, 381 insertions(+), 52 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 3b29799d3f89..dbddc556c9ba 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -35,7 +35,7 @@ from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor from vllm.utils.mem_constants import GiB_bytes -from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, @@ -61,6 +61,7 @@ ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, HiSparseHotSpec, @@ -324,6 +325,45 @@ def test_kv_cache_config_selects_only_transferable_groups(): ) +def test_kv_cache_config_selects_prefix_cacheable_groups(): + """Prefix stores exclude scratch state without changing transfer groups.""" + full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec()) + qsa_group = KVCacheGroupSpec( + ["qsa"], + CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=64, + head_size_v=0, + dtype=torch.float16, + ), + ) + disabled_group = KVCacheGroupSpec( + ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False + ) + config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[full_group, qsa_group, disabled_group], + ) + assert config.transfer_group_ids == (0, 1) + assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2]) + assert config.prefix_cacheable_group_ids == (0,) + assert config.prefix_cacheable_groups == (full_group,) + + +def test_kv_cache_blocks_selects_requested_groups(): + blocks = KVCacheBlocks( + ( + [KVCacheBlock(1)], + [KVCacheBlock(2)], + [KVCacheBlock(3)], + ) + ) + + assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) + + def new_sliding_window_spec( block_size=16, num_kv_heads=2, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 96315b026df1..6afb529d145f 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -30,6 +30,7 @@ MooncakeStoreConnectorStats, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -91,6 +92,53 @@ def _make_kv_cache_config() -> KVCacheConfig: ) +def _make_qsa_hybrid_kv_cache_config(mamba_mode: str): + full_spec = FullAttentionSpec( + block_size=800, num_kv_heads=8, head_size=64, dtype=None + ) + qsa_spec = CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=64, + head_size_v=0, + dtype=torch.float16, + ) + mamba_spec = MambaSpec( + block_size=800, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode=mamba_mode, + ) + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["full"], full_spec), + KVCacheGroupSpec(["qsa"], qsa_spec), + KVCacheGroupSpec(["mamba"], mamba_spec), + ], + ) + + +def test_validation_accepts_aligned_mamba_after_block_size_rewrite(): + vllm_config = _make_vllm_config() + vllm_config.cache_config.block_size = 8 + + mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config( + vllm_config, _make_qsa_hybrid_kv_cache_config("align") + ) + + +def test_validation_rejects_mamba_mode_directly(): + vllm_config = _make_vllm_config() + vllm_config.cache_config.block_size = 800 + + with pytest.raises(ValueError, match="mamba_cache_mode=\x27none\x27"): + mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config( + vllm_config, _make_qsa_hybrid_kv_cache_config("none") + ) + + def test_scheduler_requires_align_mode_for_mamba(): vllm_config = _make_vllm_config() mamba_spec = MambaSpec( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index e6b730242981..a93698bba954 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -759,8 +759,9 @@ def test_worker_setup_tolerates_finer_scratch_group(): worker.pp_size = 1 worker.num_kv_head = 8 assert worker.coord.enable_partial_hash_hits - # The scratch DB is keyed at its own block size and never probed. - assert worker.token_dbs[2].hash_block_size == 4 + # Only the participating full-attention and Mamba groups are registered. + assert len(worker.token_dbs) == 2 + assert all(db.hash_block_size == 8 for db in worker.token_dbs) for g_idx, db in enumerate(worker.token_dbs): db.set_kv_caches_base_addr([g_idx * 10_000]) @@ -772,11 +773,11 @@ def test_worker_setup_tolerates_finer_scratch_group(): block_size=worker.block_size, coord=worker.coord, tp_rank=0, - group_put_steps=[1, 1, 1], + group_put_steps=[1, 1], kv_role="kv_both", ready_event=threading.Event(), replicate_config=MagicMock(), - group_participates=[True, True, False], + group_participates=[True, True], ) # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]). @@ -784,7 +785,7 @@ def test_worker_setup_tolerates_finer_scratch_group(): req = ReqMeta( req_id="r0", token_len_chunk=0, - block_ids=([1], [2], [3]), + block_ids=([1], [2]), block_hashes=hs, can_save=True, num_prompt_tokens=20, @@ -797,8 +798,7 @@ def test_worker_setup_tolerates_finer_scratch_group(): # A 13-token prompt sharing the prefix must hit the first hash unit. assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8 # The scratch group's namespace never enters the store. - scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0] - assert not any(key.startswith(scratch_prefix) for key in store._data) + assert not any("@group:2" in key for key in store._data) def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): @@ -838,7 +838,8 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): worker.tp_size = 1 worker.pp_size = 1 worker.num_kv_head = 8 - assert worker.token_dbs[1].hash_block_size == 8 + assert len(worker.token_dbs) == 1 + assert worker.token_dbs[0].hash_block_size == 16 raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8) raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8) @@ -879,14 +880,14 @@ def _fake_thread_init(*args, **kwargs): kv_role=worker.kv_role, ready_event=threading.Event(), enable_kv_event=False, - group_participates=[True, False], + group_participates=[True], ) hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)] save_req = ReqMeta( req_id="r0", token_len_chunk=64, - # The ring group holds one block for the request's lifetime. - block_ids=([1, 2, 3, 4], [1]), + # Scheduler metadata is projected to the participating store groups. + block_ids=([1, 2, 3, 4],), block_hashes=hs, can_save=True, store_job_id=1, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 67fa0c88aab4..7ced38d65a41 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -2,12 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import patch import pytest +import torch from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( LoadSpec, MooncakeLookupResult, + MooncakeStoreConnectorMetadata, MooncakeStoreWorkerMetadata, ReqMeta, RequestTracker, @@ -16,7 +19,16 @@ MooncakeStoreScheduler, ) from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.kv_cache_utils import KVCacheBlock from vllm.v1.core.sched.output import KVConnectorBlockState +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) def _make_bare_scheduler( @@ -35,9 +47,9 @@ def _make_bare_scheduler( scheduler._block_size = 16 scheduler._hash_block_size = hash_block_size scheduler.enable_partial_hash_hits = enable_partial_hash_hits - scheduler.kv_cache_config = SimpleNamespace( - select_transfer_block_ids=lambda block_ids: tuple(block_ids) - ) + scheduler.kv_cache_config = SimpleNamespace() + scheduler._store_group_ids = (0,) + scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 1: 1} scheduler.load_specs = {} scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} @@ -228,14 +240,82 @@ def test_pending_load_for_non_chosen_connector_is_dropped(): assert "req-0" not in scheduler._request_trackers +def _make_qsa_hybrid_cache_config(): + full = FullAttentionSpec(block_size=800, num_kv_heads=8, head_size=64, dtype=None) + circular = CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=64, + head_size_v=0, + dtype=torch.float16, + ) + mamba = MambaSpec( + block_size=800, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["full"], full), + KVCacheGroupSpec(["qsa"], circular), + KVCacheGroupSpec(["mamba"], mamba), + ], + ) + + +def test_scheduler_projects_nonprefix_groups_and_mamba_ids(): + vllm_config = SimpleNamespace( + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=800, enable_prefix_caching=True, prefix_match_unit=None + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1), + ) + + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler(vllm_config, _make_qsa_hybrid_cache_config()) + + assert scheduler._store_group_ids == (0, 2) + assert scheduler._boundary_state_group_ids == frozenset({1}) + + +def test_current_save_block_ids_use_store_group_projection(): + scheduler = _make_bare_scheduler() + scheduler.kv_cache_config = _make_qsa_hybrid_cache_config() + scheduler._store_group_ids = (0, 2) + meta = MooncakeStoreConnectorMetadata(set(), set()) + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=800, + block_ids=(), + block_hashes=[b"h0"], + can_save=True, + ) + meta.add_request(req_meta) + output = SimpleNamespace( + kv_connector_block_state=_make_connector_block_state(([10], [80], [30])) + ) + + scheduler._apply_current_save_block_ids(meta, output) + + assert req_meta.block_ids == ([10], [30]) + + def test_update_state_excludes_nontransfer_groups(): """Store metadata must match the worker's registered cache groups.""" scheduler = _make_bare_scheduler() - scheduler.kv_cache_config = SimpleNamespace( - select_transfer_block_ids=lambda block_ids: (block_ids[0],) - ) + scheduler._store_group_ids = (0,) request = SimpleNamespace(request_id="req-1") - blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9])) + blocks = KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(9)])) scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32) @@ -1003,6 +1083,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): def test_finished_partial_tail_is_pre_pinned_as_store_job(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler._store_group_ids = (0, 1) scheduler.client = SimpleNamespace(discard=lambda *_: None) request = SimpleNamespace( request_id="req-0", @@ -1143,6 +1224,20 @@ def _make_offload_only_output(entries, block_ids=([0],)): ) +def test_boundary_state_group_ids_are_remapped_to_store_projection(): + scheduler = _make_bare_scheduler(hash_block_size=800, enable_partial_hash_hits=True) + scheduler.kv_cache_config = _make_qsa_hybrid_cache_config() + scheduler._store_group_ids = (0, 2) + scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 2: 1} + scheduler._boundary_state_group_ids = frozenset({1}) + _register_offload_request(scheduler, prefill_end_tokens=800, num_prompt_tokens=800) + meta = MooncakeStoreConnectorMetadata(set(), set()) + + scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta) + + assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)] + + def test_resumed_prefill_claims_boundaries_past_prompt_length(): # A resumed request re-prefills its previously generated tokens, so its # save window (`prefill_end_tokens`) extends past `num_prompt_tokens`. @@ -1195,6 +1290,7 @@ def test_boundary_state_job_pins_exact_blocks_once(): def test_store_job_pins_current_non_null_non_mamba_blocks(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler._store_group_ids = (0, 1) request = SimpleNamespace( all_token_ids=list(range(48)), block_hashes=[bytes([i]) for i in range(12)], diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index f613ba720e2e..9baf16e2610d 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -46,8 +46,11 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, + MambaSpec, RSWASpec, ) from vllm.v1.kv_cache_layout import KVCacheLayout @@ -342,6 +345,43 @@ def _make_kv_cache_config( ) +def _make_qsa_hybrid_kv_cache_config(): + return KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=800, + num_kv_heads=8, + head_size=64, + dtype=None, + ), + ), + KVCacheGroupSpec( + ["qsa"], + CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=64, + head_size_v=0, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=800, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + + def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str: config_path = tmp_path / "mooncake_config.json" config_path.write_text(json.dumps(config), encoding="utf-8") @@ -2107,6 +2147,31 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget(): assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2} +def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch) + monkeypatch.setattr( + worker.MooncakeStoreConfig, + "load_from_config", + staticmethod(lambda: _make_config()), + ) + vllm_config = _make_vllm_config() + vllm_config.cache_config.block_size = 800 + vllm_config.cache_config.enable_prefix_caching = True + vllm_config.cache_config.prefix_match_unit = None + + store_worker = worker.MooncakeStoreWorker( + vllm_config, _make_qsa_hybrid_kv_cache_config() + ) + + assert [ + group.kv_cache_spec.block_size for group in store_worker._kv_cache_groups + ] == [800, 800] + assert [db.block_size for db in store_worker.token_dbs] == [800, 800] + + def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch): store = MagicMock() store.setup.return_value = 0 @@ -4144,6 +4209,31 @@ def test_register_kv_caches_uses_transfer_group_memory_domain(): ) +def test_register_kv_caches_excludes_nonprefix_cacheable_layers(): + num_blocks = 2 + worker = _make_bare_worker(num_gpu_blocks=num_blocks) + store_layout = MagicMock() + worker.token_dbs[0].store_layout = store_layout + full_cache = torch.zeros(num_blocks, 4, dtype=torch.float16) + qsa_cache = torch.zeros(num_blocks, 4, dtype=torch.float16) + + _register_with_mocked_threads( + worker, + {"layer0": full_cache, "qsa": qsa_cache}, + ) + + store_layout.register_kv_caches.assert_called_once() + registered_caches, registered_num_blocks = ( + store_layout.register_kv_caches.call_args.args + ) + assert registered_num_blocks == num_blocks + assert len(registered_caches) == 1 + assert registered_caches[0] is full_cache + worker.store.register_buffer.assert_called_once_with( + full_cache.data_ptr(), full_cache.untyped_storage().nbytes() + ) + + def test_register_kv_caches_separate_head_groups(): # LHBNC gives each head group its own region; the K/V split doubles heads. layout = KVCacheLayout.LHBNC diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 2c8d118cdde8..134d5055b6f2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -132,20 +132,22 @@ def _validate_kv_cache_config( from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec unsupported: list[str] = [] - cache_block_size = vllm_config.cache_config.block_size - for g_idx, g in enumerate(kv_cache_config.transfer_groups): - spec = g.kv_cache_spec + store_group_ids = kv_cache_config.prefix_cacheable_group_ids + if not store_group_ids: + raise ValueError( + "MooncakeStore requires at least one prefix-cacheable KV cache group" + ) + for group_id in store_group_ids: + spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec if isinstance(spec, CrossAttentionSpec): - unsupported.append(f"group {g_idx}: CrossAttentionSpec") - # Enforce Mamba align mode - if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size: + unsupported.append(f"group {group_id}: CrossAttentionSpec") + if isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align": unsupported.append( - f"group {g_idx}: MambaSpec with block_size=" - f"{spec.block_size} != cache_config.block_size=" - f"{cache_block_size} (mamba_cache_mode != 'align')" + f"group {group_id}: mamba_cache_mode=" + f"{spec.mamba_cache_mode!r} != 'align'" ) pcp = vllm_config.parallel_config.prefill_context_parallel_size - if len(kv_cache_config.transfer_groups) > 1 and pcp > 1: + if len(store_group_ids) > 1 and pcp > 1: unsupported.append(f"PCP > 1 (pcp={pcp}) with hybrid attention") if unsupported: raise ValueError( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 375730807d5d..d05b553b64b9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -452,7 +452,7 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: def partial_hash_hits_enabled( - kv_cache_groups: list[KVCacheGroupSpec], + kv_cache_groups: Sequence[KVCacheGroupSpec], hash_block_size: int, dcp_world_size: int = 1, ) -> bool: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index ae61710e25f7..8f277d7eea37 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -71,19 +71,27 @@ def __init__( ) self.client = LookupKeyClient(vllm_config) self.kv_cache_config = kv_cache_config + self._store_group_ids = kv_cache_config.prefix_cacheable_group_ids + # Map scheduler group IDs to store group indices. Groups outside the + # store projection, such as the QSA ring, map to None and are skipped. + self._store_group_id_by_kv_cache_group_id = { + group_id: store_group_id + for store_group_id, group_id in enumerate(self._store_group_ids) + } + store_groups = kv_cache_config.prefix_cacheable_groups # Align with the engine's own scheduler_block_size and hash_block_size. self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) self.enable_partial_hash_hits = partial_hash_hits_enabled( - kv_cache_config.kv_cache_groups, + store_groups, self._hash_block_size, vllm_config.parallel_config.decode_context_parallel_size, ) mamba_groups = { group_id: group.kv_cache_spec - for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + for group_id, group in enumerate(store_groups) if isinstance(group.kv_cache_spec, MambaSpec) } assert all( @@ -173,8 +181,8 @@ def update_state_after_alloc( """Update state after block allocation.""" local_block_ids: tuple[list[int], ...] = () if num_external_tokens > 0: - local_block_ids = self.kv_cache_config.select_transfer_block_ids( - blocks.get_block_ids() + local_block_ids = blocks.get_block_ids( + group_ids=self._store_group_ids, ) self._unfinished_requests[request.request_id] = (request, local_block_ids) @@ -239,10 +247,7 @@ def build_connector_meta( request_real = request_tuple[0] # type: ignore[index] unfolded_block_ids = tuple( - blocks.copy() - for blocks in self.kv_cache_config.select_transfer_block_ids( - request.block_ids - ) + request.block_ids[group_id].copy() for group_id in self._store_group_ids ) prefill_tokens = _new_req_prefill_tokens(request) @@ -278,8 +283,8 @@ def build_connector_meta( for i, req_id in enumerate(cached_reqs.req_ids): new_block_ids = cached_reqs.new_block_ids[i] if new_block_ids: - new_block_ids = self.kv_cache_config.select_transfer_block_ids( - new_block_ids + new_block_ids = tuple( + new_block_ids[group_id] for group_id in self._store_group_ids ) req_meta = None @@ -438,7 +443,9 @@ def _apply_current_save_block_ids( assert block_ids is not None, ( f"Missing current block table for store request {req_meta.req_id}" ) - req_meta.block_ids = block_ids + req_meta.block_ids = tuple( + block_ids[group_id] for group_id in self._store_group_ids + ) def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: """Take a GPU block reference for every store job this step emits. @@ -509,12 +516,15 @@ def register_finished_partial_tail( return False pinned_block_ids: list[int] = [] - for group_id, block_id, _ in partial_tail_offloads: - if group_id not in self._boundary_state_group_ids: + remapped_offloads: list[tuple[int, int, int]] = [] + for group_id, block_id, boundary in partial_tail_offloads: + store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) + if store_group_id not in self._boundary_state_group_ids: return False if block_id == NULL_BLOCK_ID: return False pinned_block_ids.append(block_id) + remapped_offloads.append((store_group_id, block_id, boundary)) pinned_block_ids = list(dict.fromkeys(pinned_block_ids)) pool = self._gpu_block_pool @@ -530,12 +540,14 @@ def register_finished_partial_tail( self._finished_partial_tail_metas[request.request_id] = ReqMeta( req_id=request.request_id, token_len_chunk=0, - block_ids=tuple(group.copy() for group in block_ids), + block_ids=tuple( + block_ids[group_id].copy() for group_id in self._store_group_ids + ), block_hashes=list(request.block_hashes), can_save=True, num_prompt_tokens=tracker.prefill_end_tokens, store_job_id=store_job_id, - boundary_state_offloads=partial_tail_offloads, + boundary_state_offloads=remapped_offloads, ) tracker.has_pending_offload = True # The store job owns exact block refs, so request cleanup need not wait. @@ -574,9 +586,10 @@ def _handle_boundary_state_offloads( continue if block_id == NULL_BLOCK_ID: continue - if group_id not in self._boundary_state_group_ids: + store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) + if store_group_id not in self._boundary_state_group_ids: continue - accepted.append((group_id, block_id, boundary_tokens)) + accepted.append((store_group_id, block_id, boundary_tokens)) if not accepted: continue tracker.has_pending_offload = True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index ebcf5708c89d..6f154ea991f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1672,10 +1672,11 @@ def __init__( dataclasses.replace( group, kv_cache_spec=resolve_dcp_kv_cache_spec( - group.kv_cache_spec, self.dcp_size + group.kv_cache_spec, + self.dcp_size, ), ) - for group in kv_cache_config.transfer_groups + for group in kv_cache_config.prefix_cacheable_groups ] spec_cfg = getattr(vllm_config, "speculative_config", None) use_eagle_block_drop = bool( @@ -1923,10 +1924,16 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None if not use_group_regions: + store_layer_names = { + layer_name + for group in self._kv_cache_groups + for layer_name in group.layer_names + } seen_storage_ptrs: set[int] = set() cache_tensors = [ group_kernel_blocks(_repr_tensor(cache), self.num_blocks) - for cache in kv_caches.values() + for layer_name, cache in kv_caches.items() + if layer_name in store_layer_names ] for cache in cache_tensors: cache_storage = cache.untyped_storage() diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 20e037f60ea1..226da3a1ed88 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -67,30 +67,45 @@ def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks": def get_block_ids( self, allow_none: Literal[False] = False, + *, + group_ids: Sequence[int] | None = None, ) -> tuple[list[int], ...]: ... @overload def get_block_ids( self, allow_none: Literal[True] = True, + *, + group_ids: Sequence[int] | None = None, ) -> tuple[list[int], ...] | None: ... def get_block_ids( self, allow_none: bool = False, + *, + group_ids: Sequence[int] | None = None, ) -> tuple[list[int], ...] | None: """ Converts the KVCacheBlocks instance to block_ids. + Args: + allow_none: Return None when every selected group is empty. + group_ids: KV cache groups to include. Includes all groups by default. + Returns: tuple[list[int], ...]: A tuple of lists where: - the outer tuple corresponds to KV cache groups - each inner list contains the block_ids of the blocks in that group """ - if allow_none and all(len(group) == 0 for group in self.blocks): + groups = ( + self.blocks + if group_ids is None + else tuple(self.blocks[group_id] for group_id in group_ids) + ) + if allow_none and all(len(group) == 0 for group in groups): return None - return tuple([blk.block_id for blk in group] for group in self.blocks) + return tuple([blk.block_id for blk in group] for group in groups) def get_unhashed_block_ids(self) -> list[int]: """Get block_ids of unhashed blocks from KVCacheBlocks instance.""" diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index e2869e047f54..8cd867eb9c60 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1450,6 +1450,23 @@ def transfer_group_ids(self) -> tuple[int, ...]: if group.enable_kv_transfer ) + @cached_property + def prefix_cacheable_group_ids(self) -> tuple[int, ...]: + """IDs of transferable groups eligible for hash-addressed stores.""" + return tuple( + group_id + for group_id in self.transfer_group_ids + if self.kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable + ) + + @cached_property + def prefix_cacheable_groups(self) -> tuple[KVCacheGroupSpec, ...]: + """Transferable groups eligible for hash-addressed stores.""" + return tuple( + self.kv_cache_groups[group_id] + for group_id in self.prefix_cacheable_group_ids + ) + @cached_property def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]: """Cache groups that participate in external KV transfer.""" From dfa1984e58dc898610a320de26f94b6039fccec3 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 13 Sep 2026 22:36:27 -0400 Subject: [PATCH 147/774] [ROCm][Perf] Tune MiniMax-M3 decode top-k for short contexts (#55235) Signed-off-by: fai --- tests/kernels/attention/test_minimax_m3.py | 104 +++++++++++-------- vllm/models/minimax_m3/amd/ops/index_topk.py | 34 ++++-- 2 files changed, 90 insertions(+), 48 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index c45b87d04b1f..a7e6d499c176 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -1025,17 +1025,20 @@ def test_decode_index_topk_correctness( "expected", ), [ - (0, 24, 1, 16, True, (2, False, False)), - (1, 1, 1, 16, True, (16, True, True)), - (512, 4, 1, 16, True, (16, True, True)), - (513, 24, 4, 16, True, (16, True, True)), - (2760, 5, 2, 16, True, (16, True, True)), - (8192, 64, 4, 16, True, (16, True, True)), - (8193, 4, 1, 16, True, (16, False, False)), - (8193, 16, 1, 16, True, (4, False, False)), - (65534, 64, 1, 16, True, (1, False, False)), - (513, 16, 1, 13, True, (4, False, False)), - (513, 16, 4, 16, False, (1, False, False)), + (0, 24, 1, 16, True, (2, 512, 8, 2, False, False)), + (1, 1, 1, 16, True, (1, 128, 2, 1, True, True)), + (74, 4, 1, 16, True, (1, 128, 2, 1, True, True)), + (128, 4, 1, 16, True, (1, 128, 2, 1, True, True)), + (129, 4, 1, 16, True, (16, 512, 4, 2, True, True)), + (512, 4, 1, 16, True, (16, 512, 4, 2, True, True)), + (513, 24, 4, 16, True, (16, 512, 4, 2, True, True)), + (2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)), + (8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)), + (8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)), + (8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)), + (65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)), + (513, 16, 1, 13, True, (4, 512, 8, 2, False, False)), + (513, 16, 4, 16, False, (1, 512, 8, 2, False, False)), ], ) def test_amd_decode_topk_launch_policy( @@ -1044,9 +1047,9 @@ def test_amd_decode_topk_launch_policy( num_idx_heads: int, topk: int, is_gfx950: bool, - expected: tuple[int, bool, bool], + expected: tuple[int, int, int, int, bool, bool], ): - """All supported head and query layouts use adaptive selection.""" + """The short-context specialization has strict dispatch boundaries.""" from vllm.models.minimax_m3.amd.ops.index_topk import ( _decode_topk_launch_policy, ) @@ -1073,19 +1076,23 @@ def test_amd_decode_topk_launch_policy( "single_tile_guaranteed", "adaptive_final_merge", "max_blocks", + "block_size_k", + "num_warps", + "num_stages", "topk", "num_idx_heads", ), [ - (1, False, False, 65534, 16, 1), - (2, False, False, 8193, 16, 1), - (4, False, False, 8193, 16, 1), - (8, False, False, 8193, 16, 1), - (16, False, False, 8193, 16, 1), - (16, False, False, 65534, 16, 1), - (2, True, False, 513, 13, 1), - (4, False, True, 8192, 16, 2), - (16, True, True, 8192, 16, 4), + (1, False, False, 65534, 512, 8, 2, 16, 1), + (2, False, False, 8193, 512, 8, 2, 16, 1), + (4, False, False, 8193, 512, 8, 2, 16, 1), + (8, False, False, 8193, 512, 8, 2, 16, 1), + (16, False, False, 8193, 512, 8, 2, 16, 1), + (16, False, False, 65534, 512, 8, 2, 16, 1), + (2, True, False, 513, 512, 8, 2, 13, 1), + (4, False, True, 8192, 512, 4, 2, 16, 2), + (16, True, True, 8192, 512, 4, 2, 16, 4), + (1, True, True, 74, 128, 2, 1, 16, 1), ], ) def test_amd_decode_fused_topk_total_order_and_replay( @@ -1093,10 +1100,13 @@ def test_amd_decode_fused_topk_total_order_and_replay( single_tile_guaranteed: bool, adaptive_final_merge: bool, max_blocks: int, + block_size_k: int, + num_warps: int, + num_stages: int, topk: int, num_idx_heads: int, ): - """Atomic multi-chunk selection is ordered and resets graph state.""" + """Selector specializations preserve ordering and graph state.""" from vllm.models.minimax_m3.amd.ops.index_topk import ( SPARSE_BLOCK_SIZE, _decode_topk_fused_kernel, @@ -1104,18 +1114,25 @@ def test_amd_decode_fused_topk_total_order_and_replay( set_random_seed(0) block_size_t = 1 << (topk - 1).bit_length() + block_counts = tuple( + count + for count in ( + 15, + 16, + 17, + block_size_k // 2 - 1, + block_size_k // 2, + block_size_k // 2 + 1, + block_size_k - 1, + block_size_k, + block_size_k + 1, + ) + if count <= max_blocks + ) long_seq_lens_list = [ 0, 1, - 15 * SPARSE_BLOCK_SIZE, - 16 * SPARSE_BLOCK_SIZE, - 17 * SPARSE_BLOCK_SIZE, - 255 * SPARSE_BLOCK_SIZE, - 256 * SPARSE_BLOCK_SIZE, - 257 * SPARSE_BLOCK_SIZE, - 511 * SPARSE_BLOCK_SIZE, - 512 * SPARSE_BLOCK_SIZE, - 513 * SPARSE_BLOCK_SIZE, + *(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts), max_blocks * SPARSE_BLOCK_SIZE, ] short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129] @@ -1126,12 +1143,13 @@ def test_amd_decode_fused_topk_total_order_and_replay( boundary_limit = 4 if boundary_limit: boundary_blocks = [ - (active_chunks - 1) * 512 + 1 + (active_chunks - 1) * block_size_k + 1 for active_chunks in range(3, boundary_limit + 1) ] - assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list( - range(3, boundary_limit + 1) - ) + assert [ + (num_blocks + block_size_k - 1) // block_size_k + for num_blocks in boundary_blocks + ] == list(range(3, boundary_limit + 1)) long_seq_lens_list.extend( num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks ) @@ -1139,6 +1157,8 @@ def test_amd_decode_fused_topk_total_order_and_replay( num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks) ) batch = len(long_seq_lens_list) + repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list) + short_seq_lens_list = (short_seq_lens_list * repeats)[:batch] score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda") score[:, :, 0] = float("nan") score[:, :, 1] = float("inf") @@ -1147,8 +1167,9 @@ def test_amd_decode_fused_topk_total_order_and_replay( score[:, :, 4] = -0.0 score[:, :, 5] = 7.0 score[:, :, 6] = -1e30 - score[:, :, 256] = 7.0 - score[:, :, 512] = 7.0 + for index in (block_size_k // 2, block_size_k): + if index < max_blocks: + score[:, :, index] = 7.0 seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32) partial = torch.full( @@ -1198,13 +1219,13 @@ def launch() -> None: pages_per_sparse_block=8, block_page_stride=8, NUM_TOPK_CHUNKS=num_topk_chunks, - BLOCK_SIZE_K=512, + BLOCK_SIZE_K=block_size_k, BLOCK_SIZE_T=block_size_t, EMIT_SPARSE_TABLE=False, SINGLE_TILE_GUARANTEED=single_tile_guaranteed, ADAPTIVE_FINAL_MERGE=adaptive_final_merge, - num_warps=4 if adaptive_final_merge else 8, - num_stages=2, + num_warps=num_warps, + num_stages=num_stages, ) score_cpu = score.cpu() @@ -1251,6 +1272,7 @@ def reference(seq_lens_list: list[int]) -> torch.Tensor: (16, False, False, 8193), (4, False, True, 8192), (16, True, True, 8192), + (1, True, True, 74), ): graph = torch.cuda.CUDAGraph() torch.accelerator.synchronize() diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index b5c75371f016..ce27ea122057 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -37,6 +37,7 @@ MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64 MAX_DECODE_SCORE_KV_CHUNKS = 256 DECODE_TOPK_BLOCKS_PER_CHUNK = 512 +DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128 MAX_DECODE_TOPK_FAST_CHUNKS = 16 DECODE_TOPK_TARGET_GRID = 64 @@ -103,14 +104,23 @@ def _decode_topk_launch_policy( topk: int, *, is_gfx950: bool, -) -> tuple[int, bool, bool]: - """Choose the graph grid and bounded-context selector specialization.""" +) -> tuple[int, int, int, int, bool, bool]: + """Choose the selector grid and compile-time launch configuration.""" if ( is_gfx950 and topk == 16 and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK ): - return MAX_DECODE_TOPK_FAST_CHUNKS, True, True + if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK: + return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True + return ( + MAX_DECODE_TOPK_FAST_CHUNKS, + DECODE_TOPK_BLOCKS_PER_CHUNK, + 4, + 2, + True, + True, + ) target = max( 1, @@ -119,7 +129,14 @@ def _decode_topk_launch_policy( DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads), ), ) - return 1 << (target.bit_length() - 1), False, False + return ( + 1 << (target.bit_length() - 1), + DECODE_TOPK_BLOCKS_PER_CHUNK, + 8, + 2, + False, + False, + ) # --------------------------------------------------------------------------- @@ -1525,6 +1542,9 @@ def minimax_m3_index_decode( # only the number of chunks needed for its live context. ( num_topk_chunks, + topk_block_size, + selector_num_warps, + selector_num_stages, single_tile_guaranteed, adaptive_final_merge, ) = _decode_topk_launch_policy( @@ -1608,12 +1628,12 @@ def minimax_m3_index_decode( pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK, block_page_stride=selector_block_page_stride, NUM_TOPK_CHUNKS=num_topk_chunks, - BLOCK_SIZE_K=512, + BLOCK_SIZE_K=topk_block_size, BLOCK_SIZE_T=block_size_t, EMIT_SPARSE_TABLE=emit_sparse_table, SINGLE_TILE_GUARANTEED=single_tile_guaranteed, ADAPTIVE_FINAL_MERGE=adaptive_final_merge, - num_warps=4 if adaptive_final_merge else 8, - num_stages=2, + num_warps=selector_num_warps, + num_stages=selector_num_stages, ) return topk_idx From eb42686a30cddf325ffeff7b3bd5e3a7298c2c00 Mon Sep 17 00:00:00 2001 From: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:43:09 +0800 Subject: [PATCH 148/774] [Bugfix][Frontend] Support Responses text types in DeepSeek V4.1 (#56299) Signed-off-by: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/tokenizers_/test_deepseek_v41.py | 18 ++++++++++++++++++ vllm/tokenizers/deepseek_v41.py | 2 +- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index 493fdd88ba5b..13c0375fbb0a 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/tests/tokenizers_/test_deepseek_v41.py @@ -99,6 +99,24 @@ def test_raw_text_parts_preserve_reference_separator(): assert messages == original +@pytest.mark.parametrize( + ("role", "responses_type"), + [("user", "input_text"), ("assistant", "output_text")], +) +def test_responses_text_parts_match_chat_text_parts(role, responses_type): + responses_message = { + "role": role, + "content": [{"type": responses_type, "text": "hello"}], + } + chat_message = { + "role": role, + "content": [{"type": "text", "text": "hello"}], + } + assert render([responses_message], thinking=False) == render( + [chat_message], thinking=False + ) + + def test_mid_system_gets_its_own_marker_and_generation_header(): assert render( [ diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index 3a577172708e..a1449e59e203 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -32,7 +32,7 @@ def _normalize_messages( parts = [] for block in content: part_type = block.get("type") - if part_type == "text": + if part_type in ("text", "input_text", "output_text"): parts.append(block.get("text", "")) elif part_type in ("image_url", "input_image", "image_pil"): parts.append(IMAGE_PLACEHOLDER) From 52dd0d7562adb3c2d556ce6fe7ca7c3226e1976b Mon Sep 17 00:00:00 2001 From: siyu Date: Mon, 14 Sep 2026 11:05:54 +0800 Subject: [PATCH 149/774] [Fast Start] Fast loader support nnode>1 (#55468) Signed-off-by: siyu Co-authored-by: Isotr0py <2037008807@qq.com> --- .../model_loader/weight_cache/daemon.py | 86 ++++++++++++++++--- 1 file changed, 72 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 31f59a6ac5ff..96cedbe1e91d 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -17,8 +17,30 @@ vllm serve /path/to/model --tensor-parallel-size 4 \\ --load-format ipc_cache -Only tensor and expert parallelism are supported; pipeline and data -parallelism are rejected at launch. +Only tensor and expert parallelism are supported; pipeline and data parallelism +are rejected at launch. + +For multi-node tensor parallelism, run one launcher per node with a shared +rendezvous so the global TP group forms across nodes (CUDA IPC handles are +node-local, so each node serves only its local GPUs' shards). Reuse the same +``--nnodes``/``--node-rank``/``--master-addr`` flags you pass the engine, plus a +``--weight-cache-master-port`` distinct from the engine's ``--master-port``: + + # node 0 (8 local GPUs) + python -m vllm.model_executor.model_loader.weight_cache.daemon \\ + --model /path/to/model --tensor-parallel-size 16 \\ + --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\ + --weight-cache-master-port 29600 + # node 1 (8 local GPUs) + python -m vllm.model_executor.model_loader.weight_cache.daemon \\ + --model /path/to/model --tensor-parallel-size 16 \\ + --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\ + --weight-cache-master-port 29600 + +The global TP rank of local GPU ``i`` on node ``r`` is +``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank +assignment, so each engine worker maps its shard from the daemon on its own +node. """ import contextlib @@ -132,11 +154,13 @@ def __init__( self, vllm_config: VllmConfig, tp_rank: int, + local_rank: int, distributed_init_method: str, socket_dir: str | None = None, ): self.vllm_config = vllm_config self.tp_rank = tp_rank + self.local_rank = local_rank self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir self.model: torch.nn.Module | None = None @@ -152,12 +176,12 @@ def __init__( def load_model(self) -> None: tp_size = self.cache_config.tp_size - torch.accelerator.set_device_index(self.tp_rank) + torch.accelerator.set_device_index(self.local_rank) init_distributed_environment( world_size=tp_size, rank=self.tp_rank, distributed_init_method=self.distributed_init_method, - local_rank=self.tp_rank, + local_rank=self.local_rank, backend=current_platform.dist_backend, ) with set_current_vllm_config(self.vllm_config): @@ -307,13 +331,14 @@ def _gpu_uuid(self) -> str: def _run_daemon( tp_rank: int, + local_rank: int, vllm_config: VllmConfig, distributed_init_method: str, socket_dir: str | None, ready_queue: "multiprocessing.Queue[int]", ) -> None: daemon = WeightCacheDaemon( - vllm_config, tp_rank, distributed_init_method, socket_dir + vllm_config, tp_rank, local_rank, distributed_init_method, socket_dir ) daemon.load_model() daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank)) @@ -348,6 +373,15 @@ def main() -> None: default=None, help="Directory for the daemon Unix sockets (default: tempdir).", ) + parser.add_argument( + "--weight-cache-master-port", + type=int, + default=None, + help="Rendezvous port for the daemon's own TP group. Must differ from " + "the engine's --master-port (the daemon holds its group open while " + "serving) and match across nodes. Required when --nnodes > 1; defaults " + "to a free port for single-node.", + ) args = parser.parse_args() engine_args = EngineArgs.from_cli_args(args) vllm_config = engine_args.create_engine_config() @@ -363,22 +397,41 @@ def main() -> None: _reject_unsupported_parallelism(parallel_config) tp_size = parallel_config.tensor_parallel_size - distributed_init_method = get_distributed_init_method("127.0.0.1", get_open_port()) + nnodes = parallel_config.nnodes + node_rank = parallel_config.node_rank + local_world_size = tp_size // nnodes + + # The daemon forms its own TP group and holds it open while serving, so it + # needs a rendezvous port distinct from the engine's --master-port. All + # nodes must agree on it; single-node can auto-pick a free port. The master + # address is reused from the engine's --master-addr (node_rank 0). + master_addr = parallel_config.master_addr + if nnodes > 1 and args.weight_cache_master_port is None: + raise ValueError("--weight-cache-master-port is required when --nnodes > 1") + master_port = args.weight_cache_master_port or get_open_port() + distributed_init_method = get_distributed_init_method(master_addr, master_port) + ctx = multiprocessing.get_context("spawn") ready_queue: multiprocessing.Queue[int] = ctx.Queue() + # Global TP rank of local GPU i on this node; local index == device index. + global_ranks = [ + node_rank * local_world_size + local_rank + for local_rank in range(local_world_size) + ] procs = [ ctx.Process( target=_run_daemon, args=( - rank, + global_rank, + local_rank, vllm_config, distributed_init_method, args.weight_cache_socket_dir, ready_queue, ), - name=f"vllm-weight-cache-daemon-{rank}", + name=f"vllm-weight-cache-daemon-{global_rank}", ) - for rank in range(tp_size) + for local_rank, global_rank in enumerate(global_ranks) ] for proc in procs: proc.start() @@ -391,7 +444,7 @@ def _shutdown(signum, frame): signal.signal(signal.SIGTERM, _shutdown) ready_ranks: set[int] = set() - while len(ready_ranks) < tp_size: + while len(ready_ranks) < local_world_size: try: ready_ranks.add(ready_queue.get(timeout=1.0)) except queue.Empty: @@ -407,11 +460,16 @@ def _shutdown(signum, frame): for proc in procs: proc.join() sys.exit(max((p.exitcode or 0) for p in procs)) - logger.info( - "Weight cache daemon ready: all %d ranks serving in %s", - tp_size, - args.weight_cache_socket_dir or "the default socket dir", + socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir" + logger.info_once( + "===== Weight cache daemon READY: node %d/%d serving %d local rank(s) " + "in %s =====", + node_rank, + nnodes, + local_world_size, + socket_dir_msg, ) + for proc in procs: proc.join() sys.exit(max(proc.exitcode or 0 for proc in procs)) From 2d9c30a7b545e09fd118b2b9e92b5702028d7c1d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 13 Sep 2026 20:30:25 -0700 Subject: [PATCH 150/774] [CI][Bugfix] Complete FSE fixture R-SWA contract (#56710) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/model_executor/layers/test_fused_shared_expert.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 37693f18c436..507fd93e9d02 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -832,6 +832,7 @@ def test_models_fse_init( runner_type="generate", is_moe=True, logits_processors=None, + rswa_window=None, ) vllm_config.parallel_config.enable_expert_parallel = False if model_type == "deepseek_v4": From 7632f767f0d896a902566acab29775816b39a92d Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Mon, 14 Sep 2026 11:37:04 +0800 Subject: [PATCH 151/774] [Rust Frontend] Support HTTP RL weight synchronization (#56567) Co-authored-by: Bugen Zhao Co-authored-by: Codex Signed-off-by: Wang Zupeng Signed-off-by: Bugen Zhao --- rust/README.md | 11 + rust/src/server/src/error.rs | 7 + rust/src/server/src/routes.rs | 23 ++ rust/src/server/src/routes/abort_requests.rs | 2 +- rust/src/server/src/routes/collective_rpc.rs | 2 +- .../src/routes/openai/utils/validated_json.rs | 5 +- rust/src/server/src/routes/tests.rs | 249 ++++++++++++++++++ rust/src/server/src/routes/weight_transfer.rs | 190 +++++++++++++ 8 files changed, 483 insertions(+), 6 deletions(-) create mode 100644 rust/src/server/src/routes/weight_transfer.rs diff --git a/rust/README.md b/rust/README.md index 06b50dbbfe9c..0cc33d64f8d2 100644 --- a/rust/README.md +++ b/rust/README.md @@ -43,6 +43,17 @@ For example: VLLM_USE_RUST_FRONTEND=1 vllm serve Qwen/Qwen3-0.6B ``` +### RL weight synchronization + +With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer +lifecycle used by `HTTPVLLMWeightSyncClient`: initialization, starting an update, +transferring weights, and finishing the update. It also supports draft-model +updates and `/update_weight_version` and `/weight_info` for version tracking. + +For trainer-side usage, see the [IPC](../examples/rl/rlhf_http_ipc.py) and +[NCCL](../examples/rl/rlhf_http_nccl.py) examples. Configure the weight-transfer +backend on the server and pause generation during updates, as shown there. + ### External Engine `vllm-rs serve` can be run standalone with `--data-parallel-size-local 0` when the Python engines diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index b658388cbdd4..9f083a1a614e 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use axum::Json; +use axum::extract::rejection::JsonRejection; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use thiserror_ext::{AsReport as _, Construct, Macro}; @@ -70,6 +71,12 @@ impl ApiError { } } +impl From for ApiError { + fn from(error: JsonRejection) -> Self { + Self::json_parse_error(error.body_text()) + } +} + impl IntoResponse for ApiError { fn into_response(self) -> Response { (self.status_code(), Json(self.to_error_response())).into_response() diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 393a08450fbc..e7e91938bbda 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -17,6 +17,7 @@ mod server_info; mod sleep; mod tokenize; mod version; +mod weight_transfer; mod world_size; use std::sync::Arc; @@ -139,6 +140,28 @@ fn build_router_with_options( .route("/reset_mm_cache", post(cache::reset_mm_cache)) .route("/reset_encoder_cache", post(cache::reset_encoder_cache)) .route("/collective_rpc", post(collective_rpc::collective_rpc)) + .route( + "/init_weight_transfer_engine", + post(weight_transfer::init_weight_transfer_engine), + ) + .route( + "/start_weight_update", + post(weight_transfer::start_weight_update), + ) + .route( + "/start_draft_weight_update", + post(weight_transfer::start_draft_weight_update), + ) + .route("/update_weights", post(weight_transfer::update_weights)) + .route( + "/finish_weight_update", + post(weight_transfer::finish_weight_update), + ) + .route( + "/update_weight_version", + post(weight_transfer::update_weight_version), + ) + .route("/weight_info", get(weight_transfer::weight_info)) .route("/abort_requests", post(abort_requests::abort_requests)) .route("/sleep", post(sleep::sleep)) .route("/wake_up", post(sleep::wake_up)) diff --git a/rust/src/server/src/routes/abort_requests.rs b/rust/src/server/src/routes/abort_requests.rs index 298770316099..d08e5af3fa31 100644 --- a/rust/src/server/src/routes/abort_requests.rs +++ b/rust/src/server/src/routes/abort_requests.rs @@ -22,7 +22,7 @@ pub async fn abort_requests( State(state): State>, body: Result, JsonRejection>, ) -> Result { - let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?; + let Json(body) = body?; // Empty/missing `request_ids` aborts all in-flight requests. let request_ids = body.request_ids.unwrap_or_default(); diff --git a/rust/src/server/src/routes/collective_rpc.rs b/rust/src/server/src/routes/collective_rpc.rs index bd769bee43a9..fa3782cbe9bf 100644 --- a/rust/src/server/src/routes/collective_rpc.rs +++ b/rust/src/server/src/routes/collective_rpc.rs @@ -36,7 +36,7 @@ pub async fn collective_rpc( State(state): State>, body: Result, JsonRejection>, ) -> Result, ApiError> { - let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?; + let Json(body) = body?; let method = body.method.ok_or_else(|| { ApiError::invalid_request( "Missing 'method' in request body".to_string(), diff --git a/rust/src/server/src/routes/openai/utils/validated_json.rs b/rust/src/server/src/routes/openai/utils/validated_json.rs index 1ddec3a9f521..a9b68bc3b407 100644 --- a/rust/src/server/src/routes/openai/utils/validated_json.rs +++ b/rust/src/server/src/routes/openai/utils/validated_json.rs @@ -5,7 +5,6 @@ //! Variation of https://github.com/lightseekorg/smg/blob/main/crates/protocols/src/validated.rs use axum::Json; -use axum::extract::rejection::JsonRejection; use axum::extract::{FromRequest, Request}; use serde::de::DeserializeOwned; use validator::Validate; @@ -30,9 +29,7 @@ where type Rejection = ApiError; async fn from_request(req: Request, state: &S) -> Result { - let Json(mut data) = Json::::from_request(req, state) - .await - .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?; + let Json(mut data) = Json::::from_request(req, state).await?; data.normalize(); diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 210ae044a026..2851ec1e88b8 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -6113,6 +6113,13 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { ("POST", "/reset_prefix_cache"), ("POST", "/reset_mm_cache"), ("POST", "/reset_encoder_cache"), + ("POST", "/init_weight_transfer_engine"), + ("POST", "/start_weight_update"), + ("POST", "/start_draft_weight_update"), + ("POST", "/update_weights"), + ("POST", "/finish_weight_update"), + ("POST", "/update_weight_version"), + ("GET", "/weight_info"), ] { let response = app .clone() @@ -6132,6 +6139,248 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { engine_task.abort_and_join().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn weight_transfer_routes_support_the_http_training_lifecycle() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let mut calls = Vec::new(); + for _ in 0..10 { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + let payload = decode_value(&utility[1]).expect("decode utility"); + let array = payload.as_array().expect("utility array"); + let call_id = array[1].as_u64().expect("call id"); + let args: serde_json::Value = + rmpv::ext::from_value(array[3].clone()).expect("decode args"); + calls.push(format!("{} {args}", array[2].as_str().expect("method"))); + let result = if array[2] == Value::from("get_weight_version") { + utility_result_value("step-2") + } else if array[2] == Value::from("collective_rpc") { + utility_result_value(vec![(), ()]) + } else { + utility_none_result() + }; + send_outputs(push, utility_outputs(call_id, result)).await; + } + expect_test::expect![[r#" + [ + "collective_rpc [\"init_weight_transfer_engine\",null,[],{\"init_info\":{\"master_address\":\"127.0.0.1\",\"master_port\":29500,\"rank_offset\":1,\"world_size\":3}}]", + "collective_rpc [\"start_weight_update\",null,[],{}]", + "collective_rpc [\"update_weights\",null,[],{\"update_info\":{\"names\":[\"model.embed_tokens.weight\"],\"dtype_names\":[\"bfloat16\"],\"shapes\":[[32,16]],\"is_checkpoint_format\":true}}]", + "collective_rpc [\"update_weights\",null,[],{\"update_info\":[{\"rank\":0,\"ipc_handles_pickled\":\"opaque-handle-0\"},{\"rank\":1,\"ipc_handles_pickled\":\"opaque-handle-1\"}]}]", + "collective_rpc [\"finish_weight_update\",null,[],{}]", + "set_weight_version [\"step-1\"]", + "collective_rpc [\"start_draft_weight_update\",null,[],{}]", + "collective_rpc [\"finish_weight_update\",null,[],{}]", + "set_weight_version [\"step-2\"]", + "get_weight_version []", + ] + "#]].assert_debug_eq(&calls); + }) + }) + .await; + + let mut responses = Vec::new(); + for (method, path, body) in [ + ( + "POST", + "/init_weight_transfer_engine", + Some(json!({ + "init_info": {"master_address": "127.0.0.1", "master_port": 29500, + "rank_offset": 1, "world_size": 3} + })), + ), + ("POST", "/start_weight_update", None), + ( + "POST", + "/update_weights", + Some(json!({ + "update_info": {"names": ["model.embed_tokens.weight"], + "dtype_names": ["bfloat16"], "shapes": [[32, 16]], + "is_checkpoint_format": true} + })), + ), + ( + "POST", + "/update_weights", + Some(json!({ + "update_info": [{"rank": 0, "ipc_handles_pickled": "opaque-handle-0"}, + {"rank": 1, "ipc_handles_pickled": "opaque-handle-1"}] + })), + ), + ( + "POST", + "/finish_weight_update", + Some(json!({"weight_version": "step-1"})), + ), + ("POST", "/start_draft_weight_update", None), + ("POST", "/finish_weight_update", None), + ( + "POST", + "/update_weight_version", + Some(json!({"new_version": "step-2"})), + ), + ("GET", "/weight_info", None), + ] { + let mut request = Request::builder().method(method).uri(path); + let body = match body { + Some(body) => { + request = request.header("content-type", "application/json"); + Body::from(body.to_string()) + } + None => Body::empty(), + }; + let response = + app.call(request.body(body).expect("build request")).await.expect("call app"); + assert_eq!(response.status(), StatusCode::OK, "{path}"); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + responses.push( + serde_json::from_slice::(&body) + .expect("decode json") + .to_string(), + ); + } + engine_task.await.expect("mock engine task"); + expect_test::expect![[r#" + [ + "{\"message\":\"Weight transfer initialized\"}", + "{\"message\":\"Weight update started\"}", + "{\"message\":\"Weights updated\"}", + "{\"message\":\"Weights updated\"}", + "{\"message\":\"Weight update finished\"}", + "{\"message\":\"Draft weight update started\"}", + "{\"message\":\"Weight update finished\"}", + "{\"success\":true,\"new_version\":\"step-2\"}", + "{\"weight_version\":\"step-2\"}", + ] + "#]] + .assert_debug_eq(&responses); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, _push| { + boxed_test_future(async move { + let message = recv_engine_message(dealer).await; + panic!("invalid payload reached engine: {message:?}"); + }) + }) + .await; + + for (path, body) in [ + ("/init_weight_transfer_engine", "{"), + ("/init_weight_transfer_engine", "{}"), + ("/init_weight_transfer_engine", r#"{"init_info":null}"#), + ("/init_weight_transfer_engine", r#"{"init_info":[]}"#), + ("/update_weights", "{"), + ("/update_weights", "{}"), + ("/update_weights", r#"{"update_info":null}"#), + ("/update_weights", r#"{"update_info":1}"#), + ("/update_weights", r#"{"update_info":[{},null]}"#), + ("/finish_weight_update", "{"), + ("/finish_weight_update", r#"{"weight_version":1}"#), + ("/update_weight_version", "{}"), + ("/update_weight_version", r#"{"new_version":null}"#), + ] { + let response = app + .call( + Request::builder() + .method("POST") + .uri(path) + .header("content-type", "application/json") + .body(Body::from(body)) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}"); + } + engine_task.abort_and_join().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn weight_transfer_routes_do_not_commit_version_when_finish_fails() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]).expect("decode utility"); + let array = payload.as_array().expect("utility array"); + assert_eq!(array[2], Value::from("collective_rpc")); + assert_eq!( + array[3].as_array().expect("args")[0], + Value::from("finish_weight_update") + ); + send_outputs( + push, + UtilityCallOutput { + output: UtilityOutput { + call_id: array[1].as_u64().expect("call id").into(), + failure_message: Some("weight reload failed".to_string()), + result: None, + }, + ..Default::default() + } + .into(), + ) + .await; + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]).expect("decode utility"); + let array = payload.as_array().expect("utility array"); + assert_eq!(array[2], Value::from("get_weight_version")); + send_outputs( + push, + utility_outputs( + array[1].as_u64().expect("call id"), + utility_result_value("step-0"), + ), + ) + .await; + }) + }) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/finish_weight_update") + .header("content-type", "application/json") + .body(Body::from(r#"{"weight_version":"step-1"}"#)) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let error: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert!( + error["error"]["message"] + .as_str() + .expect("error message") + .contains("weight reload failed") + ); + + let response = app + .call( + Request::builder() + .uri("/weight_info") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + assert_eq!( + serde_json::from_slice::(&body).expect("decode json"), + json!({"weight_version": "step-0"}) + ); + engine_task.await.expect("mock engine task"); +} + // ========================= Stop string tests ========================= #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs new file mode 100644 index 000000000000..df55e25593af --- /dev/null +++ b/rust/src/server/src/routes/weight_transfer.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use axum::Json; +use axum::extract::State; +use axum::extract::rejection::JsonRejection; +use serde::{Deserialize, Serialize}; +use serde_json::Value as JsonValue; + +use crate::error::{ApiError, invalid_request}; +use crate::state::AppState; +use crate::utils::utility_call_error; + +#[derive(Deserialize)] +pub(crate) struct InitWeightTransferRequest { + init_info: Option, +} + +#[derive(Deserialize)] +pub(crate) struct UpdateWeightsRequest { + update_info: Option, +} + +#[derive(Default, Deserialize)] +pub(crate) struct FinishWeightUpdateRequest { + weight_version: Option, +} + +#[derive(Deserialize)] +pub(crate) struct UpdateWeightVersionRequest { + new_version: String, +} + +#[derive(Serialize)] +pub(crate) struct MessageResponse { + message: &'static str, +} + +#[derive(Serialize)] +pub(crate) struct UpdateWeightVersionResponse { + success: bool, + new_version: String, +} + +#[derive(Serialize)] +pub(crate) struct WeightInfoResponse { + weight_version: String, +} + +/// Initialize weight transfer for RL training. +pub async fn init_weight_transfer_engine( + State(state): State>, + body: Result, JsonRejection>, +) -> Result, ApiError> { + let Json(body) = body?; + let init_info = body.init_info.filter(JsonValue::is_object).ok_or_else(|| { + invalid_request!( + param = Some("init_info"), + "'init_info' must be a JSON object" + ) + })?; + + state + .engine_core_client() + .init_weight_transfer_engine(init_info) + .await + .map_err(|error| utility_call_error("init_weight_transfer_engine", error))?; + + Ok(Json(MessageResponse { + message: "Weight transfer initialized", + })) +} + +/// Start a weight update transaction. +pub async fn start_weight_update( + State(state): State>, +) -> Result, ApiError> { + state + .engine_core_client() + .start_weight_update() + .await + .map_err(|error| utility_call_error("start_weight_update", error))?; + + Ok(Json(MessageResponse { + message: "Weight update started", + })) +} + +/// Start a draft-model weight update transaction. +pub async fn start_draft_weight_update( + State(state): State>, +) -> Result, ApiError> { + state + .engine_core_client() + .start_draft_weight_update() + .await + .map_err(|error| utility_call_error("start_draft_weight_update", error))?; + + Ok(Json(MessageResponse { + message: "Draft weight update started", + })) +} + +/// Update model weights with backend-specific transfer metadata. +pub async fn update_weights( + State(state): State>, + body: Result, JsonRejection>, +) -> Result, ApiError> { + let Json(body) = body?; + let update_info = body + .update_info + .filter(|info| { + info.is_object() + || info.as_array().is_some_and(|items| items.iter().all(JsonValue::is_object)) + }) + .ok_or_else(|| { + invalid_request!( + param = Some("update_info"), + "'update_info' must be a JSON object or a list of per-worker JSON objects", + ) + })?; + + state + .engine_core_client() + .update_weights(update_info) + .await + .map_err(|error| utility_call_error("update_weights", error))?; + + Ok(Json(MessageResponse { + message: "Weights updated", + })) +} + +/// Finish a weight update transaction and optionally set the weight version. +pub async fn finish_weight_update( + State(state): State>, + body: Result>, JsonRejection>, +) -> Result, ApiError> { + // HTTPVLLMWeightSyncClient omits the body when no version is supplied. + let request = body?.map(|Json(request)| request).unwrap_or_default(); + + let client = state.engine_core_client(); + client + .finish_weight_update() + .await + .map_err(|error| utility_call_error("finish_weight_update", error))?; + if let Some(version) = request.weight_version { + client + .set_weight_version(&version) + .await + .map_err(|error| utility_call_error("set_weight_version", error))?; + } + + Ok(Json(MessageResponse { + message: "Weight update finished", + })) +} + +/// Update the weight version identifier. +pub async fn update_weight_version( + State(state): State>, + body: Result, JsonRejection>, +) -> Result, ApiError> { + let Json(body) = body?; + state + .engine_core_client() + .set_weight_version(&body.new_version) + .await + .map_err(|error| utility_call_error("set_weight_version", error))?; + + Ok(Json(UpdateWeightVersionResponse { + success: true, + new_version: body.new_version, + })) +} + +/// Get current weight metadata. +pub async fn weight_info( + State(state): State>, +) -> Result, ApiError> { + let weight_version = state + .engine_core_client() + .get_weight_version() + .await + .map_err(|error| utility_call_error("get_weight_version", error))?; + + Ok(Json(WeightInfoResponse { weight_version })) +} From 39545e475d3627287ff69c25465dc0bd405f67e1 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Mon, 14 Sep 2026 11:40:13 +0800 Subject: [PATCH 152/774] Revert "[CI][XPU] Disable model runner V2 for XPU quantization test for some partially pre-quantized models" (#56394) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/test-intel.yaml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 8f4e83eee8ab..a76331f4d5c4 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -148,16 +148,10 @@ steps: source_file_dependencies: - vllm/ - .buildkite/intel_jobs/test-intel.yaml - # TODO: XPU fused MoE + model runner V2 currently hangs on some - # pre-quantized models, e.g. - # quantization/test_online.py::test_online_quantization_loads_real_weights[nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable] - # Remove the VLLM_USE_V2_MODEL_RUNNER=0 override below once the hang is - # root-caused and fixed. commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - export VLLM_USE_V2_MODEL_RUNNER=0 && pytest -v -s quantization/test_auto_round.py && pytest -v -s quantization/test_online.py' - label: "XPU compressed tensors FP8 test" From a2685f2cdace04138d10719b9bb612a67bf20886 Mon Sep 17 00:00:00 2001 From: Hotragn Pettugani <103170876+Hotragn@users.noreply.github.com> Date: Mon, 14 Sep 2026 00:14:26 -0400 Subject: [PATCH 153/774] [Bugfix][Multimodal] Validate base64 video payloads, matching image and audio (#54323) Signed-off-by: hotragn Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com> --- vllm/multimodal/media/video.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index f1f5783b0e8a..7528331cdfd0 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -190,7 +190,7 @@ def load_base64( } return MediaWithBytes((frames, metadata), data.encode()) - return self.load_bytes(pybase64.b64decode(data)) + return self.load_bytes(pybase64.b64decode(data, validate=True)) def load_file( self, filepath: Path From 214248c7d896a6d9430ffe6870ee78c6f0ffb556 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 13 Sep 2026 21:30:11 -0700 Subject: [PATCH 154/774] [CI] Move smaller H200 workloads to 18GB MIG slices (#56695) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/cuda.yaml | 4 ++-- .buildkite/test_areas/engine.yaml | 20 +++++++++++++++++--- .buildkite/test_areas/misc.yaml | 4 ++-- .buildkite/test_areas/samplers.yaml | 23 +++++++++++++++++++---- 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index f0b1f4af2b80..230781cd5da8 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -30,8 +30,8 @@ steps: - vllm/ - tests/cuda -- label: ":nvidia: (H200 MIG 35GB) CUDAGraph" - device: h200_35gb +- label: ":nvidia: (H200 MIG 18GB) CUDAGraph" + device: h200_18gb key: cudagraph timeout_in_minutes: 30 source_file_dependencies: diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index a6b98bf9313b..bb848cd9e5e4 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -74,18 +74,20 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) E2E Core" - device: h200_35gb +- label: ":nvidia: (H200 MIG 18GB) E2E Core" + device: h200_18gb key: e2e-core-1-gpu timeout_in_minutes: 50 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" mirror: amd: label: ":amd: (MI250) E2E Core" + commands: + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py device: mi250_1 timeout_in_minutes: 50 depends_on: @@ -95,6 +97,18 @@ steps: - tests/v1/e2e/general/ - vllm/platforms/rocm.py +# Keep the large Gemma-3n and Mamba MRV2 warmup allocations on 35GB slices. +- label: ":nvidia: (H200 MIG 35GB) E2E Core Large Memory" + device: h200_35gb + key: e2e-core-large-memory + timeout_in_minutes: 50 + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + commands: + - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py + - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + - label: ":nvidia: (L4) V1 E2E" key: v1-e2e-2-gpus timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index d24ec5b57f46..b45b26b2f221 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -103,8 +103,8 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) V1 Executor + Worker" - device: h200_35gb +- label: ":nvidia: (H200 MIG 18GB) V1 Executor + Worker" + device: h200_18gb key: v1-executor-worker timeout_in_minutes: 45 source_file_dependencies: *v1-core-kv-metrics-deps diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index dae7e9646846..18f857f65a5f 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -2,8 +2,8 @@ group: Samplers depends_on: - image-build steps: -- label: ":nvidia: (H200 MIG 35GB) Samplers" - device: h200_35gb +- label: ":nvidia: (H200 MIG 18GB) Samplers" + device: h200_18gb key: samplers-test timeout_in_minutes: 60 source_file_dependencies: @@ -15,8 +15,8 @@ steps: commands: # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both # values explicitly to still cover the PyTorch-native (Triton) path. - - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers - - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers + - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data" + - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data" mirror: amd: label: ":amd: (MI250) Samplers" @@ -35,3 +35,18 @@ steps: - vllm/platforms/rocm.py commands: - pytest -v -s samplers + +# Qwen2-Audio-7B leaves no room for KV cache on a 16 GiB slice. +- label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search" + device: h200_35gb + key: samplers-multimodal-beam-search + timeout_in_minutes: 30 + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - tests/samplers + - tests/conftest.py + - vllm/entrypoints/generate/beam_search + commands: + - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data + - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data From dc36fcce902a63eab06c1b93a5c4a5ee178a0c56 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 13 Sep 2026 21:30:29 -0700 Subject: [PATCH 155/774] [CI] Collect GPU memory telemetry for MIG slices (#56747) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: khluu Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-otel/README.md | 14 +- .buildkite/scripts/ci-otel/ci_gpu.py | 174 +++++++++++++++- .../scripts/ci-otel/tests/test_ci_otel.py | 195 +++++++++++++++++- 3 files changed, 371 insertions(+), 12 deletions(-) diff --git a/.buildkite/scripts/ci-otel/README.md b/.buildkite/scripts/ci-otel/README.md index 3ad552b5dbf0..318cec5f2169 100644 --- a/.buildkite/scripts/ci-otel/README.md +++ b/.buildkite/scripts/ci-otel/README.md @@ -27,9 +27,17 @@ These are device readings during each test's time interval, not process-level kernel attribution. Other processes and overlapping tests can contribute to the readings. Subsecond tests may have no sample. Missing readings never become zero utilization. MIG parent-device memory and utilization are omitted because -they cannot be attributed to the job's partition. The initial implementation -covers NVIDIA commands in the existing tracing rollout; CPU jobs and AMD -mirrors do not collect GPU samples. +they cannot be attributed to the job's partition. When the environment identifies +an assigned MIG UUID, the collector queries that instance's used/total memory +through vLLM's standalone bundled NVML binding in a bounded subprocess. +The binding is loaded directly to avoid vLLM/Torch import side effects. Numeric +`CUDA_VISIBLE_DEVICES` indices are resolved against explicit MIG UUIDs in +`NVIDIA_VISIBLE_DEVICES`. For CDI assignments without an explicit UUID, a bounded +helper initializes the CUDA driver and reads only CUDA-visible device UUIDs, +without creating CUDA contexts or importing Torch. NVML verifies each MIG handle. +Discovery or permission failures remain unknown; NVML sibling instances are never +enumerated. MIG utilization stays absent. +CPU jobs and AMD mirrors do not collect GPU samples. Run the helper tests without loading vLLM or requiring a GPU: diff --git a/.buildkite/scripts/ci-otel/ci_gpu.py b/.buildkite/scripts/ci-otel/ci_gpu.py index 304d738b886b..0c77f5170b19 100644 --- a/.buildkite/scripts/ci-otel/ci_gpu.py +++ b/.buildkite/scripts/ci-otel/ci_gpu.py @@ -6,7 +6,10 @@ from __future__ import annotations import csv +import ctypes +import importlib.util import io +import json import math import os import secrets @@ -15,6 +18,8 @@ import sys import threading import time +import uuid as uuid_module +from pathlib import Path from ci_otel import Span, export_spans, record_spans @@ -67,6 +72,118 @@ def parse_samples(output: str, timestamp_ns: int) -> list[dict]: return events +def visible_mig_devices() -> list[str]: + """Resolve explicit MIG assignments without enumerating sibling instances.""" + nvidia = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").split(",") + nvidia = [value.strip() for value in nvidia] + cuda = os.environ.get("CUDA_VISIBLE_DEVICES") + devices = nvidia if cuda is None else [value.strip() for value in cuda.split(",")] + if devices and all(value.isdigit() for value in devices): + if not all(value.startswith("MIG-") for value in nvidia): + return [] + try: + devices = [nvidia[int(value)] for value in devices] + except IndexError: + return [] + if not devices or not all(value.startswith("MIG-") for value in devices): + return [] + return list(dict.fromkeys(devices))[:MAX_DEVICES] + + +def load_nvml(): + """Load the standalone bundled binding without importing vLLM or Torch.""" + # Test images install vLLM as a wheel and keep only tests/helpers in the + # checkout. Finding the top-level package does not execute its __init__. + package = importlib.util.find_spec("vllm") + if package is None or package.origin is None: + raise ImportError("Installed vLLM package unavailable") + path = Path(package.origin).parent / "third_party/pynvml.py" + spec = importlib.util.spec_from_file_location("ci_nvml", path) + if spec is None or spec.loader is None: + raise ImportError("Bundled NVML binding unavailable") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def cuda_visible_mig_devices() -> list[str]: + """Discover CDI assignments in a subprocess without creating CUDA contexts.""" + pynvml = load_nvml() + + driver = ctypes.CDLL("libcuda.so.1") + count = ctypes.c_int() + if driver.cuInit(0) or driver.cuDeviceGetCount(ctypes.byref(count)): + return [] + devices = [] + pynvml.nvmlInit() + try: + for index in range(min(count.value, MAX_DEVICES)): + device = ctypes.c_int() + identifier = (ctypes.c_ubyte * 16)() + if driver.cuDeviceGet(ctypes.byref(device), index): + continue + # The v2 entry point returns the compute instance UUID under MIG. + if driver.cuDeviceGetUuid_v2(ctypes.byref(identifier), device): + continue + uuid = "MIG-" + str(uuid_module.UUID(bytes=bytes(identifier))) + try: + handle = pynvml.nvmlDeviceGetHandleByUUID(uuid) + if pynvml.nvmlDeviceIsMigDeviceHandle(handle): + devices.append(uuid) + except pynvml.NVMLError: + continue + finally: + pynvml.nvmlShutdown() + return devices + + +def query_mig_samples(devices: list[str], timestamp_ns: int) -> list[dict]: + # Use the binding already shipped with vLLM; never initialize CUDA/Torch. + pynvml = load_nvml() + + events = [] + pynvml.nvmlInit() + try: + for index, uuid in enumerate(devices[:MAX_DEVICES]): + if not uuid.startswith("MIG-"): + continue + attributes: dict[str, str | int | bool] = { + "gpu.uuid": uuid, + "gpu.index": index, + "gpu.name": "NVIDIA MIG", + "gpu.status": "unavailable_mig_memory", + } + try: + handle = pynvml.nvmlDeviceGetHandleByUUID(uuid) + if not pynvml.nvmlDeviceIsMigDeviceHandle(handle): + continue + memory = pynvml.nvmlDeviceGetMemoryInfo(handle) + used = _number(str(memory.used)) + total = _number(str(memory.total)) + if used is not None and total and used <= total: + attributes.update( + { + "gpu.memory.used": used, + "gpu.memory.total": total, + "gpu.status": "mig_memory_only", + } + ) + # NVML does not support per-MIG utilization on these devices. + except pynvml.NVMLError: + pass + events.append( + { + "time_ns": timestamp_ns, + "name": "ci.gpu.sample", + "attributes": attributes, + } + ) + finally: + pynvml.nvmlShutdown() + return events + + def sample_span(events: list[dict]) -> Span: return Span( trace_id=os.environ["CI_INFRA_TRACE_ID"], @@ -88,16 +205,23 @@ def collect(parent_pid: int, stop: threading.Event) -> None: events: list[dict] = [] batch_started = time.monotonic() failures = 0 + discovery_attempted = False + mig_devices = visible_mig_devices() + query = ( + [sys.executable, __file__, "--query-mig", *mig_devices] + if mig_devices + else [ + "nvidia-smi", + f"--query-gpu={QUERY}", + "--format=csv,noheader,nounits", + ] + ) try: while not stop.is_set() and os.getppid() == parent_pid: started = time.monotonic() try: result = subprocess.run( - [ - "nvidia-smi", - f"--query-gpu={QUERY}", - "--format=csv,noheader,nounits", - ], + query, capture_output=True, text=True, timeout=QUERY_TIMEOUT_SECONDS, @@ -105,12 +229,41 @@ def collect(parent_pid: int, stop: threading.Event) -> None: ) if stop.is_set(): break - samples = parse_samples(result.stdout, time.time_ns()) + samples = ( + json.loads(result.stdout) + if mig_devices + else parse_samples(result.stdout, time.time_ns()) + ) + if ( + not mig_devices + and not discovery_attempted + and any( + sample["attributes"].get("gpu.status") == "unsupported_mig" + for sample in samples + ) + ): + discovery_attempted = True + # CDI can assign a MIG device without exporting its UUID. + # Only CUDA-visible devices are inspected, never NVML siblings. + try: + discovered = subprocess.run( + [sys.executable, __file__, "--discover-mig"], + capture_output=True, + text=True, + timeout=QUERY_TIMEOUT_SECONDS, + check=True, + ) + mig_devices = json.loads(discovered.stdout) + except (OSError, subprocess.SubprocessError, json.JSONDecodeError): + mig_devices = [] + if mig_devices: + query = [sys.executable, __file__, "--query-mig", *mig_devices] + continue # Do not create an extra parent-device series. if not samples: break events.extend(samples) failures = 0 - except (OSError, subprocess.SubprocessError): + except (OSError, subprocess.SubprocessError, json.JSONDecodeError): failures += 1 if failures >= 3: print( @@ -142,4 +295,9 @@ def main() -> None: if __name__ == "__main__": - main() + if len(sys.argv) > 1 and sys.argv[1] == "--discover-mig": + print(json.dumps(cuda_visible_mig_devices())) + elif len(sys.argv) > 1 and sys.argv[1] == "--query-mig": + print(json.dumps(query_mig_samples(sys.argv[2:], time.time_ns()))) + else: + main() diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index f636e7ed1826..ba54f9f27349 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -11,6 +11,7 @@ import threading import time from pathlib import Path +from types import SimpleNamespace import pytest @@ -814,6 +815,85 @@ def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics(): assert "gpu.memory.total" not in events[2]["attributes"] +@pytest.mark.parametrize( + "cuda,nvidia,expected", + [ + ("MIG-a", "GPU-parent", ["MIG-a"]), + ("0", "MIG-a", ["MIG-a"]), + ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]), + (None, "MIG-a", ["MIG-a"]), + ("", "MIG-a", []), + ("-1", "MIG-a", []), + ("2", "MIG-a", []), + ("GPU-parent", "MIG-a", []), + ("0", "all", []), + ], +) +def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected): + monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + if cuda is not None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda) + assert ci_gpu.visible_mig_devices() == expected + + +@pytest.mark.parametrize("used", [0, 3 * 1024**3]) +def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used): + handles = [] + closed = [] + + def get_handle(uuid): + handles.append(uuid) + return "slice-handle" + + def get_memory(handle): + assert handle == "slice-handle" + return SimpleNamespace(used=used, total=16 * 1024**3) + + nvml = SimpleNamespace( + nvmlInit=lambda: None, + nvmlShutdown=lambda: closed.append(True), + nvmlDeviceGetHandleByUUID=get_handle, + nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle", + nvmlDeviceGetMemoryInfo=get_memory, + NVMLError=RuntimeError, + ) + monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) + samples = ci_gpu.query_mig_samples(["MIG-a"], 1234) + assert handles == ["MIG-a"] + assert closed == [True] + assert samples[0]["time_ns"] == 1234 + attributes = samples[0]["attributes"] + assert attributes["gpu.uuid"] == "MIG-a" + assert attributes["gpu.memory.used"] == used + assert attributes["gpu.memory.total"] == 16 * 1024**3 + assert "gpu.utilization" not in attributes + + +@pytest.mark.parametrize("is_mig", [True, False]) +def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig): + def unavailable(handle): + if not is_mig: + pytest.fail("queried parent memory") + raise RuntimeError("Insufficient Permissions") + + nvml = SimpleNamespace( + nvmlInit=lambda: None, + nvmlShutdown=lambda: None, + nvmlDeviceGetHandleByUUID=lambda uuid: "handle", + nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig, + nvmlDeviceGetMemoryInfo=unavailable, + NVMLError=RuntimeError, + ) + monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) + samples = ci_gpu.query_mig_samples(["MIG-a"], 1234) + if is_mig: + assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory" + assert "gpu.memory.used" not in samples[0]["attributes"] + else: + assert samples == [] + + def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip( monkeypatch, tmp_path ): @@ -830,7 +910,11 @@ def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip( assert b"ci.gpu.sample" in ci_otel.encode_request([span]) -def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch): +@pytest.mark.parametrize("mig", [False, True]) +def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples( + monkeypatch, mig +): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0") calls = [] def query(*args, **kwargs): @@ -918,3 +1002,112 @@ def test_gpu_sampler_stops_with_command_and_preserves_failure_status( command["start_ns"] <= event["time_ns"] <= command["end_ns"] for event in batches[0]["events"] ) + + +@pytest.mark.parametrize("is_mig", [True, False]) +def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig): + identifiers = [] + + def get_count(output): + output._obj.value = 1 + return 0 + + def get_device(output, index): + assert index == 0 + output._obj.value = 7 + return 0 + + def get_uuid(output, device): + assert device.value == 7 + output._obj[:] = bytes(range(16)) + return 0 + + driver = SimpleNamespace( + cuInit=lambda flags: 0, + cuDeviceGetCount=get_count, + cuDeviceGet=get_device, + cuDeviceGetUuid_v2=get_uuid, + ) + monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver) + nvml = SimpleNamespace( + nvmlInit=lambda: None, + nvmlShutdown=lambda: None, + nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle", + nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig, + NVMLError=RuntimeError, + ) + monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) + expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f" + assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else []) + assert identifiers == [expected] + + +@pytest.mark.parametrize("discovery_fails", [False, True]) +def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails): + monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) + monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void") + monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) + stop = threading.Event() + commands = [] + records = [] + + def query(command, **kwargs): + assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS + commands.append(command) + if "--discover-mig" in command: + if discovery_fails: + raise subprocess.TimeoutExpired(command, 2) + return SimpleNamespace(stdout='["MIG-assigned"]') + if len(commands) == 4: + stop.set() + if "--query-mig" in command: + assert command[-1] == "MIG-assigned" + return SimpleNamespace( + stdout=json.dumps( + [ + { + "time_ns": 2, + "name": "ci.gpu.sample", + "attributes": { + "gpu.uuid": "MIG-assigned", + "gpu.memory.used": 123, + }, + } + ] + ) + ) + return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled") + + monkeypatch.setattr(ci_gpu.subprocess, "run", query) + monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans)) + ci_gpu.collect(os.getppid(), stop) + assert sum("--discover-mig" in command for command in commands) == 1 + attributes = [event["attributes"] for span in records for event in span.events] + if discovery_fails: + assert all("gpu.memory.used" not in event for event in attributes) + else: + assert any(event.get("gpu.memory.used") == 123 for event in attributes) + assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes) + + +def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path): + # Match test images: the collector is outside the installed vLLM package. + helper = tmp_path / "ci_gpu_relocated.py" + helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text()) + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; import ci_gpu_relocated as ci_gpu; " + "binding = ci_gpu.load_nvml(); " + "assert callable(binding.nvmlDeviceGetMemoryInfo); " + "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules", + ], + env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"}, + capture_output=True, + text=True, + timeout=5, + ) + assert result.returncode == 0, result.stderr From 67111973eeb104e83bc6707a2d0f139737b711a0 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 14 Sep 2026 13:03:08 +0800 Subject: [PATCH 156/774] [Frontend] Move grpc_server to launchers. (#56746) Signed-off-by: wang.yuqi --- .../test_grpc_health.py | 0 vllm/entrypoints/cli/serve.py | 2 +- vllm/entrypoints/grpc_server.py | 207 ++---------------- vllm/entrypoints/launchers/grpc_server.py | 200 +++++++++++++++++ 4 files changed, 218 insertions(+), 191 deletions(-) rename tests/entrypoints/{unit_tests => launchers}/test_grpc_health.py (100%) create mode 100644 vllm/entrypoints/launchers/grpc_server.py diff --git a/tests/entrypoints/unit_tests/test_grpc_health.py b/tests/entrypoints/launchers/test_grpc_health.py similarity index 100% rename from tests/entrypoints/unit_tests/test_grpc_health.py rename to tests/entrypoints/launchers/test_grpc_health.py diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index a284b4bf39f5..e43b273e5e96 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -54,7 +54,7 @@ def cmd(args: argparse.Namespace) -> None: args.model = args.model_tag if getattr(args, "grpc", False): - from vllm.entrypoints.grpc_server import serve_grpc + from vllm.entrypoints.launchers.grpc_server import serve_grpc uvloop.run(serve_grpc(args)) return diff --git a/vllm/entrypoints/grpc_server.py b/vllm/entrypoints/grpc_server.py index 59269dd18020..913d084e4a81 100644 --- a/vllm/entrypoints/grpc_server.py +++ b/vllm/entrypoints/grpc_server.py @@ -1,200 +1,27 @@ -#!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# mypy: ignore-errors -""" -vLLM gRPC Server +import warnings -Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer -from the smg-grpc-servicer package. +from vllm.entrypoints.launchers.grpc_server import main, serve_grpc -Usage: - python -m vllm.entrypoints.grpc_server --model +warnings.warn( + "`vllm.entrypoints.grpc_server is deprecated and will likely be " + "unsupported in a future version. Use the corresponding function from " + "`vllm.entrypoints.launchers.grpc_server` instead.", + DeprecationWarning, + stacklevel=1, +) -Example: - python -m vllm.entrypoints.grpc_server \ - --model meta-llama/Llama-2-7b-hf \ - --host 0.0.0.0 \ - --port 50051 -""" +__all__ = ["main", "serve_grpc"] -import argparse -import asyncio -import signal -import sys -import time -try: - import grpc - from grpc_health.v1 import health_pb2_grpc - from grpc_reflection.v1alpha import reflection - from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc - from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer - from smg_grpc_servicer.vllm.servicer import VllmEngineServicer -except ImportError as e: - raise ImportError( - "gRPC mode requires smg-grpc-servicer. " - "If not installed, run: pip install vllm[grpc]. " - "If already installed, there may be a broken import due to a " - "version mismatch — see the chained exception above for details." - ) from e - -import uvloop - -from vllm import envs -from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.entrypoints.serve.utils.api_utils import log_version_and_model -from vllm.logger import init_logger -from vllm.usage.usage_lib import UsageContext -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.v1.engine.async_llm import AsyncLLM -from vllm.version import __version__ as VLLM_VERSION - -logger = init_logger(__name__) - - -async def serve_grpc(args: argparse.Namespace): - """ - Main gRPC serving function. - - Args: - args: Parsed command line arguments - """ - log_version_and_model(logger, VLLM_VERSION, args.model) - logger.info("vLLM gRPC server args: %s", args) - - start_time = time.time() - - # Create engine args - engine_args = AsyncEngineArgs.from_cli_args(args) - - # Build vLLM config - vllm_config = engine_args.create_engine_config( - usage_context=UsageContext.OPENAI_API_SERVER, - ) - - # Create AsyncLLM - async_llm = AsyncLLM.from_vllm_config( - vllm_config=vllm_config, - usage_context=UsageContext.OPENAI_API_SERVER, - enable_log_requests=args.enable_log_requests, - disable_log_stats=args.disable_log_stats, - ) - - # Create servicer - servicer = VllmEngineServicer(async_llm, start_time) - - # Create gRPC server - server = grpc.aio.server( - options=[ - ("grpc.max_send_message_length", -1), - ("grpc.max_receive_message_length", -1), - # Tolerate client keepalive pings every 10s (default 300s is too - # strict for non-streaming requests where no DATA frames flow - # during generation) - ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000), - ("grpc.keepalive_permit_without_calls", True), - ], - ) - - # Add servicer to server - vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server) - - # Add standard gRPC health service for Kubernetes probes - health_servicer = VllmHealthServicer(async_llm) - health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) - - # Enable reflection for grpcurl and other tools - service_names = ( - vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name, - "grpc.health.v1.Health", - reflection.SERVICE_NAME, - ) - reflection.enable_server_reflection(service_names, server) - - # Bind to address - host = args.host or "0.0.0.0" - address = f"{host}:{args.port}" - server.add_insecure_port(address) - - try: - # Start server - await server.start() - logger.info("vLLM gRPC server started on %s", address) - logger.info("Server is ready to accept requests") - - # Start periodic stats logging (mirrors the HTTP server's lifespan task) - if not args.disable_log_stats: - - async def _force_log(): - while True: - await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL) - await async_llm.do_log_stats() - - stats_task = asyncio.create_task(_force_log()) - else: - stats_task = None - - # Handle shutdown signals - loop = asyncio.get_running_loop() - stop_event = asyncio.Event() - - def signal_handler(): - logger.info("Received shutdown signal") - stop_event.set() - - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, signal_handler) - - try: - await stop_event.wait() - except KeyboardInterrupt: - logger.info("Interrupted by user") - finally: - logger.info("Shutting down vLLM gRPC server...") - if stats_task is not None: - stats_task.cancel() - try: - health_servicer.set_not_serving() - except Exception: # broad: must not prevent server.stop() / shutdown() - logger.warning("Failed to set health status to NOT_SERVING", exc_info=True) - await server.stop(grace=5.0) - logger.info("gRPC server stopped") - async_llm.shutdown() - logger.info("AsyncLLM engine stopped") - logger.info("Shutdown complete") - - -def main(): - """Main entry point for python -m vllm.entrypoints.grpc_server.""" - parser = FlexibleArgumentParser( - description="vLLM gRPC Server", - ) - - # Server args - parser.add_argument( - "--host", - type=str, - default="0.0.0.0", - help="Host to bind gRPC server to", - ) - parser.add_argument( - "--port", - type=int, - default=50051, - help="Port to bind gRPC server to", +if __name__ == "__main__": + warnings.warn( + "`python -m vllm.entrypoints.grpc_server is deprecated and " + "will likely be unsupported in a future version. " + "Please use `vllm serve --grpc` instead.", + DeprecationWarning, + stacklevel=1, ) - parser = AsyncEngineArgs.add_cli_args(parser) - args = parser.parse_args() - - # Run server - try: - uvloop.run(serve_grpc(args)) - except Exception as e: - logger.exception("Server failed: %s", e) - sys.exit(1) - - -if __name__ == "__main__": main() diff --git a/vllm/entrypoints/launchers/grpc_server.py b/vllm/entrypoints/launchers/grpc_server.py new file mode 100644 index 000000000000..3dc51f2e771d --- /dev/null +++ b/vllm/entrypoints/launchers/grpc_server.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# mypy: ignore-errors +""" +vLLM gRPC Server + +Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer +from the smg-grpc-servicer package. + +Usage: + python -m vllm.entrypoints.launchers.grpc_server --model + +Example: + python -m vllm.entrypoints.launchers.grpc_server \ + --model meta-llama/Llama-2-7b-hf \ + --host 0.0.0.0 \ + --port 50051 +""" + +import argparse +import asyncio +import signal +import sys +import time + +try: + import grpc + from grpc_health.v1 import health_pb2_grpc + from grpc_reflection.v1alpha import reflection + from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc + from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer + from smg_grpc_servicer.vllm.servicer import VllmEngineServicer +except ImportError as e: + raise ImportError( + "gRPC mode requires smg-grpc-servicer. " + "If not installed, run: pip install vllm[grpc]. " + "If already installed, there may be a broken import due to a " + "version mismatch — see the chained exception above for details." + ) from e + +import uvloop + +from vllm import envs +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.entrypoints.serve.utils.api_utils import log_version_and_model +from vllm.logger import init_logger +from vllm.usage.usage_lib import UsageContext +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.version import __version__ as VLLM_VERSION + +logger = init_logger(__name__) + + +async def serve_grpc(args: argparse.Namespace): + """ + Main gRPC serving function. + + Args: + args: Parsed command line arguments + """ + log_version_and_model(logger, VLLM_VERSION, args.model) + logger.info("vLLM gRPC server args: %s", args) + + start_time = time.time() + + # Create engine args + engine_args = AsyncEngineArgs.from_cli_args(args) + + # Build vLLM config + vllm_config = engine_args.create_engine_config( + usage_context=UsageContext.OPENAI_API_SERVER, + ) + + # Create AsyncLLM + async_llm = AsyncLLM.from_vllm_config( + vllm_config=vllm_config, + usage_context=UsageContext.OPENAI_API_SERVER, + enable_log_requests=args.enable_log_requests, + disable_log_stats=args.disable_log_stats, + ) + + # Create servicer + servicer = VllmEngineServicer(async_llm, start_time) + + # Create gRPC server + server = grpc.aio.server( + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + # Tolerate client keepalive pings every 10s (default 300s is too + # strict for non-streaming requests where no DATA frames flow + # during generation) + ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000), + ("grpc.keepalive_permit_without_calls", True), + ], + ) + + # Add servicer to server + vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server) + + # Add standard gRPC health service for Kubernetes probes + health_servicer = VllmHealthServicer(async_llm) + health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) + + # Enable reflection for grpcurl and other tools + service_names = ( + vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name, + "grpc.health.v1.Health", + reflection.SERVICE_NAME, + ) + reflection.enable_server_reflection(service_names, server) + + # Bind to address + host = args.host or "0.0.0.0" + address = f"{host}:{args.port}" + server.add_insecure_port(address) + + try: + # Start server + await server.start() + logger.info("vLLM gRPC server started on %s", address) + logger.info("Server is ready to accept requests") + + # Start periodic stats logging (mirrors the HTTP server's lifespan task) + if not args.disable_log_stats: + + async def _force_log(): + while True: + await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL) + await async_llm.do_log_stats() + + stats_task = asyncio.create_task(_force_log()) + else: + stats_task = None + + # Handle shutdown signals + loop = asyncio.get_running_loop() + stop_event = asyncio.Event() + + def signal_handler(): + logger.info("Received shutdown signal") + stop_event.set() + + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, signal_handler) + + try: + await stop_event.wait() + except KeyboardInterrupt: + logger.info("Interrupted by user") + finally: + logger.info("Shutting down vLLM gRPC server...") + if stats_task is not None: + stats_task.cancel() + try: + health_servicer.set_not_serving() + except Exception: # broad: must not prevent server.stop() / shutdown() + logger.warning("Failed to set health status to NOT_SERVING", exc_info=True) + await server.stop(grace=5.0) + logger.info("gRPC server stopped") + async_llm.shutdown() + logger.info("AsyncLLM engine stopped") + logger.info("Shutdown complete") + + +def main(): + """Main entry point for python -m vllm.entrypoints.launchers.grpc_server.""" + parser = FlexibleArgumentParser( + description="vLLM gRPC Server", + ) + + # Server args + parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Host to bind gRPC server to", + ) + parser.add_argument( + "--port", + type=int, + default=50051, + help="Port to bind gRPC server to", + ) + parser = AsyncEngineArgs.add_cli_args(parser) + + args = parser.parse_args() + + # Run server + try: + uvloop.run(serve_grpc(args)) + except Exception as e: + logger.exception("Server failed: %s", e) + sys.exit(1) + + +if __name__ == "__main__": + main() From 73d2a8cf86c46e21af723ffd34ce0d56668fdd14 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:42:33 +0800 Subject: [PATCH 157/774] [2/2][Model Runner V2] FULL CUDA graph capture for microbatched steps (DBO) (#51700) Signed-off-by: specture724 Co-authored-by: Codex --- tests/v1/worker/test_gpu_ubatch_slicing.py | 163 ++++++++++++++++++++- vllm/config/vllm.py | 5 +- vllm/v1/worker/gpu/cudagraph_utils.py | 146 ++++++++++++++---- vllm/v1/worker/gpu/dp_utils.py | 46 ++++-- vllm/v1/worker/gpu/model_runner.py | 7 +- vllm/v1/worker/gpu/ubatch_utils.py | 120 +++++++++++---- 6 files changed, 407 insertions(+), 80 deletions(-) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index f54e74671071..e735953d8b15 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -20,13 +20,22 @@ import torch from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata -from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig +from vllm.config import ( + CompilationConfig, + CUDAGraphMode, + ModelConfig, + ParallelConfig, + VllmConfig, +) from vllm.forward_context import create_forward_context from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils -from vllm.v1.worker.gpu import dp_utils -from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu import cudagraph_utils, dp_utils +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.ubatch_utils import ( @@ -272,6 +281,7 @@ def _sync_dp( num_tokens_per_rank: list[int], uniform_token_count_per_rank: list[int] | None = None, allow_ubatching: bool = True, + cudagraph_manager: Any = None, ) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]: """Run the DP handshake with the all-reduce stubbed out. @@ -295,7 +305,7 @@ def _sync_dp( patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)), ): return dp_utils.sync_cudagraph_and_dp_padding( - cudagraph_manager=None, + cudagraph_manager=cudagraph_manager, desired_batch_desc=BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens_per_rank[0], @@ -718,3 +728,148 @@ def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off(): ] assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") +def test_capturable_run_replays_as_a_cudagraph(): + """Microbatched graph replay reads updated persistent input buffers.""" + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, all2all_backend="deepep_low_latency" + ), + ) + device = torch.device("cuda:0") + runner = _make_execution_runner(vllm_config) + + num_tokens = 16 + input_ids = torch.arange(num_tokens, device=device) + positions = torch.arange(num_tokens, device=device) + model_inputs = { + "input_ids": input_ids, + "positions": positions, + "inputs_embeds": None, + "intermediate_tensors": None, + } + ubatch_state = _make_ubatch_state( + vllm_config, + [ + UBatchSlice(slice(0, 4), slice(0, 8)), + UBatchSlice(slice(4, 8), slice(8, 16)), + ], + ) + + model = _YieldingModel([]) + graph = torch.cuda.CUDAGraph() + finish = runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + with torch.cuda.graph(graph, stream=runner.capture_stream): + captured = finish() + + input_ids.fill_(3) + positions.fill_(5) + graph.replay() + torch.accelerator.synchronize() + + expected = torch.full( + (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device + ) + torch.testing.assert_close(captured, expected) + + +def _request_slices(input_batch: InputBatch) -> list[slice]: + return [s.request_slice for s in create_ubatch_slices(input_batch, 2)] + + +def test_captured_split_survives_a_replay_with_enough_requests(): + """Replay preserves the captured split only if real requests reach it.""" + buffers = _make_buffers() + captured = _request_slices(InputBatch.make_dummy(16, 16, buffers)) + + def replay(num_reqs: int) -> list[slice]: + return _request_slices( + _make_input_batch( + [1] * num_reqs, + [128] * num_reqs, + buffers, + num_reqs_padded=16, + num_tokens_padded=16, + ) + ) + + assert replay(16) == captured + assert replay(9) == captured + assert replay(4) != captured + + +def _make_cudagraph_manager(capture_sizes: list[int]) -> CudaGraphManager: + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, + all2all_backend="deepep_low_latency", + dbo_decode_token_threshold=DECODE_THRESHOLD, + dbo_prefill_token_threshold=PREFILL_THRESHOLD, + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=capture_sizes, + ), + ) + with patch.object( + cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ): + manager = CudaGraphManager( + vllm_config, + torch.device("cuda:0"), + CUDAGraphMode.FULL, + decode_query_len=1, + ubatch_runner=cast(UBatchRunner, object()), + ) + manager._graphs_captured = True + return manager + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graphs_are_only_offered_to_uniform_batches(): + """Mixed batches cannot reuse graphs captured with uniform query lengths.""" + manager = _make_cudagraph_manager([64, 128]) + + desc = manager.dispatch( + num_reqs=64, + num_tokens=64, + uniform_token_count=1, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + desc = manager.dispatch( + num_reqs=8, + num_tokens=64, + uniform_token_count=None, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.NONE + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graph_needs_every_rank_to_reach_the_split(): + """Fall back to eager on all ranks if any rank cannot reach the captured split.""" + manager = _make_cudagraph_manager([64, 128]) + uniform = [1, 1] + + desc, _ = _sync_dp([64, 64], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33. + desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.NONE + assert desc.num_ubatches == 2 + assert dp_sync is not None and dp_sync.eager diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index fbe021385524..2fcc144af0c3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3034,8 +3034,9 @@ def _get_dbo_unsupported_features(self) -> list[str]: unsupported.append("dual batch overlap with multimodal models") if model_config is not None and model_config.is_hybrid: unsupported.append("dual batch overlap with hybrid models") - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - unsupported.append("dual batch overlap with CUDA graphs") + if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE: + # DBO captures FULL graphs only. + unsupported.append("dual batch overlap with PIECEWISE CUDA graphs") if self.is_mm_encoder_only: unsupported.append("dual batch overlap with encoder only models") diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 12354f3c229e..d197dbbc6df1 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -4,7 +4,8 @@ import itertools from collections import defaultdict from collections.abc import Callable, Iterable -from dataclasses import dataclass +from contextlib import ExitStack +from dataclasses import dataclass, replace from itertools import groupby, product from typing import TYPE_CHECKING, Any, NamedTuple, Protocol @@ -42,11 +43,13 @@ from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.pcp_manager import PCPManager + from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner logger = init_logger(__name__) @@ -144,6 +147,7 @@ def __init__( decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): self.vllm_config = vllm_config self.device = device @@ -153,6 +157,8 @@ def __init__( self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len self.varlen_decode = varlen_decode + # DBO supports FULL CUDA graphs only. + self.ubatch_runner = ubatch_runner self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -212,6 +218,33 @@ def _resolve_effective_loras(self, num_active_loras: int) -> int: # Counts above the largest captured case clamp to it. return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case) + def _maybe_ubatch_twin( + self, desc: BatchExecutionDescriptor + ) -> BatchExecutionDescriptor | None: + """Return a microbatched capture candidate when eligible. + + Uniform query lengths preserve the captured request split. Use the DP + dispatch thresholds so all ranks generate the same candidates. + """ + if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL: + return None + if desc.num_reqs is None: + return None + uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs) + if remainder or desc.uniform_token_count not in (None, uniform_token_count): + return None + parallel_config = self.vllm_config.parallel_config + num_ubatches = get_num_ubatches(parallel_config) + if desc.num_tokens < num_ubatches: + return None + if not check_ubatch_thresholds( + parallel_config, desc.num_tokens, uniform_decode=True + ): + return None + return replace( + desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count + ) + def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" capture_sizes = self.compilation_config.cudagraph_capture_sizes @@ -300,6 +333,12 @@ def _init_candidates(self) -> None: if desc not in descs_by_mode[decode_mode]: descs_by_mode[decode_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None and ( + ubatch_desc not in descs_by_mode[decode_mode] + ): + descs_by_mode[decode_mode].append(ubatch_desc) + # recoverSSM cannot capture a dummy query wider than its workspace. if mixed_mode and ( not self.vllm_config.cache_config.use_kda_recoverssm @@ -321,6 +360,10 @@ def _init_candidates(self) -> None: ) descs_by_mode[mixed_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None: + descs_by_mode[mixed_mode].append(ubatch_desc) + for mode, descs in descs_by_mode.items(): descs.sort(key=lambda d: d.num_tokens, reverse=True) self._capture_descs[mode] = descs @@ -344,6 +387,13 @@ def _init_candidates(self) -> None: def needs_capture(self) -> bool: return len(self._capture_descs) > 0 + def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream: + """Capture on the stream used by the microbatch threads.""" + if desc.num_ubatches > 1: + assert self.ubatch_runner is not None + return self.ubatch_runner.capture_stream + return current_stream() + @torch.inference_mode() def capture( self, @@ -359,7 +409,10 @@ def capture( because attention backends may mutate or lazily initialize metadata during warmup. """ - with graph_capture(device=self.device): + with graph_capture(device=self.device), ExitStack() as stack: + if self.ubatch_runner is not None: + # Join parked threads on failure to avoid blocking later captures. + stack.callback(self.ubatch_runner.abort_pending_run) # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger # activations so FULL activations should fit in already allocated # buffers in the graph pool. @@ -387,7 +440,9 @@ def capture( # Capture logger.debug( - "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + "CG Capture: mode=%s, batch_desc=%s", + desc.cg_mode.name, + desc, ) if ( desc.cg_mode == CUDAGraphMode.PIECEWISE @@ -415,7 +470,7 @@ def capture( torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] with torch.cuda.graph( - graph, self.pool, stream=current_stream() + graph, self.pool, stream=self._capture_stream(desc) ): forward_fn(CUDAGraphMode.NONE) # Join offloader's copy stream after forward to avoid @@ -429,6 +484,7 @@ def capture( self._capture_mem_samples.append(free_before - free_after) self.graphs[desc] = graph compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True def captured_token_counts(self) -> list[int]: @@ -510,6 +566,7 @@ def __init__( decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): super().__init__( vllm_config, @@ -518,6 +575,7 @@ def __init__( decode_query_len, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, + ubatch_runner=ubatch_runner, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] @@ -554,6 +612,35 @@ def capture( "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL." ) + def store_capture_output(num_tokens: int, model_output: Any) -> None: + """Copy outputs to persistent buffers, allocating on first use.""" + if self.is_last_pp_rank: + # Last PP rank (common case). + if self.use_aux_hidden_state_outputs: + hidden_states, aux_hidden_states = model_output + else: + hidden_states = model_output + aux_hidden_states = [] + if self.hidden_states is None: + self.hidden_states = torch.empty_like(hidden_states) + self.hidden_states[:num_tokens] = hidden_states + if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: + self.aux_hidden_states = [ + torch.empty_like(x) for x in aux_hidden_states + ] + for i, aux in enumerate(aux_hidden_states): + self.aux_hidden_states[i][:num_tokens] = aux + else: + # Non-last PP rank. + assert isinstance(model_output, IntermediateTensors) + intermediate_tensors = model_output + if self.intermediate_tensors is None: + self.intermediate_tensors = IntermediateTensors.empty_like( + intermediate_tensors + ) + for k, v in intermediate_tensors.tensors.items(): + self.intermediate_tensors[k][:num_tokens] = v + def create_forward_fn( desc: BatchExecutionDescriptor, warmup: bool, @@ -583,6 +670,30 @@ def create_forward_fn( assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] + if desc.num_ubatches > 1: + # Prepare and park threads before capture; finish runs inside it. + assert self.ubatch_runner is not None + ubatch_state = self.ubatch_runner.prepare( + InputBatch.make_dummy(num_reqs, num_tokens, input_buffers), + block_tables.get_dummy_block_tables(num_reqs), + block_tables.get_dummy_slot_mappings(num_tokens), + cg_mode=CUDAGraphMode.FULL, + for_capture=True, + ) + # Capture with dummy rows marked as padding. + input_buffers.is_padding.fill_(True) + finish = self.ubatch_runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + + def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: + assert cg_mode != CUDAGraphMode.PIECEWISE, ( + "DBO does not support PIECEWISE cudagraphs" + ) + store_capture_output(num_tokens, finish()) + + return ubatch_forward_fn + attn_metadata, slot_mappings = prepare_inputs_to_capture( num_reqs, num_tokens, @@ -629,32 +740,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: # model outputs. No need to keep track of the hidden states. return None - if self.is_last_pp_rank: - # Last PP rank (common case). - if self.use_aux_hidden_state_outputs: - hidden_states, aux_hidden_states = model_output - else: - hidden_states = model_output - aux_hidden_states = [] - if self.hidden_states is None: - self.hidden_states = torch.empty_like(hidden_states) - self.hidden_states[:num_tokens] = hidden_states - if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: - self.aux_hidden_states = [ - torch.empty_like(x) for x in aux_hidden_states - ] - for i, aux in enumerate(aux_hidden_states): - self.aux_hidden_states[i][:num_tokens] = aux - else: - # Non-last PP rank. - assert isinstance(model_output, IntermediateTensors) - intermediate_tensors = model_output - if self.intermediate_tensors is None: - self.intermediate_tensors = IntermediateTensors.empty_like( - intermediate_tensors - ) - for k, v in intermediate_tensors.tensors.items(): - self.intermediate_tensors[k][:num_tokens] = v + store_capture_output(num_tokens, model_output) return forward_fn diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index b7c73789de6f..80309bc009c9 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -112,23 +112,45 @@ def sync_cudagraph_and_dp_padding( int(num_tokens_across_dp.min()), uniform_decode=uniform_decode_across_dp, ): - # Microbatching is all-or-nothing: the expert all-to-all is - # collective, so every rank splits and pads to the same token count. - # A rank too small to fill a microbatch does no work in it, like a - # dummy run. Microbatched steps run eager; nothing is captured yet. + # Expert all-to-all requires every rank to split and pad equally. + # Empty microbatches run as dummy batches. ubatch_num_tokens = int(num_tokens_across_dp.max()) - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=ubatch_num_tokens, - num_reqs=num_reqs, - num_ubatches=get_num_ubatches(parallel_config), - ), DPSyncState( + num_ubatches = get_num_ubatches(parallel_config) + ubatch_desc = None + if cudagraph_manager is not None: + # Match a FULL microbatched CUDA graph, falling back to an eager + # descriptor if no graph matches. + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + ubatch_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens: + # If one rank has an empty second microbatch, run without + # CUDA graphs. + ubatch_desc = None + if ubatch_desc is None: + ubatch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=ubatch_num_tokens, + num_reqs=num_reqs, + num_ubatches=num_ubatches, + ) + # Refresh the token count to include CUDA graph padding. + ubatch_num_tokens = ubatch_desc.num_tokens + num_reqs = int(num_reqs_across_dp.max()) + if ubatch_desc.cg_mode == CUDAGraphMode.FULL: + assert ubatch_desc.num_reqs is not None + num_reqs = ubatch_desc.num_reqs + return ubatch_desc, DPSyncState( num_tokens_across_dp=torch.full_like( num_tokens_across_dp, ubatch_num_tokens ), uniform_token_count=synced_uniform_token_count, - eager=True, - num_reqs=int(num_reqs_across_dp.max()), + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, + num_reqs=num_reqs, ) synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9d7918126f62..d66fc6c90f7c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -723,6 +723,7 @@ def initialize_kv_cache( decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, + ubatch_runner=self.ubatch_runner, ) if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: self.fast_prefill = FastPrefillHelper( @@ -1775,7 +1776,11 @@ def execute_model( assert self.ubatch_runner is not None assert block_tables is not None and slot_mappings is not None ubatch_state = self.ubatch_runner.prepare( - input_batch, block_tables, slot_mappings + input_batch, + block_tables, + slot_mappings, + cg_mode=batch_desc.cg_mode, + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) elif not (dummy_run and skip_attn_for_dummy_run): assert slot_mappings is not None diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 8007e2dbb44c..e3cade2d0b17 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -3,6 +3,8 @@ """Microbatching (DBO) helpers for the V2 GPU model runner.""" import threading +from collections.abc import Callable +from contextlib import ExitStack from dataclasses import replace from typing import Any, NamedTuple @@ -317,8 +319,11 @@ def __init__( for _ in range(self.num_ubatches) ] self.comm_stream = torch.cuda.Stream(device=device) + # Threads and graph capture must use the same compute stream. + self.capture_stream = torch.cuda.Stream(device=device) # The microbatch threads plus the thread that starts them. self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + self._pending_finish: Callable[[], Any] | None = None self.sm_control = create_sm_control_context(self.parallel_config) def prepare( @@ -326,11 +331,16 @@ def prepare( input_batch: InputBatch, block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, ) -> UBatchState: """Split the batch into the microbatches the step will run on. Attention metadata is built per microbatch and carried in the forward contexts the threads install, not in the caller's. + + FULL uses padded sizes; `for_capture` refreshes metadata from dummy + block tables. """ ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) @@ -355,12 +365,13 @@ def prepare( attn_metadata.append( self.model_state.prepare_attn( ubatch, - CUDAGraphMode.NONE, + cg_mode, ubatch_block_tables, ubatch_slot_mappings, self.attn_groups, self.kv_cache_config, ubatch_idx=i, + for_capture=for_capture, ) ) slot_mappings_by_layer.append( @@ -415,14 +426,34 @@ def run( model_inputs: dict[str, Any], ubatch_state: UBatchState, ) -> Any: + finish = self.begin_capturable_run(model, model_inputs, ubatch_state) + return finish() + + def begin_capturable_run( + self, + model: Any, + model_inputs: dict[str, Any], + ubatch_state: UBatchState, + for_capture: bool = False, + ) -> Callable[[], Any]: + """Start microbatch threads and return a callback to run and join them. + + For capture, start threads outside the graph and call the callback inside. + `for_capture` selects the capture stream and initializes cuBLAS workspaces + before capture. + + Call the callback exactly once on this thread, or call `abort_pending_run` + on failure to release parked threads. + """ ubatch_slices = ubatch_state.slices assert len(ubatch_slices) == len(ubatch_state.forward_contexts) assert len(ubatch_slices) == self.num_ubatches + compute_stream = self.capture_stream if for_capture else current_stream() ubatch_contexts = make_ubatch_contexts( num_micro_batches=self.num_ubatches, comm_stream=self.comm_stream, - compute_stream=current_stream(), + compute_stream=compute_stream, forward_contexts=ubatch_state.forward_contexts, ready_barrier=self.ready_barrier, ) @@ -433,6 +464,13 @@ def run( @torch.inference_mode() def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None: try: + # New threads must select the worker's device. + torch.accelerator.set_device_index(self.device.index) + if for_capture: + # Allocate cuBLAS workspaces before capture. + for stream in (compute_stream, self.comm_stream): + with torch.cuda.stream(stream): + torch.cuda.current_blas_handle() with ubatch_context: outputs[ubatch_context.id] = model(**inputs) except BaseException as e: # noqa: BLE001 @@ -445,33 +483,56 @@ def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None: # is deliberately out of scope here. errors[ubatch_context.id] = e - # The threads manage the forward context themselves; clear it here so - # it is restored correctly once they are done. - with override_forward_context(None), self.sm_control: - threads = [] - for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): - thread = threading.Thread( - target=run_ubatch, - args=( - ubatch_context, - slice_model_inputs(model_inputs, ubatch_slice.token_slice), - ), - ) - threads.append(thread) - thread.start() + # Threads own the forward context; keep it cleared and SMs reserved + # until finish joins them. + stack = ExitStack() + stack.enter_context(override_forward_context(None)) + stack.enter_context(self.sm_control) + + threads = [] + for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): + thread = threading.Thread( + target=run_ubatch, + args=( + ubatch_context, + slice_model_inputs(model_inputs, ubatch_slice.token_slice), + ), + ) + threads.append(thread) + thread.start() - # Wait for every thread to reach its context, then start the first. - self.ready_barrier.wait() - ubatch_contexts[0].cpu_wait_event.set() - for thread in threads: - thread.join() + # Park all threads before the caller starts capture. + self.ready_barrier.wait() - if errors: - failed = min(errors) - raise RuntimeError( - f"Microbatch {failed} of {self.num_ubatches} failed" - ) from errors[failed] - return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + def finish() -> Any: + try: + ubatch_contexts[0].cpu_wait_event.set() + for thread in threads: + thread.join() + finally: + self._pending_finish = None + stack.close() + + if errors: + failed = min(errors) + raise RuntimeError( + f"Microbatch {failed} of {self.num_ubatches} failed" + ) from errors[failed] + return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + + self._pending_finish = finish + return finish + + def abort_pending_run(self) -> None: + """Release and join parked threads after a failed capture.""" + if self._pending_finish is None: + return + try: + self._pending_finish() + except Exception: + logger.exception("Failed to join microbatch threads after an abort") + finally: + self._pending_finish = None def maybe_build_ubatch_runner( @@ -491,10 +552,7 @@ def maybe_build_ubatch_runner( if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1: return None - logger.info_once( - "Dual batch overlap is enabled. Microbatched steps run without " - "CUDA graphs on the V2 model runner." - ) + logger.info_once("Dual batch overlap is enabled.") return UBatchRunner( vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs ) From 10e8d5614c50ad180acbe61878c93d98ca7098a9 Mon Sep 17 00:00:00 2001 From: Alexander Lee Date: Mon, 14 Sep 2026 01:49:44 -0400 Subject: [PATCH 158/774] [Bugfix][Beam Search] Respect skip_special_tokens during decoding (#56211) Signed-off-by: lxy-alexander Co-authored-by: OpenAI Codex --- tests/samplers/test_beam_search_online.py | 38 ++++++++++++++++++- .../generate/beam_search/offline.py | 4 +- .../generate/beam_search/online.py | 4 +- .../openai/chat_completion/protocol.py | 1 + .../entrypoints/openai/completion/protocol.py | 1 + vllm/sampling_params.py | 1 + 6 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 14481f79c68a..102e66dcbe79 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -11,8 +11,13 @@ class _Tokenizer: eos_token_id = 0 + special_token_id = eos_token_id - def decode(self, token_ids: list[int]) -> str: + def decode(self, token_ids: list[int], skip_special_tokens: bool = False) -> str: + if skip_special_tokens: + token_ids = [ + token_id for token_id in token_ids if token_id != self.special_token_id + ] return " ".join(str(token_id) for token_id in token_ids) @@ -72,3 +77,34 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None: assert outputs[0].outputs[0].finish_reason == "stop" assert outputs[0].outputs[0].token_ids == [] assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("skip_special_tokens", "expected_text"), + [ + pytest.param(True, "", id="skip"), + pytest.param(False, "0", id="keep"), + ], +) +async def test_beam_search_respects_skip_special_tokens( + skip_special_tokens: bool, expected_text: str +) -> None: + prompt = { + "type": "token", + "prompt": "prompt", + "prompt_token_ids": [1], + } + params = BeamSearchParams( + beam_width=1, + max_tokens=1, + ignore_eos=True, + skip_special_tokens=skip_special_tokens, + ) + + outputs = [ + output async for output in _Serving().beam_search(prompt, "request", params) + ] + + assert outputs[0].outputs[0].text == expected_text + assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id] diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index f412f5aebed8..b919435f65a1 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -186,7 +186,9 @@ def beam_search( best_beams = sorted_completed[:beam_width] for beam in best_beams: - beam.text = tokenizer.decode(beam.tokens) + beam.text = tokenizer.decode( + beam.tokens, skip_special_tokens=params.skip_special_tokens + ) outputs.append(BeamSearchOutput(sequences=best_beams)) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 02ff61d62936..ac1a585e8f62 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -197,7 +197,9 @@ async def beam_search( tokens = beam.tokens[tokenized_length:-1] else: tokens = beam.tokens[tokenized_length:] - beam.text = tokenizer.decode(tokens) + beam.text = tokenizer.decode( + tokens, skip_special_tokens=params.skip_special_tokens + ) yield RequestOutput( request_id=request_id, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index b7b292e8355d..badd5a0344c6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -659,6 +659,7 @@ def to_beam_search_params( temperature=temperature, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, + skip_special_tokens=self.skip_special_tokens, ) def extract_structured_outputs(self) -> StructuredOutputsParams | None: diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 6d7a49bacd79..bf0a94eea716 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -303,6 +303,7 @@ def to_beam_search_params( temperature=temperature, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, + skip_special_tokens=self.skip_special_tokens, ) def extract_structured_outputs(self) -> StructuredOutputsParams | None: diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 52a65aefbce1..c0ab45ced702 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1313,6 +1313,7 @@ class BeamSearchParams( length_penalty: float = 1.0 include_stop_str_in_output: bool = False structured_outputs: StructuredOutputsParams | None = None + skip_special_tokens: bool = True def __post_init__(self) -> None: _verify_num_sequences(self.beam_width, "beam_width") From cf1584f373d06f03c5662608fb80cf66c03fcb1d Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 14 Sep 2026 14:17:50 +0800 Subject: [PATCH 159/774] [Refactor] Normalize DeepSeek V4.1 model package naming (#56741) Signed-off-by: Bugen Zhao --- .buildkite/test_areas/distributed.yaml | 2 +- .github/mergify.yml | 5 ++--- tests/distributed/test_engram_dp_shard.py | 10 +++++----- tests/kernels/core/test_fused_q_kv_rmsnorm.py | 8 ++++---- tests/kernels/test_compressor_kv_cache.py | 10 +++++----- tests/kernels/test_engram.py | 12 ++++++------ tests/kernels/test_fused_indexer_q_rope_quant.py | 6 +++--- tests/kernels/test_mhc_kernels.py | 6 +++--- tests/models/test_deepseek_v4_mega_moe.py | 4 ++-- tests/models/test_dspark_mla.py | 2 +- tests/quantization/test_fp8.py | 12 ++++++------ .../test_indexer_deepseek_v4_slot_mapping.py | 2 +- .../model_executor/layers/quantization/__init__.py | 2 +- vllm/model_executor/models/registry.py | 4 ++-- .../{deepseek_v4_1 => deepseek_v41}/__init__.py | 0 .../amd/__init__.py | 0 .../{deepseek_v4_1 => deepseek_v41}/amd/dspark.py | 0 .../{deepseek_v4_1 => deepseek_v41}/amd/model.py | 4 ++-- .../{deepseek_v4_1 => deepseek_v41}/amd/rocm.py | 6 +++--- .../amd/vl_model.py | 0 .../{deepseek_v4_1 => deepseek_v41}/attention.py | 14 +++++++------- .../common/__init__.py | 0 .../common/engram.py | 0 .../common/mm_preprocess.py | 0 .../common/ops/__init__.py | 0 .../common/ops/cache_utils.py | 0 .../common/ops/fused_compress_quant_cache.py | 0 .../common/ops/indexer_k_store.py | 0 .../common/ops/query_quant.py | 0 .../{deepseek_v4_1 => deepseek_v41}/common/rope.py | 0 .../{deepseek_v4_1 => deepseek_v41}/compressor.py | 2 +- .../nvidia/__init__.py | 0 .../nvidia/dspark.py | 0 .../nvidia/engram.py | 6 +++--- .../nvidia/flashinfer_sparse.py | 6 +++--- .../nvidia/flashmla.py | 6 +++--- .../nvidia/model.py | 6 +++--- .../nvidia/model_state.py | 0 .../nvidia/vl_model.py | 0 .../quant_config.py | 0 .../{deepseek_v4_1 => deepseek_v41}/sparse_mla.py | 0 vllm/transformers_utils/configs/deepseek_v41.py | 2 +- vllm/v1/attention/backends/mla/sparse_swa.py | 2 +- vllm/v1/attention/backends/registry.py | 4 ++-- 44 files changed, 71 insertions(+), 72 deletions(-) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/dspark.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/model.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/rocm.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/vl_model.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/attention.py (98%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/engram.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/mm_preprocess.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/cache_utils.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/fused_compress_quant_cache.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/indexer_k_store.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/query_quant.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/rope.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/compressor.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/dspark.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/engram.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/flashinfer_sparse.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/flashmla.py (98%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/model.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/model_state.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/vl_model.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/quant_config.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/sparse_mla.py (100%) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 070664ac7ac3..3ea26ce89c1d 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -208,7 +208,7 @@ steps: - tests/distributed/test_utils - tests/distributed/test_engram_dp_shard.py - vllm/config/engram.py - - vllm/models/deepseek_v4_1/ + - vllm/models/deepseek_v41/ commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 diff --git a/.github/mergify.yml b/.github/mergify.yml index ba3a69ee3575..1267f0264dfe 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -343,9 +343,8 @@ pull_request_rules: conditions: - label != stale - or: - # The 4.1 sources spell the model three ways: deepseek_v4_1 (models), - # deepseek_v41 (most) and deepseekv41 (tool parsers). - - files~=^vllm/models/deepseek_v4_1/ + # The 4.1 sources use deepseek_v41 and deepseekv41 (tool parsers). + - files~=^vllm/models/deepseek_v41/ - files~=^vllm/(?:parser|tokenizers|reasoning|transformers_utils/configs)/deepseek_v41.*\.py - files~=^vllm/tool_parsers/deepseekv41.*\.py - files~=^rust/src/(?:chat|parser)/src/.*deepseek_v41 diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index ebb02542898c..35ffc659d442 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -39,9 +39,9 @@ initialize_model_parallel, ) from vllm.forward_context import set_forward_context -from vllm.models.deepseek_v4_1.common.engram import EngramLayout -from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops -from vllm.models.deepseek_v4_1.nvidia.engram import ( +from vllm.models.deepseek_v41.common.engram import EngramLayout +from vllm.models.deepseek_v41.nvidia import engram as engram_ops +from vllm.models.deepseek_v41.nvidia.engram import ( Engram, ParallelEngramEmbedding, engram_head_shard_rank, @@ -473,8 +473,8 @@ def _check_dummy_hash_model_forward( vllm_config, engram, head_sizes, weight, scales, dp_shared_memory ): """A metadata-less replica joins DP sharding but skips a shared lookup.""" - from vllm.models.deepseek_v4_1.common.engram import NgramHashState - from vllm.models.deepseek_v4_1.nvidia import model as model_ops + from vllm.models.deepseek_v41.common.engram import NgramHashState + from vllm.models.deepseek_v41.nvidia import model as model_ops dp_rank = vllm_config.parallel_config.data_parallel_rank if dp_shared_memory and dp_rank != 1: diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 7e5e41048791..3d0e2ac11e35 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -125,7 +125,7 @@ def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype): from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, ) - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( fused_q_kv_rmsnorm_quant, ) @@ -198,15 +198,15 @@ def test_shared_query_quant_preserves_projection( mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic - from vllm.models.deepseek_v4_1.attention import ( + from vllm.models.deepseek_v41.attention import ( DeepseekV4Attention, DeepseekV4Indexer, ) - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( can_fuse_query_quant, fused_q_kv_rmsnorm_quant, ) - from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config cls: ( type[FlashInferCutedslMxfp8LinearKernel] diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 082a60b7a5a6..0f1976cf683e 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -65,7 +65,7 @@ def test_v41_fused_save_compress_and_insert( a Torch reference with BF16 rounding before RoPE. Replays change both inputs and history so stale scratch/cache values cannot satisfy the test. """ - from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, ) @@ -226,7 +226,7 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): Non-boundary positions, negative slots and untouched rows must keep their prior contents. """ - from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( rope_quant_insert, ) @@ -301,7 +301,7 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): capacity`` and keep padding tokens at PAD.""" from unittest.mock import MagicMock - from vllm.models.deepseek_v4_1.compressor import CompressorMetadataBuilder + from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.kv_cache_interface import CircularBufferSpec @@ -351,11 +351,11 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph): """Both reused-event joins must publish this forward's states and cache rows.""" from vllm.forward_context import ForwardContext, override_forward_context - from vllm.models.deepseek_v4_1.attention import ( + from vllm.models.deepseek_v41.attention import ( DeepseekV4Attention, DeepseekV4Indexer, ) - from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor + from vllm.models.deepseek_v41.compressor import DeepseekCompressor torch.manual_seed(43) raw = torch.randn(19, 1024, device="cuda") diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c7e08ced82ae..c2a7607cbac0 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -7,15 +7,15 @@ import pytest import torch -from vllm.models.deepseek_v4_1.common import engram as engram_ops -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common import engram as engram_ops +from vllm.models.deepseek_v41.common.engram import ( Engram as CommonEngram, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( NgramHashState, ) -from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops -from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding +from vllm.models.deepseek_v41.nvidia import engram as nvidia_engram_ops +from vllm.models.deepseek_v41.nvidia.engram import Engram, ParallelEngramEmbedding from vllm.platforms import current_platform @@ -503,7 +503,7 @@ def fresh_cache(): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") def test_v2_model_state_gathers_lookback_window(): """The window is gathered on device from the runner's token history.""" - from vllm.models.deepseek_v4_1.nvidia.model_state import DeepseekV41ModelState + from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState depth, max_num_reqs, max_model_len = 3, 4, 16 state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 73b147584c51..8770882ddf7c 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -359,7 +359,7 @@ def test_indexer_k_inserts_only_valid_groups_into_padded_pages( num_tokens, compress_ratio, use_fp4, cache_dtype ): """Insert group ends and preserve skipped slots, graph rows, and page padding.""" - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) @@ -442,7 +442,7 @@ def test_indexer_k_inserts_only_valid_groups_into_padded_pages( @torch.inference_mode() def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4): """Replay must consume updated keys and slot mapping through the captured API.""" - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) @@ -501,7 +501,7 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat ``block_size > 1``, so a row-major store would feed the indexer permuted key bytes and randomize its top-k. Assert the write/read pair is exact. """ - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 7f0515d11ade..46f311dfe69b 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -32,7 +32,7 @@ DeepseekV4DecoderLayer, DeepseekV4Model, ) -from vllm.models.deepseek_v4_1.nvidia.model import ( +from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) from vllm.platforms import current_platform @@ -104,7 +104,7 @@ def test_hc_collapse_custom_op_supports_compile(): @pytest.mark.parametrize("num_tokens", [1, 7]) def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm.""" - from vllm.models.deepseek_v4_1.nvidia import dspark + from vllm.models.deepseek_v41.nvidia import dspark set_random_seed(0) hidden_size, hc_mult = 5120, 4 @@ -296,7 +296,7 @@ def reference(*args, norm_weight, norm_eps, **kwargs): return post, res, decoder.attn_norm(collapsed), pre monkeypatch.setattr( - "vllm.models.deepseek_v4_1.nvidia.model.mhc_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", reference, ) expected = decoder(x, positions, None, **kwargs) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index d3d9ccabad09..0d3c2b658369 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,8 +18,8 @@ ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.models.deepseek_v4_1.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID -from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE +from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID +from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 0d1507a03382..d30eb9e7c91d 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -216,7 +216,7 @@ def test_v41_dspark_loads_linear_scales( """Checkpoint ``.scale`` maps to the quant method's scale parameter and loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static loader (see tests/quantization/test_modelopt.py), not in load_weights.""" - from vllm.models.deepseek_v4_1.nvidia import dspark + from vllm.models.deepseek_v41.nvidia import dspark mxfp8 = scale_dtype != torch.float32 scale_name = "weight_scale" if mxfp8 else "weight_scale_inv" diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index a91c99bf5378..65e119b083fd 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -77,8 +77,8 @@ def test_deepseek_v41_mxfp8_scale_loading( MergedColumnParallelLinear, RowParallelLinear, ) - from vllm.models.deepseek_v4_1 import quant_config as quant_module - from vllm.models.deepseek_v4_1.nvidia import model as model_module + from vllm.models.deepseek_v41 import quant_config as quant_module + from vllm.models.deepseek_v41.nvidia import model as model_module default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) quant_config = quant_module.DeepseekV4FP8Config( @@ -214,8 +214,8 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales( """The VL wrapper must map ``.scale`` keys to the parameter the linear quant method registers, as the text model does. A hardcoded ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" - from vllm.models.deepseek_v4_1.nvidia import model as model_module - from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module + from vllm.models.deepseek_v41.nvidia import model as model_module + from vllm.models.deepseek_v41.nvidia import vl_model as vl_module vllm_config = SimpleNamespace( quant_config=SimpleNamespace(weight_block_size=weight_block_size) @@ -318,8 +318,8 @@ def test_mxfp8_bmm_loads_and_projects_grouped_weights( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) - from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model - from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.models.deepseek_v41.nvidia.model import DeepseekV4Model + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config from vllm.utils.deep_gemm import is_deep_gemm_supported if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 49462ea8fba6..87df1e409cf3 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -8,7 +8,7 @@ from tests.v1.attention.utils import create_vllm_config from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, ) from vllm.v1.attention.backend import CommonAttentionMetadata diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index fd1d0470745f..57dc2371318d 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -120,7 +120,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # The v4.1 class is the v4 one extended to accept model_type # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has # platform-specific implementations for both CUDA and ROCm. - from vllm.models.deepseek_v4_1 import ( + from vllm.models.deepseek_v41 import ( DeepseekV4FP8Config as DeepseekV41FP8Config, ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 69dd88f34918..1a5f8399b6ce 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -377,7 +377,7 @@ "DeepseekV4ForConditionalGeneration", ), "DeepseekV41ForCausalLM": ( - "vllm.models.deepseek_v4_1", + "vllm.models.deepseek_v41", "DeepseekV41ForCausalLM", ), "Dots3NoteForCausalLM": ( @@ -645,7 +645,7 @@ "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "DSparkV41DraftModel": ( - "vllm.models.deepseek_v4_1", + "vllm.models.deepseek_v41", "DSparkDeepseekV4ForCausalLM", ), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), diff --git a/vllm/models/deepseek_v4_1/__init__.py b/vllm/models/deepseek_v41/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/__init__.py rename to vllm/models/deepseek_v41/__init__.py diff --git a/vllm/models/deepseek_v4_1/amd/__init__.py b/vllm/models/deepseek_v41/amd/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/__init__.py rename to vllm/models/deepseek_v41/amd/__init__.py diff --git a/vllm/models/deepseek_v4_1/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/dspark.py rename to vllm/models/deepseek_v41/amd/dspark.py diff --git a/vllm/models/deepseek_v4_1/amd/model.py b/vllm/models/deepseek_v41/amd/model.py similarity index 99% rename from vllm/models/deepseek_v4_1/amd/model.py rename to vllm/models/deepseek_v41/amd/model.py index 898d3ab7a025..eaed15214f50 100644 --- a/vllm/models/deepseek_v4_1/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -54,8 +54,8 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v4_1.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention +from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum diff --git a/vllm/models/deepseek_v4_1/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py similarity index 99% rename from vllm/models/deepseek_v4_1/amd/rocm.py rename to vllm/models/deepseek_v41/amd/rocm.py index 9271e7ae5a14..802a59e340fc 100644 --- a/vllm/models/deepseek_v4_1/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -13,9 +13,9 @@ ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import dequantize_and_gather_k_cache -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, DeepseekV4SparseMLAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/vl_model.py rename to vllm/models/deepseek_v41/amd/vl_model.py diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v41/attention.py similarity index 98% rename from vllm/models/deepseek_v4_1/attention.py rename to vllm/models/deepseek_v41/attention.py index 9bb53526c99b..131cf4d97e7e 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -26,7 +26,7 @@ ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant, indexer_k_norm_rope_store, @@ -49,8 +49,8 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index -from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope -from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor +from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v41.compressor import DeepseekCompressor from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -156,7 +156,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ``DeepseekV4FlashInferSM120Attention`` / ``DeepseekV4FlashInferMLAAttention`` (CUDA) or ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific - deepseek_v4_1 model module. The base is never instantiated directly. + deepseek_v41 model module. The base is never instantiated directly. """ # Provided by the platform subclass. @@ -536,7 +536,7 @@ def __init__( "FLASHMLA_SPARSE_DSV41", "ROCM_FLASHMLA_SPARSE_DSV4", ): - from vllm.models.deepseek_v4_1.common.ops.cache_utils import ( + from vllm.models.deepseek_v41.common.ops.cache_utils import ( _COMBINE_TOPK_SWA_INDICES_KERNEL, ) @@ -581,7 +581,7 @@ def forward( @cached_property def _can_fuse_query_quant(self) -> bool: - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( can_fuse_query_quant, ) @@ -600,7 +600,7 @@ def _split_qkv_and_norm( """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant: - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( fused_q_kv_rmsnorm_quant, ) diff --git a/vllm/models/deepseek_v4_1/common/__init__.py b/vllm/models/deepseek_v41/common/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/__init__.py rename to vllm/models/deepseek_v41/common/__init__.py diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v41/common/engram.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/engram.py rename to vllm/models/deepseek_v41/common/engram.py diff --git a/vllm/models/deepseek_v4_1/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/mm_preprocess.py rename to vllm/models/deepseek_v41/common/mm_preprocess.py diff --git a/vllm/models/deepseek_v4_1/common/ops/__init__.py b/vllm/models/deepseek_v41/common/ops/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/__init__.py rename to vllm/models/deepseek_v41/common/ops/__init__.py diff --git a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/cache_utils.py rename to vllm/models/deepseek_v41/common/ops/cache_utils.py diff --git a/vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py rename to vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py diff --git a/vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py rename to vllm/models/deepseek_v41/common/ops/indexer_k_store.py diff --git a/vllm/models/deepseek_v4_1/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/query_quant.py rename to vllm/models/deepseek_v41/common/ops/query_quant.py diff --git a/vllm/models/deepseek_v4_1/common/rope.py b/vllm/models/deepseek_v41/common/rope.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/rope.py rename to vllm/models/deepseek_v41/common/rope.py diff --git a/vllm/models/deepseek_v4_1/compressor.py b/vllm/models/deepseek_v41/compressor.py similarity index 99% rename from vllm/models/deepseek_v4_1/compressor.py rename to vllm/models/deepseek_v41/compressor.py index d09d97d04c27..44440fcb6114 100644 --- a/vllm/models/deepseek_v4_1/compressor.py +++ b/vllm/models/deepseek_v41/compressor.py @@ -12,7 +12,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear -from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( +from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, ) diff --git a/vllm/models/deepseek_v4_1/nvidia/__init__.py b/vllm/models/deepseek_v41/nvidia/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/__init__.py rename to vllm/models/deepseek_v41/nvidia/__init__.py diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/dspark.py rename to vllm/models/deepseek_v41/nvidia/dspark.py diff --git a/vllm/models/deepseek_v4_1/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/engram.py rename to vllm/models/deepseek_v41/nvidia/engram.py index b141b6c34975..624a36c214fc 100644 --- a/vllm/models/deepseek_v4_1/nvidia/engram.py +++ b/vllm/models/deepseek_v41/nvidia/engram.py @@ -23,16 +23,16 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( DEAD_ID, EngramLayout, _engram_head_shard_weight_loader, _engram_select_rows, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( Engram as BaseEngram, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( ParallelEngramEmbedding as BaseParallelEngramEmbedding, ) from vllm.utils.platform_utils import is_uva_available diff --git a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py rename to vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 802ddede9c64..7bed50bde71a 100644 --- a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -13,12 +13,12 @@ compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import ( build_flashinfer_mixed_sparse_indices, compute_global_topk_indices_and_lens, ) -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, DeepseekV4SparseMLAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py similarity index 98% rename from vllm/models/deepseek_v4_1/nvidia/flashmla.py rename to vllm/models/deepseek_v41/nvidia/flashmla.py index 86d8fb2031ca..ec834aa80dd9 100644 --- a/vllm/models/deepseek_v4_1/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -10,13 +10,13 @@ compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import ( combine_topk_swa_indices, compute_global_topk_indices_and_lens, dequantize_and_gather_k_cache, ) -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLABackend, DeepseekV4FlashMLAMetadata, DeepseekV41SparseSWAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/model.py rename to vllm/models/deepseek_v41/nvidia/model.py index 0868089da64b..df12a9d9bbc9 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -63,12 +63,12 @@ from vllm.models.deepseek_v4.nvidia.model import ( make_deepseek_v4_expert_params_mapping, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, ) -from vllm.models.deepseek_v4_1.nvidia.flashmla import DeepseekV4FlashMLAAttention +from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv diff --git a/vllm/models/deepseek_v4_1/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/model_state.py rename to vllm/models/deepseek_v41/nvidia/model_state.py diff --git a/vllm/models/deepseek_v4_1/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/vl_model.py rename to vllm/models/deepseek_v41/nvidia/vl_model.py diff --git a/vllm/models/deepseek_v4_1/quant_config.py b/vllm/models/deepseek_v41/quant_config.py similarity index 100% rename from vllm/models/deepseek_v4_1/quant_config.py rename to vllm/models/deepseek_v41/quant_config.py diff --git a/vllm/models/deepseek_v4_1/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py similarity index 100% rename from vllm/models/deepseek_v4_1/sparse_mla.py rename to vllm/models/deepseek_v41/sparse_mla.py diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py index 0be36d2480db..8f33245cd68f 100644 --- a/vllm/transformers_utils/configs/deepseek_v41.py +++ b/vllm/transformers_utils/configs/deepseek_v41.py @@ -12,7 +12,7 @@ class DeepseekV41Config(PretrainedConfig): tower under ``vision_config``. The vLLM model code (ported from ``deepseek_v4``) reads flat attributes, so both sub-configs are flattened onto the top level here: text fields are exposed as-is, vision fields - with the ``vision_*`` naming used by ``deepseek_v4_1.common.vision``. + with the ``vision_*`` naming used by ``deepseek_v41.common.vision``. """ model_type = "deepseek_v41" diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index f47554d81f0c..62c3921b5e12 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -49,7 +49,7 @@ # v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but # their compressed page block is block_size // ratio, so each gets its own plan. # v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see -# deepseek_v4_1/sparse_mla.py. +# deepseek_v41/sparse_mla.py. _LAYER_TYPE_C1A = "c1a" _LAYER_TYPE_C2A = "c2a" diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 085f2358fabb..571566db1f25 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -108,10 +108,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): # layer). Separate names from DSV4 so a V4.1 model never resolves the # V4.0 backend classes through this enum. FLASHMLA_SPARSE_DSV41 = ( - "vllm.models.deepseek_v4_1.sparse_mla.DeepseekV4FlashMLABackend" + "vllm.models.deepseek_v41.sparse_mla.DeepseekV4FlashMLABackend" ) FLASHINFER_MLA_SPARSE_DSV41 = ( - "vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse." + "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend" From dbf49dad11e30ec5d464be74b715a9a9e31c2c5a Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 14 Sep 2026 14:28:13 +0800 Subject: [PATCH 160/774] [Fast Start] Use GPU uuid as socket folder identifier (#56669) Signed-off-by: Isotr0py --- .../model_loader/test_weight_cache.py | 8 ++-- .../model_loader/weight_cache/daemon.py | 26 ++++--------- .../model_loader/weight_cache/ipc_loader.py | 27 ++++--------- .../model_loader/weight_cache/protocol.py | 38 ++++++------------- 4 files changed, 30 insertions(+), 69 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index fa0b6b418343..1195b8b09d1c 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -7,6 +7,7 @@ identical outputs. """ +import glob import os import shutil import subprocess @@ -21,7 +22,6 @@ from vllm import SamplingParams from vllm.assets.image import ImageAsset -from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path from vllm.platforms import current_platform DAEMON_TIMEOUT_S = 600 @@ -79,9 +79,7 @@ def _wait_ready(self, timeout_s: float) -> None: # Poll for the socket files rather than a log line: model loading can # pull in JIT compilers that swap the process's stderr and swallow # everything logged afterwards, making log-based readiness flaky. - expected = [ - get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size) - ] + pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock") deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if self._proc.poll() is not None: @@ -89,7 +87,7 @@ def _wait_ready(self, timeout_s: float) -> None: f"Weight cache daemon exited with {self._proc.returncode}:\n" f"{self._logs()}" ) - if all(os.path.exists(path) for path in expected): + if len(glob.glob(pattern)) >= self.tp_size: return time.sleep(1.0) raise TimeoutError( diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 96cedbe1e91d..b19fb9627237 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -60,6 +60,7 @@ ensure_model_parallel_initialized, init_distributed_environment, ) +from vllm.engine.arg_utils import EngineArgs from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader.utils import process_weights_after_loading @@ -70,13 +71,15 @@ check_ipc_platform_support, check_ipc_quant_support, ensure_private_socket_dir, - get_physical_device_id, + get_current_device_uuid, get_socket_path, recv_msg, send_msg, verify_peer_is_owner, ) from vllm.platforms import current_platform +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.network_utils import get_distributed_init_method, get_open_port from vllm.utils.torch_utils import set_default_torch_dtype logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") @@ -209,7 +212,7 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non ready_callback: Invoked once the socket is bound and listening, so the launcher can report overall readiness. """ - socket_path = self._socket_path() + socket_path = self._socket_path ensure_private_socket_dir( os.path.dirname(socket_path), strict_perms=self.socket_dir is None ) @@ -268,12 +271,9 @@ def _acquire_gpu_lock(self, socket_path: str) -> int: ) from e return lock_fd + @property def _socket_path(self) -> str: - device_index = torch.accelerator.current_device_index() - gpu_id = get_physical_device_id(device_index) - if gpu_id is None: - gpu_id = device_index - return get_socket_path(gpu_id, self.socket_dir) + return get_socket_path(get_current_device_uuid(), self.socket_dir) def _handle_connection(self, conn: socket.socket) -> None: request = recv_msg(conn) @@ -304,7 +304,7 @@ def _handle_get_state(self, conn: socket.socket, request: dict) -> None: "status": "ok", "entries": self.entries, "aliases": self.aliases, - "gpu_uuid": self._gpu_uuid(), + "gpu_uuid": get_current_device_uuid(), }, ) logger.info_once( @@ -322,12 +322,6 @@ def _handle_release(self, conn: socket.socket) -> None: logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank) send_msg(conn, {"status": "ok"}) - def _gpu_uuid(self) -> str: - props = torch.cuda.get_device_properties( - torch.accelerator.current_device_index() - ) - return str(props.uuid) - def _run_daemon( tp_rank: int, @@ -359,10 +353,6 @@ def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: def main() -> None: - from vllm.engine.arg_utils import EngineArgs - from vllm.utils.argparse_utils import FlexibleArgumentParser - from vllm.utils.network_utils import get_distributed_init_method, get_open_port - parser = FlexibleArgumentParser( description="Launch weight cache daemons (one per TP rank)." ) diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index be2e5ebae4cf..6c3fb742cdf7 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -12,6 +12,10 @@ from vllm.config import ModelConfig, VllmConfig from vllm.config.load import LoadConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader @@ -27,7 +31,7 @@ WeightCacheUnavailableError, check_ipc_platform_support, check_ipc_quant_support, - get_physical_device_id, + get_current_device_uuid, get_socket_path, recv_msg, send_msg, @@ -56,7 +60,7 @@ class IpcModelLoader(BaseModelLoader): Extra config keys (via --model-loader-extra-config): - socket_path: explicit daemon socket path. Defaults to a per-GPU path - derived from the physical GPU id. + derived from the physical GPU uuid. - socket_dir: directory containing the daemon sockets. - mode: "zero_copy" (default) or "copy". - fallback: fall back to disk loading when the daemon is unavailable or @@ -261,11 +265,6 @@ def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None: def _fetch_entries( self, model_config: ModelConfig ) -> tuple[dict[str, TensorEntry], dict[str, str]]: - from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - ) - cache_config = WeightCacheKey.from_model_config( model_config, tp_size=get_tensor_model_parallel_world_size(), @@ -317,22 +316,12 @@ def _connect(self, timeout: float) -> socket.socket: def _resolve_socket_path(self) -> str: if self.socket_path is not None: return self.socket_path - device_index = torch.accelerator.current_device_index() - gpu_id = get_physical_device_id(device_index) - if gpu_id is None: - raise WeightCacheUnavailableError( - "Cannot infer the physical GPU id from CUDA_VISIBLE_DEVICES; " - "pass socket_path via --model-loader-extra-config" - ) - return get_socket_path(gpu_id, self.socket_dir) + return get_socket_path(get_current_device_uuid(), self.socket_dir) def _check_gpu_uuid(self, daemon_uuid: str | None) -> None: if daemon_uuid is None: return - props = torch.cuda.get_device_properties( - torch.accelerator.current_device_index() - ) - local_uuid = str(props.uuid) + local_uuid = get_current_device_uuid() if daemon_uuid != local_uuid: raise CacheConfigMismatchError( f"Daemon GPU {daemon_uuid} != engine GPU {local_uuid}; " diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index fff1fcca2a0b..00d3dbfa8ba9 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -23,14 +23,19 @@ from typing import Any import torch +from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME import vllm.version from vllm.config import ModelConfig from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.model_loader.weight_utils import ( + filter_duplicate_safetensors_files, +) from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash -SOCKET_NAME_TEMPLATE = "vllm_weight_cache_gpu{gpu_id}.sock" +SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}.sock" SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}" _LEN_STRUCT = struct.Struct("!Q") @@ -102,20 +107,9 @@ def check_ipc_quant_support(model: torch.nn.Module) -> None: ) -def get_physical_device_id(device_index: int) -> int | None: - """Map a local CUDA device index to the physical GPU id. - - Returns None if CUDA_VISIBLE_DEVICES contains non-integer entries - (e.g. GPU UUIDs), in which case an explicit socket path is required. - """ - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if not visible: - return device_index - entries = visible.split(",") - try: - return int(entries[device_index]) - except (IndexError, ValueError): - return None +def get_current_device_uuid() -> str: + """UUID of the physical GPU backing the current accelerator device.""" + return current_platform.get_device_uuid(torch.accelerator.current_device_index()) def get_socket_dir(socket_dir: str | None = None) -> str: @@ -132,9 +126,9 @@ def get_socket_dir(socket_dir: str | None = None) -> str: ) -def get_socket_path(gpu_id: int, socket_dir: str | None = None) -> str: +def get_socket_path(gpu_uuid: str, socket_dir: str | None = None) -> str: directory = get_socket_dir(socket_dir) - return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_id=gpu_id)) + return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_uuid=gpu_uuid)) def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None: @@ -245,12 +239,6 @@ def hash_checkpoint(model: str) -> str | None: """ if not os.path.isdir(model): return None - from transformers.utils import SAFE_WEIGHTS_INDEX_NAME - - from vllm.model_executor.model_loader.weight_utils import ( - filter_duplicate_safetensors_files, - ) - files = glob.glob(os.path.join(model, "*.safetensors")) if os.path.isfile(os.path.join(model, SAFE_WEIGHTS_INDEX_NAME)): files = filter_duplicate_safetensors_files( @@ -337,8 +325,6 @@ class TensorEntry: @classmethod def from_tensor(cls, tensor: torch.Tensor, kind: str) -> "TensorEntry": - from torch.multiprocessing.reductions import reduce_tensor - tensor = tensor.detach() if tensor.is_cuda: _, ipc_args = reduce_tensor(tensor) @@ -349,8 +335,6 @@ def rebuild(self, device_index: int) -> torch.Tensor: if self.ipc_args is None: assert self.cpu_tensor is not None return self.cpu_tensor - from torch.multiprocessing.reductions import rebuild_cuda_tensor - args = list(self.ipc_args) # Index 6 of the args from reduce_tensor is the device index. It must # be retargeted to the local index since the daemon and the engine may From 663d7f679eda78a8d48613dfede8a4b4bcff2b74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hung=20Hoang=20=28=E9=BB=84=E5=85=B4=EF=BC=89?= Date: Mon, 14 Sep 2026 13:55:21 +0700 Subject: [PATCH 161/774] [Bugfix] Fix --lora-modules name=path parsing when path contains '=' (#53353) Signed-off-by: hungh --- vllm/entrypoints/launchers/cli_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index c0dd64164e95..0da96b43e1dc 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -47,7 +47,7 @@ def __call__( if item in [None, ""]: # Skip if item is None or empty string continue if "=" in item and "," not in item: # Old format: name=path - name, path = item.split("=") + name, path = item.split("=", 1) lora_list.append(LoRAModulePath(name, path)) else: # Assume JSON format try: From e6b4e47d2d894a1a0815fbf03669adb1b3405083 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Mon, 14 Sep 2026 00:05:33 -0700 Subject: [PATCH 162/774] [Bugfix][Rust Frontend] Preserve selected-token logprob mode (#56406) Co-authored-by: OpenAI Codex Signed-off-by: Alec Flowers --- rust/src/server/src/grpc/convert.rs | 98 ++++++++++++++++++++++++----- 1 file changed, 84 insertions(+), 14 deletions(-) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 61eb21ef70ee..aea04495eb39 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -193,7 +193,7 @@ fn build_sampling_params( if let Some(r) = response { if r.output_logprobs { let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref()); - params.logprobs = Some(count); + params.logprobs = count; params.logprob_token_ids = token_ids; } if r.prompt_logprobs { @@ -210,7 +210,7 @@ fn build_sampling_params( )); } let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref()); - params.prompt_logprobs = Some(count); + params.prompt_logprobs = count; } } @@ -220,18 +220,22 @@ fn build_sampling_params( /// Map the proto `CandidateTokens` selector to a `(logprobs_count, /// logprob_token_ids)` pair. /// -/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability -/// - `all` → `(-1, None)` — return the full vocabulary -/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the -/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs -/// are carried via `logprob_token_ids` on `SamplingParams`) -/// - absent → `(1, None)` — just the sampled/scored token -fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option>) { +/// - `top_n(k)` → `(Some(k), None)` — return top-k candidates by probability +/// - `all` → `(Some(-1), None)` — return the full vocabulary +/// - nonempty `token_ids` → `(None, Some(ids))` — let the engine derive the count +/// - absent or empty `token_ids` → `(Some(0), None)` — return only the sampled/scored token +fn candidate_logprob_spec( + candidates: Option<&pb::CandidateTokens>, +) -> (Option, Option>) { match candidates.and_then(|c| c.select.as_ref()) { - Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None), - Some(pb::candidate_tokens::Select::All(true)) => (-1, None), - Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())), - _ => (1, None), + Some(pb::candidate_tokens::Select::TopN(n)) => (Some(*n as i32), None), + Some(pb::candidate_tokens::Select::All(true)) => (Some(-1), None), + Some(pb::candidate_tokens::Select::TokenIds(ids)) if !ids.ids.is_empty() => { + // Match Python HTTP: selected IDs have their own limit, independent + // of the numeric top-logprobs count and its max_logprobs cap. + (None, Some(ids.ids.clone())) + } + _ => (Some(0), None), } } @@ -518,7 +522,10 @@ impl ResponseOpts { #[cfg(test)] mod tests { use vllm_engine_core_client::protocol::output::StopReason; - use vllm_text::{FinishReason, Finished, Prompt}; + use vllm_text::{ + FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, + }; + use vllm_tokenizer::test_utils::TestTokenizer; use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request}; @@ -590,6 +597,69 @@ mod tests { assert_eq!(text.sampling_params.seed, Some(0)); } + #[test] + fn output_logprob_selectors_survive_request_lowering() { + use pb::candidate_tokens::Select; + + let selected_ids: Vec = (100..121).collect(); + let cases = [ + ( + Some(Select::TokenIds(pb::TokenIds { + ids: selected_ids.clone(), + })), + None, + Some(selected_ids), + ), + ( + Some(Select::TokenIds(pb::TokenIds { + ids: vec![198, 198], + })), + None, + Some(vec![198, 198]), + ), + ( + Some(Select::TokenIds(pb::TokenIds { ids: vec![] })), + Some(0), + None, + ), + (None, Some(0), None), + (Some(Select::TopN(0)), Some(0), None), + (Some(Select::TopN(2)), Some(2), None), + ]; + for (select, expected_count, expected_ids) in cases { + let req = pb::GenerateRequest { + response: Some(pb::ResponseOptions { + output_logprobs: true, + output_candidates: select.map(|select| pb::CandidateTokens { + select: Some(select), + }), + ..Default::default() + }), + ..base_request() + }; + let text = + to_text_request(req, false, &["test-model".to_string()]).expect("convert ok"); + let params = lower_sampling_params( + text.sampling_params, + SamplingHints::default(), + SamplingLimits { + max_model_len: 32, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + 1, + &TestTokenizer::new(), + ) + .expect("logprob selector should lower to an engine request"); + + assert_eq!( + (params.logprobs, params.logprob_token_ids), + (expected_count, expected_ids) + ); + } + } + #[test] fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() { let req = pb::GenerateRequest { From 78e84261ab2ffb691a8079d698d1f3f96ecb7635 Mon Sep 17 00:00:00 2001 From: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:42:15 +0900 Subject: [PATCH 163/774] [Performance] Add Triton kernel for Gemma3n sparse GELU (#48498) Signed-off-by: BWAAEEEK Co-authored-by: Andreas Karatzas Co-authored-by: Cyrus Leung --- .../kernels/benchmark_gelu_and_mul_sparse.py | 90 +++++++++ tests/compile/passes/ir/test_lowering.py | 34 ++++ tests/compile/test_aot_compile.py | 65 +++++++ tests/kernels/ir/test_activation.py | 183 ++++++++++++++++++ tests/test_config.py | 16 ++ vllm/config/kernel.py | 3 + vllm/ir/ops/__init__.py | 3 +- vllm/ir/ops/activation.py | 34 ++++ vllm/kernels/__init__.py | 4 +- vllm/kernels/triton/__init__.py | 4 + vllm/kernels/triton/activation.py | 110 +++++++++++ vllm/model_executor/layers/activation.py | 18 +- vllm/platforms/cuda.py | 5 +- vllm/platforms/rocm.py | 5 +- vllm/platforms/xpu.py | 2 +- 15 files changed, 555 insertions(+), 21 deletions(-) create mode 100644 benchmarks/kernels/benchmark_gelu_and_mul_sparse.py create mode 100644 tests/kernels/ir/test_activation.py create mode 100644 vllm/ir/ops/activation.py create mode 100644 vllm/kernels/triton/activation.py diff --git a/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py new file mode 100644 index 000000000000..9f8d97612161 --- /dev/null +++ b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import itertools + +import torch + +import vllm.kernels # noqa: F401 +from vllm import ir +from vllm.benchmarks.lib.utils import default_vllm_config +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed + +NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048] +INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768] +CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES)) +STD_MULTIPLIER = 1.6448536269514722 + + +@default_vllm_config() +def benchmark_gelu_and_mul_sparse( + num_tokens: int, + intermediate_size: int, + provider: str, + dtype: torch.dtype, +): + set_random_seed(0) + torch.set_default_device("cuda") + x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype) + native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn + args = (x, STD_MULTIPLIER, "tanh") + + expected = native(*args) + actual = triton_impl(*args) + torch.testing.assert_close( + actual, + expected, + **ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype), + ) + + if provider == "triton": + fn = lambda: triton_impl(*args) + elif provider == "native_compiled": + compiled = torch.compile(native, fullgraph=True) + compiled(*args) + fn = lambda: compiled(*args) + else: + fn = lambda: native(*args) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=[0.5, 0.2, 0.8] + ) + return ms, max_ms, min_ms + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark the Gemma3n sparse GELU activation." + ) + parser.add_argument( + "--dtype", + type=str, + choices=["half", "bfloat16", "float"], + default="bfloat16", + ) + args = parser.parse_args() + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype] + + perf_report = triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "intermediate_size"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=["triton", "native_compiled", "native_eager"], + line_names=["Triton", "Native (torch.compile)", "Native (eager)"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="ms", + plot_name="gelu-and-mul-sparse-performance", + args={}, + ) + ) + perf_report( + lambda num_tokens, intermediate_size, provider: ( + benchmark_gelu_and_mul_sparse( + num_tokens, intermediate_size, provider, dtype + ) + ) + ).run(print_data=True) diff --git a/tests/compile/passes/ir/test_lowering.py b/tests/compile/passes/ir/test_lowering.py index b7ca55e7d1e6..0ac605e7961c 100644 --- a/tests/compile/passes/ir/test_lowering.py +++ b/tests/compile/passes/ir/test_lowering.py @@ -11,6 +11,7 @@ ) from vllm.config import get_current_vllm_config from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse from vllm.platforms import current_platform from ...backend import TestBackend @@ -34,6 +35,15 @@ def forward(self, x): return x6 + 3.0 +class SparseActivationModel(nn.Module): + def __init__(self): + super().__init__() + self.activation = GeluAndMulSparse(0.95, "tanh") + + def forward(self, x): + return self.activation(x) + + @pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers()) def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.set_default_device(current_platform.device_type) @@ -67,3 +77,27 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.testing.assert_close(output_unlowered, output) torch.testing.assert_close(output_unlowered, output2) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_lowering_gelu_and_mul_sparse(default_vllm_config): + torch.set_default_device("cuda") + lowering_pass = VllmIRLoweringPass(get_current_vllm_config()) + backend = TestBackend(lowering_pass) + model = SparseActivationModel() + x = torch.randn(8, 32768, dtype=torch.bfloat16) + + with ( + ops.gelu_and_mul_sparse.set_priority(["triton", "native"]), + ir.enable_torch_wrap(True), + ): + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + output = compiled_model(x) + + selected = lowering_pass.selected_impls["gelu_and_mul_sparse"] + assert selected == {"gelu_and_mul_sparse": "triton"} + + expected = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + x, 1.6448536269514722, "tanh" + ) + torch.testing.assert_close(output, expected, rtol=1.6e-2, atol=1e-3) diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index b3a71b4c3017..90c132bee1ab 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -13,6 +13,8 @@ import torch import vllm.envs as envs +import vllm.kernels # noqa: F401 to register kernels +from vllm import ir from vllm.compilation.backends import VllmBackend from vllm.compilation.caching import ( StandaloneCompiledArtifacts, @@ -28,6 +30,9 @@ ) from vllm.envs import disable_envs_cache from vllm.forward_context import set_forward_context +from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse +from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer from ..utils import create_new_process_for_each_test @@ -77,6 +82,18 @@ def forward(self, x: torch.Tensor): return reference_fn_tuple(x) +@support_torch_compile(dynamic_arg_dims={"x": 0}) +class SparseActivationAotModel(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.up_proj = torch.nn.Linear(32, 256, bias=False) + self.activation = GeluAndMulSparse(0.95, "tanh") + self.down_proj = torch.nn.Linear(128, 16, bias=False) + + def forward(self, x: torch.Tensor): + return self.down_proj(self.activation(self.up_proj(x))) + + def make_vllm_config() -> VllmConfig: return VllmConfig( compilation_config=CompilationConfig( @@ -175,6 +192,54 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch): assert torch.allclose(ret, expected) +@pytest.mark.skipif( + not current_platform.is_cuda() or not is_torch_equal_or_newer("2.10.0"), + reason="requires CUDA and torch 2.10", +) +def test_sparse_activation_standalone_aot_preserves_dependencies( + monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path +): + def native_reference(model: SparseActivationAotModel, x: torch.Tensor): + projected = torch.nn.functional.linear(x, model.up_proj.weight) + activated = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + projected, model.activation.std_multiplier, "tanh" + ) + return torch.nn.functional.linear(activated, model.down_proj.weight) + + with monkeypatch.context() as m: + m.setenv("VLLM_USE_AOT_COMPILE", "1") + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1") + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1") + m.setenv("VLLM_CACHE_ROOT", str(vllm_tmp_cache / "vllm_cache")) + + vllm_config = make_vllm_config() + torch.manual_seed(0) + inputs = [ + torch.randn(4, 32, device="cuda", dtype=torch.bfloat16), + torch.randn(7, 32, device="cuda", dtype=torch.bfloat16), + ] + + with ( + use_vllm_config(vllm_config), + vllm_config.kernel_config.ir_op_priority.set_priority(), + ir.enable_torch_wrap(True), + ): + compiled = SparseActivationAotModel(vllm_config=vllm_config).to( + device="cuda", dtype=torch.bfloat16 + ) + compiled.requires_grad_(False) + expected = [native_reference(compiled, x) for x in inputs] + actual = [compiled(x) for x in inputs] + + for result, reference in zip(actual, expected): + assert torch.isfinite(result).all() + torch.testing.assert_close( + result, + reference, + **ops.gelu_and_mul_sparse.get_tolerance(result.dtype), + ) + + @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch): from torch._subclasses import FakeTensorMode diff --git a/tests/kernels/ir/test_activation.py b/tests/kernels/ir/test_activation.py new file mode 100644 index 000000000000..a36ff8ee339b --- /dev/null +++ b/tests/kernels/ir/test_activation.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.kernels # noqa: F401 +from tests.ir.ir_test_utils import assert_close, clone_args +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + + +def test_gelu_and_mul_sparse_registration(): + expected = { + "native": True, + "triton": HAS_TRITON and current_platform.is_cuda(), + } + actual = { + provider: impl.supported + for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items() + } + assert actual == expected + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_tokens", [1, 8, 64]) +@pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768]) +def test_gelu_and_mul_sparse_triton( + dtype: torch.dtype, num_tokens: int, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=num_tokens, + hidden_size=intermediate_size, + dtype=dtype, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + assert impl.supports_args(*args) + + expected = gelu_and_mul_sparse_native(*clone_args(args)) + actual = impl.impl_fn(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0)) + assert mask_mismatch / actual.numel() < 1e-3 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_second_seed(): + torch.manual_seed(1) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=32, + hidden_size=8192, + dtype=torch.bfloat16, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("intermediate_size", [1, 31, 128]) +def test_gelu_and_mul_sparse_triton_small_sizes( + dtype: torch.dtype, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=5, + hidden_size=intermediate_size, + dtype=dtype, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + assert torch.equal(actual == 0, expected == 0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_nd_and_empty(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + x = torch.randn(2, 3, 16384, dtype=torch.bfloat16) + args = (x, 1.6448536269514722, "tanh") + actual = impl.impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + assert actual.shape == (2, 3, 8192) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + empty = torch.empty(0, 16384, dtype=torch.bfloat16) + empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh") + assert empty_out.shape == (0, 8192) + assert empty_out.dtype == empty.dtype + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype): + torch.set_default_device("cuda") + gate = torch.full((4, 8192), 3.0, dtype=dtype) + up = torch.randn_like(gate) + x = torch.cat((gate, up), dim=-1) + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + assert torch.count_nonzero(actual) == 0 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")]) +def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float): + torch.set_default_device("cuda") + x = torch.randn(2, 16384, dtype=torch.bfloat16) + x[0, 0] = nonfinite + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, equal_nan=True) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_fallback_conditions(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + x = torch.randn(4, 16384, dtype=torch.bfloat16) + + assert impl.supports_args(x, 1.0, "tanh") + assert not impl.supports_args(x, 1.0, "none") + assert not impl.supports_args(x[:, ::2], 1.0, "tanh") + assert not impl.supports_args(x[:, :-1], 1.0, "tanh") + assert not impl.supports_args( + torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh" + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck(): + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=8, + hidden_size=4096, + dtype=torch.bfloat16, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]): + dispatched = ir.ops.gelu_and_mul_sparse(*args) + torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args) + + direct = impl.impl_fn(*args) + torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_cudagraph(): + torch.set_default_device("cuda") + x = torch.randn(16, 32768, dtype=torch.bfloat16) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + for _ in range(3): + expected = impl.impl_fn(x, 1.6448536269514722, "tanh") + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = impl.impl_fn(x, 1.6448536269514722, "tanh") + graph.replay() + torch.accelerator.synchronize() + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) diff --git a/tests/test_config.py b/tests/test_config.py index c2629a75ab0a..4d39b56d265f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3322,6 +3322,22 @@ def test_ir_op_priority_default(): assert priority_config.fused_add_rms_norm == ["native"] +@pytest.mark.parametrize("mode", [CompilationMode.NONE, CompilationMode.VLLM_COMPILE]) +@pytest.mark.parametrize("backend", ["inductor", "eager"]) +def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend): + """Worker initialization must not select an unregistered sparse GELU provider.""" + from vllm import ir + + config = SimpleNamespace( + compilation_config=CompilationConfig(mode=mode, backend=backend) + ) + priority = current_platform.get_default_ir_op_priority(config) + expected = ["triton", "native"] if current_platform.is_cuda() else ["native"] + + with priority.set_priority(): + assert ir.ops.gelu_and_mul_sparse.get_priority() == expected + + def test_ir_op_priority_str(): """Test that passing a comma-delimited string works""" from vllm.config.kernel import IrOpPriorityConfig diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index d44e83bf67c5..4505e0cfd9a0 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -34,6 +34,9 @@ class IrOpPriorityConfig: fused_add_rms_norm: list[str] = Field(default_factory=list) """Priority list for vllm.ir.ops.fused_add_rms_norm""" + gelu_and_mul_sparse: list[str] = Field(default_factory=list) + """Priority list for vllm.ir.ops.gelu_and_mul_sparse""" + def compute_hash(self) -> str: """ Produces a hash unique to the pass configuration. diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py index d4d71afef723..13e2b7dda988 100644 --- a/vllm/ir/ops/__init__.py +++ b/vllm/ir/ops/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .activation import gelu_and_mul_sparse from .layernorm import fused_add_rms_norm, rms_norm -__all__ = ["rms_norm", "fused_add_rms_norm"] +__all__ = ["rms_norm", "fused_add_rms_norm", "gelu_and_mul_sparse"] diff --git a/vllm/ir/ops/activation.py b/vllm/ir/ops/activation.py new file mode 100644 index 000000000000..149197d73550 --- /dev/null +++ b/vllm/ir/ops/activation.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import Tensor + +from ..op import register_op + + +@register_op +def gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + """Apply Gaussian sparsification, GELU, and gated multiplication.""" + d = x.shape[-1] // 2 + gate = x[..., :d] + # Statistics intentionally remain local to each tensor-parallel shard. + mean = torch.mean(gate, dim=-1, keepdim=True) + std = torch.std(gate, dim=-1, keepdim=True, unbiased=False) + sparse_gate = F.relu(gate - (mean + std * std_multiplier)) + return F.gelu(sparse_gate, approximate=approximate) * x[..., d:] + + +@gelu_and_mul_sparse.register_input_generator +def _gelu_and_mul_sparse_input_generator( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + std_multiplier: float = 1.6448536269514722, + approximate: str = "tanh", +) -> tuple: + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) + return x, std_multiplier, approximate diff --git a/vllm/kernels/__init__.py b/vllm/kernels/__init__.py index e83e2772f06d..4821ec43eccd 100644 --- a/vllm/kernels/__init__.py +++ b/vllm/kernels/__init__.py @@ -2,6 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kernel implementations for vLLM.""" -from . import aiter_ops, oink_ops, vllm_c +from . import aiter_ops, oink_ops, triton, vllm_c -__all__ = ["vllm_c", "aiter_ops", "oink_ops"] +__all__ = ["vllm_c", "aiter_ops", "oink_ops", "triton"] diff --git a/vllm/kernels/triton/__init__.py b/vllm/kernels/triton/__init__.py index 6626213145ca..3a3833505062 100644 --- a/vllm/kernels/triton/__init__.py +++ b/vllm/kernels/triton/__init__.py @@ -1,3 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Triton kernel implementations.""" + +from . import activation + +__all__ = ["activation"] diff --git a/vllm/kernels/triton/activation.py b/vllm/kernels/triton/activation.py new file mode 100644 index 000000000000..2bda1f51b4d1 --- /dev/null +++ b/vllm/kernels/triton/activation.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import Tensor +from torch.library import triton_op, wrap_triton + +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_MAX_INTERMEDIATE_SIZE = 32768 +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} + + +@triton.jit +def _gelu_and_mul_sparse_kernel( + x_ptr, + output_ptr, + d: tl.constexpr, + std_multiplier, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < d + row_start = row * 2 * d + + gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32) + # TODO: Investigate delaying this load to reduce spills for large reductions + # without regressing smaller shapes. + up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32) + + input_dtype = x_ptr.dtype.element_ty + mean = tl.sum(gate, axis=0) / d + centered = tl.where(mask, gate - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / d + # Match the input-dtype rounding at each native PyTorch op boundary. + mean = mean.to(input_dtype).to(tl.float32) + std = tl.sqrt(variance).to(input_dtype).to(tl.float32) + scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32) + cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32) + sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32) + sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate) + + inner = 0.7978845608028654 * ( + sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate + ) + tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0 + activated = 0.5 * sparse_gate * (1.0 + tanh) + activated = activated.to(input_dtype).to(tl.float32) + tl.store(output_ptr + row * d + offsets, activated * up, mask=mask) + + +def _supports_gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> bool: + del std_multiplier + if ( + x.device.type != "cuda" + or x.dtype not in _SUPPORTED_DTYPES + or not x.is_contiguous() + or x.ndim == 0 + or approximate != "tanh" + ): + return False + last_dim = x.shape[-1] + return ( + last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE + ) + + +@triton_op( + "vllm::gelu_and_mul_sparse_triton", + mutates_args=(), +) +def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor: + d = x.shape[-1] // 2 + output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device) + if output.numel() == 0: + return output + + n_rows = x.numel() // x.shape[-1] + + def grid(meta): + return (n_rows,) + + block_size = triton.next_power_of_2(d) + num_warps = min(max(block_size // 256, 1), 8) + wrap_triton(_gelu_and_mul_sparse_kernel)[grid]( + x, + output, + d=d, + std_multiplier=std_multiplier, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output + + +@ir.ops.gelu_and_mul_sparse.register_impl( + "triton", + supported=HAS_TRITON and current_platform.is_cuda(), + supports_args=_supports_gelu_and_mul_sparse, +) +def gelu_and_mul_sparse_triton( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + assert approximate == "tanh" + return _gelu_and_mul_sparse_triton_op(x, std_multiplier) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 159c521cee2e..f159d2dff8b6 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -8,6 +8,7 @@ import torch.nn as nn import torch.nn.functional as F +from vllm import ir from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.platforms import CpuArchEnum, current_platform @@ -332,24 +333,11 @@ def __init__(self, activation_sparsity: float, approximate: str = "none"): raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.") target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32) normal_dist = torch.distributions.normal.Normal(0, 1) - self.std_multiplier = normal_dist.icdf(target_sparsity_tensor) - - def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor: - """Get % sparse percentile of the Gaussian distribution.""" - # NOTE(rob): for TP>1, we could all-gather to get the means/std. - # But we do not do this because in expectation they are the same - # and in practice the eval scores are good without gathering. - mean = torch.mean(x, dim=-1, keepdim=True) - std = torch.std(x, dim=-1, keepdim=True, unbiased=False) - cutoff_x = mean + std * self.std_multiplier - return nn.functional.relu(x - cutoff_x) + self.std_multiplier = normal_dist.icdf(target_sparsity_tensor).item() def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" - d = x.shape[-1] // 2 - out = self._gaussian_topk(x[..., :d]) - out = F.gelu(out, approximate=self.approximate) - return out * x[..., d:] + return ir.ops.gelu_and_mul_sparse(x, self.std_multiplier, self.approximate) def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 6d9444567a04..98bddbde296d 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -755,7 +755,10 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf rms_norm = ["oink"] + default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["triton", "native"], ) @classmethod diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index a9cf534716f3..b3b7fe8b089c 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1161,7 +1161,10 @@ def get_default_ir_op_priority( rms_norm = default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["native"], ) @classmethod diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 8cb29fbdbd32..6920c55e71b7 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -491,7 +491,7 @@ def get_default_ir_op_priority( using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE default = ["native"] if using_inductor else ["vllm_c", "native"] - return IrOpPriorityConfig.with_default(default) + return IrOpPriorityConfig.with_default(default, gelu_and_mul_sparse=["native"]) @classmethod def device_count(cls) -> int: From 23cfaad49701c497def53552b23317335431f72a Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 14 Sep 2026 15:59:51 +0800 Subject: [PATCH 164/774] [CI] Update entrypoints CI (#56763) Signed-off-by: wang.yuqi --- .buildkite/test-amd.yaml | 12 ++++++------ .buildkite/test_areas/entrypoints.yaml | 14 +++++++------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index efe0a0f8bce3..78130fa3c719 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1067,7 +1067,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD +- label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1080,9 +1080,10 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD +- label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1096,7 +1097,6 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 65 @@ -3217,7 +3217,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD +- label: ":amd: (MI355) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3230,9 +3230,10 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD +- label: ":amd: (MI355) Entrypoints Integration (OpenAI API chat_completion)" # TBD timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3246,7 +3247,6 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 40 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 6ad579f82ad3..00091707d5b4 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -73,9 +73,9 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 1)" +- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API completion)" device: h200_35gb - key: entrypoints-integration-api-server-openai-part-1 + key: entrypoints-integration-api-server-openai-completion timeout_in_minutes: 68 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -84,19 +84,20 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness mirror: amd: - label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" + label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 2)" +- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API chat_completion)" device: h200_35gb - key: entrypoints-integration-api-server-openai-part-2 + key: entrypoints-integration-api-server-openai-chat_completion timeout_in_minutes: 83 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -106,10 +107,9 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py mirror: amd: - label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" + label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" dind: false device: mi300_1 timeout_in_minutes: 70 From b443c1cc4e12171be0b0bbb7cc107919bda4634f Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 14 Sep 2026 16:36:19 +0800 Subject: [PATCH 165/774] [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) (#55737) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/glm5next/nvidia/kda.py | 172 ++++++++++++++++++++++++----- 1 file changed, 144 insertions(+), 28 deletions(-) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 192327cf4a5a..436b9685cb12 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -37,6 +37,7 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.third_party.kda import ( @@ -123,6 +124,32 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: return x.float().sigmoid() +def _resolve_kda_prefill_backend( + backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None +) -> str: + """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on + SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton + ``chunk_kda_with_fused_gate`` path. ``backend`` comes from + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + capability = current_platform.get_device_capability() + supported = ( + current_platform.is_cuda() + and capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 " + "and a bounded KDA gate." + ) + return "flashkda" if supported and backend != "triton" else "triton" + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -291,6 +318,85 @@ def _a_log_weight_loader(param, loaded_weight): # every _forward call (it reads an env-derived flag each time). self._conv_state_dim_first = is_conv_state_dim_first() + additional_config = vllm_config.additional_config + self.kda_prefill_backend = _resolve_kda_prefill_backend( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto", + self.head_dim, + vllm_config.model_config.dtype, + self.kda_lower_bound, + ) + self._flashkda_buffer_specs: ( + tuple[tuple[tuple[int, ...], torch.dtype], ...] | None + ) = None + if self.kda_prefill_backend == "flashkda": + import vllm._flashkda_C # noqa: F401 + + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_seqs = vllm_config.scheduler_config.max_num_seqs + workspace_size = torch.ops._flashkda_C.get_workspace_size( + max_tokens, self.local_num_heads, max_seqs + ) + self._flashkda_buffer_specs = ( + ( + (max_seqs, self.local_num_heads, self.head_dim, self.head_dim), + self.get_state_dtype()[1], + ), + ((workspace_size,), torch.uint8), + # Output buffer for steps that also carry spec-decode tokens: + # the non-spec tokens are then scattered by non_spec_token_indx. + ( + (1, max_tokens, self.local_num_heads, self.head_dim), + vllm_config.model_config.dtype, + ), + ) + + def _flashkda_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g`` + and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the + bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, + matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the + attention output into ``out`` (a workspace buffer when ``None``) and + returns ``(out, final_state)``.""" + assert self._flashkda_buffer_specs is not None + final_state, workspace, workspace_out = ( + current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) + ) + final_state = final_state[: initial_state.shape[0]] + if out is None: + out = workspace_out[:, : q.shape[1]] + # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + self.head_dim**-0.5, + out, + workspace, + self.A_log.view(-1), + self.dt_bias.view(-1, self.head_dim), + self.kda_lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + None, + None, + ) + return out, final_state + def forward( self, hidden_states: torch.Tensor, @@ -542,26 +648,42 @@ def _rearr(x): initial_state = gather_initial_states( recurrent_state, non_spec_state_indices_tensor, has_initial_state ) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=_rearr(q_ns), - k=_rearr(k_ns), - v=_rearr(v_ns), - raw_g=g1_ns, - # Chunk path wants the pre-sigmoided fp32 beta (its kernels - # don't sigmoid); beta_ns is raw bf16 from forward. - beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), - A_log=self.A_log, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) + if self.kda_prefill_backend == "flashkda": + # Non-spec step: write straight into the layer output buffer + # (dense token order, no merge copy). Step with spec-decode + # tokens: write to the workspace buffer and scatter below. + ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens] + core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + out=ns_out, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + raw_g=g1_ns, + # Chunk path wants the pre-sigmoided fp32 beta (its + # kernels don't sigmoid); beta_ns is raw bf16 from forward. + beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), + A_log=self.A_log, + g_bias=self.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) # Init cache scatter_states( recurrent_state, @@ -600,14 +722,8 @@ def _rearr(x): # --- merge spec / non-spec outputs back into token order --- if use_spec and core_attn_out_non_spec is not None: assert core_attn_out_spec is not None - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, - ) - merged.index_copy_(1, spec_token_indx, core_attn_out_spec) - merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[0, :num_actual_tokens] = merged.squeeze(0) + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif use_spec: assert core_attn_out_spec is not None if spec_out is None: From dc89fdfb0ede2918a884faae771a49bbbe7b6cf9 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Mon, 14 Sep 2026 14:46:48 +0530 Subject: [PATCH 166/774] [CPU][Profiler] Group torch profiler tables by input shape when record_shapes is on (#56016) Signed-off-by: Chinmay Kulkarni Co-authored-by: Li, Jiang --- vllm/profiler/wrapper.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index 877a390d6d50..fb8963198635 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -19,6 +19,7 @@ from vllm.config import ProfilerConfig from vllm.config.profiler import _is_uri_path from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -259,9 +260,14 @@ def _build_profiler_table( sort_key: str, row_limit: int | None = None, ) -> str: + group_by_input_shape = ( + current_platform.is_cpu() + and self.profiler_config.torch_profiler_record_shapes + ) + averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape) if row_limit is None: # use profiler default row limit of 100 - return self.profiler.key_averages().table(sort_by=sort_key) - return self.profiler.key_averages().table( + return averages.table(sort_by=sort_key) + return averages.table( sort_by=sort_key, row_limit=row_limit, ) From 79f0be21ff4044d53f0a6f911422b9f52064639c Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Mon, 14 Sep 2026 17:30:28 +0800 Subject: [PATCH 167/774] [Bugfix][CPU] Fix DeepSeek-R1 (FP8 MLA + MoE) correctness on CPU backend (#56773) Signed-off-by: jiang1.li --- vllm/model_executor/kernels/linear/scaled_mm/cpu.py | 2 ++ vllm/model_executor/layers/fused_moe/experts/cpu_moe.py | 1 + 2 files changed, 3 insertions(+) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 4ef1af0a7059..92ea073a1abb 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -261,6 +261,8 @@ def can_implement( return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + return # Skip the base class process (FP8 padding / fnuz normalization) # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. params = self._get_layer_params(layer) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 44a5bda21e79..f7d82c691bf1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -468,6 +468,7 @@ def _supports_routing_method( RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.DeepseekV4, + RoutingMethodType.DeepSeekV3, ] @staticmethod From 238cb2b191630e8b02965adc772e3236d26936fe Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 14 Sep 2026 17:36:39 +0800 Subject: [PATCH 168/774] [Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat (#55738) Signed-off-by: Jared Wen Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_backends.py | 20 +++++++++++ .../v1/attention/test_mla_prefill_selector.py | 34 +++++++++++++++++++ .../v1/attention/test_sparse_mla_backends.py | 29 ++++++++++++++++ .../layers/attention/mla_attention.py | 4 +++ .../layers/attention/sparse_mla_attention.py | 2 ++ .../backends/mla/prefill/flash_attn.py | 7 ++++ 6 files changed, 96 insertions(+) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index ee5737a99573..f724949f54d8 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.attention import mla_attention as mla_attention_module from vllm.model_executor.layers.attention.mla_attention import ( MLAAttention, + MLACommonBaseImpl, QueryLenSupport, _DecodeConcatQuantFP8, _use_masked_mha, @@ -166,6 +167,25 @@ def test_glm5_flashinfer_masked_mha_routing( ) +@pytest.mark.parametrize("qk_rope_head_dim", [64, 0], ids=["rope", "nope"]) +def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim): + """The K concat used by the MLA prefill context loop must equal torch.cat of + k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself + instead of allocating and copying.""" + torch.manual_seed(0) + num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256 + k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, 1, qk_rope_head_dim, dtype=torch.bfloat16) + impl = SimpleNamespace(_use_flashinfer_concat_mla_k=False) + + k = MLACommonBaseImpl._concat_k_nope_k_pe(impl, k_nope, k_pe) + + expected = torch.cat([k_nope, k_pe.expand(-1, num_heads, -1)], dim=-1) + assert k.shape == (num_tokens, num_heads, qk_nope_head_dim + qk_rope_head_dim) + torch.testing.assert_close(k, expected, rtol=0, atol=0) + assert (k.data_ptr() == k_nope.data_ptr()) == (qk_rope_head_dim == 0) + + def test_masked_mha_routing_is_dimension_specific(): assert _use_masked_mha( backend_name="FLASHMLA_SPARSE", diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d5856ab6bdb8..098eddfbf6eb 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -313,6 +313,40 @@ def test_backend_supported_dimension_validation(self): ) assert invalid_reasons == [] + def test_flash_attn_accepts_glm53_flash_nope_dimensions(self): + """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a + RoPE-carrying 320-wide query does not.""" + try: + from vllm.v1.attention.backends.mla.prefill.flash_attn import ( + FlashAttnPrefillBackend, + ) + except ImportError: + pytest.skip("MLA prefill backend not available") + return + + capability = DeviceCapability(major=10, minor=0) + + def validate(qk_rope_head_dim: int) -> list[str]: + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=256, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=256, + ), + ) + with patch.object( + FlashAttnPrefillBackend, "is_available", return_value=True + ): + return FlashAttnPrefillBackend.validate_configuration( + capability, selector_config + ) + + assert validate(qk_rope_head_dim=0) == [] + invalid_reasons = validate(qk_rope_head_dim=64) + assert len(invalid_reasons) == 1 + assert "supported MLA dimensions" in invalid_reasons[0] + @pytest.mark.skipif( not current_platform.is_cuda_alike(), diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 0766584f5451..2ca839d59c50 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -32,6 +32,7 @@ GLOBAL_TOPK_MASK_MAX_BYTES, SparseMLACommonImpl, SparseMLAPrefillMetadata, + _is_masked_mha_available, _masked_mha_workspace_fits, _topk_mask_shape, _use_dense_mha_prefill, @@ -1165,6 +1166,33 @@ def test_masked_mha_workspace_guards_long_routing_policy( ) +@pytest.mark.parametrize( + ("model_dims", "kv_cache_dtype", "fa_version", "expected"), + [ + pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"), + pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"), + pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"), + pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"), + pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"), + pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"), + pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"), + ], +) +def test_is_masked_mha_available_model_dims( + monkeypatch, model_dims, kv_cache_dtype, fa_version, expected +): + """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2, + GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an + unquantized KV cache, nothing else.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as mod + + monkeypatch.setattr( + mod.current_platform, "is_device_capability_family", lambda family: True + ) + monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version) + assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected + + def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): """Request count and context chunk length are independent multipliers.""" base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048) @@ -1198,6 +1226,7 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): [ pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"), pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"), + pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"), ], ) def test_sparse_backend_prefill_correctness( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1c2e5b396263..874c43f8cd44 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2871,6 +2871,10 @@ def _concat_k_nope_k_pe( Returns: Tensor of shape [..., nope_dim + pe_dim] """ + if k_pe.shape[-1] == 0: + # NoPE MLA: nothing to append, so no copy either. + return k_nope + k = torch.empty( (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]), dtype=k_nope.dtype, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 79ca61203255..aed5afdcf905 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -111,6 +111,8 @@ def _is_masked_mha_available( if model_dims not in ( (128, 512, 128, 64, 128), (64, 512, 192, 64, 256), + # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel. + (64, 512, 256, 0, 256), ): return False qk_head_dim = qk_nope_head_dim + qk_rope_head_dim diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 3e0d544195e1..3af623b1f13b 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -327,6 +327,13 @@ def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool: qk_rope_head_dim=64, v_head_dim=128, ), + # GLM5Next NoPE layout: qk_head_dim 256 + 0 and v_head_dim 256 run + # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout. + MLADimensions( + qk_nope_head_dim=256, + qk_rope_head_dim=0, + v_head_dim=256, + ), ] def __init__( From dc2e8f1157173d0f0fc09712f6897e9f6eee52a0 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Mon, 14 Sep 2026 11:40:25 +0200 Subject: [PATCH 169/774] [ROCm][Perf] W4A16: keep skinny GEMM zero-points packed 4-bit (#54965) Signed-off-by: Matthias Gehre --- csrc/rocm/skinny_gemms_int4.cu | 59 ++++++++++++----- csrc/rocm/torch_bindings.cpp | 3 +- .../quantization/test_rdna_hybrid_w4a16.py | 65 ++++++++++++++++--- .../mixed_precision/rdna_hybrid_w4a16.py | 35 +++++----- 4 files changed, 117 insertions(+), 45 deletions(-) diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 8b7077d14ed0..a20ca41d413e 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) { return min(a, b); } +// Reads row `row` of packed zero points [M/8, num_groups] uint32. +// +// The parameters must be signed and the indexing must use shift/mask: with +// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8) +// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239. +__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed, + const int row, const int group, + const int num_groups) { + const uint32_t w = zp_packed[(row >> 3) * num_groups + group]; + return (w >> (4 * (row & 7))) & 0xFu; +} + // W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations // Targets the "sml" case where activations fit in LDS. // A_CHUNK: number of K-elements processed per thread per step. @@ -116,7 +128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { constexpr int max_lds_len = LDS_SIZE / 2; @@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if constexpr (!std::is_same_v) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; - scalar_t zp = zero_points[(m + y) * num_groups + group_idx]; + scalar_t zp = __float2s(static_cast( + zp_nibble(zero_points, m + y, group_idx, num_groups))); #pragma unroll for (uint32_t b = 0; b < A_CHUNK; b++) { cvtB.h[b] = cvtB.h[b] - zp; @@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if constexpr (HAS_ZERO_POINTS) { uint32_t group_idx_zp = k_ / GROUP_SIZE; - float zp_f = __s2float( - zero_points[(m + y) * num_groups + group_idx_zp]); + float zp_f = static_cast( + zp_nibble(zero_points, m + y, group_idx_zp, num_groups)); partial -= (128.0f + zp_f) * act_sum; } else { partial -= 136.0f * act_sum; @@ -335,7 +348,7 @@ __global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { @@ -352,7 +365,7 @@ template ) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; - scalar_t zp = zero_points[(m + y) * num_groups + group_idx]; + scalar_t zp = __float2s(static_cast( + zp_nibble(zero_points, m + y, group_idx, num_groups))); #pragma unroll for (uint32_t b = 0; b < A_CHUNK; b++) { cvtB.h[b] = cvtB.h[b] - zp; @@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if constexpr (HAS_ZERO_POINTS) { uint32_t group_idx_zp = k_ / GROUP_SIZE; - float zp_f = __s2float( - zero_points[(m + y) * num_groups + group_idx_zp]); + float zp_f = static_cast( + zp_nibble(zero_points, m + y, group_idx_zp, num_groups)); partial -= (128.0f + zp_f) * act_sum; } else { partial -= 136.0f * act_sum; @@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { @@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) { // in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32) // in_b: activations [N, K] (fp16/bf16) // in_scale: group scales [M, K/group_size] (fp16/bf16) -// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16) +// in_zero_points: optional raw zero points [M/8, K/group_size] (int32, +// 8 uint4 per word along dim 0) // If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric). // If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8). // group_size: 32, 64, or 128 @@ -664,14 +679,22 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups, "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]"); if (in_zero_points.has_value()) { - TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(), - "Zero points dtype must match activation dtype"); + // The kernel reads the words as uint32, so either signedness is accepted. + TORCH_CHECK(in_zero_points->dtype() == at::kInt || + in_zero_points->dtype() == at::kUInt32, + "Zero points must be int32 or uint32 (packed 8x uint4 along " + "dim 0), got ", + in_zero_points->dtype()); TORCH_CHECK(in_zero_points->dim() == 2, - "Zero points must be 2D [M, K/group_size], got shape ", + "Zero points must be 2D [M/8, K/group_size], got shape ", in_zero_points->sizes()); - TORCH_CHECK(in_zero_points->size(0) == M_in && + TORCH_CHECK(M_in % 8 == 0, + "M must be divisible by 8 for packed zero points, got ", M_in); + TORCH_CHECK(in_zero_points->is_contiguous(), + "Zero points must be contiguous"); + TORCH_CHECK(in_zero_points->size(0) == M_in / 8 && in_zero_points->size(1) == num_groups, - "Zero points must be [M, K/group_size] = [", M_in, ", ", + "Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ", num_groups, "] but got [", in_zero_points->size(0), ", ", in_zero_points->size(1), "]"); } @@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, const fptype* aptr = reinterpret_cast(in_b.data_ptr()); const fptype* sptr = reinterpret_cast(in_scale.data_ptr()); - const fptype* zpptr = + const uint32_t* zpptr = in_zero_points.has_value() - ? reinterpret_cast(in_zero_points->data_ptr()) + ? reinterpret_cast(in_zero_points->data_ptr()) : nullptr; const fptype* biasptr = (in_bias.has_value() && in_bias->numel() > 0) diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index a8114ce1969a..14d1822707ea 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK); // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales, - // optional zero points for asymmetric quantization + // optional zero points [M/8, K/group_size] int32 for asymmetric + // quantization rocm_ops.def( "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, " "Tensor? in_zero_points, Tensor? in_bias, int CuCount, " diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index 562d1ef2ef71..fe3db8dab360 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -37,6 +37,20 @@ # --------------------------------------------------------------------------- +def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor: + """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32. + + Row n's nibble lands in word[n//8] at bits 4*(n%8). + """ + assert zp_nkg.dtype == torch.int32 + N, G = zp_nkg.shape + assert N % 8 == 0 + shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None] + return torch.sum( + (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32 + ).contiguous() + + def _rdna_hybrid_w4a16_reference( x_mk: torch.Tensor, w_int4_nk: torch.Tensor, @@ -50,8 +64,8 @@ def _rdna_hybrid_w4a16_reference( x_mk: [M, K] fp16/bf16 w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15] scales_nkg: [N, K//G] fp16/bf16 - zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype), - or None for symmetric (uint4b8, dequant subtracts 8) + zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for + symmetric (uint4b8, dequant subtracts 8) """ G = group_size N, K = w_int4_nk.shape @@ -110,13 +124,16 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) 0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32) ).to(dtype) - # Optional raw zero points [N, K//G] in act dtype. + # Optional raw zero points, [N, K//G] int32 for the reference and packed + # along N for the op. if has_zp: zp_nkg = torch.randint( 0, 16, (N, K // group_size), device=device, dtype=torch.int32 - ).to(dtype) + ) + w_zp = _pack_zp_rows_for_kernel(zp_nkg) else: zp_nkg = None + w_zp = None from vllm.utils.platform_utils import num_compute_units @@ -124,7 +141,7 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) x_mk, w_q, scales_nkg, - zp_nkg, + w_zp, None, # bias num_compute_units(), group_size, @@ -332,7 +349,7 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES) def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init): - """uint4 (asymmetric): zero points unpacked to raw values in act dtype.""" + """uint4 (asymmetric): zero points are packed [N//8, K//G] int32.""" if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -374,10 +391,11 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in ) kernel.process_weights_after_loading(layer) - # Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15]. - assert layer.weight_zero_point.dtype == torch.float16 - assert tuple(layer.weight_zero_point.shape) == (N, K // G) - expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G] + # Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8] + # bits 4*(n%8). + assert layer.weight_zero_point.dtype == torch.int32 + assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G) + expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous()) torch.testing.assert_close(layer.weight_zero_point, expected_zp) # Quantized weights match symmetric path's layout regardless of zp. @@ -518,3 +536,30 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G): ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.skipif( + not hasattr(torch.ops, "_rocm_C") + or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"), + reason="wvSplitK_int4_g not built", +) +def test_wvsplitk_int4_g_rejects_unpacked_zero_points(): + """Act-dtype zero points raise instead of being read as packed words. + + Both layouts are 2D with compatible extents, so only the dtype separates + them; misreading one as the other returns wrong numbers silently. + """ + import vllm._custom_ops as ops + from vllm.utils.platform_utils import num_compute_units + + K, N, G, M = 256, 64, 128, 1 + a = torch.randn((M, K), device=device, dtype=torch.float16) + w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view( + torch.int8 + ) + scales = torch.rand((N, K // G), device=device, dtype=torch.float16) + zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16) + + with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"): + ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 9964c891b6b3..eb36308bccab 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -77,7 +77,7 @@ def _triton_w4a16_skinny_fmt_kernel( a_ptr, # [M, K] fp16/bf16 activations b_ptr, # [N, K//8] int32 packed (ExLlama shuffle, K is packed dim) scales_ptr, # [N, K//G] fp16/bf16 scales (skinny layout) - zp_ptr, # [N, K//G] fp16/bf16 raw zero-points (when HAS_ZP=True) + zp_ptr, # [N//8, K//G] int32 zero-points (when HAS_ZP=True) c_ptr, # [M, N] fp16/bf16 output # Dimensions M, @@ -103,8 +103,9 @@ def _triton_w4a16_skinny_fmt_kernel( | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28) Scales are [N, K//G] (skinny layout, NOT transposed). - When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G] - and subtracted directly: (nibble - zp_raw) * scale. + When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw + zero-point at word[n//8] bits 4*(n%8), and dequant is + (nibble - zp_raw) * scale. When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric). """ pid_m = tl.program_id(0) @@ -148,9 +149,10 @@ def _triton_w4a16_skinny_fmt_kernel( scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0) if HAS_ZP: - zp_ptrs = zp_ptr + offs_n * num_groups + group_idx - zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0) - b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None] + zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + group_idx + zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0) + zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF + b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None] else: b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None] @@ -192,7 +194,7 @@ def triton_w4a16_skinny_fmt_gemm( scales: torch.Tensor, # [N, K//G] fp16/bf16 group_size: int, zp_bias: int = 8, - zp: torch.Tensor | None = None, # [N, K//G] per-group zero-points + zp: torch.Tensor | None = None, # [N//8, K//G] int32 zero-points ) -> torch.Tensor: """ Fused W4A16 GEMM reading from skinny weight format [N, K//8]. @@ -203,8 +205,8 @@ def triton_w4a16_skinny_fmt_gemm( scales: Per-group scales [N, K//G], same dtype as a. group_size: Quantization group size (resolved from -1 to K by caller). zp_bias: Constant zero bias (default 8 for unsigned int4). - zp: Raw per-group zero-points [N, K//G] (asymmetric), - stored as zp_raw in activation dtype. When provided, + zp: Raw per-group zero-points [N//8, K//G] int32, row n at + word[n//8] bits 4*(n%8) (asymmetric). When provided, dequant is (nibble - zp_raw) * scale. Returns: @@ -225,8 +227,9 @@ def triton_w4a16_skinny_fmt_gemm( ) if zp is not None: assert zp.is_contiguous(), "Zero-points must be contiguous" - assert zp.shape == (N, num_groups), ( - f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})" + assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}" + assert zp.shape == (N // 8, num_groups), ( + f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})" ) has_zp = zp is not None @@ -381,7 +384,10 @@ def _rdna_hybrid_w4a16_apply_impl( cu_count: int, group_size: int, ) -> torch.Tensor: - """Dispatch between skinny GEMM and Triton based on batch size M.""" + """Dispatch between skinny GEMM and Triton based on batch size M. + + ``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric. + """ import vllm._custom_ops as ops M = x_2d.shape[0] @@ -520,10 +526,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: assert self.w_zp_name is not None w_zp_raw = getattr(layer, self.w_zp_name) permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0) - zp_unpacked = unpack_quantized_values_into_int32( - w_zp_raw.data, c.weight_type, packed_dim=0 - ) - w_zp = zp_unpacked.to(c.act_type).contiguous() + w_zp = w_zp_raw.data.contiguous() self._transform_param(layer, self.w_zp_name, lambda x: w_zp) self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny) From 5236bef721223a45d52b3649530e88b5e824e0d9 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 02:52:45 -0700 Subject: [PATCH 170/774] [CI] Keep OTel bytecode out of mounted checkouts (#56766) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-otel/ci_otel.sh | 16 ++++++---- .../scripts/ci-otel/tests/test_ci_otel.py | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index 78b18207ad27..c9897e732e51 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -26,7 +26,8 @@ if [ -z "${CI_INFRA_OTEL_RUNTIME_DIR}" ] || [ -z "${_CI_INFRA_OTEL_PYTHON}" ] || [ ! -f "${CI_INFRA_OTEL_DIR}/ci_otel.py" ] || ! mkdir -p "${CI_INFRA_OTEL_SPOOL_DIR}" || - ! PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \ + ! PYTHONDONTWRITEBYTECODE=1 \ + PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \ "${_CI_INFRA_OTEL_PYTHON}" -c "import ci_otel" >/dev/null 2>&1; then echo "vLLM CI OTel: tracing unavailable; test command will run normally" >&2 || : if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then @@ -58,9 +59,9 @@ fi _ci_otel_python() { if command -v timeout >/dev/null 2>&1; then - timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@" + PYTHONDONTWRITEBYTECODE=1 timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@" else - "${_CI_INFRA_OTEL_PYTHON}" "$@" + PYTHONDONTWRITEBYTECODE=1 "${_CI_INFRA_OTEL_PYTHON}" "$@" fi } @@ -86,7 +87,8 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" /dev/null 2>&1; then - timeout 4s "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : + PYTHONDONTWRITEBYTECODE=1 timeout 4s \ + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : else - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : + PYTHONDONTWRITEBYTECODE=1 \ + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : fi if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then rm -rf -- "${CI_INFRA_OTEL_RUNTIME_DIR}" || : diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ba54f9f27349..ae51010f9eeb 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -6,6 +6,7 @@ import json import os import shlex +import shutil import subprocess import sys import threading @@ -396,6 +397,35 @@ def test_shell_wrapper_preserves_failure_status(tmp_path): assert result.returncode == 1 +def test_shell_wrapper_does_not_write_bytecode_to_helper_tree(tmp_path): + helper_dir = tmp_path / "ci-otel" + shutil.copytree( + SCRIPTS_DIR, + helper_dir, + ignore=shutil.ignore_patterns("__pycache__"), + ) + shell = ( + f'. "{helper_dir / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "CI_INFRA_OTEL_DIR": str(helper_dir), + "CI_INFRA_GPU_SAMPLING": "0", + "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), + }, + ) + + assert result.returncode == 0, result.stderr + assert list(helper_dir.rglob("*.pyc")) == [] + + def test_ci_otel_run_records_command_and_preserves_status(tmp_path): shell = ( f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' From 2c7ee87223f0c94f614a0b8600f6f193e6e4388f Mon Sep 17 00:00:00 2001 From: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:55:39 +0800 Subject: [PATCH 171/774] [Bugfix][Mooncake] Report failed remote KV loads to the scheduler (#50984) Signed-off-by: liuqihao Co-authored-by: liuqihao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/mooncake/mooncake_connector.py | 80 +++++++++++++++---- 1 file changed, 66 insertions(+), 14 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 4009af7009ea..4451f0566da1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import logging +import queue import threading import time from collections import defaultdict +from collections.abc import Collection from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from enum import IntEnum @@ -433,6 +435,9 @@ class PullReqMeta: expire_time: float = float("inf") # Designed for one D pairing to multiple P pull_tasks_count: int = 0 + # Set once any worker reports a failure, so a success from another worker + # for the same request is not counted afterwards. + failed: bool = False @dataclass @@ -582,6 +587,11 @@ def get_finished( assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_block_ids_with_load_errors(self) -> set[int]: + """Get block IDs whose remote KV load failed.""" + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) @@ -1034,6 +1044,8 @@ def __init__( self.finished_sending_reqs: set[ReqId] = set() self.finished_recving_reqs: set[ReqId] = set() + # Written from the receiver loop, drained from the worker thread. + self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue() self.xfer_stats = MooncakeKVConnectorStats() @@ -1907,12 +1919,11 @@ async def receive_kv_from_single_worker( ret_msg = await sock.recv() response = self._xfer_resp_decoder.decode(ret_msg) if response.status == MooncakeXferResponseStatus.ERROR: - logger.error( - "Error happens during transferring kvcache for %s: %s", + self._handle_failed_recv( + pull_metas, req_ids, - response.err_msg, + response.err_msg or "transfer error", ) - self.xfer_stats.record_failed_recv() return self.process_pulling_result(response, pull_metas) if response.status == MooncakeXferResponseStatus.FINISH: @@ -1920,10 +1931,50 @@ async def receive_kv_from_single_worker( except zmq.ContextTerminated: logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.") except Exception as e: - logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) - self.xfer_stats.record_failed_recv() + self._handle_failed_recv(pull_metas, req_ids, f"transfer failed: {e}") return + def _handle_failed_recv( + self, + pull_metas: dict[ReqId, PullReqMeta], + req_ids: Collection[ReqId], + reason: str, + ) -> None: + """Report a failed remote KV load so the scheduler can fail or recompute it.""" + failed: list[ReqId] = [] + for req_id in req_ids: + pull_meta = pull_metas.get(req_id) + if pull_meta is None or pull_meta.failed: + continue + pull_meta.failed = True + failed.append(req_id) + self.xfer_stats.record_failed_recv() + + invalid = {b for group in pull_meta.local_block_ids for b in group} + if not invalid: + # A pull with no local blocks only asks P to release its blocks + # for a request that never reached the scheduler (see + # AsyncLLM.notify_kv_transfer_request_rejected, which submits an + # abort_immediately request just to run request_finished). No D + # request is waiting on a load, and reporting one here would trip + # the scheduler's `assert req_id in self.requests`. + continue + self._invalid_block_ids.put(invalid) + self.finished_recving_reqs.add(pull_meta.d_req_id) + + if failed: + logger.error("pulling kv_caches for %s failed: %s", failed, reason) + + def get_block_ids_with_load_errors(self) -> set[int]: + """Drain the blocks whose remote KV load failed since the last call.""" + result: set[int] = set() + while True: + try: + result.update(self._invalid_block_ids.get_nowait()) + except queue.Empty: + break + return result + def process_pulling_result( self, response: MooncakeXferResponse, @@ -1933,6 +1984,8 @@ def process_pulling_result( for req_id in ok_reqs: pull_meta = pull_metas[req_id] + if pull_meta.failed: + continue # No race because we are in async loop. pull_meta.pull_tasks_count -= 1 if pull_meta.pull_tasks_count == 0: @@ -1942,10 +1995,8 @@ def process_pulling_result( logger.debug("pulling kv_caches for %s finished", ok_reqs) if response.err_reqs: - logger.error( - "pulling kv_caches for %s failed: %s", - response.err_reqs, - response.err_msg, + self._handle_failed_recv( + pull_metas, response.err_reqs, response.err_msg or "unknown error" ) async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): @@ -2023,10 +2074,11 @@ async def handle_new_engine_id( await self._pending_bootstrap_queries[remote_bootstrap_addr].wait() if remote_engine_id not in self._remote_agents: - logger.error( - "Failed to find remote engine_id %s from bootstrap server %s", - remote_engine_id, - remote_bootstrap_addr, + self._handle_failed_recv( + pull_metas, + list(pull_metas), + f"remote engine_id {remote_engine_id} not found from bootstrap " + f"server {remote_bootstrap_addr}", ) return From 934b1fcbcbbac6e68da7a4246928ac797fe54690 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 14 Sep 2026 11:02:14 +0100 Subject: [PATCH 172/774] [Frontend] Replace `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` with `--enable-scale-out` (#55176) Signed-off-by: Martin Hickey --- docs/serving/online_serving/README.md | 7 +-- docs/serving/online_serving/derenderer.md | 3 +- docs/serving/online_serving/renderer.md | 6 +-- docs/usage/security.md | 18 +++---- examples/scale_out/example_mm_serve.py | 6 +-- examples/scale_out/token_generation_client.py | 2 +- rust/src/cmd/src/cli.rs | 10 ++++ rust/src/cmd/src/cli/tests.rs | 6 +++ rust/src/server/src/config.rs | 2 + rust/src/server/src/routes.rs | 28 ++-------- rust/src/server/src/routes/tests.rs | 48 ++++++++++------- tests/entrypoints/cohere/test_api_router.py | 23 ++++---- tests/entrypoints/cohere/test_chat_v2.py | 3 +- .../derender/test_derender_parity.py | 7 +-- .../render/test_render_multimodal.py | 9 +--- tests/entrypoints/scale_out/test_factories.py | 53 ++++++------------- .../test_return_routed_experts.py | 7 +-- .../test_serving_multimodal_tokens.py | 7 +-- .../token_in_token_out/test_serving_tokens.py | 7 +-- tests/test_envs.py | 35 ------------ vllm/entrypoints/cohere/api_router.py | 10 ++-- vllm/entrypoints/launchers/cli_args.py | 7 +++ vllm/entrypoints/scale_out/factories.py | 24 +++------ vllm/envs.py | 17 ------ 24 files changed, 130 insertions(+), 215 deletions(-) diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 5c93ca17877c..1f5d4bdb6a93 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -121,12 +121,7 @@ For further details on profiling vLLM, please refer to [this page](../../contrib ## Scale-Out APIs -Scale-out APIs are disabled by default on `vllm serve`. The environment -variable accepts only `0` or `1`; set `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1` to -register the endpoints below. The -dedicated `vllm launch render` and `vllm serve --tokens-only` modes are explicit -opt-ins and enable their required endpoints when the variable is unset; an -explicit value of `0` is rejected for those modes. +Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` to register the endpoints below. The dedicated `vllm launch render` and `vllm serve --tokens-only` modes always register their required endpoints regardless of `--enable-scale-out`. ### Tokens IN <> Tokens OUT APIs diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index d536498f5c60..298510062aaa 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -60,8 +60,7 @@ Launch the two servers first: ```bash vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve \ - meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200 +vllm serve meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200 ``` ```python diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 1f5e893570ef..9be23142c49c 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -7,16 +7,14 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and - Tokens-in / tokens-out engine: Make the engine a pure token-in / token-out service, decoupled from request preprocessing. The dedicated `vllm launch render` server always exposes the `/render` and -`/derender` endpoints when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` is unset or set to -`1`. An explicit value of `0` conflicts with the renderer command and is -rejected at startup. +`/derender` endpoints. Scale-out endpoints, including `/render`, `/derender`, and `/inference/v1/generate`, are disabled by default on a standard inference server. To expose them with `vllm serve`, opt in explicitly: ```bash -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve +vllm serve --enable-scale-out ``` ## API Reference diff --git a/docs/usage/security.md b/docs/usage/security.md index 559f58702690..3c7f1e3f9ba8 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -158,16 +158,16 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/models` - List available models - `/v1/chat/completions` - Chat completions - `/v1/chat/completions/batch` - Batch chat completions -- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) -- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) +- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/completions` - Text completions -- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) -- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) +- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/embeddings` - Generate embeddings - `/v1/audio/transcriptions` - Audio transcription - `/v1/audio/translations` - Audio translation - `/v1/messages` - Anthropic-compatible messages API -- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/messages/count_tokens` - Count tokens for Anthropic messages - `/v1/responses` - Create a response - `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) @@ -177,8 +177,8 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/rerank` - Reranking API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) -- `/inference/v1/generate` - Generate completions (available when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or with `--tokens-only` unless explicitly disabled) -- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires both `VLLM_ENABLE_COHERE_API=1` and `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`) +- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) +- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`) - `/v2/embed` - Cohere Embed API - `/v2/rerank` - Cohere Rerank API @@ -206,11 +206,11 @@ The following endpoints **do not require authentication** even when `--api-key` - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF - `/update_weights` - Update model weights (can alter model behavior) - `/get_world_size` - Get distributed world size -- `/abort_requests` - Abort in-flight requests (available with `--tokens-only` unless `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0`) +- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`) **Utility endpoints:** -- `/tokenize` - Tokenize text (not disabled by `VLLM_ENABLE_SCALE_OUT_ENDPOINTS`) +- `/tokenize` - Tokenize text (not gated by `--enable-scale-out`) - `/detokenize` - Detokenize tokens - `/health` - Health check - `/ping` - SageMaker health check diff --git a/examples/scale_out/example_mm_serve.py b/examples/scale_out/example_mm_serve.py index 157a32d3be8b..558a4f930247 100644 --- a/examples/scale_out/example_mm_serve.py +++ b/examples/scale_out/example_mm_serve.py @@ -13,9 +13,9 @@ Launch the server first: - VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ - vllm serve Qwen/Qwen3-VL-2B-Instruct \ - --dtype bfloat16 --max-model-len 4096 --enforce-eager + vllm serve Qwen/Qwen3-VL-2B-Instruct \ + --dtype bfloat16 --max-model-len 4096 --enforce-eager \ + --enable-scale-out Then run this script: diff --git a/examples/scale_out/token_generation_client.py b/examples/scale_out/token_generation_client.py index 922018f43d16..a4792a68f9a3 100644 --- a/examples/scale_out/token_generation_client.py +++ b/examples/scale_out/token_generation_client.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Call a token-in/token-out server started with: -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve Qwen/Qwen3-0.6B +vllm serve Qwen/Qwen3-0.6B --enable-scale-out """ import httpx diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index f7fcf30022b9..70b28018a3f1 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -329,6 +329,15 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub enable_request_id_headers: bool, + /// Register the scale-out `/inference/v1/generate` endpoint. + #[arg( + long, + default_missing_value = "true", + num_args = 0..=1 + )] + #[serde(default)] + pub enable_scale_out: bool, + /// If provided, the server will require one of these keys to be presented /// in the Authorization header. #[educe(Debug(ignore))] @@ -577,6 +586,7 @@ impl SharedRuntimeArgs { enable_log_requests: self.enable_log_requests, enable_prompt_tokens_details: self.enable_prompt_tokens_details, enable_request_id_headers: self.enable_request_id_headers, + enable_scale_out: self.enable_scale_out, } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 884da135d719..411d2562959c 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -195,6 +195,7 @@ fn serve_args_forward_python_flags_with_separator() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1030,6 +1031,7 @@ fn frontend_args_accept_json() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1629,6 +1631,7 @@ fn serve_args_accept_handshake_aliases() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1780,6 +1783,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ @@ -1868,6 +1872,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ @@ -1977,6 +1982,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 4928cddf31e9..0f0deac86d2c 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -55,6 +55,8 @@ pub struct ApiServerOptions { pub enable_prompt_tokens_details: bool, /// When `true`, set `X-Request-Id` on every HTTP response. pub enable_request_id_headers: bool, + /// When `true`, register the scale-out `/inference/v1/generate` route. + pub enable_scale_out: bool, } /// CORS settings mirroring Python's `CORSMiddleware`; the default is permissive. diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index e7e91938bbda..41e49878820b 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -46,33 +46,14 @@ fn runtime_lora_updating_enabled() -> bool { .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true")) } -fn parse_scale_out_endpoints_flag(value: Option<&str>) -> Result { - let Some(value) = value else { - return Ok(false); - }; - if value.is_empty() { - return Ok(false); - } - - match value.trim() { - "0" => Ok(false), - "1" => Ok(true), - _ => Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string()), - } -} - -fn scale_out_endpoints_enabled() -> bool { - let value = std::env::var("VLLM_ENABLE_SCALE_OUT_ENDPOINTS").ok(); - parse_scale_out_endpoints_flag(value.as_deref()).unwrap_or_else(|message| panic!("{message}")) -} - /// Build the minimal OpenAI-compatible router for one configured model. pub fn build_router(state: Arc) -> Router { + let scale_out_endpoints_enabled = state.api_server_options.enable_scale_out; build_router_with_options( state, server_dev_mode_enabled(), runtime_lora_updating_enabled(), - scale_out_endpoints_enabled(), + scale_out_endpoints_enabled, ) } @@ -121,10 +102,7 @@ fn build_router_with_options( if scale_out_endpoints_enabled { router = router.route("/inference/v1/generate", post(inference::generate)); } else { - info!( - "scale-out endpoints are disabled; set \ - VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them" - ); + info!("scale-out endpoints are disabled; pass --enable-scale-out to enable them"); } if runtime_lora_updating_enabled { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 2851ec1e88b8..5ee2b51f3f5c 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -54,8 +54,7 @@ use zeromq::{DealerSocket, PushSocket, ZmqMessage}; use super::{ build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora, - build_router_with_scale_out_endpoints, parse_scale_out_endpoints_flag, - render::build_router as build_render_router, + build_router_with_scale_out_endpoints, render::build_router as build_render_router, }; use crate::config::{ApiServerOptions, CorsConfig, LoraModulePath}; use crate::lora::LoadLoraError; @@ -731,23 +730,36 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router { ) } -#[test] -fn scale_out_flag_accepts_unset_empty_zero_one_and_whitespace() { - assert_eq!(parse_scale_out_endpoints_flag(None), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("")), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("0")), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("1")), Ok(true)); - assert_eq!(parse_scale_out_endpoints_flag(Some(" 1 ")), Ok(true)); -} +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn scale_out_generate_route_follows_api_server_options() { + let (chat, _engine_task) = test_models_with_engine_outputs_and_backend( + b"engine-scale-out-api-server-options", + default_stream_output_specs(), + Arc::new(FakeChatBackend::new()), + ) + .await; + let state = Arc::new( + AppState::new(vec!["model".to_string()], chat).with_api_server_options(ApiServerOptions { + enable_scale_out: true, + ..Default::default() + }), + ); + let mut app = build_router(state); -#[test] -fn scale_out_flag_rejects_values_other_than_zero_or_one() { - for value in ["-1", "2", "01", "+1", "invalid", " "] { - assert_eq!( - parse_scale_out_endpoints_flag(Some(value)), - Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string()) - ); - } + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from("{}")) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_ne!(response.status(), StatusCode::NOT_FOUND); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index a90e93b6b165..94324dbed76e 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -58,7 +58,6 @@ def _enable_cohere_api(monkeypatch): flag inside the test body. """ monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1") - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") # ---------------------------------------------------------------------- @@ -128,15 +127,19 @@ def _generate_request() -> GenerateRequest: ) -def _build_app(handler: _Handler | None) -> FastAPI: +def _build_app(handler: _Handler | None, *, enable_scale_out: bool = True) -> FastAPI: app = FastAPI() + app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False) attach_router(app) app.state.cohere_serving_chat_v2 = handler return app -def _build_render_app(chat_handler, render_handler) -> FastAPI: +def _build_render_app( + chat_handler, render_handler, *, enable_scale_out: bool = True +) -> FastAPI: app = FastAPI() + app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False) attach_router(app) app.state.cohere_serving_chat_v2 = chat_handler app.state.serving_render = render_handler @@ -151,11 +154,13 @@ def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI: into the ``CohereError`` wire shape. """ app = FastAPI() - attach_router(app) - app.state.cohere_serving_chat_v2 = handler # ``validation_exception_handler`` reads ``req.app.state.args``; the # real cli builds this via argparse. - app.state.args = Namespace(log_error_stack=False) + app.state.args = Namespace( + log_error_stack=False, enable_scale_out=True, tokens_only=False + ) + attach_router(app) + app.state.cohere_serving_chat_v2 = handler app.exception_handler(RequestValidationError)(validation_exception_handler) app.exception_handler(HTTPException)(http_exception_handler) app.exception_handler(VLLMError)(vllm_error_handler) @@ -402,10 +407,8 @@ def test_route_registered(self): paths = [getattr(r, "path", None) for r in app.routes] assert "/cohere/v2/chat/render" in paths - def test_route_not_registered_when_flag_unset(self, monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - app = FastAPI() - attach_router(app) + def test_route_not_registered_when_flag_unset(self): + app = _build_app(handler=None, enable_scale_out=False) paths = [getattr(r, "path", None) for r in app.routes] assert "/cohere/v2/chat/render" not in paths assert "/cohere/v2/chat" in paths diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py index 36acecb6224d..4f07d3d87355 100644 --- a/tests/entrypoints/cohere/test_chat_v2.py +++ b/tests/entrypoints/cohere/test_chat_v2.py @@ -57,6 +57,8 @@ def server(): # the conversation to surface a thinking block (SmolLM2 doesn't # emit Cohere-style reasoning tokens out of the box). "--no-cohere-is-reasoning-model", + # Registers /cohere/v2/chat/render alongside /cohere/v2/chat. + "--enable-scale-out", ] # Cap the CPU KV-cache pool the vLLM CPU backend reserves at @@ -71,7 +73,6 @@ def server(): env_dict = { "VLLM_CPU_KVCACHE_SPACE": "1", "VLLM_ENABLE_COHERE_API": "1", - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1", } with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server: diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 5b313464dad1..adece09c0a1a 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -34,6 +34,7 @@ "hermes", "--reasoning-parser", "deepseek_r1", + "--enable-scale-out", ] TOOLS = [ @@ -54,11 +55,7 @@ @pytest.fixture(scope="module") def server(): - with RemoteOpenAIServer( - MODEL, - ARGS, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL, ARGS) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py index 70bbc95016ce..c6a6931948b9 100644 --- a/tests/entrypoints/scale_out/render/test_render_multimodal.py +++ b/tests/entrypoints/scale_out/render/test_render_multimodal.py @@ -27,15 +27,10 @@ def vision_server(): "1", "--limit-mm-per-prompt.video", "0", + "--enable-scale-out", ] - env_overrides = {"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"} - - with RemoteOpenAIServer( - VISION_MODEL_NAME, - args, - env_dict=env_overrides, - ) as remote_server: + with RemoteOpenAIServer(VISION_MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py index 9811bca8a22c..6d8335395ed8 100644 --- a/tests/entrypoints/scale_out/test_factories.py +++ b/tests/entrypoints/scale_out/test_factories.py @@ -3,7 +3,6 @@ from argparse import Namespace -import pytest from fastapi import FastAPI from vllm.entrypoints.launchers.api_server.routers import register_api_routers @@ -22,69 +21,49 @@ def registered_paths( - supported_tasks: tuple[SupportedTask, ...], *, tokens_only: bool = False + supported_tasks: tuple[SupportedTask, ...], + *, + tokens_only: bool = False, + enable_scale_out: bool = False, ) -> set[str]: app = FastAPI() - args = Namespace(tokens_only=tokens_only, enable_fault_tolerance=False) + args = Namespace( + tokens_only=tokens_only, + enable_scale_out=enable_scale_out, + enable_fault_tolerance=False, + ) app.state.args = args register_api_routers(args, app, supported_tasks) return {route.path for route in app.routes} -def test_scale_out_routes_are_disabled_by_default(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_scale_out_routes_are_disabled_by_default(): paths = registered_paths(("generate",)) assert SCALE_OUT_PATHS.isdisjoint(paths) -def test_disabled_scale_out_routes_are_logged(monkeypatch, caplog): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_disabled_scale_out_routes_are_logged(caplog): with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"): registered_paths(("generate",)) - assert any( - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1" in record.message - for record in caplog.records - ) + assert any("--enable-scale-out" in record.message for record in caplog.records) -def test_scale_out_routes_can_be_enabled_for_generate_server(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") - - paths = registered_paths(("generate",)) +def test_scale_out_routes_can_be_enabled_for_generate_server(): + paths = registered_paths(("generate",), enable_scale_out=True) assert paths >= SCALE_OUT_PATHS -def test_render_routes_remain_enabled_for_render_server(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_render_routes_remain_enabled_for_render_server(): paths = registered_paths(("render",)) assert paths >= RENDER_PATHS assert GENERATE_PATH not in paths -def test_render_server_rejects_explicitly_disabled_scale_out_routes(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - - with pytest.raises(ValueError, match="VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0"): - registered_paths(("render",)) - - -def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(): paths = registered_paths(("generate",), tokens_only=True) assert paths >= {GENERATE_PATH, "/abort_requests"} - - -def test_tokens_only_mode_rejects_explicitly_disabled_scale_out_routes(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - - with pytest.raises(ValueError, match="--tokens-only"): - registered_paths(("generate",), tokens_only=True) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py index 963e416e7ea8..e7527a071b75 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py @@ -34,12 +34,9 @@ def server(): "--enable-return-routed-experts", "--hf-overrides", '{"sliding_window": null}', + "--enable-scale-out", ] - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 9491623fd115..892d3c8285da 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -37,13 +37,10 @@ def server(): "4096", "--enforce-eager", "--no-enable-prefix-caching", + "--enable-scale-out", ] - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index fb741df061e1..1c349e8d6724 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -69,6 +69,7 @@ def server(request): # Option B: deterministic prefix caching # "--enable-prefix-caching", # "--deterministic-prefix-caching", + "--enable-scale-out", ] extra_args = getattr(request, "param", None) @@ -79,11 +80,7 @@ def server(request): else [str(extra_args)] ) - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/test_envs.py b/tests/test_envs.py index aaa53d2a5448..1ef40cde0cc0 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -43,41 +43,6 @@ def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch): assert "VLLM_API_KEY" not in envs.compile_factors() -def test_scale_out_endpoints_flag_is_runtime_only(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") - - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is True - assert "VLLM_ENABLE_SCALE_OUT_ENDPOINTS" not in envs.compile_factors() - - -def test_scale_out_endpoints_flag_distinguishes_unset_from_disabled( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None - - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is False - - -def test_scale_out_endpoints_flag_treats_empty_as_unset( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "") - - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None - - -@pytest.mark.parametrize("value", ["-1", "2", "01", "+1", "invalid", " "]) -def test_scale_out_endpoints_flag_rejects_values_other_than_zero_or_one( - monkeypatch: pytest.MonkeyPatch, value: str -): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", value) - - with pytest.raises(ValueError, match="must be 0 or 1"): - _ = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS - - def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 445964e6aa47..eaac36f1336f 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -17,9 +17,10 @@ Even when the SDK is installed, :func:`attach_router` also requires ``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose -the chat route. The render route additionally requires -``VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1``. These gates keep unrelated -deployments from accidentally exposing the APIs. +the chat route. The render route additionally requires scale-out +endpoints to be enabled, i.e. ``--enable-scale-out`` or +``--tokens-only``. These gates keep unrelated deployments from +accidentally exposing the APIs. Note: the handlers must live at module scope (not inside ``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves @@ -306,6 +307,7 @@ def attach_router(app: FastAPI) -> None: ) return app.include_router(router) - if envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS: + args = getattr(app.state, "args", None) + if getattr(args, "enable_scale_out", False) or getattr(args, "tokens_only", False): app.include_router(render_router) app.add_middleware(CohereErrorEnvelopeMiddleware) diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 0da96b43e1dc..89af22cdc80f 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -181,6 +181,13 @@ class BaseFrontendArgs: If set to True, only enable the Tokens In<>Out endpoint. This is intended for use in a Disaggregated Everything setup. """ + enable_scale_out: bool = False + """ + If set to True, register the scale-out endpoints (`/render`, `/derender`, + and `/inference/v1/generate`) on `vllm serve`. Has no effect on + `vllm launch render` or `vllm serve --tokens-only`, which always register + their required endpoints regardless of this flag. + """ fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full" """Controls the ``system_fingerprint`` field on responses. diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py index 30c329c51f02..bd48c709ddc0 100644 --- a/vllm/entrypoints/scale_out/factories.py +++ b/vllm/entrypoints/scale_out/factories.py @@ -5,7 +5,6 @@ from fastapi import FastAPI -from vllm import envs from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm.tasks import SupportedTask @@ -67,22 +66,15 @@ def register_scale_out_api_routers( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], ): - dedicated_mode = None - if "render" in supported_tasks: - dedicated_mode = "`vllm launch render`" - elif getattr(app.state.args, "tokens_only", False): - dedicated_mode = "`--tokens-only`" - - enabled = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS - if dedicated_mode is not None and enabled is False: - raise ValueError( - "`VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0` conflicts with " - f"{dedicated_mode}; unset it or set it to 1" - ) - if dedicated_mode is None and not enabled: + args = getattr(app.state, "args", None) + # `vllm launch render` and `vllm serve --tokens-only` are dedicated to + # serving these endpoints, so they always register them regardless of + # `--enable-scale-out`. + dedicated_mode = "render" in supported_tasks or getattr(args, "tokens_only", False) + enabled = dedicated_mode or getattr(args, "enable_scale_out", False) + if not enabled: logger.info( - "Scale-out endpoints are disabled. Set " - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them." + "Scale-out endpoints are disabled. Set --enable-scale-out to enable them." ) return diff --git a/vllm/envs.py b/vllm/envs.py index 0269a96409fe..6158f86f896b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -261,7 +261,6 @@ VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True VLLM_ENABLE_RESPONSES_API_STORE: bool = False VLLM_ENABLE_COHERE_API: bool = False - VLLM_ENABLE_SCALE_OUT_ENDPOINTS: bool | None = None VLLM_HAS_FLASHINFER_CUBIN: bool = False VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True @@ -352,15 +351,6 @@ def maybe_convert_bool(value: str | None) -> bool | None: return bool(int(value)) -def maybe_convert_scale_out_endpoints(value: str | None) -> bool | None: - if value is None: - return None - normalized = value.strip() - if normalized not in ("0", "1"): - raise ValueError("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1") - return maybe_convert_bool(normalized) - - def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None: if value is None: return None @@ -1858,12 +1848,6 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ENABLE_COHERE_API": lambda: bool( int(os.getenv("VLLM_ENABLE_COHERE_API", "0")) ), - # If set to 1, expose the scale-out endpoints on `vllm serve`. - # The dedicated `vllm launch render` server exposes render and derender - # endpoints when this variable is unset or set to 1. - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": lambda: maybe_convert_scale_out_endpoints( - os.getenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS") or None - ), # If set, use the fp8 mfma in rocm paged attention. "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool( int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0")) @@ -2289,7 +2273,6 @@ def compile_factors() -> dict[str, object]: "VLLM_CONFIG_ROOT", "LD_LIBRARY_PATH", "VLLM_SERVER_DEV_MODE", - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "VLLM_DP_MASTER_IP", "VLLM_DP_MASTER_PORT", "VLLM_NIXL_SIDE_CHANNEL_HOST", From 1678b396270406c27fcab8f5b86b21fd305ac605 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Mon, 14 Sep 2026 18:05:13 +0800 Subject: [PATCH 173/774] [Bugfix][EPD] Preserve media processing options in encoder requests (#56786) Signed-off-by: jiangkuaixue123 --- .../disaggregated_encoder/disagg_epd_proxy.py | 11 ++++- .../ec_connector/unit/test_epd_proxy_retry.py | 41 ++++++++++++++++--- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 3912c1c67a4b..830da9b5d64e 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -278,7 +278,6 @@ async def fanout_encoder_primer( item_transfer_ids[idx] = transfer_id encoder_req = { - # You *may* need to keep additional fields "model": orig_request.get("model"), "messages": [ { @@ -292,6 +291,14 @@ async def fanout_encoder_primer( # once the prompt is encoded and its embeddings are published. "stream": False, } + for key in ( + "mm_processor_kwargs", + "media_io_kwargs", + "priority", + "session_id", + ): + if key in orig_request: + encoder_req[key] = orig_request[key] if consumer_zmq is not None: # No mm_hash here on purpose. The encoder's own # `mm_features[i].identifier` is derived from the uuid *and* the @@ -378,7 +385,7 @@ async def fanout_encoder_primer( # connector's own handle on the published embedding (for NIXL, # peer_host/peer_port/size_bytes). The decoder's connector # looks it up by mm_hash on the request, so carry it through. - ec_params[item_uuids.get(idx, ec_mm_hash)] = reported + ec_params[ec_mm_hash] = reported if NO_REWRITE and consumer_zmq is not None: ec_params.setdefault("ec_items", []).append( {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]} diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 9a92b1fdbcca..1a8c877c72e2 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -116,7 +116,9 @@ async def _no_prefill(req_data, p_url, req_id): first, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args)) reported = first["ec_transfer_params"] - assert [handle] == [value for key, value in reported.items() if key != "ec_items"] + assert {key: value for key, value in reported.items() if key != "ec_items"} == { + "encoder-side-hash": handle + } assert "ec_transfer_params" not in body # Attempt 2's encode reports nothing: decode must be told nothing. @@ -154,10 +156,22 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize( + "request_options", + [ + {}, + { + "mm_processor_kwargs": {"max_pixels": 262144}, + "media_io_kwargs": {"image": {"image_mode": "RGB"}}, + "priority": -3, + "session_id": "session-123", + }, + ], +) async def test_http_roundtrip_preserves_payload_and_response_bytes( - proxy, monkeypatch, stream, prefill + proxy, monkeypatch, stream, prefill, request_options ): - """Exercise real HTTP hops, rewrite and a decode retry without model servers.""" + """Preserve media and scheduling context across HTTP hops and retries.""" seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} prefix = b'data: {"content":"' # Split a multibyte character at the old 1024-byte forwarding boundary. @@ -174,10 +188,13 @@ async def backend(request): seen[stage].append(body) if stage == "encode": item = body["messages"][0]["content"][0] + mm_hash = item["uuid"] + if body.get("mm_processor_kwargs"): + mm_hash += "-processed" return web.json_response( { "ec_transfer_params": { - item["uuid"]: { + mm_hash: { "metadata": {"image_grid_thw": [[1, 2, 2]]}, "peer_port": len(seen[stage]), } @@ -230,6 +247,7 @@ async def backend(request): "max_tokens": 32, "seed": 42, "structured_outputs": {"choice": ["A", "B"]}, + **request_options, } await proxy.on_startup() try: @@ -247,6 +265,15 @@ async def backend(request): assert len(seen["encode"]) == 4 assert len(seen["decode"]) == 2 + for requests in seen.values(): + for forwarded in requests: + for key in ( + "mm_processor_kwargs", + "media_io_kwargs", + "priority", + "session_id", + ): + assert forwarded.get(key) == body.get(key) final = seen["decode"][-1] for key in [ "model", @@ -265,7 +292,11 @@ async def backend(request): "image_embeds": {"image_grid_thw": [1, 2, 2]}, "uuid": proxy.content_uuid(item), } - assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4) + mm_hash = proxy.content_uuid(item) + if request_options: + mm_hash += "-processed" + assert set(final["ec_transfer_params"]) - {"ec_items"} == {mm_hash} + assert final["ec_transfer_params"][mm_hash]["peer_port"] in (3, 4) if prefill: assert final["kv_transfer_params"] == {"remote_block_ids": [2]} From ea723c81c3ea26425cb69503a5d5e90822a04a45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:22:18 +0200 Subject: [PATCH 174/774] [Security] Cap Qwen-VL video sampling knobs (#56729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Cursor --- vllm/multimodal/video.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 4142febc3f92..c7391919794f 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -357,6 +357,9 @@ def load_bytes( video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"), ) class Qwen3VLVideoBackend(VideoBackend): + _MAX_FRAMES: ClassVar[int] = 768 + _MAX_FPS: ClassVar[int] = 30 + @classmethod def compute_frames_index_to_sample( cls, @@ -366,10 +369,10 @@ def compute_frames_index_to_sample( ) -> list[int]: total_frames_num = source.total_frames_num original_fps = source.original_fps - fps = target.fps + fps = min(target.fps, cls._MAX_FPS) max_frame_idx = source.total_frames_num - 1 min_frames = kwargs.get("min_frames", 4) - max_frames = kwargs.get("max_frames", 768) + max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES) # Refer to: # https://github.com/huggingface/transformers/blob/v5.9.0/src/transformers/models/qwen3_vl/video_processing_qwen3_vl.py#L119-L125 @@ -420,6 +423,9 @@ class Qwen2VLVideoBackend(VideoBackend): clip); it is clamped to the last valid frame. """ + _MAX_FRAMES: ClassVar[int] = 768 + _MAX_FPS: ClassVar[int] = 30 + @classmethod def compute_frames_index_to_sample( cls, @@ -433,7 +439,7 @@ def compute_frames_index_to_sample( original_fps = source.original_fps temporal_patch_size = kwargs.get("temporal_patch_size", 2) min_frames = kwargs.get("min_frames", 4) - max_frames = kwargs.get("max_frames", 768) + max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES) # vLLM reports original_fps == 0 for clips with unknown/variable fps # (VFR, malformed, streaming); fail loudly instead of dividing by zero. @@ -447,7 +453,7 @@ def compute_frames_index_to_sample( math.floor(min(max_frames, total_frames_num) / temporal_patch_size) * temporal_patch_size ) - n = total_frames_num / original_fps * target.fps + n = total_frames_num / original_fps * min(target.fps, cls._MAX_FPS) n = min(max(n, min_frames), max_frames, total_frames_num) n = math.floor(n / temporal_patch_size) * temporal_patch_size From 0ca7aef4e341cb4f4d7233291c304a168f4384aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 12:37:12 +0200 Subject: [PATCH 175/774] [Bugfix][NixlPush] Guard _remote_agents read in _do_send_reg_notif (X1) (#56317) Signed-off-by: NickLucche --- .../unit/test_nixl_push_connector.py | 104 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 5 +- .../kv_connector/v1/nixl/push_worker.py | 4 +- 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 1cdafcfe482f..2c7f0cb866bf 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -362,6 +362,7 @@ def fresh(cls) -> _StubWriterWorker: w.world_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} + w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False w.region_group_ids = [0] @@ -418,6 +419,84 @@ def _registration_data( } +class TestPushWriterRegSend: + """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of + the P engine, and the agent map is read under ``_handshake_lock`` so the + send can't race the handshake done-callback's insert or a concurrent + stale-engine eviction (X1).""" + + def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]: + w = _StubWriterWorker.fresh() + w.nixl_wrapper = MagicMock() + failures: list = [] + w._handle_failed_transfer = lambda *a: failures.append(a) + return w, failures + + def test_sends_to_all_agents_with_prefix_and_payload(self): + w, failures = self._reg_send_worker() + w._remote_agents["prefill-engine"] = {(0, 0): "agent-0", (0, 1): "agent-1"} + rd = _registration_data("req-1") + + w._do_send_reg_notif("req-1", rd) + + assert failures == [] + calls = w.nixl_wrapper.send_notif.call_args_list + assert len(calls) == 2 + assert {c.args[0] for c in calls} == {"agent-0", "agent-1"} + for call in calls: + notif = call.kwargs["notif_msg"] + assert notif.startswith(PUSH_REG_NOTIF_PREFIX) + # (Tuples decode back as lists, so compare identity fields.) + decoded = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :]) + assert decoded["request_id"] == "req-1" + assert decoded["decode_engine_id"] == rd["decode_engine_id"] + assert decoded["remote_engine_id"] == rd["remote_engine_id"] + assert decoded["local_block_ids"] == [ + list(group) for group in rd["local_block_ids"] + ] + + def test_fails_request_when_engine_not_handshaked(self): + w, failures = self._reg_send_worker() + + w._do_send_reg_notif("req-1", _registration_data("req-1")) + + assert failures == [("req-1", None)] + w.nixl_wrapper.send_notif.assert_not_called() + + def test_agent_read_is_serialized_with_handshake_lock(self): + """While another thread holds ``_handshake_lock`` (as the handshake + done-callback does while writing ``_remote_agents``), the writer's + check-then-send must not run ahead of the critical section.""" + w, failures = self._reg_send_worker() + w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"} + rd = _registration_data("req-1") + + w._handshake_lock.acquire() + done = threading.Event() + errors: list[BaseException] = [] + + def _send(): + try: + w._do_send_reg_notif("req-1", rd) + except BaseException as e: # pragma: no cover + errors.append(e) + finally: + done.set() + + t = threading.Thread(target=_send) + t.start() + assert not done.wait(timeout=0.2) + w.nixl_wrapper.send_notif.assert_not_called() + assert failures == [] + + w._handshake_lock.release() + assert done.wait(timeout=2.0) + t.join(timeout=2.0) + assert errors == [] + assert failures == [] + assert w.nixl_wrapper.send_notif.call_count == 1 + + class TestPushWriterMatching: def test_handle_push_reg_matches_existing_finished_blocks(self): """PUSH_REG arrives second (P finished first): match + fire.""" @@ -727,6 +806,31 @@ def test_stale_engine_evicted_on_push(): w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old") +def test_cleanup_remote_engine_pops_under_handshake_lock(): + """Regression (X1): eviction must not remove ``_remote_agents`` entries + while another thread is inside the lock's critical section (e.g. the + push writer reading the map to send a registration).""" + w = _eviction_worker(engine_ttl=30.0) + w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} + + w._handshake_lock.acquire() + done = threading.Event() + + def _cleanup(): + w._cleanup_remote_engine("D-old") + done.set() + + t = threading.Thread(target=_cleanup) + t.start() + assert "D-old" in w._remote_agents + + w._handshake_lock.release() + assert done.wait(timeout=2.0) + t.join(timeout=2.0) + assert "D-old" not in w._remote_agents + w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old") + + class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 4edfaefa2f5f..58ab5850eaf7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -3264,7 +3264,10 @@ def _cleanup_remote_engine( # Notif-only engines (push-mode D side) have no descriptor state. for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values(): self.nixl_wrapper.release_dlist_handle(handle) - for agent_name in self._remote_agents.pop(engine_id).values(): + # Pop under the handshake lock; NIXL teardown stays outside it. + with self._handshake_lock: + agents = self._remote_agents.pop(engine_id) + for agent_name in agents.values(): self.nixl_wrapper.remove_remote_agent(agent_name) self.kv_caches_base_addr.pop(engine_id, None) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index e506aff1f425..603321260de2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -354,7 +354,9 @@ def _on_handshake( def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: engine_id = reg_data["remote_engine_id"] notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data) - agents = self._remote_agents.get(engine_id) + # _remote_agents is mutated on other threads; snapshot under the lock. + with self._handshake_lock: + agents = dict(self._remote_agents.get(engine_id) or {}) if not agents: logger.error( "No remote agents for engine %s; cannot send registration for %s", From 7b1ea3f524445be0545ed962cb06ce21a93fcc33 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:37:59 +0800 Subject: [PATCH 176/774] [Bugfix][Rust Frontend][Multimodal] Align DeepSeek V4.1 and Kimi K3 media with rendered placeholders (#56366) Co-authored-by: Bugen Zhao Co-authored-by: Codex Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 10 +- rust/src/chat/src/multimodal.rs | 205 +++++++++++++----- rust/src/chat/src/renderer/deepseek.rs | 65 +++++- .../src/chat/src/renderer/deepseek_v32/mod.rs | 1 + rust/src/chat/src/renderer/deepseek_v4/mod.rs | 1 + .../src/chat/src/renderer/deepseek_v41/mod.rs | 5 +- .../chat/src/renderer/deepseek_v41/tests.rs | 78 ++++++- rust/src/chat/src/renderer/harmony/mod.rs | 1 + rust/src/chat/src/renderer/hf/mod.rs | 1 + rust/src/chat/src/renderer/inkling/mod.rs | 1 + .../src/chat/src/renderer/kimi_k3/encoding.rs | 51 +++-- rust/src/chat/src/renderer/kimi_k3/mod.rs | 5 +- rust/src/chat/src/renderer/kimi_k3/tests.rs | 59 +++++ rust/src/chat/src/renderer/mod.rs | 13 ++ rust/src/chat/tests/chat.rs | 1 + rust/src/server/src/grpc/tests.rs | 1 + .../server/src/routes/http_client_tests.rs | 1 + rust/src/server/src/routes/tests.rs | 1 + 18 files changed, 415 insertions(+), 85 deletions(-) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1b5d5aea6f92..9c59d1402896 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -34,8 +34,8 @@ pub use parser::{ParserSelection, validate_parser_overrides}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer, - DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, - RendererSelection, + DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource, + RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -112,6 +112,10 @@ impl ChatRequestProcessor { request: &ChatRequest, rendered: RenderedPrompt, ) -> Result<(Prompt, Option)> { + let media_is_empty = match &rendered.media_order { + Some(media_order) => media_order.is_empty(), + None => !request.has_multimodal(), + }; match self.model_dtype { Some(model_dtype) => { multimodal::finalize_rendered_prompt( @@ -122,7 +126,7 @@ impl ChatRequestProcessor { ) .await } - None if !request.has_multimodal() => Ok((rendered.prompt, None)), + None if media_is_empty => Ok((rendered.prompt, None)), None => Err(Error::UnsupportedMultimodalRenderer), } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 674979e6ce63..a5dd3ba8ea65 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -16,7 +16,7 @@ use std::fs; use std::path::Path; use std::sync::{Arc, LazyLock}; -use itertools::izip; +use itertools::{Either, izip}; use llm_multimodal::{ AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, @@ -33,7 +33,7 @@ use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; use crate::error::{Error, Result, bail_multimodal, multimodal}; -use crate::renderer::RenderedPrompt; +use crate::renderer::{MediaPartSource, RenderedPrompt}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; mod audio; @@ -549,7 +549,8 @@ pub(crate) async fn finalize_rendered_prompt( info: Option<&MultimodalModelInfo>, model_dtype: ModelDtype, ) -> Result<(Prompt, Option)> { - if !request.has_multimodal() { + let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?; + if media_parts.is_empty() { return Ok((rendered.prompt, None)); } let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?; @@ -561,63 +562,92 @@ pub(crate) async fn finalize_rendered_prompt( .map_err(|error| multimodal!("{error}"))?, Prompt::TokenIds(token_ids) => token_ids, }; - let media_parts = extract_media_parts(request)?; let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?; Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract media parts from chat messages in message/content order. -/// -/// Assistant history is skipped because generated assistant blocks are already -/// represented as text for prompt rendering in this crate. -fn extract_media_parts(request: &ChatRequest) -> Result> { - let mut all_parts = Vec::new(); - for message in &request.messages { - let content = match message { - ChatMessage::System { content } - | ChatMessage::Developer { content, .. } - | ChatMessage::User { content } - | ChatMessage::ToolResponse { content, .. } => content, - ChatMessage::Assistant { .. } => continue, - }; - let ChatContent::Parts(parts) = content else { - continue; - }; - for part in parts { - match part { - ChatContentPart::Text { .. } => {} - ChatContentPart::ImageUrl { - image_url, - detail, - uuid, - } => all_parts.push(MediaContentPart::ImageUrl { - url: image_url.clone(), - detail: *detail, - uuid: uuid.clone(), - }), - ChatContentPart::VideoUrl { video_url, uuid } => { - all_parts.push(MediaContentPart::VideoUrl { - url: video_url.clone(), - uuid: uuid.clone(), - }) - } - ChatContentPart::InputAudio { data, format, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: input_audio_data_url(data, format.as_deref())?, - uuid: uuid.clone(), - }) - } - ChatContentPart::AudioUrl { audio_url, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: audio_url.clone(), - uuid: uuid.clone(), - }) +/// Resolve media parts in the placeholder order reported by the renderer. +fn extract_media_parts( + request: &ChatRequest, + media_order: Option<&[MediaPartSource]>, +) -> Result> { + let parts = match media_order { + // Resolve exactly the sources selected by the renderer, in placeholder order. + Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> { + let message = request.messages.get(source.message_index).ok_or_else(|| { + multimodal!( + "renderer reported missing chat message {} for multimodal content", + source.message_index + ) + })?; + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => content, + ChatMessage::Assistant { .. } => { + bail_multimodal!("renderer reported multimodal assistant content") } + }; + let ChatContent::Parts(parts) = content else { + bail_multimodal!("renderer reported multimodal content in a text-only chat message") + }; + let part = parts.get(source.content_part_index).ok_or_else(|| { + multimodal!( + "renderer reported missing content part {} in chat message {}", + source.content_part_index, + source.message_index + ) + })?; + Ok(part) + })), + // Fall back to all media in request message and content-part order (e.g. Jinja). + None => Either::Right(request.messages.iter().flat_map(|message| { + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => Some(content), + ChatMessage::Assistant { .. } => None, + }; + match content { + Some(ChatContent::Parts(parts)) => parts.as_slice(), + _ => &[], } - } - } - Ok(all_parts) + .iter() + .filter(|part| part.is_multimodal()) + .map(Ok) + })), + }; + parts + .map(|part| match part? { + ChatContentPart::Text { .. } => { + bail_multimodal!("renderer reported a text part as multimodal content") + } + ChatContentPart::ImageUrl { + image_url, + detail, + uuid, + } => Ok(MediaContentPart::ImageUrl { + url: image_url.clone(), + detail: *detail, + uuid: uuid.clone(), + }), + ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }), + ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl { + url: input_audio_data_url(data, format.as_deref())?, + uuid: uuid.clone(), + }), + ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl { + url: audio_url.clone(), + uuid: uuid.clone(), + }), + }) + .collect() } /// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`. @@ -1043,6 +1073,75 @@ mod tests { assert!(input_audio_data_url("AAAA", Some("flac")).is_err()); } + #[test] + fn extract_media_parts_uses_request_order_or_explicit_selection() { + let request = ChatRequest { + messages: vec![ + ChatMessage::user(vec![ + ChatContentPart::text("before image"), + ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-b".to_string(), + detail: None, + uuid: Some("image-b".to_string()), + }, + ]), + ChatMessage::user(vec![ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-a".to_string(), + detail: None, + uuid: Some("image-a".to_string()), + }]), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ], + ..ChatRequest::for_test() + }; + let media_order = [ + MediaPartSource { + message_index: 1, + content_part_index: 0, + }, + MediaPartSource { + message_index: 0, + content_part_index: 1, + }, + ]; + + let selections = [None, Some(&media_order[..]), Some(&[][..])]; + let uuids = selections.map(|order| { + extract_media_parts(&request, order) + .unwrap() + .into_iter() + .map(|part| match part { + MediaContentPart::ImageUrl { uuid, .. } => uuid, + _ => panic!("expected image URL"), + }) + .collect::>() + }); + + expect_test::expect![[r#" + [ + [ + Some( + "image-b", + ), + Some( + "image-a", + ), + ], + [ + Some( + "image-a", + ), + Some( + "image-b", + ), + ], + [], + ] + "#]] + .assert_debug_eq(&uuids); + } + fn image_url_part() -> MediaContentPart { MediaContentPart::ImageUrl { url: "https://example.com/image.png".to_string(), diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5f900ab88b59..5925c5fd1655 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -16,6 +16,7 @@ use serde_json_fmt::JsonFormat; use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER; +use super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort, @@ -87,6 +88,14 @@ struct RenderedToolSchema<'a> { /// Render one chat request into the final prompt string. pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result { + render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt) +} + +/// Render one chat request and record media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + dialect: DsDialect, +) -> Result<(String, Vec)> { let (thinking_mode, reasoning_effort_prompt) = match dialect { DsDialect::V4 => { resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))? @@ -112,6 +121,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul dialect, ); let mut out = String::from(BOS_TOKEN); + let mut media_order = Vec::new(); if dialect == DsDialect::V41 && (thinking_mode == ThinkingMode::Thinking || synthetic_tool_system @@ -126,7 +136,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul let mut request_tools_attached = false; let mut render_index = 0isize; if synthetic_tool_system { - render_system_message(&mut out, None, request_tools, dialect)?; + render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?; request_tools_attached = true; render_index += 1; } @@ -157,14 +167,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } else { &[] }; - render_system_message(&mut out, Some(content), tools, dialect)?; + render_system_message( + &mut out, + Some((message_index, content)), + tools, + dialect, + &mut media_order, + )?; } ChatMessage::Developer { content, tools } => { render_developer_message( &mut out, + message_index, content, tools.as_deref().unwrap_or(&[]), dialect, + &mut media_order, )?; } ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => { @@ -174,6 +192,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul message_index, &last_tool_call_order, dialect, + &mut media_order, )?; } ChatMessage::Assistant { content } => { @@ -223,7 +242,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } } - Ok(out) + Ok((out, media_order)) } /// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer @@ -488,12 +507,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> { /// Render a system turn, optionally followed by the V4 tool preamble. fn render_system_message( out: &mut String, - content: Option<&ChatContent>, + content: Option<(usize, &ChatContent)>, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { - if let Some(content) = content { - write_chat_content(out, content, dialect)?; + if let Some((message_index, content)) = content { + write_chat_content(out, content, dialect, message_index, media_order)?; } if !tools.is_empty() { out.push_str("\n\n"); @@ -505,9 +525,11 @@ fn render_system_message( /// Developer messages are rendered as user-like turns with optional tools. fn render_developer_message( out: &mut String, + message_index: usize, content: &ChatContent, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -516,7 +538,7 @@ fn render_developer_message( } out.push_str(USER_SP_TOKEN); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; if !tools.is_empty() { out.push_str("\n\n"); render_tools(out, tools, dialect)?; @@ -531,6 +553,7 @@ fn render_user_content_block( message_index: usize, tool_call_order: &HashMap, dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { let (block_start, block_end) = user_content_block_bounds(messages, message_index); let mut sorted_tool_indices = @@ -542,7 +565,9 @@ fn render_user_content_block( out.push_str("\n\n"); } match &messages[message_index] { - ChatMessage::User { content } => write_chat_content(out, content, dialect)?, + ChatMessage::User { content } => { + write_chat_content(out, content, dialect, message_index, media_order)?; + } ChatMessage::ToolResponse { .. } => { let sorted_index = sorted_tool_indices .next() @@ -550,7 +575,7 @@ fn render_user_content_block( let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else { unreachable!("sorted tool response index should reference a tool message"); }; - write_tool_result(out, content, dialect)?; + write_tool_result(out, content, dialect, sorted_index, media_order)?; } _ => unreachable!("user content block should only contain user content messages"), } @@ -597,9 +622,15 @@ fn sorted_tool_response_indices( } /// Render one tool response payload inside a V4 `` block. -fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_tool_result( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { out.push_str(""); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; out.push_str(""); Ok(()) } @@ -727,7 +758,13 @@ fn encode_arguments_to_dsml( /// The V4.1 dialect inlines the image placeholder at each image part's /// position unconditionally (matching the Python encoding's /// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts. -fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_chat_content( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { match content { ChatContent::Text(text) => out.push_str(text), ChatContent::Parts(parts) => { @@ -738,6 +775,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec match part { ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => { out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER); + media_order.push(MediaPartSource { + message_index, + content_part_index: index, + }); } _ => out.push_str(part.as_text()?), } diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs index 1a41e14f8fc1..03c2765d1913 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs @@ -26,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(encoding::render_request(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs index 155a6ee1403b..54b111e5050f 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs @@ -24,6 +24,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/mod.rs b/rust/src/chat/src/renderer/deepseek_v41/mod.rs index 4e5ddb37257e..7737f7ebda4f 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/mod.rs @@ -21,9 +21,12 @@ impl DeepSeekV41ChatRenderer { impl ChatRenderer for DeepSeekV41ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (prompt, media_order) = + deepseek::render_request_with_media_order(request, DsDialect::V41)?; Ok(RenderedPrompt { - prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?), + prompt: Prompt::Text(prompt), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index d7b0c538bb54..3babb5c269f8 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -8,7 +8,7 @@ use serde_json::{Value, json}; use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; -use crate::event::AssistantContentBlock; +use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; @@ -83,6 +83,82 @@ fn inlines_image_placeholder_at_image_part_positions() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +#[test] +fn media_order_excludes_dropped_historical_developer_content() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::developer(vec![image_part("historical-image")], None), + ChatMessage::user("question"), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(!prompt.contains("<｜deepseek_image｜>")); + assert!(media_order.is_empty()); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn maps_reasoning_effort_to_reference_numeric_budget() { for (effort, budget) in [ diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs index b07d8458d739..71431cc92b46 100644 --- a/rust/src/chat/src/renderer/harmony/mod.rs +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -116,6 +116,7 @@ impl ChatRenderer for HarmonyChatRenderer { fn render(&self, request: &ChatRequest) -> Result { Ok(RenderedPrompt { prompt: Prompt::TokenIds(self.render_token_ids(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 5e7ba75e0a7f..ba2b76d10003 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -217,6 +217,7 @@ impl HfChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs index a33e3414f2db..e8afe66396dd 100644 --- a/rust/src/chat/src/renderer/inkling/mod.rs +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -346,6 +346,7 @@ impl ChatRenderer for InklingChatRenderer { Ok(RenderedPrompt { prompt: Prompt::TokenIds(out), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index 392b6b642d29..88b4e3ae6101 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -10,6 +10,7 @@ use std::collections::HashMap; use serde_json::{Map, Value, json}; use vllm_tokenizer::Tokenizer; +use super::super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, @@ -29,6 +30,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; pub(super) struct K3TokenWriter<'a> { tokenizer: &'a dyn Tokenizer, token_ids: Vec, + media_order: Vec, } impl<'a> K3TokenWriter<'a> { @@ -36,6 +38,7 @@ impl<'a> K3TokenWriter<'a> { Self { tokenizer, token_ids: Vec::new(), + media_order: Vec::new(), } } @@ -55,13 +58,25 @@ impl<'a> K3TokenWriter<'a> { Ok(()) } - pub(super) fn finish(self) -> Vec { - self.token_ids + fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> { + self.control(IMAGE_PLACEHOLDER)?; + self.media_order.push(MediaPartSource { + message_index, + content_part_index, + }); + Ok(()) + } + + pub(super) fn finish(self) -> (Vec, Vec) { + (self.token_ids, self.media_order) } } -/// Render and tokenize one chat request using K3's segment-aware contract. -pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { +/// Render and tokenize one chat request while recording media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + tokenizer: &dyn Tokenizer, +) -> Result<(Vec, Vec)> { let thinking = thinking_enabled(request)?; let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; let tools = request_tools(request); @@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - resolved.sort_by_key(|item| (item.0, item.1)); } - for (position, _, content, name) in resolved { + for (position, message_index, content, name) in resolved { let tool_name = name.as_deref().ok_or_else(|| { Error::ChatTemplate( "Kimi K3 tool messages need a resolvable tool name: \ @@ -135,7 +150,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - .to_string(), ) })?; - write_tool_message(out, tool_name, position, &content)?; + write_tool_message(out, tool_name, position, message_index, &content)?; } Ok(()) }; @@ -161,14 +176,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - } if !local_tools.is_empty() => { write_tool_declare(&mut out, local_tools, true)?; if !content_is_empty(content) { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } } ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } ChatMessage::User { content } => { - write_role_message(&mut out, "user", None, content)?; + write_role_message(&mut out, "user", None, message_index, content)?; } ChatMessage::Assistant { content } => { tool_call_id_index.clear(); @@ -364,6 +379,7 @@ fn write_role_message( out: &mut K3TokenWriter<'_>, role: &str, name: Option<&str>, + message_index: usize, content: &ChatContent, ) -> Result<()> { let mut attrs = vec![("role", role.to_string())]; @@ -372,7 +388,7 @@ fn write_role_message( } let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); write_open_tag(out, "message", &attr_refs)?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -381,6 +397,7 @@ fn write_tool_message( out: &mut K3TokenWriter<'_>, tool_name: &str, index: usize, + message_index: usize, content: &ChatContent, ) -> Result<()> { let index_str = index.to_string(); @@ -389,7 +406,7 @@ fn write_tool_message( "message", &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], )?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -471,14 +488,20 @@ fn write_assistant_tool_call( write_close_tag(out, "call") } -fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { +fn write_content( + out: &mut K3TokenWriter<'_>, + message_index: usize, + content: &ChatContent, +) -> Result<()> { match content { ChatContent::Text(text) => write_text_with_images(out, text), ChatContent::Parts(parts) => { - for part in parts { + for (content_part_index, part) in parts.iter().enumerate() { match part { ChatContentPart::Text { text } => write_text_with_images(out, text)?, - ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::ImageUrl { .. } => { + out.image(message_index, content_part_index)?; + } ChatContentPart::VideoUrl { .. } => { return Err(Error::UnsupportedMultimodalContent("video_url")); } diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs index 09bf356f5d64..90750eeac406 100644 --- a/rust/src/chat/src/renderer/kimi_k3/mod.rs +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -30,9 +30,12 @@ impl KimiK3ChatRenderer { impl ChatRenderer for KimiK3ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (token_ids, media_order) = + encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?; Ok(RenderedPrompt { - prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + prompt: Prompt::TokenIds(token_ids), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 1d9f984dbbd3..d0ecbe44d38b 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -284,6 +284,65 @@ fn partial_tool_results_keep_assistant_call_position() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())); + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let Prompt::TokenIds(token_ids) = rendered.prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + let prompt = test_tokenizer().decode(&token_ids, false).unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn defaults_thinking_effort_to_max() { let rendered = render_request(&crate::request::ChatRequest::for_test()); diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index a4649a34f70a..bc1831c375a9 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer; pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; +/// Location of one multimodal content part in the source chat request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MediaPartSource { + /// Zero-based message index in the source request. + pub message_index: usize, + /// Zero-based content-part index within the source message. + pub content_part_index: usize, +} + /// Rendered chat prompt submitted to the text backend. #[derive(Debug, Clone, PartialEq)] pub struct RenderedPrompt { /// The rendered prompt, either as text or already tokenized. pub prompt: Prompt, + /// Media sources in rendered placeholder order. + /// `None` uses the message and content-part order in the request (e.g. Jinja). + /// `Some` uses exactly the listed sources; an empty list omits all media. + pub media_order: Option>, /// Effective chat-template kwargs visible to the renderer after applying /// server defaults, request overrides, and typed reasoning controls. pub effective_template_kwargs: HashMap, diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ab99b4d8ba73..ed96134cdc23 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -250,6 +250,7 @@ impl ChatRenderer for FakeChatBackend { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: request.chat_options.template_kwargs.clone(), }) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f13dc11b277d..f5f74d557e69 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -252,6 +252,7 @@ impl ChatRenderer for FakeTextBackend { fn render(&self, _request: &ChatRequest) -> vllm_chat::Result { Ok(RenderedPrompt { prompt: Prompt::Text(String::new()), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index 216670b695b6..a0e0872e64ce 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -186,6 +186,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 5ee2b51f3f5c..6fbe22e501c3 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -540,6 +540,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(vllm_chat::RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) } From e82794f48ff0812b01543f7c19c34eb0c5b3b1ee Mon Sep 17 00:00:00 2001 From: Xiaozhen Bi Date: Mon, 14 Sep 2026 03:48:53 -0700 Subject: [PATCH 177/774] [CI][Test] Mock CPU backend block sizes in kv_connector unit conftest (#56671) Signed-off-by: Xiaozhen Bi Co-authored-by: Claude Opus 5 --- tests/v1/kv_connector/unit/conftest.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/v1/kv_connector/unit/conftest.py diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py new file mode 100644 index 000000000000..58c6e632d3d9 --- /dev/null +++ b/tests/v1/kv_connector/unit/conftest.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from vllm.platforms import current_platform + + +@pytest.fixture(autouse=True) +def _mock_cpu_backend_block_sizes(monkeypatch): + """Mock the CPU attention backend's kernel block-size declaration. + + Connector unit tests exercise transfer logic, never the attention + kernel, and hard-code block_size=16. On CPU-only machines + `CPUAttentionBackend` declares `MultipleOf(32)` (#54042), so every + worker construction would fail in `select_common_block_size` before + reaching the code under test. + """ + if not current_platform.is_cpu(): + return + from vllm.v1.attention.backend import MultipleOf + from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend + + monkeypatch.setattr( + CPUAttentionBackend, + "get_supported_kernel_block_sizes", + staticmethod(lambda: [MultipleOf(16)]), + ) From a6c5d6d0fcd7dd5b49e3146991c201b16caebef2 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Mon, 14 Sep 2026 19:21:30 +0800 Subject: [PATCH 178/774] [ROCm][Perf] Enable CSA multi-stream overlap for DeepSeek-V4 (#51794) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- vllm/models/deepseek_v4/amd/model.py | 11 +- vllm/models/deepseek_v4/amd/rocm.py | 194 ++++++++++++++++++++++ vllm/models/deepseek_v4/attention.py | 31 ++-- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 3 +- 4 files changed, 214 insertions(+), 25 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 349ff5267981..c1a6903937b6 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -957,16 +957,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps - # Three aux streams: one per non-default input GEMM in - # DeepseekV4Attention._run_parallel_input_projections - # (compressor kv_score, indexer.weights_proj, indexer.compressor - # kv_score). fused_wqa_wkv stays on the default stream. - # Disable them on ROCm because of hang issues. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] self.device = current_platform.device_type # Reserved topk indices buffer for all Indexer layers to reuse. diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index b39ad953e75d..3cc9827c1fbe 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -23,6 +23,7 @@ from vllm.platforms import current_platform from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton +from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) @@ -571,6 +572,199 @@ def __init__(self, *args, **kwargs): self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None + if self.indexer is None: + # Only enable multi-stream overlap for CSA layer now. + self.aux_stream_list = None + else: + # Disable indexer inner overlap. + self.indexer.aux_stream = None + + def _enable_csa_multi_stream(self) -> bool: + """All CSA multi-stream gates: env var, streams, and capture region. + + Dict metadata marks piecewise cudagraph, whose eager breaks rebuild + the attention inputs on the owning stream. Forking side streams + there would rely on runtime HIP event sync, which is unreliable in + this overlap on ROCm (event waits can hang), so multi-stream only + runs where the fork/join becomes static graph edges: inside capture, + or with non-dict metadata (full cudagraph or the profile run), which + has no eager breaks. + """ + attn_metadata = get_forward_context().attn_metadata + return self.aux_stream_list is not None and ( + torch.cuda.is_current_stream_capturing() + or not isinstance(attn_metadata, dict) + ) + + def _run_sequential_pipeline( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Disable ROCm streams when the current execution region cannot overlap.""" + aux_streams = self.aux_stream_list + self.aux_stream_list = None + try: + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self._run_parallel_input_projections(hidden_states) + ) + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + finally: + self.aux_stream_list = aux_streams + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + if self._enable_csa_multi_stream(): + # The ROCm override consumes these sentinels inside the capture + # boundary, moving the stream fan-out ahead of the projections. + self._prepare_and_attn_fn( + hidden_states, + None, + None, + None, + None, + None, + None, + positions, + o_padded, + ) + else: + self._run_sequential_pipeline(hidden_states, positions, o_padded) + + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor | None, + kv: torch.Tensor | None, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor | None, + indexer_kv_score: torch.Tensor | None, + indexer_weights: torch.Tensor | None, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Run the ROCm CSA fork/join inside the graph capture boundary.""" + aux_streams = self.aux_stream_list + # The sequential pipeline disables aux_stream_list before calling + # back with real projection inputs; aux_streams is None ends that + # recursion here. + if aux_streams is None: + saved_streams = self.aux_stream_list + self.aux_stream_list = None + try: + super()._prepare_and_attn( + hidden_states, + cast(torch.Tensor, qr), + cast(torch.Tensor, kv), + qr_scale, + cast(torch.Tensor, kv_score), + cast(torch.Tensor, indexer_kv_score), + cast(torch.Tensor, indexer_weights), + positions, + o_padded, + ) + finally: + self.aux_stream_list = saved_streams + return + + # Re-check: forward's gate ran inside a captured segment that + # _prepare_and_attn_eager (MRV1) then broke, making this region eager. + if not self._enable_csa_multi_stream(): + self._run_sequential_pipeline(hidden_states, positions, o_padded) + return + + indexer = self.indexer + compressor = self.compressor + assert indexer is not None and compressor is not None + + def default_chain(): + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv) + q = self._wq_b_proj(qr_out, qr_scale_out).view( + -1, self.n_local_heads, self.head_dim + ) + attn_metadata = get_forward_context().attn_metadata + q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata) + return q, qr_out, qr_scale_out, kv_out + + def main_compressor_chain() -> None: + score = torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + compressor(score, positions, self.rotary_emb) + + def indexer_compressor_chain() -> None: + score = torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + indexer.compressor(score, positions, self.indexer_rotary_emb) + + (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel( + default_chain, + [main_compressor_chain, indexer_compressor_chain], + self.ln_events[0], + self.ln_events[1:3], + aux_streams[:2], + enable=True, + ) + + indexer_weights_out, _ = indexer.weights_proj(hidden_states) + # The indexer compressor already ran on aux stream 1; build queries only. + index_q, index_q_scale, weights = indexer( + hidden_states, + qr_out, + None, + indexer_weights_out, + positions, + self.indexer_rotary_emb, + qr_scale_out, + skip_compressor=True, + ) + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + weights, + q, + kv_out, + positions, + o_padded, + ) + @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return num_heads diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index d32a48cbe43f..9c779809650c 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -318,7 +318,6 @@ def __init__( # graph and MRV1 produces garbage (#51430). self._prepare_and_attn_fn = self._prepare_and_attn_eager - # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins @@ -584,7 +583,7 @@ def project_query_and_cache_kv() -> torch.Tensor: # Keep Q projection and KV insertion on the default stream. The indexer # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the - # indexer. ROCm runs the same work sequentially without aux streams. + # indexer. if indexer is not None: assert compressor is not None q, (indexer_inputs, _) = execute_in_parallel( @@ -671,7 +670,6 @@ def _run_parallel_input_projections( # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs # on aux streams 0..2 when their owning module exists. ln_events[0] # is the fan-out start event; ln_events[1..3] are per-aux done events. - # On ROCm, aux_streams is None and execute_in_parallel runs serially. aux_fns: list[Callable[[], Any] | None] = [None, None, None] if self.compressor is not None: @@ -1055,11 +1053,12 @@ def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: compressor = self.compressor @@ -1072,7 +1071,8 @@ def forward( ): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache - compressor(compressed_kv_score, positions, rotary_emb) + if not skip_compressor: + compressor(compressed_kv_score, positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( indexer_metadata.num_decode_tokens @@ -1102,15 +1102,18 @@ def wq_b_and_q_quant(): use_fp4=self.use_fp4_kv, ) - # compressor returns None and writes K to the indexer KV cache; the - # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), _ = maybe_execute_in_parallel( - wq_b_and_q_quant, - lambda: compressor(compressed_kv_score, positions, rotary_emb), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) + if not skip_compressor: + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), _ = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + else: + q_quant, weights = wq_b_and_q_quant() if isinstance(q_quant, tuple): q, q_scale = q_quant else: diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index d51e589652f4..a2c3b00493e6 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -198,11 +198,12 @@ def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """CPU override: no aux streams, so wq_b_and_q_quant and the compressor run straight-line instead of through From a5f6f61a8a8eb9d8e839cf7e322b23cc8ae406e5 Mon Sep 17 00:00:00 2001 From: Ganesh R Date: Mon, 14 Sep 2026 17:21:38 +0530 Subject: [PATCH 179/774] [CI][CPU] Add speculative-decoding coverage to CPU CI (#54934) Signed-off-by: R Signed-off-by: jiang1.li Co-authored-by: Cursor Co-authored-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 46 +++++ .../scripts/hardware_ci/run-cpu-test.sh | 7 +- docs/features/README.md | 2 +- tests/v1/e2e/test_cpu_spec_decode.py | 167 ++++++++++++++++++ 4 files changed, 219 insertions(+), 3 deletions(-) create mode 100644 tests/v1/e2e/test_cpu_spec_decode.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 769e3d4adc3a..b50a082727f5 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -52,6 +52,52 @@ steps: tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp" parallelism: 2 +# Speculative decoding under a CPU platform. The unit tests already pass on +# CPU but no CPU job runs them; the e2e tests gate the engine path itself, +# which unit tests cannot reach. +- label: CPU-Spec Decode Tests + key: cpu-spec-decode-tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/cpu/ + - vllm/v1/worker/cpu_model_runner.py + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/attention/backends/cpu_attn.py + - vllm/config/speculative.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/model_executor/models/llama_eagle3.py + - vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py + - vllm/model_executor/layers/mamba/mamba_mixer2.py + - vllm/model_executor/layers/mamba/ops/cpu/ + - tests/v1/e2e/test_cpu_spec_decode.py + - tests/v1/spec_decode/ + - tests/v1/worker/ + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + pytest -x -v -s \ + tests/v1/worker/test_gpu_autoregressive_speculator.py \ + tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \ + tests/v1/worker/test_gpu_batch_ordering.py \ + tests/v1/worker/test_gpu_rejection_sampler_chunking.py \ + tests/v1/worker/test_eagle3_aux_hidden_states_pp.py \ + tests/v1/worker/test_spec_decode_embed_sharing_pp.py \ + tests/v1/spec_decode/test_adaptive_verification.py \ + tests/v1/spec_decode/test_dflash2.py \ + tests/v1/spec_decode/test_dflash_causality.py \ + tests/v1/spec_decode/test_eagle_draft_attn_metadata.py \ + tests/v1/spec_decode/test_draft_attention_backend_override.py \ + tests/v1/spec_decode/test_draft_moe_backend_override.py \ + tests/v1/spec_decode/test_ngram.py \ + tests/v1/spec_decode/test_llm_base_proposer_sampling.py + unset VLLM_CPU_CI_ENV + pytest -x -v -s tests/v1/e2e/test_cpu_spec_decode.py" + # Note: SDE can't be downloaded from CI host because of AWS WAF # - label: CPU-Compatibility Tests # depends_on: [] diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index e47cc22c8fea..74c7749d50ef 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -83,8 +83,11 @@ rm -f "$build_log" # Run the image, setting --shm-size=4g for tensor parallel. Default to # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network; -# retry once online if the cache is missing something. -OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)' +# retry once online if the cache is missing something. vllm's get_config() +# wraps the raw huggingface_hub offline-mode errors in a generic ValueError +# (see transformers_utils/config.py), so match that message too or the +# fallback never triggers for a config-lookup cache miss. +OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified' run_test() { local hf_offline=$1 docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ diff --git a/docs/features/README.md b/docs/features/README.md index 28362f401477..1df8960ee354 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -64,7 +64,7 @@ th:not(:first-child) { | [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | -| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | +| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) | | [pooling](../models/pooling_models/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | enc-dec | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py new file mode 100644 index 000000000000..2d34c1165625 --- /dev/null +++ b/tests/v1/e2e/test_cpu_spec_decode.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU speculative-decoding end-to-end correctness.""" + +import os + +import pytest + +from tests.models.utils import check_logprobs_close +from tests.v1.e2e.spec_decode.utils import compute_acceptance_len +from vllm import LLM, SamplingParams +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +pytestmark = pytest.mark.cpu_model + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +# Bound the KV cache so the run does not scale with host memory. +os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1") + +MODEL = "Qwen/Qwen3-0.6B" +NUM_LOGPROBS = 5 +SP = SamplingParams(max_tokens=48, temperature=0, logprobs=NUM_LOGPROBS) + +# Repetitive prompts so the ngram proposer finds matches and actually drafts; +# without hits the run is indistinguishable from plain decoding. +PROMPTS = [ + "The capital of France is Paris. The capital of Italy is Rome. " + "The capital of Spain is Madrid. The capital of France is", + "def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n\n" + "def mul(a, b):\n return", + "one two three four five one two three four five one two three four", +] + + +def _make_llm(**overrides) -> LLM: + base = dict( + model=MODEL, + dtype="bfloat16", + max_model_len=1024, + disable_log_stats=False, + ) + base.update(overrides) + return LLM(**base) + + +def _tuples(outputs) -> list[tuple[list[int], str, object]]: + return [ + (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs) + for o in outputs + ] + + +@pytest.fixture(scope="module") +def baseline_refs(): + llm = _make_llm() + refs = _tuples(llm.generate(PROMPTS, SP)) + del llm + return refs + + +def _spec_metric(llm: LLM, name: str) -> float: + for metric in llm.get_metrics(): + if metric.name == name: + return float(getattr(metric, "value", 0.0)) + return 0.0 + + +def test_ngram_spec_decode_matches_baseline(baseline_refs): + """Greedy output must be unchanged by speculative decoding. + + The drafted-token check keeps the test from passing vacuously when the + proposer silently produces nothing. + """ + llm = _make_llm( + speculative_config={ + "method": "ngram", + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": 3, + }, + ) + got = _tuples(llm.generate(PROMPTS, SP)) + drafted = _spec_metric(llm, "vllm:spec_decode_num_draft_tokens") + accepted = _spec_metric(llm, "vllm:spec_decode_num_accepted_tokens") + del llm + + assert drafted > 0, "ngram proposer drafted no tokens; the test would be vacuous" + assert accepted > 0, f"no draft tokens accepted out of {drafted} drafted" + + check_logprobs_close( + outputs_0_lst=baseline_refs, + outputs_1_lst=got, + name_0="no_spec_decode", + name_1="ngram_spec_decode", + ) + + +SMOKE_PROMPTS = [ + "The capital of France is", + "2 + 2 equals", + "In one word, the color of the sky is", + "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:", +] +SMOKE_SP = SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True) + +# Methods routed to the V2 runner, using the same target/draft pairs as the GPU +# suite. MTP needs no separate draft checkpoint: its speculator layer is +# trained jointly with, and shipped inside, the target model's own weights. +SMOKE_CONFIGS = [ + pytest.param( + "meta-llama/Llama-3.2-1B-Instruct", + { + "method": "eagle3", + "model": "nm-testing/Llama3_2_1B_speculator.eagle3", + "num_speculative_tokens": 3, + }, + id="eagle3", + ), + pytest.param( + "Qwen/Qwen3-4B", + { + "method": "dflash", + "model": "z-lab/Qwen3-4B-DFlash-b16", + # The checkpoint allows 16, but accepted tokens plateau around 64 + # regardless, so a larger draft only burns CPU on rejected tokens. + "num_speculative_tokens": 4, + }, + id="dflash", + ), + pytest.param( + "Qwen/Qwen3.5-0.8B-Base", + { + "method": "mtp", + "num_speculative_tokens": 3, + }, + id="mtp", + ), +] + +# Measured 1.85 (eagle3), 1.99 (dflash), 2.29 (mtp) on CPU; this floor leaves +# room for numeric drift while still failing if speculation stops paying off. +MIN_ACCEPTANCE_LEN = 1.5 + + +@pytest.mark.skipif( + not HAS_TRITON, reason="the V2 runner needs triton-cpu, which CI builds" +) +@pytest.mark.parametrize("model,speculative_config", SMOKE_CONFIGS) +def test_v2_speculator_smoke(model: str, speculative_config: dict): + """Gates that drafts are accepted often enough for speculation to pay off.""" + llm = _make_llm(model=model, speculative_config=speculative_config) + try: + assert llm.llm_engine.vllm_config.use_v2_model_runner, ( + f"{speculative_config['method']} fell back to the V1 runner; " + "this test is meant to cover the V2 path" + ) + llm.generate(SMOKE_PROMPTS, SMOKE_SP) + acceptance = compute_acceptance_len(llm.get_metrics()) + assert acceptance >= MIN_ACCEPTANCE_LEN, ( + f"acceptance length {acceptance:.3f} below {MIN_ACCEPTANCE_LEN}; " + "drafts are being rejected far more than expected" + ) + finally: + del llm From c612e2bff0d045a5b2304371793985c330fba436 Mon Sep 17 00:00:00 2001 From: AlexHuang Date: Mon, 14 Sep 2026 20:02:45 +0800 Subject: [PATCH 180/774] [Bugfix][KV Offload] Reuse in-flight async lookup probes (#55823) Signed-off-by: Alex Co-authored-by: OpenAI Codex Co-authored-by: Or Ozeri --- .../kv_offload/tiering/test_async_lookup.py | 76 +++++++++++++++++-- vllm/v1/kv_offload/tiering/async_lookup.py | 57 ++++++++++---- 2 files changed, 112 insertions(+), 21 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index ff5cbcd5dd44..3649c8df0488 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -8,7 +8,7 @@ import pytest from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key -from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager +from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager, LookupPhase def _key(i: int) -> OffloadKey: @@ -46,10 +46,13 @@ def test_new_key_returns_none(self): def test_found_key_returns_true(self): mgr = InMemoryLookupManager(existing_keys={_key(1)}) assert mgr.lookup(_key(1), _ctx()) is None + assert mgr._lookup_state[_key(1)].phase is LookupPhase.PENDING mgr.flush() + assert mgr._lookup_state[_key(1)].phase is LookupPhase.IN_FLIGHT mgr._results_ready.wait() mgr._results_ready.clear() assert mgr.lookup(_key(1), _ctx()) is True + assert mgr._lookup_state[_key(1)].phase is LookupPhase.RESOLVED mgr.shutdown() def test_not_found_key_returns_false(self): @@ -106,6 +109,68 @@ def test_cleanup_preserves_shared_entries(self): assert _key(1) not in mgr._lookup_state mgr.shutdown() + def test_cleanup_reuses_in_flight_probe(self, monkeypatch: pytest.MonkeyPatch): + """A replacement request shares the in-flight probe and its verdict.""" + key = _key(1) + mgr = InMemoryLookupManager(existing_keys={key}) + ctx_b = _ctx("req_b") + probe_started = threading.Event() + release_probe = threading.Event() + batch_lookup = mgr.batch_lookup + + def blocking_lookup(keys, req_context): + probe_started.set() + if not release_probe.wait(timeout=5): + raise TimeoutError("Test did not release the backend probe") + return batch_lookup(keys, req_context) + + monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup) + try: + assert mgr.lookup(key, _ctx("req_a")) is None + mgr.flush() + assert probe_started.wait(timeout=5) + assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT + mgr.cleanup("req_a") + assert mgr.lookup(key, ctx_b) is None + assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT + mgr.flush() + + release_probe.set() + batch = mgr._pending_results.get(timeout=5) + mgr._pending_results.put(batch) + replacement_result = mgr.lookup(key, ctx_b) + finally: + release_probe.set() + mgr.shutdown() + + assert mgr.batch_lookup_calls == 1 + assert replacement_result is True + + @pytest.mark.parametrize( + "reclaim_at_shutdown", [False, True], ids=["flush", "shutdown"] + ) + def test_unclaimed_probe_reclaimed_without_lookup(self, reclaim_at_shutdown: bool): + """A completed orphan is released by flush or shutdown without a lookup.""" + key = _key(1) + mgr = InMemoryLookupManager(existing_keys={key}) + try: + mgr.lookup(key, _ctx("req_a")) + mgr.flush() + mgr.cleanup("req_a") + assert key in mgr._lookup_state + assert not mgr._req_keys + + if not reclaim_at_shutdown: + batch = mgr._pending_results.get(timeout=5) + mgr._pending_results.put(batch) + mgr.flush() + assert key not in mgr._lookup_state + finally: + mgr.shutdown() + + assert not mgr._lookup_state + assert mgr._pending_results.empty() + def test_stale_result_ignored_after_cleanup_and_key_reuse(self): key = _key(1) mgr = InMemoryLookupManager() @@ -119,11 +184,12 @@ def test_stale_result_ignored_after_cleanup_and_key_reuse(self): generation = mgr._lookup_state[key].generation assert generation != stale_generation + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + current_result = mgr._pending_results.get(timeout=5) mgr._pending_results.put([(key, stale_generation, True)]) - mgr.drain_results() - assert mgr.lookup(key, ctx_b) is None - - mgr._pending_results.put([(key, generation, False)]) + mgr._pending_results.put(current_result) mgr.drain_results() assert mgr.lookup(key, ctx_b) is False mgr.shutdown() diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index f111f233b613..c08c63d01875 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -27,8 +27,10 @@ flush() is called once per step from the tier's on_schedule_end(), posting the entire batch as a single queue item so the background thread sees one batch per step. -drain_results() is called before any lookup() calls in the same step, so -lookup() is a pure OrderedDict operation. +Results are drained on the first lookup after each flush, at flush(), and +after worker shutdown. In-flight lookups with no remaining request references +are retained until their results are drained, allowing new requests to share +the same probe. """ import queue @@ -36,6 +38,7 @@ from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass, field +from enum import Enum, auto from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadKey, ReqContext @@ -43,10 +46,21 @@ logger = init_logger(__name__) +class LookupPhase(Enum): + """Lifecycle phase of a lookup probe.""" + + PENDING = auto() # Accumulated in _lookup_batch, but not yet submitted. + IN_FLIGHT = auto() # Submitted to the worker, but not yet resolved. + RESOLVED = auto() # A final verdict has been applied to the state. + + @dataclass(slots=True) class LookupState: generation: int - result: bool | None = None # True (found), False (not found), None + phase: LookupPhase = LookupPhase.PENDING + # None while pending/in flight; True if the key exists; False if absent or + # explicitly marked missing after a failed load. + result: bool | None = None request_ids: set[str] = field(default_factory=set) # requests asking for the lookup @@ -155,24 +169,28 @@ def flush(self) -> None: Called once per step from on_schedule_end() after all lookup() calls are done. The worker receives the full batch and processes it during the model-execution window, maximising time available before the next - step's drain_results(). Safe to call with an empty batch (no-op). + step's drain_results(). Also drains completed lookups when there + are no new keys to submit. """ + self.drain_results() self._need_to_drain = True batch = self._lookup_batch self._lookup_batch = [] - batch = [ - (key, req_context, generation) - for key, req_context, generation in batch - if (state := self._lookup_state.get(key)) is not None - and state.generation == generation - ] - if batch: - self._lookup_queue.put(batch) + in_flight_batch = [] + for key, req_context, generation in batch: + state = self._lookup_state.get(key) + if state is None or state.generation != generation: + continue + assert state.phase is LookupPhase.PENDING + state.phase = LookupPhase.IN_FLIGHT + in_flight_batch.append((key, req_context, generation)) + if in_flight_batch: + self._lookup_queue.put(in_flight_batch) def drain_results(self) -> None: """Apply pending worker results to _lookup_state. - Called from lookup() before checking state. + Called from lookup(), flush(), and shutdown() on the scheduler thread. """ while True: try: @@ -183,6 +201,10 @@ def drain_results(self) -> None: state = self._lookup_state.get(key) if state is None or state.generation != generation: continue + if not state.request_ids: + del self._lookup_state[key] + continue + assert state.phase is LookupPhase.IN_FLIGHT # Each lookup generation is enqueued exactly once. A matching # generation must not receive a second result; stale # generations were discarded above. @@ -192,6 +214,7 @@ def drain_results(self) -> None: "failed-load livelock" ) state.result = result + state.phase = LookupPhase.RESOLVED def mark_miss(self, keys: Collection[OffloadKey]) -> None: """Force the cached verdict for ``keys`` to False after a failed load, so @@ -201,9 +224,10 @@ def mark_miss(self, keys: Collection[OffloadKey]) -> None: state = self._lookup_state.get(key) if state is not None: state.result = False + state.phase = LookupPhase.RESOLVED def cleanup(self, req_id: str) -> None: - """Remove entries no longer needed by any active request. + """Release request references, retaining in-flight lookups. Called from the tier's on_request_finished(). Uses the reverse index to visit only keys associated with this request. @@ -211,13 +235,14 @@ def cleanup(self, req_id: str) -> None: for key in self._req_keys.pop(req_id, ()): state = self._lookup_state[key] state.request_ids.discard(req_id) - if not state.request_ids: + if not state.request_ids and state.phase is not LookupPhase.IN_FLIGHT: del self._lookup_state[key] def shutdown(self) -> None: - """Stop the worker thread.""" + """Stop the worker thread and drain completed lookups.""" self._lookup_queue.put(None) # unblock _worker from _lookup_queue.get() self._thread.join() + self.drain_results() # ------------------------------------------------------------------ # Internal helpers From 3f55ad2f073a7d1f573db5a5a3fe3dc9c688ddee Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 14 Sep 2026 20:16:25 +0800 Subject: [PATCH 181/774] [Qwen3.8-Flash-Next] Fuse PLE residual and QSA output gate (#55309) Signed-off-by: Canlin Signed-off-by: zjy0516 Co-authored-by: zjy0516 Co-authored-by: Kimi Code --- tests/models/qwen4_exp/test_ple.py | 15 ++++-- tests/models/qwen4_exp/test_qsa_reference.py | 3 ++ vllm/models/qwen4_exp/nvidia/model.py | 2 +- vllm/models/qwen4_exp/nvidia/ops/ple.py | 17 ++++++- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 48 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/ple_layer.py | 30 ++++++++++-- vllm/models/qwen4_exp/nvidia/qsa.py | 8 +++- 7 files changed, 111 insertions(+), 12 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 020e1cc8fc78..e7d0e72d8a3e 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1610,6 +1610,12 @@ def test_fused_conv_correctness( dtype=torch.bfloat16, generator=rng, ) + outer_residual = torch.randn( + inputs.shape, + device=device, + dtype=torch.bfloat16, + generator=rng, + ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() residual_reference = residual.clone() @@ -1617,6 +1623,7 @@ def test_fused_conv_correctness( module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, + outer_residual=outer_residual, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1630,15 +1637,17 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) + residual_reference = outer_residual + residual_reference - torch.testing.assert_close( - residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2 + assert torch.equal( + residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] ) assert torch.equal(conv_state, state_reference) assert torch.equal(conv_state[NULL_BLOCK_ID], null_state) if case.graph_padding: assert torch.equal( - residual_kernel[num_real_tokens:], residual[num_real_tokens:] + residual_kernel[num_real_tokens:], + (outer_residual + residual)[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 9259972ac9eb..16294a30d769 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -944,6 +944,7 @@ def test_qsa_sparse_paged_attention_correctness( q = torch.randn( num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) + output_gate = torch.randn_like(q) kv_cache = torch.randn( num_cache_blocks, page_size, @@ -1023,6 +1024,7 @@ def test_qsa_sparse_paged_attention_correctness( block_table, token_to_req, use_prefill_config=use_prefill_config, + output_gate=output_gate, ) expected = _qsa_sparse_paged_attention_reference( q, @@ -1033,6 +1035,7 @@ def test_qsa_sparse_paged_attention_correctness( token_to_req, scale, ) + expected = expected * torch.sigmoid(output_gate) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 28c3ea69d11b..54ee6f02b4f5 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -297,7 +297,7 @@ def forward( if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - hidden_states = hidden_states + self.ple( + hidden_states = self.ple( hidden_states, input_ids, query_start_loc, diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/nvidia/ops/ple.py index d76e2705a7bf..bc4c99957d90 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/ple.py +++ b/vllm/models/qwen4_exp/nvidia/ops/ple.py @@ -322,6 +322,7 @@ def _ple_conv_kernel( state_ptr, w_ptr, residual_ptr, + outer_residual_ptr, state_idx_ptr, qsl_ptr, num_acc_ptr, @@ -442,11 +443,21 @@ def _ple_conv_kernel( mask=c_mask, other=0.0, ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( residual_ptr + output_t * C + c_offs, - residual + conv_output, + ple_output, mask=c_mask, ) @@ -564,6 +575,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, + outer_residual: torch.Tensor, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -573,7 +585,7 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short-convolution output to ``residual`` and update its state.""" + """Add short convolution and the outer residual; update state.""" BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -624,6 +636,7 @@ def ple_conv( conv_state, conv_weights, residual, + outer_residual, state_indices, query_start_loc, num_accepted_tokens, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 709acbde00a4..d3986ba38ae0 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -24,6 +24,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_q_row, stride_q_head, stride_k_block, @@ -36,6 +37,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_table_req, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, num_cache_blocks, num_requests, @@ -151,6 +154,18 @@ def _qsa_sparse_paged_gqa_splitk_kernel( ) output_mask = head_offsets[:, None] < GROUP_SIZE if NUM_SPLITS == 1: + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + normalized_output = normalized_output.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + (first_head + head_offsets[:, None]) * stride_output_gate_head + + dim_offsets[None, :], + mask=output_mask, + other=0.0, + ).to(tl.float32) + normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row @@ -186,8 +201,11 @@ def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, HEAD_DIM: tl.constexpr, NUM_QUERY_HEADS: tl.constexpr, @@ -219,6 +237,16 @@ def _qsa_merge_splitk_kernel( ) merged = tl.sum(partial_output * weights[:, None], axis=0) merged = tl.where(denominator > 0, merged / denominator, 0.0) + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + merged = merged.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + head * stride_output_gate_head + + dim_offsets + ).to(tl.float32) + merged = merged.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, merged, @@ -455,6 +483,8 @@ def qsa_sparse_paged_attention( token_to_req: torch.Tensor, use_prefill_config: bool, out: torch.Tensor | None = None, + *, + output_gate: torch.Tensor, ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 K/V caches. @@ -496,6 +526,8 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse output must match its query") assert out.dtype == q.dtype and out.device == q.device assert out.stride(2) == 1 + assert output_gate.is_contiguous() + output_gate_view = output_gate.view_as(q) if not q.shape[0]: return out @@ -533,6 +565,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, q.stride(0), q.stride(1), k_cache.stride(0), @@ -545,6 +578,8 @@ def qsa_sparse_paged_attention( block_table.stride(0), out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], k_cache.shape[0], block_table.shape[0], @@ -568,8 +603,11 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], HEAD_DIM=q.shape[2], NUM_QUERY_HEADS=q.shape[1], @@ -632,6 +670,10 @@ def warmup_qsa_sparse_paged_attention( output_ptr = TritonWarmupTensor( torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) ) + # The output gate is mandatory at runtime; warm the gated specialization. + output_gate_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) head_stride = head_dim row_stride = num_query_heads * head_dim num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0]) @@ -659,6 +701,7 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, row_stride, head_stride, key_cache.stride(0), @@ -671,6 +714,8 @@ def warmup_qsa_sparse_paged_attention( block_table.stride(0), row_stride, head_stride, + row_stride, + head_stride, num_rows, num_cache_blocks, num_requests, @@ -693,6 +738,9 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, + row_stride, + head_stride, row_stride, head_stride, num_rows, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 26622357ca50..2ec8790e8d3b 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -195,6 +195,7 @@ def _short_conv_dilated_dispatch( self, inputs: torch.Tensor, residual: torch.Tensor, + outer_residual: torch.Tensor, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -209,6 +210,7 @@ def _short_conv_dilated_dispatch( has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -235,6 +237,7 @@ def _short_conv_dilated_dispatch( conv_state=conv_state, conv_weights=conv_weights, state_indices=spec_state_indices, + outer_residual=outer_residual, mode="spec", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -259,11 +262,17 @@ def _short_conv_dilated_dispatch( residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: inputs_d = inputs_p = inputs residual_d = residual_p = residual + outer_residual_d = outer_residual_p = outer_residual token_indices_d, token_indices_p = torch.split( non_spec_token_indices, [num_decode_tokens, num_prefill_tokens], @@ -277,6 +286,7 @@ def _short_conv_dilated_dispatch( conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_d, + outer_residual=outer_residual_d, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -298,6 +308,7 @@ def _short_conv_dilated_dispatch( conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_p, + outer_residual=outer_residual_p, mode="prefill", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -316,6 +327,7 @@ def _short_conv_dilated_dispatch( conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices[:num_decode_rows], + outer_residual=outer_residual, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -324,12 +336,18 @@ def _short_conv_dilated_dispatch( # State routing consumes the current request metadata on every replay. @eager_break_during_capture - def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None: + def _short_conv( + self, + inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, + ) -> None: forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata - # Profiling omits all metadata or this Mamba entry. The residual - # already contains the gated output, so short convolution is a no-op. + # Profiling omits all metadata or this Mamba entry. Short convolution + # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -340,6 +358,7 @@ def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None: layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -367,6 +386,7 @@ def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None: self._short_conv_dilated_dispatch( inputs=inputs, residual=residual, + outer_residual=outer_residual, metadata=layer_attn_metadata, conv_state=conv_state, conv_weights=conv_weights.to(dtype=inputs.dtype), @@ -404,7 +424,9 @@ def forward( self.norm_conv.weight, self.norm_key.eps, ) - self._short_conv(conv_input, gated_output) + # State routing depends on runtime request metadata and remains outside + # the graph; short convolution accumulates into gated_output. + self._short_conv(conv_input, gated_output, hidden_states) return gated_output diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 90449d5706fb..1c2808a6c47f 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -121,6 +121,7 @@ def forward_qsa( output: torch.Tensor, token_to_req: torch.Tensor, use_prefill_config: bool, + output_gate: torch.Tensor, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: @@ -157,6 +158,7 @@ def forward_qsa( token_to_req, use_prefill_config, output[:num_tokens], + output_gate=output_gate[:num_tokens], ) return output @@ -353,6 +355,7 @@ def _run_qsa( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor, ) -> None: metadata = get_forward_context().attn_metadata if isinstance(metadata, list): @@ -396,6 +399,7 @@ def _run_qsa( output, token_to_req=side_metadata.token_to_req, use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len, + output_gate=output_gate, ) def forward( @@ -405,6 +409,7 @@ def forward( ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) + assert gate is not None num_tokens = hidden_states.shape[0] query = q.view(num_tokens, self.num_heads, self.head_dim) key = k.view(num_tokens, self.num_kv_heads, self.head_dim) @@ -419,10 +424,9 @@ def forward( key, value, attn_output, + gate, ) flat_output = attn_output.view(num_tokens, -1) - if gate is not None: - flat_output = flat_output * torch.sigmoid(gate) output, _ = self.o_proj(flat_output) return output From 1d0d1081c409176a0e7c24742f7c60b9943657aa Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:52:56 +0200 Subject: [PATCH 182/774] [ROCm][Connector] SWA+HMA-support in MoRI-IO connector (Gemma4) (#53721) Signed-off-by: simondanielsson --- .../unit/test_moriio_connector.py | 232 +++++++++++++++++- .../kv_connector/unit/test_moriio_routing.py | 7 +- .../kv_connector/v1/moriio/moriio_common.py | 7 +- .../v1/moriio/moriio_connector.py | 208 +++++++++++++--- 4 files changed, 408 insertions(+), 46 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index d169779c5df1..836f9407bc66 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -39,10 +39,12 @@ make_zmq_path, ) from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + SlidingWindowSpec, compute_layer_kv_cache_shape_bytes, ) @@ -338,7 +340,7 @@ def test_write_mode_saves_local_block_ids(): req_meta = kv_connector_metadata.reqs_to_save[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -393,7 +395,7 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids(): req_meta = kv_connector_metadata.reqs_to_save[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -432,7 +434,7 @@ def test_read_mode_loads_remote_block_ids(): ].req_to_blocks[request_id] # Set remote block ids to be fetched. - request.kv_transfer_params["remote_block_ids"] = block_list + request.kv_transfer_params["remote_block_ids"] = [block_list] # Remote Prefill, triggers MoRIIOConnectorMetadata. @@ -457,7 +459,7 @@ def test_read_mode_loads_remote_block_ids(): req_meta = kv_connector_metadata.reqs_to_recv[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -699,3 +701,225 @@ def test_resolve_host_ip_prefers_extra_config(): fallback = get_ip() assert resolve_host_ip({}) == fallback assert resolve_host_ip({"host_ip": ""}) == fallback + + +def _make_hybrid_kv_cache_config() -> KVCacheConfig: + """One full-attention group + one sliding-window group (Gemma-like).""" + full_spec = FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ) + sw_spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=32, + ) + num_blocks = 2 + page = full_spec.page_size_bytes + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=2 * num_blocks * page, + layers=["full0", "sw0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec), + KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec), + ], + ) + + +def _read_scheduler( + kv_cache_config: KVCacheConfig, disable_hma: bool = False +) -> MoRIIOConnectorScheduler: + vllm_config = create_vllm_config(role="kv_producer", read_mode=True) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + assert connector.connector_scheduler is not None + return connector.connector_scheduler + + +def test_hma_blocks_per_sw_two_groups(): + """A Full + SlidingWindow config runs with HMA and computes block budgets + correctly""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + assert scheduler._is_hma_required is True + # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention. + assert scheduler.blocks_per_sw == [0, 3] + + +@pytest.mark.parametrize( + "swa_enabled, disable_hma, expected_is_hma", + [ + (True, False, True), # sliding-window group present, HMA enabled + (True, True, False), # sliding-window group present but HMA disabled + (False, False, False), # full-attention only, HMA not needed + ], +) +def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): + """_is_hma_required tracks both the KV cache groups and the + --disable-hybrid-kv-cache-manager flag. When HMA is off, + get_exchange_clipped_blocks must be a no-op.""" + config = ( + _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config() + ) + scheduler = _read_scheduler(config, disable_hma=disable_hma) + assert scheduler._is_hma_required is expected_is_hma + if not expected_is_hma: + blocks = [[1, 2, 3, 4, 5]] + assert scheduler.get_exchange_clipped_blocks(blocks) == blocks + + +def test_non_sliding_window_hybrid_is_rejected(): + """A hybrid group that is not sliding-window (e.g. chunked-local + attention) must fail closed rather than be silently mistransferred.""" + full_spec = FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ) + local_spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + attention_chunk_size=32, + ) + num_blocks = 2 + page = full_spec.page_size_bytes + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=2 * num_blocks * page, + layers=["full0", "local0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec), + KVCacheGroupSpec(layer_names=["local0"], kv_cache_spec=local_spec), + ], + ) + with pytest.raises(NotImplementedError, match="sliding-window hybrid"): + _read_scheduler(config) + + +def test_token_count_basis_uses_full_attention_group(): + """Chunked-prefill token counting must use an unclipped full-attention + group (blocks_per_sw == 0), not a clipped sliding-window group.""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + # Group 0 is the full-attention group + assert scheduler._full_attn_group_idx == 0 + assert scheduler._full_attn_block_size == 16 + + +def test_get_exchange_clipped_blocks_clips_only_sw_group(): + """get_exchange_clipped_blocks keeps the full attn group intact and clips the + sliding-window group to its window tail.""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + full = [10, 11, 12, 13, 14] + sw = [20, 21, 22, 23, 24] + clipped = scheduler.get_exchange_clipped_blocks([full, sw]) + assert clipped[0] == full + assert clipped[1] == [22, 23, 24] + + +def test_metadata_hma_block_ids_preserved_per_group(): + """add_new_req stores per-group (BlockIds) block lists unchanged for both + read and write, so the hybrid group structure is retained.""" + metadata = MoRIIOConnectorMetadata() + + # Assume: + # - Full-attention group (6 blocks) + + # - sliding-window group already clipped to its window tail (3 blocks). + fa_blocks = [0, 1, 2, 3, 4, 5] + sw_blocks = [10, 11, 12] + local_block_ids = [fa_blocks, sw_blocks] + remote_block_ids = [[100, 101, 102, 103, 104, 105], [200, 201, 202]] + + base_params = { + "remote_engine_id": "remote-engine", + "remote_host": "127.0.0.1", + "remote_handshake_port": 6301, + "remote_notify_port": 61005, + "remote_block_ids": remote_block_ids, + } + + # Read mode: both local and remote ids stay per-group. + metadata.add_new_req( + request_id="recv-req", + local_block_ids=local_block_ids, + kv_transfer_params={**base_params, "transfer_id": "recv-req"}, + ) + recv_meta = metadata.reqs_to_recv["recv-req"] + assert recv_meta.local_block_ids == [fa_blocks, sw_blocks] + assert recv_meta.remote_block_ids == remote_block_ids + + # Write mode: the decode peer allocates its own blocks, so #remote_block_ids may + # be empty, Local group structure is kept. + metadata.add_new_req( + request_id="save-req", + local_block_ids=local_block_ids, + kv_transfer_params={ + **base_params, + "transfer_id": "save-req", + "remote_block_ids": [], + }, + write_mode=True, + ) + save_meta = metadata.reqs_to_save["save-req"] + assert save_meta.local_block_ids == [fa_blocks, sw_blocks] + assert save_meta.remote_block_ids == [] + + +def test_single_group_path_unchanged(): + """A full attn config does not enable HMA and never clips blocks.""" + scheduler = _read_scheduler(_make_test_kv_cache_config()) + assert scheduler._is_hma_required is False + assert scheduler.blocks_per_sw == [0] + blocks = [[1, 2, 3, 4, 5]] + assert scheduler.get_exchange_clipped_blocks(blocks) == blocks + + +def test_hybrid_write_mode_rejected(): + """Hybrid KV cache groups are unsupported in WRITE mode and fail closed.""" + vllm_config = create_vllm_config(role="kv_producer", read_mode=False) + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError), + ): + MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_hybrid_kv_cache_config(), + ) + + +def test_worker_layer_to_group_routing(mock_parallel_groups): + """The worker maps every layer to its KV cache group correctly.""" + vllm_config = create_vllm_config(role="kv_consumer", read_mode=True) + # Building the worker directly bypasses MoRIIOConnector._set_port_defaults, + # so provide the ports manually. + vllm_config.kv_transfer_config.kv_connector_extra_config.update( + { + "http_port": 12346, + "handshake_port": 12347, + "notify_port": 12348, + } + ) + with set_current_vllm_config(vllm_config): + worker = FakeMoRIIOConnectorWorker( + vllm_config, + "engine0", + hand_shake_latency=0, + kv_cache_config=_make_hybrid_kv_cache_config(), + ) + assert worker.layer_to_group == {"full0": 0, "sw0": 1} diff --git a/tests/v1/kv_connector/unit/test_moriio_routing.py b/tests/v1/kv_connector/unit/test_moriio_routing.py index a9a3521bebd7..fc8b323218c2 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing.py @@ -20,12 +20,15 @@ from vllm.v1.request import RequestStatus _request_finished = MoRIIOConnectorScheduler.request_finished +_clip_blocks = MoRIIOConnectorScheduler.get_exchange_clipped_blocks def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): """Minimal prefill-side (producer) scheduler stand-in for request_finished.""" - return SimpleNamespace( + s = SimpleNamespace( is_producer=True, + _is_hma_required=False, + blocks_per_sw=[0], _global_dp_rank=global_dp_rank, engine_id="engine-abc", host_ip="10.0.0.1", @@ -39,6 +42,8 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): ) ), ) + s.get_exchange_clipped_blocks = _clip_blocks.__get__(s) + return s def _req( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index a3010203ee76..7d06b925ba29 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -14,6 +14,7 @@ import zmq from vllm.config import KVTransferConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) @@ -444,8 +445,8 @@ class ReqMeta: """Metadata for a single request.""" transfer_id: TransferId - local_block_ids: list[int] - remote_block_ids: list[int] + local_block_ids: BlockIds + remote_block_ids: BlockIds remote_host: str remote_port: int remote_handshake_port: int @@ -481,7 +482,7 @@ def __repr__(self): def add_new_req( self, request_id: ReqId, - local_block_ids: list[int], + local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], write_mode=False, ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 349cf50775ad..ecdd59fca3ef 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -18,10 +18,12 @@ import zmq from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( ROLE, @@ -67,19 +69,24 @@ ) from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import ( make_zmq_path, make_zmq_socket, ) from vllm.v1.attention.selector import get_attn_backend from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + SlidingWindowSpec, +) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus if TYPE_CHECKING: from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks - from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request logger = init_logger(__name__) @@ -189,7 +196,7 @@ def resolve_moriio_transfer_ack( return transfer_id -class MoRIIOConnector(KVConnectorBase_V1): +class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): def __init__( self, vllm_config: VllmConfig, @@ -225,7 +232,7 @@ def __init__( ) if role == KVConnectorRole.SCHEDULER: self.connector_scheduler: MoRIIOConnectorScheduler | None = ( - MoRIIOConnectorScheduler(vllm_config, self.engine_id) + MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) ) self.connector_worker: MoRIIOConnectorWorker | None = None elif role == KVConnectorRole.WORKER: @@ -283,6 +290,14 @@ def request_finished( self, request: "Request", block_ids: list[int], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + + def request_finished_all_groups( + self, + request: "Request", + block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None return self.connector_scheduler.request_finished(request, block_ids) @@ -368,7 +383,12 @@ def has_connector_metadata(self) -> bool: class MoRIIOConnectorScheduler: """Implementation of Scheduler side methods""" - def __init__(self, vllm_config: VllmConfig, engine_id: str): + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): self.vllm_config = vllm_config assert vllm_config.kv_transfer_config is not None, ( @@ -378,6 +398,66 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.block_size = vllm_config.cache_config.block_size self.engine_id: EngineId = engine_id self.mode = get_moriio_mode(self.kv_transfer_config) + + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.kv_cache_groups + ) + ) + if self._is_hma_required: + # TODO(simondanielsson): support non-sliding window hybrids + unsupported = [ + type(g.kv_cache_spec).__name__ + for g in kv_cache_config.kv_cache_groups + if not isinstance( + g.kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec) + ) + ] + if unsupported: + raise NotImplementedError( + "MoRIIO only supports sliding-window hybrid models (e.g. " + "Gemma) under the hybrid KV cache manager. Unsupported KV " + f"cache group spec(s): {sorted(set(unsupported))}. Pass " + "--disable-hybrid-kv-cache-manager to run without HMA." + ) + if self.mode == MoRIIOMode.WRITE: + # TODO(simondanielsson): support HMA in WRITE mode + raise NotImplementedError( + "MoRIIO WRITE mode does not support hybrid KV cache groups " + "(sliding-window attention). Use READ mode, or pass " + "--disable-hybrid-kv-cache-manager." + ) + + sw_sizes_tokens: list[tuple[int, int]] = [ + (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) + if isinstance(g.kv_cache_spec, SlidingWindowSpec) + else (0, self.block_size) + for g in kv_cache_config.kv_cache_groups + ] + # add 1 to conservatively account for boundary overlap eg window isn't fully + # aligned with blocks. + self.blocks_per_sw = [ + cdiv(n_tokens, block_size) + 1 if n_tokens else 0 + for n_tokens, block_size in sw_sizes_tokens + ] + # In WRITE mode and chunked prefill, we perform the write after the last chunk. + # Hence we need to track once we are in the last chunk. We do this by computing + # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups + # may be clipped and hence not reflect the full sequence length. Here we store + # g and block_size[g] for later. + self._full_attn_group_idx = 0 + self._full_attn_block_size = self.block_size + for gi, group in enumerate(kv_cache_config.kv_cache_groups): + is_full_attn = self.blocks_per_sw[gi] == 0 + if is_full_attn: + self._full_attn_group_idx = gi + self._full_attn_block_size = getattr( + group.kv_cache_spec, "block_size", self.block_size + ) + break + self.host_ip = resolve_host_ip( self.kv_transfer_config.kv_connector_extra_config ) @@ -409,14 +489,14 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): # Requests that need to start recv/send. # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. - self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {} - self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {} + self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} + self._reqs_need_save: dict[ReqId, tuple[Request, BlockIds]] = {} # Snapshot of kv_transfer_params for chunked prefill recovery. self._req_kv_params: dict[ReqId, dict] = {} # For chunked prefill, we perform layer-wise access within the final chunk. # TODO: Perform transfer at end chunk. - self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {} + self._reqs_need_pending_save: dict[ReqId, tuple[Request, BlockIds]] = {} if self.is_producer: set_role(ROLE.PRODUCER) @@ -442,6 +522,23 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + def get_exchange_clipped_blocks( + self, + block_ids: BlockIds, + ) -> BlockIds: + """Clip the block list based on the layer type. + + Behavior: + - sliding window: keep only the last window size blocks + - full attention: keep all blocks + """ + if len(block_ids) == 0 or not self._is_hma_required: + return [list(blocks) for blocks in block_ids] + return [ + blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks + for i, blocks in enumerate(block_ids) + ] + def map_request_id(self, request_id: ReqId, transfer_id: TransferId): self.transfer_id_to_request_id[transfer_id] = request_id self.request_id_to_transfer_id[request_id] = transfer_id @@ -642,7 +739,7 @@ def update_state_after_alloc( request_id = request.request_id self.map_request_id(request_id, transfer_id) if params.get("do_remote_decode"): - local_block_ids = blocks.get_block_ids()[0] + local_block_ids: BlockIds = blocks.get_block_ids() self._reqs_need_save[request.request_id] = (request, local_block_ids) # Snapshot params now so chunked-prefill build_connector_meta # can recover them on the final chunk even if the live @@ -657,17 +754,22 @@ def update_state_after_alloc( if "remote_engine_id" in params: if num_external_tokens > 0: # Get unhashed blocks to pull from remote. - local_block_ids = blocks.get_block_ids()[0] - assert len(local_block_ids) <= len(remote_block_ids) - if len(local_block_ids) != len(remote_block_ids): - local_block_ids = remote_block_ids[ - -len(local_block_ids) : - ] + local_block_ids = self.get_exchange_clipped_blocks( + blocks.get_block_ids() + ) + # Partial cache hits can lead to fewer local blocks vs + # remote blocks, but never more + assert all( + len(local) <= len(remote) + for local, remote in zip( + local_block_ids, remote_block_ids, strict=True + ) + ) else: # If remote_blocks and num_external_tokens = 0, we have # a full prefix cache hit on the D worker. We need to call # send_notify in _read_blocks to free the memory on the P. - local_block_ids = [] + local_block_ids = [[] for _ in remote_block_ids] self._reqs_need_recv[request.request_id] = ( request, @@ -818,8 +920,6 @@ def build_connector_meta( new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] if new_block_ids is not None: - block_ids = new_block_ids[0] - # TODO : hybrid attn, etc # A non-disagg request (no kv_transfer_params, e.g. smoke # test) is never registered in _reqs_need_pending_save; # indexing it unconditionally would KeyError and crash the @@ -827,12 +927,16 @@ def build_connector_meta( if req_id not in self._reqs_need_pending_save: continue req, existing_blocks = self._reqs_need_pending_save[req_id] - updated_blocks = list(existing_blocks) + (block_ids) + updated_blocks = [ + existing_blocks[g] + new_block_ids[g] + for g in range(len(new_block_ids)) + ] self._reqs_need_pending_save[req_id] = (req, updated_blocks) - if ( - len(self._reqs_need_pending_save[req_id][1]) * self.block_size - >= req.num_prompt_tokens - ): + saved_tokens = ( + len(updated_blocks[self._full_attn_group_idx]) + * self._full_attn_block_size + ) + if saved_tokens >= req.num_prompt_tokens: # Final chunk: live kv_transfer_params may be cleared, # so prefer the snapshot from update_state_after_alloc. kv_params = self._req_kv_params.pop( @@ -857,7 +961,10 @@ def build_connector_meta( for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) - if req.num_prompt_tokens > len(block_ids) * self.block_size: + tokens_covered = ( + len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size + ) + if req.num_prompt_tokens > tokens_covered: # not last chunk prefill self._reqs_need_pending_save[req_id] = (req, block_ids) continue @@ -898,7 +1005,7 @@ def shutdown(self): def request_finished( self, request: "Request", - block_ids: list[int], + block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: """ Once a request is finished, determine whether request blocks @@ -936,7 +1043,8 @@ def request_finished( if self.mode == MoRIIOMode.WRITE: self._release_write_prefill_blocks(request.request_id, params) else: - self._reqs_need_recv[request.request_id] = (request, []) + empty_groups: BlockIds = [[] for _ in self.blocks_per_sw] + self._reqs_need_recv[request.request_id] = (request, empty_groups) params["do_remote_prefill"] = False return False, None @@ -946,10 +1054,9 @@ def request_finished( ): return False, None - # computed_block_ids = block_ids if all_full else block_ids[:-1] - computed_block_ids = block_ids - # If prompt < block_size, no xfer so free blocks immediately. - delay_free_blocks = len(computed_block_ids) > 0 + # On producer + computed_block_ids = self.get_exchange_clipped_blocks(block_ids) + delay_free_blocks = any(len(group) > 0 for group in computed_block_ids) if delay_free_blocks: # Prefill request on remote. It will be read from D upon completion @@ -1102,6 +1209,12 @@ def __init__( self.kv_transfer_config = vllm_config.kv_transfer_config self.is_producer = self.kv_transfer_config.is_kv_producer self.layer_to_spec = build_layer_to_spec(kv_cache_config) + # we transfer by layer, but store by kv cache group + self.layer_to_group: dict[str, int] = { + layer_name: group_idx + for group_idx, group in enumerate(kv_cache_config.kv_cache_groups) + for layer_name in group.layer_names + } if self.is_producer: set_role(ROLE.PRODUCER) @@ -1789,8 +1902,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.block_size, ) self.block_size = first_geometry.block_size - # TODO(tms): self.block_len needs to be per-layer for sliding window, - # hybrid attn, etc # block size in bytes self.block_len = first_geometry.block_len self.kv_cache_shape = first_kv_cache.shape @@ -1803,14 +1914,30 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.layer_base_addr_index = {} base_addr_idx = 0 + hma_enabled = ( + not self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + ) for layer_name in kv_caches: self.layer_base_addr_index[layer_name] = base_addr_idx geometry = self._get_layer_transfer_geometry(layer_name) - if geometry.block_size != self.block_size: + if hma_enabled: + if geometry.block_len != self.block_len: + raise ValueError( + "MoRIIO KV cache block length mismatch for layer " + f"{layer_name}: {geometry.block_len} != {self.block_len}" + ) + elif geometry.block_size != self.block_size: raise ValueError( "MoRIIO KV cache block size mismatch for layer " f"{layer_name}: {geometry.block_size} != {self.block_size}" ) + # num_blocks is advertised as a single scalar to the peer, so must it + # be uniform + if geometry.num_blocks != self.num_blocks: + raise ValueError( + "MoRIIO KV cache num_blocks mismatch for layer " + f"{layer_name}: {geometry.num_blocks} != {self.num_blocks}" + ) self.block_lens[layer_name] = geometry.block_len for cache, region_len in self._iter_layer_registration_regions(layer_name): base_addr = cache.data_ptr() @@ -2455,12 +2582,16 @@ def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_lay self.remote_dp_size_local = int(meta.remote_dp_size_local) else: self.remote_dp_size_local = int(meta.remote_dp_size) + # WRITE does not support HMA, so unwrap blocks into flat lists. + # remote_block_ids can itself be empty so need to be careful when unwrapping. + local_block_ids = meta.local_block_ids[0] + remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else [] self.schedule_write_blocks( request_id=req_id, transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, - local_block_ids=meta.local_block_ids, - remote_block_ids=meta.remote_block_ids, + local_block_ids=local_block_ids, + remote_block_ids=remote_block_ids, layer_name=layer_name, kv_layer=kv_layer, remote_notify_port=meta.remote_notify_port, @@ -2607,8 +2738,8 @@ def _is_sq_full_status(status) -> bool: def _read_blocks( self, - local_block_ids: list[int], - remote_block_ids: list[int], + local_block_ids: BlockIds, + remote_block_ids: BlockIds, dst_engine_id: str, request_id: str, transfer_id: str, @@ -2653,10 +2784,11 @@ def _read_blocks( sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( layer_name ) + group_idx = self.layer_to_group[layer_name] offs = self._compute_block_transfer_offsets( layer_name, - local_block_ids, - remote_block_ids, + local_block_ids[group_idx], + remote_block_ids[group_idx], remote_moriio_meta, remote_tp_size=remote_tp_size, remote_engine_id=remote_dp_engine_id, From e0c04c7b4dddceb25eb9d23d87c3914bfc2ad3f9 Mon Sep 17 00:00:00 2001 From: Ronald Date: Mon, 14 Sep 2026 21:05:47 +0800 Subject: [PATCH 183/774] [Bugfix][Models] Fix OpenPangu sleep mode with static sinks (#53696) Signed-off-by: Ronald1995 --- .../layers/attention/static_sink_attention.py | 2 -- vllm/model_executor/models/openpangu.py | 18 +++++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index d27ccc728ad7..d321d675a5fe 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -152,8 +152,6 @@ def __init__( attn_backend=attn_backend, **kwargs, ) - CustomOp.__init__(self) - self.sink_len = sink_len self.sink_populated = False self.sink_key = None diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 60b5b301d56e..b8937f2e09fb 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -677,14 +677,18 @@ def __init__( }, ) else: - self.param_sink_value = torch.zeros( - ( - self.param_sink_number, - self.num_kv_heads, - self.v_channels, + self.register_buffer( + "param_sink_value", + torch.zeros( + ( + self.param_sink_number, + self.num_kv_heads, + self.v_channels, + ), + device=current_platform.current_device(), + dtype=config.torch_dtype, ), - device=current_platform.current_device(), - dtype=config.torch_dtype, + persistent=False, ) # To enable dummy run with out weight self.post_weight_load() From 47fbd36e6dc908277e386424beeb8647355c1c8b Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 14 Sep 2026 21:06:26 +0800 Subject: [PATCH 184/774] [Bugfix] Trim stale consequence claims from unannotated-eagle warning (#56791) Signed-off-by: zjy0516 Co-authored-by: Kimi Code --- tests/v1/core/test_kv_cache_utils.py | 21 +++++++++++++++++---- vllm/v1/core/kv_cache_utils.py | 14 ++------------ 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index dbddc556c9ba..365338ad4f3f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -4135,9 +4135,9 @@ def test_draft_group_not_annotated_without_spec_decode(): def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so every consumer falls back to - # flagging all groups -- including Mamba ones, which then can never report - # a hit. That is silent today; it must at least be visible. + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) ) @@ -4146,7 +4146,20 @@ def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): assert "no KV cache group could be identified as the draft model's" in ( caplog_vllm.text ) - assert "Mamba groups" in caplog_vllm.text + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text def test_no_warning_when_draft_group_is_identified(caplog_vllm): diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 996cce8c43f5..4e4f34bd52ad 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2182,12 +2182,7 @@ def _warn_if_unannotated_eagle_mamba( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> None: - """Warn when the flag-all eagle fallback will silently disable reuse. - - With no group annotated, consumers flag every group as a draft group. That - widens a Mamba group's required lookup window to two consecutive chunks, - which align-mode checkpointing never produces, so reuse drops to zero with - no error and no metric to show it. + """Warn when no KV cache group could be identified as the draft model's. Args: vllm_config: Config supplying the speculative method, if any. @@ -2210,13 +2205,8 @@ def _warn_if_unannotated_eagle_mamba( return logger.warning( "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's, so every group -- " - "including Mamba groups %s -- will be treated as a draft group. A " - "Mamba group cannot satisfy the widened lookup window that implies, " - "so prefix-cache reuse across requests will be disabled and any " - "external KV offload tier will store without ever serving a hit.", + "could be identified as the draft model's.", spec_config.method, - mamba_groups, ) From 6623fe5b4ec783e7264f140e0caea20bdc45a341 Mon Sep 17 00:00:00 2001 From: arkohut <39525455+arkohut@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:18:55 +0800 Subject: [PATCH 185/774] [Bugfix] MiniCPM-V 4.6: fix ViT self-attn qkv weight loading (#49417) Signed-off-by: arkohut <39525455+arkohut@users.noreply.github.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/minicpmv4_6.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index e43f13f79f7a..d8e10d349884 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -674,14 +674,6 @@ def get_num_image_tokens( class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - } - ) - def __init__( self, config, @@ -729,10 +721,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: out, _ = self.out_proj(attn_out) return out - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - class MiniCPMV4_6ViTWindowAttentionMerger(nn.Module): def __init__( @@ -971,6 +959,11 @@ class MiniCPMV4_6ForConditionalGeneration( supports_encoder_tp_data = True hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + }, orig_to_new_prefix={ # transformers v5.7+ uses `vision_tower` and nests `vit_merger` # inside it. Order matters: more specific prefix must come first. @@ -982,7 +975,7 @@ class MiniCPMV4_6ForConditionalGeneration( "model.language_model.": "language_model.model.", "lm_head.": "language_model.lm_head.", "mtp.": None, - } + }, ) packed_modules_mapping = { From 9d4d9aa5bced3fc84c9e7c89b1257f206fbff6e2 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Mon, 14 Sep 2026 21:19:15 +0800 Subject: [PATCH 186/774] [LoRA][Refactor] Unify multimodal LoRA token count hooks (#55071) Signed-off-by: linitra24 --- .../processing/test_llava_next_video.py | 37 ++++++++------ vllm/model_executor/models/blip2.py | 35 ++++++-------- vllm/model_executor/models/dots_ocr.py | 16 ++++--- vllm/model_executor/models/gemma3_mm.py | 48 +++++-------------- vllm/model_executor/models/glm4_1v.py | 20 ++++---- vllm/model_executor/models/h2ovl.py | 29 +++++------ vllm/model_executor/models/idefics3.py | 23 ++++----- vllm/model_executor/models/interfaces.py | 27 ++++------- vllm/model_executor/models/internvl.py | 28 ++++++----- vllm/model_executor/models/keye.py | 16 ++++--- vllm/model_executor/models/lfm2_vl.py | 18 +++---- vllm/model_executor/models/llava.py | 23 ++++----- .../model_executor/models/llava_next_video.py | 42 ++++++++-------- vllm/model_executor/models/mllama4.py | 28 +++++------ vllm/model_executor/models/moondream3.py | 15 ++++-- vllm/model_executor/models/openvla.py | 20 +++++--- vllm/model_executor/models/paligemma.py | 15 ++++-- vllm/model_executor/models/pixtral.py | 20 ++++---- vllm/model_executor/models/qwen2_5_vl.py | 23 ++++----- vllm/model_executor/models/qwen2_vl.py | 23 ++++----- vllm/model_executor/models/qwen3_asr.py | 16 ++++--- vllm/model_executor/models/qwen3_vl.py | 22 ++++----- vllm/model_executor/models/ultravox.py | 33 +++++++------ 23 files changed, 279 insertions(+), 298 deletions(-) diff --git a/tests/models/multimodal/processing/test_llava_next_video.py b/tests/models/multimodal/processing/test_llava_next_video.py index 2dcbb7d4b55f..c4fdecd21ac7 100644 --- a/tests/models/multimodal/processing/test_llava_next_video.py +++ b/tests/models/multimodal/processing/test_llava_next_video.py @@ -14,7 +14,7 @@ class _StubModel: - """Carries only the two attributes the token helpers read from + """Carries only the attributes the token-count method reads from `self`, so the real methods can be exercised without constructing the full `nn.Module` (vision tower, language model, etc.).""" @@ -22,18 +22,15 @@ class _StubModel: pooled_grid_length: int -get_num_mm_encoder_tokens = ( - LlavaNextVideoForConditionalGeneration.get_num_mm_encoder_tokens -) -get_num_mm_connector_tokens = ( - LlavaNextVideoForConditionalGeneration.get_num_mm_connector_tokens +get_mm_lora_token_counts = ( + LlavaNextVideoForConditionalGeneration.get_mm_lora_token_counts ) @pytest.mark.parametrize("model_id", ["llava-hf/LLaVA-NeXT-Video-7B-hf"]) def test_num_mm_tokens_match_real_config(model_id): """The stored grid lengths must match what `__init__` derives from - the real HF config, and the two helpers must invert each other's + the real HF config, and the LoRA token counts must preserve the frame-level scaling exactly.""" ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1}) config = ctx.model_config.hf_config @@ -49,10 +46,13 @@ def test_num_mm_tokens_match_real_config(model_id): for num_frames in (1, 2, 8, 16, 32): num_video_tokens = num_frames * pooled_grid_length**2 - encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens) + encoder_tokens, connector_tokens = get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=num_video_tokens, + ) assert encoder_tokens == num_frames * patch_grid_length**2 - - connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens) assert connector_tokens == num_video_tokens @@ -72,10 +72,13 @@ def test_num_mm_tokens_roundtrip(patch_grid_length, pooled_grid_length, num_fram num_video_tokens = num_frames * pooled_grid_length**2 - encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens) + encoder_tokens, connector_tokens = get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=num_video_tokens, + ) assert encoder_tokens == num_frames * patch_grid_length**2 - - connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens) assert connector_tokens == num_video_tokens @@ -84,5 +87,9 @@ def test_num_mm_tokens_zero(): stub.patch_grid_length = 24 stub.pooled_grid_length = 12 - assert get_num_mm_encoder_tokens(stub, 0) == 0 - assert get_num_mm_connector_tokens(stub, 0) == 0 + assert get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=0, + ) == (0, 0) diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index 091001fc0173..de74d1f5ace0 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -22,6 +22,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems @@ -687,28 +688,22 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_model", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - if num_image_tokens <= 0: - return 0 - assert num_image_tokens % self.config.num_query_tokens == 0, ( + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0: + return 0, 0 + assert num_mm_embeds % self.config.num_query_tokens == 0, ( "The number of image tokens must be a multiple of " "the number of query tokens." ) - num_images = num_image_tokens / self.config.num_query_tokens - return num_images * self._vision_tokens_per_image - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - if num_vision_tokens <= 0: - return 0 - assert num_vision_tokens % self._vision_tokens_per_image == 0, ( - "The number of vision tokens must be a multiple of " - "the number of tokens per image." + num_images = num_mm_embeds // self.config.num_query_tokens + return ( + num_images * self._vision_tokens_per_image, + num_images * self.config.num_query_tokens, ) - num_images = num_vision_tokens / self._vision_tokens_per_image - return num_images * self.config.num_query_tokens diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py index e91d6ce94c9e..c939c2bef5b3 100644 --- a/vllm/model_executor/models/dots_ocr.py +++ b/vllm/model_executor/models/dots_ocr.py @@ -54,6 +54,7 @@ ) from vllm.model_executor.models.vision import get_vit_attn_backend from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig, DotsVisionConfig from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -718,13 +719,16 @@ def _process_image_input( return image_embeds.split(sizes) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - merge_size = self.vision_tower.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.vision_tower.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: image_input = self._parse_and_validate_image_input(**kwargs) diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index b45327fbebf7..abe1f5f1bd55 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -19,6 +19,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -686,43 +687,16 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - """ - Calculate the number of tokens output by the vision encoder. - - The vision encoder processes images into patch embeddings. For Gemma3, - the relationship between prompt placeholder tokens and actual vision - encoder output tokens depends on the patch grid size. - - Args: - num_image_tokens: Number of image placeholder tokens in the prompt - (typically mm_tokens_per_image per image) - - Returns: - Number of tokens output by the vision encoder - """ - # For Gemma3, the vision encoder outputs tokens_per_side x tokens_per_side - # tokens per image. Since num_image_tokens represents the number of - # connector output tokens (mm_tokens_per_image = 256), and tokens_per_side - # is sqrt(256) = 16, we need to account for the token expansion. - # Based on empirical testing, the multiplier of 16 works correctly. - return num_image_tokens * 16 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - """ - Calculate the number of tokens output by the multimodal connector. - - The connector applies projection and normalization but maintains the - token count for Gemma3. - - Args: - num_vision_tokens: Number of tokens from vision encoder - - Returns: - Number of tokens after connector processing - """ - # The Gemma3 connector maintains a 1:1 token mapping - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + tower_tokens = num_mm_embeds * 16 + return tower_tokens, tower_tokens def get_encoder_cudagraph_config(self): from vllm.v1.worker.encoder_cudagraph_defs import ( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 163115ff3d5e..730a015cfe53 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -83,6 +83,7 @@ from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -2378,19 +2379,16 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - merge_size = self.config.vision_config.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.config.vision_config.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm/model_executor/models/h2ovl.py b/vllm/model_executor/models/h2ovl.py index 5c4f984c13f8..5be2d7b46d99 100644 --- a/vllm/model_executor/models/h2ovl.py +++ b/vllm/model_executor/models/h2ovl.py @@ -14,7 +14,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import BatchedTensorInputs +from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem from vllm.multimodal.parse import ( ImageEmbeddingItems, ImageProcessorItems, @@ -178,16 +178,17 @@ def _init_vision_model( prefix=prefix, ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - if num_image_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_image_tokens // self.num_image_token - return num_patches * (self.patch_tokens + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - if num_vision_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_vision_tokens // (self.patch_tokens + 1) - return num_patches * self.num_image_token + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0 or self.num_image_token <= 0: + return 0, 0 + + num_patches = num_mm_embeds // self.num_image_token + tower_tokens = num_patches * (self.patch_tokens + 1) + return tower_tokens, num_patches * self.num_image_token diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index baaa4106cb6b..4ac8eb5eeb43 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -39,6 +39,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems @@ -837,20 +838,14 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="model.vision_model", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - hf_config = self.config - scale_factor = hf_config.scale_factor - - return num_image_tokens * scale_factor**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config scale_factor = hf_config.scale_factor - - return num_vision_tokens // scale_factor**2 + return num_mm_embeds * scale_factor**2, num_mm_embeds diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 254f2c635bc1..2e79001add82 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -31,7 +31,7 @@ import torch.nn as nn from torch import Tensor from transformers.models.whisper.tokenization_whisper import LANGUAGES -from typing_extensions import Self, TypeIs +from typing_extensions import Self, TypeIs, deprecated from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig @@ -420,23 +420,16 @@ def _mark_composite_model( yield - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - """ - Implement this function to enable LoRA support - for the tower module of the multi-modal model. - Given the number of image tokens, output the number of - multi-modal encoder tokens. - """ - ... + @deprecated( + "get_num_mm_encoder_tokens is deprecated; use get_mm_lora_token_counts instead." + ) + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: ... - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - """ - Implement this function to enable LoRA support - for the connector module of the multi-modal model. - Given the number of vision tokens, output the number of - multi-modal connector tokens. - """ - ... + @deprecated( + "get_num_mm_connector_tokens is deprecated; use " + "get_mm_lora_token_counts instead." + ) + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: ... def get_mm_lora_token_counts( self, diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index 1fcb6b8681d6..c57359b7b5ad 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -29,6 +29,7 @@ from vllm.multimodal.inputs import ( BatchedTensorInputs, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -914,19 +915,20 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_model", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - if num_image_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_image_tokens // self.num_image_token - return num_patches * (self.patch_tokens + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - if num_vision_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_vision_tokens // (self.patch_tokens + 1) - return num_patches * self.num_image_token + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0 or self.num_image_token <= 0: + return 0, 0 + + num_patches = num_mm_embeds // self.num_image_token + tower_tokens = num_patches * (self.patch_tokens + 1) + return tower_tokens, num_patches * self.num_image_token # -- SupportsEncoderCudaGraph protocol methods -- diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 53cda0ed171b..d117b91d53b6 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -40,6 +40,7 @@ ImageItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -1495,13 +1496,16 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="visual.", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - merge_size = self.config.vision_config.spatial_merge_size - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.config.vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index a8b57e7c28a2..15872a1f3f39 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -35,6 +35,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -1284,12 +1285,13 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - downsample_factor = self.config.downsample_factor - - return num_image_tokens * downsample_factor**2 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs downsample_factor = self.config.downsample_factor - - return num_vision_tokens // downsample_factor**2 + return num_mm_embeds * downsample_factor**2, num_mm_embeds diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index 1f1f45bfea73..a32e0e6c63cf 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -28,6 +28,7 @@ from vllm.multimodal.cache import BaseMultiModalProcessorCache from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -739,18 +740,12 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_tower", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - # LLaVA's vision encoder outputs one token per patch without - # spatial merging or pixel shuffle - return num_image_tokens - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - # LLaVA's MLP projector outputs the same number of tokens - # as it receives from the vision encoder (1:1 mapping) - return num_vision_tokens + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds diff --git a/vllm/model_executor/models/llava_next_video.py b/vllm/model_executor/models/llava_next_video.py index 038da17b75bf..96ebdaea8eb1 100644 --- a/vllm/model_executor/models/llava_next_video.py +++ b/vllm/model_executor/models/llava_next_video.py @@ -17,6 +17,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -494,31 +495,30 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_tower", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_video_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + # Invert the spatial pooling done by `vision_resampler`: each frame # contributes `pooled_grid_length ** 2` tokens after the language # model's placeholder count, but `patch_grid_length ** 2` tokens # when it leaves the vision encoder. pooled_tokens_per_frame = self.pooled_grid_length**2 - if num_video_tokens <= 0 or pooled_tokens_per_frame <= 0: - return 0 - - num_frames = num_video_tokens // pooled_tokens_per_frame - return num_frames * (self.patch_grid_length**2) - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - # The projector is length-preserving; it runs after - # `vision_resampler` has already pooled each frame down to - # `pooled_grid_length ** 2` tokens. patch_tokens_per_frame = self.patch_grid_length**2 - if num_vision_tokens <= 0 or patch_tokens_per_frame <= 0: - return 0 - - num_frames = num_vision_tokens // patch_tokens_per_frame - return num_frames * (self.pooled_grid_length**2) + if ( + num_mm_embeds <= 0 + or pooled_tokens_per_frame <= 0 + or patch_tokens_per_frame <= 0 + ): + return 0, 0 + + num_frames = num_mm_embeds // pooled_tokens_per_frame + return ( + num_frames * patch_tokens_per_frame, + num_frames * pooled_tokens_per_frame, + ) diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 5a6e65441703..0096eb42418f 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -60,6 +60,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -1281,21 +1282,18 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_model.", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs vision_config = self.config.vision_config patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config) - if num_image_tokens <= 0 or patches_per_chunk <= 0: - return 0 - raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2 - num_chunks = num_image_tokens // patches_per_chunk - # Encoder processes raw_patches + 1 (CLS) per chunk - return num_chunks * (raw_patches + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - vision_config = self.config.vision_config + if num_mm_embeds <= 0 or patches_per_chunk <= 0: + return 0, 0 raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2 - if num_vision_tokens <= 0: - return 0 - num_chunks = num_vision_tokens // (raw_patches + 1) - patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config) - return num_chunks * patches_per_chunk + num_chunks = num_mm_embeds // patches_per_chunk + return num_chunks * (raw_patches + 1), num_chunks * patches_per_chunk diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py index 818015e1ba7c..47a240b58953 100644 --- a/vllm/model_executor/models/moondream3.py +++ b/vllm/model_executor/models/moondream3.py @@ -47,6 +47,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageSize, MultiModalDataItems @@ -1104,11 +1105,15 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None: def get_language_model(self) -> nn.Module: return self.text - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds def _split_pixel_values( self, diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py index a676b3447cf4..ecbfb8a54c27 100644 --- a/vllm/model_executor/models/openvla.py +++ b/vllm/model_executor/models/openvla.py @@ -24,7 +24,11 @@ SupportsPP, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItem, + MultiModalKwargsItems, +) from vllm.multimodal.parse import ( ImageEmbeddingItems, ImageProcessorItems, @@ -487,11 +491,15 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_backbone", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def maybe_rename_vision_weights( diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py index 2fa498ee261c..8e3524eeec4f 100644 --- a/vllm/model_executor/models/paligemma.py +++ b/vllm/model_executor/models/paligemma.py @@ -14,6 +14,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -401,8 +402,12 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index b6848b09ded0..1485b90e39a2 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -42,6 +42,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsOptionalItems, NestedTensors, ) @@ -623,17 +624,18 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="vision_encoder", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - if getattr(self, "patch_merger", None) is None: - return num_image_tokens - merge_size = self.vision_args.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs if getattr(self, "patch_merger", None) is None: - return num_vision_tokens + return num_mm_embeds, num_mm_embeds merge_size = self.vision_args.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds # Vision encoder diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 61d0153f86e8..6cfacc1421fc 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -70,6 +70,7 @@ from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems @@ -2050,21 +2051,15 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 80abdc62b576..c3761f8893f5 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -66,6 +66,7 @@ ImageItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -1747,21 +1748,15 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index e98f37ba37c2..3bea939340e5 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -67,6 +67,7 @@ AudioItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -625,16 +626,19 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model=["audio_tower."], ) - def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int: - """Return the number of tokens processed by the audio tower encoder. - - Required for LoRA support on the tower module. - """ + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs # For Qwen3-ASR, the audio tower produces one embedding per audio # placeholder token inserted into the prompt (no additional # merge/downsample step like vision towers). Therefore, the encoder # token budget is identity. - return num_audio_tokens + return num_mm_embeds, None @classmethod def get_speech_to_text_config( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index d99cfe7aa296..644fd134adce 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -3077,24 +3077,18 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds @lru_cache diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index a394e5f02e05..bffe94fdce26 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -29,6 +29,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, NestedTensors, ) @@ -373,9 +374,9 @@ def _build_chunk_attn_metadata( valid/padding boundary (equivalent to the key-padding mask the HF implementation uses), while every row still flows through the (potentially LoRA-wrapped) linears, keeping the per-chunk token counts - constant as required by `get_num_mm_encoder_tokens` / - `get_num_mm_connector_tokens`. Queries at padding positions produce - (garbage) outputs, which are trimmed by `audio_token_len` downstream. + constant as required by `get_mm_lora_token_counts`. Queries at padding + positions produce (garbage) outputs, which are trimmed by `audio_token_len` + downstream. """ batch_size = feature_lens.shape[0] starts = np.arange(batch_size, dtype=np.int64) * seq_len @@ -684,7 +685,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # LoRA on the tower/connector requires per-item token counts that are # predictable from the placeholder count alone (see - # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input + # `get_mm_lora_token_counts`), so pad every audio chunk's mel input # to the tower's full context instead of the batch's max length. self.pad_audio_to_max_context = bool( lora_config is not None and lora_config.enable_tower_connector_lora @@ -769,22 +770,26 @@ def _get_max_tokens_per_chunk(self) -> int: self.config.audio_config.max_source_positions / self.config.stack_factor ) - def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs # With `pad_audio_to_max_context` (enforced when tower/connector LoRA # is enabled), the tower's LoRA-wrapped linears always run on # `max_source_positions` conv-downsampled tokens per chunk, regardless # of the valid frame count. - num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk()) - return num_chunks * self.config.audio_config.max_source_positions - - def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int: + num_chunks = math.ceil(num_mm_embeds / self._get_max_tokens_per_chunk()) + tower_tokens = num_chunks * self.config.audio_config.max_source_positions # The connector runs on the frame-stacked tower output. Stacking pads # each chunk to a multiple of `stack_factor`, so this is # ceil(max_source_positions / stack_factor) tokens per chunk (188 for - # whisper's 1500), not `num_encoder_tokens // stack_factor` (187). - max_source_positions = self.config.audio_config.max_source_positions - num_chunks = num_encoder_tokens // max_source_positions - return num_chunks * self._get_max_tokens_per_chunk() + # whisper's 1500), not `tower_tokens // stack_factor` (187). + connector_tokens = num_chunks * self._get_max_tokens_per_chunk() + return tower_tokens, connector_tokens def _audio_features_to_embeddings( self, @@ -866,7 +871,7 @@ def _process_audio_input( # Pad every chunk to the tower's full context so the number of # tokens processed by the tower/connector per chunk is constant, # keeping the LoRA token mappings computed by - # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact. + # `get_mm_lora_token_counts` exact. # Attention to the extra padding is masked via `audio_lens`. # Over-long inputs are not trimmed here; the tower raises on them. pad_len = self.audio_tower.max_context_length - audio_features.shape[-1] From 767d1c4d473a3e81429563aac995f22f7151c325 Mon Sep 17 00:00:00 2001 From: Russell Bryant Date: Mon, 14 Sep 2026 09:44:50 -0400 Subject: [PATCH 187/774] [Docs] Move russellb to emeritus committer (#56467) Signed-off-by: Russell Bryant Co-authored-by: Claude Opus 5 --- .github/CODEOWNERS | 16 ++++++++-------- docs/contributing/vulnerability_management.md | 2 -- docs/governance/committers.md | 9 +++++---- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 10fd12576893..aa5b2a9fc9dd 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -34,9 +34,9 @@ # Entrypoints /vllm/entrypoints/anthropic @mgoin @DarkLight1337 -/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb +/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 /vllm/entrypoints/mcp @heheda12345 -/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb +/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 /vllm/entrypoints/speech_to_text/realtime @njhill /vllm/entrypoints/speech_to_text @NickLucche /vllm/entrypoints/pooling @noooop @@ -73,7 +73,7 @@ /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni -/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett @arpera +/vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera /vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium @@ -107,8 +107,8 @@ /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas -/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm -/tests/v1/structured_output @mgoin @russellb @aarnphm @arpera +/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera +/tests/v1/structured_output @mgoin @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee @@ -232,6 +232,6 @@ mkdocs.yaml @hmellor /docs/training/weight_transfer @aoshen02 # Security guide and policies -/docs/usage/security.md @russellb @jperezdealgaba -/SECURITY.md @russellb @jperezdealgaba -/docs/contributing/vulnerability_management.md @russellb @jperezdealgaba +/docs/usage/security.md @jperezdealgaba +/SECURITY.md @jperezdealgaba +/docs/contributing/vulnerability_management.md @jperezdealgaba diff --git a/docs/contributing/vulnerability_management.md b/docs/contributing/vulnerability_management.md index 8dd908a7ce5d..a73d451a6217 100644 --- a/docs/contributing/vulnerability_management.md +++ b/docs/contributing/vulnerability_management.md @@ -33,9 +33,7 @@ on GitHub. However, if you need to contact the VMT directly for an urgent issue, you may contact the following individuals: - Simon Mo - -- Russell Bryant - - Juan Pérez de Algaba - -- Huzaifa Sidhpurwala - ## Slack Discussion diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 0ca6e24ded0a..bd3675912217 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -32,6 +32,7 @@ Sorted alphabetically by GitHub handle: - [@ivanium](https://github.com/ivanium): KV Connector and offload - [@jeejeelee](https://github.com/jeejeelee): LoRA, new model support - [@jikunshang](https://github.com/jikunshang): Intel CPU/XPU integration +- [@jperezdealgaba](https://github.com/jperezdealgaba): Security - [@khluu](https://github.com/khluu): CI infrastructure - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance - [@markmc](https://github.com/markmc): Observability @@ -44,7 +45,6 @@ Sorted alphabetically by GitHub handle: - [@pavanimajety](https://github.com/pavanimajety): NVIDIA GPU integration - [@ProExpertProg](https://github.com/ProExpertProg): Compilation, startup UX - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg -- [@russellb](https://github.com/russellb): Structured output, engine core, security - [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser - [@shen-shanshan](https://github.com/shen-shanshan): AMD GPU / ROCm integration, multimodality, ViT CUDA graph - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community @@ -82,6 +82,7 @@ Committers who have contributed to vLLM significantly in the past (thank you!) b - [@pcmoritz](https://github.com/pcmoritz): MoE - [@rkooo567](https://github.com/rkooo567): Chunked prefill - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support +- [@russellb](https://github.com/russellb): Structured output, engine core, security - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support - [@sroy745](https://github.com/sroy745): Speculative decoding - [@Yard1](https://github.com/Yard1): kernels and performance @@ -100,7 +101,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - KV Cache Manager: memory management layer within scheduler maintaining KV cache logical block data - @heheda12345, @WoosukKwon - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints - - @robertgshaw2-redhat, @njhill, @russellb + - @robertgshaw2-redhat, @njhill - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni - KV Connector: Connector interface and implementation for KV cache offload and transfer @@ -153,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP. - @WoosukKwon, @benchislett, @MatthewBonanni - Structured Output: The structured output implementation - - @russellb, @aarnphm + - @aarnphm - RL: The RL related features such as collective rpc, sleep mode, etc. - @youkaichao, @zhuohan123 - LoRA: @jeejeelee @@ -167,7 +168,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Documentation: @hmellor, @DarkLight1337, @simon-mo - Benchmarks: @ywang96, @simon-mo - CI, Build, Release Process: @khluu, @njhill, @simon-mo, @Harry-Chen, @vadiklyutiy -- Security: @russellb +- Security: @jperezdealgaba ### External Kernels Integration From c676e4930bb3c87b97790520e7979d43887ca5dd Mon Sep 17 00:00:00 2001 From: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:07:07 +0200 Subject: [PATCH 188/774] [Spec Decode] Fix Qwen3 DSpark d2t requirement for padded-vocab drafts (#55133) Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis Co-authored-by: Benjamin Chislett Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 --- tests/model_executor/test_qwen3_omni.py | 46 ++++++++++++++++------ vllm/model_executor/models/qwen3_dspark.py | 7 +++- 2 files changed, 38 insertions(+), 15 deletions(-) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 92d527e6252c..42accbef85c5 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -293,6 +293,25 @@ def forward(self, hidden_states, residual): torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]])) +def _dspark_vocab_stub(input_vocab_size: int, draft_vocab_size: int): + """A Qwen3 DSpark model with only what `load_weights` reads.""" + from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM + + model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM) + nn.Module.__init__(model) + object.__setattr__( + model, + "config", + SimpleNamespace( + vocab_size=input_vocab_size, + draft_vocab_size=draft_vocab_size, + ), + ) + object.__setattr__(model, "target_vocab_size", 100) + object.__setattr__(model, "model", Mock(confidence_head=None)) + return model + + @pytest.mark.skip_global_cleanup @pytest.mark.parametrize( ("input_vocab_size", "draft_vocab_size", "weights", "error"), @@ -311,24 +330,25 @@ def forward(self, hidden_states, residual): def test_qwen3_dspark_rejects_incomplete_vocab_weights( input_vocab_size, draft_vocab_size, weights, error ): - from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM - - model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM) - nn.Module.__init__(model) - object.__setattr__( - model, - "config", - SimpleNamespace( - vocab_size=input_vocab_size, - draft_vocab_size=draft_vocab_size, - ), - ) - object.__setattr__(model, "target_vocab_size", 100) + model = _dspark_vocab_stub(input_vocab_size, draft_vocab_size) with pytest.raises(ValueError, match=error): model.load_weights(weights) +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("draft_vocab_size", [100, 101]) +def test_qwen3_dspark_accepts_draft_vocab_at_least_target(draft_vocab_size): + """A draft vocab >= the target's maps identically, so needs no lm_head/d2t. + + Guards the padded case (draft > target): the draft's unembedding is padded + to its physical size, which is not a reduced vocabulary. + """ + model = _dspark_vocab_stub(100, draft_vocab_size) + + model.load_weights([]) + + @pytest.mark.skip_global_cleanup def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index c2c252280d6a..7318d225f8d2 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -254,7 +254,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.config.draft_vocab_size, scale=logit_scale ) self.target_vocab_size = vllm_config.model_config.get_vocab_size() - if self.config.draft_vocab_size != self.target_vocab_size: + # Only a pruned draft (subset of the target vocab) needs a d2t mapping to + # scatter ids up. A draft >= the target vocab (e.g. padded to the target's + # physical embedding size) maps identically, so no d2t is needed. + if self.config.draft_vocab_size < self.target_vocab_size: self.draft_id_to_target_id = nn.Parameter( torch.zeros(self.config.draft_vocab_size, dtype=torch.long), requires_grad=False, @@ -334,7 +337,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): # embed_tokens / lm_head are optional; when omitted they are shared from # the target by load_dspark_model, so skip the unloaded params here. uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size - uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size + uses_reduced_vocab = self.config.draft_vocab_size < self.target_vocab_size if uses_expanded_input_vocab and not includes_embed_tokens: raise ValueError( "Qwen3 DSpark checkpoints whose input vocab_size is larger than " From ff5f6d41b1ccf1e305e425db41bc7f52dbce8621 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:07:09 +0200 Subject: [PATCH 189/774] [Bugfix] Scale KV page size for hidden states extraction with TP (#50894) Signed-off-by: Claude Signed-off-by: Orestis Zambounis Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Speculators Bot Signed-off-by: Speculators Bot Co-authored-by: Orestis Zambounis Co-authored-by: Claude Opus 4.6 Co-authored-by: Benjamin Chislett Co-authored-by: Speculators Bot Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 39 ++++++++++++++++++++++++++++ vllm/v1/core/kv_cache_utils.py | 34 ++++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 365338ad4f3f..c6354cdc0fa0 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3208,6 +3208,45 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): assert "attention compute is unchanged" in caplog_vllm.text +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4e4f34bd52ad..b47df61add03 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2217,6 +2217,38 @@ def _largest_divisor_at_most(value: int, limit: int) -> int: return 1 +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """Scale up group block sizes so the common page is at least as large as + the biggest hidden-state per-token cost. + + This protects the hidden state extraction feature, where we store hidden states in + the KV cache, and return them to clients using a custom hidden state connector. + + Returns the (possibly scaled) groups and updated common page size. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, int] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + def get_kv_cache_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], @@ -2288,6 +2320,8 @@ def get_kv_cache_groups( # Add hidden-state layers back with page aligned to the common page. if hidden_specs: common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) for name, spec in hidden_specs.items(): per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) From 4be3dcf0fc7a9086d978eaea3c5a4d78fb98e44a Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 14 Sep 2026 10:20:46 -0400 Subject: [PATCH 190/774] [PCP][DCP] Declare FlashMLASparse MTP support at CP interleave > 1 (#56722) Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/attention/backends/mla/flashmla_sparse.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 433cded5679d..d8c4368401c6 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -634,6 +634,7 @@ def build( class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): can_return_lse_for_decode: bool = True supports_dcp: bool = True + supports_mtp_with_cp_non_trivial_interleave_size: bool = True @staticmethod def _compute_fp8_decode_padded_heads(num_heads: int) -> int: From b3124a8237f21fcd3a4510002a5e2925b0dbedfe Mon Sep 17 00:00:00 2001 From: zql <37731799+zqlcode@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:06:01 +0800 Subject: [PATCH 191/774] [Model] Add support for Nanbeige4.2 (transformers backend) (#56071) Signed-off-by: zqlcode Signed-off-by: zql <37731799+zqlcode@users.noreply.github.com> Co-authored-by: zqlcode Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/models/supported_models.md | 1 + tests/models/registry.py | 4 ++++ vllm/model_executor/models/registry.py | 1 + 3 files changed, 6 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 81f4fb7df1f6..d55348ad724a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -490,6 +490,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | +| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 874fa9c11ca5..cb39ed7ce4f3 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -461,6 +461,10 @@ def check_available_online( trust_remote_code=True, is_available_online=False, ), + "NanbeigeForCausalLM": _HfExamplesInfo( + "Nanbeige/Nanbeige4.2-3B", + trust_remote_code=True, + ), "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"), "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"), "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"), diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1a5f8399b6ce..6fc9f9a08037 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -713,6 +713,7 @@ "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"), "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"), From 1713a094915d643c4ffaa25f421353108ee2049f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 17:17:02 +0200 Subject: [PATCH 192/774] [Rust Frontend] Forward per-request watermarking controls (#56338) Co-authored-by: Codex Signed-off-by: NickLucche --- docs/features/watermarking.md | 7 ++- rust/proto/inference.proto | 3 ++ .../src/protocol/sampling.rs | 9 ++++ .../engine-core-client/src/tests/client.rs | 2 + .../src/tests/python_compat.py | 2 + rust/src/server/src/grpc/convert.rs | 22 ++++++++++ .../src/routes/inference/generate/convert.rs | 24 +++++++++++ .../routes/openai/chat_completions/convert.rs | 1 + .../routes/openai/chat_completions/types.rs | 5 +++ .../src/routes/openai/completions/convert.rs | 1 + .../src/routes/openai/completions/types.rs | 4 ++ rust/src/server/src/routes/tests.rs | 43 +++++++++++++++++++ rust/src/text/src/lower.rs | 8 ++++ rust/src/text/src/request.rs | 3 ++ 14 files changed, 133 insertions(+), 1 deletion(-) diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index f9ceda054773..54c7e58ce7d9 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -29,7 +29,12 @@ from vllm import SamplingParams sampling_params = SamplingParams(watermarking=False) ``` -The OpenAI-compatible APIs accept the same `watermarking: false` request field. +The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the +same `watermarking: false` request field. The Rust token API accepts it in +`sampling_params`, and the Rust gRPC `GenerateRequest` accepts +the optional `watermarking` field. Omitting the field keeps watermarking enabled +when the engine has a watermark configuration. + Deployments that require watermarking must restrict this field to trusted callers, or strip and validate it at the ingress boundary, so untrusted clients cannot opt out. diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 858d8a7a7e01..46313f915a40 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -53,6 +53,9 @@ message GenerateRequest { // Loaded LoRA name; empty selects the base model. string lora_name = 15; + + // Apply the engine's configured watermark. Defaults to true when omitted. + optional bool watermarking = 16; } message RandomSampling { diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 19c579ac33a0..92f354be6fe0 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -20,6 +20,10 @@ fn default_temperature() -> f32 { 1.0 } +fn default_watermarking() -> bool { + true +} + fn default_max_tokens() -> u32 { 16 } @@ -68,6 +72,9 @@ pub struct EngineCoreSamplingParams { /// greedy sampling. #[serde(default = "default_temperature")] pub temperature: f32, + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_watermarking")] + pub watermarking: bool, /// Cumulative probability threshold for nucleus sampling. #[serde(default = "default_top_p")] pub top_p: f32, @@ -155,6 +162,7 @@ impl EngineCoreSamplingParams { pub fn for_test() -> Self { Self { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -229,6 +237,7 @@ mod tests { // Omitted fields -> Python defaults. assert_eq!(sampling.temperature, 1.0); + assert!(sampling.watermarking); assert_eq!(sampling.top_p, 1.0); assert_eq!(sampling.top_k, 0); assert_eq!(sampling.seed, None); diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 13525375f81d..27fbbe9abd93 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -151,6 +151,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest { prompt_token_ids: Some(vec![11, 22]), sampling_params: Some(EngineCoreSamplingParams { temperature: 0.8, + watermarking: false, top_p: 0.9, top_k: 8, max_tokens: 32, @@ -2662,6 +2663,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { sampling, EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index fd6cb0c86be3..20cec81808cc 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -34,6 +34,7 @@ class FinishReason(IntEnum): # Mirror of real SamplingParams; omit_defaults makes fixtures match real maps. class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True): temperature: float = 1.0 + watermarking: bool = True top_p: float = 1.0 top_k: int = 0 seed: int | None = None @@ -123,6 +124,7 @@ class EngineCoreOutputs( mm_features=None, sampling_params=EngineCoreSamplingParams( temperature=0.8, + watermarking=False, top_p=0.9, top_k=8, seed=None, diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index aea04495eb39..5f6ce5d42bc7 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -65,6 +65,7 @@ pub fn to_text_request( let mut sampling_params = build_sampling_params(req.temperature, sampling, decoding, stopping, response)?; + sampling_params.watermarking = req.watermarking.unwrap_or(true); // Thread KVCacheParameters → SamplingParams fields. if let Some(kv) = kv { @@ -521,6 +522,7 @@ impl ResponseOpts { #[cfg(test)] mod tests { + use prost::Message as _; use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{ FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, @@ -539,6 +541,26 @@ mod tests { } } + #[test] + fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() { + for watermarking in [None, Some(true), Some(false)] { + let request = pb::GenerateRequest { + watermarking, + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(decoded, stream, &["test-model".to_string()]) + .expect("convert request"); + assert_eq!( + text.sampling_params.watermarking, + watermarking.unwrap_or(true) + ); + } + } + } + #[test] fn temperature_propagates_from_top_level_request_field() { let req = pb::GenerateRequest { diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index cb06cfe33c43..08c4ce18be7e 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -164,6 +164,30 @@ mod tests { ); } + #[test] + fn prepare_generate_request_preserves_watermarking_defaults_and_opt_out() { + for watermarking in [None, Some(true), Some(false)] { + let mut body = json!({ + "token_ids": [11, 22], + "sampling_params": {} + }); + if let Some(watermarking) = watermarking { + body["sampling_params"]["watermarking"] = json!(watermarking); + } + let prepared = prepare_generate_request( + serde_json::from_value(body).expect("parse request"), + &served(&["test-model"]), + ResolvedRequestContext::default(), + None, + ) + .expect("prepare request"); + assert_eq!( + prepared.text_request.sampling_params.watermarking, + watermarking.unwrap_or(true) + ); + } + } + #[test] fn prepare_generate_request_forwards_thinking_token_budget() { let request: GenerateRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2c64ffbef48b..51b0fb68ad14 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -150,6 +150,7 @@ pub(super) fn prepare_chat_request( messages, sampling_params: SamplingParams { temperature: request.temperature, + watermarking: request.watermarking, top_p: request.top_p, top_k: request.top_k, seed: request.seed, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index e4da86c22d9b..a371462694ca 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -116,6 +116,10 @@ pub struct ChatCompletionRequest { pub user: Option, // -------- vLLM Sampling Parameters -------- + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_true")] + pub watermarking: bool, + /// Use beam search instead of sampling #[serde(default)] pub use_beam_search: bool, @@ -282,6 +286,7 @@ impl Default for ChatCompletionRequest { include_reasoning: true, parallel_tool_calls: None, user: None, + watermarking: true, use_beam_search: false, top_k: None, min_p: None, diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index d944dca65740..a4e8b6d3e423 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -115,6 +115,7 @@ pub(super) fn prepare_completion_request( mm_features: None, sampling_params: SamplingParams { temperature: request.temperature, + watermarking: request.watermarking, top_p: request.top_p, top_k: request.top_k, seed: request.seed, diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 8d93a89ed471..2a6901370682 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -90,6 +90,10 @@ pub struct CompletionRequest { pub user: Option, // -------- vLLM Sampling Parameters -------- + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_true")] + pub watermarking: bool, + /// Options for streaming response pub stream_options: Option, diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 6fbe22e501c3..4731958e2f87 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -2736,6 +2736,49 @@ async fn non_stream_chat_returns_json_response() { assert!(!message.contains_key("tool_calls"), "{json}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() { + for endpoint in ["/v1/chat/completions", "/v1/completions"] { + for stream in [false, true] { + for watermarking in [None, Some(true), Some(false)] { + let (mut app, engine_task) = test_app_with_backend_and_engine_request_check( + Arc::new(FakeChatBackend::new()), + move |request| { + let params = request.sampling_params.as_ref().expect("sampling params"); + assert_eq!(params.watermarking, watermarking.unwrap_or(true)); + }, + ) + .await; + + let mut body = json!({"stream": stream, "max_tokens": 2}); + if endpoint == "/v1/chat/completions" { + body["messages"] = json!([{"role": "user", "content": "hi"}]); + } else { + body["prompt"] = json!("hi"); + } + if let Some(watermarking) = watermarking { + body["watermarking"] = json!(watermarking); + } + let response = app + .call( + Request::builder() + .method("POST") + .uri(endpoint) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + to_bytes(response.into_body(), usize::MAX).await.expect("drain response"); + engine_task.await.expect("mock engine task"); + } + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_image_url_reaches_engine_mm_features() { diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index bff4a6d48bb1..466b9d5b31ad 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -92,6 +92,7 @@ pub fn lower_sampling_params( ) -> Result { let SamplingParams { temperature, + watermarking, top_p, top_k, seed, @@ -164,6 +165,7 @@ pub fn lower_sampling_params( let params = EngineCoreSamplingParams { temperature, + watermarking, top_p, top_k, seed, @@ -632,6 +634,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -686,6 +689,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -854,6 +858,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.6, + watermarking: true, top_p: 0.95, top_k: 20, seed: None, @@ -918,6 +923,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -990,6 +996,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.2, + watermarking: true, top_p: 0.3, top_k: 4, seed: None, @@ -1240,6 +1247,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.8, + watermarking: true, top_p: 0.9, top_k: 12, seed: None, diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 182e20c66022..f0dd2fb1f88c 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -69,6 +69,8 @@ pub struct SamplingParams { /// Controls randomness. Lower values are more deterministic; zero means /// greedy sampling. `None` means no explicit user override. pub temperature: Option, + /// Whether to apply the engine's configured watermark to this request. + pub watermarking: bool, /// Cumulative probability threshold for nucleus sampling. pub top_p: Option, /// Maximum number of top tokens to consider. `Some(0)` means all tokens. @@ -144,6 +146,7 @@ impl Default for SamplingParams { fn default() -> Self { Self { temperature: None, + watermarking: true, top_p: None, top_k: None, seed: None, From b7e8dd8f372c1324c376ba5dc03d92264eae8f24 Mon Sep 17 00:00:00 2001 From: Tomas Ruiz Date: Mon, 14 Sep 2026 17:29:13 +0200 Subject: [PATCH 193/774] [watermarking hardening] Add e2e test and basic GSM8K quality tests (#56309) Signed-off-by: Tomas Ruiz Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/lm_eval.yaml | 16 +++++++ .../gsm8k/configs/Qwen3.5-4B-watermark.yaml | 22 ++++++++++ .../evals/gsm8k/configs/models-watermark.txt | 1 + tests/watermarking/test_watermarking_e2e.py | 42 +++++++++++++++++++ 4 files changed, 81 insertions(+) create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml create mode 100644 tests/evals/gsm8k/configs/models-watermark.txt create mode 100644 tests/watermarking/test_watermarking_e2e.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 3c6561196610..d9afefc82e4f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -30,6 +30,22 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py +- label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking" + device: h200_35gb + key: lm-eval-watermarking + # Recorded runs took 80–140 seconds + timeout_in_minutes: 10 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt + # - label: LM Eval Large Models (4xA100) # key: lm-eval-large-models-4xa100 # device: a100 diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml new file mode 100644 index 000000000000..bdd42a0f1a78 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml @@ -0,0 +1,22 @@ +model_name: "Qwen/Qwen3.5-4B" +# full GSM8K set is 1319 questions +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 128 +max_tokens: 1024 +use_chat_completions: true +temperature: 0.7 +seed: 42 + +# Calibrated by running non-watermaked 13 times +# the threshold is the average accuracy, and the tolerance is set to 3 sigma (standard deviations) +accuracy_threshold: 0.885 +tolerance: 0.015 + +server_args: >- + --language-model-only + --max-model-len 8192 + --max-num-seqs 128 + --gpu-memory-utilization 0.4 + --default-chat-template-kwargs '{"enable_thinking": false}' + --watermark-config '{"key": 42, "context_width": 4}' diff --git a/tests/evals/gsm8k/configs/models-watermark.txt b/tests/evals/gsm8k/configs/models-watermark.txt new file mode 100644 index 000000000000..a761b41c1b3f --- /dev/null +++ b/tests/evals/gsm8k/configs/models-watermark.txt @@ -0,0 +1 @@ +Qwen3.5-4B-watermark.yaml diff --git a/tests/watermarking/test_watermarking_e2e.py b/tests/watermarking/test_watermarking_e2e.py new file mode 100644 index 000000000000..c73d48f6b6b7 --- /dev/null +++ b/tests/watermarking/test_watermarking_e2e.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm import SamplingParams +from vllm.platforms import current_platform +from vllm.v1.watermarking import GumbelWatermarkDetector + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_llm_generated_sequence_is_watermarked(vllm_runner): + """Engine-level watermarking must produce tokens the detector recognizes.""" + watermark_config = {"key": 42, "context_width": 4} + detector = GumbelWatermarkDetector(**watermark_config) + + runner = vllm_runner( + "facebook/opt-125m", + watermark_config=watermark_config, + seed=0, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.2, + ) + request = "Tell me a story about an explorer who discovers a mysterious island." + params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True) + params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False) + + with runner: + output_use_wm = runner.llm.generate(request, params_use_wm) + output_no_wm = runner.llm.generate(request, params_no_wm) + + result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids)) + result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids)) + + assert result_use_wm.is_watermarked, result_use_wm + recorded_p_value = 9.83e-20 + assert result_use_wm.p_value < recorded_p_value * 10 # tolerance + + assert not result_no_wm.is_watermarked, result_no_wm From 435c96f9dbdd29258cb8e0f433c5b54a00cf6b16 Mon Sep 17 00:00:00 2001 From: Thillai Chithambaram <79466435+thillai-c@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:04:55 -0400 Subject: [PATCH 194/774] [V1][Metrics] Support Sliding Window Attention (SWA) and hybrid layers in MFU/MBU estimation (#55624) Signed-off-by: Thillai Chithambaram Co-authored-by: Mark McLoughlin --- tests/v1/metrics/test_perf_metrics.py | 374 +++++++++++++++++++++++++- vllm/v1/metrics/perf.py | 314 +++++++++++++++++---- 2 files changed, 640 insertions(+), 48 deletions(-) diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index ab30f1bb9e26..c94ec1b701ca 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -8,6 +8,7 @@ from types import SimpleNamespace import pytest +from transformers import Gemma2Config, MistralConfig from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config from transformers.models.llama4.configuration_llama4 import ( Llama4Config, @@ -31,6 +32,7 @@ MLAAttentionMetrics, ModelMetrics, ParsedArgs, + SlidingWindowAttentionParser, UnembedMetrics, ) @@ -38,7 +40,8 @@ class MockModelConfig: """Mock ModelConfig that implements the getter methods used by parsers.""" - def __init__(self, hf_config, dtype): + def __init__(self, hf_config, dtype, disable_sliding_window=False): + """Initialize MockModelConfig with HF config and options.""" self.hf_config = hf_config self.hf_text_config = get_hf_text_config(hf_config) convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get( @@ -49,8 +52,10 @@ def __init__(self, hf_config, dtype): ).convert() self.dtype = dtype self.is_attention_free = False + self.disable_sliding_window = disable_sliding_window def __getattr__(self, name): + """Delegate attribute lookups to ModelConfig.""" # 1. Check if ModelConfig actually has this attribute if not hasattr(ModelConfig, name): raise AttributeError( @@ -86,9 +91,13 @@ def create_mock_vllm_config( tensor_parallel_size=1, pipeline_parallel_size=1, enable_expert_parallel=False, + disable_sliding_window=False, ) -> SimpleNamespace: + """Create a mock VllmConfig object for unit tests.""" vllm_config = SimpleNamespace() - vllm_config.model_config = MockModelConfig(hf_config, model_dtype) + vllm_config.model_config = MockModelConfig( + hf_config, model_dtype, disable_sliding_window=disable_sliding_window + ) vllm_config.cache_config = SimpleNamespace() vllm_config.cache_config.cache_dtype = cache_dtype @@ -1336,3 +1345,364 @@ def test_mla_attention_scaling_with_layers(): assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx) assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx) assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx) + + +#### Sliding Window Attention Tests #### + + +def test_sliding_window_parser_no_swa(): + """Test parser defaults to no SWA when sliding_window is unset.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=12, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 12 + result = parser.parse(args, vllm_config) + + assert result.sliding_window is None + assert result.num_swa_layers == 0 + + +def test_sliding_window_parser_uniform_swa(): + """Test SlidingWindowAttentionParser for full SWA model (e.g. Mistral).""" + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=32, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 32 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 4096 + assert result.num_swa_layers == 32 + + +def test_sliding_window_parser_explicit_layer_types(): + """Test SlidingWindowAttentionParser with explicit layer_types list.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=6, + ) + hf_config.sliding_window = 2048 + hf_config.layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ] + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 6 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 2048 + assert result.num_swa_layers == 3 + + +def test_sliding_window_parser_gemma2(): + """Test SlidingWindowAttentionParser detects alternating SWA for Gemma 2.""" + hf_config = Gemma2Config( + hidden_size=2304, + num_attention_heads=8, + num_key_value_heads=4, + num_hidden_layers=26, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 26 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 4096 + # Gemma 2 alternates: even layers are sliding window (26 layers -> 13 SWA) + assert result.num_swa_layers == 13 + + +def test_sliding_window_parser_disabled(): + """Test SlidingWindowAttentionParser respects disable_sliding_window.""" + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=32, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config, disable_sliding_window=True) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 32 + result = parser.parse(args, vllm_config) + + assert result.sliding_window is None + assert result.num_swa_layers == 0 + + +def test_attention_metrics_sliding_window_flops_bounds(): + """Test that SWA clamps attention core FLOPs when context_len > sliding_window.""" + sw = 1024 + dense_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=8, + sliding_window=None, + ) + swa_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=8, + sliding_window=sw, + ) + + dense_vllm = create_mock_vllm_config(dense_config) + swa_vllm = create_mock_vllm_config(swa_config) + + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) + + # 1. When context <= sliding_window: FLOPs should be identical + short_ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=512, is_prefill=False + ) + assert swa_metrics.get_num_flops(short_ctx) == dense_metrics.get_num_flops( + short_ctx + ) + + # 2. When context > sliding_window: FLOPs for SWA must be strictly smaller + long_ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=8192, is_prefill=False + ) + dense_flops = dense_metrics.get_num_flops_breakdown(long_ctx) + swa_flops = swa_metrics.get_num_flops_breakdown(long_ctx) + + # Projections (qkv_proj, out_proj) are unaffected by sliding window + assert swa_flops["qkv_proj"] == dense_flops["qkv_proj"] + assert swa_flops["out_proj"] == dense_flops["out_proj"] + + # Attention core FLOPs are clamped: 8192 vs 1024 (8x difference) + assert swa_flops["attn_qk"] < dense_flops["attn_qk"] + assert swa_flops["attn_av"] < dense_flops["attn_av"] + assert swa_flops["attn_qk"] * 8 == dense_flops["attn_qk"] + assert swa_flops["attn_av"] * 8 == dense_flops["attn_av"] + + +def test_attention_metrics_sliding_window_read_bytes(): + """Test that SWA bounds decode KV-cache read bytes to sliding_window.""" + sw = 2048 + dense_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=4, + num_hidden_layers=4, + sliding_window=None, + ) + swa_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=4, + num_hidden_layers=4, + sliding_window=sw, + ) + + dense_vllm = create_mock_vllm_config(dense_config) + swa_vllm = create_mock_vllm_config(swa_config) + + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) + + # Long decode context: 8192 tokens + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=8192, is_prefill=False + ) + dense_reads = dense_metrics.get_read_bytes_breakdown(ctx) + swa_reads = swa_metrics.get_read_bytes_breakdown(ctx) + + # Attention input read traffic is lower for SWA + assert swa_reads["attn_input"] < dense_reads["attn_input"] + + # Weight reads and projection inputs are identical + assert swa_reads["qkv_weight"] == dense_reads["qkv_weight"] + assert swa_reads["out_weight"] == dense_reads["out_weight"] + assert swa_reads["qkv_input"] == dense_reads["qkv_input"] + + +def test_attention_metrics_hybrid_alternating_layers(): + """Test hybrid model where half layers are full attention and half are SWA.""" + sw = 1024 + hybrid_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=4, + ) + hybrid_config.sliding_window = sw + hybrid_config.layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ] + + vllm_config = create_mock_vllm_config(hybrid_config) + metrics = AttentionMetrics.from_vllm_config(vllm_config) + + assert metrics.num_swa_layers == 2 + assert metrics.sliding_window == sw + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + flops = metrics.get_num_flops_breakdown(ctx) + + # 2 full layers attend to 4096; 2 SWA layers attend to 1024 + # Expected TC total = 2 * 4096 + 2 * 1024 = 8192 + 2048 = 10240 + # Formula: 2 * q * d * (L_full * 4096 + L_swa * 1024) + q = metrics.num_attention_heads + d = metrics.head_dim + expected_core_flops = 2 * q * d * (2 * 4096 + 2 * 1024) + assert flops["attn_qk"] == expected_core_flops + assert flops["attn_av"] == expected_core_flops + + +def test_attention_metrics_sliding_window_per_gpu(): + """Test that SWA metrics scale properly with tensor and pipeline parallelism.""" + sw = 1024 + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=16, + sliding_window=sw, + ) + single_vllm = create_mock_vllm_config(hf_config) + tp_vllm = create_mock_vllm_config(hf_config, tensor_parallel_size=2) + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) + + single_metrics = AttentionMetrics.from_vllm_config(single_vllm) + tp_metrics = AttentionMetrics.from_vllm_config(tp_vllm) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + + # Both TP=2 and PP=2 should halve the per-GPU attention core FLOPs + single_flops = single_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + tp_flops = tp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + + assert tp_flops["attn_qk"] == single_flops["attn_qk"] // 2 + assert pp_flops["attn_qk"] == single_flops["attn_qk"] // 2 + + +def test_attention_metrics_hybrid_sliding_window_per_gpu(): + """Test that hybrid SWA models preserve L_full + L_swa == L under PP.""" + sw = 1024 + hybrid_config = Gemma2Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=30, + sliding_window=sw, + ) + pp_vllm = create_mock_vllm_config(hybrid_config, pipeline_parallel_size=2) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + pp_read_bytes = pp_metrics.get_read_bytes_breakdown(ctx, per_gpu=True) + + # Under PP=2, each global count is halved: 30 layers -> 15 per GPU, and the + # 15 SWA / 15 full split becomes 7.5 each. Gemma 2 alternates, so an even + # split is the correct average stage; deriving it after reducing L instead + # produced a lopsided 8/7. + q = pp_metrics.num_attention_heads + d = pp_metrics.head_dim + expected_pp_attn = int(2 * q * d * (7.5 * 4096 + 7.5 * sw)) + assert pp_flops["attn_qk"] == expected_pp_attn + assert pp_flops["attn_av"] == expected_pp_attn + assert "attn_input" in pp_read_bytes + + +def test_attention_metrics_minority_swa_survives_pipeline_parallelism(): + """A minority SWA sublayer must not round away under PP. + + Regression test for a 6-layer model with a single sliding-window layer at + pp_size=2: deriving the split after reducing L gave + round(3 * 1 / 6) == 0, billing all three per-GPU layers as full attention + and dropping the sliding-window contribution entirely. + """ + sw = 512 + context_len = 4096 + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=6, + ) + hf_config.sliding_window = sw + hf_config.layer_types = ["sliding_attention"] + ["full_attention"] * 5 + + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + assert pp_metrics.num_swa_layers == 1 + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=context_len, is_prefill=False + ) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + + # 5 full and 1 SWA layer, each halved by PP=2: 2.5 full and 0.5 SWA. + q = pp_metrics.num_attention_heads + d = pp_metrics.head_dim + expected = int(2 * q * d * (2.5 * context_len + 0.5 * sw)) + assert pp_flops["attn_qk"] == expected + assert pp_flops["attn_av"] == expected + + # The sliding-window layer must change the result, i.e. it is not being + # billed at the full context length. + all_full = int(2 * q * d * 3 * context_len) + assert pp_flops["attn_qk"] < all_full + + +def test_attention_metrics_per_gpu_layers_sum_to_whole_model(): + """Per-GPU counts times pp_size must account for every layer. + + Flooring the layer count per stage silently drops the remainder, so a model + whose layers do not divide evenly across stages under-reports its FLOPs. + """ + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=7, + ) + pp_size = 2 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=pp_size) + metrics = AttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=8, context_len=1024, is_prefill=True + ) + per_gpu = metrics.get_num_flops_breakdown(ctx, per_gpu=True) + whole = metrics.get_num_flops_breakdown(ctx, per_gpu=False) + + for key, value in whole.items(): + assert per_gpu[key] * pp_size == pytest.approx(value, rel=1e-9) diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 75524ef12858..ab90ffcbe142 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -10,7 +10,7 @@ import time from abc import ABC, abstractmethod from collections.abc import Iterable -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from typing import Any, Protocol import prometheus_client @@ -124,8 +124,12 @@ class ExecutionContext: decode_context_len: int = 0 # sum of context_len for decode requests decode_token_context_product: int = 0 # sum of (num_tokens * context_len) + # Per-request records: (num_tokens, context_len, is_prefill) + requests: list[tuple[int, int, bool]] = field(default_factory=list) + def add(self, num_tokens: int, context_len: int, is_prefill: bool) -> None: """Add a single request's statistics to this batch context.""" + self.requests.append((num_tokens, context_len, is_prefill)) if is_prefill: self.num_prefill_requests += 1 self.prefill_num_tokens += num_tokens @@ -141,9 +145,40 @@ def total_num_tokens(self) -> int: """Total number of tokens across all requests in the batch.""" return self.prefill_num_tokens + self.decode_num_tokens - def total_token_context_product(self) -> int: - """Total sum of (num_tokens * context_len) across all requests.""" - return self.prefill_token_context_product + self.decode_token_context_product + def total_token_context_product(self, sliding_window: int | None = None) -> int: + """Total sum of (num_tokens * context_len) across all requests. + + If sliding_window is provided, context length per request is clamped to + sliding_window. + """ + if sliding_window is None: + return ( + self.prefill_token_context_product + self.decode_token_context_product + ) + return sum( + num_tokens * min(context_len, sliding_window) + for num_tokens, context_len, _ in self.requests + ) + + def prefill_context_len_for_window(self, sliding_window: int | None = None) -> int: + """Sum of context_len for prefill requests, clamped to sliding_window if set.""" + if sliding_window is None: + return self.prefill_context_len + return sum( + min(context_len, sliding_window) + for _, context_len, is_prefill in self.requests + if is_prefill + ) + + def decode_context_len_for_window(self, sliding_window: int | None = None) -> int: + """Sum of context_len for decode requests, clamped to sliding_window if set.""" + if sliding_window is None: + return self.decode_context_len + return sum( + min(context_len, sliding_window) + for _, context_len, is_prefill in self.requests + if not is_prefill + ) def num_logits_tokens(self) -> int: """Number of tokens that require logits computation (unembedding). @@ -214,6 +249,7 @@ def add_parser(self, parser: Parser) -> None: self.parsers.append(parser) def parse(self, vllm_config: VllmConfig) -> ParsedArgs: + """Apply all parsers in sequence and return the accumulated ParsedArgs.""" args = ParsedArgs() for parser in self.parsers: args = parser.parse(args, vllm_config) @@ -234,7 +270,9 @@ class ComponentMetrics(BaseModel, ABC): @classmethod @abstractmethod - def component_type(cls) -> str: ... + def component_type(cls) -> str: + """Return the unique string key identifying this component in the registry.""" + ... @classmethod @abstractmethod @@ -273,17 +311,23 @@ def registered_metrics(cls) -> Iterable[type["ComponentMetrics"]]: @abstractmethod def get_num_flops_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation FLOPs breakdown for this component.""" + ... @abstractmethod def get_read_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation read-memory-traffic breakdown for this component.""" + ... @abstractmethod def get_write_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation write-memory-traffic breakdown for this component.""" + ... def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int: return sum(self.get_num_flops_breakdown(ctx, per_gpu).values()) @@ -306,6 +350,7 @@ class BaseConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse base model dimensions, dtype byte sizes, and parallelism config.""" model_config = vllm_config.model_config args.vocab_size = model_config.get_vocab_size() @@ -357,6 +402,7 @@ class BaseAttentionConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse KV head count, head dimension, and KV-cache dtype byte size.""" model_config = vllm_config.model_config args.num_key_value_heads = model_config.get_total_num_kv_heads() @@ -371,6 +417,54 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: return args +class SlidingWindowAttentionParser(Parser): + """ + Parses sliding window attention configuration and layer breakdown. + Provides: sliding_window, num_swa_layers + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Determine sliding_window size and count of SWA layers. + + Detects per-layer type lists (e.g. Llama-4), Gemma 2 alternating + pattern, and uniform SWA models (Mistral, Qwen2.5). + """ + model_config = vllm_config.model_config + + sliding_window: int | None = None + if not getattr(model_config, "disable_sliding_window", False): + sliding_window = model_config.get_sliding_window() + if sliding_window is None and hasattr(vllm_config, "cache_config"): + sliding_window = getattr( + vllm_config.cache_config, "sliding_window", None + ) + + args.sliding_window = sliding_window + + num_layers = args.num_hidden_layers + if sliding_window is not None and sliding_window > 0: + hf_text_config = getattr(model_config, "hf_text_config", None) + layer_types = getattr(hf_text_config, "layer_types", None) + if layer_types is not None and isinstance(layer_types, list): + num_swa = sum( + 1 for lt in layer_types[:num_layers] if "sliding" in lt.lower() + ) + args.num_swa_layers = num_swa + elif ( + getattr(getattr(model_config, "hf_config", None), "model_type", None) + == "gemma2" + ): + # Gemma 2 alternates sliding window: even layers (0, 2, ...) are sliding + args.num_swa_layers = (num_layers + 1) // 2 + else: + # Full SWA model (e.g. Mistral, Qwen2.5 with sliding_window) + args.num_swa_layers = num_layers + else: + args.num_swa_layers = 0 + + return args + + class AttentionQuantizationConfigParser(Parser): """ Parses quantization configuration for attention layers. @@ -378,6 +472,7 @@ class AttentionQuantizationConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Override weight_byte_size based on the active quantization method.""" cfg = vllm_config.quant_config if cfg is None: @@ -401,6 +496,7 @@ class AttentionDetectionParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Raise InvalidComponent for MLA models so AttentionMetrics is skipped.""" if vllm_config.model_config.is_deepseek_mla: raise InvalidComponent( "Model uses MLA attention; use MLAAttentionMetrics instead" @@ -425,53 +521,129 @@ class AttentionMetrics(ComponentMetrics): # From BaseConfig Parser, overridden by AttentionQuantizationConfigParser weight_byte_size: int | float = Field(..., gt=0) - # TODO: discern cases where we have mixture of different attention layer types - # such as SWA, MLA, etc. + # From SlidingWindowAttentionParser + sliding_window: int | None = Field(None) + num_swa_layers: int = Field(0, ge=0) @classmethod def component_type(cls) -> str: + """Return the component registry key for standard (non-MLA) attention.""" return "attn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for AttentionMetrics. + + Parsers run in order: MLA detection guard, base model config, + attention-specific config, sliding-window breakdown, quantization + weight-size override. + """ return ParserChain( AttentionDetectionParser(), BaseConfigParser(), BaseAttentionConfigParser(), + SlidingWindowAttentionParser(), AttentionQuantizationConfigParser(), ) + def _layer_counts(self, per_gpu: bool) -> tuple[float, float, float]: + """Return ``(L, L_full, L_swa)``, the layer counts to bill for. + + Per-GPU counts divide each global count by ``pp_size`` independently. + Deriving the sublayer split after reducing ``L`` instead rounds a + minority sublayer away: a 6-layer model with 1 sliding-window layer at + ``pp_size=2`` yields ``round(3 * 1 / 6) == 0``, billing all three layers + as full attention. Dividing each count separately keeps + ``L_full + L_swa == L`` and keeps the per-GPU values summing back to the + whole model across stages. + + The counts are fractional because these metrics are built from config + alone (see ``from_vllm_config``); no pipeline rank is available, so a + per-GPU figure is the average stage rather than a specific one. A model + whose layers do not divide evenly across stages has no single integer + answer here. + """ + L = float(self.num_hidden_layers) + L_swa = float(self.num_swa_layers) + L_full = L - L_swa + + if per_gpu: + L /= self.pp_size + L_swa /= self.pp_size + L_full /= self.pp_size + + return L, L_full, L_swa + def get_num_flops_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - L, D, q, kv, d = ( - self.num_hidden_layers, + """Compute FLOPs breakdown for attention layers. + + Accounts for hybrid models by splitting total layers into full-attention + (L_full) and sliding-window (L_swa) sublayers. See ``_layer_counts`` for + how those are scaled under pipeline parallelism. + + Args: + ctx: Execution context describing the current batch. + per_gpu: If True, scale counts down by tensor/pipeline parallelism. + + Returns: + Dict with keys: qkv_proj, out_proj, attn_qk, attn_av. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, self.head_dim, ) T = ctx.total_num_tokens() - TC = ctx.total_token_context_product() + TC_full = ctx.total_token_context_product() + + L, L_full, L_swa = self._layer_counts(per_gpu) if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) - return { - "qkv_proj": 2 * T * D * (q + 2 * kv) * d * L, - "attn_qk": 2 * q * TC * d * L, - "attn_av": 2 * q * TC * d * L, - "out_proj": 2 * T * D * q * d * L, + flops = { + "qkv_proj": int(2 * T * D * (q + 2 * kv) * d * L), + "out_proj": int(2 * T * D * q * d * L), } + attn_qk = 0.0 + attn_av = 0.0 + if L_full > 0: + attn_qk += 2 * q * TC_full * d * L_full + attn_av += 2 * q * TC_full * d * L_full + + if L_swa > 0 and self.sliding_window is not None: + TC_swa = ctx.total_token_context_product(self.sliding_window) + attn_qk += 2 * q * TC_swa * d * L_swa + attn_av += 2 * q * TC_swa * d * L_swa + + flops["attn_qk"] = int(attn_qk) + flops["attn_av"] = int(attn_av) + + return flops + def get_read_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - L, D, q, kv, d = ( - self.num_hidden_layers, + """Compute read-memory-traffic breakdown for attention layers. + + Splits layers into full-attention and sliding-window sublayers using the + same counts as get_num_flops_breakdown; see ``_layer_counts``. + + Args: + ctx: Execution context describing the current batch. + per_gpu: If True, scale counts down by tensor/pipeline parallelism. + + Returns: + Dict with keys: qkv_input, qkv_weight, attn_input (conditional), + out_input, out_weight. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, @@ -479,35 +651,63 @@ def get_read_bytes_breakdown( ) T = ctx.total_num_tokens() + L, L_full, L_swa = self._layer_counts(per_gpu) + if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) read_bytes = {} - read_bytes["qkv_input"] = T * D * self.activation_byte_size * L + read_bytes["qkv_input"] = int(T * D * self.activation_byte_size * L) read_bytes["qkv_weight"] = int(D * (q + 2 * kv) * d * self.weight_byte_size * L) - # Attention input reads differ between prefill and decode - # Prefill: read Q, K, V activations (all in activation_byte_size) - if ctx.prefill_num_tokens > 0: - read_bytes["attn_input"] = ( - (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv) - * d - * self.activation_byte_size - * L - ) + attn_input = 0.0 - # Decode: read Q activations + read K, V from cache (in cache_byte_size) - if ctx.decode_num_tokens > 0: - read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + ( - ctx.decode_num_tokens * q * d * self.activation_byte_size * L - + 2 * ctx.decode_context_len * kv * d * self.cache_byte_size * L - ) + # Full attention layers + if L_full > 0: + if ctx.prefill_num_tokens > 0: + attn_input += ( + (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv) + * d + * self.activation_byte_size + * L_full + ) + if ctx.decode_num_tokens > 0: + attn_input += ( + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_full + + 2 + * ctx.decode_context_len + * kv + * d + * self.cache_byte_size + * L_full + ) + + # Sliding window attention layers + if L_swa > 0: + sw = self.sliding_window + prefill_sw_len = ctx.prefill_context_len_for_window(sw) + decode_sw_len = ctx.decode_context_len_for_window(sw) + + if ctx.prefill_num_tokens > 0: + attn_input += ( + (ctx.prefill_num_tokens * q + 2 * prefill_sw_len * kv) + * d + * self.activation_byte_size + * L_swa + ) + if ctx.decode_num_tokens > 0: + attn_input += ( + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_swa + + 2 * decode_sw_len * kv * d * self.cache_byte_size * L_swa + ) + + if ctx.prefill_num_tokens > 0 or ctx.decode_num_tokens > 0: + read_bytes["attn_input"] = int(attn_input) - read_bytes["out_input"] = T * q * d * self.activation_byte_size * L + read_bytes["out_input"] = int(T * q * d * self.activation_byte_size * L) read_bytes["out_weight"] = int(q * d * D * self.weight_byte_size * L) return read_bytes @@ -515,9 +715,13 @@ def get_read_bytes_breakdown( def get_write_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - """Calculate write memory traffic for attention layers.""" - L, D, q, kv, d = ( - self.num_hidden_layers, + """Calculate write memory traffic for attention layers. + + Writes do not depend on the attention window, so only the total layer + count is needed; it comes from ``_layer_counts`` so all three breakdowns + bill the same number of layers. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, @@ -525,16 +729,17 @@ def get_write_bytes_breakdown( ) T = ctx.total_num_tokens() + L, _, _ = self._layer_counts(per_gpu) + if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) return { - "qkv_output": T * (q + 2 * kv) * d * self.activation_byte_size * L, - "kv_cache": 2 * T * kv * d * self.cache_byte_size * L, - "out_output": T * D * self.activation_byte_size * L, + "qkv_output": int(T * (q + 2 * kv) * d * self.activation_byte_size * L), + "kv_cache": int(2 * T * kv * d * self.cache_byte_size * L), + "out_output": int(T * D * self.activation_byte_size * L), } @@ -549,6 +754,7 @@ class MLADetectionParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Raise InvalidComponent if the model does not use MLA attention.""" if not vllm_config.model_config.is_deepseek_mla: raise InvalidComponent("Model does not use MLA attention") return args @@ -562,6 +768,7 @@ class MLAConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse MLA-specific compression dimensions and KV-cache dtype byte size.""" model_config = vllm_config.model_config cfg = model_config.hf_text_config @@ -616,10 +823,16 @@ class MLAAttentionMetrics(ComponentMetrics): @classmethod def component_type(cls) -> str: + """Return the component registry key for MLA attention.""" return "mla_attn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for MLAAttentionMetrics. + + Parsers run in order: MLA detection guard, base model config, + MLA-specific config, quantization weight-size override. + """ return ParserChain( MLADetectionParser(), BaseConfigParser(), @@ -923,6 +1136,7 @@ class FfnQuantizationConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Override weight_byte_size based on the active FFN quantization method.""" cfg = vllm_config.quant_config if cfg is None: @@ -980,10 +1194,16 @@ def validate_moe_fields(self) -> Self: @classmethod def component_type(cls) -> str: + """Return the component registry key for FFN layers.""" return "ffn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for FfnMetrics. + + Parsers run in order: base model config, FFN parallelism, base FFN + config, MoE interleave step, MoE layer frequency, quantization override. + """ return ParserChain( BaseConfigParser(), FfnParallelParser(), @@ -1207,10 +1427,12 @@ class UnembedMetrics(ComponentMetrics): @classmethod def component_type(cls) -> str: + """Return the component registry key for the unembedding (LM head) layer.""" return "unembed" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for UnembedMetrics (base model config only).""" return ParserChain( BaseConfigParser(), ) From 03dc26e639e1927cda51d7b62639a11a354e1c0d Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Mon, 14 Sep 2026 19:26:15 +0200 Subject: [PATCH 195/774] [Perf][Kernel][Quantization] Fuse ReLU2 with static FP8 activation quantization (#53793) Signed-off-by: mgoin Signed-off-by: ElizaWszola Signed-off-by: snordmann Co-authored-by: mgoin Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: ElizaWszola Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_silu_mul_quant_manual_fusion.py | 11 +- .../fusion/test_quant_activation_contract.py | 5 +- tests/kernels/test_fused_quant_activation.py | 30 ++- tests/kernels/test_relu2_fp8_quant.py | 229 ++++++++++++++++++ tests/lora/test_layers.py | 10 + .../test_nemotron_h_quantization.py | 27 +++ vllm/lora/layers/base_linear.py | 4 + .../layers/fusion/fused_act_quant.py | 66 ++++- .../layers/fusion/quant_activation.py | 22 +- .../layers/fusion/relu2_fp8_quant.py | 76 ++++++ vllm/model_executor/layers/linear.py | 2 + vllm/model_executor/models/nemotron_h.py | 9 +- 12 files changed, 459 insertions(+), 32 deletions(-) create mode 100644 tests/kernels/test_relu2_fp8_quant.py create mode 100644 vllm/model_executor/layers/fusion/relu2_fp8_quant.py diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py index 5ee749659b3f..ee9496084420 100644 --- a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py @@ -34,6 +34,7 @@ from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, expose_input_quant_key, + get_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, @@ -45,7 +46,7 @@ # Mock linear layer for testing fusion paths that don't have real kernel support class MockLinearForFusion(torch.nn.Module): - """Mock linear layer that exposes input_quant_key for fusion testing.""" + """Mock linear layer that exposes an input key for fusion testing.""" def __init__( self, @@ -55,7 +56,7 @@ def __init__( input_global_scale_inv=None, ): super().__init__() - self.input_quant_key = quant_key + self._input_quant_key = quant_key if input_scale is not None: self.input_scale = input_scale if input_global_scale is not None: @@ -118,7 +119,7 @@ def test_manual_fusion_fp8_static_with_linear( # Enable fusion expose_input_quant_key(fp8_linear, fp8_linear.kernel) - if not hasattr(fp8_linear, "input_quant_key"): + if get_input_quant_key(fp8_linear) is None: pytest.skip( f"Kernel {force_kernel.__name__} doesn't support input_quant_key" ) @@ -315,7 +316,7 @@ def test_manual_fusion_nvfp4_dynamic( envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" ) def test_manual_fusion_fallback_no_key(dtype: torch.dtype): - """Test that maybe_fused_act_quant falls back when no input_quant_key.""" + """Test that maybe_fused_act_quant falls back without an input key.""" torch.set_default_device("cuda") torch.set_default_dtype(dtype) @@ -327,7 +328,7 @@ def test_manual_fusion_fallback_no_key(dtype: torch.dtype): with set_current_vllm_config(config): silu_and_mul = SiluAndMul() - # Linear without input_quant_key attribute + # Linear without an input quantization key mock_linear = torch.nn.Linear(128, 128) result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index f698a272f5ea..ab12393692a5 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -42,6 +42,7 @@ QuantizedActivation, as_quantized_activation, expose_input_quant_key, + get_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, @@ -118,13 +119,13 @@ def test_bridge_marks_supporting_and_skips_others(): supported = _probe(FlashInferCutlassNvFp4LinearKernel) layer = torch.nn.Module() expose_input_quant_key(layer, supported) - assert layer.input_quant_key == kNvfp4Dynamic + assert get_input_quant_key(layer) == kNvfp4Dynamic unsupported = _probe(FlashInferTrtllmNvFp4LinearKernel) assert unsupported.input_quant_key() is None layer = torch.nn.Module() expose_input_quant_key(layer, unsupported) - assert not hasattr(layer, "input_quant_key") + assert get_input_quant_key(layer) is None def test_as_quantized_activation_validates_key(): diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 267484e43361..321cb6978e3d 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -91,7 +91,7 @@ class MockLinearFp8Static(torch.nn.Module): def __init__(self, input_scale: torch.Tensor): super().__init__() - self.input_quant_key = kFp8StaticTensorSym + self._input_quant_key = kFp8StaticTensorSym self.input_scale = input_scale @@ -100,15 +100,21 @@ class MockLinearFp8Dynamic128(torch.nn.Module): def __init__(self): super().__init__() - self.input_quant_key = kFp8Dynamic128Sym + self._input_quant_key = kFp8Dynamic128Sym class MockLinearNoQuant(torch.nn.Module): - """Mock linear layer with no input_quant_key (no fusion).""" + """Mock linear layer with no input quantization key.""" pass +class MockLinearRequiresUnquantized(MockLinearFp8Static): + """Mock consumer that needs the original activation for another branch.""" + + requires_unquantized_input = True + + @pytest.mark.parametrize("num_tokens", [1, 16, 128]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @@ -199,7 +205,7 @@ def test_maybe_fused_act_quant_fallback( hidden_size: int, dtype: torch.dtype, ) -> None: - """Test maybe_fused_act_quant falls back when no input_quant_key.""" + """Test maybe_fused_act_quant falls back without an input quantization key.""" device = "cuda:0" torch.set_default_device(device) @@ -216,3 +222,19 @@ def test_maybe_fused_act_quant_fallback( ref_out = act_fn(x) torch.testing.assert_close(result, ref_out) + + +@torch.inference_mode() +def test_maybe_fused_act_quant_preserves_required_unquantized_input( + default_vllm_config, +) -> None: + device = "cuda:0" + act_fn = SiluAndMul() + scale = torch.tensor([0.5], device=device, dtype=torch.float32) + linear = MockLinearRequiresUnquantized(scale) + x = torch.randn(17, 256, dtype=torch.bfloat16, device=device) + + result = maybe_fused_act_quant(act_fn, x, linear) + + assert isinstance(result, torch.Tensor) + torch.testing.assert_close(result, act_fn(x)) diff --git a/tests/kernels/test_relu2_fp8_quant.py b/tests/kernels/test_relu2_fp8_quant.py new file mode 100644 index 000000000000..cb605914c7f5 --- /dev/null +++ b/tests/kernels/test_relu2_fp8_quant.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config +from vllm.config.vllm import OptimizationLevel +from vllm.model_executor.layers.activation import ReLUSquaredActivation +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform + +requires_sm90 = pytest.mark.skipif( + not (current_platform.is_cuda() and current_platform.has_device_capability(90)), + reason="This test requires SM90 or newer", +) + + +class _StaticFp8Linear(torch.nn.Module): + def __init__(self, input_scale: torch.Tensor) -> None: + super().__init__() + self._input_quant_key = kFp8StaticTensorSym + self.input_scale = input_scale + + +def _assert_fp8_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + assert actual.dtype == expected.dtype == current_platform.fp8_dtype() + torch.testing.assert_close( + actual.view(torch.uint8), expected.view(torch.uint8), rtol=0, atol=0 + ) + + +def _assert_quantized_activation( + result: QuantizedActivation, x: torch.Tensor, scale: torch.Tensor +) -> None: + assert result.scale is scale + assert result.orig_dtype == x.dtype + assert result.orig_shape == x.shape + assert result.quant_key == kFp8StaticTensorSym + + +@pytest.fixture(scope="module") +def o2_relu2_fp8_ops(): + config = VllmConfig( + optimization_level=OptimizationLevel.O2, + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + with set_current_vllm_config(config): + yield ( + ReLUSquaredActivation(), + QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR), + ) + + +@requires_sm90 +@pytest.mark.parametrize("shape", [(1, 1), (1, 16), (17, 5120), (128, 10240)]) +@torch.inference_mode() +def test_relu2_static_fp8_quant(o2_relu2_fp8_ops, shape) -> None: + torch.manual_seed(0) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.1495535671710968, device="cuda", dtype=torch.float32) + + relu2, quant_fp8 = o2_relu2_fp8_ops + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@pytest.mark.parametrize( + "enabled_ops", + [(), ("relu2",), ("quant_fp8",), ("relu2", "quant_fp8")], +) +@torch.inference_mode() +def test_relu2_static_fp8_quant_ignores_custom_op_dispatch(enabled_ops) -> None: + custom_ops = ["none", *(f"+{op}" for op in enabled_ops)] + config = VllmConfig( + optimization_level=OptimizationLevel.O2, + compilation_config=CompilationConfig(custom_ops=custom_ops), + ) + with set_current_vllm_config(config): + torch.manual_seed(0) + relu2 = ReLUSquaredActivation() + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_without_compilation() -> None: + config = VllmConfig(optimization_level=OptimizationLevel.O0) + with set_current_vllm_config(config): + torch.manual_seed(0) + relu2 = ReLUSquaredActivation() + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@pytest.mark.parametrize( + "scale_value", + [ + 0.00435965, + 0.007149832788854837, + 0.03125, + 0.0714285746216774, + 0.1495535671710968, + 0.1517857164144516, + 6.03571415, + ], +) +@torch.inference_mode() +def test_relu2_static_fp8_quant_all_bf16_values(o2_relu2_fp8_ops, scale_value) -> None: + bits = torch.arange(1 << 16, dtype=torch.int32).to(torch.uint16) + x = bits.view(torch.bfloat16).reshape(1, -1).to("cuda") + scale = torch.tensor(scale_value, device="cuda", dtype=torch.float32) + + relu2, quant_fp8 = o2_relu2_fp8_ops + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_empty(o2_relu2_fp8_ops) -> None: + x = torch.empty(0, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + relu2, _ = o2_relu2_fp8_ops + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + assert result.data.shape == x.shape + assert result.data.dtype == current_platform.fp8_dtype() + + +@requires_sm90 +@pytest.mark.parametrize("unsupported", ["dtype", "layout", "scale"]) +@torch.inference_mode() +def test_relu2_static_fp8_quant_falls_back(o2_relu2_fp8_ops, unsupported: str) -> None: + x = torch.randn((17, 32), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + if unsupported == "dtype": + x = x.to(torch.float16) + elif unsupported == "layout": + x = x.T + else: + scale = scale.to(torch.float16) + + relu2, _ = o2_relu2_fp8_ops + expected = relu2(x) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedActivation) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_torch_compile_dynamic( + o2_relu2_fp8_ops, +) -> None: + torch.manual_seed(1) + scale = torch.tensor(0.00435965, device="cuda", dtype=torch.float32) + relu2, quant_fp8 = o2_relu2_fp8_ops + linear = _StaticFp8Linear(scale) + + def fused(x: torch.Tensor) -> torch.Tensor: + result = maybe_fused_act_quant(relu2, x, linear) + assert isinstance(result, QuantizedActivation) + return result.data + + compile_count = 0 + + def counting_inductor_backend(gm, example_inputs): + nonlocal compile_count + compile_count += 1 + return torch._inductor.compile(gm, example_inputs) + + torch._dynamo.reset() + try: + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + torch._dynamo.mark_dynamic(x, 0) + compiled_fused = torch.compile( + fused, backend=counting_inductor_backend, fullgraph=True + ) + + for candidate in ( + x, + torch.randn((31, 5120), device="cuda", dtype=torch.bfloat16), + torch.randn((63, 5120), device="cuda", dtype=torch.bfloat16), + ): + expected, _ = quant_fp8(relu2(candidate), scale) + _assert_fp8_bitwise_equal(compiled_fused(candidate), expected) + + assert compile_count == 1 + finally: + torch._dynamo.reset() diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index bce91510d428..2dab86bc0d57 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -31,6 +31,9 @@ ) from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights from vllm.lora.punica_wrapper import get_punica_wrapper +from vllm.model_executor.layers.fusion.quant_activation import ( + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -56,6 +59,13 @@ torch.bfloat16: (3e-2, 2e-2), } + +def test_lora_linear_requires_unquantized_input() -> None: + layer = RowParallelLinearWithLoRA.__new__(RowParallelLinearWithLoRA) + layer._input_quant_key = object() + assert get_input_quant_key(layer) is None + + pytestmark = [ pytest.mark.skipif( not ( diff --git a/tests/model_executor/test_nemotron_h_quantization.py b/tests/model_executor/test_nemotron_h_quantization.py index 81a3fb5b13fc..e4e1c87db411 100644 --- a/tests/model_executor/test_nemotron_h_quantization.py +++ b/tests/model_executor/test_nemotron_h_quantization.py @@ -3,6 +3,8 @@ from unittest.mock import Mock, patch +import torch + def test_nemotron_h_lm_head_receives_quant_config(): from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM @@ -32,3 +34,28 @@ def test_nemotron_h_lm_head_receives_quant_config(): MockLMHead.assert_called_once() call_kwargs = MockLMHead.call_args.kwargs assert call_kwargs["quant_config"] is mock_quant_config + + +def test_relu2_fp8_fusion_uses_registry(): + from vllm.model_executor.models.nemotron_h import NemotronHMLP + + projected = torch.empty((1, 1), dtype=torch.bfloat16) + fused = Mock() + act_fn = Mock() + down_proj = Mock(side_effect=lambda x: (x, None)) + + mlp = NemotronHMLP.__new__(NemotronHMLP) + torch.nn.Module.__init__(mlp) + mlp.up_proj = Mock(return_value=(projected, None)) + mlp.down_proj = down_proj + mlp.act_fn = act_fn + + with patch( + "vllm.model_executor.models.nemotron_h.maybe_fused_act_quant", + return_value=fused, + ) as maybe_fused: + result = mlp(Mock()) + + maybe_fused.assert_called_once_with(act_fn, projected, down_proj) + act_fn.assert_not_called() + assert result is fused diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index 5c8e829b2997..df7cfe7122da 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -68,6 +68,10 @@ def lora_linear_async_fake( class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): + # The adapter branch consumes the original activation, independently of + # whether the wrapped base layer can consume a pre-quantized activation. + requires_unquantized_input = True + def __init__(self, base_layer: LinearBase): super().__init__() diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index 9fe253f2d9cd..e18a35a70da6 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -5,21 +5,27 @@ Given an activation module and the downstream linear it feeds, fuse the activation with that linear's input quantization into a single kernel when the -linear advertises a consumable input_quant_key (see quant_activation.py). +linear advertises a consumable input quantization key (see quant_activation.py). Falls back to the plain activation when nothing matches, so a model forward can always call maybe_fused_act_quant unconditionally. This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion -fires here the silu_and_mul pattern is already consumed, so the compiler pass -finds nothing to rewrite and the two never double-fuse. +fires here, the activation and quantization are already consumed, so a compiler +pass cannot fuse the same boundary again. """ from collections.abc import Callable import torch -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) +from vllm.model_executor.layers.fusion.relu2_fp8_quant import ( + relu_squared_static_fp8_quant, +) from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -139,12 +145,46 @@ def _silu_and_mul_nvfp4_dynamic( ) -# (activation module type, consumer input_quant_key) -> fused producer. -# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme. +# (activation module type, consumer input quantization key) -> fused producer. +# Add a row for each supported manual producer. _FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = { (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static, } + +def _relu_squared_static_fp8_quant_supported( + act_fn: torch.nn.Module, + x: torch.Tensor, + linear: LinearBase, +) -> bool: + """Return whether the ReLU2 static-FP8 producer can consume this input.""" + scale = getattr(linear, "input_scale", None) + return ( + isinstance(act_fn, ReLUSquaredActivation) + and x.is_cuda + and x.dtype == torch.bfloat16 + and x.is_contiguous() + and isinstance(scale, torch.Tensor) + and scale.dtype == torch.float32 + and scale.device == x.device + and scale.numel() == 1 + ) + + +# Optional per-entry predicates preserve the unconditional fallback contract +# for producers whose supported inputs are narrower than their registry key. +_FUSED_ACT_QUANT_SUPPORT: dict[ + tuple[type, QuantKey], Callable[[torch.nn.Module, torch.Tensor, LinearBase], bool] +] = {} + +# This path is validated on SM90 and newer CUDA devices. +if current_platform.is_cuda() and current_platform.has_device_capability(90): + _RELU2_STATIC_FP8_KEY = (ReLUSquaredActivation, kFp8StaticTensorSym) + _FUSED_ACT_QUANT[_RELU2_STATIC_FP8_KEY] = relu_squared_static_fp8_quant + _FUSED_ACT_QUANT_SUPPORT[_RELU2_STATIC_FP8_KEY] = ( + _relu_squared_static_fp8_quant_supported + ) + # Add CUDA-specific entries for dynamic block quantization if current_platform.is_cuda_alike(): _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128 @@ -161,12 +201,14 @@ def maybe_fused_act_quant( ) -> "torch.Tensor | QuantizedActivation": """Apply act_fn, fusing the downstream linear's input quant when possible. - Returns a QuantizedActivation when a fused kernel matches - (act_fn, linear.input_quant_key), else the plain activated tensor. + Returns a QuantizedActivation when a fused kernel matches the activation and + the consumer's effective input quantization key, else the plain activation. """ - key = getattr(linear, "input_quant_key", None) + key = get_input_quant_key(linear) if key is not None: - producer = _FUSED_ACT_QUANT.get((type(act_fn), key)) - if producer is not None: + registry_key = (type(act_fn), key) + producer = _FUSED_ACT_QUANT.get(registry_key) + support = _FUSED_ACT_QUANT_SUPPORT.get(registry_key) + if producer is not None and (support is None or support(act_fn, x, linear)): return producer(x, linear) return act_fn(x) diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index bd6cab668e8b..d69415448fa6 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -5,7 +5,9 @@ and consumed directly by a linear layer, letting the layer skip its own input quantization. A linear advertises the key its kernel can consume via expose_input_quant_key; the kernel validates and reads the activation via -as_quantized_activation. +as_quantized_activation. Producers query the effective capability through +get_input_quant_key, which hides the key when another consumer branch needs the +original activation. """ from dataclasses import dataclass, replace @@ -46,12 +48,11 @@ def weak_ref(self) -> "QuantizedActivation": def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: - """Advertise the kernel's pre-quantized input key on the layer, if any. + """Store the kernel's pre-quantized input key on the layer, if any. - This is the bridge from a kernel's input_quant_key() to the - layer.input_quant_key attribute that fusion call sites read. The attribute - is left unset when the kernel quantizes its own input, so non-supporting - backends never receive a QuantizedActivation. + This is the bridge from a kernel's input_quant_key() to the layer capability + that fusion call sites read through get_input_quant_key. The raw key is left + unset when the kernel quantizes its own input. TODO(mgoin): Producers also need the consumer's quantization scales (e.g. static input scale, global scale). Expose those here as well so producers @@ -59,7 +60,14 @@ def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: """ key = kernel.input_quant_key() if key is not None: - layer.input_quant_key = key + layer._input_quant_key = key + + +def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None: + """Return the pre-quantized input key the complete layer can consume.""" + if getattr(layer, "requires_unquantized_input", False): + return None + return getattr(layer, "_input_quant_key", None) def as_quantized_activation( diff --git a/vllm/model_executor/layers/fusion/relu2_fp8_quant.py b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py new file mode 100644 index 000000000000..4dc0fed4cc41 --- /dev/null +++ b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +_FP8_MIN, _FP8_MAX = get_fp8_min_max() + + +@triton.jit +def _relu_squared_static_fp8_quant_kernel( + x_ptr, + scale_ptr, + output_ptr, + num_elements, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + input_is_nan = x != x + relu = tl.maximum(x, 0.0) + # Preserve the standalone ReLU2 kernel's BF16 output boundary. + activated = (relu * relu).to(tl.bfloat16).to(tl.float32) + # Match O2's Inductor-compiled QuantFP8 reciprocal-then-multiply. + quantized = activated * (1.0 / tl.load(scale_ptr)) + quantized = tl.maximum(tl.minimum(quantized, FP8_MAX), FP8_MIN) + # Match the O2 native chain by preserving NaN through the FP8 conversion. + quantized = tl.where(input_is_nan, x, quantized) + tl.store(output_ptr + offsets, quantized, mask=mask) + + +def relu_squared_static_fp8_quant( + x: torch.Tensor, linear: LinearBase +) -> QuantizedActivation: + """BF16 ReLU2 followed by static per-tensor FP8 quantization.""" + assert x.dtype == torch.bfloat16 + assert x.is_contiguous() + scale = linear.input_scale + assert scale.dtype == torch.float32 + assert x.device == scale.device + assert scale.numel() == 1 + + output = torch.empty_like(x, dtype=current_platform.fp8_dtype()) + if x.numel() != 0: + block_size = min(triton.next_power_of_2(x.shape[-1]), 2048) + num_warps = min(max(block_size // 256, 1), 4) + grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) + _relu_squared_static_fp8_quant_kernel[grid]( + x, + scale, + output, + x.numel(), + FP8_MIN=_FP8_MIN, + FP8_MAX=_FP8_MAX, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + + return QuantizedActivation( + data=output, + scale=scale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kFp8StaticTensorSym, + ) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 212dd9537583..95dc852fb307 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -252,6 +252,8 @@ class LinearBase(PluggableLayer): disable_tp: If true, tensor parallelism will be disabled for this layer. """ + requires_unquantized_input = False + def __init__( self, input_size: int, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e65e51f3bb9d..c1bc4c6823f7 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -25,7 +25,11 @@ from torch import nn from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.config import ( + CacheConfig, + ModelConfig, + VllmConfig, +) from vllm.config.parallel import ParallelConfig from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size from vllm.distributed.communication_op import tensor_model_parallel_all_gather @@ -38,6 +42,7 @@ activation_without_mul, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -118,7 +123,7 @@ def __init__( def forward(self, x: torch.Tensor): x, _ = self.up_proj(x) - x = self.act_fn(x) + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) x, _ = self.down_proj(x) return x From e85c8826ce2a810367e8a70eedb987e31b140800 Mon Sep 17 00:00:00 2001 From: "Zhiyao(Ian) Yu" <105995086+Ianniu123@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:00:20 -0400 Subject: [PATCH 196/774] [Bugfix] Make KV cache and MFU log lines backend-neutral (#55650) Signed-off-by: Ianniu <985683179yzy@gmail.com> Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/core/kv_cache_utils.py | 4 +++- vllm/v1/metrics/loggers.py | 8 ++++++-- vllm/v1/metrics/perf.py | 2 +- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index b47df61add03..28478b69f5e6 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2391,9 +2391,11 @@ def update_kv_cache_capacity( vllm_config.cache_config.kv_cache_size_tokens = num_tokens vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency max_model_len = vllm_config.model_config.max_model_len + device_type = vllm_config.device_config.device_type logger.info_once( - "GPU KV cache size: %s tokens, " + "%s KV cache size: %s tokens, " "Maximum concurrency for %s tokens per request: %.2fx", + "GPU" if device_type == "cuda" else device_type.upper(), f"{num_tokens:,}", f"{max_model_len:,}", max_concurrency, diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 9e3f9c63a628..3a8adb6e3244 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -101,6 +101,8 @@ class LoggingStatLogger(StatLoggerBase): def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): self.engine_index = engine_index self.vllm_config = vllm_config + device_type = vllm_config.device_config.device_type + self.kv_cache_device = "GPU" if device_type == "cuda" else device_type.upper() self._reset(time.monotonic()) self.last_scheduler_stats = SchedulerStats() @@ -184,7 +186,7 @@ def _log_iteration_details( "%sIteration(%d): %d context requests, %d context tokens, " "%d generation requests, %d generation tokens, " "iteration elapsed time: %.2f ms%s, " - "GPU KV cache usage: %.1f%%%s", + "%s KV cache usage: %.1f%%%s", self._log_prefix_for_engine(engine_idx), details.iteration_index, details.num_ctx_requests, @@ -193,6 +195,7 @@ def _log_iteration_details( details.num_generation_tokens, details.elapsed_ms, " (dummy)" if details.is_dummy else "", + self.kv_cache_device, scheduler_stats.kv_cache_usage * 100, encoder_msg, ) @@ -287,12 +290,13 @@ def log(self): log_parts.extend( [ - "GPU KV cache usage: %.1f%%", + "%s KV cache usage: %.1f%%", "Prefix cache hit rate: %.1f%%", ] ) log_args.extend( [ + self.kv_cache_device, self.last_scheduler_stats.kv_cache_usage * 100, self.prefix_caching_metrics.hit_rate * 100, ] diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index ab90ffcbe142..ebb676e66dcd 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -1749,7 +1749,7 @@ def log(self, log_fn=logger.info, log_prefix: str = "") -> None: ) log_fn( - "%sMFU: %.1f TF/s/GPU %.1f GB/s/GPU", + "%sMFU: %.1f TF/s/device %.1f GB/s/device", log_prefix, avg_tflops_per_gpu, avg_gbps_per_gpu, From 56ca9904ce76975fe3aad28d1af7af961a07f054 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 11:37:08 -0700 Subject: [PATCH 197/774] [CI] Initialize ubatch runner in cudagraph unit test (#56808) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/cudagraph/test_breakable_cudagraph.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 0e10e4a51f1a..e792d68d89e0 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -43,6 +43,7 @@ def test_piecewise_capture_builds_fresh_metadata_for_both_passes(): manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]} manager._graphs_captured = False manager.use_breakable_cg = True + manager.ubatch_runner = None create_calls = [] forward_calls = [] From 5372e72a9884c141874465f0eb8054a4a81da6bb Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 11:57:37 -0700 Subject: [PATCH 198/774] [Perf][DSv4.1] Fold the mHC post block into the delayed pre projection (#56633) Signed-off-by: Yongye Zhu --- tests/distributed/test_engram_dp_shard.py | 7 +- tests/kernels/test_mhc_jit_warmup.py | 47 ++- tests/kernels/test_mhc_kernels.py | 248 ++++++++++++++- vllm/model_executor/kernels/mhc/tilelang.py | 285 +++++++++++++++++- .../kernels/mhc/tilelang_kernels.py | 140 ++++++++- vllm/model_executor/kernels/mhc/warmup.py | 22 +- vllm/models/deepseek_v41/nvidia/dspark.py | 2 +- vllm/models/deepseek_v41/nvidia/model.py | 178 +++++++---- 8 files changed, 837 insertions(+), 92 deletions(-) diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 35ffc659d442..c2fa0f857fc8 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -491,7 +491,9 @@ def _check_dummy_hash_model_forward( state.forward = lambda *args: ids.unsqueeze(1) class Decoder(SimpleNamespace): - def __call__(self, hidden, positions, input_ids, *args): + # Engram hashes and the mask stay the last positional arguments; + # **kwargs absorbs the decoder's keyword-only flags. + def __call__(self, hidden, positions, input_ids, *args, **kwargs): hashes, keep = args[-2:] if dp_shared_memory and dp_rank == 1: assert hashes is None and keep is None @@ -504,7 +506,8 @@ def __call__(self, hidden, positions, input_ids, *args): if dp_rank == 1: assert torch.all(hashes == engram_ops.DEAD_ID) output = engram.embed(hashes[:, 0]) - return output, None, None, None, None + # Trailing None is previous_aux; this stub captures no aux states. + return output, None, None, None, None, None model = SimpleNamespace( use_mega_moe=False, diff --git a/tests/kernels/test_mhc_jit_warmup.py b/tests/kernels/test_mhc_jit_warmup.py index 28e746c09295..512b3e6ba53a 100644 --- a/tests/kernels/test_mhc_jit_warmup.py +++ b/tests/kernels/test_mhc_jit_warmup.py @@ -17,6 +17,8 @@ MhcFusedTileLangKernel, MhcPostTileLangKernel, MhcPreBigFuseTileLangKernel, + mhc_fused_post_pre_split_config, + require_fused_post_pre_config, ) from vllm.model_executor.warmup import jit_warmup_tilelang_helper @@ -113,17 +115,32 @@ def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config( @pytest.mark.parametrize( - ("num_tokens", "hidden_size", "expected_n_splits", "expected_tile_n"), - [(4, 4096, 8, 2), (4, 8192, 4, 2), (8, 4096, 4, 3)], + ("num_tokens", "hidden_size", "expected_tile_n", "expected_n_splits"), + [ + (1, 4096, 2, 8), + (4, 8192, 2, 8), + (8, 5120, 2, 8), + (16, 5120, 6, 8), + (32, 7168, 6, 8), + ], ) -def test_mhc_fused_dispatch_matches_legacy_runtime_config( +def test_mhc_fused_dispatch_matches_the_launch_config( num_tokens: int, hidden_size: int, - expected_n_splits: int, expected_tile_n: int, + expected_n_splits: int, ) -> None: + """The compile key must be the config the launch will actually use. + + Both read mhc_fused_post_pre_split_config, so this pins the tuned bands + and guards against the two drifting apart again. + """ kernel = MhcFusedTileLangKernel() + tile_n, n_splits, n_thr = require_fused_post_pre_config( + num_tokens, hidden_size, hc_mult=4 + ) + assert (tile_n, n_splits) == (expected_tile_n, expected_n_splits) assert kernel.dispatch( num_tokens=num_tokens, hidden_size=hidden_size, @@ -131,11 +148,20 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( ) == kernel.CompileKey( hidden_size=hidden_size, hc_mult=4, - n_splits=expected_n_splits, - tile_n=expected_tile_n, + n_splits=n_splits, + tile_n=tile_n, + n_thr=n_thr, ) +def test_mhc_fused_declines_shapes_it_cannot_tile() -> None: + """Above the token cutoff, and for a hidden size the block cannot split.""" + assert mhc_fused_post_pre_split_config(33, 5120, 4) is None + assert mhc_fused_post_pre_split_config(1, 5137, 4) is None + with pytest.raises(ValueError, match="does not cover num_tokens"): + require_fused_post_pre_config(33, 5120, 4) + + @pytest.mark.parametrize( ("kernel", "compile_key"), [ @@ -235,8 +261,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( MhcFusedTileLangKernel.CompileKey( hidden_size=4096, hc_mult=4, - n_splits=4, - tile_n=3, + n_splits=8, + tile_n=6, + n_thr=128, ), ), ( @@ -246,6 +273,7 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( hc_mult=4, n_splits=8, tile_n=2, + n_thr=128, ), ), ( @@ -253,8 +281,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( MhcFusedTileLangKernel.CompileKey( hidden_size=8192, hc_mult=4, - n_splits=4, + n_splits=8, tile_n=2, + n_thr=128, ), ), ( diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 46f311dfe69b..504981538013 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -15,8 +15,12 @@ ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _HC_PRENORM_GEMM_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, +) +from vllm.model_executor.kernels.mhc.torch import ( + mhc_post_torch, + mhc_pre_delayed_torch, ) -from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton from vllm.model_executor.layers.mhc import ( HAS_AITER_MHC, @@ -115,9 +119,9 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): carried_mixes = [] def make_layer(mix): - def forward(hidden, positions, ids, pre, post, res, residual): + def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): carried_mixes.append(pre) - return hidden, None, None, None, mix + return hidden, None, None, None, mix, None return forward @@ -235,6 +239,171 @@ def test_deepseek_v41_mhc_pre_delayed( torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2) +def mhc_fused_post_pre_delayed_ref( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + *mix_args, + pre_mix: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Project the post result before it is rounded to BF16. + + The fused kernel keeps the updated residual streams in registers, so its + projection sees the FP32 post result while the streams it stores are the + BF16 rounding of it. + """ + post = torch.einsum("tij,tih->tjh", comb_res_mix, residual.float()) + post = post + post_layer_mix * x.float().unsqueeze(-2) + return mhc_pre_delayed_torch( + post.bfloat16(), *mix_args, pre_mix=pre_mix, x=post.flatten(1) + ) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("num_tokens", [0, 1, 8, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +@pytest.mark.parametrize("carried", [False, True]) +def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carried): + """Fold the post block into the delayed pre without changing its outputs.""" + set_random_seed(0) + hc_mult = 4 + mix_size = hc_mult * (hc_mult + 2) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + post_layer_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + comb_res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(mix_size, hc_mult * hidden_size, device=DEVICE) * 0.02 + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) + base = torch.randn(mix_size, device=DEVICE) + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) if carried else None + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE) + weight.uniform_(0.5, 1.5) + mix_args = (fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20) + + residual_cur, post, comb, layer_input, next_pre, aux = ( + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 + ) + ) + assert aux.shape == (0, hidden_size) + assert [ + residual_cur.shape, + post.shape, + comb.shape, + layer_input.shape, + next_pre.shape, + ] == [ + (num_tokens, hc_mult, hidden_size), + (num_tokens, hc_mult, 1), + (num_tokens, hc_mult, hc_mult), + (num_tokens, hidden_size), + (num_tokens, hc_mult), + ] + if num_tokens == 0: + return + + # The post mapping and the collapse are what the unfused kernels produce. + residual_ref = torch.ops.vllm.mhc_post_tilelang( + x, residual, post_layer_mix, comb_res_mix + ) + layer_input_ref = mhc_pre_delayed_tilelang( + residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 + )[2] + torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0) + torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0) + + post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix=pre_mix + ) + # Above the fused kernel's token count the projection falls back to a + # split-k TF32 GEMM over the stored BF16 streams. + tol = ( + 1e-5 + if mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) is not None + else 5e-3 + ) + for actual, expected in ( + (post, post_ref), + (comb, comb_ref), + (next_pre, next_pre_ref), + ): + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + + # Folding the stream mean into the collapse must not change it, or anything + # else the op returns. + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + *mix_args, + pre_mix, + weight, + 1e-6, + True, + ) + torch.testing.assert_close(captured[5], residual_cur.mean(dim=1), atol=0, rtol=0) + unchanged = (residual_cur, post, comb, layer_input, next_pre) + for with_aux, without in zip(captured[:5], unchanged, strict=True): + torch.testing.assert_close(with_aux, without, atol=0, rtol=0) + + # The unnormalized epilogue has its own fused aux path. + bare = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, None, 1e-6, True + ) + torch.testing.assert_close(bare[5], bare[0].mean(dim=1), atol=0, rtol=0) + + # Decode replays this op from a captured graph. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 + ) + graph.replay() + eager = (residual_cur, post, comb, layer_input, next_pre, aux) + for replayed, expected in zip(captured, eager, strict=True): + torch.testing.assert_close(replayed, expected, atol=0, rtol=0) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("carried", [False, True]) +@pytest.mark.parametrize("capture_aux", [False, True]) +def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_aux): + set_random_seed(0) + x = torch.randn(2, 5120, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + post_layer_mix = torch.rand(2, 4, 1, device=DEVICE) + comb_res_mix = torch.rand(2, 4, 4, device=DEVICE) + fn = torch.randn(24, 20480, device=DEVICE) * 0.02 + pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None + scale = torch.ones(3, device=DEVICE) + base = torch.zeros(24, device=DEVICE) + torch.library.opcheck( + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang.default, + ( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + scale, + base, + 1e-20, + 1e-6, + 1e-6, + 2.0, + 20, + pre_mix, + None, + 1e-6, + capture_aux, + ), + ) + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) def test_deepseek_v41_decoder_mixes_match_torch( @@ -295,15 +464,86 @@ def reference(*args, norm_weight, norm_eps, **kwargs): post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) return post, res, decoder.attn_norm(collapsed), pre + def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, **kw): + residual = mhc_post_torch(x, residual, post_mix, res_mix) + aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0) + return residual, *reference(residual, *args, **kw), aux + monkeypatch.setattr( "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", reference, ) + monkeypatch.setattr( + "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + fused_reference, + ) expected = decoder(x, positions, None, **kwargs) - for result, ref in zip(actual, expected, strict=True): + assert actual[-1] is None and expected[-1] is None + for result, ref in zip(actual[:-1], expected[:-1], strict=True): torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("entry", ["fused", "engram"]) +def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_config): + """Read the aux hidden state back out of the seam that already computed it. + + The draft model's input is the mean over hc streams of the previous + layer's post. That post now runs inside this layer's fused pre, so the + captured value must equal a standalone ``mhc_post_tilelang`` on the same + inputs -- and Engram must not be folded into it, since the injection + happens after the aux consumers read the stream. + """ + set_random_seed(0) + decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) + nn.Module.__init__(decoder) + decoder.hc_mult = 4 + decoder.hc_sinkhorn_iters = 20 + decoder.hc_eps = 1e-6 + decoder.rms_norm_eps = 1e-6 + decoder.hc_post_alpha = 2.0 + decoder.use_sequence_parallel = False + decoder.engram = None + from vllm.model_executor.layers.layernorm import RMSNorm + + decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to( + device=DEVICE, dtype=torch.bfloat16 + ) + decoder.attn = lambda positions, x, _: x * 0.5 + decoder.ffn = lambda x, input_ids: x * 0.25 + with torch.device(DEVICE): + decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) + decoder.hc_attn_base = torch.randn(24) + decoder.hc_ffn_base = torch.randn(24) + x = torch.randn(3, 5120, dtype=torch.bfloat16) + kwargs = dict( + pre_mix=torch.rand(3, 4), + residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.rand(3, 4, 1), + res_mix=torch.rand(3, 4, 4), + ) + if entry == "engram": + + class FakeEngram(nn.Module): + layer_hash_index = 0 + + def forward(self, residual, hashes, mask): + return residual + 0.125 + + decoder.engram = FakeEngram() + kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32) + positions = torch.arange(3) + + expected = torch.ops.vllm.mhc_post_tilelang( + x, kwargs["residual"], kwargs["post_mix"], kwargs["res_mix"] + ).mean(dim=1) + previous_aux = decoder(x, positions, None, **kwargs, capture_previous_aux=True)[-1] + torch.testing.assert_close(previous_aux, expected, atol=0, rtol=0) + assert decoder(x, positions, None, **kwargs)[-1] is None + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("carried", [False, True]) def test_mhc_pre_delayed_custom_op_supports_compile(carried): diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 97ca23b5418c..cea186e49e0d 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -194,6 +194,7 @@ def mhc_pre_delayed_tilelang( norm_weight, pre_mix if pre_mix is not None else post, next_pre_mix, + layer_input, hidden_size=hidden_size, rms_eps=rms_eps, hc_pre_eps=hc_pre_eps, @@ -218,6 +219,7 @@ def mhc_pre_delayed_tilelang( layer_input, pre_mix if pre_mix is not None else post, next_pre_mix, + layer_input, hidden_size, rms_eps, hc_pre_eps, @@ -263,6 +265,278 @@ def _mhc_pre_delayed_tilelang_fake( ) +def mhc_fused_post_pre_delayed_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Run one mHC post block followed by the next delayed mHC pre block. + + Within the fused kernel's token range the post mapping is folded into the + pre-norm GEMM, so the updated residual streams feed the projection from + registers instead of a second pass over global memory. Above it this runs + the same post kernel and split-k GEMM as the unfused pair. + + Args: + x: BF16 sublayer output of shape (tokens, hidden_size). + residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size). + post_layer_mix: FP32 post coefficients of shape (tokens, hc_mult, 1). + comb_res_mix: FP32 residual coefficients, (tokens, hc_mult, hc_mult). + fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size). + hc_scale: FP32 scales of shape (3,). + hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),). + rms_eps: RMS normalization epsilon. + hc_pre_eps: Pre-mix epsilon. + hc_sinkhorn_eps: Sinkhorn epsilon. + hc_post_mult_value: Post-mix multiplier. + sinkhorn_repeat: Number of Sinkhorn iterations. + pre_mix: FP32 coefficients from the previous sublayer, or None to + select residual stream zero. + norm_weight: Optional BF16 RMSNorm weight for the collapsed input. + norm_eps: RMSNorm epsilon for the collapsed input. + capture_aux: Also return the mean over the post-mapped streams, which + draft models consume as the target's hidden state. It is folded + into the collapse, which already reads those streams. + + Returns: + The post-mapped residual streams, the post and residual coefficients, + the optionally normalized BF16 layer input, the next FP32 pre-mix, and + the BF16 stream mean (empty unless capture_aux), with shapes + (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1), + (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult), + and (tokens, hidden_size). + """ + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, + _MHC_FUSED_TILELANG_KERNEL, + _MHC_POST_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, + mhc_pre_big_fuse_tilelang, + ) + from vllm.model_executor.kernels.mhc.warmup import ( + MHC_PRE_NORM_KERNEL, + compute_mhc_pre_num_splits, + ) + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + tf32_hc_prenorm_gemm, + ) + + assert residual.ndim == 3 and residual.dtype == torch.bfloat16 + assert residual.is_contiguous() + num_tokens, hc_mult, hidden_size = residual.shape + input_size = hc_mult * hidden_size + mix_size = hc_mult * (hc_mult + 2) + assert x.shape == (num_tokens, hidden_size) and x.dtype == torch.bfloat16 + assert x.is_contiguous() + assert post_layer_mix.shape[:2] == (num_tokens, hc_mult) + assert post_layer_mix.dtype == torch.float32 and post_layer_mix.is_contiguous() + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult) + assert comb_res_mix.dtype == torch.float32 and comb_res_mix.is_contiguous() + assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32 + assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32 + assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32 + if pre_mix is not None: + assert pre_mix.shape == (num_tokens, hc_mult) + assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous() + + next_pre_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + post = torch.empty_like(next_pre_mix) + comb = torch.empty( + num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device + ) + layer_input = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ) + aux = torch.empty( + num_tokens if capture_aux else 0, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ) + if num_tokens == 0: + return ( + torch.empty_like(residual), + post.unsqueeze(-1), + comb.view(num_tokens, hc_mult, hc_mult), + layer_input, + next_pre_mix, + aux, + ) + + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + if fused_config is not None: + mixes, sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + fn.view(mix_size, hc_mult, hidden_size), + hc_mult, + hidden_size, + mix_size, + ) + else: + residual_cur = torch.empty_like(residual) + _MHC_POST_TILELANG_KERNEL( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + residual_cur, + hc_mult, + hidden_size, + ) + # The delayed epilogue is compiled per bucketed split count, so the + # projection has to use the same bucket rather than the raw estimate. + use_deep_gemm = is_deep_gemm_supported() + n_splits = ( + compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1 + ) + mixes = torch.empty( + n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device + ) + sqrsum = torch.empty( + n_splits, num_tokens, dtype=torch.float32, device=residual.device + ) + residual_cur_2d = residual_cur.view(num_tokens, input_size) + if use_deep_gemm: + tf32_hc_prenorm_gemm(residual_cur_2d, fn, mixes, sqrsum, n_splits) + else: + _HC_PRENORM_GEMM_TILELANG_KERNEL( + residual_cur_2d, + fn, + mixes, + sqrsum, + input_size, + 1, + ) + + outputs = ( + residual_cur, + post.unsqueeze(-1), + comb.view(num_tokens, hc_mult, hc_mult), + layer_input, + next_pre_mix, + aux, + ) + if norm_weight is not None: + assert norm_weight.shape == (hidden_size,) + assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous() + MHC_PRE_NORM_KERNEL( + mixes, + sqrsum, + hc_scale, + hc_base, + residual_cur, + post, + comb, + layer_input, + norm_weight, + pre_mix if pre_mix is not None else post, + next_pre_mix, + aux if capture_aux else layer_input, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=norm_eps, + hc_mult=hc_mult, + use_pre_mix_in=pre_mix is not None, + save_pre_mix=True, + rms_numel=input_size, + write_aux=capture_aux, + ) + return outputs + mhc_pre_big_fuse_tilelang( + mixes, + sqrsum, + hc_scale, + hc_base, + residual_cur, + post, + comb, + layer_input, + pre_mix if pre_mix is not None else post, + next_pre_mix, + aux if capture_aux else layer_input, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + mixes.shape[0], + hc_mult, + use_pre_mix_in=pre_mix is not None, + save_pre_mix=True, + rms_numel=input_size, + write_aux=capture_aux, + ) + return outputs + + +def _mhc_fused_post_pre_delayed_tilelang_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + num_tokens, hc_mult, hidden_size = residual.shape + return ( + torch.empty_like(residual), + torch.empty( + num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device + ), + torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device + ), + torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ), + torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device), + torch.empty( + num_tokens if capture_aux else 0, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ), + ) + + def mhc_pre_tilelang( residual: torch.Tensor, fn: torch.Tensor, @@ -574,6 +848,7 @@ def mhc_fused_post_pre_tilelang( _MHC_FUSED_TILELANG_KERNEL, _MHC_POST_TILELANG_KERNEL, _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, ) assert residual.dtype == torch.bfloat16 @@ -614,7 +889,7 @@ def mhc_fused_post_pre_tilelang( post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult) comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) - use_small_fma = num_tokens <= 16 + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) post_mix_cur = torch.empty( num_tokens, @@ -635,7 +910,7 @@ def mhc_fused_post_pre_tilelang( device=residual.device, ) - if use_small_fma: + if fused_config is not None: gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL( comb_res_mix_flat, residual_flat, @@ -815,6 +1090,12 @@ def _hc_head_fused_kernel_tilelang_fake( mutates_args=[], fake_impl=_mhc_pre_delayed_tilelang_fake, ) +direct_register_custom_op( + op_name="mhc_fused_post_pre_delayed_tilelang", + op_func=mhc_fused_post_pre_delayed_tilelang, + mutates_args=[], + fake_impl=_mhc_fused_post_pre_delayed_tilelang_fake, +) direct_register_custom_op( op_name="mhc_pre_tilelang", op_func=mhc_pre_tilelang, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index b4d3190c0eff..9e803d123e29 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -37,8 +37,67 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k -def _mhc_fused_n_splits(num_tokens: int, hidden_size: int) -> int: - return 8 if num_tokens < 8 and hidden_size <= 4096 else 4 +# ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so +# every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300 +# (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k +# GEMM it replaces; the bands are coarse because the loss against a per-shape +# optimum is ~1-2%, while one fixed config gives up to 18% at the top of the +# range. +_FUSED_POST_PRE_N_THR = 128 +_FUSED_POST_PRE_MAX_TOKENS = 32 + + +def mhc_fused_post_pre_split_config( + num_tokens: int, hidden_size: int, hc_mult: int +) -> tuple[int, int, int] | None: + """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM. + + Returns None when a separate post kernel followed by a split-k GEMM is the + better choice, or when the hidden size does not divide evenly across the + fused kernel's block. One source of truth for the dispatch decision, the + compile key and the launch, which must agree. + """ + if num_tokens > _FUSED_POST_PRE_MAX_TOKENS: + return None + n_thr = _FUSED_POST_PRE_N_THR + if hidden_size % n_thr: + return None + # The projection tiles supply the parallelism a small token count cannot, + # so the tile grows with the batch: 1.3x at one token and 1.03x at 32 + # against a split-k GEMM given the split its own estimator computes. + if num_tokens < 16: + tile_n, n_splits = 2, 8 + else: + tile_n, n_splits = 6, 8 + # The kernel drops projection tiles and k-slices it cannot fill evenly. + mix_size = hc_mult * (hc_mult + 2) + while tile_n > 1 and mix_size % tile_n: + tile_n -= 1 + while n_splits > 1 and hidden_size % (n_splits * n_thr): + n_splits //= 2 + return tile_n, n_splits, n_thr + + +def require_fused_post_pre_config( + num_tokens: int, hidden_size: int, hc_mult: int +) -> tuple[int, int, int]: + """The config for a shape the caller has already committed to fusing.""" + config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + if config is None: + raise ValueError( + "the fused mHC post + pre-norm GEMM does not cover num_tokens=" + f"{num_tokens} at hidden_size={hidden_size}" + ) + return config + + +def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]: + """Every split-k factor the fused post + pre-norm GEMM path can pick.""" + configs = ( + mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1) + ) + return tuple(sorted({config[1] for config in configs if config is not None})) @tilelang_jit @@ -53,6 +112,7 @@ def mhc_pre_big_fuse_tilelang( layer_input, pre_mix_in, pre_mix_out, + aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, @@ -64,6 +124,7 @@ def mhc_pre_big_fuse_tilelang( use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, + write_aux: bool = False, ): """Fuse coefficient generation and residual collapse after the projection. @@ -88,6 +149,7 @@ def mhc_pre_big_fuse_tilelang( pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: if ENABLE_PDL: @@ -188,11 +250,24 @@ def mhc_pre_big_fuse_tilelang( ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) + if write_aux: + # Aux consumers (draft models) take the plain stream mean + # of the same residual this collapse reads. + aux = T.alloc_fragment(hidden_block, T.float32) + T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] + if write_aux: + aux[i1_h] += xl[i_hc, i1_h] + + if write_aux: + for i1_h in T.Parallel(hidden_block): + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( + aux[i1_h] / hc_mult + ) T.copy(ol, layer_input[i, i0_h * hidden_block]) @@ -216,6 +291,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( norm_weight, pre_mix_in, pre_mix_out, + aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, @@ -229,6 +305,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, + write_aux: bool = False, ): num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) @@ -250,6 +327,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: rms = T.alloc_fragment(1, T.float32) @@ -345,11 +423,24 @@ def mhc_pre_big_fuse_with_norm_tilelang( ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) + if write_aux: + # Aux consumers (draft models) take the plain stream mean + # of the same residual this collapse reads. + aux = T.alloc_fragment(hidden_block, T.float32) + T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] + if write_aux: + aux[i1_h] += xl[i_hc, i1_h] + + if write_aux: + for i1_h in T.Parallel(hidden_block): + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( + aux[i1_h] / hc_mult + ) if save_pre_mix: # Keep the BF16 boundary before the delayed input RMSNorm. @@ -1211,9 +1302,17 @@ def _kernel_args( middle_args = (residual_out, *output_args, norm_weight) elif compile_key.use_norm_weight: assert norm_weight is not None - middle_args = (*output_args, norm_weight, post_mix, post_mix) + # layer_input stands in for the aux buffer on both epilogues: same + # shape and dtype, and write_aux is off here so it is never written. + middle_args = ( + *output_args, + norm_weight, + post_mix, + post_mix, + layer_input, + ) else: - middle_args = (*output_args, post_mix, post_mix) + middle_args = (*output_args, post_mix, post_mix, layer_input) norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else () return ( *common_args, @@ -1247,8 +1346,14 @@ def dispatch( # type: ignore[override] if use_pre_gemm_splits else n_splits ) - fused_n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) - actual_n_splits = fused_n_splits if use_fused_tilelang else pre_gemm_n_splits + # The epilogue reduces over whatever the fused kernel split the GEMM + # into, so both have to read the same config. + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + actual_n_splits = ( + fused_config[1] + if use_fused_tilelang and fused_config is not None + else pre_gemm_n_splits + ) actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps actual_use_norm_weight = use_norm_weight or is_broadcast return self.CompileKey( @@ -1540,6 +1645,7 @@ class CompileKey: hc_mult: int n_splits: int tile_n: int + n_thr: int @staticmethod def kernel() -> Any: @@ -1552,14 +1658,15 @@ def dispatch( # type: ignore[override] hidden_size: int, hc_mult: int, ) -> CompileKey: - # TODO(gnovack): investigate autotuning these heuristics - tile_n = 2 if num_tokens < 8 else 3 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + # Subscripts rather than an unpack: the warmup tracer parses this body + # and allows only plain assignments before the return. + config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult) return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, - n_splits=n_splits, - tile_n=tile_n, + n_splits=config[1], + tile_n=config[0], + n_thr=config[2], ) def get_warmup_keys( @@ -1574,11 +1681,11 @@ def get_warmup_keys( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, - _when=lambda *, num_tokens: num_tokens <= 16, + _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: - num_tokens = 1 if compile_key.tile_n == 2 else 8 + num_tokens = 1 if compile_key.tile_n == 2 else 16 hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult hc_mult3 = hc_mult * (2 + hc_mult) @@ -1617,8 +1724,9 @@ def __call__( hc_mult3: int, ) -> TileLangLaunchSpec: num_tokens = residual_in.shape[0] - tile_n = 2 if num_tokens < 8 else 3 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + tile_n, n_splits, n_thr = require_fused_post_pre_config( + num_tokens, hidden_size, hc_mult + ) yp_out = residual_in.new_empty( (n_splits, num_tokens, hc_mult3), dtype=torch.float32 ) @@ -1639,7 +1747,7 @@ def __call__( hidden_size, hc_mult3, ), - dict(tile_n=tile_n, split_k=n_splits), + dict(n_thr=n_thr, tile_n=tile_n), (yp_out, rp_out, residual_out), ) diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 0abe1a2161c2..4e6a19781c82 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -36,23 +36,33 @@ class CompileKey: use_pre_mix_in: bool rms_numel: int save_pre_mix: bool = True + write_aux: bool = False kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang) def dispatch(self, *, n_splits, **fields) -> CompileKey: # type: ignore[override] return self.CompileKey(n_splits=n_splits, **fields) - def get_warmup_keys(self, *, max_tokens: int, **fields) -> list[CompileKey]: - # The split heuristic changes only at 64-token boundaries. - splits = ( - sorted( + def get_warmup_keys( + self, + *, + max_tokens: int, + extra_splits: Iterable[int] = (), + **fields, + ) -> list[CompileKey]: + # The split heuristic changes only at 64-token boundaries. Callers that + # reach this kernel through another GEMM (the fused post + pre-norm + # GEMM) pass the splits that GEMM picks in extra_splits. + splits = sorted( + set(extra_splits) + | ( { compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens) for num_tokens in range(1, max_tokens + 1, 64) } + if is_deep_gemm_supported() + else {1} ) - if is_deep_gemm_supported() - else [1] ) return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 46fe2558ab7f..3c711681aa75 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -207,7 +207,7 @@ def forward( residual = post_mix = res_mix = pre_mix = None for layer in self.layers: - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + hidden_states, residual, post_mix, res_mix, pre_mix, _ = layer( hidden_states, positions, input_ids, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index df12a9d9bbc9..f3530544da7b 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,6 +20,7 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -214,17 +215,45 @@ def __init__( self.hc_eps = config.hc_eps self.hc_post_alpha = 2.0 if vllm_config.kernel_config.enable_jit_warmup and current_platform.is_cuda(): + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_fused_post_pre_splits, + ) from vllm.model_executor.kernels.mhc.warmup import MHC_PRE_NORM_KERNEL max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if self.use_sequence_parallel: tp_size = vllm_config.parallel_config.tensor_parallel_size max_tokens = (max_tokens + tp_size - 1) // tp_size - for input_size, use_pre_mix in ( - (self.hidden_size, False), - (self.hc_mult * self.hidden_size, False), - (self.hc_mult * self.hidden_size, True), - ): + # The epilogue compiles per projection width and per pre-mix mode. + # The first layer projects the broadcast embedding, so it reads one + # hidden_size-wide row and selects stream zero; every later sublayer + # projects the full hc stream, and collapses it with the pre-mix the + # previous sublayer carried in. Those later shapes also arrive + # through the fused post + pre-norm GEMM, which picks split-k + # factors the token sweep below never produces. + broadcast_embedding = { + "rms_numel": self.hidden_size, + "use_pre_mix_in": False, + "extra_splits": (), + } + hc_stream = { + "rms_numel": self.hc_mult * self.hidden_size, + "extra_splits": mhc_fused_post_pre_splits( + self.hidden_size, self.hc_mult + ), + } + variants = [ + broadcast_embedding, + {**hc_stream, "use_pre_mix_in": False}, + {**hc_stream, "use_pre_mix_in": True}, + ] + if vllm_config.speculative_config is not None: + # Draft setups read aux hidden states out of the same collapse, + # which is its own specialization of the kernel. + variants.append( + {**hc_stream, "use_pre_mix_in": True, "write_aux": True} + ) + for variant in variants: MHC_PRE_NORM_KERNEL.register_warmup( max_tokens=max_tokens, hidden_size=self.hidden_size, @@ -235,8 +264,7 @@ def __init__( sinkhorn_repeat=self.hc_sinkhorn_iters, norm_eps=self.rms_norm_eps, hc_mult=self.hc_mult, - use_pre_mix_in=use_pre_mix, - rms_numel=input_size, + **variant, ) mix_hc = (2 + self.hc_mult) * self.hc_mult hc_dim = self.hc_mult * self.hidden_size @@ -295,7 +323,17 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + *, + capture_previous_aux: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + ]: + previous_aux: torch.Tensor | None = None # The reference collapses each sublayer's input with the *previous* # sublayer's pre-mix: attention uses the pre-mix carried in (identity # for the first layer), the FFN uses this layer's attention pre-mix. @@ -335,17 +373,19 @@ def forward( norm_weight=self.attn_norm.weight, norm_eps=self.attn_norm.variance_epsilon, ) - else: - residual = mhc_post_tilelang(x, residual, post_mix, res_mix) - if self.engram is not None and engram_hashes is not None: - # Engram injection happens between the previous sublayer's - # post and this block's pre, on the full hc stream, so the - # mix coefficients see the injected stream. - residual = self.engram( - residual, - engram_hashes[:, self.engram.layer_hash_index], - engram_mask, - ) + elif self.engram is not None and engram_hashes is not None: + # Engram injection happens between the previous sublayer's post + # and this block's pre, on the full hc stream, so the mix + # coefficients see the injected stream. The injection also keeps + # the post out of the pre-norm GEMM's fused prologue. + previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix) + if capture_previous_aux: + previous_aux = previous_post.mean(dim=1) + residual = self.engram( + previous_post, + engram_hashes[:, self.engram.layer_hash_index], + engram_mask, + ) post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang( residual, self.hc_attn_fn, @@ -360,6 +400,31 @@ def forward( norm_weight=self.attn_norm.weight, norm_eps=self.attn_norm.variance_epsilon, ) + else: + # The collapse already reads the post-mapped streams, so the mean + # aux consumers want comes out of the same kernel. + residual, post_mix, res_mix, x, attn_pre, aux = ( + mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, + ) + ) + if capture_previous_aux: + previous_aux = aux if self.use_sequence_parallel: x = sp_all_gather(x)[: positions.shape[0]] @@ -368,23 +433,27 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual = mhc_post_tilelang(x, residual, post_mix, res_mix) - post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_tilelang( - residual, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, + residual, post_mix, res_mix, x, ffn_pre, _ = ( + mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, + ) ) x = self.ffn(x, input_ids) - return x, residual, post_mix, res_mix, ffn_pre + return x, residual, post_mix, res_mix, ffn_pre, previous_aux class DeepseekV4Model(nn.Module, EagleModelMixin): @@ -631,13 +700,14 @@ def forward( if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call + # Every layer's post runs inside the next layer's fused pre, so aux + # hidden states are read back from there instead of recomputed. + aux_hidden_by_layer: dict[int, torch.Tensor] = {} for idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer( hidden_states, positions, input_ids, @@ -647,25 +717,29 @@ def forward( residual, engram_hashes, engram_mask, + capture_previous_aux=idx in self.aux_hidden_state_layers, ) - if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) - aux_hidden_state = aux_recon.mean(dim=1) + if previous_aux is not None: + # idx is the one-based id of the layer whose post this is. if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) - final_aux_recon = aux_recon + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] + aux_hidden_by_layer[idx] = previous_aux if layer is not None: - # Reuse if the last layer was captured as an aux hidden state + # The last layer has no successor to fold its post into. + hidden_states = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) if self.end_layer in self.aux_hidden_state_layers: - hidden_states = final_aux_recon - else: - hidden_states = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) + final_aux = hidden_states.mean(dim=1) + if self.use_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_by_layer[self.end_layer] = final_aux + + aux_hidden_states = [ + aux_hidden_by_layer[layer_id] + for layer_id in self.aux_hidden_state_layers + if layer_id in aux_hidden_by_layer + ] if not get_pp_group().is_last_rank: return IntermediateTensors( From 0a5747d410da9b08145dba2b6f3e6614bab35753 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=90=98=E5=A4=A9=E6=A5=BD?= Date: Mon, 14 Sep 2026 15:02:07 -0400 Subject: [PATCH 199/774] [Profiler] Add Proton CUDA graph attribution for MRV2 (#51084) Signed-off-by: Luossu Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- docs/contributing/profiling.md | 46 ++- tests/v1/worker/test_gpu_profiler.py | 398 +++++++++++++++++++++- vllm/config/profiler.py | 16 + vllm/config/vllm.py | 51 ++- vllm/profiler/wrapper.py | 149 +++++++- vllm/v1/worker/gpu/model_runner.py | 10 + vllm/v1/worker/gpu_worker.py | 29 +- vllm/v1/worker/mm_encoder_model_runner.py | 3 + 8 files changed, 657 insertions(+), 45 deletions(-) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index e420a3aa68e1..880d065a2c5e 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -90,18 +90,18 @@ $ curl -X POST http://localhost:8000/stop_profile is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a Chrome trace and works through the same vLLM profiling controls as the PyTorch and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and -requires eager execution. +supports CUDA graph attribution. -Start a server with a local output directory: +Start a server with a local output directory and graph attribution: ```bash vllm serve meta-llama/Llama-3.1-8B-Instruct \ - --enforce-eager \ --profiler-config '{ "profiler": "proton", "proton_profiler_dir": "./proton_profile", "proton_output_format": "hatchet", - "proton_hook": "triton" + "proton_hook": "triton", + "proton_graph_attribution": true }' ``` @@ -111,7 +111,7 @@ rank-qualified output name, such as `proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`, so distributed workers, restarted servers, and repeated profiling runs do not overwrite one another. A `profile_prefix` is included when supplied. Each -profile is finalized by `/stop_profile` and is ready to inspect immediately. +profile is written by `/stop_profile` and is ready to inspect immediately. The Proton-specific options are: @@ -121,21 +121,39 @@ The Proton-specific options are: - `proton_mode`: an optional backend mode string - `proton_hook`: `triton` to record Triton launch metadata, or unset - `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset +- `proton_graph_attribution`: observe CUDA graph capture for replay attribution; + disabled by default and requires `proton_data: "tree"` `hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while -`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet -supported, so `--enforce-eager` is required for every Proton mode. +`chrome_trace` requires `proton_data: "trace"`. Automatic backend selection is recommended. vLLM currently supports Proton's `cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not expose Proton's experimental instrumentation backend because current upstream -Triton builds can produce profiles without timing metrics. Backend-specific -modes, including `pcsampling`, can be selected with `proton_mode`. - -Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The -`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or -newer; vLLM rejects these options with the detected version before starting a -session. +can produce profiles without timing metrics. When `proton_graph_attribution` is +enabled, Proton observes vLLM's CUDA graph capture with the configured profiling +session active, then deactivates that same session until profiling starts. This +lets later profiles attribute replayed kernels without retaining model-startup +activity. Backend-specific modes can be selected with `proton_mode`; +`pcsampling` synchronizes the CUDA context and therefore requires +`--enforce-eager`. When CUDA graphs are enabled (including encoder graphs), +Proton requires `proton_graph_attribution: true` to collect replayed kernels. +For Chrome traces, disable CUDA graphs with `--enforce-eager`. + +CUDA graph-attributed profiles support repeated `start_profile`/`stop_profile` +runs. Each stop flushes and writes one tree-data phase while preserving the +graph-aware session. Without `proton_graph_attribution`, each `stop_profile` +instead finalizes and writes an independent Proton session. + +CUDA graph attribution requires Triton 3.7 or newer. It uses the phase data API +to discard graph-capture activity and separate profiling runs. The `hatchet_msgpack` +output format and `periodic_flushing` mode also require Triton 3.7 or newer. +`periodic_flushing` cannot be combined with graph attribution because both +manage the session's data phases. Ordinary Proton profiling remains available +with Triton 3.6. + +Graph attribution retains capture metadata for the worker lifetime. With eager +execution or no graphs to capture, profiling uses ordinary independent sessions. Inspect tree profiles with: diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index 347b5b379adf..f6263b68da0a 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -350,17 +350,24 @@ def test_profiler_entered_during_capture(): def make_proton(session_id: int | None = 7): + data = SimpleNamespace( + advance_phase=Mock(side_effect=range(1, 100)), + clear=Mock(), + get=Mock(return_value={"traceEvents": []}), + get_msgpack=Mock(return_value=b"profile"), + ) return SimpleNamespace( start=Mock(return_value=session_id), activate=Mock(), deactivate=Mock(), finalize=Mock(), scope=Mock(return_value=nullcontext()), + data=data, ) def make_proton_wrapper( - tmp_path, proton=None, triton_version="3.6.0", **config_overrides + tmp_path, proton=None, triton_version="3.7.0", **config_overrides ): proton = proton or make_proton() config = ProfilerConfig( @@ -415,6 +422,13 @@ def test_normalizes_local_output_directory(self, tmp_path, monkeypatch): }, "requires proton_data", ), + ( + { + "proton_data": "trace", + "proton_graph_attribution": True, + }, + "requires proton_data='tree'", + ), ], ) def test_rejects_invalid_option_combinations(self, tmp_path, options, message): @@ -450,6 +464,7 @@ def test_rejects_invalid_typed_options(self, field, tmp_path): ("proton_mode", "pcsampling"), ("proton_hook", "triton"), ("proton_output_format", "chrome_trace"), + ("proton_graph_attribution", True), ], ) def test_rejects_proton_options_for_other_profilers(self, field, value): @@ -478,17 +493,103 @@ def test_rejects_proton_on_non_cuda_platforms(self, tmp_path): compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), ) - def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path): - with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"): + @pytest.mark.parametrize("encoder_only", [False, True]) + @pytest.mark.parametrize("attribution", [False, True]) + def test_cuda_graphs_require_attribution(self, tmp_path, encoder_only, attribution): + # Encoder graphs are independent of the decoder's cudagraph_mode. + expected = ( + nullcontext() + if attribution + else pytest.raises( + ValueError, match="requires proton_graph_attribution=True" + ) + ) + with expected: + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=attribution, + ), + compilation_config=CompilationConfig( + cudagraph_mode=( + CUDAGraphMode.NONE if encoder_only else CUDAGraphMode.FULL + ), + cudagraph_mm_encoder=encoder_only, + ), + ) + + def test_validates_default_cuda_graph_mode_after_resolution(self, tmp_path): + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): VllmConfig( profiler_config=ProfilerConfig( profiler="proton", proton_profiler_dir=str(tmp_path) ), + ) + + @pytest.mark.parametrize( + "mode", + [ + "periodic_flushing", + "periodic_flushing:format=hatchet", + "PERIODIC_FLUSHING:format=hatchet", + ], + ) + def test_graph_attribution_rejects_periodic_flushing(self, tmp_path, mode): + # Reject before Proton's native phase manager can abort the worker. + with pytest.raises(ValueError, match="incompatible with periodic_flushing"): + ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + proton_mode=mode, + ) + + @pytest.mark.parametrize("attribution", [False, True]) + @pytest.mark.parametrize( + "mode", ["pcsampling", "pcsampling:interval=100", "PcSampling:interval=100"] + ) + @pytest.mark.parametrize("encoder_only", [False, True]) + def test_pcsampling_requires_graphs_disabled( + self, tmp_path, attribution, mode, encoder_only + ): + with pytest.raises(ValueError, match="PC sampling requires CUDA graphs"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=attribution, + proton_mode=mode, + ), compilation_config=CompilationConfig( - cudagraph_mode=CUDAGraphMode.PIECEWISE + cudagraph_mode=CUDAGraphMode.NONE + if encoder_only + else CUDAGraphMode.FULL, + cudagraph_mm_encoder=encoder_only, ), ) + def test_ordinary_proton_keeps_mrv1_graph_restriction(self, tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", proton_profiler_dir=str(tmp_path) + ), + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL), + ) + + def test_graph_attribution_rejects_mrv1(self, tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + with pytest.raises(ValueError, match="requires the V2 model runner"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + ) + ) + @_requires_cuda_for_proton class TestProtonProfilerWrapper: @@ -564,7 +665,20 @@ def test_output_names_are_unique_across_worker_restarts(self, tmp_path): ) def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature): with pytest.raises(RuntimeError, match=feature): - make_proton_wrapper(tmp_path, **{option: value}) + make_proton_wrapper(tmp_path, triton_version="3.6.0", **{option: value}) + + @pytest.mark.parametrize("version", ["3.6.0", "unknown"]) + def test_graph_attribution_requires_phase_api(self, tmp_path, version): + with pytest.raises(RuntimeError, match="requires Triton >= 3.7"): + make_proton_wrapper( + tmp_path, triton_version=version, proton_graph_attribution=True + ) + + def test_ordinary_profiling_still_supports_triton_3_6(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, triton_version="3.6.0") + wrapper.start() + wrapper.stop() + proton.finalize.assert_called_once_with(session=7) @pytest.mark.parametrize( ("option", "value"), @@ -583,6 +697,125 @@ def test_rejects_output_format_when_finalize_lacks_capability(self, tmp_path): with pytest.raises(RuntimeError, match="does not support selecting"): make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet") + def test_cuda_graph_tree_phase_is_written_at_stop(self, tmp_path): + wrapper, proton = make_proton_wrapper( + tmp_path, + proton_context="python", + proton_graph_attribution=True, + ) + with wrapper.capture_cuda_graphs(): + proton.start.assert_called_once() + proton.deactivate.assert_not_called() + + capture_args = proton.start.call_args.kwargs + assert capture_args["context"] == "python" + assert capture_args["data"] == "tree" + proton.data.advance_phase.assert_called_once_with(7) + proton.deactivate.assert_called_once_with(session=7, flushing=True) + proton.data.clear.assert_called_once_with(7, 0) + + wrapper.start() + wrapper.stop() + wrapper.start() + wrapper.stop() + + assert proton.data.get.call_args_list == [call(7, 1), call(7, 2)] + proton.data.clear.assert_has_calls([call(7, 0), call(7, 1), call(7, 2)]) + output_names = sorted(tmp_path.glob("proton_rank_3_*.hatchet")) + assert len(output_names) == 2 + assert output_names[0].name.endswith("_run0.hatchet") + assert output_names[1].name.endswith("_run1.hatchet") + + def test_capture_is_noop_without_opt_in(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path) + with wrapper.capture_cuda_graphs(): + pass + proton.start.assert_not_called() + + @pytest.mark.parametrize("delay", [0, 2]) + def test_duplicate_start_preserves_output_prefix(self, tmp_path, delay): + wrapper, proton = make_proton_wrapper( + tmp_path, proton_graph_attribution=True, delay_iterations=delay + ) + with wrapper.capture_cuda_graphs(): + pass + wrapper.set_output_name("first") + wrapper.start() + wrapper.set_output_name("duplicate") + wrapper.start() + for _ in range(delay): + wrapper.step() + wrapper.stop() + assert len(list(tmp_path.glob("proton_first_*.hatchet"))) == 1 + assert not list(tmp_path.glob("proton_duplicate_*")) + wrapper.set_output_name("second") + wrapper.start() + for _ in range(delay): + wrapper.step() + wrapper.stop() + assert len(list(tmp_path.glob("proton_second_*.hatchet"))) == 1 + + def test_failed_export_clears_activity_and_allows_next_interval(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + proton.data.get.side_effect = [OSError("disk full"), []] + wrapper.start() + wrapper.stop() + proton.data.clear.assert_called_with(7, 1) + wrapper.start() + wrapper.stop() + proton.data.clear.assert_called_with(7, 2) + assert len(list(tmp_path.glob("*_run1.hatchet"))) == 1 + + def test_cuda_graph_context_deactivates_after_capture_error(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + + with ( + pytest.raises(RuntimeError, match="capture failed"), + wrapper.capture_cuda_graphs(), + ): + raise RuntimeError("capture failed") + + proton.deactivate.assert_called_once_with(session=7, flushing=True) + + def test_cuda_graph_capture_deactivates_when_phase_advance_fails(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + proton.data.advance_phase.side_effect = RuntimeError("advance failed") + + with ( + pytest.raises(RuntimeError, match="advance failed"), + wrapper.capture_cuda_graphs(), + ): + pass + + proton.deactivate.assert_called_once_with(session=7, flushing=True) + proton.data.clear.assert_not_called() + + def test_cuda_graph_stop_deactivates_when_phase_advance_fails(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + proton.data.clear.reset_mock() + proton.data.advance_phase.side_effect = RuntimeError("advance failed") + + wrapper._start() + with pytest.raises(RuntimeError, match="advance failed"): + wrapper._stop() + + proton.deactivate.assert_called_with(session=7, flushing=True) + proton.data.clear.assert_not_called() + + def test_shutdown_finalizes_cuda_graph_capture_session(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + + wrapper.shutdown() + wrapper.shutdown() + + proton.finalize.assert_called_once_with(session=7) + def test_missing_proton_has_actionable_error(self, tmp_path): config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path)) with ( @@ -638,6 +871,7 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run(): ), patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, ): + wrapper.return_value.has_cuda_graph_session = False Worker.profile(worker, profile_prefix="first") Worker.profile(worker, is_start=False) Worker.profile(worker, profile_prefix="second") @@ -647,3 +881,157 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run(): call(worker.profiler_config, worker_name="second_rank1"), ] assert wrapper.return_value.start.call_count == 2 + + +@_requires_cuda_for_proton +def test_gpu_worker_reuses_cuda_graph_proton_session(): + worker = MagicMock() + worker.rank = 1 + worker.profiler = MagicMock(spec=ProtonProfilerWrapper) + worker.profiler.has_cuda_graph_session = True + worker.profiler_config.profiler = "proton" + + with patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank1", + ): + Worker.profile(worker, profile_prefix="first") + Worker.profile(worker, is_start=False) + + worker.profiler.set_output_name.assert_called_once_with("first_rank1") + worker.profiler.start.assert_called_once_with() + worker.profiler.stop.assert_called_once_with() + + +@_requires_cuda_for_proton +@pytest.mark.parametrize("runner", ["attribution_off", "v1", "no_capture"]) +def test_proton_not_initialized_without_capture(runner): + worker = MagicMock() + worker.profiler = None + worker.profiler_config.profiler = "proton" + worker.profiler_config.proton_graph_attribution = runner != "attribution_off" + worker.use_v2_model_runner = runner != "v1" + worker.model_runner.needs_cudagraph_capture.return_value = runner != "no_capture" + + context = Worker._get_cudagraph_capture_context(worker) + + assert worker.profiler is None + with context: + pass + + +@_requires_cuda_for_proton +def test_proton_initializes_before_cuda_graph_capture(): + class FakeProtonProfiler: + def __init__(self, config, worker_name): + self.config = config + self.worker_name = worker_name + self.capture_context = nullcontext() + + def capture_cuda_graphs(self): + return self.capture_context + + worker = MagicMock() + worker.rank = 2 + worker.profiler = None + worker.profiler_config.profiler = "proton" + worker.profiler_config.proton_graph_attribution = True + worker.use_v2_model_runner = True + worker.model_runner.needs_cudagraph_capture.return_value = True + + with ( + patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank2", + ), + patch( + "vllm.profiler.wrapper.ProtonProfilerWrapper", + FakeProtonProfiler, + ), + ): + context = Worker._get_cudagraph_capture_context(worker) + + assert worker.profiler.config is worker.profiler_config + assert worker.profiler.worker_name == "rank2" + assert context is worker.profiler.capture_context + + +@_requires_cuda_for_proton +@pytest.mark.parametrize("context", ["shadow", "python"]) +@pytest.mark.parametrize("output_format", ["hatchet", "hatchet_msgpack"]) +def test_proton_cuda_graph_replay_attribution_on_gpu(tmp_path, context, output_format): + """Both intervals contain replay kernels, without capture-only activity.""" + import json + + import torch + import triton + import triton.profiler as proton + from packaging.version import Version + + if Version(triton.__version__) < Version("3.7"): + pytest.skip("Graph attribution requires Triton >= 3.7") + + wrapper = ProtonProfilerWrapper( + ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + proton_context=context, + proton_output_format=output_format, + ), + worker_name="gpu", + ) + worker = SimpleNamespace( + profiler=wrapper, + use_v2_model_runner=True, + ) + x = torch.ones(1024, device="cuda") + graph = torch.cuda.CUDAGraph() + + def capture_only(): + x.add_(1) + + def captured_add(): + x.add_(1) + + def kernel_metrics(value): + if isinstance(value, list): + return [metric for child in value for metric in kernel_metrics(child)] + if isinstance(value, dict): + metrics = value.get("metrics", {}) + return ([metrics] if metrics.get("count", 0) else []) + [ + metric for child in value.values() for metric in kernel_metrics(child) + ] + return [] + + try: + with Worker._get_cudagraph_capture_context(worker): + with proton.scope("capture_only"): + capture_only() + torch.accelerator.synchronize() + with torch.cuda.graph(graph), proton.scope("captured_add"): + captured_add() + for run in range(2): + wrapper.start() + with wrapper.annotate_context_manager(f"replay_{run}"): + graph.replay() + wrapper.stop() + (path,) = tmp_path.glob(f"*_run{run}.{output_format}") + if output_format == "hatchet_msgpack": + import msgpack + + data = msgpack.unpackb(path.read_bytes()) + else: + data = json.loads(path.read_text()) + serialized = json.dumps(data) + assert "" in serialized + assert "captured_add" in serialized + assert "capture_only" not in serialized + assert f"replay_{1 - run}" not in serialized + metrics = kernel_metrics(data) + assert sum(metric["count"] for metric in metrics) == 1 + assert sum(metric["time (ns)"] for metric in metrics) > 0 + torch.testing.assert_close(x, torch.full_like(x, 4)) + finally: + wrapper.shutdown() + assert not list(tmp_path.glob(".proton_cuda_graph_session*")) diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index b3b232490f01..535960a8d7cd 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -76,6 +76,10 @@ class ProfilerConfig: """Optional format passed to Proton when finalizing a profile. ``None`` uses the default format for ``proton_data``.""" + proton_graph_attribution: bool = False + """Observe CUDA graph capture so replayed kernels can be attributed. + Requires Triton >= 3.7 and ``proton_data='tree'``.""" + torch_profiler_with_stack: bool = True """If `True`, enables stack tracing in the torch profiler. Enabled by default as it is useful for debugging. Can be disabled via @@ -199,6 +203,7 @@ def _validate_profiler_config(self) -> Self: ("proton_mode", self.proton_mode, None), ("proton_hook", self.proton_hook, None), ("proton_output_format", self.proton_output_format, None), + ("proton_graph_attribution", self.proton_graph_attribution, False), ) if value != default ] @@ -220,6 +225,17 @@ def _validate_profiler_config(self) -> Self: if self.profiler == "proton": output_format = self.proton_output_format + if self.proton_graph_attribution and self.proton_data != "tree": + raise ValueError("proton_graph_attribution requires proton_data='tree'") + if ( + self.proton_graph_attribution + and self.proton_mode + and self.proton_mode.split(":", 1)[0].lower() == "periodic_flushing" + ): + raise ValueError( + "proton_graph_attribution is incompatible with periodic_flushing: " + "both manage Proton data phases." + ) if output_format == "chrome_trace" and self.proton_data != "trace": raise ValueError("chrome_trace output requires proton_data='trace'") if ( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2fcc144af0c3..7d2baea26ce9 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1551,18 +1551,6 @@ def __post_init__(self): self.compilation_config.mode = CompilationMode.NONE self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - if self.profiler_config.profiler == "proton": - if not current_platform.is_cuda(): - raise ValueError( - "The Proton profiler currently supports NVIDIA CUDA only" - ) - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - raise ValueError( - "The Proton profiler requires CUDA graphs to be disabled. " - "Use --enforce-eager or set " - "--compilation-config.cudagraph_mode=none." - ) - if os.environ.get("TORCH_COMPILE_DISABLE") == "1": logger.warning_once( "TORCH_COMPILE_DISABLE is set, disabling torch.compile. " @@ -1860,6 +1848,7 @@ def has_blocked_weights(): else: self._validate_v1_model_runner() + self._validate_profiler_config() self._validate_batch_sharded_sampling() self._validate_adaptive_verification() @@ -3042,6 +3031,44 @@ def _get_dbo_unsupported_features(self) -> list[str]: return unsupported + def _validate_profiler_config(self) -> None: + if self.profiler_config.profiler != "proton": + return + + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + raise ValueError("The Proton profiler currently supports NVIDIA CUDA only") + if ( + self.profiler_config.proton_graph_attribution + and not self.use_v2_model_runner + ): + raise ValueError( + "Proton CUDA graph attribution requires the V2 model runner." + ) + + has_cuda_graphs = ( + self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + or ( + self.compilation_config.cudagraph_mm_encoder + and not (self.model_config and self.model_config.enforce_eager) + ) + ) + if not has_cuda_graphs: + return + mode = self.profiler_config.proton_mode + if mode and mode.split(":", 1)[0].lower() == "pcsampling": + raise ValueError( + "Proton PC sampling requires CUDA graphs to be disabled. " + "Use --enforce-eager." + ) + if not self.profiler_config.proton_graph_attribution: + raise ValueError( + "Proton profiling with CUDA graphs requires " + "proton_graph_attribution=True to capture replayed kernels. " + "Enable attribution or use --enforce-eager." + ) + def _validate_v2_model_runner(self) -> None: """Check for features not yet supported by the V2 model runner.""" if not HAS_TRITON: diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index fb8963198635..abe8a3431798 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -6,8 +6,9 @@ import json import os from abc import ABC, abstractmethod -from collections.abc import Callable -from contextlib import nullcontext +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress +from glob import glob from typing import Literal from uuid import uuid4 @@ -160,6 +161,15 @@ def shutdown(self) -> None: if self._running: self.stop() + @property + def has_cuda_graph_session(self) -> bool: + """Whether a capture-time session is retained to attribute replays.""" + return False + + def capture_cuda_graphs(self) -> AbstractContextManager[None]: + """Observe graph creation for backends that attribute replay activity.""" + return nullcontext() + def annotate_context_manager(self, name: str): """Return a context manager to annotate profiler traces.""" return nullcontext() @@ -387,6 +397,7 @@ def __init__( self._mode = profiler_config.proton_mode self._hook = profiler_config.proton_hook self._output_format = profiler_config.proton_output_format + self._graph_attribution = profiler_config.proton_graph_attribution self._triton_version_string = getattr(triton, "__version__", "unknown") try: self._triton_version = Version(self._triton_version_string) @@ -398,6 +409,13 @@ def __init__( # worker cannot overwrite profiles left by an earlier server process. self._instance_id = f"pid{os.getpid()}_{uuid4().hex}" self._run_id = 0 + self._graph_session = False + self._phase = 0 + self._active_output_path: str | None = None + self._session_storage_path = os.path.join( + self._output_dir, + f".proton_cuda_graph_session_{worker_name}_{self._instance_id}", + ) logger.info_once( "Proton profiling enabled. Output will be saved under: %s", @@ -412,6 +430,10 @@ def _require_triton_version(self, feature: str, minimum: Version) -> None: ) def _validate_capabilities(self) -> None: + if self._graph_attribution: + self._require_triton_version( + "CUDA graph attribution", _TRITON_PROTON_3_7_VERSION + ) if self._output_format is not None: parameters = inspect.signature(self._proton.finalize).parameters supports_output_format = "output_format" in parameters or any( @@ -428,7 +450,7 @@ def _validate_capabilities(self) -> None: self._require_triton_version( "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION ) - if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing": + if self._mode and self._mode.split(":", 1)[0].lower() == "periodic_flushing": self._require_triton_version( "periodic flushing", _TRITON_PROTON_3_7_VERSION ) @@ -447,28 +469,112 @@ def _create_session(self, output_path: str) -> int: raise RuntimeError("Proton did not create a profiling session") return session_id + @property + def has_cuda_graph_session(self) -> bool: + return self._graph_session + + def set_output_name(self, worker_name: str) -> None: + """Set the next run's output name after startup graph capture.""" + if self._active: + return + self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}") + + @override + @contextmanager + def capture_cuda_graphs(self) -> Iterator[None]: + """Keep a Proton session active while vLLM captures CUDA graphs.""" + if not self._graph_attribution: + yield + return + if self._session_id is None: + self._session_id = self._create_session(self._session_storage_path) + else: + self._proton.activate(session=self._session_id) + self._graph_session = True + + try: + yield + finally: + captured_phase = self._phase + try: + self._phase = self._proton.data.advance_phase(self._session_id) + finally: + self._proton.deactivate(session=self._session_id, flushing=True) + self._proton.data.clear(self._session_id, captured_phase) + @override def _start(self) -> None: - output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}" - self._session_id = self._create_session(output_path) + self._active_output_path = ( + f"{self._output_path}_{self._instance_id}_run{self._run_id}" + ) self._run_id += 1 + if self._graph_session: + assert self._session_id is not None + self._proton.activate(session=self._session_id) + else: + self._session_id = self._create_session(self._active_output_path) + + def _write_graph_phase(self, phase: int) -> None: + assert self._active_output_path is not None + output_format = self._output_format or "hatchet" + output_path = f"{self._active_output_path}.{output_format}" + if output_format == "hatchet_msgpack": + with open(output_path, "wb") as output_file: + output_file.write( + self._proton.data.get_msgpack(self._session_id, phase) + ) + else: + with open(output_path, "w", encoding="utf-8") as output_file: + json.dump(self._proton.data.get(self._session_id, phase), output_file) + + def _finalize_session(self, session_id: int) -> None: + if self._output_format is None: + self._proton.finalize(session=session_id) + else: + self._proton.finalize(session=session_id, output_format=self._output_format) @override def _stop(self) -> None: assert self._session_id is not None session_id = self._session_id + if self._graph_session: + completed_phase = self._phase + try: + self._phase = self._proton.data.advance_phase(session_id) + finally: + self._proton.deactivate(session=session_id, flushing=True) + try: + self._write_graph_phase(completed_phase) + finally: + self._proton.data.clear(session_id, completed_phase) + self._active_output_path = None + return + try: self._proton.deactivate(session=session_id) finally: try: - if self._output_format is None: - self._proton.finalize(session=session_id) - else: - self._proton.finalize( - session=session_id, output_format=self._output_format - ) + self._finalize_session(session_id) finally: self._session_id = None + self._active_output_path = None + + @override + def shutdown(self) -> None: + super().shutdown() + if self._graph_session and self._session_id is not None: + session_id = self._session_id + self._session_id = None + try: + self._finalize_session(session_id) + except Exception: + logger.exception( + "Failed to finalize Proton CUDA graph session during shutdown." + ) + finally: + for output_path in glob(f"{self._session_storage_path}.*"): + with suppress(FileNotFoundError): + os.remove(output_path) @override def annotate_context_manager(self, name: str): @@ -496,3 +602,24 @@ def _stop(self) -> None: @override def annotate_context_manager(self, name: str): return torch.cuda.nvtx.range(name) + + +def create_graph_capture_profiler( + profiler_config: ProfilerConfig, global_rank: int +) -> WorkerProfiler | None: + """Create a profiler to observe CUDA graph capture, if configured. + + Applies only to backends that attribute graph replay activity which + need a session around capture (Proton with ``proton_graph_attribution``). + """ + if ( + profiler_config.profiler == "proton" + and profiler_config.proton_graph_attribution + ): + from vllm.distributed.utils import get_worker_rank_suffix + + return ProtonProfilerWrapper( + profiler_config, + worker_name=get_worker_rank_suffix(global_rank=global_rank), + ) + return None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d66fc6c90f7c..4b1520f276b4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -987,6 +987,16 @@ def profile_cudagraph_memory(self) -> int: """Estimate the GPU memory required to capture CUDA graphs.""" return _profile_cudagraph_memory(self) + def needs_cudagraph_capture(self) -> bool: + """Whether capture_model() has any CUDA graphs left to capture.""" + return ( + self.model_state.supports_mm_inputs + and self.model_state.encoder_runner.has_cudagraph() + ) or ( + self.cudagraph_manager is not None + and self.cudagraph_manager.needs_capture() + ) + @torch.inference_mode() def capture_model(self, *, profile_only: bool = False) -> int: assert self.cudagraph_manager is not None diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 2e84fd511a22..3f457b9b7b4a 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -9,7 +9,7 @@ from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta from types import NoneType -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np import regex as re @@ -64,6 +64,7 @@ CudaProfilerWrapper, ProtonProfilerWrapper, TorchProfilerWrapper, + create_graph_capture_profiler, ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask @@ -142,6 +143,7 @@ def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int if TYPE_CHECKING: from vllm.device_allocator.sleep_mode_backend import SleepModeBackend from vllm.model_executor.model_loader.tensorizer import TensorizerConfig + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -810,7 +812,8 @@ def compile_or_warm_up_model(self) -> CompilationTimes: cuda_graph_memory_bytes = 0 if not self.model_config.enforce_eager: - cuda_graph_memory_bytes = self.model_runner.capture_model() + with self._get_cudagraph_capture_context(): + cuda_graph_memory_bytes = self.model_runner.capture_model() # Compare actual vs estimated CUDA graph memory (if we did profiling) if ( @@ -949,6 +952,21 @@ def compile_or_warm_up_model(self) -> CompilationTimes: encoder=self.compilation_config.encoder_compilation_time, ) + def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]: + """Let the configured profiler observe CUDA graph capture.""" + if not self.use_v2_model_runner: + return nullcontext() + if self.profiler is None: + model_runner = cast("GPUModelRunnerV2", self.model_runner) + if not model_runner.needs_cudagraph_capture(): + return nullcontext() + self.profiler = create_graph_capture_profiler( + self.profiler_config, global_rank=self.rank + ) + if self.profiler is None: + return nullcontext() + return self.profiler.capture_cuda_graphs() + def reset_mm_cache(self) -> None: self.model_runner.reset_mm_cache() @@ -1248,6 +1266,9 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): else: trace_name = rank_suffix + if profiler_type == "proton" and self.profiler is not None: + self.profiler.set_output_name(trace_name) + # Create the profiler wrapper only on the first start call if self.profiler is None: if profiler_type == "torch": @@ -1284,7 +1305,9 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): try: self.profiler.stop() finally: - if self.profiler_config.profiler == "proton": + if self.profiler_config.profiler == "proton" and not ( + self.profiler.has_cuda_graph_session + ): # Proton output names are fixed when the wrapper is constructed. # Recreate it so the next profile_prefix is honored. self.profiler = None diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index 8208a6ce3df9..e19f1b332a68 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/vllm/v1/worker/mm_encoder_model_runner.py @@ -70,6 +70,9 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: def capture_model(self, *, profile_only: bool = False) -> int: return 0 + def needs_cudagraph_capture(self) -> bool: + return False + def _dummy_run( self, *args: Any, **kwargs: Any ) -> tuple[torch.Tensor, torch.Tensor]: From 995e8581f462a13e32f30cfba946c63d48cf31d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 21:16:39 +0200 Subject: [PATCH 200/774] [Model] Voxtral Realtime: add support for `CUDAGraphMode.FULL_DECODE_ONLY` (#51167) Signed-off-by: NickLucche Co-authored-by: Jorge Gallego Feliciano --- .../generation/test_voxtral_realtime.py | 106 ++++++++++------ .../model_executor/models/voxtral_realtime.py | 12 +- vllm/model_executor/models/whisper_causal.py | 118 +++++++++++++++--- 3 files changed, 174 insertions(+), 62 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 0ef3c33aebc0..1bd9f6230353 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -14,7 +14,9 @@ from vllm import LLM, SamplingParams from vllm.assets.audio import AudioAsset +from vllm.config import CUDAGraphMode from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.kv_cache_interface import SlidingWindowSpec @@ -130,6 +132,45 @@ async def async_engine(): wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) +def _build_inputs( + audio_assets: list[AudioAsset], tokenizer: MistralTokenizer +) -> tuple[list[dict], list[SamplingParams]]: + audio_config = tokenizer.instruct_tokenizer.tokenizer.audio + inputs = [] + sampling_params = [] + for audio_asset in audio_assets: + audio = Audio.from_file(audio_asset.get_local_path(), strict=False) + req = TranscriptionRequest( + audio=audio.to_base64(audio.format), + streaming=StreamingMode.OFFLINE, + language=None, + ) + tokenized = tokenizer.instruct_tokenizer.encode_transcription(req) + tokens = tokenized.tokens + audio_array = tokenized.audios[0].audio_array + + num_samples = audio_array.shape[0] + max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1 + sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens)) + inputs.append( + { + "multi_modal_data": {"audio": [(audio_array, None)]}, + "prompt_token_ids": tokens, + } + ) + return inputs, sampling_params + + +def _assert_expected_text(outputs) -> None: + texts = _normalize([out.outputs[0].text for out in outputs]) + for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)): + assert got == expected, ( + f"Output mismatch at index {i}:\n" + f" got: {got!r}\n" + f" expected: {expected!r}" + ) + + def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch): monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") @@ -138,51 +179,38 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa with vllm_runner(**vllm_kwargs) as vllm_model: assert_encoder_kv_cache_spec(vllm_model.llm) - audio_config = tokenizer.instruct_tokenizer.tokenizer.audio - - def from_file(file_path: str): - audio = Audio.from_file(file_path, strict=False) - req = TranscriptionRequest( - audio=audio.to_base64(audio.format), - streaming=StreamingMode.OFFLINE, - language=None, - ) - tokenized = tokenizer.instruct_tokenizer.encode_transcription(req) + inputs, sampling_params = _build_inputs(audio_assets, tokenizer) - return (tokenized.tokens, tokenized.audios[0].audio_array) - - tokenized_list = [ - from_file(audio_asset.get_local_path()) for audio_asset in audio_assets - ] + outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) + _assert_expected_text(outputs) - inputs = [] - sampling_params = [] - for tokens, audio_array in tokenized_list: - num_samples = audio_array.shape[0] - max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1 - sampling_params.append( - SamplingParams(temperature=0.0, max_tokens=max_tokens) - ) +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="Full cudagraphs require CUDA" +) +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_voxtral_realtime_cudagraph( + audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode +): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - input_dict = { - "multi_modal_data": {"audio": [(audio_array, None)]}, - "prompt_token_ids": tokens, - } - inputs.append(input_dict) + vllm_kwargs = {**ENGINE_CONFIG} + vllm_kwargs["model_name"] = vllm_kwargs.pop("model") + vllm_kwargs["enforce_eager"] = False + vllm_kwargs["compilation_config"] = { + "cudagraph_mode": cudagraph_mode, + # max_num_seqs is 4, so this is the only decode size worth capturing. + "cudagraph_capture_sizes": [4], + } - outputs = vllm_model.llm.generate( - inputs, - sampling_params=sampling_params, - ) + with vllm_runner(**vllm_kwargs) as vllm_model: + inputs, sampling_params = _build_inputs(audio_assets, tokenizer) - texts = _normalize([out.outputs[0].text for out in outputs]) - for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)): - assert got == expected, ( - f"Output mismatch at index {i}:\n" - f" got: {got!r}\n" - f" expected: {expected!r}" - ) + outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) + _assert_expected_text(outputs) @pytest.mark.asyncio diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 738185a7aecf..961aa3af9c9a 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -217,9 +217,15 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - assert ( - not vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs() - ), "Voxtral realtime doesn't support full cudagraphs yet. Please use PIECEWISE." + # Full cudagraphs are supported for decode-only batches (the encoder's + # block-pooling attention builder is capture-safe for uniform + # single-token decode). Mixed/prefill batches are not, so block only + # pure FULL; FULL_DECODE_ONLY and FULL_AND_PIECEWISE are allowed. + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + assert not cudagraph_mode.mixed_mode().has_full_cudagraphs(), ( + "Voxtral realtime supports full cudagraphs for decode-only batches. " + "Use cudagraph_mode=FULL_DECODE_ONLY or FULL_AND_PIECEWISE (not FULL)." + ) self.time_embedding: TimeEmbedding = TimeEmbedding( dim=self.config.text_config.hidden_size diff --git a/vllm/model_executor/models/whisper_causal.py b/vllm/model_executor/models/whisper_causal.py index 7004bba2838e..ec84cae8cfcf 100644 --- a/vllm/model_executor/models/whisper_causal.py +++ b/vllm/model_executor/models/whisper_causal.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy import functools import logging import math from dataclasses import replace from fractions import Fraction from functools import partial +from typing import ClassVar import torch import torch.nn.functional as F @@ -27,6 +27,7 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import ( AttentionBackend, + AttentionCGSupport, AttentionMetadata, AttentionType, CommonAttentionMetadata, @@ -123,6 +124,16 @@ def create_whisper_attention_backend_with_block_pooling( underlying_impl = underlying_attn_backend.get_impl_cls() class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore + # Full cudagraphs only for uniform single-token decode: capture bakes in + # tensor addresses, so `build` writes metadata into persistent buffers + # (below). Prefill/mixed batches fall back to piecewise. + _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) + # Disable the update_block_table fast path: it would bypass block pooling + # by splicing the raw slot_mapping/block_table into cached metadata. + supports_update_block_table: bool = False + def __init__( self, kv_cache_spec: AttentionSpec, @@ -150,6 +161,38 @@ def __init__( # num_heads_q for the encoder is the same as num_kv_heads # (no GQA in whisper encoder) self.num_heads_q = kv_cache_spec.num_kv_heads + # Persistent capture-safe buffers for pooled metadata; lazily sized + # on first `build` to match the incoming tensors' dtype/device. + self._pool_buffers_ready = False + + def _maybe_init_pool_buffers( + self, common_attn_metadata: CommonAttentionMetadata + ) -> None: + if self._pool_buffers_ready: + return + device = common_attn_metadata.slot_mapping.device + sched = self.vllm_config.scheduler_config + self._qsl_buf = torch.empty( + sched.max_num_seqs + 1, + dtype=common_attn_metadata.query_start_loc.dtype, + device=device, + ) + self._seq_buf = torch.empty( + sched.max_num_seqs, + dtype=common_attn_metadata.seq_lens.dtype, + device=device, + ) + self._slot_buf = torch.empty( + sched.max_num_batched_tokens * block_pool_size, + dtype=common_attn_metadata.slot_mapping.dtype, + device=device, + ) + self._pool_arange = torch.arange( + block_pool_size, + dtype=common_attn_metadata.slot_mapping.dtype, + device=device, + ) + self._pool_buffers_ready = True def build( self, @@ -157,27 +200,62 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> AttentionMetadata: - new_common_attn_metadata = copy.deepcopy(common_attn_metadata) - new_common_attn_metadata.query_start_loc *= block_pool_size - new_common_attn_metadata.query_start_loc_cpu *= block_pool_size - new_common_attn_metadata.seq_lens *= block_pool_size - if new_common_attn_metadata.seq_lens_cpu_upper_bound is not None: - new_common_attn_metadata.seq_lens_cpu_upper_bound *= block_pool_size - new_common_attn_metadata.num_actual_tokens *= block_pool_size - new_common_attn_metadata.max_query_len *= block_pool_size - new_common_attn_metadata.max_seq_len *= block_pool_size - original_slot_mapping = common_attn_metadata.slot_mapping - common_prefix_len *= block_pool_size - new_common_attn_metadata.slot_mapping = ( - ( - original_slot_mapping.unsqueeze(1) * block_pool_size - + torch.arange(block_pool_size, device=original_slot_mapping.device) - ) - .flatten() - .clamp(min=-1) + cm = common_attn_metadata + self._maybe_init_pool_buffers(cm) + + # Scale query_start_loc / seq_lens into the persistent buffers. + n = cm.query_start_loc.numel() + assert n <= self._qsl_buf.numel(), ( + f"{n=} exceeds pooled query_start_loc buffer {self._qsl_buf.numel()}" + ) + qsl_buf = self._qsl_buf[:n] + qsl_buf.copy_(cm.query_start_loc) + qsl_buf.mul_(block_pool_size) + + r = cm.seq_lens.numel() + assert r <= self._seq_buf.numel(), ( + f"{r=} exceeds pooled seq_lens buffer {self._seq_buf.numel()}" + ) + seq_buf = self._seq_buf[:r] + seq_buf.copy_(cm.seq_lens) + seq_buf.mul_(block_pool_size) + + # Expand each slot into block_pool_size contiguous slots. Padding + # slots (-1) stay -1 after clamp so the KV write is skipped. + m = cm.slot_mapping.numel() + assert m * block_pool_size <= self._slot_buf.numel(), ( + f"pooled slot_mapping {m * block_pool_size} exceeds buffer " + f"{self._slot_buf.numel()}" + ) + slot_buf = self._slot_buf[: m * block_pool_size] + slot_buf.view(m, block_pool_size).copy_( + cm.slot_mapping.unsqueeze(1) * block_pool_size + self._pool_arange + ) + slot_buf.clamp_(min=-1) + + # CPU tensors aren't captured; fresh scaled copies avoid mutating + # the runner's shared buffers. + qsl_cpu = cm.query_start_loc_cpu * block_pool_size + seq_lens_cpu_upper_bound = ( + cm.seq_lens_cpu_upper_bound * block_pool_size + if cm.seq_lens_cpu_upper_bound is not None + else None + ) + + new_common_attn_metadata = cm.replace( + query_start_loc=qsl_buf, + query_start_loc_cpu=qsl_cpu, + seq_lens=seq_buf, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + num_actual_tokens=cm.num_actual_tokens * block_pool_size, + max_query_len=cm.max_query_len * block_pool_size, + max_seq_len=cm.max_seq_len * block_pool_size, + slot_mapping=slot_buf, ) return super().build( - common_prefix_len, new_common_attn_metadata, fast_build + common_prefix_len * block_pool_size, + new_common_attn_metadata, + fast_build, ) # NOTE: We need a custom impl so we can use the transformed slot_mapping From d4ee7fe7a98b3012f64309e18dd425e565069a01 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:24:18 +0300 Subject: [PATCH 201/774] [Quantization] Select linear backends per quantization (#51204) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- docs/features/quantization/README.md | 17 +++++ .../test_scaled_mm_kernel_selection.py | 70 +++++++++++++++++++ .../kernels/test_b12x_linear.py | 8 ++- vllm/config/kernel.py | 9 +++ .../model_executor/kernels/linear/__init__.py | 65 +++++++++++++---- .../kernels/linear/mxfp4/aiter.py | 2 +- 6 files changed, 154 insertions(+), 17 deletions(-) diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 69ece3607610..ba1a5b86fe45 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM: - [TorchAO](torchao.md) - [FP8 ViT Encoder Attention](fp8_vit_attn.md) +## Selecting Linear Backends per Quantization + +`--linear-backend` selects one backend for all quantized linear layers. For +mixed-precision models that use more than one linear quantization scheme, use +`linear_backend_per_quant` to override the backend for individual schemes: + +```bash +vllm serve \ + --linear-backend cutlass \ + --kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}' +``` + +Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear +layers use CUTLASS. Per-quantization overrides take precedence over +`--linear-backend`; schemes without an override continue to use the global +setting, including automatic selection when it is `auto`. + ## Supported Hardware The table below shows the compatibility of various quantization implementations with different hardware platforms in vLLM: diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index bedebdb59b85..827571474481 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -12,20 +12,90 @@ import pytest import torch +from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.kernels.linear import ( AiterInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, + HummingFP8ScaledMMLinearKernel, Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, ScaledMMLinearKernel, + _get_linear_backend, + _resolve_backend_kernels, + init_fp8_linear_kernel, init_int8_linear_kernel, register_linear_kernel, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, +) from vllm.platforms import PlatformEnum pytestmark = pytest.mark.cpu_test +def test_auto_linear_backend_is_unchanged_without_overrides(): + kernels = [object] + quantizations = ( + "fp8_block_w8a8", + "fp8_w8a8", + "int8_w8a8", + "mixed_precision", + "mxfp8", + "mxfp4", + "mxfp6", + "w8a16_fp8", + "nvfp4_w4a4", + "nvfp4_w4a16", + ) + config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + + with set_current_vllm_config(config): + for quantization in quantizations: + assert _get_linear_backend(quantization=quantization) == "auto" + assert ( + _resolve_backend_kernels(kernels, "test", quantization=quantization) + is kernels + ) + + +def test_linear_backend_override_is_quantization_specific(): + config = VllmConfig( + kernel_config=KernelConfig( + linear_backend="cutlass", + linear_backend_per_quant={"nvfp4_w4a16": "humming"}, + ) + ) + + with set_current_vllm_config(config): + assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming" + assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass" + assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass" + + +@patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None)) +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_fp8_linear_backend_override(platform_mock, _): + platform_mock._enum = PlatformEnum.CUDA + config = VllmConfig( + kernel_config=KernelConfig( + linear_backend="cutlass", + linear_backend_per_quant={"fp8_w8a8": "humming"}, + ) + ) + + with set_current_vllm_config(config): + kernel = init_fp8_linear_kernel( + activation_quant_key=kFp8StaticTensorSym, + weight_quant_key=kFp8StaticTensorSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=(128, 128), + ) + + assert type(kernel) is HummingFP8ScaledMMLinearKernel + + def test_is_supported_is_abstract(): """Test that is_supported() is properly defined as abstract.""" assert issubclass(ScaledMMLinearKernel, ABC) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 27368ead27dd..9d4f9beeb9a3 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -115,7 +115,9 @@ def test_b12x_backend_registration_priority_and_selection( assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after) monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x" + ) monkeypatch.setattr( kernel_cls, "is_supported", @@ -784,7 +786,9 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: import vllm.model_executor.kernels.linear as linear_mod monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x" + ) monkeypatch.setitem( linear_mod._POSSIBLE_NVFP4_KERNELS, PlatformEnum.CUDA, diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 4505e0cfd9a0..0ac5bf43913d 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -326,6 +326,13 @@ class KernelConfig: - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16) """ + linear_backend_per_quant: dict[str, LinearBackend] | None = Field( + default=None, min_length=1 + ) + """Backend overrides keyed by linear quantization scheme. Overrides take + precedence over ``linear_backend``; for example, + ``{"nvfp4_w4a16": "humming"}``.""" + @field_validator("moe_backend", mode="before") @classmethod def _normalize_moe_backend(cls, value: Any) -> Any: @@ -359,6 +366,8 @@ def compute_hash(self) -> str: "enable_flashinfer_autotune", "ir_op_priority", # handled separately below } + if self.linear_backend_per_quant is None: + ignored_factors.add("linear_backend_per_quant") factors = get_hash_factors(self, ignored_factors) factors["ir_op_priority"] = self.ir_op_priority.compute_hash() return hash_factors(factors) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1df73273ca39..d3c41f97d8e5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -235,14 +235,19 @@ logger = init_logger(__name__) -def _get_linear_backend() -> str: +def _get_linear_backend(*, quantization: str) -> str: """Get the linear_backend setting from the current vllm config.""" from vllm.config import get_current_vllm_config_or_none - config = get_current_vllm_config_or_none() - if config is not None: - return config.kernel_config.linear_backend - return "auto" + if (config := get_current_vllm_config_or_none()) is None: + return "auto" + overrides = config.kernel_config.linear_backend_per_quant + if overrides is not None and (override := overrides.get(quantization)): + logger.info_once( + "Applied linear backend override for %r: %r", quantization, override + ) + return override + return config.kernel_config.linear_backend # Kernel classes covered by each --linear-backend value. @@ -370,6 +375,8 @@ def _filter_kernels_by_backend( def _resolve_backend_kernels( kernels: list[type], layer_desc: str, + *, + quantization: str, ) -> list[type]: """Apply --linear-backend filtering to one layer type's kernel list. @@ -380,7 +387,7 @@ def _resolve_backend_kernels( layer types (e.g. NVFP4 MoE projections next to FP8 attention projections). """ - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization=quantization) if linear_backend == "auto": return kernels filtered = _filter_kernels_by_backend(linear_backend, kernels) @@ -625,6 +632,8 @@ def choose_scaled_mm_linear_kernel( possible_kernels: dict[PlatformEnum, list[type[_KernelT]]], compute_capability: int | None = None, force_kernel: type[_KernelT] | None = None, + *, + quantization: str, ) -> type[_KernelT]: """ Choose a _KernelT that can implement the given config for the @@ -642,6 +651,7 @@ def choose_scaled_mm_linear_kernel( force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override the possible_kernels if it can be implemented. If None, it will only try the possible kernels. + quantization: Quantization scheme used to select a backend override. Raises: ValueError: If no kernel can implement the given config. @@ -668,7 +678,11 @@ def choose_scaled_mm_linear_kernel( platform_kernels = possible_kernels.get(current_platform._enum, []) # Apply --linear-backend filtering when set. - platform_kernels = _resolve_backend_kernels(platform_kernels, "scaled-mm") + platform_kernels = _resolve_backend_kernels( + platform_kernels, + "scaled-mm", + quantization=quantization, + ) for kernel in platform_kernels: is_supported_and_can_implement, failure_reason = ( @@ -705,6 +719,7 @@ def init_fp8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] + quantization="fp8_block_w8a8", force_kernel=force_kernel, ) if module_name: @@ -736,6 +751,7 @@ def init_fp8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] + quantization="fp8_w8a8", force_kernel=force_kernel, ) if module_name: @@ -772,6 +788,7 @@ def init_int8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config, _POSSIBLE_INT8_KERNELS, + quantization="int8_w8a8", ) logger.info_once( @@ -824,7 +841,11 @@ def choose_mp_linear_kernel( platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, []) # Apply --linear-backend filtering when set. - platform_kernels = _resolve_backend_kernels(platform_kernels, "mixed-precision") + platform_kernels = _resolve_backend_kernels( + platform_kernels, + "mixed-precision", + quantization="mixed_precision", + ) failure_reasons = [] for kernel in platform_kernels: @@ -875,7 +896,11 @@ def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8Linea possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, [])) # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "MXFP8") + possible = _resolve_backend_kernels( + possible, + "MXFP8", + quantization="mxfp8", + ) failure_reasons = [] for kernel_cls in possible: @@ -917,7 +942,11 @@ def init_mxfp4_linear_kernel( possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, [])) # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "MXFP4") + possible = _resolve_backend_kernels( + possible, + "MXFP4", + quantization="mxfp4", + ) failure_reasons = [] for kernel_cls in possible: @@ -957,7 +986,7 @@ def init_mxfp6_linear_kernel( activation_quant_key=activation_quant_key, ) - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization="mxfp6") platform = current_platform._enum possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, [])) @@ -1017,7 +1046,10 @@ def init_wfp8_a16_linear_kernel( ) kernel_type = choose_scaled_mm_linear_kernel( - config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel + config, + _POSSIBLE_WFP8A16_KERNELS, + quantization="w8a16_fp8", + force_kernel=force_kernel, ) if module_name: @@ -1048,7 +1080,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: # batch-invariant CUTLASS implementation when available, otherwise fall # back to emulation. It overrides --linear-backend. force_kernel: type[NvFp4LinearKernel] | None = None - linear_backend = _get_linear_backend() + quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4" + linear_backend = _get_linear_backend(quantization=quantization) if envs.VLLM_BATCH_INVARIANT: bi_supported, reason = CutlassNvFp4LinearKernel.is_supported() if bi_supported: @@ -1110,7 +1143,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: possible = [kernel for kernel in possible if kernel in a16_kernels] # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "NVFP4") + possible = _resolve_backend_kernels( + possible, + "NVFP4", + quantization=quantization, + ) failure_reasons = [] for kernel_cls in possible: diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 13fd5db2750a..420e7f9493e1 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -158,7 +158,7 @@ def is_supported( from vllm._aiter_ops import is_aiter_found_and_supported from vllm.model_executor.kernels.linear import _get_linear_backend - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization="mxfp4") if ( current_platform.is_rocm() From 7702ee87dba0d8eed7f201e77a0a6613aac738a8 Mon Sep 17 00:00:00 2001 From: sychen52 <41452870+sychen52@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:25:33 -0700 Subject: [PATCH 202/774] [Bugfix][DSA] Write nvfp4_ds_mla from the fused norm+rope kernel (#55538) Signed-off-by: Shiyang Chen --- .../nvfp4_ds_mla_cache_kernels.cu | 14 +- .../test_fused_deepseek_v32_norm_rope.py | 120 ++++++++++++++++++ vllm/models/deepseek_v32/attention.py | 4 +- vllm/models/deepseek_v32/common/kernels.py | 68 +++++++++- 4 files changed, 186 insertions(+), 20 deletions(-) diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu index c2faf7b76023..b1f369e5ab00 100644 --- a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu +++ b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu @@ -109,9 +109,6 @@ __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw, #endif } -// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload -// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which -// would saturate the tile's largest values. __device__ __forceinline__ float nvfp4_tile_scale(const float* vals, float max_val, uint8_t* sf_out) { @@ -121,15 +118,8 @@ __device__ __forceinline__ float nvfp4_tile_scale(const float* vals, amax = fmaxf(amax, fabsf(vals[i])); } const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9 - __nv_fp8_e4m3 sf8(sf_f); - uint8_t sf_bits = *reinterpret_cast(&sf8); - if (float(sf8) < sf_f && sf_bits < 0x7E) { - // Bump to the next e4m3 value (bit patterns of positive e4m3 values are - // monotonic; 0x7E = 448 is the max finite value). - sf_bits += 1; - sf8 = *reinterpret_cast(&sf_bits); - } - *sf_out = sf_bits; + const __nv_fp8_e4m3 sf8(sf_f); + *sf_out = *reinterpret_cast(&sf8); return float(sf8); } diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index fbe22af2f897..7f74cde05d5b 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -573,6 +573,126 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" +E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor: + """Round to nearest e2m1, saturating to +-6 like cvt.rn.satfinite.e2m1x2.f32.""" + mags = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x.device) + a = x.float().abs().clamp_max(6.0) + mids = (mags[:-1] + mags[1:]) / 2 + code = torch.bucketize(a, mids, right=True) + on_tie = (a.unsqueeze(-1) == mids).any(dim=-1) + tie_code = torch.bucketize(a, mids, right=False) + code = torch.where(on_tie, tie_code + (tie_code & 1), code) + return (torch.signbit(x).to(torch.uint8) << 3) | code.to(torch.uint8) + + +def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor: + """Scale-factor byte permutation shared with FlashMLA: 8*(s&3) + (s>>2).""" + return 8 * (s % 4) + (s // 4) + + +@pytest.mark.skipif( + not current_platform.is_device_capability(100), + reason="nvfp4_ds_mla requires SM100 (Blackwell)", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) +def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int): + """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only). + + Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble + = even element) | 64 B unscaled e4m3 RoPE | 32 B byte-permuted e4m3 tile + scales, one per 16 NoPE elements. + """ + torch.manual_seed(5) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3 + # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale + # BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a + # future change to the scale rule cannot silently stop covering this. + # amax must land in (6, 9) * 2^-9 after rms_norm; 7.5 sits mid-band, with + # margin for the per-token RMS to vary. + kv_c[:, :16] = ( + 7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16) + ) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.ones(KV_LORA, device=dev, dtype=torch.bfloat16) + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos + mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + None, + None, + None, + EPS, + None, + topk, + slot_mapping=slot, + indexer_k_cache=None, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype="nvfp4_ds_mla", + mla_k_scale=None, + has_indexer=False, + index_rope_interleave=False, + ) + + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (nvfp4_ds_mla)") + + kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] + tiles = kv_ref.view(num_tokens, 32, 16) + amax = tiles.abs().amax(dim=-1) + scale_target = torch.clamp_min(amax / 6.0, 2.0**-9) + ref_scale = scale_target.to(FP8) # round-to-nearest + + # The crafted tile must actually exercise the saturating path. + assert (ref_scale[:, 0].float() < scale_target[:, 0]).all(), ( + "tile 0 should round its scale DOWN; adjust the crafted magnitude" + ) + assert (amax[:, 0] / ref_scale[:, 0].float() > 6.0).all(), ( + "tile 0 should saturate e2m1; adjust the crafted magnitude" + ) + + ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1)) + ref_codes = ref_codes.reshape(num_tokens, KV_LORA) + + cache = mla_cache[0, :num_tokens] # [N, 352] uint8 + packed = cache[:, : KV_LORA // 2] + got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1) + got_codes = got_codes.reshape(num_tokens, KV_LORA) # low nibble = even elem + torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0) + + got_rope = cache[:, KV_LORA // 2 : KV_LORA // 2 + ROPE_DIM].view(FP8) + assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3") + + perm = nvfp4_sf_byte(torch.arange(32, device=dev)) + got_scale = cache[:, KV_LORA // 2 + ROPE_DIM :].view(FP8)[:, perm] + torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0) + + # No indexer on this call: top-k buffer must be untouched. + assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" + + def test_fused_norm_rope_supports_large_token_count(): """Keep the token count off CUDA grid-y at its 65,536-block boundary.""" num_tokens = 65536 diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 00dd2e18a8d2..87cd618524a0 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -237,7 +237,9 @@ def __init__( fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) self._fp8_query = fp8_attention and self.impl.supports_quant_query_input - self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla" + self._fp8_kv_needs_view = fp8_attention and not self.kv_cache_dtype.endswith( + "_ds_mla" + ) self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index ef7fcde5d341..fd02a51076ed 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -150,6 +150,7 @@ def _fused_norm_rope_kernel( mla_cache_ds_scale_ptr, mla_cache_ds_rope_ptr, MLA_CACHE_DS_MLA: tl.constexpr, + MLA_CACHE_NVFP4: tl.constexpr, MLA_NUM_TILES: tl.constexpr, MLA_TILE_DIM: tl.constexpr, # Top k indices @@ -261,6 +262,56 @@ def _fused_norm_rope_kernel( mla_block_idx = slot_idx // mla_block_size mla_block_off = slot_idx % mla_block_size + if MLA_CACHE_NVFP4: + # nvfp4_ds_mla layout (352 B/token, KV_DIM == 512): + # [0, 256) 512 e2m1 NoPE, packed 2/byte (low nibble = even) + # [256, 320) 64 e4m3 RoPE, unscaled + # [320, 352) 32 e4m3 NoPE scales, byte-permuted + # Keep in lockstep with concat_and_cache_nvfp4_ds_mla (CUDA) and + # nvfp4_sf_byte() in FlashMLA: byte(s) = 8*(s&3) + (s>>2). + byte_base = ( + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + kv_2d = tl.reshape(kv_c.to(tl.float32), (MLA_NUM_TILES, MLA_TILE_DIM)) + amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) + # sf = e4m3(max(amax/6, 2^-9)), round-to-nearest like + # cvt_warp_fp16_to_fp4. + sf8 = tl.maximum(amax * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv) + q = kv_2d * (1.0 / sf8.to(tl.float32)) + qp = tl.reshape(q, (KV_DIM // 2, 2)) + sel = tl.arange(0, 2)[None, :] + even = tl.sum(tl.where(sel == 0, qp, 0.0), axis=1) + odd = tl.sum(tl.where(sel == 1, qp, 0.0), axis=1) + # Same instruction the CUDA writer uses, so the two are bit-exact + # by construction: dest byte = {high: odd, low: even}. + packed = tl.inline_asm_elementwise( + "{ .reg .b8 t; .reg .b32 w;" + " cvt.rn.satfinite.e2m1x2.f32 t, $2, $1;" + " mov.b32 w, {t, t, t, t}; and.b32 $0, w, 255; }", + "=r,f,f", + [even, odd], + dtype=tl.int32, + is_pure=True, + pack=1, + ).to(tl.uint8) + tl.store( + mla_cache_ptr + byte_base + tl.arange(0, KV_DIM // 2), + packed.to(tl.float8e4nv, bitcast=True), + ) + # RoPE: unscaled e4m3, interleaved like the fp8_ds_mla path + rope_base = mla_cache_ptr + byte_base + KV_DIM // 2 + tl.store(rope_base + dim_off * 2, r1.to(tl.float8e4nv)) + tl.store(rope_base + dim_off * 2 + 1, r2.to(tl.float8e4nv)) + # scales, permuted: byte(s) = 8*(s&3) + (s>>2) + sf_tile = tl.arange(0, MLA_NUM_TILES) + sf_perm = 8 * (sf_tile % 4) + (sf_tile // 4) + tl.store( + mla_cache_ptr + byte_base + (KV_DIM // 2 + 64) + sf_perm, + tl.reshape(sf8, (MLA_NUM_TILES,)), + ) + return + if MLA_CACHE_DS_MLA: # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element # tile of the NoPE is dynamically quantized to fp8 with its own @@ -491,7 +542,8 @@ def fused_norm_rope( idx_cache_stride = 0 # --- MLA KV cache setup --- - mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" + mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla" + mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" or mla_cache_nvfp4 mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla mla_num_tiles = 1 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device) @@ -499,11 +551,12 @@ def fused_norm_rope( if mla_kv_cache is not None: mla_block_size = mla_kv_cache.shape[1] if mla_cache_ds_mla: - # 656-byte custom layout addressed in bytes; mla_cache_ptr is the - # 1-byte fp8 view, so block/entry strides are byte offsets and the - # fp32/bf16 views share the same buffer. - assert kv_dim == 512, "fp8_ds_mla requires kv_lora_rank == 512" - mla_num_tiles = kv_dim // 128 + # Custom byte-addressed layout: 656 B/token for fp8_ds_mla, 352 B + # for nvfp4_ds_mla. mla_cache_ptr is the 1-byte fp8 view, so + # block/entry strides are byte offsets and the fp32/bf16 views + # (fp8_ds_mla only) share the same buffer. + assert kv_dim == 512, "ds_mla layouts require kv_lora_rank == 512" + mla_num_tiles = kv_dim // 16 if mla_cache_nvfp4 else kv_dim // 128 u8_cache = mla_kv_cache.view(torch.uint8) mla_block_stride = u8_cache.stride(0) mla_entry_stride = u8_cache.stride(1) @@ -596,7 +649,8 @@ def fused_norm_rope( mla_k_scale, mla_ds_scale_view, mla_ds_rope_view, - mla_cache_ds_mla, + mla_cache_ds_mla and not mla_cache_nvfp4, + mla_cache_nvfp4, mla_num_tiles, kv_dim // mla_num_tiles if mla_cache_ds_mla else 1, # Top k indices buffer From dabc4362b47ad2665b0802b0b42b13f660efe634 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:47:32 -0400 Subject: [PATCH 203/774] [ROCm][DSV4.1][Perf] Stride the DSA decode candidate mask over the live context (#56628) Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- .../test_rocm_aiter_candidate_mask.py | 226 ++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 120 +++++++++- 2 files changed, 344 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_candidate_mask.py diff --git a/tests/kernels/attention/test_rocm_aiter_candidate_mask.py b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py new file mode 100644 index 000000000000..85e742a09636 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm correctness tests for the strided DSA candidate mask. + +``_apply_candidate_mask_strided`` is a decode-path variant of the shared +``apply_candidate_mask``. It differs in two ways that need pinning down: + +* it launches a fixed number of programs that stride the column axis, rather + than one program per 1024-column tile, so at ``max_model_len`` widths the + grid no longer scales with a workspace that is mostly padding; +* it stops at each row's ``end`` instead of sanitizing the whole width. + +The second is only sound because every consumer bounds its scan by the same +ends -- ``top_k_per_row_decode`` takes ``rowEnd`` from ``seq_lens``. So the +contract is agreement with the shared kernel over ``[0, end)`` and nothing +beyond it, which is what these tests assert rather than comparing whole rows. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + +BLOCK_SIZE = 8 +TOPK_BLOCKS = 64 +# Above 128 tiles the grid saturates and programs start striding; below it the +# min() clamp applies instead. Both regimes need covering, so the widths here +# deliberately straddle 128 * 1024. +NARROW = 8192 +WIDE = 262144 + + +def _inputs(rows, width, ends, starts, *, block_size, col_stride, dtype): + """Logits, bounds and candidate blocks shared by both kernels.""" + torch.manual_seed(0) + # A column stride > 1 keeps the kernels honest about stride_col; the decode + # caller slices a workspace, so contiguous rows are not guaranteed. + logits = torch.randn(rows, width * col_stride, device="cuda")[:, ::col_stride] + logits[0, : min(16, width)] = 0.0 + if rows > 2: + logits[2, 5] = float("nan") + row_ke = torch.tensor(ends, device="cuda", dtype=dtype) + row_ks = torch.tensor(starts, device="cuda", dtype=dtype) if starts else None + nblocks = (width + block_size - 1) // block_size + candidates = torch.randint( + 0, nblocks, (rows, TOPK_BLOCKS), device="cuda", dtype=torch.int32 + ) + # -1 is the "no block" sentinel the selector emits for short rows. + candidates[:, ::7] = -1 + return logits, row_ks, row_ke, candidates + + +def _run_both(logits, row_ks, row_ke, candidates, block_size, row_repeat): + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_candidate_mask_strided, + ) + + ref = logits.clone() + got = logits.clone() + apply_candidate_mask(ref, row_ks, row_ke, candidates, block_size, row_repeat) + _apply_candidate_mask_strided( + got, row_ks, row_ke, candidates, block_size, row_repeat + ) + return ref, got + + +def _assert_agrees_within_ends(ref, got, row_ke, row_repeat): + """Compare only the region a consumer can read.""" + for row in range(ref.shape[0]): + end = int(row_ke[row // row_repeat]) + if end <= 0: + continue + torch.testing.assert_close( + got[row, :end], + ref[row, :end], + rtol=0, + atol=0, + equal_nan=True, + msg=lambda m, r=row, e=end: f"row {r} differs within [0, {e}): {m}", + ) + + +@pytest.mark.parametrize("width", [NARROW, WIDE]) +@pytest.mark.parametrize("col_stride", [1, 2]) +def test_matches_shared_kernel_ragged_ends(width, col_stride): + """Ends on and around tile boundaries, plus the empty and full rows.""" + ends = [0, 1, 1023, 1024, 1025, width - 1, width, width // 2] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + width, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=col_stride, + dtype=torch.int32, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + _assert_agrees_within_ends(ref, got, row_ke, 1) + + +@pytest.mark.parametrize("block_size", [8, 16]) +def test_matches_shared_kernel_with_starts(block_size): + """Non-zero starts: the decode caller passes zeros, prefill does not.""" + ends = [0, 64, 4096, 8192] + starts = [0, 7, 1024, 3] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + NARROW, + ends, + starts, + block_size=block_size, + col_stride=1, + dtype=torch.int64, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, block_size, 1) + _assert_agrees_within_ends(ref, got, row_ke, 1) + + +def test_matches_shared_kernel_row_repeat(): + """Speculative decode: next_n logit rows share one bound.""" + row_repeat = 3 + ends = [0, 1025, WIDE] + logits, row_ks, row_ke, candidates = _inputs( + len(ends) * row_repeat, + WIDE, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, row_repeat) + _assert_agrees_within_ends(ref, got, row_ke, row_repeat) + + +def test_leaves_columns_past_end_untouched(): + """The deliberate difference from the shared kernel, pinned explicitly. + + Skipping the tail is where the speedup comes from, so a change that + quietly restores full-width sanitizing should fail here rather than just + get slower. Work stops at the tile boundary containing ``end``, not at + ``end`` itself, so that is the bound asserted. + """ + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _MASK_TILE, + _apply_candidate_mask_strided, + ) + + ends = [0, 1, 1025, WIDE // 2] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + WIDE, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + sentinel = 1.5 + logits.fill_(sentinel) + _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + for row, end in enumerate(ends): + touched = min(-(-end // _MASK_TILE) * _MASK_TILE, WIDE) + tail = logits[row, touched:] + assert torch.equal(tail, torch.full_like(tail, sentinel)), ( + f"row {row} (end {end}) was written past column {touched}" + ) + # And the bound is tight: a row that ends early must not have cost a + # full-width pass, which is the whole point of the kernel. + assert touched <= end + _MASK_TILE + + +def test_cudagraph_replay_tracks_changing_ends(): + """The grid is static so a FULL capture stays valid when ends change. + + Capturing bakes the grid, so a data-dependent program count would freeze + whatever lengths were resident at capture. Replaying against longer and + shorter contexts than the captured ones is the property that matters. + """ + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_candidate_mask_strided, + ) + + rows = 4 + captured_ends = [1024, 4096, 8192, 2048] + logits, row_ks, row_ke, candidates = _inputs( + rows, + WIDE, + captured_ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + base = logits.clone() + + # Warm up outside the capture: Triton compiles on first call, and autotune + # launches would otherwise happen inside the graph. + _apply_candidate_mask_strided( + logits.clone(), row_ks, row_ke, candidates, BLOCK_SIZE, 1 + ) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + + for replay_ends in ([1, 131072, WIDE, 0], [WIDE, 1025, 3, 65536]): + row_ke.copy_(torch.tensor(replay_ends, device="cuda", dtype=row_ke.dtype)) + logits.copy_(base) + graph.replay() + torch.accelerator.synchronize() + + ref = base.clone() + apply_candidate_mask(ref, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + _assert_agrees_within_ends(ref, logits, row_ke, 1) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 598baee0b0b0..44b6017a057b 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -819,6 +819,123 @@ def rocm_fp8_mqa_logits( return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) +# Programs along the column axis of the decode mask grid. The shared kernel in +# model_executor/kernels/attention/dsa/candidate_blocks.py launches one program +# per 1024-column tile, which is 1024 of them per row at max_model_len 1048576 +# with only the first few doing work. This is a fixed count that strides +# instead, measured on gfx950; it is not a portable choice, which is why this +# variant lives here rather than replacing the shared one. +_MASK_GRID_COLS = 128 +# Columns each program handles per iteration. Work stops at the tile boundary +# containing a row's end rather than at the end itself; the overshoot is +# masked the same way the shared kernel masks it, so it costs a tile and +# changes nothing. +_MASK_TILE = 1024 + + +@triton.jit(do_not_specialize=["width", "nblocks"]) +def _mask_candidates_strided_kernel( + logits, + starts, + ends, + flags, + stride_row, + stride_col, + stride_start, + stride_end, + width, + nblocks, + BLOCK_SIZE: tl.constexpr, + HAS_STARTS: tl.constexpr, + ROW_REPEAT: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0 + end = tl.load(ends + row // ROW_REPEAT * stride_end) + edge = tl.load(flags + row * (nblocks + 1) + nblocks) + step = tl.num_programs(1) * TILE + tile_start = tl.program_id(1) * TILE + # The shared kernel also sanitizes columns at or past `end`. Nothing reads + # them: top_k_per_row_decode bounds its scan by the same row ends passed + # here, and persistent_topk/cooperative_topk clamp by the same lengths. + # Stopping at `end` is what makes the cost track the live context rather + # than max_model_len. + while tile_start < end: + cols = tile_start + tl.arange(0, TILE) + valid = (cols >= start) & (cols < end) & (cols < width) + block = (cols - start) // BLOCK_SIZE + keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0) + keep = (keep != 0) | ((cols == width - 1) & (edge != 0)) + tl.store( + logits + row * stride_row + cols * stride_col, + -float("inf"), + (cols < width) & ~(valid & keep), + ) + tile_start += step + + +def _apply_candidate_mask_strided( + logits: torch.Tensor, + row_ks: torch.Tensor | None, + row_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + block_size: int, + row_repeat: int = 1, +) -> None: + """ROCm decode variant of ``apply_candidate_mask``. + + Same masking semantics over ``[0, end)``, but the grid is sized by a fixed + program count rather than by the logits width. Only worth using where the + width is the ``max_model_len`` workspace and the live context is far + shorter, i.e. the paged decode path below; the prefill chunks pass + chunk-sized logits and stay on the shared kernel. + """ + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + _candidate_flags_kernel, + ) + + rows, width = logits.shape + if not rows or not width: + return + nblocks = triton.cdiv(width, block_size) + flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8) + start_stride = row_ks.stride(0) if row_ks is not None else 0 + _candidate_flags_kernel[(rows,)]( + candidate_blocks, + row_ks, + flags, + *candidate_blocks.stride(), + start_stride, + width, + nblocks, + block_size, + candidate_blocks.shape[1], + row_ks is not None, + row_repeat, + ) + # Derived from width, which is a tensor shape, so the grid stays static and + # a FULL cudagraph capture remains valid across replays; only the loop trip + # count inside the kernel is data-dependent. The min keeps narrow widths + # from launching programs that would only fall through. + grid_cols = min(_MASK_GRID_COLS, triton.cdiv(width, _MASK_TILE)) + _mask_candidates_strided_kernel[(rows, grid_cols)]( + logits, + row_ks, + row_ke, + flags, + *logits.stride(), + start_stride, + row_ke.stride(0), + width, + nblocks, + block_size, + row_ks is not None, + row_repeat, + _MASK_TILE, + ) + + def _max_decode_logits_rows(num_batched_tokens: int) -> int: """Upper bound on decode rows the paged-MQA logits buffer can ever hold. @@ -1102,7 +1219,6 @@ def rocm_aiter_sparse_attn_indexer( if candidate_blocks is not None: from vllm.model_executor.layers.sparse_attn_indexer import ( - _apply_candidate_mask, _select_candidate_blocks, ) @@ -1123,7 +1239,7 @@ def rocm_aiter_sparse_attn_indexer( decode_candidates, ) else: - _apply_candidate_mask( + _apply_candidate_mask_strided( logits, row_starts, visible, From d392ac836e17e87f5fec6594ef5ee5b56db72a91 Mon Sep 17 00:00:00 2001 From: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:15:28 +0200 Subject: [PATCH 204/774] [Doc][Metrics] Fix spec-decode PromQL examples to use the exposed names (#55016) Signed-off-by: Michele Campi --- vllm/v1/spec_decode/metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py index a3ccfb29e737..826dfbd16a48 100644 --- a/vllm/v1/spec_decode/metrics.py +++ b/vllm/v1/spec_decode/metrics.py @@ -187,12 +187,12 @@ class SpecDecodingProm: 1 + ( rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) / - rate(vllm:spec_decode_num_drafts[$interval])) + rate(vllm:spec_decode_num_drafts_total[$interval])) A per-position acceptance rate vector can be computed using - vllm:spec_decode_num_accepted_tokens_per_pos[$interval] / - vllm:spec_decode_num_drafts[$interval] + vllm:spec_decode_num_accepted_tokens_per_pos_total[$interval] / + vllm:spec_decode_num_drafts_total[$interval] """ _counter_cls = prometheus_client.Counter From d320a214316b52b08cb84a4745475370a1564e87 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 14:37:44 -0700 Subject: [PATCH 205/774] [Bugfix][Rust Frontend] Restore prost test dependency (#56866) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- rust/Cargo.lock | 1 + rust/src/server/Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 80edfe4ebc91..1b7dae65685b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5972,6 +5972,7 @@ dependencies = [ "libc", "llm-multimodal", "openssl", + "prost", "prost-types", "rmp-serde", "rmpv", diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index b9abe19c63b4..caccbdefc5ac 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -59,6 +59,7 @@ async-openai = { workspace = true, features = ["full"] } bytes.workspace = true clap.workspace = true expect-test.workspace = true +prost.workspace = true serial_test.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["test-util"] } From bdad63c90aec1c97682b20f282ba2a959f8229c9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 14 Sep 2026 16:40:35 -0500 Subject: [PATCH 206/774] [CI] Check target branch freshness before starting CI (#56169) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .github/workflows/run-ci-command.yml | 6 + .github/workflows/scripts/run_ci_command.py | 104 ++++++++-- .../workflows/scripts/test_run_ci_command.py | 196 ++++++++++++++++++ docs/contributing/README.md | 28 +++ 4 files changed, 322 insertions(+), 12 deletions(-) diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 747831c7d4c8..3435148ea4e3 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -21,11 +21,17 @@ jobs: (github.event.comment.body == '/ci run' || github.event.comment.body == '/ci run all' || github.event.comment.body == '/ci run nightly' || + github.event.comment.body == '/ci run --allow-stale' || + github.event.comment.body == '/ci run all --allow-stale' || + github.event.comment.body == '/ci run nightly --allow-stale' || github.event.comment.body == '/ci retry' || github.event.comment.body == '/ci cancel' || github.event.comment.body == '/amd-ci run' || github.event.comment.body == '/amd-ci run all' || github.event.comment.body == '/amd-ci run nightly' || + github.event.comment.body == '/amd-ci run --allow-stale' || + github.event.comment.body == '/amd-ci run all --allow-stale' || + github.event.comment.body == '/amd-ci run nightly --allow-stale' || github.event.comment.body == '/amd-ci retry' || github.event.comment.body == '/amd-ci cancel') runs-on: [self-hosted, linux, x64, vllm-runners] diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 6454f041ff06..530bd697d008 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import http.client import json import os import random @@ -22,6 +23,7 @@ COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly" COMMAND_RETRY_AMD_FAILED = "/amd-ci retry" COMMAND_CANCEL_AMD_CI = "/amd-ci cancel" +ALLOW_STALE_SUFFIX = " --allow-stale" RUN_CI_COMMAND_ENV = { COMMAND_RUN_CI: {}, COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"}, @@ -30,20 +32,25 @@ COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"}, COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"}, } +ALLOW_STALE_COMMANDS = frozenset( + f"{command}{ALLOW_STALE_SUFFIX}" for command in RUN_CI_COMMAND_ENV +) +RUN_CI_COMMAND_ENV.update( + { + command: RUN_CI_COMMAND_ENV[command.removesuffix(ALLOW_STALE_SUFFIX)] + for command in ALLOW_STALE_COMMANDS + } +) UPSTREAM_CI_COMMANDS = frozenset( { - COMMAND_RUN_CI, - COMMAND_RUN_CI_ALL, - COMMAND_RUN_CI_NIGHTLY, + *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/ci ")), COMMAND_RETRY_FAILED, COMMAND_CANCEL_CI, } ) AMD_CI_COMMANDS = frozenset( { - COMMAND_RUN_AMD_CI, - COMMAND_RUN_AMD_CI_ALL, - COMMAND_RUN_AMD_CI_NIGHTLY, + *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/amd-ci ")), COMMAND_RETRY_AMD_FAILED, COMMAND_CANCEL_AMD_CI, } @@ -79,6 +86,10 @@ def __init__(self, status: int | None, message: str) -> None: self.status = status +class CiPreparationError(RuntimeError): + pass + + def rate_limit_jitter() -> float: return random.uniform(1, 5) @@ -137,6 +148,8 @@ def request( None, f"API request failed: {error.reason}", ) from error + except (OSError, http.client.HTTPException) as error: + raise ApiError(None, f"API request failed: {error}") from error if not response_body: return None @@ -234,6 +247,17 @@ def _paginate(self, path: str) -> list[dict[str, Any]]: def get_pr(self, number: int) -> dict[str, Any]: return self._request(self._repo_path(f"/pulls/{number}")) + def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int: + base_ref = urllib.parse.quote(f"refs/heads/{base_ref}", safe="") + head_sha = urllib.parse.quote(head_sha, safe="") + response = self._request( + self._repo_path(f"/compare/{base_ref}...{head_sha}?per_page=1") + ) + behind = response.get("behind_by") if isinstance(response, dict) else None + if type(behind) is not int or behind < 0: + raise ApiError(None, "GitHub returned an invalid commits-behind count.") + return behind + def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]: commit = urllib.parse.quote(commit, safe="") response = self._request(self._repo_path(f"/commits/{commit}/pulls")) @@ -747,6 +771,10 @@ def notify_authorized( ( f"✅ @{author}, CI is now available for this PR.\n\n" "- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n" + "- Your branch must contain every commit currently on its upstream " + "target branch. Merge or rebase onto the latest target branch, then " + "rerun the command. Append `--allow-stale` to a run command to test " + "an outdated branch at your own risk.\n" "- `/ci retry` retries failed jobs in the CI build for the current " "PR head. If the current head has no CI build, it starts a new CI " "build for the current head containing only jobs that failed in " @@ -794,6 +822,56 @@ def find_matching_pr( return find_matching_pr(github.list_pulls_for_commit(head_sha)) +def prepare_pr_for_ci( + github: GitHubClient, + pr: Mapping[str, Any], + command: str, +) -> tuple[dict[str, Any], str]: + base_ref = pr["base"]["ref"] + no_action = "No new CI build was started." + try: + behind = github.get_commits_behind_base(base_ref, pr["head"]["sha"]) + current_pr = github.get_pr(pr["number"]) + except ApiError as error: + raise CiPreparationError( + f"Could not check the PR against upstream `{base_ref}`. " + f"{no_action} Comment `{command}` again. {error}" + ) from error + if ( + current_pr["state"] != "open" + or current_pr["head"]["sha"] != pr["head"]["sha"] + or current_pr["base"]["ref"] != pr["base"]["ref"] + ): + raise CiPreparationError( + f"The PR changed while checking its branch. {no_action} " + f"Comment `{command}` again." + ) + commits = "commit" if behind == 1 else "commits" + lag = f"This PR is {behind} {commits} behind upstream `{base_ref}`." + if command in ALLOW_STALE_COMMANDS: + warning = "" + if behind: + run_command = command.removesuffix(ALLOW_STALE_SUFFIX) + warning = ( + f"⚠️ {lag} Running CI at your own risk because " + "`--allow-stale` was requested; outdated CI configuration may " + "cause failures. Before merging, merge or rebase onto the latest " + f"`{base_ref}`, then rerun `{run_command}` on the latest PR commit." + ) + print(warning) + return current_pr, warning + + if behind: + raise CiPreparationError( + f"{lag} Your branch must contain every commit currently on upstream " + f"`{base_ref}`. {no_action} " + f"Merge or rebase onto the latest `{base_ref}`, then rerun `{command}`." + " To test this branch at your own risk, " + f"use `{command}{ALLOW_STALE_SUFFIX}`." + ) + return current_pr, "" + + def handle_run_ci( *, actor: str, @@ -828,12 +906,7 @@ def handle_run_ci( f"{ci_name} is already running for this commit: {active_build['web_url']}" ) - current_pr = github.get_pr(pr["number"]) - if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: - return ( - "The PR head changed while processing the command. " - f"Comment `{command}` again." - ) + current_pr, warning = prepare_pr_for_ci(github, pr, command) build = buildkite.create_build( create_build_payload( @@ -846,6 +919,7 @@ def handle_run_ci( return ( f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) " f"for commit `{current_pr['head']['sha'][:12]}`." + + (f"\n\n{warning}" if warning else "") ) @@ -1099,6 +1173,12 @@ def run( raise ValueError(f"Unsupported CI command: {command}") add_reaction_safely(github, comment_id, "rocket") github.add_comment(issue_number, f"✅ {message}") + except CiPreparationError as error: + add_reaction_safely(github, comment_id, "-1") + github.add_comment( + issue_number, + f"❌ {error}\n\n{command_comment_marker(comment_id)}", + ) except Exception: add_reaction_safely(github, comment_id, "confused") raise diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index ef93ea8559c8..9750cefaa394 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -22,8 +22,10 @@ COMMAND_RUN_CI_ALL, COMMAND_RUN_CI_NIGHTLY, RETRY_STATES, + RUN_CI_COMMAND_ENV, ApiError, BuildkiteClient, + GitHubClient, HttpTransport, authorize, create_build_payload, @@ -85,6 +87,7 @@ def __init__( comments: list[str] | None = None, pulls_for_commit: list[dict[str, Any]] | None = None, prs: dict[int, dict[str, Any]] | None = None, + behind: int = 0, ) -> None: self.comments = comments or [] self.permission = permission @@ -95,6 +98,12 @@ def __init__( self.reviews = reviews or [] self.pulls_for_commit = pulls_for_commit or [] self.prs = prs or {self.pr["number"]: self.pr} + self.behind = behind + self.lag_queries: list[tuple[str, str]] = [] + + def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int: + self.lag_queries.append((base_ref, head_sha)) + return self.behind def get_pr(self, number: int) -> dict[str, Any]: try: @@ -322,6 +331,11 @@ def test_only_exact_ci_commands_are_accepted(self) -> None: for command in commands: with self.subTest(command=command): self.assertEqual(parse_command(command), command) + override = f"{command} --allow-stale" + if " run" in command: + self.assertEqual(parse_command(override), override) + else: + self.assertIsNone(parse_command(override)) self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) @@ -331,6 +345,8 @@ def test_only_exact_ci_commands_are_accepted(self) -> None: self.assertIsNone(parse_command("/amd-ci retry ")) self.assertIsNone(parse_command("/AMD-CI run")) self.assertIsNone(parse_command("/amdci run")) + self.assertIsNone(parse_command("/ci run --allow-stale all")) + self.assertIsNone(parse_command("/amd-ci run --allow-stale please")) def test_commands_select_only_their_configured_pipeline(self) -> None: cases = ( @@ -530,6 +546,177 @@ def test_write_reviewer_runs_ci_without_delegation(self) -> None: self.assertTrue(github.comments[0].startswith("✅ ")) self.assertIn("Buildkite CI #123", github.comments[0]) + def test_run_variants_require_current_base_without_changing_commit(self) -> None: + for command in RUN_CI_COMMAND_ENV: + if command.endswith(" --allow-stale"): + continue + for behind in (0, 1, 5): + with self.subTest(command=command, behind=behind): + github = FakeGitHub( + behind=behind, pr=make_pr(base={"ref": "release/1.0"}) + ) + buildkite = FakeBuildkite([[], []]) + run(make_event(command), github, buildkite) + + self.assertEqual( + github.lag_queries, + [("release/1.0", github.pr["head"]["sha"])], + ) + if behind == 0: + self.assertEqual( + buildkite.created_builds[0]["commit"], + github.pr["head"]["sha"], + ) + self.assertIn( + "](https://buildkite.example/builds/123)", + github.comments[0], + ) + else: + self.assertEqual(buildkite.created_builds, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + self.assertIn(str(behind), github.comments[0]) + self.assertIn("upstream `release/1.0`", github.comments[0]) + self.assertIn("every commit", github.comments[0]) + self.assertIn( + "No new CI build was started.", github.comments[0] + ) + self.assertIn(command, github.comments[0]) + + def test_allow_stale_preserves_author_permissions_and_run_configuration( + self, + ) -> None: + cases = ( + (COMMAND_RUN_CI, "ci", {}), + (COMMAND_RUN_CI_ALL, "ci", {"RUN_ALL": "1"}), + (COMMAND_RUN_CI_NIGHTLY, "ci", {"RUN_ALL": "1", "NIGHTLY": "1"}), + (COMMAND_RUN_AMD_CI, "amd-ci", {}), + (COMMAND_RUN_AMD_CI_ALL, "amd-ci", {"RUN_ALL": "1"}), + ( + COMMAND_RUN_AMD_CI_NIGHTLY, + "amd-ci", + {"RUN_ALL": "1", "NIGHTLY": "1"}, + ), + ) + for original_command, pipeline, expected_env in cases: + command = f"{original_command} --allow-stale" + for ready, behind in ((False, 100), (True, 0), (True, 100)): + with self.subTest(command=command, ready=ready, behind=behind): + github = FakeGitHub( + permission="read", + behind=behind, + pr=make_pr( + base={"ref": "feature/stack"}, + labels=[{"name": "ready"}] if ready else [], + ), + ) + buildkite = FakeBuildkite([[], []]) + + run(make_event(command, "author"), github, buildkite) + + self.assertEqual(pipeline_for_command(command), pipeline) + if not ready: + self.assertEqual(buildkite.created_builds, []) + self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.lag_queries, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + continue + self.assertEqual(len(buildkite.created_builds), 1) + payload = buildkite.created_builds[0] + self.assertEqual(payload["commit"], github.pr["head"]["sha"]) + self.assertEqual( + payload["pull_request_base_branch"], "feature/stack" + ) + self.assertEqual(payload["message"], f"PR #42 {command} by @author") + self.assertEqual( + payload["env"], + { + "VLLM_CI_GITHUB_COMMENT_ID": "99", + "VLLM_CI_TRIGGERED_BY": "author", + **expected_env, + }, + ) + self.assertEqual( + github.lag_queries, + [("feature/stack", github.pr["head"]["sha"])], + ) + self.assertIn( + "Buildkite AMD CI #123" + if pipeline == "amd-ci" + else "Buildkite CI #123", + github.comments[0], + ) + if behind: + self.assertIn("100 commits behind", github.comments[0]) + self.assertIn("upstream `feature/stack`", github.comments[0]) + self.assertIn("Before merging", github.comments[0]) + self.assertIn("latest PR commit", github.comments[0]) + self.assertIn(original_command, github.comments[0]) + + def test_base_lag_query_qualifies_branch_refs_and_validates_count(self) -> None: + for base_ref, encoded_ref in ( + ("main", "main"), + ("release/1.0", "release%2F1.0"), + ("feature/stack", "feature%2Fstack"), + ): + for count in (0, 5, None, -1, True, "5"): + with self.subTest(base_ref=base_ref, count=count): + transport = FakeTransport( + {} if count is None else {"behind_by": count} + ) + github = GitHubClient("token", "vllm-project/vllm", transport) + if type(count) is int and count >= 0: + self.assertEqual( + github.get_commits_behind_base(base_ref, "head-sha"), count + ) + else: + with self.assertRaises(ApiError): + github.get_commits_behind_base(base_ref, "head-sha") + self.assertEqual( + transport.calls[0]["url"], + "https://api.github.com/repos/vllm-project/vllm/compare/" + f"refs%2Fheads%2F{encoded_ref}...head-sha?per_page=1", + ) + self.assertEqual(transport.calls[0]["method"], "GET") + + def test_unknown_lag_or_changed_pr_prevents_runs(self) -> None: + for command in ( + COMMAND_RUN_CI, + COMMAND_RUN_AMD_CI, + f"{COMMAND_RUN_CI} --allow-stale", + f"{COMMAND_RUN_AMD_CI} --allow-stale", + ): + for fault in ("api", "head", "closed", "base"): + with self.subTest(command=command, fault=fault): + github = FakeGitHub() + buildkite = FakeBuildkite([[], []]) + + def check( + base_ref: str, + sha: str, + fault: str = fault, + github: FakeGitHub = github, + ) -> int: + if fault == "api": + raise ApiError(503, "Unavailable") + changed = make_pr() + if fault == "head": + changed["head"]["sha"] = "another-head" + elif fault == "closed": + changed["state"] = "closed" + else: + changed["base"]["ref"] = "another-base" + github.prs[42] = changed + return 0 + + with patch.object( + github, "get_commits_behind_base", side_effect=check + ): + run(make_event(command), github, buildkite) + self.assertEqual(buildkite.created_builds, []) + self.assertEqual(buildkite.retry_calls, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + self.assertNotIn("Triggered", github.comments[0]) + def test_amd_run_ignores_blocked_builds(self) -> None: for metadata in ({}, {"github-comment-id": "98"}): with self.subTest(metadata=metadata): @@ -627,6 +814,7 @@ def test_unapproved_authors_are_denied_without_buildkite(self) -> None: run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.lag_queries, []) self.assertEqual(github.reactions, ["eyes"]) self.assertTrue(github.comments[0].startswith("❌ ")) self.assertIn("approve the PR", github.comments[0]) @@ -848,6 +1036,7 @@ def test_notification_skips_draft_prs(self) -> None: def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -867,10 +1056,12 @@ def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None: run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)]) + self.assertEqual(github.lag_queries, []) self.assertIn("Queued 3 failed job", github.comments[0]) def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -891,6 +1082,7 @@ def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None: run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite) self.assertEqual(buildkite.retry_calls, [(321, RETRY_STATES)]) + self.assertEqual(github.lag_queries, []) self.assertIn("Buildkite AMD CI #321", github.comments[0]) def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None: @@ -912,6 +1104,7 @@ def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None: def test_ci_retry_creates_filtered_build_for_new_head(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -953,6 +1146,7 @@ def test_ci_retry_creates_filtered_build_for_new_head(self) -> None: self.assertEqual(len(buildkite.created_builds), 1) payload = buildkite.created_builds[0] self.assertEqual(payload["commit"], "0123456789abcdef") + self.assertEqual(github.lag_queries, []) self.assertEqual(payload["message"], "PR #42 /ci retry by @author") self.assertEqual( json.loads(payload["env"]["VLLM_CI_ONLY_STEP_KEYS"]), @@ -1104,6 +1298,7 @@ def test_ci_cancel_cancels_active_builds_for_pr_branch(self) -> None: run(make_event(COMMAND_CANCEL_CI), github, buildkite) self.assertEqual(buildkite.cancel_calls, [123, 124, 126]) + self.assertEqual(github.lag_queries, []) self.assertEqual( [request["branch"] for request in buildkite.list_requests], ["contributor:feature", "feature"], @@ -1152,6 +1347,7 @@ def test_amd_ci_cancel_handles_command_and_fork_webhook_branches(self) -> None: run(make_event(COMMAND_CANCEL_AMD_CI), github, buildkite) self.assertEqual(buildkite.cancel_calls, [322, 321]) + self.assertEqual(github.lag_queries, []) self.assertEqual( [request["branch"] for request in buildkite.list_requests], ["contributor:feature", "feature"], diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 55966f2854fd..7a1af9c1a7b2 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -315,6 +315,34 @@ review process: the `ready` label, the PR author can use `/ci run`, `/ci retry`, `/ci cancel`, or the corresponding `/amd-ci` variants. New commits do not start upstream CI automatically. +- Before creating a Buildkite build, `/ci run` and `/amd-ci run`, including + `all` and `nightly`, check the requested PR commit against its target branch + in the upstream repository. Your branch must contain every commit currently + on that branch: **zero commits behind**. This also applies to PRs targeting + release branches or another PR's branch. +- If the PR is behind, the bot reports the count without creating a new build. + Merge or rebase onto the latest target branch, then rerun the command. If the + target branch advances in the meantime, update your branch and try again. +- To run CI on an outdated branch at your own risk, append `--allow-stale` to + any run command, for example `/ci run --allow-stale` or + `/amd-ci run all --allow-stale`. The same authorization requirements apply. + The bot reports the lag and warns that outdated CI configuration may cause + failures. Before merging, merge or rebase onto the latest target branch and + rerun CI without `--allow-stale` on the latest PR commit. +- Retrying or cancelling builds does not check branch freshness. +- These commands do not modify your branch or enforce merge requirements. + Direct Buildkite launches are outside this workflow. No additional token is + needed. + +Maintainers must configure required Buildkite status checks on the latest PR +commit and enable GitHub's +[Require branches to be up to date before merging](https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/managing-protected-branches/about-protected-branches#require-status-checks-before-merging) +protection on the relevant target branches to enforce the merge policy. Select +statuses that confirm the required tests completed, using Buildkite as their +expected source. An aggregate status that reports "passed and blocked" does not +by itself guarantee those tests ran; verify the pipeline's required and optional +blocked steps before choosing statuses or changing blocked-build reporting. +The CI command workflow does not configure these repository or pipeline settings. ### Pull Request Limits and Escalation From 00972dfd72988942138a7a6089eaee08580210b8 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:44:56 -0400 Subject: [PATCH 207/774] [ROCm][DSV4.1][Perf] Fold the mHC post step into the delayed pre projection (#56513) Signed-off-by: Fangzhou Ai Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- tests/kernels/test_mhc_kernels.py | 78 +++++++++- vllm/_aiter_ops.py | 175 ++++++++++++++++++----- vllm/model_executor/kernels/mhc/aiter.py | 21 ++- vllm/model_executor/layers/mhc.py | 82 ++++++++++- vllm/models/deepseek_v41/amd/model.py | 54 ++++--- 5 files changed, 346 insertions(+), 64 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 504981538013..f805fe39a57b 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -983,8 +983,12 @@ def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried): ) expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix) - actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix) + residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip( + *args, pre_mix=pre_mix + ) + # No post was requested, so the residual comes straight back. + assert residual_out is residual for i in (0, 1, 3): torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) # The collapse is the same FP32 multiply-and-sum in both paths. @@ -1026,7 +1030,12 @@ def spy(*spy_args, **spy_kwargs): broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1) broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x) - actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x) + # forward_hip leads with the residual; no post was requested here, so it is + # the caller's own tensor and the reference's four outputs follow it. + residual_out, *actual = op.forward_hip( + broadcast_residual, broadcast_fn, *args[2:], x=x + ) + assert residual_out is broadcast_residual assert not took_aiter, "the broadcast seam must not reach AITER" for i in range(4): torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) @@ -1268,3 +1277,68 @@ def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): assert layer.hc_attn_fn_broadcast is buffer expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) assert torch.equal(layer.hc_attn_fn_broadcast, expected) + + +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128]) +@pytest.mark.parametrize("carried", [False, True]) +def test_mhc_fused_post_pre_delayed_rocm_aiter( + num_tokens, carried, default_vllm_config +): + """Folding the post into the pre projection must match doing them apart. + + The residual is bfloat16, so the two accumulation orders are allowed to + land a rounding step apart; the tolerance here is one bfloat16 ULP. + """ + set_random_seed(0) + hc_mult, hidden_size = 4, 5120 + residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs( + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult + ) + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE) + post_layer_mix = ( + torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5 + ) + comb_res_mix = ( + torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE) + + 0.2 + ) + pre_mix = ( + torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 + if carried + else None + ) + rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6 + pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20) + + next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) + expected = ( + next_residual, + *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix), + ) + actual = MHCPreDelayedOp().forward_hip( + residual, + *pre_args, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_layer_mix, + comb_res_mix=comb_res_mix, + ) + + for got, want in zip(actual, expected, strict=True): + torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3) + + +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches(): + """AITER's heuristic hands large batches back to the unfused path.""" + from vllm._aiter_ops import rocm_aiter_ops + + assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1) + assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 454aba175e80..4cfb3b388486 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1676,6 +1676,75 @@ def _rocm_aiter_fp8_attn_fake( ) +def _mhc_delayed_pre_tail( + gemm_out: torch.Tensor, + sqrsum: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + pre_mix: torch.Tensor | None, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Finish a delayed mHC pre once AITER has produced the projection. + + Shared by the plain and the post-fused entry points, which differ only in + which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they + project. ``mhc_pre_big_fuse`` also writes a collapse against its own + pre-mix, which the delayed formulation cannot use and so discards. + Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``. + """ + from aiter.ops.mhc import mhc_pre_big_fuse + + num_tokens, hc_mult, hidden_size = residual.shape + device = residual.device + with torch.device(device): + post_mix = torch.empty( + num_tokens, hc_mult, 1, dtype=torch.float32, device=device + ) + comb_mix = torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device + ) + unused_collapse = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=device + ) + mhc_pre_big_fuse( + post_mix, + comb_mix, + unused_collapse, + gemm_out, + sqrsum, + hc_scale, + hc_base, + residual, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + + next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton( + gemm_out, + sqrsum, + hc_scale, + hc_base, + hc_mult, + hc_mult * hidden_size, + rms_eps, + hc_pre_eps, + ) + if pre_mix is None: + # Model entry selects residual stream zero. + layer_input = residual[:, 0] + else: + layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix) + return post_mix, comb_mix, layer_input, next_pre_mix + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -3611,6 +3680,10 @@ def mhc_pre_delayed( hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """mHC pre using the pre-mix carried from the previous sublayer. @@ -3625,6 +3698,15 @@ def mhc_pre_delayed( that redundant store is the price of not having a native delayed kernel, and is small next to the ~140 launches it replaces. + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` + also applies the preceding post block. AITER then projects the new + residual in the same kernel that computes it + (``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every + later stage is identical either way. That new residual is written to + ``residual_out``, which the caller must supply in that case; it is an + output buffer rather than a return value so the op never has to hand + back one of its own inputs on the unfused path. + Returns: post_mix: shape (..., hc_mult, 1), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 @@ -3632,8 +3714,9 @@ def mhc_pre_delayed( next_pre_mix: shape (..., hc_mult), dtype torch.float32 """ from aiter.ops.mhc import ( + get_mhc_fused_post_pre_config, get_mhc_pre_splitk, - mhc_pre_big_fuse, + mhc_fused_post_pre_gemm_sqrsum, mhc_pre_gemm_sqrsum, ) @@ -3666,9 +3749,16 @@ def mhc_pre_delayed( pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult) + fuse_post = sublayer_out is not None # AITER's Python wrappers allocate without explicit device arguments. with torch.device(device): - splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size) + if fuse_post: + # Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it. + splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config( + num_tokens, hidden_size + ) + else: + splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size) # AITER pads the GEMM output to a multiple of 32 columns. gemm_pad = torch.empty( splitk, @@ -3679,50 +3769,46 @@ def mhc_pre_delayed( ) gemm_out = gemm_pad[:, :, :hc_mult3] sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device) - mhc_pre_gemm_sqrsum(gemm_out, sqrsum, residual_flat, fn, tile_k, 0) - - post_mix = torch.empty( - num_tokens, hc_mult, 1, dtype=torch.float32, device=device - ) - comb_mix = torch.empty( - num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device - ) - unused_collapse = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=device - ) - mhc_pre_big_fuse( - post_mix, - comb_mix, - unused_collapse, - gemm_out, - sqrsum, - hc_scale, - hc_base, - residual_flat, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - ) + if fuse_post: + # Restated for the type checker, which cannot narrow the + # optional arguments through the fuse_post flag. + assert sublayer_out is not None + assert post_layer_mix is not None and comb_res_mix is not None + assert residual_out is not None + projected = residual_out.view(-1, hc_mult, hidden_size) + mhc_fused_post_pre_gemm_sqrsum( + gemm_out, + sqrsum, + projected, + sublayer_out.view(-1, hidden_size), + residual_flat, + # The kernel wants the post mix squeezed to (tokens, hc_mult). + post_layer_mix.view(-1, hc_mult, 1).squeeze(-1), + comb_res_mix.view(-1, hc_mult, hc_mult), + fn, + tile_m, + tile_n, + tile_k, + 0, + ) + else: + projected = residual_flat + mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0) - next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton( + post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail( gemm_out, sqrsum, + projected, hc_scale, hc_base, - hc_mult, - hc_hidden_size, + pre_mix_flat, rms_eps, hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, ) - if pre_mix_flat is None: - # Model entry selects residual stream zero. - layer_input = residual_flat[:, 0] - else: - layer_input = torch.ops.vllm.hc_collapse_triton(residual_flat, pre_mix_flat) - return ( post_mix.view(*outer_shape, hc_mult, 1), comb_mix.view(*outer_shape, hc_mult, hc_mult), @@ -3730,6 +3816,21 @@ def mhc_pre_delayed( next_pre_mix.view(*outer_shape, hc_mult), ) + @staticmethod + def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool: + """True when AITER's own heuristic favours separate post and pre. + + Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the + fused GEMM wins for decode-sized batches and loses for large ones. + """ + from aiter.jit.utils.chip_info import get_gfx_runtime + + return num_tokens >= { + "gfx950": 1024, + "gfx942": 128, + "gfx1250": 1024, + }.get(get_gfx_runtime(), 1024) + @staticmethod def hc_head( hs_flat: torch.Tensor, diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index e04ce3322bb7..d34e29314d75 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -115,6 +115,10 @@ def mhc_pre_delayed_aiter( hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """mHC pre with the pre-mix carried in from the previous sublayer. @@ -134,6 +138,13 @@ def mhc_pre_delayed_aiter( sinkhorn_repeat: number of sinkhorn iterations pre_mix: shape (..., hc_mult) from the previous sublayer, or None at model entry to select residual stream zero. + sublayer_out: attention or FFN output, shape (..., hidden_size). When + given with the two mixes below, the preceding post block is + applied here so AITER can fold it into the pre projection. + post_layer_mix: shape (..., hc_mult, 1), post gate for that block. + comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it. + residual_out: shape (..., hc_mult, hidden_size), written with the post + block's new residual. Required exactly when the post is folded in. Returns: post_mix: shape (..., hc_mult, 1), dtype torch.float32 @@ -156,6 +167,10 @@ def mhc_pre_delayed_aiter( hc_post_mult_value, sinkhorn_repeat, pre_mix, + sublayer_out, + post_layer_mix, + comb_res_mix, + residual_out, ) @@ -170,6 +185,10 @@ def _mhc_pre_delayed_aiter_fake( hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: hc_mult = residual.shape[-2] hidden_size = residual.shape[-1] @@ -313,7 +332,7 @@ def _mhc_fused_post_pre_aiter_fake( direct_register_custom_op( op_name="mhc_pre_delayed_aiter", op_func=mhc_pre_delayed_aiter, - mutates_args=[], + mutates_args=["residual_out"], fake_impl=_mhc_pre_delayed_aiter_fake, ) direct_register_custom_op( diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 14f1018597d9..73993fbc984a 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -299,7 +299,13 @@ class MHCPreDelayedOp(CustomOp): Same gates as :class:`MHCPreOp`, but the stream collapse applies the caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next - sublayer seam. Returns post_mix, comb_mix, layer_input, next_pre_mix. + sublayer seam. + + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also + applies the preceding post block, which lets AITER fold it into the pre + projection and saves a launch. Returns residual, post_mix, comb_mix, + layer_input, next_pre_mix; ``residual`` is the caller's own tensor when + no post is requested. """ # --8<-- [end:mhc_pre_delayed] @@ -307,6 +313,24 @@ class MHCPreDelayedOp(CustomOp): def enabled(cls) -> bool: return True + def __init__(self) -> None: + super().__init__() + # Built here, not lazily: CustomOp resolves dispatch against the + # current vLLM config, which is only set during model construction. + self._post = MHCPostOp() + + def _maybe_post( + self, + residual: torch.Tensor, + sublayer_out: torch.Tensor | None, + post_layer_mix: torch.Tensor | None, + comb_res_mix: torch.Tensor | None, + ) -> torch.Tensor: + if sublayer_out is None: + return residual + assert post_layer_mix is not None and comb_res_mix is not None + return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix) + def forward_cuda( self, residual: torch.Tensor, @@ -322,8 +346,14 @@ def forward_cuda( x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - return torch.ops.vllm.mhc_pre_delayed_tilelang( + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang( residual, fn, hc_scale, @@ -354,7 +384,10 @@ def forward_hip( x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual` # and folds no RMSNorm, so it cannot serve the model-entry broadcast # (which projects a narrower `x`) or a fused norm. Both are handled by @@ -364,7 +397,26 @@ def forward_hip( and norm_weight is None and _aiter_mhc_supported(residual, None, supports_norm=False) ): - return torch.ops.vllm.mhc_pre_delayed_aiter( + from vllm._aiter_ops import rocm_aiter_ops + + num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2]) + # Folding the post in only pays while the residual still fits in + # cache; past that AITER's heuristic wants the separate kernels, + # which can ask for a non-temporal store and a large-m split-k. + if sublayer_out is not None and ( + rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens) + ): + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) + sublayer_out = None + # The op writes the folded post's residual into this buffer rather + # than returning it, since on the unfused path it has none of its + # own and an op must not hand back one of its inputs. + residual_out = ( + torch.empty_like(residual) if sublayer_out is not None else None + ) + rest = torch.ops.vllm.mhc_pre_delayed_aiter( residual, fn, hc_scale, @@ -375,9 +427,17 @@ def forward_hip( hc_post_mult_value, sinkhorn_repeat, pre_mix, + sublayer_out, + post_layer_mix, + comb_res_mix, + residual_out, ) + return (residual if residual_out is None else residual_out), *rest + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) if HAS_TILELANG_MHC: - return torch.ops.vllm.mhc_pre_delayed_tilelang( + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang( residual, fn, hc_scale, @@ -392,6 +452,7 @@ def forward_hip( norm_weight, norm_eps, ) + # The post is already applied, so it must not be requested again. return self.forward_native( residual, fn, @@ -423,7 +484,13 @@ def forward_native( x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) post_mix, comb_mix, layer_input, next_pre_mix = ( mhc_kernels.mhc_pre_delayed_torch( residual, @@ -440,6 +507,7 @@ def forward_native( ) ) return ( + residual, post_mix, comb_mix, _apply_mhc_norm(layer_input, norm_weight, norm_eps), diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index eaed15214f50..b3eddc202452 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -274,7 +274,7 @@ def forward( # copies and the identity pre-mix selects copy 0. assert self.hc_attn_fn_broadcast is not None residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous() - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn_broadcast, self.hc_attn_scale, @@ -288,7 +288,7 @@ def forward( ) else: residual = x - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -301,18 +301,7 @@ def forward( pre_mix=pre_mix, ) else: - residual = self.mhc_post(x, residual, post_mix, res_mix) - if self.engram is not None and engram_hashes is not None: - # Engram injection happens between the previous sublayer's - # post and this block's pre, on the full hc stream, so the - # mix coefficients see the injected stream. - residual = self.engram( - residual, - engram_hashes[:, self.engram.layer_hash_index], - engram_mask, - ) - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( - residual, + pre_args = ( self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, @@ -321,8 +310,37 @@ def forward( self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, - pre_mix=pre_mix, ) + if self.engram is not None and engram_hashes is not None: + # Engram injection happens between the previous sublayer's + # post and this block's pre, on the full hc stream, so the + # mix coefficients see the injected stream. That read of the + # residual is what stops the post from fusing into the pre + # here, unlike the attention seam below. + residual = self.mhc_post(x, residual, post_mix, res_mix) + residual = self.engram( + residual, + engram_hashes[:, self.engram.layer_hash_index], + engram_mask, + ) + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, *pre_args, pre_mix=pre_mix + ) + else: + ( + residual, + post_mix, + res_mix, + x, + attn_pre, + ) = self.mhc_pre_delayed( + residual, + *pre_args, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, + ) x = self.attn_norm(x) if self.use_sequence_parallel: @@ -332,8 +350,7 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual = self.mhc_post(x, residual, post_mix, res_mix) - post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, self.hc_ffn_fn, self.hc_ffn_scale, @@ -344,6 +361,9 @@ def forward( self.hc_post_alpha, self.hc_sinkhorn_iters, pre_mix=attn_pre, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, ) x = self.ffn_norm(x) x = self.ffn(x, input_ids) From 8e08cef46ebeef501a1b1593fd7da77909afe4ae Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 14 Sep 2026 17:58:33 -0500 Subject: [PATCH 208/774] [Bugfix] Redact credentials from benchmark logs (#56662) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- benchmarks/multi_turn/benchmark_serving_multi_turn.py | 5 ++++- vllm/benchmarks/lib/utils.py | 11 +++++++++++ vllm/benchmarks/serve.py | 11 ++++++++--- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index 5a60d9c66883..f084a4d84c5f 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -32,8 +32,11 @@ from bench_utils import TEXT_SEPARATOR, Color, logger from transformers import AutoTokenizer # type: ignore +from vllm.benchmarks.lib.utils import redact_sensitive_namespace + NUM_TOKENS_FROM_DATASET = 0 TERM_SIGNAL = None +_SENSITIVE_ARG_FIELDS = ("api_key", "header") class ConversationSampling(str, Enum): @@ -1548,7 +1551,7 @@ async def main() -> None: args = parser.parse_args() - logger.info(args) + logger.info(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS)) logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}") logger.info(f"url={args.url}") diff --git a/vllm/benchmarks/lib/utils.py b/vllm/benchmarks/lib/utils.py index 7737621a4e7c..22409101e443 100644 --- a/vllm/benchmarks/lib/utils.py +++ b/vllm/benchmarks/lib/utils.py @@ -9,6 +9,17 @@ from typing import Any +def redact_sensitive_namespace( + args: argparse.Namespace, fields: tuple[str, ...] +) -> argparse.Namespace: + """Return a copy of CLI arguments with selected values redacted.""" + redacted_args = argparse.Namespace(**vars(args)) + for field in fields: + if getattr(redacted_args, field, None) is not None: + setattr(redacted_args, field, "***") + return redacted_args + + def extract_field( args: argparse.Namespace, extra_info: dict[str, Any], field_name: str ) -> str: diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index aef7c7e7b4d9..0292ff939a55 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -51,13 +51,18 @@ RequestFuncOutput, ) from vllm.benchmarks.lib.ready_checker import wait_for_endpoint -from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json +from vllm.benchmarks.lib.utils import ( + convert_to_pytorch_benchmark_format, + redact_sensitive_namespace, + write_to_json, +) from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.gc_utils import freeze_gc_heap from vllm.utils.network_utils import join_host_port MILLISECONDS_TO_SECONDS_CONVERSION = 1000 +_SENSITIVE_ARG_FIELDS = ("header",) def _merge_overrides(base: dict | None, override: dict | None) -> dict | None: @@ -1536,7 +1541,7 @@ def save_to_pytorch_benchmark_format( # later if needed ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"] pt_records = convert_to_pytorch_benchmark_format( - args=args, + args=redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS), metrics={k: [results[k]] for k in metrics if k in results}, extra_info={ k: results[k] @@ -2013,7 +2018,7 @@ def main(args: argparse.Namespace) -> dict[str, Any]: async def main_async(args: argparse.Namespace) -> dict[str, Any]: - print(args) + print(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS)) if args.max_concurrency is not None and args.max_concurrency <= 0: raise ValueError("--max-concurrency must be greater than 0") From 3bb0a03f35269185b6358753e64932bb946cc531 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:58:48 -0700 Subject: [PATCH 209/774] [Model Runner V2] Acceptance estimation for adaptive verification (#52228) Signed-off-by: Giancarlo Delfin --- .../spec_decode/test_acceptance_estimator.py | 294 ++++++++++ .../test_gpu_autoregressive_speculator.py | 5 + ...st_gpu_extract_hidden_states_speculator.py | 2 + vllm/config/speculative.py | 12 +- vllm/v1/worker/gpu/model_runner.py | 4 + .../gpu/spec_decode/acceptance_estimator.py | 501 ++++++++++++++++++ .../gpu/spec_decode/dflash2/speculator.py | 6 + .../gpu/spec_decode/dspark/speculator.py | 34 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 70 ++- 9 files changed, 900 insertions(+), 28 deletions(-) create mode 100644 tests/v1/spec_decode/test_acceptance_estimator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py new file mode 100644 index 000000000000..a1f3da09402f --- /dev/null +++ b/tests/v1/spec_decode/test_acceptance_estimator.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``predict`` writes one probability into two index spaces at once. + +Adaptive verification reads the estimate in this step's batch order, while the +next step's refit reads it in persistent request-state slots. A mix-up between +the two is silent -- every value stays a plausible probability, it just belongs +to another request -- so it is asserted here rather than left to acceptance +metrics to hint at. +""" + +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( + OnlineAcceptanceEstimator, +) + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for acceptance estimator tests", allow_module_level=True) + +MAX_NUM_REQS = 16 +NUM_STEPS = 4 +VOCAB_SIZE = 2048 + + +def _make_estimator(device: torch.device) -> OnlineAcceptanceEstimator: + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device) + # Distinct per-position intercepts, so reading the wrong one shows up. + estimator.slope.fill_(0.7) + estimator.intercepts[:] = torch.tensor([1.5, 1.0, 0.5, 0.0], device=device) + return estimator + + +def _expected_probs( + estimator: OnlineAcceptanceEstimator, + logits: torch.Tensor, + steps: torch.Tensor, + temperature: torch.Tensor, +) -> torch.Tensor: + """logit(q_max) by a different route than the kernel takes. + + The kernel shifts by the max and drops the unit term; here the argmax is + masked out and the complement is reduced with logsumexp, so agreement is + evidence about the value rather than about a shared formula. + """ + temp = torch.where(temperature > 0, temperature, torch.ones_like(temperature)) + scaled = logits.double() / temp[:, None].double() + top = scaled.max(dim=-1) + rest = scaled.scatter(1, top.indices[:, None], float("-inf")) + feature = (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0) + weight, bias = estimator.slope, estimator.intercepts + return torch.sigmoid(weight.double() * feature + bias[steps].double()).float() + + +@pytest.mark.parametrize("tokens_per_req", [1, NUM_STEPS]) +def test_predict_agrees_across_slot_and_batch_order(tokens_per_req: int): + """One row per request (per-step drafting) and a whole block at once. + + Which one it is is inferred from whether the draft step is per-row, so the + two cases differ only in how ``draft_step`` is shaped. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + num_reqs = 5 + + # Batch position i holds a shuffled persistent slot, as a served batch would. + slots = torch.randperm(MAX_NUM_REQS, device=device)[:num_reqs].to(torch.int32) + if tokens_per_req == 1: + idx_mapping = slots + draft_step = torch.tensor(2, device=device, dtype=torch.int32) + steps = draft_step.expand(num_reqs) + else: + idx_mapping = slots.repeat_interleave(NUM_STEPS) + draft_step = torch.arange(NUM_STEPS, device=device, dtype=torch.int32).repeat( + num_reqs + ) + steps = draft_step + batch_rows = torch.arange(num_reqs, device=device).repeat_interleave(tokens_per_req) + logits = torch.randn(idx_mapping.shape[0], VOCAB_SIZE, device=device) * 4.0 + + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + + expected = _expected_probs( + estimator, logits, steps.long(), temperature[idx_mapping.long()] + ) + torch.testing.assert_close( + estimator.predictions[idx_mapping.long(), steps.long()], + expected, + atol=2e-3, + rtol=2e-3, + ) + torch.testing.assert_close( + confidence_probs[batch_rows, steps.long()], expected, atol=2e-3, rtol=2e-3 + ) + + +def test_predict_skips_cudagraph_padded_rows(): + """Padded rows carry slot -1 and must not land on a live request.""" + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + num_reqs, num_padded = 3, 6 + + slots = torch.tensor([7, 2, 11], device=device, dtype=torch.int32) + idx_mapping = torch.cat( + [slots, torch.full((num_padded,), -1, device=device, dtype=torch.int32)] + ) + draft_step = torch.tensor(1, device=device, dtype=torch.int32) + logits = torch.randn(num_reqs + num_padded, VOCAB_SIZE, device=device) * 4.0 + + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + + live = (slots.long(), torch.full((num_reqs,), 1, device=device)) + expected = _expected_probs( + estimator, logits[:num_reqs], live[1], temperature[slots.long()] + ) + torch.testing.assert_close( + estimator.predictions[live], expected, atol=2e-3, rtol=2e-3 + ) + # Nothing outside the live (request, position) cells was touched. + untouched = torch.ones(MAX_NUM_REQS, NUM_STEPS, dtype=torch.bool, device=device) + untouched[live] = False + assert not estimator.predictions[untouched].any() + assert not confidence_probs[num_reqs:].any() + + +# A refit takes one undamped Newton step, so it needs a realistic number of +# observations per round to be well conditioned -- a handful of requests can +# send the step far past the optimum. This is the batch a served step sees. +REFIT_NUM_REQS = 512 + + +def _drive_round(estimator, features, num_accepted, slots): + """Feed one step's graded drafts through predict's buffers into `step`.""" + estimator.features[slots.long()] = features + estimator.predictions[slots.long()] = torch.sigmoid( + estimator.slope * features + estimator.intercepts + ) + estimator.step( + slots, + (num_accepted + 1).to(torch.int32), + (features.shape[1] - num_accepted).to(torch.int32), + ) + + +def test_refit_recovers_a_shared_slope_and_per_position_intercepts(): + """The fit must find the logistic that generated the observations. + + One slope is shared across draft positions and each keeps its own intercept, + so the arrowhead solve is only correct if it recovers both from data whose + positions differ in level but not in slope. + """ + device = torch.device("cuda") + torch.manual_seed(0) + true_slope = 0.6 + true_bias = torch.tensor([1.4, 0.9, 0.4, -0.1], device=device) + + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 5 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32) + for _ in range(60): + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0 + accepted = torch.rand_like(features) < torch.sigmoid( + true_slope * features + true_bias + ) + # Verification stops at the first rejection, so a position's label is + # conditional on every shallower one having been accepted. + num_accepted = (~accepted).float().argmax(dim=1) + num_accepted[accepted.all(dim=1)] = NUM_STEPS + _drive_round(estimator, features, num_accepted, slots) + + slope, bias = estimator.slope, estimator.intercepts + torch.testing.assert_close( + slope[0], torch.tensor(true_slope, device=device), atol=0.1, rtol=0 + ) + torch.testing.assert_close(bias, true_bias, atol=0.3, rtol=0) + + +def test_refit_survives_a_position_with_no_observations(): + """A draft position nobody reached must not stall the shared slope. + + Its statistics stay zero while the others accumulate, which is the degenerate + case the arrowhead denominator has to tolerate. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 5 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32) + before_slope = estimator.slope.clone() + before_intercepts = estimator.intercepts.clone() + + # Every request is rejected at position 1, so positions 2+ are never graded. + num_accepted = torch.ones(REFIT_NUM_REQS, device=device) + for _ in range(20): + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0 + _drive_round(estimator, features, num_accepted, slots) + + assert estimator.slope.isfinite().all() + assert estimator.intercepts.isfinite().all() + assert not torch.equal(estimator.slope, before_slope), "slope never moved" + # Positions 2 and 3 were never observed, so their intercepts must be untouched. + torch.testing.assert_close( + estimator.intercepts[2:], before_intercepts[2:], rtol=0, atol=0 + ) + + +def test_step_damping_scales_with_the_evidence(): + """A round moves a coefficient in proportion to the samples behind it. + + The Newton step itself barely changes with how much data a round carries -- + curvature and score both scale with it -- so scaling a round's statistics + isolates the damping factor n / (n + DAMPING_OBSERVATIONS). The slope/ + intercept coupling is zeroed so the two parameters can be checked apart. + """ + device = torch.device("cuda") + damping = OnlineAcceptanceEstimator.DAMPING_OBSERVATIONS + per_position_n = 20.0 + + def refit_with(scale): + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device) + estimator.info[:, 1] = 0.0 # sum(w*margin): no slope/intercept coupling + estimator.info[:, 2] = 1.0 * scale # sum(w) + estimator.grad[:, 1] = 0.5 * scale # sum(y - p) + # The slope's information and score are accumulated as per-position + # partials that the refit sums, so spread each total over the positions. + estimator.info[:, 0] = 30.0 * scale / NUM_STEPS # sum(w*margin^2) + estimator.grad[:, 0] = 2.0 * scale / NUM_STEPS # sum((y - p)*margin) + estimator.counts[:] = per_position_n * scale + estimator._steps_since_refit = estimator.REFIT_INTERVAL - 1 + before_slope = estimator.slope.clone() + before_intercepts = estimator.intercepts.clone() + empty = torch.zeros(0, dtype=torch.int32, device=device) + estimator.step(empty, empty, empty) + return ( + (estimator.slope - before_slope)[0].item(), + (estimator.intercepts - before_intercepts)[0].item(), + ) + + for scale in (1.0, 10.0, 100.0): + total_n = per_position_n * scale * NUM_STEPS + slope_delta, bias_delta = refit_with(scale) + assert slope_delta == pytest.approx( + (2.0 / 30.0) * total_n / (total_n + damping), rel=1e-3 + ) + assert bias_delta == pytest.approx( + 0.5 * (per_position_n * scale) / (per_position_n * scale + damping), + rel=1e-3, + ) + + # A round with no observations at all must leave the coefficients alone. + assert refit_with(0.0) == (0.0, 0.0) + + +def test_feature_ignores_which_token_was_drawn(): + """The score must depend on the draft distribution, not on the draw. + + This is what keeps rejection sampling lossless: if the score read the drawn + token, trimming a draft would condition on the token being trimmed and tilt + the emitted distribution toward the drafter's confident modes. ``predict`` + takes no token argument, and this pins that -- two rows with identical + logits must score identically no matter what was sampled from them. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + + row = torch.randn(VOCAB_SIZE, device=device) * 4.0 + logits = row.expand(2, VOCAB_SIZE).contiguous() + idx_mapping = torch.tensor([0, 1], device=device, dtype=torch.int32) + draft_step = torch.tensor(0, device=device, dtype=torch.int32) + temperature = torch.ones(MAX_NUM_REQS, device=device) + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + assert estimator.features[0, 0] == estimator.features[1, 0] + + # And the value is the top token's log-odds, not the top-2 margin. + scaled = row.double() + top = scaled.max(dim=-1) + rest = scaled.scatter(0, top.indices[None], float("-inf")) + torch.testing.assert_close( + estimator.features[0, 0].double(), + (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0), + atol=2e-3, + rtol=2e-3, + ) diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 920c9a25f062..b9b3570b8a7b 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -123,6 +123,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected) draft_model_config=draft_model_config, use_local_argmax_reduction=False, draft_sample_method="greedy", + enable_adaptive_verification=False, ) vllm_config = SimpleNamespace( speculative_config=speculative_config, @@ -161,6 +162,7 @@ def init_base(speculator, vllm_config, device): speculator.dtype = torch.float32 speculator.draft_model_config = draft_model_config speculator.supports_mm_inputs = False + speculator.use_acceptance_estimator = False checked_configs = [] @@ -195,6 +197,7 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch): speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.vllm_config = SimpleNamespace(model_config=object()) speculator.max_num_tokens = 4 speculator.hidden_size = 3 @@ -219,6 +222,7 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch): speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.vllm_config = SimpleNamespace(model_config=object()) draft_model = _TextOnlyDraftModel() speculator.test_draft_model = draft_model @@ -251,6 +255,7 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch): speculator = object.__new__(MultiModuleMTPSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.cached_draft_input_embeds = None speculator.vllm_config = SimpleNamespace(model_config=object()) speculator.max_num_tokens = 4 diff --git a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py index d1aa0798feb2..74464f542ffc 100644 --- a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py +++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py @@ -60,6 +60,7 @@ def fake_set_forward_context(*args, **kwargs): layer_name = "cache_only_layers.2" speculator = object.__new__(ExtractHiddenStatesSpeculator) speculator.vllm_config = cast(Any, SimpleNamespace()) + speculator.acceptance_estimator = None speculator.num_hidden_states = 2 speculator.hidden_states = torch.zeros(4, 2, 3) speculator.draft_attn_layer_names = {layer_name} @@ -111,6 +112,7 @@ def fake_set_forward_context(*args, **kwargs): def test_propose_requires_aux_hidden_states(): speculator = object.__new__(ExtractHiddenStatesSpeculator) + speculator.acceptance_estimator = None speculator.num_hidden_states = 2 input_batch = cast( Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32)) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 186126682093..2396b26cb01a 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1561,8 +1561,16 @@ def __post_init__(self): self.index_share_for_mtp_iteration ) - if self.method != "dspark" and self.enable_adaptive_verification: - raise ValueError("Adaptive verification only supported with DSpark") + if ( + self.method != "dspark" + and self.enable_adaptive_verification + and self.use_local_argmax_reduction + ): + raise ValueError( + "Adaptive verification estimates per-position acceptance from " + "the draft logits, which use_local_argmax_reduction never " + "materializes. Disable one of them." + ) return self diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4b1520f276b4..761893303a5c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2120,6 +2120,10 @@ def sample_tokens( if self.speculator is not None: assert self.sampler is not None + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.observe_verification( + input_batch.idx_mapping, num_sampled, num_rejected + ) # Let the target override the hidden state fed to the drafter # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py new file mode 100644 index 000000000000..46065d0546e7 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + _compute_max_and_sumexp, +) + +# The feature is a log-odds, log(q / (1 - q)), which a row carrying all of its +# mass on one token sends to infinity. Clamp it symmetrically. +_MAX_LOG_ODDS = 40.0 + +# Cold-start coefficients: the median slope and intercept fitted online across +# DeepSeek-V4-Flash, Muse-Glimmer-30B, MiMo-V2.5-Pro, Inkling and Kimi-K2.5. +# Log-loss is flat near the optimum, so one starting point serves all of them +# and a real fit replaces it within a refit or two. +_INIT_SLOPE = 0.5 +_INIT_BIAS = -0.2 + + +@triton.jit +def _accumulate_kernel( + info_ptr, + grad_ptr, + idx_mapping_ptr, + num_sampled_ptr, + num_rejected_ptr, + features_ptr, + features_stride, + pred_ptr, + pred_stride, + counts_ptr, + num_reqs, + BLOCK_R: tl.constexpr, +): + step = tl.program_id(0).to(tl.int64) + req_block = tl.arange(0, BLOCK_R) + req_mask = req_block < num_reqs + req_state_block = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + + num_sampled = tl.load(num_sampled_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + # Subtract bonus token. + num_accepted = tl.maximum(num_sampled - 1, 0) + num_rejected = tl.load(num_rejected_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + num_admitted = num_accepted + num_rejected + + feature = tl.load( + features_ptr + req_state_block * features_stride + step, + mask=req_mask, + other=0.0, + ).to(tl.float32) + pred = tl.load( + pred_ptr + req_state_block * pred_stride + step, + mask=req_mask, + other=0.0, + ).to(tl.float32) + + # The accepted tokens and the first rejected draft token (if one exists) are + # "observed" and given labels. + observed = req_mask & (step <= num_accepted) & (step < num_admitted) + label = tl.where(step < num_accepted, 1.0, 0.0) + w = tl.where(observed, pred * (1.0 - pred), 0.0) + resid = tl.where(observed, label - pred, 0.0) + + # This block's contribution to the normal equations for design row + # x = [feature, eₖ]: the XᵀWX entries aₖ, bₖ, cₖ and the Xᵀr entries g₀ₖ, g₁ₖ. + # aₖ and g₀ₖ are partials of scalars shared by every position, which + # _refit_kernel sums. + a_k = tl.sum(w * feature * feature, axis=0) + b_k = tl.sum(w * feature, axis=0) + c_k = tl.sum(w, axis=0) + g0_k = tl.sum(resid * feature, axis=0) + g1_k = tl.sum(resid, axis=0) + count = tl.sum(tl.where(observed, 1.0, 0.0), axis=0) + tl.store(info_ptr + step * 3 + 0, tl.load(info_ptr + step * 3 + 0) + a_k) + tl.store(info_ptr + step * 3 + 1, tl.load(info_ptr + step * 3 + 1) + b_k) + tl.store(info_ptr + step * 3 + 2, tl.load(info_ptr + step * 3 + 2) + c_k) + tl.store(grad_ptr + step * 2 + 0, tl.load(grad_ptr + step * 2 + 0) + g0_k) + tl.store(grad_ptr + step * 2 + 1, tl.load(grad_ptr + step * 2 + 1) + g1_k) + tl.store(counts_ptr + step, tl.load(counts_ptr + step) + count) + + +@triton.jit +def _refit_kernel( + slope_ptr, + intercepts_ptr, + info_ptr, + info_row_stride, + grad_ptr, + grad_row_stride, + counts_ptr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + L2: tl.constexpr, + DAMPING: tl.constexpr, + BLOCK: tl.constexpr, +): + k = tl.arange(0, BLOCK) + mask = k < NUM_SPECULATIVE_STEPS + + # Load the stats accumulated over this round to solve (XᵀWX + λI) Δθ = Xᵀr. + a_k = tl.load(info_ptr + k * info_row_stride + 0, mask=mask, other=0.0) + b_k = tl.load(info_ptr + k * info_row_stride + 1, mask=mask, other=0.0) + c_k = tl.load(info_ptr + k * info_row_stride + 2, mask=mask, other=0.0) + L2 + a = tl.sum(a_k, axis=0) + L2 + g0_k = tl.load(grad_ptr + k * grad_row_stride + 0, mask=mask, other=0.0) + g1_k = tl.load(grad_ptr + k * grad_row_stride + 1, mask=mask, other=0.0) + g0 = tl.sum(g0_k, axis=0) + n = tl.load(counts_ptr + k, mask=mask, other=0.0) + + # Load the current shared slope and per-position intercepts. + w = tl.load(slope_ptr) + bias = tl.load(intercepts_ptr + k, mask=mask, other=0.0) + + # Solve for the shared slope delta, Δω, in the following arrowhead system: + # ⎡ a+λ b₀ b₁ … bₙ₋₁ ⎤ ⎡ Δω ⎤ ⎡ g₀ ⎤ + # ⎢ b₀ c₀+λ ⎥ ⎢ Δβ₀ ⎥ ⎢ g₁₀ ⎥ + # ⎢ b₁ c₁+λ ⎥ ⎢ Δβ₁ ⎥ = ⎢ g₁₁ ⎥ + # ⎢ ⋮ ⋱ ⎥ ⎢ ⋮ ⎥ ⎢ ⋮ ⎥ + # ⎣ bₙ₋₁ cₙ₋₁+λ ⎦ ⎣ Δβₙ₋₁ ⎦ ⎣ g₁ₙ₋₁ ⎦ + # XᵀWX + λI Δθ Xᵀ(Y-P) + shrink = b_k * b_k / c_k + coupling = b_k * g1_k / c_k + step_w = (g0 - tl.sum(coupling, axis=0)) / (a - tl.sum(shrink, axis=0)) + + # The slope learns from every position, so it is damped by the round's + # total sample count. + total_n = tl.sum(n, axis=0) + step_w *= total_n / (total_n + DAMPING) + # Mask out NaNs and steps that would drive the slope negative. + step_w = tl.where((step_w == step_w) & (w + step_w >= 0.0), step_w, 0.0) + # Update the shared slope. + new_w = w + step_w + + # Solve for the per-position intercept deltas, Δβₖ, in the same system. + step_bias = (g1_k - b_k * step_w) / c_k + # Damp by the position's sample count. + step_bias *= n / (n + DAMPING) + # Mask out NaNs. + step_bias = tl.where(step_bias == step_bias, step_bias, 0.0) + # Update the per-position intercepts. + new_bias = bias + step_bias + + tl.store(slope_ptr, new_w) + tl.store(intercepts_ptr + k, new_bias, mask=mask) + + # Start the round clean by clearing the per-round statistics. + tl.store(info_ptr + k * info_row_stride + 0, 0.0, mask=mask) + tl.store(info_ptr + k * info_row_stride + 1, 0.0, mask=mask) + tl.store(info_ptr + k * info_row_stride + 2, 0.0, mask=mask) + tl.store(grad_ptr + k * grad_row_stride + 0, 0.0, mask=mask) + tl.store(grad_ptr + k * grad_row_stride + 1, 0.0, mask=mask) + tl.store(counts_ptr + k, 0.0, mask=mask) + + +@triton.jit +def _local_max_sumexp_kernel( + local_max_ptr, + local_max_stride, + local_sumexp_ptr, + local_sumexp_stride, + logits_ptr, + logits_stride, + idx_mapping_ptr, + idx_mapping_stride, + temperature_ptr, + num_tokens, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + if token_idx >= num_tokens: + return + + block_idx = tl.program_id(1) + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to( + tl.int64 + ) + if req_state_idx < 0: + # Cudagraph-padded requests carry -1. Write an empty partial so the + # buffer holds nothing uninitialized; _predict_kernel skips these rows. + tl.store( + local_max_ptr + token_idx * local_max_stride + block_idx, float("-inf") + ) + tl.store(local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, 0.0) + return + + # Draft logits are pre-temperature, and acceptance is decided on the scaled + # distribution the drafter actually sampled from, so scale first. A request + # with temp == 0 drafts greedily, so leave its logits alone rather than dividing + # by zero, matching what the sampler does. + temp = tl.load(temperature_ptr + req_state_idx).to(tl.float32) + temp = tl.where(temp > 0.0, temp, 1.0) + + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < vocab_size + logits = ( + tl.load( + logits_ptr + token_idx * logits_stride + offsets, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + / temp + ) + block_max, block_sumexp = _compute_max_and_sumexp(logits) + tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, block_max) + tl.store( + local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, block_sumexp + ) + + +@triton.jit +def _predict_kernel( + features_ptr, + features_stride, + pred_ptr, + pred_stride, + conf_ptr, + conf_stride, + slope_ptr, + intercepts_ptr, + local_max_ptr, + local_max_stride, + local_sumexp_ptr, + local_sumexp_stride, + idx_mapping_ptr, + idx_mapping_stride, + step_ptr, + num_tokens, + vocab_num_blocks, + per_token_step: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + MAX_LOG_ODDS: tl.constexpr, + PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + if token_idx >= num_tokens: + return + + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to( + tl.int64 + ) + if req_state_idx < 0: + # Cudagraph-padded requests carry -1. Skip them so that they don't + # scatter garbage over a live request's features. + return + + # A per-row step means the speculator sampled the whole block in one pass, so + # each request owns NUM_SPECULATIVE_STEPS consecutive rows; a step shared by + # every row means one row per request. Batch position, unlike req_state_idx, + # is only meaningful for the step that computed it, and is what adaptive + # verification indexes by. + if per_token_step: + step = tl.load(step_ptr + token_idx).to(tl.int64) + batch_idx = token_idx // NUM_SPECULATIVE_STEPS + else: + step = tl.load(step_ptr).to(tl.int64) + batch_idx = token_idx + + blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS) + blocks_mask = blocks < vocab_num_blocks + maxes = tl.load( + local_max_ptr + token_idx * local_max_stride + blocks, + mask=blocks_mask, + other=float("-inf"), + ) + sumexps = tl.load( + local_sumexp_ptr + token_idx * local_sumexp_stride + blocks, + mask=blocks_mask, + other=0.0, + ) + + # The feature, logit(q(m*)) for the row's max draft logit, m*, is: + # q(x) = e^(x - m*) / Σ_z e^(z - m*) + # logit(q(m*)) = log(q(m*) / (1 - q(m*))) + # = m* - log(Σ_{z≠m*} e^(z - m*)). + # It scores the draft distribution, not the token drawn from it, so it is + # a function of the prefix alone. Trimming a draft therefore never + # conditions on the token being trimmed, and rejection sampling stays + # lossless. Scoring the drawn token instead would select on the proposal + # and bias the emitted distribution toward the drafter's confident modes. + + # Rescale each block's partial to the global max and sum. The most likely + # token contributes exactly e^(m - m) = 1, so removing it leaves the sum + # over every other token. + global_max = tl.max(maxes, axis=0) + total_sumexp = tl.sum( + tl.where(blocks_mask, sumexps * tl.exp(maxes - global_max), 0.0), axis=0 + ) + complement_sumexp = tl.maximum(total_sumexp - 1.0, 0.0) + + # Shifted by m*, the numerator is e^(m* - m*) = 1, so the feature is just the + # negated complement log-sum-exponential. Clamp to the min/max log-odds. + feature = -tl.log(complement_sumexp) + feature = tl.minimum(tl.maximum(feature, -MAX_LOG_ODDS), MAX_LOG_ODDS) + tl.store(features_ptr + req_state_idx * features_stride + step, feature) + + # Predict the acceptance probability from the feature and current coefficients. + weight = tl.load(slope_ptr) + bias = tl.load(intercepts_ptr + step) + prob = tl.sigmoid(weight * feature + bias) + tl.store(pred_ptr + req_state_idx * pred_stride + step, prob) + tl.store(conf_ptr + batch_idx * conf_stride + step, prob) + + +class OnlineAcceptanceEstimator: + """Predicts per-position acceptance, and calibrates itself while serving. + + Lifecycle per step, driven by ``DraftModelSpeculator``: + + 1. ``step`` folds the previous step's drafts, now graded by the target, + into the IRLS accumulators, and periodically solves for new coefficients, + independently but identically on every rank. + 2. ``predict`` runs inside the captured draft graph, turning this step's + draft logits into acceptance probabilities for adaptive verification. + + Trimming starts immediately: survival is a running product of sigmoids, so it + decreases with draft position whatever the coefficients are, and an unfitted + estimator degrades to uniform-depth truncation rather than to anything harmful. + """ + + # Refit every this many steps, accumulating samples in between. + REFIT_INTERVAL = 100 + # Newton steps are damped by n / (n + DAMPING_OBSERVATIONS), so a round + # carrying this many observations moves a parameter half of a full step. It + # replaces a hard minimum-sample gate: a data-poor position keeps learning, + # just slowly, instead of freezing until it crosses a threshold. Lower values + # track a drifting workload faster, higher ones are steadier on thin data. + DAMPING_OBSERVATIONS = 50.0 + # Ridge on the Newton solve. + L2 = 1e-3 + + def __init__( + self, + max_num_reqs: int, + num_speculative_steps: int, + device: torch.device, + ): + self.num_speculative_steps = num_speculative_steps + self.device = device + self._steps_since_refit = 0 + self._refits = 0 + + # The learned coefficients used for predicting acceptance. + self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device) + self.intercepts = torch.full( + (num_speculative_steps,), _INIT_BIAS, dtype=torch.float32, device=device + ) + # Holds logit(max q), which is used as the feature for the logistic. + # Stored in stable slots keyed by persistent request-state index. + self.features = torch.zeros( + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device + ) + # Holds the acceptance predictions made at draft time, computed as: + # p = σ(ω * f + βₖ) + # where f is feature (logit(max q)) at draft position k, ω is the learned + # shared slope, and βₖ is the learned per-position intercept. Stored in + # stable slots keyed by the persistent request-state index. + self.predictions = torch.zeros( + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device + ) + + # Per-round Newton-IRLS statistics, cleared after each refit. For observation + # i the design row is xᵢ = [fᵢ, eₖ], the IRLS weight is wᵢ = pᵢ(1-pᵢ), and the + # residual is yᵢ - pᵢ, where yᵢ is the label (1 for accepted, 0 for rejected). + # XᵀWX is then the negative Hessian of the Bernoulli log-likelihood loss and + # Xᵀ(Y-P) is its gradient, so the Newton step for θ = [ω, β₀ … βₙ₋₁] is the + # the solution of (XᵀWX + λI) Δθ = Xᵀ(Y-P). + # + # Columns [aₖ, bₖ, cₖ] are each position's contribution to the XᵀWX arrowhead: + # ⎡ a b₀ b₁ … bₙ₋₁ ⎤ + # ⎢ b₀ c₀ ⎥ + # ⎢ b₁ c₁ ⎥ + # ⎢ ⋮ ⋱ ⎥ + # ⎣ bₙ₋₁ cₙ₋₁ ⎦ + # bₖ and cₖ are per-position, but a = Σₖ aₖ is one scalar shared by all of + # them. + self.info = torch.zeros( + num_speculative_steps, 3, dtype=torch.float32, device=device + ) + # Columns [g₀ₖ, g₁ₖ] are each position's contribution to the Xᵀ(Y-P) vector: + # ⎡ g₀ ⎤ + # ⎢ g₁₀ ⎥ + # ⎢ g₁₁ ⎥ + # ⎢ ⋮ ⎥ + # ⎣ g₁ₙ₋₁ ⎦ + # As above, g₁ₖ is per-position and g₀ = Σₖ g₀ₖ is shared. + self.grad = torch.zeros( + num_speculative_steps, 2, dtype=torch.float32, device=device + ) + # Number of graded drafts per-position in the current round, zeroed after + # each refit. _refit_kernel uses this to damp coefficient updates by: + # n / (n + DAMPING_OBSERVATIONS) + self.counts = torch.zeros( + num_speculative_steps, dtype=torch.float32, device=device + ) + + def step( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + # Accumulate the previous step's graded drafts into the IRLS statistics. + num_reqs = idx_mapping.shape[0] + _accumulate_kernel[(self.num_speculative_steps,)]( + self.info, + self.grad, + idx_mapping, + num_sampled, + num_rejected, + self.features, + self.features.stride(0), + self.predictions, + self.predictions.stride(0), + self.counts, + num_reqs, + BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)), + ) + + self._steps_since_refit += 1 + if self._steps_since_refit < self.REFIT_INTERVAL: + return + self._steps_since_refit = 0 + + # Fit the coefficients to the accumulated statistics gathered over the + # course of the last REFIT_INTERVAL steps. + _refit_kernel[(1,)]( + self.slope, + self.intercepts, + self.info, + self.info.stride(0), + self.grad, + self.grad.stride(0), + self.counts, + NUM_SPECULATIVE_STEPS=self.num_speculative_steps, + L2=self.L2, + DAMPING=self.DAMPING_OBSERVATIONS, + BLOCK=triton.next_power_of_2(self.num_speculative_steps), + ) + self._refits += 1 + + def predict( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + draft_step: torch.Tensor, + confidence_probs: torch.Tensor, + temperature: torch.Tensor, + ) -> None: + num_tokens, vocab_size = logits.shape + VOCAB_BLOCK_SIZE = 4096 + num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE) + local_max = torch.empty( + num_tokens, num_blocks, dtype=torch.float32, device=self.device + ) + local_sumexp = torch.empty_like(local_max) + _local_max_sumexp_kernel[(num_tokens, num_blocks)]( + local_max, + local_max.stride(0), + local_sumexp, + local_sumexp.stride(0), + logits, + logits.stride(0), + idx_mapping, + idx_mapping.stride(0), + temperature, + num_tokens, + vocab_size, + BLOCK_SIZE=VOCAB_BLOCK_SIZE, + ) + _predict_kernel[(num_tokens,)]( + self.features, + self.features.stride(0), + self.predictions, + self.predictions.stride(0), + confidence_probs, + confidence_probs.stride(0), + self.slope, + self.intercepts, + local_max, + local_max.stride(0), + local_sumexp, + local_sumexp.stride(0), + idx_mapping, + idx_mapping.stride(0), + draft_step, + num_tokens, + num_blocks, + per_token_step=draft_step.dim() > 0, + NUM_SPECULATIVE_STEPS=self.num_speculative_steps, + MAX_LOG_ODDS=_MAX_LOG_ODDS, + PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks), + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 1cd19acdab23..7db52844979e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -213,5 +213,11 @@ def _generate_draft( anchor_token_ids, ) self._sample_path(candidate_ids, scores, num_reqs) + if self.enable_adaptive_verification: + self._maybe_predict_acceptance( + self._selector_scores[:num_reqs].flatten(0, 1), + self.sample_idx_mapping[:num_sample], + self.sample_col[:num_sample], + ) if self.draft_logits is not None: self._cache_draft_logits(candidate_ids, num_sample) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 036c3014f3cb..411aa7f0415e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -79,12 +79,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.draft_model_config.hf_config, "dspark_draft_topk", None ) - self.draft_token_confidence_probs = torch.empty_like( - self.draft_tokens, dtype=torch.float32 - ) - self.enable_adaptive_verification = ( - self.speculative_config.enable_adaptive_verification - ) + self.use_confidence_head: bool = False def load_draft_model( self, @@ -108,13 +103,14 @@ def load_draft_model( dtype=self.draft_logits.dtype, device=self.device, ) - if self.enable_adaptive_verification and model.model.confidence_head is None: - raise ValueError( - "Adaptive verification needs a DSpark checkpoint with a confidence " - "head, and this one has none. Pass " - "enable_adaptive_verification=false in the speculative config to verify" - " a fixed number of drafts instead." - ) + self.use_confidence_head = ( + self.enable_adaptive_verification + and model.model.confidence_head is not None + ) + if self.use_confidence_head: + # The acceptance estimator is not needed when a trained confidence head + # is available. + self.use_acceptance_estimator = False return model def _sample_logits( @@ -124,8 +120,10 @@ def _sample_logits( sample_pos: torch.Tensor, step: int, ) -> torch.Tensor: + self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step]) if self.draft_logits is None: - return self.model.map_draft_to_target(logits.argmax(dim=-1)) + draft_ids = logits.argmax(dim=-1) + return self.model.map_draft_to_target(draft_ids) # Probabilistic sampling and rejection operate in target-vocabulary # space. A reduced draft vocabulary is scattered into its target rows. @@ -184,7 +182,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: for i in range(n_spec): # Sequential stage: Markov bias from the previously sampled token. markov_embed = self.model.markov_embed(prev) - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence_markov_embeds.append(markov_embed) bias = self.model.markov_bias(markov_embed) logits_i = base_logits[:, i] + bias @@ -194,7 +192,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence = self.model.compute_confidence( sample_hidden, torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), @@ -229,7 +227,7 @@ def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> N for i in range(n_spec): markov_embed = self.model.markov_embed(prev) - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence_markov_embeds.append(markov_embed) logits_i = self.model.apply_markov_bias_gathered( markov_embed, @@ -243,7 +241,7 @@ def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> N self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence = self.model.compute_confidence( sample_hidden, torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index f6704412191a..299cac641620 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -33,6 +33,9 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( + OnlineAcceptanceEstimator, +) from vllm.v1.worker.utils import AttentionGroup if TYPE_CHECKING: @@ -158,6 +161,15 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32) self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) + self.draft_token_confidence_probs = torch.empty_like( + self.draft_tokens, dtype=torch.float32 + ) + self.enable_adaptive_verification = ( + self.speculative_config.enable_adaptive_verification + ) + self.use_acceptance_estimator = self.enable_adaptive_verification + self.acceptance_estimator: OnlineAcceptanceEstimator | None = None + self.draft_logits: torch.Tensor | None = None if self.speculative_config.draft_sample_method == "probabilistic": # Pre-temperature logits, cached from the previous decode step. @@ -227,6 +239,20 @@ def load_model(self, target_model: nn.Module) -> None: type(self.model).__name__, ) + if self.use_acceptance_estimator: + if self.use_local_argmax_reduction: + raise ValueError( + "Adaptive verification without a confidence head estimates " + "per-position acceptance from the draft logits, which " + "use_local_argmax_reduction never materializes. Disable one " + "of them." + ) + self.acceptance_estimator = OnlineAcceptanceEstimator( + self.max_num_reqs, + self.num_speculative_steps, + self.device, + ) + def set_eplb_state(self, eplb_state: EplbState) -> None: """Inject EPLB state after construction.""" self.eplb_state = eplb_state @@ -371,12 +397,6 @@ def _validate_local_argmax_reduction(self) -> None: "(communication: O(2*tp_size) vs O(vocab_size))." ) - def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_local_argmax_reduction: - return self.model.get_top_tokens(hidden_states) - logits = self.model.compute_logits(hidden_states) - return logits.argmax(dim=-1) - def sample_draft( self, hidden_states: torch.Tensor, @@ -408,8 +428,42 @@ def sample_draft( idx_mapping, temperature, ) - return sampled - return self._greedy_sample_draft(hidden_states) + elif self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states) + else: + logits = self.model.compute_logits(hidden_states) + sampled = logits.argmax(dim=-1) + self._maybe_predict_acceptance(logits, idx_mapping, draft_step) + return sampled + + def _maybe_predict_acceptance( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + draft_step: torch.Tensor, + ) -> None: + if self.acceptance_estimator is not None: + self.acceptance_estimator.predict( + logits, + idx_mapping, + draft_step, + self.draft_token_confidence_probs, + self.temperature, + ) + + def observe_verification( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + """Fold the target's verdict on the last drafts into the estimator. + + Must run before the next `propose`, which overwrites the per-slot + features the verdict grades. + """ + if self.acceptance_estimator is not None: + self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected) def prepare_watermarking( self, contexts: torch.Tensor, watermarking: torch.Tensor From ba2ae9f23961ac67bc5c055da8c26fbd660989c6 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 14 Sep 2026 18:11:11 -0500 Subject: [PATCH 210/774] [MRV2] Run pooling post processing on non-final PP ranks (#56666) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/v1/worker/gpu_worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3f457b9b7b4a..2a3531b61b38 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1238,6 +1238,8 @@ def execute_model( ) self._pp_send_work = handles[1:] + if self.use_v2_model_runner and self.model_runner.is_pooling_model: + return self.model_runner.pool() # type: ignore return None def take_draft_token_ids(self) -> DraftTokenIds | None: From f0a61bd4386bf0d694cfe848e853f2450462e454 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Tue, 15 Sep 2026 07:48:23 +0800 Subject: [PATCH 211/774] [Rust][Benchmark] Support HF ShareGPT datasets in multi-turn mode (#51104) Signed-off-by: xiaguan <751080330@qq.com> Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/src/bench/README.md | 1 + rust/src/bench/src/config.rs | 38 +++++++++++++++++++++++++++----- rust/src/bench/src/multi_turn.rs | 26 +++++++++++++++++++--- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index d4e0c69b80b0..c8724d6cd675 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -667,6 +667,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n - `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`. - `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns. +- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`. **Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0. diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 3a0936b5f4c1..b5a24d942f50 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -549,13 +549,15 @@ impl BenchConfig { )); } - // Normalize and validate min/max turns. ShareGPT only consumes max_turns - // (the loader walks all available turns up to the cap), so the - // min/num/max coupling used for synthetic generation does not apply. - if args.dataset_name == DatasetName::ShareGpt { + // Normalize and validate min/max turns. Conversation datasets only + // consume max_turns (the loader walks all available turns up to the + // cap), so the min/num/max coupling used for synthetic generation + // does not apply. + if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) { if args.multi_turn_max_turns == 1 { return Err(BenchError::Config( - "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), + "--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn" + .into(), )); } } else { @@ -722,7 +724,7 @@ impl BenchConfig { multi_turn_min_turns, multi_turn_max_turns, sharegpt_multi_turn_max_turns: if args.multi_turn - && args.dataset_name == DatasetName::ShareGpt + && matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) && args.multi_turn_max_turns != 0 { Some(args.multi_turn_max_turns) @@ -1067,6 +1069,30 @@ mod tests { assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); } + #[test] + fn test_hf_multi_turn_uses_conversation_turn_cap() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "hf", + "--dataset-path", + "org/sharegpt-dataset", + "--hf-subset", + "sharegpt", + "--multi-turn-max-turns", + "2", + ]; + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); + + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); + } + #[test] fn test_sweep_summary_percentiles_default_empty() { let args = base_multi_turn_args(); diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index e0a31381705f..8eb175a07f9a 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -186,9 +186,29 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { - return Err(BenchError::Config( - "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), - )); + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?; + let dataset_id = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config("--dataset-path is required for --dataset-name hf".into()) + })?; + let (downloaded_path, _config, _split) = + crate::datasets::hf_dataset::download_hf_dataset( + dataset_id, + config.hf_subset.as_deref(), + config.hf_split.as_deref(), + config.num_prompts, + ) + .await?; + crate::datasets::multi_turn::load_sharegpt_multi_turn( + tok, + &downloaded_path, + config.num_prompts, + config.hf_output_len, + config.sharegpt_multi_turn_max_turns, + config.seed, + &config.request_id_prefix, + )? } }; From bbbd0a02c9d1bec965ee818af29a1319baefa69c Mon Sep 17 00:00:00 2001 From: zhaoguochun1995 <109069909+zhaoguochun1995@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:52:20 +0800 Subject: [PATCH 212/774] [Bugfix] Carry over queued work when materializing the dedicated stream (#56382) Signed-off-by: zhaoguochun1995 --- vllm/utils/torch_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b453487a01a8..6f7dc1434637 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -816,8 +816,11 @@ def current_stream() -> torch.cuda.Stream: # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77 # for more details. Therefore, we create a dedicated stream per process. if current_platform.is_rocm() or current_platform.is_cuda(): + # Ensure new stream is ordered w.r.t. replaced stream's work. + new_stream = torch.cuda.Stream() + new_stream.wait_stream(torch.cuda.current_stream()) # torch.cuda.set_stream here is the alias of _pathed_set_stream - torch.cuda.set_stream(torch.cuda.Stream()) + torch.cuda.set_stream(new_stream) elif current_platform.is_cpu(): _current_stream_tls.value = _StreamPlaceholder() else: From a7576447b86454f5c5729958d921271392ced47f Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 14 Sep 2026 19:34:21 -0500 Subject: [PATCH 213/774] [ROCm][Bugfix] Initialize Ray NIXL agents for sharded RDT (#56687) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: Codex --- .../distributed/test_sharded_rdt_producer.py | 114 ++++++++++++++++++ vllm/distributed/nixl_utils.py | 17 +++ .../weight_transfer/sharded_rdt_common.py | 10 ++ .../weight_transfer/sharded_rdt_engine.py | 5 +- .../weight_transfer/sharded_rdt_trainer.py | 5 +- 5 files changed, 144 insertions(+), 7 deletions(-) diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 62fc63876262..5cfab1957d65 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -20,17 +20,23 @@ import contextlib import gc +import os +import sys import threading import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError as FuturesTimeout +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock import pytest import torch import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod +from vllm.distributed.nixl_utils import alias_nixl_for_ray from vllm.distributed.weight_transfer.sharded_rdt_common import ( buffer_alloc_bytes, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( DEFAULT_GATHER_LOOKAHEAD, @@ -45,6 +51,114 @@ GROUP_B = ("model.layers.1.w",) +@pytest.fixture +def ray_nixl_transport(monkeypatch): + import ray + from ray.experimental.rdt import util + from ray.experimental.rdt.nixl_tensor_transport import NixlTensorTransport + + from vllm.distributed import nixl_utils + from vllm.platforms import current_platform + + transport = NixlTensorTransport() + package = ModuleType("nixl_rocm") + api = ModuleType("nixl_rocm._api") + api.__dict__.update(nixl_agent=Mock(), nixl_agent_config=Mock()) + package.__dict__["_api"] = api + modules = {package.__name__: package, api.__name__: api} + imported_rcache_limits = [] + + def import_nixl(name): + imported_rcache_limits.append(os.environ.get("UCX_RCACHE_MAX_UNRELEASED")) + monkeypatch.setitem(sys.modules, name, modules[name]) + return modules[name] + + nixl_names = ("nixl", "nixl._api", *modules) + for name in nixl_names: + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.delenv("UCX_RCACHE_MAX_UNRELEASED", raising=False) + importer = Mock(side_effect=import_nixl) + monkeypatch.setattr( + nixl_utils, "importlib", SimpleNamespace(import_module=importer) + ) + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: True) + monkeypatch.setattr(util, "transport_managers", {"NIXL": transport}) + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + monkeypatch.setattr( + ray, + "get_runtime_context", + lambda: SimpleNamespace(get_actor_id=lambda: "rdt-actor"), + ) + try: + yield transport, api, importer, imported_rcache_limits + finally: + # The helper inserts aliases directly; remove them before monkeypatch + # restores any modules and environment values present before this test. + for name in nixl_names: + sys.modules.pop(name, None) + os.environ.pop("UCX_RCACHE_MAX_UNRELEASED", None) + + +def test_ray_registers_memory_through_rocm_alias(ray_nixl_transport): + """Ray's registration must reach the ROCm agent through the alias.""" + _, api, importer, _ = ray_nixl_transport + register_nixl_memory(torch.empty(4, dtype=torch.uint8)) + api.nixl_agent.return_value.register_memory.assert_called() + assert sys.modules["nixl"] is sys.modules["nixl_rocm"] + assert sys.modules["nixl._api"] is api + + importer.reset_mock() + alias_nixl_for_ray() + importer.assert_not_called() + + +@pytest.mark.parametrize("rcache_limit", [None, "4096"]) +def test_ray_nixl_sets_ucx_limit_before_import( + monkeypatch, ray_nixl_transport, rcache_limit +): + """UCX must see the default or the user's override during the first import.""" + _, _, _, imported_rcache_limits = ray_nixl_transport + if rcache_limit is not None: + monkeypatch.setenv("UCX_RCACHE_MAX_UNRELEASED", rcache_limit) + alias_nixl_for_ray() + assert imported_rcache_limits == [rcache_limit or "1024"] * 2 + + +@pytest.mark.parametrize("existing_api", [False, True]) +def test_ray_nixl_preserves_non_rocm_or_existing_imports( + monkeypatch, ray_nixl_transport, existing_api +): + from vllm.platforms import current_platform + + _, _, importer, _ = ray_nixl_transport + if existing_api: + api = ModuleType("nixl._api") + monkeypatch.setitem(sys.modules, "nixl._api", api) + else: + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + alias_nixl_for_ray() + importer.assert_not_called() + assert "UCX_RCACHE_MAX_UNRELEASED" not in os.environ + if existing_api: + assert sys.modules["nixl._api"] is api + else: + assert "nixl._api" not in sys.modules + + +def test_ray_nixl_alias_reports_missing_rocm_package(monkeypatch, ray_nixl_transport): + from vllm.distributed import nixl_utils + + transport, api, importer, _ = ray_nixl_transport + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: False) + with pytest.raises(ImportError, match="nixl_rocm"): + register_nixl_memory(torch.empty(1)) + assert transport._nixl_agent is None + api.nixl_agent.assert_not_called() + importer.assert_not_called() + assert "nixl" not in sys.modules + assert "nixl._api" not in sys.modules + + @pytest.fixture def server_factory(monkeypatch): """Build real `_RDTProducerServer` instances without a GPU. diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py index 3879bee0ce29..c063b3ea7426 100644 --- a/vllm/distributed/nixl_utils.py +++ b/vllm/distributed/nixl_utils.py @@ -91,9 +91,26 @@ def is_nixl_available() -> bool: return pkg in sys.modules or importlib.util.find_spec(pkg) is not None +def alias_nixl_for_ray() -> None: + """Let Ray's ``nixl._api`` import resolve to the ROCm implementation.""" + if not current_platform.is_rocm() or "nixl._api" in sys.modules: + return + + package_name = _get_nixl_package_name() + if not is_nixl_available(): + raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}") + + _maybe_set_ucx_rcache_limit() + package = importlib.import_module(package_name) + api = importlib.import_module(f"{package_name}._api") + sys.modules.setdefault("nixl", package) + sys.modules.setdefault("nixl._api", api) + + __all__ = [ "NixlWrapper", "nixl_agent_config", "nixlXferTelemetry", "is_nixl_available", + "alias_nixl_for_ray", ] diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py index a5a7ed923369..70110517d3fc 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_common.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py @@ -77,6 +77,16 @@ def check_ray_rdt_version() -> None: ) +def register_nixl_memory(tensor: torch.Tensor) -> None: + """Register a tensor with Ray using the platform's NIXL package.""" + from vllm.distributed.nixl_utils import alias_nixl_for_ray + + alias_nixl_for_ray() + from ray.experimental import register_nixl_memory as ray_register_nixl_memory + + ray_register_nixl_memory(tensor) + + def assign_producer_indices( num_producers: int, num_consumers: int, consumer_idx: int ) -> list[int]: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index 2b8187007b8f..96dd53f0fdb9 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -66,6 +66,7 @@ RdtRouter, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_fake import ( BakeSink, @@ -588,8 +589,6 @@ def _preregister_at_init(self) -> None: plan = self._cached_plan if plan is None or not plan.chunks: return - from ray.experimental import register_nixl_memory - # (a) consumer receive buffers — one per ring slot, at the largest chunk. max_pack = max(c.pack_bytes for c in plan.chunks) alloc = buffer_alloc_bytes(max_pack, self._buffer_presize) @@ -740,7 +739,7 @@ def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull": any time after the metadata push. See the doc's "slot generation handshake". """ - from ray.experimental import register_nixl_memory, set_target_for_ref + from ray.experimental import set_target_for_ref if self._slot_read_done: with self._slot_cv: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 60d0e60fb02d..1be17c4fa088 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -47,6 +47,7 @@ ALLOWED_OPS, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.logger import init_logger @@ -334,8 +335,6 @@ def warmup_nixl(self) -> None: doc). The warmup buffer stays registered so the agent's CUDA-HMEM path stays initialized. """ - from ray.experimental import register_nixl_memory - self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda") with self._reg_lock: register_nixl_memory(self._nixl_warmup_buf) @@ -510,8 +509,6 @@ def _new_serve_buffer(self, nbytes: int) -> torch.Tensor: """Allocate + NIXL-register one serve slot: the single allocation seam, so registration cannot be skipped on either of the two paths that make buffers (the init-time reservation and the serve-path backstop).""" - from ray.experimental import register_nixl_memory - t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) with self._reg_lock: register_nixl_memory(t) From 772e4f05fa929491f027e61a73876a95e3d86e64 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 18:04:46 -0700 Subject: [PATCH 214/774] [Bugfix][Rust] Fix HF multi-turn dataset integration (#56915) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- rust/src/bench/src/datasets/multi_turn.rs | 57 +++++++++++++++++++++++ rust/src/bench/src/multi_turn.rs | 21 +++++---- 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 97ec522e29ad..48ff11ff7bf6 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn( .as_array() .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?; + load_sharegpt_multi_turn_from_rows( + tokenizer, + entries, + num_conversations, + output_len_override, + max_turns, + seed, + request_id_prefix, + ) +} + +/// Load multi-turn conversations from deserialized ShareGPT rows. +pub fn load_sharegpt_multi_turn_from_rows( + tokenizer: &TokenizerKind, + entries: &[serde_json::Value], + num_conversations: usize, + output_len_override: Option, + max_turns: Option, + seed: u64, + request_id_prefix: &str, +) -> Result> { // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant) let mut filtered: Vec<&serde_json::Value> = entries .iter() @@ -507,6 +528,42 @@ fn gen_prompt_to_target_len( mod tests { use super::*; + #[test] + fn test_load_sharegpt_multi_turn_from_hf_rows() { + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); + let rows = vec![serde_json::json!({ + "conversations": [ + {"from": "human", "value": "hello"}, + {"from": "gpt", "value": "hi"}, + {"from": "human", "value": "how are you?"}, + {"from": "gpt", "value": "great"} + ] + })]; + + let conversations = + load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-") + .unwrap(); + let actual = conversations + .iter() + .map(|conversation| { + ( + conversation.conversation_id.as_str(), + conversation + .turns + .iter() + .map(|turn| (turn.user_message.as_ref(), turn.expected_output_len)) + .collect::>(), + ) + }) + .collect::>(); + + assert_eq!( + actual, + vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])] + ); + } + fn common_prefix_bytes(strings: &[&str]) -> usize { if strings.is_empty() { return 0; diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 8eb175a07f9a..366fd4f6fbc1 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -192,17 +192,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Date: Mon, 14 Sep 2026 18:07:26 -0700 Subject: [PATCH 215/774] [ROCm][DSv4.1][Perf] Dequantize the MXFP8 weight once when dot_scaled cannot be used (#56560) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) --- tests/kernels/test_minimax_m3_amd_ops.py | 63 +++++++++++++++++++ .../kernels/linear/mxfp8/rocm_native.py | 18 ++++-- 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 182672597a9b..ec8f1fd4eaad 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -424,6 +424,69 @@ def test_mxfp8_linear_emulation_bf16_at_load( assert _relerr(out.float(), out_ref.float()) < 2e-2 +# --------------------------------------------------------------------------- # +# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load +# --------------------------------------------------------------------------- # +@requires_gfx950 +# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8. +@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)]) +@torch.inference_mode() +def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape): + """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it. + + ``process_weights_after_loading`` converts those to BF16 once; a weight that + never went through it has to fall back per step rather than take the + invalid path, which is what keying the dispatch on the weight dtype alone + would do. + """ + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + + N, K = shape + aligned = K % 128 == 0 + torch.manual_seed(0) + w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1 + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False) + x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5 + out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale)) + # dot_scaled re-quantizes the activation, the dequant paths do not. + tol = 5e-2 if aligned else 2e-3 + + def make_layer(): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False) + return layer + + kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + layer = make_layer() + kernel.process_weights_after_loading(layer) + assert (layer.weight.element_size() >= 2) is not aligned + assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol + + # An unaligned weight that is still FP8 must not reach dot_scaled. + unprocessed = make_layer() + unprocessed.weight_scale = torch.nn.Parameter( + w_scale[:N, : K // 32].contiguous(), requires_grad=False + ) + calls = [] + real = rocm_native._mxfp8_dot_scaled_linear + + def _spy(*a, **k): + calls.append(1) + return real(*a, **k) + + rocm_native._mxfp8_dot_scaled_linear = _spy + try: + out = kernel.apply_weights(unprocessed, x) + finally: + rocm_native._mxfp8_dot_scaled_linear = real + assert bool(calls) is aligned + assert _relerr(out.float(), out_ref.float()) < tol + + # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ──────────────── # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 9714f92dffd4..29739afae3f2 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -4,9 +4,9 @@ Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16); activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling -matrix cores. Falls back (via the kernel selector) to the BF16 -``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with -``K % 128 != 0``. +matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple +of that is dequantized to BF16 once at load and served by a plain linear +instead, since ``can_implement`` does not filter on K. """ import torch @@ -194,6 +194,9 @@ def _select_cfg(M, N, K): return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3) +_DOT_SCALED_K_ALIGN = 128 + + class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" @@ -219,6 +222,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() + if K % _DOT_SCALED_K_ALIGN != 0: + weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) layer.weight = Parameter(weight.contiguous(), requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) @@ -235,10 +240,13 @@ def apply_weights( ) out_shape = (*x.shape[:-1], layer.weight.shape[0]) x2d = x.reshape(-1, x.shape[-1]) - if x2d.shape[-1] % 128 == 0: + if layer.weight.element_size() >= 2: + out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype)) + elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0: out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale) else: - # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise. + # process_weights_after_loading dequantizes these weights once, so + # this runs only if it did not; dot_scaled would be invalid here. w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale) out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype) out = out.reshape(out_shape) From 84e5f3edac29d4b8cf7266d7882d7cb5f8dd0245 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 15 Sep 2026 09:17:41 +0800 Subject: [PATCH 216/774] [Bugfix] Validate routed-expert prompt offsets before engine submission (#56844) Signed-off-by: aoshen02 Co-authored-by: Bhushan Asati Co-authored-by: Megha Agarwal <19240983+meghaagr13@users.noreply.github.com> --- tests/test_request_input_bounds.py | 30 +++++++++++++++++++++++ tests/v1/core/test_scheduler.py | 39 ++++++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 2 +- vllm/v1/engine/input_processor.py | 22 ++++++++++------- 4 files changed, 83 insertions(+), 10 deletions(-) diff --git a/tests/test_request_input_bounds.py b/tests/test_request_input_bounds.py index 44f536e7fbe7..990f10cfebe3 100644 --- a/tests/test_request_input_bounds.py +++ b/tests/test_request_input_bounds.py @@ -12,6 +12,7 @@ import sys from collections.abc import Callable from typing import Protocol +from unittest.mock import Mock import pytest from pydantic import ValidationError @@ -27,10 +28,39 @@ from vllm.exceptions import VLLMValidationError from vllm.sampling_params import BeamSearchParams from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding +from vllm.v1.engine.input_processor import InputProcessor pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] +@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127]) +def test_routed_experts_offset_validated_before_engine_submission(offset): + """Reject invalid offsets even when routed-expert output is disabled.""" + processor = Mock(spec=InputProcessor) + processor.tokenizer = None + processor.generation_config_fields = {} + processor.renderer = Mock() + processor.renderer.get_eos_token_id.return_value = None + processor.vllm_config = Mock() + params = SamplingParams(routed_experts_prompt_start=offset) + prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]} + + if not 0 <= offset <= 3: + with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"): + InputProcessor.process_inputs( + processor, "invalid", prompt, params, ("generate",) + ) + params.routed_experts_prompt_start = 0 + + request = InputProcessor.process_inputs( + processor, "valid", prompt, params, ("generate",) + ) + assert ( + request.sampling_params.routed_experts_prompt_start + == params.routed_experts_prompt_start + ) + + # --- Stop strings: public requests cap the number of stop strings --------- diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3a972e682e4e..6ca6a8cf87aa 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -51,6 +51,7 @@ ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, + RoutedExpertsLists, SamplingMaskLists, make_empty_encoder_model_runner_output, ) @@ -131,6 +132,44 @@ def test_add_requests(): assert len(scheduler.waiting) == i + 1 +def test_routed_experts_prompt_start_at_prompt_end(): + """A prompt-end offset should return empty routing data, not crash.""" + scheduler = create_scheduler() + (request,) = create_requests(num_requests=1) + request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + num_tokens = scheduler_output.num_scheduled_tokens[request.request_id] + + manager = Mock() + manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8) + manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8) + scheduler.enable_return_routed_experts = True + scheduler.routed_experts_mgr = manager + scheduler._re_block_ids = {request.request_id: []} + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + routed_experts=RoutedExpertsLists( + routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8), + slot_mapping=np.arange(num_tokens), + ), + ), + ) + + manager.get.assert_called_once_with( + [], request.num_prompt_tokens, token_start=request.num_prompt_tokens + ) + assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1) + + def test_finish_request(): scheduler = create_scheduler() requests = create_requests(num_requests=10) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1b17eb3b1518..3af527124d73 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2101,7 +2101,7 @@ def update_from_output( prompt_start = ( request.sampling_params.routed_experts_prompt_start ) - assert prompt_start < request.num_prompt_tokens + assert prompt_start <= request.num_prompt_tokens else: prompt_start = 0 routed_experts = self.routed_experts_mgr.get( diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 4cc433cbfe38..3f86f3c8b239 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -354,12 +354,21 @@ def process_inputs( if isinstance(params, SamplingParams): # TODO: can we avoid cloning here in multiproc case? sampling_params = params.clone() + prompt_len = length_from_prompt_token_ids_or_embeds( + prompt_token_ids, prompt_embeds + ) + if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len: + raise VLLMValidationError( + f"routed_experts_prompt_start must be between 0 and " + f"the prompt length ({prompt_len}), inclusive.", + parameter="routed_experts_prompt_start", + value=sampling_params.routed_experts_prompt_start, + ) # If unset max tokens, then generate up to the max_model_len. if sampling_params.max_tokens is None: - seq_len = length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds + sampling_params.max_tokens = ( + self.model_config.max_model_len - prompt_len ) - sampling_params.max_tokens = self.model_config.max_model_len - seq_len sampling_params.update_from_generation_config( self.generation_config_fields, @@ -368,12 +377,7 @@ def process_inputs( if self.tokenizer is not None: sampling_params.update_from_tokenizer(self.tokenizer) if sampling_params.trace_decode_token_ids: - self._normalize_trace_replay_params( - sampling_params, - length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds - ), - ) + self._normalize_trace_replay_params(sampling_params, prompt_len) else: pooling_params = params.clone() From ceaa9f95ddbae0203ae444ec864bbbfb68b9b046 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 18:56:04 -0700 Subject: [PATCH 217/774] [Minor] Invert MRV1 specdec method fallback logic (#56899) Signed-off-by: Nick Hill --- vllm/config/vllm.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7d2baea26ce9..ad8f077148c3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2829,16 +2829,16 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported.append("pipeline parallelism with external_launcher") if speculative_config is not None: - # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet - if speculative_config.method in ("ngram", "ngram_gpu"): - unsupported.append("ngram/ngram_gpu speculative decoding") - elif speculative_config.method not in ( - "eagle", - "eagle3", - "mtp", - "dflash", - "dspark", - "extract_hidden_states", + if speculative_config.method in ( + # https://github.com/vllm-project/vllm/pull/40704 + "ngram", + "ngram_gpu", + # https://github.com/vllm-project/vllm/pull/43091 + "draft_model", + "suffix", + "medusa", + "mlp_speculator", + "custom_class", ): unsupported.append(f"speculative method '{speculative_config.method}'") From 31928987546ff6454f9a7d945f3ff132c5f9c531 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:57:59 -0700 Subject: [PATCH 218/774] [ROCm][CI] Prepare TheRock image for CI (#56921) Signed-off-by: Matthew Wong Co-authored-by: Andreas Karatzas --- docker/Dockerfile.rock | 9 +++++++++ docker/Dockerfile.rock_base | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index c3ce785c679d..14ba523b072f 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -875,6 +875,15 @@ RUN mkdir src && mv vllm src/vllm # Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. ENV VLLM_GPU_SYNC_CHECK=error +FROM test AS test_smoke +ARG ROCM_SMOKE_ID +RUN --network=none \ + bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \ + && printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok + +FROM scratch AS export_test_smoke +COPY --from=test_smoke /vllm-smoke-ok / + # ----------------------- # Final vLLM image FROM mori_base AS final_common diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index 1b1d08988123..80b6aaca2964 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -146,6 +146,15 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ fi; \ done \ + && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ + librocprofiler-sdk-roctx.so.1; do \ + if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ + && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ + rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ + "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + fi; \ + done \ && ldconfig # Install sccache if USE_SCCACHE is enabled (for release builds) From a08dffe21534916f218dcc93012dbaa01f942691 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Tue, 15 Sep 2026 10:23:14 +0800 Subject: [PATCH 219/774] [XPU][CI] update gpt-oss package version (#56783) Signed-off-by: zhenwei-intel --- .buildkite/intel_jobs/lm_eval_intel.yaml | 3 ++- requirements/test/xpu.txt | 19 ++----------------- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 11d3edc90e10..1ef00a47945a 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -23,5 +23,6 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + 'uv pip install "gpt-oss[eval]==0.0.5" && + cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index b7cf49649b7e..2e2fbd88ede6 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -257,7 +257,7 @@ googleapis-common-protos==1.75.1 # google-api-core # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.9 +gpt-oss==0.0.8 # via -r requirements/test/xpu.in graphql-core==3.2.11 # via hypothesis-graphql @@ -406,10 +406,6 @@ jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema -jupyter-client==8.10.0 - # via gpt-oss -jupyter-core==5.9.1 - # via jupyter-client kiwisolver==1.5.1 # via matplotlib lark==1.2.2 @@ -672,9 +668,7 @@ pillow==12.3.0 # mistral-common # torchvision platformdirs==4.11.3 - # via - # jupyter-core - # pooch + # via pooch pluggy==1.6.0 # via # pytest @@ -813,7 +807,6 @@ pytest-timeout==2.4.0 python-dateutil==2.9.0.post0 # via # botocore - # jupyter-client # matplotlib # pandas # typepy @@ -849,7 +842,6 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # jupyter-client rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in @@ -1074,8 +1066,6 @@ torchvision==0.28.0+xpu # via # open-clip-torch # timm -tornado==6.5.8 - # via jupyter-client tqdm==4.70.0 # via # -r requirements/test/../common.txt @@ -1091,10 +1081,6 @@ tqdm==4.70.0 # pqdm # sentence-transformers # transformers -traitlets==5.16.1 - # via - # jupyter-client - # jupyter-core transformers==5.16.1 # via # -c requirements/common.txt @@ -1146,7 +1132,6 @@ typing-extensions==4.16.0 # grpcio # httpx2 # huggingface-hub - # jupyter-client # lm-eval # mcp # mcp-types From 17bc3e106acbaeb552892893373659bb68d5020e Mon Sep 17 00:00:00 2001 From: liranschour Date: Tue, 15 Sep 2026 05:38:18 +0300 Subject: [PATCH 220/774] [KVConnector][P2P] Configurable unbound-store timeout and one-RTT rejection of a late fetch (#53453) Signed-off-by: Liran Schour --- docs/features/kv_offloading_usage.md | 22 ++++ .../v1/kv_offload/tiering/p2p/test_manager.py | 79 ++++++++++++- vllm/v1/kv_offload/tiering/p2p/manager.py | 106 ++++++++++++++++-- 3 files changed, 193 insertions(+), 14 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 54a54d89887d..a5c8684b4214 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -203,9 +203,31 @@ Block content hashes must match across instances for peers to exchange blocks (s | `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. | | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. | | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. | +| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. | The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process. +A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives +within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier +slots. Raise it when prefills legitimately take longer than the default: + +```bash +vllm serve \ + --kv-transfer-config '{ + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "spec_name": "TieringOffloadingSpec", + "secondary_tiers": [ + { + "type": "p2p", + "unbound_store_timeout_s": 180 + } + ] + } + }' +``` + #### Environment Variables Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence. diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index dc0e384d7c2f..aceec1cd8fa1 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager: mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._reaped_stores = {} + mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S mgr._failed_serve_ctxs = [] return mgr @@ -759,9 +761,9 @@ def test_unbound_store_kept_within_timeout(self): assert "req-fresh" in mgr._unbound_stores def test_unbound_store_reaped_after_timeout(self): - """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed - and their kv_request_id lands in _failed_req_ids so a late - FetchMsg/lookup doesn't try to satisfy them.""" + """Unbound stores past the reap deadline surface as failed and their + kv_request_id lands in _reaped_stores so a late FetchMsg is rejected + instead of parking demand nothing can satisfy.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = self._make() @@ -779,7 +781,7 @@ def test_unbound_store_reaped_after_timeout(self): # 2 baseline + 2 buffered stores assert JobResult(job_id=10, success=False) in results assert JobResult(job_id=11, success=False) in results - assert "req-stale" in mgr._failed_req_ids + assert "req-stale" in mgr._reaped_stores def test_submit_store_parks_unbound_batch(self): """submit_store on an unbound id appends a batch with a fresh @@ -1204,6 +1206,75 @@ def test_both_loads_succeed(self): assert 201 in b_ok, f"B loads succeeded: {b_ok}" assert 200 in b_ok, f"B stores succeeded: {b_ok}" + def test_late_fetch_after_reap_fails_immediately(self): + """A fetch for a reaped kv_request_id fails in one round trip. + + The producer drops parked blocks once unbound_store_timeout_s + expires. A consumer whose fetch arrives after that must learn on the + next tick via TransferDoneMsg(success=False) — not stall for + _LOAD_TIMEOUT_S and then take the abort path. + """ + from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S + + mgr_a, mgr_b = _build_paired_managers() + kv_id = "req-late" + + # Record every message type leaving the consumer, so the assertion + # below can prove no abort was ever needed. + sent_from_a: list[str] = [] + drain = mgr_a._control._drain_outbound_to + + def recording_drain(peer_local_id: str): + out = drain(peer_local_id) + sent_from_a.extend(msg.get("type") for _, msg in out) + return out + + mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign] + + # Producer parks the blocks, then reaps them before any fetch lands. + mgr_b.submit_store( + _job_metadata( + job_id=200, + keys=[b"b-block"], + chunk_ids=[0], + kv_params={"remote_decoder": {"kv_request_id": kv_id}}, + ) + ) + mgr_b._unbound_store_timeout_s = 0.0 + reap_results = list(mgr_b.get_finished_jobs()) + assert JobResult(job_id=200, success=False) in reap_results + assert kv_id in mgr_b._reaped_stores + + # Consumer now asks for the blocks that no longer exist. + consumer_params = { + "remote_prefiller": { + "kv_request_id": kv_id, + "remote_host": "B", + "remote_port": 2, + }, + } + mgr_a.on_new_request(_req_context(consumer_params)) + mgr_a.submit_load( + _job_metadata( + job_id=101, + keys=[b"b-block"], + chunk_ids=[0], + kv_params=consumer_params, + ) + ) + + started = time.monotonic() + all_a: list[JobResult] = [] + for _ in range(8): + all_a.extend(list(mgr_a.get_finished_jobs())) + list(mgr_b.get_finished_jobs()) + elapsed = time.monotonic() - started + + assert JobResult(job_id=101, success=False) in all_a, all_a + assert kv_id in mgr_a._failed_req_ids + assert "abort_fetch" not in sent_from_a, sent_from_a + assert elapsed < _LOAD_TIMEOUT_S + # --------------------------------------------------------------------------- # _accept_new_peers — duplicate connection rejection diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 2cd490171006..ba1973b76784 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -47,10 +47,18 @@ # Reap unbound store batches that have been parked without a FetchMsg # binding them to a session for longer than this. Protects against the # prefiller buffering blocks for a decoder that never asks (decoder died, -# network partition, lost kv_request_id). Must be longer than the per-store -# deadline so the store-timeout path fires first for individual jobs. +# network partition, lost kv_request_id). Applies only while a batch is +# unbound: binding hands it to the session, which times the job out under +# _STORE_TIMEOUT_S instead. +# Overridable per tier via the ``unbound_store_timeout_s`` config key. _UNBOUND_STORE_TIMEOUT_S = 60.0 +# How long a reaped kv_request_id is remembered so a fetch that arrives +# after the reap can be rejected immediately. Must exceed the consumer's +# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has +# already given up on its own, so there is nobody left to reject. +_REAPED_ID_RETENTION_S = 60.0 + # Time we wait during shutdown for inflight transfers to drain via # cancel(mode="wait") before falling back to mode="immediate". Bounded # so a wedged peer can't hang shutdown. @@ -208,6 +216,7 @@ def __init__( port: int | None = None, backends: list[str] | None = None, num_threads: int = 4, + unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S, **kwargs: Any, ) -> None: """Initialize the P2P secondary tier manager. @@ -245,9 +254,29 @@ def __init__( num_threads: NIXL agent worker threads for the UCX-only branch. Ignored when ``backends`` contains a non-UCX entry. + unbound_store_timeout_s: Seconds a producer holds stored blocks + for a consumer that has not fetched them yet, before + ``_reap_unbound_stores`` drops them. Raise it for + deployments whose prefills outlast the default; the blocks + keep primary-tier slots pinned for the whole window. **kwargs: Reserved for future tier-specific options. + + Raises: + ValueError: If ``unbound_store_timeout_s`` is not a positive + number, or anything convertible to one. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + try: + timeout_s = float(unbound_store_timeout_s) + except (TypeError, ValueError): + timeout_s = float("nan") + if not timeout_s > 0: + raise ValueError( + f"unbound_store_timeout_s must be a positive number, got " + f"{unbound_store_timeout_s!r}" + ) + self._unbound_store_timeout_s = timeout_s + # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py). # Peers whose seeds differ compute different hashes for identical # content, so lookups silently miss and no KV crosses the wire. The @@ -303,8 +332,16 @@ def __init__( # kv_request_id → list of batches submit_store'd before any peer # asked for that id. Drained into a session by _on_session_fetch # when the corresponding FetchMsg arrives, or surfaced as failures - # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S. + # by _reap_unbound_stores after _unbound_store_timeout_s. self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {} + # kv_request_id → time its parked batches were reaped. Nothing under + # one of these ids can be served any more, so a FetchMsg for it is + # rejected on the spot by _poll_once instead of parking demand until + # the consumer's own load timeout, and submit_store fails rather than + # re-parking under it — a prefill that outran the timeout would + # otherwise keep pinning fresh primary-tier slots, one timeout at a + # time. Entries are pruned after _REAPED_ID_RETENTION_S. + self._reaped_stores: dict[str, float] = {} self._finished_jobs: list[JobResult] = [] # kv_request_ids that hit a transport/session failure; On load lookup() @@ -390,7 +427,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: kv_transfer_params; if a session has bound the id, finish it. If no session has bound the id yet, this is a no-op: parked batches in `_unbound_stores` are left in place and cleaned up only by - `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`. + `_reap_unbound_stores` after `_unbound_store_timeout_s`. """ source = req_context.get_state(P2PSourceInfo) dest = req_context.get_state(P2PDestInfo) @@ -447,6 +484,22 @@ def submit_store(self, job_metadata: TransferJob) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return + # This id's earlier batches were already reaped, so the consumer has + # either given up or will be rejected by the fetch path. Parking more + # would pin primary-tier slots for another full timeout on blocks + # nobody can fetch, so fail the job now — a prefill that outran the + # timeout falls back to local recompute on the decoder either way. + if kv_request_id in self._reaped_stores: + logger.warning( + "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d " + "— failing without parking", + self._local_id, + kv_request_id, + job_id, + ) + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + return + # Fast path: a session has already received FetchMsg for this id, # so we can route the batch straight into its ServerRole. session = self._kv_to_session.get(kv_request_id) @@ -707,13 +760,27 @@ def _reap_unbound_stores(self) -> None: """Time out submit_store batches that no peer has ever fetched. Walks `_unbound_stores` for entries whose oldest batch is older - than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces - every batched job as failed, and adds the id to `_failed_req_ids` - so a late inbound FetchMsg short-circuits to a clean rejection. + than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces + every batched job as failed, and records the id in `_reaped_stores` + so a late inbound FetchMsg is rejected in one round trip and a late + `submit_store` fails instead of re-parking under the reaped id. + + Also expires `_reaped_stores` entries whose consumer can no longer + be waiting, which is why this runs before the empty check below. """ + now = time.monotonic() + if self._reaped_stores: + retention_deadline = now - _REAPED_ID_RETENTION_S + for kid in [ + kid + for kid, reaped_at in self._reaped_stores.items() + if reaped_at <= retention_deadline + ]: + del self._reaped_stores[kid] + if not self._unbound_stores: return - deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S + deadline = now - self._unbound_store_timeout_s expired: list[str] | None = None for kid, batches in self._unbound_stores.items(): # Batches are appended in arrival order, so the head is oldest. @@ -725,7 +792,7 @@ def _reap_unbound_stores(self) -> None: return for kid in expired: batches = self._unbound_stores.pop(kid) - self._failed_req_ids.add(kid) + self._reaped_stores[kid] = now for batch in batches: self._finished_jobs.append( JobResult(job_id=batch.job_id, success=False) @@ -735,7 +802,7 @@ def _reap_unbound_stores(self) -> None: "without a fetch — failing %d job(s)", self._local_id, kid, - _UNBOUND_STORE_TIMEOUT_S, + self._unbound_store_timeout_s, len(batches), ) @@ -779,6 +846,24 @@ def _poll_once(self) -> None: # inline in dispatch, so the replayed add_stored_blocks calls # match that demand and submit transfers immediately. for kv_request_id in result.new_fetch_ids: + if kv_request_id in self._reaped_stores: + # This id's blocks were already reaped, so no + # submit_store will ever satisfy the demand on_fetch + # just recorded. Finalize the round now: the peer gets + # TransferDoneMsg(success=False) on this tick and falls + # back to local prefill, instead of parking until its + # own load timeout expires. The id is deliberately left + # unbound so any later submit_store keeps taking the + # reject path. Batches re-parked between the reap and the + # retention prune are failed here rather than left to a + # second reap, which would pin their slots for another + # full timeout. + session.finish_request(kv_request_id) + for batch in self._unbound_stores.pop(kv_request_id, ()): + self._finished_jobs.append( + JobResult(job_id=batch.job_id, success=False) + ) + continue self._kv_to_session[kv_request_id] = session for batch in self._unbound_stores.pop(kv_request_id, ()): session.add_stored_blocks( @@ -810,6 +895,7 @@ def shutdown(self) -> None: JobResult(job_id=batch.job_id, success=False) ) self._unbound_stores.clear() + self._reaped_stores.clear() self._control.close() self._data.close() From 567f745c1b14ffaf09bfa71b7c4d0177a7994c04 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Mon, 14 Sep 2026 19:43:42 -0700 Subject: [PATCH 221/774] [Bugfix] Export weight-cache IPC tensors separately for each client (#56472) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../model_loader/weight_cache/daemon.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index b19fb9627237..f0af5a284a4d 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -93,11 +93,14 @@ def export_entries( ``named_parameters``/``named_buffers`` are iterated with ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing storage with ``embed_tokens.weight``) are not silently dropped. Each unique - tensor is exported once; every additional name that refers to the same + tensor is exported once per call; every additional name that refers to the same tensor object is recorded in the returned alias map so the client can re-establish the shared identity instead of allocating uninitialized memory for it. + CUDA reduction arguments must be exported separately for each consumer so + that PyTorch registers a reference for each IPC mapping's lifetime. + Returns: A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to its ``TensorEntry`` and ``aliases`` maps each duplicate name to its @@ -167,8 +170,6 @@ def __init__( self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir self.model: torch.nn.Module | None = None - self.entries: dict[str, TensorEntry] = {} - self.aliases: dict[str, str] = {} # Fingerprint before loading: process_weights_after_loading may # mutate hf_config.quantization_config. self.cache_config = WeightCacheKey.from_model_config( @@ -190,17 +191,11 @@ def load_model(self) -> None: with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(tp_size, 1) self.model = get_daemon_model(self.vllm_config) - self._export_entries() logger.info( - "Weight cache daemon rank %d cached %d tensors", + "Weight cache daemon rank %d loaded model", self.tp_rank, - len(self.entries), ) - def _export_entries(self) -> None: - assert self.model is not None - self.entries, self.aliases = export_entries(self.model) - def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None: """Serve requests until terminated. @@ -295,28 +290,28 @@ def _handle_get_state(self, conn: socket.socket, request: dict) -> None: logger.warning("WeightCacheKey mismatch on fields: %s", mismatched) send_msg(conn, {"status": "mismatch", "fields": mismatched}) return - if not self.entries: + if self.model is None: send_msg(conn, {"status": "error", "message": "Weights were released"}) return + gpu_uuid = get_current_device_uuid() + entries, aliases = export_entries(self.model) send_msg( conn, { "status": "ok", - "entries": self.entries, - "aliases": self.aliases, - "gpu_uuid": get_current_device_uuid(), + "entries": entries, + "aliases": aliases, + "gpu_uuid": gpu_uuid, }, ) logger.info_once( "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine", self.tp_rank, - len(self.entries), - len(self.aliases), + len(entries), + len(aliases), ) def _handle_release(self, conn: socket.socket) -> None: - self.entries.clear() - self.aliases.clear() self.model = None torch.accelerator.empty_cache() logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank) From e45bb43984b60fbbddd6d2ba4acd6e4a5525f49c Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Tue, 15 Sep 2026 10:48:38 +0800 Subject: [PATCH 222/774] [Bugfix][KV Offload] Ignore pending chunks in invalid sliding windows (#56486) Signed-off-by: Liqian Lin --- .../offloading_connector/test_scheduler.py | 78 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 12 ++- 2 files changed, 86 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 26664e3ac550..e7a83cb7b507 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -74,7 +74,7 @@ make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager -from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -1708,6 +1708,15 @@ def test_retry_stops_at_miss(self): class TestSlidingWindowLookup: + def test_pending_chunk_behind_gap_does_not_defer(self): + """A pending suffix too short for the window cannot improve the hit.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING} + ) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2 + ) + def test_all_hit_exact_window(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 @@ -1826,6 +1835,73 @@ def test_hit_pending_does_not_break_streak(self): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_swa_load_does_not_wait_for_unusable_pending_suffix( + request_runner, async_scheduling +): + """A CPU store beyond a missing SWA chunk must not block a ready prefix.""" + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=8, + ), + ), + ] + runner = request_runner( + block_size=4, + num_gpu_blocks=32, + async_scheduling=async_scheduling, + kv_cache_groups=groups, + ) + sched = runner.connector_scheduler + manager = CPUOffloadingManager(num_chunks=16) + sched.manager = manager + runner.new_request(token_ids=list(range(17))) + state = sched._req_status["0"] + state.update_offload_keys() + full, swa = (g.offload_keys for g in state.group_states) + ready = manager.prepare_store(full[:4] + swa[:2], state.req_context) + assert ready is not None + manager.complete_store(ready.keys_to_store, state.req_context) + pending = manager.prepare_store(swa[3:4], state.req_context) + assert pending is not None + + # The last SWA chunk is still being written; the preceding one is missing. + output = runner.scheduler.schedule() + metadata = output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + [(job_id, load)] = metadata.load_jobs.items() + assert isinstance(load.dst_spec, GPULoadStoreSpec) + assert load.dst_spec.group_sizes == [2, 2] + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + runner.scheduler.update_from_output( + output, + ModelRunnerOutput.with_kv_conn_output_only( + KVConnectorOutput( + finished_recving={"0"}, + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ), + ) + ), + ) + resumed = runner.scheduler.schedule() + assert resumed.num_scheduled_tokens == {"0": 9} + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + @pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64]) @pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3]) @pytest.mark.parametrize("is_eagle_group", [False, True]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 86e351093917..2fb3c5bb905d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -685,6 +685,7 @@ def _sliding_window_lookup( The first run may need a larger window for a partial rightmost chunk. Returns 0 on miss, None if the backend deferred a lookup.""" defer_lookup = False + pending_in_window = False consecutive_hits = 0 required_window = initial_window_size or sliding_window_size for idx in range(len(keys) - 1, -1, -1): @@ -695,7 +696,7 @@ def _sliding_window_lookup( # Block is in cache, just not readable yet — counts # as hit for the consecutive streak. Don't break: # keep scanning to let manager kick off async lookups. - defer_lookup = True + pending_in_window = True consecutive_hits += 1 case LookupResult.RETRY: # Block location uncertain — does not count as hit. @@ -703,13 +704,18 @@ def _sliding_window_lookup( # async lookups. defer_lookup = True consecutive_hits = 0 + pending_in_window = False required_window = sliding_window_size case LookupResult.MISS: consecutive_hits = 0 + # This gap rules out the incomplete window to its right. + pending_in_window = False required_window = sliding_window_size if consecutive_hits == required_window: - return idx + required_window if not defer_lookup else None - return consecutive_hits if not defer_lookup else None + return ( + None if defer_lookup or pending_in_window else idx + required_window + ) + return None if defer_lookup or pending_in_window else consecutive_hits def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip( From a529c1a748d3e3e21b3f97883b578e126892abf1 Mon Sep 17 00:00:00 2001 From: wuhangxian <1391938827@qq.com> Date: Tue, 15 Sep 2026 10:52:43 +0800 Subject: [PATCH 223/774] [Bugfix] Handle bare and malformed tool call openers in Gemma4 parser (#53444) Signed-off-by: wuhangxian <1391938827@qq.com> Co-authored-by: Chauncey --- tests/tool_parsers/test_gemma4_tool_parser.py | 108 ++++++++++++++++++ vllm/parser/gemma4.py | 8 ++ 2 files changed, 116 insertions(+) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 8d74f0431934..97ed0cd69d4c 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -897,3 +897,111 @@ def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request) } assert args_text.count("replace_all") == 1 + + +class TestBareAndMalformedOpeners: + """Tests for bare and malformed tool call openers. + + See https://github.com/vllm-project/vllm/issues/53431. + """ + + def test_bare_opener_non_streaming(self, parser, mock_request): + model_output = '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "London"} + + def test_documented_form_still_works(self, parser, mock_request): + model_output = ( + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + + def test_bare_opener_multiple_args(self, parser, mock_request): + model_output = ( + "<|tool_call>:get_weather{" + 'location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "Tokyo", "unit": "celsius"} + + def test_malformed_no_brace_bounded(self, parser, mock_request): + model_output = "<|tool_call>call:bad_func no brace" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + name = result.tool_calls[0].function.name + assert "" not in name + + def test_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>call:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + assert "London" in result.tool_calls[1].function.arguments + + def test_bare_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + + def test_three_mixed_calls(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad no brace" + '<|tool_call>call:search{input:<|"|>hello<|"|>}' + "<|tool_call>:set{flag:true,count:5}" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 3 + assert result.tool_calls[0].function.name == "bad no brace" + assert result.tool_calls[1].function.name == "search" + assert result.tool_calls[2].function.name == "set" + + def test_colon_in_content_not_broken(self, parser, mock_request): + model_output = ( + "Hello: world" + '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.content is not None + assert "Hello" in result.content + + def test_bare_opener_empty_args(self, parser, mock_request): + model_output = "<|tool_call>:set_status{}" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "set_status" diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index da1718332b21..51c46966364b 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -303,6 +303,7 @@ def gemma4_config() -> ParserEngineConfig: "TOOL_START": TOOL_CALL_START, "TOOL_END": TOOL_CALL_END, "CALL_PREFIX": "call:", + "COLON": ":", "OPEN_BRACE": "{", "OPEN_PAREN": "(", }, @@ -348,6 +349,13 @@ def gemma4_config() -> ParserEngineConfig: ParserState.TOOL_NAME, (), ), + # Bare opener: some checkpoints emit "<|tool_call>:name{...}" + # without the "call" prefix. Match the colon directly so the + # tool call is not silently dropped. + (ParserState.TOOL_PREAMBLE, "COLON"): Transition( + ParserState.TOOL_NAME, + (), + ), (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition( ParserState.TOOL_ARGS, (), From e6eb0d120c480aa5116aafc2d39f19c8ebcf3181 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 20:12:55 -0700 Subject: [PATCH 224/774] [Model][DSv4.1] Store the whole KV in MXFP8 (FlashMLA V4.1 record) (#56893) Signed-off-by: Yongye Zhu Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- cmake/external_projects/flashmla.cmake | 85 +++++--- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 94 ++++++--- csrc/libtorch_stable/ops.h | 2 +- csrc/libtorch_stable/torch_bindings.cpp | 2 +- tests/kernels/test_compressor_kv_cache.py | 130 ++++++++++++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 98 +++++++++ .../nvidia/ops/dequant_gather_k_cutedsl.py | 62 +++--- vllm/models/deepseek_v41/amd/dspark.py | 2 + vllm/models/deepseek_v41/attention.py | 43 +++- .../deepseek_v41/common/ops/cache_utils.py | 194 +++++++++++++++++- .../common/ops/fused_compress_quant_cache.py | 76 ++++++- vllm/models/deepseek_v41/nvidia/dspark.py | 2 + .../attention/backends/mla/flashmla_sparse.py | 11 + vllm/v1/attention/backends/mla/sparse_swa.py | 18 +- 14 files changed, 709 insertions(+), 110 deletions(-) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 382fd0289a2a..cfa8243f726c 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183 + GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_SOURCES ${flashmla_SOURCE_DIR}/csrc/api/api.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp # Misc kernels for decoding - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu # sm90 dense decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu # sm90 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu # sm90 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu - # sm100 dense prefill & backward - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + # sm100 dense prefill (inference-only: dense backward stays out) + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu # sm100 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu ) set(FlashMLA_Extension_SOURCES @@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/sm90 ${flashmla_SOURCE_DIR}/csrc/cutlass/include ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include ) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 06fe7108bd61..254b9783bc5a 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -23,9 +23,15 @@ * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin * second along last dim; each half is rope_dim/2 = 32 values) * - * Cache layout per paged-cache block (block_size tokens): - * [0, bs*576): token data, 448 fp8 + 128 bf16 each - * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * Cache layout per paged-cache block (block_size tokens), selected by the + * `kv_mxfp8` argument: + * V4 (584 B/token, kv_mxfp8=false) + * [0, bs*576): token data, 448 fp8 + 128 bf16 each + * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims) + * [0, bs*512): token data, 512 fp8 each + * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token */ #include "torch_utils.h" @@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64; constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad) constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 +// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims. +constexpr int kMxfp8QuantBlock = 32; +constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 +constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw // 240.0) to match the rest of vLLM's FNUZ pipeline. #if defined(USE_ROCM) && defined(__gfx942__) @@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values, // ──────────────────────────────────────────────────────────────────────────── // Small inline helpers // ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warp2MaxAbs(float val) { + // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block). + return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1)); +} + __device__ __forceinline__ float warp4MaxAbs(float val) { // Reduce absolute max across 4 consecutive lanes (lane id & 3 group). float peer = __shfl_xor_sync(FINAL_MASK, val, 1); @@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) { // slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant // + paged-cache // insert) +// +// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a +// runtime argument rather than doubling every (head count, q-norm) +// instantiation of a kernel whose hot path is the Q branch. template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, @@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { - // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert. + // ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16; + // V4.1 (kv_mxfp8) quantizes them too. int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { + int const token_data_bytes = + kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes; + int const scale_bytes = + kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken; int64_t const block_idx = slot_id / cache_block_size; int64_t const pos_in_block = slot_id % cache_block_size; uint8_t* block_base = k_cache + block_idx * static_cast(kv_block_stride); - uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes; uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = block_base + - static_cast(cache_block_size) * kTokenDataBytes + - pos_in_block * kScaleBytesPerToken; + static_cast(cache_block_size) * token_data_bytes + + pos_in_block * scale_bytes; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot( for (int i = 0; i < kElemsPerLane; i++) { local_absmax = fmaxf(local_absmax, fabsf(elements[i])); } - float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f); + float const absmax = fmaxf( + kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax), + 1e-4f); float const exponent = ceilf(log2f(absmax / kFp8Max)); float const inv_scale = exp2f(-exponent); - if (!is_rope_lane) { + if (kv_mxfp8 || !is_rope_lane) { uint8_t out_bytes[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(token_fp8_ptr + dim_base) = *reinterpret_cast(out_bytes); - if ((laneId & 3) == 0) { - int const q_block_idx = laneId >> 2; - float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); - token_scale_ptr[q_block_idx] = static_cast(encoded); - } - if (laneId == 0) { - token_scale_ptr[kNumQuantBlocks] = 0; + float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); + if (kv_mxfp8) { + // One scale per 32 dims = one per lane pair. + if ((laneId & 1) == 0) { + token_scale_ptr[laneId >> 1] = static_cast(encoded); + } + } else { + if ((laneId & 3) == 0) { + token_scale_ptr[laneId >> 2] = static_cast(encoded); + } + if (laneId == 0) { + token_scale_ptr[kNumQuantBlocks] = 0; + } } } else { uint4 out0, out1; @@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block - int const kv_block_stride) { // bytes per paged-cache block + int const kv_block_stride, // bytes per paged-cache block + bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated( uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, - int const cache_block_size, int const kv_block_stride, + int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; @@ -630,7 +663,8 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride); + num_heads_q, cache_block_size, kv_block_stride, + kv_mxfp8); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( @@ -639,7 +673,7 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) @@ -649,7 +683,7 @@ static void launchFusedDeepseekV4Templated( <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride); + cache_block_size, kv_block_stride, kv_mxfp8); #endif } @@ -663,19 +697,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, cudaStream_t stream) { + bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { #define DISPATCH(N) \ case N: \ if (apply_q_norm) { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } else { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } \ return; @@ -967,7 +1001,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm) { + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1036,7 +1070,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, stream); + apply_q_norm, kv_mxfp8, stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 68a24ca4340a..e19f49d6035f 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -264,7 +264,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 110f8595ee99..54548dfd0571 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -430,7 +430,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 0f1976cf683e..7e092229fe70 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -215,6 +215,136 @@ def run(): torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _mxfp8_record_reference( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales. + + Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the + RoPE tail first, then scale every 32-dim tile, RoPE tiles included. + """ + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + quantized, scales = _ue8m0_reference(rotated, 32, 448.0) + return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_mxfp8_record(compress_ratio: int): + """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them + in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0]. + + Byte-exact against the reference quantizer, including the rows and page + padding the kernel must leave alone. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(7) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 528 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1)) + # Leave one token unmapped so the kernel's negative-slot skip is covered. + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + expected = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + values, scales = _mxfp8_record_reference( + latent[t], cos_sin[pos // compress_ratio * compress_ratio] + ) + expected[page, row * 512 : (row + 1) * 512] = values + scale_offset = cache_block * 512 + row * 16 + expected[page, scale_offset : scale_offset + 16] = scales + + torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_mxfp8_cache_round_trip(): + """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record. + + The RoPE dims now go through fp8, so they carry quantization error too -- + the check is that every dim survives within one MXFP8 tile step, and that + the bytes on the way match the reference quantizer. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache, + quantize_and_insert_k_cache, + ) + + torch.manual_seed(11) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 528 / 512) * 512 + + k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + quantize_and_insert_k_cache( + k, cache, slot_mapping, block_size=block_size, bytes_per_token=528 + ) + + for t in (0, 1, block_size, num_tokens - 1): + values, scales = _ue8m0_reference(k[t].float(), 32, 448.0) + page, row = divmod(t, block_size) + torch.testing.assert_close( + cache[page, row * 512 : (row + 1) * 512], + values.view(torch.uint8), + rtol=0, + atol=0, + ) + scale_offset = block_size * 512 + row * 16 + torch.testing.assert_close( + cache[page, scale_offset : scale_offset + 16], + (scales.log2() + 127).to(torch.uint8), + rtol=0, + atol=0, + ) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache.view(num_blocks, block_size, 528), + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + # Half an e4m3 ULP at the top of the range (16 units) bounds the error. + tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) + tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2() + error = (out[0].float() - k.float()).abs() + assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 0e25884368c1..68eabfb94e84 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -424,6 +424,104 @@ def _dequant(k_cache_2d): ) +# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ────────────────────── + +V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32 # 512 fp8 + 16 UE8M0 scales = 528 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 64]) +@pytest.mark.parametrize("block_size", [32, 64]) +def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int): + """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3 + with one UE8M0 scale per 32 dims, no bf16 RoPE tail. + + Both sides quantize the same bf16-rounded rotated row, so the bytes must + match the Triton encoder exactly apart from the documented 1-ULP bf16 + rounding split on the rotation itself. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache, + ) + from vllm.models.deepseek_v41.common.ops import ( + quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache, + ) + + torch.manual_seed(3) + device = "cuda" + dtype = torch.bfloat16 + max_pos = 4096 + + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + k_cache_ref = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + v41_quantize_and_insert_k_cache( + apply_rope_gptj_last_k(kv, positions, cos_sin_cache), + k_cache_ref, + slot_mapping, + block_size=block_size, + use_fnuz=USE_FNUZ, + bytes_per_token=V41_HEAD_BYTES, + ) + + k_cache_fused = torch.zeros_like(k_cache_ref) + q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_dummy, + kv, + k_cache_fused, + slot_mapping, + positions, + cos_sin_cache, + 64, + 1e-6, + block_size, + True, + True, # kv_mxfp8 + ) + + # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3 + # codes, so compare after dequantizing and allow one step at the tile scale. + def _dequant(cache_2d): + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device) + v41_dequantize_and_gather_k_cache( + out, + cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES), + torch.tensor([num_tokens], dtype=torch.int32, device=device), + None, + torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0), + block_size, + offset=0, + use_fnuz=USE_FNUZ, + ) + return out[0] + + rec_ref = _dequant(k_cache_ref).float() + rec_fused = _dequant(k_cache_fused).float() + tile_scale = ( + (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX) + .log2() + .ceil() + .exp2() + ) + # One e4m3 ULP at the top of the range is 32 quantized units. + step = 32.0 * tile_scale.repeat_interleave(32, dim=-1) + assert ((rec_fused - rec_ref).abs() <= step).all() + # The NoPE dims see no rotation at all, so they must be byte-identical. + nope_bytes = slice(0, block_size * HEAD_DIM) + ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + torch.testing.assert_close( + fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 + ) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index fcb12740f27f..5920ed290cad 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -23,10 +23,11 @@ class DequantGatherKCacheKernel( VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] ): - # Hard-coded for DSv4. + # DSv4 geometry. The record format (how much of the row is quantized and + # how wide a scale group is) comes from the compile key: V4 quantizes the + # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes + # all 512 dims in groups of 32 (MXFP8) and has no tail. head_dim = 512 - fp8_dim = 448 - group_size = 64 # 1 scale per 64 elems num_warps = 4 tb_size = num_warps * 32 num_stages = 4 @@ -35,17 +36,20 @@ class DequantGatherKCacheKernel( class CompileKey: block_size: int has_gather_lens: bool + fp8_dim: int = 448 + group_size: int = 64 # dims per scale @staticmethod def kernel(compile_key: CompileKey) -> Any: head_dim = DequantGatherKCacheKernel.head_dim - fp8_dim = DequantGatherKCacheKernel.fp8_dim - group_size = DequantGatherKCacheKernel.group_size + fp8_dim = compile_key.fp8_dim + group_size = compile_key.group_size num_warps = DequantGatherKCacheKernel.num_warps tb_size = DequantGatherKCacheKernel.tb_size num_stages = DequantGatherKCacheKernel.num_stages bf16_dim = head_dim - fp8_dim data_dim = fp8_dim + bf16_dim * 2 + scale_bytes = head_dim // group_size block_size = compile_key.block_size @cute.jit @@ -62,8 +66,8 @@ def load_g2s( ): op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + cp_scale_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8 ) page_id = block_table[req_id, pos // block_size] block_offset = pos % block_size @@ -88,7 +92,7 @@ def load_g2s( ) elif idx == cutlass.const_expr(data_dim // 16): cute.copy( - cp8_atom, + cp_scale_atom, k_scale[page_id, block_offset, None], s_kscale[None, stage_id], ) @@ -119,8 +123,8 @@ def device_kernel( )[None, warp_id, None] s_kscale = smem.allocate_tensor( Uint8, - cute.make_layout((8, num_warps, num_stages)), - byte_alignment=8, + cute.make_layout((scale_bytes, num_warps, num_stages)), + byte_alignment=scale_bytes, )[None, warp_id, None] # Prepare for 16B cp.async, also for BF16 smem loads later. @@ -239,14 +243,16 @@ def device_kernel( # Last 64 elems are BF16 tail, corresponds to dequant1 of last # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= fp8_dim // 8: - idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 - cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, - ) + # 8x 16B for BF16. The V4.1 record is fully quantized and has + # no tail. + if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102 + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) # Store two 16B BF16 chunks per lane: first half, then second half. dst = out_slice[req_id, offset + i, (None, lane_id)] @@ -269,7 +275,7 @@ def host_entrypoint( ): # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( @@ -282,8 +288,8 @@ def host_entrypoint( k_scale = cute.make_tensor( k_cache.iterator + (block_size * data_dim), layout=cute.make_layout( - (k_cache.shape[0], block_size, 8), - stride=(k_cache.stride[0], 8, 1), + (k_cache.shape[0], block_size, scale_bytes), + stride=(k_cache.stride[0], scale_bytes, 1), ), ) @@ -311,10 +317,14 @@ def dispatch( # type: ignore[override] *, block_size: int, has_gather_lens: bool, + fp8_dim: int = 448, + group_size: int = 64, ) -> CompileKey: return self.CompileKey( block_size=block_size, has_gather_lens=has_gather_lens, + fp8_dim=fp8_dim, + group_size=group_size, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -345,9 +355,11 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() head_dim = self.head_dim - fp8_dim = self.fp8_dim + fp8_dim = compile_key.fp8_dim block_size = compile_key.block_size - head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 + head_bytes = ( + fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size + ) out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) k_cache = cute.runtime.make_fake_tensor( @@ -384,10 +396,14 @@ def __call__( block_table: torch.Tensor, block_size: int, offset: int, + fp8_dim: int = 448, + group_size: int = 64, ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: compile_key = self.dispatch( block_size=block_size, has_gather_lens=gather_lens is not None, + fp8_dim=fp8_dim, + group_size=group_size, ) launch_args = ( out, diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 3a6cef8aa74d..78f07b6b11db 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -259,6 +259,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 131cf4d97e7e..fa26351681c3 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -51,6 +51,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v41.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -109,6 +110,16 @@ def _fill_short_context_topk_indices( ) +# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes +# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one +# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages +# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode +# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE +# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages). +def _use_v41_mxfp8_kv_record() -> bool: + return current_platform.is_device_capability_family(100) + + def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, @@ -445,6 +456,9 @@ def __init__( self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config ) + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -454,6 +468,8 @@ def __init__( cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, + packed_bytes_per_token=self.kv_bytes_per_token, + packed_page_alignment=self.kv_page_alignment, ) # The attention layer itself was already registered with the @@ -879,6 +895,7 @@ def _fused_qnorm_rope_kv_insert( self.eps, swa_metadata.block_size, False, + self.kv_mxfp8, ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV @@ -934,9 +951,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # DeepseekV4SWACache. if not self.is_kv_source: return None - # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B - # alignment; plain bf16 / per-tensor fp8 rows use natural element-size - # pages. + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up + # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows + # use natural element-size pages. uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, @@ -945,12 +962,13 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, tokens_per_state=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else 512, + alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token; - # head_size stays semantic (512). - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + # Packed record width; head_size stays semantic (512). + state_content_bytes=( + self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + ), ) def _compressed_kv_cache(self) -> torch.Tensor: @@ -998,8 +1016,15 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: head_size=self.head_dim, dtype=self.dtype, tokens_per_state=self.compress_ratio, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + # The indexer page shares a physical block with the main KV cache + # (the block stride is the sum of every page in the group), so it + # has to carry the same padding: the packed record's TMA stride, + # else 512B for FlashInfer sparse (#44577). + alignment=( + 576 + if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() + else 512 + ), ) def forward(self): ... diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 64462eb73b4a..51743016bf07 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,6 +36,18 @@ from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import next_power_of_2 +# Per-token byte width of the two paged fp8 records. Both put a page's whole +# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the +# width the KV-cache spec publishes -- identifies the record: +# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 +# dims each plus a pad byte. +# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 +# scales of 32 dims each. FlashMLA's ``ModelType::V41``. +V4_BYTES_PER_TOKEN = 584 +V41_BYTES_PER_TOKEN = 528 +V41_QUANT_BLOCK = 32 +V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 + @triton.jit def quantize_and_insert_k_kernel( @@ -163,6 +175,67 @@ def quantize_and_insert_k_kernel( tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) +@triton.jit +def _quantize_and_insert_k_mxfp8_kernel( + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + num_tokens, + head_dim: tl.constexpr, # 512, RoPE dims included + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 dims per UE8M0 scale + cache_block_size: tl.constexpr, # paged cache block size, in tokens + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache. + + One program per token. Unlike the V4 record there is no bf16 tail, so every + dim is scaled. + """ + pid = tl.program_id(0) + if pid >= num_tokens: + return + + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + d = tl.arange(0, head_dim) + x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) + + tiles = tl.reshape(x, (scale_dim, quant_block)) + # Floor the absmax like the CUDA encoder so an all-zero tile still lands on + # a finite scale. + amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max))) + scaled = tl.clamp( + tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max + ) + if use_fnuz: # noqa: SIM108 + fp8 = scaled.to(tl.float8e4b8) + else: + fp8 = scaled.to(tl.float8e4nv) + tl.store( + token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,)) + ) + + # UE8M0 encoding: stored_value = exponent + 127 (bias). + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0) + tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8)) + + def quantize_and_insert_k_cache( k: torch.Tensor, # [num_tokens, 512] bf16 k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 @@ -170,16 +243,13 @@ def quantize_and_insert_k_cache( block_size: int = 64, is_ue8m0: bool = True, use_fnuz: bool = False, + bytes_per_token: int = V4_BYTES_PER_TOKEN, ): """ Quantize K tensor and insert into paged K cache. - K Cache block layout (block_size=64 tokens): - - First 64 * 576 = 36864 bytes: Token data - - Each token: 448 bytes (fp8) + 128 bytes (bf16) - - Next 64 * 8 = 512 bytes: Scales - - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) - - Padded to multiple of 576 + ``bytes_per_token`` picks the record (see the module header): the V4 one, + or V4.1's all-dims MXFP8 one. ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, @@ -196,6 +266,26 @@ def quantize_and_insert_k_cache( num_tokens = slot_mapping.shape[0] block_stride = k_cache.stride(0) # bytes per block + if bytes_per_token == V41_BYTES_PER_TOKEN: + if use_fnuz and not current_platform.is_fp8_fnuz(): + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") + _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)]( + k, + slot_mapping, + k_cache, + num_tokens, + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=block_stride, + fp8_max=get_fp8_min_max()[1] + if use_fnuz + else torch.finfo(torch.float8_e4m3fn).max, + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -343,6 +433,69 @@ def _dequantize_and_gather_k_kernel( tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) +@triton.jit +def _dequantize_and_gather_k_mxfp8_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + d = tl.arange(0, head_dim) + s = tl.arange(0, scale_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31. + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + x_uint8 = tl.load(token_data_ptr + d) + if use_fnuz: + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) + + # UE8M0: scale = 2^(stored_value - 127). + encoded = tl.load(token_scale_ptr + s) + scale = tl.exp2(encoded.to(tl.float32) - 127.0) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -358,6 +511,28 @@ def dequantize_and_gather_k_cache_triton( offset: int, use_fnuz: bool = False, ) -> None: + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: + _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -365,8 +540,6 @@ def dequantize_and_gather_k_cache_triton( FP8_MAX = 448.0 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - num_reqs = seq_lens.shape[0] - NUM_WORKERS = 128 _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( out, out.stride(0), @@ -408,6 +581,8 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. + The record is read off ``k_cache.shape[-1]``; see the module header. + ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder @@ -419,6 +594,7 @@ def dequantize_and_gather_k_cache( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) + mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( out=out, k_cache=k_cache, @@ -427,6 +603,8 @@ def dequantize_and_gather_k_cache( block_table=block_table, block_size=block_size, offset=offset, + fp8_dim=512 if mxfp8 else 448, + group_size=V41_QUANT_BLOCK if mxfp8 else 64, ) return diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 5950c4b6ba5c..96f15849ddd3 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -232,10 +232,13 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight - segregated UE8M0 scale bytes per token, including one zero padding scale); - ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows - read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. + layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the + per-token byte width names -- 584 B for V4 (576 value bytes and eight + segregated UE8M0 scale bytes, including one zero padding scale) or 528 B + for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 + scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain + [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the + per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -246,8 +249,15 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] == 584 - _rope_quant_insert_kernel[(num_tokens,)]( + assert kv_cache.shape[-1] in (584, 528), ( + f"unsupported paged KV record width {kv_cache.shape[-1]}" + ) + kernel = ( + _rope_quant_insert_mxfp8_kernel + if kv_cache.shape[-1] == 528 + else _rope_quant_insert_kernel + ) + kernel[(num_tokens,)]( latent, positions, cos_sin_cache, @@ -338,6 +348,60 @@ def _rope_quant_insert_kernel( tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448) +@triton.jit +def _rope_quant_insert_mxfp8_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims. + + The RoPE dims are quantized here too, so unlike the V4 kernel the rotation + has to happen before the scales are picked. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + values = page + (slot % CACHE_BLOCK) * 512 + scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16 + + quant = tl.reshape(rotated, (16, 32)) + amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0))) + scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1)) + fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv) + tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))) + + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded) + tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 3c711681aa75..88646404f44e 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -264,6 +264,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d8c4368401c6..9c440776ae15 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -91,6 +91,17 @@ The first `ue8m0` is the scale for the first 64 `float8_e4m3` values, the second for the next 64, and so on. +For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: +- **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE + dims are quantized too, so there is no `bfloat16` part. +- **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 + consecutive `float8_e4m3` values (i.e. MXFP8). + +The V4 and V4.1 records are not laid out token-by-token within a page: a page +holds all its data rows first and all its scale rows after, so a page is +`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA +stride (576 B for V4, 512 B for V4.1). + In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: - **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble = diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 62c3921b5e12..c4f3cd8b4bdd 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -79,6 +79,8 @@ def __init__( cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, block_size: int = 64, + packed_bytes_per_token: int = 584, + packed_page_alignment: int = 576, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend @@ -88,6 +90,8 @@ def __init__( self.prefix = prefix self.cache_config = cache_config self.dtype = dtype + self.packed_bytes_per_token = packed_bytes_per_token + self.packed_page_alignment = packed_page_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -105,8 +109,9 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: self.kv_cache = kv_cache.squeeze(1) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous - # bf16/fp8 cache uses the natural element-size page. + # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode + # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural + # element-size page. uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, @@ -115,10 +120,11 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + state_content_bytes=( + self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None + ), + # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) From 381c61e008682945d43d6c8808ed633477364b4b Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Mon, 14 Sep 2026 20:15:19 -0700 Subject: [PATCH 225/774] [Bugfix][NIXL] Report a full prefix cache hit as finished (#56640) Signed-off-by: Yizheng Jiao Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- .../kv_connector/unit/test_nixl_connector.py | 72 +++++++++++++++++++ .../unit/test_nixl_connector_hma.py | 43 ++++++++++- .../kv_connector/v1/nixl/base_scheduler.py | 10 ++- .../kv_connector/v1/nixl/metadata.py | 7 ++ .../kv_connector/v1/nixl/pull_scheduler.py | 5 +- .../kv_connector/v1/nixl/pull_worker.py | 7 ++ .../kv_connector/v1/nixl/push_scheduler.py | 4 +- 7 files changed, 142 insertions(+), 6 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 9b54f626f265..298886d16095 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -341,6 +341,10 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): assert req_meta.remote.request_id == f"prefill-{42}" # do_remote_prefill is consumed by request_finished to prevent re-issuing. assert request.kv_transfer_params["do_remote_prefill"] is False + # The scheduler is not waiting on this recv -- the request is already gone + # from self.requests, so reporting it would trip `assert req_id in + # self.requests` in _update_from_kv_xfer_finished. + assert req_meta.awaiting_kvs is False @patch( @@ -2695,6 +2699,74 @@ def check_xfer_state(self, handle: int) -> str: return super().check_xfer_state(handle) +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FailingNixlWrapper, +) +@pytest.mark.parametrize("awaiting_kvs", [True, False]) +@pytest.mark.parametrize("notification_fails", [False, True]) +def test_empty_recv_is_reported_only_when_awaited( + default_vllm_config, dist_init, awaiting_kvs, notification_fails +): + """An empty recv is reported iff the scheduler parked the request on it. + + _read_blocks returns early when the local block list is empty -- D already + holds the KV and only the producer needs telling. Two kinds of caller reach + that path, and they need opposite handling: + + Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the + request is named in neither _recving_transfers nor _failed_recv_reqs and + never reaches finished_recving. The scheduler has no other way to release a + WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in + skipped_waiting holding its blocks for the life of the process. + + Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to + free P's blocks for a request aborted before it was scheduled, or a readback + on a request that stays RUNNING. Reporting either one crashes the engine + core -- _update_from_kv_xfer_finished asserts the id is still in + self.requests and that the request is parked or finished. + + A failed notification changes neither: the KV is local either way, and the + producer frees its own blocks on a timeout. + """ + vllm_config = create_vllm_config() + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, + connector.engine_id, + hand_shake_latency=0.0, + kv_cache_config=connector._kv_cache_config, + ) + connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails + + request_id = f"test_empty_recv_awaited_{awaiting_kvs}" + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id=request_id, + local_block_ids=(), # empty: the whole prompt is already cached locally + kv_transfer_params={ + "remote_block_ids": [[20, 21, 22]], + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": f"prefill-{request_id}", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + }, + awaiting_kvs=awaiting_kvs, + ) + connector.bind_connector_metadata(metadata) + dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + connector.start_load_kv(dummy_ctx) + connector.bind_connector_metadata(NixlConnectorMetadata()) + time.sleep(0.1) + connector.start_load_kv(dummy_ctx) + + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert (request_id in done_recving) is awaiting_kvs + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 259e9536c1ca..60eb2d6614a9 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -149,8 +149,49 @@ def uncached(block_id): scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2) - _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id] + _, _, local_num_computed_blocks, awaiting_kvs = scheduler._reqs_need_recv[ + request.request_id + ] assert local_num_computed_blocks == (2, 1) + # There are external tokens to pull, so the scheduler parks the request and + # is waiting to see it in finished_recving. + assert awaiting_kvs is True + + +@pytest.mark.cpu_test +def test_full_local_hit_is_not_awaited_by_the_scheduler(): + """A recv with num_external_tokens == 0 must not be flagged as awaited. + + The request keeps running -- there is nothing to pull, only a notification + owed to the producer -- so the worker must not report it in + finished_recving. Doing so trips `assert RequestStatus.is_finished` in + _update_from_kv_xfer_finished, which only tolerates a finished recv for a + request that is parked or already done.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_utils import KVCacheBlock + + scheduler = object.__new__(NixlPullConnectorScheduler) + scheduler._reqs_in_batch = set() + scheduler._reqs_need_save = {} + scheduler._reqs_need_recv = {} + scheduler.use_host_buffer = False + scheduler.is_bidirectional_kv_xfer_enabled = False + scheduler._is_hma_required = False + scheduler.kv_cache_config = MagicMock( + select_transfer_block_ids=lambda block_ids: block_ids + ) + + blocks = KVCacheBlocks(blocks=([KVCacheBlock(block_id=0, _block_hash=object())],)) + request = create_request(do_remote_prefill=True) + + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0) + + _, local_block_ids, _, awaiting_kvs = scheduler._reqs_need_recv[request.request_id] + assert not local_block_ids # full local hit: nothing to pull + assert awaiting_kvs is False @pytest.mark.cpu_test diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 5c01d996968e..9cfd100730a5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -113,7 +113,7 @@ def __init__( # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. self._reqs_need_recv: dict[ - ReqId, tuple[Request, BlockIds, tuple[int, ...]] + ReqId, tuple[Request, BlockIds, tuple[int, ...], bool] ] = {} self._reqs_need_save: dict[ReqId, Request] = {} # Reqs to send and their expiration time @@ -459,13 +459,19 @@ def build_connector_meta( meta = NixlConnectorMetadata() # Loop through scheduled reqs and convert to ReqMeta. - for req_id, (req, block_ids, cached) in self._reqs_need_recv.items(): + for req_id, ( + req, + block_ids, + cached, + awaiting_kvs, + ) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req_to_recv( request_id=req_id, local_block_ids=block_ids, kv_transfer_params=req.kv_transfer_params, local_num_computed_blocks=cached, + awaiting_kvs=awaiting_kvs, ) if self.use_host_buffer: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index f61fdfa83512..be32a6969d66 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -246,6 +246,9 @@ class ReqMeta: remote_block_size: int | None = None # Remote producer pipeline-parallel size (push mode, D side). pp_size: int = 1 + # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS + # and expects it in finished_recving; notify-only recvs must not be reported. + awaiting_kvs: bool = False class NixlConnectorMetadata(KVConnectorMetadata): @@ -275,6 +278,7 @@ def _add_new_req( local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ) -> ReqMeta: return ReqMeta( local_block_ids=local_block_ids, @@ -285,6 +289,7 @@ def _add_new_req( remote_block_size=kv_transfer_params.get("remote_block_size"), pp_size=kv_transfer_params.get("pp_size", 1), local_num_computed_blocks=local_num_computed_blocks, + awaiting_kvs=awaiting_kvs, ) def add_new_req_to_save( @@ -303,11 +308,13 @@ def add_new_req_to_recv( local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ): req = self._add_new_req( local_block_ids, kv_transfer_params, local_num_computed_blocks, + awaiting_kvs, ) req.remote = RemoteMeta( block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 446126c328ea..9b9012a532ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -172,6 +172,9 @@ def update_state_after_alloc( request, local_block_ids, local_num_computed_blocks, + # Parked in WAITING_FOR_REMOTE_KVS only when there is + # something to pull; a full local hit stays RUNNING. + num_external_tokens > 0, ) else: @@ -224,7 +227,7 @@ def request_finished( # To avoid stranding the prefill blocks in the prefill instance, # we must add empty block_ids to _reqs_need_recv so that our # worker side will notify and free blocks in the prefill instance. - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 59205b943a5e..f787639cfe77 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -303,6 +303,7 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]: local_dram_handle=local_dram_handle, remote_xfer_side_handle=remote_xfer_side_handle, expected_consumers=plan.local_consumers, + awaiting_kvs=meta.awaiting_kvs, ): return @@ -327,6 +328,7 @@ def _read_blocks( local_dram_handle: int | None, remote_xfer_side_handle: int, expected_consumers: int, + awaiting_kvs: bool, ) -> bool: """ Post a READ point-to-point xfer request from a single local worker to @@ -387,6 +389,11 @@ def _read_blocks( remote_agent_name=agent_name, ) self.xfer_stats.record_failed_notification() + # Report even on notification failure: the KV is already local, and + # an unreported parked request would hold its blocks forever. + # Notify-only recvs must stay unreported (scheduler asserts). + if awaiting_kvs: + self._recving_transfers.setdefault(request_id, []) return True if read_spec.block_ids_by_region: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index c11a63bb3414..a007e0c5a3ae 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -198,7 +198,7 @@ def update_state_after_alloc( # ReqMeta without a KeyError — the actual remote block IDs are # learned by P over the NIXL handshake at WRITE time. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, local_block_ids, ()) + self._reqs_need_recv[request.request_id] = (request, local_block_ids, (), False) # Mark as processed so a re-entry (e.g. preemption + reschedule) # doesn't re-stage the registration. @@ -240,7 +240,7 @@ def request_finished( # recv so the worker emits a notif that lets P free them. # Seed remote_block_ids so add_new_req_to_recv won't KeyError. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None From a72c046463cecc35737f353b2b6cc0f44bad3f29 Mon Sep 17 00:00:00 2001 From: Miguel <60073809+migarci2@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:19:08 +0200 Subject: [PATCH 226/774] [Bugfix] Avoid nested rerank cancellation handlers (#55602) Signed-off-by: Miguel Garcia --- vllm/entrypoints/pooling/scoring/api_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index 90be32ba248a..83ed600fdd21 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -91,7 +91,6 @@ async def do_rerank(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v1(request: RerankRequest, raw_request: Request): logger.warning_once( "To indicate that the rerank API is not part of the standard OpenAI" @@ -110,6 +109,5 @@ async def do_rerank_v1(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v2(request: RerankRequest, raw_request: Request): return await do_rerank(request, raw_request) From 9d14fd3095eb169044e2da0a2a98dd1f2c77ce22 Mon Sep 17 00:00:00 2001 From: 404mario <18063257856@163.com> Date: Tue, 15 Sep 2026 12:04:22 +0800 Subject: [PATCH 227/774] [Kernel][MoE] DeepEP v2: async finalize to overlap shared experts with combine (#52781) Signed-off-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Claude Fable 5 --- .../moe/test_deepep_v2_async_finalize.py | 122 ++++++++++++++++++ .../fused_moe/prepare_finalize/deepep_v2.py | 36 ++++-- 2 files changed, 150 insertions(+), 8 deletions(-) create mode 100644 tests/kernels/moe/test_deepep_v2_async_finalize.py diff --git a/tests/kernels/moe/test_deepep_v2_async_finalize.py b/tests/kernels/moe/test_deepep_v2_async_finalize.py new file mode 100644 index 000000000000..193d8906f37c --- /dev/null +++ b/tests/kernels/moe/test_deepep_v2_async_finalize.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async. + +Hermetic: fake buffer/event, single process, no GPU collectives — only +requires deep_ep to be importable. + + T1 finalize_async -> combine is issued with async_with_compute_stream=True + and the join is deferred to the receiver, which waits + on the combine event (device-side) before the copy. + T2 DBO active -> combine falls back to async_with_compute_stream=False; + the receiver still produces the correct output. + T3 finalize (sync) -> combine is issued synchronously and the output is + copied immediately. +""" + +import pytest +import torch + +from vllm.utils.import_utils import has_deep_ep_v2 + +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="Requires DeepEP v2 (ElasticBuffer)", +) + +if has_deep_ep_v2(): + import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + ) + + +class _FakeEvent: + def __init__(self, has_event: bool): + self.event = object() if has_event else None + self.waited = False + + def current_stream_wait(self): + assert self.event is not None + self.waited = True + + +class _FakeBuffer: + def __init__(self, out: torch.Tensor): + self.out = out + self.calls: list[dict] = [] + self.last_event: _FakeEvent | None = None + + def combine( + self, + x, + handle, + topk_weights, + async_with_compute_stream, + allocate_on_comm_stream, + ): + self.calls.append({"async_with_compute_stream": async_with_compute_stream}) + # DeepEP returns a completion event only for an async combine. + self.last_event = _FakeEvent(has_event=async_with_compute_stream) + return self.out, None, self.last_event + + +def _make_pf(out: torch.Tensor): + pf = DeepEPV2PrepareAndFinalize( + buffer=_FakeBuffer(out), + num_dispatchers=1, + dp_size=1, + rank_expert_offset=0, + num_experts=8, + num_topk=2, + ) + pf.handles[0] = object() + return pf + + +def _run(pf, output: torch.Tensor, do_async: bool): + empty = torch.empty(0, dtype=torch.bfloat16) + weights = torch.empty(0, 2) + ids = torch.empty(0, 2, dtype=torch.int64) + fn = pf.finalize_async if do_async else pf.finalize + return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous()) + + +@requires_deep_ep_v2 +def test_finalize_async_defers_join_to_receiver(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is True + assert not torch.equal(dst, out_src), "join must not happen before receiver" + recv() + assert pf.buffer.last_event.waited, "receiver must join via event wait" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch): + monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True) + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + recv() + assert not pf.buffer.last_event.waited, "no event to wait on for sync combine" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_sync(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + ret = _run(pf, dst, do_async=False) + assert ret is None + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + assert torch.equal(dst, out_src) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index d834c6eff73c..b3a146274d30 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -21,6 +21,7 @@ from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, + dbo_enabled, ) @@ -89,8 +90,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): use cudagraphs anyway - Provides expert_tokens_meta for efficient batched expert kernels - Both modes use async_with_compute_stream=False (synchronous from - caller's perspective). The ElasticBuffer handles comm internally. + Dispatch always uses async_with_compute_stream=False. finalize_async + issues the combine with async_with_compute_stream=True (except under + DBO) so the modular kernel can overlap the shared-expert FFN with the + combine a2a; the returned receiver joins via a device-side event wait. """ @staticmethod @@ -465,15 +468,31 @@ def _finalize( f"got {fused_expert_output.dtype}" ) + # DBO drives its own hook/receiver schedule; keep the combine + # synchronous there (the receiver then only performs the copy). + combine_async = do_async and not dbo_enabled() combined_x, _, event = self.buffer.combine( x=fused_expert_output, handle=handle, topk_weights=None, - async_with_compute_stream=False, + async_with_compute_stream=combine_async, + allocate_on_comm_stream=combine_async, ) - output.copy_(combined_x, non_blocking=True) - return None + if do_async: + # The combine ran on DeepEP's comm stream; the modular kernel + # issues the shared-expert FFN before calling the receiver, which + # joins via a device-side cudaStreamWaitEvent (no host sync, so + # this is safe inside a captured region). + def _receiver(): + if event.event is not None: + event.current_stream_wait() + output.copy_(combined_x, non_blocking=True) + + return _receiver + else: + output.copy_(combined_x, non_blocking=True) + return None def finalize_async( self, @@ -484,16 +503,17 @@ def finalize_async( apply_router_weight_on_input: bool, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> Callable: - self._finalize( + receiver = self._finalize( output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl, - False, + True, ) - return lambda: None + assert receiver is not None + return receiver def finalize( self, From 876595275548dff6056ff0942592dbe93b070193 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:18:10 -0700 Subject: [PATCH 228/774] [CI/Build] Retry CPU image stream interruptions (#56597) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Sherlock Co-authored-by: OpenAI Codex --- .buildkite/scripts/hardware_ci/run-cpu-test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 74c7749d50ef..e8d0a24bc7d1 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -59,7 +59,7 @@ docker pull ubuntu:25.04 || true # building the docker image echo "--- :docker: Building Docker image" -BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:' +BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output' BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3 build_log="$(mktemp)" From a491b6039d8d23314efc30a2113000ea987e409e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:20:14 -0700 Subject: [PATCH 229/774] [CI] Drop root privileges for the OTel GPU sampler (#56941) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-otel/ci_otel.sh | 53 ++++++++++++- .../scripts/ci-otel/tests/test_ci_otel.py | 79 +++++++++++++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index c9897e732e51..7f1b700fe6bd 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -65,6 +65,39 @@ _ci_otel_python() { fi } +# The Docker test commands run as root, while /workdir is owned by the +# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so +# any future project-code import cannot leave root-owned files in the checkout. +_ci_otel_prepare_gpu_user() { + _CI_INFRA_GPU_UID="" + _CI_INFRA_GPU_GID="" + [ "$(id -u 2>/dev/null || :)" = "0" ] || return 0 + command -v setpriv >/dev/null 2>&1 || return 1 + + _CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}" + [ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1 + # The owner fields are numeric and whitespace-free on GNU stat. + # shellcheck disable=SC2046 + set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :) + [ "$#" -eq 2 ] || return 1 + case "$1:$2" in + *[!0-9:]* | 0:* | *:0) return 1 ;; + esac + + _CI_INFRA_GPU_UID="$1" + _CI_INFRA_GPU_GID="$2" + chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \ + "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then + # Permit traversal to the private child spool, but not directory listing. + chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1 + fi + setpriv --reuid="${_CI_INFRA_GPU_UID}" \ + --regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \ + -- test -w "${CI_INFRA_OTEL_SPOOL_DIR}" +} + ci_otel_start() { _CI_INFRA_OTEL_COMMAND_INDEX="$1" _CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}" @@ -87,9 +120,23 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - PYTHONDONTWRITEBYTECODE=1 \ - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" &2 || : + fi fi } diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ae51010f9eeb..c66530989f97 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -987,6 +987,85 @@ def upload(spans, timeout_seconds): assert [len(batch[0].events) for batch in batches] == [1, 1, 1] +def _write_executable(path: Path, body: str) -> None: + path.write_text(f"#!/bin/sh\n{body}") + path.chmod(0o755) + + +@pytest.mark.parametrize( + "workspace_owner,expect_sampler", + [("997 996", True), ("0 0", False)], +) +def test_root_shell_runs_gpu_sampler_as_checkout_owner( + tmp_path, workspace_owner, expect_sampler +): + bin_dir = tmp_path / "bin" + workspace = tmp_path / "workdir" + runtime = tmp_path / "runtime" + log = tmp_path / "calls.log" + bin_dir.mkdir() + workspace.mkdir() + + _write_executable( + bin_dir / "id", + '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n', + ) + _write_executable( + bin_dir / "stat", + f"printf '%s\\n' '{workspace_owner}'\n", + ) + for command in ("chown", "chmod"): + _write_executable(bin_dir / command, "exit 0\n") + _write_executable( + bin_dir / "setpriv", + 'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n' + 'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n' + '[ "$#" -gt 0 ] && shift\n' + 'exec "$@"\n', + ) + _write_executable( + bin_dir / "python3", + '[ "$1" = "-c" ] && exit 0\n' + '[ "$2" = "new-context" ] && {\n' + ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' + " exit 0\n" + "}\n" + 'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n' + "exit 0\n", + ) + _write_executable(bin_dir / "nvidia-smi", "exit 0\n") + shell = ( + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), + "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), + "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace), + "CI_OTEL_TEST_LOG": str(log), + }, + ) + + assert result.returncode == 0, result.stderr + calls = log.read_text() if log.exists() else "" + if expect_sampler: + expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --" + assert calls.count("setpriv") == 2 + assert calls.count(expected) == 2 + assert "gpu" in calls + else: + assert calls == "" + assert "non-root GPU sampler unavailable" in result.stderr + + @pytest.mark.parametrize("sampling", ["1", "0"]) def test_gpu_sampler_stops_with_command_and_preserves_failure_status( tmp_path, sampling From caf56ce9dac8614b11d682d1cfd288cfbe473a43 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:25:30 -0700 Subject: [PATCH 230/774] [CI] Pass the scale-out endpoint flag in EC E2E (#56819) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex Co-authored-by: Tyler Michael Smith --- docs/usage/security.md | 2 +- .../scale_out/ec_integration/run_scale_out_ec_e2e_test.sh | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 3c7f1e3f9ba8..5d4d429c6fed 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -170,7 +170,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/messages/count_tokens` - Count tokens for Anthropic messages - `/v1/responses` - Create a response -- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render` unless explicitly disabled) - `/v1/responses/{response_id}` - Retrieve a response - `/v1/responses/{response_id}/cancel` - Cancel a response - `/v1/score` - Scoring API diff --git a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh index b99b1c3ca5b0..462ada22d712 100755 --- a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh +++ b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh @@ -181,8 +181,9 @@ run_scale_out_ec() { # Encode-only EC producer instance. It runs the vision encoder and # publishes embeddings through the ECExampleConnector shared storage. echo "Starting encode instance on GPU $GPU_E, port $ENCODE_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_E" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_E" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$ENCODE_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ @@ -204,8 +205,9 @@ run_scale_out_ec() { # features plus ec_transfer_params and loads the embeddings that the # encode instance published. echo "Starting prefill instance on GPU $GPU_PD, port $PREFILL_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_PD" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_PD" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$PREFILL_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ From 45779ef99f8c7f4bd59f86a48acccc860fc3c97c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 15 Sep 2026 06:27:10 +0200 Subject: [PATCH 231/774] [Bugfix][NIXL] Fix multiple handles xfer race (#56104) Signed-off-by: NickLucche Co-authored-by: Nick Hill --- .../kv_connector/unit/test_nixl_connector.py | 161 +++++++++--------- .../unit/test_nixl_push_connector.py | 73 +++++++- .../kv_connector/v1/nixl/base_worker.py | 113 ++++++------ .../kv_connector/v1/nixl/pull_worker.py | 17 +- .../kv_connector/v1/nixl/push_worker.py | 20 ++- 5 files changed, 233 insertions(+), 151 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 298886d16095..39fb0bb54fae 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -4,6 +4,7 @@ import contextlib import inspect import os +import queue import tempfile import textwrap import time @@ -1952,17 +1953,43 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert memory_types == ["DRAM", "VRAM"] -def test_mixed_memory_read_notifies_after_both_transfers_finish(): +@pytest.fixture +def recv_worker(): + """Receive lifecycle state without distributed or device initialization.""" worker = object.__new__(NixlConnectorWorker) - worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} - worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} - worker._dram_src_handles_by_block_size = {16: 10} - worker._recving_metadata = {"request": MagicMock()} + worker.transfer_topo = MagicMock() + worker.transfer_topo.block_size_ratio.return_value = 1 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_block_size=16, remote_physical_blocks_per_logical=1 + ) + worker._physical_blocks_per_logical_kv_block = 1 + worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))} worker._recving_transfers = defaultdict(list) + worker._failed_recv_reqs = queue.Queue() + worker._recv_failures = set() + worker._replicated_pcp_done_sending = set() + worker._invalid_block_ids = queue.Queue() worker._pending_recv_notifs = {} - worker._failed_inflight_recvs = set() - worker.xfer_stats = MagicMock() + worker._reqs_to_send = {} + worker._replicated_pcp_done_sending = set() + worker._is_hma_required = False + worker._has_mamba = False + worker.use_host_buffer = False + worker.enable_permute_local_kv = False + worker.enable_heterogeneous_attn_post_process = False + worker.tp_rank = 0 + worker._log_failure = MagicMock() # type: ignore[method-assign] + worker.xfer_stats = NixlKVConnectorStats() worker.nixl_wrapper = MagicMock() + worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry() + return worker + + +def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker): + worker = recv_worker + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + worker._dram_src_handles_by_block_size = {16: 10} worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102] worker.nixl_wrapper.check_xfer_state.return_value = "DONE" @@ -1989,27 +2016,20 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish(): np.testing.assert_array_equal(device_read.args[4], [7]) worker.nixl_wrapper.send_notif.assert_not_called() - assert worker._pop_done_transfers(worker._recving_transfers) == {"request"} + assert worker.get_finished() == (set(), {"request"}) worker.nixl_wrapper.send_notif.assert_called_once_with( "prefill", notif_msg=b"request:1" ) -def test_mixed_memory_read_failure_does_not_notify_producer(): +def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker): """A failed half of a split READ must suppress its success notification.""" - worker = object.__new__(NixlConnectorWorker) - worker._recving_metadata = {"request": MagicMock()} + worker = recv_worker worker._recving_transfers = {"request": [101, 102]} worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} - worker._failed_inflight_recvs = set() - worker._is_hma_required = True - worker._failed_recv_reqs = MagicMock() - worker._log_failure = MagicMock() # type: ignore[method-assign] - worker.xfer_stats = MagicMock() - worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] - worker._pop_done_transfers(worker._recving_transfers) + assert worker.get_finished() == (set(), {"request"}) assert "request" not in worker._pending_recv_notifs worker.nixl_wrapper.send_notif.assert_not_called() @@ -2969,13 +2989,10 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): # Wait for handshake to fail time.sleep(0.3) - # Check that blocks were marked invalid - invalid_blocks = connector.get_block_ids_with_load_errors() - assert invalid_blocks == {1, 2, 3} - - # Check that request appears in get_finished + assert connector.get_block_ids_with_load_errors() == set() _, done_recving = connector.get_finished(finished_req_ids=set()) assert request_id in done_recving + assert connector.get_block_ids_with_load_errors() == {1, 2, 3} @patch( @@ -3026,6 +3043,7 @@ def test_transfer_setup_failure_returns_finished( time.sleep(0.1) connector.start_load_kv(dummy_ctx) + assert connector.get_block_ids_with_load_errors() == set() results = connector.get_transfer_results(finished_req_ids=set()) assert request_id in results.finished_recving assert results.failed_recving == {request_id} @@ -3248,74 +3266,49 @@ def test_failed_request_skips_kv_postprocessing( assert invalid_blocks == {1, 2, 3} -@patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", - FailingNixlWrapper, -) -def test_handles_failing_in_separate_polls_do_not_kill_the_engine( - default_vllm_config, dist_init -): - """Failures across polls report once, after the last writer is terminal.""" - vllm_config = create_vllm_config() - connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) - ) - connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 - ) - worker = connector.connector_worker +@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"]) +def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state): + """Failure must not release blocks while a sibling can still write to them.""" + worker = recv_worker + worker._recving_transfers["request"] = [101, 102] + worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} + worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state] - request_id = "test_two_handles_failing_in_separate_polls" - metadata = NixlConnectorMetadata() - metadata.add_new_req_to_recv( - request_id=request_id, - local_block_ids=([1, 2, 3],), - kv_transfer_params={ - "remote_block_ids": ([4, 5, 6],), - "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, - "remote_request_id": f"prefill-{request_id}", - "remote_host": "localhost", - "remote_port": 1234, - "remote_tp_size": 1, - }, - ) - connector.bind_connector_metadata(metadata) - dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) - connector.start_load_kv(dummy_ctx) - connector.bind_connector_metadata(NixlConnectorMetadata()) - time.sleep(0.1) - connector.start_load_kv(dummy_ctx) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata + worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101) - # Give the request a second handle, as a second remote rank would. - handles = worker._recving_transfers[request_id] - assert len(handles) == 1 - first, second = handles[0], object() - worker._recving_transfers[request_id] = [first, second] + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + worker.nixl_wrapper.send_notif.assert_not_called() - # Poll 1: the first handle has failed, the second is still running. - def first_failed(handle): - return "ERR" if handle is first else "PROC" + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() - with patch.object( - worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed - ): - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id not in done_recving - assert request_id in worker._recving_metadata - assert worker._recving_transfers[request_id] == [second] - assert connector.get_block_ids_with_load_errors() == set() +def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker): + """A polling exception and failed release do not prove that DMA has stopped.""" + worker = recv_worker + worker._recving_transfers["request"] = [101] + worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"] + worker.nixl_wrapper.release_xfer_handle.side_effect = [ + RuntimeError("release"), + None, + ] - # Only the last terminal handle permits the scheduler to reuse the blocks. - with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"): - _, done_recving = connector.get_finished(finished_req_ids=set()) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata - assert done_recving == {request_id} - assert request_id not in worker._recving_transfers - assert request_id not in worker._recving_metadata - assert connector.get_block_ids_with_load_errors() == {1, 2, 3} - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert done_recving == set() + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + assert worker.nixl_wrapper.release_xfer_handle.call_count == 2 def _set_test_speculative_config( diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 2c7f0cb866bf..423b0e9deca3 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -351,7 +351,7 @@ def fresh(cls) -> _StubWriterWorker: # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} w._pending_recv_notifs = {} - w._failed_inflight_recvs = set() + w._recv_failures = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() w._reqs_to_send = {} @@ -877,6 +877,68 @@ def test_get_transfer_results_evicts_completed_state(self): assert evicted == ["req-done"] assert w._push_writer_wake.is_set() + @staticmethod + def _pollable_worker() -> _StubWriterWorker: + """Stub worker with the extra base-worker state the real + ``get_transfer_results`` needs to poll ``_sending_transfers``.""" + w = _StubWriterWorker.fresh() + w.transfer_topo = MagicMock() + w.nixl_wrapper = MagicMock() + w.xfer_stats = MagicMock() + w._log_failure = MagicMock() # type: ignore[method-assign] + w._failed_recv_reqs = queue.Queue() + w._invalid_block_ids = queue.Queue() + w._pending_recv_notifs = {} + w._replicated_pcp_done_sending = set() + w.use_host_buffer = False + return w + + def _make_sending_req(self, w: _StubWriterWorker) -> str: + request_id = "req-send-1" + w._sending_transfers[request_id] = [101, 102] + w._reqs_to_send[request_id] = time.perf_counter() + 60 + w._reqs_to_process.add(request_id) + return request_id + + def test_failed_push_send_is_not_reported_done_sending(self): + """A failed P-side WRITE must not surface as finished sending. + + Reporting it would free P-side blocks as if the KV had been + delivered while D still waits for the data; the request must be + left to the lease / watchdog instead. + """ + w = self._pollable_worker() + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + # Lease tracking is preserved so the watchdog can reschedule. + assert request_id in w._reqs_to_send + assert request_id in w._reqs_to_process + # Both handles were still cleaned up. + assert request_id not in w._sending_transfers + + def test_completed_push_send_waits_for_consumer_notif(self): + """P-side sends are completed by consumer notifs / lease expiry, + not by local WRITE-handle completion.""" + w = self._pollable_worker() + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.return_value = "DONE" + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + assert request_id in w._reqs_to_send + assert request_id not in w._sending_transfers + + # The consumer notif is what completes the send. + w._pending_completion_notifs.put(f"{request_id}:1".encode()) + results = w.get_transfer_results() + assert request_id in results.finished_sending + assert request_id not in w._reqs_to_send + # ----------------------------------------------------------------- # # Negative / error-path tests # @@ -1079,7 +1141,7 @@ def test_empty_transfer_results_do_not_enqueue_eviction(self): # Wake set so the writer drains NIXL notifs even when idle. assert w._push_writer_wake.is_set() - def test_failed_push_transfer_finishes_sending(self): + def test_failed_push_transfer_does_not_evict_writer_state(self): w = _StubWriterWorker.fresh() w.nixl_wrapper = MagicMock() w.nixl_wrapper.check_xfer_state.return_value = "ERR" @@ -1096,11 +1158,12 @@ def test_failed_push_transfer_finishes_sending(self): ): results = w.get_transfer_results() - assert results.finished_sending == {"req-failed"} + assert results.finished_sending == set() assert results.finished_recving == set() assert "req-failed" not in w._sending_transfers - assert "req-failed" not in w._reqs_to_send - assert "req-failed" not in w._reqs_to_process + assert "req-failed" in w._reqs_to_send + assert "req-failed" in w._reqs_to_process + assert w._evict_finished_inbox.empty() w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17) def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 58ab5850eaf7..6f8f1e55dc6b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -692,7 +692,7 @@ def __init__( # Uses Queue for thread-safe cross-thread coordination with the # background handshake thread, matching the _ready_requests pattern. self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() - self._failed_inflight_recvs: set[ReqId] = set() + self._recv_failures: set[ReqId] = set() self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {} # Handshake metadata of this worker for NIXL transfers. @@ -1212,7 +1212,7 @@ def request_ready(f: Future[Any], entry=(req_id, meta)): error=e, meta=meta, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) fut.add_done_callback(request_ready) @@ -2618,22 +2618,32 @@ def get_transfer_results(self) -> KVConnectorTransferResults: """ assert self.transfer_topo is not None done_sending = self._get_new_notifs() - done_recving = self._pop_done_transfers(self._recving_transfers) + done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers) + if newly_failed: + self._recv_failures.update(newly_failed) done_sending.update(self._replicated_pcp_done_sending) self._replicated_pcp_done_sending.clear() - # Drain queue of requests where handshake or transfer setup failed. - failed_recv_reqs = set[ReqId]() + # Process receive failures reported by background threads. while not self._failed_recv_reqs.empty(): try: - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait()) + req_id = self._failed_recv_reqs.get_nowait() except queue.Empty: break + self._handle_failed_transfer(req_id, None, self._recv_failures) + + failed_recv_reqs: set[ReqId] = set() + if self._recv_failures: + # Keep failure bookkeeping independent of healthy request count. + for req_id in tuple(self._recv_failures): + if req_id not in self._recving_metadata: + self._recv_failures.remove(req_id) + elif req_id not in self._recving_transfers: + self._recv_failures.remove(req_id) + failed_recv_reqs.add(req_id) + done_recving.update(failed_recv_reqs) - # Add failed requests to done_recving for scheduler tracking - # (blocks are already marked invalid, scheduler will handle recompute) - done_recving.update(failed_recv_reqs) if len(done_sending) > 0 or len(done_recving) > 0: logger.debug( "Rank %s, get_finished: %s requests done sending " @@ -2654,12 +2664,17 @@ def get_transfer_results(self) -> KVConnectorTransferResults: # Skip KV sync and post-processing for failed requests if req_id in failed_recv_reqs: + self._pending_recv_notifs.pop(req_id, None) + # TODO (NickLucche) handle failed transfer for HMA. + if not self._is_hma_required: + self._invalid_block_ids.put(set(meta.local_block_ids[0])) logger.warning( "Skipping KV post-processing for failed request %s", req_id, ) continue + self._send_pending_recv_notifs(req_id) assert meta.remote is not None if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) @@ -2783,22 +2798,26 @@ def _handle_heartbeat(self, payload: str) -> None: new_expiry, ) - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: - """ - Pop completed xfers by checking for DONE state. + def _pop_done_transfers( + self, transfers: dict[str, list[int]] + ) -> tuple[set[str], set[str]]: + """Poll transfers, retaining handles until they can be released. + Args: - transfers: dict of req_id -> list[running_xfer] + transfers: Outstanding handles by request, updated in place. + Returns: - set of req_ids that have all done xfers + Requests with no outstanding handles and requests with observed + failures. Failed requests may still have outstanding handles. """ done_req_ids: set[str] = set() + failed_req_ids: set[str] = set() for req_id, handles in list(transfers.items()): in_progress = [] for handle in handles: try: xfer_state = self.nixl_wrapper.check_xfer_state(handle) if xfer_state == "DONE": - # Get telemetry from NIXL res = self.nixl_wrapper.get_xfer_telemetry(handle) self.xfer_stats.record_transfer(res) self.nixl_wrapper.release_xfer_handle(handle) @@ -2807,58 +2826,52 @@ def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: else: self._log_failure( failure_type="transfer_failed", - msg="Deferring request completion until its last " - "xfer is terminal", req_id=req_id, xfer_state=xfer_state, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer( + req_id, handle, failed_req_ids + ): + in_progress.append(handle) except Exception as e: self._log_failure( failure_type="transfer_exception", - msg="Handle is unpollable; treating it as terminal", req_id=req_id, error=e, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer(req_id, handle, failed_req_ids): + in_progress.append(handle) if not in_progress: - del transfers[req_id] done_req_ids.add(req_id) - if req_id in self._failed_inflight_recvs: - self._failed_inflight_recvs.remove(req_id) - self._report_failed_recv(req_id) - else: - self._send_pending_recv_notifs(req_id) + del transfers[req_id] else: transfers[req_id] = in_progress - return done_req_ids + return done_req_ids, failed_req_ids - def _handle_failed_transfer(self, req_id: str, handle: int | None): - """ - Handle a failed transfer by marking all (logical) blocks as invalid and - recording the failure. - - Args: - req_id: The request ID. - handle: The transfer handle. - """ - # A sibling READ may still be writing these blocks. Retain metadata - # and defer invalidation until every handle is terminal. - self._pending_recv_notifs.pop(req_id, None) - if req_id in self._recving_transfers: - self._failed_inflight_recvs.add(req_id) - else: - self._report_failed_recv(req_id) - if handle is not None: + def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool: + """Release a handle, returning False if the caller must retain it.""" + try: self.nixl_wrapper.release_xfer_handle(handle) + except Exception as e: + # A status error does not guarantee that the backend stopped DMA. + self._log_failure( + failure_type="transfer_release_failed", + msg="Retaining handle and blocks until release succeeds", + req_id=req_id, + error=e, + ) + return False + return True + + def _handle_failed_transfer( + self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None + ) -> bool: + """Record a failure and release its handle, returning False to retain it.""" self.xfer_stats.record_failed_transfer() - - def _report_failed_recv(self, req_id: str) -> None: - if (meta := self._recving_metadata.get(req_id)) is not None: - if not self._is_hma_required: - self._invalid_block_ids.put(set(meta.local_block_ids[0])) - self._failed_recv_reqs.put(req_id) + if failed_req_ids is not None: + failed_req_ids.add(req_id) + return handle is None or self._try_release_xfer_handle(req_id, handle) def _send_pending_recv_notifs(self, req_id: str) -> None: """Send notifications deferred by split DRAM/VRAM reads.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index f787639cfe77..04caebbd1f21 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -149,7 +149,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): engine_id, ) self.xfer_stats.record_kv_expired_req() - self._handle_failed_transfer(req_id, None) + self._handle_failed_transfer(req_id, None, self._recv_failures) return if any(len(group) > 0 for group in meta.local_block_ids): @@ -493,16 +493,19 @@ def _read_blocks( self._recving_transfers[request_id].append(handle) return True except Exception as e: - # mark all (logical) blocks for this request as invalid self._log_failure( failure_type="transfer_setup_failed", req_id=request_id, - msg="Marking blocks as invalid", + msg="Deferring failure reporting until outstanding transfers finish", error=e, dst_engine_id=dst_engine_id, remote_rank=remote_rank, ) - self._handle_failed_transfer(request_id, handle) + if not self._handle_failed_transfer( + request_id, handle, self._recv_failures + ): + assert handle is not None + self._recving_transfers[request_id].append(handle) return False def _read_blocks_mixed( @@ -544,7 +547,8 @@ def _read_blocks_mixed( ) except Exception: for handle in handles: - self.nixl_wrapper.release_xfer_handle(handle) + if not self._try_release_xfer_handle(request_id, handle): + self._recving_transfers[request_id].append(handle) raise self._pending_recv_notifs.setdefault(request_id, []).append( @@ -555,7 +559,8 @@ def _read_blocks_mixed( self.nixl_wrapper.transfer(handle) except Exception: for unstarted in handles[i:]: - self.nixl_wrapper.release_xfer_handle(unstarted) + if not self._try_release_xfer_handle(request_id, unstarted): + self._recving_transfers[request_id].append(unstarted) raise self._recving_transfers[request_id].append(handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 603321260de2..8b34e7ffbb4c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -341,7 +341,7 @@ def _on_handshake( self._log_failure( failure_type="push_reg_handshake_failed", req_id=rid, error=e ) - self._handle_failed_transfer(rid, None) + self._failed_recv_reqs.put(rid) return # Re-queue for the writer to send now that the handshake is done. self._reg_send_inbox.put((rid, rd)) @@ -363,7 +363,7 @@ def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: engine_id, req_id, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) return for rank, agent_name in agents.items(): try: @@ -722,9 +722,8 @@ def _xfer_blocks( # don't have a ``_recving_metadata`` entry to invalidate, so # we just release the handle and let the engine reschedule # via the lease / watchdog. - if handle is not None: - self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() + if not self._handle_failed_transfer(request_id, handle): + return handle return None # --- Notification handling on engine main thread ------------------ # @@ -803,7 +802,16 @@ def get_transfer_results(self) -> KVConnectorTransferResults: # ``_pop_done_transfers`` mutates ``_sending_transfers``; the # writer thread also appends to it, so guard the pop. with self._sending_transfers_lock: - done_pushing = self._pop_done_transfers(self._sending_transfers) + done_pushing, failed_pushing = self._pop_done_transfers( + self._sending_transfers + ) + # A failed send must never be reported as done: its blocks + # are freed via the lease / watchdog instead. + done_pushing = { + req_id + for req_id in done_pushing - failed_pushing + if req_id in self._recving_metadata + } for req_id in done_pushing: self._reqs_to_send.pop(req_id, None) self._reqs_to_process.discard(req_id) From f84b0c4bceee120936030dfb0c90af0a818263ed Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 21:28:32 -0700 Subject: [PATCH 232/774] [Model Runner V2] Tolerate lack of pinned memory support (#56908) Signed-off-by: Nick Hill --- vllm/v1/worker/cpu/buffer_utils.py | 8 +++--- vllm/v1/worker/gpu/buffer_utils.py | 43 +++++++++++++++++++++++++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/cpu/buffer_utils.py b/vllm/v1/worker/cpu/buffer_utils.py index 8ee257a28f40..cdc76ab8b47f 100644 --- a/vllm/v1/worker/cpu/buffer_utils.py +++ b/vllm/v1/worker/cpu/buffer_utils.py @@ -4,13 +4,11 @@ import torch -from vllm.utils.platform_utils import is_uva_available - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): - if not is_uva_available(): - raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu") self.np = self.cpu.numpy() - self.uva = self.cpu + + def uva(self, n: int | None = None): + return self.cpu[:n] if n is not None else self.cpu diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2c1dac2d0710..0066f80907ec 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -6,13 +6,17 @@ import numpy as np import torch +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( + PIN_MEMORY, async_tensor_h2d, get_accelerator_view_from_cpu_tensor, ) +logger = init_logger(__name__) + # Default round-robin depth for the UVA buffer pools. Must be >= the number of # concurrent in-flight steps (engine batch_queue_size). _DEFAULT_MAX_CONCURRENCY = 2 @@ -37,7 +41,7 @@ def async_copy_to_gpu( out = torch.empty_like(x, device=device) # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() + pinned = x.pin_memory() if PIN_MEMORY else x return out.copy_(pinned, non_blocking=True) @@ -47,7 +51,30 @@ def __init__(self, size: int | Sequence[int], dtype: torch.dtype): raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True) self.np = self.cpu.numpy() - self.uva = get_accelerator_view_from_cpu_tensor(self.cpu) + self._uva = get_accelerator_view_from_cpu_tensor(self.cpu) + + def uva(self, n: int | None = None) -> torch.Tensor: + return self._uva[:n] if n is not None else self._uva + + +class NonUvaBuffer: + """Explicit-copy fallback for platforms without pinned memory.""" + + def __init__(self, size: int | Sequence[int], dtype: torch.dtype): + from vllm.platforms import current_platform + + logger.warning_once( + "Pinned memory is not available on this platform; falling back " + "to device memory for UVA buffers." + ) + self.cpu = torch.zeros(size, dtype=dtype, device="cpu") + self.np = self.cpu.numpy() + self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type) + + def uva(self, n: int | None = None) -> torch.Tensor: + if n is None: + return self._uva.copy_(self.cpu) + return self._uva[:n].copy_(self.cpu[:n]) class UvaBufferPool: @@ -69,7 +96,8 @@ def __init__( self.max_concurrency = max_concurrency # UVA buffers for concurrency - self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)] + self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer + self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)] # Current buffer index self._curr = 0 @@ -80,12 +108,12 @@ def copy_to_uva(self, x: torch.Tensor | np.ndarray | list) -> torch.Tensor: n = len(x) if n > buf.cpu.shape[0]: capacity = 1 << (n - 1).bit_length() - buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype) + buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype) self._uva_bufs[self._curr] = buf # CPU-to-CPU copy dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np dst[:n] = x - return buf.uva[:n] + return buf.uva(n) def copy_to_gpu( self, @@ -141,6 +169,9 @@ def __init__( self.device = device self.max_concurrency = max_concurrency + # Without UVA, large-but-cold tensors live on the GPU directly. + uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available() + if not uva_instead_of_gpu: # Create a GPU tensor (default) self.gpu = torch.zeros(size, dtype=dtype, device=device) @@ -148,7 +179,7 @@ def __init__( # For a large but not-frequently-accessed tensor, we can use UVA instead of # GPU to save GPU memory self._uva_buf = UvaBuffer(size, dtype) - self.gpu = self._uva_buf.uva + self.gpu = self._uva_buf.uva() self._staged_write_indices: list[int] = [] self._staged_write_starts: list[int] = [] From 7dbc3d62c388df67d2764f28c90c51cc11ff0329 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 21:29:15 -0700 Subject: [PATCH 233/774] [Cleanup] Remove vestigial `tpu_input_batch.py` (#56898) Signed-off-by: Nick Hill --- .../v1/e2e/general/test_mamba_prefix_cache.py | 7 +- vllm/v1/worker/lora_model_runner_mixin.py | 6 +- vllm/v1/worker/mamba_utils.py | 11 +- vllm/v1/worker/tpu_input_batch.py | 559 ------------------ 4 files changed, 9 insertions(+), 574 deletions(-) delete mode 100644 vllm/v1/worker/tpu_input_batch.py diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index aa1f730afbad..e388b73468d3 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -31,9 +31,8 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker import mamba_utils -from vllm.v1.worker.gpu_input_batch import CachedRequestState +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch from vllm.v1.worker.mamba_utils import get_mamba_groups @@ -315,7 +314,7 @@ def check_copy_info( action: tuple[int, int], kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - input_batch: GPUInputBatch, + input_batch: InputBatch, ): assert copy_info is not None if action == (-1, -1): @@ -371,7 +370,7 @@ def fake_preprocess_mamba_fn( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 27e850b66a67..5e7f1f70c56d 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/vllm/v1/worker/lora_model_runner_mixin.py @@ -6,7 +6,6 @@ from collections.abc import Callable from contextlib import contextmanager -from typing import TypeAlias import numpy as np import torch @@ -19,10 +18,7 @@ from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor.models import supports_lora -from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch -from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch - -InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch +from vllm.v1.worker.gpu_input_batch import InputBatch logger = init_logger(__name__) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 724c714057f7..1deb89bda65c 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -26,8 +26,7 @@ UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer -from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch logger = init_logger(__name__) @@ -1436,7 +1435,7 @@ def preprocess_mamba( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, @@ -1541,7 +1540,7 @@ def preprocess_mamba( def postprocess_mamba_all( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], num_spec_tokens: int, @@ -1572,7 +1571,7 @@ def postprocess_mamba_all( def preprocess_mamba_all_specdec( scheduler_output: SchedulerOutput, - input_batch: GPUInputBatch, + input_batch: InputBatch, mamba_state_idx: dict[str, int], num_reqs: int, prev_last_scheduled_idx_buf: CpuGpuBuffer, @@ -1591,7 +1590,7 @@ def postprocess_mamba_align_gpu( num_reqs: int, num_accepted_tokens_gpu: torch.Tensor, num_accepted_tokens_cpu_tensor: torch.Tensor, - input_batch: GPUInputBatch, + input_batch: InputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py deleted file mode 100644 index c16156a5e3ef..000000000000 --- a/vllm/v1/worker/tpu_input_batch.py +++ /dev/null @@ -1,559 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Datastructures defining a TPU input batch - -from typing import cast - -import numpy as np -import torch - -from vllm.lora.request import LoRARequest -from vllm.sampling_params import SamplingType -from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.utils.collection_utils import swap_dict_values -from vllm.v1.outputs import LogprobsTensors -from vllm.v1.worker.block_table import MultiGroupBlockTable -from vllm.v1.worker.gpu_input_batch import CachedRequestState - -_SAMPLING_EPS = 1e-5 - - -class InputBatch: - def __init__( - self, - max_num_reqs: int, - max_model_len: int, - max_num_batched_tokens: int, - device: torch.device, - pin_memory: bool, - vocab_size: int, - block_sizes: list[int], # The block_size of each kv cache group - kernel_block_sizes: list[int], - max_num_blocks_per_req: list[int], - ): - self.max_num_reqs = max_num_reqs - self.max_model_len = max_model_len - self.max_num_batched_tokens = max_num_batched_tokens - self.device = device - self.pin_memory = pin_memory - self.vocab_size = vocab_size - - self._req_ids: list[str | None] = [] - self.req_id_to_index: dict[str, int] = {} - - # TODO(woosuk): This buffer could be too large if max_model_len is big. - # Find a way to reduce the CPU memory usage. - # This buffer is not directly transferred to the GPU, so it does not - # need to be pinned. - self.token_ids_cpu_tensor = torch.zeros( - (max_num_reqs, max_model_len), - device="cpu", - dtype=torch.int32, - pin_memory=False, - ) - self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() - self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32) - self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32) - self.num_computed_tokens_cpu_tensor = torch.zeros( - (max_num_reqs,), - device="cpu", - dtype=torch.int32, - pin_memory=pin_memory, - ) - self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() - - # Block table. - self.block_table = MultiGroupBlockTable( - max_num_reqs=max_num_reqs, - max_num_batched_tokens=max_num_batched_tokens, - pin_memory=pin_memory, - device=device, - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks=max_num_blocks_per_req, - ) - - # Sampling-related. - self.temperature = torch.empty( - (max_num_reqs,), dtype=torch.float32, device=device - ) - self.temperature_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.temperature_cpu = self.temperature_cpu_tensor.numpy() - self.greedy_reqs: set[str] = set() - self.random_reqs: set[str] = set() - - self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.top_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.top_p_cpu = self.top_p_cpu_tensor.numpy() - self.top_p_reqs: set[str] = set() - - self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) - self.top_k_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory - ) - self.top_k_cpu = self.top_k_cpu_tensor.numpy() - self.top_k_reqs: set[str] = set() - - self.min_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.min_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.min_p_cpu = self.min_p_cpu_tensor.numpy() - self.min_p_reqs: set[str] = set() - - # Frequency penalty related data structures - self.frequency_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.frequency_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() - self.frequency_penalties_reqs: set[str] = set() - - # Presence penalty related data structures - self.presence_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.presence_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() - self.presence_penalties_reqs: set[str] = set() - - # Repetition penalty related data structures - self.repetition_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.repetition_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() - self.repetition_penalties_reqs: set[str] = set() - - # req_index -> (min_tokens, stop_token_ids) - self.min_tokens: dict[int, tuple[int, set[int]]] = {} - - # lora related - self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64) - self.lora_id_to_request_ids: dict[int, set[str]] = {} - self.lora_id_to_lora_request: dict[int, LoRARequest] = {} - - # req_index -> generator - # NOTE(woosuk): The indices of the requests that do not have their own - # generator should not be included in the dictionary. - self.generators: dict[int, torch.Generator] = {} - - self.num_logprobs: dict[str, int] = {} - - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - - self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs - self.has_allowed_token_ids: set[str] = set() - # NOTE(lufang): In the mask tensor, if the corresponding token allowed, - # the value is False. Since we use masked_fill_ to set -inf. - self.allowed_token_ids_mask: torch.Tensor | None = None - self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None - - # req_index -> bad_words_token_ids - self.bad_words_token_ids: dict[int, list[list[int]]] = {} - - self.req_output_token_ids: list[list[int] | None] = [] - - @property - def req_ids(self) -> list[str]: - # None elements should only be present transiently - # while performing state updates to the batch. - return cast(list[str], self._req_ids) - - def add_request( - self, - request: "CachedRequestState", - req_index: int | None = None, - ) -> None: - if req_index is None: - req_index = self.num_reqs - assert req_index < self.max_num_reqs - - req_id = request.req_id - if req_index == len(self._req_ids): - self._req_ids.append(req_id) - self.req_output_token_ids.append(request.output_token_ids) - else: - self._req_ids[req_index] = req_id - self.req_output_token_ids[req_index] = request.output_token_ids - - self.req_id_to_index[req_id] = req_index - - # Copy the prompt token ids and output token ids. - num_prompt_tokens = length_from_prompt_token_ids_or_embeds( - request.prompt_token_ids, request.prompt_embeds - ) - # TODO: copy prompt_embeds - self.num_prompt_tokens[req_index] = num_prompt_tokens - self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids - start_idx = num_prompt_tokens - end_idx = start_idx + len(request.output_token_ids) - self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids - # Number of tokens without spec decode tokens. - self.num_tokens_no_spec[req_index] = request.num_tokens - - self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens - self.block_table.add_row(request.block_ids, req_index) - - sampling_params = request.sampling_params - assert sampling_params is not None, "pooling requests not supported yet" - if sampling_params.sampling_type == SamplingType.GREEDY: - # Should avoid division by zero later when apply_temperature. - self.temperature_cpu[req_index] = 0.0 - self.greedy_reqs.add(req_id) - else: - self.temperature_cpu[req_index] = sampling_params.temperature - self.random_reqs.add(req_id) - - self.top_p_cpu[req_index] = sampling_params.top_p - if sampling_params.top_p < 1: - self.top_p_reqs.add(req_id) - top_k = sampling_params.top_k - if 0 < top_k < self.vocab_size: - self.top_k_reqs.add(req_id) - else: - top_k = self.vocab_size - self.top_k_cpu[req_index] = top_k - self.min_p_cpu[req_index] = sampling_params.min_p - self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty - if sampling_params.min_p > _SAMPLING_EPS: - self.min_p_reqs.add(req_id) - if sampling_params.frequency_penalty != 0.0: - self.frequency_penalties_reqs.add(req_id) - self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty - if sampling_params.presence_penalty != 0.0: - self.presence_penalties_reqs.add(req_id) - self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty - if sampling_params.repetition_penalty != 1.0: - self.repetition_penalties_reqs.add(req_id) - if sampling_params.min_tokens: - self.min_tokens[req_index] = ( - sampling_params.min_tokens, - sampling_params.all_stop_token_ids, - ) - - # NOTE(woosuk): self.generators should not include the requests that - # do not have their own generator. - if request.generator is not None: - self.generators[req_index] = request.generator - - if sampling_params.logprobs is not None: - self.num_logprobs[req_id] = sampling_params.logprobs - if sampling_params.logit_bias is not None: - self.logit_bias[req_index] = sampling_params.logit_bias - - if sampling_params.allowed_token_ids: - self.has_allowed_token_ids.add(req_id) - if self.allowed_token_ids_mask_cpu_tensor is None: - # Lazy allocation for this tensor, which can be large. - # False means we don't fill with -inf. - self.allowed_token_ids_mask = torch.zeros( - self.max_num_reqs, - self.vocab_size, - dtype=torch.bool, - device=self.device, - ) - self.allowed_token_ids_mask_cpu_tensor = torch.zeros( - self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu" - ) - self.allowed_token_ids_mask_cpu_tensor[req_index] = True - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index][ - sampling_params.allowed_token_ids - ] = False - - if sampling_params.bad_words_token_ids: - self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids - - # Add request lora ID - if request.lora_request: - lora_id = request.lora_request.lora_int_id - if lora_id not in self.lora_id_to_request_ids: - self.lora_id_to_request_ids[lora_id] = set() - - self.request_lora_mapping[req_index] = lora_id - self.lora_id_to_request_ids[lora_id].add(request.req_id) - self.lora_id_to_lora_request[lora_id] = request.lora_request - else: - # No LoRA - self.request_lora_mapping[req_index] = 0 - - def remove_request(self, req_id: str) -> int | None: - """This method must always be followed by a call to condense().""" - - req_index = self.req_id_to_index.pop(req_id, None) - if req_index is None: - return None - self._req_ids[req_index] = None - self.req_output_token_ids[req_index] = None - - self.greedy_reqs.discard(req_id) - self.random_reqs.discard(req_id) - self.top_p_reqs.discard(req_id) - self.top_k_reqs.discard(req_id) - self.min_p_reqs.discard(req_id) - self.min_tokens.pop(req_index, None) - self.frequency_penalties_reqs.discard(req_id) - self.presence_penalties_reqs.discard(req_id) - self.repetition_penalties_reqs.discard(req_id) - self.generators.pop(req_index, None) - self.num_logprobs.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) - - # LoRA - lora_id = self.request_lora_mapping[req_index] - if lora_id != 0: - self.lora_id_to_request_ids[lora_id].discard(req_id) - if len(self.lora_id_to_request_ids[lora_id]) == 0: - self.lora_id_to_request_ids.pop(lora_id) - self.lora_id_to_lora_request.pop(lora_id) - self.request_lora_mapping[req_index] = 0 - - self.logit_bias[req_index] = None - self.has_allowed_token_ids.discard(req_id) - if self.allowed_token_ids_mask_cpu_tensor is not None: - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False) - self.bad_words_token_ids.pop(req_index, None) - return req_index - - def swap_states(self, i1: int, i2: int) -> None: - old_id_i1 = self._req_ids[i1] - old_id_i2 = self._req_ids[i2] - self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa - self.req_output_token_ids[i1], self.req_output_token_ids[i2] = ( - self.req_output_token_ids[i2], - self.req_output_token_ids[i1], - ) - assert old_id_i1 is not None and old_id_i2 is not None - self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = ( - self.req_id_to_index[old_id_i2], - self.req_id_to_index[old_id_i1], - ) - self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = ( - self.num_tokens_no_spec[i2], - self.num_tokens_no_spec[i1], - ) - self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = ( - self.num_prompt_tokens[i2], - self.num_prompt_tokens[i1], - ) - self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( - self.num_computed_tokens_cpu[i2], - self.num_computed_tokens_cpu[i1], - ) - self.temperature_cpu[i1], self.temperature_cpu[i2] = ( - self.temperature_cpu[i2], - self.temperature_cpu[i1], - ) - self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] - self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] - self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( - self.frequency_penalties_cpu[i2], - self.frequency_penalties_cpu[i1], - ) - self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = ( - self.presence_penalties_cpu[i2], - self.presence_penalties_cpu[i1], - ) - self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = ( - self.repetition_penalties_cpu[i2], - self.repetition_penalties_cpu[i1], - ) - self.min_p_cpu[i1], self.min_p_cpu[i2] = self.min_p_cpu[i2], self.min_p_cpu[i1] - - # NOTE: the following is unsafe - # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\ - # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...] - # instead, we need to temporarily copy the data for one of the indices - n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2]) - tmp = self.token_ids_cpu[i1, :n].copy() - self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n] - self.token_ids_cpu[i2, :n] = tmp - - swap_dict_values(self.generators, i1, i2) - swap_dict_values(self.min_tokens, i1, i2) - swap_dict_values(self.bad_words_token_ids, i1, i2) - - self.request_lora_mapping[i1], self.request_lora_mapping[i2] = ( - self.request_lora_mapping[i2], - self.request_lora_mapping[i1], - ) - self.logit_bias[i1], self.logit_bias[i2] = ( - self.logit_bias[i2], - self.logit_bias[i1], - ) - - if self.allowed_token_ids_mask_cpu_tensor is not None: - ( - self.allowed_token_ids_mask_cpu_tensor[i1], - self.allowed_token_ids_mask_cpu_tensor[i2], - ) = ( - self.allowed_token_ids_mask_cpu_tensor[i2], - self.allowed_token_ids_mask_cpu_tensor[i1], - ) - self.block_table.swap_row(i1, i2) - - def condense(self, empty_req_indices: list[int]) -> None: - """Move non-empty requests down into lower, empty indices. - - Args: - empty_req_indices: empty batch indices, sorted descending. - """ - num_reqs = self.num_reqs - if num_reqs == 0: - # The batched states are empty. - self._req_ids.clear() - self.req_output_token_ids.clear() - return - - # NOTE(woosuk): This function assumes that the empty_req_indices - # is sorted in descending order. - last_req_index = num_reqs + len(empty_req_indices) - 1 - while empty_req_indices: - # Find the largest non-empty index. - while last_req_index in empty_req_indices: - last_req_index -= 1 - - # Find the smallest empty index. - empty_index = empty_req_indices.pop() - if empty_index >= last_req_index: - break - - # Swap the states. - req_id = self._req_ids[last_req_index] - output_token_ids = self.req_output_token_ids[last_req_index] - assert req_id is not None - self._req_ids[empty_index] = req_id - self._req_ids[last_req_index] = None - self.req_output_token_ids[empty_index] = output_token_ids - self.req_output_token_ids[last_req_index] = None - self.req_id_to_index[req_id] = empty_index - - num_tokens = self.num_tokens_no_spec[last_req_index] - self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ - last_req_index, :num_tokens - ] - self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[ - last_req_index - ] - self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] - self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ - last_req_index - ] - self.block_table.move_row(last_req_index, empty_index) - self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] - self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] - self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] - self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ - last_req_index - ] - self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[ - last_req_index - ] - self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[ - last_req_index - ] - self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index] - generator = self.generators.pop(last_req_index, None) - if generator is not None: - self.generators[empty_index] = generator - - min_token = self.min_tokens.pop(last_req_index, None) - if min_token is not None: - self.min_tokens[empty_index] = min_token - - self.request_lora_mapping[empty_index] = self.request_lora_mapping[ - last_req_index - ] - - self.logit_bias[empty_index] = self.logit_bias[last_req_index] - - if self.allowed_token_ids_mask_cpu_tensor is not None: - self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( - self.allowed_token_ids_mask_cpu_tensor[last_req_index] - ) - - bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) - if bad_words_token_ids is not None: - self.bad_words_token_ids[empty_index] = bad_words_token_ids - # Decrement last_req_index since it is now empty. - last_req_index -= 1 - - # Trim lists to the batch size. - del self._req_ids[self.num_reqs :] - del self.req_output_token_ids[self.num_reqs :] - - def make_lora_inputs( - self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray - ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: - """ - Given the num_scheduled_tokens for each request in the batch, return - datastructures used to activate the current LoRAs. - Returns: - 1. prompt_lora_mapping: A tuple of size self.num_reqs where, - prompt_lora_mapping[i] is the LoRA id to use for the ith prompt. - 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) - where, token_lora_mapping[i] is the LoRA id to use for ith token. - 3. lora_requests: Set of relevant LoRA requests. - """ - - req_lora_mapping = self.request_lora_mapping[: self.num_reqs] - prompt_lora_mapping = tuple(req_lora_mapping) - token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) - active_lora_requests: set[LoRARequest] = set( - self.lora_id_to_lora_request.values() - ) - - return prompt_lora_mapping, token_lora_mapping, active_lora_requests - - @property - def num_reqs(self) -> int: - return len(self.req_id_to_index) - - @property - def all_greedy(self) -> bool: - return len(self.random_reqs) == 0 - - @property - def all_random(self) -> bool: - return len(self.greedy_reqs) == 0 - - @property - def no_top_p(self) -> bool: - return len(self.top_p_reqs) == 0 - - @property - def no_top_k(self) -> bool: - return len(self.top_k_reqs) == 0 - - @property - def no_min_p(self) -> bool: - return len(self.min_p_reqs) == 0 - - @property - def no_penalties(self) -> bool: - return ( - len(self.presence_penalties_reqs) == 0 - and len(self.frequency_penalties_reqs) == 0 - and len(self.repetition_penalties_reqs) == 0 - ) - - @property - def max_num_logprobs(self) -> int | None: - return max(self.num_logprobs.values()) if self.num_logprobs else None - - @property - def no_allowed_token_ids(self) -> bool: - return len(self.has_allowed_token_ids) == 0 From eb61707851f35281aa82c3ba7714f1ce9c1860e0 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 15 Sep 2026 12:34:41 +0800 Subject: [PATCH 234/774] [Bugfix] Fix stale HPC QK-norm weights after weight refit (#56706) Signed-off-by: aoshen02 --- .../model_loader/test_reload.py | 107 ++++++++++++++++++ vllm/model_executor/layers/hpc/hpc_ihc.py | 51 +++------ vllm/model_executor/layers/hpc/hpc_module.py | 3 - vllm/model_executor/layers/hpc/rope_norm.py | 63 +++-------- vllm/model_executor/model_loader/utils.py | 10 -- vllm/model_executor/models/hy_v3.py | 7 +- 6 files changed, 143 insertions(+), 98 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index f19ec35560f1..b4a544d2fc0f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc +import importlib.machinery import inspect +import sys +import types from unittest.mock import Mock from weakref import WeakKeyDictionary, ref @@ -817,6 +820,110 @@ def materialize_with_sentinel(meta_tensor): assert "0.scale" not in model.state_dict() +@pytest.fixture +def hpc_rope_norm(monkeypatch, default_vllm_config): + """Import HpcRopeNorm with the external ``hpc`` package stubbed out.""" + if "hpc" not in sys.modules: + stub = types.ModuleType("hpc") + stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None) + stub.QuantType = types.SimpleNamespace( # type: ignore[attr-defined] + QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0) + ) + monkeypatch.setitem(sys.modules, "hpc", stub) + from vllm.model_executor.layers.hpc import rope_norm + + monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {}) + return rope_norm + + +def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm): + """The fused HPC kernel is handed the live QK-norm weights after a refit. + + Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The + Q/K norm weights it receives must be the model's own float32 parameters, + so a layerwise reload that rewrites them in place (same storage) is what + the kernel sees. Previously the kernel read separate mirrors that no + reload path refreshed. + """ + from vllm.model_executor.layers.layernorm import RMSNorm + + head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4 + layer = torch.nn.Module() + layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + cos_sin_cache=torch.ones(16, head_dim), + use_qk_norm=True, + fallback_qnorm=layer.q_norm, + fallback_knorm=layer.k_norm, + kv_cache_dtype="auto", + layer_name="hpc_test_layer", + ) + model = torch.nn.Sequential(layer) + rnorm = layer.hpc_rope_norm + + calls: list[dict] = [] + + def record(*args, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False) + + def kernel_norm_weights(): + q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim + qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16) + kv_cache = torch.zeros( + 2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16 + ) + attn_layer = types.SimpleNamespace( + _k_scale=torch.ones(1), _v_scale=torch.ones(1) + ) + attn_metadata = types.SimpleNamespace( + num_actual_tokens=1, + num_decodes=1, + num_decode_tokens=1, + num_prefills=0, + num_prefill_tokens=0, + max_query_len=1, + decode_query_len=1, + qo_indptr=None, + qo_indptr_decode=None, + slot_mapping=torch.tensor([4]), + seq_lens=torch.tensor([1]), + block_table_tensor=torch.tensor([[1]]), + hpc_kv_written=False, + ) + output = torch.zeros(1, q_size, dtype=torch.bfloat16) + rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output) + return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"] + + def loaded(value): + return torch.full((head_dim,), value, dtype=torch.bfloat16) + + default_weight_loader(layer.q_norm.weight, loaded(0.5)) + default_weight_loader(layer.k_norm.weight, loaded(0.25)) + q, k = kernel_norm_weights() + assert torch.equal(q, loaded(0.5).float()) + assert torch.equal(k, loaded(0.25).float()) + q_ptr, k_ptr = q.data_ptr(), k.data_ptr() + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0)) + layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0)) + finalize_layerwise_reload(model, model_config=None) + + q, k = kernel_norm_weights() + assert q.dtype == k.dtype == torch.float32 + assert torch.equal(q, loaded(2.0).float()) + assert torch.equal(k, loaded(3.0).float()) + assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr) + assert not hasattr(rnorm, "qnorm_weight") + + @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) diff --git a/vllm/model_executor/layers/hpc/hpc_ihc.py b/vllm/model_executor/layers/hpc/hpc_ihc.py index 1a8681aa8ead..9a8065a052e3 100644 --- a/vllm/model_executor/layers/hpc/hpc_ihc.py +++ b/vllm/model_executor/layers/hpc/hpc_ihc.py @@ -20,9 +20,8 @@ - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint keeps them out of fp8 quantization via modules_to_not_convert) -The fused modules own no parameters: they hold references to the weights of the -eager layer they replace, bound in ``process_weights_after_loading`` (invoked -generically for every ``HpcModule`` by the model loader). +The fused modules own no parameters: their forwards read the weights of the +eager layer they replace directly via the captured owner reference. NOTE: The reference implementation also offers a cross-layer post+pre fusion (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre. @@ -147,41 +146,26 @@ def __init__( object.__setattr__(self, "_fallback_op", fallback_op) # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module). object.__setattr__(self, "_norm_owner", norm_owner) - self.weight = None - self.hc_scale = None - self.hc_base = None - self.rms_weight = None - self.rms_eps = 0.0 @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d]. - # It stays float32 -- the HPC kernel consumes fp32 weights directly, so - # there is no conversion (and no precision loss) here. - owner = self._fallback_op - self.weight = owner.hc_fn.weight - self.hc_scale = owner.hc_scale - self.hc_base = owner.hc_base - if self._norm_owner is not None: - self.rms_weight = self._norm_owner.weight - self.rms_eps = self._norm_owner.variance_epsilon - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: import hpc + owner = self._fallback_op + norm_owner = self._norm_owner return hpc.fuse_ihc_pre( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_fn.weight, + owner.hc_scale, + owner.hc_base, self.norm_eps, self.hc_eps, self.magnitude, - self.rms_weight, - self.rms_eps, + norm_owner.weight if norm_owner is not None else None, + norm_owner.variance_epsilon if norm_owner is not None else 0.0, ) @@ -233,29 +217,20 @@ def __init__( self.norm_eps = norm_eps # See HpcIHCPre: keep the owner out of the module tree. object.__setattr__(self, "_fallback_op", fallback_op) - self.weight = None - self.hc_scale = None - self.hc_base = None @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_head_fn weight is [hc_mult, hc_mult * d], float32. - owner = self._fallback_op - self.weight = owner.hc_head_fn.weight - self.hc_scale = owner.hc_head_scale - self.hc_base = owner.hc_head_base - def forward(self, x: torch.Tensor) -> torch.Tensor: import hpc + owner = self._fallback_op return hpc.fuse_ihc_head( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_head_fn.weight, + owner.hc_head_scale, + owner.hc_head_base, self.norm_eps, self.hc_eps, ) diff --git a/vllm/model_executor/layers/hpc/hpc_module.py b/vllm/model_executor/layers/hpc/hpc_module.py index e1efad832eb9..c3a0dab676a1 100644 --- a/vllm/model_executor/layers/hpc/hpc_module.py +++ b/vllm/model_executor/layers/hpc/hpc_module.py @@ -11,8 +11,5 @@ def __init__(self): def support(cls, *args, **kwargs): return True - def process_weights_after_loading(self, model): - pass - def forward(self, *args, **kwargs): pass diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index c1503c754c83..02f06ed89cbd 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -90,13 +90,13 @@ class HpcRopeNorm(CustomOp, HpcModule): """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). Registered as a sub-module in model layers (e.g. HunYuanAttention). - Norm weights are extracted from fallback norm modules via - process_weights_after_loading() after all weights are loaded. + The QK-Norm weights are read directly from the fallback norm modules + (built in float32 when HPC is active), eliminating the derived copy. forward() is dispatched by CustomOp framework: - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward as a splitting point — internal Python control flow is opaque - to torch.compile and not captured by CUDA Graph. + to torch.compile (its kernel launches can still be CUDA-graph captured). - In eager/native mode: forward_native() falls back to forward_cuda(). """ @@ -124,8 +124,6 @@ def __init__( self.num_kv_heads = num_kv_heads self.head_dim = head_dim - self.use_qk_norm = use_qk_norm - self.q_size = num_heads * head_dim self.kv_size = num_kv_heads * head_dim @@ -139,31 +137,18 @@ def __init__( self.head_per_group = num_heads // num_kv_heads - # Pre-allocate norm weight tensors as Parameters so they are tracked by - # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA - # Graph replay. process_weights_after_loading() updates them inplace via - # copy_() so refit does not invalidate captured graph tensor pointers. - # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout. - if use_qk_norm and fallback_qnorm is not None: - self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.qnorm_weight = None - if use_qk_norm and fallback_knorm is not None: - self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.knorm_weight = None - self.use_fp8 = "fp8" in kv_cache_dtype # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the # caller. When QK-Norm is disabled the policy is forced to NONE. self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE + if self.qk_norm_policy != QkNormPolicy.NONE and ( + fallback_qnorm is None or fallback_knorm is None + ): + raise ValueError( + "QK-Norm is enabled but fallback_qnorm/fallback_knorm were not " + "supplied; the fused kernel reads their weights directly." + ) # Register layer_name + add self to the global instance registry so the # module-level custom op (hpc_rope_norm_forward) can route back here. @@ -215,20 +200,6 @@ def support( logger.info_once("enable hpc rope_norm") return True - def process_weights_after_loading(self, model: torch.nn.Module = None) -> None: - """Copy norm weights (float32) from fallback norm modules inplace. - - Uses copy_() to preserve tensor addresses for CUDA Graph / refit - compatibility. Called by the model's load_weights() after all weights - are loaded (and generically from the model loader for DummyModelLoader - / sleep-wake_up reload paths). - """ - if self.use_qk_norm: - if self.fallback_qnorm is not None and self.qnorm_weight is not None: - self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float()) - if self.fallback_knorm is not None and self.knorm_weight is not None: - self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float()) - def register_layer_name(self, layer_name: str) -> None: """Register layer_name and add self to the global registry. @@ -430,12 +401,14 @@ def _forward_impl( k_scale = attn_layer._k_scale.reshape(1) v_scale = attn_layer._v_scale.reshape(1) - q_norm_weight = ( - self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) - k_norm_weight = ( - self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) + if self.qk_norm_policy != QkNormPolicy.NONE: + assert self.fallback_qnorm is not None + assert self.fallback_knorm is not None + q_norm_weight = self.fallback_qnorm.weight + k_norm_weight = self.fallback_knorm.weight + else: + q_norm_weight = None + k_norm_weight = None # --- Prefill --- if num_prefill_reqs > 0: diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index e12910865ec2..11e6be5cf8a1 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -15,7 +15,6 @@ from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import is_deferred_attention_layer -from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -154,15 +153,6 @@ def process_weights_after_loading( with device_loading_context(module, target_device): module.process_weights_after_loading(model_config.dtype) - # Process HPC modules (HpcRopeNorm, etc.) that rely on - # process_weights_after_loading being called from the model's - # load_weights(). When using DummyModelLoader (e.g. profiling or - # sleep/wake_up reload), the model's load_weights() is not called, so we - # must handle HPC modules here generically. - for _, module in model.named_modules(): - if isinstance(module, HpcModule): - module.process_weights_after_loading(model) - # Model-level post-load hook, after the per-layer quant finalize. if hasattr(model, "process_weights_after_loading"): model.process_weights_after_loading() diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 0b54c26f47a4..248ae04747ae 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -300,8 +300,11 @@ def __init__( prefix=f"{prefix}.attn", ) if self.use_qk_norm: - self.q_norm = RMSNorm(self.head_dim, rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, rms_norm_eps) + # The HPC fused kernel reads norm weights directly in float32; + # the fallback path keeps the model default dtype. + norm_dtype = torch.float32 if rope_support else None + self.q_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) + self.k_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE. From 7b7e117c291c023ede8fb504b00913566271e509 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:54:57 -0700 Subject: [PATCH 235/774] [CI] Extend DSv4 engine readiness timeout on main (#56735) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/disaggregated.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 5ab0e836ebd8..226d7ac94ae7 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -287,6 +287,7 @@ steps: PREFILL_BLOCK_SIZE: "256" DECODE_BLOCK_SIZE: "256" MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" + VLLM_ENGINE_READY_TIMEOUT_S: "900" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" # The V4 indexer only supports block-outermost layouts; let it pick. PREFILLER_KV_LAYOUT: "" From cd10ed6f9f6b37a8ace9cf380007e66fe12ec0c3 Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:58:52 +0900 Subject: [PATCH 236/774] [feature] [xgrammar] support `patternProperties`/`propertyNames`/`unevaluatedProperties` kw for object types (#42904) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: Artem Perevedentsev Co-authored-by: Claude Sonnet 5 --- tests/v1/structured_output/test_utils.py | 138 +++++++++++++++++- tests/v1/structured_output/test_validation.py | 9 +- vllm/v1/structured_output/backend_xgrammar.py | 41 +++++- 3 files changed, 172 insertions(+), 16 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index 11b3e5c01b38..f5f80202205b 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, @@ -10,13 +12,19 @@ pytestmark = pytest.mark.cpu_test +def grammar_accepts(schema: dict, text: str) -> bool: + return _is_grammar_accept_string(Grammar.from_json_schema(schema), text) + + @pytest.fixture def unsupported_string_schemas(): return [ {"type": "string", "format": "non_existing_format"}, + # TODO(arpera): # pattern/format is compiled but length bounds are silently dropped, # so the combination must be rejected instead of producing quietly # wrong output + # https://github.com/mlc-ai/xgrammar/issues/749 {"type": "string", "pattern": "^a+$", "maxLength": 2}, {"type": "string", "pattern": "^a+$", "minLength": 3}, {"type": "string", "format": "email", "maxLength": 10}, @@ -48,15 +56,93 @@ def unsupported_array_schemas(): @pytest.fixture -def unsupported_object_schemas(): +def unsupported_property_names_combinations(): + return [ + # TODO(arpera): + # this case is not covered by xgrammar, so we should report this bug to xgrammar + # xgrammar drops propertyNames whenever patternProperties is present, + # so "grade_12" matches the pattern and escapes the name constraint + { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + "propertyNames": {"pattern": "^grade_[0-9]$"}, # does NOT match "grade_12" + }, + # propertyNames makes xgrammar discard the sibling additionalProperties + # value schema: https://github.com/mlc-ai/xgrammar/issues/826 + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "additionalProperties": {"type": "integer"}, + }, + # propertyNames is a string schema that conventionally omits "type", so + # it escapes the string check while xgrammar drops its length bound all + # the same: https://github.com/mlc-ai/xgrammar/issues/749 + { + "type": "object", + "propertyNames": {"pattern": "^a+$", "maxLength": 2}, + }, + # xgrammar emits named `properties` separately from `propertyNames` and + # only applies `propertyNames` in its additional-properties branch, so a + # key declared in `properties` escapes the name constraint regardless of + # what `propertyNames` contains. + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "required": ["Bad"], + "propertyNames": {"pattern": "^[a-z]+$"}, + }, + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "propertyNames": {"enum": ["good"]}, + }, + # xgrammar's propertyNames-only branch falls back to an unconstrained + # value type, so it drops unevaluatedProperties whether it restricts + # values (a schema) or forbids them (false). + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": {"type": "integer"}, + }, + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": False, + }, + ] + + +@pytest.fixture +def unsupported_pattern_properties_combinations(): return [ - {"type": "object", "propertyNames": {"pattern": "^[a-z]+$"}}, - {"type": "object", "patternProperties": {"^S": {"type": "string"}}}, + # JSON Schema requires a property matched by both `properties` and + # `patternProperties` to satisfy both (conjunction), but xgrammar + # compiles them as alternative branches, so satisfying either one is + # enough. + { + "type": "object", + "properties": {"x": {"type": "string"}}, + "patternProperties": {"^x$": {"type": "integer"}}, + }, + # The same alternative-branches problem applies to overlapping + # patternProperties patterns: a key matching both patterns only has to + # satisfy one of them. + { + "type": "object", + "patternProperties": { + "^a[a-z]*$": {"type": "string"}, + "^[a-z]*z$": {"type": "integer"}, + }, + }, ] @pytest.fixture -def supported_schema(): +def supported_frankenstein_schema(): + # IMPORTANT(arpera): + # Do NOT add more keywords here! This schema is overcrowded enough that a new + # keyword can end up having no effect on the compiled grammar while the + # test still passes. Give a new keyword its own fixture and named test. return { "type": "object", "properties": { @@ -88,6 +174,19 @@ def supported_schema(): } +@pytest.fixture +def property_names_schema(): + return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} + + +@pytest.fixture +def pattern_properties_schema(): + return { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + } + + @pytest.mark.parametrize( "schema_type", [ @@ -95,7 +194,8 @@ def supported_schema(): "unsupported_integer_schemas", "unsupported_number_schemas", "unsupported_array_schemas", - "unsupported_object_schemas", + "unsupported_property_names_combinations", + "unsupported_pattern_properties_combinations", ], ) def test_unsupported_json_features_by_type(schema_type, request): @@ -106,7 +206,29 @@ def test_unsupported_json_features_by_type(schema_type, request): ) -def test_supported_json_features(supported_schema): - assert not has_xgrammar_unsupported_json_features(supported_schema), ( - "Schema should be supported" +@pytest.mark.parametrize( + "schema_type", + [ + "supported_frankenstein_schema", + "pattern_properties_schema", + "property_names_schema", + ], +) +def test_supported_json_features(schema_type, request): + schema = request.getfixturevalue(schema_type) + assert not has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be supported: {schema}" ) + + +def test_property_names_constrains_keys(property_names_schema): + assert grammar_accepts(property_names_schema, '{"score": 5}') + assert grammar_accepts(property_names_schema, '{"score": "seven"}') + assert not grammar_accepts(property_names_schema, '{"Score": 5}') + assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + + +def test_pattern_properties_constrains_keys_and_values(pattern_properties_schema): + assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index af417bc03e6f..a23189438475 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -137,7 +137,7 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): @pytest.mark.parametrize( "schema, expected_backend", [ - # multipleOf is unsupported by xgrammar, patternProperties also by guidance. + # multipleOf is unsupported by xgrammar. ( { "type": "object", @@ -145,8 +145,13 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): }, "guidance", ), + # patternProperties + properties is also unsupported by guidance. ( - {"type": "object", "patternProperties": {"^a": {"type": "string"}}}, + { + "type": "object", + "properties": {"a": {"type": "string"}}, + "patternProperties": {"^a$": {"type": "string"}}, + }, "outlines", ), ], diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 3809ba2d91b0..1ba7e2734d9e 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -236,6 +236,12 @@ def reset(self): } +def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool: + return ("pattern" in schema or "format" in schema) and ( + "minLength" in schema or "maxLength" in schema + ) + + def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: """Check if JSON schema contains features unsupported by xgrammar.""" @@ -269,16 +275,39 @@ def check_object(obj: dict[str, Any]) -> bool: # the compiled EBNF: pattern/format grammars come out byte-identical # with and without the length keywords, while maxLength alone lowers # to {0, N} correctly. + if obj.get("type") == "string" and _has_pattern_and_length_bounds(obj): + return True + + # propertyNames validates names, so it is a string schema even when it + # omits "type", which is the form that escapes the check above. if ( - obj.get("type") == "string" - and ("pattern" in obj or "format" in obj) - and ("minLength" in obj or "maxLength" in obj) + obj.get("type") == "object" + and isinstance(obj.get("propertyNames"), dict) + and _has_pattern_and_length_bounds(obj["propertyNames"]) ): return True - # Unsupported keywords for objects - if obj.get("type") == "object" and any( - key in obj for key in ("patternProperties", "propertyNames") + # FIXME: propertyNames conflicts with properties/patternProperties/ + # additionalProperties/unevaluatedProperties under xgrammar. + # https://github.com/mlc-ai/xgrammar/issues/826 + if ( + obj.get("type") == "object" + and "propertyNames" in obj + and ( + "properties" in obj + or "patternProperties" in obj + or isinstance(obj.get("additionalProperties"), dict) + or obj.get("unevaluatedProperties", True) is not True + ) + ): + return True + + # FIXME: multiple patternProperties, or patternProperties alongside + # properties, conflict under xgrammar. + if ( + obj.get("type") == "object" + and isinstance(obj.get("patternProperties"), dict) + and ("properties" in obj or len(obj["patternProperties"]) > 1) ): return True From ec6b14c0448cef6173ec1c331dc2f251ff41591a Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 23:05:59 -0700 Subject: [PATCH 237/774] [CI] Select the supported DCP backend for PCP eval (#56677) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml index f57467ac7418..b9d2675431e7 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -11,6 +11,7 @@ server_args: >- --moe-backend flashinfer_cutlass --prefill-context-parallel-size 4 --decode-context-parallel-size 4 + --dcp-comm-backend ag_rs --enable-expert-parallel --kv-cache-dtype fp8 env: From f547c23ec971fd1c78f0d89d15ac1957ae407471 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Mon, 14 Sep 2026 23:18:24 -0700 Subject: [PATCH 238/774] [Bugfix][Kimi-K3] Keep transient checkpoints out of prefix-cache eviction (#56794) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: zjy0516 Co-authored-by: zjy0516 Co-authored-by: Kimi Code --- .../test_mamba_eagle_resume_checkpoint.py | 9 ++- .../test_partial_prefix_cache_hits.py | 71 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 16 ++++- 3 files changed, 94 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index fd003fdce6ff..059b4005f0db 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing(): running = request.num_computed_tokens assert not mamba.req_to_blocks["r"][running // block_size].is_null request.shared_prefix_boundary = running - assert mamba._cache_partial_tail_block(request, running) is None + assert ( + mamba._cache_partial_tail_block( + request, + running, + retention_interval=manager.coordinator.retention_interval, + ) + is None + ) @pytest.mark.parametrize("eagle_group", [None, 0]) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 8012c82f7841..393ec484651e 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -615,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 +@pytest.mark.parametrize( + "retention_interval,transient_published", + [(None, True), (0, False)], +) +def test_internal_checkpoint_publication_respects_retention( + retention_interval, transient_published +): + """With retention_interval=0, a mid-prompt internal checkpoint is + request-local and must not enter the prefix cache; the prompt-end + checkpoint keeps the existing publication behavior. Dense retention + (None) is unchanged. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = retention_interval + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + + # Mid-prompt chunk: exports a transient checkpoint at 112. + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + transient_hash = request.block_hashes[112 // hash_block_size - 1] + transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) + assert (transient_hit is not None) == transient_published + + request.num_computed_tokens = 128 + manager.new_step_starts() + + # Final chunk: the prompt-end checkpoint at 224 stays published. + assert manager.allocate_slots(request, 112) is not None + end_hash = request.block_hashes[224 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + + +def test_transient_checkpoint_evicts_retained_boundary_hash(): + """With retention_interval=0, the checkpoint slot can coincide with a + reachable boundary block that this step's full-block pass just hashed. + The checkpoint state overwrites the slot, so the stale boundary hash + must be evicted rather than left pointing at the wrong state. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = 0 + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + # Block 2 ends at the shared-prefix junction 96, so the full-block pass + # hashes it as a retained boundary. + request.shared_prefix_boundary = 96 + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + + # The same slot is the transient checkpoint block (state@112): both the + # boundary hash and the checkpoint itself must stay unpublished. + boundary_hash = request.block_hashes[96 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None + checkpoint_hash = request.block_hashes[112 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None + + def test_eagle_block_aligned_checkpoint_replaces_newer_hash(): hash_block_size = mamba_block_size = 32 manager = make_full_mamba_manager( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index d4e68d8adc86..fc5dc410485d 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2003,7 +2003,9 @@ def cache_blocks( ) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if self.mamba_cache_mode == "align": - partial_hash = self._cache_partial_tail_block(request, num_tokens) + partial_hash = self._cache_partial_tail_block( + request, num_tokens, retention_interval=retention_interval + ) if partial_hash is not None: self.cached_blocks_this_step.add(partial_hash) if num_cached_blocks_after > num_cached_blocks_before: @@ -2038,6 +2040,7 @@ def _cache_partial_tail_block( self, request: Request, num_tokens: int, + retention_interval: int | None, ) -> BlockHashWithGroupId | None: hash_block_size = self.block_pool.hash_block_size # Re-key the reserved block at its exported checkpoint boundary. @@ -2047,6 +2050,17 @@ def _cache_partial_tail_block( blocks = self.req_to_blocks[request.request_id] assert 0 <= checkpoint_idx < len(blocks) checkpoint_block = blocks[checkpoint_idx] + if ( + retention_interval == 0 + and num_tokens < request.num_prompt_tokens + and checkpoint_position != request.shared_prefix_boundary + ): + # retention_interval == 0 keeps this transient checkpoint + # request-local. The slot may carry a hash from this step's + # full-block pass; that must go too, since the checkpoint + # state is about to overwrite the block. + self.block_pool._maybe_evict_cached_block(checkpoint_block) + return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: return None return self.block_pool.cache_partial_block( From 79e205e8cdf7864b8b9f2d23c908b887f2f5badb Mon Sep 17 00:00:00 2001 From: Jyotirmoy Roy Date: Tue, 15 Sep 2026 12:05:13 +0530 Subject: [PATCH 239/774] [Doc] Clarify ITL vs TPOT Prometheus metrics (#55283) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Jyotirmoy Roy Signed-off-by: Cenab Batu Bora <40250311+cenab@users.noreply.github.com> Signed-off-by: Itay Alroy Signed-off-by: Rita Brugarolas Brufau Signed-off-by: Josiah Davis Signed-off-by: simondanielsson Signed-off-by: Nick Hill Signed-off-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Signed-off-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> Signed-off-by: Wei Zhao Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 Signed-off-by: Matthew Wong Signed-off-by: Micah Williamson Signed-off-by: JiataiWang Signed-off-by: 子华 Signed-off-by: Colin McNamara Signed-off-by: Tyler Michael Smith Signed-off-by: Marceli Fylcek Signed-off-by: LearningMachine621 <147347705+LearningMachine621@users.noreply.github.com> Signed-off-by: jamesETsmith Signed-off-by: Cursor Signed-off-by: Yichao Zhu Signed-off-by: wangyicong Signed-off-by: Kushal <72650064+kushaldabbe@users.noreply.github.com> Signed-off-by: Tianyu Guo Signed-off-by: Summer Yang Signed-off-by: Lucas Wilkinson Signed-off-by: hungh Signed-off-by: wang.yuqi Signed-off-by: zhicheng Liu Signed-off-by: Andrii Skliar Signed-off-by: Alexander Salikov Signed-off-by: Luca Motz Signed-off-by: pmanczak Signed-off-by: Fadi Arafeh Signed-off-by: Julien Darve Signed-off-by: Bugen Zhao Signed-off-by: NickLucche Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Signed-off-by: khluu Signed-off-by: Kevin H. Luu Signed-off-by: Tony Lin Signed-off-by: Yongye Zhu Signed-off-by: Jee Jee Li Signed-off-by: andyluo7 Signed-off-by: zjy0516 Signed-off-by: JinYan Su <751080330@qq.com> Signed-off-by: nicole-lihui Signed-off-by: waizuichougou <2082431897@qq.com> Signed-off-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Signed-off-by: zRzRzRzRzRzRzR Signed-off-by: 김재억 Signed-off-by: xiaoming <1259730330@qq.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Guy Stone Signed-off-by: Djordje Ramic Signed-off-by: Isotr0py Signed-off-by: yewentao256 Signed-off-by: Lucain Pouget Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Signed-off-by: Sören Dréano Signed-off-by: Chris Eastwood Signed-off-by: Andy Lo Signed-off-by: mgoin Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Signed-off-by: Lyle Lai Signed-off-by: Vincent Signed-off-by: Giancarlo Delfin Signed-off-by: xawang Signed-off-by: Alec Flowers Signed-off-by: Fangzhou Ai Signed-off-by: Yizheng Jiao Signed-off-by: Tahsin Tunan Signed-off-by: Divakar Verma Signed-off-by: zhenwei-intel Signed-off-by: Biswa Panda Signed-off-by: LopezCastroRoberto Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Signed-off-by: Nicolò Lucchesi Signed-off-by: Canlin Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com> Signed-off-by: walterbm Signed-off-by: arcusbuilds Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Kunshang Ji Signed-off-by: Taneem Ibrahim Signed-off-by: JartX Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Signed-off-by: Artem Perevedentsev Signed-off-by: Joe Cotant Signed-off-by: yangzeyu <532183776@qq.com> Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Signed-off-by: Raphael Rialland Signed-off-by: Simon Veitner Signed-off-by: Tomas Ruiz Signed-off-by: patryk.saffer Signed-off-by: patryk.saffer Signed-off-by: PatrykSaffer Signed-off-by: Chao Lei Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Signed-off-by: Ayushman Singh Signed-off-by: Markus Hartikainen Signed-off-by: Woosuk Kwon Signed-off-by: surajm20061998 Signed-off-by: Roger Wang Signed-off-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Signed-off-by: Juntian Liu Signed-off-by: Kebe Signed-off-by: linitra24 Signed-off-by: Shenglei Fu Signed-off-by: pavelzak Signed-off-by: fai Signed-off-by: Stu Cao Signed-off-by: <> Signed-off-by: Zhewen Li Signed-off-by: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Signed-off-by: siyu Signed-off-by: Wang Zupeng Signed-off-by: Chaojun Zhang Signed-off-by: hotragn Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: specture724 Signed-off-by: lxy-alexander Signed-off-by: BWAAEEEK Signed-off-by: Jared Wen Signed-off-by: Chinmay Kulkarni Signed-off-by: jiang1.li Signed-off-by: Jared Wen Signed-off-by: Matthias Gehre Signed-off-by: liuqihao Signed-off-by: Martin Hickey Signed-off-by: jiangkuaixue123 Signed-off-by: Juan Pérez de Algaba Signed-off-by: reidliu41 Signed-off-by: Xiaozhen Bi Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: R Signed-off-by: Alex Signed-off-by: Ronald1995 Signed-off-by: arkohut <39525455+arkohut@users.noreply.github.com> Signed-off-by: Russell Bryant Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis Signed-off-by: Claude Signed-off-by: Orestis Zambounis Signed-off-by: Speculators Bot Signed-off-by: Speculators Bot Signed-off-by: zqlcode Signed-off-by: zql <37731799+zqlcode@users.noreply.github.com> Signed-off-by: Thillai Chithambaram Signed-off-by: ElizaWszola Signed-off-by: snordmann Signed-off-by: Ianniu <985683179yzy@gmail.com> Signed-off-by: Luossu Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Signed-off-by: Shiyang Chen Signed-off-by: Fangzhou Ai Signed-off-by: Michele Campi Signed-off-by: xiaguan <751080330@qq.com> Signed-off-by: zhaoguochun1995 Signed-off-by: JohnQinAMD Signed-off-by: aoshen02 Signed-off-by: Liran Schour Signed-off-by: Liqian Lin Signed-off-by: wuhangxian <1391938827@qq.com> Signed-off-by: Yongye Zhu Signed-off-by: Miguel Garcia Signed-off-by: Mario <404mario@users.noreply.github.com> Signed-off-by: aoshen02 Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: Cenab Batu Bora <40250311+cenab@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Co-authored-by: Tyler Michael Smith Co-authored-by: Rita Brugarolas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Josiah Davis Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: Nick Hill Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Wei Zhao Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com> Co-authored-by: Micah Williamson Co-authored-by: Jiatai Wang Co-authored-by: OpenAI Codex Co-authored-by: shaohuaxi Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Colin McNamara Co-authored-by: Marceli Fylcek Co-authored-by: 焦恺文 <15529091557@163.com> Co-authored-by: James E T Smith Co-authored-by: Cursor Co-authored-by: YukioZzz Co-authored-by: Yicong Wang Co-authored-by: Kushal <72650064+kushaldabbe@users.noreply.github.com> Co-authored-by: Tianyu Guo Co-authored-by: Lucas Wilkinson Co-authored-by: Summer Yang Co-authored-by: Hung Hoang (黄兴） Co-authored-by: wang.yuqi Co-authored-by: Zhicheng Liu <2084597944@qq.com> Co-authored-by: Andrii Skliar Co-authored-by: Andrii Skliar Co-authored-by: Alexander Co-authored-by: lucamotz Co-authored-by: Cyrus Leung Co-authored-by: pmanczak Co-authored-by: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Co-authored-by: Li, Jiang Co-authored-by: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Co-authored-by: Bugen Zhao Co-authored-by: Nicolò Lucchesi Co-authored-by: Andreas Karatzas Co-authored-by: Lucas Wilkinson Co-authored-by: Kevin H. Luu Co-authored-by: Tony Lin Co-authored-by: Yongye Zhu Co-authored-by: Jee Jee Li Co-authored-by: andyluo7 Co-authored-by: zjy0516 Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: JinYan Su <751080330@qq.com> Co-authored-by: Nicole LiHui 🥜 Co-authored-by: Chauncey Co-authored-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py Co-authored-by: Yuxuan Zhang Co-authored-by: 김재억 Co-authored-by: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> Co-authored-by: Guy Stone Co-authored-by: djramic Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Lucain Co-authored-by: Sören Dréano <71752785+SorenDreano@users.noreply.github.com> Co-authored-by: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com> Co-authored-by: Andy Lo Co-authored-by: Misha Goin Co-authored-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Lyle Lai Co-authored-by: mevince <26759093+mevince@users.noreply.github.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: amd-xavierwang Co-authored-by: Alec <35311602+alec-flowers@users.noreply.github.com> Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Co-authored-by: Fangzhou Ai Co-authored-by: Yizheng Jiao Co-authored-by: Tahsin Tunan Co-authored-by: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Co-authored-by: liuzhenwei Co-authored-by: Biswa Panda Co-authored-by: Woosuk Kwon Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Lucas Wilkinson Co-authored-by: Canlin Guo Co-authored-by: junuxyz <216036880+junuxyz@users.noreply.github.com> Co-authored-by: Walter Beller-Morales Co-authored-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Srijan Keshri Co-authored-by: Shengqi Chen Co-authored-by: Kunshang Ji Co-authored-by: Taneem Ibrahim Co-authored-by: JartX Co-authored-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> Co-authored-by: Artem Perevedentsev Co-authored-by: jcotant-inferact Co-authored-by: Zeyu Yang <532183776@qq.com> Co-authored-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: Raphaël Rialland <36076211+TQCB@users.noreply.github.com> Co-authored-by: Simon Veitner Co-authored-by: Tomas Ruiz Co-authored-by: PatrykSaffer Co-authored-by: patryk.saffer Co-authored-by: patryk.saffer Co-authored-by: Patryk Saffer Co-authored-by: Chao Lei Co-authored-by: Ayushman Singh Co-authored-by: Markus Hartikainen Co-authored-by: Jee Jee Li Co-authored-by: Suraj Co-authored-by: Kimi Code Co-authored-by: Roger Wang Co-authored-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Co-authored-by: Juntian Liu Co-authored-by: Kebe Co-authored-by: linitra24 Co-authored-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Co-authored-by: Pavel Zakharov Co-authored-by: stu-cao Co-authored-by: Stu Cao Co-authored-by: Zhewen Li Co-authored-by: Zhewen Li Co-authored-by: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Co-authored-by: siyu Co-authored-by: Zupeng Wang Co-authored-by: Chaojun Zhang Co-authored-by: Hotragn Pettugani <103170876+Hotragn@users.noreply.github.com> Co-authored-by: Anjie Hou <149605198+specture724@users.noreply.github.com> Co-authored-by: Alexander Lee Co-authored-by: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Co-authored-by: Jared Wen Co-authored-by: Chinmay-Kulkarni-AMD Co-authored-by: Matthias Gehre Co-authored-by: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Co-authored-by: liuqihao Co-authored-by: Martin Hickey Co-authored-by: jiangkuaixue123 Co-authored-by: Juan Pérez de Algaba <124347725+jperezdealgaba@users.noreply.github.com> Co-authored-by: Reid <61492567+reidliu41@users.noreply.github.com> Co-authored-by: Xiaozhen Bi Co-authored-by: Ganesh R Co-authored-by: AlexHuang Co-authored-by: Or Ozeri Co-authored-by: Kimi Code Co-authored-by: Ronald Co-authored-by: arkohut <39525455+arkohut@users.noreply.github.com> Co-authored-by: Russell Bryant Co-authored-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Co-authored-by: Benjamin Chislett Co-authored-by: Orestis Zambounis Co-authored-by: Speculators Bot Co-authored-by: zql <37731799+zqlcode@users.noreply.github.com> Co-authored-by: zqlcode Co-authored-by: Thillai Chithambaram <79466435+thillai-c@users.noreply.github.com> Co-authored-by: Mark McLoughlin Co-authored-by: Samuel Nordmann Co-authored-by: ElizaWszola Co-authored-by: Zhiyao(Ian) Yu <105995086+Ianniu123@users.noreply.github.com> Co-authored-by: 鐘天楽 Co-authored-by: Jorge Gallego Feliciano Co-authored-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: sychen52 <41452870+sychen52@users.noreply.github.com> Co-authored-by: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: zhaoguochun1995 <109069909+zhaoguochun1995@users.noreply.github.com> Co-authored-by: JohnQinAMD Co-authored-by: aoshen02 Co-authored-by: Bhushan Asati Co-authored-by: Megha Agarwal <19240983+meghaagr13@users.noreply.github.com> Co-authored-by: liranschour Co-authored-by: Liqian Lin Co-authored-by: wuhangxian <1391938827@qq.com> Co-authored-by: Yongye Zhu Co-authored-by: Miguel <60073809+migarci2@users.noreply.github.com> Co-authored-by: 404mario <18063257856@163.com> Co-authored-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- docs/design/metrics.md | 23 ++- .../prometheus_grafana/grafana.json | 168 +++++++++++++++++- 2 files changed, 186 insertions(+), 5 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index c7bc85756e0d..3d108890d702 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -33,7 +33,8 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts. - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts. - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT). -- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency. +- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency (time between consecutive streamed outputs). +- `vllm:request_time_per_output_token_seconds` (Histogram) - Per-request Time Per Output Token (TPOT). - `vllm:e2e_request_latency_seconds` (Histogram) - End-to-end request latency. - `vllm:request_prefill_time_seconds` (Histogram) - Request prefill time. - `vllm:request_decode_time_seconds` (Histogram) - Request decode time. @@ -49,7 +50,8 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of - `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds. - `vllm:prompt_tokens` - Prompt tokens. - `vllm:generation_tokens` - Generation tokens. -- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds. +- `vllm:inter_token_latency_seconds` - Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains exactly one token, but differs from request-level TPOT when an output contains multiple tokens or when aggregation weights differ. +- `vllm:request_time_per_output_token_seconds` - Per-request Time Per Output Token (TPOT) in seconds, computed as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. It is recorded as zero for requests that generate no more than one token. - `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds. - `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states. - `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM. @@ -63,6 +65,20 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of See [the PR which added this Dashboard](https://github.com/vllm-project/vllm/pull/2316) for interesting and useful background on the choices made here. +`vllm:inter_token_latency_seconds` records one sample per streamed output event +(the wall-clock gap between successive outputs), whereas +`vllm:request_time_per_output_token_seconds` is recorded once per finished request +as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. These two +metrics differ when an output bundles multiple tokens (e.g. speculative +decoding) or when mean ITL and mean per-request TPOT are aggregated over +differing weights. See [Understanding the latency metrics](../benchmarking/cli.md#understanding-the-latency-metrics) +for the benchmark definitions. Requests that generate no more than one token +are recorded with a TPOT of zero, while `vllm bench serve` excludes them from +its TPOT statistics. Therefore, the Prometheus histogram and benchmark TPOT +statistics can differ. Use `vllm:request_time_per_output_token_seconds` for +request-level TPOT and `vllm:inter_token_latency_seconds` when you specifically +want the inter-output latency distribution. + ### Prometheus Client Library Prometheus support was initially added [using the aioprometheus library](https://github.com/vllm-project/vllm/pull/1890), but a switch was made quickly to [prometheus_client](https://github.com/vllm-project/vllm/pull/2730). The rationale is discussed in both linked PRs. @@ -238,8 +254,7 @@ statistics relating to that iteration: iteration. - The prefill intervals for any requests that completed prefill in this iteration. -- The inter-token intervals (Time Per Output Token, TPOT), for all - requests included in this iteration. +- The inter-token intervals, for all requests included in this iteration. - The Time-To-First-Token (TTFT) for any requests that completed prefill in this iteration. However, we calculate this interval relative to when the request was first received by the frontend diff --git a/examples/observability/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json index 1c89d4593830..6144f2bfd3ad 100644 --- a/examples/observability/prometheus_grafana/grafana.json +++ b/examples/observability/prometheus_grafana/grafana.json @@ -318,7 +318,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "description": "Inter token latency in seconds.", + "description": "Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains one token, but differs from request-level TPOT under speculative decoding or varying aggregation weights.", "fieldConfig": { "defaults": { "color": { @@ -479,6 +479,172 @@ "title": "Inter Token Latency", "type": "timeseries" }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Per-request Time Per Output Token (TPOT) in seconds, computed as (end-to-end latency - TTFT) divided by the number of output tokens minus one. Requests with no more than one output token are recorded as zero; vllm bench serve excludes these requests from its TPOT statistics.", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "s" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 48 + }, + "id": 17, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.95, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P95", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.9, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.5, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "D", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "rate(vllm:request_time_per_output_token_seconds_sum{model_name=\"$model_name\"}[$__rate_interval])\n/\nrate(vllm:request_time_per_output_token_seconds_count{model_name=\"$model_name\"}[$__rate_interval])", + "hide": false, + "instant": false, + "legendFormat": "Mean", + "range": true, + "refId": "E" + } + ], + "title": "Time Per Output Token (TPOT)", + "type": "timeseries" + }, { "datasource": { "type": "prometheus", From 51918e252af4dfff381001f42b7ebe6a0454501b Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Tue, 15 Sep 2026 14:51:35 +0800 Subject: [PATCH 240/774] [Bugfix][SparseMLA] Fix piecewise cudagraph capture crash in index group (#56825) Signed-off-by: Jared Wen Co-authored-by: Kimi Code --- .../v1/attention/test_sparse_mla_backends.py | 78 ++++++++++++++++++- vllm/forward_context.py | 9 +++ vllm/v1/attention/backends/mla/index_group.py | 15 +++- vllm/v1/hisparse/runtime.py | 4 +- 4 files changed, 100 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 2ca839d59c50..6d2c4fbd5880 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -22,10 +22,17 @@ create_vllm_config, ) from vllm import _custom_ops as ops -from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.config import ( + CUDAGraphMode, + HiSparseConfig, + SpeculativeConfig, + set_current_vllm_config, +) from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( HiSparseConnectorWorker, ) +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( @@ -897,6 +904,75 @@ def test_triton_convert_req_index_to_global_index_decode_only( torch.testing.assert_close(result, reference_result, rtol=0, atol=0) +def test_index_group_convert_during_piecewise_capture(): + """In piecewise cudagraph mode the indexer is captured while the convert + runs in the following eager break. The convert's side stream must not + wait on ``logical_topk_ready`` there: an event recorded inside a captured + segment is graph-local, and an eager wait on it raises + cudaErrorInvalidValue.""" + device = torch.device(DEVICE_TYPE) + num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16 + + logical_topk_indices = torch.randint( + 0, + block_size * blocks_per_req, + (num_tokens, num_topk), + dtype=torch.int32, + device=device, + ) + builder = SparseMLAIndexGroupBuilder(logical_topk_indices) + group, layer_index = builder.register_layer(is_index_producing_layer=True) + assert layer_index == 0 and group.has_indexer + + req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests + block_table = torch.arange( + num_requests * blocks_per_req, dtype=torch.int32, device=device + ).view(num_requests, blocks_per_req) + attn_metadata = SimpleNamespace( + req_id_per_token=req_id, block_table=block_table, block_size=block_size + ) + expected = triton_convert_req_index_to_global_index( + req_id, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk, + ) + + vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B") + marker = torch.zeros(1, device=device) + capture_stream = torch.cuda.Stream(device=device) + # torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream, + # so current_stream() tracks the capture stream as in production. + with ( + torch.cuda.stream(capture_stream), + set_forward_context( + None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE + ), + ): + capture = BreakableCUDAGraphCapture() + with capture: + marker.add_(1) + # Record inside a captured segment, convert in the eager + # break — mirrors mla.py -> unified_mla_attention_with_output. + group.set_logical_topk_ready(layer_index) + capture.add_eager( + lambda: group.convert_logical_to_physical_topk( + layer_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + ) + marker.add_(1) + capture.replay() + capture_stream.synchronize() + + result = group.physical_topk_indices[:num_tokens] + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("block_size", [16]) @pytest.mark.skipif( torch.cuda.get_device_capability() < (9, 0), diff --git a/vllm/forward_context.py b/vllm/forward_context.py index f57fc2c950bd..8d9eb5d7db3b 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -209,6 +209,15 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def in_piecewise_cudagraph() -> bool: + """Whether the current forward runs in piecewise cudagraph mode (graph + segments separated by eager breaks), at capture or replay time.""" + return ( + is_forward_context_available() + and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + + def create_forward_context( attn_metadata: Any, vllm_config: VllmConfig, diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 8ba0171835fd..1b8455a2328a 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -10,6 +10,7 @@ from vllm import _custom_ops as ops from vllm.config import VllmConfig +from vllm.forward_context import in_piecewise_cudagraph from vllm.utils.torch_utils import current_stream from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, @@ -56,8 +57,16 @@ def register_layer( return layer_index def set_logical_topk_ready(self, layer_index: int) -> None: - if layer_index == 0 and self.has_indexer: - self.logical_topk_ready.record(current_stream()) + if layer_index != 0 or not self.has_indexer: + return + # In piecewise cudagraph mode the consumer converts on the side stream + # from an eager break; an event recorded inside a captured segment is + # graph-local and any eager wait on it raises cudaErrorInvalidValue. + # The convert side falls back to wait_stream there instead (and never + # touches this event), so skip the record as well. + if in_piecewise_cudagraph(): + return + self.logical_topk_ready.record(current_stream()) def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: pass @@ -129,7 +138,7 @@ def _convert_once( valid_topk_counts = self.valid_topk_counts[:num_tokens] compute_stream = current_stream() if layer_index == 0: - if self.has_indexer: + if self.has_indexer and not in_piecewise_cudagraph(): self.side_stream.wait_event(self.logical_topk_ready) else: self.side_stream.wait_stream(compute_stream) diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 314d3ba86b7a..5100379083c3 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -16,7 +16,7 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed import get_tp_group -from vllm.forward_context import get_forward_context +from vllm.forward_context import get_forward_context, in_piecewise_cudagraph from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import round_up @@ -975,7 +975,7 @@ def _resolve_and_stage_group( ) -> None: group = self.index_group compute_stream = current_stream() - if group.logical_topk_ready is not None: + if group.logical_topk_ready is not None and not in_piecewise_cudagraph(): group.copy_stream.wait_event(group.logical_topk_ready) else: group.copy_stream.wait_stream(compute_stream) From 000c7df9ffd3e470980fd4cd6b8ec1b0585500ff Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:55:04 -0700 Subject: [PATCH 241/774] [KV Offload] Add KVCR secondary-tier adapter (#53624) Signed-off-by: Moein Khazraee Signed-off-by: Moein Co-authored-by: Olga Andreeva <124622579+oandreeva-nv@users.noreply.github.com> Co-authored-by: Olga Andreeva --- .buildkite/scripts/install-kv-offload.sh | 18 + .buildkite/test_areas/misc.yaml | 1 + tests/v1/kv_offload/tiering/test_factory.py | 6 +- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 544 +++++++++++++++ vllm/v1/kv_offload/tiering/factory.py | 6 + vllm/v1/kv_offload/tiering/kvcr/__init__.py | 2 + vllm/v1/kv_offload/tiering/kvcr/manager.py | 635 ++++++++++++++++++ 7 files changed, 1211 insertions(+), 1 deletion(-) create mode 100644 .buildkite/scripts/install-kv-offload.sh create mode 100644 tests/v1/kv_offload/tiering/test_kvcr_tier.py create mode 100644 vllm/v1/kv_offload/tiering/kvcr/__init__.py create mode 100644 vllm/v1/kv_offload/tiering/kvcr/manager.py diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh new file mode 100644 index 000000000000..46d756b621b1 --- /dev/null +++ b/.buildkite/scripts/install-kv-offload.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +if [ -z "${CUDA_VERSION:-}" ]; then + exit 0 +fi + +# TODO: Use the PyPI wheel once it is published. +uv pip install --system \ + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + +# Keep only the NIXL wheel matching the CI image's CUDA runtime. +NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') +uv pip uninstall --system nixl-cu12 nixl-cu13 +uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}" diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index b45b26b2f221..ed82c7d71d68 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -130,6 +130,7 @@ steps: source_file_dependencies: *v1-core-kv-metrics-deps commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh - pytest -v -s v1/kv_offload - pytest -v -s v1/simple_kv_offload mirror: diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f5a576d25429..c391ac201631 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -10,6 +10,7 @@ 3. Error paths — missing tier_type, unknown tier_type, duplicate registration. """ +from importlib.util import find_spec from unittest.mock import MagicMock import pytest @@ -42,8 +43,11 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): - """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager.""" + """CI sentinel: registered tiers import and yield SecondaryTierManager.""" for tier_type in SecondaryTierFactory._registry: + # KVCR is an optional external dependency and may not be installed. + if tier_type == "kvcr" and find_spec("kvcr") is None: + continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py new file mode 100644 index 000000000000..38a2ae044ce2 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Collection, Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +pytest.importorskip("kvcr") + +from kvcr import ROUTER_HINT_KEY, KVCRBindings +from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig +from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpEntryResult, + OpEntryStatus, + OpHandle, + QueryStatus, +) + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadKey, + ReqContext, + make_offload_key, +) +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager +from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager + + +def _op_entries( + entries: Mapping[BlockKey, bool], +) -> dict[BlockKey, OpEntryResult]: + return { + key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED) + for key, success in entries.items() + } + + +class _StubControlChannel: + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + + def send(self, endpoint: str, message: bytes) -> bool: + return True + + def recv(self) -> list[bytes]: + return [] + + def close(self) -> None: + pass + + +class RecordingKVCR: + def __init__(self) -> None: + self.config: KVCRConfig | None = None + self.guard_config: KVCRGuardConfig | None = None + self.backend_configs = KVCRBackendConfigs() + self.constructor_bindings: KVCRBindings | None = None + self.nixl_agent_name = "recording" + self.framework_control: _StubControlChannel | None = None + self.inventory_sink = None + self.query_status = QueryStatus.MISS + self.stats: OffloadingConnectorStats | None = None + self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.discard_hint_calls: list[str] = [] + self.deliver_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] + ] = [] + self.deposit_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + ] = [] + self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] + self._next_op_handle = 1 + + def submit_hint(self, *args: Any, **kwargs: Any) -> None: + self.submit_hint_calls.append((args, kwargs)) + + def discard_hint(self, request_id: str) -> None: + self.discard_hint_calls.append(request_id) + + def query( + self, + keys: Collection[BlockKey], + request_id: str | None = None, + ) -> list[tuple[QueryStatus, CacheTier | None]]: + tier = { + QueryStatus.HIT: CacheTier.LOCAL_G2, + QueryStatus.FETCHING: CacheTier.LOCAL_G2, + QueryStatus.FETCHABLE: CacheTier.REMOTE_G2, + QueryStatus.MISS: None, + }[self.query_status] + return [(self.query_status, tier) for _ in keys] + + def deliver( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + request_id: str | None = None, + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deliver_calls.append((op_handle, dict(blocks), request_id)) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def deposit( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deposit_calls.append((op_handle, dict(blocks))) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def complete( + self, + op_handle: OpHandle, + entries: Mapping[BlockKey, bool], + ) -> None: + self.completed.append((op_handle, _op_entries(entries))) + + def poll_completed( + self, + ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]: + completed = self.completed + self.completed = [] + return completed + + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.stats + self.stats = None + return stats + + def close(self) -> None: + pass + + +class _ExternalPolicy(FIFOPolicy): + pass + + +def _job( + job_id: int, + req_context: ReqContext, + key: OffloadKey | None = None, + chunk_id: int = 0, +) -> TransferJob: + return TransferJob( + job_id=job_id, + keys=[key if key is not None else OffloadKey(b"k0")], + chunk_ids=np.array([chunk_id], dtype=np.int64), + is_promotion=True, + req_context=req_context, + ) + + +def _make_tier( + monkeypatch, + kvcr: RecordingKVCR, + *, + enable_telemetry: bool = False, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + enable_kv_cache_events: bool = False, + self_describing_kv_events: bool = False, + policy: str | None = None, + g3: dict[str, object] | None = None, + control_ports: list[int] | None = None, + data_parallel_rank_local: int | None = None, +) -> KVCRSecondaryTierManager: + def make_control(_bind_host, bind_port, advertise_host): + return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") + + def make_kvcr(config, bindings, backend_configs, guard_config): + kvcr.config = config + kvcr.guard_config = guard_config + kvcr.nixl_agent_name = config.nixl_agent_name + kvcr.backend_configs = backend_configs + kvcr.constructor_bindings = bindings + kvcr.framework_control = bindings.framework_control + kvcr.inventory_sink = bindings.inventory_sink + return kvcr + + monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) + monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) + return KVCRSecondaryTierManager( + offloading_spec=SimpleNamespace( + config=SimpleNamespace( + parallel=SimpleNamespace( + data_parallel_rank_local=data_parallel_rank_local, + ) + ), + kv_events_config=SimpleNamespace( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=self_describing_kv_events, + ), + ), + primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), + tier_type="kvcr", + router_capabilities=["router_hint"], + control_host="127.0.0.1", + control_ports=control_ports if control_ports is not None else [7777], + control_advertise_host="127.0.0.1", + enable_telemetry=enable_telemetry, + secondary_g2_slots=secondary_g2_slots, + kvcr_service_socket_path=kvcr_service_socket_path, + compatibility_digest=compatibility_digest, + policy=policy, + g3=g3, + local_dram_backend="UCX", + remote_fw_dram_backend="UCX", + ) + + +def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): + """Keep the control endpoint and guard pool aligned to the local DP rank.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + kvcr_service_socket_path="/tmp/kvcr.sock", + compatibility_digest="Opaque-Digest", + secondary_g2_slots=1, + control_ports=[7001, 7002], + data_parallel_rank_local=1, + ) + + assert kvcr.framework_control is not None + assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002" + assert kvcr.guard_config == KVCRGuardConfig( + kvcr_service_socket_path="/tmp/kvcr.sock", + guard_index=1, + compatibility_digest="Opaque-Digest", + ) + assert kvcr.config is not None + assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.backend_configs.local_dram is None + + +def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): + """Convert user-provided G3 paths to KVCR's typed configuration.""" + kvcr = RecordingKVCR() + path = tmp_path / "g3.data" + + _make_tier( + monkeypatch, + kvcr, + g3={"paths": [str(path)], "capacity_bytes_per_file": 64}, + ) + + assert kvcr.backend_configs.g3 == G3Options( + paths=(path,), + capacity_bytes_per_file=64, + ) + + +def test_kvcr_tier_adapts_request_and_load(monkeypatch): + """Check hint forwarding, key conversion, load descriptors, and cleanup.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + router_hint = {"opaque": True} + ctx = ReqContext( + req_id="req", + kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()}, + ) + key = make_offload_key((123).to_bytes(8, "big"), 0) + same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) + other_key = make_offload_key((124).to_bytes(8, "big"), 0) + + tier.on_new_request(ctx) + + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})] + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.key_adapter is not None + monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0") + decode = bindings.key_adapter.decode + hash_123 = (123).to_bytes(8, "big") + assert decode(BlockKey(bytes(key))) == hash_123 + assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123 + assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big") + + tier.submit_load(_job(7, ctx, key=key, chunk_id=2)) + + # Check load destinations and job completion independently of hint parsing. + assert len(kvcr.submit_hint_calls) == 1 + _, blocks, request_id = kvcr.deliver_calls[0] + assert request_id == "req" + assert list(blocks) == [key] + (descriptor,) = blocks[key] + assert descriptor.end_point_name == kvcr.nixl_agent_name + assert descriptor.addr == tier._primary_base_addr + 2 * 16 + assert descriptor.size == 16 + assert list(tier.get_finished_jobs()) == [JobResult(7, True)] + + # Here we verify that request cleanup discards the request-scoped hint in + # KVCR. + tier.on_request_finished(ctx) + assert kvcr.discard_hint_calls == ["req"] + + +def test_kvcr_tier_allows_request_without_router_hint(monkeypatch): + """Keep router hints optional for requests from non-hint-aware routers.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + + tier.on_new_request(ReqContext(req_id="req")) + + assert kvcr.submit_hint_calls == [] + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (QueryStatus.MISS, LookupResult.MISS), + (QueryStatus.HIT, LookupResult.HIT), + (QueryStatus.FETCHABLE, LookupResult.HIT), + (QueryStatus.FETCHING, LookupResult.RETRY), + ], +) +def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): + """Map KVCR cache states to the scheduler's lookup contract.""" + kvcr = RecordingKVCR() + kvcr.query_status = status + tier = _make_tier(monkeypatch, kvcr) + + assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected + + +def test_kvcr_tier_serves_primary_pin_request(monkeypatch): + """Hold primary-tier hits until KVCR releases the corresponding pin.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) + hit_keys = (keys[0], keys[2]) + chunk_ids = {keys[0]: 1, keys[2]: 5} + lifecycle: list[str] = [] + + class Parent: + def on_new_request(self, req_context): + lifecycle.append("new") + return SimpleNamespace() + + def lookup(self, key, req_context): + return LookupResult.HIT if key in hit_keys else LookupResult.MISS + + def create_store_job(self, requested_keys, req_context): + return TransferJob( + job_id=11, + keys=requested_keys, + chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), + is_promotion=True, + req_context=req_context, + ) + + def on_request_finished(self, req_context): + lifecycle.append("finished") + + bindings = kvcr.constructor_bindings + assert bindings is not None + request = bindings.request_pin(keys) + + tier.serve_external_requests(Parent()) + + [(queued_request, result)] = bindings.poll_pin_results() + assert queued_request == request + assert result is not None + pin_handle, descriptors = result + assert descriptors[keys[1]] is None + block_descriptors = descriptors[keys[2]] + assert block_descriptors is not None + (descriptor,) = block_descriptors + assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert lifecycle == ["new", "finished"] + + polls = 0 + + def release_pin(): + nonlocal polls + polls += 1 + bindings.release_pin(pin_handle) + return [] + + monkeypatch.setattr(kvcr, "poll_completed", release_pin) + tier.drain_jobs() + + assert polls == 1 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + +def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch): + """Keep telemetry opt-in and namespace metrics only at the vLLM boundary.""" + assert KVCRSecondaryTierManager.build_metric_definitions({}) == {} + definitions = KVCRSecondaryTierManager.build_metric_definitions( + {"enable_telemetry": True} + ) + assert "vllm:kvcr_duration_seconds" in definitions + assert "vllm:kvcr_transfer_blocks" in definitions + + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True) + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.stats_factory is not None + stats = bindings.stats_factory() + stats.increase_counter( + "kvcr_transfer_blocks", + 2, + ("remote_deliver",), + ) + kvcr.stats = stats + + returned = tier.get_stats() + + assert returned is stats + assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2} + + +@pytest.mark.parametrize( + ("policy", "expected_type"), + [ + ("fifo", FIFOPolicy), + ("lru", LRUPolicy), + ("g3_fifo", G3FIFOPolicy), + ("g3_lru", G3LRUPolicy), + (f"{__name__}._ExternalPolicy", _ExternalPolicy), + ], +) +def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type): + """Resolve every built-in and fully qualified external policy.""" + kvcr = RecordingKVCR() + _make_tier(monkeypatch, kvcr, policy=policy) + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert type(bindings.policy) is expected_type + + +@pytest.mark.parametrize( + ("socket_path", "digest"), + [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")], +) +def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest): + """Reject partial guard configuration before connecting to the service.""" + with pytest.raises(ValueError, match="configured together"): + _make_tier( + monkeypatch, + RecordingKVCR(), + kvcr_service_socket_path=socket_path, + compatibility_digest=digest, + ) + + +def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): + """Cover store descriptors and inventory translation at the KVCR boundary.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + secondary_g2_slots=2, + enable_kv_cache_events=True, + self_describing_kv_events=True, + ) + local_dram = kvcr.backend_configs.local_dram + assert local_dram is not None + assert [(name, size) for name, _, size in local_dram.pools] == [ + ("", 2 * tier._primary_row_stride) + ] + + key = OffloadKey(b"k0") + tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) + + _, blocks = kvcr.deposit_calls[0] + assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + assert kvcr.inventory_sink is not None + kvcr.inventory_sink( + InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False) + ) + kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False)) + events = list(tier.take_events()) + assert [ + (event.keys, event.medium, event.ownership, event.removed) for event in events + ] == [ + ([key], Medium.CPU, "kvcr", False), + ([key], Medium.STORAGE, "kvcr", False), + ] + assert all(event.removal_expected for event in events) + tier.shutdown() + + +def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): + """Require tier-aware events whenever KVCR owns local cache inventory.""" + with pytest.raises(ValueError, match="self_describing_kv_events"): + _make_tier( + monkeypatch, + RecordingKVCR(), + secondary_g2_slots=1, + enable_kv_cache_events=True, + ) + + +def test_kvcr_tier_accumulates_block_results(monkeypatch): + """Wait for every block result and preserve partial success while draining.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = [OffloadKey(b"k0"), OffloadKey(b"k1")] + tier.submit_load( + TransferJob( + job_id=13, + keys=keys, + chunk_ids=np.array([0, 1], dtype=np.int64), + is_promotion=True, + req_context=ReqContext(req_id="req"), + ) + ) + + op_handle, _, _ = kvcr.deliver_calls[0] + kvcr.completed.clear() + kvcr.complete(op_handle, {keys[0]: True}) + assert list(tier.get_finished_jobs()) == [] + + kvcr.complete(op_handle, {keys[1]: False}) + tier.drain_jobs() + + [result] = tier.get_finished_jobs() + assert result.job_id == 13 + assert not result.success + assert result.successful_keys == {keys[0]} diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 17f8430204d9..3310167b870b 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -123,3 +123,9 @@ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]: "vllm.v1.kv_offload.tiering.obj.manager", "ObjectStoreSecondaryTierManager", ) + +SecondaryTierFactory.register_tier( + "kvcr", + "vllm.v1.kv_offload.tiering.kvcr.manager", + "KVCRSecondaryTierManager", +) diff --git a/vllm/v1/kv_offload/tiering/kvcr/__init__.py b/vllm/v1/kv_offload/tiering/kvcr/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py new file mode 100644 index 000000000000..00186279a4bc --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM secondary-tier adapter for KVCR.""" + +import ctypes +import mmap +import socket +import time +import uuid +from collections.abc import Collection, Iterable, Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from kvcr import ( + DURATION_METRIC, + KVCR, + ROUTER_HINT_CAPABILITIES, + ROUTER_HINT_KEY, + STATE_METRIC, + TRANSFER_BLOCKS_METRIC, + TRANSFER_BYTES_METRIC, + KVCRBindings, +) +from kvcr.config import ( + FrameworkDramInput, + G3Options, + KeyAdapter, + KVCRBackendConfigs, + KVCRConfig, + KVCRGuardConfig, + LocalDramOptions, + RemoteFWDramOptions, +) +from kvcr.control_channels import ZmqPeerControlChannel +from kvcr.policy import ( + FIFOPolicy, + G3FIFOPolicy, + G3LRUPolicy, + KVCachePolicy, + LRUPolicy, +) +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpHandle, + PinRequestId, + PinResult, + QueryStatus, +) +from typing_extensions import override + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.core.kv_cache_utils import ( + BlockHash, + maybe_convert_block_hash, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadingCounterMetadata, + OffloadingEvent, + OffloadingGaugeMetadata, + OffloadingHistogramMetadata, + OffloadingMetricMetadata, + OffloadKey, + ReqContext, + RequestOffloadingContext, + get_offload_block_hash, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + ParentManager, + SecondaryTierManager, + TransferJob, +) + +if TYPE_CHECKING: + from vllm.v1.kv_offload.base import OffloadingSpec + + +_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES + +logger = init_logger(__name__) + +_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = { + "fifo": FIFOPolicy, + "lru": LRUPolicy, + "g3_fifo": G3FIFOPolicy, + "g3_lru": G3LRUPolicy, +} + + +def _resolve_policy(name: str | None) -> KVCachePolicy | None: + if name is None: + return None + policy_type = _BUILTIN_POLICIES.get(name) + if policy_type is None: + if "." not in name: + raise ValueError( + f"Unknown KVCR policy: {name!r}. " + f"Supported: {list(_BUILTIN_POLICIES)}; external policies " + "must use a fully qualified module.Class path." + ) + policy_type = resolve_obj_by_qualname(name) + if not isinstance(policy_type, type) or not issubclass( + policy_type, KVCachePolicy + ): + raise TypeError(f"{name} is not a KVCachePolicy") + return policy_type() + + +def _vllm_metric_name(name: str) -> str: + return f"vllm:{name}" + + +def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]: + return { + _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata( + documentation="KVCR operation and stage duration in seconds.", + labelnames=("scope", "result"), + ), + _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata( + documentation="Bytes transferred by KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata( + documentation="KV blocks transferred by successful KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata( + documentation="Current KVCR metadata and operation counts.", + labelnames=("resource",), + ), + } + + +class _FrameworkPinAdapter: + """Adapt KVCR framework pin callbacks to primary-tier store jobs.""" + + def __init__(self, tier: "KVCRSecondaryTierManager") -> None: + self._tier = tier + self._next_request_id = 0 + self._next_pin_request_id = 0 + self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {} + self._completed_pins: list[tuple[PinRequestId, PinResult]] = [] + self._pin_jobs: dict[str, int] = {} + self._pin_job_results: list[JobResult] = [] + + def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId: + # The parent callback is valid only from serve_external_requests(). + request = PinRequestId(self._next_pin_request_id) + self._next_pin_request_id += 1 + self._pending_pins[request] = tuple(keys) + return request + + def process_pending(self, parent: ParentManager) -> None: + pending = self._pending_pins + self._pending_pins = {} + self._completed_pins.extend( + (request, self._resolve_pin(parent, keys)) + for request, keys in pending.items() + ) + + def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]: + completed = self._completed_pins + self._completed_pins = [] + return completed + + def cancel_pin_request(self, request: PinRequestId) -> None: + self._pending_pins.pop(request, None) + + def _resolve_pin( + self, parent: ParentManager, keys: Collection[BlockKey] + ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) + request_id = f"kvcr-source:{self._next_request_id}" + self._next_request_id += 1 + req_context = ReqContext(req_id=request_id) + started = False + job: TransferJob | None = None + pin_handle: str | None = None + try: + parent.on_new_request(req_context) + started = True + hit_keys = tuple( + key + for key in offload_keys + if parent.lookup(key, req_context) is LookupResult.HIT + ) + if not hit_keys: + return None + + job = parent.create_store_job(hit_keys, req_context) + job_keys = tuple(job.keys) + chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) + if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): + return None + + descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + BlockKey(bytes(key)): None for key in offload_keys + } + descriptors.update( + ( + BlockKey(bytes(key)), + [self._tier._make_descriptor(chunk_id)], + ) + for key, chunk_id in zip(job_keys, chunk_ids) + ) + pin_handle = f"kvcr-job-{job.job_id}" + self._pin_jobs[pin_handle] = job.job_id + return pin_handle, descriptors + except Exception: + logger.warning("KVCR source acquisition failed", exc_info=True) + return None + finally: + if started: + try: + parent.on_request_finished(req_context) + except Exception: + logger.warning("KVCR source request cleanup failed", exc_info=True) + if job is not None and pin_handle is None: + self._pin_job_results.append( + JobResult(job_id=job.job_id, success=False) + ) + + def release_pin(self, pin_handle: str) -> bool: + job_id = self._pin_jobs.pop(pin_handle, None) + if job_id is not None: + self._pin_job_results.append(JobResult(job_id=job_id, success=True)) + return True + + def has_active_pins(self) -> bool: + return bool(self._pin_jobs) + + def take_pin_job_results(self) -> list[JobResult]: + results = self._pin_job_results + self._pin_job_results = [] + return results + + +class _VllmKeyAdapter: + """Adapt vLLM's local key format to KVCR.""" + + def encode(self, framework_key: object) -> BlockKey: + if not isinstance(framework_key, bytes): + raise TypeError("vLLM offload keys must be bytes") + return BlockKey(framework_key) + + def decode(self, key: BlockKey) -> int | bytes: + block_hash = BlockHash(get_offload_block_hash(OffloadKey(key))) + return maybe_convert_block_hash(block_hash) + + +# Job ID, remaining blocks, aggregate success, and successful load keys. +_JobState = tuple[int, int, bool, set[OffloadKey] | None] + + +class KVCRSecondaryTierManager(SecondaryTierManager): + """Secondary tier wrapper around the KVCR KV P2P API.""" + + @classmethod + @override + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + if not extra_config.get("enable_telemetry", False): + return {} + return _kvcr_metric_definitions() + + def __init__( + self, + offloading_spec: "OffloadingSpec", + primary_kv_view: memoryview, + tier_type: str, + router_capabilities: Iterable[str] | None = None, + control_host: str = "0.0.0.0", + control_ports: list[int] | None = None, + control_advertise_host: str | None = None, + eager_ctrl_connect: bool = True, + opportunistic_query: bool = True, + enable_telemetry: bool = False, + operation_timeout_ms: int = 1000, + abandon_timeout_ms: int = 5000, + metadata_retry_interval_ms: int = 100, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + policy: str | None = None, + g3: dict[str, Any] | None = None, + local_dram_backend: str = "UCX", + remote_fw_dram_backend: str = "UCX", + ) -> None: + super().__init__(offloading_spec, primary_kv_view, tier_type) + selected_policy = _resolve_policy(policy) + if (kvcr_service_socket_path is None) != (compatibility_digest is None): + raise ValueError( + "kvcr_service_socket_path and compatibility_digest must be " + "configured together" + ) + missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ()) + if missing: + raise ValueError( + f"KVCR secondary tier requires router capabilities: {sorted(missing)}" + ) + events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events + self_describing_events = ( + offloading_spec.kv_events_config.self_describing_kv_events + ) + if ( + events_enabled + and ( + secondary_g2_slots > 0 + or kvcr_service_socket_path is not None + or g3 is not None + ) + and not self_describing_events + ): + raise ValueError( + "KVCR local G2/G3 inventory requires self_describing_kv_events" + ) + self._key_adapter: KeyAdapter = _VllmKeyAdapter() + advertise_host = control_advertise_host or socket.gethostname() + dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local + if dp_local_rank is None: + if control_ports is None or len(control_ports) != 1: + raise ValueError( + "control_ports must contain exactly one port when " + "data_parallel_rank_local is unset" + ) + dp_local_rank = 0 + if control_ports is None or not 0 <= dp_local_rank < len(control_ports): + raise ValueError( + "control_ports must contain a port for " + f"local data-parallel rank {dp_local_rank}" + ) + if any(not 1 <= port <= 65535 for port in control_ports): + raise ValueError("control_ports must be between 1 and 65535") + control = ZmqPeerControlChannel( + control_host, + control_ports[dp_local_rank], + advertise_host, + ) + # Give colocated workers distinct transport listen ports. + with socket.socket() as _s: + _s.bind(("", 0)) + _nixl_listen_port = _s.getsockname()[1] + self._primary_base_addr = ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ) + if primary_kv_view.strides is None: + raise ValueError("primary KV memoryview must expose strides") + self._primary_row_stride = int(primary_kv_view.strides[0]) + if secondary_g2_slots < 0: + raise ValueError("secondary_g2_slots must be non-negative") + if ( + secondary_g2_slots or kvcr_service_socket_path is not None + ) and not events_enabled: + logger.warning( + "KVCR local DRAM is enabled, but KV cache events are disabled; " + "local DRAM inventory will not be published" + ) + local_mapping: mmap.mmap | None = None + local_dram: LocalDramOptions | None = None + if secondary_g2_slots: + if kvcr_service_socket_path is not None: + logger.warning( + "secondary_g2_slots is ignored when " + "kvcr_service_socket_path is configured" + ) + else: + local_mapping = mmap.mmap( + -1, secondary_g2_slots * self._primary_row_stride + ) + local_dram = LocalDramOptions( + pools=[ + ( + "", + ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)), + len(local_mapping), + ) + ], + backend=local_dram_backend, + ) + guard_config = ( + KVCRGuardConfig( + kvcr_service_socket_path=kvcr_service_socket_path, + guard_index=dp_local_rank, + compatibility_digest=compatibility_digest, + ) + if kvcr_service_socket_path is not None and compatibility_digest is not None + else None + ) + g3_config = None + if g3 is not None: + g3_values = dict(g3) + g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"]) + g3_config = G3Options(**g3_values) + nixl_agent_name = f"KVCR-{uuid.uuid4()}" + self._framework_pin_adapter = _FrameworkPinAdapter(self) + self._inventory_events: list[OffloadingEvent] = [] + + try: + self._kvcr = KVCR( + KVCRConfig( + pool_layouts=[("", self._primary_row_stride)], + enable_telemetry=enable_telemetry, + operation_timeout_ms=operation_timeout_ms, + abandon_timeout_ms=abandon_timeout_ms, + nixl_agent_name=nixl_agent_name, + nixl_listen_port=_nixl_listen_port, + ), + KVCRBindings( + request_pin=self._framework_pin_adapter.request_pin, + poll_pin_results=self._framework_pin_adapter.poll_pin_results, + release_pin=self._framework_pin_adapter.release_pin, + cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request), + framework_control=control, + key_adapter=self._key_adapter, + inventory_sink=(self._record_inventory if events_enabled else None), + stats_factory=( + OffloadingConnectorStats if enable_telemetry else None + ), + policy=selected_policy, + ), + KVCRBackendConfigs( + framework_dram=FrameworkDramInput( + self._primary_base_addr, primary_kv_view.nbytes + ), + local_dram=local_dram, + g3=g3_config, + remote_fw_dram=RemoteFWDramOptions( + eager_ctrl_connect=eager_ctrl_connect, + opportunistic_query=opportunistic_query, + metadata_retry_interval_ms=metadata_retry_interval_ms, + backend=remote_fw_dram_backend, + ), + ), + guard_config, + ) + except BaseException: + control.close() + if local_mapping is not None: + local_mapping.close() + raise + self._local_dram_mmap = local_mapping + self._finished_jobs: list[JobResult] = [] + self._jobs_by_op: dict[OpHandle, _JobState] = {} + + @override + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + block_key = self._key_adapter.encode(key) + status, _ = self._kvcr.query((block_key,), req_context.req_id)[0] + if status is QueryStatus.FETCHING: + return LookupResult.RETRY + if status in (QueryStatus.HIT, QueryStatus.FETCHABLE): + return LookupResult.HIT + return LookupResult.MISS + + @override + def submit_load(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + + op_handle = self._kvcr.deliver( + blocks, request_id=job_metadata.req_context.req_id + ) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + set(), + ) + + @override + def submit_store(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + op_handle = self._kvcr.deposit(blocks) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + None, + ) + + @override + def get_finished_jobs(self) -> Iterable[JobResult]: + results = self._finished_jobs + self._finished_jobs = [] + results.extend(self._poll_finished_jobs()) + return results + + @override + def get_stats(self) -> OffloadingConnectorStats | None: + stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats()) + if stats is not None: + # Add the framework namespace only at the vLLM boundary. + for metrics in stats.data.values(): + for name in tuple(metrics): + metrics[_vllm_metric_name(name)] = metrics.pop(name) + return stats + + @override + def serve_external_requests(self, parent: ParentManager) -> None: + self._framework_pin_adapter.process_pending(parent) + + def _poll_finished_jobs(self) -> list[JobResult]: + results: list[JobResult] = [] + for op_handle, entries in self._kvcr.poll_completed(): + job_state = self._jobs_by_op.get(op_handle) + if job_state is None: + continue + job_id, remaining, success, successful_keys = job_state + remaining -= len(entries) + success = success and all(entry.success for entry in entries.values()) + if successful_keys is not None: + successful_keys.update( + OffloadKey(bytes(key)) + for key, entry in entries.items() + if entry.success + ) + if remaining > 0: + self._jobs_by_op[op_handle] = ( + job_id, + remaining, + success, + successful_keys, + ) + continue + self._jobs_by_op.pop(op_handle, None) + results.append( + JobResult( + job_id=job_id, + success=success, + successful_keys=successful_keys if not success else None, + ) + ) + results.extend(self._framework_pin_adapter.take_pin_job_results()) + return results + + @override + def has_pending_work(self) -> bool: + # Keep scheduler polling: control messages may arrive without a load job. + return True + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + events = self._inventory_events + self._inventory_events = [] + return events + + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + params = getattr(req_context, "kv_transfer_params", None) + if isinstance(params, Mapping): + hint = params.get(ROUTER_HINT_KEY) + if hint is not None: + self._kvcr.submit_hint( + request_id=req_context.req_id, + hints=hint, + ) + return RequestOffloadingContext() + + @override + def on_request_finished(self, req_context: ReqContext) -> None: + self._kvcr.discard_hint(req_context.req_id) + + @override + def drain_jobs(self) -> None: + while self._jobs_by_op or self._framework_pin_adapter.has_active_pins(): + drained = self._poll_finished_jobs() + if drained: + self._finished_jobs.extend(drained) + else: + time.sleep(0.001) + + @override + def shutdown(self) -> None: + self.drain_jobs() + # TODO: Keep registered buffers alive if KVCR.close() fails. + self._kvcr.close() + if self._local_dram_mmap is not None: + self._local_dram_mmap.close() + self._local_dram_mmap = None + + def _record_inventory(self, event: InventoryEvent) -> None: + if event.tier is CacheTier.LOCAL_G2: + medium = Medium.CPU + elif event.tier is CacheTier.G3: + medium = Medium.STORAGE + else: + raise ValueError(f"unsupported KVCR inventory tier: {event.tier}") + self._inventory_events.append( + OffloadingEvent( + keys=[OffloadKey(bytes(key)) for key in event.keys], + medium=medium, + removed=event.removed, + ownership="kvcr", + removal_expected=True, + ) + ) + + def _make_descriptor(self, chunk_id: int) -> MemDescriptor: + return MemDescriptor( + end_point_name=self._kvcr.config.nixl_agent_name, + mem_type="DRAM", + addr=self._primary_base_addr + chunk_id * self._primary_row_stride, + size=self._primary_row_stride, + device_Id=0, + info="", + ) From 8263ea12bd8fa7584277f5200d521193259707b7 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 15 Sep 2026 00:57:18 -0700 Subject: [PATCH 242/774] [Build] Move DeepGEMM pin to the vLLM fork and update the Mega MoE call convention (#56876) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- cmake/external_projects/deepgemm.cmake | 16 ++++++++++------ docker/Dockerfile | 1 + tools/build_deepgemm_C.py | 2 +- tools/install_deepgemm.sh | 8 +++++--- vllm/models/kimi_k3/nvidia/model.py | 6 ++++-- vllm/utils/deep_gemm.py | 6 +++++- 6 files changed, 26 insertions(+), 13 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 72f7f5eb66a8..2ff69bd8b6b3 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,11 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # Pinned to the tip of the nv_dev branch (SM120 support). - set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 + # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU + # alpha/beta, and SiTU for FP8/FP4 Mega MoE. + set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -51,7 +53,7 @@ else() BINARY_DIR "${_deepgemm_bin}" GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit" GIT_PROGRESS TRUE ) endif() @@ -119,13 +121,15 @@ if(DEEPGEMM_ARCHS) message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") # add_custom_command does no implicit header scanning; glob explicitly so - # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild. + # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild. file(GLOB_RECURSE _dg_headers "${deepgemm_SOURCE_DIR}/csrc/*.h" "${deepgemm_SOURCE_DIR}/csrc/*.hpp" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp" - "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh") + "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp") set(_dg_markers) set(_dg_seen_soabis) diff --git a/docker/Dockerfile b/docker/Dockerfile index 9b344552c5ae..b3e32777e7bb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ git \ curl \ sudo \ + elfutils-devel \ rdma-core-devel \ && dnf clean all \ && rm -rf /var/cache/dnf diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 67a527405e2d..f4438c67552f 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -50,7 +50,7 @@ str(src / "deep_gemm/include"), str(src / "third-party/cutlass/include"), str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/fmt/include"), + str(src / "third-party/deep_jit/include"), *cpp_extension.include_paths(device_type="cuda"), ] diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index f6a64d8796c1..db81bf4be310 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,11 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -# NOTE: This is currently targeting the nv_dev branch tip due to sm120 support -DEEPGEMM_GIT_REF="8b1392b978f5a03c828dd1711090d7fb50958b8a" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" +# NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 +# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE +# activations +DEEPGEMM_GIT_REF="ad1f1726aa540a76c1d26d6a120effb8de21eaa4" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c4f0092f8d9b..6a26ce78b4c4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -532,8 +532,10 @@ def forward( symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - situ_beta=self.activation_beta, - situ_linear_beta=self.activation_linear_beta, + # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the + # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation_alpha=self.activation_beta or 1.0, + activation_beta=self.activation_linear_beta or 0.0, fast_math=fast_math, ) return y diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3fc70c0f5da3..d515d5131e75 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -280,7 +280,11 @@ def _lazy_init() -> None: _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None) _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None) _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None) - _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) + # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so + # an externally pinned older DeepGEMM still resolves. + _grouped_masked_impl = getattr( + _dg, "m_grouped_fp8_gemm_nt_masked", None + ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`. From 2e2fdaa99a512b58e6fa3b9fd24a07f7699d694c Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 01:00:40 -0700 Subject: [PATCH 243/774] [CI] Extend Arm CPU kernel shard timeout (#56601) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 052063f2f69e..77e0f3e07bad 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -112,7 +112,9 @@ print_packages export CONTAINER_NAME export -f kernel_tests model_tests serving_tests case "$SHARD_ID" in - 0) timeout 30m bash -c kernel_tests ;; + # The kernel shard contains several large parametrized suites and can exceed + # 30 minutes on otherwise healthy Arm runners. + 0) timeout 60m bash -c kernel_tests ;; 1) timeout 30m bash -c model_tests ;; 2) timeout 30m bash -c serving_tests ;; esac From ca67438c088d5ada7da901b77f9aeaebc1b9542c Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Tue, 15 Sep 2026 01:17:30 -0700 Subject: [PATCH 244/774] [Bugfix][Frontend] Handle aborted requests in beam search (#56249) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- tests/samplers/test_beam_search.py | 85 +++++++++++++++++++ tests/samplers/test_beam_search_online.py | 54 +++++++++++- .../generate/beam_search/offline.py | 14 ++- .../generate/beam_search/online.py | 15 ++-- .../entrypoints/generate/beam_search/utils.py | 3 +- 5 files changed, 162 insertions(+), 9 deletions(-) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 91dee5634b52..8bcf4c48ab59 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -6,13 +6,17 @@ """ import json +from types import SimpleNamespace +from unittest.mock import Mock import jsonschema import pytest from transformers import AutoModelForSeq2SeqLM +from vllm import CompletionOutput, RequestOutput from vllm.assets.audio import AudioAsset from vllm.entrypoints.llm import LLM +from vllm.logprobs import Logprob, SampleLogprobs from vllm.platforms import current_platform from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams @@ -45,6 +49,87 @@ MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"] +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Return tokens and scores from the last completed step for aborted prompts. + + Other prompts in the batch continue generating. + """ + + def run_requests(prompts, **kwargs): + results = [] + for prompt in prompts: + tokens = prompt["prompt_token_ids"] + if tokens[0] == prompt_token: + assert len(tokens) <= abort_after + 1, "Continued after abort" + aborted = ( + tokens[0] == prompt_token + and len(tokens) == abort_after + 1 + and tokens[-1] != 12 + ) + results.append( + RequestOutput( + request_id="inner", + prompt=None, + prompt_token_ids=tokens, + prompt_logprobs=None, + finished=True, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[] if aborted else [11], + cumulative_logprob=None, + logprobs=terminal_logprobs + if aborted + else [{11: Logprob(-0.1), 12: Logprob(-0.2)}], + finish_reason="abort" if aborted else "length", + ) + ], + ) + ) + return results + + llm = LLM.__new__(LLM) + llm.llm_engine = Mock() + tokenizer = SimpleNamespace( + eos_token_id=0, + decode=lambda tokens, skip_special_tokens=False: " ".join( + str(token) for token in tokens if not skip_special_tokens or token != 0 + ), + ) + llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer)) + monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts) + monkeypatch.setattr(llm, "_render_and_run_requests", run_requests) + prompts = [ + {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2] + ] + params = BeamSearchParams(beam_width=2, max_tokens=3) + outputs = llm.beam_search(prompts, params) + + aborted, normal = outputs + expected_tokens = ( + [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]] + ) + assert [beam.tokens for beam in aborted.sequences] == expected_tokens + for beam in aborted.sequences: + assert beam.finish_reason == "abort" + assert beam.text == tokenizer.decode( + beam.tokens, skip_special_tokens=params.skip_special_tokens + ) + assert len(beam.logprobs) == abort_after + assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after) + assert len(normal.sequences) == 2 + assert all(len(beam.tokens) == 4 for beam in normal.sequences) + assert all(beam.finish_reason != "abort" for beam in normal.sequences) + + @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", MAX_TOKENS) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 102e66dcbe79..ccb40cc65906 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -5,7 +5,7 @@ from vllm import CompletionOutput, RequestOutput from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin -from vllm.logprobs import Logprob +from vllm.logprobs import Logprob, SampleLogprobs from vllm.sampling_params import BeamSearchParams @@ -108,3 +108,55 @@ async def test_beam_search_respects_skip_special_tokens( assert outputs[0].outputs[0].text == expected_text assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +async def test_beam_search_abort_returns_partial_outputs( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Abort returns each beam's tokens and scores from the last completed step.""" + calls = 0 + + async def generate(prompt, *args, **kwargs): + nonlocal calls + calls += 1 + depth = len(prompt["prompt_token_ids"]) - 1 + assert depth <= abort_after, "Beam search continued after abort" + result = await anext(_EngineClient().generate(prompt, *args, **kwargs)) + output = result.outputs[0] + output.token_ids = [11] + output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}] + output.finish_reason = "length" + if depth == abort_after and prompt["prompt_token_ids"][-1] != 12: + output.token_ids = [] + output.logprobs = terminal_logprobs + output.finish_reason = "abort" + yield result + + serving = _Serving() + monkeypatch.setattr(serving.engine_client, "generate", generate) + prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]} + outputs = [ + output + async for output in serving.beam_search( + prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4) + ) + ] + + assert calls == 1 + 2 * abort_after + assert len(outputs) == 1 + assert outputs[0].finished + assert outputs[0].request_id == "request" + expected_tokens = [[11], [12]] if abort_after else [[]] + assert [output.token_ids for output in outputs[0].outputs] == expected_tokens + for output in outputs[0].outputs: + assert output.finish_reason == "abort" + assert output.text == _Tokenizer().decode(list(output.token_ids)) + assert output.logprobs is not None + assert len(output.logprobs) == abort_after + assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after) diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index b919435f65a1..73b9ebe06b8f 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -293,6 +293,18 @@ def _beam_search_step( allowed_sets[i] = set(entry[1]) for (start, end), instance in zip(instance_start_and_end, instances_batch): + instance_output = output[start:end] + if any( + result is not None and result.outputs[0].finish_reason == "abort" + for result in instance_output + ): + for beam, result in zip(all_beams[start:end], instance_output): + if result is not None: + beam.finish_reason = "abort" + instance.completed.append(beam) + instance.beams = [] + continue + instance_new_beams = [] for i in range(start, end): current_beam = all_beams[i] @@ -302,8 +314,6 @@ def _beam_search_step( continue if result.outputs[0].logprobs is not None: - # if logprobs is None, the sequence completed - # due to max-model-len or abortion. logprobs = result.outputs[0].logprobs[0] allowed = allowed_sets[i] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index ac1a585e8f62..da02b974c0a5 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -99,11 +99,7 @@ async def beam_search( output = [x[0] for x in await asyncio.gather(*tasks)] - candidates = [] - # Iterate through all beam inference results - for i, result in enumerate(output): - current_beam = all_beams[i] - + for result in output: # check for error finish reason and abort beam search if result.outputs[0].finish_reason == "error": # yield error output and terminate beam search @@ -126,6 +122,15 @@ async def beam_search( ) return + if any(result.outputs[0].finish_reason == "abort" for result in output): + for beam in all_beams: + beam.finish_reason = "abort" + break + + candidates = [] + for i, result in enumerate(output): + current_beam = all_beams[i] + if result.outputs[0].logprobs is not None: logprobs = result.outputs[0].logprobs[0] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 47f69cbf7abe..61f97963b6f3 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -150,7 +150,8 @@ def get_beam_search_score( if tokens[-1] == eos_token_id: seq_len -= 1 - return cumulative_logprob / (seq_len**length_penalty) + # An aborted beam may contain only an EOS prompt token. + return cumulative_logprob / (max(seq_len, 1) ** length_penalty) def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): From 676650397ff71cd802e7e0e30c2e46a215ab86ee Mon Sep 17 00:00:00 2001 From: Qiu Chunshuo Date: Tue, 15 Sep 2026 16:27:46 +0800 Subject: [PATCH 245/774] [Feature][PCP] Support decode-only FULL CUDA graphs (#53867) Signed-off-by: QiuChunshuo Signed-off-by: QiuChunshuo Co-authored-by: Codex Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_pcp_manager.py | 159 +++++++++++++++++++++++- vllm/v1/worker/gpu/cudagraph_utils.py | 9 +- vllm/v1/worker/gpu/input_batch.py | 8 +- vllm/v1/worker/gpu/model_runner.py | 35 +++--- vllm/v1/worker/gpu/pcp_manager.py | 114 ++++++++++++++--- 5 files changed, 283 insertions(+), 42 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index dc935929d4c4..d3ec562aa7a6 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock import numpy as np import pytest import torch +from vllm.config import CUDAGraphMode from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -20,6 +24,58 @@ def _copy_to_cpu(value, out=None, device=None): return tensor +def _make_config(cudagraph_mode: CUDAGraphMode): + return SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=2, + decode_context_parallel_size=1, + pipeline_parallel_size=1, + dcp_comm_backend="ag_rs", + ), + model_config=SimpleNamespace( + use_mla=True, + is_encoder_decoder=False, + hf_text_config=SimpleNamespace(), + ), + lora_config=None, + speculative_config=None, + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + ) + + +def _make_capture_manager(block_table: torch.Tensor): + block_tables = SimpleNamespace( + input_block_tables=(block_table,), + num_kv_cache_groups=1, + kernel_block_sizes=(2,), + blocks_per_kv_block=(1,), + ) + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + max_num_reqs=4, + max_num_tokens=8, + block_tables=block_tables, + ) + return manager, block_tables + + +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode): + PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False) + + +def test_validate_config_rejects_full_graph_for_prefills(): + with pytest.raises(NotImplementedError, match="decode-only routines"): + PCPManager.validate_config( + _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False + ) + + def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager = PCPManager( pcp_world_size=2, @@ -114,6 +170,89 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +@pytest.mark.parametrize( + ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"), + [ + (CUDAGraphMode.NONE, 4, None, None), + (CUDAGraphMode.PIECEWISE, None, 8, None), + (CUDAGraphMode.FULL, 4, 8, 4), + ], +) +def test_partition_padding_is_derived_from_batch_descriptor( + cg_mode, num_reqs, expected_tokens, expected_reqs +): + manager = MagicMock() + input_batch = MagicMock() + manager.partition_batch.return_value = input_batch + batch_desc = BatchExecutionDescriptor( + cg_mode=cg_mode, + num_tokens=8, + num_reqs=num_reqs, + ) + + result = pcp_manager_module.maybe_partition_pcp_batch( + manager, + input_batch, + batch_desc, + ) + + assert result is input_batch + manager.partition_batch.assert_called_once_with( + input_batch, + padded_num_tokens=expected_tokens, + padded_num_reqs=expected_reqs, + ) + + +def test_capture_uses_pcp_persistent_inputs(): + manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) + + dummy_batch = InputBatch.make_dummy( + num_reqs=4, + num_tokens=4, + input_buffers=InputBuffers(4, 8, torch.device("cpu")), + max_query_len=1, + ) + dummy_batch.input_ids.copy_(torch.arange(4)) + dummy_batch.positions.fill_(7) + input_batch = manager.prepare_inputs_to_capture(dummy_batch) + + assert input_batch is not dummy_batch + assert input_batch.req_ids == dummy_batch.req_ids + for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"): + actual = getattr(input_batch, name) + torch.testing.assert_close(actual, getattr(dummy_batch, name)) + assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr() + + +def test_dummy_context_updates_pcp_local_block_tables(): + global_block_table = torch.full((4, 4), -1, dtype=torch.int32) + manager, block_tables = _make_capture_manager(global_block_table) + input_batch = InputBatch.make_dummy( + num_reqs=2, + num_tokens=2, + input_buffers=manager.input_buffers, + max_query_len=1, + ) + input_batch = manager.prepare_inputs_to_capture(input_batch) + local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs) + + set_dummy_context( + input_batch, + block_tables, + context_len=3, + num_kv_blocks=16, + max_model_len=16, + input_block_tables=local_block_tables, + ) + + torch.testing.assert_close( + local_block_tables[0][:2, :2], + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + ) + assert torch.all(global_block_table == -1) + + def _rank_rows( pcp_rank: int, pcp_world_size: int, @@ -331,10 +470,22 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2] global_batch.dcp_local_seq_lens.fill_(-1) - local_batch = manager.partition_batch(global_batch, padded_num_tokens=2) + local_batch = manager.partition_batch( + global_batch, + padded_num_tokens=4, + padded_num_reqs=4, + ) assert local_batch.dcp_local_seq_lens is None - assert local_batch.seq_lens.tolist() == [17, 25] + assert local_batch.seq_lens.tolist() == [17, 25, 0, 0] + assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None + expected_dcp_upper_bound = get_dcp_local_seq_lens( + torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + ) + assert torch.equal( + local_batch.dcp_local_seq_lens_cpu_upper_bound, + torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]), + ) # What execute_model does next: derive DCP metadata from the final batch # on the PCP-owned buffers. @@ -348,7 +499,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): num_reqs_padded=local_batch.num_reqs_after_padding, ) expected = get_dcp_local_seq_lens( - torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1 ) assert local_batch.dcp_local_seq_lens is not None assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d197dbbc6df1..9657c28fe601 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -777,11 +777,12 @@ def prepare_inputs_to_capture( input_batch = InputBatch.make_dummy( num_reqs, num_tokens, input_buffers, max_query_len=max_query_len ) - input_block_tables = block_tables.get_dummy_block_tables(num_reqs) - slot_mapping_provider: BlockTables | PCPManager = block_tables if pcp_manager is not None: - slot_mapping_provider = pcp_manager - slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens) + input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + + block_table_provider = pcp_manager or block_tables + input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) + slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, kv_cache_config ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5cbd642d2529..d3c65bd06f11 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING @@ -217,9 +218,12 @@ def set_dummy_context( context_len: int, num_kv_blocks: int, max_model_len: int, + input_block_tables: Sequence[torch.Tensor] | None = None, ) -> None: """Give each dummy request context_len of context, used when profiling step cost.""" - if not block_tables.input_block_tables: + if input_block_tables is None: + input_block_tables = block_tables.input_block_tables + if not input_block_tables: # Attention-free models have no KV context to fabricate. return num_reqs = input_batch.num_reqs @@ -241,7 +245,7 @@ def set_dummy_context( seq_len = context_len + query_len for block_table, block_size, bpk in zip( - block_tables.input_block_tables, + input_block_tables, block_tables.kernel_block_sizes, block_tables.blocks_per_kv_block, ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 761893303a5c..25395c3330eb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1468,7 +1468,7 @@ def prepare_inputs( input_batch = pcp.maybe_partition_pcp_batch( self.pcp_manager, input_batch, - padded_num_tokens=batch_desc.num_tokens, + batch_desc, ) return input_batch @@ -1496,7 +1496,8 @@ def prepare_attn( def prepare_dummy_attn( self, input_batch: InputBatch, valid_state_slots: bool = False ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: - block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + block_table_provider = self.pcp_manager or self.block_tables + block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs) if valid_state_slots: state_slots = torch.arange( 1, @@ -1745,25 +1746,29 @@ def execute_model( self.input_buffers, max_query_len=batch_desc.max_query_len, ) - if not skip_attn_for_dummy_run: - block_tables, slot_mappings = self.prepare_dummy_attn( - input_batch, valid_dummy_state_slots - ) - if context_len: - set_dummy_context( - input_batch, - self.block_tables, - context_len, - self.kv_cache_config.num_blocks, - self.max_model_len, - ) - else: + if self.pcp_manager is not None: + input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) + if skip_attn_for_dummy_run: assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( "Attention metadata must be prepared for dummy runs when using " "FULL cudagraph mode." ) block_tables = None slot_mappings = None + else: + block_tables, slot_mappings = self.prepare_dummy_attn( + input_batch, valid_dummy_state_slots + ) + if not skip_attn_for_dummy_run and context_len: + assert block_tables is not None + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + input_block_tables=block_tables, + ) attn_metadata = None slot_mappings_by_layer = None diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 3a51532f57a1..adfbf2f9f0ba 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterator from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np import torch @@ -17,6 +18,9 @@ InputBuffers, ) +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + logger = init_logger(__name__) @@ -169,8 +173,14 @@ def validate_config( "MRV2 sparse MLA PCP does not support CUDA graphs yet. " "Set -cc.cudagraph_mode=NONE." ) - if cudagraph_mode.has_full_cudagraphs(): - raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + raise NotImplementedError( + "MRV2 PCP supports full CUDA graphs for decode-only routines. " + "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE." + ) if ( parallel_config.decode_context_parallel_size > 1 and parallel_config.dcp_comm_backend != "ag_rs" @@ -375,6 +385,22 @@ def get_num_tokens_for_dispatch( for rank in range(self.pcp_world_size) ) + @staticmethod + def _resolve_num_reqs_after_padding( + input_batch: InputBatch, + padded_num_reqs: int | None, + num_local_reqs: int, + ) -> int: + if padded_num_reqs is None: + return num_local_reqs + if input_batch.has_prefill: + raise RuntimeError("PCP FULL graphs require a decode-only batch.") + assert padded_num_reqs >= num_local_reqs, ( + "PCP graph request capacity must cover the rank-local batch: " + f"{padded_num_reqs} < {num_local_reqs}." + ) + return padded_num_reqs + @property def input_buffers(self) -> InputBuffers: assert self._input_buffers is not None @@ -384,6 +410,7 @@ def partition_batch( self, input_batch: InputBatch, padded_num_tokens: int | None = None, + padded_num_reqs: int | None = None, ) -> InputBatch: assert self._input_buffers is not None input_buffers = self._input_buffers @@ -414,10 +441,15 @@ def partition_batch( ] num_local_reqs = len(local_segments) - if num_local_reqs > input_buffers.max_num_reqs: + num_reqs_after_padding = self._resolve_num_reqs_after_padding( + global_batch, + padded_num_reqs, + num_local_reqs, + ) + if num_reqs_after_padding > input_buffers.max_num_reqs: raise RuntimeError( - "PCP local request count exceeds the MRV2 input buffer size: " - f"{num_local_reqs} > {input_buffers.max_num_reqs}." + "PCP padded local request count exceeds the MRV2 input buffer size: " + f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}." ) local_to_global_batch_req_idx_np = np.fromiter( @@ -505,21 +537,25 @@ def partition_batch( np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) - local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1] + local_query_start_loc = input_buffers.query_start_loc[ + : num_reqs_after_padding + 1 + ] local_to_global_req_idx = async_copy_to_gpu( local_to_global_req_idx_np, device=self.device ) - seq_lens = input_buffers.seq_lens[:num_local_reqs] + seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] + real_seq_lens = seq_lens[:num_local_reqs] if num_local_tokens > 0: local_end_positions = torch.index_select( input_buffers.positions, 0, - local_query_start_loc[1:] - 1, + local_query_start_loc[1 : num_local_reqs + 1] - 1, ) - seq_lens.copy_(local_end_positions + 1) + real_seq_lens.copy_(local_end_positions + 1) else: - seq_lens.zero_() + real_seq_lens.zero_() + seq_lens[num_local_reqs:].zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -552,11 +588,15 @@ def partition_batch( local_num_computed_prefill_tokens_np = np.minimum( local_start_pos_np, local_prefill_len_np ) - local_is_prefilling_np = ( + real_local_is_prefilling_np = ( local_num_computed_prefill_tokens_np < local_prefill_len_np ) - seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) - seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens + local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_) + local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32) + seq_lens_cpu_upper_bound_np[:num_local_reqs] = ( + local_start_pos_np + local_num_scheduled_tokens + ) dcp_local_seq_lens_cpu_upper_bound = None if self.dcp_world_size > 1: # The largest DCP shard of each row's whole request, identical on @@ -564,18 +604,24 @@ def partition_batch( request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ local_to_global_batch_req_idx_np ] - dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( torch.from_numpy(request_seq_lens.astype(np.int32)), self.dcp_world_size, 0, self.cp_interleave, ) + dcp_local_seq_lens_cpu_upper_bound = torch.zeros( + num_reqs_after_padding, dtype=torch.int32 + ) + dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_( + real_dcp_local_seq_lens_cpu_upper_bound + ) self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, - num_reqs_after_padding=num_local_reqs, + num_reqs_after_padding=num_reqs_after_padding, idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, @@ -588,7 +634,7 @@ def partition_batch( num_draft_tokens=0, num_draft_tokens_per_req=None, query_start_loc=local_query_start_loc, - query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1], + query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1], seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, @@ -608,6 +654,31 @@ def partition_batch( ) return self._local_batch + def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch: + """Stage a capture or dummy batch in persistent PCP input buffers.""" + input_buffers = self.input_buffers + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + input_batch = replace( + input_batch, + input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids), + positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions), + is_padding=input_buffers.is_padding[:num_tokens].copy_( + input_batch.is_padding + ), + query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_( + input_batch.query_start_loc + ), + seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens), + ) + return input_batch + + def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: + assert self._local_block_tables is not None + return tuple( + block_table[:num_reqs].zero_() for block_table in self._local_block_tables + ) + def prepare_attn( self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: @@ -728,13 +799,22 @@ def restore_for_sampling( def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, - padded_num_tokens: int | None = None, + batch_desc: "BatchExecutionDescriptor", ) -> InputBatch: if manager is None: return input_batch + + padded_num_tokens = None + padded_num_reqs = None + if batch_desc.cg_mode != CUDAGraphMode.NONE: + padded_num_tokens = batch_desc.num_tokens + if batch_desc.cg_mode == CUDAGraphMode.FULL: + padded_num_reqs = batch_desc.num_reqs + return manager.partition_batch( input_batch, padded_num_tokens=padded_num_tokens, + padded_num_reqs=padded_num_reqs, ) From 2c2cbd84fdf199901f3f4e6b0e708ee0664aadd6 Mon Sep 17 00:00:00 2001 From: Yixin Dong Date: Tue, 15 Sep 2026 04:36:16 -0400 Subject: [PATCH 246/774] [Frontend] Use XGrammar schema constraints for DeepSeek V4.1 (#56408) Signed-off-by: Ubospica Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structural_tag_registry.py | 80 ++++++++++++++++++- vllm/tool_parsers/structural_tag_registry.py | 14 +++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index b77b4efaffe6..88e1f3f3832c 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock import pytest -from xgrammar import Grammar, StructuralTag +from xgrammar import Grammar, StructuralTag, builtin_structural_tag from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -81,10 +81,14 @@ def test_supported_structural_tag_models_include_vllm_builtins(): @pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) -def test_deepseek_v41_format_constraints_ignore_parameter_schema( +def test_deepseek_v41_older_xgrammar_keeps_format_constraints( choice, sample_tools_strict, + monkeypatch, ): + monkeypatch.setattr( + builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None, raising=False + ) tool_choice = ( ChatCompletionNamedToolChoiceParam( function=ChatCompletionNamedFunction(name=choice) @@ -128,7 +132,11 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): reasoning=False, ) grammar = Grammar.from_structural_tag(tag) - call = '<｜DSML｜ invoke name="get_weather">\n\n' + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="city" string="true">Paris\n' + "\n" + ) assert _is_grammar_accept_string( grammar, "\n\n<｜DSML｜ calls>\n" + call + "" ) @@ -141,6 +149,72 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): ) +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) +@pytest.mark.parametrize("reasoning", [False, True]) +def test_deepseek_v41_constrains_parameters_after_reasoning( + sample_tools_strict, choice, reasoning +): + """Strict calls must enforce the schema without requiring another think close.""" + tool_choice = ( + ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=choice) + ) + if choice == "get_weather" + else choice + ) + sample_tools_strict[0].function.parameters["additionalProperties"] = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools_strict, tool_choice, reasoning=reasoning + ) + grammar = Grammar.from_structural_tag(tag) + begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n' + end = "\n" + parameter = ( + '<｜DSML｜ parameter name="city" string="true">Paris\n' + ) + assert _is_grammar_accept_string(grammar, begin + parameter + end) + for invalid in ( + "", # missing required city + parameter * 2, + parameter.replace('name="city"', 'name="unknown"'), + parameter.replace('string="true">Paris', 'string="false">42'), + parameter.replace("｜ parameter", "｜parameter"), + '{"city":"Paris"}', # the encoder emits DSML parameters, not a JSON body + ): + assert not _is_grammar_accept_string(grammar, begin + invalid + end) + assert not _is_grammar_accept_string( + grammar, "reason" + begin + parameter + end + ) + assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto") + + +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools): + sample_tools[0].function.strict = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools, "required", reasoning=False + ) + grammar = Grammar.from_structural_tag(tag) + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="extra" string="false">' + '{"nested":[true,null,1.5]}\n' + "\n" + ) + output = "\n\n<｜DSML｜ calls>\n" + call * 2 + "" + assert _is_grammar_accept_string(grammar, output) + assert not _is_grammar_accept_string( + grammar, output.replace('{"nested":[true,null,1.5]}', "invalid") + ) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_all_xgrammar_builtins( model: str, diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 34e57840bed7..989250e2ea5b 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -10,6 +10,7 @@ from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from openai.types.responses.tool_choice_function import ToolChoiceFunction from xgrammar import StructuralTag, normalize_tool_choice +from xgrammar import builtin_structural_tag as xgrammar_builtin_structural_tag from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag from xgrammar.openai_tool_call_schema import ( BuiltinToolParam, @@ -242,9 +243,20 @@ def get_deepseek_v41_structural_tag( token_suffix: str = "", ) -> StructuralTag: # Serving enables this visible-text grammar after the reasoning boundary. + builder = getattr( + xgrammar_builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None + ) + if builder is not None: + return builder( + tools=tools, + builtin_tools=builtin_tools, + tool_choice=tool_choice, + reasoning="disabled", + ) + del builtin_tools, reasoning, token_suffix - # TODO: lower parameter schemas into DSML constraints. This builder constrains + # Compatibility with xgrammar releases without the V4.1 builtin. This constrains # tool names and DSML/value syntax only. Parameter names, presence, uniqueness # and value schemas remain unconstrained, including for strict=true. The request # layer still uses strict to decide whether auto tool choice activates a grammar. From 2983fc62869de7943bd68ea58730eec7a39ff597 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 02:00:58 -0700 Subject: [PATCH 247/774] [CI] Fix NIXL push worker test stub after failure deferral (#56955) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/v1/kv_connector/unit/test_nixl_push_connector.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 423b0e9deca3..913ed0dcba0d 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -351,6 +351,7 @@ def fresh(cls) -> _StubWriterWorker: # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} w._pending_recv_notifs = {} + w._failed_recv_reqs = queue.Queue() w._recv_failures = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() @@ -455,12 +456,14 @@ def test_sends_to_all_agents_with_prefix_and_payload(self): list(group) for group in rd["local_block_ids"] ] - def test_fails_request_when_engine_not_handshaked(self): + def test_queues_failed_request_when_engine_not_handshaked(self): w, failures = self._reg_send_worker() w._do_send_reg_notif("req-1", _registration_data("req-1")) - assert failures == [("req-1", None)] + assert failures == [] + assert w._failed_recv_reqs.get_nowait() == "req-1" + assert w._failed_recv_reqs.empty() w.nixl_wrapper.send_notif.assert_not_called() def test_agent_read_is_serialized_with_handshake_lock(self): @@ -886,7 +889,6 @@ def _pollable_worker() -> _StubWriterWorker: w.nixl_wrapper = MagicMock() w.xfer_stats = MagicMock() w._log_failure = MagicMock() # type: ignore[method-assign] - w._failed_recv_reqs = queue.Queue() w._invalid_block_ids = queue.Queue() w._pending_recv_notifs = {} w._replicated_pcp_done_sending = set() From abb4cb9fc5f1b1039876ed8a657b025f65643d11 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:58:42 -0400 Subject: [PATCH 248/774] =?UTF-8?q?[CI]=20Shard=20V1=20KV=20Connectors=201?= =?UTF-8?q?=E2=86=924=20(#56324)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Thang Nguyen Co-authored-by: Thang Nguyen Co-authored-by: Kimi Code --- .buildkite/test_areas/misc.yaml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index ed82c7d71d68..940447bf1b52 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -142,19 +142,20 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors" +- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors Shard %N" device: h200_35gb key: v1-kv-connectors timeout_in_minutes: 45 + parallelism: 4 source_file_dependencies: *v1-core-kv-metrics-deps commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: - label: ":amd: (MI300) V1 KV Connectors" + label: ":amd: (MI300) V1 KV Connectors Shard %N" dind: false device: mi300_1 timeout_in_minutes: 45 From 073f883b56a5b585366da874473db100bc6fff7e Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 02:58:57 -0700 Subject: [PATCH 249/774] [Perf][DSpark] Collapse DeepSeek-V4.1 draft states before SP all-gather (#56903) Signed-off-by: Juntian Liu Co-authored-by: Codex Co-authored-by: Woosuk Kwon --- tests/kernels/test_mhc_kernels.py | 21 +++++++++++++++++++++ vllm/models/deepseek_v41/nvidia/dspark.py | 5 ++--- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index f805fe39a57b..ff72e2a96c17 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -147,6 +147,27 @@ def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("num_tokens", [1, 7, 32]) +@pytest.mark.parametrize("tp_size", [4, 8]) +def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size): + """Simulate SP on one GPU, including padding and ranks without valid tokens.""" + set_random_seed(0) + padded_tokens = num_tokens + (-num_tokens) % tp_size + streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + pre_mix = torch.rand(padded_tokens, 4, device=DEVICE) + stream_shards = streams.chunk(tp_size) + mix_shards = pre_mix.chunk(tp_size) + before = hc_collapse_triton( + torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens] + ) + after = torch.cat( + [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)] + )[:num_tokens] + assert after.shape == (num_tokens, 5120) + torch.testing.assert_close(after, before, atol=0, rtol=0) + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize( "num_tokens,sinkhorn_repeat,norm_eps", diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 88646404f44e..c19dca800555 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -217,15 +217,14 @@ def forward( residual, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) - if self.use_sequence_parallel: - hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] - pre_mix = sp_all_gather(pre_mix)[:full_num_tokens] # Collapse the hc copies with the pre-mix from the last layer's FFN # mixes — the mix the reference forward_head applies (hc_pre with the # last block's ffn pre-mix). Return the PRE-norm head hidden; # compute_logits applies self.norm. assert pre_mix is not None hidden_states = hc_collapse_triton(hidden_states, pre_mix) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] return hidden_states From 142020c6c2c0d4fec4d36df28290e3d9b654a385 Mon Sep 17 00:00:00 2001 From: Yeonwoo Sung Date: Tue, 15 Sep 2026 19:03:26 +0900 Subject: [PATCH 250/774] [Bugfix][Rust Frontend] Honor `add_generation_prompt` in DeepSeek V3.2/V4/V4.1 renderers (#56593) Co-authored-by: Bugen Zhao Signed-off-by: YeonwooSung Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/deepseek.rs | 47 ++++++---- .../src/renderer/deepseek_v32/encoding.rs | 18 +++- .../chat/src/renderer/deepseek_v32/tests.rs | 38 ++++++++ .../chat/src/renderer/deepseek_v4/tests.rs | 91 +++++++++++++++++++ .../chat/src/renderer/deepseek_v41/tests.rs | 47 +++++++++- 5 files changed, 217 insertions(+), 24 deletions(-) diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5925c5fd1655..01576a7c4341 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -224,21 +224,30 @@ pub(super) fn render_request_with_media_order( } } - if (is_user_like_entry(message, current_render_index as usize, dialect) - || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. }))) - && next_rendered_entry_is_assistant_or_end( + if is_user_like_entry(message, current_render_index as usize, dialect) + || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. })) + { + // Skip the opener only at the last rendered entry when + // add_generation_prompt is off. Inspect the next rendered + // entry, not the raw index: V4 merges trailing tool results. + let should_write_transition = match next_rendered_entry( request.messages.as_slice(), message_index, drop_historical_developers, last_user_like_message_index, - ) - { - write_assistant_transition( - &mut out, - thinking_mode, - drop_thinking, - current_render_index >= last_user_render_index, - ); + ) { + Some(ChatMessage::Assistant { .. }) => true, + None => request.chat_options.add_generation_prompt(), + Some(_) => false, + }; + if should_write_transition { + write_assistant_transition( + &mut out, + thinking_mode, + drop_thinking, + current_render_index >= last_user_render_index, + ); + } } } @@ -420,14 +429,17 @@ fn is_historical_developer( && last_user_like_message_index.is_some_and(|last_index| message_index < last_index) } -/// Return whether the next rendered entry is assistant, or there is no next -/// entry. -fn next_rendered_entry_is_assistant_or_end( +/// Return the next rendered message after `message_index`. +/// +/// Skips merged user-content tails and dropped historical developers so the +/// caller can tell a following assistant turn apart from the end of the +/// conversation. +fn next_rendered_entry( messages: &[ChatMessage], message_index: usize, drop_historical_developers: bool, last_user_like_message_index: Option, -) -> bool { +) -> Option<&ChatMessage> { let mut next_index = message_index + 1; while next_index < messages.len() && (is_following_user_content(messages, next_index) @@ -441,10 +453,7 @@ fn next_rendered_entry_is_assistant_or_end( next_index += 1; } - messages - .get(next_index) - .map(|message| matches!(message, ChatMessage::Assistant { .. })) - .unwrap_or(true) + messages.get(next_index) } /// Render the tool preamble shown to the model for one DeepSeek dialect. diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 3ccf4d11db7f..4fc5f26530f8 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -61,6 +61,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { find_last_user_render_index(request.messages.as_slice(), render_offset); let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice()); let continue_final_message = request.chat_options.continue_final_message(); + let add_generation_prompt = request.chat_options.add_generation_prompt(); let mut prompt = String::from(BOS_TOKEN); if !initial_tools.is_empty() { @@ -79,6 +80,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { thinking_mode, drop_thinking, continue_final_message, + add_generation_prompt, )?; } @@ -110,11 +112,13 @@ fn render_message( thinking_mode: ThinkingMode, drop_thinking: bool, continue_final_message: bool, + add_generation_prompt: bool, ) -> Result<()> { let render_index = message_index as isize + render_offset; let opens_thinking = render_index == last_user_render_index; let after_last_user_turn = render_index > last_user_render_index; let after_or_at_last_user_turn = render_index >= last_user_render_index; + let add_generation_prompt = message_index + 1 != messages.len() || add_generation_prompt; match message { ChatMessage::System { content } => render_system_message(out, Some(content), &[]), @@ -123,11 +127,13 @@ fn render_message( content, tools.as_deref().unwrap_or(&[]), thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::User { content } => render_user_message( out, content, thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::Assistant { content } => render_assistant_message( out, @@ -195,6 +201,7 @@ fn render_developer_message( content: &ChatContent, tools: &[ChatTool], opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -209,7 +216,7 @@ fn render_developer_message( } out.push_str("\n\n# The user's message is: "); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } @@ -219,17 +226,20 @@ fn render_user_message( out: &mut String, content: &ChatContent, opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { out.push_str("<｜User｜>"); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } /// Shared trailing wrapper used by both real user turns and native developer /// turns after their content has already been written. -// TODO: respect `add_generation_prompt` option -fn write_user_like_suffix(out: &mut String, opens_thinking: bool) { +fn write_user_like_suffix(out: &mut String, opens_thinking: bool, add_generation_prompt: bool) { + if !add_generation_prompt { + return; + } out.push_str("<｜Assistant｜>"); if opens_thinking { out.push_str(THINKING_START_TOKEN); diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 3f93e3653efd..56c03a92cab4 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -281,6 +281,44 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() { .assert_eq(&error.to_report_string()); } +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>"] + .assert_eq(&rendered); +} + #[test] fn continue_final_assistant_omits_final_eos() { let mut request = ChatRequest { diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 172b67463c42..9404b54296a7 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -633,3 +633,94 @@ fn continue_final_assistant_omits_final_eos() { expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜>partial answer"] .assert_eq(&rendered); } + +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), Value::Bool(false)); + + let rendered = render_request(&request); + + expect![ + "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>" + ] + .assert_eq(&rendered); +} + +#[test] +fn trailing_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("rules"), + ChatMessage::user("question"), + ChatMessage::system("final rules"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), Value::Bool(false)); + request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜>rules<｜User｜>questionfinal rules"].assert_eq(&rendered); +} + +#[test] +fn last_merged_tool_result_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ + ChatMessage::user("question"), + ChatMessage::tool_response("result", "call-1"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question\n\nresult"] + .assert_eq(&rendered); + assert!(!rendered.contains("<｜Assistant｜>")); +} + +#[test] +fn last_developer_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::developer("latest instruction", None)]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered); +} diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index 3babb5c269f8..3d621c2246fd 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -10,7 +10,9 @@ use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, +}; fn render(request: &ChatRequest) -> String { DeepSeekV41ChatRenderer::new() @@ -281,6 +283,49 @@ fn mid_system_advances_last_user_boundary_and_drops_old_reasoning() { .assert_eq(&render(&request)); } +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question"] + .assert_eq(&render(&request)); +} + +#[test] +fn last_mid_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("sys"), + ChatMessage::user("q1"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { text: "r1".into() }, + AssistantContentBlock::Text { text: "a1".into() }, + ]), + ChatMessage::developer("old policy", None), + ChatMessage::system("mid sys"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("reasoning_effort".into(), json!(88)); + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜>a1<｜end▁of▁sentence｜><｜System｜>mid sys"] + .assert_eq(&render(&request)); +} + +#[test] +fn chat_mode_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request.chat_options.template_kwargs.insert("thinking".into(), json!(false)); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question"].assert_eq(&render(&request)); +} + #[test] fn joins_text_parts_with_reference_separator_and_preserves_literal_tags() { let mut request = ChatRequest { From 24bfbd5d4d530a2657a8925e5d12e86b1ccb9cca Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Tue, 15 Sep 2026 06:07:58 -0400 Subject: [PATCH 251/774] [Refactor] Derive is_reasoning_end from the engine grammar (#56200) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../engine/test_gemma4_streaming_reasoning.py | 33 +++++++++++++++ tests/parser/engine/test_qwen3_reasoning.py | 42 +++++++++++++++++-- vllm/parser/deepseek_v32.py | 1 + vllm/parser/deepseek_v4.py | 1 + vllm/parser/engine/parser_engine.py | 36 +++++++++------- vllm/parser/engine/parser_engine_config.py | 2 + vllm/parser/gemma4.py | 34 ++++----------- vllm/parser/glm47_moe.py | 6 +-- vllm/parser/inkling.py | 1 + vllm/parser/kimi_k2.py | 19 +-------- vllm/parser/minimax_m2.py | 1 + vllm/parser/mistral.py | 19 +-------- vllm/parser/qwen3.py | 33 ++------------- 13 files changed, 111 insertions(+), 117 deletions(-) diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index 74224c15b749..ae21d69fc303 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -1616,3 +1616,36 @@ def test_non_streaming_multiple_commas(self, multi_comma_parser, request_obj): assert result.tools_called is True args = json.loads(result.tool_calls[0].function.arguments) assert args["destination"] == "456 Oakwood Avenue, Rivermist, 83214" + + +class TestGemma4IsReasoningEnd: + @staticmethod + def _parser(thinking: bool): + return Gemma4Parser( + _make_tokenizer(_PLAIN_ANSWER_TOKENS), + chat_template_kwargs={"enable_thinking": thinking}, + ) + + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + ([CHANNEL_START_ID, 3000, CHANNEL_END_ID], True, True), + ([CHANNEL_START_ID, 3000, TOOL_CALL_START_ID], True, True), + ([CHANNEL_START_ID, 3000], True, False), + ([NEW_TURN_ID, 9100], True, False), + ([CHANNEL_END_ID, NEW_TURN_ID, 9100], True, False), + ([NEW_TURN_ID, 9100], False, True), + # The bundled template's thinking-off prompt: closed channel. + ([NEW_TURN_ID, 9100, CHANNEL_START_ID, 3000, CHANNEL_END_ID], False, True), + # A marker still wins with thinking off: the grammar honours it. + ([CHANNEL_START_ID, 3000], False, False), + # A tool-call opener after the channel closed: still ended. + ( + [CHANNEL_START_ID, 3000, CHANNEL_END_ID, TOOL_CALL_START_ID, 9200], + True, + True, + ), + ], + ) + def test_is_reasoning_end(self, ids, thinking, ended): + assert self._parser(thinking).is_reasoning_end(ids) is ended diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index 37e6b5d13090..3f285e4d0315 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -180,9 +180,8 @@ def test_prompt_tool_example_before_generation_think_not_end(self, parser): """Tool examples before the generation must not end reasoning.""" assert not parser.is_reasoning_end([_TOOL_CALL_ID, _TEXT_ID, _THINK_START_ID]) - def test_paired_tool_call_not_end(self, parser): - """Paired ... (from template) is NOT end.""" - assert not parser.is_reasoning_end( + def test_completed_tool_call_after_open_think_is_end(self, parser): + assert parser.is_reasoning_end( [_THINK_START_ID, 1, _TOOL_CALL_ID, 2, _TOOL_CALL_END_ID] ) @@ -196,6 +195,43 @@ def test_empty_ids(self, parser): assert not parser.is_reasoning_end([]) +class TestIsReasoningEndDerivedFromGrammar: + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + # Explicit close, and the implicit close on . + ([_THINK_START_ID, 1, _THINK_END_ID], True, True), + ([_THINK_START_ID, 1, _TOOL_CALL_ID], True, True), + # Still inside the block. + ([_THINK_START_ID, 1], True, False), + ([_THINK_END_ID, _THINK_START_ID, 1], True, False), + # Uncommitted: a bare turn, an empty sequence. Thinking on starts + # in REASONING, so reasoning is expected. + ([_IM_END_ID, _IM_START_ID, 1], True, False), + ([], True, False), + # Thinking off starts in CONTENT with no way into REASONING, so + # a the template left behind is inert, and a silent tail + # is ended because no block is expected. + ([_THINK_START_ID, 1], False, True), + ([_THINK_END_ID, _THINK_START_ID, 1], False, True), + ([_IM_END_ID, _IM_START_ID, 1], False, True), + ([], False, True), + ], + ) + def test_is_reasoning_end(self, mock_tokenizer, ids, thinking, ended): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.is_reasoning_end(ids) is ended + + @pytest.mark.parametrize("thinking", [True, False]) + def test_wait_for_reasoning_follows_thinking(self, mock_tokenizer, thinking): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.parser_engine_config.wait_for_reasoning is thinking + + class TestIsReasoningEndTurnBoundaries: """Reasoning markers from earlier turns must not mark the new turn's reasoning as finished (think-marker injection / interleaved-thinking diff --git a/vllm/parser/deepseek_v32.py b/vllm/parser/deepseek_v32.py index f76c954a93d6..de6634cf7368 100644 --- a/vllm/parser/deepseek_v32.py +++ b/vllm/parser/deepseek_v32.py @@ -53,6 +53,7 @@ def deepseek_v32_config() -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v32", initial_state=ParserState.CONTENT, + wait_for_reasoning=False, terminals={ "TOOL_START": DSML_FUNC_START, "TOOL_END": DSML_FUNC_END, diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py index 13eb83e6b804..79ddd93f6b2e 100644 --- a/vllm/parser/deepseek_v4.py +++ b/vllm/parser/deepseek_v4.py @@ -140,6 +140,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v4", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ "THINK_START": DSML_THINK_START, "THINK_END": DSML_THINK_END, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 4817dedb4663..f873e36342f1 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -648,24 +648,28 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: return None def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id + config = self.parser_engine_config + wait_for_reasoning = config.wait_for_reasoning + if wait_for_reasoning is None: + wait_for_reasoning = config.initial_state is ParserState.REASONING + start_transition = config.transitions.get((config.initial_state, "THINK_START")) + start_opens_reasoning = ( + start_transition is not None + and start_transition.next_state is ParserState.REASONING + ) + end_ids = self._reasoning_end_token_ids start_id = self._reasoning_start_token_id - if end_id is not None: - if not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - boundary_ids = self._turn_boundary_token_ids - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if token_id == end_id: - return True - if start_id is not None and token_id == start_id: + boundary_ids = self._turn_boundary_token_ids + for token_id in reversed(input_ids): + if token_id in end_ids: + return True + if token_id == start_id: + if start_opens_reasoning: return False - if token_id in boundary_ids: - return ( - self.parser_engine_config.initial_state != ParserState.REASONING - ) - return False - return self._reasoning_ended + break + if token_id in boundary_ids: + break + return not wait_for_reasoning def extract_content_ids(self, input_ids: list[int]) -> list[int]: end_id = self._reasoning_end_token_id diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 5c5cd2b467fd..b8d7dc6e465a 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -74,6 +74,8 @@ class ParserEngineConfig: initial_state: ParserState = ParserState.CONTENT + wait_for_reasoning: bool | None = None + arg_converter: Callable[[str, bool], str] | None = None stream_arg_deltas: bool = True diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index 51c46966364b..e3baaf488a8d 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -293,10 +293,14 @@ def _gemma4_arg_converter(raw_args: str, partial: bool) -> str: @functools.cache -def gemma4_config() -> ParserEngineConfig: +def gemma4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="gemma4", initial_state=ParserState.CONTENT, + # Decides a silent tail: the E-series HF templates write no marker + # with thinking off, and no template does after a tool response. + wait_for_reasoning=thinking, + turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}), terminals={ "THINK_START": CHANNEL_START, "THINK_END": CHANNEL_END, @@ -417,11 +421,11 @@ def __init__( **kwargs, ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - self._thinking_enabled = chat_kwargs.get("enable_thinking", False) + thinking = chat_kwargs.get("enable_thinking", False) super().__init__( tokenizer, tools, - parser_engine_config=gemma4_config(), + parser_engine_config=gemma4_config(thinking=thinking), **kwargs, ) vocab = self.vocab @@ -472,30 +476,6 @@ def _preprocess_feed( return delta_text, delta_token_ids - def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id - start_id = self._reasoning_start_token_id - tool_call_id = self._tool_call_token_id - new_turn_id = self._new_turn_token_id - tool_response_id = self._tool_response_token_id - - if end_id is not None and not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - - for i in range(len(input_ids) - 1, -1, -1): - tid = input_ids[i] - if start_id is not None and tid == start_id: - return False - if tool_call_id is not None and tid == tool_call_id: - return True - if new_turn_id is not None and tid == new_turn_id: - return not self._thinking_enabled - if tool_response_id is not None and tid == tool_response_id: - return not self._thinking_enabled - if end_id is not None and tid == end_id: - return True - return True - def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool: """Whether the prompt tail is inside an open ``<|channel>`` block. diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 8aa4feef2599..4a804a4c9c94 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -123,6 +123,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="glm47_moe", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_START": TOOL_CALL_START, @@ -206,11 +207,6 @@ def _handle_tool_end(self, event, deltas) -> None: self._tool_slots[idx].name = self._tool_slots[idx].name.strip() super()._handle_tool_end(event, deltas) - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - return super().is_reasoning_end(input_ids) - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 16883b58aa7f..2081fb278464 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -282,6 +282,7 @@ def inkling_config() -> ParserEngineConfig: # `<|message_model|>`. Non-streaming parsing receives only the generated # suffix, so begin in the corresponding message-header state as well. initial_state=ParserState.MESSAGE_HEADER, + wait_for_reasoning=True, terminals=terminals, # Inkling content-kind markers are the grammar. When the engine is # used through DelegatingParser, the reasoning pass can hand the tool diff --git a/vllm/parser/kimi_k2.py b/vllm/parser/kimi_k2.py index 1e1067a8c2ee..7c5f30aafef4 100644 --- a/vllm/parser/kimi_k2.py +++ b/vllm/parser/kimi_k2.py @@ -81,6 +81,7 @@ def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="kimi_k2", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_SECTION_START": TOOL_SECTION_START, @@ -236,24 +237,6 @@ def _handle_arg_chunk(self, event, deltas) -> None: def _extract_args_json(self, raw_args: str, func_name: str) -> str: return raw_args.strip() or "{}" - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - - start_id = self._start_token_id - end_id = self._end_token_id - tool_section_id = self._tool_section_start_token_id - - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if start_id is not None and token_id == start_id: - return False - if end_id is not None and token_id == end_id: - return True - if tool_section_id is not None and token_id == tool_section_id: - return True - return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py index ab1b7eb11427..c96f852b9df6 100644 --- a/vllm/parser/minimax_m2.py +++ b/vllm/parser/minimax_m2.py @@ -95,6 +95,7 @@ def minimax_m2_config() -> ParserEngineConfig: return ParserEngineConfig( name="minimax_m2", initial_state=ParserState.REASONING, + wait_for_reasoning=True, terminals={ "THINK_START": THINK_START, "THINK_END": THINK_END, diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index fc9b8fbefe37..ce516eea1419 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -193,6 +193,7 @@ def mistral_config( return ParserEngineConfig( name=name, initial_state=ParserState.CONTENT, + wait_for_reasoning=reasoning_encoding != "none", terminals={ **reasoning_terminals, "TOOL_CALLS": _TOOL_CALLS, @@ -891,24 +892,6 @@ def _split_delta( return (delta_to_be_parsed, "") - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if self._reasoning_encoding == "none": - return True - if super().is_reasoning_end(input_ids): - return True - # [TOOL_CALLS] acts as an implicit reasoning-end marker - if self.bot_token_id is not None: - reasoning_start_id = self._reasoning_start_token_id - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] == self.bot_token_id: - return True - return False - def extract_reasoning( self, model_output: str, diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index 8edf42e26363..a11551a209a0 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -100,6 +100,7 @@ def qwen3_config( return ParserEngineConfig( name=name, initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, turn_boundary_tokens=turn_boundary_tokens, terminals={ # Reasoning terminals @@ -205,8 +206,8 @@ class Qwen3Parser(ParserEngine): """Qwen3 parser: ````/```` reasoning + ```` XML tool calls in a single engine. - - ```` as implicit reasoning end - - Unpaired ```` token ID detection for ``is_reasoning_end`` + - ```` as implicit reasoning end (a grammar transition, so it + also feeds ``is_reasoning_end`` and the structured-output gate) Subclasses that share the grammar but differ only in the four wrapper token strings (reasoning + tool-call) override the class attributes @@ -245,9 +246,6 @@ def __init__( tools, **kwargs, ) - vocab = self.vocab - self._tool_call_token_id: int | None = vocab.get(self.TOOL_START) - self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END) def extract_reasoning( self, @@ -257,28 +255,3 @@ def extract_reasoning( if not self.thinking_enabled: return None, model_output return super().extract_reasoning(model_output, request) - - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if super().is_reasoning_end(input_ids): - return True - tool_call_id = self._tool_call_token_id - tool_call_end_id = self._tool_call_end_token_id - reasoning_start_id = self._reasoning_start_token_id - boundary_ids = self._turn_boundary_token_ids - if tool_call_id is not None: - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] in boundary_ids: - return False - if input_ids[i] == tool_call_id: - if tool_call_end_id is not None and any( - input_ids[j] == tool_call_end_id - for j in range(i + 1, len(input_ids)) - ): - continue - return True - return False From 5be911ee871604895c5a950ac1a5675c71896fd2 Mon Sep 17 00:00:00 2001 From: ALEX LILUZ <49665315+alexliluz@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:31:25 +0800 Subject: [PATCH 252/774] [Bugfix] Load reasoning parser plugins before headless engine config (#53124) Signed-off-by: alexliluz <49665315+alexliluz@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- tests/entrypoints/cli/test_serve.py | 42 +++++++++++++++++++++++++++++ vllm/entrypoints/cli/serve.py | 4 +++ 2 files changed, 46 insertions(+) create mode 100644 tests/entrypoints/cli/test_serve.py diff --git a/tests/entrypoints/cli/test_serve.py b/tests/entrypoints/cli/test_serve.py new file mode 100644 index 000000000000..9ba9ac7455c3 --- /dev/null +++ b/tests/entrypoints/cli/test_serve.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the ``vllm serve`` CLI subcommand.""" + +import argparse +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.entrypoints.cli.serve import run_headless + + +class _StopAfterEngineConfig(Exception): + pass + + +def test_headless_imports_reasoning_parser_plugin_before_engine_config(): + plugin_path = "/tmp/custom_reasoning_parser.py" + args = argparse.Namespace( + api_server_count=0, + reasoning_parser_plugin=plugin_path, + ) + engine_args = MagicMock() + + with ( + patch( + "vllm.entrypoints.cli.serve.vllm.AsyncEngineArgs.from_cli_args", + return_value=engine_args, + ), + patch( + "vllm.reasoning.ReasoningParserManager.import_reasoning_parser" + ) as import_plugin, + ): + + def create_engine_config(**kwargs): + import_plugin.assert_called_once_with(plugin_path) + raise _StopAfterEngineConfig + + engine_args.create_engine_config.side_effect = create_engine_config + + with pytest.raises(_StopAfterEngineConfig): + run_headless(args) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index e43b273e5e96..6a8492f4c4b7 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -18,6 +18,7 @@ from vllm.entrypoints.launchers.dp_supervisor import run_dp_supervisor from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger +from vllm.reasoning import ReasoningParserManager from vllm.usage.usage_lib import UsageContext from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import get_tcp_uri @@ -179,6 +180,9 @@ def run_headless(args: argparse.Namespace): if args.api_server_count > 1: raise ValueError("api_server_count can't be set in headless mode") + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + # Create the EngineConfig. engine_args = vllm.AsyncEngineArgs.from_cli_args(args) usage_context = UsageContext.OPENAI_API_SERVER From bb5507741155a698750bc422cda1f7c57f7d114d Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 04:01:06 -0700 Subject: [PATCH 253/774] [Perf][Kernel] Integrate Mega-mHC from DeepGEMM for DeepSeek V4.1 (reopen of #56255) (#56962) Signed-off-by: Yongye Zhu Signed-off-by: Juntian Liu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Thien Tran --- tests/kernels/quantization/test_block_fp8.py | 7 +- tests/kernels/test_mhc_kernels.py | 105 +++++++++++- vllm/models/deepseek_v41/nvidia/model.py | 72 ++++---- .../deepseek_v41/nvidia/ops/__init__.py | 2 + .../deepseek_v41/nvidia/ops/mega_mhc.py | 156 ++++++++++++++++++ vllm/utils/deep_gemm.py | 15 +- 6 files changed, 316 insertions(+), 41 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/__init__.py create mode 100644 vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index fd530e7ffb53..d02527aba048 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -27,6 +27,7 @@ from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_tma_aligned_size, + is_deep_gemm_e8m0_used, per_block_cast_to_fp8, should_use_deepgemm_for_fp8_linear, ) @@ -291,7 +292,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed): A_fp8, As_fp8 = per_token_group_quant_fp8( A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True ) - B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size) + B_fp8, Bs_fp8 = per_block_cast_to_fp8( + B_fp32, + block_size=block_size, + use_ue8m0=is_deep_gemm_e8m0_used(), + ) As = As_fp8.to(torch.float32) Bs = Bs_fp8.to(torch.float32) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index ff72e2a96c17..5fab150d9fce 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -11,6 +11,7 @@ import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( _torch_hc_prenorm_gemm, + mhc_post_tilelang, mhc_pre_delayed_tilelang, ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -39,6 +40,10 @@ from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) +from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import ( + is_mega_mhc_supported, + mhc_shifted_post_pre_deep_gemm, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -495,7 +500,7 @@ def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, ** reference, ) monkeypatch.setattr( - "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", fused_reference, ) expected = decoder(x, positions, None, **kwargs) @@ -580,6 +585,104 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): ) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="DeepGEMM Mega mHC requires SM100-family CUDA", +) +@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024]) +@pytest.mark.parametrize("hidden_size", [5120, 7168]) +def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size): + set_random_seed(0) + hc_mult = 4 + if not is_mega_mhc_supported(hidden_size, 4): + pytest.skip("DeepGEMM Mega mHC is not available") + + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01 + scale = torch.randn(3, device=DEVICE) * 0.1 + base = torch.randn(24, device=DEVICE) * 0.1 + norm_weight = torch.empty( + hidden_size, dtype=torch.bfloat16, device=DEVICE + ).uniform_(0.9, 1.1) + + expected_residual = ( + mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual + ) + ( + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) = mhc_pre_delayed_tilelang( + expected_residual, + fn, + scale, + base, + 2e-5, + 3e-4, + 2e-6, + 1.25, + 10, + pre_mix=previous_mix, + norm_weight=norm_weight, + norm_eps=7e-6, + ) + args = ( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + 2e-5, + 3e-4, + 1.25, + 2e-6, + 10, + norm_weight, + 7e-6, + ) + actual = mhc_shifted_post_pre_deep_gemm(*args) + expected = ( + expected_residual, + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) + for i, (result, ref) in enumerate(zip(actual, expected, strict=True)): + # Match the existing mHC BF16 tolerance: post rounding differences + # can be amplified by cancellation in the carried collapse. + atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3) + torch.testing.assert_close(result, ref, atol=atol, rtol=rtol) + + if num_tokens: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # DeepGEMM initializes barriers per stream before capture. + mhc_shifted_post_pre_deep_gemm(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = mhc_shifted_post_pre_deep_gemm(*args) + for multiplier in (0.5, -1.0): + x.mul_(multiplier) + graph.replay() + eager = mhc_shifted_post_pre_deep_gemm(*args) + for result, ref in zip(captured, eager, strict=True): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.cuda.current_stream().wait_stream(stream) + + def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: x = x.softmax(-1) + eps x = x / (x.sum(-2, keepdim=True) + eps) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index f3530544da7b..9c9719e97729 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,7 +20,6 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -79,6 +78,7 @@ from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes +from .ops.mega_mhc import mhc_shifted_post_pre if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -403,25 +403,23 @@ def forward( else: # The collapse already reads the post-mapped streams, so the mean # aux consumers want comes out of the same kernel. - residual, post_mix, res_mix, x, attn_pre, aux = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=pre_mix, - norm_weight=self.attn_norm.weight, - norm_eps=self.attn_norm.variance_epsilon, - capture_aux=capture_previous_aux, - ) + residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, ) if capture_previous_aux: previous_aux = aux @@ -433,24 +431,22 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual, post_mix, res_mix, x, ffn_pre, _ = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, - ) + residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, ) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre, previous_aux diff --git a/vllm/models/deepseek_v41/nvidia/ops/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py new file mode 100644 index 000000000000..f2c74902616f --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import torch + +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + _import_deep_gemm, + is_deep_gemm_supported, + mega_mhc, +) + + +@functools.cache +def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: + """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel.""" + if not ( + is_deep_gemm_supported() + and current_platform.is_device_capability_family(100) + and hidden_size > 0 + and hidden_size % 1024 == 0 + and hc_mult == 4 + ): + return False + deep_gemm = _import_deep_gemm() + return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) + + +def mhc_shifted_post_pre_deep_gemm( + x: torch.Tensor, + residual: torch.Tensor, + shifted_prev_mix: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_weight: torch.Tensor, + rmsnorm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC.""" + num_tokens, hidden_size = x.shape + hc_mult = residual.shape[1] + new_residual = torch.empty_like(residual) + new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32) + new_post_mix = torch.empty_like(post_mix) + new_comb_res_mix = torch.empty_like(comb_res_mix) + y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16) + mega_mhc( + x=x, + residual=residual, + shifted_prev_mix=shifted_prev_mix.unsqueeze(-1), + post_mix=post_mix, + comb_res_mix=comb_res_mix, + fn=fn, + mix_scales=mix_scales, + mix_bases=mix_bases, + hc_mult=hc_mult, + hc_norm_eps=hc_norm_eps, + hc_pre_eps=hc_pre_eps, + hc_post_scale=hc_post_scale, + sinkhorn_eps=sinkhorn_eps, + num_sinkhorn_iters=num_sinkhorn_iters, + rmsnorm_weight=rmsnorm_weight, + rmsnorm_eps=rmsnorm_eps, + rmsnorm_scale=1.0, + new_residual=new_residual, + new_prev_mix=new_prev_mix, + new_post_mix=new_post_mix, + new_comb_res_mix=new_comb_res_mix, + y_bf16=y_bf16, + ) + return ( + new_residual, + new_post_mix, + new_comb_res_mix, + y_bf16, + new_prev_mix.squeeze(-1), + ) + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Use Mega-mHC for carried mixing, retaining TileLang for aux capture.""" + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( + x, + residual, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + return *outputs, x.new_empty(0, x.shape[1]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + capture_aux=capture_aux, + ) diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index d515d5131e75..5d05c29a84c5 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -159,6 +159,7 @@ def _missing(*_: Any, **__: Any) -> NoReturn: _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None +_mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -231,7 +232,7 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl - global _tf32_hc_prenorm_gemm_impl + global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -251,6 +252,7 @@ def _lazy_init() -> None: or _fp8_fp4_paged_mqa_logits_impl is not None or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None + or _mega_mhc_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -294,6 +296,7 @@ def _lazy_init() -> None: _dg, "get_paged_mqa_logits_metadata", None ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) + _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -704,6 +707,15 @@ def tf32_hc_prenorm_gemm( ) +def mega_mhc(*args: Any, **kwargs: Any) -> None: + """Run DeepGEMM Mega mHC with caller-owned output tensors.""" + _lazy_init() + if _mega_mhc_impl is None: + _missing() + return + _mega_mhc_impl(*args, **kwargs) + + def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -797,6 +809,7 @@ def should_use_deepgemm_for_fp8_linear( "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", "is_deep_gemm_supported", + "mega_mhc", "get_num_sms", "set_num_sms", "should_use_deepgemm_for_fp8_linear", From 78eaaad2c7a7f8c7864d9c9c72772e3f9f2669d0 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 15 Sep 2026 20:12:15 +0800 Subject: [PATCH 254/774] [Rust Frontend] Add iteration token histogram (#56990) Co-authored-by: Codex Signed-off-by: Bugen Zhao --- rust/src/engine-core-client/src/client/imp.rs | 26 +++++++++++++++++- rust/src/engine-core-client/src/metrics.rs | 21 +++++++++++++-- rust/src/metrics/src/request.rs | 27 ++++++++++++++----- rust/src/mock-engine/src/engine.rs | 19 +++++++++---- rust/src/mock-engine/src/tests.rs | 11 ++++++++ 5 files changed, 90 insertions(+), 14 deletions(-) diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 117473ce47cc..7318fa01b8ce 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; +use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; @@ -47,6 +47,7 @@ pub(crate) struct ClientInner { /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created. msgpack_zero_copy_threshold: usize, scheduler_stats_recorder: SchedulerStatsRecorder, + iteration_metrics: BTreeMap, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, @@ -64,12 +65,21 @@ impl ClientInner { ) -> Self { let scheduler_stats_recorder = SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines); + let iteration_metrics = engines + .iter() + .filter_map(|engine| { + let engine = engine.engine_id.engine_index()?; + let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine); + Some((engine, handles)) + }) + .collect(); Self { input_send, handle, model_name, msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(), scheduler_stats_recorder, + iteration_metrics, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), @@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { + let has_outputs = !batch.outputs.is_empty(); + let mut iteration_tokens = 0_u64; let senders = inner.take_senders_for_outputs(&batch.outputs); for (output, sender) in batch.outputs.into_iter().zip(senders) { let request_id = output.request_id.clone(); @@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop( continue; }; + iteration_tokens += output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + iteration_tokens += u64::from(prefill.num_computed_tokens); + } + let wrapped_output = EngineCoreStreamOutput { engine_index: batch.engine_index, timestamp: batch.timestamp, @@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop( } } + if has_outputs + && let Some(handles) = inner.iteration_metrics.get(&batch.engine_index) + { + handles.iteration_tokens_total.observe(iteration_tokens as f64); + } + // The sender for normally-finished requests should have already been removed // from the registry when their final output was dispatched // above. This serves as a safety net to capture any diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e15258d6081d..b392d5260185 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, - MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, - WaitingReasonLabels, + MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, + U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::{ @@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; +/// Cached output-batch metric handles for one model and engine index. +pub(crate) struct IterationMetricHandles { + pub iteration_tokens_total: HistogramMetric, +} + +impl IterationMetricHandles { + pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self { + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; + Self { + iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels), + } + } +} + /// Cached scheduler-stats metric handles for all engines connected to one /// frontend client. pub(crate) struct SchedulerStatsRecorder { diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index c1135781b23c..691c789936b2 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -21,6 +21,9 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [ 480.0, 960.0, 1920.0, 7680.0, ]; const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0]; +const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ + 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0, +]; fn build_1_2_5_buckets(max_value: u32) -> Vec { let mut buckets = Vec::new(); @@ -66,6 +69,10 @@ fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } +fn iteration_tokens_histogram() -> Histogram { + Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied()) +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct FinishedReasonLabels { pub model_name: String, @@ -83,8 +90,12 @@ pub struct PromptTokenSourceLabels { pub(crate) type FinishedReasonCounterFamily = Family; pub(crate) type PromptTokenSourceCounterFamily = Family; -/// Request-lifecycle Prometheus families exported from the `llm` layer. +/// Request and output-batch Prometheus families. pub struct RequestMetrics { + /// Computed prompt and generated tokens per nonempty engine output batch, + /// recorded by the engine-core client for requests tracked by this frontend. + pub iteration_tokens_total: HistogramFamily, + // Request-derived counters. pub num_preemptions: Family, pub prompt_tokens: Family, @@ -92,11 +103,6 @@ pub struct RequestMetrics { pub prompt_tokens_cached: Family, pub generation_tokens: Family, - // We intentionally don't support iteration-level histograms for now, since it seems to make - // more sense if the engine maintains these metrics and frontend simply forwards. - // - // pub iteration_tokens_total: HistogramFamily, - // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, pub request_prompt_tokens: HistogramFamily, @@ -118,6 +124,14 @@ pub struct RequestMetrics { impl RequestMetrics { /// Register the request-oriented metric families into the shared registry. pub(crate) fn register(registry: &mut Registry) -> Self { + let iteration_tokens_total = + Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram); + registry.register( + "vllm:iteration_tokens_total", + "Histogram of computed prompt and generated tokens per nonempty engine output batch.", + iteration_tokens_total.clone(), + ); + // Request-derived counters. let num_preemptions = Family::default(); registry.register( @@ -276,6 +290,7 @@ impl RequestMetrics { ); Self { + iteration_tokens_total, num_preemptions, prompt_tokens, prompt_tokens_by_source, diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index c21682915910..38f9b21be1d7 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -19,6 +19,7 @@ use vllm_engine_core_client::protocol::output::{ UtilityCallOutput, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::protocol::utility::{ EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope, }; @@ -187,6 +188,11 @@ impl ActiveRequest { /// Advance this request by one mock engine step. fn step(&mut self, opt: &Opt) -> EngineCoreOutput { + let prefill_stats = (self.generated == 0).then(|| PrefillStats { + num_prompt_tokens: self.prompt_len as u32, + num_computed_tokens: self.prompt_len as u32, + ..Default::default() + }); let remaining = self.max_tokens - self.generated; let chunk_len = remaining.min(opt.output_token_chunk_size); let mut new_token_ids = Vec::with_capacity(chunk_len); @@ -196,11 +202,14 @@ impl ActiveRequest { self.generated += chunk_len; let finished = self.generated >= self.max_tokens; - request_output( - self.request_id.clone(), - new_token_ids, - finished.then_some(EngineCoreFinishReason::Length), - ) + EngineCoreOutput { + prefill_stats, + ..request_output( + self.request_id.clone(), + new_token_ids, + finished.then_some(EngineCoreFinishReason::Length), + ) + } } } diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 9bfee620ec74..bee710fd0c15 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -10,6 +10,7 @@ use tokio_util::sync::CancellationToken; use vllm_engine_core_client::protocol::output::EngineCoreFinishReason; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::IpcNamespace; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; @@ -119,12 +120,22 @@ async fn chunk_size_one_outputs_one_token_per_update() { let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); assert_eq!(first.finish_reason, None); + assert_eq!( + first.prefill_stats, + Some(PrefillStats { + num_prompt_tokens: 3, + num_computed_tokens: 3, + ..Default::default() + }) + ); let second = stream.next().await.expect("second").expect("second ok"); assert_eq!(second.new_token_ids.len(), 1); assert_eq!(second.finish_reason, None); + assert!(second.prefill_stats.is_none()); let third = stream.next().await.expect("third").expect("third ok"); assert_eq!(third.new_token_ids.len(), 1); assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length)); + assert!(third.prefill_stats.is_none()); assert!(stream.next().await.is_none()); shutdown_mock(client, shutdown, task).await; From 64856080b05054fc62754f37234f7b483b753445 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 08:21:05 -0500 Subject: [PATCH 255/774] [Bugfix] Measure complete pooling responses (#56760) Signed-off-by: Taneem Ibrahim --- vllm/benchmarks/lib/endpoint_request_func.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index b55c5f4e164d..54c6ce63b0fa 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -606,6 +606,7 @@ async def _run_pooling_request( headers: dict[str, Any], pbar: tqdm | None = None, num_input_sequences: int = 1, + prompt_len: int = 0, ) -> RequestFuncOutput: output = RequestFuncOutput(num_input_sequences=num_input_sequences) st = time.perf_counter() @@ -613,11 +614,19 @@ async def _run_pooling_request( try: async with session.post(url=api_url, headers=headers, json=payload) as response: if response.status == 200: + encoding_format = payload.get("encoding_format", "float") + if encoding_format in ("bytes", "bytes_only"): + async for _ in response.content.iter_any(): + pass + else: + await response.read() output.ttft = output.latency = time.perf_counter() - st - if payload.get("encoding_format", "float") == "bytes": + if encoding_format == "bytes": metadata = json.loads(response.headers["metadata"]) usage = metadata.get("usage", {}) + elif encoding_format == "bytes_only": + usage = {"prompt_tokens": prompt_len} else: data = await response.json() usage = data.get("usage", {}) @@ -672,6 +681,7 @@ async def async_request_openai_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -709,6 +719,7 @@ async def async_request_vllm_rerank( headers=headers, pbar=pbar, num_input_sequences=len(request_func_input.prompt) - 1, + prompt_len=request_func_input.prompt_len, ) @@ -743,6 +754,7 @@ async def async_request_openai_embeddings_chat( payload=payload, headers=headers, pbar=pbar, + prompt_len=request_func_input.prompt_len, ) @@ -847,6 +859,7 @@ async def async_request_infinity_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -896,6 +909,7 @@ async def async_request_vllm_pooling( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) From e6960af33b379d502f409e3e2241bbf2b2c2f68d Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Tue, 15 Sep 2026 09:48:49 -0400 Subject: [PATCH 256/774] [Refactor] StructuredOutputManager x Speculative Decoding Refactor (#48200) --- tests/parser/engine/test_deepseek_v4.py | 2 +- tests/parser/engine/test_parser_engine.py | 4 +- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_scheduler.py | 28 +- .../spec_decode/test_mtp_structured_output.py | 359 +---------- .../test_reasoning_structured_output.py | 557 ------------------ .../test_structured_output_manager.py | 414 +++++++++++++ vllm/parser/abstract_parser.py | 8 +- vllm/parser/engine/parser_engine.py | 2 +- vllm/v1/core/sched/scheduler.py | 53 +- vllm/v1/structured_output/__init__.py | 347 +++++------ vllm/v1/structured_output/request.py | 6 - 12 files changed, 595 insertions(+), 1191 deletions(-) delete mode 100644 tests/v1/structured_output/test_reasoning_structured_output.py create mode 100644 tests/v1/structured_output/test_structured_output_manager.py diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 1c1029b51a01..7cbf9791a7b0 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -574,7 +574,7 @@ def test_reasoning_end_token_ids_drop_multi_token_tool_opener( # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) is None + assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index f27ce52cb37b..8fd7c4d2623f 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -185,8 +185,8 @@ def test_config_without_reasoning_has_empty_set(self): def test_find_reasoning_end_offset_returns_first_match(self): engine = _make_engine() assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) is None - assert engine.find_reasoning_end_offset([]) is None + assert engine.find_reasoning_end_offset([5, 6]) == 2 + assert engine.find_reasoning_end_offset([]) == 0 # Rejected-draft placeholders never match. assert engine.find_reasoning_end_offset([-1, 201]) == 1 diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7ff63d019c85..7e27a4517e03 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request = create_requests(num_requests=1, num_tokens=1)[0] request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 @@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 6ca6a8cf87aa..061c825fc895 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3674,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens @@ -3682,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() @@ -3734,8 +3730,8 @@ def free_request(req: Request, delay_free_blocks: bool = False): ) engine_core_outputs = scheduler.update_from_output(output, model_runner_output) - request.structured_output_request.grammar.accept_tokens.assert_called_once_with( - request.request_id, [123] + scheduler.structured_output_manager.accept_tokens.assert_called_once_with( + request, [123] ) assert request.resumable is False assert request.status == RequestStatus.FINISHED_ERROR @@ -6573,16 +6569,15 @@ def _decode_ready_request(scheduler): return request -def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): - """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must +def test_update_draft_token_ids_strips_ngram_padding(): + """ngram_gpu pads unfilled draft slots with -1; manager validation must strip them before grammar.validate_tokens (which otherwise raises).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler.update_draft_token_ids( @@ -6593,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): assert request.spec_token_ids == [10, 11] -def test_update_draft_token_ids_in_output_strips_padding(monkeypatch): +def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to validate_tokens is stripped).""" + is preserved (only the input to manager validation is stripped).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler_output = SimpleNamespace( diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 330e10438f0b..4e12a6f49dd8 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""grammar_bitmask under spec-decode draft padding (#44006).""" +"""Backend and engine-reasoner checks that are not manager-flow tests.""" import pytest from transformers import AutoTokenizer @@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'): return tokenizer, manager, request, prompt -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_with_padded_invalid_drafts(backend): - """Bitmask handles -1 padded drafts and returns N+1 rows.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]] - padded = valid_drafts + [-1, -1] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_when_grammar_terminates_mid_window(backend): - """Drafts following an EOS that terminates the grammar are a no-op.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - eos = tokenizer.eos_token_id - drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_idempotent_across_calls(backend): - """Repeated calls with the same input return the same bitmask.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - drafts = [tokenizer.encode(":")[0], -1, -1, -1] - - first = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - second = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert first is not None and second is not None - assert (first == second).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bonus_position_constrained_after_invalid_drafts(backend): - """Regression for #44006: bonus row stays constrained after -1 padding.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid = tokenizer.encode(":")[0] - drafts = [valid, -1, -1, -1] - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_constrained_when_reasoning_ends_midwindow(backend): - """Drafts after a mid-window reasoning-end marker stay constrained.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, marker, post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog): - """Post-marker drafts predate the bitmask and may be grammar-invalid; - grammar_bitmask must skip the grammar advance instead of asserting. - """ - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str="{" - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - assert not grammar.is_terminated() - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - # A token that the JSON grammar would reject as the first post-marker - # token; without the fix grammar.accept_tokens fires the assertion. - invalid_post = tokenizer.encode("z")[0] - drafts = [pre, marker, invalid_post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - # Post-marker position is still bitmask-constrained. - assert not (bitmask[2] == -1).all() - # Grammar must not have advanced through the unvalidated draft. - assert not grammar.is_terminated() - assert "Failed to advance FSM" not in caplog.text - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_validate_tokens_then_bitmask_round_trip(backend): - """validate_tokens -> pad with -1 -> grammar_bitmask must not assert.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890] - accepted = grammar.validate_tokens(raw_drafts) - assert len(accepted) <= len(raw_drafts) - - padded = accepted + [-1] * (len(raw_drafts) - len(accepted)) - assert len(padded) == len(raw_drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - def test_xgrammar_accept_tokens_stops_at_termination(capfd): """Tokens after a terminating EOS do not reach the matcher.""" tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") @@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.validate_tokens([eos, trailing]) == [eos] assert "trying to accept new token" not in capfd.readouterr().err - # Check matcher state directly to verify validation rolled it back. assert not grammar.matcher.is_terminated() assert grammar.accept_tokens(request.request_id, [eos]) @@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert "trying to accept new token" not in capfd.readouterr().err -class _MarkerReasoner: - """Stub reasoner whose reasoning-end marker is a single fixed token.""" - - def __init__(self, marker: int): - self.marker = marker - - def is_reasoning_end(self, input_ids): - return self.marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return self.marker in list(delta_ids) - - -def _setup_boundary_request(backend: str): - """Request with a structural-tag key and reasoning not yet ended.""" - from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - marker = tokenizer.encode("\n")[0] - structured_req = request.structured_output_request - # The grammar itself is JSON (cheap to build); only the key kind matters - # for the should_advance structural-tag branch, so pre-seed the cached - # property. - structured_req.__dict__["structured_output_key"] = ( - StructuredOutputOptions.STRUCTURAL_TAG, - "", - ) - manager.reasoner_cls = _MarkerReasoner - structured_req.reasoner = _MarkerReasoner(marker) - structured_req.reasoning_ended = False - return tokenizer, manager, request, prompt, marker - - -def test_should_advance_records_reasoning_end_index(): - """Regression for #44006 on post-#42452 main: the boundary step must - record where reasoning ends so the scheduler can trim before advancing. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - request.append_output_token_ids([pre, marker, post]) - - assert manager.should_advance(request) - assert structured_req.reasoning_ended - # Marker sits at absolute index len(prompt) + 1. - assert structured_req.reasoning_end_token_index == len(prompt) + 1 - - -def test_trim_reasoning_for_advance(): - """trim drops the marker and everything before it; later steps and - requests without a recorded boundary pass through unchanged. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - - # No boundary recorded yet: pass-through. - assert manager.trim_reasoning_for_advance(request, [pre]) == [pre] - - # Boundary step: marker mid-step keeps only the suffix. - step_tokens = [pre, marker, post] - request.append_output_token_ids(step_tokens) - assert manager.should_advance(request) - assert manager.trim_reasoning_for_advance(request, step_tokens) == [post] - - # Boundary step variant: marker last (the #44006 crash shape - # [198, ]) trims to empty -> scheduler skips accept_tokens. - structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1 - assert manager.trim_reasoning_for_advance(request, step_tokens) == [] - - # Later steps: tokens are past the boundary, returned unchanged. - structured_req.reasoning_end_token_index = len(prompt) + 1 - next_step = [post, post] - request.append_output_token_ids(next_step) - assert manager.trim_reasoning_for_advance(request, next_step) == next_step - - class _EngineReasonerStub(ParserEngineReasoningAdapter): """Adapter-typed reasoner with a fixed end-token set and no real engine.""" @@ -409,7 +112,7 @@ def find_reasoning_end_offset(self, token_ids): for offset, token in enumerate(token_ids): if token in self._end_token_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids): return any(token in self._end_token_ids for token in input_ids) @@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend): assert not (bitmask[2] == -1).all() assert reasoner.windows == [[pre, marker, post]] assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend): - """Legacy predicates are probed one draft at a time, never on the whole - window. A predicate that reports which marker is *newest* (KimiK3) would - otherwise answer False for a window that closes and then reopens - reasoning, leaving the post-marker drafts and the bonus row - unconstrained. - """ - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - close = tokenizer.encode("\n")[0] - reopen = tokenizer.encode("#")[0] - assert close != reopen - - class NewestMarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end(self, input_ids): - return self.is_reasoning_end_streaming(input_ids, input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - for token in reversed(list(delta_ids)): - if token == close: - return True - if token == reopen: - return False - return False - - manager.reasoner_cls = NewestMarkerReasoner - request.structured_output_request.reasoner = NewestMarkerReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, close, post, reopen] - # Sanity: the whole window says "reopened", so only per-token probing - # can find the close marker. - assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py deleted file mode 100644 index a57e5da72b26..000000000000 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ /dev/null @@ -1,557 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for reasoning-aware structured output functionality (PR #25515).""" - -from collections.abc import Sequence -from unittest.mock import Mock - -import pytest - -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig -from vllm.v1.request import Request -from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - -class UntouchableTokenIds(Sequence[int]): - """A token history that fails the test if anything reads it.""" - - def __init__(self, length: int) -> None: - self._length = length - - def __len__(self) -> int: - return self._length - - def __getitem__(self, item): - raise AssertionError("token history must not be read") - - def __iter__(self): - raise AssertionError("token history must not be iterated") - - -def make_qwen3_reasoner(): - """A real engine-based reasoning adapter over a tiny mock vocabulary.""" - from tests.parser.engine.conftest import make_mock_tokenizer - from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter - - tokenizer = make_mock_tokenizer( - {"": 200, "": 201, "": 202, "": 203} - ) - return Qwen3ParserReasoningAdapter(tokenizer) - - -class MockReasoner: - def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) - self.is_reasoning_end_streaming = Mock(return_value=False) - - -class TestReasoningStructuredOutput: - """Test reasoning-aware structured output functionality.""" - - @pytest.fixture - def mock_model_config(self): - """Create a mock ModelConfig.""" - config = Mock(spec=ModelConfig) - config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls - config.get_vocab_size = Mock(return_value=50000) - # Add missing runner_type attribute that tokenizer initialization expects - config.runner_type = "generate" - # Add other attributes that tokenizer initialization might need - config.tokenizer = "test-tokenizer" - config.tokenizer_mode = "auto" - config.trust_remote_code = False - config.tokenizer_revision = None - return config - - @pytest.fixture - def mock_scheduler_config(self): - """Create a mock SchedulerConfig.""" - config = Mock(spec=SchedulerConfig) - config.max_num_seqs = 128 - return config - - @pytest.fixture - def mock_vllm_config(self, mock_model_config, mock_scheduler_config): - """Create a mock VllmConfig.""" - config = Mock(spec=VllmConfig) - config.model_config = mock_model_config - config.scheduler_config = mock_scheduler_config - config.structured_outputs_config = Mock() - config.structured_outputs_config.reasoning_parser = None - config.structured_outputs_config.enable_in_reasoning = False - config.speculative_config = None - return config - - @pytest.fixture - def mock_request_with_structured_output(self): - """Create a mock request with structured output.""" - request = Mock(spec=Request) - request.structured_output_request = Mock() - request.structured_output_request.reasoning_ended = None - request.structured_output_request.grammar = Mock() - request.structured_output_request.reasoning_parser_kwargs = None - request.structured_output_request.reasoner = None - request.structured_output_request.grammar.is_terminated = Mock( - return_value=False - ) - request.use_structured_output = True - request.prompt_token_ids = [1, 2, 3, 4, 5] - request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8] - request.num_computed_tokens = 5 - request.num_output_placeholders = 0 - request.request_id = "mock_req" - return request - - @pytest.fixture - def manager_with_reasoner(self, mock_vllm_config): - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = MockReasoner - manager.tokenizer = Mock() - return manager - - def test_should_fill_bitmask_with_enable_in_reasoning( - self, manager_with_reasoner, mock_request_with_structured_output - ): - """Test should_fill_bitmask when enable_in_reasoning is True.""" - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_fill_bitmask_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """An active grammar does not need a request-local reasoner.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - assert result is True - assert structured_req.reasoner is None - - def test_grammar_bitmask_skips_reasoner_when_already_active( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Bitmask generation skips reasoning-boundary detection.""" - manager_with_reasoner.vllm_config.num_speculative_tokens = 1 - manager_with_reasoner._grammar_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.shape = (1, 1) - expected_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask - manager_with_reasoner._fill_bitmasks = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - request_id = mock_request_with_structured_output.request_id - - bitmask = manager_with_reasoner.grammar_bitmask( - requests={request_id: mock_request_with_structured_output}, - structured_output_request_ids=[request_id], - scheduled_spec_decode_tokens={}, - ) - - assert bitmask is expected_bitmask - assert structured_req.reasoner is None - - def test_should_fill_bitmask_without_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_fill_bitmask when enable_in_reasoning is False.""" - # Keep enable_in_reasoning as False (default) - config = manager_with_reasoner.vllm_config.structured_outputs_config - assert config.enable_in_reasoning is False - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - # Should set reasoning_ended and return its value - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is False - ) - assert result is False - - def test_should_fill_bitmask_no_reasoner( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test should_fill_bitmask when no reasoner is configured.""" - manager = StructuredOutputManager(mock_vllm_config) - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - # Should default to True when no reasoner - assert result is True - - def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test request-level parser kwargs override the default reasoner.""" - - class KwargReasoner: - def __init__(self, tokenizer, chat_template_kwargs=None): - self.chat_template_kwargs = chat_template_kwargs or {} - - def is_reasoning_end(self, input_ids): - return not self.chat_template_kwargs.get("enable_thinking", False) - - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = KwargReasoner - manager.tokenizer = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_parser_kwargs = { - "chat_template_kwargs": {"enable_thinking": True} - } - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - assert result is False - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is not None - ) - - def test_should_advance_with_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_reasoning_not_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has not ended.""" - # Set reasoning as not ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return False since reasoning hasn't ended - assert result is False - - def test_should_advance_reasoning_just_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning ends in current step.""" - # Set reasoning as not ended initially, but ends in this step - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = True - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoner = reasoner - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # The scheduler trims the reasoning prefix before advancing the grammar. - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is True - ) - assert result is True - - def test_should_advance_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has already ended.""" - # Set reasoning as already ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = True - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return True since reasoning has ended - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_uses_new_token_ids_when_provided( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Regression for #43388: when caller passes new_token_ids, the - reasoner sees the exact multi-token delta rather than the - placeholder-derived window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - - end_token_id = 248069 - - reasoner = MockReasoner(tokenizer=Mock()) - # Detection mirrors the real Qwen3 parser: end token in the delta. - reasoner.is_reasoning_end_streaming = Mock( - side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids) - ) - structured_req.reasoner = reasoner - - # Scenario from #43388: async + spec decode K=4, 4 tokens accepted - # but only 1 placeholder remains (some drafts were rejected). - # The placeholder math would yield delta=[271] and miss . - # Passing new_token_ids must override that. - new_token_ids = [9, 198, end_token_id, 271] - mock_request_with_structured_output.all_token_ids = [ - 1, - 2, - 3, - 4, - 5, - ] + new_token_ids - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - # First call to is_reasoning_end_streaming was with the full - # new_token_ids (not the truncated placeholder window). - first_call = reasoner.is_reasoning_end_streaming.call_args_list[0] - _, called_delta = first_call.args - assert list(called_delta) == new_token_ids - - assert structured_req.reasoning_ended is True - assert result is True - - def test_should_advance_without_new_token_ids_falls_back( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Backward compat: callers that don't pass new_token_ids keep - the original placeholder-derived delta window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = False - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - mock_request_with_structured_output.num_computed_tokens = 5 - mock_request_with_structured_output.num_output_placeholders = 2 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # placeholder window: start = 5 - 2 = 3, delta = [4, 5] - _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0] - assert list(called_delta) == [4, 5] - assert result is False - - def test_should_advance_whole_delta_fallback_pins_last_sequence_index( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """When only the whole-delta predicate fires, the boundary is the - last index of the sequence, even if the placeholder window - overshoots it and the delta is empty. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - # Base-default shape: rescans the whole sequence, ignores the delta. - reasoner.is_reasoning_end_streaming.return_value = True - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - # placeholder window: start = 9 - 1 = 8, past the end -> delta = [] - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_trims_reasoning_prefix_for_json( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """JSON uses the common trim-then-advance path at the boundary.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.structured_output_key = ( - StructuredOutputOptions.JSON_OBJECT, - "{}", - ) - - marker = 248069 - - class MarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - structured_req.reasoner = MarkerReasoner() - - new_token_ids = [9, 198, marker, 271, 5005] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - structured_req.grammar.accept_tokens.assert_not_called() - assert structured_req.reasoning_ended is True - assert result is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271, 5005] - - def test_should_advance_engine_adapter_scans_only_the_delta( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Engine-based parsers locate the boundary from the delta alone.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = make_qwen3_reasoner() - assert reasoner.reasoning_end_token_ids == {201, 202} - reasoner.is_reasoning_end_streaming = Mock( - side_effect=AssertionError("engine path must not rescan the prefix") - ) - structured_req.reasoner = reasoner - - new_token_ids = [9, 198, 201, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271] - - def test_should_advance_engine_adapter_tool_start_ends_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 202, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_engine_adapter_no_boundary( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 198, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert not manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is False - - def test_should_advance_engine_adapter_never_reads_history( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Per-step cost must not depend on how much has been generated.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 201] - total = (1 << 20) + len(new_token_ids) - mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total) - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == total - 1 - - def test_should_advance_engine_adapter_without_new_tokens_persists( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Draft-validation callers persist the boundary like the legacy path.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7] - # placeholder window: start = 7 - 2 = 5, delta = [201, 7] - mock_request_with_structured_output.num_computed_tokens = 7 - mock_request_with_structured_output.num_output_placeholders = 2 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py new file mode 100644 index 000000000000..17bda09b3efa --- /dev/null +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from dataclasses import dataclass + +import pytest +from transformers import AutoTokenizer + +from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config.model import ModelConfig +from vllm.config.speculative import SpeculativeConfig +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.request import Request +from vllm.v1.structured_output import StructuredOutputManager + +TOKENIZER = "gpt2" +THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +EOS = "<|eos|>" # resolved to tokenizer.eos_token_id +JSON_SCHEMA = '{"type": "object"}' +BACKENDS = ("xgrammar", "guidance") +MAX_WAIT_SECONDS = 5 +NUM_SPEC_TOKENS = 8 + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +@dataclass(frozen=True) +class FlowCase: + """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" + + raw_drafts: tuple[str, ...] + expected_row_pattern: str + expected_reasoning: bool | None + expect_terminated: bool + prefix: str = "" + expected_validated: tuple[str, ...] | None = None + reasoning_ended: bool | None = False + xfail_guidance: str | None = None + + +class MockReasoner: + def __init__(self, tokenizer, marker: int | None = None): + self.marker = marker + + def is_reasoning_end(self, input_ids): + if self.marker is None: + return True + return self.marker in list(input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return self.is_reasoning_end(delta_ids) + + +def _single_token(tokenizer, text: str) -> int: + token_ids = tokenizer.encode(text) + assert len(token_ids) == 1, (text, token_ids) + return token_ids[0] + + +def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: + """Convert literal text (or EOS sentinel) to token IDs.""" + return [ + tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) + for t in texts + ] + + +def _wait_for_grammar(request: Request) -> None: + structured_req = request.structured_output_request + assert structured_req is not None + + deadline = time.time() + MAX_WAIT_SECONDS + while not structured_req._check_grammar_completion(): + if time.time() > deadline: + pytest.fail("Grammar compilation timed out") + time.sleep(0.01) + + +def _build_harness( + tokenizer, + backend: str, + prefix: str = "", + use_reasoner: bool = True, + reasoning_ended: bool | None = None, + enable_in_reasoning: bool = False, + reasoning_parser_kwargs: dict | None = None, +) -> tuple[StructuredOutputManager, Request]: + vllm_config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + device_config=DeviceConfig(device="cpu"), + structured_outputs_config=StructuredOutputsConfig( + backend=backend, + enable_in_reasoning=enable_in_reasoning, + ), + speculative_config=SpeculativeConfig( + model="[ngram]", + num_speculative_tokens=NUM_SPEC_TOKENS, + ), + ) + manager = StructuredOutputManager(vllm_config) + if use_reasoner: + manager.reasoner_cls = MockReasoner + + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + sampling_params.structured_outputs._backend = backend + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + + prompt_ids = tokenizer.encode(prefix) if prefix else [] + request = Request( + request_id=f"{backend}-flow", + prompt_token_ids=prompt_ids, + sampling_params=sampling_params, + pooling_params=None, + reasoning_ended=reasoning_ended, + reasoning_parser_kwargs=reasoning_parser_kwargs, + ) + manager.grammar_init(request) + _wait_for_grammar(request) + + structured_req = request.structured_output_request + assert request.prompt_token_ids is not None + assert structured_req is not None + assert structured_req.grammar is not None + if prompt_ids: + assert structured_req.grammar.accept_tokens( + request.request_id, + prompt_ids, + ) + return manager, request + + +def _row_pattern(bitmask) -> str: + return "".join("U" if (row == -1).all() else "C" for row in bitmask) + + +def _run_real_flow( + manager: StructuredOutputManager, + request: Request, + raw_drafts: list[int], + expected_validated: list[int], + expected_row_pattern: str, + expected_reasoning: bool | None, + expect_terminated: bool, +) -> None: + assert len(expected_validated) <= len(raw_drafts) + + validated = manager.validate_tokens(request, list(raw_drafts)) + assert validated == expected_validated + + padded = validated + [-1] * (len(raw_drafts) - len(validated)) + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: padded}, + ) + assert bitmask is not None + assert bitmask.shape[0] == len(raw_drafts) + 1 + assert _row_pattern(bitmask) == expected_row_pattern + + # grammar_bitmask() must rollback any speculative state it advanced. + assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated + + structured_req = request.structured_output_request + assert structured_req is not None + grammar = structured_req.grammar + assert grammar is not None + assert not grammar.is_terminated() + # `bitmask[i]` is the grammar state after the first `i` scheduled tokens. + # These tests commit the validated prefix, so the matching post-accept + # state is already present in the scheduled bitmask. + expected_post_accept_row = bitmask[len(expected_validated)].copy() + + # Mirror the scheduler flow: sampled tokens are appended before + # StructuredOutputManager.accept_tokens() trims them to the + # grammar-constrained suffix. + request.append_output_token_ids(list(expected_validated)) + assert manager.accept_tokens(request, list(expected_validated)) + + post_accept_bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert post_accept_bitmask is not None + assert post_accept_bitmask.shape[0] == 1 + assert (post_accept_bitmask[0] == expected_post_accept_row).all() + assert structured_req.reasoning_ended == expected_reasoning + assert grammar.is_terminated() is expect_terminated + + +FLOW_CASES = [ + pytest.param( + FlowCase( + raw_drafts=(" ", "z", " "), + expected_row_pattern="UUUU", + expected_reasoning=False, + expect_terminated=False, + ), + id="inactive", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}'), + expected_row_pattern="CCCCC", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_all_valid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=("z", ":", ' "', "b"), + expected_validated=(), + expected_row_pattern="CUUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_first_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", "z", ' "', "b"), + expected_validated=(":",), + expected_row_pattern="CCUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_later_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a": "b"}', + raw_drafts=(EOS, " ", " "), + expected_validated=(EOS,), + expected_row_pattern="CUUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + xfail_guidance=( + "guidance validate_tokens rejects EOS on an already-complete " + "object even though accept_tokens would accept it" + ), + ), + id="active_first_token_terminates", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}', EOS, " "), + expected_validated=(":", ' "', "b", '"}', EOS), + expected_row_pattern="CCCCCUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + ), + id="active_later_token_terminates", + ), + pytest.param( + FlowCase( + raw_drafts=(THINK_END, "{", ' "', "b"), + expected_row_pattern="UCCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_first_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", ' "'), + expected_row_pattern="UUCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_middle_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", " ", THINK_END), + expected_row_pattern="UUUC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_last_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "z", "{"), + expected_validated=(" ", THINK_END), + expected_row_pattern="UUCUU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_first_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "z"), + expected_validated=(" ", THINK_END, "{"), + expected_row_pattern="UUCCU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_later_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "}", EOS, " "), + expected_validated=(" ", THINK_END, "{", "}", EOS), + expected_row_pattern="UUCCCUU", + expected_reasoning=True, + expect_terminated=True, + ), + id="becomes_active_terminates", + ), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("case", FLOW_CASES) +def test_real_flow( + tokenizer, + backend: str, + case: FlowCase, +): + if backend == "guidance" and case.xfail_guidance: + pytest.xfail(case.xfail_guidance) + + reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + manager, request = _build_harness( + tokenizer, + backend, + prefix=case.prefix, + reasoning_ended=case.reasoning_ended, + reasoning_parser_kwargs=reasoner_kwargs, + ) + + raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) + expected_texts = ( + case.raw_drafts if case.expected_validated is None else case.expected_validated + ) + expected_validated = _to_token_ids(tokenizer, expected_texts) + + _run_real_flow( + manager, + request, + raw_drafts=raw_drafts, + expected_validated=expected_validated, + expected_row_pattern=case.expected_row_pattern, + expected_reasoning=case.expected_reasoning, + expect_terminated=case.expect_terminated, + ) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), + [ + pytest.param(True, False, True, id="enable_in_reasoning"), + pytest.param(True, True, False, id="reasoning_ended"), + pytest.param(False, False, False, id="no_reasoner"), + pytest.param(True, None, False, id="inferred"), + ], +) +def test_initial_constraint_activation( + tokenizer, + backend: str, + use_reasoner: bool, + reasoning_ended: bool | None, + enable_in_reasoning: bool, +): + manager, request = _build_harness( + tokenizer, + backend, + use_reasoner=use_reasoner, + reasoning_ended=reasoning_ended, + enable_in_reasoning=enable_in_reasoning, + reasoning_parser_kwargs={ + "marker": _single_token(tokenizer, THINK_END) + if reasoning_ended is not None + else None + }, + ) + + # "{" is valid JSON start; "z" is not. Truncation proves grammar is active. + open_brace = _single_token(tokenizer, "{") + z = _single_token(tokenizer, "z") + assert manager.validate_tokens(request, [open_brace, z]) == [open_brace] + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert bitmask is not None + request.append_output_token_ids([open_brace]) + assert manager.accept_tokens(request, [open_brace]) + + # Short-circuit paths must not construct a request-local reasoner (#52573). + if enable_in_reasoning or reasoning_ended is True: + structured_req = request.structured_output_request + assert structured_req is not None + assert structured_req.reasoner is None diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index cf83bb996e63..f099d7752ae5 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -197,6 +197,8 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: Used by structured engines like `xgrammar` to check if the reasoning content ends in the model output. + Must be a pure function of the input_ids. + Args: input_ids: The token IDs of the model output. @@ -210,6 +212,8 @@ def is_reasoning_end_streaming( """ Check if the reasoning content ends during a decode step. + Must be a pure function of the input_ids and delta_ids. + Args: input_ids: The entire model output token IDs. delta_ids: The last few computed tokens at the current decode step. @@ -228,8 +232,8 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: Returns: The offset within ``token_ids`` of the last reasoning token, or - ``None`` when reasoning does not end in this window. Parsers that - cannot answer from a window alone always return ``None``; callers + ``len(token_ids)`` when reasoning does not end in this window. Parsers + that cannot answer from a window alone always return ``None``; callers fall back to :meth:`is_reasoning_end_streaming`. """ return None diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index f873e36342f1..fa306a5da4a1 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -645,7 +645,7 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: for offset, token_id in enumerate(token_ids): if token_id in end_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids: list[int]) -> bool: config = self.parser_engine_config diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 3af527124d73..daf2b316d075 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,8 +70,7 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager -from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext logger = init_logger(__name__) @@ -2052,33 +2051,18 @@ def update_from_output( request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance( - request, new_token_ids=new_token_ids + if new_token_ids and not self.structured_output_manager.accept_tokens( + request, new_token_ids ): - struct_output_request = request.structured_output_request - assert struct_output_request is not None - grammar = struct_output_request.grammar - assert isinstance(grammar, StructuredOutputGrammar) - # new_token_ids can be a mixed block of reasoning content, then - # the reasoning end marker, then the start of the grammar content. - # Trim the reasoning content so the grammar only sees grammar content. - advance_token_ids = ( - self.structured_output_manager.trim_reasoning_for_advance( - request, new_token_ids - ) + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, ) - if advance_token_ids and not grammar.accept_tokens( - req_id, advance_token_ids - ): - logger.error( - "Unexpected: grammar rejected tokens %s for request %s. " - "Terminating request.", - advance_token_ids, - req_id, - ) - request.status = RequestStatus.FINISHED_ERROR - request.resumable = False - stopped = True + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True routed_experts = None if ( @@ -2434,11 +2418,9 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: continue # Add newly generated spec token ids to the request. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] - request.spec_token_ids = spec_token_ids + request.spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput @@ -2464,10 +2446,9 @@ def update_draft_token_ids_in_output( # (needed for chunked prefill case for example). del spec_token_ids[orig_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) # Pad to original number of spec tokens. num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) if num_invalid_tokens: diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index ca6adc736d06..7fd311c873c1 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import itertools import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor @@ -18,6 +17,7 @@ StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import strip_speculative_padding if TYPE_CHECKING: import numpy as np @@ -217,6 +217,100 @@ def _async_submit_fill_bitmask( ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _get_constraint_start( + self, + request: "Request", + spec_tokens: Sequence[int], + spec_tokens_committed: bool = False, + ) -> int: + """Return the index into `spec_tokens` where tokens should start being + constrained by the grammar. + Assumes `spec_tokens` are stripped of -1 padding tokens. + Returns `len(spec_tokens) + 1` if no token should be constrained after + accepting all `spec_tokens`. + + `spec_tokens_committed` is True if spec_tokens is already in + `request.all_token_ids`. + """ + if self.enable_in_reasoning: + return 0 + + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.reasoning_ended: + return 0 + + reasoner = self._get_reasoner(request) + if reasoner is None: + return 0 + + if structured_req.reasoning_ended is None: + # This should be removed here, but since `openai_gptoss` + # is an independent code path, it is kept for now. + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, + # it can be removed. + if reasoner.is_reasoning_end(request.prompt_token_ids or []): + return 0 + structured_req.reasoning_ended = False + + num_spec_tokens = len(spec_tokens) + if num_spec_tokens <= 0: + return num_spec_tokens + 1 + + # Use `find_reasoning_end_offset` to find constraint start if supported + if ( + isinstance(reasoner, ParserEngineReasoningAdapter) + and reasoner.reasoning_end_token_ids + ): + offset = reasoner.find_reasoning_end_offset(spec_tokens) + if offset is not None: + return offset + 1 + + # Fallback to `find_reasoning_end_offset` + # TODO: Build a read-only Sequence view over all_token_ids + # instead of copying the entire all_token_ids into input_ids. + if spec_tokens_committed: + if not spec_tokens or not reasoner.is_reasoning_end_streaming( + request.all_token_ids, spec_tokens + ): + return num_spec_tokens + 1 + + input_ids = request.all_token_ids.copy() + delta_ids = list(spec_tokens) + else: + input_ids = request.all_token_ids.copy() + input_ids.extend(spec_tokens) + if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): + return num_spec_tokens + 1 + delta_ids = list(spec_tokens) + + for i in range(num_spec_tokens - 1, 0, -1): + input_ids.pop() + delta_ids.pop() + if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): + return i + 1 + return 1 + + def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: + """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" + if not request.use_structured_output: + return spec_tokens + + spec_tokens = strip_speculative_padding(spec_tokens) + constraint_start = self._get_constraint_start(request, spec_tokens) + if constraint_start >= len(spec_tokens): + return spec_tokens + + structured_req = request.structured_output_request + if TYPE_CHECKING: + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) + prefix = spec_tokens[:constraint_start] + validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + return prefix + validated + def grammar_bitmask( self, requests: dict[str, "Request"], @@ -264,7 +358,7 @@ def grammar_bitmask( if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) + apply_bitmask = self._get_constraint_start(request, ()) == 0 batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -288,76 +382,47 @@ def grammar_bitmask( grammar = structured_output_request.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) - - reasoner = None if apply_bitmask else self._get_reasoner(request) + req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) + constraint_start = self._get_constraint_start( + request, strip_speculative_padding(req_tokens) + ) state_advancements = 0 - post_reasoning_end_in_window = False - req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - # Locate the reasoning-end marker in the draft window once. - # Rejected-draft padding (-1) is trailing, so the window is - # cut there and the reasoner never sees a placeholder. - window_start = len(request.all_token_ids) - reasoning_end_index: int | None = None - if reasoner is not None and req_tokens: - window = req_tokens - if -1 in window: - window = window[: window.index(-1)] - reasoning_end_index = self._find_reasoning_end_index( - reasoner, - request.all_token_ids, - window_start, - window, - delta_appended=False, - ) + seen_padding = False + failed = False for i, token in enumerate(req_tokens): + apply_bitmask = ( + not failed and not seen_padding and i >= constraint_start + ) self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) - advance_grammar = apply_bitmask if token == -1: - apply_bitmask = False - advance_grammar = False - elif not apply_bitmask and window_start + i == reasoning_end_index: - # Reasoning ended mid-window. Constrain the rest - # of the window via bitmask. Skip grammar advance - # through the marker (it is reasoning content); - # try to advance through subsequent drafts so the - # next bitmask row reflects the post-advance state, - # but tolerate rejection since those drafts predate - # the bitmask and are not guaranteed valid. - apply_bitmask = True - advance_grammar = False - post_reasoning_end_in_window = True - if advance_grammar and not grammar.is_terminated(): - if post_reasoning_end_in_window: - accepted = bool(grammar.validate_tokens([token])) - if accepted: - accepted = grammar.accept_tokens(req_id, [token]) - else: - accepted = grammar.accept_tokens(req_id, [token]) - if accepted: + seen_padding = True + elif apply_bitmask: + if not grammar.is_terminated() and grammar.accept_tokens( + req_id, [token] + ): state_advancements += 1 - elif not post_reasoning_end_in_window: - raise AssertionError( - (token, req_id, scheduled_spec_decode_tokens) + else: + failed = True + logger.error( + "Unexpected: grammar terminated or rejected draft " + "token %s for request %s during bitmask fill.", + token, + req_id, ) cumulative_index += 1 + # Diffusion LLMs don't sample a bonus token after the # scheduled positions, so skip its bitmask in that case. if not (self.vllm_config.model_config.is_diffusion and req_tokens): - # bonus_apply must be True when the bonus-row position - # should be grammar-constrained. Two triggers: - # - should_fill_bitmask(request): reasoning was already - # over at step start (or no reasoner / - # enable_in_reasoning). - # - apply_bitmask: reasoning ended mid-window in this - # call and was flipped True after the marker; - # should_fill_bitmask still returns False here because - # reasoning_ended is only persisted later by - # should_advance. - bonus_apply = self.should_fill_bitmask(request) or apply_bitmask + bonus_apply = ( + not failed + and not seen_padding + and constraint_start <= len(req_tokens) + ) self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 + if state_advancements > 0: grammar.rollback(state_advancements) @@ -370,161 +435,29 @@ def grammar_bitmask( # and deserialization when sending this to the GPU workers. return bitmask_tensor.numpy() - def should_fill_bitmask(self, request: "Request") -> bool: - # NOTE (Hanchen) if enable_in_reasoning is True, it means that - # the model needs to be constrained in reasoning. So we should always - # enable the bitmask filling. - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): + def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: + """Advance grammar with accepted tokens. Returns False on rejection.""" + if not request.use_structured_output or not new_token_ids: return True - reasoner = self._get_reasoner(request) - if reasoner is not None: + structured_req = request.structured_output_request + if TYPE_CHECKING: assert structured_req is not None - if structured_req.reasoning_ended is None: - # This should be removed here, but since `openai_gptoss` - # is an independent code path, it is kept for now. - # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, - # it can be removed. - structured_req.reasoning_ended = reasoner.is_reasoning_end( - request.prompt_token_ids or [] - ) - return structured_req.reasoning_ended - return True - - def should_advance( - self, - request: "Request", - new_token_ids: list[int] | None = None, - ) -> bool: - if not request.use_structured_output: - return False - - # To determine whether we can advance the FSM. - # Supports thinking usage where we skip the reasoning components. + grammar = structured_req.grammar if TYPE_CHECKING: - assert request.structured_output_request is not None - assert request.structured_output_request.grammar is not None - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): - return True - - # by default, we should always advance - # for cases that don't use thinking mode. - reasoner = self._get_reasoner(request) - if reasoner is None: - return True + assert isinstance(grammar, StructuredOutputGrammar) - assert structured_req is not None - - # Check if reasoning ends in *this* step. - # When the caller passes new_token_ids (the tokens that were just - # appended this step), use it directly as the delta window. The - # placeholder-derived fallback assumes num_output_placeholders == - # len(new_token_ids), which breaks under async scheduling + spec - # decode when some drafts are rejected (#43388): the placeholder - # count remains > 0 after the step and the computed delta window - # starts past the reasoning-end marker. - all_token_ids = request.all_token_ids - if new_token_ids: - # The tokens were already appended this step, so the step window - # starts exactly len(new_token_ids) from the end. - start = len(all_token_ids) - len(new_token_ids) - delta_ids: Sequence[int] = new_token_ids - else: - # Draft-validation callers arrive here after update_from_output - # has already persisted any boundary found this step, so this - # window only needs to cover tokens not examined yet. - delta_from = request.num_computed_tokens - request.num_output_placeholders - start = ( - delta_from - if delta_from >= 0 - else max(len(all_token_ids) + delta_from, 0) - ) - delta_ids = all_token_ids[start:] - end_index = self._find_reasoning_end_index( - reasoner, all_token_ids, start, delta_ids, delta_appended=True + constraint_start = self._get_constraint_start( + request, new_token_ids, spec_tokens_committed=True ) - if end_index is not None: - structured_req.reasoning_ended = True - # Record the boundary so the scheduler can exclude reasoning tokens. - structured_req.reasoning_end_token_index = end_index + # Early return only when the constraint hasn't started. + # Otherwise, latch structured_req.reasoning_ended. + if constraint_start > len(new_token_ids): return True - - return False - - @staticmethod - def _find_reasoning_end_index( - reasoner: "ReasoningParser", - all_token_ids: Sequence[int], - start: int, - delta_ids: Sequence[int], - *, - delta_appended: bool, - ) -> int | None: - """Locates the token that ends reasoning within ``delta_ids``. - - ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says - whether they are already in ``all_token_ids`` (accepted tokens) or - not (speculative drafts). - - Legacy parsers are probed one token at a time. Accepted tokens get a - whole-delta check as well, and a marker seen only that way reports - the final token. Drafts skip it: some predicates are order-sensitive - over a multi-token window and would miss the marker. - - Returns: - Absolute index of the last reasoning token, or ``None``. - """ - if ( - isinstance(reasoner, ParserEngineReasoningAdapter) - and reasoner.reasoning_end_token_ids - ): - offset = reasoner.find_reasoning_end_offset(delta_ids) - return None if offset is None else start + offset - - if delta_appended and not reasoner.is_reasoning_end_streaming( - all_token_ids, delta_ids - ): - return None - prefix = list(itertools.islice(all_token_ids, start)) - for offset, token in enumerate(delta_ids): - prefix.append(token) - if reasoner.is_reasoning_end_streaming(prefix, [token]): - return start + offset - if delta_appended: - return len(all_token_ids) - 1 - return None - - def trim_reasoning_for_advance( - self, request: "Request", new_token_ids: list[int] - ) -> list[int]: - """Drops reasoning content from tokens about to advance the grammar. - - When reasoning ends mid-step (see should_advance), the step's output - still contains reasoning tokens up to and including the end marker. - Those are not grammar content: feeding them to accept_tokens makes - the grammar reject the marker and kills the request (#44006). - - Returns: - The suffix of ``new_token_ids`` that follows the reasoning-end - marker. Steps fully after the boundary are returned unchanged. - """ - structured_req = request.structured_output_request - if structured_req is None: - return new_token_ids - end_idx = structured_req.reasoning_end_token_index - if end_idx is None: - return new_token_ids - first_idx = len(request.all_token_ids) - len(new_token_ids) - num_reasoning = end_idx + 1 - first_idx - if num_reasoning <= 0: - return new_token_ids - return new_token_ids[num_reasoning:] + structured_req.reasoning_ended = True + return grammar.accept_tokens( + request.request_id, new_token_ids[constraint_start:] + ) def clear_backend(self) -> None: if self.backend is not None: diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index b0f4f3c3a8fd..dce0038782ce 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -25,12 +25,6 @@ class StructuredOutputRequest: Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None ) = None reasoning_ended: bool | None = None - # Absolute index into the request's all_token_ids of the last reasoning - # token (the reasoning-end marker). Tokens at or before this index are - # reasoning content and must never be fed to the grammar. Only set when - # reasoning ends in a step whose tokens the scheduler advances immediately - # (structural tags + speculative decoding, see #42452). - reasoning_end_token_index: int | None = None reasoning_parser_kwargs: dict[str, Any] | None = None # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. From a2e8f1ea94b1332428e2282744cec12ce7c75d88 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:04:09 +0800 Subject: [PATCH 257/774] [XPU][CI]Skip test_chat_completion_with_tools in Intel GPU CI (#56934) Signed-off-by: zengxian Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/entrypoints_intel.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index efed04c1b383..e8006bac8f66 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -134,7 +134,8 @@ steps: 'cd tests && pytest -v -s tool_use --deselect="tests/tool_use/mistral/test_mistral_tool_calls.py::test_tool_call_none_with_tools[ministral-3b]" - --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" && + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[llama3.2]" && pytest -v -s entrypoints/tool_parsers && pytest -v -s entrypoints/generate && pytest -v -s entrypoints/anthropic && From 4bac767695d2dfecd4098a77f71dacd100061247 Mon Sep 17 00:00:00 2001 From: Junhao Shen Date: Tue, 15 Sep 2026 22:33:51 +0800 Subject: [PATCH 258/774] [Perf][Nemotron] Skip redundant latent-MoE all-reduce at TP>1 (~13% decode win) (#52301) Signed-off-by: Junhao Shen Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Misha Goin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/fused_moe/runner/moe_runner.py | 9 +++++++++ vllm/model_executor/models/nemotron_h.py | 12 ++++++++++++ 2 files changed, 21 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 2868690ebeb1..6e63563ab861 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -457,9 +457,18 @@ def _maybe_reduce_routed_output_before_transform( Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm. TP partial routed outputs must be summed in latent space before such transforms are applied. + + A transform that commutes with the TP sum is exempt: if + ``sum_r T(x_r) == T(sum_r x_r)``, applying the transform to the local + partial output and letting the existing late all-reduce sum the + combined result is equivalent, and costs one collective instead of two. + Such a transform opts out by setting ``reduce_commutative = True``. + The default is False, so transforms that do not declare themselves + keep being reduced early. """ if ( self.routed_output_transform is not None + and not getattr(self.routed_output_transform, "reduce_commutative", False) and not self.moe_config.is_sequence_parallel and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not fused_output_is_reduced diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index c1bc4c6823f7..ea8a28ebd9d7 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -49,6 +49,7 @@ QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 @@ -204,6 +205,17 @@ def __init__( disable_tp=self.is_sequence_parallel, prefix=f"{prefix}.fc2_latent_proj", ) + # A bias-free, unquantized linear commutes with the TP sum + # (sum_r W x_r == W sum_r x_r), so one reduce after the transform + # suffices. Test the layer, not the model-wide `quant_config`: ModelOpt + # excludes the latent projections, so quantized checkpoints still get + # an UnquantizedLinearMethod here. + self.fc2_latent_proj.reduce_commutative = ( + not config.mlp_bias + and isinstance( + self.fc2_latent_proj.quant_method, UnquantizedLinearMethod + ) + ) else: self.fc1_latent_proj = None self.fc2_latent_proj = None From 0136df94b0d75732be6c66f620f51289c888e41d Mon Sep 17 00:00:00 2001 From: zack <51604064+luoyuctl@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:35:47 +0800 Subject: [PATCH 259/774] [Bugfix][Spec Decode][MoE] Avoid uninitialized EPLB state in DeepSeek V4.1 DSpark drafter (#56387) Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com> --- tests/v1/worker/test_dspark_utils.py | 51 +++++++++++++++++++ .../v1/worker/gpu/spec_decode/dspark/utils.py | 32 +++++++++--- 2 files changed, 77 insertions(+), 6 deletions(-) create mode 100644 tests/v1/worker/test_dspark_utils.py diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py new file mode 100644 index 000000000000..94655f192c89 --- /dev/null +++ b/tests/v1/worker/test_dspark_utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, field + +from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config + + +@dataclass +class _FakeEPLBConfig: + num_redundant_experts: int = 0 + + +@dataclass +class _FakeParallelConfig: + pipeline_parallel_size: int = 2 + tensor_parallel_size: int = 8 + enable_eplb: bool = True + eplb_config: _FakeEPLBConfig = field( + default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32) + ) + enable_elastic_ep: bool = True + + def __post_init__(self) -> None: + if not self.enable_eplb and self.eplb_config.num_redundant_experts: + raise ValueError("redundant experts require EPLB") + if self.enable_elastic_ep and not self.enable_eplb: + raise ValueError("elastic EP requires EPLB") + + +def test_dspark_parallel_config_disables_eplb_atomically(): + target_config = _FakeParallelConfig() + + draft_config = _get_dspark_parallel_config( + target_config, + tensor_parallel_size=4, + ) + + assert target_config.pipeline_parallel_size == 2 + assert target_config.tensor_parallel_size == 8 + assert target_config.enable_eplb + assert target_config.eplb_config.num_redundant_experts == 32 + assert target_config.enable_elastic_ep + + assert draft_config is not target_config + assert draft_config.pipeline_parallel_size == 1 + assert draft_config.tensor_parallel_size == 4 + assert not draft_config.enable_eplb + assert draft_config.eplb_config.num_redundant_experts == 0 + assert not draft_config.enable_elastic_ep + assert draft_config.eplb_config is not target_config.eplb_config diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index db1e17a9f23b..bf1a5761b2ef 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -3,7 +3,7 @@ import torch.nn as nn -from vllm.config import ModelConfig, VllmConfig, replace +from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -31,6 +31,29 @@ def _resolve_dspark_attention_backend( return None +def _get_dspark_parallel_config( + parallel_config: ParallelConfig, + tensor_parallel_size: int, +) -> ParallelConfig: + if parallel_config.enable_eplb: + logger.warning_once( + "EPLB is disabled for the DSpark draft model. EPLB remains enabled " + "for the target model." + ) + + return replace( + parallel_config, + pipeline_parallel_size=1, + tensor_parallel_size=tensor_parallel_size, + enable_eplb=False, + eplb_config=replace( + parallel_config.eplb_config, + num_redundant_experts=0, + ), + enable_elastic_ep=False, + ) + + def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: speculative_config = vllm_config.speculative_config assert speculative_config is not None @@ -54,12 +77,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, - parallel_config=replace( + parallel_config=_get_dspark_parallel_config( vllm_config.parallel_config, - pipeline_parallel_size=1, - tensor_parallel_size=( - speculative_config.draft_parallel_config.tensor_parallel_size - ), + speculative_config.draft_parallel_config.tensor_parallel_size, ), attention_config=replace( vllm_config.attention_config, From f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 22:50:26 +0800 Subject: [PATCH 260/774] [Frontend] Share max_num_queued_reqs across API server processes (#54746) Signed-off-by: chaunceyjiang --- .../test_api_server_process_manager.py | 39 ++++++++++++++ tests/v1/engine/test_admission_control.py | 39 ++++++++++++++ vllm/v1/engine/admission_control.py | 52 +++++++++++++++++++ vllm/v1/engine/async_llm.py | 24 +++++++-- vllm/v1/engine/output_processor.py | 14 ++++- vllm/v1/utils.py | 11 ++++ 6 files changed, 173 insertions(+), 6 deletions(-) create mode 100644 vllm/v1/engine/admission_control.py diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 40fefc237404..668f11544858 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -5,6 +5,7 @@ import socket import threading import time +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -14,6 +15,7 @@ exit_before_report_worker, ) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, get_engine_client_zmq_addr, @@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None): print("Mock worker completed successfully") +def update_admission_stats_worker(listen_address, sock, args, client_config): + stats = SharedAdmissionStats( + client_config, + client_config["client_count"], + client_config["client_index"], + ) + value = client_config["client_index"] + 1 + stats.set_num_requests(value) + args.barrier.wait() + args.results.put(stats.get_num_requests()) + + # Module-level stub for the gather_actual_addresses test. Must be # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): @@ -131,6 +145,31 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +@pytest.mark.skip_global_cleanup +def test_api_server_processes_share_admission_stats(api_server_args): + args = api_server_args.copy() + spawn_context = multiprocessing.get_context("spawn") + worker_args = SimpleNamespace( + barrier=spawn_context.Barrier(args["num_servers"]), + results=spawn_context.Queue(), + max_num_queued_reqs=10, + ) + args["args"] = worker_args + args["target_server_fn"] = update_admission_stats_worker + manager = APIServerProcessManager(**args) + + try: + for process in manager.processes: + process.join(timeout=30) + assert process.exitcode == 0 + + assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3 + finally: + manager.shutdown() + worker_args.results.close() + worker_args.results.join_thread() + + @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker) def test_wait_for_completion_or_failure(api_server_args): """Test that wait_for_completion_or_failure works with failures.""" diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index a03718ab5a15..219850f5d489 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -14,6 +14,7 @@ import argparse import asyncio +import multiprocessing from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import human_readable_int +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor @@ -66,6 +68,7 @@ def _make_async_llm( llm.output_processor = MagicMock() llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + llm.admission_stats = None return llm @@ -84,6 +87,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig: ) +def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]: + counters = multiprocessing.RawArray( + "q", SharedAdmissionStats.num_counters(client_count) + ) + return [ + SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index) + for index in range(client_count) + ] + + +@pytest.mark.skip_global_cleanup +def test_shared_admission_stats_aggregate_api_servers(): + server_0, server_1 = _make_shared_stats() + + server_0.set_num_requests(2) + assert server_0.get_num_requests() == 2 + server_1.set_num_requests(3) + assert server_0.get_num_requests() == server_1.get_num_requests() == 5 + + server_0.set_num_requests(1) + assert server_0.get_num_requests() == 4 + + +@pytest.mark.skip_global_cleanup +def test_shared_request_snapshot_enforces_global_limit(): + server_0, server_1 = _make_shared_stats() + server_0.set_num_requests(2) + llm = _make_async_llm(max_num_queued_reqs=2) + llm.admission_stats = server_1 + + with pytest.raises(QueueOverflowError): + llm.check_admission() + + assert server_0.get_num_requests() == 2 + + # --------------------------------------------------------------------------- # Exception classes # --------------------------------------------------------------------------- diff --git a/vllm/v1/engine/admission_control.py b/vllm/v1/engine/admission_control.py new file mode 100644 index 000000000000..b72d50f3465d --- /dev/null +++ b/vllm/v1/engine/admission_control.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any, Protocol, cast + + +class AdmissionCounterBuffer(Protocol): + def __getitem__(self, index: int) -> int: ... + + def __setitem__(self, index: int, value: int) -> None: ... + + +class SharedAdmissionStats: + """Lock-free admission stats shared by API server processes. + + Each process writes exclusively to its own cache-line-sized slot. Readers + aggregate the slots, avoiding both inter-process locks and cache-line + contention on the request path. + """ + + # Space per-process counters 64 bytes apart to avoid false sharing on CPUs + # with 64-byte cache lines: one process's writes would otherwise invalidate + # the cache line containing another process's counter. + CACHE_LINE_BYTES = 64 + COUNTER_BYTES = 8 + COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES + + def __init__( + self, + client_addresses: dict[str, Any], + client_count: int = 1, + client_index: int = 0, + ): + self._counters = cast( + AdmissionCounterBuffer, + client_addresses.pop("mp_admission_counters"), + ) + self._client_count = client_count + self._slot_offset = client_index * self.COUNTERS_PER_SLOT + + @classmethod + def num_counters(cls, client_count: int) -> int: + return client_count * cls.COUNTERS_PER_SLOT + + def set_num_requests(self, count: int) -> None: + self._counters[self._slot_offset] = count + + def get_num_requests(self) -> int: + return sum( + self._counters[client_index * self.COUNTERS_PER_SLOT] + for client_index in range(self._client_count) + ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 91618c409362..bf3f5f420ae8 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,6 +44,7 @@ from vllm.utils.async_utils import cancel_task_threadsafe from vllm.utils.collection_utils import as_list from vllm.v1.engine import EngineCoreRequest, PauseMode +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor @@ -144,12 +145,20 @@ def __init__( # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) + self.admission_stats = ( + SharedAdmissionStats(client_addresses, client_count, client_index) + if client_addresses is not None + and "mp_admission_counters" in client_addresses + else None + ) + # Converts EngineCoreOutputs --> RequestOutput. self.output_processor = OutputProcessor( renderer.tokenizer, log_stats=self.log_stats, stream_interval=self.vllm_config.scheduler_config.stream_interval, tracing_enabled=tracing_endpoint is not None, + admission_stats=self.admission_stats, ) # EngineCore (starts the engine in background process). @@ -291,8 +300,9 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: Both limits return HTTP 503 (Service Unavailable) so that load balancers and client SDKs retry on a different instance. - - ``max_num_queued_reqs``: hard cap on the number of unfinished requests - (waiting + running). A request with ``n > 1`` counts as ``n`` slots. + - ``max_num_queued_reqs``: approximate global cap on the number of + unfinished requests (waiting + running). A request with ``n > 1`` + counts as ``n`` slots. - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt tokens of requests still in prefill. @@ -312,13 +322,17 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: """ max_num_reqs = self.scheduler_config.max_num_queued_reqs if max_num_reqs is not None: - current = self.get_num_unfinished_requests() - if current + n > max_num_reqs: + current_requests = ( + self.admission_stats.get_num_requests() + if self.admission_stats is not None + else self.get_num_unfinished_requests() + ) + if current_requests + n > max_num_reqs: logger.info( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, - current, + current_requests, n, max_num_reqs, ) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 419e3aa8e99d..257287243dbd 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -5,7 +5,7 @@ from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import numpy as np import torch @@ -41,6 +41,9 @@ ) from vllm.v1.outputs import SamplingMaskLists +if TYPE_CHECKING: + from vllm.v1.engine.admission_control import SharedAdmissionStats + # shared empty CPU tensor used as a placeholder pooling output EMPTY_CPU_TENSOR = torch.empty(0, device="cpu") @@ -455,6 +458,7 @@ def __init__( log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False, + admission_stats: "SharedAdmissionStats | None" = None, ): self.log_stats = log_stats self.tokenizer = tokenizer @@ -464,6 +468,7 @@ def __init__( self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) self.lora_states = LoRARequestStates(log_stats) self.tracing_enabled = tracing_enabled + self.admission_stats = admission_stats def get_num_unfinished_requests(self): return len(self.request_states) @@ -553,6 +558,7 @@ def abort_requests(self, request_ids: Iterable[str], internal: bool) -> list[str child_reqs = self.abort_requests(child_reqs, internal=True) request_ids_to_abort.extend(child_reqs) self.parent_requests.pop(request_id, None) + self._update_admission_stats() return request_ids_to_abort def add_request( @@ -585,6 +591,7 @@ def add_request( # Track the external_req_id -> [internal_req_id, ...] mapping self.external_req_ids[req_state.external_req_id].append(request_id) + self._update_admission_stats() def _update_streaming_request_state( self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None @@ -763,6 +770,11 @@ def _finish_request(self, req_state: RequestState) -> None: parent_req = req_state.parent_req if parent_req and not parent_req.child_requests: self.parent_requests.pop(parent_req.request_id, None) + self._update_admission_stats() + + def _update_admission_stats(self) -> None: + if self.admission_stats is not None: + self.admission_stats.set_num_requests(self.get_num_unfinished_requests()) def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None): self.lora_states.update_scheduler_stats(scheduler_stats) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b0411dc9e8b1..da17110d80af 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -209,6 +209,15 @@ def __init__( self.processes: list[BaseProcess] = [] self._address_pipes: list[connection.Connection] = [] + admission_counters = None + if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: + from vllm.v1.engine.admission_control import SharedAdmissionStats + + admission_counters = spawn_context.RawArray( + "q", + SharedAdmissionStats.num_counters(num_servers), + ) + for i, in_addr, out_addr in zip( range(num_servers), input_addresses, output_addresses ): @@ -218,6 +227,8 @@ def __init__( "client_count": num_servers, "client_index": i, } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: client_config["stats_update_address"] = stats_update_address if tensor_queue is not None: From 1e47ec00d2e6ad911969508068b72f7d1aae1968 Mon Sep 17 00:00:00 2001 From: Abhinav Verma Date: Tue, 15 Sep 2026 20:27:32 +0530 Subject: [PATCH 261/774] [Bugfix][Gemma 4] Don't read fft_length when profiling unified audio (#56721) Signed-off-by: Abhinav Verma Co-authored-by: Claude Opus 5 --- .../processing/test_gemma4_unified.py | 44 ++++++++++++++++++- vllm/model_executor/models/gemma4_mm.py | 25 ++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 67a81ddb7b41..fe544e233884 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -6,9 +6,13 @@ import pytest import torch from PIL import Image as PILImage +from transformers.models.gemma4 import Gemma4AudioFeatureExtractor from vllm.exceptions import VLLMValidationError -from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs +from vllm.model_executor.models.gemma4_mm import ( + Gemma4ImagePixelInputs, + _dummy_audio_num_samples, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -204,3 +208,41 @@ def test_limit_mm_per_prompt( mm_items=processor.info.parse_mm_data(mm_data), hf_processor_mm_kwargs={}, ) + + +def test_dummy_audio_length_falls_back_for_the_unified_extractor(): + """The unified audio feature extractor does not define ``fft_length``. + + ``Gemma4DummyInputsBuilder.get_dummy_mm_data`` read that attribute + unconditionally, so memory profiling raised ``AttributeError`` for every + unified Gemma 4 deployment in which audio is the modality with the largest + per-item token count -- which is the case as soon as ``audio_seq_length`` + is raised above the video budget. + """ + + class _UnifiedFeatureExtractor: + sampling_rate = 16000 + + class _UnifiedProcessor: + feature_extractor = _UnifiedFeatureExtractor() + audio_seq_length = 3000 + audio_ms_per_token = 40 + + # 3000 tokens * 40 ms = 120 s at 16 kHz + assert _dummy_audio_num_samples(_UnifiedProcessor()) == 120 * 16000 + + +def test_dummy_audio_length_still_uses_fft_length_for_the_tower_extractor(): + """The mel/tower extractor keeps its existing behaviour.""" + + class _TowerFeatureExtractor(Gemma4AudioFeatureExtractor): + def __init__(self): # noqa: D107 - bypass the real feature-extractor init + self.fft_length = 512 + self.sampling_rate = 16000 + + class _TowerProcessor: + feature_extractor = _TowerFeatureExtractor() + audio_seq_length = 750 + audio_ms_per_token = 40 + + assert _dummy_audio_num_samples(_TowerProcessor()) == 512 diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 2dddc72b218d..2fe2752475a0 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -24,6 +24,7 @@ from torch import nn from transformers import AutoModel, BatchFeature from transformers.models.gemma4 import ( + Gemma4AudioFeatureExtractor, Gemma4Config, Gemma4Processor, Gemma4VisionConfig, @@ -456,6 +457,28 @@ def get_video_repl( # --------------------------------------------------------------------------- +def _dummy_audio_num_samples(processor: Any) -> int: + """Length in samples of a worst-case dummy audio item. + + ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not + defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free + Gemma 4 variant uses. Reading it unconditionally therefore raises + ``AttributeError`` during memory profiling on that variant, which is reached + whenever audio is the modality with the largest per-item token count. + + For any other extractor, fall back to the audio budget the processor itself + declares -- the same pair of values used to warn about over-long audio + elsewhere in this file. + """ + feature_extractor = processor.feature_extractor + if isinstance(feature_extractor, Gemma4AudioFeatureExtractor): + return feature_extractor.fft_length + + sampling_rate = feature_extractor.sampling_rate + max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0 + return int(max_duration_s * sampling_rate) + + class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) @@ -508,7 +531,7 @@ def get_dummy_mm_data( } if num_audios > 0: - audio_len = processor.feature_extractor.fft_length + audio_len = _dummy_audio_num_samples(processor) data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, From 9446ea168067c141992f86c3a796e8551f37139e Mon Sep 17 00:00:00 2001 From: drslark <96540755+drslark@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:59:02 +0800 Subject: [PATCH 262/774] [Bugfix][Spec Decode] Only create draft_id_to_target_id when draft vocab differs (#53458) Signed-off-by: drslark Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: GLM-5.3 Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py --- vllm/model_executor/models/deepseek_eagle3.py | 12 ++++++++---- vllm/model_executor/models/llama_eagle3.py | 12 ++++++++---- vllm/model_executor/models/qwen3_eagle3.py | 12 ++++++++---- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 80c41501e6e9..48d019cd357b 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -305,10 +305,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None def embed_input_ids( self, diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index c3d956a946d1..68e21c9dc284 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -335,10 +335,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = speculative_config.parallel_drafting diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py index 39d0173b3ce9..745c77a97730 100644 --- a/vllm/model_executor/models/qwen3_eagle3.py +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -309,10 +309,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting From 836bb3839ffefcda8283ea7d41671a89e1a613df Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 23:03:58 +0800 Subject: [PATCH 263/774] [Bugfix] Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches (#56969) Signed-off-by: chaunceyjiang --- .../test_flashinfer_mla_sparse_sm90.py | 26 ++++++++++++++----- .../mla/flashinfer_mla_sparse_sm90.py | 8 ++++++ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index b20c16330659..8dc082b08bbc 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -151,32 +151,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert kwargs == {} -def test_builder_attaches_its_state(monkeypatch): +@pytest.mark.parametrize("use_mha", [False, True]) +@pytest.mark.parametrize("num_decodes", [0, 1]) +def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes): + """MHA prefill rows must not make the MQA kernel read beyond its query.""" builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 builder._async_scheduling = False builder.state = FakeState(TOPK) + builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha) metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata) metadata.state = None + metadata.num_prefills = 1 + metadata.num_decode_tokens = num_decodes monkeypatch.setattr( sm90_mod.FlashInferMLASparseMetadataBuilder, "build", lambda *_args, **_kwargs: metadata, ) cam = SimpleNamespace( - num_reqs=1, - query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), - seq_lens=torch.tensor([1], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32), + num_reqs=num_decodes + 1, + query_start_loc_cpu=torch.tensor( + [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32 + ), + seq_lens_cpu_upper_bound=torch.tensor( + [1402, 5] if num_decodes else [5], dtype=torch.int32 + ), positions=None, ) result = builder.build(0, cam) assert result.state is builder.state - assert builder.state.plan_calls[0][0] == 1 - assert builder.state.plan_calls[0][1].tolist() == [1] + expected_lens = [1402] if num_decodes else [] + if not use_mha: + expected_lens += [1, 2, 3, 4, 5] + assert builder.state.plan_calls[0][0] == len(expected_lens) + assert builder.state.plan_calls[0][1].tolist() == expected_lens def test_plan_uses_state_params(monkeypatch): diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 90dc4cd94738..830bc52ab939 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -277,6 +277,7 @@ def __init__( attention_layer = vllm_config.compilation_config.static_forward_context[ layer_names[0] ] + self._attention_layer = attention_layer impl = attention_layer.impl if not isinstance(impl, FlashInferMLASparseSM90Impl): raise TypeError( @@ -376,6 +377,13 @@ def build( # Replan every step outside any CUDA graph capture with this step's # exact per-row lengths; captured runs read the refreshed buffers. num_rows, kv_lens = self._kv_lens_host(common_attn_metadata) + # MHA prefills are removed from q before forward_mqa. Plan only the + # decode prefix in that case, using the same routing decision as MLA. + if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( + metadata + ): + num_rows = metadata.num_decode_tokens + kv_lens = kv_lens[:num_rows] self.state.plan(num_rows, kv_lens) metadata.state = self.state return metadata From 241e9391ed85f6c967432135427717a529776a63 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 09:58:30 -0700 Subject: [PATCH 264/774] [Bugfix][Attention] Stabilize sparse-MLA DCP for GLM PCP evals (#55879) Signed-off-by: khluu Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Tyler Michael Smith Co-authored-by: Jason Yao --- .../GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml | 5 +- .../configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 2 + tests/v1/attention/test_dcp_a2a_pack_mask.py | 16 +- .../test_flashinfer_sparse_mla_workspace.py | 166 ++++++++++++++++++ .../v1/attention/test_indexer_dcp_localize.py | 17 ++ .../test_indexer_deepseek_v4_slot_mapping.py | 32 ++++ .../backends/mla/flashinfer_mla_sparse.py | 125 ++++++++++++- vllm/v1/attention/backends/mla/indexer.py | 46 ++++- vllm/v1/attention/backends/utils.py | 8 +- vllm/v1/attention/ops/dcp.py | 7 + 10 files changed, 411 insertions(+), 13 deletions(-) create mode 100644 tests/v1/attention/test_flashinfer_sparse_mla_workspace.py diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml index b9d2675431e7..8f3e9ce918ef 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -3,10 +3,13 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Preserve KV-cache headroom after CUDA graph memory profiling and bound the +# DCP-aware FlashInfer workspace to 2.39 GiB at the 4,096-token scheduler cap. server_args: >- --compilation-config '{"cudagraph_mode":"PIECEWISE"}' + --gpu-memory-utilization 0.96 --max-model-len 4096 - --max-num-batched-tokens 32768 + --max-num-batched-tokens 4096 --safetensors-load-strategy prefetch --moe-backend flashinfer_cutlass --prefill-context-parallel-size 4 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index 72f7fdf1ae91..bc5c11fa53e2 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml @@ -3,8 +3,10 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation. server_args: >- --compilation-config '{"cudagraph_mode":"NONE"}' + --gpu-memory-utilization 0.85 --max-model-len 4096 --max-num-batched-tokens 32768 --safetensors-load-strategy prefetch diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py index 36032e7a69a0..cff1228baf03 100644 --- a/tests/v1/attention/test_dcp_a2a_pack_mask.py +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -13,7 +13,7 @@ mask_dcp_empty_shards_, ) -pytestmark = pytest.mark.skipif( +requires_accelerator = pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="needs a CUDA or ROCm device for the Triton pack kernel", ) @@ -44,6 +44,7 @@ def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc) @pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)]) @pytest.mark.parametrize("tokens_per_req", [1, 2, 3]) @pytest.mark.parametrize("num_pad_rows", [0, 2]) +@requires_accelerator @pytest.mark.parametrize( "seq_lens_list", [ @@ -97,6 +98,7 @@ def test_fused_mask_matches_eager( @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("seed", range(8)) +@requires_accelerator def test_fused_mask_matches_eager_ragged(dtype, seed): """Non-uniform per-request query lengths. @@ -137,6 +139,7 @@ def test_fused_mask_matches_eager_ragged(dtype, seed): ) +@requires_accelerator def test_mask_disabled_is_unmasked(): """Passing no seq_lens/query_start_loc must leave the LSE untouched.""" torch.manual_seed(0) @@ -149,3 +152,14 @@ def test_mask_disabled_is_unmasked(): packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None) assert torch.isfinite(packed.float()).all() + + +def test_eager_mask_handles_rank_with_no_local_sequences(): + """Padded graph rows are all empty when a DCP rank has no sequences.""" + lse = torch.randn(8, 4, dtype=torch.float32) + seq_lens = torch.empty(0, dtype=torch.int32) + query_start_loc = torch.tensor([0], dtype=torch.int32) + + mask_dcp_empty_shards_(lse, seq_lens, query_start_loc) + + assert torch.isneginf(lse).all() diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py new file mode 100644 index 000000000000..5b63c2c609f9 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for FlashInfer sparse MLA backend constraints.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + FlashInferMLASparseTRTLLMBackend, + _required_workspace_bytes, + compute_trtllm_sparse_mla_workspace_bytes, +) + +# Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8, +# max_num_batched_tokens=16384, FlashInfer 0.6.14. +_REPORTER_HEADS_PER_RANK = 64 // 8 +_REPORTER_DCP = 8 +_REPORTER_MAX_TOKENS = 16384 +# From the reported crash: trtllm-gen requested exactly this many bytes for +# trtllm_gen_softmax_workspace after an 8 MiB counter carve had left +# 404,750,336 of the 413,138,944-byte default buffer. +_REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312 +_REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608 + + +def test_reporter_softmax_carve_is_reproduced_byte_exactly(): + # The crashing step scheduled 12288 tokens (a single 12K-token request); + # the softmax slab formula must reproduce the kernel's request exactly: + # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard. + softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=0, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=12288, + ) + assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES + + +def test_reporter_config_covers_observed_overflow(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert computed >= ( + _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES + ) + + +def test_non_dcp_size_is_unchanged(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=1, + num_heads_per_rank=128, + max_num_batched_tokens=65536, + ) + assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE + + +def test_combined_pcp_dcp_rejects_flashinfer_sparse(): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=4, + decode_context_parallel_size=4, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is not None + assert "use FLASHMLA_SPARSE" in reason + + +@pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)]) +def test_single_context_parallel_mode_allows_flashinfer_sparse( + pcp_size: int, dcp_size: int +): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=dcp_size, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is None + + +def test_default_constant_matches_envs_default(monkeypatch): + from vllm import envs + + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + + +def test_env_unset_returns_computed(monkeypatch): + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + _REPORTER_DCP, + _REPORTER_HEADS_PER_RANK, + _REPORTER_MAX_TOKENS, + ) + + +def test_env_override_below_computed_is_respected_with_warning(monkeypatch): + override = 100 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_called_once() + + +def test_env_override_above_computed_is_respected_without_warning(monkeypatch): + override = 8 * 1024 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_not_called() diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index eab4759640a6..812bce39d889 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -307,6 +307,23 @@ def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape(): torch.testing.assert_close(actual, expected) +def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch): + seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32) + world = 4 + rank = 2 + rank_tensor = torch.tensor(rank, dtype=torch.int32) + expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2) + + def fail_tensor_construction(*args, **kwargs): + raise AssertionError("rank tensor must be reused") + + with monkeypatch.context() as patch: + patch.setattr(torch, "tensor", fail_tensor_construction) + actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2) + + torch.testing.assert_close(actual, expected) + + def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): world = 2 rank = 1 diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 87df1e409cf3..f3b4b7f5ba24 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -179,6 +179,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, + prefill_context_parallel_size=1, cp_kv_cache_interleave_size=1, ), ) @@ -193,6 +194,37 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): } +def test_indexer_warmup_includes_pcp_normalized_dcp_key(monkeypatch): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.indexer.get_dcp_group", + lambda: SimpleNamespace(rank_in_group=2), + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(compress_ratios=[32], index_kpool=32) + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + ) + + keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config) + + # Triton's compile key normalizes generic i32 values to 2 and divisible + # i32 values (including zero) to 16. + assert {(key.dcp_rank, key.dcp_world) for key in keys} == {(2, 2), (16, 1)} + assert { + ( + key.input_variant.is_aligned("uncompressed_seq_lens"), + key.input_variant.is_aligned("cu_compressed_seq_lens"), + ) + for key in keys + } == {(False, False), (False, True), (True, False), (True, True)} + + def test_compressed_slot_mapping_warmup_includes_index_kpool(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=256), diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index ae2e51284015..ba10a2e91cea 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -113,6 +113,16 @@ def supports_combination( from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 + ): + return ( + "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " + "running the rank-local PCP prefill queries" + ) if kv_cache_dtype == "fp8_ds_mla": return ( "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" @@ -265,6 +275,31 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS hisparse_supports_multi_token_decode: ClassVar[bool] = True + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + # Under DCP the workspace must hold the trtllm-gen softmax-stats slab. + # The buffer address is baked into CUDA graphs, so it has to reach its + # final size here, before the first forward/capture (no lazy regrow). + if self.dcp_world_size > 1: + num_q_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + _get_workspace_buffer( + device, + _required_workspace_bytes( + self.dcp_world_size, + num_q_heads, + vllm_config.scheduler_config.max_num_batched_tokens, + ), + ) + def _build_req_id_per_token( self, common_attn_metadata: "CommonAttentionMetadata", @@ -275,14 +310,98 @@ def _build_req_id_per_token( # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None +# trtllm-gen carves a softmax-stats slab from the workspace whenever LSE is +# requested (FlashInfer csrc/trtllm_fmha_kernel_launcher.cu, unchanged from +# v0.6.14 through current main): +# sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256) +# + 1 MiB guard +# forward_mqa always passes q_len == 1, so each (head, token) pair costs +# round_up(1, 256) == 256 slots. +_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8 # sizeof(float2) +_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256 # round_up(q_len=1, 256) +_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024 + +# Keep in sync with the VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE default in +# vllm/envs.py. +_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024 + + +def compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes: int, + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Workspace bytes needed by the trtllm-gen sparse MLA decode kernel. + + Under DCP the query is all-gathered across the DCP group in the head dim + and the kernel is asked for LSE, so trtllm-gen carves the softmax-stats + slab described above from the workspace before its (batch-independent) + counter and scratch regions. ``base_workspace_bytes`` must stay available + for those regions, so the slab is added on top of it (see #50781). + + Without DCP no LSE is requested, no slab is carved, and the base size is + returned unchanged. + """ + if dcp_world_size <= 1: + return base_workspace_bytes + softmax_bytes = ( + _TRTLLM_GEN_SOFTMAX_STAT_BYTES + * (num_heads_per_rank * dcp_world_size) + * max_num_batched_tokens + * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN + + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES + ) + return base_workspace_bytes + softmax_bytes + + +def _required_workspace_bytes( + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Resolve the workspace size, honoring an explicit env override.""" + computed = compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + ) + if not envs.is_set("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE"): + return computed + env_bytes = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + if env_bytes < computed: + logger.warning_once( + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the %d bytes " + "computed for the sparse MLA workspace with dcp_world_size=%d, " + "%d heads/rank and max_num_batched_tokens=%d. Respecting the " + "override, but the trtllm-gen kernel may crash with a workspace " + "overflow; set VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d or unset " + "it to use the computed size.", + env_bytes, + computed, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + computed, + ) + return env_bytes + -def _get_workspace_buffer(device: torch.device) -> torch.Tensor: +def _get_workspace_buffer( + device: torch.device, min_bytes: int | None = None +) -> torch.Tensor: global _fi_sparse_workspace - if _fi_sparse_workspace is None: + required = ( + min_bytes + if min_bytes is not None + else envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + ) + if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required: # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; # the trtllm-gen path views it as uint8, so int8 is safe for all backends. _fi_sparse_workspace = torch.zeros( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, + required, dtype=torch.int8, device=device, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2321cf8d13d4..28bd450da1ca 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -10,7 +10,7 @@ from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonPointerInputVariant, @@ -493,6 +493,14 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: dcp_world = parallel_config.decode_context_parallel_size dcp_interleave = parallel_config.cp_kv_cache_interleave_size dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 + dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)] + if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1: + # PCP+DCP prefill constructs a global PCP chunk plan and therefore + # dispatches this kernel with already-localized row starts. The + # runtime key is normalized to rank 0/world 1 in that path; warm it + # explicitly so one PCP rank cannot enter a collective while + # another rank is still compiling the first real inference batch. + dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1)) compress_ratios = tuple( dict.fromkeys( max(1, int(ratio)) @@ -506,17 +514,33 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios: compress_ratios = compress_ratios + (index_kpool,) return self._trace_dispatch(self.dispatch)( + zip_inputs(*dcp_cases), # Cover Triton's divisible, exact-one, and generic i32 classes. query_slice_start=(0, 1, 2), query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), - DCP_RANK=dcp_rank, - DCP_WORLD=dcp_world, DCP_INTERLEAVE=dcp_interleave, BLOCK_SIZE=self.BLOCK_SIZE, COMPRESS_RATIO=list(compress_ratios), + # PCP's global cumulative lengths are the second row of one packed + # allocation, so their pointer is not always 16-byte aligned. + # Prefill chunking can independently unalign uncompressed lengths. input_variant=( - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True), - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=False, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=False, + ), ), ) @@ -527,7 +551,9 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: uncompressed_seq_lens=compile_key.input_variant.pointer( "uncompressed_seq_lens", torch.int32 ), - cu_compressed_seq_lens=int32_ptr, + cu_compressed_seq_lens=compile_key.input_variant.pointer( + "cu_compressed_seq_lens", torch.int32 + ), row_start_cu_compressed_seq_lens=int32_ptr, token_to_seq=int32_ptr, cu_compressed_seq_len_ks=int32_ptr, @@ -836,6 +862,12 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: dtype=torch.int32, device=self.device, ) + # Materialize the rank on device during builder initialization. Creating + # this scalar in build() would introduce a GPU<->CPU sync in the decode + # hot path. + self.dcp_rank_tensor = torch.tensor( + self.dcp_rank, dtype=torch.int32, device=self.device + ) self.expanded_block_table_buffer = torch.zeros( (scheduler_config.max_num_batched_tokens, block_table_width), dtype=torch.int32, @@ -889,7 +921,7 @@ def _dcp_localize_decode_seq_lens( local_seq_lens = get_dcp_local_seq_lens( seq_lens, self.dcp_world_size, - self.dcp_rank, + self.dcp_rank_tensor, self.cp_kv_cache_interleave_size, ) if seq_lens_is_buffer_view: diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index a3adf70d2588..e656343a9189 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -1094,7 +1094,7 @@ def compute_causal_conv1d_metadata( def get_dcp_local_seq_lens( seq_lens: torch.Tensor, dcp_size: int = 1, - dcp_rank: int | None = None, + dcp_rank: int | torch.Tensor | None = None, cp_kv_cache_interleave_size: int = 1, ) -> torch.Tensor: """While using dcp, kv_cache size stored on each rank may be different, @@ -1112,6 +1112,12 @@ def get_dcp_local_seq_lens( dcp_size, ) seq_lens_tiled = seq_lens_i32.unsqueeze(-1) + elif isinstance(dcp_rank, torch.Tensor): + assert dcp_rank.dtype == torch.int32 + assert dcp_rank.device == seq_lens.device + assert dcp_rank.numel() == 1 + rank_offsets = dcp_rank + seq_lens_tiled = seq_lens_i32 else: rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device) seq_lens_tiled = seq_lens_i32 diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index bd20c9020d3e..58ac1e14a09d 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -77,6 +77,13 @@ def mask_dcp_empty_shards_( return assert seq_lens is not None and query_start_loc is not None + # A DCP rank can receive no local sequences during CUDA graph warmup even + # though the padded LSE buffer still has rows. In that case every row is an + # empty shard; avoid indexing the empty seq_lens tensor below. + if seq_lens.shape[0] == 0: + lse.fill_(float("-inf")) + return + row_indices = torch.arange( lse.shape[0], device=lse.device, dtype=query_start_loc.dtype ) From 35dc273072ff133743bc5fcfa306bf030fa8730b Mon Sep 17 00:00:00 2001 From: ppalanga Date: Tue, 15 Sep 2026 10:18:09 -0700 Subject: [PATCH 265/774] [ROCm][Spec Decode] Add Aiter MLA decode support non-causal draft block (#55966) Signed-off-by: Poovaiah Palangappa Co-authored-by: Poovaiah Palangappa Co-authored-by: John Qin Co-authored-by: Cursor Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- .../test_rocm_aiter_mla_fp8_support.py | 117 +++++++++++++- .../test_rocm_aiter_mla_fp8_decode_routing.py | 5 + .../test_rocm_aiter_mla_mtp_split.py | 98 ++++++++++++ vllm/_aiter_ops.py | 27 +++- .../attention/backends/mla/rocm_aiter_mla.py | 148 +++++++++++++++++- 5 files changed, 390 insertions(+), 5 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py index d67df8ee2b4d..d5c37724aacd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py @@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None: def _install_fake_aiter_modules( - monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool + monkeypatch: pytest.MonkeyPatch, + *, + supports_fp8: bool, + supports_causal: bool = False, ) -> None: aiter_mod: Any = types.ModuleType("aiter") mla_mod: Any = types.ModuleType("aiter.mla") @@ -65,6 +68,14 @@ def mla_decode_fwd_without_fp8( mla_decode_fwd = mla_decode_fwd_without_fp8 + if supports_causal: + inner = mla_decode_fwd + + def mla_decode_fwd_with_causal(*args, causal=True, **kwargs): + return inner(*args, **kwargs) + + mla_decode_fwd = mla_decode_fwd_with_causal + mla_mod.mla_decode_fwd = mla_decode_fwd aiter_mod.mla = mla_mod @@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch): assert _check_aiter_mla_fp8_support() is True assert _check_aiter_mla_fp8_support() is True assert signature_mock.call_count == 1 + + +def test_missing_causal_arg_fails_closed(monkeypatch): + """Old aiter wheels have no causal=; do not report that as a valid draft path.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False) + with pytest.raises(RuntimeError, match="causal-only"): + rocm_aiter_ops.mla_decode_supports_non_causal() + with pytest.raises(RuntimeError, match="causal-only"): + AiterMLABackend.supports_non_causal() + + +@pytest.mark.parametrize("has_kernels", [False, True]) +def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels): + """The Python argument alone does not guarantee an arch has the kernel.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, + "_aiter_mla_non_causal_asm_kernels", + lambda: has_kernels, + ) + assert rocm_aiter_ops.mla_decode_supports_non_causal() is True + assert AiterMLABackend.supports_non_causal() is has_kernels + + +@pytest.mark.parametrize("causal", [True, False]) +def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal): + """The decode op must pass the block mask through; omitting it is causal-only.""" + seen: list = [] + + def mla_decode_fwd(*args, causal=None, **kwargs): + seen.append(causal) + + mla_mod: Any = types.ModuleType("aiter.mla") + mla_mod.mla_decode_fwd = mla_decode_fwd + aiter_mod: Any = types.ModuleType("aiter") + aiter_mod.mla = mla_mod + monkeypatch.setitem(sys.modules, "aiter", aiter_mod) + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod) + + import torch + + import vllm._aiter_ops as aiter_ops + + aiter_ops._rocm_aiter_mla_decode_fwd_impl( + torch.zeros(2, 1, 8), + torch.zeros(2, 1, 1, 8), + torch.zeros(2, 1, 8), + torch.zeros(3, dtype=torch.int32), + 1, + causal=causal, + ) + assert seen == [causal] + + +def _aiter_mla_validate( + monkeypatch, + *, + dtype, + kv_cache_dtype, + use_non_causal, +): + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True + ) + return AiterMLABackend.validate_configuration( + head_size=576, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=1, + use_mla=True, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(9, 5), + attn_type=AttentionType.DECODER, + use_non_causal=use_non_causal, + ) + + +def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch): + """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton.""" + import torch + + reasons = _aiter_mla_validate( + monkeypatch, + dtype=torch.float16, + kv_cache_dtype="auto", + use_non_causal=True, + ) + assert any("non-causal fp16" in reason for reason in reasons) diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py index 06684302cf39..a86c7e617be1 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py +++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py @@ -176,3 +176,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads): assert unpadded.shape == q.shape torch.testing.assert_close(unpadded, q) + + +def test_a_non_causal_block_never_routes_to_gluon(gluon_available): + """A small-head block that normally uses Gluon must use ASM when non-causal.""" + assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 112cdff324df..efffcad927db 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -130,6 +130,7 @@ def _builder( _mla_q_dtype=torch.bfloat16, _mla_kv_dtype=torch.bfloat16, decode_attn_out_dtype=torch.bfloat16, + _decode_causal=True, ) # Bound to the stub rather than faked: the verify flatten's per-row view is # part of what _build_decode is being tested for. @@ -155,6 +156,25 @@ def test_backend_declares_uniform_batch_support(): ) +@pytest.mark.parametrize("causal", [True, False]) +def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal): + seen = [] + + def fake_common_build(self, *args): + seen.append(self._decode_causal) + return SimpleNamespace(decode=None, prefill=None) + + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build + ) + builder = object.__new__(AiterMLAMetadataBuilder) + builder._fp8_prefill_enabled = False + + builder.build(0, SimpleNamespace(causal=causal)) + + assert seen == [causal] + + def test_dcp_verify_row_view_is_causal_per_row(): """Row t must cover this rank's share of positions [0, seq_len - qlen + t]. @@ -949,3 +969,81 @@ def test_persistent_metadata_gate_without_gluon_build( assert metadata.has_persistent_metadata is expect_persistent assert get_mla_metadata_v1.called is expect_persistent + + +def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen): + """Drive _build_decode with causal=False and hand back the aiter mock.""" + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + builder = _builder( + mtp_decode_qlen=mtp_qlen, + kv_cache_dtype=kv_cache_dtype, + num_heads=num_heads, + ) + builder._decode_causal = False + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), + max_seq_len=7, + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + num_decode_tokens=2 * qlen, + dcp_tot_seq_lens_device=None, + ) + return metadata, get_mla_metadata_v1 + + +def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch): + """The schedule carries the mask, so is_causal must follow the block.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata + # 6th positional arg of get_mla_metadata_v1 is is_causal. + assert get_mla_metadata_v1.call_args.args[5] is False + + +@pytest.mark.parametrize("num_heads", [12, 16, 48]) +def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads): + """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel. + + 48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they + must be refused the same way as native 16-head qlen 2. + """ + with pytest.raises(ValueError, match="2-token"): + _build_non_causal( + monkeypatch, + num_heads=num_heads, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + + +def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch): + """32 heads at qlen 2 folds onto the 16-head / 4-token kernel.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, + num_heads=32, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.called + + +def test_a_two_token_bf16_block_is_allowed(monkeypatch): + """The bf16 fold has no such hole, so the guard must not fire there.""" + metadata, _ = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4cfb3b388486..9bea4f8b3a6e 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -583,10 +583,11 @@ def _rocm_aiter_mla_decode_fwd_impl( reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ) -> None: from aiter.mla import mla_decode_fwd - kwargs: dict[str, float | torch.Tensor | None] = { + kwargs: dict[str, float | torch.Tensor | None | bool] = { "sm_scale": sm_scale, "logit_cap": logit_cap, } @@ -596,6 +597,8 @@ def _rocm_aiter_mla_decode_fwd_impl( kwargs["q_scale"] = q_scale kwargs["kv_scale"] = kv_scale + kwargs["causal"] = causal + if work_meta_data is not None: assert work_indptr is not None, ( "work_indptr must be provided with work_meta_data" @@ -2036,6 +2039,26 @@ def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool: def is_mla_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._MLA_ENABLED + @classmethod + @if_aiter_supported + def mla_decode_supports_non_causal(cls) -> bool: + """Whether installed aiter.mla.mla_decode_fwd accepts `causal`. + + Added in aiter v0.1.20. A missing argument is a hard error rather + than a boolean the caller can use to pick another MLA backend. + """ + import inspect + + from aiter.mla import mla_decode_fwd + + if "causal" not in inspect.signature(mla_decode_fwd).parameters: + raise RuntimeError( + "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). " + "The installed aiter is causal-only; upgrade aiter rather than " + "falling back to another MLA backend." + ) + return True + @classmethod @if_aiter_supported def is_mha_enabled(cls) -> bool: @@ -2858,6 +2881,7 @@ def mla_decode_fwd( reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ): torch.ops.vllm.rocm_aiter_mla_decode_fwd( q, @@ -2878,6 +2902,7 @@ def mla_decode_fwd( reduce_indptr=reduce_indptr, reduce_final_map=reduce_final_map, reduce_partial_map=reduce_partial_map, + causal=causal, ) @staticmethod diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index b850deb07091..c9ba2b0a566e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -4,7 +4,7 @@ import functools from dataclasses import dataclass from pathlib import Path -from typing import ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final import torch @@ -38,6 +38,9 @@ ) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +if TYPE_CHECKING: + from vllm.platforms.interface import DeviceCapability + logger = init_logger(__name__) @@ -144,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool: ) +def _aiter_mla_non_causal_asm_kernels() -> bool: + """Whether this arch ships non-causal MLA decode ASM kernels. + + The Python `causal=` probe cannot see the per-arch manifest. gfx950 has + the kernels; gfx942 and everything else must fall through to another + backend. + """ + try: + from vllm.platforms.rocm import on_gfx950 + except Exception: # noqa: BLE001 + return False + return bool(on_gfx950()) + + @functools.lru_cache(maxsize=1) def _gluon_mla_decode_supported() -> bool: """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. @@ -252,6 +269,72 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: def get_name() -> str: return "ROCM_AITER_MLA" + @classmethod + def supports_non_causal(cls) -> bool: + # causal= on mla_decode_fwd is necessary but not sufficient: pinned + # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM + # manifest has no non-causal decode kernels. Selecting this backend + # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA. + return bool( + rocm_aiter_ops.mla_decode_supports_non_causal() + and _aiter_mla_non_causal_asm_kernels() + ) + + @classmethod + def validate_configuration( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + use_per_head_quant_scales: bool, + device_capability: "DeviceCapability", + attn_type: str, + has_sliding_window: bool = False, + use_non_causal: bool = False, + use_batch_invariant: bool = False, + use_kv_connector: bool = False, + use_pcp: bool = False, + use_adaptive_verification: bool = False, + use_dcp: bool = False, + use_rswa: bool = False, + ) -> list[str]: + invalid_reasons = super().validate_configuration( + head_size=head_size, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=block_size, + use_mla=use_mla, + has_sink=has_sink, + use_sparse=use_sparse, + use_mm_prefix=use_mm_prefix, + use_per_head_quant_scales=use_per_head_quant_scales, + device_capability=device_capability, + attn_type=attn_type, + has_sliding_window=has_sliding_window, + use_non_causal=use_non_causal, + use_batch_invariant=use_batch_invariant, + use_kv_connector=use_kv_connector, + use_pcp=use_pcp, + use_adaptive_verification=use_adaptive_verification, + use_dcp=use_dcp, + use_rswa=use_rswa, + ) + # Arch/signature cannot express this: the backend still advertises + # fp16, so a DSpark draft with kv_cache_dtype="auto" would select + # AITER and abort on fp16 query. Causal fp16 is fine. + if ( + use_non_causal + and dtype == torch.float16 + and kv_cache_dtype in (None, "auto", "float16") + ): + invalid_reasons.append("non-causal fp16 MLA decode not supported") + return invalid_reasons + @staticmethod def get_impl_cls() -> type["AiterMLAImpl"]: return AiterMLAImpl @@ -347,6 +430,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # https://github.com/vllm-project/vllm/issues/22945 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + # Served by passing the mask to the kernel; _build_decode turns away the + # shapes AITER has no non-causal kernel for. + supports_non_causal_multi_token_decode: ClassVar[bool] = True + # Set from the common metadata every build; a batch is causal unless the + # drafter says otherwise. + _decode_causal: bool = True @staticmethod def _uniform_padded_mtp_qo_len( @@ -919,6 +1008,7 @@ def _build_decode( num_decode_tokens: int, dcp_tot_seq_lens_device: torch.Tensor | None, ) -> AiterMLADecodeMetadata: + causal = self._decode_causal device = self.device num_reqs = seq_lens_device.size(0) qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] @@ -970,6 +1060,7 @@ def _build_decode( int(max_qo_len), self._kv_cache_dtype_str, self.dcp_world_size, + causal, ) use_segmented_dcp_verify = ( self._supports_segmented_dcp_verify and max_qo_len > 1 @@ -1047,14 +1138,36 @@ def _build_decode( # A padded rank has no bf16 persistent kernel past qlen 4 where the # gfx950 fold is absent; the non-persistent entry covers it. fp8 # keeps the schedule -- its fold rejects non-persistent outright. + # A non-causal block keeps it too: what the fold drops past qlen 4 + # is the block's causal staircase, which a non-causal block does + # not have, and the schedule is the only thing carrying its mask. and ( - self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + not causal + or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN or is_quantized_kv_cache(self._kv_cache_dtype_str) ) and max_qo_len >= 1 and max_qo_len <= self._mtp_decode_qlen ) + if ( + not causal + and max_qo_len == 2 + and is_quantized_kv_cache(self._kv_cache_dtype_str) + and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel( + self._decode_num_heads + ) + ): + # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the + # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads + # at qlen 2 fold onto that 16-head / 4-token non-causal kernel; + # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the + # process rather than raising. The bf16 fold has no such hole. + raise ValueError( + "AITER has no non-causal fp8 MLA kernel for this 2-token " + "query block. Pin the draft to TRITON_MLA for this " + "speculative config." + ) if use_persistent_metadata: from aiter import get_mla_metadata_v1 @@ -1067,7 +1180,7 @@ def _build_decode( paged_kv_last_page_len, self._num_attention_heads, 1, - True, + causal, self._mla_work_meta_data, self._mla_work_info_set, self._mla_work_indptr, @@ -1137,6 +1250,11 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> AiterMLAMetadata: + # The common MLA builder owns the causal flag but its _build_decode + # contract is shared with CUDA backends. Keep this ROCm-specific state + # on the per-ubatch Aiter builder instead of changing that shared API; + # super().build invokes _build_decode synchronously. + self._decode_causal = common_attn_metadata.causal attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) @@ -1226,6 +1344,10 @@ class AiterMLAHelper: # fold that reaches a persistent one is gfx950-only. _ASM_PADDED_MAX_PS_QLEN: Final = 4 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = () + # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the + # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112) + # fold to H16 while keeping Q2, which has no non-causal fp8 entry. + _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128) @staticmethod def qo_indptr_for_uniform_qlen( @@ -1263,6 +1385,18 @@ def is_valid_num_heads(num_heads: int) -> bool: ) ) + @staticmethod + def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool: + """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel. + + Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the + non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2 + block has no matching entry, and padded 48/80/112 keep Q2 after the + H16 fold. + """ + kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads) + return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS + @staticmethod def get_actual_mla_num_heads(num_heads: int) -> int: if num_heads == 24 and _aiter_mla_native_h24_supported(): @@ -1364,6 +1498,7 @@ def use_gluon_verify( max_qo_len: int, kv_cache_dtype: str, dcp_world_size: int = 1, + causal: bool = True, ) -> bool: """Whether a small-head multi-token verify uses native Gluon MTP. @@ -1376,7 +1511,13 @@ def use_gluon_verify( DCP verify is excluded: its per-row causal windows are served by the segmented path, which Gluon's MTP entry cannot express. + + Gluon masks the block causally with no way to turn it off, so a + non-causal block takes the padded asm decode whatever the head count + -- padding is what gives it a kernel there. """ + if not causal: + return False if max_qo_len <= 1 or dcp_world_size > 1: return False if is_quantized_kv_cache(kv_cache_dtype): @@ -1974,6 +2115,7 @@ def forward_mqa( decode.paged_kv_indptr, decode.paged_kv_indices, decode.paged_kv_last_page_len, + causal=attn_metadata.causal, **mla_kwargs, ) From 2f46a1d9c6ab33c4ea8196c375ff85b661982bff Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 11:19:14 -0700 Subject: [PATCH 266/774] [CI] Keep Qwen3 Omni DSpark config fixture complete (#57044) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/model_executor/test_qwen3_omni.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 42accbef85c5..30d9909a5a2e 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -10,6 +10,7 @@ import torch.nn as nn from transformers import PretrainedConfig +from vllm.config import ParallelConfig from vllm.multimodal.processing import InputProcessingContext @@ -371,7 +372,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): attention_backend=None, kv_cache_dtype=None, ), - parallel_config=SimpleNamespace(tensor_parallel_size=1), + parallel_config=ParallelConfig(), attention_config=SimpleNamespace(backend=None), cache_config=SimpleNamespace(), load_config=SimpleNamespace(), From dffbb714e4e8e4b95ccc888df98d47c7d2cef78d Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:34:24 -0400 Subject: [PATCH 267/774] [ROCm][Perf] Optimize DSV4.1 K=512 decode top-k on gfx950 (#56743) Signed-off-by: fai --- csrc/libtorch_stable/sampler.cu | 263 ++++++++++++++---- tests/kernels/test_top_k_per_row.py | 189 +++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 12 + 3 files changed, 406 insertions(+), 58 deletions(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index ada1c4ba39c0..b8d64f7af130 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads, } template __device__ bool processHistogramStep( const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, @@ -180,10 +181,15 @@ __device__ bool processHistogramStep( << patternShift; } + int negativeInfinityCount = 0; auto distributeToBins = [&](float logit, int /* idx */ = 0) { if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - atomicAdd(&smemFinal.histo.data[binIdx], 1); + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { + ++negativeInfinityCount; + } else { + atomicAdd(&smemFinal.histo.data[binIdx], 1); + } } }; @@ -198,6 +204,12 @@ __device__ bool processHistogramStep( distributeToBins(logit, idx); } } + if constexpr (useTopK512Optimization) { + if (negativeInfinityCount > 0) { + const auto bin = extractBinIdx(-INFINITY); + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); + } + } // Make sure the histogram is ready. __syncthreads(); @@ -254,6 +266,12 @@ __device__ bool processHistogramStep( // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. + const bool finalBinFits = + smemFinalBinSize[0] <= kNumFinalItems && + !(useTopK512Optimization && step < 3 && + isPartialMatch(-INFINITY, logitPattern) && + thresholdBinIdx == extractBinIdx(-INFINITY)); auto processBins = [&](float logit, int idx) { if (isPartialMatch(logit, logitPattern)) { @@ -262,8 +280,7 @@ __device__ bool processHistogramStep( // 1. This is step 0 and the threshold bin is small enough (no step 1) // 2. This is step >= 1 (where pattern matching filters correctly) // This prevents duplicates when step 0 and step 1 both run. - bool shouldWriteDirectly = - (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1); + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); if (binIdx < thresholdBinIdx && shouldWriteDirectly) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -279,8 +296,7 @@ __device__ bool processHistogramStep( } if constexpr (step < 3) { // Only fill the final items for sorting if the threshold bin fits - if (binIdx == thresholdBinIdx && - smemFinalBinSize[0] <= kNumFinalItems) { + if (binIdx == thresholdBinIdx && finalBinFits) { int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); smemFinal.items.logits[dstIdx] = logit; if constexpr (mergeBlocks) { @@ -325,20 +341,21 @@ __device__ bool processHistogramStep( __syncthreads(); // Check if we should continue to next step - return smemFinalBinSize[0] > kNumFinalItems; + return !finalBinFits; } // Follows half - 11 - 11 - 10 bit iterations // Keep the adaptive kernel's device instantiation separate from legacy calls. template + bool deviceLengthAware = false, bool useTopK512Optimization = false> static __device__ void topKPerRowJob(const int* indices, const float* logits, int rowStart, int rowEnd, int* outIndices, float* outLogits, int stride1, int topK) { static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); // The number of slots for the final pass. - static constexpr int kNumFinalItems = 2048; + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048; // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; @@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, rowIt += kNumThreadsPerBlock) { outIndices[rowIt] = -1; if constexpr (multipleBlocksPerRow) { - outLogits[rowIt] = -FLT_MAX; + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; } } @@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 0: Process first 11 bits of half representation bool continueToNextStep = processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 1: Process next 11 bits continueToNextStep = processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 2: Process next 11 bits continueToNextStep = processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, if (continueToNextStep) { // Step 3: Process last 10 bits processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // The histogram did not proceed to the final 10 bits, therefore we need to // sort the final items The logits of the elements to be sorted in the final // pass. - if constexpr (useRadixSort) { - // Sorting with radix sort - float finalLogits[kNumFinalItemsPerThread]; - // The indices of the elements to be sorted in the final pass. - int finalIndices[kNumFinalItemsPerThread]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - finalLogits[ii] = -FLT_MAX; - } - - // Read the elements from SMEM. -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) { - finalLogits[ii] = smemFinal.items.logits[srcIdx]; - finalIndices[ii] = smemFinal.items.indices[srcIdx]; - } - } - // Make sure the shared memory has been read. - __syncthreads(); - - // Sort the elements. - FinalSort(smemFinal.finalSort) - .SortDescendingBlockedToStriped(finalLogits, finalIndices); - - // Copy the data back to the shared memory storage. - int baseIdx = smemFoundTopKValues[0]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - int dstIdx = baseIdx + srcIdx; - - if (dstIdx < topK) { - smemOutput[dstIdx] = finalIndices[ii]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[dstIdx] = - finalLogits[ii]; - } - } - } - } else { + auto insertionSort = [&]() { // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; for (int i = threadIdx.x; i < smemFinalDstIdx[0]; @@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, } } } + }; + if constexpr (useRadixSort) { + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { + insertionSort(); + } else { + // Sorting with radix sort + float finalLogits[kNumFinalItemsPerThread]; + // The indices of the elements to be sorted in the final pass. + int finalIndices[kNumFinalItemsPerThread]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + if constexpr (useTopK512Optimization) finalIndices[ii] = -1; + } + + // Read the elements from SMEM. +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + if (srcIdx < smemFinalDstIdx[0]) { + finalLogits[ii] = smemFinal.items.logits[srcIdx]; + finalIndices[ii] = smemFinal.items.indices[srcIdx]; + } + } + // Make sure the shared memory has been read. + __syncthreads(); + + // Sort the elements. + FinalSort(smemFinal.finalSort) + .SortDescendingBlockedToStriped(finalLogits, finalIndices); + + // Copy the data back to the shared memory storage. + int baseIdx = smemFoundTopKValues[0]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + int dstIdx = baseIdx + srcIdx; + + if (dstIdx < topK) { + smemOutput[dstIdx] = finalIndices[ii]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[dstIdx] = + finalLogits[ii]; + } + } + } + } + } else { + insertionSort(); } __syncthreads(); } @@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } +#ifdef USE_ROCM +// The grid and workspace stay fixed while graph replays change device lengths. +__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, + int maxBlocks) { + int activeBlocks = 1; + if (numRows <= 32 && rowLength > 16 * 1024) { + activeBlocks = rowLength <= 128 * 1024 ? 8 + : rowLength <= 256 * 1024 ? 10 + : 16; + } else if (numRows <= 128 && rowLength > 64 * 1024) { + activeBlocks = rowLength <= 512 * 1024 ? 4 : 10; + } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) { + activeBlocks = numRows <= 160 ? 3 : 2; + } + // Every block in a merged row must contain at least k real indices. + return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512)); +} + +template +static __global__ +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware( + const float* logits, const int* seqLens, int* outIndices, + int* outIndicesAux, float* outLogitsAux, int stride0, int next_n, + int seqLensIs2D, int maxBlocks) { + #if defined(__gfx950__) + constexpr int topK = 512; + constexpr int kNumBins = 2048; + const int rowIdx = blockIdx.x; + const int rowEnd = + seqLensIs2D + ? max(0, seqLens[rowIdx]) + : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1); + const int activeBlocks = + gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks); + outIndices += static_cast(rowIdx) * topK; + if constexpr (mergeBlocks) { + if (activeBlocks == 1) return; + const int64_t offset = static_cast(rowIdx) * maxBlocks * topK; + topKPerRowJob(outIndicesAux + offset, outLogitsAux + offset, 0, + activeBlocks * topK, outIndices, nullptr, 1, topK); + } else { + if (blockIdx.y >= activeBlocks) return; + logits += static_cast(rowIdx) * stride0; + if (activeBlocks == 1) { + topKPerRowJob(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1, + topK); + return; + } + const int blockSize = rowEnd / activeBlocks; + const int rowStart = blockSize * blockIdx.y; + const int blockRowEnd = + blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize; + const int64_t offset = + (static_cast(rowIdx) * maxBlocks + blockIdx.y) * topK; + topKPerRowJob(nullptr, logits, rowStart, blockRowEnd, + outIndicesAux + offset, outLogitsAux + offset, 1, topK); + } + #endif +} +#endif + } // namespace vllm void apply_repetition_penalties_( @@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, // the same seq_len and the kernel computes the per-row offset itself. int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; +#ifdef USE_ROCM + const bool useGfx950DeviceLengthAwareTopK512 = + topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 && + numColumns <= 1024 * 1024 && + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0; + if (useGfx950DeviceLengthAwareTopK512) { + int multipleBlocksPerRowConfig = 1; + if (numRows <= 32 && numColumns > 16 * 1024) { + multipleBlocksPerRowConfig = 16; + } else if (numRows <= 128 && numColumns > 64 * 1024) { + multipleBlocksPerRowConfig = 10; + } else if (numRows <= 160 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 3; + } else if (numRows <= 256 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 2; + } + + constexpr int kNumThreadsPerSplitBlock = 1024; + if (multipleBlocksPerRowConfig == 1) { + vllm::topKPerRowDecode512DeviceLengthAware + <<>>(logits.const_data_ptr(), + seqLens.const_data_ptr(), + indices.mutable_data_ptr(), nullptr, nullptr, + static_cast(stride0), static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } + const auto outIndicesAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Int, std::nullopt, logits.device()); + const auto outLogitsAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Float, std::nullopt, logits.device()); + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), static_cast(stride0), + static_cast(next_n), seqLensIs2D, multipleBlocksPerRowConfig); + constexpr int kNumThreadsPerBlockMerge = 1024; + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + nullptr, seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), 0, static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } +#endif + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index c629ae15dea9..591362020ee7 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -22,11 +22,22 @@ def _has_device_capability(major: int) -> bool: return current_platform.is_cuda() and current_platform.has_device_capability(major) +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + # DeepSelect is compiled for sm_100a/sm_103a only. requires_sm100 = pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="DeepSelect requires SM100a/SM103a", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), reason="This test exercises the gfx950 launch configuration" +) COOPERATIVE_TOPK_BACKEND = pytest.param( @@ -487,6 +498,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None: ) +def _assert_exact_topk( + logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor +) -> None: + ends = row_ends.reshape(-1, 1).clamp_min(0) + valid = (indices >= 0) & (indices < ends) + assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten()) + assert torch.all(valid | (indices == -1)) + ordered = indices.sort(1).values + assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any() + + width = int(ends.max()) + visible = logits[:, :width].clone() + columns = torch.arange(width, device=logits.device) + visible.masked_fill_(columns[None, :] >= ends, float("-inf")) + expected = visible.topk(min(indices.shape[1], width), dim=1).values + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, float("-inf")) + selected = selected.sort(1, descending=True).values + assert torch.equal(selected[:, : expected.shape[1]], expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None: + """Honor causal offsets and exact compressed lengths, including short rows.""" + next_n, width, top_k = 6, 16_384, 512 + lengths = torch.tensor( + [0, 1, 511, 512, 513, 1023, 1024, 16_000], + dtype=torch.int32, + device="cuda", + ) + offsets = torch.arange(next_n, dtype=torch.int32, device="cuda") + row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten() + rows = row_ends.numel() + storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda") + logits = storage[:, 1 : width + 1] + logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda") + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + _assert_exact_topk(logits, indices, row_ends) + + +@pytest.mark.parametrize( + ("rows", "row_length"), + [ + pytest.param(128, 65_536, id="single-before-64k"), + pytest.param(128, 65_537, id="four-splits-after-64k"), + pytest.param(129, 262_144, id="single-before-three-split-length"), + pytest.param(129, 262_145, id="three-splits-after-256k"), + pytest.param(160, 262_145, id="three-splits-upper-row-bound"), + pytest.param(161, 262_144, id="single-before-two-split-length"), + pytest.param(161, 262_145, id="two-splits-after-256k"), + pytest.param(256, 262_145, id="two-splits-upper-row-bound"), + pytest.param(257, 262_145, id="large-row-single-block"), + pytest.param(384, 262_145, id="optimized-row-limit"), + pytest.param(385, 262_145, id="native-fallback"), + ], +) +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries( + rows: int, row_length: int +) -> None: + width, top_k = 262_145, 512 + row = torch.arange(width, dtype=torch.float32, device="cuda") + logits = row.expand(rows, -1) + lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda") + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + + _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width) + + expected = torch.arange( + row_length - top_k, row_length, dtype=torch.int32, device="cuda" + ).expand_as(indices) + assert torch.equal(indices.sort(1).values, expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: + """Preserve visible top-k scores when masks and lengths change on replay.""" + rows, next_n, width, top_k = 24, 6, 262_145, 512 + logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda") + lengths = torch.full( + (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda" + ) + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + logits[:, :70_000] = float("-inf") + bin_sizes = [128, 129, 1024, 1025] + for row in range(rows): + pattern = row % 9 + if pattern < 4: + count = [0, 1, 511, 513][pattern] + logits[row, :count] = torch.arange( + count, dtype=torch.float32, device="cuda" + ) + elif pattern == 4: + logits[row, :511] = 0 + logits[row, 511] = torch.finfo(torch.float32).min + else: + count = bin_sizes[pattern - 5] + higher = max(0, top_k - count // 2) + logits[row, :higher] = 2 + logits[row, higher : higher + count] = 1 + logits[row, higher] = 1 + 2**-23 + + def run() -> None: + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + + run() + _assert_exact_topk(logits, indices, lengths) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + bounds = [ + 0, + 511, + 512, + 513, + 16_384, + 16_385, + 131_072, + 131_073, + 262_144, + 262_145, + ] + row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda") + repeats = (rows + len(bounds) - 1) // len(bounds) + lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows]) + logits[:] = torch.arange(width, dtype=torch.float32, device="cuda") + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: @@ -620,6 +775,31 @@ def decode_kernel(*args, **kwargs) -> None: compress_ratio=4, num_rows=257, max_valid_seq_len=250_000, + num_columns=524_288, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=385, + max_valid_seq_len=250_000, + num_columns=524_288, + topk_tokens=512, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=50_000, + num_columns=1_048_577, + topk_tokens=512, on_gfx950=True, ), decode_kernel, @@ -638,6 +818,15 @@ def fail_if_imported() -> None: ) ineligible_cases = [ dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=1_048_576, + num_columns=1_048_576, + topk_tokens=512, + on_gfx950=True, + ), dict( is_prefill=False, compress_ratio=4, diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 44b6017a057b..0e342cb0e064 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -58,6 +58,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 _GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -69,6 +70,8 @@ def _get_aiter_top_k_kernel( compress_ratio: int, num_rows: int, max_valid_seq_len: int | None = None, + num_columns: int | None = None, + topk_tokens: int = 1024, on_gfx950: bool = _ON_GFX950, ) -> Callable[..., None] | None: if compress_ratio <= 1 or not on_gfx950: @@ -76,6 +79,13 @@ def _get_aiter_top_k_kernel( if not is_prefill: assert max_valid_seq_len is not None + if ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns is not None + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ): + return None # AITER v0.1.19 decode is one-block only. This measured gfx950 # FP32/k=1024 compressed-row boundary is independent of the native # split-count boundary in sampler.cu. @@ -1261,6 +1271,8 @@ def rocm_aiter_sparse_attn_indexer( compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + num_columns=logits.shape[1], + topk_tokens=topk_tokens, ) if aiter_topk_kernel is not None: _launch_aiter_top_k_per_row_decode( From 2fcc5247882849b675088b3146b27c006e87d87f Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 14:41:54 -0500 Subject: [PATCH 268/774] [Bugfix] Avoid nested score cancellation handlers for /v1/score alias (#57024) Signed-off-by: Taneem Ibrahim --- vllm/entrypoints/pooling/scoring/api_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index 83ed600fdd21..8e99e9a1feca 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -57,8 +57,6 @@ async def create_score(request: ScoreRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation -@load_aware_call async def create_score_v1(request: ScoreRequest, raw_request: Request): logger.warning( "To indicate that Score API is not part of standard OpenAI API, we " From df42d112ee88dd4a9b64efbad55621af6a66a44b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 15 Sep 2026 21:44:26 +0200 Subject: [PATCH 269/774] [Config] Infer HiSparse attention config from HiSparseConnector (#57041) Signed-off-by: NickLucche --- tests/test_config.py | 102 +++++++++++++++++++++++++++++++++++++++ vllm/config/attention.py | 4 +- vllm/config/vllm.py | 9 +++- 3 files changed, 113 insertions(+), 2 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 4d39b56d265f..f75592f089e2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -337,6 +337,108 @@ def test_hisparse_rejects_non_cuda(monkeypatch): VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) +@pytest.mark.parametrize( + "kv_transfer_config", + [ + KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, + }, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_consumer", + kv_connector_extra_config={ + "connectors": [ + {"kv_connector": "NixlConnector", "kv_role": "kv_consumer"}, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + ], + ids=["standalone", "multi-connector", "pd-decode"], +) +def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_config): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + kv_transfer_config=kv_transfer_config, + # Skip the HMA auto-detect block: it imports the connector class, + # which pulls the platform's compiled attention extensions. + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert isinstance(config.attention_config.hisparse_config, HiSparseConfig) + + +def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig(device_buffer_size=512) + ), + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert config.attention_config.hisparse_config.device_buffer_size == 512 + + +def test_hisparse_connector_without_cuda_still_rejected(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + with pytest.raises(ValueError, match="requires NVIDIA CUDA"): + VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ) + ) + + +def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + config = VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="OffloadingConnector", kv_role="kv_both" + ) + ) + assert config.attention_config.hisparse_config is None + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 05da048bdbd6..129fbcbb0869 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -88,7 +88,9 @@ class AttentionConfig: hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental - Model Runner V2-only HiSparse sparse-MLA decode hot-buffering.""" + Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred + with defaults when HiSparseConnector is configured (directly or via + MultiConnector); set it explicitly only to tune its fields.""" use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ad8f077148c3..8e0846982b74 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -27,7 +27,7 @@ from vllm.utils import random_uuid from vllm.utils.hashing import safe_hash -from .attention import AttentionConfig +from .attention import AttentionConfig, HiSparseConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig @@ -1630,6 +1630,13 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if ( + self.attention_config.hisparse_config is None + and self.kv_transfer_config is not None + and self.kv_transfer_config.has_connector("HiSparseConnector") + ): + self.attention_config.hisparse_config = HiSparseConfig() + if self.attention_config.hisparse_config is not None: if not current_platform.is_cuda(): raise ValueError("HiSparse currently requires NVIDIA CUDA.") From 18f8aa04656de0a37eb5c7e9efc8ebb595ac73cf Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:17:53 -0400 Subject: [PATCH 270/774] [Refactor] Remove dead kernel code (#56845) Signed-off-by: yewentao256 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/mamba/ops/replayssm_config.py | 15 - .../models/glm5next/amd/ops/kpool_compress.py | 44 - vllm/models/inkling/amd/ops/lamport.py | 766 ------------------ vllm/models/inkling/amd/ops/mm_towers.py | 190 ----- .../cute_dsl/latent_moe_tail/primitives.py | 31 - .../v1/attention/ops/triton_fp8_mqa_logits.py | 262 ------ .../triton_turboquant_decode_v2.py | 572 ------------- .../triton_turboquant_unified_attention.py | 15 - 8 files changed, 1895 deletions(-) delete mode 100644 vllm/models/inkling/amd/ops/lamport.py delete mode 100644 vllm/models/inkling/amd/ops/mm_towers.py delete mode 100644 vllm/v1/attention/ops/triton_fp8_mqa_logits.py diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index aff4fc1c3f2d..3cee44ba5015 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -8,7 +8,6 @@ """ import functools -from contextlib import contextmanager from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -26,20 +25,6 @@ def _is_blackwell() -> bool: _overrides: dict[str, tuple] = {} -@contextmanager -def override_replayssm_config(kernel: str, config: tuple): - """Pin ``kernel``'s launch config for the duration of the context.""" - prev = _overrides.get(kernel) - _overrides[kernel] = config - try: - yield - finally: - if prev is None: - _overrides.pop(kernel, None) - else: - _overrides[kernel] = prev - - def _dstate_tile(dstate: int, tile: int) -> int: return max(16, min(tile, triton.next_power_of_2(dstate))) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 64e6fd46d73f..09e7a40aa5e0 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -65,50 +65,6 @@ def _hadamard128(x): return x * 0.08838834764831845 # 1/sqrt(128) -# Map pool ids to physical cache slots. - - -def compute_pooled_write_locs( - page_table_64: torch.Tensor, - pool_ids: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Map logical pooled-K ids to physical flat cache slots. - - ``pool_size`` consecutive tokens share one pool slot that lives at the - *first* token page of each page-group. ``page_table_64`` maps token pages - to physical block ids; we gather the block id of each pool's page-group - and add the in-block pool offset. - """ - assert page_table_64.ndim == 1 - pool_ids = pool_ids.to(torch.int64) - block_size = 64 - pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor") - token_page_row = pool_page_group * pool_size - packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64)) - return packed_page.to(torch.int64) * block_size + torch.remainder( - pool_ids, block_size - ) - - -def build_pooled_page_table( - page_table: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Build a pool-granular page table by taking every ``pool_size``-th - token-page column (one pool maps to ``pool_size`` token pages). - - Uses gather (not strided slicing) so the result is always a fresh - row-major tensor — some downstream kernels require stride(-1) == 1. - """ - block_size = page_table.shape[-1] - assert block_size % pool_size == 0, ( - f"pool_size ({pool_size}) must divide page columns ({block_size})" - ) - idx = torch.arange(0, block_size, pool_size, device=page_table.device) - return page_table[..., idx].contiguous() - - # Fused pool compression and cache write. diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py deleted file mode 100644 index 9705a9a3df37..000000000000 --- a/vllm/models/inkling/amd/ops/lamport.py +++ /dev/null @@ -1,766 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. - -The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. - -Liveness --------- -Large grids are deliberately split at the two communication dependencies: - - 1. ``_publish_input_kernel`` only publishes rank partials. - 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. - 3. ``_sconv_publish_kernel`` only computes and publishes the local result. - 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. - -Without PDL, CUDA stream order completes a producer before its consumer. With -PDL, every producer CTA posts all peer stores before triggering its dependent, -and the consumer executes ``gdc_wait`` before polling. A consumer therefore -waits only for stores from a producer that is already running (or complete) on -another GPU. It never waits for another block in its own grid. Consequently -spinning consumers cannot occupy resources needed by any producer, and the -wait-for graph has no cycle. The proof is independent of grid size, block -dispatch order, and occupancy. - -For one token, the first three phases use eight independent channel slices to -expose enough CTA parallelism for decode latency. Each slice has exclusive -ownership of its cache and Lamport columns; the gather/RMSNorm phase retains -one CTA per token so no cross-CTA reduction or completion counter is needed. - -Immediate buffer reuse (including replay of a captured CUDA graph) is also -safe. Rank R cannot republish input for call n+1 until its gather for n has -finished; that gather waited for owner O's output, which O publishes only -after consuming R's input for n. Likewise, R cannot republish output for n+1 -until its reduction for n+1 has observed destination D's input; D publishes -that input only after its gather consumed R's output for n. Thus every prior -read happens-before a same-slot rewrite. Three generations reduce incidental -coupling for ordinary launches, but correctness does not depend on rotation. - -The payload itself is the Lamport flag. A 32-bit store publishes two bf16s -atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real -negative zeroes are changed to positive zero before publication. Consumers -use volatile 32-bit loads and restore the sentinel after consuming a slot. -""" - -from __future__ import annotations - -import os - -import torch - -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import in_the_same_node_as -from vllm.logger import init_logger -from vllm.triton_utils import tl, triton - -logger = init_logger(__name__) - -_MAX_TOKENS = 16384 -_EMPTY_PAIR = tl.constexpr(0x80008000) - - -@triton.jit -def _pack_bf16_pairs(values): - """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" - lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) - lo = lo.to(tl.uint16, bitcast=True) - hi = hi.to(tl.uint16, bitcast=True) - lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) - hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) - return lo | (hi << 16) - - -@triton.jit -def _unpack_bf16_pairs(values): - lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - return tl.interleave(lo, hi) - - -@triton.jit -def _wait_pairs(ptr, offsets, mask): - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - return values - - -@triton.jit -def _publish_input_kernel( - stage_ptr, - peer_ptrs, - peer_offset_u32, - stride_stage_t, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Publish this rank's full partial row into every shard owner's slots.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - elem = tl.arange(0, CS_P2) - elem_mask = elem < CSS - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - # Address generation is independent of the preceding stream kernel. The - # acquire remains before stage/peer loads, which may consume its writes. - if USE_PDL: - tl.extra.cuda.gdc_wait() - - for owner in tl.static_range(WORLD): - values = tl.load( - stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, - mask=elem_mask, - other=0.0, - ) - packed = _pack_bf16_pairs(values) - base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) - dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair - tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _reduce_insert_kernel( - input_peer_ptrs, - input_peer_offset_u32, - cache_ptr, - slot_ptr, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume all partials for this rank and insert the reduced cache row.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - source = tl.arange(0, WORLD) - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - offsets = ( - (token * WORLD + source)[:, None] * (CS // 2) - + split * (CSS // 2) - + pair[None, :] - ) - mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] - input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - input_u32 += input_peer_offset_u32 - # Slot metadata and the cache destination do not depend on input publish. - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - safe_slot = tl.maximum(slot, 0).to(tl.int64) - dst = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(input_u32, offsets, mask) - - lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - reduced = tl.interleave( - tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), - tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), - ) - tl.store( - input_u32 + offsets, - tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=mask, - ) - - tl.store(dst, reduced, mask=valid & channel_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _sconv_publish_kernel( - peer_ptrs, - peer_offset_u32, - residual_ptr, - weight_ptr, - cache_ptr, - position_ptr, - sequence_ptr, - slot_ptr, - block_table_ptr, - stride_residual_t, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - stride_block_table_r, - max_blocks, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - WINDOW: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Compute this rank's shard, then publish it to every rank.""" - tl.static_assert(WINDOW == 4) - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - position = tl.load(position_ptr + token) - sequence = tl.load(sequence_ptr + token) - safe_slot = tl.maximum(slot, 0).to(tl.int64) - own_ptr = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - # These loads were made visible before reduce/insert could signal us and - # are independent of its cache store. Hoisting them gives PDL useful work - # to overlap while retaining the acquire before every cache read. - residual = tl.load( - residual_ptr + token * stride_residual_t + RANK * CS + global_channel, - mask=channel_mask, - other=0.0, - ) - weight0 = tl.load( - weight_ptr + global_channel * WINDOW, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight1 = tl.load( - weight_ptr + global_channel * WINDOW + 1, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight2 = tl.load( - weight_ptr + global_channel * WINDOW + 2, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight3 = tl.load( - weight_ptr + global_channel * WINDOW + 3, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - if USE_PDL: - tl.extra.cuda.gdc_wait() - current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) - - conv = tl.zeros([CS_P2], tl.float32) - for tap_idx in tl.static_range(WINDOW): - source_position = position - (WINDOW - 1) + tap_idx - take = valid & (source_position >= 0) - if tap_idx == WINDOW - 1: - value = tl.where(take, current.to(tl.float32), 0.0) - else: - safe_position = tl.maximum(source_position, 0) - logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) - physical_block = tl.load( - block_table_ptr + sequence * stride_block_table_r + logical_block, - mask=take, - other=0, - ).to(tl.int64) - source_ptr = ( - cache_ptr - + physical_block * stride_cache_block - + head * stride_cache_head - + (safe_position % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) - value = tl.where(take, cached.to(tl.float32), 0.0) - if tap_idx == 0: - weight = weight0 - elif tap_idx == 1: - weight = weight1 - elif tap_idx == 2: - weight = weight2 - else: - weight = weight3 - conv += value * weight - - # Preserve both bf16 rounding points of the original sublayer. - short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) - output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( - tl.bfloat16 - ) - - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - packed = _pack_bf16_pairs(output) - row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair - - for destination in tl.static_range(WORLD): - base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) - tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _gather_norm_kernel( - output_peer_ptrs, - output_peer_offset_u32, - norm_weight_ptr, - normed_ptr, - residual_out_ptr, - eps, - stride_output_t, - C: tl.constexpr, - C_P2: tl.constexpr, - RANK: tl.constexpr, - HAS_NORM: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume a complete row; one CTA owns all outputs for one token.""" - token = tl.program_id(0).to(tl.int64) - pair = tl.arange(0, C_P2 // 2) - pair_mask = pair < C // 2 - output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - output_u32 += output_peer_offset_u32 - offsets = token * (C // 2) + pair - channel = tl.arange(0, C_P2) - channel_mask = channel < C - # Gamma is independent of the preceding sconv publication. Keep the - # acquire immediately before polling the Lamport output slots. - if HAS_NORM: - weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(output_u32, offsets, pair_mask) - row = _unpack_bf16_pairs(packed) - tl.store( - residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask - ) - if HAS_NORM: - row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) - inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) - tl.store( - normed_ptr + token * stride_output_t + channel, - (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), - mask=channel_mask, - ) - tl.store( - output_u32 + offsets, - tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=pair_mask, - ) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _validate_lamport_init_kernel( - input_peer_ptrs, - output_peer_ptrs, - bad_ptr, - num_pairs, - RANK: tl.constexpr, - BLOCK: tl.constexpr, -): - """Validate both complete local allocations through their fabric pointers.""" - offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_pairs - ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) - local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) - value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) - value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) - bad = tl.max( - tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) - ) - if bad != 0: - tl.atomic_max(bad_ptr, 1) - - -class LamportRSConv: - """Persistent symmetric buffers for one TP group.""" - - def _initialize_lamport_buffers(self) -> None: - """Arm every slot and collectively verify the exact sentinel bits. - - Do not replace this with a zero-fill: Lamport readiness distinguishes - the bf16 bit pattern for -0.0 (0x8000) from every published payload. - The validation is deliberately collective so one rank cannot enter the - first polling kernel while another rank still has an unarmed buffer. - """ - if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): - raise RuntimeError("Lamport symmetric buffers must be contiguous") - if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: - raise RuntimeError("Lamport buffers must contain whole uint32 pairs") - - # Fill through int16 so each bf16 lane receives the exact -0.0 bits. - # This covers all three generations and all max-token slots, including - # slots that a smaller first invocation does not touch. - self.buf_in.view(torch.int16).fill_(-0x8000) - self.buf_out.view(torch.int16).fill_(-0x8000) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - # Scan the complete local allocations. Since every rank performs the - # scan and participates in MAX, success proves every symmetric backing - # allocation was armed before any rank is allowed to use generation 0. - bad = torch.logical_or( - self.buf_in.view(torch.int16).ne(-0x8000).any(), - self.buf_out.view(torch.int16).ne(-0x8000).any(), - ).to(dtype=torch.float32) - bad = self.tp.all_reduce(bad) - if int(bad.item()) != 0: - raise RuntimeError("Lamport sentinel initialization failed on a TP rank") - self.tp.barrier() - - def _initialize_mnnvl_buffers(self) -> None: - """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" - self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) - self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - bad = torch.zeros((), dtype=torch.int32, device=self.device) - num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 - _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( - self.input_peer_ptrs, - self.output_peer_ptrs, - bad, - num_pairs, - RANK=self.rank, - BLOCK=256, - num_warps=4, - ) - bad = self.tp.all_reduce(bad.to(torch.float32)) - if int(bad.item()) != 0: - raise RuntimeError("MNNVL Lamport sentinel initialization failed") - self.tp.barrier() - - def __init__( - self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS - ) -> None: - import torch.distributed._symmetric_memory as symm_mem - - tp = get_tp_group() - self.tp = tp - self.group = tp.device_group - self.world_size = tp.world_size - self.rank = tp.rank_in_group - self.device = torch.device(tp.device) - if self.world_size not in (2, 4, 8): - raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") - if hidden_size % (2 * self.world_size) != 0: - raise ValueError("hidden size must produce an even shard on every rank") - if window_size != 4: - raise ValueError(f"short-conv window size must be 4, got {window_size}") - if max_tokens < 1 or max_tokens > _MAX_TOKENS: - raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") - - is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) - self.hidden_size = hidden_size - self.window_size = window_size - self.max_tokens = max_tokens - self.shard_size = hidden_size // self.world_size - # Three generations follow FlashInfer's Lamport layout. A generation - # is reused only after two intervening collective calls. - self.num_buffers = 3 - self.input_generation_bytes = max_tokens * hidden_size * 2 - self.output_generation_bytes = max_tokens * hidden_size * 2 - self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 - if is_cross_node: - try: - from flashinfer.comm.mnnvl import ( - McastGPUBuffer, - TorchDistBackend, - is_mnnvl_fabric_supported, - ) - except ImportError as error: - raise RuntimeError( - "cross-node TP requires FlashInfer MNNVL support" - ) from error - - local_supported = int( - is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) - ) - unsupported = torch.tensor( - 1 - local_supported, dtype=torch.float32, device=self.device - ) - unsupported = tp.all_reduce(unsupported) - if int(unsupported.item()) != 0: - raise RuntimeError("cross-node TP is supported only on MNNVL fabric") - - comm_backend = TorchDistBackend(self.group) - allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 - self._mnnvl_input_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self._mnnvl_output_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() - self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() - self._initialize_mnnvl_buffers() - logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") - else: - self.buf_in = symm_mem.empty( - self.num_buffers, - max_tokens, - self.world_size, - self.shard_size, - dtype=torch.bfloat16, - device=self.device, - ) - self.buf_out = symm_mem.empty( - self.num_buffers, - max_tokens, - hidden_size, - dtype=torch.bfloat16, - device=self.device, - ) - group_name = self.group.group_name - input_handle = symm_mem.rendezvous(self.buf_in, group_name) - output_handle = symm_mem.rendezvous(self.buf_out, group_name) - self.input_peer_ptrs = input_handle.buffer_ptrs_dev - self.output_peer_ptrs = output_handle.buffer_ptrs_dev - self._input_handle = input_handle - self._output_handle = output_handle - self._initialize_lamport_buffers() - self.generation = 0 - - def usable(self, num_tokens: int) -> bool: - return 0 < num_tokens <= self.max_tokens - - def rs_sconv_ag_add_norm( - self, - input_tensor: torch.Tensor, - residual: torch.Tensor, - conv_weight: torch.Tensor, - norm_weight: torch.Tensor | None, - eps: float, - cache: torch.Tensor, - positions: torch.Tensor, - block_table: torch.Tensor, - seq_idx: torch.Tensor, - slot_mapping: torch.Tensor, - off_s: int, - ws: int, - block_size: int, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" - tokens, hidden_size = residual.shape - if not self.usable(tokens): - raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") - if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: - raise ValueError("residual must be bf16 [T, 6144]") - if ( - input_tensor.shape != residual.shape - or input_tensor.dtype != torch.bfloat16 - or input_tensor.stride(1) != 1 - ): - raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") - shard_size = hidden_size // self.world_size - if conv_weight.shape != (shard_size, self.window_size): - raise ValueError( - f"conv_weight must have shape [{shard_size}, {self.window_size}]" - ) - if ( - conv_weight.dtype != torch.bfloat16 - or conv_weight.stride(0) != self.window_size - ): - raise ValueError("conv_weight must be contiguous bf16") - if norm_weight is not None and ( - norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 - ): - raise ValueError("norm_weight must be bf16 [6144] or None") - if cache.dtype != torch.bfloat16 or cache.ndim != 4: - raise ValueError("cache must be a 4-D bf16 tensor") - if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: - raise ValueError("cache head layout is inconsistent with ws") - if off_s < 0 or off_s + ws > cache.shape[3]: - raise ValueError("cache channel offset is out of bounds") - - index = self.generation - input_offset = index * self.input_generation_bytes // 4 - output_offset = index * self.output_generation_bytes // 4 - normed = torch.empty_like(residual) if norm_weight is not None else None - residual_out = torch.empty_like(residual) - phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 - phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) - phase_grid = (tokens, phase_splits) - # Wide CTAs win before the grid saturates; smaller CTAs reduce register - # pressure once high-throughput batches provide enough parallelism. - if 128 <= tokens <= 2048: - phase_warps = 16 - elif tokens > 2048: - phase_warps = 8 - else: - phase_warps = 4 - gather_warps = 4 if tokens >= 256 else 8 - _publish_input_kernel[phase_grid]( - input_tensor, - self.input_peer_ptrs, - input_offset, - input_tensor.stride(0), - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _reduce_insert_kernel[phase_grid]( - self.input_peer_ptrs, - input_offset, - cache, - slot_mapping, - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _sconv_publish_kernel[phase_grid]( - self.output_peer_ptrs, - output_offset, - residual, - conv_weight, - cache, - positions, - seq_idx, - slot_mapping, - block_table, - residual.stride(0), - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_table.stride(0), - block_table.shape[1], - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - WINDOW=self.window_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _gather_norm_kernel[(tokens,)]( - self.output_peer_ptrs, - output_offset, - norm_weight if norm_weight is not None else residual, - normed if normed is not None else residual_out, - residual_out, - eps, - residual.stride(0), - C=hidden_size, - C_P2=triton.next_power_of_2(hidden_size), - RANK=self.rank, - HAS_NORM=norm_weight is not None, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=gather_warps, - ) - self.generation = (index + 1) % self.num_buffers - return normed, residual_out - - -_STATE: LamportRSConv | None = None -_STATE_FAILED = False - - -def initialize_lamport_rs_conv( - hidden_size: int, window_size: int, max_num_batched_tokens: int -) -> None: - """Collectively initialize the TP-group state during model construction.""" - global _STATE, _STATE_FAILED - if _STATE is not None: - if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: - raise RuntimeError("all Lamport users must share hidden and window sizes") - return - if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": - return - try: - max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) - _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) - except Exception: - _STATE_FAILED = True - logger.exception("fused collective unavailable; use the NCCL fallback") - - -def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: - """Return the state initialized with the model, or ``None`` for fallback.""" - if _STATE is not None and ( - _STATE.hidden_size != hidden_size or _STATE.window_size != window_size - ): - raise RuntimeError("all Lamport users must share hidden and window sizes") - return _STATE diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py deleted file mode 100644 index 021cf5274bd2..000000000000 --- a/vllm/models/inkling/amd/ops/mm_towers.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused CUDA kernels for the Inkling vision/audio towers. - -Both kernels keep the reference paths' fp32 accumulation and per-op bf16 -rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently -bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order -(real-checkpoint-weight cosine vs reference > 0.9999998). -""" - -from __future__ import annotations - -import torch - -from vllm.triton_utils import tl, tldevice, triton - -from .norm import _get_num_warps_from_block_size - - -@triton.jit -def _dmel_embed_sum_norm_kernel( - idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) - w_ptr, # [NB * VOCAB, D] bf16 embedding table - norm_w_ptr, # [D] (unused if HAS_NORM=False) - out_ptr, # [T, D] bf16 - eps, - D, - stride_idx_t, - NB: tl.constexpr, - VOCAB: tl.constexpr, - D_P2: tl.constexpr, - HAS_NORM: tl.constexpr, -): - t = tl.program_id(0).to(tl.int64) - offs = tl.arange(0, D_P2) - mask = offs < D - # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), - # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). - acc = tl.zeros([D_P2], dtype=tl.float32) - for b in tl.static_range(NB): - v = tl.load(idx_ptr + t * stride_idx_t + b) - row = (b * VOCAB + v).to(tl.int64) - acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) - h = acc.to(tl.bfloat16) - if HAS_NORM: - # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded - # bf16 multiply with the bf16 weight. - x32 = h.to(tl.float32) - var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D - xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) - w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) - h = xn * w - tl.store(out_ptr + t * D + offs, h, mask=mask) - - -def dmel_embed_sum_norm( - dmel_idx: torch.Tensor, # [T, NB] int32 - weight: torch.Tensor, # [NB * VOCAB, D] bf16 - norm_weight: torch.Tensor | None, - eps: float, -) -> torch.Tensor: - """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no - [T, NB, D] intermediate).""" - T, nb = dmel_idx.shape - D = weight.shape[1] - assert weight.shape[0] % nb == 0 - vocab = weight.shape[0] // nb - out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) - if T == 0: - return out - d_p2 = triton.next_power_of_2(D) - _dmel_embed_sum_norm_kernel[(T,)]( - dmel_idx, - weight, - norm_weight if norm_weight is not None else weight, - out, - eps, - D, - dmel_idx.stride(0), - NB=nb, - VOCAB=vocab, - D_P2=d_p2, - HAS_NORM=norm_weight is not None, - # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% - # at large T (the 80-row gather chain is latency- not lane-bound). - num_warps=8, - ) - return out - - -@triton.jit -def _rmsnorm_gelu_kernel( - x_ptr, # [R, D] bf16 - w_ptr, # [D] - out_ptr, # [R, D] bf16 (or the folded layout when FOLD) - eps, - R, - D, - D_P2: tl.constexpr, - BLOCK_M: tl.constexpr, # rows per block (>1 for small D) - HAS_GELU: tl.constexpr, - FOLD: tl.constexpr, - # fold geometry: input rows index [N, T, H, W]; the store scatters each - # row to (out_row, slot) of fold_timespace_to_depth's output layout. - FT: tl.constexpr, - FH: tl.constexpr, - FW: tl.constexpr, - TF: tl.constexpr, - HF: tl.constexpr, -): - pid = tl.program_id(0).to(tl.int64) - rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) - rmask = rows < R - offs = tl.arange(0, D_P2) - mask = rmask[:, None] & (offs < D)[None, :] - x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - var = tl.sum(x32 * x32, axis=1) / D - xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) - w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) - h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm - if HAS_GELU: - # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, - # matching F.gelu's opmath on a bf16 tensor. - g32 = h.to(tl.float32) - h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) - if FOLD: - # Store directly into the next layer's folded layout (a pure - # permutation — replaces the separate fold copy pass). - t = (rows // (FH * FW)) % FT - hh = (rows // FW) % FH - ww = rows % FW - n = rows // (FT * FH * FW) - slot = ((t % TF) * HF + hh % HF) * HF + ww % HF - out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( - FW // HF - ) + ww // HF - base = (out_row * (TF * HF * HF) + slot) * D - tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) - else: - tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) - - -def rmsnorm_gelu( - x: torch.Tensor, # [..., D] bf16 contiguous - weight: torch.Tensor, - eps: float, - gelu: bool = True, - fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold -) -> torch.Tensor: - """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block - when D is small. With ``fold``, x must be [N, T, H, W, D] and the output - comes back as ``fold_timespace_to_depth(result, *fold)``.""" - D = x.shape[-1] - flat = x.reshape(-1, D) - assert flat.stride(1) == 1 and flat.stride(0) == D - R = flat.shape[0] - if fold is None: - out = torch.empty_like(flat) - ft = fh = fw = tf = hf = 1 - out_shape = x.shape - else: - tf, hf = fold - N, ft, fh, fw, _ = x.shape - out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - if R == 0: - return out.reshape(out_shape) - d_p2 = triton.next_power_of_2(D) - block_m = max(1, 4096 // d_p2) - _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( - flat, - weight, - out, - eps, - R, - D, - D_P2=d_p2, - BLOCK_M=block_m, - HAS_GELU=gelu, - FOLD=fold is not None, - FT=ft, - FH=fh, - FW=fw, - TF=tf, - HF=hf, - num_warps=_get_num_warps_from_block_size(d_p2 * block_m), - ) - return out.reshape(out_shape) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py index 82782fcb96d2..7baec0b42b44 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -587,34 +587,3 @@ def warp_sum_specialized( mask_and_clamp=31, ) return value - - -@cute.jit -def block_sum_specialized( - value: Float32, - warp_sums: cute.Tensor, - tidx: Int32, - warps: cutlass.Constexpr[int], - last_warp_lanes: cutlass.Constexpr[int], - last_warp_mask: cutlass.Constexpr[int], -) -> Float32: - """Upstream-equivalent FP32 block reduction.""" - - lane = cute.arch.lane_idx() - warp_idx = cute.arch.warp_idx() - value = warp_sum_specialized( - value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask - ) - if lane == 0: - warp_sums[warp_idx] = value - cute.arch.barrier() - - block_sum = Float32(0.0) - if warp_idx == 0: - if lane < Int32(warps): - block_sum = warp_sums[lane] - block_sum = cute.arch.warp_reduction_sum(block_sum) - if lane == 0: - warp_sums[0] = block_sum - cute.arch.barrier() - return warp_sums[0] diff --git a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py b/vllm/v1/attention/ops/triton_fp8_mqa_logits.py deleted file mode 100644 index 619d0ec50a9d..000000000000 --- a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py +++ /dev/null @@ -1,262 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel. - -This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942 -tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's -pinned AITER version lacks that fix. - -TODO: Remove this vendored copy once vLLM pins an AITER version that includes -ROCm/aiter#3257 bugfix for gfx942. -""" - -import torch - -from vllm.triton_utils import tl, triton - -# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default -# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold: -# -# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two -# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton -# keeps Q in registers (loop-invariant) and the fp32 scores accumulator -# in VGPRs (heavy VALU), so only the double-buffered KV tile is -# expected to live in LDS. A 0.9 safety factor leaves headroom for any -# LDS overhead the compiler may add. -# -# 2. Hardware ceiling. Defensive upper bound that also counts Q and -# scores against the 64 KiB CU limit, in case a Triton version (older -# or future) decides to spill them to LDS. False positives here only -# shrink the tile; false negatives are JIT-aborts, so we lean -# conservative. -_GFX942_CU_LDS_BYTES = 64 * 1024 -_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB - - -def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool: - """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS.""" - BLOCK_KV = 128 - NUM_STAGES = 2 - kv_bytes = head_size * BLOCK_KV * NUM_STAGES - scores_bytes = num_heads * BLOCK_KV * 4 - q_bytes = num_heads * head_size - fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES - fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES - return fits_occupancy and fits_hardware - - -@triton.jit -def _fp8_mqa_logits_kernel( - Q_ptr, # fp8e4m3 [seq_len, H, D] - KV_ptr, # fp8e4m3 [seq_len_kv, D] - kv_scales_ptr, # fp32 [seq_len_kv] - weights_ptr, # fp32 [seq_len, H] - cu_start_ptr, # int32 [seq_len] - cu_end_ptr, # int32 [seq_len] - logits_ptr, # fp32 [seq_len, seq_len_kv] - seq_len, - seq_len_kv, - NUM_HEADS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - # strides - stride_q_s: tl.int64, - stride_q_h: tl.constexpr, - stride_q_d: tl.constexpr, - stride_kv_s: tl.int64, - stride_kv_d: tl.constexpr, - stride_w_s: tl.int64, - stride_w_h: tl.constexpr, - stride_logits_s: tl.int64, - stride_logits_k: tl.int64, - # block sizes - BLOCK_KV: tl.constexpr, -): - row_id = tl.program_id(0) - # go from larger to smaller in terms of work - # to reduce the tail effect - row_id = tl.num_programs(0) - row_id - 1 - tl.assume(row_id >= 0) - tl.assume(stride_q_s > 0) - tl.assume(stride_q_h > 0) - tl.assume(stride_q_d > 0) - tl.assume(stride_kv_s > 0) - tl.assume(stride_kv_d > 0) - tl.assume(stride_w_s > 0) - tl.assume(stride_w_h > 0) - - logits_row_ptrs = logits_ptr + row_id * stride_logits_s - - h_inds = tl.arange(0, NUM_HEADS)[:, None] - d_inds = tl.arange(0, HEAD_SIZE) - - # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE] - q_ptrs = ( - Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d - ) - - q_block = tl.load(q_ptrs, cache_modifier=".cg") - w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h - w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32) - - # Load start/end for each row in this block - start_ind = tl.load(cu_start_ptr + row_id) - end_ind = tl.load(cu_end_ptr + row_id) - - start_ind = tl.maximum(start_ind, 0) - end_ind = tl.minimum(end_ind, seq_len_kv) - shifted_end = end_ind - start_ind - shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV - - kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind - kv_ptrs = ( - KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d - ) - - kv_scales_ptrs = kv_scales_ptr + kv_col_offsets - - logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k - - # Loop over KV tiles - for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV): - kv_block = tl.load(kv_ptrs) - kv_scales = tl.load(kv_scales_ptrs) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - tl.store(logits_ptrs, scores) - - kv_ptrs += BLOCK_KV * stride_kv_s - kv_scales_ptrs += BLOCK_KV - logits_ptrs += BLOCK_KV * stride_logits_k - kv_col_offsets += BLOCK_KV - - # masked load - kv_col_mask = kv_col_offsets < end_ind - kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0) - kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - # masked store - in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind) - tl.store(logits_ptrs, scores, mask=in_window) - - -def fp8_mqa_logits_gfx942( - q: torch.Tensor, - k_fp8: torch.Tensor, - kv_scales: torch.Tensor, - weights: torch.Tensor, - cu_starts: torch.Tensor, - cu_ends: torch.Tensor, -) -> torch.Tensor: - """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel. - - Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits. - fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on - whether the default tile fits within the 64 KiB LDS budget of a gfx942 - CU (see module docstring). - - Args: - q: Query tensor of shape ``[M, H, D]``, FP8 dtype. - k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype. - kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as - ``[N]``), float32. - weights: Per-head weights of shape ``[M, H]``, float32. - cu_starts: Start indices (inclusive) of shape ``[M]``, int32. - cu_ends: End indices (exclusive) of shape ``[M]``, int32. - - Returns: - Logits of shape ``[M, N]``, float32 -- positions outside - ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with - ``-inf`` so the caller can run a top-k without masking. - """ - seq_len, num_heads, head_size = q.shape - seq_len_kv = k_fp8.shape[0] - assert num_heads & (num_heads - 1) == 0, ( - f"num_heads must be a power of two (got {num_heads})" - ) - assert head_size & (head_size - 1) == 0, ( - f"head_size must be a power of two (got {head_size})" - ) - - # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N - # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller - # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as - # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten - # explicitly to keep the kernel's pointer arithmetic intent clear. - kv_scales_1d = kv_scales.reshape(-1) - - # Initialise with -inf so positions outside [cu_starts, cu_ends) read - # as ``-inf`` after the masked store path -- this matches AITER's - # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects. - logits = torch.full( - (seq_len, seq_len_kv), - fill_value=-float("inf"), - dtype=torch.float32, - device=q.device, - ) - - if _gfx942_default_tile_fits_lds(num_heads, head_size): - block_kv = 128 - num_stages = 2 - else: - # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here: - # default tile spills past gfx942's 64 KiB LDS budget. (64, 1) - # needs ~33 KiB and clears the per-WG budget with margin. - block_kv = 64 - num_stages = 1 - - # heuristic for MFMA instruction shape, identical to AITER's choice - matrix_instr_nonkdim = 32 - if seq_len <= 1024: - matrix_instr_nonkdim = 16 - - stride_q_s, stride_q_h, stride_q_d = q.stride() - stride_kv_s, stride_kv_d = k_fp8.stride() - stride_w_s, stride_w_h = weights.stride() - stride_logits_s, stride_logits_k = logits.stride() - - _fp8_mqa_logits_kernel[(seq_len,)]( - Q_ptr=q, - KV_ptr=k_fp8, - kv_scales_ptr=kv_scales_1d, - weights_ptr=weights, - cu_start_ptr=cu_starts, - cu_end_ptr=cu_ends, - logits_ptr=logits, - seq_len=seq_len, - seq_len_kv=seq_len_kv, - NUM_HEADS=num_heads, - HEAD_SIZE=head_size, - stride_q_s=stride_q_s, - stride_q_h=stride_q_h, - stride_q_d=stride_q_d, - stride_kv_s=stride_kv_s, - stride_kv_d=stride_kv_d, - stride_w_s=stride_w_s, - stride_w_h=stride_w_h, - stride_logits_s=stride_logits_s, - stride_logits_k=stride_logits_k, - BLOCK_KV=block_kv, - num_warps=4, - num_stages=num_stages, - waves_per_eu=2, - matrix_instr_nonkdim=matrix_instr_nonkdim, - ) - - return logits diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py index 321f42bcb199..68ce0e0de2e2 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py @@ -21,15 +21,11 @@ """ import math -from typing import Any import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2 - -from .triton_turboquant_decode import _use_fp8_e4b15 # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory # pressure (mirrors the pattern used in triton_decode_attention.py and @@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return lut.reshape(N * N, 2).contiguous() -# --------------------------------------------------------------------------- -# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2 -# --------------------------------------------------------------------------- - - -@triton.jit -def _tq_decode_stage1_v2( - Q_rot_ptr, - KV_cache_ptr, - KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads - Block_table_ptr, - Seq_lens_ptr, - Centroids_ptr, - Pair_lut_ptr, - Mid_o_ptr, - stride_qb, - stride_qh, - stride_cache_block, # bytes per block (bs*H*slot_aligned) - stride_bt_b, - stride_mid_b, - stride_mid_h, - stride_mid_s, - NUM_Q_HEADS: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - PADDED_SLOT: tl.constexpr, - MAX_NUM_BLOCKS: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - KV_GROUP_SIZE: tl.constexpr, - MSE_BITS: tl.constexpr, - MSE_BYTES: tl.constexpr, - VQB: tl.constexpr, - VAL_DATA_BYTES: tl.constexpr, - # Opt#3 SoA layout - KEY_DATA_BYTES: tl.constexpr, - META_REGION_OFFSET: tl.constexpr, - NUM_SOA_FIELDS: tl.constexpr, - SOA_K_NORM: tl.constexpr, - SOA_V_SCALE: tl.constexpr, - SOA_V_ZERO: tl.constexpr, - N_CENTROIDS: tl.constexpr, - ATTN_SCALE: tl.constexpr, - BLOCK_D: tl.constexpr, - TILE_SIZE: tl.constexpr, - BLOCK_M: tl.constexpr, - KEY_FP8: tl.constexpr, - NORM_CORRECTION: tl.constexpr = 0, - FP8_E4B15: tl.constexpr = 0, - USE_PAIR_LUT: tl.constexpr = 0, - USE_BF16_DOT: tl.constexpr = 0, -): - # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group - bid = tl.program_id(0) - kv_hid = tl.program_id(1) - sid = tl.program_id(2) - - seq_len = tl.load(Seq_lens_ptr + bid) - split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) - split_start = split_len * sid - split_end = tl.minimum(split_start + split_len, seq_len) - - if split_start >= split_end: - return - - # Dimension and tile offsets - d_offs = tl.arange(0, BLOCK_D) - d_mask = d_offs < HEAD_DIM - m_offs = tl.arange(0, BLOCK_M) - t_offs = tl.arange(0, TILE_SIZE) - - # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS) - q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs - q_mask = m_offs < KV_GROUP_SIZE - - # Load Q tile: [BLOCK_M, BLOCK_D] float32 - q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :] - Q = tl.load( - Q_rot_ptr + q_addrs, - mask=q_mask[:, None] & d_mask[None, :], - other=0.0, - ) - - # Pre-warm centroids in L1 - if not KEY_FP8: - _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841 - - # Loop-invariant bit-extraction constants for MSE keys - if not KEY_FP8: - if MSE_BITS == 4: - pass # 4-bit byte extraction is computed inline in the tile loop - else: - # Generic bit extraction - mse_bit_off = d_offs * MSE_BITS - mse_byte_idx = mse_bit_off // 8 - mse_bit_shift = mse_bit_off % 8 - mse_mask_val = (1 << MSE_BITS) - 1 - - # Loop-invariant bit-extraction for 3-bit values - if VQB == 3: - val_bit_off = d_offs * 3 - val_byte_idx = val_bit_off // 8 - val_bit_shift = val_bit_off % 8 - - # exp2 pre-scaling - RCP_LN2: tl.constexpr = 1.4426950408889634 - LN2: tl.constexpr = 0.6931471805599453 - QK_SCALE = ATTN_SCALE * RCP_LN2 - - # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L - M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - L = tl.zeros([BLOCK_M], dtype=tl.float32) - - acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) - - bt_base = bid * stride_bt_b - - # ================================================================ - # TILED LOOP: TILE_SIZE KV tokens per iteration - # ================================================================ - for start_n in range(split_start, split_end, TILE_SIZE): - kv_offs = start_n + t_offs - kv_mask_1d = kv_offs < split_end - - page_idx = kv_offs // BLOCK_SIZE - page_off = kv_offs % BLOCK_SIZE - block_nums = tl.load( - Block_table_ptr + bt_base + page_idx, - mask=kv_mask_1d, - other=0, - ).to(tl.int64) - - # Opt#3 SoA addressing: data region then per-block SoA metadata. - slot_within_block = page_off.to(tl.int64) - block_base = block_nums * stride_cache_block - DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES - data_bases = ( - block_base - + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT) - + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT - ) - head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast( - kv_hid, tl.int64 - ) * (NUM_SOA_FIELDS * BLOCK_SIZE) - knorm_u16_addrs = ( - head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block - ) - vscale_u16_addrs = ( - head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block - ) - vzero_u16_addrs = ( - head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block - ) - - # ============================================================ - # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T) - # ============================================================ - if KEY_FP8: - k_addrs = data_bases[:, None] + d_offs[None, :] - k_raw = tl.load( - KV_cache_ptr + k_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ) - if FP8_E4B15: - K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) - else: - K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) - # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot - K_T = tl.trans(K_f) - # S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range) - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - else: - # --- MSE dequantization --- - if MSE_BITS == 4 and USE_PAIR_LUT: - # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each - # packed byte loaded EXACTLY ONCE instead of twice (the - # previous code built byte_addrs from d_offs // 2 which has - # the pattern [0,0,1,1,2,2,...], duplicating every byte - # load). We load the packed bytes as a [TILE_SIZE, HALF_D] - # tile, decode both nibbles, then issue a single 3-D gather - # against pair_lut which returns (T[lo], T[hi]) in slots - # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D] - # interleaves the pairs into the expected dim order - # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]). - HALF_D: tl.constexpr = BLOCK_D // 2 - half_offs = tl.arange(0, HALF_D) - # Each byte k covers dims (2k, 2k+1); it is in-range iff - # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts, - # so this is equivalent to k < HEAD_DIM // 2. - byte_mask = (half_offs * 2) < HEAD_DIM - - byte_addrs = data_bases[:, None] + half_offs[None, :] - byte_raw = tl.load( - KV_cache_ptr + byte_addrs, - mask=kv_mask_1d[:, None] & byte_mask[None, :], - other=0, - ).to(tl.int32) - - lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D] - hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D] - pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D] - - # Single 3-D gather: pair_lut[pair_key*2 + {0,1}] - # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element. - pair_slot = tl.arange(0, 2) - c_pair = tl.load( - Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :], - mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]), - other=0.0, - ) - # Row-major reshape interleaves pairs: - # out[t, 2k] = c_pair[t, k, 0] = T[lo_k] - # out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k] - c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D]) - elif MSE_BITS == 4: - # 4-bit without pair LUT: exploit byte alignment - half_idx = d_offs // 2 - nibble_shift = (d_offs % 2) * 4 - mse_addrs = data_bases[:, None] + half_idx[None, :] - mse_raw = tl.load( - KV_cache_ptr + mse_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - else: - # Generic 3-bit path: two byte loads → 16-bit → shift/mask - mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :] - mse_raw0 = tl.load( - KV_cache_ptr + mse_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_raw1 = tl.load( - KV_cache_ptr + mse_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = mse_raw0 | (mse_raw1 << 8) - mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - - # Opt#1: norm-correction is pre-folded into the stored per-token - # scalar at store time, so no per-tile sum+sqrt+divide here. The - # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals - # below gives the original c_vec/||c_vec|| * ||k||. - - # Opt#3: K-norms from per-block SoA region. Single u16 load per - # token; contiguous across tokens within a block. - norm_u16 = tl.load( - KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0 - ) - vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32) - - # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d] - K_recon = c_vals * vec_norms[:, None] - # K_T: [BLOCK_D, TILE_SIZE] - K_T = tl.trans(K_recon) - - # S: [BLOCK_M, TILE_SIZE] via tensor core MMA - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - - # Mask out-of-range positions - S = tl.where(kv_mask_1d[None, :], S, float("-inf")) - - # ============================================================ - # ONLINE SOFTMAX with exp2 - # ============================================================ - m_j = tl.maximum(M, tl.max(S, axis=1)) - m_j = tl.where(m_j > float("-inf"), m_j, 0.0) - P = tl.math.exp2(S - m_j[:, None]) - l_j = tl.sum(P, axis=1) - alpha = tl.math.exp2(M - m_j) - - acc = acc * alpha[:, None] - L = L * alpha + l_j - M = m_j - - # ============================================================ - # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V) - # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero - # move to the SoA metadata region (one u16 load each per token). - # ============================================================ - val_bases = data_bases + KEY_DATA_BYTES - - if VQB == 3: - val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] - val_raw0 = tl.load( - KV_cache_ptr + val_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - val_raw1 = tl.load( - KV_cache_ptr + val_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = val_raw0 | (val_raw1 << 8) - v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) - else: # VQB == 4 - # OPT#3 (load-halving, value path): same pattern as OPT#1 for - # keys -- load each byte exactly once as [TILE_SIZE, HALF_D], - # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D] - # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...]. - V_HALF_D: tl.constexpr = BLOCK_D // 2 - v_half_offs = tl.arange(0, V_HALF_D) - v_byte_mask = (v_half_offs * 2) < HEAD_DIM - - val_addrs = val_bases[:, None] + v_half_offs[None, :] - val_byte = tl.load( - KV_cache_ptr + val_addrs, - mask=kv_mask_1d[:, None] & v_byte_mask[None, :], - other=0, - ).to(tl.int32) - v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D] - v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D] - v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D] - - scale_u16 = tl.load( - KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0 - ) - zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0) - v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32) - v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32) - V = v_idx * v_scales[:, None] + v_zeros[:, None] - - # P·V accumulation via tensor core MMA - # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D] - if USE_BF16_DOT: - acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16)) - else: - acc += tl.dot(P.to(tl.float16), V.to(tl.float16)) - - # ================================================================ - # EPILOGUE: Store per-Q-head partial results for stage2 - # ================================================================ - # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M] - # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1] - safe_L = tl.where(L > 0.0, L, 1.0) - one_over_L = 1.0 / safe_L - acc = acc * one_over_L[:, None] - - # Convert M from log2 space back to natural log for stage2 compatibility - lse = M * LN2 + tl.log(safe_L) - - # Store each Q head's result using 2D scatter - # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs - out_addrs = ( - bid * stride_mid_b - + q_head_offs[:, None] * stride_mid_h - + sid * stride_mid_s - + d_offs[None, :] - ) - tl.store( - Mid_o_ptr + out_addrs, - acc, - mask=q_mask[:, None] & d_mask[None, :], - ) - # Store LSE per Q head - lse_addrs = ( - bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM - ) - tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask) - - # --------------------------------------------------------------------------- # Launcher v2 # --------------------------------------------------------------------------- @@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: shows up on a profile, cache by a hash-of-values fingerprint instead. """ return build_pair_lut(centroids) - - -def triton_turboquant_decode_attention_v2( - query: torch.Tensor, - kv_cache: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - Pi: torch.Tensor, - centroids: torch.Tensor, - scale: float, - mse_bits: int, - key_packed_size: int, - value_quant_bits: int, - value_packed_size: int, - max_seq_len: int = 0, # unused; kept for backward compatibility - key_fp8: bool = False, - norm_correction: bool = False, - PiT: torch.Tensor | None = None, - # Pre-allocated buffers (optional, required for CUDA-graph stability). - mid_o_buf: torch.Tensor | None = None, - output_buf: torch.Tensor | None = None, - lse_buf: torch.Tensor | None = None, - buf_holder: Any = None, - # Fixed split count — MUST be a compile-time constant across iterations - # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract. - max_num_kv_splits: int = 32, -) -> torch.Tensor: - """Launch optimized TQ decode attention (v2 stage1 + shared stage2). - - Follows the same buffer-reuse + fixed-grid contract as the v1 launcher - (triton_turboquant_decode_attention) so the backend can capture a CUDA - graph across both versions. - """ - B, Hq, D = query.shape - Hk = kv_cache.shape[2] - block_size = kv_cache.shape[1] - padded_slot = kv_cache.shape[3] - max_num_blocks = block_table.shape[1] - n_centroids = centroids.shape[0] - kv_group_size = Hq // Hk - device = query.device - del max_seq_len # no longer used: splits is fixed via max_num_kv_splits - - cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size) - - # Opt#3 SoA layout constants (match store-side computation). - key_data_bytes = D if key_fp8 else cfg["mse_bytes"] - data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"] - meta_region_offset = block_size * Hk * data_bytes_per_slot - num_soa_fields = 2 if key_fp8 else 3 - soa_k_norm = 0 - soa_v_scale = 0 if key_fp8 else 1 - soa_v_zero = 1 if key_fp8 else 2 - kv_cache_u16 = kv_cache.view(torch.uint16) - - # Compute q_rot = q @ Pi.T - if key_fp8: - q_rot = query.float().contiguous() - else: - q_float = query.float() - if PiT is None: - PiT = Pi.T.contiguous() - q_rot = (q_float @ PiT).contiguous() - - # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores - BLOCK_M = max(16, triton.next_power_of_2(kv_group_size)) - - # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration - TILE_SIZE = 16 - - # Fixed split count — must be constant across calls for cudagraph replay. - NUM_KV_SPLITS = max_num_kv_splits - - # --- mid_o buffer reuse (same pattern as v1 launcher) --- - if ( - mid_o_buf is not None - and mid_o_buf.shape[0] >= B - and mid_o_buf.shape[2] >= NUM_KV_SPLITS - ): - mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] - else: - mid_o = torch.empty( - B, - Hq, - NUM_KV_SPLITS, - D + 1, - dtype=torch.float32, - device=device, - ) - if buf_holder is not None: - buf_holder._tq_mid_o_buf = mid_o - - fp8_e4b15 = _use_fp8_e4b15(device.index or 0) - - # Build pair LUT for 4-bit MSE (FLUTE §3.2) - use_pair_lut = mse_bits == 4 and not key_fp8 - pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids - - # Platform-dependent pipelining depth: ROCm prefers num_stages=1. - stage1_num_stages = 1 if _is_hip else 2 - - # --- Stage 1: v2 kernel --- - # Grid over KV heads (not Q heads) — each program handles all Q heads - grid = (B, Hk, NUM_KV_SPLITS) - _tq_decode_stage1_v2[grid]( - q_rot, - kv_cache, - kv_cache_u16, - block_table, - seq_lens, - centroids, - pair_lut, - mid_o, - q_rot.stride(0), - q_rot.stride(1), - kv_cache.stride(0), - block_table.stride(0), - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - NUM_Q_HEADS=Hq, - NUM_KV_HEADS=Hk, - HEAD_DIM=D, - BLOCK_SIZE=block_size, - PADDED_SLOT=padded_slot, - MAX_NUM_BLOCKS=max_num_blocks, - NUM_KV_SPLITS=NUM_KV_SPLITS, - KV_GROUP_SIZE=kv_group_size, - MSE_BITS=mse_bits, - MSE_BYTES=cfg["mse_bytes"], - VQB=value_quant_bits, - VAL_DATA_BYTES=cfg["val_data_bytes"], - KEY_DATA_BYTES=key_data_bytes, - META_REGION_OFFSET=meta_region_offset, - NUM_SOA_FIELDS=num_soa_fields, - SOA_K_NORM=soa_k_norm, - SOA_V_SCALE=soa_v_scale, - SOA_V_ZERO=soa_v_zero, - N_CENTROIDS=n_centroids, - ATTN_SCALE=scale, - BLOCK_D=cfg["BLOCK_D"], - TILE_SIZE=TILE_SIZE, - BLOCK_M=BLOCK_M, - KEY_FP8=1 if key_fp8 else 0, - NORM_CORRECTION=1 if norm_correction else 0, - FP8_E4B15=fp8_e4b15, - USE_PAIR_LUT=1 if use_pair_lut else 0, - USE_BF16_DOT=1 if _is_hip else 0, - num_warps=4, - num_stages=stage1_num_stages, - ) - - # --- output / lse buffer reuse (same pattern as v1 launcher) --- - if output_buf is not None and output_buf.shape[0] >= B: - output = output_buf[:B, :Hq, :D] - else: - output = torch.empty(B, Hq, D, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_output_buf = output - if lse_buf is not None and lse_buf.shape[0] >= B: - lse = lse_buf[:B, :Hq] - else: - lse = torch.empty(B, Hq, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_lse_buf = lse - - # --- Stage 2: reduce across KV splits (unchanged) --- - grid2 = (B, Hq) - _fwd_kernel_stage2[grid2]( - mid_o, - output, - lse, - seq_lens, - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - output.stride(0), - output.stride(1), - lse.stride(0), - NUM_KV_SPLITS=NUM_KV_SPLITS, - BLOCK_DV=cfg["BLOCK_D"], - Lv=D, - num_warps=4, - num_stages=2, - ) - - return output.to(query.dtype) diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py index 03ef36e5b68c..d77592010a4b 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py @@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return build_pair_lut(centroids) -# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every -# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is -# tied to the caller's tensor. -_pit_qdtype_cache: dict = {} - - -def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor: - key = (PiT.data_ptr(), qdtype) - cached = _pit_qdtype_cache.get(key) - if cached is None: - cached = PiT.to(qdtype).contiguous() - _pit_qdtype_cache[key] = cached - return cached - - def triton_turboquant_unified_attention( query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 From 38ca7a899cd3a9aa6f1e948b1b63f4fae7a685ce Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 15 Sep 2026 16:27:11 -0400 Subject: [PATCH 271/774] [Bugfix][HiSparse] Preserve per-layer offsets in KV cache bindings (#57027) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .buildkite/test_areas/lm_eval.yaml | 18 +++++++++++++++++ .../gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml | 20 +++++++++++++++++++ tests/evals/gsm8k/configs/models-hisparse.txt | 1 + vllm/v1/hisparse/binding.py | 15 ++++++++------ 4 files changed, 48 insertions(+), 6 deletions(-) create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml create mode 100644 tests/evals/gsm8k/configs/models-hisparse.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d9afefc82e4f..ec45f7e69e73 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -116,6 +116,24 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt +- label: ":nvidia: (B200) LM Eval HiSparse Nightly" + key: lm-eval-hisparse-nightly-4xb200 + timeout_in_minutes: 60 + device: b200-k8s + num_devices: 4 + optional: true + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/v1/hisparse/ + - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/ + - vllm/v1/attention/backends/mla/ + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/v1/worker/gpu/ + - csrc/libtorch_stable/hisparse_kernels.cu + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt + - label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 timeout_in_minutes: 360 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml new file mode 100644 index 000000000000..009871925925 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml @@ -0,0 +1,20 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +tolerance: 0.0 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 256 +max_concurrency: 32 +request_timeout_seconds: 1800 +startup_max_wait_seconds: 1200 +server_args: >- + --tensor-parallel-size 4 + --max-model-len 8192 + --max-num-seqs 32 + --max-num-batched-tokens 4096 + --enforce-eager + --kv-cache-dtype bfloat16 + --attention-config '{"hisparse_config":{}}' + --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}' +env: + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/models-hisparse.txt b/tests/evals/gsm8k/configs/models-hisparse.txt new file mode 100644 index 000000000000..53ff181f4f66 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-hisparse.txt @@ -0,0 +1 @@ +GLM-5.2-NVFP4-HiSparse.yaml diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index c23abc9cb4b8..716211fbb30a 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -205,9 +205,12 @@ def bind_hisparse_kv_caches( """Bind existing cache storage and block tables; return the bound handles.""" assert host_pool.registered is not None tensor_configs = { - name: tensor_config + name: ( + tensor_config, + tensor_config.offset + layer_index * tensor_config.layer_stride, + ) for tensor_config in kv_cache_config.kv_cache_tensors - for name in tensor_config.layers + for layer_index, name in enumerate(tensor_config.layers) } resident_source_index = 0 for group_id, group in enumerate(kv_cache_config.kv_cache_groups): @@ -216,12 +219,12 @@ def bind_hisparse_kv_caches( for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle = _get_hisparse_cache(forward_context, layer_name) cache_handle.bind_cache( kv_caches[cache_name], - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, @@ -249,11 +252,11 @@ def bind_hisparse_kv_caches( raise RuntimeError("HiSparse hot tensors must share one GPU backing.") layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] cache_handle = _get_hisparse_cache(forward_context, layer_name) - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle.runtime.bind_hot_cache( raw_tensor, - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, From b6ce7143548a937b10764eaddfd03e459d0dab33 Mon Sep 17 00:00:00 2001 From: zwang86 <46699021+zwang86@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:35:54 -0700 Subject: [PATCH 272/774] [Bugfix] Format kernel-import errors eagerly so warning_once does not retain them (#54098) Signed-off-by: Zeyu Wang Co-authored-by: Zeyu Wang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/platforms/cpu.py | 12 ++++++++---- vllm/platforms/cuda.py | 4 +++- vllm/platforms/interface.py | 2 +- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index ecd2dc5fe623..89b68a401e26 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -605,26 +605,30 @@ def import_kernels(cls) -> None: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once( + "Failed to import from vllm._C: %s", repr(e) + ) else: try: import vllm._C_AVX512 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: logger.warning_once( - "Failed to import from vllm._C_AVX512: %r", e + "Failed to import from vllm._C_AVX512: %s", repr(e) ) else: try: import vllm._C_AVX2 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: - logger.warning_once("Failed to import from vllm._C_AVX2: %r", e) + logger.warning_once( + "Failed to import from vllm._C_AVX2: %s", repr(e) + ) else: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) @classmethod def pack_kv_cache( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 98bddbde296d..30eed82ccb16 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -239,7 +239,9 @@ def import_kernels(cls) -> None: try: import vllm._C_stable_libtorch # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e) + logger.warning_once( + "Failed to import from vllm._C_stable_libtorch: %s", repr(e) + ) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 with contextlib.suppress(ImportError): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7ad0b42da774..bd5d98191201 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -365,7 +365,7 @@ def import_kernels(cls) -> None: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 From f2aad6aa70749ccc4604d31d6f8a04c6728c3361 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 15 Sep 2026 13:36:51 -0700 Subject: [PATCH 273/774] [MRV2] Buffer util simplifications (#56888) Signed-off-by: Nick Hill --- tests/utils_/test_torch_utils.py | 25 +++++++++++++ .../v1/attention/test_indexer_dcp_localize.py | 2 +- tests/v1/worker/test_gpu_pcp_manager.py | 6 ++-- .../layers/attention/mla_attention.py | 35 +++++++++---------- vllm/model_executor/models/diffusion_gemma.py | 6 ++-- vllm/models/dots3_note/nvidia/attention.py | 22 ++++++++---- .../flash_linear_attention/ops/index.py | 5 +-- vllm/utils/torch_utils.py | 17 +++++++-- vllm/v1/attention/backends/flashinfer.py | 17 ++------- .../attention/backends/mla/flashmla_sparse.py | 9 +++-- vllm/v1/attention/backends/mla/indexer.py | 7 ++-- vllm/v1/sample/rejection_sampler.py | 5 ++- vllm/v1/spec_decode/llm_base_proposer.py | 1 - vllm/v1/utils.py | 7 ++-- vllm/v1/worker/cpu/shm.py | 9 +++-- vllm/v1/worker/gpu/buffer_utils.py | 19 ---------- vllm/v1/worker/gpu/model_runner.py | 13 +++---- vllm/v1/worker/gpu/pcp_manager.py | 12 +++---- vllm/v1/worker/gpu/pp_utils.py | 6 ++-- .../gpu/spec_decode/adaptive_verification.py | 9 ++--- vllm/v1/worker/gpu/structured_outputs.py | 5 ++- 21 files changed, 125 insertions(+), 112 deletions(-) diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 6901a162a407..79dd3e320ae6 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch from vllm.utils.torch_utils import ( OMP_NUM_THREADS_SET_BY_VLLM, + async_tensor_h2d, available_cpu_count, common_broadcastable_dtype, current_stream, @@ -192,3 +194,26 @@ def test_runtime_threads_override_vllm_set_omp_num_threads( torch.set_num_threads(3) set_torch_threads_for_runtime() assert torch.get_num_threads() == 1 + + +@pytest.mark.parametrize( + "device", ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"] +) +def test_async_tensor_h2d_staging(device): + """Staging must preserve values while normalizing dtype and layout.""" + # Non-contiguous source with a dtype conversion: would previously sync. + src = torch.arange(12, dtype=torch.int64).reshape(3, 4).T + result = async_tensor_h2d(src, device=device, dtype=torch.int32) + assert result.dtype == torch.int32 + assert result.is_contiguous() + assert torch.equal(result.cpu(), src.to(torch.int32)) + + # numpy and list sources still work. + assert torch.equal( + async_tensor_h2d(np.arange(4, dtype=np.int64), device=device), + torch.arange(4, device=device), + ) + assert torch.equal( + async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32), + torch.tensor([1, 2, 3], dtype=torch.int32, device=device), + ) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 812bce39d889..8bc409adb704 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -996,7 +996,7 @@ def test_pcp_plan_deinterleave_restores_global_order( monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False) monkeypatch.setattr( - "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu", + "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d", lambda x, device: x.to(device), ) scheduled = np.array(req_lens, dtype=np.int64) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index d3ec562aa7a6..2bd7be2cfa34 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -17,7 +17,7 @@ from vllm.v1.worker.gpu.pcp_manager import PCPManager -def _copy_to_cpu(value, out=None, device=None): +def _copy_to_cpu(value, device=None, dtype=None, out=None): tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value if out is not None: return out.copy_(tensor) @@ -83,7 +83,7 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) segments_by_rank, per_rank_num_tokens = manager._build_batch_layout( num_scheduled_tokens=np.ones(3, dtype=np.int32), @@ -158,7 +158,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) with pytest.raises(ValueError, match="smaller than the largest rank-local batch"): manager._build_batch_layout( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 874c43f8cd44..1e673bd2da14 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -269,11 +269,11 @@ LayerNameType, _encode_layer_name, _resolve_layer_name, + async_tensor_h2d, direct_register_custom_op, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, - np_to_pinned_tensor, ) from vllm.v1.attention.backend import ( AttentionBackend, @@ -1980,11 +1980,6 @@ def aligned_split_len(start: int, remaining: int, available: int) -> int: return plans -def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor: - """Pinned int32 CPU tensor backing one concatenated per-chunk field.""" - return np_to_pinned_tensor(np.asarray(values, dtype=np.int32)) - - def align_mla_chunked_context_workspace_size( vllm_config: VllmConfig, workspace_size: int, @@ -2161,24 +2156,28 @@ def padded_rows(rows: int) -> int: token_offset += num_tokens local_token_offset += num_local_tokens - seq_lens_cpu = _flat_int32(seq_lens_flat) - starts_and_context_lens = _flat_int32(starts_flat + context_lens).to( - device, non_blocking=True + seq_lens_cpu = torch.tensor( + seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY + ) + starts_and_context_lens = async_tensor_h2d( + starts_flat + context_lens, device=device, dtype=torch.int32 ) starts = starts_and_context_lens[: len(starts_flat)] context_lens_gpu = starts_and_context_lens[len(starts_flat) :] - cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True) - cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True) - token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to( - device, non_blocking=True + cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32) + cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32) + token_to_seq = async_tensor_h2d( + np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32 ) if use_dcp: - padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to( - device, non_blocking=True + padded_local_cu_seq_lens = async_tensor_h2d( + padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32 + ) + padded_local_token_to_seq = async_tensor_h2d( + np.concatenate(padded_local_token_to_seq_parts), + device=device, + dtype=torch.int32, ) - padded_local_token_to_seq = _flat_int32( - np.concatenate(padded_local_token_to_seq_parts) - ).to(device, non_blocking=True) chunks: list[MLACommonPrefillMetadata.ContextChunk] = [] for index, (plan, layout) in enumerate(zip(plans, layouts)): diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 603d68410df6..c90d78bd6dd1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -53,7 +53,7 @@ from vllm.v1.outputs import LogprobsTensors from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.attn_utils import build_attn_metadata -from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores @@ -1159,7 +1159,7 @@ def _finish_prefills( return # Move the slot indices across once, up front: indexing a device # tensor with a numpy array copies them over synchronously each time. - ps_gpu = async_copy_to_gpu( + ps_gpu = async_tensor_h2d( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) @@ -1266,7 +1266,7 @@ def __call__( # was truncated near max_model_len, in which case the scheduler gave us # fewer than CL logits for that request. valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] - valid_canvas_len = async_copy_to_gpu( + valid_canvas_len = async_tensor_h2d( valid_canvas_len_np.astype(np.int64), device=device ) diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index cbd9ad7f4d78..27928b6e6f43 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -19,7 +19,11 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + PIN_MEMORY, + async_tensor_h2d, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionCGSupport, AttentionLayer, @@ -231,8 +235,12 @@ def _build_sliding_window_metadata( query_lens = query_lens_cpu[req_start:req_end] kv_lens = kv_lens_cpu[req_start:req_end] num_reqs = req_end - req_start - cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) - cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) + cu_seq_lens_q_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) + cu_seq_lens_k_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:]) torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:]) token_to_seq_cpu = torch.repeat_interleave( @@ -248,8 +256,8 @@ def _build_sliding_window_metadata( query_end=query_end, cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True), cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True), - starts=starts_cpu[req_start:req_end].to(device, non_blocking=True), - token_to_seq=token_to_seq_cpu.to(device, non_blocking=True), + starts=async_tensor_h2d(starts_cpu[req_start:req_end], device), + token_to_seq=async_tensor_h2d(token_to_seq_cpu, device), num_kv_tokens=num_kv_tokens, max_seq_len_q=int(query_lens.max().item()), max_seq_len_k=int(kv_lens.max().item()), @@ -396,8 +404,8 @@ def build( seq_lens = common_attn_metadata.seq_lens[ reqs_start + req_start : reqs_start + req_end ] - query_lens = query_lens_cpu[req_start:req_end].to( - self.device, non_blocking=True + query_lens = async_tensor_h2d( + query_lens_cpu[req_start:req_end], self.device ) kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1) chunk.cu_seq_lens_k.zero_() diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py index cc6f54d8af0b..880be7747dcc 100644 --- a/vllm/third_party/flash_linear_attention/ops/index.py +++ b/vllm/third_party/flash_linear_attention/ops/index.py @@ -10,6 +10,7 @@ import torch from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.triton_utils import triton from .utils import tensor_cache @@ -27,8 +28,8 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() indices = torch.cat([torch.arange(n) for n in chunk_counts]) chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1) - return chunk_indices.to( - device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True + return async_tensor_h2d( + chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype ) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 6f7dc1434637..0b330edf0038 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -705,17 +705,30 @@ def create_kv_caches_with_random( def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Copy list/numpy array/tensor async from host to device.""" + if dtype is None and out is not None: + dtype = out.dtype if isinstance(data, np.ndarray): data = torch.from_numpy(data) if isinstance(data, torch.Tensor): - t = data.pin_memory() if PIN_MEMORY else data + t = data + if PIN_MEMORY and not t.is_pinned(): + # Stage in pinned, contiguous buffer to ensure fully async copy. + t = torch.empty( + t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True + ).copy_(t) else: t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") assert t.is_cpu + + if out is not None: + assert out.dtype == dtype + return out.copy_(t, non_blocking=True) + assert device is not None, "must provide destination tensor or device" return t.to(device=device, dtype=dtype, non_blocking=True) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 0da68e3e3d93..7ab2b2018e09 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1264,15 +1264,9 @@ def _compute_flashinfer_kv_metadata( """ # write self.paged_kv_indptr_cpu inplace (0-index is always 0) np.cumsum( - num_blocks_np, - dtype=np.int32, - out=self.paged_kv_indptr.np[1 : num_reqs + 1], + num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1] ) - paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1] - paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1] - if PIN_MEMORY: - paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory() - paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True) + paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1) # write self.paged_kv_indices inplace num_actual_pages = self.paged_kv_indptr.np[num_reqs] @@ -1292,12 +1286,7 @@ def _compute_flashinfer_kv_metadata( page_size, paged_kv_last_page_len_np, ) - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs] - if PIN_MEMORY: - paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory() - self.paged_kv_last_page_len.gpu[:num_reqs].copy_( - paged_kv_last_page_len_cpu, non_blocking=True - ) + self.paged_kv_last_page_len.copy_to_gpu(num_reqs) return paged_kv_indices def build( diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 9c440776ae15..3a9ca2ffdb06 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -20,7 +20,11 @@ from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache +from vllm.utils.torch_utils import ( + async_tensor_h2d, + current_stream, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -48,7 +52,6 @@ get_mla_metadata, ) from vllm.v1.kv_cache_interface import AttentionSpec -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -519,7 +522,7 @@ def _build_fp8_separate_prefill_decode( ) workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) max_prefill_buffer_size //= self.dcp_world_size - entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device) + entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device) prefill_block_table = prefill_block_table.index_select(0, entry_rows) prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) num_entries = len(workspace_rows) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 28bd450da1ca..837dbca817c0 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -27,7 +27,7 @@ ) from vllm.utils.math_utils import round_down from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -47,7 +47,6 @@ KVCacheSpec, MLAAttentionSpec, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) @@ -338,14 +337,14 @@ def build_pcp_global_chunk_plan( + local ) - cu = async_copy_to_gpu(cu, device=device) + cu = async_tensor_h2d(cu, device=device) return PCPGlobalChunkPlan( row_start_cu=cu[0], global_cu=cu[1], padded_local_cu=cu[2], padded_local_total=padded_total, total=total, - deinterleave_idx=async_copy_to_gpu(idx, device=device), + deinterleave_idx=async_tensor_h2d(idx, device=device), ) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 5aa83454739e..85b623bd84fa 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -19,6 +19,7 @@ triton_kernel_dispatcher_with_warmup, ) from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor from vllm.v1.sample.metadata import SamplingMetadata @@ -327,9 +328,7 @@ def apply_logits_processors( num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu") original_indices = torch.arange(num_requests, device="cpu") repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens) - repeat_indices = repeat_indices_cpu.to( - device=logits.device, non_blocking=True - ) + repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device) logits = self.apply_penalties( logits, sampling_metadata, metadata, repeat_indices, output_token_ids ) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index c0e456521ca0..1bd1bdc45bbd 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -255,7 +255,6 @@ def __init__( self.backup_next_token_ids = CpuGpuBuffer( self.max_batch_size, dtype=torch.int32, - pin_memory=PIN_MEMORY, device=device, with_numpy=True, ) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index da17110d80af..facafd28608f 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -137,9 +137,10 @@ def __init__( self.np = self.cpu.numpy() def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: - if n is None: - return self.gpu.copy_(self.cpu, non_blocking=True) - return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True) + cpu, gpu = self.cpu, self.gpu + if n is not None: + cpu, gpu = cpu[:n], gpu[:n] + return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True) def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: """NOTE: Because this method is non-blocking, explicit synchronization diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index f92c6c3bf28d..b615f4b5b5c8 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -76,14 +76,17 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]: def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(data, np.ndarray): data = torch.from_numpy(data) - if isinstance(data, torch.Tensor): + if not isinstance(data, torch.Tensor): + data = torch.tensor(data, dtype=dtype, device="cpu") + elif out is None: return data.to(dtype=dtype) - return torch.tensor(data, dtype=dtype, device="cpu") + return data if out is None else out.copy_(data) torch_utils.async_tensor_h2d = async_tensor_h2d diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 0066f80907ec..0731072e5968 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -10,7 +10,6 @@ from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( - PIN_MEMORY, async_tensor_h2d, get_accelerator_view_from_cpu_tensor, ) @@ -27,24 +26,6 @@ def set_default_max_concurrency(n: int) -> None: _DEFAULT_MAX_CONCURRENCY = max(2, n) -def async_copy_to_gpu( - x: torch.Tensor | np.ndarray, - out: torch.Tensor | None = None, - device: torch.device | None = None, -) -> torch.Tensor: - if isinstance(x, np.ndarray): - x = torch.from_numpy(x) - assert x.is_cpu - - if out is None: - assert device is not None - out = torch.empty_like(x, device=device) - - # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() if PIN_MEMORY else x - return out.copy_(pinned, non_blocking=True) - - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): if not is_uva_available(): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 25395c3330eb..b66d552cb487 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -64,7 +64,7 @@ from vllm.tasks import SupportedTask from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -99,10 +99,7 @@ init_kv_cache, ) from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import ( - async_copy_to_gpu, - set_default_max_concurrency, -) +from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -1289,7 +1286,7 @@ def prepare_inputs( req_ids = batch_req_state.req_ids num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens idx_mapping_np = batch_req_state.idx_mapping_np - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) num_reqs = len(req_ids) # Get the number of draft tokens for each request. @@ -1324,7 +1321,7 @@ def prepare_inputs( cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) - cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) + cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) adaptive_verification = ( self.adaptive_verification if num_draft_tokens_per_req is not None else None @@ -1352,7 +1349,7 @@ def prepare_inputs( # Some attention backends like FA3 require query_start_loc to be non-decreasing. query_start_loc_np[num_reqs + 1 :] = num_tokens query_start_loc = self.input_buffers.query_start_loc - async_copy_to_gpu(query_start_loc_np, out=query_start_loc) + async_tensor_h2d(query_start_loc_np, out=query_start_loc) if adaptive_verification is not None: cu_num_logits, query_start_loc, total_num_draft_tokens = ( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index adfbf2f9f0ba..376e0f251012 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -10,9 +10,9 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import ( InputBatch, InputBuffers, @@ -358,13 +358,13 @@ def _build_batch_layout( dtype=np.int64, ) - self._hidden_restore_idx = async_copy_to_gpu( + self._hidden_restore_idx = async_tensor_h2d( hidden_restore_idx, device=self.device ) - self._padded_gather_idx = async_copy_to_gpu( + self._padded_gather_idx = async_tensor_h2d( padded_gather_idx, device=self.device ) - self._gathered_kv_write_mask = async_copy_to_gpu( + self._gathered_kv_write_mask = async_tensor_h2d( gathered_kv_write_mask, device=self.device ) return segments_by_rank, per_rank_num_tokens @@ -536,12 +536,12 @@ def partition_batch( local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1] np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens - async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) + async_tensor_h2d(local_query_start_loc_np, out=input_buffers.query_start_loc) local_query_start_loc = input_buffers.query_start_loc[ : num_reqs_after_padding + 1 ] - local_to_global_req_idx = async_copy_to_gpu( + local_to_global_req_idx = async_tensor_h2d( local_to_global_req_idx_np, device=self.device ) seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 8c65a23ef7ec..ab1b991d1a8d 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -11,7 +11,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -138,7 +138,7 @@ def get_prev_sampled_outputs( return None # Filter excluded request indices. idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) self.main_stream.wait_event(slot.event) if slot.draft_tokens is not None and draft_tokens_to_update is not None: @@ -148,7 +148,7 @@ def get_prev_sampled_outputs( keep = ~exclude_mask keep_t = torch.as_tensor(keep, device=self.device) draft_tokens = draft_tokens[keep_t] - draft_idx_mapping = async_copy_to_gpu( + draft_idx_mapping = async_tensor_h2d( slot.idx_mapping_np[keep], device=self.device ) draft_tokens_to_update[draft_idx_mapping] = draft_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index af58b8993bfa..245f304ac4e3 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -14,6 +14,7 @@ from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream @@ -21,7 +22,6 @@ get_attn_cg_support, get_query_lens_mismatch_unsupported_backend, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) _PROFILE_REPLAYS = 5 @@ -413,7 +413,7 @@ def reallocate_drafts( if draft_budget == 0: capacities.zero_() else: - async_copy_to_gpu(scheduled_drafts, out=capacities) + async_tensor_h2d(scheduled_drafts, out=capacities) if draft_budget < int(scheduled_drafts.sum()): _assign_draft_token_budget_compiled( self._confidence_probs, @@ -424,10 +424,7 @@ def reallocate_drafts( ) num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs] - async_copy_to_gpu( - num_non_draft_tokens, - out=num_non_draft_tokens_gpu, - ) + async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu) self._cu_num_logits[:1].zero_() torch.cumsum( capacities + self.num_bonus_tokens, diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 8437b0aed9a4..057d6f7154c7 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -5,8 +5,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import PIN_MEMORY -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -68,7 +67,7 @@ def apply_grammar_bitmask( # Asynchronously copy the bitmask to GPU. with torch.cuda.stream(self.copy_stream): - bitmask = async_copy_to_gpu( + bitmask = async_tensor_h2d( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) From 12575123059c2142e39c9af4e217a4792c3c8176 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Wed, 16 Sep 2026 05:07:14 +0800 Subject: [PATCH 274/774] [DSA] Wire DeepGEMM sparse MQA logits into the DeepSeek V4.1 indexer (#56254) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 --- .../test_fused_indexer_q_rope_quant.py | 52 ++ .../layers/test_mla_short_prefill_indexer.py | 544 ++++++++++++++++++ vllm/config/attention.py | 9 + .../attention/dsa/sparse_mqa_logits.py | 469 +++++++++++++++ .../layers/sparse_mqa_indexer.py | 209 +++++++ .../deepseek_v4/common/ops/fused_indexer_q.py | 24 +- .../nvidia/ops/fused_indexer_q_cutedsl.py | 17 +- vllm/models/deepseek_v41/attention.py | 87 ++- vllm/utils/deep_gemm.py | 183 ++++++ vllm/v1/attention/backends/mla/indexer.py | 7 +- .../attention/backends/mla/sparse_indexer.py | 235 ++++++++ vllm/v1/core/kv_cache_utils.py | 9 + vllm/v1/kv_cache_interface.py | 11 +- 13 files changed, 1830 insertions(+), 26 deletions(-) create mode 100644 vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py create mode 100644 vllm/model_executor/layers/sparse_mqa_indexer.py create mode 100644 vllm/v1/attention/backends/mla/sparse_indexer.py diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 8770882ddf7c..7464c9076973 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -560,3 +560,55 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat torch.testing.assert_close( k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 ) + + +@pytest.mark.parametrize("use_cutedsl", [False, True]) +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", +) +@torch.inference_mode() +def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): + """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's + sparse MQA-logits kernels take) instead of fp32; the values are the fp32 + result rounded once, so the scoring kernel needs no cast.""" + if use_cutedsl and not has_cutedsl(): + pytest.skip("cutedsl (cutlass) not installed") + + device = "cuda" + torch.manual_seed(0) + num_tokens, n_head = 257, 32 + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) + positions = torch.randint( + 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device + ) + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) + softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above + + _, weights_ref = _reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True + ) + cutedsl_patch = ( + mock.patch( + "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", + return_value=False, + ) + if not use_cutedsl + else contextlib.nullcontext() + ) + with cutedsl_patch: + _, weights_fused = fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=True, + weights_out_dtype=torch.bfloat16, + ) + assert weights_fused.dtype == torch.bfloat16 + assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16)) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index f8e7a89ccad2..41a06d0546f0 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -450,3 +450,547 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) reference = torch.where(keep, 3.0, -torch.inf) torch.testing.assert_close(logits, reference, rtol=0, atol=0) + + +# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) --- + + +def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): + """Production-faithful candidates: top min(avail, K) blocks per row — + full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + avail = ((lens.cpu() + cbk - 1) // cbk).tolist() + rows = [] + for a in avail: + n = min(a, num_candidates) + picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist() + rows.append(picks + [-1] * (num_candidates - len(picks))) + return torch.tensor(rows, dtype=torch.int32) + + +def _candidate_token_mask( + candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None +): + """[rows, total] bool mask of the candidate-covered token positions. + + ``ks`` converts request-local candidate blocks to absolute positions + (prefill's packed workspace); pass None when already request-local. + """ + rows = candidates.shape[0] + device = candidates.device + # Sentinel block keeps -1 padding from scattering onto real columns. + width = max(total, (candidates.shape[1] + 1) * cbk) + if ks is not None: + width += int(ks.max()) + cbk + in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device) + block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1]) + tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device) + if ks is not None: + tok = tok + ks.unsqueeze(1).unsqueeze(2) + keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk) + in_cand.scatter_(1, tok.flatten(1), keep.flatten(1)) + return in_cand[:, :total] + + +def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: + """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid + count are -1 (the indexer's "no token" sentinel).""" + values, indices = masked.topk(topk, dim=1) + return torch.where(values > float("-inf"), indices.int(), -1) + + +def _assert_same_selection( + logits: torch.Tensor, + got: torch.Tensor, + expected: torch.Tensor, + row_starts: torch.Tensor | None = None, +) -> None: + """Both top-k index sets must pick the same multiset of logit values per + row (-1 padding must match too). Indices themselves may differ: the + top-k kernels break ties among equal values in nondeterministic order.""" + for r in range(logits.shape[0]): + base = int(row_starts[r]) if row_starts is not None else 0 + g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] + assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected" + vg = logits[r][(base + g).long()].sort().values + ve = logits[r][(base + e).long()].sort().values + assert torch.equal(vg, ve), f"row {r}: selected values differ" + + +def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: + """Caller-owned buffers the sparse path writes into (the metadata builder + owns these in production).""" + i32 = dict(dtype=torch.int32, device="cuda") + return dict( + sparse_indices=torch.zeros(rows, num_sparse, **i32), + end=torch.zeros(rows, **i32), + col_indices=torch.zeros(rows, topk, **i32), + ) + + +def _expand_candidates_reference( + candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand + each candidate block into ``cbk // sbk`` sparse blocks, keep those inside + ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by + repeating the last valid block, and count the valid sparse columns.""" + ratio = cbk // sbk + rows, num_candidates = candidates.shape + num_sparse = num_candidates * ratio + big = torch.iinfo(torch.int32).max + offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32) + blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1) + max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk + valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1)) + sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1) + n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32) + base = (ks // sbk).unsqueeze(1) + kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base + last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long() + last_block = kernel_blocks.gather(1, last_slot) + pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1) + sparse_indices = torch.where(pad, last_block, kernel_blocks) + last_start = last_block.squeeze(1) * sbk + ks % sbk + last_fill = (ke - last_start).clamp(min=1, max=sbk) + end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0) + return sparse_indices.int(), end.int() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +def test_candidate_blocks_to_sparse_indices_math(): + """Candidate->sparse-block expansion: ks anchoring, range filtering, + repeat-last padding, and the per-row valid column count (``end``).""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + # Rows cover: exactly one block, an unaligned ks with a partial tail + # block, an empty row, duplicates / -1 padding / out-of-range garbage + # (warmup runs touch uninitialized buffers), and int32-overflow junk. + ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda") + ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda") + candidates = torch.tensor( + [ + [0, -1, -1, -1], + [2, 0, -1, -1], + [3, -1, -1, -1], + [1, 1, 7, 0], + [2**30, -(2**30), 0, -1], + ], + dtype=torch.int32, + device="cuda", + ) + indices, end = candidate_blocks_to_sparse_indices( + candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8 + ) + assert indices.tolist() == [ + [8, 8, 8, 8], + [1, 3, 3, 3], + [12, 12, 12, 12], + [1, 2, 2, 2], + [0, 0, 0, 0], + ] + # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token; + # row 3 keeps the duplicated block (production candidates are unique). + assert end.tolist() == [8, 9, 0, 24, 8] + + # Wider candidate blocks expand into multiple sparse blocks. + idx2, end2 = candidate_blocks_to_sparse_indices( + torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([48], dtype=torch.int32, device="cuda"), + candidate_block_size=16, + sparse_block_kv=8, + ) + assert idx2.tolist() == [[0, 1, 2, 3]] + assert end2.tolist() == [32] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)]) +def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): + """The Triton expansion equals the PyTorch reference on production-shaped + random inputs, including strided (column-sliced) candidate rows and + strided output buffers.""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + torch.manual_seed(0) + rows, num_candidates = 41, 2048 + ratio = cbk // sbk + lens = torch.randint(1, 40000, (rows,), dtype=torch.int32) + lens[0], lens[1] = 1, 3 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke = ks.cuda(), (ks + lens).cuda() + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph + # warm-up on uninitialized buffers would. + candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32) + + ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk) + + cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda() + cand_buf[:, :num_candidates] = candidates + si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda() + end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf[::2] = ks + got_indices, got_end = candidate_blocks_to_sparse_indices( + cand_buf[:, :num_candidates], + ks_buf[::2], + ke, + cbk, + sbk, + out=(si_buf[:, : num_candidates * ratio], end_buf[::2]), + ) + assert torch.equal(got_indices, ref_indices) + assert torch.equal(got_end, ref_end) + assert torch.equal( + cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7) + ) + assert not si_buf[:, num_candidates * ratio :].any() + assert not end_buf[1::2].any() + + +def _skip_unless_sm100_sparse_kernels(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 required") + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"): + pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits") + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + ) + + if not has_deep_select(): + pytest.skip("DeepSelect extension (vllm._deepselect_C) required") + return deep_gemm + + +def _quant_fp4(x: torch.Tensor): + from deep_gemm.utils import per_token_cast_to_fp4 + + return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): + """The sparse-MQA candidate path selects exactly the same logit values as + ``torch.topk`` over dense bf16 logits masked to the candidates (the two + DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA + graph capture, and a later layer reusing the step's expansion/schedule + selects the same values.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_prefill_chunk, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 37, 32, 128 + cbk, topk, num_candidates = 8, 512, 2048 # V4.1-Flash production values + lens = torch.randint(1, 30000, (rows,), dtype=torch.int32) + lens[0], lens[1], lens[2] = 1, 5, 2048 # single token, sub-block, long + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + + # Reference: dense bf16 logits masked to the candidates, then torch.topk. + dense = deep_gemm.fp8_fp4_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + ks, + ke, + clean_logits=False, + logits_dtype=torch.bfloat16, + ) + cols = torch.arange(total, device="cuda") + in_range = (cols - ks.unsqueeze(1) >= 0) & ( + cols - ks.unsqueeze(1) < lens.unsqueeze(1) + ) + masked = dense.masked_fill( + ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), + float("-inf"), + ).float() + ref_abs = _reference_topk(masked, topk) # packed-workspace positions + ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(kernel_metadata=None): + return sparse_mqa_logits_prefill_chunk( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + k_fp.view(torch.int8), + k_sf.view(total), + weights, + ks, + ke, + candidates, + cbk, + cbk, + topk, + out, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse() # warm up JIT + the DeepGEMM workspace pre-capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse() + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse() + _assert_same_selection(dense, out, ref, row_starts=ks) + out.fill_(-1) + run_sparse(kernel_metadata) + + _assert_same_selection(dense, out, ref, row_starts=ks) + + +def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): + """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' + fp4 data first, then their packed UE8M0 scales.""" + device = req_lens.device + pages_per = (req_lens + page_kv - 1) // page_kv + num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + kv_cache = torch.zeros( + num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device + ) + block_table = torch.zeros( + len(req_lens), max_pages, dtype=torch.int32, device=device + ) + page_pool = torch.randperm(num_pages, device=device) + offset = 0 + for b in range(len(req_lens)): + n = int(pages_per[b]) + block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32) + offset += n + k_fp, k_sf = _quant_fp4( + torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16) + ) + rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device) + rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8) + rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4) + by_page = rows68.view(n, page_kv, 68) + split = torch.cat( + [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1 + ).view(n, page_kv, 1, head_dim // 2 + 4) + kv_cache[block_table[b, :n].long()] = split + return kv_cache, block_table + + +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): + """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks + from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to + one query row each), CUDA graph capture/replay and schedule reuse by a + later layer.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_paged_decode, + ) + + torch.manual_seed(0) + batch, num_heads, head_dim, page_kv = 9, 32, 128, 128 + cbk, topk, num_candidates = 8, 512, 2048 + rows = batch * next_n + req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda") + req_lens[0], req_lens[1] = 1, 30000 # tiny row + a filtered (subset) row + # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout). + row_lens = ( + (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda")) + .clamp(min=0) + .to(torch.int32) + ) + kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) + assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2) + q_scale = q_sf.view(batch, next_n, num_heads) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() + + def reference_topk(): + context_lens_2d = row_lens.view(batch, next_n) + sched = deep_gemm.get_paged_mqa_logits_metadata( + context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + (q, q_scale), + kv_cache, + weights, + context_lens_2d, + block_table, + sched, + int(req_lens.max()), + False, + torch.bfloat16, + None, + ) + lens_flat = row_lens.flatten() + cols = torch.arange(logits.shape[1], device="cuda") + masked = logits.masked_fill( + ~( + (cols.unsqueeze(0) < lens_flat.unsqueeze(1)) + & _candidate_token_mask(candidates, cbk, logits.shape[1]) + ), + float("-inf"), + ).float() + return _reference_topk(masked, topk), logits + + expected, dense_logits = reference_topk() + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + + # The metadata builder flattens every query to one row: per-row context + # lengths, an expanded block table and a row -> request map. + context_lens = row_lens.flatten().contiguous() + row_block_table = block_table.repeat_interleave(next_n, dim=0) + row_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + row_indices = row_indices.repeat_interleave(next_n) + row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(o: torch.Tensor, kernel_metadata=None): + return sparse_mqa_logits_paged_decode( + q.view(rows, 1, num_heads, head_dim // 2), + q_scale.view(rows, 1, num_heads), + kv_cache, + weights, + context_lens, + row_block_table, + row_indices, + candidates, + cbk, + cbk, + topk, + o, + row_ks=row_ks, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse(out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse(out) + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse(out) + _assert_same_selection(dense_logits, out, expected) + out.fill_(-1) + run_sparse(out, kernel_metadata) + + _assert_same_selection(dense_logits, out, expected) + + +def test_sparse_mqa_logits_accuracy_vs_torch(): + """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized + MXFP4 inputs.""" + _skip_unless_sm100_sparse_kernels() + from deep_gemm.utils import cast_back_from_fp4 + + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + from vllm.utils.deep_gemm import ( + fp8_fp4_sparse_mqa_logits, + get_sparse_mqa_logits_metadata, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 17, 32, 128 + sbk, num_candidates = 8, 2048 + lens = torch.randint(1, 1500, (rows,), dtype=torch.int32) + lens[0] = 1 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, sbk).cuda() + + sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk) + metadata = get_sparse_mqa_logits_metadata( + ks, ke, total, sparse_idx, torch.int8, sbk + ) + logits = fp8_fp4_sparse_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + metadata, + sparse_idx.shape[1], + sbk, + ) + + # PyTorch fp32 reference on the dequantized inputs. + q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True) + k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True) + dots = torch.einsum( + "rhd,td->rht", + q_deq.view(rows, num_heads, head_dim).float(), + k_deq.float(), + ) + ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1) # [rows, total] + + # Compare on valid sparse columns (absolute packed-workspace positions). + width = sparse_idx.shape[1] * sbk + cols = torch.arange(width, device="cuda") + valid = cols.unsqueeze(0) < end.unsqueeze(1).long() + pos = ( + sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk + + (ks % sbk).unsqueeze(1) + + (cols % sbk).unsqueeze(0) + ).clamp(min=0, max=total - 1) + abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid] + row_scale = ref.abs().amax(dim=1) + norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[ + valid + ] + print( + f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; " + f"rowmax-normalized p50={norm_err.median().item():.5f} " + f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}" + ) + # bf16-level numerics: far below the value gaps that change top-k selection. + assert norm_err.quantile(0.99).item() < 1e-2 + assert norm_err.max().item() < 2e-2 diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 129fbcbb0869..3646677728b5 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -86,6 +86,15 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_sparse_logits: bool = False + """DeepSeek V4.1 two-level indexer: score only the candidate blocks with + DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over + the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`, + an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension + (the top-k runs on the kernels' bf16 logits). The sparse path costs + O(candidate blocks) per query regardless of context length, so it pays off + for long contexts (roughly 32K tokens and beyond) and is slower below.""" + hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py new file mode 100644 index 000000000000..839c581583a7 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer. + +The candidate-source indexer publishes the top candidate blocks (in units of +``candidate_block_size`` compressed positions). Instead of scoring every KV +position and masking the dense logits down to those blocks, this path scores +only the candidate tokens with DeepGEMM's sparse MQA logits kernels +(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs +DeepSelect's top-k directly on the bf16 logits they produce, and remaps the +selected columns back to request-local positions. + +Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed +(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The +FP8 indexer cache stores one fp32 scale per quant block instead. + +The functions here take caller-owned output buffers (``sparse_indices``, +``end``, ``col_indices``) so the hot path allocates nothing; the indexer +metadata builder owns them. +""" + +import importlib.util + +import torch + +from vllm.model_executor.layers.indexer_topk import ( + deep_select_topk, + get_deep_select_stride_requirement, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import ( + fp8_fp4_paged_sparse_mqa_logits, + fp8_fp4_sparse_mqa_logits, + get_paged_sparse_mqa_logits_metadata, + get_sparse_mqa_logits_metadata, +) + +# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels. +SPARSE_BLOCK_KV_CHOICES = (16, 8) + + +def has_deep_select() -> bool: + """Whether the DeepSelect top-k extension is built and this GPU runs it.""" + return ( + importlib.util.find_spec("vllm._deepselect_C") is not None + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) + + +def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: + """Raise unless DeepSelect accepts the sparse logits layout. + + The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]`` + logits and the top-k writes int32 ``[rows, topk_tokens]`` indices; + DeepSelect needs both row strides aligned (1024B / 32B on SM100). + """ + in_align, out_align = get_deep_select_stride_requirement() + logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize + if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0: + raise ValueError( + f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 " + f"columns, {logits_row_bytes} B) to be {in_align}B-aligned." + ) + if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0: + raise ValueError( + f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 " + f"output row, got topk={topk_tokens}." + ) + + +_INT32_BIG = tl.constexpr(2**31 - 1) + + +@triton.jit +def _expand_candidates_kernel( + cand_ptr, + cand_stride, + ks_ptr, + ks_stride, + ke_ptr, + ke_stride, + si_ptr, + si_stride, + end_ptr, + end_stride, + K: tl.constexpr, + S: tl.constexpr, + RATIO: tl.constexpr, + SBK: tl.constexpr, +): + """One program per row: expand + filter + sort candidate blocks. + + Each row's valid sparse blocks come out sorted ascending (no dedup pass: + production candidates are unique, and the kernel tolerates repeats), + padded by repeating the last valid block — DeepGEMM's padding convention. + ``end`` counts valid sparse-token columns; valid columns form a prefix + because block ids are sorted. + """ + row = tl.program_id(0) + ks = tl.load(ks_ptr + row * ks_stride) + ke = tl.load(ke_ptr + row * ke_stride) + cols = tl.arange(0, K) + cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask + if RATIO == 2: + cand = tl.interleave(cand * RATIO, cand * RATIO + 1) + + # Keep only blocks inside the row's KV range [0, ke - ks); the lower + # bound is free since candidates are non-negative. + kv_len = tl.maximum(ke - ks, 0) + max_block = (kv_len + SBK - 1) // SBK + valid = (cand >= 0) & (cand < max_block) + vals = tl.where(valid, cand, _INT32_BIG) + vals = tl.sort(vals) + + # Anchor at the row's ks rounded down to a sparse-block boundary + # (identity for the paged path, where ks == 0). + base = ks // SBK + si = tl.where(vals != _INT32_BIG, vals, 0) + base + n_valid = tl.sum(valid.to(tl.int32)) + last_rel = tl.max(tl.where(valid, cand, -1)) + pad_val = tl.where(n_valid > 0, last_rel + base, base) + tl.store( + si_ptr + row * si_stride + tl.arange(0, S), + tl.where(tl.arange(0, S) < n_valid, si, pad_val), + ) + + # Valid sparse-token column count; only the last block can be partial. + # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK, + # which is exactly last_rel * SBK + ks. + last_start = last_rel * SBK + ks + fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK) + end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0) + tl.store(end_ptr + row * end_stride, end) + + +@triton.jit +def _sparse_topk_remap_kernel( + col_ptr, + col_stride, + si_ptr, + si_stride, + ks_ptr, + ks_stride, + out_ptr, + out_stride, + k, + SBK: tl.constexpr, + K_POW2: tl.constexpr, +): + """Remap sparse top-k columns to request-local positions. + + Column ``j * SBK + o`` scores the token at + ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1. + """ + row = tl.program_id(0) + ks_base = tl.load(ks_ptr + row * ks_stride) // SBK + cols = tl.arange(0, K_POW2) + c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1) + valid = c >= 0 + cc = tl.where(valid, c, 0) + blocks = tl.load(si_ptr + row * si_stride + cc // SBK) + pos = (blocks - ks_base) * SBK + cc % SBK + tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k) + + +def pick_sparse_block_kv(candidate_block_size: int) -> int: + """Largest sparse-block size (tokens) dividing the candidate block size.""" + for sbk in SPARSE_BLOCK_KV_CHOICES: + if candidate_block_size % sbk == 0: + return sbk + raise ValueError( + f"candidate_block_size={candidate_block_size} is not a multiple of " + f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}." + ) + + +def candidate_blocks_to_sparse_indices( + candidate_blocks: torch.Tensor, + row_ks: torch.Tensor, + row_ke: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand request-local candidate blocks into DeepGEMM sparse block ids. + + Args: + candidate_blocks: [rows, K] int32 request-local candidate block ids + (-1 padded), in units of ``candidate_block_size`` positions. K + must be a power of two (the in-kernel sort); rows may be strided. + row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + (packed-workspace) coordinates the sparse kernel iterates over. + Pass zeros for the paged path, whose blocks are context-relative. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + out: Optional ``(sparse_indices, end)`` buffers to write into, + shaped [rows, K * ratio] and [rows], both int32. + + Returns: + sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM, + sorted ascending per row with the valid prefix first (unique in + production, where candidates are unique) and the remaining slots + repeating the last valid block (the kernel's padding convention). + With the unaligned-ks variant, block ``i`` starts at + ``i * sparse_block_kv + ks % sparse_block_kv``. + end: [rows] int32 count of valid sparse-token columns; valid columns + always form a prefix because the block ids are sorted. + """ + rows, num_candidates = candidate_blocks.shape + ratio = candidate_block_size // sparse_block_kv + assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), ( + candidate_block_size, + sparse_block_kv, + ) + assert num_candidates & (num_candidates - 1) == 0, num_candidates + num_sparse = num_candidates * ratio + if out is None: + device = candidate_blocks.device + out = ( + torch.empty(rows, num_sparse, dtype=torch.int32, device=device), + torch.empty(rows, dtype=torch.int32, device=device), + ) + sparse_indices, end = out + assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,) + assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1 + for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end): + assert t.dtype == torch.int32, t.dtype + + _expand_candidates_kernel[(rows,)]( + candidate_blocks, + candidate_blocks.stride(0), + row_ks, + row_ks.stride(0), + row_ke, + row_ke.stride(0), + sparse_indices, + sparse_indices.stride(0), + end, + end.stride(0), + K=num_candidates, + S=num_sparse, + RATIO=ratio, + SBK=sparse_block_kv, + num_warps=16, + ) + return sparse_indices, end + + +def sparse_topk_remap( + logits: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + row_ks: torch.Tensor, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + col_indices: torch.Tensor, +) -> None: + """Row top-k over sparse logits, remapped to request-local positions. + + Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at + request-local position + ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``. + + DeepSelect consumes the bf16 logits as they come out of the sparse + kernels; the valid columns form a prefix (block ids are sorted), so the + per-row ``end`` bounds drive it directly. Slots beyond a row's valid + count come back as -1. + + Args: + col_indices: [rows, topk_tokens] int32 scratch for the sparse-column + top-k result (row stride 32B-aligned for DeepSelect). + """ + rows, width = logits.shape + assert logits.dtype == torch.bfloat16 and width >= topk_tokens + assert end.dtype == torch.int32 and end.shape == (rows,) + cols = col_indices[:rows] + deep_select_topk(logits, topk_tokens, end=end, output_idx=cols) + _sparse_topk_remap_kernel[(rows,)]( + cols, + cols.stride(0), + sparse_indices, + sparse_indices.stride(0), + row_ks, + row_ks.stride(0), + topk_indices, + topk_indices.stride(0), + topk_tokens, + SBK=sparse_block_kv, + K_POW2=triton.next_power_of_2(topk_tokens), + ) + + +def sparse_mqa_logits_prefill_chunk( + q: torch.Tensor, + q_scale: torch.Tensor, + k_quant: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace). + + Args: + q: [rows, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, H] int32 packed UE8M0 Q scales. + k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8). + k_scale: [total_kv] int32 packed UE8M0 K scales. + weights: [rows, H] bf16 per-head weights; the sparse kernels take + bf16 and do not fold the Q scale in. + cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the + packed workspace. + candidate_blocks: [rows, K] int32 request-local candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + sparse_indices/end/col_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + kernel_metadata: DeepGEMM schedule from a previous call with the same + candidates and bounds (i.e. another indexer layer in the same + step). When given, the candidate expansion is skipped and + ``sparse_indices``/``end`` are assumed to be up to date. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + assert weights.dtype == torch.bfloat16, weights.dtype + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_sparse_mqa_logits_metadata( + cu_seqlen_ks, + cu_seqlen_ke, + k_quant.shape[0], + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_sparse_mqa_logits( + (q, q_scale), + (k_quant, k_scale), + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + cu_seqlen_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata + + +def sparse_mqa_logits_paged_decode( + q: torch.Tensor, + q_scale: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + row_indices: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + row_ks: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for paged decode rows. + + Every query is one row (spec-decode rows are flattened by the metadata + builder), so ``block_table`` is already per row. + + Args: + q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. + kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache + view; the page stride must be 512B-aligned. + weights: [rows, H] bf16 per-head weights. + context_lens: [rows] int32 per-row (compressed) context lengths. + block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. + row_indices: [rows] int32 row -> request map (pairing only affects + scheduling). + candidate_blocks: [rows, K] int32 candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + row_ks: [rows] int32 zeros (paged blocks are context-relative). + sparse_indices/end/col_indices: Caller-owned scratch. + kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + rows = q.shape[0] + page_kv = kv_cache.shape[1] + assert weights.dtype == torch.bfloat16, weights.dtype + assert kv_cache.stride(0) % 512 == 0, ( + "paged cache page stride must be 512B-aligned: " + f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" + ) + assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) + assert block_table.shape[0] == rows and block_table.stride(-1) == 1 + assert context_lens.shape == (rows,) and context_lens.is_contiguous() + + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + row_ks, + context_lens, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_paged_sparse_mqa_logits_metadata( + context_lens, + block_table, + row_indices, + page_kv, + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_paged_sparse_mqa_logits( + (q, q_scale), + kv_cache, + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + row_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py new file mode 100644 index 000000000000..78db1b3df5a4 --- /dev/null +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse Attention Indexer that scores only the candidate blocks. + +DeepSeek V4.1 two-level selection: the candidate-source indexer publishes +the top candidate blocks and later indexers pick their top-k inside them. +`SparseAttnIndexer` does that by computing dense logits over the whole +context and masking; this layer instead calls DeepGEMM's sparse MQA-logits +kernels on the candidate blocks only, so the work is O(candidate blocks) +instead of O(context). It requires the `DeepseekV41SparseIndexerBackend` +metadata (see ``AttentionConfig.indexer_sparse_logits``). +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + sparse_mqa_logits_paged_decode, + sparse_mqa_logits_prefill_chunk, +) +from vllm.model_executor.layers.sparse_attn_indexer import ( + _gather_workspace_shapes, + kv_cache_as_quant_view, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerMetadata, +) +from vllm.v1.worker.workspace import current_workspace_manager + + +def _prefill_k_workspaces( + total_seq_lens: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with + the dense indexer layers of the same model.""" + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True + ) + k_quant, k_scale = current_workspace_manager().get_simultaneous( + values_spec, scales_spec + ) + return k_quant, k_scale + + +def _gather_prefill_chunk_k( + kv_cache: torch.Tensor, + k_quant_full: torch.Tensor, + k_scale_full: torch.Tensor, + chunk: DeepseekV32IndexerPrefillChunkMetadata, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather one prefill chunk's paged K into the packed workspace.""" + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.local_cu_seq_lens, + ) + return k_quant, k_scale + + +class SparseMQAIndexer(nn.Module): + """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels. + + Only valid for indexer layers that read candidate blocks with the MXFP4 + indexer cache on SM100. The K cache is written by the model before this + runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the + attention layer can call either. + """ + + weights_dtype = torch.bfloat16 + """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant + kernel writes it directly so no cast runs per step.""" + + def __init__( + self, + k_cache, + topk_tokens: int, + head_dim: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + ): + super().__init__() + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_deep_gemm_sparse_mqa() + ): + raise ValueError( + "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8." + ) + if not has_deep_select(): + raise ValueError( + "SparseMQAIndexer requires the DeepSelect top-k extension " + "(vllm._deepselect_C)." + ) + self.k_cache = k_cache + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + + def _reserve_workspaces(self, device: torch.device) -> None: + """Profiling run: claim the K-gather workspace and the peak sparse + logits allocation so the memory estimate covers them.""" + _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) + + def forward( + self, + hidden_states: torch.Tensor, + q_quant: tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ) -> torch.Tensor: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + self._reserve_workspaces(hidden_states.device) + return self.topk_indices_buffer + + metadata = attn_metadata[self.k_cache.prefix] + assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), ( + "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata" + ) + assert k is None, "the model writes the indexer K cache" + q_values, q_scale = q_quant + kv_cache = self.k_cache.kv_cache + topk_tokens = self.topk_tokens + topk_indices_buffer = self.topk_indices_buffer + topk_indices_buffer[: hidden_states.shape[0]] = -1 + sparse_block_kv = metadata.sparse_block_kv + + if metadata.num_prefills > 0: + prefill = metadata.prefill + sparse_chunks = metadata.sparse_prefill + assert prefill is not None and sparse_chunks is not None + k_quant_full, k_scale_full = _prefill_k_workspaces( + self.max_total_seq_len, self.head_dim + ) + for chunk, rows in zip(prefill.chunks, sparse_chunks): + k_quant, k_scale = _gather_prefill_chunk_k( + kv_cache, k_quant_full, k_scale_full, chunk + ) + if chunk.local_total_seq_lens == 0: + continue # buffer already holds -1 + start, end = chunk.token_start, chunk.token_end + rows.kernel_metadata = sparse_mqa_logits_prefill_chunk( + q_values[start:end].view(torch.int8), + q_scale[start:end], + k_quant.view(torch.int8), + k_scale.view(torch.int32).squeeze(-1), + weights[start:end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + self.candidate_blocks[start:end], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[start:end, :topk_tokens], + sparse_indices=rows.sparse_indices, + end=rows.end, + col_indices=rows.col_indices, + kernel_metadata=rows.kernel_metadata, + ) + + if metadata.num_decodes > 0: + decode_rows = metadata.sparse_decode + assert decode_rows is not None and decode_rows.block_table is not None + assert decode_rows.row_indices is not None + num_rows = decode_rows.row_ke.shape[0] + if num_rows > 0: + decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode( + q_values[:num_rows].view(torch.int8).unsqueeze(1), + q_scale[:num_rows].unsqueeze(1), + kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True), + weights[:num_rows], + decode_rows.row_ke, + decode_rows.block_table, + decode_rows.row_indices, + self.candidate_blocks[:num_rows], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[:num_rows, :topk_tokens], + row_ks=decode_rows.row_ks, + sparse_indices=decode_rows.sparse_indices, + end=decode_rows.end, + col_indices=decode_rows.col_indices, + kernel_metadata=decode_rows.kernel_metadata, + ) + return topk_indices_buffer diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index cf6e7235f9ef..669186edbb1b 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -309,12 +309,22 @@ def __call__( ) +def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: + """Weights dtypes the model's indexer layers ask for: fp32 for the dense + scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer + (`SparseMQAIndexer`) is enabled.""" + if vllm_config.attention_config.indexer_sparse_logits: + return (torch.float32, torch.bfloat16) + return (torch.float32,) + + class FusedIndexerQRopeMxFp4TritonKernel( VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: dtype: torch.dtype + weights_out_dtype: torch.dtype num_heads: int index_q_half_rot_dim: int index_q_head_dim: int @@ -434,12 +444,14 @@ def dispatch( # type: ignore[override] self, *, dtype: torch.dtype, + weights_out_dtype: torch.dtype, num_heads: int, head_dim: int, rope_dim: int, ) -> CompileKey: return self.CompileKey( dtype=dtype, + weights_out_dtype=weights_out_dtype, num_heads=num_heads, index_q_half_rot_dim=rope_dim // 2, index_q_head_dim=head_dim, @@ -456,6 +468,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: return self._trace_dispatch(self.dispatch)( dtype=vllm_config.model_config.dtype, + weights_out_dtype=_indexer_weights_out_dtypes(vllm_config), num_heads=num_heads, head_dim=head_dim, rope_dim=rope_dim, @@ -491,7 +504,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: shape=(1, compile_key.num_heads, scale_head_dim), ), index_weights_out=TritonWarmupTensor( - torch.float32, + compile_key.weights_out_dtype, shape=(1, compile_key.num_heads), ), ) @@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + weights_out_dtype: torch.dtype = torch.float32, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, ]: """Fused RoPE + quantize Q for the sparse indexer. + ``weights_out_dtype`` is the dtype the downstream scoring kernel takes: + fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse + MQA-logits kernels (CUDA MXFP4 path only). + Weight-fold semantics (important — the two paths differ): FP8 path (use_fp4=False, default): @@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype) if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -613,6 +631,7 @@ def fused_indexer_q_rope_quant( weights_out=index_weights_out, ) elif current_platform.is_xpu(): + assert weights_out_dtype == torch.float32, weights_out_dtype torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( index_q, positions, @@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant( index_q_scale.view(torch.int32).squeeze(-1), ), index_weights_out + assert weights_out_dtype == torch.float32, weights_out_dtype fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 8052178a03d9..f67a77e2541f 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -165,6 +165,7 @@ class CompileKey: rope_dim: int num_heads: int cos_sin_dtype: type[cutlass.Numeric] + weights_out_dtype: type[cutlass.Numeric] coarsen: int @staticmethod @@ -177,6 +178,7 @@ def kernel(compile_key: CompileKey) -> Any: # process multiple heads at the same time to armotize RoPE load costs coarsen = compile_key.coarsen tb_size = IndexerQMxFp4Kernel.tb_size + weights_out_dtype = compile_key.weights_out_dtype assert num_heads % coarsen == 0 # later we will use 32B load = 16 BF16 elems @@ -287,7 +289,7 @@ def device_kernel( weight_head_id = global_tid % num_heads weights_out[weight_token_id, weight_head_id] = ( weights[weight_token_id, weight_head_id].to(Float32) * scale - ) + ).to(weights_out_dtype) @cute.jit def host_entrypoint( @@ -320,11 +322,13 @@ def dispatch( # type: ignore[override] self, *, cos_sin_dtype: type[cutlass.Numeric], + weights_out_dtype: type[cutlass.Numeric], **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, cos_sin_dtype=cos_sin_dtype, + weights_out_dtype=weights_out_dtype, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -340,11 +344,17 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: return [] + # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the + # dense scoring kernels take fp32. + weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,) + if vllm_config.attention_config.indexer_sparse_logits: + weights_out_dtypes = (Float32, BFloat16) return self._trace_dispatch(self.dispatch)( head_dim=head_dim, rope_dim=rope_dim, num_heads=num_heads, cos_sin_dtype=Float32, + weights_out_dtype=weights_out_dtypes, coarsen=(1, 4), ) @@ -380,7 +390,9 @@ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: divisibility=4, ) weights_out = make_fake_tensor( - Float32, (num_tokens, compile_key.num_heads), divisibility=4 + compile_key.weights_out_dtype, + (num_tokens, compile_key.num_heads), + divisibility=4, ) return ( positions, @@ -413,6 +425,7 @@ def __call__( rope_dim=cos_sin_cache.shape[-1], num_heads=num_heads, cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + weights_out_dtype=torch_to_cute_dtype(weights_out.dtype), coarsen=1 if num_tokens < 512 else 4, ) launch_args = ( diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index fa26351681c3..134328dbe908 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -4,6 +4,7 @@ DeepseekV4 MLA Attention Layer """ +import math from abc import ABC, abstractmethod from collections.abc import Callable from functools import cached_property @@ -25,6 +26,7 @@ RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, @@ -63,6 +65,9 @@ dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerBackend, +) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache from vllm.v1.kv_cache_interface import ( KVCacheSpec, @@ -997,7 +1002,9 @@ def __init__( self.cache_config = cache_config self.dtype = dtype self.compress_ratio = compress_ratio - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self @@ -1010,6 +1017,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + page_alignment = ( + 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 + ) return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, @@ -1020,16 +1030,21 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # (the block stride is the sum of every page in the group), so it # has to carry the same padding: the packed record's TMA stride, # else 512B for FlashInfer sparse (#44577). - alignment=( - 576 - if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() - else 512 + alignment=page_alignment, + # DeepGEMM's paged sparse MQA logits address pages by the block + # stride, which packs every layer's page in this layout, and need + # it 512B-aligned; the main KV record read from the same blocks + # needs its own TMA stride, so keep both. + block_stride_alignment=( + math.lcm(512, page_alignment) if self.sparse_logits else None ), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: + if self.sparse_logits: + return DeepseekV41SparseIndexerBackend return DeepseekV41IndexerBackend @@ -1129,21 +1144,52 @@ def __init__( self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.k_cache = k_cache - self.indexer_op = SparseAttnIndexer( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=True, - use_fp4_cache=self.use_fp4_kv, - compress_ratio=self.compress_ratio, - candidate_blocks=candidate_block_buffer, - candidate_block_size=candidate_block_size, - candidate_write=candidate_write, + # Candidate consumers can score only the candidate blocks (opt-in); + # the candidate source and pre-candidate indexers stay dense. + use_sparse_logits = ( + vllm_config.attention_config.indexer_sparse_logits + and candidate_block_buffer is not None + and not candidate_write + ) + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer + if use_sparse_logits: + if not self.use_fp4_kv: + raise ValueError( + "attention_config.indexer_sparse_logits requires the MXFP4 " + "indexer cache (attention_config.indexer_kv_dtype='mxfp4')." + ) + assert candidate_block_buffer is not None + assert topk_indices_buffer is not None + self.indexer_op = SparseMQAIndexer( + self.k_cache, + self.topk_tokens, + self.head_dim, + self.max_total_seq_len, + topk_indices_buffer, + candidate_block_buffer, + candidate_block_size, + ) + else: + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + compress_ratio=self.compress_ratio, + candidate_blocks=candidate_block_buffer, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + # The fused Q kernel writes the per-head weights in the dtype the + # scoring kernels take, so no cast runs per step. + self.indexer_weights_dtype = ( + SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32 ) def _produce_k( @@ -1231,6 +1277,7 @@ def forward( self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, ) if isinstance(q_quant, tuple): q, q_scale = q_quant diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 5d05c29a84c5..91e315afb9e7 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -158,6 +158,10 @@ def _missing(*_: Any, **__: Any) -> NoReturn: _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None @@ -232,6 +236,9 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl + global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl + global _get_sparse_mqa_logits_metadata_impl + global _get_paged_sparse_mqa_logits_metadata_impl global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl @@ -295,6 +302,17 @@ def _lazy_init() -> None: _get_paged_mqa_logits_metadata_impl = getattr( _dg, "get_paged_mqa_logits_metadata", None ) + # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only). + _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None) + _fp8_fp4_paged_sparse_mqa_logits_impl = getattr( + _dg, "fp8_fp4_paged_sparse_mqa_logits", None + ) + _get_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_sparse_mqa_logits_metadata", None + ) + _get_paged_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_paged_sparse_mqa_logits_metadata", None + ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( @@ -681,6 +699,171 @@ def fp8_fp4_paged_mqa_logits( ) +def has_deep_gemm_sparse_mqa() -> bool: + """Whether the installed DeepGEMM provides the sparse-indexer kernels + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" + _lazy_init() + return ( + _fp8_fp4_sparse_mqa_logits_impl is not None + and _fp8_fp4_paged_sparse_mqa_logits_impl is not None + and _get_sparse_mqa_logits_metadata_impl is not None + and _get_paged_sparse_mqa_logits_metadata_impl is not None + ) + + +def get_sparse_mqa_logits_metadata( + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + num_kv_tokens: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`. + + Always uses the unaligned-ks variant, which degenerates to the aligned + variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never + need a host-side sync to pick one. + + Args: + cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + workspace, shape [num_q_tokens], dtype int32. + num_kv_tokens: Total KV tokens in the packed workspace. + sparse_kv_block_indices: Per-row candidate block ids, shape + [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's + valid prefix must be sorted ascending (repeats are tolerated but + wasteful); pad by repeating the last valid block. Block ``i`` + covers tokens + ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``. + qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, + ``torch.float8_e4m3fn`` for FP8). + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_sparse_mqa_logits_metadata_impl( + cu_seqlen_ks, + cu_seqlen_ke, + num_kv_tokens, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def get_paged_sparse_mqa_logits_metadata( + context_lens: torch.Tensor, + block_table: torch.Tensor, + indices: torch.Tensor, + page_kv: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`. + + Args: + context_lens: Per-row context length, shape [num_q_tokens], int32. + block_table: Per-row page table, shape [num_q_tokens, max_pages], + int32; rows of queries paired with the same request must be + identical (expand it per query row before calling). + indices: Request index for each query row, shape [num_q_tokens]. + page_kv: Tokens per KV page; must be a multiple of sparse_block_kv. + sparse_kv_block_indices: Per-row logical (context-relative) candidate + block ids, [num_q_tokens, num_max_sparse_blocks], int32. + qk_dtype: dtype of the packed Q values. + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_paged_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_paged_sparse_mqa_logits_metadata_impl( + context_lens, + block_table, + indices, + page_kv, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + ) + + +def fp8_fp4_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (prefill path). + + Args: + q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed + as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H] + int32 (mandatory, unlike the dense kernels). + kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0 + int32 scales. + weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in). + metadata: From `get_sparse_mqa_logits_metadata`. + num_max_sparse_blocks: Width of each sparse-index row. + sparse_block_kv: Tokens per sparse block, 8 or 16. + + Returns: + bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; + column ``j * sparse_block_kv + o`` scores the token at + ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. + """ + _lazy_init() + if _fp8_fp4_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_sparse_mqa_logits_impl( + q, + kv, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def fp8_fp4_paged_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (paged decode path). + + Args: + q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. + kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, + page stride 512B-aligned. + weights: [num_q_tokens, H] ``torch.bfloat16``. + metadata: From `get_paged_sparse_mqa_logits_metadata`. + + Returns: + bf16 logits of shape + [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. + """ + _lazy_init() + if _fp8_fp4_paged_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_paged_sparse_mqa_logits_impl( + q, + kv_cache, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + ) + + def tf32_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 837dbca817c0..53206afd9337 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1116,6 +1116,11 @@ def _split_pcp_dcp_prefill_chunks( for request_slice, query_slice in chunk_specs ] + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + """Per-request KV lengths the prefill chunker budgets logits with; + subclasses whose logits rows are wider than the context override.""" + return seq_lens_cpu + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1278,7 +1283,7 @@ def build( ) else: chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], + self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, diff --git a/vllm/v1/attention/backends/mla/sparse_indexer.py b/vllm/v1/attention/backends/mla/sparse_indexer.py new file mode 100644 index 000000000000..c0b445534d22 --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_indexer.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path. + +Candidate-consuming indexer layers score only the candidate blocks published +by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead +of dense logits over the whole context. This backend extends the dense +indexer metadata with the per-row scratch that path needs, so the layer op +allocates nothing per step, and lets all consumer layers of a step share one +candidate expansion and DeepGEMM schedule (they all read the same candidate +blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``. +""" + +from dataclasses import dataclass, fields + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + check_deep_select_layout, + has_deep_select, + pick_sparse_block_kv, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, + dsa_indexer_uses_fp4, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +@dataclass +class SparseMQARowsMetadata: + """Per-row state for one sparse-logits call: the decode rows, or one + prefill chunk. Buffers are builder-owned views sized to the rows. + + ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first + consumer indexer layer of the step together with ``sparse_indices`` / + ``end``; later consumer layers reuse all three. + """ + + row_ks: torch.Tensor + """[rows] int32 K-range start per row; zeros for paged decode.""" + row_ke: torch.Tensor + """[rows] int32 K-range end (prefill) or compressed context length.""" + sparse_indices: torch.Tensor + """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids.""" + end: torch.Tensor + """[rows] int32 valid sparse-column count per row.""" + col_indices: torch.Tensor + """[rows, topk] int32 scratch for the sparse-column top-k.""" + block_table: torch.Tensor | None = None + """Decode only: [rows, pages] per-row page table.""" + row_indices: torch.Tensor | None = None + """Decode only: [rows] row -> request index.""" + kernel_metadata: torch.Tensor | None = None + + +@dataclass +class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata): + sparse_block_kv: int = 0 + sparse_decode: SparseMQARowsMetadata | None = None + sparse_prefill: list[SparseMQARowsMetadata] | None = None + """Parallel to ``prefill.chunks``.""" + + +class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_SPARSE_INDEXER" + + @staticmethod + def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]: + return DeepseekV41SparseIndexerMetadataBuilder + + +class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """Dense indexer metadata plus the sparse-logits row scratch. + + Every requirement of the sparse kernels is checked here, at engine start, + rather than falling back to the dense path per step. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + hf_config = vllm_config.model_config.hf_text_config + candidate_block_size = getattr(hf_config, "candidate_block_size", 0) + candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) + topk_tokens = getattr(hf_config, "index_topk", 0) + prefix = "attention_config.indexer_sparse_logits" + if candidate_block_size <= 0 or candidate_topk_blocks <= 0: + raise ValueError( + f"{prefix} requires a two-level (candidate block) indexer model " + "such as DeepSeek-V4.1." + ) + if not dsa_indexer_uses_fp4(vllm_config): + raise ValueError( + f"{prefix} requires the MXFP4 indexer cache " + "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels " + "take UE8M0-packed granularity-32 scales." + ) + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + raise ValueError(f"{prefix} requires an SM100-class GPU.") + if not has_deep_gemm_sparse_mqa(): + raise ValueError( + f"{prefix} requires DeepGEMM >= 2.8 " + "(fp8_fp4_sparse_mqa_logits not found)." + ) + if self.dcp_world_size > 1 or self.use_pcp: + raise NotImplementedError( + f"{prefix} is not supported with decode/prefill context parallel." + ) + if not self.supports_varlen: + raise ValueError( + f"{prefix} requires the varlen paged MQA logits decode layout." + ) + # The top-k runs on the kernels' bf16 logits with DeepSelect; there is + # no fp32 fallback, so the other top-k backends cannot be honored. + topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend + if not has_deep_select(): + raise ValueError( + f"{prefix} requires the DeepSelect top-k extension " + "(vllm._deepselect_C) for this GPU." + ) + if topk_backend not in ("auto", "deep_select"): + raise ValueError( + f"{prefix} uses DeepSelect for the top-k; " + f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' " + "is not supported with it." + ) + self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size) + if kv_cache_spec.block_size % self.sparse_block_kv != 0: + raise ValueError( + f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} " + f"tokens) must be a multiple of {self.sparse_block_kv} tokens." + ) + self.num_sparse_blocks = candidate_topk_blocks * ( + candidate_block_size // self.sparse_block_kv + ) + num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv + check_deep_select_layout(num_sparse_cols, topk_tokens) + # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill + # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every + # row were this wide. + self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4 + + # Rows are indexed by batch token (decode rows first, then prefill + # chunks by token range), like ``arange_buffer``. + max_rows = self.arange_buffer.shape[0] + int32 = dict(dtype=torch.int32, device=device) + self.sparse_indices_buffer = torch.zeros( + (max_rows, self.num_sparse_blocks), **int32 + ) + self.end_buffer = torch.zeros(max_rows, **int32) + self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32) + self.zero_row_ks = torch.zeros(max_rows, **int32) + logger.info_once( + "DeepSeek V4.1 indexer: candidate consumers score only the candidate " + "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d " + "tokens per row).", + self.num_sparse_blocks, + self.sparse_block_kv, + ) + + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + return seq_lens_cpu.clamp(min=self.min_split_seq_len) + + def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs): + return SparseMQARowsMetadata( + row_ks=row_ks, + row_ke=row_ke, + sparse_indices=self.sparse_indices_buffer[start:end], + end=self.end_buffer[start:end], + col_indices=self.col_indices_buffer[start:end], + **kwargs, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV41SparseIndexerMetadata: + base = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = DeepseekV41SparseIndexerMetadata( + **{f.name: getattr(base, f.name) for f in fields(base)}, + sparse_block_kv=self.sparse_block_kv, + ) + if base.decode is not None: + decode = base.decode + rows = decode.seq_lens.shape[0] + # The varlen layout flattens every query to one row, so the base + # decode tensors are already per row. + assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding + assert decode.block_table.shape[0] == rows + metadata.sparse_decode = self._rows( + 0, + rows, + row_ks=self.zero_row_ks[:rows], + row_ke=decode.seq_lens.view(-1), + block_table=decode.block_table, + row_indices=( + decode.indices + if decode.indices is not None + else self.arange_buffer[:rows] + ), + ) + if base.prefill is not None: + metadata.sparse_prefill = [ + self._rows( + chunk.token_start, + chunk.token_end, + row_ks=chunk.cu_seqlen_ks, + row_ke=chunk.cu_seqlen_ke, + ) + for chunk in base.prefill.chunks + ] + return metadata diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 28478b69f5e6..1c6e102c161c 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1608,6 +1608,15 @@ def _get_kv_cache_bytes_per_block( ] if hot_page_sizes: bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) + stride_alignments = [ + spec.block_stride_alignment + for group in kv_cache_groups + for layer_name in group.layer_names + if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec) + and spec.block_stride_alignment + ] + if stride_alignments: + bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments)) return bytes_per_block diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8cd867eb9c60..7e91300b2fe2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -653,6 +653,12 @@ class MLAAttentionSpec(FullAttentionSpec): is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" + block_stride_alignment: int | None = None + """Required alignment, in bytes, of the distance between consecutive + blocks of this cache. In block-major layouts that distance is the whole + block (all layers' pages), so the allocator rounds the block up to it. + DeepGEMM's paged sparse MQA-logits kernels address pages as + ``base + page * stride`` and need it 512B-aligned.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -674,6 +680,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) + block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 @@ -681,10 +688,11 @@ def merge(cls, specs: list[Self]) -> Self: and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 + and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, and storage block size." + "index-sharing role, storage block size and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -701,6 +709,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), + block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs ), From c6fa1f05d1bc627171c3a58c05f23a1c41202ac1 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Tue, 15 Sep 2026 17:15:59 -0400 Subject: [PATCH 275/774] [Perf][Kernel] Add sampled filtering for persistent top-k (#56346) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- .../kernels/benchmark_persistent_topk.py | 134 ++++++++++ csrc/libtorch_stable/persistent_topk.cuh | 86 +++--- csrc/libtorch_stable/sampled_topk.cuh | 253 ++++++++++++++++++ csrc/libtorch_stable/topk.cu | 18 +- tests/kernels/test_top_k_per_row.py | 100 +++++++ 5 files changed, 556 insertions(+), 35 deletions(-) create mode 100644 benchmarks/kernels/benchmark_persistent_topk.py create mode 100644 csrc/libtorch_stable/sampled_topk.cuh diff --git a/benchmarks/kernels/benchmark_persistent_topk.py b/benchmarks/kernels/benchmark_persistent_topk.py new file mode 100644 index 000000000000..a77a04cfd625 --- /dev/null +++ b/benchmarks/kernels/benchmark_persistent_topk.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows. + +Run on a reserved GPU, for example: + chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py + +Use --capacity to emulate graphs captured at a larger maximum context length. +Run the same command on the base and candidate builds to compare JSON results. +""" + +import argparse +import itertools +import json +import statistics +from pathlib import Path + +import torch +from flashinfer.testing import bench_gpu_time_with_cupti + +import vllm._custom_ops # noqa: F401 +from vllm.platforms import current_platform + + +def benchmark(rows: int, width: int, k: int, args) -> dict: + capacity = args.capacity or width + if capacity < width: + raise ValueError("Capacity must be at least the valid-length upper bound") + torch.manual_seed(args.seed) + logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if args.full_length: + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) + else: + lengths = torch.randint( + int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32 + ) + invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None] + logits.masked_fill_(invalid, float("nan")) + out = torch.empty((rows, k), device="cuda", dtype=torch.int32) + workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8) + cooperative = ( + args.backend == "auto" + and rows <= 64 + and capacity % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk + + def run(): + op(logits, lengths, out, workspace, k, capacity) + + run() + valid = torch.arange(k, device="cuda")[None] < lengths[:, None] + assert torch.all(out[~valid] == -1) + assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid) + sorted_indices = out.sort(dim=1).values + assert torch.all( + (sorted_indices[:, 1:] != sorted_indices[:, :-1]) + | (sorted_indices[:, 1:] == -1) + ) + selected = logits.gather(1, out.clamp_min(0).long()).masked_fill( + ~valid, -float("inf") + ) + reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0 + ) + trials = [ + statistics.median( + bench_gpu_time_with_cupti( + run, + use_cuda_graph=True, + cold_l2_cache=not args.warm, + dry_run_iters=5, + repeat_iters=args.repeat, + ) + ) + * 1000 + for _ in range(args.trials) + ] + us = statistics.median(trials) + return dict( + rows=rows, + width=width, + capacity=capacity, + k=k, + backend="cooperative" if cooperative else "persistent", + us=us, + trials_us=trials, + # Minimum useful traffic: one valid-score read and one index write. + effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024]) + parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840]) + parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048]) + parser.add_argument("--capacity", type=int) + parser.add_argument( + "--full-length", action="store_true", help="Use exactly --widths valid scores" + ) + parser.add_argument("--backend", choices=["auto", "persistent"], default="auto") + parser.add_argument( + "--warm", action="store_true", help="Keep L2 warm between replays" + ) + parser.add_argument("--repeat", type=int, default=60) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + metadata = dict( + gpu=str(torch.cuda.get_device_properties(0)), + torch=torch.__version__, + cuda=torch.version.cuda, + args=vars(args), + ) + print(json.dumps(metadata, default=str)) + results = [] + for rows, width, k in itertools.product(args.rows, args.widths, args.top_k): + result = benchmark(rows, width, k, args) + results.append(result) + print(json.dumps(result), flush=True) + if args.output: + args.output.write_text( + json.dumps(dict(metadata=metadata, results=results), indent=2, default=str) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 9ae0fd4bcbaa..bf5b46ab1875 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB -/*! - * \brief Filtered Top-K kernel for ragged sequences. - * - * \tparam DType Data type (float, half, nv_bfloat16) - * \tparam IdType Index type (int32_t) - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) - */ -template -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) - FilteredTopKUnifiedKernel(const DType* __restrict__ input, - IdType* __restrict__ output, - const IdType* __restrict__ lengths, - uint32_t num_rows, uint32_t top_k, - uint32_t max_len) { +template +struct FilteredTopKStorage { + alignas(128) int histogram[2][256 + 128]; + alignas(128) int counter; + alignas(128) int threshold_bin; + alignas(128) int num_input[2]; + alignas(128) int indices[MAX_K]; + int last_remain; +}; + +// With CheckOverflow, return false before using a truncated stash so the caller +// can retry with exact full-row selection. +template +__device__ __forceinline__ bool filtered_topk_row( + const DType* score, IdType* dst, int length, uint32_t top_k, + FilteredTopKStorage& storage) { constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; - - const uint32_t bid = blockIdx.x; const int tx = threadIdx.x; - if (bid >= num_rows) return; - - const int length = - (lengths != nullptr) ? lengths[bid] : static_cast(max_len); - const DType* score = input + bid * max_len; - IdType* dst = output + bid * top_k; - // Trivial case: length <= top_k if (length <= static_cast(top_k)) { for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { dst[i] = (i < length) ? static_cast(i) : static_cast(-1); } - return; + return true; } // Short path @@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) hist4096::histogram_4096_topk(score, dst, length, _smem_reg); } - return; + return true; } - // Static shared memory - alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; - alignas(128) __shared__ int s_counter; - alignas(128) __shared__ int s_threshold_bin_id; - alignas(128) __shared__ int s_num_input[2]; - alignas(128) __shared__ int s_indices[MAX_K]; - + auto& s_histogram_buf = storage.histogram; + auto& s_counter = storage.counter; + auto& s_threshold_bin_id = storage.threshold_bin; + auto& s_num_input = storage.num_input; + auto& s_indices = storage.indices; auto& s_histogram = s_histogram_buf[0]; // Dynamic shared memory for input double buffer @@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) // Stage 2: refine with 8bit radix passes #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { - __shared__ int s_last_remain; + auto& s_last_remain = storage.last_remain; const auto r_idx = round % 2; const auto _raw_num_input = s_num_input[r_idx]; + if constexpr (CheckOverflow) { + if (_raw_num_input > SMEM_INPUT_SIZE) return false; + } const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; @@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const int idx = s_indices[base]; dst[base] = static_cast(idx); } + return true; +} + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + const uint32_t bid = blockIdx.x; + if (bid >= num_rows) return; + const int length = lengths ? lengths[bid] : static_cast(max_len); + __shared__ FilteredTopKStorage storage; + filtered_topk_row( + input + bid * max_len, output + bid * top_k, length, top_k, storage); } // Helper to compute GCD for VEC_SIZE selection diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/libtorch_stable/sampled_topk.cuh new file mode 100644 index 000000000000..867abbd63542 --- /dev/null +++ b/csrc/libtorch_stable/sampled_topk.cuh @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect: +// https://github.com/deepseek-ai/DeepSelect +// A sample estimates a cutoff, then survivors are compacted in shared memory +// for exact FP32 selection. Too few or too many candidates trigger an exact +// full-row fallback. + +#pragma once + +#include +#include +#include + +#include "persistent_topk.cuh" + +namespace vllm::sampled_topk { + +constexpr int kThreads = 1024; +// Conservative crossover bounds from the B300 batch/length sweep. +template +constexpr int kMinSampledLength = K == 512 ? 98304 : 65536; +// Target half the buffer to leave room for sampling error. +constexpr int kSample = 4096; +constexpr int kCapacity = 8192; + +struct Storage { + uint2 candidates[kCapacity]; + int histogram[2048]; + typename cub::BlockScan::TempStorage scan; + int count; + int remaining; + int bin; + uint32_t prefix; +}; + +__device__ __forceinline__ uint32_t ordered(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ uint32_t coarse(uint32_t key) { + const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key; + return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits)); +} + +// Select the bin containing the rank and update the rank within that bin. +template +__device__ void select_bin(Storage& s) { + constexpr int kItems = (Bins + kThreads - 1) / kThreads; + int sums[kItems]; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + sums[j] = bin >= 0 ? s.histogram[bin] : 0; + } + cub::BlockScan(s.scan).InclusiveSum(sums, sums); + const int remaining = s.remaining; + __syncthreads(); +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + if (bin >= 0 && sums[j] >= remaining && + sums[j] - s.histogram[bin] < remaining) { + s.bin = bin; + s.remaining = remaining - (sums[j] - s.histogram[bin]); + } + } + __syncthreads(); +} + +// Find a coarse bin, then refine its full FP32 keys one byte at a time. +// Read shared candidates normally, or the full row for exact fallback. +template +__device__ uint32_t threshold(Storage& s, const float* row, int count, + int rank) { + const int tid = threadIdx.x; + if (tid == 0) { + s.remaining = rank; + s.prefix = 0; + } + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + atomicAdd(&s.histogram[coarse(key)], 1); + } + __syncthreads(); + select_bin<2048>(s); + const int coarse_bin = s.bin; + uint32_t mask = 0; + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) s.histogram[tid] = 0; + __syncthreads(); + const uint32_t prefix = s.prefix; + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + if (coarse(key) == coarse_bin && (key & mask) == prefix) { + atomicAdd(&s.histogram[(key >> shift) & 255], 1); + } + } + __syncthreads(); + select_bin<256>(s); + if (tid == 0) s.prefix = prefix | (static_cast(s.bin) << shift); + mask |= 255u << shift; + __syncthreads(); + } + return s.prefix; +} + +__device__ __forceinline__ int reserve(bool hit, int* counter) { + const uint32_t mask = __ballot_sync(0xffffffffu, hit); + if (mask == 0) return 0; + const int lane = threadIdx.x & 31; + int base = 0; + if (lane == 0) base = atomicAdd(counter, __popc(mask)); + base = __shfl_sync(0xffffffffu, base, 0); + return base + __popc(mask & ((1u << lane) - 1)); +} + +template +__device__ void emit(Storage& s, const float* row, int32_t* dst, int count, + uint32_t cutoff) { + const int tid = threadIdx.x; + if (tid == 0) s.count = 0; + __syncthreads(); + for (int base = 0; base < count; base += kThreads) { + const int i = base + tid; + const uint2 pair = i < count ? (Buffered ? s.candidates[i] + : make_uint2(i, ordered(row[i]))) + : make_uint2(0, 0); + bool keep = i < count && pair.y > cutoff; + if (i < count && pair.y == cutoff) { + keep = atomicSub(&s.remaining, 1) > 0; + } + const int offset = reserve(keep, &s.count); + if (keep) dst[offset] = pair.x; + } +} + +template +__global__ void __launch_bounds__(kThreads) + sampled_topk_kernel(const float* __restrict__ input, + const int32_t* __restrict__ lengths, + int32_t* __restrict__ output, int64_t stride, + int max_length) { + extern __shared__ __align__(16) unsigned char smem[]; + auto& s = *reinterpret_cast(smem); + const int tid = threadIdx.x; + const int length = max(0, min(lengths[blockIdx.x], max_length)); + const float* row = input + blockIdx.x * stride; + int32_t* dst = output + blockIdx.x * K; + // Graph allocations can hold rows much shorter than the host length bound. + if (length < kMinSampledLength) { + __shared__ filtered_topk::FilteredTopKStorage fallback_storage; + bool selected; + if ((stride & 3) == 0 && (reinterpret_cast(row) & 15) == 0) { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } else { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } + if (!selected) { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } + return; + } + if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length); + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + // Sample contiguous warps spread across the row to keep loads coalesced. + for (int i = tid; i < kSample; i += kThreads) { + const int index = (i / 32) * (length / (kSample / 32)) + i % 32; + atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1); + } + __syncthreads(); + select_bin<2048>(s); + const uint16_t half_key = s.bin << 5; + const uint16_t half_bits = + (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key; + const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits))); + if (tid == 0) s.count = 0; + __syncthreads(); + constexpr int kItems = 16; + // Compact survivors with one reservation per warp per tile. Keep counting + // beyond capacity so overflow triggers fallback instead of using a subset. + for (int base = 0; base < length; base += kItems * kThreads) { + uint32_t keys[kItems]; + uint32_t hit_mask = 0; + if ((reinterpret_cast(row) & 15) == 0 && + base + kItems * kThreads <= length) { +#pragma unroll + for (int j = 0; j < kItems / 4; ++j) { + const float4 values = + reinterpret_cast(row + base)[tid + j * kThreads]; + keys[4 * j] = ordered(values.x); + keys[4 * j + 1] = ordered(values.y); + keys[4 * j + 2] = ordered(values.z); + keys[4 * j + 3] = ordered(values.w); + } + } else { +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + keys[j] = i < length ? ordered(row[i]) : 0; + } + } +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j; + } + const int lane = tid & 31; + const int hits = __popc(hit_mask); + int inclusive = hits; +#pragma unroll + for (int delta = 1; delta < 32; delta *= 2) { + const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta); + if (lane >= delta) inclusive += previous; + } + const int total = __shfl_sync(0xffffffffu, inclusive, 31); + int warp_base = 0; + if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total); + warp_base = __shfl_sync(0xffffffffu, warp_base, 31); + int offset = warp_base + inclusive - hits; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + if (hit_mask & (1u << j)) { + if (offset < kCapacity) { + s.candidates[offset] = make_uint2( + base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]); + } + ++offset; + } + } + } + __syncthreads(); + const int count = s.count; + if (count >= K && count <= kCapacity) { + const uint32_t exact_cutoff = threshold(s, row, count, K); + emit(s, row, dst, count, exact_cutoff); + } else { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } +} + +} // namespace vllm::sampled_topk diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index b3e1a8623f3e..fb3d946b96b0 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -9,6 +9,7 @@ #ifndef USE_ROCM #include "persistent_topk.cuh" + #include "sampled_topk.cuh" #endif namespace { @@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, max_smem_per_block = device_prop->sharedMemPerBlockOptin; } - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + // Allow static fallback storage in addition to the 128 KiB dynamic buffer. + if (num_rows > 64 && + max_seq_len >= vllm::sampled_topk::kMinSampledLength && + max_smem_per_block >= 144 * 1024) { + auto kernel = vllm::sampled_topk::sampled_topk_kernel; + constexpr size_t smem_size = + vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; + cudaError_t status = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + STD_TORCH_CHECK(status == cudaSuccess, + "sampled_topk smem failed: ", cudaGetErrorString(status)); + kernel<<>>( + logits.const_data_ptr(), lengths.const_data_ptr(), + output.mutable_data_ptr(), stride, + static_cast(std::min(max_seq_len, logits.size(1)))); + } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { cudaError_t status = vllm::FilteredTopKRaggedTransform( logits.const_data_ptr(), output.mutable_data_ptr(), diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 591362020ee7..289539470631 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -932,6 +932,106 @@ def test_deepseek_workspace_topk( ) +@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+") +@pytest.mark.parametrize("rows", [65, 128]) +@pytest.mark.parametrize("top_k", [512, 1024, 2048]) +@pytest.mark.parametrize( + "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"] +) +@torch.inference_mode() +def test_persistent_topk_sampled_graph( + rows: int, top_k: int, distribution: str +) -> None: + """Sampling and both fallbacks preserve exact values across graph replays.""" + if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024: + pytest.skip("Sampled top-k requires at least 144 KiB of shared memory") + set_random_seed(42) + min_sampled_length = 98304 if top_k == 512 else 65536 + width = min_sampled_length + 3 + lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda") + logits = create_random_logits( + torch.zeros_like(lengths), + lengths, + torch.float32, + 42, + False, + "10LSBits" if distribution == "10LSBits" else "random", + ) + if distribution == "ascending": + logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32)) + elif distribution == "constant": + logits.fill_(1.0) + elif distribution == "sampled_peaks": + # A biased sample leaves fewer than k survivors and must fall back. + logits.zero_() + sample = torch.arange(4096, device="cuda") + positions = (sample // 32) * (width // 128) + sample % 32 + logits[:, positions] = 1 + sample.float() / 4096 + original = logits.clone() + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") + + def run() -> None: + torch.ops._C.persistent_topk( + logits, + lengths.view(-1, 4 if rows == 128 else 1), + indices, + workspace, + top_k, + width, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + positions = torch.arange(width, device="cuda") + slots = torch.arange(top_k, device="cuda") + for step in range(3): + bounds = torch.tensor( + [ + -1, + 0, + 1, + top_k - 1, + top_k, + 32769, + min_sampled_length - 1, + min_sampled_length, + width, + width + 17, + ], + dtype=torch.int32, + device="cuda", + ) + lengths.copy_( + bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()] + ) + logits.copy_(original if step != 1 else original.flip(1)) + logits.masked_fill_(positions[None] >= lengths[:, None], float("nan")) + indices.fill_(-2) + graph.replay() + valid = slots[None] < lengths.clamp(max=top_k)[:, None] + assert torch.all(indices[~valid] == -1) + assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid) + ordered_indices = indices.sort(dim=1).values + assert torch.all( + (ordered_indices[:, 1:] != ordered_indices[:, :-1]) + | (ordered_indices[:, 1:] == -1) + ) + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, -float("inf")) + expected = logits.masked_fill( + positions[None] >= lengths[:, None], -float("inf") + ) + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, + expected.topk(top_k, dim=1).values, + atol=0, + rtol=0, + ) + + def run_large_context_topk_test( batch_size: int, seq_lens: list[int], From 0fefffc9346644a01dc132e544b881ad646ea81b Mon Sep 17 00:00:00 2001 From: qizixi <22851944+zixi-qi@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:43:56 -0700 Subject: [PATCH 276/774] [KVConnector][NIXL] Support attention-HMA layouts in pipeline-parallel push prefill (#50494) Signed-off-by: zixi-qi Co-authored-by: Claude Co-authored-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: Codex --- docs/design/nixl_kv_push_connector.md | 44 ++ docs/features/nixl_connector_compatibility.md | 18 + .../unit/test_nixl_connector_hma.py | 1 + .../unit/test_nixl_desc_geometry.py | 32 +- .../unit/test_nixl_push_connector.py | 400 +++++++++++++++++- .../kv_connector/v1/nixl/base_worker.py | 256 +++++++++-- .../kv_connector/v1/nixl/metadata.py | 7 +- .../kv_connector/v1/nixl/push_worker.py | 4 +- 8 files changed, 713 insertions(+), 49 deletions(-) diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index 6f66be5eb9f5..a6debd7dba81 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing is D's own request id, taken from the registration), so the D-side accounting code is unchanged. +## Pipeline parallelism and hybrid KV caches + +The pull connector requires the local and remote workers to expose a +congruent list of KV regions — region *i* here corresponds to region +*i* there. That assumption breaks under pipeline parallelism (PP) +combined with a hybrid (HMA) KV layout: + +* a PP-sharded prefiller (**P**) holds only a slice of the model's + layers while the `PP=1` decoder (**D**) holds them all, so region + counts differ; +* with HMA, several layer names are pooled into one region and the layer + that represents a pooled region can differ between P and D. + +`NixlPushConnector` handles this for PP-sharded producers by routing +**by layer-name (member) identity** instead of by region index. Each +worker advertises which layer names back each of its NIXL regions in the +handshake metadata (`NixlAgentMetadata.region_members`). A producer that +needs member routing derives its member-major layout once, when it +registers its KV caches, and every transfer it issues uses that order. +`add_remote_agent` then selects exactly the remote regions this stage +owns and reorders them to match, so both sides stay paired regardless of +how each remote rank happens to order its metadata. + +Addresses, block lengths, strides, and per-region capacities follow the same +member order. Descriptor offsets use each member's region capacity, so P and +D need not allocate the same number of blocks. Physical allocations are still +registered once, even when multiple members share them. + +Invariants enforced when the remote regions are aligned: + +* every locally owned member must be advertised exactly once by the + remote; a missing member fails the handshake rather than silently + leaving that layer's KV stale, and remote-only members (owned by other + PP stages) are ignored; +* a remote that omits member metadata while the local layout requires + member routing fails the handshake instead of falling back to + region-index routing; +* member order is a property of the local layout alone, so the same local + source descriptors serve every remote engine and TP rank; only the + remote descriptor list is rebuilt per rank. + +Decode-side PP is unsupported because completions are counted per +consumer rank. Mamba/SSM hybrids are unsupported under PP. + ## Scheduler-side responsibilities `NixlPushConnectorScheduler` extends the base scheduler with: diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 031fba5b2f04..a41d5868b169 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting. - Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA. +### Pipeline parallelism + +The default (pull) NixlConnector does **not** support +`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA) +layout: region indices are not a stable identity across the +prefill/decode layer split, so the connector raises at startup. Use the +push connector (`NixlPushConnector`) for pipeline parallelism with +hybrid KV caches — it routes transfers by layer-name (member) identity, +so a PP-sharded prefiller can write into a `PP=1` decoder. See +[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md). + +Current push PP + HMA limitations: + +- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported. +- Hybrid SSM/Mamba layouts are not supported under PP. +- HMA requires the same block size on P and D. +- Attention-HMA member routing requires decode TP to be no greater than prefill TP. + ### Quantized KV cache [Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake. diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 60eb2d6614a9..ae2ddece7f77 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -918,6 +918,7 @@ def _make_mock_worker_for_desc_ids( worker = MagicMock(spec=NixlConnectorWorker) worker.num_regions = num_regions worker._uses_region_group_mapping = False + worker._transfer_layer_group_ids = () worker._has_mamba = has_mamba worker._group_spec_types = group_spec_types worker.block_len_per_layer = block_len_per_layer or [100] diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index bfec51da6901..a4bfa9ab83ae 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -12,6 +12,7 @@ """ from collections import defaultdict +from threading import Event, Lock from unittest.mock import MagicMock, patch import numpy as np @@ -112,6 +113,7 @@ def test_local_descriptors_follow_each_region_pool_capacity(): worker.block_len_per_layer = [16, 16] worker.block_stride_per_layer = [16, 16] worker.region_num_blocks = [2, 3] + worker._transfer_layer_region_indices = () descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1) @@ -119,7 +121,8 @@ def test_local_descriptors_follow_each_region_pool_capacity(): @pytest.mark.cpu_test -def test_overlaid_transfer_groups_share_region_geometry(): +@pytest.mark.parametrize("push_pp", [False, True]) +def test_overlaid_transfer_groups_share_region_geometry(push_pp): """Groups overlaid on one allocation share its transfer region.""" import msgspec @@ -127,6 +130,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlAgentMetadata, ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -153,7 +159,14 @@ def test_overlaid_transfer_groups_share_region_geometry(): } groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches] - worker = object.__new__(NixlConnectorWorker) + worker_cls = NixlPushConnectorWorker if push_pp else NixlConnectorWorker + worker = object.__new__(worker_cls) + if push_pp: + worker._push_writer_stop = Event() + worker._push_writer_wake = Event() + worker._push_writer_thread = MagicMock() + worker._sending_transfers_lock = Lock() + worker._sending_transfers = defaultdict(list) worker.tp_rank = 0 worker.world_size = 1 worker.transfer_tp_rank = 0 @@ -198,6 +211,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker._desc_is_dram_by_block_size = {} worker._desc_pos_by_block_size = {} worker._dram_src_handles_by_block_size = {} + worker._transfer_layer_names = () + worker._transfer_layer_region_indices = () + worker._transfer_layer_group_ids = () worker._region_is_mla = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] @@ -205,7 +221,8 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker.use_host_buffer = False worker.host_xfer_buffers = {} worker.device_kv_caches = {} - worker.pp_size = 1 + worker.pp_size = 2 if push_pp else 1 + worker._is_hma_required = True worker.dcp_size = 1 worker.pcp_size = 1 worker.kv_buffer_device = "cuda" @@ -240,7 +257,13 @@ def test_overlaid_transfer_groups_share_region_geometry(): expected_addrs = [ backing.data_ptr() + block * block_stride for block in range(num_blocks) ] - assert worker.src_blocks_data[:, 0].tolist() == expected_addrs + num_desc_regions = 2 if push_pp else 1 + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions + assert worker.num_descs == num_blocks * num_desc_regions + assert ( + worker.dst_region_num_blocks[worker.engine_id] + == [num_blocks] * num_desc_regions + ) metadata = msgspec.msgpack.decode( worker.xfer_handshake_metadata.agent_metadata_bytes, @@ -248,6 +271,7 @@ def test_overlaid_transfer_groups_share_region_geometry(): ) assert metadata.region_group_ids == [-1] assert metadata.region_num_blocks == [num_blocks] + assert metadata.region_members == ([["layer.0", "layer.1"]] if push_pp else []) assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]] diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 913ed0dcba0d..a908d3246261 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -32,9 +32,13 @@ import msgspec import pytest +from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorTransferResults, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( PUSH_REG_NOTIF_PREFIX, NixlAgentMetadata, @@ -43,6 +47,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( NixlPushConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) @@ -354,6 +359,7 @@ def fresh(cls) -> _StubWriterWorker: w._failed_recv_reqs = queue.Queue() w._recv_failures = set() w._recving_transfers = defaultdict(list) + w._is_hma_required = False w._reqs_to_process = set() w._reqs_to_send = {} w.consumer_notification_counts_by_req = defaultdict(int) @@ -361,12 +367,16 @@ def fresh(cls) -> _StubWriterWorker: w.pcp_rank = 0 w.pcp_dcp_sharded = False w.world_size = 1 + w.pp_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False w.region_group_ids = [0] + w._transfer_layer_names = () + w._transfer_layer_region_indices = () + w._transfer_layer_group_ids = () w.dst_region_num_blocks = {} w.dst_region_group_ids = {} w.dst_uses_region_group_mapping = {} @@ -922,12 +932,20 @@ def test_failed_push_send_is_not_reported_done_sending(self): # Both handles were still cleaned up. assert request_id not in w._sending_transfers - def test_completed_push_send_waits_for_consumer_notif(self): + @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)]) + def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma): """P-side sends are completed by consumer notifs / lease expiry, not by local WRITE-handle completion.""" w = self._pollable_worker() + w.pp_size = pp_size + w._is_hma_required = is_hma request_id = self._make_sending_req(w) - w.nixl_wrapper.check_xfer_state.return_value = "DONE" + w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"] + + results = w.get_transfer_results() + assert request_id not in results.finished_sending + assert w._sending_transfers[request_id] == [102] + assert request_id in w._reqs_to_send results = w.get_transfer_results() @@ -940,6 +958,12 @@ def test_completed_push_send_waits_for_consumer_notif(self): results = w.get_transfer_results() assert request_id in results.finished_sending assert request_id not in w._reqs_to_send + assert request_id not in w._reqs_to_process + assert w._evict_finished_inbox.get_nowait() == request_id + assert w.get_transfer_results().finished_sending == set() + assert w._evict_finished_inbox.empty() + assert w.nixl_wrapper.release_xfer_handle.call_count == 2 + w.xfer_stats.record_kv_expired_req.assert_not_called() # ----------------------------------------------------------------- # @@ -1576,3 +1600,375 @@ def test_no_prefix_hit_leaves_blocks_untrimmed(self): local, remote = self._written_block_ids(w) assert local == [10, 11, 12] assert remote == [500, 501, 502] + + @pytest.mark.parametrize("layer_name_routing", [False, True]) + def test_different_region_groups_require_layer_routing(self, layer_name_routing): + """A layer mapped to a pooled remote region can still receive its blocks.""" + w, engine_id = self._worker_driving_xfer() + w.dst_region_group_ids[engine_id] = [-1] + if layer_name_routing: + w._transfer_layer_group_ids = (0,) + reg = _registration_data("req-groups", local_block_ids=([500, 501],)) + + if not layer_name_routing: + with pytest.raises(NotImplementedError, match="cache-group layouts"): + NixlPushConnectorWorker._do_start_push_kv( + w, "req-groups", ([10, 11],), reg + ) + w.nixl_wrapper.make_prepped_xfer.assert_not_called() + return + + NixlPushConnectorWorker._do_start_push_kv(w, "req-groups", ([10, 11],), reg) + + local, remote = self._written_block_ids(w) + assert local == [10, 11] + assert remote == [500, 501] + + +def _agent_metadata( + region_layers: list[list[str]], + base_addresses: list[int], + block_lens: list[int], + block_strides: list[int] | None = None, +) -> NixlAgentMetadata: + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"agent", + kv_caches_base_addr=base_addresses, + device_id=7, + num_blocks=2, + block_lens=block_lens, + block_strides=list(block_lens if block_strides is None else block_strides), + kv_cache_layout="LBHNC", + block_size=16, + ssm_sizes=(0, 0), + attn_backend_name="FLASH_ATTN", + physical_blocks_per_logical_kv_block=1, + region_members=region_layers, + ) + + +def _layer_routing_worker( + region_layers: list[list[str]], + group_by_layer: dict[str, int], + pp_size: int = 2, + is_hma: bool = True, +) -> _StubWriterWorker: + worker = _StubWriterWorker.fresh() + worker.pp_size = pp_size + worker.dcp_size = 1 + worker.block_size = 16 + worker._has_mamba = False + worker._is_hma_required = is_hma + worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer) + worker.transfer_topo = MagicMock() + worker._set_region_layers(region_layers) + return worker + + +@pytest.mark.parametrize( + ("region_num_blocks", "expected"), + [(None, [1, 2, 15, 21, 22]), ([4, 7, 5], [1, 2, 9, 12, 13])], +) +def test_layer_group_ids_route_descriptor_blocks(region_num_blocks, expected): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + desc_ids = worker._compute_desc_ids( + block_ids=[[1, 2], [5]], + dst_num_blocks=10, + block_size_ratio=None, + physical_blocks_per_logical=1, + region_num_blocks=region_num_blocks, + ) + assert desc_ids.tolist() == expected + + +def test_layer_metadata_round_trip(): + metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256]) + metadata.region_num_blocks = [4] + metadata.region_group_ids = [-1] + metadata.region_names = ["L0"] + metadata.region_mem_types = ["VRAM"] + + encoded = msgspec.msgpack.encode(metadata) + assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata + + +def test_layer_identity_gate_preserves_the_non_hma_path(): + assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == ( + 0, + ) + + # A base worker (pull) never routes by layer name. + pull = object.__new__(NixlBaseConnectorWorker) + pull._transfer_layer_names = () + pull._transfer_layer_region_indices = () + pull._transfer_layer_group_ids = () + pull._set_region_layers([["a"]]) + assert pull._transfer_layer_region_indices == () + + # A non-HMA local layout does not require layer-name routing. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, is_hma=False + )._transfer_layer_region_indices + == () + ) + + # PP=1 has congruent local/remote regions and keeps the region route even + # when its allocator is hybrid. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, pp_size=1 + )._transfer_layer_region_indices + == () + ) + + +def test_layer_alignment_fails_loud_when_remote_omits_layers(): + # Falling back to region-index routing here would silently transfer stale + # KV, so an unannounced peer must fail the handshake instead. + worker = _layer_routing_worker([["a"]], {"a": 0}) + with pytest.raises(AssertionError, match="no region_members"): + worker._align_remote_regions_by_layer(_agent_metadata([], [0xA000], [128])) + + +def test_layer_alignment_expands_pooled_regions(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + assert worker._transfer_layer_names == ("a", "a.swa", "b") + assert worker._transfer_layer_region_indices == (0, 0, 1) + assert worker._transfer_layer_group_ids == (0, 1, 0) + + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + worker._align_remote_regions_by_layer(metadata) + + assert metadata.kv_caches_base_addr == [0xA000, 0xA000, 0xB000] + assert metadata.block_lens == [128, 128, 128] + assert metadata.block_strides == [128, 128, 128] + assert metadata.region_members == [["a"], ["a.swa"], ["b"]] + assert metadata.region_num_blocks == [4, 4, 6] + assert metadata.region_group_ids == [-1, -1, 0] + assert metadata.region_names == ["a", "a.swa", "b"] + assert metadata.region_mem_types == ["VRAM", "VRAM", "VRAM"] + + +def test_layer_alignment_filters_and_reorders_a_pp_stage(): + worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1}) + metadata = _agent_metadata( + [["l2"], ["l0"], ["l3"], ["l1"]], + [0xC000, 0xA000, 0xD000, 0xB000], + [65536, 65536, 32768, 32768], + [131072, 131072, 65536, 65536], + ) + metadata.region_num_blocks = [4, 6, 8, 10] + metadata.region_group_ids = [0, 0, 1, 1] + metadata.region_names = ["l2", "l0", "l3", "l1"] + metadata.region_mem_types = ["VRAM"] * 4 + + worker._align_remote_regions_by_layer(metadata) + + assert worker._transfer_layer_region_indices == (0, 1) + assert worker._transfer_layer_group_ids == (0, 1) + assert metadata.kv_caches_base_addr == [0xC000, 0xD000] + assert metadata.block_lens == [65536, 32768] + assert metadata.block_strides == [131072, 65536] + assert metadata.region_num_blocks == [4, 8] + assert metadata.region_group_ids == [0, 1] + assert metadata.region_names == ["l2", "l3"] + assert metadata.region_mem_types == ["VRAM", "VRAM"] + + +@pytest.mark.parametrize( + ("local_num_blocks", "remote_num_blocks"), + [([4, 4], [4, 4, 4]), ([4, 6], [3, 5, 7])], +) +def test_layer_descriptors_pair_layers_across_asymmetric_pp_split( + local_num_blocks, remote_num_blocks +): + """A PP split can leave each stage a different mix of attention types. + + Stage 1 owns L3 (full) pooled with L4 (sliding) in region 0, and L5 (full) + alone in region 1. The consumer holds every layer and pools them + differently, so region indices cannot be paired positionally. + """ + worker = _layer_routing_worker([["L3", "L4"], ["L5"]], {"L3": 0, "L4": 1, "L5": 0}) + worker.block_len_per_layer = [128, 128] + worker.block_stride_per_layer = [256, 512] + worker._region_is_mla = [False, False] + worker.num_blocks = 4 + worker.region_num_blocks = local_num_blocks + worker.device_id = 0 + + consumer = _agent_metadata( + [["L0", "L1"], ["L3", "L2"], ["L5", "L4"]], + [0xA000, 0xB000, 0xC000], + [128, 128, 128], + [256, 512, 1024], + ) + consumer.num_blocks = 4 + consumer.region_num_blocks = remote_num_blocks + worker._align_remote_regions_by_layer(consumer) + + # L3 -> remote region 1, L4 -> 2, L5 -> 2. Pairing by index would have sent + # L5 to region 1 and L4 to region 2's sibling. + assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000] + + plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0) + local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1) + remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1) + local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices] + local_ids = worker._compute_desc_ids( + [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts + ) + remote_ids = worker._compute_desc_ids( + [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks + ) + assert len(local_descs) == sum(local_counts) + assert len(remote_descs) == sum(consumer.region_num_blocks) + + # Each pair addresses the same layer's blocks despite different pooling + # and strides. L3 and L5 use group 0; L4 uses group 1. + assert local_descs[local_ids].tolist() == [ + [0x1100, 128, 0], + [0x1200, 128, 0], + [0x1300, 128, 0], + [0x2200, 128, 0], + [0x2400, 128, 0], + ] + assert remote_descs[remote_ids].tolist() == [ + [0xB000, 128, 7], + [0xB200, 128, 7], + [0xC800, 128, 7], + [0xC000, 128, 7], + [0xC400, 128, 7], + ] + + +def test_layer_alignment_is_canonical_across_remote_orderings(): + local, groups = [["x"], ["y"]], {"x": 0, "y": 1} + rank0 = _agent_metadata([["x"], ["y"]], [0x1000, 0x2000], [64, 128], [256, 512]) + rank1 = _agent_metadata([["y"], ["x"]], [0x2000, 0x1000], [128, 64], [512, 256]) + + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank0) + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank1) + + assert rank0.kv_caches_base_addr == rank1.kv_caches_base_addr == [0x1000, 0x2000] + assert rank0.block_lens == rank1.block_lens == [64, 128] + assert rank0.block_strides == rank1.block_strides == [256, 512] + + +def test_layer_alignment_is_idempotent(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + + worker._align_remote_regions_by_layer(metadata) + aligned = msgspec.msgpack.encode(metadata) + worker._align_remote_regions_by_layer(metadata) + + assert msgspec.msgpack.encode(metadata) == aligned + + +def test_layer_alignment_rejects_missing_local_layer(): + worker = _layer_routing_worker([["l0"], ["l1"]], {"l0": 0, "l1": 1}) + metadata = _agent_metadata([["l0"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="missing locally owned layers"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_duplicate_remote_layer(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["a"]], [0xA000, 0xB000], [128, 128]) + + with pytest.raises(AssertionError, match="in multiple regions"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_inconsistent_remote_metadata(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + + +@pytest.mark.parametrize( + "field", + ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"], +) +def test_layer_alignment_rejects_inconsistent_region_geometry(field): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000, 0xB000], [128, 128]) + setattr(metadata, field, []) + original = msgspec.msgpack.encode(metadata) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + assert msgspec.msgpack.encode(metadata) == original + + +def test_set_region_layers_rejects_duplicate_local_layer(): + with pytest.raises(AssertionError, match="spans multiple NIXL regions"): + _layer_routing_worker([["a"], ["a"]], {"a": 0}) + + +def test_set_region_layers_rejects_layer_outside_any_kv_group(): + with pytest.raises(AssertionError, match="outside any local group"): + _layer_routing_worker([["a"], ["b"]], {"a": 0}) + + +@pytest.mark.parametrize( + ("local_block_size", "remote_block_size", "remote_tp_size", "error"), + [ + (32, 16, 1, "identical P/D block sizes"), + (16, 32, 1, "identical P/D block sizes"), + (16, 16, 2, "decode TP greater"), + ], +) +def test_layer_handshake_rejects_unsupported_geometry( + local_block_size: int, remote_block_size: int, remote_tp_size: int, error: str +): + """Reject unsupported peers without registering agents or transfer state.""" + metadata = _agent_metadata([["a"]], [0xA000], [128]) + metadata.block_size = remote_block_size + worker = _layer_routing_worker([["a"]], {"a": 0}) + worker.block_size = local_block_size + worker.block_len_per_layer = [128] + worker.use_mla = False + worker.nixl_wrapper = MagicMock() + worker.kv_caches_base_addr = defaultdict(dict) + worker.dst_num_blocks = {} + worker.tp_mappings = {} + worker.transfer_topo = TransferTopology( + tp_rank=0, + tp_size=1, + block_size=local_block_size, + engine_id=worker.engine_id, + is_mla=False, + is_mamba=False, + total_num_kv_heads=8, + attn_backends=[], + ) + + with pytest.raises(NotImplementedError, match=error): + worker.add_remote_agent(metadata, remote_tp_size=remote_tp_size) + worker.nixl_wrapper.add_remote_agent.assert_not_called() + worker.nixl_wrapper.prep_xfer_dlist.assert_not_called() + assert not worker.tp_mappings + assert not worker.dst_num_blocks + assert not worker.kv_caches_base_addr + with pytest.raises(KeyError): + worker.transfer_topo.get_engine_info(metadata.engine_id) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 6f8f1e55dc6b..2704de589f05 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -164,6 +164,10 @@ class NixlBaseConnectorWorker: # Overridden by NixlPushConnectorWorker. _TRANSFER_MODE: str = "pull" + # Layer-name routing is supported only by NixlPushConnector for HMA KV + # caches under PP. + _supports_pp_hma = False + def _compute_desc_ids( self, block_ids: BlockIds, @@ -191,28 +195,42 @@ def _compute_desc_ids( num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - num_regions = self.num_regions or len(self.block_len_per_layer) + use_layer_name_routing = bool(self._transfer_layer_group_ids) + num_desc_regions = ( + len(self._transfer_layer_group_ids) + if use_layer_name_routing + else self.num_regions or len(self.block_len_per_layer) + ) if region_num_blocks is None: - region_num_blocks = [num_blocks] * num_regions + region_num_blocks = [num_blocks] * num_desc_regions elif block_size_ratio is not None: region_num_blocks = [ int(count * block_size_ratio) for count in region_num_blocks ] - if region_num_blocks: - num_regions = len(region_num_blocks) - assert len(region_num_blocks) == num_regions + if region_num_blocks and not use_layer_name_routing: + num_desc_regions = len(region_num_blocks) + assert len(region_num_blocks) == num_desc_regions region_offsets = np.cumsum([0, *region_num_blocks[:-1]]) num_fa_descs = sum(region_num_blocks) + if use_layer_name_routing: + # Descriptors in layer order, using each layer's KV-cache group. + return np.concatenate( + [ + np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k] + for k, g in enumerate(self._transfer_layer_group_ids) + ] + ) + # All-attention fast path. if num_ssm_regions == 0: if region_group_ids is None: region_group_ids = self.region_group_ids if uses_region_group_mapping is None: uses_region_group_mapping = self._uses_region_group_mapping - if not region_group_ids and num_regions == 1: + if not region_group_ids and num_desc_regions == 1: region_group_ids = [0] - assert len(region_group_ids) == num_regions + assert len(region_group_ids) == num_desc_regions if uses_region_group_mapping is None: uses_region_group_mapping = len(set(region_group_ids)) > 1 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32) @@ -392,6 +410,103 @@ def _is_region_replicated(self, region_idx: int) -> bool: """ return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + def _set_region_layers(self, region_layers: list[list[str]]) -> None: + layer_names = [layer_name for region in region_layers for layer_name in region] + assert len(layer_names) == len(set(layer_names)), ( + "A KV cache layer spans multiple NIXL regions" + ) + self.region_members: list[list[str]] = region_layers + if not self._requires_layer_name_routing(): + return + group_by_layer = self.kv_cache_config.transfer_group_index_by_layer + ungrouped = [name for name in layer_names if name not in group_by_layer] + assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}" + self._transfer_layer_names = tuple(layer_names) + self._transfer_layer_region_indices = tuple( + region_idx + for region_idx, region in enumerate(region_layers) + for _ in region + ) + self._transfer_layer_group_ids = tuple( + group_by_layer[name] for name in layer_names + ) + + def _requires_layer_name_routing(self) -> bool: + """Whether PP push must match HMA layers by name, not region index.""" + return ( + self._supports_pp_hma + and self.pp_size > 1 + and self._is_hma_required + and not self._has_mamba + ) + + def _align_remote_regions_by_layer( + self, nixl_agent_meta: NixlAgentMetadata + ) -> None: + """Select remote regions for this PP stage in local layer order.""" + remote_region_layers = nixl_agent_meta.region_members + # Region-index routing here would silently transfer stale KV. + assert remote_region_layers, "Remote advertised no region_members" + assert ( + len(nixl_agent_meta.kv_caches_base_addr) + == len(nixl_agent_meta.block_lens) + == len(nixl_agent_meta.block_strides) + == len(remote_region_layers) + ), "Remote region metadata lengths disagree" + assert all( + values is None or len(values) == len(remote_region_layers) + for values in ( + nixl_agent_meta.region_num_blocks, + nixl_agent_meta.region_group_ids, + nixl_agent_meta.region_names, + nixl_agent_meta.region_mem_types, + ) + ), "Remote region metadata lengths disagree" + + remote_region_by_layer: dict[str, int] = {} + for region_idx, layer_names in enumerate(remote_region_layers): + for layer_name in layer_names: + assert layer_name not in remote_region_by_layer, ( + f"Remote advertised layer {layer_name!r} in multiple regions" + ) + remote_region_by_layer[layer_name] = region_idx + + missing = [ + name + for name in self._transfer_layer_names + if name not in remote_region_by_layer + ] + assert not missing, f"Remote is missing locally owned layers: {missing}" + + remote_regions = [ + remote_region_by_layer[name] for name in self._transfer_layer_names + ] + nixl_agent_meta.kv_caches_base_addr = [ + nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions + ] + nixl_agent_meta.block_lens = [ + nixl_agent_meta.block_lens[i] for i in remote_regions + ] + nixl_agent_meta.block_strides = [ + nixl_agent_meta.block_strides[i] for i in remote_regions + ] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = [ + nixl_agent_meta.region_num_blocks[i] for i in remote_regions + ] + if nixl_agent_meta.region_group_ids is not None: + nixl_agent_meta.region_group_ids = [ + nixl_agent_meta.region_group_ids[i] for i in remote_regions + ] + if nixl_agent_meta.region_mem_types is not None: + nixl_agent_meta.region_mem_types = [ + nixl_agent_meta.region_mem_types[i] for i in remote_regions + ] + if nixl_agent_meta.region_names is not None: + nixl_agent_meta.region_names = list(self._transfer_layer_names) + # One layer per region keeps a second alignment pass a no-op. + nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names] + def __init__( self, vllm_config: "VllmConfig", @@ -638,18 +753,26 @@ def __init__( # transfers into the matching sub-range of a PP=1 remote's regions. self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self._remote_region_offset = 0 - # PP push slices regions per layer (uniform count); HMA breaks that. - if self.pp_size > 1 and self._is_hma_required: - raise NotImplementedError( - "NixlPushConnector does not support pipeline_parallel_size > 1 " - "with hybrid KV cache layouts (HMA) yet." - ) - # Decode-side PP is unsupported (completions counted per consumer rank). - if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1: - raise NotImplementedError( - "NixlPushConnector consumer (decode) does not support " - "pipeline_parallel_size > 1." - ) + if self.pp_size > 1: + if self._is_hma_required and not self._supports_pp_hma: + raise NotImplementedError( + "NixlConnector (pull) does not support pipeline_parallel_size " + "> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector " + "for PP + HMA." + ) + # PP push routes HMA (hybrid) attention layouts by layer name; + # Mamba/SSM hybrids are not yet supported under PP. + if self._has_mamba: + raise NotImplementedError( + "NixlPushConnector does not support pipeline_parallel_size > 1 " + "with Mamba/SSM hybrid KV cache layouts yet." + ) + # Decode-side PP is unsupported (completions counted per consumer rank). + if vllm_config.kv_transfer_config.kv_role == "kv_consumer": + raise NotImplementedError( + "NixlPushConnector consumer (decode) does not support " + "pipeline_parallel_size > 1." + ) # Keep heartbeat handshakes to a PP-sharded producer notif-only. self._hb_handshake_notif_only = False @@ -773,6 +896,13 @@ def __init__( # within a block (BLHNC/BHLNC), where stride > block_len. self.block_stride_per_layer = list[int]() + # Region layer membership is populated for HMA layouts. + self.region_members = [] + # Local layer order, set by ``_set_region_layers`` only when + # routing by layer name. Empty means region-index routing. + self._transfer_layer_names = tuple[str, ...]() + self._transfer_layer_region_indices = tuple[int, ...]() + self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} @@ -1306,6 +1436,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ): compressed_region_owners.setdefault(cache.data_ptr(), cache) + track_region_layers = ( + self._supports_pp_hma and self._is_hma_required and not self._has_mamba + ) + region_layers: list[list[str]] = [] + # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). @@ -1508,6 +1643,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._region_is_mla.append(is_mla_region) region_mem_types.append(mem_type) + if track_region_layers: + if region_index == len(region_layers): + region_layers.append([layer_name]) + else: + region_layers[region_index].append(layer_name) + if isinstance(layer_spec, MambaSpec): if layer_spec.tp_replicated: assert self._is_csa_linear @@ -1563,7 +1704,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): if self._has_mamba: self.region_num_blocks = [self.num_blocks] * self.num_regions - if self.pp_size > 1: + self._set_region_layers(region_layers) + + if self.pp_size > 1 and not self._is_hma_required: start_layer, end_layer = self.model_config.get_layers_start_end_indices( self.vllm_config.parallel_config ) @@ -1573,7 +1716,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._remote_region_offset = regions_per_layer * start_layer # Total local FA descriptors (boundary between FA and mamba descs). - self.num_descs = sum(self.region_num_blocks) + xfer_region_num_blocks = ( + [self.region_num_blocks[i] for i in self._transfer_layer_region_indices] + if self._transfer_layer_region_indices + else self.region_num_blocks + ) + self.num_descs = sum(xfer_region_num_blocks) self._mixed_mem_types = len(set(region_mem_types)) > 1 if self._mixed_mem_types: @@ -1597,7 +1745,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.device_kv_caches = kv_caches self.dst_num_blocks[self.engine_id] = self.num_blocks - self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks + self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks self.dst_region_group_ids[self.engine_id] = self.region_group_ids self.dst_uses_region_group_mapping[self.engine_id] = ( self._uses_region_group_mapping @@ -1647,6 +1795,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_mem_types=self.region_mem_types, dcp_size=self.dcp_size, pcp_size=self.pcp_size, + region_members=self.region_members, ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -1805,12 +1954,16 @@ def _build_fa_local( base_addresses: list[int], block_size_ratio: int, ) -> np.ndarray: - """Build local FA descriptors for all layers as an Nx3 uint64 array.""" + """Build local FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert base_addresses, "Local KV cache base addresses must not be empty." device_id = self.device_id + region_indices = self._transfer_layer_region_indices or range( + len(base_addresses) + ) parts: list[np.ndarray] = [] - for i, base_addr in enumerate(base_addresses): + for i in region_indices: + base_addr = base_addresses[i] block_len = self.block_len_per_layer[i] // block_size_ratio logical_blocks = np.repeat( np.arange(self.region_num_blocks[i], dtype=np.uint64), @@ -1834,7 +1987,7 @@ def _build_fa_remote( nixl_agent_meta: NixlAgentMetadata, block_size_ratio: int, ) -> np.ndarray: - """Build remote FA descriptors for all layers as an Nx3 uint64 array.""" + """Build remote FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert nixl_agent_meta.kv_caches_base_addr, ( "Remote KV cache base addresses must not be empty." @@ -1851,10 +2004,14 @@ def _build_fa_remote( device_id = nixl_agent_meta.device_id block_strides = nixl_agent_meta.block_strides parts: list[np.ndarray] = [] + use_layer_name_routing = bool(self._transfer_layer_region_indices) for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - replicated = self._is_region_replicated(i) + local_region = ( + self._transfer_layer_region_indices[i] if use_layer_name_routing else i + ) + replicated = self._is_region_replicated(local_region) # Read our whole local region size from remote.. - local_block_len = self.block_len_per_layer[i] + local_block_len = self.block_len_per_layer[local_region] remote_kv_block_len = local_block_len // block_size_ratio if block_size_ratio > 1: # ..using remote kv_block_len as transfer unit @@ -2021,9 +2178,22 @@ def add_remote_agent( ) return self._remote_agents[engine_id][(0, remote_tp_rank)] + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo # Number of physical regions registered locally (one per layer/tensor). num_local_regions = len(self.block_len_per_layer) - if ( + if self._transfer_layer_region_indices: + if self.block_size != nixl_agent_meta.block_size: + raise NotImplementedError( + "Attention-HMA push requires identical P/D block sizes." + ) + if remote_tp_size > transfer_topo.tp_size and not self.use_mla: + raise NotImplementedError( + "Attention-HMA push does not support decode TP greater than " + "prefill TP yet" + ) + self._align_remote_regions_by_layer(nixl_agent_meta) + elif ( self.pp_size > 1 and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions ): @@ -2038,6 +2208,10 @@ def add_remote_agent( ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ + start:end + ] if nixl_agent_meta.region_group_ids is not None: nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[ start:end @@ -2048,14 +2222,8 @@ def add_remote_agent( nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] - if nixl_agent_meta.region_num_blocks is not None: - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ - start:end - ] ### Register remote engine in TransferTopology (idempotent). - assert self.transfer_topo is not None - transfer_topo = self.transfer_topo physical_blocks_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block ) @@ -2349,6 +2517,9 @@ def _validate_remote_agent_handshake( # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. + local_regions = self._transfer_layer_region_indices or range( + len(self.block_len_per_layer) + ) if self._has_mamba and self.use_mla: # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are # kernel-granularity views of the mamba-unified page. The MLA @@ -2365,7 +2536,7 @@ def _validate_remote_agent_handshake( f"remote={nixl_agent_meta.block_lens}." ) elif not self._has_mamba: - assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( + assert len(local_regions) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated( @@ -2374,8 +2545,9 @@ def _validate_remote_agent_handshake( total_kv_heads = self.transfer_topo.total_num_kv_heads local_heads = self.transfer_topo.local_physical_heads remote_heads = max(1, total_kv_heads // remote_tp_size) - for i, local_len in enumerate(self.block_len_per_layer): - replicated = model_replicated or self._is_region_replicated(i) + for i, region_idx in enumerate(local_regions): + local_len = self.block_len_per_layer[region_idx] + replicated = model_replicated or self._is_region_replicated(region_idx) remote_len = nixl_agent_meta.block_lens[i] if replicated: assert local_len // block_size_ratio == remote_len, ( @@ -2406,8 +2578,7 @@ def _validate_remote_agent_handshake( # TP workers that handhshake with same remote have same #blocks. assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks - # Same number of regions/~layers. - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) + assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions) num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr) if nixl_agent_meta.region_num_blocks is not None: assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions @@ -2423,7 +2594,12 @@ def _validate_remote_agent_handshake( if nixl_agent_meta.region_group_ids is not None: assert len(nixl_agent_meta.region_group_ids) == num_remote_regions if nixl_agent_meta.region_names is not None: - assert nixl_agent_meta.region_names == self.region_names, ( + expected_names = ( + list(self._transfer_layer_names) + if self._transfer_layer_names + else self.region_names + ) + assert nixl_agent_meta.region_names == expected_names, ( "NIXL transfer regions must name the same model caches in the " "same order" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index be32a6969d66..c2b4d9dce753 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Metadata dataclasses and helpers for the NIXL connector.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from vllm.config import VllmConfig @@ -47,8 +47,9 @@ # 9: Add block_strides # 10: Add dense virtual transfer pages for compressed MLA caches # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata +# 12: Add per-region member names for PP push # -NIXL_CONNECTOR_VERSION: int = 11 +NIXL_CONNECTOR_VERSION: int = 12 @dataclass @@ -71,6 +72,8 @@ class NixlAgentMetadata: region_mem_types: list[str] | None = None dcp_size: int = 1 pcp_size: int = 1 + # Layer names sharing each advertised region, in region order. + region_members: list[list[str]] = field(default_factory=list) @dataclass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 8b34e7ffbb4c..01c956a5251d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -83,6 +83,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # Distinguishes push from pull in the NIXL compatibility hash. _TRANSFER_MODE: str = "push" + _supports_pp_hma = True + def __init__( self, vllm_config: "VllmConfig", @@ -530,7 +532,7 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): local_region_groups = self.region_group_ids remote_region_groups = self.dst_region_group_ids[engine_id] groups_differ = local_region_groups != remote_region_groups - if groups_differ: + if groups_differ and not self._transfer_layer_group_ids: raise NotImplementedError( "NixlPushConnector does not support different producer and " "consumer cache-group layouts" From 031f5810c1957b6ec2a7666f65a3fe68c1f7e534 Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Tue, 15 Sep 2026 16:29:13 -0700 Subject: [PATCH 277/774] [Build][NVIDIA] Update public Rubin dependencies and MSA compatibility (#56545) Signed-off-by: Shang Wang Co-authored-by: Codex --- cmake/external_projects/fmha_sm100.cmake | 2 +- docker/Dockerfile | 2 +- requirements/rubin-prerelease.txt | 22 ++++++++++++++++------ 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 442097dec7d3..ae5bfec7b47d 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa + GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/docker/Dockerfile b/docker/Dockerfile index b3e32777e7bb..d1a7afcbfb70 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -641,7 +641,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ python3 use_existing_torch.py --prefix; \ fi \ && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \ + sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \ requirements/cuda.txt; \ fi diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt index 18a644f6b1b4..8b9f7d31550b 100644 --- a/requirements/rubin-prerelease.txt +++ b/requirements/rubin-prerelease.txt @@ -2,15 +2,25 @@ # Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile). # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels # (vllm-project/vllm#49735). ---pre --extra-index-url https://download.pytorch.org/whl/nightly/cu134 +--extra-index-url https://flashinfer.ai/whl/cu134 +--extra-index-url https://flashinfer.ai/whl -torch==2.15.0.dev20260903+cu134 -torchvision==0.30.0.dev20260903+cu134 -torchaudio==2.11.0.dev20260903+cu134 +# Match each architecture's torch pin to the latest torchvision wheel. +torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64" +torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64" +torchvision==0.30.0.dev20260909+cu134 +torchaudio==2.11.0.dev20260911+cu134 -# docker/Dockerfile strips these two from requirements/cuda.txt whenever +# docker/Dockerfile strips these packages from requirements/cuda.txt whenever # INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to # still get installed. nvidia-cutlass-dsl[cu13]==4.8.0.dev0 -flashinfer-python==0.6.18 +flashinfer-python==0.6.18.post1 +quack-kernels==0.6.5 + +cuda-python==13.4.1 +cuda-bindings==13.4.1 +flashinfer-jit-cache==0.6.18.post1+cu134 +flashinfer-cubin==0.6.18.post1 +nixl==1.4.1 From 85c1f58d50b80a6668f1f8158f0d8ef7b1d9112a Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:31:25 -0700 Subject: [PATCH 278/774] [BugFix][Model Runner V2][Spec Decode] Fix decode instance's multi-layer MTP kv caches during P/D (#55055) Signed-off-by: Giancarlo Delfin --- .../unit/test_nixl_push_connector.py | 2 +- tests/v1/kv_connector/unit/utils.py | 7 ++- vllm/config/vllm.py | 23 ++++++++ .../kv_connector/v1/nixl/base_scheduler.py | 52 +++++++++++++------ vllm/v1/core/kv_cache_utils.py | 7 +-- vllm/v1/core/sched/scheduler.py | 14 +---- vllm/v1/worker/gpu/model_runner.py | 10 +--- 7 files changed, 69 insertions(+), 46 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a908d3246261..7ce84bf685b1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -138,7 +138,7 @@ def test_p_side_mamba_truncates_before_cache_lookup(self): with patch.object( sched, - "_truncate_mamba_request_for_prefill", + "_truncate_request_for_prefill", side_effect=AssertionError("must not truncate after cache lookup"), ): assert sched.get_num_new_matched_tokens(request, 0) == (0, False) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 1ee1e817766a..98633c233bb2 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -521,6 +521,8 @@ def make_nixl_scheduler( Only sets the flags needed by the tests. When *heartbeat=True* the scheduler-side heartbeat bookkeeping fields are also initialised. """ + from types import SimpleNamespace + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) @@ -532,6 +534,7 @@ def make_nixl_scheduler( block_size=16, mamba_enabled=has_mamba, ) + sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0) if heartbeat: sched._heartbeat_by_engine = {} @@ -598,9 +601,11 @@ def make_nixl_push_scheduler( mamba_enabled=has_mamba, ) - # vllm_config is consulted for parallel_config.tensor_parallel_size. + # vllm_config is consulted for parallel_config.tensor_parallel_size, and by + # `_prefill_backoff` on both the P and D prefill paths. vllm_config = MagicMock() vllm_config.parallel_config.tensor_parallel_size = 1 + vllm_config.num_prefill_lookahead_tokens = 0 sched.vllm_config = vllm_config # Push-specific state. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8e0846982b74..7167ca5a9063 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -650,6 +650,29 @@ def num_lookahead_tokens(self) -> int: return self.num_speculative_tokens return 0 + @property + def num_prefill_lookahead_tokens(self) -> int: + """Prefill tokens past the computed range that the drafter reads. + + Mid-prefill the drafter consumes tokens the target model has not been + scheduled for yet, so every component that has to keep them available + must apply this margin: the scheduler, which never ends a chunk within + it and shifts encoder scheduling by it, and the KV cache manager, which + treats the trailing `this - 1` tokens as re-prefillable rather than + finalized. Consumers must read this property rather than re-deriving + their own per-method lookahead, so those components cannot drift apart. + """ + speculative_config = self.speculative_config + if speculative_config is None or not speculative_config.use_eagle(): + return 0 + if speculative_config.use_multi_module_mtp(): + # Each MTP module reads one token further ahead than the one before + # it, so the chain needs num_speculative_tokens of runway at a + # chunked-prefill boundary. + return self.num_speculative_tokens + # Eagle-family drafters read only the immediate next token. + return 1 + @property def uniform_decode_query_len(self) -> int: """Query length of every request in a uniform decode batch. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 9cfd100730a5..54184aa33945 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -199,8 +199,8 @@ def shutdown(self): def on_new_request(self, request: "Request") -> None: """Track a request that may need heartbeats.""" params = request.kv_transfer_params - if params is not None and params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) + if params is not None and params.get("do_remote_decode"): + self._truncate_request_for_prefill(request) # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are # effectively disabled for Bidirectional KV transfer. @@ -382,36 +382,56 @@ def _nixl_handshake_listener( (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts) ) + def _prefill_backoff(self) -> int: + """Trailing prompt tokens the prefiller must not compute; the decoder + recomputes them locally. + + Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module + MTP needs to keep its whole lookahead window off of the prefiller, which + would otherwise embed the unverified drafts in the MTP layer's KV cache. The + decoder would never rebuild them, because the update is sized by the rejection + count, which is zero for the first decode. + """ + return max( + 1 if self._has_mamba else 0, + self.vllm_config.num_prefill_lookahead_tokens - 1, + ) + def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: - return num_prompt_tokens - 1 + """D-side only. The number of prompt tokens to load from the prefiller. + Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder + will recompute locally.""" + backoff = self._prefill_backoff() + if backoff and num_prompt_tokens > backoff: + return num_prompt_tokens - backoff return num_prompt_tokens - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. + def _truncate_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the trailing ``_prefill_backoff()`` prompt tokens + so the prefiller stops short of what the decoder recomputes locally. + For Mamba that is the single token needed to yield h(N-1); for + multi-module MTP it is the drafter's whole lookahead window. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the request is preempted and rescheduled.""" + backoff = self._prefill_backoff() params = request.kv_transfer_params if ( - params is not None + backoff + and params is not None # Guard against repeated truncation after preemption/reschedule. and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 + and request.num_prompt_tokens > backoff ): if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() + del request.prompt_token_ids[-backoff:] elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] + request.prompt_embeds = request.prompt_embeds[:-backoff] else: return - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 + del request._all_token_ids[-backoff:] + request.num_prompt_tokens -= backoff request.max_tokens = 1 params["_p_side_truncated"] = True diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 1c6e102c161c..c55c06abf1cc 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2678,12 +2678,7 @@ def get_kv_cache_configs( # When speculating with more than 1 speculative module (e.g. multi-layered MTP) # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = ( - vllm_config.speculative_config.num_speculative_tokens - 1 - if vllm_config.speculative_config is not None - and vllm_config.speculative_config.use_multi_module_mtp() - else 0 - ) + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) for layer_name, layer_spec in merged_kv_cache_specs.items(): if isinstance(layer_spec, SlidingWindowSpec): merged_kv_cache_specs[layer_name] = replace( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index daf2b316d075..f1e11e9b1f3d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -262,13 +262,7 @@ def __init__( self.use_eagle_block_drop = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = vllm_config.num_lookahead_tokens - # Positions past the computed tokens that the drafter reads mid-prefill. - # Eagle-family drafters read 1 ahead, but multi-module MTP reads - # num_spec_tokens ahead at chunked-prefill boundaries. Determines the - # encoder scheduling shift, the deferred encoder free, the KV cache - # manager's re-prefillable window (this minus 1), and how many tokens to - # reserve between a chunk boundary and the prefill end. - self.num_prefill_lookahead = 0 + self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: if speculative_config.num_speculative_tokens_per_batch_size: @@ -278,12 +272,6 @@ def __init__( vllm_num_speculative_tokens=self.num_spec_tokens, ) self.use_eagle = speculative_config.use_eagle() - if self.use_eagle: - self.num_prefill_lookahead = ( - self.num_spec_tokens - if speculative_config.use_multi_module_mtp() - else 1 - ) self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() if self.use_eagle and not self.use_eagle_block_drop: logger.warning( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b66d552cb487..82cc697e731c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -296,15 +296,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None - # Multi-module MTP feeds its modules the next num_speculative_steps prefill - # tokens during chunked prefill. Other speculators only read the immediate - # next one. - num_prefill_lookahead = ( - self.num_speculative_steps - if self.speculative_config is not None - and self.speculative_config.use_multi_module_mtp() - else 1 - ) + num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) self.step_timing = StepTimingCollector() From 1fd119def5a841ada882fa3f33919b96783f9d50 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 19:31:51 -0500 Subject: [PATCH 279/774] [Pooling] Cap max-length padding for chunked embeddings (#56505) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/entrypoints/pooling/base/protocol.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index e30195a7b1f9..c1366734f7c1 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -171,6 +171,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin): + padding: Literal["max_length", "do_not_pad"] | None + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: default_max_total_tokens = model_config.max_model_len max_total_tokens: int | None = default_max_total_tokens @@ -184,7 +186,7 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: max_embed_len = pooler_config.max_embed_len or default_max_total_tokens max_output_tokens = default_max_total_tokens - max_embed_len - return self._build_pooling_tok_params( + tok_params = self._build_pooling_tok_params( model_config, add_special_tokens=self.add_special_tokens, max_total_tokens=max_total_tokens, @@ -192,6 +194,20 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: max_output_tokens_param="max_model_len - max_embed_len", ) + if ( + self.padding == "max_length" + and pooler_config is not None + and pooler_config.enable_chunked_processing + ): + max_padding_len = pooler_config.max_embed_len + if max_padding_len is None: + max_padding_len = default_max_total_tokens + return tok_params.with_kwargs( + pad_prompt_tokens=min(default_max_total_tokens, max_padding_len) + ) + + return tok_params + class CompletionRequestMixin(OpenAIBaseModel): # --8<-- [start:completion-params] From c8d1cf077a7878ea30040a33ac61262003a3953a Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:25:29 -0700 Subject: [PATCH 280/774] [ROCm] [Bugfix] Enable Load and Inference of GLM-5.3-Flash Quark MXFP4 Checkpoint (#56176) Signed-off-by: Colin Zeng Co-authored-by: Andreas Karatzas --- tests/quantization/test_quark.py | 145 ++++++++++++++++++ vllm/models/glm5next/__init__.py | 4 +- vllm/models/glm5next/common/__init__.py | 2 + .../glm5next/{nvidia => common}/attention.py | 0 .../models/glm5next/{nvidia => common}/kda.py | 0 .../glm5next/{nvidia => common}/model.py | 15 +- .../models/glm5next/{nvidia => common}/mtp.py | 2 +- .../glm5next/{nvidia => common}/multimodal.py | 0 8 files changed, 164 insertions(+), 4 deletions(-) create mode 100644 vllm/models/glm5next/common/__init__.py rename vllm/models/glm5next/{nvidia => common}/attention.py (100%) rename vllm/models/glm5next/{nvidia => common}/kda.py (100%) rename vllm/models/glm5next/{nvidia => common}/model.py (98%) rename vllm/models/glm5next/{nvidia => common}/mtp.py (99%) rename vllm/models/glm5next/{nvidia => common}/multimodal.py (100%) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index ae236d71f805..9f00adf78314 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -1668,3 +1668,148 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected): ) is expected ) + + +_GLM5_MXFP4_WEIGHT = { + "dtype": "fp4", + "qscheme": "per_group", + "group_size": 32, + "scale_format": "e8m0", + "is_dynamic": False, +} +_GLM5_BLOCK_FP8_WEIGHT = { + "dtype": "fp8_e4m3", + "qscheme": "per_block", + "is_dynamic": False, + "block_size": [128, 128], + "symmetric": True, +} +_GLM5_BLOCK_FP8_INPUT = { + "dtype": "fp8_e4m3", + "qscheme": "per_group", + "is_dynamic": True, + "group_size": 128, + "symmetric": True, +} +_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj" + + +def _glm5_mixed_precision_config() -> QuarkConfig: + """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8.""" + fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT} + return QuarkConfig( + { + "global_quant_config": { + "weight": _GLM5_MXFP4_WEIGHT, + "input_tensors": None, + }, + "layer_type_quant_config": {}, + "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8}, + "exclude": [], + } + ) + + +class _GLM5RecordingParam: + """Stand-in for a BF16 target param that records what gets loaded.""" + + def __init__(self): + self.loaded_weight: torch.Tensor | None = None + self.loaded_shard: object = "unset" + + def weight_loader(self, param, loaded_weight, shard_id=None): + self.loaded_weight = loaded_weight + self.loaded_shard = shard_id + + +def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128): + """Return a (fp8 weight, f32 per-block scale) pair.""" + weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn) + scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5 + return weight, scale + + +def test_glm5next_gate_up_proj_mapping_is_genuine_fusion(): + from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration + + mapping = Glm5NextForConditionalGeneration.packed_modules_mapping + assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"] + + +def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8(): + # The fused module must expand to its real shard names so each resolves to + # the per-layer block-FP8 entry (rather than the global MXFP4 scheme). + config = _glm5_mixed_precision_config() + config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} + _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP) + assert scheme_cls is QuarkW8A8Fp8PerBlock + + +def test_glm5next_attn_loader_accepts_quark_weight_scale(): + # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused + # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale + # name must be recognized and routed to the correct fused shard. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.3.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(256, 256) + # fp8 weight arrives first -> buffered, nothing loaded yet. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_weight is None + # Quark-named scale completes the pair -> dequantize + load. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight_scale", + scale, + buf, + params_dict, + loaded, + 0, + ) + is True + ) + assert target.loaded_weight is not None + assert target.loaded_weight.dtype == torch.bfloat16 + assert target.loaded_shard == 1 # kv_a is shard 1 of fused_qkv_a_proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) + assert f"{prefix}.fused_qkv_a_proj.weight" in loaded + + +def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv(): + # DeepSeek scale name keeps working unchanged. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.7.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.o_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(128, 256) + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0 + ) + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_shard is None # o_proj is a direct (non-fused) proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py index a1146e2602eb..b38806260bd4 100644 --- a/vllm/models/glm5next/__init__.py +++ b/vllm/models/glm5next/__init__.py @@ -6,8 +6,8 @@ if current_platform.is_xpu(): raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.") -from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration -from .nvidia.mtp import Glm5NextMTP +from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration +from .common.mtp import Glm5NextMTP __all__ = [ "Glm5NextForCausalLM", diff --git a/vllm/models/glm5next/common/__init__.py b/vllm/models/glm5next/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/glm5next/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/common/attention.py similarity index 100% rename from vllm/models/glm5next/nvidia/attention.py rename to vllm/models/glm5next/common/attention.py diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/common/kda.py similarity index 100% rename from vllm/models/glm5next/nvidia/kda.py rename to vllm/models/glm5next/common/kda.py diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/common/model.py similarity index 98% rename from vllm/models/glm5next/nvidia/model.py rename to vllm/models/glm5next/common/model.py index 21918bdf6bed..e4f669835e49 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/common/model.py @@ -1036,6 +1036,17 @@ class Glm5NextForConditionalGeneration( has_inner_state: ClassVar[Literal[True]] = True is_hybrid: ClassVar[Literal[True]] = True + # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like + # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``), + # so the fused ``gate_up_proj`` must expand to its real shard names for + # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry + # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj) + # would otherwise route the module to ``global_quant_config`` and mismatch + # at load for mixed-precision Quark checkpoints. + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + } + # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration # (``model.visual.`` -> ``visual.``, ``model.language_model.`` -> # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``), @@ -1259,7 +1270,9 @@ def _try_load_fp8_attn_proj( return False suffix, (key, target_base, shard_id, is_kva) = matched is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn - is_scale = "weight_scale_inv" in name + # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark + # ``weight_scale`` names before feeding the shared block dequant below. + is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale") if not is_weight and not is_scale: return False diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/common/mtp.py similarity index 99% rename from vllm/models/glm5next/nvidia/mtp.py rename to vllm/models/glm5next/common/mtp.py index 35b452f44f79..33a5402b8fd1 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -22,6 +22,7 @@ from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -33,7 +34,6 @@ _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) -from .ops.fused_eh_norm import fused_eh_norm class Glm5NextMultiTokenPredictorLayer(nn.Module): diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/common/multimodal.py similarity index 100% rename from vllm/models/glm5next/nvidia/multimodal.py rename to vllm/models/glm5next/common/multimodal.py From a4d2d9d95edd2db39cc7dec0a3433c764ff2c578 Mon Sep 17 00:00:00 2001 From: ppalanga Date: Tue, 15 Sep 2026 18:38:02 -0700 Subject: [PATCH 281/774] [ROCm][Kimi-K3] Enable a4w4 flydsl kernels for KimiK3 (#53940) Signed-off-by: ppalanga Signed-off-by: Poovaiah Palangappa Co-authored-by: Cursor Co-authored-by: Poovaiah Palangappa Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- tests/kernels/moe/test_rocm_aiter_moe.py | 68 +++++++++++++++++++ vllm/_aiter_ops.py | 38 +++++++++-- vllm/envs.py | 20 ++++-- .../fused_moe/experts/rocm_aiter_moe.py | 12 ++-- .../layers/fused_moe/oracle/mxfp4.py | 19 +++--- 5 files changed, 127 insertions(+), 30 deletions(-) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 89244bed1b0f..0967f9499e0f 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env( assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected +def test_aiter_moe_situv2_syncs_aiter_a4w4_env( + monkeypatch: pytest.MonkeyPatch, +): + """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4( + monkeypatch: pytest.MonkeyPatch, +): + """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_clears_aiter_env_when_disabled( + monkeypatch: pytest.MonkeyPatch, +): + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.setenv("AITER_SITUV2_A4W4", "1") + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert "AITER_SITUV2_A4W4" not in os.environ + + @pytest.mark.parametrize("moe_padding", [True, False]) def test_aiter_moe_padding_env_var( moe_padding: bool, diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9bea4f8b3a6e..8f2b2a24adac 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1752,6 +1752,30 @@ def _mhc_delayed_pre_tail( _OPS_REGISTERED = False +def _sync_aiter_situv2_moe_env() -> None: + """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env. + + AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 / + AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first). + When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4 + (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override. + + Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4 + dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs + afp4 FlyDSL but falls back to heuristic configs, which is not the + tuned a4w4 path this recipe is meant to use. + """ + import os + + import vllm.envs as envs + + if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2: + os.environ["AITER_SITUV2_A4W4"] = "1" + os.environ.pop("AITER_SITUV2_A8W4", None) + else: + os.environ.pop("AITER_SITUV2_A4W4", None) + + class rocm_aiter_ops: """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM. @@ -1774,7 +1798,7 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant. + VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4). VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM. Note: @@ -1841,7 +1865,7 @@ def get_moe_dispatch_policy(cls) -> int: # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 # TODO: Consolidate under _LINEAR_ENABLED _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM # Lazily probed: whether aiter.topk_softmax supports the @@ -1870,7 +1894,8 @@ def refresh_env_variables(cls): cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 + _sync_aiter_situv2_moe_env() cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM @staticmethod @@ -1983,10 +2008,10 @@ def is_fusion_moe_shared_experts_enabled(cls) -> bool: @classmethod @if_aiter_supported - def is_fused_moe_situv2_a8w4_enabled(cls) -> bool: - # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter + def is_fused_moe_situv2_enabled(cls) -> bool: + # _MOE_SITUV2 is a variant of aiter fused moe, so aiter # fused moe must be enabled as well. - return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2 @classmethod @if_aiter_supported @@ -4029,3 +4054,4 @@ def mhc_fused_post_pre( rocm_aiter_ops.register_ops_once() +_sync_aiter_situv2_moe_env() diff --git a/vllm/envs.py b/vllm/envs.py index 6158f86f896b..78af9bcb72f1 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -137,7 +137,7 @@ VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False + VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" @@ -1251,12 +1251,18 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_USE_AITER_MOE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") ), - # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up- - # interleaved flydsl kernels instead of the default a16w4 separated path. - # This is the only flag users need: vLLM picks the kernels by passing - # gate_mode to AITER and sets the AITER-side workaround env at init. - "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1") + # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4 + # activations, separated gate/up layout) instead of default a16w4. vLLM + # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any + # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first). + # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing + # recipes; it does not select a8w4 kernels. + # Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag + # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics. + "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1") + or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() + in ("true", "1") ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index cb4225c92a0f..62df395f692c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -385,14 +385,10 @@ def rocm_aiter_fused_experts( gate_mode = "" if activation == MoEActivation.SITU: - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up- - # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays - # separated. - gate_mode = ( - GateMode.INTERLEAVE.value - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() - else GateMode.SEPARATED.value - ) + # SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4 + # fp4 activations with separated gate/up weights (AITER #4463); + # default a16w4 SiTU also stays separated. + gate_mode = GateMode.SEPARATED.value elif quant_config.use_mxfp4_w4a16: gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3b415acd815d..1c38e4378b26 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1572,12 +1572,6 @@ def convert_weight_to_mxfp4_moe_kernel_format( if w2_bias is not None: w2_bias = w2_bias.data.to(torch.float32) - import os - - # TODO: Remove this once AITER is fixed - # Necessary for AITER side from crashing - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" - if activation == MoEActivation.SITU: from aiter.utility.fp4_utils import e8m0_shuffle @@ -1585,9 +1579,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( fp4_dtype = torch.float4_e2m1fn_x2 e8m0_dtype = torch.float8_e8m0fnu - # a8w4 uses gate/up-interleaved flydsl kernels; - # default a16w4 keeps the separated layout. - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() + # SiTUv2 flydsl uses separated gate/up layout (a4w4). + guinterleave = False w13 = rocm_aiter_ops.shuffle_weight_a16w4( w13_weight.data.view(fp4_dtype), 16, guinterleave ) @@ -1606,6 +1599,14 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2.is_shuffled = True return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) + import os + + # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to + # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2 + # a4w4 is separated and selects q_dtype_a independently, so the bound + # is unused on that path. + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" + from aiter.ops.shuffle import shuffle_scale as _shuf_s from aiter.ops.shuffle import shuffle_weight as _shuf_w From b0898a493737ed34864efb2dbebc36a73ea043c5 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 16 Sep 2026 10:02:37 +0800 Subject: [PATCH 282/774] [Misc] Allow install empty package from pip (#41074) Signed-off-by: wangxiyuan --- setup.py | 60 ++++++++++++++++++++++++++++------------------------ vllm/envs.py | 5 +++++ 2 files changed, 37 insertions(+), 28 deletions(-) diff --git a/setup.py b/setup.py index 13692f3d20ca..4012000e634e 100644 --- a/setup.py +++ b/setup.py @@ -1257,35 +1257,39 @@ def get_vllm_version() -> str: version = get_version(write_to="vllm/_version.py") sep = "+" if "+" not in version else "." # dev versions might contain + - if _no_device(): - if envs.VLLM_TARGET_DEVICE == "empty": - version += f"{sep}empty" - elif _is_cuda(): - if USE_PRECOMPILED_EXTENSIONS and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: - version += f"{sep}precompiled" + if not envs.VLLM_SKIP_VERSION_SUFFIX: + if _no_device(): + if envs.VLLM_TARGET_DEVICE == "empty": + version += f"{sep}empty" + elif _is_cuda(): + if ( + USE_PRECOMPILED_EXTENSIONS + and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX + ): + version += f"{sep}precompiled" + else: + cuda_version = str(get_nvcc_cuda_version()) + if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: + cuda_version_str = cuda_version.replace(".", "")[:3] + # skip this for source tarball, required for pypi + if "sdist" not in sys.argv: + version += f"{sep}cu{cuda_version_str}" + elif _is_hip(): + # Get the Rocm Version + rocm_version = get_rocm_version() or torch.version.hip + if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: + version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" + elif _is_tpu(): + version += f"{sep}tpu" + elif _is_cpu(): + # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), + # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. + if VLLM_TARGET_DEVICE == "cpu": + version += f"{sep}cpu" + elif _is_xpu(): + version += f"{sep}xpu" else: - cuda_version = str(get_nvcc_cuda_version()) - if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: - cuda_version_str = cuda_version.replace(".", "")[:3] - # skip this for source tarball, required for pypi - if "sdist" not in sys.argv: - version += f"{sep}cu{cuda_version_str}" - elif _is_hip(): - # Get the Rocm Version - rocm_version = get_rocm_version() or torch.version.hip - if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: - version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" - elif _is_tpu(): - version += f"{sep}tpu" - elif _is_cpu(): - # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), - # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. - if VLLM_TARGET_DEVICE == "cpu": - version += f"{sep}cpu" - elif _is_xpu(): - version += f"{sep}xpu" - else: - raise RuntimeError("Unknown runtime environment") + raise RuntimeError("Unknown runtime environment") return version diff --git a/vllm/envs.py b/vllm/envs.py index 78af9bcb72f1..bc4724e2e4a7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2176,6 +2176,10 @@ def _resolve_rust_cli_path() -> str | None: # Each op additionally checks its own shape / dtype constraints and falls # back to the eager path when they do not hold. "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))), + # Whether to skip version suffix when building package + "VLLM_SKIP_VERSION_SUFFIX": lambda: bool( + int(os.getenv("VLLM_SKIP_VERSION_SUFFIX", "0")) + ), } @@ -2344,6 +2348,7 @@ def compile_factors() -> dict[str, object]: "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", "NO_COLOR", + "VLLM_SKIP_VERSION_SUFFIX", } from vllm.config.utils import normalize_value From 7d4bba54985589176a60cbf3ad4424be3ee84999 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Tue, 15 Sep 2026 22:48:18 -0400 Subject: [PATCH 283/774] [Doc] Show how to get Responses prompt token IDs (#55912) Signed-off-by: Francisco Javier Arceo Co-authored-by: OpenAI Codex --- docs/serving/online_serving/README.md | 7 +++- .../openai_compatible_server.md | 5 +++ docs/serving/online_serving/renderer.md | 35 +++++++++++++++++-- 3 files changed, 44 insertions(+), 3 deletions(-) diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 1f5d4bdb6a93..a292a4a3a436 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -130,7 +130,12 @@ Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` ### Renderer APIs -For further details on renderer APIs, please refer to [this page](renderer.md). +Renderer APIs preprocess completion, chat, and Responses requests without running +inference. They handle tokenization, model-specific prompt formatting, and +multimodal preprocessing, returning prompt token IDs, sampling parameters, and +any processed multimodal inputs for generation. + +See the [renderer guide](renderer.md) for setup instructions and examples. - [Completions Render API](renderer.md) (`/v1/completions/render`) - Render completion requests diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index 22e2e9a60844..c50910040d32 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -199,6 +199,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../../examples/tool_calling/openai_responses_client_with_tools.py) +To get prompt token IDs before generation, use +[`/v1/responses/render`](renderer.md#get-responses-prompt-token-ids). It preprocesses +a self-contained Responses request without inference and returns the prompt +token IDs, so a separate `/tokenize` call is not needed. + #### Extra parameters The following extra parameters in the request object are supported: diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 9be23142c49c..9d494df22efe 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -26,27 +26,58 @@ vllm serve --enable-scale-out - [Responses Render API](renderer.md) (`/v1/responses/render`) - Render a self-contained Responses request +## Get Responses prompt token IDs + +Use `/v1/responses/render` to get prompt token IDs before choosing a model +replica. Rendering applies prompt construction and preprocessing without running +inference. + The Responses render endpoint uses the same prompt construction as `/v1/responses` and returns one token-in `GenerateRequest`. It is stateless: inline history is supported, but `previous_response_id` is not. Callers must resolve stored response state and include the resulting history in the request before rendering. +Configure the renderer and generation workers with the same model, tokenizer, +chat template, and preprocessing options. Render the full request, including +instructions, history, tools, and any template or truncation options, so the +returned IDs reflect the prompt the model will receive. + For multimodal requests, the `GenerateRequest` contains the model-processed multimodal payload, which can be substantially larger than the source image or video. The caller must forward that payload unchanged to the generation service and provision transport limits and memory accordingly. +For example, start a standard inference server with scale-out endpoints enabled: + +```bash +VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + vllm serve meta-llama/Llama-3.1-8B-Instruct +``` + +Send a Responses request and use `jq` to extract the `token_ids` field: + ```bash -curl http://localhost:8000/v1/responses/render \ +curl --fail --silent --show-error http://localhost:8000/v1/responses/render \ -H "Content-Type: application/json" \ -d '{ "model": "meta-llama/Llama-3.1-8B-Instruct", "input": "Explain prefix caching in one sentence.", "max_output_tokens": 32 - }' + }' | jq '.token_ids' ``` +To count the rendered prompt tokens for this text request, replace the `jq` +filter with `'.token_ids | length'`. The endpoint returns the full +`GenerateRequest`; `jq` filters the response on the client. Keep the complete +response if you will forward it to `/inference/v1/generate`, including any +multimodal features. + +If the server has `--api-key` or `VLLM_API_KEY` configured, add +`-H "Authorization: Bearer "` to the request. See +[API key authentication limitations](../../usage/security.md#api-key-authentication-limitations) +for the existing authentication boundaries. + For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md). ## Multimodal Render Features From 0673de980055a4283fcc152fc1ecc837ff534f1d Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:51:17 +0800 Subject: [PATCH 284/774] [Rust Frontend][Metrics] Add per-request preemption histogram (#57033) Co-authored-by: Bugen Zhao Signed-off-by: reidliu41 --- rust/src/llm/src/request_metrics.rs | 6 ++++++ rust/src/llm/tests/generate.rs | 12 ++++++++++++ rust/src/metrics/src/request.rs | 15 +++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 280fa0f10fbb..43802eaf4abe 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -45,6 +45,7 @@ pub(crate) struct RequestMetricsTracker { last_token_ts: f64, first_token_latency: f64, num_generation_tokens: u32, + num_preemptions: u64, latest_num_cached_tokens: u32, } @@ -69,6 +70,7 @@ struct RequestMetricHandles { request_max_num_generation_tokens: HistogramMetric, request_params_max_tokens: HistogramMetric, request_params_n: HistogramMetric, + request_num_preemptions: HistogramMetric, request_prefill_kv_computed_tokens: HistogramMetric, time_to_first_token_seconds: HistogramMetric, inter_token_latency_seconds: HistogramMetric, @@ -104,6 +106,7 @@ impl RequestMetricsTracker { last_token_ts: 0.0, first_token_latency: 0.0, num_generation_tokens: 0, + num_preemptions: 0, latest_num_cached_tokens: 0, } } @@ -184,6 +187,7 @@ impl RequestMetricsTracker { self.handles.request_params_max_tokens.observe(max_tokens_param as f64); } self.handles.request_params_n.observe(self.n_param as f64); + self.handles.request_num_preemptions.observe(self.num_preemptions as f64); self.handles .request_prefill_kv_computed_tokens .observe(prefill_kv_computed_tokens as f64); @@ -223,6 +227,7 @@ impl RequestMetricsTracker { } } EngineCoreEventType::Preempted => { + self.num_preemptions += 1; self.handles.num_preemptions.inc(); } } @@ -276,6 +281,7 @@ fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetri .get_or_create_owned(&labels), request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels), request_params_n: metrics.request_params_n.get_or_create_owned(&labels), + request_num_preemptions: metrics.request_num_preemptions.get_or_create_owned(&labels), request_prefill_kv_computed_tokens: metrics .request_prefill_kv_computed_tokens .get_or_create_owned(&labels), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 2ff3faa8d459..2d09aff889a8 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -790,6 +790,12 @@ async fn generate_records_request_metrics_in_prometheus_output() { assert!(rendered.contains(&format!( "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" + ))); assert!(rendered.contains(&format!( "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); @@ -883,6 +889,12 @@ async fn dropping_stream_records_abort_terminal_request_metrics() { assert!(rendered.contains(&format!( "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" + ))); llm.shutdown().await.unwrap(); } diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index 691c789936b2..5bfc32e89534 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -21,6 +21,7 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [ 480.0, 960.0, 1920.0, 7680.0, ]; const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0]; +const REQUEST_NUM_PREEMPTIONS_BUCKETS: [f64; 7] = [1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0]; const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0, ]; @@ -69,6 +70,10 @@ fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } +fn request_num_preemptions_histogram() -> Histogram { + Histogram::new(REQUEST_NUM_PREEMPTIONS_BUCKETS.iter().copied()) +} + fn iteration_tokens_histogram() -> Histogram { Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied()) } @@ -110,6 +115,7 @@ pub struct RequestMetrics { pub request_max_num_generation_tokens: HistogramFamily, pub request_params_max_tokens: HistogramFamily, pub request_params_n: HistogramFamily, + pub request_num_preemptions: HistogramFamily, pub request_prefill_kv_computed_tokens: HistogramFamily, pub time_to_first_token_seconds: HistogramFamily, pub inter_token_latency_seconds: HistogramFamily, @@ -216,6 +222,14 @@ impl RequestMetrics { request_params_n.clone(), ); + let request_num_preemptions = + Family::new_with_constructor(request_num_preemptions_histogram as fn() -> Histogram); + registry.register( + "vllm:request_num_preemptions", + "Histogram of the number of times a request was preempted.", + request_num_preemptions.clone(), + ); + let request_prefill_kv_computed_tokens = Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); registry.register( @@ -302,6 +316,7 @@ impl RequestMetrics { request_max_num_generation_tokens, request_params_max_tokens, request_params_n, + request_num_preemptions, request_prefill_kv_computed_tokens, time_to_first_token_seconds, inter_token_latency_seconds, From bfd713bf8b3869a7b65a8f3198095c3d73c88875 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 16 Sep 2026 10:52:55 +0800 Subject: [PATCH 285/774] [Rust Frontend] Support HF config overrides `--hf-overrides` (#56931) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 + rust/src/chat/src/backend/hf.rs | 22 ++--- rust/src/chat/src/backend/mod.rs | 3 + rust/src/cmd/Cargo.toml | 1 + rust/src/cmd/src/cli.rs | 17 ++++ rust/src/cmd/src/cli/tests.rs | 85 +++++++++++++++++ rust/src/cmd/src/cli/unsupported.rs | 5 - rust/src/managed-engine/src/cli.rs | 5 + .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 3 + rust/src/server/src/lib.rs | 1 + rust/src/server/src/render.rs | 4 + rust/src/text/Cargo.toml | 2 +- rust/src/text/src/backend/hf/mod.rs | 2 + rust/src/text/src/backend/hf/model_files.rs | 72 ++++++++++++++ rust/src/text/src/backend/hf/overrides.rs | 93 +++++++++++++++++++ 16 files changed, 298 insertions(+), 19 deletions(-) create mode 100644 rust/src/text/src/backend/hf/overrides.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1b7dae65685b..54a1640926fb 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5813,6 +5813,7 @@ dependencies = [ "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", + "vllm-text", "vllm-tracing", ] diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 37212a96505c..5b5b8772c55f 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,7 +129,8 @@ pub(super) async fn load_model_backends( model_id: &str, options: LoadModelBackendsOptions, ) -> Result { - let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + let mut files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + files.apply_overrides(&options.hf_overrides)?; let text_backend = HfTextBackend::from_resolved_model_files( files.clone(), model_id.to_string(), @@ -166,13 +167,11 @@ fn resolve_multimodal_render_info( #[cfg(test)] mod tests { use std::collections::HashMap; - use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; use thiserror_ext::AsReport as _; use vllm_text::Prompt; - use vllm_text::backend::hf::TokenizerSource; use vllm_text::tokenizer::DynTokenizer; use vllm_tokenizer::test_utils::TestTokenizer; @@ -206,16 +205,12 @@ mod tests { write_json(&config_path, config_json); write_json(&tokenizer_config_path, tokenizer_config_json); - vllm_text::backend::hf::ResolvedModelFiles { - tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")), - tokenizer_config_path: Some(tokenizer_config_path), - generation_config_path: None, - preprocessor_config_path: None, - video_preprocessor_config_path: None, - processor_config_path: None, - chat_template_path: None, - config_path: Some(config_path), - } + write_json(&root.join("tokenizer.json"), "{}"); + futures::executor::block_on(vllm_text::backend::hf::ResolvedModelFiles::new( + root.to_str().unwrap(), + None, + )) + .unwrap() } fn test_tokenizer() -> DynTokenizer { @@ -232,6 +227,7 @@ mod tests { "test-model".to_string(), LoadModelBackendsOptions { revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), renderer, language_model_only: false, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 028c0c40c8ec..eed4e95cbe0c 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::sync::Arc; use serde_json::Value; +use vllm_text::backend::hf::HfOverrides; use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend}; use crate::error::Result; @@ -63,6 +64,8 @@ pub type DynChatTextBackend = Arc; pub struct LoadModelBackendsOptions { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before loading any backend. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index bf92e05ee186..e61c48ee3c83 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -33,6 +33,7 @@ vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true vllm-server.workspace = true +vllm-text.workspace = true vllm-tracing.workspace = true [dev-dependencies] diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 70b28018a3f1..cda2e6c1a8b9 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -33,6 +33,7 @@ use vllm_server::{ DEFAULT_KEEP_ALIVE_TIMEOUT, HttpListenerMode, LoraModulePath, ParserSelection, RenderConfig, RendererSelection, }; +use vllm_text::backend::hf::HfOverrides; use crate::cli::ssl::SslArgs; use crate::cli::unsupported::UnsupportedArgs; @@ -113,6 +114,9 @@ pub struct RenderArgs { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). #[arg(long)] revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + hf_overrides: HfOverrides, /// HTTP bind host. #[arg(long, default_value = "127.0.0.1")] host: String, @@ -161,6 +165,7 @@ impl RenderArgs { RenderConfig { model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, served_model_name: self.served_model_name, host: self.host, port: self.port, @@ -204,6 +209,12 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + /// Objects merge recursively and arrays/scalars replace existing values. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + #[serde(default)] + pub hf_overrides: HfOverrides, + /// The source of generation-config sampling defaults. `"auto"` loads the /// model's defaults, while `"vllm"` uses vLLM's neutral defaults. #[arg(long, default_value_t)] @@ -500,6 +511,7 @@ impl SharedRuntimeArgs { }, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, @@ -556,6 +568,7 @@ impl SharedRuntimeArgs { coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode, @@ -761,6 +774,9 @@ impl ServeArgs { let reasoning_parser = effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model); let profiler_config = self.runtime.profiler_config_json(); + let hf_overrides = (!self.runtime.hf_overrides.is_empty()).then(|| { + serde_json::to_string(&self.runtime.hf_overrides).expect("JSON object serializes") + }); self.managed_engine.clone().into_config( self.runtime.model.clone(), @@ -773,6 +789,7 @@ impl ServeArgs { self.runtime.shutdown_timeout, handshake_port, self.runtime.limit_mm_per_prompt_json(), + hf_overrides, ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 411d2562959c..5cc6b1921450 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -177,6 +177,9 @@ fn serve_args_forward_python_flags_with_separator() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1013,6 +1016,9 @@ fn frontend_args_accept_json() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: None, @@ -1120,6 +1126,73 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); } +#[test] +fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { + let patch = r#"{"text_config":{"rope_parameters":{"factor":4}},"sliding_window":null}"#; + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap(); + let Command::Serve(args) = cli.command else { + panic!("expected serve args") + }; + let expected: serde_json::Value = serde_json::from_str(patch).unwrap(); + assert_eq!( + serde_json::to_value(&args.runtime.hf_overrides).unwrap(), + expected + ); + let engine = args.to_managed_engine_config(1234); + let index = engine.python_args.iter().position(|arg| arg == "--hf-overrides").unwrap(); + assert_eq!( + serde_json::from_str::(&engine.python_args[index + 1]).unwrap(), + expected + ); + assert!(!engine.python_args.iter().any(|arg| arg == "--hf-config-path")); + let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string()); + assert_eq!(serde_json::to_value(config.hf_overrides).unwrap(), expected); + + let payload = serde_json::json!({"model_tag":"Qwen/Qwen3-0.6B", "hf_overrides":expected}); + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + &payload.to_string(), + ]) + .unwrap(); + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args") + }; + assert_eq!( + serde_json::to_value(args.into_config().hf_overrides).unwrap(), + expected + ); +} + +#[test] +fn hf_overrides_reject_non_object_inputs() { + for patch in ["null", "[]", "1", r#""callable""#] { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); + } +} + #[test] fn frontend_args_json_accepts_supported_non_default_fields() { let cli = Cli::try_parse_from([ @@ -1613,6 +1686,9 @@ fn serve_args_accept_handshake_aliases() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1763,6 +1839,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1852,6 +1931,9 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1963,6 +2045,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present }, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: InheritedFd { diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 29ea5d06b0e7..8e514c95349c 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -235,11 +235,6 @@ pub struct EngineUnsupportedArgs { #[arg(long, default_missing_value = "true", num_args = 0..=1)] pub hf_token: Option, - /// If a dictionary, contains arguments to be forwarded to the Hugging Face - /// config. If a callable, it is called to update the HuggingFace config. - #[arg(long)] - pub hf_overrides: Option, - /// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If /// used with `--generation-config auto`, the override parameters will be /// merged with the default config from the model. If used with diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index f0cd7bf1715d..0a81ae598e9c 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -93,6 +93,7 @@ impl ManagedEngineArgs { shutdown_timeout: u64, handshake_port: u16, limit_mm_per_prompt: Option, + hf_overrides: Option, ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. @@ -136,6 +137,10 @@ impl ManagedEngineArgs { python_args.push("--limit-mm-per-prompt".to_string()); python_args.push(limit_mm_per_prompt); } + if let Some(hf_overrides) = hf_overrides { + python_args.push("--hf-overrides".to_string()); + python_args.push(hf_overrides); + } ManagedEngineConfig { python: self.python, diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index fc9914e3e2f2..ad3872843a1f 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -61,6 +61,7 @@ async fn main() -> Result<()> { coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), served_model_name: vec![], listener_mode: HttpListenerMode::BindTcp { diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 0f0deac86d2c..4c337adbef8f 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -16,6 +16,7 @@ use vllm_chat::{ ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, }; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; +use vllm_text::backend::hf::HfOverrides; /// Default keep-alive idle timeout (seconds); also the head-read bound /// when keep-alive is disabled (`0`). @@ -214,6 +215,8 @@ pub struct Config { pub model: String, /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before backend initialization. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 95ae359f35b3..49d3d57f3264 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -99,6 +99,7 @@ async fn build_state(config: &Config) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: config.generation_config, renderer: config.renderer, language_model_only: config.language_model_only, diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index fe7033fa8cad..4fc55fe60be9 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -13,6 +13,7 @@ use vllm_chat::{ ParserSelection, RendererSelection, load_model_backends, }; use vllm_text::TextRequestProcessor; +use vllm_text::backend::hf::HfOverrides; use crate::{ HttpListenerMode, TlsConfig, @@ -25,6 +26,7 @@ use crate::{ pub struct RenderConfig { pub model: String, pub revision: Option, + pub hf_overrides: HfOverrides, pub served_model_name: Vec, pub host: String, pub port: u16, @@ -69,6 +71,7 @@ async fn build_state(config: &RenderConfig) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: Default::default(), renderer: config.renderer, language_model_only: true, @@ -153,6 +156,7 @@ mod tests { RenderConfig { model: "test-model".to_string(), revision: None, + hf_overrides: Default::default(), served_model_name: Vec::new(), host: "127.0.0.1".to_string(), port: 8000, diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml index 7bda7f976e1c..36152d7cc4a1 100644 --- a/rust/src/text/Cargo.toml +++ b/rust/src/text/Cargo.toml @@ -16,6 +16,7 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true serde_with.workspace = true +tempfile.workspace = true thiserror.workspace = true thiserror-ext.workspace = true tracing.workspace = true @@ -28,7 +29,6 @@ vllm-tokenizer.workspace = true expect-test.workspace = true futures.workspace = true serial_test.workspace = true -tempfile.workspace = true tokio.workspace = true vllm-llm = { workspace = true, features = ["test-util"] } vllm-tokenizer = { workspace = true, features = ["test-utils"] } diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index ae768767d99a..2f32f2cb483d 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -3,6 +3,7 @@ mod config; mod model_files; +mod overrides; use std::collections::BTreeSet; use std::sync::Arc; @@ -18,6 +19,7 @@ pub use self::config::{ load_tokenizer_config, }; pub use self::model_files::{ResolvedModelFiles, TokenizerSource}; +pub use self::overrides::HfOverrides; use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend}; use crate::error::Result; diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index de802fb4390c..3ffd7b2bb79c 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -2,11 +2,13 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::path::{Path, PathBuf}; +use std::sync::Arc; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; use hf_hub::{Cache, Repo, RepoType}; use thiserror_ext::AsReport as _; +use super::HfOverrides; use super::config::{HfTokenizerConfig, load_tokenizer_config}; use crate::error::{Error, Result}; @@ -50,9 +52,27 @@ pub struct ResolvedModelFiles { pub processor_config_path: Option, pub chat_template_path: Option, pub config_path: Option, + /// Keep a patched config alive across clones and backend ownership. + config_temp_path: Option>, } impl ResolvedModelFiles { + /// Materialize a merge-patched config owned by these resolved model files. + pub fn apply_overrides(&mut self, overrides: &HfOverrides) -> Result<()> { + if overrides.is_empty() { + return Ok(()); + } + let path = materialize_config(self.config_path.as_deref(), overrides).map_err(|error| { + Error::Tokenizer(format!( + "failed to materialize HF overrides: {}", + error.as_report() + )) + })?; + self.config_path = Some(path.to_path_buf()); + self.config_temp_path = Some(Arc::new(path)); + Ok(()) + } + /// Resolve tokenizer/config files from a local model directory first when /// `model_id` points to one, otherwise consult the local HF cache and /// finally the Hub. @@ -84,6 +104,24 @@ impl ResolvedModelFiles { } } +fn materialize_config( + path: Option<&Path>, + overrides: &HfOverrides, +) -> anyhow::Result { + let mut config: serde_json::Value = match path { + Some(path) => serde_json::from_slice(&std::fs::read(path)?)?, + None => serde_json::json!({}), + }; + anyhow::ensure!(config.is_object(), "HF model config must be a JSON object"); + overrides.apply(&mut config); + let file = tempfile::Builder::new() + .prefix("vllm-rs-hf-config-") + .suffix(".json") + .tempfile()?; + serde_json::to_writer(file.as_file(), &config)?; + Ok(file.into_temp_path()) +} + fn resolve_local_model_files(model_dir: &Path) -> Result { let tokenizer_config_path = local_file_if_exists(model_dir, "tokenizer_config.json"); let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?; @@ -101,6 +139,7 @@ fn resolve_local_model_files(model_dir: &Path) -> Result { processor_config_path: local_file_if_exists(model_dir, "processor_config.json"), chat_template_path: discover_chat_template_in_dir(model_dir), config_path: local_file_if_exists(model_dir, "config.json"), + config_temp_path: None, }) } @@ -164,6 +203,7 @@ async fn resolve_remote_model_files( processor_config_path, chat_template_path, config_path, + config_temp_path: None, }) } @@ -204,6 +244,7 @@ fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result); + +impl HfOverrides { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Apply the patch to an in-memory config, preserving the source file. + pub fn apply(&self, config: &mut Value) { + merge_object(config, &self.0); + } +} + +/// Apply RFC 7396 object-member updates at every nesting level. +fn merge_object(target: &mut Value, patch: &Map) { + if !target.is_object() { + *target = Value::Object(Map::new()); + } + let target = target.as_object_mut().unwrap(); + for (key, value) in patch { + match value { + Value::Null => { + target.remove(key); + } + Value::Object(object) => { + merge_object(target.entry(key.clone()).or_insert(Value::Null), object) + } + _ => { + target.insert(key.clone(), value.clone()); + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::HfOverrides; + + #[test] + fn merge_patch_preserves_siblings_replaces_arrays_and_deletes_null_members() { + let mut config = json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 2}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 2}}, + "architectures": ["Old", "Other"], + "sliding_window": 4096 + }); + let patch: HfOverrides = serde_json::from_value(json!({ + "rope_parameters": {"factor": 4}, + "text_config": {"rope_parameters": {"factor": 4}}, + "architectures": ["New"], + "sliding_window": null, + "missing": null + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 4}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 4}}, + "architectures": ["New"] + }) + ); + } + + #[test] + fn merge_patch_creates_objects_and_preserves_nulls_inside_replacement_arrays() { + let mut config = json!({"a": 1, "b": [1, 2]}); + let patch: HfOverrides = serde_json::from_value(json!({ + "a": {"nested": {"removed": null}}, + "b": {}, + "c": [null, {"literal": null}] + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({"a": {"nested": {}}, "b": {}, "c": [null, {"literal": null}]}) + ); + } +} From 22bb158c21fcba630f267a12ebab4e86296ccc42 Mon Sep 17 00:00:00 2001 From: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:59:43 +0900 Subject: [PATCH 286/774] [Bugfix][ROCm] Fix MiniMax-M3 fused MXFP8 block scale (#53674) Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> --- tests/kernels/test_minimax_m3_amd_ops.py | 24 ++++++++++++++++++-- vllm/models/minimax_m3/amd/ops/swiglu_oai.py | 8 +++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index ec8f1fd4eaad..997c61980913 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -42,6 +42,7 @@ from vllm.models.minimax_m3.amd.ops import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, + swiglu_oai_quantize_mxfp8, swiglu_oai_split, ) from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402 @@ -167,6 +168,25 @@ def test_swiglu_oai_split(m, inter, limit, dtype): assert _relerr(got, ref) < 5e-3 +@torch.inference_mode() +def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale(): + # Keep a small value in two MX blocks next to their maxima. The scale must + # use E4M3's finite range (448), otherwise that value underflows to zero. + # Keep the third block empty to cover the reference's finite tiny clamp. + m, inter = 3, 96 + gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16) + up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16) + up[:, :64:32] = 1.0 + up[:, 64:] = 0.0 + gate_up = torch.cat((gate, up), dim=-1) + + got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None) + ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False) + + assert torch.equal(got_s, ref_s) + assert torch.equal(got_q, ref_q) + + # --------------------------------------------------------------------------- # # Fused MXFP8 activation quant (Triton vs torch reference) # --------------------------------------------------------------------------- # @@ -179,8 +199,8 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype): xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False) xq_k, s_k = _mxfp8_e4m3_quantize_triton(x) assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32) - # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at - # most a 1-step difference at exact powers of two. + # Both paths use the E4M3-aware scale calculation; allow a 1-step + # difference at exact powers of two due to floating-point rounding. assert (s_k.int() - s_t.int()).abs().max().item() <= 1 # Dequantized values agree to fp8 granularity. deq_t = dequant_mxfp8_to_bf16(xq_t, s_t) diff --git a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py index 836649b725b4..df7abc0760db 100644 --- a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py +++ b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py @@ -83,6 +83,8 @@ def _swiglu_oai_quant_kernel( limit, HAS_LIMIT: tl.constexpr, BLOCK_M: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, ): """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0) quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one @@ -109,8 +111,8 @@ def _swiglu_oai_quant_kernel( gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32 - amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M] - sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0) + amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M] + sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0) descale = tl.exp2(sb - 127.0) aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty) tl.store( @@ -173,6 +175,8 @@ def swiglu_oai_quantize_mxfp8( 0.0 if limit is None else float(limit), HAS_LIMIT=limit is not None, BLOCK_M=block_m, + FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max), + TINY=float(torch.finfo(torch.float32).tiny), num_warps=4, ) return aq, asc From fdfcbbcac79b0d8134a3a447f2c33001f5de07c1 Mon Sep 17 00:00:00 2001 From: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:02:25 -0700 Subject: [PATCH 287/774] [CI/Build][Rust Frontend] Retire Buf schema publishing (#57046) Co-authored-by: OpenAI Codex Co-authored-by: Bugen Zhao Signed-off-by: Julien Darve --- .github/workflows/buf.yml | 29 ----------------------------- rust/proto/README.md | 15 +++++---------- rust/proto/buf.md | 8 -------- rust/proto/buf.yaml | 1 - 4 files changed, 5 insertions(+), 48 deletions(-) delete mode 100644 rust/proto/buf.md diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml index fda4add1d726..dd695d28d529 100644 --- a/.github/workflows/buf.yml +++ b/.github/workflows/buf.yml @@ -5,12 +5,6 @@ on: paths: - "rust/proto/**" - ".github/workflows/buf.yml" - push: - tags: - - "v*" - schedule: - - cron: "0 8 * * *" - workflow_dispatch: permissions: contents: read @@ -21,37 +15,14 @@ concurrency: jobs: proto: - if: github.repository_owner == 'vllm-project' runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - with: - fetch-depth: 0 - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 with: version: 1.72.0 input: rust/proto - token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }} format: false breaking: false pr_comment: false push: false - - name: Publish nightly schema - if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label nightly - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" - - name: Publish release schema - if: github.repository == 'vllm-project/vllm' && github.event_name == 'push' - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label main - --label "${GITHUB_REF_NAME}" - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" diff --git a/rust/proto/README.md b/rust/proto/README.md index 62b134280bd9..3181cdc6a07b 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -2,14 +2,10 @@ This directory is the canonical source for vLLM's gRPC schema. -The schema is published to `buf.build/vllm-project/vllm`: - -- A daily workflow publishes the latest Git `main` schema to the `nightly` label. -- The workflow can be run manually to retry nightly publication. -- Tags matching `v*` update the Buf `main` label and publish the corresponding release label. -- Buf commits and generated SDK versions are immutable and can be pinned by consumers. - -Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically. +Schema updates are no longer published to the Buf Schema Registry. Rust consumers +should use `vllm-proto` from crates.io; consumers in other languages can generate +bindings from the `.proto` files in this directory. Buf still builds and lints +the schemas on pull requests, including on forks. ## Rust crate @@ -47,8 +43,7 @@ change PR. This check becomes available after the first manual publication. The `proto-crate.yml` workflow checks that the tag matches the crate version, verifies the package, and publishes it to crates.io. -Use `proto-v*` tags so they do not trigger Buf's `v*` release workflow. There is -no scheduled crate publication. Pull requests and manual workflow runs only +There is no scheduled crate publication. Pull requests and manual workflow runs only verify packaging, including on forks. Retry a failed release by rerunning its tag workflow; crates.io versions cannot be overwritten. diff --git a/rust/proto/buf.md b/rust/proto/buf.md deleted file mode 100644 index 065124feee65..000000000000 --- a/rust/proto/buf.md +++ /dev/null @@ -1,8 +0,0 @@ -# vLLM gRPC API - -Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend. - -- `Inference` provides unary and streaming generation. -- `Control` provides discovery and request control. - -The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds. diff --git a/rust/proto/buf.yaml b/rust/proto/buf.yaml index 64764f561ef3..d30dc3297d76 100644 --- a/rust/proto/buf.yaml +++ b/rust/proto/buf.yaml @@ -1,7 +1,6 @@ version: v2 modules: - path: . - name: buf.build/vllm-project/vllm lint: use: - MINIMAL From 4ace5cd9d80a8d56bf935c4a058a1369f1b94da3 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Tue, 15 Sep 2026 20:05:43 -0700 Subject: [PATCH 288/774] [XPU] fix_test_worker_memory_snapshot_for_xpu (#50097) Signed-off-by: mayuyuace Co-authored-by: Kunshang Ji --- .../v1/worker/test_worker_memory_snapshot.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/v1/worker/test_worker_memory_snapshot.py b/tests/v1/worker/test_worker_memory_snapshot.py index fe8a5a21f8dc..6ed62cb771fa 100644 --- a/tests/v1/worker/test_worker_memory_snapshot.py +++ b/tests/v1/worker/test_worker_memory_snapshot.py @@ -12,6 +12,7 @@ from vllm.config import set_current_vllm_config from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment @@ -66,8 +67,22 @@ def worker_process( model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy" ).create_engine_config() + # Select worker class and patch target based on platform + if current_platform.is_xpu(): + from vllm.v1.worker.xpu_worker import XPUWorker + + worker_cls = XPUWorker + init_patch_target = ( + "vllm.v1.worker.xpu_worker.init_worker_distributed_environment" + ) + else: + worker_cls = Worker + init_patch_target = ( + "vllm.v1.worker.gpu_worker.init_worker_distributed_environment" + ) + # Create worker - worker = Worker( + worker = worker_cls( vllm_config=vllm_config, local_rank=rank, rank=rank, @@ -81,7 +96,7 @@ def worker_process( # Apply minimal patches to track operation order init_patch = patch( - "vllm.v1.worker.gpu_worker.init_worker_distributed_environment", + init_patch_target, side_effect=make_operation_tracker( "init_distributed", original_init_worker ), @@ -93,7 +108,7 @@ def worker_process( ) all_reduce_patch = patch( "torch.distributed.all_reduce", - side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce), + side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce), ) with ( @@ -182,13 +197,13 @@ def test_init_distributed_is_called_before_memory_snapshot(): # Raises ValueError if the operation is not found init_distributed = rank_ops.index("init_distributed") - nccl_all_reduce = rank_ops.index("nccl_all_reduce") + dist_all_reduce = rank_ops.index("dist_all_reduce") memory_snapshot = rank_ops.index("memory_snapshot") # Verify order: init_distributed should happen before memory_snapshot - assert init_distributed < nccl_all_reduce < memory_snapshot, ( + assert init_distributed < dist_all_reduce < memory_snapshot, ( f"Rank {rank}: init_distributed (index {init_distributed}) " - f"must happen before nccl_all_reduce (index {nccl_all_reduce}) " + f"must happen before dist_all_reduce (index {dist_all_reduce}) " f"and memory_snapshot (index {memory_snapshot})" ) From ccbaed5b91e09f528180db363af7ebb9d4fafe32 Mon Sep 17 00:00:00 2001 From: Marceli Fylcek Date: Wed, 16 Sep 2026 05:05:54 +0200 Subject: [PATCH 289/774] [XPU] Route to fused_qk_rmsnorm_rope_gate triton kernel (#56096) Signed-off-by: Marceli Fylcek Co-authored-by: Kunshang Ji --- vllm/model_executor/models/qwen3_next.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 495650eb1989..7ea3ed131919 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -378,7 +378,7 @@ def __init__( self.use_fused_qk_norm_rope_gate = ( self.attn_output_gate and getattr(self.rotary_emb, "is_neox_style", False) - and current_platform.is_cuda() + and (current_platform.is_cuda() or current_platform.is_xpu()) and supports_dtype and (text_only or supports_mrope) ) From c16bb6068f70878fb8a2f7c4d6cda95cd03a778b Mon Sep 17 00:00:00 2001 From: Islam <43927639+almersawi@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:10:58 +0400 Subject: [PATCH 290/774] [Bugfix][Rust Frontend] Tolerate NaN-corrupted logprobs in engine-core (#51026) Co-authored-by: Bugen Zhao Signed-off-by: Islam Signed-off-by: Bugen Zhao --- .../src/protocol/logprobs.rs | 7 +- .../src/protocol/logprobs/tests.rs | 118 ++++++++++++++++++ .../src/routes/openai/utils/logprobs.rs | 47 +++++++ 3 files changed, 168 insertions(+), 4 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index d5bc40e06159..646c106a04ed 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct TokenLogprob { pub token_id: u32, + /// Preserves the engine's value, including NaN and infinities. pub logprob: f32, /// The sampled/selected token uses its actual vocab rank. Remaining entries /// use 1-based top-k ranks matching the engine's returned candidate /// order. + /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's + /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values. pub rank: u32, } @@ -53,10 +56,6 @@ impl PositionLogprobs { logprobs.len() ); } - if sampled_rank == 0 { - bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs"); - } - let mut entries = Vec::with_capacity(token_ids.len()); for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() { let rank = if index == 0 { diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 2470460b52a6..f69ff22f1ae0 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() { "new_logprobs: zero-column logprobs payload with 2 rows" ); } + +fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value { + let rows = ranks.len(); + let cols = ids.len() / rows; + Value::Array(vec![ + ndarray_value( + " Date: Wed, 16 Sep 2026 11:53:23 +0800 Subject: [PATCH 291/774] [XPU] Honor device_ids for worker device placement (#56015) Signed-off-by: Chaojun Zhang Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- vllm/v1/engine/core.py | 11 +++--- vllm/v1/worker/xpu_worker.py | 50 +++++++++++++++++++++------ 3 files changed, 44 insertions(+), 19 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index a76331f4d5c4..ef2b4ba07396 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -97,7 +97,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + 'cd tests && pip install dpctl && pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ffc9a5de7e68..62b6ef880e29 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -2452,13 +2452,10 @@ def _set_nixl_side_channel_host(): def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int): from vllm.platforms import current_platform - if current_platform.is_xpu(): - pass - else: - device_control_env_var = current_platform.device_control_env_var - self._set_assigned_physical_gpu_ids( - vllm_config, local_dp_rank, device_control_env_var - ) + device_control_env_var = current_platform.device_control_env_var + self._set_assigned_physical_gpu_ids( + vllm_config, local_dp_rank, device_control_env_var + ) def _set_assigned_physical_gpu_ids( self, diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 68d6bde3a049..5c5d125f4bfa 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -115,27 +115,55 @@ def init_device(self): self.local_rank += dp_local_rank * replica_world_size - assert self.local_rank < visible_device_count, ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - assert parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." - ) - device = self.device_config.device if ( isinstance(device, torch.device) and device.type == "xpu" and current_platform.is_xpu() ): - self.device = torch.device(f"xpu:{self.local_rank}") + assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + # NOTE: local_world_size is derived from parallel_config.nnodes, + # which is only set for the "mp" multi-node backend. With the + # "ray"/"external_launcher" backends nnodes stays 1, so + # local_world_size collapses to the full world_size and this + # check wrongly fires on cross-node deployments. + # assigned_physical_gpu_ids is already per-node and the + # local_rank bound above fully validates the mapping for + # these backends, so skip the check for them. + if parallel_config.distributed_executor_backend not in ( + "ray", + "external_launcher", + ): + assert parallel_config.local_world_size <= len( + assigned_physical_gpu_ids + ), ( + f"local_world_size ({parallel_config.local_world_size})" + " exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: + assert self.local_rank < torch.accelerator.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." + ) + + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id(self.local_rank) + ) + self.device = torch.device(f"xpu:{visible_device_index}") torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) torch.accelerator.empty_cache() self.init_gpu_memory = torch.xpu.get_device_properties( - self.local_rank + visible_device_index ).total_memory else: raise RuntimeError(f"Unsupported device type: {self.device_config.device}") From 75dc5882699be65eaee144cb7d0b06c1f86d3cec Mon Sep 17 00:00:00 2001 From: alberto Date: Wed, 16 Sep 2026 04:57:38 +0100 Subject: [PATCH 292/774] [KV Offload] Add per-request `max_load_tokens control` (#55885) Signed-off-by: Alberto Perdomo --- docs/features/kv_offloading_usage.md | 29 ++++++ .../offloading_connector/test_scheduler.py | 93 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 30 ++++++ 3 files changed, 151 insertions(+), 1 deletion(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index a5c8684b4214..22a2350d8095 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -22,6 +22,35 @@ flowchart LR CPU <--> SN["..."] ``` +## Per-request load control + +Individual requests can cap how many tokens are loaded from offloaded storage +by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens +beyond those already available in the GPU prefix cache. Set it to `0` to +disable external loading for the request: + +```json +{ + "kv_transfer_params": { + "max_load_tokens": 0 + } +} +``` + +GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap +are recomputed. Omitting the field leaves loading uncapped. The value must be a +non-negative integer; invalid values are ignored. Positive caps are rounded +down to a boundary supported by the configured KV cache groups. Offloading +remains enabled unless it is controlled separately with `max_offload_tokens`. + +!!! warning + `max_load_tokens` is experimental and subject to change. + +`kv_load_tiers` continues to select secondary tiers. CPU is always included: +it can satisfy a resident hit directly and is the required staging tier for +secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but +does not query any secondary tier. + ## Terminology: Chunks The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e7a83cb7b507..968474dd74a6 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, call import pytest @@ -195,15 +196,17 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: def _make_partial_tail_request( scheduler: OffloadingConnectorScheduler, + kv_transfer_params: dict[str, Any] | None = None, ) -> MagicMock: request = MagicMock() request.request_id = "req" - request.kv_transfer_params = None + request.kv_transfer_params = kv_transfer_params request.num_prompt_tokens = 30 request.num_tokens = 30 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)] request.all_token_ids = list(range(30)) request.lora_request = None + request.skip_reading_prefix_cache = False request.is_finished.return_value = False scheduler.on_new_request(request) return request @@ -427,6 +430,53 @@ def test_lookup_cap_stops_at_authoritative_prefix_boundary(): assert scheduler._req_status["req"].partial_tail_boundary == 20 +def test_max_load_tokens_zero_skips_partial_tail_lookup(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 0} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False) + scheduler.manager.lookup.assert_not_called() + assert all( + state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states + ) + + +def test_max_load_tokens_caps_tokens_beyond_gpu_prefix(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 8} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True) + assert scheduler._req_status["req"].partial_tail_boundary == 24 + + +def test_max_load_tokens_rounds_down_without_partial_tail(): + scheduler = _make_partial_tail_scheduler() + scheduler.config = scheduler.config._replace(supports_partial_tail=False) + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 20} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True) + + +@pytest.mark.parametrize("value", ["8", 8.5, -1, True]) +def test_invalid_max_load_tokens_is_ignored(value): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": value} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True) + + def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): """A recurrent group may hold unhashed blocks inside the computed prefix. @@ -2964,6 +3014,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): runner.manager.lookup.assert_not_called() +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool): + """The load cap limits external reads without changing the store path.""" + block_size = 4 + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + blocks_per_chunk=blocks_per_chunk, + ) + + runner.new_request(token_ids=[0] * tokens_per_chunk) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + ) + + runner.scheduler.reset_prefix_cache() + runner.manager.lookup.reset_mock() + runner.new_request( + token_ids=[0] * tokens_per_chunk, + kv_transfer_params={"max_load_tokens": 0}, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_loaded=(), + expected_stored=(0, 1, 2), + ) + + runner.manager.lookup.assert_not_called() + + # --------------------------------------------------------------------------- # Eagle/MTP test class # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 2fb3c5bb905d..8cd5e90b6fac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -354,6 +354,8 @@ class RequestOffloadState: req_context: ReqContext offloading_context: RequestOffloadingContext group_states: tuple[RequestGroupState, ...] = field(init=False) + # upper bound on tokens to load for this request; None means no cap + max_load_tokens: int | None = None # upper bound on tokens to offload for this request; None means no cap max_offload_tokens: int | None = None # number of hits in the GPU cache @@ -376,6 +378,20 @@ def __post_init__(self) -> None: ) params = self.req.kv_transfer_params + # NOTE: This field is experimental and subject to change in the future. + raw = params.get("max_load_tokens") if params else None + if type(raw) is int and raw >= 0: + self.max_load_tokens = raw + logger.debug( + "Request %s: max_load_tokens set to %d", + self.req.request_id, + raw, + ) + elif raw is not None: + logger.warning( + "max_load_tokens must be a non-negative int, got %r; ignoring", raw + ) + # NOTE: This field is experimental and subject to change in the future. raw = params.get("max_offload_tokens") if params else None if type(raw) is int and raw >= 0: @@ -767,6 +783,11 @@ def _lookup_complete_chunks( max_hit_size_tokens = min( max_hit_size_tokens, num_computed_tokens + max_num_new_tokens ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = min( + max_hit_size_tokens, + num_computed_tokens + req_status.max_load_tokens, + ) if self._sliding_window_groups: # the last prompt token has to be recomputed to get the logprobs # for sliding window attention, we must reduce by 1 to make sure @@ -810,6 +831,10 @@ def _lookup_complete_chunks( max_hit_size_tokens = min( max_hit_size_tokens, len(offload_keys) * tokens_per_chunk ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = round_down( + max_hit_size_tokens, tokens_per_chunk + ) if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk: # We can only load less than a chunk, so skip. return 0 @@ -965,6 +990,11 @@ def _lookup( max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) if max_num_new_tokens is not None: max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) + if req_status.max_load_tokens is not None: + max_boundary = min( + max_boundary, + local_tokens + req_status.max_load_tokens, + ) max_boundary = round_down(max_boundary, tokens_per_hash) if max_boundary <= complete_boundary: return complete_hit From a31ec3a68bbbedd4b5d59490c762bcf32abf1a17 Mon Sep 17 00:00:00 2001 From: danisereb Date: Wed, 16 Sep 2026 07:20:28 +0300 Subject: [PATCH 293/774] [LoRA][Nemotron] Add LoRA support for Nemotron VL models (for the language model only) (#56231) Signed-off-by: Daniel Serebrenik --- docs/models/supported_models.md | 7 +++++ .../model_executor/models/nano_nemotron_vl.py | 29 +++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d55348ad724a..ddc50ad41cdb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -602,6 +602,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I+ + V+ | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ | +| `NemotronH_Nano_Omni_Reasoning_V3`, `NemotronH_Nano_VL_V2` | Nemotron Nano 3 Omni, Nemotron Nano V2 VL | T + IE+ + V+ + A* | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | ✅︎ | | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ | | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + IE+ + VE+ | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ | @@ -701,6 +702,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor `meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through its [DFlash](../features/speculative_decoding/README.md) path. +!!! note + Both Nemotron Nano architecture names map to the same vLLM implementation. + Audio inputs are only supported by the Omni variants; `NemotronH_Nano_VL_V2` + checkpoints such as `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` accept text, + image and video only. + !!! note The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now. For more details, please see: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 2edf56871dbf..6f5e073bd623 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -23,6 +23,7 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.logger import init_logger +from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.layers.activation import ReLUSquaredActivation from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -30,6 +31,7 @@ HasInnerState, IsHybrid, MultiModalEmbeddings, + SupportsLoRA, SupportsMultiModal, SupportsMultiModalPruning, ) @@ -897,11 +899,22 @@ def get_dummy_mm_data( dummy_inputs=NanoNemotronVLDummyInputsBuilder, ) class NemotronH_Nano_VL_V2( - nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning + nn.Module, + HasInnerState, + IsHybrid, + SupportsMultiModal, + SupportsMultiModalPruning, + SupportsLoRA, ): requires_sequential_video_encoding = True """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet""" + # LoRA covers the language model only + is_non_gated_moe = NemotronHForCausalLM.is_non_gated_moe + packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping + embedding_modules = NemotronHForCausalLM.embedding_modules + lora_skip_prefixes = NemotronHForCausalLM.lora_skip_prefixes + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "language_model.backbone": "language_model.model", @@ -1331,7 +1344,19 @@ def _create_final_video_embeddings( # Create final video embeddings, merging text embeddings for indicator # tokens with video embeddings - text_embeddings = self.get_language_model().embed_input_ids(repl_token_ids) + + # LoRA support - + # These replacement tokens are produced inside the encoder, so they are + # absent from the request-token batch that the LoRA adapter index + # mapping is built from, and there can be more + # of them than max_num_batched_tokens. + # Embed them with the base weights - + # the LoRA delta is undefined for tokens with no mapping entry. + embed_tokens = self.get_language_model().model.embed_tokens + if isinstance(embed_tokens, BaseLayerWithLoRA): + embed_tokens = embed_tokens.base_layer + text_embeddings = embed_tokens(repl_token_ids) + final_video_embeddings = _merge_multimodal_embeddings( inputs_embeds=text_embeddings, multimodal_embeddings=video_embeddings, From 6ecd97f1b7fb13991065370f6d6be3dffe9fe1db Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 16 Sep 2026 12:34:55 +0800 Subject: [PATCH 294/774] [Perf][DSpark] Add KV-only context insertion across V4.1 cache formats (#56441) Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: Yongye Zhu Co-authored-by: 0z5a <0z5a@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 87 ++++++++++++- csrc/libtorch_stable/ops.h | 7 ++ csrc/libtorch_stable/torch_bindings.cpp | 8 ++ tests/kernels/test_compressor_kv_cache.py | 116 ++++++++++++++++++ vllm/models/deepseek_v41/nvidia/dspark.py | 77 +++--------- 5 files changed, 234 insertions(+), 61 deletions(-) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 254b9783bc5a..344b82fbda2c 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -657,7 +657,12 @@ static void launchFusedDeepseekV4Templated( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - if (num_tokens_full < NUM_TOKEN_CUTOFF) { + // The reduced grid pays off by iterating a token's Q heads inside one CTA. + // With zero query heads there is nothing to iterate, so one warp per block + // would do all the work and the other seven idle; keep the warp-per-slot + // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single + // insert, where it is worth ~2x. + if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { cudaLaunchKernelEx(&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, @@ -993,6 +998,86 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() && + kv.dim() == 2 && kv.size(1) == 512 && + kv.scalar_type() == ScalarType::BFloat16, + "kv must be contiguous CUDA bfloat16 [tokens, 512]"); + auto const device = kv.get_device_index(); + STD_TORCH_CHECK(k_cache.device().is_cuda() && + k_cache.get_device_index() == device && + k_cache.dim() == 3 && cache_block_size > 0 && + k_cache.size(1) == cache_block_size && + k_cache.stride(2) == 1, + "k_cache must be a paged cache on the KV device"); + for (auto const* tensor : {&slot_mapping, &position_ids}) { + STD_TORCH_CHECK(tensor->device().is_cuda() && + tensor->get_device_index() == device && + tensor->scalar_type() == ScalarType::Long && + tensor->dim() == 1 && tensor->is_contiguous(), + "slots and positions must be contiguous int64 on the KV device"); + } + STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) && + slot_mapping.size(0) <= kv.size(0), + "positions must match KV rows; slots may omit padded rows"); + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() && + cos_sin_cache.get_device_index() == device && + cos_sin_cache.scalar_type() == ScalarType::Float && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(), + "cos_sin_cache must be contiguous float32 [max_pos, 64]"); + bool const packed = k_cache.scalar_type() == ScalarType::Byte; + bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK(packed || fp8 || + k_cache.scalar_type() == ScalarType::BFloat16, + "cache must be uint8, bfloat16 or float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512), + "cache row width does not match its dtype"); + if (fp8) { + STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() && + fp8_scale->get_device_index() == device && + fp8_scale->scalar_type() == ScalarType::Float && + fp8_scale->numel() == 1, + "fp8 cache requires one float32 KV scale on the KV device"); + } + int const num_tokens = static_cast(slot_mapping.size(0)); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard(device); + const cudaStream_t stream = get_current_cuda_stream(device); + using scalar_t = c10::BFloat16; + auto const* input = reinterpret_cast(kv.const_data_ptr()); + auto* cache = reinterpret_cast(k_cache.mutable_data_ptr()); + auto const* slots = slot_mapping.const_data_ptr(); + auto const* positions = position_ids.const_data_ptr(); + auto const* cos_sin = cos_sin_cache.const_data_ptr(); + // Zero query heads schedules only the existing KV branch, preserving its + // RoPE rounding and quantization without allocating any query tensors. + if (packed) { + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, + num_tokens, num_tokens, 0, static_cast(cache_block_size), + static_cast(k_cache.stride(0)), kv_mxfp8, stream); + } else if (fp8) { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, + fp8_scale->const_data_ptr(), nullptr, 0.0f, num_tokens, num_tokens, + 0, static_cast(cache_block_size), k_cache.stride(0), + k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream); + } else { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr, + nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), + k_cache.stride(0) * 2, k_cache.stride(1) * 2, false, + "fused_deepseek_v4_kv_rope_insert", stream); + } +} + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index e19f49d6035f..bc08d90a4836 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -259,6 +259,13 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 54548dfd0571..3918ab8498cd 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); ops.def( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" "Tensor q_in, Tensor kv, Tensor! k_cache, " @@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Positional encoding kernels (shared CUDA/ROCm) ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); ops.impl( diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 7e092229fe70..c1de94be92fc 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -42,6 +42,122 @@ from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize( + "cache_dtype,kv_mxfp8", + [ + (torch.uint8, False), + (torch.uint8, True), + (torch.bfloat16, False), + (torch.float8_e4m3fn, False), + ], +) +@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024]) +@pytest.mark.parametrize("use_graph", [False, True]) +def test_dspark_context_kv_matches_query_insert( + cache_dtype, kv_mxfp8, num_tokens, use_graph +): + """KV-only insertion must preserve every cache byte, including graph replay.""" + from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv + + torch.manual_seed(42) + block_size = 256 + num_blocks = math.ceil((num_tokens + 7) / block_size) + row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(num_tokens + 3, device="cuda") + 7 + angles = torch.randn(num_tokens + 16, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens] + slots[::5] = -1 + scale = torch.tensor([0.7], device="cuda") + # Only what _insert_context_kv actually reads: the record width comes off + # the cache tensor, so no per-record width attribute is stubbed here. + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size), + kv_mxfp8=kv_mxfp8, + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=scale, + ) + + def legacy_insert(): + q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda") + if cache_dtype == torch.uint8: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + reference.view(num_blocks, -1), + slots, + positions, + cos_sin, + 8, + 1e-20, + block_size, + True, + kv_mxfp8, + ) + elif cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + reference, + slots, + positions, + cos_sin, + 1e-20, + block_size, + ) + else: + q_fp8 = torch.empty_like(q, dtype=cache_dtype) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + reference, + slots, + positions, + cos_sin, + scale, + scale, + 1e-20, + block_size, + ) + + def insert(): + _insert_context_kv(attn, kv, positions, slots) + + insert() + graph = None + if use_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + insert() + for iteration in range(3): + kv.normal_() + positions.add_(1) + if iteration == 2: + slots.fill_(-1) + backing.view(torch.uint8).fill_(165) + reference_backing.copy_(backing) + before = kv.clone() + legacy_insert() + insert() if graph is None else graph.replay() + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close(kv, before, rtol=0, atol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index c19dca800555..214681f3a932 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -234,67 +234,24 @@ def _insert_context_kv( positions: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: - """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. - - Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy - query and discard it, since context tokens have no query). Mirrors - ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. - """ - swa_cache = attn.swa_cache_layer.kv_cache - block_size = attn.swa_cache_layer.block_size - cos_sin_cache = attn.rotary_emb.cos_sin_cache - cache_dtype = swa_cache.dtype - n_ctx = kv.shape[0] - dummy_q = torch.zeros( - (n_ctx, attn.n_local_heads, attn.head_dim), - dtype=kv.dtype, - device=kv.device, + """Insert normalized context KV without constructing an unused query.""" + # The bound cache already is [num_blocks, block_size, bytes_per_token], so + # the record width needs no per-layer attribute -- one whose name the + # attention layer is free to change (kv_bytes_per_token has to become + # swa_bytes_per_token once the sliding-window and compressed records can + # differ). block_size stays the layer's, so the op's size(1) check still + # cross-checks the two rather than restating the tensor. + cache = attn.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + slot_mapping, + positions, + attn.rotary_emb.cos_sin_cache, + attn.swa_cache_layer.block_size, + attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None, + attn.kv_mxfp8, ) - if cache_dtype == torch.uint8: - # fp8_ds_mla UE8M0 paged layout - swa_2d = swa_cache.view(swa_cache.shape[0], -1) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - dummy_q, - kv, - swa_2d, - slot_mapping, - positions, - cos_sin_cache, - attn.padded_heads, - attn.eps, - block_size, - True, # apply_q_norm (unused: the query is a discarded dummy) - attn.kv_mxfp8, - ) - elif cache_dtype == torch.bfloat16: - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - dummy_q, - kv, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn.eps, - block_size, - ) - else: # per-tensor fp8 (torch.float8_e4m3fn) - # TODO(ben): double-check if this is being dispatched correctly for FI backend - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - dummy_q, - kv, - dummy_q_fp8, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn._flashinfer_fp8_kv_scale, - attn._flashinfer_fp8_q_scale_inv, - attn.eps, - block_size, - ) class DSparkDeepseekV4ForCausalLM(nn.Module): From af1c01499b289be555c475669ba50a88e96d846e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 21:37:08 -0700 Subject: [PATCH 295/774] [Bugfix] Make GPU sync checks safe under torch.compile (#56904) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../generation/test_voxtral_realtime.py | 3 ++- tests/utils_/test_gpu_sync_debug.py | 18 ++++++++++++++++++ vllm/utils/gpu_sync_debug.py | 14 ++++++++++++-- 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 1bd9f6230353..2c13e07ff367 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -21,7 +21,7 @@ from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.kv_cache_interface import SlidingWindowSpec -from ....utils import ROCM_ENGINE_KWARGS +from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn" @@ -192,6 +192,7 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa "cudagraph_mode", [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], ) +@create_new_process_for_each_test(method="spawn") def test_voxtral_realtime_cudagraph( audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode ): diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index dc8b12280010..10b550c0a174 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -123,6 +123,24 @@ def test_suppressing_works_while_compiling(monkeypatch): with_gpu_sync_check(suppressed)() +@create_new_process_for_each_test() +def test_copy_checker_is_traceable_while_compiling(monkeypatch): + """The copy wrapper must bypass ContextVars while Dynamo is tracing.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + gsd._install_copy_checkers() + + def convert_like(x, like): + return x.to(device=like.device, dtype=like.dtype) + + compiled = torch.compile(convert_like, fullgraph=True) + result = compiled( + torch.ones(4, device="cuda", dtype=torch.float64), + torch.ones(4, device="cuda", dtype=torch.float32), + ) + assert result.dtype == torch.float32 + + @create_new_process_for_each_test() def test_sync_debug_mode_restored_after_checked_call(monkeypatch): """The mode is armed only for the duration of a checked call. Leaving it diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 53fcdfe71004..8d0e7ae1871f 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -225,6 +225,16 @@ def _install_compile_time_sync_suppressors() -> None: _torch_cuda = torch.Tensor.cuda _torch_copy_ = torch.Tensor.copy_ + +# Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor +# through this module global instead of as a recognized tensor method. Put the +# descriptor behind a weak-referenceable Python callable so ``allow_in_graph`` +# can admit it and the downstream compiler can lower the original operation. +@torch.compiler.allow_in_graph +def _torch_to_in_graph(self, *args, **kwargs): + return _torch_to(self, *args, **kwargs) + + # torch's own parser for `Tensor.to`'s overloaded signature, returning # (device, dtype, non_blocking, memory_format) (also used by torch._dynamo). _parse_to = getattr(torch._C._nn, "_parse_to", None) @@ -232,7 +242,7 @@ def _install_compile_time_sync_suppressors() -> None: def _active_check_mode() -> str | None: """The mode the calling thread is being checked in, if any.""" - if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling(): + if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get(): return None return _checking.get() @@ -296,7 +306,7 @@ def _checked_to(self, *args, **kwargs): else: if non_blocking and (reason := _cpu_copy_stall_reason(self, device)): _report_implicit_sync("to", "H2D", reason, mode) - return _torch_to(self, *args, **kwargs) + return _torch_to_in_graph(self, *args, **kwargs) def _checked_cuda(self, *args, **kwargs): From 711768fc17a82c530a1d3c3f070daf40a76caba4 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 16 Sep 2026 00:09:45 -0500 Subject: [PATCH 296/774] [CI][ROCm] Add opt-in TheRock builds for AMD CI (#56351) Signed-off-by: Andreas Karatzas Co-authored-by: Codex --- .buildkite/scripts/ci-bake-rocm.sh | 78 +++++++++++++---- .../scripts/hardware_ci/run-amd-test.sh | 10 +++ .buildkite/scripts/rocm/build-config.sh | 85 +++++++++++++++++++ .buildkite/scripts/rocm/refresh-base-image.sh | 29 ++++++- docker/docker-bake-rocm.hcl | 6 +- 5 files changed, 188 insertions(+), 20 deletions(-) create mode 100644 .buildkite/scripts/rocm/build-config.sh diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 15f8274fc62b..6c5f3308b785 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -540,6 +540,26 @@ prepare_ci_build_context() { echo "Using canonical CI Docker context: ${ROCM_BUILD_CONTEXT_ROOT}" } +configure_custom_rocm_stages() { + using_custom_rocm_dockerfiles || return 0 + local context_root="${ROCM_BUILD_CONTEXT_ROOT:-.}" + local dockerfile="${context_root}/${CI_BASE_DOCKERFILE}" + local stages="" + + (cd "${context_root}" && \ + validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \ + ci_base test export_vllm export_test_smoke csrc-build rust-build) \ + || return $? + # Hash all stages for custom layouts, including additional helper stages. + stages=$(rocm_dockerfile_stages "${dockerfile}" | paste -sd ' ') + CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${stages}}" + ROCM_CSRC_DOCKERFILE_STAGES="${ROCM_CSRC_DOCKERFILE_STAGES:-${stages}}" + ROCM_RUST_DOCKERFILE_STAGES="${ROCM_RUST_DOCKERFILE_STAGES:-${stages}}" + if [[ " ${stages} " != *" build_rocshmem "* ]]; then + unset ROCSHMEM_BRANCH ROCSHMEM_CACHE_KEY DEEPEP_CACHE_KEY + fi +} + compose_dependency_cache_key() { local prefix="$1" local material="$2" @@ -569,11 +589,11 @@ hash_dockerfile_stages() { } emit = 0 } - $1 == "FROM" { + toupper($1) == "FROM" { stage = "" for (idx = 1; idx <= NF; idx++) { if (tolower($idx) == "as" && idx < NF) { - stage = $(idx + 1) + stage = tolower($(idx + 1)) } } emit = (stage in wanted) @@ -613,11 +633,11 @@ discover_dockerfile_stage_args() { } { line = $0 - if ($1 == "FROM") { + if (toupper($1) == "FROM") { stage = "" for (idx = 1; idx <= NF; idx++) { if (tolower($idx) == "as" && idx < NF) { - stage = $(idx + 1) + stage = tolower($(idx + 1)) } } emit = (stage in wanted) @@ -630,7 +650,7 @@ discover_dockerfile_stage_args() { for (idx = 1; idx <= line_count; idx++) { line = lines[idx] arg_name = line - sub(/^[[:space:]]*ARG[[:space:]]+/, "", arg_name) + sub(/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+/, "", arg_name) if (arg_name != line) { sub(/[=[:space:]].*/, "", arg_name) if (arg_name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) { @@ -716,7 +736,7 @@ extract_dockerfile_arg_default() { if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" fi - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ "${physical_dockerfile}" | head -1 } @@ -1000,7 +1020,21 @@ create_and_bootstrap_builder() { } init_config() { + # shellcheck source=.buildkite/scripts/rocm/build-config.sh + source "$(dirname "${BASH_SOURCE[0]}")/rocm/build-config.sh" + configure_rocm_build + TARGET="${1:-test-ci}" + if using_custom_rocm_dockerfiles; then + if [[ -z "${BASE_IMAGE:-}" ]]; then + echo "Custom ROCm builds require BASE_IMAGE from the selected base handoff" >&2 + return 1 + fi + if ! is_ci_base_target && [[ -z "${CI_BASE_IMAGE:-}" ]]; then + echo "Custom ROCm builds require CI_BASE_IMAGE from the selected ci_base handoff" >&2 + return 1 + fi + fi BAKE_TARGETS=("${TARGET}") DEPENDENCY_CACHE_TARGETS=() CI_HCL_SOURCE="${CI_HCL_SOURCE:-${CI_HCL_FILE:-${DEFAULT_CI_HCL_SOURCE}}}" @@ -1010,10 +1044,12 @@ init_config() { PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950}" CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" - CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + if ! using_custom_rocm_dockerfiles; then + CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + fi CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" - export PYTORCH_ROCM_ARCH + export PYTORCH_ROCM_ARCH CI_BASE_DOCKERFILE SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX) CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl" @@ -1042,6 +1078,11 @@ print_header() { } validate_inputs() { + if using_custom_rocm_dockerfiles; then + validate_rocm_dockerfile "${ROCM_BASE_DOCKERFILE}" || return $? + validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \ + ci_base test export_vllm export_test_smoke csrc-build rust-build || return $? + fi if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}" echo "Make sure you're running from the vLLM repository root" @@ -1069,6 +1110,11 @@ load_ci_hcl() { init_bake_files() { BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") + if using_custom_rocm_dockerfiles; then + # Custom stacks reuse local cache without replacing the + # standard ROCm registry caches for the same commit or branch. + BAKE_FILES+=(--set '*.cache-to=') + fi } compute_ci_base_hash_if_needed() { @@ -1804,7 +1850,7 @@ compute_rocm_csrc_content_hash() { local -a content_args=() bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" read -r -a content_paths <<< "${content_files}" mapfile -t content_args < <( get_content_arg_names "${dockerfile_rocm}" "${stages}" "${ROCM_CSRC_CONTENT_ARGS:-}" @@ -1853,7 +1899,7 @@ compute_rocm_rust_content_hash() { local -a content_args=() bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" read -r -a content_paths <<< "${content_files}" mapfile -t content_args < <( get_content_arg_names \ @@ -1943,7 +1989,7 @@ write_rocm_build_arg_override() { local arg_value="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" mapfile -t arg_names < <( { get_content_arg_names \ @@ -2259,7 +2305,7 @@ extract_dependency_pins() { local val="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" physical_dockerfile="${dockerfile_rocm}" if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile_rocm}" != /* ]]; then physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile_rocm}" @@ -2275,12 +2321,12 @@ extract_dependency_pins() { fi val=$( - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ "${physical_dockerfile}" | head -1 ) if [[ -n "${val}" ]]; then export "${var}=${val}" - echo "Extracted ${var}=${val} from Dockerfile.rocm" + echo "Extracted ${var}=${val} from ${dockerfile_rocm}" fi done } @@ -2297,7 +2343,7 @@ compute_dependency_cache_keys() { local deepep_material="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH") ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") @@ -2687,6 +2733,7 @@ main() { init_bake_files if is_ci_base_target; then prepare_ci_build_context + configure_custom_rocm_stages fi compute_ci_base_hash_if_needed configure_ci_base_image_refs @@ -2697,6 +2744,7 @@ main() { # version metadata only after that lookup sees the available tag history. if ! is_ci_base_target; then prepare_ci_build_context + configure_custom_rocm_stages fi extract_dependency_pins write_rocm_build_arg_override diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 62fda27bf0be..541a8f55258c 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -34,6 +34,9 @@ ############################################################################### set -o pipefail +# shellcheck source=.buildkite/scripts/rocm/build-config.sh +source "$(dirname "${BASH_SOURCE[0]}")/../rocm/build-config.sh" || exit $? + : "${BUILDKIT_PROGRESS:=plain}" : "${TERM:=xterm-256color}" : "${FORCE_COLOR:=1}" @@ -106,6 +109,10 @@ clear_ci_orchestration_env() { VLLM_CI_REQUIRE_WORKSPACE_MOUNT \ VLLM_TEST_COMMANDS \ VLLM_CI_BRANCH \ + CI_ROCM_DOCKERFILE_BASE \ + CI_ROCM_DOCKERFILE \ + ROCM_BASE_DOCKERFILE \ + CI_BASE_DOCKERFILE \ VLLM_CI_BASE_IMAGE \ VLLM_CI_FALLBACK_IMAGE \ VLLM_CI_DOCKER_DISABLED \ @@ -225,6 +232,9 @@ prepare_artifact_image() { metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1) if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then base_image=$(tr -d '[:space:]' < "${metadata_file}") + elif using_custom_rocm_dockerfiles; then + echo "Custom ROCm ci_base metadata is missing; using the full CI image" + return 1 fi echo "--- Preparing local ROCm test image" diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh new file mode 100644 index 000000000000..ad0179b16359 --- /dev/null +++ b/.buildkite/scripts/rocm/build-config.sh @@ -0,0 +1,85 @@ +#!/usr/bin/env bash +# Shared Dockerfile selection for the normal AMD CI build chain. +# CI_ROCM_DOCKERFILE_BASE and CI_ROCM_DOCKERFILE accept repository-relative paths. +# For Rock, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rock_base and +# CI_ROCM_DOCKERFILE=docker/Dockerfile.rock. +# Custom recipes must preserve the CI targets and artifact layout. Additional +# COPY inputs also need the corresponding *_CONTENT_FILES overrides. + +rocm_base_dockerfile() { + local default="docker/Dockerfile.rocm_base" + printf '%s\n' "${CI_ROCM_DOCKERFILE_BASE:-${ROCM_BASE_DOCKERFILE:-${default}}}" +} + +rocm_ci_dockerfile() { + local default="docker/Dockerfile.rocm" + printf '%s\n' "${CI_ROCM_DOCKERFILE:-${CI_BASE_DOCKERFILE:-${default}}}" +} + +using_custom_rocm_dockerfiles() { + [[ "$(rocm_base_dockerfile)" != "docker/Dockerfile.rocm_base" \ + || "$(rocm_ci_dockerfile)" != "docker/Dockerfile.rocm" ]] +} + +rocm_dockerfile_stages() { + awk ' + toupper($1) == "FROM" { + for (idx = 2; idx < NF; idx++) { + if (tolower($idx) == "as") print tolower($(idx + 1)) + } + } + ' "$1" +} + +validate_rocm_dockerfile() { + local dockerfile="$1" + local stage="" + local stages="" + shift + + case "${dockerfile}" in + /*|../*|*/../*|*/..|*[[:space:]]*) + echo "ROCm Dockerfile must be a repository-relative path without '..' or whitespace: ${dockerfile}" >&2 + return 2 + ;; + esac + if [[ ! -f "${dockerfile}" ]]; then + echo "ROCm Dockerfile not found: ${dockerfile}" >&2 + return 2 + fi + if (($#)) && ! awk ' + toupper($1) == "FROM" { + named = 0 + for (idx = 2; idx < NF; idx++) { + if (tolower($idx) == "as") named = 1 + } + if (!named) exit 1 + } + ' "${dockerfile}"; then + echo "ROCm Dockerfile must name every build stage: ${dockerfile}" >&2 + return 2 + fi + stages=$(rocm_dockerfile_stages "${dockerfile}") + for stage in "$@"; do + if ! grep -Fxq -- "${stage}" <<< "${stages}"; then + echo "ROCm Dockerfile ${dockerfile} is missing required stage: ${stage}" >&2 + return 2 + fi + done +} + +configure_rocm_build() { + ROCM_BASE_DOCKERFILE=$(rocm_base_dockerfile) + CI_BASE_DOCKERFILE=$(rocm_ci_dockerfile) + export ROCM_BASE_DOCKERFILE CI_BASE_DOCKERFILE + using_custom_rocm_dockerfiles || return 0 + + # Custom stacks may supply dependencies without separate build stages. + export ROCM_DEP_CACHE_EXPORT_MODE=never + + # Preserve the standard runtime tags, which are unique to each build, + # while preventing experimental builds from promoting stable images. + export ROCM_BASE_PUSH_STABLE_TAG=0 + export CI_BASE_PUSH_STABLE_TAG=0 + echo "Custom ROCm Dockerfiles: ${ROCM_BASE_DOCKERFILE} and ${CI_BASE_DOCKERFILE}" +} diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh index 0e10b98c9eb2..10e594577a87 100755 --- a/.buildkite/scripts/rocm/refresh-base-image.sh +++ b/.buildkite/scripts/rocm/refresh-base-image.sh @@ -4,6 +4,10 @@ set -euo pipefail +# shellcheck source=.buildkite/scripts/rocm/build-config.sh +source "$(dirname "${BASH_SOURCE[0]}")/build-config.sh" +configure_rocm_build + DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}" BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}" CACHE_REPO="${ROCM_BASE_CACHE_REPO:-${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}}" @@ -12,7 +16,12 @@ DEFAULT_ROCM_BASE_METADATA_VERSION="2" DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}" ROCM_BASE_LAYER_CACHE_REF="" -ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-main" +BASE_CACHE_PREFIX="rocm-base" +if using_custom_rocm_dockerfiles; then + BASE_CACHE_PREFIX="rocm-base-$(printf '%s\n%s\n' \ + "${ROCM_BASE_DOCKERFILE}" "${CI_BASE_DOCKERFILE}" | sha256sum | cut -c1-12)" +fi +ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-main" ROCM_BASE_STABLE_TAG_UPDATED=0 declare -a ROCM_BASE_CACHE_ARGS=() @@ -120,7 +129,7 @@ configure_rocm_base_layer_cache() { fi scope=$(rocm_base_layer_cache_scope) - ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-${scope}" + ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-${scope}" ROCM_BASE_CACHE_ARGS+=( --cache-from "type=registry,ref=${ROCM_BASE_LAYER_CACHE_REF}" ) @@ -139,7 +148,7 @@ configure_rocm_base_layer_cache() { extract_arg_default() { local arg_name="$1" - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ "${DOCKERFILE}" | head -1 } @@ -419,11 +428,16 @@ build_base_image() { read -r -a content_paths <<< "${content_files}" content_files_hash="$(compute_content_hash "${content_paths[@]}")" base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}") - rocm_version="$(rocm_version_from_base_image "${base_image_arg}")" + rocm_version="$(extract_arg_default ROCM_SDK_VERSION)" + rocm_version="${rocm_version:-$(rocm_version_from_base_image "${base_image_arg}")}" triton_arg="$(extract_arg_default TRITON_BRANCH)" pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)" pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)" pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)" + triton_arg="${triton_arg:-$(extract_arg_default TRITON_VERSION)}" + pytorch_arg="${pytorch_arg:-$(extract_arg_default TORCH_VERSION)}" + pytorch_vision_arg="${pytorch_vision_arg:-$(extract_arg_default TORCHVISION_VERSION)}" + pytorch_audio_arg="${pytorch_audio_arg:-$(extract_arg_default TORCHAUDIO_VERSION)}" fa_arg="$(extract_arg_default FA_BRANCH)" aiter_arg="$(extract_arg_default AITER_BRANCH)" mori_arg="$(extract_arg_default MORI_BRANCH)" @@ -569,6 +583,13 @@ build_base_image() { } main() { + if using_custom_rocm_dockerfiles; then + if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then + echo "Custom ROCm Dockerfiles require ROCM_BASE_REFRESH_SKIP=0" >&2 + return 2 + fi + validate_rocm_dockerfile "${DOCKERFILE}" || return $? + fi metadata_set "rocm-base-refresh" "0" if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl index 9d73a6505eab..ab523a296a93 100644 --- a/docker/docker-bake-rocm.hcl +++ b/docker/docker-bake-rocm.hcl @@ -33,6 +33,10 @@ variable "CI_BASE_CONTENT_HASH" { default = "" } +variable "CI_BASE_DOCKERFILE" { + default = "docker/Dockerfile.rocm" +} + # REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/) # REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source) variable "REMOTE_VLLM" { @@ -74,7 +78,7 @@ group "default" { } target "_common-rocm" { - dockerfile = "docker/Dockerfile.rocm" + dockerfile = CI_BASE_DOCKERFILE context = "." args = { max_jobs = MAX_JOBS From 961c5b6b7ea98b765a6d436434d18ef945596f32 Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Wed, 16 Sep 2026 13:14:09 +0800 Subject: [PATCH 297/774] [Bugfix][KV Offload] Restore MTP-retained sliding-window history (#56709) Signed-off-by: jacklin78911-collab Co-authored-by: Codex Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 92 +++++++++---------- .../kv_connector/v1/offloading/scheduler.py | 30 +++++- 2 files changed, 70 insertions(+), 52 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 968474dd74a6..89b77ae43fa8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -79,20 +79,31 @@ from vllm.v1.request import RequestStatus -@pytest.mark.parametrize("boundary", [3840, 3904]) +@pytest.mark.parametrize( + "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)] +) +@pytest.mark.parametrize("boundary", [3840, 3904, 3968]) @pytest.mark.parametrize("eagle", [False, True]) @pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"]) -def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): +def test_swa_offload_window_covers_unaligned_hit( + boundary, eagle, left_state, window, extra_retained +): """A smaller SWA group can move the hit inside another group's CPU chunk.""" groups = [] - for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]): + for i, (block_size, group_window) in enumerate( + [(256, None), (64, window), (8, 128)] + ): kwargs = dict( block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 ) kv_spec = ( FullAttentionSpec(**kwargs) - if window is None - else SlidingWindowSpec(**kwargs, sliding_window=window) + if group_window is None + else SlidingWindowSpec( + **kwargs, + sliding_window=group_window, + extra_retained_tokens=extra_retained if i == 1 else 0, + ) ) groups.append( KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1) @@ -153,7 +164,9 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): sched._chunks_being_loaded.add(left_key) external, load_async = sched.get_num_new_matched_tokens(request, 0) - if boundary == 3904 and left_state != "present": + first_retained_block = (boundary - window + 1 - extra_retained) // 64 + needs_left = first_retained_block // 4 < end - 1 + if needs_left and left_state != "present": assert external == (0 if left_state == "missing" else None) assert not load_async return @@ -177,8 +190,8 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): sched.update_state_after_alloc(request, blocks, external) assert len(sched._jobs) == 1 job = next(iter(sched._jobs.values())) - if boundary == 3904: - assert {swa[14], swa[15]} <= job.keys + if needs_left: + assert {left_key, swa[end - 1]} <= job.keys manager.complete_load(job.keys, state.req_context) @@ -2629,27 +2642,14 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner assert req_status.transfer_jobs == pending_store_jobs +@pytest.mark.parametrize("extra_retained", [0, 1, 3]) @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_swa_alignment_skip(request_runner, async_scheduling: bool): - """SWA blocks unreachable by the load path are skipped during store. - - Simulates a DeepSeek V4-like hybrid architecture where SWA groups have - much smaller block sizes than the full-attention (MLA) group, causing - most SWA blocks to be unreachable by the alignment-based load path. - - Setup: - - Group 0: full attention (MLA-like), block_size=16 - - Group 1: SWA, block_size=4, sliding_window=8 - - alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment. - sliding_window_size_in_chunks = ceil(8 / 4) = 2. - Within each segment of 4 SWA blocks, only the trailing 2 are stored. - - With 32 tokens (2 full-attn blocks, 8 SWA blocks): - - Group 0 stores: blocks 0, 1 (all full-attn blocks) - - Group 1 stores: blocks 2, 3, 6, 7 (skip 0,1,4,5) +def test_swa_alignment_skip(request_runner, async_scheduling: bool, extra_retained): + """Store the window a cold consumer must restore at each FA boundary. - For real DeepSeek V4 (100K tokens), this reduces SWA stores by ~78%. + With FA blocks of 16 tokens and SWA blocks of 4 tokens, a window of 8 + needs the last two blocks per segment. Retaining three extra tokens for + MTP requires one more block; retaining just one token does not. """ full_attn_block_size = 16 swa_block_size = 4 @@ -2674,6 +2674,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): head_size=1, dtype=torch.float32, sliding_window=sliding_window, + extra_retained_tokens=extra_retained, ), ), ] @@ -2691,16 +2692,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): # Group 0: full attention assert kv_group_configs[0].sliding_window_size_in_chunks is None assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size - # Group 1: SWA with sliding_window_size of 2 chunks - assert kv_group_configs[1].sliding_window_size_in_chunks == 2 + # Group 1: SWA, including any additional retained blocks + assert kv_group_configs[1].sliding_window_size_in_chunks == ( + 3 if extra_retained > 1 else 2 + ) assert kv_group_configs[1].tokens_per_chunk == swa_block_size # alignment_tokens = full_attn_block_size = 16 assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size - # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks - # (block_size=4). Decode 1 token to kick off processing (stores are - # deferred to next step). - num_tokens = 32 + # Keep a partial last block so decoding does not reuse the producer's + # freed blocks before the harness records its deferred stores. + num_tokens = 33 runner.new_request(token_ids=[0] * num_tokens) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) @@ -2713,12 +2715,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): ) runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0 (full attn, block_size=16): 2 offloaded chunks - # -> GPU blocks (0, 0) and (0, 1) - # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2 - # per segment of 4: - # Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3) - # Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7) + # Retain each segment's tail, including MTP history when needed. expected_stored=( (0, 0), (0, 1), @@ -2726,27 +2723,26 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): (1, 3), (1, 6), (1, 7), - ), + ) + + (((1, 1), (1, 5)) if extra_retained > 1 else ()), ) # Verify that loads still work correctly for the stored SWA blocks. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) - runner.manager.lookup.return_value = LookupResult.HIT - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2 + runner.manager.lookup.side_effect = lambda key, ctx: ( + LookupResult.HIT if (key, 0) in runner.offloaded else LookupResult.MISS + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0: full prefix lookup hits 2 offloaded chunks - # -> loads GPU blocks (0, 0), (0, 1) - # Group 1: sliding window lookup finds trailing 2 from last segment - # (blocks 6, 7 which were stored) - # -> loads GPU blocks (1, 6), (1, 7) + # The consumer restores the full FA prefix and the final SWA tail. expected_loaded=( (0, 0), (0, 1), (1, 6), (1, 7), - ), + ) + + (((1, 5),) if extra_retained > 1 else ()), ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 8cd5e90b6fac..5cc66e4bf0ff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time from collections.abc import Iterable, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import chain, islice from typing import Any, NamedTuple @@ -114,7 +114,10 @@ def load_window_size_in_chunks(self, num_tokens: int) -> int | None: window = max( window, cdiv( - self.kv_cache_spec.sliding_window - 1 + right_padding, + self.kv_cache_spec.sliding_window + - 1 + + self.kv_cache_spec.extra_retained_tokens + + right_padding, self.tokens_per_chunk, ), ) @@ -126,7 +129,13 @@ def get_sliding_window_size_in_chunks( ) -> int | None: if isinstance(kv_cache_spec, SlidingWindowSpec): assert kv_cache_spec.sliding_window > 0 - return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk) + return max( + cdiv(kv_cache_spec.sliding_window, tokens_per_chunk), + cdiv( + kv_cache_spec.sliding_window - 1 + kv_cache_spec.extra_retained_tokens, + tokens_per_chunk, + ), + ) if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec): # Attention never reaches back past one chunk @@ -1461,11 +1470,24 @@ def _reachable_store_block_mask( ) -> list[bool] | None: """Build the block mask for a range of candidate offload chunks.""" blocks_per_chunk = self.config.blocks_per_chunk + kv_cache_spec = group_config.kv_cache_spec + if isinstance(kv_cache_spec, SlidingWindowSpec) and ( + kv_cache_spec.extra_retained_tokens + ): + # Offload restores all allocated history, including MTP re-prefill + # tokens. Widen only the store mask, not the model's attention window. + kv_cache_spec = replace( + kv_cache_spec, + sliding_window=( + kv_cache_spec.sliding_window + kv_cache_spec.extra_retained_tokens + ), + extra_retained_tokens=0, + ) return group_config.manager_cls.reachable_block_mask( start_block=start_chunk_idx * blocks_per_chunk, end_block=end_chunk_idx * blocks_per_chunk, alignment_tokens=self.config.alignment_tokens, - kv_cache_spec=group_config.kv_cache_spec, + kv_cache_spec=kv_cache_spec, use_eagle=group_config.is_eagle_group, retention_interval=self.config.retention_interval, reachable_boundaries=reachable_boundaries, From 4fe9e6f6e5647a159c9e65e4f3bf33dec5d12e1c Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 16 Sep 2026 13:50:35 +0800 Subject: [PATCH 298/774] [Refactor][GLM-5.3-Flash] Move sparse_attn_indexer_kpool into the model folder and split AMD/NVIDIA (#55358) Signed-off-by: zjy0516 Co-authored-by: TJian --- .../test_sparse_indexer_decode_seq_lens.py | 16 +- vllm/models/glm5next/amd/sparse_indexer.py | 729 ++++++++++++++++++ vllm/models/glm5next/common/attention.py | 2 +- vllm/models/glm5next/common/sparse_indexer.py | 170 ++++ .../glm5next/nvidia/sparse_indexer.py} | 467 ++--------- vllm/models/glm5next/sparse_indexer.py | 22 + 6 files changed, 1008 insertions(+), 398 deletions(-) create mode 100644 vllm/models/glm5next/amd/sparse_indexer.py create mode 100644 vllm/models/glm5next/common/sparse_indexer.py rename vllm/{model_executor/layers/sparse_attn_indexer_kpool.py => models/glm5next/nvidia/sparse_indexer.py} (67%) create mode 100644 vllm/models/glm5next/sparse_indexer.py diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py index 94fdadd6e95d..d87a82c37b2b 100644 --- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py +++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py @@ -20,10 +20,10 @@ import torch -# Bootstrap the glm5next package before entering the indexer module: its -# kpool_compress import runs glm5next/__init__, which pulls model -> -# attention -> back into sparse_attn_indexer_kpool (attention.py imports the -# class at module scope). Production always enters via attention.py first. +# Bootstrap the glm5next package before importing kpool_compress directly: its +# import runs glm5next/__init__, which pulls model -> attention -> the indexer +# dispatcher (sparse_indexer) -> back into the nvidia/amd sparse_indexer +# modules. Production always enters via attention.py first. # isort: off import vllm.models.glm5next # noqa: F401 from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402 @@ -32,13 +32,17 @@ ) # isort: on -import vllm.model_executor.layers.sparse_attn_indexer_kpool as indexer_mod -from vllm.model_executor.layers.sparse_attn_indexer_kpool import ( +from vllm.models.glm5next.common.sparse_indexer import ( _decode_topk_seq_lens, _fill_short_decode_causal_indices, ) from vllm.platforms import current_platform +if current_platform.is_rocm(): + from vllm.models.glm5next.amd import sparse_indexer as indexer_mod +else: + from vllm.models.glm5next.nvidia import sparse_indexer as indexer_mod + KPOOL = 4 TOPK_TOKENS = 16 SELECT_K = TOPK_TOKENS // KPOOL diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py new file mode 100644 index 000000000000..f47c9164bc8b --- /dev/null +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -0,0 +1,729 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom Sparse Attention Indexer layers.""" + +import torch + +import vllm.envs as envs +from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels +from vllm._aiter_ops import rocm_aiter_ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.custom_op import CustomOp +from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops +from vllm.models.glm5next.common.sparse_indexer import ( + RADIX_TOPK_WORKSPACE_SIZE, + _build_decode_scatter_indices, + _decode_topk_seq_lens, + _fill_causal_indices, + _fill_short_decode_causal_indices, + _gather_workspace_shapes, + _scatter_decode_tokens_by_request, + kv_cache_as_quant_view, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, +) +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +# kpool write helper: form pools from the current token batch and compress them +# into the index K cache via the fused Triton kernel. + + +def _kpool_compress_insert( + k: torch.Tensor, + gate_score: torch.Tensor, + ape: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kpool: int, + head_dim: int, + round_scale: bool, +) -> None: + """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots. + + ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec): + only the *last* token of each complete pool carries a valid (>=0) slot; + intra-pool tokens are -1. Every position is treated as a pool-completion + candidate and non-completions are masked off inside the kernel. Compacting + the valid rows first costs two device syncs on the eager prefill path and + buys nothing numerically. Assumes pool-aligned chunk starts. + """ + n = slot_mapping.shape[0] + # No pool can complete in a batch smaller than one pool; also keeps the + # clamped gather indices below in bounds. + if n < kpool: + return + pos = torch.arange(n, device=k.device) + valid = slot_mapping >= 0 + # Drop pools whose start falls before the batch (leading padding); their + # gate/k data is undefined anyway. + write_mask = valid & (pos >= kpool - 1) + offs = torch.arange(kpool, device=k.device) + idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :] + kpool_ops.kpool_compress_and_write_cache( + kv_cache, + k[idx], # [n, kpool, head_dim] + gate_score[idx], + ape, + slot_mapping.to(torch.int64), + pool_size=kpool, + head_dim=head_dim, + write_mask=write_mask, + round_scale=round_scale, + write_cache=True, + return_compressed=False, + ) + + +@eager_break_during_capture +def sparse_attn_indexer_kpool( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool, + use_fp4_cache: bool = False, + # kpool params (Plan-A: gate is consumed at write time and read back at + # topk time to softmax-weight the pool). + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + # Paged tail cache (in-progress pool's raw K + gate score), replacing the + # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix] + # for the tail group's token-granular slot_mapping. None on the dummy/profiling + # path and when the tail cache is disabled. + tail_kv_cache: torch.Tensor | None = None, + tail_prefix: str | None = None, +) -> torch.Tensor: + # careful! this will be None in dummy run + attn_metadata = get_forward_context().attn_metadata + fp8_dtype = current_platform.fp8_dtype() + k_cache_prefix = _resolve_layer_name(k_cache_prefix) + + # assert isinstance(attn_metadata, dict) + if not isinstance(attn_metadata, dict): + # Reserve workspace for indexer during profiling run + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + current_workspace_manager().get_simultaneous( + values_spec, + scales_spec, + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + + # Reserve profiler-visible memory for the worst-case decode logits, + # whose shape is [B * next_n, max_model_len]. This profiling branch + # returns before invoking the logits kernel itself. + cfg = get_current_vllm_config_or_none() + worst_decode_tokens = 0 + if cfg is not None: + sched = cfg.scheduler_config + num_spec = ( + cfg.speculative_config.num_speculative_tokens + if cfg.speculative_config is not None + else 0 + ) + worst_decode_tokens = min( + sched.max_num_seqs * (num_spec + 1), + sched.max_num_batched_tokens, + ) + # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes. + decode_logits_elems = worst_decode_tokens * max_model_len * 4 + prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + max_logits_elems = max(decode_logits_elems, prefill_cap_elems) + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) + + return topk_indices_buffer + attn_metadata_narrowed = attn_metadata[k_cache_prefix] + assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) + slot_mapping = attn_metadata_narrowed.slot_mapping + has_decode = attn_metadata_narrowed.num_decodes > 0 + has_prefill = attn_metadata_narrowed.num_prefills > 0 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + + # q_scale is required iff the FP4 cache path is enabled; the FP8 path + # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. + if use_fp4_cache: + assert q_scale is not None, "use_fp4_cache=True requires q_scale" + else: + assert q_scale is None, "q_scale must be None when use_fp4_cache=False" + + # During speculative decoding, k may be padded to the CUDA graph batch + # size while slot_mapping only covers actual tokens. Truncate k to avoid + # out-of-bounds reads in the kernel. + num_tokens = slot_mapping.shape[0] + if k is not None: + k = k[:num_tokens] + + if not skip_k_cache_insert: + assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" + if index_kpool > 1 and gate_score is not None and compress_ape is not None: + # kpool prefill write: pool kpool consecutive prefill tokens via + # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots. + # Decode tokens (the first num_decode_tokens in the batch) cannot be + # pooled here — their pool's earlier tokens are not in this batch — + # so they are deferred to the tail-buffer kernel in has_decode. + # compress_ratio == index_kpool makes slot_mapping pool-granular. + n_prefill = num_tokens - num_decode_tokens + if n_prefill > 0: + # decode tokens are batched first; prefill tokens follow. + prefill_slice = slice(num_decode_tokens, num_tokens) + _kpool_compress_insert( + k[prefill_slice], + gate_score[prefill_slice], + compress_ape, + kv_cache, + slot_mapping[prefill_slice], + index_kpool, + head_dim, + round_scale=(scale_fmt is not None), + ) + # Persist each request's incomplete prefill pool so decode can + # finish it, including after PD transfer. Tail slots use + # ``pos % kpool`` within the request's tail block. Processing + # only the batch's trailing tokens would miss all but the last + # request in a multi-request prefill. + if tail_kv_cache is not None and tail_prefix is not None: + tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix)) + if tail_meta is not None: + assert isinstance(tail_meta, DeepseekV32IndexerMetadata) + kpool_ops.kpool_seed_tail_cache( + tail_kv_cache, + k[prefill_slice], + gate_score[prefill_slice], + tail_meta.slot_mapping[prefill_slice], + index_kpool, + head_dim, + ) + else: + # standard: per-token fp8 quant + scatter (all tokens). + assert scale_fmt is not None + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + indexer_k_quant_and_cache_triton, + ) + + indexer_k_quant_and_cache_triton( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + ) + + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if has_prefill: + prefill_metadata = attn_metadata_narrowed.prefill + assert prefill_metadata is not None + + # Short sequences select every pool, so skip sparse scoring and fill + # the top-k buffer with all causal token indices. The index-K cache was + # already written above. + n_prefill_sf = num_tokens - num_decode_tokens + # Host-side short-prefill predicate: max_prefill_seq_len is computed + # in the metadata builder (exact for prefill rows) and equals + # positions[prefill_slice].max() + 1, so this replaces a + # positions.max().item() device sync per layer. -1 (unknown metadata) + # falls back to the device-side check. + if prefill_metadata.max_prefill_seq_len >= 0: + short_prefill = ( + n_prefill_sf > 0 + and positions is not None + and prefill_metadata.max_prefill_seq_len <= topk_tokens + ) + else: + short_prefill = ( + n_prefill_sf > 0 + and positions is not None + and int(positions[num_decode_tokens:num_tokens].max().item()) + 1 + <= topk_tokens + ) + if short_prefill: + # short_prefill is only True when positions is not None (above), + # but narrow explicitly for the indexer below. + assert positions is not None + _pos = positions[num_decode_tokens:num_tokens].to(torch.int32) + _buf = topk_indices_buffer[num_decode_tokens:num_tokens] + _fill_causal_indices(_buf, _pos) + + # Get the full shared workspace buffers once (will allocate on first use). + # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and + # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 + # scales) based on use_fp4_cache. + workspace_manager = current_workspace_manager() + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + values_spec, + scales_spec, + ) + for chunk in prefill_metadata.chunks if not short_prefill else (): + k_quant = k_quant_full[: chunk.total_seq_lens] + k_scale = k_scale_full[: chunk.total_seq_lens] + + if not chunk.skip_kv_gather: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, + ) + + cp_gather_indexer_k_quant_cache_triton( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.cu_seq_lens, + token_to_seq=chunk.token_to_seq, + ) + + q_slice = q_quant[chunk.token_start : chunk.token_end] + q_scale_slice = ( + q_scale[chunk.token_start : chunk.token_end] + if q_scale is not None + else None + ) + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_fp8_mqa_logits, + ) + + assert q_scale_slice is None + logits = rocm_fp8_mqa_logits( + q_slice_cast, + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + ) + num_rows = logits.shape[0] + + # kpool: logits are pool-granular (compress_ratio == index_kpool), + # so topk selects pools. We pick topk_tokens // kpool pools then + # expand each pool back to its kpool constituent tokens. + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens + if index_kpool > 1: + pool_topk = torch.empty( + (num_rows, select_k), dtype=torch.int32, device=logits.device + ) + topk_dst = pool_topk + else: + topk_dst = topk_indices_buffer[ + chunk.token_start : chunk.token_end, :topk_tokens + ] + + torch.ops._C.top_k_per_row_prefill( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + + if index_kpool > 1: + pool_ids = pool_topk.to(torch.int64) + if positions is not None: + # Fused expand-pools + append-tail into one Triton kernel + # (replaces ~25 elementwise ops). seq_len is token-granular + # (pos+1); the kernel derives pool_len internally. + q_seq = ( + positions[chunk.token_start : chunk.token_end].to(torch.int32) + + 1 + ) + expanded = kpool_ops.expand_pools_and_append_tail( + pool_ids, q_seq, index_kpool + ) + else: + valid = pool_ids >= 0 + expanded = kpool_ops.expand_pools_to_tokens( + pool_ids, valid, topk_tokens, index_kpool + ) + topk_indices_buffer[ + chunk.token_start : chunk.token_end, : expanded.shape[-1] + ] = expanded + + if has_decode: + decode_metadata = attn_metadata_narrowed.decode + assert decode_metadata is not None + kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes + kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) + + # Update the tail before reading logits; completed pools are compressed + # into the slot supplied by slot_mapping. + # Spec verification groups tokens by request and preserves position + # order so each token is stashed before the next completes its pool. + # Positions must remain token-granular because the kernel derives the + # pool phase and tail index from ``pos % kpool``. + if ( + index_kpool > 1 + and gate_score is not None + and compress_ape is not None + and positions is not None + and not skip_k_cache_insert + ): + num_requests = attn_metadata_narrowed.num_decodes + # Kpool writes must recover the original request grouping after the + # indexer's flattened decode path. Host metadata avoids a CUDA graph + # sync when choosing the uniform or padded layout. + per_req_lens = decode_metadata.per_req_decode_lens + if per_req_lens is not None: + use_uniform = ( + decode_metadata.decode_is_uniform + and num_decode_tokens + == num_requests * decode_metadata.write_max_decode_len + ) + group_lens = per_req_lens + lmax = decode_metadata.write_max_decode_len + else: + # Legacy metadata without per-request lens: fall back to the + # host-side requires_padding flag. Unreached now (per-request + # lens is always populated for decode), kept defensive. + use_uniform = not decode_metadata.requires_padding + group_lens = decode_metadata.decode_lens + lmax = int(decode_metadata.decode_lens.max().item()) + if not use_uniform: + # Non-uniform decode_lens (mixed plain-decode + spec-verify, or + # a variable MTP-verify batch): scatter actual tokens into a + # padded [B, lmax] layout. int32 tensors can't go through + # pack_seq_triton (float/uint8 only). The scatter indices are + # shared by all five scatters below (and the tail slot one). + scatter_idx = _build_decode_scatter_indices( + group_lens, num_requests, num_decode_tokens + ) + dec_k = _scatter_decode_tokens_by_request( + k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx + ) + dec_gate = _scatter_decode_tokens_by_request( + gate_score[:num_decode_tokens], + 0, + num_requests, + lmax, + scatter_idx, + ) + dec_slot = _scatter_decode_tokens_by_request( + slot_mapping[:num_decode_tokens], + -1, + num_requests, + lmax, + scatter_idx, + ) + dec_pos = _scatter_decode_tokens_by_request( + positions[:num_decode_tokens].to(torch.int32), + -1, + num_requests, + lmax, + scatter_idx, + ) + else: + next_n = num_decode_tokens // num_requests + shape2 = (num_requests, next_n) + dec_k = k[:num_decode_tokens].view(*shape2, head_dim) + dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim) + dec_slot = slot_mapping[:num_decode_tokens].view(shape2) + dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2) + tail_meta = ( + attn_metadata.get(_resolve_layer_name(tail_prefix)) + if tail_prefix is not None + else None + ) + # Paged tail cache replaces the transient _DECODE_TAIL ring. Group + # the tail group's token-granular slot_mapping per-request, mirroring + # dec_slot / dec_pos, so the kernel gets each request's current-token + # tail slot (block * kpool + pos % kpool). + if tail_meta is not None: + assert isinstance(tail_meta, DeepseekV32IndexerMetadata) + if tail_meta is None or tail_kv_cache is None: + dec_tail_slot = None + elif not use_uniform: + dec_tail_slot = _scatter_decode_tokens_by_request( + tail_meta.slot_mapping[:num_decode_tokens], + -1, + num_requests, + lmax, + scatter_idx, + ) + else: + dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2) + # The compress kernel writes the raw fp8 cache (not the quant view); + # pass the underlying kv_cache, not kv_cache_quant_view. + if dec_tail_slot is not None: + # Single batched launch over [num_requests, next_n] replaces the + # per-token sequential loop. The kernel iterates each request's + # tokens in position order internally, preserving the + # pool-completion read-after-stash dependency that the loop + # provided. Inputs are already grouped per request (uniform: + # view; non-uniform: _scatter_decode_tokens_by_request padded to + # [B, lmax]) — no per-token .contiguous() copies needed. + kpool_ops.kpool_decode_update_and_maybe_write_cache_batched( + kv_cache_raw, + tail_kv_cache, + dec_tail_slot, + dec_k, + dec_gate, + compress_ape, + dec_slot, + dec_pos, + index_kpool, + head_dim, + round_scale=(scale_fmt is not None), + ) + if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices( + topk_indices_buffer, + positions, + num_decode_tokens, + attn_metadata_narrowed.max_seq_len, + topk_tokens, + ): + return topk_indices_buffer + decode_lens = decode_metadata.decode_lens + if decode_metadata.requires_padding: + # Padding also covers short chunked prefills classified as decode. + # MXFP4 uses zero-byte padding so padded slots dequantize to zero. + if q_scale is not None: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens, pad_value=0 + ) + padded_q_scale = pack_seq_triton( + q_scale[:num_decode_tokens], decode_lens, pad_value=0 + ) + else: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens + ) + padded_q_scale = None + padded_weights = pack_seq_triton( + weights[:num_decode_tokens], decode_lens, pad_value=0 + ).reshape(-1, *weights.shape[1:]) + else: + padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_quant.shape[1:] + ) + if q_scale is not None: + padded_q_scale = q_scale[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_scale.shape[1:] + ) + else: + padded_q_scale = None + padded_weights = weights[:num_decode_tokens] + # TODO: move and optimize below logic with triton kernels + batch_size = padded_q_quant_decode_tokens.shape[0] + next_n = padded_q_quant_decode_tokens.shape[1] + num_padded_tokens = batch_size * next_n + seq_lens = decode_metadata.seq_lens[:batch_size] + # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) + # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; + # the downstream topk kernels accept both 1D and 2D. + padded_q_quant_cast = ( + padded_q_quant_decode_tokens.view(torch.int8) + if use_fp4_cache + else padded_q_quant_decode_tokens + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_fp8_paged_mqa_logits, + ) + + assert padded_q_scale is None + logits = rocm_fp8_paged_mqa_logits( + padded_q_quant_cast, + kv_cache, + padded_weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + ) + num_rows = logits.shape[0] + # kpool: logits are pool-granular -> select topk_tokens//kpool pools, + # then expand each pool back to its kpool tokens. + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens + if index_kpool > 1: + pool_topk = torch.empty( + (num_rows, select_k), dtype=torch.int32, device=logits.device + ) + topk_dst = pool_topk + else: + topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] + + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + + # Resolve to token-level indices in the output buffer. + if index_kpool > 1: + pool_ids = pool_topk.to(torch.int64) + n = pool_topk.shape[0] + # Decode seq_lens are pool-granular; recover token lengths from + # positions using the padded [B, next_n] row layout when needed. + if positions is not None: + dec_seq = _decode_topk_seq_lens( + positions, + decode_lens, + num_decode_tokens, + batch_size, + next_n, + decode_metadata.requires_padding, + ) + else: + dec_seq = decode_metadata.seq_lens[:n] + if dec_seq.ndim == 2: + dec_seq = dec_seq[:, -1] + dec_seq = dec_seq.to(torch.int32) + out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool) + else: + out = topk_dst + + if decode_metadata.requires_padding: + # Drop padded query rows introduced by the next_n padding above. + out = unpack_seq_triton( + out.reshape(batch_size, -1, out.shape[-1]), decode_lens + ) + topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out + + return topk_indices_buffer + + +@CustomOp.register("sparse_attn_indexer_kpool") +class SparseAttnIndexerKpool(CustomOp): + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a + separate custom op since it involves heavy custom kernels like `mqa_logits`, + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires + specific memory layout or implementation for different hardware backends to + achieve optimal performance. + + For now, the default native path will use CUDA backend path. Other platform + may requires add the corresponding Custom Op name `sparse_attn_indexer` to + `custom_ops` in `CompilationConfig` to enable the platform specific path. + """ + + def __init__( + self, + k_cache, + quant_block_size: int, + scale_fmt: str, + topk_tokens: int, + head_dim: int, + max_model_len: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool = False, + use_fp4_cache: bool = False, + tail_cache=None, + ): + super().__init__() + self.k_cache = k_cache + self.tail_cache = tail_cache + self.quant_block_size = quant_block_size + self.scale_fmt = scale_fmt + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_model_len = max_model_len + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.skip_k_cache_insert = skip_k_cache_insert + self.use_fp4_cache = use_fp4_cache + + def forward_native( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + *, + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + ): + assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" + assert isinstance(q_quant, torch.Tensor), ( + "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" + ) + if not rocm_aiter_ops.is_enabled(): + raise RuntimeError( + "Sparse attention indexer ROCm path is only supported on AITER. " + "Please enable aiter with VLLM_ROCM_USE_AITER=1" + ) + if index_kpool <= 1: + return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_quant, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=self.skip_k_cache_insert, + ) + return sparse_attn_indexer_kpool( + hidden_states, + self.k_cache.prefix, + self.k_cache.kv_cache, + q_quant, + None, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + self.skip_k_cache_insert, + self.use_fp4_cache, + gate_score, + compress_ape, + index_kpool, + positions, + self.tail_cache.kv_cache if self.tail_cache is not None else None, + self.tail_cache.prefix if self.tail_cache is not None else None, + ) diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index 0be749348494..6bd9ec44a152 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -23,7 +23,6 @@ from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope -from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool from vllm.model_executor.models.deepseek_v2 import ( DeepSeekV2FusedQkvAProjLinear, DeepseekV32IndexerCache, @@ -31,6 +30,7 @@ ) from vllm.model_executor.utils import maybe_disable_graph_partition from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8 +from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool from vllm.platforms import current_platform from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py new file mode 100644 index 000000000000..ad3845457063 --- /dev/null +++ b/vllm/models/glm5next/common/sparse_indexer.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared helpers for the glm5next sparse attention indexer (kpool) layers.""" + +import torch + +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + +# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. +MXFP4_BLOCK_SIZE = 32 + + +def _build_decode_scatter_indices( + decode_lens: torch.Tensor, + num_requests: int, + n: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token (request id, intra-request index) for a non-uniform decode + batch, with ``n == decode_lens.sum()`` as a host int (avoids a + device sync and keeps both repeat_interleaves sync-free). + + Shared by every ``_scatter_decode_tokens_by_request`` call in a step: + building it per call would repeat the same repeat_interleave/cumsum chain + up to 5x per layer on the eager decode break. + """ + device = decode_lens.device + dl = decode_lens.to(torch.int64) + req_id = torch.repeat_interleave( + torch.arange(num_requests, device=device, dtype=torch.int64), + dl, + output_size=n, + ) + req_starts = torch.cumsum( + torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]), + dim=0, + ) + # Broadcast the per-request start offsets to per-token (length n == + # dl.sum()) so each token's intra-request index subtracts its own + # request's start. + starts = torch.repeat_interleave(req_starts, dl, output_size=n) + intra = torch.arange(n, device=device, dtype=torch.int64) - starts + return req_id, intra + + +def _scatter_decode_tokens_by_request( + tokens: torch.Tensor, + pad_value, + num_requests: int, + lmax: int, + scatter_indices: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]`` + layout: request ``r``'s tokens at row ``r`` in order; short requests padded. + + Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32 + slot/pos tensors) — it scatters with the shared per-step indices from + ``_build_decode_scatter_indices``. Used only for the non-uniform + (``requires_padding``) decode batch; uniform batches use a zero-copy + reshape. + """ + req_id, intra = scatter_indices + out = torch.full( + (num_requests, lmax, *tokens.shape[1:]), + pad_value, + dtype=tokens.dtype, + device=tokens.device, + ) + out[req_id, intra] = tokens + return out + + +def _decode_topk_seq_lens( + positions: torch.Tensor, + decode_lens: torch.Tensor, + num_decode_tokens: int, + batch_size: int, + next_n: int, + requires_padding: bool, +) -> torch.Tensor: + """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware. + + ``pool_topk`` (and the logits it comes from) follow the padded + ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row + ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT + ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]`` + (the uniform-layout shortcut) misaligns every row after the first + non-uniform request and, past the decode region, reads prefill tokens' + positions; ``expand_pools_and_append_tail`` then anchors the tail at + another request's length, dropping the row's real tail tokens or emitting + indices past its sequence (out-of-bounds block-table reads). Padded rows + get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway. + """ + n = batch_size * next_n + if not requires_padding: + return positions[:n].to(torch.int32) + 1 + scatter_idx = _build_decode_scatter_indices( + decode_lens, batch_size, num_decode_tokens + ) + padded = _scatter_decode_tokens_by_request( + positions[:num_decode_tokens].to(torch.int32), + -1, + batch_size, + next_n, + scatter_idx, + ) + return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail + + +def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None: + causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32) + positions = positions.to(torch.int32) + rows[:] = causal_range[None, :] + rows[causal_range[None, :] > positions[:, None]] = -1 + + +def _fill_short_decode_causal_indices( + topk_indices_buffer: torch.Tensor, + positions: torch.Tensor | None, + num_decode_tokens: int, + max_seq_len: int, + topk_tokens: int, +) -> bool: + """Fill exact causal rows when sparse decode would select every token.""" + if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens: + return False + _fill_causal_indices( + topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens] + ) + return True + + +def _gather_workspace_shapes( + total_seq_lens: int, + head_dim: int, + fp8_dtype: torch.dtype, + use_fp4_cache: bool, +) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: + """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for + the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 + scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + if use_fp4_cache: + return ( + ((total_seq_lens, head_dim // 2), torch.uint8), + ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), + ) + return ( + ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, 4), torch.uint8), + ) + + +def kv_cache_as_quant_view( + kv_cache: torch.Tensor, + head_dim: int, + use_fp4_cache: bool, +) -> torch.Tensor: + """4D ``[num_blocks, block_size, 1, head_width]`` view expected by + DeepGEMM, from the 3D indexer kv-cache allocation.""" + if use_fp4_cache: + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 + num_blocks, block_size, _ = kv_cache.shape + page_bytes = int(kv_cache.stride(0)) + fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE + return torch.as_strided( + kv_cache, + size=(num_blocks, block_size, 1, fp4_bytes), + stride=(page_bytes, fp4_bytes, fp4_bytes, 1), + ) + return kv_cache.unsqueeze(-2) diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/models/glm5next/nvidia/sparse_indexer.py similarity index 67% rename from vllm/model_executor/layers/sparse_attn_indexer_kpool.py rename to vllm/models/glm5next/nvidia/sparse_indexer.py index 2a314042deef..1fc42c9e4114 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -2,30 +2,29 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Custom Sparse Attention Indexer layers.""" -from typing import TYPE_CHECKING - import torch import vllm.envs as envs -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config_or_none from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.models.glm5next.common.sparse_indexer import ( + RADIX_TOPK_WORKSPACE_SIZE, + _build_decode_scatter_indices, + _decode_topk_seq_lens, + _fill_causal_indices, + _fill_short_decode_causal_indices, + _gather_workspace_shapes, + _scatter_decode_tokens_by_request, + kv_cache_as_quant_view, +) +from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform - -if TYPE_CHECKING: - from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops -elif current_platform.is_rocm(): - from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops -else: - from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops - from vllm.utils.deep_gemm import has_deep_gemm from vllm.utils.torch_utils import ( LayerNameType, - _encode_layer_name, _resolve_layer_name, ) from vllm.v1.attention.backends.mla.indexer import ( @@ -36,16 +35,9 @@ if current_platform.is_cuda_alike(): from vllm import _custom_ops as ops -elif current_platform.is_xpu(): - from vllm._xpu_ops import xpu_ops logger = init_logger(__name__) -RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 - -# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. -MXFP4_BLOCK_SIZE = 32 - # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -96,166 +88,6 @@ def _kpool_compress_insert( ) -def _build_decode_scatter_indices( - decode_lens: torch.Tensor, - num_requests: int, - n: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Per-token (request id, intra-request index) for a non-uniform decode - batch, with ``n == decode_lens.sum()`` as a host int (avoids a - device sync and keeps both repeat_interleaves sync-free). - - Shared by every ``_scatter_decode_tokens_by_request`` call in a step: - building it per call would repeat the same repeat_interleave/cumsum chain - up to 5x per layer on the eager decode break. - """ - device = decode_lens.device - dl = decode_lens.to(torch.int64) - req_id = torch.repeat_interleave( - torch.arange(num_requests, device=device, dtype=torch.int64), - dl, - output_size=n, - ) - req_starts = torch.cumsum( - torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]), - dim=0, - ) - # Broadcast the per-request start offsets to per-token (length n == - # dl.sum()) so each token's intra-request index subtracts its own - # request's start. - starts = torch.repeat_interleave(req_starts, dl, output_size=n) - intra = torch.arange(n, device=device, dtype=torch.int64) - starts - return req_id, intra - - -def _scatter_decode_tokens_by_request( - tokens: torch.Tensor, - pad_value, - num_requests: int, - lmax: int, - scatter_indices: tuple[torch.Tensor, torch.Tensor], -) -> torch.Tensor: - """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]`` - layout: request ``r``'s tokens at row ``r`` in order; short requests padded. - - Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32 - slot/pos tensors) — it scatters with the shared per-step indices from - ``_build_decode_scatter_indices``. Used only for the non-uniform - (``requires_padding``) decode batch; uniform batches use a zero-copy - reshape. - """ - req_id, intra = scatter_indices - out = torch.full( - (num_requests, lmax, *tokens.shape[1:]), - pad_value, - dtype=tokens.dtype, - device=tokens.device, - ) - out[req_id, intra] = tokens - return out - - -def _decode_topk_seq_lens( - positions: torch.Tensor, - decode_lens: torch.Tensor, - num_decode_tokens: int, - batch_size: int, - next_n: int, - requires_padding: bool, -) -> torch.Tensor: - """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware. - - ``pool_topk`` (and the logits it comes from) follow the padded - ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row - ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT - ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]`` - (the uniform-layout shortcut) misaligns every row after the first - non-uniform request and, past the decode region, reads prefill tokens' - positions; ``expand_pools_and_append_tail`` then anchors the tail at - another request's length, dropping the row's real tail tokens or emitting - indices past its sequence (out-of-bounds block-table reads). Padded rows - get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway. - """ - n = batch_size * next_n - if not requires_padding: - return positions[:n].to(torch.int32) + 1 - scatter_idx = _build_decode_scatter_indices( - decode_lens, batch_size, num_decode_tokens - ) - padded = _scatter_decode_tokens_by_request( - positions[:num_decode_tokens].to(torch.int32), - -1, - batch_size, - next_n, - scatter_idx, - ) - return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail - - -def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None: - causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32) - positions = positions.to(torch.int32) - rows[:] = causal_range[None, :] - rows[causal_range[None, :] > positions[:, None]] = -1 - - -def _fill_short_decode_causal_indices( - topk_indices_buffer: torch.Tensor, - positions: torch.Tensor | None, - num_decode_tokens: int, - max_seq_len: int, - topk_tokens: int, -) -> bool: - """Fill exact causal rows when sparse decode would select every token.""" - if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens: - return False - _fill_causal_indices( - topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens] - ) - return True - - -def _gather_workspace_shapes( - total_seq_lens: int, - head_dim: int, - fp8_dtype: torch.dtype, - use_fp4_cache: bool, -) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: - """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for - the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 - scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + - (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" - if use_fp4_cache: - return ( - ((total_seq_lens, head_dim // 2), torch.uint8), - ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), - ) - return ( - ((total_seq_lens, head_dim), fp8_dtype), - ((total_seq_lens, 4), torch.uint8), - ) - - -def kv_cache_as_quant_view( - kv_cache: torch.Tensor, - head_dim: int, - use_fp4_cache: bool, -) -> torch.Tensor: - """4D ``[num_blocks, block_size, 1, head_width]`` view expected by - DeepGEMM, from the 3D indexer kv-cache allocation.""" - if use_fp4_cache: - assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 - num_blocks, block_size, _ = kv_cache.shape - page_bytes = int(kv_cache.stride(0)) - fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE - return torch.as_strided( - kv_cache, - size=(num_blocks, block_size, 1, fp4_bytes), - stride=(page_bytes, fp4_bytes, fp4_bytes, 1), - ) - return kv_cache.unsqueeze(-2) - - @eager_break_during_capture def sparse_attn_indexer_kpool( hidden_states: torch.Tensor, @@ -393,26 +225,13 @@ def sparse_attn_indexer_kpool( else: # standard: per-token fp8 quant + scatter (all tokens). assert scale_fmt is not None - if current_platform.is_rocm(): - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - indexer_k_quant_and_cache_triton, - ) - - indexer_k_quant_and_cache_triton( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, - ) - else: - ops.indexer_k_quant_and_cache( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, - ) + ops.indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + ) topk_indices_buffer[: hidden_states.shape[0]] = -1 if has_prefill: @@ -466,27 +285,13 @@ def sparse_attn_indexer_kpool( k_scale = k_scale_full[: chunk.total_seq_lens] if not chunk.skip_kv_gather: - if current_platform.is_rocm(): - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - cp_gather_indexer_k_quant_cache_triton, - ) - - cp_gather_indexer_k_quant_cache_triton( - kv_cache, - k_quant, - k_scale, - chunk.block_table, - chunk.cu_seq_lens, - token_to_seq=chunk.token_to_seq, - ) - else: - ops.cp_gather_indexer_k_quant_cache( - kv_cache, - k_quant, - k_scale, - chunk.block_table, - chunk.cu_seq_lens, - ) + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.cu_seq_lens, + ) q_slice = q_quant[chunk.token_start : chunk.token_end] q_scale_slice = ( @@ -504,30 +309,16 @@ def sparse_attn_indexer_kpool( q_slice_cast = q_slice k_quant_cast = k_quant k_scale_cast = k_scale.view(torch.float32).squeeze(-1) - if current_platform.is_rocm(): - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - rocm_fp8_mqa_logits, - ) - - assert q_scale_slice is None - logits = rocm_fp8_mqa_logits( - q_slice_cast, - (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - ) - else: - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits - - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - clean_logits=False, - ) + from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) num_rows = logits.shape[0] # kpool: logits are pool-granular (compress_ratio == index_kpool), @@ -544,28 +335,16 @@ def sparse_attn_indexer_kpool( chunk.token_start : chunk.token_end, :topk_tokens ] - if current_platform.is_xpu(): - xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined] - logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) - else: - torch.ops._C.top_k_per_row_prefill( - logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) + torch.ops._C.top_k_per_row_prefill( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) if index_kpool > 1: pool_ids = pool_topk.to(torch.int64) @@ -765,34 +544,18 @@ def sparse_attn_indexer_kpool( if use_fp4_cache else padded_q_quant_decode_tokens ) - if current_platform.is_rocm(): - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - rocm_fp8_paged_mqa_logits, - ) - - assert padded_q_scale is None - logits = rocm_fp8_paged_mqa_logits( - padded_q_quant_cast, - kv_cache, - padded_weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_model_len, - ) - else: - from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits + from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - padded_weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_model_len, - clean_logits=False, - ) + logits = fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + padded_weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + clean_logits=False, + ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools, # then expand each pool back to its kpool tokens. @@ -819,28 +582,16 @@ def sparse_attn_indexer_kpool( attn_metadata_narrowed.max_seq_len, ) else: - if current_platform.is_xpu(): - xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined] - logits, - next_n, - seq_lens, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) - else: - torch.ops._C.top_k_per_row_decode( - logits, - next_n, - seq_lens, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) # Resolve to token-level indices in the output buffer. if index_kpool > 1: @@ -942,33 +693,16 @@ def forward_native( index_kpool: int = 1, positions: torch.Tensor | None = None, ): - if current_platform.is_cuda() or current_platform.is_xpu(): - return self.forward_cuda( - hidden_states, - q_quant, - k, - weights, - gate_score=gate_score, - compress_ape=compress_ape, - index_kpool=index_kpool, - positions=positions, - ) - elif current_platform.is_rocm(): - return self.forward_hip( - hidden_states, - q_quant, - k, - weights, - gate_score=gate_score, - compress_ape=compress_ape, - index_kpool=index_kpool, - positions=positions, - ) - else: - raise NotImplementedError( - "SparseAttnIndexer native forward is only implemented for " - "CUDA, ROCm and XPU platforms." - ) + return self.forward_cuda( + hidden_states, + q_quant, + k, + weights, + gate_score=gate_score, + compress_ape=compress_ape, + index_kpool=index_kpool, + positions=positions, + ) def forward_cuda( self, @@ -1012,52 +746,3 @@ def forward_cuda( self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, ) - - def forward_hip( - self, - hidden_states: torch.Tensor, - q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - weights: torch.Tensor, - *, - gate_score: torch.Tensor | None = None, - compress_ape: torch.Tensor | None = None, - index_kpool: int = 1, - positions: torch.Tensor | None = None, - ): - assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" - assert isinstance(q_quant, torch.Tensor), ( - "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" - ) - if rocm_aiter_ops.is_enabled(): - if index_kpool <= 1: - return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( - hidden_states, - _encode_layer_name(self.k_cache.prefix), - self.k_cache.kv_cache, - q_quant, - k, - weights, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=self.skip_k_cache_insert, - ) - return self.forward_cuda( - hidden_states, - q_quant, - k, - weights, - gate_score=gate_score, - compress_ape=compress_ape, - index_kpool=index_kpool, - positions=positions, - ) - raise RuntimeError( - "Sparse attention indexer ROCm path is only supported on AITER. " - "Please enable aiter with VLLM_ROCM_USE_AITER=1" - ) diff --git a/vllm/models/glm5next/sparse_indexer.py b/vllm/models/glm5next/sparse_indexer.py new file mode 100644 index 000000000000..5aa0e47367f5 --- /dev/null +++ b/vllm/models/glm5next/sparse_indexer.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse attention indexer (kpool) — hardware-isolated entry point. + +The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers +in ``common/``); this module picks the right one for the current platform. +""" + +from vllm.platforms import current_platform + +if current_platform.is_rocm(): + from .amd.sparse_indexer import ( + SparseAttnIndexerKpool, + ) +else: + from .nvidia.sparse_indexer import ( + SparseAttnIndexerKpool, + ) + +__all__ = [ + "SparseAttnIndexerKpool", +] From 323a29d8f2b2a50ed77a6173f120822376a8b14a Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Wed, 16 Sep 2026 00:53:36 -0500 Subject: [PATCH 299/774] [Bugfix] Unreadable `prompt_embeds` payload should be a 400, not a 500 (#55451) Signed-off-by: lzhan011 <35493221+lzhan011@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Cyrus Leung --- .../test_prompt_embeds_payload_validation.py | 163 ++++++++++++++++++ .../test_sparse_tensor_concurrent_race.py | 11 +- .../test_sparse_tensor_validation.py | 14 +- vllm/renderers/embed_utils.py | 59 ++++++- 4 files changed, 230 insertions(+), 17 deletions(-) create mode 100644 tests/renderers/test_prompt_embeds_payload_validation.py diff --git a/tests/renderers/test_prompt_embeds_payload_validation.py b/tests/renderers/test_prompt_embeds_payload_validation.py new file mode 100644 index 000000000000..1e425805e1c0 --- /dev/null +++ b/tests/renderers/test_prompt_embeds_payload_validation.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""`prompt_embeds` payloads that are not tensor files must be client errors. + +`prompt_embeds` is a base64 blob the caller writes, and every other way of +getting it wrong -- wrong rank, wrong hidden_size, wrong dtype, not a tensor at +all -- already answers 400. A payload that `torch.load` cannot read at all was +the exception: torch has no single error type for that (the zip reader raises +`RuntimeError`, the legacy pickle path raises an unpickling error, `KeyError` or +`IndexError`, an empty payload raises `EOFError`), none of them is a +`ValueError`, and the entrypoints' fallback therefore mapped them to 500. + +The assertions go through `create_error_response`, which is what actually turns +an exception into a status code, so they pin the HTTP answer rather than the +exception class alone. +""" + +import io +import zipfile +from unittest.mock import patch + +import pybase64 as base64 +import pytest +import torch + +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) +from vllm.exceptions import VLLMValidationError +from vllm.renderers.embed_utils import ( + _MAX_EMBED_ERROR_REASON_CHARS, + safe_load_prompt_embeds, +) + + +@pytest.fixture +def model_config(): + from vllm.config import ModelConfig + + return ModelConfig( + model="facebook/opt-125m", + tokenizer="facebook/opt-125m", + tokenizer_mode="auto", + trust_remote_code=False, + dtype="float32", + seed=0, + enable_prompt_embeds=True, + ) + + +def _encode(raw: bytes) -> bytes: + return base64.b64encode(raw) + + +def _zip_without_a_tensor() -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("junk.txt", "hello") + return buffer.getvalue() + + +def _saved_tensor(tensor: torch.Tensor) -> bytes: + buffer = io.BytesIO() + torch.save(tensor, buffer) + return buffer.getvalue() + + +UNPARSABLE_PAYLOADS = [ + pytest.param(b"", id="empty"), + pytest.param(b"\x80", id="single-byte"), + pytest.param(b"\x00\x01\x02\x03not-a-tensor", id="random-bytes"), + pytest.param(b"hello world this is not a torch file", id="ascii-text"), + pytest.param(b"PK\x03\x04" + b"\x00" * 20, id="truncated-zip"), + pytest.param(_zip_without_a_tensor(), id="zip-without-a-tensor"), + # pickle.dumps(42): a valid pickle stream that is not a torch checkpoint. + # Spelled out rather than built, so this file does not import pickle. + pytest.param(b"\x80\x04K*.", id="legacy-pickle-of-an-int"), +] + + +@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS) +def test_unreadable_prompt_embeds_payload_is_a_client_error(model_config, payload): + with pytest.raises(VLLMValidationError) as excinfo: + safe_load_prompt_embeds(model_config, _encode(payload)) + + assert excinfo.value.parameter == "prompt_embeds" + assert int(create_error_response(excinfo.value).error.code) == 400 + + +@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS) +def test_the_error_body_stays_bounded(model_config, payload): + """torch's reason is built from the caller's bytes, so it must be truncated. + + An unpickler that reports the offending global, for instance, quotes a name + the caller chose; without a bound a large one would be reflected whole. + """ + padding = b"A" * 100_000 + with pytest.raises(VLLMValidationError) as excinfo: + safe_load_prompt_embeds(model_config, _encode(payload + padding)) + + assert len(str(excinfo.value)) < 500 + + +def test_a_long_reason_is_truncated_and_says_how_much_was_left_out(model_config): + """The cap is pinned against a reason of known length, not a real payload. + + Which of torch's messages happen to exceed the cap is torch's business and + changes between releases, so the reason is injected here: that keeps the + assertion about *this* code's truncation rule rather than about whichever + error text the installed torch produces. + """ + reason = "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 37) + + with ( + patch( + "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason) + ), + pytest.raises(VLLMValidationError) as excinfo, + ): + safe_load_prompt_embeds(model_config, _encode(b"ignored")) + + message = str(excinfo.value) + assert "E" * _MAX_EMBED_ERROR_REASON_CHARS in message + assert "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 1) not in message + assert "... (37 more characters truncated)" in message + + +def test_a_short_reason_is_passed_through_whole(model_config): + """A reason that fits is not annotated -- the count only appears when it is + telling the reader something.""" + reason = "invalid magic number; corrupt file?" + assert len(reason) <= _MAX_EMBED_ERROR_REASON_CHARS + + with ( + patch( + "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason) + ), + pytest.raises(VLLMValidationError) as excinfo, + ): + safe_load_prompt_embeds(model_config, _encode(b"ignored")) + + message = str(excinfo.value) + assert reason in message + assert "truncated" not in message + + +def test_a_well_formed_payload_still_loads(model_config): + """Positive control: the healthy path is untouched.""" + hidden_size = model_config.get_hidden_size() + tensor = torch.zeros(3, hidden_size, dtype=torch.float32) + + loaded = safe_load_prompt_embeds(model_config, _encode(_saved_tensor(tensor))) + + assert loaded.shape == (3, hidden_size) + assert loaded.dtype == model_config.dtype + + +def test_a_payload_that_is_not_base64_is_still_a_client_error(model_config): + """Positive control: the base64 stage already answered 400 and still does.""" + with pytest.raises(ValueError) as excinfo: + safe_load_prompt_embeds(model_config, b"!!!!not base64!!!!") + + assert int(create_error_response(excinfo.value).error.code) == 400 diff --git a/tests/renderers/test_sparse_tensor_concurrent_race.py b/tests/renderers/test_sparse_tensor_concurrent_race.py index 22925a9f1961..a1683eee0bff 100644 --- a/tests/renderers/test_sparse_tensor_concurrent_race.py +++ b/tests/renderers/test_sparse_tensor_concurrent_race.py @@ -11,6 +11,7 @@ import torch from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO from vllm.renderers.embed_utils import safe_load_prompt_embeds from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe @@ -54,7 +55,7 @@ class TestNegativeControlWithoutRace: def test_malicious_sparse_rejected_by_prompt_loader(self, model_config): encoded = _encode_tensor(_create_malicious_sparse_tensor()) - with pytest.raises((RuntimeError, ValueError)): + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)): safe_load_prompt_embeds(model_config, encoded) def test_malicious_sparse_rejected_by_image_loader(self): @@ -147,7 +148,7 @@ def attempt_load(_): try: safe_load_prompt_embeds(model_config, malicious_encoded) bypassed.append(True) - except (RuntimeError, ValueError): + except (RuntimeError, ValueError, VLLMValidationError): rejected.append(True) with ThreadPoolExecutor(max_workers=num_workers) as pool: @@ -192,7 +193,7 @@ def test_flag_restored_after_exception(self, model_config): initial = torch.sparse.check_sparse_tensor_invariants.is_enabled() malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor()) - with pytest.raises((RuntimeError, ValueError)): + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)): safe_load_prompt_embeds(model_config, malicious_encoded) assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial @@ -203,7 +204,7 @@ def test_flag_restored_after_concurrent_exceptions(self, model_config): malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor()) def attempt_load(_): - with contextlib.suppress(RuntimeError, ValueError): + with contextlib.suppress(RuntimeError, ValueError, VLLMValidationError): safe_load_prompt_embeds(model_config, malicious_encoded) with ThreadPoolExecutor(max_workers=4) as pool: @@ -244,7 +245,7 @@ def load_image(): # serialization — one finishes before the other starts), AND invalid # tensors are still rejected. malicious = _encode_tensor(_create_malicious_sparse_tensor()) - with pytest.raises((RuntimeError, ValueError)): + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)): safe_load_prompt_embeds(model_config, malicious) with pytest.raises((RuntimeError, ValueError)): io_handler.load_base64("", malicious.decode("utf-8")) diff --git a/tests/renderers/test_sparse_tensor_validation.py b/tests/renderers/test_sparse_tensor_validation.py index c119b090c0d0..e49ac560dc71 100644 --- a/tests/renderers/test_sparse_tensor_validation.py +++ b/tests/renderers/test_sparse_tensor_validation.py @@ -117,8 +117,10 @@ def test_malicious_sparse_tensor_rejected(self, model_config): malicious_tensor = _create_malicious_sparse_tensor() encoded = _encode_tensor(malicious_tensor) - # Should raise RuntimeError due to invalid sparse tensor - with pytest.raises((RuntimeError, ValueError)) as exc_info: + # The invariant check raises RuntimeError from inside torch.load; + # safe_load_prompt_embeds now reports that as a client error so the + # caller gets a 400 rather than the 500 a bare RuntimeError produced. + with pytest.raises(VLLMValidationError) as exc_info: safe_load_prompt_embeds(model_config, encoded) # Error should indicate sparse tensor validation failure @@ -137,7 +139,7 @@ def test_extremely_large_indices_rejected(self, model_config): ) encoded = _encode_tensor(malicious_tensor) - with pytest.raises((RuntimeError, ValueError)): + with pytest.raises(VLLMValidationError): safe_load_prompt_embeds(model_config, encoded) def test_negative_indices_rejected(self, model_config): @@ -152,7 +154,7 @@ def test_negative_indices_rejected(self, model_config): ) encoded = _encode_tensor(malicious_tensor) - with pytest.raises((RuntimeError, ValueError)): + with pytest.raises(VLLMValidationError): safe_load_prompt_embeds(model_config, encoded) def test_hidden_size_mismatch_rejected(self, model_config): @@ -379,8 +381,8 @@ def test_attack_scenario_completions_api(self, model_config): # Step 1-2: Attacker creates malicious payload attack_payload = _encode_tensor(_create_malicious_sparse_tensor()) - # Step 3-4: Server processes and should reject - with pytest.raises((RuntimeError, ValueError)): + # Step 3-4: Server processes and should reject, as a client error + with pytest.raises(VLLMValidationError): safe_load_prompt_embeds(model_config, attack_payload) def test_attack_scenario_chat_api_image(self): diff --git a/vllm/renderers/embed_utils.py b/vllm/renderers/embed_utils.py index c4079005a54a..981aeee3fb09 100644 --- a/vllm/renderers/embed_utils.py +++ b/vllm/renderers/embed_utils.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from io import BytesIO -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final import pybase64 import torch @@ -17,6 +17,27 @@ from vllm.config import ModelConfig +# How much of torch's own reason is repeated back to the caller. The reason is +# built from bytes the caller sent, so it is bounded rather than echoed whole. +_MAX_EMBED_ERROR_REASON_CHARS: Final = 200 + + +def _truncated_reason(exc: Exception) -> str: + """`exc`'s message, capped, saying how much was left out when it is capped. + + Without the count a capped reason reads like the whole one, so a caller + debugging a large payload cannot tell that torch said more than this. + """ + reason: Final = str(exc).strip() or type(exc).__name__ + omitted: Final = len(reason) - _MAX_EMBED_ERROR_REASON_CHARS + if omitted <= 0: + return reason + return ( + f"{reason[:_MAX_EMBED_ERROR_REASON_CHARS]}" + f"... ({omitted} more characters truncated)" + ) + + def safe_load_prompt_embeds( model_config: "ModelConfig", embed: bytes, @@ -27,12 +48,38 @@ def safe_load_prompt_embeds( parameter="prompt_embeds", ) + # Decoded outside the try below so that a body which is not base64 keeps + # raising binascii.Error, which is a ValueError and already answered 400. + payload: Final = pybase64.b64decode(embed, validate=True) + with check_sparse_tensor_invariants_threadsafe(): - tensor = torch.load( - BytesIO(pybase64.b64decode(embed, validate=True)), - weights_only=True, - map_location=torch.device("cpu"), - ) + try: + tensor = torch.load( + BytesIO(payload), + weights_only=True, + map_location=torch.device("cpu"), + ) + except (torch.OutOfMemoryError, MemoryError): + # Running out of memory is the server's condition, not the caller's. + raise + except Exception as exc: + # Everything else here is about the bytes the caller sent, and torch + # has no single error type for "this is not a tensor file": the zip + # reader raises `RuntimeError`, the legacy pickle path raises + # `UnpicklingError`, `KeyError` or `IndexError`, an empty payload + # raises `EOFError`. None of them is a `ValueError`, so the + # entrypoints' fallback mapped them all to 500. Catching by + # behaviour rather than by type also survives torch changing which + # error it raises for a given malformed payload. + # + # torch's reason is worth keeping -- for a malformed sparse tensor + # it names the offending index -- but it is built from the caller's + # own bytes, so it is truncated rather than echoed whole. + raise VLLMValidationError( + "`prompt_embeds` could not be deserialized as a torch tensor: " + f"{_truncated_reason(exc)}", + parameter="prompt_embeds", + ) from exc tensor = safe_to_dense(tensor, parameter="prompt_embeds") if tensor.dim() > 2: From b3b13c12923f0b68f36ee515ebbc3efc7679eb4f Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:57:40 -0700 Subject: [PATCH 300/774] [ROCm][CI] Fix MLA RoPE fused-kernel tests for TheRock image (#57112) Signed-off-by: Matthew Wong Co-authored-by: Andreas Karatzas --- .../test_rotary_embedding_mla_cache_fused.py | 137 ++++++++++++++++-- 1 file changed, 121 insertions(+), 16 deletions(-) diff --git a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py index 2a2525c307bc..1ac17335b69c 100644 --- a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py +++ b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py @@ -17,6 +17,88 @@ from vllm.utils.torch_utils import set_random_seed +def assert_rope_close_rocm( + actual: torch.Tensor, + ref: torch.Tensor, + rope_input: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + is_neox: bool, + *, + max_fail_fraction: float | None = None, + atol: float | None = None, + rtol: float | None = None, +) -> None: + """assert_close for RoPE fused-vs-reference on ROCm, with a per-element + product-ULP allowance for catastrophic-cancellation elements. + + The fused and reference kernels may contract the RoPE multiply-add into an + FMA with either product pre-rounded to the working dtype. The two valid + orderings differ by at most 2 ULPs of the product at that element, which + cancellation (x*cos ~= -y*sin) can amplify past the default tolerance. + + Tier 1: |actual - ref| <= atol + rtol * |ref| (default tolerance). + Tier 2: elements failing tier 1 must each satisfy + |actual_e - ref_e| <= 2 * eps * |product_e|, where product_e is the larger + of the two products combined at that element, and their fraction must be + <= max_fail_fraction (1e-4 for fp16; 1e-2 for bf16, whose coarser eps + puts 1-ULP divergences past the default tolerance far more often). + """ + if atol is None: + atol = get_default_atol(ref) + if rtol is None: + rtol = get_default_rtol(ref) + if max_fail_fraction is None: + max_fail_fraction = 1e-2 if ref.dtype == torch.bfloat16 else 1e-4 + eps = torch.finfo(ref.dtype).eps + + # per-element product magnitude, in the (seq, heads, rot_dim) layout + seq_len, rot_dim = rope_input.shape[0], rope_input.shape[-1] + half = rot_dim // 2 + inp = rope_input.reshape(seq_len, -1, rot_dim).to(torch.float32) + cs = cos_sin_cache[positions].to(torch.float32) # (seq_len, 2*half) + cos = cs[:, :half].unsqueeze(1) # broadcast over heads + sin = cs[:, half:].unsqueeze(1) + x, y = inp[..., :half], inp[..., half:] + if is_neox: + # out = [x*cos - y*sin, y*cos + x*sin] + prod = torch.cat( + [ + torch.maximum((x * cos).abs(), (y * sin).abs()), + torch.maximum((y * cos).abs(), (x * sin).abs()), + ], + dim=-1, + ) + else: + # out[2i] = x[2i]*cos[i] - x[2i+1]*sin[i], interleaved + xe, xo = inp[..., 0::2], inp[..., 1::2] + prod = torch.stack( + [ + torch.maximum((xe * cos).abs(), (xo * sin).abs()), + torch.maximum((xo * cos).abs(), (xe * sin).abs()), + ], + dim=-1, + ).flatten(-2) + prod = prod.reshape(rope_input.shape) + + diff = (actual.to(torch.float32) - ref.to(torch.float32)).abs() + tol = atol + rtol * ref.to(torch.float32).abs() + fail = diff > tol + n_fail = int(fail.sum()) + if n_fail == 0: + return + assert n_fail <= max_fail_fraction * diff.numel(), ( + f"{n_fail} elements beyond default tolerance " + f"({n_fail / diff.numel():.2%} > {max_fail_fraction:.2%}); " + f"max diff {diff.max().item()}" + ) + excess = diff[fail] - 2 * eps * prod[fail] + assert (excess <= 0).all(), ( + f"{int((excess > 0).sum())} cancellation elements exceed 2 product-ULPs; " + f"worst excess {excess.max().item()}" + ) + + @pytest.fixture def default_vllm_config(monkeypatch): """Enable the AITER triton rope on ROCm for fp16-consistent numerics. @@ -99,6 +181,8 @@ def test_concat_and_cache_mla_rope_fused( k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device) kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device) + k_pe_orig = k_pe.clone() + query_orig = query.clone() if current_platform.is_rocm(): # We use forward_hip for the same numerics as the fused custom kernel on ROCm @@ -181,12 +265,9 @@ def test_concat_and_cache_mla_rope_fused( kv_cache_scale, ) - # On ROCm the AITER Triton rope diverges by ~1 ULP from the fused kernel, - # which the tight CUDA tolerances don't cover for the low-precision paths: - # - bf16: up to ~1 bf16 ULP (0.0156-0.03125 at values ~2-3), so relax the - # kv-cache atol (bounded ~0.02) and the query atol (~0.04). - # - neox-style fp8: one e4m3 ULP (~12.5%), so relax rtol. - # Other paths use the CUDA defaults. + # On ROCm the AITER Triton rope diverges by up to 2 product-ULPs from the + # fused kernel, which the default tolerance misses at the rare + # catastrophic-cancellation elements. fp8 keeps its calibrated tolerance. rocm_neox = current_platform.is_rocm() and is_neox_style rocm_bf16 = current_platform.is_rocm() and dtype == torch.bfloat16 if kv_cache_dtype == "fp8": @@ -207,16 +288,40 @@ def test_concat_and_cache_mla_rope_fused( atol=0.004 if rocm_bf16 else 0.001, rtol=0.15 if rocm_neox or rocm_bf16 else 0.1, ) - elif rocm_bf16: - torch.testing.assert_close(kv_cache, ref_kv_cache, atol=0.02, rtol=1e-3) - elif rocm_neox: - torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3) + elif current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16): + # gather the rope region of the cache back to token order + block_idx = slot_mapping // block_size + block_offset = slot_mapping % block_size + fused_k_rope = kv_cache[block_idx, block_offset, kv_lora_rank:] + ref_k_rope_gathered = ref_kv_cache[block_idx, block_offset, kv_lora_rank:] + assert_rope_close_rocm( + fused_k_rope, + ref_k_rope_gathered, + k_pe_orig, + rope.cos_sin_cache, + positions, + is_neox_style, + ) + # the nope region is a byte-exact copy + torch.testing.assert_close( + kv_cache[..., :kv_lora_rank], ref_kv_cache[..., :kv_lora_rank] + ) else: torch.testing.assert_close(kv_cache, ref_kv_cache) - torch.testing.assert_close( - query, - ref_q_pe, - atol=0.04 if rocm_bf16 else get_default_atol(query), - rtol=get_default_rtol(query), - ) + if current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16): + assert_rope_close_rocm( + query, + ref_q_pe, + query_orig, + rope.cos_sin_cache, + positions, + is_neox_style, + ) + else: + torch.testing.assert_close( + query, + ref_q_pe, + atol=get_default_atol(query), + rtol=get_default_rtol(query), + ) From e0301afb2a1364fe594a52d9ce649264a62fa243 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Tue, 15 Sep 2026 23:07:06 -0700 Subject: [PATCH 301/774] [Bugfix][Responses] Clean up MCP tool sessions once before closing them (#56988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- .../responses/test_serving_responses.py | 84 ++++++++++++++++++- vllm/entrypoints/openai/responses/context.py | 8 +- 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 041e517e29dd..bd41b8cda282 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -41,6 +41,7 @@ from vllm.entrypoints.openai.responses.context import ( ConversationContext, HarmonyContext, + ParsableContext, SimpleContext, ) from vllm.entrypoints.openai.responses.protocol import ( @@ -551,6 +552,87 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs): class TestInitializeToolSessions: """Test class for _initialize_tool_sessions method""" + @pytest.fixture(params=[ParsableContext, HarmonyContext]) + def tool_context(self, request): + context = request.param.__new__(request.param) + context.available_tools = ["browser", "python"] + context._tool_sessions = {} + context.called_tools = set() + return context + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("available_tools", "called_tools"), + [ + ([], []), + (["browser"], ["browser"]), + (["browser", "python"], []), + (["browser", "python"], ["browser"]), + (["browser", "python"], ["python"]), + (["browser", "python"], ["browser", "python"]), + ], + ids=["empty", "single", "unused", "first", "last", "both"], + ) + async def test_mcp_sessions_cleaned_once_before_close( + self, tool_context, available_tools, called_tools + ): + events = [] + sessions = {} + + @asynccontextmanager + async def new_session(name, *_args): + session = MagicMock() + session.call_tool = AsyncMock( + side_effect=lambda *_: events.append(("cleanup", name)) + ) + sessions[name] = session + try: + yield session + finally: + events.append(("close", name)) + + tool_server = MagicMock(spec=ToolServer) + tool_server.new_session.side_effect = new_session + tool_context.available_tools = available_tools + async with AsyncExitStack() as stack: + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + # Reinitializing existing sessions must not duplicate cleanup either. + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + tool_context.called_tools.update(called_tools) + + for name, session in sessions.items(): + if name in called_tools: + session.call_tool.assert_awaited_once_with("cleanup_session", {}) + assert events.index(("cleanup", name)) < events.index(("close", name)) + else: + session.call_tool.assert_not_awaited() + assert [name for event, name in events if event == "close"] == list( + reversed(available_tools) + ) + + @pytest.mark.asyncio + async def test_mcp_partial_initialization_closes_opened_session(self, tool_context): + closed = [] + session = MagicMock(call_tool=AsyncMock()) + + @asynccontextmanager + async def new_session(name, *_args): + if name == "python": + raise RuntimeError("session initialization failed") + try: + yield session + finally: + closed.append(name) + + tool_server = MagicMock(spec=ToolServer) + tool_server.new_session.side_effect = new_session + with pytest.raises(RuntimeError, match="session initialization failed"): + async with AsyncExitStack() as stack: + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + + assert closed == ["browser"] + session.call_tool.assert_not_awaited() + @pytest_asyncio.fixture async def serving_responses_instance(self): """Create a real OpenAIServingResponses instance for testing""" diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 2268bb0700a0..6b93d39b8803 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -565,6 +565,7 @@ async def init_tool_sessions( mcp_tools: dict[str, Mcp], ): if tool_server: + initialized_session = False for tool_name in self.available_tools: if tool_name in self._tool_sessions: continue @@ -577,6 +578,8 @@ async def init_tool_sessions( tool_server.new_session(tool_name, request_id, headers) ) self._tool_sessions[tool_name] = tool_session + initialized_session = True + if initialized_session: exit_stack.push_async_exit(self.cleanup_session) async def cleanup_session(self, *args, **kwargs) -> None: @@ -863,6 +866,7 @@ async def init_tool_sessions( mcp_tools: dict[str, Mcp], ): if tool_server: + initialized_session = False for tool_name in self.available_tools: if tool_name not in self._tool_sessions: tool_type = _map_tool_name_to_tool_type(tool_name) @@ -873,7 +877,9 @@ async def init_tool_sessions( tool_server.new_session(tool_name, request_id, headers) ) self._tool_sessions[tool_name] = tool_session - exit_stack.push_async_exit(self.cleanup_session) + initialized_session = True + if initialized_session: + exit_stack.push_async_exit(self.cleanup_session) async def call_container_tool( self, tool_session: Union["ClientSession", Tool], last_msg: Message From 0f8fa53acf43a9a6a7a1f7e6b196950c78cd6d7a Mon Sep 17 00:00:00 2001 From: AlexHuang Date: Wed, 16 Sep 2026 14:17:50 +0800 Subject: [PATCH 302/774] [Bugfix][KV Offload] Check cgroup memory before SHM allocation (#54014) Signed-off-by: Alex Signed-off-by: AlexHuang Signed-off-by: Alex Co-authored-by: OpenAI Codex Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- tests/distributed/test_shm_broadcast.py | 35 +++- tests/utils_/test_cpu_resource_utils.py | 189 ++++++++++++++++++ .../cpu/test_shared_offload_region.py | 11 +- .../device_communicators/shm_broadcast.py | 40 ++-- .../model_loader/weight_utils.py | 8 +- vllm/utils/cpu_resource_utils.py | 89 ++++++++- .../kv_offload/cpu/shared_offload_region.py | 5 +- 7 files changed, 343 insertions(+), 34 deletions(-) create mode 100644 tests/utils_/test_cpu_resource_utils.py diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index c3833b6b9f82..ef7f523b78dc 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -694,14 +694,43 @@ def test_check_shm_free_space_raises_when_insufficient(tmp_path): def test_check_shm_free_space_passes_when_sufficient(tmp_path): - with mock.patch.object( - shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(512 << 20) + with ( + mock.patch.object( + shm_broadcast.shutil, + "disk_usage", + return_value=_fake_disk_usage(512 << 20), + ), + mock.patch.object(shm_broadcast, "check_cgroup_memory_available"), ): check_shm_free_space(240 << 20, shm_path=str(tmp_path)) def test_check_shm_free_space_skipped_when_path_missing(tmp_path): - check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) + with mock.patch.object(shm_broadcast, "check_cgroup_memory_available"): + check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) + + +def test_check_shm_free_space_checks_cgroup(tmp_path): + with ( + mock.patch.object( + shm_broadcast.shutil, + "disk_usage", + return_value=_fake_disk_usage(512 << 20), + ), + mock.patch.object( + shm_broadcast, "check_cgroup_memory_available" + ) as check_cgroup, + ): + check_shm_free_space( + 240 << 20, + shm_path=str(tmp_path), + allocation_name="SHM mmap", + ) + + check_cgroup.assert_called_once_with( + 240 << 20, + "SHM mmap", + ) def test_shm_ring_buffer_creation_checks_free_space(): diff --git a/tests/utils_/test_cpu_resource_utils.py b/tests/utils_/test_cpu_resource_utils.py new file mode 100644 index 000000000000..ff314d2231d2 --- /dev/null +++ b/tests/utils_/test_cpu_resource_utils.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cgroup memory readers and SHM allocation preflight checks.""" + +import sys +from io import StringIO +from unittest import mock + +import pytest + +from vllm.utils import cpu_resource_utils as cru +from vllm.utils.mem_constants import GiB_bytes + +_V2_LIMIT_PATH = "/sys/fs/cgroup/memory.max" +_V2_USAGE_PATH = "/sys/fs/cgroup/memory.current" +_V1_LIMIT_PATH = "/sys/fs/cgroup/memory/memory.limit_in_bytes" +_V1_USAGE_PATH = "/sys/fs/cgroup/memory/memory.usage_in_bytes" + + +@pytest.fixture(autouse=True) +def _clear_cgroup_cache(): + cru.get_cgroup_memory_limit.cache_clear() + yield + cru.get_cgroup_memory_limit.cache_clear() + + +def _stub_files(monkeypatch, files: dict): + """Stub ``open()`` for a fixed set of paths; ``None`` -> OSError.""" + real_open = open + + def fake_open(path, *args, **kwargs): + if path in files: + content = files[path] + if content is None: + raise OSError(f"no such file: {path}") + return StringIO(content) + return real_open(path, *args, **kwargs) + + monkeypatch.setattr("builtins.open", fake_open) + + +def test_cgroup_v2_limit_and_usage(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: f"{20 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{5 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 20 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 5 * GiB_bytes + + +def test_cgroup_v2_unlimited_falls_back_to_v1(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: "max\n", + _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n", + _V1_USAGE_PATH: f"{1 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 1 * GiB_bytes + + +def test_cgroup_v1_limit_and_usage(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: None, + _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n", + _V1_USAGE_PATH: f"{2 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 2 * GiB_bytes + + +def test_cgroup_v1_unlimited_sentinel_is_ignored(monkeypatch): + """An unlimited cgroup v1 sentinel must not be treated as a real limit.""" + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: f"{(1 << 63) - 1}\n"}, + ) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_no_limit_files_present(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files(monkeypatch, {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: None}) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_skipped_on_non_linux(monkeypatch): + monkeypatch.setattr(sys, "platform", "darwin") + _stub_files(monkeypatch, {_V2_LIMIT_PATH: f"{1 * GiB_bytes}\n"}) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_limit_is_cached_until_explicitly_cleared(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + files = { + _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{512 << 20}\n", + } + _stub_files(monkeypatch, files) + + assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes + files[_V2_LIMIT_PATH] = f"{2 * GiB_bytes}\n" + assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes + + cru.get_cgroup_memory_limit.cache_clear() + assert cru.get_cgroup_memory_limit() == 2 * GiB_bytes + + +def test_cgroup_usage_is_read_without_cache(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + files = { + _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{512 << 20}\n", + } + _stub_files(monkeypatch, files) + + assert cru.get_cgroup_memory_usage() == 512 << 20 + files[_V2_USAGE_PATH] = f"{768 << 20}\n" + assert cru.get_cgroup_memory_usage() == 768 << 20 + + +def test_check_cgroup_memory_available_warns_on_low_headroom(monkeypatch): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(600 << 20, "mmap") + + log_debug.assert_not_called() + log_warning.assert_called_once() + assert "current headroom is below" in log_warning.call_args.args[-1] + + +def test_check_cgroup_memory_available_logs_success_at_debug(monkeypatch): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(256 << 20, "mmap") + + log_debug.assert_called_once() + log_warning.assert_not_called() + + +@pytest.mark.parametrize( + ("limit", "usage"), + [(None, 512 << 20), (1 << 30, None)], +) +def test_check_cgroup_memory_available_skips_without_snapshot( + monkeypatch, limit, usage +): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: limit) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: usage) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(1 << 60, "mmap") + + log_debug.assert_not_called() + log_warning.assert_not_called() diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index c448e76a9c51..4f22d122dbec 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -906,12 +906,13 @@ def test_insufficient_space_raises_clear_error(monkeypatch): monkeypatch.setattr(region.os, "open", mock_open) monkeypatch.setattr(region.os, "unlink", mock_unlink) monkeypatch.setattr(region.os, "close", mock_close) + mock_check = MagicMock( + side_effect=RuntimeError("Insufficient space in /dev/shm: 30 GB required.") + ) monkeypatch.setattr( region, "check_shm_free_space", - lambda *a, **kw: (_ for _ in ()).throw( - RuntimeError("Insufficient space in /dev/shm: 30 GB required.") - ), + mock_check, ) with pytest.raises(RuntimeError, match="Insufficient space"): @@ -925,6 +926,10 @@ def test_insufficient_space_raises_clear_error(monkeypatch): mock_unlink.assert_called_once_with(mmap_path) mock_close.assert_called_once_with(9999) + mock_check.assert_called_once_with( + 4 * PAGE_SIZE, + allocation_name="CPU KV offload shared region in /dev/shm", + ) def test_ftruncate_failure_cleans_up_creator(monkeypatch): diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 6caacea6ec5d..2171f79fb7ce 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -34,6 +34,7 @@ from vllm.distributed.utils import StatelessProcessGroup, sched_yield from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.cpu_resource_utils import check_cgroup_memory_available from vllm.utils.network_utils import ( get_ip, get_open_zmq_inproc_path, @@ -224,26 +225,37 @@ def notify(self): SHM_PATH = "/dev/shm" -def check_shm_free_space(required_bytes: int, shm_path: str = SHM_PATH) -> None: - """Raise if ``shm_path`` cannot fit a ``required_bytes`` shared segment. +def check_shm_free_space( + required_bytes: int, + shm_path: str = SHM_PATH, + *, + allocation_name: str = "shared-memory allocation", +) -> None: + """Raise if SHM cannot fit a shared segment and log cgroup headroom. Args: required_bytes: Size of the shared-memory segment to be created. - shm_path: Mount point backing POSIX shared memory; skipped if absent. + shm_path: Mount point backing POSIX shared memory; its filesystem + check is skipped if absent. + allocation_name: Human-readable name used in errors and logs. Raises: - RuntimeError: If ``required_bytes`` exceeds the free space. + RuntimeError: If the SHM filesystem has insufficient space. """ - if not os.path.isdir(shm_path): - return - free_bytes = shutil.disk_usage(shm_path).free - if required_bytes <= free_bytes: - return - mib = 1 << 20 - raise RuntimeError( - f"Insufficient space in {shm_path}: {required_bytes / mib:.0f} MiB " - f"required, {free_bytes / mib:.0f} MiB free. Increase {shm_path} " - "(e.g. --shm-size or --ipc=host)." + if os.path.isdir(shm_path): + free_bytes = shutil.disk_usage(shm_path).free + if required_bytes > free_bytes: + mib = 1 << 20 + raise RuntimeError( + f"Insufficient space in {shm_path} for {allocation_name}: " + f"{required_bytes / mib:.0f} MiB required, " + f"{free_bytes / mib:.0f} MiB free. Increase {shm_path} " + "(e.g. --shm-size or --ipc=host)." + ) + + check_cgroup_memory_available( + required_bytes, + allocation_name, ) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index abe24fa06eca..d291f1929ca0 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -727,11 +727,15 @@ def _get_available_ram_bytes() -> int: host_available = psutil.virtual_memory().available - from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit + from vllm.utils.cpu_resource_utils import ( + get_cgroup_memory_limit, + get_cgroup_memory_usage, + ) - cgroup_limit, cgroup_usage = get_cgroup_memory_limit() + cgroup_limit = get_cgroup_memory_limit() if cgroup_limit is None: return host_available + cgroup_usage = get_cgroup_memory_usage() cgroup_available = ( cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage) ) diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index 53e74610659b..a1e7ab0ee30d 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -12,8 +12,12 @@ import psutil import regex as re +from vllm.logger import init_logger + DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES" +logger = init_logger(__name__) + @dataclass class LogicalCPUInfo: @@ -63,21 +67,20 @@ def _read_int_file(path: str) -> int | None: @cache -def get_cgroup_memory_limit() -> tuple[int | None, int | None]: - """Return (limit, usage) in bytes from cgroup, or (None, None). +def get_cgroup_memory_limit() -> int | None: + """Return the cgroup memory limit in bytes, or None. - Supports both cgroup v2 (unified) and v1. Returns (None, None) when + Supports both cgroup v2 (unified) and v1. Returns None when not running under a constrained cgroup (e.g. bare metal, or limit reported as `max`/an unrealistically large value). """ if sys.platform != "linux": - return None, None + return None # cgroup v2 unified hierarchy v2_limit = _read_int_file("/sys/fs/cgroup/memory.max") if v2_limit is not None: - v2_usage = _read_int_file("/sys/fs/cgroup/memory.current") - return v2_limit, v2_usage + return v2_limit # cgroup v1 v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes") @@ -85,11 +88,74 @@ def get_cgroup_memory_limit() -> tuple[int | None, int | None]: # cgroup v1 reports a huge sentinel (close to PAGE_COUNTER_MAX) # when unlimited. Treat absurdly large values as "no limit". if v1_limit >= (1 << 62): - return None, None - v1_usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes") - return v1_limit, v1_usage + return None + return v1_limit + + return None + - return None, None +def get_cgroup_memory_usage() -> int | None: + """Return the current cgroup memory usage in bytes, or None. + + The usage value is intentionally read on every call because cgroup + memory usage changes while the process is running. + """ + if sys.platform != "linux": + return None + + # cgroup v2 unified hierarchy + if _read_int_file("/sys/fs/cgroup/memory.max") is not None: + return _read_int_file("/sys/fs/cgroup/memory.current") + + # cgroup v1 + v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes") + if v1_limit is not None and v1_limit < (1 << 62): + return _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes") + + return None + + +def check_cgroup_memory_available( + required_bytes: int, + allocation_name: str, +) -> None: + """Log cgroup memory headroom for an upcoming allocation. + + Args: + required_bytes: Bytes required by the allocation. + allocation_name: Human-readable name used in log messages. + + Low headroom logs a warning, but does not reject the allocation because + cgroup usage can include reclaimable memory. If the cgroup limit or usage + cannot be read, the check is skipped. + """ + cgroup_limit = get_cgroup_memory_limit() + cgroup_usage = get_cgroup_memory_usage() + if cgroup_limit is None or cgroup_usage is None: + return + + cgroup_available = max(0, cgroup_limit - cgroup_usage) + mib = 1 << 20 + remaining_bytes = cgroup_available - required_bytes + log_fn = logger.debug if remaining_bytes >= 0 else logger.warning + status = ( + "current headroom meets the requested allocation" + if remaining_bytes >= 0 + else "current headroom is below the requested allocation; allocation " + "will still be attempted because cgroup usage may be reclaimable" + ) + log_fn( + "Cgroup memory preflight for %s: %.0f MiB required, %.0f MiB current " + "usage, %.0f MiB available under %.0f MiB limit, %.0f MiB remaining " + "after allocation based on current usage; %s.", + allocation_name, + required_bytes / mib, + cgroup_usage / mib, + cgroup_available / mib, + cgroup_limit / mib, + remaining_bytes / mib, + status, + ) def get_memory_affinity(pid: int = 0) -> list[int]: @@ -161,7 +227,8 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo: # would be applied to host RAM instead of the pod's limit. cgroup # does not expose per-NUMA-node limits, so we just clamp the totals # against the pod-wide limit here. - cgroup_limit, cgroup_usage = get_cgroup_memory_limit() + cgroup_limit = get_cgroup_memory_limit() + cgroup_usage = get_cgroup_memory_usage() if cgroup_limit is not None and cgroup_limit < total_memory: total_memory = cgroup_limit cgroup_available = cgroup_limit - (cgroup_usage or 0) diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 46a1e7ae6460..284f1de68e8e 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -135,7 +135,10 @@ def __init__( self._creator = True if creator_memory_check is not None: creator_memory_check(self.total_size_bytes) - check_shm_free_space(self.total_size_bytes) + check_shm_free_space( + self.total_size_bytes, + allocation_name="CPU KV offload shared region in /dev/shm", + ) os.ftruncate(self.fd, self.total_size_bytes) logger.info( "Created mmap file %s (%.2f GB)", From 18553c5c9ff02acb84bd0f8698e90609e1aec2e9 Mon Sep 17 00:00:00 2001 From: freyfwt <110148731+freyfwt@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:52:26 +0800 Subject: [PATCH 303/774] [Kernel] Use Murmur3 RNG for Gumbel sampling (#51367) Signed-off-by: freyfwt <110148731+freyfwt@users.noreply.github.com> Signed-off-by: freyfwt Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- .../llm/test_struct_output_generate.py | 2 +- .../chat_completion/test_include_reasoning.py | 12 +- .../openai/completion/test_completion.py | 12 +- tests/v1/worker/test_gpu_gumbel_sample.py | 108 +++++++++++++++++- vllm/v1/worker/gpu/sample/gumbel.py | 108 ++++++++++++++++-- 5 files changed, 216 insertions(+), 26 deletions(-) diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 0f55ad7adc3e..030b0207e4c7 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -948,7 +948,7 @@ def test_structured_output_batched_with_non_structured_outputs_requests( prompts = [structured_outputs_prompt, non_structured_outputs_prompt] sampling_params = [ SamplingParams( - temperature=1.0, + temperature=0, max_tokens=400, structured_outputs=StructuredOutputsParams(json=sample_json_schema), ), diff --git a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py index 1d4e1ebc4055..5eeca438fe64 100644 --- a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py +++ b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py @@ -45,7 +45,8 @@ async def test_include_reasoning_true_non_streaming(client: openai.AsyncOpenAI): response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, extra_body={"include_reasoning": True}, ) @@ -63,7 +64,8 @@ async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI) response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, extra_body={"include_reasoning": False}, ) @@ -99,7 +101,8 @@ async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI): stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, stream=True, extra_body={"include_reasoning": True}, ) @@ -130,7 +133,8 @@ async def test_include_reasoning_false_streaming(client: openai.AsyncOpenAI): stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, stream=True, extra_body={"include_reasoning": False}, ) diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index af9f1def3614..1bc6260cdebc 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -353,7 +353,7 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str prompt = "What is an LLM?" n = 3 - max_tokens = 50 # we want some to finish earlier than others + max_tokens = 50 # High temperature to maximize chance of unique completions. completion = await client.completions.create( @@ -371,16 +371,12 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str num_completions = len(completion.choices) assert num_completions == n, f"Num completions {num_completions} but expected {n}." completion_repeats: dict[str, int] = {} - output_token_lengths = set() for idx, choice in enumerate(completion.choices): # Assert correct completion index & some finish reason. assert choice.index == idx, f"Index {choice.index} but expected {idx}." assert choice.finish_reason is not None, "None finish_reason is invalid." text = choice.text completion_repeats[text] = completion_repeats.get(text, 0) + 1 - output_token_lengths.add(len(choice.logprobs.tokens)) - # Assert subrequests finished at different times - assert len(output_token_lengths) > 1 # Assert `n` unique completions num_unique = len(completion_repeats) if num_unique != n: @@ -403,7 +399,7 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): prompt = "What is an LLM?" n = 3 - max_tokens = 50 # we want some to finish earlier than others + max_tokens = 50 stream = await client.completions.create( model=model_name, @@ -427,19 +423,15 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): f"Expected {n} completions with valid indices and finish_reason." ) completion_repeats: dict[str, int] = {} - chunk_lengths = set() for chunk in chunks: chunk_len = len(chunk) # Assert correct number of completion tokens - chunk_lengths.add(chunk_len) assert chunk_len <= max_tokens, ( f"max_tokens={max_tokens} but chunk len is {chunk_len}." ) text = "".join(chunk) completion_repeats[text] = completion_repeats.get(text, 0) + 1 print(text) - # Assert subrequests finished at different times - assert len(chunk_lengths) > 1 # Assert `n` unique completions num_unique = len(completion_repeats) if num_unique != n: diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index 10a09b9de4a6..a66a1395d2c7 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/tests/v1/worker/test_gpu_gumbel_sample.py @@ -21,7 +21,12 @@ if not torch.cuda.is_available(): pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True) -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.sample.gumbel import ( + _uniform64_from_random53, + gumbel_sample, + murmur3_hash32, +) DEVICE = "cuda" VOCAB_SIZE = 200_000 @@ -33,6 +38,103 @@ Z_TOLERANCE = 10.0 +@triton.jit +def _murmur3_known_answer_kernel( + seeds_ptr, + positions_ptr, + offsets_ptr, + hashes_ptr, + domain_hashes_ptr, + numel: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + idx = tl.arange(0, BLOCK_SIZE) + mask = idx < numel + seed = tl.load(seeds_ptr + idx, mask=mask) + pos = tl.load(positions_ptr + idx, mask=mask) + offset = tl.load(offsets_ptr + idx, mask=mask) + hash32 = murmur3_hash32(seed, pos, offset).to(tl.int64) + domain_hash32 = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.int64) + tl.store(hashes_ptr + idx, hash32, mask=mask) + tl.store(domain_hashes_ptr + idx, domain_hash32, mask=mask) + + +@triton.jit +def _uniform64_endpoint_kernel( + random53_ptr, + uniform_ptr, + numel: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + idx = tl.arange(0, BLOCK_SIZE) + mask = idx < numel + random53 = tl.load(random53_ptr + idx, mask=mask).to(tl.uint64) + tl.store(uniform_ptr + idx, _uniform64_from_random53(random53), mask=mask) + + +def test_murmur3_known_answers(): + seeds = torch.tensor( + [0, 1, -1, 0x123456789ABCDEF, -(1 << 63), (1 << 63) - 1], + dtype=torch.int64, + device=DEVICE, + ) + positions = torch.tensor( + [0, 2, 17, 123456789, 1 << 30, (1 << 30) + 7], + dtype=torch.int64, + device=DEVICE, + ) + offsets = torch.tensor( + [0, 3, 1023, 199999, 4096, 0xFFFFFFFF], + dtype=torch.int64, + device=DEVICE, + ) + hashes = torch.empty_like(seeds) + domain_hashes = torch.empty_like(seeds) + + _murmur3_known_answer_kernel[(1,)]( + seeds, + positions, + offsets, + hashes, + domain_hashes, + numel=seeds.numel(), + BLOCK_SIZE=8, + ) + + assert hashes.tolist() == [ + 0x8134CDF8, + 0x23B53F91, + 0xC8C13C14, + 0x5DDE974F, + 0xF1A9355B, + 0xD952F9AA, + ] + assert domain_hashes.tolist() == [ + 0xD28AD962, + 0xE3AAF818, + 0xD223C51C, + 0x0547D2FE, + 0x3AED160F, + 0x49E8EA7F, + ] + + +def test_uniform64_excludes_endpoints(): + random53 = torch.tensor([0, (1 << 53) - 1], dtype=torch.int64, device=DEVICE) + uniform = torch.empty(2, dtype=torch.float64, device=DEVICE) + + _uniform64_endpoint_kernel[(1,)]( + random53, + uniform, + numel=random53.numel(), + BLOCK_SIZE=2, + ) + + assert torch.isfinite(uniform).all() + assert ((uniform > 0.0) & (uniform < 1.0)).all() + assert uniform.tolist() == [2.0**-54, 1.0 - 2.0**-53] + + def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor: """Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N.""" gen = torch.Generator(device=DEVICE).manual_seed(seed) @@ -180,7 +282,7 @@ def test_full_vocab_distribution_fidelity(): def test_drafting_uses_a_separate_noise_stream(): - """is_drafting salts the Philox offset: same inputs, different draws. + """is_drafting salts the RNG position: same inputs, different draws. The draft proposal and the residual resample after a rejection must be independent. They key noise by the same (seed, pos), so only the salt keeps @@ -189,7 +291,7 @@ def test_drafting_uses_a_separate_noise_stream(): tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional consequence. - Relocating the offset must not distort the draw either, so both streams are + Salting the position must not distort the draw either, so both streams are checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits below the head -- the regime where fp32 Gumbel precision matters. """ diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 4bf1f47280cf..722021236926 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton +from vllm.triton_utils import HAS_TRITON, tl, triton -# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp -# zero draws before the flipped Gumbel transform below. +# Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the +# retained Philox helper for rejection sampling. # # Triton requires globals accessed from `@triton.jit` functions to be wrapped # in `tl.constexpr(...)`. We can only do that when Triton is actually @@ -87,6 +87,97 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr): return u +@triton.jit +def _murmur3_rotl32(value, shift: tl.constexpr): + return (value << shift) | (value >> (32 - shift)) + + +@triton.jit +def _murmur3_mix(h, key): + key *= 0xCC9E2D51 + key = _murmur3_rotl32(key, 15) + key *= 0x1B873593 + h ^= key + h = _murmur3_rotl32(h, 13) + return h * 5 + 0xE6546B64 + + +@triton.jit +def _murmur3_fmix32(h): + h ^= h >> 16 + h *= 0x85EBCA6B + h ^= h >> 13 + h *= 0xC2B2AE35 + return h ^ (h >> 16) + + +@triton.jit +def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0): + seed = seed.to(tl.int64) + pos = pos.to(tl.int64) + offset = offset.to(tl.uint32) + # Keep the request-wide prefix scalar until the token offset is mixed in. + h = (seed ^ seed).to(tl.uint32) + h ^= domain + h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, offset) + return _murmur3_fmix32(h ^ 16) + + +@triton.jit +def murmur3_uniform32(seed, pos, offset): + random32 = murmur3_hash32(seed, pos, offset) + # Split the uint32 before converting to fp32 so backends without a native + # uint32-to-float conversion can still use all 32 source bits. Both 16-bit + # halves convert exactly; their sum is the correctly rounded fp32 value of + # (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps + # the full 32-bit source resolution instead of being truncated to 24 bits. + hi16 = (random32 >> 16).to(tl.int32) + lo16 = (random32 & 0xFFFF).to(tl.int32) + return ( + hi16.to(tl.float32) * 1.52587890625e-05 + + (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10 + ) + + +@triton.jit +def _uniform64_from_random53(random53): + uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16 + # The largest midpoint rounds to 1.0 in fp64; keep the uniform open without + # relying on a near-one literal that Triton may materialize in fp32. + return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform) + + +@triton.jit +def murmur3_uniform64(seed, pos, offset): + lo = murmur3_hash32(seed, pos, offset).to(tl.uint64) + hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64) + random53 = ((hi << 32) | lo) >> 11 + return _uniform64_from_random53(random53) + + +@triton.jit +def _log1p_neg_stable(value): + # Preserve precision for the positive Gumbel tail without relying on a + # backend-specific libdevice log1p. The degree-8 series has absolute error + # below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for + # the part of the distribution that can win the argmax. + polynomial = 1.0 / 8.0 + polynomial = 1.0 / 7.0 + value * polynomial + polynomial = 1.0 / 6.0 + value * polynomial + polynomial = 1.0 / 5.0 + value * polynomial + polynomial = 1.0 / 4.0 + value * polynomial + polynomial = 1.0 / 3.0 + value * polynomial + polynomial = 1.0 / 2.0 + value * polynomial + polynomial = 1.0 + value * polynomial + series = -value * polynomial + + direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08)) + return tl.where(value < 0.25, series, direct) + + @triton.jit def gumbel_noised_argmax( logits, @@ -117,14 +208,15 @@ def gumbel_noised_argmax( if temp != 0.0: if IS_DRAFTING: pos = pos + _DRAFT_NOISE_SALT - gumbel_seed = tl.randint(seed, pos) if USE_FP64: - u = tl_rand64(gumbel_seed, keys, includes_zero=False) + u = murmur3_uniform64(seed, pos, keys) gumbel_noise = -tl.log(-tl.log(u)) else: - u = tl_rand32(gumbel_seed, keys, includes_zero=False) - # log1p keeps the winning tail at u -> 0, where fp32 resolves it. - gumbel_noise = -tl.log(-tldevice.log1p(-u)) + u = murmur3_uniform32(seed, pos, keys) + # Draw the large-noise tail (which decides the argmax winner) from + # u -> 0, where fp32 has fine resolution. Avoid backend-specific + # log1p while preserving precision in the winning tail. + gumbel_noise = -tl.log(-_log1p_neg_stable(u)) logits = tl.where(mask, logits + gumbel_noise, float("-inf")) return tl.max(logits, axis=0, return_indices=True) From f30a195bbb15b920d9c2c40e6a3466d8961ab101 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Wed, 16 Sep 2026 03:23:20 -0400 Subject: [PATCH 304/774] [Bugfix] Fix incorrect Mamba block allocation estimate that prevents request admission (#57050) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../core/test_single_type_kv_cache_manager.py | 57 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 8 +-- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 14e46d7ba07c..f20a96f5cc31 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -179,6 +179,63 @@ def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks): assert pool.get_num_free_blocks() == initial_free +@pytest.mark.parametrize("block_size", [896, 1536]) +@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4]) +@pytest.mark.parametrize("prompt_tokens", [25121, 704547]) +def test_mamba_checkpoint_admission_matches_allocation( + block_size, num_speculative_blocks, prompt_tokens +): + """Checkpoint admission must match the subsequent physical allocation.""" + spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + num_prefill_checkpoint_blocks=1, + prefill_checkpoint_alignment=64, + ) + pool = BlockPool( + num_gpu_blocks=2048, + enable_caching=True, + hash_block_size=128, + ) + manager = MambaManager( + spec, + block_pool=pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=block_size, + ) + request_id = "prefill" + computed_tokens = 23040 + + def estimate(num_tokens, total_computed_tokens, apply_admission_cap): + return manager.get_num_blocks_to_allocate( + request_id=request_id, + num_tokens=num_tokens, + new_computed_blocks=[], + total_computed_tokens=total_computed_tokens, + num_local_computed_tokens=total_computed_tokens, + num_tokens_main_model=num_tokens, + apply_admission_cap=apply_admission_cap, + ) + + estimate(computed_tokens, 0, False) + manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens) + assert request_id in manager._allocated_block_reqs + + admission_estimate = estimate(prompt_tokens, computed_tokens, True) + allocation_estimate = estimate(prompt_tokens, computed_tokens, False) + assert request_id in manager._checkpoints + + free_before = pool.get_num_free_blocks() + manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens) + allocated = free_before - pool.get_num_free_blocks() + + assert admission_estimate == allocation_estimate == allocated + + def get_sliding_window_manager( sliding_window_spec, block_pool, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index fc5dc410485d..08eab7628768 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1796,10 +1796,10 @@ def get_num_blocks_to_allocate( self._checkpoints.pop(request_id, None) if num_new_blocks > 0: blocks_allocated = request_id in self._allocated_block_reqs - if not (checkpoint_block and blocks_allocated): - num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block - if not blocks_allocated: - num_new_blocks += self.num_speculative_blocks + physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block + if not blocks_allocated or checkpoint_block: + physical_block_cap += self.num_speculative_blocks + num_new_blocks = min(num_new_blocks, physical_block_cap) num_evictable_computed_blocks = self._get_num_evictable_blocks( new_computed_blocks From ab35354c21cc3c36439e79e18070f85cfafc3af2 Mon Sep 17 00:00:00 2001 From: Zheng Gong Date: Wed, 16 Sep 2026 15:26:56 +0800 Subject: [PATCH 305/774] [ROCm][AITER] Skip AITER norm kernels when flattening to 2D would copy (#55991) Signed-off-by: Zheng Gong Co-authored-by: Cursor Co-authored-by: TJian --- tests/kernels/ir/test_aiter_norm_dispatch.py | 60 ++++++++++++++++++++ vllm/kernels/aiter_ops.py | 32 ++++++++++- 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/ir/test_aiter_norm_dispatch.py diff --git a/tests/kernels/ir/test_aiter_norm_dispatch.py b/tests/kernels/ir/test_aiter_norm_dispatch.py new file mode 100644 index 000000000000..19cf5d3112dc --- /dev/null +++ b/tests/kernels/ir/test_aiter_norm_dispatch.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the AITER norm support predicates. + +The AITER norm wrappers flatten >2-D activations with ``Tensor.reshape``, which +silently copies when the flattened shape is not expressible with the input's +strides. The support predicates must reject those inputs so dispatch falls +through to a provider that handles arbitrary strides. +""" + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.kernels.aiter_ops import flatten_to_2d_is_free + +pytestmark = pytest.mark.skipif( + not is_aiter_found_and_supported(), + reason="Only test on ROCm with AITER installed and supported", +) + + +def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim): + """q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" + q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim + qkv = torch.empty(num_tokens, q_size + 2 * kv_size) + q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] + return q.view(*q.shape[:-1], num_q_heads, head_dim) + + +@pytest.mark.parametrize( + "x", + [ + torch.empty(8, 16), + torch.empty(16), + torch.empty(2, 4, 16), + torch.empty(2, 1, 16), + # A contiguous tensor stays flattenable after a leading-dim slice. + torch.empty(8, 4, 16)[2:6], + ], +) +def test_flattenable(x): + assert flatten_to_2d_is_free(x) + assert x.reshape(-1, x.shape[-1]).data_ptr() == x.data_ptr() + + +@pytest.mark.parametrize( + "x", + [ + _qkv_slice_by_head(8, 32, 8, 128), + # Non-unit last-dim stride. + torch.empty(4, 8, 32).transpose(-1, -2), + # Leading dims cannot be merged: a slice along the middle dim. + torch.empty(4, 8, 32)[:, :4], + ], +) +def test_not_flattenable(x): + assert not flatten_to_2d_is_free(x) + # reshape has to copy, which is exactly what the predicate guards against. + assert x.reshape(-1, x.shape[-1]).data_ptr() != x.data_ptr() diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index de8a53ae8b77..402c9cac1752 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -34,13 +34,38 @@ def is_aiter_found() -> bool: AITER_SUPPORTED = is_aiter_found() """Most kernels in this file are supported if AITER is installed.""" + +def flatten_to_2d_is_free(x: Tensor) -> bool: + """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows. + + The AITER norm kernels only take dense 2D inputs, so the wrappers below + flatten the leading dims with ``Tensor.reshape``, which silently falls back + to ``contiguous()`` when the flattened shape is not expressible with the + existing strides, adding a whole-tensor device copy that costs more than + the kernel saves. A strided last dim flattens without a copy but not into a + dense buffer, so it is rejected too. + """ + if x.dim() <= 2: + return True + if x.stride(-1) != 1: + return False + expected_stride = x.size(-1) + for i in range(x.dim() - 2, -1, -1): + if x.size(i) != 1 and x.stride(i) != expected_stride: + return False + expected_stride *= x.size(i) + return True + + rms_no_var_16bit_only = ( lambda x, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) + and flatten_to_2d_is_free(x) ) """AITER rms_norm only supports float16 and bfloat16 acts, no var_size override, -and requires weight dtype to match x dtype.""" +requires weight dtype to match x dtype, and requires flattening the activation +to 2D to be free.""" @ir.ops.rms_norm.register_impl( @@ -79,10 +104,13 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) + and flatten_to_2d_is_free(x) + and flatten_to_2d_is_free(x_residual) ) """ AITER fused_add_rms_norm only supports 16-bit activations and no var_size override. -Requires weight dtype to match x dtype. +Requires weight dtype to match x dtype, and flattening both the activation and the +residual to 2D to be free. """ From f12fe10c5c377ccd401b66c39df3ffb394fa067c Mon Sep 17 00:00:00 2001 From: MINJUN GIL Date: Wed, 16 Sep 2026 16:47:26 +0900 Subject: [PATCH 306/774] [Bugfix][KV Offload] Track cache recency once per request (#51787) Signed-off-by: mindungil Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 71 +-- tests/v1/kv_offload/cpu/test_manager.py | 436 +++++++++++++++++- .../tiering/test_tiering_offloading.py | 61 ++- .../kv_connector/v1/offloading/scheduler.py | 70 ++- vllm/v1/kv_offload/base.py | 12 + vllm/v1/kv_offload/cpu/manager.py | 146 +++++- vllm/v1/kv_offload/cpu/policies/arc.py | 81 +++- vllm/v1/kv_offload/cpu/policies/base.py | 52 ++- vllm/v1/kv_offload/cpu/policies/lru.py | 142 ++++-- vllm/v1/kv_offload/tiering/manager.py | 18 +- 10 files changed, 918 insertions(+), 171 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 89b77ae43fa8..035fb6b0054f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -258,6 +258,9 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources(): jobs = scheduler._build_partial_tail_store_jobs(output) assert len(jobs) == 1 + offered_keys = scheduler.manager.prepare_store.call_args.args[0] + req_context = scheduler._req_status["req"].req_context + assert all(req_context.get_offload_key_position(key) == 28 for key in offered_keys) [job_id] = jobs src_spec = jobs[job_id].src_spec assert isinstance(src_spec, GPULoadStoreSpec) @@ -380,6 +383,13 @@ def test_normal_store_excludes_align_mode_mamba_sources(): req_status.group_states[0].block_ids[:] = [11] req_status.group_states[1].block_ids[:] = [99] req_status.update_offload_keys() + for group_config, group_state in zip( + scheduler.config.kv_group_configs, req_status.group_states + ): + assert ( + req_status.req_context.get_offload_key_position(group_state.offload_keys[0]) + == group_config.tokens_per_chunk + ) scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -921,7 +931,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1)) # 1 more block (+ token for kicking off offloading) - # now check touch was called with all 6 blocks runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -929,9 +938,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool): decoded_tokens=[0] * (tokens_per_chunk + 1), expected_stored=(15, 16, 17), ) - runner.manager.touch.assert_called() - block_hashes1 = list(runner.manager.touch.call_args.args[0]) - assert len(block_hashes1) == 6 # terminate request runner.run(decoded_tokens=[EOS_TOKEN_ID]) @@ -939,13 +945,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # create a new request differing only on the last token runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1]) runner.run(decoded_tokens=[0]) - runner.manager.touch.assert_called() - block_hashes2 = list(runner.manager.touch.call_args.args[0]) - assert len(block_hashes2) == 6 - - # verify hashes are the same, except for the last block - assert block_hashes1[:5] == block_hashes2[:5] - assert block_hashes1[5] != block_hashes2[5] + runner.manager.touch.assert_not_called() # terminate request runner.run( @@ -1353,14 +1353,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo generate_store_output(keys) ) runner.run(decoded_tokens=[0]) - # _touch called from get_num_new_matched_tokens (2 groups) and - # _get_reqs_to_store (2 groups) → 4 touch calls total. - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 - assert len(touch_calls[0].args[0]) == 3 - assert len(touch_calls[1].args[0]) == 3 - assert len(touch_calls[2].args[0]) == 3 - assert len(touch_calls[3].args[0]) == 3 + runner.manager.touch.assert_not_called() # store 3 more block runner.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -1371,9 +1364,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo expected_stored=(0, 1, 2, 3, 4, 5), ) - # touch called from _get_reqs_to_store * 3 blocks, once for each group - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 6 + runner.manager.touch.assert_not_called() # EOS lands in the last slot of the 7th block (offset 6). No forward pass # writes that slot, so the finishing step declines the block. @@ -1392,13 +1383,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), ) - # 2 touch calls from get_num_new_matched_tokens (2 groups) - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - # full attention group touched all 3 blocks - assert len(touch_calls[0].args[0]) == 3 - # sliding window group touched just the last 2 blocks - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # 3 blocks are hit on GPU [0, 1, 2] # 1 block loaded [3,] @@ -1464,15 +1449,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool generate_store_output(keys) ) runner.run(decoded_tokens=[0]) - # _touch called from get_num_new_matched_tokens (2 groups) and - # _get_reqs_to_store (2 groups) → 4 touch calls total. - # Group 0 has 2 offload keys, group 1 has 1. - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 - assert len(touch_calls[0].args[0]) == 2 - assert len(touch_calls[1].args[0]) == 1 - assert len(touch_calls[2].args[0]) == 2 - assert len(touch_calls[3].args[0]) == 1 + runner.manager.touch.assert_not_called() # Get to 31 tokens # No further blocks offloaded @@ -1482,11 +1459,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1], ending_token_offset = 24 # Group 1 blocks: [0, 1], ending_token_offset = 32 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: only group 1 has a new block to store - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 2 - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # Get to 35 tokens # No further blocks offloaded @@ -1496,11 +1469,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1, 2], ending_token_offset = 36 # Group 1 blocks: [0, 1], ending_token_offset = 32 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: only group 0 has a new block to store - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 3 - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # Get to 47 tokens # No further blocks offloaded @@ -1510,11 +1479,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4 # Group 1 blocks: [0, 1, 2], ending_token_offset = 48 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: both groups have a new block, each with 1 key - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 4 - assert len(touch_calls[1].args[0]) == 3 + runner.manager.touch.assert_not_called() runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2))) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index a79f480ff551..4fce65fb2a4f 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -23,6 +23,7 @@ ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy def make_req_context( @@ -430,31 +431,34 @@ def test_cpu_manager(): assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare load [2, 3] - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) + load_ctx = make_req_context("load-2-3") + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), load_ctx) verify_load_output(prepare_load_output, [1, 2]) # prepare store with no space ([2, 3] is being loaded) assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent. - cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) + cpu_manager.complete_load(to_keys([2, 3]), load_ctx) + cpu_manager.on_request_finished(load_ctx) - # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest) + # prepare store [6, 7, 8] -> evicts [4, 5, 3] (oldest). Within + # the accessed prefix [2, 3], the tail is less valuable than the head. prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[6, 7, 8], - store_chunk_ids=[1, 0, 3], - evicted_keys=[4, 5, 2], + store_chunk_ids=[2, 0, 3], + evicted_keys=[4, 5, 3], ), ) # complete store [6, 7, 8] cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) - # touch [3, 6, 7] (move to end of LRU order) - cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX) + # touch [2, 6, 7] (move to end of LRU order) + cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX) # prepare store [7, 9] -> evicts [8] (oldest following previous touch) prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) @@ -477,7 +481,7 @@ def test_cpu_manager(): verify_events( cpu_manager.take_events(), expected_stores=({3, 4, 5}, {6, 7, 8}), - expected_evictions=({4, 5, 2}, {8}), + expected_evictions=({4, 5, 3}, {8}), ) @@ -518,6 +522,65 @@ def test_prepare_load_preserves_key_order(): manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant +def test_lru_batch_eviction_failure_is_atomic(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + policy = manager._policy + assert isinstance(policy, LRUCachePolicy) + keys = to_keys([1, 2, 3, 4]) + assert manager.prepare_store(keys, _EMPTY_REQ_CTX) is not None + manager.complete_store(keys, _EMPTY_REQ_CTX) + + protected = set(keys[1:]) + assert policy.evict(2, protected) is None + assert all(manager.lookup(key, _EMPTY_REQ_CTX) is LookupResult.HIT for key in keys) + + evicted = policy.evict(1, protected) + assert evicted is not None + assert [key for key, _ in evicted] == [keys[0]] + + +def test_lru_repeated_pin_cycles_compact_lazy_heap_entries(): + manager = make_cpu_manager(num_chunks=1, cache_policy="lru") + policy = manager._policy + assert isinstance(policy, LRUCachePolicy) + key = to_key(1) + assert manager.prepare_store([key], _EMPTY_REQ_CTX) is not None + manager.complete_store([key], _EMPTY_REQ_CTX) + + for _ in range(128): + manager.prepare_load([key], _EMPTY_REQ_CTX) + manager.complete_load([key], _EMPTY_REQ_CTX) + + assert len(policy._heap) < 64 + evicted = policy.evict(1, set()) + assert evicted is not None + assert [evicted_key for evicted_key, _ in evicted] == [key] + + +def test_reset_discards_stale_request_access_classification(): + manager = make_cpu_manager(num_chunks=1, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key = to_key(1) + resumed_ctx = make_req_context("resumed") + + assert manager.prepare_store([key], resumed_ctx) is not None + manager.complete_store([key], resumed_ctx) + manager.reset_cache() + + replacement_ctx = make_req_context("replacement") + assert manager.prepare_store([key], replacement_ctx) is not None + manager.complete_store([key], replacement_ctx) + manager.on_request_finished(replacement_ctx) + assert key in policy.t1 + + manager.prepare_load([key], resumed_ctx) + manager.complete_load([key], resumed_ctx) + manager.on_request_finished(resumed_ctx) + + assert key in policy.t2 + + class TestARCPolicy: """Unit tests for CPUOffloadingManager with ARC eviction policy.""" @@ -1042,6 +1105,48 @@ def test_filter_reused_manager_oversized_offer_makes_progress(): assert stored_keys == set(keys) +def test_store_threshold_does_not_hide_ready_reuse(): + manager = make_cpu_manager( + num_chunks=2, + cache_policy="lru", + store_threshold=2, + max_tracker_size=2, + ) + head, tail, noise_1, noise_2, replacement = to_keys(list(range(5))) + + seed_ctx = make_req_context("seed") + output = manager.prepare_store([head, tail], seed_ctx) + assert output is not None + assert not output.keys_to_store + output = manager.prepare_store([head, tail], seed_ctx) + assert output is not None + assert output.keys_to_store == [head, tail] + manager.complete_store([head, tail], seed_ctx) + manager.on_request_finished(seed_ctx) + + # Age both resident keys out of the admission tracker without storing the + # noise keys. Reusing the tail recreates its counter at one, below the + # threshold, but must still refresh its cache recency. + output = manager.prepare_store([noise_1, noise_2], make_req_context("noise")) + assert output is not None + assert not output.keys_to_store + reuse_ctx = make_req_context("reuse") + skipped_before_reuse = manager.stores_skipped_in_current_batch + output = manager.prepare_store([tail], reuse_ctx) + assert output is not None + assert not output.keys_to_store + assert manager.stores_skipped_in_current_batch == skipped_before_reuse + manager.on_request_finished(reuse_ctx) + + replacement_ctx = make_req_context("replacement") + output = manager.prepare_store([replacement], replacement_ctx) + assert output is not None + assert not output.keys_to_store + output = manager.prepare_store([replacement], replacement_ctx) + assert output is not None + assert output.evicted_keys == [head] + + def test_evictable_cache_chunk_count(): """ Verifies _num_evictable_cache_chunks is maintained correctly through the @@ -1156,3 +1261,318 @@ def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: assert len(received) == 1 assert received[0][0] == keys assert received[0][1] is ctx + + +@pytest.mark.parametrize("finish_before_completion", [False, True]) +def test_request_finish_orders_lru_prefix_independent_of_store_completion( + finish_before_completion: bool, +): + """A transfer's completion order must not make the prefix head LRU.""" + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + prefix = to_keys([1, 2, 3, 4]) + ctx = make_req_context("prefix") + + output = manager.prepare_store(prefix, ctx) + assert output is not None + if finish_before_completion: + manager.on_request_finished(ctx) + manager.complete_store(set(prefix), ctx) + if not finish_before_completion: + manager.on_request_finished(ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == to_keys([4]) + + +@pytest.mark.parametrize("finish_before_completion", [False, True]) +def test_request_finish_orders_lru_prefix_independent_of_load_completion( + finish_before_completion: bool, +): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + prefix = to_keys([1, 2, 3, 4]) + + seed_ctx = make_req_context("seed") + output = manager.prepare_store(prefix, seed_ctx) + assert output is not None + manager.complete_store(prefix, seed_ctx) + + reuse_ctx = make_req_context("reuse") + manager.prepare_load(prefix, reuse_ctx) + if finish_before_completion: + manager.on_request_finished(reuse_ctx) + manager.complete_load(set(prefix), reuse_ctx) + if not finish_before_completion: + manager.on_request_finished(reuse_ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == to_keys([4]) + + +def test_late_old_completion_does_not_override_newer_request_recency(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + old_keys = to_keys([1, 2]) + new_keys = to_keys([3, 4]) + + old_ctx = make_req_context("old") + assert manager.prepare_store(old_keys, old_ctx) is not None + manager.on_request_finished(old_ctx) + + new_ctx = make_req_context("new") + assert manager.prepare_store(new_keys, new_ctx) is not None + manager.complete_store(new_keys, new_ctx) + manager.on_request_finished(new_ctx) + + # Physical completion is late, but the logical access is still older. + manager.complete_store(old_keys, old_ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == [old_keys[-1]] + + +def test_arc_counts_reuse_once_and_keeps_insertions_in_t1(): + manager = make_cpu_manager(num_chunks=4, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + prefix = to_keys([1, 2, 3, 4]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store(prefix, seed_ctx) is not None + manager.complete_store(prefix, seed_ctx) + # Repeated store offers from the same request must not become ARC hits. + assert manager.prepare_store(prefix, seed_ctx) is not None + # Nor should an internal read used to cascade the new chunks to a tier. + manager.prepare_load(prefix, seed_ctx) + manager.complete_load(prefix, seed_ctx) + manager.on_request_finished(seed_ctx) + + assert list(policy.t1) == list(reversed(prefix)) + assert not policy.t2 + + reuse_ctx = make_req_context("reuse") + manager.prepare_load(prefix, reuse_ctx) + manager.complete_load(prefix, reuse_ctx) + # A duplicate internal pin, as used by tiering cascades, is deduplicated. + manager.prepare_load(prefix, reuse_ctx) + manager.complete_load(prefix, reuse_ctx) + manager.on_request_finished(reuse_ctx) + + assert not policy.t1 + assert list(policy.t2) == list(reversed(prefix)) + + +def test_arc_reuse_wins_if_same_request_later_reinserts_key(): + manager = make_cpu_manager(num_chunks=1, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key_1, key_2 = to_keys([1, 2]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([key_1], seed_ctx) is not None + manager.complete_store([key_1], seed_ctx) + manager.on_request_finished(seed_ctx) + + reuse_ctx = make_req_context("reuse-and-reinsert") + manager.prepare_load([key_1], reuse_ctx) + manager.complete_load([key_1], reuse_ctx) + + other_ctx = make_req_context("evict-reused-key") + assert manager.prepare_store([key_2], other_ctx) is not None + manager.complete_store([key_2], other_ctx) + manager.on_request_finished(other_ctx) + + assert manager.prepare_store([key_1], reuse_ctx) is not None + manager.complete_store([key_1], reuse_ctx) + manager.on_request_finished(reuse_ctx) + + assert key_1 not in policy.t1 + assert list(policy.t2) == [key_1] + + +def test_request_finish_forwards_grouped_deduplicated_accesses(monkeypatch): + manager = make_cpu_manager(num_chunks=5) + existing = make_offload_key(b"existing", 0) + group_0_new = make_offload_key(b"group-0-new", 0) + group_1_head = make_offload_key(b"group-1-head", 1) + group_1_tail = make_offload_key(b"group-1-tail", 1) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([existing], seed_ctx) is not None + manager.complete_store([existing], seed_ctx) + + ctx = make_req_context("grouped") + # Group 1 is observed first, but policy finalization follows group index. + offered = [group_1_head, existing, group_0_new] + assert manager.prepare_store(offered, ctx) is not None + # Later offers extend a group's prefix and may repeat earlier keys. + assert ( + manager.prepare_store([group_0_new, group_1_head, group_1_tail], ctx) + is not None + ) + + received = [] + + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): + received.append((key_groups, insertion_only_keys, reused_keys, req_context)) + + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) + manager.on_request_finished(ctx) + manager.on_request_finished(ctx) + + assert received == [ + ( + ((existing, group_0_new), (group_1_head, group_1_tail)), + {group_0_new, group_1_head, group_1_tail}, + {existing}, + ctx, + ) + ] + + +def test_failed_store_records_access_once_and_uses_store_miss_hook(monkeypatch): + manager = make_cpu_manager(num_chunks=1) + resident, candidate = to_keys([1, 2]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([resident], seed_ctx) is not None + manager.complete_store([resident], seed_ctx) + + pin_ctx = make_req_context("pin") + manager.prepare_load([resident], pin_ctx) + + record_calls = 0 + original_record = manager._record_request_cache_access + + def record_access(*args, **kwargs): + nonlocal record_calls + record_calls += 1 + return original_record(*args, **kwargs) + + store_miss_calls = [] + + def on_store_miss(keys, req_context): + store_miss_calls.append((list(keys), req_context)) + + finalized = [] + + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): + finalized.append((key_groups, insertion_only_keys, reused_keys, req_context)) + + monkeypatch.setattr(manager, "_record_request_cache_access", record_access) + monkeypatch.setattr(manager._policy, "on_store_miss", on_store_miss) + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) + + ctx = make_req_context("failed-store") + assert manager.prepare_store([candidate, resident], ctx) is None + assert record_calls == 1 + assert store_miss_calls == [([candidate], ctx)] + + manager.on_request_finished(ctx) + assert finalized == [ + (((candidate, resident),), set(), {resident}, ctx), + ] + + manager.complete_load([resident], pin_ctx) + + +def test_arc_ghost_hit_adapts_once_per_request_before_insertion(): + manager = make_cpu_manager(num_chunks=2, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + keys = to_keys([1, 2, 3]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store(keys[:2], seed_ctx) is not None + manager.complete_store(keys[:2], seed_ctx) + manager.on_request_finished(seed_ctx) + + evict_ctx = make_req_context("create-ghost") + output = manager.prepare_store([keys[2]], evict_ctx) + assert output is not None + [ghost_key] = output.evicted_keys + manager.complete_store([keys[2]], evict_ctx) + manager.on_request_finished(evict_ctx) + assert ghost_key in policy.b1 + + resident_keys = [key for key in keys if key != ghost_key] + pin_ctx = make_req_context("pin-residents") + manager.prepare_load(resident_keys, pin_ctx) + + ghost_ctx = make_req_context("ghost-hit") + target_before = policy.target_t1_size + assert manager.prepare_store([ghost_key], ghost_ctx) is None + target_after_first_offer = policy.target_t1_size + assert target_after_first_offer > target_before + assert manager.prepare_store([ghost_key], ghost_ctx) is None + assert policy.target_t1_size == target_after_first_offer + + manager.complete_load(resident_keys, pin_ctx) + output = manager.prepare_store([ghost_key], ghost_ctx) + assert output is not None + assert policy.target_t1_size == target_after_first_offer + manager.complete_store([ghost_key], ghost_ctx) + manager.on_request_finished(ghost_ctx) + assert ghost_key in policy.t1 + assert ghost_key not in policy.t2 + + +def test_arc_pending_chunk_is_not_a_frequency_hit(): + manager = make_cpu_manager(num_chunks=2, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key = to_key(1) + + writer_ctx = make_req_context("writer") + assert manager.prepare_store([key], writer_ctx) is not None + + observer_ctx = make_req_context("pending-observer") + output = manager.prepare_store([key], observer_ctx) + assert output is not None and not output.keys_to_store + manager.on_request_finished(observer_ctx) + + chunk = policy.get(key) + assert chunk is not None and not chunk.is_ready + assert key in policy.t1 + assert key not in policy.t2 + + +def test_request_key_positions_override_store_observation_order(): + manager = make_cpu_manager(num_chunks=2, cache_policy="lru") + head, tail = to_keys([1, 2]) + ctx = make_req_context("out-of-order-store") + ctx.set_offload_key_position(head, 16) + ctx.set_offload_key_position(tail, 32) + + # A backfill or partial-tail path may offer a later key first. + assert manager.prepare_store([tail], ctx) is not None + assert manager.prepare_store([head], ctx) is not None + manager.on_request_finished(ctx) + manager.complete_store({head, tail}, ctx) + + output = manager.prepare_store([to_key(3)], make_req_context("evict-tail")) + assert output is not None + assert output.evicted_keys == [tail] + + +def test_request_key_positions_order_tails_across_kv_groups(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + group_0_head = make_offload_key(b"group-0-head", 0) + group_0_tail = make_offload_key(b"group-0-tail", 0) + group_1_head = make_offload_key(b"group-1-head", 1) + group_1_tail = make_offload_key(b"group-1-tail", 1) + keys = [group_0_head, group_0_tail, group_1_head, group_1_tail] + ctx = make_req_context("hybrid-groups") + for key in (group_0_head, group_1_head): + ctx.set_offload_key_position(key, 16) + for key in (group_0_tail, group_1_tail): + ctx.set_offload_key_position(key, 32) + + assert manager.prepare_store(keys, ctx) is not None + manager.complete_store(keys, ctx) + manager.on_request_finished(ctx) + + output = manager.prepare_store(to_keys([10, 11]), make_req_context("evict")) + assert output is not None + assert set(output.evicted_keys) == {group_0_tail, group_1_tail} diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index f737685621b4..771153ec4734 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -40,6 +40,7 @@ TierMatcher, make_offload_key, ) +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy from vllm.v1.kv_offload.tiering.base import ( JobResult, SecondaryTierManager, @@ -753,14 +754,64 @@ def test_touch_propagates_to_all_tiers(self, manager_setup): # Touch chunks self.manager.touch(chunks, _CTX) - # Verify touch was called on primary tier (check LRU order) - primary_keys = list(self.primary_tier._policy.evictable_chunks.keys()) - assert primary_keys[-3:] == list(reversed(chunks)) + # Verify touch was called on the primary tier (head is most recent). + policy = self.primary_tier._policy + assert isinstance(policy, LRUCachePolicy) + ranks = policy._ranks + assert ranks[chunks[2]] < ranks[chunks[1]] < ranks[chunks[0]] # Verify touch was propagated to all secondary tiers self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX) self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX) + def test_request_order_survives_late_cascade_completion(self, manager_setup): + """Cascade completion must not replace request order with I/O order.""" + chunks = to_keys(range(5)) + self._start_request() + output = self.manager.prepare_store(chunks, _CTX) + assert output is not None + self.manager.complete_store(chunks, _CTX, success=True) + + # Finalize while both secondary-tier cascades still pin primary chunks. + self.manager.on_request_finished(_CTX) + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + + evict_ctx = ReqContext(req_id="evict") + self._start_request(evict_ctx) + output = self.manager.prepare_store(to_keys([5]), evict_ctx) + assert output is not None + assert output.evicted_keys == [chunks[-1]] + + def test_late_old_store_submission_does_not_override_newer_request_order( + self, manager_setup + ): + """Cascade submission after finish is not a new primary access.""" + old_blocks = to_keys(range(2)) + old_ctx = ReqContext(req_id="old") + self._start_request(old_ctx) + assert self.manager.prepare_store(old_blocks, old_ctx) is not None + self.manager.on_request_finished(old_ctx) + + new_blocks = to_keys(range(2, 5)) + new_ctx = ReqContext(req_id="new") + self._start_request(new_ctx) + assert self.manager.prepare_store(new_blocks, new_ctx) is not None + self.manager.complete_store(new_blocks, new_ctx, success=True) + self.manager.on_request_finished(new_ctx) + + # The old GPU->primary write lands after the newer request finished. + # Its cascade pins must not turn it into the more recent request. + self.manager.complete_store(old_blocks, old_ctx, success=True) + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + + evict_ctx = ReqContext(req_id="evict") + self._start_request(evict_ctx) + output = self.manager.prepare_store(to_keys([5]), evict_ctx) + assert output is not None + assert output.evicted_keys == [old_blocks[-1]] + def test_failed_store_no_cascade(self, manager_setup): """Test that failed GPU→primary store doesn't cascade.""" chunks = to_keys(range(3)) @@ -873,10 +924,10 @@ def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup job_metadata = self.secondary_tier1.submit_store.call_args.args[0] assert job_metadata.req_context is ctx - def test_on_request_finished_delays_secondary_until_store_submitted( + def test_on_request_finished_delays_only_secondary_until_store_submitted( self, manager_setup ): - """Manager hook is eager; secondary hooks wait for cascade submission.""" + """Primary order commits immediately; secondary cleanup waits.""" chunks = to_keys(range(2)) ctx = ReqContext(req_id="req_delayed_secondary") calls: list[tuple[str, str]] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 5cc66e4bf0ff..b253a2399867 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -427,8 +427,11 @@ def update_offload_keys(self) -> None: None, group_config.hashes_per_chunk, ): - group_state.offload_keys.append( - make_offload_key(req_block_hash, group_config.group_idx) + key = make_offload_key(req_block_hash, group_config.group_idx) + group_state.offload_keys.append(key) + self.req_context.set_offload_key_position( + key, + len(group_state.offload_keys) * group_config.tokens_per_chunk, ) def update_block_id_groups( @@ -742,37 +745,6 @@ def _sliding_window_lookup( ) return None if defer_lookup or pending_in_window else consecutive_hits - def _touch(self, req_status: RequestOffloadState): - for group_config, group_state in zip( - self.config.kv_group_configs, req_status.group_states - ): - if group_config.sliding_window_size_in_chunks is None: - self.manager.touch(group_state.offload_keys, req_status.req_context) - else: - # Keep only chunks needed to hit the original request, plus - # decoded chunks. - chunks_to_skip = max( - 0, - group_state.num_hit_chunks - - group_config.sliding_window_size_in_chunks, - ) - self.manager.touch( - group_state.offload_keys[chunks_to_skip:], - req_status.req_context, - ) - if req_status.partial_tail_boundary is not None: - self.manager.touch( - tuple( - self._make_boundary_key( - req_status.req, - group.group_idx, - req_status.partial_tail_boundary, - ) - for group in self.config.kv_group_configs - ), - req_status.req_context, - ) - def _lookup_complete_chunks( self, req_status: RequestOffloadState, @@ -977,10 +949,16 @@ def _lookup_complete_chunks( return num_hit_tokens def _make_boundary_key( - self, request: Request, group_idx: int, boundary_tokens: int + self, + request: Request, + group_idx: int, + boundary_tokens: int, + req_context: ReqContext, ) -> OffloadKey: hash_idx = boundary_tokens // self.config.tokens_per_hash - 1 - return make_offload_key(request.block_hashes[hash_idx], group_idx) + key = make_offload_key(request.block_hashes[hash_idx], group_idx) + req_context.set_offload_key_position(key, boundary_tokens) + return key def _lookup( self, @@ -1015,7 +993,10 @@ def _lookup( boundary_keys = [] for group_config in self.config.kv_group_configs: key = self._make_boundary_key( - req_status.req, group_config.group_idx, boundary + req_status.req, + group_config.group_idx, + boundary, + req_status.req_context, ) boundary_keys.append(key) result = self.manager.lookup(key, req_status.req_context) @@ -1108,8 +1089,6 @@ def get_num_new_matched_tokens( self._maybe_observe_lookup_async_delay(req_status) req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0)) - self._touch(req_status) - return num_hit_tokens, bool(num_hit_tokens) def update_state_after_alloc( @@ -1183,7 +1162,10 @@ def update_state_after_alloc( if partial_tail_boundary is not None: keys_to_load.append( self._make_boundary_key( - request, group_config.group_idx, partial_tail_boundary + request, + group_config.group_idx, + partial_tail_boundary, + req_status.req_context, ) ) @@ -1307,7 +1289,9 @@ def _build_aligned_boundary_store_jobs( ): continue - key = self._make_boundary_key(req, group_idx, boundary) + key = self._make_boundary_key( + req, group_idx, boundary, req_status.req_context + ) store_output = self.manager.prepare_store([key], req_status.req_context) if store_output is None: self._connector_stats.increase_counter( @@ -1402,7 +1386,9 @@ def _build_partial_tail_store_jobs( ): continue keys = [ - self._make_boundary_key(req, group.group_idx, boundary) + self._make_boundary_key( + req, group.group_idx, boundary, req_status.req_context + ) for group in self.config.kv_group_configs ] block_ids = [ @@ -1689,8 +1675,6 @@ def _build_store_jobs( req_status.advance_stored_idx(num_offloadable_tokens) continue - self._touch(req_status) - keys_to_store = set(store_output.keys_to_store) group_sizes: list[int] = [] diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 93f8e4bda4e9..988747e06e1e 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -96,6 +96,12 @@ class ReqContext: # kv_transfer_params once (in on_new_request) and read the result back # on later calls for the same request. _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False) + # End-token position for each key in this request. The scheduler records + # these positions so managers can recover prefix order even when store + # calls arrive out of order (for example, SWA backfills). + _offload_key_positions: dict[OffloadKey, int] = field( + default_factory=dict, repr=False, init=False + ) def set_state(self, val: Any) -> None: self._state[type(val)] = val @@ -103,6 +109,12 @@ def set_state(self, val: Any) -> None: def get_state(self, cls: type[_T]) -> _T | None: return self._state.get(cls) + def set_offload_key_position(self, key: OffloadKey, end_token: int) -> None: + self._offload_key_positions[key] = end_token + + def get_offload_key_position(self, key: OffloadKey) -> int | None: + return self._offload_key_positions.get(key) + class LookupResult(Enum): """Result of OffloadingManager.lookup().""" diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 67b469b6a394..9d85090f2d8f 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict from collections.abc import Collection, Iterable +from dataclasses import dataclass, field from typing_extensions import override @@ -18,6 +19,7 @@ PrepareStoreOutput, ReqContext, RequestOffloadingContext, + get_offload_group_idx, ) from vllm.v1.kv_offload.cpu.common import ( CPULoadStoreSpec, @@ -27,6 +29,20 @@ from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +@dataclass(slots=True) +class _RequestCacheAccess: + """Cache keys observed by one request, grouped in prefix order.""" + + owner: object + cache_generation: int + key_groups: dict[int, list[OffloadKey]] = field(default_factory=dict) + seen_keys: set[OffloadKey] = field(default_factory=set) + inserted_keys: set[OffloadKey] = field(default_factory=set) + reused_keys: set[OffloadKey] = field(default_factory=set) + store_miss_keys: set[OffloadKey] = field(default_factory=set) + finished: bool = False + + class CPUOffloadingManager(OffloadingManager): """ An OffloadingManager with a pluggable CachePolicy, resolved by name via @@ -66,6 +82,7 @@ def __init__( self.max_tracker_size: int = max_tracker_size self.stores_skipped_in_current_batch: int = 0 self.allocation_sizes_in_current_batch: list[int] = [] + self._cache_generation = 0 # Number of chunk references. It is ordered so can evict the LRU entry in O(1). self.counts: OrderedDict[OffloadKey, int] | None = ( @@ -124,10 +141,48 @@ def _record_accesses(self, keys: Collection[OffloadKey]) -> None: num_unprotected -= 1 self.counts[key] = 1 + def _get_request_cache_access(self, req_context: ReqContext) -> _RequestCacheAccess: + state = req_context.get_state(_RequestCacheAccess) + if ( + state is None + or state.owner is not self + or state.cache_generation != self._cache_generation + ): + state = _RequestCacheAccess( + owner=self, cache_generation=self._cache_generation + ) + req_context.set_state(state) + return state + + def _record_request_cache_access( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + inserted_keys: Iterable[OffloadKey] = (), + reused_keys: Iterable[OffloadKey] = (), + ) -> None: + state = self._get_request_cache_access(req_context) + for key in keys: + if key in state.seen_keys: + continue + assert not state.finished, ( + "New cache keys observed after request finalization" + ) + group_idx = get_offload_group_idx(key) + state.key_groups.setdefault(group_idx, []).append(key) + state.seen_keys.add(key) + state.inserted_keys.update(inserted_keys) + # Re-reading a chunk inserted by this request is an internal transfer + # (for example, a tiering cascade), not a second cache access. + state.reused_keys.update( + key for key in reused_keys if key not in state.inserted_keys + ) + # --- OffloadingManager interface --- @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + self._get_request_cache_access(req_context) return RequestOffloadingContext() @override @@ -144,6 +199,15 @@ def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext, + ) -> LoadStoreSpec: + return self._prepare_load(keys, req_context, record_access=True) + + def _prepare_load( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + *, + record_access: bool, ) -> LoadStoreSpec: chunks = [] for key in keys: @@ -156,6 +220,8 @@ def prepare_load( assert self._num_evictable_cache_chunks >= 0 chunk.ref_cnt += 1 chunks.append(chunk) + if record_access: + self._record_request_cache_access(keys, req_context, reused_keys=keys) return self._get_load_store_spec(keys, chunks) @override @@ -181,13 +247,52 @@ def prepare_store( keys: Collection[OffloadKey], req_context: ReqContext, ) -> PrepareStoreOutput | None: + keys = list(keys) if self.counts is not None: - num_keys = len(keys) self._record_accesses(keys) - keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold] - self.stores_skipped_in_current_batch += num_keys - len(keys) - # filter out chunks that are already stored - keys_to_store = [k for k in keys if self._policy.get(k) is None] + # Partition the original offer before admission filtering. Ready + # resident chunks are request reuses even if their tracker entry was + # aged out; the threshold applies only to new store candidates. + keys_to_store: list[OffloadKey] = [] + ready_existing_keys: list[OffloadKey] = [] + for key in keys: + chunk = self._policy.get(key) + if chunk is None: + keys_to_store.append(key) + else: + if chunk.is_ready: + ready_existing_keys.append(key) + + if self.counts is not None: + num_store_candidates = len(keys_to_store) + keys_to_store = [ + key + for key in keys_to_store + if self.counts.get(key, 0) >= self.store_threshold + ] + self.stores_skipped_in_current_batch += num_store_candidates - len( + keys_to_store + ) + + state = self._get_request_cache_access(req_context) + new_store_misses = [ + key for key in keys_to_store if key not in state.store_miss_keys + ] + if new_store_misses: + # ARC learns from B1/B2 before insert() removes the ghost entry. + # Deduplication makes this one policy observation per request. + self._policy.on_store_miss(new_store_misses, req_context) + state.store_miss_keys.update(new_store_misses) + + recorded_keys = set(keys_to_store) + recorded_keys.update(ready_existing_keys) + # Record once before any allocation-related early return. Store + # candidates are classified as insertions only after allocation succeeds. + self._record_request_cache_access( + (key for key in keys if key in recorded_keys), + req_context, + reused_keys=ready_existing_keys, + ) if not keys_to_store: return PrepareStoreOutput( @@ -239,6 +344,7 @@ def prepare_store( for key, chunk in zip(keys_to_store, chunks): self._policy.insert(key, chunk) self._num_write_pending_chunks += len(keys_to_store) + state.inserted_keys.update(keys_to_store) # build store specs for allocated chunks store_spec = self._get_load_store_spec(keys_to_store, chunks) @@ -284,6 +390,35 @@ def complete_store( ) ) + @override + def on_request_finished(self, req_context: ReqContext) -> None: + state = req_context.get_state(_RequestCacheAccess) + if ( + state is None + or state.owner is not self + or state.cache_generation != self._cache_generation + or state.finished + ): + return + state.finished = True + key_groups = [] + for group_idx in sorted(state.key_groups): + keys = state.key_groups[group_idx] + positions = { + key: position + for key in keys + if (position := req_context.get_offload_key_position(key)) is not None + } + if len(positions) == len(keys): + keys = sorted(keys, key=positions.__getitem__) + key_groups.append(tuple(keys)) + self._policy.on_request_finished( + tuple(key_groups), + state.inserted_keys - state.reused_keys, + state.reused_keys, + req_context, + ) + @override def reset_cache(self) -> None: # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold @@ -292,6 +427,7 @@ def reset_cache(self) -> None: # flushes in-flight load job IDs to the workers before any new stores # can begin, preventing a cross-direction data race on reused offload chunk IDs. self._policy.clear() + self._cache_generation += 1 self._num_evictable_cache_chunks = 0 self._num_write_pending_chunks = 0 diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index e413bd603eb4..882ee02219c2 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Sequence from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus +from vllm.v1.kv_offload.cpu.policies.base import ( + CachePolicy, + ChunkStatus, + order_request_keys, +) class ARCCachePolicy(CachePolicy): @@ -24,12 +28,11 @@ class ARCCachePolicy(CachePolicy): Searches T1 and T2 for chunk hashes and counts consecutive hits until a miss or non-ready chunk is encountered. - 2. Cache touch (touch) - Adaptive Learning: - For each key (in reverse order): - - If in T1: Move to T2 (promotion from recent to frequent). - - If in T2: Move to MRU position (end of queue). - - If in B1 ghost list: Increase target_t1_size. - - If in B2 ghost list: Decrease target_t1_size. + 2. Request access - Adaptive Learning: + - Ready chunks reused by a request move from T1 to T2 once, or + move to the MRU end of T2. + - B1/B2 misses adjust target_t1_size once per request before + insertion removes their ghost entries. 3. Chunk eviction (evict) - Adaptive Replacement: Determines eviction source based on adaptive target: @@ -39,7 +42,7 @@ class ARCCachePolicy(CachePolicy): 4. Chunk insertion (insert): New chunks are always inserted into T1 and removed from B1/B2 if - present. Chunks may later be promoted to T2 during touch operations. + present. A later request reuse may promote them to T2. Adaptive Behavior: The algorithm self-tunes the recency vs. frequency trade-off: @@ -71,6 +74,26 @@ def remove(self, key: OffloadKey) -> None: if self.t1.pop(key, None) is None: self.t2.pop(key, None) + def _adapt_to_ghost_hit(self, key: OffloadKey) -> bool: + if key in self.b1: + delta = max(1, len(self.b2) / len(self.b1)) + self.target_t1_size = min(self.target_t1_size + delta, self.cache_capacity) + self.b1.move_to_end(key) + return True + if key in self.b2: + delta = max(1, len(self.b1) / len(self.b2)) + self.target_t1_size = max(self.target_t1_size - delta, 0) + self.b2.move_to_end(key) + return True + return False + + @override + def on_store_miss( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> None: + for key in reversed(list(keys)): + self._adapt_to_ghost_hit(key) + @override def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): @@ -86,19 +109,33 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: elif key in self.t2: self.t2.move_to_end(key) - elif key in self.b1: - delta = max(1, len(self.b2) / len(self.b1)) - self.target_t1_size = min( - self.target_t1_size + delta, self.cache_capacity - ) - # move to MRU position (end) to keep it fresh in the ghost list - self.b1.move_to_end(key) - - elif key in self.b2: - delta = max(1, len(self.b1) / len(self.b2)) - self.target_t1_size = max(self.target_t1_size - delta, 0) - # move to MRU position (end) to keep it fresh in the ghost list - self.b2.move_to_end(key) + else: + self._adapt_to_ghost_hit(key) + + @override + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + for key in reversed(order_request_keys(key_groups, req_context)): + if key in insertion_only_keys: + # A store is not a frequency hit. Preserve T1 membership + # while still restoring tail-to-head recency. + if key in self.t1: + self.t1.move_to_end(key) + elif key in self.t2: + self.t2.move_to_end(key) + continue + + # Ready chunks reused by this request count as one access. + if key in reused_keys and key in self.t1: + chunk = self.t1.pop(key) + self.t2[key] = chunk + elif key in reused_keys and key in self.t2: + self.t2.move_to_end(key) @override def clear(self) -> None: diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 4059dd742f7c..3efffbd3961a 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -2,11 +2,26 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import ctypes from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from vllm.v1.kv_offload.base import OffloadKey, ReqContext +def order_request_keys( + key_groups: Sequence[Sequence[OffloadKey]], req_context: ReqContext +) -> list[OffloadKey]: + """Return one head-to-tail order across all KV cache groups.""" + keys = [key for group in key_groups for key in group] + positions = { + key: position + for key in keys + if (position := req_context.get_offload_key_position(key)) is not None + } + if len(positions) == len(keys): + keys.sort(key=positions.__getitem__) + return keys + + class ChunkStatus(ctypes.Structure): """ Offloading status for a single chunk of KV data. @@ -66,6 +81,41 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: req_context: Per-request context for the request touching these chunks. """ + def on_store_miss( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> None: + """Observe store misses before their cache entries are inserted. + + The default delegates to ``touch`` for compatibility with external + policies. Policies may override this hook when a miss has distinct + semantics, such as ARC adapting to a ghost-list hit. + """ + self.touch(keys, req_context) + + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + """Apply one request-scoped cache access in prefix order. + + ``key_groups`` contains keys observed for each KV cache group in + head-to-tail order. ``insertion_only_keys`` and ``reused_keys`` + distinguish chunks only created by this request from ready chunks it + actually reused, which matters for policies such as ARC where reuse + changes frequency but insertion and pending observations do not. + + The default forwards one touch per group, preserving compatibility + for experimental out-of-tree policies while moving those touches to + request finalization. Policies that distinguish insertion from reuse + can override this hook and inspect the access classifications. + """ + del insertion_only_keys, reused_keys + for keys in key_groups: + self.touch(keys, req_context) + @abstractmethod def evict( self, n: int, protected: set[OffloadKey] diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index c1dc631d4fe9..5b5f65dec779 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -1,28 +1,74 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections import OrderedDict -from collections.abc import Iterable +import heapq +from collections.abc import Iterable, Sequence from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus +from vllm.v1.kv_offload.cpu.policies.base import ( + CachePolicy, + ChunkStatus, + order_request_keys, +) class LRUCachePolicy(CachePolicy): """ - LRU Caching policy that keeps a dedicated evictable list for fast eviction. + LRU caching policy with logical recency independent of transfer pinning. + + Evictable chunks live in a lazy-invalidating min-heap. A chunk's recency + can therefore be updated while it is pinned; when it later becomes + evictable, it enters the heap with the order assigned by the request + rather than its transfer-completion time. + A use is indicated by, - First time the key is added (store). - - Load job completion - - touch + - A request-scoped access. """ def __init__(self, cache_capacity: int): super().__init__(cache_capacity) - # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU - self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict() self.chunks: dict[OffloadKey, ChunkStatus] = {} + self._ranks: dict[OffloadKey, int] = {} + self._evictable: set[OffloadKey] = set() + self._heap: list[tuple[int, OffloadKey]] = [] + self._next_rank = 0 + + def _assign_new_rank(self, key: OffloadKey) -> bool: + if key not in self.chunks: + return False + self._next_rank += 1 + self._ranks[key] = self._next_rank + return key in self._evictable + + def _push_evictable(self, key: OffloadKey) -> None: + heapq.heappush( + self._heap, + (self._ranks[key], key), + ) + + def _is_current(self, entry: tuple[int, OffloadKey]) -> bool: + rank, key = entry + return key in self._evictable and self._ranks.get(key) == rank + + def _maybe_compact_heap(self) -> None: + if len(self._heap) < max(64, 2 * len(self._evictable)): + return + self._heap = [(self._ranks[key], key) for key in self._evictable] + heapq.heapify(self._heap) + + def _update_recency(self, keys: Iterable[OffloadKey]) -> None: + updated_evictable = [key for key in keys if self._assign_new_rank(key)] + # Rebuilding is linear and substantially cheaper than k heap pushes + # for a long prefix. Small updates retain the incremental path. + if len(updated_evictable) >= max(64, len(self._evictable) // 4): + self._heap = [(self._ranks[key], key) for key in self._evictable] + heapq.heapify(self._heap) + else: + for key in updated_evictable: + self._push_evictable(key) + self._maybe_compact_heap() @override def get(self, key: OffloadKey) -> ChunkStatus | None: @@ -31,26 +77,40 @@ def get(self, key: OffloadKey) -> ChunkStatus | None: @override def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: self.chunks[key] = chunk + self._next_rank += 1 + self._ranks[key] = self._next_rank if chunk.ref_cnt == 0: - self.evictable_chunks[key] = None + self._evictable.add(key) + self._push_evictable(key) @override def remove(self, key: OffloadKey) -> None: del self.chunks[key] - self.evictable_chunks.pop(key, None) + self._ranks.pop(key, None) + self._evictable.discard(key) @override def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: - for key in reversed(list(keys)): - if key in self.evictable_chunks: - self.evictable_chunks.move_to_end(key) - # active chunks are untouched as they are non-evictable now. They - # will eventually reach the end of evictable_chunks when they finish. + self._update_recency(reversed(list(keys))) + + @override + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + del insertion_only_keys, reused_keys + self._update_recency(reversed(order_request_keys(key_groups, req_context))) @override def clear(self) -> None: - self.evictable_chunks.clear() self.chunks.clear() + self._ranks.clear() + self._evictable.clear() + self._heap.clear() + self._next_rank = 0 @override def evict( @@ -59,22 +119,44 @@ def evict( if n == 0: return [] - candidates: list[tuple[OffloadKey, ChunkStatus]] = [] - for key, _ in self.evictable_chunks.items(): + selected: list[tuple[tuple[int, OffloadKey], ChunkStatus]] = [] + selected_keys: set[OffloadKey] = set() + deferred: list[tuple[int, OffloadKey]] = [] + while self._heap and len(selected) < n: + entry = heapq.heappop(self._heap) + if not self._is_current(entry): + continue + key = entry[1] + # Re-pinning without a recency change can leave an equivalent + # lazy entry behind. Select each cache key at most once. + if key in selected_keys: + continue if key in protected: + deferred.append(entry) continue - chunk = self.chunks[key] assert chunk.ref_cnt == 0 - candidates.append((key, chunk)) - if len(candidates) == n: - break - - if len(candidates) < n: + selected.append((entry, chunk)) + selected_keys.add(key) + + if len(selected) < n: + for entry, _ in selected: + heapq.heappush(self._heap, entry) + for entry in deferred: + heapq.heappush(self._heap, entry) return None - for key, _ in candidates: - del self.evictable_chunks[key] + + for entry in deferred: + heapq.heappush(self._heap, entry) + + candidates: list[tuple[OffloadKey, ChunkStatus]] = [] + for entry, chunk in selected: + key = entry[1] + self._evictable.remove(key) del self.chunks[key] + del self._ranks[key] + candidates.append((key, chunk)) + self._maybe_compact_heap() return candidates @override @@ -82,9 +164,11 @@ def mark_evictable(self, key: OffloadKey) -> None: # chunks can become evictable when, # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list - self.evictable_chunks[key] = None + self._evictable.add(key) + self._push_evictable(key) + self._maybe_compact_heap() @override def mark_non_evictable(self, key: OffloadKey) -> None: - # key must have been in the evictable list. - del self.evictable_chunks[key] + # key must have been evictable before it was pinned. + self._evictable.remove(key) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 0940aaeae07b..f19616828daa 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -110,13 +110,22 @@ def __init__( # read/write is for CPU<->secondary transfers, # load/store is for CPU<->GPU transfers. # These aliases avoid calling prepare_load inside a store path. - self.prepare_read = self.prepare_load self.complete_read = self.complete_load self.prepare_write = self.prepare_store self.complete_write = self.complete_store self._kv_memoryview = mmap_region.create_kv_memoryview() + def prepare_read( + self, keys: Collection[OffloadKey], req_context: ReqContext + ) -> LoadStoreSpec: + """Pin chunks for a CPU-to-secondary transfer. + + Cascade reads are implementation details of tiering, not additional + request accesses, so they must not alter request-scoped recency. + """ + return self._prepare_load(keys, req_context, record_access=False) + def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. @@ -781,11 +790,10 @@ def _maybe_finalize_request( req_id: str, exclude_tier_idx: int | None = None, ) -> None: - """Finalize secondary tiers once no more store cascades can be submitted. + """Finalize secondary tiers once no more cascades can be submitted. - Finalization means forwarding on_request_finished() to secondary tiers. - It is delayed until pending GPU->primary stores finish, since their - complete_store() callbacks may still submit primary->secondary stores. + Their finalization is delayed until pending GPU->primary stores + finish, since those callbacks may still submit secondary stores. """ state = self._req_state[req_id] if not state.is_finished: From b68408043d80faf9c99fc16c6aec1174336e9626 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Wed, 16 Sep 2026 15:49:17 +0800 Subject: [PATCH 307/774] [Bugfix][CPU] Add per-tensor FP8 W8A16 kernel to fix Ministral crash on CPU (#56985) Signed-off-by: jiang1.li --- .../scripts/hardware_ci/run-cpu-test.sh | 21 +++- .../quantization/test_cpu_fp8_scaled_mm.py | 113 ++++++++++++++++++ tests/v1/e2e/test_cpu_spec_decode.py | 2 +- .../model_executor/kernels/linear/__init__.py | 2 + .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/cpu.py | 100 ++++++++++++++++ 6 files changed, 233 insertions(+), 7 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index e8d0a24bc7d1..43555a19ab58 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -83,14 +83,23 @@ rm -f "$build_log" # Run the image, setting --shm-size=4g for tensor parallel. Default to # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network; -# retry once online if the cache is missing something. vllm's get_config() -# wraps the raw huggingface_hub offline-mode errors in a generic ValueError -# (see transformers_utils/config.py), so match that message too or the -# fallback never triggers for a config-lookup cache miss. -OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified' +# retry once online if the cache is missing something. vllm wraps offline +# cache-miss errors in ways that don't preserve the raw huggingface_hub +# exception name: get_config() re-raises as a generic ValueError +# (transformers_utils/config.py), and default_loader.py raises its own +# RuntimeError when a snapshot dir has a config but no weight files. Match +# those messages too or the fallback never triggers for those cache misses. +# +# Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via +# vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets +# by default), a directory distinct from the HF hub cache above. Without a +# persistent mount there, those assets never survive past one container's +# lifetime, so every run re-fetches them from the network on the "offline" +# attempt and unconditionally falls back to the online retry. +OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with' run_test() { local hf_offline=$1 - docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ + docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" } diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index 3154e2cb98bb..b70a1a8543ba 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -9,6 +9,15 @@ import torch from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( + CPUFp8PerTensorScaledMMLinearKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, +) from vllm.platforms import current_platform if not current_platform.is_cpu(): @@ -17,6 +26,9 @@ if not ops._supports_cpu_fp8_w8a16: pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True) +if not torch.cpu._is_amx_tile_supported(): + pytest.skip("requires AMX tile support", allow_module_level=True) + BLOCK_SIZE = [128, 128] @@ -160,3 +172,104 @@ def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool): assert kernel_out.dtype == out_dtype torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01) + + +def quantize_weight_per_tensor_fp8( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with a single per-tensor scale.""" + fp8_max = torch.finfo(torch.float8_e4m3fn).max + scale = weight.abs().amax() / fp8_max + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + q = (weight / scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + return q, scale + + +def ref_fp8_per_tensor_scaled_mm( + x: torch.Tensor, + fp8_weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + w_dq = fp8_weight.float() * scale + out = torch.mm(x.float(), w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +NK_SIZES_PER_TENSOR = [ + (32, 64), + (5120, 5120), +] + + +@pytest.mark.parametrize("M", [1, 64]) +@pytest.mark.parametrize("N,K", NK_SIZES_PER_TENSOR) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_per_tensor_scaled_mm_kernel( + M: int, N: int, K: int, use_bias: bool, default_vllm_config +): + """CPUFp8PerTensorScaledMMLinearKernel correctness against float reference. + + Exercises the full kernel class (not just the raw op), including the + weight-orientation fixup: Fp8LinearMethod stores `layer.weight` as + [K, N] (torch._scaled_mm convention) before calling + process_weights_after_loading, so the kernel must transpose back to + [N, K] before VNNI-packing. + """ + torch.manual_seed(0) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=out_dtype) / (K**0.5) + w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5) + fp8_weight, scale = quantize_weight_per_tensor_fp8(w_f32) + bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None + + ref_out = ref_fp8_per_tensor_scaled_mm(x, fp8_weight, scale, bias, out_dtype) + + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=kFp8StaticTensorSym, + weight_shape=(N, K), + input_dtype=out_dtype, + out_dtype=out_dtype, + ) + kernel = CPUFp8PerTensorScaledMMLinearKernel( + config, + layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"], + ) + + layer = torch.nn.Module() + # Fp8LinearMethod stores the weight transposed to [K, N] before calling + # process_weights_after_loading. + layer.register_parameter( + "weight", torch.nn.Parameter(fp8_weight.t().contiguous(), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(scale.clone(), requires_grad=False) + ) + kernel.process_weights_after_loading(layer) + + kernel_out = kernel.apply_weights(layer, x, bias) + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01) + + +@pytest.mark.parametrize("n", [16, 48]) +def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int): + """The AMX tinygemm kernel tiles N in chunks of 32 and cannot handle a + remainder tile smaller than that, so can_implement must reject any N + that isn't a multiple of 32 rather than let it crash the process. + """ + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=kFp8StaticTensorSym, + weight_shape=(n, 64), + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + ) + supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config) + assert not supported diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py index 2d34c1165625..a7ba5580c33b 100644 --- a/tests/v1/e2e/test_cpu_spec_decode.py +++ b/tests/v1/e2e/test_cpu_spec_decode.py @@ -131,7 +131,7 @@ def test_ngram_spec_decode_matches_baseline(baseline_refs): id="dflash", ), pytest.param( - "Qwen/Qwen3.5-0.8B-Base", + "Qwen/Qwen3.5-0.8B", { "method": "mtp", "num_speculative_tokens": 3, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index d3c41f97d8e5..c610e9b9c8f0 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -187,6 +187,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, + CPUFp8PerTensorScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -438,6 +439,7 @@ def _resolve_backend_kernels( ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.CPU: [ + CPUFp8PerTensorScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, ], diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 39f9abd460e1..f4a062b7f2b3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -9,6 +9,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, + CPUFp8PerTensorScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -67,5 +68,6 @@ "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", + "CPUFp8PerTensorScaledMMLinearKernel", "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 92ea073a1abb..e72dffcbbc50 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -21,10 +21,14 @@ FP8ScaledMMLinearLayerConfig, ) from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, ) +# BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size. +_BLOCK_SIZE_K = 128 + class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): @classmethod @@ -330,3 +334,99 @@ def apply_block_scaled_mm( raise NotImplementedError( "CPUFp8BlockScaledMMKernel overrides apply_weights directly." ) + + +class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU. + + Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single + synthetic block spanning the whole weight, so activations stay BF16/FP32 + — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel. + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops._supports_cpu_fp8_w8a16: + return False, "fp8_scaled_mm_cpu op not available." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + if not c.weight_quant_key.scale.group_shape.is_per_tensor(): + return False, "requires a per-tensor weight scale." + if not c.activation_quant_key.scale.group_shape.is_per_tensor(): + return False, "requires a per-tensor activation scale." + if c.out_dtype not in (torch.bfloat16, torch.float32): + return False, "Only bfloat16/float32 output dtype supported." + n = c.weight_shape[0] + if n % 32 != 0: + # The AMX tinygemm kernel tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_name, w_s_name, _, _ = self.layer_param_names + # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style + # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as + # produced by CUDA nn.Linear-style checkpoints. + weight = getattr(layer, w_name).t().contiguous() + n, k = weight.shape + + packed_weight = torch.ops._C.convert_weight_packed(weight) + replace_parameter( + layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False) + ) + + # Synthesize a single-block "block scale" tensor from the scalar + # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is. + weight_scale = getattr(layer, w_s_name) + self._block_size_n = -(-n // 32) * 32 # round up to a multiple of 32 + num_k_blocks = -(-k // _BLOCK_SIZE_K) + block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous() + replace_parameter( + layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False) + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + w_name, w_s_name, _, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, w_s_name) + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + out = torch.ops._C.fp8_scaled_mm_cpu( + x_2d, + weight, + weight_scale, + [self._block_size_n, _BLOCK_SIZE_K], + bias, + x.dtype, + True, # is_vnni (weight already prepacked) + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + raise NotImplementedError( + "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly." + ) From 9ca6dbba71329a7b08711e2acfa17b31e542553c Mon Sep 17 00:00:00 2001 From: Ama Senevirathne <97525823+amasen02@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:26:50 +0530 Subject: [PATCH 308/774] fix(multimodal): tolerate malformed EXIF metadata during hashing (#56527) (#56576) Signed-off-by: amasen02 Co-authored-by: amasen02 --- tests/multimodal/test_hasher.py | 31 ++++++++++++++++++++++++++++++- vllm/multimodal/hasher.py | 29 +++++++++++++++++++---------- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index d0be316790d2..1b738eeeb89f 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import hashlib +import struct import uuid from io import BytesIO from pathlib import Path @@ -8,7 +10,7 @@ import numpy as np import pytest import torch -from PIL import Image, ImageDraw +from PIL import Image, ImageDraw, ImageOps from vllm.config.multimodal import MMHasherAlgorithm from vllm.multimodal.hasher import MultiModalHasher @@ -199,6 +201,33 @@ def test_hash_image_exif_id(): ) +def test_hash_image_malformed_exif(): + # Test that images with malformed EXIF headers (e.g. invalid TIFF header) + # do not raise an unhandled exception during hashing and fall back to image data. + buf = BytesIO() + Image.new("RGB", (64, 48)).save(buf, "JPEG") + jpg = buf.getvalue() + rest = jpg[2:] + rest = rest[2 + struct.unpack(">H", rest[2:4])[0] :] + payload = b"Exif\x00\x00XXXX\x00\x00\x00\x08" + bytes(32) + data = b"\xff\xd8\xff\xe1" + struct.pack(">H", len(payload) + 2) + payload + rest + + image = Image.open(BytesIO(data)) + with contextlib.suppress(Exception): + image = ImageOps.exif_transpose(image) + image.load() + + hasher = MultiModalHasher + # Should hash without raising SyntaxError or any other exception + hash_val = hasher.hash_kwargs("blake3", image=image) + assert isinstance(hash_val, str) and len(hash_val) > 0 + + # Also verify MediaWithBytes wrapping the image with malformed EXIF + media_item = MediaWithBytes(image, data) + hash_media = hasher.hash_kwargs("blake3", image=media_item) + assert isinstance(hash_media, str) and len(hash_media) > 0 + + def _rgba_png_bytes() -> bytes: image = Image.new("RGBA", (8, 8), (255, 0, 0, 128)) buf = BytesIO() diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index a432afdc9bbd..3750ae162bb0 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -48,6 +48,19 @@ def _get_hasher_factory( raise ValueError(f"Unsupported hash algorithm: {algorithm}") +def _get_image_id_bytes(image: Image.Image) -> bytes | None: + try: + exif = image.getexif() + image_id = exif.get(Image.ExifTags.Base.ImageID) + if isinstance(image_id, uuid.UUID): + return image_id.bytes + except Exception: + # Tolerate malformed EXIF metadata (e.g. invalid TIFF header) + # and fall back to serializing raw image data or bytes. + pass + return None + + class MultiModalHasher: @classmethod def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: @@ -60,11 +73,9 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: return (np.array(obj).tobytes(),) if isinstance(obj, Image.Image): - exif = obj.getexif() - if Image.ExifTags.Base.ImageID in exif and isinstance( - exif[Image.ExifTags.Base.ImageID], uuid.UUID - ): - return (exif[Image.ExifTags.Base.ImageID].bytes,) + image_id = _get_image_id_bytes(obj) + if image_id is not None: + return (image_id,) data = {"mode": obj.mode, "data": np.asarray(obj)} palette = obj.palette @@ -76,11 +87,9 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: return cls.iter_item_to_bytes("image", data) if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image): - exif = obj.media.getexif() - if Image.ExifTags.Base.ImageID in exif and isinstance( - exif[Image.ExifTags.Base.ImageID], uuid.UUID - ): - return (exif[Image.ExifTags.Base.ImageID].bytes,) + image_id = _get_image_id_bytes(obj.media) + if image_id is not None: + return (image_id,) if obj.io_config: return cls.iter_item_to_bytes( From 8be5205abbabf4c377c603d6c4180a99373f6415 Mon Sep 17 00:00:00 2001 From: dev <92231400+devtyagi3909@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:47:22 +0530 Subject: [PATCH 309/774] [Parser] Fix: correct parser frontend handling of reasoning end and boundary tokens (#56635) Signed-off-by: DEV TYAGI --- tests/parser/engine/test_glm47_moe.py | 108 ++++++++++++++++++++++++++ vllm/parser/engine/parser_engine.py | 19 +++++ vllm/parser/glm47_moe.py | 14 ++-- 3 files changed, 135 insertions(+), 6 deletions(-) create mode 100644 tests/parser/engine/test_glm47_moe.py diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py new file mode 100644 index 000000000000..58fe7aa28cbd --- /dev/null +++ b/tests/parser/engine/test_glm47_moe.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning-boundary tests for the GLM-4.7 parser. + +GLM can transition straight from reasoning into a tool call without ever +emitting ````, so ```` is an implicit reasoning +terminator. These cover that path plus the multi-turn prompts where an +earlier turn's markers must not be read as the current turn's state. +""" + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from vllm.parser.glm47_moe import ( + THINK_END, + THINK_START, + TOOL_CALL_END, + TOOL_CALL_START, + Glm47MoeParser, +) + +THINK_S, THINK_E = 1000, 1001 +TOOL_S, TOOL_E = 1002, 1003 +ASSISTANT, OBSERVATION, USER = 1004, 1005, 1006 +TEXT = 42 # stand-in for an ordinary reasoning/content token + +VOCAB = { + THINK_START: THINK_S, + THINK_END: THINK_E, + TOOL_CALL_START: TOOL_S, + TOOL_CALL_END: TOOL_E, + "<|assistant|>": ASSISTANT, + "<|observation|>": OBSERVATION, + "<|user|>": USER, +} + + +@pytest.fixture +def parser(): + return Glm47MoeParser(make_mock_tokenizer(VOCAB)) + + +@pytest.fixture +def no_thinking_parser(): + return Glm47MoeParser( + make_mock_tokenizer(VOCAB), + chat_template_kwargs={"thinking": False}, + ) + + +class TestIsReasoningEnd: + def test_open_reasoning(self, parser): + assert not parser.is_reasoning_end([THINK_S, TEXT]) + + def test_think_end(self, parser): + assert parser.is_reasoning_end([THINK_S, TEXT, THINK_E, TEXT]) + + def test_tool_call_without_think_end(self, parser): + """Reasoning -> tool call with no ```` still ends reasoning.""" + assert parser.is_reasoning_end([THINK_S, TEXT, TOOL_S, TEXT]) + + def test_previous_turn_tool_call_ignored(self, parser): + """A finished tool call from an earlier turn says nothing about the + turn currently being generated.""" + prompt = [THINK_S, TEXT, TOOL_S, TEXT, TOOL_E, OBSERVATION, TEXT, ASSISTANT] + assert not parser.is_reasoning_end(prompt) + + def test_previous_turn_think_end_ignored(self, parser): + prompt = [THINK_S, TEXT, THINK_E, TEXT, USER, TEXT, ASSISTANT] + assert not parser.is_reasoning_end(prompt) + + def test_empty_input(self, parser): + assert not parser.is_reasoning_end([]) + + def test_thinking_disabled(self, no_thinking_parser): + assert no_thinking_parser.is_reasoning_end([THINK_S, TEXT]) + + +class TestExtractContentIds: + def test_think_end_wins_over_later_tool_call(self, parser): + """Content between ```` and a tool call must survive.""" + ids = [THINK_S, TEXT, THINK_E, 20, 21, TOOL_S, 30, TOOL_E] + assert parser.extract_content_ids(ids) == [20, 21, TOOL_S, 30, TOOL_E] + + def test_every_tool_call_after_think_end_kept(self, parser): + ids = [THINK_E, 20, TOOL_S, 30, TOOL_E, 21, TOOL_S, 31, TOOL_E] + assert parser.extract_content_ids(ids) == ids[1:] + + def test_falls_back_to_tool_call(self, parser): + """Without ````, content starts at the opener itself so the + tool parser still receives a well-formed call.""" + ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E] + assert parser.extract_content_ids(ids) == [TOOL_S, 30, TOOL_E] + + def test_falls_back_to_first_tool_call_of_turn(self, parser): + ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E, TOOL_S, 31, TOOL_E] + assert parser.extract_content_ids(ids) == ids[2:] + + def test_previous_turn_tool_call_ignored(self, parser): + ids = [TOOL_S, 30, TOOL_E, OBSERVATION, TEXT, ASSISTANT, THINK_S, TEXT] + assert parser.extract_content_ids(ids) == ids + + def test_no_markers_returns_input_ids(self, parser): + assert parser.extract_content_ids([20, 21]) == [20, 21] + + def test_thinking_disabled(self, no_thinking_parser): + ids = [THINK_S, TEXT, TOOL_S] + assert no_thinking_parser.extract_content_ids(ids) == ids diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index fa306a5da4a1..527c310c1485 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -672,11 +672,30 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: return not wait_for_reasoning def extract_content_ids(self, input_ids: list[int]) -> list[int]: + config = self.parser_engine_config + wait_for_reasoning = config.wait_for_reasoning + if wait_for_reasoning is None: + wait_for_reasoning = config.initial_state is ParserState.REASONING + if not wait_for_reasoning: + return input_ids + end_id = self._reasoning_end_token_id if end_id is not None: for i in range(len(input_ids) - 1, -1, -1): if input_ids[i] == end_id: return input_ids[i + 1 :] + + end_ids = self._reasoning_end_token_ids + if end_ids: + turn_start = 0 + boundary_ids = self._turn_boundary_token_ids + for i in range(len(input_ids) - 1, -1, -1): + if input_ids[i] in boundary_ids: + turn_start = i + 1 + break + for i in range(turn_start, len(input_ids)): + if input_ids[i] in end_ids: + return input_ids[i:] return input_ids def get_streaming_fallback_content( diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 4a804a4c9c94..00e517f2ea68 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -41,6 +41,13 @@ ARG_VALUE_START = "" ARG_VALUE_END = "" +# Special tokens that delimit conversation turns in a rendered GLM prompt. +# Reasoning markers belonging to earlier turns must not be mistaken for the +# state of the turn currently being generated. +GLM_TURN_BOUNDARIES = frozenset( + ("<|system|>", "<|user|>", "<|assistant|>", "<|observation|>") +) + _ARG_RE = re.compile( r"(?P.*?)\s*" r"(?P.*?)", @@ -123,7 +130,6 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="glm47_moe", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, - wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_START": TOOL_CALL_START, @@ -166,6 +172,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: ), **arg_tag_transitions, }, + turn_boundary_tokens=GLM_TURN_BOUNDARIES, arg_converter=_glm47_arg_converter, stream_arg_deltas=True, tool_args_json=False, @@ -207,11 +214,6 @@ def _handle_tool_end(self, event, deltas) -> None: self._tool_slots[idx].name = self._tool_slots[idx].name.strip() super()._handle_tool_end(event, deltas) - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - if not self.thinking_enabled: - return input_ids - return super().extract_content_ids(input_ids) - def extract_reasoning( self, model_output: str, From 0d8173d1537f60331084e12af6b72a6565978c9d Mon Sep 17 00:00:00 2001 From: wenjinhust <33047235+wenjinhust@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:25:49 +0800 Subject: [PATCH 310/774] [Bugfix][Frontend] Lazy-import model_hosting_container_standards to prevent log suppression (#51366) Signed-off-by: wenjinhust --- vllm/entrypoints/serve/lora/api_router.py | 30 ++++++++- .../entrypoints/serve/sagemaker/api_router.py | 61 ++++++++++++++++++- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/vllm/entrypoints/serve/lora/api_router.py b/vllm/entrypoints/serve/lora/api_router.py index 95dff9792b2f..61436371c27d 100644 --- a/vllm/entrypoints/serve/lora/api_router.py +++ b/vllm/entrypoints/serve/lora/api_router.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging - -import model_hosting_container_standards.sagemaker as sagemaker_standards from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse, Response @@ -22,6 +21,12 @@ def attach_router(app: FastAPI): + """Attach the LoRA adapter load/unload endpoints to the API server. + + Does nothing when dynamic LoRA updating is disabled. Handler levels are + snapshotted and restored because importing + model_hosting_container_standards may reconfigure root logging. + """ if not envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING: """If LoRA dynamic loading & unloading is not enabled, do nothing.""" return @@ -30,6 +35,23 @@ def attach_router(app: FastAPI): "This should ONLY be used for local development!" ) + snapshot = [ + (h, h.level) + for lg in (logging.getLogger(), logging.getLogger("vllm")) + for h in lg.handlers + ] + + try: + _attach_router(app) + finally: + for handler, level in snapshot: + handler.setLevel(level) + + +def _attach_router(app: FastAPI): + """Register the LoRA adapter load/unload routes on the app.""" + import model_hosting_container_standards.sagemaker as sagemaker_standards + @sagemaker_standards.register_load_adapter_handler( request_shape={ "lora_name": "body.name", @@ -40,6 +62,8 @@ def attach_router(app: FastAPI): ) @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)]) async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request): + """Handle POST /v1/load_lora_adapter: load a LoRA adapter into the + serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.load_lora_adapter(request) if isinstance(response, ErrorResponse): @@ -60,6 +84,8 @@ async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Reques async def unload_lora_adapter( request: UnloadLoRAAdapterRequest, raw_request: Request ): + """Handle POST /v1/unload_lora_adapter: unload a LoRA adapter from + the serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.unload_lora_adapter(request) if isinstance(response, ErrorResponse): diff --git a/vllm/entrypoints/serve/sagemaker/api_router.py b/vllm/entrypoints/serve/sagemaker/api_router.py index 142256c26002..60576f289a0d 100644 --- a/vllm/entrypoints/serve/sagemaker/api_router.py +++ b/vllm/entrypoints/serve/sagemaker/api_router.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import logging from collections.abc import Awaitable, Callable from http import HTTPStatus from typing import Any -import model_hosting_container_standards.sagemaker as sagemaker_standards import pydantic from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response @@ -27,11 +27,55 @@ EndpointFn = Callable[[RequestType, Request], Awaitable[Any]] +def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]: + """Snapshot handler levels for loggers that may be affected by third-party + logging configuration side effects (e.g. model_hosting_container_standards + calling configure_root_logger() at import time).""" + loggers = [logging.getLogger(), logging.getLogger("vllm")] + seen: set[int] = set() + snapshot: list[tuple[logging.Handler, int]] = [] + for logger in loggers: + for handler in logger.handlers: + if id(handler) not in seen: + seen.add(id(handler)) + snapshot.append((handler, handler.level)) + return snapshot + + +def _restore_handler_levels( + snapshot: list[tuple[logging.Handler, int]], +) -> None: + """Restore handler levels from a snapshot.""" + for handler, level in snapshot: + handler.setLevel(level) + + def attach_router( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], model_config: ModelConfig | None = None, ): + """Attach the SageMaker hosting endpoints to the API server. + + Handler levels are snapshotted and restored because importing + model_hosting_container_standards may reconfigure root logging. + """ + snapshot = _snapshot_handler_levels() + try: + _attach_router(app, supported_tasks, model_config) + finally: + _restore_handler_levels(snapshot) + + +def _attach_router( + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + """Register the SageMaker hosting routes (/ping, /invocations) on the + app.""" + import model_hosting_container_standards.sagemaker as sagemaker_standards + router = APIRouter() # NOTE: Construct the TypeAdapters only once @@ -99,4 +143,17 @@ async def invocations(raw_request: Request): def sagemaker_standards_bootstrap(app: FastAPI) -> FastAPI: - return sagemaker_standards.bootstrap(app) + """Bootstrap the app with the SageMaker hosting standards. + + Handler levels are restored right after the import because importing + model_hosting_container_standards may reconfigure root logging, and + bootstrap must run with the original levels. + """ + snapshot = _snapshot_handler_levels() + try: + import model_hosting_container_standards.sagemaker as sagemaker_standards + + app = sagemaker_standards.bootstrap(app) + finally: + _restore_handler_levels(snapshot) + return app From f37c550bf6353d7d2a7289cbf256943e8c282fad Mon Sep 17 00:00:00 2001 From: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Date: Wed, 16 Sep 2026 01:29:43 -0700 Subject: [PATCH 311/774] [Feature][Frontend] Expose effective attention block size for DCP (#56538) Co-authored-by: Bugen Zhao Co-authored-by: OpenAI Codex Signed-off-by: Julien Darve Signed-off-by: Bugen Zhao --- docs/serving/context_parallel_deployment.md | 9 +++++ rust/Cargo.toml | 2 +- rust/proto/README.md | 6 ++- rust/proto/control.proto | 2 + rust/src/engine-core-client/src/client.rs | 9 +++++ .../src/engine-core-client/src/mock_engine.rs | 1 + .../src/protocol/handshake.rs | 3 ++ .../engine-core-client/src/tests/client.rs | 6 +++ .../src/tests/python_compat.py | 2 + rust/src/server/src/grpc/control.rs | 1 + rust/src/server/src/grpc/tests.rs | 13 ++++++- tests/config/test_config_utils.py | 3 ++ tests/v1/core/test_kv_cache_utils.py | 38 +++++++++++++++++++ tests/v1/engine/test_engine_core_client.py | 31 +++++++++++++-- vllm/config/cache.py | 3 ++ vllm/v1/engine/__init__.py | 2 + vllm/v1/engine/core.py | 21 +++++++++- vllm/v1/engine/core_client.py | 18 +++++++-- 18 files changed, 159 insertions(+), 11 deletions(-) diff --git a/docs/serving/context_parallel_deployment.md b/docs/serving/context_parallel_deployment.md index 3df2b06635f5..cb9a1ece71a4 100644 --- a/docs/serving/context_parallel_deployment.md +++ b/docs/serving/context_parallel_deployment.md @@ -20,6 +20,15 @@ Both approaches are under active development. Due to the auto-regressive nature of decoding, every decoding step needs to compute a small amount of query tokens w.r.t. a large number of key/value tokens stored in the paged KV cache. The core of decode context parallel is how to shard the KV cache across GPUs. +`engine_client.vllm_config.cache_config.effective_attention_block_size` reports the +initialized full-attention block size in tokens, including DCP: a physical block of +16 tokens with DCP=4 represents 64 tokens. The same value is available through +gRPC `Control.GetServerInfo.effective_attention_block_size`. Physical block-size +fields retain their existing meaning. The new value is unavailable (`None` in +Python, absent in gRPC) if there is no common full-attention size, the scheduler +does not expose it, or an older engine omits it. gRPC also omits the value when +engine replicas disagree. Partial cache events carry their own actual size. + For a model with `H` kv-heads, a request with `T` tokens in the context needs to store `H * T` key/value tensors in the KV cache. 1. If one GPU can hold them all, and the performance is good enough, then no parallelization is needed. diff --git a/rust/Cargo.toml b/rust/Cargo.toml index bd9d99de623c..dde065e79044 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -147,7 +147,7 @@ vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } -vllm-proto = { version = "0.2.0", path = "proto" } +vllm-proto = { path = "proto" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } diff --git a/rust/proto/README.md b/rust/proto/README.md index 3181cdc6a07b..b8a20b0c349b 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -2,6 +2,9 @@ This directory is the canonical source for vLLM's gRPC schema. +See [context parallel deployment](../../docs/serving/context_parallel_deployment.md) +for effective attention block-size metadata in Python and `Control.GetServerInfo`. + Schema updates are no longer published to the Buf Schema Registry. Rust consumers should use `vllm-proto` from crates.io; consumers in other languages can generate bindings from the `.proto` files in this directory. Buf still builds and lints @@ -34,8 +37,7 @@ On pull requests and releases, `cargo-semver-checks` compares the crate with its latest published version. Include any required version bump in the protocol change PR. This check becomes available after the first manual publication. -1. Update the crate version and the `vllm-proto` workspace dependency together, - and update `rust/Cargo.lock`. +1. Update the crate version in `rust/proto/Cargo.toml` and update `rust/Cargo.lock`. 2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run` and the frontend gRPC tests. Record the tested vLLM releases or revisions in the release notes; matching crate versions alone do not establish runtime compatibility. diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 3050035fdcf6..d2816f68d347 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -50,6 +50,8 @@ message ServerInfo { uint64 max_batched_tokens = 9; uint32 max_loras = 10; RlCapabilities rl_capabilities = 11; + // Full-attention block size in tokens after initialization; absent if unavailable. + optional uint64 effective_attention_block_size = 14; } message RlCapabilities { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 63689ea3db83..4cd7420e561e 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -510,6 +510,15 @@ impl EngineCoreClient { self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum() } + /// Return the effective attention block size if all engines report the same value. + pub fn effective_attention_block_size(&self) -> Option { + let size = self.ready_response().effective_attention_block_size?; + self.engines + .iter() + .all(|engine| engine.ready_response.effective_attention_block_size == Some(size)) + .then_some(size) + } + /// Return the minimum engine-reported `max_model_len` across all engines. /// /// This is the auto-fitted value after KV cache profiling and may differ diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 77d72b525fbd..958fcf8aa1b3 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -72,6 +72,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { weight_transfer_backend: None, enable_sleep_mode: false, supports_draft_weight_updates: false, + effective_attention_block_size: None, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index c14c19a50373..52d897be52dc 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -104,6 +104,9 @@ pub struct EngineCoreReadyResponse { /// Whether the engine has a speculative draft model that can be updated. #[serde(default)] pub supports_draft_weight_updates: bool, + /// Full-attention block size in tokens after initialization, or unavailable. + #[serde(default)] + pub effective_attention_block_size: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 27fbbe9abd93..3db573d2aadc 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2846,6 +2846,11 @@ fn python_msgpack_fixtures_match_rust_encoding() { let ready_response: EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + let mut legacy_ready = serde_json::to_value(&ready_response).unwrap(); + legacy_ready.as_object_mut().unwrap().remove("effective_attention_block_size"); + let legacy_ready: EngineCoreReadyResponse = + rmp_serde::from_slice(&rmp_serde::to_vec_named(&legacy_ready).unwrap()).unwrap(); + assert!(legacy_ready.effective_attention_block_size.is_none()); assert!(ready_response.supports_lora); assert_eq!(ready_response.max_loras, 8); assert_eq!( @@ -2854,6 +2859,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { ); assert!(ready_response.enable_sleep_mode); assert!(ready_response.supports_draft_weight_updates); + assert_eq!(ready_response.effective_attention_block_size, Some(64)); let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); assert!(kv_events_config.enable_kv_cache_events); assert_eq!(kv_events_config.publisher, "zmq"); diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 20cec81808cc..0cebca972d4d 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -423,6 +423,7 @@ class EngineCoreReadyResponse: weight_transfer_backend: str | None = None enable_sleep_mode: bool = False supports_draft_weight_updates: bool = False + effective_attention_block_size: int | None = None ready_response = EngineCoreReadyResponse( @@ -446,6 +447,7 @@ class EngineCoreReadyResponse: weight_transfer_backend="nccl", enable_sleep_mode=True, supports_draft_weight_updates=True, + effective_attention_block_size=64, kv_events_config=KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index ddb74bf40f1d..de9c5ae65fb9 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -195,6 +195,7 @@ impl pb::control_server::Control for ControlServiceImpl { max_batched_tokens: ready.max_num_batched_tokens, max_loras: ready.max_loras, rl_capabilities: Some(self.rl_capabilities()), + effective_attention_block_size: self.client().effective_attention_block_size(), })) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f5f74d557e69..a21e013323d4 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1911,8 +1911,16 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_reports_server_and_model_info() { + let mut ready = default_ready_response(); + ready.effective_attention_block_size = Some(64); let (generate_service, control_service, engine_health, _engine_task) = - setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await; + setup_grpc_service_with_engine_script( + b"engine-grpc-info".to_vec(), + ready, + Arc::new(FakeTextBackend), + |_, _| boxed_test_future(async {}), + ) + .await; let (channel, server_task) = start_grpc_test_server( generate_service, control_service, @@ -1928,6 +1936,7 @@ async fn control_reports_server_and_model_info() { .expect("get server info") .into_inner(); assert_eq!(server.engine_version, "test-vllm-version"); + assert_eq!(server.effective_attention_block_size, Some(64)); assert_eq!(server.api_version, "vllm"); assert_eq!(server.instance_id, "test-instance"); assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32); @@ -2195,6 +2204,7 @@ async fn control_aggregates_multi_engine_capacity() { ready_0.weight_transfer_backend = Some("nccl".to_string()); ready_0.enable_sleep_mode = true; ready_0.supports_draft_weight_updates = true; + ready_0.effective_attention_block_size = Some(64); let mut ready_1 = default_ready_response(); ready_1.max_model_len = 4_096; @@ -2248,6 +2258,7 @@ async fn control_aggregates_multi_engine_capacity() { .into_inner(); assert_eq!(server.max_model_len, 4_096); assert_eq!(server.total_kv_blocks, 30); + assert!(server.effective_attention_block_size.is_none()); let rl = server.rl_capabilities.expect("RL capabilities"); assert!(!rl.weight_transfer_enabled); assert!(rl.weight_transfer_backend.is_empty()); diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 94060d1724a8..590e97c1eb54 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -215,6 +215,9 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): base_hash = CacheConfig().compute_hash() assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash + config = CacheConfig() + config.effective_attention_block_size = 64 + assert config.compute_hash() == base_hash def test_scheduler_config_hash_includes_max_num_seqs(): diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c6354cdc0fa0..86e66e1941f1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -279,6 +279,44 @@ def make_request( ) +@pytest.mark.parametrize("dcp", [1, 4]) +def test_effective_attention_block_size_matches_events(dcp): + from vllm.distributed.kv_events import BlockStored + from vllm.v1.engine.core import EngineCore + + config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], new_kv_cache_spec()), + ], + ) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([config]), + max_model_len=256, + scheduler_block_size=16 * dcp, + hash_block_size=16 * dcp, + dcp_world_size=dcp, + enable_kv_cache_events=True, + ) + core = EngineCore.__new__(EngineCore) + core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16)) + core.scheduler = SimpleNamespace(kv_cache_manager=manager) + core._initialize_effective_attention_block_size() + block_size = core.vllm_config.cache_config.effective_attention_block_size + assert block_size == 16 * dcp + + request = make_request( + "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256 + ) + assert manager.allocate_slots(request, 64) is not None + assert [ + event.block_size + for event in manager.take_events() + if isinstance(event, BlockStored) + ] == [block_size] + + def new_kv_cache_spec( block_size=16, num_kv_heads=2, diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index f97ed76c5c48..3f407457b22a 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -158,7 +158,10 @@ def setsockopt(self, *_args, **_kwargs): local_engines_only=False, enable_elastic_ep=False, ) - vllm_config = SimpleNamespace(parallel_config=parallel_config) + vllm_config = SimpleNamespace( + parallel_config=parallel_config, + cache_config=SimpleNamespace(effective_attention_block_size=None), + ) client = core_client_mod.MPClient( asyncio_mode=False, @@ -368,10 +371,15 @@ def test_dplb_finished_requests_release_inflight(): assert req.request_id not in client.reqs_in_flight -def test_apply_ready_response_syncs_block_size(): +@pytest.mark.parametrize( + ("effective_size", "other_size"), + [(None, None), (4224, 4224), (4224, 1056), (4224, None)], +) +def test_apply_ready_response_syncs_block_size(effective_size, other_size): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), @@ -399,14 +407,31 @@ def test_apply_ready_response_syncs_block_size(): max_loras=0, ) ) - client._apply_ready_response(payload) + fields = msgspec.msgpack.decode(payload) + if effective_size is None: + del fields["effective_attention_block_size"] + else: + fields["effective_attention_block_size"] = effective_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) assert client.vllm_config.cache_config.block_size == 1056 + cache_config = client.vllm_config.cache_config + assert cache_config.effective_attention_block_size == effective_size + + fields["effective_attention_block_size"] = other_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) + assert cache_config.effective_attention_block_size == ( + effective_size if effective_size == other_size else None + ) + + client._apply_ready_response(b"") + assert cache_config.effective_attention_block_size is None def test_apply_ready_response_syncs_mamba_block_size(): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 9bd898731eb1..ce3199c52a25 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -211,6 +211,8 @@ class CacheConfig: """The number of blocks to allocate for CPU memory.""" # Set after KV cache initialization. + effective_attention_block_size: int | None = field(default=None, init=False) + """Full-attention block size in tokens, including DCP, or None if unavailable.""" kv_cache_size_tokens: int | None = field(default=None, init=False) """Per-DP-engine KV cache capacity in tokens (group-aware). Uses group-aware capacity since num_gpu_blocks * block_size can be wrong @@ -277,6 +279,7 @@ def compute_hash(self) -> str: # Post-init/derived counters "num_gpu_blocks", "num_cpu_blocks", + "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", # WIP feature toggle not impacting compiled graph shape diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 5ae9ee0cac83..7d68b6e7978b 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -102,6 +102,8 @@ class EngineCoreReadyResponse: weight_transfer_backend: str | None = None enable_sleep_mode: bool = False supports_draft_weight_updates: bool = False + # Full-attention block size in tokens after initialization, or unavailable. + effective_attention_block_size: int | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 62b6ef880e29..b3d94de3e7db 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -85,7 +85,7 @@ EngineCoreSentinel, fault_tolerant_wrapper, ) -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus @@ -168,6 +168,7 @@ def __init__( block_size=scheduler_block_size, hash_block_size=hash_block_size, ) + self._initialize_effective_attention_block_size() self.use_spec_decode = vllm_config.speculative_config is not None self.check_for_draft_tokens = ( self.use_spec_decode or vllm_config.model_config.is_diffusion @@ -382,6 +383,21 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: ) return scheduler_kv_cache_config + def _initialize_effective_attention_block_size(self) -> None: + cache_config = self.vllm_config.cache_config + cache_config.effective_attention_block_size = None + cache_manager = getattr(self.scheduler, "kv_cache_manager", None) + if cache_manager is None: + return + block_sizes = { + manager.block_size + for manager in cache_manager.coordinator.single_type_managers + if is_full_attention_spec(manager.kv_cache_spec) + } + cache_config.effective_attention_block_size = ( + block_sizes.pop() if len(block_sizes) == 1 else None + ) + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: supported_tasks = self.model_executor.supported_tasks self._log_pooler_config(supported_tasks) @@ -1649,6 +1665,9 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, block_size=self.vllm_config.cache_config.block_size, mamba_block_size=self.vllm_config.cache_config.mamba_block_size, + effective_attention_block_size=( + self.vllm_config.cache_config.effective_attention_block_size + ), dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index c354ba699c75..b63e8016c470 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -566,6 +566,7 @@ def __init__( ): self.vllm_config = vllm_config self._renderer: BaseRenderer | None = renderer + self._effective_attention_block_sizes: set[int | None] = set() # ZMQ setup. sync_ctx = zmq.Context(io_threads=2) @@ -803,10 +804,21 @@ def monitor_engine_cores(): def _apply_ready_response(self, payload: bytes) -> None: """Decode an EngineCoreReadyResponse and sync any post-initialization config changes (e.g. auto-fitted max_model_len) back to the frontend.""" - if not payload: - return vllm_config = self.vllm_config - response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + response = ( + msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + if payload + else None + ) + self._effective_attention_block_sizes.add( + response.effective_attention_block_size if response is not None else None + ) + sizes = self._effective_attention_block_sizes + vllm_config.cache_config.effective_attention_block_size = ( + next(iter(sizes)) if len(sizes) == 1 else None + ) + if response is None: + return vllm_config.model_config.max_model_len = min( vllm_config.model_config.max_model_len, response.max_model_len ) From 35f6047c1a2df9bb26ebbefc2b4e64717010e0fe Mon Sep 17 00:00:00 2001 From: Jakub Byczkowski Date: Wed, 16 Sep 2026 10:54:02 +0200 Subject: [PATCH 312/774] [XPU][Bugfix] Fix Qwen2-Audio ValueError on audio clips longer than 30s (#56912) Signed-off-by: Jakub Byczkowski Signed-off-by: Kunshang Ji Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Kunshang Ji --- vllm/model_executor/models/qwen2_audio.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 74341575e744..639db2884706 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -253,6 +253,7 @@ def _preprocess_hf_mm_data( hf_processor_mm_kwargs = dict( **hf_processor_mm_kwargs, sampling_rate=feature_extractor.sampling_rate, + truncation=True, ) return mm_data, hf_processor_mm_kwargs From 0384e72693985d0eddae70bd330b9cb283ff4313 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 16 Sep 2026 04:54:10 -0400 Subject: [PATCH 313/774] [UX] Add thinking support to `vllm chat` (#57045) --- vllm/entrypoints/cli/openai.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/vllm/entrypoints/cli/openai.py b/vllm/entrypoints/cli/openai.py index fa01ff534b4c..b4e286838524 100644 --- a/vllm/entrypoints/cli/openai.py +++ b/vllm/entrypoints/cli/openai.py @@ -46,21 +46,35 @@ def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]: def _print_chat_stream(stream, stats: bool = False) -> str: - output = "" - start = time.perf_counter() + output: str = "" + in_reasoning: bool = False + start: float = time.perf_counter() ttft: float | None = None - completion_tokens = 0 + completion_tokens: int = 0 for chunk in stream: if chunk.usage is not None: completion_tokens = chunk.usage.completion_tokens if not chunk.choices: continue delta = chunk.choices[0].delta + reasoning = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if ttft is None and (reasoning or delta.content): + ttft = time.perf_counter() - start + if reasoning: + if not in_reasoning: + print("", flush=True) + in_reasoning = True + print(reasoning, end="", flush=True) if delta.content: - if ttft is None: - ttft = time.perf_counter() - start + if in_reasoning: + print("\n", flush=True) + in_reasoning = False output += delta.content print(delta.content, end="", flush=True) + if in_reasoning: + print("\n", end="", flush=True) print() if stats: _print_metrics(start, ttft, completion_tokens) From 03f67b3ad1e61d12c21d42e20e69cb5f99bdd82e Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 16 Sep 2026 17:08:42 +0800 Subject: [PATCH 314/774] [Perf][DSV4.1] Pad shared experts for native MegaMoE fusion (#56568) Signed-off-by: Canlin Co-authored-by: Codex --- tests/models/test_deepseek_v4_mega_moe.py | 98 ++++++++++++++++++----- vllm/models/deepseek_v4/nvidia/model.py | 53 ++++++++++-- 2 files changed, 128 insertions(+), 23 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 0d3c2b658369..9a5723d233a1 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -367,17 +367,21 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat [ pytest.param(128, 512, 128, True, False, id="aligned-block128"), pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"), - pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"), + pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash"), + pytest.param(128, 2304, 128, True, False, id="padded-block128"), + pytest.param(5120, 2304, 32, True, True, id="deepseek-v41-flash-mxfp8"), + pytest.param(128, 2304, 128, False, False, id="padded-packed-scale-fallback"), ], ) def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8 ): - """V4.1-Flash's routed padding must preserve the serial shared-expert fallback.""" + """Shared fusion preserves checkpoint channels and zeros padded channels.""" class FakeDeepGemm: transformed_dims: list[tuple[int, int]] = [] scale_inputs: list[tuple[int, ...]] = [] + scale_values: list[torch.Tensor] = [] @staticmethod def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs): @@ -402,6 +406,7 @@ def fp8_fp4_mega_moe( @classmethod def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs): cls.scale_inputs.append(tuple(sf.shape)) + cls.scale_values.append(sf) return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32) @classmethod @@ -429,7 +434,10 @@ def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): def fp8_parameter(*shape): return torch.nn.Parameter( - torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False + torch.empty(*shape, dtype=torch.uint8) + .random_(1, 119) + .view(torch.float8_e4m3fn), + requires_grad=False, ) def scale_parameter(*shape, dtype=torch.int32): @@ -460,21 +468,34 @@ def scale_parameter(*shape, dtype=torch.int32): linear.weight_block_size = (1, 32) linear.weight_scale = torch.nn.Parameter( linear.weight_scale_inv.view(torch.uint8) - .repeat_interleave(128, dim=0) - .repeat_interleave(4, dim=1), + .repeat_interleave(block_size, dim=0) + .repeat_interleave(block_size // 32, dim=1), requires_grad=False, ) del linear.weight_scale_inv + originals = [] + for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): + scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv + if not fused: + # Already packed scales cannot be padded as checkpoint UE8M0 bytes. + linear.weight_scale_inv = scale_parameter( + linear.weight.shape[0], linear.weight.shape[1] // 128 + ) + continue + scale.data.view(torch.uint8).random_(124, 130) + block_m, block_k = linear.weight_block_size + scale_1x32 = ( + experts._ue8m0_uint8_to_float(scale.view(torch.uint8)) + .repeat_interleave(block_m, dim=0) + .repeat_interleave(block_k // 32, dim=1) + ) + originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32)) monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() original_down_ptr = shared_experts.down_proj.weight.data_ptr() experts.finalize_weights(shared_experts) - assert FakeDeepGemm.scale_inputs[-2:] == [ - (1, 2 * intermediate_size, hidden_size // 32), - (1, hidden_size, intermediate_size // 32), - ] assert experts.has_fused_shared_experts is fused if not fused: assert experts.intermediate_size == 2560 @@ -490,15 +511,56 @@ def scale_parameter(*shape, dtype=torch.int32): return assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] - assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr - assert ( - experts._transformed_shared_l1_weights[0].data_ptr() - == shared_experts.gate_up_proj.weight.data_ptr() - ) - assert ( - experts._transformed_shared_l2_weights[0].data_ptr() - == shared_experts.down_proj.weight.data_ptr() - ) + padded_size = experts.intermediate_size + assert FakeDeepGemm.scale_inputs[-2:] == [ + (1, 2 * padded_size, hidden_size // 32), + (1, hidden_size, padded_size // 32), + ] + for actual, original, fill in ( + ( + experts._transformed_shared_l1_weights[0].view(torch.uint8), + originals[0][0], + 0, + ), + (FakeDeepGemm.scale_values[-2][0], originals[0][1], 1), + ): + actual = actual.unflatten(0, (2, padded_size)) + assert torch.equal( + actual[:, :intermediate_size], original.unflatten(0, (2, intermediate_size)) + ) + assert torch.all(actual[:, intermediate_size:] == fill) + for actual, original, fill in ( + ( + experts._transformed_shared_l2_weights[0].view(torch.uint8), + originals[1][0], + 0, + ), + (FakeDeepGemm.scale_values[-1][0], originals[1][1], 1), + ): + assert torch.equal(actual[:, : original.shape[1]], original) + assert torch.all(actual[:, original.shape[1] :] == fill) + if padded_size != intermediate_size: + # Generic linear post-load hooks still receive checkpoint-shaped weights. + for linear, (weight, _) in zip( + (shared_experts.gate_up_proj, shared_experts.down_proj), originals + ): + assert torch.equal(linear.weight.view(torch.uint8), weight) + assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr + assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr + else: + assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr + assert ( + experts._transformed_shared_l1_weights[0].data_ptr() + == shared_experts.gate_up_proj.weight.data_ptr() + ) + assert ( + experts._transformed_shared_l2_weights[0].data_ptr() + == shared_experts.down_proj.weight.data_ptr() + ) + + transformed_l1 = experts._transformed_shared_l1_weights + experts.finalize_weights(shared_experts) + assert experts._transformed_shared_l1_weights is transformed_l1 @pytest.mark.parametrize("fused", [False, True]) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b1379cc9eb45..4aba14cd03b1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -204,6 +204,7 @@ def __init__( self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.unpadded_intermediate_size = intermediate_size self.num_shared_experts = num_shared_experts self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -386,6 +387,17 @@ def _finalize_shared_expert_weights( # the generic linear post-load hook replaces the raw checkpoint scales # with its 128x128 DeepGEMM layout. checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8) + unpadded_size = self.unpadded_intermediate_size * self.num_shared_experts + padding = self.intermediate_size * self.num_shared_experts - unpadded_size + pad_weights = ( + padding > 0 + and gate_up_weight.dtype == torch.float8_e4m3fn + and down_weight.dtype == torch.float8_e4m3fn + and gate_up_weight.shape == (2 * unpadded_size, self.hidden_size) + and down_weight.shape == (self.hidden_size, unpadded_size) + and gate_up_scale.dtype in checkpoint_scale_dtypes + and down_scale.dtype in checkpoint_scale_dtypes + ) if ( gate_up_scale.dtype in checkpoint_scale_dtypes and down_scale.dtype in checkpoint_scale_dtypes @@ -396,6 +408,7 @@ def _finalize_shared_expert_weights( gate_up_scale, gate_up_weight.shape[0], gate_up_weight.shape[1], + padding=(padding, 0) if pad_weights else (0, 0), ) down_scale = self._prepare_shared_expert_scale( deep_gemm, @@ -403,12 +416,27 @@ def _finalize_shared_expert_weights( down_scale, down_weight.shape[0], down_weight.shape[1], + padding=(0, padding) if pad_weights else (0, 0), ) if gate_up_scale is None or down_scale is None: self.num_shared_experts = 0 return + if pad_weights: + # Pad gate/up separately so the SwiGLU split stays at the midpoint. + gate_up_weight = ( + torch.nn.functional.pad( + gate_up_weight.view(torch.uint8).unflatten(0, (2, unpadded_size)), + (0, 0, 0, padding), + ) + .flatten(0, 1) + .view(gate_up_weight.dtype) + ) + down_weight = torch.nn.functional.pad( + down_weight.view(torch.uint8), (0, padding) + ).view(down_weight.dtype) + shared_intermediate_size = self.intermediate_size * self.num_shared_experts expected_gate_up_shape = ( 2 * shared_intermediate_size, @@ -449,11 +477,11 @@ def _finalize_shared_expert_weights( # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. # The generic linear post-load hook may still repack the serial scales, # but this shared MLP is never called after native fusion is enabled. - gate_up.weight.data = transformed_l1[0] - self._transformed_shared_l1_weights = ( - gate_up.weight.data, - transformed_l1[1], - ) + # Padded weights need separate storage: the generic linear post-load + # hooks still expect the original checkpoint weight/scale shapes. + if not pad_weights: + gate_up.weight.data = transformed_l1[0] + self._transformed_shared_l1_weights = transformed_l1 self._transformed_shared_l2_weights = transformed_l2 def _prepare_shared_expert_scale( @@ -463,6 +491,8 @@ def _prepare_shared_expert_scale( scale: torch.Tensor, mn: int, k: int, + *, + padding: tuple[int, int] = (0, 0), ) -> torch.Tensor | None: block_size = getattr(linear, "weight_block_size", None) if block_size is None or len(block_size) != 2: @@ -497,6 +527,19 @@ def _prepare_shared_expert_scale( .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32] .contiguous() ) + pad_m, pad_k = padding + if pad_m: + scale_1x32 = torch.nn.functional.pad( + scale_1x32.unflatten(0, (2, mn // 2)), + (0, 0, 0, pad_m), + value=1.0, + ).flatten(0, 1) + mn += 2 * pad_m + if pad_k: + scale_1x32 = torch.nn.functional.pad( + scale_1x32, (0, pad_k // 32), value=1.0 + ) + k += pad_k # The grouped API is used with a singleton dimension to request the # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the # 2D layout required for a shared expert. From 82731931ae1a770ef0439a2afa88a54cb850cdfa Mon Sep 17 00:00:00 2001 From: Ganesh R Date: Wed, 16 Sep 2026 14:45:12 +0530 Subject: [PATCH 315/774] [Bugfix][CPU] Fall back to CpuPlatform when zentorch fails to import (#54923) Signed-off-by: R Co-authored-by: Cursor --- tests/test_zen_cpu_platform_detection.py | 52 +++++++++++++++++++++++- vllm/platforms/__init__.py | 9 ++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/test_zen_cpu_platform_detection.py b/tests/test_zen_cpu_platform_detection.py index 9f386f70cd98..0c99c6568a2b 100644 --- a/tests/test_zen_cpu_platform_detection.py +++ b/tests/test_zen_cpu_platform_detection.py @@ -65,7 +65,7 @@ def test_cpu_target_selects_cpu_platform_from_non_cpu_wheel( rocm_plugin.assert_not_called() -def test_platform_detection_logs_zentorch_import_failure(caplog): +def test_broken_zentorch_falls_back_to_cpu_platform(caplog): original_import = builtins.__import__ def import_with_broken_zentorch(name, *args, **kwargs): @@ -88,6 +88,54 @@ def import_with_broken_zentorch(name, *args, **kwargs): ): platform = resolve_current_platform_cls_qualname() + assert platform == "vllm.platforms.cpu.CpuPlatform" + assert "zentorch failed to import" in caplog.text + assert "OSError: incompatible shared library" in caplog.text + + +def test_zentorch_failure_other_than_oserror_is_not_recovered(caplog): + original_import = builtins.__import__ + + def import_with_broken_zentorch(name, *args, **kwargs): + if name == "zentorch": + raise RuntimeError("unknown zentorch failure") + return original_import(name, *args, **kwargs) + + with ( + patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"), + patch.dict( + "vllm.platforms.builtin_platform_plugins", + {"cpu": cpu_platform_plugin}, + clear=True, + ), + patch("vllm.platforms.load_plugins_by_group", return_value={}), + patch("vllm.platforms.vllm_version_matches_substr", return_value=True), + patch("vllm.platforms._is_amd_zen_cpu", return_value=True), + patch.object(builtins, "__import__", side_effect=import_with_broken_zentorch), + caplog.at_level(logging.DEBUG, logger="vllm.platforms"), + ): + platform = resolve_current_platform_cls_qualname() + + assert platform == "vllm.platforms.interface.UnspecifiedPlatform" + assert "RuntimeError: unknown zentorch failure" in caplog.text + + +def test_platform_plugin_failure_is_logged(caplog): + def failing_plugin(): + raise RuntimeError("plugin exploded") + + with ( + patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"), + patch.dict( + "vllm.platforms.builtin_platform_plugins", + {"cpu": failing_plugin}, + clear=True, + ), + patch("vllm.platforms.load_plugins_by_group", return_value={}), + caplog.at_level(logging.DEBUG, logger="vllm.platforms"), + ): + platform = resolve_current_platform_cls_qualname() + assert platform == "vllm.platforms.interface.UnspecifiedPlatform" assert "Platform plugin cpu failed during detection" in caplog.text - assert "OSError: incompatible shared library" in caplog.text + assert "RuntimeError: plugin exploded" in caplog.text diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index b34b0b9a7dd9..7cf0843669c3 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -217,6 +217,15 @@ def cpu_platform_plugin() -> str | None: "AMD Zen CPU detected but zentorch not installed, " "falling back to CpuPlatform." ) + except OSError: + # An ABI-mismatched build fails here with an undefined-symbol + # error; other failures are not known to be safe to recover from. + logger.warning( + "AMD Zen CPU detected but zentorch failed to import, falling " + "back to CpuPlatform. This usually means the zentorch build " + "does not match the installed torch version.", + exc_info=True, + ) return "vllm.platforms.cpu.CpuPlatform" From 5f203baedb29bd97f96243f2b83e73fb1bfec7f0 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:16:32 +0100 Subject: [PATCH 316/774] [Docs] Split slash-combined docstring parameters (#57141) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- vllm/_custom_ops.py | 10 ++++++++-- .../attention/dsa/sparse_mqa_logits.py | 20 +++++++++++++------ vllm/utils/deep_gemm.py | 4 +++- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 3311215016a7..99577766e7a8 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1482,8 +1482,14 @@ def cutlass_w4a8_moe_mm( Cumulative token offsets problem_sizes: Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher. - a/b/c/group_scale_strides: - Strides describing the memory layout of the input tensors. + a_strides: + Strides describing the memory layout of a_tensors. + b_strides: + Strides describing the memory layout of b_tensors. + c_strides: + Strides describing the memory layout of out_tensors. + group_scale_strides: + Strides describing the memory layout of b_group_scales. maybe_schedule: Optional override to choose a specific kernel or epilogue schedule. diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index 839c581583a7..138833ad149e 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -190,9 +190,11 @@ def candidate_blocks_to_sparse_indices( candidate_blocks: [rows, K] int32 request-local candidate block ids (-1 padded), in units of ``candidate_block_size`` positions. K must be a power of two (the in-kernel sort); rows may be strided. - row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + row_ks: [rows] int32 per-row K range start; bounds are in the same (packed-workspace) coordinates the sparse kernel iterates over. Pass zeros for the paged path, whose blocks are context-relative. + row_ke: [rows] int32 per-row K range end, in the same coordinates as + ``row_ks``. candidate_block_size: Positions per candidate block. sparse_block_kv: Positions per sparse block (8 or 16). out: Optional ``(sparse_indices, end)`` buffers to write into, @@ -322,12 +324,16 @@ def sparse_mqa_logits_prefill_chunk( k_scale: [total_kv] int32 packed UE8M0 K scales. weights: [rows, H] bf16 per-head weights; the sparse kernels take bf16 and do not fold the Q scale in. - cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the - packed workspace. + cu_seqlen_ks: [rows] int32 per-token K start bounds in the packed + workspace. + cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed + workspace. candidate_blocks: [rows, K] int32 request-local candidate block ids. topk_indices: [rows, topk_tokens] output buffer. - sparse_indices/end/col_indices: Caller-owned scratch, see - `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + sparse_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices`. + end: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`. + col_indices: Caller-owned scratch, see `sparse_topk_remap`. kernel_metadata: DeepGEMM schedule from a previous call with the same candidates and bounds (i.e. another indexer layer in the same step). When given, the candidate expansion is skipped and @@ -413,7 +419,9 @@ def sparse_mqa_logits_paged_decode( candidate_blocks: [rows, K] int32 candidate block ids. topk_indices: [rows, topk_tokens] output buffer. row_ks: [rows] int32 zeros (paged blocks are context-relative). - sparse_indices/end/col_indices: Caller-owned scratch. + sparse_indices: Caller-owned scratch. + end: Caller-owned scratch. + col_indices: Caller-owned scratch. kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. Returns: diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 91e315afb9e7..6fef0072243f 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -726,8 +726,10 @@ def get_sparse_mqa_logits_metadata( need a host-side sync to pick one. Args: - cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + cu_seqlen_ks: Per-row K range start bounds in the packed KV workspace, shape [num_q_tokens], dtype int32. + cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace, + shape [num_q_tokens], dtype int32. num_kv_tokens: Total KV tokens in the packed workspace. sparse_kv_block_indices: Per-row candidate block ids, shape [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's From ceb87de065b9cbdaffd34a0e6aed86c99e772b95 Mon Sep 17 00:00:00 2001 From: Adababy Date: Wed, 16 Sep 2026 17:17:07 +0800 Subject: [PATCH 317/774] [Misc] Remove no-op self-assignments across vLLM (#55988) Signed-off-by: shaolila Co-authored-by: shaolila --- .../layers/fused_moe/routed_experts.py | 15 ++++++++------- vllm/model_executor/models/deepencoder.py | 4 +--- vllm/model_executor/models/mimo_v2_omni.py | 1 - vllm/model_executor/models/step3p5.py | 1 - 4 files changed, 9 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 2dafb6c493d6..257e19b034a4 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -615,14 +615,15 @@ def weight_loader( # compressed-tensors checkpoints with packed weights are stored flipped # TODO (mgoin): check self.quant_method.quant_config.quant_format # against known CompressionFormat enum values that have this quality - if quant_method_name in ( - "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsW4A16FlydslMoEMethod", + if ( + quant_method_name + in ( + "CompressedTensorsWNA16MoEMethod", + "CompressedTensorsW4A16FlydslMoEMethod", + ) + and is_transposed ): - if is_transposed: - loaded_weight = loaded_weight.t().contiguous() - else: - loaded_weight = loaded_weight + loaded_weight = loaded_weight.t().contiguous() if shard_id not in ("w1", "w2", "w3"): raise ValueError(f"shard_id must be ['w1','w2','w3'] but got {shard_id}.") diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py index 189e8b55e510..d75e5bb77d2f 100644 --- a/vllm/model_executor/models/deepencoder.py +++ b/vllm/model_executor/models/deepencoder.py @@ -780,9 +780,7 @@ def forward( self, pixel_values: torch.Tensor, patch_embeds: torch.Tensor | None = None ) -> torch.Tensor: batch_size = pixel_values.shape[0] - if patch_embeds is not None: - patch_embeds = patch_embeds - else: + if patch_embeds is None: patch_embeds = self.patch_embedding(pixel_values) patch_embeds = patch_embeds.flatten(2).transpose(1, 2) diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 93bd6a44b556..9281a1689ae1 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -964,7 +964,6 @@ def _apply_hf_processor_main( assert isinstance(va_item, (list, tuple)) and len(va_item) == 2 vid, audio_src = va_item va_item = VideoAudioInput(video=vid, audio=audio_src) - vid = vid # Convert video frames to (TCHW, timestamps) if needed if ( isinstance(vid, tuple) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 3a928aa59daa..03343a89a144 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -173,7 +173,6 @@ def __init__( rope_scaling = None if sliding_window is not None and enable_sliding_window: - sliding_window = sliding_window if swa_num_attention_heads is not None: num_heads = swa_num_attention_heads self.total_num_heads = swa_num_attention_heads From d6a1677d5504244c566eb900ca605cb5511f4ab4 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 16 Sep 2026 03:07:50 -0700 Subject: [PATCH 318/774] [Model][DSv4.1] FlashMLA mega attention and the NVFP4 compressed KV cache (#56935) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Yongye Zhu Signed-off-by: 云挚 Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Yunxiao Ning <73625538+foraxe@users.noreply.github.com> Co-authored-by: AI Assistant --- .../kernels/benchmark_dsv41_mega_attn.py | 272 ++++++++++ ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 206 +++++--- csrc/libtorch_stable/ops.h | 3 +- csrc/libtorch_stable/torch_bindings.cpp | 9 +- tests/kernels/test_compressor_kv_cache.py | 139 ++++- tests/kernels/test_dsv41_mega_attn_layouts.py | 71 +++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 183 +++++++ vllm/models/deepseek_v41/amd/rocm.py | 3 +- vllm/models/deepseek_v41/attention.py | 204 +++++-- .../deepseek_v41/common/ops/cache_utils.py | 98 +++- .../common/ops/fused_compress_quant_cache.py | 84 ++- .../deepseek_v41/common/ops/fused_layout.py | 86 +++ vllm/models/deepseek_v41/nvidia/dspark.py | 14 + .../nvidia/flash_mla_mega_attn.py | 497 ++++++++++++++++++ .../deepseek_v41/nvidia/flashinfer_sparse.py | 6 +- vllm/models/deepseek_v41/nvidia/flashmla.py | 3 +- vllm/models/deepseek_v41/nvidia/model.py | 30 +- vllm/models/deepseek_v41/sparse_mla.py | 51 ++ vllm/v1/attention/backends/mla/sparse_swa.py | 5 +- vllm/v1/attention/backends/registry.py | 3 + 20 files changed, 1811 insertions(+), 156 deletions(-) create mode 100644 benchmarks/kernels/benchmark_dsv41_mega_attn.py create mode 100644 tests/kernels/test_dsv41_mega_attn_layouts.py create mode 100644 vllm/models/deepseek_v41/common/ops/fused_layout.py create mode 100644 vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py diff --git a/benchmarks/kernels/benchmark_dsv41_mega_attn.py b/benchmarks/kernels/benchmark_dsv41_mega_attn.py new file mode 100644 index 000000000000..1b3b5d9c3848 --- /dev/null +++ b/benchmarks/kernels/benchmark_dsv41_mega_attn.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashMLA mega attention vs the split-KV decode pipeline (DeepSeek V4.1). + +Times the full set of per-layer ops each path runs in a decode step, captured +in a CUDA graph -- eager timing is dominated by launch overhead and flatters +whichever path launches fewer kernels. + + mega: fused_qnorm_rope_kv_rope_quant_insert (Q pad + KV insert) + + mega kernel + split-KV: fused_qnorm_rope_kv_rope_quant_insert (Q RoPE + pad + KV insert) + + flash_mla_with_kvcache + fused_inv_rope_fp8_quant + +Both arms fuse their Q preparation into the KV insert, so that op is charged +to both; the mega arm's Q side is only a zero-pad (nothing at all when the +shard is already at the kernel head count) because its kernel does the Q RoPE, +the inverse RoPE and the FP8 cast. + +``--local-heads`` is the variable that decides the outcome: the mega kernel +absorbs work proportional to the live head count while paying for the padded +one, so it wins at 64 live heads and is a wash at 16 (TP4 on a 64-head model). + +Both arms use the same KV records so the comparison is kernel-for-kernel: +a V4.1 fp8 (528 B) sliding-window cache and a compressed cache that is either +V4.1 fp8 or V4.1 NVFP4 (288 B) -- FlashMLA's SM100 sparse decode reads the +NVFP4 compressed record in both the fused and unfused kernels. + +Run: .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py +""" + +import argparse + +import torch + +import vllm.v1.attention.ops.flashmla as fm +from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant, +) +from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache +from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, +) +from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused +from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + alloc_mega_attn_output, + is_flashmla_mega_attn_supported, +) +from vllm.utils.math_utils import round_up + +HEAD_DIM, ROPE_DIM = 512, 64 +V41_BYTES, V41_FP4_BYTES = 528, 288 +SWA_BLOCK, COMPRESSED_BLOCK = 32, 128 + + +def make_cos_sin_cache(max_pos, device): + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, ROPE_DIM, 2, device=device).float() / ROPE_DIM) + ) + freqs = torch.outer(torch.arange(max_pos, device=device).float(), inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], -1) + + +def empty_paged_cache(num_rows, block_size, bytes_per_token, device): + """A zeroed paged cache view [num_blocks, block_size, bytes_per_token].""" + num_blocks = (num_rows + block_size - 1) // block_size + 1 + page = round_up(block_size * bytes_per_token, 512) + backing = torch.zeros(num_blocks, page, dtype=torch.uint8, device=device) + return backing, backing.as_strided( + (num_blocks, block_size, bytes_per_token), (page, bytes_per_token, 1) + ) + + +def fill_cache(view, bytes_per_token, cos_sin, device): + n = view.shape[0] * view.shape[1] + rows = torch.randn(n, HEAD_DIM, device=device, dtype=torch.bfloat16) + slots = torch.arange(n, dtype=torch.int64, device=device) + if bytes_per_token == V41_FP4_BYTES: + # A page-aligned cache has more slots than the cos_sin table has rows; + # wrap so the insert never indexes past it. + positions = slots % cos_sin.shape[0] + rope_quant_insert(rows, positions, cos_sin, view, slots, 1) + else: + quantize_and_insert_k_cache( + rows, + view.reshape(view.shape[0], -1), + slots, + block_size=view.shape[1], + bytes_per_token=bytes_per_token, + ) + + +def time_graph(fn, iters=100, warmup=5): + """Capture ``fn`` in a CUDA graph and time its replay.""" + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + fn() + torch.cuda.current_stream().wait_stream(s) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + for _ in range(warmup): + g.replay() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + g.replay() + end.record() + torch.accelerator.synchronize() + return start.elapsed_time(end) * 1000 / iters + + +def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads): + padded_heads, scale = 64, HEAD_DIM**-0.5 + n_wv_group = padded_heads // 8 + cos_sin = make_cos_sin_cache(65536, device) + positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64) + pos32 = positions.to(torch.int32) + + # Per-step inputs: the wq_b output (local heads) and this step's KV row. + q_local = torch.randn( + s_q, local_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 + ) + q_fused = permute_q_to_fused(q_local) + kv = torch.randn(s_q, HEAD_DIM, device=device, dtype=torch.bfloat16) + + n_swa = s_q * 128 + 128 + swa_backing, swa = empty_paged_cache(n_swa, SWA_BLOCK, V41_BYTES, device) + fill_cache(swa, V41_BYTES, cos_sin, device) + _, extra = empty_paged_cache(65536, COMPRESSED_BLOCK, extra_bytes, device) + fill_cache(extra, extra_bytes, cos_sin, device) + + slot_mapping = torch.arange(s_q, dtype=torch.int64, device=device) + swa_idx = torch.randint(0, n_swa, (s_q, 128), device=device, dtype=torch.int32) + swa_len = torch.full((s_q,), 128, device=device, dtype=torch.int32) + ex_idx = torch.randint( + 0, 65536, (s_q, topk_extra), device=device, dtype=torch.int32 + ) + ex_len = torch.full((s_q,), topk_extra, device=device, dtype=torch.int32) + sink = torch.zeros(padded_heads, device=device) + sched = fm.FlashMLASchedMeta() + out = alloc_mega_attn_output(s_q, n_wv_group, device) + swa4 = swa.unsqueeze(-2) + extra4 = extra.unsqueeze(-2) + swa_2d = swa_backing + + def mega(): + # One launch zero-pads the fused-layout Q and inserts the SWA KV. + q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_fused, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin, + 0 if padded_heads == local_heads else padded_heads, + 1e-6, + SWA_BLOCK, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope: the mega kernel rotates Q itself + True, # is_q_interleaved + ) + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode( + q_fused if padded_heads == local_heads else q_pad, + swa4, + swa_idx, + scale, + HEAD_DIM, + sink, + swa_len, + extra4, + ex_idx, + ex_len, + False, + 0.0, + pos32, + False, + 64, + cos_sin, + n_wv_group, + 32, + True, + True, + True, + out.data, + out.scale, + ) + + def split_kv(): + q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_local, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin, + padded_heads, + 1e-6, + SWA_BLOCK, + False, + True, + ) + o = fm.flash_mla_with_kvcache( + q=q_pad.unsqueeze(1), + k_cache=swa4, + block_table=None, + head_dim_v=HEAD_DIM, + tile_scheduler_metadata=sched, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_idx.view(s_q, 1, -1), + topk_length=swa_len, + softmax_scale=scale, + attn_sink=sink, + extra_k_cache=extra4, + extra_indices_in_kvcache=ex_idx.view(s_q, 1, -1), + extra_topk_length=ex_len, + )[0] + fused_inv_rope_fp8_quant( + o.squeeze(1)[:, :local_heads], + positions, + cos_sin, + n_groups=local_heads // 8, + heads_per_group=8, + quant_group_size=32, + tma_aligned_scales=True, + ) + + return time_graph(mega), time_graph(split_kv) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--decode-s-q", + nargs="+", + type=int, + default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + ) + parser.add_argument("--topk-extra", type=int, default=512) + parser.add_argument("--local-heads", type=int, default=16) + parser.add_argument( + "--extra-bytes", + type=int, + choices=[V41_BYTES, V41_FP4_BYTES], + default=V41_BYTES, + help="compressed-cache record: 528 (V4.1 fp8) or 288 (V4.1 NVFP4)", + ) + args = parser.parse_args() + ok, reason = is_flashmla_mega_attn_supported() + if not ok: + raise SystemExit(reason) + device = torch.device("cuda") + torch.manual_seed(0) + print( + f"decode, cudagraph (topk_swa=128, topk_extra={args.topk_extra}, " + f"local_heads={args.local_heads}, extra={args.extra_bytes}B)" + ) + print(f"{'s_q':>6} {'mega_us':>9} {'split_kv_us':>12} {'speedup':>8}") + for s_q in args.decode_s_q: + m, s = bench_decode( + s_q, device, args.topk_extra, args.extra_bytes, args.local_heads + ) + print(f"{s_q:>6d} {m:9.1f} {s:12.1f} {s / m:7.2f}x") + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 344b82fbda2c..5a5caacce7a5 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -3,7 +3,10 @@ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project * * Horizontally-fused DeepseekV4-MLA kernel: - * - Q side: optional per-head RMSNorm + GPT-J RoPE on last ROPE_DIM + * - Q side: optional per-head RMSNorm (apply_q_norm) and GPT-J RoPE on the + * last ROPE_DIM (apply_q_rope), plus the zero-pad to + * q_head_padded, in the head-major layout or (is_q_interleaved) + * FlashMLA's mega-attention chunk-interleaved one * - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged * cache insert * @@ -174,6 +177,30 @@ __device__ __forceinline__ float warpSum(float val) { return val; } +// Offset of the 16 elements a lane owns inside a Q tensor of `num_heads` +// heads. The default layout is head-major: a head's 512 dims are contiguous. +// `Q_INTERLEAVED` is FlashMLA's mega-attention layout, 32 chunks of +// [num_heads, 16] per token, which changes the address only -- lane L still +// owns dims [16L, 16L+16) of its head either way, so `dim_base` stays the +// true dim offset and the norm and RoPE below are equally valid in both. +// What it does change is where the padding heads land: at the tail of every +// chunk rather than after the live heads. +template +__device__ __forceinline__ int64_t qElemOffset(int const tokenIdx, + int const slotIdx, + int const laneId, + int const dim_base, + int const num_heads) { + if constexpr (Q_INTERLEAVED) { + return static_cast(tokenIdx) * num_heads * kHeadDim + + static_cast(laneId) * num_heads * kElemsPerLane + + static_cast(slotIdx) * kElemsPerLane; + } else { + return (static_cast(tokenIdx) * num_heads + slotIdx) * kHeadDim + + dim_base; + } +} + // ──────────────────────────────────────────────────────────────────────────── // Per-slot inner pipeline // ──────────────────────────────────────────────────────────────────────────── @@ -192,10 +219,16 @@ __device__ __forceinline__ float warpSum(float val) { // + paged-cache // insert) // -// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a -// runtime argument rather than doubling every (head count, q-norm) -// instantiation of a kernel whose hot path is the Q branch. -template +// `kv_mxfp8` and `apply_q_rope` are grid-uniform, so they stay runtime +// arguments rather than multiplying every (head count, q-norm, q-layout) +// instantiation of a kernel whose hot path is the Q branch. `Q_INTERLEAVED` +// cannot: it picks the address arithmetic, which has to fold. +// +// The three Q knobs are independent. The mega-attention layer uses +// Q_INTERLEAVED with both norm and RoPE off, because its attention kernel +// applies those itself -- not because the layout forbids them. +template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, int const dim_base, int const laneId, int const num_heads_q, @@ -203,7 +236,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride, bool const kv_mxfp8) { + int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -213,10 +246,8 @@ __device__ __forceinline__ void processDeepseekV4Slot( // zero literal is correct. Matches the live-Q branch's vectorized store. if (isPadQ) { scalar_t_in* dst = - q_out + - (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * - kHeadDim + - dim_base; + q_out + qElemOffset(tokenIdx, slotIdx, laneId, dim_base, + kNumHeadsQPadded); uint4 const zero4 = {0u, 0u, 0u, 0u}; *reinterpret_cast(dst) = zero4; *reinterpret_cast(dst + 8) = zero4; @@ -264,7 +295,9 @@ __device__ __forceinline__ void processDeepseekV4Slot( // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────── // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage). - bool const is_rope_lane = dim_base >= kNopeDim; + // KV is always rotated; Q only when the caller has not already had it + // rotated elsewhere (the mega-attention kernel does its own). + bool const is_rope_lane = dim_base >= kNopeDim && (isKV || apply_q_rope); if (is_rope_lane) { int64_t const pos = position_ids[tokenIdx]; constexpr int kHalfRope = kRopeDim / 2; @@ -314,10 +347,8 @@ __device__ __forceinline__ void processDeepseekV4Slot( make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1])); } scalar_t_in* dst = - q_out + - (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * - kHeadDim + - dim_base; + q_out + qElemOffset(tokenIdx, slotIdx, laneId, dim_base, + kNumHeadsQPadded); *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { @@ -432,14 +463,17 @@ __device__ __forceinline__ void processDeepseekV4Slot( // `kNumHeadsQPadded` is a template parameter (compile-time constant) so the // divisions in the grid math and the KV-sentinel comparison fold to fast // constant operations. The launch wrapper dispatches the runtime value to -// the matching instantiation. +// the matching instantiation. `kNumHeadsQPadded == 0` leaves one slot per +// token and is the KV-insert-only instantiation, for a Q that is already in +// the shape its attention kernel wants. // // With DP padding, q/kv/position_ids can have more rows than slot_mapping. // The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream // attention uses them). The KV branch only inserts the first // `num_tokens_insert` tokens (= slot_mapping length) into the paged cache. // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512] scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512] @@ -454,7 +488,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block int const kv_block_stride, // bytes per paged-cache block - bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record + bool const kv_mxfp8, // V4.1 all-dims MXFP8 record + bool const apply_q_rope) { // rotate Q too, not just KV #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -497,20 +532,18 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( if (isKV) { src_ptr = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; } else { - int64_t const q_row_offset = - (static_cast(tokenIdx) * num_heads_q + slotIdx) * - kHeadDim + - dim_base; - src_ptr = q_in + q_row_offset; + src_ptr = q_in + qElemOffset(tokenIdx, slotIdx, laneId, + dim_base, num_heads_q); } v0 = *reinterpret_cast(src_ptr); v1 = *reinterpret_cast(src_ptr + 8); } - processDeepseekV4Slot( + processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride, kv_mxfp8); + kv_block_stride, kv_mxfp8, apply_q_rope); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -530,7 +563,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( // Q branch (optional RMSNorm + RoPE, in place) head_slot == num_heads_q // KV branch (RoPE + UE8M0 quant + insert) // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out, scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache, @@ -539,7 +573,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride, bool const kv_mxfp8) { + int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -569,11 +603,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( if (s == kNumHeadsQPadded) { src = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; } else { - src = q_in + - (static_cast(tokenIdx) * num_heads_q + - static_cast(s)) * - kHeadDim + - dim_base; + src = q_in + qElemOffset(tokenIdx, s, laneId, dim_base, + num_heads_q); } va = *reinterpret_cast(src); vb = *reinterpret_cast(src + 8); @@ -594,10 +625,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( load_slot(next_slot, v0_next, v1_next); } - processDeepseekV4Slot( + processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8, + apply_q_rope); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -617,14 +650,15 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( // ──────────────────────────────────────────────────────────────────────────── // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── -template +template static void launchFusedDeepseekV4Templated( scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, - cudaStream_t stream) { + bool const apply_q_rope, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; int64_t const total_warps = @@ -663,37 +697,40 @@ static void launchFusedDeepseekV4Templated( // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single // insert, where it is worth ~2x. if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { - cudaLaunchKernelEx(&config, - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride, - kv_mxfp8); + cudaLaunchKernelEx( + &config, + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>, + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, + kv_block_stride, kv_mxfp8, apply_q_rope); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride, kv_mxfp8); + kv_block_stride, kv_mxfp8, apply_q_rope); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) // clang-format off fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM> + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED> <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride, kv_mxfp8); + cache_block_size, kv_block_stride, kv_mxfp8, apply_q_rope); #endif } // Runtime dispatch into one of the precompiled `kNumHeadsQPadded` -// instantiations. Supported padded head counts: 8, 16, 32, 64, 128. +// instantiations. Supported padded head counts: 8, 16, 32, 64, 128, plus 0 +// for a KV-only launch that writes no Q at all. `apply_q_norm` and +// `is_q_interleaved` are template parameters, `apply_q_rope` is not (see +// processDeepseekV4Slot), so each head count carries four instantiations. template void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, @@ -702,22 +739,32 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { -#define DISPATCH(N) \ - case N: \ - if (apply_q_norm) { \ - launchFusedDeepseekV4Templated( \ - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ - } else { \ - launchFusedDeepseekV4Templated( \ - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ - } \ + bool const apply_q_norm, bool const kv_mxfp8, bool const apply_q_rope, + bool const is_q_interleaved, cudaStream_t stream) { +#define LAUNCH(N, NORM, INTERLEAVED) \ + launchFusedDeepseekV4Templated( \ + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, \ + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, \ + kv_block_stride, kv_mxfp8, apply_q_rope, stream) +#define DISPATCH(N) \ + case N: \ + if (is_q_interleaved) { \ + if (apply_q_norm) { \ + LAUNCH(N, true, true); \ + } else { \ + LAUNCH(N, false, true); \ + } \ + } else if (apply_q_norm) { \ + LAUNCH(N, true, false); \ + } else { \ + LAUNCH(N, false, false); \ + } \ return; + if (num_heads_q_padded == 0) { + LAUNCH(0, false, false); + return; + } switch (num_heads_q_padded) { DISPATCH(8) DISPATCH(16) @@ -729,9 +776,10 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: " "unsupported num_heads_q_padded=", num_heads_q_padded, - " (compiled instantiations: 8, 16, 32, 64, 128)."); + " (compiled instantiations: 0, 8, 16, 32, 64, 128)."); } #undef DISPATCH +#undef LAUNCH } // ──────────────────────────────────────────────────────────────────────────── @@ -1059,10 +1107,15 @@ void fused_deepseek_v4_kv_rope_insert( // Zero query heads schedules only the existing KV branch, preserving its // RoPE rounding and quantization without allocating any query tensors. if (packed) { - vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + // With zero query heads no Q slot is scheduled, so neither Q knob is + // reachable: Q_INTERLEAVED picks addressing nothing uses, and apply_q_rope + // gates a branch only a Q slot takes (KV is rotated unconditionally). + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), - static_cast(k_cache.stride(0)), kv_mxfp8, stream); + static_cast(k_cache.stride(0)), kv_mxfp8, /*apply_q_rope=*/true, + stream); } else if (fp8) { vllm::deepseek_v4_fused_ops::launchFullCacheKernel( nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, @@ -1085,8 +1138,9 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& slot_mapping, // [N] int64 torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 - int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { + int64_t q_head_padded, // padded Q head count for output, 0 for no Q + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1106,8 +1160,11 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(), "q_in and kv dtype must match"); - STD_TORCH_CHECK(q_head_padded >= q_in.size(1), - "q_head_padded must be >= q_in.size(1) (num_heads_q)"); + // q_head_padded == 0 asks for the KV insert alone: Q is left untouched and + // an empty tensor comes back, for a caller whose Q is already in the shape + // its attention kernel reads. + STD_TORCH_CHECK(q_head_padded == 0 || q_head_padded >= q_in.size(1), + "q_head_padded must be 0 or >= q_in.size(1) (num_heads_q)"); STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte, "k_cache must be uint8"); STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, @@ -1138,8 +1195,12 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( // Allocate the padded q output. The kernel writes every element (live // region gets optional RMSNorm+RoPE; pad region gets zeros), so `empty` is // safe. - auto q_out = torch::stable::new_empty( - q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); + auto q_out = + q_head_padded == 0 + ? torch::stable::new_empty(q_in, {int64_t{0}}, q_in.scalar_type()) + : torch::stable::new_empty( + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, + q_in.scalar_type()); VLLM_STABLE_DISPATCH_HALF_TYPES( q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { @@ -1155,7 +1216,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, kv_mxfp8, stream); + apply_q_norm, kv_mxfp8, apply_q_rope, is_q_interleaved, + stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index bc08d90a4836..6b17815a07c7 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -271,7 +271,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 3918ab8498cd..a76852df89e5 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + // q_head_padded is the padded Q head count of the returned tensor, or 0 to + // do the KV insert alone and return an empty tensor. The Q knobs are + // independent: apply_q_norm and apply_q_rope each drop that step for Q + // alone (KV is always rotated), and is_q_interleaved reads and writes Q in + // FlashMLA's mega-attention chunk-interleaved layout, which moves the + // padding heads to the tail of every head-dim chunk. ops.def( "fused_deepseek_v4_kv_rope_insert(" "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " @@ -436,7 +442,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " + "bool is_q_interleaved=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1de94be92fc..73c2027a7866 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -331,6 +331,19 @@ def run(): torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _rotate_rope_tail( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> torch.Tensor: + """GPT-J RoPE of one latent row's last 64 dims, in fp32.""" + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + return rotated + + def _mxfp8_record_reference( latent_row: torch.Tensor, cos_sin_row: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -339,12 +352,7 @@ def _mxfp8_record_reference( Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the RoPE tail first, then scale every 32-dim tile, RoPE tiles included. """ - row = latent_row.float() - c, s = cos_sin_row.chunk(2) - rotated = row.clone() - even, odd = row[448::2], row[449::2] - rotated[448::2] = even * c - odd * s - rotated[449::2] = odd * c + even * s + rotated = _rotate_rope_tail(latent_row, cos_sin_row) quantized, scales = _ue8m0_reference(rotated, 32, 448.0) return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) @@ -461,6 +469,125 @@ def test_v41_mxfp8_cache_round_trip(): assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() +# e2m1 magnitudes indexed by the low 3 bits of a code; bit 3 is the sign. +_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Unpack one 256-byte e2m1 row and apply its 32 e4m3 tile scales.""" + mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device) + codes = torch.empty(512, dtype=torch.uint8, device=packed.device) + codes[0::2] = packed & 0xF # even element in the low nibble + codes[1::2] = packed >> 4 + vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0) + return ( + vals.view(32, 16) * scales.view(torch.float8_e4m3fn).float().view(32, 1) + ).flatten() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_nvfp4_record(compress_ratio: int): + """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales. + + Scales are byte-exact against the reference (``amax / 6`` clamped to the + e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose + coarsest step is 2 between 4 and 6, so half a step is ``1.0 * scale``. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(13) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 288 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 288), (cache_stride, 288, 1)) + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 # covers the negative-slot skip + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + untouched = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + cs = cos_sin[pos // compress_ratio * compress_ratio] + rotated = _rotate_rope_tail(latent[t], cs) + amax = rotated.view(32, 16).abs().amax(-1) + scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn) + + got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32] + torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0) + decoded = _decode_nvfp4_row( + cache_backing[page, row * 256 : (row + 1) * 256], got_scales + ) + step = scale.float().repeat_interleave(16) + assert ((decoded - rotated).abs() <= step).all() + untouched[page, row * 256 : (row + 1) * 256] = cache_backing[ + page, row * 256 : (row + 1) * 256 + ] + untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales + # Rows the kernel must not have touched, page padding included. + torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_nvfp4_gather_matches_insert(): + """The NVFP4 gather dequantizes exactly what the insert kernel wrote.""" + from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(17) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 288 / 512) * 512 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(128, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1)) + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + rope_quant_insert(latent, positions, cos_sin, cache, slots, 1) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache, + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + for t in (0, 1, block_size, num_tokens - 1): + page, row = divmod(t, block_size) + expected = _decode_nvfp4_row( + backing[page, row * 256 : (row + 1) * 256], + backing[page, block_size * 256 + row * 32 :][:32], + ) + torch.testing.assert_close(out[0, t].float(), expected, rtol=0, atol=0) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_dsv41_mega_attn_layouts.py b/tests/kernels/test_dsv41_mega_attn_layouts.py new file mode 100644 index 000000000000..6f9ff56bd236 --- /dev/null +++ b/tests/kernels/test_dsv41_mega_attn_layouts.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Layout plumbing for DeepSeek V4.1 mega attention. + +The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and +``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs +agree with a reference that permutes the activation instead. +""" + +import torch + +from vllm.models.deepseek_v41.common.ops.fused_layout import ( + WV_GROUP_SIZE, + o_fused_chunk_permutation, + o_fused_permutation, + permute_q_to_fused, + permute_wo_a_, + permute_wq_b_, + q_fused_permutation, +) + +HEAD_DIM = 512 + + +def test_permuted_wq_b_gemm_matches_permuted_activation(): + """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly.""" + torch.manual_seed(0) + num_heads, q_lora_rank, num_tokens = 16, 64, 5 + weight = torch.randn(num_heads * HEAD_DIM, q_lora_rank) + # One scale row per weight row, as an MXFP8 wq_b shard carries. + scale = torch.randint( + 0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8 + ) + qr = torch.randn(num_tokens, q_lora_rank) + orig_weight, orig_scale = weight.clone(), scale.clone() + + standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) + expected = permute_q_to_fused(standard) + + permute_wq_b_(weight, scale, num_heads) + fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) + torch.testing.assert_close(fused, expected) + # The scale has to follow its row, or dequant reads another row's scale. + perm = q_fused_permutation(num_heads, HEAD_DIM) + torch.testing.assert_close(weight, orig_weight[perm]) + torch.testing.assert_close(scale, orig_scale[perm]) + + +def test_permuted_wo_a_consumes_fused_output(): + """Permuting wo_a's input columns lets it read the kernel's O layout.""" + torch.manual_seed(1) + out_features, num_tokens = 32, 4 + in_features = WV_GROUP_SIZE * HEAD_DIM + weight = torch.randn(out_features, in_features) + scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8) + standard_o = torch.randn(num_tokens, in_features) + orig_weight, orig_scale = weight.clone(), scale.clone() + + expected = standard_o @ weight.T + # The kernel emits the same values in the fused chunk order. + perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM) + fused_o = standard_o[:, perm] + + permute_wo_a_(weight, scale, WV_GROUP_SIZE) + torch.testing.assert_close(weight, orig_weight[:, perm]) + torch.testing.assert_close( + scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)] + ) + # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the + # summation order changes and the result differs in the last few ulps. + torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 68eabfb94e84..a0dcfc5bb7da 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -522,6 +522,189 @@ def _dequant(cache_2d): ) +# ── Test 2d: mega-attention Q layout (is_q_interleaved) ────────────────────── + +Q_CHUNK = 16 +NUM_Q_CHUNKS = HEAD_DIM // Q_CHUNK # 32 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 2048]) +@pytest.mark.parametrize("n_heads,q_head_padded", [(16, 64), (64, 128), (64, 0)]) +def test_q_interleaved_pads_without_rope( + num_tokens: int, n_heads: int, q_head_padded: int +): + """``is_q_interleaved`` with no norm or RoPE only zero-pads Q. + + The mega-attention kernel norms and rotates Q itself, so every live + element must survive bit-exact while the padding heads -- which interleave + with the live ones inside each of the 32 head-dim chunks -- come back + zero. ``q_head_padded=0`` drops the Q pass entirely and returns nothing, + for a shard already at the kernel's head count. + """ + torch.manual_seed(5) + device = "cuda" + dtype = torch.bfloat16 + block_size = 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + k_cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + + q_out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope + True, # is_q_interleaved + ) + + # The KV half is unaffected by the Q mode: it must still be inserted. + assert k_cache.any() + + if q_head_padded == 0: + assert q_out.numel() == 0 + return + assert q_out.shape == (num_tokens, q_head_padded, HEAD_DIM) + got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK) + want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK) + torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0) + assert (got[:, :, n_heads:] == 0).all() + + +@pytest.mark.parametrize("n_heads", [8, 16, 32]) +def test_q_interleaved_survives_graph_replay(n_heads: int): + """A captured fused insert refills its padded Q buffer on every replay. + + The op allocates ``q_out`` itself, so the tensor a capture hands back has + to keep receiving the new padding when the graph is replayed over changed + input -- that is what lets the layer prepare Q inside the captured region + rather than in ``forward_mqa``. + """ + torch.manual_seed(6) + device = "cuda" + dtype = torch.bfloat16 + block_size, num_tokens, q_head_padded = 64, 7, 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + k_cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + + def prepare(): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope + True, # is_q_interleaved + ) + + # Capture takes its warmup on a side stream. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + prepare() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q_out = prepare() + + for factor in (0.25, -0.75): + q.copy_(torch.randn_like(q) * factor) + graph.replay() + torch.accelerator.synchronize() + got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK) + want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK) + torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0) + assert (got[:, :, n_heads:] == 0).all() + + +@pytest.mark.parametrize("n_heads", [16, 64]) +@pytest.mark.parametrize("apply_q_norm", [False, True]) +@pytest.mark.parametrize("apply_q_rope", [False, True]) +def test_q_interleaved_is_orthogonal_to_norm_and_rope( + n_heads: int, apply_q_norm: bool, apply_q_rope: bool +): + """The Q layout composes with whatever norm and RoPE it is given. + + A lane owns the same 512-dim slice of its head in either layout -- only + the address differs -- so the head-major result permuted into the + interleaved layout must equal running the interleaved path on a permuted + input, bit for bit, for every norm/RoPE combination. + """ + from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused + + torch.manual_seed(7) + device = "cuda" + dtype = torch.bfloat16 + block_size, num_tokens, q_head_padded = 64, 5, 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + def run(q_in, is_q_interleaved): + cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_in, + kv, + cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + apply_q_norm, + True, # kv_mxfp8 + apply_q_rope, + is_q_interleaved, + ) + return out, cache + + head_major, cache_major = run(q, False) + interleaved, cache_inter = run(permute_q_to_fused(q), True) + + torch.testing.assert_close( + interleaved, permute_q_to_fused(head_major), rtol=0, atol=0 + ) + # The Q layout must not leak into the KV half. + torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 802a59e340fc..b3d80f8d2591 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -636,7 +636,8 @@ def _split_qkv_and_norm( transpose_scale=False, ) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. z = rocm_inv_rope_einsum( self.rotary_emb, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 134328dbe908..752b335eafd8 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -44,6 +44,7 @@ VllmConfig, get_current_vllm_config, ) +from vllm.config.cache import CacheDType from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -127,34 +128,40 @@ def _use_v41_mxfp8_kv_record() -> bool: def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, - kv_cache_dtype: str, + kv_cache_dtype: CacheDType, cache_config: CacheConfig | None, -) -> tuple[str, torch.dtype]: + packed_kv_cache_dtype: CacheDType = "fp8_ds_mla", +) -> tuple[CacheDType, torch.dtype]: """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. Both layouts are paged; they differ in the per-token block format. The - ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the - canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the - page-size specs pick the 576B per-token slot). Plain-row backends store each - token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. + packed formats are ``uint8``-backed: ``fp8_ds_mla`` is UE8M0 block-scaled + fp8 throughout, ``nvfp4_ds_mla`` keeps that sliding-window record and + stores the compressed cache as NVFP4. An unspecific ``--kv-cache-dtype`` + (``auto`` / ``fp8``) resolves to ``packed_kv_cache_dtype``, the record this + layer's kernel prefers, and the canonical string is written back onto + ``cache_config`` so the page-size specs pick the right per-token slot. + Plain-row backends store each token's KV row in its element dtype: bf16 or + per-tensor FP8 E4M3. """ if use_fp8_ds_mla_layout: - # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. - if kv_cache_dtype == "auto": - kv_cache_dtype = "fp8" - if not kv_cache_dtype.startswith("fp8"): + if kv_cache_dtype in ("auto", "fp8"): + kv_cache_dtype = packed_kv_cache_dtype + elif not kv_cache_dtype.endswith("_ds_mla"): raise ValueError( - "DeepseekV4 fp8_ds_mla layout only supports fp8 " - f"kv-cache, got {kv_cache_dtype}. Please set " - "`--kv-cache-dtype fp8` or select a backend that supports " - "bfloat16 KV cache." + "DeepseekV4 packed KV layouts only support fp8 kv-cache, got " + f"{kv_cache_dtype}. Please set `--kv-cache-dtype fp8` or " + "select a backend that supports bfloat16 KV cache." ) - if kv_cache_dtype != "fp8_ds_mla": - if cache_config is not None: - cache_config.cache_dtype = "fp8_ds_mla" - kv_cache_dtype = "fp8_ds_mla" - logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + if cache_config is not None and cache_config.cache_dtype != kv_cache_dtype: + cache_config.cache_dtype = kv_cache_dtype + logger.info_once("Using DeepSeek's %s KV cache format.", kv_cache_dtype) return kv_cache_dtype, torch.uint8 + if kv_cache_dtype.endswith("_ds_mla"): + raise ValueError( + f"{kv_cache_dtype} is a packed FlashMLA DeepSeek V4.1 KV cache " + "format; the selected backend stores plain KV rows." + ) # Plain bf16 / per-tensor fp8 KV row (FlashInfer). if kv_cache_dtype.startswith("fp8"): @@ -189,6 +196,28 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # path to pre-reserve that workspace. PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + # ---- attention-interface contract, declared by the platform subclass ---- + + @property + def accepts_unnormed_unroped_query(self) -> bool: + """Whether ``forward_mqa``'s ``q`` is the raw ``wq_b`` output. + + True when the attention kernel applies the Q norm and RoPE itself, and + reads Q in its own chunk-interleaved layout, so the layer only + zero-pads Q to ``padded_heads`` and inserts KV before calling it. + """ + return False + + @property + def packed_kv_cache_dtype(self) -> CacheDType: + """The packed KV record this layer's kernel prefers. + + What an unspecific ``--kv-cache-dtype`` (``auto`` / ``fp8``) resolves + to. Mega attention overrides it: its kernel is the one that can read + an NVFP4 compressed cache. + """ + return "fp8_ds_mla" + @classmethod @abstractmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -213,8 +242,19 @@ def forward_mqa( raise NotImplementedError @abstractmethod - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + def _o_proj( + self, + attn_out: "torch.Tensor | QuantizedActivation", + positions: torch.Tensor, + ) -> torch.Tensor: + """Project whatever ``_alloc_attn_out`` produced through wo_a and wo_b. + + Takes the buffer whole, so each layer owns the shape it allocated: the + bf16 layers slice off their padding heads and apply the inverse RoPE, + while a layer whose attention kernel already did the inverse RoPE and + the FP8 cast gets a QuantizedActivation and has only wo_a and wo_b + left. + """ raise NotImplementedError def _uses_fp8_ds_mla_layout(self) -> bool: @@ -459,11 +499,27 @@ def __init__( # Resolve the kv-cache dtype from this backend's block format. The same # resolution drives the SWA cache tensor dtype below. self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( - self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config + self._uses_fp8_ds_mla_layout(), + cache_config.cache_dtype, + cache_config, + self.packed_kv_cache_dtype, ) self.kv_mxfp8 = _use_v41_mxfp8_kv_record() - self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.swa_bytes_per_token = 528 if self.kv_mxfp8 else 584 + # nvfp4_ds_mla keeps the MXFP8 sliding-window record and stores the + # compressed cache as NVFP4 (256 B of e2m1 pairs + 32 e4m3 scales). + self.compressed_bytes_per_token = ( + 288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token + ) + # One alignment for every page in the block: the block stride is their + # sum, and 512 satisfies both TMA strides in play (512 for the V4.1 + # fp8 record, 256 for NVFP4). self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 + if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8: + raise ValueError( + "nvfp4_ds_mla needs the V4.1 KV records, which FlashMLA " + "decodes only on SM100." + ) self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -473,7 +529,7 @@ def __init__( cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, - packed_bytes_per_token=self.kv_bytes_per_token, + packed_bytes_per_token=self.swa_bytes_per_token, packed_page_alignment=self.kv_page_alignment, ) @@ -553,8 +609,12 @@ def __init__( block_size=self.swa_cache_layer.block_size, ) + # Every backend that gathers a chunk's KV through + # combine_topk_swa_indices needs its Triton kernel warmed, mega + # attention included -- it calls it from _forward_prefill_mega. if self.backend_cls.get_name() in ( "FLASHMLA_SPARSE_DSV41", + "FLASHMLA_MEGA_ATTN_DSV41", "ROCM_FLASHMLA_SPARSE_DSV4", ): from vllm.models.deepseek_v41.common.ops.cache_utils import ( @@ -569,14 +629,10 @@ def forward( hidden_states: torch.Tensor, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: - # Pre-allocate attention output with FlashMLA-padded head count. - # The op writes into `o_padded`; we slice to n_local_heads after. - num_tokens = hidden_states.shape[0] - o_padded = torch.empty( - (num_tokens, self.padded_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + # The eager attention region writes into a caller-owned buffer + # (breakable_cudagraph needs in-place outputs); its shape and how it is + # projected afterwards follow the interface contract above. + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states) # Keep the attention input preparation in the captured graph. Only the # sparse indexer and MLA attention run in the eager break below. @@ -593,12 +649,19 @@ def forward( kv_score, indexer_weights, positions, - o_padded, + attn_out, ) - o = o_padded[:, : self.n_local_heads, :] + return self._o_proj(attn_out, positions) - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> "torch.Tensor | QuantizedActivation": + """The buffer ``forward_mqa`` fills, per the interface contract.""" + return torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) @cached_property def _can_fuse_query_quant(self) -> bool: @@ -652,7 +715,7 @@ def _prepare_and_attn_eager( kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + attn_out: "torch.Tensor | QuantizedActivation", ) -> None: """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. @@ -667,7 +730,7 @@ def _prepare_and_attn_eager( kv_score, indexer_weights, positions, - o_padded, + attn_out, ) def _prepare_and_attn( @@ -679,7 +742,7 @@ def _prepare_and_attn( kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + attn_out: "torch.Tensor | QuantizedActivation", ) -> None: """Attention input preparation followed by the sparse indexer and MLA. @@ -751,7 +814,7 @@ def prepare_indexer(): q, kv, positions, - o_padded, + attn_out, ) def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -858,16 +921,28 @@ def _fused_qnorm_rope_kv_insert( dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None ), ) -> torch.Tensor: + """Ready ``q`` for the attention kernel and publish this step's KV. + + One launch does both. With ``accepts_unnormed_unroped_query`` the + attention kernel norms and rotates Q itself and reads it in its own + chunk-interleaved layout, so the Q half of the launch is a zero-pad to + ``padded_heads`` -- and nothing at all once the shard is that wide. + """ if not isinstance(attn_metadata, dict): # Profile run: kernel doesn't fire; produce a padded tensor so # downstream FlashMLA gets the right shape. - if self.n_local_heads < self.padded_heads: - return F.pad( - q, - (0, 0, 0, self.padded_heads - self.n_local_heads), - value=0.0, - ) - return q + if self.n_local_heads >= self.padded_heads: + return q + if self.accepts_unnormed_unroped_query: + # Padding heads sit at the tail of every head-dim chunk in + # that layout, so no head-major pad of `q` reproduces it -- + # and nothing reads it on a profile run. + return q.new_zeros((q.shape[0], self.padded_heads, q.shape[2])) + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) swa_metadata = cast( "DeepseekSparseSWAMetadata | None", @@ -886,23 +961,38 @@ def _fused_qnorm_rope_kv_insert( if cache_dtype == torch.uint8: # fp8_ds_mla UE8M0 paged path. Horizontally fused: # Q side: GPT-J RoPE, zero-filling the padding head slots; the - # kernel allocates and returns the padded q tensor. + # kernel allocates and returns the padded q tensor. An + # interleaved Q skips the RoPE its attention kernel owns + # and keeps only the pad, which q_head_padded=0 drops + # too once the shard is already padded_heads wide. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + pad_to = ( + 0 + if self.accepts_unnormed_unroped_query + and self.n_local_heads == self.padded_heads + else self.padded_heads + ) + q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, swa_kv_cache_2d, swa_metadata.slot_mapping, positions, cos_sin_cache, - self.padded_heads, + pad_to, self.eps, swa_metadata.block_size, - False, + False, # apply_q_norm: qr is normed before wq_b self.kv_mxfp8, + not self.accepts_unnormed_unroped_query, # apply_q_rope + self.accepts_unnormed_unroped_query, # is_q_interleaved ) + return q if pad_to == 0 else q_padded + assert not self.accepts_unnormed_unroped_query, ( + "the chunk-interleaved Q layout only pairs with a packed KV record" + ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV # row in its element dtype (no Q padding). bf16 rewrites q in place; # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the @@ -959,7 +1049,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows # use natural element-size pages. - uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" + uses_fp8_ds_mla_layout = self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, @@ -972,7 +1062,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), # Packed record width; head_size stays semantic (512). state_content_bytes=( - self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + self.compressed_bytes_per_token if uses_fp8_ds_mla_layout else None ), ) @@ -1016,7 +1106,13 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. - uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + # nvfp4_ds_mla is packed too: its compressed record is NVFP4 but the + # sliding-window record stays the V4.1 MXFP8 one, so the indexer page + # takes the same alignment either way. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype in ( + "fp8_ds_mla", + "nvfp4_ds_mla", + ) page_alignment = ( 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 ) diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 51743016bf07..4a9da21c036b 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -43,10 +43,16 @@ # dims each plus a pad byte. # V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 # scales of 32 dims each. FlashMLA's ``ModelType::V41``. +# V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in +# the low nibble), then 32 e4m3 scales of 16 dims each. +# FlashMLA's ``ModelType::V41_FP4``, compressed cache only. V4_BYTES_PER_TOKEN = 584 V41_BYTES_PER_TOKEN = 528 V41_QUANT_BLOCK = 32 V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 +V41_NVFP4_BYTES_PER_TOKEN = 288 +V41_NVFP4_QUANT_BLOCK = 16 +V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK # 32 @triton.jit @@ -496,6 +502,73 @@ def _dequantize_and_gather_k_mxfp8_kernel( tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) +@triton.jit +def _dequantize_and_gather_k_nvfp4_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 32 + quant_block: tl.constexpr, # 16 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, +): + """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + packed_bytes: tl.constexpr = head_dim // 2 + d = tl.arange(0, head_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + physical_block_idx = tl.load( + block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq + ) + page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + + packed = tl.load( + page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes) + ) + # Even element in the low nibble, odd in the high nibble. + codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) + # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. + mag_code = codes & 7 + e = (mag_code >> 1).to(tl.float32) + m = (mag_code & 1).to(tl.float32) + mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) + vals = tl.where(codes >= 8, -mag, mag) + + sf = tl.load( + page + + cache_block_size * packed_bytes + + pos_in_block * scale_dim + + tl.arange(0, scale_dim) + ) + scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) + tiles = tl.reshape(vals, (scale_dim, quant_block)) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -513,6 +586,25 @@ def dequantize_and_gather_k_cache_triton( ) -> None: num_reqs = seq_lens.shape[0] NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN: + _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NVFP4_NUM_SCALES, + quant_block=V41_NVFP4_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + ) + return + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( out, @@ -581,14 +673,16 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. - The record is read off ``k_cache.shape[-1]``; see the module header. + The record is read off ``k_cache.shape[-1]``; see the module header. Only + the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton + path. ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder writes FNUZ on gfx942 and OCP on gfx950). """ - if has_cutedsl(): + if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN: # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 96f15849ddd3..0f7830cf58c0 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -4,6 +4,7 @@ import torch +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import _fp32x2_to_fp4x2 from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -232,13 +233,14 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the - per-token byte width names -- 584 B for V4 (576 value bytes and eight - segregated UE8M0 scale bytes, including one zero padding scale) or 528 B - for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 - scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain - [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the - per-tensor ``fp8_scale``. + layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token + byte width names -- 584 B for V4 (576 value bytes and eight segregated + UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1 + (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of + 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32 + e4m3 scales of 16 dims each), which only the compressed cache uses. + ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows + read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -249,14 +251,14 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] in (584, 528), ( + kernel = { + 584: _rope_quant_insert_kernel, + 528: _rope_quant_insert_mxfp8_kernel, + 288: _rope_quant_insert_nvfp4_kernel, + }.get(kv_cache.shape[-1]) + assert kernel is not None, ( f"unsupported paged KV record width {kv_cache.shape[-1]}" ) - kernel = ( - _rope_quant_insert_mxfp8_kernel - if kv_cache.shape[-1] == 528 - else _rope_quant_insert_kernel - ) kernel[(num_tokens,)]( latent, positions, @@ -402,6 +404,62 @@ def _rope_quant_insert_mxfp8_kernel( tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) +@triton.jit +def _rope_quant_insert_nvfp4_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims. + + The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the + e4m3 range, with no per-tensor scale on top. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + tiles = tl.reshape(rotated, (32, 16)) + amax = tl.max(tl.abs(tiles), 1) + # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest. + scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv) + # Round-to-nearest division: Triton's default div.full misplaces values + # that land exactly on an e2m1 tie. + scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1))) + lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2))) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + tl.store( + page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256), + _fp32x2_to_fp4x2(lo, hi), + ) + scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32 + tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/common/ops/fused_layout.py b/vllm/models/deepseek_v41/common/ops/fused_layout.py new file mode 100644 index 000000000000..071afde6bca4 --- /dev/null +++ b/vllm/models/deepseek_v41/common/ops/fused_layout.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight permutations for FlashMLA's mega-attention kernel. + +The kernel reads Q with 16-element head-dim chunks interleaved across heads and +writes O with 32-element chunks interleaved across the 8 heads of a ``wo_a`` +group. Every permutation here satisfies ``fused = standard[perm]``, and is +applied once to ``wq_b`` rows and ``wo_a`` columns at load time so the +surrounding GEMMs produce and consume the kernel's layouts directly -- no +per-step shuffle. +""" + +import torch + +HEAD_DIM = 512 +Q_CHUNK = 16 +O_CHUNK = 32 +WV_GROUP_SIZE = 8 + + +def q_fused_permutation(num_heads: int, head_dim: int = HEAD_DIM) -> torch.Tensor: + """``fused[(d // 16) * (H * 16) + h * 16 + d % 16] = standard[h * D + d]``.""" + h = torch.arange(num_heads).view(num_heads, 1) + d = torch.arange(head_dim).view(1, head_dim) + fused_index = (d // Q_CHUNK) * (num_heads * Q_CHUNK) + h * Q_CHUNK + d % Q_CHUNK + perm = torch.empty(num_heads * head_dim, dtype=torch.long) + perm[fused_index.reshape(-1)] = torch.arange(num_heads * head_dim) + return perm + + +def o_fused_permutation( + heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM +) -> torch.Tensor: + """``fused[(c * G + h) * 32 + j] = standard[h * D + c * 32 + j]`` per group.""" + h = torch.arange(heads_per_group).view(-1, 1, 1) + c = torch.arange(head_dim // O_CHUNK).view(1, -1, 1) + j = torch.arange(O_CHUNK).view(1, 1, -1) + fused_index = (c * heads_per_group + h) * O_CHUNK + j + perm = torch.empty(heads_per_group * head_dim, dtype=torch.long) + perm[fused_index.reshape(-1)] = torch.arange(heads_per_group * head_dim) + return perm + + +def o_fused_chunk_permutation( + heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM +) -> torch.Tensor: + """Per-32-element-chunk form of :func:`o_fused_permutation`, for scales.""" + return o_fused_permutation(heads_per_group, head_dim)[::O_CHUNK] // O_CHUNK + + +def permute_q_to_fused(q: torch.Tensor) -> torch.Tensor: + """``[N, H, D]`` standard layout -> the same shape in the fused layout.""" + n, h, d = q.shape + perm = q_fused_permutation(h, d).to(q.device) + return q.reshape(n, h * d)[:, perm].view(n, h, d) + + +def _bytes_view(t: torch.Tensor) -> torch.Tensor: + """Byte view of a 1-byte-element tensor, so fp8/ue8m0 can be gathered.""" + return t.view(torch.uint8) if t.element_size() == 1 else t + + +def permute_wq_b_( + weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int +) -> None: + """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place.""" + head_dim = weight.shape[0] // num_local_heads + perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device) + for t in (weight, weight_scale): + b = _bytes_view(t) + b.copy_(b[perm]) + + +def permute_wo_a_( + weight: torch.Tensor, + weight_scale: torch.Tensor, + heads_per_group: int = WV_GROUP_SIZE, +) -> None: + """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place.""" + head_dim = weight.shape[1] // heads_per_group + perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device) + w = _bytes_view(weight) + w.copy_(w[:, perm]) + chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim) + s = _bytes_view(weight_scale) + s.copy_(s[:, chunk_perm.to(weight.device)]) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 214681f3a932..9a07189bc3bb 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -477,8 +477,22 @@ def _finalize_moe(self) -> None: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() + def _finalize_attn(self) -> None: + """Run the attention backend's post-load weight step on the draft layers. + + They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide + attention backend and owe it whatever the target model owes it -- mega + attention permutes wq_b / wo_a here and refuses to run without it. + Idempotent, like the target's. + """ + for layer in self.model.layers: + finalize = getattr(layer.attn, "finalize_loaded_weights", None) + if finalize is not None: + finalize() + def process_weights_after_loading(self) -> None: self._finalize_moe() + self._finalize_attn() def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py new file mode 100644 index 000000000000..233d2f80f903 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100). + +One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and +its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the +layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself, +leaving the fused KV insert only the zero-pad to the kernel's head count -- and +its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead +of bf16, since the output needs no rotation or quantization here. + +``wq_b`` rows and ``wo_a`` columns are permuted once at load so the surrounding +GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill +and decode segments write disjoint token ranges of one output buffer pair, so +a single ``wo_a`` einsum covers the whole step. +""" + +from dataclasses import replace +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.models.deepseek_v41.common.ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, +) +from vllm.models.deepseek_v41.common.ops.fused_layout import ( + WV_GROUP_SIZE, + permute_wo_a_, + permute_wq_b_, +) +from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention +from vllm.models.deepseek_v41.sparse_mla import ( + DeepseekV4FlashMLAMetadata, + FlashMLAMegaAttnBackend, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size +from vllm.utils.math_utils import round_up +from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + +_HEAD_DIM_V = 512 +_ROPE_DIM = 64 +# The kernel always emits per-32 scales, whatever num_per_channels says. +_QUANT_GROUP = 32 +# Fixed kernel conventions: GPT-J (non-neox) RoPE, TMA-aligned col-major +# scales, round-to-nearest, packed ue8m0 -- i.e. what DeepGEMM's fp8_einsum +# consumes. Q norm is off: the V4.1 layer norms qr before wq_b. +_ROPE_ARGS = (False, _ROPE_DIM) +_SF_ARGS = (_QUANT_GROUP, True, True, True) + +# Decode always goes through the mega kernel: there is no batch size at which +# falling back to the split-KV decode pays, so there is no minimum-token +# threshold and no fallback path to maintain. What decides whether this layer +# is worth selecting at all is TP, not the batch size. +# +# Measured on one GB300 with benchmarks/kernels/benchmark_dsv41_mega_attn.py, +# CUDA graphs on, counting every per-step op each path runs (64-head model, +# topk_swa=128, topk_extra=512, NVFP4 compressed cache), as mega / split-KV +# microseconds: +# +# s_q 1 32 128 256 512 +# TP1 (64 live heads) 29/29 31/31 31/40 53/72 96/140 +# TP2 (32 live) 29/29 31/30 35/35 61/63 110/119 +# TP4 (16 live) 29/29 31/29 34/34 59/58 109/109 +# TP8 (8 live) 29/29 31/29 33/33 59/57 107/104 +# +# The kernel absorbs the inverse RoPE and the FP8 cast, whose cost scales with +# *live* heads, while its own cost scales with *padded* heads -- always 64, +# since FlashMLA's FP8 decode takes only h_q in {64, 128}. So it wins at TP1 +# (up to 1.45x, and TP1 is also where the fused insert skips the Q pad +# entirely) and from TP2 down the two cancel to within a few percent, with a +# fixed ~2 us penalty at small s_q where the padded-head cost dominates. +# Measure before assuming a win at high TP -- and measure under CUDA graphs: +# in eager mode launch overhead swamps both paths and overstates the fused one +# several-fold. + + +def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]: + """Whether FlashMLA's fused mega-attention kernels are usable here.""" + is_available, maybe_reason = is_flashmla_sparse_supported() + if not is_available: + return False, maybe_reason + if not current_platform.is_device_capability_family(100): + return False, "FlashMLA mega attention requires sm_10x GPUs." + if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"): + return False, "vllm._flashmla_C was built without the mega attention ops." + if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"): + return False, "vllm._flashmla_C predates the mega attention weight permutes." + return True, None + + +def alloc_mega_attn_output( + num_tokens: int, + n_wv_group: int, + device: torch.device, +) -> QuantizedActivation: + """Allocate the output buffer pair the mega-attention kernels write into. + + One pair per forward step: the prefill and decode segments fill disjoint + token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can + consume the result instead of one call per segment. + + The scale buffer is MN-major -- stride 1 along tokens, head-dim stride + ``ceil4(num_tokens)`` -- which is both what the kernel requires of a + caller-provided buffer and what DeepGEMM expects for this N. + """ + d = WV_GROUP_SIZE * _HEAD_DIM_V + data = torch.empty( + (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device + ) + aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize) + scale = torch.empty( + (n_wv_group, d // (_QUANT_GROUP * 4), aligned), + dtype=torch.int32, + device=device, + ).permute(2, 0, 1)[:num_tokens] + return QuantizedActivation( + data=data, + scale=scale, + orig_dtype=torch.bfloat16, + orig_shape=data.shape, + quant_key=kMxfp8Dynamic, + ) + + +def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation: + """The ``[start, end)`` token slice of a mega-attention output buffer.""" + data = out.data[start:end] + return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape) + + +class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): + """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100).""" + + backend_cls = FlashMLAMegaAttnBackend + + @classmethod + def is_available_for(cls, vllm_config: VllmConfig) -> bool: + """Whether this layer can serve the configured model on this device. + + The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a`` + group, which stops holding once TP divides the head count far enough + (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the + default-backend selector asks here first rather than turning an + unsupported topology into a startup crash. + """ + if not is_flashmla_mega_attn_supported()[0]: + return False + config = vllm_config.model_config.hf_config + config = getattr(config, "text_config", config) + n_heads = getattr(config, "num_attention_heads", 0) or 0 + n_groups = getattr(config, "o_groups", 0) or 0 + tp_size = vllm_config.parallel_config.tensor_parallel_size + if n_heads % tp_size or n_groups % tp_size: + return False + n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size + if not n_local_groups or n_local_heads % WV_GROUP_SIZE: + return False + return n_local_heads // n_local_groups == WV_GROUP_SIZE + + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: + super().__init__(vllm_config, *args, **kwargs) + if self.n_local_heads % WV_GROUP_SIZE: + raise ValueError( + f"{self.prefix}: mega attention needs the local head count " + f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}." + ) + if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE: + raise ValueError( + f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per " + "wo_a group." + ) + self.n_wv_group = self.padded_heads // WV_GROUP_SIZE + self._fused_layouts_ready = False + + # ---- interface contract ------------------------------------------------ + + @property + def accepts_unnormed_unroped_query(self) -> bool: + return True + + @property + def packed_kv_cache_dtype(self) -> CacheDType: + # This kernel is the only one that reads an NVFP4 compressed cache, so + # it is what an unspecific --kv-cache-dtype resolves to here. + return "nvfp4_ds_mla" + + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> QuantizedActivation: + return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device) + + def _o_proj( + self, attn_out: QuantizedActivation, positions: torch.Tensor + ) -> torch.Tensor: + """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output. + + The inverse RoPE and the FP8 cast happened inside the attention + kernel, and ``wo_a`` is permuted for its output layout, so there is + nothing to rotate or quantize here and ``positions`` is unused. The + scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the + einsum consumes the kernel's output with no repacking. + """ + groups = self.n_local_groups + z = torch.empty( + (attn_out.data.shape[0], groups, self.o_lora_rank), + dtype=torch.bfloat16, + device=attn_out.data.device, + ) + fp8_einsum( + "bhr,hdr->bhd", + (attn_out.data[:, :groups], attn_out.scale[:, :groups]), + (self.wo_a.weight, self.wo_a.weight_scale), + z, + recipe=self._einsum_recipe, + ) + return self.wo_b(z.flatten(1)) + + # ---- weights ----------------------------------------------------------- + + def finalize_loaded_weights(self) -> None: + """Permute wq_b rows / wo_a columns into the kernel's layouts. + + Idempotent: a second post-load pass must not permute twice. + """ + if self._fused_layouts_ready: + return + permute_wq_b_( + self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads + ) + permute_wo_a_( + self.wo_a.weight.data, + self.wo_a.weight_scale.data, + self.n_local_heads // self.n_local_groups, + ) + self._fused_layouts_ready = True + + # ---- forward ----------------------------------------------------------- + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: QuantizedActivation, + ) -> None: + if not self._fused_layouts_ready: + raise RuntimeError( + f"{self.prefix}: wq_b / wo_a were never permuted for the mega " + "attention kernel; refusing to run with mismatched layouts." + ) + attn_metadata = get_forward_context().attn_metadata + if attn_metadata is None: + # Warmup dummy run: reserve the prefill workspace, produce zeros. + self._reserve_prefill_workspace(q) + output.data.zero_() + output.scale.zero_() + return + + assert isinstance(attn_metadata, dict) + # q was zero-padded to the kernel's head count by the fused KV insert. + assert q.shape[1] == self.padded_heads + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, + attn_metadata.get(self.compressed_cache_prefix) + if self.compressed_cache_prefix is not None + else None, + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix] + ) + # The kernel takes int32 RoPE positions. + positions_int32 = positions.to(torch.int32) + num_decode_tokens = swa_metadata.num_decode_tokens + + if swa_metadata.num_prefills > 0: + self._forward_prefill_mega( + q[num_decode_tokens:], + positions_int32[num_decode_tokens:], + flashmla_metadata, + swa_metadata, + output, + num_decode_tokens, + ) + if swa_metadata.num_decodes > 0: + self._forward_decode_mega( + q[:num_decode_tokens], + positions_int32[:num_decode_tokens], + flashmla_metadata, + swa_metadata, + _token_slice(output, 0, num_decode_tokens), + ) + + def _reserve_prefill_workspace(self, q: torch.Tensor) -> None: + swa_only = self.compress_ratio == 0 + n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio) + m = n + self.window_size + self.max_num_batched_tokens + if swa_only: + top_k = 0 + else: + assert self.topk_indices_buffer is not None + top_k = self.topk_indices_buffer.shape[-1] + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + current_workspace_manager().get_simultaneous( + ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16), + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), + ) + + def _decode_compressed_kv_and_topk( + self, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """The compressed cache and its ``[n, topk]`` slot indices / lengths.""" + if self.compress_ratio == 0: + return None, None, None + assert flashmla_metadata is not None + assert swa_metadata.is_valid_token is not None + assert self.topk_indices_buffer is not None + num_decode_tokens = swa_metadata.num_decode_tokens + indices, lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + flashmla_metadata.block_table[: swa_metadata.num_decodes], + flashmla_metadata.block_size // self.compress_ratio, + swa_metadata.is_valid_token[:num_decode_tokens], + ) + return self._compressed_kv_cache().unsqueeze(-2), indices, lens + + def _forward_decode_mega( + self, + q: torch.Tensor, + positions_int32: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + out: QuantizedActivation, + ) -> None: + """Mega attention over the paged quantized caches, in place into ``out``. + + The caches are ``[num_blocks, page, 1, bytes]``, the record taken from + ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed + cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32 + slot ids (``block * page + offset``, ``-1`` invalid). + """ + extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk( + flashmla_metadata, swa_metadata + ) + assert swa_metadata.decode_swa_indices is not None + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode( + q, + self.swa_cache_layer.kv_cache.unsqueeze(-2), + swa_metadata.decode_swa_indices.view(q.shape[0], -1), + self.scale, + _HEAD_DIM_V, + self.attn_sink, + swa_metadata.decode_swa_lens, + extra_cache, + extra_idx, + extra_len, + False, # enable_q_norm + 0.0, # rms_norm_eps + positions_int32, + *_ROPE_ARGS, + self.rotary_emb.cos_sin_cache, + self.n_wv_group, + *_SF_ARGS, + out.data, + out.scale, + ) + + def _forward_prefill_mega( + self, + q: torch.Tensor, + positions_int32: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + out: QuantizedActivation, + token_base: int, + ) -> None: + swa_only = self.compress_ratio == 0 + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert seq_lens is not None and gather_lens is not None + assert query_start_loc_cpu is not None and query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:][ + : swa_metadata.num_prefill_tokens + ] + top_k = 0 if swa_only else topk_indices.shape[-1] + chunk_plan = swa_metadata.get_prefill_chunk_plan( + compress_ratio=self.compress_ratio, + prefill_chunk_size=self.PREFILL_CHUNK_SIZE, + has_compressed=not swa_only, + ) + assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" + workspace_manager = current_workspace_manager() + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: + chunk_size = chunk_end - chunk_start + kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous( + ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16), + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), + ) + if not swa_only: + assert flashmla_metadata is not None + dequantize_and_gather_k_cache( + kv_ws[:chunk_size], + self._compressed_kv_cache(), + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=flashmla_metadata.block_table[num_decodes:][ + chunk_start:chunk_end + ], + block_size=flashmla_metadata.block_size // self.compress_ratio, + offset=0, + ) + dequantize_and_gather_k_cache( + kv_ws[:chunk_size], + self.swa_cache_layer.kv_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_metadata.block_table[num_decodes:][ + chunk_start:chunk_end + ], + block_size=swa_metadata.block_size, + offset=chunk_n, + ) + first, last = num_decodes + chunk_start, num_decodes + chunk_end + qs = int(query_start_loc_cpu[first] - prefill_token_base) + qe = int(query_start_loc_cpu[last] - prefill_token_base) + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[qs:qe], + query_start_loc[first : last + 1], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + chunk_m, + chunk_n, + out=(idx_ws[: qe - qs], lens_ws[: qe - qs]), + left_visible=( + swa_metadata.prefill_left_visible[ + num_decode_tokens + qs : num_decode_tokens + qe + ] + if swa_metadata.prefill_left_visible is not None + else None + ), + right_visible=( + swa_metadata.prefill_right_visible[ + num_decode_tokens + qs : num_decode_tokens + qe + ] + if swa_metadata.prefill_right_visible is not None + else None + ), + max_image_tokens=self.max_image_tokens, + ) + chunk_out = _token_slice(out, token_base + qs, token_base + qe) + # Mega attention over the gathered non-paged bf16 KV (RoPE already + # applied), in place into this chunk's token range. `indices` + # entries outside [0, s_kv) skip. + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd( + q[qs:qe], + kv_ws.view(-1, 1, q.shape[-1]), + combined_indices.unsqueeze(1), + self.scale, + _HEAD_DIM_V, + self.attn_sink, + combined_lens, + False, # enable_q_norm + 0.0, # rms_norm_eps + positions_int32[qs:qe], + *_ROPE_ARGS, + self.rotary_emb.cos_sin_cache, + self.n_wv_group, + *_SF_ARGS, + chunk_out.data, + chunk_out.scale, + ) diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 7bed50bde71a..0b43cfad190b 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -201,7 +201,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, @@ -562,7 +563,8 @@ def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index ec834aa80dd9..47a7b35c76d6 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -58,7 +58,8 @@ def __init__(self, *args, **kwargs) -> None: self._o_proj_block_size ) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 9c9719e97729..21fa00e5c963 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -64,6 +64,9 @@ make_deepseek_v4_expert_params_mapping, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + DeepseekV4MegaAttnAttention, +) from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -119,8 +122,9 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: The generic CUDA backend selector does not instantiate DSv4 layers directly, so map generic sparse-MLA choices to the DSv4-specialized attention class. - Without an explicit backend, SM12 defaults to FlashInfer while the other - CUDA arches keep the FlashMLA path. + Without an explicit backend: SM12 takes FlashInfer, SM100 takes mega + attention where the topology allows it, and everything else keeps the + FlashMLA path. """ backend = vllm_config.attention_config.backend device_capability = current_platform.get_device_capability() @@ -140,6 +144,8 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: if device_capability is not None and device_capability.major == 12: return DeepseekV4FlashInferSM120Attention return DeepseekV4FlashInferMLAAttention + if backend is AttentionBackendEnum.FLASHMLA_MEGA_ATTN_DSV41: + return DeepseekV4MegaAttnAttention if backend in ( AttentionBackendEnum.FLASHMLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE_DSV4, @@ -149,6 +155,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: if device_capability is not None and device_capability.major == 12: return DeepseekV4FlashInferSM120Attention + # Mega attention is the SM100 default: it fuses Q RoPE, sparse attention, + # the output's inverse RoPE and its FP8 cast into one launch, and brings + # the 288 B NVFP4 compressed record -- the format the reference + # implementation itself stores. It declines topologies it cannot serve + # (non-SM100, TP that leaves fewer than WV_GROUP_SIZE heads per wo_a + # group, a build without the kernel), which then fall through to FlashMLA. + if DeepseekV4MegaAttnAttention.is_available_for(vllm_config): + return DeepseekV4MegaAttnAttention return DeepseekV4FlashMLAAttention @@ -936,6 +950,17 @@ def finalize_mega_moe_weights(self) -> None: for layer in islice(self.layers, self.start_layer, self.end_layer): layer.ffn.finalize_mega_moe_weights() + def finalize_mega_attn_weights(self) -> None: + """Permute wq_b / wo_a into FlashMLA's mega-attention layouts. + + A no-op for every other attention layer, and idempotent, so a second + post-load pass cannot permute twice. + """ + for layer in islice(self.layers, self.start_layer, self.end_layer): + finalize = getattr(layer.attn, "finalize_loaded_weights", None) + if finalize is not None: + finalize() + def finalize_mhc_broadcast_weights(self) -> None: if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer: return @@ -1190,6 +1215,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() + self.model.finalize_mega_attn_weights() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 6a0f09f952e4..3d69499d2a80 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -229,3 +229,54 @@ def get_name() -> str: @staticmethod def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]: return DeepseekV4FlashMLAMetadataBuilder + + +class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend): + """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse + RoPE + FP8 cast of the output, in one launch. + + Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what + differs is the attention layer's interface contract (it takes an unnormed, + unroped Q and returns an already-inverse-RoPE'd, quantized output) and the + extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read. + SM100 only; the kernel has no SM90 instantiation. + """ + + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "fp8_ds_mla", + "fp8", # alias for fp8_ds_mla + "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache + ] + + @staticmethod + def get_name() -> str: + return "FLASHMLA_MEGA_ATTN_DSV41" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [128] + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major == 10 + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + # Imported here: the layer module imports this backend class. + from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + is_flashmla_mega_attn_supported, + ) + + return is_flashmla_mega_attn_supported()[1] diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c4f3cd8b4bdd..700f328a37c0 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -112,7 +112,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural # element-size page. - uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" + uses_fp8_ds_mla_layout = self.cache_config.cache_dtype in ( + "fp8_ds_mla", + "nvfp4_ds_mla", + ) return SlidingWindowMLASpec( block_size=self.block_size, num_kv_heads=1, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 571566db1f25..dcaaffaaa8eb 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -114,6 +114,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) + FLASHMLA_MEGA_ATTN_DSV41 = ( + "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend" + ) B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend" FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" FLASH_ATTN_MLA_SPARSE = ( From 118e17f5ff301f67f1f4f5bdae2c9ae5e50242c5 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 16 Sep 2026 18:08:36 +0800 Subject: [PATCH 319/774] [Platform] Move env check function to platform interface (#48599) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/engine/arg_utils.py | 2 +- vllm/envs.py | 9 --------- vllm/platforms/interface.py | 18 ++++++++++++++++++ 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b420e84f542b..6128bc625589 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2051,7 +2051,7 @@ def create_engine_config( device_config = DeviceConfig(device=cast(Device, current_platform.device_type)) - envs.validate_environ(self.fail_on_environ_validation) + current_platform.validate_environ(self.fail_on_environ_validation) # Check if the model is a speculator and override model/tokenizer/config # BEFORE creating ModelConfig, so the config is created with the target model diff --git a/vllm/envs.py b/vllm/envs.py index bc4724e2e4a7..333d75b0ff0e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2249,15 +2249,6 @@ def is_set(name: str): raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -def validate_environ(hard_fail: bool) -> None: - for env in os.environ: - if env.startswith("VLLM_") and env not in environment_variables: - if hard_fail: - raise ValueError(f"Unknown vLLM environment variable detected: {env}") - else: - logger.warning("Unknown vLLM environment variable detected: %s", env) - - def compile_factors() -> dict[str, object]: """Return env vars used for torch.compile cache keys. diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index bd5d98191201..98db974a02e4 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1334,6 +1334,24 @@ def is_arch_support_pdl(cls) -> bool: """ return False + @classmethod + def validate_environ(cls, hard_fail: bool) -> None: + """ + Validate environment variables for the current platform. + """ + from vllm import envs + + for env in os.environ: + if env.startswith("VLLM_") and env not in envs.environment_variables: + if hard_fail: + raise ValueError( + f"Unknown vLLM environment variable detected: {env}" + ) + else: + logger.warning( + "Unknown vLLM environment variable detected: %s", env + ) + class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED From 3bb782621492711485dc86791b5978128783814a Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Wed, 16 Sep 2026 06:55:26 -0400 Subject: [PATCH 320/774] [Agent] Add agent instructions for parser directories (#56883) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- vllm/parser/AGENTS.md | 41 ++++++++++++++++++++++++++++ vllm/parser/CLAUDE.md | 1 + vllm/reasoning/AGENTS.md | 54 +++++++++++++++++++++++++++++++++++++ vllm/reasoning/CLAUDE.md | 1 + vllm/tool_parsers/AGENTS.md | 45 +++++++++++++++++++++++++++++++ vllm/tool_parsers/CLAUDE.md | 1 + 6 files changed, 143 insertions(+) create mode 100644 vllm/parser/AGENTS.md create mode 100644 vllm/parser/CLAUDE.md create mode 100644 vllm/reasoning/AGENTS.md create mode 100644 vllm/reasoning/CLAUDE.md create mode 100644 vllm/tool_parsers/AGENTS.md create mode 100644 vllm/tool_parsers/CLAUDE.md diff --git a/vllm/parser/AGENTS.md b/vllm/parser/AGENTS.md new file mode 100644 index 000000000000..d9e6b79707a4 --- /dev/null +++ b/vllm/parser/AGENTS.md @@ -0,0 +1,41 @@ +# vllm/parser + +Model-output parsing for the serving layer. `Parser` (`abstract_parser.py`) +unifies reasoning extraction and tool-call extraction behind one object; the +serving layer calls `parse` / `parse_delta` and `is_reasoning_end` on it. + +## Layout + +- `abstract_parser.py` — `Parser`, and `DelegatingParser` for wrapping a legacy + `ReasoningParser` + `ToolParser` pair. +- `parser_manager.py` — resolves parser names to classes. +- `engine/` — the declarative parser engine. A `ParserEngineConfig` + (`parser_engine_config.py`) is a set of terminals, transitions and content + events; `parser_engine.py` and `streaming_parser_engine.py` drive it for the + non-streaming and streaming paths; `token_id_scanner.py` pre-lexes + special-token ids and `incremental_lexer.py` lexes text; `adapters.py` + exposes an engine through the legacy reasoning/tool parser interfaces. +- `.py` — one file per model family. Most subclass `ParserEngine` + (`deepseek_v4.py`, `qwen3.py`, ...); a few are `DelegatingParser` wrappers + over legacy parsers (`harmony.py`, `kimi_k3.py`). + +## Guidelines + +- A new parser either subclasses `ParserEngine` with a declarative config, or + subclasses `Parser` and implements `parse` / `parse_delta` directly. Use the + engine when the format is expressible as terminals and transitions, and + subclass an existing config when it is a variant of one (`seed_oss.py`, + `ling3.py`). `DelegatingParser` exists for backward compatibility with the + separate `ReasoningParser` / `ToolParser` classes; do not use it for new + formats. +- Keep a model-specific fix inside that model's `.py` where possible. + Change `engine/` or `abstract_parser.py` only when the behavior is wrong for + every parser, and say so in the PR. +- Text emitted in a state with no `content_events` entry is discarded. Make + that deliberate when adding a state. +- Special tokens reach the engine as token ids, not text. A text terminal that + spans a special token never matches in streaming; list it in + `token_id_terminals`. +- Comment only where the code alone would leave a reader with a question — + a non-obvious transition, a deliberate swallow, a model quirk being worked + around. Keep it minimal and concise. diff --git a/vllm/parser/CLAUDE.md b/vllm/parser/CLAUDE.md new file mode 100644 index 000000000000..43c994c2d361 --- /dev/null +++ b/vllm/parser/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/vllm/reasoning/AGENTS.md b/vllm/reasoning/AGENTS.md new file mode 100644 index 000000000000..1f7c8b3557b1 --- /dev/null +++ b/vllm/reasoning/AGENTS.md @@ -0,0 +1,54 @@ +# vllm/reasoning + +Every `--reasoning-parser` name resolves here. The parser has two consumers: +`ParserManager.get_parser` composes it with the tool parser into the unified +`Parser` the serving layer uses, and the structured-output manager +(`vllm/v1/structured_output/__init__.py`) instantiates it directly and calls +`is_reasoning_end` / `is_reasoning_end_streaming` every decode step to decide +when the grammar bitmask starts applying. The `ReasoningParser` interface +itself is legacy: it is kept so existing in-tree parsers and out-of-tree +plugins keep working, not as the base for new ones. + +## Layout + +- `abs_reasoning_parsers.py` — `ReasoningParser` (`extract_reasoning`, + `extract_reasoning_streaming`, `is_reasoning_end`, + `is_reasoning_end_streaming`, `extract_content_ids`, `adjust_request`) and + `ReasoningParserManager`. +- `__init__.py` — `_REASONING_PARSERS_TO_REGISTER`, the name → (file, class) + table that lazily registers every built-in parser. +- `basic_parsers.py` — `BaseThinkingReasoningParser`, the base for parsers + that only need a start/end token pair. +- `_reasoning_parser.py` — one file per family. Either a subclass of a + `ParserReasoningAdapter` generated from a `vllm/parser/.py` engine + grammar (`mistral_reasoning_parser.py`, `kimi_k2_reasoning_parser.py`; the + `_engine_` suffix is not applied consistently), or a self-contained + `ReasoningParser` subclass. + +## Guidelines + +- Do not change the `ReasoningParser` interface in `abs_reasoning_parsers.py`. + Out-of-tree plugins loaded through + `ReasoningParserManager.import_reasoning_parser` subclass it. New + capabilities go on `Parser` in `vllm/parser/`; if a legacy parser needs the + same capability, `DelegatingParser` supplies the default. +- For a new model, implement a unified `Parser` in `vllm/parser/` first, then + expose it here as a thin adapter subclass plus an entry in + `_REASONING_PARSERS_TO_REGISTER`. +- The `_REASONING_PARSERS_TO_REGISTER` key is what users pass on the command + line; treat renames as breaking. +- Keep adapters thin: registration, `adjust_request` tweaks, and class + attributes. Parsing logic belongs in `vllm/parser/`. +- Keep a model-specific fix inside that model's `_reasoning_parser.py`, + or its grammar in `vllm/parser/.py`, where possible. Touch + `abs_reasoning_parsers.py` or `basic_parsers.py` only for behavior every + parser needs, and say so in the PR. +- `is_reasoning_end_streaming` runs on the engine core for every decode step + of every structured-output request. Work on `delta_ids`; re-scanning the + full sequence makes the step cost grow with sequence length. +- The structured-output gate and the frontend's extraction must agree on + where reasoning ends, whether the turn continues with content or a tool + call. Test the end-detection methods with the same inputs as extraction. +- Comment only where the code alone would leave a reader with a question — + a request tweak in `adjust_request`, a marker handled unusually, a model + quirk being worked around. Keep it minimal and concise. diff --git a/vllm/reasoning/CLAUDE.md b/vllm/reasoning/CLAUDE.md new file mode 100644 index 000000000000..43c994c2d361 --- /dev/null +++ b/vllm/reasoning/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/vllm/tool_parsers/AGENTS.md b/vllm/tool_parsers/AGENTS.md new file mode 100644 index 000000000000..2bd487997c57 --- /dev/null +++ b/vllm/tool_parsers/AGENTS.md @@ -0,0 +1,45 @@ +# vllm/tool_parsers + +Every `--tool-call-parser` name resolves here. `ParserManager.get_parser` +composes the tool parser chosen by name with the reasoning parser into the +unified `Parser` the serving layer uses. The `ToolParser` interface itself is +legacy: it is kept so existing in-tree parsers and out-of-tree plugins keep +working, not as the base for new ones. + +## Layout + +- `abstract_tool_parser.py` — `ToolParser` (`extract_tool_calls`, + `extract_tool_calls_streaming`, `adjust_request`) and `ToolParserManager`. +- `__init__.py` — `_TOOL_PARSERS_TO_REGISTER`, the name → (file, class) table + that lazily registers every built-in parser. +- `structural_tag_registry.py` — structural-tag builders used to constrain + tool-call output for `tool_choice="required"` / named tools. +- `_tool_parser.py` — one file per family. Either a subclass of a + `ParserToolAdapter` generated from a `vllm/parser/.py` engine grammar + (`mistral_tool_parser.py`, `kimi_k2_tool_parser.py`; the `_engine_` suffix + is not applied consistently), or a self-contained `ToolParser` subclass. + +## Guidelines + +- Do not change the `ToolParser` interface in `abstract_tool_parser.py`. + Out-of-tree plugins loaded through `ToolParserManager.import_tool_parser` + subclass it. New capabilities go on `Parser` in `vllm/parser/`; if a legacy + parser needs the same capability, `DelegatingParser` supplies the default. +- For a new model, implement a unified `Parser` in `vllm/parser/` first, then + expose it here as a thin adapter subclass plus an entry in + `_TOOL_PARSERS_TO_REGISTER`. +- The `_TOOL_PARSERS_TO_REGISTER` key is what users pass on the command line; + treat renames as breaking. +- Keep adapters thin: registration, `adjust_request` tweaks, and class + attributes. Parsing logic belongs in `vllm/parser/`. +- Keep a model-specific fix inside that model's `_tool_parser.py`, or + its grammar in `vllm/parser/.py`, where possible. Touch + `abstract_tool_parser.py` or `structural_tag_registry.py` only for behavior + every parser needs, and say so in the PR. +- Set `structural_tag_model` only to a key in `SUPPORTED_STRUCTURAL_TAG_MODELS` + (an xgrammar builtin, or one registered with `register_vllm_structural_tag`). + It only takes effect under `VLLM_ENFORCE_STRICT_TOOL_CALLING`, where it also + disables the generic required/named handling for that parser. +- Comment only where the code alone would leave a reader with a question — + a request tweak in `adjust_request`, a marker handled unusually, a model + quirk being worked around. Keep it minimal and concise. diff --git a/vllm/tool_parsers/CLAUDE.md b/vllm/tool_parsers/CLAUDE.md new file mode 100644 index 000000000000..43c994c2d361 --- /dev/null +++ b/vllm/tool_parsers/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md From 6a2fdf9ac6f929f937dbe4617f1993c068045f81 Mon Sep 17 00:00:00 2001 From: Harshit Kedia <140254510+harshit-sarvam@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:32:47 +0530 Subject: [PATCH 321/774] [Model] Optimize Sarvam MLA routing and preserve FP32 router logits (#56034) Signed-off-by: Harshit Kedia Co-authored-by: Claude Opus 5 (1M context) --- vllm/model_executor/models/sarvam.py | 31 +++++++++++++++++++--------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 14d68fec568f..2cc8a27f4dcf 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -36,7 +36,11 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, + GateLinear, + MoERunner, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -306,11 +310,11 @@ def __init__( else: self.router_dtype = torch.bfloat16 - self.gate = nn.Linear( + self.gate = GateLinear( self.hidden_size, self.num_experts, - bias=False, - dtype=self.router_dtype, + out_dtype=self.router_dtype, + prefix=f"{prefix}.gate", ) if getattr(config, "moe_router_enable_expert_bias", True): @@ -365,12 +369,7 @@ def maybe_get_fused_moe(self) -> MoERunner: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) - router_logits = self.gate( - hidden_states.to(self.router_dtype) - if self.router_dtype is not None - else hidden_states - ) - router_logits = router_logits.to(hidden_states.dtype) + router_logits, _ = self.gate(hidden_states) final_hidden = self.experts( hidden_states=hidden_states, router_logits=router_logits, @@ -643,9 +642,21 @@ class SarvamMLAForCausalLM( "gate_up_proj": ["gate_proj", "up_proj"], } + @staticmethod + def _remap_config(config) -> None: + """Default the routing keys the released checkpoints omit.""" + defaults = { + "n_group": 1, + "topk_group": 1, + } + for attr, default in defaults.items(): + if getattr(config, attr, None) is None: + setattr(config, attr, default) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() config = vllm_config.model_config.hf_config + self._remap_config(config) quant_config = vllm_config.quant_config self.config = config self.quant_config = quant_config From 903285fbcc4e51332d562499bd90ec9a1513c813 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Wed, 16 Sep 2026 05:14:54 -0700 Subject: [PATCH 322/774] [Bugfix] Let an optional Literal flag accept the None it advertises (#52370) Signed-off-by: Vineeth Sai Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 27 +++++++++++++++++++++++---- vllm/engine/arg_utils.py | 4 ++-- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index d6ce5b36b514..e58379874601 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -254,10 +254,13 @@ def test_get_type_hints(type_hint, expected): assert get_type_hints(type_hint) == expected -def test_get_kwargs(): - kwargs = get_kwargs(DummyConfig) - print(kwargs) +@pytest.fixture +def dummy_config_kwargs(): + return get_kwargs(DummyConfig) + +def test_get_kwargs(dummy_config_kwargs): + kwargs = dummy_config_kwargs # bools should not have their type set assert kwargs["regular_bool"].get("type") is None assert kwargs["optional_bool"].get("type") is None @@ -267,7 +270,7 @@ def test_get_kwargs(): assert kwargs["optional_bool_or_str"]["const"] is True assert "action" not in kwargs["optional_bool_or_str"] # optional literals should have None as a choice - assert kwargs["optional_literal"]["choices"] == ["x", "y", "None"] + assert kwargs["optional_literal"]["choices"] == ["x", "y", None] # tuples should have the correct nargs assert kwargs["tuple_n"]["nargs"] == "+" assert kwargs["tuple_2"]["nargs"] == 2 @@ -293,6 +296,22 @@ def test_get_kwargs(): assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2) # type: ignore[call-arg] +@pytest.mark.parametrize( + ("args", "expected"), + [ + (["--optional-literal", "None"], None), + (["--optional-literal", ""], None), + (["--optional-literal", "x"], "x"), + ], +) +def test_optional_handling(args, expected, dummy_config_kwargs): + parser = FlexibleArgumentParser() + parser.add_argument("--optional-literal", **dummy_config_kwargs["optional_literal"]) + + assert parser.parse_args(args).optional_literal is expected + assert "None" in parser.format_help() + + def test_jit_monitor_verbose_arg(): parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = parser.parse_args(["--jit-monitor-verbose"]) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 6128bc625589..a977b64649eb 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -411,7 +411,7 @@ def parse_dataclass(val: str, cls=dataclass_cls) -> Any: if type(None) in type_hints and not contains_type(type_hints, bool): kwargs[name]["type"] = optional_type(kwargs[name]["type"]) if kwargs[name].get("choices"): - kwargs[name]["choices"].append("None") + kwargs[name]["choices"].append(None) return kwargs @@ -1526,7 +1526,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: ) # TODO: generalise this special case choices = observability_kwargs["collect_detailed_traces"]["choices"] - metavar = f"{{{','.join(choices)}}}" + metavar = f"{{{','.join(str(c) for c in choices)}}}" observability_kwargs["collect_detailed_traces"]["metavar"] = metavar observability_kwargs["collect_detailed_traces"]["choices"] += [ ",".join(p) for p in permutations(get_args(DetailedTraceModules), r=2) From 8b1d188046034b42024d22ee52fa6c3948b4acbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 16 Sep 2026 14:55:40 +0200 Subject: [PATCH 323/774] [Bugfix][Mooncake] Report request-level KV load failures under HMA (#56855) Signed-off-by: NickLucche --- .../unit/test_kv_load_failure_recovery.py | 345 +----------------- .../unit/test_mooncake_connector_hma.py | 77 ++++ .../unit/test_mooncake_store_connector.py | 33 ++ .../unit/test_mooncake_store_worker.py | 82 +++++ tests/v1/kv_connector/unit/utils.py | 2 - .../v1/mooncake/mooncake_connector.py | 46 ++- .../v1/mooncake/store/connector.py | 9 + .../kv_connector/v1/mooncake/store/worker.py | 60 ++- vllm/v1/core/sched/scheduler.py | 52 +-- 9 files changed, 329 insertions(+), 377 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 74bac5559823..a54a023a13c4 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -2,21 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from typing import Literal -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest -import torch -from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler -from vllm.v1.kv_cache_interface import ( - KVCacheConfig, - KVCacheGroupSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, -) -from vllm.v1.request import FinishReason, Request, RequestStatus +from vllm.v1.request import Request, RequestStatus from .utils import ( create_model_runner_output, @@ -44,18 +35,6 @@ def scheduler(): return create_scheduler(vllm_config) -@pytest.fixture -def hybrid_scheduler(): - """Scheduler for a hybrid model: full attention plus a Mamba group.""" - vllm_config = create_vllm_config(kv_load_failure_policy="recompute") - kv_cache_config = make_kv_cache_config( - block_size=vllm_config.cache_config.block_size, - mamba_enabled=True, - num_blocks=10000, - ) - return create_scheduler(vllm_config, kv_cache_config=kv_cache_config) - - @pytest.mark.parametrize( "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs", [ @@ -361,317 +340,19 @@ def test_async_progressive_load_failure( assert scheduler.connector.get_num_new_matched_tokens.call_count == 1 -def _schedule_hybrid_async_load( - scheduler: Scheduler, num_prompt_blocks: int, num_external_computed_blocks: int -) -> Request: - """Put one request into an async external load on a hybrid model.""" - request = create_request(num_tokens=num_prompt_blocks * scheduler.block_size) - scheduler.add_request(request=request) - - scheduler.connector = Mock() - scheduler.connector.get_num_new_matched_tokens.side_effect = ( - _make_get_num_new_matched_tokens( - {request.request_id: num_external_computed_blocks * scheduler.block_size}, - async_load=True, - ) - ) - scheduler.connector.take_events.return_value = () - return request - - -def _report_invalid_blocks( - scheduler: Scheduler, scheduler_output, invalid_block_ids: set[int] -) -> None: - scheduler.update_from_output( - scheduler_output, - create_model_runner_output( - reqs=[], - finished_recving=set(), - invalid_block_ids=invalid_block_ids, - use_eos=True, - ), - ) +@pytest.mark.cpu_test +def test_block_level_load_failure_fails_closed_with_multiple_groups(): + """Block-level invalid-block reporting is rejected on multi-group layouts. - -@pytest.mark.parametrize("failed_group_idx", [0, 1]) -def test_hybrid_load_failure_recomputes_whole_request( - hybrid_scheduler: Scheduler, failed_group_idx: int -): - """A hybrid model must recover, not crash. - - Its groups sit on different block grids, so there is no single truncation - point; the request restarts from scratch instead. + Block IDs are only unique within a group, so the scheduler cannot map + them back to requests; connectors must report failed_recving instead. """ - num_prompt_blocks, num_external_computed_blocks = 10, 9 - request = _schedule_hybrid_async_load( - hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks - ) - scheduler_output = hybrid_scheduler.schedule() - - block_ids_per_group = hybrid_scheduler.kv_cache_manager.get_block_ids( - request.request_id - ) - assert len(block_ids_per_group) == 2, "expected a hybrid (multi-group) layout" - null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id - # The Mamba group backs only its aligned snapshot with a real block, so - # pick a block that actually belongs to this request. - failed_block = next( - block_id - for block_id in block_ids_per_group[failed_group_idx] - if block_id != null_block_id - ) - - _report_invalid_blocks(hybrid_scheduler, scheduler_output, {failed_block}) - - assert request.num_computed_tokens == 0 - assert hybrid_scheduler.failed_recving_kv_req_ids == {request.request_id} - assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS - - -def test_hybrid_load_failure_ignores_null_block(hybrid_scheduler: Scheduler): - """The null block is shared by every request and holds no KV, so reporting - it must not restart anyone.""" - num_prompt_blocks, num_external_computed_blocks = 10, 9 - request = _schedule_hybrid_async_load( - hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks - ) - scheduler_output = hybrid_scheduler.schedule() - - null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id - _, mamba_block_ids = hybrid_scheduler.kv_cache_manager.get_block_ids( - request.request_id - ) - # The Mamba group parks its unaligned positions on the null block, so a - # request that never failed would otherwise be restarted by this report. - assert mamba_block_ids[0] == null_block_id - - num_computed_tokens = request.num_computed_tokens - _report_invalid_blocks(hybrid_scheduler, scheduler_output, {null_block_id}) - - assert request.num_computed_tokens == num_computed_tokens - assert not hybrid_scheduler.failed_recving_kv_req_ids - - -def _make_deepseek_v4_five_group_config(num_blocks: int) -> KVCacheConfig: - """Build the scheduler-side projection of V4-Flash's packed KV layout.""" - - def mla_spec() -> MLAAttentionSpec: - return MLAAttentionSpec( - block_size=256, - num_kv_heads=1, - head_size=512, - dtype=torch.bfloat16, - tokens_per_state=4, - model_version="deepseek_v4", - ) - - def swa_spec( - *, block_size: int, sliding_window: int, head_size: int - ) -> SlidingWindowMLASpec: - return SlidingWindowMLASpec( - block_size=block_size, - num_kv_heads=1, - head_size=head_size, - dtype=torch.float32, - sliding_window=sliding_window, - model_version="deepseek_v4", - ) - - # DeepSeek-V4 packs these into five physical groups. The two SWA groups - # have identical scheduling semantics but are split to align layer tuples. - groups = [ - KVCacheGroupSpec(["full_mla_and_indexer"], mla_spec()), - KVCacheGroupSpec( - ["c4_state"], - swa_spec(block_size=4, sliding_window=8, head_size=2048), - ), - KVCacheGroupSpec( - ["c128_state"], - swa_spec(block_size=8, sliding_window=128, head_size=2048), - ), - KVCacheGroupSpec( - ["swa.0"], - swa_spec(block_size=64, sliding_window=512, head_size=512), - ), - KVCacheGroupSpec( - ["swa.1"], - swa_spec(block_size=64, sliding_window=512, head_size=512), - ), - ] - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[], - kv_cache_groups=groups, - ) - - -def _create_deepseek_v4_scheduler( - kv_load_failure_policy: Literal["recompute", "fail"] = "recompute", -) -> Scheduler: - num_blocks = 10000 - vllm_config = create_vllm_config( - block_size=256, - max_num_batched_tokens=4096, - max_model_len=8192, - # max_model_len only reaches SchedulerConfig; opt-125m still caps the - # model at 2048, which is below the 9-block requests used here. - hf_overrides={"max_position_embeddings": 8192}, - kv_load_failure_policy=kv_load_failure_policy, - ) - return create_scheduler( + vllm_config = create_vllm_config() + scheduler = create_scheduler( vllm_config, - num_blocks=num_blocks, - kv_cache_config=_make_deepseek_v4_five_group_config(num_blocks), - hash_block_size=4, + kv_cache_config=make_kv_cache_config(block_size=16, swa_enabled=True), ) + assert len(scheduler.kv_cache_config.kv_cache_groups) == 2 - -@pytest.fixture -def deepseek_v4_scheduler() -> Scheduler: - return _create_deepseek_v4_scheduler() - - -def _schedule_deepseek_v4_async_load( - scheduler: Scheduler, -) -> tuple[Request, SchedulerOutput]: - num_external_tokens = 8 * scheduler.block_size - request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) - scheduler.add_request(request) - - scheduler.connector = Mock() - scheduler.connector.get_num_new_matched_tokens.side_effect = ( - _make_get_num_new_matched_tokens( - {request.request_id: num_external_tokens}, async_load=True - ) - ) - scheduler.connector.take_events.return_value = () - return request, scheduler.schedule() - - -@pytest.mark.parametrize("failed_group_idx", range(5)) -def test_deepseek_v4_load_failure_recomputes_whole_request( - deepseek_v4_scheduler: Scheduler, failed_group_idx: int -): - """A failed load in any V4-Flash KV group must not kill EngineCore.""" - request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) - - block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( - request.request_id - ) - assert len(block_ids_per_group) == 5 - null_block_id = ( - deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id - ) - failed_block_id = next( - block_id - for block_id in block_ids_per_group[failed_group_idx] - if block_id != null_block_id - ) - - deepseek_v4_scheduler.update_from_output( - scheduler_output, - create_model_runner_output( - reqs=[], - finished_recving=set(), - invalid_block_ids={failed_block_id}, - use_eos=True, - ), - ) - - assert request.num_computed_tokens == 0 - assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS - assert deepseek_v4_scheduler.failed_recving_kv_req_ids == {request.request_id} - - -def test_deepseek_v4_load_failure_ignores_shared_null_block( - deepseek_v4_scheduler: Scheduler, -): - """The shared null block contains no request KV and cannot fail a load.""" - request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) - null_block_id = ( - deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id - ) - block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( - request.request_id - ) - assert any(null_block_id in group for group in block_ids_per_group[1:]) - num_computed_tokens = request.num_computed_tokens - - deepseek_v4_scheduler.update_from_output( - scheduler_output, - create_model_runner_output( - reqs=[], - finished_recving=set(), - invalid_block_ids={null_block_id}, - use_eos=True, - ), - ) - - assert request.num_computed_tokens == num_computed_tokens - assert not deepseek_v4_scheduler.failed_recving_kv_req_ids - - -def test_deepseek_v4_sync_load_failure_fails_and_evicts_real_blocks(): - """Fail policy must terminate cleanly and never evict the null block.""" - scheduler = _create_deepseek_v4_scheduler(kv_load_failure_policy="fail") - num_external_tokens = 8 * scheduler.block_size - request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) - scheduler.add_request(request) - - scheduler.connector = Mock() - scheduler.connector.get_num_new_matched_tokens.side_effect = ( - _make_get_num_new_matched_tokens( - {request.request_id: num_external_tokens}, async_load=False - ) - ) - scheduler.connector.request_finished.return_value = (False, None) - scheduler.connector.take_events.return_value = () - scheduler_output = scheduler.schedule() - - block_ids_per_group = scheduler.kv_cache_manager.get_block_ids(request.request_id) - assert len(block_ids_per_group) == 5 - null_block_id = scheduler.kv_cache_manager.block_pool.null_block.block_id - real_block_ids = { - block_id - for group_block_ids in block_ids_per_group - for block_id in group_block_ids - if block_id != null_block_id - } - failed_block_id = next( - block_id for block_id in block_ids_per_group[2] if block_id != null_block_id - ) - - evicted_block_ids: list[set[int]] = [] - original_evict_blocks = scheduler.kv_cache_manager.evict_blocks - - def evict_blocks_spy(block_ids: set[int]) -> None: - evicted_block_ids.append(set(block_ids)) - original_evict_blocks(block_ids) - - with ( - patch.object( - scheduler.kv_cache_manager, - "evict_blocks", - side_effect=evict_blocks_spy, - ), - patch.object(scheduler, "_connector_finished", return_value=(False, None)), - ): - outputs = scheduler.update_from_output( - scheduler_output, - create_model_runner_output( - [request], - invalid_block_ids={failed_block_id}, - use_eos=True, - ), - ) - - assert request.status == RequestStatus.FINISHED_ERROR - assert request.get_finished_reason() == FinishReason.ERROR - assert request.request_id not in scheduler.requests - assert evicted_block_ids == [real_block_ids] - assert null_block_id not in evicted_block_ids[0] - - engine_outputs = next(iter(outputs.values())) - assert len(engine_outputs.outputs) == 1 - assert engine_outputs.outputs[0].request_id == request.request_id - assert engine_outputs.outputs[0].finish_reason == FinishReason.ERROR + with pytest.raises(RuntimeError, match="failed_recving"): + scheduler._handle_invalid_blocks({1}, {}) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py index 54b97d482eff..f3739ef4301a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py @@ -18,6 +18,7 @@ MooncakeConnectorMetadata, MooncakeConnectorScheduler, MooncakeXferMetadata, + PullReqMeta, SendBlockMeta, TransferRegion, ) @@ -436,3 +437,79 @@ def test_request_finished_with_hma_groups(): assert stored_blocks[0] == fa_blocks # SW: clipped to last 9 blocks (sw_size=128, block_size=16 → 8+1=9) assert stored_blocks[1] == sw_blocks[-9:] + + +# --------------------------------------------------------------------------- +# Worker-side load-failure reporting (HMA vs non-HMA) +# --------------------------------------------------------------------------- +def _make_kv_consumer_worker(swa_enabled: bool, disable_hma: bool = False): + block_size = 16 + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", + kv_role="kv_consumer", + block_size=block_size, + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma + kv_cache_config = make_kv_cache_config( + block_size=block_size, swa_enabled=swa_enabled + ) + with set_current_vllm_config(vllm_config), patch_worker_dependencies(): + connector = MooncakeConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + return connector.connector_worker + + +def _make_pull_meta(d_req_id: str, local_block_ids: list[list[int]]) -> PullReqMeta: + return PullReqMeta( + d_req_id=d_req_id, + transfer_id="tx-1", + local_block_ids=local_block_ids, + remote_engine_id="engine-0", + remote_bootstrap_addr="127.0.0.1:1234", + ) + + +# The tests below construct a MooncakeConnectorWorker, whose __init__ +# requires an accelerator device, so (like the worker tests in +# test_mooncake_connector.py) they run on the GPU KV-connectors job only. +@pytest.mark.parametrize( + "swa_enabled,disable_hma,expected_is_hma", + [ + (True, False, True), # SWA group present, HMA enabled + (True, True, False), # SWA group present, but HMA disabled + (False, False, False), # FA only, HMA not needed + ], +) +def test_worker_is_hma_required(swa_enabled, disable_hma, expected_is_hma): + """Worker-side _is_hma_required mirrors the scheduler's derivation.""" + worker = _make_kv_consumer_worker(swa_enabled, disable_hma) + assert worker._is_hma_required is expected_is_hma + + +def test_worker_failed_recv_reports_request_level_failure_with_hma(): + """With HMA, load failures report the request, not ambiguous block IDs.""" + worker = _make_kv_consumer_worker(swa_enabled=True) + assert worker._is_hma_required + + pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]]) + worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom") + + assert worker.get_block_ids_with_load_errors() == set() + results = worker.get_transfer_results() + assert results.failed_recving == {"d-req-1"} + assert results.finished_recving == {"d-req-1"} + + +def test_worker_failed_recv_reports_block_ids_without_hma(): + """Without HMA, load failures keep reporting block-level errors.""" + worker = _make_kv_consumer_worker(swa_enabled=False) + assert not worker._is_hma_required + + pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]]) + worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom") + + assert worker.get_block_ids_with_load_errors() == {1, 2, 3, 4} + results = worker.get_transfer_results() + assert results.failed_recving == set() + assert results.finished_recving == {"d-req-1"} diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 6afb529d145f..f3ffa55a5439 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -12,6 +12,7 @@ from vllm.distributed.kv_events import BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorRole, + KVConnectorTransferResults, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( connector as mooncake_store_connector, @@ -245,6 +246,38 @@ def test_worker_methods_delegate_to_store_worker(): assert invalid_block_ids == {3, 4} +def test_get_transfer_results_delegates_to_store_worker(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + metadata = MooncakeStoreConnectorMetadata(set(), set()) + finished_req_ids = {"req-1"} + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + transfer_results = KVConnectorTransferResults( + finished_sending={"req-1"}, + finished_recving={"req-2"}, + failed_recving={"req-3"}, + ) + worker.get_transfer_results.return_value = transfer_results + connector.bind_connector_metadata(metadata) + + result = connector.get_transfer_results(finished_req_ids) + + worker.get_transfer_results.assert_called_once_with(finished_req_ids, metadata) + assert result is transfer_results + + def test_get_kv_connector_kv_cache_events_returns_none_when_disabled(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 9baf16e2610d..dea4178ff3e8 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -177,6 +177,7 @@ def _make_store_recving_thread( *, tp_rank: int = 0, disk_offload_buffer_budget_bytes: int | None = None, + is_hma_required: bool = False, ) -> mooncake_store_worker.KVCacheStoreRecvingThread: from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec @@ -199,6 +200,7 @@ def _make_store_recving_thread( ready_event=threading.Event(), coord=coord, disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes, + is_hma_required=is_hma_required, ) thread.request_queue.task_done = MagicMock() return thread @@ -303,6 +305,7 @@ def _make_vllm_config( kv_role: str = "kv_both", pipeline_parallel_size: int = 1, kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC, + disable_hybrid_kv_cache_manager: bool = True, ) -> SimpleNamespace: cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10) cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout @@ -319,6 +322,9 @@ def _make_vllm_config( kv_role=kv_role, extra_config=extra_config ), cache_config=cache_config, + scheduler_config=SimpleNamespace( + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager + ), kv_events_config=SimpleNamespace(enable_kv_cache_events=False), speculative_config=None, ) @@ -2147,6 +2153,48 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget(): assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2} +def test_recv_thread_partial_load_failure_reports_block_ids_without_hma(): + store = MagicMock() + # First key fails, second succeeds. + store.batch_get_into_multi_buffers.return_value = [-1, 0] + thread = _make_store_recving_thread(store) + + thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32)) + + assert thread.get_and_clear_block_ids_with_load_errors() == {0} + assert thread.get_and_clear_failed_requests() == set() + assert thread.get_and_clear_finished_requests() == {"req-a"} + + +def test_recv_thread_partial_load_failure_reports_request_with_hma(): + store = MagicMock() + store.batch_get_into_multi_buffers.return_value = [-1, 0] + thread = _make_store_recving_thread(store, is_hma_required=True) + + thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32)) + + assert thread.get_and_clear_block_ids_with_load_errors() == set() + assert thread.get_and_clear_failed_requests() == {"req-a"} + assert thread.get_and_clear_finished_requests() == {"req-a"} + + +def test_recv_thread_oversized_key_reports_request_with_hma(): + store = MagicMock() + thread = _make_store_recving_thread( + store, + tp_rank=2, + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL, + is_hma_required=True, + ) + + thread._handle_request(_make_load_req("req-a", [b"a0", b"a1", b"a2"], token_len=48)) + + assert store.batch_get_into_multi_buffers.call_count == 0 + assert thread.get_and_clear_block_ids_with_load_errors() == set() + assert thread.get_and_clear_failed_requests() == {"req-a"} + assert thread.get_and_clear_finished_requests() == {"req-a"} + + def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch): store = MagicMock() store.setup.return_value = 0 @@ -2172,6 +2220,39 @@ def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch): assert [db.block_size for db in store_worker.token_dbs] == [800, 800] +@pytest.mark.parametrize( + ("disable_hybrid_kv_cache_manager", "expected_is_hma_required"), + [ + (True, False), + (False, True), + ], +) +def test_worker_is_hma_required_from_kv_cache_groups( + monkeypatch, disable_hybrid_kv_cache_manager, expected_is_hma_required +): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch) + monkeypatch.setattr( + worker.MooncakeStoreConfig, + "load_from_config", + staticmethod(lambda: _make_config()), + ) + vllm_config = _make_vllm_config( + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager + ) + vllm_config.cache_config.block_size = 800 + vllm_config.cache_config.enable_prefix_caching = True + vllm_config.cache_config.prefix_match_unit = None + + store_worker = worker.MooncakeStoreWorker( + vllm_config, _make_qsa_hybrid_kv_cache_config() + ) + + assert store_worker._is_hma_required is expected_is_hma_required + + def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch): store = MagicMock() store.setup.return_value = 0 @@ -3345,6 +3426,7 @@ def _make_bare_worker( worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache worker._capacity_only = False worker.block_size = block_size + worker._is_hma_required = False worker.tp_rank = 0 worker.enable_kv_events = False worker.load_async = True diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 98633c233bb2..62e13fa43834 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -155,7 +155,6 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, - hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -186,7 +185,6 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, - hash_block_size=hash_block_size, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 4451f0566da1..5a3f462577c4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -30,6 +30,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats @@ -587,6 +588,13 @@ def get_finished( assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + """Get finished and failed sends/recvs from the worker.""" + assert self.connector_worker is not None + return self.connector_worker.get_transfer_results() + def get_block_ids_with_load_errors(self) -> set[int]: """Get block IDs whose remote KV load failed.""" assert self.connector_worker is not None @@ -1046,6 +1054,16 @@ def __init__( self.finished_recving_reqs: set[ReqId] = set() # Written from the receiver loop, drained from the worker thread. self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue() + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.transfer_groups + ) + ) + # Block IDs are only unique within a group; with HMA the scheduler + # tracks a single merged group, so failures are reported per request. + self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() self.xfer_stats = MooncakeKVConnectorStats() @@ -1866,6 +1884,29 @@ def get_finished(self) -> tuple[set[str] | None, set[str] | None]: return finished_sending_reqs or None, finished_recving_reqs or None + def get_transfer_results(self) -> KVConnectorTransferResults: + """ + Get transfers that completed on this specific worker, including + requests whose remote KV load failed. + + The scheduler process (via the MultiprocExecutor) will use this output + to track which workers are done. + """ + finished_sending_reqs, finished_recving_reqs = self.get_finished() + + failed_recving_reqs: set[ReqId] = set() + while True: + try: + failed_recving_reqs.add(self._failed_recv_reqs.get_nowait()) + except queue.Empty: + break + + return KVConnectorTransferResults( + finished_sending=set(finished_sending_reqs or ()), + finished_recving=set(finished_recving_reqs or ()), + failed_recving=failed_recving_reqs, + ) + def get_kv_connector_stats(self) -> KVConnectorStats | None: """Return transfer stats collected since the last call, or None if nothing has been recorded in this interval.""" @@ -1959,7 +2000,10 @@ def _handle_failed_recv( # request is waiting on a load, and reporting one here would trip # the scheduler's `assert req_id in self.requests`. continue - self._invalid_block_ids.put(invalid) + if self._is_hma_required: + self._failed_recv_reqs.put(pull_meta.d_req_id) + else: + self._invalid_block_ids.put(invalid) self.finished_recving_reqs.add(pull_meta.d_req_id) if failed: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 134d5055b6f2..eb4a7fafe92a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -27,6 +27,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, KVConnectorWorkerMetadata, SupportsHMA, ) @@ -359,6 +360,14 @@ def get_finished( assert isinstance(metadata, MooncakeStoreConnectorMetadata) return self.connector_worker.get_finished(finished_req_ids, metadata) + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + metadata = self._get_connector_metadata() + assert isinstance(metadata, MooncakeStoreConnectorMetadata) + return self.connector_worker.get_transfer_results(finished_req_ids, metadata) + def get_block_ids_with_load_errors(self) -> set[int]: assert self.connector_worker is not None return self.connector_worker.get_block_ids_with_load_errors() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 6f154ea991f0..0231ba73948e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -36,6 +36,9 @@ get_tensor_model_parallel_world_size, ) from vllm.distributed.kv_events import BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 ExternalCachedBlockPool, @@ -1265,6 +1268,7 @@ def __init__( record_operation: Callable[..., None] | None = None, request_queue: queue.Queue[Any] | None = None, group_participates: Sequence[bool] | None = None, + is_hma_required: bool = False, ): super().__init__( store, @@ -1284,6 +1288,10 @@ def __init__( # _invalid_block_ids can be access by both the Worker and RecvingThread self._invalid_block_ids_lock = threading.Lock() self._invalid_block_ids: set[int] = set() + # With HMA the scheduler tracks a single merged group while block IDs + # are only unique within a group, so failures are reported per request. + self._is_hma_required = is_hma_required + self._failed_requests: set[str] = set() self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes self.usable_disk_offload_buffer_budget_bytes = ( None @@ -1298,6 +1306,22 @@ def _add_load_error_block_ids(self, block_ids: list[int]) -> None: with self._invalid_block_ids_lock: self._invalid_block_ids.update(block_ids) + def set_failed_request(self, req_id: str): + with self.done_task_lock: + self._failed_requests.add(req_id) + + def get_and_clear_failed_requests(self) -> set[str]: + with self.done_task_lock: + failed = self._failed_requests.copy() + self._failed_requests.clear() + return failed + + def _report_load_error(self, req_id: str, block_ids: list[int]) -> None: + if self._is_hma_required: + self.set_failed_request(req_id) + else: + self._add_load_error_block_ids(block_ids) + def get_and_clear_block_ids_with_load_errors(self) -> set[int]: with self._invalid_block_ids_lock: invalid_block_ids = self._invalid_block_ids.copy() @@ -1392,7 +1416,7 @@ def _handle_request(self, req_meta: ReqMeta): # Mark every block: we skip the whole request, and the # tp_rank rotation means oversized_key isn't necessarily # the first block in the request's original order. - self._add_load_error_block_ids(block_id_list_c) + self._report_load_error(req_id, block_id_list_c) oversized_key_bytes = _estimate_disk_offload_staging_bytes( size_list_c[oversized_key_index] ) @@ -1465,8 +1489,8 @@ def _handle_request(self, req_meta: ReqMeta): num_failed_keys=len(failed), ) if failed: - self._add_load_error_block_ids( - [block_id for _, _, block_id in failed] + self._report_load_error( + req_id, [block_id for _, _, block_id in failed] ) logger.warning( "Failed to get %d Mooncake keys from sub-batch " @@ -1477,7 +1501,7 @@ def _handle_request(self, req_meta: ReqMeta): ) break except Exception as e: - self._add_load_error_block_ids(current_batch_block_ids) + self._report_load_error(req_id, current_batch_block_ids) self._record_operation( "load_get", load_get_start, @@ -1551,6 +1575,13 @@ def __init__( and not self.can_put ) self.cache_config = vllm_config.cache_config + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.transfer_groups + ) + ) self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) @@ -2096,6 +2127,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: group.kv_cache_spec.prefix_cacheable for group in self._kv_cache_groups ], + is_hma_required=self._is_hma_required, ) recv_thread.name = f"KVCacheStoreRecvingThread-{i}" recv_thread.start() @@ -2185,6 +2217,26 @@ def get_block_ids_with_load_errors(self) -> set[int]: block_ids |= recv_thread.get_and_clear_block_ids_with_load_errors() return block_ids + def get_transfer_results( + self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata + ) -> KVConnectorTransferResults: + """Get completed sends/recvs plus requests whose remote KV load failed.""" + done_sending, done_recving = self.get_finished(finished_req_ids, meta) + + if self._capacity_only: + return KVConnectorTransferResults(done_sending, done_recving) + + failed_recving: set[str] = set() + if self.load_async: + for recv_thread in self.kv_recv_threads: + failed_recving |= recv_thread.get_and_clear_failed_requests() + + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recving, + ) + def _record_kv_connector_operation( self, operation: str, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f1e11e9b1f3d..c1e956b9eca1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3027,11 +3027,6 @@ def _update_requests_with_invalid_blocks( For observability, it also accumulates the total number of tokens that will need to be recomputed across all affected requests. - A hybrid model spreads one token span over KV cache groups that disagree - on block size, so a block index maps to a different token position in - each group and there is no single longest valid prefix. Those requests - recompute in full instead. - Args: requests: The set of requests to scan for invalid blocks. invalid_block_ids: IDs of invalid blocks. @@ -3055,49 +3050,19 @@ def _update_requests_with_invalid_blocks( # these requests must be rescheduled, but only the first will recompute # it. This set tracks blocks already marked for recomputation. marked_invalid_block_ids: set[int] = set() - # Positions a group holds no KV for -- skipped sliding-window blocks, - # and every Mamba block but the aligned snapshot -- all point at this - # one block, shared by every request. It carries no request's data, so - # it can neither fail for a request nor be evicted on its behalf. - null_block_id = self.kv_cache_manager.block_pool.null_block.block_id for request in requests: is_affected = False marked_invalid_block = False req_id = request.request_id - # One block-id list per KV cache group. - req_block_ids_per_group = self.kv_cache_manager.get_block_ids(req_id) + # One block-id list per KV cache group; single-group layouts are + # enforced by _handle_invalid_blocks. + (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) # We iterate only over blocks that may contain externally computed # tokens req_num_computed_tokens = ( request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) ) - if len(req_block_ids_per_group) > 1: - # No single truncation point exists across groups, so the whole - # request is recomputed. Blocks are still marked, so a request - # sharing one of them is rescheduled too. - request_invalid_block_ids = { - block_id - for req_block_ids in req_block_ids_per_group - for block_id in req_block_ids - if block_id != null_block_id and block_id in invalid_block_ids - } - if request_invalid_block_ids: - marked_invalid_block_ids |= request_invalid_block_ids - total_affected_tokens += req_num_computed_tokens - request.num_computed_tokens = 0 - if evict_blocks: - for req_block_ids in req_block_ids_per_group: - blocks_to_evict.update( - block_id - for block_id in req_block_ids - if block_id != null_block_id - ) - affected_req_ids.add(req_id) - continue - - (req_block_ids,) = req_block_ids_per_group - req_num_computed_blocks = ( req_num_computed_tokens + self.block_size - 1 ) // self.block_size @@ -3175,6 +3140,17 @@ def _handle_invalid_blocks( Returns: Set of affected request IDs to skip in update_from_output main loop. """ + if len(self.kv_cache_config.kv_cache_groups) > 1: + # Block IDs are only unique within a group, so a flat set of + # invalid block IDs cannot be mapped back to requests. + raise RuntimeError( + "A KV connector reported block-level load failures " + "(invalid_block_ids) on a layout with multiple KV cache " + "groups, where block IDs are only unique within a group. " + "Connectors must report failed requests via " + "KVConnectorTransferResults.failed_recving instead." + ) + should_fail = not self.recompute_kv_load_failures # handle async KV loads (not cached yet, evict_blocks=False) From 975dca5bb5db302077674cfa9afe851ee700ad73 Mon Sep 17 00:00:00 2001 From: Arpan Tripathi <42506819+tripathiarpan20@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:03:57 +0100 Subject: [PATCH 324/774] [Perf][GDN] Scatter mixed speculative outputs into the caller buffer (#57140) Signed-off-by: Arpan Tripathi Signed-off-by: zjy0516 Co-authored-by: Xavier Lyu <9938360+xavierlyu@users.noreply.github.com> Co-authored-by: Bohdan Podziubanchuk <114604629+Danbog32@users.noreply.github.com> Co-authored-by: zjy0516 Co-authored-by: Codex --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index bbbb1c8bd5ad..a57c984db4c7 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1575,14 +1575,10 @@ def _forward_core( # 3. Merge core attention output if spec_sequence_masks is not None and core_attn_out_non_spec is not None: - merged_out = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, + core_attn_out.index_copy_(0, spec_token_indx, core_attn_out_spec.squeeze(0)) + core_attn_out.index_copy_( + 0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0) ) - merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec) - merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) else: From f8b5c11468f665c75968e3a7c12f16ca074f3a30 Mon Sep 17 00:00:00 2001 From: sashko-zakharchuk Date: Wed, 16 Sep 2026 15:24:48 +0200 Subject: [PATCH 325/774] [Bugfix][Kimi-K3] Do not classify a stateless first chunk as a decode (#51483) Signed-off-by: Oleksandr Zakharchuk Signed-off-by: zjy0516 Co-authored-by: zjy0516 --- tests/models/kimi_k3/test_kda_metadata.py | 88 ++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/kda_metadata.py | 27 ++++++- 2 files changed, 113 insertions(+), 2 deletions(-) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index c3f423804a90..fa810217b188 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -934,3 +934,91 @@ def test_aligned_block_table_matches_shared_gdn(): ) torch.testing.assert_close(actual, expected) + + +def _build_non_spec(batch, is_prefilling, full_cuda_graph=False): + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace( + is_prefilling=None + if is_prefilling is None + else torch.tensor(is_prefilling, dtype=torch.bool) + ) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=0, + full_cuda_graph=full_cuda_graph, + ) + return builder, common_attn_metadata, builder.build(0, common_attn_metadata) + + +def test_one_token_first_chunk_excludes_padding(): + """Neither padding requests nor padding tokens count as prefill work.""" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]), + BLOCK_SIZE, + DEVICE, + ).replace( + is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool), + num_actual_tokens=4, + ) + builder = _make_builder( + KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False + ) + actual = builder.build(0, common) + + assert actual.num_decodes == 1 + assert actual.num_prefills == 1 + assert actual.num_decode_tokens == 1 + assert actual.num_prefill_tokens == 1 + + +@pytest.mark.parametrize( + ("seq_len", "query_len", "is_prefilling", "num_prefills"), + [ + pytest.param(1, 1, True, 1, id="first-chunk"), + pytest.param(65, 1, True, 0, id="resumed-chunk"), + pytest.param(0, 0, True, 0, id="padding"), + pytest.param(1, 1, None, 0, id="missing-prefill-flag"), + ], +) +def test_one_token_chunk_classification( + seq_len, query_len, is_prefilling, num_prefills +): + """Only a real first chunk with a prefill flag needs state initialization.""" + _, _, actual = _build_non_spec( + BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]), + is_prefilling=None if is_prefilling is None else [False, is_prefilling], + ) + + assert actual.num_prefills == num_prefills + assert actual.num_decodes == 2 - num_prefills + assert actual.num_prefill_tokens == num_prefills + assert actual.num_decode_tokens == 1 + query_len - num_prefills + if num_prefills: + assert actual.has_initial_state is not None + assert actual.has_initial_state.tolist() == [True, False] + else: + assert actual.has_initial_state is None + + +def test_cudagraph_capture_batch_stays_decode_only(): + """Capture rows have no history, but must still select decode kernels.""" + batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.zeros(4, dtype=torch.bool)) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=0, + full_cuda_graph=True, + ) + actual = builder.build_for_cudagraph_capture(common_attn_metadata) + + assert actual.num_prefills == 0 + assert actual.num_decodes == 4 + assert actual.has_initial_state is None + staged = actual.non_spec_state_indices_tensor + assert staged is not None + assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr() + torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0]) diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 0b395b0c6c8f..d342adb6a688 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -427,10 +427,33 @@ def build( # type: ignore[override] spec_token_start = None non_spec_token_start = None if num_spec_decodes == 0: - # The runner orders ordinary decodes before prefills. + # V2 already excludes prefills from full decode graphs via has_prefill. + # Classify first chunks as prefills to mask recycled state; + # resumed one-token chunks can still use the decode kernels. + assert m.seq_lens_cpu_upper_bound is not None + query_lens_cpu = query_start_loc_cpu.diff() + no_prior_state = (query_lens_cpu > 0) & ( + m.seq_lens_cpu_upper_bound <= query_lens_cpu + ) + # Capture batches also have seq_len == query_len, but are not prefills. + if m.is_prefilling is not None: + no_prior_state &= m.is_prefilling + else: + no_prior_state = torch.zeros_like(no_prior_state) num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills(m, decode_threshold=1) + split_decodes_and_prefills( + m.replace(is_prefilling=no_prior_state), + decode_threshold=1, + treat_short_extends_as_decodes=False, + ) ) + # Exclude trailing padding from both prefill counts. + if num_prefills: + num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum()) + num_prefill_tokens = ( + int(query_start_loc_cpu[num_decodes + num_prefills]) + - num_decode_tokens + ) num_spec_decode_tokens = 0 spec_token_indx = None non_spec_token_indx = None From ec4a3a537068db40afbc9374a67da719c8c9b964 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:34:16 +0100 Subject: [PATCH 326/774] [CI] Bump Transformers version to 5.17.0 (#56108) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 --- requirements/test/cpu.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.in | 2 +- requirements/test/xpu.txt | 2 +- vllm/model_executor/models/laguna.py | 27 +++-- vllm/model_executor/models/pixtral.py | 45 +++++++- vllm/models/qwen4_exp/amd/model.py | 18 ++- vllm/models/qwen4_exp/config.py | 13 ++- vllm/models/qwen4_exp/nvidia/model.py | 18 ++- vllm/transformers_utils/config.py | 51 ++++++++- vllm/transformers_utils/configs/__init__.py | 4 - vllm/transformers_utils/configs/laguna.py | 120 -------------------- vllm/transformers_utils/configs/mellum.py | 7 -- 17 files changed, 148 insertions(+), 171 deletions(-) delete mode 100644 vllm/transformers_utils/configs/laguna.py delete mode 100644 vllm/transformers_utils/configs/mellum.py diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 052896c9f22b..996caab96a93 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1174,7 +1174,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.16.1 +transformers==5.17.0 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index b16c64ed346f..5e3faaebcc92 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -39,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.16.1 +transformers==5.17.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index db6e35918289..ea94558000ec 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1281,7 +1281,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.16.1 +transformers==5.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 35c4156c5d00..cebf3b782b26 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.16.1 +transformers==5.17.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 8397c9817110..c8c08fabe5f3 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.16.1 +transformers==5.17.0 schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.50.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index a6b9271989af..631c3952fd95 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1233,7 +1233,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.16.1 +transformers==5.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 3aab5bdb00e1..796e76eb441a 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -19,7 +19,7 @@ lm_eval[api]>=0.4.12 modelscope<1.38 runai-model-streamer[s3,gcs,azure]>=0.15.7 tensorizer==2.10.1 -transformers==5.16.1 +transformers==5.17.0 # --- Audio Processing --- av diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 2e2fbd88ede6..4e4fd628b656 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1081,7 +1081,7 @@ tqdm==4.70.0 # pqdm # sentence-transformers # transformers -transformers==5.16.1 +transformers==5.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py index 12046379fd77..7fe412eb230d 100644 --- a/vllm/model_executor/models/laguna.py +++ b/vllm/model_executor/models/laguna.py @@ -54,8 +54,19 @@ logger = init_logger(__name__) +def get_mlp_layer_types(config) -> list[str]: + """Return which layers are MoE, rejecting configs predating Transformers support.""" + mlp_layer_types = getattr(config, "mlp_layer_types", None) + if mlp_layer_types is None: + raise ValueError( + "Laguna requires the `LagunaConfig` from Transformers, which is " + "available from v5.17.0. Please update Transformers to run this model." + ) + return mlp_layer_types + + class LagunaMLP(nn.Module): - """Dense MLP for Laguna (used in mlp_only_layers).""" + """Dense MLP for Laguna.""" def __init__( self, @@ -206,7 +217,7 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - renormalize=config.norm_topk_prob, + renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", scoring_func="sigmoid", @@ -293,13 +304,12 @@ def __init__( else: self.sliding_window = None - # QKV projection (no bias for Laguna) self.qkv_proj = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, - bias=config.qkv_bias, + bias=config.attention_bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) @@ -506,14 +516,9 @@ def __init__( ), ) - # Check if this layer uses MoE or dense MLP (matches Qwen2/Qwen3 convention) - mlp_only_layers = ( - [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers - ) self.is_moe_layer = ( - (layer_idx not in mlp_only_layers) - and (config.num_experts > 0) - and ((layer_idx + 1) % config.decoder_sparse_step == 0) + get_mlp_layer_types(config)[layer_idx] == "sparse" + and config.num_experts > 0 ) if self.is_moe_layer: diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index 1485b90e39a2..c948fb5ad72e 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata import math from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, fields @@ -12,15 +13,12 @@ from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk from mistral_common.protocol.instruct.messages import UserMessage from mistral_common.protocol.instruct.request import ChatCompletionRequest +from packaging.version import Version from transformers import BatchFeature, PixtralVisionConfig from transformers.models.pixtral.image_processing_pixtral import ( _num_image_tokens as _get_pixtral_hf_num_image_tokens, ) -from transformers.models.pixtral.modeling_pixtral import ( - PixtralRotaryEmbedding, - apply_rotary_pos_emb, - position_ids_in_meshgrid, -) +from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions @@ -92,6 +90,39 @@ PATCH_MERGE = "patch_merge" +# Transformers 5.17 renamed Pixtral's rotary embedding and switched it to axial +# RoPE, which takes 2D (height, width) positions instead of flattened grid ids. +TRANSFORMERS_VERSION = importlib.metadata.version("transformers") +TRANSFORMERS_WITH_AXIAL_ROPE = Version(TRANSFORMERS_VERSION) >= Version("5.17.0.dev0") + +if TRANSFORMERS_WITH_AXIAL_ROPE: + from transformers.models.pixtral.modeling_pixtral import ( + PixtralVisionRotaryEmbedding, + ) +else: + from transformers.models.pixtral.modeling_pixtral import ( + PixtralRotaryEmbedding as PixtralVisionRotaryEmbedding, + ) + from transformers.models.pixtral.modeling_pixtral import ( + position_ids_in_meshgrid as flat_position_ids_in_meshgrid, + ) + + +def position_ids_in_meshgrid( + patch_embeds_list: list[torch.Tensor], + max_width: int, +) -> torch.Tensor: + if not TRANSFORMERS_WITH_AXIAL_ROPE: + return flat_position_ids_in_meshgrid(patch_embeds_list, max_width) + positions = [] + for patch in patch_embeds_list: + height, width = patch.shape[-2:] + h_ids, w_ids = torch.meshgrid( + torch.arange(height), torch.arange(width), indexing="ij" + ) + positions.append(torch.stack([h_ids.flatten(), w_ids.flatten()], dim=-1)) + return torch.cat(positions, dim=0) + def _make_packed_sequence_metadata( sequence_lengths: list[int], @@ -1467,7 +1498,9 @@ def __init__( self.dtype = next(self.parameters()).dtype self.device = next(self.parameters()).device - self.patch_positional_embedding = PixtralRotaryEmbedding(config, self.device) + self.patch_positional_embedding = PixtralVisionRotaryEmbedding(config).to( + self.device + ) def forward( self, diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index aeb3a14dadb8..b62179055e83 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -76,7 +76,7 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MambaSpec -from ..config import Qwen4ExpConfig +from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .ple_layer import Qwen4ExpPLELayer @@ -214,8 +214,11 @@ def __init__( prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, ) - elif layer_type == "full_attention": - use_qsa = getattr(config, "indexer_n_heads", None) is not None + elif layer_type in ATTENTION_LAYER_TYPES: + use_qsa = ( + layer_type == QSA_LAYER_TYPE + or getattr(config, "indexer_n_heads", None) is not None + ) if not use_qsa: self.self_attn = Qwen3NextAttention( config, @@ -311,7 +314,7 @@ def forward( if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) - elif self.layer_type == "full_attention": + elif self.layer_type in ATTENTION_LAYER_TYPES: attn_out = self.self_attn( hidden_states=block_input, positions=positions, @@ -402,8 +405,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self._qsa_layer_ids = frozenset( layer_idx for layer_idx, layer_type in enumerate(config.layer_types) - if layer_type == "full_attention" - and getattr(config, "indexer_n_heads", None) is not None + if layer_type == QSA_LAYER_TYPE + or ( + layer_type == "full_attention" + and getattr(config, "indexer_n_heads", None) is not None + ) ) self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py index a8eedab83f6b..341ddb4d3c8f 100644 --- a/vllm/models/qwen4_exp/config.py +++ b/vllm/models/qwen4_exp/config.py @@ -19,6 +19,12 @@ "indexer_compress_ratio", ) +# Transformers labels Qwen4Exp's sparse-attention layers `qwen_sparse_attention`. +# Older checkpoints label them `full_attention` and mark QSA by setting the +# indexer fields for the whole model. +QSA_LAYER_TYPE = "qwen_sparse_attention" +ATTENTION_LAYER_TYPES = ("full_attention", QSA_LAYER_TYPE) + class Qwen4ExpVisionConfig(Qwen3VLVisionConfig): model_type = "qwen4_exp" @@ -134,6 +140,11 @@ def _validate_ple_layer_ids(self) -> None: def _validate_qsa_config(self) -> None: configured = {name: getattr(self, name, None) for name in _QSA_CONFIG_FIELDS} if all(value is None for value in configured.values()): + if QSA_LAYER_TYPE in (self.layer_types or ()): + raise ValueError( + f"{QSA_LAYER_TYPE} layers require the QSA config fields: " + f"{list(_QSA_CONFIG_FIELDS)}" + ) return missing = [name for name, value in configured.items() if value is None] @@ -165,7 +176,7 @@ def _validate_qsa_config(self) -> None: @property def layers_block_type(self) -> list[str]: return [ - "attention" if layer_type == "full_attention" else layer_type + "attention" if layer_type in ATTENTION_LAYER_TYPES else layer_type for layer_type in self.layer_types ] diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 54ee6f02b4f5..386987e18fc8 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -75,7 +75,7 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MambaSpec -from ..config import Qwen4ExpConfig +from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .ple_layer import Qwen4ExpPLELayer @@ -215,8 +215,11 @@ def __init__( prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, ) - elif layer_type == "full_attention": - use_qsa = getattr(config, "indexer_n_heads", None) is not None + elif layer_type in ATTENTION_LAYER_TYPES: + use_qsa = ( + layer_type == QSA_LAYER_TYPE + or getattr(config, "indexer_n_heads", None) is not None + ) if not use_qsa: self.self_attn = Qwen3NextAttention( config, @@ -314,7 +317,7 @@ def forward( if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) - elif self.layer_type == "full_attention": + elif self.layer_type in ATTENTION_LAYER_TYPES: attn_out = self.self_attn( hidden_states=block_input, positions=positions, @@ -394,8 +397,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self._qsa_layer_ids = frozenset( layer_idx for layer_idx, layer_type in enumerate(config.layer_types) - if layer_type == "full_attention" - and getattr(config, "indexer_n_heads", None) is not None + if layer_type == QSA_LAYER_TYPE + or ( + layer_type == "full_attention" + and getattr(config, "indexer_n_heads", None) is not None + ) ) self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index d34e391a2626..565caeebc286 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -13,6 +13,7 @@ import huggingface_hub import torch +import transformers.configuration_utils as hf_configuration_utils from huggingface_hub import constants from packaging.version import Version from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE @@ -126,7 +127,6 @@ def __getitem__(self, key): RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) mlp_speculator="MLPSpeculatorConfig", medusa="MedusaConfig", - mellum="MellumConfig", midashenglm="MiDashengLMConfig", minimax_m3_vl="MiniMaxM3Config", minimax_m3_mtp="MiniMaxM3MTPConfig", @@ -151,7 +151,6 @@ def __getitem__(self, key): qwen3_5_text="Qwen3_5TextConfig", qwen3_5_moe="Qwen3_5MoeConfig", qwen3_5_moe_text="Qwen3_5MoeTextConfig", - laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, **{"deepseek_v41": "DeepseekV41Config"}, @@ -163,6 +162,12 @@ def __getitem__(self, key): _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"} +# Model types whose checkpoints store shared RoPE parameters alongside the +# per-layer-type dicts (e.g. Laguna's `original_max_position_embeddings`). +# Since transformers 5.17, `validate_rope` treats every top-level value of such +# a dict as a layer's parameters and raises on the shared ones. +_PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna"} + # Model types whose checkpoints declare `layer_types` entries that upstream # transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config # validation rejects them (e.g. GLM-5.2 `glm_moe_dsa` use @@ -236,6 +241,11 @@ def _patch_hf_transformers_validate_rope(): def patched_validate_rope(self, *args, **kwargs): ignore_keys_param = kwargs.pop("ignore_keys", None) original_ignore_keys = self.ignore_keys_at_rope_validation + if ignore_keys_param is not None: + logger.warning( + "validate_rope() was called with the legacy 'ignore_keys' " + "argument; use 'ignore_keys_at_rope_validation' instead" + ) self.ignore_keys_at_rope_validation = original_ignore_keys or ignore_keys_param result = _original_validate_rope(self, *args, **kwargs) return result @@ -244,6 +254,40 @@ def patched_validate_rope(self, *args, **kwargs): PretrainedConfig.validate_rope = patched_validate_rope +def _patch_hf_transformers_nested_rope_validation() -> None: + """Drop shared entries sitting alongside a nested ``rope_parameters`` dict. + + Transformers validates a dict with any layer-type key by iterating all of + its values, so a shared entry next to the per-layer dicts raises an + ``AttributeError``. The per-layer dicts already carry their own defaults by + the time validation runs, so the shared entries can be dropped. + """ + if hasattr(PretrainedConfig.validate_rope, "__vllm_nested_rope_patched__"): + return + + _original_validate_rope = PretrainedConfig.validate_rope + + @wraps(_original_validate_rope) + def patched_validate_rope(self, *args, **kwargs): + rope_parameters = getattr(self, "rope_parameters", None) + if isinstance(rope_parameters, dict): + layer_types = set(rope_parameters) & set( + hf_configuration_utils.ALLOWED_LAYER_TYPES + ) + if shared_keys := set(rope_parameters) - layer_types: + for key in shared_keys: + del rope_parameters[key] + logger.warning( + "Dropping %s from rope_parameters; entries shared across " + "layer types are not supported", + sorted(shared_keys), + ) + return _original_validate_rope(self, *args, **kwargs) + + patched_validate_rope.__vllm_nested_rope_patched__ = True # type: ignore[attr-defined] + PretrainedConfig.validate_rope = patched_validate_rope + + def _patch_hf_transformers_allowed_layer_types( extra_layer_types: tuple[str, ...], ) -> None: @@ -304,6 +348,9 @@ def parse( if model_type in _PATCH_HF_VALIDATE_ROPE: _patch_hf_transformers_validate_rope() + if model_type in _PATCH_HF_NESTED_ROPE_VALIDATION: + _patch_hf_transformers_nested_rope_validation() + if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type): _patch_hf_transformers_allowed_layer_types(extra_layer_types) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 971b6ffcb2b5..3a6874820ea1 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -55,10 +55,8 @@ # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the # `FalconConfig` class from the official HuggingFace transformers library. "RWConfig": "vllm.transformers_utils.configs.falcon", - "LagunaConfig": "vllm.transformers_utils.configs.laguna", "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe", "MedusaConfig": "vllm.transformers_utils.configs.medusa", - "MellumConfig": "vllm.transformers_utils.configs.mellum", "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm", "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3", "MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3", @@ -151,10 +149,8 @@ "HyperCLOVAXConfig", "IsaacConfig", "RWConfig", - "LagunaConfig", "Lfm2MoeConfig", "MedusaConfig", - "MellumConfig", "MiDashengLMConfig", "MiniMaxM3Config", "MiniMaxM3MTPConfig", diff --git a/vllm/transformers_utils/configs/laguna.py b/vllm/transformers_utils/configs/laguna.py deleted file mode 100644 index 2702d3af5aa1..000000000000 --- a/vllm/transformers_utils/configs/laguna.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from transformers.configuration_utils import PretrainedConfig - - -class LagunaConfig(PretrainedConfig): - model_type = "laguna" - keys_to_ignore_at_inference = ["past_key_values"] - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - def __init__( - self, - vocab_size: int = 100352, - hidden_size: int = 2048, - intermediate_size: int = 8192, - num_hidden_layers: int = 40, - num_attention_heads: int = 48, - num_key_value_heads: int = 8, - head_dim: int = 128, - qkv_bias: bool = False, - attention_bias: bool = False, - gating: bool | str = True, - hidden_act: str = "silu", - max_position_embeddings: int = 131072, - initializer_range: float = 0.02, - rms_norm_eps: float = 1e-6, - use_cache: bool = True, - tie_word_embeddings: bool = False, - rope_theta: float = 500000.0, - rope_scaling: dict | None = None, - rope_parameters: dict | None = None, - partial_rotary_factor: float = 1.0, - attention_dropout: float = 0.0, - sliding_window: int | None = None, - layer_types: list[str] | None = None, - swa_attention_sink_enabled: bool = False, - swa_rope_parameters: dict | None = None, - num_attention_heads_per_layer: list[int] | None = None, - num_experts: int = 256, - num_experts_per_tok: int = 8, - moe_intermediate_size: int = 512, - shared_expert_intermediate_size: int = 512, - norm_topk_prob: bool = True, - decoder_sparse_step: int = 1, - mlp_only_layers: list[int] | None = None, - router_aux_loss_coef: float = 0.001, - output_router_logits: bool = False, - moe_routed_scaling_factor: float = 1.0, - moe_apply_router_weight_on_input: bool = False, - **kwargs, - ): - if mlp_only_layers is None: - mlp_only_layers = [0] - - # Accept either v4-style (rope_theta + rope_scaling) or v5-style - # (rope_parameters). Translate v5 → v4 so downstream code has one path. - if rope_parameters is not None: - rp = dict(rope_parameters) - rope_theta = float(rp.pop("rope_theta", rope_theta)) - rt = rp.pop("rope_type", None) - if rt is not None and rt != "default": - rope_scaling = {"rope_type": rt, **rp} - elif rp and rope_scaling is None: - rope_scaling = {"rope_type": "default", **rp} - - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.qkv_bias = qkv_bias - self.attention_bias = attention_bias - self.gating = gating - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.partial_rotary_factor = partial_rotary_factor - self.attention_dropout = attention_dropout - self.sliding_window = sliding_window - self.layer_types = layer_types - self.swa_attention_sink_enabled = swa_attention_sink_enabled - self.swa_rope_parameters = swa_rope_parameters - self.num_attention_heads_per_layer = num_attention_heads_per_layer - self.num_experts = num_experts - self.num_experts_per_tok = num_experts_per_tok - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.norm_topk_prob = norm_topk_prob - self.decoder_sparse_step = decoder_sparse_step - self.mlp_only_layers = mlp_only_layers - self.router_aux_loss_coef = router_aux_loss_coef - self.output_router_logits = output_router_logits - self.moe_routed_scaling_factor = moe_routed_scaling_factor - self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input - - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["LagunaConfig"] diff --git a/vllm/transformers_utils/configs/mellum.py b/vllm/transformers_utils/configs/mellum.py deleted file mode 100644 index 2bed53394b25..000000000000 --- a/vllm/transformers_utils/configs/mellum.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from transformers import Qwen3MoeConfig - - -class MellumConfig(Qwen3MoeConfig): - model_type = "mellum" From c81eb3e69c40ab17cb3c4cd8350e4a60b1be70d4 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Wed, 16 Sep 2026 22:15:42 +0800 Subject: [PATCH 327/774] [LoRA] Support modules_to_save for sequence classification (#53555) Signed-off-by: linitra24 Signed-off-by: linitra24 Co-authored-by: Jee Jee Li --- .buildkite/test_areas/lora.yaml | 2 + docs/features/lora.md | 12 ++ .../classification_with_lora_offline.py | 32 ++++ .../serve/lora/test_lora_adapters.py | 2 +- tests/lora/conftest.py | 24 +++ tests/lora/test_peft_helper.py | 2 +- tests/lora/test_sequence_classification.py | 167 ++++++++++++++++++ vllm/lora/layers/__init__.py | 2 + vllm/lora/layers/classifier.py | 97 ++++++++++ vllm/lora/lora_model.py | 41 ++++- vllm/lora/lora_weights.py | 15 ++ vllm/lora/model_manager.py | 156 +++++++++++++++- vllm/lora/peft_helper.py | 10 +- vllm/lora/punica_wrapper/punica_base.py | 28 +++ vllm/lora/punica_wrapper/punica_cpu.py | 20 +++ vllm/lora/punica_wrapper/punica_gpu.py | 27 +++ vllm/lora/punica_wrapper/punica_xpu.py | 20 +++ vllm/lora/utils.py | 27 ++- vllm/lora/worker_manager.py | 4 + vllm/model_executor/layers/pooler/special.py | 14 ++ 20 files changed, 681 insertions(+), 21 deletions(-) create mode 100644 examples/pooling/classify/classification_with_lora_offline.py create mode 100644 tests/lora/test_sequence_classification.py create mode 100644 vllm/lora/layers/classifier.py diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index d63fb567eaf1..93fb5f5b81f0 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -54,6 +54,7 @@ steps: - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp mirror: amd: label: ":amd: (MI300) LoRA TP (Distributed) %N" @@ -79,3 +80,4 @@ steps: - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp diff --git a/docs/features/lora.md b/docs/features/lora.md index 839624542ec5..78004722a935 100644 --- a/docs/features/lora.md +++ b/docs/features/lora.md @@ -422,6 +422,18 @@ vllm serve ibm-granite/granite-speech-3.3-2b \ Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions. +## Sequence-Classification LoRA Adapters + +vLLM supports PEFT sequence-classification adapters that save a complete, single-layer linear classification head through `modules_to_save`. The saved module must be named `score` or `classifier`. + +See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter. + +This support has the following limitations: + +- All adapters in one engine must have the same `num_labels` and hidden size as the base classification head. An incompatible adapter is rejected when it is loaded. +- A classification head stored as float32 is converted to the runtime head dtype when it is loaded. +- Token-classification adapters are not supported by this feature. + ## Using Tips ### Configuring `max_lora_rank` diff --git a/examples/pooling/classify/classification_with_lora_offline.py b/examples/pooling/classify/classification_with_lora_offline.py new file mode 100644 index 000000000000..a44066a653ef --- /dev/null +++ b/examples/pooling/classify/classification_with_lora_offline.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Example of offline classification with a LoRA adapter.""" + +from vllm import LLM +from vllm.lora.request import LoRARequest +from vllm.transformers_utils.repo_utils import hf_api + + +def main(): + lora_path = hf_api().snapshot_download( + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora" + ) + llm = LLM( + model="Skywork/Skywork-Reward-V2-Qwen3-0.6B", + runner="pooling", + enable_lora=True, + max_lora_rank=8, + max_model_len=512, + ) + + prompt = "Which response is more helpful and correct?" + (output,) = llm.classify( + prompt, + lora_request=LoRARequest("reward-lora", 1, lora_path), + ) + print(f"Prompt: {prompt!r}") + print(f"Class Probabilities: {output.outputs.probs}") + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index a22f0b38991b..8a860dc6251a 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -26,7 +26,7 @@ ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, - "only supports modules_to_save being None", + "Unsupported modules_to_save", ), ] diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 2c64fb224534..8635779cff44 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -270,6 +270,30 @@ def qwen3_woofing_lora_files(): return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") +@pytest.fixture(scope="session") +def qwen3_guard_star_trek_lora_files(): + return hf_api().snapshot_download( + repo_id="geoffmunn/Qwen3Guard-StarTrek-Classification-0.6B", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + +@pytest.fixture(scope="session") +def qwen3_guard_new_zealand_lora_files(): + return hf_api().snapshot_download( + repo_id="geoffmunn/Qwen3Guard-NewZealand-Classification-0.6B", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + +@pytest.fixture(scope="session") +def skywork_qwen3_reward_lora_files(): + return hf_api().snapshot_download( + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + @pytest.fixture(scope="session") def deepseekv2_lora_files(): return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index 23a8aad1a58c..9fb52ea7275c 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -20,7 +20,7 @@ ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, - "only supports modules_to_save being None", + "Unsupported modules_to_save", ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py new file mode 100644 index 000000000000..86721bfb6ef0 --- /dev/null +++ b/tests/lora/test_sequence_classification.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.utils import multi_gpu_test +from vllm import LLM, PoolingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.lora.request import LoRARequest + +MODEL_NAME = "Qwen/Qwen3-0.6B" +NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B" +PROMPTS = { + "star_trek": "Does warp drive appear in Star Trek?", + "new_zealand": "Wellington is the capital of New Zealand.", +} +EXPECTED_OUT = { + "star_trek": [13.6640625, -13.6640625], + "new_zealand": [-4.80078125, 8.6171875], + "native": [-1.52734375], +} + + +def _classify_logits( + llm: LLM, + prompts: list[str], + requests: LoRARequest | list[LoRARequest], +) -> torch.Tensor: + params = [PoolingParams(task="classify", use_activation=False) for _ in prompts] + outputs = llm.classify( + prompts, + pooling_params=params, + lora_request=requests, + use_tqdm=False, + ) + return torch.tensor([output.outputs.probs for output in outputs]) + + +def test_converted_model_with_modules_to_save( + qwen3_guard_star_trek_lora_files: str, +) -> None: + prompt = PROMPTS["star_trek"] + llm = LLM( + model=MODEL_NAME, + runner="pooling", + convert="classify", + hf_overrides={"num_labels": 2}, + dtype="float16", + enable_lora=True, + max_lora_rank=16, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + ) + actual = _classify_logits( + llm, + [prompt], + LoRARequest("star-trek", 1, qwen3_guard_star_trek_lora_files), + )[0] + + torch.testing.assert_close( + actual, + torch.tensor(EXPECTED_OUT["star_trek"]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def test_native_classification_model_with_modules_to_save( + skywork_qwen3_reward_lora_files: str, +) -> None: + prompt = "Which response is more helpful and correct?" + llm = LLM( + model=NATIVE_MODEL_NAME, + runner="pooling", + dtype="float16", + enable_lora=True, + max_lora_rank=8, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + ) + actual = _classify_logits( + llm, + [prompt], + LoRARequest("native", 1, skywork_qwen3_reward_lora_files), + )[0] + + torch.testing.assert_close( + actual, + torch.tensor(EXPECTED_OUT["native"]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def _test_batched_loras( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, + tp_size: int, +) -> None: + names = ["star_trek", "new_zealand"] + prompts = [PROMPTS[name] for name in names] + adapter_paths = [ + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + ] + requests = [ + LoRARequest(name, index, adapter_path) + for index, (name, adapter_path) in enumerate(zip(names, adapter_paths), start=1) + ] + + llm = LLM( + model=MODEL_NAME, + runner="pooling", + convert="classify", + hf_overrides={"num_labels": 2}, + dtype="float16", + enable_lora=True, + max_loras=2, + max_cpu_loras=2, + max_lora_rank=16, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + tensor_parallel_size=tp_size, + ) + actual = _classify_logits(llm, prompts, requests) + + torch.testing.assert_close( + actual, + torch.tensor([EXPECTED_OUT[name] for name in names]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def test_batched_loras( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, +) -> None: + _test_batched_loras( + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + tp_size=1, + ) + + +@multi_gpu_test(num_gpus=2) +def test_batched_loras_tp( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, +) -> None: + _test_batched_loras( + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + tp_size=2, + ) diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py index e0c5f202bd01..2f88640e5fd1 100644 --- a/vllm/lora/layers/__init__.py +++ b/vllm/lora/layers/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.lora.layers.base import BaseLayerWithLoRA +from vllm.lora.layers.classifier import ClassificationHeadWithLoRA from vllm.lora.layers.column_parallel_linear import ( ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA, @@ -42,4 +43,5 @@ "LoRAMappingType", "FusedMoEWithLoRA", "FusedMoE3DWithLoRA", + "ClassificationHeadWithLoRA", ] diff --git a/vllm/lora/layers/classifier.py b/vllm/lora/layers/classifier.py new file mode 100644 index 000000000000..aa42023763e5 --- /dev/null +++ b/vllm/lora/layers/classifier.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from transformers import PretrainedConfig + +from vllm.config.lora import LoRAConfig +from vllm.platforms import current_platform + +from .replicated_linear import ReplicatedLinearWithLoRA + + +class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA): + def create_lora_weights( + self, + max_loras: int, + lora_config: LoRAConfig, + model_config: PretrainedConfig | None = None, + ) -> None: + # Preserve ordinary LoRA A/B support for classification heads. + super().create_lora_weights(max_loras, lora_config, model_config) + self.full_weight_stacked = torch.zeros( + max_loras, + 1, + self.output_size, + self.input_size, + dtype=lora_config.lora_dtype, + device=self.device, + ) + self.full_bias_stacked = torch.zeros( + max_loras, + self.output_size, + dtype=self.base_layer.params_dtype, + device=self.device, + ) + self.full_module_enabled = torch.zeros( + max_loras, dtype=torch.bool, device=self.device + ) + + def reset_module_to_save(self, index: int) -> None: + self.full_weight_stacked[index].zero_() + self.full_bias_stacked[index].zero_() + self.full_module_enabled[index] = False + + def set_module_to_save( + self, + index: int, + weight: torch.Tensor, + bias: torch.Tensor | None, + ) -> None: + self.reset_module_to_save(index) + self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True) + if bias is not None: + self.full_bias_stacked[index].copy_(bias, non_blocking=True) + self.full_module_enabled[index] = True + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + # TODO base_result maybe don't need compute truly. + base_result = super().forward(input_) + output = base_result[0] if isinstance(base_result, tuple) else base_result + + full_output = self.punica_wrapper.apply_lora_full_linear( + output, + input_.to(self.full_weight_stacked.dtype), + self.full_weight_stacked, + self.full_bias_stacked, + self.full_module_enabled, + ) + + if full_output is not None: + output = full_output + return output + + def _apply_lora_to_output( + self, x: torch.Tensor, output: torch.Tensor + ) -> torch.Tensor: + original_shape = output.shape if output.ndim == 3 else None + if x.ndim == 3 and output.ndim == 3: + output = output.flatten(0, 1) + x = x.flatten(0, 1) + + # Classification outputs are request-level, so use prompt-based LoRA routing. + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits( + output, + x.to(self.lora_a_stacked[0].dtype), + self.lora_a_stacked[0], + self.lora_b_stacked[0], + 1.0, + ) + if not current_platform.can_update_inplace(): + output = lora_output + + if original_shape is not None: + output = output.reshape(original_shape) + + return output diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index f094f11f09e8..609f14f85e72 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -8,7 +8,7 @@ import torch from vllm.logger import init_logger -from vllm.lora.lora_weights import LoRALayerWeights +from vllm.lora.lora_weights import LoRAFullModuleWeights, LoRALayerWeights from vllm.lora.peft_helper import PEFTHelper from vllm.lora.utils import ( get_lora_id, @@ -66,6 +66,8 @@ def __init__( rank: int, loras: dict[str, LoRALayerWeights], is_3d_lora_weight: bool = False, + *, + modules_to_save: dict[str, LoRAFullModuleWeights] | None = None, ) -> None: """ Args: @@ -85,6 +87,7 @@ def __init__( ) self.rank = rank self.loras: dict[str, LoRALayerWeights] = loras + self.modules_to_save = modules_to_save or {} self.is_3d_lora_weight = is_3d_lora_weight def clone(self, lora_model_id: int) -> "LoRAModel": @@ -96,12 +99,16 @@ def clone(self, lora_model_id: int) -> "LoRAModel": rank=self.rank, loras=self.loras.copy(), is_3d_lora_weight=self.is_3d_lora_weight, + modules_to_save=self.modules_to_save.copy(), ) def get_lora(self, module_name: str) -> LoRALayerWeights | None: """Get LoRA for a given module by name""" return self.loras.get(module_name, None) + def get_module_to_save(self, module_name: str) -> LoRAFullModuleWeights | None: + return self.modules_to_save.get(module_name) + def check_lora_name(self, lora_name: str) -> bool: return lora_name in self.loras @@ -127,6 +134,9 @@ def from_lora_tensors( ) -> "LoRAModel": """Create a LoRAModel from a dictionary of tensors.""" pin_memory = str(device) == "cpu" and PIN_MEMORY + + modules_to_save_names = peft_helper.modules_to_save or [] + full_parameters: dict[str, dict[str, torch.Tensor]] = {} loras: dict[str, LoRALayerWeights] = {} for tensor_name, tensor in tensors.items(): if is_base_embedding_weights(tensor_name): @@ -137,6 +147,12 @@ def from_lora_tensors( module_name, is_lora_a = parse_fine_tuned_lora_name( tensor_name, weights_mapper ) + if module_name in modules_to_save_names: + full_parameters.setdefault(module_name, {})[ + tensor_name.split(".")[-1] + ] = tensor.to(device=device) + continue + if module_name not in loras: loras[module_name] = LoRALayerWeights.from_config( module_name, peft_helper @@ -161,7 +177,23 @@ def from_lora_tensors( if pin_memory: loras[module_name].lora_b = loras[module_name].lora_b.pin_memory() - return cls(lora_model_id, peft_helper.r, loras) + modules_to_save = {} + for module_name, parameters in full_parameters.items(): + weight = parameters.get("weight") + if weight is None: + raise ValueError(f"Full module {module_name!r} is missing weight.") + modules_to_save[module_name] = LoRAFullModuleWeights( + module_name=module_name, + weight=weight, + bias=parameters.get("bias"), + ) + if len(modules_to_save) > 1: + raise ValueError( + "Only one full classification module is supported per " + f"adapter, received {sorted(modules_to_save)}." + ) + + return cls(lora_model_id, peft_helper.r, loras, modules_to_save=modules_to_save) @classmethod def from_local_checkpoint( @@ -223,6 +255,7 @@ def check_unexpected_modules(modules: dict): ): continue module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper) + base_name = module_name.rsplit(".", 1)[-1] # Case for expert lora weights if ".experts" in module_name: expert_idx = module_name.find(".experts") @@ -230,7 +263,9 @@ def check_unexpected_modules(modules: dict): if expert_suffix not in expected_lora_modules: unexpected_modules.append(module_name) - elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules: + elif base_name not in expected_lora_modules and base_name not in ( + peft_helper.modules_to_save or () + ): unexpected_modules.append(module_name) if unexpected_modules: diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index bb241db77e13..d9b200f0bba5 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence as GenericSequence +from dataclasses import dataclass import torch import torch.types @@ -281,3 +282,17 @@ def output_dim(self) -> int: @property def is_packed(self) -> bool: return True + + +@dataclass +class LoRAFullModuleWeights: + """LoRA weights for classification layers.""" + + module_name: str + weight: torch.Tensor + bias: torch.Tensor | None = None + + def pin_memory(self) -> None: + self.weight = self.weight.pin_memory() + if self.bias is not None: + self.bias = self.bias.pin_memory() diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 7791e1b762e1..2f3b069f0774 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -13,16 +13,22 @@ from vllm.logger import init_logger from vllm.lora.layers import ( BaseLayerWithLoRA, + ClassificationHeadWithLoRA, FusedMoE3DWithLoRA, FusedMoEWithLoRA, LoRAMapping, LoRAMappingType, ) from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec -from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.lora.lora_weights import ( + LoRAFullModuleWeights, + LoRALayerWeights, + PackedLoRALayerWeights, +) from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper from vllm.lora.utils import ( from_layer, + from_layer_classification, from_layer_logits_processor, get_supported_lora_modules, is_in_target_modules, @@ -114,6 +120,13 @@ def __init__( self.vocab_size = vocab_size self.is_pooling_model = is_pooling_model(self.model) + self._pooling_task = ( + vllm_config.model_config.get_pooling_task( + tuple(self.model.pooler.get_supported_tasks()) + ) + if self.is_pooling_model + else None + ) self.packed_modules: dict[str, list[str]] = {} self.modules: dict[str, BaseLayerWithLoRA] = {} self._last_mapping: LoRAMapping | None = None @@ -121,6 +134,13 @@ def __init__( is_moe = is_moe_model(self.model) self._is_moe = is_moe + self.supported_modules_to_save = ( + {"score", "classifier"} + if self.is_pooling_model + and vllm_config.model_config.score_type == "cross-encoder" + else set() + ) + # When the engine is started with enable_mixed_moe_lora_format=True # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of # the model's 3D flag, so 2D and 3D adapters can coexist. @@ -147,6 +167,20 @@ def __init__( self._init_punica_wrapper(max_num_batched_tokens, vllm_config) self._create_lora_modules() + self._classification_head: tuple[str, ClassificationHeadWithLoRA] | None = None + if self.is_pooling_model: + classification_heads = [ + (module_name, module) + for module_name, module in self.modules.items() + if isinstance(module, ClassificationHeadWithLoRA) + ] + if classification_heads: + assert len(classification_heads) == 1, ( + "Expected 1 classification head, but found " + f"{len(classification_heads)}." + ) + self._classification_head = classification_heads[0] + self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec() self.model.lora_manager = self @@ -338,6 +372,15 @@ def activate_adapter( self.lora_index_to_id[index] = lora_model.id num_lora_weights_applied = 0 for module_name, module in self.modules.items(): + if isinstance(module, ClassificationHeadWithLoRA): + full_module = self._get_module_to_save_weights(lora_model, module_name) + if full_module is None: + module.reset_module_to_save(index) + else: + module.set_module_to_save( + index, full_module.weight, full_module.bias + ) + module_lora = self._get_lora_layer_weights(lora_model, module_name) if not module_lora: module.reset_lora(index) @@ -427,6 +470,8 @@ def _parent_module(module_name: str) -> str: if isinstance(module, PPMissingLayer): continue + is_classifier_head = module_name in self.supported_modules_to_save + if self.lora_config.target_modules is None: if not is_supported_lora_module( module_name, @@ -434,7 +479,8 @@ def _parent_module(module_name: str) -> str: self.supported_lora_modules, ): continue - elif not self._match_target_modules(module_name): + elif not self._match_target_modules(module_name) and not is_classifier_head: + # Full classification heads bypass the A/B LoRA target filter. continue punica_wrapper = self._get_punica_wrapper(module_name) @@ -486,17 +532,25 @@ def _parent_module(module_name: str) -> str: # LoRA weights of w1 and w3 have already been fused on disk. packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"] - new_module = replace_submodule( - self.model, - module_name, - from_layer( + if is_classifier_head: + new_module = from_layer_classification( + module, + self.lora_slots, + self.lora_config, + self.model.config, + ) + else: + new_module = from_layer( module, self.lora_slots, self.lora_config, packed_moduled_lst, self.model.config, - ), - ) + ) + new_module = replace_submodule(self.model, module_name, new_module) + + if is_classifier_head: + self.model.pooler.replace_classifier(new_module) if isinstance(new_module, BaseLayerWithLoRA): wrapped_by_id[id(module)] = new_module wrapped_by_id[id(new_module)] = new_module @@ -746,6 +800,12 @@ def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None: if not self.supports_mm: return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] + if module_name in self.supported_modules_to_save: + lm_wrapper = self.punica_wrapper_mapping.get( + self.mm_mapping.language_model[0] + ) + return lm_wrapper + # For multimodal model # NOTE Sort by prefix length (descending) to match the longest prefix first # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.' @@ -871,6 +931,10 @@ def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: lora.lora_a = lora.lora_a.pin_memory() lora.lora_b = lora.lora_b.pin_memory() + if lora_model.modules_to_save: + for full_module in lora_model.modules_to_save.values(): + full_module.pin_memory() + def _stack_moe_lora_weights( self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str ): @@ -1164,6 +1228,82 @@ def _get_lora_layer_weights( ) return lora_model.get_lora(org_module_name) + def _get_module_to_save_weights( + self, lora_model: LoRAModel, module_name: str + ) -> LoRAFullModuleWeights | None: + weights = lora_model.get_module_to_save(module_name) + if weights is not None: + return weights + + unprefixed_module_name = module_name.removeprefix("model.") + weights = lora_model.get_module_to_save(unprefixed_module_name) + if weights is not None and unprefixed_module_name != module_name: + logger.info_once( + "For the pool model, successfully loaded the full module weights " + "after removing the prefix 'model.'." + ) + return weights + + def _validate_modules_to_save(self, lora_model: LoRAModel) -> None: + if not lora_model.modules_to_save: + return + saved_module_name = next(iter(lora_model.modules_to_save)) + if self._classification_head is None: + raise ValueError( + f"Cannot load full module {saved_module_name!r}: the model does " + "not expose a unique ClassificationHeadWithLoRA." + ) + + module_name, wrapper = self._classification_head + full_module = self._get_module_to_save_weights(lora_model, module_name) + if full_module is None: + raise ValueError( + f"Full module {saved_module_name!r} does not match the model's " + f"classification head {module_name!r}." + ) + + expected_weight_shape = (wrapper.output_size, wrapper.input_size) + received_weight_shape = tuple(full_module.weight.shape) + if received_weight_shape != expected_weight_shape: + raise ValueError( + f"Full module {saved_module_name!r} for {module_name!r} has " + "an incompatible weight shape: expected " + f"{expected_weight_shape}, received {received_weight_shape}." + ) + + if full_module.bias is None: + return + expected_bias_shape = (wrapper.output_size,) + received_bias_shape = tuple(full_module.bias.shape) + if received_bias_shape != expected_bias_shape: + raise ValueError( + f"Full module {saved_module_name!r} for {module_name!r} has " + "an incompatible bias shape: expected " + f"{expected_bias_shape}, received {received_bias_shape}." + ) + + def _validate_token_classification_lora(self, lora_model: LoRAModel) -> None: + if self._pooling_task != "token_classify": + return + if self._classification_head is None: + return + + module_name, _ = self._classification_head + has_full_module = ( + self._get_module_to_save_weights(lora_model, module_name) is not None + ) + has_lora_weights = ( + self._get_lora_layer_weights(lora_model, module_name) is not None + ) + if not has_full_module and not has_lora_weights: + return + + raise ValueError( + f"LoRA adapter {lora_model.id} contains weights for classification " + f"head {module_name!r}, but token_classify only supports LoRA on " + "the model backbone." + ) + def deactivate_adapter(self, adapter_id: int) -> bool: if adapter_id not in self._active_adapters: return False diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index a0bb8bf3c283..8b40af03456e 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -45,7 +45,15 @@ def _validate_features(self) -> list[str]: """ error_msg = [] if self.modules_to_save: - error_msg.append("vLLM only supports modules_to_save being None.") + unsupported_modules = [ + m for m in self.modules_to_save if m not in ["classifier", "score"] + ] + if unsupported_modules: + error_msg.append( + "vLLM only supports modules_to_save being either None " + 'or ["classifier", "score"] for classification models. ' + f"Unsupported modules_to_save: {unsupported_modules}" + ) if self.use_dora: error_msg.append("vLLM does not yet support DoRA.") return error_msg diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py index f5c7ca9cb73b..7ab309515c1a 100644 --- a/vllm/lora/punica_wrapper/punica_base.py +++ b/vllm/lora/punica_wrapper/punica_base.py @@ -120,6 +120,18 @@ def add_lora_logits( """ raise NotImplementedError + @abstractmethod + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> torch.Tensor | None: + """Apply request-routed full linear weights to selected rows.""" + raise NotImplementedError + class PunicaWrapperBase(PunicaWrapperABC): """ @@ -217,6 +229,22 @@ def _update_prefill_metadata(self, token_lora_tensor: torch.Tensor) -> None: self.token_nums = token_nums self.no_lora = no_lora + def _select_full_linear_output( + self, + y: torch.Tensor, + adapter_y: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> torch.Tensor: + indices = self.sampler_indices + assert indices.size(0) == y.size(0), ( + "Full linear rows do not match LoRA request mapping" + ) + safe_indices = indices.clamp_min(0) + use_full = (indices >= 0) & module_enabled[safe_indices] + adapter_y = adapter_y.to(y.dtype) + bias_stacked[safe_indices].to(y.dtype) + return torch.where(use_full.unsqueeze(-1), adapter_y, y) + @property def prefill_metadata( self, diff --git a/vllm/lora/punica_wrapper/punica_cpu.py b/vllm/lora/punica_wrapper/punica_cpu.py index 1a700d9bf1f0..d462ce4476d8 100644 --- a/vllm/lora/punica_wrapper/punica_cpu.py +++ b/vllm/lora/punica_wrapper/punica_cpu.py @@ -349,3 +349,23 @@ def add_lora_logits( bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale) bgmv_expand(buffer, lora_b_stacked, y, self.sampler_indices, add_inputs=True) y = y.view_as(y_org) + + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + indices = self.sampler_indices + adapter_y = torch.zeros( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 2f746f967e37..7deb97e5134e 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -327,6 +327,33 @@ def add_lora_logits( ) y = y.view_as(y_org) + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + adapter_y = torch.empty( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + lora_shrink( + x, + [weight_stacked], + adapter_y.unsqueeze(0), + *self.prompt_mapping_meta.meta_args( + x.size(0), self.lora_config.specialize_active_lora + ), + 1.0, + ) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) + def moe_lora_align_block_size( self, topk_ids: torch.Tensor, diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py index 7f040732ee4c..3d1cf29683fc 100755 --- a/vllm/lora/punica_wrapper/punica_xpu.py +++ b/vllm/lora/punica_wrapper/punica_xpu.py @@ -333,6 +333,26 @@ def add_lora_logits( bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True) y = y.view_as(y_org) + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + indices = self.sampler_indices + adapter_y = torch.zeros( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) + def moe_lora_align_block_size( self, topk_ids: torch.Tensor, diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index f2a9dc902fce..fba78eddcf25 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -15,6 +15,7 @@ # being imported for _all_lora_classes below from vllm.lora.layers import ( BaseLayerWithLoRA, + ClassificationHeadWithLoRA, ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA, FusedMoE3DWithLoRA, @@ -143,6 +144,17 @@ def from_layer_logits_processor( return ret +def from_layer_classification( + layer: nn.Module, + max_loras: int, + lora_config: LoRAConfig, + model_config: PretrainedConfig | None = None, +) -> ClassificationHeadWithLoRA: + instance_layer = ClassificationHeadWithLoRA(layer) + instance_layer.create_lora_weights(max_loras, lora_config, model_config) + return instance_layer + + def replace_submodule( model: nn.Module, module_name: str, new_module: nn.Module ) -> nn.Module: @@ -193,13 +205,14 @@ def parse_fine_tuned_lora_name( start_index = 2 if name.startswith("base_model.model.") else 0 parts = name.split(".") - if ( - parts[-1] == "weight" - and len(parts) >= 2 - and (parts[-2] == "lora_A" or parts[-2] == "lora_B") - ): - new_name = ".".join(parts[start_index:-2]) - return new_name, parts[-2] == "lora_A" + if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2: + if parts[-2] in ["lora_A", "lora_B"]: + new_name = ".".join(parts[start_index:-2]) + return new_name, parts[-2] == "lora_A" + # For modules_to_save in classification. + elif parts[-2] in ["score", "classifier"]: + new_name = parts[-2] + return new_name, False if parts[-1] == "lora_embedding_A" or parts[-1] == "lora_embedding_B": new_name = ".".join(parts[start_index:-1]) diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 4c00c6531bc8..93624523e2f8 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -157,6 +157,10 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: # 3D->2D conversion when running under the universal 2D wrapper. lora.is_3d_lora_weight = lora_request.is_3d_lora_weight + # Validate classification-head weights. + self._adapter_manager._validate_modules_to_save(lora) + self._adapter_manager._validate_token_classification_lora(lora) + except FileNotFoundError as e: # FileNotFoundError should be raised if both # - No adapter found to download from huggingface (or in diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 090991bedb54..23354d66f3c4 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -74,6 +74,20 @@ def get_supported_tasks(self) -> Set[PoolingTask]: def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: return self.poolers_by_task[task].get_pooling_updates(task) + def replace_classifier( + self, + new_classifier: ClassifierFn, + ) -> None: + """Replaces the classifier to the LoRA-wrapped version.""" + for task, pooler in self.poolers_by_task.items(): + if task not in {"classify"}: + # Now LoRA only supports classification tasks, + # so we only replace the classifier for "classify" task. + continue + head = getattr(pooler, "head", None) + if head is not None and hasattr(head, "classifier"): + head.classifier = new_classifier + def forward( self, hidden_states: torch.Tensor, From ca542ee154e91c9a7757fb0b8f9d90ee578fb613 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Wed, 16 Sep 2026 22:16:09 +0800 Subject: [PATCH 328/774] [Model][LoRA] Enable LoRA support for ModernBertModel (#57148) Signed-off-by: linitra24 --- docs/models/pooling_models/embed.md | 2 +- vllm/model_executor/models/modernbert.py | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index a60bd4307a1e..13a2e0d38147 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -47,7 +47,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `JinaEmbeddingsV5Model`C | Qwen3-decoder or EuroBERT-encoder backbone with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small`, `jinaai/jina-embeddings-v5-text-nano` (see note) | ✅︎ | ✅︎ | | `LlamaBidirectionalModel`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ | | `LlamaModel`C, `LlamaForCausalLM`C, `MistralModel`C, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | -| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | +| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | ✅︎ | | | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | | `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ | diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 8fe80bbe19f1..df901a6566d9 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -32,7 +32,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors -from .interfaces import SupportsCrossEncoding +from .interfaces import SupportsCrossEncoding, SupportsLoRA from .interfaces_base import attn_type, default_pooling_type from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -262,7 +262,7 @@ def forward( @support_torch_compile @default_pooling_type(seq_pooling_type="CLS") -class ModernBertModel(nn.Module): +class ModernBertModel(nn.Module, SupportsLoRA): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "model.layers.": "encoder_layer.layers.", @@ -271,6 +271,11 @@ class ModernBertModel(nn.Module): } ) + packed_modules_mapping = { + "Wqkv": ["Wqkv"], + "Wi": ["Wi"], + } + def __init__( self, vllm_config: VllmConfig, From e97ff8061344d1dbf954ad3c3174ab3b576aaf33 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:24:59 +0200 Subject: [PATCH 329/774] [BugFix][PCP] Handle missing DP metadata in one-sided EP (#54016) Signed-off-by: LopezCastroRoberto Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../prepare_finalize/flashinfer_nvlink_one_sided.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index 0b57afb6062c..d7a84c931b8d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -13,9 +13,10 @@ from vllm.utils.flashinfer import nvfp4_block_scale_interleave -def get_local_sizes(): +def get_local_sizes() -> list[int] | None: dp_metadata = get_forward_context().dp_metadata - assert dp_metadata is not None + if dp_metadata is None: # PCP with DP=1 + return None return dp_metadata.get_chunk_sizes_across_dp_rank() From dff1bde84dd6e34a49c150116d0f212507280910 Mon Sep 17 00:00:00 2001 From: semerandre <150122900+semerandre@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:07:52 +0200 Subject: [PATCH 330/774] [Model] Qwen4Exp: fp8_e4m3 main KV cache on the QSA path (#55557) Signed-off-by: Andrea Semeraro Co-authored-by: Claude Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- tests/models/qwen4_exp/test_qsa_reference.py | 75 +++++++--- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 140 ++++++++++++++++--- vllm/models/qwen4_exp/nvidia/qsa.py | 117 ++++++++++++++-- 3 files changed, 287 insertions(+), 45 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 16294a30d769..b3bc0d1cae6d 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -203,7 +203,14 @@ def _qsa_sparse_paged_attention_reference( block_table: torch.Tensor, token_to_req: torch.Tensor, softmax_scale: float, + k_scale: float = 1.0, + v_scale: float = 1.0, ) -> torch.Tensor: + """Dense reference for QSA sparse paged attention. + + Mirrors the kernel's dequant: fp8-e4m3 K/V caches are dequantized with the + per-tensor k_scale/v_scale host floats; bf16 caches use unit scales. + """ output = torch.zeros_like(q) repeats = q.shape[1] // k_cache.shape[2] page_size = k_cache.shape[1] @@ -215,13 +222,15 @@ def _qsa_sparse_paged_attention_reference( request = token_to_req[row].long() pages = block_table[request, logical // page_size].long() offsets = logical % page_size - keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1) - values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1) - scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float()) - probabilities = torch.softmax(scores * softmax_scale, dim=-1) - output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to( - q.dtype + keys = (k_cache[pages, offsets].float() * k_scale).repeat_interleave( + repeats, dim=1 ) + values = (v_cache[pages, offsets].float() * v_scale).repeat_interleave( + repeats, dim=1 + ) + scores = torch.einsum("hd,khd->hk", q[row].float(), keys) + probabilities = torch.softmax(scores * softmax_scale, dim=-1) + output[row] = torch.einsum("hk,khd->hd", probabilities, values).to(q.dtype) return output @@ -906,22 +915,28 @@ def test_qsa_block_expansion_correctness() -> None: "page_size", "use_prefill_config", "num_requests", + "fp8", ), [ # Production page sizes from hybrid-cache block alignment: 784/800 # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head # splits are per-rank TP1/TP2/TP4; the largest batch runs both # use_prefill_config variants. - pytest.param(1, 24, 2, 1600, True, 2, id="tp1_r1"), - pytest.param(16, 12, 1, 1600, True, 3, id="tp2_r16"), - pytest.param(32, 6, 1, 800, True, 5, id="tp4_r32"), - pytest.param(128, 24, 2, 1568, True, 7, id="tp1_r128"), - pytest.param(257, 6, 1, 800, True, 13, id="tp4_r257"), - pytest.param(513, 6, 1, 784, True, 17, id="tp4_r513"), - pytest.param(700, 6, 1, 800, True, 23, id="tp4_r700"), - pytest.param(1024, 24, 2, 1600, True, 33, id="tp1_r1024"), - pytest.param(2048, 24, 2, 1600, True, 63, id="tp1_r2048_prefill"), - pytest.param(2048, 24, 2, 1600, False, 63, id="tp1_r2048_uniform"), + pytest.param(1, 24, 2, 1600, True, 2, False, id="tp1_r1"), + pytest.param(16, 12, 1, 1600, True, 3, False, id="tp2_r16"), + pytest.param(32, 6, 1, 800, True, 5, False, id="tp4_r32"), + pytest.param(128, 24, 2, 1568, True, 7, False, id="tp1_r128"), + pytest.param(257, 6, 1, 800, True, 13, False, id="tp4_r257"), + pytest.param(513, 6, 1, 784, True, 17, False, id="tp4_r513"), + pytest.param(700, 6, 1, 800, True, 23, False, id="tp4_r700"), + pytest.param(1024, 24, 2, 1600, True, 33, False, id="tp1_r1024"), + pytest.param(2048, 24, 2, 1600, True, 63, False, id="tp1_r2048_prefill"), + pytest.param(2048, 24, 2, 1600, False, 63, False, id="tp1_r2048_uniform"), + # fp8_e4m3 K/V caches on the TP1 head split. + pytest.param(1, 24, 2, 1600, True, 2, True, id="tp1_r1_fp8"), + pytest.param(128, 24, 2, 1568, True, 7, True, id="tp1_r128_fp8"), + pytest.param(2048, 24, 2, 1600, True, 63, True, id="tp1_r2048_prefill_fp8"), + pytest.param(2048, 24, 2, 1600, False, 63, True, id="tp1_r2048_uniform_fp8"), ], ) def test_qsa_sparse_paged_attention_correctness( @@ -931,8 +946,18 @@ def test_qsa_sparse_paged_attention_correctness( page_size: int, use_prefill_config: bool, num_requests: int, + fp8: bool, ) -> None: + """QSA sparse paged attention matches the dense reference. + + fp8 only changes the K/V cache dtype (e4m3 with a per-tensor scale pair) and + the scales; the reference dequantizes the same cache with those scales, so + both paths compare the production kernel against the reference on identical + inputs. fp8=True additionally covers the host-side scale folding. + """ torch.manual_seed(2) + # One QSA attention problem: bf16 Q and paged K/V, a packed selection with + # the trailing count column, block table and row-to-request map. head_dim = 256 num_selected_pages = 64 # Keep the newest page outside the synthetic top-k as causal headroom. @@ -1013,8 +1038,18 @@ def test_qsa_sparse_paged_attention_correctness( indexer_budget, logical_indices, ) - assert logical_indices.shape == (num_rows, selection_width + 1) - scale = q.shape[-1] ** -0.5 + + scale = head_dim**-0.5 + + if fp8: + # A fixed non-unit pair (k != v) exercises the host-side scale folding + # and catches a k/v swap; scales are host floats, as the layer exposes + # them. Stored values are the scaled ones, as reshape_and_cache does. + k_scale, v_scale = 0.5, 2.0 + k_cache = (k_cache / k_scale).to(torch.float8_e4m3fn) + v_cache = (v_cache / v_scale).to(torch.float8_e4m3fn) + else: + k_scale, v_scale = 1.0, 1.0 actual = qsa_ops.qsa_sparse_paged_attention( q, @@ -1024,6 +1059,8 @@ def test_qsa_sparse_paged_attention_correctness( block_table, token_to_req, use_prefill_config=use_prefill_config, + k_scale=k_scale, + v_scale=v_scale, output_gate=output_gate, ) expected = _qsa_sparse_paged_attention_reference( @@ -1034,6 +1071,8 @@ def test_qsa_sparse_paged_attention_correctness( block_table, token_to_req, scale, + k_scale=k_scale, + v_scale=v_scale, ) expected = expected * torch.sigmoid(output_gate) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index d3986ba38ae0..6d0df0e741b7 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -4,15 +4,24 @@ from __future__ import annotations +from functools import lru_cache + import torch from vllm.model_executor.warmup.jit_warmup_triton_helper import ( TritonWarmupTensor, triton_scalar_specialization_rep, ) +from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton +@lru_cache(maxsize=1) +def _is_sm120() -> bool: + """True on sm_120 (RTX PRO 6000 Blackwell): selects the sm_120 tuning table.""" + return current_platform.get_device_capability() == (12, 0) + + @triton.jit(do_not_specialize=["num_rows", "num_requests"]) def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, @@ -24,6 +33,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + softmax_scale, + output_scale, output_gate_ptr, stride_q_row, stride_q_head, @@ -52,6 +63,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -81,7 +93,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32) normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32) accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) - softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 + # softmax_scale is the host-side attention scale (1/sqrt(head_dim), with the + # fp8 K dequant scale already pre-multiplied in); convert to log2 units once + # here for the exp2-based online softmax. + score_scale = softmax_scale * 1.4426950408889634 tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N)) @@ -129,15 +144,26 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) + if IS_FP8: + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK + # measured slower here and less accurate). + keys = keys.to(query.dtype) scores = tl.dot(query, keys) - # Scaling scores avoids re-quantizing a scaled query to BF16. - scores *= softmax_scale_log2 + # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 + # caches the K dequant scale is already folded into softmax_scale on the + # host. + scores *= score_scale scores = tl.where(valid[None, :], scores, -1.0e20) next_max = tl.maximum(max_value, tl.max(scores, axis=1)) alpha = tl.math.exp2(max_value - next_max) probabilities = tl.where( valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 ) + if IS_FP8: + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online + # softmax) so fp16 has the range, its wider mantissa is more + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. + values = values.to(tl.float16) accumulator = tl.dot( probabilities.to(values.dtype), values, @@ -147,9 +173,14 @@ def _qsa_sparse_paged_gqa_splitk_kernel( max_value = next_max has_values = normalizer > 0 + # Fold the fp8 V dequant scale (output_scale, 1.0 for bf16) into a per-row + # reciprocal normalizer, so the output is a per-row multiply rather than a + # HEAD_DIM-wide scale. The split-K LSE below keeps the unscaled normalizer, + # so the merge stays correct. + inv_normalizer = output_scale / tl.maximum(normalizer, 1.0e-20) normalized_output = tl.where( has_values[:, None], - accumulator / tl.maximum(normalizer[:, None], 1.0e-20), + accumulator * inv_normalizer[:, None], 0.0, ) output_mask = head_offsets[:, None] < GROUP_SIZE @@ -439,18 +470,65 @@ def _compress_qsa_groups_kernel( ) +def _select_sm120_config( + base_programs: int, use_prefill_config: bool, is_fp8: bool +) -> tuple[int, int, int]: + """(block_n, target_splits, num_warps) retuned on sm_120 (RTX PRO 6000 + Blackwell), split by cache dtype: bf16 and fp8 favour different configs on + sm_120, most visibly on the large-prefill region. Each entry is the fastest + config that stays correct on its own path; the main win is more warps.""" + if is_fp8: + if base_programs > 2048: + return (32, 2, 1) if use_prefill_config else (64, 1, 2) + if base_programs <= 24: + return 64, 64, 8 + if base_programs <= 32: + return 128, 8, 4 + if base_programs <= 64: + return 64, 8, 4 + if base_programs <= 128: + return 32, 4, 4 + if base_programs <= 256: + return 128, 4, 4 + if base_programs <= 512: + return 64, 4, 4 + return 32, 2, 1 + # bf16 K/V cache. + if base_programs > 2048: + return 64, 2, 2 + if base_programs <= 24: + return 64, 64, 8 + if base_programs <= 64: + return 64, 8, 8 + if base_programs <= 128: + return 32, 8, 8 + if base_programs <= 256: + return 32, 8, 1 + if base_programs <= 512: + return 64, 4, 2 + return 32, 4, 1 + + def _select_config( - num_rows: int, num_kv_heads: int, use_prefill_config: bool, num_columns: int + num_rows: int, + num_kv_heads: int, + use_prefill_config: bool, + num_columns: int, + is_fp8: bool = False, ) -> tuple[int, int, int, int]: """Select (block_n, num_warps, num_tiles, num_splits) for the kernel. - Tuned on GB300 for the Qwen3.8-Flash-Next TP1/TP2/TP4 shapes, keyed on - base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on - use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the - uniform decode/verify length). + Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits + on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is + the uniform decode/verify length). This default table was tuned on GB300; + sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead. """ base_programs = num_rows * num_kv_heads - if base_programs > 2048: + if _is_sm120(): + BLOCK_N, target_splits, num_warps = _select_sm120_config( + base_programs, use_prefill_config, is_fp8 + ) + elif base_programs > 2048: BLOCK_N, target_splits, num_warps = ( (32, 1, 1) if use_prefill_config else (64, 1, 2) ) @@ -483,10 +561,17 @@ def qsa_sparse_paged_attention( token_to_req: torch.Tensor, use_prefill_config: bool, out: torch.Tensor | None = None, + k_scale: float | None = None, + v_scale: float | None = None, *, output_gate: torch.Tensor, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches. + """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. + + With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales + as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is + pre-multiplied into the softmax scale and v_scale becomes the kernel's + output scale, so the kernel needs no device scale buffers. logical_indices is the PACKED selection buffer: [rows, selection_width + 1] with the trailing column holding each row's valid-entry count (written by @@ -510,7 +595,19 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 - assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16 + assert q.dtype == torch.bfloat16 + assert k_cache.dtype == v_cache.dtype + is_fp8 = k_cache.dtype == torch.float8_e4m3fn + if is_fp8: + assert k_scale is not None and v_scale is not None + # Host pre-multiply: fold the K dequant scale into the attention scale + # and pass V's dequant scale as the kernel's output scale. + softmax_scale = (head_dim**-0.5) * float(k_scale) + output_scale = float(v_scale) + else: + assert k_cache.dtype == torch.bfloat16 + softmax_scale = head_dim**-0.5 + output_scale = 1.0 assert logical_indices.dtype == block_table.dtype == torch.int32 assert token_to_req.dtype == torch.int32 assert q.device == k_cache.device == v_cache.device @@ -535,7 +632,7 @@ def qsa_sparse_paged_attention( block_m = triton.next_power_of_2(group_size) selection_width = logical_indices.shape[1] - 1 # trailing column is the count block_n, partial_warps, num_tiles, num_splits = _select_config( - q.shape[0], k_cache.shape[2], use_prefill_config, selection_width + q.shape[0], k_cache.shape[2], use_prefill_config, selection_width, is_fp8 ) # Split=1 writes output directly and compiles out all workspace accesses. @@ -565,6 +662,8 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + softmax_scale, + output_scale, output_gate_view, q.stride(0), q.stride(1), @@ -593,6 +692,7 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + IS_FP8=is_fp8, num_warps=partial_warps, num_stages=2, ) @@ -630,13 +730,18 @@ def warmup_qsa_sparse_paged_attention( head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) + # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. + is_fp8 = kv_cache.dtype == torch.uint8 + cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype num_kv_heads = key_cache.shape[2] group_size = num_query_heads // num_kv_heads block_m = triton.next_power_of_2(group_size) # Every config the dispatch can pick for this group size. profiles = { - _select_config(num_rows, num_kv_heads, use_prefill_config, selection_width) + _select_config( + num_rows, num_kv_heads, use_prefill_config, selection_width, is_fp8 + ) for num_rows in range(1, 8193) for use_prefill_config in (False, True) } @@ -650,12 +755,12 @@ def warmup_qsa_sparse_paged_attention( torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) ) k_cache_ptr = TritonWarmupTensor( - key_cache.dtype, + cache_dtype, shape=tuple(key_cache.shape), strides=tuple(key_cache.stride()), ) v_cache_ptr = TritonWarmupTensor( - value_cache.dtype, + cache_dtype, shape=tuple(value_cache.shape), strides=tuple(value_cache.stride()), ) @@ -701,6 +806,8 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + 1.0, + 1.0, output_gate_ptr, row_stride, head_stride, @@ -729,6 +836,7 @@ def warmup_qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + IS_FP8=is_fp8, num_warps=warps, num_stages=2, grid=(num_rows, num_kv_heads, num_splits), diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 1c2808a6c47f..e480142a2218 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, ) @@ -64,7 +65,38 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): """FullAttentionSpec backend used by the merged QSA owner.""" supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] + # fp8/fp8_e4m3: e4m3 bytes in a uint8 cache, written by reshape_and_cache + # with the layer's per-tensor scales and dequantized on load inside the QSA + # Triton kernel. flash-attn never runs over this cache, so its fp8 probe + # does not apply (see supports_kv_cache_dtype and the impl constructor). + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: + return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + # QSA dequantizes the fp8 KV in its own Triton kernel and never runs + # flash-attn over the quantized cache, so the parent's fp8-KV rejection + # does not apply and every combination it is handed is accepted here. + return None @staticmethod def get_name() -> str: @@ -98,16 +130,53 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): supports_dcp: bool = False supports_pcp: bool = False - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + # The parent constructor probes flash-attn for quantized-KV support and + # raises where it is unavailable (sm120), but QSA dequantizes fp8 inside + # its own Triton kernel and never runs flash-attn over the cache. Hand + # the parent "auto" for that probe and restore the real dtype afterwards: + # the parent only uses it there, and do_kv_cache_update reads the + # attribute at call time. + real_kv_cache_dtype = kv_cache_dtype + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + kv_cache_dtype = "auto" + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + sinks, + ) + self.kv_cache_dtype = real_kv_cache_dtype if not is_flash_attn_varlen_func_available(): raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") if self.dcp_world_size != 1: raise NotImplementedError( "Qwen4Exp QSA does not support decode context parallelism" ) - if self.kv_cache_dtype not in ("auto", "bfloat16"): - raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache") + if self.kv_cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"): + raise NotImplementedError( + "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache" + ) self.supports_quant_query_input = False def forward_qsa( @@ -144,8 +213,23 @@ def forward_qsa( logical_indices = topk_buffer[:num_tokens] token_to_req = token_to_req[:num_tokens] key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) - if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16: - raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V") + k_scale = v_scale = None + if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): + # The cache is allocated as uint8; reinterpret the e4m3 bytes + # (same itemsize, so shape and strides are preserved). + key_cache = key_cache.view(torch.float8_e4m3fn) + value_cache = value_cache.view(torch.float8_e4m3fn) + # Host-side per-tensor dequant scales (Python floats), as used by + # other host-scale backends; folded into the kernel's scales. + k_scale = layer._k_scale_float + v_scale = layer._v_scale_float + if query.dtype != torch.bfloat16 or key_cache.dtype not in ( + torch.bfloat16, + torch.float8_e4m3fn, + ): + raise NotImplementedError( + "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V" + ) from .ops.qsa import qsa_sparse_paged_attention @@ -158,6 +242,8 @@ def forward_qsa( token_to_req, use_prefill_config, output[:num_tokens], + k_scale=k_scale, + v_scale=v_scale, output_gate=output_gate[:num_tokens], ) return output @@ -185,8 +271,15 @@ def __init__( raise ValueError("Qwen4Exp QSA requires a paged KV cache") if model_config.dtype != torch.bfloat16: raise NotImplementedError("Qwen4Exp QSA currently requires BF16") - if cache_config.cache_dtype not in ("auto", "bfloat16"): - raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache") + if cache_config.cache_dtype not in ( + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + ): + raise NotImplementedError( + "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache" + ) if getattr(quant_config, "kv_cache_scheme", None) is not None: raise NotImplementedError("Qwen4Exp QSA does not support KV quantization") parallel_config = vllm_config.parallel_config @@ -285,8 +378,10 @@ def __init__( self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( self.kv_cache_dtype, model_config ) - if self.kv_cache_torch_dtype != torch.bfloat16: - raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage") + if self.kv_cache_torch_dtype not in (torch.bfloat16, torch.uint8): + raise NotImplementedError( + "Qwen4Exp QSA requires BF16 or FP8-e4m3 (uint8) cache storage" + ) self.kv_sharing_target_layer_name = None self.kv_cache = torch.tensor([]) set_default_quant_scales(self, register_buffer=True) From 42919b49c573aa0b06710616a3b9785674f36ef5 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:23:31 -0400 Subject: [PATCH 331/774] [P/D] Report prefill worker cache hits in prompt_tokens_details (#54222) Signed-off-by: Robert Shaw Co-authored-by: Claude --- tests/v1/engine/test_output_processor.py | 54 ++++++++++++++++++- .../kv_connector/unit/test_nixl_connector.py | 53 ++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 11 ++++ vllm/v1/engine/output_processor.py | 16 ++++++ 4 files changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index f578aae7f015..64ffc6b26675 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -23,6 +23,7 @@ from vllm.v1.engine import ( EngineCoreEvent, EngineCoreEventType, + EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, FinishReason, @@ -32,7 +33,7 @@ RequestOutputCollector, RequestState, ) -from vllm.v1.metrics.stats import IterationStats, SchedulerStats +from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats @pytest.mark.parametrize("flat_logprobs", [False, True]) @@ -171,6 +172,57 @@ def test_incremental_detokenization( assert not output_processor.has_unfinished_requests() +@pytest.mark.parametrize("do_remote_prefill", [True, False]) +def test_remote_prefill_cached_tokens_override(do_remote_prefill: bool): + """P/D disaggregation: num_cached_tokens should report the P worker's + cache hits (passed via kv_transfer_params) instead of the local count, + which sees the KVs pulled from the remote prefill as a ~100% hit.""" + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + + prompt_tokens = [1, 2, 3, 4, 5, 6, 7, 8] + kv_transfer_params = { + "do_remote_prefill": do_remote_prefill, + "remote_prefill_cached_tokens": 5, + } + request = EngineCoreRequest( + request_id="request-0-int", + external_req_id="request-0", + prompt_token_ids=prompt_tokens, + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=SamplingParams( + detokenize=False, + extra_args={"kv_transfer_params": kv_transfer_params}, + ), + pooling_params=None, + ) + output_processor.add_request(request, prompt=None) + + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=len(prompt_tokens), + num_local_cached_tokens=0, + num_external_cached_tokens=len(prompt_tokens) - 1, + ) + processed = output_processor.process_outputs( + [ + EngineCoreOutput( + request_id="request-0-int", + new_token_ids=[42], + prefill_stats=prefill_stats, + ) + ] + ) + request_output = processed.request_outputs[0] + if do_remote_prefill: + assert request_output.num_cached_tokens == 5 + else: + assert request_output.num_cached_tokens == len(prompt_tokens) - 1 + + def test_request_stream_interval_raises_but_not_below_engine_default( dummy_test_vectors, ): diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 39fb0bb54fae..e41bd2644a9a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -74,6 +74,7 @@ from vllm.v1.request import RequestStatus from .utils import ( + create_model_runner_output, create_request, create_scheduler, create_vllm_config, @@ -348,6 +349,58 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): assert req_meta.awaiting_kvs is False +def test_prefill_exports_cached_tokens_in_kv_transfer_params(): + """The P worker reports its own prefix-cache hits in the returned + kv_transfer_params so the D worker can surface them in + prompt_tokens_details instead of the ~100% local hit it measures + when pulling the KVs from the remote.""" + vllm_config = create_vllm_config() + scheduler = create_scheduler(vllm_config) + + BLOCK_SIZE = vllm_config.cache_config.block_size + NUM_TOKENS = BLOCK_SIZE * 3 + + # Warm the prefix cache with a request sharing the full prompt. + warmup = create_request( + request_id=1, + num_tokens=NUM_TOKENS, + common_prefix_len=NUM_TOKENS, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(warmup) + scheduler_output = scheduler.schedule() + model_runner_output = create_model_runner_output([warmup], use_eos=True) + scheduler.update_from_output(scheduler_output, model_runner_output) + + # P-side request with the same prompt hits the local cache on all but + # the last block (recomputed to obtain logits). + request = create_request( + request_id=2, + num_tokens=NUM_TOKENS, + common_prefix_len=NUM_TOKENS, + block_size=BLOCK_SIZE, + do_remote_decode=True, + ) + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + assert request.prefill_stats is not None + assert request.prefill_stats.num_cached_tokens == NUM_TOKENS - BLOCK_SIZE + + # max_tokens=1, so the request finishes and returns kv_transfer_params. + model_runner_output = create_model_runner_output([request]) + engine_core_outputs = scheduler.update_from_output( + scheduler_output, model_runner_output + ) + + output = engine_core_outputs[0].outputs[0] + assert output.finish_reason is not None + assert output.kv_transfer_params is not None + assert ( + output.kv_transfer_params["remote_prefill_cached_tokens"] + == NUM_TOKENS - BLOCK_SIZE + ) + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c1e956b9eca1..b662bcc10b44 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2110,6 +2110,17 @@ def update_from_output( finished = self._handle_stopped_request(request) if finished: kv_transfer_params, ec_transfer_params = self._free_request(request) + if ( + prefill_stats is not None + and kv_transfer_params is not None + and kv_transfer_params.get("do_remote_prefill") + ): + # P-side cache hits, so D can report them in + # prompt_tokens_details instead of its own (~100%) + # hit rate from the KV transfer. + kv_transfer_params["remote_prefill_cached_tokens"] = ( + prefill_stats.num_cached_tokens + ) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 257287243dbd..b8fbddf72ff8 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -155,6 +155,7 @@ def __init__( n: int | None = None, temperature: float | None = None, stream_input: bool = False, + remote_prefill_cached_tokens: int | None = None, ): self.request_id = request_id self.external_req_id = external_req_id @@ -179,6 +180,7 @@ def __init__( self.queue = queue self.num_cached_tokens = 0 self.num_cache_creation_tokens = 0 + self.remote_prefill_cached_tokens = remote_prefill_cached_tokens # Per-sequence spec-decode accumulator; arrives once (on finish) via # EngineCoreOutput, then attached to this sequence's CompletionOutput. self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None @@ -230,7 +232,18 @@ def from_new_request( log_stats: bool, stream_interval: int, ) -> "RequestState": + remote_prefill_cached_tokens = None if sampling_params := request.sampling_params: + # In a remote prefill scenario, report cached tokens + # as cache hit rate on the remote P worker. + if sampling_params.extra_args: + kv_transfer_params = sampling_params.extra_args.get( + "kv_transfer_params" + ) + if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"): + cached = kv_transfer_params.get("remote_prefill_cached_tokens") + if isinstance(cached, int): + remote_prefill_cached_tokens = cached if not sampling_params.detokenize: tokenizer = None output_kind = sampling_params.output_kind @@ -281,6 +294,7 @@ def from_new_request( log_stats=log_stats, stream_interval=stream_interval, stream_input=request.resumable, + remote_prefill_cached_tokens=remote_prefill_cached_tokens, ) def make_request_output( @@ -687,6 +701,8 @@ def process_outputs( req_state.num_cache_creation_tokens = ( engine_core_output.prefill_stats.num_cache_creation_tokens ) + if req_state.remote_prefill_cached_tokens is not None: + req_state.num_cached_tokens = req_state.remote_prefill_cached_tokens req_state.is_prefilling = False if engine_core_output.spec_decode_metrics is not None: From 804e5377aa0b0cea9544061beb2ebff162543c72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 16 Sep 2026 18:39:07 +0200 Subject: [PATCH 332/774] [Bugfix][HiSparse][PD] Align region-mapped pulls across logical block sizes (#57077) Signed-off-by: NickLucche --- .../unit/test_nixl_connector_hma.py | 169 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 46 ++++- .../kv_connector/v1/nixl/metadata.py | 5 + .../kv_connector/v1/nixl/pull_worker.py | 58 ++++-- 4 files changed, 267 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index ae2ddece7f77..1e34340d0f10 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -3,6 +3,8 @@ """Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill.""" import gc +import queue +from types import SimpleNamespace from unittest.mock import MagicMock, patch import msgspec @@ -39,6 +41,173 @@ ) +@pytest.fixture +def region_pull_worker(): + """Real page mapping and descriptors, with only the transport mocked.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping + + worker = object.__new__(NixlConnectorWorker) + worker.engine_id = "D" + worker.tp_rank = 0 + worker.block_size = 64 + worker._physical_blocks_per_logical_kv_block = 1 + worker.enable_permute_local_kv = False + worker.enable_heterogeneous_attn_post_process = False + worker._engine_last_active = {} + worker._bidirectional_kv_xfer_enabled = False + worker._recving_transfers = {} + worker.use_mla, worker._has_mamba = True, False + worker.dcp_size = 1 + spec = MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 + ) + worker.kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["resident"], spec, enable_kv_transfer=False), + KVCacheGroupSpec(["host"], spec), + KVCacheGroupSpec(["indexer"], spec), + ], + ) + worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec) + worker.region_group_ids = [0, 1] + worker.dst_region_group_ids = {"P": [0, 0]} + worker.num_regions = 2 + worker.block_len_per_layer = [1024, 1024] + worker.transfer_topo = MagicMock() + worker.transfer_topo.tp_ratio.return_value = 1 + worker.transfer_topo.block_size_ratio.return_value = 1 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_tp_size=4, + remote_dcp_size=1, + remote_block_size=64, + remote_physical_blocks_per_logical=2, + ) + worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)} + worker._mixed_mem_types = True + worker.src_xfer_handles_by_block_size = {64: 1} + worker._dram_src_handles_by_block_size = {64: 2} + worker.dst_xfer_side_handles = {"P": {0: 3}} + worker.dst_num_blocks = {"P": 100, "D": 100} + worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]} + worker._remote_agents = {"P": {(0, 0): "P-rank0"}} + worker._read_blocks_mixed = MagicMock() + worker.nixl_wrapper = MagicMock() + return worker + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("local_ratio,remote_ratio", [(1, 2), (2, 1), (2, 2)]) +@pytest.mark.parametrize("num_tokens", [160, 224]) +@pytest.mark.parametrize("cached", [0, 1]) +def test_region_pull_ignores_allocation_padding( + region_pull_worker, local_ratio, remote_ratio, num_tokens, cached +): + """Padded capacity must not shift either the host or indexer token range.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + from vllm.utils.math_utils import cdiv + + worker = region_pull_worker + worker._physical_blocks_per_logical_kv_block = local_ratio + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_tp_size=4, + remote_dcp_size=1, + remote_block_size=64, + remote_physical_blocks_per_logical=remote_ratio, + ) + local_count = cdiv(num_tokens, 64 * local_ratio) + remote_count = cdiv(num_tokens, 64 * remote_ratio) + # The host has a cached prefix, while the indexer still needs the full range. + local = [ + list(range(30 + cached, 30 + local_count)), + list(range(40, 40 + local_count)), + ] + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="request", + local_block_ids=local, + local_num_computed_blocks=(0, cached, 0), + awaiting_kvs=True, + kv_transfer_params={ + "remote_engine_id": "P", + "remote_request_id": "request-P", + "remote_host": "localhost", + "remote_port": 1, + "remote_num_tokens": num_tokens, + "remote_block_ids": [list(range(10, 10 + remote_count))], + }, + ) + meta = metadata.reqs_to_recv["request"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + local, local_ratio + ) + worker._read_blocks_for_req("request", meta) + read = worker._read_blocks_mixed.call_args.kwargs + valid_pages = cdiv(num_tokens, 64) + source_start = 10 * remote_ratio + assert read["remote_block_descs_ids"].tolist() == ( + list(range(source_start + cached * local_ratio, source_start + valid_pages)) + + list(range(100 + source_start, 100 + source_start + valid_pages)) + ) + assert read["local_block_descs_ids"].tolist() == ( + list(range((30 + cached) * local_ratio, 30 * local_ratio + valid_pages)) + + list(range(100 + 40 * local_ratio, 100 + 40 * local_ratio + valid_pages)) + ) + assert meta.region_blocks_to_zero == [ + list( + range(base * local_ratio + valid_pages, (base + local_count) * local_ratio) + ) + for base in (30, 40) + ] + + +@pytest.mark.cpu_test +def test_region_pull_completion_zeros_only_own_padding(region_pull_worker): + """Completion must preserve neighboring layers and already imported pages.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + RemoteMeta, + ReqMeta, + ) + + worker = region_pull_worker + # Two strided layer regions in one backing, with nonzero layer offsets. + backing = torch.full((4, 3, 8), 7, dtype=torch.uint8) + caches = {"host": backing[:, 1], "indexer": backing[:, 2]} + worker.device_kv_caches = caches + worker.region_names = list(caches) + worker.kv_caches_base_addr = {"D": {0: [c.data_ptr() for c in caches.values()]}} + worker.region_num_blocks = [4, 4] + worker.block_len_per_layer = [8, 8] + worker.block_stride_per_layer = [24, 24] + worker._recving_metadata = { + "request": ReqMeta( + local_block_ids=[[1, 2], [0, 1]], + local_physical_block_ids=[[1, 2], [0, 1]], + tp_size=1, + remote=RemoteMeta([[3]], "localhost", 1, "P", "request-P"), + region_blocks_to_zero=[[2], [1]], + ) + } + worker._get_new_notifs = MagicMock(return_value=set()) + worker._pop_done_transfers = MagicMock(return_value=({"request"}, set())) + worker._replicated_pcp_done_sending = set() + worker._failed_recv_reqs = queue.Queue() + worker._recv_failures = set() + worker._send_pending_recv_notifs = MagicMock() + worker._sync_device_after_direct_recv = MagicMock() + worker.use_host_buffer = False + worker._reqs_to_send = {} + result = worker.get_transfer_results() + expected = torch.full_like(backing, 7) + expected[2, 1] = 0 + expected[1, 2] = 0 + assert torch.equal(backing, expected) + assert result.finished_recving == {"request"} + + @pytest.mark.cpu_test @pytest.mark.parametrize( "swa_enabled,expected_sw_sizes", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 2704de589f05..db392078ca0c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2785,6 +2785,25 @@ def post_process_device_kv_on_receive_heterogeneous_attn( indices=indices, ) + def _zero_region_blocks(self, block_ids: BlockIds) -> None: + """Clear clipped physical pages in their owning region only.""" + if not any(block_ids): + return + bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank] + for region, blocks in enumerate(block_ids): + if not blocks: + continue + cache = self.device_kv_caches[self.region_names[region]] + storage = cache.untyped_storage() + pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_( + storage, + bases[region] - storage.data_ptr(), + (self.region_num_blocks[region], self.block_len_per_layer[region]), + (self.block_stride_per_layer[region], 1), + ) + for block in blocks: + pages[block].zero_() + def get_transfer_results(self) -> KVConnectorTransferResults: """ Get transfers that completed on this specific worker. @@ -2857,6 +2876,12 @@ def get_transfer_results(self) -> KVConnectorTransferResults: direct_device_recving.add(req_id) + if meta.region_blocks_to_zero is not None: + # P/D group positions differ. Use the actual region read plan, + # including any local allocation padding the read did not cover. + self._zero_region_blocks(meta.region_blocks_to_zero) + continue + # Post processing for heteroblocksize/layout, and for blocks the # transfer clipped. The latter happens either at remote-block # granularity (block_size_ratio > 1) or at kernel-block @@ -3254,13 +3279,32 @@ def _block_ids_by_region( @staticmethod def _apply_prefix_caching_by_region( - decode_block_ids: BlockIds, prefill_block_ids: BlockIds + decode_block_ids: BlockIds, + prefill_block_ids: BlockIds, + *, + num_computed_blocks: list[int] | None = None, + num_remote_blocks: int | None = None, ) -> tuple[BlockIds, BlockIds]: """Pair an uncached decode suffix with the same prefill regions.""" assert len(decode_block_ids) == len(prefill_block_ids) if not any(decode_block_ids): return [], prefill_block_ids + if num_computed_blocks is not None: + assert num_remote_blocks is not None + matched_decode, matched_prefill = [], [] + for decode_region, prefill_region, start in zip( + decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True + ): + count = min(len(decode_region), max(num_remote_blocks - start, 0)) + if start + count > len(prefill_region): + raise ValueError("Remote KV pages do not cover the requested range") + # An extra producer page can be allocation padding, not a + # cached prefix. Select by token position instead of list length. + matched_decode.append(list(decode_region[:count])) + matched_prefill.append(list(prefill_region[start : start + count])) + return matched_decode, matched_prefill + trimmed_prefill: list[list[int]] = [] for decode_region, prefill_region in zip( decode_block_ids, prefill_block_ids, strict=True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index c2b4d9dce753..1ae0ed72cecf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -230,6 +230,7 @@ class RemoteMeta: engine_id: str request_id: str blocks_expiry_time: float | None = None + num_tokens: int | None = None @dataclass @@ -252,6 +253,9 @@ class ReqMeta: # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS # and expects it in finished_recving; notify-only recvs must not be reported. awaiting_kvs: bool = False + # Worker-only, per-region physical pages to zero after a successful pull. + # None selects group-based completion; empty lists mean no zeroing. + region_blocks_to_zero: BlockIds | None = None class NixlConnectorMetadata(KVConnectorMetadata): @@ -326,5 +330,6 @@ def add_new_req_to_recv( host=kv_transfer_params["remote_host"], port=kv_transfer_params["remote_port"], blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"), + num_tokens=kv_transfer_params.get("remote_num_tokens"), ) self.reqs_to_recv[request_id] = req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 04caebbd1f21..1530b9fb6a69 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -20,6 +20,7 @@ _is_attention_spec, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.config import VllmConfig @@ -176,6 +177,10 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): "pure MLA models" ) assert len(plan.all_source_ranks) == 1 + if self.block_size != remote_info.remote_block_size: + raise NotImplementedError( + "Region-mapped NIXL transfers require matching physical block sizes" + ) remote_physical_block_ids = self._logical_to_kernel_block_ids( meta.remote.block_ids, remote_info.remote_physical_blocks_per_logical, @@ -183,13 +188,51 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_by_region = self._block_ids_by_region( remote_physical_block_ids, remote_region_groups ) + local_by_region = self._block_ids_by_region( + local_block_ids, local_region_groups + ) + num_computed_blocks = None + num_remote_blocks = None + if ( + meta.remote.num_tokens is not None + and meta.local_num_computed_blocks + and all(group >= 0 for group in local_region_groups) + and all(group >= 0 for group in remote_region_groups) + and not dcp_active + ): + transfer_groups = self.kv_cache_config.transfer_group_ids + num_computed_blocks = [ + meta.local_num_computed_blocks[transfer_groups[group]] + * self._physical_blocks_per_logical_kv_block + for group in local_region_groups + ] + num_remote_blocks = cdiv( + meta.remote.num_tokens, remote_info.remote_block_size + ) + elif ( + remote_info.remote_physical_blocks_per_logical + != self._physical_blocks_per_logical_kv_block + ): + raise NotImplementedError( + "Region-mapped pulls with different logical block sizes require " + "remote_num_tokens, per-group prefix counts, unshared regions " + "and DCP=1" + ) + matched_local, matched_remote = self._apply_prefix_caching_by_region( + local_by_region, + remote_by_region, + num_computed_blocks=num_computed_blocks, + num_remote_blocks=num_remote_blocks, + ) + meta.region_blocks_to_zero = [ + list(blocks[len(matched) :]) + for blocks, matched in zip(local_by_region, matched_local, strict=True) + ] read_specs = [ ReadSpec( remote_rank=plan.all_source_ranks[0], - local_block_ids=self._block_ids_by_region( - local_block_ids, local_region_groups - ), - remote_block_ids=remote_by_region, + local_block_ids=matched_local, + remote_block_ids=matched_remote, block_ids_by_region=True, ) ] @@ -396,12 +439,7 @@ def _read_blocks( self._recving_transfers.setdefault(request_id, []) return True - if read_spec.block_ids_by_region: - local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region( - decode_block_ids=local_block_ids, - prefill_block_ids=remote_block_ids, - ) - else: + if not read_spec.block_ids_by_region: assert ( len(remote_block_ids) == len(local_block_ids) From 8c1557a79c539ffe82d004d2a0c8d7b5e71159ce Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:44:29 +0100 Subject: [PATCH 333/774] Add `pydocstyle` to the `ruff` rules (#52136) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor Co-authored-by: Claude Opus 5 (1M context) --- .../test_lm_eval_correctness.py | 3 +- .../scripts/compare-json-results.py | 6 +- .../convert-results-json-to-markdown.py | 3 +- .buildkite/scripts/generate-nightly-index.py | 21 +- benchmarks/attention_benchmarks/batch_spec.py | 28 +- benchmarks/attention_benchmarks/benchmark.py | 20 +- benchmarks/attention_benchmarks/common.py | 15 +- benchmarks/attention_benchmarks/mla_runner.py | 51 +- benchmarks/attention_benchmarks/runner.py | 12 +- benchmarks/benchmark_batch_invariance.py | 6 +- benchmarks/benchmark_hash.py | 3 +- .../benchmark_hidden_state_extraction.py | 6 +- .../benchmark_long_document_qa_throughput.py | 11 +- benchmarks/benchmark_prefix_block_hash.py | 4 +- benchmarks/benchmark_prefix_caching.py | 3 +- .../benchmark_serving_structured_output.py | 5 +- benchmarks/benchmark_topk_topp.py | 9 +- .../merge_attn_states_benchmarks.py | 6 +- benchmarks/kernels/bench_concat_mla_q.py | 6 +- .../benchmark_2d_silu_mul_fp8_quant.py | 3 +- .../kernels/benchmark_cutlass_moe_fp8.py | 5 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 3 +- .../kernels/benchmark_device_communicators.py | 5 +- .../kernels/benchmark_fused_collective.py | 9 +- .../kernels/benchmark_kimi_k3_attn_res.py | 1 + .../kernels/benchmark_kimi_k3_kda_decode.py | 1 + .../benchmark_kimi_k3_latent_moe_tail.py | 1 + benchmarks/kernels/benchmark_lora.py | 55 +- benchmarks/kernels/benchmark_mla_k_concat.py | 3 +- benchmarks/kernels/benchmark_moe_defaults.py | 6 +- .../benchmark_per_token_group_quant.py | 2 +- .../kernels/benchmark_per_token_quant_fp8.py | 6 +- .../benchmark_rdna_hybrid_w4a16_gemm.py | 3 +- .../kernels/benchmark_reshape_and_cache.py | 1 - .../benchmark_reshape_and_cache_flash.py | 1 - .../benchmark_selective_state_update.py | 9 +- .../kernels/benchmark_silu_mul_fp8_quant.py | 7 +- .../kernels/benchmark_vit_aiter_fp8_attn.py | 1 + .../benchmark_vocab_parallel_embedding.py | 3 +- .../kernels/benchmark_w8a8_block_fp8.py | 2 + benchmarks/kernels/ir/bench_ir_ops.py | 3 +- benchmarks/kernels/ir/shapes.py | 4 +- benchmarks/kernels/utils.py | 3 +- benchmarks/multi_turn/bench_dataset.py | 3 +- .../multi_turn/convert_sharegpt_to_openai.py | 3 +- benchmarks/replayssm/e2e_decode_speedup.py | 1 + csrc/cpu/generate_cpu_attn_dispatch.py | 4 +- docs/mkdocs/gen_files/generate_argparse.py | 6 +- .../gen_files/generate_attention_backends.py | 6 +- docs/mkdocs/gen_files/generate_examples.py | 14 +- docs/mkdocs/gen_files/generated_content.py | 1 + docs/mkdocs/hooks/autoref_code.py | 3 +- docs/mkdocs/hooks/url_schemes.py | 11 +- examples/applications/api_server/server.py | 3 +- .../streamlit_openai_chatbot_webserver.py | 7 +- ...val_augmented_generation_with_langchain.py | 37 +- ...al_augmented_generation_with_llamaindex.py | 18 +- examples/deployment/async_llm_streaming.py | 6 +- examples/deployment/llm_engine_example.py | 3 +- .../disaggregated_encoder/disagg_epd_proxy.py | 20 +- .../disagg_proxy_demo.py | 3 +- .../disagg_proxy_multiturn.py | 6 +- .../disagg_proxy_pushconnector_demo.py | 3 +- .../example_connector/decode_example.py | 2 +- .../flexkv_connector/prefix_caching_flexkv.py | 3 +- .../decode_example.py | 2 +- .../load_recovery_example_connector.py | 3 +- .../lmcache/cpu_offload_lmcache.py | 3 +- .../disagg_proxy_server.py | 12 +- .../lmcache/kv_cache_sharing_lmcache_v1.py | 3 +- .../mooncake_connector_proxy.py | 18 +- .../automatic_prefix_caching_offline.py | 3 +- .../reproducibility_offline.py | 3 +- .../context_extension_offline.py | 3 +- .../data_parallel/data_parallel_offline.py | 3 +- .../kv_events/kv_events_subscriber.py | 2 +- examples/features/logits_processor/custom.py | 2 +- .../features/logits_processor/custom_req.py | 9 +- .../logits_processor/custom_req_init.py | 9 +- .../lora/lora_with_quantization_offline.py | 5 +- examples/features/lora/multilora_offline.py | 3 +- .../data_parallel_pause_resume.py | 3 +- .../prompt_embed/prompt_embed_offline.py | 3 +- .../features/reset_kv/reset_kv_offline.py | 3 +- .../load_sharded_state_offline.py | 3 +- .../save_sharded_state_offline.py | 3 +- .../mlpspeculator_offline.py | 3 +- .../structured_outputs_offline.py | 3 +- .../torchrun/torchrun_dp_example_offline.py | 3 +- .../torchrun/torchrun_example_offline.py | 3 +- .../multimodal/audio_language_offline.py | 6 +- .../encoder_decoder_multimodal_offline.py | 6 +- ...i_chat_completion_client_for_multimodal.py | 2 - .../multimodal/qwen2_5_omni/only_thinker.py | 3 +- .../multimodal/qwen3_omni/only_thinker.py | 3 +- .../vision_language_multi_image_offline.py | 7 +- .../multimodal/vision_language_offline.py | 9 +- examples/generate/trace_replay_offline.py | 3 +- .../openai_embedding_long_text/client.py | 3 +- .../pooling/embed/vision_embedding_offline.py | 3 +- .../pooling/embed/vision_embedding_online.py | 22 +- .../plugin/prithvi_geospatial_mae_offline.py | 13 +- .../pooling/reward/sequence_reward_offline.py | 3 +- .../pooling/reward/sequence_reward_online.py | 3 +- .../pooling/reward/token_reward_offline.py | 3 +- .../pooling/reward/token_reward_online.py | 3 +- .../pooling/score/cohere_rerank_client.py | 3 +- .../pooling/score/colbert_rerank_online.py | 3 +- .../score/colmodernvbert_rerank_online.py | 3 +- .../pooling/score/colqwen3_5_rerank_online.py | 3 +- .../pooling/score/colqwen3_rerank_online.py | 3 +- .../pooling/score/convert_model_to_seq_cls.py | 15 +- .../pooling/score/qwen3_reranker_offline.py | 7 +- .../pooling/score/qwen3_reranker_online.py | 3 +- examples/pooling/score/rerank_api_online.py | 3 +- examples/pooling/score/score_api_online.py | 3 +- .../pooling/score/using_template_offline.py | 5 +- .../pooling/score/using_template_online.py | 6 +- .../pooling/score/vision_rerank_api_online.py | 3 +- .../pooling/score/vision_reranker_offline.py | 3 +- .../pooling/score/vision_score_api_online.py | 3 +- .../forced_alignment_offline.py | 3 +- .../token_classify/forced_alignment_online.py | 3 +- examples/pooling/token_classify/ner_online.py | 3 +- .../colqwen3_token_embed_online.py | 3 +- .../token_embed/jina_reranker_v3_online.py | 3 +- .../multi_vector_retrieval_online.py | 3 +- examples/ray_serving/batch_llm_inference.py | 3 +- examples/ray_serving/ray_serve_deepseek.py | 3 +- ...at_completion_tool_calls_with_reasoning.py | 3 +- .../openai_chat_completion_with_reasoning.py | 3 +- ...hat_completion_with_reasoning_streaming.py | 3 +- examples/reasoning/openai_responses_client.py | 3 +- examples/rl/rlhf_async_new_apis.py | 3 +- examples/rl/rlhf_http_ipc.py | 3 +- examples/rl/rlhf_http_nccl.py | 3 +- examples/rl/rlhf_ipc_fsdp_ep.py | 3 +- examples/rl/rlhf_nccl_fsdp_ep.py | 6 +- examples/rl/rlhf_sharded_rdt_small_ep.py | 3 +- examples/rl/routed_experts_e2e.py | 3 +- .../openai/openai_transcription_client.py | 14 +- .../realtime/openai_realtime_client.py | 11 +- .../openai_realtime_microphone_client.py | 3 +- ...penai_chat_completion_client_with_tools.py | 3 +- ...t_completion_client_with_tools_required.py | 3 +- ..._chat_completion_client_with_tools_xlam.py | 8 +- ...letion_client_with_tools_xlam_streaming.py | 7 +- .../openai_responses_client_with_mcp_tools.py | 3 +- .../openai_responses_client_with_tools.py | 6 +- pyproject.toml | 11 + scripts/autotune_helion_kernels.py | 3 +- scripts/benchmark_helion_kernels.py | 4 +- setup.py | 10 +- .../test_basic_correctness.py | 2 +- tests/basic_correctness/test_mem.py | 6 +- tests/benchmarks/test_bfcl_dataset.py | 18 +- ...est_custom_dataset_chat_template_kwargs.py | 3 +- .../test_endpoint_request_func_timing.py | 2 +- tests/benchmarks/test_sampling_params.py | 9 +- tests/benchmarks/test_skip_tokenizer_init.py | 3 +- tests/ci_envs.py | 4 +- tests/compile/backend.py | 6 +- tests/compile/conftest.py | 6 +- .../correctness_e2e/test_sequence_parallel.py | 11 +- .../compile/fullgraph/test_full_cudagraph.py | 10 +- .../compile/fullgraph/test_multiple_graphs.py | 3 +- tests/compile/fullgraph/test_simple.py | 6 +- tests/compile/fullgraph/test_toy_llama.py | 3 +- .../passes/distributed/test_async_tp.py | 24 +- tests/compile/passes/ir/test_clone_cleanup.py | 3 +- .../ir/test_inplace_functionalization.py | 9 +- .../test_double_aiter_rms_quant_fusion.py | 6 +- .../passes/test_fuse_mla_dual_rms_norm.py | 29 +- tests/compile/passes/test_fusion.py | 3 +- tests/compile/passes/test_fusion_attn.py | 3 +- .../passes/test_mla_attn_quant_fusion.py | 3 +- .../test_silu_mul_quant_manual_fusion.py | 3 +- tests/compile/passes/test_split_coalescing.py | 3 +- tests/compile/silly_attention.py | 12 +- tests/compile/test_aot_compile.py | 19 +- tests/compile/test_codegen.py | 30 +- tests/compile/test_compile_ranges.py | 6 +- tests/compile/test_config.py | 9 +- .../test_dynamic_shapes_compilation.py | 3 +- tests/compile/test_graph_partition.py | 43 +- ...test_inductor_fallback_allow_list_patch.py | 3 +- tests/compile/test_structured_logging.py | 1 + tests/config/test_config_utils.py | 2 +- tests/config/test_model_arch_config.py | 6 +- tests/config/test_mp_reducer.py | 4 +- tests/config/test_multimodal_config.py | 15 +- .../test_speculative_draft_hf_overrides.py | 12 +- ...eculative_draft_max_position_embeddings.py | 3 +- tests/cuda/test_cuda_context.py | 4 +- tests/cuda/test_platform_no_cuda_init.py | 3 +- ...stop_string_while_stop_model_terminates.py | 4 +- tests/distributed/conftest.py | 21 +- tests/distributed/test_context_parallel.py | 11 +- tests/distributed/test_eplb_algo.py | 24 +- tests/distributed/test_eplb_events.py | 5 +- tests/distributed/test_eplb_execute.py | 19 +- .../test_eplb_quant_scale_consistency.py | 12 +- tests/distributed/test_eplb_spec_decode.py | 8 +- tests/distributed/test_eplb_utils.py | 6 +- tests/distributed/test_events.py | 30 +- tests/distributed/test_expert_placement.py | 3 - tests/distributed/test_mnnvl_alltoall.py | 11 +- tests/distributed/test_mq_connect_ip.py | 12 +- tests/distributed/test_multiproc_executor.py | 6 +- tests/distributed/test_pipeline_parallel.py | 11 +- tests/distributed/test_pynccl.py | 3 +- tests/distributed/test_quick_all_reduce.py | 3 +- tests/distributed/test_ray_v2_executor.py | 4 +- tests/distributed/test_ray_v2_executor_e2e.py | 10 +- tests/distributed/test_sharded_rdt_plan.py | 189 +- .../distributed/test_sharded_rdt_producer.py | 119 +- tests/distributed/test_sharded_rdt_trainer.py | 36 +- tests/distributed/test_shm_broadcast.py | 14 +- tests/distributed/test_shm_buffer.py | 18 +- tests/distributed/test_utils.py | 3 +- tests/distributed/test_weight_transfer.py | 66 +- .../test_weight_transfer_nccl_uid.py | 9 +- tests/engine/test_arg_utils.py | 22 +- .../test_anthropic_messages_conversion.py | 34 +- .../cohere/test_serving_streaming.py | 6 +- .../test_api_server_process_manager.py | 12 +- tests/entrypoints/launchers/test_cli_args.py | 32 +- .../launchers/test_dp_supervisor.py | 24 +- tests/entrypoints/launchers/test_run_batch.py | 11 +- tests/entrypoints/launchers/test_shutdown.py | 1 - .../llm/offline_mode/test_offline_mode.py | 2 +- tests/entrypoints/llm/test_accuracy.py | 6 +- tests/entrypoints/llm/test_chat.py | 9 +- .../llm/test_mm_cache_external_injection.py | 3 +- .../chat_completion/test_audio_in_video.py | 9 +- ...chat_completion_with_mixed_audio_embeds.py | 3 +- ...chat_completion_with_mixed_image_embeds.py | 3 +- .../openai/chat_completion/test_vision.py | 9 +- .../test_batched_chat_completions.py | 2 +- .../openai/chat_completion/test_chat_error.py | 10 +- .../test_chat_logit_bias_validation.py | 9 +- .../test_completion_with_function_calling.py | 2 +- .../test_extra_content_fields.py | 10 +- .../chat_completion/test_logprob_token_ids.py | 3 +- .../chat_completion/test_serving_chat.py | 17 +- .../test_thinking_token_budget.py | 7 +- .../openai/completion/test_completion.py | 2 - .../completion/test_completion_error.py | 13 +- .../completion/test_token_in_token_out.py | 3 +- .../openai/correctness/test_lmeval.py | 5 +- .../parser/test_harmony_render_parity.py | 12 +- .../openai/parser/test_harmony_utils.py | 11 +- .../openai/responses/test_errors.py | 4 +- .../openai/responses/test_function_call.py | 3 +- .../responses/test_function_call_parsing.py | 12 +- .../openai/responses/test_harmony_utils.py | 3 +- .../test_response_input_to_harmony.py | 12 +- .../openai/responses/test_responses_utils.py | 9 +- .../responses/test_serving_responses.py | 50 +- .../openai/responses/test_simple.py | 3 +- .../openai/responses/test_streaming_events.py | 6 +- .../entrypoints/openai/test_openai_schema.py | 3 +- .../openai/test_render_token_offsets.py | 9 +- .../openai/test_return_token_ids.py | 3 +- .../entrypoints/openai/test_stop_token_ids.py | 3 +- tests/entrypoints/openai/utils.py | 3 +- .../pooling/basic/test_tiling_engine.py | 12 +- .../pooling/embed/test_cohere_online.py | 6 +- .../embed/test_cohere_openai_parity.py | 3 +- .../pooling/embed/test_io_processor.py | 3 +- .../pooling/embed/test_online_dimensions.py | 4 +- .../pooling/embed/test_online_long_text.py | 17 +- .../pooling/embed/test_protocol.py | 3 +- .../scale_out/derender/test_derender.py | 5 +- .../derender/test_derender_parity.py | 9 +- .../derender/test_derender_stream.py | 3 - .../scale_out/render/test_render.py | 13 +- .../render/test_render_multimodal.py | 3 - .../token_in_token_out/test_mm_serde.py | 3 +- .../token_in_token_out/test_protocol.py | 9 +- tests/entrypoints/serve/dev/rlhf/conftest.py | 4 +- .../serve/dev/rpc/test_collective_rpc.py | 8 +- tests/entrypoints/serve/dev/test_sleep.py | 3 +- .../test_error_sanitization.py | 6 +- .../test_validation_exception_handler.py | 15 +- .../serve/instrumentator/test_basic.py | 2 +- .../serve/instrumentator/test_metrics.py | 3 +- .../serve/lora/test_lora_adapters.py | 3 +- .../test_authentication_middleware.py | 9 +- .../middleware/test_optional_middleware.py | 3 +- .../test_sagemaker_handler_overrides.py | 18 +- .../serve/tokenize/test_tokenization_vlm.py | 3 +- .../tokenize/test_tokenize_then_chat_vlm.py | 7 +- .../entrypoints/serve/utils/test_api_utils.py | 12 +- .../serve/utils/test_request_logger.py | 3 +- .../serve/utils/test_sse_keep_alive.py | 24 +- tests/entrypoints/speech_to_text/conftest.py | 1 + .../test_transcription_api_correctness.py | 3 +- .../realtime/test_realtime_validation.py | 3 +- .../test_chunk_timestamp_offset.py | 5 +- .../test_transcription_validation_whisper.py | 3 +- .../tool_parsers/test_hermes_tool_parser.py | 4 - .../tool_parsers/test_openai_tool_parser.py | 11 +- .../entrypoints/unit_tests/test_chat_utils.py | 3 - tests/entrypoints/unit_tests/test_context.py | 22 +- .../test_weight_transfer_llm.py | 3 +- tests/evals/gpt_oss/conftest.py | 4 +- tests/evals/gpt_oss/test_gpqa_correctness.py | 4 +- tests/evals/gsm8k/gsm8k_eval.py | 10 +- tests/evals/gsm8k/test_gsm8k_correctness.py | 3 +- tests/evals/gsm8k/test_gsm8k_offloading.py | 3 +- tests/evals/mrcr/test_mrcr_correctness.py | 3 +- .../fusion/test_quant_activation_contract.py | 3 +- tests/ir/ir_test_utils.py | 4 +- tests/ir/test_op.py | 4 +- tests/kernels/attention/test_amx_mla.py | 9 +- .../attention/test_attention_selector.py | 6 +- .../attention/test_cutlass_mla_decode.py | 3 +- tests/kernels/attention/test_flashinfer.py | 4 +- .../attention/test_flashinfer_mla_decode.py | 2 +- .../test_flashinfer_trtllm_attention.py | 3 +- .../test_kimi_k3_mla_fused_epilogue.py | 3 +- .../kernels/attention/test_lightning_attn.py | 4 +- .../attention/test_merge_attn_states.py | 3 +- tests/kernels/attention/test_mha_attn.py | 10 +- tests/kernels/attention/test_minimax_m3.py | 42 +- ...test_mla_cross_layer_kernel_equivalence.py | 24 +- tests/kernels/attention/test_rocm_aiter_fa.py | 21 +- .../test_rocm_aiter_mla_op_registration.py | 2 +- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 6 +- .../attention/test_triton_decode_attention.py | 6 +- .../test_triton_unified_attention_diffkv.py | 4 +- .../attention/test_trtllm_kvfp8_dequant.py | 9 +- .../kernels/attention/test_xpu_mla_sparse.py | 4 +- tests/kernels/core/test_apply_rotary_emb.py | 6 +- tests/kernels/core/test_fused_embed_norm.py | 6 +- tests/kernels/core/test_fused_q_kv_rmsnorm.py | 6 +- .../kernels/core/test_fused_rms_norm_gated.py | 3 +- tests/kernels/core/test_minimax_reduce_rms.py | 1 - tests/kernels/core/test_mrope.py | 2 +- tests/kernels/core/test_opcheck.py | 4 +- tests/kernels/core/test_rocm_aiter_ops.py | 3 +- tests/kernels/core/test_rotary_embedding.py | 4 +- .../test_rotary_embedding_mla_cache_fused.py | 4 +- .../core/test_vocab_parallel_embedding.py | 3 +- tests/kernels/helion/test_autotune.py | 3 +- tests/kernels/helion/test_config_manager.py | 3 +- ...test_dynamic_per_token_scaled_fp8_quant.py | 2 +- .../kernels/helion/test_fused_qk_norm_rope.py | 2 +- tests/kernels/helion/test_helion_available.py | 3 +- .../helion/test_per_token_group_fp8_quant.py | 2 +- tests/kernels/helion/test_register.py | 13 +- .../test_rms_norm_dynamic_per_token_quant.py | 2 +- .../helion/test_rms_norm_per_block_quant.py | 2 +- tests/kernels/helion/utils.py | 2 +- tests/kernels/ir/test_aiter_norm_dispatch.py | 2 +- tests/kernels/ir/test_ir_ops.py | 3 +- tests/kernels/ir/test_layernorm.py | 3 +- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 9 +- tests/kernels/mamba/test_causal_conv1d.py | 9 +- tests/kernels/mamba/test_mamba_ssm.py | 5 +- tests/kernels/mamba/test_mamba_ssm_configs.py | 12 +- tests/kernels/mamba/test_mamba_ssm_ssd.py | 5 +- ...ayssm_prefill_decode_equivalence_mamba2.py | 3 +- .../test_replayssm_standard_decode_mamba2.py | 6 +- tests/kernels/mamba/utils.py | 5 +- .../moe/modular_kernel_tools/common.py | 11 +- tests/kernels/moe/parallel_utils.py | 4 +- tests/kernels/moe/test_b12x.py | 3 +- tests/kernels/moe/test_batched_deepgemm.py | 1 - tests/kernels/moe/test_batched_moe.py | 9 +- tests/kernels/moe/test_block_int8.py | 6 +- .../moe/test_count_expert_num_tokens.py | 4 +- tests/kernels/moe/test_cpu_fused_moe.py | 7 +- tests/kernels/moe/test_cpu_int4_moe.py | 1 + tests/kernels/moe/test_cpu_quant_fused_moe.py | 1 + tests/kernels/moe/test_cutedsl_moe.py | 8 +- tests/kernels/moe/test_deepep_deepgemm_moe.py | 16 +- tests/kernels/moe/test_deepep_moe.py | 8 +- tests/kernels/moe/test_deepep_v2_moe.py | 3 +- tests/kernels/moe/test_deepgemm.py | 15 +- tests/kernels/moe/test_flashinfer.py | 3 +- .../moe/test_flashinfer_cutedsl_nvfp4_moe.py | 2 +- tests/kernels/moe/test_fused_topk.py | 2 +- .../moe/test_gpt_oss_triton_kernels.py | 3 +- tests/kernels/moe/test_grouped_topk.py | 2 +- .../moe/test_marlin_vs_trtllm_mxint4.py | 6 +- .../moe/test_modular_kernel_combinations.py | 3 +- .../moe/test_modular_oai_triton_moe.py | 4 +- tests/kernels/moe/test_moe.py | 11 +- .../kernels/moe/test_moe_align_block_size.py | 10 +- tests/kernels/moe/test_moe_kernel_oracle.py | 3 +- tests/kernels/moe/test_moe_layer.py | 10 +- .../kernels/moe/test_moe_permute_unpermute.py | 2 +- tests/kernels/moe/test_mxfp4_moe.py | 13 +- .../moe/test_mxfp8_aiter_backend_selection.py | 8 +- tests/kernels/moe/test_ocp_mx_moe.py | 10 +- tests/kernels/moe/test_rocm_aiter_moe.py | 33 +- .../test_routed_experts_capture_monolithic.py | 24 +- tests/kernels/moe/test_routing.py | 18 +- tests/kernels/moe/test_routing_simulator.py | 9 +- .../test_shared_fused_moe_routed_transform.py | 3 +- .../moe/test_silu_mul_fp8_quant_deep_gemm.py | 10 +- ..._mul_per_token_group_quant_fp8_colmajor.py | 7 +- tests/kernels/moe/test_triton_moe_ptpc_fp8.py | 7 +- tests/kernels/moe/test_trtllm_bf16_moe.py | 3 +- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 3 +- .../moe/test_unquantized_backend_selection.py | 6 +- tests/kernels/moe/test_zero_expert_moe.py | 3 +- tests/kernels/moe/utils.py | 3 +- tests/kernels/quant_utils.py | 3 +- tests/kernels/quantization/nvfp4_utils.py | 1 + .../test_aiter_hipb_mm_linear_kernel.py | 3 +- .../quantization/test_cpu_fp8_scaled_mm.py | 1 + .../quantization/test_cutlass_scaled_mm.py | 2 +- .../kernels/quantization/test_cutlass_w4a8.py | 3 +- .../quantization/test_cutlass_w4a8_moe.py | 16 +- .../quantization/test_fp8_min_max_helper.py | 3 +- .../kernels/quantization/test_int8_kernel.py | 7 +- .../quantization/test_marlin_tile_padding.py | 15 +- .../test_mxfp4_kernel_selection.py | 20 +- .../quantization/test_mxfp4_triton_ep.py | 6 +- .../test_mxfp6_kernel_selection.py | 11 +- .../quantization/test_nvfp4_emulation.py | 3 +- .../kernels/quantization/test_nvfp4_quant.py | 3 +- .../test_per_token_group_quant.py | 14 +- .../quantization/test_quantized_embedding.py | 3 +- .../test_rocm_aiter_grouped_quant.py | 6 +- tests/kernels/quantization/test_rocm_mxfp4.py | 5 +- .../test_scaled_mm_kernel_selection.py | 2 +- .../quantization/test_triton_scaled_mm.py | 2 +- .../test_apply_repetition_penalties.py | 6 +- tests/kernels/test_awq_int4_to_int8.py | 4 +- tests/kernels/test_cache_kernels.py | 4 +- tests/kernels/test_compressor_kv_cache.py | 18 +- tests/kernels/test_cp_gather_fp8.py | 4 +- tests/kernels/test_deepseek_v4_cpu_kernels.py | 57 +- tests/kernels/test_engram.py | 9 +- tests/kernels/test_fla_layernorm_guard.py | 3 +- tests/kernels/test_flex_attention.py | 2 +- tests/kernels/test_fp32_router_gemm.py | 3 +- .../test_fused_deepseek_v32_norm_rope.py | 3 +- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 9 +- .../test_fused_indexer_q_rope_quant.py | 4 +- .../kernels/test_fused_inv_rope_fp8_quant.py | 17 +- tests/kernels/test_mhc_kernels.py | 10 +- tests/kernels/test_minimax_m3_amd_ops.py | 12 +- tests/kernels/test_shuffle_rows.py | 8 +- tests/kernels/test_top_k_per_row.py | 51 +- tests/kernels/utils.py | 262 +- tests/lora/conftest.py | 4 +- tests/lora/test_add_lora.py | 3 +- tests/lora/test_default_mm_loras.py | 7 +- tests/lora/test_fused_moe_lora_kernel.py | 22 +- tests/lora/test_layers.py | 9 +- tests/lora/test_lora_functions.py | 4 +- tests/lora/test_mixtral.py | 2 +- tests/lora/test_moe_lora_align_sum.py | 6 +- tests/lora/test_moe_lora_ep_load.py | 9 +- tests/lora/test_punica_ops.py | 23 +- tests/lora/test_punica_ops_fp8.py | 6 +- tests/lora/test_punica_xpu_ops.py | 12 +- tests/lora/test_qwen3_unembed.py | 3 +- tests/lora/test_qwen3_with_multi_loras.py | 12 +- tests/lora/test_qwenvl.py | 9 +- tests/lora/test_whisper.py | 3 +- tests/lora/utils.py | 3 +- .../layers/test_fused_shared_expert.py | 7 +- .../layers/test_mla_short_prefill_indexer.py | 36 +- .../model_loader/test_ep_weight_filter.py | 9 +- .../model_loader/test_reload.py | 9 +- tests/model_executor/test_deep_gemm_warmup.py | 8 +- .../test_model_load_with_params.py | 12 +- tests/model_executor/test_oink_integration.py | 1 - tests/model_executor/test_qwen3_omni.py | 4 +- tests/model_executor/test_qwen3_vl_mrope.py | 4 +- .../test_routed_experts_capture.py | 7 +- tests/model_executor/test_weight_utils.py | 32 +- tests/models/inkling/test_mm_towers.py | 3 +- .../kimi_k3/test_amd_latent_moe_runner.py | 3 +- .../kimi_k3/test_aux_attn_res_stream.py | 6 +- tests/models/kimi_k3/test_kda_metadata.py | 3 +- .../models/language/generation/test_gemma.py | 6 +- .../models/language/generation/test_hybrid.py | 12 +- .../language/generation/test_mistral.py | 3 +- .../pooling/test_max_tokens_per_doc.py | 4 +- .../pooling/test_splade_sparse_pooler.py | 3 +- .../multimodal/generation/test_interleaved.py | 4 +- .../multimodal/generation/test_maverick.py | 24 +- .../generation/test_moondream3_moe.py | 3 +- .../generation/test_nemotron_parse.py | 3 +- .../multimodal/generation/test_qwen2_vl.py | 6 +- .../generation/test_transformers_audio.py | 3 +- .../multimodal/generation/test_ultravox.py | 1 - .../multimodal/generation/test_voxtral.py | 1 - .../generation/vlm_utils/builders.py | 3 +- .../generation/vlm_utils/case_filtering.py | 1 + .../multimodal/generation/vlm_utils/core.py | 2 +- .../generation/vlm_utils/custom_inputs.py | 2 + .../generation/vlm_utils/model_utils.py | 10 +- .../multimodal/pooling/test_dse_qwen2_vl.py | 2 +- .../pooling/test_jinavl_reranker.py | 13 +- .../pooling/test_llama_nemotron_vl.py | 3 +- .../processing/test_audio_in_video.py | 9 +- .../multimodal/processing/test_common.py | 7 +- .../processing/test_deepseek_ocr.py | 9 +- .../multimodal/processing/test_glm4_1v.py | 3 +- .../multimodal/processing/test_llava_next.py | 3 +- .../processing/test_llava_next_video.py | 6 +- .../processing/test_llava_onevision.py | 3 +- .../multimodal/processing/test_molmo2.py | 3 +- .../processing/test_qwen2_5_omni_embed.py | 29 +- .../multimodal/processing/test_qwen2_vl.py | 1 - .../multimodal/processing/test_qwen3_omni.py | 6 +- .../multimodal/processing/test_qwen3_vl.py | 12 +- .../processing/test_transformers_audio.py | 3 +- .../processing/test_transformers_image.py | 24 +- .../processing/transformers_backend.py | 6 +- tests/models/multimodal/test_cohere_asr.py | 3 +- tests/models/multimodal/test_mapping.py | 13 +- .../multimodal/test_nano_nemotron_vl.py | 9 +- tests/models/quantization/test_gpt_oss.py | 6 +- tests/models/qwen4_exp/test_config.py | 3 +- tests/models/registry.py | 7 +- tests/models/test_deepseek_v4_fi_moe_ep.py | 17 +- tests/models/test_deepseek_v4_mega_moe.py | 6 +- tests/models/test_dspark_mla.py | 3 +- tests/models/test_initialization.py | 5 +- tests/models/test_qwen3_5_mtp_config.py | 9 +- tests/models/test_utils.py | 3 +- tests/models/test_vision.py | 14 +- .../models/transformers/fusers/test_linear.py | 69 +- tests/models/transformers/fusers/test_mla.py | 15 +- tests/models/transformers/fusers/test_moe.py | 18 +- .../transformers/fusers/test_rms_norm.py | 27 +- .../transformers/test_layer_registry.py | 6 +- tests/models/utils.py | 8 +- tests/multimodal/media/test_audio.py | 22 +- tests/multimodal/media/test_connector.py | 3 +- tests/multimodal/media/test_image.py | 1 - tests/multimodal/media/test_video.py | 28 +- tests/multimodal/test_audio.py | 6 - .../test_embedding_shape_validation.py | 12 +- .../test_embedding_shape_validation_unit.py | 13 +- tests/multimodal/test_hasher.py | 3 +- tests/multimodal/test_image.py | 3 +- tests/multimodal/test_parse.py | 6 +- tests/multimodal/test_processing.py | 21 +- tests/multimodal/test_registry.py | 6 +- .../test_sparse_tensor_validation_unit.py | 6 +- tests/multimodal/test_utils.py | 3 +- tests/multimodal/test_video.py | 45 +- tests/parser/cohere/test_reasoning.py | 3 +- tests/parser/cohere/utils.py | 9 +- tests/parser/engine/conftest.py | 1 + tests/parser/engine/replay_harness.py | 1 + tests/parser/engine/test_deepseek_v32.py | 3 +- tests/parser/engine/test_deepseek_v4.py | 18 +- tests/parser/engine/test_delegating_replay.py | 3 +- tests/parser/engine/test_engine.py | 57 +- .../engine/test_gemma4_streaming_reasoning.py | 27 +- tests/parser/engine/test_glm47_moe.py | 6 +- tests/parser/engine/test_inkling.py | 30 +- tests/parser/engine/test_nemotron_v3.py | 18 +- tests/parser/engine/test_parser_engine.py | 69 +- tests/parser/engine/test_qwen3.py | 15 +- tests/parser/engine/test_qwen3_reasoning.py | 18 +- tests/parser/engine/test_replay.py | 12 +- tests/parser/engine/test_seed_oss.py | 6 +- .../engine/test_ufffd_reasoning_transition.py | 3 +- tests/parser/engine/trace_builder.py | 3 +- tests/parser/mistral/test_tool_calls.py | 18 +- tests/parser/test_streaming.py | 36 +- .../query_embedding_processor.py | 2 +- .../prithvi_io_processor/prithvi_processor.py | 5 +- .../prithvi_io_processor/types.py | 4 +- .../dummy_stat_logger/dummy_stat_logger.py | 3 +- .../gguf/test_gguf_plugin_generate.py | 4 +- tests/plugins_tests/test_endpoint_plugins.py | 12 +- .../test_io_processor_plugins.py | 3 +- tests/quantization/test_auto_round.py | 15 +- tests/quantization/test_compressed_tensors.py | 4 +- tests/quantization/test_fp8.py | 3 +- tests/quantization/test_gfx950_moe.py | 6 +- .../test_humming_mxfp4_block_fp8.py | 6 +- tests/quantization/test_int8_moe_oracle.py | 6 +- tests/quantization/test_modelopt.py | 12 +- tests/quantization/test_online.py | 5 +- tests/quantization/test_online_mxfp4.py | 12 +- tests/quantization/test_per_token_kv_cache.py | 3 +- .../test_register_quantization_config.py | 2 - .../test_base_thinking_reasoning_parser.py | 12 +- .../test_kimi_k3_reasoning_parser.py | 12 +- .../test_chat_utils_prompt_embeds.py | 9 +- tests/renderers/test_completions.py | 2 +- tests/renderers/test_gemma4_chat_template.py | 31 +- tests/renderers/test_hf.py | 13 +- .../test_process_multi_modal_uuids.py | 3 +- .../test_prompt_embeds_payload_validation.py | 5 +- .../test_sparse_tensor_concurrent_race.py | 3 +- .../test_sparse_tensor_validation.py | 48 +- tests/renderers/test_token_offsets.py | 15 +- tests/spec_decode/test_custom_proposer.py | 2 + tests/test_access_log_filter.py | 4 +- tests/test_cmake_utils.py | 5 +- tests/test_config.py | 36 +- tests/test_logger.py | 24 +- tests/test_ray_env.py | 6 +- tests/tokenizers_/test_deepseek_v4.py | 3 +- tests/tokenizers_/test_hf.py | 3 +- tests/tool_parsers/common_tests.py | 2 + .../tool_parsers/test_apertus_tool_parser.py | 12 +- .../test_deepseekv32_tool_parser.py | 30 +- tests/tool_parsers/test_gemma4_tool_parser.py | 3 +- .../test_gigachat3_tool_parser.py | 7 +- .../test_internlm2_tool_parser.py | 1 - .../tool_parsers/test_kimi_k2_tool_parser.py | 1 + .../test_kimi_k3_named_tool_choice.py | 3 +- tests/tool_parsers/test_lfm2_tool_parser.py | 27 +- .../test_llama3_json_tool_parser.py | 2 +- .../test_llama4_pythonic_tool_parser.py | 2 +- tests/tool_parsers/test_olmo3_tool_parser.py | 2 +- .../tool_parsers/test_phi4mini_tool_parser.py | 1 - .../tool_parsers/test_pythonic_tool_parser.py | 2 +- .../test_qwen3coder_tool_parser.py | 14 +- .../tool_parsers/test_step3p5_tool_parser.py | 14 +- tests/tool_use/test_chat_completions.py | 3 +- .../test_gemma4_responses_adjust_request.py | 10 +- tests/tool_use/test_muse_glimmer.py | 3 +- tests/tool_use/test_parallel_tool_calls.py | 5 +- tests/tool_use/test_tool_choice_required.py | 3 +- .../processors/test_glm5next.py | 6 +- tests/transformers_utils/test_config.py | 9 +- .../test_hf_overrides_model_type.py | 4 +- tests/utils.py | 61 +- tests/utils_/__init__.py | 3 +- tests/utils_/test_gpu_sync_debug.py | 15 +- tests/utils_/test_jsontree.py | 1 - tests/utils_/test_network_utils.py | 15 +- tests/utils_/test_numa_utils.py | 18 +- tests/v1/attention/test_attention_backends.py | 8 +- .../test_attention_backends_selection.py | 3 +- .../v1/attention/test_attention_splitting.py | 31 +- tests/v1/attention/test_dcp_a2a_pack_mask.py | 3 +- .../test_dspark_noncausal_sparse_mla.py | 3 +- .../test_flashinfer_mla_sparse_sm90.py | 3 +- .../v1/attention/test_gdn_metadata_builder.py | 3 +- .../v1/attention/test_indexer_dcp_localize.py | 18 +- .../attention/test_kpool_tail_slot_mapping.py | 21 +- .../test_mamba_update_block_table.py | 27 +- tests/v1/attention/test_mla_backends.py | 9 +- .../test_mla_prefill_quant_output.py | 9 +- .../v1/attention/test_mla_prefill_selector.py | 3 +- .../test_replayssm_metadata_builder.py | 3 +- .../test_rocm_aiter_mla_mtp_split.py | 3 - .../test_rocm_attention_backends_selection.py | 3 +- .../test_sparse_indexer_decode_seq_lens.py | 12 +- .../v1/attention/test_sparse_mla_backends.py | 31 +- .../test_trtllm_attention_integration.py | 10 +- tests/v1/attention/utils.py | 1 - .../test_partial_prefix_cache_hits.py | 51 +- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_contiguous_kv_packing.py | 18 +- tests/v1/core/test_encoder_cache_manager.py | 9 +- tests/v1/core/test_kv_cache_utils.py | 43 +- tests/v1/core/test_kv_sharing.py | 9 +- tests/v1/core/test_output.py | 9 +- tests/v1/core/test_prefix_caching.py | 120 +- tests/v1/core/test_repetition_detection.py | 42 +- tests/v1/core/test_scheduler.py | 120 +- .../v1/core/test_swa_inflight_window_free.py | 12 +- tests/v1/core/utils.py | 1 + .../v1/cudagraph/test_breakable_cudagraph.py | 13 +- tests/v1/cudagraph/test_cudagraph_dispatch.py | 1 - tests/v1/determinism/test_batch_invariance.py | 16 +- .../test_matmul_batch_invariant.py | 11 +- .../test_online_batch_invariance.py | 3 +- .../test_rms_norm_batch_invariant.py | 25 +- tests/v1/distributed/test_async_llm_dp.py | 16 +- tests/v1/distributed/test_dbo.py | 7 +- tests/v1/distributed/test_external_lb_dp.py | 3 +- tests/v1/distributed/test_hybrid_lb_dp.py | 3 +- tests/v1/distributed/test_internal_lb_dp.py | 11 +- tests/v1/distributed/test_pp_dp_v2.py | 14 +- tests/v1/e2e/general/test_async_scheduling.py | 9 +- .../test_attention_backend_per_kind.py | 3 +- tests/v1/e2e/general/test_context_length.py | 11 +- .../test_correctness_sliding_window.py | 7 +- .../general/test_kv_sharing_fast_prefill.py | 4 +- tests/v1/e2e/general/test_min_tokens.py | 33 +- .../general/test_pooling_chunked_prefill.py | 2 - tests/v1/e2e/general/test_sharded_sampling.py | 6 +- tests/v1/e2e/general/test_streaming_input.py | 3 +- .../mtp_other/test_synthetic.py | 3 +- .../e2e/spec_decode/acceptance_rates/utils.py | 3 +- .../draft_model/test_draft_model.py | 22 +- .../e2e/spec_decode/draft_model/test_lora.py | 6 +- tests/v1/e2e/spec_decode/eagle/utils.py | 3 +- tests/v1/e2e/spec_decode/mtp/test_mtp.py | 3 +- .../ngram_suffix/test_ngram_suffix.py | 6 +- .../speculators/test_speculators.py | 3 +- .../e2e/spec_decode/test_sharded_sampling.py | 3 +- tests/v1/e2e/test_hybrid_chunked_prefill.py | 4 +- .../integration/test_epd_correctness.py | 5 +- .../cpu/scheduler/test_embedding_cache.py | 3 +- .../unit/cpu/scheduler/test_scheduler.py | 36 +- .../unit/cpu/test_ec_shared_region.py | 12 +- .../unit/cpu/worker/test_worker.py | 18 +- .../unit/test_ec_cpu_connector.py | 3 +- .../unit/test_ec_example_connector.py | 4 +- .../unit/test_epd_proxy_round_robin.py | 3 +- tests/v1/ec_connector/unit/test_metadata.py | 5 +- .../unit/test_scheduler_nixl_producer.py | 3 +- tests/v1/ec_connector/unit/test_session.py | 7 +- tests/v1/ec_connector/unit/utils.py | 1 + tests/v1/engine/conftest.py | 7 +- tests/v1/engine/test_abort_final_step.py | 10 +- tests/v1/engine/test_async_llm.py | 2 - .../engine/test_core_engine_actor_manager.py | 3 +- tests/v1/engine/test_engine_core.py | 28 +- tests/v1/engine/test_engine_core_client.py | 13 +- .../v1/engine/test_fast_incdec_prefix_err.py | 3 +- tests/v1/engine/test_llm_engine.py | 3 +- tests/v1/engine/test_logprobs_processor.py | 3 +- tests/v1/engine/test_output_processor.py | 5 +- .../engine/test_preprocess_error_handling.py | 1 - tests/v1/engine/utils.py | 22 +- .../predictable_llama.py | 1 + .../nixl_integration/test_disagg_accuracy.py | 7 +- .../test_multi_connector_edge_cases.py | 3 +- .../nixl_integration/toy_proxy_server.py | 16 +- .../test_canonical_mapping.py | 3 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/offloading_connector/test_metrics.py | 6 +- .../offloading_connector/test_scheduler.py | 48 +- .../unit/offloading_connector/test_worker.py | 7 +- .../unit/offloading_connector/utils.py | 10 +- .../unit/test_bidirectional_kv_transfer.py | 57 +- .../unit/test_cache_pollution_prevention.py | 8 +- tests/v1/kv_connector/unit/test_config.py | 18 +- .../unit/test_decode_bench_connector.py | 3 +- .../unit/test_error_propagation.py | 6 +- .../unit/test_example_connector.py | 14 +- .../unit/test_flexkv_connector.py | 9 +- .../unit/test_handshake_pp_aggregation.py | 9 +- .../v1/kv_connector/unit/test_hf3fs_client.py | 9 +- .../kv_connector/unit/test_hf3fs_connector.py | 7 +- .../unit/test_hf3fs_metadata_server.py | 9 +- .../unit/test_invalid_blocks_correctness.py | 16 +- .../unit/test_lmcache_connector.py | 6 +- .../unit/test_mooncake_connector.py | 32 +- .../unit/test_mooncake_connector_hma.py | 2 - .../kv_connector/unit/test_mooncake_stats.py | 3 +- .../unit/test_mooncake_store_connector.py | 9 +- .../unit/test_mooncake_store_coordinator.py | 54 +- .../unit/test_mooncake_store_hma_e2e.py | 24 +- .../unit/test_mooncake_store_scheduler.py | 6 +- .../unit/test_mooncake_store_worker.py | 62 +- .../unit/test_moriio_connector.py | 24 +- .../unit/test_moriio_routing_fairness.py | 6 +- .../kv_connector/unit/test_multi_connector.py | 22 +- .../kv_connector/unit/test_nixl_connector.py | 66 +- .../unit/test_nixl_connector_hma.py | 54 +- .../unit/test_nixl_desc_geometry.py | 30 +- .../unit/test_nixl_push_connector.py | 105 +- .../unit/test_nixl_rocm_gpu_mem_diag.py | 3 +- .../unit/test_nixl_simple_cpu_offload.py | 9 +- .../unit/test_offloading_connector.py | 14 +- .../unit/test_remote_decode_lifecycle.py | 4 - .../unit/test_remote_prefill_lifecycle.py | 18 +- .../unit/test_simple_cpu_offload_connector.py | 3 +- tests/v1/kv_connector/unit/test_tp_mapping.py | 6 +- tests/v1/kv_connector/unit/utils.py | 1 - .../kv_offload/cpu/policies/test_factory.py | 15 +- .../kv_offload/cpu/test_canonical_layout.py | 9 +- tests/v1/kv_offload/cpu/test_manager.py | 41 +- .../cpu/test_shared_offload_region.py | 42 +- tests/v1/kv_offload/test_file_mapper.py | 3 +- .../tiering/p2p/p2p_connector_proxy.py | 3 +- .../tiering/p2p/test_data_transport.py | 8 +- .../v1/kv_offload/tiering/p2p/test_manager.py | 93 +- .../kv_offload/tiering/p2p/test_sessions.py | 122 +- .../kv_offload/tiering/test_async_lookup.py | 6 +- tests/v1/kv_offload/tiering/test_factory.py | 3 +- tests/v1/kv_offload/tiering/test_fs_tier.py | 27 +- tests/v1/kv_offload/tiering/test_obj_tier.py | 9 +- .../tiering/test_tiering_offloading.py | 36 +- .../v1/logits_processors/test_correctness.py | 30 +- .../logits_processors/test_custom_offline.py | 11 +- .../logits_processors/test_custom_online.py | 7 +- tests/v1/logits_processors/utils.py | 19 +- tests/v1/metrics/test_engine_logger_apis.py | 11 +- tests/v1/metrics/test_perf_metrics.py | 18 +- tests/v1/metrics/test_ray_metrics.py | 22 +- .../sample/test_batched_count_greater_than.py | 3 +- tests/v1/sample/test_logprobs.py | 20 +- tests/v1/sample/test_rejection_sampler.py | 26 +- tests/v1/sample/test_sampler.py | 19 +- tests/v1/sample/test_sampling_params_e2e.py | 8 - tests/v1/sample/test_topk_topp_sampler.py | 9 +- tests/v1/sample/utils.py | 18 +- tests/v1/shutdown/test_delete.py | 2 + tests/v1/shutdown/test_startup_error.py | 1 - tests/v1/shutdown/utils.py | 2 +- tests/v1/simple_kv_offload/test_kv_events.py | 9 +- tests/v1/simple_kv_offload/test_scheduler.py | 19 +- .../v1/spec_decode/test_acceptance_length.py | 3 +- .../test_backup_token_async_spec.py | 3 +- tests/v1/spec_decode/test_dynamic_sd_cug.py | 6 - tests/v1/spec_decode/test_eagle.py | 22 +- .../test_eagle_draft_attn_metadata.py | 9 +- .../v1/spec_decode/test_eagle_step_kernel.py | 3 +- .../spec_decode/test_extract_hidden_states.py | 6 +- .../v1/spec_decode/test_llm_base_proposer.py | 6 +- tests/v1/spec_decode/test_max_len.py | 3 +- tests/v1/spec_decode/test_mtp.py | 4 +- .../test_rejection_sampler_utils.py | 22 +- .../test_speculators_correctness.py | 6 +- .../v1/spec_decode/test_speculators_eagle3.py | 3 +- .../test_gpu_model_runner_streaming.py | 1 - .../test_gpu_model_runner_v2_streaming.py | 1 - .../test_scheduler_streaming.py | 3 +- tests/v1/structured_output/test_validation.py | 12 +- tests/v1/test_kv_cache_spec_registry.py | 16 +- tests/v1/test_outputs.py | 24 +- tests/v1/test_serial_utils.py | 5 +- tests/v1/utils.py | 3 + .../test_dsv4_packed_zeroer_geometry.py | 3 +- tests/v1/worker/test_encoder_runner.py | 15 +- tests/v1/worker/test_gpu_bad_words.py | 9 +- tests/v1/worker/test_gpu_batch_shard.py | 27 +- tests/v1/worker/test_gpu_block_table.py | 6 +- tests/v1/worker/test_gpu_gumbel_sample.py | 2 +- tests/v1/worker/test_gpu_input_batch.py | 13 +- tests/v1/worker/test_gpu_logit_bias.py | 3 +- tests/v1/worker/test_gpu_model_runner.py | 14 +- .../worker/test_gpu_model_runner_mm_gather.py | 9 +- tests/v1/worker/test_gpu_model_runner_v2.py | 13 +- ...gpu_model_runner_v2_cudagraph_profiling.py | 3 +- tests/v1/worker/test_gpu_profiler.py | 10 +- .../worker/test_gpu_rejection_sampler_i64.py | 3 +- tests/v1/worker/test_gpu_thinking_budget.py | 3 +- tests/v1/worker/test_gpu_ubatch_slicing.py | 3 +- tests/v1/worker/test_gpu_worker.py | 15 +- tests/v1/worker/test_kv_block_zeroer.py | 3 +- tests/v1/worker/test_mamba_utils.py | 45 +- tests/v1/worker/test_mixed_warmup_gate.py | 3 +- tests/v1/worker/test_prompt_embeds_state.py | 12 +- .../v1/worker/test_worker_memory_snapshot.py | 1 + tests/v1/worker/test_workspace.py | 3 +- .../vllm_test_utils/__init__.py | 3 +- .../vllm_test_utils/vllm_test_utils/blame.py | 3 +- .../vllm_test_utils/monitor.py | 3 +- tests/weight_loading/test_weight_loading.py | 5 +- tools/generate_cmake_presets.py | 1 - tools/generate_versions_json.py | 3 +- tools/install_nixl_from_source_ubuntu.py | 3 +- tools/pre_commit/check_spdx_header.py | 8 +- .../pre_commit/generate_nightly_torch_test.py | 3 +- tools/pre_commit/mypy.py | 12 +- tools/pre_commit/validate_config.py | 7 +- .../nsys_profile_tools/gputrc2graph.py | 27 +- tools/recipes/recipe_json_to_vllm_config.py | 7 +- tools/report_build_time_ninja.py | 10 +- tools/vllm-rocm/pin_rocm_dependencies.py | 14 +- vllm/__init__.py | 2 +- vllm/_aiter_ops.py | 43 +- vllm/_custom_ops.py | 110 +- vllm/_xpu_ops.py | 3 + vllm/assets/base.py | 3 +- vllm/assets/image.py | 8 +- vllm/assets/video.py | 6 +- .../datasets/create_txt_slices_dataset.py | 7 +- vllm/benchmarks/datasets/datasets.py | 163 +- vllm/benchmarks/datasets/utils.py | 8 +- vllm/benchmarks/lib/endpoint_request_func.py | 8 +- vllm/benchmarks/lib/ready_checker.py | 5 +- vllm/benchmarks/lib/utils.py | 7 +- vllm/benchmarks/mm_processor.py | 20 +- vllm/benchmarks/plot.py | 13 +- vllm/benchmarks/serve.py | 10 +- vllm/benchmarks/startup.py | 10 +- vllm/benchmarks/sweep/param_sweep.py | 10 +- vllm/benchmarks/sweep/plot.py | 3 +- vllm/benchmarks/throughput.py | 9 +- vllm/compilation/backends.py | 21 +- vllm/compilation/base_static_graph.py | 12 +- vllm/compilation/caching.py | 9 +- vllm/compilation/codegen.py | 2 + vllm/compilation/compiler_interface.py | 32 +- vllm/compilation/cuda_graph.py | 2 +- vllm/compilation/decorators.py | 14 +- vllm/compilation/partition_rules.py | 5 +- .../passes/fusion/act_quant_fusion.py | 17 +- .../passes/fusion/allreduce_rms_fusion.py | 21 +- .../passes/fusion/attn_quant_fusion.py | 9 +- .../passes/fusion/matcher_utils.py | 2 +- .../passes/fusion/mla_attn_quant_fusion.py | 12 +- .../passes/fusion/qk_norm_rope_fusion.py | 3 +- .../fusion/qk_norm_rope_kvcache_fusion.py | 6 +- .../passes/fusion/rms_quant_fusion.py | 6 +- .../passes/fusion/rocm_aiter_fusion.py | 39 +- .../passes/fusion/rope_kvcache_fusion.py | 12 +- .../passes/fusion/sequence_parallelism.py | 9 +- vllm/compilation/passes/inductor_pass.py | 20 +- .../passes/ir/clone_elimination.py | 3 +- .../passes/ir/inplace_functionalization.py | 3 +- vllm/compilation/passes/ir/lowering_pass.py | 7 +- vllm/compilation/passes/pass_manager.py | 9 +- .../passes/utility/fix_functionalization.py | 21 +- .../passes/utility/noop_elimination.py | 7 +- .../passes/utility/post_cleanup.py | 3 +- .../passes/utility/scatter_split_replace.py | 3 +- .../passes/utility/split_coalescing.py | 9 +- vllm/compilation/passes/vllm_inductor_pass.py | 20 +- vllm/compilation/piecewise_backend.py | 3 +- vllm/compilation/wrapper.py | 10 +- vllm/config/attention.py | 3 +- vllm/config/cache.py | 3 +- vllm/config/compilation.py | 39 +- vllm/config/device.py | 3 +- vllm/config/ec_manager_config.py | 3 +- vllm/config/ec_transfer.py | 3 +- vllm/config/kernel.py | 22 +- vllm/config/kv_transfer.py | 3 +- vllm/config/load.py | 3 +- vllm/config/lora.py | 3 +- vllm/config/model.py | 36 +- vllm/config/model_arch.py | 4 +- vllm/config/multimodal.py | 15 +- vllm/config/observability.py | 6 +- vllm/config/offload.py | 3 +- vllm/config/parallel.py | 10 +- vllm/config/pooler.py | 3 +- vllm/config/profiler.py | 3 +- vllm/config/reasoning.py | 6 +- vllm/config/scheduler.py | 7 +- vllm/config/speculative.py | 17 +- vllm/config/structured_outputs.py | 3 +- vllm/config/utils.py | 28 +- vllm/config/vllm.py | 47 +- vllm/device_allocator/cumem.py | 30 +- vllm/device_allocator/sleep_mode_backend.py | 15 +- vllm/device_allocator/xpumem.py | 1 + .../aiter_custom_all_reduce.py | 4 +- .../device_communicators/all2all.py | 52 +- .../device_communicators/all_reduce_utils.py | 7 +- .../base_device_communicator.py | 18 +- .../device_communicators/cpu_communicator.py | 16 +- .../device_communicators/cuda_communicator.py | 12 +- .../device_communicators/custom_all_reduce.py | 7 +- .../flashinfer_all_reduce.py | 7 +- .../device_communicators/pynccl.py | 4 +- .../device_communicators/quick_all_reduce.py | 11 +- .../device_communicators/ray_communicator.py | 30 +- .../device_communicators/shm_broadcast.py | 20 +- .../shm_object_storage.py | 41 +- .../device_communicators/xpu_communicator.py | 10 +- .../ec_transfer/ec_connector/base.py | 65 +- .../ec_connector/cpu/ec_shared_region.py | 3 +- .../ec_connector/cpu/scheduler/__init__.py | 3 + .../cpu/scheduler/embedding_cache.py | 3 +- .../ec_connector/example_connector.py | 22 +- .../ec_transfer/ec_transfer_state.py | 5 +- vllm/distributed/eplb/async_worker.py | 4 +- vllm/distributed/eplb/eplb_communicator.py | 12 +- vllm/distributed/eplb/eplb_state.py | 51 +- vllm/distributed/eplb/eplb_utils.py | 13 +- vllm/distributed/eplb/policy/abstract.py | 7 +- vllm/distributed/eplb/policy/default.py | 30 +- vllm/distributed/eplb/rebalance_execute.py | 33 +- vllm/distributed/kv_events.py | 38 +- .../kv_transfer/kv_connector/factory.py | 1 + .../kv_transfer/kv_connector/utils.py | 23 +- .../kv_transfer/kv_connector/v1/base.py | 140 +- .../kv_connector/v1/decode_bench_connector.py | 32 +- .../kv_connector/v1/example_connector.py | 14 +- .../v1/example_hidden_states_connector.py | 21 +- .../kv_connector/v1/flexkv_connector.py | 10 + .../kv_connector/v1/hf3fs/hf3fs_client.py | 8 +- .../kv_connector/v1/hf3fs/hf3fs_connector.py | 16 +- .../v1/hf3fs/hf3fs_metadata_server.py | 8 +- .../v1/hf3fs/utils/gather_scatter_helper.py | 2 + .../kv_connector/v1/lmcache_connector.py | 66 +- .../multi_process_adapter.py | 95 +- .../v1/lmcache_integration/utils.py | 11 +- .../v1/lmcache_integration/vllm_v1_adapter.py | 21 +- .../kv_connector/v1/lmcache_mp_connector.py | 103 +- .../kv_transfer/kv_connector/v1/metrics.py | 21 +- .../v1/mooncake/mooncake_connector.py | 39 +- .../v1/mooncake/mooncake_utils.py | 3 +- .../v1/mooncake/store/coordinator.py | 3 +- .../kv_connector/v1/mooncake/store/data.py | 3 + .../kv_connector/v1/mooncake/store/worker.py | 4 +- .../kv_connector/v1/moriio/moriio_common.py | 3 +- .../v1/moriio/moriio_connector.py | 33 +- .../kv_connector/v1/moriio/moriio_engine.py | 10 +- .../kv_connector/v1/multi_connector.py | 20 +- .../kv_connector/v1/nixl/base_scheduler.py | 10 +- .../kv_connector/v1/nixl/base_worker.py | 55 +- .../kv_connector/v1/nixl/connector.py | 8 +- .../kv_connector/v1/nixl/metadata.py | 7 +- .../kv_connector/v1/nixl/pull_scheduler.py | 8 +- .../kv_connector/v1/nixl/pull_worker.py | 9 +- .../kv_connector/v1/nixl/push_worker.py | 6 +- .../kv_transfer/kv_connector/v1/nixl/stats.py | 2 +- .../kv_transfer/kv_connector/v1/nixl/utils.py | 3 +- .../v1/offloading/canonical_mapping.py | 12 +- .../kv_connector/v1/offloading/events.py | 10 +- .../kv_connector/v1/offloading/metrics.py | 6 +- .../kv_connector/v1/offloading/scheduler.py | 30 +- .../kv_connector/v1/offloading/worker.py | 18 +- .../v1/ssm_conv_transfer_utils.py | 3 + .../kv_transfer/kv_transfer_state.py | 6 +- vllm/distributed/parallel_state.py | 51 +- vllm/distributed/stateless_coordinator.py | 3 +- vllm/distributed/utils.py | 18 +- vllm/distributed/weight_transfer/__init__.py | 3 +- vllm/distributed/weight_transfer/base.py | 37 +- vllm/distributed/weight_transfer/factory.py | 6 +- .../distributed/weight_transfer/ipc_engine.py | 17 +- .../weight_transfer/nccl_common.py | 10 +- .../weight_transfer/nccl_engine.py | 17 +- .../weight_transfer/packed_tensor.py | 4 + .../weight_transfer/sharded_rdt_common.py | 12 +- .../weight_transfer/sharded_rdt_engine.py | 48 +- .../weight_transfer/sharded_rdt_fake.py | 6 +- .../weight_transfer/sharded_rdt_trainer.py | 45 +- .../weight_transfer/sparse_nccl_engine.py | 9 +- vllm/engine/arg_utils.py | 4 +- vllm/engine/protocol.py | 29 +- vllm/entrypoints/anthropic/protocol.py | 30 +- vllm/entrypoints/anthropic/serving.py | 33 +- vllm/entrypoints/chat_utils.py | 26 +- vllm/entrypoints/cli/benchmark/base.py | 1 + vllm/entrypoints/cli/main.py | 5 +- vllm/entrypoints/generate/base/serving.py | 6 +- .../generate/beam_search/offline.py | 6 +- .../generate/beam_search/online.py | 2 +- .../generate/generative_scoring/serving.py | 6 + .../entrypoints/launchers/api_server/entry.py | 7 +- vllm/entrypoints/launchers/cli_args.py | 3 +- vllm/entrypoints/launchers/dp_supervisor.py | 27 +- vllm/entrypoints/launchers/grpc_server.py | 8 +- vllm/entrypoints/launchers/launcher.py | 10 +- .../entrypoints/launchers/render/app_state.py | 1 - vllm/entrypoints/launchers/render/entry.py | 1 - vllm/entrypoints/launchers/run_batch.py | 38 +- .../launchers/utils/server_utils.py | 3 +- vllm/entrypoints/llm.py | 38 +- vllm/entrypoints/mcp/tool.py | 6 +- vllm/entrypoints/mcp/tool_server.py | 11 +- vllm/entrypoints/offline_utils.py | 10 +- .../openai/chat_completion/batch_serving.py | 1 + .../openai/chat_completion/serving.py | 7 +- vllm/entrypoints/openai/completion/serving.py | 4 +- vllm/entrypoints/openai/models/serving.py | 7 +- .../openai/parser/harmony_utils.py | 24 +- .../openai/responses/api_router.py | 2 +- vllm/entrypoints/openai/responses/context.py | 21 +- vllm/entrypoints/openai/responses/harmony.py | 6 +- vllm/entrypoints/openai/responses/protocol.py | 11 +- .../openai/responses/streaming_events.py | 15 +- vllm/entrypoints/openai/responses/utils.py | 14 +- vllm/entrypoints/pooling/offline.py | 17 +- .../pooling/pooling/io_processor.py | 3 +- vllm/entrypoints/pooling/scoring/typing.py | 3 +- vllm/entrypoints/pooling/scoring/utils.py | 7 +- .../token_in_token_out/api_router.py | 3 +- .../scale_out/token_in_token_out/mm_serde.py | 3 +- .../entrypoints/serve/dev/cache/api_router.py | 10 +- vllm/entrypoints/serve/dev/rlhf/api_router.py | 10 +- .../serve/elastic_ep/middleware.py | 3 +- vllm/entrypoints/serve/engine/serving.py | 2 +- .../exception_handling/handlers/validation.py | 3 +- .../exception_handling/handlers/vllm_error.py | 4 +- .../serve/exception_handling/register.py | 3 +- .../serve/instrumentator/metrics.py | 1 - vllm/entrypoints/serve/lora/api_router.py | 6 +- .../serve/middleware/authenticate.py | 4 +- .../serve/middleware/log_response.py | 1 - .../serve/middleware/x_request_id.py | 6 +- .../entrypoints/serve/sagemaker/api_router.py | 8 +- vllm/entrypoints/serve/tokenize/protocol.py | 3 +- vllm/entrypoints/serve/utils/api_utils.py | 2 +- vllm/entrypoints/serve/utils/fingerprint.py | 6 +- vllm/entrypoints/serve/utils/orca_metrics.py | 35 +- .../serve/utils/tool_calls_utils.py | 1 - .../speech_to_text/base/serving.py | 9 +- vllm/entrypoints/speech_to_text/base/utils.py | 1 + .../speech_to_text/realtime/protocol.py | 14 +- .../speech_to_text/realtime/serving.py | 1 + vllm/env_override.py | 15 +- vllm/envs.py | 30 +- vllm/exceptions.py | 5 +- vllm/forward_context.py | 9 +- vllm/inputs/engine.py | 19 +- vllm/inputs/llm.py | 10 +- vllm/ir/op.py | 53 +- vllm/ir/ops/layernorm.py | 4 +- vllm/ir/util.py | 10 +- vllm/kernels/helion/case_key.py | 4 +- vllm/kernels/helion/config_manager.py | 3 +- vllm/kernels/helion/ops/__init__.py | 1 + .../ops/dynamic_per_token_scaled_fp8_quant.py | 1 - vllm/kernels/helion/ops/fused_qk_norm_rope.py | 1 - .../helion/ops/per_token_group_fp8_quant.py | 1 - .../ops/rms_norm_dynamic_per_token_quant.py | 1 - .../helion/ops/rms_norm_per_block_quant.py | 1 - .../ops/silu_and_mul_per_block_quant.py | 1 - vllm/kernels/helion/register.py | 11 +- vllm/kernels/helion/utils.py | 3 +- vllm/logger.py | 33 +- vllm/logging_utils/access_log_filter.py | 15 +- vllm/logging_utils/formatter.py | 5 +- vllm/logging_utils/log_time.py | 7 +- vllm/logprobs.py | 24 +- vllm/lora/layers/base_linear.py | 3 +- vllm/lora/layers/column_parallel_linear.py | 31 +- vllm/lora/layers/fused_moe.py | 31 +- vllm/lora/layers/logits_processor.py | 4 +- vllm/lora/layers/replicated_linear.py | 3 +- vllm/lora/layers/row_parallel_linear.py | 6 +- vllm/lora/layers/utils.py | 6 +- vllm/lora/lora_model.py | 35 +- vllm/lora/lora_weights.py | 1 - vllm/lora/model_manager.py | 12 +- vllm/lora/ops/triton_ops/fp8_kernel_utils.py | 21 +- .../ops/triton_ops/fused_moe_lora_fp8_op.py | 13 +- vllm/lora/ops/triton_ops/fused_moe_lora_op.py | 13 +- vllm/lora/ops/triton_ops/kernel_utils.py | 34 +- .../lora/ops/triton_ops/lora_expand_fp8_op.py | 16 +- vllm/lora/ops/triton_ops/lora_expand_op.py | 52 +- .../ops/triton_ops/lora_kernel_metadata.py | 16 +- .../lora/ops/triton_ops/lora_shrink_fp8_op.py | 40 +- vllm/lora/ops/triton_ops/lora_shrink_op.py | 52 +- vllm/lora/ops/triton_ops/utils.py | 11 +- vllm/lora/peft_helper.py | 10 +- vllm/lora/punica_wrapper/punica_base.py | 72 +- vllm/lora/punica_wrapper/punica_cpu.py | 50 +- vllm/lora/punica_wrapper/punica_gpu.py | 56 +- vllm/lora/punica_wrapper/punica_xpu.py | 59 +- vllm/lora/punica_wrapper/utils.py | 6 +- vllm/lora/request.py | 9 +- vllm/lora/resolver.py | 7 + vllm/lora/utils.py | 27 +- vllm/lora/worker_manager.py | 9 +- vllm/model_executor/custom_op.py | 12 +- .../determinism/batch_invariant.py | 25 +- vllm/model_executor/hw_agnostic/custom_op.py | 15 +- .../hw_agnostic/layers/layernorm.py | 3 +- .../attention/dsa/sparse_mqa_logits.py | 18 + .../model_executor/kernels/linear/__init__.py | 28 +- vllm/model_executor/kernels/linear/base.py | 14 +- .../mixed_precision/rdna_hybrid_w4a16.py | 10 +- .../linear/mixed_precision/triton_w4a16.py | 16 +- .../kernels/linear/mxfp4/base.py | 1 + .../kernels/linear/mxfp6/base.py | 1 + .../kernels/linear/nvfp4/flashinfer.py | 3 +- .../scaled_mm/BlockScaledMMLinearKernel.py | 3 +- .../kernels/linear/scaled_mm/aiter.py | 3 +- .../kernels/linear/scaled_mm/cutlass.py | 6 +- .../kernels/linear/scaled_mm/flashinfer.py | 15 +- .../kernels/linear/scaled_mm/marlin.py | 3 +- .../kernels/linear/scaled_mm/pytorch.py | 3 +- vllm/model_executor/kernels/mhc/aiter.py | 8 +- vllm/model_executor/kernels/mhc/tilelang.py | 11 +- .../kernels/mhc/tilelang_kernels.py | 2 +- vllm/model_executor/kernels/mhc/torch.py | 5 +- vllm/model_executor/layers/activation.py | 13 +- .../layers/attention/attention.py | 12 +- .../layers/attention/cross_attention.py | 4 +- .../attention/encoder_only_attention.py | 4 +- .../layers/attention/mla_attention.py | 45 +- .../layers/attention/mm_encoder_attention.py | 19 +- .../layers/attention/static_sink_attention.py | 4 +- .../layers/attention_layer_base.py | 6 +- vllm/model_executor/layers/conv.py | 4 +- .../layers/fused_allreduce_gemma_rms_norm.py | 3 +- .../model_executor/layers/fused_embed_norm.py | 3 +- .../layers/fused_moe/activation.py | 4 +- .../layers/fused_moe/all2all_utils.py | 4 +- .../model_executor/layers/fused_moe/config.py | 88 +- .../layers/fused_moe/deep_gemm_utils.py | 3 +- .../layers/fused_moe/expert_map_manager.py | 40 +- .../experts/aiter_mxfp4_w4a16_moe.py | 4 +- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 6 +- .../fused_moe/experts/aiter_mxfp8_moe.py | 3 +- .../experts/batched_deep_gemm_moe.py | 6 +- .../layers/fused_moe/experts/cpu_int4_moe.py | 2 + .../layers/fused_moe/experts/cpu_moe.py | 11 +- .../layers/fused_moe/experts/cutlass_moe.py | 3 +- .../layers/fused_moe/experts/deep_gemm_moe.py | 3 +- .../layers/fused_moe/experts/fallback.py | 3 +- .../experts/flashinfer_cutedsl_batched_moe.py | 10 +- .../experts/flashinfer_cutedsl_moe.py | 3 +- .../experts/flashinfer_cutlass_moe.py | 6 +- .../fused_moe/experts/fused_batched_moe.py | 6 +- .../fused_moe/experts/fused_humming_moe.py | 15 +- .../experts/gpt_oss_triton_kernels_moe.py | 9 +- .../layers/fused_moe/experts/lora_context.py | 3 +- .../fused_moe/experts/lora_experts_mixin.py | 3 +- .../layers/fused_moe/experts/marlin_moe.py | 64 +- .../fused_moe/experts/mxfp8_native_moe.py | 3 +- .../fused_moe/experts/nvfp4_emulation_moe.py | 12 +- .../fused_moe/experts/ocp_mx_emulation_moe.py | 10 +- .../fused_moe/experts/trtllm_bf16_moe.py | 11 +- .../fused_moe/experts/trtllm_fp8_moe.py | 11 +- .../fused_moe/experts/trtllm_lora_moe.py | 5 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 10 +- .../fused_moe/experts/trtllm_mxint4_moe.py | 3 +- .../fused_moe/experts/trtllm_nvfp4_moe.py | 12 +- .../layers/fused_moe/fused_moe.py | 14 +- .../layers/fused_moe/fused_moe_method_base.py | 20 +- .../layers/fused_moe/hpc_moe.py | 3 +- vllm/model_executor/layers/fused_moe/layer.py | 6 + .../layers/fused_moe/modular_kernel.py | 273 +- .../layers/fused_moe/moe_align_block_size.py | 56 +- .../layers/fused_moe/moe_fused_mul_sum.py | 4 +- .../layers/fused_moe/moe_output.py | 2 +- .../layers/fused_moe/moe_permute_unpermute.py | 51 +- .../layers/fused_moe/oracle/base.py | 15 +- .../layers/fused_moe/oracle/fp8.py | 18 +- .../layers/fused_moe/oracle/int8.py | 11 +- .../layers/fused_moe/oracle/int_wna16.py | 25 +- .../layers/fused_moe/oracle/mxfp4.py | 20 +- .../layers/fused_moe/oracle/mxfp8.py | 2 +- .../layers/fused_moe/oracle/nvfp4.py | 4 +- .../layers/fused_moe/oracle/unquantized.py | 7 +- .../layers/fused_moe/oracle/w4a8_int8.py | 24 +- .../fused_moe/prepare_finalize/batched.py | 3 +- .../fused_moe/prepare_finalize/deepep_ht.py | 4 +- .../fused_moe/prepare_finalize/deepep_ll.py | 8 +- .../fused_moe/prepare_finalize/deepep_v2.py | 6 +- .../layers/fused_moe/prepare_finalize/mori.py | 7 +- .../fused_moe/prepare_finalize/naive_dp_ep.py | 8 +- .../fused_moe/prepare_finalize/nixl_ep.py | 8 +- .../layers/fused_moe/routed_experts.py | 30 +- .../fused_moe/routed_experts_capturer.py | 4 +- .../aiter_shared_routed_fused_moe_router.py | 3 +- .../layers/fused_moe/router/base_router.py | 22 +- .../fused_moe/router/fused_moe_router.py | 8 +- .../layers/fused_moe/router/gate_linear.py | 3 +- .../layers/fused_moe/router/router_factory.py | 5 +- .../router/routing_simulator_router.py | 31 +- .../fused_moe/router/zero_expert_router.py | 3 +- .../layers/fused_moe/runner/moe_runner.py | 7 +- .../fused_moe/runner/moe_runner_interface.py | 3 +- .../fused_moe/topk_weight_and_reduce.py | 12 +- vllm/model_executor/layers/fused_moe/utils.py | 28 +- .../layers/fused_qk_norm_rope.py | 1 + .../layers/fusion/fused_act_quant.py | 3 +- .../layers/fusion/quant_activation.py | 3 +- vllm/model_executor/layers/hpc/gated_mla.py | 1 + vllm/model_executor/layers/hpc/hpc_ihc.py | 1 + vllm/model_executor/layers/indexer_topk.py | 10 +- vllm/model_executor/layers/layernorm.py | 5 +- vllm/model_executor/layers/lightning_attn.py | 11 +- vllm/model_executor/layers/linear.py | 18 +- .../model_executor/layers/logits_processor.py | 6 +- vllm/model_executor/layers/mamba/abstract.py | 6 +- .../layers/mamba/gdn/olmo_gdn_linear_attn.py | 7 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 25 +- .../mamba/linear/bailing_linear_attn.py | 2 +- .../layers/mamba/mamba_mixer.py | 7 +- .../layers/mamba/mamba_mixer2.py | 4 +- .../layers/mamba/mamba_utils.py | 8 +- .../layers/mamba/ops/causal_conv1d.py | 5 +- .../layers/mamba/ops/cpu/causal_conv1d.py | 3 +- .../mamba/ops/gdn_chunk_cutedsl/__init__.py | 1 + .../mamba/ops/gdn_chunk_cutedsl/kernel_h.py | 1 - .../layers/mamba/ops/mamba_ssm.py | 64 +- .../layers/mamba/ops/ssd_bmm.py | 5 +- .../layers/mamba/ops/ssd_combined.py | 4 +- .../layers/mamba/ops/ssu_dispatch.py | 3 +- .../minimax_rms_norm/lamport_workspace.py | 20 +- vllm/model_executor/layers/pooler/abstract.py | 4 +- .../layers/pooler/seqwise/poolers.py | 3 +- .../layers/pooler/tokwise/poolers.py | 3 +- .../layers/quantization/__init__.py | 1 + .../layers/quantization/auto_awq.py | 3 +- .../layers/quantization/auto_gptq.py | 2 +- .../layers/quantization/base_config.py | 28 +- .../compressed_tensors/compressed_tensors.py | 47 +- .../compressed_tensors_embedding.py | 3 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 4 +- .../compressed_tensors_moe_w4a8_int8.py | 3 +- .../compressed_tensors_moe_wna16.py | 3 +- .../schemes/compressed_tensors_scheme.py | 19 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 3 +- .../schemes/compressed_tensors_w8a8_mxfp8.py | 3 +- .../schemes/compressed_tensors_wNa8o8.py | 3 +- .../compressed_tensors/transform/linear.py | 10 +- .../compressed_tensors/transform/module.py | 3 +- .../transform/schemes/linear_qutlass_nvfp4.py | 3 +- .../quantization/compressed_tensors/utils.py | 14 +- .../model_executor/layers/quantization/fp8.py | 6 +- .../layers/quantization/fp_quant.py | 1 + .../layers/quantization/humming.py | 3 +- .../layers/quantization/inc/config_parser.py | 3 +- .../quantization/inc/schemes/inc_scheme.py | 3 +- .../layers/quantization/input_quant_fp8.py | 27 +- .../layers/quantization/kv_cache.py | 7 +- .../layers/quantization/modelopt.py | 42 +- .../layers/quantization/moe_wna16.py | 1 + .../layers/quantization/mxfp4.py | 3 +- .../layers/quantization/online/base.py | 6 +- .../layers/quantization/online/fp8.py | 18 +- .../layers/quantization/quark/quark.py | 22 +- .../layers/quantization/quark/quark_moe.py | 5 +- .../quantization/quark/schemes/quark_nvfp4.py | 3 +- .../quark/schemes/quark_scheme.py | 19 +- .../quark/schemes/quark_w4a8_mxfp4_fp8.py | 3 +- .../layers/quantization/qutlass_utils.py | 12 +- .../layers/quantization/torchao.py | 11 +- .../quantization/turboquant/centroids.py | 1 + .../layers/quantization/turboquant/config.py | 1 + .../layers/quantization/utils/config_utils.py | 5 +- .../quantization/utils/flashinfer_fp4_moe.py | 2 +- .../utils/flashinfer_mxint4_moe.py | 12 +- .../quantization/utils/flashinfer_utils.py | 8 +- .../layers/quantization/utils/fp8_utils.py | 36 +- .../quantization/utils/humming_utils.py | 30 +- .../layers/quantization/utils/int8_utils.py | 1 + .../quantization/utils/machete_utils.py | 4 +- .../layers/quantization/utils/marlin_utils.py | 12 +- .../quantization/utils/marlin_utils_fp4.py | 13 +- .../quantization/utils/marlin_utils_fp8.py | 12 +- .../quantization/utils/marlin_utils_test.py | 2 +- .../layers/quantization/utils/mxfp4_utils.py | 5 +- .../utils/nvfp4_emulation_utils.py | 3 +- .../layers/quantization/utils/nvfp4_utils.py | 28 +- .../layers/quantization/utils/quant_utils.py | 38 +- vllm/model_executor/layers/resampler.py | 12 +- .../layers/rotary_embedding/common.py | 16 +- .../rotary_embedding/deepseek_scaling_rope.py | 1 - .../rotary_embedding/ernie45_vl_rope.py | 2 +- .../layers/rotary_embedding/fope.py | 2 +- .../layers/rotary_embedding/mrope.py | 6 + .../rotary_embedding/mrope_interleaved.py | 6 +- .../rotary_embedding/ntk_scaling_rope.py | 3 +- .../layers/sparse_attn_indexer.py | 7 +- .../layers/sparse_mqa_indexer.py | 6 +- vllm/model_executor/layers/utils.py | 4 +- .../layers/vocab_parallel_embedding.py | 5 +- vllm/model_executor/model_loader/__init__.py | 1 + .../model_loader/base_loader.py | 3 +- .../model_loader/checkpoint_weight_patch.py | 3 +- .../model_loader/default_loader.py | 3 +- .../model_loader/ep_weight_filter.py | 10 +- .../model_loader/reload/__init__.py | 3 +- .../model_loader/reload/layerwise.py | 28 +- .../model_loader/reload/meta.py | 12 +- .../model_loader/reload/sanitize.py | 8 +- .../model_loader/reload/torchao_decorator.py | 3 +- .../model_loader/reload/utils.py | 8 +- .../model_loader/runai_streamer_loader.py | 9 +- .../model_loader/sharded_state_loader.py | 6 +- .../model_executor/model_loader/tensorizer.py | 15 +- .../model_loader/tensorizer_loader.py | 3 +- vllm/model_executor/model_loader/utils.py | 3 +- .../model_loader/weight_cache/daemon.py | 2 + .../model_loader/weight_cache/ipc_loader.py | 3 +- .../model_loader/weight_cache/protocol.py | 2 + .../model_loader/weight_tying.py | 3 +- .../model_loader/weight_utils.py | 29 +- vllm/model_executor/models/AXK1.py | 3 +- vllm/model_executor/models/adapters.py | 15 +- vllm/model_executor/models/arcee.py | 15 +- vllm/model_executor/models/aria.py | 34 +- vllm/model_executor/models/audioflamingo3.py | 27 +- vllm/model_executor/models/bagel.py | 21 +- .../models/bailing_moe_linear.py | 5 +- vllm/model_executor/models/bert.py | 8 +- vllm/model_executor/models/blip.py | 12 +- vllm/model_executor/models/blip2.py | 26 +- vllm/model_executor/models/clip.py | 18 +- vllm/model_executor/models/cohere2_moe.py | 3 +- vllm/model_executor/models/cohere2_vision.py | 23 +- vllm/model_executor/models/cohere_asr.py | 54 +- vllm/model_executor/models/cohere_compass.py | 21 +- vllm/model_executor/models/colbert.py | 5 +- vllm/model_executor/models/colpali.py | 3 +- vllm/model_executor/models/colqwen3.py | 4 +- vllm/model_executor/models/colqwen3_5.py | 4 +- vllm/model_executor/models/commandr.py | 3 +- vllm/model_executor/models/config.py | 8 +- .../models/conformer_encoder.py | 7 +- vllm/model_executor/models/cosmos3_edge.py | 6 +- vllm/model_executor/models/deepencoder.py | 122 +- vllm/model_executor/models/deepencoder2.py | 18 +- vllm/model_executor/models/deepseek_eagle3.py | 3 +- vllm/model_executor/models/deepseek_mtp.py | 3 +- vllm/model_executor/models/deepseek_ocr.py | 32 +- vllm/model_executor/models/deepseek_ocr2.py | 4 +- vllm/model_executor/models/deepseek_v2.py | 9 +- vllm/model_executor/models/deepseek_vl2.py | 22 +- vllm/model_executor/models/diffusion_gemma.py | 4 +- vllm/model_executor/models/dots_ocr.py | 24 +- vllm/model_executor/models/eagle2_5_vl.py | 33 +- vllm/model_executor/models/ernie45_vl.py | 27 +- .../models/extract_hidden_states.py | 6 +- vllm/model_executor/models/falcon_h1.py | 7 +- vllm/model_executor/models/fireredasr2.py | 9 +- vllm/model_executor/models/funasr.py | 9 +- vllm/model_executor/models/gemma3_mm.py | 19 +- vllm/model_executor/models/gemma3n.py | 13 +- .../models/gemma3n_audio_utils.py | 4 +- vllm/model_executor/models/gemma3n_mm.py | 41 +- vllm/model_executor/models/gemma4.py | 2 + vllm/model_executor/models/gemma4_mm.py | 19 +- vllm/model_executor/models/gemma4_mtp.py | 1 + vllm/model_executor/models/glm4_1v.py | 60 +- vllm/model_executor/models/glm4_moe.py | 3 +- .../models/glm4_moe_lite_mtp.py | 3 +- vllm/model_executor/models/glm4_moe_mtp.py | 6 +- vllm/model_executor/models/glm4v.py | 32 +- vllm/model_executor/models/glm_ocr_mtp.py | 3 +- vllm/model_executor/models/glmasr.py | 78 +- vllm/model_executor/models/glmasr_utils.py | 4 +- vllm/model_executor/models/gpt2.py | 1 + vllm/model_executor/models/gpt_oss.py | 1 + vllm/model_executor/models/granite.py | 1 + vllm/model_executor/models/granite4_vision.py | 1 + vllm/model_executor/models/granite_speech.py | 12 +- .../model_executor/models/granitemoehybrid.py | 1 + vllm/model_executor/models/hrm_text.py | 3 +- .../models/hyperclovax_vision_v2.py | 18 +- .../models/idefics2_vision_model.py | 32 +- vllm/model_executor/models/idefics3.py | 26 +- vllm/model_executor/models/interfaces.py | 107 +- vllm/model_executor/models/intern_vit.py | 2 +- vllm/model_executor/models/interns1.py | 48 +- vllm/model_executor/models/interns1_pro.py | 2 +- vllm/model_executor/models/interns1_vit.py | 9 +- vllm/model_executor/models/internvl.py | 54 +- .../model_executor/models/iquest_loopcoder.py | 1 + vllm/model_executor/models/isaac.py | 24 +- vllm/model_executor/models/kanana_v.py | 9 +- vllm/model_executor/models/keye.py | 61 +- vllm/model_executor/models/keye_vl1_5.py | 69 +- vllm/model_executor/models/kimi_k25.py | 6 +- vllm/model_executor/models/kimi_k25_vit.py | 15 +- vllm/model_executor/models/kimi_vl.py | 11 +- vllm/model_executor/models/lfm2.py | 1 + vllm/model_executor/models/lfm2_moe.py | 1 + vllm/model_executor/models/lfm2_siglip2.py | 29 +- vllm/model_executor/models/lfm2_vl.py | 17 +- vllm/model_executor/models/llama4.py | 5 +- vllm/model_executor/models/llava.py | 23 +- vllm/model_executor/models/llava_next.py | 19 +- .../model_executor/models/llava_next_video.py | 19 +- vllm/model_executor/models/llava_onevision.py | 61 +- .../model_executor/models/llava_onevision2.py | 3 +- .../models/longcat_flash_mtp.py | 3 +- .../models/longcat_flash_ngram.py | 1 + vllm/model_executor/models/mamba2.py | 1 + vllm/model_executor/models/medusa.py | 3 +- vllm/model_executor/models/mellum.py | 12 +- vllm/model_executor/models/midashenglm.py | 8 +- vllm/model_executor/models/mimo_audio.py | 2 + vllm/model_executor/models/mimo_v2.py | 4 +- vllm/model_executor/models/mimo_v2_mtp.py | 2 +- vllm/model_executor/models/mimo_v2_omni.py | 25 +- vllm/model_executor/models/minicpmo.py | 19 +- vllm/model_executor/models/minicpmv.py | 36 +- vllm/model_executor/models/minicpmv4_6.py | 16 +- vllm/model_executor/models/mistral3.py | 23 +- vllm/model_executor/models/mllama4.py | 35 +- vllm/model_executor/models/mlp_speculator.py | 10 +- vllm/model_executor/models/molmo.py | 25 +- vllm/model_executor/models/molmo2.py | 64 +- vllm/model_executor/models/moondream3.py | 4 +- vllm/model_executor/models/moonvit.py | 71 +- vllm/model_executor/models/moss_audio.py | 11 +- .../models/moss_transcribe_diarize.py | 20 +- vllm/model_executor/models/muse_glimmer.py | 13 +- .../model_executor/models/nano_nemotron_vl.py | 68 +- vllm/model_executor/models/nemotron_h.py | 1 + vllm/model_executor/models/nemotron_parse.py | 33 +- vllm/model_executor/models/nemotron_vl.py | 7 +- vllm/model_executor/models/openpangu_mtp.py | 3 +- vllm/model_executor/models/openpangu_vl.py | 8 +- vllm/model_executor/models/openvla.py | 11 +- vllm/model_executor/models/ovis.py | 18 +- vllm/model_executor/models/ovis2_5.py | 29 +- vllm/model_executor/models/paddleocr_vl.py | 5 +- vllm/model_executor/models/paligemma.py | 20 +- vllm/model_executor/models/parakeet.py | 4 +- vllm/model_executor/models/param2moe.py | 12 +- vllm/model_executor/models/phi3.py | 2 +- vllm/model_executor/models/phi3v.py | 37 +- vllm/model_executor/models/phi4mm.py | 74 +- vllm/model_executor/models/phi4mm_audio.py | 27 +- vllm/model_executor/models/phi4mm_utils.py | 83 +- vllm/model_executor/models/phi4siglip.py | 9 +- vllm/model_executor/models/pixtral.py | 29 +- .../models/qwen2_5_omni_thinker.py | 47 +- vllm/model_executor/models/qwen2_5_vl.py | 38 +- vllm/model_executor/models/qwen2_audio.py | 18 +- vllm/model_executor/models/qwen2_vl.py | 24 +- vllm/model_executor/models/qwen3_5.py | 2 +- vllm/model_executor/models/qwen3_asr.py | 7 +- vllm/model_executor/models/qwen3_dflash.py | 12 +- .../models/qwen3_omni_moe_thinker.py | 35 +- vllm/model_executor/models/qwen3_vl.py | 27 +- vllm/model_executor/models/radio.py | 7 +- vllm/model_executor/models/registry.py | 16 +- vllm/model_executor/models/roberta.py | 1 + vllm/model_executor/models/sarvam.py | 1 + vllm/model_executor/models/siglip.py | 16 +- vllm/model_executor/models/siglip2navit.py | 66 +- vllm/model_executor/models/skyworkr1v.py | 24 +- vllm/model_executor/models/stablelm.py | 3 +- vllm/model_executor/models/step3_vl.py | 30 +- vllm/model_executor/models/step3p5.py | 4 +- vllm/model_executor/models/step3p5_mtp.py | 3 +- vllm/model_executor/models/step_vl.py | 2 +- vllm/model_executor/models/terratorch.py | 2 +- .../models/transformers/__init__.py | 2 +- .../models/transformers/base.py | 27 +- .../models/transformers/causal.py | 1 - .../models/transformers/fusers/base.py | 27 +- .../models/transformers/fusers/mla.py | 3 +- .../models/transformers/fusers/moe.py | 12 +- .../models/transformers/fusers/rms_norm.py | 2 +- .../models/transformers/fx_utils.py | 37 +- .../models/transformers/layers.py | 3 +- .../model_executor/models/transformers/moe.py | 6 +- .../models/transformers/multimodal.py | 37 +- .../models/transformers/pooling.py | 3 +- .../models/transformers/utils.py | 27 +- vllm/model_executor/models/ultravox.py | 23 +- vllm/model_executor/models/utils.py | 62 +- vllm/model_executor/models/vision.py | 27 +- vllm/model_executor/models/voxtral.py | 6 +- .../model_executor/models/voxtral_realtime.py | 4 +- vllm/model_executor/models/voyage.py | 3 +- vllm/model_executor/models/whisper.py | 18 +- vllm/model_executor/models/zamba2.py | 33 +- vllm/model_executor/offloader/base.py | 1 + vllm/model_executor/offloader/prefetch.py | 1 + vllm/model_executor/offloader/prefetch_ops.py | 2 + vllm/model_executor/offloader/uva.py | 1 + vllm/model_executor/parameter.py | 50 +- vllm/model_executor/utils.py | 10 +- .../model_executor/warmup/deep_gemm_warmup.py | 15 +- .../warmup/jit_warmup_tilelang_helper.py | 1 - vllm/model_executor/warmup/kernel_warmup.py | 6 +- .../warmup/qwen4_exp_qsa_warmup.py | 4 +- vllm/models/deepseek_v32/common/kernels.py | 3 +- vllm/models/deepseek_v4/amd/model.py | 3 +- vllm/models/deepseek_v4/amd/mtp.py | 3 +- vllm/models/deepseek_v4/attention.py | 4 +- .../deepseek_v4/common/mm_preprocess.py | 9 +- .../deepseek_v4/common/ops/cache_utils.py | 9 +- .../common/ops/fused_compress_quant_cache.py | 7 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 3 +- .../common/ops/fused_inv_rope_fp8_quant.py | 4 +- vllm/models/deepseek_v4/common/vision.py | 1 + vllm/models/deepseek_v4/cpu/cpu_compressor.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_mla.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 24 +- vllm/models/deepseek_v4/cpu/cpu_utils.py | 1 + vllm/models/deepseek_v4/cpu/model.py | 3 +- vllm/models/deepseek_v4/nvidia/fi_moe.py | 3 +- .../deepseek_v4/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v4/nvidia/model.py | 3 +- vllm/models/deepseek_v4/nvidia/mtp.py | 3 +- vllm/models/deepseek_v4/xpu/model.py | 3 +- vllm/models/deepseek_v4/xpu/mtp.py | 3 +- vllm/models/deepseek_v41/amd/model.py | 6 +- vllm/models/deepseek_v41/amd/vl_model.py | 3 +- vllm/models/deepseek_v41/attention.py | 7 +- vllm/models/deepseek_v41/common/engram.py | 9 +- .../deepseek_v41/common/mm_preprocess.py | 9 +- .../deepseek_v41/common/ops/cache_utils.py | 9 +- .../common/ops/fused_compress_quant_cache.py | 1 + .../common/ops/indexer_k_store.py | 2 + .../deepseek_v41/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v41/nvidia/model.py | 6 +- vllm/models/deepseek_v41/nvidia/vl_model.py | 3 +- .../models/dots3_note/nvidia/audio_encoder.py | 10 +- .../models/glm5next/amd/ops/kpool_compress.py | 14 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- .../amd/ops/third_party/kda/kernels.py | 29 +- vllm/models/glm5next/common/kda.py | 6 +- vllm/models/glm5next/common/multimodal.py | 3 +- vllm/models/glm5next/common/sparse_indexer.py | 6 +- .../glm5next/nvidia/ops/fused_eh_norm.py | 3 +- .../glm5next/nvidia/ops/kpool_compress.py | 15 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- .../nvidia/ops/third_party/kda/kernels.py | 29 +- vllm/models/hy_v4/nvidia/attention.py | 3 + vllm/models/hy_v4/nvidia/flashmla_sparse.py | 2 + vllm/models/hy_v4/nvidia/hc.py | 4 + vllm/models/hy_v4/nvidia/model.py | 2 +- vllm/models/hy_v4/nvidia/mtp.py | 4 + vllm/models/inkling/amd/logits_processor.py | 1 + vllm/models/inkling/amd/model.py | 3 +- vllm/models/inkling/amd/moe.py | 9 +- vllm/models/inkling/amd/mtp.py | 9 +- vllm/models/inkling/amd/ops/norm.py | 6 +- .../models/inkling/nvidia/logits_processor.py | 1 + vllm/models/inkling/nvidia/model.py | 3 +- vllm/models/inkling/nvidia/moe.py | 9 +- vllm/models/inkling/nvidia/mtp.py | 9 +- vllm/models/inkling/nvidia/ops/mm_towers.py | 6 +- vllm/models/inkling/nvidia/ops/norm.py | 6 +- vllm/models/kimi_k3/amd/latent_moe_runner.py | 4 +- vllm/models/kimi_k3/amd/linear.py | 4 +- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 9 + vllm/models/kimi_k3/amd/ops/kda_prefill.py | 17 + .../kimi_k3/amd/ops/third_party/kda/chunk.py | 13 +- .../amd/ops/third_party/kda/chunk_intra.py | 3 +- .../kda/chunk_intra_token_parallel.py | 5 +- .../kimi_k3/nvidia/latent_moe_runner.py | 4 +- .../models/kimi_k3/nvidia/low_latency_gemm.py | 1 + vllm/models/kimi_k3/nvidia/mla.py | 3 +- vllm/models/kimi_k3/nvidia/model.py | 2 + ...educe_rmsnorm_reduce_scatter_early_exit.py | 3 - .../fused_add_multicast_gemm.py | 1 - .../cute_dsl/latent_moe_tail/primitives.py | 13 - .../nvidia/ops/third_party/kda/chunk.py | 17 +- .../nvidia/ops/third_party/kda/chunk_intra.py | 300 +- .../kda/chunk_intra_token_parallel.py | 56 +- .../ops/third_party/kda/fused_recurrent.py | 4 +- vllm/models/minimax_m3/amd/model.py | 3 +- vllm/models/minimax_m3/common/indexer.py | 6 +- .../minimax_m3/common/sparse_attention.py | 3 +- vllm/models/minimax_m3/nvidia/indexer_msa.py | 6 +- vllm/models/minimax_m3/nvidia/model.py | 6 +- vllm/models/qwen4_exp/amd/indexer_qsa.py | 5 - vllm/models/qwen4_exp/amd/low_latency_gemm.py | 1 - vllm/models/qwen4_exp/amd/model.py | 2 - vllm/models/qwen4_exp/amd/mtp.py | 4 - vllm/models/qwen4_exp/amd/ops/qsa.py | 6 - vllm/models/qwen4_exp/amd/ple_layer.py | 1 - vllm/models/qwen4_exp/amd/qsa.py | 1 - vllm/models/qwen4_exp/common/ple.py | 2 - vllm/models/qwen4_exp/common/qsa_cache.py | 3 - vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 2 - vllm/models/qwen4_exp/nvidia/model.py | 2 - vllm/models/qwen4_exp/nvidia/mtp.py | 4 - .../qwen4_exp/nvidia/ngram_embedding.py | 1 - vllm/models/qwen4_exp/nvidia/ops/qsa.py | 3 - .../qwen4_exp/nvidia/ops/qsa_indexer.py | 6 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 1 - vllm/multimodal/audio.py | 7 + vllm/multimodal/cache.py | 68 +- vllm/multimodal/encoder_budget.py | 3 +- vllm/multimodal/gpu_ipc_memory.py | 1 + vllm/multimodal/hasher.py | 5 +- vllm/multimodal/image.py | 3 +- vllm/multimodal/inputs.py | 77 +- vllm/multimodal/media/audio.py | 15 +- vllm/multimodal/media/base.py | 6 +- vllm/multimodal/media/connector.py | 66 +- vllm/multimodal/parse.py | 50 +- vllm/multimodal/processing/context.py | 48 +- vllm/multimodal/processing/dummy_inputs.py | 15 +- vllm/multimodal/processing/inputs.py | 3 +- vllm/multimodal/processing/processor.py | 85 +- vllm/multimodal/registry.py | 26 +- vllm/multimodal/utils.py | 33 +- vllm/multimodal/video.py | 13 +- vllm/multimodal/video_decoders/opencv.py | 4 +- vllm/multimodal/video_prune/evs.py | 17 +- vllm/multimodal/video_prune/vidcom2.py | 1 + vllm/outputs.py | 11 +- vllm/parser/abstract_parser.py | 43 +- vllm/parser/cohere_command.py | 12 +- vllm/parser/engine/__init__.py | 3 +- vllm/parser/engine/incremental_lexer.py | 3 +- vllm/parser/engine/streaming_parser_engine.py | 3 +- vllm/parser/engine/token_id_scanner.py | 3 +- vllm/parser/gemma4.py | 1 + vllm/parser/inkling.py | 3 +- vllm/parser/mistral.py | 2 + vllm/parser/parser_manager.py | 8 +- vllm/platforms/__init__.py | 4 +- vllm/platforms/cpu.py | 15 +- vllm/platforms/cuda.py | 11 +- vllm/platforms/interface.py | 186 +- vllm/platforms/rocm.py | 21 +- vllm/platforms/xpu.py | 8 +- vllm/plugins/__init__.py | 1 + .../lora_resolvers/filesystem_resolver.py | 3 +- .../plugins/lora_resolvers/hf_hub_resolver.py | 6 +- vllm/pooling_params.py | 2 + vllm/profiler/layerwise_profile.py | 6 +- vllm/profiler/wrapper.py | 5 +- vllm/ray/lazy_utils.py | 1 - vllm/ray/ray_env.py | 1 + vllm/reasoning/abs_reasoning_parsers.py | 64 +- vllm/reasoning/basic_parsers.py | 13 +- .../reasoning/deepseek_r1_reasoning_parser.py | 3 +- .../reasoning/deepseek_v3_reasoning_parser.py | 7 +- vllm/reasoning/ernie45_reasoning_parser.py | 9 +- vllm/reasoning/gemma4_utils.py | 1 + vllm/reasoning/gptoss_reasoning_parser.py | 3 +- vllm/reasoning/granite_reasoning_parser.py | 11 +- .../hunyuan_a13b_reasoning_parser.py | 7 +- vllm/reasoning/hy_v3_reasoning_parser.py | 3 +- vllm/reasoning/hy_v4_reasoning_parser.py | 2 +- vllm/reasoning/identity_reasoning_parser.py | 3 +- vllm/reasoning/minimax_m2_reasoning_parser.py | 7 +- vllm/reasoning/olmo3_reasoning_parser.py | 13 +- .../reasoning/poolside_v1_reasoning_parser.py | 3 +- vllm/reasoning/step3_reasoning_parser.py | 6 +- vllm/reasoning/step3p5_reasoning_parser.py | 3 +- vllm/renderers/base.py | 9 +- vllm/renderers/hf.py | 8 +- vllm/renderers/inputs/preprocess.py | 15 +- vllm/renderers/inputs/tokenize.py | 7 +- vllm/renderers/online_derenderer.py | 6 +- vllm/renderers/params.py | 6 +- vllm/sampling_params.py | 13 +- vllm/scalar_type.py | 37 +- .../ops/fused_norm_gate.py | 1 + .../flash_linear_attention/ops/kda.py | 1 + .../ops/layernorm_guard.py | 4 +- vllm/third_party/pynvml.py | 4490 ++++++++++------- vllm/tilelang_utils/__init__.py | 1 + vllm/tokenizers/deepseek_v32.py | 4 +- vllm/tokenizers/deepseek_v4.py | 4 +- vllm/tokenizers/detokenizer_utils.py | 1 + vllm/tokenizers/hf.py | 6 +- vllm/tool_parsers/abstract_tool_parser.py | 22 +- vllm/tool_parsers/apertus_tool_parser.py | 46 +- vllm/tool_parsers/ernie45_tool_parser.py | 3 +- .../tool_parsers/functiongemma_tool_parser.py | 3 +- vllm/tool_parsers/gemma4_utils.py | 3 + vllm/tool_parsers/gptoss_tool_parser.py | 3 +- .../granite_20b_fc_tool_parser.py | 3 +- vllm/tool_parsers/granite_tool_parser.py | 3 +- vllm/tool_parsers/hermes_tool_parser.py | 3 +- vllm/tool_parsers/hunyuan_a13b_tool_parser.py | 9 +- vllm/tool_parsers/hy_v3_tool_parser.py | 1 + vllm/tool_parsers/hy_v4_tool_parser.py | 5 +- vllm/tool_parsers/lfm2_tool_parser.py | 6 +- .../llama4_pythonic_tool_parser.py | 8 +- vllm/tool_parsers/llama_tool_parser.py | 6 +- vllm/tool_parsers/olmo3_tool_parser.py | 7 +- vllm/tool_parsers/phi4mini_tool_parser.py | 7 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 3 +- vllm/tool_parsers/pythonic_tool_parser.py | 7 +- vllm/tool_parsers/step3_tool_parser.py | 3 +- vllm/tool_parsers/step3p5_tool_parser.py | 75 +- vllm/tool_parsers/structural_tag_registry.py | 4 +- vllm/tool_parsers/utils.py | 18 +- vllm/tracing/__init__.py | 8 +- vllm/tracing/otel.py | 11 +- vllm/tracing/utils.py | 6 +- vllm/transformers_utils/config.py | 43 +- vllm/transformers_utils/configs/AXK1.py | 5 +- vllm/transformers_utils/configs/__init__.py | 3 +- vllm/transformers_utils/configs/colpali.py | 3 +- vllm/transformers_utils/configs/colqwen3.py | 3 +- .../configs/extract_hidden_states.py | 3 +- vllm/transformers_utils/configs/falcon.py | 2 +- vllm/transformers_utils/configs/hy_v3.py | 4 +- vllm/transformers_utils/configs/hy_v4.py | 1 + .../transformers_utils/configs/hyperclovax.py | 4 +- vllm/transformers_utils/configs/kimi_k25.py | 4 +- vllm/transformers_utils/configs/lfm2_moe.py | 7 +- .../configs/mlp_speculator.py | 4 +- vllm/transformers_utils/configs/moondream3.py | 1 + vllm/transformers_utils/configs/nemotron.py | 13 +- vllm/transformers_utils/configs/nemotron_h.py | 7 +- .../transformers_utils/configs/olmo_hybrid.py | 6 +- vllm/transformers_utils/configs/ovis.py | 2 + vllm/transformers_utils/configs/qwen3_5.py | 2 +- .../transformers_utils/configs/qwen3_5_moe.py | 2 +- vllm/transformers_utils/configs/qwen3_next.py | 6 +- vllm/transformers_utils/configs/radio.py | 6 +- .../configs/speculators/algos.py | 16 +- .../configs/speculators/base.py | 8 +- vllm/transformers_utils/configs/ultravox.py | 4 +- vllm/transformers_utils/dynamic_module.py | 3 +- vllm/transformers_utils/processor.py | 17 +- .../transformers_utils/processors/__init__.py | 3 +- vllm/transformers_utils/processors/bagel.py | 15 +- .../processors/cohere_asr.py | 5 +- .../processors/cosmos3_edge.py | 4 +- .../processors/deepseek_ocr.py | 7 +- .../processors/deepseek_vl2.py | 12 +- .../processors/fireredasr2.py | 13 +- vllm/transformers_utils/processors/funasr.py | 16 +- vllm/transformers_utils/processors/glm4v.py | 3 +- .../transformers_utils/processors/glm5next.py | 6 +- .../transformers_utils/processors/internvl.py | 3 +- vllm/transformers_utils/processors/isaac.py | 5 + .../transformers_utils/processors/kimi_k25.py | 4 +- .../processors/mimo_v2_omni.py | 3 + .../transformers_utils/processors/minicpmo.py | 22 +- .../transformers_utils/processors/minicpmv.py | 14 +- .../processors/moondream3.py | 15 +- .../processors/nano_nemotron_vl.py | 15 +- .../processors/nemotron_vl.py | 6 +- vllm/transformers_utils/processors/ovis.py | 28 +- vllm/transformers_utils/processors/ovis2_5.py | 64 +- vllm/transformers_utils/processors/pixtral.py | 3 +- vllm/transformers_utils/processors/voxtral.py | 3 +- vllm/transformers_utils/repo_utils.py | 18 +- vllm/transformers_utils/runai_utils.py | 11 +- vllm/transformers_utils/s3_utils.py | 8 +- vllm/transformers_utils/utils.py | 8 +- vllm/utils/argparse_utils.py | 8 +- vllm/utils/async_utils.py | 10 +- vllm/utils/cache.py | 10 +- vllm/utils/collection_utils.py | 9 +- vllm/utils/counter.py | 8 +- vllm/utils/cpu_resource_utils.py | 6 +- vllm/utils/cpu_triton_utils.py | 4 +- vllm/utils/deep_gemm.py | 20 +- vllm/utils/flashinfer.py | 5 +- vllm/utils/func_utils.py | 8 +- vllm/utils/gc_utils.py | 24 +- vllm/utils/gpu_sync_debug.py | 6 +- vllm/utils/hashing.py | 4 + vllm/utils/import_utils.py | 32 +- vllm/utils/jit_monitor.py | 3 +- vllm/utils/jsontree.py | 3 +- vllm/utils/math_utils.py | 2 +- vllm/utils/mem_utils.py | 4 +- vllm/utils/multi_stream_utils.py | 2 + vllm/utils/network_utils.py | 8 +- vllm/utils/nvtx_pytorch_hooks.py | 5 +- vllm/utils/ompmultiprocessing.py | 13 +- vllm/utils/platform_utils.py | 3 +- vllm/utils/registry.py | 15 +- vllm/utils/sparse_utils.py | 3 +- vllm/utils/system_utils.py | 9 +- vllm/utils/tensor_schema.py | 1 - vllm/utils/torch_utils.py | 44 +- vllm/v1/attention/backend.py | 49 +- vllm/v1/attention/backends/cpu_attn.py | 11 +- vllm/v1/attention/backends/fa_utils.py | 1 + vllm/v1/attention/backends/flash_attn.py | 12 +- .../attention/backends/flash_attn_diffkv.py | 8 + vllm/v1/attention/backends/flashinfer.py | 18 +- vllm/v1/attention/backends/flex_attention.py | 21 +- vllm/v1/attention/backends/gdn_attn.py | 3 +- vllm/v1/attention/backends/mamba2_attn.py | 3 +- vllm/v1/attention/backends/mamba_attn.py | 19 +- vllm/v1/attention/backends/mla/amx_mla.py | 3 +- .../attention/backends/mla/flashmla_sparse.py | 4 +- vllm/v1/attention/backends/mla/indexer.py | 3 +- .../backends/mla/prefill/aiter_flash_attn.py | 2 +- .../v1/attention/backends/mla/prefill/base.py | 3 +- .../backends/mla/prefill/registry.py | 3 + .../backends/mla/prefill/selector.py | 3 + .../attention/backends/mla/rocm_aiter_mla.py | 4 +- .../backends/mla/rocm_aiter_mla_sparse.py | 3 +- .../v1/attention/backends/mla/sparse_utils.py | 4 +- vllm/v1/attention/backends/registry.py | 10 +- vllm/v1/attention/backends/rocm_aiter_fa.py | 11 +- .../backends/rocm_aiter_unified_attn.py | 10 +- vllm/v1/attention/backends/rocm_attn.py | 11 +- vllm/v1/attention/backends/triton_attn.py | 11 +- vllm/v1/attention/backends/utils.py | 37 +- vllm/v1/attention/ops/common.py | 6 +- vllm/v1/attention/ops/dcp.py | 33 +- vllm/v1/attention/ops/flashmla.py | 8 +- vllm/v1/attention/ops/merge_attn_states.py | 2 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- .../attention/ops/triton_decode_attention.py | 3 +- .../attention/ops/triton_prefill_attention.py | 6 +- .../turboquant_soa/triton_turboquant_store.py | 3 +- vllm/v1/attention/ops/vit_attn_wrappers.py | 6 +- vllm/v1/attention/ops/xpu_mla_sparse.py | 9 +- vllm/v1/attention/selector.py | 2 +- vllm/v1/core/block_pool.py | 39 +- vllm/v1/core/encoder_cache_manager.py | 10 +- vllm/v1/core/kv_cache_coordinator.py | 62 +- vllm/v1/core/kv_cache_manager.py | 29 +- vllm/v1/core/kv_cache_utils.py | 117 +- vllm/v1/core/sched/interface.py | 14 +- vllm/v1/core/sched/request_queue.py | 12 +- vllm/v1/core/sched/scheduler.py | 31 +- vllm/v1/core/sched/utils.py | 4 + vllm/v1/core/single_type_kv_cache_manager.py | 92 +- vllm/v1/cudagraph_dispatcher.py | 14 +- vllm/v1/engine/__init__.py | 10 +- vllm/v1/engine/async_llm.py | 46 +- vllm/v1/engine/core.py | 43 +- vllm/v1/engine/core_client.py | 46 +- vllm/v1/engine/detokenizer.py | 7 +- vllm/v1/engine/input_processor.py | 3 +- vllm/v1/engine/llm_engine.py | 2 - vllm/v1/engine/logprobs.py | 8 +- vllm/v1/engine/output_processor.py | 8 +- vllm/v1/engine/parallel_sampling.py | 4 +- vllm/v1/engine/tensor_ipc.py | 4 +- vllm/v1/engine/utils.py | 26 +- vllm/v1/executor/abstract.py | 57 +- vllm/v1/executor/multiproc_executor.py | 19 +- vllm/v1/executor/ray_executor.py | 9 +- vllm/v1/executor/ray_executor_v2.py | 5 +- vllm/v1/executor/ray_utils.py | 13 +- vllm/v1/kv_cache_interface.py | 63 +- vllm/v1/kv_cache_layout.py | 3 +- vllm/v1/kv_cache_spec_registry.py | 27 +- vllm/v1/kv_offload/base.py | 71 +- vllm/v1/kv_offload/cpu/gpu_worker.py | 33 +- vllm/v1/kv_offload/cpu/manager.py | 3 +- vllm/v1/kv_offload/cpu/policies/arc.py | 3 +- vllm/v1/kv_offload/cpu/policies/base.py | 20 +- vllm/v1/kv_offload/cpu/policies/factory.py | 1 + vllm/v1/kv_offload/cpu/policies/lru.py | 3 +- .../kv_offload/cpu/shared_offload_region.py | 5 +- vllm/v1/kv_offload/cpu/spec.py | 6 +- vllm/v1/kv_offload/file_mapper.py | 10 +- vllm/v1/kv_offload/tiering/async_lookup.py | 16 +- vllm/v1/kv_offload/tiering/base.py | 48 +- vllm/v1/kv_offload/tiering/example/manager.py | 32 +- vllm/v1/kv_offload/tiering/factory.py | 1 + vllm/v1/kv_offload/tiering/fs/io.py | 13 +- vllm/v1/kv_offload/tiering/fs/manager.py | 40 +- vllm/v1/kv_offload/tiering/fs/thread_pool.py | 19 +- vllm/v1/kv_offload/tiering/manager.py | 55 +- vllm/v1/kv_offload/tiering/obj/manager.py | 32 +- .../v1/kv_offload/tiering/p2p/control/base.py | 2 + vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 3 +- vllm/v1/kv_offload/tiering/p2p/data/base.py | 5 + vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 4 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 10 +- .../tiering/p2p/session/protocol.py | 3 +- .../kv_offload/tiering/p2p/session/session.py | 3 +- vllm/v1/kv_offload/tiering/spec.py | 13 +- vllm/v1/metrics/buckets.py | 3 + vllm/v1/metrics/loggers.py | 22 +- vllm/v1/metrics/perf.py | 81 +- vllm/v1/metrics/prometheus.py | 2 +- vllm/v1/metrics/ray_wrappers.py | 22 +- vllm/v1/metrics/reader.py | 1 + vllm/v1/metrics/stats.py | 12 +- vllm/v1/outputs.py | 4 +- vllm/v1/request.py | 3 +- vllm/v1/sample/logits_processor/__init__.py | 9 +- vllm/v1/sample/logits_processor/builtin.py | 9 +- vllm/v1/sample/logits_processor/interface.py | 3 +- vllm/v1/sample/logits_processor/state.py | 9 +- vllm/v1/sample/ops/logprobs.py | 4 +- vllm/v1/sample/ops/penalties.py | 8 +- vllm/v1/sample/ops/topk_topp_sampler.py | 12 +- vllm/v1/sample/ops/topk_topp_triton.py | 7 +- vllm/v1/sample/rejection_sampler.py | 19 +- vllm/v1/sample/sampler.py | 9 +- vllm/v1/serial_utils.py | 15 +- vllm/v1/simple_kv_offload/manager.py | 13 +- vllm/v1/simple_kv_offload/metadata.py | 3 +- vllm/v1/simple_kv_offload/worker.py | 2 + vllm/v1/spec_decode/dflash.py | 3 +- vllm/v1/spec_decode/extract_hidden_states.py | 10 +- vllm/v1/spec_decode/gemma4.py | 6 +- vllm/v1/spec_decode/llm_base_proposer.py | 33 +- vllm/v1/spec_decode/medusa.py | 4 +- vllm/v1/spec_decode/ngram_proposer.py | 6 +- vllm/v1/spec_decode/ngram_proposer_gpu.py | 28 +- vllm/v1/spec_decode/step3p5.py | 1 - vllm/v1/spec_decode/suffix_decoding.py | 3 +- vllm/v1/spec_decode/utils.py | 25 +- vllm/v1/structured_output/backend_guidance.py | 1 - vllm/v1/structured_output/backend_outlines.py | 5 +- vllm/v1/structured_output/backend_types.py | 36 +- vllm/v1/structured_output/utils.py | 23 +- vllm/v1/utils.py | 26 +- vllm/v1/worker/block_table.py | 29 +- vllm/v1/worker/dp_utils.py | 9 +- .../worker/ec_connector_model_runner_mixin.py | 4 +- vllm/v1/worker/encoder_cudagraph.py | 9 +- vllm/v1/worker/gpu/cudagraph_utils.py | 15 +- vllm/v1/worker/gpu/dp_utils.py | 10 +- vllm/v1/worker/gpu/lora_utils.py | 3 +- vllm/v1/worker/gpu/mm/encoder_cache.py | 3 +- vllm/v1/worker/gpu/model_runner.py | 6 +- vllm/v1/worker/gpu/model_states/interface.py | 3 +- vllm/v1/worker/gpu/pp_utils.py | 7 +- vllm/v1/worker/gpu/sample/output.py | 3 +- vllm/v1/worker/gpu/sample/sampler.py | 3 +- vllm/v1/worker/gpu/sample/watermark.py | 1 + .../gpu/spec_decode/adaptive_verification.py | 11 +- .../spec_decode/autoregressive/speculator.py | 3 +- .../gpu/spec_decode/dflash/cudagraph.py | 3 +- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 4 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 3 +- vllm/v1/worker/gpu_input_batch.py | 27 +- vllm/v1/worker/gpu_model_runner.py | 151 +- vllm/v1/worker/gpu_ubatch_wrapper.py | 3 +- vllm/v1/worker/gpu_worker.py | 20 +- .../worker/kv_connector_model_runner_mixin.py | 4 +- vllm/v1/worker/lora_model_runner_mixin.py | 14 +- vllm/v1/worker/mamba_utils.py | 20 +- vllm/v1/worker/startup_plan.py | 9 +- vllm/v1/worker/ubatch_utils.py | 15 +- vllm/v1/worker/ubatching.py | 4 +- vllm/v1/worker/utils.py | 27 +- vllm/v1/worker/worker_base.py | 17 +- vllm/v1/worker/workspace.py | 6 + vllm/vllm_flash_attn/flash_attn_interface.py | 39 + 2002 files changed, 15050 insertions(+), 14295 deletions(-) diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index fe9e332bdfa7..5576c84b05cb 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -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 \ diff --git a/.buildkite/performance-benchmarks/scripts/compare-json-results.py b/.buildkite/performance-benchmarks/scripts/compare-json-results.py index c9f8139fe62f..823c7488ec59 100644 --- a/.buildkite/performance-benchmarks/scripts/compare-json-results.py +++ b/.buildkite/performance-benchmarks/scripts/compare-json-results.py @@ -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 diff --git a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py index e694c38780cb..deb14a65e4a8 100644 --- a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py +++ b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py @@ -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' diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/scripts/generate-nightly-index.py index 67adf9c92b80..8f92996ff9ab 100644 --- a/.buildkite/scripts/generate-nightly-index.py +++ b/.buildkite/scripts/generate-nightly-index.py @@ -18,8 +18,7 @@ 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.+)-(?P[^-]+?)(-(?P[^-]+))?-(?P[^-]+)-(?P[^-]+)-(?P[^-]+)\.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: diff --git a/benchmarks/attention_benchmarks/batch_spec.py b/benchmarks/attention_benchmarks/batch_spec.py index 9f15f1d8096e..85c689b48944 100644 --- a/benchmarks/attention_benchmarks/batch_spec.py +++ b/benchmarks/attention_benchmarks/batch_spec.py @@ -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: (?) q(k?) (s(k?))? @@ -24,6 +23,7 @@ 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: (?) q(k?) (s(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) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 9c1ea1add70d..943133252be0 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -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 @@ # 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 = [] diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 45858f9148b3..b6fd26c80344 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -21,8 +21,7 @@ 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 @@ def print_table( 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 @@ def save_json(self, results: list[BenchmarkResult], path: str): 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 diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 8a6ae3f87d4a..04dcd73719a3 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -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 @@ def get_prefill_backend_config(prefill_backend: str) -> dict: 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): diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index 01fd558043b7..a1b7d045050e 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -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 @@ 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 benchmark_fn(): 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) diff --git a/benchmarks/benchmark_batch_invariance.py b/benchmarks/benchmark_batch_invariance.py index 7473a41e5140..dd22523ce201 100755 --- a/benchmarks/benchmark_batch_invariance.py +++ b/benchmarks/benchmark_batch_invariance.py @@ -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. """ diff --git a/benchmarks/benchmark_hash.py b/benchmarks/benchmark_hash.py index 08cdc012d652..c5af4da9998b 100644 --- a/benchmarks/benchmark_hash.py +++ b/benchmarks/benchmark_hash.py @@ -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) diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py index f0a35a0cf150..1f5da88c2cf8 100644 --- a/benchmarks/benchmark_hidden_state_extraction.py +++ b/benchmarks/benchmark_hidden_state_extraction.py @@ -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() diff --git a/benchmarks/benchmark_long_document_qa_throughput.py b/benchmarks/benchmark_long_document_qa_throughput.py index 6106b9014952..1f5b5f5504b0 100644 --- a/benchmarks/benchmark_long_document_qa_throughput.py +++ b/benchmarks/benchmark_long_document_qa_throughput.py @@ -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 @@ 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": diff --git a/benchmarks/benchmark_prefix_block_hash.py b/benchmarks/benchmark_prefix_block_hash.py index 8bcd8af0d310..522bf9d1ab31 100644 --- a/benchmarks/benchmark_prefix_block_hash.py +++ b/benchmarks/benchmark_prefix_block_hash.py @@ -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 diff --git a/benchmarks/benchmark_prefix_caching.py b/benchmarks/benchmark_prefix_caching.py index 0145f6b7c8e5..5abdd65b2aaf 100644 --- a/benchmarks/benchmark_prefix_caching.py +++ b/benchmarks/benchmark_prefix_caching.py @@ -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 diff --git a/benchmarks/benchmark_serving_structured_output.py b/benchmarks/benchmark_serving_structured_output.py index 664fa58dd49f..dc8e30268e7f 100644 --- a/benchmarks/benchmark_serving_structured_output.py +++ b/benchmarks/benchmark_serving_structured_output.py @@ -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) diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index f727f16ea29c..14111f733b9a 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -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. """ diff --git a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py index 26b04299b353..a4b4533fa3b3 100644 --- a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py +++ b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py @@ -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 diff --git a/benchmarks/kernels/bench_concat_mla_q.py b/benchmarks/kernels/bench_concat_mla_q.py index 8d940484d6b3..c528ab86b295 100644 --- a/benchmarks/kernels/bench_concat_mla_q.py +++ b/benchmarks/kernels/bench_concat_mla_q.py @@ -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") diff --git a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py index 0dd5c6d84882..d0880b3f03cd 100644 --- a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py +++ b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py @@ -106,8 +106,7 @@ def make_impl_kwargs(self, impl_type: ImplType) -> dict[str, Any]: 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() diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 03d7fb386f74..bff4f73e2c0f 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -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() diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index 7379bf858889..a14242732a48 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -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. diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 6c3a726a9753..9f9a150471af 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -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 @@ 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 @@ def benchmark_allreduce( 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 diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index c999c16021bf..aaa39fd9fe06 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/benchmarks/kernels/benchmark_fused_collective.py @@ -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, diff --git a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py index c8cf6f9acc77..dba82a872226 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py +++ b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py @@ -10,6 +10,7 @@ .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \ --output /tmp/attn_res.json + """ from __future__ import annotations diff --git a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py index fcce64015256..6b6d944455e9 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py +++ b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py @@ -16,6 +16,7 @@ Example: python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \ --tokens 1 8 32 64 128 --heads 12 + """ import argparse diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py index 2669dadb348d..855769263d84 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -27,6 +27,7 @@ For multi-node runs, launch one ``torchrun`` agent per node and use a shared rendezvous endpoint. + """ from __future__ import annotations diff --git a/benchmarks/kernels/benchmark_lora.py b/benchmarks/kernels/benchmark_lora.py index ab930c59d219..606741d0474e 100644 --- a/benchmarks/kernels/benchmark_lora.py +++ b/benchmarks/kernels/benchmark_lora.py @@ -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 @@ def mkn( 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 @@ def matmul_shapes( 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 @@ def run_ref_group_gemm( @dataclass class BenchmarkContext: - """ - LoRA benchmark context - """ + """LoRA benchmark context.""" batch_size: int hidden_size: int @@ -467,9 +454,7 @@ def bench_sublabel(self, op_type: OpType) -> str: @dataclass class BenchmarkTensors: - """ - Input/Output tensors used for benchmarks - """ + """Input/Output tensors used for benchmarks.""" # matmul tensors input: torch.Tensor @@ -555,9 +540,7 @@ def make( ) 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 @@ def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None: ) 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 @@ def to_device(tensor: torch.Tensor): ) 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 @@ def moe_lora_align_block_size( 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 @@ def bench_fn_kwargs( 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, diff --git a/benchmarks/kernels/benchmark_mla_k_concat.py b/benchmarks/kernels/benchmark_mla_k_concat.py index 7debf3634804..3cae613e82c7 100644 --- a/benchmarks/kernels/benchmark_mla_k_concat.py +++ b/benchmarks/kernels/benchmark_mla_k_concat.py @@ -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 diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py index f6ad59366dca..f21fbaa703a9 100644 --- a/benchmarks/kernels/benchmark_moe_defaults.py +++ b/benchmarks/kernels/benchmark_moe_defaults.py @@ -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 @@ 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, diff --git a/benchmarks/kernels/benchmark_per_token_group_quant.py b/benchmarks/kernels/benchmark_per_token_group_quant.py index f2195a6d780b..fda3d2a652b4 100644 --- a/benchmarks/kernels/benchmark_per_token_group_quant.py +++ b/benchmarks/kernels/benchmark_per_token_group_quant.py @@ -15,7 +15,7 @@ @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 diff --git a/benchmarks/kernels/benchmark_per_token_quant_fp8.py b/benchmarks/kernels/benchmark_per_token_quant_fp8.py index 6ce97e30368b..0069b06c6456 100644 --- a/benchmarks/kernels/benchmark_per_token_quant_fp8.py +++ b/benchmarks/kernels/benchmark_per_token_quant_fp8.py @@ -16,7 +16,7 @@ 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 diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 4ecd917f31ac..2d5ec80c6aae 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -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 diff --git a/benchmarks/kernels/benchmark_reshape_and_cache.py b/benchmarks/kernels/benchmark_reshape_and_cache.py index 97af4ac976ee..f49bf61036d0 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache.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.") diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py index 55c203725186..84d9cc707c40 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py @@ -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.") diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index ddf1efd018c8..476dce7dff8d 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -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}. """ diff --git a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py index 9fcf278f2ef3..22ef378b0763 100644 --- a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py +++ b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py @@ -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() diff --git a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py index 6bb504b7f090..29ff156b6035 100644 --- a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py +++ b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py @@ -10,6 +10,7 @@ Example: python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \ --seq-lens 2304 4096 8192 16384 --head-dim 72 + """ from types import SimpleNamespace diff --git a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py index 4925e6c02442..2922c449eaf0 100644 --- a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py +++ b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py @@ -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( diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py index 590d4cfdc6d9..3ad7ec888728 100644 --- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py +++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py @@ -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] diff --git a/benchmarks/kernels/ir/bench_ir_ops.py b/benchmarks/kernels/ir/bench_ir_ops.py index b23c4e8ae327..80fdd87acb21 100644 --- a/benchmarks/kernels/ir/bench_ir_ops.py +++ b/benchmarks/kernels/ir/bench_ir_ops.py @@ -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 diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py index 6cc44cf6cec1..bac2920103c2 100644 --- a/benchmarks/kernels/ir/shapes.py +++ b/benchmarks/kernels/ir/shapes.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 diff --git a/benchmarks/kernels/utils.py b/benchmarks/kernels/utils.py index a9af811bbe9c..64d23691bdbb 100644 --- a/benchmarks/kernels/utils.py +++ b/benchmarks/kernels/utils.py @@ -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 diff --git a/benchmarks/multi_turn/bench_dataset.py b/benchmarks/multi_turn/bench_dataset.py index 8cb8a2f386a9..9b159a9ee999 100644 --- a/benchmarks/multi_turn/bench_dataset.py +++ b/benchmarks/multi_turn/bench_dataset.py @@ -148,8 +148,7 @@ def __init__( 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 diff --git a/benchmarks/multi_turn/convert_sharegpt_to_openai.py b/benchmarks/multi_turn/convert_sharegpt_to_openai.py index fccab4d0ce21..b415744873b8 100644 --- a/benchmarks/multi_turn/convert_sharegpt_to_openai.py +++ b/benchmarks/multi_turn/convert_sharegpt_to_openai.py @@ -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: diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py index b7084a21315b..ff7c30669a12 100644 --- a/benchmarks/replayssm/e2e_decode_speedup.py +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -15,6 +15,7 @@ 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 diff --git a/csrc/cpu/generate_cpu_attn_dispatch.py b/csrc/cpu/generate_cpu_attn_dispatch.py index a82a27b844a4..12a14ef58774 100644 --- a/csrc/cpu/generate_cpu_attn_dispatch.py +++ b/csrc/cpu/generate_cpu_attn_dispatch.py @@ -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 diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index 1c7025366f4b..ecdce7761fb4 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -98,8 +98,7 @@ class MockModule: # 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) diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py index 942165eaa940..725d7072ca92 100644 --- a/docs/mkdocs/gen_files/generate_attention_backends.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -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: diff --git a/docs/mkdocs/gen_files/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py index c9c10018ecfb..c051653cf949 100644 --- a/docs/mkdocs/gen_files/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -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 @@ def main_file(self) -> Path | None: 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 @@ def other_files(self) -> list[Path]: """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 @@ def title(self) -> str: 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 @@ def fix_relative_links(self, content: str) -> str: 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 # diff --git a/docs/mkdocs/gen_files/generated_content.py b/docs/mkdocs/gen_files/generated_content.py index fa3f31b294bf..d6691944246d 100644 --- a/docs/mkdocs/gen_files/generated_content.py +++ b/docs/mkdocs/gen_files/generated_content.py @@ -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(): diff --git a/docs/mkdocs/hooks/autoref_code.py b/docs/mkdocs/hooks/autoref_code.py index 647f74f202d8..58cec14d45b3 100644 --- a/docs/mkdocs/hooks/autoref_code.py +++ b/docs/mkdocs/hooks/autoref_code.py @@ -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] diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py index d208b9d6779e..cee418936f47 100644 --- a/docs/mkdocs/hooks/url_schemes.py +++ b/docs/mkdocs/hooks/url_schemes.py @@ -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 @@ def run(self, lines): 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 @@ def replace_relative_link(match: re.Match) -> str: 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") diff --git a/examples/applications/api_server/server.py b/examples/applications/api_server/server.py index e56a7931f666..36474e210184 100644 --- a/examples/applications/api_server/server.py +++ b/examples/applications/api_server/server.py @@ -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. diff --git a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py index d60dbf4d753a..eb18b41a5ad2 100644 --- a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py +++ b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py @@ -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, diff --git a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py index d9a4cadb036e..1681ea146665 100644 --- a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py +++ b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py @@ -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 @@ - 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 @@ 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 @@ def get_qa_prompt(): 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, diff --git a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py index be4796acd1b6..0df32206bf07 100644 --- a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py +++ b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py @@ -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 @@ - 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 @@ 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 diff --git a/examples/deployment/async_llm_streaming.py b/examples/deployment/async_llm_streaming.py index ef69089a0460..bf9b36d05cb1 100644 --- a/examples/deployment/async_llm_streaming.py +++ b/examples/deployment/async_llm_streaming.py @@ -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 @@ 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 diff --git a/examples/deployment/llm_engine_example.py b/examples/deployment/llm_engine_example.py index d9215255a808..0d25a157b416 100644 --- a/examples/deployment/llm_engine_example.py +++ b/examples/deployment/llm_engine_example.py @@ -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. """ diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 830da9b5d64e..c44c9026906b 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -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( diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py index 57deef6a15d2..ae35ec6d7695 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py @@ -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: diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py index cc1cc402d29e..3102f344b027 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py @@ -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 @@ 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 --served-model-name \\ --url http://:8000 \\ @@ -67,6 +66,7 @@ Dependencies: pip install fastapi uvicorn httpx + """ from __future__ import annotations diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py index a35a44fbba2d..53d8b210d677 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py @@ -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: diff --git a/examples/disaggregated/example_connector/decode_example.py b/examples/disaggregated/example_connector/decode_example.py index 2d575840e6a7..9b76e2f00178 100644 --- a/examples/disaggregated/example_connector/decode_example.py +++ b/examples/disaggregated/example_connector/decode_example.py @@ -6,7 +6,7 @@ def read_prompts(): - """Read prompts from output.txt""" + """Read prompts from output.txt.""" prompts = [] try: with open("output.txt") as f: diff --git a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py index b67c2459319f..55f4fe52007d 100644 --- a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py +++ b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py @@ -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. diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py index db9c5a85f7f0..3c3d5dbff985 100644 --- a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py +++ b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py @@ -7,7 +7,7 @@ def read_prompts(): - """Read prompts from prefill_output.txt""" + """Read prompts from prefill_output.txt.""" prompts = [] try: with open("prefill_output.txt") as f: diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py index 5bae04426746..2b2e94aa746b 100644 --- a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py +++ b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py @@ -119,8 +119,7 @@ def get_num_new_matched_tokens( 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. diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache.py b/examples/disaggregated/lmcache/cpu_offload_lmcache.py index b67a929e5d96..8a84675952ab 100644 --- a/examples/disaggregated/lmcache/cpu_offload_lmcache.py +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache.py @@ -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. diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py index e048aecff9f7..3a64139c637a 100644 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py +++ b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py @@ -14,9 +14,7 @@ @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 @@ def parse_args(): 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 diff --git a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py index 489ff1321229..a50d6cbbc30d 100644 --- a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py +++ b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py @@ -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: diff --git a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py index 09880a32aa35..6815f420616c 100644 --- a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py +++ b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py @@ -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, diff --git a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py index 801b4b769792..db033e1861e8 100644 --- a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py +++ b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py @@ -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 diff --git a/examples/features/batch_invariance/reproducibility_offline.py b/examples/features/batch_invariance/reproducibility_offline.py index 72c1e841dca4..54afcdee9f41 100644 --- a/examples/features/batch_invariance/reproducibility_offline.py +++ b/examples/features/batch_invariance/reproducibility_offline.py @@ -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 """ diff --git a/examples/features/context_extension/context_extension_offline.py b/examples/features/context_extension/context_extension_offline.py index 3874288b5e11..cf270a22305f 100644 --- a/examples/features/context_extension/context_extension_offline.py +++ b/examples/features/context_extension/context_extension_offline.py @@ -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. diff --git a/examples/features/data_parallel/data_parallel_offline.py b/examples/features/data_parallel/data_parallel_offline.py index c38ff7297afc..86e61e04f0d7 100644 --- a/examples/features/data_parallel/data_parallel_offline.py +++ b/examples/features/data_parallel/data_parallel_offline.py @@ -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" \ diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index 9c1b5138bc0e..91d3829c54ad 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -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): diff --git a/examples/features/logits_processor/custom.py b/examples/features/logits_processor/custom.py index ce000872dc96..c5028aa797cc 100644 --- a/examples/features/logits_processor/custom.py +++ b/examples/features/logits_processor/custom.py @@ -48,7 +48,7 @@ class object. # 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): diff --git a/examples/features/logits_processor/custom_req.py b/examples/features/logits_processor/custom_req.py index 5763fff5410d..36b83ec95876 100644 --- a/examples/features/logits_processor/custom_req.py +++ b/examples/features/logits_processor/custom_req.py @@ -55,10 +55,11 @@ 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 @@ def __call__( 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 @@ def new_req_logits_processor( Returns: `Callable` request logits processor, or None + """ target_token: Any | None = params.extra_args and params.extra_args.get( "target_token" diff --git a/examples/features/logits_processor/custom_req_init.py b/examples/features/logits_processor/custom_req_init.py index acd2c47f230f..4f29bd271417 100644 --- a/examples/features/logits_processor/custom_req_init.py +++ b/examples/features/logits_processor/custom_req_init.py @@ -56,10 +56,11 @@ 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 @@ def __call__( 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 @@ def new_req_logits_processor( Returns: `Callable` request logits processor, or None + """ if ( not self.is_cuda diff --git a/examples/features/lora/lora_with_quantization_offline.py b/examples/features/lora/lora_with_quantization_offline.py index ef6dbd4534fc..23f165cc5f75 100644 --- a/examples/features/lora/lora_with_quantization_offline.py +++ b/examples/features/lora/lora_with_quantization_offline.py @@ -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) { diff --git a/examples/features/lora/multilora_offline.py b/examples/features/lora/multilora_offline.py index 92021f9fb226..d870ca7562a1 100644 --- a/examples/features/lora/multilora_offline.py +++ b/examples/features/lora/multilora_offline.py @@ -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. diff --git a/examples/features/pause_resume/data_parallel_pause_resume.py b/examples/features/pause_resume/data_parallel_pause_resume.py index 1f11536e5366..fa26d0aa3c66 100644 --- a/examples/features/pause_resume/data_parallel_pause_resume.py +++ b/examples/features/pause_resume/data_parallel_pause_resume.py @@ -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. diff --git a/examples/features/prompt_embed/prompt_embed_offline.py b/examples/features/prompt_embed/prompt_embed_offline.py index 9e90aa46b5b7..2b5d3121c89b 100644 --- a/examples/features/prompt_embed/prompt_embed_offline.py +++ b/examples/features/prompt_embed/prompt_embed_offline.py @@ -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. diff --git a/examples/features/reset_kv/reset_kv_offline.py b/examples/features/reset_kv/reset_kv_offline.py index 3fbe7fa7545e..c19d88a35fbc 100644 --- a/examples/features/reset_kv/reset_kv_offline.py +++ b/examples/features/reset_kv/reset_kv_offline.py @@ -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. """ diff --git a/examples/features/sharded_state/load_sharded_state_offline.py b/examples/features/sharded_state/load_sharded_state_offline.py index e867db5d12fe..66d8a75107aa 100644 --- a/examples/features/sharded_state/load_sharded_state_offline.py +++ b/examples/features/sharded_state/load_sharded_state_offline.py @@ -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: diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py index 675f2e35a53f..7ed0de1f1054 100644 --- a/examples/features/sharded_state/save_sharded_state_offline.py +++ b/examples/features/sharded_state/save_sharded_state_offline.py @@ -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. diff --git a/examples/features/speculative_decoding/mlpspeculator_offline.py b/examples/features/speculative_decoding/mlpspeculator_offline.py index 6a533eb5c937..b863b28222e2 100644 --- a/examples/features/speculative_decoding/mlpspeculator_offline.py +++ b/examples/features/speculative_decoding/mlpspeculator_offline.py @@ -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. diff --git a/examples/features/structured_outputs/structured_outputs_offline.py b/examples/features/structured_outputs/structured_outputs_offline.py index 6b6099f71b12..e716935755c4 100644 --- a/examples/features/structured_outputs/structured_outputs_offline.py +++ b/examples/features/structured_outputs/structured_outputs_offline.py @@ -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. diff --git a/examples/features/torchrun/torchrun_dp_example_offline.py b/examples/features/torchrun/torchrun_dp_example_offline.py index f18f6042e9c6..1f8ac603777d 100644 --- a/examples/features/torchrun/torchrun_dp_example_offline.py +++ b/examples/features/torchrun/torchrun_dp_example_offline.py @@ -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. diff --git a/examples/features/torchrun/torchrun_example_offline.py b/examples/features/torchrun/torchrun_example_offline.py index e41bcd420c20..b4b7cbd6a9d2 100644 --- a/examples/features/torchrun/torchrun_example_offline.py +++ b/examples/features/torchrun/torchrun_example_offline.py @@ -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`, diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py index a2e9b1975602..8b2215d8c1eb 100644 --- a/examples/generate/multimodal/audio_language_offline.py +++ b/examples/generate/multimodal/audio_language_offline.py @@ -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") diff --git a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py index 1ab087de3ff2..b4ea893fc03e 100644 --- a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py +++ b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py @@ -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 diff --git a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py index e1c4fd76d7d1..60aba7abd066 100644 --- a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py +++ b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py @@ -43,7 +43,6 @@ 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") diff --git a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py index 476ca91d9f5e..66bfd490037f 100644 --- a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py +++ b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py @@ -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). """ diff --git a/examples/generate/multimodal/qwen3_omni/only_thinker.py b/examples/generate/multimodal/qwen3_omni/only_thinker.py index 84dd5228e525..57e845ac4f8b 100644 --- a/examples/generate/multimodal/qwen3_omni/only_thinker.py +++ b/examples/generate/multimodal/qwen3_omni/only_thinker.py @@ -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). """ diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py index 8b12df3991f0..cccb75d838ce 100644 --- a/examples/generate/multimodal/vision_language_multi_image_offline.py +++ b/examples/generate/multimodal/vision_language_multi_image_offline.py @@ -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. diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 33e1a8a36b84..7cd471d266d0 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -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 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData: def get_multi_modal_input(args): - """ - return { + """Return { "data": image or video, "question": question, } diff --git a/examples/generate/trace_replay_offline.py b/examples/generate/trace_replay_offline.py index 2baa03307d19..498a46cbd4ae 100644 --- a/examples/generate/trace_replay_offline.py +++ b/examples/generate/trace_replay_offline.py @@ -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 diff --git a/examples/pooling/embed/openai_embedding_long_text/client.py b/examples/pooling/embed/openai_embedding_long_text/client.py index 7cc33b1f24e2..44e415f670bf 100644 --- a/examples/pooling/embed/openai_embedding_long_text/client.py +++ b/examples/pooling/embed/openai_embedding_long_text/client.py @@ -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 diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py index 2c82d0374dfe..5175a6cd46ce 100644 --- a/examples/pooling/embed/vision_embedding_offline.py +++ b/examples/pooling/embed/vision_embedding_offline.py @@ -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 diff --git a/examples/pooling/embed/vision_embedding_online.py b/examples/pooling/embed/vision_embedding_online.py index 79d6e397c8b6..765528978bda 100644 --- a/examples/pooling/embed/vision_embedding_online.py +++ b/examples/pooling/embed/vision_embedding_online.py @@ -37,8 +37,7 @@ def create_chat_embeddings( add_special_tokens: bool = False, padding: str | NotGiven = NOT_GIVEN, ) -> CreateEmbeddingResponse: - """ - Convenience function for accessing vLLM's Chat Embeddings API, + """Convenience function for accessing vLLM's Chat Embeddings API, which is an extension of OpenAI's existing Embeddings API. """ return client.post( @@ -56,13 +55,11 @@ def create_chat_embeddings( def run_clip(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve openai/clip-vit-base-patch32 \ --runner pooling """ - response = create_chat_embeddings( client, messages=[ @@ -97,8 +94,7 @@ def run_clip(client: OpenAI, model: str): def run_dse_qwen2_vl(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve MrLight/dse-qwen2-2b-mrl-v1 \ --runner pooling \ @@ -159,14 +155,12 @@ def run_dse_qwen2_vl(client: OpenAI, model: str): def run_qwen3_vl(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve Qwen/Qwen3-VL-Embedding-2B \ --runner pooling \ --max-model-len 8192 """ - default_instruction = "Represent the user's input." print("Text embedding output:") @@ -266,14 +260,12 @@ def run_qwen3_vl(client: OpenAI, model: str): def run_siglip(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve google/siglip-base-patch16-224 \ --runner pooling \ --chat-template template_basic.jinja """ - response = create_chat_embeddings( client, messages=[ @@ -311,8 +303,7 @@ def run_siglip(client: OpenAI, model: str): def run_vlm2vec(client: OpenAI, model: str): - """ - Start the server using: + """Start the server using: vllm serve TIGER-Lab/VLM2Vec-Full \ --runner pooling \ @@ -320,7 +311,6 @@ def run_vlm2vec(client: OpenAI, model: str): --max-model-len 4096 \ --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja """ - response = create_chat_embeddings( client, messages=[ diff --git a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py index f7b30d9313ba..6bbde74264d8 100644 --- a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py +++ b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py @@ -85,8 +85,7 @@ def generate_datamodule(): def process_channel_group(orig_img, channels): - """ - Args: + """Args: orig_img: torch.Tensor representing original image (reference) with shape = (bands, H, W). channels: list of indices representing RGB channels. @@ -94,8 +93,8 @@ def process_channel_group(orig_img, channels): Returns: torch.Tensor with shape (num_channels, height, width) for original image - """ + """ orig_img = orig_img[channels, ...] valid_mask = torch.ones_like(orig_img, dtype=torch.bool) valid_mask[orig_img == NO_DATA_FLOAT] = False @@ -121,8 +120,8 @@ def read_geotiff(file_path: str): Returns: np.ndarray with shape (bands, height, width) meta info dict - """ + """ with rasterio.open(file_path) as src: img = src.read() meta = src.meta @@ -142,8 +141,8 @@ def save_geotiff(image, output_path: str, meta: dict): image: np.ndarray with shape (bands, height, width) output_path: path where to save the image meta: dict with meta info. - """ + """ with rasterio.open(output_path, "w", **meta) as dest: for i in range(image.shape[0]): dest.write(image[i, :, :], i + 1) @@ -172,12 +171,14 @@ def load_example( images in *file_paths*. std: list containing std values for each band in the images in *file_paths*. + indices: optional band indices to select from each image. All bands + are kept when None. Returns: np.array containing created example list of meta info for each image in *file_paths* - """ + """ imgs = [] metas = [] temporal_coords = [] diff --git a/examples/pooling/reward/sequence_reward_offline.py b/examples/pooling/reward/sequence_reward_offline.py index 0727bceee111..4453f4a390b3 100644 --- a/examples/pooling/reward/sequence_reward_offline.py +++ b/examples/pooling/reward/sequence_reward_offline.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example offline usage of sequence reward models. +"""Example offline usage of sequence reward models. The key distinction between sequence classification and token classification lies in their output granularity: sequence classification produces a single diff --git a/examples/pooling/reward/sequence_reward_online.py b/examples/pooling/reward/sequence_reward_online.py index 40d8d28e3908..c8362265ddeb 100644 --- a/examples/pooling/reward/sequence_reward_online.py +++ b/examples/pooling/reward/sequence_reward_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of sequence reward models. +"""Example online usage of sequence reward models. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/reward/token_reward_offline.py b/examples/pooling/reward/token_reward_offline.py index 4705c0491241..1820c390f8a9 100644 --- a/examples/pooling/reward/token_reward_offline.py +++ b/examples/pooling/reward/token_reward_offline.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example offline usage of token reward models. +"""Example offline usage of token reward models. The key distinction between sequence classification and token classification lies in their output granularity: sequence classification produces a single diff --git a/examples/pooling/reward/token_reward_online.py b/examples/pooling/reward/token_reward_online.py index 64ee0c9dfdcc..38f23944515e 100644 --- a/examples/pooling/reward/token_reward_online.py +++ b/examples/pooling/reward/token_reward_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of token reward models. +"""Example online usage of token reward models. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/score/cohere_rerank_client.py b/examples/pooling/score/cohere_rerank_client.py index b32209967be9..e7bf720cac6d 100644 --- a/examples/pooling/score/cohere_rerank_client.py +++ b/examples/pooling/score/cohere_rerank_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using the OpenAI entrypoint's rerank API which is compatible with +"""Example of using the OpenAI entrypoint's rerank API which is compatible with the Cohere SDK: https://github.com/cohere-ai/cohere-python Note that `pip install cohere` is needed to run this example. diff --git a/examples/pooling/score/colbert_rerank_online.py b/examples/pooling/score/colbert_rerank_online.py index 4cc509b95ac4..b8bc05b9e03c 100644 --- a/examples/pooling/score/colbert_rerank_online.py +++ b/examples/pooling/score/colbert_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColBERT late interaction models for reranking and scoring. +"""Example of using ColBERT late interaction models for reranking and scoring. ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings and MaxSim scoring for document reranking, providing better accuracy than diff --git a/examples/pooling/score/colmodernvbert_rerank_online.py b/examples/pooling/score/colmodernvbert_rerank_online.py index de827ae06260..95ca36ed717c 100644 --- a/examples/pooling/score/colmodernvbert_rerank_online.py +++ b/examples/pooling/score/colmodernvbert_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColModernVBERT late interaction model for reranking. +"""Example of using ColModernVBERT late interaction model for reranking. ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP vision encoder with a ModernBERT text encoder. It produces per-token diff --git a/examples/pooling/score/colqwen3_5_rerank_online.py b/examples/pooling/score/colqwen3_5_rerank_online.py index 00746634d5d1..a2cd635361fd 100644 --- a/examples/pooling/score/colqwen3_5_rerank_online.py +++ b/examples/pooling/score/colqwen3_5_rerank_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using ColQwen3.5 late interaction model for reranking. +"""Example of using ColQwen3.5 late interaction model for reranking. ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5. It produces per-token embeddings and uses MaxSim scoring for retrieval diff --git a/examples/pooling/score/colqwen3_rerank_online.py b/examples/pooling/score/colqwen3_rerank_online.py index 0e61531bfd34..22d5fee0795d 100644 --- a/examples/pooling/score/colqwen3_rerank_online.py +++ b/examples/pooling/score/colqwen3_rerank_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example of using ColQwen3 late interaction model for reranking and scoring. +"""Example of using ColQwen3 late interaction model for reranking and scoring. ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL. It produces per-token embeddings and uses MaxSim scoring for retrieval diff --git a/examples/pooling/score/convert_model_to_seq_cls.py b/examples/pooling/score/convert_model_to_seq_cls.py index 21e294166456..494f683ae709 100644 --- a/examples/pooling/score/convert_model_to_seq_cls.py +++ b/examples/pooling/score/convert_model_to_seq_cls.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Script to convert Large Language Models (LLMs) to Sequence Classification models. +"""Script to convert Large Language Models (LLMs) to Sequence Classification models. This is particularly useful for converting reranker models that use next-token prediction to a sequence classification format for compatibility with standard classification and rerank pipelines. @@ -35,8 +34,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): - """ - This method extracts the difference between weights for 'true' and 'false' tokens + """This method extracts the difference between weights for 'true' and 'false' tokens from the language model head to create a single classification weight vector. Args: @@ -47,6 +45,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): device: Target device (cpu/cuda) Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3 + """ assert len(tokens) == 2, ( "Method requires exactly two tokens for binary classification" @@ -73,8 +72,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device): def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device): - """ - Directly use token weights from the language model head for classification. + """Directly use token weights from the language model head for classification. This method maps each classification label directly to a corresponding token in the vocabulary without additional transformation. @@ -85,6 +83,7 @@ def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device): tokenizer: Model tokenizer tokens: List of tokens representing class labels device: Target device (cpu/cuda) + """ # Get the language model head weights (vocabulary_size x hidden_size) lm_head_weights = causal_lm.lm_head.weight @@ -110,8 +109,7 @@ def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device): def converting( model_name, classifier_from_tokens, path, method, use_sep_token=False, device="cpu" ): - """ - Main conversion function to transform a CausalLM model to SequenceClassification. + """Main conversion function to transform a CausalLM model to SequenceClassification. Args: model_name: Name or path of the pretrained model @@ -120,6 +118,7 @@ def converting( method: Conversion method ('from_2_way_softmax' or 'no_post_processing') use_sep_token: Whether to use separating token in the sequence classification model device: Device to load the model on ('cpu' or 'cuda') + """ assert method in method_map, f"Unknown method: {method}" diff --git a/examples/pooling/score/qwen3_reranker_offline.py b/examples/pooling/score/qwen3_reranker_offline.py index c79ebf97fe64..dbe8d47522c1 100644 --- a/examples/pooling/score/qwen3_reranker_offline.py +++ b/examples/pooling/score/qwen3_reranker_offline.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -What is the difference between the official original version and one +"""What is the difference between the official original version and one that has been converted into a sequence classification model? Qwen3-Reranker is a language model that doing reranker by using the @@ -34,8 +33,7 @@ def get_llm() -> LLM: - """ - Initializes and returns the LLM model for Qwen3-Reranker. + """Initializes and returns the LLM model for Qwen3-Reranker. Returns: LLM: Configured vLLM instance for reranking tasks. @@ -43,6 +41,7 @@ def get_llm() -> LLM: Note: This function loads the ORIGINAL Qwen3-Reranker model with specific overrides to make it compatible with vLLM's score API. + """ return LLM( # Specify the original model from HuggingFace diff --git a/examples/pooling/score/qwen3_reranker_online.py b/examples/pooling/score/qwen3_reranker_online.py index f117b3b5b84d..324e3f5bc349 100644 --- a/examples/pooling/score/qwen3_reranker_online.py +++ b/examples/pooling/score/qwen3_reranker_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -What is the difference between the official original version and one +"""What is the difference between the official original version and one that has been converted into a sequence classification model? Qwen3-Reranker is a language model that doing reranker by using the diff --git a/examples/pooling/score/rerank_api_online.py b/examples/pooling/score/rerank_api_online.py index 908d6a9240aa..caf05862825c 100644 --- a/examples/pooling/score/rerank_api_online.py +++ b/examples/pooling/score/rerank_api_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example of using the OpenAI entrypoint's rerank API which is compatible with +"""Example of using the OpenAI entrypoint's rerank API which is compatible with Jina and Cohere https://jina.ai/reranker run: vllm serve BAAI/bge-reranker-base diff --git a/examples/pooling/score/score_api_online.py b/examples/pooling/score/score_api_online.py index af2886f136ea..a60cf04bf8fb 100644 --- a/examples/pooling/score/score_api_online.py +++ b/examples/pooling/score/score_api_online.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of Score API. +"""Example online usage of Score API. Run `vllm serve --runner pooling` to start up the server in vLLM. """ diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py index f434e699ff0e..abe8157d4c2f 100644 --- a/examples/pooling/score/using_template_offline.py +++ b/examples/pooling/score/using_template_offline.py @@ -68,13 +68,13 @@ def get_chat_template(model: str) -> str: def get_hf_overrides(model: str) -> dict[str, Any]: """Convert Large Language Models (LLMs) to Sequence Classification models. - note: + Note: Some reranking models require special configuration overrides to work correctly with vLLM's score API. Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py - """ + """ model_name_to_hf_overrides_map = { "BAAI/bge-reranker-v2-gemma": { "architectures": ["GemmaForSequenceClassification"], @@ -122,7 +122,6 @@ def get_hf_overrides(model: str) -> dict[str, Any]: def main(args: Namespace): """Main execution function for the reranking example.""" - # Get the overrides for the specified model args.hf_overrides = get_hf_overrides(args.model) diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py index f0bfa7d15769..f7dc109dd81f 100644 --- a/examples/pooling/score/using_template_online.py +++ b/examples/pooling/score/using_template_online.py @@ -1,15 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example of using the rerank API with template. +"""Example of using the rerank API with template. This script demonstrates how to interact with a vLLM server running a reranking model via the REST API. Before running this script, start the vLLM server with one of the supported reranking models using the commands below. -note: +Note: Some reranking models require special configuration overrides to work correctly with vLLM's score API. Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py @@ -21,6 +20,7 @@ vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja + """ import json diff --git a/examples/pooling/score/vision_rerank_api_online.py b/examples/pooling/score/vision_rerank_api_online.py index dce2efd1dc07..cc1f2e97bf2f 100644 --- a/examples/pooling/score/vision_rerank_api_online.py +++ b/examples/pooling/score/vision_rerank_api_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example Python client for multimodal rerank API which is compatible with +"""Example Python client for multimodal rerank API which is compatible with Jina and Cohere https://jina.ai/reranker Run `vllm serve --runner pooling` to start up the server in vLLM. diff --git a/examples/pooling/score/vision_reranker_offline.py b/examples/pooling/score/vision_reranker_offline.py index ef2b35487771..24c8b9068f45 100644 --- a/examples/pooling/score/vision_reranker_offline.py +++ b/examples/pooling/score/vision_reranker_offline.py @@ -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 vision language reranker models for multimodal scoring tasks. Vision language rerankers score the relevance between a text query and diff --git a/examples/pooling/score/vision_score_api_online.py b/examples/pooling/score/vision_score_api_online.py index 543d4bfa2147..3a3ac62cf7fd 100644 --- a/examples/pooling/score/vision_score_api_online.py +++ b/examples/pooling/score/vision_score_api_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of Score API. +"""Example online usage of Score API. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/token_classify/forced_alignment_offline.py b/examples/pooling/token_classify/forced_alignment_offline.py index e97a3735767e..13facab4806d 100644 --- a/examples/pooling/token_classify/forced_alignment_offline.py +++ b/examples/pooling/token_classify/forced_alignment_offline.py @@ -3,8 +3,7 @@ # Adapted from Qwen3-ForcedAligner inference: # https://github.com/QwenLM/Qwen3-ASR -""" -Offline forced alignment example using Qwen3-ForcedAligner-0.6B. +"""Offline forced alignment example using Qwen3-ForcedAligner-0.6B. Forced alignment takes audio and reference text as input and produces word-level timestamps. The model predicts a time bin at each diff --git a/examples/pooling/token_classify/forced_alignment_online.py b/examples/pooling/token_classify/forced_alignment_online.py index e967a8ba70d8..ab3026160d5f 100644 --- a/examples/pooling/token_classify/forced_alignment_online.py +++ b/examples/pooling/token_classify/forced_alignment_online.py @@ -3,8 +3,7 @@ # Adapted from Qwen3-ForcedAligner inference: # https://github.com/QwenLM/Qwen3-ASR -""" -Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf. +"""Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf. Forced alignment takes audio and reference text as input and produces word-level timestamps. The model predicts a time bin at each diff --git a/examples/pooling/token_classify/ner_online.py b/examples/pooling/token_classify/ner_online.py index 9ec2bd45a0fe..bd60d6862600 100644 --- a/examples/pooling/token_classify/ner_online.py +++ b/examples/pooling/token_classify/ner_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://huggingface.co/boltuix/NeuroBERT-NER -""" -Example online usage of Pooling API for Named Entity Recognition (NER). +"""Example online usage of Pooling API for Named Entity Recognition (NER). Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/pooling/token_embed/colqwen3_token_embed_online.py b/examples/pooling/token_embed/colqwen3_token_embed_online.py index cac11188e87e..c1a8f3a56ef8 100644 --- a/examples/pooling/token_embed/colqwen3_token_embed_online.py +++ b/examples/pooling/token_embed/colqwen3_token_embed_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of Pooling API for ColQwen3 multi-vector retrieval. +"""Example online usage of Pooling API for ColQwen3 multi-vector retrieval. ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that produces per-token embeddings (320-dim, L2-normalized) for both text and diff --git a/examples/pooling/token_embed/jina_reranker_v3_online.py b/examples/pooling/token_embed/jina_reranker_v3_online.py index 8350aee2f143..b9018f0f23f8 100644 --- a/examples/pooling/token_embed/jina_reranker_v3_online.py +++ b/examples/pooling/token_embed/jina_reranker_v3_online.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Example online usage of the Jina Reranker v3 score and rerank APIs with a task +"""Example online usage of the Jina Reranker v3 score and rerank APIs with a task instruction. Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the diff --git a/examples/pooling/token_embed/multi_vector_retrieval_online.py b/examples/pooling/token_embed/multi_vector_retrieval_online.py index c2aae88294e8..f36e3dc31d10 100644 --- a/examples/pooling/token_embed/multi_vector_retrieval_online.py +++ b/examples/pooling/token_embed/multi_vector_retrieval_online.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example online usage of Pooling API for multi vector retrieval. +"""Example online usage of Pooling API for multi vector retrieval. Run `vllm serve --runner pooling` to start up the server in vLLM. e.g. diff --git a/examples/ray_serving/batch_llm_inference.py b/examples/ray_serving/batch_llm_inference.py index 22408dc95033..1b265efb1766 100644 --- a/examples/ray_serving/batch_llm_inference.py +++ b/examples/ray_serving/batch_llm_inference.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This example shows how to use Ray Data for data parallel batch inference. +"""This example shows how to use Ray Data for data parallel batch inference. Ray Data is a data processing framework that can process very large datasets with first-class support for vLLM. diff --git a/examples/ray_serving/ray_serve_deepseek.py b/examples/ray_serving/ray_serve_deepseek.py index af53443b9101..a1e9348115a4 100644 --- a/examples/ray_serving/ray_serve_deepseek.py +++ b/examples/ray_serving/ray_serve_deepseek.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Deploy DeepSeek R1 or V3 with Ray Serve LLM. +"""Deploy DeepSeek R1 or V3 with Ray Serve LLM. Ray Serve LLM is a scalable and production-grade model serving library built on the Ray distributed computing framework and first-class support for the vLLM engine. diff --git a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py index 1dfc3084646d..53b30f534d47 100644 --- a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py +++ b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example demonstrates how to use tool calling with reasoning models +"""An example demonstrates how to use tool calling with reasoning models like QwQ-32B. The reasoning will not be parsed by the tool calling process; only the final output will be parsed. diff --git a/examples/reasoning/openai_chat_completion_with_reasoning.py b/examples/reasoning/openai_chat_completion_with_reasoning.py index 87043897b058..a72606108f8e 100644 --- a/examples/reasoning/openai_chat_completion_with_reasoning.py +++ b/examples/reasoning/openai_chat_completion_with_reasoning.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example shows how to generate chat completions from reasoning models +"""An example shows how to generate chat completions from reasoning models like DeepSeekR1. To run this example, you need to start the vLLM server diff --git a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py index 8e262701b720..36f4c1f8e6e2 100644 --- a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py +++ b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -An example shows how to generate chat completions from reasoning models +"""An example shows how to generate chat completions from reasoning models like DeepSeekR1. To run this example, you need to start the vLLM server with the reasoning diff --git a/examples/reasoning/openai_responses_client.py b/examples/reasoning/openai_responses_client.py index b4eb24671507..94ecd84468d5 100644 --- a/examples/reasoning/openai_responses_client.py +++ b/examples/reasoning/openai_responses_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server. +"""Set up this example by starting a vLLM OpenAI-compatible server. Reasoning models can be used through the Responses API as seen here https://platform.openai.com/docs/api-reference/responses For example: diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index ea7099962f21..081023156edf 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Demonstrates async reinforcement learning using vLLM and Ray, +"""Demonstrates async reinforcement learning using vLLM and Ray, with native weight syncing APIs and batch-invariant generation. The script separates training and inference workloads onto distinct GPUs diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py index c688707df984..0ba9350216f6 100644 --- a/examples/rl/rlhf_http_ipc.py +++ b/examples/rl/rlhf_http_ipc.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the +"""RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the data plane. * OpenAI-compatible API for inference requests diff --git a/examples/rl/rlhf_http_nccl.py b/examples/rl/rlhf_http_nccl.py index 97b0db0b0b83..d095b647aafb 100644 --- a/examples/rl/rlhf_http_nccl.py +++ b/examples/rl/rlhf_http_nccl.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the +"""RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the data plane. * OpenAI-compatible API for inference requests diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py index e3fa1c44f15e..4eb3fbcbbb31 100644 --- a/examples/rl/rlhf_ipc_fsdp_ep.py +++ b/examples/rl/rlhf_ipc_fsdp_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC** +"""RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC** weight transfer and **packed** tensors. Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors diff --git a/examples/rl/rlhf_nccl_fsdp_ep.py b/examples/rl/rlhf_nccl_fsdp_ep.py index e4380bea948d..b9b72c012ec8 100644 --- a/examples/rl/rlhf_nccl_fsdp_ep.py +++ b/examples/rl/rlhf_nccl_fsdp_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). +"""RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). 8-GPU layout: Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors @@ -74,8 +73,7 @@ @ray.remote(num_gpus=1) class FSDPTrainWorker: - """ - One FSDP2 training worker per GPU. Four of these form the FSDP group. + """One FSDP2 training worker per GPU. Four of these form the FSDP group. Rank 0 additionally drives weight transfer to the vLLM server. """ diff --git a/examples/rl/rlhf_sharded_rdt_small_ep.py b/examples/rl/rlhf_sharded_rdt_small_ep.py index 213eae55ef22..ed31a32ad8f1 100644 --- a/examples/rl/rlhf_sharded_rdt_small_ep.py +++ b/examples/rl/rlhf_sharded_rdt_small_ep.py @@ -101,7 +101,8 @@ class FSDPTrainWorker: """One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray - actor with no producer mixin.""" + actor with no producer mixin. + """ def __init__(self, model_name, rank, world_size, master_addr, master_port): self.rank = rank diff --git a/examples/rl/routed_experts_e2e.py b/examples/rl/routed_experts_e2e.py index 1666bc3ffe16..04b5871fd1e9 100644 --- a/examples/rl/routed_experts_e2e.py +++ b/examples/rl/routed_experts_e2e.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end example for routed experts capture with hybrid models. +"""End-to-end example for routed experts capture with hybrid models. Validates that: 1. routed_experts is returned in CompletionOutput for MoE models. diff --git a/examples/speech_to_text/openai/openai_transcription_client.py b/examples/speech_to_text/openai/openai_transcription_client.py index f928c06d45ea..cc6757bd9635 100644 --- a/examples/speech_to_text/openai/openai_transcription_client.py +++ b/examples/speech_to_text/openai/openai_transcription_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script demonstrates how to use the vLLM API server to perform audio +"""This script demonstrates how to use the vLLM API server to perform audio transcription with the `openai/whisper-large-v3` model. Before running this script, you must start the vLLM server with the following command: @@ -35,8 +34,7 @@ def sync_openai( hotwords: str = None, prompt: str | None = None, ): - """ - Perform synchronous transcription using OpenAI-compatible API. + """Perform synchronous transcription using OpenAI-compatible API. The optional ``prompt`` is the OpenAI-API ``prompt`` field (style / vocabulary hint). It is wired through model-by-model: Whisper uses it @@ -69,9 +67,7 @@ async def stream_openai_response( hotwords: str = None, prompt: str | None = None, ): - """ - Perform asynchronous transcription using OpenAI-compatible API. - """ + """Perform asynchronous transcription using OpenAI-compatible API.""" print("\ntranscription result [stream]:", end=" ") with open(audio_path, "rb") as f: transcription = await client.audio.transcriptions.create( @@ -98,9 +94,7 @@ async def stream_openai_response( def stream_api_response(audio_path: str, model: str, openai_api_base: str): - """ - Perform streaming transcription using raw HTTP requests to the vLLM API server. - """ + """Stream a transcription using raw HTTP requests to the vLLM API server.""" import json import os diff --git a/examples/speech_to_text/realtime/openai_realtime_client.py b/examples/speech_to_text/realtime/openai_realtime_client.py index fda3d7cb4564..a290b63e83dc 100644 --- a/examples/speech_to_text/realtime/openai_realtime_client.py +++ b/examples/speech_to_text/realtime/openai_realtime_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script demonstrates how to use the vLLM Realtime WebSocket API to perform +"""This script demonstrates how to use the vLLM Realtime WebSocket API to perform audio transcription by uploading an audio file. Before running this script, you must start the vLLM server with a realtime-capable @@ -34,9 +33,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str: - """ - Load an audio file and convert it to base64-encoded PCM16 @ 16kHz. - """ + """Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.""" # Load audio and resample to 16kHz mono audio, _ = load_audio(audio_path, sr=16000, mono=True) # Convert to PCM16 @@ -46,9 +43,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str: async def realtime_transcribe(audio_path: str, host: str, port: int, model: str): - """ - Connect to the Realtime API and transcribe an audio file. - """ + """Connect to the Realtime API and transcribe an audio file.""" uri = f"ws://{host}:{port}/v1/realtime" async with websockets.connect(uri) as ws: diff --git a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py index a3c07673ffbe..c29486120c03 100644 --- a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py +++ b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API. +"""Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API. Start the vLLM server first: diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools.py b/examples/tool_calling/openai_chat_completion_client_with_tools.py index 0bd1d05322f8..508c6900b614 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled. For example: IMPORTANT: for mistral, you must use one of the provided mistral tool call diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py index c00d712b351d..f38011e10d9c 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -To run this example, you can start the vLLM server +"""To run this example, you can start the vLLM server without any specific flags: ```bash diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py index f0b0a2db44ed..004b5fc7fcce 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled for xLAM-2 models: vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam @@ -105,8 +104,7 @@ def translate_text(text: str, target_language: str): def process_response(response, tool_functions, original_query): - """Process a non-streaming response with possible tool calls""" - + """Process a non-streaming response with possible tool calls.""" print("\n--- Response Output ---") # Check if the response has content @@ -195,7 +193,7 @@ def process_response(response, tool_functions, original_query): def run_test_case(query, test_name): - """Run a single test case with the given query""" + """Run a single test case with the given query.""" print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}") print(f"Query: '{query}'") diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py index 94e664c9ec3d..43e4ed21a0d6 100644 --- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py +++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled for xLAM-2 models: vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam @@ -107,7 +106,7 @@ def translate_text(text: str, target_language: str): def process_stream(response, tool_functions, original_query): - """Process a streaming response with possible tool calls""" + """Process a streaming response with possible tool calls.""" # Track multiple tool calls tool_calls = {} # Dictionary to store tool calls by ID @@ -223,7 +222,7 @@ def process_stream(response, tool_functions, original_query): def run_test_case(query, test_name): - """Run a single test case with the given query""" + """Run a single test case with the given query.""" print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}") print(f"Query: '{query}'") diff --git a/examples/tool_calling/openai_responses_client_with_mcp_tools.py b/examples/tool_calling/openai_responses_client_with_mcp_tools.py index 1e685a4f3440..0b6ea412bf56 100644 --- a/examples/tool_calling/openai_responses_client_with_mcp_tools.py +++ b/examples/tool_calling/openai_responses_client_with_mcp_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Example demonstrating MCP (Model Context Protocol) tools with the Responses API. +"""Example demonstrating MCP (Model Context Protocol) tools with the Responses API. This example shows how to use MCP tools with different allowed_tools configurations: 1. No filter (allows all tools from the MCP server) diff --git a/examples/tool_calling/openai_responses_client_with_tools.py b/examples/tool_calling/openai_responses_client_with_tools.py index a47210f60020..6530f331e45d 100644 --- a/examples/tool_calling/openai_responses_client_with_tools.py +++ b/examples/tool_calling/openai_responses_client_with_tools.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Set up this example by starting a vLLM OpenAI-compatible server with tool call +"""Set up this example by starting a vLLM OpenAI-compatible server with tool call options enabled. Reasoning models can be used through the Responses API as seen here https://platform.openai.com/docs/api-reference/responses @@ -17,8 +16,7 @@ def get_weather(latitude: float, longitude: float) -> str: - """ - Mock function to simulate getting weather data. + """Mock function to simulate getting weather data. In a real application, this would call an external weather API. """ return f"Current temperature at ({latitude}, {longitude}) is 20°C." diff --git a/pyproject.toml b/pyproject.toml index 1a77433f9cb8..83949db8fe89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,9 +60,13 @@ include = ["vllm*"] "vllm/third_party/**" = ["ALL"] "vllm/version.py" = ["F401"] "vllm/_version.py" = ["ALL"] +# test args are mostly pytest fixtures, which don't benefit from descriptions +"tests/**" = ["D417"] [tool.ruff.lint] select = [ + # pydocstyle + "D", # pycodestyle "E", # Pyflakes @@ -83,6 +87,10 @@ select = [ "INP", ] ignore = [ + # don't enforce that everything has a docstring + "D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107", + # unwanted pydocstyle rules + "D205", "D301", "D400", "D401", "D404", "D415", # star imports "F405", "F403", # lambda expression assignment @@ -95,6 +103,9 @@ ignore = [ "UP032", ] +[tool.ruff.lint.pydocstyle] +ignore-var-parameters = true + [tool.ruff.format] docstring-code-format = true diff --git a/scripts/autotune_helion_kernels.py b/scripts/autotune_helion_kernels.py index 864aabfc0964..e6352a40af41 100644 --- a/scripts/autotune_helion_kernels.py +++ b/scripts/autotune_helion_kernels.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Autotune registered Helion kernels for optimal configurations. +"""Autotune registered Helion kernels for optimal configurations. Usage: # Autotune all registered kernels diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py index fea30d6f9cfa..04350896af8b 100644 --- a/scripts/benchmark_helion_kernels.py +++ b/scripts/benchmark_helion_kernels.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark a registered Helion kernel against a baseline. +"""Benchmark a registered Helion kernel against a baseline. For each input case produced by the kernel's registered input generator, this checks the Helion kernel's numerics once against its eager reference, measures @@ -394,6 +393,7 @@ def check_kernel_correctness( ``passed=False`` with the exception text in ``error``; iteration continues regardless. An empty input mapping returns an empty list, which the CLI reports as a skipped check. + """ if inputs_dict is None: inputs_dict = kernel.get_inputs() diff --git a/setup.py b/setup.py index 4012000e634e..48cbdcf763b4 100644 --- a/setup.py +++ b/setup.py @@ -512,8 +512,7 @@ def fetch_metadata_for_variant( *, rocm: bool = False, ) -> tuple[list[dict], str]: - """ - Fetches metadata for a specific variant of the precompiled wheel. + """Fetches metadata for a specific variant of the precompiled wheel. For non-ROCm, fetches vllm metadata. @@ -591,7 +590,6 @@ def is_rocm_system() -> bool: @staticmethod def detect_system_cuda_variant() -> str: """Auto-detect CUDA variant from torch, nvidia-smi, or env default.""" - # Map CUDA major version to hosted wheel variants on wheels.vllm.ai supported = {12: "cu129", 13: "cu130"} @@ -726,7 +724,8 @@ def warn_if_rocm_torch_version_mismatch( wheels: list[dict], repo_url: str, arch: str ) -> None: """Warn if installed torch differs from the custom ROCm build on - wheels.vllm.ai and suggest the correct install command.""" + wheels.vllm.ai and suggest the correct install command. + """ try: installed = torch.__version__ except Exception: @@ -886,8 +885,7 @@ def determine_wheel_url_rocm() -> tuple[str, str | None]: @staticmethod def determine_wheel_url() -> tuple[str, str | None]: - """ - Try to determine the precompiled wheel URL or path to use. + """Try to determine the precompiled wheel URL or path to use. The order of preference is: 1. user-specified wheel location (can be either local or remote, via VLLM_PRECOMPILED_WHEEL_LOCATION) diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 3412b30814e3..1b3c92ee5b60 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/tests/basic_correctness/test_basic_correctness.py @@ -89,7 +89,7 @@ def _resolve_target_test_suite() -> str: condition=current_platform.is_rocm(), ) def test_vllm_gc_ed(): - """Verify vllm instance is GC'ed when it is deleted""" + """Verify vllm instance is GC'ed when it is deleted.""" llm = LLM("hmellor/tiny-random-LlamaForCausalLM") weak_llm = weakref.ref(llm) del llm diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index c3f66e24abc5..f04bb397f068 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -51,8 +51,7 @@ def create_and_map_with_poison(handle) -> None: @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_python_error(): - """ - Test if Python error occurs when there's low-level + """Test if Python error occurs when there's low-level error happening from the C++ side. """ allocator = get_mem_allocator_instance() @@ -113,7 +112,8 @@ def test_basic_cumem(): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_discard_tags(): """Test that discard(tags) selectively frees GPU memory for specific - tags while keeping other tags mapped and usable.""" + tags while keeping other tags mapped and usable. + """ allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): diff --git a/tests/benchmarks/test_bfcl_dataset.py b/tests/benchmarks/test_bfcl_dataset.py index d19192231979..982b93ad8e89 100644 --- a/tests/benchmarks/test_bfcl_dataset.py +++ b/tests/benchmarks/test_bfcl_dataset.py @@ -13,7 +13,8 @@ def _patch_hf_api(side_effect): """Return a patch context that swaps `hf_api()` to a stub whose - `.hf_hub_download` attribute uses `side_effect`.""" + `.hf_hub_download` attribute uses `side_effect`. + """ fake_api = MagicMock() fake_api.hf_hub_download.side_effect = side_effect return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api) @@ -115,7 +116,8 @@ def test_bfcl_dataset_translates_schema_and_attaches_tools( ) -> None: """BFCLDataset should translate schemas to OpenAI tool format, set `messages` directly on SampleRequest, and attach tools/tool_choice via - request_overrides.""" + request_overrides. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -166,7 +168,8 @@ def test_bfcl_dataset_missing_category_raises_clear_error( hf_tokenizer: PreTrainedTokenizerBase, ) -> None: """A typo'd category should produce an actionable ValueError, not an - opaque huggingface_hub exception.""" + opaque huggingface_hub exception. + """ from huggingface_hub.errors import EntryNotFoundError args = _args_for_bfcl(categories=["simpl"]) # typo @@ -184,7 +187,8 @@ def raise_missing(_repo, filename, **_kwargs): @pytest.mark.benchmark def test_chat_backend_uses_messages_field_when_set() -> None: """When RequestFuncInput.chat_messages is set, the chat backend must use - it verbatim and skip default content construction from `prompt`.""" + it verbatim and skip default content construction from `prompt`. + """ import asyncio from vllm.benchmarks.lib.endpoint_request_func import ( @@ -244,7 +248,8 @@ def post(self, url, json, headers): # noqa: A002 def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None: """prompt_len must reflect tokens from both messages *and* tool schemas, so percentile buckets and input-distribution summaries aren't biased - low for tool-heavy traffic.""" + low for tool-heavy traffic. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -291,7 +296,8 @@ def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools( tmp_path: Path, ) -> None: """Older tokenizers don't accept tools=; fallback must still produce a - non-zero prompt_len without crashing.""" + non-zero prompt_len without crashing. + """ paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): diff --git a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py index 149ea6625a9a..28ecb6e443ec 100644 --- a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py +++ b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py @@ -12,7 +12,8 @@ class _RecordingTokenizer: """Minimal tokenizer stub that records the kwargs forwarded to apply_chat_template, so we can assert chat_template_kwargs propagation - without loading a real model/template.""" + without loading a real model/template. + """ def __init__(self) -> None: self.captured_kwargs: dict | None = None diff --git a/tests/benchmarks/test_endpoint_request_func_timing.py b/tests/benchmarks/test_endpoint_request_func_timing.py index 34394e3b014a..57ac5824ab49 100644 --- a/tests/benchmarks/test_endpoint_request_func_timing.py +++ b/tests/benchmarks/test_endpoint_request_func_timing.py @@ -228,7 +228,7 @@ def test_decode_span_identity_is_exact( api_url: str, chat: bool, ) -> None: - """latency - ttft must equal sum(itl) exactly. + """Latency - ttft must equal sum(itl) exactly. On a scripted clock this can only drift if a code path reads the clock an extra time or attributes a read to the wrong event. diff --git a/tests/benchmarks/test_sampling_params.py b/tests/benchmarks/test_sampling_params.py index 3bc34a84b377..13dcf8ad208c 100644 --- a/tests/benchmarks/test_sampling_params.py +++ b/tests/benchmarks/test_sampling_params.py @@ -10,7 +10,8 @@ class _FakeTokenizer(TokenizerLike): """Minimal tokenizer implementing the TokenizerLike protocol - for testing get_sampling_params.""" + for testing get_sampling_params. + """ def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None: self._vocab_size = vocab_size @@ -171,7 +172,8 @@ def test_output_range_bounds(self): def test_output_low_clamped_to_one(self): """Even with a high ratio that would push output_low to 0, - the function clamps it to 1.""" + the function clamps it to 1. + """ rng = np.random.default_rng(0) # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1 _, output_lens, _ = get_sampling_params( @@ -233,7 +235,8 @@ def test_invalid_dict_missing_keys(self): def test_input_len_zero_with_special_tokens(self): """input_len < num_special_tokens → real_input_len = 0, which is fine - (range [0, 0]).""" + (range [0, 0]). + """ rng = np.random.default_rng(0) input_lens, _, _ = get_sampling_params( rng, 5, 0.0, 5, 50, self._tok(num_special=10) diff --git a/tests/benchmarks/test_skip_tokenizer_init.py b/tests/benchmarks/test_skip_tokenizer_init.py index 741db2e0d626..1eb9a2536eb6 100644 --- a/tests/benchmarks/test_skip_tokenizer_init.py +++ b/tests/benchmarks/test_skip_tokenizer_init.py @@ -52,7 +52,8 @@ def _write_dataset(path: Path) -> None: def _args(dataset_path: str) -> argparse.Namespace: """Reproduce the argparse.Namespace that serve.py builds from the - failing command, including skip_tokenizer_init=True.""" + failing command, including skip_tokenizer_init=True. + """ return argparse.Namespace( # dataset dataset_name="custom", diff --git a/tests/ci_envs.py b/tests/ci_envs.py index f3a54f308cd8..171fac996a68 100644 --- a/tests/ci_envs.py +++ b/tests/ci_envs.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -These envs only work for a small part of the tests, fix what you need! -""" +"""These envs only work for a small part of the tests, fix what you need!""" import os from collections.abc import Callable diff --git a/tests/compile/backend.py b/tests/compile/backend.py index 87f98946a8ad..167de6ad68a3 100644 --- a/tests/compile/backend.py +++ b/tests/compile/backend.py @@ -29,8 +29,7 @@ class LazyInitPass(InductorPass): - """ - If there's a pass that we want to initialize lazily in a test, + """If there's a pass that we want to initialize lazily in a test, we can wrap it in LazyInitPass, which will initialize the pass when invoked and then immediately invoke it. """ @@ -45,8 +44,7 @@ def __call__(self, graph: fx.Graph) -> None: class TestBackend: - """ - This class provides a simple Inductor backend that can be used for testing. + """This class provides a simple Inductor backend that can be used for testing. It takes a list of custom passes and runs them after Inductor's passes. It also saves the graph before and after the custom passes for inspection. diff --git a/tests/compile/conftest.py b/tests/compile/conftest.py index 7d15b5c47e55..972ec3fef52d 100644 --- a/tests/compile/conftest.py +++ b/tests/compile/conftest.py @@ -11,8 +11,7 @@ @pytest.fixture def mock_cuda_platform(): - """ - Fixture that returns a factory for creating mocked CUDA platforms. + """Fixture that returns a factory for creating mocked CUDA platforms. Usage: def test_something(mock_cuda_platform): @@ -51,8 +50,7 @@ def is_device_capability_family( @pytest.fixture def mock_xpu_platform(): - """ - Fixture that returns a factory for creating mocked XPU platforms. + """Fixture that returns a factory for creating mocked XPU platforms. Usage: def test_something(mock_xpu_platform): diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py index eed5efb7b36b..e5f8d1870eb0 100644 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ b/tests/compile/correctness_e2e/test_sequence_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py index c2af7be048af..d91cd84816e5 100644 --- a/tests/compile/fullgraph/test_full_cudagraph.py +++ b/tests/compile/fullgraph/test_full_cudagraph.py @@ -17,8 +17,7 @@ @contextlib.contextmanager def temporary_environ(env_vars): - """ - Temporarily set environment variables and restore them afterward. + """Temporarily set environment variables and restore them afterward. We have to do this vs monkeypatch because monkeypatch doesn't work with "module" scoped fixtures. """ @@ -125,8 +124,7 @@ def llm_pair(request): indirect=True, ) class TestFullCUDAGraph: - """ - Use a class such that an llm pair is constructed once for all + """Use a class such that an llm pair is constructed once for all batch_size/max_tokens combinations and released immediately after. Module-scope fixtures would stick around the whole time, @@ -149,11 +147,9 @@ class TestFullCUDAGraph: ], ) def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]): - """ - Test various batch sizes and max_tokens to ensure that the + """Test various batch sizes and max_tokens to ensure that the full cudagraph compilation works for padded cases too. """ - full_cudagraph_llm, piecewise_llm = llm_pair prompts = ["the quick brown fox"] * batch_size diff --git a/tests/compile/fullgraph/test_multiple_graphs.py b/tests/compile/fullgraph/test_multiple_graphs.py index 6d3788af9de0..e2656633344f 100644 --- a/tests/compile/fullgraph/test_multiple_graphs.py +++ b/tests/compile/fullgraph/test_multiple_graphs.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test (piecewise) compilation with a simple model where multiple submodules +"""Test (piecewise) compilation with a simple model where multiple submodules are compiled and graph captured separately. """ diff --git a/tests/compile/fullgraph/test_simple.py b/tests/compile/fullgraph/test_simple.py index f1ea0e414d76..4d10a92a076f 100644 --- a/tests/compile/fullgraph/test_simple.py +++ b/tests/compile/fullgraph/test_simple.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test the piecewise compilation with a simple model so that we +"""Test the piecewise compilation with a simple model so that we can exactly calculate the expected output and side effects. """ @@ -52,8 +51,7 @@ def __init__( self.intermediate_unbacked = intermediate_unbacked def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Overall effect: + """Overall effect: x = 3 * x + 19 global_counter += 2 """ diff --git a/tests/compile/fullgraph/test_toy_llama.py b/tests/compile/fullgraph/test_toy_llama.py index 69c758702e8a..7823c362f6c6 100644 --- a/tests/compile/fullgraph/test_toy_llama.py +++ b/tests/compile/fullgraph/test_toy_llama.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test the piecewise compilation with a simple model, comparing the output +"""Test the piecewise compilation with a simple model, comparing the output with and without the piecewise compilation. This is a tractable model, the weights and computation are specially designed diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 7b547c9720d1..ff0b04896280 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -56,10 +56,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16): torch.nn.init.normal_(self.gate_proj, std=0.02) def forward(self, hidden_states): - """ - Forward pass implementing the mm + reduce scatter in the FX graph - - """ + """Forward pass implementing the mm + reduce scatter in the FX graph.""" # Reshape input view = hidden_states.reshape(-1, self.hidden_size) @@ -88,9 +85,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16): torch.nn.init.normal_(self.weight, std=0.02) def forward(self, hidden_states): - """ - Forward pass implementing the mm + all gather in the FX graph - """ + """Forward pass implementing the mm + all gather in the FX graph.""" # Reshape input view = hidden_states.reshape(-1, self.hidden_size) all_gather = tensor_model_parallel_all_gather(view, dim=0) @@ -122,10 +117,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16): class TestScaledMMRSModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the scaled_mm + reduce scatter in the FX graph - - """ + """Forward pass implementing the scaled_mm + reduce scatter in the FX graph.""" fp8_input = input.to(FP8_DTYPE) scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) scaled_mm = torch._scaled_mm( @@ -147,9 +139,7 @@ def ops_in_model_after(self): class TestAGScaledMMModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the all gather + scaled_mm in the FX graph - """ + """Forward pass implementing the all gather + scaled_mm in the FX graph.""" # Reshape input fp8_input = input.to(FP8_DTYPE) all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) @@ -173,8 +163,7 @@ def ops_in_model_after(self): class TestCutlassScaledMMRSModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the cutlass_scaled_mm + reduce scatter + """Forward pass implementing the cutlass_scaled_mm + reduce scatter in the FX graph """ @@ -200,8 +189,7 @@ def ops_in_model_after(self): class TestAGCutlassScaledMMModel(_BaseScaledMMModel): def forward(self, input: torch.Tensor): - """ - Forward pass implementing the all gather + cutlass_scaled_mm + """Forward pass implementing the all gather + cutlass_scaled_mm in the FX graph """ # Reshape input diff --git a/tests/compile/passes/ir/test_clone_cleanup.py b/tests/compile/passes/ir/test_clone_cleanup.py index b6626a5cb70a..8a51e6d53988 100644 --- a/tests/compile/passes/ir/test_clone_cleanup.py +++ b/tests/compile/passes/ir/test_clone_cleanup.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Comprehensive tests for UnsafeCloneEliminationPass. +"""Comprehensive tests for UnsafeCloneEliminationPass. This test suite exercises all possible valid FX graph patterns involving clones: 1. Clone with no users (dead code) diff --git a/tests/compile/passes/ir/test_inplace_functionalization.py b/tests/compile/passes/ir/test_inplace_functionalization.py index 1e8d5662162f..773ea93ce609 100644 --- a/tests/compile/passes/ir/test_inplace_functionalization.py +++ b/tests/compile/passes/ir/test_inplace_functionalization.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for IR inplace functionalization pass integration. +"""Tests for IR inplace functionalization pass integration. This test suite verifies that the inplace functionalization pass, lowering pass, and clone cleanup pass work together correctly with donated buffer tracking. @@ -108,8 +107,7 @@ def forward( class ModelWithTritonAfterMaybeInplace(nn.Module): - """ - Model using maybe_inplace followed by a Triton kernel. + """Model using maybe_inplace followed by a Triton kernel. Test clone elimination can handle Triton in the graph """ @@ -410,8 +408,7 @@ def inner(*args): reason="Only test on cuda and rocm platform", ) def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache): - """ - Test piecewise compilation with donated buffers across graph splits. + """Test piecewise compilation with donated buffers across graph splits. Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching. """ torch.set_default_device(current_platform.device_type) diff --git a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py index 6a620d11a492..f05f821852c8 100644 --- a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py +++ b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the DoubleQuant fan-out variants registered by +"""Unit tests for the DoubleQuant fan-out variants registered by ``RocmAiterRMSNormQuantFusionPass``. Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds @@ -90,8 +89,7 @@ def test_double_aiter_rms_fp8_group_quant_fusion( model_cls: type[torch.nn.Module], monkeypatch: pytest.MonkeyPatch, ) -> None: - """ - Both fan-out shapes (with and without an intermediate view) must fuse + """Both fan-out shapes (with and without an intermediate view) must fuse into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via ``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling. diff --git a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py index 6f20d4e15874..0f631b9472ad 100644 --- a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py +++ b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit test for the MLADualRMSNormFusionPass. +"""Unit test for the MLADualRMSNormFusionPass. The pass fuses paired q/kv RMS norms in MLA attention into a single fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel. @@ -38,13 +37,12 @@ class MLADualRMSNormTestModel(torch.nn.Module): - """ - Minimal model reproducing the MLA dual RMS norm pattern: - linear -> split([q_dim, kv_dim]) - +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear - +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) - +-- kv_c (getitem 0) -> rms_norm(kv_w, eps) - +-- k_pe + """Minimal model reproducing the MLA dual RMS norm pattern: + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm(kv_w, eps) + +-- k_pe """ def __init__( @@ -155,13 +153,12 @@ def test_fuse_mla_dual_rms_norm( class MLADualRMSNormFp8PerTokenTestModel(torch.nn.Module): - """ - Minimal model reproducing the FP8 MLA attention path with *per-token* quant: - linear -> split([q_dim, kv_dim]) - +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant - +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) - +-- kv_c (getitem 0) -> rms_norm (bf16) - +-- k_pe + """Minimal model reproducing the FP8 MLA attention path with *per-token* quant: + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm (bf16) + +-- k_pe """ def __init__( diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 843052f4afdc..0945bd622a6c 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -732,7 +732,8 @@ def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers( monkeypatch: pytest.MonkeyPatch, ): """Verify that without GDN layers in static_forward_context, - the gated pattern is not registered and no matches occur.""" + the gated pattern is not registered and no matches occur. + """ group_shape = GroupShape(1, 128) vllm_config = VllmConfig( model_config=ModelConfig(dtype=dtype), diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 4df02cab9543..37f0d4126cce 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -107,7 +107,6 @@ def __init__( def build_attn_metadata(self, batch_size: int) -> AttentionMetadata: """Initialize attention metadata.""" - # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py # Create common attn metadata @@ -304,7 +303,7 @@ def test_attention_quant_pattern( dist_init, disable_vllm_compile_cache, ): - """Test AttentionStaticQuantPattern fusion pass""" + """Test AttentionStaticQuantPattern fusion pass.""" if backend == AttentionBackendEnum.FLASHINFER and ( not current_platform.is_device_capability((10, 0)) or not has_flashinfer() ): diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 11c455bfa2b5..f1de3dde0cbb 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -141,7 +141,6 @@ def build_attn_metadata(self, batch_size: int) -> AttentionMetadata: the full output[:num_actual_toks] buffer after both forward_mha and forward_mqa have written their results. """ - batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) common_attn_metadata = create_common_attn_metadata( batch_spec, self.block_size, self.device, arange_block_indices=True @@ -408,7 +407,7 @@ def test_mla_attention_quant_pattern( dist_init, disable_vllm_compile_cache, ): - """Test MLA AttentionQuantPattern fusion pass""" + """Test MLA AttentionQuantPattern fusion pass.""" if ( model_class is TestMLAAttentionNvfp4QuantPatternModel and not is_nvfp4_supported() diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py index ee9496084420..10dc401a26a1 100644 --- a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for manual fusion via maybe_fused_act_quant. +"""Tests for manual fusion via maybe_fused_act_quant. Tests all fusion paths in _FUSED_ACT_QUANT: - kFp8StaticTensorSym: all platforms diff --git a/tests/compile/passes/test_split_coalescing.py b/tests/compile/passes/test_split_coalescing.py index ab7a0be1a215..52e531e12742 100644 --- a/tests/compile/passes/test_split_coalescing.py +++ b/tests/compile/passes/test_split_coalescing.py @@ -15,7 +15,8 @@ class SplitCoalescingModel(torch.nn.Module): """Model with 3 separate split_with_sizes calls on the same input, - simulating the B200+FP8 graph where CSE fails to merge them.""" + simulating the B200+FP8 graph where CSE fails to merge them. + """ def __init__(self, q_size: int, kv_size: int) -> None: super().__init__() diff --git a/tests/compile/silly_attention.py b/tests/compile/silly_attention.py index 29c02f6e6a1d..0d3313512912 100644 --- a/tests/compile/silly_attention.py +++ b/tests/compile/silly_attention.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared PyTorch custom silly attention for compilation tests. +"""Shared PyTorch custom silly attention for compilation tests. Centralizes custom operation definitions to avoid duplicate registrations. """ @@ -21,12 +20,12 @@ def get_global_counter(): - """Get the current global counter value""" + """Get the current global counter value.""" return _global_counter def reset_global_counter(): - """Reset the global counter to 0""" + """Reset the global counter to 0.""" global _global_counter _global_counter = 0 @@ -34,8 +33,7 @@ def reset_global_counter(): def silly_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor ) -> None: - """ - Unified attention implementation that depends on + """Unified attention implementation that depends on all inputs and affects the output. Always increments a global counter that tests can use or ignore. """ @@ -51,7 +49,7 @@ def silly_attention( def silly_attention_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor ) -> None: - """Fake implementation for testing""" + """Fake implementation for testing.""" return diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index 90c132bee1ab..e5c11b8acc8a 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -268,8 +268,7 @@ def foo(x: torch.Tensor): def test_cache_load_returns_tuple_consistency_tuple_output( monkeypatch: pytest.MonkeyPatch, ): - """ - Test that cache loading correctly handles models that return tuples. + """Test that cache loading correctly handles models that return tuples. This verifies that when a model returns a tuple of tensors, the output type is preserved as a tuple between fresh compilation and cache load. @@ -335,8 +334,7 @@ def test_cache_load_returns_tuple_consistency_tuple_output( @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_shape_env(monkeypatch: pytest.MonkeyPatch): - """ - Test that the shape environment is correctly serialized and preserved + """Test that the shape environment is correctly serialized and preserved when loading from cache. """ with monkeypatch.context() as m: @@ -374,8 +372,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch): def test_partition_wrapper_applied_on_aot_load( monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker ): - """ - Test that partition wrappers are applied when loading AOT cached functions. + """Test that partition wrappers are applied when loading AOT cached functions. This test verifies the fix for GitHub issue #31439 where AOT compile caused 2x latency regression when use_inductor_graph_partition=True. @@ -498,14 +495,12 @@ def test_standalone_compile_correctness(): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") @create_new_process_for_each_test("spawn") def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch): - """ - Test that compiling gpt2 twice results in a cache hit. + """Test that compiling gpt2 twice results in a cache hit. Counter values are read from the EngineCore subprocess via ``LLM.collective_rpc`` so the test works under default V1 multiprocessing (no shared memory between test and engine). """ - from vllm import LLM def _snap(self): @@ -641,7 +636,7 @@ def test_num_artifacts_and_entries(self): @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_load_all_success(self, mock_deserialize): - """Test successful loading of all artifacts""" + """Test successful loading of all artifacts.""" cache = StandaloneCompiledArtifacts() mock_artifact1 = Mock() @@ -658,7 +653,7 @@ def test_load_all_success(self, mock_deserialize): @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_load_all_already_loaded(self, mock_deserialize): - """Test that load_all skips if already loaded""" + """Test that load_all skips if already loaded.""" cache = StandaloneCompiledArtifacts() mock_artifact = Mock() @@ -671,7 +666,7 @@ def test_load_all_already_loaded(self, mock_deserialize): @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize") def test_get_loaded_artifact(self, mock_deserialize): - """Test retrieving loaded artifacts""" + """Test retrieving loaded artifacts.""" cache = StandaloneCompiledArtifacts() mock_artifact = Mock() diff --git a/tests/compile/test_codegen.py b/tests/compile/test_codegen.py index 21db287a2478..a0086cb702e5 100644 --- a/tests/compile/test_codegen.py +++ b/tests/compile/test_codegen.py @@ -77,7 +77,8 @@ def test_non_primitive_kwargs_lifted_to_consts( into the namespace, so ``device`` is unbound. The fix collects such objects into ``__vllm_consts__`` and references them by index. The unqualified ``device(type=...)`` form must never appear in the - generated source.""" + generated source. + """ split_gm = _trace_and_split(model_fn, (x,), split_ops) code, submod_names, consts = generate_execution_code(split_gm) @@ -98,7 +99,8 @@ def test_dtype_singleton_deduped(x: torch.Tensor) -> None: """``torch.float16`` is a process-wide singleton, so two ops referring to it in the traced graph share a single consts slot via ``id()``-based dedup. Distinct expressions (``x.to(...)`` vs ``(x*2).to(...)``) ensure - the tracer can't CSE the two ops into a single node.""" + the tracer can't CSE the two ops into a single node. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.to(dtype=torch.float16) + (x * 2).to(dtype=torch.float16) @@ -147,7 +149,8 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_consts_ordering_deterministic(x: torch.Tensor) -> None: """Two independent traces of the same model must produce equal consts lists *in the same order*. Cache artifacts identify const slots by - index, so a non-deterministic order would invalidate cached code.""" + index, so a non-deterministic order would invalidate cached code. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: # Multiple distinct non-primitives encountered in a fixed graph order. @@ -185,7 +188,8 @@ def test_consts_shared_across_split_submods(x: torch.Tensor) -> None: The function below splits into three inlined submods, two of which independently reference ``torch.float16``. The shared ``const_index`` threaded through recursive ``generate_execution_code_with_name`` calls - must collapse the dtype to a single slot used from both submods.""" + must collapse the dtype to a single slot used from both submods. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: a = x.to(dtype=torch.float16) # submod_0: _to_copy(fp16) @@ -230,7 +234,8 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None: """When a child of split_gm is *not* a ``torch.fx.GraphModule`` — as happens in production once ``PiecewiseBackend`` replaces submods — codegen emits ``__vllm_submods__[idx](...)`` instead of inlining, and - the runtime callable is bound from ``submod_callables``.""" + the runtime callable is bound from ``submod_callables``. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -280,7 +285,8 @@ def forward(self, *args, **kwargs): def test_getitem_in_stitching_graph(x: torch.Tensor) -> None: """``operator.getitem`` on submod tuple returns is the ``call_function`` special case at codegen.py — emitted as ``name = source[index]`` - rather than a function call.""" + rather than a function call. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -298,7 +304,8 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_del_emitted_for_intermediate_values(x: torch.Tensor) -> None: """The codegen schedules ``del`` after a value's last use to free memory early. Multi-submod splits naturally have intermediates whose - last use is not the output node.""" + last use is not the output node. + """ def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid().tanh() @@ -318,7 +325,8 @@ def test_with_submod_false_rejects_call_module() -> None: """``generate_execution_code_with_name(with_submod=False)`` is the recursive entry for inlining a GraphModule into its parent. It must refuse a graph that itself contains ``call_module`` nodes — the parent - is responsible for handling those.""" + is responsible for handling those. + """ g = fx.Graph() x_node = g.placeholder("x") root = torch.nn.Module() @@ -335,7 +343,8 @@ def test_node_ref_recurses_through_containers() -> None: """``_node_ref`` is the recursive walker that lifts non-primitives nested inside list/tuple/dict args. Real aten ops rarely produce such structures, but the path is needed for DTensor placement lists and - other future cases — unit-test the walker directly.""" + other future cases — unit-test the walker directly. + """ consts: list = [] const_index: dict[int, int] = {} cpu = torch.device("cpu") @@ -359,7 +368,8 @@ def test_node_ref_recurses_through_containers() -> None: def test_legacy_code_without_consts() -> None: """``compile_execution_fn(consts=None)`` must still load code that has no ``__vllm_consts__`` reference, so older serialized cache artifacts - keep working.""" + keep working. + """ # Pre-consts codegen: no __vllm_consts__ reference, only torch/operator. legacy_code = ( "import torch\n" diff --git a/tests/compile/test_compile_ranges.py b/tests/compile/test_compile_ranges.py index 4dfea42a6b4a..b3e2ec4731a8 100644 --- a/tests/compile/test_compile_ranges.py +++ b/tests/compile/test_compile_ranges.py @@ -129,7 +129,8 @@ def test_compile_config_get_compile_ranges(): class PostGradStaticShapeChecker(InductorPass): """Asserts that compile_sizes entries produce graphs with fully concrete - (non-symbolic) shapes, and compile_ranges entries have symbolic shapes.""" + (non-symbolic) shapes, and compile_ranges entries have symbolic shapes. + """ def __init__(self): self.num_static_calls = 0 @@ -170,7 +171,8 @@ def uuid(self) -> str: def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache): """Verify that compile_sizes entries are compiled with fully concrete - shapes (no SymInts), while compile_ranges entries retain dynamic shapes.""" + shapes (no SymInts), while compile_ranges entries retain dynamic shapes. + """ checker = PostGradStaticShapeChecker() torch.set_default_device("cuda") vllm_config = VllmConfig( diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5c9132ccbfe2..fc3295621806 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -1229,7 +1229,8 @@ def test_compile_sizes_padding_validation(): def test_inductor_asserts_default_disabled(monkeypatch): """Test that inductor runtime asserts are disabled by default - (INFO logging level) on torch < 2.12.""" + (INFO logging level) on torch < 2.12. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib @@ -1247,7 +1248,8 @@ def test_inductor_asserts_default_disabled(monkeypatch): def test_inductor_asserts_enabled_in_debug(monkeypatch): """Test that VLLM_LOGGING_LEVEL=DEBUG enables inductor runtime asserts - on torch < 2.12.""" + on torch < 2.12. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "DEBUG") import importlib @@ -1283,7 +1285,8 @@ def test_get_inductor_factors_includes_configs(): def test_inductor_asserts_user_override(monkeypatch): """Test that explicit inductor_compile_config overrides the - debug-logging default.""" + debug-logging default. + """ monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 71f66fce4045..ea2ccafbe7e0 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -22,7 +22,7 @@ def get_test_models(): - """Get list of models to test based on PyTorch version""" + """Get list of models to test based on PyTorch version.""" models = [ "Qwen/Qwen3-0.6B", "openai-community/gpt2", @@ -172,7 +172,6 @@ def test_model_specialization_with_evaluate_guards( """Test that evaluate_guards correctly detects shape specialization violations. """ - if ( use_aot_compile and dynamic_shapes_type == DynamicShapesType.BACKED diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index bb2a6f2aee55..f5aaffe00e4d 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -24,8 +24,7 @@ def test_getitem_moved_to_producer_subgraph(): - """ - Test that getitem operations are moved to the same subgraph as their input, + """Test that getitem operations are moved to the same subgraph as their input, preventing tuple inputs to submodules. """ @@ -79,8 +78,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_no_tuple_inputs_with_multiple_consumers(): - """ - Test that when a tuple is consumed by multiple split operations, + """Test that when a tuple is consumed by multiple split operations, getitem operations are properly moved to avoid tuple inputs. """ @@ -138,13 +136,10 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_consecutive_ops_in_split(): - """ - Test that consecutive splitting operations are grouped into the same subgraph - """ + """Test that consecutive splitting operations are grouped into the same subgraph.""" def model_fn(x: torch.Tensor) -> torch.Tensor: - """ - Define a simple model where consecutive operations create opportunities + """Define a simple model where consecutive operations create opportunities for splitting subgraphs. """ # Apply silly attention followed by consecutive operations @@ -211,8 +206,7 @@ def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph): def test_empty_only_partition_stays_separate_after_splitting_predecessor(): - """ - Empty-only subgraphs should not be merged when the only predecessor is + """Empty-only subgraphs should not be merged when the only predecessor is a splitting-op subgraph. """ @@ -248,8 +242,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_empty_only_partition_is_merged(): - """ - Empty-only subgraphs should still be merged when a non-splitting predecessor + """Empty-only subgraphs should still be merged when a non-splitting predecessor exists. The merged empty node must remain outside splitting-op subgraphs. """ @@ -294,8 +287,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_builtin_empty_only_partition_is_merged(): - """ - In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets + """In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets (not aten OpOverload). Ensure empty-only partitions are still merged. """ @@ -340,8 +332,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_sym_size_whole_shape_boundary(): - """ - Test that using x.size() (whole shape) across a split boundary can be + """Test that using x.size() (whole shape) across a split boundary can be compiled by standalone_compile. The dynamo graph looks like: @@ -400,8 +391,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_symint_crosses_split_boundary(): - """ - Test that SymInt placeholders from torch.compile + mark_dynamic + """Test that SymInt placeholders from torch.compile + mark_dynamic cross split boundaries safely via split_module's natural threading. SymInt values are threaded through subgraphs by split_module and @@ -459,8 +449,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_shape_boundary_standalone_compile(): - """ - Repro for the original production bug: + """Repro for the original production bug: AssertionError: out_spec mismatch TreeSpec(tuple, None, [*, *, TreeSpec(Size, None, [*, *]), *]) @@ -529,8 +518,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_size_used_in_multiple_consumer_subgraphs(): - """ - Validates that x.size() (whole shape) used by multiple downstream subgraphs + """Validates that x.size() (whole shape) used by multiple downstream subgraphs does not cause torch.Size to cross split boundaries. Model: @@ -584,8 +572,7 @@ def model_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_sym_size_metadata_propagated(): - """ - Validates that new sym_size.int nodes created by the pre-pass have + """Validates that new sym_size.int nodes created by the pre-pass have example_value metadata set. Without it, placeholder metadata in consumer subgraphs would be None, breaking any code that dynamically builds example inputs from metadata (e.g. standalone_compile per-submodule). @@ -634,8 +621,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_decompose_size_with_getitem_user(): - """ - Regression test: _decompose_size_nodes must handle getitem users of size() + """Regression test: _decompose_size_nodes must handle getitem users of size() correctly. When a graph contains x.shape[i], it can appear as: @@ -704,8 +690,7 @@ def test_decompose_size_with_getitem_user(): def test_decompose_size_leaves_scalar_size_with_dim(): - """ - Regression test: _decompose_size_nodes must leave x.size(dim) alone. + """Regression test: _decompose_size_nodes must leave x.size(dim) alone. x.size() returns a torch.Size tuple that can't cross split boundaries and must be decomposed. x.size(dim), however, already returns a scalar diff --git a/tests/compile/test_inductor_fallback_allow_list_patch.py b/tests/compile/test_inductor_fallback_allow_list_patch.py index 29fe9962e345..c5705da21171 100644 --- a/tests/compile/test_inductor_fallback_allow_list_patch.py +++ b/tests/compile/test_inductor_fallback_allow_list_patch.py @@ -236,7 +236,8 @@ def test_kimi_k2_6_style_op_stream(self): def test_inner_set_membership_still_works_for_standard_ops(self): """The patch must not break Inductor's existing fallback decisions - for non-vllm ops such as ``torchvision::roi_align``.""" + for non-vllm ops such as ``torchvision::roi_align``. + """ import torch._inductor.lowering as _lowering allow_list = _lowering.FALLBACK_ALLOW_LIST diff --git a/tests/compile/test_structured_logging.py b/tests/compile/test_structured_logging.py index 10b0ed139cfe..98f092bfce44 100644 --- a/tests/compile/test_structured_logging.py +++ b/tests/compile/test_structured_logging.py @@ -60,6 +60,7 @@ def get(self, event_type: str, name_pattern: str) -> list[dict]: Args: event_type: The event type to filter by (e.g., "artifact", "graph_dump") name_pattern: Regex pattern to match against the artifact name + """ regex = re.compile(name_pattern) return [ diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 590e97c1eb54..1d0f4d915bcf 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -40,7 +40,7 @@ class DummyLogprobsMode(Enum): def test_hash_factors_deterministic(): - """Test that hash_factors produces consistent SHA-256 hashes""" + """Test that hash_factors produces consistent SHA-256 hashes.""" factors = {"a": 1, "b": "test"} hash1 = hash_factors(factors) hash2 = hash_factors(factors) diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index b68820c0d51c..f71d4875d75c 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -372,7 +372,8 @@ class _HeterogeneousConfig(PretrainedConfig): """A heterogeneous config with no convertor of its own. Mirrors the parts vLLM uses: per-layer configs are shallow copies with the - varying attributes applied and heterogeneity stripped, so they do not recurse.""" + varying attributes applied and heterogeneity stripped, so they do not recurse. + """ is_heterogeneous = True @@ -471,7 +472,8 @@ def test_deepseek_v4_convertor_splits_vision_architecture(): """The vision checkpoint shares model_type/architectures with the text-only DeepSeek-V4; the convertor routes it to the VL wrapper class by rewriting hf_config.architectures (model-class resolution reads the - raw attribute).""" + raw attribute). + """ from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.transformers_utils.model_arch_config_convertor import ( DeepseekV4ModelArchConfigConvertor, diff --git a/tests/config/test_mp_reducer.py b/tests/config/test_mp_reducer.py index 56dc542f1c76..56e4e474740e 100644 --- a/tests/config/test_mp_reducer.py +++ b/tests/config/test_mp_reducer.py @@ -9,12 +9,10 @@ def test_mp_reducer(): - """ - Test that _reduce_config reducer is registered when AsyncLLM is instantiated + """Test that _reduce_config reducer is registered when AsyncLLM is instantiated without transformers_modules. This is a regression test for https://github.com/vllm-project/vllm/pull/18640. """ - # Ensure transformers_modules is not in sys.modules if "transformers_modules" in sys.modules: del sys.modules["transformers_modules"] diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 84cf092c2159..abf834608a55 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -42,7 +42,8 @@ def test_mm_encoder_attn_backend_hash_updates(): def test_language_model_only_does_not_affect_mm_hash(): """language_model_only does not affect the ViT computation graph, - so it should not change the multimodal config hash.""" + so it should not change the multimodal config hash. + """ base_hash = MultiModalConfig().compute_hash() lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash() assert base_hash == lm_only_hash @@ -50,7 +51,8 @@ def test_language_model_only_does_not_affect_mm_hash(): def test_language_model_only_affects_model_hash(): """language_model_only affects the LM computation graph, - so it should change the model config hash.""" + so it should change the model config hash. + """ model = "llava-hf/llava-1.5-7b-hf" base_hash = ModelConfig(model).compute_hash() lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash() @@ -386,7 +388,8 @@ def _resolve_mm_video_decode_device( torchcodec_available: bool = True, ) -> dict: """Run processor-device then video-decode-device resolution and report - the resulting video media IO kwargs.""" + the resulting video media IO kwargs. + """ mm_config = MultiModalConfig( mm_processor_kwargs={} if device is None else {"device": device}, mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type] @@ -416,7 +419,8 @@ def _resolve_mm_video_decode_device( def test_auto_video_decode_uses_nvdec_on_encoder_instance(): """The processor runs on the accelerator there, so decoded frames should - stay on-device too.""" + stay on-device too. + """ assert _resolve_mm_video_decode_device(ec_role="ec_producer") == { "backend": "torchcodec", "device": "cuda", @@ -428,7 +432,8 @@ def test_auto_video_decode_stays_on_cpu_off_encoder_instance( ec_role: ECRole | None, ): """The processor stays on CPU off encode-only instances, so video - decoding should too.""" + decoding should too. + """ assert _resolve_mm_video_decode_device(ec_role=ec_role) == {} diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 8505172da387..10f8c994f44b 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -34,7 +34,8 @@ def _make_hf_config(**kwargs) -> PretrainedConfig: @pytest.mark.cpu_test def test_dict_overrides_are_not_forwarded_to_draft(): """Dict overrides are target-specific key patches; the draft must get - only the architecture-mapping override.""" + only the architecture-mapping override. + """ composed = SpeculativeConfig.compose_draft_hf_overrides( {"max_position_embeddings": 1234} ) @@ -50,7 +51,8 @@ def test_none_overrides_fall_back_to_arch_mapping(): @pytest.mark.cpu_test def test_callable_overrides_reach_the_draft_config(): """A callable override (config-to-config transform) composes with the - architecture-mapping override and is applied to the draft config.""" + architecture-mapping override and is applied to the draft config. + """ def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.num_hidden_layers = 1 @@ -67,7 +69,8 @@ def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: @pytest.mark.cpu_test def test_arch_mapping_applies_before_callable_override(): """The static arch-mapping override runs first, so the user callable - observes (and may adjust) the post-mapping config.""" + observes (and may adjust) the post-mapping config. + """ seen_architectures: list[str] = [] def record(hf_config: PretrainedConfig) -> PretrainedConfig: @@ -128,7 +131,8 @@ def test_composed_override_is_picklable(): the composed override must be picklable. A nested local closure is not (it raised ``Can't get local object`` on DFlashDraftModel); a ``functools.partial`` over a module-referenceable static method is. - Guard against regressing to a closure.""" + Guard against regressing to a closure. + """ composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) assert isinstance(composed, functools.partial) diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py index 117ae5a66d53..711aaf9b036b 100644 --- a/tests/config/test_speculative_draft_max_position_embeddings.py +++ b/tests/config/test_speculative_draft_max_position_embeddings.py @@ -99,7 +99,8 @@ def test_independent_draft_model_keeps_its_own_limit( vllm_caplog: pytest.LogCaptureFixture, ): """An independent AR draft may genuinely have a smaller context than the - target; its max_position_embeddings must not be resized.""" + target; its max_position_embeddings must not be resized. + """ target_model_config = ModelConfig( AR_MODEL, hf_overrides={"max_position_embeddings": 8192} ) diff --git a/tests/cuda/test_cuda_context.py b/tests/cuda/test_cuda_context.py index aca490d15644..cb2ba1b19922 100644 --- a/tests/cuda/test_cuda_context.py +++ b/tests/cuda/test_cuda_context.py @@ -11,7 +11,7 @@ def check_cuda_context(): - """Check CUDA driver context status""" + """Check CUDA driver context status.""" try: cuda = ctypes.CDLL("libcuda.so") device = ctypes.c_int() @@ -22,7 +22,7 @@ def check_cuda_context(): def run_cuda_test_in_thread(device_input, expected_device_id): - """Run CUDA context test in separate thread for isolation""" + """Run CUDA context test in separate thread for isolation.""" try: # New thread should have no CUDA context initially valid_before, device_before = check_cuda_context() diff --git a/tests/cuda/test_platform_no_cuda_init.py b/tests/cuda/test_platform_no_cuda_init.py index 697d2cfec8e1..6f50ce606b76 100644 --- a/tests/cuda/test_platform_no_cuda_init.py +++ b/tests/cuda/test_platform_no_cuda_init.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that platform imports do not prematurely initialize CUDA. +"""Test that platform imports do not prematurely initialize CUDA. This is critical for Ray-based multi-GPU setups where workers need to set CUDA_VISIBLE_DEVICES after importing vLLM but before CUDA is initialized. diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py index 44215cb72ae1..3b72d612dd1c 100644 --- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py +++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py @@ -44,8 +44,7 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0): def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: bool): - """ - This test verifies that the detokenizer correctly handles the case where + """This test verifies that the detokenizer correctly handles the case where the generated token sequence contains both: - a stop token - an token @@ -67,7 +66,6 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo where stop-string evaluation occurs before the early-return on stop_terminated. """ - # Generate text "abcdeZ" and tokenize it. generated_text = "abcde" eos_token = "Z" diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index 48df856f5e75..876ee9dcf518 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -18,13 +18,13 @@ @pytest.fixture def random_port(): - """Generate a random port number for testing""" + """Generate a random port number for testing.""" return random.randint(10000, 59900) @pytest.fixture def publisher_config(random_port, request): - """Create a publisher config with inproc transport""" + """Create a publisher config with inproc transport.""" how = request.param if hasattr(request, "param") else "inproc" if how == "inproc": @@ -47,7 +47,7 @@ def publisher_config(random_port, request): @pytest.fixture def publisher(publisher_config): - """Create and return a publisher instance""" + """Create and return a publisher instance.""" pub = EventPublisherFactory.create(publisher_config, DP_RANK) yield pub pub.shutdown() @@ -55,7 +55,7 @@ def publisher(publisher_config): @pytest.fixture def subscriber(publisher_config): - """Create and return a subscriber for testing""" + """Create and return a subscriber for testing.""" endpoint = publisher_config.endpoint replay_endpoint = publisher_config.replay_endpoint @@ -74,7 +74,7 @@ def subscriber(publisher_config): class MockSubscriber: - """Helper class to receive and verify published events""" + """Helper class to receive and verify published events.""" def __init__( self, @@ -113,7 +113,7 @@ def __init__( self.decoder = msgspec.msgpack.Decoder(type=decode_type) def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None: - """Receive a single message with timeout""" + """Receive a single message with timeout.""" if not self.sub.poll(timeout): return None @@ -127,7 +127,7 @@ def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None: return seq, data def request_replay(self, start_seq: int, socket_idx: int = 0) -> None: - """Request replay of messages starting from start_seq""" + """Request replay of messages starting from start_seq.""" if not self.replay_sockets: raise ValueError("Replay sockets not initialized") if socket_idx >= len(self.replay_sockets): @@ -138,7 +138,7 @@ def request_replay(self, start_seq: int, socket_idx: int = 0) -> None: ) def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: - """Receive replayed messages from a specific replay socket""" + """Receive replayed messages from a specific replay socket.""" if not self.replay_sockets: raise ValueError("Replay sockets not initialized") if socket_idx >= len(self.replay_sockets): @@ -170,7 +170,7 @@ def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: return replayed def close(self): - """Clean up resources""" + """Clean up resources.""" self.sub.close() for replay in self.replay_sockets: replay.close() @@ -179,7 +179,8 @@ def close(self): @pytest.fixture def enable_ray_v2_backend(): """Set env vars for the Ray V2 executor backend and shut down Ray - between tests.""" + between tests. + """ import ray saved = { diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index 007948aeba82..6d5d454b4a8d 100644 --- a/tests/distributed/test_context_parallel.py +++ b/tests/distributed/test_context_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/distributed/test_eplb_algo.py b/tests/distributed/test_eplb_algo.py index 20360c8c61b9..aaa987b9fc42 100644 --- a/tests/distributed/test_eplb_algo.py +++ b/tests/distributed/test_eplb_algo.py @@ -29,7 +29,7 @@ def test_eplb_load_stats_reduce_across_ranks(): def test_basic_rebalance(): - """Test basic rebalancing functionality""" + """Test basic rebalancing functionality.""" # Example from https://github.com/deepseek-ai/eplb weight = torch.tensor( [ @@ -92,7 +92,7 @@ def test_basic_rebalance(): def test_single_gpu_case(): - """Test single GPU case""" + """Test single GPU case.""" weight = torch.tensor([[10, 20, 30, 40]]) num_replicas = 4 num_groups = 1 @@ -115,7 +115,7 @@ def test_single_gpu_case(): def test_equal_weights(): - """Test case with equal weights""" + """Test case with equal weights.""" weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]]) num_replicas = 8 num_groups = 2 @@ -139,7 +139,7 @@ def test_equal_weights(): def test_extreme_weight_imbalance(): - """Test extreme weight imbalance case""" + """Test extreme weight imbalance case.""" weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]]) num_replicas = 12 num_groups = 2 @@ -162,7 +162,7 @@ def test_extreme_weight_imbalance(): def test_multiple_layers(): - """Test multiple layers case""" + """Test multiple layers case.""" weight = torch.tensor( [ [10, 20, 30, 40, 50, 60], # First layer @@ -195,7 +195,7 @@ def test_multiple_layers(): def test_parameter_validation(): - """Test parameter validation""" + """Test parameter validation.""" weight = torch.tensor([[10, 20, 30, 40]]) # Test non-divisible case - this should handle normally without throwing @@ -213,7 +213,7 @@ def test_parameter_validation(): def test_small_scale_hierarchical(): - """Test small-scale hierarchical load balancing""" + """Test small-scale hierarchical load balancing.""" weight = torch.tensor( [ [100, 50, 200, 75, 150, 25, 300, 80], # 8 experts @@ -243,7 +243,7 @@ def test_small_scale_hierarchical(): def test_global_load_balance_fallback(): - """Test global load balancing fallback case""" + """Test global load balancing fallback case.""" # When num_groups % num_nodes != 0, should fall back to global load # balancing weight = torch.tensor([[10, 20, 30, 40, 50, 60]]) @@ -265,7 +265,7 @@ def test_global_load_balance_fallback(): @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_device_compatibility(device): - """Test device compatibility""" + """Test device compatibility.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") @@ -287,8 +287,7 @@ def test_device_compatibility(device): def test_additional_cases(): - """Test more edge cases and different parameter combinations""" - + """Test more edge cases and different parameter combinations.""" # Test case 1: Large-scale distributed setup weight1 = torch.tensor( [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]] @@ -320,8 +319,7 @@ def test_additional_cases(): def test_compute_logical_maps_with_negative_indices(): - """ - Test that compute_logical_maps correctly handles physical slots containing + """Test that compute_logical_maps correctly handles physical slots containing -1 (unused slots). """ # 2 layers, 6 physical slots, 4 logical experts. diff --git a/tests/distributed/test_eplb_events.py b/tests/distributed/test_eplb_events.py index 40323dc89809..963e3cbb77f8 100644 --- a/tests/distributed/test_eplb_events.py +++ b/tests/distributed/test_eplb_events.py @@ -57,9 +57,8 @@ def waiter(): def test_producer_consumer(): - """ - This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU - tensor on multiple CPU threads. + """This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared + GPU tensor on multiple CPU threads. """ worker_stream = torch.cuda.Stream() # Create a single element counter that will be shared between two threads diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 1ff41f5df269..aed0decc6d35 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -56,8 +56,7 @@ def create_expert_indices_with_redundancy( total_physical_experts: int, redundancy_config: list[int], # redundancy for each logical expert ) -> torch.Tensor: - """ - Create expert indices with redundancy. + """Create expert indices with redundancy. Args: num_layers: number of layers @@ -67,6 +66,7 @@ def create_expert_indices_with_redundancy( Returns: indices: Shape (num_layers, total_physical_experts) + """ assert sum(redundancy_config) == total_physical_experts assert len(redundancy_config) == num_logical_experts @@ -95,8 +95,7 @@ def create_expert_weights( device: torch.device, physical_to_logical_mapping: torch.Tensor, ) -> list[list[torch.Tensor]]: - """ - Create fake expert weights tensor for testing. + """Create fake expert weights tensor for testing. Use `arange` to generate predictable weights values, based on logical expert ID. @@ -105,6 +104,7 @@ def create_expert_weights( Args: physical_to_logical_mapping: Shape (num_layers, num_local_experts) mapping[layer, physical_pos] = logical_expert_id + """ expert_weights = [] @@ -209,9 +209,7 @@ def verify_redundant_experts_have_same_weights( world_size: int, num_local_experts: int, ) -> bool: - """ - Verify that all replicas of the same logical expert have the same weights. - """ + """Verify that all replicas of the same logical expert have the same weights.""" num_layers = len(expert_weights) total_physical_experts = world_size * num_local_experts @@ -568,7 +566,6 @@ def test_rearrange_expert_weights_with_redundancy( eplb_communicator, ): """Test the functionality of rearranging expert weights with redundancy.""" - if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: @@ -678,7 +675,6 @@ def test_async_transfer_layer_without_mtp( eplb_communicator: str, ): """Exercise async EPLB transfer path without MTP/spec decode.""" - if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: @@ -696,11 +692,9 @@ def test_async_transfer_layer_without_mtp( @pytest.mark.parametrize("world_size", [2, 4]) def test_rearrange_expert_weights_no_change(world_size): - """ - Test that when the indices do not change, the weights should remain + """Test that when the indices do not change, the weights should remain unchanged. """ - if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -800,7 +794,6 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: @pytest.mark.parametrize("world_size", [2, 4]) def test_rearrange_expert_weights_profile_mode(world_size): """Test profile mode (should not copy actual weights)""" - if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( diff --git a/tests/distributed/test_eplb_quant_scale_consistency.py b/tests/distributed/test_eplb_quant_scale_consistency.py index 3b50b315f0a1..49dc8afbf311 100644 --- a/tests/distributed/test_eplb_quant_scale_consistency.py +++ b/tests/distributed/test_eplb_quant_scale_consistency.py @@ -124,7 +124,8 @@ def _build_processed_layer( backend: str, raw: dict[str, torch.Tensor], device: torch.device ) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]: """Create the method + layer, load ``raw`` and run the real - process_weights_after_loading (which builds the flashinfer kernel).""" + process_weights_after_loading (which builds the flashinfer kernel). + """ moe_config = _make_moe_config(backend) method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False) layer = _RoutedExpertsStub(moe_config).to(device) @@ -157,7 +158,8 @@ def _assert_tensors_equal( def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None: """Permute experts of every registered Parameter in place, the way - EPLB's in-place rearrangement moves expert slices along dim 0.""" + EPLB's in-place rearrangement moves expert slices along dim 0. + """ with torch.no_grad(): for _, param in layer.named_parameters(): param.copy_(param[perm]) @@ -165,7 +167,8 @@ def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> def _ensure_world1_distributed() -> None: """The enable_eplb config makes weight processing all-reduce activation - scale amaxes over the EP group, so a (single-rank) EP group must exist.""" + scale amaxes over the EP group, so a (single-rank) EP group must exist. + """ from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, init_distributed_environment, @@ -253,7 +256,8 @@ def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None: def test_fp8_per_tensor_alphas_registered_and_aliased() -> None: """CPU-only: the fp8 per-tensor oracle must register the fused (w_scale * a_scale) products as layer Parameters aliased by the quant - config, so EPLB rearrangement keeps kernels consistent.""" + config, so EPLB rearrangement keeps kernels consistent. + """ from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, make_fp8_moe_quant_config, diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py index b4211c74c171..ced206a3056a 100644 --- a/tests/distributed/test_eplb_spec_decode.py +++ b/tests/distributed/test_eplb_spec_decode.py @@ -75,8 +75,7 @@ def test_eplb_spec_decode( monkeypatch: pytest.MonkeyPatch, model_setup: tuple[str, str, str, int, float], ): - """ - Test the correctness of EPLB speculative decoding with GSM8K dataset. + """Test the correctness of EPLB speculative decoding with GSM8K dataset. Applicable to MoE models with mtp or eagle spec decode. """ method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup @@ -109,10 +108,7 @@ def test_eplb_spec_decode( @large_gpu_mark(min_gb=80) def test_eplb_spec_decode_qwen3_next_mtp_async() -> None: - """ - Ensure async EPLB works with MTP speculative decoding for Qwen3-Next. - """ - + """Ensure async EPLB works with MTP speculative decoding for Qwen3-Next.""" TASK = "gsm8k" FILTER = "exact_match,strict-match" RTOL = 0.03 diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py index a6cfb0d29d8e..dd2bb2760b1e 100644 --- a/tests/distributed/test_eplb_utils.py +++ b/tests/distributed/test_eplb_utils.py @@ -30,8 +30,7 @@ def _make_model_state( def test_commit_eplb_maps_shape_change(): - """ - When the number of physical experts changes, resize the active map within its + """When the number of physical experts changes, resize the active map within its preallocated storage. """ num_layers, num_logical, num_physical = 2, 4, 6 @@ -65,8 +64,7 @@ def test_commit_eplb_maps_shape_change(): def test_commit_eplb_maps_for_layer_logical_padding(): - """ - Test that logical_to_physical_map is padded with -1 to fill the + """Test that logical_to_physical_map is padded with -1 to fill the pre-allocated slots when the new map has fewer replicas than the max. """ num_layers, num_logical, num_physical = 2, 4, 6 diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index 894f3f7f8b83..aae80d9cce31 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -20,27 +20,26 @@ class EventSample( tag=True, # type: ignore array_like=True, # type: ignore ): - """Test event for publisher testing""" + """Test event for publisher testing.""" id: int value: str class SampleBatch(EventBatch): - """Test event batch for publisher testing""" + """Test event batch for publisher testing.""" events: list[EventSample] def create_test_events(count: int) -> SampleBatch: - """Create a batch of test events""" + """Create a batch of test events.""" events = [EventSample(id=i, value=f"test-{i}") for i in range(count)] return SampleBatch(ts=time.time(), events=events) def test_basic_publishing(publisher, subscriber): - """Test basic event publishing works""" - + """Test basic event publishing works.""" test_batch = create_test_events(5) publisher.publish(test_batch) @@ -58,7 +57,7 @@ def test_basic_publishing(publisher, subscriber): def test_multiple_events(publisher, subscriber): - """Test publishing and receiving multiple event batches""" + """Test publishing and receiving multiple event batches.""" for _ in range(10): batch = create_test_events(2) publisher.publish(batch) @@ -75,7 +74,7 @@ def test_multiple_events(publisher, subscriber): def test_replay_mechanism(publisher, subscriber): - """Test the replay mechanism works correctly""" + """Test the replay mechanism works correctly.""" for _ in range(19): batch = create_test_events(1) publisher.publish(batch) @@ -96,7 +95,7 @@ def test_replay_mechanism(publisher, subscriber): def test_replay_includes_topic(publisher, subscriber, publisher_config): - """Test that replay responses include the topic, matching PUB format""" + """Test that replay responses include the topic, matching PUB format.""" for _ in range(5): publisher.publish(create_test_events(1)) @@ -115,7 +114,7 @@ def test_replay_includes_topic(publisher, subscriber, publisher_config): def test_buffer_limit(publisher, subscriber, publisher_config): - """Test buffer limit behavior""" + """Test buffer limit behavior.""" buffer_size = publisher_config.buffer_steps # Publish more events than the buffer can hold @@ -139,9 +138,7 @@ def test_buffer_limit(publisher, subscriber, publisher_config): def test_topic_filtering(publisher_config): - """ - Test that a subscriber only receives messages matching its topic filter - """ + """Test that a subscriber only receives messages matching its topic filter.""" publisher_config.replay_endpoint = None publisher_config.topic = "foo" @@ -174,7 +171,7 @@ def test_topic_filtering(publisher_config): def test_high_volume(publisher, subscriber): - """Test publishing and receiving a high volume of events""" + """Test publishing and receiving a high volume of events.""" num_batches = 10_000 events_per_batch = 100 @@ -210,7 +207,7 @@ def publish_events(): def test_null_publisher(): - """Test that NullEventPublisher can be used without errors""" + """Test that NullEventPublisher can be used without errors.""" publisher = NullEventPublisher(DP_RANK) # This should not raise any errors @@ -220,8 +217,7 @@ def test_null_publisher(): def test_data_parallel_rank_tagging(publisher_config): - """Test that events are properly tagged with their data parallel rank""" - + """Test that events are properly tagged with their data parallel rank.""" publisher_config.topic = "foo" pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK) pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1) @@ -285,7 +281,7 @@ def test_data_parallel_rank_tagging(publisher_config): def test_event_publisher_factory(random_port): - """Test event publisher factory creation behavior under different configurations""" + """Test event publisher factory creation behavior under different configurations.""" from vllm.config.kv_events import KVEventsConfig from vllm.distributed.kv_events import ZmqEventPublisher diff --git a/tests/distributed/test_expert_placement.py b/tests/distributed/test_expert_placement.py index 46f63408f467..4174b5c6ccff 100644 --- a/tests/distributed/test_expert_placement.py +++ b/tests/distributed/test_expert_placement.py @@ -49,7 +49,6 @@ def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts) @pytest.mark.parametrize("world_size", [2, 4]) def test_expert_placement_various_sizes(expert_placement_strategy, world_size): """Test round_robin expert placement with various expert counts.""" - # Test with different global_num_experts values # Include both divisible and non-divisible cases if world_size == 2: @@ -116,7 +115,6 @@ def test_expert_placement_various_sizes(expert_placement_strategy, world_size): @pytest.mark.parametrize("world_size", [2, 4]) def test_expert_placement_edge_cases(expert_placement_strategy, world_size): """Test edge cases for round_robin expert placement.""" - # Test case 1: ep_size = 1 (should return None for expert_map) local_num_experts, expert_map, _ = determine_expert_map( ep_size=1, @@ -139,7 +137,6 @@ def test_expert_placement_edge_cases(expert_placement_strategy, world_size): def test_determine_expert_map_comprehensive(): """Test of determine_expert_map function with various configurations.""" - # Test cases: (ep_size, ep_rank, global_num_experts, # expert_placement_strategy, expected_local, expected_map_pattern) test_cases = [ diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 636f7cb525f5..0258c781c721 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for MNNVL AllToAll operations. +"""Tests for MNNVL AllToAll operations. Requires: docker run ... --cap-add=SYS_PTRACE ... Run: pytest tests/distributed/test_mnnvl_alltoall.py -v @@ -87,6 +86,7 @@ def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue): Otherwise use tp=world_size (default for EP-based tests). dp_port: Separate port for the DP master (only used when dp_size is set). err_queue: Queue for propagating tracebacks to the parent process. + """ try: os.environ.pop("CUDA_VISIBLE_DEVICES", None) @@ -115,6 +115,7 @@ def _init_dp_environment(world_size, rank, port, dp_size, dp_port): Args: port: Port for torch.distributed init. dp_port: Separate port for the DP master group init. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.parallel import ParallelConfig @@ -159,7 +160,8 @@ def _make_forward_context(rank, world_size, num_tokens_per_rank): class _AttnMeta: """Minimal placeholder so set_forward_context's ``attn_metadata is not None`` guard (forward_context.py:334) - is satisfied. The real DPMetadata is built from num_tokens_across_dp.""" + is satisfied. The real DPMetadata is built from num_tokens_across_dp. + """ dp_metadata = None @@ -458,7 +460,8 @@ def _one_sided_workspace_grow_worker(rank, world_size): def test_one_sided_manager_workspace_grow(world_size): """A later initialize() with a larger per-token payload must grow the workspace and rebuild MoeAlltoAll; a later initialize() with a smaller - payload must no-op.""" + payload must no-op. + """ _spawn_workers( _one_sided_workspace_grow_worker, world_size, diff --git a/tests/distributed/test_mq_connect_ip.py b/tests/distributed/test_mq_connect_ip.py index 4b0cdda3ad9e..701dc8f248c7 100644 --- a/tests/distributed/test_mq_connect_ip.py +++ b/tests/distributed/test_mq_connect_ip.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that MessageQueue uses the local node's IP for binding, +"""Test that MessageQueue uses the local node's IP for binding, not a remote master_addr. This validates the fix for cross-node data-parallel where each DP group leader must bind to its own IP. @@ -22,7 +21,8 @@ def test_mq_bind_with_local_ip(): """MessageQueue with remote readers should successfully bind - when connect_ip is the local node's IP.""" + when connect_ip is the local node's IP. + """ # n_reader=2, n_local_reader=1 means 1 remote reader, # which triggers the remote ZMQ socket bind. mq = MessageQueue( @@ -44,7 +44,8 @@ def test_mq_bind_with_local_ip(): def test_mq_bind_with_non_local_ip_fails(): """MessageQueue should fail to bind when connect_ip is a non-local IP address (simulating the bug where master_addr - from a different node was used).""" + from a different node was used). + """ # Use a non-local IP that we definitely can't bind to. # 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned. non_local_ip = "198.51.100.1" @@ -58,7 +59,8 @@ def test_mq_bind_with_non_local_ip_fails(): def test_mq_bind_defaults_to_local_ip(): """When connect_ip is None, MessageQueue should auto-detect - the local IP and bind successfully.""" + the local IP and bind successfully. + """ mq = MessageQueue( n_reader=2, n_local_reader=1, diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 599025d1c8ec..d0a417a7886d 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for MultiprocExecutor at the executor level. +"""Integration tests for MultiprocExecutor at the executor level. This test directly tests the executor without going through the LLM interface, focusing on executor initialization, RPC calls, and distributed execution. """ @@ -377,8 +376,7 @@ def test_multiproc_executor_properties(): @multi_gpu_test(num_gpus=4) def test_multiproc_executor_multi_node(): - """ - Test MultiprocExecutor with multi-node configuration. + """Test MultiprocExecutor with multi-node configuration. This simulates 2 nodes with TP=4: - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2 - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index a464a61683b5..8ba93df4f504 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -WARNING: This test runs in both single-node (4 GPUs) and multi-node - (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is - important to set the distributed backend to "mp" to avoid Ray scheduling - all workers in a node other than the head node, which can cause the test - to fail. +"""WARNING: This test runs in both single-node (4 GPUs) and multi-node +(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is +important to set the distributed backend to "mp" to avoid Ray scheduling +all workers in a node other than the head node, which can cause the test +to fail. """ import json diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 91dccd579869..0eda9a14e05d 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -495,7 +495,8 @@ def test_pynccl_suspend_resume_idempotent(): def test_pynccl_suspend_noop_without_symbol(): """RCCL / NCCL < 2.29.7 lack ncclCommSuspend; suspend() must no-op, not - crash, and resume() must then no-op too since nothing was suspended.""" + crash, and resume() must then no-op too since nothing was suspended. + """ from unittest.mock import Mock comm = object.__new__(PyNcclCommunicator) diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index fb2b04abf0d5..f9228fbd6832 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -396,8 +396,7 @@ def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch): def qr_variable_input(rank, world_size): - """ - When the tensor parallelism is set to 4 or 8, frequent changes + """When the tensor parallelism is set to 4 or 8, frequent changes in the input shape can cause QuickReduce to hang (this issue has been observed with the gpt_oss model). """ diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py index bd4a253b96a4..f9b9e57fc54c 100644 --- a/tests/distributed/test_ray_v2_executor.py +++ b/tests/distributed/test_ray_v2_executor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for RayExecutorV2 at the executor level. +"""Integration tests for RayExecutorV2 at the executor level. Validates executor initialization, placement group support, RPC calls, and distributed execution with various TP/PP configurations. """ @@ -299,7 +298,6 @@ def test_ray_v2_run_refs_stored_for_monitoring(executor): @pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) def test_ray_v2_single_node_generation(tp_size, pp_size): """End-to-end LLM generation with RayExecutorV2.""" - llm = LLM( model=MODEL, tensor_parallel_size=tp_size, diff --git a/tests/distributed/test_ray_v2_executor_e2e.py b/tests/distributed/test_ray_v2_executor_e2e.py index fb5830132698..27de410e3892 100644 --- a/tests/distributed/test_ray_v2_executor_e2e.py +++ b/tests/distributed/test_ray_v2_executor_e2e.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Orchestration-level integration tests for RayExecutorV2. -""" +"""Orchestration-level integration tests for RayExecutorV2.""" import gc import os @@ -25,7 +23,8 @@ def _ray_init(): Without this, workers cannot unpickle actor classes defined in the ``tests`` package, causing FunctionActorManager to fall back to - TemporaryActor which drops async method signatures.""" + TemporaryActor which drops async method signatures. + """ project_root = str(pathlib.Path(__file__).resolve().parents[2]) ray.init( ignore_reinit_error=True, @@ -158,8 +157,7 @@ def test_multi_replicas_with_bundle_indices(ray_init): def test_env_var_and_runtime_env_propagation(): - """ - Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env + """Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env propagate to RayWorkerProc actors. """ sentinel_vars = { diff --git a/tests/distributed/test_sharded_rdt_plan.py b/tests/distributed/test_sharded_rdt_plan.py index 7a75f0b7fcd6..5cf0a82203b1 100644 --- a/tests/distributed/test_sharded_rdt_plan.py +++ b/tests/distributed/test_sharded_rdt_plan.py @@ -53,7 +53,8 @@ class _FakeLayer: """Stand-in for a leaf module. The planner only stores it and keys on - ``id()``; nothing in the plan path touches its attributes.""" + ``id()``; nothing in the plan path touches its attributes. + """ def __init__(self, tag: str) -> None: self.tag = tag @@ -64,7 +65,8 @@ def __repr__(self) -> str: def _module(layer, copies): """What the bake records for one leaf module: a scatter list whose entries - all carry that module as their ``layer``.""" + all carry that module as their ``layer``. + """ for c in copies: c.layer = layer return copies @@ -81,7 +83,8 @@ def _copy( dtype=torch.bfloat16, ): """A `_Scatter` as the bake would record it: source key, owning layer, and - the meta destination region.""" + the meta destination region. + """ stride = ( (1,) if len(shape) == 1 @@ -140,7 +143,8 @@ def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None): def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): """`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model - where every gather group is one fully-baked leaf module.""" + where every gather group is one fully-baked leaf module. + """ names = ["embed.weight"] names += [f"model.layers.{i}.w" for i in range(n_layers)] names += ["norm.weight"] @@ -161,7 +165,8 @@ def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): class TestFakeOpChains: """Every allowlisted op must append itself to the chain and hand back a child whose shape/dtype PyTorch itself computed. The chain is the wire - format the producer replays, so its exact contents are load-bearing.""" + format the producer replays, so its exact contents are load-bearing. + """ def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None): return FakeRDTTensor( @@ -222,7 +227,8 @@ def test_kwargs_are_frozen_sorted_for_hashability(self): @pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)]) def test_multi_return_op_emits_one_child_per_output(self, op, n): """chunk/unbind hand back a tuple; each child carries the base op plus a - trailing __getitem__(i) so the producer can index the replayed result.""" + trailing __getitem__(i) so the producer can index the replayed result. + """ parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0) assert isinstance(parts, tuple) assert len(parts) == n @@ -238,7 +244,8 @@ def test_op_chain_is_hashable_as_a_fetch_key(self): class TestFakeUnsupportedOps: """Anything that needs real data must fail loudly at bake time rather than - silently transferring the wrong bytes.""" + silently transferring the wrong bytes. + """ def _fake(self): return FakeRDTTensor( @@ -278,7 +285,8 @@ def test_unsupported_is_a_notimplementederror(self): class TestBakeRecording: """During the dry run the fake's ``copy_`` is the data sink: it records the - source chain plus the meta destination's strided region and moves nothing.""" + source chain plus the meta destination's strided region and moves nothing. + """ def _recorder_and_fake(self, shape=(4, 6)): rec = BakeSink() @@ -314,7 +322,8 @@ def test_copy_marks_the_source_name_live(self): def test_unattributed_copy_is_live_but_unrecorded(self): """A copy_ with no loader stamp cannot be attributed to a param, so its - module must fall back to the plain load — but the name still moved data.""" + module must fall back to the plain load — but the name still moved data. + """ rec, fake = self._recorder_and_fake() rec.current = None torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake) @@ -354,7 +363,8 @@ def test_the_recorded_dtype_is_the_chains_output_not_the_sources(self): producer packs what the replay yields, so the record must carry the POST-chain dtype: taking it from the source name's metadata instead sizes the slice with the wrong itemsize and shifts every later slice in the - chunk, carving the packed blob differently on the two sides.""" + chunk, carving the packed blob differently on the two sides. + """ rec, fake = self._recorder_and_fake(shape=(4,)) layer = _FakeLayer("reinterpreted") rec.current = (layer, "weight") @@ -374,7 +384,8 @@ def test_a_broadcasting_copy_is_refused_at_bake(self): the producer packs what the chain yields, and nothing downstream can catch a mismatch — the consumer's buffer view is exactly prod(dest.shape) elements, so it reshapes cleanly over bytes laid out at other offsets. - So the bake refuses rather than recording a slice it cannot carve.""" + So the bake refuses rather than recording a slice it cannot carve. + """ rec, fake = self._recorder_and_fake(shape=(1, 6)) rec.current = (_FakeLayer("broadcast"), "weight") param = torch.empty((4, 6), dtype=torch.bfloat16, device=META) @@ -390,7 +401,8 @@ def test_a_broadcasting_copy_is_refused_at_bake(self): class TestConsumerIdentity: """Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the worker's producer block and keys the producer's per-consumer serve ring, so - a collision silently serves two workers out of one ring.""" + a collision silently serves two workers out of one ring. + """ def _engine(self, *, dp_index, rank, world_size): eng = object.__new__(ShardedRDTWeightTransferEngine) @@ -427,7 +439,8 @@ def test_one_engine_indexes_by_dp_and_tp(self): def test_independent_engines_offset_into_distinct_ranges(self): """A fleet of separate engines restarts _global_worker_index at 0 in each, so without the replica offset every engine would claim 0..w-1. - The driver sets replica_rank on the payload; these are the ids it buys.""" + The driver sets replica_rank on the payload; these are the ids it buys. + """ ids = [ self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2) for r in (0, 1) @@ -453,7 +466,8 @@ def _router( group_lens=None, ): """A router over ``n`` single-name groups unless told otherwise, so a test - can state expectations in group indices.""" + can state expectations in group indices. + """ if names is None: n = len(name_owner_class or []) or (len(group_lens or []) or 1) names = [f"g{i}" for i in range(n)] @@ -482,7 +496,8 @@ def test_identity_when_fleets_match(self): def test_every_producer_nic_carries_traffic(self): """16 producers / 8 consumers: the block rule spreads each consumer's pulls over its block, alternating by group, so no producer sits idle. - Plain `consumer_id % len(owners)` would idle half of them.""" + Plain `consumer_id % len(owners)` would idle half of them. + """ r = _router(16, 8, group_lens=[1] * 95) names = [f"g{i}" for i in range(95)] for c in range(8): @@ -500,7 +515,8 @@ def test_fan_in_shares_one_producer(self): def test_pipeline_stages_route_to_the_owning_stage(self): """2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner - sets, one per stage; each consumer reaches both.""" + sets, one per stage; each consumer reaches both. + """ r = _router( 16, 8, @@ -516,7 +532,8 @@ def test_pipeline_stages_route_to_the_owning_stage(self): def test_expert_names_route_to_the_holding_ranks(self): """A name held by only some ranks routes to those; a replicated name - keeps the full set.""" + keeps the full set. + """ r = RdtRouter( 8, 4, @@ -535,7 +552,8 @@ def test_expert_names_route_to_the_holding_ranks(self): def test_consumers_spread_over_an_owner_set(self): """Several ranks hold a name (its TP peers); the block rule must spread - consumers across them, not funnel through one.""" + consumers across them, not funnel through one. + """ r = _router( 4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6 ) @@ -544,7 +562,8 @@ def test_consumers_spread_over_an_owner_set(self): def test_a_group_frees_at_every_owner_of_any_of_its_names(self): """The free barrier is per group, so its fan-out is the UNION over the - group's names: a rank holding none of them is not signalled.""" + group's names: a rank holding none of them is not signalled. + """ r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1]) assert r.group_owners(0) == [0, 1, 2, 3] assert r.group_owners(1) == [0, 1] @@ -614,7 +633,8 @@ def test_every_replica_of_a_worker_resolves_the_same_producer(self): def test_adding_a_deployment_does_not_move_the_first_one(self): """The strongest form: an existing consumer's routes are byte-identical - before and after another deployment is provisioned.""" + before and after another deployment is provisioned. + """ one, names = self._fleet(16, 8, 1) two, _ = self._fleet(16, 8, 2) for consumer_id in range(8): @@ -624,7 +644,8 @@ def test_adding_a_deployment_does_not_move_the_first_one(self): def test_one_deployment_is_the_historical_carve(self): """Default (no width) and width == the whole fleet must agree with the - plain block rule, so a single-deployment fleet is untouched.""" + plain block rule, so a single-deployment fleet is untouched. + """ plain = _router(16, 8, group_lens=[1] * 6) explicit, names = self._fleet(16, 8, 1) for consumer_id in range(8): @@ -638,7 +659,8 @@ def test_one_deployment_is_the_historical_carve(self): def test_the_overlay_holds_for_partially_owned_names(self): """Pipeline stages and experts too: two owner sets, one narrow. The narrow one already forced agreement (one owner, no block to carve); the - wide one is what the carve decides.""" + wide one is what the carve decides. + """ r, names = self._fleet( 8, 4, @@ -654,14 +676,16 @@ def test_the_overlay_holds_for_partially_owned_names(self): def test_a_width_that_does_not_divide_the_fleet_raises(self): """A uniform fleet is asserted upstream in ``get_world_size``; if it ever is not, the overlay would map two workers of one deployment onto the same - block index, so refuse rather than serve the wrong bytes.""" + block index, so refuse rather than serve the wrong bytes. + """ with pytest.raises(ValueError, match="does not divide"): RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4) def test_pulls_still_carry_the_global_consumer_id(self): """The carve uses the index WITHIN a deployment, but the wire must keep the fleet-global id: it is what lets the producer tell the sharers of a - slot apart, and count their arrivals separately.""" + slot apart, and count their arrivals separately. + """ class _Method: def __init__(self): @@ -704,7 +728,8 @@ def _producer_pack_offsets(slices): class TestPackedLayout: def test_consumer_layout_matches_the_producer_rule(self): """Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style - group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment.""" + group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment. + """ specs = [ ("w.fp8", "float8_e4m3fn", [17]), ("w.scale", "float32", [3]), @@ -779,7 +804,8 @@ def test_duplicate_source_keys_are_packed_once(self): class TestChunkModuleScatters: """The chunk cut: one chunk per distinct owner class present in the copies, ascending by class index. Derived purely from the bake plus the ownership - table, so vLLM's expert placement needs no cases.""" + table, so vLLM's expert placement needs no cases. + """ def test_unstamped_copies_form_one_chunk(self): layer = _FakeLayer("l") @@ -793,7 +819,8 @@ def test_one_chunk_per_owner_class_ascending(self): """Classes are numbered by first appearance in metadata order, so the replicated class (0) leads and the rest follow — the ordering the old "-1 first, then ascending stamp" rule produced, now falling out of the - table instead of a special case.""" + table instead of a special case. + """ layer = _FakeLayer("moe") names = ["norm", "e5", "e0", "e2"] copies = [_copy(n, shape=(4,)) for n in names] @@ -813,7 +840,8 @@ def test_one_chunk_per_owner_class_ascending(self): def test_copy_order_within_a_chunk_is_bake_order(self): """Scatter order within a chunk must follow the bake, not the stamp - sort: the packed layout and the producer's replay agree on keys order.""" + sort: the packed layout and the producer's replay agree on keys order. + """ layer = _FakeLayer("moe") names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names copies = [_copy(n, shape=(4,)) for n in names] @@ -827,7 +855,8 @@ def test_copy_order_within_a_chunk_is_bake_order(self): def test_stamps_interleave_across_modules_without_reordering_within(self): """Two modules' copies bucket by stamp; within each bucket the order is - module-then-bake order.""" + module-then-bake order. + """ a, b = _FakeLayer("a"), _FakeLayer("b") mods = [ _module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]), @@ -845,9 +874,10 @@ def test_stamps_interleave_across_modules_without_reordering_within(self): ] def test_scatters_carry_their_own_dtype(self): - """dtype rides the record from the bake, where it is the fake's dtype + """Dtype rides the record from the bake, where it is the fake's dtype AFTER its op chain — not a plan-time lookup of the source name, which - would be wrong for any chain that reinterprets dtype.""" + would be wrong for any chain that reinterprets dtype. + """ layer = _FakeLayer("l") eng = _planner({}, name_meta={"w": ("float32", [10])}) ((_er, scatters),) = eng._chunk_module_scatters( @@ -886,7 +916,8 @@ def test_stamps_multiply_chunks_within_a_group(self): def test_materialize_fires_on_a_modules_first_chunk_only(self): """Empty HF params are allocated once per module, by construction — - including a FusedMoE-like module whose copies span owner-class chunks.""" + including a FusedMoE-like module whose copies span owner-class chunks. + """ layer = _FakeLayer("spanning") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -921,7 +952,8 @@ def test_free_signal_fires_on_each_groups_last_chunk(self): def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self): """The load-bearing case: with the group cut across owner-class chunks the signal must hang off the LAST one. Signaling earlier lets the producers - drop the gather buffers while a later chunk's RDMA is still reading.""" + drop the gather buffers while a later chunk's RDMA is still reading. + """ layer = _FakeLayer("l") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -937,7 +969,8 @@ def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self): def test_free_signal_timing_across_two_split_groups(self): """Two groups x 2 owner-class chunks each: signals land on each group's - last chunk only.""" + last chunk only. + """ names_a = [f"model.layers.0.w{i}" for i in range(4)] names_b = [f"model.layers.1.w{i}" for i in range(4)] baked, name_meta, stamps = {}, {}, {} @@ -955,7 +988,8 @@ def test_free_signal_timing_across_two_split_groups(self): def test_a_group_with_nothing_local_is_signaled_at_sync_start(self): """The owners still published it, so this consumer must signal it — immediately, before the pipeline (signal-before-publish is tolerated), - never hung off another group's chunk.""" + never hung off another group's chunk. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["model.layers.1.w"] # nothing baked, and not live => no pull eng = _planner( @@ -976,7 +1010,8 @@ def test_leading_groups_with_nothing_local_go_to_pre_free(self): def test_modules_are_deduped_within_a_group(self): """Several names of one fused module map to the same scatter list; the - plan must not chunk it twice.""" + plan must not chunk it twice. + """ layer = _FakeLayer("fused") names = ["qkv.q", "qkv.k", "qkv.v"] copies = [_copy(n, shape=(4,)) for n in names] @@ -992,7 +1027,8 @@ def test_modules_are_deduped_within_a_group(self): def test_plan_is_a_pure_function_of_its_inputs(self): """Called twice with the same arguments it must produce the same plan — - that is what lets the engine build it once and reuse it every sync.""" + that is what lets the engine build it once and reuse it every sync. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) eng = _planner(baked, name_meta=name_meta) a = eng._build_call_plan(names, group_lens) @@ -1012,7 +1048,8 @@ def test_single_producer_routes_every_chunk_to_local_owner_zero(self): class TestUnbakedNamesGuard: """A live name with no baked plan has no way to load — there is no fallback, and its pull would target groups the pipeline already freed — so the plan - build fails at init, naming the names.""" + build fails at init, naming the names. + """ def test_a_live_unbaked_name_fails_the_plan_build(self): baked, name_meta, names, group_lens = _one_module_per_layer(2) @@ -1023,7 +1060,8 @@ def test_a_live_unbaked_name_fails_the_plan_build(self): def test_never_copied_names_are_dropped_entirely(self): """Experts owned by another EP rank no-op in their loader: not an error, - and not a chunk — their group is pre-freed if nothing else fills it.""" + and not a chunk — their group is pre-freed if nothing else fills it. + """ baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["norm.weight"] eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"}) @@ -1035,13 +1073,15 @@ def test_never_copied_names_are_dropped_entirely(self): class TestCallPlanRouting: """Under partial ownership each chunk must be pulled from a producer that actually holds every name in it. Consumers bind EVERY producer (signals fan - out to all owners), so local index == trainer rank.""" + out to all owners), so local index == trainer rank. + """ def _routed_planner( self, held_by_group, num_consumers, consumer_id, *, held_by=None ): """``held_by_group[g]`` = the ranks holding group g's names, so a test - can still speak in groups; ``held_by`` overrides individual names.""" + can still speak in groups; ``held_by`` overrides individual names. + """ baked, name_meta, names, group_lens = _one_module_per_layer( len(held_by_group) - 2 ) @@ -1068,14 +1108,16 @@ def _routed_planner( def test_two_stage_ownership_splits_chunks_between_owners(self): """Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the - second. Every chunk goes to a holder of its names.""" + second. Every chunk goes to a holder of its names. + """ eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0) plan = eng._build_call_plan(names, group_lens) assert [c.owner for c in plan.chunks] == [0, 0, 1, 1] def test_a_multi_owner_group_still_resolves_to_one_producer(self): """Every chunk is served by exactly ONE producer: splitting a pull only - multiplies produce calls, since the consumer's NIC bounds it.""" + multiplies produce calls, since the consumer's NIC bounds it. + """ eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0) plan = eng._build_call_plan(names, group_lens) assert all(0 <= c.owner < 2 for c in plan.chunks) @@ -1083,7 +1125,8 @@ def test_a_multi_owner_group_still_resolves_to_one_producer(self): def test_an_expert_name_routes_to_the_rank_holding_it(self): """PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert name inside a stage-1 group is held by rank 3 alone. The name's own owner - set is what routes it.""" + set is what routes it. + """ eng, names, group_lens = self._routed_planner( [[0, 1], [0, 1], [2, 3], [2, 3]], 1, @@ -1098,7 +1141,8 @@ def test_an_expert_name_routes_to_the_rank_holding_it(self): def test_an_unheld_name_is_rejected_at_init(self): """A name no rank holds can never be served. The trainer raises when it transposes the holdings; the router's validate() is the consumer-side - backstop, and either way it fails at init rather than at first pull.""" + backstop, and either way it fails at init rather than at first pull. + """ with pytest.raises(ValueError, match="empty"): RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate() @@ -1107,14 +1151,16 @@ class TestSignalCompleteness: """Every gather group is signaled exactly once — after its last chunk when the worker pulls from it, at sync start otherwise — for ANY expert placement. A group signaled twice over-credits the barrier; a group never - signaled parks a producer credit and hangs end_sync.""" + signaled parks a producer credit and hangs end_sync. + """ def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2): - """pre / n_layers MoE layers / post. Each layer: one norm (its own + """Pre / n_layers MoE layers / post. Each layer: one norm (its own module, stamp -1) + n_experts expert names stamped ``e // n_local``. The worker's bake covers only ``worker_experts`` (its placement); foreign expert names never copied => dropped entirely, as in the real - bake.""" + bake. + """ n_local = n_experts // ep_size names = ["embed.weight"] group_lens = [1] @@ -1173,8 +1219,9 @@ def test_round_robin_placement_signals_every_group_exactly_once(self): assert self._signals(plan) == list(range(len(group_lens))) def test_placement_changes_only_the_chunk_count_never_the_signals(self): - """linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6 - hits all 4. More chunks per group, identical signal set.""" + """Linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6 + hits all 4. More chunks per group, identical signal set. + """ plans = {} for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])): eng, names, group_lens = self._moe_planner(worker_experts=experts) @@ -1188,7 +1235,8 @@ def test_placement_changes_only_the_chunk_count_never_the_signals(self): def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self): """The FusedMoE shape: one module's copies span every owner-class chunk of - its group; materialize on its first, quant on its last.""" + its group; materialize on its first, quant on its last. + """ eng, names, group_lens = self._moe_planner( worker_experts=[0, 2, 4, 6], n_layers=1 ) @@ -1246,7 +1294,8 @@ def test_partition_is_one_group_per_layer_between_unindexed_blocks(self): def test_flattening_a_partition_returns_the_input_order(self): """Un-indexed names group by POSITION, not by name class — so one that looks like a pre-block name still lands after the layers it follows. - Group index therefore means the same thing on every rank.""" + Group index therefore means the same thing on every rank. + """ names = [ "embed.weight", "model.layers.0.a", @@ -1264,7 +1313,8 @@ def test_layers_are_ordered_by_index_not_appearance(self): def test_stacks_keep_prefix_appearance_order_and_sort_within(self): """Two stacks stay in the order they appear — the vision tower before the - text stack — while each sorts by index internally.""" + text stack — while each sorts by index internally. + """ names = [ "visual.blocks.1.w", "visual.blocks.0.w", @@ -1281,7 +1331,8 @@ def test_stacks_keep_prefix_appearance_order_and_sort_within(self): def test_the_post_block_lands_last_however_early_it_arrives(self): """Megatron-Bridge streams the last pipeline stage's output block before its layers. Sweeping un-indexed names after the first layer into a - trailing group is what lets the gather loop walk groups ascending.""" + trailing group is what lets the gather loop walk groups ascending. + """ names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"] assert layerwise_groups(names) == [ ["model.layers.0.w"], @@ -1291,7 +1342,8 @@ def test_the_post_block_lands_last_however_early_it_arrives(self): def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self): """A raw-checkpoint source may yield in shard-packing order. Coalescing on - the key keeps that from shattering one layer into many groups.""" + the key keeps that from shattering one layer into many groups. + """ names = [ "model.layers.0.a", "model.layers.1.a", @@ -1333,13 +1385,15 @@ def test_a_model_with_no_layers_is_one_pre_group(self): def test_every_supported_naming_convention_partitions_per_layer(self, prefix): """One decoder layer per group under any prefix, with no per-architecture table. A literal `model.layers.` match silently yields one whole-model - group for most of these.""" + group for most of these. + """ names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)] assert layerwise_groups(names) == [[n] for n in names] def test_a_moe_layer_stays_whole(self): """The OUTERMOST index wins, so per-expert names group by layer. Splitting - on the expert would cut FusedMoE — one leaf module — across groups.""" + on the expert would cut FusedMoE — one leaf module — across groups. + """ names = [ "model.layers.0.mlp.experts.0.w1.weight", "model.layers.0.mlp.experts.1.w1.weight", @@ -1356,7 +1410,8 @@ def test_a_moe_layer_stays_whole(self): def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self): """Qwen-VL layout. Today's literal `model.layers.` match puts every one of these in ONE group — the whole model — because the text stack ships as - `model.language_model.layers.`.""" + `model.language_model.layers.`. + """ names = [ "visual.patch_embed.proj.weight", "visual.blocks.0.attn.qkv.weight", @@ -1375,7 +1430,8 @@ def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self): def test_two_stacks_sharing_an_index_do_not_merge(self): """The key carries the prefix, so vision block 0 and text layer 0 are - different groups even though both are index 0.""" + different groups even though both are index 0. + """ names = [ "visual.blocks.0.attn.qkv.weight", "model.language_model.layers.0.self_attn.q_proj.weight", @@ -1412,20 +1468,23 @@ def test_the_two_allowlists_describe_the_same_contract(self): def test_every_bake_recordable_op_is_replayable(self): """Anything the bake can record must be serveable, or the failure lands a - whole sync later than the mistake.""" + whole sync later than the mistake. + """ assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS) def test_the_producer_allows_nothing_the_consumer_cannot_emit(self): """The allowlist is also a guard against a misbehaving or spoofed consumer invoking arbitrary methods on trainer tensors, so it must not be wider than the ops the bake can actually produce. ``to`` in particular - would let a replay change dtype or device.""" + would let a replay change dtype or device. + """ assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values()) assert "to" not in ALLOWED_OPS def test_transpose_via_t_is_serveable(self): """The specific regression: a chain recorded from ``.t()`` must pass the - producer's guard.""" + producer's guard. + """ fake = FakeRDTTensor( name="w", shape=torch.Size((4, 6)), @@ -1470,7 +1529,8 @@ def test_transpose_via_t_is_serveable(self): class TestRequiresTheRayExecutor: """The engine's data plane is Ray's, so a non-Ray executor cannot work. The check is at construction because the alternative is an opaque failure during - the first handshake, long after the misconfiguration.""" + the first handshake, long after the misconfiguration. + """ def _construct(self, backend): cfg = SimpleNamespace( @@ -1497,7 +1557,8 @@ def test_ray_is_accepted(self): def test_a_custom_executor_class_is_left_alone(self): """A `type[Executor]` override is deliberate and unjudgeable here, so it - must not be rejected for merely not being the string "ray".""" + must not be rejected for merely not being the string "ray". + """ class _CustomExecutor: pass diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 5cfab1957d65..7036d102897a 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -236,7 +236,8 @@ def test_rebuilt_views_have_the_requested_dtype_and_geometry(self, server_factor def test_views_of_one_storage_do_not_overlap(self, server_factory): """One IPC export per storage, one as_strided view per name — the whole point of the storage/view split. Overlapping views would serve the same - bytes for two different weights.""" + bytes for two different weights. + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A, nbytes=64) @@ -259,7 +260,8 @@ def test_publish_records_the_group_as_unfreed(self, server_factory): def test_the_default_lookahead_is_one(self): """1 is the sweet spot: group N+1 is gathered AND pullable while N is being pulled, with resident memory at its floor of 2 groups — the bound - the larger-model runs size against.""" + the larger-model runs size against. + """ assert DEFAULT_GATHER_LOOKAHEAD == 1 @@ -274,7 +276,8 @@ def test_single_consumer_signal_releases_the_group(self, server_factory): def test_group_is_held_until_the_last_live_consumer_signals(self, server_factory): """Every live consumer signals every owner; releasing on the first - signal would drop storage other consumers are still reading.""" + signal would drop storage other consumers are still reading. + """ server = server_factory() server.begin_sync(3) _publish(server, GI_A, GROUP_A) @@ -301,7 +304,8 @@ def test_a_signal_that_beats_its_publish_is_completed_at_publish( ): """A consumer with nothing to pull from a group signals it at sync start, which can precede the publish. The publish must then release the - group rather than wait for a signal that will never come again.""" + group rather than wait for a signal that will never come again. + """ server = server_factory() server.begin_sync(1) server.free_group(GI_A) @@ -324,7 +328,8 @@ def test_freed_groups_are_reported_back_to_the_engine(self, server_factory): the server says the group is gone — and it hears that ONLY through ``wait_freed`` / ``end_sync``, never a publish return (a freed notice riding an unharvested async publish while the engine blocks in - ``wait_freed`` would wedge the loop).""" + ``wait_freed`` would wedge the loop). + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -353,7 +358,8 @@ def test_release_drops_cache_and_count_entries(self, server_factory): class _CountingSource: """`iter_groups` yielding one CPU-tensor group per call, built lazily so a - group's tensors exist only once the engine's credit gate lets it gather.""" + group's tensors exist only once the engine's credit gate lets it gather. + """ def __init__(self, groups): self._groups = groups @@ -366,7 +372,8 @@ def iter_groups(self): class _ResidencyDict(dict): """The engine's `_inflight`, instrumented: `max_len` is the high-water mark of groups whose CUDA-IPC export refs were alive at once — the trainer's - resident-memory bound in groups.""" + resident-memory bound in groups. + """ def __init__(self): super().__init__() @@ -380,7 +387,8 @@ def __setitem__(self, key, value): def _loop_engine(server, n_groups, *, lookahead): """A real engine wired straight to a real (non-Ray) server, with only the state `_run_gather_loop` reads. The `_rpc` seam dispatches inline; - `_publish_async` sees no `.remote` and runs inline too.""" + `_publish_async` sees no `.remote` and runs inline too. + """ groups = [(f"model.layers.{gi}.w",) for gi in range(n_groups)] e = ShardedRDTTrainerWeightTransferEngine.__new__( ShardedRDTTrainerWeightTransferEngine @@ -403,14 +411,16 @@ class TestGatherCredit: serveable immediately, and the ENGINE's loop stops gathering while more than `gather_lookahead` groups are unfreed. So at most `lookahead + 1` groups are resident — two at the default — while group N+1 is pulled the instant N's - pulls finish.""" + pulls finish. + """ @pytest.fixture def gather_engine(self, monkeypatch): """Neutralize the two CUDA touches in the export path so the REAL gather loop runs on CPU: `.cuda()` becomes identity, and the storage export ships the `(nbytes, ..., device, ...)` tuple the fixture's - stubbed `rebuild_cuda_tensor` expects.""" + stubbed `rebuild_cuda_tensor` expects. + """ monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self) monkeypatch.setattr( trainer_mod, @@ -422,7 +432,8 @@ def gather_engine(self, monkeypatch): def test_publish_never_blocks(self, server_factory): """Publishing must not park on a credit: the whole point is that a gathered group becomes serveable immediately. Three publishes with no - frees, inline on this thread — a block here hangs the test.""" + frees, inline on this thread — a block here hangs the test. + """ server = server_factory(gather_lookahead=1) server.begin_sync(1) for gi in range(3): @@ -436,7 +447,8 @@ def test_lookahead_one_holds_at_most_two_groups_resident( than 2 groups of gathered weights, no matter how the consumer paces — and the pipeline still overlaps, because the consumer here refuses to free group N until N+1 is already published (pullable). A gate an - off-by-one too tight deadlocks this test; too loose fails max_len.""" + off-by-one too tight deadlocks this test; too loose fails max_len. + """ server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -478,7 +490,8 @@ def test_lookahead_zero_holds_exactly_one_group( ): """0 is the serialized baseline: the loop waits for group N to be freed before gathering N+1, so exactly ONE group is ever resident and nothing - overlaps. A floor of 1 on the bound would silently pipeline instead.""" + overlaps. A floor of 1 on the bound would silently pipeline instead. + """ server = server_factory(gather_lookahead=0) engine = gather_engine(server, 4, lookahead=0) @@ -503,8 +516,9 @@ def _run(): def test_the_residency_bound_scales_with_the_lookahead( self, server_factory, gather_engine ): - """lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing - consumer sees exactly 3 groups gathered before the loop parks.""" + """Lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing + consumer sees exactly 3 groups gathered before the loop parks. + """ server = server_factory(gather_lookahead=2) engine = gather_engine(server, 8, lookahead=2) @@ -540,7 +554,8 @@ def test_wait_freed_returns_the_freed_backlog(self, server_factory): def test_a_gather_error_releases_a_blocked_wait_freed(self, server_factory): """Otherwise a trainer-side failure on another rank deadlocks this one inside its credit gate. It must RAISE, not return empty — an empty - return would spin the engine straight back into the wait.""" + return would spin the engine straight back into the wait. + """ server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -621,7 +636,8 @@ def test_begin_sync_sets_the_barrier_target(self, server_factory): def test_the_target_is_floored_at_one(self, server_factory): """A bare/zero call must not make every group free instantly (or divide - the barrier by zero).""" + the barrier by zero). + """ server = server_factory() server.begin_sync(0) assert server._live_count == 1 @@ -629,7 +645,8 @@ def test_the_target_is_floored_at_one(self, server_factory): def test_a_degraded_sync_lowers_the_target(self, server_factory): """FT degraded sync: the live count is the WHOLE degraded-sync mechanism on this side — the group releases after the live consumers - alone.""" + alone. + """ server = server_factory() server.begin_sync(2) # 4 provisioned, 2 alive _publish(server, GI_A, GROUP_A) @@ -653,7 +670,8 @@ def test_begin_sync_clears_per_sync_state(self, server_factory): def test_a_straggler_signal_cannot_credit_the_next_sync(self, server_factory): """begin_sync resets the counts, so a signal that leaked past the sync boundary would over-credit — the reason the consumer drains its fired - signals before finishing. Pin the reset side of that contract.""" + signals before finishing. Pin the reset side of that contract. + """ server = server_factory() server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -675,7 +693,8 @@ def test_the_packed_destination_cache_survives_the_sync_boundary( self, server_factory ): """The packed layout repeats every sync — that is what makes caching the - destination views worth ~5ms per 384-spec group.""" + destination views worth ~5ms per 384-spec group. + """ server = server_factory() server._pack_dsts[("sentinel",)] = (0, []) server.begin_sync(1) @@ -686,7 +705,8 @@ class TestServedNamesGuard: """A producer serves only the names it holds: its stage's groups, and within them its own EP coordinate's experts plus the replicated names. Without this guard a misrouted pull would block forever in the cache wait - for a name this rank never publishes.""" + for a name this rank never publishes. + """ def test_a_pull_for_an_unserved_name_fails_loudly(self, server_factory): server = server_factory(served_names=list(GROUP_A)) @@ -702,7 +722,8 @@ def test_the_error_names_the_unserved_weights(self, server_factory): def test_served_names_none_accepts_anything(self, server_factory): """Gather-to-all: every producer holds every group, so there is nothing - to guard. The call proceeds to the cache wait (not exercised here).""" + to guard. The call proceeds to the cache wait (not exercised here). + """ server = server_factory(served_names=None) assert server._served_names is None @@ -749,7 +770,8 @@ def _wait(): def test_end_sync_fails_when_a_consumer_never_signals(self, server_factory): """The exact mid-sync death signature: the group is published, one of - the two live consumers that owed a signal is gone.""" + the two live consumers that owed a signal is gone. + """ server = server_factory(stall_timeout_s=0.3) server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -763,7 +785,8 @@ def test_the_error_reaches_every_blocked_waiter(self, server_factory): """The watchdog fires on the existing `set_gather_error` channel, so one stall unwinds the whole rank through one path rather than each waiter timing out separately — here both a consumer parked in the serve cache - wait and the engine parked in its credit gate.""" + wait and the engine parked in its credit gate. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -796,7 +819,8 @@ def test_progress_anywhere_keeps_a_slow_sync_alive(self, server_factory): """A consumer that is slow but signaling must never trip it: the stamp is global to the producer, so a steady trickle of signals holds the whole rank open. Nine sequential publishes with a 0.3s timeout each take - longer in total than the timeout, and none may fire.""" + longer in total than the timeout, and none may fire. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) for gi in range(9): @@ -809,7 +833,8 @@ def test_progress_anywhere_keeps_a_slow_sync_alive(self, server_factory): def test_begin_sync_resets_the_progress_stamp(self, server_factory): """Syncs are minutes apart. Without the reset the first publish of sync N+1 - would measure its stall from somewhere inside sync N and fire immediately.""" + would measure its stall from somewhere inside sync N and fire immediately. + """ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -831,7 +856,8 @@ def test_the_lifecycle_survives_concurrent_publishes_and_signals( self, server_factory ): """Every group published must end up freed and out of the cache, with - end_sync returning cleanly — the property whose violation is a hang.""" + end_sync returning cleanly — the property whose violation is a hang. + """ server = server_factory(gather_lookahead=2) server.begin_sync(1) groups = [(gi, (f"model.layers.{gi}.w",)) for gi in range(12)] @@ -866,7 +892,8 @@ def test_concurrent_signals_of_one_group_release_it_exactly_once( ): """All live consumers signaling at once: the release must happen on the last signal and only once, or the engine drops IPC refs twice (or - never).""" + never). + """ server = server_factory() server.begin_sync(8) _publish(server, GI_A, GROUP_A) @@ -893,7 +920,8 @@ class TestFakeServerAgreesWithTheRealOne: """`test_sharded_rdt_trainer.py`'s `_FakeProducerServer` is a second, independent model of this protocol that the engine-side tests assert on. Pin the two against each other so the fake cannot silently drift from the - semantics the real server enforces.""" + semantics the real server enforces. + """ @staticmethod def _fake(): @@ -924,7 +952,8 @@ class TestExportRing: """The packed export path. Its one hard invariant is that the trainer packs a group exactly the way ``publish_group`` rebuilds it — a layout disagreement would serve the wrong bytes with nothing downstream to catch - it.""" + it. + """ @staticmethod def _engine(): @@ -1016,7 +1045,8 @@ def test_a_slot_is_never_rewritten_while_its_group_is_still_live( ): """Free the later group of each pair first and hold the earlier one live until the NEXT group is gathered -- the window a modular slot counter got - wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's).""" + wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's). + """ server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -1140,7 +1170,8 @@ class TestSharedServeSlots: def test_one_deployment_never_shares(self, sharing_server): """With the width equal to the live count every group is a singleton, so - each call packs for itself out of its own ring.""" + each call packs for itself out of its own ring. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(2, [0, 1]) _publish(server, GI_A, GROUP_A) @@ -1220,7 +1251,8 @@ def test_the_slot_follows_the_consumers_issue_order(self, sharing_server): def test_a_seqless_call_is_refused(self, sharing_server): """The slot comes from ``seq``, so a caller that sends none cannot be - served safely.""" + served safely. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(1, [0]) _publish(server, GI_A, GROUP_A) @@ -1230,7 +1262,8 @@ def test_a_seqless_call_is_refused(self, sharing_server): def test_a_failed_pack_reaches_every_sharer(self, sharing_server): """Only the packer touches the buffer, so its failure has to be - republished: a waiter must raise rather than read a half-written slot.""" + republished: a waiter must raise rather than read a half-written slot. + """ server = sharing_server(workers_per_replica=2) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1254,7 +1287,8 @@ def test_a_failed_pack_reaches_every_sharer(self, sharing_server): def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server): """A dead deployment must not be waited for: a survivor left alone in - its group packs for itself instead of blocking on the corpse.""" + its group packs for itself instead of blocking on the corpse. + """ server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0) server.begin_sync(3, [0, 1, 3]) # consumer 2 died _publish(server, GI_A, GROUP_A) @@ -1268,7 +1302,8 @@ def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server): def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server): """A consumer that dies INSIDE the sync window has no detector, exactly - as for the free barrier; the rendezvous waits on the same channel.""" + as for the free barrier; the rendezvous waits on the same channel. + """ server = sharing_server(workers_per_replica=2, stall_timeout_s=0.3) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1280,7 +1315,8 @@ def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server): class TestSharedReservation: """``reserve_serve_buffer`` under sharing: one ring per GROUP, and the - init-time check that the group's members pull the same chunks.""" + init-time check that the group's members pull the same chunks. + """ def test_matching_plans_reserve_one_ring_between_them(self, sharing_server): server = sharing_server(workers_per_replica=2) @@ -1302,7 +1338,8 @@ def test_sharing_off_reserves_a_ring_per_consumer(self, sharing_server): def test_mismatched_plans_fail_at_reservation(self, sharing_server): """Sharing rests on the sharers pulling the same chunks. Checked here, where it names both consumers, rather than at serve time, where it is a - rendezvous nobody completes.""" + rendezvous nobody completes. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") @@ -1311,7 +1348,8 @@ def test_mismatched_plans_fail_at_reservation(self, sharing_server): def test_different_groups_may_differ(self, sharing_server): """Only the members of ONE group have to agree; worker 0 and worker 1 of - the same deployment hold different slices and never share a slot.""" + the same deployment hold different slices and never share a slot. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") server.reserve_serve_buffer(1, 1024, "plan-b") @@ -1320,7 +1358,8 @@ def test_different_groups_may_differ(self, sharing_server): def test_a_missing_digest_skips_the_check(self, sharing_server): """The digest is optional on the wire, so an engine that does not send - one still reserves.""" + one still reserves. + """ server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, None) server.reserve_serve_buffer(2, 1024, None) diff --git a/tests/distributed/test_sharded_rdt_trainer.py b/tests/distributed/test_sharded_rdt_trainer.py index 9339e503c45f..df9628cc8501 100644 --- a/tests/distributed/test_sharded_rdt_trainer.py +++ b/tests/distributed/test_sharded_rdt_trainer.py @@ -34,7 +34,8 @@ class _ListSource(WeightSource): """A WeightSource over an explicit ordered (name, cpu-tensor) list, so the - sharded-RDT group/order logic can be tested without a real model.""" + sharded-RDT group/order logic can be tested without a real model. + """ def __init__(self, pairs): self._pairs = list(pairs) @@ -54,7 +55,8 @@ class _FakeProducerServer: per-group barrier: publish/free are keyed by GROUP INDEX, signals count to the ``begin_sync`` live total, publish_group returns nothing, and freed groups flow back only through wait_freed / end_sync (see - test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne).""" + test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne). + """ def __init__(self, auto_free=True): self.order: list[str] = [] @@ -98,7 +100,8 @@ def wait_freed(self): must already have a banked credit when the gate asks (auto_free, or a test's own free_group) — anything else is the deadlock the real watchdog would kill, so fail loudly. Not appended to ``order``: it is - pacing, not a lifecycle milestone.""" + pacing, not a lifecycle milestone. + """ assert self._pending_freed, ( "wait_freed with nothing freed: the gather loop would deadlock" ) @@ -121,7 +124,8 @@ def _rdt_engine_with_fake_server( ): """Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the - group-major metadata, and route _rpc to the fake.""" + group-major metadata, and route _rpc to the fake. + """ import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod # reduce_tensor needs CUDA; the fake server never rebuilds, so stub it. @@ -321,7 +325,8 @@ def update_weights(self, update_info): class _OwnedSource(_ListSource): """A source holding only some groups' names, like a pipeline-parallel rank. - Takes group indices for convenience and declares the names inside them.""" + Takes group indices for convenience and declares the names inside them. + """ def __init__(self, pairs, owned_group_idx): super().__init__(pairs) @@ -387,7 +392,8 @@ def __iter__(self): ) def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): """The free barrier's target is one integer per sync: the live consumer - count, defaulting to the whole provisioned fleet.""" + count, defaulting to the whole provisioned fleet. + """ server = _FakeProducerServer(auto_free=True) engine = _rdt_engine_with_fake_server( _rdt_source_two_layers(), @@ -405,7 +411,8 @@ def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): class TestLiveCountPlumbing: """``send_weights(live_consumer_ids)`` -> the barrier target. The provisioned geometry is frozen; a degraded sync only lowers the target and - narrows each slot-sharing group's rendezvous.""" + narrows each slot-sharing group's rendezvous. + """ @staticmethod def _engine(num_consumers, world=2, rank=0): @@ -438,7 +445,8 @@ def test_a_full_live_set_matches_the_provisioned_count(self): def test_which_consumers_died_does_not_matter_only_how_many(self): """The whole point of the barrier: no routed per-producer targets, so - the identity of the dead consumer is irrelevant to the producers.""" + the identity of the dead consumer is irrelevant to the producers. + """ counts = [] for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]): engine, got = self._engine(num_consumers=8) @@ -449,7 +457,8 @@ def test_which_consumers_died_does_not_matter_only_how_many(self): def test_the_live_ids_travel_with_the_count(self): """The count sizes the free barrier, the ids size the slot-sharing rendezvous, so a producer that shares slots can tell WHICH consumers it - is still waiting for. They must describe the same set.""" + is still waiting for. They must describe the same set. + """ engine, got = self._engine(num_consumers=8) engine.send_weights([6, 0, 2, 2]) ((count, ids),) = got @@ -517,7 +526,8 @@ def held_names(self): def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): """Every name must be held somewhere or it can never be served — caught - when the holdings are transposed, naming the orphan.""" + when the holdings are transposed, naming the orphan. + """ with pytest.raises(ValueError, match="no trainer rank holds"): _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]), @@ -531,7 +541,8 @@ def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch): """Only the sender's metadata reaches the consumers, so a rank describing - just its own share must fail loudly rather than silently drop the rest.""" + just its own share must fail loudly rather than silently drop the rest. + """ engine = _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]), is_sender=False, @@ -554,7 +565,8 @@ def _serve_ring_server(src_name, src): """A producer server with one cached tensor and a pre-seeded serve ring, so a pull needs no Ray and no NIXL registration. Returns (server, serve) where ``serve(chain)`` packs one spec into the SAME ring slot every time — which is - what puts the destination-view cache, and only it, under test.""" + what puts the destination-view cache, and only it, under test. + """ from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( _RDTProducerServer, ) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index ef7f523b78dc..2274f15d97bd 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -43,6 +43,7 @@ def distributed_run(fn, world_size, timeout=60): fn: Function to run in each process world_size: Number of processes to spawn timeout: Maximum time in seconds to wait for processes (default: 60) + """ number_of_processes = world_size processes = [] @@ -418,7 +419,8 @@ def test_tensor_broadcast(): def _dumps_oob(obj) -> tuple[bytes, list]: """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor - dispatch table + out-of-band buffers >= 1MiB.""" + dispatch table + out-of-band buffers >= 1MiB. + """ buffers = [] def callback(buf: pickle.PickleBuffer) -> bool: @@ -489,7 +491,8 @@ def test_tensor_pickle_roundtrip(case: str): @pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) def test_reduce_tensor_fallback(case: str): """Tensors the zero-copy reducer can't safely alias must fall back to - torch's default reduction.""" + torch's default reduction. + """ if case == "cuda": if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -634,11 +637,9 @@ def test_acquire_read_releases_slot_when_reader_raises(): def test_warning_logs(caplog_vllm): - """ - Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals + """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals when indefinite=False, and are not emitted when indefinite=True. """ - # Patch the warning log interval to every 1 ms during reads with mock.patch( "vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL", @@ -754,7 +755,8 @@ def test_remote_subscribe_addr_unique_concurrent_writers( Pre-fix, the writer probed a port with get_open_port() and bound it afterwards; pinning the probe to one free port makes every writer bind the same port and fail deterministically on that code path, - while the late-binding implementation never consults the probe.""" + while the late-binding implementation never consults the probe. + """ from vllm.distributed.device_communicators import shm_broadcast colliding_port = get_open_port() diff --git a/tests/distributed/test_shm_buffer.py b/tests/distributed/test_shm_buffer.py index e56c71b1c241..6f61c28322f5 100644 --- a/tests/distributed/test_shm_buffer.py +++ b/tests/distributed/test_shm_buffer.py @@ -12,20 +12,20 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase): - """Test suite for the ring buffer implementation""" + """Test suite for the ring buffer implementation.""" def setUp(self): - """Set up test fixtures""" + """Set up test fixtures.""" self.buffer_size = 4096 self.ring_buffer = None def tearDown(self): - """Clean up after tests""" + """Clean up after tests.""" if self.ring_buffer: self.ring_buffer.close() def test_buffer_opening(self): - """Test opening an existing buffer""" + """Test opening an existing buffer.""" # First create a buffer self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True @@ -39,7 +39,7 @@ def test_buffer_opening(self): ) def test_buffer_access(self): - """Test accessing allocated buffers""" + """Test accessing allocated buffers.""" self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True ) @@ -61,7 +61,7 @@ def test_buffer_access(self): self.assertEqual(read_id, monotonic_id) def test_memory_error_on_full_buffer(self): - """Test that MemoryError is raised when buffer is full""" + """Test that MemoryError is raised when buffer is full.""" small_buffer_size = 200 self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=small_buffer_size, create=True @@ -76,7 +76,7 @@ def test_memory_error_on_full_buffer(self): self.ring_buffer.allocate_buf(1) # Would exceed buffer capacity def test_allocation_and_free(self): - """Test allocation and freeing of buffers""" + """Test allocation and freeing of buffers.""" small_buffer_size = 200 self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=small_buffer_size, create=True @@ -96,7 +96,7 @@ def test_allocation_and_free(self): self.assertEqual(freed_ids[0], i) def test_clear_buffer(self): - """Test clearing the buffer""" + """Test clearing the buffer.""" self.ring_buffer = SingleWriterShmRingBuffer( data_buffer_size=self.buffer_size, create=True ) @@ -180,7 +180,7 @@ def ring_allocate(allocate_size): def main(): - """Main function demonstrating usage and running tests""" + """Main function demonstrating usage and running tests.""" print("=== SingleWriterShmRingBuffer Test Suite ===\n") # Run unit tests diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py index 031937d7f5e1..fad1d2a00666 100644 --- a/tests/distributed/test_utils.py +++ b/tests/distributed/test_utils.py @@ -31,7 +31,8 @@ def get_cuda_visible_devices(self): def test_cuda_device_count_stateless(): """Test that cuda_device_count_stateless changes return value if - CUDA_VISIBLE_DEVICES is changed.""" + CUDA_VISIBLE_DEVICES is changed. + """ if current_platform.is_rocm(): pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.") actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 1353def73862..ced71ead76c5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1166,7 +1166,8 @@ def finish_weight_update(self, weight_version: str | None = None) -> None: def _module_with(*pairs): """A tiny nn.Module exposing the given (name, tensor) pairs as parameters, - so trainer tests can build a ModuleSource without a real model.""" + so trainer tests can build a ModuleSource without a real model. + """ module = torch.nn.Module() for name, tensor in pairs: module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) @@ -1287,7 +1288,8 @@ def test_source_is_reiterable(self): def test_metadata_agrees_with_iteration(self): """The two channels must line up element-for-element: engines declare - the round from `metadata()` and then send what iteration yields.""" + the round from `metadata()` and then send what iteration yields. + """ source = ModuleSource( _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3))) ) @@ -1301,11 +1303,13 @@ def test_metadata_agrees_with_iteration(self): class TestWeightSourceGroupContract: """`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what backends gather and free by, so the default must agree with - `layerwise_groups` over `metadata()`, restricted to what this rank holds.""" + `layerwise_groups` over `metadata()`, restricted to what this rank holds. + """ class _Source(WeightSource): """Minimal source over an ordered (name, tensor) list, optionally holding - only some names (in which case it iterates only their groups).""" + only some names (in which case it iterates only their groups). + """ def __init__(self, names, held=None, reverse=False): self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)] @@ -1360,7 +1364,8 @@ def test_groups_keeps_a_partially_held_group_whole(self): def test_groups_order_follows_metadata_not_the_declaration(self): """The declaration is a SET of names; the groups it selects still come out in metadata order, so each pairs with the right ``iter_groups()`` - batch however the source listed them.""" + batch however the source listed them. + """ names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"] source = self._source( names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"] @@ -1394,14 +1399,16 @@ def test_iter_groups_yields_only_held_groups(self): def test_out_of_order_iteration_raises(self): """Materializing is usually a collective, so a rank that iterates out of - order deadlocks its peers -- fail loudly instead.""" + order deadlocks its peers -- fail loudly instead. + """ source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True) with pytest.raises(RuntimeError, match="iteration order must match"): list(source.iter_groups()) def test_a_source_may_override_iter_groups(self): """The extension point: materialize a whole group in one step instead of - one generator resume per tensor.""" + one generator resume per tensor. + """ calls = [] class _Batched(TestWeightSourceGroupContract._Source): @@ -1422,7 +1429,8 @@ class TestDeferredProcessingContract: caller that takes over the update tail (running its own `finalize_layerwise_reload` instead of going through `finish_weight_update`) reads the flag and calls the method. Both must be answerable on any engine, or - that caller ends up reaching through a getattr.""" + that caller ends up reaching through a getattr. + """ def _engines(self): return { @@ -1436,7 +1444,8 @@ def test_every_engine_declares_whether_it_defers(self): def test_every_engine_can_be_drained(self): """The default is a no-op, so a caller never has to check whether the - method exists before calling it.""" + method exists before calling it. + """ for name, cls in self._engines().items(): assert callable(cls.drain_pending), name @@ -1530,7 +1539,8 @@ def test_shutdown_default_is_noop(self): ) def test_ipc_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; - the packed wire param rides the init info, not the per-round update_info.""" + the packed wire param rides the init info, not the per-round update_info. + """ client = RecordingClient() engine = IPCTrainerWeightTransferEngine( client=client, @@ -1549,7 +1559,8 @@ def test_ipc_trainer_send_weights_drives_client_in_order(): def test_ipc_trainer_init_ships_packed_to_worker(): """trainer_init drives the inference-side init handshake and propagates the - must-agree `packed` flag to the worker.""" + must-agree `packed` flag to the worker. + """ if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU (CUDA IPC handles).") @@ -1570,7 +1581,8 @@ def test_ipc_trainer_init_ships_packed_to_worker(): def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the inference-side init handshake with the worker-shaped init info (rank_offset=1) while opening its own endpoint, - and propagates the must-agree wire params to the worker.""" + and propagates the must-agree wire params to the worker. + """ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod # Bypass the real NCCL rendezvous. @@ -1611,7 +1623,8 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): """The worker engine reads packed + buffer geometry from the trainer-supplied init info at the handshake, not from the config or the - per-round update info.""" + per-round update info. + """ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod monkeypatch.setattr( @@ -1644,7 +1657,8 @@ def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): """Non-sender trainer ranks build an engine without opening an endpoint or - touching the client; they only join the collectives in send_weights.""" + touching the client; they only join the collectives in send_weights. + """ client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( init_info=NCCLTrainerInitInfo( @@ -1674,7 +1688,8 @@ def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): def test_nccl_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; the packed wire params ride the init handshake, not the - per-round update_info.""" + per-round update_info. + """ client = RecordingClient() engine = NCCLTrainerWeightTransferEngine( client=client, @@ -1695,7 +1710,8 @@ def test_nccl_trainer_send_weights_drives_client_in_order(): class _ScriptedSource(WeightSource): """Declares `meta` but yields whatever `pairs` says — used to drive the - metadata/iteration agreement checks.""" + metadata/iteration agreement checks. + """ def __init__(self, meta, pairs): self._meta = meta @@ -1731,7 +1747,8 @@ def test_nccl_trainer_init_requires_source(): def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): """The worker sizes its buffers (and cuts packed chunks) from metadata(), so - iteration disagreeing with it must raise rather than corrupt the stream.""" + iteration disagreeing with it must raise rather than corrupt the stream. + """ meta = [ ParamMeta("w", torch.float32, (4,)), ParamMeta("b", torch.float32, (2,)), @@ -1745,7 +1762,8 @@ def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch): """A source that declares one wire dtype and materializes another would make - the two sides disagree on every byte offset.""" + the two sides disagree on every byte offset. + """ meta = [ParamMeta("w", torch.float32, (4,))] engine = _mock_group_engine( _ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]), @@ -1772,7 +1790,8 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch): def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch): """NCCL sends numel elements from data_ptr(), so a non-contiguous view must - be linearized first or the worker receives unrelated memory.""" + be linearized first or the worker receives unrelated memory. + """ base = torch.arange(6, dtype=torch.float32).reshape(2, 3) view = base.t() # non-contiguous meta = [ParamMeta("w", torch.float32, tuple(view.shape))] @@ -1842,7 +1861,8 @@ def _sparse_patch(device: str = "cpu") -> SparseWeightPatch: def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the init handshake with the worker-shaped init info; sparse ships no packed wire params, so the worker - keeps its unpacked defaults. Sparse takes no `source`.""" + keeps its unpacked defaults. Sparse takes no `source`. + """ import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock()) @@ -1895,7 +1915,8 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch): def test_sparse_nccl_trainer_send_weights_empty_round_is_noop(): """A round with no patches must not touch the client (an empty sparse - update info is invalid by construction).""" + update info is invalid by construction). + """ client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() @@ -1931,7 +1952,8 @@ def test_sparse_nccl_trainer_rejects_source(): def test_sparse_nccl_trainer_validates_patch_before_any_rpc(): """Malformed patches must fail on the trainer, before start_weight_update: the worker's own checks only run once the broadcasts are already in flight, - where a size mismatch wedges both sides instead of raising.""" + where a size mismatch wedges both sides instead of raising. + """ client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() diff --git a/tests/distributed/test_weight_transfer_nccl_uid.py b/tests/distributed/test_weight_transfer_nccl_uid.py index 18a6442a2a94..58201376041a 100644 --- a/tests/distributed/test_weight_transfer_nccl_uid.py +++ b/tests/distributed/test_weight_transfer_nccl_uid.py @@ -403,7 +403,8 @@ def trainer_broadcast_tensor_torch_free( def _run_uid_transfer(trainer_task) -> None: """Mint a unique id, run the worker and the given trainer peer concurrently, and assert the tensor round-trips. Both ranks must enter init together -- - there is no store barrier on the UID path.""" + there is no store barrier on the UID path. + """ from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary _init_ray_for_weight_transfer() @@ -450,7 +451,8 @@ def _run_uid_transfer(trainer_task) -> None: ) def test_nccl_weight_transfer_between_processes_uid(): """Weight transfer over a pre-shared ncclUniqueId (no TCPStore), with both - ranks on vLLM's ``PyNcclCommunicator``.""" + ranks on vLLM's ``PyNcclCommunicator``. + """ _run_uid_transfer(trainer_broadcast_tensor_uid) @@ -460,6 +462,7 @@ def test_nccl_weight_transfer_between_processes_uid(): ) def test_nccl_weight_transfer_torch_free_trainer(): """UID weight transfer against a foreign peer that uses no torch and no vLLM - communicator, only a cupy buffer and raw ``ncclBroadcast``.""" + communicator, only a cupy buffer and raw ``ncclBroadcast``. + """ pytest.importorskip("cupy") _run_uid_transfer(trainer_broadcast_tensor_torch_free) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index e58379874601..9adfedc9f489 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -394,8 +394,7 @@ def test_media_io_kwargs_parser(arg, expected): ], ) def test_optimization_level(args, expected): - """ - Test space-separated optimization levels (-O 1, -O 2, -O 3) map to + """Test space-separated optimization levels (-O 1, -O 2, -O 3) map to optimization_level. """ parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) @@ -414,9 +413,7 @@ def test_optimization_level(args, expected): ], ) def test_mode_parser(args, expected): - """ - Test compilation config modes (-cc.mode=int) map to compilation_config. - """ + """Test compilation config modes (-cc.mode=int) map to compilation_config.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) parsed_args = parser.parse_args(args) assert parsed_args.compilation_config.mode == expected @@ -573,7 +570,8 @@ def test_attention_config(): def test_multi_node_world_size_includes_pcp(monkeypatch): """PCP expands the process world size, so the --nnodes divisibility check must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world - size of 1 and the launch is rejected before the engine starts.""" + size of 1 and the launch is rejected before the engine starts. + """ import vllm.config.vllm # PCP requires the V2 model runner, which is gated on Triton. @@ -799,7 +797,8 @@ def test_expand_json_human_readable_numbers(input_json, expected_json): ) def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): """Cloud storage URIs should not be passed to get_model_path() - when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.""" + when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs. + """ import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -811,7 +810,8 @@ def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): """Cloud storage tokenizer URI should not be passed to - get_model_path() when HF_HUB_OFFLINE=1.""" + get_model_path() when HF_HUB_OFFLINE=1. + """ import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -973,7 +973,8 @@ class TestDpDeviceIdSharding: def test_dp_supervisor_device_ids_stay_env_relative(self): """Regression test: the DP supervisor must pass env-relative indices, not physical IDs, because each child re-resolves --device-ids - against its inherited device-control env var.""" + against its inherited device-control env var. + """ import argparse from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids @@ -1000,7 +1001,8 @@ def test_dp_supervisor_shards_user_device_ids(self): def test_dp_rank_shards_user_assigned_gpu_ids(self): """get_physical_gpu_ids_for_local_dp_rank slices the user-provided - --device-ids list instead of recomputing from the env var.""" + --device-ids list instead of recomputing from the env var. + """ from vllm.platforms import current_platform from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 0c986f129139..63e44eacebb4 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -241,7 +241,8 @@ def _make_tool_result_request( self, tool_result_content ) -> AnthropicMessagesRequest: """Build a request with assistant tool_use followed by user - tool_result.""" + tool_result. + """ return _make_request( [ { @@ -326,7 +327,8 @@ def test_tool_result_with_image(self): def test_tool_result_with_text_and_image(self): """Mixed text+image tool_result: text in tool msg, image in user - msg.""" + msg. + """ request = self._make_tool_result_request( [ {"type": "text", "text": "Here is the screenshot"}, @@ -423,7 +425,8 @@ def test_tool_result_no_follow_up_when_no_images(self): class TestAttributionHeaderStripping: def test_billing_header_stripped_from_system(self): """Claude Code's x-anthropic-billing-header block should be - stripped to preserve prefix caching.""" + stripped to preserve prefix caching. + """ request = _make_request( [{"role": "user", "content": "Hello"}], system=[ @@ -480,7 +483,7 @@ class TestThinkingBlockConversion: """ def test_thinking_plus_text_in_assistant_message(self): - """thinking + text → reasoning field + plain-string content.""" + """Thinking + text → reasoning field + plain-string content.""" request = _make_request( [ {"role": "user", "content": "Write me some code."}, @@ -542,7 +545,7 @@ def test_thinking_only_in_assistant_message(self): assert asst.get("content") is None def test_thinking_plus_tool_use_in_assistant_message(self): - """thinking + tool_use: reasoning field set, tool_calls populated.""" + """Thinking + tool_use: reasoning field set, tool_calls populated.""" request = _make_request( [ {"role": "user", "content": "What is 2+2?"}, @@ -1147,7 +1150,8 @@ async def sse_input(): async def test_buffered_content_flushed_on_done_without_usage_chunk(self): """Content buffered during tool_use must be emitted even if the stream jumps straight from finish_reason to [DONE], skipping the - empty-choices usage chunk.""" + empty-choices usage chunk. + """ async def sse_input(): yield _make_stream_chunk( @@ -1251,7 +1255,8 @@ class TestStreamingCacheUsageSemantics: async def test_streaming_cache_fields_absent_then_populated(self): """First chunk lacks prompt_tokens_details (vLLM contract); message_start omits cache fields. The final chunk carries - prompt_tokens_details, so message_delta carries resolved values.""" + prompt_tokens_details, so message_delta carries resolved values. + """ async def sse_input(): yield _make_stream_chunk( @@ -1296,7 +1301,8 @@ async def sse_input(): @pytest.mark.asyncio async def test_streaming_no_cache_hit(self): """When the final chunk reports cached_tokens=0, message_delta carries - cache fields = 0 (cache miss); message_start still omits them.""" + cache fields = 0 (cache miss); message_start still omits them. + """ async def sse_input(): yield _make_stream_chunk( @@ -1337,7 +1343,8 @@ async def sse_input(): @pytest.mark.asyncio async def test_streaming_no_prompt_tokens_details_at_all(self): """If --enable-prompt-tokens-details is off, no chunk carries cache - info; both message_start and message_delta omit cache fields.""" + info; both message_start and message_delta omit cache fields. + """ async def sse_input(): yield _make_stream_chunk( @@ -1458,7 +1465,8 @@ def test_empty_completion_emits_one_text_block(self): class TestCacheSalt: def test_cache_salt_passed_through(self): """cache_salt on the Anthropic request reaches the converted - ChatCompletionRequest so prefix-cache isolation works via /v1/messages.""" + ChatCompletionRequest so prefix-cache isolation works via /v1/messages. + """ request = _make_request( [{"role": "user", "content": "Hello"}], cache_salt="tenant-abc-secret-salt", @@ -1663,7 +1671,8 @@ def _request_body() -> dict: @staticmethod def _conversion_error() -> ValidationError: """A real pydantic ValidationError like the one ChatCompletionRequest - construction raises when Anthropic input violates the OpenAI schema.""" + construction raises when Anthropic input violates the OpenAI schema. + """ class _StubRequest(BaseModel): stop: Annotated[list[str], Field(max_length=4)] | None = None @@ -1674,7 +1683,8 @@ class _StubRequest(BaseModel): def test_validation_error_returns_bad_request(self): """A pydantic ValidationError during Anthropic->OpenAI conversion is - surfaced as a 400 BadRequestError, not a 500.""" + surfaced as a 400 BadRequestError, not a 500. + """ handler = MagicMock(spec=AnthropicServingMessages) handler.create_messages.side_effect = self._conversion_error() diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py index 8c2e8623ad3c..8723c6ecc792 100644 --- a/tests/entrypoints/cohere/test_serving_streaming.py +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -87,7 +87,8 @@ def _make_chunk( citations: list[Any] | None = None, ) -> str: """Build the ``data: {...}\\n\\n`` SSE frame the production code - consumes.""" + consumes. + """ delta: dict[str, Any] = {} if role is not None: delta["role"] = role @@ -135,7 +136,8 @@ async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect - the emitted SSE frames.""" + the emitted SSE frames. + """ request = CohereChatV2Request( model="m", messages=[{"role": "user", "content": "hi"}], stream=True ) diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 668f11544858..7f2f771107c0 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -50,7 +50,8 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit.""" + endpoints back via the pipe, then exit. + """ ctx = zmq.Context() try: in_sock = make_zmq_socket( @@ -357,7 +358,8 @@ def run_with_exception_capture(): def test_gather_actual_addresses_end_to_end(): """Each child binds ROUTER/PULL with a kernel-picked port and reports the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`.""" + them via :py:meth:`gather_actual_addresses`. + """ host = "127.0.0.1" num_servers = 4 @@ -407,7 +409,8 @@ def test_gather_actual_addresses_end_to_end(): @pytest.mark.timeout(30) def test_gather_actual_addresses_child_crash_before_report(): """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots.""" + clear ``RuntimeError`` rather than hang or return ``None`` slots. + """ host = "127.0.0.1" num_servers = 2 placeholder_inputs = [ @@ -444,7 +447,8 @@ def test_gather_actual_addresses_child_crash_before_report(): def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): """The Rust frontend command carries every non-default arg as JSON. - Credentials must not reach the log line.""" + Credentials must not reach the log line. + """ import subprocess as subprocess_mod from vllm.entrypoints.launchers.cli_args import make_arg_parser diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index e84abe5d2617..cf2f678e40da 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -152,7 +152,7 @@ def test_multiple_valid_inputs(serve_parser): ### Tests for serve argument validation that run prior to loading def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser): - """Ensure validation fails if tool choice is enabled with no call parser""" + """Ensure validation fails if tool choice is enabled with no call parser.""" # If we enable-auto-tool-choice, explode with no tool-call-parser args = serve_parser.parse_args(args=["--enable-auto-tool-choice"]) with pytest.raises(TypeError): @@ -160,7 +160,7 @@ def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser): def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser): - """Ensure validation passes with tool choice enabled with a call parser""" + """Ensure validation passes with tool choice enabled with a call parser.""" args = serve_parser.parse_args( args=[ "--enable-auto-tool-choice", @@ -172,7 +172,7 @@ def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser): def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): - """Ensure validation fails if reasoning is enabled with auto tool choice""" + """Ensure validation fails if reasoning is enabled with auto tool choice.""" args = serve_parser.parse_args( args=[ "--enable-auto-tool-choice", @@ -186,7 +186,8 @@ def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): def test_passes_with_reasoning_parser(serve_parser): """Ensure validation passes if reasoning is enabled - with a reasoning parser""" + with a reasoning parser + """ args = serve_parser.parse_args( args=[ "--reasoning-parser", @@ -197,7 +198,7 @@ def test_passes_with_reasoning_parser(serve_parser): def test_chat_template_validation_for_happy_paths(serve_parser): - """Ensure validation passes if the chat template exists""" + """Ensure validation passes if the chat template exists.""" args = serve_parser.parse_args( args=["--chat-template", CHATML_JINJA_PATH.absolute().as_posix()] ) @@ -205,7 +206,7 @@ def test_chat_template_validation_for_happy_paths(serve_parser): def test_chat_template_validation_for_sad_paths(serve_parser): - """Ensure validation fails if the chat template doesn't exist""" + """Ensure validation fails if the chat template doesn't exist.""" args = serve_parser.parse_args(args=["--chat-template", "does/not/exist"]) with pytest.raises(VLLMValidationError): validate_parsed_serve_args(args) @@ -304,13 +305,13 @@ def test_sse_keep_alive_interval_non_integer(serve_parser, value): ], ) def test_middleware(serve_parser, cli_args, expected_middleware): - """Ensure multiple middleware args are parsed properly""" + """Ensure multiple middleware args are parsed properly.""" args = serve_parser.parse_args(args=cli_args) assert args.middleware == expected_middleware def test_default_chat_template_kwargs_parsing(serve_parser): - """Ensure default_chat_template_kwargs JSON is parsed correctly""" + """Ensure default_chat_template_kwargs JSON is parsed correctly.""" args = serve_parser.parse_args( args=["--default-chat-template-kwargs", '{"enable_thinking": false}'] ) @@ -318,7 +319,7 @@ def test_default_chat_template_kwargs_parsing(serve_parser): def test_default_chat_template_kwargs_complex(serve_parser): - """Ensure complex default_chat_template_kwargs JSON is parsed correctly""" + """Ensure complex default_chat_template_kwargs JSON is parsed correctly.""" kwargs_json = '{"enable_thinking": false, "custom_param": "value", "num": 42}' args = serve_parser.parse_args(args=["--default-chat-template-kwargs", kwargs_json]) assert args.default_chat_template_kwargs == { @@ -329,7 +330,7 @@ def test_default_chat_template_kwargs_complex(serve_parser): def test_default_chat_template_kwargs_default_none(serve_parser): - """Ensure default_chat_template_kwargs defaults to None""" + """Ensure default_chat_template_kwargs defaults to None.""" args = serve_parser.parse_args(args=[]) assert args.default_chat_template_kwargs is None @@ -355,7 +356,7 @@ def test_resolve_default_chat_template_kwargs(default_kwargs, cohere_format, exp def test_default_chat_template_kwargs_invalid_json(serve_parser): - """Ensure invalid JSON raises an error""" + """Ensure invalid JSON raises an error.""" with pytest.raises(SystemExit): serve_parser.parse_args( args=["--default-chat-template-kwargs", "not valid json"] @@ -381,7 +382,8 @@ def test_default_chat_template_kwargs_invalid_json(serve_parser): ) def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises): """Ensure that users don't misuse --served-model-name and end up with the default - model tag instead of the one they intended to serve.""" + model tag instead of the one they intended to serve. + """ # Call the serve subparser args.insert(0, "serve") # Create a dummy config file if the test case includes it @@ -402,7 +404,7 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises): ### Tests for LoRA target modules parsing def test_lora_target_modules_single(serve_parser): - """Test parsing single lora-target-modules argument""" + """Test parsing single lora-target-modules argument.""" args = serve_parser.parse_args( args=["--enable-lora", "--lora-target-modules", "o_proj"] ) @@ -410,7 +412,7 @@ def test_lora_target_modules_single(serve_parser): def test_lora_target_modules_multiple(serve_parser): - """Test parsing multiple lora-target-modules arguments""" + """Test parsing multiple lora-target-modules arguments.""" args = serve_parser.parse_args( args=[ "--enable-lora", @@ -424,6 +426,6 @@ def test_lora_target_modules_multiple(serve_parser): def test_lora_target_modules_default_none(serve_parser): - """Test that lora-target-modules defaults to None""" + """Test that lora-target-modules defaults to None.""" args = serve_parser.parse_args(args=[]) assert args.lora_target_modules is None diff --git a/tests/entrypoints/launchers/test_dp_supervisor.py b/tests/entrypoints/launchers/test_dp_supervisor.py index eb5bf0a08875..f96b020d10d8 100644 --- a/tests/entrypoints/launchers/test_dp_supervisor.py +++ b/tests/entrypoints/launchers/test_dp_supervisor.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for DPSupervisor: unit tests and lifecycle integration tests. +"""Tests for DPSupervisor: unit tests and lifecycle integration tests. Lifecycle integration tests replace child vLLM servers with lightweight aiohttp "fake" servers controlled by the test, so the suite runs without GPUs. @@ -463,8 +462,7 @@ async def serve(self): class MockVLLMServer: - """ - Minimal FastAPI server that mimics one vLLM replica. + """Minimal FastAPI server that mimics one vLLM replica. GET /health returns 200 when healthy, 503 otherwise. Health state is toggled by the test via set_healthy(). """ @@ -582,8 +580,7 @@ def launch_mock_vllm_with_drain( async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool: - """ - GET /health on the supervisor once and check for expected_status. + """GET /health on the supervisor once and check for expected_status. Pass expected_status=-1 to assert the supervisor is not listening yet (a connection error is expected). The supervisor only starts its HTTP @@ -705,8 +702,7 @@ async def _run_supervisor( @pytest.mark.asyncio async def test_basic_lifecycle(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) SIGTERM and shutdown """ @@ -777,8 +773,7 @@ async def test_basic_lifecycle_with_ssl(monkeypatch): @pytest.mark.asyncio async def test_failed_startup(monkeypatch): - """ - A) One of the vLLM servers crashes during startup. + """A) One of the vLLM servers crashes during startup. B) DPSupervisor detects this, and cleans up resources. """ args = _make_args() @@ -801,8 +796,7 @@ async def test_failed_startup(monkeypatch): @pytest.mark.asyncio async def test_becomes_unhealthy(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) Child process becomes unhealthy. D) Detected and shutdown. @@ -842,8 +836,7 @@ async def test_becomes_unhealthy(monkeypatch): @pytest.mark.asyncio async def test_dp_server_fails(monkeypatch): - """ - A) Supervisor is not listening while children are unhealthy. + """A) Supervisor is not listening while children are unhealthy. B) /health returns 200 once every child reports healthy. C) Child process fails. D) Detected and shutdown. @@ -886,8 +879,7 @@ async def test_dp_server_fails(monkeypatch): @pytest.mark.asyncio async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch): - """ - Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight + """Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight request drain). The supervisor is configured with shutdown_timeout=10, so its total wait budget is 10 + CHILD_EXIT_GRACE_S seconds. The children exit naturally within that window, so no force-kill should occur diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py index 7e6ba3d78d72..f7a250580727 100644 --- a/tests/entrypoints/launchers/test_run_batch.py +++ b/tests/entrypoints/launchers/test_run_batch.py @@ -509,9 +509,7 @@ def test_completions(): def test_completions_invalid_input(): - """ - Ensure that we fail when the input doesn't conform to the openai api. - """ + """Ensure that we fail when the input doesn't conform to the openai api.""" with ( tempfile.NamedTemporaryFile("w") as input_file, tempfile.NamedTemporaryFile("r") as output_file, @@ -602,9 +600,7 @@ def test_score(input_batch): def test_reasoning_parser(): - """ - Test that reasoning_parser parameter works correctly in run_batch. - """ + """Test that reasoning_parser parameter works correctly in run_batch.""" with ( tempfile.NamedTemporaryFile("w") as input_file, tempfile.NamedTemporaryFile("r") as output_file, @@ -767,8 +763,7 @@ def test_translation(): def test_tool_calling(): - """ - Test that tool calling works correctly in run_batch. + """Test that tool calling works correctly in run_batch. Verifies that requests with tools return tool_calls in the response. """ with ( diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py index 643eff7594de..6d18d825eba9 100644 --- a/tests/entrypoints/launchers/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -134,7 +134,6 @@ async def test_shutdown_on_engine_failure(tmp_path: Path): Starts a vLLM server, kills it to simulate a crash, then verifies that subsequent API calls fail appropriately. """ - port = get_open_port() # Redirect server output to a file rather than a pipe: nothing drains the diff --git a/tests/entrypoints/llm/offline_mode/test_offline_mode.py b/tests/entrypoints/llm/offline_mode/test_offline_mode.py index 555e8b9b7d07..ab7edd019ba9 100644 --- a/tests/entrypoints/llm/offline_mode/test_offline_mode.py +++ b/tests/entrypoints/llm/offline_mode/test_offline_mode.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for HF_HUB_OFFLINE mode""" +"""Tests for HF_HUB_OFFLINE mode.""" import importlib import sys diff --git a/tests/entrypoints/llm/test_accuracy.py b/tests/entrypoints/llm/test_accuracy.py index af607720c8b0..49d83b9cf6e8 100644 --- a/tests/entrypoints/llm/test_accuracy.py +++ b/tests/entrypoints/llm/test_accuracy.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file test accuracy of the vLLM server via LMEval. +"""This file test accuracy of the vLLM server via LMEval. It uses local-completions, which interacts with vLLM through the OAI API with N concurrent connections. This simulates real work usage of the API and makes @@ -33,7 +32,6 @@ def run_test(model_name, more_args=None): """Run the end to end accuracy test.""" - model_args = f"pretrained={model_name},max_model_len=4096" if more_args is not None: @@ -64,7 +62,6 @@ def run_test(model_name, more_args=None): @pytest.mark.parametrize("model", MODEL_NAMES) def test_lm_eval_accuracy_v1_engine(model): """Run with the V1 Engine.""" - more_args = None if current_platform.is_tpu(): # Limit compilation time for TPU V1 @@ -81,7 +78,6 @@ def test_lm_eval_accuracy_v1_engine(model): @pytest.mark.parametrize("model", FP8_KV_MODEL_NAMES) def test_lm_eval_accuracy_v1_engine_fp8_kv_cache(model): """Run with the V1 Engine.""" - more_args = None if current_platform.is_tpu(): # Limit compilation time for TPU V1 diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 6ee8d34ce138..276426e692c7 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -20,8 +20,7 @@ def text_llm(vllm_runner): @pytest.fixture(scope="function") def llm_for_failure_test(vllm_runner): - """ - Fixture for testing issue #26081. + """Fixture for testing issue #26081. Uses a small max_model_len to easily trigger length errors. """ with vllm_runner( @@ -67,8 +66,7 @@ def test_multi_chat(text_llm): def test_llm_chat_tokenization_no_double_bos(text_llm): - """ - LLM.chat() should not add special tokens when using chat templates. + """LLM.chat() should not add special tokens when using chat templates. Check we get a single BOS token for llama chat. """ messages = [ @@ -127,8 +125,7 @@ def test_chat_extra_kwargs(thinking_llm, enable_thinking): def test_chat_batch_failure_cleanup(llm_for_failure_test): - """ - Tests that if a batch call to llm.chat() fails mid-way + """Tests that if a batch call to llm.chat() fails mid-way (e.g., due to one invalid prompt), the requests that were already enqueued are properly aborted and do not pollute the queue for subsequent calls. diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py index 5c378de35dbe..a3511b9528b1 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects +"""Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects pre-processed mm_kwargs into the processor cache and reports MM cache hit rate metrics accurately. diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py index 14ac7ce0439b..5031cbea1a25 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py @@ -47,8 +47,7 @@ async def client(server): async def test_online_audio_in_video( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test video input with `audio_in_video=True`""" - + """Test video input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: @@ -97,8 +96,7 @@ async def test_online_audio_in_video( async def test_online_audio_in_video_multi_videos( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test multi-video input with `audio_in_video=True`""" - + """Test multi-video input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: @@ -151,8 +149,7 @@ async def test_online_audio_in_video_multi_videos( async def test_online_audio_in_video_interleaved( client: openai.AsyncOpenAI, video_assets: VideoTestAssets ): - """Test interleaved video/audio input with `audio_in_video=True`""" - + """Test interleaved video/audio input with `audio_in_video=True`.""" # we don't use video_urls above because they missed audio stream. video_path = video_assets[0].video_path with open(video_path, "rb") as f: diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index aac68334bf94..b12bf3a60b8c 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """E2E test for mixing `prompt_embeds` with `audio_embeds` in a single -Chat Completions request.""" +Chat Completions request. +""" import json diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py index c6ce165cf390..b72d94071425 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -68,7 +68,8 @@ async def client(server): def image_url() -> str: """Stable real image as a data URL, kept identical across both the text and prompt_embeds requests so any output difference must come from - how the text content is delivered.""" + how the text content is delivered. + """ return encode_image_url(ImageAsset("stop_sign").pil_image) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py index e6b8dc23e9f2..6d1c2f24e98f 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py @@ -34,8 +34,7 @@ def check_output_matches_terms(content: str, term_groups: list[list[str]]) -> bool: - """ - Check if content matches all required term groups. + """Check if content matches all required term groups. Each term group requires at least one of its terms to be present. All term groups must be satisfied. """ @@ -146,7 +145,8 @@ def describe_image_messages( ) -> list[dict]: """Build the system + user messages used by the completions-with-image family of tests. *extra_image_fields* is merged into the top-level - image content block (for uuid / bad-key tests).""" + image content block (for uuid / bad-key tests). + """ image_block: dict = { "type": "image_url", "image_url": {"url": image_url}, @@ -176,7 +176,8 @@ async def complete_and_check( temperature: float = 0.0, ) -> str: """Run a chat completion and assert the output is non-empty. - Returns the content string.""" + Returns the content string. + """ chat_completion = await client.chat.completions.create( model=model_name, messages=messages, diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index 20218ac6c875..467e23e7d3d2 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -259,7 +259,7 @@ async def _generator(prompt_idx: int, text: str) -> AsyncGenerator[RequestOutput @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_batched_echo_does_not_prepend_user_prompt() -> None: - """echo must only echo the assistant turn, never the user's prompt. + """Echo must only echo the assistant turn, never the user's prompt. With `add_generation_prompt` (the default) the response role is `assistant` while the last conversation message is the user's, so there is nothing to diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index aafb899f005c..e0f71287441e 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -121,7 +121,7 @@ async def _fake_preprocess_chat(*args, **kwargs): @pytest.mark.asyncio async def test_chat_error_non_stream(): - """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (non-streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -246,7 +246,7 @@ async def test_renderer_only_chat_request_skips_mm_cache(): @pytest.mark.asyncio async def test_chat_error_stream(): - """test finish_reason='error' returns 500 InternalServerError (streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -465,7 +465,8 @@ def test_system_message_warns_on_video(video_content): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500.""" + so the API returns 400 instead of 500. + """ with pytest.raises(Exception, match="json_schema.*must be provided"): ChatCompletionRequest( model=MODEL_NAME, @@ -519,7 +520,8 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): @pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + TypeError from the mode='before' comparison (which surfaces as HTTP 500). + """ with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): ChatCompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py index b415fa116dac..f1801701e3c5 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py @@ -81,7 +81,8 @@ async def test_chat_logit_bias_invalid(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_integer_key(client): """Test that a non-integer logit_bias key is rejected with a clean, - informative error instead of a raw 'invalid literal for int()' message.""" + informative error instead of a raw 'invalid literal for int()' message. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -101,7 +102,8 @@ async def test_chat_logit_bias_non_integer_key(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_numeric_value(client): """Test that a non-numeric logit_bias value is rejected with a message - that names the specific offending token, not just a generic TypeError.""" + that names the specific offending token, not just a generic TypeError. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -120,7 +122,8 @@ async def test_chat_logit_bias_non_numeric_value(client): @pytest.mark.asyncio async def test_chat_logit_bias_multiple_non_integer_keys(client): """Test that ALL invalid logit_bias keys are reported together, - not just the first one encountered.""" + not just the first one encountered. + """ with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 3e59e7402ea4..61fd2f1576bc 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -451,7 +451,7 @@ async def test_inconsistent_tool_choice_and_tools( async def test_max_tokens_with_tool_choice_required( client: openai.AsyncOpenAI, tool_choice ): - """ """ + """Requesting a single token with a required tool choice must not crash.""" models = await client.models.list() model_name: str = models.data[0].id diff --git a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py index 8c03fc1daa01..9c0a4db68b1e 100644 --- a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py +++ b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py @@ -52,9 +52,10 @@ def stop_sign_image_url(): @pytest.mark.asyncio async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI): - """en -> es translation through TranslateGemma's bundled chat template's + """En -> es translation through TranslateGemma's bundled chat template's ``text`` content branch, which depends on ``source_lang_code`` / - ``target_lang_code`` extra fields being preserved on the content part.""" + ``target_lang_code`` extra fields being preserved on the content part. + """ completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ @@ -92,13 +93,14 @@ async def test_translategemma_image_extra_lang_code_fields( client: openai.AsyncOpenAI, stop_sign_image_url: str, ): - """en -> es OCR-translation through TranslateGemma's bundled chat + """En -> es OCR-translation through TranslateGemma's bundled chat template's ``image`` content branch. Exercises the multimodal branch of ``_collect_extra_fields`` in ``_parse_chat_message_content_part``: the ``source_lang_code`` / ``target_lang_code`` extras must survive parsing and end up alongside the ``{"type": "image"}`` placeholder so the template's image branch (``content["type"] == 'image'``) renders the - right OCR-translation prompt.""" + right OCR-translation prompt. + """ completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py index 1eea6db48a7c..62f915b054c0 100644 --- a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -191,7 +191,8 @@ async def test_logprob_token_ids_stream_with_no_top_k(server): async def test_logprob_token_ids_default_behavior_unchanged(server): """Without `logprob_token_ids`, the response carries the natural top-k most-likely tokens. This guards against the new field accidentally - changing the default-path output.""" + changing the default-path output. + """ async with server.get_async_client() as client: resp = await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index e7ea15877c22..73cffa76cfb5 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1232,7 +1232,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): """Regression test: when the Mistral tokenizer path returns token IDs directly, we must still apply input length + max_tokens validation. """ - mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) @@ -1271,7 +1270,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(): """Regression test: MistralTokenizer token-id prompts must still enforce the max context length for the input itself (token_num >= max_model_len). """ - mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) @@ -1408,7 +1406,8 @@ async def render_chat_request(request): @pytest.mark.asyncio async def test_serving_chat_data_parallel_rank_extraction(): """Test that data_parallel_rank is properly extracted from header and - passed to engine.""" + passed to engine. + """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -1512,8 +1511,7 @@ async def _render_chat_prompt_token_ids( class TestServingChatWithHarmony: - """ - These tests ensure Chat Completion requests are being properly converted into + """These tests ensure Chat Completion requests are being properly converted into Harmony messages and Harmony response messages back into Chat Completion responses. These tests are not exhaustive, but each one was created to cover a specific case that we got wrong but is now fixed. @@ -1731,7 +1729,8 @@ async def test_simple_chat(self, serving_chat, stream, include_reasoning): @pytest.mark.asyncio async def test_system_message_without_tools(self, serving_chat, stream): """Leading system message produces a developer message with - DeveloperContent (# Instructions header).""" + DeveloperContent (# Instructions header). + """ messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello"}, @@ -1753,7 +1752,8 @@ async def test_system_message_without_tools(self, serving_chat, stream): @pytest.mark.asyncio async def test_system_message_with_tools(self, serving_chat, stream, weather_tools): """Leading system message is folded into the developer message - alongside tool definitions.""" + alongside tool definitions. + """ messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather?"}, @@ -2168,7 +2168,8 @@ async def test_non_tool_reasoning_empty_content_list(self, serving_chat): @pytest.mark.asyncio async def test_tool_choice_validation_without_parser(): """Test that tool_choice='required' or named tool without tool_parser - returns an appropriate error message.""" + returns an appropriate error message. + """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index ec7a977f1a77..b7ed422223bd 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -33,7 +33,6 @@ def _count_reasoning_decode_token_ids_between_markers( reasoning_end_ids: list[int], ) -> int | None: """Count decode tokens in the thinking span (after last start, before first end).""" - if not reasoning_start_ids or not reasoning_end_ids: raise ValueError("reasoning marker token id lists must be non-empty") @@ -152,8 +151,8 @@ async def client(request, server, server_with_auto_reasoning_config): @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True) async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): """Test that mixed requests (some with thinking_token_budget, some without) - complete successfully without errors.""" - + complete successfully without errors. + """ response_with_budget = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, @@ -182,7 +181,6 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI grouped into streamed chunks (a single chunk can carry several tokens under async scheduling / stream_interval > 1). Counting chunks under-counts. """ - tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False)) end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False)) @@ -227,7 +225,6 @@ async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_ Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with ``_count_reasoning_decode_token_ids_between_markers`` on full token ids. """ - _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [ ("budget", 1), ("budget", 12), diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 1bc6260cdebc..287664224123 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -350,7 +350,6 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str """Parallel sampling without streaming. A single request output contains a list of completions. """ - prompt = "What is an LLM?" n = 3 max_tokens = 50 @@ -396,7 +395,6 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): The tokens from multiple samples, are flattened into a single stream, with an index to indicate which sample the token belongs to. """ - prompt = "What is an LLM?" n = 3 max_tokens = 50 diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 513bde465d71..1fee99be5762 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -303,7 +303,7 @@ def test_completion_metrics_carries_both_timing_and_spec_decode(): @pytest.mark.asyncio async def test_completion_error_non_stream(): - """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (non-streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -439,7 +439,7 @@ async def test_renderer_only_completion_request_skips_mm_cache(): @pytest.mark.asyncio async def test_completion_error_stream(): - """test finish_reason='error' returns 500 InternalServerError (streaming)""" + """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -521,7 +521,8 @@ async def mock_generate(*args, **kwargs): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500.""" + so the API returns 400 instead of 500. + """ with pytest.raises(Exception, match="json_schema.*must be provided"): CompletionRequest( model=MODEL_NAME, @@ -584,7 +585,8 @@ def test_negative_prompt_token_ids_flat(): def test_logprobs_minus_one_allowed(): """logprobs=-1 means "return all logprobs". The sampling layer and the chat top_logprobs / prompt_logprobs validators all accept -1, so the completion - logprobs validator must accept it too instead of rejecting it as negative.""" + logprobs validator must accept it too instead of rejecting it as negative. + """ request = CompletionRequest( model=MODEL_NAME, prompt="Test prompt", @@ -744,7 +746,8 @@ def test_bounded_prompt_embeds_list_allowed(self, monkeypatch): @pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" + TypeError from the mode='before' comparison (which surfaces as HTTP 500). + """ with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): CompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/completion/test_token_in_token_out.py b/tests/entrypoints/openai/completion/test_token_in_token_out.py index 8882ae624428..45024d827740 100644 --- a/tests/entrypoints/openai/completion/test_token_in_token_out.py +++ b/tests/entrypoints/openai/completion/test_token_in_token_out.py @@ -39,8 +39,7 @@ def server(): @pytest.mark.asyncio async def test_token_in_token_out_and_logprobs(server): - """ - Test token-in-token-out and token_ids align with prompt_logprobs + """Test token-in-token-out and token_ids align with prompt_logprobs & logprobs when return_tokens_as_token_ids is enabled. """ tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py index aad1b5e0624e..c6ce3ebc2b21 100644 --- a/tests/entrypoints/openai/correctness/test_lmeval.py +++ b/tests/entrypoints/openai/correctness/test_lmeval.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file test accuracy of the vLLM server via LMEval. +"""This file test accuracy of the vLLM server via LMEval. It uses local-completions, which interacts with vLLM through the OAI API with N concurrent connections. This simulates real work usage of the API and makes @@ -37,7 +36,6 @@ def run_test(more_args): """Run the end to end accuracy test.""" - args = list(DEFAULT_ARGS) args.extend(more_args) print(f"Running with: {args}") @@ -68,7 +66,6 @@ def run_test(more_args): def test_lm_eval_accuracy_v1_engine(): """Run with the V1 Engine.""" - more_args = [] # Limit compilation time for V1 on TPU diff --git a/tests/entrypoints/openai/parser/test_harmony_render_parity.py b/tests/entrypoints/openai/parser/test_harmony_render_parity.py index 5cb446122361..ce6e4c80f89d 100644 --- a/tests/entrypoints/openai/parser/test_harmony_render_parity.py +++ b/tests/entrypoints/openai/parser/test_harmony_render_parity.py @@ -43,7 +43,8 @@ def _system(): class TestResponseInputToHarmonyRenderParity: """Each test drives the same conversation through both APIs and asserts - identical Harmony messages and rendered token sequences.""" + identical Harmony messages and rendered token sequences. + """ # ----------------------------------------------------------------------- # Single-message cases @@ -51,7 +52,8 @@ class TestResponseInputToHarmonyRenderParity: def test_developer_message(self): """Both APIs must render developer messages identically using - DeveloperContent (with the '# Instructions' header).""" + DeveloperContent (with the '# Instructions' header). + """ chat_msgs = parse_chat_input_to_harmony_message( {"role": "developer", "content": "Be concise."} ) @@ -246,7 +248,8 @@ def test_tool_output(self): def test_reasoning_combined_with_function_call(self): """Chat API packs reasoning + tool_calls into one dict; responses API represents them as two separate items. Both must produce the same two - Harmony messages in the same order: analysis then commentary.""" + Harmony messages in the same order: analysis then commentary. + """ chat_msgs = parse_chat_input_to_harmony_message( { "role": "assistant", @@ -529,7 +532,8 @@ def test_multi_turn_two_tool_calls_with_reasoning_between(self): def test_completed_turns_drop_reasoning(self): """Validates that reasoning from completed turns is dropped, while reasoning from the current in-progress tool-call turn is preserved - in both chat completions and responses previous_input_messages.""" + in both chat completions and responses previous_input_messages. + """ first_turn_reasoning = "FIRST_TURN_REASONING" second_turn_reasoning = "SECOND_TURN_REASONING" diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index 1b19714ac222..816e879342e4 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -170,7 +170,7 @@ class TestIsFunctionRecipientWithAllowedNames: """Tests for is_function_recipient with allowed_function_tool_names.""" def test_prefixed_always_accepted(self): - """functions. prefix is always accepted regardless of allowed names.""" + """Functions. prefix is always accepted regardless of allowed names.""" fn_names = frozenset({"other_tool"}) assert is_function_recipient("functions.get_weather", fn_names) is True @@ -252,8 +252,7 @@ def test_dotted_function_name_extraction(self, recipient, expected): class TestCommonParseInputToHarmonyMessage: - """ - Tests for scenarios that are common to both Chat Completion + """Tests for scenarios that are common to both Chat Completion parse_chat_input_to_harmony_message and Responses API response_previous_input_to_harmony functions. """ @@ -324,7 +323,8 @@ def test_assistant_message_with_empty_tool_call_arguments(self, parse_function): def test_system_message(self, parse_function): """Test parsing system messages, which are parsed into developer messages - with DeveloperContent.""" + with DeveloperContent. + """ chat_msg = { "role": "system", "content": "You are a helpful assistant", @@ -448,8 +448,7 @@ def test_array_content_with_missing_text(self, parse_function): class TestParseChatInputToHarmonyMessage: - """ - Tests for scenarios that are specific to the Chat Completion API + """Tests for scenarios that are specific to the Chat Completion API parse_chat_input_to_harmony_message function. """ diff --git a/tests/entrypoints/openai/responses/test_errors.py b/tests/entrypoints/openai/responses/test_errors.py index 0ea7833da9a7..1e0182ba6b8c 100644 --- a/tests/entrypoints/openai/responses/test_errors.py +++ b/tests/entrypoints/openai/responses/test_errors.py @@ -14,7 +14,7 @@ @pytest.mark.asyncio async def test_raise_if_error_raises_generation_error(): - """test _raise_if_error raises GenerationError""" + """Test _raise_if_error raises GenerationError.""" # create a minimal GenerateBaseServing instance mock_engine = MagicMock() mock_engine.model_config = MagicMock() @@ -42,7 +42,7 @@ async def test_raise_if_error_raises_generation_error(): @pytest.mark.asyncio async def test_convert_generation_error_to_streaming_response(): - """test _convert_generation_error_to_streaming_response output""" + """Test _convert_generation_error_to_streaming_response output.""" mock_engine = MagicMock() mock_engine.model_config = MagicMock() mock_engine.model_config.max_model_len = 100 diff --git a/tests/entrypoints/openai/responses/test_function_call.py b/tests/entrypoints/openai/responses/test_function_call.py index 9dcbd74c890b..006ee0f6f24e 100644 --- a/tests/entrypoints/openai/responses/test_function_call.py +++ b/tests/entrypoints/openai/responses/test_function_call.py @@ -168,8 +168,7 @@ async def test_max_tokens_with_tool_choice_required( @pytest.mark.asyncio async def test_named_tool_use(client: openai.AsyncOpenAI): def get_weather(latitude: float, longitude: float) -> str: - """ - Mock function to simulate getting weather data. + """Mock function to simulate getting weather data. In a real application, this would call an external weather API. """ return f"Current temperature at ({latitude}, {longitude}) is 20°C." diff --git a/tests/entrypoints/openai/responses/test_function_call_parsing.py b/tests/entrypoints/openai/responses/test_function_call_parsing.py index f90a641db7fd..c83fe0f61cce 100644 --- a/tests/entrypoints/openai/responses/test_function_call_parsing.py +++ b/tests/entrypoints/openai/responses/test_function_call_parsing.py @@ -12,7 +12,8 @@ def test_function_call_dict_converted_to_object(): """Test that function_call dictionaries are correctly parsed into - ResponseFunctionToolCall objects.""" + ResponseFunctionToolCall objects. + """ # Create a request with function_call as dict request_data = { "model": "gpt-oss", @@ -57,7 +58,6 @@ def test_direct_function_call_object_preservation(): def test_mixed_input_types_with_function_calls(): """Test parsing with mixed input types including function calls.""" - request_data = { "model": "gpt-oss", "input": [ @@ -177,7 +177,6 @@ def test_empty_list_input(): def test_function_call_output_not_affected(): """Test that FunctionCallOutput is not affected by the function_call parsing.""" - # Test with FunctionCallOutput as dict (should not be parsed) request_data = { "model": "gpt-oss", @@ -265,7 +264,6 @@ def test_function_call_validation_failure_logs_debug(caplog): def test_validator_handles_iterator_input(): """Test that validator can handle ValidatorIterator input (Pydantic internal).""" - # This test simulates when Pydantic passes a ValidatorIterator instead of a list # This happened with complex nested structures containing reasoning + function_call @@ -332,7 +330,8 @@ def test_validator_handles_empty_iterator(): def test_assistant_string_content_stays_easyinput(): """EasyInput assistant message with plain string content is not - coerced into a ResponseOutputMessage.""" + coerced into a ResponseOutputMessage. + """ request_data = { "model": "test-model", "input": [ @@ -354,7 +353,8 @@ def test_assistant_string_content_stays_easyinput(): def test_assistant_output_style_content_coerced(): """Assistant message whose content is output-message-shaped (list of - output_text items) should be coerced to ResponseOutputMessage.""" + output_text items) should be coerced to ResponseOutputMessage. + """ request_data = { "model": "test-model", "input": [ diff --git a/tests/entrypoints/openai/responses/test_harmony_utils.py b/tests/entrypoints/openai/responses/test_harmony_utils.py index 54c034de4420..b524842467ab 100644 --- a/tests/entrypoints/openai/responses/test_harmony_utils.py +++ b/tests/entrypoints/openai/responses/test_harmony_utils.py @@ -20,8 +20,7 @@ class TestResponsePreviousInputToHarmony: - """ - Tests for scenarios that are specific to the Responses API + """Tests for scenarios that are specific to the Responses API response_previous_input_to_harmony function. """ diff --git a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py index e2be8477e435..2485983bf3ae 100644 --- a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py +++ b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py @@ -67,7 +67,8 @@ def test_no_type_key_defaults_to_message_branch(self): def test_system_message(self): """System messages carry developer instructions and must be rendered - as developer messages with DeveloperContent.""" + as developer messages with DeveloperContent. + """ msg = response_input_to_harmony( {"type": "message", "role": "system", "content": "Be helpful."}, prev_responses=[], @@ -89,7 +90,8 @@ def test_assistant_message_gets_final_channel(self): def test_developer_message_gets_instructions_prefix(self): """Developer messages must use DeveloperContent which adds the - '# Instructions' header the model was trained on.""" + '# Instructions' header the model was trained on. + """ msg = response_input_to_harmony( {"type": "message", "role": "developer", "content": "Be concise."}, prev_responses=[], @@ -119,7 +121,8 @@ def test_message_with_array_content(self): def test_developer_message_array_content_concatenated(self): """Array content in developer messages is flattened and rendered - via DeveloperContent with the '# Instructions' header.""" + via DeveloperContent with the '# Instructions' header. + """ msg = response_input_to_harmony( { "type": "message", @@ -210,7 +213,8 @@ def test_function_call_output_channel_recipient_and_author_name(self): def test_function_call_output_uses_most_recent_matching_call(self): """When multiple prev_responses share a call_id, the last one wins - because the search is reversed.""" + because the search is reversed. + """ earlier = ResponseFunctionToolCall( id="fc_old", call_id="call_test", diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 967ddacda6de..1e59c8224429 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -839,11 +839,13 @@ def test_merge_chain_breaks(self, items, num_expected_messages): class TestConstructInputMessagesInstructionsLeak: """Regression tests for #37697: instructions from a prior response - should NOT leak through previous_response_id.""" + should NOT leak through previous_response_id. + """ def test_old_instructions_stripped_from_prev_msg(self): """System message in prev_msg must be dropped so the new request's - instructions are the only system message in the conversation.""" + instructions are the only system message in the conversation. + """ prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "What is 2+2?"}, @@ -860,7 +862,8 @@ def test_old_instructions_stripped_from_prev_msg(self): def test_no_instructions_in_new_request(self): """If the new request has no instructions, old ones should still - be stripped -- they must not carry over.""" + be stripped -- they must not carry over. + """ prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "Hi"}, diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index bd41b8cda282..1f9d95ea2e1c 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -81,7 +81,7 @@ def _new_online_renderer() -> OnlineRenderer: class MockConversationContext(ConversationContext): - """Mock conversation context for testing""" + """Mock conversation context for testing.""" def __init__(self): self.init_tool_sessions_called = False @@ -123,7 +123,7 @@ def test_serialize_message_pydantic_model_returns_dict() -> None: @pytest.fixture def mock_serving_responses(): - """Create a mock OpenAIServingResponses instance""" + """Create a mock OpenAIServingResponses instance.""" serving_responses = MagicMock(spec=OpenAIServingResponses) serving_responses.tool_server = MagicMock(spec=ToolServer) return serving_responses @@ -131,13 +131,13 @@ def mock_serving_responses(): @pytest.fixture def mock_context(): - """Create a mock conversation context""" + """Create a mock conversation context.""" return MockConversationContext() @pytest.fixture def mock_exit_stack(): - """Create a mock async exit stack""" + """Create a mock async exit stack.""" return MagicMock(spec=AsyncExitStack) @@ -550,7 +550,7 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs): class TestInitializeToolSessions: - """Test class for _initialize_tool_sessions method""" + """Test class for _initialize_tool_sessions method.""" @pytest.fixture(params=[ParsableContext, HarmonyContext]) def tool_context(self, request): @@ -635,7 +635,7 @@ async def new_session(name, *_args): @pytest_asyncio.fixture async def serving_responses_instance(self): - """Create a real OpenAIServingResponses instance for testing""" + """Create a real OpenAIServingResponses instance for testing.""" # Create minimal mocks for required dependencies engine_client = MagicMock() @@ -802,8 +802,7 @@ async def generate_output(): async def test_initialize_tool_sessions( self, serving_responses_instance, mock_context, mock_exit_stack ): - """Test that method works correctly with only MCP tools""" - + """Test that method works correctly with only MCP tools.""" request = ResponsesRequest(input="test input", tools=[]) # Call the method @@ -852,11 +851,11 @@ def test_validate_create_responses_input( class TestValidateGeneratorInput: - """Test class for _validate_generator_input method""" + """Test class for _validate_generator_input method.""" @pytest_asyncio.fixture async def serving_responses_instance(self): - """Create a real OpenAIServingResponses instance for testing""" + """Create a real OpenAIServingResponses instance for testing.""" # Create minimal mocks for required dependencies engine_client = MagicMock() @@ -884,7 +883,7 @@ async def serving_responses_instance(self): return instance def test_validate_generator_input(self, serving_responses_instance): - """Test _validate_generator_input with valid prompt length""" + """Test _validate_generator_input with valid prompt length.""" # Create an engine prompt with valid length (less than max_model_len) valid_prompt_token_ids = list(range(5)) # 5 tokens < 100 max_model_len engine_input = tokens_input(valid_prompt_token_ids) @@ -999,7 +998,7 @@ async def dummy_result_generator(): class TestExtractAllowedToolsFromMcpRequests: - """Test class for _extract_allowed_tools_from_mcp_requests function""" + """Test class for _extract_allowed_tools_from_mcp_requests function.""" def test_extract_allowed_tools_basic_formats(self): """Test extraction with list format, object format, and None.""" @@ -1113,7 +1112,7 @@ def _make_previous_item(*, channel, recipient, text="preamble text"): return item def test_preamble_delta_emits_text_events(self) -> None: - """commentary + recipient=None should emit output_text.delta events.""" + """Commentary + recipient=None should emit output_text.delta events.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_content_delta_events, ) @@ -1146,7 +1145,7 @@ def test_preamble_delta_second_token_no_added(self) -> None: assert "response.output_item.added" not in type_names def test_commentary_with_function_recipient_not_preamble(self) -> None: - """commentary + recipient='functions.X' must NOT use preamble path.""" + """Commentary + recipient='functions.X' must NOT use preamble path.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_content_delta_events, ) @@ -1164,7 +1163,8 @@ def test_commentary_with_function_recipient_not_preamble(self) -> None: def test_preamble_done_emits_text_done_events(self) -> None: """Completed preamble should emit text done + content_part done + - output_item done, same shape as final channel.""" + output_item done, same shape as final channel. + """ from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1184,8 +1184,9 @@ def test_preamble_done_emits_text_done_events(self) -> None: assert "response.output_item.done" in type_names def test_commentary_with_recipient_no_preamble_done(self) -> None: - """commentary + recipient='functions.X' should route to function call - done, not preamble done.""" + """Commentary + recipient='functions.X' should route to function call + done, not preamble done. + """ from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1331,7 +1332,8 @@ def mock_parse_delta(**kwargs): class TestStreamingReasoningToContentTransition: """Tests for _process_simple_streaming_events reasoning-to-content transition, specifically the fix for mixed deltas that carry both - reasoning and content simultaneously.""" + reasoning and content simultaneously. + """ @pytest.mark.asyncio async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta( @@ -1341,8 +1343,8 @@ async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta( and content set (e.g. reasoning end and content start in the same chunk), the trailing reasoning text must be emitted as a ResponseReasoningTextDeltaEvent and included in the - ResponseReasoningTextDoneEvent text.""" - + ResponseReasoningTextDoneEvent text. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1411,8 +1413,8 @@ async def test_transition_without_mixed_delta_no_extra_reasoning_event( self, monkeypatch ): """When the transition from reasoning to content is clean (no mixed - delta), no extra reasoning delta event should be emitted.""" - + delta), no extra reasoning delta event should be emitted. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1473,8 +1475,8 @@ async def result_generator(): async def test_reasoning_only_stream_no_content(self, monkeypatch): """When the stream has only reasoning deltas and no content, the reasoning done event should be emitted at finalization with the - full accumulated text, and no text delta events should appear.""" - + full accumulated text, and no text delta events should appear. + """ monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index ba55a58e7d3c..1880c06392d1 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -143,7 +143,8 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_streaming_logprobs(client: OpenAI, model_name: str): """Test that streaming with logprobs returns valid logprob data on - output_text.delta events and that top_logprobs has the requested count.""" + output_text.delta events and that top_logprobs has the requested count. + """ response = await client.responses.create( model=model_name, input="Say hello.", diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index e1a8bf46f640..2e9b47133dbb 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -124,7 +124,8 @@ def test_parallel_tool_calls(self): def test_split_name_and_args_same_index(self): """Regression: parsers like KimiK2 emit name and args as separate - DeltaToolCalls at the same index within one DeltaMessage.""" + DeltaToolCalls at the same index within one DeltaMessage. + """ tc_name = _make_tool_call(0, name="get_weather") tc_args = _make_tool_call(0, arguments='{"city":"SF"}') delta = DeltaMessage(tool_calls=[tc_name, tc_args]) @@ -140,7 +141,8 @@ def test_split_name_and_args_same_index(self): def test_reasoning_to_content_transition(self): """Regression: the old special case in emit_delta handled this; - now split_delta handles it generically.""" + now split_delta handles it generically. + """ processor = SimpleStreamingEventProcessor() _run_through_processor(processor, DeltaMessage(reasoning="think")) assert processor.state.current_state == _StateType.REASONING diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 6d3fc2f44746..18b01b85a8a3 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -83,8 +83,7 @@ def before_generate_case(context: schemathesis.HookContext, strategy): assert op is not None def no_invalid_types(case: schemathesis.Case): - """ - Skips tool_calls with `"type": "custom"` which schemathesis incorrectly + """Skips tool_calls with `"type": "custom"` which schemathesis incorrectly generates instead of the valid `"type": "function"`. Example test case that is skipped: diff --git a/tests/entrypoints/openai/test_render_token_offsets.py b/tests/entrypoints/openai/test_render_token_offsets.py index a2e66b7bd6ce..6ef66ee37bc0 100644 --- a/tests/entrypoints/openai/test_render_token_offsets.py +++ b/tests/entrypoints/openai/test_render_token_offsets.py @@ -25,7 +25,8 @@ def _model_config() -> Mock: def test_completion_flag_forwarded_to_tok_params(): """build_tok_params must forward return_token_offsets, defaulting to False (zero behavioral change for existing callers) and coercing JSON - null to False via the bool() guard.""" + null to False via the bool() guard. + """ cfg = _model_config() default = CompletionRequest(model="m", prompt="hi") @@ -40,7 +41,8 @@ def test_completion_flag_forwarded_to_tok_params(): def test_chat_flag_forwarded_to_tok_params(): """Chat build_tok_params has its own (max_completion_tokens) branch, so - its return_token_offsets forwarding is verified independently.""" + its return_token_offsets forwarding is verified independently. + """ cfg = _model_config() messages = [{"role": "user", "content": "hi"}] @@ -64,7 +66,8 @@ def test_generate_request_token_offsets_default_none(): def test_generate_request_token_offsets_survive_json_round_trip(): """GenerateRequest crosses the disagg serialization boundary; the - tuple[int, int] offsets must survive model_dump and re-validate.""" + tuple[int, int] offsets must survive model_dump and re-validate. + """ req = GenerateRequest( token_ids=[10, 20], sampling_params=SamplingParams(), diff --git a/tests/entrypoints/openai/test_return_token_ids.py b/tests/entrypoints/openai/test_return_token_ids.py index 8537082e3f8d..9fc8c102d4e4 100644 --- a/tests/entrypoints/openai/test_return_token_ids.py +++ b/tests/entrypoints/openai/test_return_token_ids.py @@ -189,8 +189,7 @@ async def test_chat_completion_with_tool_use(server): @pytest.mark.asyncio async def test_comparison_with_prompt_logprobs_and_logprobs(server): - """ - Test that token_ids align with prompt_logprobs and + """Test that token_ids align with prompt_logprobs and logprobs when return_tokens_as_token_ids is enabled. """ async with server.get_async_client() as client: diff --git a/tests/entrypoints/openai/test_stop_token_ids.py b/tests/entrypoints/openai/test_stop_token_ids.py index 5e9d79f68649..d6ac2cf1906f 100644 --- a/tests/entrypoints/openai/test_stop_token_ids.py +++ b/tests/entrypoints/openai/test_stop_token_ids.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for stop_token_ids propagation from default_sampling_params +"""Unit tests for stop_token_ids propagation from default_sampling_params to SamplingParams in ChatCompletionRequest and CompletionRequest. Regression test for https://github.com/vllm-project/vllm/issues/22519 diff --git a/tests/entrypoints/openai/utils.py b/tests/entrypoints/openai/utils.py index 4cc4800bf8dd..02a4c51f94e9 100644 --- a/tests/entrypoints/openai/utils.py +++ b/tests/entrypoints/openai/utils.py @@ -16,8 +16,7 @@ async def accumulate_streaming_response( stream_generator: AsyncGenerator[str, None], ) -> ChatCompletionResponse: - """ - Accumulate streaming SSE chunks into a complete ChatCompletionResponse. + """Accumulate streaming SSE chunks into a complete ChatCompletionResponse. This helper parses the SSE format and builds up the complete response by combining all the delta chunks. diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py index 88c9b434b38f..c620816dc56a 100644 --- a/tests/entrypoints/pooling/basic/test_tiling_engine.py +++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py @@ -30,8 +30,7 @@ def llm(vllm_runner): @pytest.mark.skip_global_cleanup def test_tiling_engine_basic(llm): - """ - Basic test with a small number of prompts (less than max_num_seqs). + """Basic test with a small number of prompts (less than max_num_seqs). No tiling should be triggered, but the engine still processes correctly. """ prompts = ["Hello", "World"] @@ -41,8 +40,7 @@ def test_tiling_engine_basic(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_many_requests(llm): - """ - Test with a large number of prompts that exceeds max_num_seqs. + """Test with a large number of prompts that exceeds max_num_seqs. This verifies that _run_tiling_engine correctly chunks requests, processes all of them, and returns outputs in the correct order. """ @@ -54,8 +52,7 @@ def test_tiling_engine_many_requests(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_with_pooling_params(llm): - """ - Test the tiling engine when different PoolingParams are provided. + """Test the tiling engine when different PoolingParams are provided. The engine must handle a list of params that matches the number of prompts. """ num_prompts = 10 @@ -77,8 +74,7 @@ def test_tiling_engine_with_pooling_params(llm): @pytest.mark.skip_global_cleanup def test_tiling_engine_abort_on_exception(llm): - """ - Test that abort_request IS called with the correct arguments when an + """Test that abort_request IS called with the correct arguments when an exception occurs inside the engine's step() loop. """ prompts = ["Prompt 0", "Prompt 1", "Prompt 2"] diff --git a/tests/entrypoints/pooling/embed/test_cohere_online.py b/tests/entrypoints/pooling/embed/test_cohere_online.py index 4964d99e0c66..29276cb0501f 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_online.py +++ b/tests/entrypoints/pooling/embed/test_cohere_online.py @@ -111,7 +111,8 @@ def test_basic_embed(server: RemoteOpenAIServer, model_name: str): def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str): """An input_type not defined in the model's prompt config should be - rejected with a 400 error.""" + rejected with a 400 error. + """ body = { "model": model_name, "input_type": "nonexistent_type", @@ -138,7 +139,8 @@ def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str): """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same - float embeddings for a generic model.""" + float embeddings for a generic model. + """ texts = ["hello world"] v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"]) v1 = _openai_embed(server, model_name, texts) diff --git a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py index 3e168adbdba9..5608cbcb4e06 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py +++ b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py @@ -76,7 +76,8 @@ def test_single_text_parity(server: RemoteOpenAIServer): def test_batch_parity(server: RemoteOpenAIServer): """A batch of texts should produce equivalent embeddings via both APIs, - in the same order.""" + in the same order. + """ texts = [ "machine learning", "deep learning", diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index f06e43d143ea..b002aae8db31 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -241,7 +241,8 @@ def test_truncate_none(self): def test_truncate_none_with_max_tokens(self): """truncate=NONE should NOT set truncate_prompt_tokens; the - max_tokens limit is enforced separately via _check_max_tokens.""" + max_tokens limit is enforced separately via _check_max_tokens. + """ req = self._make_request(truncate="NONE", max_tokens=10) tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req) assert tokens is None diff --git a/tests/entrypoints/pooling/embed/test_online_dimensions.py b/tests/entrypoints/pooling/embed/test_online_dimensions.py index 7abec99ec355..67c6d7b1345c 100644 --- a/tests/entrypoints/pooling/embed/test_online_dimensions.py +++ b/tests/entrypoints/pooling/embed/test_online_dimensions.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`. -""" +"""Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.""" import openai import pytest diff --git a/tests/entrypoints/pooling/embed/test_online_long_text.py b/tests/entrypoints/pooling/embed/test_online_long_text.py index 60cacb17f8fe..aeb82336fa48 100644 --- a/tests/entrypoints/pooling/embed/test_online_long_text.py +++ b/tests/entrypoints/pooling/embed/test_online_long_text.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test cases for long text embedding with automatic chunking mechanism. +"""Test cases for long text embedding with automatic chunking mechanism. This test suite validates vLLM's automatic chunking functionality for handling text inputs that exceed the model's maximum token length, specifically targeting @@ -241,8 +240,8 @@ async def test_long_text_embedding_1500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~1500 character long text - (~1028 tokens, exceeding 512 token limit).""" - + (~1028 tokens, exceeding 512 token limit). + """ # Verify text length # Verify text has sufficient word count (approximately 1500 words) word_count = len(LONG_TEXT_1500_WORDS.split()) @@ -290,8 +289,8 @@ async def test_long_text_embedding_2500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~2500 character long text - (~2048 tokens, requiring multiple chunks).""" - + (~2048 tokens, requiring multiple chunks). + """ # Verify text length # Verify text has sufficient word count (approximately 2500 words) word_count = len(LONG_TEXT_2500_WORDS.split()) @@ -339,7 +338,6 @@ async def test_batch_long_text_embedding( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test batch long text embedding processing.""" - input_texts = [ LONG_TEXT_1500_WORDS, LONG_TEXT_2500_WORDS, @@ -384,8 +382,8 @@ async def test_chunked_vs_normal_consistency( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test consistency between chunked and - normal processing (using short text).""" - + normal processing (using short text). + """ # Use a short text within the 512 token limit short_text = ( "Artificial intelligence technology is changing our world, " @@ -424,7 +422,6 @@ async def test_chunked_processing_response_format( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test response format and structure during chunked processing.""" - # Test with long text to trigger chunking embedding_response = await client_with_chunked_processing.embeddings.create( model=model_name, diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 91c853f91098..24d4f9e6d730 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for Cohere embed protocol: build_typed_embeddings and its -underlying packing helpers, plus Cohere-specific serving helpers.""" +underlying packing helpers, plus Cohere-specific serving helpers. +""" import struct diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index a30e48745018..db702b4b07be 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -454,7 +454,8 @@ async def test_derender_completion_empty_generate_responses(client): @pytest.mark.asyncio async def test_derender_completion_logprobs(client): """token_id:N placeholders in logprobs are resolved; CompletionLogProbs - flat-list structure is returned with non-empty tokens and text_offsets.""" + flat-list structure is returned with non-empty tokens and text_offsets. + """ gr1 = await _render_completion(client, "Hello world") ids1 = gr1["token_ids"][:3] token_id = ids1[0] @@ -555,7 +556,7 @@ async def test_derender_chat_oversized_token_ids_rejected(client): @pytest.mark.asyncio async def test_derender_chat_too_many_choices_rejected(client): - """choices count exceeding VLLM_MAX_N_SEQUENCES returns 400.""" + """Choices count exceeding VLLM_MAX_N_SEQUENCES returns 400.""" # Default VLLM_MAX_N_SEQUENCES is 16384; use a larger count. oversized_choices = [ {"index": i, "token_ids": [42], "finish_reason": "stop"} for i in range(20_000) diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index adece09c0a1a..6bcede31992d 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -121,7 +121,8 @@ async def _disagg( def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: """[(name, json normalized args)] so key ordering / whitespace don't - cause false negatives.""" + cause false negatives. + """ return [ (tc["function"]["name"], json.loads(tc["function"]["arguments"])) for tc in (response_choice["message"].get("tool_calls") or []) @@ -143,7 +144,8 @@ async def _run_parity_case( client: httpx.AsyncClient, messages: list[dict], **extra ) -> tuple[dict, dict]: """Run the coupled request then feed its generated tokens into the - disaggregated derender endpoint. Returns (coupled, disagg).""" + disaggregated derender endpoint. Returns (coupled, disagg). + """ coupled = await _coupled(client, messages, **extra) ch = coupled["choices"][0] chat_request = {"model": MODEL, "messages": messages, **extra} @@ -203,7 +205,8 @@ async def test_parity_tool_call(client): @pytest.mark.asyncio async def test_parity_reasoning_and_tool_call(client): """Combined reasoning + tool call parity means the highest drift risk - since it exercises both parser branches on the same output.""" + since it exercises both parser branches on the same output. + """ messages = [{"role": "user", "content": "What's the weather in Paris?"}] coupled, disagg = await _run_parity_case( client, diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index b783f40ae3c6..df10b31d2847 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -41,7 +41,6 @@ def _make_stream_chunk( usage: dict | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" - return GenerateStreamResponse( request_id=request_id, choices=[ @@ -61,7 +60,6 @@ def _make_usage_chunk( request_id: str = "test-req", ) -> GenerateStreamResponse: """Build a usage only final SSE chunk (empty choices).""" - return GenerateStreamResponse( request_id=request_id, choices=[], @@ -531,7 +529,6 @@ async def client(server): async def _render_chat(client) -> dict: """Render a minimal chat request and return the GenerateRequest dict.""" - resp = await client.post( "/v1/chat/completions/render", json={ diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 6aaf119f16a0..bb6a521ef14b 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -337,7 +337,6 @@ async def test_chat_completion_render_multi_turn(client): @pytest.mark.asyncio async def test_chat_completion_render_with_stream_true(client): """Render accepts stream params but still returns JSON (non-streamed).""" - response = await client.post( "/v1/chat/completions/render", json={ @@ -456,7 +455,8 @@ async def test_chat_completion_render_with_sampling_params(client): @pytest.mark.asyncio async def test_completion_render_emits_token_offsets(client): """With return_token_offsets, /v1/completions/render returns per-token - (start, end) char offsets aligned with token_ids.""" + (start, end) char offsets aligned with token_ids. + """ prompt = "Hello, world." response = await client.post( "/v1/completions/render", @@ -481,7 +481,8 @@ async def test_completion_render_emits_token_offsets(client): @pytest.mark.asyncio async def test_completion_render_default_no_token_offsets(client): """Without the flag, token_offsets must be null (existing responses - unchanged).""" + unchanged). + """ response = await client.post( "/v1/completions/render", json={ @@ -498,7 +499,8 @@ async def test_completion_render_default_no_token_offsets(client): @pytest.mark.asyncio async def test_chat_render_emits_token_offsets(client): """With return_token_offsets, /v1/chat/completions/render returns - per-token offsets relative to the templated prompt string.""" + per-token offsets relative to the templated prompt string. + """ response = await client.post( "/v1/chat/completions/render", json={ @@ -672,7 +674,8 @@ async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags client, ): """With a ``{% generation %}``-enabled template, the mask marks assistant - tokens and the masked tokens decode to the assistant content.""" + tokens and the masked tokens decode to the assistant content. + """ response = await client.post( "/v1/chat/completions/render", json={ diff --git a/tests/entrypoints/scale_out/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py index c6a6931948b9..eb45ce68ea9b 100644 --- a/tests/entrypoints/scale_out/render/test_render_multimodal.py +++ b/tests/entrypoints/scale_out/render/test_render_multimodal.py @@ -16,7 +16,6 @@ @pytest.fixture(scope="module") def vision_server(): """Vision-capable server used for multimodal /render tests.""" - args = [ "--enforce-eager", "--max-model-len", @@ -48,7 +47,6 @@ async def test_chat_completion_render_with_base64_image_url( local_asset_server, ): """Render a multimodal chat request and verify tokens are returned.""" - image = local_asset_server.get_image_asset("RGBA_comp.png") data_url = encode_image_url(image, format="PNG") @@ -112,7 +110,6 @@ async def test_tokenize_matches_render_for_multimodal_input( local_asset_server, ): """`/tokenize` should match `/v1/chat/completions/render` token output.""" - image = local_asset_server.get_image_asset("RGBA_comp.png") data_url = encode_image_url(image, format="PNG") diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index 764d684445d8..5cd60a959fb3 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Roundtrip tests for multimodal serde used by the +"""Roundtrip tests for multimodal serde used by the token_in_token_out generate endpoint. """ diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index e51530ce456f..1dcf93330edf 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -25,7 +25,8 @@ def _base_payload() -> dict: def test_omitted_max_tokens_is_not_provided(): """Body without ``max_tokens`` must surface as 'not provided' so the - server can apply its own default instead of the dataclass 16.""" + server can apply its own default instead of the dataclass 16. + """ req = GenerateRequest.model_validate(_base_payload()) # SamplingParams' dataclass default leaks through the parsed instance — # this is exactly the bug the server-side defaulting works around. @@ -35,7 +36,8 @@ def test_omitted_max_tokens_is_not_provided(): def test_explicit_max_tokens_is_provided(): """Even when the client picks the same value as the dataclass default, - it must register as explicitly set so the server won't override it.""" + it must register as explicitly set so the server won't override it. + """ payload = _base_payload() payload["sampling_params"] = {"max_tokens": 16} req = GenerateRequest.model_validate(payload) @@ -67,7 +69,8 @@ def test_json_roundtrip_preserves_provided_keys(): def test_internal_instance_construction_treats_all_as_provided(): """When internal callers build ``GenerateRequest`` from a pre-resolved ``SamplingParams`` instance, every field is considered explicitly set - so server-side defaulting can't clobber values resolved upstream.""" + so server-side defaulting can't clobber values resolved upstream. + """ sp = SamplingParams(max_tokens=500, temperature=0.0) req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp) assert req.is_sampling_param_provided("max_tokens") diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/serve/dev/rlhf/conftest.py index 1d2c598b1960..c636b5157291 100644 --- a/tests/entrypoints/serve/dev/rlhf/conftest.py +++ b/tests/entrypoints/serve/dev/rlhf/conftest.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared fixtures and helpers for the RL lifecycle test suite. +"""Shared fixtures and helpers for the RL lifecycle test suite. All test modules under this directory import from here to avoid duplication. @@ -92,6 +91,7 @@ def server( port: HTTP port to bind (caller is responsible for uniqueness). timeout: Seconds to wait for /health before giving up. dummy_weights: If True, use --load-format dummy (fast, no real weights). + """ env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"} base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS diff --git a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py index eb9aa7663c93..7cf6f4490c62 100644 --- a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py +++ b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py @@ -25,7 +25,7 @@ def echo_args_kwargs(self, *args, **kwargs) -> dict[str, Any]: ) def return_none(self, *args, **kwargs) -> None: - """Test method that does not return anything""" + """Test method that does not return anything.""" return @@ -48,7 +48,7 @@ def server(): def test_get_model_name(server): - """Test basic response""" + """Test basic response.""" response = requests.post( server.url_for("collective_rpc"), json={"method": "get_model_name"} ) @@ -59,7 +59,7 @@ def test_get_model_name(server): def test_return_none(server): - """Test return none""" + """Test return none.""" response = requests.post( server.url_for("collective_rpc"), json={"method": "return_none"} ) @@ -69,7 +69,7 @@ def test_return_none(server): def test_echo_args_kwargs(server): - """Test args, kwargs, and dict response""" + """Test args, kwargs, and dict response.""" args = ["arg1", "arg2"] kwargs = {"key1": "value1", "key2": "value2"} response = requests.post( diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py index 260dcd00bae9..5ebee506002a 100644 --- a/tests/entrypoints/serve/dev/test_sleep.py +++ b/tests/entrypoints/serve/dev/test_sleep.py @@ -87,8 +87,7 @@ def test_sleep_mode(): def _get_sleep_metrics_from_api(response: requests.Response): - """Return (awake, weights_offloaded, discard_all)""" - + """Return (awake, weights_offloaded, discard_all).""" awake, weights_offloaded, discard_all = None, None, None for family in text_string_to_metric_families(response.text): diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py index 426d1ede745d..cb14260ab82c 100644 --- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -22,7 +22,8 @@ def test_sanitize_message(): class TestSanitizeMessageFilePaths: """sanitize_message should also strip file paths and traceback - frames, not just memory addresses - see #31683.""" + frames, not just memory addresses - see #31683. + """ def test_strips_traceback_style_frame(self): msg = ( @@ -44,7 +45,8 @@ def test_strips_arbitrary_absolute_path(self): def test_strips_single_parent_container_path(self): """Regression: /app/server.py and /workspace/server.py (common in - container deployments) were missed by the original {2,} quantifier.""" + container deployments) were missed by the original {2,} quantifier. + """ assert "/app/" not in sanitize_message("Error in /app/server.py") assert "/workspace/" not in sanitize_message("Error in /workspace/server.py") diff --git a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py index 942dd95c182a..bd4e87a15543 100644 --- a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py +++ b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py @@ -24,7 +24,8 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: """Minimal stand-in for a FastAPI Request - just enough for the - handler to read req.app.state.args.log_error_stack.""" + handler to read req.app.state.args.log_error_stack. + """ return SimpleNamespace( app=SimpleNamespace( state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack)) @@ -35,7 +36,8 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: class TestValidationErrorParamFallback: """Ensure `param` falls back to the Pydantic error's `loc` when no - custom VLLMValidationError context is present.""" + custom VLLMValidationError context is present. + """ @pytest.mark.parametrize( ("error_type", "msg"), @@ -59,7 +61,8 @@ async def test_param_falls_back_to_loc(self, error_type: str, msg: str): async def test_param_fallback_does_not_crash_on_non_dict_error(self): """Schemathesis fuzzing found that errors[0] isn't always a dict. The fallback must not crash in that case - it should just leave - param as None instead of raising.""" + param as None instead of raising. + """ exc = RequestValidationError(["some unexpected non-dict error"]) response = await validation_exception_handler(_fake_request(), exc) @@ -97,7 +100,8 @@ def test_strips_internal_markers(self, loc, expected): def test_all_internal_falls_back_to_raw_join(self): """If every loc segment looks internal, fall back to the raw - dot-join rather than returning an empty string.""" + dot-join rather than returning an empty string. + """ loc = ("function-wrap[__log_extra_fields__()]",) assert clean_loc_for_param(loc) == "function-wrap[__log_extra_fields__()]" @@ -105,7 +109,8 @@ def test_all_internal_falls_back_to_raw_join(self): class TestValidationErrorDoesNotLeakServerPaths: """FastAPI sets endpoint_file/line/function/path on the exception during routing, not in the constructor - so we set them manually - here to reproduce the real leak. See #31683.""" + here to reproduce the real leak. See #31683. + """ @pytest.mark.asyncio async def test_handler_strips_endpoint_file_context(self): diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py index 73a97c4fa264..0711cbee957d 100644 --- a/tests/entrypoints/serve/instrumentator/test_basic.py +++ b/tests/entrypoints/serve/instrumentator/test_basic.py @@ -21,7 +21,7 @@ @pytest.fixture(scope="module") def server_args(request: pytest.FixtureRequest) -> list[str]: - """Provide extra arguments to the server via indirect parametrization + """Provide extra arguments to the server via indirect parametrization. Usage: diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index 297fe233b5fb..ac4576fbb95d 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -403,8 +403,7 @@ async def _poll_until( def _get_running_metrics_from_api(server: RemoteOpenAIServer): - """Return (running_count, waiting_count, kv_cache_usage)""" - + """Return (running_count, waiting_count, kv_cache_usage).""" response = requests.get(server.url_for("metrics")) assert response.status_code == HTTPStatus.OK diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index 8a860dc6251a..79356fc951b7 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -245,7 +245,8 @@ async def test_multiple_lora_adapters( client: openai.AsyncOpenAI, tmp_path, qwen3_lora_files ): """Validate that many loras can be dynamically registered and inferenced - with concurrently""" + with concurrently + """ # This test file configures the server with --max-cpu-loras=2 and this test # will concurrently load 10 adapters, so it should flex the LRU cache diff --git a/tests/entrypoints/serve/middleware/test_authentication_middleware.py b/tests/entrypoints/serve/middleware/test_authentication_middleware.py index 528261e644db..ee22407fdd0b 100644 --- a/tests/entrypoints/serve/middleware/test_authentication_middleware.py +++ b/tests/entrypoints/serve/middleware/test_authentication_middleware.py @@ -72,7 +72,8 @@ def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]): @pytest.fixture(params=get_args(SupportedTask)) def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: """For each supported task, build an app with only that task's routers, - extract all routes, and return the task name and routes.""" + extract all routes, and return the task name and routes. + """ task = request.param # Enable development mode to register all routes (including dev-only routes). monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") @@ -98,7 +99,8 @@ def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]] def test_auto_discovered_protected_routes_require_auth(task_routes): """For every auto-discovered route that starts with a guarded prefix, - verify that authentication is enforced.""" + verify that authentication is enforced. + """ task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) @@ -132,7 +134,8 @@ def test_auto_discovered_protected_routes_require_auth(task_routes): def test_auto_discovered_unprotected_routes_no_auth(task_routes): """For every auto-discovered route that does NOT start with a guarded - prefix, verify that no authentication is required.""" + prefix, verify that no authentication is required. + """ task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) diff --git a/tests/entrypoints/serve/middleware/test_optional_middleware.py b/tests/entrypoints/serve/middleware/test_optional_middleware.py index 865bd445c812..160f1e248dcc 100644 --- a/tests/entrypoints/serve/middleware/test_optional_middleware.py +++ b/tests/entrypoints/serve/middleware/test_optional_middleware.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for middleware that's off by default and can be toggled through +"""Tests for middleware that's off by default and can be toggled through server arguments, mainly --api-key and --enable-request-id-headers. """ diff --git a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py index 20b917400b3c..b6137dae93bd 100644 --- a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py +++ b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py @@ -114,7 +114,8 @@ def _clear_env_vars(self): def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path): """Test customer scenario: script functions automatically override - framework defaults.""" + framework defaults. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -175,7 +176,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_customer_decorator_usage(self, monkeypatch, tmp_path): """Test customer scenario: using @custom_ping_handler and - @custom_invocation_handler decorators.""" + @custom_invocation_handler decorators. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -231,7 +233,8 @@ async def my_invoke(request: Request): def test_handler_priority_order(self, monkeypatch, tmp_path): """Test priority: @custom_ping_handler/@custom_invocation_handler - decorators vs script functions.""" + decorators vs script functions. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -304,7 +307,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_environment_variable_script_loading(self, monkeypatch, tmp_path): """Test that environment variables correctly specify script location - and loading.""" + and loading. + """ try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -416,7 +420,8 @@ async def test_framework_default_handlers(self): def test_handler_env_var_override(self, monkeypatch, tmp_path): """Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER - environment variable overrides.""" + environment variable overrides. + """ try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, @@ -510,7 +515,8 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path): """Test that environment variables have highest priority over decorators - and script functions for both ping and invocation handlers.""" + and script functions for both ping and invocation handlers. + """ try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, diff --git a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py index 6b226c6999ef..c4317973cb23 100644 --- a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py +++ b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test: ``/tokenize`` must expand image placeholders for VLM models. +"""Regression test: ``/tokenize`` must expand image placeholders for VLM models. Fixed by PR #34560 ("Move InputPreprocessor into Renderer (2/2)"). Before that change, ``/tokenize`` returned ~26 tokens for a message with an diff --git a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py index 50083e345538..49e181e36e3b 100644 --- a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py +++ b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test: calling ``/tokenize`` with multimodal data followed by +"""Regression test: calling ``/tokenize`` with multimodal data followed by ``/v1/chat/completions`` with the same data must not cause an error. Ensures that the ``/tokenize`` endpoint does not pollute internal caches @@ -52,8 +51,8 @@ async def test_tokenize_then_chat_completion_with_image( local_asset_server, ): """Tokenize a multimodal message, then send the same message to chat - completions. The chat completion must succeed (not 500).""" - + completions. The chat completion must succeed (not 500). + """ image_url = local_asset_server.url_for("stop_sign.jpg") messages = [ { diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index a41f39187f4f..051978b60c47 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -36,11 +36,13 @@ def test_should_include_usage_force_enables_continuous_usage(stream_options, exp class TestGetMaxTokens: """Tests for get_max_tokens() to ensure generation_config's max_tokens acts as a default when from model author, and as a ceiling when - explicitly set by the user.""" + explicitly set by the user. + """ def test_default_sampling_params_used_when_no_request_max_tokens(self): """When user doesn't specify max_tokens, generation_config default - should apply.""" + should apply. + """ result = get_max_tokens( max_model_len=24000, max_tokens=None, @@ -51,7 +53,8 @@ def test_default_sampling_params_used_when_no_request_max_tokens(self): def test_request_max_tokens_not_capped_by_default_sampling_params(self): """When user specifies max_tokens in request, model author's - generation_config max_tokens must NOT cap it (fixes #34005).""" + generation_config max_tokens must NOT cap it (fixes #34005). + """ result = get_max_tokens( max_model_len=24000, max_tokens=5000, @@ -94,7 +97,8 @@ def test_max_model_len_still_caps_output(self): def test_request_max_tokens_smaller_than_default(self): """When user explicitly requests fewer tokens than gen_config default, - that should be respected.""" + that should be respected. + """ result = get_max_tokens( max_model_len=24000, max_tokens=512, diff --git a/tests/entrypoints/serve/utils/test_request_logger.py b/tests/entrypoints/serve/utils/test_request_logger.py index 717ff190b3a7..5df3c3cd8425 100644 --- a/tests/entrypoints/serve/utils/test_request_logger.py +++ b/tests/entrypoints/serve/utils/test_request_logger.py @@ -239,7 +239,8 @@ def test_request_logger_log_outputs_integration(): def test_streaming_complete_logs_full_text_content(): """Test that streaming complete logging includes - full accumulated text, not just token count.""" + full accumulated text, not just token count. + """ mock_logger = MagicMock() with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger): diff --git a/tests/entrypoints/serve/utils/test_sse_keep_alive.py b/tests/entrypoints/serve/utils/test_sse_keep_alive.py index 519b9bb4a129..6c61080f8f51 100644 --- a/tests/entrypoints/serve/utils/test_sse_keep_alive.py +++ b/tests/entrypoints/serve/utils/test_sse_keep_alive.py @@ -17,7 +17,8 @@ def test_disabled_returns_same_generator(): """For 0, negative, and NaN intervals the exact same generator object must be returned (identity), proving no wrapper or per-chunk overhead when - disabled.""" + disabled. + """ async def gen(): yield "x" @@ -33,7 +34,8 @@ async def gen(): @pytest.mark.asyncio async def test_keep_alive_during_silent_start(): """A silent start (queue wait / prefill) emits keep-alive comments before - the first real chunk.""" + the first real chunk. + """ release = asyncio.Event() async def gen(): @@ -55,7 +57,8 @@ async def gen(): @pytest.mark.asyncio async def test_no_keep_alive_when_streaming_is_fast(): """If chunks arrive faster than the interval, no comment is emitted and all - data is forwarded.""" + data is forwarded. + """ async def gen(): for chunk in ["a", "b", "c"]: @@ -69,7 +72,8 @@ async def gen(): @pytest.mark.asyncio async def test_keep_alive_between_chunks_does_not_drop_data(): """Keep-alives interleaved with real chunks must never drop or reorder - data.""" + data. + """ release_b = asyncio.Event() release_c = asyncio.Event() @@ -135,7 +139,8 @@ async def gen(): @pytest.mark.asyncio async def test_aclose_closes_upstream(): """Closing a wrapper that is mid-stream must complete promptly and close - the upstream generator (regression for the earlier leaked-task hang).""" + the upstream generator (regression for the earlier leaked-task hang). + """ finalized = asyncio.Event() async def gen(): @@ -154,7 +159,8 @@ async def gen(): @pytest.mark.asyncio async def test_client_disconnect_cancels_upstream(): """Cancelling the consumer mid-stream must cancel the pending anext and - close the upstream generator so its cleanup runs.""" + close the upstream generator so its cleanup runs. + """ finalized = asyncio.Event() async def gen(): @@ -181,7 +187,8 @@ async def gen(): @pytest.mark.asyncio async def test_streaming_response_emits_comment(): """A keep-alive comment must reach the client through the real - StreamingResponse used by the routers.""" + StreamingResponse used by the routers. + """ release = asyncio.Event() async def gen(): @@ -220,7 +227,8 @@ async def receive(): async def test_streaming_response_disconnect_closes_upstream(): """On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette cancels streaming and the wrapper's finally must close the upstream - generator.""" + generator. + """ finalized = asyncio.Event() async def gen(): diff --git a/tests/entrypoints/speech_to_text/conftest.py b/tests/entrypoints/speech_to_text/conftest.py index 94505bd054f4..bf65a72ab57d 100644 --- a/tests/entrypoints/speech_to_text/conftest.py +++ b/tests/entrypoints/speech_to_text/conftest.py @@ -11,6 +11,7 @@ def add_attention_backend(server_args, attention_config): Args: server_args: List of server arguments to extend in-place. attention_config: Dict with 'backend' key, or None. + """ if attention_config and "backend" in attention_config: server_args.extend(["--attention-backend", attention_config["backend"]]) diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index 713fa48eee48..855dafaaba09 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Evaluate Transcription API correctness by computing Word Error Rate (WER) +"""Evaluate Transcription API correctness by computing Word Error Rate (WER) on a given ASR dataset. When provided, it will also compare the WER against a baseline. This simulates real work usage of the API and makes sure that the frontend and diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index 7fb7a5522d29..b9ea18bbaea9 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -295,7 +295,8 @@ async def test_session_update_invalid_model_returns_error(model_name): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_commit_without_session_update_returns_error(model_name): """Test that committing before validating the model returns an error - and does not fall through to processing.""" + and does not fall through to processing. + """ server_args = ["--enforce-eager", "--max-model-len", "2048"] if model_name.startswith("mistralai"): diff --git a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py index cd61507728a6..5534845326ab 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py +++ b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py @@ -18,9 +18,8 @@ @pytest.mark.asyncio async def test_chunk_offsets_are_cumulative_not_nominal(): - """ - chunk_start_offsets must be cumulative actual chunk lengths, not the old approach of - 'idx * max_audio_clip_s'. When split_audio places a boundary before the + """chunk_start_offsets must be cumulative actual chunk lengths, not the old approach + of 'idx * max_audio_clip_s'. When split_audio places a boundary before the nominal 30 s mark, the old formula drifts; the fixed formula stays exact. """ # Chunks shorter than exactly 30 s, as split_audio produces when a quiet diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index 51599d359e8f..d5658e7cf0b4 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -238,8 +238,7 @@ async def test_stream_options(winning_call, whisper_client): @pytest.mark.asyncio async def test_sampling_params(mary_had_lamb, whisper_client): - """ - Compare sampling with params and greedy sampling to assert results + """Compare sampling with params and greedy sampling to assert results are different when extreme sampling parameters values are picked. """ transcription = await whisper_client.audio.transcriptions.create( diff --git a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py index f03b6e8a3818..79b46bea51ab 100644 --- a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py @@ -150,7 +150,6 @@ async def test_non_streaming_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call in non-streaming mode.""" - response = await client.chat.completions.create( model=server_config["model_arg"], messages=MESSAGES, @@ -183,7 +182,6 @@ async def test_streaming_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call in streaming mode.""" - stream = await client.chat.completions.create( model=server_config["model_arg"], messages=MESSAGES, @@ -230,7 +228,6 @@ async def test_non_streaming_product_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call integer and boolean parameters in non-streaming mode.""" - response = await client.chat.completions.create( model=server_config["model_arg"], messages=PRODUCT_MESSAGES, @@ -272,7 +269,6 @@ async def test_streaming_product_tool_call( client: openai.AsyncOpenAI, server_config: ServerConfig ): """Test tool call integer and boolean parameters in streaming mode.""" - stream = await client.chat.completions.create( model=server_config["model_arg"], messages=PRODUCT_MESSAGES, diff --git a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py index d99b66d9ac6b..de402b435c23 100644 --- a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py @@ -127,8 +127,7 @@ async def client(server): # Utility to extract reasoning and tool calls # ========================================================== def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str]]: - """ - Extract accumulated reasoning text and tool call arguments + """Extract accumulated reasoning text and tool call arguments from streaming chunks. """ reasoning: str = "" @@ -165,7 +164,6 @@ def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str] @pytest.mark.asyncio async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOpenAI): """Verify calculator tool call is made and arguments are accurate.""" - response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES_CALC, @@ -205,7 +203,6 @@ async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOp @pytest.mark.asyncio async def test_streaming_tool_call_get_time_with_reasoning(client: openai.AsyncOpenAI): """Verify streamed reasoning and tool call behavior for get_time.""" - stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES_GET_TIME, @@ -258,8 +255,7 @@ async def test_streaming_multiple_tools(client: openai.AsyncOpenAI): @pytest.mark.asyncio async def test_invalid_tool_call(client: openai.AsyncOpenAI): - """ - Verify that ambiguous instructions that should not trigger a tool + """Verify that ambiguous instructions that should not trigger a tool do not produce any tool calls. """ response = await client.chat.completions.create( @@ -283,8 +279,7 @@ async def test_invalid_tool_call(client: openai.AsyncOpenAI): @pytest.mark.asyncio async def test_tool_call_with_temperature(client: openai.AsyncOpenAI): - """ - Verify model produces valid tool or text output + """Verify model produces valid tool or text output under non-deterministic sampling. """ response = await client.chat.completions.create( diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py index 630f641df59d..ab57140c2b07 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/entrypoints/unit_tests/test_chat_utils.py @@ -940,7 +940,6 @@ def test_parse_chat_messages_audio_embeds_with_string( audio_embeds_model_config, ): """Test audio_embeds with base64 string embedding data.""" - import torch # Create a sample audio embedding tensor @@ -983,7 +982,6 @@ async def test_parse_chat_messages_audio_embeds_async( audio_embeds_model_config, ): """Test audio_embeds with async futures.""" - import torch # Create a sample audio embedding tensor @@ -2986,7 +2984,6 @@ def test_postprocess_messages_null_arguments_string(): @pytest.mark.asyncio async def test_resolve_items_runs_modalities_concurrently_and_preserves_order(): """Media fetches overlap while modality and item order are preserved.""" - active_fetches = 0 max_active_fetches = 0 diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index 1c1f6ed23593..e09687425bb9 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -51,7 +51,8 @@ async def generate_mock_outputs( num_turns, prompt_token_counts, output_token_counts, cached_token_counts=None ): """Generate a sequence of mock RequestOutput objects to simulate multiple - turns.""" + turns. + """ if cached_token_counts is None: cached_token_counts = [0] * num_turns @@ -270,7 +271,8 @@ def test_reasoning_tokens_counting(): def test_preamble_tokens_not_counted_as_reasoning(): """Preambles (commentary with no recipient) are visible user text, - not hidden reasoning. They must NOT inflate num_reasoning_tokens.""" + not hidden reasoning. They must NOT inflate num_reasoning_tokens. + """ context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=0) @@ -287,7 +289,8 @@ def test_preamble_tokens_not_counted_as_reasoning(): def test_commentary_with_recipient_counted_as_reasoning(): """Commentary directed at a tool (recipient != None) is hidden from - the user, so it should still count as reasoning tokens.""" + the user, so it should still count as reasoning tokens. + """ context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=3) @@ -346,7 +349,8 @@ async def test_single_turn_no_tool_output(): @pytest.mark.asyncio async def test_negative_tool_tokens_edge_case(): """Test edge case where calculation could result in negative tool - tokens. We should log an error and clamp the value to 0.""" + tokens. We should log an error and clamp the value to 0. + """ # Use patch to check if logger.error was called with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log: context, _ = make_harmony_context(available_tools=["browser"]) @@ -544,7 +548,6 @@ async def test_streaming_multi_turn_token_counting(): @pytest.mark.asyncio async def test_streaming_message_synchronization(): """Completed messages from append-local and flush segments sync into context.""" - # Create a streaming context with some initial messages initial_messages = [ Message( @@ -687,7 +690,8 @@ def create_simple_context_output( finished=True, ): """Helper to create a RequestOutput with customizable text for - SimpleContext tests.""" + SimpleContext tests. + """ if token_ids is None: token_ids = [] return RequestOutput( @@ -871,7 +875,8 @@ def test_simple_context_final_output(): def test_simple_context_output_messages_empty_text_with_tokens(): """output_messages should be returned when tokens exist even if text is - empty (e.g. special tokens).""" + empty (e.g. special tokens). + """ context = SimpleContext() context.append_output( create_simple_context_output( @@ -889,7 +894,8 @@ def test_simple_context_output_messages_empty_text_with_tokens(): def test_simple_context_output_messages_no_mutation(): """Each call to output_messages returns a fresh list; callers can't - corrupt internal state.""" + corrupt internal state. + """ context = SimpleContext() context.append_output( create_simple_context_output( diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index e5c3616d6e65..9b36ab398d3d 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -121,7 +121,8 @@ def test_get_world_size_tp1(vllm_runner): @create_new_process_for_each_test() def test_init_weight_transfer_engine_calls_engine(vllm_runner): """Test that init_weight_transfer_engine calls the engine's - init_transfer_engine method.""" + init_transfer_engine method. + """ if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") diff --git a/tests/evals/gpt_oss/conftest.py b/tests/evals/gpt_oss/conftest.py index d35dec4831a3..28fbb88d527e 100644 --- a/tests/evals/gpt_oss/conftest.py +++ b/tests/evals/gpt_oss/conftest.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Pytest configuration for GPT-OSS evaluation tests. -""" +"""Pytest configuration for GPT-OSS evaluation tests.""" from pathlib import Path diff --git a/tests/evals/gpt_oss/test_gpqa_correctness.py b/tests/evals/gpt_oss/test_gpqa_correctness.py index 63188ec40767..99d799698a37 100644 --- a/tests/evals/gpt_oss/test_gpqa_correctness.py +++ b/tests/evals/gpt_oss/test_gpqa_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GPQA evaluation using vLLM server and GPT-OSS evaluation package. +"""GPQA evaluation using vLLM server and GPT-OSS evaluation package. Usage: pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \ @@ -48,7 +47,6 @@ def ensure_tiktoken_files(): def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float: """Run GPQA evaluation using the gpt-oss evaluation package.""" - # Build the command to run the evaluation cmd = [ sys.executable, diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 83e54afd9e7f..7e2534b55a95 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -1,9 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Isolated GSM8K evaluation script for vLLM serve endpoint. -""" +"""Isolated GSM8K evaluation script for vLLM serve endpoint.""" import argparse import ast @@ -45,7 +43,7 @@ def download_and_cache_file(url: str, filename: str | None = None) -> str: def load_gsm8k_data() -> tuple[list[dict], list[dict]]: - """Load GSM8K train and test data""" + """Load GSM8K train and test data.""" train_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/train.jsonl" test_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/test.jsonl" @@ -91,6 +89,7 @@ async def call_vllm_api( Returns: Tuple of (response_text, completion_tokens) + """ data = { "prompt": prompt, @@ -221,8 +220,7 @@ def evaluate_gsm8k( gen_prefix: str = "", max_concurrency: int | None = None, ) -> dict[str, float | int]: - """ - Evaluate GSM8K accuracy using vLLM serve endpoint. + """Evaluate GSM8K accuracy using vLLM serve endpoint. Returns dict with accuracy, invalid_rate, latency, etc. """ diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index 5a237cf94764..477a69a16c67 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GSM8K evaluation using vLLM server and isolated GSM8K script. +"""GSM8K evaluation using vLLM server and isolated GSM8K script. Replacement for lm-eval-harness with better performance and control. Usage: diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index c140a41e2644..bcc001e2bd43 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GSM8K correctness test for CPU KV offloading connectors. +"""GSM8K correctness test for CPU KV offloading connectors. Regression guard for stride computation bugs in the offloading worker (e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV diff --git a/tests/evals/mrcr/test_mrcr_correctness.py b/tests/evals/mrcr/test_mrcr_correctness.py index 3adfd4dc8799..f381a592dde6 100644 --- a/tests/evals/mrcr/test_mrcr_correctness.py +++ b/tests/evals/mrcr/test_mrcr_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -MRCR long-context accuracy test. +"""MRCR long-context accuracy test. Usage: pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \ diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index ab12393692a5..ec9bb6110dcd 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -78,7 +78,8 @@ def _all_kernel_classes() -> list[type]: def _probe(cls: type): """A bare kernel instance with a plausible config, so input_quant_key() - can be queried without the hardware-gated constructor.""" + can be queried without the hardware-gated constructor. + """ obj = cls.__new__(cls) # type: ignore[call-overload] if issubclass(cls, NvFp4LinearKernel): obj.config = NvFp4LinearLayerConfig() diff --git a/tests/ir/ir_test_utils.py b/tests/ir/ir_test_utils.py index a82206b6f72e..5cd05b9bcf09 100644 --- a/tests/ir/ir_test_utils.py +++ b/tests/ir/ir_test_utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared test utilities for vLLM IR op correctness tests. -""" +"""Shared test utilities for vLLM IR op correctness tests.""" import torch diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py index ca8185b2d3dc..5033b0ec10f3 100644 --- a/tests/ir/test_op.py +++ b/tests/ir/test_op.py @@ -34,7 +34,7 @@ def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: @_custom_add.register_impl("impl_b", inplace=True) def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Computes x+y+20""" + """Computes x+y+20.""" x.add_(y) x.add_(20) return x @@ -101,7 +101,7 @@ def _custom_kwarg_op( class TestIrOpCustomAdd: # Registration invariants def test_decorated_object(self, custom_add_op): - """Make sure that referring directly to an op is correct""" + """Make sure that referring directly to an op is correct.""" _custom_add = custom_add_op assert isinstance(_custom_add, IrOp) assert "_custom_add" in IrOp.registry diff --git a/tests/kernels/attention/test_amx_mla.py b/tests/kernels/attention/test_amx_mla.py index f5fa2fad16c2..7dab61f1c4d1 100644 --- a/tests/kernels/attention/test_amx_mla.py +++ b/tests/kernels/attention/test_amx_mla.py @@ -247,7 +247,8 @@ def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens): class _FakeLinear: """Minimal duck-typed stand-in for kv_b_proj: only .weight and - .quant_method are read by get_and_maybe_dequant_weights().""" + .quant_method are read by get_and_maybe_dequant_weights(). + """ def __init__(self, weight: torch.Tensor): self.weight = weight @@ -318,7 +319,8 @@ def __init__(self, decode=None, prefill=None, max_seq_len=0): def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl would produce via layer.W_UK_T) and must reproduce plain per-head MLA - decode attention through W_UK/W_UV.""" + decode attention through W_UK/W_UV. + """ set_random_seed(1) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( @@ -384,7 +386,8 @@ def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config): """forward_mha receives raw unabsorbed Q and must attend correctly for a prefill batch with a mix of fresh and continued (cached-prefix) - sequences.""" + sequences. + """ set_random_seed(2) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 53867a5edf93..0511fe7b07f4 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -616,7 +616,8 @@ def test_non_causal_autoselect_backend(): ) def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str): """FlashAttentionBackend must not claim support for kv_cache dtypes - that it cannot handle.""" + that it cannot handle. + """ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype) @@ -627,7 +628,8 @@ def test_flash_attn_accepts_handled_fp8_variants( kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch ): """FlashAttentionBackend must accept the two fp8 dtypes it can actually - handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'.""" + handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'. + """ import vllm.v1.attention.backends.fa_utils as fa_utils_mod from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend diff --git a/tests/kernels/attention/test_cutlass_mla_decode.py b/tests/kernels/attention/test_cutlass_mla_decode.py index 7a5011b5caec..d56d72adba79 100644 --- a/tests/kernels/attention/test_cutlass_mla_decode.py +++ b/tests/kernels/attention/test_cutlass_mla_decode.py @@ -225,7 +225,8 @@ def test_cutlass_mla_decode_cross_layer_view(): """The kernel must read the cache's page-dim stride instead of assuming pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must - match a contiguous cache holding the same data exactly.""" + match a contiguous cache holding the same data exactly. + """ device = torch.device("cuda:0") torch.set_default_dtype(torch.bfloat16) torch.set_default_device(device) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 109f59bc8508..b786fac095d8 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -96,6 +96,7 @@ def _make_paged_kv_metadata( kv_indices – CUDA int32, shape [total_blocks] kv_last_page_lens – CPU int32, shape [num_seqs] block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq] + """ num_seqs = len(kv_lens) max_blocks = (max(kv_lens) + block_size - 1) // block_size @@ -161,7 +162,8 @@ def test_fast_decode_plan_importable() -> None: @torch.inference_mode def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: """On the first call fast_plan_decode must route through self.plan() and - flip vllm_first_call to False on the wrapper object.""" + flip vllm_first_call to False on the wrapper object. + """ from unittest.mock import patch from vllm.v1.attention.backends.flashinfer import fast_plan_decode diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index a49ae3ae1ad3..0d3791964398 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -286,7 +286,7 @@ def run( @requires_sm10x def test_flashinfer_mla_decode_workspace_supports_autotune(): - """vLLM's FlashInfer MLA decode workspace must be int8 for autotuning. + """VLLM's FlashInfer MLA decode workspace must be int8 for autotuning. Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py index 48eb58b5e1d9..b35005d36112 100644 --- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py +++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py @@ -134,7 +134,8 @@ def make_quantized_kv_cache( head_size: int, ) -> tuple: """Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf, - kv_scale, ref_kv_cache, is_nvfp4_kv).""" + kv_scale, ref_kv_cache, is_nvfp4_kv). + """ is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE if is_nvfp4_kv: data, scales, kv_scale, ref = make_nvfp4_kv_cache( diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py index aa7c836666f7..e86686ee8668 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -73,7 +73,8 @@ def _strided_context_inputs( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The real prefill-context layouts: ``k_nope``/``v`` as the two strided halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather - workspace (left in the fp8 cache layout for a plain fp8 cache).""" + workspace (left in the fp8 cache layout for a plain fp8 cache). + """ kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype) k_nope, v = kv_nope.split((128, 128), dim=-1) workspace = _randn(num_tokens, 576, dtype=dtype) diff --git a/tests/kernels/attention/test_lightning_attn.py b/tests/kernels/attention/test_lightning_attn.py index 61e13166808c..d4b1a359d28e 100644 --- a/tests/kernels/attention/test_lightning_attn.py +++ b/tests/kernels/attention/test_lightning_attn.py @@ -23,7 +23,7 @@ def reference_lightning_attention(q, k, v, ed, block_size, kv_history): - """Reference implementation of lightning attention core algorithm + """Reference implementation of lightning attention core algorithm. The difference from the main implementation is that this processes each step sequentially, instead of using parallelized triton kernels @@ -72,7 +72,7 @@ def reference_lightning_attention(q, k, v, ed, block_size, kv_history): def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx): - """Reference implementation: linear attention decode function""" + """Reference implementation: linear attention decode function.""" B, H, _, D = q.shape output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device) diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 1b89e4db66b3..07295ed0b871 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -89,7 +89,8 @@ def merge_attn_states_torch( @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16]) def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: """When a token is empty on both sides (both LSE -inf), the 0/0 softmax - scales must not surface as NaN in the merged output.""" + scales must not surface as NaN in the merged output. + """ num_tokens, num_heads, head_size = 6, 8, 128 prefix_output = torch.zeros( num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index b6578fa2f915..8fa4963a6a19 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test: +"""Test: * Tests for MMEncoderAttention layer """ @@ -49,9 +48,7 @@ def reset_test_state(): @pytest.mark.parametrize("device", devices) def test_mha_attn_platform(default_vllm_config, device: str): - """ - Test the attention selector between different platform and device. - """ + """Test the attention selector between different platform and device.""" torch.set_default_dtype(torch.float16) if device == "cpu": @@ -113,8 +110,7 @@ def ref_attention( value: torch.Tensor, scale: float, ) -> torch.Tensor: - """ - Native implementation of scaled dot product attention without mask: + """Native implementation of scaled dot product attention without mask: - query, key, value: [batch_size, seq_len, num_heads, head_size] - attn_mask: [batch_size, seq_len, seq_len] """ diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index a7e6d499c176..c7b763e7c7e1 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -45,7 +45,8 @@ def kv_layout(request) -> KVCacheLayout: def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]: """Per-layer physical stride order for the given layout; the 3-dim - indexer side cache (H=1) is contiguous, so identity.""" + indexer side cache (H=1) is contiguous, so identity. + """ if ndim == 3: return (0, 1, 2) stride_order = layout.layer_view_order @@ -65,7 +66,8 @@ def _main_spec() -> FullAttentionSpec: def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]: """Standardized per-layer logical shape [B, H, N, C] for the main cache, - derived the same way the production allocator does.""" + derived the same way the production allocator does. + """ shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages) return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize) @@ -75,7 +77,8 @@ def _allocate_main_kv_via_contract( ) -> torch.Tensor: """Build the main KV cache exactly as the production allocator does for the given layout: allocate the physical (permuted) tensor, then expose the - inverse-permuted logical [B, H, N, C] view the backend sees.""" + inverse-permuted logical [B, H, N, C] view the backend sees. + """ logical_shape = _main_kv_logical_shape(num_pages) stride_order = _layer_stride_order(layout, len(logical_shape)) physical_shape = tuple(logical_shape[i] for i in stride_order) @@ -1376,7 +1379,8 @@ def test_amd_decode_index_topk_end_to_end( @pytest.mark.parametrize("num_idx_heads", [1, 4]) def test_decode_index_topk_fp8(num_idx_heads: int): """The standalone Triton path must score FP8 inputs in FP32 so its top-k - matches a reference computed from the dequantized FP8 values.""" + matches a reference computed from the dequantized FP8 values. + """ torch.manual_seed(0) topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128 decode_query_len = 1 @@ -1649,7 +1653,8 @@ def test_prefill_sparse_attention_correctness( def test_main_cache_layout_contract(): """The standardized per-layer logical shape is [B, H, N, C] with packed K/V content, and the legacy layout aliases resolve to the expected - per-layer stride orders.""" + per-layer stride orders. + """ nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) assert logical == (nb, h, bs, 2 * d) @@ -1773,7 +1778,8 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): """The indexer side cache is standardized 4D with H=1: under both layouts the allocator's logical view stays contiguous and squeezes (as `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim - [num_blocks, block_size, head_dim] cache the kernels consume.""" + [num_blocks, block_size, head_dim] cache the kernels consume. + """ nb = 5 ispec = MLAAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE @@ -1793,7 +1799,8 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): def test_hnd_allocation_is_packed_head_major(): """Under HND the backend-visible logical view is the packed head-major - physical allocation.""" + physical allocation. + """ nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) stride_order = KVCacheLayout.LBHNC.layer_view_order @@ -1821,7 +1828,8 @@ def test_hnd_allocation_is_packed_head_major(): def test_main_cache_is_block_first_and_unpadded(): """The allocator's contiguous-view branch (not the padded-strided branch) is used for the main GQA cache: its spec is unpadded and the physical - layout keeps num_blocks as the first dimension under both layouts.""" + layout keeps num_blocks as the first dimension under both layouts. + """ spec = _main_spec() # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided(). assert spec.page_size_padded is None @@ -1843,7 +1851,8 @@ def _build_decode_inputs( ): """Shared decode setup: uniform query tokens per request, a non-identity block table, and topk indices selecting the current block plus older causal - blocks for each query token.""" + blocks for each query token. + """ active_batch = len(seq_lens_list) batch = active_batch + num_padded_reqs pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list] @@ -2179,7 +2188,8 @@ def test_decode_sparse_attention_correctness( ): """Decode (split-K) parity under both layouts: this is the only coverage of the decode-site cache feed, and the strided HND case fails if the kernel - ignores the cache strides.""" + ignores the cache strides. + """ torch.manual_seed(0) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs( seq_lens_list, decode_query_len, num_padded_reqs @@ -2229,7 +2239,8 @@ def test_decode_wrong_layout_breaks_parity(): already contiguous-NHD (i.e. skipping the allocator's inverse permute) reorders the K/V content, so the decode output no longer matches the reference computed on the correct logical view. The mislabeled tensor keeps - the same shape as the correct view, so the kernel stays in bounds.""" + the same shape as the correct view, so the kernel stays in bounds. + """ torch.manual_seed(0) seq_lens_list = (130, 257) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list) @@ -2264,7 +2275,8 @@ def test_decode_wrong_layout_breaks_parity(): def test_main_cache_byte_identical_through_production_allocator(): """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main `FullAttentionSpec` under HND and assert the backend-visible view has the - same shape, stride, and storage offset as the packed-HND allocation.""" + same shape, stride, and storage offset as the packed-HND allocation. + """ nb = 4 spec = _main_spec() raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) @@ -2279,7 +2291,8 @@ def test_main_cache_byte_identical_through_production_allocator(): def test_padded_main_cache_is_flagged(): """AC-2.1 negative: the M3 main cache relies on the allocator's contiguous-view branch (`page_size_padded is None`). A spec that sets - `page_size_padded` is explicitly flagged rather than silently wrong-strided.""" + `page_size_padded` is explicitly flagged rather than silently wrong-strided. + """ def _require_unpadded_block_first(spec, stride_order): inv_order = [stride_order.index(i) for i in range(len(stride_order))] @@ -2316,7 +2329,8 @@ def _require_unpadded_block_first(spec, stride_order): def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout): """AC-5 write path: the `reshape_and_cache_flash` write site now consumes packed-content K/V split views. Writing through those views must persist - into the bound storage under both layouts.""" + into the bound storage under both layouts. + """ torch.manual_seed(0) num_pages = 4 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index ffdb76e176f8..8106ba37a6e9 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -23,7 +23,8 @@ def test_concat_and_cache_mla_into_unified_slot_view(): """concat_and_cache_mla must write correctly into a per-layer view whose block stride is the full unified slot (block-major), with zero bleed into - the other layers' segments of the same slot.""" + the other layers' segments of the same slot. + """ from vllm import _custom_ops as ops torch.manual_seed(0) @@ -85,7 +86,8 @@ def write(cache): def test_flashmla_dense_decode_unified_slot_view(): """FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA on Hopper) must read a unified-slot block-major view bit-identically to a - contiguous per-layer cache.""" + contiguous per-layer cache. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -150,7 +152,8 @@ def run(kc): def test_flashinfer_mla_dense_decode_unified_slot_view(): """FlashInfer MLA dense decode must read a unified-slot block-major view (inflated stride(0), non-zero storage offset) bit-identically to a - contiguous per-layer cache.""" + contiguous per-layer cache. + """ try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: @@ -221,7 +224,8 @@ def run(kv): def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a unified-slot block-major view bit-identically to a contiguous fp8_ds_mla - cache, with finite nonzero output.""" + cache, with finite nonzero output. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_sparse_supported() @@ -310,7 +314,8 @@ def run(kc): def test_indexer_k_quant_and_cache_into_unified_slot_view(): """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must write correctly into a per-layer view whose block stride is the full - unified slot, with zero bleed into the other layers' segments.""" + unified slot, with zero bleed into the other layers' segments. + """ from vllm import _custom_ops as ops torch.manual_seed(0) @@ -355,7 +360,8 @@ def write(cache): def test_flashattn_mla_dense_decode_unified_slot_view(): """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot - block-major view bit-identically to a contiguous per-layer cache.""" + block-major view bit-identically to a contiguous per-layer cache. + """ try: from vllm.vllm_flash_attn import flash_attn_varlen_func except ImportError: @@ -428,7 +434,8 @@ def run(cache): def test_flashmla_dense_fp8_decode_unified_slot_view(): """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache) must read a unified-slot block-major view bit-identically to a contiguous - per-layer fp8 cache.""" + per-layer fp8 cache. + """ import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -499,7 +506,8 @@ def run(kc): def test_flashinfer_mla_dense_fp8_decode_unified_slot_view(): """FlashInfer MLA dense decode with an fp8 KV cache must read a unified-slot block-major view bit-identically to a contiguous per-layer - cache.""" + cache. + """ try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index dd2870ccf178..0fbe7bced2da 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -283,7 +283,8 @@ def _run_direct_flash_attn_case( # Backend contract tests -------------------------------------------------- def test_aiter_mha_backend_contract(): """The ROCm backend advertises the dtypes, shapes, and attention types it - is designed to handle.""" + is designed to handle. + """ from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -311,7 +312,8 @@ def test_aiter_mha_backend_contract(): def test_aiter_mha_backend_validates_kv_cache_block_size(): """The backend should reject KV cache specs that cannot be gathered - correctly by the ROCm kernel.""" + correctly by the ROCm kernel. + """ from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -331,7 +333,8 @@ def test_aiter_mha_backend_validates_kv_cache_block_size(): def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe(): """The backend should trust the ROCm MI3xx probe instead of the raw torch - capability tuple.""" + capability tuple. + """ from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -356,7 +359,8 @@ def test_aiter_mha_env_flags_control_enablement( use_aiter, use_mha, expected, monkeypatch ): """Both the global AITER flag and the MHA-specific flag must be enabled - before the ROCm MHA path is considered active.""" + before the ROCm MHA path is considered active. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -374,7 +378,8 @@ def test_aiter_mha_env_flags_control_enablement( def test_aiter_mha_platform_gate_matches_install_and_arch(): """The global AITER availability check should only open on ROCm MI3xx when - the aiter package is installed.""" + the aiter package is installed. + """ from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported assert is_aiter_found_and_supported() is ( @@ -645,7 +650,8 @@ def test_aiter_mha_multi_batch(num_heads, head_size, dtype): @pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS) def test_aiter_fa_large_block_table_matches_reference(num_blocks): """The direct paged-KV path should stay stable for both normal and very - large block tables.""" + large block tables. + """ atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() @@ -667,7 +673,8 @@ def test_aiter_fa_large_block_table_matches_reference(num_blocks): @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") def test_aiter_fa_sliding_window_matches_reference(): """The direct kernel should respect the same sliding-window causal mask as - the naive reference implementation.""" + the naive reference implementation. + """ atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py index 659211521a07..0944be8937c8 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py @@ -35,7 +35,7 @@ def _require_aiter(): @torch.inference_mode() def test_mla_decode_fwd_op_schema() -> None: - """opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. + """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. A single opcheck call covers that the op is registered/callable, that its fake implementation matches the real op (torch.compile tracing), and that diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index aabce97fac9e..17b59c074332 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -111,7 +111,8 @@ def _make_mixed_common_metadata(): def _patch_build_deps(monkeypatch, events=None): """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs - on CPU.""" + on CPU. + """ def fake_generate_sparse_seqlen_triton( query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len @@ -152,7 +153,8 @@ def test_build_populates_decode_only_split_fields(monkeypatch): def test_build_populates_mixed_split_fields(monkeypatch): """Mixed decode+prefill batch: split is reported, prefill fields stay - default because this impl always runs the MQA path.""" + default because this impl always runs the MQA path. + """ builder = _make_builder() _patch_build_deps(monkeypatch) diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index 5ac1d6cdcbf8..a46916e14ef3 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -125,7 +125,8 @@ def _quantize_to_fp8(tensor: torch.Tensor): @pytest.mark.parametrize("PAGE_SIZE", [1, 16]) def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE): """Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales, - and compare against BF16 reference output.""" + and compare against BF16 reference output. + """ assert CACHE_SIZE % PAGE_SIZE == 0 dtype = torch.bfloat16 seq_len = L @@ -248,7 +249,8 @@ def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SI """The kernel must honor the cache's page-dim stride, not assume pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must match a - contiguous cache holding the same data exactly.""" + contiguous cache holding the same data exactly. + """ B = 3 seq_len = 1027 CACHE_SIZE = 16384 diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 513811a09c04..04b7755853eb 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the Triton DiffKV unified-attention kernel. -""" +"""Unit tests for the Triton DiffKV unified-attention kernel.""" import pytest import torch diff --git a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py index bd267c43bebe..8368869b65e3 100644 --- a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py +++ b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant. +"""Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant. Tests both contiguous and non-contiguous (cross-layer unified) KV cache layouts against a pure-PyTorch reference implementation. @@ -43,8 +42,7 @@ def make_random_kv_cache( layout=KVCacheLayout.LBHNC, num_layers=None, ): - """ - Create a random fp8 KV cache mimicking the production allocation. + """Create a random fp8 KV cache mimicking the production allocation. Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order`` Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size) @@ -382,8 +380,7 @@ def test_large_block_size(): @torch.inference_mode() def test_cross_layer_many_layers(): - """ - Non-contiguous with 36 layers -- matches real gpt-oss-120b. + """Non-contiguous with 36 layers -- matches real gpt-oss-120b. Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps). """ from vllm.v1.attention.backends.flashinfer import ( diff --git a/tests/kernels/attention/test_xpu_mla_sparse.py b/tests/kernels/attention/test_xpu_mla_sparse.py index 965c6f735f0f..47139fc465d5 100644 --- a/tests/kernels/attention/test_xpu_mla_sparse.py +++ b/tests/kernels/attention/test_xpu_mla_sparse.py @@ -30,12 +30,12 @@ def reference_mla_sparse_prefill( topk_length: torch.Tensor | None = None, attn_sink: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Returns: + """Returns: - o: [s_q, h_q, dv] - o_fp32: [s_q, h_q, dv] - max_logits: [s_q, h_q] - lse: [s_q, h_q] + """ s_q, h_q, d_qk = q.shape s_kv, _, _ = kv.shape diff --git a/tests/kernels/core/test_apply_rotary_emb.py b/tests/kernels/core/test_apply_rotary_emb.py index 5c765bbb4a21..fd92de16bb39 100644 --- a/tests/kernels/core/test_apply_rotary_emb.py +++ b/tests/kernels/core/test_apply_rotary_emb.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for ApplyRotaryEmb CustomOp dispatch behavior. +"""Tests for ApplyRotaryEmb CustomOp dispatch behavior. This test ensures that RotaryEmbedding classes correctly call the appropriate ApplyRotaryEmb methods based on the calling context: @@ -179,8 +178,7 @@ def test_rotary_embedding_dispatch( test_case: RotaryEmbeddingTestCase, device: str, ): - """ - Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method. + """Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method. - forward_native methods should call ApplyRotaryEmb.forward_native() - forward_cuda/forward methods should call ApplyRotaryEmb.forward() diff --git a/tests/kernels/core/test_fused_embed_norm.py b/tests/kernels/core/test_fused_embed_norm.py index 9da43dba8361..fc30ab3acaca 100644 --- a/tests/kernels/core/test_fused_embed_norm.py +++ b/tests/kernels/core/test_fused_embed_norm.py @@ -41,7 +41,8 @@ def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: def test_fused_embed_norm_matches_reference(): """Main-model fusion: the residual is the exact gather and the second output is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16 - ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel.""" + ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel. + """ set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") @@ -60,7 +61,8 @@ def test_fused_embed_norm_matches_reference(): @torch.inference_mode() def test_fused_embed_eh_norm_matches_reference(): """MTP fusion (folded gather) is bit-exact vs gathering the embeds and - feeding the untouched model-local ``fused_eh_norm``.""" + feeding the untouched model-local ``fused_eh_norm``. + """ set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 3d0e2ac11e35..81382bdda91d 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -62,7 +62,8 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): a [1, q_size] slice of a [1, q_size + kv_size] buffer used to produce a qr_out with row stride q_size + kv_size. Downstream dispatchers that require packed row-major inputs then reject the tensor and silently fall - back to a slower GEMM path on every decode step.""" + back to a slower GEMM path on every decode step. + """ device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 @@ -88,7 +89,8 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int): """Regression guard: grid used to be (2, num_tokens), hitting CUDA's 65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2) - lifts that bound to 2**31-1.""" + lifts that bound to 2**31-1. + """ device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 86b8e7965673..74fdb00c30bf 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests that FusedRMSNormGated decomposes correctly under torch.compile, -matching the eager triton kernel output.""" +matching the eager triton kernel output. +""" import pytest import torch diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py index cd8efc14a3c2..b92a08a317dc 100644 --- a/tests/kernels/core/test_minimax_reduce_rms.py +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -36,7 +36,6 @@ def _worker_forward_qk( eps, ): """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel.""" - if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): cleanup_dist_env_and_memory() return diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 81e0a916a304..0e355d7cd352 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -229,7 +229,7 @@ def test_mrope_torch_compile_tracing( # Create a wrapper that makes the in-place function appear functional def functional_forward_cuda(pos, q, k): - """Wrapper that converts in-place operation to functional style + """Wrapper that converts in-place operation to functional style. CUDA Graph does not support in-place operations. This wrapper creates working copies of the diff --git a/tests/kernels/core/test_opcheck.py b/tests/kernels/core/test_opcheck.py index 40ced08b933a..45dd01d44e50 100644 --- a/tests/kernels/core/test_opcheck.py +++ b/tests/kernels/core/test_opcheck.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for miscellaneous utilities -""" +"""Tests for miscellaneous utilities.""" import torch diff --git a/tests/kernels/core/test_rocm_aiter_ops.py b/tests/kernels/core/test_rocm_aiter_ops.py index c15afd72d7fb..b2bed636e9f1 100644 --- a/tests/kernels/core/test_rocm_aiter_ops.py +++ b/tests/kernels/core/test_rocm_aiter_ops.py @@ -579,7 +579,8 @@ def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m): def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view(): """Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match - contiguous inputs.""" + contiguous inputs. + """ require_aiter() require_fp8() from vllm._aiter_ops import rocm_aiter_ops diff --git a/tests/kernels/core/test_rotary_embedding.py b/tests/kernels/core/test_rotary_embedding.py index 1cbb5dbd1881..353aefa11aae 100644 --- a/tests/kernels/core/test_rotary_embedding.py +++ b/tests/kernels/core/test_rotary_embedding.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for miscellaneous utilities -""" +"""Tests for miscellaneous utilities.""" import pytest import torch diff --git a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py index 1ac17335b69c..3f192ba2e7f1 100644 --- a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py +++ b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for fused MLA KV-cache write and RoPE fused kernel -""" +"""Tests for fused MLA KV-cache write and RoPE fused kernel.""" import random diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index b0038c084e3c..fa50d2b3eb39 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -87,7 +87,8 @@ def _shard(full_table, si, rows_per_rank, dtype): @pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) def test_matches_masked_gather(tp_size, vocab_size, org_vocab_size, hidden, id_dtype): """Every rank's partial output is bit-exact with the eager path, including - the added-vocab (LoRA) rows and the ids owned by another rank.""" + the added-vocab (LoRA) rows and the ids owned by another rank. + """ set_random_seed(7) dtype = torch.bfloat16 rows = pad_vocab_size(vocab_size, 64) // tp_size diff --git a/tests/kernels/helion/test_autotune.py b/tests/kernels/helion/test_autotune.py index 8b42e145d484..f28e8921093b 100644 --- a/tests/kernels/helion/test_autotune.py +++ b/tests/kernels/helion/test_autotune.py @@ -59,7 +59,8 @@ def teardown_method(self): def test_autotune_disabled_kernel_produces_valid_config(self): """Register a kernel with no configs (disabled), run autotune, - verify it produces a valid helion.Config.""" + verify it produces a valid helion.Config. + """ with dummy_kernel_registry(configs={}) as register: wrapper = register( "autotune_test_kernel", diff --git a/tests/kernels/helion/test_config_manager.py b/tests/kernels/helion/test_config_manager.py index f8e5eae6f106..63f813084926 100644 --- a/tests/kernels/helion/test_config_manager.py +++ b/tests/kernels/helion/test_config_manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for Helion ConfigManager and ConfigSet. +"""Unit tests for Helion ConfigManager and ConfigSet. Tests the simplified configuration management system for Helion custom kernels. """ diff --git a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py index 6f5a40fad861..14a4356e14a2 100644 --- a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py +++ b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel +"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel. Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`. """ diff --git a/tests/kernels/helion/test_fused_qk_norm_rope.py b/tests/kernels/helion/test_fused_qk_norm_rope.py index 4ced522d9193..a973533169a9 100644 --- a/tests/kernels/helion/test_fused_qk_norm_rope.py +++ b/tests/kernels/helion/test_fused_qk_norm_rope.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the fused_qk_norm_rope helion kernel +"""Tests for the fused_qk_norm_rope helion kernel. Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`. """ diff --git a/tests/kernels/helion/test_helion_available.py b/tests/kernels/helion/test_helion_available.py index 560a671f149b..a4de68bd6716 100644 --- a/tests/kernels/helion/test_helion_available.py +++ b/tests/kernels/helion/test_helion_available.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for Helion kernel availability and basic functionality. +"""Tests for Helion kernel availability and basic functionality. This module demonstrates the pattern for testing optional Helion kernels. Tests in this directory will be skipped if Helion is not installed. diff --git a/tests/kernels/helion/test_per_token_group_fp8_quant.py b/tests/kernels/helion/test_per_token_group_fp8_quant.py index d9d96f9287c2..d9dd8a881763 100644 --- a/tests/kernels/helion/test_per_token_group_fp8_quant.py +++ b/tests/kernels/helion/test_per_token_group_fp8_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the per_token_group_fp8_quant helion kernel +"""Tests for the per_token_group_fp8_quant helion kernel. Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`. """ diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index 9876135056b2..572a4126cf24 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for Helion kernel registration. +"""Unit tests for Helion kernel registration. Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch including config picker registration and custom autotuner integration. @@ -599,7 +598,8 @@ def default_picker(args, config_keys): ) def test_init_eagerly_initializes_hop_path(self): """Test that register_kernel eagerly builds the configured kernel - on the HOP path (no custom op registration needed).""" + on the HOP path (no custom op registration needed). + """ from vllm.kernels.helion.utils import get_canonical_gpu_name configs: dict[CaseKey, helion.Config] = { @@ -634,7 +634,8 @@ def test_init_eagerly_initializes_hop_path(self): ) def test_init_eagerly_initializes(self): """Test that register_kernel eagerly loads configs and detects GPU - during construction so __call__ needs no further initialization.""" + during construction so __call__ needs no further initialization. + """ from vllm.kernels.helion.utils import get_canonical_gpu_name with ( @@ -962,7 +963,8 @@ class TestTorchCompileHOP: def test_compiled_graph_contains_helion_hop(self): """Verify torch.compile on a HelionKernelWrapper emits a - helion_kernel_wrapper_mutation HOP node in the FX graph.""" + helion_kernel_wrapper_mutation HOP node in the FX graph. + """ configs: dict[CaseKey, helion.Config] = { CaseKey.default(): helion.Config(block_sizes=[4, 4]) } @@ -1020,7 +1022,6 @@ def f(x, y): ) def test_inductor_backend_compiles_helion_hop(self): """Test torch.compile with inductor backend and Helion fusion enabled.""" - configs: dict[CaseKey, helion.Config] = { CaseKey.default(): helion.Config(block_sizes=[4, 4]) } diff --git a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py index 50591ae34a61..d0650dfc7afe 100644 --- a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py +++ b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the rms_norm_dynamic_per_token_quant helion kernel +"""Tests for the rms_norm_dynamic_per_token_quant helion kernel. Run `pytest tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py`. """ diff --git a/tests/kernels/helion/test_rms_norm_per_block_quant.py b/tests/kernels/helion/test_rms_norm_per_block_quant.py index 4ffc759c63d2..e9c003f4e303 100644 --- a/tests/kernels/helion/test_rms_norm_per_block_quant.py +++ b/tests/kernels/helion/test_rms_norm_per_block_quant.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the rms_norm_per_block_quant helion kernel +"""Tests for the rms_norm_per_block_quant helion kernel. Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`. """ diff --git a/tests/kernels/helion/utils.py b/tests/kernels/helion/utils.py index 38893fc8fecb..4a282f275d10 100644 --- a/tests/kernels/helion/utils.py +++ b/tests/kernels/helion/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Helion Kernel test utils""" +"""Helion Kernel test utils.""" import pytest import torch diff --git a/tests/kernels/ir/test_aiter_norm_dispatch.py b/tests/kernels/ir/test_aiter_norm_dispatch.py index 19cf5d3112dc..f84a7615b167 100644 --- a/tests/kernels/ir/test_aiter_norm_dispatch.py +++ b/tests/kernels/ir/test_aiter_norm_dispatch.py @@ -21,7 +21,7 @@ def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim): - """q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" + """Q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim qkv = torch.empty(num_tokens, q_size + 2 * kv_size) q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] diff --git a/tests/kernels/ir/test_ir_ops.py b/tests/kernels/ir/test_ir_ops.py index 730db9833b1d..560d047c0f01 100644 --- a/tests/kernels/ir/test_ir_ops.py +++ b/tests/kernels/ir/test_ir_ops.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Meta-tests for vLLM IR op infrastructure. +"""Meta-tests for vLLM IR op infrastructure. Ensures all registered ops have input generators defined. Per-op correctness tests live alongside their op definitions diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index c324b189b9df..e41b7c09d3fe 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -368,7 +368,8 @@ def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): @pytest.mark.parametrize("provider", ["vllm_c"]) def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider): """Test that inplace implementations reuse inputs, - for maybe_inplace overload but not for default overload.""" + for maybe_inplace overload but not for default overload. + """ impl = ir.ops.fused_add_rms_norm.impls[provider] if not impl.supported: pytest.skip(f"{provider} impl not supported on this platform") diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index c239f831c614..9495465b9ec3 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -1055,7 +1055,8 @@ def test_batch_memcpy_cpu_fallback() -> None: def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"): """High-precision conv reference: everything in float32, no bf16 rounding. - x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim].""" + x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim]. + """ xf = x.float() wf = weight.float().unsqueeze(1) bf = bias.float() @@ -1164,7 +1165,8 @@ def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens): def test_conv_cpp_vnni_packed_matches_torch(seq_lens): """The exact runtime prefill sequence (VNNI-packed weight + SD-layout conv_state view + is_vnni=True) must match the torch fallback. Validates - the packing + layout handoff on any AVX-512BF16 CPU.""" + the packing + layout handoff on any AVX-512BF16 CPU. + """ x, weight, bias = _conv_inputs(sum(seq_lens)) out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True) @@ -1178,7 +1180,8 @@ def test_conv_cpp_vnni_packed_matches_torch(seq_lens): @torch.inference_mode() def test_conv_update_cpp_matches_torch(batch): """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch, - including the in-place conv_state update (the next-step handoff).""" + including the in-place conv_state update (the next-step handoff). + """ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_update_torch, ) diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index d2a2981edf2c..d4f5fac59ca6 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -36,8 +36,7 @@ def causal_conv1d_ref( final_states_out: torch.Tensor | None = None, activation: str | None = "silu", ): - """ - x: (batch, dim, seqlen) + """x: (batch, dim, seqlen) weight: (dim, width) bias: (dim,) initial_states: (batch, dim, width - 1) @@ -72,8 +71,7 @@ def causal_conv1d_ref( def causal_conv1d_update_ref( x, conv_state, weight, bias=None, activation=None, cache_seqlens=None ): - """ - x: (batch, dim) or (batch, dim, seqlen) + """x: (batch, dim) or (batch, dim, seqlen) conv_state: (batch, dim, state_len), where state_len >= width - 1 weight: (dim, width) bias: (dim,) @@ -136,8 +134,7 @@ def causal_conv1d_opcheck_fn( activation: str | None = "silu", null_block_id: int = NULL_BLOCK_ID, ): - """ - x: (batch, dim, seqlen) + """x: (batch, dim, seqlen) weight: (dim, width) bias: (dim,) seq_idx: (batch, seqlen) diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index f8af3db56be6..c24f4c0f8719 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -51,8 +51,7 @@ def selective_scan_ref( prev_state=None, final_state_out=None, ): - """ - u: r(B D L) + """u: r(B D L) delta: r(B D L) A: c(D N) or r(D N) B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L) @@ -135,7 +134,7 @@ def selective_scan_opcheck_fn( cu_chunk_seqlen=None, last_chunk_indices=None, ): - """if return_last_state is True, returns (out, last_state) + """If return_last_state is True, returns (out, last_state) last_state has shape (batch, dim, dstate). """ if u.stride(-1) != 1: diff --git a/tests/kernels/mamba/test_mamba_ssm_configs.py b/tests/kernels/mamba/test_mamba_ssm_configs.py index d35ce95746c7..4330caa6115a 100644 --- a/tests/kernels/mamba/test_mamba_ssm_configs.py +++ b/tests/kernels/mamba/test_mamba_ssm_configs.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the JSON-based config loader added to selective_state_update. +"""Unit tests for the JSON-based config loader added to selective_state_update. Tests cover: - Flat MoE-style filename generation @@ -123,7 +122,8 @@ def test_fallback_when_no_config(monkeypatch, tmp_path): def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): """When effective_batch = batch*nheads is not an exact key, the closest - key should be selected.""" + key should be selected. + """ _write_config( tmp_path, dstate=32, @@ -168,7 +168,8 @@ def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): def test_non_dict_json_returns_none(monkeypatch, tmp_path): """A valid JSON file that is not a dict (e.g. a list) must be ignored - and return None rather than raising AttributeError.""" + and return None rather than raising AttributeError. + """ device_name = get_ssm_device_name() config_path = tmp_path / get_ssm_config_file_name( _HEADDIM, 16, _CACHE_DTYPE, device_name @@ -190,7 +191,8 @@ def test_non_dict_json_returns_none(monkeypatch, tmp_path): def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path): """An empty JSON object {} must not crash min() — should fall back - to the hard-coded heuristic.""" + to the hard-coded heuristic. + """ _write_config(tmp_path, dstate=64, payload={}) monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path)) diff --git a/tests/kernels/mamba/test_mamba_ssm_ssd.py b/tests/kernels/mamba/test_mamba_ssm_ssd.py index 1de25780eac4..22abeab1ed8f 100644 --- a/tests/kernels/mamba/test_mamba_ssm_ssd.py +++ b/tests/kernels/mamba/test_mamba_ssm_ssd.py @@ -41,14 +41,15 @@ def segsum(x): def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None): - """ - Arguments: + """Arguments: X: (batch, length, n_heads, d_head) A: (batch, length, n_heads) B: (batch, length, n_heads, d_state) C: (batch, length, n_heads, d_state) + Return: Y: (batch, length, n_heads, d_head) + """ assert X.dtype == A.dtype == B.dtype == C.dtype assert X.shape[1] % block_len == 0 diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py index 063ff17dd337..fbc7b691870e 100644 --- a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -68,7 +68,8 @@ def _run_prefill_decode_equivalence( the exact fp32 recurrence (and each other). All paths use the production dt flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` - is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" + is the recurrent-state precision; ``act_dtype`` the activation/buffer one. + """ device = current_platform.device_type rtol, atol = _prefill_tolerances(act_dtype) set_random_seed(seed) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py index 588ce52dbeb6..6e5ada19bc8f 100644 --- a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -95,7 +95,8 @@ def _run_standard_decode( ) -> None: """Drive ``num_steps`` decode steps and check every path against the trusted recurrence. ``state_dtype`` is the recurrent-state precision; - ``act_dtype`` is the activation/buffer precision.""" + ``act_dtype`` is the activation/buffer precision. + """ device = "cuda" both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) @@ -539,7 +540,8 @@ def test_replayssm_standard_decode_tp_head_shard_equivalence( separable: one run over all heads must equal concatenating ``tp`` independent per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This guards the per-rank divisors and group->head mapping the state-shape wiring - relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite. + """ state_dtype, act_dtype = precision nheads, headdim, dstate, ngroups = geometry assert nheads % tp == 0 and ngroups % tp == 0 diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index b82c4d82c33e..c7cf47342778 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -9,8 +9,7 @@ def selective_state_update_ref( state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False ): - """ - Argument: + """Argument: state: (batch, dim, dstate) or (batch, nheads, dim, dstate) x: (batch, dim) or (batch, nheads, dim) dt: (batch, dim) or (batch, nheads, dim) @@ -20,8 +19,10 @@ def selective_state_update_ref( D: (dim,) or (nheads, dim) z: (batch, dim) or (batch, nheads, dim) dt_bias: (dim,) or (nheads, dim) + Return: out: (batch, dim) or (batch, nheads, dim) + """ has_heads = state.dim() > 3 if state.dim() == 3: diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 8041db68d752..dd17298cb367 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -148,9 +148,7 @@ def num_local_experts(self) -> int: return self.E // self.world_size def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: - """ - make env data for vllm launch. - """ + """Make env data for vllm launch.""" vllm_config = VllmConfig() vllm_config.model_config = SimpleNamespace( enforce_eager=True, @@ -169,7 +167,8 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: def fe_supports_quant_scheme(self) -> bool: """Check if the fused experts class supports this quant config. - See https://github.com/ROCm/aiter/issues/2419 for AITER gaps.""" + See https://github.com/ROCm/aiter/issues/2419 for AITER gaps. + """ if self.quant_config is None or self.quant_dtype is None: return True if self.quant_dtype != torch.float8_e4m3fn: @@ -444,9 +443,7 @@ def describe(self): def make_hidden_states( config: Config, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - Return hidden_states - """ + """Return hidden_states.""" m, k, dtype = (config.M, config.K, config.dtype) device = torch.accelerator.current_device_index() a = torch.randn((m, k), device=device, dtype=dtype) / 15.0 diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index f926b034b3c2..972ebe6512eb 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DeepEP test utilities -""" +"""DeepEP test utilities.""" import dataclasses import os diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index d0b1ce7cae46..de8f75d0e9a1 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -468,7 +468,8 @@ def test_deepseek_v4_flashinfer_cutlass_falls_through_to_w4a8( monkeypatch: pytest.MonkeyPatch, ) -> None: """Explicit flashinfer_cutlass must try the W4A8 variant when the BF16 - variant is unsupported (the BF16 variant is gated to SM90).""" + variant is unsupported (the BF16 variant is gated to SM90). + """ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( FlashInferExperts, ) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 3fca9190dc31..4c4c82d54e17 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -34,7 +34,6 @@ def test_batched_deepgemm_vs_triton( E: int, T: int, K: int, N: int, topk: int, monkeypatch, workspace_init ): """Compare BatchedDeepGemmExperts to BatchedTritonExperts.""" - monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1") device = "cuda" diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index 9be58e6337bd..c4ebc5ff4077 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -129,7 +129,8 @@ def test_batched_mm( per_act_token_quant: bool, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -268,7 +269,8 @@ def test_fused_moe_batched_experts( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -649,7 +651,8 @@ def test_batched_experts_masked_activation_end_to_end( def test_batched_experts_end_to_end(m, n, k, e, topk): """End-to-end BatchedTritonExperts via the reference (no-comms) BatchedPrepareAndFinalize, validated against a torch reference. Exercises - the device-enablement path.""" + the device-enablement path. + """ from vllm.v1.worker.workspace import init_workspace_manager set_random_seed(7) diff --git a/tests/kernels/moe/test_block_int8.py b/tests/kernels/moe/test_block_int8.py index 415faafb3ab0..84a4a0d2d102 100644 --- a/tests/kernels/moe/test_block_int8.py +++ b/tests/kernels/moe/test_block_int8.py @@ -54,7 +54,8 @@ # For test def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape): """This function performs fused moe with block-wise quantization using - native torch.""" + native torch. + """ B, D = a.shape a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) @@ -97,7 +98,8 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed): """Tests the fused_moe kernel with W8A8 INT8 block quantization against a - native torch reference.""" + native torch reference. + """ torch.manual_seed(seed) a = torch.randn((M, K), dtype=dtype) / 10 diff --git a/tests/kernels/moe/test_count_expert_num_tokens.py b/tests/kernels/moe/test_count_expert_num_tokens.py index 8014076192e2..b2b6abc5b489 100644 --- a/tests/kernels/moe/test_count_expert_num_tokens.py +++ b/tests/kernels/moe/test_count_expert_num_tokens.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests compute_expert_num_tokens kernels -""" +"""Tests compute_expert_num_tokens kernels.""" import dataclasses diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index dccda70c0b9e..e7af163b04d9 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -397,7 +397,8 @@ def test_cpu_fused_moe_int8( class _StubMoELayer(torch.nn.Module): """Minimal stand-in for the real MoE layer module, exposing just what the unquantized CPU experts read/replace (w13_weight, w2_weight, the - router configuration, and optionally w13_bias/w2_bias).""" + router configuration, and optionally w13_bias/w2_bias). + """ def __init__( self, @@ -500,8 +501,8 @@ def test_cpu_fused_moe_unaligned_intermediate_size( act: MoEActivation, ): """CPU kernels handle unaligned intermediate sizes by zero-padding the - weights before prepacking.""" - + weights before prepacking. + """ set_random_seed(0) batch_size = 64 intermediate_size = UNALIGNED_INTERMEDIATE_DIM diff --git a/tests/kernels/moe/test_cpu_int4_moe.py b/tests/kernels/moe/test_cpu_int4_moe.py index 04931e386f14..2d5534ba086e 100644 --- a/tests/kernels/moe/test_cpu_int4_moe.py +++ b/tests/kernels/moe/test_cpu_int4_moe.py @@ -62,6 +62,7 @@ def _make_int4_moe_weights( Returns: (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias) where *_ref are the dequantized float reference weights + """ # Generate INT4 weights as int8 values in [-8, 7] w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8) diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 84ff1ae266db..99cae4534be7 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -826,6 +826,7 @@ def _make_int4_moe_weights(E, N, K, group_size, quant_algo): w1_zeros, w2_zeros, w1_zeros_packed, w2_zeros_packed, w1_s, w2_s + """ w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32) w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32) diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py index 3a371968fc51..09a7256a5cdc 100644 --- a/tests/kernels/moe/test_cutedsl_moe.py +++ b/tests/kernels/moe/test_cutedsl_moe.py @@ -91,11 +91,9 @@ def break_fp4_bytes(a, dtype): def generate_balanced_routing( hidden_states: torch.Tensor, num_experts: int, top_k: int ): - """ - Generate routing weights and topk indices such that every expert is active. + """Generate routing weights and topk indices such that every expert is active. Returns routing_weights, topk_idx """ - num_tokens, hidden_dim = hidden_states.shape # num_tokens = batch_size * seq_len @@ -232,13 +230,13 @@ def grouped_gemm_ref( *, block_size: int = 16, ) -> torch.Tensor: - """ - Computes the reference grouped GEMM (fp4 quantized per-expert loop), + """Computes the reference grouped GEMM (fp4 quantized per-expert loop), computes flashinfer grouped GEMM (for scale consistency), and returns ONLY the repacked reference output: out_ref. Returns: out_ref: Tensor [num_experts, max_m, n_out] + """ device_hs = hidden_states_expanded.device device_w = weights.device diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index efb1e2f2969d..b755741380ed 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test DeepEP + DeepGEMM integration +"""Test DeepEP + DeepGEMM integration DeepGEMM are gemm kernels specialized for the fp8 block-quantized case. """ @@ -92,9 +91,7 @@ def make_block_quant_fp8_weights( k: int, block_size: list[int], ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Return weights w1q, w2q, w1_scale, w2_scale - """ + """Return weights w1q, w2q, w1_scale, w2_scale.""" (_, w1q, w1_scale, _), (_, w2q, w2_scale, _) = make_test_weights( e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size ) @@ -459,10 +456,7 @@ def test_ht_deepep_deepgemm_moe( disable_deepgemm_ue8m0, workspace_init, ): - """ - Tests for High-Throughput DeepEP + DeepGemm integration. - """ - + """Tests for High-Throughput DeepEP + DeepGemm integration.""" m, n, k = mnk set_random_seed(7) @@ -533,9 +527,7 @@ def test_ll_deepep_deepgemm_moe( disable_deepgemm_ue8m0, workspace_init, ): - """ - Tests for Low-Latency DeepEP + DeepGemm integration. - """ + """Tests for Low-Latency DeepEP + DeepGemm integration.""" assert not is_deep_gemm_e8m0_used() m, n, k = mnk diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index ce9e4f343d62..ea2113d64b43 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test deepep dispatch-combine logic -""" +"""Test deepep dispatch-combine logic.""" import dataclasses @@ -56,9 +54,7 @@ def make_weights( e, n, k, dtype ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Return weights w1, w2, w1_scale, w2_scale - """ + """Return weights w1, w2, w1_scale, w2_scale.""" if dtype in [torch.float16, torch.bfloat16]: w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index f1e34d5bd780..2fbde2f86e06 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. +"""Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. Compares against a pure-PyTorch reference MoE implementation. """ diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index ae6c9685a459..f6df64c51bf7 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP). +"""Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP). Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts. """ @@ -79,8 +78,7 @@ def make_block_quant_fp8_weights( k: int, block_size: list[int], ): - """ - Generate (w1, w2) expert weights and their per-block scale tensors + """Generate (w1, w2) expert weights and their per-block scale tensors in FP8 block-quantized format. w1 shape: (E, 2N, K) @@ -121,8 +119,7 @@ def make_block_quant_fp8_weights( def run_single_case(m, n, k, topk, num_experts, block_size): - """ - Run one (M,N,K) configuration on a single GPU and assert DeepGEMM == + """Run one (M,N,K) configuration on a single GPU and assert DeepGEMM == Triton baseline within tolerance. """ tokens_bf16 = ( @@ -248,8 +245,7 @@ def make_mxfp4_weights( n: int, k: int, ): - """ - Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales, + """Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales, plus BF16 reference weights for validation. w1 shape: (E, 2N, K//2) uint8 — packed FP4 @@ -310,8 +306,7 @@ def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids): def run_single_fp4_case(m, n, k, topk, num_experts): - """ - Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert + """Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert DeepGEMM FP4 == BF16 reference within tolerance. """ tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2e022ede3627..9627e2c67d4d 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -464,7 +464,8 @@ def test_trtllm_fp8_swiglu_clamp_support( ): """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the - per-tensor kernel has no clamp, and Relu2 rejects the parameters.""" + per-tensor kernel has no clamp, and Relu2 rejects the parameters. + """ class _Experts(TrtLlmFp8ExpertsModular): @staticmethod diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py index d8e449ec0aff..3d637258d3f2 100644 --- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -83,7 +83,7 @@ def _reference_activation( beta: float | None, limit: float | None, ): - """vLLM's own op for this activation, so the reference is not re-derived.""" + """VLLM's own op for this activation, so the reference is not re-derived.""" if activation == MoEActivation.RELU2_NO_MUL: return lambda x: torch.square(torch.relu(x)) if activation == MoEActivation.SWIGLUOAI: diff --git a/tests/kernels/moe/test_fused_topk.py b/tests/kernels/moe/test_fused_topk.py index 9c196a884d9f..18d1d9b4dfb0 100644 --- a/tests/kernels/moe/test_fused_topk.py +++ b/tests/kernels/moe/test_fused_topk.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MoE fused topk kernel +"""Tests for the MoE fused topk kernel. Run `pytest tests/kernels/moe/test_fused_topk.py`. """ diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 4afcba706ccd..0a16a3179739 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -377,7 +377,8 @@ def test_unit_shuffle(): @pytest.mark.parametrize("n_experts,topk", [(32, 4), (128, 4)]) def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): """routing_data_from_sparse_topk must produce routing structures - identical to make_routing_data for the same topk result.""" + identical to make_routing_data for the same topk result. + """ from vllm.model_executor.layers.fused_moe.experts import ( gpt_oss_triton_kernels_moe as gptoss_moe, ) diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index 1ab8550d64e8..c65cd381db86 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MoE grouped topk kernel +"""Tests for the MoE grouped topk kernel. Run `pytest tests/kernels/moe/test_grouped_topk.py`. """ diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index 6756d479febb..b647439c49f7 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -31,6 +31,7 @@ def mxint4_quantize( Returns: - uint8 packed (2 INT4/byte): [..., k//2] - stores SIGNED INT4 [-8, 7] - scales in BF16: [..., k//sf_vec_size] + """ x_reshaped = x.reshape(-1, sf_vec_size) x_max = x_reshaped.max(dim=-1, keepdim=True)[0].to(torch.float32) @@ -61,6 +62,7 @@ def mxint4_quantize_moe_weights( Returns: - weights_mxint4: Quantized weights [e, n, k//2] uint8 - scales_mxint4: Quantization scales [e, n, k//group_size] bf16 + """ e = weights_bf16.shape[0] weight_list = [] @@ -101,6 +103,7 @@ def marlin_quantize_moe_weights( Returns: - weights_marlin: Marlin quantized weights [e, k//8, n] int32 - scales_marlin: Marlin quantization scales [e, k//group_size, n] bf16 + """ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( marlin_quantize, @@ -286,7 +289,8 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group @torch.inference_mode() def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk): """Test that the flashinfer_trtllm_mxint4_moe wrapper matches the raw - trtllm_mxint4_block_scale_moe kernel call.""" + trtllm_mxint4_block_scale_moe kernel call. + """ pytest.importorskip("flashinfer") from flashinfer import RoutingMethodType from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 0c0e1d61f90e..6c21eb0d9ade 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -352,7 +352,8 @@ def test_modular_kernel_combinations_singlegpu( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware.""" + and those tests will be skipped on unsupported hardware. + """ config = Config( Ms=Ms, K=k, diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index d0fd4df23ccd..47be82dd2ab1 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test modular OAI Triton MoE -""" +"""Test modular OAI Triton MoE.""" from __future__ import annotations diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index ea732d68876a..8b4efaf91379 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -119,8 +119,7 @@ def iterative_moe( expert_map: torch.Tensor = None, renormalize: bool = False, ) -> torch.Tensor: - """ - Baseline implementation of fused moe. + """Baseline implementation of fused moe. Args: hidden_states: [*, hidden_size] @@ -128,6 +127,7 @@ def iterative_moe( w2: [num_experts, hidden_size, intermediate_size] gating_output: [*, num_experts] expert_map: [num_experts] + """ orig_shape = hidden_states.shape hidden_size = hidden_states.shape[-1] @@ -1956,8 +1956,7 @@ def __init__( ) def is_valid(self): - """ - Return True only if the input can be represented in a Batched + """Return True only if the input can be represented in a Batched format. """ return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch) @@ -2078,9 +2077,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend( dtype: torch.dtype, workspace_init, ): - """ - Test BF16 unquantized MoE with FlashInfer TRTLLM backend. - """ + """Test BF16 unquantized MoE with FlashInfer TRTLLM backend.""" set_random_seed(7) from vllm.model_executor.layers.fused_moe.config import ( diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index 8c16690785c3..a1599f5ad3da 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -54,8 +54,7 @@ def _verify_expert_level_sorting( valid_length: int, total_tokens: int, ): - """ - Verify that actual_sorted_ids follows the correct expert-level sorting. + """Verify that actual_sorted_ids follows the correct expert-level sorting. The kerne limplementation may or may not preserve original token order in topk_ids in the final sorted_ids however this does not impact quality. """ @@ -96,8 +95,7 @@ def torch_moe_align_block_size( expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Golden torch implementation of moe_align_block_size. + """Golden torch implementation of moe_align_block_size. This function aligns the token distribution across experts to be compatible with block size for matrix multiplication by sorting tokens by expert and @@ -189,7 +187,7 @@ def torch_moe_align_block_size( def test_moe_align_block_size( m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool ): - """Test moe_align_block_size without expert mapping""" + """Test moe_align_block_size without expert mapping.""" topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32) for i in range(m): experts = torch.randperm(num_experts, device="cuda")[:topk] @@ -254,7 +252,7 @@ def test_moe_align_block_size_with_expert_map( block_size: int, mask_inactive_experts: bool, ): - """Test moe_align_block_size with expert mapping (EP scenario)""" + """Test moe_align_block_size with expert mapping (EP scenario).""" expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32) local_experts = list(range(0, num_experts, 2)) for i, expert_id in enumerate(local_experts): diff --git a/tests/kernels/moe/test_moe_kernel_oracle.py b/tests/kernels/moe/test_moe_kernel_oracle.py index fbdf804a3b78..6ba5753f8d19 100644 --- a/tests/kernels/moe/test_moe_kernel_oracle.py +++ b/tests/kernels/moe/test_moe_kernel_oracle.py @@ -21,7 +21,8 @@ class TestUnquantizedDelegation: """UnquantizedMoEKernelOracle methods must delegate to the existing - module-level functions; behaviour is bit-identical.""" + module-level functions; behaviour is bit-identical. + """ def test_make_kernel_delegates(self) -> None: quant_config = object() diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index e1a152a23325..38b9dd025228 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -240,8 +240,7 @@ def maybe_roundup_layer_hidden_size( act_dtype: torch.dtype, backend: str | None, ) -> int: - """ - Given layer hidden size and MoE configurations, round up hidden_size + """Given layer hidden size and MoE configurations, round up hidden_size if necessary. Args: @@ -253,6 +252,7 @@ def maybe_roundup_layer_hidden_size( Rounded up hidden_size if rounding up is required based on the configs and all2all backend. Original hidden size otherwise. + """ if backend == "deepep_high_throughput": from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( @@ -315,7 +315,8 @@ def tp_chunk_gate_up( device: torch.device | int | None = None, ) -> torch.Tensor: """TP-chunk a combined [gate; up] weight, splitting each half separately - so every rank gets a portion of both gate and up.""" + so every rank gets a portion of both gate and up. + """ half = w.shape[dim] // 2 gate = chunk_by_rank( w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device @@ -945,6 +946,7 @@ def create_shared_experts_from_config( Returns: TestMLP instance or None if config is None + """ if shared_experts_config is None: return None @@ -991,6 +993,7 @@ def setup_moe_test_data( Returns: MoETestData containing all test data and transforms + """ # For latent MoE: latent_size = k // 2 latent_size = k // 2 @@ -1710,7 +1713,6 @@ def test_moe_layer_no_parallel( monkeypatch, ): """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False).""" - if os.environ.get("VLLM_LOGGING_LEVEL") is None: monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 4089bbfe843a..76d1e57f58d0 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MOE permute/unpermute kernel +"""Tests for the MOE permute/unpermute kernel. Run `pytest tests/kernels/test_moe_permute_unpermute.py`. """ diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py index 16b233b935e0..7c7ca1088574 100644 --- a/tests/kernels/moe/test_mxfp4_moe.py +++ b/tests/kernels/moe/test_mxfp4_moe.py @@ -67,8 +67,7 @@ def compute_ref_output( @pytest.mark.parametrize("num_experts", [8, 16, 32]) @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): - """ - Test the MXFP4 grouped GEMM kernel by: + """Test the MXFP4 grouped GEMM kernel by: 1. Creating random per-expert inputs and weights 2. Quantizing both to MXFP4 using the CUDA kernel 3. Running the CUTLASS grouped GEMM @@ -205,9 +204,7 @@ def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): reason="mxfp4_experts_quant requires CUDA SM100", ) def test_mxfp4_experts_quant_basic(): - """ - Basic smoke test for the MXFP4 experts quantization kernel. - """ + """Basic smoke test for the MXFP4 experts quantization kernel.""" device = "cuda" num_experts = 4 k = 256 @@ -295,8 +292,7 @@ def compute_reference_e8m0_scale(block_max: float) -> int: @pytest.mark.parametrize("k", [256, 7168]) @pytest.mark.parametrize("m", [16, 64]) def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): - """ - Test that mxfp4_experts_quant computes E8M0 block scales correctly + """Test that mxfp4_experts_quant computes E8M0 block scales correctly per OCP MX spec (not the NVFP4 formula). The old buggy kernel used: floor(log2(max/6)) + 127 @@ -416,8 +412,7 @@ def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): reason="mxfp4_experts_quant requires CUDA SM100", ) def test_mxfp4_experts_quant_no_saturation(): - """ - Test that the E8M0 scale is large enough to avoid unexpected saturation. + """Test that the E8M0 scale is large enough to avoid unexpected saturation. With the buggy NVFP4 formula, the scale was too small causing most values to saturate to ±6 in FP4. The fixed OCP MX formula should ensure that diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 830ff48734e1..e94882a310a2 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -130,7 +130,8 @@ def test_is_supported_config(present, ep_size, supported, reason_substr): def test_explicit_moe_backend_aiter(): """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear - ValueError when the flydsl package is missing.""" + ValueError when the flydsl package is missing. + """ with _gfx950(), _flydsl_installed(True): assert ( _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) @@ -161,8 +162,9 @@ def test_gfx950_picks_aiter(): def test_gfx942_picks_emulation(): - """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton - dot_scaled backend wins instead.""" + """Flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead. + """ with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False): backend, experts_cls = select_mxfp8_moe_backend(_config()) assert backend is Fp8MoeBackend.EMULATION diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 9157d4eaa254..d513a14f775a 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -214,8 +214,7 @@ def reference_moe( activation: str = "swiglu", use_interleaved_layout: bool = False, ): - """ - Reference MoE implementation for accuracy testing. + """Reference MoE implementation for accuracy testing. Args: activation: One of "swiglu", "silu", "relu2". Controls the activation @@ -224,6 +223,7 @@ def reference_moe( (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI. If False, uses chunked layout (gate, up = chunk(x, 2)) as used by standard swiglu/silu. + """ # renormalize routing experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True) @@ -1441,8 +1441,7 @@ def test_rocm_mxfp4_moe_oracle( intermediate_size: int, monkeypatch: pytest.MonkeyPatch, ): - """ - Test ROCm MXFP4 MoE using oracle functions. + """Test ROCm MXFP4 MoE using oracle functions. This test validates that the oracle functions work end-to-end: - select_mxfp4_moe_backend() selects a valid backend @@ -1825,8 +1824,7 @@ def test_mxfp4_emulation_rounds_up_to_block_size( def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( monkeypatch: pytest.MonkeyPatch, ): - """ - select_mxfp4_moe_backend() must raise NotImplementedError, with the + """select_mxfp4_moe_backend() must raise NotImplementedError, with the collected per-backend unsupported reasons in the message, when no backend supports the requested deployment configuration. """ diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 0967f9499e0f..de4ae5ae31d8 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -249,7 +249,8 @@ def _make_topk_ids( device: str = "cuda", ) -> torch.Tensor: """Generate distinct expert IDs per token with the same top-k shape as - production routing.""" + production routing. + """ router_logits = torch.randn(num_tokens, num_experts, device=device) _, topk_ids = torch.topk(torch.softmax(router_logits, dim=-1), k=topk, dim=-1) return topk_ids.to(torch.int32) @@ -443,7 +444,8 @@ def test_aiter_fused_moe_custom_op_registered(): def test_aiter_asm_moe_tkw1_custom_op_registered(): """The tkw1 custom op should stay registered for FP8 apply-router-weight - paths.""" + paths. + """ _assert_aiter_supported() import vllm._aiter_ops as aiter_ops # noqa: F401 @@ -453,7 +455,8 @@ def test_aiter_asm_moe_tkw1_custom_op_registered(): def test_aiter_fused_moe_fake_tensor_support(): """The fused-MoE op should preserve fake-tensor compatibility for - torch.compile-style tracing.""" + torch.compile-style tracing. + """ _assert_aiter_supported() import vllm._aiter_ops # noqa: F401 @@ -516,7 +519,8 @@ def test_aiter_moe_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """The fused-MoE gate should depend only on the main AITER toggle and the - MoE-specific toggle.""" + MoE-specific toggle. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -547,7 +551,8 @@ def test_aiter_moe_shared_experts_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """Shared-expert fusion should only be enabled when the fused-MoE path is - enabled too.""" + enabled too. + """ from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -688,7 +693,8 @@ def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): """The gfx950 AITER MXFP4 W4A16 MoE path should match the dequantized - MXFP4 reference.""" + MXFP4 reference. + """ from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -737,7 +743,8 @@ def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_determinism(): """The gfx950 AITER MXFP4 W4A16 MoE path should stay bitwise - deterministic.""" + deterministic. + """ from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -830,7 +837,8 @@ def test_aiter_fused_moe_bf16_accuracy( intermediate_dim: int, ): """The ROCm AITER fused-MoE BF16 path should match the float32 reference - on representative shapes.""" + on representative shapes. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -875,7 +883,8 @@ def test_aiter_fused_moe_bf16_accuracy( def test_aiter_fused_moe_gelu_accuracy(): """The GELU activation variant should stay aligned with the float32 - reference.""" + reference. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -920,7 +929,8 @@ def test_aiter_fused_moe_gelu_accuracy(): def test_aiter_fused_moe_determinism(): """The BF16 fused-MoE kernel should stay bitwise deterministic for the - same inputs.""" + same inputs. + """ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -1075,7 +1085,8 @@ def test_aiter_fused_moe_mi3xx_bf16_accuracy(): ) def test_aiter_fused_moe_mi3xx_fp8_accuracy(): """The MI3xx FP8 per-tensor MoE path should stay within the measured FP8 - error budget.""" + error budget. + """ from tests.kernels.moe.utils import make_test_weights from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py index eec6180550dc..ba5eef662454 100644 --- a/tests/kernels/moe/test_routed_experts_capture_monolithic.py +++ b/tests/kernels/moe/test_routed_experts_capture_monolithic.py @@ -185,7 +185,8 @@ def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts( ) -> None: """The capture callback should receive the int16 routed-expert IDs the kernel actually used, the values should be valid expert indices, and - each token should pick ``top_k`` distinct experts.""" + each token should pick ``top_k`` distinct experts. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -264,7 +265,8 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( routing_method: RoutingMethodType, ) -> None: """Routing replay works for non-DeepSeekV3 routing methods too. - FlashInfer's ``routing_replay_out`` is routing-method-agnostic.""" + FlashInfer's ``routing_replay_out`` is routing-method-agnostic. + """ torch.manual_seed(0) device = torch.device("cuda:0") @@ -317,7 +319,8 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: """With no callback installed the kernel should not see a - ``routing_replay_out`` tensor — verify the helper short-circuits.""" + ``routing_replay_out`` tensor — verify the helper short-circuits. + """ torch.manual_seed(0) device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( @@ -339,7 +342,8 @@ def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: """FlashInfer's ``routing_replay_out`` is supported by all routing methods, so ``supports_routing_replay_capture`` should be True - regardless of routing method.""" + regardless of routing method. + """ device = torch.device("cuda:0") for routing_method in ( RoutingMethodType.DeepSeekV3, @@ -361,7 +365,8 @@ def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None: """When capture is installed, the allocated buffer is int16 and shaped - ``(num_tokens, experts_per_token)``.""" + ``(num_tokens, experts_per_token)``. + """ device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( num_experts=_DSV3_NUM_EXPERTS, @@ -592,7 +597,8 @@ def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale - path, DSV3 routing) captures valid expert IDs.""" + path, DSV3 routing) captures valid expert IDs. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -778,7 +784,8 @@ def _run_nvfp4_monolithic( routing_bias: torch.Tensor, ) -> torch.Tensor: """The monolithic NVFP4 apply expects packed fp4 hidden states + the - matching fp8 per-block scale stored in the ``a1q_scale`` slot.""" + matching fp8 per-block scale stored in the ``a1q_scale`` slot. + """ # Stash the weight tensors on the experts in the locations the apply() # implementation reads from (it pulls them from quant_config / scales # already; w1/w2 come in as args). @@ -806,7 +813,8 @@ def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs - on the DSV3 routing path.""" + on the DSV3 routing path. + """ if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 0eded09b1ae5..f922a3ebd336 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -214,8 +214,7 @@ def assert_routing_results_close( rtol: float = 1e-3, atol: float = 1e-3, ): - """ - Compare routing results, sorting by expert ID first to handle non-deterministic + """Compare routing results, sorting by expert ID first to handle non-deterministic ordering from sorted=False in topk. """ # Sort both results by expert IDs for consistent comparison @@ -251,7 +250,8 @@ def assert_aiter_routing_valid( the Python baseline (different group selection, scoring internals), so numerical comparison is not meaningful. Instead we verify the outputs satisfy the routing contract: correct shapes, valid expert - IDs, non-negative weights, and proper normalization.""" + IDs, non-negative weights, and proper normalization. + """ n_tokens = topk_weights.shape[0] # Shape @@ -294,8 +294,7 @@ def assert_aiter_routing_valid( def baseline_fused_topk( router_logits: torch.Tensor, top_k: int, renormalize: bool ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for standard fused top-k routing. + """Baseline for standard fused top-k routing. Algorithm: 1. Apply softmax to router logits @@ -320,8 +319,7 @@ def baseline_fused_topk_bias( e_score_correction_bias: torch.Tensor, routed_scaling_factor: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for fused top-k with bias correction. + """Baseline for fused top-k with bias correction. Algorithm: 1. Apply softmax to router logits @@ -364,8 +362,7 @@ def baseline_grouped_topk( e_score_correction_bias: torch.Tensor | None, routed_scaling_factor: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for grouped top-k routing (e.g., DeepSeek). + """Baseline for grouped top-k routing (e.g., DeepSeek). Algorithm: 1. Apply scoring function (softmax or sigmoid) @@ -436,8 +433,7 @@ def baseline_grouped_topk( def baseline_custom_llama4( router_logits: torch.Tensor, top_k: int ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Baseline for Llama4 custom routing. + """Baseline for Llama4 custom routing. Algorithm: 1. Select top-k expert indices (without softmax) diff --git a/tests/kernels/moe/test_routing_simulator.py b/tests/kernels/moe/test_routing_simulator.py index 35b1af618cf7..b1adcd53d93d 100644 --- a/tests/kernels/moe/test_routing_simulator.py +++ b/tests/kernels/moe/test_routing_simulator.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test script for the token-to-expert routing simulator. +"""Test script for the token-to-expert routing simulator. This script demonstrates how to use the routing simulator to test different routing strategies and analyze their performance, including @@ -77,7 +76,8 @@ def test_basic_functionality( def test_routing_strategy_integration(monkeypatch, device): """Test that the routing strategy environment variable works with - FusedMoEFactory.""" + FusedMoEFactory. + """ pytest.importorskip("vllm.model_executor.layers.fused_moe.layer") import vllm.envs as envs @@ -158,7 +158,8 @@ def test_routing_strategy_integration(monkeypatch, device): def test_distribution_based_routing_with_custom_strategy(): """Test registering and using DistributionBasedRouting with custom - parameters.""" + parameters. + """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Register custom distribution-based strategy diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index af9455b96039..895e168b2da3 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for FusedMoEFactory with routed_input_transform. +"""Tests for FusedMoEFactory with routed_input_transform. Verifies that applying routed_input_transform inside FusedMoEFactory produces the same results as applying the transform manually outside. diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py index 85b992ad6544..0ca08c50b324 100644 --- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py +++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py @@ -142,9 +142,7 @@ def silu_mul_quant( def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: - """ - pack float32 scales into a int32 tensor - """ + """Pack float32 scales into a int32 tensor.""" assert x.dtype == torch.float32 E, T, G = x.size() @@ -170,8 +168,7 @@ def ref_with_scale_fmt( up: torch.Tensor, scale_fmt: DeepGemmQuantScaleFMT, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - The precision types of the operations triggered by this function + """The precision types of the operations triggered by this function match closely with the kernel implementation so we compare more accurately. """ @@ -204,8 +201,7 @@ def ref_with_scale_fmt( def token_random(E, T, H2, tokens_per_expert): - """ - Initialize each token in a random range so we test a range of + """Initialize each token in a random range so we test a range of scale values. """ y = torch.empty((E, T, H2), dtype=torch.bfloat16, device=DEVICE) diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index bfe763b44ab2..76744da8f4df 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -19,11 +19,9 @@ def reference_quant(x: torch.Tensor, use_ue8m0: bool): - """ - Reference triton quant kernel from, + """Reference triton quant kernel from, vllm.model_executor.layers.quantization.utils.fp8_utils """ - x_q = torch.empty_like(x, device=x.device, dtype=FLOAT8_DTYPE) # Allocate the scale tensor in column-major format. @@ -72,7 +70,8 @@ def reference_with_clamp( ) -> tuple[torch.Tensor, torch.Tensor]: """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to match the C++ compute() template, then run the standard silu_and_mul + - quant reference.""" + quant reference. + """ N_2 = x.size(1) // 2 dtype = x.dtype gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype) diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py index 01246325c803..15cef0df38ee 100644 --- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py +++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py @@ -28,7 +28,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization""" + quantization and per-column weight quantization + """ A = A.to(torch.float32) B = B.to(torch.float32) @@ -56,8 +57,8 @@ def fp8_mask(a, mask): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk): """This function performs fused moe with per-column int8 - quantization using native torch.""" - + quantization using native torch. + """ B, D = a.shape # Perform per-token quantization a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True) diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index 7b6b7d81ffcf..9e50b4fc3bab 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/tests/kernels/moe/test_trtllm_bf16_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the FlashInfer TRTLLM BF16 MoE backend +"""Tests for the FlashInfer TRTLLM BF16 MoE backend (`TrtLlmBf16ExpertsModular`). This mirrors the TRTLLM NvFP4 modular test shape: construct the modular diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 59eaa4c1390c..b9684396b245 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the FlashInfer TRTLLM NvFP4 MoE backend +"""Tests for the FlashInfer TRTLLM NvFP4 MoE backend (`TrtLlmNvFp4ExpertsModular`). Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated), diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index d4c037d978a3..48e98f9c1457 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -150,7 +150,8 @@ def test_select_default_backend_by_platform( expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled.""" + accelerators (FlashInfer, AITER) disabled. + """ with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -444,7 +445,8 @@ def test_select_cuda_flashinfer_cutlass_backend( mock_has_flashinfer, ): """Test CUDA backend selection when FlashInfer TRTLLM is not available - and FlashInfer CUTLASS is available.""" + and FlashInfer CUTLASS is available. + """ with ( patch.object(current_platform, "is_cuda", return_value=True), patch.object(current_platform, "is_rocm", return_value=False), diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index c2084800f979..0202b380c526 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -83,8 +83,7 @@ def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_token @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): - """ - Run a forward pass through FusedMoEFactory with zero experts + """Run a forward pass through FusedMoEFactory with zero experts and verify output shape. """ layer, vllm_config = zero_expert_moe diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index e91fa58f3e81..e7624729502d 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -64,8 +64,7 @@ def make_dummy_moe_config( max_num_tokens: int = 512, activation: MoEActivation = MoEActivation.SILU, ) -> FusedMoEConfig: - """ - This is a dummy config for the mk constructor interface + """This is a dummy config for the mk constructor interface as most kernels like DeepGEMM, CUTLASSFp4, Triton, MARLIN do not actually use this config. diff --git a/tests/kernels/quant_utils.py b/tests/kernels/quant_utils.py index 7e14ec95d367..93a21596ca70 100644 --- a/tests/kernels/quant_utils.py +++ b/tests/kernels/quant_utils.py @@ -159,7 +159,8 @@ def native_per_token_group_quant_fp8( x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn ): """Function to perform per-token-group quantization on an input tensor - `x` using native torch.""" + `x` using native torch. + """ assert x.shape[-1] % group_size == 0, ( "the last dimension of `x` must be divisible by `group_size`" ) diff --git a/tests/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py index e0120301f7e3..04f7943ae630 100644 --- a/tests/kernels/quantization/nvfp4_utils.py +++ b/tests/kernels/quantization/nvfp4_utils.py @@ -118,6 +118,7 @@ def dequant_nvfp4_kv_cache( Returns: [..., num_heads, block_size, head_size] float32. + """ data_dim = head_size // 2 scale_dim = head_size // 16 diff --git a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py index c855d0e819dc..70abcbfc8ea4 100644 --- a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py +++ b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py @@ -304,7 +304,8 @@ def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel) def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): """Kernel output should match a dequantized fp32 reference within - fp8 per-token / per-channel quantization noise.""" + fp8 per-token / per-channel quantization noise. + """ weight_shape = (512, 4096) # (N, K) num_tokens = 32 _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index b70a1a8543ba..e4dadb53bd66 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -45,6 +45,7 @@ def quantize_weight_block_fp8( Returns: fp8_weight: [N, K] float8_e4m3fn scales: [n_tiles, k_tiles] float32 + """ N, K = weight.shape block_n, block_k = block_size diff --git a/tests/kernels/quantization/test_cutlass_scaled_mm.py b/tests/kernels/quantization/test_cutlass_scaled_mm.py index a522d3171450..9e1d6d2cc76c 100644 --- a/tests/kernels/quantization/test_cutlass_scaled_mm.py +++ b/tests/kernels/quantization/test_cutlass_scaled_mm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for cutlass kernels +"""Tests for cutlass kernels. Run `pytest tests/kernels/quantization/test_cutlass_scaled_mm.py`. """ diff --git a/tests/kernels/quantization/test_cutlass_w4a8.py b/tests/kernels/quantization/test_cutlass_w4a8.py index 8cfc993fe8e8..33bb7bccc74e 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8.py +++ b/tests/kernels/quantization/test_cutlass_w4a8.py @@ -300,8 +300,7 @@ def test_w4a8_cuda_graph(): ) @pytest.mark.parametrize("shape", MNK_SHAPES) def test_convert_packed_uint4b8_to_signed_int4_inplace(shape): - """ - The W4A16 checkpoints encode the weights as int4b8 packed to int32. + """The W4A16 checkpoints encode the weights as int4b8 packed to int32. The CUTLASS kernels expect signed int4 packed to int32. This tests checks that the runtime int4b8 -> signed int4 conversion matches the offline conversion step exactly. diff --git a/tests/kernels/quantization/test_cutlass_w4a8_moe.py b/tests/kernels/quantization/test_cutlass_w4a8_moe.py index 5e6c170db644..10a76e553e25 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8_moe.py +++ b/tests/kernels/quantization/test_cutlass_w4a8_moe.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer. -""" +"""Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer.""" import random from dataclasses import dataclass @@ -37,8 +35,7 @@ def cutlass_quantize( group_size: int | None, zero_points: bool = False, ): - """ - Quantize weights into W4 and compute reference dequantized weights. + """Quantize weights into W4 and compute reference dequantized weights. Encoding/reordering of weights and packing of scales is deferred until after all experts are combined. @@ -66,13 +63,13 @@ def cutlass_quantize( def cutlass_preprocess( w_q_experts: list[torch.Tensor], w_s_experts: list[torch.Tensor] ): - """ - Reorder/encode expert weights and pack scales. + """Reorder/encode expert weights and pack scales. Returns: w_q_packed: Packed/encoded int4 weights for all experts. w_s_packed: Packed fp8 scales for all experts. packed_layout: Layout/stride metadata for grouped GEMM. + """ w_s_packed = ops.cutlass_pack_scale_fp8(torch.stack(w_s_experts)) w_q_packed, packed_layout = ops.cutlass_encode_and_reorder_int4b_grouped( @@ -130,7 +127,6 @@ def make_moe_test_setup( random_zero: bool = False, ) -> MoETestSetup: """Create a full set of tensors for testing cutlass_w4a8_moe_mm.""" - assert K % GROUP_SIZE == 0 # Token counts per expert (multiples of `alignment`). Ms = [alignment * random.randint(1, max_blocks) for _ in range(num_experts)] @@ -276,9 +272,7 @@ def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero): class W4A8MoELayer(torch.nn.Module): - """ - Minimal wrapper module to test cuda graphs - """ + """Minimal wrapper module to test cuda graphs.""" def __init__(self, setup: MoETestSetup): super().__init__() diff --git a/tests/kernels/quantization/test_fp8_min_max_helper.py b/tests/kernels/quantization/test_fp8_min_max_helper.py index 8cd68a3fef7e..4f1beb857674 100644 --- a/tests/kernels/quantization/test_fp8_min_max_helper.py +++ b/tests/kernels/quantization/test_fp8_min_max_helper.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the get_fp8_min_max() helper function. +"""Unit tests for the get_fp8_min_max() helper function. These tests verify the FP8 min/max value logic for both standard and fnuz (ROCm MI300) dtype handling. diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index 704bb90f84a8..c01929c5d130 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -30,7 +30,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization""" + quantization and per-column weight quantization + """ A = A.to(torch.float32) B = B.to(torch.float32) @@ -53,8 +54,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids): """This function performs fused moe with per-column int8 quantization - using native torch.""" - + using native torch. + """ B, D = a.shape # Perform per-token quantization a_q, a_s = per_token_quant_int8(a) diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c96..f160125809e7 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -360,7 +360,8 @@ def test_gptq_marlin_padded_round_trip(shape, group_size): def test_fp8_block_marlin_padded_round_trip(shape): """Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers): group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the - weight_scale_inv group-wise scale padding.""" + weight_scale_inv group-wise scale padding. + """ size_n, size_k = shape block = 128 dtype = torch.float16 @@ -410,7 +411,8 @@ def test_fp8_block_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 288), (4640, 512)]) def test_mxfp8_marlin_padded_round_trip(shape): """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to - 2^-127 instead of zero and must still contribute nothing.""" + 2^-127 instead of zero and must still contribute nothing. + """ size_n, size_k = shape group_size = 32 # The e8m0-scale Marlin kernels are only instantiated for bf16 activations. @@ -457,7 +459,8 @@ def test_mxfp8_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 512), (4640, 512)]) def test_awq_zp_marlin_padded_round_trip(shape): """AWQ-style uint4 with runtime zero-points, padded the way - MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0.""" + MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0. + """ size_n, size_k = shape group_size = 128 dtype = torch.float16 @@ -692,7 +695,8 @@ def make_layer(hidden, intermediate): @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_fp8_marlin_moe_padded_round_trip(shape, quant): """FP8 weight-only MoE: pad a tile-misaligned intermediate and check the - real prepare + fused_marlin_moe against the dequantized reference.""" + real prepare + fused_marlin_moe against the dequantized reference. + """ from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk @@ -777,7 +781,8 @@ def q(w): # (out, in) -> fp8 weight, scale, dequant reference @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_mxfp8_marlin_moe_padded_round_trip(shape): """MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with - unit e8m0 scales so the reference is the exact fp8 dequant.""" + unit e8m0 scales so the reference is the exact fp8 dequant. + """ from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 28ce615ce655..de9feb7dfa05 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for MXFP4 linear kernel selection logic (CPU-only) +"""Tests for MXFP4 linear kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`. """ @@ -72,7 +72,8 @@ def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an - explicit request for a different activation format must be rejected.""" + explicit request for a different activation format must be rejected. + """ config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) can_implement, reason = kernel_cls.can_implement(config) assert not can_implement @@ -87,7 +88,8 @@ def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( """Marlin/Humming never quantize activations, so an unset activation key, or one that already describes MXFP4-shaped data, is tolerated. When an activation key is explicitly set, a warning must be logged noting that it - is ignored, since these kernels are weight-only (A16).""" + is ignored, since these kernels are weight-only (A16). + """ config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: can_implement, reason = kernel_cls.can_implement(config) @@ -116,7 +118,8 @@ def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): ) def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): """EmulationMxfp4LinearKernel is the universal fallback: it must accept - every supported activation format.""" + every supported activation format. + """ # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -134,7 +137,8 @@ def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): """quant_dequant_func must be derived purely from the config's activation - QuantKey, not set externally.""" + QuantKey, not set externally. + """ # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -158,7 +162,8 @@ def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): def test_aiter_kernel_is_supported_requires_native_mx_support(): """AiterMxfp4LinearKernel must not be selected on platforms without - native MX compute, even if AITER itself is importable.""" + native MX compute, even if AITER itself is importable. + """ with patch( "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", return_value=False, @@ -194,7 +199,8 @@ def apply_weights( @patch("vllm.model_executor.kernels.linear.current_platform") def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp4_linear_kernel should select a registered kernel that - reports itself as supported, and construct it with a fresh config.""" + reports itself as supported, and construct it with a fresh config. + """ platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py index 93d30f459816..74996d64e368 100644 --- a/tests/kernels/quantization/test_mxfp4_triton_ep.py +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests that triton_kernel_moe_forward correctly applies expert_map +"""Tests that triton_kernel_moe_forward correctly applies expert_map remapping when expert parallelism (EP) is enabled. When expert_map is provided, global expert IDs are remapped to local IDs @@ -17,7 +16,8 @@ class TestTritonMoeForwardExpertMap: """Test that triton_kernel_moe_forward applies expert_map remapping - when expert_map is provided (EP active).""" + when expert_map is provided (EP active). + """ def test_expert_map_remap(self): device = "cuda" if torch.cuda.is_available() else "cpu" diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py index 2c7b565bea57..2c0e94c0f519 100644 --- a/tests/kernels/quantization/test_mxfp6_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for MXFP6 linear kernel selection logic (CPU-only) +"""Tests for MXFP6 linear kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`. """ @@ -41,7 +41,8 @@ def test_can_implement_is_abstract(): def test_emulation_kernel_rejects_non_mxfp6_weights(): """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight - format.""" + format. + """ config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static) can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) assert not can_implement @@ -57,7 +58,8 @@ def test_emulation_kernel_accepts_any_supported_config( weight_quant_key, activation_quant_key ): """EmulationMxfp6LinearKernel is the only backend today: it must accept - every supported weight/activation format combination.""" + every supported weight/activation format combination. + """ config = MxFp6LinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key ) @@ -102,7 +104,8 @@ def apply_weights( def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp6_linear_kernel should select a registered kernel that reports itself as supported/able to implement the given config, and - construct it with that exact config.""" + construct it with that exact config. + """ platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6") diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index eb057cebd912..33295b8784f3 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -64,8 +64,7 @@ def loaded_model_files(): class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts): - """ - Extension of TritonExperts to support emulated NVFP4 MoE experts. + """Extension of TritonExperts to support emulated NVFP4 MoE experts. It may be used for NVFP4 models when the device does not have native support for this dtype. diff --git a/tests/kernels/quantization/test_nvfp4_quant.py b/tests/kernels/quantization/test_nvfp4_quant.py index 120dc2774a5d..b8a1cdda4612 100644 --- a/tests/kernels/quantization/test_nvfp4_quant.py +++ b/tests/kernels/quantization/test_nvfp4_quant.py @@ -176,8 +176,7 @@ def test_python_util_matches_cpp_allocation( shape: tuple[int, int], is_sf_swizzled_layout: bool, ) -> None: - """ - Verify that the Python utility (create_fp4_output_tensors) allocates + """Verify that the Python utility (create_fp4_output_tensors) allocates tensors with the same shapes and dtypes as the C++ functional variant (scaled_fp4_quant_func). """ diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index f2326bfd25fc..c1226efcc803 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -118,8 +118,8 @@ def test_per_token_group_quant_fp8_packed( num_tokens, hidden_dim, group_size, poisoned_scales ): """Test the packed DeepGEMM quantization kernel against the Triton - reference (row-major, UE8M0 scales).""" - + reference (row-major, UE8M0 scales). + """ device = "cuda" torch.manual_seed(42) @@ -206,7 +206,6 @@ def test_per_token_group_quant_fp8_packed_all_zero(): For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to 1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94). """ - device = "cuda" num_tokens, hidden_dim, group_size = 4, 7168, 128 x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) @@ -258,8 +257,8 @@ def test_per_token_group_quant_fp8_packed_all_zero(): def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): """Inputs whose absmax/max_8bit produces a non-power-of-2 force the mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior - before optimization.""" - + before optimization. + """ device = "cuda" num_tokens, hidden_dim, group_size = 4, 7168, 128 @@ -331,8 +330,8 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( ): """When output_q is allocated with shape (tma_aligned_mn, k) instead of (mn, k), the kernel must overwrite the padded mn rows with zeros so - callers can use ``torch.empty`` instead of ``torch.zeros``.""" - + callers can use ``torch.empty`` instead of ``torch.zeros``. + """ device = "cuda" group_size = 128 torch.manual_seed(42) @@ -396,7 +395,6 @@ def test_per_token_group_quant_fp8_packed_large_mn(): This is a differential test that compares fp8 output against Triton output reference when token size sits just above the gridDim.y 2^16 - 1 limit. """ - device = "cuda" group_size = 128 # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn diff --git a/tests/kernels/quantization/test_quantized_embedding.py b/tests/kernels/quantization/test_quantized_embedding.py index 0e4af0a0c1a6..35c0e74713cb 100644 --- a/tests/kernels/quantization/test_quantized_embedding.py +++ b/tests/kernels/quantization/test_quantized_embedding.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the Triton dequant-gather kernel used by -``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup).""" +``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup). +""" import pytest import torch diff --git a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py index c7f0f1eda355..2a5a8081a2f5 100644 --- a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py +++ b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py @@ -29,7 +29,8 @@ def test_rocm_aiter_group_fp8_quant_fake_implementation(): """Test that the fake implementation is correctly - defined for torch.ops.vllm.rocm_aiter_group_fp8_quant.""" + defined for torch.ops.vllm.rocm_aiter_group_fp8_quant. + """ # Create test tensors M = 128 N = 4096 @@ -49,7 +50,8 @@ def test_rocm_aiter_group_fp8_quant_fake_implementation(): def test_rocm_aiter_group_fp8_quant_torch_compile_with_cudagraph(): """Test that rocm_aiter_ops.group_fp8_quant with group size 128 can be used with - torch.compile in cudagraph mode.""" + torch.compile in cudagraph mode. + """ # Create test tensors M = 128 N = 4096 diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 43b5daa45a7e..7002165af093 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -209,7 +209,8 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch( @pytest.mark.parametrize("shape", LLAMA_SHAPES) def test_vllm_quant_dequant_mxfp4_matches_reference_on_large_shapes(shape, dtype): """The public vLLM MXFP4 QDQ helper should stay exact against the torch - reference even on larger Llama-like shapes.""" + reference even on larger Llama-like shapes. + """ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( quant_dequant_mxfp4, ) @@ -504,7 +505,7 @@ def run_gemm(): ], ) def test_aiter_hardware_fp4_dynamic_quant_format(shape): - """aiter hardware FP4 dynamic quant produces correct output format. + """Aiter hardware FP4 dynamic quant produces correct output format. Tests gfx950 hardware-accelerated FP4 quantization (OCP MXFP4 E2M1). Parity with B200 scaled_fp4_quant: block_size=32, packed uint8 output. diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index 827571474481..d682ed05b0d8 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for ScaledMM kernel selection logic (CPU-only) +"""Tests for ScaledMM kernel selection logic (CPU-only). Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`. """ diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index e829c63a8462..55d1f76b2d1e 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the triton_scaled_mm kernel +"""Tests for the triton_scaled_mm kernel. Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`. """ diff --git a/tests/kernels/test_apply_repetition_penalties.py b/tests/kernels/test_apply_repetition_penalties.py index 8270cf885f60..c1551e2e82e6 100644 --- a/tests/kernels/test_apply_repetition_penalties.py +++ b/tests/kernels/test_apply_repetition_penalties.py @@ -35,8 +35,7 @@ def test_apply_repetition_penalties( dtype: torch.dtype, seed: int, ) -> None: - """ - Test the apply_repetition_penalties custom op + """Test the apply_repetition_penalties custom op against a reference implementation. """ set_random_seed(seed) @@ -86,8 +85,7 @@ def test_apply_repetition_penalties( ) @torch.inference_mode() def test_apply_repetition_penalties_zero_seqs() -> None: - """ - Test the apply_repetition_penalties custom op with num_seqs=0 + """Test the apply_repetition_penalties custom op with num_seqs=0 against a reference implementation. """ num_seqs = 0 diff --git a/tests/kernels/test_awq_int4_to_int8.py b/tests/kernels/test_awq_int4_to_int8.py index f94214785d22..dc7124e9ae70 100644 --- a/tests/kernels/test_awq_int4_to_int8.py +++ b/tests/kernels/test_awq_int4_to_int8.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration). +"""Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration). Part 1: Weight packing tests - convert_weight_packed_scale_zp correctness @@ -44,6 +43,7 @@ def make_awq_checkpoint_data(K, N, group_size, seed=42): float_ref: [K, N] float32, reference dequantized weights weight_int4_orig: [K, N] int32, original int4 values (0-15) zeros_int4_orig: [num_groups, N] int32, original zero points (0-15) + """ rng = np.random.RandomState(seed) num_groups = K // group_size diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py index eff4508b0d83..a3cb91242af4 100644 --- a/tests/kernels/test_cache_kernels.py +++ b/tests/kernels/test_cache_kernels.py @@ -15,12 +15,10 @@ @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="Need CUDA device") def test_gather_cache_oob(): - """ - Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909). + """Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909). This test constructs a boundary case identified in the issue where seq_starts causes the block_table offset to read out of bounds. """ - block_size = 64 # The kernel only supports the MLA entry sizes. entry_size = 576 diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 73c2027a7866..3dc658250fec 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant. +"""Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant. These tests cover: A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64 @@ -671,7 +670,8 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): def test_v41_compressor_metadata_maps_tokens_to_their_ring(): """The ring group's generic slot mapping is disabled (all PAD), so the builder must map every real token to ``ring_block * capacity + pos % - capacity`` and keep padding tokens at PAD.""" + capacity`` and keep padding tokens at PAD. + """ from unittest.mock import MagicMock from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder @@ -1096,8 +1096,8 @@ def test_get_c128_boundary(starts, query_start_loc, expected): @pytest.mark.parametrize("block_size", [16, 64]) def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int): """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache - → compare against original.""" - + → compare against original. + """ HEAD_DIM = 512 NOPE_DIM = 448 HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale @@ -1308,9 +1308,9 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): - """k → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache - → manual dequant → compare against original.""" - + """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache + → manual dequant → compare against original. + """ HEAD_DIM = 128 QUANT_BLOCK_SIZE = 128 # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132 @@ -1374,7 +1374,6 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): def test_indexer_gather_accepts_upper_bound_output(): """Gather only exact cu_seq_lens even when dst is over-allocated.""" - head_dim = 128 quant_block_size = 128 cache_stride = head_dim + head_dim * 4 // quant_block_size @@ -1474,7 +1473,6 @@ def test_indexer_gather_accepts_upper_bound_output(): def test_deepseek_v4_quant_magnitude_range(): """Test that quantization handles a range of magnitudes correctly.""" - HEAD_DIM = 512 NOPE_DIM = 448 HEAD_BYTES = 584 diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index fe539abb4f72..d9157c411066 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -29,6 +29,7 @@ def _build_test_case(seq_lens, block_size, seed=42): Returns: Tuple of (cache, block_table, workspace_starts_t, num_reqs, total_tokens, expected_output). + """ torch.manual_seed(seed) @@ -431,7 +432,8 @@ def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts): ) def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): """Correctness test with large sequence lengths matching benchmark - scenarios (8K-128K prefill).""" + scenarios (8K-128K prefill). + """ ( cache, block_table, diff --git a/tests/kernels/test_deepseek_v4_cpu_kernels.py b/tests/kernels/test_deepseek_v4_cpu_kernels.py index 81c03a0bca9f..9c3fddb6c171 100644 --- a/tests/kernels/test_deepseek_v4_cpu_kernels.py +++ b/tests/kernels/test_deepseek_v4_cpu_kernels.py @@ -85,7 +85,8 @@ def _decode_token_regions( cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int ): """Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte - regions this layout defines.""" + regions this layout defines. + """ valid = slot_mapping >= 0 idx = valid.nonzero(as_tuple=True)[0] slots = slot_mapping[idx] @@ -110,7 +111,8 @@ def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size): """NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an un-rotated bf16 value, so they match tightly; the RoPE region is rotated in fp32 before its bf16 store and eager/kernel can round a tie to - opposite sides, so allow <=1 ULP there.""" + opposite sides, so allow <=1 ULP there. + """ nope_ref, rope_ref, scale_ref = _decode_token_regions( cache_ref, slot_mapping, block_size ) @@ -190,7 +192,8 @@ def _dup_padded_block_cache( num_blocks: int, block_size: int, width: int, pad_elems: int ) -> tuple[torch.Tensor, torch.Tensor]: """Two independently-stored padded caches with identical values, for ops - that mutate the cache in place.""" + that mutate the cache in place. + """ page_stride = block_size * width + pad_elems template = torch.randn(num_blocks, page_stride, dtype=torch.float32) @@ -207,7 +210,8 @@ def _run_save_partial_states_eager( ): """Reference: state_cache[block_idx, pos_in_block, :W] = kv; state_cache[..., W:] = score + ape[position % compress_ratio]; skip - slot < 0. Mutates state_cache in place, matching the op's contract.""" + slot < 0. Mutates state_cache in place, matching the op's contract. + """ state_width = kv.shape[-1] block_size = state_cache.shape[1] valid = slot_mapping >= 0 @@ -267,7 +271,8 @@ def _run_compress_eager( ): """Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0 quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla - layout.""" + layout. + """ num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -444,7 +449,8 @@ def test_compress_norm_rope_store_cpu_matches_eager( def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride(): """The main (output) kv_cache is a separate physical buffer from the - state cache and can independently be a strided view.""" + state cache and can independently be a strided view. + """ eps = 1e-6 compress_ratio, overlap = 128, False num_tokens = 5 @@ -520,7 +526,8 @@ def _run_compress_indexer_eager( """Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512 path, but the indexer's cache quantizes all 128 post-RoPE dims as one FP8 block with a single raw fp32 scale per token (matching - _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks.""" + _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks. + """ num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -648,7 +655,8 @@ def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int): """Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and return (cache, dequant_ref) -- the exact post-quantization value the kernel's dequant should reproduce. Slot id == row index, packed - contiguously from block 0.""" + contiguously from block 0. + """ n = kv_f32.shape[0] num_blocks = max(1, -(-n // block_size)) cache = _make_cache(num_blocks, block_size) @@ -706,7 +714,8 @@ def _flash_mla_eager( """MQA attention: gather the (already-dequantized) window+compressed KV rows referenced by *_slots (-1 = invalid/masked), append a learned per-head sink logit with zero value contribution, softmax, weighted - sum.""" + sum. + """ num_tokens, num_heads, _ = q.shape q_f32 = q.float() @@ -805,7 +814,8 @@ def test_flash_mla_cpu_kernel_matches_eager( def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink(): """When every KV slot is masked out, softmax degenerates to the sink - logit alone and the output must be all-zero (no valid value rows).""" + logit alone and the output must be all-zero (no valid value rows). + """ torch.manual_seed(2) dtype = torch.bfloat16 num_tokens, num_heads, num_window = 3, 4, 5 @@ -845,7 +855,8 @@ def _qnorm_kv_insert_eager( q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs ): """Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE - + UE8M0 FP8 quant + paged insert on KV.""" + + UE8M0 FP8 quant + paged insert on KV. + """ num_quant_blocks = NOPE_DIM // QUANT_BLOCK num_tokens_full, num_heads_q, _ = q.shape @@ -971,7 +982,8 @@ def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager( @pytest.mark.parametrize("pad", [1, 5]) def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int): """slot_mapping shorter than q/kv rows: the KV branch must only touch - the first num_tokens rows while Q-norm+RoPE still runs on all rows.""" + the first num_tokens rows while Q-norm+RoPE still runs on all rows. + """ torch.manual_seed(1) dtype = torch.bfloat16 eps = 1e-6 @@ -1032,7 +1044,8 @@ def _inverse_gptj_rope_o_proj_eager( rope_dim: int, ) -> torch.Tensor: """NoPE dims pass through as fp32; the RoPE segment is de-rotated by - R(-pos) in GPT-J (interleaved even/odd) convention.""" + R(-pos) in GPT-J (interleaved even/odd) convention. + """ nope_dim = o.shape[-1] - rope_dim half_dim = rope_dim // 2 @@ -1098,7 +1111,8 @@ def _write_page( ) -> None: """Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size] fp32) into one page row: K-region then scale-region, matching - compress_norm_rope_store_indexer_cpu's layout.""" + compress_norm_rope_store_indexer_cpu's layout. + """ k_region = page_row[: block_size * INDEXER_HEAD_DIM] k_region.copy_(k_fp8.view(torch.uint8).flatten()) scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32) @@ -1134,7 +1148,8 @@ def _build_paged_mqa_logits_batch( ): """Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly non-contiguous in dim 0) with deterministic-per-seed random K/scale - data.""" + data. + """ torch.manual_seed(seed) batch_size = len(seq_lens) max_seq_len = max(seq_lens) @@ -1200,7 +1215,8 @@ def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride(): """The paged K-cache in production is a per-layer view into a shared multi-layer allocation, so its page stride generally exceeds the page byte width -- check the kernel reads a padded/strided buffer identically - to an equivalent tightly-packed one.""" + to an equivalent tightly-packed one. + """ block_size = 64 num_heads = 4 seq_lens = [1, block_size + 7, 3 * block_size] @@ -1326,7 +1342,8 @@ def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel(): CPU op on a CPU tensor, not silently fall through to the Triton kernel (which would also "work" under triton-cpu, masking a dispatch regression). The dispatcher does its CPU import inside the CPU branch, - so the patch target is the source module.""" + so the patch target is the source module. + """ torch.manual_seed(0) num_tokens = 1 @@ -1393,7 +1410,8 @@ def _hc_head_torch( ) -> torch.Tensor: """Eager reference for hc_head_fused_cpu: round-trips the RMSNorm output through bf16 before the linear, matching the kernel's bf16 - arithmetic.""" + arithmetic. + """ hc_mult, hidden_size = hidden_states.shape[-2:] outer_shape = hidden_states.shape[:-2] hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size) @@ -1475,7 +1493,8 @@ def test_mhc_post_cpu(num_tokens: int, hidden_size: int): @pytest.mark.parametrize("hidden_size", [4096, 7168]) def test_hc_head_cpu(num_tokens: int, hidden_size: int): """hc_head_fused_cpu vs. the _hc_head_torch eager reference it - replaces.""" + replaces. + """ torch.set_default_device(DEVICE) set_random_seed(0) hc_mult = 4 diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c2a7607cbac0..8deba1241ae0 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -383,7 +383,8 @@ def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False): """Run the op on a batch of (token_ids, start, window) requests; `window` is the runner's [depth] lookback for that request (None = all unknown). Token IDs 14 and 22 stand in for image tokens that break n-grams. - Returns one [len(token_ids), ...] hash tensor per request.""" + Returns one [len(token_ids), ...] hash tensor per request. + """ token_map, multipliers, primes, offsets, block_size = tables_meta depth = multipliers.shape[1] - 1 @@ -442,7 +443,8 @@ def test_lookback_window_reproduces_single_instance(runner, capture): V2 supplies every lookback from its device token history and needs no slot cache. V1 supplies prompt positions only and reads generated positions from the slot cache it fills itself. Without any window the - first decode token hashes stale slots (negative control).""" + first decode token hashes stale slots (negative control). + """ block_size = 4 block_table = torch.tensor( [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda" @@ -708,7 +710,8 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): """The fused gather must be bit-exact with the torch dequant path it replaces, from HBM and from pinned host memory alike, and must contribute - zeros for rows another TP rank owns.""" + zeros for rows another TP rank owns. + """ layer = _make_embedding(cpu_offload) cols, rows = 24, layer.part_num_embeddings ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py index 492c099bf1c5..fad8d90805eb 100644 --- a/tests/kernels/test_fla_layernorm_guard.py +++ b/tests/kernels/test_fla_layernorm_guard.py @@ -336,7 +336,8 @@ def test_layer_norm_rows_per_block( @torch.inference_mode() def test_strided_input(dtype: torch.dtype) -> None: """Test that the kernel handles non-contiguous (strided) - inputs correctly.""" + inputs correctly. + """ set_random_seed(42) device = torch.device(DEVICE) num_tokens = 128 diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index 12d00ec4054b..6862f8ffb181 100644 --- a/tests/kernels/test_flex_attention.py +++ b/tests/kernels/test_flex_attention.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Integration tests for FlexAttention backend vs default backend""" +"""Integration tests for FlexAttention backend vs default backend.""" import pytest import torch diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index d330abba532a..b6c2a53ed873 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -97,7 +97,8 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: """The gate feeds top-k expert selection: the kernel's top-8 must match an fp64 reference's top-8 per token (ties tolerated). This is the business-level correctness of the router — numeric error only matters - if it flips the argsort.""" + if it flips the argsort. + """ _requires_sm90() top_k = 8 device = torch.device("cuda") diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 7f74cde05d5b..6691de7ebfff 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -854,7 +854,8 @@ def test_fused_q_no_indexer(num_tokens: int): @pytest.mark.parametrize("has_indexer", [True, False]) def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool): """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is - produced (bf16, unquantized); ql_nope is consumed directly by the caller.""" + produced (bf16, unquantized); ql_nope is consumed directly by the caller. + """ torch.manual_seed(6) dev = "cuda" max_pos = 8192 diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index a0dcfc5bb7da..3e95913a017a 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel: +"""Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel: fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert - Q side: optional per-head RMSNorm + GPT-J RoPE on last 64 dims @@ -177,7 +176,8 @@ def _call_fused( def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor: """Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on - gfx942) encoding the kernel actually wrote, without touching the bytes.""" + gfx942) encoding the kernel actually wrote, without touching the bytes. + """ return t.contiguous().view(torch.uint8).view(FP8_STORE_DTYPE) @@ -713,7 +713,8 @@ def run(q_in, is_q_interleaved): @pytest.mark.parametrize("block_size", [16, 64]) def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int): """slot_mapping.size(0) < q.size(0): the kernel must skip padded - tokens in the KV branch while still processing Q on all rows.""" + tokens in the KV branch while still processing Q on all rows. + """ torch.manual_seed(3) device = "cuda" dtype = torch.bfloat16 diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 7464c9076973..20ac8f19f1a6 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -45,6 +45,7 @@ def quantize_to_mxfp4( Returns: packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index scales: [..., head_dim//32] uint8 1 ue8m0 byte + """ MXFP4_BLOCK_SIZE = 32 orig_shape = x.shape @@ -573,7 +574,8 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's sparse MQA-logits kernels take) instead of fp32; the values are the fp32 - result rounded once, so the scoring kernel needs no cast.""" + result rounded once, so the scoring kernel needs no cast. + """ if use_cutedsl and not has_cutedsl(): pytest.skip("cutedsl (cutlass) not installed") diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index d27144a38e9b..97ea8f7e3443 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel. +"""Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel. Tests compare the fused kernel against a reference implementation built from the existing separate operations (inverse RoPE via rotate_neox + FP8 quant @@ -132,6 +131,7 @@ def reference_inv_rope( Returns: o with inverse RoPE applied on the rope portion (bf16). + """ assert cos_sin_cache.dtype == torch.float32 cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32 @@ -164,6 +164,7 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten Returns: x_fp8: same shape, float8_e4m3fn scales: [...] float32, one scale per block + """ absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS) scale_raw = absmax * (1.0 / FP8_MAX) @@ -192,6 +193,7 @@ def reference_inv_rope_fp8_quant( Returns: o_fp8: [T, G, D] FP8 with strides (D, T*D, 1) o_scale: [T, G, S] FP32 with strides (S, T*S, 1) + """ assert cos_sin_cache.dtype == torch.float32 T, _H, head_dim = o.shape @@ -402,7 +404,6 @@ def test_output_strides(num_tokens, num_heads, n_groups): - FP8: logical [T, G, D] backed by contiguous [G, T, D]. - Scale: MN-major TMA-aligned (column-major: T-stride=1). """ - heads_per_group = num_heads // n_groups max_pos = 4096 device = "cuda" @@ -445,7 +446,8 @@ def test_output_strides(num_tokens, num_heads, n_groups): def test_per_group_contiguity(num_tokens): """FP8 per-group slices must be contiguous. Scale per-group slices are column-major (T-stride=1) — not row-major contiguous, which is - correct for TMA loads.""" + correct for TMA loads. + """ num_heads, n_groups = 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -511,7 +513,8 @@ def test_scales_are_power_of_two(): def test_nope_dims_unchanged(): """Nope dimensions (first 448 per head) should only be quantized, not rotated. Verify by dequantizing and comparing against - quantize-only reference (no RoPE).""" + quantize-only reference (no RoPE). + """ num_tokens, num_heads, n_groups = 16, 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -902,8 +905,8 @@ def calc_diff(x, y): @torch.inference_mode() def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config): """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style, - mscale=0, YaRN scaling) matching the production config.""" - + mscale=0, YaRN scaling) matching the production config. + """ num_heads = 64 n_groups = 8 heads_per_group = num_heads // n_groups diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 5fab150d9fce..219e42c74294 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -703,7 +703,7 @@ def mhc_pre_ref( hc_post_mult_value: float, sinkhorn_repeat: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """mHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303""" + """MHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303.""" hc_mult = residual.shape[-2] residual_flat = residual.flatten(-2, -1).float() @@ -742,7 +742,7 @@ def mhc_post_ref( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - """mHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68""" + """MHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68.""" term2 = torch.bmm(comb_res_mix.mT, residual.float()) return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16() @@ -1372,7 +1372,8 @@ def _patch_first_rank_pp_group(monkeypatch): def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): """First finalize (at the end of load_weights) allocates - hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" + hc_attn_fn_broadcast as hc_attn_fn summed over hc streams. + """ _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) @@ -1387,7 +1388,8 @@ def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): """Re-finalizing after a weight refit must copy into the existing broadcast tensor so its address stays stable for captured CUDA graphs, - while picking up the new hc_attn_fn values.""" + while picking up the new hc_attn_fn values. + """ _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 997c61980913..45fafb0433ca 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -403,7 +403,8 @@ def test_mxfp8_linear_emulation_bf16_at_load( ): """EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the ``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the - same result; the dtype-match (BF16/FP16 activations) must also hold.""" + same result; the dtype-match (BF16/FP16 activations) must also hold. + """ from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) @@ -517,7 +518,8 @@ def _spy(*a, **k): # produced EP garbage). def _capture_expert_mask(expert_mask, *, global_num_experts): """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the - ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``.""" + ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``. + """ from types import SimpleNamespace from unittest import mock @@ -568,7 +570,8 @@ def _fake_fused_moe(hidden_states, w1, w2, tw, ti, *, expert_mask, **kw): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is - (no re-derivation that would collapse an already-0/1 mask to all-ones).""" + (no re-derivation that would collapse an already-0/1 mask to all-ones). + """ # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3). ep_mask = torch.tensor( [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE @@ -582,7 +585,8 @@ def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): def test_routed_experts_expert_map_delegates_to_kernel(): """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed - on the resolved kernel, not the global aiter switch.""" + on the resolved kernel, not the global aiter switch. + """ from types import SimpleNamespace from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts diff --git a/tests/kernels/test_shuffle_rows.py b/tests/kernels/test_shuffle_rows.py index 08d5883a1e70..35c2f2644290 100644 --- a/tests/kernels/test_shuffle_rows.py +++ b/tests/kernels/test_shuffle_rows.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the shuffle_rows function +"""Tests for the shuffle_rows function. Run `pytest tests/kernels/test_shuffle_rows.py`. """ @@ -17,7 +17,8 @@ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_shuffle_rows_basic(num_tokens: int, hidden_size: int, dtype: torch.dtype): """Test basic functionality of shuffle_rows with various tensor sizes and - dtypes.""" + dtypes. + """ if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") @@ -72,7 +73,8 @@ def test_shuffle_rows_permutation( @pytest.mark.parametrize("hidden_size", [256, 512]) def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int): """Test shuffle_rows with expansion (more output tokens than input - tokens).""" + tokens). + """ if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 289539470631..f249adeaffec 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -157,8 +157,7 @@ def compare_top_k_results( top_k: int, tolerance: float = 1e-5, ) -> bool: - """ - Compare results from CUDA top_k_per_row with torch.topk. + """Compare results from CUDA top_k_per_row with torch.topk. Both results should be sorted and contain the same top-k elements. """ num_rows = cuda_indices.shape[0] @@ -213,8 +212,7 @@ def validate_topk_against_reference( top_k: int, kernel_name: str, ) -> None: - """ - Validate CUDA top-k results against PyTorch reference implementation. + """Validate CUDA top-k results against PyTorch reference implementation. Args: logits: Input logits tensor @@ -223,6 +221,7 @@ def validate_topk_against_reference( row_ends: Row end positions top_k: Number of top elements to select kernel_name: Name of the kernel being tested (for error messages) + """ num_rows = cuda_indices.shape[0] torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") @@ -248,9 +247,7 @@ def test_top_k_per_row( top_k: int, clean_logits: bool, ) -> None: - """ - Test top_k_per_row. - """ + """Test top_k_per_row.""" set_random_seed(0) torch.set_default_device("cuda:0") @@ -298,9 +295,7 @@ def _run_top_k_per_row_decode_test( clean_logits: bool, data_generation: str, ) -> None: - """ - Helper function to run top_k_per_row_decode test with given parameters. - """ + """Helper function to run top_k_per_row_decode test with given parameters.""" torch.set_default_device("cuda:0") # Create test data @@ -365,9 +360,7 @@ def test_top_k_per_row_decode( clean_logits: bool, data_generation: str, ) -> None: - """ - Test top_k_per_row with seq_lens tensor. - """ + """Test top_k_per_row with seq_lens tensor.""" set_random_seed(0) vocab_size = 20000 _run_top_k_per_row_decode_test( @@ -379,9 +372,7 @@ def test_top_k_per_row_decode( @pytest.mark.parametrize("clean_logits", [True, False]) @torch.inference_mode() def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None: - """ - Test top_k_per_row_decode with large vocabulary size. - """ + """Test top_k_per_row_decode with large vocabulary size.""" set_random_seed(0) top_k = 2048 batch_size = 2 @@ -886,8 +877,7 @@ def test_deepseek_workspace_topk( next_n: int, backend: str, ) -> None: - """ - Test workspace top-k backends with varying sequence lengths and speculative + """Test workspace top-k backends with varying sequence lengths and speculative decoding. Supports speculative decoding with next_n > 1. """ @@ -1040,8 +1030,7 @@ def run_large_context_topk_test( seed: int = 42, backend: str = "cooperative_topk", ) -> None: - """ - Helper to run a top-k backend test with given parameters. + """Helper to run a top-k backend test with given parameters. Args: batch_size: Number of rows/sequences @@ -1050,6 +1039,7 @@ def run_large_context_topk_test( data_type: Type of test data to generate seed: Random seed for reproducibility backend: Top-k backend to test + """ torch.set_default_device("cuda:0") set_random_seed(seed) @@ -1274,8 +1264,7 @@ def test_cooperative_topk_cs2(top_k: int) -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: - """ - Comprehensive correctness tests covering: + """Comprehensive correctness tests covering: - Sequence length edge cases (trivial, boundary, varied) - Very small sequences (< 100 elements) - Mixed sequence lengths in same batch @@ -1339,8 +1328,7 @@ def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None: - """ - Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): + """Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): - Batch size scalability (1, 4, 32, 256) - Single-CTA vs Multi-CTA execution - Extreme configurations (large batch, max context length) @@ -1357,8 +1345,7 @@ def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_stress(backend: str) -> None: - """ - Stress test with random configurations to catch edge cases. + """Stress test with random configurations to catch edge cases. Capped at 163840 (DeepSeek V3.2 max context) for realistic testing. """ torch.set_default_device("cuda:0") @@ -1505,8 +1492,7 @@ def test_cooperative_topk_512_tie_workspace_is_per_row() -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None: - """ - Tests specific to workspace top-k backends: + """Tests specific to workspace top-k backends: - Mixed medium/large rows in the same batch (dynamic per-row dispatch) - Boundary around LARGE_THRESHOLD (32K) - Trivial + medium + large rows in a single batch @@ -1567,8 +1553,7 @@ def test_persistent_topk_reused_group_after_short_row() -> None: @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: - """ - Test workspace top-k backends with padded logits (large stride, small seq_len) + """Test workspace top-k backends with padded logits (large stride, small seq_len) to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits returns [B, max_model_len] with max_model_len=163840. """ @@ -1755,7 +1740,8 @@ def test_sparse_indexer_decode_topk_explicit_backends( ) -> None: """Every explicit sparse-indexer decode top-k backend must restrict each row to its seq_len (dirty data past the end must never be selected) and - match torch.topk on the valid region.""" + match torch.topk on the valid region. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk @@ -1826,7 +1812,8 @@ def test_sparse_indexer_decode_topk_short_seq_lens( backend: str, workspace_init ) -> None: """seq_len < next_n: derived per-row ends must clamp at 0 (the reference - kernels do max(0, ...)); negative ends are OOB for the ragged kernels.""" + kernels do max(0, ...)); negative ends are OOB for the ragged kernels. + """ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk diff --git a/tests/kernels/utils.py b/tests/kernels/utils.py index 8bd63872f3d6..cf68565e9d5b 100644 --- a/tests/kernels/utils.py +++ b/tests/kernels/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kernel test utils""" +"""Kernel test utils.""" import itertools import random @@ -94,16 +94,15 @@ def _collect(result: Any) -> list[torch.Tensor]: class QKVInputs(NamedTuple): - """ - Data structure for representing unpacked attention inputs, + """Data structure for representing unpacked attention inputs, query/key/values and their sequence lengths. Attributes: - * {query,key,value}: unpacked (batch_size x padded_seq_len x num_heads x head_size) attention inputs * q_seq_lens: query sequence lengths list * kv_seq_lens: shared key/value sequence lengths list + """ query: torch.Tensor @@ -114,16 +113,15 @@ class QKVInputs(NamedTuple): class QKVO(NamedTuple): - """ - Data structure for representing unpacked attention inputs, + """Data structure for representing unpacked attention inputs, alongside unpacked known-correct attention output Attributes: - * qkv: unpacked (batch_size x padded_seq_len x num_heads x head_size) attention inputs * ideal_output: unpacked (batch_size x padded_seq_len x num_heads x head_size) known-correct attention output + """ qkv: QKVInputs @@ -131,11 +129,9 @@ class QKVO(NamedTuple): class PackedQKVInputs(NamedTuple): - """ - Data structure for representing packed attention inputs + """Data structure for representing packed attention inputs. Attributes: - * {query,key,value}: packed (number_of_tokens x num_heads x head_size) attention inputs * q_start_loc_list: list of query start locations within packed tensor @@ -143,6 +139,7 @@ class PackedQKVInputs(NamedTuple): packed tensor * q_seq_lens: query sequence lengths list * kv_seq_lens: shared key/value sequence lengths list + """ query: torch.Tensor @@ -155,16 +152,15 @@ class PackedQKVInputs(NamedTuple): class PackedQKVO(NamedTuple): - """ - Data structure for representing packed attention inputs, + """Data structure for representing packed attention inputs, alongside packed known-correct attention output Attributes: - * packed_qkv: packed (number_of_tokens x num_heads x head_size) attention inputs * ideal_output: packed (number_of_tokens x num_heads x head_size) known-correct attention output + """ packed_qkv: PackedQKVInputs | None @@ -172,13 +168,12 @@ class PackedQKVO(NamedTuple): class KVMemoryMap(NamedTuple): - """ - Data structure for encapsulating KV cache memory mapping. + """Data structure for encapsulating KV cache memory mapping. Attributes: - * block_tables: KV cache block tables * slot_mapping: mapping of sequence offset to physical address + """ block_tables: torch.Tensor @@ -186,18 +181,17 @@ class KVMemoryMap(NamedTuple): class PhaseTestParameters(NamedTuple): - """ - Data structure for encapsulating the test parameters + """Data structure for encapsulating the test parameters for a given test "phase" (prefill or decode phase) and attention scenario (encoder, decoder-self, encoder/decoder-cross) Attributes: - * packed_qkvo: packed (number_of_tokens x num_heads x head_size) attention inputs & known-correct output * kv_mmap: KV cache memory mapping, specific to this test phase & attention scenario + """ packed_qkvo: PackedQKVO @@ -208,13 +202,11 @@ def maybe_make_int_tensor( _list: list[int] | None, device: torch.device | str, ) -> torch.Tensor: - """ - Convert Python int list to a 1D int torch.Tensor on `device` + """Convert Python int list to a 1D int torch.Tensor on `device`. Returns: + If _list is not None: 1D int torch.Tensor on `device`, `None` otherwise. - * If _list is not None: 1D int torch.Tensor on `device` - * None otherwise """ return ( None if _list is None else torch.tensor(_list, dtype=torch.int, device=device) @@ -225,13 +217,11 @@ def maybe_make_long_tensor( _list: list[int] | None, device: torch.device | str, ) -> torch.Tensor: - """ - Convert Python int list to a 1D long torch.Tensor on `device` + """Convert Python int list to a 1D long torch.Tensor on `device`. Returns: + If _list is not None: 1D long torch.Tensor on `device`, `None` otherwise. - * If _list is not None: 1D long torch.Tensor on `device` - * None otherwise """ return ( None if _list is None else torch.tensor(_list, dtype=torch.long, device=device) @@ -239,11 +229,11 @@ def maybe_make_long_tensor( def maybe_max(_list: list | None) -> Number | None: - """ + """Compute the maximum of a list, if it exists. + Returns: + If _list is not None: max(_list), `None` otherwise. - * If _list is not None: max(_list) - * None otherwise """ return None if _list is None else max(_list) @@ -252,19 +242,16 @@ def make_causal_mask( q_max_seq_len: int, kv_max_seq_len: int, ) -> torch.Tensor: - """ - Create a q_max_seq_len x kv_max_seq_len causal mask + """Create a q_max_seq_len x kv_max_seq_len causal mask. Arguments: - - * q_max_seq_len: query max seq len - * kv_max_seq_len: key/value max seq len + q_max_seq_len: query max seq len + kv_max_seq_len: key/value max seq len Returns: + 2D tensor, q_max_seq_len x kv_max_seq_len - * 2D tensor, q_max_seq_len x kv_max_seq_len """ - # Create a matrix where entry (i, j) is True if i >= j mask = torch.triu(torch.ones(q_max_seq_len, kv_max_seq_len), diagonal=1) # Replace True with float('-inf') and False with 0 @@ -281,8 +268,7 @@ def ref_masked_attention( q_seq_lens: list | None = None, kv_seq_lens: list | None = None, ) -> torch.Tensor: - """ - "Golden" masked attention reference. Supports two types of masking: + """Reference ("golden") masked attention. Supports two types of masking: * Basic attention mask, utilizing {q,kv}_seq_lens args to mask out padding elements @@ -290,21 +276,18 @@ def ref_masked_attention( causal Arguments: - - * query: batch_size x q_padded_seq_len x num_heads x head_size - * key: batch_size x kv_padded_seq_len x num_heads x head_size - * value: batch_size x kv_padded_seq_len x num_heads x head_size - * scale: Attention scale factor - * custom_mask: custom attention mask; good place to inject a causal - attention mask - * q_seq_lens: list of unpadded query seq_lens for each batch index - * kv_seq_lens: list of unpadded key/value seq_lens for each batch index + query: batch_size x q_padded_seq_len x num_heads x head_size + key: batch_size x kv_padded_seq_len x num_heads x head_size + value: batch_size x kv_padded_seq_len x num_heads x head_size + scale: Attention scale factor + custom_mask: custom attention mask; good place to inject a causal attention mask + q_seq_lens: list of unpadded query seq_lens for each batch index + kv_seq_lens: list of unpadded key/value seq_lens for each batch index Returns: + Attention result, batch_size x q_padded_seq_len x num_heads x head_size - * Attention result, batch_size x q_padded_seq_len x num_heads x head_size """ - assert q_seq_lens is not None assert kv_seq_lens is not None @@ -346,8 +329,7 @@ def make_qkv( attn_type: AttentionType = AttentionType.ENCODER_DECODER, force_max_len: bool = False, ) -> tuple[QKVInputs, QKVInputs, QKVInputs]: - """ - Construct QKV test tensors for self- and cross-attention. + """Construct QKV test tensors for self- and cross-attention. Generates three query/key/value triplets: @@ -361,30 +343,28 @@ def make_qkv( seqlens Arguments: - - * batch_size - * max_q_seq_len: max query seq len - * max_kv_seq_len: max key/value seq len - * num_heads - * head_size - * is_encoder_decoder_attn: if True, query seqlen may differ from - key/value seqlen (as is often the case for cross-attention); - o/w, query/key/value seqlens match at each batch index - (max_kv_seq_len is unused) - * force_kv_seq_lens: if not None, overrides kv sequence lengths - * attn_type: encoder, decoder self, or enc/dec cross attention - * force_max_len: if True, all query seqlens are max_q_seq_len; o/w query - seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens - and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False - * device: CPU or CUDA device + batch_size: number of sequences in the batch + max_q_seq_len: max query seq len + max_kv_seq_len: max key/value seq len + num_heads: number of attention heads + head_size: head dimension + is_encoder_decoder_attn: if True, query seqlen may differ from + key/value seqlen (as is often the case for cross-attention); + o/w, query/key/value seqlens match at each batch index + (max_kv_seq_len is unused) + force_kv_seq_lens: if not None, overrides kv sequence lengths + attn_type: encoder, decoder self, or enc/dec cross attention + force_max_len: if True, all query seqlens are max_q_seq_len; o/w query + seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens + and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False + device: CPU or CUDA device Returns: + Overall QKVInputs structure (containing full unpacked Q/K/V tensors) + Prefill QKVInputs structure (containing all but the last sequence offset) + Decode QKVInputs structure (containing all only the last sequence offset) - * Overall QKVInputs structure (containing full unpacked Q/K/V tensors) - * Prefill QKVInputs structure (containing all but the last sequence offset) - * Decode QKVInputs structure (containing all only the last sequence offset) """ - if force_max_len: q_seq_lens = [max_q_seq_len for _ in range(batch_size)] else: @@ -474,24 +454,21 @@ def make_qkv( def pack_tensor( unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str ) -> tuple[torch.Tensor, list[int]]: - """ - Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an + """Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an unpadded number_of_tokens x num_heads x head_size tensor, where number_of_tokens = sum(seq_lens) Arguments: + unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size + seq_lens: list of token counts for each seq + device: CPU or CUDA device - * unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size - * seq_lens: list of token counts for each seq - * device: CPU or CUDA device - - Returns + Returns: + packed_tensor: number_of_tokens x num_heads x head_size + start_loc_list: start idx of each batch elt in packed_tensor; [0] + + list(itertools.accumulate(seq_lens)) - * packed_tensor: number_of_tokens x num_heads x head_size - * start_loc_list: start idx of each batch elt in packed_tensor; [0] + - list(itertools.accumulate(seq_lens)) """ - num_tok = sum(seq_lens) num_heads = unpacked_tensor.shape[-2] head_size = unpacked_tensor.shape[-1] @@ -507,8 +484,7 @@ def pack_tensor( def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs: - """ - Individually pack each of Q, K and V, each with dimensions batch_size x + """Individually pack each of Q, K and V, each with dimensions batch_size x padded_seq_len x num_heads x head_size, into respective number_of_tokens x num_heads x head_size tensors. @@ -517,17 +493,15 @@ def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs: For K and V, number_of_tokens = sum(kv_seq_lens) Arguments: + qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size) + attention inputs + device: CPU or CUDA device - * qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size) - attention inputs - * device: CPU or CUDA device - - Returns + Returns: + Packed (number_of_tokens x num_heads x head_size) QKV inputs + derived from unpacked inputs - * Packed (number_of_tokens x num_heads x head_size) QKV inputs - derived from unpacked inputs """ - if qkv.query is None: packed_query = None q_start_loc_list = None @@ -563,26 +537,24 @@ def _make_metadata_tensors( torch.Tensor, int | None, ]: - """ - Build scalar & tensor values required to build attention metadata structure. + """Build scalar & tensor values required to build attention metadata structure. Arguments: - - * seq_lens: list of token-counts for each decoder input seq - * context_lens: list of context length values for each seq - * encoder_seq_lens: list of token-counts for each encoder input seq - * device: CPU or CUDA device + seq_lens: list of token-counts for each decoder input seq + context_lens: list of context length values for each seq + encoder_seq_lens: list of token-counts for each encoder input seq + device: CPU or CUDA device Returns: + seq_lens_tensor: decoder seq_lens list, as tensor + context_lens_tensor: context_lens list, as tensor + max_context_len: max(context_lens) + max_seq_len: max(seq_lens) + seq_start_loc: start idx of each sequence + encoder_seq_lens_tensor: encoder seq_lens list, as tensor + encoder_seq_start_loc: start idx of each encoder sequence + max_encoder_seq_len: encoder seq_lens list, as tensor - * seq_lens_tensor: decoder seq_lens list, as tensor - * context_lens_tensor: context_lens list, as tensor - * max_context_len: max(context_lens) - * max_seq_len: max(seq_lens) - * seq_start_loc: start idx of each sequence - * encoder_seq_lens_tensor: encoder seq_lens list, as tensor - * encoder_seq_start_loc: start idx of each encoder sequence - * max_encoder_seq_len: encoder seq_lens list, as tensor """ seq_lens_tensor = maybe_make_int_tensor(seq_lens, device) context_lens_tensor = maybe_make_int_tensor(context_lens, device) @@ -637,22 +609,21 @@ def make_kv_cache( backend: str, default_val: float = 0.0, ) -> torch.Tensor: - """ - Create a fake KV cache. + """Create a fake KV cache. Arguments: - - * num_blocks: number of blocks in the KV cache - * num_heads: number of attention heads - * head_size: head dimension - * block_size: number of offsets within a block - * device: CPU or CUDA device - * default_val: initialization value for KV cache elements + num_blocks: number of blocks in the KV cache + num_heads: number of attention heads + head_size: head dimension + block_size: number of offsets within a block + device: CPU or CUDA device + backend: attention backend name; only 'FLASH_ATTN' is supported + default_val: initialization value for KV cache elements Returns: + kv_cache: 2 x num_blocks x block_size x num_heads x head_size + for backend 'FLASH_ATTN' - * kv_cache: 2 x num_blocks x block_size x num_heads x head_size - * for backend 'FLASH_ATTN' """ if backend != "FLASH_ATTN": raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.") @@ -663,8 +634,7 @@ def make_kv_cache( def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int: - """ - Compute the minimum number of blocks required to hold num_tokens tokens, + """Compute the minimum number of blocks required to hold num_tokens tokens, given block_size """ return (num_tokens + block_size - 1) // block_size @@ -683,8 +653,7 @@ def split_slot_mapping( seq_lens: list[int], device: torch.device | str, ): - """ - Split a slot mapping into valid prefill- and decode-phase slot mappings. + """Split a slot mapping into valid prefill- and decode-phase slot mappings. Context: * Your goal is to test (1) prefill of N prompts, with prompt-lengths @@ -708,21 +677,19 @@ def split_slot_mapping( The N excised entries are appended to obtain the decode-phase slot mapping Arguments: - - * slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N - post-decode sequences - * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the - description above) - * device: cuda, cpu, etc. + slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N + post-decode sequences + seq_lens: list of N post-decode sequence lengths (K_i + 1 in the + description above) + device: cuda, cpu, etc. Returns: + prefill_slot_mapping: Length-M 1D slot mapping (as Tensor) + reflecting all N prefill prompts + decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting + all N decoded tokens - * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor) - reflecting all N prefill prompts - * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting - all N decoded tokens """ - prefill_slot_mapping = [] decode_slot_mapping = [] @@ -746,8 +713,7 @@ def make_block_tables_slot_mapping( device: torch.device | str, block_base_addr: int = 0, ) -> tuple[torch.Tensor, list[int], int]: - """ - Construct fake block tables & slot mappings. + """Construct fake block tables & slot mappings. For a sequence with num_tokens tokens the minimum number of required KV cache blocks is @@ -772,19 +738,17 @@ def make_block_tables_slot_mapping( i.e. the total of prefill prompt tokens + decoded tokens. Arguments: - - * block_size: number of offsets per block - * seq_lens: list of token-counts for each sequence - * block_base_addr: the block table base address - * device: CPU or CUDA device + block_size: number of offsets per block + seq_lens: list of token-counts for each sequence + block_base_addr: the block table base address + device: CPU or CUDA device Return: - * block_tables_tensor: block table for sequence * slot_mapping_list: slot mapping for sequence * max_block_idx: the highest block address within this block table - """ + """ # Provision minimum number of KV cache blocks num_blocks_list = [ _num_tokens_to_min_blocks(num_tokens, block_size) for num_tokens in seq_lens @@ -824,14 +788,14 @@ def make_block_tables_slot_mapping( def assert_actual_matches_ideal( test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str ) -> None: - """ - Assert that observed output matches the ideal output + """Assert that observed output matches the ideal output contained in the test parameters data structure. Arguments: + test_params: Test parameters including packed ideal output + output_under_test: actually observed output value + backend: attention backend name; only 'FLASH_ATTN' is supported - * test_params: Test parameters including packed ideal output - * output_under_test: actually observed output value """ ideal_output = test_params.packed_qkvo.ideal_output if backend != "FLASH_ATTN": @@ -852,9 +816,7 @@ def fp8_allclose( atol: float = 1e-08, equal_nan: bool = False, ) -> bool: - """ - Reference implementation of torch.allclose - """ + """Reference implementation of torch.allclose.""" torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol) return bool( diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 8635779cff44..f5ca44f174e2 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -33,7 +33,6 @@ def should_do_global_cleanup_after_test(request) -> bool: This can provide a ~10x speedup for non-GPU unit tests since they don't need to initialize torch. """ - return not request.node.get_closest_marker("skip_global_cleanup") @@ -367,8 +366,7 @@ def qwen36_moe_3d_lora_files(): @pytest.fixture def reset_default_device(): - """ - Some tests, such as `test_punica_ops.py`, explicitly set the + """Some tests, such as `test_punica_ops.py`, explicitly set the default device, which can affect subsequent tests. Adding this fixture helps avoid this problem. """ diff --git a/tests/lora/test_add_lora.py b/tests/lora/test_add_lora.py index 628f5bc28eb5..97c37c189522 100644 --- a/tests/lora/test_add_lora.py +++ b/tests/lora/test_add_lora.py @@ -55,8 +55,7 @@ async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> flo @pytest.mark.asyncio async def test_add_lora(chatglm3_lora_files): - """ - The add_lora function is used to preload some LoRA adapters into the + """The add_lora function is used to preload some LoRA adapters into the engine in anticipation of future requests using these adapters. To test this functionality, we use the async engine to process some requests - We do it twice, once with add_lora() preloading and once without. diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py index 5e3c762f79b9..3c9505adba7c 100644 --- a/tests/lora/test_default_mm_loras.py +++ b/tests/lora/test_default_mm_loras.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for applying default registered multimodal loras. -""" +"""Tests for applying default registered multimodal loras.""" import os import unittest.mock as mock @@ -134,7 +132,8 @@ def test_default_mm_lora_fails_with_overridden_lora_request( monkeypatch: pytest.MonkeyPatch, ): """Ensure that if the lora_request conflicts with default_mm_loras, - we use the lora_request.""" + we use the lora_request. + """ run_test( vllm_runner, audio_assets, diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index 444b609492ba..9ee0ed8584c8 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -41,8 +41,7 @@ def CEILDIV(x, y): def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): - """ - Split `num_tokens` into `num_sequences` sequences. + """Split `num_tokens` into `num_sequences` sequences. Each sequence randomly selects 1 LoRA index from [0, max_loras), and all tokens in that sequence are assigned this LoRA index. @@ -54,6 +53,7 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): Returns: torch.Tensor: 1D tensor of shape [num_tokens], where each value is the LoRA index assigned to that token. + """ assert num_sequences > 0 and max_loras > 0 assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences" @@ -81,9 +81,8 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int): def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int): - """ - For each token, randomly select `top_k_num` distinct experts out of `num_experts`, - and assign normalized random weights that sum to 1. + """For each token, randomly select `top_k_num` distinct experts out of + `num_experts`, and assign normalized random weights that sum to 1. Args: num_tokens (int): Total number of tokens. @@ -95,6 +94,7 @@ def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int): expert index for each token. expert_weights (torch.Tensor): shape [num_tokens, top_k_num], normalized weights (sum = 1 per row). + """ assert top_k_num <= num_experts, "top_k_num must be <= num_experts" @@ -379,8 +379,7 @@ def use_fused_moe_lora_kernel_naive( offset=0, add_inputs=True, ): - """ - Test helper for naive_block_assignment path. + """Test helper for naive_block_assignment path. Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids. """ config = { @@ -472,8 +471,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment( device, seed, ): - """ - Test the naive_block_assignment path of the fused_moe_lora kernel. + """Test the naive_block_assignment path of the fused_moe_lora kernel. This path is triggered when batch_size * top_k is much smaller than num_experts * max_loras, and skips the moe_lora_align_block_size kernel. """ @@ -1125,7 +1123,8 @@ def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device): @pytest.mark.parametrize("device", DEVICES) def test_fused_moe_lora_kernel_zero_grid_no_crash(device): """num_active_loras=0 (or num_slices=0) would otherwise launch a grid - with a zero dimension. one-shot wrapper must short-circuit before launch.""" + with a zero dimension. one-shot wrapper must short-circuit before launch. + """ torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024 @@ -1187,7 +1186,8 @@ def test_fused_moe_lora_kernel_zero_grid_no_crash(device): def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device): """one-shot must surface a clear assertion when shrink_block_size_m is not a power of 2 / less than 16, instead of the cryptic Triton compile - failure (`arange's range must be a power of 2`).""" + failure (`arange's range must be a power of 2`). + """ torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024 diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 2dab86bc0d57..e6f84cf0bc6c 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -126,8 +126,7 @@ def clean_cache_reset_device(reset_default_device): @pytest.fixture(autouse=True) def skip_cuda_with_stage_false(request): - """ - On cuda-like platforms, we use the same kernels for prefill and decode + """On cuda-like platforms, we use the same kernels for prefill and decode stage, and 'stage' is generally ignored, so we only need to test once. """ if current_platform.is_cuda_alike() or current_platform.is_xpu(): @@ -153,8 +152,8 @@ def get_random_id_to_index( num_slots: The number of slots in the mapping. Must be larger than num_loras. log: Whether to log the output. - """ + """ if num_loras > num_slots: raise ValueError( f"num_loras is higher than num_slots: {num_loras} > {num_slots}. " @@ -190,8 +189,8 @@ def populate_loras( repeats: must only be set for column parallel packed layers. Indicates the number of loras to compose together to create a single lora layer. - """ + """ # Dictionary that maps the lora ID to the # corresponding lora weights. lora_dict: dict[int, LoRALayerWeights] = dict() @@ -246,8 +245,8 @@ def create_random_inputs( input_range: the range of values to include in the input. input_range[0] <= possible input values < input_range[1] input_type: the type of values in the input. - """ + """ low, high = input_range inputs: list[torch.Tensor] = [] diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py index ecf2ad9ee654..1b841b70cea2 100644 --- a/tests/lora/test_lora_functions.py +++ b/tests/lora/test_lora_functions.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Script to test add_lora, remove_lora, pin_lora, list_loras functions. -""" +"""Script to test add_lora, remove_lora, pin_lora, list_loras functions.""" import pytest diff --git a/tests/lora/test_mixtral.py b/tests/lora/test_mixtral.py index 3868bff79663..a8f1fdda64c3 100644 --- a/tests/lora/test_mixtral.py +++ b/tests/lora/test_mixtral.py @@ -32,7 +32,7 @@ def do_sample( @pytest.mark.parametrize("tp_size", [4]) def test_mixtral_lora(mixtral_lora_files, tp_size): - """Original test, the LoRA model has the common target modules, not all""" + """Original test, the LoRA model has the common target modules, not all.""" if ( torch.accelerator.device_count() < tp_size and tp_size > 1 diff --git a/tests/lora/test_moe_lora_align_sum.py b/tests/lora/test_moe_lora_align_sum.py index 80ba0a56ca40..df27752d2c37 100644 --- a/tests/lora/test_moe_lora_align_sum.py +++ b/tests/lora/test_moe_lora_align_sum.py @@ -249,7 +249,8 @@ def _build_and_run_align( ) def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras): """Regression test for issue #32235: real LoRA slot must not be skipped - when ``active_lora_ids`` has -1 at position 0.""" + when ``active_lora_ids`` has -1 at position 0. + """ out = _build_and_run_align( num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras ) @@ -288,7 +289,8 @@ def test_moe_lora_align_block_size_disabled_adapter_untouched(): kernels. Pins the invariant protected by the ``adapter_enabled`` guard in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort kernel reads uninitialized ``token_mask`` values for disabled slots and - pollutes ``sorted_token_ids`` / ``cumsum_buffer``.""" + pollutes ``sorted_token_ids`` / ``cumsum_buffer``. + """ max_loras = 1 out = _build_and_run_align( num_lora_tokens=16, diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py index ead40b25522b..52cccc6b06ce 100644 --- a/tests/lora/test_moe_lora_ep_load.py +++ b/tests/lora/test_moe_lora_ep_load.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end correctness test for 2D MoE LoRA expert-parallel +"""End-to-end correctness test for 2D MoE LoRA expert-parallel load-time slicing """ @@ -20,7 +19,8 @@ def _expected_lora_modules() -> set[str]: """Replicate the set ``WorkerLoRAManager._load_adapter`` would build - from this model's ``packed_modules_mapping``.""" + from this model's ``packed_modules_mapping``. + """ expected: set[str] = set(NON_EXPERT_MODULES) for expert in range(GLOBAL_NUM_EXPERTS): for proj in EXPERT_PROJECTIONS: @@ -43,7 +43,8 @@ def _load(lora_dir, peft_helper, *, moe_ep_spec, lora_id): def test_moe_lora_ep2_real_qwen3moe(qwen3moe_lora_files, ep_rank): """ep_size=2 against the real Qwen3-MoE adapter: each rank's loaded LoRA has the right size, the right expert membership, and the - right tensor values.""" + right tensor values. + """ peft_helper = PEFTHelper.from_local_dir( qwen3moe_lora_files, max_position_embeddings=4096 ) diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 9e1f99b3f6a3..f2bb18aac26e 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -45,7 +45,8 @@ def _bgmv_shrink( inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 ): """Memory-efficient shrink reference: per-LoRA matmul loop. - output[mask] = scaling * inputs[mask] @ weight.T""" + output[mask] = scaling * inputs[mask] @ weight.T + """ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: @@ -66,7 +67,8 @@ def _bgmv_expand( add_inputs=False, ): """Memory-efficient expand reference: per-LoRA matmul loop. - output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" + output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T + """ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: @@ -160,8 +162,7 @@ def check_lora_shrink_kernel( seq_length: int, scaling: float, ): - """ - Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink + """Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink kernels. """ data: PunicaTensors = generate_data_for_nslices( @@ -234,8 +235,7 @@ def check_lora_expand_kernel( seq_length: int, add_inputs: bool, ): - """ - Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand + """Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand kernels. """ data: PunicaTensors = generate_data_for_nslices( @@ -438,9 +438,7 @@ def test_kernels( seed: int, op_type: str, ): - """ - Tests LoRA kernels. - """ + """Tests LoRA kernels.""" torch.set_default_device(device) torch.accelerator.set_device_index(device) set_random_seed(seed) @@ -491,9 +489,7 @@ def test_kernels_hidden_size( seed: int, op_type: str, ): - """ - Tests SGMV and LoRA kernels. - """ + """Tests SGMV and LoRA kernels.""" torch.set_default_device(device) torch.accelerator.set_device_index(device) set_random_seed(seed) @@ -526,8 +522,7 @@ def test_kernels_hidden_size( @pytest.mark.parametrize("device", DEVICES) def test_add_lora_fused_moe_early_exit(device): - """ - Ensures add_lora_fused_moe does not invoke the LoRA kernel or + """Ensures add_lora_fused_moe does not invoke the LoRA kernel or modify the output tensor when no_lora_flag_cpu is True """ from types import SimpleNamespace diff --git a/tests/lora/test_punica_ops_fp8.py b/tests/lora/test_punica_ops_fp8.py index e447cef9d8a8..18ac6fcd7aad 100644 --- a/tests/lora/test_punica_ops_fp8.py +++ b/tests/lora/test_punica_ops_fp8.py @@ -877,7 +877,8 @@ def test_lora_shrink_fp8( quant_mode: str, ): """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline.""" + quantization, comparing against the bf16 baseline. + """ torch.set_default_device(device) set_random_seed(seed) @@ -932,7 +933,8 @@ def test_lora_expand_fp8( quant_mode: str, ): """Test FP8 expand kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline.""" + quantization, comparing against the bf16 baseline. + """ torch.set_default_device(device) set_random_seed(seed) diff --git a/tests/lora/test_punica_xpu_ops.py b/tests/lora/test_punica_xpu_ops.py index 585c97cfa547..36966809ebce 100644 --- a/tests/lora/test_punica_xpu_ops.py +++ b/tests/lora/test_punica_xpu_ops.py @@ -86,9 +86,7 @@ def check_bgmv_shrink( device: str, scaling: float, ): - """ - Compare vllm.bgmv_shrink against a reference implementation. - """ + """Compare vllm.bgmv_shrink against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data( batches, @@ -130,9 +128,7 @@ def check_bgmv_expand( device: str, add_inputs: bool, ): - """ - Compare vllm.bgmv_expand against a reference implementation. - """ + """Compare vllm.bgmv_expand against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data( batches, @@ -172,9 +168,7 @@ def check_bgmv_expand_slice( device: str, add_inputs: bool, ): - """ - Compare vllm.bgmv_expand_slice against a reference implementation. - """ + """Compare vllm.bgmv_expand_slice against a reference implementation.""" seq_length = 1 data: PunicaTensors = generate_data_for_expand_nslices( batches, diff --git a/tests/lora/test_qwen3_unembed.py b/tests/lora/test_qwen3_unembed.py index 4064b443ac6a..a70df902f719 100644 --- a/tests/lora/test_qwen3_unembed.py +++ b/tests/lora/test_qwen3_unembed.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for Qwen3 unembed LoRA support. +"""Tests for Qwen3 unembed LoRA support. This test creates synthetic LoRA weights that include lm_head (output embedding) to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer. diff --git a/tests/lora/test_qwen3_with_multi_loras.py b/tests/lora/test_qwen3_with_multi_loras.py index 9870da6959db..185ab91c7a3d 100644 --- a/tests/lora/test_qwen3_with_multi_loras.py +++ b/tests/lora/test_qwen3_with_multi_loras.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script contains: +"""This script contains: 1. test multi loras service with tp >= 2 2. test multi loras request """ @@ -125,8 +124,7 @@ def call_llm_get_outputs(prompt: str, lora_name: str): return output_text def reload_lora(name: str): - """ - reload a lora to simulate the case: + """Reload a lora to simulate the case: setting `VLLM_ALLOW_RUNTIME_LORA_UPDATING=true` for dynamic lora loading and unloading """ @@ -205,8 +203,7 @@ def test_multiple_lora_requests(): def test_load_inplace_offline_reload( qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str ) -> None: - """ - Test that load_inplace=True allows reloading LoRA adapters with the same ID + """Test that load_inplace=True allows reloading LoRA adapters with the same ID in offline mode (using LLM class directly). """ llm = LLM( @@ -255,8 +252,7 @@ def test_load_inplace_offline_reload( def test_load_inplace_false_no_reload( qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str ) -> None: - """ - Test that load_inplace=False prevents reloading when an adapter + """Test that load_inplace=False prevents reloading when an adapter with the same ID already exists. """ llm = LLM( diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index 3362aa46ed26..063de2b5875d 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -46,7 +46,7 @@ def __post_init__(self): class Qwen2VLTester: - """Test helper for Qwen2 VL models with LoRA""" + """Test helper for Qwen2 VL models with LoRA.""" PROMPT_TEMPLATE = ( "<|im_start|>system\nYou are a helpful assistant.<|im_end|>" @@ -199,7 +199,7 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: def test_qwen2vl_lora(qwen2vl_lora_files): - """Test Qwen 2.0 VL model with LoRA""" + """Test Qwen 2.0 VL model with LoRA.""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) with Qwen2VLTester(config) as tester: # Test with different LoRA IDs @@ -230,7 +230,7 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests" ) def test_qwen25vl_lora(qwen25vl_lora_files): - """Test Qwen 2.5 VL model with LoRA""" + """Test Qwen 2.5 VL model with LoRA.""" config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files) with Qwen2VLTester(config) as tester: # Test with different LoRA IDs @@ -292,8 +292,7 @@ def test_qwen2vl_multiple_lora_types( qwen2vl_vision_tower_lora_files, monkeypatch: pytest.MonkeyPatch, ): - """ - Test multiple LoRA adapter types (language, vision tower + connector, + """Test multiple LoRA adapter types (language, vision tower + connector, vision tower only) using the same LLM instance to verify mm_encoder_cache behavior with different LoRA requests. diff --git a/tests/lora/test_whisper.py b/tests/lora/test_whisper.py index 6f1a894cf916..73cc550b294b 100644 --- a/tests/lora/test_whisper.py +++ b/tests/lora/test_whisper.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration tests for Whisper models with LoRA adapters. +"""Integration tests for Whisper models with LoRA adapters. These tests verify that Whisper models can correctly load and use LoRA adapters for speech-to-text transcription tasks. diff --git a/tests/lora/utils.py b/tests/lora/utils.py index e5ce7a88464c..8ee2f8af8224 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -116,8 +116,7 @@ class PunicaTensors: token_lora_mapping: torch.Tensor def meta(self) -> tuple[int, int]: - """ - Infer max_seq_length and token_nums from the tensors + """Infer max_seq_length and token_nums from the tensors and return them. """ max_seq_length = self.seq_len_tensor.max() diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 507fd93e9d02..6fdd99deaa44 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -808,7 +808,6 @@ def test_models_fse_init( monkeypatch: pytest.MonkeyPatch, ) -> None: """Model construction resolves FSE consistently with Quark quantization.""" - quantization_config: dict[str, Any] = ( get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else []) if model_type == "deepseek_v4" @@ -1320,7 +1319,8 @@ def test_fp8_shared_expert_fse_allows_symmetric_ignored_layers() -> None: def test_fp8_shared_expert_fse_expands_packed_projections() -> None: """Both shards of a fused projection ignored is symmetric within that - projection, so the mismatch reported is against the routed experts.""" + projection, so the mismatch reported is against the routed experts. + """ quant_config = _fp8_config( ignored_layers=[ "model.layers.0.mlp.shared_experts.gate_proj", @@ -1344,7 +1344,8 @@ def test_fp8_shared_expert_fse_expands_packed_projections() -> None: def test_fp8_shared_expert_fse_propagates_partial_shard_exclusion() -> None: """Half a fused projection excluded is rejected by `is_layer_skipped` - itself, exactly as it is in `Fp8Config.get_quant_method`.""" + itself, exactly as it is in `Fp8Config.get_quant_method`. + """ quant_config = _fp8_config( ignored_layers=["model.layers.0.mlp.shared_experts.gate_proj"] ) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 41a06d0546f0..91fc039569dc 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -346,7 +346,8 @@ def record_forward_impl(*args): def test_select_candidate_blocks_tolerates_empty_rows(): """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest block pin must not index -1 for them (device-side assert); they select - no candidate blocks and real rows still pin their newest block.""" + no candidate blocks and real rows still pin their newest block. + """ block_size, topk_blocks = 8, 3 logits = torch.zeros(3, 64, device="cuda") logits[0, 3] = 5.0 # block 0 scores highest for row 0 @@ -457,7 +458,8 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): """Production-faithful candidates: top min(avail, K) blocks per row — - full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + full coverage for lens <= K*cbk, a fixed-size subset beyond that. + """ avail = ((lens.cpu() + cbk - 1) // cbk).tolist() rows = [] for a in avail: @@ -493,7 +495,8 @@ def _candidate_token_mask( def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid - count are -1 (the indexer's "no token" sentinel).""" + count are -1 (the indexer's "no token" sentinel). + """ values, indices = masked.topk(topk, dim=1) return torch.where(values > float("-inf"), indices.int(), -1) @@ -506,7 +509,8 @@ def _assert_same_selection( ) -> None: """Both top-k index sets must pick the same multiset of logit values per row (-1 padding must match too). Indices themselves may differ: the - top-k kernels break ties among equal values in nondeterministic order.""" + top-k kernels break ties among equal values in nondeterministic order. + """ for r in range(logits.shape[0]): base = int(row_starts[r]) if row_starts is not None else 0 g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] @@ -518,7 +522,8 @@ def _assert_same_selection( def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: """Caller-owned buffers the sparse path writes into (the metadata builder - owns these in production).""" + owns these in production). + """ i32 = dict(dtype=torch.int32, device="cuda") return dict( sparse_indices=torch.zeros(rows, num_sparse, **i32), @@ -533,7 +538,8 @@ def _expand_candidates_reference( """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand each candidate block into ``cbk // sbk`` sparse blocks, keep those inside ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by - repeating the last valid block, and count the valid sparse columns.""" + repeating the last valid block, and count the valid sparse columns. + """ ratio = cbk // sbk rows, num_candidates = candidates.shape num_sparse = num_candidates * ratio @@ -559,7 +565,8 @@ def _expand_candidates_reference( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") def test_candidate_blocks_to_sparse_indices_math(): """Candidate->sparse-block expansion: ks anchoring, range filtering, - repeat-last padding, and the per-row valid column count (``end``).""" + repeat-last padding, and the per-row valid column count (``end``). + """ from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -611,7 +618,8 @@ def test_candidate_blocks_to_sparse_indices_math(): def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): """The Triton expansion equals the PyTorch reference on production-shaped random inputs, including strided (column-sliced) candidate rows and - strided output buffers.""" + strided output buffers. + """ from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -681,7 +689,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): ``torch.topk`` over dense bf16 logits masked to the candidates (the two DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA graph capture, and a later layer reusing the step's expansion/schedule - selects the same values.""" + selects the same values. + """ deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_prefill_chunk, @@ -771,7 +780,8 @@ def run_sparse(kernel_metadata=None): def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' - fp4 data first, then their packed UE8M0 scales.""" + fp4 data first, then their packed UE8M0 scales. + """ device = req_lens.device pages_per = (req_lens + page_kv - 1) // page_kv num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) @@ -807,7 +817,8 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to one query row each), CUDA graph capture/replay and schedule reuse by a - later layer.""" + later layer. + """ deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_paged_decode, @@ -916,7 +927,8 @@ def run_sparse(o: torch.Tensor, kernel_metadata=None): def test_sparse_mqa_logits_accuracy_vs_torch(): """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized - MXFP4 inputs.""" + MXFP4 inputs. + """ _skip_unless_sm100_sparse_kernels() from deep_gemm.utils import cast_back_from_fp4 diff --git a/tests/model_executor/model_loader/test_ep_weight_filter.py b/tests/model_executor/model_loader/test_ep_weight_filter.py index 0ff53d24e71d..9d177f59910f 100644 --- a/tests/model_executor/model_loader/test_ep_weight_filter.py +++ b/tests/model_executor/model_loader/test_ep_weight_filter.py @@ -232,7 +232,8 @@ def test_fused_3d_expert_not_skipped(self): class TestSafetensorsWeightsIteratorWithEpFilter: """Verify that EP filtering produces a strict subset of unfiltered loading - and that all expected dense + local expert weights are present.""" + and that all expected dense + local expert weights are present. + """ @pytest.fixture(scope="class") def gpt2_files(self): @@ -262,7 +263,8 @@ def test_no_filter_returns_all(self, gpt2_files): def test_empty_filter_skips_experts_only(self, gpt2_files): """GPT-2 has no expert weights, so even an empty local_expert_ids - set should return all weights (all are dense).""" + set should return all weights (all are dense). + """ all_weights = dict(safetensors_weights_iterator(gpt2_files, False)) filtered_weights = dict( safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set()) @@ -273,7 +275,8 @@ def test_empty_filter_skips_experts_only(self, gpt2_files): class TestEpFilterOnSyntheticMoeWeights: """Create synthetic safetensors files with expert-like naming and verify - that the filter correctly skips non-local experts.""" + that the filter correctly skips non-local experts. + """ @pytest.fixture def synthetic_moe_files(self, tmp_path): diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b4a544d2fc0f..d123dcb7508b 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -350,7 +350,8 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold.""" + reuse the workspace storage whose address captured CUDA graphs hold. + """ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -428,7 +429,8 @@ def load_checkpoint_format_weights(): def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that.""" + graphs is already unusable; allocating a replacement would hide that. + """ from vllm.model_executor.layers.quantization.utils import marlin_utils monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) @@ -609,7 +611,8 @@ def process_weights_after_loading(self, layer): class _LateBiasLayer(torch.nn.Module): """Mimics an online-quantized linear: `weight` is created on meta by `create_weights()`, which wraps the loaders, and the linear base registers - `bias` afterwards.""" + `bias` afterwards. + """ def __init__(self, quant_method): super().__init__() diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index 6f2a31cf2787..ada073706606 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -42,7 +42,8 @@ def _run_warmup(monkeypatch, layers) -> list[dict]: @pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"]) def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None: """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the - quantization method that owns it or the name of its scale parameter.""" + quantization method that owns it or the name of its scale parameter. + """ layer = _block_fp8_layer(256, 128, scale_name) calls = _run_warmup(monkeypatch, [layer]) @@ -54,7 +55,8 @@ def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> Non def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None: """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576), - so warmup must not require N % 128 == 0.""" + so warmup must not require N % 128 == 0. + """ assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1 @@ -71,7 +73,7 @@ def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None: @pytest.mark.parametrize("is_bmm", [False, True]) def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: - """bmm layers bypass the kernel at runtime, so they must not be warmed.""" + """Bmm layers bypass the kernel at runtime, so they must not be warmed.""" kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel) kernel.is_deep_gemm_supported = False layer = torch.nn.Module() diff --git a/tests/model_executor/test_model_load_with_params.py b/tests/model_executor/test_model_load_with_params.py index 785335e99111..7d1bddab0db2 100644 --- a/tests/model_executor/test_model_load_with_params.py +++ b/tests/model_executor/test_model_load_with_params.py @@ -23,9 +23,7 @@ current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test parameter weight loading with tp>1. - """ + """Test parameter weight loading with tp>1.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") with vllm_runner( @@ -68,9 +66,7 @@ def check_model(model): current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_roberta_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test parameter weight loading with tp>1. - """ + """Test parameter weight loading with tp>1.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") with vllm_runner( @@ -113,9 +109,7 @@ def check_model(model): current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." ) def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch): - """ - Test loading roberta-base model with no lm_head. - """ + """Test loading roberta-base model with no lm_head.""" # to use apply_model monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") model_name = "FacebookAI/roberta-base" diff --git a/tests/model_executor/test_oink_integration.py b/tests/model_executor/test_oink_integration.py index 2f37472b73ef..b16d71641369 100644 --- a/tests/model_executor/test_oink_integration.py +++ b/tests/model_executor/test_oink_integration.py @@ -62,7 +62,6 @@ def test_oink_availability_checks( expected_fused: bool, ): """Test OINK support detection with clean import state for each parameter set.""" - # Use spawn to run function in fresh process with clean imports # TODO migrate to spawn utility: # https://github.com/vllm-project/vllm/issues/41415 diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 30d9909a5a2e..cf3df85caa58 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -16,8 +16,7 @@ # Helper function to print input IDs with coalesced audio/video tokens. def print_input_ids(input_ids): - """ - Print input IDs, compressing consecutive special tokens. + """Print input IDs, compressing consecutive special tokens. - 151675: <|audio_pad|> - 151656: <|video_pad|> """ @@ -137,7 +136,6 @@ def test_qwen3_omni_get_updates_use_audio_in_video( mock_image_processor, ): """Test the get_updates_use_audio_in_video method directly.""" - from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerMultiModalProcessor, Qwen3OmniMoeThinkerProcessingInfo, diff --git a/tests/model_executor/test_qwen3_vl_mrope.py b/tests/model_executor/test_qwen3_vl_mrope.py index 90d9fd6e4ff8..a50bf319c938 100644 --- a/tests/model_executor/test_qwen3_vl_mrope.py +++ b/tests/model_executor/test_qwen3_vl_mrope.py @@ -51,8 +51,7 @@ class DummyConfig: def make_video_embedding( t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0 ): - """ - Helper function to make a video embedding for a given video size and pruning rate. + """Make a video embedding for a given video size and pruning rate. Args: t: Number of frames. @@ -64,6 +63,7 @@ def make_video_embedding( Returns: Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask) + """ unpruned_tokens_sequence = [] population = list(range(1, 100)) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 49986a991e2b..ed10955f963a 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -278,7 +278,7 @@ def test_routed_experts_capturer_single_dp_no_metadata(): def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): - """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" + """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) ctx = SimpleNamespace( @@ -295,7 +295,7 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): def test_routed_experts_capturer_dp_modular_local_tokens(): - """n == token_num_per_dp: topk is already local to this DP rank.""" + """N == token_num_per_dp: topk is already local to this DP rank.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) ctx = SimpleNamespace( @@ -326,7 +326,8 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises(): def test_routed_experts_attention_group_is_shared_and_fail_closed(): """Both sides key routing data by this gid, so it must skip non-full-attention - groups rather than defaulting to 0, and fail closed when none exists.""" + groups rather than defaulting to 0, and fail closed when none exists. + """ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) config = SimpleNamespace( kv_cache_groups=[ diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index e577690c92fc..d19ea0ed05b3 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -57,7 +57,8 @@ class TestMaybeRemapKvScaleName: def test_qkv_proj_k_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047""" + Regression test for https://github.com/vllm-project/vllm/issues/25047 + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -65,35 +66,36 @@ def test_qkv_proj_k_scale(self): def test_qkv_proj_v_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047""" + Regression test for https://github.com/vllm-project/vllm/issues/25047 + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.v_scale" def test_modelopt_k_proj_k_scale(self): - """ModelOpt format: k_proj.k_scale -> attn.k_scale""" + """ModelOpt format: k_proj.k_scale -> attn.k_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.k_scale" def test_modelopt_v_proj_v_scale(self): - """ModelOpt format: v_proj.v_scale -> attn.v_scale""" + """ModelOpt format: v_proj.v_scale -> attn.v_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.v_scale" def test_deprecated_kv_scale(self): - """Old format: kv_scale -> attn.k_scale (deprecated)""" + """Old format: kv_scale -> attn.k_scale (deprecated).""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.kv_scale", self.PARAMS_DICT ) assert result == "model.layers.0.self_attn.attn.k_scale" def test_default_bare_k_scale(self): - """Default format: .k_scale -> .attn.k_scale""" + """Default format: .k_scale -> .attn.k_scale.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_scale", self.PARAMS_DICT ) @@ -108,7 +110,8 @@ def test_non_scale_name_unchanged(self): def test_nvfp4_modelopt_k_proj_k_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): k_proj.k_scale -> attn.k_scale. - Validates that NVFP4 checkpoints are not broken by this change.""" + Validates that NVFP4 checkpoints are not broken by this change. + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT ) @@ -117,7 +120,8 @@ def test_nvfp4_modelopt_k_proj_k_scale(self): def test_nvfp4_modelopt_v_proj_v_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): v_proj.v_scale -> attn.v_scale. - Validates that NVFP4 checkpoints are not broken by this change.""" + Validates that NVFP4 checkpoints are not broken by this change. + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT ) @@ -125,7 +129,8 @@ def test_nvfp4_modelopt_v_proj_v_scale(self): def test_qwen3_vl_moe_qkv_proj_k_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" + Regression test for qwen3_vl_moe.py fix (same bug as #25047). + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -133,7 +138,8 @@ def test_qwen3_vl_moe_qkv_proj_k_scale(self): def test_qwen3_vl_moe_qkv_proj_v_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" + Regression test for qwen3_vl_moe.py fix (same bug as #25047). + """ result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) @@ -175,7 +181,8 @@ class TestKvCacheScaleMapper: """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent - so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`.""" + so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`. + """ def _mapper(self): # `get_cache_scale_mapper` does not use `self`; call it on the base @@ -265,7 +272,8 @@ def test_idempotent(self, name): def test_composes_with_qkv_mapper(self): """Applied together with a model's qkv/gate_up mapper, the regex scale rules run before the substr rename, so scales are normalized to `.attn.` - and regular projections are still fused correctly.""" + and regular projections are still fused correctly. + """ from vllm.model_executor.models.utils import WeightsMapper model_mapper = WeightsMapper( diff --git a/tests/models/inkling/test_mm_towers.py b/tests/models/inkling/test_mm_towers.py index d969c8101a29..51daa506c2b4 100644 --- a/tests/models/inkling/test_mm_towers.py +++ b/tests/models/inkling/test_mm_towers.py @@ -42,7 +42,8 @@ def _assert_close_ulps( compare against the reference's local ulp spacing. ``pre_act`` (the reference pre-activation) adds derivative-propagated slack: near ``gelu(x) ~ 0`` an ulp-level input flip legitimately moves the output by - many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps.""" + many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps. + """ g, r = got.float(), ref.float() tol = torch.clamp(max_ulps * _bf16_spacing(ref), min=atol) if pre_act is not None: diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index ec270c04fe6c..4346a7292846 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -123,7 +123,8 @@ def _check_writes_only_its_own_shard( device: torch.device, tp_size: int, rank: int ) -> None: """Two ranks that swapped both weight rows and write offsets still sum to - the right total, so inspect each slice before the final collective.""" + the right total, so inspect each slice before the final collective. + """ transform = _build_transform(device) runner = _tail_runner(transform, tp_size) group = get_tp_group().device_group diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py index 26cb76a55039..64cec452d25f 100644 --- a/tests/models/kimi_k3/test_aux_attn_res_stream.py +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -130,7 +130,8 @@ def test_disabled_reproduces_the_plain_residual_sum( def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): """The value the next layer reads is the mixture against *its* weights, - so the tap has to reach forward rather than use the current layer's.""" + so the tap has to reach forward rather than use the current layer's. + """ _set_last_rank(monkeypatch, True) _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) @@ -181,7 +182,8 @@ def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( ): """The kernel writes an applied delta back into the prefix in place, which would double-add it into the live residual stream, so the pending output - has to arrive already summed into the prefix with `delta` left None.""" + has to arrive already summed into the prefix with `delta` left None. + """ _set_last_rank(monkeypatch, True) prefix_sum = torch.tensor([1.0, 2.0]) pending = torch.tensor([0.5, 0.25]) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index fa810217b188..5dad5754b634 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -727,7 +727,8 @@ def test_cudagraph_capture_metadata_avoids_device_to_host_copy( """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing scheduled re-stages its FULL-graph metadata through it on every dummy step, so it must not synchronize the device. The host-side draft counts have to - come from `query_start_loc_cpu` and match the device-derived values.""" + come from `query_start_loc_cpu` and match the device-derived values. + """ num_speculative_tokens = 2 batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index c2c8e88e427c..307a572afd1f 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -60,7 +60,8 @@ def __init__(self, *, vllm_config, prefix): def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: """KV-shared layers have q_proj instead of a packed qkv_proj; their redundant K/V tensors in original checkpoints have no parameter and are - skipped rather than failing the load.""" + skipped rather than failing the load. + """ model = torch.nn.Module() model.config = SimpleNamespace(num_experts=0) model.start_layer, model.end_layer = 0, 2 @@ -92,7 +93,8 @@ def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: @pytest.mark.cpu_test def test_gemma3n_kv_shared_layer_mapper() -> None: """Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep - q_proj and drop the redundant K/V tensors original checkpoints ship.""" + q_proj and drop the redundant K/V tensors original checkpoints ship. + """ config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2) mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config) weights = [ diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index c5406fae48e0..b77670956aaf 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -164,8 +164,7 @@ def test_chunked_prefill_with_parallel_sampling( max_tokens: int, conv_state_layout: str, ) -> None: - """ - Tests chunked prefill in conjunction with n > 1. + """Tests chunked prefill in conjunction with n > 1. In this case, prefill is populated with decoding tokens and we test that it doesn't fail. @@ -199,8 +198,7 @@ def test_mamba_cache_cg_padding( max_tokens: int, conv_state_layout: str, ) -> None: - """ - This test is for verifying that mamba cache is padded to CG captured + """This test is for verifying that mamba cache is padded to CG captured batch size. If it's not, a torch RuntimeError will be raised because tensor dimensions aren't compatible. """ @@ -236,8 +234,7 @@ def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks( example_prompts, model: str, ) -> None: - """ - This test is for verifying that the hybrid inner state management doesn't + """This test is for verifying that the hybrid inner state management doesn't collapse in case where the number of incoming requests and finished_requests_ids is larger than the maximum mamba block capacity. @@ -263,8 +260,7 @@ def test_state_cleanup( example_prompts, model: str, ) -> None: - """ - This test is for verifying that the Hybrid state is cleaned up between + """This test is for verifying that the Hybrid state is cleaned up between steps. If it's not cleaned, an error would be expected. diff --git a/tests/models/language/generation/test_mistral.py b/tests/models/language/generation/test_mistral.py index fc774bc6e835..a21c5c355a66 100644 --- a/tests/models/language/generation/test_mistral.py +++ b/tests/models/language/generation/test_mistral.py @@ -274,7 +274,8 @@ def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None: def test_mistral_function_call_nested_json(): """Ensure that the function-name regex captures the entire outermost - JSON block, including nested braces.""" + JSON block, including nested braces. + """ # Create a minimal stub tokenizer that provides the few attributes the # parser accesses (`version` and `get_vocab`). diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py index d99aa4990982..b109b935dff6 100644 --- a/tests/models/language/pooling/test_max_tokens_per_doc.py +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for max_tokens_per_doc and max_tokens_per_query. -""" +"""Tests for max_tokens_per_doc and max_tokens_per_query.""" import json import os diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index 0b3561f90fd1..d470637477c0 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/tests/models/language/pooling/test_splade_sparse_pooler.py @@ -29,7 +29,8 @@ @torch.inference_mode def test_splade_pooler_matches_reference_formula(B, T, H, V): """Ensure SPLADESparsePooler forward() matches the mathematical formula: - log1p(relu(logits)) -> max over sequence length (after masking).""" + log1p(relu(logits)) -> max over sequence length (after masking). + """ torch.manual_seed(0) # Prepare [B] sequences of shape [T, H] diff --git a/tests/models/multimodal/generation/test_interleaved.py b/tests/models/multimodal/generation/test_interleaved.py index a773db19825e..07ce235853b6 100644 --- a/tests/models/multimodal/generation/test_interleaved.py +++ b/tests/models/multimodal/generation/test_interleaved.py @@ -23,11 +23,9 @@ def base_prompt(modalities_str: str) -> str: @pytest.mark.parametrize("dtype", ["float16"]) @pytest.mark.parametrize("max_tokens", [128]) def test_models(vllm_runner, model, dtype: str, max_tokens: int) -> None: - """ - This is a simple test to check if interleaved and non-interleaved prompts + """This is a simple test to check if interleaved and non-interleaved prompts give the same result. """ - image_cherry = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") image_stop = convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB") images = [image_cherry, image_stop] diff --git a/tests/models/multimodal/generation/test_maverick.py b/tests/models/multimodal/generation/test_maverick.py index d9afd93faaf3..0f6bec107fc2 100644 --- a/tests/models/multimodal/generation/test_maverick.py +++ b/tests/models/multimodal/generation/test_maverick.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Create a reduced-layer version of the Maverick model for testing purposes. +"""Create a reduced-layer version of the Maverick model for testing purposes. This script creates a new model with fewer layers by: 1. Loading the original Maverick model configuration @@ -39,7 +38,6 @@ def run_maverick_serving(model: str): """Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent options with reduced layers. """ - try: sampling_params = SamplingParams(temperature=0.8, top_p=0.95) @@ -71,8 +69,7 @@ def run_maverick_serving(model: str): def get_rope_layers_config(model_path: str) -> list[int]: - """ - Get the interleaved RoPE configuration from HuggingFace config + """Get the interleaved RoPE configuration from HuggingFace config. Args: model_path: Path to the local directory containing the reduced @@ -81,6 +78,7 @@ def get_rope_layers_config(model_path: str) -> list[int]: Returns: List of 0 or 1 indicating whether each layer uses RoPE and local attn 0 indicates that RoPE is not used while 1 indicates that RoPE is used. + """ config_path = Path(model_path) / "config.json" model_config = json.loads(config_path.read_text()) @@ -98,8 +96,7 @@ def create_reduced_maverick_model( vision_layers: int = 2, force_recreate: bool = False, ) -> str: - """ - Create a reduced-layer version of the Maverick model. + """Create a reduced-layer version of the Maverick model. Args: original_model_name: Name of the original Maverick model @@ -111,8 +108,8 @@ def create_reduced_maverick_model( Returns: Path to the created reduced model directory - """ + """ print( f"Creating reduced Maverick model with {text_layers} text layers and " f"{vision_layers} vision layers..." @@ -178,7 +175,6 @@ def create_reduced_config( original_config: Any, text_layers: int, num_experts: int, vision_layers: int ) -> dict[str, Any]: """Create a reduced configuration based on the original.""" - # Convert config to dictionary config_dict = original_config.to_dict() @@ -220,7 +216,6 @@ def create_reduced_config( def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None: """Copy tokenizer files from the original model.""" - try: tokenizer = AutoTokenizer.from_pretrained( original_model_name, trust_remote_code=True @@ -233,7 +228,6 @@ def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None: def create_preprocessor_config(original_config: Any, output_path: Path) -> None: """Create preprocessor_config.json for multimodal model.""" - # Try to load the original preprocessor config try: processor = AutoProcessor.from_pretrained( @@ -253,7 +247,6 @@ def create_reduced_safetensors( original_config: Any, reduced_config: dict[str, Any], output_path: Path ) -> None: """Create safetensors files with weights for the reduced model.""" - print("Generating synthetic weights for reduced model...") text_config = reduced_config["text_config"] @@ -276,7 +269,6 @@ def create_reduced_safetensors( def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Tensor]: """Create synthetic weights for the text model with MoE structure.""" - weights = {} vocab_size = text_config["vocab_size"] @@ -405,7 +397,6 @@ def create_vision_model_weights( vision_config: dict[str, Any], ) -> dict[str, torch.Tensor]: """Create synthetic weights for the vision model.""" - weights = {} hidden_size = vision_config["hidden_size"] @@ -473,8 +464,7 @@ def create_vision_model_weights( def create_shared_weights( text_config: dict[str, Any], vision_config: dict[str, Any] ) -> dict[str, torch.Tensor]: - """Create weights for shared components (vision-language connector)""" - + """Create weights for shared components (vision-language connector).""" weights = {} text_hidden_size = text_config["hidden_size"] @@ -492,7 +482,6 @@ def save_weights_to_safetensors( weights: dict[str, torch.Tensor], output_path: Path ) -> None: """Save weights to safetensors files and create index.""" - # Determine how to shard the weights max_shard_size = 5 * 1024 * 1024 * 1024 # 5GB per shard @@ -655,7 +644,6 @@ def test_dummy_maverick( def main(): """Main function to create and test the reduced model.""" - import argparse parser = argparse.ArgumentParser( diff --git a/tests/models/multimodal/generation/test_moondream3_moe.py b/tests/models/multimodal/generation/test_moondream3_moe.py index 79d18627656b..41654d4a7209 100644 --- a/tests/models/multimodal/generation/test_moondream3_moe.py +++ b/tests/models/multimodal/generation/test_moondream3_moe.py @@ -33,7 +33,8 @@ class _StubGate(torch.nn.Module): """Router stand-in, so the test does not depend on platform-specific - ``ReplicatedLinear`` weight post-processing.""" + ``ReplicatedLinear`` weight post-processing. + """ def __init__(self, hidden_size: int, num_experts: int): super().__init__() diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py index db0d3d06840a..69e41b21e342 100644 --- a/tests/models/multimodal/generation/test_nemotron_parse.py +++ b/tests/models/multimodal/generation/test_nemotron_parse.py @@ -31,8 +31,7 @@ def mask_bbox_tokens( output: tuple[list[int], str, SampleLogprobs], tokenizer: TokenizerLike, ) -> tuple[list[int], str, SampleLogprobs]: - """ - Always pass check_logprobs_close check for bounding box tokens + """Always pass check_logprobs_close check for bounding box tokens because it is reasonable for them to differ slightly. """ ignore_pattern = r"<[xy]_[\d.]+>" diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py index e163a25eb82f..4508135b2bb9 100644 --- a/tests/models/multimodal/generation/test_qwen2_vl.py +++ b/tests/models/multimodal/generation/test_qwen2_vl.py @@ -97,7 +97,7 @@ def batch_make_image_embeddings( processor, llm: VllmRunner, ) -> list[Qwen2VLPromptImageEmbeddingInput]: - """batched image embeddings for Qwen2-VL + """Batched image embeddings for Qwen2-VL. This will infer all images' embeddings in a single batch, and split the result according to input batches. @@ -108,7 +108,6 @@ def batch_make_image_embeddings( returns: `list[Qwen2VLPromptImageEmbeddingInput]` """ - image_batches_: list[Any] = image_batches[:] # convert single-image batches to multiple-image batches @@ -183,7 +182,7 @@ def get_image_embeds(model): def batch_make_video_embeddings( video_batches: PromptVideoInput, processor, llm: VllmRunner ) -> list[Qwen2VLPromptVideoEmbeddingInput]: - """batched video embeddings for Qwen2-VL + """Batched video embeddings for Qwen2-VL. A NDArray represents a single video's all frames. @@ -194,7 +193,6 @@ def batch_make_video_embeddings( - Single-video batches: `list[NDArray]` - Multiple-video batches: `list[list[NDArray]]` """ - video_batches_: list[Any] = video_batches[:] for idx in range(len(video_batches_)): diff --git a/tests/models/multimodal/generation/test_transformers_audio.py b/tests/models/multimodal/generation/test_transformers_audio.py index 919557987796..11ed8e51786a 100644 --- a/tests/models/multimodal/generation/test_transformers_audio.py +++ b/tests/models/multimodal/generation/test_transformers_audio.py @@ -81,7 +81,8 @@ def test_transformers_audio_generation( under pytest can cause (possibly flaky) hangs, so they are grouped under the same config. Using VLLM_WORKER_MULTIPROC_METHOD=spawn avoids the deadlock and allows worker processes to terminate cleanly, and release - GPU memory between test runs until the issue is fixed.""" + GPU memory between test runs until the issue is fixed. + """ # TODO: Remove monkeypatch once # https://github.com/vllm-project/vllm/issues/17676 is fixed. disable_envs_cache() diff --git a/tests/models/multimodal/generation/test_ultravox.py b/tests/models/multimodal/generation/test_ultravox.py index 964454331232..72f473835800 100644 --- a/tests/models/multimodal/generation/test_ultravox.py +++ b/tests/models/multimodal/generation/test_ultravox.py @@ -204,7 +204,6 @@ def test_variable_length_audio_batching( @pytest.mark.asyncio async def test_online_serving(client, audio_assets: AudioTestAssets): """Exercises online serving with/without chunked prefill enabled.""" - messages = [ { "role": "user", diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py index a6e8f3ff18a4..544c06020f66 100644 --- a/tests/models/multimodal/generation/test_voxtral.py +++ b/tests/models/multimodal/generation/test_voxtral.py @@ -78,7 +78,6 @@ def test_online_serving(vllm_runner, audio_assets: AudioTestAssets): Steps run sequentially so each releases the GPU before the next starts. """ - question = f"What's happening in these {len(audio_assets)} audio clips?" max_tokens = 10 audio_data = [asset.audio_and_sample_rate for asset in audio_assets] diff --git a/tests/models/multimodal/generation/vlm_utils/builders.py b/tests/models/multimodal/generation/vlm_utils/builders.py index 1b7e2347be2f..1f4783473d87 100644 --- a/tests/models/multimodal/generation/vlm_utils/builders.py +++ b/tests/models/multimodal/generation/vlm_utils/builders.py @@ -294,7 +294,8 @@ def sample_frames_with_video_metadata( def apply_image_size_scaling(image, size: float | tuple[int, int], size_type: SizeType): """Applies a size scaler to one image; this can be an image size factor, - which scales the image while maintaining the aspect ratio""" + which scales the image while maintaining the aspect ratio + """ # Special case for embeddings; if it's a tensor, it's only valid if we # are considering size factors at constant scale, i.e., we just clone # the tensor diff --git a/tests/models/multimodal/generation/vlm_utils/case_filtering.py b/tests/models/multimodal/generation/vlm_utils/case_filtering.py index cbe660425751..39d9d119bfa4 100644 --- a/tests/models/multimodal/generation/vlm_utils/case_filtering.py +++ b/tests/models/multimodal/generation/vlm_utils/case_filtering.py @@ -156,6 +156,7 @@ def get_wrapped_test_sizes( Args: test_info: Test configuration to be expanded. test_type: The type of test being filtered for. + """ # If it is an embedding test, we always use the EMBEDDING_SIZE_FACTORS if test_type == VLMTestType.EMBEDDING: diff --git a/tests/models/multimodal/generation/vlm_utils/core.py b/tests/models/multimodal/generation/vlm_utils/core.py index 207d3a3202a1..d665b6ad353b 100644 --- a/tests/models/multimodal/generation/vlm_utils/core.py +++ b/tests/models/multimodal/generation/vlm_utils/core.py @@ -200,7 +200,7 @@ def process_runner_outputs( def process_outputs(output_processor, model, outputs_per_image): - """Applies a model specific post-processor function to a runner's output""" + """Applies a model specific post-processor function to a runner's output.""" return [ [output_processor(res, model) for res in outputs] for outputs in outputs_per_image diff --git a/tests/models/multimodal/generation/vlm_utils/custom_inputs.py b/tests/models/multimodal/generation/vlm_utils/custom_inputs.py index 84109233685b..11ef99ea6443 100644 --- a/tests/models/multimodal/generation/vlm_utils/custom_inputs.py +++ b/tests/models/multimodal/generation/vlm_utils/custom_inputs.py @@ -22,6 +22,7 @@ def multi_image_multi_aspect_ratio_inputs(formatter: Callable[[str], str]): Args: formatter: model-specific prompt formatter. + """ stop_sign = IMAGE_ASSETS[0].pil_image cherry_blossom = IMAGE_ASSETS[1].pil_image @@ -65,6 +66,7 @@ def multi_video_multi_aspect_ratio_inputs( Args: formatter: model-specific prompt formatter. + """ video = sample_frames_from_video(VIDEO_ASSETS[0].np_ndarrays, num_frames) # Apply the selected formatter to the base prompts diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py index e85a49311a77..77beb64601af 100644 --- a/tests/models/multimodal/generation/vlm_utils/model_utils.py +++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py @@ -252,6 +252,7 @@ def qwen_prompt_path_encoder( tmp_path: Tempdir for test under consideration. prompt: Prompt with image placeholders. assets: list of image assets whose len equals the num placeholders. + """ # Ensure that the number of placeholders matches the number of assets; # If this is not true, the test is probably written incorrectly. @@ -507,15 +508,12 @@ def isaac_patch_hf_runner(hf_model: HfRunner) -> HfRunner: 2) Ensure IsaacModel.forward returns hidden_states for compatibility with hidden_states_to_seq_logprobs() """ - from perceptron.tensorstream import TextType from perceptron.tensorstream.ops import compute_mrope_pos_tensor, modality_mask from transformers.modeling_outputs import BaseModelOutputWithPast def compute_position_ids_input_ids(input_ids: torch.Tensor) -> torch.Tensor: - """ - Create 3D positional indices for token input. - """ + """Create 3D positional indices for token input.""" batch_size, seq_length = input_ids.shape position_ids = torch.arange(seq_length, device=input_ids.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) @@ -570,8 +568,7 @@ def patched_forward( cache_position=None, **kwargs, ): - """ - Forward pass with MRoPE position embeddings. + """Forward pass with MRoPE position embeddings. Computes position embeddings once and passes them through all layers. """ output_hidden_states = ( @@ -1184,7 +1181,6 @@ def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner": HfRunner.generate calls batch_decode on the full sequence (prompt + generated). """ - import io import pybase64 as base64 diff --git a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py index ac3eb6e61723..0fff8f40ac74 100644 --- a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py +++ b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py @@ -87,7 +87,7 @@ def _run_test( *, dtype: str, ) -> None: - """SET PYTHONPATH""" + """SET PYTHONPATH.""" # NOTE: take care of the order. run vLLM first, and then run HF. # vLLM needs a fresh new process without cuda initialization. # if we run HF first, the cuda initialization will be done and it diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index aa039735ced7..c71a2f08cbf1 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -108,8 +108,7 @@ def _normalize_image(image_val: str) -> str: def create_score_multimodal_param( content_parts: list[dict], ) -> list[ScoreMultiModalParam]: - """ - Create a ScoreMultiModalParam from a list of content dictionaries. + """Create a ScoreMultiModalParam from a list of content dictionaries. Each dict supports the following formats: - Text: {'text': 'content'} @@ -257,7 +256,7 @@ def test_model_text_image( model: str, dtype: str, ) -> None: - """Visual Documents Reranking""" + """Visual Documents Reranking.""" _run_test( hf_runner, vllm_runner, @@ -280,7 +279,7 @@ def test_model_text_text( model: str, dtype: str, ) -> None: - """Textual Documents Reranking""" + """Textual Documents Reranking.""" _run_test( hf_runner, vllm_runner, @@ -303,7 +302,7 @@ def test_model_image_text( model: str, dtype: str, ) -> None: - """Image Querying for Textual Documents""" + """Image Querying for Textual Documents.""" _run_test( hf_runner, vllm_runner, @@ -326,7 +325,7 @@ def test_model_image_image( model: str, dtype: str, ) -> None: - """Image Querying for Image Documents""" + """Image Querying for Image Documents.""" _run_test( hf_runner, vllm_runner, @@ -349,7 +348,7 @@ def test_model_text_mixed_documents( model: str, dtype: str, ) -> None: - """Text Query for Mixed Text and Image Documents""" + """Text Query for Mixed Text and Image Documents.""" _run_test( hf_runner, vllm_runner, diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index 9516ab95d52d..52802939cea5 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the LlamaNemotronVL model family: +"""Tests for the LlamaNemotronVL model family: - nvidia/llama-nemotron-embed-vl-1b-v2 (LlamaNemotronVLForCausalLM / embed) - nvidia/llama-nemotron-rerank-vl-1b-v2 (LlamaNemotronVLForSequenceClassification / rerank) diff --git a/tests/models/multimodal/processing/test_audio_in_video.py b/tests/models/multimodal/processing/test_audio_in_video.py index df6eea0f3865..f696d7a561a7 100644 --- a/tests/models/multimodal/processing/test_audio_in_video.py +++ b/tests/models/multimodal/processing/test_audio_in_video.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression tests for Qwen2.5-Omni and Qwen3-Omni audio-in-video processor +"""Regression tests for Qwen2.5-Omni and Qwen3-Omni audio-in-video processor caching. Tests the use_audio_in_video feature where audio is extracted from video and @@ -50,8 +49,7 @@ def create_mm_data(num_videos: int) -> dict[str, list]: @pytest.mark.parametrize("model_id", MODELS) @pytest.mark.parametrize("num_videos", [1, 2]) def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None: - """ - Regression test for https://github.com/vllm-project/vllm/pull/36800 + """Regression test for https://github.com/vllm-project/vllm/pull/36800. MultiModalProcessorSenderCache.get_and_update_item returns (None, updates) on a cache hit, so mm_kwargs["video"] items become None on the second call. @@ -117,8 +115,7 @@ def run(processor): @pytest.mark.parametrize("model_id", MODELS) def test_use_audio_in_video_without_audio_track(model_id: str) -> None: - """ - A video with no audio track, combined with `use_audio_in_video=True`, + """A video with no audio track, combined with `use_audio_in_video=True`, must raise a clear `ValueError` naming the mismatch, rather than an opaque `StopIteration`-derived error from the underlying HF processor. """ diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index bbc39ca171d3..1fbb2149a60d 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -33,9 +33,7 @@ def add_video_metadata(mm_data: MultiModalDataDict) -> MultiModalDataDict: - """ - Add metadata to video mm_data - """ + """Add metadata to video mm_data.""" def create_metadata(frames: np.ndarray): num_frames = len(frames) @@ -61,8 +59,7 @@ def create_metadata(frames: np.ndarray): def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict: - """ - Patch the multimodal data for GLM-ASR model. + """Patch the multimodal data for GLM-ASR model. GLM-ASR requires text and audio to match 1:1, so we limit audio to 1. """ if "audio" in mm_data: diff --git a/tests/models/multimodal/processing/test_deepseek_ocr.py b/tests/models/multimodal/processing/test_deepseek_ocr.py index 860f15ef3800..1851244be085 100644 --- a/tests/models/multimodal/processing/test_deepseek_ocr.py +++ b/tests/models/multimodal/processing/test_deepseek_ocr.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test for DeepSeek-OCR TensorSchema validation with empty images_crop. +"""Regression test for DeepSeek-OCR TensorSchema validation with empty images_crop. When using the Gundam preset (BASE_SIZE=1024, IMAGE_SIZE=640, CROP_MODE=True), images that are small enough to not require cropping produce an empty @@ -109,7 +108,8 @@ def test_populated_images_crop_large_image(self, processor): def test_mismatched_image_size_raises(self, processor): """Deliberately wrong image_size binding should still be caught - by TensorSchema validation.""" + by TensorSchema validation. + """ small_image = Image.new("RGB", (100, 100), color="green") result = processor( @@ -143,7 +143,8 @@ class TestDeepseekOCRInputValidation: def test_mismatched_image_count_raises_value_error(self, processor): """Two ```` tokens with one image must raise ValueError, not AssertionError (which would yield HTTP 500 and, under ``-O``, silently - drop the middle text segment).""" + drop the middle text segment). + """ image = Image.new("RGB", (100, 100), color="red") with pytest.raises(ValueError, match="does not match number of images"): processor(prompt=" and ", images=[image]) diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index ce2ab2adfb41..adf7219db1b7 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -152,8 +152,7 @@ def test_video_loader_consistency( fps: int, backend: str, ): - """ - Ensure dynamic video loader (pre-sampled by loader) and normal video + """Ensure dynamic video loader (pre-sampled by loader) and normal video loader (post-sampled by processor) produce same video processing outputs. """ ctx = build_model_context( diff --git a/tests/models/multimodal/processing/test_llava_next.py b/tests/models/multimodal/processing/test_llava_next.py index d7eaa35dd8d9..dc1ad425fb7d 100644 --- a/tests/models/multimodal/processing/test_llava_next.py +++ b/tests/models/multimodal/processing/test_llava_next.py @@ -116,8 +116,7 @@ def _test_image_prompt_replacements( num_imgs: int, image_sizes: list[ImageSize], ) -> None: - """ - Ensure LlavaNextMultiModalProcessor + """Ensure LlavaNextMultiModalProcessor handles prompt replacement properly for input images. """ failed_size_excs = list[tuple[ImageSize, Exception]]() diff --git a/tests/models/multimodal/processing/test_llava_next_video.py b/tests/models/multimodal/processing/test_llava_next_video.py index c4fdecd21ac7..512440ee8ee6 100644 --- a/tests/models/multimodal/processing/test_llava_next_video.py +++ b/tests/models/multimodal/processing/test_llava_next_video.py @@ -16,7 +16,8 @@ class _StubModel: """Carries only the attributes the token-count method reads from `self`, so the real methods can be exercised without constructing - the full `nn.Module` (vision tower, language model, etc.).""" + the full `nn.Module` (vision tower, language model, etc.). + """ patch_grid_length: int pooled_grid_length: int @@ -31,7 +32,8 @@ class _StubModel: def test_num_mm_tokens_match_real_config(model_id): """The stored grid lengths must match what `__init__` derives from the real HF config, and the LoRA token counts must preserve the - frame-level scaling exactly.""" + frame-level scaling exactly. + """ ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1}) config = ctx.model_config.hf_config diff --git a/tests/models/multimodal/processing/test_llava_onevision.py b/tests/models/multimodal/processing/test_llava_onevision.py index 4e8a8103a592..388316aca289 100644 --- a/tests/models/multimodal/processing/test_llava_onevision.py +++ b/tests/models/multimodal/processing/test_llava_onevision.py @@ -122,8 +122,7 @@ def _test_image_prompt_replacements( num_imgs: int, image_sizes: list[ImageSize], ) -> None: - """ - Ensure LlavaOnevisionMultiModalProcessor + """Ensure LlavaOnevisionMultiModalProcessor handles prompt replacement properly for input images. """ failed_size_excs = list[tuple[ImageSize, Exception]]() diff --git a/tests/models/multimodal/processing/test_molmo2.py b/tests/models/multimodal/processing/test_molmo2.py index 5633d47a97dc..98ad241650e0 100644 --- a/tests/models/multimodal/processing/test_molmo2.py +++ b/tests/models/multimodal/processing/test_molmo2.py @@ -70,7 +70,8 @@ def test_dummy_video_num_frames_override_honors_min_of_two( ): """A ``num_frames`` override below 2 must be ignored (the model needs at least 2 frames), matching the "cannot be less than 2, will be ignored" - warning.""" + warning. + """ builder = object.__new__(Molmo2DummyInputsBuilder) builder.info = SimpleNamespace( get_hf_processor=lambda: SimpleNamespace( diff --git a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py index af12a14a45a0..ebbc7fab421d 100644 --- a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py +++ b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are +"""Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are correctly assigned to audio/image/video token positions. Regression test for: https://github.com/vllm-project/vllm/issues/34506 @@ -40,8 +39,7 @@ def _mm_embed(shape: tuple[int, ...], value: float, modality: str) -> torch.Tens def make_token_seq( audio_n: int, image_n: int, video_n: int, text_prefix: int = 3, text_sep: int = 2 ): - """ - Build a flat token sequence: + """Build a flat token sequence: [text_prefix] [AUDIO * audio_n] [text_sep] [IMAGE * image_n] [text_sep] [VIDEO * video_n] [text_sep] Returns (input_ids tensor, is_multimodal mask, positions dict). @@ -67,9 +65,8 @@ def make_token_seq( def make_interleaved_seq( video_chunks: list[int], audio_chunks: list[int], text_prefix: int = 2 ): - """ - Build an interleaved sequence like use_audio_in_video: - [text] [V*v0] [A*a0] [V*v1] [A*a1] ... + """Build an interleaved sequence like use_audio_in_video: + [text] [V*v0] [A*a0] [V*v1] [A*a1] ... """ tokens = [TEXT_TOKEN_ID] * text_prefix for v, a in zip(video_chunks, audio_chunks): @@ -140,8 +137,7 @@ def test_multi_video_with_boundary_tokens(self): ) def test_batched_non_interleaved_no_false_positive(self): - """ - Regression test for https://github.com/vllm-project/vllm/issues/35394. + """Regression test for https://github.com/vllm-project/vllm/issues/35394. 5 identical non-interleaved mixed-modality requests batched together: each has [audio][image][video] in separate blocks with text between them. @@ -172,8 +168,7 @@ def test_batched_non_interleaved_no_false_positive(self): def make_mock_model(hidden: int = 8): - """ - Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration + """Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration that has enough structure to run embed_input_ids. """ from vllm.model_executor.models.qwen2_5_omni_thinker import ( @@ -236,8 +231,7 @@ def fake_super_embed( def build_mm_embeds( audio_n, image_n, video_n, hidden, audio_val=10.0, image_val=20.0, video_val=30.0 ): - """ - Build multimodal_embeddings list in position order (audio, image, video). + """Build multimodal_embeddings list in position order (audio, image, video). Each embedding is filled with a distinct constant so we can verify placement. """ embs = [] @@ -252,8 +246,7 @@ def build_mm_embeds( class TestEmbedInputIds: def _run(self, audio_n, image_n, video_n, hidden=8): - """ - Run embed_input_ids for a non-interleaved mixed-modality sequence. + """Run embed_input_ids for a non-interleaved mixed-modality sequence. Returns (result_embeds, input_ids, is_multimodal). """ input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n) @@ -288,8 +281,7 @@ def test_video_only(self): ) def test_mixed_modalities_audio_goes_to_audio_pos(self): - """ - Regression test for GitHub issue #34506: + """Regression test for GitHub issue #34506: With audio + image + video (non-interleaved), audio positions must receive audio embeddings (not image or video embeddings). """ @@ -329,8 +321,7 @@ def test_text_positions_unchanged(self): ) def test_interleaved_use_audio_in_video(self): - """ - Interleaved (use_audio_in_video): video chunks interleaved with audio. + """Interleaved (use_audio_in_video): video chunks interleaved with audio. Video embeddings must go to video positions, audio to audio positions. """ hidden = 8 diff --git a/tests/models/multimodal/processing/test_qwen2_vl.py b/tests/models/multimodal/processing/test_qwen2_vl.py index 5f3f5b05e106..4976b05c3e0a 100644 --- a/tests/models/multimodal/processing/test_qwen2_vl.py +++ b/tests/models/multimodal/processing/test_qwen2_vl.py @@ -232,7 +232,6 @@ def test_mm_device_do_normalize( num_imgs: int, ): """Ensure that enable mm_device_do_normalize yields the correct result.""" - ctx = build_model_context( model_id, limit_mm_per_prompt={"image": num_imgs}, diff --git a/tests/models/multimodal/processing/test_qwen3_omni.py b/tests/models/multimodal/processing/test_qwen3_omni.py index e7a7e2de87a0..301471f45a7b 100644 --- a/tests/models/multimodal/processing/test_qwen3_omni.py +++ b/tests/models/multimodal/processing/test_qwen3_omni.py @@ -25,8 +25,7 @@ def test_processor_with_audio_sample_rate( audio_sample_rate: int, audio_duration_sec: float, ) -> None: - """ - Test that vLLM's processor generates expected outputs with audio_sample_rate. + """Test that vLLM's processor generates expected outputs with audio_sample_rate. This validates that the processor correctly handles audio_sample_rate passed via hf_processor_mm_kwargs and generates audio tokens. @@ -70,8 +69,7 @@ def test_processor_with_audio_sample_rate( @pytest.mark.parametrize("model_id", ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]) def test_longer_audio_generates_more_tokens(model_id: str) -> None: - """ - Test that longer audio generates more tokens than shorter audio. + """Test that longer audio generates more tokens than shorter audio. This validates that audio_sample_rate is being used correctly by checking that audio duration affects token count as expected. diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index c58904a9e2c9..47071c2f1bed 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -31,7 +31,8 @@ def _build_video_mm_data( ``total_num_frames`` is set equal to the ndarray frame count so that HF's ``sample_frames`` indices stay within bounds of the - actual tensor that is passed.""" + actual tensor that is passed. + """ video = np.zeros((num_frames, height, width, 3), dtype=np.uint8) metadata = { "fps": original_fps, @@ -237,7 +238,8 @@ def _build_video_embeds_mm_data( grid_thw: tuple[int, int, int] = (2, 4, 4), ) -> dict[str, Any]: """Create an embeds-only video item as an EPD consumer receives it: - pre-computed embeddings plus the metadata published by the encoder.""" + pre-computed embeddings plus the metadata published by the encoder. + """ import torch t, h, w = grid_thw @@ -256,7 +258,8 @@ def test_processor_video_embeds_with_timestamps(model_id: str) -> None: """Embeds-only video input must size the placeholder range from the grid and the real timestamps published by the encoder (EC consumer path); synthesized or missing timestamps would change the token count - and break embedding merging downstream.""" + and break embedding merging downstream. + """ # `build_model_context` forces enable_mm_embeds off, so build the # config directly; keep the same online-availability skip behavior. HF_EXAMPLE_MODELS.find_hf_info(model_id).check_available_online(on_fail="skip") @@ -301,7 +304,8 @@ def test_processor_video_embeds_with_timestamps(model_id: str) -> None: def test_processor_video_embeds_missing_timestamps(model_id: str) -> None: """Timestamps are required metadata for video embeds: they size the placeholder range, so omitting them must fail loudly at parse time - instead of silently producing a wrong prompt.""" + instead of silently producing a wrong prompt. + """ ctx = build_model_context( model_id, limit_mm_per_prompt={"image": 0, "video": 1}, diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py index 61517ea0af19..14007c87403d 100644 --- a/tests/models/multimodal/processing/test_transformers_audio.py +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -116,7 +116,8 @@ def test_audio_multimodal_processor(model_id, processor_cls): @pytest.mark.parametrize("separator", [" and ", ""]) def test_audio_multiple_inputs(separator, processor_cls): """Multiple audios per prompt are each detected as a separate placeholder - and multi-modal item by the Transformers modelling backend.""" + and multi-modal item by the Transformers modelling backend. + """ model_id = "ibm-granite/granite-speech-3.3-2b" mm_processor = create_processor(model_id, processor_cls) diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index f5de7d66f4ff..731794da2667 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -83,7 +83,8 @@ def _process_two_images(processor_cls, separator: str): @pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) def test_image_multiple_inputs(processor_cls): """Multiple images per prompt are each detected as a separate placeholder - and multi-modal item by the Transformers modelling backend.""" + and multi-modal item by the Transformers modelling backend. + """ result = _process_two_images(processor_cls, separator="\n and ") assert len(result["mm_placeholders"]["image"]) == 2 @@ -102,7 +103,8 @@ def test_image_adjacent_inputs(processor_cls): @pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) def test_batch_padding_removed_from_image_items(processor_cls): """Emu3 pads every image up to the largest in the batch, which would leave an - item's data dependent on what it was processed with and so uncacheable.""" + item's data dependent on what it was processed with and so uncacheable. + """ mm_processor = create_processor("BAAI/Emu3-Chat-hf", processor_cls) image_token = mm_processor.info.get_hf_processor().image_token @@ -144,7 +146,8 @@ def _process_one_gemma3_image(processor_cls): @offsets_only def test_non_embedding_tokens_excluded_from_placeholders(): """Gemma3 wraps each image in text that carries no embeddings, which must be - inside the placeholder range but masked out of it.""" + inside the placeholder range but masked out of it. + """ _, result = _process_one_gemma3_image(OffsetsMultiModalProcessor) (placeholder,) = result["mm_placeholders"]["image"] @@ -156,7 +159,8 @@ def test_legacy_placeholders_hold_only_image_tokens(): """The legacy path spans whatever `mm_token_type_ids` attributes to the image, which for Gemma3 excludes the text wrapping it, unlike the replacement the offsets path spans. Gemma3 is also the sharp case for the mask: its `image_token_id` is the - marker in the unexpanded prompt, not the token the expansion repeats.""" + marker in the unexpanded prompt, not the token the expansion repeats. + """ hf_processor, result = _process_one_gemma3_image(LegacyMultiModalProcessor) (placeholder,) = result["mm_placeholders"]["image"] @@ -172,7 +176,8 @@ def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): embeddings. Those belong inside the placeholder and masked out, because the token count the processor reports is over the whole span. Idefics3 also refuses a prompt holding `` when no images are passed, which is how the offsets path has to - tokenize it before splicing in the expansion.""" + tokenize it before splicing in the expansion. + """ mm_processor = create_processor( "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls ) @@ -192,7 +197,8 @@ def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): @offsets_only def test_missing_replacement_offsets_names_the_processor(): """A processor that reports no replacement offsets cannot be served, which must - be said plainly rather than surfacing later as a field config mismatch.""" + be said plainly rather than surfacing later as a field config mismatch. + """ model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" mm_processor = create_processor(model_id, OffsetsMultiModalProcessor) hf_processor_cls = type(mm_processor.info.get_hf_processor()) @@ -276,7 +282,8 @@ def process(): def test_spliced_prompt_matches_hf_expansion(model_id, prompt): """The offsets path splices the expansion into a prompt tokenized without any multi-modal data, so its token ids have to come out the same as the ones the HF - processor produces itself, which is what the legacy path returns.""" + processor produces itself, which is what the legacy path returns. + """ prompt_ids = [] for processor_cls in (LegacyMultiModalProcessor, OffsetsMultiModalProcessor): mm_processor = create_processor(model_id, processor_cls) @@ -298,7 +305,8 @@ def test_spliced_prompt_matches_hf_expansion(model_id, prompt): def test_nested_image_fields_split_per_image(processor_cls): """Idefics3 returns image fields with a leading batch dimension, putting the rows belonging to each image one dimension further in. Slicing the batch dimension - instead handed the first image every row and the second an empty tensor.""" + instead handed the first image every row and the second an empty tensor. + """ mm_processor = create_processor( "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls ) diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py index dc49036e8352..d748e89f20f9 100644 --- a/tests/models/multimodal/processing/transformers_backend.py +++ b/tests/models/multimodal/processing/transformers_backend.py @@ -29,7 +29,8 @@ def create_processor(model_id: str, processor_cls): """Build a processor directly, because the registry only ever builds the one the - installed transformers version selects, leaving the other path untested.""" + installed transformers version selects, leaving the other path untested. + """ model_config = ModelConfig(model=model_id, model_impl="transformers") ctx = InputProcessingContext( model_config, cached_tokenizer_from_config(model_config) @@ -40,7 +41,8 @@ def create_processor(model_id: str, processor_cls): def create_cached_processor(model_id: str, processor_cls): """Build a processor backed by a real multi-modal processor cache, and hand the - cache back so a test can check it was actually used.""" + cache back so a test can check it was actually used. + """ model_config = ModelConfig(model=model_id, model_impl="transformers") model_config.multimodal_config.mm_processor_cache_gb = 4 ctx = InputProcessingContext( diff --git a/tests/models/multimodal/test_cohere_asr.py b/tests/models/multimodal/test_cohere_asr.py index 85a9b59dd513..a9c45f84055c 100644 --- a/tests/models/multimodal/test_cohere_asr.py +++ b/tests/models/multimodal/test_cohere_asr.py @@ -18,7 +18,8 @@ def test_get_num_audio_tokens_streaming_estimate(audio_duration_s, expected_tokens): """The duration-based estimate must convert ``window_stride`` (seconds) to a sample hop and divide by the encoder subsampling factor. Values are pinned to - concrete numbers so a rounding regression is actually caught.""" + concrete numbers so a rounding regression is actually caught. + """ model_config = SimpleNamespace( hf_config=SimpleNamespace( preprocessor={"window_stride": 0.01, "sample_rate": 16000}, diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py index b762ab2a13b6..a167fe205db7 100644 --- a/tests/models/multimodal/test_mapping.py +++ b/tests/models/multimodal/test_mapping.py @@ -78,7 +78,8 @@ def test_cosmos3_modelopt_quantizer_weights_mapper(): """ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers (``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable ``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former - (it has no parameter for them) while keeping the latter.""" + (it has no parameter for them) while keeping the latter. + """ from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper @@ -215,7 +216,7 @@ def test_bailing_vl_mapper_handles_module_and_parameter_names(): def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]: - """Create weights from safetensors checkpoint metadata""" + """Create weights from safetensors checkpoint metadata.""" metadata = try_get_safetensors_metadata(repo) weight_names = list(metadata.weight_map.keys()) with torch.device("meta"): @@ -223,9 +224,7 @@ def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]: def create_dummy_base_model(repo: str, model_arch: str) -> PreTrainedModel: - """ - Create weights from a dummy meta deserialized hf base model with name conversion - """ + """Create weights from a dummy meta-deserialized HF base model, with renaming.""" config = AutoConfig.from_pretrained(repo) with torch.device("meta"): model = AutoModel.from_config(config) @@ -233,9 +232,7 @@ def create_dummy_base_model(repo: str, model_arch: str) -> PreTrainedModel: def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel: - """ - Create weights from a dummy meta deserialized hf model with name conversion - """ + """Create weights from a dummy meta deserialized hf model with name conversion.""" model_cls: PreTrainedModel = getattr(transformers, model_arch) config = AutoConfig.from_pretrained(repo) with torch.device("meta"): diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 75900cde53a3..c64073ffe2be 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -81,7 +81,8 @@ class _FakeTensor: """Sentinel stand-in for torch.Tensor in load_weights tests. Supports the .detach().clone() chain used by load_weights for buffered mm weights; both methods return self so identity (and the existing equality - assertions) are preserved through cloning.""" + assertions) are preserved through cloning. + """ def detach(self): return self @@ -165,7 +166,8 @@ def _make_mm_items_with_video_bytes( def test_extract_audio_from_videos_passes_max_duration(): """_extract_audio_from_videos must forward VLLM_MAX_AUDIO_DECODE_DURATION_S - to load_audio_pyav so decompression-bomb audio is rejected.""" + to load_audio_pyav so decompression-bomb audio is rejected. + """ dummy_audio = (np.zeros(16000, dtype=np.float32), 16000.0) mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64) @@ -185,7 +187,8 @@ def test_extract_audio_from_videos_passes_max_duration(): def test_extract_audio_from_videos_rejects_oversized_audio(): """When load_audio_pyav raises due to duration limit the video is - marked as having no audio instead of crashing the server.""" + marked as having no audio instead of crashing the server. + """ mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64) processor = object.__new__(NanoNemotronVLMultiModalProcessor) diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index 209f15275902..10b0777f5e3a 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end accuracy test for GPT-OSS model quantization. +"""End-to-end accuracy test for GPT-OSS model quantization. Config: Task: gsm8k_platinum @@ -170,8 +169,7 @@ def test_aiter_mxfp4_moe_ignores_padded_rows( monkeypatch: pytest.MonkeyPatch, dist_init, ) -> None: - """ - Garbage in cudagraph padding rows must not reach the unpadded outputs. + """Garbage in cudagraph padding rows must not reach the unpadded outputs. A cudagraph replay of a size-`TOKENS_PADDED` graph driven by `TOKENS_UNPADDED` real tokens leaves the trailing padding row containing diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 9d30365b8151..c5ed11d0c693 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -144,7 +144,8 @@ def test_qwen4_exp_mtp_override_sets_draft_config( def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> None: """PLE needs raw input_ids, which non-first pipeline ranks never see. The rest of the architecture is PP-capable, so the refusal must be conditional - -- and must land before the engine spends time loading weights.""" + -- and must land before the engine spends time loading weights. + """ vllm_config = SimpleNamespace( model_config=SimpleNamespace( hf_text_config=_text_config(ple_layer_ids=ple_layer_ids), diff --git a/tests/models/registry.py b/tests/models/registry.py index cb39ed7ce4f3..8adb9ed15241 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -127,8 +127,7 @@ def check_transformers_version( check_min_version: bool = True, check_max_version: bool = True, ) -> str | None: - """ - If the installed transformers version does not meet the requirements, + """If the installed transformers version does not meet the requirements, perform the given action. """ if ( @@ -185,9 +184,7 @@ def check_available_online( *, on_fail: Literal["error", "skip"], ) -> None: - """ - If the model is not available online, perform the given action. - """ + """If the model is not available online, perform the given action.""" if not self.is_available_online: msg = "Model is not available online" diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 77b05101d685..7007eb5aca59 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -103,7 +103,8 @@ def test_fi_backend_strings_are_registered_mega_moe_backends(): @pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend): """An FI moe_ep backend with a non-DSv4 model must fail at config time - instead of silently falling through to the generic FusedMoE path.""" + instead of silently falling through to the generic FusedMoE path. + """ with pytest.raises(ValueError, match="only supported for DeepSeek-V4"): validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"]) @@ -118,8 +119,9 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]], ) def test_native_deep_gemm_mega_moe_not_arch_gated(architectures): - """vLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it); - models validate their own constraints at construction time.""" + """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it); + models validate their own constraints at construction time. + """ validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures) @@ -132,7 +134,8 @@ def test_all_mega_backends_get_sequence_parallel_moe(moe_backend): """Every mega backend must qualify for sequence-parallel MoE at TP>1/EP: the predicate once matched only the native backend string, which silently ran the fi backends full-batch with an all-reduce on - every rank — 0.42-0.65x native e2e at TP8.""" + every rank — 0.42-0.65x native e2e at TP8. + """ from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel vllm_config = SimpleNamespace( @@ -176,7 +179,8 @@ def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch): """The runtime must not rederive the device from LOCAL_RANK/rank: under a remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the - device it already bound via BootstrapConfig.device.""" + device it already bound via BootstrapConfig.device. + """ import vllm.utils.flashinfer_moe_ep as mod pg = object() @@ -241,7 +245,8 @@ def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo(): def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales(): """One 32-element scale group per row: low nibble is the even element, - high nibble the odd one, ue8m0 scale applies to the whole group.""" + high nibble the odd one, ue8m0 scale applies to the whole group. + """ packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16) sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 9a5723d233a1..31aad5395f4b 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -783,7 +783,8 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): """The loader invokes the model-level PWAL hook for every load format, so it must finalize megamoe + mhc broadcast weights to cover dummy - load, which skips load_weights().""" + load, which skips load_weights(). + """ calls = [] stub = SimpleNamespace( model=SimpleNamespace( @@ -800,7 +801,8 @@ def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): """MTP/DSpark drafters load as their own top-level models, so each needs its own PWAL hook now that the megamoe forward no longer finalizes - weights lazily on first use.""" + weights lazily on first use. + """ calls = [] mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) DeepSeekV4MTP.process_weights_after_loading(mtp) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index d30eb9e7c91d..ee381ee36b22 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -215,7 +215,8 @@ def test_v41_dspark_loads_linear_scales( ): """Checkpoint ``.scale`` maps to the quant method's scale parameter and loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static - loader (see tests/quantization/test_modelopt.py), not in load_weights.""" + loader (see tests/quantization/test_modelopt.py), not in load_weights. + """ from vllm.models.deepseek_v41.nvidia import dspark mxfp8 = scale_dtype != torch.float32 diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index cc982f699686..e3a08980e90c 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -64,7 +64,6 @@ def can_initialize( The spawn process causes the _initialize_kv_caches_v1 function below to become ineffective. """ - model_info = EXAMPLE_MODELS.get_hf_info(model_arch) model_info.check_available_online(on_fail="skip") model_info.check_transformers_version( @@ -199,13 +198,13 @@ def _initialize_kv_caches_v1(self, vllm_config): @pytest.mark.parametrize("model_arch", MINIMAL_MODEL_ARCH_LIST) def test_can_initialize_small_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch): - """Test initializing small subset of supported models""" + """Test initializing small subset of supported models.""" can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS) @pytest.mark.parametrize("model_arch", OTHER_MODEL_ARCH_LIST) def test_can_initialize_large_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch): - """Test initializing large subset of supported models + """Test initializing large subset of supported models. This test covers the complement of the tests covered in the "small subset" test. diff --git a/tests/models/test_qwen3_5_mtp_config.py b/tests/models/test_qwen3_5_mtp_config.py index 6cf54dbc95a8..3c098d3ad55e 100644 --- a/tests/models/test_qwen3_5_mtp_config.py +++ b/tests/models/test_qwen3_5_mtp_config.py @@ -73,7 +73,8 @@ def test_mtp_override_extracts_n_predict_from_multimodal_wrapper( model_type: str, expected_arch: str ) -> None: """Verify that multimodal wrapper checkpoints with mtp_num_hidden_layers - in text_config resolve n_predict and architecture correctly.""" + in text_config resolve n_predict and architecture correctly. + """ cfg = SpeculativeConfig.hf_config_override( _multimodal_wrapper_mtp_config(model_type, mtp_layers=2) ) @@ -93,7 +94,8 @@ def test_mtp_override_top_level_precedence_over_nested_text_config( model_type: str, expected_arch: str ) -> None: """Verify that an explicit top-level mtp_num_hidden_layers takes precedence - over a nested text_config value.""" + over a nested text_config value. + """ cfg = _multimodal_wrapper_mtp_config(model_type, mtp_layers=2) cfg.mtp_num_hidden_layers = 3 overridden = SpeculativeConfig.hf_config_override(cfg) @@ -113,7 +115,8 @@ def test_mtp_override_downloads_real_hf_hub_configs( model_id: str, expected_arch: str ) -> None: """Verify that unmodified real-world checkpoints downloaded via - AutoConfig.from_pretrained resolve n_predict=1 from text_config.""" + AutoConfig.from_pretrained resolve n_predict=1 from text_config. + """ hf_config = AutoConfig.from_pretrained(model_id) cfg = SpeculativeConfig.hf_config_override(hf_config) assert cfg.model_type == "qwen3_5_mtp" diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 02903504c684..f9e7d691769b 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -197,7 +197,8 @@ def test_merge_multimodal_embeddings_no_sync(): def test_get_rename_mapper_keeps_only_renames(): """`None` means "do not load", which is meaningless to the consumers of this mapper (LoRA name parsing, quantization config layer lists), and - applying it would silently shrink their lists.""" + applying it would silently shrink their lists. + """ mapper = WeightsMapper( orig_to_new_regex={re.compile(r"^drop_regex\."): None}, orig_to_new_substr={"drop_substr": None, "keep_substr": "kept"}, diff --git a/tests/models/test_vision.py b/tests/models/test_vision.py index 24260e2ef804..b08661abc958 100644 --- a/tests/models/test_vision.py +++ b/tests/models/test_vision.py @@ -41,9 +41,7 @@ def test_resolve_visual_encoder_outputs( select_layers, num_layers_loaded, max_possible_layers, expected_features ): - """ - Test that offsets are correctly handled for vision feature layers. - """ + """Test that offsets are correctly handled for vision feature layers.""" encoder_outputs = [torch.tensor([idx]) for idx in range(num_layers_loaded + 1)] output_tensor = resolve_visual_encoder_outputs( encoder_outputs=encoder_outputs, @@ -94,11 +92,9 @@ def test_run_dp_sharded_vision_model(batch_size: int): def run_dp_sharded_vision_model_vs_direct( local_rank: int, world_size: int, batch_size: int, master_port: int ): - """ - Test that run_dp_sharded_vision_model produces the same results as + """Test that run_dp_sharded_vision_model produces the same results as calling the model directly. """ - # Set random seed for reproducibility set_random_seed(0) @@ -282,8 +278,7 @@ def test_run_dp_sharded_mrope_vision_model(batch_size: int): def run_dp_sharded_mrope_vision_model_vs_direct( local_rank: int, world_size: int, batch_size: int, master_port: int ): - """ - Test that run_dp_sharded_mrope_vision_model produces the same results as + """Test that run_dp_sharded_mrope_vision_model produces the same results as calling the model directly. """ # Set random seed for reproducibility @@ -517,7 +512,8 @@ def _reference_input_norm( def test_fused_input_norm_matches_reference(num_patches: int): """FusedInputNorm must equal the plain affine, including for num_patches above the cuDNN batch-norm grid limit (~65535) that previously raised - CUDNN_STATUS_INTERNAL_ERROR (issue #51717).""" + CUDNN_STATUS_INTERNAL_ERROR (issue #51717). + """ channel = 3 patch_size = 14 * 14 image_mean = [0.48145466, 0.4578275, 0.40821073] diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index bcad62765a76..d26d742dd3d2 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -251,7 +251,8 @@ def forward(self, x): class PlainQKV(nn.Module): """q/k/v returned as a bare tuple: no reshape stands between the calls and - the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it.""" + the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it. + """ def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8): super().__init__() @@ -338,7 +339,8 @@ class GuardedVAttention(FakeAttention): """Gemma 4-style: `v_proj` is called once but also read in a `None` guard. The guard is a fusion invariant (the fuser only binds instances where the - projection exists), so it folds to `True` and the call still fuses.""" + projection exists), so it folds to `True` and the call still fuses. + """ def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -368,7 +370,8 @@ class TruthyGuardAttention(FakeAttention): An `nn.Linear` is always truthy, so the guard folds to `True` and the call still fuses (the projection is a proven invariant, exactly as for a `None` - guard).""" + guard). + """ def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -399,7 +402,8 @@ class StatementGuardVAttention(FakeAttention): The `if`'s test is not an `ast.Compare`, so the identity-check scan skips it; `_in_boolean_context` folds the test to `True` (an `nn.Linear` is always truthy) and the call inside the now-`if True:` body still fuses, with the - live `else` surviving the rewrite.""" + live `else` surviving the rewrite. + """ def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -428,7 +432,8 @@ class BranchedKVAttention(FakeAttention): The three calls do not share a block, so the fused GEMM must be hoisted to the innermost block that dominates all of them (the function body). The - live `if`/`else` must survive the rewrite.""" + live `if`/`else` must survive the rewrite. + """ is_kv_shared_layer = False @@ -459,7 +464,8 @@ class RebindArgAttention(FakeAttention): The branch is dead at trace time so the match succeeds, but the source rewrite must refuse: hoisting the GEMM above the rebind would feed k/v a - different input. Result: no fusion.""" + different input. Result: no fusion. + """ recompute = False @@ -490,7 +496,8 @@ class InPlaceMutatedArgAttention(FakeAttention): Like `RebindArgAttention`, but the mutation keeps the name in a `Load` context (`hidden_states.mul_(2)`), so a bare `Name`-store check would miss it. The rewrite must still refuse: hoisting the GEMM above the mutation changes - the value k/v see. Result: no fusion.""" + the value k/v see. Result: no fusion. + """ recompute = False @@ -520,7 +527,8 @@ class AliasMutatedArgAttention(FakeAttention): `hidden_states.view(-1)` shares storage with its base, so `alias.mul_(2)` changes what k/v read without ever naming `hidden_states` as a target. Only - tracking names that may alias the input catches it. Result: no fusion.""" + tracking names that may alias the input catches it. Result: no fusion. + """ recompute = False @@ -551,7 +559,8 @@ class FunctionalInplaceArgAttention(FakeAttention): `F.relu(hidden_states, inplace=True)` writes through an argument, so the name appears only in a `Load` context inside a call that is not a method on - it -- invisible to a trailing-underscore method check. Result: no fusion.""" + it -- invisible to a trailing-underscore method check. Result: no fusion. + """ recompute = False @@ -581,7 +590,8 @@ class SubscriptMutatedArgAttention(FakeAttention): `hidden_states[..., 0] = 0` leaves the name in a `Load` context on the subscript's value, so it too evades a bare `Name`-store check. The rewrite - must refuse for the same reason. Result: no fusion.""" + must refuse for the same reason. Result: no fusion. + """ recompute = False @@ -611,7 +621,8 @@ class NonFoldableRefAttention(FakeAttention): Deleting `v_proj` would break `isinstance(self.v_proj, nn.Linear)`: it reads the projection's type, not just whether it exists, so the fuser cannot fold - it and must refuse rather than rewrite it.""" + it and must refuse rather than rewrite it. + """ def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -836,7 +847,8 @@ def test_glu_identifies_down_projection(): It is forced to `RowParallelLinear` in `update_attrs` so its sharded input matches the column-parallel merged gate/up; `None` when there is no such - projection to force (fusion of gate/up still applies).""" + projection to force (fusion of gate/up still applies). + """ with torch.device("meta"): assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj" assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj" @@ -997,7 +1009,8 @@ def test_bypass_existence_guard_folds_identity_none_check(expr, expected): The projection is a proven invariant (it exists as an `nn.Linear`), so an `is (not) None` guard has a constant truth value; folding it lets the call be rewritten away without changing the guard's outcome. `==`/`!=` are excluded - (see the refusal test) because a subclass may override `__eq__`/`__ne__`.""" + (see the refusal test) because a subclass may override `__eq__`/`__ne__`. + """ funcdef, ref = _guard_funcdef(expr) bypass_existence_guard(funcdef, ref, "v_proj") test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.IfExp)).test @@ -1016,7 +1029,8 @@ def test_bypass_existence_guard_folds_bare_truthiness(expr): An `nn.Linear` has no `__bool__`/`__len__`, so it is always truthy; in a pure boolean position (a test, or a `not` operand) the reference can be replaced by - `True`, letting the projection be deleted.""" + `True`, letting the projection be deleted. + """ funcdef, ref = _guard_funcdef(expr) bypass_existence_guard(funcdef, ref, "v_proj") # No `self.v_proj` reference survives; a literal `True` took its place. @@ -1031,7 +1045,8 @@ def test_bypass_existence_guard_folds_bare_if_statement(): The `len(node.ops) == 1` `Compare` scan skips it (it is not a comparison at all); `_in_boolean_context` then matches the `if`'s test and folds it, so the - bare-truthiness statement form is handled, not silently dropped.""" + bare-truthiness statement form is handled, not silently dropped. + """ funcdef = ast.parse( "def f(self, a, b):\n if self.v_proj:\n return a\n return b" ).body[0] @@ -1063,7 +1078,8 @@ def test_bypass_existence_guard_refuses_non_guard_reference(expr): Rewriting it is unsafe (its value depends on more than the projection's existence), so the fuser refuses rather than change semantics. `==`/`!=` are refused too: `is_linear` accepts `nn.Linear` subclasses that could override - `__eq__`/`__ne__`, so equality need not agree with `is (not) None`.""" + `__eq__`/`__ne__`, so equality need not agree with `is (not) None`. + """ funcdef, ref = _guard_funcdef(expr) with pytest.raises(ValueError, match="outside an existence guard"): bypass_existence_guard(funcdef, ref, "v_proj") @@ -1103,7 +1119,8 @@ def test_merged_column_fuser_folds_existence_guard(): Guard folding lives on `StackedFuser`, not just `QKVFuser`, because every stacked fuser deletes the projections it merges -- a surviving reference to - one would fail at runtime.""" + one would fail at runtime. + """ with torch.device("meta"): module = GuardedParallelLinears() fuser = MergedColumnParallelFuser.match(trace(module), module) @@ -1129,7 +1146,8 @@ def test_merged_column_fuser_rejects_input_mutation(cls): The mutation may be direct (`x.add_(1)`), through a view that shares the input's storage, or through a call that writes an argument in place (`F.relu(x, inplace=True)`, `torch.relu_(x)`). Each one silently changes the - fused result, so the rewrite must refuse rather than fuse.""" + fused result, so the rewrite must refuse rather than fuse. + """ with torch.device("meta"): module = cls() fuser = MergedColumnParallelFuser.match(trace(module), module) @@ -1187,7 +1205,8 @@ def test_merged_column_fuser_supports_any_number_of_linears( def test_merged_name_is_unique_per_linear_names(): """The computed name must not collide across different projections, so - stacking two unrelated fusions can never clobber packed_modules_mapping.""" + stacking two unrelated fusions can never clobber packed_modules_mapping. + """ ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c")) assert ab_c.merged_name != a_bc.merged_name @@ -1203,7 +1222,8 @@ def test_heterogeneous_instances_fall_back_to_merged_column_fuser( layer whose out_features aren't a multiple of its own head_dim (as a misconfigured or heterogeneous checkpoint might have) fails QKVFuser's validation and must still fuse via the generic fallback, while a - well-formed sibling of the same class keeps using QKVFuser.""" + well-formed sibling of the same class keeps using QKVFuser. + """ with torch.device("meta"): container = nn.Module() container.compatible = PlainQKV(head_dim=8) @@ -1224,7 +1244,8 @@ def test_detects_and_rewrites_packed_qkv(kv_heads): """A single projection split into q/k/v must be re-sharded, not merged. Only the split sizes change: `QKVParallelLinear` loads the packed - checkpoint weight as-is, and shards q by heads while replicating k/v.""" + checkpoint weight as-is, and shards q by heads while replicating k/v. + """ with torch.device("meta"): meta = PackedQKVAttention(kv_heads=kv_heads) fuser = get_fuser(meta, PackedQKVFuser) @@ -1255,7 +1276,8 @@ def test_detects_and_rewrites_packed_qkv(kv_heads): def test_per_head_split_is_not_packed_qkv(): """The split must consume the whole projection, else its sizes are head - widths and re-sharding by them would be wrong.""" + widths and re-sharding by them would be wrong. + """ with torch.device("meta"): assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None @@ -1375,7 +1397,8 @@ def test_act_and_mul_derived_from_module(default_vllm_config): def _wider_model_config(head_dim: int) -> SimpleNamespace: """A model whose global head size is twice `head_dim`, as a wider layer - elsewhere in a heterogeneous checkpoint would make it.""" + elsewhere in a heterogeneous checkpoint would make it. + """ return SimpleNamespace( model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim), quant_config=None, diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py index 9f51ab56621b..366123389fdf 100644 --- a/tests/models/transformers/fusers/test_mla.py +++ b/tests/models/transformers/fusers/test_mla.py @@ -45,7 +45,8 @@ def _attention(q_lora_rank: int | None): def _match(q_lora_rank: int | None) -> MLAFuser | None: """Match a meta `DeepseekV2Attention` directly (bypassing the per-class - `get_fuser` cache, so both q_lora variants of the same class are seen).""" + `get_fuser` cache, so both q_lora variants of the same class are seen). + """ attn = _attention(q_lora_rank) return MLAFuser.match(trace(attn), attn) @@ -81,7 +82,8 @@ def test_discovers_modules_with_q_lora(): def test_q_lora_stacks_qkv_a_proj(): """The MLA layer reads `q_a_proj` and `kv_a_proj_with_mqa` fused into one - down-projection, so both checkpoint weights must remap into it.""" + down-projection, so both checkpoint weights must remap into it. + """ fuser = _match(q_lora_rank=64) assert isinstance(fuser, MLAFuser) prefix = "model.layers.0.self_attn" @@ -102,7 +104,8 @@ def test_update_forward_rewrites_real_attention(q_lora_rank): The q-LoRA path hoists the merged down-projection to a top-level statement while its calls sit inside the `q_lora_rank` branches, so the region it spans is wide; a check that over-approximates aliasing would silently refuse - here and cost the fusion rather than fail loudly.""" + here and cost the fusion rather than fail loudly. + """ fuser = _match(q_lora_rank) assert isinstance(fuser, MLAFuser) fuser.update_forward(_attention(q_lora_rank)) @@ -126,7 +129,8 @@ def forward(self, x): class RenamedMLA(nn.Module): """An MLA-shaped attention whose children have non-standard names, proving - discovery is by structure and not attribute name.""" + discovery is by structure and not attribute name. + """ def __init__(self): super().__init__() @@ -153,7 +157,8 @@ def forward(self, hidden_states): def test_discovers_modules_under_arbitrary_names(): """Discovery is purely structural: `RenamedMLA` gives its children non-standard - names, and `match` still locates each projection by dataflow.""" + names, and `match` still locates each projection by dataflow. + """ with torch.device("meta"): module = RenamedMLA() fuser = MLAFuser.match(trace(module), module) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 7de4589eef91..a465116920d3 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -277,7 +277,8 @@ def test_moe_fuser_matches_scaled_router(): def test_moe_fuser_matches_grouped_router(): """Group masking between the score and the routing top-k still matches: the - scorer is anchored on the routing (last) top-k, not the group one.""" + scorer is anchored on the routing (last) top-k, not the group one. + """ with torch.device("meta"): block = MoEBlock(GroupedRouter) fuser = MoEBlockFuser.match(block, "experts") @@ -287,7 +288,8 @@ def test_moe_fuser_matches_grouped_router(): def test_moe_fuser_matches_correction_router(): """A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the - rebuilt gate carries it in fp32 for FusedMoE's biased routers.""" + rebuilt gate carries it in fp32 for FusedMoE's biased routers. + """ with torch.device("meta"): block = MoEBlock(CorrectionRouter) fuser = MoEBlockFuser.match(block, "experts") @@ -298,7 +300,8 @@ def test_moe_fuser_matches_correction_router(): def test_moe_fuser_reads_router_dtype_from_the_gate(): """A router that computes its logits in fp32 must keep routing in fp32 when rebuilt, even though no config field names the dtype. The cast back to the - activation dtype after the top-k must not be mistaken for the routing dtype.""" + activation dtype after the top-k must not be mistaken for the routing dtype. + """ with torch.device("meta"): fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts") default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts") @@ -317,7 +320,8 @@ def test_moe_fuser_detects_shared_experts(): def test_moe_fuser_skips_shared_detection_without_extra_children(): """With only experts and gate, shared-expert detection (and its block trace) - is skipped, so a gate-derived add is not misread as a shared expert.""" + is skipped, so a gate-derived add is not misread as a shared expert. + """ with torch.device("meta"): block = MoEBlockNoShared() fuser = MoEBlockFuser.match(block, "experts") @@ -362,7 +366,8 @@ def test_moe_fuser_declines_unsupported(block_cls): def test_moe_fuser_ignores_nested_returns(): """A tuple `return` inside a nested helper must not decline a block whose own - forward returns a single tensor.""" + forward returns a single tensor. + """ with torch.device("meta"): block = MoEBlockNestedTupleReturn() assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) @@ -370,7 +375,8 @@ def test_moe_fuser_ignores_nested_returns(): def test_moe_fuser_router_requires_connected_dataflow(): """A gate with linear + softmax + top-k present but not wired as a router - (top-k selects over the input, not the scored logits) is not detected.""" + (top-k selects over the input, not the scored logits) is not detected. + """ with torch.device("meta"): block = MoEBlock(DisconnectedRouter) assert MoEBlockFuser.match(block, "experts") is None diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index 153f7863bde5..dacf79f77af2 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -83,7 +83,8 @@ class ToleranceRMSNorm(NamedEpsRMSNorm): class LiteralEpsRMSNorm(RMSNorm): """eps is a literal in the source, and an unrelated attribute happens to - hold the same value: matching on the value alone would bind to it.""" + hold the same value: matching on the value alone would bind to it. + """ def __init__(self, hidden: int = 16, eps: float = 1e-5): super().__init__(hidden, eps) @@ -95,7 +96,8 @@ def forward(self, x): class AmbiguousEpsRMSNorm(nn.Module): """Two attributes hold the eps value, and the forward reads the second, so - their order cannot pick it.""" + their order cannot pick it. + """ def __init__(self, hidden: int = 16, eps: float = 1e-6): super().__init__() @@ -231,7 +233,8 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config): def test_fused_rms_norm_op_default_eps(default_vllm_config): """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path; - its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`.""" + its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`. + """ from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm with torch.device("meta"): @@ -247,7 +250,8 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): def test_eps_is_derived_per_instance(default_vllm_config): """Two instances of the same norm class with different eps must fuse to their - own eps: the type-cached fuser holds only structure, not this value.""" + own eps: the type-cached fuser holds only structure, not this value. + """ with torch.device("meta"): for eps in (1e-5, 1e-6): module = RMSNorm(16, eps=eps) @@ -259,7 +263,8 @@ def test_eps_is_derived_per_instance(default_vllm_config): @pytest.mark.parametrize("cls", [NamedEpsRMSNorm, ToleranceRMSNorm]) def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): """The eps attribute is identified by holding the traced value, so a norm - stays per-instance correct whatever it names it.""" + stays per-instance correct whatever it names it. + """ with torch.device("meta"): for eps in (1e-5, 1e-6): module = cls(16, eps=eps) @@ -271,7 +276,8 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplog): """A literal eps is recognised as coming from no attribute, even when one - holds the same value, and is taken from the traced source instead.""" + holds the same value, and is taken from the traced source instead. + """ logger = "vllm.model_executor.models.transformers.fusers.rms_norm" with caplog.at_level("DEBUG", logger=logger), torch.device("meta"): module = LiteralEpsRMSNorm() @@ -284,7 +290,8 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): """When several attributes hold the eps value, the one the forward actually - reads is identified, and they are all left as they were found.""" + reads is identified, and they are all left as they were found. + """ with torch.device("meta"): module = AmbiguousEpsRMSNorm(16, eps=1e-6) before = dict(vars(module)) @@ -301,7 +308,8 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): def test_fused_norm_is_gather_capable(default_vllm_config): """Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded projection (OLMoE-style) self-corrects at runtime with no QKV-specific - plumbing. A full-width input skips the gather and equals a plain norm.""" + plumbing. A full-width input skips the gather and equals a plain norm. + """ from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.model_executor.models.transformers.fusers import rms_norm @@ -323,7 +331,8 @@ def test_fused_norm_is_gather_capable(default_vllm_config): def test_gathered_norm_rejects_uneven_sharding(default_vllm_config): """A sharded input (narrower than the full-width weight) that does not tile - the weight evenly across ranks is rejected before any collective.""" + the weight evenly across ranks is rejected before any collective. + """ from vllm.model_executor.models.transformers.fusers import rms_norm norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6) diff --git a/tests/models/transformers/test_layer_registry.py b/tests/models/transformers/test_layer_registry.py index 9e20e2ba54ad..cc809837ae7e 100644 --- a/tests/models/transformers/test_layer_registry.py +++ b/tests/models/transformers/test_layer_registry.py @@ -28,7 +28,8 @@ def fake_hw_layernorm(monkeypatch): A `SimpleNamespace` stands in for the module: `importlib.import_module` returns it from `sys.modules` and `getattr` resolves `RMSNorm`, while its attributes are set at construction (no `ModuleType` attribute-set that mypy - rejects, no constant `setattr` that ruff rejects).""" + rejects, no constant `setattr` that ruff rejects). + """ module = types.SimpleNamespace(RMSNorm=type("HwRMSNorm", (), {})) monkeypatch.setitem(sys.modules, HW_MODULE, module) return module @@ -81,7 +82,8 @@ def test_act_and_mul_falls_back_for_unknown_activation( @pytest.fixture(scope="module") def tiny_llama_path(tmp_path_factory): """A randomly-initialized microscopic Llama saved to disk (with an ungated - tokenizer) so vLLM can load it like any local checkpoint.""" + tokenizer) so vLLM can load it like any local checkpoint. + """ from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") diff --git a/tests/models/utils.py b/tests/models/utils.py index aff9fb0f5a66..b52855bcdd30 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -29,8 +29,7 @@ def check_outputs_equal( name_0: str, name_1: str, ): - """ - Compare the two sequences generated by different models, + """Compare the two sequences generated by different models, which should be equal. """ assert len(outputs_0_lst) == len(outputs_1_lst) @@ -131,6 +130,7 @@ def check_logprobs_close( warn_on_mismatch: Issue a warning if there is token-wise or text-wise mismatch between the two sequences always_check_logprobs: If true, check logprobs even when tokens match + """ assert len(outputs_0_lst) == len(outputs_1_lst) @@ -290,6 +290,7 @@ def build_model_context( Returns: InputProcessingContext for the model being considered. + """ model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) model_info.check_available_online(on_fail="skip") @@ -447,8 +448,7 @@ def dummy_hf_overrides( exist_overrides: dict[str, Any] | None = None, use_original_num_layers: bool = False, ) -> PretrainedConfig: - """ - Dummy HF overrides function used to create dummy model + """Dummy HF overrides function used to create dummy model with only minimum nums of layer. """ # Copy because this helper is called more than once diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index 3473a911eab5..55736fbe2ac8 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -54,7 +54,8 @@ def test_audio_media_io_load_base64(dummy_audio_bytes): def test_audio_media_io_load_base64_rejects_malformed(dummy_audio_bytes): """Malformed base64 must surface as a ValueError so the server answers 400. Without strict decoding the bad characters are dropped, the garbage reaches - libsndfile, and the client gets a 500 instead.""" + libsndfile, and the client gets a 500 instead. + """ encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8") malformed = encoded[:8] + "!!!@@@###" + encoded[8:] @@ -133,7 +134,8 @@ def test_small_file_passes_memory_guard(): def test_memory_guard_rejects_large_allocation(): """A file whose frames*channels*4 exceeds the byte limit must be - rejected before allocating the buffer.""" + rejected before allocating the buffer. + """ # 100_000 frames * 8 channels * 4 bytes = 3.2 MB payload = _make_flac_bytes(frames=100_000, channels=8, samplerate=48000) # Set limit to 1 MiB — should reject @@ -148,7 +150,8 @@ def test_memory_guard_rejects_large_allocation(): def test_forged_samplerate_rejected_by_memory_guard(): """The PoC scenario: high sample rate fools the duration guard but - the memory guard catches the large frame*channel allocation.""" + the memory guard catches the large frame*channel allocation. + """ # Forged high sample rate: 655350 Hz, 8 channels, 1M frames # Duration guard sees: 1_000_000 / 655_350 = 1.5s → passes # Memory: 1_000_000 * 8 * 4 = 32 MB @@ -209,7 +212,8 @@ def test_load_audio_unknown_backend_rejected(dummy_audio_bytes): def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes): """`auto` must fall back to the soundfile → PyAV chain when torchcodec - is not importable.""" + is not importable. + """ ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) with patch.object(audio_module, "load_audio_torchcodec", side_effect=ImportError): audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="auto") @@ -218,8 +222,9 @@ def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes): def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes): - """torchcodec installed but system ffmpeg missing (`AudioDecoder is None`) - must surface as ImportError so `auto` falls back to soundfile → PyAV.""" + """Torchcodec installed but system ffmpeg missing (`AudioDecoder is None`) + must surface as ImportError so `auto` falls back to soundfile → PyAV. + """ ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) with patch.object(audio_module, "AudioDecoder", None): with pytest.raises(ImportError, match="torchcodec audio backend"): @@ -232,9 +237,10 @@ def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes): def test_load_audio_auto_falls_back_when_libtorchcodec_unloadable( dummy_audio_bytes, ): - """torchcodec loads its ffmpeg-backed core lazily at decoder construction; + """Torchcodec loads its ffmpeg-backed core lazily at decoder construction; a "Could not load libtorchcodec" RuntimeError there (no system ffmpeg) - must also surface as ImportError so `auto` falls back.""" + must also surface as ImportError so `auto` falls back. + """ class _UnloadableDecoder: def __init__(self, *args, **kwargs): diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index b222d57d3c24..64145e0fea4a 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -228,7 +228,8 @@ async def test_fetch_image_local_files_with_space_in_name(image_url: str): @pytest.mark.asyncio async def test_fetch_image_data_url_with_params(): """RFC 2397 allows parameters between the mediatype and the base64 - marker; they must not be rejected or leak into the media type.""" + marker; they must not be rejected or leak into the media type. + """ connector = MediaConnector() image = Image.new("RGB", (4, 4), color=(255, 0, 0)) diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index 5bbf54b7872e..6b7e4ad407b7 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -105,7 +105,6 @@ def test_image_media_io_no_mode_conversion(tmp_path): def test_image_media_io_rgba_background_color_validation(): """Test that invalid rgba_background_color values are properly rejected.""" - # Test invalid types with pytest.raises( ValueError, match="rgba_background_color must be a list or tuple" diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index 79070ac50d5d..2205b7c7ff3d 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -77,8 +77,7 @@ def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch): @pytest.mark.parametrize("is_color", [True, False]) @pytest.mark.parametrize("fourcc, ext", [("mp4v", "mp4"), ("XVID", "avi")]) def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext: str): - """ - Test all functions that use OpenCV for video I/O return RGB format. + """Test all functions that use OpenCV for video I/O return RGB format. Both RGB and grayscale videos are tested. """ image_path = get_vllm_public_assets( @@ -168,8 +167,7 @@ def load_bytes( def test_video_media_io_backend_kwarg_override(monkeypatch: pytest.MonkeyPatch): - """ - Test that video_backend kwarg can override the VLLM_VIDEO_LOADER_BACKEND + """Test that video_backend kwarg can override the VLLM_VIDEO_LOADER_BACKEND environment variable. This allows users to dynamically select a different video backend @@ -201,8 +199,7 @@ def test_video_media_io_backend_kwarg_override(monkeypatch: pytest.MonkeyPatch): def test_video_media_io_backend_kwarg_not_passed_to_loader( monkeypatch: pytest.MonkeyPatch, ): - """ - Test that video_backend kwarg is consumed by VideoMediaIO and NOT passed + """Test that video_backend kwarg is consumed by VideoMediaIO and NOT passed through to the underlying video loader's load_bytes method. This ensures the kwarg is properly popped from kwargs before forwarding. @@ -245,8 +242,7 @@ def load_bytes( def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch): - """ - Test that when video_backend kwarg is None or not provided, + """Test that when video_backend kwarg is None or not provided, VideoMediaIO falls back to VLLM_VIDEO_LOADER_BACKEND env var. """ with monkeypatch.context() as m: @@ -285,7 +281,6 @@ def test_load_base64_jpeg_returns_metadata(): metadata, which broke downstream consumers that rely on fields like total_num_frames and fps. See PR #37301. """ - num_test_frames = 3 b64_frames = _make_jpeg_b64_frames(num_test_frames) @@ -403,7 +398,8 @@ def raise_unrelated_error(cls, file_path, nvc): class TestMergeKwargsGpuBackendPolicy: """Verify that merge_kwargs blocks request-level GPU backend selection - when the static (engine-level) config did not configure that backend.""" + when the static (engine-level) config did not configure that backend. + """ def test_pynvvideocodec_requires_gpu(self): assert VIDEO_LOADER_REGISTRY.backend_requires_gpu(PYNVVIDEOCODEC_VIDEO_BACKEND) @@ -496,7 +492,8 @@ def test_other_kwargs_preserved_when_gpu_backend_stripped(self): def test_static_pynv_with_different_runtime_gpu_backend(self): """If static sets pynv via video_backend but runtime tries to set it - via the codec-level 'backend' key (without a static match), strip it.""" + via the codec-level 'backend' key (without a static match), strip it. + """ result = VideoMediaIO.merge_kwargs( default_kwargs={"video_backend": "pynvvideocodec"}, runtime_kwargs={"backend": "pynvvideocodec"}, @@ -535,7 +532,8 @@ def test_unknown_backend_not_treated_as_gpu(self): def test_strips_request_level_device(self): """The decode device is a startup-only knob: a request must not move decoding onto the GPU when the startup config did not opt in, nor off - it when it did.""" + it when it did. + """ result = VideoMediaIO.merge_kwargs( default_kwargs={"backend": "torchcodec"}, runtime_kwargs={"device": "cuda"}, @@ -560,7 +558,8 @@ def test_strips_request_level_device(self): def test_switching_backend_drops_stale_codec_options(default_kwargs): """Codec-specific options from the static config must not leak into a different codec backend selected per-request, where they would fail the - new backend's option validation.""" + new backend's option validation. + """ result = VideoMediaIO.merge_kwargs( default_kwargs={**default_kwargs, "num_frames": 8}, runtime_kwargs={"backend": "opencv"}, @@ -582,7 +581,8 @@ def test_pynvvc_frames_normalized_to_nhwc(layout: str): """PyNvVideoCodec frame batches are normalized to NHWC regardless of the per-frame layout the decoder emits (it has varied across versions), so the HF video processors (which materialize a PIL image per frame) receive the - same NHWC shape as every other video backend.""" + same NHWC shape as every other video backend. + """ torch = pytest.importorskip("torch") n, h, w, c = 4, 5, 6, 3 diff --git a/tests/multimodal/test_audio.py b/tests/multimodal/test_audio.py index 21b39710d4f4..c816e979abce 100644 --- a/tests/multimodal/test_audio.py +++ b/tests/multimodal/test_audio.py @@ -695,7 +695,6 @@ class TestAudioChunking: def test_split_audio_short_clip(self): """Audio shorter than max_clip_duration_s should not be split.""" - # 10 seconds of audio at 16kHz audio = np.linspace(-1.0, 1.0, 160000, dtype=np.float32) @@ -712,7 +711,6 @@ def test_split_audio_short_clip(self): def test_split_audio_exact_length(self): """Audio exactly at max_clip_duration_s should not be split.""" - # Exactly 30 seconds at 16kHz audio = np.linspace(-1.0, 1.0, 480000, dtype=np.float32) @@ -729,7 +727,6 @@ def test_split_audio_exact_length(self): def test_split_audio_long_clip(self): """Long audio should be split into multiple chunks.""" - # 65 seconds of audio at 16kHz audio = np.linspace(-1.0, 1.0, 1040000, dtype=np.float32) @@ -749,7 +746,6 @@ def test_split_audio_long_clip(self): def test_split_audio_chunks_have_correct_length(self): """Each chunk (except last) should be approximately max_clip_duration_s.""" - # 65 seconds of audio at 16kHz audio = np.linspace(-1.0, 1.0, 1040000, dtype=np.float32) @@ -841,7 +837,6 @@ def test_split_audio_rejects_multi_channel(self): def test_split_audio_preserves_boundaries(self): """Verify first and last samples are preserved when chunking.""" - audio = np.arange(1120000, dtype=np.float32) # 70s at 16kHz chunks = split_audio( @@ -893,7 +888,6 @@ def test_split_audio_nan_input_terminates(self): def test_split_audio_with_different_sample_rates(self): """Test chunking works with different sample rates.""" - # 40 seconds at 8kHz audio_8k = np.linspace(-1.0, 1.0, 320000, dtype=np.float32) diff --git a/tests/multimodal/test_embedding_shape_validation.py b/tests/multimodal/test_embedding_shape_validation.py index fcac5be5b5a7..3d7151271065 100644 --- a/tests/multimodal/test_embedding_shape_validation.py +++ b/tests/multimodal/test_embedding_shape_validation.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Embedding shape validation in multimodal APIs. +"""Embedding shape validation in multimodal APIs. Tests verify that embeddings with correct ndim but incorrect hidden_size are rejected before they can cause crashes during model inference. @@ -178,8 +177,7 @@ class TestShapeValidationIntegration: """Integration tests verifying attack scenarios are blocked.""" def test_attack_scenario_multimodal_image(self): - """ - Simulate attack through Chat API with image embeddings. + """Simulate attack through Chat API with image embeddings. Verifies validation occurs in multimodal parser path. """ @@ -193,8 +191,7 @@ def test_attack_scenario_multimodal_image(self): parser.parse_mm_data({"image": attack_tensor}) def test_attack_scenario_multimodal_audio(self): - """ - Simulate attack through Chat API with audio embeddings. + """Simulate attack through Chat API with audio embeddings. Verifies validation occurs in multimodal parser path. """ @@ -208,8 +205,7 @@ def test_attack_scenario_multimodal_audio(self): parser.parse_mm_data({"audio": attack_tensor}) def test_attack_scenario_multimodal_video(self): - """ - Simulate attack through Chat API with video embeddings. + """Simulate attack through Chat API with video embeddings. Verifies validation occurs in multimodal parser path. """ diff --git a/tests/multimodal/test_embedding_shape_validation_unit.py b/tests/multimodal/test_embedding_shape_validation_unit.py index 7966aad4e988..c8f785046eb5 100644 --- a/tests/multimodal/test_embedding_shape_validation_unit.py +++ b/tests/multimodal/test_embedding_shape_validation_unit.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for embedding shape validation. +"""Unit tests for embedding shape validation. Simple, fast unit tests that can run without server fixtures. Run with: pytest tests/multimodal/test_embedding_shape_validation_unit.py -v @@ -167,16 +166,14 @@ def test_hidden_size_validation_wrong_size_rejected(self): class TestShapeValidationDoSPrevention: - """ - Tests for DoS prevention through shape validation. + """Tests for DoS prevention through shape validation. Verifies that embeddings with incorrect shapes are rejected early, preventing crashes during model inference. """ def test_prevent_crash_from_wrong_shape_image_embeds(self): - """ - Prevent crash scenario: wrong hidden size in image embeddings. + """Prevent crash scenario: wrong hidden size in image embeddings. Without validation, this would pass initial checks but crash later during model forward pass when dimensions don't match. @@ -198,9 +195,7 @@ def test_prevent_crash_from_wrong_shape_image_embeds(self): assert str(wrong_hidden_size) in error_msg # Received def test_prevent_crash_from_wrong_shape_audio_embeds(self): - """ - Prevent crash scenario: wrong hidden size in audio embeddings. - """ + """Prevent crash scenario: wrong hidden size in audio embeddings.""" expected_hidden_size = 768 wrong_hidden_size = 4096 diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 1b738eeeb89f..34ecc1c1f657 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -141,7 +141,8 @@ def item_for_hash(num_frames: int): def test_hash_video_tensor_frames(): """Videos holding tensor frames (e.g. NVDEC-decoded) hash like - array-framed ones, from the original bytes without a D2H copy.""" + array-framed ones, from the original bytes without a D2H copy. + """ source = b"x" * 100 def item_for_hash(frames): diff --git a/tests/multimodal/test_image.py b/tests/multimodal/test_image.py index ea7f18611c3f..ad1ce3162866 100644 --- a/tests/multimodal/test_image.py +++ b/tests/multimodal/test_image.py @@ -80,7 +80,8 @@ def test_l_mode_no_trns_to_rgb(): def test_exif_transpose_normalizes_orientation(): """Image with EXIF orientation 3 (180-degree rotation) should be - normalized so pixel data matches visual display.""" + normalized so pixel data matches visual display. + """ # Built synthetically because the existing assets are PNGs which # don't carry EXIF orientation metadata; we need a JPEG with an # injected EXIF orientation=3 tag. diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index 2649f80777c1..9c4dbed0f02c 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -59,7 +59,8 @@ def test_frame_size_hwc_chw(frame): def test_video_with_metadata_tensor_passthrough(): """Tensor frames pass through unchanged regardless of device: HF video processors accept tensors, and device-resident frames (e.g. NVDEC-decoded) - must not be copied back to host.""" + must not be copied back to host. + """ frames = torch.zeros((4, H, W, 3), dtype=torch.uint8) video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) @@ -70,7 +71,8 @@ def test_video_with_metadata_tensor_passthrough(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_video_with_metadata_keeps_device_tensor(): """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors, - so a device-side HF processor can consume them without a D2H copy.""" + so a device-side HF processor can consume them without a D2H copy. + """ frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda") video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index 4a41013b21a5..7d6a7fd7e168 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -968,8 +968,7 @@ def test_hf_processor_call_kwargs( def test_apply_matches_no_match_exits_quickly(): - """ - Test that _apply_matches exits quickly when no matches are found. + """Test that _apply_matches exits quickly when no matches are found. Previously, _apply_matches had O(n²) behavior when no match was found because it would increment start_idx by 1 each iteration while @@ -1027,8 +1026,7 @@ def test_iter_token_matches_rejects_negative_start_idx(): def test_find_mm_placeholders_avoids_quadratic_false_prefixes(): - """ - Test that placeholder scanning stays linear under adversarial candidates. + """Test that placeholder scanning stays linear under adversarial candidates. The fast-forward scan must not rescan the prompt tail per position when one candidate's first token never occurs (forcing a full search) while @@ -1061,8 +1059,7 @@ def test_find_mm_placeholders_avoids_quadratic_false_prefixes(): ], ) def test_find_mm_placeholders_stops_at_missing_item(prompt): - """ - Test that the scan returns no placeholders once it fails to find + """Test that the scan returns no placeholders once it fails to find an item's placeholder, leaving later items unresolved. """ result = find_mm_placeholders( @@ -1079,8 +1076,7 @@ def test_find_mm_placeholders_stops_at_missing_item(prompt): class _FakeTokenizer: - """ - Character-level tokenizer where "foo" merges into one token differently + """Character-level tokenizer where "foo" merges into one token differently depending on whether it is followed by "d", like BPE merging "foo" in "food" across the search-text boundary. """ @@ -1145,8 +1141,7 @@ def _text_fallback_processor() -> BaseMultiModalProcessor: def test_apply_prompt_updates_falls_back_to_text_matching(): - """ - Test that the fallback in `_apply_prompt_updates` finds targets that + """Test that the fallback in `_apply_prompt_updates` finds targets that tokenize differently inside the prompt ("foo" in "food"). """ processor = _text_fallback_processor() @@ -1166,8 +1161,7 @@ def test_apply_prompt_updates_falls_back_to_text_matching(): def test_apply_prompt_updates_falls_back_with_prefix_target(): - """ - Test that `PromptIndexTargets.prefix` targets are resolved against the + """Test that `PromptIndexTargets.prefix` targets are resolved against the decoded text in the fallback path of `_apply_prompt_updates`. """ processor = _text_fallback_processor() @@ -1192,8 +1186,7 @@ def test_apply_prompt_updates_falls_back_with_prefix_target(): def test_apply_prompt_updates_falls_back_with_index_targets(): - """ - Test that the text resolvers of `PromptIndexTargets.start`/`end` + """Test that the text resolvers of `PromptIndexTargets.start`/`end` match against the decoded text when another item forces the fallback in `_apply_prompt_updates`. """ diff --git a/tests/multimodal/test_registry.py b/tests/multimodal/test_registry.py index 99f0689c7da3..4a08d4fdf66b 100644 --- a/tests/multimodal/test_registry.py +++ b/tests/multimodal/test_registry.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for MultiModalRegistry.supports_multimodal_inputs and +"""Unit tests for MultiModalRegistry.supports_multimodal_inputs and Qwen2.5-VL visual component loading behavior. """ @@ -58,7 +57,8 @@ def test_dummy_inputs_scheduler_budget( @pytest.mark.core_model def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected): """Test supports_multimodal_inputs returns correct boolean for various - configs.""" + configs. + """ ctx = build_model_context( model_id, limit_mm_per_prompt=limit_mm_per_prompt, diff --git a/tests/multimodal/test_sparse_tensor_validation_unit.py b/tests/multimodal/test_sparse_tensor_validation_unit.py index 2eec8ea8283a..4e35e4abc81c 100644 --- a/tests/multimodal/test_sparse_tensor_validation_unit.py +++ b/tests/multimodal/test_sparse_tensor_validation_unit.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for sparse tensor validation. +"""Unit tests for sparse tensor validation. Simple, fast unit tests that can run without server fixtures. Run with: pytest tests/multimodal/test_sparse_tensor_validation_unit.py -v @@ -56,8 +55,7 @@ def test_negative_indices_rejected(self): tensor.to_dense() def test_without_context_manager_allows_invalid(self): - """ - WITHOUT validation, invalid tensors may not immediately error. + """WITHOUT validation, invalid tensors may not immediately error. This demonstrates the vulnerability: PyTorch 2.8.0+ doesn't validate by default, which can lead to memory corruption. diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py index 97f46de13c18..994e1afc8bc4 100644 --- a/tests/multimodal/test_utils.py +++ b/tests/multimodal/test_utils.py @@ -231,7 +231,8 @@ def test_group_and_batch_mm_items_splits_shared_data_by_dtype(): def test_encode_image_url_uncommon_format_has_valid_mimetype(): """A format missing from mimetypes.types_map must still produce a - well-formed type/subtype media type, not a bare 'image'.""" + well-formed type/subtype media type, not a bare 'image'. + """ # ".tga" is not registered in mimetypes.types_map, so it exercises the # fallback path; Pillow can still encode it. url = encode_image_url(Image.new("RGB", (2, 2)), format="TGA") diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 0bc3b21b42b4..1bc65e86a97f 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -842,8 +842,7 @@ def test_video_processor_from_model_repo( def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch): - """ - Regression test for handling videos with broken frames. + """Regression test for handling videos with broken frames. This test uses a pre-corrupted video file (assets/corrupted.mp4) that contains broken frames to verify the video loader handles them gracefully without crashing and returns accurate metadata. @@ -882,8 +881,7 @@ def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch): def test_video_backend_handles_edit_list_trimmed_video( monkeypatch: pytest.MonkeyPatch, ): - """ - An mp4 edit list (e.g. from a lossless ``ffmpeg -ss ... -c copy`` cut) + """An mp4 edit list (e.g. from a lossless ``ffmpeg -ss ... -c copy`` cut) hides the decode lead-in: the header still counts every physical sample while sequential decode only yields the visible frames. Sampling over the header count used to collapse such videos to the few indices below the @@ -923,8 +921,7 @@ def test_video_backend_handles_edit_list_trimmed_video( def test_video_recovery_simulated_failures(monkeypatch: pytest.MonkeyPatch): - """ - Test that frame recovery correctly uses the next valid frame when + """Test that frame recovery correctly uses the next valid frame when target frames fail to load. Uses corrupted.mp4 and mocks VideoCapture.grab() to fail on specific @@ -1019,8 +1016,7 @@ def release(self): def test_video_recovery_with_corrupted_file(monkeypatch: pytest.MonkeyPatch): - """ - Test frame recovery with an actual corrupted video file using sparse sampling. + """Test frame recovery with an actual corrupted video file using sparse sampling. This test uses corrupted.mp4 which has genuine H.264 codec errors on frame 17. With num_frames=8, the target frames are [0, 3, 7, 10, 14, 17, 21, 25]. @@ -1088,8 +1084,7 @@ def test_video_recovery_with_corrupted_file(monkeypatch: pytest.MonkeyPatch): def test_video_recovery_dynamic_backend(monkeypatch: pytest.MonkeyPatch): - """ - Test that frame_recovery works with the dynamic video backend. + """Test that frame_recovery works with the dynamic video backend. The dynamic backend samples frames based on fps/duration rather than loading all frames. This test verifies recovery works in that context. @@ -1219,7 +1214,8 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path): """With device="cuda", torchcodec decodes via NVDEC and keeps the frames - on the GPU instead of returning a host-side numpy array.""" + on the GPU instead of returning a host-side numpy array. + """ pytest.importorskip("torchcodec") with open(dummy_video_path, "rb") as f: @@ -1512,7 +1508,8 @@ def test_glm46v_duration_estimation_from_fps(): class TestGLMGASamplingCaps: """Regression tests for GHSA-58v5-2m8f-94pr: request-controlled fps and max_frames must not create intermediate allocations larger than - the actual frame count.""" + the actual frame count. + """ @staticmethod def _source( @@ -1566,7 +1563,8 @@ def test_glm5next_backend_selected_for_processor(): """Glm5NextVideoProcessor maps to the glm5next loader so only the sampled frames are decoded instead of the whole container. Both the borrowed-config spelling and the dedicated Glm5next class name (landing - with the new checkpoint) must resolve.""" + with the new checkpoint) must resolve. + """ for name in ("Glm5NextVideoProcessor", "Glm5nextVideoProcessor"): assert VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(name) == "glm5next" @@ -1588,7 +1586,8 @@ def test_glm5next_backend_indices_match_sampler( total_frames, original_fps, duration, fps, max_frames ): """The loader must select exactly the frames the processor's sampler - would, with target.fps mapping onto the raw-fps override.""" + would, with target.fps mapping onto the raw-fps override. + """ from vllm.transformers_utils.processors.glm5next import ( glm_sample_frame_indices, ) @@ -1616,7 +1615,8 @@ def test_glm5next_backend_indices_match_sampler( def test_glm5next_backend_metadata_contract(): """create_hf_metadata reports the subset so the processor skips - re-sampling (do_sample_frames=False) and keeps the original totals.""" + re-sampling (do_sample_frames=False) and keeps the original totals. + """ source = VideoSourceMetadata(total_frames_num=900, original_fps=30.0, duration=30.0) target = VideoTargetMetadata(num_frames=-1, fps=-1, max_duration=-1) indices = Glm5NextVideoBackend.compute_frames_index_to_sample(source, target) @@ -1645,7 +1645,8 @@ def test_glm5next_backend_metadata_contract(): def _write_gray_video(tmp_path, total_frames, fps, size=(32, 32)): """Synthetic clip whose frame i is flat gray level i (near-lossless under - mp4v), so a decoded frame's level maps back to its source index.""" + mp4v), so a decoded frame's level maps back to its source index. + """ cv2 = pytest.importorskip("cv2") path = tmp_path / f"gray_{total_frames}_{fps}.mp4" @@ -1686,7 +1687,8 @@ def read(self): def test_glm5next_backend_codec_parity(tmp_path, backend): """Every codec samples the same GLM indices and decodes the same frames; the OpenCV seek reader and torchcodec batched index-exact decode - must agree.""" + must agree. + """ if backend == "torchcodec": pytest.importorskip("torchcodec") @@ -1716,7 +1718,8 @@ def test_glm5next_backend_codec_parity(tmp_path, backend): def test_glm5next_backend_decodes_only_sampled_frames(tmp_path): """End to end over a synthetic clip: load_bytes returns exactly the - sampler's frame count, with the right frame content at each index.""" + sampler's frame count, with the right frame content at each index. + """ pytest.importorskip("cv2") total_frames, fps = 60, 10 @@ -1742,7 +1745,8 @@ def test_glm5next_backend_decodes_only_sampled_frames(tmp_path): def test_glm5next_read_frames_seeks_past_large_gaps(tmp_path): """Sparse targets must not walk the container: the stock reader grabs - every frame up to the last index; the GLM reader seeks instead.""" + every frame up to the last index; the GLM reader seeks instead. + """ cv2 = pytest.importorskip("cv2") total_frames, fps = 200, 10 @@ -1769,7 +1773,8 @@ def test_glm5next_read_frames_seeks_past_large_gaps(tmp_path): def test_glm5next_read_frames_dense_walk_matches_stock(tmp_path): """Dense targets keep the sequential walk (seeking would be slower) and - return the same frames as the stock reader.""" + return the same frames as the stock reader. + """ cv2 = pytest.importorskip("cv2") total_frames, fps = 120, 10 diff --git a/tests/parser/cohere/test_reasoning.py b/tests/parser/cohere/test_reasoning.py index 8c3311e8388c..374c65aaaef1 100644 --- a/tests/parser/cohere/test_reasoning.py +++ b/tests/parser/cohere/test_reasoning.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Reasoning extraction, reasoning-end gating and parser selection for the -unified Cohere Command parser.""" +unified Cohere Command parser. +""" from __future__ import annotations diff --git a/tests/parser/cohere/utils.py b/tests/parser/cohere/utils.py index f26df5296c6d..0696b9c24a68 100644 --- a/tests/parser/cohere/utils.py +++ b/tests/parser/cohere/utils.py @@ -76,7 +76,8 @@ def token_deltas( ) -> list[tuple[str, list[int]]]: """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing - U+FFFD) into the next step as real streaming does.""" + U+FFFD) into the next step as real streaming does. + """ ids = tokenizer.encode(text) deltas: list[tuple[str, list[int]]] = [] prev = "" @@ -100,7 +101,8 @@ def make_parser( model_config=None, ) -> Parser: """Resolve the unified Cohere parser through ``ParserManager`` by registry - name (``cohere_command3`` / ``cohere_command4``).""" + name (``cohere_command3`` / ``cohere_command4``). + """ cls = ParserManager.get_parser(name, name, enable_auto_tools=True) assert cls is not None return cls( @@ -117,7 +119,8 @@ def drive_parser( deltas: list[tuple[str, list[int]]], ) -> list[DeltaMessage]: """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and - return the non-``None`` messages in order.""" + return the non-``None`` messages in order. + """ out: list[DeltaMessage] = [] for i, (text, ids) in enumerate(deltas): delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1) diff --git a/tests/parser/engine/conftest.py b/tests/parser/engine/conftest.py index 2522d146c895..1c482434eb5f 100644 --- a/tests/parser/engine/conftest.py +++ b/tests/parser/engine/conftest.py @@ -30,6 +30,7 @@ def make_mock_tokenizer( The returned mock supports get_vocab(), encode(), and decode(). decode() maps known token IDs back to their text and falls back to chr(id) for ASCII IDs or ```` for others. + """ id_to_text = {v: k for k, v in vocab.items()} tokenizer = MagicMock() diff --git a/tests/parser/engine/replay_harness.py b/tests/parser/engine/replay_harness.py index 824a1ed86d88..5eec8d213a02 100644 --- a/tests/parser/engine/replay_harness.py +++ b/tests/parser/engine/replay_harness.py @@ -183,6 +183,7 @@ def replay_streaming( Returns: List of ``DeltaMessage`` results from each ``parse_delta()`` call. + """ if chunk_size is None: chunk_size = len(tokens) diff --git a/tests/parser/engine/test_deepseek_v32.py b/tests/parser/engine/test_deepseek_v32.py index 8db4e0cc76b4..aacde642fc89 100644 --- a/tests/parser/engine/test_deepseek_v32.py +++ b/tests/parser/engine/test_deepseek_v32.py @@ -289,7 +289,8 @@ def test_missing_invoke_end(self, mock_tokenizer, mock_request): class TestMissingFunctionCallsWrapper: """An invoke without the ``<｜DSML｜function_calls>`` wrapper is still a - tool call (V3.2 counterpart of #48931).""" + tool call (V3.2 counterpart of #48931). + """ def test_non_streaming(self, mock_tokenizer, mock_request): text = _invoke("fn", _param("k", "true", "v")) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 7cbf9791a7b0..0c8a375ac3b9 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -334,7 +334,8 @@ def test_trailing_content_dropped_consistently( ): """Text after a tool block is dropped whether or not the closing wrapper is present, so a missing ```` does not - change what the client sees.""" + change what the client sees. + """ text = "pre\n" + self._ORPHAN if trailing_end: text += DSML_TOOL_END @@ -915,7 +916,8 @@ def test_unwrap_single_tool_still_works(self, mock_tokenizer, mock_request): class TestStreamingWrapperConsistency: """Streamed arg deltas must stay consistent with final extraction - when wrapper params like 'arguments' are unwrapped.""" + when wrapper params like 'arguments' are unwrapped. + """ def test_streaming_wrapper_unwrap_consistency(self, mock_tokenizer, mock_request): tool = _make_tool("get_weather", {"location": {"type": "string"}}) @@ -1196,7 +1198,8 @@ def test_eos_drop_token_does_not_swallow_tool_calls(self): ) def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size): """EOS must not leak into reasoning_content when the model never - emits (generation ends while still in REASONING state).""" + emits (generation ends while still in REASONING state). + """ eos_text = "<｜end▁of▁sentence｜>" eos_id = 128801 vocab = { @@ -1285,7 +1288,8 @@ def test_corrupted_wrapper_ends_reasoning( self, mock_tokenizer, mock_request, wrapper ): """Like the real wrapper, a corrupted one inside ```` closes - the reasoning block and starts the tool call.""" + the reasoning block and starts the tool call. + """ text = f"thinking{wrapper}\n{self._INNER}" parser = self._parser(mock_tokenizer, thinking=True) reasoning, content = parser.extract_reasoning(text, mock_request) @@ -1329,7 +1333,8 @@ def test_prose_mention_matches_canonical_wrapper( prose mention of the real wrapper: the opener starts a tool block and the rest of the message is dropped. That is pre-existing behavior of ``<｜DSML｜tool_calls>``; the variants must not differ from it either - way.""" + way. + """ template = "The opener looks like {} and then params follow. Done." text = template.format(wrapper) reference = template.format(DSML_TOOL_START) @@ -1353,7 +1358,8 @@ def test_prose_mention_matches_canonical_wrapper( @pytest.mark.parametrize("chunk_size", [1, 3, None], ids=lambda c: f"chunk={c}") def test_delegating_parser_corrupted_wrapper(self, wrapper, chunk_size): """The serving-layer shape: reasoning and tool adapters on separate - engines, with the corrupted opener arriving as plain text.""" + engines, with the corrupted opener arriving as plain text. + """ tokens = _dsv4_tokens( reasoning="Checking the weather.", tool_name="get_weather", diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py index 2bc4338785cb..7aaad7c4dfc1 100644 --- a/tests/parser/engine/test_delegating_replay.py +++ b/tests/parser/engine/test_delegating_replay.py @@ -174,7 +174,8 @@ def test_delegating_replay(parser_cls, sample, chunk_size): ) def test_delegating_parse_tool_choice_none(parser_cls, parser_name, sample): """Non-streaming parse() with tool_choice='none' via DelegatingParser - must not leak special tokens into content.""" + must not leak special tokens into content. + """ tokenizer = make_mock_tokenizer(sample) validated_tools = ( _TOOLS_VALIDATOR.validate_python(sample.tools) if sample.tools else None diff --git a/tests/parser/engine/test_engine.py b/tests/parser/engine/test_engine.py index 36fa1ad8d09a..f638cd599f83 100644 --- a/tests/parser/engine/test_engine.py +++ b/tests/parser/engine/test_engine.py @@ -275,7 +275,8 @@ def test_text_between_tool_calls(self): def test_unmatched_close_brace_does_not_poison_depth(self): """A stray } in malformed JSON must not kill streaming for - all subsequent content.""" + all subsequent content. + """ engine = StreamingParserEngine(_hermes_config(), tokenizer=None) engine.feed("", []) @@ -291,7 +292,8 @@ def test_unmatched_close_brace_does_not_poison_depth(self): def test_json_args_no_premature_close_brace(self): """Closing braces of the top-level JSON shouldn't be streamed - until confirmed by the end tag.""" + until confirmed by the end tag. + """ engine = StreamingParserEngine(_hermes_config(), tokenizer=None) engine.feed("", []) @@ -346,7 +348,8 @@ class TestLexerBufferFlush: def test_buffered_prefix_emitted_in_current_state(self): """Text buffered by the lexer (e.g. '<') must be emitted as - REASONING_CHUNK before THINK_END transitions to CONTENT.""" + REASONING_CHUNK before THINK_END transitions to CONTENT. + """ engine = StreamingParserEngine(_think_config(), _make_think_tokenizer()) events = engine.feed("", [_START_ID]) @@ -384,11 +387,13 @@ def test_empty_buffer_no_extra_events(self): class TestTokenIdFiltering: """When token IDs are available, lex-matched terminals that also - have token_id_terminal entries should be demoted to content.""" + have token_id_terminal entries should be demoted to content. + """ def test_lex_matched_terminal_demoted_after_token_ids_seen(self): """After receiving token IDs, text that matches a token-ID - terminal should be treated as content, not trigger a transition.""" + terminal should be treated as content, not trigger a transition. + """ engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) # First feed with a non-special token ID to set _ever_had_token_ids @@ -409,7 +414,8 @@ def test_lex_matched_terminal_demoted_after_token_ids_seen(self): def test_scanner_matched_terminal_bypasses_filter(self): """PreLexedTerminals from the scanner bypass the filter and - still trigger state transitions.""" + still trigger state transitions. + """ engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events = engine.feed("", [_TOOL_START_ID]) @@ -422,7 +428,8 @@ def test_scanner_matched_terminal_bypasses_filter(self): def test_no_filtering_without_token_ids(self): """When no token IDs are ever provided (non-streaming), - text matching still triggers transitions.""" + text matching still triggers transitions. + """ engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events = engine.feed('{"name": "f"}', []) @@ -434,7 +441,8 @@ def test_no_filtering_without_token_ids(self): def test_mixed_text_then_real_tool_call(self): """Text mentioning tool syntax followed by a real special-token - tool call.""" + tool call. + """ engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events1 = engine.feed("Mention in text. ", [1, 2, 3, 4]) @@ -494,7 +502,8 @@ def test_deferred_reasoning_tokens_stay_before_next_end_terminal(self): def _func_prefix_config() -> ParserEngineConfig: """Config mixing token-ID terminals (TOOL_START/END) with - text-only terminals (FUNC_PREFIX) and fallback transitions.""" + text-only terminals (FUNC_PREFIX) and fallback transitions. + """ return ParserEngineConfig( name="func_prefix_test", terminals={ @@ -563,11 +572,13 @@ def _make_func_prefix_tokenizer(): class TestTextOnlyFallbackFiltering: """When token IDs are available, transitions marked - skip_in_token_id_mode should be skipped.""" + skip_in_token_id_mode should be skipped. + """ def test_func_prefix_in_prose_demoted_in_strict_mode(self): """ in prose should NOT trigger a tool call - when strict mode is active.""" + when strict mode is active. + """ engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -584,7 +595,8 @@ def test_func_prefix_in_prose_demoted_in_strict_mode(self): def test_normal_flow_after_tool_start_still_works(self): """TOOL_START (special token) -> FUNC_PREFIX (text) should - still parse a tool call normally in strict mode.""" + still parse a tool call normally in strict mode. + """ engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -604,7 +616,8 @@ def test_normal_flow_after_tool_start_still_works(self): def test_fallback_fires_without_token_ids(self): """When no token IDs are provided, fallback transitions should - still fire normally.""" + still fire normally. + """ engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -618,7 +631,8 @@ def test_fallback_fires_without_token_ids(self): def test_tool_between_fallback_blocked_in_strict_mode(self): """The (TOOL_BETWEEN, FUNC_PREFIX) fallback should also be - blocked in strict mode.""" + blocked in strict mode. + """ engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -648,7 +662,8 @@ def test_no_tokenizer_attribute(self): class TestArgsResetOnReentry: """When leaving TOOL_ARGS and later re-entering (e.g. two tool calls), the entering-TOOL_ARGS block resets args tracking. The - redundant reset on exit was removed.""" + redundant reset on exit was removed. + """ @staticmethod def _multi_tool_config() -> ParserEngineConfig: @@ -723,7 +738,8 @@ def test_brace_depth_resets_on_reentry(self): class TestToolPreambleFinish: """finish() in TOOL_PREAMBLE state emits TOOL_CALL_END when a tool - call was started (tool_index >= 0), but not when tool_index is -1.""" + call was started (tool_index >= 0), but not when tool_index is -1. + """ @staticmethod def _preamble_with_tool_call_start_config() -> ParserEngineConfig: @@ -782,7 +798,8 @@ def test_finish_no_tool_call_end_without_tool_index(self): class TestRegexTerminalInfraRemoved: """TerminalDef.priority, LexerShape.regex_terminals, and the regex - matching loop were removed.""" + matching loop were removed. + """ def test_terminal_def_no_priority(self): import regex as re @@ -807,7 +824,8 @@ def test_terminals_from_literals_still_works(self): class TestMultiCharTerminalInArgs: """Regression: multi-char terminals falling through in TOOL_ARGS - must be fed char-by-char via _feed_args_text, not _feed_args_char.""" + must be fed char-by-char via _feed_args_text, not _feed_args_char. + """ @staticmethod def _newline_config() -> ParserEngineConfig: @@ -899,7 +917,8 @@ def test_reset_preserves_skip_tool_parsing(self): def _message_header_config(reasoning_end: bool) -> ParserEngineConfig: """MESSAGE_HEADER-based grammar whose tool opener optionally carries - REASONING_END, mirroring the two shapes a transition table can take.""" + REASONING_END, mirroring the two shapes a transition table can take. + """ return ParserEngineConfig( name=f"message_header_test_{reasoning_end}", initial_state=ParserState.MESSAGE_HEADER, diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index ae21d69fc303..ab83bdf5ed4c 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -1097,7 +1097,8 @@ def test_bug_report_scenario(self, tool_call_parser, mock_request): def test_both_extractions_independent(self, parser, request_obj): """Calling extract_reasoning then extract_tool_calls on the same - parser instance should both work (each resets the engine).""" + parser instance should both work (each resets the engine). + """ model_output = FULL_MODEL_OUTPUT reasoning, _ = parser.extract_reasoning(model_output, request_obj) @@ -1111,7 +1112,8 @@ def test_both_extractions_independent(self, parser, request_obj): class TestAdapterExtractReasoning: """The reasoning adapter's extract_reasoning uses skip_tool_parsing - so tool call text is preserved as content for the tool adapter.""" + so tool call text is preserved as content for the tool adapter. + """ @pytest.fixture def adapter(self, mock_tokenizer): @@ -1152,7 +1154,8 @@ def test_no_reasoning_returns_none(self, adapter, request_obj): class TestGemma4SchemaAwareTypeCoercion: """Verify that streaming and non-streaming produce identical type-fixed arguments when tool schemas declare string parameters - but the model outputs bare numbers/booleans.""" + but the model outputs bare numbers/booleans. + """ @pytest.fixture def tools(self): @@ -1182,7 +1185,8 @@ def parser_with_tools(self, tool_call_tokenizer, tools): def test_streaming_string_param_not_coerced(self, parser_with_tools, mock_request): """A numeric value for a string-typed param must remain a string - in the streamed output, matching the non-streaming result.""" + in the streamed output, matching the non-streaming result. + """ chunks = [ "<|tool_call>", "call:update_record{", @@ -1213,7 +1217,8 @@ def test_streaming_mixed_types(self, parser_with_tools, mock_request): def test_streaming_matches_non_streaming(self, parser_with_tools, mock_request): """Concatenated streaming deltas must produce the same arguments - as non-streaming extraction.""" + as non-streaming extraction. + """ text = "<|tool_call>call:update_record{zipcode:12345}" non_streaming = parser_with_tools.extract_tool_calls(text, mock_request) @@ -1234,7 +1239,8 @@ def test_streaming_matches_non_streaming(self, parser_with_tools, mock_request): class TestGemma4SchemaCoercionBoolNumberNull: """Verify that _fix_arg_types coerces string values to non-string - schema types for the Gemma4 parser.""" + schema types for the Gemma4 parser. + """ @pytest.fixture def tools(self): @@ -1295,7 +1301,8 @@ def test_null_stays_string_without_null_schema( def test_streaming_type_stability(self, parser_with_tools, mock_request): """Values streamed incrementally must not cause prefix - incompatibility when types are coerced.""" + incompatibility when types are coerced. + """ text = ( "<|tool_call>call:configure{" "enabled:true," @@ -1395,7 +1402,8 @@ def test_nested_object_coerced(self, parser_with_tools, mock_request): class TestBareThoughtWithoutChannelOpener: """When the model omits <|channel> and starts with bare ``thought``, the parser should auto-inject the channel opener so reasoning is - captured correctly.""" + captured correctly. + """ @pytest.fixture def bare_thought_tokenizer(self): @@ -1466,7 +1474,8 @@ def test_normal_content_not_classified_as_reasoning(self, request_obj): def test_bare_thought_token_at_end_of_stream(self, request_obj): """When the stream ends with just "thought" (no \\n), the parser - should treat it as the thought prefix token, not real reasoning.""" + should treat it as the thought prefix token, not real reasoning. + """ seq: list[tuple[int, str]] = [ (CHANNEL_START_ID, "<|channel>"), (3000, "thought"), diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py index 58fe7aa28cbd..fe6d0ec5c5ad 100644 --- a/tests/parser/engine/test_glm47_moe.py +++ b/tests/parser/engine/test_glm47_moe.py @@ -61,7 +61,8 @@ def test_tool_call_without_think_end(self, parser): def test_previous_turn_tool_call_ignored(self, parser): """A finished tool call from an earlier turn says nothing about the - turn currently being generated.""" + turn currently being generated. + """ prompt = [THINK_S, TEXT, TOOL_S, TEXT, TOOL_E, OBSERVATION, TEXT, ASSISTANT] assert not parser.is_reasoning_end(prompt) @@ -88,7 +89,8 @@ def test_every_tool_call_after_think_end_kept(self, parser): def test_falls_back_to_tool_call(self, parser): """Without ````, content starts at the opener itself so the - tool parser still receives a well-formed call.""" + tool parser still receives a well-formed call. + """ ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E] assert parser.extract_content_ids(ids) == [TOOL_S, 30, TOOL_E] diff --git a/tests/parser/engine/test_inkling.py b/tests/parser/engine/test_inkling.py index 9026d548dfbb..c38e74544c31 100644 --- a/tests/parser/engine/test_inkling.py +++ b/tests/parser/engine/test_inkling.py @@ -70,7 +70,8 @@ def _tool_block(name: str, args: str) -> str: def _tokenize(text: str) -> list[tuple[int, str]]: """Tokenize like the real stream: markers are atomic special tokens, plain text becomes one token per character (matching the mock - tokenizer's ``chr``-based decode).""" + tokenizer's ``chr``-based decode). + """ tokens: list[tuple[int, str]] = [] i = 0 while i < len(text): @@ -87,7 +88,8 @@ def _tokenize(text: str) -> list[tuple[int, str]]: def _stream(parser, request, text: str, chunk_size: int): """Stream production-shaped deltas: ``chunk_size`` tokens per delta, - with delta_token_ids covering every token (specials and text).""" + with delta_token_ids covering every token (specials and text). + """ tokens = _tokenize(text) results = [] previous_text = "" @@ -119,7 +121,8 @@ def _stream(parser, request, text: str, chunk_size: int): def _stream_text_only(parser, request, text: str, chunk_size: int): """Stream text-only deltas (no token ids), chunked at arbitrary character boundaries — exercises the text-lexing fallback path, - including markers split across chunks.""" + including markers split across chunks. + """ results = [] previous_text = "" for start in range(0, len(text), chunk_size): @@ -170,7 +173,8 @@ def _stream_delegating(parser, request, text, chunk_size, prompt_token_ids): """Stream ``text`` through ``DelegatingParser.parse_delta``, ``chunk_size`` tokens per delta; return ``(content, reasoning, ordered tool names, ordered tool arguments)``. Arguments arrive in fragments, so they are - concatenated per tool index.""" + concatenated per tool index. + """ tokens = _tokenize(text) content, reasoning = "", "" tools: dict[int, str] = {} @@ -522,7 +526,8 @@ class TestToolCallFiltering: """Inkling equivalents of the generic tool-call-filtering replay tests (Inkling is excluded from those in test_replay.py: its structural role/kind tokens and shared block-end token don't fit the generic - reasoning/tool split model).""" + reasoning/tool split model). + """ def test_skip_tool_parsing_round_trip(self, mock_tokenizer, mock_request): # First pass (reasoning adapter, skip_tool_parsing): reasoning is @@ -820,7 +825,8 @@ def test_function_name_header_before_tool_start_streaming( ): """The optional function name between ``<|message_model|>`` and the content-kind marker is metadata: the buffered header must be - discarded on the way out, not flushed into content.""" + discarded on the way out, not flushed into content. + """ tools = [_function_tool()] mock_request.tools = tools content, _, names, args = _stream_delegating( @@ -837,7 +843,8 @@ def test_function_name_header_before_tool_start_streaming( def test_content_state_tool_start_streaming(self, mock_tokenizer, mock_request): """Same opener reached from CONTENT rather than MESSAGE_HEADER: a text block closed with no ``<|message_model|>`` before the tool - block. Preceding text must reach content exactly once, unmarked.""" + block. Preceding text must reach content exactly once, unmarked. + """ tools = [_function_tool()] mock_request.tools = tools content, _, names, args = _stream_delegating( @@ -859,7 +866,8 @@ def test_content_tool_start_emits_reasoning_end_in_reasoning_pass(): boundary (#49876), but the header-flush path (MESSAGE_HEADER --END_MESSAGE--> CONTENT) reaches CONTENT without one, so a tool block opening from there relies on this transition - alone to hand off to the tool pass.""" + alone to hand off to the tool pass. + """ engine = StreamingParserEngine(inkling_config(), tokenizer=None) engine.skip_tool_parsing = True engine.reset(initial_state=ParserState.CONTENT) @@ -876,7 +884,8 @@ def test_skip_reasoning_parsing_inert_for_shared_markers(): markers. Inkling has none — ``<|end_message|>`` is labelled THINK_END yet also closes text, header, and tool blocks — so the flag must be inert: a tool block still closes through its transition and following - text returns to CONTENT instead of leaking into the argument stream.""" + text returns to CONTENT instead of leaking into the argument stream. + """ engine = StreamingParserEngine(inkling_config(), tokenizer=None) engine.skip_reasoning_parsing = True engine.reset(initial_state=ParserState.CONTENT) @@ -896,7 +905,8 @@ class TestToolParserWithoutReasoningParser: labels ``<|end_message|>`` as THINK_END; a bypass keyed on the label would neutralize the closer of every block kind, leaking markers into content and leaving tool calls unterminated. Structure must keep - parsing exactly as with the reasoning parser attached.""" + parsing exactly as with the reasoning parser attached. + """ GEN_PROMPT = [_TML_VOCAB[MSG_MODEL]] diff --git a/tests/parser/engine/test_nemotron_v3.py b/tests/parser/engine/test_nemotron_v3.py index a7f166f9f3a9..201ddc4e781c 100644 --- a/tests/parser/engine/test_nemotron_v3.py +++ b/tests/parser/engine/test_nemotron_v3.py @@ -56,7 +56,8 @@ def parser(): class TestNemotronSwap: def test_enable_thinking_false_swaps(self, parser): """When enable_thinking=False, model output without think tags - should have reasoning swapped to content.""" + should have reasoning swapped to content. + """ text = "The answer is 42." request = _make_request(enable_thinking=False) reasoning, content = parser.extract_reasoning(text, request) @@ -73,7 +74,8 @@ def test_force_nonempty_content_swaps(self, parser): def test_no_swap_when_content_exists(self, parser): """With enable_thinking=False but real giving content, - no swap occurs.""" + no swap occurs. + """ text = "Some reasoning.Actual content here." request = _make_request(enable_thinking=False) reasoning, content = parser.extract_reasoning(text, request) @@ -180,11 +182,13 @@ def test_streaming_tool_calls(self, parser): class TestParseDeltaTokenIdFiltering: """parse_delta must not trigger tool call parsing when - appears as regular text rather than as a special token ID.""" + appears as regular text rather than as a special token ID. + """ def test_tool_call_text_in_reasoning_is_not_parsed(self, parser): """Literal in model reasoning should be content, - not a tool call.""" + not a tool call. + """ request = _make_request() text = ( @@ -210,7 +214,8 @@ def test_tool_call_text_in_reasoning_is_not_parsed(self, parser): def test_special_token_id_still_triggers_tool_call(self, parser): """When the scanner matches a special token ID, the tool call - must still be parsed correctly.""" + must still be parsed correctly. + """ request = _make_request() parser.parse_delta( @@ -250,7 +255,8 @@ def test_special_token_id_still_triggers_tool_call(self, parser): def test_text_discussion_then_real_tool_call(self, parser): """Model discusses tool syntax in reasoning, then makes a real - tool call via special tokens.""" + tool call via special tokens. + """ request = _make_request() r1 = parser.parse_delta( diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 8fd7c4d2623f..a64ab2ecd541 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -323,7 +323,8 @@ def test_kimi_k2_tool_call_id_includes_func_name(self): def test_multiple_arg_chunks_same_batch_coalesced(self): """Multiple events for the same tool in one batch must produce - at most one DeltaToolCall per index.""" + at most one DeltaToolCall per index. + """ engine = _make_engine() events = [ SemanticEvent(EventType.TOOL_CALL_START, tool_index=0), @@ -551,7 +552,8 @@ class TestPostToolContentDeferral: """Regression: content after TOOL_CALL_END in the same batch must not produce a mixed DeltaMessage(content=..., tool_calls=...) — that causes split_delta to reorder content before tool_calls, breaking the Responses - API state machine.""" + API state machine. + """ def test_text_after_tool_end_deferred(self): engine = _make_engine() @@ -605,7 +607,8 @@ def test_deferred_content_not_flushed_during_arg_continuation(self): """Deferred content from batch N must not mix with arg-continuation tool events in batch N+1 — that creates a DeltaMessage with both content and nameless tool_calls, which crashes the Responses API - state machine (name=None → Pydantic ValidationError).""" + state machine (name=None → Pydantic ValidationError). + """ engine = _make_engine() engine._content_has_nonws = True @@ -834,7 +837,8 @@ def test_content_passthrough(self): class TestEngineBasedPath: """Tests for the _engine_based accumulation behavior in - DelegatingParser.parse_delta.""" + DelegatingParser.parse_delta. + """ def test_engine_based_true_when_both_parsers_engine(self): r = SimpleNamespace(engine_based_streaming=True) @@ -880,7 +884,8 @@ def test_parse_delta_tool_call(self, mock_request): class TestParseTokenIdPassthrough: """parse() must forward model_output_token_ids to _single_pass_parse - so that token-ID-based strict terminal matching is active.""" + so that token-ID-based strict terminal matching is active. + """ def test_literal_tool_tag_in_content_preserved_with_token_ids(self, mock_request): engine = _make_engine(_hermes_config()) @@ -1104,7 +1109,8 @@ class TestAdapterFinishOnStreamEnd: def test_lexer_buffer_flushed_on_finished(self): """Text buffered as a potential terminal prefix must be emitted - as content when the stream ends.""" + as content when the stream ends. + """ tokenizer = make_mock_tokenizer(_VOCAB) parser = _CombinedDelegating(tokenizer) request = _make_delegating_request() @@ -1122,7 +1128,8 @@ def test_lexer_buffer_flushed_on_finished(self): def test_args_buffer_flushed_on_finished(self): """Pending arg buffer text must be emitted when stream ends - mid-tool-call (closing brace held back in buffer).""" + mid-tool-call (closing brace held back in buffer). + """ tokenizer = make_mock_tokenizer(_VOCAB) parser = _CombinedDelegating(tokenizer) request = _make_delegating_request() @@ -1358,7 +1365,8 @@ def test_content_after_transitionless_closer_observable(self): class TestToolAdapterForwardsKwargs: """ParserEngineToolAdapter.__init__ must forward **kwargs to the - parser engine class so chat_template_kwargs reach model parsers.""" + parser engine class so chat_template_kwargs reach model parsers. + """ @pytest.mark.parametrize( "enable_thinking,expected_state", @@ -1388,7 +1396,8 @@ def test_kwargs_forwarded_to_parser_engine(self, enable_thinking, expected_state class TestExtractContentIdsNoEmptyReturn: """extract_content_ids must return input_ids (not []) when there is - no THINK_END token ID and _reasoning_ended is True.""" + no THINK_END token ID and _reasoning_ended is True. + """ _NO_THINK_CONFIG = ParserEngineConfig(name="no_think_end", token_id_terminals={}) @@ -1786,7 +1795,8 @@ class TestDropSpecialTokens: def test_drops_special_token_by_id_from_content(self): """A special token (not a configured terminal) is dropped when - it arrives as its actual token ID.""" + it arrives as its actual token ID. + """ config = ParserEngineConfig( name="drop_content_test", terminals={}, @@ -1821,7 +1831,8 @@ def test_drops_special_token_by_id_from_reasoning(self): def test_drops_via_text_fallback_when_no_token_ids(self): """When no token IDs are provided, text-based lexer catches - drop tokens as a fallback.""" + drop tokens as a fallback. + """ engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1834,7 +1845,8 @@ def test_drops_via_text_fallback_when_no_token_ids(self): def test_regular_tokens_spelling_special_survive(self): """Regular tokens that spell out a drop-token string survive - when token IDs prove they are not the special token.""" + when token IDs prove they are not the special token. + """ vocab = {**_DROP_VOCAB, "h": 72, "<": 73, "bos": 74, ">": 75, "w": 76} engine = _make_engine( vocab=vocab, @@ -1850,7 +1862,8 @@ def test_regular_tokens_spelling_special_survive(self): def test_configured_terminal_not_treated_as_drop(self): """Tokens already in config.terminals (like ) are handled - by the state machine, not the drop mechanism.""" + by the state machine, not the drop mechanism. + """ engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1905,7 +1918,8 @@ def test_no_special_tokens_means_no_drops(self): def test_drops_applied_with_skip_tool_parsing(self): """Drop tokens are always dropped, even with skip_tool_parsing. DROP_TERMINALs have no transitions by construction, so no parser - pass can use them.""" + pass can use them. + """ for initial_state in (ParserState.REASONING, ParserState.CONTENT): engine = _make_engine( vocab=_DROP_VOCAB, @@ -1921,7 +1935,8 @@ def test_drops_applied_with_skip_tool_parsing(self): def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self): """With skip_tool_parsing in REASONING state, drop tokens are - removed but configured terminals still fire their transitions.""" + removed but configured terminals still fire their transitions. + """ engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1957,7 +1972,8 @@ def test_drops_in_tool_args_state(self): def test_mixed_configured_and_drop_terminals(self): """Configured terminals trigger transitions while drop terminals - are silently removed in the same stream.""" + are silently removed in the same stream. + """ engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1981,7 +1997,8 @@ class TestTruncatedToolOpenerStreamParity: """Regression tests for #47137: when generation terminates inside a ```` opener that has not been promoted to a tool call (via ``max_tokens`` or a ``stop`` string), the non-streaming path - must drop the incomplete markup, matching the streaming path.""" + must drop the incomplete markup, matching the streaming path. + """ _QWEN3_VOCAB = { "": 100, @@ -2027,7 +2044,8 @@ def _stream_content(self, request, chunks: list[str]) -> str: ) def test_truncated_opener_dropped_in_both_paths(self, mock_request, chunks): """A cutoff inside the opener yields no content and no tool calls - in both the non-streaming and streaming paths.""" + in both the non-streaming and streaming paths. + """ text = "".join(chunks) _, content, tool_calls = self._make_parser().parse( @@ -2040,7 +2058,8 @@ def test_truncated_opener_dropped_in_both_paths(self, mock_request, chunks): def test_content_before_truncated_opener_preserved(self, mock_request): """Only the incomplete markup is dropped; content generated before - the opener is returned identically by both paths.""" + the opener is returned identically by both paths. + """ chunks = ["Checking the weather. ", "", "\n", "", "garbage", "", " B"] text = "".join(chunks) @@ -2068,7 +2088,8 @@ def test_content_around_unpromoted_tool_block_stays_ordered(self, mock_request): def test_complete_tool_call_still_promoted(self, mock_request): """Sanity check: a complete tool call still parses in the - non-streaming path after the truncation fix.""" + non-streaming path after the truncation fix. + """ text = ( "\n" "\n" @@ -2088,7 +2109,8 @@ class TestThinkMarkupWithoutReasoningParser: """With no reasoning parser configured, reasoning markup is plain content: engine-backed tool parsers must return it verbatim in both the non-streaming and streaming paths, not consume the markers or - reclassify the block as reasoning and drop it.""" + reclassify the block as reasoning and drop it. + """ _VOCAB = { "": 90, @@ -2154,7 +2176,8 @@ def test_think_block_passes_through_both_paths( def test_think_block_preserved_alongside_tool_call(self, mock_request): """The tools-called branch also returns the tool parser's content; - a think block preceding a promoted tool call must survive in it.""" + a think block preceding a promoted tool call must survive in it. + """ think = "plan" text = think + ( "\n" diff --git a/tests/parser/engine/test_qwen3.py b/tests/parser/engine/test_qwen3.py index 07a025876ae7..b8f96294cb79 100644 --- a/tests/parser/engine/test_qwen3.py +++ b/tests/parser/engine/test_qwen3.py @@ -318,7 +318,8 @@ def test_streaming_multi_param(self, parser, mock_request): def test_streaming_args_arrive_incrementally(self, parser, mock_request): """Arguments must stream as intermediate deltas, not batch at - tool-end.""" + tool-end. + """ chunks = [ "\n", "\n", @@ -523,7 +524,8 @@ def test_streaming_value_split_across_chunks(self, parser, mock_request): def test_streaming_split_tool_call_tag(self, parser, mock_request): """ arrives as a single special token; the rest of - the content is split into fine-grained chunks.""" + the content is split into fine-grained chunks. + """ chunks = [ "\n", "\n", @@ -700,7 +702,8 @@ def test_partial_value_with_angle_bracket_and_complete_param(self): class TestSchemaAwareTypeCoercion: """Verify that _fix_arg_types corrects miscoerced values using the - tool schema.""" + tool schema. + """ @pytest.fixture def tools(self): @@ -842,7 +845,8 @@ def parser_with_anyof(self, mock_tokenizer, tools_with_anyof): def test_anyof_string_param_not_coerced(self, parser_with_anyof, mock_request): """A param with anyOf including 'string' must not be coerced - to integer.""" + to integer. + """ text = ( "\n" "\n" @@ -858,7 +862,8 @@ def test_anyof_string_param_not_coerced(self, parser_with_anyof, mock_request): class TestSchemaCoercionBoolNumberNull: """Verify that _fix_arg_types coerces string values to non-string - schema types using coerce_to_schema_type.""" + schema types using coerce_to_schema_type. + """ @pytest.fixture def tools(self): diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index 3f285e4d0315..d8e8096fd3b6 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -276,7 +276,8 @@ def test_continued_turn_with_open_thinking_not_end(self, parser): def test_empty_boundary_config_keeps_global_walk(self): """Configs without turn_boundary_tokens keep the pre-existing - behavior: any in the sequence ends reasoning.""" + behavior: any in the sequence ends reasoning. + """ cfg = qwen3_config(turn_boundary_tokens=frozenset()) parser = Qwen3Parser( make_mock_tokenizer(_QWEN3_VOCAB), parser_engine_config=cfg @@ -288,7 +289,8 @@ def test_empty_boundary_config_keeps_global_walk(self): def test_boundary_tokens_missing_from_vocab_keep_global_walk(self): """Boundary tokens absent from the vocabulary resolve to nothing, so subclasses reusing the default on a non-ChatML vocab keep the - pre-existing behavior.""" + pre-existing behavior. + """ vocab = { token: token_id for token, token_id in _QWEN3_VOCAB.items() @@ -299,7 +301,8 @@ def test_boundary_tokens_missing_from_vocab_keep_global_walk(self): def test_boundary_with_thinking_disabled_is_end(self, mock_tokenizer): """With thinking disabled (initial state CONTENT) a walk that stops at a turn - boundary must report reasoning as ended: the model never emits a marker.""" + boundary must report reasoning as ended: the model never emits a marker. + """ parser = Qwen3Parser( mock_tokenizer, chat_template_kwargs={"enable_thinking": False} ) @@ -307,7 +310,8 @@ def test_boundary_with_thinking_disabled_is_end(self, mock_tokenizer): def test_boundary_with_thinking_enabled_not_end(self, parser): """With thinking disabled (initial state REASONING) a walk that stops at a turn - boundary must keep reasoning open""" + boundary must keep reasoning open + """ assert not parser.is_reasoning_end([_IM_START_ID, _TEXT_ID]) @@ -668,7 +672,8 @@ def test_streaming_trailing_newlines_before_tool_call(self, parser_with_strip): class TestWhitespaceStrippingDisabled: """When strip_trailing_reasoning_whitespace is False, - trailing whitespace in reasoning must be preserved.""" + trailing whitespace in reasoning must be preserved. + """ @pytest.fixture def parser_no_strip(self): @@ -726,7 +731,8 @@ def test_default_initial_state_is_reasoning(self, mock_tokenizer): def test_thinking_disabled_streaming_content_only(self, mock_tokenizer): """Plain text with thinking disabled must stream as content, not reasoning. Before the fix, the REASONING initial state caused all - output to be emitted as reasoning chunks.""" + output to be emitted as reasoning chunks. + """ p = Qwen3Parser( mock_tokenizer, chat_template_kwargs={"enable_thinking": False}, diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py index fee74f8d2cca..75fdbfca0fcc 100644 --- a/tests/parser/engine/test_replay.py +++ b/tests/parser/engine/test_replay.py @@ -440,7 +440,8 @@ def test_replay(self, parser_cls, sample, think_end, tool_start, mode, chunk_siz ) class TestToolCallFilteringNonStreaming: """Non-streaming parse() with tool_choice='none' must suppress tool - calls and not leak special tokens into content.""" + calls and not leak special tokens into content. + """ def test_parse(self, parser_cls, sample, think_end, tool_start): tokenizer = make_mock_tokenizer(sample) @@ -483,7 +484,8 @@ def test_parse(self, parser_cls, sample, think_end, tool_start): ) class TestToolChoiceNoneStreamingParity: """Streaming and non-streaming must return the same content - when tool_choice='none' suppresses tool calls.""" + when tool_choice='none' suppresses tool calls. + """ def test_content_matches(self, parser_cls, sample): tokenizer = make_mock_tokenizer(sample) @@ -541,7 +543,8 @@ def _inject_drop_tokens(sample): class TestDropTokenReplay: """Verify unconfigured special tokens are silently dropped across - all parsers and chunk sizes.""" + all parsers and chunk sizes. + """ @pytest.mark.parametrize( "parser_info", @@ -598,7 +601,8 @@ def test_drop_tokens_removed_from_output(self, parser_info): class TestAdapterReferences: """Verify make_adapters sets reasoning/tool parser class refs on parser engine - parser classes so the serving layer finds them and calls adjust_request.""" + parser classes so the serving layer finds them and calls adjust_request. + """ @pytest.mark.parametrize( "parser_name", diff --git a/tests/parser/engine/test_seed_oss.py b/tests/parser/engine/test_seed_oss.py index e01f2fd9f6d4..159832135923 100644 --- a/tests/parser/engine/test_seed_oss.py +++ b/tests/parser/engine/test_seed_oss.py @@ -129,7 +129,8 @@ def test_single_tool_call(tool_parser, mock_request): def test_malformed_function_end_does_not_drop_siblings(tool_parser, mock_request): """Regression for #46314: a malformed ```` with no closing ``>`` - on the header must not discard the other, well-formed calls.""" + on the header must not discard the other, well-formed calls. + """ text = ( f"{TOOL_CALL_START}\n\n{TOOL_CALL_END}" f"{TOOL_CALL_START}\n\n" @@ -171,7 +172,8 @@ def test_reasoning_then_tool_call(parser): def test_streaming_think_end_and_tool_call_same_delta(parser): """```` and ```` arriving in one delta must - not leak the terminal tokens into the reasoning text.""" + not leak the terminal tokens into the reasoning text. + """ reasoning, content = simulate_reasoning_streaming( parser, [ diff --git a/tests/parser/engine/test_ufffd_reasoning_transition.py b/tests/parser/engine/test_ufffd_reasoning_transition.py index ffd2ead95d7b..6497a6c29f63 100644 --- a/tests/parser/engine/test_ufffd_reasoning_transition.py +++ b/tests/parser/engine/test_ufffd_reasoning_transition.py @@ -146,7 +146,8 @@ def test_no_ufffd(self, chunk_size, vocab, delegating_cls): def test_byte_fallback_tokenizer_produces_ufffd(self): """Validate the fixture: decode() returns U+FFFD for isolated - byte-fallback token IDs, proving the old code path would leak.""" + byte-fallback token IDs, proving the old code path would leak. + """ vocab = dict(_MODEL_CONFIGS[0].values[0]) tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS) assert tokenizer.decode([200]) == "�" diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 36e347fc572f..347df05b4d46 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -992,7 +992,8 @@ def _inkling_block( ) -> None: """Append one Inkling content block; blocks after the first start with the ``<|message_model|>`` role token (the first block continues the - generation prompt directly).""" + generation prompt directly). + """ if segs: segs.append(("<|message_model|>", True)) segs.append((kind_token, True)) diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py index d5e09453a035..8645a38dd3cb 100644 --- a/tests/parser/mistral/test_tool_calls.py +++ b/tests/parser/mistral/test_tool_calls.py @@ -492,7 +492,8 @@ def test_extract_tool_calls_pre_v11_regex_fallback( ): """The regex fallback path finds valid JSON via regex when the primary raw_decode fails on leading junk. It should re-serialize arguments - and return a valid tool call.""" + and return a valid tool call. + """ model_output = ( '[TOOL_CALLS] junk [{"name": "add", "arguments":{"a": 1, "b": 2}}] trail' ) @@ -2003,7 +2004,8 @@ def test_streaming_pre_v11_parallel_calls_batched_deltas( mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size, driver ): """A batched delta spanning the boundary between two parallel calls must - keep them on distinct indices (the bug collapsed both onto index 0).""" + keep them on distinct indices (the bug collapsed both onto index 0). + """ model_output = ( '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, ' '{"name": "get_current_weather", "arguments": ' @@ -2045,7 +2047,6 @@ def test_content_tool_calls_transition_emits_reasoning_end( reasoning_encoding, ): """(CONTENT, TOOL_CALLS) must emit REASONING_END when reasoning is enabled.""" - cfg = mistral_config(reasoning_encoding=reasoning_encoding) engine = StreamingParserEngine( config=cfg, tokenizer=mistral_tokenizer, vocab=mistral_tokenizer.get_vocab() @@ -2122,7 +2123,8 @@ def test_reasoning_active_no_think_block_no_leak( chunk_size, ): """Tool call without a preceding think block must not leak as content - when the reasoning parser is active.""" + when the reasoning parser is active. + """ accumulated_content = "" function_names: list[str] = [] function_args_strs: list[str] = [] @@ -2219,11 +2221,12 @@ def test_adjust_request_pre_v11_required_clears_response_format( mistral_pre_v11_tool_parser: MistralToolParser, response_format: dict, ) -> None: - """required + response_format must inject the tool schema and clear + """Required + response_format must inject the tool schema and clear response_format so the tool schema is the sole structured-output constraint, matching the base ToolParser (otherwise the request either hits the "multiple constraints" engine error or, with no constraint, - rambles to finish_reason='length').""" + rambles to finish_reason='length'). + """ request = _make_request(tool_choice="required", response_format=response_format) result = mistral_pre_v11_tool_parser.adjust_request(request) @@ -2240,7 +2243,8 @@ def test_adjust_request_non_mistral_tokenizer_required_injects_schema( ) -> None: """The guided-schema injection must also fire for non-Mistral (e.g. HF-mode) tokenizers driving the Mistral tool parser, mirroring the base ToolParser so - required + response_format does not fall back to an unconstrained ramble.""" + required + response_format does not fall back to an unconstrained ramble. + """ request = _make_request( tool_choice="required", response_format={"type": "json_object"} ) diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index cad16746be44..6c4158b81490 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -229,7 +229,8 @@ def _boundary_chunks(tokenizer, parser, end_token_id=None): def test_parse_delta_reasoning_not_dropped_on_boundary(tokenizer, request_obj): """Regression: reasoning must not be lost when a multi-token delta - spans the reasoning/tool-call boundary.""" + spans the reasoning/tool-call boundary. + """ parser = make_parser(tokenizer, reasoning=True, tool=True) chunks = _boundary_chunks(tokenizer, parser) results = stream_chunks(parser, tokenizer, chunks, request_obj) @@ -247,7 +248,8 @@ def test_parse_delta_reasoning_not_dropped_on_boundary(tokenizer, request_obj): def test_parse_delta_reasoning_boundary_no_tool_parser(tokenizer, request_obj): """When no tool parser is active, boundary-spanning chunks must still - preserve reasoning and pass post- text as content.""" + preserve reasoning and pass post- text as content. + """ parser = make_parser(tokenizer, reasoning=True, tool=False) chunks = _boundary_chunks(tokenizer, parser) results = stream_chunks(parser, tokenizer, chunks, request_obj) @@ -261,7 +263,8 @@ def test_parse_delta_reasoning_boundary_no_tool_parser(tokenizer, request_obj): def test_parse_delta_reasoning_only_no_think_leak(tokenizer, request_obj): """Regression: must not leak into content when streaming - token-by-token with reasoning=True, tool=False.""" + token-by-token with reasoning=True, tool=False. + """ parser = make_parser(tokenizer, reasoning=True, tool=False) results = stream_text( parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] @@ -303,7 +306,8 @@ def test_parse_delta_reasoning_only_thinking_disabled(tokenizer, request_obj): def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, request_obj): """When finished=True but the final parse_delta produces no - tool-call delta, unstreamed args are not flushed.""" + tool-call delta, unstreamed args are not flushed. + """ parser = make_parser(tokenizer, reasoning=False, tool=True) results = stream_text( @@ -326,7 +330,8 @@ def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, reques def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, request_obj): """When all args have been streamed, finished=True must not - produce extra or duplicate arguments.""" + produce extra or duplicate arguments. + """ parser = make_parser(tokenizer, reasoning=False, tool=True) results = stream_text( parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] @@ -346,7 +351,8 @@ def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, reque def test_parse_delta_finished_appends_remaining_args(tokenizer, request_obj): """When finished=True and the tool parser has unstreamed args, - parse_delta appends the remaining arguments to the tool-call delta.""" + parse_delta appends the remaining arguments to the tool-call delta. + """ parser = make_parser(tokenizer, reasoning=False, tool=True) token_ids = tokenizer.encode(MODEL_OUTPUT, add_special_tokens=False) @@ -493,7 +499,8 @@ def test_engine_reasoning_hermes_tool_token_by_token(tokenizer, request_obj): """Qwen3 engine reasoning + Hermes tool parser, token-by-token. Sanity check that the mixed engine/non-engine configuration works - when tokens arrive one at a time (no deferred content).""" + when tokens arrive one at a time (no deferred content). + """ parser = Qwen3ReasoningHermesToolParser(tokenizer) assert parser._reasoning_parser.engine_based_streaming is True @@ -520,7 +527,8 @@ def test_engine_reasoning_hermes_tool_boundary(tokenizer, request_obj): When and are in the same chunk with aligned text and token IDs, the engine processes both terminals and returns - the text as content.""" + the text as content. + """ parser = Qwen3ReasoningHermesToolParser(tokenizer) end_token_id = parser._reasoning_parser._parser_engine._reasoning_end_token_id chunks = _boundary_chunks(tokenizer, parser, end_token_id=end_token_id) @@ -551,7 +559,8 @@ def test_engine_reasoning_hermes_tool_text_holdback(tokenizer, request_obj): Without the fix, finish_streaming() is never called at the reasoning->tool transition when _engine_based is False, so the '<' is lost and the Hermes parser sees 'tool_call>...' instead of - '...'.""" + '...'. + """ parser = Qwen3ReasoningHermesToolParser(tokenizer) vocab = tokenizer.get_vocab() think_end_id = vocab[""] @@ -642,7 +651,8 @@ def test_engine_reasoning_no_tool_batched_content_passthrough(tokenizer, request batched delta carries ```` plus the following content (as happens with stream_interval > 1). The post-```` text must be emitted as content -- not dropped, not reclassified as - reasoning -- and the ```` marker must not leak either way.""" + reasoning -- and the ```` marker must not leak either way. + """ parser = Qwen3ReasoningNoToolParser(tokenizer) assert parser._reasoning_parser.engine_based_streaming is True assert parser._tool_parser is None @@ -670,7 +680,8 @@ def _decode_stream_deltas(tokenizer, groups): This mirrors how vLLM's detokenizer feeds ``parse_delta`` in production: byte-level UTF-8 hold-back means a character whose bytes span multiple tokens is only surfaced once complete (a naive - per-token ``decode`` would instead emit U+FFFD replacement chars).""" + per-token ``decode`` would instead emit U+FFFD replacement chars). + """ from tokenizers.decoders import DecodeStream stream = DecodeStream(skip_special_tokens=False) @@ -702,7 +713,8 @@ def test_engine_reasoning_hermes_tool_multibyte_holdback(tokenizer, request_obj) ``DecodeStream``. Its UTF-8 hold-back yields the correct character round-trip (a naive per-token decode would corrupt it), verifying the boundary stays byte-safe for multi-token - characters.""" + characters. + """ parser = Qwen3ReasoningHermesToolParser(tokenizer) vocab = tokenizer.get_vocab() think_end_id = vocab[""] diff --git a/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py index a35cb91bdb00..b2cf6f23455a 100644 --- a/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py +++ b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py @@ -121,7 +121,7 @@ def _build_document_prompt( tokenizer, content_ids: list[int], ) -> TokensPrompt: - """[CLS] [DocumentMarker] [SEP]""" + """[CLS] [DocumentMarker] [SEP].""" _, document_marker_id = self._resolve_marker_ids(tokenizer) content_ids = content_ids[: self.max_model_len - 3] diff --git a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py index d83399e07aa2..3c540f480aaa 100644 --- a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py +++ b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py @@ -58,6 +58,7 @@ def save_geotiff(image: torch.Tensor, meta: dict, out_format: str) -> str | byte image: np.ndarray with shape (bands, height, width) output_path: path where to save the image meta: dict with meta info. + """ if out_format == "path": # create temp file @@ -100,8 +101,8 @@ def read_geotiff( Returns: np.ndarray with shape (bands, height, width) meta info dict - """ + """ if all([x is None for x in [file_path, path_type, file_data]]): raise Exception("All input fields to read_geotiff are None") write_to_file: bytes | None = None @@ -168,8 +169,8 @@ def load_image( Returns: np.array containing created example list of meta info for each image in *file_paths* - """ + """ imgs: list[np.ndarray] = [] metas = [] temporal_coords = [] diff --git a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py index 3a1a9c3be41e..c75c5626161c 100644 --- a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py +++ b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py @@ -39,13 +39,13 @@ class ImagePrompt(BaseModel): class ImageRequestOutput(BaseModel): - """ - The output data of an image request to vLLM. + """The output data of an image request to vLLM. Args: type (str): The data content type [path, object] format (str): The image format (e.g., jpeg, png, etc.) data (Any): The resulting data. + """ type: Literal["path", "b64_json"] diff --git a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py index 7ba2c09b283c..6acd2a9f6dbf 100644 --- a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py +++ b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py @@ -5,8 +5,7 @@ class DummyStatLogger(StatLoggerBase): - """ - A dummy stat logger for testing purposes. + """A dummy stat logger for testing purposes. Implements the minimal interface expected by StatLoggerManager. """ diff --git a/tests/plugins_tests/gguf/test_gguf_plugin_generate.py b/tests/plugins_tests/gguf/test_gguf_plugin_generate.py index fbda46527530..2f35a272c0da 100644 --- a/tests/plugins_tests/gguf/test_gguf_plugin_generate.py +++ b/tests/plugins_tests/gguf/test_gguf_plugin_generate.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E tests for GGUF plugin functionality. -""" +"""E2E tests for GGUF plugin functionality.""" import os from typing import NamedTuple diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py index 2d921539004f..97b74b7bc0ca 100644 --- a/tests/plugins_tests/test_endpoint_plugins.py +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -81,7 +81,8 @@ def test_no_plugins_loaded_when_allowlist_is_empty_string( ): """`VLLM_PLUGINS=""` parses to `[""]`, not `None` (see `vllm.envs`), so it must be treated as a (non strict) allowlist matching no plugin name, not - as "unset".""" + as "unset". + """ monkeypatch.setenv("VLLM_PLUGINS", "") assert load_endpoint_plugins(("generate",)) == [] @@ -166,7 +167,8 @@ def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch) async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch): """`init_app_state` callers that never ran `build_app` (e.g. `run_batch.py`, which builds a bare `State()`) must not crash just - because `state.endpoint_plugins` was never set.""" + because `state.endpoint_plugins` was never set. + """ from starlette.datastructures import State monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") @@ -184,7 +186,8 @@ async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( """The CPU only render server has no `EngineClient` but a plugin eligible for the `render` task (`required_tasks` is `None` or includes `"render"`) still gets its routes attached at Phase A. Phase B passes `None` for - `engine_client` and it's up to the plugin to handle that.""" + `engine_client` and it's up to the plugin to handle that. + """ monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") args = _build_args() @@ -210,7 +213,8 @@ async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( @pytest.mark.asyncio async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch): """Phase A (attach) + Phase B (init) wired through `build_app` then - exercised with a real HTTP request against the worked example plugin.""" + exercised with a real HTTP request against the worked example plugin. + """ monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") args = _build_args() diff --git a/tests/plugins_tests/test_io_processor_plugins.py b/tests/plugins_tests/test_io_processor_plugins.py index 14e1a2d5769c..c5c226cb54ef 100644 --- a/tests/plugins_tests/test_io_processor_plugins.py +++ b/tests/plugins_tests/test_io_processor_plugins.py @@ -38,7 +38,8 @@ def my_plugin_entry_points(): """Patch importlib.metadata.entry_points to expose a single 'my_plugin' entry point backed by DummyIOProcessor, exercising the full plugin-loading code path: entry_points → plugin.load() → func() → - resolve_obj_by_qualname → IOProcessor.__init__.""" + resolve_obj_by_qualname → IOProcessor.__init__. + """ qualname = f"{DummyIOProcessor.__module__}.{DummyIOProcessor.__qualname__}" ep = MagicMock() ep.name = "my_plugin" diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 2bc463b593cb..a2fed9a2d47d 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -382,7 +382,8 @@ def test_inc_get_quant_method_unquantized_layer_with_model_prefix( layer_factory, ) -> None: """``model.``-prefixed extra_config entries trigger early exit for both - LinearBase and ParallelLMHead layers.""" + LinearBase and ParallelLMHead layers. + """ config = make_config(extra_config={"model.layer": {"bits": 16}}) layer = layer_factory() @@ -395,7 +396,8 @@ def test_inc_get_quant_method_unquantized_routed_experts_with_model_prefix( monkeypatch, ) -> None: """``model.``-prefixed extra_config entries trigger early exit for - RoutedExperts layers, returning UnquantizedFusedMoEMethod.""" + RoutedExperts layers, returning UnquantizedFusedMoEMethod. + """ class DummyUnquantizedFusedMoEMethod: def __init__(self, moe_config) -> None: @@ -418,7 +420,8 @@ def __init__(self, moe_config) -> None: def test_inc_get_quant_method_unknown_layer_with_model_prefix_returns_none() -> None: """``model.``-prefixed extra_config entries return None for unhandled - layer types (non-Linear, non-MoE).""" + layer types (non-Linear, non-MoE). + """ config = make_config(extra_config={"model.unknown": {"bits": 16}}) method = config.get_quant_method(DummyLayer(), "unknown") @@ -733,7 +736,8 @@ class DummyMoeConfig: @pytest.mark.parametrize("bits", [4, 8]) def test_wna16_cuda_high_bit_skips_humming(monkeypatch, bits) -> None: """4/8-bit int stays on the Marlin/GPTQ/AWQ path even on CUDA so a single - model can mix high-bit Marlin and low-bit humming layers.""" + model can mix high-bit Marlin and low-bit humming layers. + """ called = {"humming": False} monkeypatch.setattr(current_platform, "is_cuda", lambda: True) @@ -1664,7 +1668,8 @@ def test_inc_get_quant_method_unquantized_moe_returns_unquantized( monkeypatch, ) -> None: """Early-exit returns UnquantizedFusedMoEMethod for FusedMoE layers - when extra_config has bits >= 16.""" + when extra_config has bits >= 16. + """ config = make_config(extra_config={"layer": {"bits": 16}}) layer = object.__new__(RoutedExperts) layer.moe_config = None # UnquantizedFusedMoEMethod accepts moe_config diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index ef80287e5ff2..126ac92eefe2 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -574,8 +574,7 @@ def test_compressed_tensors_fp8_block_enabled( reason="This test is not for non-CUDA platforms", ) def test_compressed_tensors_moe_ignore_with_model(vllm_runner): - """ - Integration test for MoE layer ignore functionality with a real model. + """Integration test for MoE layer ignore functionality with a real model. This test would verify that when loading a compressed-tensors quantized MoE model where some MoE layers are in the ignore list, those layers @@ -587,7 +586,6 @@ def test_compressed_tensors_moe_ignore_with_model(vllm_runner): - Config with ignore list containing specific MoE layers - Multiple MoE layers where some are quantized and some are not """ - # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3 model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" # CT 12.2 diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 65e119b083fd..44a508ed57ce 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -213,7 +213,8 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales( ): """The VL wrapper must map ``.scale`` keys to the parameter the linear quant method registers, as the text model does. A hardcoded - ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" + ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints. + """ from vllm.models.deepseek_v41.nvidia import model as model_module from vllm.models.deepseek_v41.nvidia import vl_model as vl_module diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index 67728a2da8aa..3adcb020b654 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950). +"""Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950). These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms. """ @@ -161,7 +160,8 @@ def is_supported_config(*args, **kwargs): @pytest.fixture def mxfp4_oracle_config(): """Stub the config the oracle reads (``model_config.quantization_config``) - so backend dispatch resolves without a real model / user override.""" + so backend dispatch resolves without a real model / user override. + """ from unittest.mock import patch with patch( diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py index dbec989df3f9..bf508dda8992 100644 --- a/tests/quantization/test_humming_mxfp4_block_fp8.py +++ b/tests/quantization/test_humming_mxfp4_block_fp8.py @@ -85,7 +85,8 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): (FP8 dtype, [1, 128] block, not per-act-token). This is what makes the DeepEP prepare/finalize step quantize activations to block FP8 *before* the all-to-all dispatch instead of deferring to Humming -- and what makes - HummingExpertsBase.expects_unquantized_inputs return False.""" + HummingExpertsBase.expects_unquantized_inputs return False. + """ from vllm.model_executor.layers.quantization.utils.humming_utils import ( make_humming_moe_quant_config, ) @@ -107,7 +108,8 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): def test_default_activation_quant_config_defers_to_humming(): """Without a block activation shape the config stays per-token (the deferred path): Humming quantizes internally, preserving pre-existing behavior for - every non-block-FP8 scheme.""" + every non-block-FP8 scheme. + """ from vllm.model_executor.layers.quantization.utils.humming_utils import ( make_humming_moe_quant_config, ) diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py index 2eb6fa92a547..048b0e645abb 100644 --- a/tests/quantization/test_int8_moe_oracle.py +++ b/tests/quantization/test_int8_moe_oracle.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for INT8 (W8A8) fused-MoE oracle backend selection. +"""Tests for INT8 (W8A8) fused-MoE oracle backend selection. These exercise ``select_int8_moe_backend`` only (no kernels are launched), so they run on any platform where the Triton INT8 MoE kernel is available — CUDA @@ -71,7 +70,8 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: ) def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor - weights) select the Triton backend.""" + weights) select the Triton backend. + """ config = _make_int8_moe_config() backend, experts_cls = select_int8_moe_backend( config, weight_key=weight_key, activation_key=activation_key diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 60b207c1840f..8e590d7c3f9e 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -587,7 +587,8 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i def test_modelopt_nvfp4_config_dispatches_w4a4_method(): """``quant_method="NVFP4"`` (W4A4) resolves to a ``(kNvfp4Static, kNvfp4Dynamic)`` QuantSpec under the generic - ``ModelOptLinearMethod``.""" + ``ModelOptLinearMethod``. + """ from vllm.model_executor.layers.linear import LinearBase config = ModelOptNvFp4Config( @@ -636,7 +637,8 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): """The bespoke-method escape hatch: a format registered in ``LINEAR_METHOD_BUILDERS`` routes that algo to its own method instead of the generic ``ModelOptLinearMethod``. This is how a format that cannot be a - ``(weight, activation)`` key pair plugs into dispatch.""" + ``(weight, activation)`` key pair plugs into dispatch. + """ from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import modelopt as m @@ -667,7 +669,8 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``: ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The - generic method routes this through ``select_linear_kernel``.""" + generic method routes this through ``select_linear_kernel``. + """ from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.quantization.modelopt import ( RuntimeDtypes, @@ -938,7 +941,8 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch): def test_modelopt_fp8_pb_wo_rejects_non_128_input(): """Input width must still be a multiple of 128 (same as #53132, which only pads the output). A partial input block is refused loudly rather than - silently loading wrong scales.""" + silently loading wrong scales. + """ from vllm.model_executor.layers.quantization import modelopt as mo from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Static128BlockSym, diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index fbd5cbc56392..846b9814786a 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -449,7 +449,6 @@ def test_online_shorthand_selects_checkpoint_or_online_config( tmp_path, quantization, checkpoint_quantization_config, is_config_file ) -> None: """Online shorthands defer to checkpoint metadata when available.""" - if is_config_file: (tmp_path / "hf_quant_config.json").write_text( '{"quant_method": "mxfp8", "ignored_layers": []}' @@ -609,14 +608,12 @@ def test_online_quantization( dist_init, workspace_init, ) -> None: - """ - Tests that online quantization frontend configuration works - + """Tests that online quantization frontend configuration works - selecting quant schemes, overriding quant schemes by type, ignoring layers. Does not test performance, peak memory usage, etc. """ - # TODO: Relax this condition once there is a native MXFP4_MXFP4 # linear/moe backend supported on cuda. if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()): diff --git a/tests/quantization/test_online_mxfp4.py b/tests/quantization/test_online_mxfp4.py index c435790bce39..8d29b4602cd4 100644 --- a/tests/quantization/test_online_mxfp4.py +++ b/tests/quantization/test_online_mxfp4.py @@ -135,8 +135,7 @@ def _skip_reason_if_unavailable(backend: str, dtype: torch.dtype) -> str | None: @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("backend", ["triton", "aiter", "xpu", "quark"]) def test_mxfp4_quantization_correctness(backend: str, dtype: torch.dtype): - """ - Tests that the different implementations of mxfp4_quantize + """Tests that the different implementations of mxfp4_quantize in mxfp4_utils.py all match. """ skip_reason = _skip_reason_if_unavailable(backend, dtype) @@ -290,8 +289,7 @@ def test_online_mxfp4_moe_matches_quark( dist_init, monkeypatch, ): - """ - Ensures `Mxfp4OnlineMoEMethod` (online quantization) + """Ensures `Mxfp4OnlineMoEMethod` (online quantization) and `QuarkOCP_MX_MoEMethod` (AMD Quark checkpoints) produce the same weights, with same MOE backend used. @@ -413,7 +411,8 @@ def scatter_into_padded( ) -> tuple[torch.Tensor, torch.Tensor]: """Place the source weights into buffers of the padded size, the way the weight loader writes only the unpadded slice of a larger - allocation.""" + allocation. + """ w13 = torch.full( (num_experts, 2 * intermediate_size, hidden_size), padding_value, @@ -495,8 +494,7 @@ def scatter_into_padded( def test_online_mxfp4_dense_matches_quark( linear_backend: str, default_vllm_config, dist_init, monkeypatch ): - """ - Ensures `Mxfp4OnlineLinearMethod` (online quantization) + """Ensures `Mxfp4OnlineLinearMethod` (online quantization) and `QuarkOCP_MX` (AMD Quark checkpoints) produce the same weights, with same linear backend used. """ diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 4d6193211829..419b079d21b4 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -490,7 +490,8 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig): ) def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str): """Per-token-head should set _k_scale=1.0, _v_scale=1.0 - and delete checkpoint attrs.""" + and delete checkpoint attrs. + """ from vllm.model_executor.layers.quantization.kv_cache import ( BaseKVCacheMethod, ) diff --git a/tests/quantization/test_register_quantization_config.py b/tests/quantization/test_register_quantization_config.py index 5529a65b778f..227c604ccf87 100644 --- a/tests/quantization/test_register_quantization_config.py +++ b/tests/quantization/test_register_quantization_config.py @@ -48,7 +48,6 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: """Perform fake quantization before the linear layer.""" - # Calculate the scales dynamically max_val = torch.amax(x, dim=(0, -1), keepdims=True) min_val = torch.amin(x, dim=(0, -1), keepdims=True) @@ -108,7 +107,6 @@ def get_quant_method( def test_register_quantization_config(caplog_vllm): """Test register custom quantization config.""" - # The quantization method `custom_quant` should be registered. assert get_quantization_config("custom_quant") == CustomQuantConfig diff --git a/tests/reasoning/test_base_thinking_reasoning_parser.py b/tests/reasoning/test_base_thinking_reasoning_parser.py index f4d74ceeec0a..fb506860621c 100644 --- a/tests/reasoning/test_base_thinking_reasoning_parser.py +++ b/tests/reasoning/test_base_thinking_reasoning_parser.py @@ -51,8 +51,7 @@ def test_tokenizer(): class TestBaseThinkingReasoningParserInit: - """ - Test initialization and basic properties of + """Test initialization and basic properties of BaseThinkingReasoningParser. """ @@ -274,8 +273,7 @@ class TestBaseThinkingReasoningParserStreaming: @pytest.mark.parametrize("streaming", [True, False]) def test_simple_reasoning_extraction(self, test_tokenizer, streaming): - """ - Test basic reasoning extraction in both + """Test basic reasoning extraction in both streaming and non-streaming modes. """ parser = TestThinkingReasoningParser(test_tokenizer) @@ -369,14 +367,12 @@ def test_streaming_only_end_token(self, test_tokenizer): class TestBaseThinkingReasoningParserMultipleImplementations: - """ - Test that multiple implementations of + """Test that multiple implementations of BaseThinkingReasoningParser work correctly. """ def test_different_token_implementations(self, test_tokenizer): - """ - Test that different implementations + """Test that different implementations with different tokens work independently. """ parser1 = TestThinkingReasoningParser(test_tokenizer) diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py index ea5b374e6e0d..dad70ea8d6f7 100644 --- a/tests/reasoning/test_kimi_k3_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -276,7 +276,8 @@ def last(needle: list[int]) -> int: def test_is_reasoning_end_streaming_only_scans_the_step_window(): """The decode-step check must not re-derive the answer from the whole sequence: an already-closed think block earlier in the sequence is the - engine's business (it latches `reasoning_ended`), not this call's.""" + engine's business (it latches `reasoning_ended`), not this call's. + """ parser = KimiK3ReasoningParser(DummyTokenizer()) prompt = [*CLOSE_IDS, *OPEN_IDS, 9, 9, 9] @@ -286,7 +287,8 @@ def test_is_reasoning_end_streaming_only_scans_the_step_window(): def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): """A 3-token marker can be split over decode steps; the check carries the - preceding len(marker)-1 tokens so the final token still completes it.""" + preceding len(marker)-1 tokens so the final token still completes it. + """ parser = KimiK3ReasoningParser(DummyTokenizer()) history = [*OPEN_IDS, 5, *CLOSE_IDS[:2]] @@ -297,7 +299,8 @@ def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): def test_is_reasoning_end_streaming_reopened_block_is_not_ended(): """Kimi K3 may close and immediately reopen the think channel inside one window (speculative decoding lands several tokens per step). The newest - marker wins, matching is_reasoning_end.""" + marker wins, matching is_reasoning_end. + """ parser = KimiK3ReasoningParser(DummyTokenizer()) window = [*CLOSE_IDS, *OPEN_IDS] @@ -324,7 +327,8 @@ def test_is_reasoning_end_streaming_thinking_disabled(): def test_reasoning_end_matches_reference_over_marker_dense_sequences(seed): """Both the full-sequence check and the per-step check must agree with the two-scan reference, on sequences where the markers (which share the suffix - [2, 3]) collide and overlap constantly.""" + [2, 3]) collide and overlap constantly. + """ import random rnd = random.Random(seed) diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index bbd4dc5c5f2e..bcaff9b94421 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -116,7 +116,8 @@ def test_prompt_embeds_keys_registered(): def test_ensure_placeholder_token_is_single_token_and_idempotent(tokenizer): """Ensure the placeholder token is a single token and that multiple calls to - "ensure" are idempotent, across all tokenizer families.""" + "ensure" are idempotent, across all tokenizer families. + """ tid1 = _ensure_prompt_embeds_placeholder_token(tokenizer) tid2 = _ensure_prompt_embeds_placeholder_token(tokenizer) assert tid1 == tid2 @@ -469,7 +470,8 @@ def test_build_mixed_prompt_embeds(stream): @pytest.mark.parametrize("role", ["user", "system"]) async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): """Full renderer pipeline: parse -> chat template -> expand -> locate - -> build mixed prompt, across tokenizers, roles, and sync/async.""" + -> build mixed prompt, across tokenizers, roles, and sync/async. + """ tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE tid = _ensure_prompt_embeds_placeholder_token(tokenizer) @@ -524,7 +526,8 @@ async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): @pytest.mark.asyncio async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn): """Full pipeline with prompt_embeds spread across system + user messages, - verifying ordering and positioning in the final token stream.""" + verifying ordering and positioning in the final token stream. + """ tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE tid = _ensure_prompt_embeds_placeholder_token(tokenizer) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 6ae8c4d5510a..3094eb6b1178 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -506,7 +506,7 @@ def test_explicit_side_text_pretokenization_guard(self): class TestRenderEmbedPrompt: def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes: - """Helper to create base64-encoded tensor bytes""" + """Helper to create base64-encoded tensor bytes.""" buffer = io.BytesIO() torch.save(tensor, buffer) buffer.seek(0) diff --git a/tests/renderers/test_gemma4_chat_template.py b/tests/renderers/test_gemma4_chat_template.py index 2c1312a84c6c..c244b2672312 100644 --- a/tests/renderers/test_gemma4_chat_template.py +++ b/tests/renderers/test_gemma4_chat_template.py @@ -33,7 +33,8 @@ def _render(template, messages, **kwargs): class TestGemma4ChatTemplate: def test_basic_multiturn_thinking_disabled(self, gemma4_template): """With enable_thinking=False (default), generation prompt ends with - an empty thought channel to suppress thinking.""" + an empty thought channel to suppress thinking. + """ messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, @@ -49,7 +50,8 @@ def test_basic_multiturn_thinking_disabled(self, gemma4_template): def test_basic_multiturn_thinking_enabled(self, gemma4_template): """With enable_thinking=True, generation prompt ends with model - turn opener (no thought suppression).""" + turn opener (no thought suppression). + """ messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, @@ -144,7 +146,8 @@ def test_tool_calls_in_assistant(self, gemma4_template): def test_tool_responses_openai_style(self, gemma4_template): """role='tool' messages are formatted as <|tool_response> blocks - with content dumped as-is.""" + with content dumped as-is. + """ messages = [ {"role": "user", "content": "Weather?"}, { @@ -202,7 +205,8 @@ def test_tool_responses_legacy_style(self, gemma4_template): def test_generation_prompt_not_after_tool_response(self, gemma4_template): """add_generation_prompt=True should NOT add <|turn>model when the - last message type was tool_response (the model turn continues).""" + last message type was tool_response (the model turn continues). + """ messages = [ {"role": "user", "content": "Weather?"}, { @@ -228,8 +232,9 @@ def test_generation_prompt_not_after_tool_response(self, gemma4_template): assert not result.strip().endswith("<|turn>model\n") def test_reasoning_in_tool_chains(self, gemma4_template): - """reasoning field on assistant with tool_calls after last user - message emits <|channel>thought\\n....""" + """Reasoning field on assistant with tool_calls after last user + message emits <|channel>thought\\n.... + """ messages = [ {"role": "user", "content": "Calculate something"}, { @@ -252,7 +257,7 @@ def test_reasoning_in_tool_chains(self, gemma4_template): assert "" in result def test_reasoning_not_before_last_user(self, gemma4_template): - """reasoning on assistant BEFORE the last user message is dropped.""" + """Reasoning on assistant BEFORE the last user message is dropped.""" messages = [ {"role": "user", "content": "First"}, { @@ -275,7 +280,8 @@ def test_reasoning_not_before_last_user(self, gemma4_template): def test_strip_thinking_in_model_content(self, gemma4_template): """<|channel>... in model content is stripped by the - strip_thinking macro.""" + strip_thinking macro. + """ messages = [ {"role": "user", "content": "Hi"}, { @@ -289,7 +295,8 @@ def test_strip_thinking_in_model_content(self, gemma4_template): def test_multi_turn_tool_chain(self, gemma4_template): """assistant->tool->assistant->tool produces exactly one - <|turn>model (later assistants continue the same turn).""" + <|turn>model (later assistants continue the same turn). + """ messages = [ {"role": "user", "content": "Do two things"}, { @@ -346,7 +353,8 @@ def test_format_argument_types(self, gemma4_template): def test_tool_response_with_multimodal_content(self, gemma4_template): """Multimodal placeholders in tool messages are emitted after the - tool_response block.""" + tool_response block. + """ messages = [ {"role": "user", "content": "Download the image and describe it."}, { @@ -380,7 +388,8 @@ def test_tool_response_with_multimodal_content(self, gemma4_template): def test_tool_response_with_all_modalities(self, gemma4_template): """All multimodal types (image, audio, video) in a single tool - response are rendered.""" + response are rendered. + """ messages = [ {"role": "user", "content": "Process media"}, { diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index 736478fbfafe..624108a20689 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -123,7 +123,7 @@ def test_no_load_chat_template_literallike(): ) @pytest.mark.parametrize("use_tools", [True, False]) def test_resolve_chat_template(sample_json_schema, model, use_tools): - """checks that chat_template is a dict type for HF models.""" + """Checks that chat_template is a dict type for HF models.""" model_info = HF_EXAMPLE_MODELS.find_hf_info(model) model_info.check_available_online(on_fail="skip") @@ -196,7 +196,7 @@ def test_resolve_chat_template(sample_json_schema, model, use_tools): ], ) def test_resolve_chat_template_kwargs(sample_json_schema, model, expected_kwargs): - """checks that chat_template is a dict type for HF models.""" + """Checks that chat_template is a dict type for HF models.""" model_info = HF_EXAMPLE_MODELS.find_hf_info(model) model_info.check_available_online(on_fail="skip") @@ -308,7 +308,8 @@ def apply_chat_template(self, conversation, **kwargs): def test_resolve_chat_template_resolves_name(): """When chat_template is a name, resolve_chat_template should return - the actual Jinja content so that kwargs detection works correctly.""" + the actual Jinja content so that kwargs detection works correctly. + """ from unittest.mock import MagicMock jinja_content = "{{ messages }}{% if tools %}{{ tools }}{% endif %}" @@ -330,7 +331,8 @@ def test_resolve_chat_template_resolves_name(): def test_resolve_chat_template_kwargs_with_template_name(): """Ensures template kwargs are not silently dropped when chat_template - was originally a template name that has been resolved to Jinja content.""" + was originally a template name that has been resolved to Jinja content. + """ from unittest.mock import MagicMock jinja_content = ( @@ -932,7 +934,8 @@ def test_developer_only_no_prior_system(self, model_config, tokenizer): class TestApplyChatTemplateEffortTolerant: """Chat templates that reject unsupported reasoning_effort values should - surface a 400-style client error instead of crashing the request.""" + surface a 400-style client error instead of crashing the request. + """ @pytest.fixture def model_config(self): diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py index 8987782d71bc..992177a3a633 100644 --- a/tests/renderers/test_process_multi_modal_uuids.py +++ b/tests/renderers/test_process_multi_modal_uuids.py @@ -49,7 +49,8 @@ def _build_text_only_renderer() -> HfRenderer: def test_text_only_model_mm_data_maps_to_bad_request(): """Sending multimodal data to a text-only model is a client mistake, so it - must surface as a ValueError and reach the client as HTTP 400, not 500.""" + must surface as a ValueError and reach the client as HTTP 400, not 500. + """ renderer = _build_text_only_renderer() with pytest.raises(ValueError, match="text-only") as exc_info: diff --git a/tests/renderers/test_prompt_embeds_payload_validation.py b/tests/renderers/test_prompt_embeds_payload_validation.py index 1e425805e1c0..fec937a4b723 100644 --- a/tests/renderers/test_prompt_embeds_payload_validation.py +++ b/tests/renderers/test_prompt_embeds_payload_validation.py @@ -89,7 +89,7 @@ def test_unreadable_prompt_embeds_payload_is_a_client_error(model_config, payloa @pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS) def test_the_error_body_stays_bounded(model_config, payload): - """torch's reason is built from the caller's bytes, so it must be truncated. + """Torch's reason is built from the caller's bytes, so it must be truncated. An unpickler that reports the offending global, for instance, quotes a name the caller chose; without a bound a large one would be reflected whole. @@ -127,7 +127,8 @@ def test_a_long_reason_is_truncated_and_says_how_much_was_left_out(model_config) def test_a_short_reason_is_passed_through_whole(model_config): """A reason that fits is not annotated -- the count only appears when it is - telling the reader something.""" + telling the reader something. + """ reason = "invalid magic number; corrupt file?" assert len(reason) <= _MAX_EMBED_ERROR_REASON_CHARS diff --git a/tests/renderers/test_sparse_tensor_concurrent_race.py b/tests/renderers/test_sparse_tensor_concurrent_race.py index a1683eee0bff..c627dc31cbad 100644 --- a/tests/renderers/test_sparse_tensor_concurrent_race.py +++ b/tests/renderers/test_sparse_tensor_concurrent_race.py @@ -218,7 +218,8 @@ class TestCrossLoaderLockSharing: def test_prompt_and_image_share_lock(self, model_config): """Prompt and image loaders cannot run their sparse guards - concurrently.""" + concurrently. + """ valid_prompt = _encode_tensor(_create_valid_dense_tensor()) valid_image = _encode_tensor(_create_valid_dense_tensor(hidden_size=10)) diff --git a/tests/renderers/test_sparse_tensor_validation.py b/tests/renderers/test_sparse_tensor_validation.py index e49ac560dc71..d90ad1e09a18 100644 --- a/tests/renderers/test_sparse_tensor_validation.py +++ b/tests/renderers/test_sparse_tensor_validation.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests verify that malicious sparse tensors are rejected before they can trigger +"""Tests verify that malicious sparse tensors are rejected before they can trigger out-of-bounds memory writes, or an unbounded allocation, during to_dense() operations. @@ -51,8 +50,7 @@ def _encode_tensor(tensor: torch.Tensor) -> bytes: def _create_malicious_sparse_tensor() -> torch.Tensor: - """ - Create a malicious sparse COO tensor with out-of-bounds indices. + """Create a malicious sparse COO tensor with out-of-bounds indices. This tensor has indices that point beyond the declared shape, which would cause an out-of-bounds write when converted to dense format without @@ -159,7 +157,8 @@ def test_negative_indices_rejected(self, model_config): def test_hidden_size_mismatch_rejected(self, model_config): """Tensors whose trailing dim doesn't match the model's hidden_size - must be rejected at parse time.""" + must be rejected at parse time. + """ # opt-125m has hidden_size=768, passing 512 triggers the check. wrong_hidden = torch.randn(10, 512, dtype=torch.float32) encoded = _encode_tensor(wrong_hidden) @@ -170,7 +169,8 @@ def test_hidden_size_mismatch_rejected(self, model_config): def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): """Tensors whose dtype doesn't match the model's dtype but are still floating-point are cast, since API clients generally can't know the - server's `--dtype` setting ahead of time.""" + server's `--dtype` setting ahead of time. + """ # Fixture pins model dtype to float32, upload a bfloat16 tensor. mismatched_float = torch.randn(10, 768, dtype=torch.bfloat16) encoded = _encode_tensor(mismatched_float) @@ -183,7 +183,8 @@ def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): def test_non_float_dtype_rejected(self, model_config): """Non-floating-point dtypes cannot be safely cast for embeddings (e.g. integer tensors almost certainly indicate caller confusion), - so they are rejected at parse time.""" + so they are rejected at parse time. + """ non_float = torch.randint(0, 100, (10, 768), dtype=torch.int32) encoded = _encode_tensor(non_float) @@ -192,7 +193,8 @@ def test_non_float_dtype_rejected(self, model_config): def test_non_2d_tensor_rejected(self, model_config): """Tensors that aren't 2D (even after squeezing a leading dim) - must be rejected with a clear error.""" + must be rejected with a clear error. + """ # A 1D tensor cannot be interpreted as (num_tokens, hidden_size). bad = torch.randn(768, dtype=torch.float32) encoded = _encode_tensor(bad) @@ -202,7 +204,8 @@ def test_non_2d_tensor_rejected(self, model_config): def test_non_tensor_payload_rejected(self, model_config): """Deserializing to a non-Tensor object must raise a clear error - instead of propagating an AssertionError.""" + instead of propagating an AssertionError. + """ # `torch.save` will serialize a plain dict; `weights_only=True` allows # loading built-in containers, so this exercises the isinstance check. buffer = io.BytesIO() @@ -266,7 +269,7 @@ def test_load_bytes_validates(self): io_handler.load_bytes(buffer.read()) def test_valid_numpy_tensor_accepted(self): - """numpy .npy format should load and return correct tensor.""" + """Numpy .npy format should load and return correct tensor.""" io_handler = ImageEmbeddingMediaIO() arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) @@ -281,8 +284,7 @@ def test_valid_numpy_tensor_accepted(self): assert torch.allclose(result, torch.from_numpy(arr)) def test_numpy_int32_tensor_accepted(self): - """numpy int32 arrays should round-trip correctly.""" - + """Numpy int32 arrays should round-trip correctly.""" io_handler = ImageEmbeddingMediaIO() arr = np.arange(280, dtype=np.int32) @@ -296,8 +298,7 @@ def test_numpy_int32_tensor_accepted(self): assert (result == torch.from_numpy(arr)).all() def test_load_file_numpy_tensor_accepted(self, tmp_path): - """numpy .npy files should load correctly via load_file.""" - + """Numpy .npy files should load correctly via load_file.""" io_handler = ImageEmbeddingMediaIO() arr = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32) @@ -364,13 +365,10 @@ def test_load_bytes_validates(self): class TestSparseTensorValidationIntegration: - """ - These tests verify the complete attack chain is blocked at all entry points. - """ + """These tests verify the complete attack chain is blocked at all entry points.""" def test_attack_scenario_completions_api(self, model_config): - """ - Simulate a complete attack through the Completions API. + """Simulate a complete attack through the Completions API. Attack scenario: 1. Attacker crafts malicious sparse tensor @@ -386,8 +384,7 @@ def test_attack_scenario_completions_api(self, model_config): safe_load_prompt_embeds(model_config, attack_payload) def test_attack_scenario_chat_api_image(self): - """ - Simulate attack through Chat API with image_embeds. + """Simulate attack through Chat API with image_embeds. Verifies the image embeddings path is protected. """ @@ -398,8 +395,7 @@ def test_attack_scenario_chat_api_image(self): io_handler.load_base64("", attack_payload.decode("utf-8")) def test_attack_scenario_chat_api_audio(self): - """ - Simulate attack through Chat API with audio_embeds. + """Simulate attack through Chat API with audio_embeds. Verifies the audio embeddings path is protected. """ @@ -467,7 +463,8 @@ def test_oversized_sparse_tensor_rejected(self, small_decode_limit): def test_default_limit_would_reject_the_real_payload(self): """The shipped default has to bound the actual attack, not just the shrunken shape used above. Checked arithmetically so no test ever asks - the allocator for 4 TiB.""" + the allocator for 4 TiB. + """ dense_bytes = BOMB_SHAPE[0] * BOMB_SHAPE[1] * 4 assert dense_bytes > envs.VLLM_MAX_EMBED_DECODE_BYTES > 0 @@ -525,7 +522,8 @@ def test_media_io_rejects_oversized_file( @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) def test_media_io_rejects_non_tensor_payload(self, io_cls, parameter): """These three had no `isinstance` check, so a non-tensor payload - surfaced as an `AttributeError` (500) rather than a 4xx.""" + surfaced as an `AttributeError` (500) rather than a 4xx. + """ buffer = io.BytesIO() torch.save({"not": "a tensor"}, buffer) encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") diff --git a/tests/renderers/test_token_offsets.py b/tests/renderers/test_token_offsets.py index fc9192bb43e7..973ea175cbb9 100644 --- a/tests/renderers/test_token_offsets.py +++ b/tests/renderers/test_token_offsets.py @@ -26,7 +26,8 @@ def fast_tokenizer(): def _make_base_renderer_with(tokenizer): """Build a minimal BaseRenderer subclass that exposes the tokenizer so we can call ``_tokenize_prompt`` directly. BaseRenderer is abstract because of - ``render_messages``; we just need a stub.""" + ``render_messages``; we just need a stub. + """ from vllm.renderers.base import BaseRenderer class _StubRenderer(BaseRenderer): @@ -91,7 +92,8 @@ def test_fast_tokenizer_with_flag_returns_offsets(self, fast_tokenizer): def test_base_renderer_without_override_yields_no_offsets(self, fast_tokenizer): """A renderer that does not override ``_can_produce_offsets`` never emits offsets, even with a fast tokenizer and the flag set. This locks - in the base-default-False / subclass-override design.""" + in the base-default-False / subclass-override design. + """ from vllm.renderers.base import BaseRenderer class _BareRenderer(BaseRenderer): @@ -125,7 +127,8 @@ def test_default_flag_no_offsets(self, fast_tokenizer): def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): """Force is_fast=False to simulate a Slow tokenizer: the flag is set - but offsets must not be returned because it cannot produce them.""" + but offsets must not be returned because it cannot produce them. + """ from unittest.mock import PropertyMock, patch renderer = _make_base_renderer_with(fast_tokenizer) @@ -144,7 +147,8 @@ def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): @pytest.mark.parametrize("mm_key", ["multi_modal_data", "multi_modal_uuids"]) def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): """Offsets index the text prompt, which is meaningless once multimodal - data is interleaved, so they are suppressed when MM inputs are present.""" + data is interleaved, so they are suppressed when MM inputs are present. + """ renderer = _make_base_renderer_with(fast_tokenizer) params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) prompt = {"prompt": "Hello.", mm_key: {"image": ["x"]}} @@ -156,7 +160,8 @@ def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): @pytest.mark.asyncio async def test_tokenize_prompt_async_returns_offsets(self, fast_tokenizer): """The async path offloads the sync tokenizer; it must yield the same - offsets as the sync path.""" + offsets as the sync path. + """ renderer = _make_base_renderer_with(fast_tokenizer) params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) diff --git a/tests/spec_decode/test_custom_proposer.py b/tests/spec_decode/test_custom_proposer.py index 1e42bde8104d..2e9282af744b 100755 --- a/tests/spec_decode/test_custom_proposer.py +++ b/tests/spec_decode/test_custom_proposer.py @@ -29,6 +29,7 @@ def __init__(self, vllm_config: VllmConfig): Args: vllm_config: vLLM configuration containing model and speculative settings. + """ assert vllm_config.speculative_config is not None self.num_speculative_tokens = ( @@ -57,6 +58,7 @@ def propose( Returns: List of draft token sequences for each request. + """ # Cross-process flag to prove this method was executed with open("proposer_called.flag", "w") as f: diff --git a/tests/test_access_log_filter.py b/tests/test_access_log_filter.py index a28771ce59ab..4f03892e7753 100644 --- a/tests/test_access_log_filter.py +++ b/tests/test_access_log_filter.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the UvicornAccessLogFilter class. -""" +"""Tests for the UvicornAccessLogFilter class.""" import logging diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index dd0c6261d349..572bc76caa6b 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -38,9 +38,10 @@ def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path): def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( tmp_path: Path, ): - """torch's autodetection can emit a bogus arch=compute_* half (e.g. + """Torch's autodetection can emit a bogus arch=compute_* half (e.g. capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS - target must win, while PTX-only entries keep the virtual arch.""" + target must win, while PTX-only entries keep the virtual arch. + """ repo_root = Path(__file__).parents[1] script = tmp_path / "test_extract_archs.cmake" script.write_text( diff --git a/tests/test_config.py b/tests/test_config.py index f75592f089e2..4b37b4f56e03 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -441,7 +441,8 @@ def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled - MRV1 paths and off breakable cudagraphs by default.""" + MRV1 paths and off breakable cudagraphs by default. + """ from vllm.config.vllm import ( ROCM_DEFAULT_MRV1_ARCHITECTURES, default_breakable_cudagraph_architectures, @@ -656,7 +657,8 @@ def test_v2_model_runner_supports_extract_hidden_states(): def test_dflash2_draft_forces_v2_model_runner(): """A DFlash2 draft must reach the V2 speculator, the only one that runs its - candidate selector; on V1 it would draft as DFlash1 without raising.""" + candidate selector; on V1 it would draft as DFlash1 without raising. + """ def config(method, architectures): return SimpleNamespace( @@ -769,7 +771,8 @@ def test_late_piecewise_restrictions_without_compilation(monkeypatch, engine_kwa def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): """Cudagraph memory profiling uses a minimal KV cache, so the Mamba - block-count guard must only fire for the real cache sizing.""" + block-count guard must only fire for the real cache sizing. + """ kv_cache_config = SimpleNamespace(has_mamba_layers=True, num_blocks=4) compilation_config = CompilationConfig( @@ -1094,7 +1097,8 @@ def test_v1_model_runner_rejects_v2_only_features(): def test_batch_sharded_sampling_rejects_return_sampling_mask(): """The batch-sharded gather drops sampling masks, so the combination must - fail loudly instead of returning ``sampling_mask=None``.""" + fail loudly instead of returning ``sampling_mask=None``. + """ config = SimpleNamespace( parallel_config=SimpleNamespace( enable_batch_sharded_sampling=True, tensor_parallel_size=2 @@ -2173,7 +2177,7 @@ def test_get_and_verify_max_len_yarn_is_already_scaled( class MockConfig: - """Simple mock object for testing maybe_pull_model_tokenizer_for_runai""" + """Simple mock object for testing maybe_pull_model_tokenizer_for_runai.""" def __init__(self, model: str, tokenizer: str): self.model = model @@ -2191,7 +2195,8 @@ def __init__(self, model: str, tokenizer: str): @patch("vllm.transformers_utils.runai_utils.ObjectStorageModel.pull_files") def test_s3_url_model_tokenizer_paths(mock_pull_files, s3_url): """Test that S3 URLs create deterministic local directories for model and - tokenizer.""" + tokenizer. + """ # Mock pull_files to avoid actually downloading files during tests mock_pull_files.return_value = None @@ -2289,7 +2294,8 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files): @patch("vllm.transformers_utils.runai_utils.ObjectStorageModel.pull_files") def test_s3_url_different_model_and_tokenizer(mock_pull_files): """Test that when model and tokenizer are different cloud URIs, - pull_files receives the correct URI for each.""" + pull_files receives the correct URI for each. + """ mock_pull_files.return_value = None model_url = "s3://bucket/model/" @@ -2651,7 +2657,6 @@ def test_validate_mamba_align_subblock_prefill(): ) def test_vllm_config_defaults(model_id, compilation_config, optimization_level): """Test that optimization-level defaults are correctly applied.""" - model_config = None if model_id is not None: model_config = ModelConfig(model_id) @@ -2835,8 +2840,8 @@ def test_vllm_config_explicit_overrides(): def test_fusion_pass_op_priority(): """This test checks that custom op enablement & IR op priority - correctly control default fusions""" - + correctly control default fusions + """ # Default config, O2, rms_norm+quant fusion disabled cfg1 = VllmConfig() assert not cfg1.compilation_config.pass_config.fuse_norm_quant @@ -2920,7 +2925,8 @@ def test_needs_dp_coordination( def test_fault_tolerance_requires_single_api_server(): """Fault tolerance assumes one AsyncMPClient manages all engines, so it - is incompatible with API server scale-out (_api_process_count > 1).""" + is incompatible with API server scale-out (_api_process_count > 1). + """ with pytest.raises(ValueError, match="single API server"): ParallelConfig(enable_fault_tolerance=True, _api_process_count=2) @@ -2930,7 +2936,8 @@ def test_fault_tolerance_requires_single_api_server(): def test_renderer_num_workers_with_mm_cache(): """Disallow renderer_num_workers > 1 with the mm processor cache only for - pooling models, whose preprocessing runs on the renderer workers.""" + pooling models, whose preprocessing runs on the renderer workers. + """ mm_model = "Qwen/Qwen2-VL-2B-Instruct" # Should raise: pooling + multi-worker + cache enabled (default cache_gb=4) @@ -3256,7 +3263,8 @@ def test_watermarking_forces_model_runner_v2(monkeypatch): def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls): """Regression test: MTP + runai_streamer should use model_weights (original S3 URL) for the draft model, not model (local cache dir set by - pull_runai_model_from_obj_storage).""" + pull_runai_model_from_obj_storage). + """ from unittest.mock import MagicMock s3_url = "s3://my-bucket/Qwen3-35B-A3B-FP8" @@ -3441,7 +3449,7 @@ def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend): def test_ir_op_priority_str(): - """Test that passing a comma-delimited string works""" + """Test that passing a comma-delimited string works.""" from vllm.config.kernel import IrOpPriorityConfig priority_config = IrOpPriorityConfig(rms_norm="vllm_c") diff --git a/tests/test_logger.py b/tests/test_logger.py index 2ff100151b29..f051fe205b8c 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -51,7 +51,8 @@ def test_trace_function_call(): def test_default_vllm_root_logger_configuration(monkeypatch): """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default - behavior is activated.""" + behavior is activated. + """ monkeypatch.setenv("VLLM_LOGGING_COLOR", "0") _configure_vllm_root_logger() @@ -75,7 +76,8 @@ def test_default_vllm_root_logger_configuration(monkeypatch): def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypatch): """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default - behavior is activated.""" + behavior is activated. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False) @@ -105,7 +107,8 @@ def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypa def test_logger_configuring_can_be_disabled(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur.""" + configuration occur. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0") monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False) @@ -117,7 +120,8 @@ def test_logger_configuring_can_be_disabled(monkeypatch): def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs.""" + configuration occurs. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") monkeypatch.setenv( "VLLM_LOGGING_CONFIG_PATH", @@ -133,7 +137,8 @@ def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist(monke def test_an_error_is_raised_when_custom_logging_config_is_invalid_json(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs.""" + configuration occurs. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file: @@ -160,7 +165,8 @@ def test_an_error_is_raised_when_custom_logging_config_is_unexpected_json( ): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs.""" + configuration occurs. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file: @@ -176,7 +182,8 @@ def test_an_error_is_raised_when_custom_logging_config_is_unexpected_json( def test_custom_logging_config_is_parsed_and_used_when_provided(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur.""" + configuration occur. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") valid_logging_config = { @@ -200,7 +207,8 @@ def test_custom_logging_config_is_parsed_and_used_when_provided(monkeypatch): def test_custom_logging_config_causes_an_error_if_configure_logging_is_off(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur.""" + configuration occur. + """ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0") valid_logging_config = { diff --git a/tests/test_ray_env.py b/tests/test_ray_env.py index 945b2d80b691..43ec51a8f81b 100644 --- a/tests/test_ray_env.py +++ b/tests/test_ray_env.py @@ -15,7 +15,8 @@ class TestDefaultPrefixes: """Built-in prefixes (VLLM_, LMCACHE_, NCCL_, UCX_, HF_, HUGGING_FACE_) - should be forwarded without any extra configuration.""" + should be forwarded without any extra configuration. + """ @patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False) def test_lmcache_prefix(self): @@ -53,7 +54,8 @@ class TestDefaultExtraVars: def test_pythonhashseed_in_result(self): """PYTHONHASHSEED should always be in the result set (as a name to - copy) regardless of whether it is actually set in os.environ.""" + copy) regardless of whether it is actually set in os.environ. + """ result = get_env_vars_to_copy() assert "PYTHONHASHSEED" in result diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index a55aa57666dd..aac88553e6e9 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -452,7 +452,8 @@ def test_deepseek_v4_image_blocks_become_placeholders(): def test_deepseek_v4_image_sentinel_ids_match_tokenizer(): """The borrowed sentinel ids must line up with the reserved - ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer.""" + ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer. + """ from vllm.models.deepseek_v4.common.mm_preprocess import ( IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_TOKEN_NAMES, diff --git a/tests/tokenizers_/test_hf.py b/tests/tokenizers_/test_hf.py index 61c81302f071..f1dfd046b93e 100644 --- a/tests/tokenizers_/test_hf.py +++ b/tests/tokenizers_/test_hf.py @@ -51,7 +51,8 @@ def _check_consistency(target: TokenizerLike, expected: TokenizerLike): def test_thread_pool_tokenizer_pickle(model_id: str): """Regression test for issue #45433: the thread-pool tokenizer wrapper reconstructs through maybe_make_thread_pool on unpickling, which used to - fall off the end and return None.""" + fall off the end and return None. + """ reference_tokenizer = AutoTokenizer.from_pretrained(model_id) pooled_tokenizer = maybe_make_thread_pool(deepcopy(reference_tokenizer)) diff --git a/tests/tool_parsers/common_tests.py b/tests/tool_parsers/common_tests.py index 925506aa73d4..b3f9ca7907ea 100644 --- a/tests/tool_parsers/common_tests.py +++ b/tests/tool_parsers/common_tests.py @@ -49,6 +49,7 @@ class ToolParserTestConfig: Special flags: allow_empty_or_json_empty_args: True if "" or "{}" both valid for empty args supports_typed_arguments: True if the parser supports typed function arguments + """ # Parser identification @@ -110,6 +111,7 @@ def test_config(self) -> ToolParserTestConfig: def test_mistral_specific_feature(self, tool_parser): # Custom test logic pass + """ @pytest.fixture diff --git a/tests/tool_parsers/test_apertus_tool_parser.py b/tests/tool_parsers/test_apertus_tool_parser.py index 5ab0a2bcc732..f3cf5fcaf4ea 100644 --- a/tests/tool_parsers/test_apertus_tool_parser.py +++ b/tests/tool_parsers/test_apertus_tool_parser.py @@ -341,7 +341,8 @@ def test_streaming_partial_tag_buffering(self, parser, mock_request): def test_mtp_streaming_massive_chunk(self, parser, mock_request): """Simulates MTP predicting text, tool calls, - and trailing text all in a single chunk.""" + and trailing text all in a single chunk. + """ chunks = [ "Sure! " '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>' @@ -374,7 +375,8 @@ def test_mtp_streaming_multiple_tools_burst(self, parser, mock_request): def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request): """Simulates MTP chunks that jump over entire tools - (e.g., from middle of tool 1 to middle of tool 3).""" + (e.g., from middle of tool 1 to middle of tool 3). + """ chunks = [ '<|tools_prefix|>[{"t1": {"a": 1}', '}, {"t2": {"b": 2}}, {"t3": {"c": 3', @@ -393,7 +395,8 @@ def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request): def test_vllm_streaming_character_by_character(self, parser, mock_request): """Simulates worst-case vLLM fragmentation where - chunks arrive character-by-character.""" + chunks arrive character-by-character. + """ text = ( 'Hi <|tools_prefix|>[{"get_weather": ' '{"location": "London"}}]<|tools_suffix|> ' @@ -411,7 +414,8 @@ def test_vllm_streaming_character_by_character(self, parser, mock_request): def test_vllm_streaming_empty_deltas(self, parser, mock_request): """Simulates vLLM stream producing empty string chunks - (e.g., hidden tokens or artifacts).""" + (e.g., hidden tokens or artifacts). + """ chunks = [ "Wait", "", diff --git a/tests/tool_parsers/test_deepseekv32_tool_parser.py b/tests/tool_parsers/test_deepseekv32_tool_parser.py index 40ad6033d3bf..d0d032167a37 100644 --- a/tests/tool_parsers/test_deepseekv32_tool_parser.py +++ b/tests/tool_parsers/test_deepseekv32_tool_parser.py @@ -173,7 +173,8 @@ def test_type_conversion_in_non_streaming(self): def test_string_attr_true_coerced_by_schema(self): """string="true" delivers a string, but the engine's schema-aware - type fixer coerces it to the schema type (integer).""" + type fixer coerces it to the schema type (integer). + """ tool = ChatCompletionToolsParam( function=FunctionDefinition( name="score", @@ -287,7 +288,8 @@ def test_composed_schema_converts_object_and_array_params(self): def test_string_attr_true_coerced_by_composed_schema(self): """string="true" delivers a JSON string, but the engine's schema-aware - type fixer coerces it to the composed schema type (object).""" + type fixer coerces it to the composed schema type (object). + """ tool = ChatCompletionToolsParam( function=FunctionDefinition( name="set_timer", @@ -321,7 +323,8 @@ def test_string_attr_true_coerced_by_composed_schema(self): def test_arguments_wrapper_repaired(self): """A single 'arguments' wrapper parameter must be unwrapped when it - is not part of the tool schema and the inner object matches schema fields.""" + is not part of the tool schema and the inner object matches schema fields. + """ tool = ChatCompletionToolsParam( function=FunctionDefinition( name="get_weather", @@ -492,7 +495,8 @@ def test_multi_typed_null_value(self): def test_null_coerced_back_to_string_by_schema(self): """string="false" with 'null' is json-parsed to None, but the - engine's schema fixer coerces it back to "null" for string schemas.""" + engine's schema fixer coerces it back to "null" for string schemas. + """ tool = ChatCompletionToolsParam( function=FunctionDefinition( name="echo", @@ -655,7 +659,8 @@ def test_type_conversion_in_streaming(self): def test_string_attr_true_coerced_by_schema_streaming(self): """Streaming: string='true' delivers a string but the engine's - schema fixer coerces it to the schema type (integer).""" + schema fixer coerces it to the schema type (integer). + """ tool = ChatCompletionToolsParam( function=FunctionDefinition( name="score", @@ -971,7 +976,8 @@ def test_emits_arguments_before_invoke_completes(self, parser): def test_no_marker_leak_chunked(self, parser): """Chunked streaming must NOT leak DSML start-marker fragments - as content (GitHub #40801).""" + as content (GitHub #40801). + """ full_text = build_tool_call("fn", {"k": "v"}) deltas = self._stream_chunked(parser, full_text, chunk_size=5) content = "".join(d.content for d in deltas if d.content is not None) @@ -981,7 +987,8 @@ def test_no_marker_leak_chunked(self, parser): def test_no_marker_leak_with_prefix_chunked(self, parser): """Content before a tool call must not include start-marker - fragments when chunked (GitHub #40801).""" + fragments when chunked (GitHub #40801). + """ full_text = "Hello!" + build_tool_call("fn", {"a": "b"}) deltas = self._stream_chunked(parser, full_text, chunk_size=5) content = "".join(d.content for d in deltas if d.content is not None) @@ -993,7 +1000,8 @@ def test_no_marker_leak_with_prefix_chunked(self, parser): def test_no_marker_leak_char_by_char(self, parser): """Character-by-character streaming must not leak marker - fragments (GitHub #40801).""" + fragments (GitHub #40801). + """ full_text = build_tool_call("fn", {"k": "v"}) deltas = self._stream_chunked(parser, full_text, chunk_size=1) content = "".join(d.content for d in deltas if d.content is not None) @@ -1003,7 +1011,8 @@ def test_no_marker_leak_char_by_char(self, parser): def test_no_marker_leak_all_split_points(self, parser): """Start token split at every possible boundary must not - leak (GitHub #40801).""" + leak (GitHub #40801). + """ for chunk_size in range(1, len(FC_START) + 2): p = make_parser() full_text = build_tool_call("fn", {"k": "v"}) @@ -1015,7 +1024,8 @@ def test_no_marker_leak_all_split_points(self, parser): def test_false_partial_marker_emitted(self, parser): """Text ending with a prefix of the start token that turns out - NOT to be a marker must still be emitted as content.""" + NOT to be a marker must still be emitted as content. + """ full_text = "<｜DSM some regular text" deltas = self._stream_chunked(parser, full_text, chunk_size=3) content = "".join(d.content for d in deltas if d.content is not None) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 97ed0cd69d4c..bfb12ff15786 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -124,7 +124,8 @@ def mock_request(): class TestParseGemma4Args: """Values are returned as strings; type coercion to proper JSON types - happens at the engine layer.""" + happens at the engine layer. + """ def test_empty_string(self): assert _parse_gemma4_args("") == {} diff --git a/tests/tool_parsers/test_gigachat3_tool_parser.py b/tests/tool_parsers/test_gigachat3_tool_parser.py index a011fbf080d8..be9eb06f8803 100644 --- a/tests/tool_parsers/test_gigachat3_tool_parser.py +++ b/tests/tool_parsers/test_gigachat3_tool_parser.py @@ -18,7 +18,8 @@ @pytest.fixture(scope="function") def default_tokenizer() -> TokenizerLike: """Override module-scoped default_tokenizer because gigachat tests - mutate the tokenizer via ``add_tokens``.""" + mutate the tokenizer via ``add_tokens``. + """ return AutoTokenizer.from_pretrained("openai-community/gpt2") @@ -333,9 +334,7 @@ def test_streaming_tool_call_with_large_steps( model_output_deltas: list[str], gigachat_tokenizer: TokenizerLike, ): - """ - Test that the closing braces are streamed correctly. - """ + """Test that the closing braces are streamed correctly.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")( gigachat_tokenizer ) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 7fd3860ef719..881933890dfb 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -18,7 +18,6 @@ class TestInternLM2ToolParser(ToolParserTests): @pytest.fixture def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike: """Add some internlm2 specific tokens to the default vocab.""" - tokenizer_vocab = default_tokenizer.get_vocab() default_tokenizer.get_vocab = MagicMock() tokenizer_vocab.update( diff --git a/tests/tool_parsers/test_kimi_k2_tool_parser.py b/tests/tool_parsers/test_kimi_k2_tool_parser.py index b56032b91c17..69c5abfd78a6 100644 --- a/tests/tool_parsers/test_kimi_k2_tool_parser.py +++ b/tests/tool_parsers/test_kimi_k2_tool_parser.py @@ -218,6 +218,7 @@ def _split_tool_output_to_deltas( Args: content: text before tool section tool_strs: list of (tool_id, args_json) + """ deltas = [content, SECTION_BEGIN] for tool_id, args_json in tool_strs: diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py index 36b04b67b195..740f4c30130f 100644 --- a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Named tool choice for Kimi K3: allowed when the XTML structural tag is -attached (strict tool calling), rejected otherwise.""" +attached (strict tool calling), rejected otherwise. +""" import pytest diff --git a/tests/tool_parsers/test_lfm2_tool_parser.py b/tests/tool_parsers/test_lfm2_tool_parser.py index 9ed30e159884..599aec7d9236 100644 --- a/tests/tool_parsers/test_lfm2_tool_parser.py +++ b/tests/tool_parsers/test_lfm2_tool_parser.py @@ -416,7 +416,8 @@ def test_tool_call( def test_empty_tool_call_block(body: str, lfm2_tokenizer: TokenizerLike): """An empty block ([] or [ ]) means the model called no tools; it must not report tools_called=True with zero tool calls (the invariant - run_tool_extraction asserts).""" + run_tool_extraction asserts). + """ cls = ToolParserManager.get_tool_parser("lfm2") model_output = f"{TOOL_CALL_START}[{body}]{TOOL_CALL_END}" @@ -436,7 +437,8 @@ def test_whitespace_after_start_token(lfm2_tokenizer: TokenizerLike): not break parsing. The streaming path used to feed the indented text to ast.parse verbatim, so every completion candidate raised IndentationError and the call was silently dropped (non-streaming - stripped it and succeeded).""" + stripped it and succeeded). + """ cls = ToolParserManager.get_tool_parser("lfm2") model_output = f"{TOOL_CALL_START} [{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -476,7 +478,8 @@ def test_streaming_full_block_and_trailing_in_single_delta( ): """The entire <|tool_call_start|>[...]<|tool_call_end|> block plus trailing assistant text arrive in one delta. Trailing content must - still be emitted — not silently dropped.""" + still be emitted — not silently dropped. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = f"{TOOL_CALL_START}[{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}\nDone." @@ -491,7 +494,8 @@ def test_streaming_leading_content_and_full_block_in_single_delta( lfm2_tokenizer: TokenizerLike, ): """Leading assistant text plus the entire tool block arrive in one - delta. Leading content must be emitted — not silently dropped.""" + delta. Leading content must be emitted — not silently dropped. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( f"Let me check. {TOOL_CALL_START}[{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -508,7 +512,8 @@ def test_streaming_leading_block_and_trailing_in_single_delta( lfm2_tokenizer: TokenizerLike, ): """Leading text + complete tool block + trailing text in one delta. - Both leading and trailing content must be preserved.""" + Both leading and trailing content must be preserved. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( "Let me check. " @@ -529,7 +534,8 @@ def test_echoed_tool_call_body_not_leaked_to_content( """LFM2 sometimes emits the tool call body again after the first <|tool_call_end|>, capped with a second <|tool_call_end|>. The echoed body must not surface as assistant content — neither in - streaming nor non-streaming paths.""" + streaming nor non-streaming paths. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) body = ( "[grocery.orderIngredients(" @@ -564,7 +570,8 @@ def test_streaming_char_by_char_multi_dict_list(lfm2_tokenizer: TokenizerLike): character at a time. Every prefix lands in some partial-parse state (mid-key, mid-value, open quote inside dict, empty dict, etc.). The parser must not raise — incomplete prefixes should silently wait for - more text instead of logging exceptions.""" + more text instead of logging exceptions. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( f"{TOOL_CALL_START}[grocery.orderIngredients(" @@ -595,7 +602,8 @@ def test_streaming_char_by_char_multi_dict_list(lfm2_tokenizer: TokenizerLike): def test_streaming_dotted_name_in_single_delta(lfm2_tokenizer: TokenizerLike): """A pythonic call with a dotted/attribute function name (e.g. ``domain.method(arg=...)``) must be parsed correctly in streaming mode - just as in non-streaming mode.""" + just as in non-streaming mode. + """ tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = f"{TOOL_CALL_START}[{DOTTED_NAME_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -611,7 +619,8 @@ def test_adjust_request_disables_skip_special_tokens( """When tools are present, the parser must force ``skip_special_tokens=False`` so the engine does not strip the <|tool_call_start|>/<|tool_call_end|> sentinels before they reach the - parser.""" + parser. + """ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) diff --git a/tests/tool_parsers/test_llama3_json_tool_parser.py b/tests/tool_parsers/test_llama3_json_tool_parser.py index b8e194448184..e11434c0b0fd 100644 --- a/tests/tool_parsers/test_llama3_json_tool_parser.py +++ b/tests/tool_parsers/test_llama3_json_tool_parser.py @@ -252,7 +252,7 @@ def test_extract_tool_calls_missing_parameters_and_arguments_key(parser): def test_regex_timeout_handling(parser): - """Test regex timeout is handled gracefully""" + """Test regex timeout is handled gracefully.""" fake_problematic_input = "{hello world[A(A=" + "\t)A(A=,\t" * 2 # create a mock regex that raises TimeoutError diff --git a/tests/tool_parsers/test_llama4_pythonic_tool_parser.py b/tests/tool_parsers/test_llama4_pythonic_tool_parser.py index a3995c9f4bea..351e4115a6fc 100644 --- a/tests/tool_parsers/test_llama4_pythonic_tool_parser.py +++ b/tests/tool_parsers/test_llama4_pythonic_tool_parser.py @@ -247,7 +247,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike): @pytest.mark.parametrize("streaming", [False]) def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike): - """test regex timeout is handled gracefully""" + """Test regex timeout is handled gracefully.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")( default_tokenizer ) diff --git a/tests/tool_parsers/test_olmo3_tool_parser.py b/tests/tool_parsers/test_olmo3_tool_parser.py index 1928607222f9..02b5c621f61a 100644 --- a/tests/tool_parsers/test_olmo3_tool_parser.py +++ b/tests/tool_parsers/test_olmo3_tool_parser.py @@ -229,7 +229,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike): @pytest.mark.parametrize("streaming", [False]) def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike): - """test regex timeout is handled gracefully""" + """Test regex timeout is handled gracefully.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("olmo3")( default_tokenizer ) diff --git a/tests/tool_parsers/test_phi4mini_tool_parser.py b/tests/tool_parsers/test_phi4mini_tool_parser.py index eff9fa9bb8ff..98237ff5b271 100644 --- a/tests/tool_parsers/test_phi4mini_tool_parser.py +++ b/tests/tool_parsers/test_phi4mini_tool_parser.py @@ -16,7 +16,6 @@ class TestPhi4MiniToolParser(ToolParserTests): @pytest.fixture def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike: """Add some phi4mini specific tokens to the default vocab.""" - tokenizer = default_tokenizer tokenizer_vocab = tokenizer.get_vocab() tokenizer.get_vocab = MagicMock() diff --git a/tests/tool_parsers/test_pythonic_tool_parser.py b/tests/tool_parsers/test_pythonic_tool_parser.py index 05ddc411f142..fec713d5a78d 100644 --- a/tests/tool_parsers/test_pythonic_tool_parser.py +++ b/tests/tool_parsers/test_pythonic_tool_parser.py @@ -209,7 +209,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike): @pytest.mark.parametrize("streaming", [False]) def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike): - """test regex timeout is handled gracefully""" + """Test regex timeout is handled gracefully.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("pythonic")( default_tokenizer ) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index 946d490f94aa..c397d0ffadcf 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -423,7 +423,7 @@ def test_extract_tool_calls( def test_extract_tool_calls_fallback_no_tags( qwen3_tool_parser, ): - """Test fallback parsing when XML tags are missing""" + """Test fallback parsing when XML tags are missing.""" model_output = """ Dallas @@ -444,7 +444,7 @@ def test_extract_tool_calls_fallback_no_tags( def test_extract_tool_calls_type_conversion(qwen3_tokenizer): - """Test parameter type conversion based on tool schema""" + """Test parameter type conversion based on tool schema.""" tools = [ ChatCompletionToolsParam( type="function", @@ -862,7 +862,7 @@ def test_extract_tool_calls_streaming( expected_tool_calls, expected_content, ): - """Test incremental streaming behavior including typed parameters""" + """Test incremental streaming behavior including typed parameters.""" request = ChatCompletionRequest(model=MODEL, messages=[]) other_content = "" @@ -932,7 +932,7 @@ def test_extract_tool_calls_streaming( def test_extract_tool_calls_missing_closing_parameter_tag( qwen3_tool_parser, ): - """Test handling of missing closing tag""" + """Test handling of missing closing tag.""" # Using get_current_weather from sample_tools but with malformed XML model_output = """Let me check the weather for you: @@ -974,7 +974,7 @@ def test_extract_tool_calls_missing_closing_parameter_tag( def test_extract_tool_calls_streaming_missing_closing_tag( qwen3_tool_parser, qwen3_tokenizer ): - """Test streaming with missing closing tag""" + """Test streaming with missing closing tag.""" # Using get_current_weather from sample_tools but with malformed XML model_output = """Let me check the weather for you: @@ -1046,7 +1046,7 @@ def test_extract_tool_calls_streaming_missing_closing_tag( def test_extract_tool_calls_streaming_incremental(qwen3_tool_parser, qwen3_tokenizer): - """Test that streaming is truly incremental""" + """Test that streaming is truly incremental.""" model_output = """I'll check the weather. @@ -1106,7 +1106,7 @@ def test_extract_tool_calls_streaming_incremental(qwen3_tool_parser, qwen3_token def test_extract_tool_calls_streaming_missing_opening_tag( qwen3_tool_parser, qwen3_tokenizer ): - """Test streaming with missing opening tag + """Test streaming with missing opening tag. This tests that the streaming parser correctly handles tool calls that start directly with diff --git a/tests/tool_parsers/test_step3p5_tool_parser.py b/tests/tool_parsers/test_step3p5_tool_parser.py index d8f1f9332db6..de21aab79762 100644 --- a/tests/tool_parsers/test_step3p5_tool_parser.py +++ b/tests/tool_parsers/test_step3p5_tool_parser.py @@ -366,7 +366,7 @@ def test_extract_tool_calls( def test_extract_tool_calls_fallback_no_tags(step3p5_tool_parser, sample_tools): - """Test fallback parsing when XML tags are missing""" + """Test fallback parsing when XML tags are missing.""" model_output = """ Dallas @@ -387,7 +387,7 @@ def test_extract_tool_calls_fallback_no_tags(step3p5_tool_parser, sample_tools): def test_extract_tool_calls_type_conversion(step3p5_tokenizer): - """Test parameter type conversion based on tool schema""" + """Test parameter type conversion based on tool schema.""" tools = [ ChatCompletionToolsParam( type="function", @@ -622,7 +622,7 @@ def test_extract_tool_calls_streaming( expected_tool_calls, expected_content, ): - """Test incremental streaming behavior including typed parameters""" + """Test incremental streaming behavior including typed parameters.""" request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) other_content = "" @@ -692,7 +692,7 @@ def test_extract_tool_calls_streaming( def test_extract_tool_calls_missing_closing_parameter_tag( step3p5_tool_parser, sample_tools ): - """Test handling of missing closing tag""" + """Test handling of missing closing tag.""" # Using get_current_weather from sample_tools but with malformed XML model_output = """Let me check the weather for you: @@ -734,7 +734,7 @@ def test_extract_tool_calls_missing_closing_parameter_tag( def test_extract_tool_calls_streaming_missing_closing_tag( step3p5_tool_parser, step3p5_tokenizer, sample_tools ): - """Test streaming with missing closing tag""" + """Test streaming with missing closing tag.""" # Using get_current_weather from sample_tools but with malformed XML model_output = """Let me check the weather for you: @@ -808,7 +808,7 @@ def test_extract_tool_calls_streaming_missing_closing_tag( def test_extract_tool_calls_streaming_incremental( step3p5_tool_parser, step3p5_tokenizer, sample_tools ): - """Test that streaming is truly incremental""" + """Test that streaming is truly incremental.""" model_output = """I'll check the weather. @@ -864,7 +864,7 @@ def test_extract_tool_calls_streaming_incremental( def test_extract_tool_calls_complex_type_with_single_quote(step3p5_tokenizer): - """Test parameter type conversion based on tool schema""" + """Test parameter type conversion based on tool schema.""" tools = [ ChatCompletionToolsParam( type="function", diff --git a/tests/tool_use/test_chat_completions.py b/tests/tool_use/test_chat_completions.py index e5bb475875ac..585ca9314b0d 100644 --- a/tests/tool_use/test_chat_completions.py +++ b/tests/tool_use/test_chat_completions.py @@ -166,8 +166,7 @@ async def test_chat_completion_with_tools( async def test_response_format_with_tool_choice_required( client: openai.AsyncOpenAI, server_config: ServerConfig ): - """ - Test that combining response_format: json_object with tool_choice: required + """Test that combining response_format: json_object with tool_choice: required doesn't crash the engine. Before the fix, this would cause a validation error: diff --git a/tests/tool_use/test_gemma4_responses_adjust_request.py b/tests/tool_use/test_gemma4_responses_adjust_request.py index 937be9d016a4..fabfa7476660 100644 --- a/tests/tool_use/test_gemma4_responses_adjust_request.py +++ b/tests/tool_use/test_gemma4_responses_adjust_request.py @@ -171,7 +171,7 @@ def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None: def test_gemma4_required_skips_structured_outputs_chatcompletion() -> None: - """required + ChatCompletion: ``Gemma4EngineToolParser`` must skip the + """Required + ChatCompletion: ``Gemma4EngineToolParser`` must skip the forced JSON ``structured_outputs`` so the model emits its native ``<|tool_call>`` syntax. The base parser constrained output to JSON the native parser cannot read, leaking it as content with empty @@ -187,7 +187,7 @@ def test_gemma4_required_skips_structured_outputs_chatcompletion() -> None: def test_gemma4_named_skips_structured_outputs_chatcompletion() -> None: - """named + ChatCompletion: the forced single-function JSON schema must be + """Named + ChatCompletion: the forced single-function JSON schema must be skipped, same as ``required``. """ parser = Gemma4ToolParser(_StubTokenizer()) @@ -202,7 +202,7 @@ def test_gemma4_named_skips_structured_outputs_chatcompletion() -> None: def test_gemma4_required_skips_structured_outputs_responses() -> None: - """required + Responses: the forced JSON schema (``request.text``) must be + """Required + Responses: the forced JSON schema (``request.text``) must be skipped so the native delimiters reach the extractor. """ parser = Gemma4ToolParser(_StubTokenizer()) @@ -215,7 +215,7 @@ def test_gemma4_required_skips_structured_outputs_responses() -> None: def test_gemma4_named_skips_structured_outputs_responses() -> None: - """named (``ToolChoiceFunction``) + Responses: the forced single-function + """Named (``ToolChoiceFunction``) + Responses: the forced single-function JSON schema must be skipped. """ parser = Gemma4ToolParser(_StubTokenizer()) @@ -230,7 +230,7 @@ def test_gemma4_named_skips_structured_outputs_responses() -> None: def test_gemma4_keeps_special_tokens_with_tools_thinking_disabled() -> None: - """tools active + thinking disabled: ``skip_special_tokens`` must stay + """Tools active + thinking disabled: ``skip_special_tokens`` must stay False so ``<|tool_call>`` delimiters reach the extractor. The merged enable_thinking early-return stripped them, breaking tool calling when thinking is off. diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py index 77f99a47ec38..6232f10e4353 100644 --- a/tests/tool_use/test_muse_glimmer.py +++ b/tests/tool_use/test_muse_glimmer.py @@ -220,7 +220,8 @@ def test_reasoning_then_parallel_calls(): def _stream(raw: str, chunk: int): """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming - parsers; return (reasoning, content, tool_calls).""" + parsers; return (reasoning, content, tool_calls). + """ reasoning, content, toolcalls = [], [], [] prev = "" i = 0 diff --git a/tests/tool_use/test_parallel_tool_calls.py b/tests/tool_use/test_parallel_tool_calls.py index 4cfd165f1a8e..d1dc6e47921c 100644 --- a/tests/tool_use/test_parallel_tool_calls.py +++ b/tests/tool_use/test_parallel_tool_calls.py @@ -237,10 +237,7 @@ async def test_parallel_tool_calls_with_results( async def test_parallel_tool_calls_false( client: openai.AsyncOpenAI, server_config: ServerConfig ): - """ - Ensure only one tool call is returned when parallel_tool_calls is False. - """ - + """Ensure only one tool call is returned when parallel_tool_calls is False.""" models = await client.models.list() model_name: str = models.data[0].id messages = apply_parallel_tool_system_prompt( diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index f37a3c9681f9..bead3c6f3279 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -376,7 +376,8 @@ def test_streaming_output_valid_with_trailing_extra_data(): class TestNonFunctionToolsSkipped: """Non-function tools (web_search, etc.) must be silently skipped - by the tool-schema utilities instead of raising TypeError.""" + by the tool-schema utilities instead of raising TypeError. + """ def test_find_tool_properties_skips_web_search(self): tools = [WEB_SEARCH_TOOL, FUNCTION_TOOL] diff --git a/tests/transformers_utils/processors/test_glm5next.py b/tests/transformers_utils/processors/test_glm5next.py index 917a3f627ce6..0e24d9853bf4 100644 --- a/tests/transformers_utils/processors/test_glm5next.py +++ b/tests/transformers_utils/processors/test_glm5next.py @@ -330,7 +330,8 @@ def test_sample_frames_fps_interval(video_processor): def test_defaults_mirror_checkpoint_config(): """Bare instantiation matches the checkpoint's ``processor_config.json`` - token-budget style defaults; ``size`` carries no budget.""" + token-budget style defaults; ``size`` carries no budget. + """ image_processor = Glm5NextImageProcessor() assert image_processor.patch_expand_factor == 1 assert image_processor.min_image_tokens == 16 @@ -364,7 +365,8 @@ def test_missing_token_budgets_rejected(): def test_video_config_fields_land(): """fps_interval / max_frame_count_dynamic from the dedicated config - shape sampling without any request overrides.""" + shape sampling without any request overrides. + """ proc = Glm5NextVideoProcessor( fps_interval=4, max_frame_count_dynamic=32, diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 23a0f44a5dc4..18b33147f69d 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This test file includes some cases where it is inappropriate to +"""This test file includes some cases where it is inappropriate to only get the `eos_token_id` from the tokenizer as defined by `BaseRenderer.get_eos_token_id`. """ @@ -189,7 +188,8 @@ def test_model_config_generation_fallback_forwards_code_revision(): def test_safetensors_metadata_of_repo_without_safetensors(): """A repo storing its weights in another format is an answer, not a failure, - so it must not be retried.""" + so it must not be retried. + """ from huggingface_hub.errors import LocalEntryNotFoundError, NotASafetensorsRepoError get_safetensors_metadata = MagicMock( @@ -228,7 +228,8 @@ def test_mrope_num_dims(section_key, mrope_section, expected_num_dims): @pytest.mark.parametrize("section_name", ["mrope_section", "xdrope_section"]) def test_mrope_num_dims_from_config_attribute(section_name): """Some configs expose the section as an attribute rather than under - `rope_parameters`.""" + `rope_parameters`. + """ config = PretrainedConfig() setattr(config, section_name, [16, 16, 16, 16]) diff --git a/tests/transformers_utils/test_hf_overrides_model_type.py b/tests/transformers_utils/test_hf_overrides_model_type.py index cce10d3b084e..242594d63196 100644 --- a/tests/transformers_utils/test_hf_overrides_model_type.py +++ b/tests/transformers_utils/test_hf_overrides_model_type.py @@ -22,8 +22,8 @@ def test_hf_overrides_model_type_returns_correct_config_class(): """When hf_overrides sets model_type to a registered custom type whose checkpoint has a *different* model_type on disk, get_config() must return an instance of the registered config class — not the class that matches - the on-disk model_type.""" - + the on-disk model_type. + """ # Register the custom config _CONFIG_REGISTRY["test_custom_model"] = _TestCustomConfig diff --git a/tests/utils.py b/tests/utils.py index e0166785ba5f..5b988d27401e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -72,7 +72,8 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None: """Pre-populate the HF cache for (repo_id, filename) pairs that upstream trust_remote_code modules would otherwise fetch from third-party CDNs - (often unreachable from US-based CI).""" + (often unreachable from US-based CI). + """ if HF_HUB_OFFLINE: return for repo_id, filename in assets: @@ -245,7 +246,7 @@ def _create_cli_subcommand(self): def _start_server( self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None ) -> None: - """Subclasses override this method to customize server process launch""" + """Subclasses override this method to customize server process launch.""" raise NotImplementedError def _pre_download_model(self, model: str, args) -> None: @@ -703,7 +704,7 @@ def _wait_for_gpu_memory_release( ) def _poll(self) -> int | None: - """Subclasses override this method to customize process polling""" + """Subclasses override this method to customize process polling.""" return self.proc.poll() def _wait_for_server(self, *, url: str, timeout: float): @@ -866,7 +867,7 @@ def _wait_for_gpu_memory_release( class RemoteOpenAIServerCustom(RemoteOpenAIServer): - """Launch test server with custom child process""" + """Launch test server with custom child process.""" def _start_server( self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None @@ -892,7 +893,8 @@ def __init__( max_wait_seconds: float | None = None, ) -> None: """Store custom child process function then invoke superclass - constructor which will indirectly launch it.""" + constructor which will indirectly launch it. + """ self.child_process_fxn = child_process_fxn super().__init__( model=model, @@ -1213,8 +1215,7 @@ def compare_two_settings( include_seeded_sampling: bool = True, force_v1_runner: bool = False, ) -> None: - """ - Launch API server with two different sets of arguments/environments + """Launch API server with two different sets of arguments/environments and compare the results of the API calls. Args: @@ -1228,8 +1229,8 @@ def compare_two_settings( force_v1_runner: Whether to pin all compared settings to the v1 model runner to avoid mixing model runner differences into correctness tests. - """ + """ compare_all_settings( model, [arg1, arg2], @@ -1251,9 +1252,9 @@ def compare_all_settings( include_seeded_sampling: bool = True, force_v1_runner: bool = False, ) -> None: - """ - Launch API server with several different sets of arguments/environments + """Launch API server with several different sets of arguments/environments and compare the results of the API calls with the first set of arguments. + Args: model: The model to test. all_args: A list of argument lists to pass to the API server. @@ -1263,8 +1264,8 @@ def compare_all_settings( force_v1_runner: Whether to pin all compared settings to the v1 model runner to avoid mixing model runner differences into correctness tests. - """ + """ if force_v1_runner: all_envs = [ {"VLLM_USE_V2_MODEL_RUNNER": "0", **(env or {})} for env in all_envs @@ -1394,6 +1395,7 @@ def ensure_current_vllm_config(): with ensure_current_vllm_config(): init_distributed_environment(...) ensure_model_parallel_initialized(...) + """ from vllm.config import ( VllmConfig, @@ -1575,8 +1577,7 @@ def assert_rocm_custom_allreduce_backend_state_on_worker( @contextmanager def error_on_warning(category: type[Warning] = Warning): - """ - Within the scope of this context manager, tests will fail if any warning + """Within the scope of this context manager, tests will fail if any warning of the given category is emitted. """ with warnings.catch_warnings(): @@ -2029,6 +2030,7 @@ def create_new_process_for_each_test( Returns: A decorator to run test functions in separate processes. + """ if method is None: method = "spawn" if requires_spawn_multiprocessing() else "fork" @@ -2042,8 +2044,7 @@ def create_new_process_for_each_test( def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator: - """ - Get a pytest mark, which skips the test if the GPU doesn't meet + """Get a pytest mark, which skips the test if the GPU doesn't meet a minimum memory requirement in GB. This can be leveraged via `@large_gpu_test` to skip tests in environments @@ -2078,8 +2079,7 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator: def large_gpu_test(*, min_gb: int): - """ - Decorate a test to be skipped if no GPU is available or it does not have + """Decorate a test to be skipped if no GPU is available or it does not have sufficient memory. Currently, the CI machine uses L4 GPU which has 24 GB VRAM. @@ -2104,9 +2104,7 @@ def multi_gpu_marks(*, num_gpus: int): def multi_gpu_test(*, num_gpus: int): - """ - Decorate a test to be run only when multiple GPUs are available. - """ + """Decorate a test to be run only when multiple GPUs are available.""" marks = multi_gpu_marks(num_gpus=num_gpus) def wrapper(f: Callable[_P, None]) -> Callable[_P, None]: @@ -2120,13 +2118,13 @@ def wrapper(f: Callable[_P, None]) -> Callable[_P, None]: def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None): - """ - Mark a test to only run when the GPU count falls within [min_gpus, max_gpus]. + """Mark a test to only run when the GPU count falls within [min_gpus, max_gpus]. Examples: @gpu_tier_mark(min_gpus=2) # only on multi-GPU @gpu_tier_mark(max_gpus=1) # only on single-GPU @gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only + """ gpu_count = current_platform.device_count() marks = [] @@ -2194,8 +2192,8 @@ async def completions_with_server_args( Returns: OpenAI Completion instance - """ + """ if isinstance(max_tokens, int): max_tokens = [max_tokens] * len(prompts) @@ -2253,9 +2251,7 @@ def get_client_text_logprob_generations( def has_module_attribute(module_name, attribute_name): - """ - Helper function to check if a module has a specific attribute. - """ + """Helper function to check if a module has a specific attribute.""" try: module = importlib.import_module(module_name) return hasattr(module, attribute_name) @@ -2327,14 +2323,15 @@ def _native_plain( def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)): - """ - Generate prompts which a bunch of assignments, + """Generate prompts which a bunch of assignments, then asking for the value of one of them. The prompt is just under 10k tokens; sliding window is 4k so the answer is outside sliding window, but should still be correct. + Args: batch_size: number of prompts to generate ln_range: an argument to control the length of the prompt + """ prompts: list[str] = [] answer: list[int] = [] @@ -2374,8 +2371,7 @@ def check_answers( def flat_product(*iterables: Iterable[Any]): - """ - Flatten lists of tuples of the cartesian product. + """Flatten lists of tuples of the cartesian product. Useful when we want to avoid nested tuples to allow test params to be unpacked directly from the decorator. @@ -2387,6 +2383,7 @@ def flat_product(*iterables: Iterable[Any]): (3, 4, "a"), (3, 4, "b"), ] + """ for element in itertools.product(*iterables): normalized = (e if isinstance(e, tuple) else (e,) for e in element) @@ -2394,8 +2391,7 @@ def flat_product(*iterables: Iterable[Any]): class TestFP8Layer(torch.nn.Module): - """ - Test helper for FP8 linear operations. Creates random weights and scales + """Test helper for FP8 linear operations. Creates random weights and scales based on quantization configuration. Args: @@ -2404,6 +2400,7 @@ class TestFP8Layer(torch.nn.Module): weight_quant_key: Weight quantization configuration. out_dtype: Output dtype. Defaults to current default dtype. force_kernel: Optional kernel to force use of specific implementation. + """ def __init__( diff --git a/tests/utils_/__init__.py b/tests/utils_/__init__.py index e6b4c3f6364c..57451f3bd857 100644 --- a/tests/utils_/__init__.py +++ b/tests/utils_/__init__.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This module is named `utils_` instead of `utils` to avoid obscuring +"""This module is named `utils_` instead of `utils` to avoid obscuring `tests/utils.py`. """ diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index 10b550c0a174..efcda2dfd823 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -59,7 +59,8 @@ def test_with_env_set(monkeypatch, mode): @create_new_process_for_each_test() def test_other_threads_are_not_policed(monkeypatch): """A background thread that syncs deliberately must not be broken by the - check being armed on the thread running the decorated function.""" + check being armed on the thread running the decorated function. + """ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -84,7 +85,8 @@ def worker(): def test_allow_on_other_thread_does_not_disarm(monkeypatch): """`gpu_sync_allowed()` on one thread must not suppress the check on another. It is scoped by ContextVar rather than torch's process-global - sync debug mode, which a previous implementation mutated.""" + sync debug mode, which a previous implementation mutated. + """ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -113,7 +115,8 @@ def worker(): def test_suppressing_works_while_compiling(monkeypatch): """`_suppressing` wraps torch compile entry points, which run with `torch.compiler.is_compiling()` true. `gpu_sync_allowed()` deliberately - no-ops in that state, so `_suppressing` must not route through it.""" + no-ops in that state, so `_suppressing` must not route through it. + """ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) # Emulate being inside a torch compile, as inductor passes are. @@ -145,7 +148,8 @@ def convert_like(x, like): def test_sync_debug_mode_restored_after_checked_call(monkeypatch): """The mode is armed only for the duration of a checked call. Leaving it on process-wide made every sync outside a checked region emit a - `UserWarning` whenever our handler was not the installed one.""" + `UserWarning` whenever our handler was not the installed one. + """ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -290,7 +294,8 @@ def test_genuinely_async_transfers_pass(monkeypatch): asynchronous in both directions -- even with a dtype conversion, which runs GPU-side, and even permuted (e.g. transposed), which uses pitched cudaMemcpy2D/3DAsync; D2H `Tensor.to` allocates a pinned destination. - None of these may be flagged.""" + None of these may be flagged. + """ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) gsd._install_copy_checkers() diff --git a/tests/utils_/test_jsontree.py b/tests/utils_/test_jsontree.py index 0af2751b2638..49e75c3477ef 100644 --- a/tests/utils_/test_jsontree.py +++ b/tests/utils_/test_jsontree.py @@ -5,7 +5,6 @@ def test_json_count_leaves(): """Test json_count_leaves function from jsontree utility.""" - # Single leaf values assert json_count_leaves(42) == 1 assert json_count_leaves("hello") == 1 diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index 02889c20f66a..e2935a254954 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -29,7 +29,8 @@ def test_get_file_store_init_method_is_unique(): def _call_with_timeout(func, timeout: float = 10.0): """Run func in a daemon thread so a livelock regression fails the test - quickly instead of hanging the suite.""" + quickly instead of hanging the suite. + """ result: dict = {} def target(): @@ -96,7 +97,8 @@ def test_get_open_port_skips_reserved_dp_master_ports( monkeypatch: pytest.MonkeyPatch, ): """VLLM_PORT == VLLM_DP_MASTER_PORT must not livelock by returning the - same free-but-reserved port forever (upstream issue #50024).""" + same free-but-reserved port forever (upstream issue #50024). + """ base = 42123 with monkeypatch.context() as m: m.setenv("VLLM_PORT", str(base)) @@ -109,7 +111,8 @@ def test_get_open_ports_list_skips_reserved_dp_master_ports( monkeypatch: pytest.MonkeyPatch, ): """The VLLM_PORT band scan must not hand out ports inside the window - reserved for the data parallel master process.""" + reserved for the data parallel master process. + """ base = 42223 reserved = range(base, base + 10) with monkeypatch.context() as m: @@ -125,7 +128,8 @@ def test_get_open_port_ephemeral_skips_reserved_range( monkeypatch: pytest.MonkeyPatch, ): """An ephemeral port that happens to fall in the reserved range is - replaced by a rescan starting past the range.""" + replaced by a rescan starting past the range. + """ base = 42323 with monkeypatch.context() as m: m.delenv("VLLM_PORT", raising=False) @@ -144,7 +148,8 @@ def test_get_open_port_ephemeral_without_dp_master_port( monkeypatch: pytest.MonkeyPatch, ): """Without VLLM_PORT and VLLM_DP_MASTER_PORT, an ephemeral port is - returned as before.""" + returned as before. + """ with monkeypatch.context() as m: m.delenv("VLLM_PORT", raising=False) m.delenv("VLLM_DP_MASTER_PORT", raising=False) diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 9f718703a7c1..845fdda2a122 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -157,7 +157,8 @@ def test_pct_binding_filters_cpus(monkeypatch): ) def test_pct_binding_fires_on_every_capable_sku(monkeypatch, sku, expected_cpus): """Each SKU in ``_PCT_CAPABLE_SKUS`` engages the gate at its own - expected ``highest_perf`` and uses its own priority-core stride.""" + expected ``highest_perf`` and uses its own priority-core stride. + """ sku_config = numa_utils._PCT_CAPABLE_SKUS[sku] _patch_pct_gates( monkeypatch, @@ -242,7 +243,8 @@ def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes(): def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch): """PCT: EngineCore for a multi-NUMA shard binds to the union of priority cores across all shard nodes, so worker `--physcpubind` is always a - subset of EngineCore's `cpus_allowed`.""" + subset of EngineCore's `cpus_allowed`. + """ _patch_pct_gates( monkeypatch, model_match=True, @@ -292,7 +294,8 @@ def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes( """external_launcher (or multi-node-within-DP, or Ray) hits the fallback branch. EngineCore must still span every local NUMA node so it can mp-spawn its local workers without ``--physcpubind`` - strict-validation failures.""" + strict-validation failures. + """ _patch_pct_gates( monkeypatch, model_match=True, @@ -316,7 +319,8 @@ def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes( def test_get_numactl_args_engine_core_baseline_multi_node_within_dp_spans_locals(): """Multi-node-within-DP fallback: bind EngineCore to all local NUMA - nodes that the visible ``numa_bind_nodes`` reference.""" + nodes that the visible ``numa_bind_nodes`` reference. + """ vllm_config = _make_config( numa_bind=True, numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1], @@ -339,7 +343,8 @@ def test_get_numactl_args_engine_core_skips_user_cpu_list(monkeypatch): ``--physcpubind`` spawns need. We fall back to ``--cpunodebind`` over the shard's NUMA nodes instead. PCT auto-detect is also bypassed when the user is explicit (its priority-core union may not be a superset - of the user's per-worker cores).""" + of the user's per-worker cores). + """ _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) vllm_config = _make_config( numa_bind=True, @@ -392,7 +397,8 @@ def test_get_numactl_args_requires_detectable_nodes(monkeypatch): def test_configure_subprocess_rejects_unknown_process_kind(): """configure_subprocess only knows 'worker' and 'EngineCore'; anything else must raise ValueError instead of silently routing to the worker - path.""" + path. + """ vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0]) with ( pytest.raises(ValueError, match="process_kind"), diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 6efdf2a31b0b..af9ed0c6b0d0 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -151,6 +151,7 @@ def create_and_prepopulate_kv_cache( Returns: A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size, 2 * head_size)`` order with strides determined by ``layout``. + """ batch_size = len(k_contexts) seq_lens = common_attn_metadata.seq_lens.cpu() @@ -286,7 +287,6 @@ def run_attention_backend( use_cuda_graph: bool = False, ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" - use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0") if backend == "FLEX_ATTENTION_SLOW": use_direct_block_mask = False @@ -414,8 +414,7 @@ def _test_backend_correctness( max_num_seqs: int | None = None, max_num_batched_tokens: int | None = None, ): - """ - Test that all backends produce similar outputs to a reference implementation + """Test that all backends produce similar outputs to a reference implementation using FlexAttention or an explicit attention-sink reference. This test works by: @@ -1393,7 +1392,8 @@ def test_non_causal_backend_correctness( model: str, ): """Test backend's correctness with non-causal (bidirectional) decoder - attention, as used by DFlash speculative decoding.""" + attention, as used by DFlash speculative decoding. + """ def bidirectional_mask_mod( b: torch.Tensor, diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 8486d216a125..0481ebc79094 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -115,7 +115,8 @@ def test_mamba_layers_have_unified_interface( layer_class, expected_backend, expected_mamba_type ): """Test that all Mamba layers have the unified get_attn_backend - interface.""" + interface. + """ assert hasattr(layer_class, "get_attn_backend"), ( f"{layer_class.__name__} should have get_attn_backend method" ) diff --git a/tests/v1/attention/test_attention_splitting.py b/tests/v1/attention/test_attention_splitting.py index acbc22a02793..ff0e06514e86 100644 --- a/tests/v1/attention/test_attention_splitting.py +++ b/tests/v1/attention/test_attention_splitting.py @@ -20,12 +20,12 @@ @pytest.fixture def sample_query_start_loc(): - """Sample query_start_loc tensor for testing""" + """Sample query_start_loc tensor for testing.""" return torch.tensor([0, 5, 12, 20, 35, 50]) def test_basic_slice_middle(sample_query_start_loc): - """Test slicing from middle of tensor""" + """Test slicing from middle of tensor.""" req_slice = slice(1, 3) # slice from index 1 to 3 result = slice_query_start_locs(sample_query_start_loc, req_slice) @@ -34,7 +34,7 @@ def test_basic_slice_middle(sample_query_start_loc): def test_slice_from_beginning(sample_query_start_loc): - """Test slicing from the beginning of tensor""" + """Test slicing from the beginning of tensor.""" req_slice = slice(0, 2) # slice from index 0 to 2 result = slice_query_start_locs(sample_query_start_loc, req_slice) @@ -43,7 +43,7 @@ def test_slice_from_beginning(sample_query_start_loc): def test_slice_to_end(sample_query_start_loc): - """Test slicing to the end of tensor""" + """Test slicing to the end of tensor.""" req_slice = slice(3, 5) # slice from index 3 to 5 (last index) result = slice_query_start_locs(sample_query_start_loc, req_slice) @@ -52,7 +52,7 @@ def test_slice_to_end(sample_query_start_loc): def test_single_element_slice(sample_query_start_loc): - """Test slice that results in single element""" + """Test slice that results in single element.""" req_slice = slice(2, 3) # slice from index 2 to 3 result = slice_query_start_locs(sample_query_start_loc, req_slice) @@ -61,7 +61,7 @@ def test_single_element_slice(sample_query_start_loc): def test_full_tensor_slice(sample_query_start_loc): - """Test slicing the entire tensor""" + """Test slicing the entire tensor.""" req_slice = slice(0, 5) # slice entire tensor result = slice_query_start_locs(sample_query_start_loc, req_slice) @@ -80,7 +80,7 @@ def test_slice_bounds_edge_cases(sample_query_start_loc): @pytest.fixture def small_decode_metadata(): - """Create metadata for small decode batch""" + """Create metadata for small decode batch.""" batch_spec = BATCH_SPECS["small_decode"] device = torch.device("cpu") return create_common_attn_metadata(batch_spec, block_size=16, device=device) @@ -88,7 +88,7 @@ def small_decode_metadata(): @pytest.fixture def large_decode_metadata(): - """Create metadata for small decode batch""" + """Create metadata for small decode batch.""" batch_spec = BATCH_SPECS["large_decode"] device = torch.device("cpu") return create_common_attn_metadata(batch_spec, block_size=16, device=device) @@ -96,7 +96,7 @@ def large_decode_metadata(): @pytest.fixture def mixed_small_metadata(): - """Create metadata for mixed small batch""" + """Create metadata for mixed small batch.""" batch_spec = BATCH_SPECS["mixed_small"] device = torch.device("cpu") return create_common_attn_metadata(batch_spec, block_size=16, device=device) @@ -104,7 +104,7 @@ def mixed_small_metadata(): # Tests for _make_metadata_with_slice def test_make_metadata_with_slice_decode_batch(small_decode_metadata): - """Test slicing decode batch metadata""" + """Test slicing decode batch metadata.""" # Split first request only ubatch_slice = UBatchSlice(slice(0, 1), slice(0, 1)) @@ -119,7 +119,7 @@ def test_make_metadata_with_slice_decode_batch(small_decode_metadata): def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata): - """Test slicing mixed batch metadata""" + """Test slicing mixed batch metadata.""" ubatch_slice = UBatchSlice(slice(1, 3), slice(1, 7)) # Requests 1-3, tokens 1-7 result = _make_metadata_with_slice(ubatch_slice, mixed_small_metadata) @@ -132,7 +132,7 @@ def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata): def test_split_attn_metadata_decode_batch(large_decode_metadata): - """Test splitting decode batch into two equal parts""" + """Test splitting decode batch into two equal parts.""" num_tokens = large_decode_metadata.num_reqs mid_point = num_tokens // 2 ubatch_slices = [ @@ -164,7 +164,8 @@ def apply_split_decodes_and_prefills( treat_short_extends_as_decodes: bool = True, ): """Helper function to apply split_decodes_and_prefills and return - the results.""" + the results. + """ device = torch.device("cpu") seq_lens = [10 * (i + 1) for i in range(len(query_lens))] common_metadata = create_common_attn_metadata( @@ -297,7 +298,7 @@ def test_split_decodes_and_prefills_uniform_mixed_batch_non_uniform_decodes(): def test_split_decodes_and_prefills_uniform_padded_batch_all_same(): - """uniform batch where all query lengths are identical with 0 length padded reqs.""" + """Uniform batch where all query lengths are identical with 0 length padded reqs.""" # All query lengths are 2, with decode_threshold=3 (so 2 <= 3) # This triggers the padded uniform path at line 891 query_lens = [2, 2, 2, 0] @@ -324,7 +325,7 @@ def test_split_decodes_and_prefills_uniform_padded_batch_all_same(): def test_prefill_split_across_ubatches( seq_lens, query_lens, split_point, expected_first_reqs, expected_second_reqs ): - """Test splitting a prefill across ubatches""" + """Test splitting a prefill across ubatches.""" import numpy as np device = torch.device("cpu") diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py index cff1228baf03..34621900d5a1 100644 --- a/tests/v1/attention/test_dcp_a2a_pack_mask.py +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The empty-shard mask fused into the DCP A2A pack kernel must be bit-exact -with the eager ``mask_dcp_empty_shards_`` pass it replaces.""" +with the eager ``mask_dcp_empty_shards_`` pass it replaces. +""" import pytest import torch diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index d5c794131474..6f61a6aeed86 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -495,7 +495,8 @@ def test_dspark_noncausal_sparse_mla_matches_sdpa( num_heads, ): """Non-causal (window ∪ block, future-pointing) per-token indices must match - a dense SDPA reference over the same indices, for both sparse-MLA backends.""" + a dense SDPA reference over the same indices, for both sparse-MLA backends. + """ _skip_if_backend_unavailable(backend_cls, kv_cache_dtype, block_size) window, block, topk_width = _DSPARK_CONFIGS[config_name] diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 8dc082b08bbc..c712cb5844b0 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -231,7 +231,8 @@ def test_plan_uses_state_params(monkeypatch): def test_kv_lens_host_formula(): """Per-row host lengths: context == position + 1; capped at - index_topk + trailing-pool remainder past the sparse threshold.""" + index_topk + trailing-pool remainder past the sparse threshold. + """ builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 221f933d8942..c8032d9eaade 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -187,7 +187,8 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase): def test_has_initial_state_after_reclassification(): """After reclassification, num_prefills > 0 so the prefill kernel path should compute has_initial_state. For the reclassified request with - context_lens > 0, the corresponding entry must be True.""" + context_lens > 0, the corresponding entry must be True. + """ builder = _create_gdn_builder(num_speculative_tokens=2) batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3]) meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2]) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 8bc409adb704..4175bb6ff5e1 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -56,7 +56,8 @@ def _ref_stable_topk_from_candidates_fp64( ) -> torch.Tensor: """Pure-PyTorch reference for the CuteDSL stable-topk selector order (score desc, then lowest global token id). Selects the same SET as the - kernel; only the set is compared in tests.""" + kernel; only the set is compared in tests. + """ num_rows, num_candidates = candidate_scores.shape device = candidate_scores.device select_k = min(k, num_candidates) @@ -121,7 +122,8 @@ def _dcp_lse_merge( class _FakeDCPGroup: """Single-process stand-in: ``all_gather`` returns the pre-built concatenation of every rank's packed ``(score, global_id)`` candidates, - mirroring the one packed all-gather the merge issues.""" + mirroring the one packed all-gather the merge issues. + """ def __init__(self, gathered_packed: torch.Tensor) -> None: self.gathered_packed = gathered_packed @@ -207,7 +209,8 @@ def _merge_local_topks_global_with_fake_dcp( """Run ``_merge_dcp_topk_global`` for every rank against a faked all-gather and return each rank's global-index result (all ranks should agree). The fake pre-builds the packed candidate concatenation the merge's - single ``all_gather`` would return.""" + single ``all_gather`` would return. + """ # The merge is now CuteDSL-only (no PyTorch fallback), so it runs the real # Triton pack + CuteDSL selector kernels even behind the faked all-gather. if not current_platform.is_cuda() or not has_cutedsl(): @@ -709,7 +712,8 @@ def test_dcp_filter_compaction_matches_reference( """In-kernel compaction must, for every row, produce exactly the rank-owned physical slots packed into [0, valid_count) with -1 in the tail -- the same SET a reference filter + sort/gather produces. Uses wide rows (> BLOCK_N - valid slots), with interior -1 gaps.""" + valid slots), with interior -1 gaps. + """ device = torch.device("cuda") torch.manual_seed(7) dcp_size = 2 @@ -867,7 +871,8 @@ def test_decode_topk_pads_surplus_with_negative_one(): so a non-(-1) pad would be silently attended. This is the common case under DCP, where each rank's local seq_len = global / world is usually << topk. Checked on the *set* of valid indices (the kernel may order them by score, - not position).""" + not position). + """ torch.manual_seed(0) device = torch.device("cuda") topk = 8 @@ -913,7 +918,8 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp(): """End-to-end DCP decode where the global seq_len < topk (so every rank's local top-k is surplus-padded). Exercises the kernel surplus -> merge mask -> global top-k -> physical localize -> LSE merge chain for the common - short-context decode case, vs the non-DCP reference.""" + short-context decode case, vs the non-DCP reference. + """ torch.manual_seed(4) device = torch.device("cuda") world = 2 diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 0966f9b20a6b..6377d7df4661 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -58,7 +58,8 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the request's single KpoolTailManager block, the remaining columns are never - written and stay zero.""" + written and stay zero. + """ bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) return bt @@ -66,7 +67,8 @@ def make_tail_block_table(own_blocks, width=64): def legacy_generic_tail_slots(block_table, query_start_loc, positions): """Reference of the generic ``_compute_slot_mappings_kernel`` arithmetic - (block_table.py:305-313) applied to the tail group's table.""" + (block_table.py:305-313) applied to the tail group's table. + """ slots = [] for req in range(block_table.shape[0]): for i in range(query_start_loc[req], query_start_loc[req + 1]): @@ -107,7 +109,8 @@ def make_batch(per_req_positions, padded_len=None): def test_legacy_generic_mapping_collapses_onto_block_zero(): """The bug: with the manager's 1-column block table, the generic kernel - maps every pos >= kpool onto tail block 0, and distinct requests collide.""" + maps every pos >= kpool onto tail block 0, and distinct requests collide. + """ own_blocks = [5, 9] per_req = [list(range(10)), list(range(12))] # prompts of len 10 and 12 positions, qsl, _, num_actual, num_reqs = make_batch(per_req) @@ -137,7 +140,8 @@ def test_legacy_generic_mapping_collapses_onto_block_zero(): def test_circular_mapping_isolates_requests(): """The fix: every token lands in its own request's block at pos % kpool, and no slot is ever shared by two different requests (slots do recur - within a request every kpool positions -- that is the circular design).""" + within a request every kpool positions -- that is the circular design). + """ own_blocks = [5, 9] per_req = [list(range(10)), list(range(12))] positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) @@ -165,7 +169,8 @@ def test_circular_mapping_isolates_requests(): def test_circular_mapping_matches_generic_for_short_requests(prompt_len): """For pos < kpool the generic kernel already picks the own block, so the two mappings agree while every position fits the request's first block - (single-request behavior is unchanged).""" + (single-request behavior is unchanged). + """ own_blocks = [7] per_req = [list(range(prompt_len))] positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) @@ -275,7 +280,8 @@ class TailRingMirror: """Mirror of _kpool_tail_seed_kernel / _kpool_decode_update_batched_kernel addressing: block = tail_slot // kpool, ring offset = pos % kpool; a pool completing at pos reads ring slots (pool_start + s) % kpool and uses the - current token's own K/score for the last member.""" + current token's own K/score for the last member. + """ def __init__(self, num_blocks, kpool=KPOOL): self.kpool = kpool @@ -318,7 +324,8 @@ def tail_slot_for(mapping, req, pos, own_block): def run_scenario(mapping, interleave): """Requests A (block 5, prompt len 9) and B (block 9, prompt len 11) - decode concurrently; returns A's boundary pool [8, 9, 10, 11].""" + decode concurrently; returns A's boundary pool [8, 9, 10, 11]. + """ ring = TailRingMirror(num_blocks=16) blocks = {"A": 5, "B": 9} prompts = {"A": 9, "B": 11} diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 4ec138270203..f281213221bc 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Regression test for https://github.com/vllm-project/vllm/issues/34865 +"""Regression test for https://github.com/vllm-project/vllm/issues/34865. When multiple KV cache groups share the same MambaSpec (as in Nemotron hybrid models), the metadata caching optimization reuses metadata from @@ -29,7 +28,8 @@ def _make_vllm_config( block_size: int | None = None, ): """Create a minimal mock VllmConfig with only the fields the builder - accesses, avoiding any model download / HF config inspection.""" + accesses, avoiding any model download / HF config inspection. + """ speculative_config = ( SimpleNamespace( num_speculative_tokens=num_speculative_tokens, @@ -75,8 +75,8 @@ def test_mamba_single_token_prompt_runs_as_prefill(): def test_update_block_table_copies_block_idx_to_persistent_buffers(): """update_block_table() must write block_idx tensors to the current builder's persistent buffers, not leave them pointing to a different - builder's buffers.""" - + builder's buffers. + """ block_size = 16 max_model_len = 256 num_reqs = 4 @@ -180,8 +180,8 @@ def test_state_indices_tensor_d_includes_num_speculative_blocks(): bug 1: with mamba_cache_mode='all' and speculative decoding enabled, the cudagraph buffer for state_indices_tensor_d must allocate the same per-request column count as the runtime block table, which includes - num_speculative_blocks trailing scratch columns.""" - + num_speculative_blocks trailing scratch columns. + """ block_size = 16 max_model_len = 256 max_num_seqs = 4 @@ -214,8 +214,8 @@ def test_block_idx_cudagraph_capture_padded_by_num_reqs(): bug 2: with mamba_cache_mode='all' and spec decode, _update_metadata_for _cudagraph_capture must slice block_idx_last_{scheduled,computed}_token by the request count (padded_bs == num_reqs), not by num_decode_tokens. - Past num_decodes, the slice must be zero-filled.""" - + Past num_decodes, the slice must be zero-filled. + """ block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -295,7 +295,8 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): block_idx_last_scheduled_token_prev_step as a persistent buffer with the same shape as the existing block_idx_last_{scheduled,computed}_token buffers, so cudagraph capture records a stable pointer for the prev-step - input anchor consumed by mamba_mixer2's input gather.""" + input anchor consumed by mamba_mixer2's input gather. + """ block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -323,7 +324,8 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode(): """Without spec decode, the prev-step buffer is unused and must not be - allocated — the input anchor reduces to last_computed_token.""" + allocated — the input anchor reduces to last_computed_token. + """ block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -347,7 +349,8 @@ def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer(): """_update_metadata_for_cudagraph_capture must copy the prev-step anchor into the builder's persistent buffer (so cudagraph replay reads from the same underlying memory), pad past num_decodes with zero, and return a - slice of the persistent buffer in the metadata.""" + slice of the persistent buffer in the metadata. + """ block_size = 16 max_model_len = 256 max_num_seqs = 8 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index f724949f54d8..5a8fdea20464 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -171,7 +171,8 @@ def test_glm5_flashinfer_masked_mha_routing( def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim): """The K concat used by the MLA prefill context loop must equal torch.cat of k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself - instead of allocating and copying.""" + instead of allocating and copying. + """ torch.manual_seed(0) num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256 k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16) @@ -515,6 +516,7 @@ def create_and_prepopulate_kv_cache( Returns: MLA KV cache tensor + """ batch_size = len(kv_c_contexts) seq_lens = common_attn_metadata.seq_lens.cpu() @@ -1307,7 +1309,6 @@ def run_attention_backend( chunked_prefill_workspace_size: int | None = None, ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" - builder_cls, impl_cls = try_get_attention_backend(backend) # Force the prefill backend selection (None means auto-select). @@ -1434,8 +1435,7 @@ def _run_backend_correctness( v_head_dim: int, chunked_prefill_workspace_size: int | None = None, ): - """ - Test that all backends produce similar outputs to a reference implementation + """Test that all backends produce similar outputs to a reference implementation using torch.nn.functional.scaled_dot_product_attention. This test works by: @@ -1454,7 +1454,6 @@ def _run_backend_correctness( multiple GPUs. This tests that backends work correctly with different head counts. """ - # Filter backends to those that support the requested kv_cache_dtype backends_to_test = [ b diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py index 220f6a198dba..211cde663bc3 100644 --- a/tests/v1/attention/test_mla_prefill_quant_output.py +++ b/tests/v1/attention/test_mla_prefill_quant_output.py @@ -134,7 +134,8 @@ def test_flash_attn_registers_warmup_only_for_fa4( def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): """run_prefill_new_tokens must accept out/output_scale so the direct - (non-**kwargs) call in forward_mha type- and runtime-checks.""" + (non-**kwargs) call in forward_mha type- and runtime-checks. + """ import inspect params = inspect.signature( @@ -150,7 +151,8 @@ def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): def test_mla_impl_forward_mha_accepts_output_scale(): """The abstract MLA impl forward_mha must carry output_scale so every - override (and the unconditional forward_impl call) stays compatible.""" + override (and the unconditional forward_impl call) stays compatible. + """ import inspect from vllm.v1.attention.backend import MLAAttentionImpl @@ -177,7 +179,8 @@ def _fused_fp8_skip_reason() -> str | None: def test_fa4_fused_fp8_output_matches_post_quant(default_vllm_config): """FA4's fused FP8 write (output_scale, flash-attention#135) must match the bf16-attention + standalone static-FP8-quant path it replaces, since - production uses the same output_scale for both.""" + production uses the same output_scale for both. + """ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index 098eddfbf6eb..fa52b1484b71 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -315,7 +315,8 @@ def test_backend_supported_dimension_validation(self): def test_flash_attn_accepts_glm53_flash_nope_dimensions(self): """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a - RoPE-carrying 320-wide query does not.""" + RoPE-carrying 320-wide query does not. + """ try: from vllm.v1.attention.backends.mla.prefill.flash_attn import ( FlashAttnPrefillBackend, diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 89cb0f566467..351fe0006897 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -257,7 +257,8 @@ def test_replayssm_write_pos(case: ReplaySSMBuildCase): def test_resumed_request_differs_from_fresh(): """Same token count, different decode_base: fresh (base 100) -> write_pos 5, - resumed (base 105) -> write_pos 0.""" + resumed (base 105) -> write_pos 0. + """ builder = _create_replayssm_builder(16) batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1]) common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index efffcad927db..b8ba46d8dec6 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -514,7 +514,6 @@ def test_mtp_builder_init_sizes_native_fp8_metadata( Sweeping num_heads asserts metadata is sized for the padded decode shape, covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding. """ - dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") info_calls = [] @@ -694,7 +693,6 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr( monkeypatch, ): """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759.""" - get_mla_metadata_v1 = mock.MagicMock() monkeypatch.setitem( sys.modules, @@ -756,7 +754,6 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr( def test_decode_expands_kernel_block_page_indices(monkeypatch): """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704.""" - expand_kernel = _ExpandPageIndicesKernel() monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel) # qlen==1 now takes the persistent-metadata path, which imports diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 2342a279a397..01018a630c86 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -517,7 +517,8 @@ def test_auto_selection_for_kv_connector( aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version ): """Auto-selection with a KV connector and AITER enabled resolves to unified attn, - and to triton attn if AITER not enabled.""" + and to triton attn if AITER not enabled. + """ from vllm.platforms.rocm import RocmPlatform with patch( diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py index d87a82c37b2b..9cbb7a6f8f45 100644 --- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py +++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py @@ -102,7 +102,8 @@ def expected_row_seq_lens(per_req_positions, batch_size, next_n): def test_legacy_flat_layout_misaligns_and_bleeds(): """The bug: ``positions[:n] + 1`` with padded-row coordinates reads other - requests' positions and prefill positions.""" + requests' positions and prefill positions. + """ decode_lens, per_req, positions = make_non_uniform_batch() batch_size = decode_lens.shape[0] next_n = int(decode_lens.max()) @@ -137,7 +138,8 @@ def test_helper_uniform_layout_matches_flat_slice(): def test_helper_padded_layout_per_row(): """Non-uniform batches map every padded row to its own token's position; - pad rows get 0 (empty tail).""" + pad rows get 0 (empty tail). + """ decode_lens, per_req, positions = make_non_uniform_batch() out = _decode_topk_seq_lens(positions, decode_lens, 8, 3, 4, requires_padding=True) expected = expected_row_seq_lens(per_req, 3, 4) @@ -149,7 +151,8 @@ def test_helper_padded_layout_per_row(): def expand_tail_region(dec_seq): """Run the production pure-torch expand + append pair (the fused kernel replicates it exactly on the identity path) and return the tail columns - [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1).""" + [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1). + """ rows = dec_seq.shape[0] pool_ids = torch.arange(SELECT_K, dtype=torch.int64).expand(rows, SELECT_K) valid = torch.ones_like(pool_ids, dtype=torch.bool) @@ -163,7 +166,8 @@ def expand_tail_region(dec_seq): def test_tail_expansion_legacy_vs_fixed(): """End-to-end tail consequence: the fixed mapping appends exactly the request's trailing incomplete pool; the legacy flat slice drops real tail - tokens and emits indices far past the request's own sequence.""" + tokens and emits indices far past the request's own sequence. + """ decode_lens, per_req, positions = make_non_uniform_batch() n = 12 fixed = _decode_topk_seq_lens( diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 6d2c4fbd5880..ecaf71d55b67 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -245,8 +245,8 @@ def _dequantize_fp8_ds_mla_entry( Args: simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's float -> e8m0 -> bf16 scale conversion path. - """ + """ # The first kv_lora_rank bytes store FP8 latent values with one scale per # 128 element tile written as float32 right after the latent payload. scales = cache_slice.view(torch.float32)[kv_lora_rank // 4 : kv_lora_rank // 4 + 4] @@ -283,8 +283,8 @@ def _quantize_dequantize_fp8_ds_mla( Args: simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's float -> e8m0 -> bf16 scale conversion in dequantization. - """ + """ if kv_c.numel() == 0: return kv_c.clone(), k_pe.clone() @@ -909,7 +909,8 @@ def test_index_group_convert_during_piecewise_capture(): runs in the following eager break. The convert's side stream must not wait on ``logical_topk_ready`` there: an event recorded inside a captured segment is graph-local, and an eager wait on it raises - cudaErrorInvalidValue.""" + cudaErrorInvalidValue. + """ device = torch.device(DEVICE_TYPE) num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16 @@ -1048,7 +1049,8 @@ def test_triton_convert_rejects_req_id_longer_than_token_indices(): req_id but the output is allocated like token_indices, so a full-batch req_id combined with an MQA-subset token_indices wrote past the end of the output buffer. The wrapper must reject the length mismatch instead - of corrupting memory.""" + of corrupting memory. + """ device = torch.device(DEVICE_TYPE) num_topk_tokens = 128 block_size = 64 @@ -1103,7 +1105,8 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): is active, forward_mqa only receives the leading decode tokens, but _forward_bf16_kv passed the full-batch req_id_per_token to the index conversion, making it write past the end of its output buffer. The call - site must slice req_id_per_token to the MQA tokens.""" + site must slice req_id_per_token to the MQA tokens. + """ device = torch.device(DEVICE_TYPE) num_topk_tokens = 128 block_size = 64 @@ -1201,7 +1204,8 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected): ) def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expected): """A 32K prefill needs the default workspace exactly; shrinking it would - push a supported request onto MQA.""" + push a supported request onto MQA. + """ assert ( _masked_mha_workspace_fits( batch_size=1, @@ -1259,7 +1263,8 @@ def test_is_masked_mha_available_model_dims( ): """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2, GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an - unquantized KV cache, nothing else.""" + unquantized KV cache, nothing else. + """ import vllm.model_executor.layers.attention.sparse_mla_attention as mod monkeypatch.setattr( @@ -3599,7 +3604,8 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of the cross-rank LSE merge, or a NaN would survive the merge even at zero - weight (0 * NaN = NaN).""" + weight (0 * NaN = NaN). + """ num_tokens, num_heads, head_dim = 3, 2, 3 q = torch.empty(num_tokens, num_heads, head_dim, device=DEVICE_TYPE) local_indices = torch.tensor( @@ -3772,7 +3778,8 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): def test_sparse_impl_observes_repointed_indexer_buffer(): """The MTP proposer repoints the draft's indexer at the target model's buffer after the backend impl is built, so the impl must resolve the buffer per read. - Snapshotting it in __init__ leaves the layer reading indices nothing writes.""" + Snapshotting it in __init__ leaves the layer reading indices nothing writes. + """ impl = object.__new__(FlashInferMLASparseImpl) own = torch.zeros(4, 8, dtype=torch.int32) target = torch.ones(4, 8, dtype=torch.int32) @@ -3788,7 +3795,8 @@ def test_sparse_impl_observes_repointed_indexer_buffer(): def test_explicit_topk_buffer_supersedes_indexer(): """Backbone skip-topk layers have no indexer, and the proposer also assigns the - shared buffer directly onto draft submodules.""" + shared buffer directly onto draft submodules. + """ impl = object.__new__(FlashInferMLASparseImpl) indexer = SimpleNamespace( topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32), @@ -3803,7 +3811,8 @@ def test_explicit_topk_buffer_supersedes_indexer(): def test_sparse_mla_common_impl_resolves_buffer_lazily(): """Guards the whole SparseMLACommonImpl family at once: a backend that - snapshots the buffer instead silently loses MTP buffer sharing.""" + snapshots the buffer instead silently loses MTP buffer sharing. + """ assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer) assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( "topk_indices_buffer must stay a lazily-resolved property" diff --git a/tests/v1/attention/test_trtllm_attention_integration.py b/tests/v1/attention/test_trtllm_attention_integration.py index d81bc2de4ea1..36c2a710d532 100644 --- a/tests/v1/attention/test_trtllm_attention_integration.py +++ b/tests/v1/attention/test_trtllm_attention_integration.py @@ -124,6 +124,7 @@ def _create_nvfp4_hnd_kv_cache( Returns: ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided). + """ # First create a bf16 cache so block tables are populated. bf16_cache = create_and_prepopulate_kv_cache( @@ -197,7 +198,8 @@ def _create_nvfp4_hnd_kv_cache( def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL): """Run TRTLLM attention through the full FlashInfer pipeline - and compare against an SDPA reference.""" + and compare against an SDPA reference. + """ set_random_seed(42) device = torch.device(f"{DEVICE_TYPE}:0") @@ -425,7 +427,8 @@ def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len): def test_trtllm_gen_full_attention_integration(batch_spec_name: str): """Test TRTLLM gen-full attention through the full FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline, - with real TRTLLM kernels on Blackwell.""" + with real TRTLLM kernels on Blackwell. + """ _run_trtllm_integration(BATCH_SPECS[batch_spec_name]) @@ -436,7 +439,8 @@ def test_trtllm_gen_full_attention_integration(batch_spec_name: str): @torch.inference_mode() def test_trtllm_gen_nvfp4_kv_integration(batch_spec_name: str): """Test TRTLLM attention with nvfp4 KV cache through the full - FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline.""" + FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline. + """ _run_trtllm_integration( BATCH_SPECS[batch_spec_name], kv_cache_dtype="nvfp4", diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 48e618565c72..cd062cb486b8 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -193,7 +193,6 @@ def create_vllm_config( hf_config_override: dict | None = None, ) -> VllmConfig: """Create a VllmConfig for testing with reasonable defaults.""" - model_config = ModelConfig( model=model_name, tokenizer=model_name, diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 393ec484651e..eeb977adef65 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Fine-grained partial prefix-cache hits for hybrid (full attention + mamba "align") models: scheduler chunk splitting, partial tail registration, CoW -on partial hits, and same-step deferral.""" +on partial hits, and same-step deferral. +""" from math import lcm from types import SimpleNamespace @@ -174,7 +175,8 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: - 0 -> 8192 -> 9728 -> 9984 -> 10000.""" + 0 -> 8192 -> 9728 -> 9984 -> 10000. + """ block_size = 512 scheduler_block_size = block_size * dcp_world_size hash_block_size = 32 @@ -742,7 +744,8 @@ def test_hash_aligned_query_end_uses_regular_partial_tail(): def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): """An external mid-block hit must become a running request even when its - first continuation does not need another Mamba block.""" + first continuation does not need another Mamba block. + """ hash_block_size = 2 mamba_block_size = 4 * hash_block_size kv_cache_config = KVCacheConfig( @@ -814,7 +817,8 @@ def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): def test_boundary_state_offloads_returns_cow_target(): """Boundary hand-offs expose aligned snapshots and the partial-tail CoW - target, never the overwritten CoW source.""" + target, never the overwritten CoW source. + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1000,7 +1004,8 @@ def request_finished_all_groups(self, request, block_ids): def test_boundary_state_offload_dropped_when_request_freed_before_drain(): """A hand-off recorded in the same scheduling pass as the request's death must not be drained: its release hook has already run, so draining would - leak a pinned block.""" + leak a pinned block. + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1047,7 +1052,8 @@ def test_boundary_state_offload_dropped_when_request_freed_before_drain(): def test_boundary_state_offloads_block_aligned_prompt(): """A prompt ending on a block boundary registers no CoW partial tail; its - boundary state block is handed off as a snapshot instead (once).""" + boundary state block is handed off as a snapshot instead (once). + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1101,7 +1107,8 @@ def test_boundary_state_offloads_block_aligned_prompt(): def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): """truncate_computed_blocks slices each group by its own block size, keeps null placeholders in the retained prefix, and leaves the original - lookup result untouched (pure view, no refcount changes).""" + lookup result untouched (pure view, no refcount changes). + """ hash_block_size = 2 kv_cache_config = KVCacheConfig( num_blocks=24, @@ -1156,7 +1163,8 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): def test_truncate_computed_blocks_allows_short_mamba_group_only(): """External state may replace a short Mamba hit, but other groups must - cover the aligned local endpoint.""" + cover the aligned local endpoint. + """ hash_block_size = 2 kv_cache_config = KVCacheConfig( num_blocks=24, @@ -1298,7 +1306,8 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): """The owner's move re-arms the same-step guard: the moved entry is filled by this step's copy, and chained same-step copies read stale - sources, so a request hitting it in the move step must be deferred.""" + sources, so a request hitting it in the move step must be deferred. + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1588,7 +1597,8 @@ def test_hybrid_partial_hash_truncates_full_attention_hit_length(): def test_cow_retained_blocks_returned_for_release(): """new_step_starts returns the CoW copy retentions instead of freeing - them; the scheduler owns releasing them once the copy has run.""" + them; the scheduler owns releasing them once the copy has run. + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1641,7 +1651,8 @@ def test_cow_retained_blocks_returned_for_release(): def test_free_cow_retained_blocks_defers_until_copy_step_processed(): """Scheduler releases CoW retentions immediately when the copy's step has - been processed (or deferral is off), and defers them otherwise.""" + been processed (or deferral is off), and defers them otherwise. + """ from collections import deque freed: list = [] @@ -1678,7 +1689,8 @@ def test_free_cow_retained_blocks_defers_until_copy_step_processed(): def test_full_attention_eagle_drops_one_hash_unit(): """With fine-grained partial hits, eagle rewinds the hit by one hash unit instead of a whole cache block: the tail block's KV is append-only, so it - still covers the reduced length and stays in the hit as a partial block.""" + still covers the reduced length and stays in the hit as a partial block. + """ from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager @@ -1754,7 +1766,8 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks(): """Regression: with eagle, the mamba group must not receive the eagle lookup margin — its finder never applies the drop, so it could return a hit past the blocks the (dropped) full-attention group covers, crashing - the consumer's CoW with block_idx >= len(req_blocks).""" + the consumer's CoW with block_idx >= len(req_blocks). + """ hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -2356,7 +2369,8 @@ def test_retention_snapshots_handed_off_with_exact_block_ids(): """Under sparse retention, each retained mamba boundary state block is handed to the connector with its exact block id. Connectors must not resolve align-mode state blocks positionally because the block table is not - append-only.""" + append-only. + """ hash_block_size = 2 block_size = 4 manager = make_kv_cache_manager( @@ -2397,7 +2411,8 @@ def test_snapshot_handoff_dense_default_retention(): """With dense (default) retention, every materialized mamba boundary state block is handed off — regular mamba-align + prefix-match-unit deployments get store-able boundary snapshots without setting a - prefix-cache retention interval.""" + prefix-cache retention interval. + """ hash_block_size = 2 block_size = 4 manager = make_kv_cache_manager( @@ -2429,7 +2444,8 @@ def test_snapshot_handoff_dense_default_retention(): def _run_chunked_prefill(manager, req, chunk_size, num_chunks): """Prefill ``req`` one ``chunk_size`` chunk per scheduler step, yielding the - boundary-state hand-offs offered (and claimed) at each step.""" + boundary-state hand-offs offered (and claimed) at each step. + """ computed_blocks, num_computed, _ = manager.get_computed_blocks(req) for step in range(num_chunks): manager.new_step_starts() @@ -2493,7 +2509,8 @@ def test_boundary_states_offered_past_prompt_for_resumed_prefill(): ``num_prompt_tokens``. A resumed request re-prefills its generated tokens and every group re-saves them, so filtering on the original prompt length would silently strip the mamba key for boundaries full attention still - stores; only the connector knows where its save window ends.""" + stores; only the connector knows where its save window ends. + """ block_size = 8 prompt_len = block_size manager = make_kv_cache_manager( diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7e27a4517e03..097878ff7d50 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -552,7 +552,8 @@ def _create_async_pp_scheduler( def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None: """Delivered tokens must be an order-preserving subset of the emitted - tokens with no duplicates (tokens are globally unique).""" + tokens with no duplicates (tokens are globally unique). + """ it = iter(emitted) for token in delivered: assert token in it, f"token {token} delivered out of order or twice" @@ -562,7 +563,8 @@ def _assert_positions_consistent(req, engine: PipelinedEngine) -> None: """The i-th delivered output token must be one the runner sampled for exactly sequence position prompt_len + i: catches a preempted request's stale output landing on a position the resumed request resampled (or - vice versa), which token-stream equality alone cannot see.""" + vice versa), which token-stream equality alone cannot see. + """ for i, token in enumerate(req.output_token_ids): expected = req.num_prompt_tokens + i actual = engine.emitted_position[token] diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 2fcf70c6efa7..f3692672e970 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -178,7 +178,8 @@ def _shared_layout_config(): class TestCSALinearGrouping: """A CSA + linear-attention model (sparse attention with a compressor ring, plus sharded GDN and one TP-replicated PLE state) goes through the generic - packed-group path; no model-specific branch is involved.""" + packed-group path; no model-specific branch is involved. + """ @staticmethod def _mamba_groups(groups): @@ -193,7 +194,8 @@ def _mamba_groups(groups): def test_replicated_state_gets_its_own_group(self): """The PLE state is TP-replicated and a different size, so it must not share a manager group with the sharded GDN states (the NIXL worker - requires exactly one single-layer replicated group).""" + requires exactly one single-layer replicated group). + """ groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs()) mamba_groups = self._mamba_groups(groups) @@ -246,7 +248,8 @@ def test_every_group_fits_one_packed_block(self): def test_scratch_group_survives_computed_block_truncation(self): """The scratch group contributes no computed blocks, so truncating a lookup result must skip it: its block size is the ring capacity, which - neither divides the hit length nor bounds the (empty) block list.""" + neither divides the hit length nor bounds the (empty) block list. + """ config = _shared_layout_config() config.cache_config.enable_prefix_caching = True groups = get_kv_cache_groups(config, _make_csa_linear_specs()) @@ -360,7 +363,8 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid per page size, which is emitted whole, or a single-layer state that cannot be split at all -- the mamba layers must be sized against it. Sizing them against a narrower bucket splits them past what the block - already fits, spending a pool block per extra group for no saving.""" + already fits, spending a pool block per extra group for no saving. + """ config = _mock_vllm_config("BLNHC") config.speculative_config = None specs = {} @@ -405,7 +409,8 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would set the block width and the MLA group would fill under a third of every block; capped, it splits into groups no wider than the MLA - page.""" + page. + """ config = _shared_layout_config() specs: dict[str, KVCacheSpec] = {} for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)): @@ -606,7 +611,8 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring claims one block per request, never hashes, and does not disturb the - scheduler/hash block sizes the paged groups imply.""" + scheduler/hash block sizes the paged groups imply. + """ config = _shared_layout_config() config.cache_config.enable_prefix_caching = True specs: dict[str, KVCacheSpec] = { diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index 3948f7c6e300..7f0747abe0c7 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -367,7 +367,8 @@ def test_free_request_with_duplicate_mm_hashes(): """Freeing a request whose two inputs share the same mm_hash must fully clean up request_cached_ids. After the first free_encoder_input call, cached[mm_hash] becomes empty; the second call must still remove the - remaining input_id from request_cached_ids.""" + remaining input_id from request_cached_ids. + """ manager = EncoderCacheManager(cache_size=20) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) @@ -383,7 +384,8 @@ def test_free_request_with_duplicate_mm_hashes(): def test_duplicate_mm_hash_stays_referenced_until_last_free(): """`cached` holds one reference per request, not per position, so freeing - the first of two occurrences must not release the entry.""" + the first of two occurrences must not release the entry. + """ manager = EncoderCacheManager(cache_size=20) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) @@ -404,7 +406,8 @@ def test_duplicate_mm_hash_stays_referenced_until_last_free(): def test_duplicate_mm_hash_is_not_evicted_before_last_use(): """An item repeated within one request (an image carried across conversation turns) must not become reclaimable capacity while a later - occurrence still has to be spliced in, or the encoder recomputes it.""" + occurrence still has to be spliced in, or the encoder recomputes it. + """ manager = EncoderCacheManager(cache_size=8) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) other = MockRequest("r2", ["imgB"], [8]) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 86e66e1941f1..a0600abd8510 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1069,9 +1069,7 @@ def _stats(requests: int, queries: int, hits: int) -> PrefixCacheStats: def test_metrics(): - """ - Test the prefix caching metrics. - """ + """Test the prefix caching metrics.""" metrics = CachingMetrics(max_recent_requests=5) assert metrics.hit_rate == 0.0 @@ -1101,9 +1099,7 @@ def test_metrics(): def test_metrics_empty_stats(): - """ - Test the prefix caching metrics with empty stats. - """ + """Test the prefix caching metrics with empty stats.""" metrics = CachingMetrics(max_recent_requests=5) metrics.observe(_stats(0, 0, 0)) metrics.observe(_stats(1, 20, 9)) @@ -1903,7 +1899,7 @@ def test_get_max_concurrency_for_kv_cache_config(): def test_allocate_with_lookahead(): - """Verify that lookahead tokens correctly affect block allocation""" + """Verify that lookahead tokens correctly affect block allocation.""" block_size = 4 config = KVCacheConfig( num_blocks=10, @@ -2597,7 +2593,8 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): """The kpool tail parasitizes the indexer tensors instead of getting its own: sibling idx/tail tensors paired by layer order, zero standalone tail bytes, one shared block per request, and no prefix-caching leakage from - the kpool-sized scratch group.""" + the kpool-sized scratch group. + """ model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2668,7 +2665,8 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): def test_glm5_kpool_tail_does_not_drag_hash_block_size(): """The tail's kpool-sized scratch block (4 tokens) must not constrain the - prefix-cache hash granularity: participating groups alone decide it.""" + prefix-cache hash granularity: participating groups alone decide it. + """ model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2744,7 +2742,8 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer(): def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba(): """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not - disable slot sharing; the layout re-pads them to the MLA page.""" + disable slot sharing; the layout re-pads them to the MLA page. + """ model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2775,7 +2774,8 @@ def prepadded_mamba_spec(): def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection(): """Round-robin mamba grouping keeps practical PP splits balanced: every stage's largest projected mamba group slice fits its projected MLA - layers, so per-stage slot tensors all have an MLA owner.""" + layers, so per-stage slot tensors all have an MLA owner. + """ model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2819,7 +2819,8 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatc """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA layers but a round-robin slice of 3 under the minimum 4 mamba groups; grouping bumps to 5 groups so every stage's projection keeps sharing - on instead of silently falling back on stage 0.""" + on instead of silently falling back on stage 0. + """ from vllm.config import ParallelConfig from vllm.distributed.utils import get_pp_indices @@ -2868,7 +2869,8 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch): def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag(): """Hybrid slot sharing must take precedence over the experimental enable_cross_layers_blocks packed layout: generic packing would give the - MLA slots a strided view, breaking the contiguous virtual split.""" + MLA slots a strided view, breaking the contiguous virtual split. + """ model_config = ModelConfig(max_model_len=8192) kv_transfer_config = KVTransferConfig( kv_connector="NixlConnector", @@ -3170,7 +3172,8 @@ def test_multi_run_layer_compact_strides_place_hoisted_heads(): """A layer-compact run region is its own dense allocation: under LHBNC the head groups sit between the layers and the blocks, so a run's block stride is one head group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models).""" + on ROCm hybrid models). + """ full = new_kv_cache_spec() swa = new_sliding_window_spec(sliding_window=full.block_size * 2) page = full.page_size_bytes @@ -3249,7 +3252,8 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): def test_hidden_states_with_tp_scales_page_size(): """When TP shrinks KV pages below the hidden-state per-token cost, get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states.""" + common page accommodates the unsharded hidden states. + """ # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) # = 16 * 2 * 64 * 2 = 4096 bytes. @@ -3922,8 +3926,7 @@ def test_unify_kv_cache_spec_page_size_mamba(): def test_hma_not_disabled_when_kv_events_enabled(): - """ - Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. + """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. This test guards against that regression by verifying that a VllmConfig with kv_events_config set still resolves disable_hybrid_kv_cache_manager @@ -4080,7 +4083,8 @@ def test_check_enough_kv_cache_memory_reserves_null_block(): def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers - scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" + scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs. + """ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) full = FullAttentionSpec(block_size=4, **common) mla = MLAAttentionSpec(block_size=4, **common) @@ -4249,7 +4253,8 @@ def test_no_warning_when_draft_group_is_identified(caplog_vllm): def _deepseek_v4_specs(model_version="deepseek_v4"): """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last.""" + at differing page sizes, with the MTP draft layer registered last. + """ return { "model.layers.0.self_attn.attn": new_mla_spec(), "model.layers.1.self_attn.attn": new_mla_spec(), diff --git a/tests/v1/core/test_kv_sharing.py b/tests/v1/core/test_kv_sharing.py index db0e8dae3e78..ca8787e8d8c0 100644 --- a/tests/v1/core/test_kv_sharing.py +++ b/tests/v1/core/test_kv_sharing.py @@ -17,8 +17,7 @@ def new_kv_cache_spec(): def test_initialize_kv_cache_for_kv_sharing_different_attn_groups(): - """ - Test initializing KV cache sharing with different attention groups. + """Test initializing KV cache sharing with different attention groups. Layers in the same KV cache group might be placed in different attn groups if they have different attention backends. """ @@ -50,8 +49,7 @@ def test_initialize_kv_cache_for_kv_sharing_different_attn_groups(): def test_initialize_kv_cache_for_kv_sharing_same_attn_groups(): - """ - Test case assuming that all layers in the same KV cache group have the same + """Test case assuming that all layers in the same KV cache group have the same attention backends. This is true for most models. """ shared_kv_cache_layers = { @@ -79,8 +77,7 @@ def test_initialize_kv_cache_for_kv_sharing_same_attn_groups(): def test_initialize_kv_cache_for_kv_sharing_no_attn_groups(): - """ - Test KV sharing set up when no attention groups are provided. + """Test KV sharing set up when no attention groups are provided. This is the case for the TPU model runner, which doesn't have support for attention groups yet. """ diff --git a/tests/v1/core/test_output.py b/tests/v1/core/test_output.py index ee8a4a0e78f9..4e1e22eb8e42 100644 --- a/tests/v1/core/test_output.py +++ b/tests/v1/core/test_output.py @@ -55,7 +55,8 @@ def _mm_feature(offset: int, length: int) -> MultiModalFeatureSpec: def test_strip_covered_mm_data() -> None: """Items fully inside the computed prefix lose their data; items touching - the uncomputed region keep it; already-None data stays None.""" + the uncomputed region keep it; already-None data stays None. + """ from dataclasses import replace covered = _mm_feature(offset=0, length=100) @@ -107,7 +108,8 @@ def _mm_feature_mixed(offset: int, length: int) -> MultiModalFeatureSpec: def test_strip_covered_mm_data_mrope() -> None: """For M-RoPE models, covered items keep their keep_on_cpu metadata fields - (the worker needs them to compute positions); payload fields are dropped.""" + (the worker needs them to compute positions); payload fields are dropped. + """ covered = _mm_feature_mixed(offset=0, length=100) uncovered = _mm_feature_mixed(offset=300, length=100) @@ -127,7 +129,8 @@ def test_strip_covered_mm_data_shm_address_item() -> None: """SHM address descriptors must survive stripping even when covered: the worker needs the address to resolve the payload and to acknowledge the sender's writer reference. Stripping it crashed the SHM receiver cache on - the second identical request (vllm-project/vllm#54994).""" + the second identical request (vllm-project/vllm#54994). + """ address_item = MultiModalKwargsItem( { "address": MultiModalFieldElem(data=4096, field=MultiModalBatchedField()), diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index d91260e30ae0..653d485a82ed 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -126,7 +126,8 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the config (LCM of group block sizes) unless explicitly provided. This mirrors ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by - these tests, so callers don't have to pass it at every site.""" + these tests, so callers don't have to pass it at every site. + """ kwargs.setdefault( "scheduler_block_size", lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)), @@ -1826,8 +1827,7 @@ def _test_partial_request_hit( def _make_hybrid_kv_cache_config( block_size: int, num_blocks: int, spec_types: list[str] ) -> KVCacheConfig: - """ - Create a KVCacheConfig with the specified spec types. + """Create a KVCacheConfig with the specified spec types. Args: block_size: The block size for KV cache. @@ -1837,6 +1837,7 @@ def _make_hybrid_kv_cache_config( - "sliding_window": SlidingWindowSpec with window=2*block_size - "sliding_window_large": SlidingWindowSpec with window=4*block_size - "mamba": MambaSpec + """ spec_map = { "full": lambda: FullAttentionSpec( @@ -1938,8 +1939,7 @@ def _make_hybrid_kv_cache_config( @pytest.mark.parametrize("spec_types", _HYBRID_MODEL_TEST_CASES) def test_prefill_hybrid_model_combinations(spec_types: list[str]): - """ - Test prefix caching with hybrid models containing various combinations of + """Test prefix caching with hybrid models containing various combinations of KV cache spec types. This unified test covers: @@ -2021,8 +2021,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): def test_prefill_hybrid_model_combinations_eagle( spec_types: list[str], expect_hit_length: int ): - """ - Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE. + """Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE. More complex hybrid models with EAGLE are not yet supported (see issue #32802). """ block_size = 16 @@ -2489,8 +2488,7 @@ def test_evict(): def test_hash_block_correct_reuse(): - """ - This tests when a previously cached block is reused as a new block, + """This tests when a previously cached block is reused as a new block, its hash metadata should be correctly reset. """ block_size = 16 @@ -2530,8 +2528,7 @@ def test_hash_block_correct_reuse(): def test_computed_blocks_not_evicted(): - """ - Test that the computed blocks are not evicted when getting new blocks + """Test that the computed blocks are not evicted when getting new blocks for a request if there are any other free blocks. """ block_size = 16 @@ -2590,9 +2587,7 @@ def test_computed_blocks_not_evicted(): def test_basic_prefix_caching_disabled(): - """ - This tests that the prefix caching is disabled. - """ + """This tests that the prefix caching is disabled.""" block_size = 4 manager = make_kv_cache_manager( make_kv_cache_config(block_size, 5), @@ -2639,11 +2634,9 @@ def test_basic_prefix_caching_disabled(): @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_cache_blocks(hash_fn): - """ - This is a unit test that tests the correctness of the _cache_full_blocks + """This is a unit test that tests the correctness of the _cache_full_blocks function of KVCacheManager. """ - block_size = 4 block_pool = BlockPool( num_gpu_blocks=5, @@ -2688,9 +2681,7 @@ def test_cache_blocks(hash_fn): def test_cache_blocks_multi_group(): - """ - This tests that blocks are cached correctly for different kv cache groups. - """ + """This tests that blocks are cached correctly for different kv cache groups.""" block_size = 4 block_pool = BlockPool( num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size @@ -2773,10 +2764,7 @@ def test_cache_blocks_multi_group(): def test_mm_prefix_caching(): - """ - This tests that the multi-modal prefix caching is correct. - """ - + """This tests that the multi-modal prefix caching is correct.""" block_size = 16 manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), @@ -2881,8 +2869,7 @@ def test_mm_prefix_caching(): def test_cache_key_salting(): - """ - This tests that cache salts are applied during hashing and the cache + """This tests that cache salts are applied during hashing and the cache is separated cache as expected. """ block_size = 16 @@ -2961,8 +2948,7 @@ def test_cache_key_salting(): def test_prefill_not_enough_free_blocks_with_computed_blocks(): - """ - This is a unit test that tests the correctness of the allocate_slots + """This is a unit test that tests the correctness of the allocate_slots when there is not enough free blocks. Specifically, when a request has computed blocks but cannot be allocated due to not enough free blocks, the computed blocks should not be touched. @@ -3344,7 +3330,8 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int): @pytest.mark.parametrize("group_id", [0, 1, 2]) def test_block_stored_event_group_idx(group_id: int): """Test BlockStored events emitted by cache_full_blocks carry the correct - group_idx.""" + group_idx. + """ block_size = 4 num_tokens = block_size * 2 @@ -3462,8 +3449,7 @@ def test_session_id_does_not_affect_prefix_cache_identity(): def test_block_stored_event_group_idx_multiple_groups(): - """ - Test BlockStored events for separate HMA groups that each carry the + """Test BlockStored events for separate HMA groups that each carry the correct group_idx. Simulates the HMA scenario where full-attention blocks (group 0) and @@ -3585,8 +3571,7 @@ def collect_warning(message, *args, **kwargs): @pytest.mark.parametrize("group_id", [0, 1, 2]) def test_block_removed_event_group_idx(group_id: int): - """ - Test BlockRemoved events emitted on eviction carry the group_idx extracted + """Test BlockRemoved events emitted on eviction carry the group_idx extracted from the evicted block's BlockHashWithGroupId via get_group_id(). """ block_size = 4 @@ -3639,7 +3624,8 @@ def test_block_removed_event_group_idx(group_id: int): def test_emit_cached_block_events(): """emit_cached_block_events emits one BlockStored for already-cached (reused) prefix blocks, carrying the correct group_idx / - parent_block_hash / token_ids, and without mutating block state.""" + parent_block_hash / token_ids, and without mutating block state. + """ block_size = 4 num_cached_blocks = 3 kv_cache_group_id = 1 @@ -3752,7 +3738,8 @@ def test_emit_cached_block_events_zero_cached(): def test_eagle_enabled_removes_last_block(): """Verify Eagle does NOT remove blocks when request - length is divisible by block size.""" + length is divisible by block size. + """ block_size = 16 manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks=10), @@ -4155,7 +4142,8 @@ def test_hybrid_cache_blocks_swa_tail_window_only(): (its right-to-left scan stops once a contiguous match is found). Blocks earlier in the segment can never serve a hit, so ``HybridKVCacheCoordinator.cache_blocks`` should skip them rather than - polluting the prefix-cache hash map.""" + polluting the prefix-cache hash map. + """ block_size = 8 # Full attn block_size=32, SWA block_size=8, sw=8 -> lcm=32. # tail = ceil(7/8) = 1; per_segment = 32/8 = 4. @@ -4226,7 +4214,8 @@ def test_hybrid_cache_blocks_clamped_to_lcm(): Chunks past the last lcm-aligned boundary can never participate in a cache hit (find_longest_cache_hit always returns lcm-aligned hits), so caching them only pollutes the prefix-cache hash map and keeps blocks - on the LRU list that could otherwise return to the free pool.""" + on the LRU list that could otherwise return to the free pool. + """ block_size = 16 # Full attn block_size=32, SWA block_size=16 -> lcm=32. kv_cache_config = KVCacheConfig( @@ -4361,7 +4350,8 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(): ) def test_hybrid_local_kv_retention_interval_rejects_invalid(interval, expected_match): """A retention interval that is negative or not a multiple of - scheduler_block_size errors out at construction time.""" + scheduler_block_size errors out at construction time. + """ block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -4831,7 +4821,8 @@ def test_block_lookup_cache_multi_blocks_per_key(): def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): """Hybrid full+SWA model with a pool sized at the startup minimum should admit a prompt longer than the SWA cap, because SlidingWindowManager - recycles blocks during chunked prefill (issue #39734).""" + recycles blocks during chunked prefill (issue #39734). + """ block_size = 16 sliding_window = 4 * block_size # 64 tokens max_num_batched_tokens = 8 * block_size # 128 tokens @@ -4926,7 +4917,8 @@ def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk( def test_can_fit_full_sequence_full_attention_still_gates_oversized(): """The cap only loosens the SWA group; a prompt that exceeds the - full-attention pool capacity must still be rejected.""" + full-attention pool capacity must still be rejected. + """ block_size = 16 sliding_window = 4 * block_size max_num_batched_tokens = 8 * block_size @@ -5045,7 +5037,8 @@ def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheB without removing them. These ref_cnt==0 blocks stand in for evictable cache-hit blocks left behind by a previous (e.g. preempted) request, and sitting at the head guarantees a later group's external ``get_new_blocks`` - would contend for them on unpatched code (issue #33775).""" + would contend for them on unpatched code (issue #33775). + """ blocks: list[KVCacheBlock] = [] head = manager.block_pool.free_block_queue.fake_free_list_head for _ in range(num_blocks): @@ -5056,7 +5049,8 @@ def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheB def _assert_no_double_allocation(manager: KVCacheManager, req_id: str) -> None: """No physical block may be handed out twice across groups, and every - block referenced by the request must have a live ref_cnt.""" + block referenced by the request must have a live ref_cnt. + """ block_ids = manager.get_blocks(req_id).get_block_ids() flat = [block_id for group in block_ids for block_id in group] assert len(set(flat)) == len(flat), "Block IDs are not unique across groups" @@ -5083,7 +5077,8 @@ def _cross_group_cache_hit( ) -> Request: """Allocate ``req_id`` with a per-group local prefix hit plus external (connector) computed tokens, driving the coordinator's two-phase path. - Returns the allocated request so callers can free it (e.g. to preempt).""" + Returns the allocated request so callers can free it (e.g. to preempt). + """ block_size = _two_phase_block_size(manager) hit_blocks = _take_free_blocks(manager, num_groups * local_blocks_per_group) cache_hit = KVCacheBlocks( @@ -5178,7 +5173,8 @@ def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(): """Dense retention: freeing an SWA request must place its uncached scratch blocks at the front of the free queue (recycled first) and keep its cached checkpoint blocks at the back (retained for prefix - hits). This split is always-on, independent of the retention interval.""" + hits). This split is always-on, independent of the retention interval. + """ block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -5289,7 +5285,8 @@ def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs) def test_pure_swa_retention_interval_caches_sparse_tails(): """Sparse retention must work for a pure-SWA single-group model, not just hybrid models: only the per-interval tails plus the latest replay tail are - cached, and a replay still hits the latest replayable boundary.""" + cached, and a replay still hits the latest replayable boundary. + """ block_size = 16 manager = _make_pure_swa_manager( block_size, sliding_window=block_size, retention_interval=64 @@ -5361,7 +5358,8 @@ def test_pure_swa_retention_latest_only(): def test_pure_swa_dense_retention_caches_all(): """With retention set to ``None``, a pure-SWA model keeps dense behavior: - every block boundary is a potential hit, so all blocks are cached.""" + every block boundary is a potential hit, so all blocks are cached. + """ block_size = 16 manager = _make_pure_swa_manager(block_size, sliding_window=block_size) @@ -5390,7 +5388,8 @@ def test_mamba_reachable_block_mask_sparsifies_retention(): """Mamba state-snapshot retention: with a configured retention interval, the manager keeps one cached state per interval-sized segment (plus the latest replay boundary) instead of a snapshot per block, which is what - lets a small attention block_size avoid Mamba dominating the KV pool.""" + lets a small attention block_size avoid Mamba dominating the KV pool. + """ from vllm.v1.core.single_type_kv_cache_manager import MambaManager block_size = 16 @@ -5428,7 +5427,8 @@ def test_mamba_reachable_block_mask_pins_shared_prefix(): """A Marconi-detected shared prefix (``shared_prefix_boundary``) lands before ``num_prompt`` so the replay-boundary rule alone would drop it. The mask must pin the single state block ending on that boundary so sparse retention does - not defeat cross-request shared-prefix reuse.""" + not defeat cross-request shared-prefix reuse. + """ from vllm.v1.core.single_type_kv_cache_manager import MambaManager block_size = 16 @@ -5476,7 +5476,8 @@ def test_mamba_shared_prefix_survives_zero_retention(): detection) keeps its Mamba state block cached under ``prefix_cache_retention_interval=0``, which otherwise retains only the end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request.""" + before ``num_prompt``) would be recomputed by every sharing request. + """ block_size = 16 # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. @@ -5518,7 +5519,8 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): detecting request must stay reusable by a later request under ``prefix_cache_retention_interval=0``. Without the pin the junction is masked out and the later request misses; with it (and under dense) the reuse - is preserved.""" + is preserved. + """ block_size = 16 def last_req_hit(retention, pin): @@ -5601,7 +5603,8 @@ def reachable_blocks( def test_swa_reachable_block_mask_pins_shared_prefix(): """SWA analog of the Mamba pin: the shared-prefix junction must keep the ``need``-block sliding-window tail ending on that boundary (not a single - block), so a windowed hit can land there under sparse retention.""" + block), so a windowed hit can land there under sparse retention. + """ from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager block_size = 16 @@ -5646,7 +5649,8 @@ def retained(retention, boundary, window, end_block=16): def test_swa_reachable_block_mask_with_dcp_scaling(): """DCP shards each block's KV across ranks, scaling the effective block size. Verify that dcp_world_size > 1 scales block size in reachability calculations, - producing different masks than dcp_world_size=1.""" + producing different masks than dcp_world_size=1. + """ from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager block_size = 16 @@ -5705,7 +5709,8 @@ def test_swa_reachable_block_mask_sub_block_alignment_is_dense( multiple of the DCP-scaled block size) cannot be represented exactly. This happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the full-attention chunk size and need not divide the SWA block size. The mask - must fall back to dense (None) rather than raise.""" + must fall back to dense (None) rather than raise. + """ from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager spec = SlidingWindowSpec( @@ -5733,7 +5738,8 @@ def test_mamba_reachable_block_mask_ignores_dcp(): """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a state block spans kv_cache_spec.block_size tokens regardless of DCP. The mask must not scale by dcp_world_size, so retention granularity is - identical for any DCP world size.""" + identical for any DCP world size. + """ from vllm.v1.core.single_type_kv_cache_manager import MambaManager spec = MambaSpec( @@ -5767,7 +5773,8 @@ def get_mask(dcp_world_size): def test_mamba_reachable_block_mask_large_dcp_stays_sparse(): """A large DCP world size must not scale the Mamba block size, so it can - never collapse sparse retention into dense caching.""" + never collapse sparse retention into dense caching. + """ from vllm.v1.core.single_type_kv_cache_manager import MambaManager spec = MambaSpec( @@ -5802,7 +5809,8 @@ def test_swa_shared_prefix_reuse_under_zero_retention(): """SWA cross-request analog: a partial shared prefix's sliding-window tail must stay reusable under ``prefix_cache_retention_interval=0``. Without the pin the junction window is masked out and a later request misses; with - it (and under dense) reuse is preserved.""" + it (and under dense) reuse is preserved. + """ block_size = 16 def last_req_hit(retention, pin): diff --git a/tests/v1/core/test_repetition_detection.py b/tests/v1/core/test_repetition_detection.py index aae6e3b70cae..6814a6281433 100644 --- a/tests/v1/core/test_repetition_detection.py +++ b/tests/v1/core/test_repetition_detection.py @@ -14,10 +14,10 @@ class TestCheckSequenceRepetition: - """Unit tests for the check_sequence_repetition function""" + """Unit tests for the check_sequence_repetition function.""" def test_simple_repetition_detected(self): - """Test detection of simple repetitive patterns""" + """Test detection of simple repetitive patterns.""" token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 3] params = RepetitionDetectionParams( max_pattern_size=3, @@ -27,7 +27,7 @@ def test_simple_repetition_detected(self): assert check_sequence_repetition(token_ids, params) def test_repetition_below_min_count(self): - """Test that pattern below min_count is not detected""" + """Test that pattern below min_count is not detected.""" token_ids = [1, 2, 3, 1, 2, 3] params = RepetitionDetectionParams( max_pattern_size=3, @@ -37,7 +37,7 @@ def test_repetition_below_min_count(self): assert not check_sequence_repetition(token_ids, params) def test_two_token_pattern(self): - """Test detection of 2-token patterns""" + """Test detection of 2-token patterns.""" token_ids = [1, 2, 1, 2, 1, 2, 1, 2] params = RepetitionDetectionParams( max_pattern_size=5, @@ -47,7 +47,7 @@ def test_two_token_pattern(self): assert check_sequence_repetition(token_ids, params) def test_no_repetition_varied_sequence(self): - """Test that non-repetitive sequences are not flagged""" + """Test that non-repetitive sequences are not flagged.""" token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9] params = RepetitionDetectionParams( max_pattern_size=5, @@ -57,7 +57,7 @@ def test_no_repetition_varied_sequence(self): assert not check_sequence_repetition(token_ids, params) def test_partial_repetition_not_detected(self): - """Test that incomplete repetitions are not detected""" + """Test that incomplete repetitions are not detected.""" token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 4] params = RepetitionDetectionParams( max_pattern_size=3, @@ -67,7 +67,7 @@ def test_partial_repetition_not_detected(self): assert not check_sequence_repetition(token_ids, params) def test_empty_token_list(self): - """Test with empty token list""" + """Test with empty token list.""" params = RepetitionDetectionParams( max_pattern_size=3, min_pattern_size=2, @@ -76,19 +76,19 @@ def test_empty_token_list(self): assert not check_sequence_repetition([], params) def test_detection_disabled_max_size_zero(self): - """Test that zero max_pattern_size disables detection""" + """Test that zero max_pattern_size disables detection.""" token_ids = [1, 2, 1, 2, 1, 2] params = RepetitionDetectionParams() assert not check_sequence_repetition(token_ids, params) def test_invalid_min_count(self): - """Test that min_count < 2 returns False""" + """Test that min_count < 2 returns False.""" token_ids = [1, 2, 1, 2] params = RepetitionDetectionParams() assert not check_sequence_repetition(token_ids, params) def test_repetition_at_end_of_sequence(self): - """Test detection when repetition occurs at the end""" + """Test detection when repetition occurs at the end.""" token_ids = [1, 2, 3, 4, 5, 6, 5, 6, 5, 6] params = RepetitionDetectionParams( max_pattern_size=3, @@ -98,7 +98,7 @@ def test_repetition_at_end_of_sequence(self): assert check_sequence_repetition(token_ids, params) def test_large_pattern_many_repetitions(self): - """Test large pattern repeated many times""" + """Test large pattern repeated many times.""" token_ids = [1, 2, 3, 4, 5, 6, 7, 8] * 5 params = RepetitionDetectionParams( max_pattern_size=10, @@ -114,10 +114,10 @@ def test_large_pattern_many_repetitions(self): class TestRepetitionDetectionIntegration: - """Integration tests for repetition detection in check_stop""" + """Integration tests for repetition detection in check_stop.""" def test_basic_repetition_stops_generation(self): - """Test that repetition is detected and stops generation""" + """Test that repetition is detected and stops generation.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( @@ -138,7 +138,7 @@ def test_basic_repetition_stops_generation(self): assert request.stop_reason == "repetition_detected" def test_detection_disabled_no_stop(self): - """Test that disabled detection doesn't stop generation""" + """Test that disabled detection doesn't stop generation.""" params = SamplingParams( max_tokens=100, ) @@ -152,7 +152,7 @@ def test_detection_disabled_no_stop(self): assert not check_stop(request, max_model_len=1024) def test_repetition_respects_min_tokens(self): - """Test that repetition detection respects min_tokens""" + """Test that repetition detection respects min_tokens.""" params = SamplingParams( min_tokens=10, max_tokens=100, @@ -172,7 +172,7 @@ def test_repetition_respects_min_tokens(self): assert not check_stop(request, max_model_len=1024) def test_no_repetition_continues_generation(self): - """Test that non-repetitive tokens don't stop generation""" + """Test that non-repetitive tokens don't stop generation.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( @@ -191,7 +191,7 @@ def test_no_repetition_continues_generation(self): assert not check_stop(request, max_model_len=1024) def test_pattern_at_size_boundary(self): - """Test detection at exact pattern size boundary""" + """Test detection at exact pattern size boundary.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( @@ -211,7 +211,7 @@ def test_pattern_at_size_boundary(self): assert request.status == RequestStatus.FINISHED_REPETITION def test_multiple_pattern_sizes_checked(self): - """Test that function checks pattern sizes in range""" + """Test that function checks pattern sizes in range.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( @@ -231,7 +231,7 @@ def test_multiple_pattern_sizes_checked(self): assert request.status == RequestStatus.FINISHED_REPETITION def test_eos_takes_precedence_over_repetition(self): - """Test that EOS token stops before repetition check""" + """Test that EOS token stops before repetition check.""" params = SamplingParams( max_tokens=100, stop_token_ids=[999], @@ -252,7 +252,7 @@ def test_eos_takes_precedence_over_repetition(self): assert request.status == RequestStatus.FINISHED_STOPPED def test_min_pattern_size_filters_small_patterns(self): - """Test that min_pattern_size filters out smaller patterns""" + """Test that min_pattern_size filters out smaller patterns.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( @@ -271,7 +271,7 @@ def test_min_pattern_size_filters_small_patterns(self): assert not check_stop(request, max_model_len=1024) def test_high_repetition_threshold(self): - """Test that high min_count requires many repetitions""" + """Test that high min_count requires many repetitions.""" params = SamplingParams( max_tokens=100, repetition_detection=RepetitionDetectionParams( diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 061c825fc895..5774ec014698 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -295,7 +295,7 @@ def test_schedule_multimodal_requests(): def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders(): - """Async scheduling + PP: allow multi-step in-flight scheduling per request""" + """Async scheduling + PP: allow multi-step in-flight scheduling per request.""" scheduler = create_scheduler(async_scheduling=True, pipeline_parallel_size=2) (req,) = create_requests(num_requests=1, num_tokens=8) scheduler.add_request(req) @@ -538,7 +538,8 @@ def test_throttle_defers_inflight_prefill_chunk(): """DP prefill balancing throttles ALL prefill compute on a throttled step, not just new admissions: an in-progress (chunked) prefill already in the running queue is also deferred, so the step runs decode-only, while a - separate decode keeps being scheduled.""" + separate decode keeps being scheduled. + """ scheduler = create_scheduler( max_num_seqs=16, max_num_batched_tokens=50, enable_chunked_prefill=True ) @@ -592,7 +593,8 @@ def test_throttle_capacity_bound_guard_admits(): """Saturation guard: if a cadence-aligned release step cannot drain the waiting prefill queue (it ran out of token budget), the throttle backs off on the next step so the backlog cannot grow into a TTFT avalanche -- prefills are - admitted even though throttle_prefills is set.""" + admitted even though throttle_prefills is set. + """ scheduler = create_scheduler( max_num_seqs=16, max_num_batched_tokens=200, enable_chunked_prefill=True ) @@ -785,7 +787,7 @@ def test_update_from_output_routes_sampling_masks_by_request(): def test_stop_via_update_from_output(): - """Test stopping behavior through update_from_output""" + """Test stopping behavior through update_from_output.""" scheduler = create_scheduler(num_speculative_tokens=1) # Test case 1: Stop on EOS token @@ -1291,7 +1293,8 @@ def test_pending_hisparse_reclamation_defers_preemption(monkeypatch): def test_prefix_cache_query_not_inflated_by_connector_defer(): """The GPU prefix-cache query is recorded at admission, so a request the - connector defers several times is counted once, not once per retry.""" + connector defers several times is counted once, not once per retry. + """ num_defers_before_matching = 3 scheduler = create_scheduler( enable_prefix_caching=True, @@ -1319,7 +1322,8 @@ def test_prefix_cache_query_not_inflated_by_connector_defer(): def test_preemption_re_records_prefix_cache_query(): """A preempted request re-enters the lookup on resume, so its recomputation - is counted again into the preempted stats.""" + is counted again into the preempted stats. + """ scheduler = create_scheduler(enable_prefix_caching=True) request = create_requests(num_requests=1)[0] scheduler.add_request(request) @@ -1355,7 +1359,8 @@ def test_preemption_re_records_prefix_cache_query(): def test_prefix_cache_stats_not_recorded_when_caching_disabled(): """With prefix caching off there is no local lookup, so admitting a request - records no phantom miss.""" + records no phantom miss. + """ scheduler = create_scheduler(enable_prefix_caching=False) for request in create_requests(num_requests=2): scheduler.add_request(request) @@ -1369,7 +1374,8 @@ def test_prefix_cache_stats_not_recorded_when_caching_disabled(): def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request(): """A real cache hit rejected once by allocate_slots and admitted on the next - step is counted exactly once, hits included.""" + step is counted exactly once, hits included. + """ block_size = 16 scheduler = create_scheduler( enable_prefix_caching=True, @@ -2130,7 +2136,6 @@ def _assert_right_scheduler_output( expected_num_scheduled_tokens: int, ): """Check if SchedulerOutput is correct after remote KV cache hit.""" - # We should inject the kv_connector_metadata. assert len(output.kv_connector_metadata.requests) == num_requests @@ -2148,7 +2153,6 @@ def _assert_right_kv_cache_manager( num_total_blocks: int, ): """Check whether KVCacheManager is correct after allocate.""" - # Make sure the request stats are right. EXPECTED_TOTAL_BLOCKS = num_tokens // block_size for req in requests: @@ -2179,7 +2183,6 @@ def _step_until_done( model_runner_output: ModelRunnerOutput, ): """Loop over schedule(), update_from_output() until finished.""" - all_finished = False _ = scheduler.update_from_output(output, model_runner_output) while not all_finished: @@ -2207,7 +2210,6 @@ def _num_waiting_requests(scheduler: Scheduler) -> int: def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): """Cycle requests through a KV transfer cycle.""" - # Requests should first transition to WAITING_FOR_REMOTE_KVS output = scheduler.schedule() assert _num_waiting_requests(scheduler) == len(req_ids) @@ -2291,11 +2293,9 @@ def test_has_sync_kv_loads( @pytest.mark.parametrize("is_async", [False, True]) def test_kv_connector_basic(is_async: bool): - """ - Test whether Scheduler with KVConnector schedules tokens, allocates + """Test whether Scheduler with KVConnector schedules tokens, allocates memory, and cleans up requests as expected under normal operation. """ - # Setup Scheduler. BLOCK_SIZE = 16 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2 @@ -2420,11 +2420,9 @@ def test_kv_connector_basic(is_async: bool): @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.parametrize("local_cache_hits", [False, True]) def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool): - """ - Verify connector prefix cache metrics are updated + """Verify connector prefix cache metrics are updated correctly when the scheduler processes requests with KV connector hits. """ - BLOCK_SIZE = 16 if local_cache_hits: NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2 # 32 tokens @@ -2540,11 +2538,9 @@ def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool): "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")] ) def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role): - """ - Test whether scheduler with KVConnector is able to handle + """Test whether scheduler with KVConnector is able to handle unable to allocate (run out of blocks in allocate_slots(). """ - # Setup Scheduler With Mock External Cache Hit. BLOCK_SIZE = 4 NUM_BLOCKS = 10 @@ -2627,11 +2623,9 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role): def test_kv_connector_handles_preemption( is_async, use_ec_connector, ec_role, use_v2_model_runner ): - """ - Test whether scheduler with KVConnector is able to handle + """Test whether scheduler with KVConnector is able to handle unable to allocate (run out of blocks in allocate_slots(). """ - # Setup Scheduler With Mock External Cache Hit. BLOCK_SIZE = 2 # NOTE: there is 1 null block, so this is 6 blocks. @@ -2842,7 +2836,6 @@ def assert_scheduler_empty(scheduler: Scheduler): def test_memory_leak(): """Test that we do not have a memory leak.""" - scheduler = create_scheduler(enable_prefix_caching=True) NUM_REQUESTS = 5 @@ -2899,6 +2892,7 @@ def create_scheduler_with_priority( Returns: {class}`Scheduler` instance with priority scheduling + """ model_config = ModelConfig( model=model, @@ -3100,7 +3094,8 @@ def create_requests_with_priority( def test_priority_scheduling_basic_ordering(): """Test that requests are scheduled in priority order - (lower value = higher priority).""" + (lower value = higher priority). + """ scheduler = create_scheduler_with_priority() # Create requests with different priorities @@ -3129,7 +3124,8 @@ def test_priority_scheduling_basic_ordering(): def test_priority_scheduling_arrival_time_tiebreaker(): """Test that arrival time is used - as tiebreaker when priorities are equal.""" + as tiebreaker when priorities are equal. + """ scheduler = create_scheduler_with_priority() # Create requests with same priority but different arrival times @@ -3286,7 +3282,8 @@ def test_priority_scheduling_preemption(): def test_priority_scheduling_no_preemption_when_space_available(): """Test that preemption doesn't happen - when there's space for new requests.""" + when there's space for new requests. + """ scheduler = create_scheduler_with_priority( max_num_seqs=3, # Allow 3 concurrent requests max_num_batched_tokens=200, # Sufficient token budget @@ -3339,7 +3336,8 @@ def test_priority_scheduling_no_preemption_when_space_available(): def test_priority_scheduling_preemption_victim_selection(): """Test that the correct victim is selected for - preemption based on priority and arrival time.""" + preemption based on priority and arrival time. + """ # This test verifies the priority-based victim selection logic # by checking the waiting queue order after adding requests with different # priorities @@ -3459,7 +3457,8 @@ def test_priority_scheduling_waiting_queue_order(): def test_priority_scheduling_fcfs_fallback(): """Test that FCFS behavior is maintained when all - requests have same priority.""" + requests have same priority. + """ scheduler = create_scheduler_with_priority() # Create requests with same priority but different arrival times @@ -3528,7 +3527,8 @@ def test_priority_scheduling_with_limited_slots(): def test_priority_scheduling_heap_property(): """Test that the waiting queue maintains heap - property for priority scheduling.""" + property for priority scheduling. + """ scheduler = create_scheduler_with_priority( max_num_seqs=1, # Only one request can run at a time ) @@ -3753,7 +3753,8 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv( use_ec_connector, ec_role, use_v2_model_runner ): """Test that priority scheduling preempts lower priority requests - when out of KV cache space.""" + when out of KV cache space. + """ # Create scheduler with very limited memory to force preemption scheduler = create_scheduler_with_priority( max_num_seqs=2, # Allow multiple requests @@ -3898,7 +3899,8 @@ def test_chunked_prefill_disabled_for_encoder_decoder( enable_chunked_prefill: bool, is_encoder_decoder: bool, expect_enabled: bool ) -> None: """Validate that chunked prefill is appropriately disabled for - encoder-decoder models.""" + encoder-decoder models. + """ scheduler_config = SchedulerConfig( enable_chunked_prefill=enable_chunked_prefill, is_encoder_decoder=is_encoder_decoder, @@ -3927,7 +3929,8 @@ def _validate_chunked_prefill_settings_for_encoder_decoder( scheduler_config: SchedulerConfig, is_encoder_decoder: bool, expect_enabled: bool ) -> None: """Validate chunked prefill settings in the scheduler config for - encoder-decoder models.""" + encoder-decoder models. + """ assert scheduler_config.enable_chunked_prefill is expect_enabled if is_encoder_decoder: # Encoder-decoder models should automatically disable chunked multimodal @@ -3987,7 +3990,7 @@ def _assert_right_ec_connector_metadata( output: SchedulerOutput, mm_features_list: list[MultiModalFeatureSpec], ): - """Verify that ECConnector metadata EXACTLY matches the input MM data""" + """Verify that ECConnector metadata EXACTLY matches the input MM data.""" # Get the connector metadata metadata = output.ec_connector_metadata @@ -4017,8 +4020,8 @@ def _assert_right_encoder_inputs( ): """Verify that requests/mm_hashes should (not) in scheduled encoder input If check_exist is False, this function returns True - if requests are NOT in encoder inputs""" - + if requests are NOT in encoder inputs + """ # Get the scheduled encoder inputs # NOTE: scheduled_encoder_inputs is a dictionary with request id as key scheduled_encoder_inputs = output.scheduled_encoder_inputs @@ -4179,7 +4182,8 @@ def test_ec_connector_text_only_request(use_kv_connector): @pytest.mark.parametrize("use_kv_connector", [False, True]) def test_ec_connector_cache_hit_external_load(use_kv_connector): """Test ec_consumer loads from external cache when hit. - A normal basic operation for EPD disaggrgation""" + A normal basic operation for EPD disaggrgation + """ scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", enable_prefix_caching=True, @@ -4527,11 +4531,9 @@ def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector): @pytest.mark.parametrize("use_kv_connector", [False, True]) def test_ec_connector_unable_to_allocate(use_kv_connector): - """ - Test whether scheduler with ECConnector is able to handle + """Test whether scheduler with ECConnector is able to handle unable to allocate (run out of blocks). """ - # Setup Scheduler With Mock External Cache Hit. BLOCK_SIZE = 4 NUM_BLOCKS = 10 @@ -4628,7 +4630,8 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption( cache_exist, use_kv_connector ): """Test that priority scheduling preempts lower priority requests - when out of KV cache space.""" + when out of KV cache space. + """ # Create scheduler with very limited memory to force preemption scheduler = create_scheduler_with_priority( model="llava-hf/llava-1.5-7b-hf", @@ -4860,8 +4863,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption( @pytest.mark.parametrize("use_kv_connector", [False, True]) def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connector): - """ - Scenario: + """Scenario: - Encoder cache size: 32 - Request A: 1 feature (12 tokens) → NOT cached remotely. - Request B: 3 features (3 x 10 tokens) → ALL cached remotely. @@ -5298,7 +5300,8 @@ def update_connector_output( def test_ec_connector_update_connector_output_called(): """Test that worker-side EC connector output is forwarded to the - EC connector's update_connector_output hook.""" + EC connector's update_connector_output hook. + """ scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", use_ec_connector=True, @@ -5503,7 +5506,8 @@ def test_variable_length_cross_attn_block_allocation(): def test_cross_attn_blocks_not_over_allocated(): """Test that cross-attention blocks are not over-allocated compared to - what each request actually needs.""" + what each request actually needs. + """ from math import ceil block_size = 16 @@ -5599,7 +5603,8 @@ def test_cross_attn_blocks_not_under_allocated(): def test_cross_attn_zero_blocks_without_encoder_inputs(): """Test that requests without encoder inputs get zero cross-attention - blocks, even when the scheduler is configured for encoder-decoder.""" + blocks, even when the scheduler is configured for encoder-decoder. + """ block_size = 16 scheduler = _create_encoder_decoder_scheduler(block_size=block_size) @@ -5704,7 +5709,8 @@ def test_free_encoder_inputs_respects_unconfirmed_placeholders(): evicted entry and crashes the engine with "Encoder cache miss". The scheduler must retain the input until the *confirmed* progress (num_computed_tokens - num_output_placeholders) passes the range end, so - that no pending rejection can rewind into the range.""" + that no pending rejection can rewind into the range. + """ scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -5815,7 +5821,8 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead(): position so the drafter's +1 look-ahead mm-embedding gather (which reads one position past the target's computed range) still finds it cached. This is the primary mechanism that prevents the drafter "Encoder cache miss"; the - worker-side token-embedding fallback is only a backstop.""" + worker-side token-embedding fallback is only a backstop. + """ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") # create_scheduler only builds ngram spec configs; force the eagle path that # _free_encoder_inputs keys off (its read-ahead deferral). @@ -5845,7 +5852,8 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead(): def test_free_encoder_inputs_unchanged_without_spec_decode(): """Without speculative decoding, encoder inputs are freed as soon as - num_computed_tokens passes the placeholder range, as before.""" + num_computed_tokens passes the placeholder range, as before. + """ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") mm_positions = [[PlaceholderRange(offset=50, length=100)]] request = create_requests( @@ -5876,7 +5884,8 @@ def test_encoder_cache_retained_across_preemption_and_resume(): encoder input must pull the still-cached entry back out of `freeable` without scheduling a recompute, keeping the scheduler and worker consistent. The spec-rollback retention margin does not gate this path, - so it is covered separately here.""" + so it is covered separately here. + """ scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -5920,7 +5929,8 @@ def test_encoder_cache_recomputed_when_evicted_during_preemption(): pressure before it resumes, the scheduler reports the mm_hash as freed (so the worker drops it) and a resume must schedule a recompute rather than assume the worker still holds it. check_and_update_cache must - return False so the encoder input is re-scheduled.""" + return False so the encoder input is re-scheduled. + """ scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -6571,7 +6581,8 @@ def _decode_ready_request(scheduler): def test_update_draft_token_ids_strips_ngram_padding(): """ngram_gpu pads unfilled draft slots with -1; manager validation must - strip them before grammar.validate_tokens (which otherwise raises).""" + strip them before grammar.validate_tokens (which otherwise raises). + """ scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) @@ -6590,7 +6601,8 @@ def test_update_draft_token_ids_strips_ngram_padding(): def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to manager validation is stripped).""" + is preserved (only the input to manager validation is stripped). + """ scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py index 737796a4bf7d..beaf5203f24c 100644 --- a/tests/v1/core/test_swa_inflight_window_free.py +++ b/tests/v1/core/test_swa_inflight_window_free.py @@ -97,7 +97,8 @@ def test_num_in_flight_tokens_accounting(): def test_swa_free_waits_for_in_flight_step(): """Async: out-of-window blocks stay allocated until the step that still - reads them has been processed.""" + reads them has been processed. + """ scheduler = _create_swa_scheduler(async_scheduling=True) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -132,7 +133,8 @@ def test_swa_free_waits_for_in_flight_step(): def test_swa_free_immediate_when_sync(): """Sync: no in-flight step at schedule time, frees happen at the first - decode allocation as before.""" + decode allocation as before. + """ scheduler = _create_swa_scheduler(async_scheduling=False) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -171,7 +173,8 @@ def test_swa_admission_cap_accounts_for_overlapping_batches(): def test_chunked_local_free_waits_for_in_flight_step(): """Chunked-local attention frees whole chunks left of the current one, and is exposed to the same load-WAR: with async scheduling those chunks must - stay allocated until the in-flight step that still reads them settles.""" + stay allocated until the in-flight step that still reads them settles. + """ scheduler = _create_chunked_scheduler(async_scheduling=True) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -226,7 +229,8 @@ def test_connector_finish_frees_on_settled_basis(): is handed to a KV connector (simple CPU offload / NIXL store), must use the same processed-token basis. Otherwise the connector reads/hands off a block the still-in-flight step is optimistically counted as done with, which is - the load-path WAR this fix closes.""" + the load-path WAR this fix closes. + """ scheduler = create_scheduler( block_size=BLOCK_SIZE, async_scheduling=True, diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 3c4a8960a3fa..28d1f71ac492 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -89,6 +89,7 @@ def create_scheduler( Returns: {class}`Scheduler` instance + """ model_config = ModelConfig( model=model, diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index e792d68d89e0..3e87770d4e82 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for the breakable cudagraph primitives. -""" +"""Unit tests for the breakable cudagraph primitives.""" from __future__ import annotations @@ -82,7 +80,8 @@ def forward_fn(cg_mode): @pytest.fixture(autouse=True) def _reset_breakable_tls(): """Defensively clear thread-local capture state between tests so a - failure in one test can't leak "nested capture" errors into the next.""" + failure in one test can't leak "nested capture" errors into the next. + """ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture BreakableCUDAGraphCapture._tls.active = None @@ -292,7 +291,8 @@ def eager_break_op(): def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): """Verify @eager_break_during_capture correctly routes through add_eager when inside a capture context, and runs straight through - when there's no active capture.""" + when there's no active capture. + """ from vllm.compilation.breakable_cudagraph import ( BreakableCUDAGraphCapture, eager_break_during_capture, @@ -554,7 +554,8 @@ def _mock_vllm_config(): def test_wrapper_falls_through_on_runtime_mode_mismatch(): """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing - nothing, so an outer CUDAGraphWrapper can own them.""" + nothing, so an outer CUDAGraphWrapper can own them. + """ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index c10835821f58..ddf72c3b1141 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -376,7 +376,6 @@ def _run_and_monitor_call( wrapper, input_tensor, runtime_mode, batch_descriptor, vllm_config ): """Helper to run a single call and monitor the action.""" - with ( patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context, patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable, diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 4a86a5826bc0..e370f87833eb 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -31,8 +31,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( backend, rms_norm_impl, ): - """ - Ensures that the same request (the 'needle' prompt) yields identical output + """Ensures that the same request (the 'needle' prompt) yields identical output whether run alone (bs=1) or mixed into a larger batch (e.g., bs=64), using the high-level v1 LLM() API only (no manual batching). @@ -51,6 +50,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( to produce a more random-sounding phrase, yet remain deterministic by seed. - Keep max_tokens and max_model_len bounded for speed and memory use. + """ seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) @@ -395,8 +395,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( BACKENDS, ) def test_simple_generation(backend): - """ - Simple test that runs the model with a basic prompt and prints the output. + """Simple test that runs the model with a basic prompt and prints the output. Useful for quick smoke testing and debugging. """ model = TEST_MODEL @@ -447,8 +446,7 @@ def test_simple_generation(backend): def test_logprobs_without_batch_invariance_should_fail( backend, monkeypatch: pytest.MonkeyPatch ): - """ - This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN. + """This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN. It DISABLES batch invariance mode and expects to see non-deterministic behavior between BS=1 and BS=N runs. This demonstrates that batch invariance is actually doing something useful. @@ -660,8 +658,7 @@ def test_logprobs_without_batch_invariance_should_fail( def test_decode_logprobs_match_prefill_logprobs( backend, ): - """ - Test that verifies decode logprobs match prefill logprobs. + """Test that verifies decode logprobs match prefill logprobs. For each decoded token at position i: 1. Run decode to generate N tokens and collect their logprobs @@ -927,8 +924,7 @@ def LLM_with_max_seqs( attention_config: dict | None = None, kernel_config: dict | None = None, ) -> LLM: - """ - Helper to construct an LLM with a specific max_num_seqs (batch-size limit) + """Helper to construct an LLM with a specific max_num_seqs (batch-size limit) using the high-level v1 LLM API, while constraining memory usage. """ extra_kwargs: dict = {} diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index 086490526e50..af395e74e7f1 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test batch-invariant matmul against torch.matmul for various shape combinations. +"""Test batch-invariant matmul against torch.matmul for various shape combinations. Tests correctness (matches torch.matmul) and batch invariance (result for one item doesn't change based on other items in the batch). @@ -53,9 +52,7 @@ ) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_matmul_correctness(a_shape, b_shape, dtype): - """ - Compare matmul_batch_invariant against torch.matmul for various shapes. - """ + """Compare matmul_batch_invariant against torch.matmul for various shapes.""" device = torch.device(DEVICE_TYPE) torch.manual_seed(42) @@ -87,11 +84,9 @@ def test_matmul_correctness(a_shape, b_shape, dtype): @skip_unsupported @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_matmul_batch_invariance(dtype): - """ - Verify that the result for one item is bitwise identical regardless + """Verify that the result for one item is bitwise identical regardless of what other items are in the batch. """ - device = torch.device(DEVICE_TYPE) torch.manual_seed(42) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index de80e4918c1e..dcba87dc4927 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/tests/v1/determinism/test_online_batch_invariance.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -HTTP-based batch invariance test: send requests to a running +"""HTTP-based batch invariance test: send requests to a running vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs). Environment variables: diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index a9a15d1bf083..d0a4c5c5b2b4 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -41,8 +41,7 @@ def test_rms_norm_batch_invariant_vs_reference( eps: float, seed: int, ): - """ - Compare batch-invariant Triton RMS norm against a PyTorch reference. + """Compare batch-invariant Triton RMS norm against a PyTorch reference. Tests that the Triton-based batch-invariant RMS norm produces numerically equivalent results to an independent implementation across configurations. @@ -90,9 +89,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( n_extra: int, seed: int, ): - """ - Test the batch-invariant fused residual-add + RMSNorm helper directly. - """ + """Test the batch-invariant fused residual-add + RMSNorm helper directly.""" device = torch.device(DEVICE_TYPE) torch.manual_seed(seed) @@ -278,7 +275,8 @@ def test_rms_norm_per_block_quant_batch_invariant( hidden_size: int, dtype: torch.dtype, seed: int ): """C++ ``rms_norm_per_block_quant`` must be batch invariant across the - block threshold (compiled fp8 block-quant path; block pinned to 512).""" + block threshold (compiled fp8 block-quant path; block pinned to 512). + """ import vllm._custom_ops as ops device = torch.device(DEVICE_TYPE) @@ -313,8 +311,7 @@ def per_block_quant(x): def test_rms_norm_3d_input( default_vllm_config, batch_size: int, seq_len: int, hidden_size: int ): - """ - Test RMS norm with 3D input tensors (batch, seq_len, hidden_size). + """Test RMS norm with 3D input tensors (batch, seq_len, hidden_size). Ensures that the batch-invariant RMS norm correctly handles multi-dimensional inputs that are common in transformer models. @@ -349,8 +346,7 @@ def test_rms_norm_3d_input( @skip_if_not_cuda def test_rms_norm_numerical_stability(default_vllm_config): - """ - Test RMS norm numerical stability with extreme values. + """Test RMS norm numerical stability with extreme values. Ensures that both implementations handle edge cases like very small or large values without producing NaN or Inf. @@ -406,8 +402,7 @@ def test_rms_norm_numerical_stability(default_vllm_config): @skip_unsupported def test_rms_norm_formula(default_vllm_config): - """ - Test that RMS norm follows the correct mathematical formula. + """Test that RMS norm follows the correct mathematical formula. Verifies: output = input / sqrt(mean(input^2) + eps) * weight """ @@ -440,8 +435,7 @@ def test_rms_norm_formula(default_vllm_config): @skip_if_not_cuda @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): - """ - Test RMS norm with various hidden sizes to ensure block size handling. + """Test RMS norm with various hidden sizes to ensure block size handling. The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it correctly handles hidden sizes both smaller and larger than the block size. @@ -474,8 +468,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): @skip_unsupported def test_rms_norm_determinism(default_vllm_config): - """ - Test that batch-invariant RMS norm produces deterministic results. + """Test that batch-invariant RMS norm produces deterministic results. Runs the same input through the kernel multiple times and verifies identical outputs. diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index e1adc948a853..4c544fab7b21 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -493,7 +493,6 @@ async def gen(rid: str): @pytest.mark.parametrize("expert_parallel", [False, True]) async def test_dp_pause_keep_then_resume(expert_parallel: bool): """Start generation, pause after a few tokens (keep mode), resume; verify gap.""" - pause_duration = 2.0 min_tokens_before_pause = 3 @@ -548,7 +547,8 @@ async def controller_task(): async def test_dp_pause_keep_race_staggered_engines(): """Race: send pause(keep) to engine 0, then add two requests, then pause(keep) to engine 1. Ensures no deadlock when pause - requests are staggered and requests arrive in between.""" + requests are staggered and requests arrive in between. + """ if DP_SIZE != 2: pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2") @@ -607,8 +607,7 @@ async def consume_gen(req_id: str) -> None: @pytest.mark.asyncio async def test_dp_pause_barrier_request_deadlock(): - """ - Test that start_dp_wave is ignored while paused. + """Test that start_dp_wave is ignored while paused. Sequence: 1. Pause all engines (PAUSED_ALL). @@ -719,7 +718,8 @@ async def consume_gen(req_id: str) -> None: async def test_dp_pause_wait_mode_drains_in_flight(): """mode="wait" through the DP consensus path: the pause lets the in-flight request run to completion while the consensus is pending, - resolves only once drained, and the engine generates again on resume.""" + resolves only once drained, and the engine generates again on resume. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) after.callback(engine.shutdown) @@ -750,7 +750,8 @@ async def test_dp_pause_wait_mode_drains_in_flight(): async def test_dp_pause_while_asleep(): """Pausing a sleeping DP engine: sleeping ranks skip dummy batches yet must still reach the pause consensus, and the completion barrier must - be harmless on workers whose memory is unmapped.""" + be harmless on workers whose memory is unmapped. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) after.callback(engine.shutdown) @@ -786,7 +787,8 @@ async def test_dp_pause_completion_implies_device_idle( """A resolved pause future promises an idle device. Inflate the dummy batches the pause consensus manufactures so any work the pause fails to wait on stays visible, then require a quiet worker stream the moment - pause_generation() returns.""" + pause_generation() returns. + """ with ExitStack() as after: # The probes below ship functions through collective_rpc. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") diff --git a/tests/v1/distributed/test_dbo.py b/tests/v1/distributed/test_dbo.py index a4e3d92be6c4..37e47e3c6900 100644 --- a/tests/v1/distributed/test_dbo.py +++ b/tests/v1/distributed/test_dbo.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism. +"""Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism. DBO is specifically designed for DP+EP scenarios to hide communication latency by overlapping computation of two batches. This test validates that DBO works @@ -54,9 +53,7 @@ ), ) def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): - """ - Test DBO with DP+EP using GSM8K evaluation. - """ + """Test DBO with DP+EP using GSM8K evaluation.""" lm_eval = pytest.importorskip("lm_eval") required_gpus = DP_SIZE diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index b53d0d8af09d..95a31f37092c 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -24,7 +24,8 @@ class ExternalLBServerManager: """Manages data parallel vLLM server instances for external - load balancer testing.""" + load balancer testing. + """ def __init__( self, diff --git a/tests/v1/distributed/test_hybrid_lb_dp.py b/tests/v1/distributed/test_hybrid_lb_dp.py index fcd3c69af54a..748370f37c5c 100644 --- a/tests/v1/distributed/test_hybrid_lb_dp.py +++ b/tests/v1/distributed/test_hybrid_lb_dp.py @@ -30,7 +30,8 @@ class HybridLBServerManager: """Manages hybrid data parallel vLLM server instances where each node runs a single logical API server that balances requests only to the - DP engines running on that same node.""" + DP engines running on that same node. + """ def __init__( self, diff --git a/tests/v1/distributed/test_internal_lb_dp.py b/tests/v1/distributed/test_internal_lb_dp.py index e4f7b3d5a83d..9fe8c9f0afd5 100644 --- a/tests/v1/distributed/test_internal_lb_dp.py +++ b/tests/v1/distributed/test_internal_lb_dp.py @@ -107,7 +107,8 @@ async def _run_request_bursts( class MultinodeInternalLBServerManager: """Manages multi-node data parallel vLLM server instances for internal - load balancer testing using --headless mode.""" + load balancer testing using --headless mode. + """ def __init__( self, @@ -241,7 +242,8 @@ def __exit__(self, exc_type, exc_val, exc_tb): class APIOnlyServerManager: """Manages API-only server (Node 0) and headless engines server (Node 1) - for testing separated API server and engine configuration.""" + for testing separated API server and engine configuration. + """ def __init__( self, @@ -261,7 +263,6 @@ def __init__( def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]: """Start API-only server and headless engines server.""" - # Start API-only server (Node 0) - no engines, only API server api_server_args = self.base_server_args.copy() api_server_args.extend( @@ -607,7 +608,6 @@ async def test_api_only_multinode_dp_completion( model_name: str, ) -> None: """Test API-only server with all engines on separate headless server.""" - # Test single request result = await _make_completion_request(api_only_client, model_name) assert result is not None @@ -645,7 +645,8 @@ async def test_api_only_multinode_dp_completion_streaming( model_name: str, ) -> None: """Test API-only server streaming with all engines on separate - headless server.""" + headless server. + """ prompt = "What is an LLM?" async def make_streaming_request(): diff --git a/tests/v1/distributed/test_pp_dp_v2.py b/tests/v1/distributed/test_pp_dp_v2.py index b931db151005..f2a274019b12 100644 --- a/tests/v1/distributed/test_pp_dp_v2.py +++ b/tests/v1/distributed/test_pp_dp_v2.py @@ -82,7 +82,8 @@ async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int: @pytest.mark.parametrize("async_scheduling", [True, False]) async def test_pp_dp_v2_low_concurrency(async_scheduling: bool): """A single in-flight request at a time, repeated, to exercise the - PP slot ring under empty batches between decodes.""" + PP slot ring under empty batches between decodes. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) after.callback(engine.shutdown) @@ -96,7 +97,8 @@ async def test_pp_dp_v2_low_concurrency(async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool): """64 concurrent requests, staggered, to exercise the steady-state - DP all-reduce + PP slot-ring path.""" + DP all-reduce + PP slot-ring path. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) after.callback(engine.shutdown) @@ -112,8 +114,8 @@ async def _one(i: int) -> int: @pytest.mark.asyncio async def test_pp_dp_v2_abort_mid_decode(): """Cancel half the in-flight requests mid-stream and confirm the - engine survives the abort storm.""" - + engine survives the abort storm. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) after.callback(engine.shutdown) @@ -153,8 +155,8 @@ async def _maybe_cancel(i: int): @pytest.mark.asyncio async def test_pp_dp_v2_pause_resume(): """Pause an engine with a request in flight, then resume and confirm - new requests still work.""" - + new requests still work. + """ with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) after.callback(engine.shutdown) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 1024829307c7..0610463484aa 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -48,7 +48,8 @@ def test_without_spec_decoding( monkeypatch: pytest.MonkeyPatch, ): """Test consistency of combos of async scheduling, preemption, - uni/multiproc executor, prefill chunking.""" + uni/multiproc executor, prefill chunking. + """ struct_outputs = StructuredOutputsParams(json=sample_json_schema) test_sampling_params: list[dict[str, Any]] = [ dict(), @@ -111,7 +112,6 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke preemption, executor, async scheduling, prefill chunking, spec decoding model length. """ - spec_config = { "method": "eagle3", "num_speculative_tokens": 2, @@ -171,7 +171,6 @@ def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch): - Async scheduling enabled (as in production) - Different executors and chunking settings """ - # Variant with larger speculation window ngram_gpu_config = { "method": "ngram_gpu", @@ -206,8 +205,8 @@ def run_tests( test_sampling_params: list[dict[str, Any]], ): """Test consistency of combos of async scheduling, preemption, - uni/multiproc executor with spec decoding.""" - + uni/multiproc executor with spec decoding. + """ # Flex attention supports float32. attention_config = {"backend": "FLEX_ATTENTION"} diff --git a/tests/v1/e2e/general/test_attention_backend_per_kind.py b/tests/v1/e2e/general/test_attention_backend_per_kind.py index 3281247b1c4a..42f2a72554e6 100644 --- a/tests/v1/e2e/general/test_attention_backend_per_kind.py +++ b/tests/v1/e2e/general/test_attention_backend_per_kind.py @@ -72,7 +72,8 @@ def test_backend_per_kind_splits_groups(backend_per_kind, monkeypatch): ) def test_backend_per_kind_overrides_global_backend(monkeypatch): """A per-kind entry wins over the global ``backend`` for its kind; other - kinds fall back to the global backend.""" + kinds fall back to the global backend. + """ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") llm = LLM( model=MODEL, diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index 7e77b31fb1b7..48f60747e832 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()` +"""Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()` handling of maximum context length for decoder models. This test ensures: @@ -35,8 +34,8 @@ def test_decoder_max_context_length_validation( max_tokens: int, ) -> None: """Check vLLM decoder model input validation for edge cases where - the prompt length is (almost) equal to the max model length.""" - + the prompt length is (almost) equal to the max model length. + """ prompt_ids = [[43] * prompt_len] with vllm_runner( @@ -73,8 +72,8 @@ def test_auto_fit_max_model_len_rejects_oversized_input( """When max_model_len='auto' and KV cache memory is very limited, the engine auto-fits max_model_len to a small value. The frontend must see this reduced value and reject prompts that exceed it, - rather than accepting them and hanging.""" - + rather than accepting them and hanging. + """ # Use a small KV cache budget to force auto-fit to a small # max_model_len. Pin block_size=16 so the budget is independent # of the platform's default block size. One block for this model is diff --git a/tests/v1/e2e/general/test_correctness_sliding_window.py b/tests/v1/e2e/general/test_correctness_sliding_window.py index 5209d1894db7..ce71bf1e5def 100644 --- a/tests/v1/e2e/general/test_correctness_sliding_window.py +++ b/tests/v1/e2e/general/test_correctness_sliding_window.py @@ -35,8 +35,7 @@ class TestConfig: def test_sliding_window_retrieval( model, batch_size, seed, disable_hybrid_kv_cache_manager, vllm_runner ): - """ - The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then + """The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then asks for value of one of them (which is outside the sliding window). If we tell it upfront which we are going to be looking for, then it answers correctly (mostly). @@ -112,14 +111,14 @@ def test_hybrid_kv_cache_manager_output_equivalence(model, vllm_runner): def check_length(prompts: list[str], llm: LLM, sliding_window: int): - """ - Check if the prompt length is valid, i.e., longer than the sliding window + """Check if the prompt length is valid, i.e., longer than the sliding window size and shorter than the model's max length. Args: prompts: list of prompts llm: LLM object sliding_window: Sliding window size + """ tokenizer = llm.get_tokenizer() max_model_len = llm.llm_engine.model_config.max_model_len diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index d92aa5c77572..b03acc21afc0 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -17,9 +17,7 @@ @pytest.fixture def test_prompts(): - """ - Adapted from tests/v1/e2e/spec_decode/utils.py - """ + """Adapted from tests/v1/e2e/spec_decode/utils.py.""" prompt_types = ["repeat", "sentence"] # Setting higher num prompts increases the chance of numerics mismatch # due to matrix multiplication numerics depending on batch dimension diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py index 0100a38ba150..26a1770d39cd 100644 --- a/tests/v1/e2e/general/test_min_tokens.py +++ b/tests/v1/e2e/general/test_min_tokens.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Comprehensive end-to-end tests for `min_tokens` in the V1 engine. +"""Comprehensive end-to-end tests for `min_tokens` in the V1 engine. Addresses #21950: verify and add CI coverage. @@ -28,7 +27,7 @@ class MinTokensTestCase: - """Data class for min_tokens test scenarios""" + """Data class for min_tokens test scenarios.""" def __init__( self, @@ -161,7 +160,7 @@ def __str__(self): @pytest.fixture(scope="module") def llm_v1(): - """Create V1 LLM instance for testing""" + """Create V1 LLM instance for testing.""" llm = LLM( model=TEST_MODEL, tensor_parallel_size=1, @@ -173,7 +172,8 @@ def llm_v1(): def test_min_tokens_with_terminal_structured_output(llm_v1: LLM): """A grammar whose only legal continuation is EOS must stop under - min_tokens instead of sampling outside the grammar.""" + min_tokens instead of sampling outside the grammar. + """ params = SamplingParams( temperature=GREEDY, min_tokens=20, @@ -193,7 +193,7 @@ def test_min_tokens_with_terminal_structured_output(llm_v1: LLM): def get_token_count(output: RequestOutput) -> int: - """Extract token count from LLM output""" + """Extract token count from LLM output.""" if not output.outputs: return 0 return len(output.outputs[0].token_ids) @@ -202,7 +202,7 @@ def get_token_count(output: RequestOutput) -> int: def assert_min_tokens_satisfied( output: RequestOutput, test_case: MinTokensTestCase ) -> None: - """Assert that min_tokens requirement is satisfied""" + """Assert that min_tokens requirement is satisfied.""" token_count = get_token_count(output) stop_reason = output.outputs[0].stop_reason if output.outputs else "no output" @@ -228,8 +228,7 @@ def assert_min_tokens_satisfied( ids=lambda tc: tc.name, ) def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase): - """ - Comprehensive test for min_tokens functionality in V1 engine. + """Comprehensive test for min_tokens functionality in V1 engine. This test covers all critical scenarios for min_tokens: - Basic functionality (should work) @@ -240,6 +239,7 @@ def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase): Args: llm_v1: V1 LLM instance test_case: Test scenario parameters + """ # Known failing cases are handled via param-level xfail marks above. @@ -279,8 +279,7 @@ def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase): def test_min_tokens_basic_functionality(llm_v1: LLM): - """ - Test basic min_tokens functionality without stop conditions. + """Test basic min_tokens functionality without stop conditions. This is a baseline test that should always pass and validates that min_tokens works correctly in the simple case. @@ -302,8 +301,7 @@ def test_min_tokens_basic_functionality(llm_v1: LLM): strict=False, ) def test_min_tokens_stop_strings_bug(llm_v1: LLM): - """ - Test the specific bug where stop strings bypass min_tokens. + """Test the specific bug where stop strings bypass min_tokens. This test specifically reproduces the bug Calvin is fixing in PR #22014. It should fail until that fix is merged. @@ -353,8 +351,7 @@ def test_min_tokens_stop_strings_bug(llm_v1: LLM): strict=False, ) def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM): - """ - Guaranteed test for stop strings bypassing min_tokens bug. + """Guaranteed test for stop strings bypassing min_tokens bug. Strategy: Use very low temperature and multiple common stop strings to virtually guarantee early detection, combined with long min_tokens @@ -408,8 +405,7 @@ def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM): strict=False, ) def test_min_tokens_eos_behavior(llm_v1: LLM): - """ - Verify EOS handling with and without min_tokens. + """Verify EOS handling with and without min_tokens. - Without min_tokens: expect early EOS -> finish_reason == "stop", stop_reason is None, and generated tokens < max_tokens (25). @@ -492,8 +488,7 @@ def test_min_tokens_eos_behavior(llm_v1: LLM): def test_min_tokens_validation(): - """ - Test that SamplingParams correctly validates min_tokens parameters. + """Test that SamplingParams correctly validates min_tokens parameters. This tests the parameter validation logic in SamplingParams. """ diff --git a/tests/v1/e2e/general/test_pooling_chunked_prefill.py b/tests/v1/e2e/general/test_pooling_chunked_prefill.py index 976e4d17387e..e6b63eaf61aa 100644 --- a/tests/v1/e2e/general/test_pooling_chunked_prefill.py +++ b/tests/v1/e2e/general/test_pooling_chunked_prefill.py @@ -71,7 +71,6 @@ def retrieve_chunks(self): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available") def test_pooling_chunked_prefill(vllm_runner, monkeypatch): """Test chunked prefill for pooling models with LastPool.""" - with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") model_id = "Qwen/Qwen3-Embedding-0.6B" @@ -124,7 +123,6 @@ def test_pooling_chunked_prefill(vllm_runner, monkeypatch): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available") def test_pooling_prefix_cache(vllm_runner, monkeypatch): """Test chunked prefill for pooling models with LastPool.""" - verses = prompt.split("\n\n") with monkeypatch.context() as m: diff --git a/tests/v1/e2e/general/test_sharded_sampling.py b/tests/v1/e2e/general/test_sharded_sampling.py index a580605fb007..4db56337a82a 100644 --- a/tests/v1/e2e/general/test_sharded_sampling.py +++ b/tests/v1/e2e/general/test_sharded_sampling.py @@ -7,7 +7,8 @@ bitwise identical across the two modes; cross-state boots differ by up to ~0.25 logprob and can flip a near-tie token). A real sharding bug — rows routed to the wrong request, misaligned gathers — produces many-nats -divergence and scrambled top-k sets, far outside these bounds.""" +divergence and scrambled top-k sets, far outside these bounds. +""" import pytest @@ -106,7 +107,8 @@ def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None: @multi_gpu_test(num_gpus=2) def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch): """Generation, logprobs, and prompt logprobs match between batch-sharded - sampling and the replicated fallback, up to measured boot-state noise.""" + sampling and the replicated fallback, up to measured boot-state noise. + """ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") # Required for the collective_rpc mode check below. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py index 1954ce6a7dcf..c71e7ad0a6b5 100644 --- a/tests/v1/e2e/general/test_streaming_input.py +++ b/tests/v1/e2e/general/test_streaming_input.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end tests for the streaming input feature in AsyncLLM. +"""End-to-end tests for the streaming input feature in AsyncLLM. These tests verify that: 1. Streaming inputs work correctly with bunched inputs (queued) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py index 8714e4f56a1e..4fcaa2a59c84 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py @@ -11,7 +11,8 @@ @single_gpu_only def test_synthetic_acceptance_rate(vllm_runner): """Verify that synthetic rejection sampling produces an acceptance - length close to the requested mean acceptance length.""" + length close to the requested mean acceptance length. + """ num_spec_tokens = 3 expected_acceptance_len = 1.875 tolerance = 0.15 diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py index 97aa82fb5456..659d43daee89 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py @@ -43,8 +43,7 @@ def run_acceptance_length_eval( chat_template_kwargs: dict[str, Any], use_mrv2: bool, ): - """ - E2E acceptance-rate validation for speculative decoding. + """E2E acceptance-rate validation for speculative decoding. Drives one or more datasets (keyed in ``expected_acceptance_lengths``) through the spec decode engine and asserts the mean acceptance length diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py index cf7b44fb226a..f4cd06ca69d8 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -167,7 +167,8 @@ def test_draft_model_tensor_parallelism(vllm_runner): @multi_gpu_only(num_gpus=2) def test_draft_model_engine_args_tensor_parallelism(): """Ensure the vllm_config for the draft model is created correctly, - and independently of the target model (quantization, TP, etc.)""" + and independently of the target model (quantization, TP, etc.) + """ engine_args = EngineArgs( model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized tensor_parallel_size=2, @@ -198,7 +199,8 @@ def test_draft_model_engine_args_tensor_parallelism(): def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic - so we can test it without instantiating a full proposer.""" + so we can test it without instantiating a full proposer. + """ spec_cfg = vllm_config.speculative_config if spec_cfg.moe_backend is not None: return replace( @@ -223,7 +225,8 @@ def _platform_moe_backend(monkeypatch: pytest.MonkeyPatch) -> MoEBackend: def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): """When moe_backend is set in speculative_config, the draft VllmConfig - should use it while the target keeps its own setting.""" + should use it while the target keeps its own setting. + """ target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", @@ -248,7 +251,8 @@ def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch): """When moe_backend is not set in speculative_config, the draft should - inherit the target's moe_backend.""" + inherit the target's moe_backend. + """ target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", @@ -271,7 +275,8 @@ def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch def test_draft_model_moe_backend_default_auto(): """When neither target nor draft set moe_backend explicitly, both should - default to 'auto'.""" + default to 'auto'. + """ engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, @@ -292,8 +297,8 @@ def test_draft_model_moe_backend_default_auto(): def test_draft_model_engine_args_rejects_invalid_tp_argname(): """The user should pass "draft_tensor_parallel_size" rather than - "tensor_parallel_size". We enforce this with validation.""" - + "tensor_parallel_size". We enforce this with validation. + """ engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, @@ -310,7 +315,8 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname(): def assert_draft_model_correctness(args: ArgsTest, vllm_runner): """Compare the outputs using and not using speculative decoding. - In the greedy decoding case, the outputs must match EXACTLY.""" + In the greedy decoding case, the outputs must match EXACTLY. + """ test_prompts: list[Messages] = get_messages( dataset=args.dataset, n=args.num_prompts ) diff --git a/tests/v1/e2e/spec_decode/draft_model/test_lora.py b/tests/v1/e2e/spec_decode/draft_model/test_lora.py index 460ac4e596d9..36a03360c7bd 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_lora.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_lora.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This script contains: +"""This script contains: 1. test lora with speculative decoding for batch inference """ @@ -59,8 +58,7 @@ def test_batch_inference_correctness( model_setup: tuple[str, str, str, str, int], vllm_runner, ): - """ - Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora. + """Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora. Should be the same and no failure when doing batch inference. model_setup: (method, model_name, spec_model_name, lora_path, tp_size) """ diff --git a/tests/v1/e2e/spec_decode/eagle/utils.py b/tests/v1/e2e/spec_decode/eagle/utils.py index 29127df85f5c..918a06e32c08 100644 --- a/tests/v1/e2e/spec_decode/eagle/utils.py +++ b/tests/v1/e2e/spec_decode/eagle/utils.py @@ -28,8 +28,7 @@ def _run_eagle_correctness( attn_backend: str, vllm_runner, ): - """ - Compare the outputs of an original LLM and a speculative LLM + """Compare the outputs of an original LLM and a speculative LLM which should be the same when using eagle speculative decoding. """ method, model_name, spec_model_name, tp_size = model_setup diff --git a/tests/v1/e2e/spec_decode/mtp/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/test_mtp.py index acd3285bd5fa..92f4a2a773f6 100644 --- a/tests/v1/e2e/spec_decode/mtp/test_mtp.py +++ b/tests/v1/e2e/spec_decode/mtp/test_mtp.py @@ -157,8 +157,7 @@ def test_mtp_correctness( expected_accuracy_threshold: float, vllm_runner, ): - """ - Compare the outputs of a original LLM and a speculative LLM + """Compare the outputs of a original LLM and a speculative LLM which should be the same when using MTP speculative decoding. Due to some variance in the engine, it is possible for some outputs to differ, so we expect that at least 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py index d3627a8703b3..c204cc1c012d 100644 --- a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py +++ b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py @@ -64,8 +64,7 @@ def test_ngram_gpu_default_with_async_scheduling( async_scheduling: bool, vllm_runner, ): - """ - Test ngram_gpu speculative decoding (k=3) correctness with and without + """Test ngram_gpu speculative decoding (k=3) correctness with and without async scheduling, validated via GSM8K accuracy. Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%). """ @@ -99,8 +98,7 @@ def test_suffix_decoding_acceptance( model_name: str, vllm_runner, ): - """ - Check that suffix decoding caching takes effect and improves acceptance + """Check that suffix decoding caching takes effect and improves acceptance lengths and acceptance rates over multiple runs of the same prompts. """ test_prompts = get_test_prompts(mm_enabled=False) diff --git a/tests/v1/e2e/spec_decode/speculators/test_speculators.py b/tests/v1/e2e/spec_decode/speculators/test_speculators.py index 990f6190751a..c6acc5573137 100644 --- a/tests/v1/e2e/spec_decode/speculators/test_speculators.py +++ b/tests/v1/e2e/spec_decode/speculators/test_speculators.py @@ -32,8 +32,7 @@ def test_speculators_model_integration( expected_accuracy_threshold: float, vllm_runner, ): - """ - Test that speculators models work with the simplified integration. + """Test that speculators models work with the simplified integration. This verifies the `vllm serve ` use case where speculative config is automatically detected from the model config diff --git a/tests/v1/e2e/spec_decode/test_sharded_sampling.py b/tests/v1/e2e/spec_decode/test_sharded_sampling.py index cc9781ca9fdb..bcb39ad95acd 100644 --- a/tests/v1/e2e/spec_decode/test_sharded_sampling.py +++ b/tests/v1/e2e/spec_decode/test_sharded_sampling.py @@ -28,7 +28,8 @@ def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch): sampling under MTP spec decoding: the collectives move the same logits bytes, Gumbel keys derive from (request slot, position, seed), and slot assignment is rank-deterministic. Both runs here are spec decode with - identical math, so outputs must match exactly.""" + identical math, so outputs must match exactly. + """ tp_size = 2 _skip_if_insufficient_gpus_for_tp(tp_size) model_name = "Qwen/Qwen3.5-0.8B-Base" diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index dd24f3884ea2..177c2f4bdbfd 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -67,8 +67,8 @@ def test_mtp_speculative_mixed_batch_short_prefill( vllm_runner, model_name, enable_prefix_caching ): """Test to ensure MTP speculative decoding correctly handles - short prefill chunks that fall below the reorder_batch_threshold.""" - + short prefill chunks that fall below the reorder_batch_threshold. + """ # Set so large that both prefills will be classified as decodes in a mixed batch # note, with prefix caching we require chunk_size >= mamba_block_size chunk_size = 256 if not enable_prefix_caching else 2048 diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index d3ae0f1d8792..b8565a0acd8c 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/tests/v1/ec_connector/integration/test_epd_correctness.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -EPD Correctness Test +"""EPD Correctness Test. Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same outputs as a baseline single instance. @@ -104,6 +103,7 @@ def check_vllm_server(url: str, timeout=5, retries=10) -> bool: Returns: True if the server is ready, False otherwise + """ for attempt in range(retries): try: @@ -138,6 +138,7 @@ def run_chat_completion( Returns: Generated text content + """ with openai.OpenAI( api_key="EMPTY", base_url=base_url, timeout=120, max_retries=0 diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py index 844753290948..df584d3dcb8a 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py @@ -381,7 +381,8 @@ def test_pin_if_ready_none_when_absent(): def test_pin_if_ready_returns_unpinned_entry_when_not_ready(): """The producer needs "present but not ready" told apart from "absent": - the first is NACKed as retryable, the second as a miss.""" + the first is NACKed as retryable, the second as a miss. + """ cache = EmbeddingCache(num_blocks=4) cache.alloc("h", 1) # not-ready entry = cache.pin_if_ready("h") diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 24ea717cabb7..9b6338a77383 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -131,7 +131,8 @@ def test_has_cache_item_false_when_not_consumer(monkeypatch): def test_connector_keys_on_identifier_not_mm_hash(monkeypatch): """The connector must key the encoder cache on feature.identifier (what - has_cache_item is called with), NOT feature.mm_hash.""" + has_cache_item is called with), NOT feature.mm_hash. + """ s = _make_scheduler(monkeypatch) req = _Request([_Feature("PROC_KEY", length=1, identifier="ENC_KEY")]) s.update_state_after_alloc(req, 0) @@ -201,7 +202,8 @@ def test_loads_only_serves_hashes_touched_this_step(monkeypatch): def test_repeated_reload_same_step_loads_once(monkeypatch): """The same mm_hash requested twice in one step must appear in - meta.loads exactly once.""" + meta.loads exactly once. + """ s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=2) @@ -224,7 +226,8 @@ def test_load_not_emitted_for_uncached_entry(monkeypatch): def test_load_pin_protects_blocks_until_completion(monkeypatch): """Loaded blocks stay pinned until the worker reports the load memcpy complete. Observable: a save that needs eviction cannot reclaim a - still-pinned loaded entry, but can once the completion report arrives.""" + still-pinned loaded entry, but can once the completion report arrives. + """ # 2 blocks, both occupied by "a" (ready). s = _make_scheduler(monkeypatch, num_blocks=2) _seed_cached(s, "a", n_blocks=2) @@ -251,7 +254,8 @@ def test_load_pin_protects_blocks_until_completion(monkeypatch): def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool: """Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the - blocks currently held by other entries are reclaimable.""" + blocks currently held by other entries are reclaimable. + """ s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0) meta = s.build_connector_meta(scheduler_output=None) return mm_hash in meta.saves @@ -311,7 +315,8 @@ def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch): def test_region_full_skips_save_and_never_blocks(monkeypatch): """When the region is fully occupied by pinned entries, new saves are - silently skipped and ensure_cache_available never blocks.""" + silently skipped and ensure_cache_available never blocks. + """ s = _make_scheduler(monkeypatch, num_blocks=1) _seed_cached(s, "pinned", n_blocks=1) s._cache.pin("pinned") @@ -330,7 +335,8 @@ def test_region_full_skips_save_and_never_blocks(monkeypatch): def test_producer_only_never_emits_loads(monkeypatch): """A producer-only scheduler must never populate meta.loads, even when - entries are ready.""" + entries are ready. + """ s = _make_scheduler(monkeypatch, ec_role="ec_producer") req = _Request([_Feature("h1", length=1)]) @@ -372,7 +378,8 @@ def test_consumer_only_has_cache_item(monkeypatch): def test_save_not_emitted_for_already_cached_entry(monkeypatch): """An entry already in the cache (from a prior save) must not trigger - a second allocation or appear in meta.saves again.""" + a second allocation or appear in meta.saves again. + """ s = _make_scheduler(monkeypatch) req = _Request([_Feature("h1", length=1)]) @@ -393,7 +400,8 @@ def test_save_not_emitted_for_already_cached_entry(monkeypatch): def test_update_connector_output_ignores_foreign_meta(monkeypatch): """A worker output whose payload is not an ECCPUWorkerMetadata is a - no-op — no cache mutation, no crash.""" + no-op — no cache mutation, no crash. + """ s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=1) @@ -407,7 +415,8 @@ class _Foreign: def test_update_connector_output_ignores_unknown_report(monkeypatch): """Reports for a never-seen save hash or a transfer id this scheduler never - dispatched are dropped rather than mutating unrelated state.""" + dispatched are dropped rather than mutating unrelated state. + """ s = _make_scheduler(monkeypatch) s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242])) assert s.has_cache_item("gone") is False @@ -419,7 +428,8 @@ def test_update_connector_output_ignores_unknown_report(monkeypatch): def test_has_pending_push_work_tracks_inflight_save(monkeypatch): """A dispatched-but-unconfirmed save keeps push work pending until the - completion report marks it ready.""" + completion report marks it ready. + """ s = _make_scheduler(monkeypatch) assert s.has_pending_push_work() is False @@ -434,7 +444,8 @@ def test_has_pending_push_work_tracks_inflight_save(monkeypatch): def test_has_pending_push_work_tracks_inflight_load(monkeypatch): """A dispatched-but-unconfirmed load keeps push work pending until the - unpin completion report arrives.""" + unpin completion report arrives. + """ s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=1) assert s.has_pending_push_work() is False # ready + unpinned @@ -453,7 +464,8 @@ def test_has_pending_push_work_tracks_inflight_load(monkeypatch): def test_shutdown_disables_roles_and_cleans_region(monkeypatch): """After shutdown the scheduler serves no items and the region is - cleaned up.""" + cleaned up. + """ s = _make_scheduler(monkeypatch, num_blocks=2) _seed_cached(s, "a", n_blocks=2) assert s.has_cache_item("a") is True diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index 722c0d89fa21..fc5a8fc8368a 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -98,7 +98,8 @@ def test_second_instance_opens_existing_file_and_shares_memory(): def test_only_creator_unlinks_file_on_cleanup(): """Critical contract: if the non-creator unlinks, the creator's mmap path becomes a dangling backing file and a third opener would create a new one - out from under the creator.""" + out from under the creator. + """ instance_id = str(uuid.uuid4()) r1 = ECSharedRegion( engine_id=instance_id, @@ -215,7 +216,8 @@ def _spy_fallback(mm, off, ln): def test_non_creator_does_not_probe_or_populate(monkeypatch): """Only the creator pre-faults. A second opener must not probe the advice - nor re-touch pages the creator already populated.""" + nor re-touch pages the creator already populated. + """ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr instance_id = str(uuid.uuid4()) @@ -312,7 +314,8 @@ def test_madvise_unexpected_oserror_propagates(monkeypatch): def test_pin_memory_success_sets_flag(region): """When cudaHostRegister returns 0, _is_pinned flips to True - and cleanup will correspondingly call cudaHostUnregister.""" + and cleanup will correspondingly call cudaHostUnregister. + """ fake_cudart = MagicMock() success = MagicMock() success.value = 0 @@ -332,7 +335,8 @@ def test_pin_memory_success_sets_flag(region): def test_pin_memory_failure_leaves_flag_false(): """If cudaHostRegister fails (non-zero), don't pretend it succeeded — - cleanup must NOT call cudaHostUnregister on memory we never registered.""" + cleanup must NOT call cudaHostUnregister on memory we never registered. + """ r = _make_region() try: fake_cudart = MagicMock() diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index f4f5f03f6946..317c714ccae0 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -363,7 +363,8 @@ def test_save_caches_coalesces_only_adjacent_blocks( def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker): """When the scheduler hasn't pre-allocated blocks for ``mm_hash``, - ``save_caches`` + ``flush_saves`` is a pure no-op.""" + ``save_caches`` + ``flush_saves`` is a pure no-op. + """ worker = make_worker() sentinel = 0x42 worker._region.blocks.fill_(sentinel) @@ -381,7 +382,8 @@ def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker): ) def test_save_caches_noop_for_non_save_rank(make_worker, tp_rank, pcp_rank): """Only TP rank 0 + PCP rank 0 writes to mmap. All other combinations - must skip the save entirely.""" + must skip the save entirely. + """ worker = make_worker(tp_rank=tp_rank, pcp_rank=pcp_rank) sentinel = 0x42 worker._region.blocks.fill_(sentinel) @@ -396,7 +398,8 @@ def test_save_caches_noop_for_non_save_rank(make_worker, tp_rank, pcp_rank): def test_save_caches_raises_when_allocated_blocks_too_small(make_worker): """``save_caches`` must raise ``AssertionError`` when the encoder output - is larger than the allocated block space.""" + is larger than the allocated block space. + """ worker = make_worker() src = torch.zeros(3 * _HIDDEN_DIM, dtype=_DTYPE) with pytest.raises(AssertionError, match="exceeds allocated blocks"): @@ -549,7 +552,8 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker): def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): """Every hash in ``meta.loads`` is copied from mmap, including one whose key is already resident in ``encoder_cache`` — there is no worker-side - skip.""" + skip. + """ worker = make_worker() n_blocks = 3 src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape( @@ -778,7 +782,8 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker): @_requires_accelerator def test_buffer_pool_is_reused_across_load_steps(make_worker): """Once a load copy completes its descriptor buffers return to the pool and - are reused by the next call — no reallocation.""" + are reused by the next call — no reallocation. + """ worker = make_worker() worker._region.blocks[0].fill_(0x01) worker._region.blocks[1].fill_(0x02) @@ -849,7 +854,8 @@ def test_init_pins_memory_when_available(pin_available, expected_pinned): def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm): """``shutdown`` must always call ``region.cleanup`` — and must never - raise.""" + raise. + """ worker = object.__new__(ECCPUWorker) mock_region = Mock(spec=ECSharedRegion) worker._region = mock_region diff --git a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py index e2cb8b120af5..bbbb615d9de4 100644 --- a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E test for the ECCPUConnector. +"""E2E test for the ECCPUConnector. Verifies: - Accuracy: outputs from EC CPU cache match fresh encoder computation. diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index 8edb54d2a69c..dac809d1a42c 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for ECExampleConnector. -""" +"""Unit tests for ECExampleConnector.""" import os from unittest.mock import Mock, patch diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py index 558a80f9cf82..1cceddb79c1a 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -43,7 +43,8 @@ def assign(proxy): def _drive(assign, e_urls, counts): """Feed a sequence of request item-counts through the cursor, as - fanout_encoder_primer does one request at a time.""" + fanout_encoder_primer does one request at a time. + """ cursor = 0 all_urls = [] for count in counts: diff --git a/tests/v1/ec_connector/unit/test_metadata.py b/tests/v1/ec_connector/unit/test_metadata.py index cca1df017193..e9b78e09f287 100644 --- a/tests/v1/ec_connector/unit/test_metadata.py +++ b/tests/v1/ec_connector/unit/test_metadata.py @@ -40,7 +40,8 @@ def test_metadata_saves_dict_operations(): def test_metadata_loads_carry_a_transfer_id_and_blocks(): """Each load entry is a (transfer_id, block_ids) pair, so the worker can - report the transfer rather than the mm_hash.""" + report the transfer rather than the mm_hash. + """ meta = ECCPUConnectorMetadata() meta.loads["mm_hash_1"] = (0, [5, 6, 7]) @@ -54,7 +55,7 @@ def test_metadata_loads_carry_a_transfer_id_and_blocks(): def test_metadata_saves_and_loads_are_independent(): - """saves and loads dicts are independent.""" + """Saves and loads dicts are independent.""" meta = ECCPUConnectorMetadata() meta.saves["key"] = [1, 2] diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index 156d98ae794c..462dcc80a9ce 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -79,7 +79,8 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): """A save whose GPU->mmap copy hasn't been confirmed complete yet is still announced: the entry can't be evicted before it's ready, and a read arriving too early is NACKed NACK_NOT_READY for the consumer to retry — - protecting it is not this method's job.""" + protecting it is not this method's job. + """ s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index 4d78b0426afc..fff7ba2ba746 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -129,7 +129,8 @@ def test_producer_session_grant_nacks_missing_entry(): def test_producer_session_grant_nacks_not_ready_entry_as_retryable(): """An announced entry whose save is still in flight is distinguished from a - miss, so the consumer retries instead of falling back to local compute.""" + miss, so the consumer retries instead of falling back to local compute. + """ cache = EmbeddingCache(num_blocks=8) cache.alloc("h1", 2) # not ready s = _make_producer_session(cache) @@ -329,7 +330,6 @@ def test_consumer_xfer_quarantined_exception_settles(): def test_consumer_xfer_quarantined_timeout_evicts_blocks(): """Quarantine timeout should release handle and return READ_FAILED.""" - data = _make_data(xfer_state="PROC") x = _started(data) # First poll transitions to QUARANTINED (deadline extended by 60s) @@ -438,7 +438,8 @@ def test_consumer_session_poll_nack_goes_to_tombstoned(): def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): """A not-ready NACK must not tombstone: tombstoning admits the request with - no embedding, which is fatal once the media has been rewritten away.""" + no embedding, which is fatal once the media has been rewritten away. + """ import msgspec s = _make_consumer_session() diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py index 93b97cb5ef53..a718bc63b635 100644 --- a/tests/v1/ec_connector/unit/utils.py +++ b/tests/v1/ec_connector/unit/utils.py @@ -45,6 +45,7 @@ def create_ec_vllm_config( Returns: A `VllmConfig`-specced mock carrying the configs above. + """ parallel_config = ParallelConfig( tensor_parallel_size=tensor_parallel_size, diff --git a/tests/v1/engine/conftest.py b/tests/v1/engine/conftest.py index 283a76dab672..1cfe671a62fd 100644 --- a/tests/v1/engine/conftest.py +++ b/tests/v1/engine/conftest.py @@ -24,12 +24,12 @@ def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors: - """Generate output processor dummy test vectors, without logprobs + """Generate output processor dummy test vectors, without logprobs. Returns: DummyOutputProcessorTestVectors instance with no logprobs - """ + """ tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) vllm_config = EngineArgs(model=TOKENIZER_NAME).create_engine_config() # Tokenize prompts under test & create dummy generated tokens @@ -62,10 +62,11 @@ def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors: @pytest.fixture def dummy_test_vectors() -> DummyOutputProcessorTestVectors: - """Generate output processor dummy test vectors, with logprobs + """Generate output processor dummy test vectors, with logprobs. Returns: DummyOutputProcessorTestVectors instance with logprobs + """ # Build dummy test vectors without logprobs dtv = _build_test_vectors_no_logprobs() diff --git a/tests/v1/engine/test_abort_final_step.py b/tests/v1/engine/test_abort_final_step.py index d8d5b73d45b9..c97e1f743637 100644 --- a/tests/v1/engine/test_abort_final_step.py +++ b/tests/v1/engine/test_abort_final_step.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Test for the fix in PR #29987: Eagerly abort cancelled final-step requests. +"""Test for the fix in PR #29987: Eagerly abort cancelled final-step requests. This test verifies that when a request is aborted during its final execution step (when it would naturally complete), it is properly marked as aborted @@ -52,8 +51,7 @@ def __init__(self): class DummyKVConnector(KVConnectorBase_V1): - """ - Dummy KV connector that captures request finish statuses to a file. + """Dummy KV connector that captures request finish statuses to a file. This is used to verify the fix - without the fix, a request aborted during its final step would be captured as FINISHED_LENGTH_CAPPED instead of FINISHED_ABORTED. @@ -144,8 +142,7 @@ def wait_for_save(self): @pytest.mark.parametrize("async_scheduling", [False, True]) @pytest.mark.asyncio async def test_abort_during_final_step(async_scheduling: bool): - """ - Test that a request aborted during its final execution step is treated as + """Test that a request aborted during its final execution step is treated as aborted rather than completed. This test: @@ -165,7 +162,6 @@ async def test_abort_during_final_step(async_scheduling: bool): With the fix, _process_aborts_queue() runs before update_from_output(), so the abort takes precedence and the KV connector sees FINISHED_ABORTED. """ - # Create three temporary files: # 1. ready_file: deleted by execute_model to signal it has started # 2. block_file: execute_model waits for this to be deleted diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 5dd37a6f7ed4..53a68aa135d3 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -457,7 +457,6 @@ async def test_customize_loggers(monkeypatch): If a customized logger is provided at the init, it should be added to the default loggers. """ - with ExitStack() as after: with set_default_torch_num_threads(1): engine = AsyncLLM.from_engine_args( @@ -642,7 +641,6 @@ async def test_check_health(): @pytest.mark.asyncio async def test_abort_final_output(output_kind: RequestOutputKind): """Test that abort() returns a final output with correct information.""" - with ExitStack() as after: with set_default_torch_num_threads(1): engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 69f52e2a7f34..49858f685cd2 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -68,7 +68,8 @@ def get_addresses(self) -> tuple[list[str], list[str]]: # in ``tests/entrypoints/test_api_server_process_manager.py``. def _bind_and_report_worker(listen_address, sock, args, client_config): """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via ``actual_address_pipe``, then exit.""" + endpoints back via ``actual_address_pipe``, then exit. + """ ctx = zmq.Context() try: in_sock = make_zmq_socket( diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index 98eda4403271..5843ce5f55ae 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -193,8 +193,7 @@ def test_engine_core(): @create_new_process_for_each_test() def test_engine_core_advanced_sampling(): - """ - A basic end-to-end test to verify that the engine functions correctly + """A basic end-to-end test to verify that the engine functions correctly when additional sampling parameters, such as top_p, min_tokens, and presence_penalty, are set. """ @@ -242,9 +241,7 @@ def _check_engine_state(): @create_new_process_for_each_test() def test_engine_core_concurrent_batches(): - """ - Test that the engine can handle multiple concurrent batches. - """ + """Test that the engine can handle multiple concurrent batches.""" def make_request_with_max_tokens(req_id: str, max_tokens: int) -> EngineCoreRequest: request = make_request() @@ -265,7 +262,6 @@ def execute_model( non_block=False, ) -> Future[ModelRunnerOutput | None]: """Make execute_model non-blocking.""" - # DummyExecutor used only for testing async case. assert non_block @@ -282,7 +278,6 @@ def sample_tokens( self, grammar_output, non_block=False ) -> Future[ModelRunnerOutput]: """Make sample_tokens non-blocking.""" - # DummyExecutor used only for testing async case. assert non_block @@ -401,10 +396,7 @@ def shutdown(self): @multi_gpu_test(num_gpus=2) def test_engine_core_tp(): - """ - Test engine can initialize worker in tp properly - """ - + """Test engine can initialize worker in tp properly.""" """Setup the EngineCore.""" engine_args = EngineArgs( model=MODEL_NAME, @@ -490,7 +482,7 @@ def test_encoder_instance_zero_kv_cache( enable_prefix_caching: bool, use_kv_connector: bool, ): - """EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests + """EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests. This test verifies encoder-only instance initializes with 0 KV cache blocks. Under EPD disagg mode, Encoder instances (EC producer role) only execute @@ -615,7 +607,8 @@ def _cadenced_dp_engine_core( monkeypatch, results: list[tuple[bool, bool]], dp_sync_interval: int = 32 ): """A bare DPEngineCoreProc whose all-reduce is scripted by `results`; - returns the core and the step numbers at which the all-reduce ran.""" + returns the core and the step numbers at which the all-reduce ran. + """ core = object.__new__(DPEngineCoreProc) core.dp_group = object() core.dp_sync_interval = dp_sync_interval @@ -634,7 +627,8 @@ def scripted_sync(dp_group, has_unfinished, pending_pause): def test_dp_sync_interval_default_is_16(): """Regression pin: lowering the default narrows the mid-wave pause tail - for async RL, where the engine is rarely idle when pause lands.""" + for async RL, where the engine is rarely idle when pause lands. + """ assert ParallelConfig.dp_sync_interval == 16 @@ -652,7 +646,8 @@ def test_dp_sync_interval_normal_wave(monkeypatch, dp_sync_interval: int): def test_dp_sync_interval_idle_pause_consensus_on_first_step(monkeypatch): """A pause of an idle engine arms every rank before its kick-started first step, so the step-1 sync reaches consensus after one dummy batch - regardless of the configured cadence.""" + regardless of the configured cadence. + """ core, synced = _cadenced_dp_engine_core( monkeypatch, [(False, True)], dp_sync_interval=32 ) @@ -678,7 +673,8 @@ def _pausable_engine_core_proc() -> EngineCoreProc: @pytest.mark.parametrize("deferred", [False, True]) def test_pause_synchronizes_device_before_cache_reset(deferred: bool): """A resolved pause promises an idle device: the barrier must run before - caches are cleared and before the caller is unblocked.""" + caches are cleared and before the caller is unblocked. + """ core = _pausable_engine_core_proc() core.engines_running = deferred order: list[str] = [] diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 3f407457b22a..e6520bfadc8e 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -249,7 +249,8 @@ def test_dplb_non_late_interaction_still_uses_lb(): def test_dplb_burst_round_robins_despite_snapshot_rebinds(): """A stats snapshot rebind wipes the optimistic lb_engines increments; - the exact in-flight floor must keep a burst spreading round-robin.""" + the exact in-flight floor must keep a burst spreading round-robin. + """ client = _make_dplb_client(num_engines=4) for _ in range(4): @@ -324,7 +325,8 @@ async def wait_for_snapshot(): def test_dplb_snapshot_backpressure_overrides_inflight(): """An engine reported heavily loaded by the coordinator is avoided even - when this client has routed nothing to it.""" + when this client has routed nothing to it. + """ client = _make_dplb_client(num_engines=2) client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]] @@ -337,7 +339,8 @@ def test_dplb_snapshot_backpressure_overrides_inflight(): def test_dplb_kv_pressure_amplifies_waiting_penalty(): """A waiting queue on a KV-bound engine (slow drain) is penalized, while - the same queue with low KV usage is not (e.g. transient burst).""" + the same queue with low KV usage is not (e.g. transient burst). + """ client = _make_dplb_client(num_engines=2) # Engine 0 has a smaller total but is KV-bound with a queue. client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]] @@ -1392,11 +1395,9 @@ def kill_first_child(): @create_new_process_for_each_test() def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPatch): - """ - Test that EngineCoreProc can be instantiated when CUDA_VISIBLE_DEVICES + """Test that EngineCoreProc can be instantiated when CUDA_VISIBLE_DEVICES is empty. This ensures the engine frontend does not need access to GPUs. """ - from vllm.v1.engine.core import EngineCoreProc from vllm.v1.executor.abstract import Executor diff --git a/tests/v1/engine/test_fast_incdec_prefix_err.py b/tests/v1/engine/test_fast_incdec_prefix_err.py index 036a19b82579..8fa8ff789b3b 100644 --- a/tests/v1/engine/test_fast_incdec_prefix_err.py +++ b/tests/v1/engine/test_fast_incdec_prefix_err.py @@ -11,8 +11,7 @@ def test_fast_inc_detok_invalid_utf8_err_case(): - """ - Test edge case where tokenizer can produce non-monotonic, + """Test edge case where tokenizer can produce non-monotonic, invalid UTF-8 output, which breaks the internal state of tokenizers' DecodeStream. See https://github.com/vllm-project/vllm/issues/17448. diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index 65dbb0e48199..e0e6004b0e4e 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -82,7 +82,7 @@ def _get_test_sampling_params( rng = random.Random(seed) def get_mostly_n_gt1() -> int: - r"""Mostly n \in [2,20], ~1/3 n=1""" + r"""Mostly n \in [2,20], ~1/3 n=1.""" x = rng.randint(0, 28) if x < 10: return 1 @@ -125,6 +125,7 @@ def test_parallel_sampling(vllm_model, example_prompts) -> None: Args: vllm_model: VllmRunner instance under test. example_prompt: test fixture providing prompts for testing. + """ sampling_params_list, n_list = _get_test_sampling_params(example_prompts) llm: LLM = vllm_model.llm diff --git a/tests/v1/engine/test_logprobs_processor.py b/tests/v1/engine/test_logprobs_processor.py index edb8cef518ca..3a315c04d501 100644 --- a/tests/v1/engine/test_logprobs_processor.py +++ b/tests/v1/engine/test_logprobs_processor.py @@ -30,7 +30,8 @@ def _make_processor(num_logprobs: int) -> LogprobsProcessor: def test_drops_trailing_sentinel_columns(): """A request that asked for 3 custom token logprobs but ended up in a - batch padded to width 5 must not surface the trailing -inf entries.""" + batch padded to width 5 must not surface the trailing -inf entries. + """ processor = _make_processor(num_logprobs=3) sampled = 42 diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 64ffc6b26675..cdd61445bf6e 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -73,6 +73,7 @@ def _ref_convert_id_to_token( Returns: String representation of input token id + """ return tokenizer.decode([token_id]) or "" @@ -228,7 +229,8 @@ def test_request_stream_interval_raises_but_not_below_engine_default( ): """A per-request stream_interval can raise the interval above the engine default but not below it (values under the default clamp up), without - altering the generated text.""" + altering the generated text. + """ engine_stream_interval = 5 # Request 0 (below the default) clamps up to 5; request 1 raises it to 10. request_stream_intervals = [1, 10] @@ -792,6 +794,7 @@ def test_stop_token( stop_token_type: "eos_token_id" for EOS, "stop_token_ids" for stop token ignore_eos: if True, EOS stops are disabled dummy_test_vectors: dummy engine core outputs and other data structures + """ model_id = dummy_test_vectors.tokenizer.name_or_path if model_id != "meta-llama/Llama-3.2-1B": diff --git a/tests/v1/engine/test_preprocess_error_handling.py b/tests/v1/engine/test_preprocess_error_handling.py index 13649a52be9a..72c7e4d45e92 100644 --- a/tests/v1/engine/test_preprocess_error_handling.py +++ b/tests/v1/engine/test_preprocess_error_handling.py @@ -14,7 +14,6 @@ def test_preprocess_error_handling(monkeypatch: pytest.MonkeyPatch): """Test that preprocessing errors are handled gracefully.""" - if current_platform.is_rocm() or current_platform.is_xpu(): pytest.skip( "Skipped on ROCm/XPU: this test only works with 'fork', " diff --git a/tests/v1/engine/utils.py b/tests/v1/engine/utils.py index b1e8c612ae99..476359510feb 100644 --- a/tests/v1/engine/utils.py +++ b/tests/v1/engine/utils.py @@ -54,6 +54,7 @@ def _create_random_top_logprob_test_vector( Returns: 1D length-`num_logprobs` torch Tensor of float logprob values + """ return torch.rand(num_logprobs) * (upper - lower) + lower @@ -79,6 +80,7 @@ def _create_random_top_logprob_test_matrix( Returns: 2D num_tokens x num_logprobs torch Tensor of float logprob values + """ return torch.rand(*shape) * (upper - lower) + lower @@ -90,7 +92,7 @@ def _create_random_top_token_test_vector( sampled_token_id: int, adjust_num_logprobs: bool = True, ) -> tuple[torch.Tensor, int]: - """Create a random vector of top logprob token indices + """Create a random vector of top logprob token indices. Use to create fake sample logprobs for testing. The sampled token ID must always be one of the top logprobs, which this dummy test @@ -117,8 +119,8 @@ def _create_random_top_token_test_vector( sampled_token_rank: the rank of sampled_token_id in the vocab vector when sorted in descending order by logprob - """ + """ # Calculate the final number of logprobs required total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs @@ -148,7 +150,7 @@ def _create_random_top_token_test_matrix( upper: int, tokens_list: list[int], ) -> tuple[torch.Tensor, torch.Tensor]: - """Create a random matrix of top logprob token indices + """Create a random matrix of top logprob token indices. Use to create fake prompt logprobs for testing. @@ -166,6 +168,7 @@ def _create_random_top_token_test_matrix( - 2D num_tokens x num_logprobs+1 torch Tensor of token ids - 1D tensor of ranks of prompt tokens in their respective rows, or random values + """ num_elements = shape[0] * shape[1] choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower @@ -205,6 +208,7 @@ def decode_token( Returns: string representation of token + """ return tokenizer.convert_ids_to_tokens(tok_id) @@ -214,7 +218,7 @@ def generate_dummy_sample_logprobs( num_logprobs: int, tokenizer: PythonBackend, ) -> list[tuple[list[int], list[float], int]]: - """Generate dummy sample logprobs + """Generate dummy sample logprobs. Generate a test data structure which imitates the list of sample logprobs which would be assembled in the engine core during decode phase. @@ -224,13 +228,14 @@ def generate_dummy_sample_logprobs( num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token tokenizer: model tokenizer to use for detokenization - Returns + Returns: list of (top token ids vector, logprobs vector, sampled token rank) Python lists tuples; in each tuple the logprobs and top token ids vectors have the same length which is either `num_logprobs` or `num_logprobs+1`. Sampled token rank is the rank (index+1) of the sampled token within the vocab vector when sorted by logprob in descending order. + """ res = [] for sampled_token_id in sampled_tokens_list: @@ -263,7 +268,7 @@ def generate_dummy_prompt_logprobs_tensors( num_logprobs: int, tokenizer: PythonBackend, ) -> LogprobsTensors: - """Generate dummy prompt logprobs tensors + """Generate dummy prompt logprobs tensors. Generate a test data structure which imitates the torch Tensors of prompt logprobs which would be assembled in the engine core during chunked @@ -274,10 +279,11 @@ def generate_dummy_prompt_logprobs_tensors( num_logprobs: return `num_logprobs` logprobs per token tokenizer: model tokenizer to use for detokenization - Returns + Returns: Single tuple of (logprobs matrix, top token ids matrix) torch Tensor, where both matrices have dimensions num_prompt_tokens x num_logprobs + """ # For now, assume the whole prompt is processed in one chunk; thus, # the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`. @@ -308,7 +314,7 @@ def generate_dummy_prompt_logprobs_tensors( @dataclass class DummyOutputProcessorTestVectors: - """Dummy test vectors for output processor tests""" + """Dummy test vectors for output processor tests.""" tokenizer: GeneralTokenizerType vllm_config: EngineArgs diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py b/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py index f5754ecb93ad..1ccbc24e8b7f 100644 --- a/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py +++ b/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py @@ -59,6 +59,7 @@ def forward( Returns: If aux_hidden_state_layers is set: (hidden_states, aux_hidden_states) Otherwise: hidden_states + """ # Determine sequence length if inputs_embeds is not None: diff --git a/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py index caa4aab870ab..4945ab44d481 100644 --- a/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py +++ b/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py @@ -25,8 +25,7 @@ def check_vllm_server(url: str, timeout=5, retries=3) -> bool: - """ - Checks if the vLLM server is ready by sending a GET request to the + """Checks if the vLLM server is ready by sending a GET request to the /health endpoint. Args: @@ -36,6 +35,7 @@ def check_vllm_server(url: str, timeout=5, retries=3) -> bool: Returns: bool: True if the server is ready, False otherwise. + """ for attempt in range(retries): try: @@ -81,8 +81,7 @@ def run_simple_prompt( def main(): - """ - This script demonstrates how to accept two optional string arguments + """This script demonstrates how to accept two optional string arguments ("service_url" and "file_name") from the command line, each with a default value of an empty string, using the argparse module. """ diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py index f109190a4a0a..407ab3da29fb 100644 --- a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py +++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). +"""Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). Tests cover: - Output correctness across block-size boundaries (proxy vs direct prefill). diff --git a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py index 1a910c8d24dc..e207181958d2 100644 --- a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py +++ b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py @@ -18,9 +18,7 @@ @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 = [] @@ -132,8 +130,7 @@ def parse_args(): 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 @@ -141,6 +138,7 @@ def get_next_client(app, service_type: str): Returns: The next client to use + """ if service_type == "prefill": client_idx = next(app.state.prefill_iterator) @@ -155,9 +153,7 @@ def get_next_client(app, service_type: str): async def send_request_to_service( client_info: dict, 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, @@ -200,9 +196,7 @@ async def send_request_to_service( async def stream_service_response( client_info: dict, endpoint: str, 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, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py index 9c3ba6d555dd..3c2f63c66ab6 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py @@ -80,7 +80,8 @@ def _split_hnd_cache(spec) -> torch.Tensor: def _packed_nhd_cache(spec) -> torch.Tensor: """Logical (num_blocks, heads, block_size, 2 * head_size) over an NHD - physical layout — the FlashAttention/FlashInfer/Triton/Flex form.""" + physical layout — the FlashAttention/FlashInfer/Triton/Flex form. + """ return torch.zeros( NUM_BLOCKS, spec.block_size, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index f597c8fa27c8..2d87431e319d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -863,7 +863,8 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): ) def test_canonical_layout_gate(kv_cache_groups, certified): """The canonical layout certifies portability group by group; none of - these shapes are portable in the direct layout.""" + these shapes are portable in the direct layout. + """ assert not _parallelism_agnostic(kv_cache_groups) assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified @@ -871,7 +872,8 @@ def test_canonical_layout_gate(kv_cache_groups, certified): def test_canonical_layout_certifies_v2_model_runner(): """Canonical bytes are certified per layer against live tensor strides at registration, so the static gate must not depend on the model-runner - version — the v2 runner is the case the canonical layout exists for.""" + version — the v2 runner is the case the canonical layout exists for. + """ groups = _groups(_full_attention_spec()) assert _parallelism_agnostic(groups, canonical=True, v2=True) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index b7d9d6c8efac..961ce1af12b7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -117,7 +117,8 @@ def _spec_cls_with_metric_definitions( metric_definitions: dict[str, Any], ) -> type: """Build a fake offloading spec class reporting the given metric - definitions, so tests don't need to patch the real CPU spec.""" + definitions, so tests don't need to patch the real CPU spec. + """ class _FakeOffloadingSpec: @staticmethod @@ -187,7 +188,8 @@ def test_build_kv_connector_stats_with_empty_dict(): def test_build_kv_connector_stats_reconstructs_offload_stats(): """Test that OffloadingConnector stats are properly reconstructed with - correct data.""" + correct data. + """ serialized_data = { _StatsKey.TYPES: { LOAD_BYTES: _MetricType.COUNTER, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 035fb6b0054f..df69b9bb66f7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1583,7 +1583,8 @@ def _maximal_lookup(sched, keys, start_chunk_idx: int = 0): @pytest.mark.parametrize("result", list(LookupResult)) def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult): """Whether a result keeps a sliding-window streak alive decides how wide - the demanded chunk set gets, so every member needs deliberate behavior.""" + the demanded chunk set gets, so every member needs deliberate behavior. + """ assert result in _SCAN_BEHAVIOR, f"{result} has no declared scan behavior" expected_lookups, expected_end = _SCAN_BEHAVIOR[result] @@ -2034,7 +2035,8 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling def test_loads_do_not_populate_fence_index(request_runner): """Loads don't populate _block_id_to_pending_jobs (protected by - delay_free_blocks while in flight).""" + delay_free_blocks while in flight). + """ runner = request_runner( blocks_per_chunk=3, block_size=4, @@ -2104,7 +2106,8 @@ def capture_fence(): def test_fence_at_build_store_jobs(request_runner): """A new prefill (no load -> update_state_after_alloc returns early) reusing a finished request's pending-store block is flushed by - _build_store_jobs's fence.""" + _build_store_jobs's fence. + """ runner = request_runner( blocks_per_chunk=1, block_size=4, @@ -2154,7 +2157,8 @@ def capture_fence(): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_complete_store_called_per_job(request_runner, async_scheduling: bool): """complete_store fires per-job, not deferred to request finish. - Each call carries only that store's keys.""" + Each call carries only that store's keys. + """ tokens_per_block = 4 blocks_per_chunk = 3 tokens_per_chunk = tokens_per_block * blocks_per_chunk @@ -2383,7 +2387,8 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_reset_cache(request_runner, async_scheduling: bool): """reset_cache flushes in-flight loads, calls manager.reset_cache(), resets - next_stored_chunk_idx for active requests and clears job tracking.""" + next_stored_chunk_idx for active requests and clears job tracking. + """ block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -2929,7 +2934,8 @@ def test_stale_sliding_window_block_after_prepare_store_failure( @pytest.mark.parametrize("async_scheduling", [True, False]) def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): """When skip_reading_prefix_cache=True, the offloading connector must not - load any blocks from CPU even if a matching prefix is cached there.""" + load any blocks from CPU even if a matching prefix is cached there. + """ block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -3023,7 +3029,8 @@ def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: class TestEagle: """Tests for Eagle/MTP speculative decoding support in the offloading - connector scheduler — both _lookup() unit tests and integration tests.""" + connector scheduler — both _lookup() unit tests and integration tests. + """ # ------------------------------------------------------------------- # Helpers @@ -3565,7 +3572,8 @@ def test_sw_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): """Eagle sliding-window group stores all prompt blocks but excludes - the trailing chunk while decoding.""" + the trailing chunk while decoding. + """ block_size = 4 sliding_window = 8 num_gpu_blocks = 100 @@ -3617,7 +3625,8 @@ def test_single_block_stored_at_end_of_prefill( self, request_runner, async_scheduling: bool ): """An eagle group with a single-block prompt stores it at the end of - prefill: prompt blocks are stable, so no tail is held back.""" + prefill: prompt blocks are stable, so no tail is held back. + """ block_size = 4 blocks_per_chunk = 1 tokens_per_chunk = block_size * blocks_per_chunk @@ -4254,7 +4263,8 @@ def capture_fence(): def test_chunked_local_attention_reports_its_chunk_window(): """Llama 4 uses chunked local attention, which used to trip the - FullAttentionSpec assert and kill the engine at startup.""" + FullAttentionSpec assert and kill the engine at startup. + """ spec = ChunkedLocalAttentionSpec( block_size=16, num_kv_heads=2, @@ -4493,7 +4503,8 @@ def test_retention_interval_zero_stores_only_replay_boundary( def _shared_kv_mtp_config(): """Speculative config for a shared-group MTP model: eagle-family method whose drafter layer merges into a target KV-cache group, so no group - self-identifies as a drafter group (issue #52735).""" + self-identifies as a drafter group (issue #52735). + """ spec = MagicMock(name="shared_kv_mtp_spec") spec.use_eagle.return_value = True spec.use_eagle_block_drop.return_value = True @@ -4507,12 +4518,14 @@ def _shared_kv_mtp_config(): class TestSharedGroupMTPOffload: """Regression tests for issue #52735: OffloadingConnector must keep serving when speculative decoding is enabled but no KV-cache group is - annotated as a drafter group (shared-group MTP models).""" + annotated as a drafter group (shared-group MTP models). + """ def test_no_annotation_marks_no_groups(self, request_runner): """Spec decode on + zero annotated groups must NOT mark every group as a drafter group; the full store->load roundtrip must match the - non-speculative behavior of test_two_groups_full_and_sliding_window.""" + non-speculative behavior of test_two_groups_full_and_sliding_window. + """ block_size = 4 kv_cache_groups = [ KVCacheGroupSpec( @@ -4569,7 +4582,8 @@ def test_no_annotation_marks_no_groups(self, request_runner): class TestMambaHybridOffloadServing: """Store->finish->lookup flows on a full-attention + mamba-align hybrid, with a manager that only HITs keys that were actually stored. Guards the - two collapse routes of issue #52735.""" + two collapse routes of issue #52735. + """ BLOCK = 4 MAMBA_BLOCK = 16 @@ -4698,7 +4712,8 @@ def test_shared_group_mtp_serves_like_non_speculative(self): """Route 1 of #52735: with spec decode on and no drafter annotation, serving must equal the non-speculative baseline (was 0 before the fix: the all-groups fallback marked the mamba group as a drafter and - the volatile-tail pop consumed its only servable chunk).""" + the volatile-tail pop consumed its only servable chunk). + """ baseline = self._roundtrip_served_tokens(self._make_scheduler(None)) assert baseline == 16 served = self._roundtrip_served_tokens( @@ -4711,7 +4726,8 @@ def test_annotated_eagle_group_does_not_starve_coarse_sibling(self): group must not drag the confirmed boundary below a coarser sibling's chunk granularity. The widened query makes the volatile-tail pop land on an extra queried chunk, holding the boundary at 16 tokens - (was 0 before the fix: 4 hits -> pop -> 12 tokens < mamba chunk).""" + (was 0 before the fix: 4 hits -> pop -> 12 tokens < mamba chunk). + """ scheduler = self._make_scheduler(None, mamba_eagle=True) assert [c.is_eagle_group for c in scheduler.config.kv_group_configs] == [ True, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 366624e1d3b1..9fdac585276d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -80,7 +80,8 @@ def _allocate_kv_caches( def _single_rank_vllm_config(total_kv_heads: int): """A one-rank (TP=1) parallel config, as canonical mappings are derived - from it.""" + from it. + """ vllm_config = MagicMock() parallel_config = vllm_config.parallel_config parallel_config.tensor_parallel_size = 1 @@ -98,9 +99,7 @@ def _make_worker( replicated_layout: bool = False, rank: int = 0, ): - """ - Create an OffloadingConnectorWorker with mocked dependencies. - """ + """Create an OffloadingConnectorWorker with mocked dependencies.""" from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( OffloadingConnectorWorker, ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 5707003ce935..ad8d4fc362b9 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -475,8 +475,7 @@ def _run( complete_transfers: bool, post_step_fn: Callable[[], None] | None = None, ): - """ - Runs multiple engine (scheduler + worker) steps. + """Runs multiple engine (scheduler + worker) steps. Assumes a single request is running. Args: @@ -484,8 +483,8 @@ def _run( complete_transfers: complete transfers immediately post_step_fn: optional callback invoked after each step's update_from_output(), before the next schedule(). - """ + """ tokens_iter = iter(decoded_tokens) token_id = next(tokens_iter, None) prev_scheduler_output = None @@ -617,8 +616,7 @@ def run( expected_flushed: tuple[int | tuple[int, int], ...] = (), post_step_fn: Callable[[], None] | None = None, ): - """ - Runs multiple engine (scheduler + worker) steps. + """Runs multiple engine (scheduler + worker) steps. Assumes a single request is running. Args: @@ -634,8 +632,8 @@ def run( A GPU block is either a (group_idx: int, request_block_offset: int) or just request_block_offset: int. The latter case is a convenience for representing all groups. - """ + """ expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored) expected_loaded_gpu_blocks = self._to_gpu_blocks(expected_loaded) expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed) diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 78a6d930f97a..006229b8699f 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -148,7 +148,8 @@ def test_multiturn_lifecycle(): (do_remote_decode=True). Finishes LENGTH_CAPPED with remote_block_ids. Turn 2: P receives remote_block_ids from D. P pulls KV from D because remote_block_ids is not None and external tokens > 0. Computes only - new tokens, finishes LENGTH_CAPPED.""" + new tokens, finishes LENGTH_CAPPED. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -194,7 +195,8 @@ def test_multiturn_lifecycle(): def test_first_turn_no_remote_blocks(): """First turn: P has no remote_block_ids from D yet. - Standard local prefill, returns kv_transfer_params for future turns.""" + Standard local prefill, returns kv_transfer_params for future turns. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -274,7 +276,8 @@ def test_abort_p_side_non_length_capped(): def test_remote_blocks_exceed_prompt_tokens(): """D provides more remote tokens than P's prompt needs. - P caps external tokens to prompt length.""" + P caps external tokens to prompt length. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -310,7 +313,8 @@ def test_remote_blocks_exceed_prompt_tokens(): def test_p_node_pulls_partial_last_block_from_d(): """D sends remote_block_ids with partially filled last block. remote_num_tokens < len(remote_block_ids) * block_size. - P pulls only remote_num_tokens worth of external tokens.""" + P pulls only remote_num_tokens worth of external tokens. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -355,7 +359,8 @@ def test_p_node_pulls_partial_last_block_from_d(): def test_add_new_req_to_recv_populates_remote_meta(): """add_new_req_to_recv correctly populates RemoteMeta for P-node - bi-directional KV pull from D.""" + bi-directional KV pull from D. + """ meta = NixlConnectorMetadata() kv_params = { "remote_block_ids": [[0, 1, 2]], @@ -383,7 +388,8 @@ def test_add_new_req_to_recv_populates_remote_meta(): def test_build_connector_meta_recv_entries(): """P-node scheduler: do_remote_decode=True + remote_block_ids → - _reqs_need_recv populated, build_connector_meta produces reqs_to_recv.""" + _reqs_need_recv populated, build_connector_meta produces reqs_to_recv. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -457,7 +463,8 @@ def test_p_node_pull_kv_from_d(dist_init): ) def test_p_node_pull_then_send_kv(dist_init): """Full P-node bi-directional: pull KV from D → prefill → - send KV back to D via notification.""" + send KV back to D via notification. + """ connector, worker = _make_connector_with_fake_worker() meta = _make_p_node_recv_metadata("req-p2", [10, 11], [20, 21]) _do_load_kv(connector, meta) @@ -504,7 +511,8 @@ def test_p_node_deferred_pull_on_no_handshake(dist_init): def test_d_node_request_finished_returns_kv_params(): """D-node request_finished returns kv_transfer_params with do_remote_decode=True, remote_block_ids, remote_num_tokens - for P to pull. These params go directly to P node.""" + for P to pull. These params go directly to P node. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -585,7 +593,8 @@ def test_d_node_request_finished_remote_num_tokens(): def test_d_node_partial_last_block_remote_num_tokens(): """D-node: remote_num_tokens < len(remote_block_ids) * block_size - when last block is partially filled.""" + when last block is partially filled. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -619,7 +628,8 @@ def test_no_double_read_blocks_after_reschedule(): bidirectional request (once on initial schedule, once after WAITING_FOR_REMOTE_KVS → reschedule). The _remote_blocks_processed flag must prevent the request from being added to _reqs_need_recv - twice, which would cause P to read D's blocks twice.""" + twice, which would cause P to read D's blocks twice. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -672,7 +682,8 @@ def test_remote_num_tokens_bounded_by_blocks(): """Edge case 2: D-node request_finished must return remote_num_tokens <= len(remote_block_ids) * block_size. request.num_tokens includes the last sampled token which has no KV - in the cache, so remote_num_tokens must use num_computed_tokens.""" + in the cache, so remote_num_tokens must use num_computed_tokens. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -708,7 +719,8 @@ def test_remote_num_tokens_bounded_by_blocks(): def test_kv_recompute_threshold_skips_small_transfer(): """Edge case 3: When remote tokens are below kv_recompute_threshold, P should skip the remote pull and compute locally instead of - entering WAITING_FOR_REMOTE_KVS.""" + entering WAITING_FOR_REMOTE_KVS. + """ threshold = 256 vllm_config = create_vllm_config( kv_connector_extra_config={ @@ -750,7 +762,8 @@ def test_p_node_finished_holds_blocks_for_d(): """Edge case 4: P-node finishes with FINISHED_LENGTH_CAPPED and do_remote_decode=True. P must hold blocks (delay_free=True) and return kv_transfer_params with do_remote_prefill=True so D can - read P's blocks.""" + read P's blocks. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -791,7 +804,8 @@ def test_p_node_finished_holds_blocks_for_d(): def test_cache_miss_first_turn_no_remote_pull(): """Edge case 5: First turn with do_remote_decode=True but no remote_block_ids (cache MISS). P should prefill locally with - num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS.""" + num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -826,7 +840,8 @@ def test_cache_miss_first_turn_no_remote_pull(): def test_partial_remote_tokens_less_than_prompt(): """Edge case 6: D's remote_num_tokens covers only part of P's prompt. P should pull remote_num_tokens worth of external tokens - and compute the rest locally.""" + and compute the rest locally. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -871,7 +886,8 @@ def test_partial_remote_tokens_less_than_prompt(): def test_remote_blocks_processed_flag_persists(): """Edge case 7: After recv completes and request is rescheduled, the _remote_blocks_processed flag in kv_transfer_params prevents - the bidirectional path from re-entering _reqs_need_recv.""" + the bidirectional path from re-entering _reqs_need_recv. + """ vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -943,7 +959,8 @@ def test_remote_blocks_processed_flag_persists(): ) def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): """P declines on the turn-2 readback if the deadline is known, - near-expiry and a handshake clock offset is known""" + near-expiry and a handshake clock offset is known + """ connector, worker = _make_connector_with_fake_worker() worker._engine_clock_offset[_REMOTE] = offset expiry_time = None if expiry_delta is None else time.perf_counter() + expiry_delta @@ -979,7 +996,8 @@ def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): ) def test_turn2_full_prefix_hit_with_expired_deadline_skips_gate(dist_init): """A full local prefix hit issues no READ. The expiry gate must be skipped - even when D's deadline is expired.""" + even when D's deadline is expired. + """ connector, worker = _make_connector_with_fake_worker() worker._engine_clock_offset[_REMOTE] = 0.0 @@ -1032,7 +1050,8 @@ def test_d_node_request_finished_exports_blocks_expiry_time(): def test_handshake_listener_appends_perf_counter_frame(): """The handshake reply carries a 2nd frame with the listener's live - perf_counter.""" + perf_counter. + """ import threading import msgspec diff --git a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py index ec3fb8231e19..21b50a13794b 100644 --- a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py +++ b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -test that invalid blocks are evicted from prefix cache to prevent pollution. +"""test that invalid blocks are evicted from prefix cache to prevent pollution. verifies that when sync-loading fails, invalid blocks are removed from the prefix cache hash table so future requests cannot match and reuse corrupted data. @@ -39,7 +38,7 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]: @pytest.fixture def fail_scheduler(): - """scheduler with kv_load_failure_policy='fail'""" + """Scheduler with kv_load_failure_policy='fail'.""" vllm_config = create_vllm_config() vllm_config.kv_transfer_config.kv_load_failure_policy = "fail" return create_scheduler(vllm_config) @@ -48,8 +47,7 @@ def fail_scheduler(): def test_invalid_blocks_evicted_prevents_cache_pollution( fail_scheduler: Scheduler, ): - """ - verify invalid blocks are evicted to prevent future cache hits. + """Verify invalid blocks are evicted to prevent future cache hits. scenario: 1. request 1 loads externally-computed blocks (sync mode) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 019b8d1504ac..d3dabe4fb468 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -19,13 +19,15 @@ class _StubLMCacheMPConnector: test only asserts on the connector *name* and the ``extra_config`` dict produced by ``VllmConfig``, never instantiates the connector, so a bare placeholder class is sufficient. Not subclassing ``SupportsHMA`` mirrors - the real connector's HMA support (it does not support HMA either).""" + the real connector's HMA support (it does not support HMA either). + """ @pytest.fixture def stub_lmcache_mp_connector(monkeypatch): """Replace the lazy loader so VllmConfig.__post_init__ does not import - ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests.""" + ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests. + """ monkeypatch.setitem( KVConnectorFactory._registry, "LMCacheMPConnector", @@ -105,7 +107,8 @@ def _build_config( enable_cumem_allocator: bool = False, ) -> VllmConfig: """Build a VllmConfig that exercises _verify_kv_transfer_compat without - requiring a real model (avoids HF downloads in CI).""" + requiring a real model (avoids HF downloads in CI). + """ from types import SimpleNamespace kv_transfer_config = ( @@ -131,7 +134,8 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): invalidated when expandable_segments lets the CUDA VMM allocator remap the underlying physical pages. We can't enumerate every connector that does this (especially OOT ones), so reject the combination whenever any - connector is configured.""" + connector is configured. + """ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") with pytest.raises(ValueError, match="expandable_segments"): _build_config(kv_connector=kv_connector) @@ -139,7 +143,8 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch): """Sleep mode routes KV allocations through CuMemAllocator's pool, which - auto-disables expandable_segments (see #40812).""" + auto-disables expandable_segments (see #40812). + """ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") _build_config(kv_connector="NixlConnector", enable_sleep_mode=True) @@ -162,7 +167,8 @@ def test_kv_connector_allows_other_alloc_conf(monkeypatch): def test_no_kv_connector_ignores_expandable_segments(monkeypatch): """The expandable_segments check only applies when a KV connector is - configured.""" + configured. + """ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") _build_config(kv_connector=None) diff --git a/tests/v1/kv_connector/unit/test_decode_bench_connector.py b/tests/v1/kv_connector/unit/test_decode_bench_connector.py index 76bb4648a933..e1054d9ef86f 100644 --- a/tests/v1/kv_connector/unit/test_decode_bench_connector.py +++ b/tests/v1/kv_connector/unit/test_decode_bench_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for DecodeBenchConnector. +"""Unit tests for DecodeBenchConnector. Tests the functionality of the DecodeBenchConnector which fills KV cache with dummy values for decode performance benchmarking. diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py index a07364cd3ea1..2f33ab18b917 100644 --- a/tests/v1/kv_connector/unit/test_error_propagation.py +++ b/tests/v1/kv_connector/unit/test_error_propagation.py @@ -32,14 +32,14 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]: @pytest.fixture def fail_scheduler(): - """scheduler with kv_load_failure_policy='fail'""" + """Scheduler with kv_load_failure_policy='fail'.""" vllm_config = create_vllm_config() vllm_config.kv_transfer_config.kv_load_failure_policy = "fail" return create_scheduler(vllm_config) def test_error_propagation_sync_load(fail_scheduler: Scheduler): - """test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load)""" + """Test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load).""" num_prompt_blocks = 100 num_external_computed_blocks = 99 invalid_block_idx = 50 @@ -93,7 +93,7 @@ def test_error_propagation_sync_load(fail_scheduler: Scheduler): def test_error_propagation_async_load(fail_scheduler: Scheduler): - """test invalid_block_ids with fail policy -> FINISHED_ERROR (async load)""" + """Test invalid_block_ids with fail policy -> FINISHED_ERROR (async load).""" num_prompt_blocks = 100 num_external_computed_blocks = 99 invalid_block_idx = 50 diff --git a/tests/v1/kv_connector/unit/test_example_connector.py b/tests/v1/kv_connector/unit/test_example_connector.py index 923eb22b31f5..d5b1856cc8b9 100644 --- a/tests/v1/kv_connector/unit/test_example_connector.py +++ b/tests/v1/kv_connector/unit/test_example_connector.py @@ -29,12 +29,12 @@ class InputCase(NamedTuple): def _check_path_len(path): - """Return the latest length in path""" + """Return the latest length in path.""" return len(list(path.iterdir())) def _list_path(path): - """Return the list of foldername (hashes generated) under the path""" + """Return the list of foldername (hashes generated) under the path.""" return list(path.iterdir()) @@ -47,8 +47,7 @@ def run_test( expected_len: int, info: str, ): - """ - One individual test to process the prompt and output base on 1 set of input + """One individual test to process the prompt and output base on 1 set of input Then check if the length in the storage path matches the expected length `info` introduces details or purpose of the individual test """ @@ -67,9 +66,7 @@ def run_test( def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]): - """ - Form the prompt based on the text and image input, then llm generate output - """ + """Form the prompt based on the text and image input, then llm generate output.""" placeholders = [ { "type": "image_url", @@ -120,8 +117,7 @@ def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]): ), ) def test_shared_storage_connector_hashes(tmp_path, attn_backend): - """ - Tests that ExampleConnector saves KV to the storage locations + """Tests that ExampleConnector saves KV to the storage locations with proper hashes; that are unique for inputs with identical text but different images (same size), or same multiple images but different orders. """ diff --git a/tests/v1/kv_connector/unit/test_flexkv_connector.py b/tests/v1/kv_connector/unit/test_flexkv_connector.py index 8cb57366345c..a40e5f828cab 100644 --- a/tests/v1/kv_connector/unit/test_flexkv_connector.py +++ b/tests/v1/kv_connector/unit/test_flexkv_connector.py @@ -50,7 +50,8 @@ def _make_flexkv_module( impl_mock: MagicMock, ) -> tuple[types.ModuleType, types.ModuleType]: """Build a fake ``flexkv`` package hierarchy that returns *impl_mock* - when ``FlexKVConnectorV1Impl`` is instantiated.""" + when ``FlexKVConnectorV1Impl`` is instantiated. + """ flexkv_mod = types.ModuleType("flexkv") integration_mod = types.ModuleType("flexkv.integration") vllm_mod = types.ModuleType("flexkv.integration.vllm") @@ -71,7 +72,8 @@ def _make_flexkv_module( def _install_flexkv_mock(impl_mock: MagicMock): """Insert fake flexkv modules into sys.modules and return a context that - cleans them up afterwards.""" + cleans them up afterwards. + """ flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock) mods = { "flexkv": flexkv_mod, @@ -104,7 +106,8 @@ def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock): class TestFlexKVConnectorImportError: """FlexKVConnectorV1 should fail with a helpful message when flexkv is - absent.""" + absent. + """ def test_import_error_message(self): from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import ( diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index e5cbd977b6d8..b2cf0f0b64dd 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -176,7 +176,8 @@ def test_engine_unwraps_handshake_metadata_for_legacy_connector( ) -> None: """Engine core always asks workers for `(pp_rank, tp_rank)`-keyed metadata, then unwraps to `{tp_rank: metadata}` for a connector that has not opted - into PP-aware handshake (single-PP producer, all `pp_rank == 0`).""" + into PP-aware handshake (single-PP producer, all `pp_rank == 0`). + """ metadata_0 = _Metadata() metadata_1 = _Metadata() connector = _LegacyConnector() @@ -199,7 +200,8 @@ def test_engine_rejects_pp_producer_for_legacy_connector( monkeypatch: pytest.MonkeyPatch, ) -> None: """A connector that has not opted into PP-aware handshake must not silently - drop metadata from `pp_rank > 0`; engine core init raises instead.""" + drop metadata from `pp_rank > 0`; engine core init raises instead. + """ connector = _LegacyConnector() with pytest.raises(ValueError, match="does not support PP-disaggregated"): @@ -214,7 +216,8 @@ def test_engine_passes_handshake_metadata_through_for_pp_aware_connector( monkeypatch: pytest.MonkeyPatch, ) -> None: """A PP-aware connector receives the full `(pp_rank, tp_rank)`-keyed dict - unchanged.""" + unchanged. + """ metadata_0 = _Metadata() metadata_1 = _Metadata() connector = _PPAwareConnector() diff --git a/tests/v1/kv_connector/unit/test_hf3fs_client.py b/tests/v1/kv_connector/unit/test_hf3fs_client.py index d9c34a8907de..d0ae6c24acd0 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_client.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_client.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for resource management in hf3fs_client.py: constructor failure cleanup +"""Tests for resource management in hf3fs_client.py: constructor failure cleanup and idempotent close(). Tests use mock to replace real I/O operations (hf3fs_fuse.io, SharedMemory, os, CUDA). Requires hf3fs_fuse.io to be installed; skipped otherwise. @@ -132,7 +131,8 @@ def _make_client(self, tmp_path): def test_close_idempotent_and_handles_cleared(self, tmp_path): """Multiple close() calls must not raise; deregister_fd called exactly - once, all handles set to None, shm.close() invoked.""" + once, all handles set to None, shm.close() invoked. + """ client, shm_r, shm_w = self._make_client(tmp_path) with ( @@ -238,7 +238,8 @@ def test_constructor_failure_after_shm_alloc_closes_shm(self, tmp_path): def test_constructor_failure_does_not_close_unallocated_shm(self, tmp_path): """Failure before SharedMemory is created must not raise AttributeError - or TypeError from cleanup.""" + or TypeError from cleanup. + """ with ( patch(f"{self._MOD}.HF3FS_AVAILABLE", True), patch(f"{self._MOD}.register_fd"), diff --git a/tests/v1/kv_connector/unit/test_hf3fs_connector.py b/tests/v1/kv_connector/unit/test_hf3fs_connector.py index 94bb94c6fbd7..8a8f046835fc 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_connector.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_connector.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for HF3FS KV Connector high-level components: - - TestHf3fsMockClient : file-backed mock client I/O correctness - - TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation +"""Tests for HF3FS KV Connector high-level components: +- TestHf3fsMockClient : file-backed mock client I/O correctness +- TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation """ import os diff --git a/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py b/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py index f922c7c85586..a366f95caa28 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for HF3FS metadata server data structures and allocation logic: - - RankFileMetadata : page allocation / release primitives - - KeyMetadata : per-key rank-page tracking and completion detection - - GlobalMetadataState : coordinated allocation with cache-hit semantics +"""Tests for HF3FS metadata server data structures and allocation logic: +- RankFileMetadata : page allocation / release primitives +- KeyMetadata : per-key rank-page tracking and completion detection +- GlobalMetadataState : coordinated allocation with cache-hit semantics """ import pytest diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 77d629729776..41fce6e9ca0f 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for correctness in invalid block handling. +"""Tests for correctness in invalid block handling. These tests verify correct behavior in three scenarios: 1. Sync recompute case: Blocks should not be freed for running requests @@ -42,7 +41,7 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]: @pytest.fixture def fail_scheduler(): - """scheduler with kv_load_failure_policy='fail'""" + """Scheduler with kv_load_failure_policy='fail'.""" vllm_config = create_vllm_config() vllm_config.kv_transfer_config.kv_load_failure_policy = "fail" return create_scheduler(vllm_config) @@ -50,7 +49,7 @@ def fail_scheduler(): @pytest.fixture def recompute_scheduler(): - """scheduler with kv_load_failure_policy='recompute'""" + """Scheduler with kv_load_failure_policy='recompute'.""" vllm_config = create_vllm_config() vllm_config.kv_transfer_config.kv_load_failure_policy = "recompute" return create_scheduler(vllm_config) @@ -59,8 +58,7 @@ def recompute_scheduler(): def test_sync_recompute_blocks_not_freed_for_running_requests( recompute_scheduler: Scheduler, ): - """ - Test sync recompute case - blocks must not be freed for running requests. + """Test sync recompute case - blocks must not be freed for running requests. When a running request has invalid blocks and retry_policy is 'recompute': 1. Request should remain in RUNNING state @@ -180,8 +178,7 @@ def test_sync_recompute_blocks_not_freed_for_running_requests( def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler): - """ - Test sync fail case - invalid blocks must be evicted from cache. + """Test sync fail case - invalid blocks must be evicted from cache. When a request fails with policy='fail' and has invalid blocks from sync loading: 1. Request should be finished with FINISHED_ERROR @@ -296,8 +293,7 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler): def test_async_recompute_blocks_not_cached_when_invalid( recompute_scheduler: Scheduler, ): - """ - Test async recompute case - invalid blocks not cached after transfer. + """Test async recompute case - invalid blocks not cached after transfer. When async KV loading has invalid blocks and retry_policy is 'recompute': 1. Blocks are allocated but not cached yet diff --git a/tests/v1/kv_connector/unit/test_lmcache_connector.py b/tests/v1/kv_connector/unit/test_lmcache_connector.py index c3df2b68b1ff..ad8ac0566e0d 100644 --- a/tests/v1/kv_connector/unit/test_lmcache_connector.py +++ b/tests/v1/kv_connector/unit/test_lmcache_connector.py @@ -212,7 +212,8 @@ def test_does_nothing_when_kv_cache_events_is_not_lmcache_kv_events( self, mock_connector ): """Test that method returns early when kv_cache_events is not - LMCacheKVEvents.""" + LMCacheKVEvents. + """ # Create a mock object that is not LMCacheKVEvents fake_events = MagicMock() connector_output = KVConnectorOutput(kv_cache_events=fake_events) @@ -675,8 +676,7 @@ def __init__(self, cycle_num): assert mock_connector._kv_cache_events is None def test_lmcache_kv_events_aggregation(self): - """ - Test LMCacheKVEvents aggregation across TP ranks using + """Test LMCacheKVEvents aggregation across TP ranks using KVOutputAggregator (used by MultiprocExecutor). """ from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index 43de313821be..db9cf9417165 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -179,7 +179,6 @@ def batch_register_memory(self, buffer_addresses, capacities) -> int: def test_align_transfer_regions_uses_layer_name_occurrences(): """Repeated layer names should align by occurrence order.""" - local_regions = [ TransferRegion( layer_name="model.layers.1.self_attn", @@ -232,7 +231,6 @@ def test_align_transfer_regions_uses_layer_name_occurrences(): @pytest.mark.asyncio async def test_build_transfer_params_separates_prefill_pp_layers(): """Each producer PP stage should send only its registered layer shard.""" - worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) worker.async_zmq_ctx = MagicMock() worker.is_kv_consumer = True @@ -368,7 +366,6 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata( monkeypatch, ): """Producer sends its PP layer shard to the matching consumer layer address.""" - monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" @@ -447,7 +444,6 @@ async def run_in_executor(self, executor, func, *args): def test_basic_interface(): """Unit test for basic MooncakeConnector interface functionality.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -497,7 +493,6 @@ def test_basic_interface(): def test_prompt_less_than_block_size(): """Test that we can handle case where prompt is < block.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -536,7 +531,6 @@ def test_prompt_less_than_block_size(): @pytest.fixture def bootstrap_server(): """Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server.""" - port = get_open_port() server = MooncakeBootstrapServer("127.0.0.1", port) server.start() @@ -546,12 +540,10 @@ def bootstrap_server(): @pytest.mark.asyncio async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer): - """ - Tests the bootstrap server's api for worker registration and querying. + """Tests the bootstrap server's api for worker registration and querying. Validates DP/TP/PP rank indexing and error handling for duplicate registrations. """ - import httpx base_url = f"http://127.0.0.1:{bootstrap_server.port}" @@ -716,13 +708,11 @@ def test_get_mooncake_bootstrap_addr_selects_expected_host( def test_scheduler_request_finished(): - """ - Tests the scheduler-side logic when a request finishes. + """Tests the scheduler-side logic when a request finishes. Differentiates between 'Finished' (requires transfer) and 'Aborted' (immediate free). """ - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" ) @@ -751,7 +741,6 @@ def test_scheduler_request_finished(): @contextlib.contextmanager def patch_worker_dependencies(): """Helper to mock all distributed and network dependencies for Worker tests.""" - with ( patch( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine", @@ -820,7 +809,6 @@ async def test_receive_kv_selects_remote_pp_workers( expected_addrs: list[str], ): """Decode workers should not hard-code producer pp_rank 0.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -875,7 +863,6 @@ async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]): def test_resolve_need_send_accounts_for_remote_tp_fanout(): """Producer-side completion waits for every paired consumer TP pull.""" - worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) worker.async_zmq_ctx = MagicMock() worker.is_kv_consumer = True @@ -960,13 +947,11 @@ async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers(): FakeMooncakeWrapper, ) async def test_kv_producer(monkeypatch): - """ - Simulates a Producer Worker (Prefiller) receiving a transfer request + """Simulates a Producer Worker (Prefiller) receiving a transfer request from a Consumer (Decoder). Verifies memory offset calculation: ptr = base_addr + block_id * block_len. """ - monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" @@ -1133,12 +1118,10 @@ def expected_transfers(src_base, dst_base, src_blocks, dst_blocks): @pytest.mark.asyncio async def test_kv_consumuer(monkeypatch): - """ - Simulates a Consumer Worker (Decoder) initiating a pull from a Producer. + """Simulates a Consumer Worker (Decoder) initiating a pull from a Producer. Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ. """ - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -1211,7 +1194,6 @@ async def test_kv_consumuer(monkeypatch): @pytest.mark.asyncio async def test_worker_get_finished_timeout(monkeypatch): """Tests the cleanup mechanism for requests.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" ) @@ -1263,7 +1245,6 @@ async def test_worker_get_finished_timeout(monkeypatch): ) def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool): """Tests the memory registration logic with the underlying Mooncake engine.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -1344,7 +1325,6 @@ def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): """Mixed MLA+Eagle caches should register by byte length, not shape.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -1408,8 +1388,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): ) @pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"]) async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size): - """ - Tests heterogeneous TP support in the producer transfer path. + """Tests heterogeneous TP support in the producer transfer path. Verifies correct pointer and offset calculation when producer TP=2 sends to consumer with TP=1 (P>D) or TP=4 (P D TP=1: one D rank receives; dst_offset based on P rank - P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank """ - P_TP_SIZE = 2 P_TP_RANK = 0 LOCAL_BLOCK_LEN = 4096 diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py index f3739ef4301a..37f33ce35392 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py @@ -217,7 +217,6 @@ def test_metadata_hma_block_ids(): ) async def test_build_transfer_params_multi_group_trimming(monkeypatch): """_build_transfer_params trims per-group blocks when local > remote.""" - monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" @@ -314,7 +313,6 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch): ) async def test_build_transfer_params_group_count_mismatch(monkeypatch): """_build_transfer_params reports an error when group counts differ.""" - monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer" diff --git a/tests/v1/kv_connector/unit/test_mooncake_stats.py b/tests/v1/kv_connector/unit/test_mooncake_stats.py index a20fcb505330..6e2bb59fa4f2 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_stats.py +++ b/tests/v1/kv_connector/unit/test_mooncake_stats.py @@ -199,7 +199,8 @@ def _bare_worker() -> MooncakeConnectorWorker: """Construct a MooncakeConnectorWorker skipping __init__ (full init requires a live TransferEngine). Only the attributes touched by the methods under test are populated; role flags and async_zmq_ctx keep __del__'s shutdown - path a no-op.""" + path a no-op. + """ worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) worker.xfer_stats = MooncakeKVConnectorStats() worker.engine = MagicMock() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index f3ffa55a5439..7e8e532fffb9 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -77,7 +77,8 @@ def test_store_tp_does_not_override_backend_kv_cache_layout(extra_config): def _make_kv_cache_config() -> KVCacheConfig: """Single-group full-attention KVCacheConfig — enough for the connector - constructor's validate() pass.""" + constructor's validate() pass. + """ spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) return KVCacheConfig( num_blocks=4, @@ -692,7 +693,8 @@ def _poll_lookup(client, req_id, num_tokens=128, block_hashes=(), timeout=5.0): def _gated_recv(gate: threading.Event, value: int): """Mock recv side-effect that blocks until ``gate`` is set, so the - executor's lookup can be held pending deterministically.""" + executor's lookup can be held pending deterministically. + """ def recv(): gate.wait() @@ -703,7 +705,8 @@ def recv(): def test_lookup_key_client_non_block_lookup_async(): """Non-blocking lookup defers to the executor: None first, hit once the - Future resolves.""" + Future resolves. + """ vllm_config = _make_vllm_config() with patch( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 8767b1c0d983..583c24958831 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -33,7 +33,8 @@ def _mamba_align(block_size=32): def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for - the test fixtures.""" + the test fixtures. + """ block_sizes = [g.kv_cache_spec.block_size for g in groups] scheduler_block_size = lcm(*block_sizes) return MooncakeStoreCoordinator( @@ -143,7 +144,8 @@ def test_coordinator_single_full_attention_no_hits(): def test_coordinator_single_swa_tautological_pool_masks_pre_window(): """SWA tautological-pool: hit_length spans full prefix, mask is - tail-window only.""" + tail-window only. + """ groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) # 4 chunks * 16 tokens @@ -190,7 +192,8 @@ def test_coordinator_hybrid_hole_in_full_clips_both(): def test_coordinator_group_block_size_double_hash(): """Group block_size=32 over hash_block_size=16 hashes: adjacent - hashes merge before pool lookup.""" + hashes merge before pool lookup. + """ groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _full(32)), @@ -214,7 +217,8 @@ def test_coordinator_group_block_size_double_hash(): def test_coordinator_fine_grained_partial_tail_hit(): """K3 shape: FA + mamba-align, block_size=32 over hash_block_size=16. When both groups have the sub-block boundary hash, the reconciled hit lands on - the hash boundary (48), not the block boundary (32).""" + the hash boundary (48), not the block boundary (32). + """ groups = [ KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), @@ -234,7 +238,8 @@ def test_coordinator_fine_grained_partial_tail_hit(): def test_coordinator_fine_grained_clips_when_one_group_missing_tail(): """If only one group has the sub-block boundary, min-convergence clips the - reconciled hit back to the block boundary (32).""" + reconciled hit back to the block boundary (32). + """ groups = [ KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), @@ -275,7 +280,8 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): """Hybrid full-attn(block=32) + SWA(block=8, sw=8). lcm=32. With aligned=64 the SWA group should mark exactly the blocks ending at 32 and 64 (i.e. blocks 3 and 7 at block_size=8); the rest can never - participate in any future hit.""" + participate in any future hit. + """ full = _full(32) swa = _swa(block_size=8, sliding_window=8) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -321,7 +327,8 @@ def make(dcp_world_size): def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm - boundary should now span two SWA tail blocks.""" + boundary should now span two SWA tail blocks. + """ full = _full(32) swa = _swa(block_size=8, sliding_window=16) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -351,7 +358,8 @@ def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa(): """DSV4-shaped: full-MLA(B=256) + 4 SWA groups with B in {64, 64, 4, 8} and varied sliding windows. lcm=256, hash_block_size=4. Two lcm segments (aligned_len=512). Validates that the tile-once strategy produces the - expected per-segment tail-window pattern, repeated.""" + expected per-segment tail-window pattern, repeated. + """ full_mla = _full(block_size=256) swa_64_sw128 = _swa(block_size=64, sliding_window=128) swa_64_sw512 = _swa(block_size=64, sliding_window=512) @@ -385,7 +393,8 @@ def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa(): def test_store_mask_fast_path_all_block_sizes_equal_lcm(): """When every non-full-attn group already aligns to lcm_block_size, the - fast path returns all-True without invoking find_longest_cache_hit.""" + fast path returns all-True without invoking find_longest_cache_hit. + """ full = _full(block_size=64) swa = _swa(block_size=64, sliding_window=128) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -398,7 +407,8 @@ def test_store_mask_fast_path_all_block_sizes_equal_lcm(): def test_store_mask_fast_path_single_attention_group(): """Two groups sharing the same SWA spec collapse to one attention group; - no lcm filter applies, every chunk is True.""" + no lcm filter applies, every chunk is True. + """ swa = _swa(block_size=16, sliding_window=32) groups = [KVCacheGroupSpec(["L0"], swa), KVCacheGroupSpec(["L1"], swa)] coord = _make_coord(groups, hash_block_size=16) @@ -412,7 +422,8 @@ def test_store_mask_fast_path_single_attention_group(): def _retention_groups(): """Hybrid full-attn(block=32) + SWA(block=8, sw=8); lcm=32. The SWA group - densely keeps one tail block per 32-token boundary.""" + densely keeps one tail block per 32-token boundary. + """ full = _full(32) swa = _swa(block_size=8, sliding_window=8) return [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -420,7 +431,8 @@ def _retention_groups(): def test_store_mask_dense_default_matches_every_lcm_boundary(): """retention_interval=None (default) keeps the SWA tail at every lcm - boundary: tokens 32/64/96/128 -> chunks 3/7/11/15.""" + boundary: tokens 32/64/96/128 -> chunks 3/7/11/15. + """ coord = _make_coord(_retention_groups(), hash_block_size=8) masks = coord.store_mask(128) assert masks[0] is None @@ -430,7 +442,8 @@ def test_store_mask_dense_default_matches_every_lcm_boundary(): def test_store_mask_retention_interval_sparsifies_swa_tails(): """retention_interval=64 keeps an SWA tail once per 64-token segment (chunks 7 and 15) instead of every 32 tokens, dropping the mid-segment - boundaries at 32 and 96.""" + boundaries at 32 and 96. + """ coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64) masks = coord.store_mask(128) assert masks[0] is None # full attn unaffected @@ -439,7 +452,8 @@ def test_store_mask_retention_interval_sparsifies_swa_tails(): def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): """retention_interval=0 drops all segment tails; only the latest replay - boundary (capped at num_prompt-1, aligned down to lcm) is retained.""" + boundary (capped at num_prompt-1, aligned down to lcm) is retained. + """ coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=0) # No replay info -> nothing reachable for the SWA group. assert coord.store_mask(128)[1] == [False] * 16 @@ -450,7 +464,8 @@ def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): def test_store_mask_retention_interval_keeps_segment_and_replay_tails(): """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay - boundary tail (num_prompt=100 -> chunk 11) coexist.""" + boundary tail (num_prompt=100 -> chunk 11) coexist. + """ coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64) masks = coord.store_mask(128, num_prompt_tokens=100) assert masks[1] == [i in (7, 11, 15) for i in range(16)] @@ -488,7 +503,8 @@ def test_store_mask_excludes_mamba_groups_lookup_unaffected(): nulled/freed; speculative blocks relocate), so the positional normal save must never cover mamba chunks — regardless of retention. Lookups still probe mamba boundaries: their keys come from the pinned snapshot/CoW - hand-off path instead.""" + hand-off path instead. + """ groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _mamba_align(16)), @@ -509,7 +525,8 @@ def test_store_mask_excludes_mamba_groups_lookup_unaffected(): def test_lookup_with_eagle_pops_last_full_attention_block(): """Sanity: with use_eagle, find_longest_cache_hit drops the last block. - Pairs with the load_mask test below to lock the round-trip contract.""" + Pairs with the load_mask test below to lock the round-trip contract. + """ groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16, use_eagle=True) hs = _hashes(4) @@ -550,7 +567,8 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention(): def test_load_mask_with_eagle_hybrid_full_plus_swa(): """Hybrid (FullAttn + SWA) with eagle: load_mask must cover every chunk in [0, token_len) for the FullAttn group; SWA group keeps its - tail-window mask.""" + tail-window mask. + """ groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _swa(16, 32)), diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index a93698bba954..85d49700895d 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """End-to-end save->lookup test for MooncakeStoreConnector on a hybrid -(SWA + Full) attention config, using a dict-backed mock store.""" +(SWA + Full) attention config, using a dict-backed mock store. +""" import sys import threading @@ -147,8 +148,7 @@ def _build_worker_with_dict_store(vllm_config, kv_cache_config, store): def test_e2e_swa_plus_full_save_then_lookup_hits(): - """ - E2E: build a SWA+Full hybrid worker, save all blocks via the sending + """E2E: build a SWA+Full hybrid worker, save all blocks via the sending thread (synchronously), then verify lookup returns the full hit length. Also verify that evicting SWA's early blocks (outside its window) still allows a full hit because the window covers the tail. @@ -283,7 +283,8 @@ def _fake_thread_init(*args, **kwargs): def test_recv_skips_swa_blocks_before_window(): """Producer stored every block for both groups; consumer must only fetch - SWA blocks within the sliding window, not the head.""" + SWA blocks within the sliding window, not the head. + """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) # sliding_window=32, block_size=16 → 2 contiguous blocks within window. swa = SlidingWindowSpec( @@ -351,7 +352,8 @@ def batch_get_into_multi_buffers(self, keys, addrs, sizes): def test_chunked_token_database_hash_block_size_smaller_than_block_size(): """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens - keys each chunk by its ending fine hash, including a partial tail.""" + keys each chunk by its ending fine hash, including a partial tail. + """ md = KeyMetadata("m", 0, 0, 0, 0, group_id=3) db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4) db.set_kv_caches_base_addr([0]) @@ -384,7 +386,8 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the partial tail is offloaded for both groups under the boundary sub-hash. The full-attention block is read from the request block table; the mamba block - is the core-provided CoW target, not block_ids.""" + is the core-provided CoW target, not block_ids. + """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -463,7 +466,8 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): def test_offload_syncs_event_before_put(): """An offload-carrying meta synchronizes its CoW-fence event before the store put reads the blocks, then completes in one pass and drains the - completion counter.""" + completion counter. + """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -534,7 +538,8 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): offload must persist every FA block up to the boundary — the normal save floors to the lcm, so those blocks are otherwise never written and the consumer's per-group lookup would miss. The mamba boundary block still - reads the core-provided CoW target.""" + reads the core-provided CoW target. + """ full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -805,7 +810,8 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): """A per-request ring group (CircularBufferSpec, capacity as block_size) sits beside the paged group: worker setup tolerates its block size (the DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged - blocks only, and lookup hits on the paged group alone.""" + blocks only, and lookup hits on the paged group alone. + """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) ring = CircularBufferSpec( block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8 diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 7ced38d65a41..af2dfd608313 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -536,7 +536,8 @@ def _make_pending_load_unfinished_request( def _make_pending_load_scheduler_output() -> SimpleNamespace: """scheduler_output for a step where req-0 is parked on a pending load - (not in scheduled_new_reqs or scheduled_cached_reqs).""" + (not in scheduled_new_reqs or scheduled_cached_reqs). + """ return SimpleNamespace( finished_req_ids=set(), preempted_req_ids=set(), @@ -712,7 +713,8 @@ def test_running_request_not_in_resumed_req_ids_appends_blocks(): def test_resumed_request_in_resumed_req_ids_replaces_blocks(): """A request the scheduler marks resumed gets the FULL block table in new_block_ids and must REPLACE the tracker's blocks (not append), even if - a stale tracker from before preemption is still present.""" + a stale tracker from before preemption is still present. + """ scheduler = _make_bare_scheduler() _make_resumed_unfinished_request( scheduler, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index dea4178ff3e8..9646331fe014 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -432,7 +432,8 @@ def _patch_worker_runtime( def test_pool_key_to_string_without_prefix_is_unchanged(): """Default (empty) cache_prefix keeps keys byte-identical to the - historical unprefixed format so existing deployments keep their hits.""" + historical unprefixed format so existing deployments keep their hits. + """ key = PoolKey(KeyMetadata("test-model", 0, 0, 0, 0), "deadbeef") assert ( key.to_string() == "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@deadbeef" @@ -573,7 +574,8 @@ def test_tp_shared_receiving_reads_each_local_store_shard(): def test_pool_key_cache_prefix_namespaces_and_disambiguates(): """A non-empty cache_prefix is prepended, and two instances with - different prefixes never collide on identical block hashes.""" + different prefixes never collide on identical block hashes. + """ md_a = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="depA") md_b = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="depB") @@ -589,7 +591,8 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates(): def test_default_local_buffer_size_matches_pr40900(): """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode - patch preserves it (and the JSON key) so unchanged PR-40900 configs work.""" + patch preserves it (and the JSON key) so unchanged PR-40900 configs work. + """ assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 @@ -1011,7 +1014,8 @@ def test_store_sending_thread_skips_null_sparse_group_blocks(): def test_partial_tail_offload_skips_cap_omitted_mamba_group(): """A Mamba group omitted by the handoff cap must not fall back to the - connector's positional block table.""" + connector's positional block table. + """ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1139,7 +1143,8 @@ def test_partial_tail_put_failure_activates_pressure_gate(): def test_normal_save_excludes_mamba_group_and_null_blocks(): """The positional normal save must never cover mamba chunks (align-mode block tables are not append-only) and must skip chunks whose block is the - reserved null block, for any group.""" + reserved null block, for any group. + """ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1200,7 +1205,8 @@ def test_normal_save_excludes_mamba_group_and_null_blocks(): def test_block_aligned_snapshot_offload_uses_provided_block(): """A block-aligned hand-off (sparse-retention mamba checkpoint) uploads exactly the core-pinned block under the boundary-end hash key; the - positional block table entry for that chunk must be ignored.""" + positional block table entry for that chunk must be ignored. + """ store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) @@ -1228,7 +1234,8 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): def test_mixed_snapshot_and_sub_block_offloads(): """A retention snapshot and the prompt-end sub-block CoW tail can arrive in one hand-off; the sub-block path covers FA gap blocks but reads the - mamba boundary only from the CoW block, never positionally.""" + mamba boundary only from the CoW block, never positionally. + """ store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) @@ -1268,7 +1275,8 @@ def test_mixed_snapshot_and_sub_block_offloads(): def test_snapshot_offload_skips_null_handoff_block(): """A hand-off the core could not materialize (null block) carries no - committed state; persisting it would poison the boundary key.""" + committed state; persisting it would poison the boundary key. + """ store = MagicMock() thread = _make_partial_tail_send_thread(store) @@ -1596,7 +1604,8 @@ def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_segment(): # noqa: E501 """Without a preferred_segment the SendingThread still forwards a (default-constructed) ReplicateConfig so the C++ side always sees a - well-defined config object.""" + well-defined config object. + """ store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.return_value = [256, 256] @@ -2106,7 +2115,8 @@ def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch(): def test_recv_thread_skips_split_when_budget_holds_all_keys(): """PR-36 removed the count-based split trigger; with budget for 3 keys, - all three should be requested in a single call.""" + all three should be requested in a single call. + """ store = MagicMock() store.batch_get_into_multi_buffers.return_value = [256, 256, 256] thread = _make_store_recving_thread( @@ -2320,7 +2330,8 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled( monkeypatch, ): """enable_offload=False zeroes out the disk budget so we don't generate - redundant owner GET-RPCs.""" + redundant owner GET-RPCs. + """ store = MagicMock() store.setup.return_value = 0 _install_fake_mooncake(monkeypatch, store) @@ -2995,7 +3006,8 @@ def __init__(self) -> None: def test_store_sending_thread_clamps_token_len_to_lcm(): """Partial chunks past the last lcm boundary aren't stored — cache hits - are always lcm-aligned (mirrors HybridKVCacheCoordinator).""" + are always lcm-aligned (mirrors HybridKVCacheCoordinator). + """ store = MagicMock() store.batch_is_exist.return_value = [0, 0] store.batch_put_from_multi_buffers.return_value = [256, 256] @@ -3020,7 +3032,8 @@ def test_store_sending_thread_clamps_token_len_to_lcm(): def test_store_sending_thread_skips_when_token_len_below_lcm(): """Requests shorter than lcm_block_size cannot produce any aligned chunk, - so neither the existence check nor the put should be issued.""" + so neither the existence check nor the put should be issued. + """ from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec store = MagicMock() @@ -3810,7 +3823,8 @@ def test_lookup_full_hit_with_eagle_pops_once_not_twice(): """Eagle already leaves the last block for the drafter, so a full-prompt re-derivation must never fire for eagle-governed hits: firing would anchor the search one block lower and pop a second - block, regressing the hit by an extra producer boundary.""" + block, regressing the hit by an extra producer boundary. + """ worker = _make_bare_worker(block_size=16) worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( worker._kv_cache_groups, @@ -3967,7 +3981,8 @@ def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable(): the full hit into 48, whose SWA window needs the never-written block 1, so every load failed and the recompute re-entered the same lookup. The re-derivation must report that no stored boundary below the request end - is usable.""" + is usable. + """ from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec worker = _make_bare_worker(block_size=16) @@ -3989,7 +4004,8 @@ def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable(): def test_lookup_swa_single_group_returns_full_when_tail_window_present(): """Single-SWA, sliding_window=32 (= 2 blocks): producer stored only the tail. Coordinator-driven lookup returns full prefix even though the - pre-window blocks are absent.""" + pre-window blocks are absent. + """ from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec worker = _make_bare_worker(block_size=16) @@ -4385,7 +4401,8 @@ def test_config_defaults_to_embedded(): def test_config_pr40900_unchanged(tmp_path): """A literal PR-40900 config (no mode, no enable_offload, no preferred_segment) - parses without raising and resolves to embedded mode.""" + parses without raising and resolves to embedded mode. + """ config_path = _write_mooncake_config( tmp_path, { @@ -4493,7 +4510,8 @@ def test_config_zero_local_buffer(): def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch): """standalone-store + disk: global_segment_size=0, enable_offload=True, preferred_segment set. Assert setup() positional args, ReplicateConfig - wiring, and that the disk-offload buffer budget is allocated.""" + wiring, and that the disk-offload buffer budget is allocated. + """ store = MagicMock() store.setup.return_value = 0 fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store) @@ -4540,9 +4558,10 @@ def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch): def test_topology_embedded_cpu_only(tmp_path, monkeypatch): - """embedded + CPU-only: no mode key (defaults to embedded), + """Embedded + CPU-only: no mode key (defaults to embedded), global_segment_size>0, enable_offload absent, no preferred_segment. - This is the PR-40900 baseline recipe.""" + This is the PR-40900 baseline recipe. + """ store = MagicMock() store.setup.return_value = 0 fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store) @@ -4731,7 +4750,8 @@ def test_store_worker_close_swallows_store_errors(): def test_blob_block_hashes_wire_roundtrip(): """The lookup wire format sends a ``hash_len`` frame plus the raw hashes concatenated back-to-back; the server rebuilds them through a zero-copy - ``BlobBlockHashes`` view over the frame buffer.""" + ``BlobBlockHashes`` view over the frame buffer. + """ hashes = [BlockHash(bytes([i]) * 16) for i in range(5)] hash_len = len(hashes[0]) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 836f9407bc66..eae9f71ee3fa 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -294,7 +294,6 @@ def create_vllm_config( def test_write_mode_saves_local_block_ids(): """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save.""" - # Setup Scheduler and Request vllm_config = create_vllm_config(role="kv_producer") scheduler = create_scheduler(vllm_config) @@ -405,7 +404,6 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids(): def test_read_mode_loads_remote_block_ids(): """Read mode loads remote block ids into local cache mapping.""" - # Setup Scheduler and Request vllm_config = create_vllm_config(role="kv_consumer", read_mode=True) scheduler = create_scheduler(vllm_config) @@ -631,7 +629,6 @@ def test_register_kv_caches(mock_parallel_groups): ) def test_moriio_handshake_returns_metadata(mock_parallel_groups): """MoRIIO handshake socket returns valid agent metadata over ZMQ.""" - ROLE = "kv_consumer" vllm_config = create_vllm_config(role=ROLE) # Create test kv cache tensors using KVCacheSpec layout @@ -695,7 +692,8 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups): def test_resolve_host_ip_prefers_extra_config(): """An explicit ``host_ip`` in kv_connector_extra_config overrides get_ip() (so an external router can advertise a routable/internal address); an - absent or empty value falls back to get_ip().""" + absent or empty value falls back to get_ip(). + """ assert resolve_host_ip({"host_ip": "10.0.0.7"}) == "10.0.0.7" fallback = get_ip() @@ -749,7 +747,8 @@ def _read_scheduler( def test_hma_blocks_per_sw_two_groups(): """A Full + SlidingWindow config runs with HMA and computes block budgets - correctly""" + correctly + """ scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) assert scheduler._is_hma_required is True # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention. @@ -767,7 +766,8 @@ def test_hma_blocks_per_sw_two_groups(): def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): """_is_hma_required tracks both the KV cache groups and the --disable-hybrid-kv-cache-manager flag. When HMA is off, - get_exchange_clipped_blocks must be a no-op.""" + get_exchange_clipped_blocks must be a no-op. + """ config = ( _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config() ) @@ -780,7 +780,8 @@ def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): def test_non_sliding_window_hybrid_is_rejected(): """A hybrid group that is not sliding-window (e.g. chunked-local - attention) must fail closed rather than be silently mistransferred.""" + attention) must fail closed rather than be silently mistransferred. + """ full_spec = FullAttentionSpec( block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 ) @@ -814,7 +815,8 @@ def test_non_sliding_window_hybrid_is_rejected(): def test_token_count_basis_uses_full_attention_group(): """Chunked-prefill token counting must use an unclipped full-attention - group (blocks_per_sw == 0), not a clipped sliding-window group.""" + group (blocks_per_sw == 0), not a clipped sliding-window group. + """ scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) # Group 0 is the full-attention group assert scheduler._full_attn_group_idx == 0 @@ -823,7 +825,8 @@ def test_token_count_basis_uses_full_attention_group(): def test_get_exchange_clipped_blocks_clips_only_sw_group(): """get_exchange_clipped_blocks keeps the full attn group intact and clips the - sliding-window group to its window tail.""" + sliding-window group to its window tail. + """ scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) full = [10, 11, 12, 13, 14] sw = [20, 21, 22, 23, 24] @@ -834,7 +837,8 @@ def test_get_exchange_clipped_blocks_clips_only_sw_group(): def test_metadata_hma_block_ids_preserved_per_group(): """add_new_req stores per-group (BlockIds) block lists unchanged for both - read and write, so the hybrid group structure is retained.""" + read and write, so the hybrid group structure is retained. + """ metadata = MoRIIOConnectorMetadata() # Assume: diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py index e38ec3f7e9be..76a33cff82c2 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -110,7 +110,8 @@ def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0" def build_decode_workers(cfg: PDConfig) -> list: """Decode workers that issue reads for one prefill instance. A TP decode instance reads from every tp rank; a DP decode instance from each dp-rank - worker. Reused across requests so per-worker round-robin state advances.""" + worker. Reused across requests so per-worker round-robin state advances. + """ if cfg.d_tp > 1: return [ make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0) @@ -125,7 +126,8 @@ def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter: """Drive the REAL _resolve_read_source over a fair input stream; tally which prefill GPU each read targets. The only modelled assumption is the proxy contract -- each owner dp-rank delivered equally (``for owner_dp in - range(p_dp)``); no proxy algorithm is reproduced.""" + range(p_dp)``); no proxy algorithm is reproduced. + """ workers = build_decode_workers(cfg) hits: Counter = Counter() for _ in range(rounds): diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 737d03ac6900..e6a78063be3d 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -171,7 +171,7 @@ def test_register_finished_partial_tail_notifies_every_connector(): @pytest.fixture def mc() -> MultiConnector: - """MultiConnector using two mocked connectors""" + """MultiConnector using two mocked connectors.""" mock_connector_config = { "kv_connector": "MockConnector", "kv_role": "kv_both", @@ -215,8 +215,7 @@ def _compare_directories(dir1: Path, dir2: Path) -> bool: def test_multi_example_connector_consistency(): - """ - Tests that MultiConnector with two ExampleConnectors saves + """Tests that MultiConnector with two ExampleConnectors saves identical KV cache data to separate storage locations. """ storage_1_path = Path("storage_1/") @@ -455,8 +454,7 @@ def test_engine_id_conflict(): def test_multi_connector_handle_preemptions_integration(): - """ - Integration test: verify MultiConnector delegates handle_preemptions + """Integration test: verify MultiConnector delegates handle_preemptions to all sub-connectors. Uses TestExampleConnector which logs all method calls to temp files. @@ -557,7 +555,8 @@ def test_build_kv_connector_stats_with_empty_dict(self): def test_build_kv_connector_stats_reconstructs_nixl_stats(self): """Test that NixlConnector stats are properly reconstructed with - correct data.""" + correct data. + """ serialized_data = { "NixlConnector": { "transfer_duration": [1.5, 2.3], @@ -863,9 +862,7 @@ def test_is_empty_with_multiple_connectors(self): def test_multi_connector_overrides_all_base_methods(): - """ - Ensure MultiConnector overrides all public methods from KVConnectorBase_V1. - """ + """Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.""" # These are fine to inherit from KVConnectorBase_V1 # TODO(https://github.com/vllm-project/vllm/pull/31811): Remove # get_kv_connector_kv_cache_events from INHERITED_OK once implemented. @@ -1045,12 +1042,10 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector: def test_multi_connector_hma_support_detection(): - """ - At runtime, _all_support_hma is True only when every sub-connector + """At runtime, _all_support_hma is True only when every sub-connector implements SupportsHMA. Test all combinations of HMA / non-HMA sub-connectors. """ - assert supports_hma(MultiConnector) # -- All non-HMA connectors => _all_support_hma is False -- @@ -1092,8 +1087,7 @@ def test_divergent_local_hybrid_hit_capability_is_conservative(): not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM" ) def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch): - """ - When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA + """When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA (MockConnector) sub-connectors, verify that: 1. The scheduler's MultiConnector has _all_support_hma == False. 2. vLLM auto-disables the hybrid KV cache manager (no preference expressed by user) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index e41bd2644a9a..d02be5fd564e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -84,8 +84,7 @@ @pytest.fixture(scope="module", autouse=True) def clear_kv_transfer(): - """ - The test cases in this file use `VLLM_ENABLE_V1_MULTIPROCESSING=0`, + """The test cases in this file use `VLLM_ENABLE_V1_MULTIPROCESSING=0`, causing the global variable `_KV_CONNECTOR_AGENT` to be assigned but never deleted. @@ -245,7 +244,6 @@ def _make_fake_nixl_pkg(): def test_basic_interface(): """Unit test for basic NixlConnector interface functionality.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -284,8 +282,7 @@ def test_basic_interface(): def test_prompt_less_than_block_size(): - """ - Test that we can handle case where prompt is < block. + """Test that we can handle case where prompt is < block. In this case, the P worker will still send remote_block_ids of the partial block. The D worker should schedule an async read @@ -321,7 +318,8 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): """A remote-prefill request added with abort_immediately=True should be added to the scheduler's waiting queue then immediately aborted, so the NIXL connector's request_finished hook enqueues an empty recv to notify - the prefill instance to free its blocks.""" + the prefill instance to free its blocks. + """ from vllm.v1.request import RequestStatus scheduler = create_scheduler(create_vllm_config()) @@ -407,7 +405,6 @@ def test_prefill_exports_cached_tokens_in_kv_transfer_params(): ) def test_kv_transfer_handshake(dist_init): """Unit test for basic NixlConnector interface functionality.""" - # Test setup, we creates a scheduler that contains a NixlConnector # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from # all workers of the instance. @@ -830,7 +827,6 @@ def test_async_load_kv( prefill_tp_size, ): """Test that NixlConnector's start_load_kv should be non-blocking.""" - vllm_config = create_vllm_config() vllm_config.parallel_config.tensor_parallel_size = decode_tp_size @@ -885,8 +881,7 @@ def test_async_load_kv( def test_prefill_tp_size_greater_than_decode_tp_size( self, local_tp_size: int, default_vllm_config, dist_init, monkeypatch ): - """ - Verify remote TP > local TP handshake succeeds with different + """Verify remote TP > local TP handshake succeeds with different remote configurations. """ monkeypatch.setattr( @@ -959,8 +954,7 @@ def check_handshake(remote_tp_size: int): def test_prefill_tp_size_greater_than_decode_tp_size_mla( self, default_vllm_config, dist_init ): - """ - Verify remote TP > local TP handshake succeeds with different + """Verify remote TP > local TP handshake succeeds with different remote configurations for an MLA model. """ vllm_config = create_vllm_config() @@ -1066,7 +1060,6 @@ def test_concurrent_load_kv( dist_init, ): """Test that multiple start_load_kv calls should occur concurrently.""" - vllm_config = create_vllm_config() # Test worker role in decode server. @@ -1127,8 +1120,7 @@ def test_concurrent_load_kv( def test_handshake_fails_on_kv_cache_layout_mismatch( self, default_vllm_config, dist_init ): - """ - Verify that adding a remote agent fails if kv_cache_layout differs. + """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ vllm_config = create_vllm_config() @@ -1185,8 +1177,7 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( self, default_vllm_config, dist_init ): - """ - Verify that adding a remote agent fails if kv_cache_layout differs. + """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ vllm_config = create_vllm_config(enable_permute_local_kv=True) @@ -1624,11 +1615,9 @@ def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist def test_kv_connector_stats_aggregation(): - """ - Test KV transfer stats aggregation across TP ranks using + """Test KV transfer stats aggregation across TP ranks using KVOutputAggregator (used by MultiprocExecutor). """ - # Create KVOutputAggregator for 3 workers (simulating TP=3), same thing # done in MultiprocExecutor.execute_model aggregator = KVOutputAggregator(expected_finished_count=3) @@ -1691,11 +1680,9 @@ def test_kv_connector_stats_aggregation(): def test_multi_kv_connector_stats_aggregation(): - """ - Test MultiKVConnectorStats aggregation across TP ranks using + """Test MultiKVConnectorStats aggregation across TP ranks using KVOutputAggregator (used by MultiprocExecutor). """ - aggregator = KVOutputAggregator(expected_finished_count=3) from dataclasses import dataclass @@ -1836,8 +1823,7 @@ def test_scheduler_kv_connector_stats_aggregation(): FakeNixlWrapper, ) def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): - """ - Test lifecycle of an aborted Remote Prefill request hitting the timeout. + """Test lifecycle of an aborted Remote Prefill request hitting the timeout. -----> P | {process request} <-/--- | {result is NOT delivered, eg proxy is down} @@ -2123,8 +2109,7 @@ def test_register_kv_caches( layout, separate_kv_head_groups, ): - """ - Test that register_kv_caches() properly calls nixl_wrapper methods with + """Test that register_kv_caches() properly calls nixl_wrapper methods with correct data. This test verifies: @@ -2133,7 +2118,6 @@ def test_register_kv_caches( 2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing block layout info """ - vllm_config = create_vllm_config(attention_backend=attn_backend) vllm_config.cache_config.kv_cache_layout = layout @@ -2361,17 +2345,14 @@ class FakePlatform(Platform): @classmethod def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]: - """ - Returns a mapping from device_type to a tuple of supported + """Returns a mapping from device_type to a tuple of supported kv_buffer_device for nixl. """ return {"oot": ("oot",)} @classmethod def get_nixl_memory_type(cls) -> str | None: - """ - Returns the nixl memory type for the current platform. - """ + """Returns the nixl memory type for the current platform.""" return "VRAM" @@ -2384,8 +2365,7 @@ def get_nixl_memory_type(cls) -> str | None: def test_kv_buffer_to_nixl_memory_types( default_vllm_config, dist_init, kv_buffer_device, nixl_memory_type ): - """ - Test that register_kv_caches() passes the correct memory types from the + """Test that register_kv_caches() passes the correct memory types from the config to the nixl_wrapper. """ vllm_config = create_vllm_config() @@ -2657,8 +2637,7 @@ def test_transfer_topology_unregister(): FakeNixlWrapper, ) def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init): - """ - Create and schedule a request so that P adds it to in-batch tracking via + """Create and schedule a request so that P adds it to in-batch tracking via the real scheduler, then simulate an abort (request not in next scheduler iteration) and verify the worker no longer tracks it as in-batch. """ @@ -3167,7 +3146,8 @@ def test_split_read_failure_defers_report_until_last_handle( request — nor invalidate its blocks — while the sibling xfer is still in flight: a posted READ cannot be aborted and would DMA into blocks the scheduler could free and reuse. The report happens exactly once, when the - last handle is terminal.""" + last handle is terminal. + """ request_id = "split_read_partial_failure" err_handle, live_handle = 11, 22 connector, worker, wrapper = _make_split_read_connector( @@ -3562,15 +3542,16 @@ def test_compatibility_hash_validation( should_fail, enforce_handshake_compat, ): - """ - Test NIXL compatibility hash validation during handshake. + """Test NIXL compatibility hash validation during handshake. - Parameters: + Parameters + ---------- mismatch_type: description of what is being tested config_overrides: dict of config to override for the remote instance version_override: version dict e.g. {"vllm_version": "0.6.1"} should_fail: whether the handshake should fail enforce_handshake_compat: whether to enforce compatibility checking + """ local_vllm_config = create_vllm_config( model="facebook/opt-125m", @@ -3698,8 +3679,7 @@ def test_compatibility_hash_validation( FakeNixlWrapper, ) def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario): - """ - Test that msgspec decode errors are properly handled during handshake. + """Test that msgspec decode errors are properly handled during handshake. Tests both DecodeError and ValidationError for both decoders: - NixlHandshakePayload decoder diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 1e34340d0f10..d79fcab69819 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -281,7 +281,8 @@ def test_update_state_after_alloc_tracks_cached_blocks_per_group(): """Hybrid SWA+FA requests can have different prefix-cache-hit counts per KV cache group. Each group's count must be tracked independently rather than collapsed to a single scalar, or DCP read-phase alignment - (local_num_computed_blocks) would misalign one of the groups.""" + (local_num_computed_blocks) would misalign one of the groups. + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( NixlPullConnectorScheduler, ) @@ -335,7 +336,8 @@ def test_full_local_hit_is_not_awaited_by_the_scheduler(): owed to the producer -- so the worker must not report it in finished_recving. Doing so trips `assert RequestStatus.is_finished` in _update_from_kv_xfer_finished, which only tolerates a finished recv for a - request that is parked or already done.""" + request that is parked or already done. + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( NixlPullConnectorScheduler, ) @@ -440,7 +442,8 @@ def test_apply_dcp_prefix_caching_matches_global_logical_positions( """_apply_dcp_prefix_caching must operate on logical block IDs (DCP interleaves at logical-block granularity) and always return equal-length local/remote slices, so a later, unconditional _apply_prefix_caching - trim is a guaranteed no-op rather than double-processing the read.""" + trim is a guaranteed no-op rather than double-processing the read. + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -839,7 +842,8 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit( def test_apply_prefix_caching_ssm_unpairable_slots_rejected(): """Local SSM slots can only exceed the remote ones by the position D recomputes itself. A larger excess means the lists aren't - position-aligned: fail loudly rather than transfer into wrong slots.""" + position-aligned: fail loudly rather than transfer into wrong slots. + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -1029,8 +1033,7 @@ def run_test_and_cleanup(): @pytest.mark.cpu_test def test_nixl_metadata_hma_block_ids_structure(): - """ - Test that NixlConnectorMetadata correctly stores block IDs for multiple + """Test that NixlConnectorMetadata correctly stores block IDs for multiple KV cache groups when HMA is enabled. """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( @@ -1156,7 +1159,8 @@ def test_nixl_region_sort_preserves_pipeline_layer_order(): @pytest.mark.cpu_test def test_get_block_descs_ids_hybrid_ssm(): """Test _compute_desc_ids uses per-group strides for hybrid - FA+SSM when ratio=1 (no kernel block size mismatch).""" + FA+SSM when ratio=1 (no kernel block size mismatch). + """ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec worker = _make_mock_worker_for_desc_ids( @@ -1240,7 +1244,8 @@ def test_get_block_descs_ids_uses_per_region_pool_capacity(): @pytest.mark.cpu_test def test_get_block_descs_ids_kernel_block_mismatch(): """Test _compute_desc_ids uses different strides for FA - (kernel blocks) vs SSM (logical blocks) when ratio > 1.""" + (kernel blocks) vs SSM (logical blocks) when ratio > 1. + """ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec ratio = 4 @@ -1271,7 +1276,8 @@ def test_get_block_descs_ids_kernel_block_mismatch(): def test_get_block_descs_ids_hetero_block_size_hybrid(): """With a block-size ratio, FA desc ids are ratio-expanded while SSM desc ids keep the unexpanded logical stride (state blocks are never - sub-split).""" + sub-split). + """ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec worker = _make_mock_worker_for_desc_ids( @@ -1310,7 +1316,8 @@ def _bind_worker_method(worker, name): @pytest.mark.cpu_test def test_map_block_ids_for_block_size_ratio_hybrid(): """Attention groups expand to remote granularity and clip to the remote - coverage; mamba state blocks pass through 1:1.""" + coverage; mamba state blocks pass through 1:1. + """ from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( @@ -1341,7 +1348,8 @@ def test_map_block_ids_for_block_size_ratio_hybrid(): @pytest.mark.cpu_test def test_post_process_zeroes_untransferred_tail(): """The untransferred sub-blocks of the last local block are zeroed on - receive; mamba state caches are untouched by the attention permute.""" + receive; mamba state caches are untouched by the attention permute. + """ from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( @@ -1384,7 +1392,8 @@ def test_post_process_zeroes_untransferred_tail(): @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM - groups with different block counts (kernel mismatch active).""" + groups with different block counts (kernel mismatch active). + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -1465,7 +1474,8 @@ def _make_fake_kv_cache_manager(): @pytest.mark.cpu_test def test_zeroing_block_ids_cover_only_loaded_attention_blocks(): """Only zero-recorded (attention) groups contribute, sliced to the - externally-loaded token range; Mamba state blocks are never zeroed.""" + externally-loaded token range; Mamba state blocks are never zeroed. + """ manager = _make_fake_kv_cache_manager() # Tokens [0, 16) are locally cached; the load covers tokens [16, 56). @@ -1493,7 +1503,8 @@ def take_new_block_ids(self): @pytest.mark.cpu_test def test_failed_load_rezeroes_unwritten_skipped_blocks(): """A failed async load leaves zeroing-skipped blocks unwritten beyond - the valid prefix; they must be zeroed before local recompute.""" + the valid prefix; they must be zeroed before local recompute. + """ from unittest.mock import MagicMock from vllm.v1.core.sched.scheduler import Scheduler @@ -1953,7 +1964,8 @@ def test_logical_to_kernel_block_ids_with_remote_ratio( def test_exchange_clipped_blocks_ssm_single_state(): """In single-state cache modes, SSM lists are reduced to the running state slot: speculative scratch slots, null placeholders and the previous - step's state carry nothing. Attention groups pass through untouched.""" + step's state carry nothing. Attention groups pass through untouched. + """ sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) sched.blocks_per_sw = [0, 0] sched._ssm_spec_blocks = [None, 2] @@ -1980,7 +1992,8 @@ def test_exchange_clipped_blocks_ssm_single_state(): @pytest.mark.cpu_test def test_exchange_clipped_blocks_ssm_positional_states(): """In "all" mode every position holds a state, so only the speculative - slots go; placeholders stay to keep the list position-indexed.""" + slots go; placeholders stay to keep the list position-indexed. + """ sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) sched.blocks_per_sw = [0, 0] sched._ssm_spec_blocks = [None, 2] @@ -1997,7 +2010,8 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed MambaSpec) groups whose layers share the same HMA tensors, with a mamba-aligned unified page and an MLA kernel block smaller than the - logical block.""" + logical block. + """ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -2137,7 +2151,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): """Hybrid MLA+KDA registration: HMA tensors shared by both layer types must be flagged as MLA regions even when a KDA layer registers them first, expose TP-independent kernel-granularity block lens, and build - FA + mamba descriptors for every region.""" + FA + mamba descriptors for every region. + """ from unittest.mock import MagicMock from vllm.config import set_current_vllm_config @@ -2238,7 +2253,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): def test_push_write_hybrid_mla_replicates_attention(): """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be written to every covered D rank (replicated MLA latent) while SSM state - is written per-rank through the split handles.""" + is written per-rank through the split handles. + """ import threading from collections import defaultdict from unittest.mock import MagicMock diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index a4bfa9ab83ae..dab6432b5ca3 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -24,7 +24,8 @@ class _RecordingNixl: """Minimal NIXL wrapper stand-in that records descriptor lists and - prepared transfers so tests can resolve desc ids to byte ranges.""" + prepared transfers so tests can resolve desc ids to byte ranges. + """ def __init__(self, *args, **kwargs): self.dlists: dict[int, np.ndarray] = {} @@ -593,7 +594,8 @@ def _register_remote_agents(worker, metadata, tp_size): def _owned_byte_ranges(worker, group_logical_ids): """Byte ranges owned by a request: for each HMA region tensor, every - logical block id of every group maps to one unified page.""" + logical block id of every group maps to one unified page. + """ unified_page = worker._test_unified_page bases = [t.data_ptr() for t in worker._test_tensors] owned = [] @@ -631,7 +633,8 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and TP-sharded KDA state. Every local descriptor of the request's reads must - stay within its own blocks.""" + stay within its own blocks. + """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -698,7 +701,8 @@ def _resolve( """Resolve a desc id to (region, kind, token_start) where kind is 'attn' (desc-page sized, sub-block-aligned, in the request's attention blocks) or 'mamba'. token_start is the request-relative token offset, so local - and remote are comparable even when their kernel blocks differ in size.""" + and remote are comparable even when their kernel blocks differ in size. + """ addr, length, _ = (int(x) for x in desc_arr[int(idx)]) for region, base in enumerate(bases): off = addr - base @@ -892,7 +896,8 @@ def _run_hetero_case( def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): """Sweep prompt lengths across block-boundary residues for several hetero-ppl geometries; assert neighbor-safety, token alignment, and - coverage of every transferred kernel block.""" + coverage of every transferred kernel block. + """ _run_hetero_case( local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens ) @@ -911,7 +916,8 @@ def test_hetero_ppl_with_block_size_ratio(num_tokens): 2). The transfer is clipped at remote sub-block granularity by the pairing and front-trimmed by _apply_prefix_caching, so the untransferred tail can span both a partial block and whole blocks — - the case each of the two former zeroing paths handled only half of.""" + the case each of the two former zeroing paths handled only half of. + """ _run_hetero_case( local_block=24, kernel=8, @@ -960,7 +966,8 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): decode (local) logical block 5760 / kernel 64 (ppl=90), prefill (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with - replicated MLA and 8-way TP-sharded KDA state.""" + replicated MLA and 8-way TP-sharded KDA state. + """ _run_hetero_case( local_block=5760, kernel=64, @@ -974,7 +981,8 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): """The MLA per-token page is TP-independent, so kernel block lengths differing by anything other than the block-size ratio must fail the - handshake loudly rather than transfer at mismatched geometry.""" + handshake loudly rather than transfer at mismatched geometry. + """ worker = _make_mla_hybrid_worker( local_block_size=12, kernel_block_size=4, num_logical_blocks=8 ) @@ -1158,7 +1166,8 @@ def test_csa_linear_registration_discovers_shared_regions(): def test_csa_linear_scratch_descs_follow_the_pages_they_overlay(): """The scratch ring is addressed by the regions it actually registered in, not by inferring that it must sit on the MLA pages: overlaying the main KV - pages instead moves its descriptors with it.""" + pages instead moves its descriptors with it. + """ worker = _make_csa_linear_ple_worker(scratch_aliases="main_kv") # Same MLA regions as before, but the ring now lives in the non-MLA ones. @@ -1292,7 +1301,8 @@ def test_csa_linear_remote_ple_is_copied_whole(): def _make_ring_worker(): """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a - model whose sliding-window KV lives in per-request rings.""" + model whose sliding-window KV lives in per-request rings. + """ from unittest.mock import MagicMock from vllm.config import set_current_vllm_config diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 7ce84bf685b1..964cadc3ed76 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -115,7 +115,8 @@ def get_unhashed_block_ids_all_groups(self) -> tuple[list[int], ...]: def _stub_sw_clipping(scheduler) -> None: """Make ``get_exchange_clipped_blocks`` a passthrough so tests don't - need the full sliding-window machinery.""" + need the full sliding-window machinery. + """ scheduler.get_exchange_clipped_blocks = lambda block_ids, clip_ssm=True: block_ids @@ -173,7 +174,8 @@ def test_d_side_update_state_after_alloc_stages_registration(self): def test_p_side_request_finished_stages_blocks(self): """P scheduler pushes blocks into both _finished_request_blocks (lease) - and _newly_finished_push_blocks (metadata for next step).""" + and _newly_finished_push_blocks (metadata for next step). + """ sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -192,7 +194,8 @@ def test_p_side_request_finished_stages_blocks(self): def test_build_connector_meta_drains_both_sides(self): """meta.push_registrations and meta.push_finished_blocks are filled - from the staging dicts and the staging dicts are cleared.""" + from the staging dicts and the staging dicts are cleared. + """ sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -287,7 +290,8 @@ def test_update_connector_output_clears_lease_and_watchdog(self): def test_registration_watchdog_expires(self, caplog): """Stale D registrations whose deadline has passed are dropped at - ``build_connector_meta`` time.""" + ``build_connector_meta`` time. + """ # Watchdog logs a WARNING when it drops the stale entry; that's # what this test is verifying, so silence it in the test report. caplog.set_level( @@ -328,7 +332,8 @@ def test_registration_watchdog_expires(self, caplog): class _StubWriterWorker(NixlPushConnectorWorker): """Construct a worker without invoking ``__init__`` so we can drive - the matching/notif logic without bringing up NIXL or torch.""" + the matching/notif logic without bringing up NIXL or torch. + """ @classmethod def fresh(cls) -> _StubWriterWorker: @@ -434,7 +439,8 @@ class TestPushWriterRegSend: """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of the P engine, and the agent map is read under ``_handshake_lock`` so the send can't race the handshake done-callback's insert or a concurrent - stale-engine eviction (X1).""" + stale-engine eviction (X1). + """ def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]: w = _StubWriterWorker.fresh() @@ -479,7 +485,8 @@ def test_queues_failed_request_when_engine_not_handshaked(self): def test_agent_read_is_serialized_with_handshake_lock(self): """While another thread holds ``_handshake_lock`` (as the handshake done-callback does while writing ``_remote_agents``), the writer's - check-then-send must not run ahead of the critical section.""" + check-then-send must not run ahead of the critical section. + """ w, failures = self._reg_send_worker() w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"} rd = _registration_data("req-1") @@ -589,7 +596,8 @@ def test_handle_push_reg_drops_malformed(self, caplog): class TestPushWriterStartLoadKv: def test_finished_blocks_inbox_matches_stashed_registration(self): """Run the writer-loop's finished-blocks drain against a - pre-populated _pending_d_registrations entry.""" + pre-populated _pending_d_registrations entry. + """ w = _StubWriterWorker.fresh() w._pending_d_registrations["req-C"] = _registration_data("req-C") @@ -614,7 +622,8 @@ def test_finished_blocks_inbox_matches_stashed_registration(self): def test_start_load_kv_enqueues_to_writer(self): """``start_load_kv`` should hand registrations + finished blocks - to the writer queues without doing matching itself.""" + to the writer queues without doing matching itself. + """ w = _StubWriterWorker.fresh() # Stub heartbeats to a no-op; tests don't exercise the heartbeat # path here. @@ -672,7 +681,8 @@ def test_do_start_push_kv_defers_then_writes_when_handshake_ready(): NIXL op from the writer or the executor callback); once it resolves the request is re-queued on ``_deferred_push_inbox`` with the wake set; and on re-drive with the handshake ready the WRITE is issued with a correct - ReqMeta.""" + ReqMeta. + """ w = _StubWriterWorker.fresh() w._logical_to_kernel_block_ids = lambda x, ratio: x xfer_calls: list[dict[str, Any]] = [] @@ -712,7 +722,8 @@ def test_do_start_push_kv_defers_then_writes_when_handshake_ready(): def test_do_start_push_kv_drops_request_on_handshake_failure(): """Handshake raises: the request is dropped (not re-queued, no WRITE) and the failure is logged. Blocks are reclaimed by the lease/watchdog, matching - the old blocking behaviour.""" + the old blocking behaviour. + """ w = _StubWriterWorker.fresh() w._logical_to_kernel_block_ids = lambda x, ratio: x xfer_calls: list[dict[str, Any]] = [] @@ -734,7 +745,8 @@ def test_do_start_push_kv_drops_request_on_handshake_failure(): def test_writer_loop_drains_deferred_push_inbox(): """The writer loop drains ``_deferred_push_inbox`` and re-drives - ``_do_start_push_kv`` for each entry (event-driven, no polling).""" + ``_do_start_push_kv`` for each entry (event-driven, no polling). + """ w = _StubWriterWorker.fresh() w.nixl_wrapper = MagicMock() w.nixl_wrapper.get_new_notifs.return_value = {} @@ -767,7 +779,8 @@ def _tracked(rid, blocks, rd): def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker: """A push worker wired to drive the base ``_ensure_handshake`` eviction - path (``_evict_stale_engines`` + ``_cleanup_remote_engine``).""" + path (``_evict_stale_engines`` + ``_cleanup_remote_engine``). + """ w = _StubWriterWorker.fresh() w._engine_ttl = engine_ttl w._engine_last_active = {} @@ -789,7 +802,8 @@ def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker: def test_active_push_refreshes_engine_last_active(): """The base ``_ensure_handshake`` refreshes liveness only on a new handshake, so an active push to an already-connected engine must refresh - ``_engine_last_active`` itself or it is reaped mid-stream (S1).""" + ``_engine_last_active`` itself or it is reaped mid-stream (S1). + """ w = _eviction_worker(engine_ttl=3600.0) # Already-connected D -> _ensure_handshake returns None, WRITE runs inline. w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"} @@ -805,7 +819,8 @@ def test_active_push_refreshes_engine_last_active(): def test_stale_engine_evicted_on_push(): """A push drives ``_ensure_handshake`` -> ``_evict_stale_engines``, so a D engine silent past its TTL is reaped -- no _remote_agents / NIXL agent leak - across D scale up/down (S1).""" + across D scale up/down (S1). + """ w = _eviction_worker(engine_ttl=30.0) w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} w._engine_last_active["D-old"] = time.perf_counter() - 10_000.0 @@ -822,7 +837,8 @@ def test_stale_engine_evicted_on_push(): def test_cleanup_remote_engine_pops_under_handshake_lock(): """Regression (X1): eviction must not remove ``_remote_agents`` entries while another thread is inside the lock's critical section (e.g. the - push writer reading the map to send a registration).""" + push writer reading the map to send a registration). + """ w = _eviction_worker(engine_ttl=30.0) w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} @@ -847,7 +863,8 @@ def _cleanup(): class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained - on the engine main thread inside ``_get_new_notifs``.""" + on the engine main thread inside ``_get_new_notifs``. + """ w = _StubWriterWorker.fresh() # Pretend the writer thread already forwarded a completion notif # for a request whose KV is being received. @@ -893,7 +910,8 @@ def test_get_transfer_results_evicts_completed_state(self): @staticmethod def _pollable_worker() -> _StubWriterWorker: """Stub worker with the extra base-worker state the real - ``get_transfer_results`` needs to poll ``_sending_transfers``.""" + ``get_transfer_results`` needs to poll ``_sending_transfers``. + """ w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() w.nixl_wrapper = MagicMock() @@ -935,7 +953,8 @@ def test_failed_push_send_is_not_reported_done_sending(self): @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)]) def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma): """P-side sends are completed by consumer notifs / lease expiry, - not by local WRITE-handle completion.""" + not by local WRITE-handle completion. + """ w = self._pollable_worker() w.pp_size = pp_size w._is_hma_required = is_hma @@ -1006,7 +1025,8 @@ def test_update_state_after_alloc_zero_external_tokens_does_not_register(self): def test_request_finished_unfinished_status_does_not_stage(self): """If a request is still RUNNING, request_finished must not stash - blocks for the worker (no push needed).""" + blocks for the worker (no push needed). + """ sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1023,7 +1043,8 @@ def test_request_finished_unfinished_status_does_not_stage(self): def test_request_finished_empty_blocks_does_not_arm_lease(self): """Empty block-id groups should still complete cleanly without - arming the lease/finished maps.""" + arming the lease/finished maps. + """ sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1038,7 +1059,8 @@ def test_request_finished_empty_blocks_does_not_arm_lease(self): def test_update_connector_output_unknown_request_is_noop(self): """Idempotent cleanup: clearing a request that was never staged - must not raise or mutate other state.""" + must not raise or mutate other state. + """ sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1071,7 +1093,8 @@ def test_pop_matching_finished_blocks_returns_none_when_empty(self): def test_pop_matching_registration_no_match_when_base_ids_differ(self): """A registration whose base id (after stripping the random suffix) - does NOT match the lookup request_id must not be popped.""" + does NOT match the lookup request_id must not be popped. + """ w = _StubWriterWorker.fresh() # Two unrelated requests: different base UUIDs, so stripping the # trailing ``-<8 hex>`` suffix still yields different base ids. @@ -1087,7 +1110,8 @@ def test_pop_matching_registration_no_match_when_base_ids_differ(self): def test_handle_push_reg_with_non_dict_payload_is_dropped(self, caplog): """msgpack-encoded non-dict payload (e.g. a list) should be - dropped without raising.""" + dropped without raising. + """ caplog.set_level( logging.CRITICAL, logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), @@ -1115,7 +1139,8 @@ def test_handle_push_reg_with_non_string_request_id_is_dropped(self, caplog): def test_handle_push_reg_idempotent_for_same_request_id(self): """Receiving the same PUSH_REG twice (e.g. P retries after a - flake) keeps the entry staged exactly once and never fires.""" + flake) keeps the entry staged exactly once and never fires. + """ w = _StubWriterWorker.fresh() notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode( _registration_data("req-dup") @@ -1130,7 +1155,8 @@ def test_get_transfer_results_enqueues_each_completed_request(self): """``get_transfer_results`` must enqueue every completed request in ``done_sending`` so the writer can drop stale matching state. Unlike the happy-path test, this verifies the *cardinality*: N - completed requests -> N evictions, in order.""" + completed requests -> N evictions, in order. + """ w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], @@ -1153,7 +1179,8 @@ def test_get_transfer_results_enqueues_each_completed_request(self): def test_empty_transfer_results_do_not_enqueue_eviction(self): """If there's nothing newly done, no eviction should be enqueued. The wake event IS still set because ``get_finished`` always wakes - the writer to drain notifs.""" + the writer to drain notifs. + """ w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], @@ -1194,7 +1221,8 @@ def test_failed_push_transfer_does_not_evict_writer_state(self): def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): """A completion notif for a request the worker doesn't know - about should be logged but not crash.""" + about should be logged but not crash. + """ caplog.set_level( logging.CRITICAL, logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), @@ -1226,7 +1254,8 @@ def test_start_load_kv_with_empty_metadata_is_noop(self): def test_get_new_notifs_extends_lease_on_heartbeat(self): """``HB:`` notifs forwarded by the writer thread must extend the leases of tracked P-side requests on the engine main thread, and - ignore request IDs that aren't being tracked.""" + ignore request IDs that aren't being tracked. + """ w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() # _handle_heartbeat reads ``self._lease_extension`` (set in the @@ -1262,11 +1291,13 @@ def test_get_new_notifs_extends_lease_on_heartbeat(self): class TestPushPipelineParallel: """PP-sharded producer: per-stage completion counting, pp_size plumbing, - and remote-region slicing.""" + and remote-region slicing. + """ def test_completion_waits_for_one_notif_per_pp_stage(self): """Each PP stage WRITEs its own layers and sends one notif; D must - collect pp_size notifs before reporting the recv done.""" + collect pp_size notifs before reporting the recv done. + """ w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() request_id = "req-pp-2" @@ -1287,7 +1318,8 @@ def test_completion_waits_for_one_notif_per_pp_stage(self): def test_req_meta_reads_pp_size_from_kv_transfer_params(self): """D learns the producer's pp_size from kv_transfer_params (forwarded - by the proxy) and defaults to 1 when absent.""" + by the proxy) and defaults to 1 when absent. + """ metadata = NixlConnectorMetadata() params = { "remote_block_ids": ([0],), @@ -1310,7 +1342,8 @@ def test_add_remote_agent_slices_remote_regions_to_local_pp_window(self): this worker holds only a contiguous layer slice; add_remote_agent trims the remote region list to [offset : offset + num_local_regions] before building descriptors. We stop right after the slice via a - sentinel on the next collaborator call.""" + sentinel on the next collaborator call. + """ block_len = 4096 * 16 w = _StubWriterWorker.fresh() # seeds writer-thread state for teardown w.pp_size = 2 @@ -1352,7 +1385,8 @@ class TestPushWriterMlaReplication: """MLA latent KV is replicated across D's TP ranks, so when D_TP > P_TP (``tp_ratio < 0``) the producer must *WRITE* the latent into every D rank it handshook -- not write one and merely notify the rest, which would - leave the un-written ranks decoding against stale KV.""" + leave the un-written ranks decoding against stale KV. + """ @staticmethod def _mla_worker_writing_to(d_ranks): @@ -1577,7 +1611,8 @@ def _written_block_ids(w) -> tuple[list[int], list[int]]: def test_partial_prefix_hit_end_trims_producer_blocks(self): """D registered only its 2 uncomputed suffix blocks; P finished the - full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots.""" + full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots. + """ w, _ = self._worker_driving_xfer() reg = _registration_data("req-pc", local_block_ids=([500, 501],)) diff --git a/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py index da8be07157d8..bab9ed6e0fff 100644 --- a/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py +++ b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py @@ -171,7 +171,8 @@ def test_gpu_memory_nixl_hma(model_name, sw_size): @pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"]) def test_gpu_memory_no_nixl_baseline(model_name): """Same workload without NixlConnector. Comparing driver-level memory - between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact.""" + between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact. + """ from vllm import LLM, SamplingParams from vllm.distributed.parallel_state import cleanup_dist_env_and_memory diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index bdab91a21ffe..7044f9f0be52 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -213,7 +213,8 @@ def _multi_connector_config(swa_enabled: bool = False): def test_nixl_wins_load_over_cpu_offload(): """When NixlConnector (index 0) has matched tokens from a remote prefill, it should win the load: Nixl metadata tracks the recv while CPU offload metadata has no load - scheduled.""" + scheduled. + """ vllm_config, kv_cache_config = _multi_connector_config() scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) mc = scheduler.connector @@ -250,7 +251,8 @@ def test_nixl_wins_load_over_cpu_offload(): @patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper) def test_cpu_offload_wins_when_nixl_has_no_match(): """When NixlConnector returns 0 matched tokens and SimpleCPUOffloadConnector has a - CPU cache hit, the CPU offload connector (index 1) wins the load.""" + CPU cache hit, the CPU offload connector (index 1) wins the load. + """ vllm_config, kv_cache_config = _multi_connector_config() scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) mc = scheduler.connector @@ -309,7 +311,8 @@ def test_cpu_offload_wins_when_nixl_has_no_match(): def test_request_finished_no_async_save(swa_enabled: bool): """A normal request (no P/D) produces no async save from either connector. MultiConnector returns (False, None) via both request_finished and - request_finished_all_groups, and cleans up _requests_to_connector.""" + request_finished_all_groups, and cleans up _requests_to_connector. + """ from vllm.v1.request import RequestStatus vllm_config, kv_cache_config = _multi_connector_config(swa_enabled=swa_enabled) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index cde672bcadba..1d99ac0d1148 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -76,7 +76,7 @@ class MockSubscriber: - """Helper class to receive and verify published events""" + """Helper class to receive and verify published events.""" def __init__( self, @@ -121,7 +121,7 @@ def get_new_cpu_stored_events(self) -> list[BlockStored]: return cpu_stored_events def close(self): - """Clean up resources""" + """Clean up resources.""" self.sub.close() @@ -261,9 +261,7 @@ def test_cpu_offloading( cpu_block_size: int | None, uses_hma: bool, ) -> None: - """ - Tests OffloadingConnector with CPUOffloadingSpec. - """ + """Tests OffloadingConnector with CPUOffloadingSpec.""" # configure OffloadingConnector (spec_name=CPUOffloadingSpec by default) extra_config: dict = {"cpu_bytes_to_use": 500 << 20} if cpu_block_size is not None: @@ -323,7 +321,8 @@ def test_cpu_offloading( @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_cpu_offloading_metrics() -> None: """Verify that offloading Prometheus metrics (new flat and deprecated - labeled) are emitted after stores and loads.""" + labeled) are emitted after stores and loads. + """ extra_config: dict = { "cpu_bytes_to_use": 500 << 20, "block_size": CPU_BLOCK_SIZES, @@ -512,7 +511,8 @@ def test_tiering_offloading() -> None: def test_fs_tiering_offloading(tmp_path) -> None: """Tests OffloadingConnector with TieringOffloadingSpec - + fs secondary tier.""" + + fs secondary tier. + """ extra_config: dict = { "cpu_bytes_to_use": 1 << 30, "block_size": CPU_BLOCK_SIZES, diff --git a/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py index b656e0809543..8aa1e23d6fea 100644 --- a/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py @@ -20,7 +20,6 @@ def test_basic_lifecycle(): """Test lifecycle of a Remote Decode request.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -117,7 +116,6 @@ def test_basic_lifecycle(): def test_short_prompt_lifecycle(): """Test lifecycle of a Remote Decode request with short prompt.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -164,7 +162,6 @@ def test_short_prompt_lifecycle(): def test_prefix_cache_lifecycle(): """Test that remote decode params still work with a prefix cache hit.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -224,7 +221,6 @@ def test_prefix_cache_lifecycle(): def test_abort_during_kv_transfer(): """Test aborting request does not release blocks for remote decode.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 7d21d975c173..9ee93c889446 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -30,7 +30,6 @@ def _num_waiting_requests(scheduler) -> int: def test_basic_lifecycle(): """Test lifecycle of a remote prefill.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -151,7 +150,6 @@ def test_basic_lifecycle(): def test_interleaved_lifecycle(): """Test Remote Prefills Work Well With Other Requests.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -249,13 +247,11 @@ def test_interleaved_lifecycle(): def test_no_spurious_prefix_caching(): - """ - With P/D, blocks can be allocated but uncomputed for + """With P/D, blocks can be allocated but uncomputed for multiple engine steps. This test confirms that we do not accidentally have cache hits against uncomputed blocks. """ - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -320,7 +316,6 @@ def test_no_spurious_prefix_caching(): def test_full_block_prompt(): """Test that we handle a prompt that is the full block size.""" - vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) @@ -394,11 +389,9 @@ def test_full_block_prompt(): def test_cannot_schedule_after_recv(): - """ - Test that we can handle no schedule after recv due to not + """Test that we can handle no schedule after recv due to not enough remaining KV blocks. """ - # NOTE: the KVCacheManager will use 1 null block. # So there are 5 total working blocks. TOTAL_NUM_BLOCKS = 6 @@ -501,11 +494,9 @@ def test_cannot_schedule_after_recv(): def test_cannot_recv(): - """ - Test that we can handle no schedule KV block transfer due to not + """Test that we can handle no schedule KV block transfer due to not enough remaining KV blocks. """ - # NOTE: the KVCacheManager will use 1 null block. # So there are 5 total working blocks. TOTAL_NUM_BLOCKS = 6 @@ -713,7 +704,8 @@ def test_async_load_reserves_blocks_for_inflight(): def test_async_loads_both_admitted_when_pool_fits(): """Sanity: with a pool large enough, the reservation gate admits both async - loads (it is not over-conservative).""" + loads (it is not over-conservative). + """ vllm_config = create_vllm_config() BLOCK_SIZE = vllm_config.cache_config.block_size scheduler = create_scheduler(vllm_config, num_blocks=64) diff --git a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py index af4c7858fcad..f9f790232925 100644 --- a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py +++ b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py @@ -129,7 +129,8 @@ def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig: def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None: """The ring group's capacity is not a token granularity, it holds no - hashed blocks, and it never participates in CPU hits.""" + hashed blocks, and it never participates in CPU hits. + """ vllm_config = create_vllm_config( kv_connector="SimpleCPUOffloadConnector", kv_role="kv_both", diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 3ddc568a524d..201950705a73 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -235,7 +235,8 @@ def test_fa_and_ssm_different_split_factors(self): def test_hetero_block_size_splits(self): """With a block-size ratio, single-source FA sub-block descs pass - through whole; SSM descs are unexpanded and split per source.""" + through whole; SSM descs are unexpanded and split per source. + """ plan = TPMapping( source_ranks_per_group=((0,), (0, 1)), all_source_ranks=(0, 1), @@ -270,7 +271,8 @@ def test_hetero_block_size_splits(self): def test_hetero_block_size_head_sharded_asserts(self): """Head-sharded FA reads (multiple FA sources) are incompatible with - a block-size mismatch and must fail loudly.""" + a block-size mismatch and must fail loudly. + """ plan = TPMapping( source_ranks_per_group=((0, 1), (0, 1)), all_source_ranks=(0, 1), diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 62e13fa43834..80b97c6a1366 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -261,7 +261,6 @@ def create_model_runner_output( kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None, ) -> ModelRunnerOutput: """Make dummy model runner output for testing.""" - # Make request data. req_ids = [req.request_id for req in reqs] req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)} diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py index 338ec18d0e67..35897a9236ae 100644 --- a/tests/v1/kv_offload/cpu/policies/test_factory.py +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -15,7 +15,8 @@ class _DummyCachePolicy(CachePolicy): """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded by module path, so it must be importable at module scope (mirrors - tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec).""" + tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec). + """ def __init__(self, cache_capacity: int) -> None: self.cache_capacity = cache_capacity @@ -54,7 +55,8 @@ class TestCachePolicyFactory: def test_pre_registered_policies_can_be_imported(self): """If someone moves a policy module but forgets to update - factory.py, CI fails.""" + factory.py, CI fails. + """ for name in CachePolicyFactory._registry: cls = CachePolicyFactory._registry[name]() assert issubclass(cls, CachePolicy) @@ -88,7 +90,8 @@ def test_dynamic_load_via_cache_policy_module_path(self): register_cache_policy() call -- this is how external projects integrate a custom CachePolicy without forking/patching vLLM. Mirrors tests/v1/kv_offload/test_factory.py's - test_dynamic_load_via_spec_module_path.""" + test_dynamic_load_via_spec_module_path. + """ policy_cls = CachePolicyFactory.get_cache_policy_cls( "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory" ) @@ -96,7 +99,8 @@ def test_dynamic_load_via_cache_policy_module_path(self): def test_manager_resolves_policy_via_module_path(self): """End-to-end: CPUOffloadingManager resolves an unregistered policy - purely from cache_policy_module_path.""" + purely from cache_policy_module_path. + """ manager = CPUOffloadingManager( num_chunks=4, cache_policy="_DummyCachePolicy", @@ -106,6 +110,7 @@ def test_manager_resolves_policy_via_module_path(self): def test_unregistered_policy_without_module_path_raises(self): """eviction_policy not in registry + no cache_policy_module_path -> - ValueError, same shape as the OffloadingSpecFactory error path.""" + ValueError, same shape as the OffloadingSpecFactory error path. + """ with pytest.raises(ValueError, match="Unknown cache policy"): CachePolicyFactory.get_cache_policy_cls("nonexistent", None) diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py index 4594e1eb0924..d4e383801cfd 100644 --- a/tests/v1/kv_offload/cpu/test_canonical_layout.py +++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py @@ -126,7 +126,8 @@ def _canonical_handler(gpu_tensor, cpu_tensor, mapping, gpu_to_cpu): def test_gpu_roundtrip_assembles_canonical_page_across_ranks(): """Two TP2-style rank handlers must scatter into one shared canonical CPU page such that each rank reloads bit-exact and a whole-page (TP1) - reader sees both shards — the cross-topology contract.""" + reader sees both shards — the cross-topology contract. + """ torch.manual_seed(0) num_blocks = 4 gpu_rank = [ @@ -203,7 +204,8 @@ def _nhd_shard_mapping(tp: int, rank: int) -> CanonicalPageMapping: def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: """This rank's local page rows out of the (blocks, 2, tokens, heads, - head_bytes) ground truth.""" + head_bytes) ground truth. + """ local_heads = _TOTAL_HEADS // tp shard = full_kv[:, :, :, rank * local_heads : (rank + 1) * local_heads, :] return shard.reshape(full_kv.shape[0], -1).contiguous() @@ -214,7 +216,8 @@ def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int): """KV written at one tp must be readable at another: writer ranks scatter head shards into a shared canonical region, reader ranks gather their own - shards, and every reader must see the writers' ground-truth bytes.""" + shards, and every reader must see the writers' ground-truth bytes. + """ torch.manual_seed(0) num_blocks = 3 full_kv = torch.randint( diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 4fce65fb2a4f..0d78a18c1dbc 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -166,8 +166,7 @@ def test_cpu_eviction_removed_precedes_stored(): @pytest.mark.parametrize("eviction_policy", ["lru", "arc"]) def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy): - """ - Regression test: a chunk that is already stored must not be evicted + """Regression test: a chunk that is already stored must not be evicted by prepare_store() when it needs to make room for new chunks. Applies to both lru and arc policies. @@ -367,9 +366,7 @@ def test_cpu_manager_clears_write_usage_after_failed_store(): def test_cpu_manager(): - """ - Tests CPUOffloadingManager with lru policy. - """ + """Tests CPUOffloadingManager with lru policy.""" # initialize a CPU manager with a capacity of 4 chunks cpu_manager = make_cpu_manager(num_chunks=4, cache_policy="lru", enable_events=True) @@ -597,8 +594,7 @@ def _make_manager( return manager, policy def test_basic(self): - """ - Tests CPUOffloadingManager with arc policy. + """Tests CPUOffloadingManager with arc policy. Verifies that ARC handles store, load, and lookup operations correctly. """ cpu_manager, arc_policy = self._make_manager() @@ -637,8 +633,7 @@ def test_basic(self): assert len(arc_policy.t2) == 0 def test_t1_to_t2_promotion(self): - """ - Tests that accessing a chunk in T1 promotes it to T2 (frequent). + """Tests that accessing a chunk in T1 promotes it to T2 (frequent). This is a key feature of ARC's adaptive behavior. """ cpu_manager, arc_policy = self._make_manager(enable_events=False) @@ -659,8 +654,7 @@ def test_t1_to_t2_promotion(self): assert to_keys([1])[0] in arc_policy.t2 def test_eviction_with_load(self): - """ - Tests ARC eviction behavior similar to LRU test. + """Tests ARC eviction behavior similar to LRU test. Verifies that chunks being loaded (ref_cnt > 0) cannot be evicted. """ cpu_manager, _ = self._make_manager() @@ -700,8 +694,7 @@ def test_eviction_with_load(self): assert len(prepare_store_output.evicted_keys) >= 1 def test_adaptive_target(self): - """ - Tests ARC's adaptive target adjustment via ghost lists. + """Tests ARC's adaptive target adjustment via ghost lists. When a chunk in B1 (ghost list) is accessed, target_t1_size increases. When a chunk in B2 is accessed, target_t1_size decreases. """ @@ -728,8 +721,7 @@ def test_adaptive_target(self): assert arc_policy.target_t1_size > initial_target def test_t1_t2_eviction_policy(self): - """ - Tests that ARC evicts from T1 or T2 based on target_t1_size. + """Tests that ARC evicts from T1 or T2 based on target_t1_size. If |T1| >= target_t1_size, evict from T1, otherwise from T2. """ cpu_manager, arc_policy = self._make_manager(enable_events=False) @@ -917,8 +909,7 @@ def test_batch_eviction_failure_is_atomic(self): assert not arc_policy.b2 def test_ghost_list_bounds(self): - """ - Tests that ghost lists (B1, B2) don't grow unbounded. + """Tests that ghost lists (B1, B2) don't grow unbounded. They should be capped at cache_capacity. """ cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False) @@ -937,8 +928,7 @@ def test_ghost_list_bounds(self): assert len(arc_policy.b2) <= arc_policy.cache_capacity def test_touch_ordering(self): - """ - Tests that touch() correctly updates access patterns. + """Tests that touch() correctly updates access patterns. Similar to LRU test but verifies T1/T2 ordering. """ cpu_manager, arc_policy = self._make_manager() @@ -970,8 +960,7 @@ def test_touch_ordering(self): ) def test_failed_store(self): - """ - Tests that failed store operations clean up correctly. + """Tests that failed store operations clean up correctly. Similar to LRU test but for ARC. """ cpu_manager, arc_policy = self._make_manager() @@ -999,8 +988,7 @@ def test_failed_store(self): assert evicted_hash in arc_policy.b1 def test_full_scenario(self): - """ - Comprehensive test covering multiple ARC operations in sequence. + """Comprehensive test covering multiple ARC operations in sequence. Similar to the full LRU test but adapted for ARC behavior. """ cpu_manager, arc_policy = self._make_manager() @@ -1039,9 +1027,7 @@ def test_full_scenario(self): def test_filter_reused_manager(): - """ - Tests CPUOffloadingManager reuse filtering (store_threshold=2). - """ + """Tests CPUOffloadingManager reuse filtering (store_threshold=2).""" manager = make_cpu_manager( num_chunks=4, cache_policy="lru", @@ -1148,8 +1134,7 @@ def test_store_threshold_does_not_hide_ready_reuse(): def test_evictable_cache_chunk_count(): - """ - Verifies _num_evictable_cache_chunks is maintained correctly through the + """Verifies _num_evictable_cache_chunks is maintained correctly through the full store/load lifecycle, eviction, failed stores, concurrent loads, reset_cache, and the early-exit fast path in prepare_store. """ diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 4f22d122dbec..1d6a78f9e7ef 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -167,7 +167,8 @@ def _mp_race_construct_and_write( ) -> None: """Race to construct a SharedOffloadRegion, write fill_value, then wait for the parent's cleanup signal before tearing down. The wait gives the - parent a window to read the raw mmap before the creator removes the file.""" + parent a window to read the raw mmap before the creator removes the file. + """ try: region = SharedOffloadRegion( engine_id=engine_id, @@ -196,7 +197,8 @@ def _mp_barrier_construct_and_hold( replicated: bool = False, ) -> None: """Construct with a real cross-process barrier, write, then hold the - mapping (no cleanup) until the parent SIGKILLs this process.""" + mapping (no cleanup) until the parent SIGKILLs this process. + """ try: populate_calls = 0 get_populate_write_fn = region_module._get_populate_write_fn @@ -250,7 +252,8 @@ def iid(): def test_create_next_worker_view_shape_and_stride(iid): """Returned tensor must have shape (num_chunks, tensor_page_size) and - stride (row_stride, 1) where row_stride = cpu_page_size * num_workers.""" + stride (row_stride, 1) where row_stride = cpu_page_size * num_workers. + """ with _region(iid, num_chunks=4, cpu_page_size=2 * PAGE_SIZE) as r: t = r.create_next_worker_view(PAGE_SIZE) assert t.shape == (4, PAGE_SIZE) @@ -290,7 +293,8 @@ def test_create_next_worker_view_row_stride_with_multiple_workers(iid): def test_create_next_worker_view_cursor_advances(iid): """Each create_next_worker_view call must advance _worker_offset by - tensor_page_size.""" + tensor_page_size. + """ with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r: assert r._worker_offset == 0 r.create_next_worker_view(PAGE_SIZE) @@ -347,7 +351,8 @@ def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid): def test_create_next_worker_view_write_visible_in_raw_mmap(iid): """Writes into a create_next_worker_view view must appear at the correct - raw mmap offset""" + raw mmap offset + """ with _region(iid, num_chunks=4) as r: t = r.create_next_worker_view(PAGE_SIZE) t[2, :] = 42 # write to chunk row 2 @@ -380,7 +385,8 @@ def test_create_next_worker_view_multi_tensor_layout(iid): def test_create_next_worker_view_multiprocess_slots(iid): """Each worker process calls create_next_worker_view and writes distinct data; - the parent verifies each slot lands at the correct interleaved offset.""" + the parent verifies each slot lands at the correct interleaved offset. + """ num_workers = 2 num_chunks = 4 @@ -665,7 +671,8 @@ def test_multi_worker_race_shared_memory_visible(iid): def test_multiprocess_race_construct_and_write(iid): """N processes race to construct the same SharedOffloadRegion, each writes - fill_value = rank+1 into their slot; parent verifies interleaved layout.""" + fill_value = rank+1 into their slot; parent verifies interleaved layout. + """ num_workers = 4 num_chunks = 3 @@ -791,7 +798,8 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): def test_cleanup_after_create_next_worker_view_releases_mmap(iid): """cleanup() must close the mmap even after create_next_worker_view was called. create_next_worker_view returns a view that shares storage with _base; both must be - released before mmap.close() can succeed.""" + released before mmap.close() can succeed. + """ r = _make_region(iid) mmap_obj = r.mmap_obj @@ -970,7 +978,8 @@ def test_ftruncate_failure_cleans_up_creator(monkeypatch): def test_backing_file_unlinked_after_barrier(iid): """The file must exist until the barrier releases (late joiners need the - name) and be gone right after, so no exit path can leak it.""" + name) and be gone right after, so no exit path can leak it. + """ path = f"/dev/shm/vllm_offload_{iid}.mmap" seen_at_barrier = [] @@ -992,7 +1001,8 @@ def test_backing_file_unlinked_after_barrier(iid): def test_barrier_failure_unlinks_creator_and_raises(iid): """A failed rendezvous must remove the creator's file and re-raise, not - leave a stub that wedges the next start in _wait_for_file_size.""" + leave a stub that wedges the next start in _wait_for_file_size. + """ path = f"/dev/shm/vllm_offload_{iid}.mmap" with pytest.raises(RuntimeError, match="peer died"): _make_region(iid, barrier=MagicMock(side_effect=RuntimeError("peer died"))) @@ -1002,7 +1012,8 @@ def test_barrier_failure_unlinks_creator_and_raises(iid): @pytest.mark.parametrize("replicated", [False, True]) def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): """With a real cross-process barrier every worker maps one shared inode, - the name is gone while they run, and SIGKILL leaks nothing.""" + the name is gone while they run, and SIGKILL leaks nothing. + """ num_workers = 2 path = f"/dev/shm/vllm_offload_{iid}.mmap" ctx = get_mp_context() @@ -1041,7 +1052,8 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): @pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"]) def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase): """A worker that dies before the rendezvous must still arrive at the - barrier, or its peers block in the collective until it times out.""" + barrier, or its peers block in the collective until it times out. + """ failure = MagicMock(side_effect=ValueError("Insufficient space")) memory_check = failure if failure_phase == "memory" else None if failure_phase == "shm": @@ -1072,7 +1084,8 @@ def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_p def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): """A creator that fails after sizing the file must drop it before arriving - at the barrier, so the next start does not land on a stale file.""" + at the barrier, so the next start does not land on a stale file. + """ path = f"/dev/shm/vllm_offload_{iid}.mmap" monkeypatch.setattr("mmap.mmap", MagicMock(side_effect=OSError("mmap"))) seen_at_barrier = [] @@ -1085,7 +1098,8 @@ def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): def test_barrier_release_failure_keeps_original_error(iid, monkeypatch): """When releasing the peers fails too, the setup error that explains the - failure must be the one that propagates.""" + failure must be the one that propagates. + """ import vllm.v1.kv_offload.cpu.shared_offload_region as region monkeypatch.setattr( diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 8cd1ae9b10a1..cfe6660675f6 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -76,8 +76,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: def test_get_file_name_full_structure(): - """ - Path must match: _r//_g/.bin + """Path must match: _r//_g/.bin. Concretely: - The segment immediately after base_path must end with `_r3` diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py index 630fcdef4407..9f9e330b863a 100644 --- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py +++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -P2PConnector proxy server for OffloadingConnector + TieringOffloadingSpec. +"""P2PConnector proxy server for OffloadingConnector + TieringOffloadingSpec. Unlike NixlConnector (which returns remote_host/remote_port in the prefill response), OffloadingConnector does not embed connector coordinates in its diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index d3bacc8f326f..e40f16ad7393 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -132,7 +132,7 @@ def test_write_blocks_increments_transfer_id(self): assert tid1 != tid2 def test_poll_empty_when_no_inflight(self): - """poll returns empty when nothing is inflight.""" + """Poll returns empty when nothing is inflight.""" transport = self._make_transport() result = transport.poll() assert result == PollResult(done=(), failed=()) @@ -213,7 +213,7 @@ def test_poll_peer_id_scopes_to_peer(self): assert tid2 not in transport._inflight def test_cancel_removes_inflight(self): - """cancel removes transfers and releases handles.""" + """Cancel removes transfers and releases handles.""" transport = self._make_transport() transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) @@ -226,7 +226,7 @@ def test_cancel_removes_inflight(self): transport._agent.release_xfer_handle.assert_called() def test_cancel_ignores_unknown_ids(self): - """cancel with unknown IDs doesn't crash.""" + """Cancel with unknown IDs doesn't crash.""" transport = self._make_transport() assert transport.cancel([999, 1000]) == [] assert transport.cancel([999, 1000], mode="wait") == [] @@ -286,7 +286,7 @@ def test_add_and_remove_remote_peer(self): transport._agent.remove_remote_agent.assert_called() def test_close_releases_everything(self): - """close releases all handles and clears state.""" + """Close releases all handles and clears state.""" transport = self._make_transport() transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) transport.write_blocks("peer:1", [0], [1]) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index aceec1cd8fa1..2e8e9ff4be69 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -48,7 +48,8 @@ def _remote_prefiller_kv_params( kv_request_id: str = "req-1", ) -> dict: """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries - kv_request_id + remote_host + remote_port.""" + kv_request_id + remote_host + remote_port. + """ return { "remote_prefiller": { "kv_request_id": kv_request_id, @@ -64,7 +65,8 @@ def _remote_kv_source_kv_params( kv_request_id: str = "req-1", ) -> dict: """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has - the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port). + """ return { "remote_kv_source": { "kv_request_id": kv_request_id, @@ -76,7 +78,8 @@ def _remote_kv_source_kv_params( def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries - kv_request_id only.""" + kv_request_id only. + """ return {"remote_decoder": {"kv_request_id": kv_request_id}} @@ -241,7 +244,8 @@ def test_lookup_returns_hit_for_different_request_id(self): def test_lookup_returns_miss_without_prefill_key(self): """No ``remote_prefiller`` sub-dict means the request was not routed for remote prefill — local prefill should run instead, so lookup() - returns MISS even when a stale ``remote_decoder`` block is present.""" + returns MISS even when a stale ``remote_decoder`` block is present. + """ mgr = _make_manager() ctx = _req_context(kv_params=_remote_decoder_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.MISS @@ -309,7 +313,8 @@ class TestServeExternalRequests: def test_flushes_failed_serve_ctxs_then_serves_each_session(self): """serve_external_requests releases the failed serves left by reaped sessions via parent.on_request_finished (clearing the - queue), then delegates to every live session with the same parent.""" + queue), then delegates to every live session with the same parent. + """ mgr = _make_manager() ctx = ReqContext(req_id="p2p:peer:req-1:lu1") mgr._failed_serve_ctxs = [ctx] @@ -363,7 +368,8 @@ def test_missing_kv_request_id_fails(self): def test_no_binding_yet_parks_in_unbound_stores(self): """submit_store without a bound session buffers the batch keyed by kv_request_id; no session is pre-created (peer_id is unknown - to the producer at store time).""" + to the producer at store time). + """ mgr = _make_manager() job = _job_metadata( job_id=1, @@ -390,7 +396,8 @@ def test_no_binding_yet_parks_in_unbound_stores(self): def test_routes_to_bound_session(self): """If a session has already received FetchMsg for this kv_request_id (so _kv_to_session is populated), submit_store - forwards directly to that session rather than re-buffering.""" + forwards directly to that session rather than re-buffering. + """ mgr = _make_manager() bound = _FakeSession(peer_id="10.0.0.1:8000", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] @@ -413,7 +420,8 @@ def test_routes_to_bound_session(self): def test_extra_top_level_keys_are_ignored(self): """Producer-side kv_transfer_params should not pre-create a session even when a stale caller still passes a top-level - ``remote_host``/``remote_port`` next to ``remote_decoder``.""" + ``remote_host``/``remote_port`` next to ``remote_decoder``. + """ mgr = _make_manager() params = _remote_decoder_kv_params() params["remote_host"] = "stale" @@ -457,7 +465,8 @@ def test_no_session_fails(self): def test_happy_path_with_active_session(self): """When the peer's session exists, submit_load forwards to - session.request_blocks and does NOT add a finished result yet.""" + session.request_blocks and does NOT add a finished result yet. + """ mgr = _make_manager() peer_id = "10.0.0.1:8000" existing = _FakeSession(peer_id=peer_id, connected=True) @@ -477,7 +486,8 @@ def test_happy_path_with_active_session(self): def test_missing_consumer_flag_fails(self): """Peer fields present but neither do_remote_prefill nor do_p2p_fetch is set — submit_load fails the job rather than - emit a stray FetchMsg.""" + emit a stray FetchMsg. + """ mgr = _make_manager() params = { "remote_host": "10.0.0.1", @@ -521,7 +531,8 @@ def test_no_kv_request_id_does_nothing(self): def test_decoder_side_calls_session_finish_request(self): """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id because the consumer addresses the producer it loaded from. The - session's finish_request cancels the client-role load.""" + session's finish_request cancels the client-role load. + """ mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) @@ -533,7 +544,8 @@ def test_decoder_side_calls_session_finish_request(self): def test_p2p_consumer_side_calls_session_finish_request(self): """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id so the session drops any pending lookups (cancel_lookups) and - cancels any inbound load.""" + cancels any inbound load. + """ mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) @@ -544,7 +556,8 @@ def test_p2p_consumer_side_calls_session_finish_request(self): def test_prefiller_bound_id_routes_via_kv_to_session(self): """Prefiller-side finish for an id whose session is already bound - (FetchMsg received) routes via _kv_to_session and pops the entry.""" + (FetchMsg received) routes via _kv_to_session and pops the entry. + """ mgr = _make_manager() bound = _FakeSession(peer_id="some-peer:1", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] @@ -558,7 +571,8 @@ def test_prefiller_unbound_id_leaves_batches_parked(self): and no session binding is a no-op on `_unbound_stores`. The parked batches survive until a peer fetches them or the `_reap_unbound_stores` timeout fires — `on_request_finished` - must not evict them.""" + must not evict them. + """ from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = _make_manager() @@ -728,7 +742,8 @@ def remove_remote_peer(self, pid): def test_reap_fails_probes(self): """A reaped session's in-flight lookups land in _failed_req_ids so - the consumer's lookup() returns MISS instead of RETRY forever.""" + the consumer's lookup() returns MISS instead of RETRY forever. + """ class FakeData: def remove_remote_peer(self, pid): @@ -763,7 +778,8 @@ def test_unbound_store_kept_within_timeout(self): def test_unbound_store_reaped_after_timeout(self): """Unbound stores past the reap deadline surface as failed and their kv_request_id lands in _reaped_stores so a late FetchMsg is rejected - instead of parking demand nothing can satisfy.""" + instead of parking demand nothing can satisfy. + """ from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = self._make() @@ -785,7 +801,8 @@ def test_unbound_store_reaped_after_timeout(self): def test_submit_store_parks_unbound_batch(self): """submit_store on an unbound id appends a batch with a fresh - submitted_at stamp so the unbound-store sweep can age it out.""" + submitted_at stamp so the unbound-store sweep can age it out. + """ mgr = _make_manager() job = _job_metadata( job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") @@ -807,18 +824,21 @@ class TestHasPendingWork: """has_pending_work() must always return True so the engine keeps ticking the offload pipeline — that's the only thread driving _control.poll() (incoming peer connects) and session.poll() - (incoming fetch messages on existing sessions).""" + (incoming fetch messages on existing sessions). + """ def test_returns_true_unconditionally_to_keep_engine_ticking(self): """Even with no sessions and no jobs, has_pending_work() returns True so the engine keeps calling get_finished_jobs(), which is - what drives _control.poll() for inbound peer connects.""" + what drives _control.poll() for inbound peer connects. + """ mgr = _make_manager() assert mgr.has_pending_work() is True def test_returns_true_even_when_sessions_present(self): """The result is the same regardless of session state — there is - no 'idle' branch.""" + no 'idle' branch. + """ mgr = _make_manager() mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") # type: ignore[assignment] assert mgr.has_pending_work() is True @@ -831,7 +851,8 @@ def test_returns_true_even_when_sessions_present(self): class _ShutdownFakeData: """Fake DataTransport that records cancel/poll/close calls and - drives the wait-cancel loop with a scriptable `still` queue.""" + drives the wait-cancel loop with a scriptable `still` queue. + """ def __init__(self, still_queue: list[list[int]] | None = None) -> None: # Each list in still_queue is the set of ids the next @@ -875,7 +896,8 @@ def close(self) -> None: class TestShutdownDrain: """shutdown() drains inflight transfers via cancel(mode='wait') before calling _data.close(), with a 3s deadline fallback to - cancel(mode='immediate').""" + cancel(mode='immediate'). + """ def _prep( self, @@ -1294,7 +1316,8 @@ def close(self) -> None: class TestAcceptNewPeers: """A second inbound from an already-connected peer is rejected and the - new conn is closed; the existing session is left untouched.""" + new conn is closed; the existing session is left untouched. + """ def test_duplicate_connection_is_closed_and_existing_session_untouched(self): mgr = _make_manager() @@ -1316,7 +1339,8 @@ def test_creates_session_for_new_peer(self): """An inbound conn from a peer with no existing session creates a fresh connected session and registers it under conn.peer_id. The prefiller has no pre-created pending session anymore — the - first signal of a peer's existence is its inbound connection.""" + first signal of a peer's existence is its inbound connection. + """ class FakeData: block_len = 4096 @@ -1363,7 +1387,8 @@ def close(self) -> None: class TestPollOnce: """_poll_once must drain control, accept new peers, poll every session, - surface results, and reap dead sessions — in that order.""" + surface results, and reap dead sessions — in that order. + """ def test_orchestrates_accept_poll_and_reap(self): mgr = _make_manager() @@ -1419,7 +1444,8 @@ def remove_remote_peer(self_inner, pid): def test_new_fetch_id_binds_and_replays_unbound_batches(self): """When session.poll() reports a kv_request_id whose FetchMsg arrived this tick, the manager binds it to that session and - replays every parked submit_store batch via add_stored_blocks.""" + replays every parked submit_store batch via add_stored_blocks. + """ from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = _make_manager() @@ -1453,7 +1479,8 @@ def poll(self_inner): def test_new_fetch_id_with_no_unbound_still_binds(self): """A FetchMsg for a kv_request_id with no parked batches still - records the binding so subsequent submit_stores route fast.""" + records the binding so subsequent submit_stores route fast. + """ mgr = _make_manager() peer = "10.0.0.1:8000" sess = _FakeSession( @@ -1476,7 +1503,8 @@ def poll(self_inner): def test_failed_load_records_kv_request_id(self): """A LoadResult(success=False) from session.poll() must add its - kv_request_id to _failed_req_ids so future lookups return MISS.""" + kv_request_id to _failed_req_ids so future lookups return MISS. + """ mgr = _make_manager() peer = "10.0.0.1:8000" sess = _FakeSession( @@ -1540,7 +1568,8 @@ def test_returns_when_session_has_no_inflight_or_inbound(self): def test_logs_warning_after_5s_then_completes(self, monkeypatch): """A session that stays inflight past 5s triggers the warning, and - once it clears the loop returns.""" + once it clears the loop returns. + """ mgr = _make_manager() mgr._control = _DrainCtrl() # type: ignore[assignment] sess = _FakeSession(peer_id="peer:1") @@ -1598,7 +1627,8 @@ def clearing_sleep(_s): class TestOnScheduleEnd: def test_is_noop(self): """on_schedule_end is a documented no-op; just confirm it doesn't - raise and doesn't mutate state.""" + raise and doesn't mutate state. + """ mgr = _make_manager() before_sessions = dict(mgr._sessions) before_jobs = list(mgr._finished_jobs) @@ -1624,7 +1654,8 @@ class TestConnectionDeathMidTransfer: prefiller-side store no longer travels through the session at store time (it's parked in _unbound_stores keyed by kv_request_id), so its cleanup on connection death is via on_request_finished or the - unbound-store timeout — covered separately below.""" + unbound-store timeout — covered separately below. + """ def test_dead_connection_with_pending_work_surfaces_failures(self): mgr_a, mgr_b = _build_paired_managers() diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 2e4415fbb75d..ff5641134ece 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -626,7 +626,8 @@ def test_active_loads_work_list_tracks_in_flight(self): """collect_results / has_active_loads use the _active_loads work-list, armed when a fetch is issued and discarded exactly when its load clears — a probe-only request never enters it, and completion empties - it while the entry may briefly linger for GC.""" + it while the entry may briefly linger for GC. + """ session, conn, _ = _make_session() _activate(session, conn) client = session._client @@ -703,7 +704,8 @@ def test_load_abort_ack_timeout_surfaces_failure(self): def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from - _requests — covers the on_abort_ack arrival path.""" + _requests — covers the on_abort_ack arrival path. + """ session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( @@ -740,7 +742,8 @@ class TestLookupFlow: def test_aggregate_flush_resolve_round_trip(self): """register_lookup → flush sends one LookupMsg → response resolves entries → register_lookup returns the cached bool on - every call, and repeat probes never re-issue a LookupMsg.""" + every call, and repeat probes never re-issue a LookupMsg. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -793,7 +796,8 @@ def test_aggregate_flush_resolve_round_trip(self): def test_request_blocks_clears_probe_cache(self): """A resolved HIT probe is popped when its fetch is issued, so a re-scheduled request re-probes instead of trusting the stale True - (the served block is unpinned and may have been evicted).""" + (the served block is unpinned and may have been evicted). + """ session, conn, _ = _make_session() _activate(session, conn) @@ -827,7 +831,7 @@ def test_request_blocks_clears_probe_cache(self): assert fresh[0][LookupMsg.KEYS] == [b"hA"] def test_flush_uses_work_list_not_full_scan(self): - """flush drains a work-list rather than scanning every live request. + """Flush drains a work-list rather than scanning every live request. A request with no newly-registered keys is not revisited: after a flush the work-list is empty, an idle re-flush sends nothing, and a @@ -879,7 +883,8 @@ def test_separate_lookup_msg_per_kv_request_id(self): def test_multiple_lookup_msgs_across_steps(self): """A request's block set may be discovered across scheduler steps: each step that registers new keys flushes its own LookupMsg for - the same kv_request_id, carrying only the newly-probed keys.""" + the same kv_request_id, carrying only the newly-probed keys. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -907,7 +912,8 @@ def test_multiple_lookup_msgs_across_steps(self): def test_split_response_resolves_across_messages(self): """Producer may answer one LookupMsg's keys across multiple - LookupRespMsgs — pairs are self-describing so each lands.""" + LookupRespMsgs — pairs are self-describing so each lands. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -954,7 +960,8 @@ def test_finish_request_cancels_pending_lookups(self): def test_finish_after_flushed_lookup_sends_empty_fetch(self): """LookupMsg flushed but no FetchMsg sent (all-miss case) → finish_request emits an empty FetchMsg so the peer can drop its - lookup state and call parent.on_request_finished.""" + lookup state and call parent.on_request_finished. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -972,7 +979,8 @@ def test_finish_after_flushed_lookup_sends_empty_fetch(self): def test_finish_without_flushed_lookup_sends_no_fetch(self): """No LookupMsg was ever sent → finish_request must not emit an - empty FetchMsg (the peer has no state to release).""" + empty FetchMsg (the peer has no state to release). + """ session, conn, _ = _make_session() _activate(session, conn) @@ -987,7 +995,8 @@ def test_finish_without_flushed_lookup_sends_no_fetch(self): def test_finish_after_real_fetch_sends_no_second_fetch(self): """A real FetchMsg was already sent for the id → finish_request - must not emit a second (empty) FetchMsg.""" + must not emit a second (empty) FetchMsg. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -1018,7 +1027,8 @@ def test_server_lookup_deferred_until_serve_then_all_misses(self): """``poll()`` only enqueues an inbound LookupMsg — no response is sent until ``serve_external_requests``. With an all-miss parent the aggregated LookupRespMsg carries the same keys and - ``hits=[False, ...]``.""" + ``hits=[False, ...]``. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -1078,7 +1088,8 @@ class TestServerLookupHandling: def test_immediate_hits_create_one_store_job(self): """All-HIT batch: one create_store_job call with all keys, one LookupRespMsg with hits=[True]*N, on_request_finished fires at the - end of serve, and `available` is populated for the eventual fetch.""" + end of serve, and `available` is populated for the eventual fetch. + """ cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) session, conn, _ = _make_session() _activate(session, conn) @@ -1108,7 +1119,8 @@ def test_immediate_hits_create_one_store_job(self): def test_all_misses_no_store_job_finish_fires(self): """All-MISS batch: no create_store_job call; one LookupRespMsg - with hits=[False]*N; on_request_finished fires at end of serve.""" + with hits=[False]*N; on_request_finished fires at end of serve. + """ cb = FakeParent() session, conn, _ = _make_session() _activate(session, conn) @@ -1128,7 +1140,8 @@ def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): """HIT/MISS resolutions do not go out on first sight when any key is still HIT_PENDING / RETRY. The lookup parks until every key has settled (or the deadline fires), then one - LookupRespMsg carries all keys in wire order.""" + LookupRespMsg carries all keys in wire order. + """ cb = FakeParent( stored={b"hA": 1}, pending={b"hB"}, @@ -1157,7 +1170,8 @@ def test_pending_resolves_then_aggregate_response_fires(self): the deferred aggregate response: one LookupRespMsg carrying both keys in wire order, and one create_store_job call per HIT (the second HIT is pinned when it resolves, not when the - response goes out).""" + response goes out). + """ cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1193,7 +1207,8 @@ def test_pending_resolves_then_aggregate_response_fires(self): def test_pending_timeout_replies_miss_no_store_job(self): """A HIT_PENDING key that stays pending past the batch ``deadline`` is force-MISS and never pinned; the deferred - aggregate response fires with hits=[False].""" + aggregate response fires with hits=[False]. + """ cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1220,7 +1235,8 @@ def test_pending_timeout_replies_miss_no_store_job(self): def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): """Two LookupMsgs for the same kv_request_id get distinct ctxs - and two on_request_finished calls (one per batch).""" + and two on_request_finished calls (one per batch). + """ cb = FakeParent(stored={b"hA": 1, b"hB": 2}) session, conn, _ = _make_session() _activate(session, conn) @@ -1244,7 +1260,8 @@ def test_close_returns_open_batch_ctxs_as_failed_serves(self): """Tearing the session down with a parked batch returns the synthetic ctx as a failed serve (no parent handle at teardown) so the manager can release the TieringManager's state on its next - serve.""" + serve. + """ cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1265,7 +1282,8 @@ def test_close_returns_open_batch_ctxs_as_failed_serves(self): def test_wire_finish_drops_pending_batches_for_kv_request_id(self): """``ServerRole.finish(kv_request_id)`` drops every parked batch whose kv_request_id matches and queues its ctx for the next - serve's on_request_finished.""" + serve's on_request_finished. + """ cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1296,7 +1314,8 @@ def test_wire_finish_drops_pending_batches_for_kv_request_id(self): def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): """A peer FetchMsg terminates the lookup phase for its id: parked lookups with matching kv_request_id are dropped and their - ctx queued for on_request_finished; other kv_request_ids untouched.""" + ctx queued for on_request_finished; other kv_request_ids untouched. + """ cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1338,7 +1357,8 @@ def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): def test_lookup_then_fetch_round_trip_emits_store_result(self): """End-to-end: lookup pins primary slots → fetch matches them → NIXL transfer completes → StoreResult surfaces with the - create_store_job's job_id (the engine releases the pin).""" + create_store_job's job_id (the engine releases the pin). + """ cb = FakeParent(stored={b"hA": 7, b"hB": 8}) session, conn, transport = _make_session() _activate(session, conn) @@ -1460,7 +1480,8 @@ def test_abort_fetch_replies_with_ack(self): def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is - deferred and the abort is parked (abort_started_at set).""" + deferred and the abort is parked (abort_started_at set). + """ session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends @@ -1495,7 +1516,8 @@ def test_abort_fetch_defers_ack_when_cancel_pending(self): def test_abort_fetch_acks_after_drain(self): """Once the transport reports the tid as DONE the parked abort - completes and AbortAckMsg is sent.""" + completes and AbortAckMsg is sent. + """ session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1534,7 +1556,8 @@ def test_abort_fetch_acks_after_drain(self): def test_abort_fetch_force_cancels_after_timeout(self): """If wait-mode never drains, the deadline forces immediate - cancel and an ack is still sent.""" + cancel and an ack is still sent. + """ session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1576,7 +1599,8 @@ def test_abort_fetch_force_cancels_after_timeout(self): def test_abort_fetch_idempotent_while_draining(self): """Receiving AbortFetchMsg twice for the same kv_request_id - keeps a single pending entry and produces a single ack.""" + keeps a single pending entry and produces a single ack. + """ session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1634,7 +1658,8 @@ def test_store_timeout(self): def test_store_timeout_then_late_completion_no_duplicate(self): """A job timed out by _timeout_pending_store_jobs must not also emit a contradictory StoreResult(success=True) when the transport - later reports the same transfer as done.""" + later reports the same transfer as done. + """ session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1667,7 +1692,8 @@ def test_store_timeout_then_late_completion_no_duplicate(self): def test_store_timeout_then_late_failure_no_duplicate(self): """Symmetric guard: a timed-out job must not also emit a second StoreResult(success=False) when the transport later reports the - same transfer as failed.""" + same transfer as failed. + """ session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1710,7 +1736,8 @@ def _last_transfer_done(self, conn: FakeConnection) -> dict | None: def test_no_inflight_unmatched_demand_sends_failure(self): """finish_request with unmatched demand and no inflight -> - immediate TransferDoneMsg(success=False); _outbound cleared.""" + immediate TransferDoneMsg(success=False); _outbound cleared. + """ session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded a block we never stored. @@ -1736,7 +1763,8 @@ def test_no_inflight_unmatched_demand_sends_failure(self): def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the - early-fail message and clears _outbound.""" + early-fail message and clears _outbound. + """ session, conn, transport = _make_session() _activate(session, conn) # Demand 2 blocks; we store 1 (kicks one inflight transfer). @@ -1773,7 +1801,8 @@ def test_with_inflight_defers_then_fires_on_last_transfer(self): def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: - when remaining hits 0, success=True still fires.""" + when remaining hits 0, success=True still fires. + """ session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( @@ -1802,7 +1831,8 @@ def test_prefiller_first_finish_before_fetch(self): """Prefiller-first: finish_request runs before the decoder's fetch arrives. State is held until fetch, then finalized — success=True if all demand was matched against - available blocks, else success=False.""" + available blocks, else success=False. + """ # Case A: all demand satisfied by available blocks. session, conn, transport = _make_session() _activate(session, conn) @@ -1863,7 +1893,8 @@ def test_finish_request_no_inflight_emits_store_failure(self): """finish_request with a stored-but-unmatched job and no inflight -> TransferDoneMsg(success=False) AND deferred StoreResult(success=False) for the submit_store'd job, instead of - the 30s _STORE_TIMEOUT_S path.""" + the 30s _STORE_TIMEOUT_S path. + """ session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded b"demand"; we never stored it. @@ -1899,7 +1930,8 @@ def test_finish_request_no_inflight_emits_store_failure(self): def test_finish_request_remaining_zero_emits_success_via_inflight(self): """Deferred-via-inflight path: finish_request with inflight, last transfer drains remaining to 0 -> TransferDoneMsg(success=True) - AND StoreResult(success=True).""" + AND StoreResult(success=True). + """ session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( @@ -2247,7 +2279,8 @@ def test_close_returns_pending_loads_and_stores(self): def test_send_failure_marks_connection_dead(self): """A raising send must mark the connection dead, not silently drop the message — otherwise the session lingers alive, is never reaped, - and in-flight lookups/loads toward the dead peer hang forever.""" + and in-flight lookups/loads toward the dead peer hang forever. + """ session, conn, _ = _make_session() _activate(session, conn) assert session.alive @@ -2260,7 +2293,8 @@ def test_send_failure_marks_connection_dead(self): def test_close_surfaces_inflight_lookups(self): """close() reports kv_request_ids whose symmetric-P2P probe is still - unresolved; resolved probes are not reported (their answer is in).""" + unresolved; resolved probes are not reported (their answer is in). + """ session, conn, _ = _make_session() _activate(session, conn) @@ -2368,7 +2402,8 @@ class TestDispatchErrorHandling: def test_value_error_disconnects_on_first_occurrence(self): """A FetchMsg that fails validate() raises ValueError and must - terminate the session immediately, with a DisconnectMsg sent.""" + terminate the session immediately, with a DisconnectMsg sent. + """ session, conn, _ = _make_session() _activate(session, conn) # length mismatch → FetchMsg.validate raises ValueError @@ -2395,7 +2430,8 @@ def test_transfer_done_missing_field_disconnects(self): def test_internal_error_does_not_disconnect_once(self): """A non-ValueError raised by a handler is treated as an internal - bug: counter increments, session stays alive on a single hit.""" + bug: counter increments, session stays alive on a single hit. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -2418,7 +2454,8 @@ def _boom(*args, **kwargs): def test_internal_error_threshold_disconnects(self): """Once consecutive non-protocol errors hit the threshold, the - session tears down via _protocol_error.""" + session tears down via _protocol_error. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -2442,7 +2479,8 @@ def _boom(*args, **kwargs): def test_internal_error_counter_resets_on_success(self): """A successful dispatch between errors prevents the threshold - from being reached.""" + from being reached. + """ session, conn, _ = _make_session() _activate(session, conn) @@ -2485,7 +2523,8 @@ class TestInflightPerReqInvariant: def test_invariant_holds_through_lifecycle(self): """Run a full submit→complete sequence for two concurrent kv_request_ids and assert the counter matches `_inflight` at - every observable step, including the empty-after-finish case.""" + every observable step, including the empty-after-finish case. + """ session, conn, transport = _make_session() _activate(session, conn) @@ -2550,7 +2589,8 @@ def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must match the actual presence in `_inflight` for both hits and misses. The whole point of the counter is that this lookup is - constant-time, but we assert correctness, not timing.""" + constant-time, but we assert correctness, not timing. + """ session, _, _ = _make_session() for kv_id_idx in range(100): kv_id = f"req-{kv_id_idx}" diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index 3649c8df0488..616c6efa9479 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -264,7 +264,8 @@ def test_mark_miss_flips_cached_verdict_without_reprobing(self): subsequent lookup then returns False (MISS) directly, WITHOUT enqueuing a fresh batch_lookup — that is what makes the request unable to loop. The entry is retained (as False) not dropped, so cleanup()'s reverse - index stays consistent; an unknown key is a no-op.""" + index stays consistent; an unknown key is a no-op. + """ mgr = InMemoryLookupManager(existing_keys={_key(1), _key(2)}) ctx = _ctx("reqA") assert mgr.lookup(_key(1), ctx) is None @@ -304,7 +305,8 @@ def test_enqueue_once_invariant_enforced(self): receive at most one result per key. Normal operation resolves a key a single time; a second result for an already-decided key trips the assert that guards the invariant (a silent overwrite could flip a corrected - miss back to True and reopen the failed-load livelock).""" + miss back to True and reopen the failed-load livelock). + """ mgr = InMemoryLookupManager(existing_keys={_key(1)}) ctx = _ctx("reqA") diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index c391ac201631..f91f54bb0d16 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for SecondaryTierFactory. +"""Unit tests for SecondaryTierFactory. These tests verify: 1. Pre-registration integrity — registered tier module paths can import diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 4b495e1b315c..13621fb4c432 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for FileSystemTierManager. +"""Unit tests for FileSystemTierManager. These tests use real disk I/O to verify the filesystem tier implementation. The tier manager writes KV cache blocks to disk and reads them back, verifying @@ -370,7 +369,8 @@ def test_shutdown_discards_pending_tasks(fs_tier): @pytest.mark.parametrize("use_c_ext", [True, False]) def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): """Data written by store must be exactly recovered by load, for batches - of any size -- including the empty batch.""" + of any size -- including the empty batch. + """ import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -467,7 +467,8 @@ def test_wait_idle_blocks_until_tasks_complete(): def test_batch_lookup_c_extension(tmp_path): """Validates batch_lookup_C: empty, single, all-existing, all-missing, - mixed ordering, and input type validation.""" + mixed ordering, and input type validation. + """ try: from vllm.fs_io_C import batch_lookup as batch_lookup_C except ImportError: @@ -531,7 +532,8 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): @pytest.mark.parametrize("use_c_ext", [True, False]) def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): """Smoke test: a block id beyond the primary tensor's block count must - fail the job, for both the C extension and the Python fallback.""" + fail the job, for both the C extension and the Python fallback. + """ import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -613,7 +615,8 @@ def test_batched_partial_load_failure_keeps_loaded_blocks( loaded before it (#50321). Corrupt the LAST block: the earlier blocks load fine, so the job reports successful_keys for them and marks only the failed tail a miss. The earlier keys stay HIT — including for the same request — - while the corrupt block stays a MISS (its file was removed).""" + while the corrupt block stays a MISS (its file was removed). + """ import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -665,7 +668,8 @@ def test_batched_load_first_block_fails_marks_whole_batch( ): """When the FIRST block fails, nothing loaded before it: the job reports no successful_keys (None) and the whole batch is marked a miss for the - request.""" + request. + """ import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -696,7 +700,8 @@ def test_transient_load_failure_leaves_file(fs_tier, monkeypatch, use_c_ext): """A transient host error (here ELOOP on open) is NOT a short read: the job fails but the block file must survive untouched, on both the C and Python paths. Deleting on a transient error would turn a passing hiccup into - permanent data loss.""" + permanent data loss. + """ import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -800,7 +805,8 @@ def test_mixed_job_results_emit_event_only_for_successful_job( fs_tier_with_events, monkeypatch ): """With a failed and a successful store job in flight, exactly one event - is emitted and its keys belong to the successful job.""" + is emitted and its keys belong to the successful job. + """ import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod tier = fs_tier_with_events @@ -885,7 +891,8 @@ def test_events_require_global_kv_events_flag(tmp_path): def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): """A GPU->CPU->fs cascade surfaces the tier-owned FS stored event via the - TieringOffloadingManager's aggregated take_events().""" + TieringOffloadingManager's aggregated take_events(). + """ from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 13840bac3599..0b3168ec3d32 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Mock-based unit tests for ObjectStoreSecondaryTierManager. +"""Mock-based unit tests for ObjectStoreSecondaryTierManager. These tests replace the NIXL backend with an in-memory mock so they run without S3 credentials or a live object store. They verify the manager's @@ -347,7 +346,8 @@ def test_failed_load_marks_verdict_negative(self): the scheduler would re-issue the same doomed promotion every step for the life of the request. The mark is served from cache with no re-probe, so even though the mock object is still 'present' the SAME - request now resolves to MISS.""" + request now resolves to MISS. + """ ctx = ReqContext(req_id="obj-livelock") self.tier.submit_store(make_job(1, [key(1)], [0])) assert all(r.success for r in drain(self.tier)) @@ -651,7 +651,8 @@ def test_store_event_uses_configured_locality(self, locality, expected): def test_mixed_job_results_emit_event_only_for_successful_job(self): """With a failed and a successful store job resolving in the same poll, exactly one event is emitted and its keys belong to the - successful job.""" + successful job. + """ original = self.agent.check_xfer_state self.agent.check_xfer_state = lambda h: "ERR" if h._id == 0 else original(h) self.tier.submit_store(make_job(1, [key(1)], [0])) # handle 0: fails diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 771153ec4734..5cc5f48e3a03 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager. +"""Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager. These tests verify: 1. Basic tiered offloading operations (store, load, lookup) @@ -274,7 +273,8 @@ def test_basic_store_and_lookup(self): def test_every_lookup_result_has_a_declared_cascade_behavior(): """A new LookupResult must state what the cascade does with it, rather - than falling into whichever branch happens to catch it.""" + than falling into whichever branch happens to catch it. + """ assert set(_CASCADE_SUPPLY_BEHAVIOR) == set(LookupResult) @@ -321,7 +321,8 @@ def _start_request(self, req_context: ReqContext = _CTX): def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): """A failed promotion still finalizes the primary slots with - success=False; the tier corrects its own verdict.""" + success=False; the tier corrects its own verdict. + """ from unittest.mock import patch from vllm.v1.kv_offload.tiering.base import JobResult @@ -1092,7 +1093,8 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( def test_on_new_request_lifecycle(self, manager_setup): """Policy defaults to CHUNK_LEVEL, escalates when a tier requests it, - and is cleaned up on on_request_finished.""" + and is cleaned up on on_request_finished. + """ # Default: all tiers return CHUNK_LEVEL ctx = ReqContext(req_id="req_policy_lifecycle") result = self.manager.on_new_request(ctx) @@ -1169,7 +1171,8 @@ def test_prepare_store_cascades_existing_chunks_to_request_level_tiers( def _make_request_level_request(self, req_id: str) -> ReqContext: """Start a request for which tier1 asks for request-level offloading, - with tier1's submit_store wrapped so the cascade can be observed.""" + with tier1's submit_store wrapped so the cascade can be observed. + """ self.secondary_tier1.on_new_request = lambda req_context: ( RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) ) @@ -1191,7 +1194,8 @@ def _cascaded_keys_for(self, req_id: str) -> list[set[OffloadKey]]: def test_cascade_rejects_retry_from_primary(self, manager_setup): """RETRY would be parked with no guarantee of ever draining, so the cascade refuses it. No real primary tier returns it, so this is the one - disposition that has to be forced.""" + disposition that has to be forced. + """ keys = to_keys(range(3)) self._start_request() assert self.manager.prepare_store(keys, _CTX) is not None @@ -1206,7 +1210,8 @@ def test_cascade_rejects_retry_from_primary(self, manager_setup): def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: """Leave a GPU->primary write for `keys` open, so they look present to - prepare_store but are not yet readable.""" + prepare_store but are not yet readable. + """ writer_ctx = ReqContext(req_id="req_writer") self._start_request(writer_ctx) assert self.manager.prepare_store(keys, writer_ctx) is not None @@ -1215,7 +1220,8 @@ def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setup): """A key another request is still writing must reach the peer once the write lands, not be dropped: prepare_store already counts it as stored - and the scheduler advances past its chunk, so nothing re-offers it.""" + and the scheduler advances past its chunk, so nothing re-offers it. + """ keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1232,7 +1238,8 @@ def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setu def test_deferred_cascade_holds_request_from_finalization(self, manager_setup): """A request that finishes with keys still deferred must stay alive, or - its tiers are torn down before the peer is ever served.""" + its tiers are torn down before the peer is ever served. + """ keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1255,7 +1262,8 @@ def test_cascade_drops_deferred_keys_whose_primary_write_failed( self, manager_setup ): """A failed write frees the chunk, so the deferred key resolves to MISS - and the request finalizes instead of parking forever.""" + and the request finalizes instead of parking forever. + """ keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1275,7 +1283,8 @@ def test_cascade_drops_deferred_keys_whose_primary_write_failed( def test_reset_cache_clears_orchestrator_state(self, manager_setup): """reset_cache wipes every kind of orchestrator state and resets primary tier; pending submissions are dropped without being sent - to the secondary tier. Active request state is retained.""" + to the secondary tier. Active request state is retained. + """ # Cascade — populates primary chunks and leaves cascade jobs # in _jobs (the synchronous example tier has already # queued completions); reset_cache's drain loop will pick them up. @@ -1367,7 +1376,8 @@ def test_tier_filter_skips_filtered_secondary( self, manager_setup, load_tier_filter ): """Filter excluding secondary medium returns MISS from secondaries - even when they hold the chunk; primary is unaffected.""" + even when they hold the chunk; primary is unaffected. + """ chunks = to_keys(range(2)) # Put one chunk in primary, one only in secondary self._start_request() diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index 7e60eeb2e3b8..fccdf56409d6 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -98,7 +98,7 @@ def __init__(self, workload_index: int, logitproc_type: LogitprocType): self.params = _sampling_params_from_logitproc(logitproc_type) def __str__(self): - """For debugging""" + """For debugging.""" summ = ", ".join(f"{k}={v}" for k, v in vars(self).items()) return f"MyClass({summ})" @@ -117,7 +117,7 @@ def _generate_fake_sampling_metadata( vocab_size: int, device: torch.device, ) -> SamplingMetadata: - """Generate fake sampling metadata with fake logitsprocs""" + """Generate fake sampling metadata with fake logitsprocs.""" output_token_ids: list[list[int]] = [] prompt_token_ids: list[list[int]] = [] for _ in range(batch_size): @@ -176,7 +176,7 @@ def _generate_fake_sampling_metadata( def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes: - """Generate fake logits and sampling metadata""" + """Generate fake logits and sampling metadata.""" fake_logits = create_fake_logits(batch_size, VOCAB_SIZE) # Create one dominant token per batch, to support min-p test for i in range(batch_size): @@ -192,7 +192,7 @@ def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes: def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingParams: - """Customize request SamplingParams for a specified logitproc""" + """Customize request SamplingParams for a specified logitproc.""" # SamplingParams for req with no logitproc kwargs = {"min_p": 0.0, "logit_bias": None, "min_tokens": 0} if fxn := logitsprocs_test_mapping[logitproc_type].gen_request_fxn: @@ -220,6 +220,7 @@ def _generate_mixed_logitsprocs_batch_params( Returns: List of per-request params which configure the engine for that request's enabled logitproc + """ batch_size = len(logitsprocs_types) * reqs_per_logitproc # Generate multiple repeats of key params for each logitproc; @@ -251,7 +252,7 @@ def _raise_error_invalid( def _logit_bias_params(kwargs: dict) -> None: - """Logit bias config""" + """Logit bias config.""" kwargs["logit_bias"] = { random.randint(0, VOCAB_SIZE - 1): random.choice([-0.1, 0.2]) } @@ -265,7 +266,7 @@ def _logit_bias_validate( request_params: LogitsProcsRequestParams, step_idx: int, ) -> None: - """Validate logit bias logitproc applied correctly""" + """Validate logit bias logitproc applied correctly.""" logit_bias = request_params.params.logit_bias logits_old = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu() logits_new = logits_new[batch_index].cpu() @@ -301,7 +302,7 @@ def _logit_bias_validate( def _min_p_params(kwargs: dict) -> None: - """Min-p logitproc config""" + """Min-p logitproc config.""" kwargs["min_p"] = 0.1 @@ -313,7 +314,7 @@ def _min_p_validate( request_params: LogitsProcsRequestParams, step_idx: int, ) -> None: - """Validate min-p logitproc applied correctly""" + """Validate min-p logitproc applied correctly.""" for token_id in range(VOCAB_SIZE): logits_for_token = logits_new[batch_index][token_id] if token_id == 0: @@ -347,7 +348,7 @@ def _min_p_validate( def _min_tokens_params(kwargs: dict) -> None: - """Min-tokens logitproc config""" + """Min-tokens logitproc config.""" kwargs["min_tokens"] = MIN_TOKENS_LEN_THRESHOLD kwargs["stop_token_ids"] = [ np.random.randint(0, VOCAB_SIZE - 1) @@ -363,7 +364,7 @@ def _min_tokens_validate( request_params: LogitsProcsRequestParams, step_idx: int, ) -> None: - """Validate min-tokens logitsproc applied correctly""" + """Validate min-tokens logitsproc applied correctly.""" ref_num_out_tokens = len(request_params.out_tokens) min_reached = ref_num_out_tokens >= MIN_TOKENS_LEN_THRESHOLD ref_all_stop_token_ids = request_params.params.all_stop_token_ids @@ -560,7 +561,7 @@ def test_min_tokens_restores_all_masked_structured_output_stop_token_spec_decode def _thinking_budget_params(kwargs: dict) -> None: - """Set SamplingParams kwargs for thinking token budget tests""" + """Set SamplingParams kwargs for thinking token budget tests.""" kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET @@ -661,7 +662,7 @@ def _none_validate( request_params: LogitsProcsRequestParams, step_idx: int, ) -> None: - """Validate that no logits processors are applied""" + """Validate that no logits processors are applied.""" logits = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu() ref_logits = logits_new[batch_index] if not torch.all(ref_logits == logits): @@ -706,7 +707,7 @@ class LogitsprocTestHelpers(NamedTuple): def _get_test_cases() -> list[list[str]]: - """Each test case is a set of logitsprocs""" + """Each test case is a set of logitsprocs.""" logitsprocs_types = list(logitsprocs_test_mapping.keys()) # Isolate thinking-budget handling from other processors to avoid cross-talk. @@ -1343,7 +1344,8 @@ def test_thinking_budget_invalid_budget_rejected(invalid_budget): def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index(): """Test thinking budget enforced for a long thinking run, - then a natural end marker.""" + then a natural end marker. + """ h = ThinkingBudgetStateHolder( MockReasoningConfig(), 8, 0, torch.device("cpu"), False ) diff --git a/tests/v1/logits_processors/test_custom_offline.py b/tests/v1/logits_processors/test_custom_offline.py index 302a72057eaa..0d0b9531dcf3 100644 --- a/tests/v1/logits_processors/test_custom_offline.py +++ b/tests/v1/logits_processors/test_custom_offline.py @@ -86,8 +86,8 @@ def _run_test(kwargs: dict, logitproc_loaded: bool) -> None: Args: kwargs: `LLM` constructor kwargs logitproc_loaded: server has loaded dummy logitproc if True - """ + """ # Create a vLLM instance and load custom logitproc llm_logitproc = LLM( model=MODEL_NAME, @@ -131,7 +131,7 @@ def _run_test(kwargs: dict, logitproc_loaded: bool) -> None: @create_new_process_for_each_test() @pytest.mark.parametrize("logitproc_source", list(CustomLogitprocSource)) def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource): - """Test offline Python interface for passing custom logitsprocs + """Test offline Python interface for passing custom logitsprocs. Construct an `LLM` instance which loads a custom logitproc that has a well-defined behavior (mask out all tokens except one `target_token`) @@ -158,8 +158,8 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource logitproc_source: what source (entrypoint, fully-qualified class name (FQCN), class object, or None) the user pulls the logitproc from - """ + """ # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") set_random_seed(40) @@ -192,7 +192,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource @create_new_process_for_each_test() def test_custom_logitsprocs_req(monkeypatch): - """Test passing request-level logits processor to offline Python interface + """Test passing request-level logits processor to offline Python interface. Wrap a request-level logits processor to create a batch level logits processor that has a well-defined behavior (mask out all tokens except one @@ -213,8 +213,8 @@ def test_custom_logitsprocs_req(monkeypatch): Args: monkeypatch: for setting env vars - """ + """ # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") set_random_seed(40) @@ -258,6 +258,7 @@ def test_rejects_custom_logitsprocs( logitproc_source: what source (entrypoint, fully-qualified class name (FQCN), or class object) the user pulls the logitproc from + """ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") set_random_seed(40) diff --git a/tests/v1/logits_processors/test_custom_online.py b/tests/v1/logits_processors/test_custom_online.py index 3b7add3b80f4..4a7ad8cc0069 100644 --- a/tests/v1/logits_processors/test_custom_online.py +++ b/tests/v1/logits_processors/test_custom_online.py @@ -77,7 +77,6 @@ def server(default_server_args, request, monkeypatch): (2) No --logits-processors cli arg; inject a dummy logits processor entrypoint """ - # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") @@ -118,7 +117,7 @@ async def client(server): [MODEL_NAME], ) def test_custom_logitsprocs(server, model_name: str): - """Test custom logitsprocs when starting OpenAI server from CLI + """Test custom logitsprocs when starting OpenAI server from CLI. Launch vLLM OpenAI-compatible server, configured to load a custom logitproc that has a well-defined behavior (mask out all tokens except one @@ -130,7 +129,6 @@ def test_custom_logitsprocs(server, model_name: str): Validate that requests which activate the custom logitproc, repeat the same token """ - import asyncio async def _async_main(srv, mn): @@ -180,8 +178,7 @@ async def _run(client): async def test_invalid_custom_logitsproc_arg( client: openai.AsyncOpenAI, model_name: str ): - """Test that request with invalid custom logitsproc is rejected""" - + """Test that request with invalid custom logitsproc is rejected.""" prompt = "Hello, my name is" # Pass invalid (non-int) target_token value to dummy logits processor request_keyword_args: dict[str, Any] = { diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index fb45b73c9893..85ad2032cde3 100644 --- a/tests/v1/logits_processors/utils.py +++ b/tests/v1/logits_processors/utils.py @@ -37,7 +37,7 @@ class CustomLogitprocSource(Enum): - """How to source a logitproc for testing purposes""" + """How to source a logitproc for testing purposes.""" LOGITPROC_SOURCE_NONE = auto() # No custom logitproc LOGITPROC_SOURCE_ENTRYPOINT = auto() # Via entrypoint @@ -55,7 +55,7 @@ class CustomLogitprocSource(Enum): 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): @@ -73,7 +73,7 @@ def __init__( self.req_info: dict[int, int] = {} def is_argmax_invariant(self) -> bool: - """Never impacts greedy sampling""" + """Never impacts greedy sampling.""" return False def update_state(self, batch_update: BatchUpdate | None): @@ -115,7 +115,7 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: class EntryPoint: - """Dummy entrypoint class for logitsprocs testing""" + """Dummy entrypoint class for logitsprocs testing.""" def __init__(self): self.name = DUMMY_LOGITPROC_ENTRYPOINT @@ -126,7 +126,7 @@ def load(self): class EntryPoints(list): - """Dummy EntryPoints class for logitsprocs testing""" + """Dummy EntryPoints class for logitsprocs testing.""" def __init__(self, group: str): # Emulate list-like functionality @@ -138,10 +138,11 @@ def __init__(self, group: str): 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__( @@ -157,7 +158,8 @@ def __call__( class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): """Example of wrapping a fake request-level logit processor to create a - batch-level logits processor""" + batch-level logits processor + """ def is_argmax_invariant(self) -> bool: return False @@ -178,6 +180,7 @@ def new_req_logits_processor( Returns: `Callable` request logits processor, or None + """ target_token: Any | None = params.extra_args and params.extra_args.get( "target_token" diff --git a/tests/v1/metrics/test_engine_logger_apis.py b/tests/v1/metrics/test_engine_logger_apis.py index 2e243c23cbf9..5741df0d56ed 100644 --- a/tests/v1/metrics/test_engine_logger_apis.py +++ b/tests/v1/metrics/test_engine_logger_apis.py @@ -13,8 +13,7 @@ @pytest.fixture def log_stats_enabled_engine_args(): - """ - Shared fixture providing common AsyncEngineArgs configuration + """Shared fixture providing common AsyncEngineArgs configuration used across multiple tests. """ return AsyncEngineArgs( @@ -27,10 +26,7 @@ def log_stats_enabled_engine_args(): @pytest.mark.asyncio async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args): - """ - RayPrometheusStatLogger should replace the default PrometheusStatLogger - """ - + """RayPrometheusStatLogger should replace the default PrometheusStatLogger.""" engine = AsyncLLM.from_engine_args( log_stats_enabled_engine_args, stat_loggers=[RayPrometheusStatLogger] ) @@ -40,8 +36,7 @@ async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args): @pytest.mark.asyncio async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args): - """ - It's still possible to use custom stat loggers exclusively by passing + """It's still possible to use custom stat loggers exclusively by passing disable_log_stats=True in addition to a list of custom stat loggers. """ # Create engine_args with disable_log_stats=True for this test diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index c94ec1b701ca..32c203705cca 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests for the analytic estimators in metrics/flops.py. -""" +"""Tests for the analytic estimators in metrics/flops.py.""" import types from types import SimpleNamespace @@ -156,8 +154,7 @@ def test_base_attention_config_parser_with_gqa(): def test_base_attention_config_parser_without_gqa(): - """ - Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not + """Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not specified. """ hf_config = Qwen3Config( @@ -707,9 +704,7 @@ def test_ffn_per_gpu_with_pipeline_parallelism(): def test_moe_per_gpu_with_expert_parallelism(): - """ - Test MoE metrics with expert parallelism - verifies num_activated_experts bug fix. - """ + """Test MoE metrics with expert parallelism, covering num_activated_experts.""" hf_config = Qwen3MoeConfig( hidden_size=2048, intermediate_size=8192, @@ -765,9 +760,7 @@ def test_moe_per_gpu_with_expert_parallelism(): def test_moe_per_gpu_expert_activation_accounting(): - """ - Test that MoE correctly accounts for expert activations with small batch sizes. - """ + """Test MoE expert-activation accounting with small batch sizes.""" hf_config = Qwen3MoeConfig( hidden_size=2048, intermediate_size=8192, @@ -1248,7 +1241,8 @@ def test_mla_per_gpu_with_pipeline_parallelism(): def test_mla_model_metrics_excludes_standard_attention(): """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics, - for DeepSeek MLA models.""" + for DeepSeek MLA models. + """ hf_config = DeepseekV3Config( hidden_size=7168, num_attention_heads=128, diff --git a/tests/v1/metrics/test_ray_metrics.py b/tests/v1/metrics/test_ray_metrics.py index d8440d3bb02a..7bc49aae6e64 100644 --- a/tests/v1/metrics/test_ray_metrics.py +++ b/tests/v1/metrics/test_ray_metrics.py @@ -45,7 +45,8 @@ def test_engine_log_metrics_ray( max_tokens: int, ) -> None: """Simple smoke test, verifying this can be used without exceptions. - Need to start a Ray cluster in order to verify outputs.""" + Need to start a Ray cluster in order to verify outputs. + """ @ray.remote(num_gpus=1) class EngineTestActor: @@ -75,7 +76,6 @@ async def run(self): def test_sanitized_opentelemetry_name(): """Test the metric name sanitization logic for Ray.""" - # Only a-z, A-Z, 0-9, _, test valid characters are preserved valid_name = "valid_metric_123_abcDEF" assert ( @@ -120,7 +120,8 @@ def test_sanitized_opentelemetry_name(): def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock: """Swap the wrapper's underlying Ray metric for a MagicMock while preserving the real metric's ``_tag_keys`` (labels() reads them to - validate arity).""" + validate arity). + """ real_metric = wrapper.metric mock = MagicMock() mock._tag_keys = real_metric._tag_keys @@ -130,7 +131,8 @@ def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock: def test_ray_counter_labels_returns_independent_children(): """RayCounterWrapper.labels() must return distinct labeled children that - each carry their own tag set.""" + each carry their own tag set. + """ base = RayCounterWrapper( name="vllm_test_finish_reason", documentation="", @@ -150,7 +152,8 @@ def test_ray_counter_labels_returns_independent_children(): def test_ray_counter_inc_forwards_per_child_tags(): """.inc() on a labeled counter must forward that child's tags to the - underlying Ray metric (not rely on a shared set_default_tags).""" + underlying Ray metric (not rely on a shared set_default_tags). + """ wrapper = RayCounterWrapper( name="vllm_test_counter_tag_forward", documentation="", @@ -215,7 +218,8 @@ def test_ray_histogram_labels_returns_independent_children_and_forwards_tags(): def test_ray_counter_labels_accepts_non_string_label_values(): """RayPrometheusStatLogger passes ``str(idx)`` for engine indexes; this covers the coercion path for any caller that passes a non-string label - value positionally.""" + value positionally. + """ wrapper = RayCounterWrapper( name="vllm_test_nonstr_label", documentation="", @@ -238,7 +242,8 @@ def test_ray_counter_labels_arity_validation(): def test_unlabeled_inc_carries_replica_id(): """Recording on an unlabeled metric must still pass ReplicaId — it's a - declared tag_key and Ray rejects updates that omit any declared key.""" + declared tag_key and Ray rejects updates that omit any declared key. + """ wrapper = RayCounterWrapper( name="vllm_test_unlabeled_replica_id", documentation="", @@ -251,7 +256,8 @@ def test_unlabeled_inc_carries_replica_id(): def test_double_labels_raises(): """labels() on an already-labeled child should raise, mirroring the - prometheus_client contract.""" + prometheus_client contract. + """ wrapper = RayCounterWrapper( name="vllm_test_double_labels", documentation="", diff --git a/tests/v1/sample/test_batched_count_greater_than.py b/tests/v1/sample/test_batched_count_greater_than.py index c9ace93c6961..0e2d19288245 100644 --- a/tests/v1/sample/test_batched_count_greater_than.py +++ b/tests/v1/sample/test_batched_count_greater_than.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Test that batched_count_greater_than does not trigger 0/1 specialization -recompiles when batch_size varies.""" +recompiles when batch_size varies. +""" import torch diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 3bf26d51706e..372e1c815643 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -127,6 +127,7 @@ def _repeat_logprob_config( `logprob_prompt_logprob_list` enough times to match the number of `test_prompts`, or else is truncated to match the number of `test_prompts` + """ num_test_prompts = len(test_prompts) # Make sure there is a logprobs configuration for each test prompt @@ -312,7 +313,7 @@ def test_get_logprobs_and_prompt_logprobs( temperature: float, example_prompts: list[str], ) -> None: - """Test V1 Engine logprobs & prompt logprobs + """Test V1 Engine logprobs & prompt logprobs. Exercise a variety of combinations of `logprobs` and `prompt_logprobs` settings and validate that @@ -337,6 +338,7 @@ def test_get_logprobs_and_prompt_logprobs( batch_logprobs_composition: logprobs configuration for test batch temperature: "temperature" sampling parameter example_prompts: example prompt fixture + """ vllm_config = vllm_model.llm.llm_engine.vllm_config do_apc = vllm_config.cache_config.enable_prefix_caching @@ -389,7 +391,7 @@ def test_get_logprobs_and_prompt_logprobs( def test_max_logprobs(): - """vLLM v1 engine should fail a request with `logprobs > max_logprobs` + """VLLM v1 engine should fail a request with `logprobs > max_logprobs` Should also fail for `prompt_logprobs > max_logprobs` APC should not matter as this test checks basic request validation. """ @@ -431,11 +433,12 @@ def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]): def test_none_logprobs(vllm_model, example_prompts): - """Engine should return `logprobs` and `prompt_logprobs` as `None` + """Engine should return `logprobs` and `prompt_logprobs` as `None`. Args: vllm_model: vLLM model fixture example_prompts: list of example prompts (test fixture) + """ max_tokens = 5 @@ -459,11 +462,12 @@ def test_none_logprobs(vllm_model, example_prompts): def test_zero_logprobs(vllm_model, example_prompts): - """Engine should return sampled token and prompt token logprobs + """Engine should return sampled token and prompt token logprobs. Args: vllm_model: vLLM model fixture example_prompts: list of example prompts (test fixture) + """ max_tokens = 5 @@ -491,10 +495,11 @@ def test_zero_logprobs(vllm_model, example_prompts): def test_all_logprobs(example_prompts): - """Engine should return all vocabulary logprobs and prompt logprobs + """Engine should return all vocabulary logprobs and prompt logprobs. Args: example_prompts: list of example prompts (test fixture) + """ with VllmRunner( "facebook/opt-125m", @@ -568,7 +573,8 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode): def test_prompt_logprobs_mode(): """prompt_logprobs must respect logprobs_mode: *_logits and *_logprobs must return different values. Prompt tokens skip sampling processors, - so processed_* == raw_* on the prompt side.""" + so processed_* == raw_* on the prompt side. + """ from vllm import LLM values: dict[str, float] = {} @@ -1127,6 +1133,7 @@ def test_spec_decode_logprobs( logprobs_mode: logprobs mode. model_setup: Tuple of (method, base model name, speculative_config dict, top_logprobs). + """ from vllm import LLM @@ -1224,7 +1231,6 @@ def test_prompt_logprobs_with_chunking_and_preemption(): This test ensures that the num_prompt_logprobs tracking persists across preemptions and prefill chunks. """ - # Create prompts that will trigger chunking and preemption prompts = [ "The following numbers of the sequence " diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index b02e53af8dac..abb069a26a78 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/tests/v1/sample/test_rejection_sampler.py @@ -54,7 +54,8 @@ def create_logits_tensor( token_idx_to_override: int | None = None, ) -> torch.Tensor: """Helper function to create logits tensor that - will produce desired token ids on argmax""" + will produce desired token ids on argmax + """ token_ids = [tokens[:-1] for tokens in output_token_ids] num_total_tokens = sum(len(tokens) for tokens in token_ids) logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE_TYPE) @@ -131,7 +132,7 @@ def create_sampling_metadata( ########################### Tests for Greedy Sampling ################### def test_perfect_match(rejection_sampler): - """Test when output tokens perfectly match speculated tokens""" + """Test when output tokens perfectly match speculated tokens.""" spec_tokens = [[1, 2, 3]] output_tokens = [[1, 2, 3, 4]] # 4 is the bonus token @@ -152,7 +153,7 @@ def test_perfect_match(rejection_sampler): def test_early_mismatch(rejection_sampler): - """Test when there's an early mismatch in tokens""" + """Test when there's an early mismatch in tokens.""" spec_tokens = [[1, 2, 3]] output_tokens = [[1, 5, 3, 4]] # Mismatch at position 1 @@ -177,7 +178,7 @@ def test_early_mismatch(rejection_sampler): def test_multiple_sequences(rejection_sampler): - """Test handling multiple sequences of speculated tokens""" + """Test handling multiple sequences of speculated tokens.""" spec_tokens = [[1, 2], [3]] output_tokens = [[1, 2, 5], [3, 4]] # Two sequences with bonus tokens 5 and 4 @@ -202,7 +203,7 @@ def test_multiple_sequences(rejection_sampler): def test_single_token_sequence(rejection_sampler): - """Test handling sequences with single token""" + """Test handling sequences with single token.""" spec_tokens = [[1]] output_tokens = [[1, 2]] # Single token with bonus token 2 @@ -223,7 +224,7 @@ def test_single_token_sequence(rejection_sampler): def test_empty_sequence(rejection_sampler): - """Test handling empty sequence of speculated tokens""" + """Test handling empty sequence of speculated tokens.""" spec_tokens: list[list[int]] = [[]] output_tokens = [[5]] # Just the bonus token @@ -244,7 +245,7 @@ def test_empty_sequence(rejection_sampler): def test_multiple_mismatches(rejection_sampler): - """Test handling multiple sequences with mismatches""" + """Test handling multiple sequences with mismatches.""" spec_tokens = [[1, 2, 3], [4, 5, 6]] output_tokens = [[1, 2, 7, 6], [4, 8, 6, 9]] # Mismatches in both sequences @@ -286,7 +287,7 @@ def test_multiple_mismatches(rejection_sampler): ], ) def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens, expected): - """Parametrized test for various matching scenarios""" + """Parametrized test for various matching scenarios.""" metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) bonus_token_tensor = torch.tensor( @@ -486,6 +487,7 @@ def estimate_rejection_sampling_pdf( Returns: Estimated probability distribution of the output tokens. + """ mock_sampler = Mock(spec=Sampler) mock_sampler.logprobs_mode = "raw_logprobs" @@ -656,7 +658,7 @@ def _test_masked_logits( @pytest.mark.parametrize("top_k", [1, 5, 99]) def test_top_k(rejection_sampler, top_k): - """Test rejection sampling with top-k sampling""" + """Test rejection sampling with top-k sampling.""" vocab_size = 100 batch_size = 100 num_draft_tokens = 3 @@ -699,7 +701,7 @@ def test_top_k(rejection_sampler, top_k): @pytest.mark.parametrize("top_p", [0.5, 0.9, 0.99]) def test_top_p(rejection_sampler, top_p): - """Test rejection sampling with top-p sampling""" + """Test rejection sampling with top-p sampling.""" vocab_size = 100 batch_size = 100 num_draft_tokens = 3 @@ -746,7 +748,7 @@ def test_top_p(rejection_sampler, top_p): ########################### Tests for Logit Processors ################### def test_frequency_penalties(rejection_sampler): - """Test rejection sampling with frequency penalties""" + """Test rejection sampling with frequency penalties.""" spec_tokens = [[1, 1, 1], [], [1, 1, 1]] output_tokens = [[1, 1, 1, 1], [7], [1, 1, 1, 1]] # 1, 7 and 1 are the bonus tokens @@ -829,7 +831,7 @@ def test_bad_words(rejection_sampler): def test_allowed_token_ids(rejection_sampler): - """Test rejection sampling with allowed token ids""" + """Test rejection sampling with allowed token ids.""" spec_tokens = [[1, 2, 10], [10, 5, 3], [7, 10, 12]] output_tokens = [[1, 2, 10, 5], [10, 5, 10, 5], [7, 10, 12, 5]] # Not allowed tokens: diff --git a/tests/v1/sample/test_sampler.py b/tests/v1/sample/test_sampler.py index c67199fa4077..b6d4fcc9e8a3 100644 --- a/tests/v1/sample/test_sampler.py +++ b/tests/v1/sample/test_sampler.py @@ -169,8 +169,7 @@ def _create_default_sampling_metadata( def _create_weighted_output_token_list( batch_size: int, vocab_size: int ) -> tuple[list[list[int]], list[list[int]]]: - """ - Creates an output token list where each token occurs a distinct + """Creates an output token list where each token occurs a distinct number of times. For each batch, a random subset of token IDs is selected from the @@ -185,6 +184,7 @@ def _create_weighted_output_token_list( - The second element is a list of distinct token IDs for each batch, ordered by their frequency in the corresponding output list. + """ output_token_ids: list[list[int]] = [] sorted_token_ids_in_output: list[list[int]] = [] @@ -206,8 +206,7 @@ def _create_weighted_output_token_list( def test_sampler_presence_penalty( device: str, batch_size: int, presence_penalty: float ): - """ - Test to verify that if presence penalty is enabled then tokens + """Test to verify that if presence penalty is enabled then tokens are penalized as per their presence in the existing output. """ torch.set_default_device(device) @@ -256,8 +255,7 @@ def test_sampler_presence_penalty( def test_sampler_frequency_penalty( device: str, batch_size: int, frequency_penalty: float ): - """ - Test to verify that if frequency penalty is enabled then tokens are + """Test to verify that if frequency penalty is enabled then tokens are penalized as per their frequency of occurrence. """ torch.set_default_device(device) @@ -312,8 +310,7 @@ def test_sampler_frequency_penalty( def test_sampler_repetition_penalty( device: str, batch_size: int, repetition_penalty: float ): - """ - Test to verify that when the repetition penalty is enabled, tokens + """Test to verify that when the repetition penalty is enabled, tokens are penalized based on their presence in the prompt or the existing output. """ @@ -370,8 +367,7 @@ def test_sampler_repetition_penalty( def test_sampler_allowed_token_ids( device: str, batch_size: int, num_allowed_token_ids: int ): - """ - Test to verify that when the repetition penalty is enabled, tokens + """Test to verify that when the repetition penalty is enabled, tokens are penalized based on their presence in the prompt or the existing output. """ @@ -416,8 +412,7 @@ def test_sampler_allowed_token_ids( def test_sampler_bad_words( device: str, batch_size: int, bad_words_lengths: tuple[int, ...] ): - """ - Test to verify that when the bad words restriction is present, tokens + """Test to verify that when the bad words restriction is present, tokens are penalized based on their match with the bad words. """ torch.set_default_device(device) diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py index d385e96b7a21..12c1048ff2c1 100644 --- a/tests/v1/sample/test_sampling_params_e2e.py +++ b/tests/v1/sample/test_sampling_params_e2e.py @@ -17,7 +17,6 @@ def llm() -> LLM: def test_n_gt_1(llm): """ParallelSampling is supported.""" - params = SamplingParams(n=3) outputs = llm.generate(PROMPT, params) assert len(outputs[0].outputs) == 3 @@ -25,7 +24,6 @@ def test_n_gt_1(llm): def test_penalties(llm): """Check that we do not get errors if applied.""" - params = SamplingParams( temperature=1.2, presence_penalty=1.2, @@ -40,7 +38,6 @@ def test_penalties(llm): def test_stop(llm): """Check that we respect the stop words.""" - output = llm.generate(PROMPT, SamplingParams(temperature=0)) split_text = output[0].outputs[0].text.split() @@ -64,7 +61,6 @@ def test_stop(llm): def test_stop_token_ids(llm): """Check that we respect the stop token ids.""" - output = llm.generate(PROMPT, SamplingParams(temperature=0)) stop_token_id_0 = output[0].outputs[0].token_ids[5] @@ -83,7 +79,6 @@ def test_stop_token_ids(llm): def test_detokenize_false(llm): """Check that detokenize=False option works.""" - output = llm.generate(PROMPT, SamplingParams(detokenize=False)) assert len(output[0].outputs[0].token_ids) > 0 assert len(output[0].outputs[0].text) == 0 @@ -106,7 +101,6 @@ def test_detokenize_false(llm): def test_bad_words(llm): """Check that we respect bad words.""" - tokenizer = llm.get_tokenizer() def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool: @@ -147,7 +141,6 @@ def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool: def test_allowed_token_ids(llm): """Check that we can use allowed_token_ids.""" - TOKEN_ID = 10 allowed_token_ids = [TOKEN_ID] output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids)) @@ -176,7 +169,6 @@ def test_allowed_token_ids(llm): def test_seed(llm): """Check that seed impacts randomness.""" - out_1 = llm.generate(PROMPT, SamplingParams(seed=42)) out_2 = llm.generate(PROMPT, SamplingParams(seed=42)) out_3 = llm.generate(PROMPT, SamplingParams(seed=43)) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 5a0ce5cc089f..9f53a8c28760 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -50,8 +50,7 @@ def _seed_default_generator(seed: int) -> None: @pytest.fixture(autouse=True) def reset_default_device(): - """ - Explicitly set the default device, which can affect subsequent tests. + """Explicitly set the default device, which can affect subsequent tests. Adding this fixture helps avoid this problem. """ original_device = torch.get_default_device() @@ -171,8 +170,7 @@ def test_topk_impl_equivalence(): "interface differences between Python and FlashInfer implementations" ) def test_flashinfer_sampler(): - """ - This test verifies that the FlashInfer top-k and top-p sampling + """This test verifies that the FlashInfer top-k and top-p sampling implementation produces the same results as the Python implementation. NOTE: FlashInfer did not directly expose an interface for fused top-k and @@ -886,7 +884,8 @@ def setup(self): def _make_logits(self, pattern: str) -> torch.Tensor: """Build (BATCH, VOCAB) logits with `pattern` applied to row 0 - (rows 1..B-1 stay clean so we can detect cross-row corruption).""" + (rows 1..B-1 stay clean so we can detect cross-row corruption). + """ logits = ( torch.randn( self.BATCH, diff --git a/tests/v1/sample/utils.py b/tests/v1/sample/utils.py index 907be3614b9c..19c50f79b1cf 100644 --- a/tests/v1/sample/utils.py +++ b/tests/v1/sample/utils.py @@ -15,7 +15,7 @@ class BatchLogprobsComposition(Enum): - """Types of logprobs configs to include in test batch""" + """Types of logprobs configs to include in test batch.""" NONE = 0 SAMPLE = 1 @@ -29,7 +29,7 @@ class BatchLogprobsComposition(Enum): def get_test_batch( batch_logprobs_composition: BatchLogprobsComposition, ) -> BatchLogprobsSpecType: - """Generate logprobs configs for a batch of requests + """Generate logprobs configs for a batch of requests. A given request's logprobs configuration is (1) num_sample_logprobs and (2) num_prompt_logprobs. The batch logprobs configuration is the list of request @@ -53,9 +53,9 @@ def get_test_batch( batch_logprobs_composition: types of logprobs configs to include in batch Returns: - list of (Optional[num_sample_logprobs], Optional[num_prompt_logprobs]) tuples + """ if batch_logprobs_composition == BatchLogprobsComposition.NONE: # No requests with sample or prompt logprobs @@ -101,7 +101,7 @@ def assert_incr_detok_str_matches_non_incr_detok_str( non_incremental_detokenization_str: str, msg: str, ) -> None: - """Compare incrementally detok. text to non-incrementally detok. text + """Compare incrementally detok. text to non-incrementally detok. text. Fail if the strings mismatch after non-alphanumeric characters are stripped out. @@ -120,6 +120,7 @@ def assert_incr_detok_str_matches_non_incr_detok_str( non_incremental_detokenization_str: non-incrementally-detokenized logprob tokens msg: error message if `assert` fails + """ rgx = r"[^a-zA-Z0-9]+" assert re.sub(rgx, "", incremental_detokenization_str) == re.sub( @@ -128,13 +129,14 @@ def assert_incr_detok_str_matches_non_incr_detok_str( def compute_correct_cumulative_logprob(completion_output: CompletionOutput) -> float: - """Compute known-good value for evaluating cumulative logprob + """Compute known-good value for evaluating cumulative logprob. Args: completion_output: completion output from engine Returns: Known-good cumulative logprob value + """ token_ids = completion_output.token_ids logprobs = completion_output.logprobs @@ -170,7 +172,7 @@ def create_prompt_tokens_tensor( class LogitsprocsTestFakes(NamedTuple): - """Wraps fake data structures to support testing""" + """Wraps fake data structures to support testing.""" logits: torch.Tensor sampling_metadata: SamplingMetadata @@ -186,6 +188,7 @@ def get_logitsprocs_by_cls( Returns: Iterator over logits processors + """ return ( lp for lp in self.sampling_metadata.logitsprocs.all if isinstance(lp, cls) @@ -201,7 +204,8 @@ def fake_update_logitsprocs_state( batch_update: BatchUpdate | None, ) -> None: """Imitate logits processors persistent batch state update - in engine core""" + in engine core + """ for logitproc in test_fakes.get_logitsprocs(): logitproc.update_state(batch_update) holder = test_fakes.sampling_metadata.thinking_budget_state_holder diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index 3a513945eaaa..b8bf620b33f5 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/tests/v1/shutdown/test_delete.py @@ -35,6 +35,7 @@ async def test_async_llm_delete( model: model under test tensor_parallel_size: degree of tensor parallelism send_one_request: send one request to engine before deleting + """ if current_platform.device_count() < tensor_parallel_size: pytest.skip(reason="Not enough CUDA devices") @@ -87,6 +88,7 @@ def test_llm_delete( tensor_parallel_size: degree of tensor parallelism enable_multiprocessing: enable workers in separate process(es) send_one_request: send one request to engine before deleting + """ if current_platform.device_count() < tensor_parallel_size: pytest.skip(reason="Not enough CUDA devices") diff --git a/tests/v1/shutdown/test_startup_error.py b/tests/v1/shutdown/test_startup_error.py index 03982aa48e71..6d6d0d91e77f 100644 --- a/tests/v1/shutdown/test_startup_error.py +++ b/tests/v1/shutdown/test_startup_error.py @@ -23,7 +23,6 @@ def evil_method(self, *args, **kwargs): """Evil method that raises an exception.""" - if get_tensor_model_parallel_rank() == 0: raise Exception("Simulated Error in startup!") diff --git a/tests/v1/shutdown/utils.py b/tests/v1/shutdown/utils.py index 124254a41337..606928c5991b 100644 --- a/tests/v1/shutdown/utils.py +++ b/tests/v1/shutdown/utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shutdown test utils""" +"""Shutdown test utils.""" SHUTDOWN_TEST_TIMEOUT_SEC = 120 SHUTDOWN_TEST_THRESHOLD_BYTES = 2 * 2**30 diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index f21dae7e0bd8..5f5b63c93ecf 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -67,7 +67,8 @@ def _make_mixed_kv_cache_config( mamba_block_size: int = BLOCK_SIZE, ) -> KVCacheConfig: """Build a KVCacheConfig with a FullAttention group plus optional - SlidingWindow and Mamba groups for per-group metadata tests.""" + SlidingWindow and Mamba groups for per-group metadata tests. + """ register_all_kvcache_specs(vllm_config=None) groups = [] tensors = [] @@ -238,7 +239,8 @@ def test_events_disabled_no_events() -> None: # --------------------------------------------------------------------------- def test_block_stored_per_group_metadata_full_attention() -> None: """BlockStored carries medium, group_idx, kind, sliding_window, block_size, - and the correct block hash.""" + and the correct block hash. + """ fix = make_events_scheduler() sched = fix.scheduler assert sched.kv_event_medium == MEDIUM_CPU @@ -638,7 +640,8 @@ def test_mamba_block_size_not_scaled() -> None: def test_multi_group_fa_block_stored_group_idx() -> None: """Two distinct groups (FA + SW): each BlockStored carries its group_idx and kind. HybridKVCacheCoordinator requires >=2 distinct spec groups, so - we use FA (group 0) + SlidingWindow (group 1).""" + we use FA (group 0) + SlidingWindow (group 1). + """ kv_cfg = _make_mixed_kv_cache_config(num_blocks=16, sliding_window=32) fix = make_events_scheduler( num_cpu_blocks=8, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 901e076e0950..c87a5f7c956a 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -250,6 +250,7 @@ def make_scheduler_output( cached_req_new_blocks: For returning (cached) requests, maps req_id -> new_block_ids (incremental) or None. These are placed into ``scheduled_cached_reqs``. + """ scheduled_new_reqs: list[NewRequestData] = [] if new_reqs: @@ -1169,7 +1170,8 @@ def test_touched_blocks_survive_eviction() -> None: # --------------------------------------------------------------------------- def test_preemption_no_cpu_block_leak() -> None: """request_finished during in-flight load defers cleanup; - completes after load done.""" + completes after load done. + """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1264,7 +1266,8 @@ def test_eager_store_preemption_cleanup() -> None: # --------------------------------------------------------------------------- def test_inflight_finish_deferred_cleanup() -> None: """Store, then start a load, request_finished defers, - load completion fires cleanup.""" + load completion fires cleanup. + """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1398,7 +1401,8 @@ def test_chunked_prefill_reads_live_block_ids() -> None: """With chunked prefill, block IDs accumulate across scheduler steps. _prepare_eager_store_specs reads block IDs from scheduler_output via yield_req_data, so the store should reflect the updated (larger) block - list, not a stale snapshot.""" + list, not a stale snapshot. + """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1969,7 +1973,8 @@ def test_cp_eager_store_and_load_roundtrip( ) -> None: """With CP enabled, store blocks to CPU and reload them for a new request with matching tokens. Verifies that hash matching and transfer-pair - construction work with the virtual block size.""" + construction work with the virtual block size. + """ fix = _make_cp_scheduler( dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, @@ -2113,7 +2118,8 @@ def test_cp_effective_block_size_store_and_load( @pytest.mark.parametrize("cp_world_size", [1, 2, 4]) def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None: """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1 - because each virtual block covers more tokens.""" + because each virtual block covers more tokens. + """ kv_cache_config = _make_kv_cache_config(num_blocks=16) max_batched = 64 @@ -2805,7 +2811,8 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: """A sliding-window group nulls pages that left the window before the connector sees the finished request's block table, while the GPU keeps them hashed in its free queue. The eager store must copy those pages, or - the CPU lookup (which needs the whole window) misses every time.""" + the CPU lookup (which needs the whole window) misses every time. + """ fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False) sched = fix.scheduler gpu_pool = fix.gpu_block_pool diff --git a/tests/v1/spec_decode/test_acceptance_length.py b/tests/v1/spec_decode/test_acceptance_length.py index 4d093f6055ce..04db178bacfb 100644 --- a/tests/v1/spec_decode/test_acceptance_length.py +++ b/tests/v1/spec_decode/test_acceptance_length.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -EAGLE3 Acceptance Length Regression Tests. +"""EAGLE3 Acceptance Length Regression Tests. These tests verify that acceptance lengths for EAGLE3 speculative decoding do not regress across vLLM commits. Each test runs inference on the MT-Bench diff --git a/tests/v1/spec_decode/test_backup_token_async_spec.py b/tests/v1/spec_decode/test_backup_token_async_spec.py index 9340503ea4f2..3f99a5de0682 100644 --- a/tests/v1/spec_decode/test_backup_token_async_spec.py +++ b/tests/v1/spec_decode/test_backup_token_async_spec.py @@ -131,7 +131,8 @@ def test_various_spec_token_counts(self, num_spec_tokens: int): def test_buggy_code_was_always_off_by_one(self): """The original code used seq_len as index, which is always one past - the end of output_token_ids even without async inflation.""" + the end of output_token_ids even without async inflation. + """ req_ids = ["r0"] requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])} batch = _FakeInputBatch(req_ids, [5]) diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index a5cbd119bc66..8abbe91c1a35 100644 --- a/tests/v1/spec_decode/test_dynamic_sd_cug.py +++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py @@ -39,7 +39,6 @@ def _create_vllm_config_for_dsd( a schedule covering every query length in ``[1, max_decode_query_len]`` is generated. """ - max_decode_query_len = max_spec_tokens + 1 max_capture_tokens = max_num_seqs * max_decode_query_len @@ -91,7 +90,6 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch) shapes have been built, dispatch() should pick a FULL graph for every uniform decode batch shape produced by DSD up to max_num_seqs. """ - max_num_seqs = 512 max_spec_tokens = 7 max_decode_query_len = max_spec_tokens + 1 @@ -190,7 +188,6 @@ def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch): FULL candidates should be skipped in favor of the mixed-batch PIECEWISE graph under FULL_AND_PIECEWISE mode. """ - max_spec_tokens = 4 monkeypatch.setattr( @@ -236,7 +233,6 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch): by coincidence; replaying the FULL graph over them corrupts every request in the batch. See https://github.com/vllm-project/vllm/issues/49918. """ - max_spec_tokens = 7 decode_query_len = max_spec_tokens + 1 @@ -307,7 +303,6 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch): Uniform batches with smaller query lengths should therefore miss the FULL path entirely when using FULL_AND_PIECEWISE. """ - max_num_seqs = 512 max_spec_tokens = 7 max_decode_query_len = max_spec_tokens + 1 @@ -368,7 +363,6 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch): FULL graphs, while every other query length (e.g. the lower values 1, 2, 3) must fall back to the mixed-batch PIECEWISE graph. """ - max_num_seqs = 128 max_spec_tokens = 7 max_decode_query_len = max_spec_tokens + 1 diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index b62d7e90da42..9b3bfd73d06c 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -116,8 +116,7 @@ def _create_proposer( def test_prepare_next_token_ids(): - """ - Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded. + """Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded. Each will produce a device tensor of next_token_ids, taking as input either the GPU tensor of sampled_token_ids with -1 for rejected tokens, or the CPU python list[list[int]] with the rejected tokens removed. @@ -196,8 +195,7 @@ def test_prepare_next_token_ids(): def test_prepare_inputs(): - """ - cu_target_query_lens: [0, a, a + b, a + b + c] + """cu_target_query_lens: [0, a, a + b, a + b + c] num_rejected_tokens: [n1, n2, n3] num_tokens_per_req: [a - n1, b - n2, c - n3] cu_num_tokens: [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3] @@ -286,8 +284,7 @@ def test_prepare_inputs(): def test_prepare_inputs_padded(): - """ - Input scenario is 3 requests with num_speculative_tokens == 2 and: + """Input scenario is 3 requests with num_speculative_tokens == 2 and: - Request 1: query_len = 3, rejected = 1 - Request 2: query_len = 3, rejected = 0 - Request 3: query_len = 3, rejected = 2 @@ -297,7 +294,6 @@ def test_prepare_inputs_padded(): Reason: After accounting for rejections, these are the valid token positions from the original indices to sample from. """ - device = torch.device(DEVICE_TYPE) expected_token_indices_to_sample = torch.tensor( @@ -350,8 +346,7 @@ def test_prepare_inputs_padded(): def test_set_inputs_first_pass_default_eagle(): - """ - Test for set_inputs_first_pass without extra input slots (default EAGLE). + """Test for set_inputs_first_pass without extra input slots (default EAGLE). This tests the path where needs_extra_input_slots=False, which is the default EAGLE pathway. In this case: @@ -436,8 +431,7 @@ def test_set_inputs_first_pass_default_eagle(): def test_set_inputs_first_pass_draft_model(): - """ - Test for set_inputs_first_pass with a draft model (extra input slots, + """Test for set_inputs_first_pass with a draft model (extra input slots, no shift). This tests the path where needs_extra_input_slots=True and @@ -578,8 +572,7 @@ def test_set_inputs_first_pass_draft_model(): def test_set_inputs_first_pass_parallel_drafting(): - """ - Test for set_inputs_first_pass with parallel drafting (extra input slots, + """Test for set_inputs_first_pass with parallel drafting (extra input slots, with shift). This tests the path where needs_extra_input_slots=True and @@ -1107,8 +1100,7 @@ def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel): def test_set_inputs_first_pass_dflash(): - """ - Test for DFlash set_inputs_first_pass. + """Test for DFlash set_inputs_first_pass. DFlash uses cross-attention: context tokens become K/V and only query tokens (bonus + mask) are Q. This tests the DFlash-specific diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index 0cc78f3371c2..19166ee193f3 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -105,7 +105,8 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus the current draft-step offset. Padded entries are zeroed (matching the - main model runner's convention).""" + main model runner's convention). + """ fake = _make_fake_speculator() base = torch.tensor([100, 200, 300, 0], dtype=torch.int32) @@ -126,7 +127,8 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): """The padded tail of the upper-bound tensor is zeroed, so it stays a - valid tensor of length ``num_reqs_padded`` regardless of padding.""" + valid tensor of length ``num_reqs_padded`` regardless of padding. + """ fake = _make_fake_speculator() base = torch.zeros(2, dtype=torch.int32) @@ -141,7 +143,8 @@ def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): def test_build_draft_attn_metadata_clamps_to_max_model_len(): """The per-request upper bound (target bound + step) is clamped to the - model length so it never exceeds the allocated KV range.""" + model length so it never exceeds the allocated KV range. + """ fake = _make_fake_speculator(max_model_len=1024) base = torch.tensor([1023, 500], dtype=torch.int32) diff --git a/tests/v1/spec_decode/test_eagle_step_kernel.py b/tests/v1/spec_decode/test_eagle_step_kernel.py index 83a29bc0469c..4d6244eb77fe 100644 --- a/tests/v1/spec_decode/test_eagle_step_kernel.py +++ b/tests/v1/spec_decode/test_eagle_step_kernel.py @@ -132,7 +132,8 @@ def test_eagle_step_slot_mapping_kernel_exceeds_max(): def test_eagle_step_slot_mapping_kernel_cudagraph_padding(): """Test that padding threads write PADDING_SLOT_ID when - input_batch_size > batch_size (cudagraph padding).""" + input_batch_size > batch_size (cudagraph padding). + """ device = torch.device(DEVICE_TYPE) batch_size = 4 input_batch_size = 8 diff --git a/tests/v1/spec_decode/test_extract_hidden_states.py b/tests/v1/spec_decode/test_extract_hidden_states.py index 6b4e53ced679..e330d9588072 100644 --- a/tests/v1/spec_decode/test_extract_hidden_states.py +++ b/tests/v1/spec_decode/test_extract_hidden_states.py @@ -130,8 +130,7 @@ def test_proposer_initialization_missing_layer_ids(): def test_prepare_next_token_ids_padded(): - """ - Test for prepare_next_token_ids_padded with extract_hidden_states. + """Test for prepare_next_token_ids_padded with extract_hidden_states. Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1). For each request we either use the sampled token (if valid and not discarded) @@ -196,8 +195,7 @@ def test_prepare_next_token_ids_padded(): def test_propose(): - """ - Test the propose() method of ExtractHiddenStatesProposer. + """Test the propose() method of ExtractHiddenStatesProposer. This should: 1. Accept target hidden states and sampled token IDs diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py index 5510c36e60d2..622d4250655a 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer.py +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -69,7 +69,8 @@ def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace: def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): """The proposer's slot-mapping math runs against the kernel-granularity - block table, so block_size must come from kernel_block_sizes.""" + block table, so block_size must come from kernel_block_sizes. + """ layer_names = {"draft.0.self_attn.attn"} proposer = _make_proposer(monkeypatch, layer_names) @@ -97,7 +98,8 @@ def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch) def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch): """_draft_attn_layer_names is a set; the attention groups built from it - must not depend on its (process-random) iteration order.""" + must not depend on its (process-random) iteration order. + """ layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"} expected_order = sorted(layer_names) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 81f842419375..3644f0528fe0 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -93,7 +93,8 @@ def test_eagle_max_len( @pytest.mark.parametrize("spec_max_model_len", [80, 150]) def test_mtp_speculative_config_max_model_len(spec_max_model_len: int): """Regression test for #41456: max_model_len in speculative config - should be respected for the draft model.""" + should be respected for the draft model. + """ model_config = ModelConfig( model="XiaomiMiMo/MiMo-7B-Base", runner="generate", diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index e334371f6d8a..8bc791e9634d 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -66,7 +66,6 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: @mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_model") def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_group): """Test MTP-specific model loading with unified model approach.""" - # Setup mocks mock_model = mock.MagicMock() mock_model.model.embed_tokens.weight.shape = (131072, 4096) @@ -117,8 +116,7 @@ class _TargetModelStub(LlamaForCausalLM): @pytest.mark.parametrize("num_speculative_tokens", [1]) def test_mtp_propose(num_speculative_tokens, monkeypatch): - """Test that MTP's forward method returns hidden states directly""" - + """Test that MTP's forward method returns hidden states directly.""" device = torch.device(DEVICE_TYPE) batch_size = 2 seq_lens = [5, 3] diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index a64b0eb6e940..9e0052bcc59c 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -93,8 +93,7 @@ def _assert_distribution_match( label: str = "", min_expected: float = 5.0, ): - """ - Assert sampled tokens match the target distribution via a + """Assert sampled tokens match the target distribution via a chi-squared goodness-of-fit test. This is done by computing observed vs expected token counts (target_probs * num_samples), then checking that the chi-squared statistic is below a conservative @@ -242,8 +241,7 @@ def test_watermarked_recovery_supports_smaller_draft_vocabulary(): def test_stochastic_rejection_sample( num_speculative_steps: int, temperature: float, draft_logits_dtype: torch.dtype ): - """ - Verify that rejection sampling produces the target distribution. + """Verify that rejection sampling produces the target distribution. This is done by simulating many independent trials of speculative decoding (from a fixed target and draft distribution). We then run rejection sample on all of the trials (requests), and verify @@ -253,7 +251,6 @@ def test_stochastic_rejection_sample( Parametrized over the draft-logits dtype: storing them in the draft head's dtype must not bias the output distribution. """ - torch.manual_seed(42) device = "cuda" num_trials = 10 * VOCAB_SIZE @@ -383,11 +380,9 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int) @pytest.mark.parametrize("num_speculative_steps", [1, 3]) def test_greedy_rejection_sample(num_speculative_steps: int): - """ - Verify that greedy (temperature=0) always outputs the target argmax + """Verify that greedy (temperature=0) always outputs the target argmax at every accepted position. """ - torch.manual_seed(42) device = "cuda" num_trials = 10 * VOCAB_SIZE @@ -435,8 +430,7 @@ def test_synthetic_rejection_sample( temperature: float, unconditional_rates: list[float], ): - """ - Verify that synthetic rejection sampling produces the expected + """Verify that synthetic rejection sampling produces the expected per-position acceptance rates. The unconditional rate at position i is P(all draft steps 0..i accepted) = product(conditional_rates[0:i+1]). This is approximately mean(num accepted >= i + 1) over many trials. @@ -689,14 +683,12 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool): def test_block_verification_rejection_sample( num_speculative_steps: int, temperature: float, has_draft_logits: bool ): - """ - Verify that block verification (Sun et al.) preserves the target + """Verify that block verification (Sun et al.) preserves the target distribution at every accepted position, for both the full draft-logits case and the one-hot (no draft logits) case. Block verification changes *which* prefix is accepted, but the marginal of every output position must still match the target distribution p(x). """ - torch.manual_seed(42) device = "cuda" num_trials = 10 * VOCAB_SIZE @@ -733,12 +725,10 @@ def test_block_verification_rejection_sample( @pytest.mark.parametrize("num_speculative_steps", [3, 5]) def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int): - """ - Block verification is designed to accept at least as long a prefix as + """Block verification is designed to accept at least as long a prefix as token verification in expectation. Verify the mean accepted length is no worse than the standard method on the same inputs. """ - torch.manual_seed(0) device = "cuda" num_trials = 20 * VOCAB_SIZE diff --git a/tests/v1/spec_decode/test_speculators_correctness.py b/tests/v1/spec_decode/test_speculators_correctness.py index 5a92fe00f3ff..4cc66d5cd3c2 100644 --- a/tests/v1/spec_decode/test_speculators_correctness.py +++ b/tests/v1/spec_decode/test_speculators_correctness.py @@ -134,8 +134,7 @@ def print_spec_decode_stats(stats: dict) -> None: @pytest.mark.parametrize("config", SPECULATOR_CONFIGS) def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config): - """ - Test speculators model properly initializes speculative decoding. + """Test speculators model properly initializes speculative decoding. Verifies: 1. Speculative config is automatically initialized from speculators config @@ -185,8 +184,7 @@ def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config): @large_gpu_mark(min_gb=40) @pytest.mark.parametrize("config", SPECULATOR_CONFIGS) def test_speculators_correctness(monkeypatch, config): - """ - E2E correctness test via the speculators auto-detect path. + """E2E correctness test via the speculators auto-detect path. Evaluates GSM8k accuracy to ensure the speculators-format model produces correct outputs, and checks that acceptance length does not collapse under diff --git a/tests/v1/spec_decode/test_speculators_eagle3.py b/tests/v1/spec_decode/test_speculators_eagle3.py index 9a252cfffc8f..66736cce87b8 100644 --- a/tests/v1/spec_decode/test_speculators_eagle3.py +++ b/tests/v1/spec_decode/test_speculators_eagle3.py @@ -32,8 +32,7 @@ def test_eagle3_speculators_model( vllm_runner, example_prompts, model_path, monkeypatch ): - """ - Test Eagle3 speculators models properly initialize speculative decoding. + """Test Eagle3 speculators models properly initialize speculative decoding. This test verifies: 1. Eagle3 support is detected for the model diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index a8f7221da81b..9457c66050db 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -23,7 +23,6 @@ @pytest.fixture def mock_model_runner_with_input_batch(): """Create a mock GPUModelRunner with a real InputBatch for e2e testing.""" - runner = Mock(spec=GPUModelRunner) runner.uses_mrope = False runner.requests = {} diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index 31c1a6f4846a..c5f8090f943f 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py @@ -23,7 +23,6 @@ @pytest.fixture def mock_model_runner_with_req_states(): """Create a mock MRv2 GPUModelRunner with a real RequestState.""" - runner = Mock(spec=GPUModelRunner) runner.req_states = RequestState( max_num_reqs=10, diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 822b4e49da34..af858eaa8c04 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -315,8 +315,7 @@ def test_update_request_as_session_with_output_tokens(self): assert num_new_tokens == 2 def test_streaming_e2e_lifecycle(self): - """ - Comprehensive integration test covering complete streaming request lifecycle + """Comprehensive integration test covering complete streaming request lifecycle including scheduler state management and aliasing bug prevention. FULL LIFECYCLE: diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index a23189438475..ea4429270228 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -29,7 +29,8 @@ def __init__(self, is_diffusion: bool): def test_structured_outputs_rejected_for_diffusion_models(): """Diffusion LLMs denoise the canvas in parallel, which is incompatible with the token-by-token grammar FSM. The request must fail with a clear - validation error instead of an FSM rejection mid-generation (#45436).""" + validation error instead of an FSM rejection mid-generation (#45436). + """ params = SamplingParams( structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) ) @@ -62,7 +63,8 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match): """json_object=False and an empty json schema pass the `is not None` exclusivity check but resolve to no structured-output key, so they must be rejected at request validation (-> 400) instead of reaching and crashing - the engine.""" + the engine. + """ params = SamplingParams(structured_outputs=structured_outputs) with pytest.raises(VLLMValidationError, match=match): params._validate_structured_outputs( @@ -85,7 +87,8 @@ def test_regex_with_nul_byte_rejected(regex): handled by xgrammar's native regex converter. It must be rejected at request validation in every backend mode (a clean 400), instead of reaching that native code or silently falling back to another backend in the default - 'auto' mode.""" + 'auto' mode. + """ params = SamplingParams(structured_outputs=StructuredOutputsParams(regex=regex)) # Rejected before backend selection, so it is a 400 even in 'auto' mode @@ -124,7 +127,8 @@ def test_regex_with_nul_byte_rejected(regex): ) def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): """Only `VLLMClientError` survives `AsyncLLM.generate` untouched; anything else - is wrapped in `EngineGenerateError` and served as a 500 instead of a 400.""" + is wrapped in `EngineGenerateError` and served as a 500 instead of a 400. + """ params = SamplingParams(structured_outputs=structured_outputs) with pytest.raises(VLLMClientError): params._validate_structured_outputs( diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py index 6da14b4e426b..716921cdf456 100644 --- a/tests/v1/test_kv_cache_spec_registry.py +++ b/tests/v1/test_kv_cache_spec_registry.py @@ -70,8 +70,7 @@ def restore_kv_cache_spec_registry(): @dataclass(frozen=True) class _TrulyUnregisteredSpec(KVCacheSpec): - """ - A spec that inherits directly from KVCacheSpec with no registered + """A spec that inherits directly from KVCacheSpec with no registered ancestor in the MRO. Used to test that the registry correctly raises when no entry can be found. """ @@ -212,7 +211,6 @@ class _CustomSpec(spec_cls): # type: ignore[valid-type,misc] def test_custom_spec_register_requires_manager(self): """Invalid register decorator arguments fail early.""" - with pytest.raises(AssertionError, match="manager_class is required"): @register_kv_cache_spec( @@ -223,8 +221,7 @@ class _CustomFullSpecWithoutManager(FullAttentionSpec): custom_param: int = 16 def test_unregistered_spec_no_registered_parent_raises(self): - """ - A spec whose entire MRO contains no registered class resolves to None. + """A spec whose entire MRO contains no registered class resolves to None. Runtime callers should use check_kv_cache_spec_registry to fail early. Subclasses of registered specs intentionally do not fail — they inherit their parent's manager via MRO walking. @@ -243,8 +240,7 @@ def test_unregistered_spec_no_registered_parent_raises(self): UniformTypeKVCacheSpecs.is_uniform_type({"layer_0": spec}) def test_unregistered_subclass_inherits_parent_manager(self): - """ - An unregistered subclass of a registered spec resolves via MRO + """An unregistered subclass of a registered spec resolves via MRO to its parent's manager — this is intentional registry behaviour. """ @@ -350,8 +346,7 @@ def test_homogeneous_group_keeps_member_kind(self, spec_cls, expected_kind): @pytest.mark.parametrize("other_cls", [MLAAttentionSpec, SinkFullAttentionSpec]) def test_mixed_kinds_sharing_full_attention_base(self, other_cls): - """ - Members with differing kinds but a single registered + """Members with differing kinds but a single registered uniform_type_base_spec of FullAttentionSpec report FULL_ATTENTION. This is the shape produced by a model that pairs a main K/V cache with a @@ -400,8 +395,7 @@ def test_mixed_base_specs_report_unknown(self): assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN def test_uniform_non_full_attention_base_reports_unknown(self): - """ - The FULL_ATTENTION result is gated on the shared base spec being + """The FULL_ATTENTION result is gated on the shared base spec being FullAttentionSpec; a uniform base of any other type stays UNKNOWN. """ diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 50abb221ad82..b5f180710e03 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -46,7 +46,8 @@ def test_logprobs_tensors_cat(): def test_logprobs_tensors_tolists_with_tensor_boundaries(): """Adaptive verification hands over the request boundaries as a tensor (they only exist on device); tolists() must materialize it as a plain - list so slice_request splits requests correctly.""" + list so slice_request splits requests correctly. + """ tensors = LogprobsTensors( torch.tensor([[1, 2], [3, 4], [5, 6]]), torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), @@ -85,7 +86,8 @@ def test_sampling_mask_lists_to_nested_list(): def test_sampling_mask_tensors_match_finite_support(max_num_kept): """Whatever its size (empty, one, around the compact width, beyond the cap, the whole vocab), a sampled row's mask is exactly the ascending set - of its finite logits; rows that sampled nothing are empty.""" + of its finite logits; rows that sampled nothing are empty. + """ from vllm.v1.worker.gpu.sample.output import MAX_COMPACT_SUPPORT vocab_size = 20_001 @@ -112,7 +114,8 @@ def test_sampling_mask_tensors_match_finite_support(max_num_kept): def test_sampling_mask_matches_processed_top_k_top_p_support(): """The mask must exactly mirror whatever support `apply_top_k_top_p` (the real logits-processing function used by the sampler) actually - produces, whatever backend implements it.""" + produces, whatever backend implements it. + """ processed_logits = apply_top_k_top_p( logits=torch.tensor( [[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE @@ -138,7 +141,8 @@ def test_sampling_mask_matches_processed_top_k_top_p_support(): def test_sampling_mask_preserves_top_k_boundary_ties(): """When the kept support is wider than `max_num_kept` (e.g. a top-k boundary tie keeps more than k logits), the mask must fall back to the - exact bitmask instead of silently truncating to `max_num_kept` ids.""" + exact bitmask instead of silently truncating to `max_num_kept` ids. + """ processed_logits = torch.tensor( [[6.0, 5.0, 4.0, 4.0, 4.0, float("-inf"), float("-inf"), float("-inf")]], device=DEVICE_TYPE, @@ -185,7 +189,7 @@ def setUp(self): ) def test_slice_without_cu_num_generated_tokens(self): - """Test slicing without cu_num_generated_tokens""" + """Test slicing without cu_num_generated_tokens.""" logprobsLists = LogprobsLists( logprob_token_ids=[[1], [2], [3]], logprobs=[[0.1], [0.2], [0.3]], @@ -200,7 +204,7 @@ def test_slice_without_cu_num_generated_tokens(self): assert sliced.cu_num_generated_tokens is None def test_slice_from_start(self): - """Test slicing from the start position""" + """Test slicing from the start position.""" sliced = self.logprobsLists.slice_request(0, num_positions=5) assert len(sliced.logprob_token_ids) == 5 assert sliced.logprob_token_ids == [ @@ -213,7 +217,7 @@ def test_slice_from_start(self): assert sliced.cu_num_generated_tokens is None def test_slice_from_middle(self): - """Test slicing from the middle position""" + """Test slicing from the middle position.""" sliced = self.logprobsLists.slice_request(1, num_positions=7) assert len(sliced.logprob_token_ids) == 7 assert sliced.logprob_token_ids == [ @@ -228,21 +232,21 @@ def test_slice_from_middle(self): assert sliced.cu_num_generated_tokens is None def test_slice_single_request(self): - """Test slicing a single request""" + """Test slicing a single request.""" sliced = self.logprobsLists.slice_request(1, num_positions=3) assert len(sliced.logprob_token_ids) == 3 assert sliced.logprob_token_ids == [[5, 6], [7, 8], [9, 10]] assert sliced.cu_num_generated_tokens is None def test_slice_last_request(self): - """Test slicing the last request""" + """Test slicing the last request.""" sliced = self.logprobsLists.slice_request(2, num_positions=4) assert len(sliced.logprob_token_ids) == 4 assert sliced.logprob_token_ids == [[11, 12], [13, 14], [15, 16], [17, 18]] assert sliced.cu_num_generated_tokens is None def test_slice_all_requests(self): - """Test slicing all requests (full slice)""" + """Test slicing all requests (full slice).""" sliced = self.logprobsLists.slice_request(0, num_positions=9) assert len(sliced.logprob_token_ids) == 9 # All tokens assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 9f7761c22b5d..8a1af4e468ed 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -44,7 +44,6 @@ class MyType: def test_encode_decode(monkeypatch: pytest.MonkeyPatch): """Test encode/decode loop with zero-copy tensors.""" - with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") @@ -251,7 +250,6 @@ def test_custom_class_serialization_allowed_with_pickle( monkeypatch: pytest.MonkeyPatch, ): """Test that serializing a custom class succeeds when allow_pickle=True.""" - with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") encoder = MsgpackEncoder() @@ -429,7 +427,8 @@ def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends it: many requests, each with per-token tensors small enough that pyzmq copies their frames, while the accumulated payload frame is large enough - that pyzmq sends it zero-copy.""" + that pyzmq sends it zero-copy. + """ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.outputs import LogprobsTensors diff --git a/tests/v1/utils.py b/tests/v1/utils.py index 993ad8a947d0..c55ace9d572c 100644 --- a/tests/v1/utils.py +++ b/tests/v1/utils.py @@ -17,6 +17,7 @@ def get_prometheus_metrics(server: RemoteOpenAIServer) -> dict[str, dict[str, fl For example: {"vllm:request_success": { "engine=0": 5.0, "engine=1": 3.0} } + """ try: response = requests.get(server.url_for("metrics"), timeout=10) @@ -72,6 +73,7 @@ def get_engine_request_counts(metrics: dict[str, dict[str, float]]) -> dict[str, Returns: Dict mapping engine indices to request counts. For example: {"0": 15.0, "1": 12.0} + """ engine_counts = {} @@ -97,6 +99,7 @@ def check_request_balancing(server: RemoteOpenAIServer, dp_size: int): Args: server: The RemoteOpenAIServer instance dp_size: Number of data parallel ranks + """ if dp_size <= 1: return diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 4cc8f9ffe2fd..299ef7695f45 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -101,7 +101,8 @@ def test_packed_dsv4_zeroer_zeroes_only_each_layers_page(): def test_overlaid_zeroer_dedups_segments_with_max_span(): """Two groups overlay one allocation; the zeroer must emit one segment per distinct byte offset, spanning the widest overlaid page, so a newly allocated block is fully - zeroed no matter which group owns it.""" + zeroed no matter which group owns it. + """ from unittest.mock import MagicMock from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index f8bd50c6d9c4..28df60f49d67 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -97,7 +97,8 @@ def test_draft_lookahead_uses_boundary_feature_when_cached(): """The drafter's +1 look-ahead can reach the feature at offset == processed_end (the next chunk). When its encoder output is already cached (the scheduler encoded it ahead), it is used for the look-ahead position - rather than ignored.""" + rather than ignored. + """ f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end runner = _make_runner([f0, f1], cached=[f0, f1]) @@ -114,7 +115,8 @@ def test_draft_lookahead_uses_boundary_feature_when_cached(): def test_draft_lookahead_tolerates_missing_boundary_feature(): """When the +1 look-ahead feature past the processed boundary is not yet encoded, fall back to the token embedding (the draft token is verified by - the target) instead of raising.""" + the target) instead of raising. + """ f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached runner = _make_runner([f0, f1], cached=[f0]) @@ -131,7 +133,8 @@ def test_draft_lookahead_tolerates_missing_boundary_feature(): def test_draft_lookahead_raises_on_interior_miss(): """A miss for a feature within the processed range (not the look-ahead boundary) is a real invariant violation and must fail loudly, even on the - drafter path.""" + drafter path. + """ f0 = _feature("h0", offset=0, length=8) # interior, within processed range runner = _make_runner([f0], cached=[]) @@ -141,7 +144,8 @@ def test_draft_lookahead_raises_on_interior_miss(): def test_target_path_raises_on_encoder_cache_miss(): """On the target path (no look-ahead) a miss is a real invariant - violation and must fail loudly.""" + violation and must fail loudly. + """ f0 = _feature("h0", offset=0, length=8) runner = _make_runner([f0], cached=[]) @@ -152,7 +156,8 @@ def test_target_path_raises_on_encoder_cache_miss(): @pytest.mark.parametrize("draft_lookahead", [0, 1]) def test_multi_request_batch_gathers_per_request(draft_lookahead): """Two prefilling requests in one batch: per-request query bounds must be - indexed by request, not applied as whole arrays.""" + indexed by request, not applied as whole arrays. + """ a0 = _feature("a0", offset=0, length=8) b0 = _feature("b0", offset=0, length=8) cache = EncoderCache() diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py index 150c4ed00711..6081163c4b51 100644 --- a/tests/v1/worker/test_gpu_bad_words.py +++ b/tests/v1/worker/test_gpu_bad_words.py @@ -73,7 +73,8 @@ def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor: def test_v2_bad_words_prefix_inside_draft_tokens(): """A prefix matching entirely within the draft tokens must mask the bad - word's last token at the draft position that completes the prefix.""" + word's last token at the draft position that completes the prefix. + """ out = _apply([[12, 13, 40]]) expected = torch.zeros_like(out) expected[2, 40] = -float("inf") @@ -82,7 +83,8 @@ def test_v2_bad_words_prefix_inside_draft_tokens(): def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): """A prefix spanning the committed/draft boundary must mask at the row - where the prefix completes, not one draft position later.""" + where the prefix completes, not one draft position later. + """ out = _apply([[11, 12, 30]]) expected = torch.zeros_like(out) expected[1, 30] = -float("inf") @@ -91,7 +93,8 @@ def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): def test_v2_bad_words_no_spurious_match_from_last_committed_token(): """The last committed token must not be double-counted as the first draft - token; [11, 11] never occurs in output [10, 11] + drafts [12, 13].""" + token; [11, 11] never occurs in output [10, 11] + drafts [12, 13]. + """ out = _apply([[11, 11, 50]]) expected = torch.zeros_like(out) torch.testing.assert_close(out, expected) diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 1309686d53ce..9f0b6da218af 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -43,7 +43,8 @@ def _make_batch( slot_pool: np.ndarray | None = None, ) -> InputBatch: """Consistent batch with random unique request slots, 1 + k_i logits - rows per request at the tail of each query segment.""" + rows per request at the tail of each query segment. + """ if slot_pool is None: slot_pool = np.arange(max_num_reqs) slots = rng.choice(slot_pool, size=num_reqs, replace=False) @@ -141,7 +142,8 @@ def _owned_rows(cu: np.ndarray, owned: np.ndarray) -> np.ndarray: def test_local_batch_partition(tp_size: int, seed: int): """Local sub-batches partition the batch by owner, preserve batch order within an owner, and carry exactly the owned requests' entries of every - per-request and per-row array.""" + per-request and per-row array. + """ rng = np.random.default_rng(seed) max_num_reqs = 64 batch = _make_batch(rng, num_reqs=23, max_num_reqs=max_num_reqs, max_spec=3) @@ -223,7 +225,8 @@ def test_local_batch_partition(tp_size: int, seed: int): def test_all_to_all_reassembly_simulated(tp_size: int): """Vocab shards computed from owner-sorted hidden states and routed through the simulated all-to-all reassemble into exactly the full-vocab - logits of each rank's owned rows.""" + logits of each rank's owned rows. + """ rng = np.random.default_rng(7) max_num_reqs = 32 batch = _make_batch(rng, num_reqs=11, max_num_reqs=max_num_reqs, max_spec=4) @@ -261,7 +264,8 @@ def test_all_to_all_reassembly_simulated(tp_size: int): @requires_cuda def test_shard_no_spec(): """Batches without draft tokens (one logits row per request) degenerate - cleanly: rows mirror requests.""" + cleanly: rows mirror requests. + """ rng = np.random.default_rng(11) tp_size = 4 max_num_reqs = 32 @@ -280,7 +284,8 @@ def test_shard_no_spec(): @pytest.mark.parametrize("tp_size", [2, 4]) def test_shard_grammar_output(tp_size: int): """The grammar filter keeps exactly the owned requests' bitmask rows, in - grammar order, and the per-rank pieces partition the global bitmask.""" + grammar order, and the per-rank pieces partition the global bitmask. + """ rng = np.random.default_rng(5) max_num_reqs = 32 num_reqs = 12 @@ -335,7 +340,8 @@ def test_shard_grammar_output(tp_size: int): @pytest.mark.parametrize("tp_size", [2, 4]) def test_gather_src_indices_round_trip(tp_size: int): """Padded per-rank result blocks indexed by gathered_src_indices restore - batch order, including when some rank owns zero requests.""" + batch order, including when some rank owns zero requests. + """ rng = np.random.default_rng(3) max_num_reqs = 40 all_slots = np.arange(max_num_reqs) @@ -370,7 +376,8 @@ def test_gather_src_indices_round_trip(tp_size: int): def test_gather_sampler_output_kernels(tp_size: int): """The fused pack/unpack kernels round-trip per-rank sampler results into batch order: token ids zero-filled past the local width, counts - cast back to int32, identical on every rank.""" + cast back to int32, identical on every rank. + """ from dataclasses import replace as dc_replace from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -465,7 +472,8 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: def test_gather_sampler_output_logprobs_and_nans(): """num_nans is reduced per request and gathered; LogprobsTensors are packed into fixed per-request blocks, width-aligned across ranks (pad - and truncate paths), and restored to global row order.""" + and truncate paths), and restored to global row order. + """ from dataclasses import replace as dc_replace from vllm.v1.outputs import LogprobsTensors @@ -598,7 +606,8 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: def test_finish_requests_frees_slots_in_sorted_order(): """Request slots must be freed in the same order on every TP rank: ownership derives from slot indices, and `finished_req_ids` is a set - whose iteration order is per-process hash-randomized.""" + whose iteration order is per-process hash-randomized. + """ from vllm.v1.worker.gpu.model_runner import GPUModelRunner removed: list[str] = [] diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index e77156eef901..f44edabea4cc 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -217,7 +217,8 @@ def test_v1_block_table_move_row_clears_vacated_row(): not keep stale block ids. Padded dummy-run batches dereference stale rows as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode padding) and write state in place there — corrupting the blocks' new - owner once they are reallocated, e.g. to an in-flight NIXL load.""" + owner once they are reallocated, e.g. to an in-flight NIXL load. + """ from vllm.v1.worker.block_table import BlockTable block_table = BlockTable( @@ -249,7 +250,8 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): PAD-filled, state indices are not), so stale rows would direct dummy state writes at freed — possibly reallocated — blocks. get_dummy_block_tables must hand out zeroed (null block) rows while - preserving the persistent storage address for CUDA graphs.""" + preserving the persistent storage address for CUDA graphs. + """ device = torch.device("cuda") block_tables = BlockTables( block_sizes=[16], diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index a66a1395d2c7..a4b3d71698a9 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/tests/v1/worker/test_gpu_gumbel_sample.py @@ -331,7 +331,7 @@ def test_drafting_uses_a_separate_noise_stream(): def test_greedy_temperature_zero_returns_argmax(): - """temperature == 0 skips Gumbel noise and returns the exact argmax.""" + """Temperature == 0 skips Gumbel noise and returns the exact argmax.""" torch.manual_seed(0) num_reqs = 128 logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32) diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index 3ef4e0dabe01..2fe3567739aa 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -66,11 +66,9 @@ def _compare_objs(obj1, obj2, skip: Sequence = ("logitsprocs", "batch_update_bui def _remove_requests( input_batch: InputBatch, batch_size: int, reqs: list[CachedRequestState] ) -> set[str]: - """ - Remove some requests randomly from the batch and returns + """Remove some requests randomly from the batch and returns set of request removed """ - num_reqs_to_remove = np.random.randint(0, batch_size) req_indices_to_remove: set[int] = set() for _ in range(num_reqs_to_remove): @@ -90,8 +88,7 @@ def _construct_expected_sampling_metadata( req_id_index_in_input_batch: dict[str, int], device: torch.device, ) -> SamplingMetadata: - """ - Constructs and returns the expected SamplingMetadata for this + """Constructs and returns the expected SamplingMetadata for this batch. """ num_reqs = len(req_ids_retained) @@ -219,8 +216,7 @@ def _construct_cached_request_state(req_id_suffix: int): @pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32, 64]) def test_sampling_metadata_in_input_batch(device: str, batch_size: int): - """ - Tests the logic for managing sampling metadata in the InputBatch. + """Tests the logic for managing sampling metadata in the InputBatch. This test involves adding a set of requests to the InputBatch, followed by removing a subset of them. Afterward, the batch is compacted, @@ -314,8 +310,7 @@ def same(t1: torch.Tensor | None, t2: torch.Tensor | None) -> bool: @pytest.mark.parametrize("batch_size", [32]) @pytest.mark.parametrize("swap_list", [((0, 1),)]) def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list): - """ - Tests the logic for managing sampling metadata in the InputBatch. + """Tests the logic for managing sampling metadata in the InputBatch. This test involves adding a set of requests to the InputBatch, followed by removing a subset of them. Afterward, the batch is compacted, diff --git a/tests/v1/worker/test_gpu_logit_bias.py b/tests/v1/worker/test_gpu_logit_bias.py index 3bf62380bbc6..8e7148575b7b 100644 --- a/tests/v1/worker/test_gpu_logit_bias.py +++ b/tests/v1/worker/test_gpu_logit_bias.py @@ -85,7 +85,8 @@ def test_v2_min_tokens_keeps_masking_when_other_candidates_remain(): def test_v2_min_tokens_mixed_batch_gates_restore_per_request(): """CHECK_ALL_MASKED_ROWS is a batch-wide constexpr, so a mixed batch must - still consult the per-request flag.""" + still consult the per-request flag. + """ out = _apply(_only_stop_token_left(2), structured=[True, False]) assert out[0, STOP_TOKEN] == 1.0 diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index bb77f90ced2a..2dae08688995 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -73,16 +73,15 @@ def _restore_default_dtype(): """Several tests here set the process-wide default dtype to float16 and previously leaked it, corrupting later float-sensitive tests in the same - pytest process (torch.randn silently produced fp16).""" + pytest process (torch.randn silently produced fp16). + """ old = torch.get_default_dtype() yield torch.set_default_dtype(old) def initialize_kv_cache(runner: GPUModelRunner): - """ - Only perform necessary steps in GPUModelRunner.initialize_kv_cache() - """ + """Only perform necessary steps in GPUModelRunner.initialize_kv_cache().""" attn_spec = FullAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config), @@ -1193,14 +1192,12 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config): reason="Attention backend FLASHINFER is only supported on CUDA.", ) def test_hybrid_attention_mamba_tensor_shapes(): - """ - The GPU model runner creates different views into the + """The GPU model runner creates different views into the KVCacheTensors for the attention and mamba layers (via _allocate_kv_caches). This test verifies that the views are compatible: writing a mamba block will not corrupt an attention block and vice versa """ - set_random_seed(42) update_environment_variables( @@ -1457,7 +1454,8 @@ class StopInitialization(Exception): def test_hybrid_block_table_initialization(): """Test hybrid block table with different kernel and kvcache_manager block - sizes.""" + sizes. + """ from vllm.v1.worker.block_table import BlockTable # Test configuration: kvcache_manager block size = 32, diff --git a/tests/v1/worker/test_gpu_model_runner_mm_gather.py b/tests/v1/worker/test_gpu_model_runner_mm_gather.py index a86e77c8f52f..dcf754fcae28 100644 --- a/tests/v1/worker/test_gpu_model_runner_mm_gather.py +++ b/tests/v1/worker/test_gpu_model_runner_mm_gather.py @@ -62,7 +62,8 @@ def _gather(features, cached, *, num_scheduled, shift, num_computed=0): def test_draft_shift_uses_boundary_feature_when_cached(): """The drafter's +1 look-ahead reaches the feature at offset == processed_end; when it is already cached it is used for the look-ahead - position rather than ignored.""" + position rather than ignored. + """ f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end mm_embeds, is_mm_embed = _gather([f0, f1], [f0, f1], num_scheduled=8, shift=1) @@ -75,7 +76,8 @@ def test_draft_shift_uses_boundary_feature_when_cached(): def test_draft_shift_tolerates_missing_boundary_feature(): """When the +1 look-ahead feature past the processed boundary is not yet - encoded, fall back to the token embedding instead of raising.""" + encoded, fall back to the token embedding instead of raising. + """ f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached mm_embeds, is_mm_embed = _gather([f0, f1], [f0], num_scheduled=8, shift=1) @@ -87,7 +89,8 @@ def test_draft_shift_tolerates_missing_boundary_feature(): def test_draft_shift_raises_on_interior_miss(): """A miss for a feature within the processed range (not the look-ahead - boundary) is a real invariant violation, even on the drafter path.""" + boundary) is a real invariant violation, even on the drafter path. + """ f0 = _feature("h0", offset=0, length=8) # interior, within processed range with pytest.raises(RuntimeError, match="Encoder cache miss"): _gather([f0], [], num_scheduled=8, shift=1) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 86ff1a074538..891778678b9c 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -126,7 +126,6 @@ def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table( expected: int, ): """Mamba/GDN block-table rows index global positions, unlike DCP KV.""" - max_model_len = 1_048_576 attention_block_size = 1_536 mamba_block_size = 16 @@ -220,7 +219,8 @@ def stage_write(self, *_args): def _make_capture_runner(captured: bool) -> GPUModelRunner: """Minimal V2 runner for capture_model: fakes everything except the - cudagraph_manager's needs_capture decision.""" + cudagraph_manager's needs_capture decision. + """ runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_state = SimpleNamespace(supports_mm_inputs=False) runner.cudagraph_manager = SimpleNamespace( @@ -246,7 +246,8 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner: def test_capture_model_locks_workspace_after_capture(monkeypatch): """A workspace resize after capture frees the buffer the captured graphs baked in, so capture_model must lock the workspace before returning - (https://github.com/vllm-project/vllm/issues/55336).""" + (https://github.com/vllm-project/vllm/issues/55336). + """ runner = _make_capture_runner(captured=True) monkeypatch.setattr( model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext @@ -267,7 +268,8 @@ def test_capture_model_locks_workspace_after_capture(monkeypatch): def test_capture_model_skips_lock_when_nothing_captured(monkeypatch): """With no graphs to capture (e.g. enforce_eager) there is nothing baked - into the workspace, so the early return must not lock it.""" + into the workspace, so the early return must not lock it. + """ runner = _make_capture_runner(captured=False) lock_calls = [] monkeypatch.setattr( @@ -281,7 +283,8 @@ def test_capture_model_skips_lock_when_nothing_captured(monkeypatch): def test_capture_model_profile_only_skips_lock(monkeypatch): """The memory-profiling capture pass runs before kernel warmup and the real capture; locking there would stop the warmup from growing the - workspace to its scheduler-realistic size.""" + workspace to its scheduler-realistic size. + """ runner = _make_capture_runner(captured=True) monkeypatch.setattr( model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index 74191206d363..0230ff9b83c2 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py @@ -82,7 +82,8 @@ def _capture_model(*, profile_only: bool = False) -> int: class _FakePlatform: """Stands in for current_platform; the global graph pool is a class attr, - matching vllm.platforms.Platform's lazy singleton.""" + matching vllm.platforms.Platform's lazy singleton. + """ _global_graph_pool: Any = GLOBAL_POOL diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index f6263b68da0a..fce0f930771c 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -24,9 +24,7 @@ class ConcreteWorkerProfiler(WorkerProfiler): - """ - A basic implementation of a worker profiler for testing purposes. - """ + """A basic implementation of a worker profiler for testing purposes.""" def __init__(self, profiler_config: ProfilerConfig): self.start_call_count = 0 @@ -122,7 +120,8 @@ def test_max_iterations(default_profiler_config): def test_restart_after_max_iterations(default_profiler_config): """A start_profile after an auto-stop must actually restart, not be silently ignored (regression test: auto-stop used to leave _active - True forever, so start() always bailed out early).""" + True forever, so start() always bailed out early). + """ default_profiler_config.max_iterations = 2 profiler = ConcreteWorkerProfiler(default_profiler_config) @@ -333,7 +332,8 @@ def test_skips_annotation_work_after_profiler_stops(self): def test_profiler_entered_during_capture(): """Profiler is used as a context manager in _warmup_and_capture, - confirming it is active during the actual graph capture run.""" + confirming it is active during the actual graph capture run. + """ runner = MagicMock() runner.compilation_config.cudagraph_num_of_warmups = 0 mock_profiler = MagicMock() diff --git a/tests/v1/worker/test_gpu_rejection_sampler_i64.py b/tests/v1/worker/test_gpu_rejection_sampler_i64.py index 769444ed64c6..b69732c8e19c 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_i64.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_i64.py @@ -34,7 +34,8 @@ def _run( ) -> tuple[torch.Tensor, torch.Tensor]: """Run rejection_sample on a batch of uniform k=2 requests, with the request under test at request index `roi_req` and draft-logits row - `req_state_idx`. Returns that request's (sampled, num_sampled).""" + `req_state_idx`. Returns that request's (sampled, num_sampled). + """ num_logits = num_reqs * LOGITS_PER_REQ roi_row = roi_req * LOGITS_PER_REQ diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index 2bc697b08438..a7fe0608a7ce 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -288,7 +288,8 @@ def test_v2_thinking_budget_clamps_oversized_budget(): def test_v2_thinking_budget_continues_end_prefix_from_prompt(): """A resumed prompt ending with a partial forced-end marker must not - restart the marker sequence and duplicate its first token.""" + restart the marker sequence and duplicate its first token. + """ req_states = _make_req_states([1, START, 10, 11, END_A], prompt_len=5) state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) state.add_request(3, SamplingParams(thinking_token_budget=3)) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index e735953d8b15..13f1bbe45ecc 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -701,7 +701,8 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int): @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA") def test_each_microbatch_owns_its_dcp_buffer(): """Microbatch DCP lengths live in per-microbatch persistent buffers so - CPU-visible shapes stay valid under CUDA graph capture and replay.""" + CPU-visible shapes stay valid under CUDA graph capture and replay. + """ input_batch, _ = _make_cuda_input_batch([4, 4, 4, 4], [8, 9, 10, 11]) ubatches, dcp_buffers = _slice_with_dcp(input_batch, dcp_rank=0) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 88cbee680240..f780b2f78f5f 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -49,7 +49,8 @@ def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path): def test_startup_plan_fingerprint_sensitivity(plan_env): """The fingerprint is the OOM-safety key: stable for identical inputs, - different for anything the profiled value depends on.""" + different for anything the profiled value depends on. + """ fp = startup_plan.compute_plan_fingerprint base = fp(_plan_worker().vllm_config, 0, 1) assert base == fp(_plan_worker().vllm_config, 0, 1) @@ -97,7 +98,8 @@ def _snapshot(free_memory, torch_memory=0): def _profile_result(consumed, reserved_before=0, reserved_after=0): """A result whose free-memory readings agree with `consumed`, which - `memory_profiling` derives as the drop in free memory, negative when it grew.""" + `memory_profiling` derives as the drop in free memory, negative when it grew. + """ return SimpleNamespace( total_consumed=consumed, transient_peak_headroom=0, @@ -125,7 +127,8 @@ def test_profiling_fallback_declines_when_free_memory_dropped(rocm): @pytest.mark.parametrize("rocm", [True], indirect=True) def test_profiling_fallback_replaces_a_released_measurement(rocm): """A negative measurement describes the rest of the device, so it is replaced - by this process's reservation, which the rest of the device cannot move.""" + by this process's reservation, which the rest of the device cannot move. + """ result = _profile_result( consumed=-RELEASED_BY_OTHERS, reserved_after=TORCH_RESERVED, @@ -149,7 +152,8 @@ def test_profiling_fallback_never_returns_a_negative_amount(rocm): @pytest.mark.parametrize("rocm", [False], indirect=True) def test_profiling_fallback_declines_off_rocm(rocm): """Platforms that account frees eagerly keep reporting the error, so the - caller's assertion stays reachable there.""" + caller's assertion stays reachable there. + """ result = _profile_result(consumed=-RELEASED_BY_OTHERS) assert maybe_rocm_profiling_fallback(result) is None @@ -173,7 +177,8 @@ def test_execute_model_waits_previous_pp_send_before_forward( monkeypatch: pytest.MonkeyPatch, ): """Previous device handles are waited before the forward pass; the - metadata handle is left to the GroupCoordinator's reaper.""" + metadata handle is left to the GroupCoordinator's reaper. + """ import torch from vllm.sequence import IntermediateTensors diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 64a81771ddd3..de541c9df670 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -358,7 +358,8 @@ def test_warmup_respects_available_block_count(): def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout): """The zeroer must zero every byte of the target block in every layer and nothing outside it — per head-group region under LHBNC, and never past the target block's - tile under block-major layouts (no out-of-bounds writes, no clobbering).""" + tile under block-major layouts (no out-of-bounds writes, no clobbering). + """ device = torch.device("cuda") num_blocks, num_layers = 4, 2 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a162c39a8695..bce4c8c35838 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -645,7 +645,8 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: """Build a MambaSpecDecodeGPUContext stand-in exposing only the four - per-request staging buffers touched by stage_postprocess_inputs_to_gpu.""" + per-request staging buffers touched by stage_postprocess_inputs_to_gpu. + """ ctx = MagicMock() ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) @@ -657,7 +658,8 @@ def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): """stage_postprocess_inputs_to_gpu writes each request's state_idx and scheduled/computed/draft counts into slot `i` for `req_ids[i]`, leaves - slots past `num_reqs` untouched, and mirrors the values to the GPU buffers.""" + slots past `num_reqs` untouched, and mirrors the values to the GPU buffers. + """ device = torch.device("cpu") max_num_reqs = 8 ctx = _make_staging_ctx(max_num_reqs, device) @@ -729,7 +731,8 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): """If preprocess_mamba didn't populate mamba_state_idx for a req in the - batch, staging must fail loudly rather than silently writing a stale index.""" + batch, staging must fail loudly rather than silently writing a stale index. + """ device = torch.device("cpu") ctx = _make_staging_ctx(max_num_reqs=4, device=device) scheduler_output = _make_postprocess_scheduler_output( @@ -807,7 +810,8 @@ def _run_gpu_postprocess( device: torch.device, ) -> None: """Initialize the GPU context against `block_table`, run the fused - postprocess kernel for `req_ids`, and synchronize.""" + postprocess kernel for `req_ids`, and synchronize. + """ def t(values): return torch.tensor(values, dtype=torch.int32, device=device) @@ -867,8 +871,7 @@ def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device): torch.testing.assert_close(state, expected, rtol=0, atol=0) def test_matches_python_postprocess_mamba(self, device, test_config): - """ - Golden test: GPU kernel produces identical results to Python impl. + """Golden test: GPU kernel produces identical results to Python impl. This test: 1. Sets up identical initial state for both paths @@ -1146,8 +1149,7 @@ def test_various_batch_sizes(self, device, test_config, num_reqs): ) def test_block_table_with_realistic_stride(self, device, test_config): - """ - Test kernel with realistic block table strides. + """Test kernel with realistic block table strides. In real usage, the block table is pre-allocated with shape [max_num_reqs, max_num_blocks_per_req] and then sliced to @@ -1275,8 +1277,7 @@ def test_block_table_with_realistic_stride(self, device, test_config): def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( self, device, test_config ): - """ - Test the ``src_addr == dst_addr`` early-return path in + """Test the ``src_addr == dst_addr`` early-return path in postprocess_mamba_fused_kernel matches Python behavior. When src_addr == dst_addr (source and destination memory addresses are @@ -1419,8 +1420,7 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( self, device, test_config ): - """ - Test the ``src_block_idx == dest_block_idx`` post-copy update in + """Test the ``src_block_idx == dest_block_idx`` post-copy update in postprocess_mamba_fused_kernel matches Python behavior. When src_block_idx == dest_block_idx but accept_token_bias > 0, both @@ -1583,8 +1583,7 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( def test_different_block_idx_copies_without_setting_accepted_to_1( self, device, test_config ): - """ - Test that neither special-case path triggers when + """Test that neither special-case path triggers when src_block_idx != dest_block_idx, and GPU matches Python behavior. When copying between different blocks: @@ -1716,8 +1715,7 @@ def test_different_block_idx_copies_without_setting_accepted_to_1( def test_prefix_caching_shared_block_does_not_set_accepted_to_1( self, device, test_config ): - """ - Regression test: with prefix caching, different logical block indices + """Regression test: with prefix caching, different logical block indices can map to the same physical block. The kernel must NOT set num_accepted_tokens to 1 in that case. @@ -1849,8 +1847,7 @@ def test_prefix_caching_shared_block_does_not_set_accepted_to_1( ) def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_config): - """ - Regression test: non-sequential physical block IDs under prefix caching + """Regression test: non-sequential physical block IDs under prefix caching with the needs_copy boundary at exact equality. Under PC, the block allocator assigns physical block IDs in arbitrary @@ -2000,8 +1997,7 @@ def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_conf assert input_batch_py.num_accepted_tokens_cpu[1] == num_accepted_tokens[1] def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_config): - """ - Regression test: mixed batch under prefix caching where some requests + """Regression test: mixed batch under prefix caching where some requests have shared physical blocks (aliased) and others have distinct blocks, with the needs_copy boundary at various positions. @@ -2166,8 +2162,7 @@ def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_conf def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr( self, device, test_config ): - """ - Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return + """Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return guard must compare logical block indices, not physical addresses. Under prefix caching, different logical blocks (src_block_idx=0, @@ -2299,8 +2294,7 @@ def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr( ) def test_as_strided_temporal_copy_size(self, device, test_config): - """ - Regression test for 240723d46: temporal copy_size must be + """Regression test for 240723d46: temporal copy_size must be inner_size * elem_size, not state_block_stride. In production (gpu_model_runner.py), conv and temporal states share @@ -2458,8 +2452,7 @@ def make_views(raw): ) def test_temporal_copy_with_bias_ge_2(self, device, test_config): - """ - Coverage test for the temporal-state block-table stride arithmetic + """Coverage test for the temporal-state block-table stride arithmetic when ``accept_token_bias >= 2``. The kernel computes, for temporal (non-conv) states:: diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 04f456a7c9cd..d4123269a952 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -17,7 +17,8 @@ def _fail(*args, **kwargs): @pytest.mark.parametrize("max_num_reqs", [1, 0]) def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 - the warmup must be skipped without touching the worker callbacks.""" + the warmup must be skipped without touching the worker callbacks. + """ runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) assert ( diff --git a/tests/v1/worker/test_prompt_embeds_state.py b/tests/v1/worker/test_prompt_embeds_state.py index 1893d0f3470e..f28c6d2e01ae 100644 --- a/tests/v1/worker/test_prompt_embeds_state.py +++ b/tests/v1/worker/test_prompt_embeds_state.py @@ -78,7 +78,8 @@ def _apply( def test_overlay_chunked_prefill_and_decode(): """Rows within the embeds range come from prompt_embeds at the num_computed offset; requests without embeds and requests past their - embeds length (decode) keep the base embedding.""" + embeds length (decode) keep the base embedding. + """ state = _make_state() embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32) embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32) @@ -97,7 +98,8 @@ def test_overlay_chunked_prefill_and_decode(): def test_overlay_clamps_to_embeds_length(): """A window straddling the end of the prompt embeds writes only the - in-range rows (e.g. final prefill chunk + sampled token).""" + in-range rows (e.g. final prefill chunk + sampled token). + """ state = _make_state() embeds = torch.randn(4, HIDDEN, dtype=torch.float32) state.add_request(3, _NewReqData("a", embeds)) @@ -112,7 +114,8 @@ def test_overlay_clamps_to_embeds_length(): def test_overlay_respects_is_token_ids_mask(): """Mixed mode: positions marked as real token ids keep the base - embedding; only embed positions are overwritten.""" + embedding; only embed positions are overwritten. + """ state = _make_state() embeds = torch.randn(5, HIDDEN, dtype=torch.float32) is_token_ids = [True, False, False, True, False] @@ -130,7 +133,8 @@ def test_overlay_respects_is_token_ids_mask(): def test_index_reuse_clears_stale_entry(): """A request added at a previously-used index without embeds must not - inherit the prior occupant's pointer-table entry.""" + inherit the prior occupant's pointer-table entry. + """ state = _make_state() state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32))) # A second live embeds request so the kernel actually launches (the diff --git a/tests/v1/worker/test_worker_memory_snapshot.py b/tests/v1/worker/test_worker_memory_snapshot.py index 6ed62cb771fa..3232d3abcf8e 100644 --- a/tests/v1/worker/test_worker_memory_snapshot.py +++ b/tests/v1/worker/test_worker_memory_snapshot.py @@ -35,6 +35,7 @@ def make_operation_tracker(operation_name: str, original_func): Returns: A wrapper function that tracks the operation and calls the original + """ def wrapper(*args, **kwargs): diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 7824c8000d2b..cc52f28e4b5c 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -84,7 +84,8 @@ def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None: def test_workspace_lock_blocks_growth_and_unlock_restores(monkeypatch) -> None: """Once locked, oversized requests fail loudly instead of reallocating the - buffer that captured CUDA graphs point at; unlock restores growth.""" + buffer that captured CUDA graphs point at; unlock restores growth. + """ monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) diff --git a/tests/vllm_test_utils/vllm_test_utils/__init__.py b/tests/vllm_test_utils/vllm_test_utils/__init__.py index 2818428de4a7..724b57613263 100644 --- a/tests/vllm_test_utils/vllm_test_utils/__init__.py +++ b/tests/vllm_test_utils/vllm_test_utils/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -vllm_utils is a package for vLLM testing utilities. +"""vllm_utils is a package for vLLM testing utilities. It does not import any vLLM modules. """ diff --git a/tests/vllm_test_utils/vllm_test_utils/blame.py b/tests/vllm_test_utils/vllm_test_utils/blame.py index 9746c3964e21..73eddfdf02a9 100644 --- a/tests/vllm_test_utils/vllm_test_utils/blame.py +++ b/tests/vllm_test_utils/vllm_test_utils/blame.py @@ -16,8 +16,7 @@ class BlameResult: @contextlib.contextmanager def blame(func: Callable) -> Generator[BlameResult, None, None]: - """ - Trace the function calls to find the first function that satisfies the + """Trace the function calls to find the first function that satisfies the condition. The trace stack will be stored in the result. Usage: diff --git a/tests/vllm_test_utils/vllm_test_utils/monitor.py b/tests/vllm_test_utils/vllm_test_utils/monitor.py index ba22bde8795b..658347889efe 100644 --- a/tests/vllm_test_utils/vllm_test_utils/monitor.py +++ b/tests/vllm_test_utils/vllm_test_utils/monitor.py @@ -21,8 +21,7 @@ class MonitoredValues(Generic[_T]): def monitor( measure_func: Callable[[], _T], ) -> Generator[MonitoredValues[_T], None, None]: - """ - Trace the function calls to continuously monitor the change of + """Trace the function calls to continuously monitor the change of a value. Usage: diff --git a/tests/weight_loading/test_weight_loading.py b/tests/weight_loading/test_weight_loading.py index e739164a5f08..c32ce71acae2 100644 --- a/tests/weight_loading/test_weight_loading.py +++ b/tests/weight_loading/test_weight_loading.py @@ -29,10 +29,7 @@ reason="Current system does not have minimum capability.", ) def test_weight_loading(vllm_runner): - """ - Test parameter weight loading with tp>1. - """ - + """Test parameter weight loading with tp>1.""" # MoE models need fp16. NEEDS_FP16 = QUANTIZATION == "gptq" with vllm_runner( diff --git a/tools/generate_cmake_presets.py b/tools/generate_cmake_presets.py index 6bc8443c4473..76d8341843df 100644 --- a/tools/generate_cmake_presets.py +++ b/tools/generate_cmake_presets.py @@ -28,7 +28,6 @@ def get_cpu_cores(): def generate_presets(output_path="CMakeUserPresets.json", force_overwrite=False): """Generates the CMakeUserPresets.json file.""" - print("Attempting to detect your system configuration...") # Detect NVCC diff --git a/tools/generate_versions_json.py b/tools/generate_versions_json.py index f5d2893a996b..e1611457c913 100755 --- a/tools/generate_versions_json.py +++ b/tools/generate_versions_json.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Generate docker/versions.json from Dockerfile ARG defaults. +"""Generate docker/versions.json from Dockerfile ARG defaults. This script parses the Dockerfile and extracts ARG defaults to create a bake-native versions.json file that can be used directly with: diff --git a/tools/install_nixl_from_source_ubuntu.py b/tools/install_nixl_from_source_ubuntu.py index b8a55c615426..1eb73bd0cbcf 100644 --- a/tools/install_nixl_from_source_ubuntu.py +++ b/tools/install_nixl_from_source_ubuntu.py @@ -21,7 +21,7 @@ # --- Helper Functions --- def get_latest_nixl_version(): - """Helper function to get latest release version of NIXL""" + """Helper function to get latest release version of NIXL.""" try: nixl_release_url = "https://api.github.com/repos/ai-dynamo/nixl/releases/latest" with urllib.request.urlopen(nixl_release_url) as response: @@ -104,7 +104,6 @@ def install_system_dependencies(): def build_and_install_prerequisites(args): """Builds UCX and NIXL from source, creating a self-contained wheel.""" - if not args.force_reinstall and is_pip_package_installed("nixl"): print("--> NIXL is already installed. Nothing to do.", flush=True) return diff --git a/tools/pre_commit/check_spdx_header.py b/tools/pre_commit/check_spdx_header.py index 2585e82b4eba..1252ccb7adcb 100644 --- a/tools/pre_commit/check_spdx_header.py +++ b/tools/pre_commit/check_spdx_header.py @@ -16,7 +16,7 @@ class HeaderStyle: class SPDXStatus(Enum): - """SPDX header status enumeration""" + """SPDX header status enumeration.""" EMPTY = "empty" # empty __init__.py COMPLETE = "complete" @@ -58,7 +58,7 @@ def header_insertion_index(style, lines): def check_spdx_header_status(file_path): - """Check SPDX header status of the file""" + """Check SPDX header status of the file.""" license_line, copyright_line = spdx_header(file_style(file_path)) with open(file_path, encoding="UTF-8") as file: lines = file.readlines() @@ -92,7 +92,7 @@ def check_spdx_header_status(file_path): def add_header(file_path, status): - """Add or supplement SPDX header based on status""" + """Add or supplement SPDX header based on status.""" style = file_style(file_path) license_line, copyright_line = spdx_header(style) full_spdx_header = f"{license_line}\n{copyright_line}" @@ -137,7 +137,7 @@ def add_header(file_path, status): def main(): - """Main function""" + """Main function.""" files_missing_both = [] files_missing_copyright = [] files_missing_license = [] diff --git a/tools/pre_commit/generate_nightly_torch_test.py b/tools/pre_commit/generate_nightly_torch_test.py index 354da54df8f4..3b39c4a0b274 100644 --- a/tools/pre_commit/generate_nightly_torch_test.py +++ b/tools/pre_commit/generate_nightly_torch_test.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Generates specialized requirements files for nightly PyTorch testing. +"""Generates specialized requirements files for nightly PyTorch testing. This script reads the main test requirements input file (`requirements/test/cuda.in`) and splits its content into two files: diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 9da39acc7899..b7a278b9bd6b 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Run mypy on changed files. +"""Run mypy on changed files. This script is designed to be used as a pre-commit hook. It runs mypy on files that have been changed. It groups files into different mypy calls @@ -14,6 +13,7 @@ python_version: Python version to use (e.g., "3.10") or "local" to use the local Python version. changed_files: List of changed files to check. + """ import subprocess @@ -125,14 +125,14 @@ def group_files(changed_files: list[str]) -> dict[str, list[str]]: - """ - Group changed files into different mypy calls. + """Group changed files into different mypy calls. Args: changed_files: List of changed files. Returns: A dictionary mapping file group names to lists of changed files. + """ exclude_pattern = re.compile(f"^{'|'.join(EXCLUDE)}.*") silent_pattern = re.compile(f"^({'|'.join(SILENT_GROUPS)}).*") @@ -166,8 +166,7 @@ def mypy( follow_imports: str | None, file_group: str, ) -> int: - """ - Run mypy on the given targets. + """Run mypy on the given targets. Args: targets: List of files or directories to check. @@ -179,6 +178,7 @@ def mypy( Returns: The return code from mypy. + """ args = ["mypy"] if python_version is not None: diff --git a/tools/pre_commit/validate_config.py b/tools/pre_commit/validate_config.py index 7da32bc6b485..498caf7ab2ca 100644 --- a/tools/pre_commit/validate_config.py +++ b/tools/pre_commit/validate_config.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Ensures all fields in a config dataclass have default values +"""Ensures all fields in a config dataclass have default values and that each field has a docstring. """ @@ -14,13 +13,11 @@ def get_attr_docs(cls_node: ast.ClassDef) -> dict[str, str]: - """ - Get any docstrings placed after attribute assignments in a class body. + """Get any docstrings placed after attribute assignments in a class body. Adapted from https://davidism.com/attribute-docstrings/ https://davidism.com/mit-license/ """ - out = {} # Consider each pair of nodes. diff --git a/tools/profiler/nsys_profile_tools/gputrc2graph.py b/tools/profiler/nsys_profile_tools/gputrc2graph.py index fd237c0b214a..2759d029f75d 100755 --- a/tools/profiler/nsys_profile_tools/gputrc2graph.py +++ b/tools/profiler/nsys_profile_tools/gputrc2graph.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This generates gpu kernel analysis output from nsys rep. Will call nsys +"""This generates gpu kernel analysis output from nsys rep. Will call nsys stats -r cuda_gpu_kern_trace, get non-overlapped gpu cycles, then generate csv and html output for analysis """ @@ -17,7 +16,7 @@ # helper data class for annotating kernels def load_engine_model(): - """returns engine_model built from all json files in the current dir""" + """Returns engine_model built from all json files in the current dir.""" import glob import json @@ -31,9 +30,7 @@ def load_engine_model(): class GPUTrace2Graph: - """ - Parses output of nsys report, generates csv and bar chart output - """ + """Parses output of nsys report, generates csv and bar chart output.""" def __init__(self): import pandas as pd # avoid importing till needed @@ -64,8 +61,7 @@ def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file): ) def sum_non_overlapping_intervals(self, df): - """ - returns new sorted df with Elapsed Time (ns) column using + """Returns new sorted df with Elapsed Time (ns) column using vectorized operations """ logger.info("sorting %s trace records by start time", str(df.shape)) @@ -105,7 +101,7 @@ def sum_non_overlapping_intervals(self, df): # functions for generating html files def make_html(self, df, output_dir, title): - """make html graph from df""" + """Make html graph from df.""" import plotly.express as px if df.empty: @@ -163,7 +159,7 @@ def make_html(self, df, output_dir, title): ) def anno_gpu_kernname(self, df, mapping): - """add "Category" column""" + """Add "Category" column.""" def anno_gpu_kernname_helper(name): for kern_name, val in mapping.items(): @@ -173,7 +169,7 @@ def anno_gpu_kernname_helper(name): df["Category"] = df["Name"].apply(anno_gpu_kernname_helper) def make_nongpu_row(self, df, nongpu_sec): - """this will append non-gpu time entry at end of df""" + """This will append non-gpu time entry at end of df.""" nongpu_row = self.pd.DataFrame([df.iloc[-1]]) nongpu_row["Category"] = nongpu_row["Name"] = "CPU(non-GPU)" nongpu_row["Instances"] = 1 @@ -181,13 +177,13 @@ def make_nongpu_row(self, df, nongpu_sec): return nongpu_row def is_valid_file(self, base_file): - """asserts if base_file is non-existent or is empty""" + """Asserts if base_file is non-existent or is empty.""" assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, ( f"{base_file} doesn't exist or is empty" ) def should_gen_file(self, new_file, base_file): - """figure out if new file should be generated from base_file""" + """Figure out if new file should be generated from base_file.""" self.is_valid_file(base_file) if ( os.path.exists(new_file) @@ -201,8 +197,7 @@ def should_gen_file(self, new_file, base_file): return True def gen_sum_file(self, file, nsys_cmd): - """ - generates sum file from nsys trace with times per kernel and + """Generates sum file from nsys trace with times per kernel and returns the name of the sum file """ import subprocess @@ -246,7 +241,7 @@ def gen_sum_file(self, file, nsys_cmd): return sum_file def gen_graph(self, in_file, out_dir, title, nsys_cmd, engine_model): - """generates graph and csv file from in_file into out_dir""" + """Generates graph and csv file from in_file into out_dir.""" # Initialize an empty DataFrame to store combined data combined_df = self.pd.DataFrame() for idx, (file, engine, model, total_sec) in enumerate(in_file): diff --git a/tools/recipes/recipe_json_to_vllm_config.py b/tools/recipes/recipe_json_to_vllm_config.py index 5fd7b7729d26..bd947b07451c 100644 --- a/tools/recipes/recipe_json_to_vllm_config.py +++ b/tools/recipes/recipe_json_to_vllm_config.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Convert a vLLM Recipes per-hardware JSON rendering into: +"""Convert a vLLM Recipes per-hardware JSON rendering into: 1) config.yml - native `vllm serve --config` YAML 2) env.sh - environment variables required by the recipe @@ -38,6 +37,7 @@ The converter intentionally targets a single `vllm serve` process. If the recipe rendering is multi-node, PD-disaggregated, or another multi-process deployment, it exits instead of silently generating an incomplete config. + """ from __future__ import annotations @@ -551,8 +551,7 @@ def merge_value(dst: dict[str, Any], path: list[str], value: Any) -> None: def normalize_key(raw_key: str) -> list[str]: - """ - Convert the CLI key to config-file spelling. + """Convert the CLI key to config-file spelling. Only the top-level CLI option name gets underscore -> dash normalization. Nested JSON field names after a dot are preserved. diff --git a/tools/report_build_time_ninja.py b/tools/report_build_time_ninja.py index fe3f352fe153..916e29a35f92 100644 --- a/tools/report_build_time_ninja.py +++ b/tools/report_build_time_ninja.py @@ -56,7 +56,8 @@ class Target: def __init__(self, start, end): """Creates a target object by passing in the start/end times in seconds - as a float.""" + as a float. + """ self.start = start self.end = end # A list of targets, appended to by the owner of this object. @@ -104,7 +105,8 @@ def DescribeTargets(self): def ReadTargets(log, show_all): """Reads all targets from .ninja_log file |log_file|, sorted by duration. - The result is a list of Target objects.""" + The result is a list of Target objects. + """ header = log.readline() assert header == "# ninja log v5\n", "unrecognized ninja log version {!r}".format( header @@ -153,7 +155,8 @@ def GetExtension(target, extra_patterns): For targets that generate multiple outputs it is important to return a consistent 'canonical' extension. Ultimately the goal is to group build steps - by type.""" + by type. + """ for output in target.targets: if extra_patterns: for fn_pattern in extra_patterns.split(";"): @@ -197,7 +200,6 @@ def GetExtension(target, extra_patterns): def SummarizeEntries(entries, extra_step_types): """Print a summary of the passed in list of Target objects.""" - # Create a list that is in order by time stamp and has entries for the # beginning and ending of each build step (one time stamp may have multiple # entries due to multiple steps starting/stopping at exactly the same time). diff --git a/tools/vllm-rocm/pin_rocm_dependencies.py b/tools/vllm-rocm/pin_rocm_dependencies.py index 7d90d66692ad..f4c5041b0ee1 100644 --- a/tools/vllm-rocm/pin_rocm_dependencies.py +++ b/tools/vllm-rocm/pin_rocm_dependencies.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Pin vLLM dependencies to exact versions of custom ROCm wheels. +"""Pin vLLM dependencies to exact versions of custom ROCm wheels. This script modifies vLLM's requirements files to replace version constraints with exact versions of custom-built ROCm wheels (torch, triton, torchvision, amdsmi). @@ -18,12 +17,12 @@ def extract_version_from_wheel(wheel_name: str) -> str: - """ - Extract version from wheel filename. + """Extract version from wheel filename. Example: torch-2.9.0a0+git1c57644-cp312-cp312-linux_x86_64.whl -> 2.9.0a0+git1c57644 triton-3.4.0-cp312-cp312-linux_x86_64.whl -> 3.4.0 + """ # Wheel format: # {distribution}-{version}(-{build tag})?-{python}-{abi}-{platform}.whl @@ -38,11 +37,11 @@ def extract_version_from_wheel(wheel_name: str) -> str: def get_custom_wheel_versions(install_dir: str) -> dict[str, str]: - """ - Read /install directory and extract versions of custom wheels. + """Read /install directory and extract versions of custom wheels. Returns: Dict mapping package names to exact versions + """ install_path = Path(install_dir) if not install_path.exists(): @@ -92,8 +91,7 @@ def get_custom_wheel_versions(install_dir: str) -> dict[str, str]: def pin_dependencies_in_requirements(requirements_path: str, versions: dict[str, str]): - """ - Insert custom wheel pins at the TOP of requirements file. + """Insert custom wheel pins at the TOP of requirements file. This ensures that when setup.py processes the file line-by-line, custom wheels (torch, triton, etc.) are encountered FIRST, before diff --git a/vllm/__init__.py b/vllm/__init__.py index 968d1a143b16..16101ca245f6 100644 --- a/vllm/__init__.py +++ b/vllm/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""vLLM: a high-throughput and memory-efficient inference engine for LLMs""" +"""vLLM: a high-throughput and memory-efficient inference engine for LLMs.""" # The version.py should be independent library, and we always import the # version library first. Such assumption is critical for some customization. diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 8f2b2a24adac..dfcba987d60b 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1835,6 +1835,7 @@ class rocm_aiter_ops: - MLA decode: mla_decode_fwd - Quantization: per_tensor_quant, per_token_quant, group_fp8_quant - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale + """ _MOE_DISPATCH_POLICY: int | None = None @@ -1874,8 +1875,7 @@ def get_moe_dispatch_policy(cls) -> int: @classmethod def refresh_env_variables(cls): - """ - Since the environment variables are assigned when the module is imported, + """Since the environment variables are assigned when the module is imported, This is a helper function to reload all the env variables from the environment variables. for example, after monkey patching the env variables in the unit test, @@ -1900,8 +1900,7 @@ def refresh_env_variables(cls): @staticmethod def get_aiter_activation_type(activation_str: str) -> "ActivationType | None": - """ - Given an activation type as a string, returns the corresponding aiter ActivationType enum. + """Given an activation type as a string, returns the corresponding aiter ActivationType enum. Supported activation types: "no", "none", "silu", "gelu", "swiglu". Returns None if the mapping fails. @@ -1910,6 +1909,7 @@ def get_aiter_activation_type(activation_str: str) -> "ActivationType | None": Returns: Aiter ActivationType enum value, or None if not found. + """ # Import only locally, since aiter may not always be available. try: @@ -1933,8 +1933,7 @@ def get_aiter_activation_type(activation_str: str) -> "ActivationType | None": @staticmethod def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None": - """ - Given a quantization type as a string, returns the corresponding aiter QuantType enum. + """Given a quantization type as a string, returns the corresponding aiter QuantType enum. Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128". Returns None if the mapping fails. @@ -1943,6 +1942,7 @@ def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None": Returns: Aiter QuantType enum value, or None if not found. + """ try: from aiter import QuantType @@ -1974,7 +1974,8 @@ def is_rdna_aiter_enabled(cls) -> bool: enabled aiter. Only aiter's Triton kernels exist on gfx12 (no CK build), so this deliberately stays off the gfx9/CK `@if_aiter_supported` umbrella and gates only the Triton paths rdna4 uses. The gfx12 analog of - `is_enabled()`.""" + `is_enabled()`. + """ if not current_platform.is_rocm() or not IS_AITER_FOUND: return False from vllm.platforms.rocm import on_rdna4 @@ -3357,8 +3358,7 @@ def shuffle_weight_a16w4( nLane: int, gate_up: bool, ) -> "torch.Tensor": - """ - Shuffles the weight tensor into (A16W4) layout for AITER kernels. + """Shuffles the weight tensor into (A16W4) layout for AITER kernels. Args: tensor: The input weight tensor to be shuffled. @@ -3367,6 +3367,7 @@ def shuffle_weight_a16w4( Returns: torch.Tensor: The shuffled tensor. + """ from aiter.ops.shuffle import shuffle_weight_a16w4 @@ -3378,8 +3379,7 @@ def shuffle_scale_a16w4( num_experts: int, gate_up: bool, ) -> "torch.Tensor": - """ - Shuffles the scale tensor into (A16W4) layout for AITER kernels. + """Shuffles the scale tensor into (A16W4) layout for AITER kernels. Args: tensor: The input scale tensor to be shuffled. @@ -3388,6 +3388,7 @@ def shuffle_scale_a16w4( Returns: torch.Tensor: The shuffled scale tensor. + """ from aiter.ops.shuffle import shuffle_scale_a16w4 @@ -3404,8 +3405,7 @@ def shuffle_scale(tensor: torch.Tensor) -> torch.Tensor: def shuffle_weights( *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16) ) -> tuple[torch.Tensor, ...]: - """ - Applies shuffle_weight function from AITER to each + """Applies shuffle_weight function from AITER to each input tensor and returns them. Rearranges (shuffles) the input tensor/s @@ -3418,6 +3418,7 @@ def shuffle_weights( Returns: A Tuple of shuffled tensors. + """ from aiter.ops.shuffle import shuffle_weight @@ -3467,8 +3468,7 @@ def flash_attn_varlen_func( out: torch.Tensor | None = None, sink_ptr: torch.Tensor | None = None, ): - """ - Flash attention with variable length sequences. + """Flash attention with variable length sequences. This function is NOT wrapped with @is_aiter_supported decorator to allow explicit backend selection via attention_config to work @@ -3537,8 +3537,7 @@ def pa_fwd_asm( V_QScale: torch.Tensor, out_: torch.Tensor, ): - """ - Paged attention forward pass using assembly kernel. + """Paged attention forward pass using assembly kernel. This function is NOT wrapped with @is_aiter_supported decorator to allow explicit backend selection via attention_config to work @@ -3580,8 +3579,7 @@ def paged_attention_common( out_: torch.Tensor, kv_cache_dtype: str, ): - """ - Paged attention common function. + """Paged attention common function. This function is NOT wrapped with @is_aiter_supported decorator to allow explicit backend selection via attention_config to work @@ -3625,8 +3623,7 @@ def mhc_pre( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Forward pass for mHC pre block. + """Forward pass for mHC pre block. Args: residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 @@ -3645,6 +3642,7 @@ def mhc_pre( post_mix: shape (..., hc_mult), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 + """ from aiter.ops.mhc import mhc_pre @@ -3735,7 +3733,7 @@ def mhc_pre_delayed( comb_res_mix: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """mHC pre using the pre-mix carried from the previous sublayer. + """MHC pre using the pre-mix carried from the previous sublayer. Same gates as :meth:`mhc_pre`, but the stream collapse applies the caller's ``pre_mix`` instead of the one computed here, and the one @@ -3762,6 +3760,7 @@ def mhc_pre_delayed( comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 next_pre_mix: shape (..., hc_mult), dtype torch.float32 + """ from aiter.ops.mhc import ( get_mhc_fused_post_pre_config, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 99577766e7a8..96f47e2069f1 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -39,8 +39,7 @@ def create_fp4_scale_tensor( device: torch.device, is_sf_swizzled_layout: bool, ) -> torch.Tensor: - """ - Allocate the output scale tensor for scaled_fp4_quant. + """Allocate the output scale tensor for scaled_fp4_quant. When is_sf_swizzled_layout=True, we use rounded values to store the swizzled scales. Due to the requirement of the Tensor Core, the minimum @@ -73,8 +72,7 @@ def create_fp4_output_tensors( is_sf_swizzled_layout: bool, padded_n: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Allocate both output tensors for scaled_fp4_quant: + """Allocate both output tensors for scaled_fp4_quant: (quantized_output, output_scale). Must match the C++ scaled_fp4_quant_func allocation exactly when @@ -265,9 +263,18 @@ def vocab_parallel_embedding( Args: input_ids: ``[num_tokens]`` int32 or int64 token ids. weight: ``[num_embeddings_per_partition, embedding_dim]`` shard. + org_vocab_start_index: First original vocab id owned by this rank. + org_vocab_end_index: One past the last original vocab id owned by + this rank. + num_org_vocab_padding: Padding rows between this rank's original and + added vocab shards. + added_vocab_start_index: First added vocab id owned by this rank. + added_vocab_end_index: One past the last added vocab id owned by + this rank. Returns: ``[num_tokens, embedding_dim]`` partial embeddings. + """ out = torch.empty( input_ids.shape[0], @@ -399,6 +406,7 @@ def apply_repetition_penalties( prompt_mask: A boolean tensor indicating which tokens appear in the prompt. output_mask: A boolean tensor indicating which tokens appear in the output. repetition_penalties: The repetition penalties of shape (num_seqs, ). + """ if logits.is_cuda and logits.is_contiguous(): apply_repetition_penalties_cuda( @@ -787,8 +795,7 @@ def cutlass_scaled_mm( out_dtype: torch.dtype, bias: torch.Tensor | None = None, ) -> torch.Tensor: - """ - `cutlass_scaled_mm` implements a fused version of + """`cutlass_scaled_mm` implements a fused version of `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)` where scale_a * a and scale_b * b are implemented using numpy-style broadcasting. @@ -840,11 +847,11 @@ def cutlass_scaled_mm_azp( azp: torch.Tensor | None = None, bias: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Args: - azp_adj: In the per-tensor case, this should include the azp. - Always per-channel. - azp: Only set in the per-token case. Per-token if set. + """Args: + azp_adj: In the per-tensor case, this should include the azp. + Always per-channel. + azp: Only set in the per-token case. Per-token if set. + """ assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0 assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 @@ -883,8 +890,7 @@ def get_cutlass_moe_mm_data( blockscale_offsets: torch.Tensor | None = None, is_gated: bool = True, ): - """ - Prepare data necessary to perform CUTLASS grouped matrix multiplications + """Prepare data necessary to perform CUTLASS grouped matrix multiplications used in CUTLASS-based fused MoE. The function takes in topk_ids (token-expert mapping) and uses it to @@ -943,8 +949,7 @@ def get_cutlass_moe_mm_problem_sizes_from_expert_offsets( def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor): - """ - Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor. + """Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor. This is used in MoE to permute the input tensor before performing grouped matrix multiplications. """ num_tokens_permuted = dst2src_map.shape[0] @@ -967,8 +972,7 @@ def get_cutlass_batched_moe_mm_data( n: int, k: int, ): - """ - Prepare data necessary to perform CUTLASS grouped matrix multiplications + """Prepare data necessary to perform CUTLASS grouped matrix multiplications used in CUTLASS-based fused MoE. The function takes in expert_num_tokens (token count per expert) and @@ -1006,8 +1010,7 @@ def cutlass_moe_mm( per_act_token: bool, per_out_ch: bool, ): - """ - A single grouped matrix multiplication used in CUTLASS-based fused MoE. + """A single grouped matrix multiplication used in CUTLASS-based fused MoE. The function executes fp8-quantized OUT = AB matrix multiplication. - expert_offsets: Indices that mark at which token index each expert begins @@ -1044,8 +1047,7 @@ def cutlass_fp4_moe_mm( expert_offsets: torch.Tensor, sf_offsets: torch.Tensor, ): - """ - An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs + """An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs the gemms for each combination based on the specified problem sizes. This is used as the MoE gemm during NVFP4 Quantized MoERunner forward. @@ -1083,8 +1085,7 @@ def cutlass_mxfp4_moe_mm( expert_offsets: torch.Tensor, sf_offsets: torch.Tensor, ): - """ - An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4). + """An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4). Uses mx_float4_t types with E8M0 scale factors and 32-element blocks. - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte) @@ -1458,8 +1459,7 @@ def cutlass_w4a8_moe_mm( group_scale_strides: torch.Tensor, maybe_schedule: str | None = None, ) -> None: - """ - Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the + """Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8) and both per-channel + per-token scaling in the epilogue. @@ -1495,6 +1495,7 @@ def cutlass_w4a8_moe_mm( Returns: out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result. + """ return torch.ops._C.cutlass_w4a8_moe_mm( out_tensors, @@ -1535,8 +1536,7 @@ def scaled_fp4_quant( backend: str = "none", padded_n: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to FP4 and return quantized tensor and scale. + """Quantize input tensor to FP4 and return quantized tensor and scale. This function quantizes the last dimension of the given tensor `input`. For every 16 consecutive elements, a single dynamically computed scaling factor @@ -1559,6 +1559,7 @@ def scaled_fp4_quant( tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every two values are packed into a uint8 and float8_e4m3 scaling factors in the sizzled layout. + """ assert not current_platform.is_rocm() assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}." @@ -1613,17 +1614,19 @@ def scaled_fp4_experts_quant( blockscale_offsets: torch.Tensor, topk: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to NVFP4 and return quantized tensor and scale, for + """Quantize input tensor to NVFP4 and return quantized tensor and scale, for packed MoE Inputs. + Args: input_tensor: The input tensor to be quantized to NVFP4 input_global_scale: A scalar scaling factor for the entire tensor. expert_offsets: The expert offsets tensor blockscale_offsets: The blockscale offsets tensor + topk: The number of experts each token is routed to Outputs: output: The quantized tensor in NVFP4 output_scales: The blockscale tensor in FP8-E4M3 + """ assert not current_platform.is_rocm() assert input_tensor.ndim == 2, ( @@ -1675,8 +1678,7 @@ def silu_and_mul_scaled_fp4_experts_quant( blockscale_offsets: torch.Tensor, topk: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations. + """Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations. Args: input_tensor: The input tensor with gate || up layout [m_topk, k*2] @@ -1687,6 +1689,7 @@ def silu_and_mul_scaled_fp4_experts_quant( Outputs: output: The quantized tensor in NVFP4 [m_topk, k/2] output_scales: The blockscale tensor in FP8-E4M3 + """ assert not current_platform.is_rocm() assert input_tensor.ndim == 2, ( @@ -1740,8 +1743,7 @@ def mxfp4_experts_quant( n_experts: int, topk: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to MXFP4 for packed MoE inputs. + """Quantize input tensor to MXFP4 for packed MoE inputs. Uses 32-element blocks with E8M0 (power-of-two) scale factors. MXFP4 has no global scale - only block-level E8M0 scale factors. @@ -1754,6 +1756,7 @@ def mxfp4_experts_quant( Returns: output: [m_topk, k//2] packed E2M1 values (uint8) output_scales: E8M0 blockscales in swizzled layout (uint8 view) + """ assert not current_platform.is_rocm() assert input_tensor.ndim == 2 @@ -1799,8 +1802,7 @@ def silu_and_mul_mxfp4_experts_quant( n_experts: int, topk: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations. + """Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations. MXFP4 has no global scale - only block-level E8M0 scale factors. """ assert not current_platform.is_rocm() @@ -1846,8 +1848,7 @@ def scaled_fp8_quant( output: torch.Tensor | None = None, group_shape: tuple[int, int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to FP8 and return quantized tensor and scale. + """Quantize input tensor to FP8 and return quantized tensor and scale. This function supports both static and dynamic quantization: If you provide the scale, it will use static scaling and if you omit it, @@ -1869,6 +1870,8 @@ def scaled_fp8_quant( of the output to at least this value. use_per_token_if_dynamic: Whether to do per_tensor or per_token in the dynamic quantization case. + output: Optional tensor to write the quantized result into. A new + tensor is allocated when omitted. group_shape: Optional tuple (group_m, group_n) specifying the group shape for static quantization. Use -1 for "full extent" (e.g., (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.) @@ -1877,6 +1880,7 @@ def scaled_fp8_quant( Returns: tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and scaling factor. + """ # This code assumes batch_dim and num_tokens are flattened assert input.ndim == 2 @@ -1913,8 +1917,7 @@ def allspark_repack_weight( zero_point: torch.Tensor | None = None, has_zp: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format + """Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format for Ampere W8A16 Fused Gemm kernel Args: @@ -1928,6 +1931,7 @@ def allspark_repack_weight( Returns: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : rearranged weight, scale, and optionally zero_point. + """ K = qweight.shape[0] N = qweight.shape[1] @@ -1997,8 +2001,7 @@ def scaled_int8_quant( azp: torch.Tensor | None = None, symmetric: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """ - Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. Args: input: The input tensor to be quantized to int8. @@ -2010,6 +2013,7 @@ def scaled_int8_quant( Returns: Output int8 tensor, scales, and optionally azp. + """ if current_platform.is_xpu(): # XPU has no _C int8 quant op; use the torch.compile reference. @@ -2463,8 +2467,7 @@ def grouped_topk( bias: torch.Tensor, scoring_func: int = 0, ): - """ - Perform grouped top-k routing for mixture of experts. + """Perform grouped top-k routing for mixture of experts. Args: scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0) @@ -2475,6 +2478,7 @@ def grouped_topk( routed_scaling_factor: Scaling factor for routing weights bias: Bias tensor (e_score_correction_bias). Always fused in kernel. scoring_func: 0=none (no activation), 1=sigmoid + """ if not (current_platform.is_cuda() or current_platform.is_xpu()): raise NotImplementedError( @@ -2911,8 +2915,7 @@ def swap_blocks( block_size_in_bytes: int, block_mapping: torch.Tensor, ) -> None: - """ - Copy specific blocks from one tensor to another. + """Copy specific blocks from one tensor to another. This method assumes each of the two input tensors is composed of consecutive contiguous blocks, of size block_size_in_bytes. @@ -2940,8 +2943,7 @@ def swap_blocks_batch( sizes: torch.Tensor, is_src_access_order_any: bool = False, ) -> None: - """ - Batch version of swap_blocks: submit all copies in a single driver call. + """Batch version of swap_blocks: submit all copies in a single driver call. Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i] of sizes[i] bytes. All three tensors must be CPU tensors with the @@ -3030,6 +3032,7 @@ def cp_gather_and_upconvert_fp8_kv_cache( host_cache: Optional pinned host rows used for non-resident entries host_row_ids: Host source row for each remapped cache row device_row_ids: Device source row, or -1, for each remapped cache row + """ torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( src_cache, @@ -3059,6 +3062,7 @@ def cp_gather_and_upconvert_nvfp4_kv_cache( block_table: Block indices [num_reqs, max_blocks] workspace_starts: Workspace start offsets [num_reqs] batch_size: Number of requests + """ torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache( src_cache, @@ -3080,6 +3084,7 @@ def concat_mla_q( ql_nope: Query nope component [num_tokens, num_heads, nope_dim] q_pe: Query rope component [num_tokens, num_heads, rope_dim] q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim] + """ torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out) @@ -3802,8 +3807,7 @@ def onednn_scaled_int8_quant( azp: torch.Tensor | None = None, symmetric: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """ - Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. Args: input: The input tensor to be quantized to int8. @@ -3815,6 +3819,7 @@ def onednn_scaled_int8_quant( Returns: Output int8 tensor, scales, and optionally azp. + """ output = torch.empty_like(input, dtype=torch.int8) token_num = input.numel() // input.shape[-1] @@ -4674,8 +4679,7 @@ def safeFusedQuantizeNv( xh_e4m3: torch.Tensor, global_scale: torch.Tensor, ) -> None: - """ - Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place + """Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place rather than returning them, to prevent torch 2.12+ errors that outputs of custom operators may not alias any inputs to the custom operator. """ @@ -4684,8 +4688,7 @@ def safeFusedQuantizeNv( def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: - """ - Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832) + """Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832) kernels. Note that these kernels exploit the recursive properties of Sylvester Hadamards, and therefore do not require transform weight data @@ -4698,6 +4701,7 @@ def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: Returns: value after transformation + """ return torch.ops._C.hadacore_transform(x, inplace) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index c1f2bb6969b9..5c19b014082e 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -286,6 +286,7 @@ def _xpu_fp8_bmm_impl( This implementation centralizes access to ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while ``a`` and ``b`` may be non-contiguous views. + """ return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias) @@ -430,6 +431,7 @@ def _xpu_deepseek_fused_indexer_q_rope_fp8_impl( index_weights_out: (T, H) float32 output, preallocated; = index_weights * q_scale * softmax_scale * head_scale (the per-(token, head) q_scale is folded in here). [written] + """ torch.ops._xpu_C.deepseek_fused_indexer_q_rope_fp8( index_q, @@ -473,6 +475,7 @@ def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl( index_weights_out: (T, H) float32 output, preallocated; = index_weights * softmax_scale * head_scale (no q_scale folded; per-block scales live in index_q_scale). [written] + """ torch.ops._xpu_C.deepseek_fused_indexer_q_rope_mxfp4( index_q, diff --git a/vllm/assets/base.py b/vllm/assets/base.py index 5ca9de4076ad..a0bd6b1eb26d 100644 --- a/vllm/assets/base.py +++ b/vllm/assets/base.py @@ -20,8 +20,7 @@ def get_cache_dir() -> Path: @lru_cache def get_vllm_public_assets(filename: str, s3_prefix: str | None = None) -> Path: - """ - Download an asset file from `s3://vllm-public-assets` + """Download an asset file from `s3://vllm-public-assets` and return the path to the downloaded file. """ asset_directory = get_cache_dir() / "vllm_public_assets" diff --git a/vllm/assets/image.py b/vllm/assets/image.py index a91eb7d4b67d..4f87fd1a4700 100644 --- a/vllm/assets/image.py +++ b/vllm/assets/image.py @@ -34,9 +34,7 @@ class ImageAsset: name: ImageAssetName def get_path(self, ext: str) -> Path: - """ - Return s3 path for given image. - """ + """Return s3 path for given image.""" return get_vllm_public_assets( filename=f"{self.name}.{ext}", s3_prefix=VLM_IMAGES_DIR ) @@ -51,9 +49,7 @@ def pil_image_ext(self, ext: str) -> Image.Image: @property def image_embeds(self) -> torch.Tensor: - """ - Image embeddings, only used for testing purposes with llava 1.5. - """ + """Image embeddings, only used for testing purposes with llava 1.5.""" image_path = self.get_path("pt") return torch.load(image_path, map_location="cpu", weights_only=True) diff --git a/vllm/assets/video.py b/vllm/assets/video.py index e06d033e0be1..bad21663751c 100644 --- a/vllm/assets/video.py +++ b/vllm/assets/video.py @@ -21,8 +21,7 @@ def _sample_frame_indices(total_frames: int, num_frames: int) -> npt.NDArray: @lru_cache def download_video_asset(filename: str) -> str: - """ - Download and open an image from huggingface + """Download and open an image from huggingface repo: raushan-testing-hf/videos-test """ video_directory = get_cache_dir() / "video-example-data" @@ -141,8 +140,7 @@ def metadata(self) -> dict[str, Any]: return ret def get_audio(self, sampling_rate: float | None = None) -> npt.NDArray: - """ - Read audio data from the video asset, used in Qwen2.5-Omni examples. + """Read audio data from the video asset, used in Qwen2.5-Omni examples. See also: examples/generate/multimodal/qwen2_5_omni/only_thinker.py """ diff --git a/vllm/benchmarks/datasets/create_txt_slices_dataset.py b/vllm/benchmarks/datasets/create_txt_slices_dataset.py index 8d19cc25ff3d..e4555bce9ae3 100644 --- a/vllm/benchmarks/datasets/create_txt_slices_dataset.py +++ b/vllm/benchmarks/datasets/create_txt_slices_dataset.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Convert a plain-text file (local path or URL) into a JSONL dataset +"""Convert a plain-text file (local path or URL) into a JSONL dataset compatible with ``CustomDataset`` (``--dataset-name custom``), by randomly slicing the tokenized text into prompts. @@ -66,8 +65,8 @@ def create_txt_slices_jsonl( trust_remote_code: bool = False, ) -> None: """Read *input_path*, slice it into prompts, and write JSONL to - *output_path*.""" - + *output_path*. + """ tokenizer = AutoTokenizer.from_pretrained( tokenizer_name, trust_remote_code=trust_remote_code ) diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 3ca9fd670806..3456fc633786 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This module defines a framework for sampling benchmark requests from various +"""This module defines a framework for sampling benchmark requests from various datasets. Each dataset subclass of BenchmarkDataset must implement sample generation. Supported dataset types include: - ShareGPT @@ -75,9 +74,7 @@ @dataclass class SampleRequest: - """ - Represents a single inference request for benchmarking. - """ + """Represents a single inference request for benchmarking.""" prompt: str | list[str] | list[int] | list[list[int]] | list[dict] prompt_len: int @@ -112,8 +109,7 @@ def __init__( disable_shuffle: bool = False, **kwargs, ) -> None: - """ - Initialize the BenchmarkDataset with an optional dataset path and random + """Initialize the BenchmarkDataset with an optional dataset path and random seed. Args: @@ -121,6 +117,9 @@ def __init__( indicates that a default or random dataset might be used. random_seed (int): Seed value for reproducible shuffling or sampling. Defaults to DEFAULT_SEED. + disable_shuffle (bool): If True, keep the dataset in its original + order instead of shuffling it. + """ self.dataset_path = dataset_path # Set the random seed, ensuring that a None value is replaced with the @@ -134,8 +133,7 @@ def apply_multimodal_chat_transformation( prompt: str, mm_content: MultiModalDataDict | dict | list[dict] | None = None, ) -> list[dict]: - """ - Transform a prompt and optional multimodal content into a chat format. + """Transform a prompt and optional multimodal content into a chat format. This method is used for chat models that expect a specific conversation format. """ @@ -152,14 +150,14 @@ def apply_multimodal_chat_transformation( return [{"role": "user", "content": content}] def load_data(self) -> None: - """ - Load data from the dataset path into self.data. + """Load data from the dataset path into self.data. This method must be overridden by subclasses since the method to load data will vary depending on the dataset format and source. Raises: NotImplementedError: If a subclass does not implement this method. + """ # TODO (jenniferzhao): add support for downloading data raise NotImplementedError("load_data must be implemented in subclasses.") @@ -169,8 +167,7 @@ def get_random_lora_request( max_loras: int | None = None, lora_path: str | None = None, ) -> LoRARequest | None: - """ - Optionally select a random LoRA request. + """Optionally select a random LoRA request. This method is used when LoRA parameters are provided. It randomly selects a LoRA based on max_loras. @@ -184,6 +181,7 @@ def get_random_lora_request( Returns: A new [`LoRARequest`][vllm.lora.request.LoRARequest] (or `None` if not applicable). + """ if max_loras is None or lora_path is None: return None @@ -203,8 +201,7 @@ def get_round_robin_lora_request( max_loras: int | None = None, lora_path: str | None = None, ) -> LoRARequest | None: - """ - Optionally select a LoRA request using deterministic round-robin. + """Optionally select a LoRA request using deterministic round-robin. This method cycles through LoRA IDs in order based on the request index, providing reproducible LoRA assignment. @@ -219,6 +216,7 @@ def get_round_robin_lora_request( Returns: A new [`LoRARequest`][vllm.lora.request.LoRARequest] (or `None` if not applicable). + """ if max_loras is None or lora_path is None: return None @@ -239,8 +237,7 @@ def get_lora_request( lora_path: str | None = None, lora_assignment: str = "random", ) -> LoRARequest | None: - """ - Select a LoRA request using the specified assignment strategy. + """Select a LoRA request using the specified assignment strategy. Args: index (int): The request index (used for round-robin). @@ -252,6 +249,7 @@ def get_lora_request( Returns: A new [`LoRARequest`][vllm.lora.request.LoRARequest] (or `None` if not applicable). + """ if lora_assignment == "round-robin": return BenchmarkDataset.get_round_robin_lora_request( @@ -270,8 +268,7 @@ def sample( no_oversample: bool = False, **kwargs, ) -> list[SampleRequest]: - """ - Abstract method to generate sample requests from the dataset. + """Abstract method to generate sample requests from the dataset. Subclasses must override this method to implement dataset-specific logic for generating a list of SampleRequest objects. @@ -281,10 +278,13 @@ def sample( for processing the dataset's text. num_requests (int): The number of sample requests to generate. request_id_prefix (str): The prefix of request_id. + no_oversample (bool): If True, do not oversample the dataset to + reach num_requests. Returns: list[SampleRequest]: A list of sample requests generated from the dataset. + """ raise NotImplementedError("sample must be implemented in subclasses.") @@ -295,8 +295,7 @@ def maybe_oversample_requests( request_id_prefix: str = "", no_oversample: bool = False, ) -> None: - """ - Oversamples the list of requests if its size is less than the desired + """Oversamples the list of requests if its size is less than the desired number. Args: @@ -305,6 +304,8 @@ def maybe_oversample_requests( num_requests (int): The target number of requests. request_id_prefix (str): The prefix applied to generated request identifiers. + no_oversample (bool): If True, return the requests unchanged + instead of oversampling up to num_requests. """ if no_oversample: @@ -346,8 +347,7 @@ def is_valid_sequence( max_total_len: int = 2048, skip_min_output_len_check: bool = False, ) -> bool: - """ - Validate a sequence based on prompt and output lengths. + """Validate a sequence based on prompt and output lengths. Default pruning criteria are copied from the original `sample_hf_requests` and `sample_sharegpt_requests` functions in benchmark_serving.py, as well as @@ -379,8 +379,7 @@ def process_image( *, ensure_client_side_data: bool = False, ) -> Mapping[str, Any]: - """ - Process a single image input and return a multimedia content dictionary. + """Process a single image input and return a multimedia content dictionary. Supports the following input types: @@ -402,6 +401,7 @@ def process_image( Raises: ValueError: If the input is not a supported type. + """ if isinstance(image, dict) and "bytes" in image: image = Image.open(BytesIO(image["bytes"])) @@ -437,8 +437,7 @@ def process_image( def process_video(video: Any) -> Mapping[str, Any]: - """ - Process a single video input and return a multimedia content dictionary. + """Process a single video input and return a multimedia content dictionary. Supports the following input types: @@ -451,6 +450,7 @@ def process_video(video: Any) -> Mapping[str, Any]: Raises: ValueError: If the input is not a supported type. + """ if isinstance(video, dict) and "bytes" in video: video_bytes = video["bytes"] @@ -474,8 +474,7 @@ def process_video(video: Any) -> Mapping[str, Any]: def process_audio(audio: Any) -> tuple: - """ - Process a single audio input and return a (array, sample_rate) tuple. + """Process a single audio input and return a (array, sample_rate) tuple. Supports: 1. String: treated as a file path, loaded with soundfile. @@ -502,8 +501,7 @@ def gen_prompt_decode_to_target_len( add_special_tokens: bool = False, rng: np.random.Generator | None = None, ) -> tuple[str, list[int], int]: - """ - Ensure decoded-then-encoded prompt length matches the target token length. + """Ensure decoded-then-encoded prompt length matches the target token length. This function decodes an initial token sequence to text and re-encodes it , iteratively adjusting the token sequence length to match a target. @@ -557,8 +555,7 @@ def gen_prompt_decode_to_target_len( class RandomDataset(BenchmarkDataset): - """ - Synthetic text-only dataset for serving/throughput benchmarks. + """Synthetic text-only dataset for serving/throughput benchmarks. Strategy: - Sample input/output token lengths per request from integer-uniform ranges @@ -699,9 +696,7 @@ def get_prefix( allowed_tokens: np.ndarray, prefix_len: int, ) -> list[int]: - """ - Get the prefix for the dataset. - """ + """Get the prefix for the dataset.""" if prefix_len <= 0: return [] @@ -738,8 +733,7 @@ def generate_token_sequence( index: int, allowed_tokens: np.ndarray, ) -> tuple[str, int, int]: - """ - Returns (prompt, total_input_len). + """Returns (prompt, total_input_len). NOTE: After decoding the prompt we have to encode and decode it again. This is done because in some cases N consecutive tokens @@ -778,8 +772,7 @@ def generate_token_sequence( class RandomDatasetForReranking(RandomDataset): - """ - Random dataset specialized for the needs of scoring: + """Random dataset specialized for the needs of scoring: - Batches of inputs - Inputs composed of pairs """ @@ -906,8 +899,7 @@ def sample( class RandomMultiModalDataset(RandomDataset): - """ - Synthetic multimodal dataset (text + images) that extends RandomDataset. + """Synthetic multimodal dataset (text + images) that extends RandomDataset. Status: - Images: supported via synthetic RGB data. @@ -1022,8 +1014,7 @@ def map_config_to_modality(self, config: tuple[int, int, int]) -> str: def normalize_bucket_config( self, bucket_config: dict[tuple[int, int, int], float] ) -> dict[tuple[int, int, int], float]: - """ - Remove zero probability entries + """Remove zero probability entries and normalize the bucket config to sum to 1. """ # Raise error if value is negative @@ -1044,13 +1035,11 @@ def generate_mm_item( self, mm_item_config: tuple[int, int, int], ) -> Mapping[str, Any]: - """ - Create synthetic images and videos and + """Create synthetic images and videos and apply process_image/process_video respectively. This follows the OpenAI API chat completions https://github.com/openai/openai-python """ - if self.map_config_to_modality(mm_item_config) == "image": return process_image( self.generate_synthetic_image(mm_item_config[1], mm_item_config[0]) @@ -1071,9 +1060,7 @@ def get_mm_item_sampling_params( limit_mm_per_prompt: dict[str, int], bucket_config: dict[tuple[int, int, int], float], ) -> tuple[int, int, dict[str, int], dict[tuple[int, int, int], float]]: - """ - Get the sampling parameters for the multimodal items. - """ + """Get the sampling parameters for the multimodal items.""" # Enforce num_mm_items_range_ratio <= 1 if not (0.0 <= num_mm_items_range_ratio <= 1.0): raise ValueError("num_mm_items_range_ratio must be in [0, 1].") @@ -1146,8 +1133,7 @@ def get_mm_item_iterator( bucket_config: dict[tuple[int, int, int], float], limit_mm_per_prompt: dict[str, int], ) -> Iterator[tuple[int, int, int]]: - """ - Iterator over the multimodal items for each request + """Iterator over the multimodal items for each request whose size is between min_num_mm_items and max_num_mm_items. Loop over the bucket config and sample a multimodal item. @@ -1160,6 +1146,7 @@ def get_mm_item_iterator( `bucket_config` (tuple->float). The original dict passed to `sample` is not mutated. If this ever changes, a test is implemented and will fail. + """ # Get the number of multimodal items to sample request_num_mm_items = int( @@ -1346,8 +1333,7 @@ def sample( class ShareGPTDataset(BenchmarkDataset): - """ - Implements the ShareGPT dataset. Loads data from a JSON file and generates + """Implements the ShareGPT dataset. Loads data from a JSON file and generates sample requests based on conversation turns. """ @@ -1437,8 +1423,7 @@ def sample( class TimedTrace(BenchmarkDataset): - """ - Implements a base class to replay various timed traces. + """Implements a base class to replay various timed traces. Loads data from a JSON file and generates sample requests based on the timing information in the traces. """ @@ -1927,6 +1912,7 @@ def add_random_dataset_base_args( Args: parser_or_group: Either a parser or an argument group to add arguments to. + """ parser_or_group.add_argument( "--random-input-len", @@ -1988,6 +1974,7 @@ def add_random_multimodal_dataset_args( Args: parser_or_group: Either a parser or an argument group to add arguments to. + """ parser_or_group.add_argument( "--random-mm-base-items-per-request", @@ -2503,8 +2490,7 @@ def get_samples( class CustomDataset(BenchmarkDataset): - """ - Implements the Custom dataset. Loads data from a JSONL file and generates + """Implements the Custom dataset. Loads data from a JSONL file and generates sample requests based on conversation turns. E.g., ``` {"prompt": "What is the capital of India?", "output_tokens": 10} @@ -2632,8 +2618,7 @@ def sample( class CustomImageDataset(CustomDataset): - """ - Implements the Custom image dataset. Loads data from a JSONL file and generates + """Implements the Custom image dataset. Loads data from a JSONL file and generates sample requests based on conversation turns. E.g., ``` { @@ -2917,8 +2902,7 @@ def sample( class CustomAudioDataset(CustomDataset): - """ - Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g., + """Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g., {"prompt": "Transcribe the audio.", "audio": "/path/to/audio.wav"} Supports both: @@ -3027,8 +3011,7 @@ def sample( class SpecBench(CustomDataset): - """ - Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench + """Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench Download the dataset using: wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl """ # noqa: E501 @@ -3100,8 +3083,7 @@ def sample( "SonnetDataset is deprecated and will be removed in a future version.", ) class SonnetDataset(BenchmarkDataset): - """ - Simplified implementation of the Sonnet dataset. Loads poem lines from a + """Simplified implementation of the Sonnet dataset. Loads poem lines from a text file and generates sample requests. Default values here copied from `benchmark_serving.py` for the sonnet dataset. """ @@ -3198,8 +3180,7 @@ def sample( class BurstGPTDataset(BenchmarkDataset): - """ - Implements the BurstGPT dataset. Loads data from a CSV file and generates + """Implements the BurstGPT dataset. Loads data from a CSV file and generates sample requests based on synthetic prompt generation. Only rows with Model "GPT-4" and positive response tokens are used. """ @@ -3447,9 +3428,7 @@ def sample( class VisionArenaDataset(HuggingFaceDataset): - """ - Vision Arena Dataset. - """ + """Vision Arena Dataset.""" DEFAULT_OUTPUT_LEN = 128 SUPPORTED_DATASET_PATHS = { @@ -3506,8 +3485,7 @@ def sample( class MMVUDataset(HuggingFaceDataset): - """ - MMVU Dataset. + """MMVU Dataset. https://huggingface.co/datasets/yale-nlp/MMVU """ @@ -3586,8 +3564,7 @@ def sample( class InstructCoderDataset(HuggingFaceDataset): - """ - InstructCoder Dataset. + """InstructCoder Dataset. https://huggingface.co/datasets/likaixin/InstructCoder InstructCoder is the dataset designed for general code editing. It consists @@ -3655,8 +3632,7 @@ def sample_prompts(self, n: int) -> Iterator[str]: class MTBenchDataset(HuggingFaceDataset): - """ - MT-Bench Dataset. + """MT-Bench Dataset. https://huggingface.co/datasets/philschmid/mt-bench We create a single turn dataset for MT-Bench. @@ -3718,8 +3694,7 @@ def sample( class HumanEvalDataset(HuggingFaceDataset): - """ - HumanEvalDataset Dataset. + """HumanEvalDataset Dataset. https://huggingface.co/datasets/openai/openai_humaneval We create a single turn dataset for HumanEval. @@ -3779,8 +3754,7 @@ def sample( class GSM8KDataset(HuggingFaceDataset): - """ - GSM8K Dataset. + """GSM8K Dataset. https://huggingface.co/datasets/openai/gsm8k We create a single turn dataset for GSM8K. @@ -3840,8 +3814,7 @@ def sample( class BlazeditDataset(HuggingFaceDataset): - """ - Blazedit Dataset. + """Blazedit Dataset. https://github.com/ise-uiuc/blazedit 5k char version: vdaita/edit_5k_char @@ -3932,9 +3905,7 @@ def sample( class AIMODataset(HuggingFaceDataset): - """ - Dataset class for processing a AIMO dataset with reasoning questions. - """ + """Dataset class for processing a AIMO dataset with reasoning questions.""" SUPPORTED_DATASET_PATHS = { "AI-MO/aimo-validation-aime", @@ -4026,6 +3997,7 @@ def _format_zeta_prompt( Returns: A dictionary with the formatted prompts and expected outputs. + """ events = sample["events"] input = sample["input"] @@ -4042,9 +4014,7 @@ def _format_zeta_prompt( class NextEditPredictionDataset(HuggingFaceDataset): - """ - Dataset class for processing a Next Edit Prediction dataset. - """ + """Dataset class for processing a Next Edit Prediction dataset.""" SUPPORTED_DATASET_PATHS = { "zed-industries/zeta", @@ -4092,8 +4062,7 @@ def sample( class ASRDataset(HuggingFaceDataset): - """ - Dataset class for processing a ASR dataset for transcription. + """Dataset class for processing a ASR dataset for transcription. Tested on the following set: +---------------------------+----------------------------------------+--------------------------+-----------------------------+ @@ -4290,8 +4259,7 @@ def sample( class MLPerfDataset(HuggingFaceDataset): - """ - MLPerf Inference Dataset. + """MLPerf Inference Dataset. Dataset on HF: https://huggingface.co/datasets/mgoin/mlperf-inference-llama2-data @@ -4415,7 +4383,8 @@ def sample( def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]: """Generate tokens that decode and re-encode to exactly - target_length.""" + target_length. + """ # Generate random tokens tokens = np.random.randint(0, vocab_size, size=target_length).tolist() @@ -4470,8 +4439,7 @@ def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]: class MMStarDataset(HuggingFaceDataset): - """ - Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar + """Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar refer to: https://github.com/sgl-project/SpecForge/pull/106 """ @@ -4760,8 +4728,7 @@ def sample( class SpeedBench(CustomDataset): - """ - SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench + """SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench. Download the dataset using: diff --git a/vllm/benchmarks/datasets/utils.py b/vllm/benchmarks/datasets/utils.py index bc5a4340dd62..a2dd5093c50b 100644 --- a/vllm/benchmarks/datasets/utils.py +++ b/vllm/benchmarks/datasets/utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared utilities for benchmark dataset sampling. -""" +"""Shared utilities for benchmark dataset sampling.""" import logging import math @@ -46,8 +44,7 @@ def get_sampling_params( output_len: int, tokenizer: TokenizerLike, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """ - Sample per-request input/output token lengths and vocab offsets. + """Sample per-request input/output token lengths and vocab offsets. Lengths are drawn uniformly from integer ranges around the configured means, controlled by *range_ratio*. It may be a single ``float`` @@ -60,6 +57,7 @@ def get_sampling_params( Returns: (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of shape ``(num_requests,)``. + """ input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio) diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 54c6ce63b0fa..5adb54c47f16 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -22,7 +22,8 @@ class StreamedResponseHandler: """Handles streaming HTTP responses by accumulating chunks until complete - messages are available.""" + messages are available. + """ def __init__(self): self.buffer = "" @@ -30,7 +31,8 @@ def __init__(self): def add_chunk(self, chunk_bytes: bytes) -> list[str]: """Add a chunk of bytes to the buffer and return any complete - messages.""" + messages. + """ chunk_str = self._decoder.decode(chunk_bytes) self.buffer += chunk_str @@ -167,10 +169,12 @@ async def async_request_openai_completions( Args: request_func_input: The input for the request function. + session: The aiohttp session used to issue the request. pbar: The progress bar to display the progress. Returns: The output of the request function. + """ api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Completions API", "completions") diff --git a/vllm/benchmarks/lib/ready_checker.py b/vllm/benchmarks/lib/ready_checker.py index eec4a42cb670..e3b4d8ee3ea7 100644 --- a/vllm/benchmarks/lib/ready_checker.py +++ b/vllm/benchmarks/lib/ready_checker.py @@ -22,12 +22,12 @@ async def wait_for_endpoint( timeout_seconds: int = 600, retry_interval: int = 5, ) -> RequestFuncOutput: - """ - Wait for an endpoint to become available before starting benchmarks. + """Wait for an endpoint to become available before starting benchmarks. Args: request_func: The async request function to call test_input: The RequestFuncInput to test with + session: The aiohttp session used to issue the probe requests timeout_seconds: Maximum time to wait in seconds (default: 10 minutes) retry_interval: Time between retries in seconds (default: 5 seconds) @@ -36,6 +36,7 @@ async def wait_for_endpoint( Raises: ValueError: If the endpoint doesn't become available within the timeout + """ deadline = time.perf_counter() + timeout_seconds output = RequestFuncOutput(success=False) diff --git a/vllm/benchmarks/lib/utils.py b/vllm/benchmarks/lib/utils.py index 22409101e443..f5bd11439ab4 100644 --- a/vllm/benchmarks/lib/utils.py +++ b/vllm/benchmarks/lib/utils.py @@ -36,9 +36,7 @@ def extract_field( def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool: - """ - Check if the benchmark is run with torch.compile - """ + """Check if the benchmark is run with torch.compile.""" return not ( extract_field(args, extra_info, "compilation_config.mode") == "0" or "eager" in getattr(args, "output_json", "") @@ -49,8 +47,7 @@ def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool: def convert_to_pytorch_benchmark_format( args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any] ) -> list: - """ - Save the benchmark results in the format used by PyTorch OSS benchmark with + """Save the benchmark results in the format used by PyTorch OSS benchmark with on metric per record https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database """ diff --git a/vllm/benchmarks/mm_processor.py b/vllm/benchmarks/mm_processor.py index 4c840786fe72..40b3494151d3 100644 --- a/vllm/benchmarks/mm_processor.py +++ b/vllm/benchmarks/mm_processor.py @@ -44,8 +44,7 @@ def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, float]]: - """ - Get all multimodal timing stats from the LLM engine. + """Get all multimodal timing stats from the LLM engine. Collects both preprocessing stats (HF processor, hashing, cache lookup, prompt update) and encoder forward pass timing, merged by request_id. @@ -70,6 +69,7 @@ def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, f 'num_encoder_calls': 1 } } + """ observability_config = llm_engine.vllm_config.observability_config if not observability_config or not observability_config.enable_mm_processor_stats: @@ -126,8 +126,7 @@ def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, f def collect_mm_processor_stats(llm_engine: LLMEngine) -> dict[str, list[float]]: - """ - Collect multimodal processor timing stats. + """Collect multimodal processor timing stats. Returns a dictionary mapping stage names to lists of timing values (in seconds). """ all_stats = get_timing_stats_from_engine(llm_engine) @@ -147,9 +146,7 @@ def calculate_mm_processor_metrics( *, unit: Literal["us", "ms", "s"] = "ms", ) -> dict[str, dict[str, float]]: - """ - Calculate aggregate metrics from stats by stage. - """ + """Calculate aggregate metrics from stats by stage.""" unit2mult = {"us": 1000000, "ms": 1000, "s": 1} unit_mult = unit2mult[unit] @@ -181,9 +178,7 @@ def calculate_mm_processor_metrics( def validate_args(args): - """ - Validate command-line arguments for mm_processor benchmark. - """ + """Validate command-line arguments for mm_processor benchmark.""" if not getattr(args, "tokenizer", None): args.tokenizer = args.model if not hasattr(args, "dataset_path"): @@ -214,9 +209,7 @@ def validate_args(args): def benchmark_multimodal_processor( args: argparse.Namespace, ) -> dict[str, Any]: - """ - Run the multimodal processor benchmark. - """ + """Run the multimodal processor benchmark.""" from vllm import LLM, SamplingParams validate_args(args) @@ -455,7 +448,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None: def main(args: argparse.Namespace) -> None: """Main entry point for the multimodal processor benchmark.""" - print("Starting multimodal processor benchmark...") result = benchmark_multimodal_processor(args) diff --git a/vllm/benchmarks/plot.py b/vllm/benchmarks/plot.py index 3f36ede721ad..450019e80e25 100644 --- a/vllm/benchmarks/plot.py +++ b/vllm/benchmarks/plot.py @@ -29,8 +29,7 @@ def generate_timeline_plot( itl_thresholds: list[float] | None = None, labels: list[str] | None = None, ) -> None: - """ - Generate an HTML timeline plot from benchmark results. + """Generate an HTML timeline plot from benchmark results. Args: results: List of per-request result dictionaries containing: @@ -44,8 +43,8 @@ def generate_timeline_plot( colors: List of colors for ITL categories (default: green, orange, red, black) itl_thresholds: ITL thresholds in seconds (default: [1.0, 4.0, 6.0]) labels: Labels for ITL categories (default based on thresholds) - """ + """ # Set defaults if colors is None: colors = ["#109618", "#FF7F0E", "#D62728"] @@ -118,8 +117,7 @@ def construct_timeline_data( itl_thresholds: list[float], labels: list[str], ) -> list[dict[str, Any]]: - """ - Construct timeline data from request results. + """Construct timeline data from request results. Args: requests_data: List of per-request result dictionaries @@ -128,6 +126,7 @@ def construct_timeline_data( Returns: List of timeline segments for plotting + """ def tostr(sec_time: float) -> str: @@ -229,8 +228,7 @@ def generate_dataset_stats_plot( results: list[dict[str, Any]], output_path: Path, ) -> None: - """ - Generate a matplotlib figure with dataset statistics. + """Generate a matplotlib figure with dataset statistics. Creates a figure with 4 subplots: - Top-left: Prompt tokens distribution (histogram) @@ -243,6 +241,7 @@ def generate_dataset_stats_plot( - prompt_len: Number of prompt tokens - output_tokens: Number of output tokens output_path: Path where the figure will be saved + """ # Extract data prompt_tokens = [] diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 0292ff939a55..14c72386a5d0 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -407,13 +407,15 @@ async def get_request( ramp_up_end_rps: int | None = None, self_timed: bool = False, ) -> AsyncGenerator[tuple[SampleRequest, float], None]: - """ - Asynchronously generates requests at a specified rate + """Asynchronously generates requests at a specified rate with OPTIONAL burstiness and OPTIONAL ramp-up strategy. Args: input_requests: A list of input requests, each represented as a SampleRequest. + self_timed: + If True, the requests carry their own arrival timing and no + request rate, burstiness or ramp-up is applied. request_rate: The rate at which requests are generated (requests/s). burstiness (optional): @@ -431,6 +433,7 @@ async def get_request( The starting request rate for ramp-up. ramp_up_end_rps (optional): The ending request rate for ramp-up. + """ assert burstiness > 0, ( f"A positive burstiness factor is expected, but given {burstiness}." @@ -526,6 +529,7 @@ def calculate_metrics_for_embeddings( Returns: The calculated benchmark metrics. + """ total_input = 0 total_input_sequences = 0 @@ -585,6 +589,7 @@ def calculate_metrics( Returns: A tuple of the benchmark metrics and the actual output lengths. + """ actual_output_lens: list[int] = [] total_input = 0 @@ -1571,6 +1576,7 @@ def compute_result_filename( Returns: The computed filename path or None if no result saving is requested + """ if not (args.plot_timeline or args.save_result or args.append_result): return None diff --git a/vllm/benchmarks/startup.py b/vllm/benchmarks/startup.py index 86f87a32d666..1df66af622e1 100644 --- a/vllm/benchmarks/startup.py +++ b/vllm/benchmarks/startup.py @@ -107,8 +107,7 @@ def _metric_to_json(m: MetricStats) -> dict[str, Any]: @contextmanager def cold_startup(): - """ - Context manager to measure cold startup time: + """Context manager to measure cold startup time: 1. Uses a temporary directory for vLLM cache to avoid any pollution between cold startup iterations. 2. Uses inductor's fresh_cache to clear torch.compile caches. @@ -132,8 +131,7 @@ def cold_startup(): def run_startup_in_subprocess(engine_args, result_queue): - """ - Run LLM startup in a subprocess and return timing metrics via a queue. + """Run LLM startup in a subprocess and return timing metrics via a queue. This ensures complete isolation between iterations. """ try: @@ -229,11 +227,9 @@ def main(args: argparse.Namespace): engine_args = EngineArgs.from_cli_args(args) def create_llm_and_measure_startup(): - """ - Create LLM instance in a subprocess and measure startup time. + """Create LLM instance in a subprocess and measure startup time. Returns timing metrics, using subprocess for complete isolation. """ - # Create a queue for inter-process communication result_queue: multiprocessing.Queue[Any] = multiprocessing.Queue() process = multiprocessing.Process( diff --git a/vllm/benchmarks/sweep/param_sweep.py b/vllm/benchmarks/sweep/param_sweep.py index 5c2620e794ec..c7607109483d 100644 --- a/vllm/benchmarks/sweep/param_sweep.py +++ b/vllm/benchmarks/sweep/param_sweep.py @@ -18,14 +18,14 @@ def read_json(cls, filepath: os.PathLike): @classmethod def read_from_dict(cls, data: dict[str, dict[str, object]]): - """ - Read parameter sweep from a dict format where keys are names. + """Read parameter sweep from a dict format where keys are names. Example: { "experiment1": {"max_tokens": 100, "temperature": 0.7}, "experiment2": {"max_tokens": 200, "temperature": 0.9} } + """ records = [{"_benchmark_name": name, **params} for name, params in data.items()] return cls.from_records(records) @@ -66,8 +66,7 @@ def __or__(self, other: dict[str, object], /) -> "ParameterSweepItem": # type: @property def name(self) -> str: - """ - Get the name for this parameter sweep item. + """Get the name for this parameter sweep item. Returns the '_benchmark_name' field if present, otherwise returns a text representation of all parameters. @@ -103,8 +102,7 @@ def has_param(self, param_key: str) -> bool: return any(k in self for k in self._iter_param_key_candidates(param_key)) def _normalize_cmd_kv_pair(self, k: str, v: object) -> list[str]: - """ - Normalize a key-value pair into command-line arguments. + """Normalize a key-value pair into command-line arguments. Returns a list containing either: - A single element for boolean flags (e.g., ['--flag'] or ['--flag=true']) diff --git a/vllm/benchmarks/sweep/plot.py b/vllm/benchmarks/sweep/plot.py index 79e50603b8c8..59a90636f0e0 100644 --- a/vllm/benchmarks/sweep/plot.py +++ b/vllm/benchmarks/sweep/plot.py @@ -178,8 +178,7 @@ def _json_load_bytes(path: Path) -> list[dict[str, object]]: def _convert_inf_nan_strings(data: list[dict[str, object]]) -> list[dict[str, object]]: - """ - Convert string values "inf", "-inf", and "nan" to their float equivalents. + """Convert string values "inf", "-inf", and "nan" to their float equivalents. This handles the case where JSON serialization represents inf/nan as strings. """ diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 45fdcd9694f0..94b22b29567c 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -195,8 +195,7 @@ def run_vllm_chat( warmup_requests: list[SampleRequest] | None = None, prequeue_requests: bool = False, ) -> tuple[float, list[RequestOutput]]: - """ - Run vLLM chat benchmark. This function is recommended ONLY for benchmarking + """Run vLLM chat benchmark. This function is recommended ONLY for benchmarking multimodal models as it properly handles multimodal inputs and chat formatting. For non-multimodal models, use run_vllm() instead. """ @@ -536,6 +535,7 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace: Returns: A namespace satisfying get_samples's attribute reads. + """ d = vars(args).copy() # random_*: prefer --random-* over legacy --input/output/prefix-len. @@ -627,10 +627,7 @@ def assign_loras(requests, args): def validate_args(args): - """ - Validate command-line arguments. - """ - + """Validate command-line arguments.""" # === Deprecation and Defaulting === if args.dataset is not None: warnings.warn( diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 439a87e1522a..fb981a3600f0 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -73,6 +73,7 @@ def make_copy_and_call( Returns: A wrapper function that copies inputs and calls the compiled function + """ def copy_and_call(*args: Any) -> Any: @@ -122,8 +123,7 @@ def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface: class CompilerManager: - """ - A manager to manage the compilation process, including + """A manager to manage the compilation process, including caching the compiled graph, loading the compiled graph, and compiling the graph. @@ -150,7 +150,8 @@ def compute_hash(self, vllm_config: VllmConfig) -> str: def compile_context(self, compile_range: Range) -> Generator[None, None, None]: """Provide compilation context for the duration of compilation to set any torch global properties we want to scope to a single Inductor - compilation (e.g. partition rules, pass context).""" + compilation (e.g. partition rules, pass context). + """ with pass_context(compile_range): if self.compilation_config.use_inductor_graph_partition: with inductor_partition_rule_context( @@ -163,8 +164,7 @@ def compile_context(self, compile_range: Range) -> Generator[None, None, None]: def initialize_cache( self, cache_dir: str, disable_cache: bool = False, prefix: str = "" ) -> None: - """ - Initialize the cache directory for the compiler. + """Initialize the cache directory for the compiler. The organization of the cache directory is as follows: cache_dir=/path/to/hash_str/rank_i_j/prefix/ @@ -177,7 +177,6 @@ def initialize_cache( base cache dir of /path/to/hash_str/rank_i_j/ , to store some common compilation artifacts. """ - self.disable_cache = disable_cache self.cache_dir = cache_dir self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py") @@ -438,12 +437,10 @@ def _merge_empty_only_subgraphs( node_to_subgraph_id: dict[fx.Node, int], split_op_graphs: list[int], ) -> None: - """ - Merge a partition that only contains an empty allocation op into the + """Merge a partition that only contains an empty allocation op into the previous partition. This avoids generating standalone empty submodules, which can lead to empty cudagraph captures. """ - nodes_by_subgraph_id: dict[int, list[fx.Node]] = defaultdict(list) for node, subgraph_id in node_to_subgraph_id.items(): nodes_by_subgraph_id[subgraph_id].append(node) @@ -637,8 +634,7 @@ def wrap_with_cudagraph_if_needed( is_first_graph: bool, is_last_graph: bool, ) -> Any: - """ - Wrap a piecewise backend with CUDA graph wrapper if needed. + """Wrap a piecewise backend with CUDA graph wrapper if needed. This function is shared between VllmBackend and construct_serializable_fn_from_inductor_cache. @@ -651,6 +647,7 @@ def wrap_with_cudagraph_if_needed( Returns: The wrapped backend if CUDA graphs are enabled, otherwise the original backend + """ if ( not compilation_config.cudagraph_mode.has_piecewise_cudagraphs() @@ -883,8 +880,8 @@ def collect_standalone_compile_artifacts( sym_shape_indices - returns_tuple_map: dict mapping submod_name to returns_tuple - """ + """ if not envs.VLLM_USE_MEGA_AOT_ARTIFACT: return None, None, None diff --git a/vllm/compilation/base_static_graph.py b/vllm/compilation/base_static_graph.py index 12f1ff5bc044..b645af104313 100644 --- a/vllm/compilation/base_static_graph.py +++ b/vllm/compilation/base_static_graph.py @@ -8,8 +8,7 @@ class AbstractStaticGraphWrapper(Protocol): - """ - StaticGraphWrapper interface that allows platforms to wrap a callable + """StaticGraphWrapper interface that allows platforms to wrap a callable to be captured as a static graph. """ @@ -20,8 +19,7 @@ def __init__( runtime_mode: CUDAGraphMode, **kwargs: Any, ) -> None: - """ - Initializes the StaticGraphWrapper class with graph capturing and + """Initializes the StaticGraphWrapper class with graph capturing and execution-related configurations. Args: @@ -31,15 +29,16 @@ def __init__( graph runtime. See CUDAGraphMode in vllm/config.py. Note that only the subset enum `NONE`, `PIECEWISE` and `FULL` are used as concrete runtime mode for cudagraph dispatching. + Keyword Args: kwargs: Additional keyword arguments for platform-specific configurations. + """ raise NotImplementedError def __call__(self, *args: Any, **kwargs: Any) -> Any: - """ - Executes the wrapped callable. + """Executes the wrapped callable. If the current runtime mode in the ForwardContext matches the runtime mode of this instance, it replays the CUDAGraph or captures it using @@ -53,5 +52,6 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: Returns: Any: Output of the executed callable. + """ raise NotImplementedError diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py index 94456fc18271..5205c7404fcb 100644 --- a/vllm/compilation/caching.py +++ b/vllm/compilation/caching.py @@ -172,8 +172,7 @@ def patch_pytree_map_over_slice(): class VllmSerializableFunction(SerializableCallable): # type: ignore[misc] - """ - A wrapper around a compiled function by vllm. It will forward the tensor + """A wrapper around a compiled function by vllm. It will forward the tensor inputs to the compiled function and return the result. It also implements a serialization interface to support PyTorch's precompile with custom backend, so that we can save and load the compiled function on @@ -410,9 +409,7 @@ def finalize_loading(self, vllm_config: VllmConfig) -> None: @property def co_name(self) -> Literal["VllmSerializableFunction"]: - """ - Used for depyf debugging. - """ + """Used for depyf debugging.""" return "VllmSerializableFunction" @@ -445,6 +442,7 @@ def reconstruct_serializable_fn_from_mega_artifact( If modifying the backend creation/wrapping logic, consider updating both. Args: + fake_mode: FakeTensorMode used to unpickle the cached graph. state: Deserialized state dict containing graph_module, example_inputs, prefix, sym_tensor_indices, is_encoder, etc. standalone_compile_artifacts: The StandaloneCompiledArtifacts containing @@ -455,6 +453,7 @@ def reconstruct_serializable_fn_from_mega_artifact( Returns: A VllmSerializableFunction that can be called directly. + """ from vllm.compilation.backends import ( VllmBackend, diff --git a/vllm/compilation/codegen.py b/vllm/compilation/codegen.py index 67e9ac843a64..a0eb90d34dc7 100644 --- a/vllm/compilation/codegen.py +++ b/vllm/compilation/codegen.py @@ -154,6 +154,7 @@ def generate_execution_code( consts is a list of non-primitive constant objects referenced by the generated code via __vllm_consts__. These objects are kept alive for the lifetime of the compiled function. + """ code, submod_names, consts = generate_execution_code_with_name( split_gm, "execution_fn", with_submod=True @@ -181,6 +182,7 @@ def compile_execution_fn( Returns: A callable that executes the stitching logic. + """ trace_structured( "artifact", diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 742ec55e6efa..0195a0879923 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -25,9 +25,7 @@ class CompilerInterface: - """ - The interface for a compiler that can be used by vLLM. - """ + """The interface for a compiler that can be used by vLLM.""" # The name of the compiler, e.g. inductor. # This is a class-level attribute. @@ -36,8 +34,7 @@ class CompilerInterface: def initialize_cache( self, cache_dir: str, disable_cache: bool = False, prefix: str = "" ) -> None: - """ - when the vLLM process uses `cache_dir` as the cache directory, + """When the vLLM process uses `cache_dir` as the cache directory, the compiler should initialize itself with the cache directory, e.g. by re-directing its own cache directory to a sub-directory. @@ -52,8 +49,7 @@ def initialize_cache( pass def compute_hash(self, vllm_config: VllmConfig) -> str: - """ - Gather all the relevant information from the vLLM config, + """Gather all the relevant information from the vLLM config, to compute a hash so that we can cache the compiled model. See [`VllmConfig.compute_hash`][vllm.config.VllmConfig.compute_hash] @@ -71,8 +67,7 @@ def compile( compile_range: Range, key: str | None = None, ) -> tuple[Callable[..., Any] | None, Any | None]: - """ - Compile the graph with the given example inputs and compiler config, + """Compile the graph with the given example inputs and compiler config, with a range. The `compile_range` specifies the range of the inputs, it could be concrete size (if compile_sizes is provided), e.g. [4, 4] or a range [5, 8]. @@ -105,8 +100,7 @@ def load( graph_index: int, compile_range: Range, ) -> Callable[..., Any]: - """ - Load the compiled function from the handle. + """Load the compiled function from the handle. Raises an error if the handle is invalid. The handle is the second return value of the `compile` function. @@ -115,8 +109,7 @@ def load( class AlwaysHitShapeEnv: - """ - Why do we need this class: + """Why do we need this class: For normal `torch.compile` usage, every compilation will have one Dynamo bytecode compilation and one Inductor compilation. @@ -160,7 +153,8 @@ def _get_vllm_functorch_config() -> dict[str, Any]: """Return the functorch config overrides that vLLM applies at compile time. Used by both set_functorch_config() and get_inductor_factors() to ensure - the compile-time config and cache key are always consistent.""" + the compile-time config and cache key are always consistent. + """ cfg: dict[str, Any] = {} if not envs.VLLM_USE_MEGA_AOT_ARTIFACT: cfg["bundled_autograd_cache"] = False @@ -249,8 +243,7 @@ def _save( class InductorStandaloneAdaptor(CompilerInterface): - """ - The adaptor for the Inductor compiler. + """The adaptor for the Inductor compiler. Requires PyTorch 2.8+. This is not on by default yet, but we plan to turn it on by default for PyTorch 2.8. @@ -447,9 +440,7 @@ def compiled_graph_wrapper(*args: Any) -> tuple[Any, ...] | Any: class InductorAdaptor(CompilerInterface): - """ - The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7. - """ + """The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7.""" name = "inductor" @@ -727,8 +718,7 @@ def compiled_graph(*args: Any) -> tuple[Any, ...] | Any: return compiled_graph def metrics_context(self) -> contextlib.AbstractContextManager[Any]: - """ - This method returns the Dynamo metrics context (if it exists, + """This method returns the Dynamo metrics context (if it exists, otherwise a null context). It is used by various compile components. Present in torch>=2.6, it's used inside FxGraphCache in torch==2.6 (but not after). It might also be used in various other diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 55771608837d..1306cd612f4e 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -38,7 +38,7 @@ class CUDAGraphStat: class CUDAGraphLogging: - """Aggregate and log cudagraph metrics""" + """Aggregate and log cudagraph metrics.""" COLUMN_HEADERS = [ "Unpadded Tokens", diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index d5a2ac7e6813..242e7227e922 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -56,8 +56,7 @@ def should_torch_compile_mm_encoder(vllm_config: VllmConfig) -> bool: def ignore_torch_compile(cls: type[_T]) -> type[_T]: - """ - A decorator to ignore support_torch_compile decorator + """A decorator to ignore support_torch_compile decorator on the class. This is useful when a parent class has a support_torch_compile decorator, but we don't want to compile the class `cls` that inherits the parent class. @@ -76,9 +75,7 @@ def ignore_torch_compile(cls: type[_T]) -> type[_T]: def _should_ignore_torch_compile(cls: type[_T]) -> bool: - """ - Check if the class should be ignored for torch.compile. - """ + """Check if the class should be ignored for torch.compile.""" return getattr(cls, IGNORE_COMPILE_KEY, False) @@ -125,8 +122,7 @@ def support_torch_compile( enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, ) -> Callable[[type[_T]], type[_T]] | type[_T]: - """ - A decorator to add support for compiling the forward method of a class. + """A decorator to add support for compiling the forward method of a class. Usage 1: use directly as a decorator without arguments: @@ -338,7 +334,6 @@ def _support_torch_compile( is_encoder: bool = False, ) -> type[_T]: """Internal implementation of support_torch_compile decorator.""" - if TorchCompileWithNoGuardsWrapper in cls.__bases__: # support decorating multiple times return cls @@ -736,8 +731,7 @@ def save_aot_compiled_function(self: type[_T]) -> None: def maybe_use_cudagraph_partition_wrapper( vllm_config: VllmConfig, ) -> Generator[None, None, None]: - """ - Context manager to set/unset customized cudagraph partition wrappers. + """Context manager to set/unset customized cudagraph partition wrappers. If we're using Inductor-based graph partitioning, we currently have the whole `fx.Graph` before Inductor lowering and the piecewise diff --git a/vllm/compilation/partition_rules.py b/vllm/compilation/partition_rules.py index 18ebb15d1112..8c2e5d6d2b33 100644 --- a/vllm/compilation/partition_rules.py +++ b/vllm/compilation/partition_rules.py @@ -12,12 +12,10 @@ def should_split(node: torch.fx.Node, splitting_ops: list[str]) -> bool: - """ - Check if a node should be split for dynamo graph partition. + """Check if a node should be split for dynamo graph partition. It operates on dynamo graph, so the node.target can be anything. We need to check and split only on OpOverload and OpOverloadPacket. """ - if node.op != "call_function": return False @@ -50,6 +48,7 @@ def inductor_partition_rule_context( Args: splitting_ops: List of operator names to partition on. + """ if not splitting_ops: logger.debug("No partition ops provided; skipping rule registration.") diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index c58ce31bd29c..ffafb57d9d34 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -45,8 +45,7 @@ class ActivationQuantPattern(VllmPatternReplacement): - """ - Base class for Activation+Quant fusions. + """Base class for Activation+Quant fusions. Should not be used directly. """ @@ -79,9 +78,7 @@ def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor: class SiluMulFp8StaticQuantPattern(ActivationQuantPattern): - """ - Fusion for SiluMul+Fp8StaticQuant Pattern - """ + """Fusion for SiluMul+Fp8StaticQuant Pattern.""" def __init__(self) -> None: super().__init__(kFp8StaticTensorSym) @@ -126,9 +123,7 @@ def _replacement( class SiluMulNvfp4QuantPattern(ActivationQuantPattern): - """ - Fusion for SiluMul+Nvfp4Quant Pattern - """ + """Fusion for SiluMul+Nvfp4Quant Pattern.""" def __init__(self) -> None: super().__init__(kNvfp4Dynamic) @@ -182,8 +177,7 @@ def _replacement( class SiluMulBlockQuantPattern(ActivationQuantPattern): - """ - Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern. + """Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern. Supports group_size 128 and 64 via QuantKey. Parameterized on is_scale_transposed for different scale layouts. """ @@ -281,8 +275,7 @@ def _replacement( class ActivationQuantFusionPass(VllmFusionPatternMatcherPass): - """ - This pass fuses a pre-defined set of custom ops into fused ops. + """This pass fuses a pre-defined set of custom ops into fused ops. It uses the torch pattern matcher to find the patterns and replace them. Because patterns can only be registered once, the pass is a singleton. diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 5e7570ff0d81..bb41c9550d6f 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -66,7 +66,8 @@ def _view_flashinfer_nvfp4_scale_out_as_int32( def _norm_input_weight_dtype_match(match: pm.Match) -> bool: """Prevent fusion when the norm input and weight dtypes differ (e.g. a Gemma - fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm.""" + fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm. + """ for node in match.nodes: if node.target == _IR_RMS_NORM_OP: x, weight = node.args[0], node.args[1] @@ -347,8 +348,7 @@ def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: class AllReduceRMSNormPattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (without residual) + """This pattern replaces the allreduce + rms norm (without residual) with fused flashinfer implementation. Applies to allreduce + rmsnorm before attn in the first Transformer block. """ @@ -409,8 +409,7 @@ def replacement( class AllReduceFusedAddRMSNormPattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (with residual) + """This pattern replaces the allreduce + rms norm (with residual) with fused flashinfer implementation. Applies to o_proj + rmsnorm after attn and mlp + rmsnorm before attn. """ @@ -609,8 +608,7 @@ def replacement( class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (without residual) + """This pattern replaces the allreduce + rms norm (without residual) + static fp8 quant with fused flashinfer implementation. Applies to allreduce + rmsnorm + quant before attn in the first Transformer block. @@ -684,8 +682,7 @@ def replacement( class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (with residual) + """This pattern replaces the allreduce + rms norm (with residual) + static fp8 quant with fused flashinfer implementation. Applies to o_proj + rmsnorm after attn + quant and mlp + rmsnorm + quant before attn. @@ -767,8 +764,7 @@ def replacement( class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (without residual) + """This pattern replaces the allreduce + rms norm (without residual) + static nvfp4 quant with fused flashinfer implementation. Applies to allreduce + rmsnorm + quant before attn in the first Transformer block. @@ -864,8 +860,7 @@ def replacement( class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern): - """ - This pattern replaces the allreduce + rms norm (with residual) + """This pattern replaces the allreduce + rms norm (with residual) + static nvfp4 quant with fused flashinfer implementation. Applies to o_proj + rmsnorm after attn + quant and mlp + rmsnorm + quant before attn. diff --git a/vllm/compilation/passes/fusion/attn_quant_fusion.py b/vllm/compilation/passes/fusion/attn_quant_fusion.py index 3e2ed2bc707e..432d59e5aba9 100644 --- a/vllm/compilation/passes/fusion/attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/attn_quant_fusion.py @@ -36,8 +36,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): - """ - Fusion for Attention+Fp8StaticQuant. + """Fusion for Attention+Fp8StaticQuant. Only triggers when the attention implementation returns True in `fused_output_quant_supported()`. If the pattern is found, the @@ -172,8 +171,7 @@ def get_inputs(self): class AttnNvfp4QuantPattern( VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]] ): - """ - Fusion for Attention+Nvfp4Quant. + """Fusion for Attention+Nvfp4Quant. Only triggers when the attention implementation returns True in `fused_output_quant_supported()`. If the pattern is found, the @@ -360,8 +358,7 @@ def get_inputs(self): class AttnQuantFusionPass(VllmFusionPatternMatcherPass): - """ - This pass fuses post-attention quantization onto attention if supported. + """This pass fuses post-attention quantization onto attention if supported. It uses the pattern matcher and matches each layer manually, as strings cannot be wildcarded. This also lets us check support on attention layers diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index af2719ce3b2d..7fd44dd005f8 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -79,7 +79,7 @@ def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) def inputs(self) -> list[torch.Tensor]: - """Utility for inputs to the pattern""" + """Utility for inputs to the pattern.""" raise NotImplementedError diff --git a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py index 84c24bc60e53..982d90726e06 100644 --- a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py @@ -33,8 +33,7 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): - """ - Fusion for MLA Attention+Fp8StaticQuant. + """Fusion for MLA Attention+Fp8StaticQuant. Matches the pattern: MLA attention -> static FP8 quant, and replaces it with MLA attention(output_scale=scale, output=fp8_buffer). @@ -174,8 +173,7 @@ def get_inputs(self) -> list[torch.Tensor]: class MLAAttnNvfp4QuantPattern( VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]] ): - """ - Fusion for MLA Attention+Nvfp4Quant. + """Fusion for MLA Attention+Nvfp4Quant. Matches the pattern: MLA attention -> NVFP4 quant, and replaces it with MLA attention(output_scale=scale, output_block_scale=block_scale, @@ -362,8 +360,7 @@ def get_inputs(self) -> list[torch.Tensor]: class MLAAttnFp8GroupQuantPattern( VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]] ): - """ - Fusion for MLA Attention+Fp8GroupQuant (per-group dynamic FP8). + """Fusion for MLA Attention+Fp8GroupQuant (per-group dynamic FP8). Matches the pattern: MLA attention -> per_token_group_fp8_quant, and replaces it with MLA attention(output_block_scale=group_scale_buffer). @@ -572,8 +569,7 @@ def get_inputs(self) -> list[torch.Tensor]: class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass): - """ - This pass fuses post-attention quantization onto MLA attention if supported. + """This pass fuses post-attention quantization onto MLA attention if supported. It uses the pattern matcher and matches each MLA layer manually, as strings cannot be wildcarded. This also lets us check support on attention layers diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py index 2c8b9c43bdee..4841d255dbef 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py @@ -34,8 +34,7 @@ class QkNormRopePattern: - """ - Match the unfused sequence in attention blocks and replace with the fused op. + """Match the unfused sequence in attention blocks and replace with the fused op. Unfused (conceptually): q, k, v = split(qkv, [qsz, kvsz, kvsz], -1) diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py index 85d82c191b27..80eeb648f407 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py @@ -106,8 +106,7 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_fake( class QkNormRopeKvCachePattern: - """ - Match the unfused sequence: + """Match the unfused sequence: q, k, v = split(qkv, ...) q = rms_norm(q.view(heads), q_weight).view(flat) k = rms_norm(k.view(heads), k_weight).view(flat) @@ -413,8 +412,7 @@ def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache): class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass): - """ - Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel. + """Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel. Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for attention layers that support the combined operation, eliminating two diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 91c44f572cf0..90897a96435b 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -168,8 +168,7 @@ def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor: class FusedRMSQuantKey(NamedTuple): - """ - Named tuple for identifying the type of RMSNorm + quant fusion. + """Named tuple for identifying the type of RMSNorm + quant fusion. quant: type of quantization fused_add: does the op also perform the residual add """ @@ -771,8 +770,7 @@ def replacement( class RMSNormQuantFusionPass(VllmPatternMatcherPass): - """ - This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op. + """This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op. It also supports fused_add_rms_norm. """ diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index 67ed51e89322..3a097b0bf491 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -180,8 +180,7 @@ def replacement( class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): - """ - This pattern fuses aiter rms_norm & group fp8 quant custom + """This pattern fuses aiter rms_norm & group fp8 quant custom ops into an aiter rms_norm_group_fp8_quant op. """ @@ -236,8 +235,7 @@ def replacement( class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): - """ - This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops + """This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops into a aiter rms_norm_with_add_group_fp8_quant op. """ @@ -298,8 +296,7 @@ def replacement( class DoubleAiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern): - """ - Pattern matching ``rms_norm`` whose output feeds *two* distinct + """Pattern matching ``rms_norm`` whose output feeds *two* distinct ``rocm_aiter_group_fp8_quant`` consumers, replacing it with two independent fused ``rms_norm_group_fp8_quant`` ops. @@ -367,8 +364,7 @@ def replacement( class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern): - """ - View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``. + """View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``. Matches the same 1-to-2 fan-out, but with a ``view``/``reshape`` between the ``rms_norm`` output and the two ``rocm_aiter_group_fp8_quant`` @@ -453,8 +449,7 @@ def trace_with_view_to_reshape(*args: Any, **kwargs: Any) -> fx.GraphModule: class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern): - """ - Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces + """Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces with rocm_aiter_fused_rms_gated_fp8_group_quant. The norm operates per-head on (N*H, D) tensors. The compiler folds the @@ -544,8 +539,7 @@ def trace_fn(*args, **kwargs): class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): - """ - This pass fuses aiter rms_norm & vllm/aiter quant custom ops + """This pass fuses aiter rms_norm & vllm/aiter quant custom ops into a fused rms_norm_quant op. It also supports fused_add_rms_norm. """ @@ -681,8 +675,7 @@ def uuid(self) -> str: class AiterSiluMulFp8GroupQuantPattern(VllmPatternReplacement): - """ - This pattern fuses aiter silu_and_mul & group fp8 quant custom + """This pattern fuses aiter silu_and_mul & group fp8 quant custom ops into an aiter silu_and_mul_group_fp8_quant op. """ @@ -722,8 +715,7 @@ def _replacement( class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmFusionPatternMatcherPass): - """ - This pass fuses a pre-defined set of custom ops into fused ops. + """This pass fuses a pre-defined set of custom ops into fused ops. It uses the torch pattern matcher to find the patterns and replace them. Because patterns can only be registered once, the pass is a singleton. @@ -744,8 +736,7 @@ def __init__(self, config: VllmConfig) -> None: class AddAiterRMSNormPadPattern: - """ - This pattern replaces an aiter_rmsnorm_with_add & a pad op + """This pattern replaces an aiter_rmsnorm_with_add & a pad op with a custom triton_add_rmsnorm_pad op from AITER. """ @@ -820,8 +811,7 @@ def replacement( class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass): - """ - This pass replaces an AITER CK RMSNorm + residual add and a pad op + """This pass replaces an AITER CK RMSNorm + residual add and a pad op with an triton_add_rmsnorm_pad op from AITER. """ @@ -854,8 +844,7 @@ def uuid(self) -> str: class MLADualRMSNormPattern( VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]] ): - """ - Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into + """Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into AITER's ``fused_qk_rmsnorm`` HIP kernel. Target FX-graph pattern (unfused, ``vllm_ir`` stage):: @@ -946,8 +935,7 @@ class MLADualRMSPerTokenQuantPattern( ], ] ): - """ - Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant + """Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant plus kv-latent RMSNorm -- into AITER's ``fused_qk_rmsnorm_per_token_quant``. With a per-token FP8 ``q_b_proj`` (Quark / ModelOpt), the earlier @@ -1064,8 +1052,7 @@ def _replacement( class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass): - """ - Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in + """Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's ``fused_qk_rmsnorm`` HIP kernel. diff --git a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py index 2364281761f7..cc664941d8be 100644 --- a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py @@ -52,8 +52,7 @@ def fused_rope_and_unified_kv_cache_update_impl( is_neox: bool, layer_name: LayerNameType, ) -> torch.Tensor: - """ - This impl fetches the KV cache and slot mapping from the forward context, + """This impl fetches the KV cache and slot mapping from the forward context, then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method. It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`, that is passed to unified_attention to signal a side effect and @@ -98,8 +97,7 @@ def fused_rope_and_unified_kv_cache_update_fake( class RopeStaticQQuantKVCachePattern: - """ - Fuse rope + static Q fp8 quant while preserving explicit KV-cache update + """Fuse rope + static Q fp8 quant while preserving explicit KV-cache update dependency ordering. """ @@ -277,8 +275,7 @@ def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule: class RopeReshapeKVCachePattern: - """ - This pattern matches the following unfused inplace ops: + """This pattern matches the following unfused inplace ops: q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox) kv_cache_dummy = unified_kv_cache_update(k, v, layer_name) @@ -410,8 +407,7 @@ def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule: class RopeKVCacheFusionPass(VllmPatternMatcherPass): - """ - This pass fuses the rotary embedding and KV cache update operations + """This pass fuses the rotary embedding and KV cache update operations into a single fused kernel if available. It uses the pattern matcher and matches each layer manually, as strings diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index c4caaaedec22..b1a39b1e0146 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -57,8 +57,7 @@ def get_sequence_parallelism_threshold( tp_size: int, element_size: int, ) -> int | None: - """ - Calculate the minimum token threshold for applying sequence parallelism. + """Calculate the minimum token threshold for applying sequence parallelism. Returns None if sequence parallelism should not be applied based on model size. @@ -496,8 +495,7 @@ def replacement( class SequenceParallelismPass(VllmPatternMatcherPass): - """ - This pass enables sequence parallelism for models. + """This pass enables sequence parallelism for models. It identifies patterns where an AllReduce operation is followed by an RMSNorm (or RMSNorm and then Quantization) operation. These patterns are replaced with a ReduceScatter operation, followed by @@ -590,8 +588,7 @@ def __init__(self, config: VllmConfig) -> None: self.dump_patterns(config, self.patterns) def is_applicable_for_range(self, compile_range: Range) -> bool: - """ - Determines if sequence parallelism should be applied for the given + """Determines if sequence parallelism should be applied for the given compile range. SP is only beneficial for larger batch sizes where the communication diff --git a/vllm/compilation/passes/inductor_pass.py b/vllm/compilation/passes/inductor_pass.py index 29410f960cda..dadc56c5b2d3 100644 --- a/vllm/compilation/passes/inductor_pass.py +++ b/vllm/compilation/passes/inductor_pass.py @@ -64,14 +64,12 @@ def _hash_source_cached(*srcs: str | type | types.FunctionType) -> str: class InductorPass(CustomGraphPass): # type: ignore[misc] - """ - A custom graph pass that uses a hash of its source as the UUID. + """A custom graph pass that uses a hash of its source as the UUID. This is defined as a convenience and should work in most cases. """ def uuid(self) -> str: - """ - Provide a unique identifier for the pass, used in Inductor code cache. + """Provide a unique identifier for the pass, used in Inductor code cache. This should depend on the pass implementation, so that changes to the pass result in recompilation. By default, the object source is hashed. @@ -80,14 +78,14 @@ def uuid(self) -> str: @staticmethod def hash_source(*srcs: str | Any) -> str: - """ - Utility method to hash the sources of functions or objects. + """Utility method to hash the sources of functions or objects. Args: srcs: strings or objects to add to the hash. Objects and functions have their source inspected. Results are cached by resolved types to avoid repeated inspect.getsource() calls. + """ # Resolve instances to their class for a hashable cache key. cache_key = tuple( @@ -98,11 +96,11 @@ def hash_source(*srcs: str | Any) -> str: @staticmethod def hash_dict(dict_: dict[Any, Any]) -> str: - """ - Utility method to hash a dictionary, can alternatively be used for uuid. + """Utility method to hash a dictionary, can alternatively be used for uuid. Returns: A sha256 hash of the json rep of the dictionary. + """ encoded = json.dumps(dict_, sort_keys=True).encode("utf-8") return hashlib.sha256(encoded).hexdigest() @@ -112,8 +110,7 @@ def is_applicable_for_range(self, compile_range: Range) -> bool: class CallableInductorPass(InductorPass): - """ - This class is a wrapper for a callable that automatically provides an + """This class is a wrapper for a callable that automatically provides an implementation of the UUID. """ @@ -131,8 +128,7 @@ def uuid(self) -> Any: def enable_fake_mode(fn: Callable[P, R]) -> Callable[P, R]: - """ - Applies a FakeTensorMode context. This is useful when you don't want to + """Applies a FakeTensorMode context. This is useful when you don't want to create or run things with real tensors. """ diff --git a/vllm/compilation/passes/ir/clone_elimination.py b/vllm/compilation/passes/ir/clone_elimination.py index 377cd669a4eb..d35ebec97f46 100644 --- a/vllm/compilation/passes/ir/clone_elimination.py +++ b/vllm/compilation/passes/ir/clone_elimination.py @@ -70,8 +70,7 @@ def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool: class UnsafeCloneEliminationPass(VllmInductorPass): - """ - This pass removes clone nodes that are no longer needed after vLLM IR lowering. + """This pass removes clone nodes that are no longer needed after vLLM IR lowering. It uses donated_input_ids to eliminate clones of donated graph inputs, preserving contents of non-donated graph inputs. diff --git a/vllm/compilation/passes/ir/inplace_functionalization.py b/vllm/compilation/passes/ir/inplace_functionalization.py index e69351075bca..0c6e66e6d6b6 100644 --- a/vllm/compilation/passes/ir/inplace_functionalization.py +++ b/vllm/compilation/passes/ir/inplace_functionalization.py @@ -19,8 +19,7 @@ class VllmIRInplaceFunctionalizationPass(VllmInductorPass): - """ - This pass functionalizes maybe_inplace vLLM IR ops to the default overload. + """This pass functionalizes maybe_inplace vLLM IR ops to the default overload. The maybe_inplace overloads have the same signature as the default overload so the pass simply replaces the called overload. That makes the graph properly functional. diff --git a/vllm/compilation/passes/ir/lowering_pass.py b/vllm/compilation/passes/ir/lowering_pass.py index f34f1c64b76e..f40d23ac473c 100644 --- a/vllm/compilation/passes/ir/lowering_pass.py +++ b/vllm/compilation/passes/ir/lowering_pass.py @@ -23,9 +23,7 @@ class VllmIRLoweringPass(VllmInductorPass): - """ - This pass lowers vLLM IR ops to their implementations the priority list. - """ + """This pass lowers vLLM IR ops to their implementations the priority list.""" def __init__(self, vllm_config: VllmConfig) -> None: super().__init__(vllm_config) @@ -113,8 +111,7 @@ def print_count(counts: dict[str, int]) -> str: logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes)) def uuid(self) -> str: - """ - IR op priority & impl sources affect lowering pass output, + """IR op priority & impl sources affect lowering pass output, so we include them in the cache key. """ priorities = {name: op.get_priority() for name, op in IrOp.registry.items()} diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 4a65376ca09b..8ee8f0354e2f 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -71,8 +71,7 @@ def with_pattern_match_debug(fn: Callable[P, R]) -> Callable[P, R]: - """ - Function decorator that turns on inductor pattern match debug + """Function decorator that turns on inductor pattern match debug for the duration of the call. Used to avoid logging builtin Inductor pattern matching. """ @@ -89,8 +88,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: class PostGradPassManager(CustomGraphPass): # type: ignore[misc] - """ - The pass manager for post-grad passes. + """The pass manager for post-grad passes. It handles configuration, adding custom passes, and running passes. It supports uuid for the Inductor code cache. That includes torch<2.6 support using pickling (in .inductor_pass.CustomGraphPass). @@ -236,8 +234,7 @@ def add(self, pass_: InductorPass) -> None: self.passes.append(pass_) def uuid(self) -> str: - """ - The PostGradPassManager is set as a custom pass in the Inductor and + """The PostGradPassManager is set as a custom pass in the Inductor and affects compilation caching. Its uuid depends on the UUIDs of all dependent passes and the pass config. See InductorPass for more info. """ diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index c04f9abbb6a9..b7baac10e0b6 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/vllm/compilation/passes/utility/fix_functionalization.py @@ -17,8 +17,7 @@ class FixFunctionalizationPass(VllmInductorPass): - """ - This pass defunctionalizes certain nodes to avoid redundant tensor copies. + """This pass defunctionalizes certain nodes to avoid redundant tensor copies. After this pass, DCE (dead-code elimination) should never be run, as de-functionalized nodes may appear as dead code. @@ -260,9 +259,7 @@ def __call__(self, graph: torch.fx.Graph) -> None: self.nodes_to_remove.clear() def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]) -> None: - """ - Stage a node (or nodes) for removal at the end of the pass. - """ + """Stage a node (or nodes) for removal at the end of the pass.""" if isinstance(node_or_nodes, torch.fx.Node): self.nodes_to_remove.append(node_or_nodes) else: @@ -275,8 +272,7 @@ def defunctionalize( mutated_args: dict[int, torch.fx.Node | str], args: tuple[torch.fx.Node | str, ...] | None = None, ) -> None: - """ - De-functionalize a node by replacing it with a call to the original. + """De-functionalize a node by replacing it with a call to the original. It also replaces the getitem users with the mutated arguments. See replace_users_with_mutated_args and insert_defunctionalized. """ @@ -287,14 +283,14 @@ def defunctionalize( def replace_users_with_mutated_args( self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str] ) -> None: - """ - Replace mutated getitem users of the auto-functionalized node with the + """Replace mutated getitem users of the auto-functionalized node with the mutated arguments. Args: node: The auto-functionalized node mutated_args: The mutated arguments, indexed by getitem index. If the value of an arg is a string, `node.kwargs[arg]` is used. + """ for idx, user in self.getitem_users(node).items(): # Some functionalized nodes may return both a result at getitem[0] @@ -310,8 +306,7 @@ def replace_users_with_mutated_args( self._remove(user) def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: - """ - Returns the operator.getitem users of the auto-functionalized node, + """Returns the operator.getitem users of the auto-functionalized node, indexed by the index they are getting. """ users = {} @@ -327,8 +322,7 @@ def insert_defunctionalized( node: torch.fx.Node, args: tuple[torch.fx.Node | str, ...] | None = None, ) -> None: - """ - Insert a new defunctionalized node into the graph before node. + """Insert a new defunctionalized node into the graph before node. If one of the kwargs is 'out', provide args directly, as node.kwargs cannot be used. See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351 @@ -338,6 +332,7 @@ def insert_defunctionalized( node: The auto-functionalized node to defunctionalize args: If we cannot use kwargs, specify args directly. If an arg is a string, `node.kwargs[arg]` is used. + """ # noqa: E501 assert is_func(node, auto_functionalized), ( f"node must be auto-functionalized, is {node} instead" diff --git a/vllm/compilation/passes/utility/noop_elimination.py b/vllm/compilation/passes/utility/noop_elimination.py index 80bf8ecc603f..ec9df636f7b2 100644 --- a/vllm/compilation/passes/utility/noop_elimination.py +++ b/vllm/compilation/passes/utility/noop_elimination.py @@ -16,8 +16,7 @@ class NoOpEliminationPass(VllmInductorPass): - """ - This is an inductor pass that removes redundant reshape/slice operations. + """This is an inductor pass that removes redundant reshape/slice operations. It is required for RMSNorm-quant fusion to work properly. That's because apply_fp8_linear adds a reshape, which is redundant in the 2D-case. Additionally, torch internal no-op elimination pass does @@ -106,8 +105,7 @@ def __call__(self, graph: torch.fx.Graph) -> None: # ---------------------- Shape comparison helpers ---------------------- def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool: - """ - This function checks if two dimensions are equivalent. + """This function checks if two dimensions are equivalent. Args: dim: The dimension arg to reshape/slice @@ -119,6 +117,7 @@ def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool: There are two cases in which the dimensions are equivalent: 1. The dimensions are equal (both integers) 2. The dimensions both correspond to the same SymInt + """ # Case 1 return statically_known_true(dim == i_dim) # type: ignore[no-any-return] diff --git a/vllm/compilation/passes/utility/post_cleanup.py b/vllm/compilation/passes/utility/post_cleanup.py index d4ecd4d65563..5b0695c61f16 100644 --- a/vllm/compilation/passes/utility/post_cleanup.py +++ b/vllm/compilation/passes/utility/post_cleanup.py @@ -6,8 +6,7 @@ class PostCleanupPass(VllmInductorPass): - """ - This pass performs cleanup after custom passes. + """This pass performs cleanup after custom passes. It topologically sorts the graph and removes unused nodes. This is needed because the pattern matcher does not guarantee producing a topologically sorted graph, and there may be unused nodes left around. diff --git a/vllm/compilation/passes/utility/scatter_split_replace.py b/vllm/compilation/passes/utility/scatter_split_replace.py index a17a7b336d2d..de7d87d87394 100644 --- a/vllm/compilation/passes/utility/scatter_split_replace.py +++ b/vllm/compilation/passes/utility/scatter_split_replace.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single +"""Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single assignment if there are no users for the inplace tensor written to by the slice_scatter call. diff --git a/vllm/compilation/passes/utility/split_coalescing.py b/vllm/compilation/passes/utility/split_coalescing.py index 6373eacfa847..9b60d10fba37 100644 --- a/vllm/compilation/passes/utility/split_coalescing.py +++ b/vllm/compilation/passes/utility/split_coalescing.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Coalesce duplicate ``split_with_sizes`` nodes that operate on the same +"""Coalesce duplicate ``split_with_sizes`` nodes that operate on the same input tensor with the same split sizes. On certain hardware/dtype combinations (e.g. B200 + FP8) the Inductor @@ -10,9 +9,10 @@ so that downstream pattern-matching passes (e.g. QK-Norm+RoPE fusion) see a single split node with all users attached. -See also: +See Also: - vLLM #33295 (original issue) - PyTorch #174472 (upstream CSE gap) + """ import operator @@ -30,7 +30,8 @@ class SplitCoalescingPass(VllmInductorPass): """Replace duplicate ``split_with_sizes`` nodes with a single canonical - node when they share the same input tensor and split sizes.""" + node when they share the same input tensor and split sizes. + """ @VllmInductorPass.time_and_log def __call__(self, graph: fx.Graph) -> None: diff --git a/vllm/compilation/passes/vllm_inductor_pass.py b/vllm/compilation/passes/vllm_inductor_pass.py index 351d83e91855..f05c732af381 100644 --- a/vllm/compilation/passes/vllm_inductor_pass.py +++ b/vllm/compilation/passes/vllm_inductor_pass.py @@ -35,8 +35,7 @@ class InductorCompilationConfig: class VllmInductorPass(InductorPass): - """ - An inductor pass with access to vLLM PassConfig. + """An inductor pass with access to vLLM PassConfig. It provides timing, logging, and dumping utilities. """ @@ -93,8 +92,7 @@ def get_match_table() -> dict[str, int]: class VllmPatternMatcherPass(VllmInductorPass): - """ - A VllmInductorPass that uses the Inductor pattern matcher. + """A VllmInductorPass that uses the Inductor pattern matcher. Provides pattern registration with match counting, debug dumping, and logging. """ @@ -109,7 +107,7 @@ class VllmPatternMatcherPass(VllmInductorPass): ) def _replace_op_overloads(self, string: str) -> str: - """Replace with nicer formulations""" + """Replace with nicer formulations.""" return str( self._OP_OVERLOAD_PATTERN.sub( lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}", @@ -123,8 +121,7 @@ def log_match_summary(cls) -> None: logger.debug("fusion pass matches: %s", dict(cls.match_table)) def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None: - """ - If debug dumping is enabled, dump the Inductor pattern-matcher patterns + """If debug dumping is enabled, dump the Inductor pattern-matcher patterns into the debug_dump_path folder next to the dumped fx graphs. This method does its best to print something that looks like Python code @@ -195,8 +192,7 @@ def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None class VllmPatternReplacement(ABC, Generic[P, R]): - """ - A pattern/replacement pair for FX graph fusion. + """A pattern/replacement pair for FX graph fusion. Implement the three abstract members below, then pass instances to VllmFusionPatternMatcherPass.register(). The pass will @@ -216,8 +212,7 @@ def pattern(self) -> Callable[P, R]: @property @abstractmethod def replacement(self) -> Callable[P, R]: - """ - Returns a closure defining the FX subgraph to + """Returns a closure defining the FX subgraph to substitute in place of each match. """ ... @@ -294,8 +289,7 @@ def _remove_noop_permutes(gm: fx.GraphModule) -> None: class VllmFusionPatternMatcherPass(VllmPatternMatcherPass): - """ - A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects. + """A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects. Subclasses register patterns via self.register() in their own __init__. """ diff --git a/vllm/compilation/piecewise_backend.py b/vllm/compilation/piecewise_backend.py index b647e0d8581a..576cd941f6a8 100644 --- a/vllm/compilation/piecewise_backend.py +++ b/vllm/compilation/piecewise_backend.py @@ -96,8 +96,7 @@ def __init__( compiled_runnables: dict[str, Callable[..., Any]] | None = None, submod_name: str = "", ): - """ - The backend for piecewise compilation. + """The backend for piecewise compilation. It mainly handles the compilation of static shapes and dispatching based on runtime shape. diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py index 4821daf9dded..1df30a70a83f 100644 --- a/vllm/compilation/wrapper.py +++ b/vllm/compilation/wrapper.py @@ -45,8 +45,7 @@ def _compilation_context() -> Generator[None, None, None]: class TorchCompileWithNoGuardsWrapper: - """ - A wrapper class for torch.compile, it ensures that all guards are dropped + """A wrapper class for torch.compile, it ensures that all guards are dropped when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE. When guards are dropped, the first time __call__ is invoked, a single compilation is triggered. Dynamo should never be traced again after that @@ -272,8 +271,7 @@ def cleanup(self) -> None: @contextmanager def _dispatch_to_compiled_code(self) -> Generator[None, None, None]: # noqa: E501 - """ - Context manager to dispatch to internally compiled code for torch<2.8. + """Context manager to dispatch to internally compiled code for torch<2.8. Why does this work? Because Dynamo guarantees that the compiled bytecode has exactly the same arguments, cell variables, and free variables as the original code. Therefore we can directly switch @@ -291,9 +289,7 @@ def _dispatch_to_compiled_code(self) -> Generator[None, None, None]: def reset_compile_wrapper(model: torch.nn.Module) -> None: - """ - Clean up compiled model and captured CUDA graphs for elastic EP. - """ + """Clean up compiled model and captured CUDA graphs for elastic EP.""" if not isinstance(model, TorchCompileWithNoGuardsWrapper) and hasattr( model, "model" ): diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 3646677728b5..a6602ed5491b 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -150,8 +150,7 @@ def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType: return self.indexer_kv_dtype def compute_hash(self) -> str: - """ - Provide a hash that uniquely identifies all the configs + """Provide a hash that uniquely identifies all the configs that affect the structure of the computation graph from input ids/embeddings to the final hidden states, excluding anything before input ids/embeddings and after diff --git a/vllm/config/cache.py b/vllm/config/cache.py index ce3199c52a25..37ad5b2b34fe 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -248,8 +248,7 @@ class CacheConfig: KV offloading is only activated when kv_offloading_size is set.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2cd5fae77554..f1b5e90f2ac7 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -36,7 +36,8 @@ class CompilationMode(enum.IntEnum): """The compilation approach used for torch.compile-based compilation of the - model.""" + model. + """ NONE = 0 """No torch.compile compilation is applied, model runs in fully eager pytorch mode. @@ -192,12 +193,10 @@ class PassConfig: # TODO(luka) better pass enabling system. def flashinfer_max_size(self, world_size: int) -> int | None: - """ - Returns the max communication size in bytes for flashinfer + """Returns the max communication size in bytes for flashinfer allreduce fusion for the given world size. Returns None if world size is not supported by configs as it's not supported by flashinfer. """ - MiB = 1024 * 1024 FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16] if world_size not in FI_SUPPORTED_WORLD_SIZES: @@ -223,12 +222,10 @@ def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]: return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(capability.to_int(), {}) def compute_hash(self) -> str: - """ - Produces a hash unique to the pass configuration. + """Produces a hash unique to the pass configuration. Any new fields that affect compilation should be added to the hash. Any future fields that don't affect compilation should be excluded. """ - return hash_factors(get_hash_factors(self, set())) @field_validator( @@ -317,8 +314,7 @@ def __post_init__(self) -> None: self.fuse_rope_kvcache_cat_mla = False def log_enabled_passes(self) -> None: - """ - Log the enabled custom fusion passes. + """Log the enabled custom fusion passes. This is called at the end of VLLMConfig post_init, after all defaults are finalized. TODO also log the compile ranges for which this is enabled. @@ -341,7 +337,8 @@ def log_enabled_passes(self) -> None: class DynamicShapesType(str, enum.Enum): """Types of dynamic shapes handling in torch.compile(). see Dynamic shapes and vllm guard dropping in torch_compile.md - for more details.""" + for more details. + """ BACKED = "backed" """Use backed dynamic shapes. torch.compile() guards on backed dynamic @@ -391,10 +388,7 @@ class DynamicShapesConfig: """ def compute_hash(self) -> str: - """ - Provide a hash for DynamicShapesConfig - """ - + """Provide a hash for DynamicShapesConfig.""" from vllm.config.utils import get_hash_factors, hash_factors factors = get_hash_factors(self, set()) @@ -790,8 +784,7 @@ class CompilationConfig: ] def compute_hash(self) -> str: - """ - Provide a hash that uniquely identifies all the configs + """Provide a hash that uniquely identifies all the configs that affect the structure of the computation graph from input ids/embeddings to the final hidden states, excluding anything before input ids/embeddings and after @@ -856,8 +849,7 @@ def __repr__(self) -> str: @field_validator("mode", mode="before") @classmethod def validate_mode_before(cls, value: Any) -> Any: - """ - Enable parsing the `mode` field from string mode names. + """Enable parsing the `mode` field from string mode names. Accepts both integers (0-3) and string names, like NONE, STOCK_TORCH_COMPILE, DYNAMO_TRACE_ONCE, VLLM_COMPILE. """ @@ -1079,15 +1071,17 @@ def init_backend( prefix: str = "", is_encoder: bool = False, ) -> str | Callable: - """ - Initialize the backend for the compilation config from a vllm config. + """Initialize the backend for the compilation config from a vllm config. + Arguments: vllm_config: The vllm config to initialize the backend from. prefix: Cache directory prefix for this compiled module. is_encoder: Whether this module is used in an encoder (as opposed to a text backbone). + Returns: The backend for the compilation config. + """ if self.mode is None: raise ValueError( @@ -1122,7 +1116,6 @@ def post_init_cudagraph_sizes(self) -> None: configs are set. This includes: - initialize compile_sizes """ - computed_compile_sizes: list[int] = [] if self.compile_sizes is not None: # de-duplicate the sizes provided by the config @@ -1321,13 +1314,11 @@ def is_attention_compiled_piecewise(self) -> bool: return self.backend == "inductor" and self.mode != CompilationMode.NONE def custom_op_log_check(self): - """ - This method logs the enabled/disabled custom ops and checks that the + """This method logs the enabled/disabled custom ops and checks that the passed custom_ops field only contains relevant ops. It is called at the end of set_current_vllm_config, after the custom ops have been instantiated. """ - if len(self.enabled_custom_ops) + len(self.disabled_custom_ops) == 0: logger.debug("No custom ops found in model.") return diff --git a/vllm/config/device.py b/vllm/config/device.py index 2d46dcdb117e..c0bdf01ee47d 100644 --- a/vllm/config/device.py +++ b/vllm/config/device.py @@ -28,8 +28,7 @@ class DeviceConfig: `__post_init__`.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/ec_manager_config.py b/vllm/config/ec_manager_config.py index a59d2cdc62b8..db284e9cfc4f 100644 --- a/vllm/config/ec_manager_config.py +++ b/vllm/config/ec_manager_config.py @@ -31,8 +31,7 @@ def get_encoder_cache_manager_obj(self): class EncoderCacheManagerMetadata(ABC): # noqa: B024 - """ - Abstract Metadata used to communicate between the + """Abstract Metadata used to communicate between the Scheduler EncoderCacheManager and Worker EncoderCacheManager. """ diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 3948f3205934..d6129807ab01 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -63,8 +63,7 @@ class ECTransferConfig: Only supported in V1.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 0ac5bf43913d..ea0b7d60d1fc 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -18,8 +18,7 @@ @config class IrOpPriorityConfig: - """ - Configuration for vLLM IR op priority for dispatching/lowering during the + """Configuration for vLLM IR op priority for dispatching/lowering during the forward pass. Each member is a list of strings, which will be installed in worker init via vllm.ir.ops..set_default(). A single comma-separated string is accepted as well, @@ -38,8 +37,7 @@ class IrOpPriorityConfig: """Priority list for vllm.ir.ops.gelu_and_mul_sparse""" def compute_hash(self) -> str: - """ - Produces a hash unique to the pass configuration. + """Produces a hash unique to the pass configuration. Any new fields that affect compilation should be added to the hash. Any future fields that don't affect compilation should be excluded. @@ -71,8 +69,7 @@ def _to_list_str(cls, value: str | list[str]): return value def _iter_op_priorities(self): - """ - Yield (IrOp, priority_list) for each field, after importing platform + """Yield (IrOp, priority_list) for each field, after importing platform kernels and validating each entry. """ from vllm.ir.op import IrOp @@ -89,16 +86,13 @@ def _iter_op_priorities(self): yield IrOp.registry[field.name], op_priority def set_default(self) -> None: - """ - Permanently set the IR op priority for all op members. - """ + """Permanently set the IR op priority for all op members.""" for ir_op, op_priority in self._iter_op_priorities(): ir_op.set_default(op_priority) @contextlib.contextmanager def set_priority(self): - """ - Context manager to set the IR op priority for all op members. + """Context manager to set the IR op priority for all op members. It also imports IR kernel implementations for the current platform to ensure all implementations are made available. """ @@ -111,8 +105,7 @@ def set_priority(self): def with_default( cls, default: list[str], /, **kwargs: list[str] ) -> "IrOpPriorityConfig": - """ - A helper to create an IrOpPriorityConfig where fields not specified in kwargs + """A helper to create an IrOpPriorityConfig where fields not specified in kwargs use the given default list. """ for field in fields(cls): # type: ignore[arg-type] @@ -355,8 +348,7 @@ def _normalize_sparse_indexer_topk_backend(cls, value: Any) -> Any: return value def compute_hash(self) -> str: - """ - Produces a hash unique to the pass configuration. + """Produces a hash unique to the pass configuration. Any new fields that affect compilation should be added to the hash. Any future fields that don't affect compilation should be excluded. """ diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index b299aae6f0a7..dfc9e02f3d0a 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -106,8 +106,7 @@ class KVTransferConfig: 'fail': immediately fail the request with an error finish reason (default)""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/load.py b/vllm/config/load.py index 5e24a617ba85..ee13e2fccd7c 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -118,8 +118,7 @@ class LoadConfig: """ def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 4323da96d806..24e42a42fc90 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -86,8 +86,7 @@ class LoRAConfig: MoE models whose adapters use this layout; ignored otherwise.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/model.py b/vllm/config/model.py index 89978c35d051..a07fbbc7a257 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -418,8 +418,7 @@ class ModelConfig: mm_processor_device: InitVar[MMProcessorDevice | None] = None def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. @@ -938,7 +937,7 @@ def validate_quantization_before(cls, value: Any) -> Any: @model_validator(mode="after") def validate_model_config_after(self: "ModelConfig") -> "ModelConfig": - """Called after __post_init__""" + """Called after __post_init__.""" if not isinstance(self.tokenizer, str): raise ValueError( f"tokenizer must be a string, got " @@ -1006,7 +1005,8 @@ def _resolve_mm_device_do_normalize( def _get_transformers_backend_cls(self) -> str: """Determine which Transformers modeling backend class will be used if - `model_impl` is set to `transformers` or `auto`.""" + `model_impl` is set to `transformers` or `auto`. + """ cls = "Transformers" # If 'hf_config is not hf_text_config' it's a nested config, i.e. multimodal cls += "MultiModal" if self.hf_config is not self.hf_text_config else "" @@ -1083,8 +1083,8 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No Args: model: Model name or path tokenizer: Tokenizer name or path - """ + """ # Skip if model_weights is already set (model already pulled) if self.model_weights: return @@ -1662,9 +1662,7 @@ def get_num_layers_by_block_type( raise AssertionError(f"Unsupported block type: {block_type}") def get_mamba_chunk_size(self) -> int: - """ - Returns the mamba chunk size if it exists - """ + """Returns the mamba chunk size if it exists.""" # used by e.g. Bamba, FalconH1, Granite chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None) if chunk_size is None: @@ -1679,11 +1677,11 @@ def get_mamba_chunk_size(self) -> int: return chunk_size def get_multimodal_config(self) -> MultiModalConfig: - """ - Get the multimodal configuration of the model. + """Get the multimodal configuration of the model. Raises: ValueError: If the model is not multimodal. + """ if self.multimodal_config is None: raise ValueError("The model is not multimodal.") @@ -1691,8 +1689,7 @@ def get_multimodal_config(self) -> MultiModalConfig: return self.multimodal_config def try_get_generation_config(self) -> dict[str, Any]: - """ - This method attempts to retrieve the non-default values of the + """This method attempts to retrieve the non-default values of the generation config for this model. The generation config can contain information about special tokens, as @@ -1701,6 +1698,7 @@ def try_get_generation_config(self) -> dict[str, Any]: Returns: A dictionary containing the non-default generation config. + """ if self.generation_config in {"auto", "vllm"}: config = try_get_generation_config( @@ -1726,8 +1724,7 @@ def try_get_generation_config(self) -> dict[str, Any]: return config.to_diff_dict() def get_diff_sampling_param(self) -> dict[str, Any]: - """ - This method returns a dictionary containing the non-default sampling + """This method returns a dictionary containing the non-default sampling parameters with `override_generation_config` applied. The default sampling parameters are: @@ -1739,6 +1736,7 @@ def get_diff_sampling_param(self) -> dict[str, Any]: Returns: A dictionary containing the non-default sampling parameters. + """ src = self.generation_config @@ -1870,8 +1868,7 @@ def requires_raw_input_tokens(self) -> bool: @property def score_type(self) -> ScoreType: - """ - Scoring API handles score/rerank for: + """Scoring API handles score/rerank for: - "classify" task (score_type: cross-encoder models) - "embed" task (score_type: bi-encoder models) @@ -1964,8 +1961,7 @@ def use_sep_token(self) -> bool: @property def head_dtype(self) -> torch.dtype: - """ - "head" refers to the last Linear layer(s) of an LLM, + """The "head" refers to the last Linear layer(s) of an LLM, such as the lm_head in a generation model, or the score or classifier in a classification model. @@ -1976,7 +1972,6 @@ def head_dtype(self) -> torch.dtype: fp32, which is required for RL training-inference consistency (the trainer computes logits in fp32). """ - head_dtype = _get_head_dtype( config=self.hf_config, dtype=self.dtype, runner_type=self.runner_type ) @@ -2174,8 +2169,7 @@ def is_nvfp4_quantized(self) -> bool: def get_served_model_name(model: str, served_model_name: str | list[str] | None): - """ - If the input is a non-empty list, the first model_name in + """If the input is a non-empty list, the first model_name in `served_model_name` is taken. If the input is a non-empty string, it is used directly. For cases where the input is either an empty string or an diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py index 9bd449b7b1de..9aed35384a8c 100644 --- a/vllm/config/model_arch.py +++ b/vllm/config/model_arch.py @@ -14,9 +14,7 @@ @dataclass(config=ConfigDict(arbitrary_types_allowed=True)) class ModelArchitectureConfig: - """ - Configuration for model architecture that required by vLLM runtime - """ + """Configuration for model architecture that required by vLLM runtime.""" architectures: list[str] """List of model architecture class names (e.g., ['LlamaForCausalLM']). diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 0d0dc22a273b..5f93fce0aa58 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -357,6 +357,7 @@ def fold_mm_processor_device( Returns: The kwargs to build the config with, unchanged unless the flag adds a `device`. + """ if mm_processor_device in (None, "auto"): return mm_processor_kwargs @@ -390,6 +391,7 @@ def get_mm_processor_device_type(self) -> str | None: ValueError: If `device` is not something `torch.device` accepts. `validate_mm_processor_device` is what surfaces this during startup, so the value is only parsed once. + """ device = (self.mm_processor_kwargs or {}).get("device") if device is None: @@ -418,6 +420,7 @@ def validate_mm_processor_device(self, ec_config: ECTransferConfig | None) -> No ValueError: If the requested device is not a torch device, or if it is the accelerator on an instance that also runs the language model. + """ from vllm.platforms import current_platform @@ -452,8 +455,7 @@ def validate_mm_processor_device(self, ec_config: ECTransferConfig | None) -> No ) def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. @@ -476,8 +478,7 @@ def compute_hash(self) -> str: return hash_str def get_limit_per_prompt(self, modality: str) -> int: - """ - Get the maximum number of input items allowed per prompt + """Get the maximum number of input items allowed per prompt for the given modality (backward compatible). """ if self.language_model_only: @@ -495,8 +496,7 @@ def merge_mm_processor_kwargs( self, inference_kwargs: Mapping[str, object], ) -> dict[str, object]: - """ - Get the keyword arguments to pass to the multi-modal processor + """Get the keyword arguments to pass to the multi-modal processor according to the extra arguments passed during inference. """ kwargs = self.mm_processor_kwargs or {} @@ -524,7 +524,8 @@ def is_multimodal_pruning_enabled(self): def get_video_pruning_spec(self) -> tuple[VideoPruningMethod, float] | None: """Return `(method, rate)` when video pruning is enabled, else None. - `rate` is the fraction of video tokens to prune.""" + `rate` is the fraction of video tokens to prune. + """ if self.video_pruning_rate is not None and self.video_pruning_rate > 0: return (self.video_pruning_method, float(self.video_pruning_rate)) return None diff --git a/vllm/config/observability.py b/vllm/config/observability.py index 0394245f3cf2..c964fe904db5 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -112,8 +112,7 @@ def collect_model_execute_time(self) -> bool: ) def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. @@ -157,7 +156,8 @@ def _validate_collect_detailed_traces( cls, value: list[DetailedTraceModules] | None ) -> list[DetailedTraceModules] | None: """Handle the legacy case where users might provide a comma-separated - string instead of a list of strings.""" + string instead of a list of strings. + """ if value is not None and len(value) == 1 and "," in value[0]: value = cast(list[DetailedTraceModules], value[0].split(",")) return value diff --git a/vllm/config/offload.py b/vllm/config/offload.py index ad65e8acf35a..98fab0409f05 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -137,8 +137,7 @@ def validate_offload_config(self) -> "OffloadConfig": return self def compute_hash(self) -> str: - """ - Provide a hash that uniquely identifies all the offload configs. + """Provide a hash that uniquely identifies all the offload configs. All fields are included because PrefetchOffloader patches module forwards and inserts custom ops (wait_prefetch, start_prefetch) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 0213e06283df..f66af2f30b8d 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -602,15 +602,13 @@ def num_ubatches(self) -> int: @property def local_engines_only(self) -> bool: - """ - Client manages local+remote EngineCores in pure internal LB case. + """Client manages local+remote EngineCores in pure internal LB case. Client manages local EngineCores in hybrid and external LB case. """ return self.data_parallel_external_lb or self.data_parallel_hybrid_lb def get_next_dp_init_port(self) -> int: - """ - We might need to initialize process groups in multiple + """We might need to initialize process groups in multiple processes that is related to data parallelism, e.g. both in the worker and in the engine, which can live in different processes. To avoid port conflicts, we @@ -794,6 +792,7 @@ def sync_dp_state( Returns: (has_unfinished_global, pause_consensus) + """ tensor = torch.tensor( [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu" @@ -815,8 +814,7 @@ def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> i return tensor.item() def compute_hash(self): - """ - Provide a hash that uniquely identifies all the configs + """Provide a hash that uniquely identifies all the configs that affect the structure of the computation graph from input ids/embeddings to the final hidden states, excluding anything before input ids/embeddings and after diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py index 9ca0ba457ebf..a5fbcbb622ed 100644 --- a/vllm/config/pooler.py +++ b/vllm/config/pooler.py @@ -176,8 +176,7 @@ def get_tok_pooling_type(self) -> TokenPoolingType: return self.tok_pooling_type def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 535960a8d7cd..f3a69853243a 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -151,8 +151,7 @@ class ProfilerConfig: """ def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/reasoning.py b/vllm/config/reasoning.py index 0ac93be0823a..e00d9adc27ab 100644 --- a/vllm/config/reasoning.py +++ b/vllm/config/reasoning.py @@ -47,13 +47,15 @@ class ReasoningConfig: @property def enabled(self) -> bool: """Returns True if reasoning is enabled (i.e. if token IDs have been - initialized), False otherwise.""" + initialized), False otherwise. + """ return self._enabled @property def reasoning_start_token_ids(self) -> list[int] | None: """Token IDs derived from `reasoning_start_str`. Set automatically by - `initialize_token_ids`. Not intended to be configured directly.""" + `initialize_token_ids`. Not intended to be configured directly. + """ return self._reasoning_start_token_ids @property diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 8aa0ffa5dcff..d3264b2d8071 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -200,9 +200,7 @@ class SchedulerConfig: @staticmethod def default_factory(**kwargs): - """ - Factory method to create `SchedulerConfig` with default values for `InitVar`s. - """ + """Create a `SchedulerConfig` with default values for its `InitVar`s.""" if "max_model_len" not in kwargs: kwargs["max_model_len"] = 8192 if "is_encoder_decoder" not in kwargs: @@ -233,8 +231,7 @@ def get_scheduler_cls(self) -> type["SchedulerInterface"]: return resolve_obj_by_qualname(self.scheduler_cls) def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 2396b26cb01a..96340df59117 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -543,7 +543,8 @@ class SpeculativeConfig: @staticmethod def _acceptance_length_to_rates(length: float, n: int) -> list[float]: """Mean acceptance length to unconditional per-position rates, using - the minimum-variance schedule.""" + the minimum-variance schedule. + """ num_drafts = length - 1 # expected number of accepted draft tokens num_full = int(num_drafts) return ( @@ -557,7 +558,8 @@ def _resolve_synthetic_acceptance_rates( length: float | None, ) -> list[float]: """Return per-position unconditional acceptance rates from exactly one - of `rates` or `length` (validates range, length, and monotonicity).""" + of `rates` or `length` (validates range, length, and monotonicity). + """ if (rates is None) == (length is None): raise ValueError( "rejection_sample_method='synthetic' requires exactly one of " @@ -598,8 +600,7 @@ def _resolve_synthetic_acceptance_rates( top-k base-logit candidates. Requires draft tensor parallel size 1.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. @@ -1627,7 +1628,6 @@ def _maybe_override_draft_max_model_len( speculative_max_model_len is mainly used for testing that sequences can skip speculation. """ - if speculative_max_model_len is not None: if speculative_max_model_len > draft_max_model_len: raise ValueError( @@ -1671,6 +1671,7 @@ def _maybe_override_draft_max_position_embeddings( Args: draft_hf_config: The draft model's HF config, mutated in place. target_max_model_len: The target model's max_model_len. + """ draft_max_position_embeddings = getattr( draft_hf_config, "max_position_embeddings", None @@ -1695,8 +1696,7 @@ def _verify_and_get_draft_tp( speculative_draft_tensor_parallel_size: int | None, draft_hf_config: PretrainedConfig, ) -> int: - """ - Verifies and adjusts the tensor parallel size for a draft model + """Verifies and adjusts the tensor parallel size for a draft model specified using speculative_draft_tensor_parallel_size. """ # If speculative_draft_tensor_parallel_size is unset then set it @@ -1725,8 +1725,7 @@ def _verify_and_get_draft_tp( return speculative_draft_tensor_parallel_size def update_arch_(self): - """ - EagleConfig and ExtractHiddenStatesConfig update architectures, so update all + """EagleConfig and ExtractHiddenStatesConfig update architectures, so update all architectures-related fields in self.draft_model_config """ self.draft_model_config.hf_text_config = get_hf_text_config( diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py index e7afbb65bc7f..059581e37e65 100644 --- a/vllm/config/structured_outputs.py +++ b/vllm/config/structured_outputs.py @@ -42,8 +42,7 @@ class StructuredOutputsConfig: """Whether to use structured input for reasoning.""" def compute_hash(self) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 6ab346d1a030..192b8bc8707c 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -64,7 +64,9 @@ def config( cls: The class to decorate config: The pydantic ConfigDict to use. If provided, it will be merged with the default config. - **kwargs: Additional arguments to pass to pydantic.dataclass.""" + **kwargs: Additional arguments to pass to pydantic.dataclass. + + """ # Extra fields are forbidden by default merged_config = ConfigDict(extra="forbid") if config is not None: @@ -82,7 +84,8 @@ def decorator(cls: type[ConfigT]) -> type[ConfigT]: def get_field(cls: ConfigType, name: str) -> Any: """Get the default factory field of a dataclass by name. Used for getting - default factory fields in `EngineArgs`.""" + default factory fields in `EngineArgs`. + """ if not is_dataclass(cls): raise TypeError("The given class is not a dataclass.") try: @@ -119,7 +122,8 @@ def is_init_field(cls: ConfigType, name: str) -> bool: def replace(dataclass_instance: ConfigT, /, **kwargs) -> ConfigT: """Like [`dataclasses.replace`](https://docs.python.org/3/library/dataclasses.html#dataclasses.replace), but compatible with Pydantic dataclasses which use `pydantic.fields.Field` instead - of `dataclasses.field`""" + of `dataclasses.field` + """ cls = type(dataclass_instance) dataclass_dict = dataclass_instance.__dict__ dataclass_dict = {k: v for k, v in dataclass_dict.items() if is_init_field(cls, k)} @@ -134,8 +138,7 @@ def getattr_iter( default_factory: Callable[[], Any] | None = None, warn: bool = False, ) -> Any: - """ - A helper function that retrieves an attribute from an object which may + """A helper function that retrieves an attribute from an object which may have multiple possible names. This is useful when fetching attributes from arbitrary `transformers.PretrainedConfig` instances. @@ -158,12 +161,10 @@ def getattr_iter( def get_attr_docs(cls: type[Any]) -> dict[str, str]: - """ - Get any docstrings placed after attribute assignments in a class body. + """Get any docstrings placed after attribute assignments in a class body. https://davidism.com/mit-license/ """ - cls_node = ast.parse(textwrap.dedent(inspect.getsource(cls))).body[0] if not isinstance(cls_node, ast.ClassDef): @@ -385,8 +386,7 @@ def hash_factors(items: dict[str, object]) -> str: @dataclass class Range: - """ - A range of numbers. + """A range of numbers. Inclusive of start, inclusive of end. """ @@ -445,8 +445,7 @@ def get_from_deprecated_env_if_set( removal_version: str, field_name: str | None = None, ) -> str | None: - """ - Get value from deprecated environment variable with warning. + """Get value from deprecated environment variable with warning. Args: env_name: Name of the deprecated environment variable @@ -455,6 +454,7 @@ def get_from_deprecated_env_if_set( Returns: The environment variable value if set, None otherwise + """ if envs.is_set(env_name): value = os.environ.get(env_name) @@ -477,8 +477,7 @@ def set_from_deprecated_env_if_set( to_bool: bool = False, to_int: bool = False, ) -> None: - """ - Set object field from deprecated environment variable with warning. + """Set object field from deprecated environment variable with warning. Args: config: Config object to set the field on @@ -489,6 +488,7 @@ def set_from_deprecated_env_if_set( to_int: Whether to convert the environment variable value to integer Returns: None + """ if to_bool and to_int: raise ValueError("Cannot convert to both boolean and integer.") diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7167ca5a9063..174e184c538d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -153,8 +153,8 @@ class OptimizationLevel(IntEnum): def enable_norm_fusion(cfg: "VllmConfig") -> bool: """Enable if either RMS norm or quant FP8 custom op is active; - otherwise Inductor handles fusion.""" - + otherwise Inductor handles fusion. + """ return ( cfg.compilation_config.is_custom_op_enabled("rms_norm") or cfg.compilation_config.is_custom_op_enabled("quant_fp8") @@ -163,8 +163,7 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool: def enable_act_fusion(cfg: "VllmConfig") -> bool: - """ - Enable if either SiLU+Mul or quant FP8 custom op is active; + """Enable if either SiLU+Mul or quant FP8 custom op is active; otherwise Inductor handles fusion. Also enable for FP4 models as FP4 quant is always custom so Inductor cannot fuse it. """ @@ -220,7 +219,6 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool: def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool: """Enable if use_inductor_graph_partition is enabled.""" - return ( cfg.compilation_config.use_inductor_graph_partition or not cfg.compilation_config.splitting_ops_contain_kv_cache_update() @@ -229,7 +227,6 @@ def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool: def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool: """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss.""" - return ( cfg.kernel_config.ir_op_priority.fused_add_rms_norm[0] == "aiter" and cfg.model_config is not None @@ -460,8 +457,7 @@ class VllmConfig: """ def compute_hash(self, include_version: bool = True) -> str: - """ - WARNING: Whenever a new field is added to this config, + """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph. @@ -473,6 +469,7 @@ def compute_hash(self, include_version: bool = True) -> str: Args: include_version: Include the vLLM version in the hash. + """ factors: list[Any] = [] @@ -747,7 +744,8 @@ def use_v2_model_runner(self) -> bool: def _is_dflash2_draft(self) -> bool: """Whether the DFlash draft is a DFlash2 one, by the architecture the - speculator selects on (v1/worker/gpu/spec_decode/__init__.py).""" + speculator selects on (v1/worker/gpu/spec_decode/__init__.py). + """ spec = self.speculative_config if spec is None or spec.method != "dflash": return False @@ -798,8 +796,7 @@ def _maybe_enable_breakable_cudagraph(self) -> bool: @property def needs_dp_coordinator(self) -> bool: - """ - Determine if the DPCoordinator process is needed. + """Determine if the DPCoordinator process is needed. The DPCoordinator is needed in two cases: 1. For MoE models with DP > 1: to handle wave coordination @@ -809,8 +806,8 @@ def needs_dp_coordinator(self) -> bool: Returns: True if DPCoordinator process is needed, False otherwise. - """ + """ # For non-MoE models, only need coordinator in internal/hybrid LB mode # (for stats collection). return self.parallel_config.data_parallel_size > 1 and ( @@ -820,8 +817,7 @@ def needs_dp_coordinator(self) -> bool: ) def enable_trace_function_call_for_thread(self) -> None: - """ - Set up function tracing for the current thread, + """Set up function tracing for the current thread, if enabled via the `VLLM_TRACE_FUNCTION` environment variable. """ if envs.VLLM_TRACE_FUNCTION: @@ -955,6 +951,7 @@ def _set_config_default(self, config_obj: Any, key: str, value: Any) -> None: config_obj: Configuration object to update. key: Attribute name. value: Default value (static or callable). + """ if getattr(config_obj, key) is None: # Some config values are known before initialization and are @@ -976,6 +973,7 @@ def _apply_optimization_level_defaults(self, defaults: dict[str, Any]) -> None: Args: defaults: Dictionary of default values to apply. + """ def apply_recursive(config_obj: Any, config_defaults: dict[str, Any]) -> None: @@ -1276,7 +1274,6 @@ def _resolve_and_verify_engram_config(self) -> None: def __post_init__(self): """Verify configs are valid & consistent with each other.""" - # To give each torch profile run a unique instance name. self.instance_id = f"{time.time_ns()}" @@ -2133,8 +2130,7 @@ def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list: ] def _set_max_num_scheduled_tokens(self): - """ - In most cases, the scheduler may schedule a batch with as many tokens as the + """In most cases, the scheduler may schedule a batch with as many tokens as the worker is configured to handle. """ if self.speculative_config is not None: @@ -2172,8 +2168,7 @@ def _set_max_num_scheduled_tokens(self): ) def _set_cudagraph_sizes(self): - """ - vLLM defines the default candidate list of batch sizes for CUDA graph + """VLLM defines the default candidate list of batch sizes for CUDA graph capture as: ```python @@ -2219,8 +2214,8 @@ def _set_cudagraph_sizes(self): padded CUDA graph will be used. - If batch size > largest `cudagraph_capture_sizes`, cudagraph will not be used. - """ + """ if ( self.model_config is not None and not self.model_config.enforce_eager @@ -2441,9 +2436,7 @@ def request_counts(max_reqs: int) -> list[int]: self.compilation_config.post_init_cudagraph_sizes() def _set_compile_ranges(self): - """ - Set the compile ranges for the compilation config. - """ + """Set the compile ranges for the compilation config.""" compilation_config = self.compilation_config computed_compile_ranges_endpoints = [] @@ -3310,8 +3303,7 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": def set_current_vllm_config( vllm_config: VllmConfig, check_compile=False, prefix: str | None = None ): - """ - Temporarily set the current vLLM config. + """Temporarily set the current vLLM config. Used during model initialization. We save the current vLLM config in a global variable, so that all modules can access it, e.g. custom ops @@ -3392,15 +3384,14 @@ def get_layers_from_vllm_config( layer_type: type[T], layer_names: Iterable[str] | None = None, ) -> dict[str, T]: - """ - Get layers from the vLLM config. + """Get layers from the vLLM config. Args: vllm_config: The vLLM config. layer_type: The type of the layer to get. layer_names: The names of the layers to get. If None, return all layers. - """ + """ forward_context = vllm_config.compilation_config.static_forward_context if layer_names is None: layer_names = forward_context.keys() diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index a16f1b01486e..63d4a1f4e0ec 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -80,8 +80,7 @@ def use_memory_pool_with_allocator( class CuMemAllocator: - """ - A singleton class that manages a memory pool for CUDA tensors. + """A singleton class that manages a memory pool for CUDA tensors. The memory in this pool can be offloaded or discarded when the allocator sleeps. @@ -109,8 +108,7 @@ class CuMemAllocator: @staticmethod def get_instance() -> "CuMemAllocator": - """ - CuMemAllocator is a singleton class. + """CuMemAllocator is a singleton class. We cannot call the constructor directly. Call this method to get the instance. """ @@ -182,9 +180,9 @@ def close(self) -> None: self.release_pools() def _python_malloc_callback(self, allocation_handle: HandleType) -> None: + """Internal method to store the allocation data + when memory is allocated in the memory pool. """ - Internal method to store the allocation data - when memory is allocated in the memory pool.""" py_d_mem = allocation_handle[2] self.pointer_to_data[py_d_mem] = AllocationData( allocation_handle, self.current_tag @@ -198,9 +196,9 @@ def _python_malloc_callback(self, allocation_handle: HandleType) -> None: return def _python_free_callback(self, ptr: int) -> HandleType: + """Internal method to look up the allocation data + when memory is freed in the memory pool. """ - Internal method to look up the allocation data - when memory is freed in the memory pool.""" data = self.pointer_to_data.pop(ptr) if data.cpu_backup_tensor is not None: data.cpu_backup_tensor = None @@ -227,14 +225,14 @@ def _python_free_callback(self, ptr: int) -> HandleType: return data.handle def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: - """ - Put the allocator in sleep mode. + """Put the allocator in sleep mode. All data in the memory allocation with the specified tag will be offloaded to CPU memory, and others will be discarded. Args: offload_tags: The tags of the memory allocation that will be offloaded. The rest of the memory allocation will be discarded. + """ if offload_tags is None: # by default, allocated tensors are offloaded @@ -325,8 +323,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None: ) def wake_up(self, tags: list[str] | None = None) -> None: - """ - Wake up the allocator from sleep mode. + """Wake up the allocator from sleep mode. All data that is previously offloaded will be loaded back to GPU memory, and the rest of the data will have empty memory. @@ -334,6 +331,7 @@ def wake_up(self, tags: list[str] | None = None) -> None: tags: The tags of the memory allocation that will be loaded back to GPU memory. If None, all memory allocation will be loaded back to GPU memory. + """ gc.collect() torch.accelerator.empty_cache() @@ -357,14 +355,14 @@ def wake_up(self, tags: list[str] | None = None) -> None: @contextmanager def use_memory_pool(self, tag: str | None = None): - """ - A context manager to use the memory pool. + """A context manager to use the memory pool. All memory allocation created inside the context will be allocated in the memory pool, and has the specified tag. Args: tag: The tag of the memory allocation. If None, the default tag will be used. + """ if tag is None: tag = CuMemAllocator.default_tag @@ -416,9 +414,7 @@ def use_memory_pool(self, tag: str | None = None): torch.cuda.memory._set_allocator_settings("expandable_segments:True") def get_current_usage(self) -> int: - """ - Get the total number of bytes allocated in the memory pool. - """ + """Get the total number of bytes allocated in the memory pool.""" sum_bytes: int = 0 for ptr, data in self.pointer_to_data.items(): handle = data.handle diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 45a635d61389..bf28278f7cde 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -70,7 +70,8 @@ def resume(self, tags: list[str] | None = None) -> None: def state(self) -> SleepModeState: """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle - (suspended) engine from a healthy-serving one (see RFC #34303).""" + (suspended) engine from a healthy-serving one (see RFC #34303). + """ return self._state # -- Capability introspection (no instance required) -- @@ -83,26 +84,30 @@ def is_supported(cls) -> bool: @classmethod def preserves_communicators(cls) -> bool: """If False, collective communicators (e.g. NCCL) are destroyed by - ``suspend`` and the executor must re-initialize them on ``resume``.""" + ``suspend`` and the executor must re-initialize them on ``resume``. + """ return False @classmethod def preserves_compiled_artifacts(cls) -> bool: """If True, torch.compile / JIT kernels survive suspend/resume and need - not be recompiled on resume.""" + not be recompiled on resume. + """ return False @classmethod def preserves_graphs_with_communicators(cls) -> bool: """If True, CUDA graphs containing collective communicators (e.g. NCCL) stay valid after resume. False when communicators are rebuilt (embedded - comm handles go stale).""" + comm handles go stale). + """ return False @classmethod def supports_durable_storage(cls) -> bool: """If True, suspended state can be persisted beyond the process - lifetime (disk or object storage) and restored in a new process.""" + lifetime (disk or object storage) and restored in a new process. + """ return False diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 146ea965db07..17fa69ee8411 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -116,6 +116,7 @@ class XpuMemAllocator: Sleep will offload selected payloads to CPU or discard and unmap XPU physical memory. Wake-up remaps physical memory back to the same reserved virtual address and restores payload. + """ instance: "XpuMemAllocator | None" = None diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 3247b33262e2..6fbffcaaf73d 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -22,9 +22,7 @@ class AiterCustomAllreduce: @classmethod def effective_max_size(cls) -> int: - """ - Max input byte size eligible for AITER custom allreduce. - """ + """Max input byte size eligible for AITER custom allreduce.""" return cls.MAX_SIZE // 2 def __init__( diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 080e7feeff1f..caa3e3386ca0 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -44,8 +44,7 @@ class AgRsAll2AllManager(All2AllManagerBase): - """ - An implementation of all2all communication based on + """An implementation of all2all communication based on all-gather (dispatch) and reduce-scatter (combine). """ @@ -79,9 +78,7 @@ def dispatch_router_logits( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Gather hidden_states and router_logits from all dp ranks. - """ + """Gather hidden_states and router_logits from all dp ranks.""" dist_group = self._get_comm_group(is_sequence_parallel) sizes = self._get_sizes(hidden_states.shape[0], dist_group) assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] @@ -111,9 +108,7 @@ def dispatch( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Gather hidden_states and router_logits from all dp ranks. - """ + """Gather hidden_states and router_logits from all dp ranks.""" dist_group = self._get_comm_group(is_sequence_parallel) sizes = self._get_sizes(hidden_states.shape[0], dist_group) assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] @@ -140,9 +135,7 @@ def dispatch( def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: - """ - Reduce-scatter hidden_states across all dp ranks. - """ + """Reduce-scatter hidden_states across all dp ranks.""" dist_group = self._get_comm_group(is_sequence_parallel) sizes = self._get_sizes( hidden_states.shape[0] // dist_group.world_size, @@ -156,9 +149,7 @@ def destroy(self): class DeepEPAll2AllManagerBase(All2AllManagerBase): - """ - All2All communication based on DeepEP High-Throughput kernels. - """ + """All2All communication based on DeepEP High-Throughput kernels.""" def __init__(self, cpu_group, tcp_store_group=None): assert has_deep_ep(), ( @@ -210,9 +201,7 @@ def destroy(self): class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase): - """ - All2All communication based on DeepEP High-Throughput kernels. - """ + """All2All communication based on DeepEP High-Throughput kernels.""" def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) @@ -271,9 +260,7 @@ def set_num_sms(self, num_sms: int): class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): - """ - All2All communication based on DeepEP Low-Latency kernels. - """ + """All2All communication based on DeepEP Low-Latency kernels.""" _buffer: Any = None _mask: torch.Tensor | None = None @@ -293,8 +280,7 @@ def _make_all2all_kwargs( num_global_experts: int, num_local_experts: int, ) -> dict[Any, Any]: - """ - max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank + """max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank can dispatch all the ranks must hold the same value. token_hidden_size: the hidden dimension of each token. num_ep_ranks: the number of EP group ranks. @@ -330,8 +316,7 @@ def _make_all2all_kwargs( return kwargs def get_handle(self, kwargs): - """ - The kwargs for DeepEPLLAll2AllManager is dictated by + """The kwargs for DeepEPLLAll2AllManager is dictated by _make_all2all_kwargs. """ import deep_ep # type: ignore[import-not-found] @@ -385,8 +370,7 @@ class _NixlEPBufferState: class NixlEPAll2AllManager(All2AllManagerBase): - """ - All2All communication based on NIXL EP kernels. + """All2All communication based on NIXL EP kernels. This backend supports elastic EP with dynamic rank connection/disconnection. """ @@ -612,9 +596,7 @@ def clean_buffers(self) -> None: class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): - """ - All2All communication based on flashinfer all2allv/two-sided NVLink kernels. - """ + """All2All communication based on flashinfer all2allv/two-sided NVLink kernels.""" # This type lint could be removed after all of the work in # https://github.com/vllm-project/vllm/issues/26533 done. @@ -640,7 +622,7 @@ def initialize( rank: int, gpus_per_node: int, ): - """Initialize workspace""" + """Initialize workspace.""" if self.initialized: return @@ -681,7 +663,7 @@ def initialize( ) def ensure_alltoall_workspace_initialized(self): - """Ensure workspace is initialized""" + """Ensure workspace is initialized.""" if not has_flashinfer_nvlink_two_sided(): return False @@ -700,7 +682,7 @@ def get_handle(self, kwargs): return self def cleanup(self): - """Clean up workspace""" + """Clean up workspace.""" if ( self.initialized and self.workspace_tensor is not None @@ -719,8 +701,7 @@ def cleanup(self): class FlashInferNVLinkOneSidedManager(All2AllManagerBase): - """ - All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel. + """All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel. This is a newer kernel from trtllm that should perform better than the kernel used by flashinfer_nvlink_two_sided. """ @@ -1023,8 +1004,7 @@ def get_handle(self, kwargs): class DeepEPV2All2AllManager(All2AllManagerBase): - """ - All2All communication based on DeepEP v2 ElasticBuffer (unified API). + """All2All communication based on DeepEP v2 ElasticBuffer (unified API). Uses NCCL Gin backend with analytical SM calculation. """ diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index 51c28788cc01..43e615a3984b 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -119,8 +119,7 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool: - """ - Determine if NCCL symmetric memory allreduce should be used. + """Determine if NCCL symmetric memory allreduce should be used. Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for: - Small tensors (≤16K): Lower latency than custom_AR @@ -157,7 +156,8 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) def should_nccl_symm_mem_ag_rs() -> bool: """Check whether NCCL symmetric memory should be used for - AllGather / ReduceScatter collectives.""" + AllGather / ReduceScatter collectives. + """ from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_enabled, ) @@ -344,7 +344,6 @@ def can_actually_p2p( def gpu_p2p_access_check(src: int, tgt: int) -> bool: """Check if GPU src can access GPU tgt.""" - # if the cache variable is already calculated, # read from the cache instead of checking it again global _gpu_p2p_access_cache diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 4c1900590b99..f70bf0921518 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -114,6 +114,7 @@ def query_active_mask(self) -> torch.Tensor: Returns: An int32 device tensor where 0 marks a live rank and 1 marks a masked (dead/unreachable) rank. + """ raise NotImplementedError @@ -171,8 +172,7 @@ def mask_remote_ranks(self) -> Iterator[None]: class DeviceCommunicatorBase: - """ - Base class for device-specific communicator. + """Base class for device-specific communicator. It can use the `cpu_group` to initialize the communicator. If the device has PyTorch integration (PyTorch can recognize its communication backend), the `device_group` will also be given. @@ -318,8 +318,7 @@ def reduce_scatterv( def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: - """ - NOTE: We assume that the input tensor is on the same device across + """NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ @@ -347,7 +346,7 @@ def gather( return output_tensor def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: - """Sends a tensor to the destination rank in a blocking way""" + """Sends a tensor to the destination rank in a blocking way.""" """NOTE: `dst` is the local rank of the destination rank.""" if dst is None: dst = (self.rank_in_group + 1) % self.world_size @@ -385,8 +384,7 @@ def dispatch_router_logits( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and router logits to the appropriate device. + """Dispatch the hidden states and router logits to the appropriate device. This is a no-op in the base class. """ if extra_tensors is not None: @@ -404,8 +402,7 @@ def dispatch( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and topk weights/ids to the appropriate device. + """Dispatch the hidden states and topk weights/ids to the appropriate device. This is a no-op in the base class. """ if extra_tensors is not None: @@ -415,8 +412,7 @@ def dispatch( def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: - """ - Combine the hidden states and router logits from the appropriate device. + """Combine the hidden states and router logits from the appropriate device. This is a no-op in the base class. """ return hidden_states diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 8ea12d9255a9..c190c9ded691 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -69,8 +69,7 @@ def __init__( logger.info("Using allgather_reducescatter all2all manager.") def _all_group_ranks_share_shm_group_name(self) -> bool: - """ - CPUSHM requires all ranks in this group to agree on one SHM group name. + """CPUSHM requires all ranks in this group to agree on one SHM group name. This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs. """ local_name = _CPUSHMDistributed.make_group_name(self) @@ -89,8 +88,7 @@ def all_reduce(self, input_): def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: - """ - NOTE: We assume that the input tensor is on the same device across + """NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ @@ -180,11 +178,9 @@ def dispatch_router_logits( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and router logits to the appropriate device. + """Dispatch the hidden states and router logits to the appropriate device. This is a no-op in the base class. """ - assert self.all2all_manager is not None return self.all2all_manager.dispatch_router_logits( hidden_states, @@ -204,8 +200,7 @@ def dispatch( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and topk weights/ids to the appropriate device. + """Dispatch the hidden states and topk weights/ids to the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None @@ -220,8 +215,7 @@ def dispatch( def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: - """ - Combine the hidden states and router logits from the appropriate device. + """Combine the hidden states and router logits from the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 4afc28f65214..fb4ccd9f1e51 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -611,7 +611,7 @@ def _reduce_scatter_symm_mem( return symm_output def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: - """Sends a tensor to the destination rank in a blocking way""" + """Sends a tensor to the destination rank in a blocking way.""" """NOTE: `dst` is the local rank of the destination rank.""" if dst is None: dst = (self.rank_in_group + 1) % self.world_size @@ -851,11 +851,9 @@ def dispatch_router_logits( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and router logits to the appropriate device. + """Dispatch the hidden states and router logits to the appropriate device. This is a no-op in the base class. """ - assert self.all2all_manager is not None return self.all2all_manager.dispatch_router_logits( hidden_states, @@ -875,8 +873,7 @@ def dispatch( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and topk weights/ids to the appropriate device. + """Dispatch the hidden states and topk weights/ids to the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None @@ -891,8 +888,7 @@ def dispatch( def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: - """ - Combine the hidden states and router logits from the appropriate device. + """Combine the hidden states and router logits from the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 6c15a9827429..e18f9194f54e 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -129,8 +129,7 @@ def __init__( max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE, symm_mem_enabled=False, ) -> None: - """ - Args: + """Args: group: the process group to work on. If None, it will use the default process group. device: the device to bind the CustomAllreduce to. If None, @@ -138,6 +137,7 @@ def __init__( It is the caller's responsibility to make sure each communicator is bind to a unique device, and all communicators in this group are in the same node. + """ self._IS_CAPTURING = False self._ptr = 0 @@ -370,8 +370,7 @@ def _init_mnnvl_buffer(self, stage_size: int) -> None: @contextmanager def capture(self): - """ - The main responsibility of this context manager is the + """The main responsibility of this context manager is the `register_graph_buffers` call at the end of the context. It records all the buffer addresses used in the CUDA graph. """ diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index 904bf2dae0c1..f1644d699ed6 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -140,6 +140,7 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]: is True only when ``auto`` selects mnnvl for a single node, so that workspace creation can fall back to trtllm on single-node topologies without NVSwitch multicast support (where mnnvl is unavailable). + """ backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND if backend != "auto": @@ -168,8 +169,7 @@ def get_fi_ar_workspace( dtype: torch.dtype, group: ProcessGroup, ): - """ - Return the allreduce workspace for non-quant patterns, initializing if needed. + """Return the allreduce workspace for non-quant patterns, initializing if needed. Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce for standalone allreduce. Backend is controlled by @@ -234,8 +234,7 @@ def get_fi_ar_quant_workspace( dtype: torch.dtype, group: ProcessGroup, ): - """ - Return the allreduce workspace for quant patterns, initializing if needed. + """Return the allreduce workspace for quant patterns, initializing if needed. Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index ece4db82dae2..76eeb2291f90 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/vllm/distributed/device_communicators/pynccl.py @@ -98,8 +98,7 @@ def __init__( device: int | str | torch.device, library_path: str | None = None, ): - """ - Args: + """Args: group: the process group to work on. If None, it will use the default process group. device: the device to bind the PyNcclCommunicator to. If None, @@ -108,6 +107,7 @@ def __init__( use the default library path. It is the caller's responsibility to make sure each communicator is bind to a unique device. + """ if not isinstance(group, StatelessProcessGroup): assert dist.is_initialized() diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 3c9d759dd7f1..406869ea2c4b 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -62,8 +62,7 @@ class QuickAllReduce: } def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None: - """ - Custom allreduce provides non-destructive acceleration and is + """Custom allreduce provides non-destructive acceleration and is available for CUDA and ROCm MI300 series. Custom quick allreduce leverages quantization for further @@ -86,6 +85,7 @@ def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> Non It is the caller's responsibility to make sure each communicator is bind to a unique device, and all communicators in this group are in the same node. + """ self.disabled = True if not self._rocm_arch_available(): @@ -301,8 +301,7 @@ def _rocm_arch_available(self): return False def create_shared_buffer(self): - """ - Creates a shared buffer for quickreduce. + """Creates a shared buffer for quickreduce. Has to be called after init_custom_qr """ handle = ops.qr_get_handle(self._ptr) @@ -312,9 +311,7 @@ def create_shared_buffer(self): ops.qr_open_handles(self._ptr, handles) def should_quick_allreduce(self, inp: torch.Tensor): - """ - Check if quickreduce is available - """ + """Check if quickreduce is available.""" if self.disabled: return False if inp.dtype not in self._SUPPORTED_DTYPES: diff --git a/vllm/distributed/device_communicators/ray_communicator.py b/vllm/distributed/device_communicators/ray_communicator.py index d9517f51acad..ff3c2f1fd479 100644 --- a/vllm/distributed/device_communicators/ray_communicator.py +++ b/vllm/distributed/device_communicators/ray_communicator.py @@ -20,8 +20,7 @@ class RayPPCommunicator(Communicator): - """ - Communicator to be used for pipeline parallelism in Ray Compiled Graph. + """Communicator to be used for pipeline parallelism in Ray Compiled Graph. This is wraps around the vLLM _PP GroupCoordinator. This class is not thread-safe. @@ -38,8 +37,7 @@ def __init__( cuda_stream: torch.cuda.Stream | None, use_communication_streams: bool = False, ): - """ - Initialize a RayPPCommunicator that can be used to communicate with + """Initialize a RayPPCommunicator that can be used to communicate with other Ray Compiled Graph actors for pipeline parallelism. Args: @@ -54,6 +52,7 @@ def __init__( is not supported. use_communication_streams: Whether to use communication streams. This is not supported. + """ self._world_size = world_size self._rank: int | None = None @@ -87,8 +86,7 @@ def __init__( self._closed = False def _build_actor_rank_mapping(self): - """ - Use collective communication to build a mapping from actor IDs to ranks. + """Use collective communication to build a mapping from actor IDs to ranks. This should be called once during initialization. """ if self._comm is None: @@ -128,9 +126,7 @@ def get_actor_handles(self) -> list["ray.actor.ActorHandle"]: return self._actor_handles def get_rank(self, actor: ray.actor.ActorHandle) -> int: - """ - Return the given actor's rank using device communicator collective ops. - """ + """Return the given actor's rank using device communicator collective ops.""" assert hasattr(self, "_actor_id_to_rank"), ( "Actor rank mapping not built. " "This should have been done during initialization." @@ -144,20 +140,15 @@ def get_rank(self, actor: ray.actor.ActorHandle) -> int: raise ValueError(f"Actor {actor} not found in communicator group") def get_self_rank(self) -> int | None: - """ - Return this actor's rank. - """ + """Return this actor's rank.""" return self._rank def get_world_size(self) -> int: - """ - Return the number of ranks in the RayPPCommunicator group. - """ + """Return the number of ranks in the RayPPCommunicator group.""" return self._world_size def send(self, buf: "torch.Tensor", peer_rank: int) -> None: - """ - Send a torch.Tensor to a peer. + """Send a torch.Tensor to a peer. This returns when the send kernel has been queued, but the kernel may not have completed. Therefore, the caller should ensure that there are @@ -170,6 +161,7 @@ def send(self, buf: "torch.Tensor", peer_rank: int) -> None: buf: The torch.Tensor to send. It should already be on this actor's default device. peer_rank: The rank of the actor to send to. + """ if self._closed: raise RayChannelError("RayPPCommunicator has been destroyed.") @@ -184,8 +176,7 @@ def recv( peer_rank: int, allocator: TorchTensorAllocator, ) -> "torch.Tensor": - """ - Receive a torch.Tensor from a peer and synchronize the current stream. + """Receive a torch.Tensor from a peer and synchronize the current stream. After this call returns, the receive buffer is safe to read from any stream. An RayChannelError will be raised if an error occurred @@ -197,6 +188,7 @@ def recv( peer_rank: The rank of the actor to receive from. allocator: The allocator to use to create the received tensor. This is ignored for this implementation. + """ if self._closed: raise RayChannelError("RayPPCommunicator has been destroyed.") diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 2171f79fb7ce..f457cf61e8a9 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -77,8 +77,7 @@ def memory_fence(): - """ - Full memory barrier for shared memory synchronization. + """Full memory barrier for shared memory synchronization. Ensures all prior memory writes are visible to other processes before any subsequent reads. This is critical for lock-free producer-consumer @@ -111,8 +110,7 @@ def to_bytes_big(value: int, size: int) -> bytes: class SpinCondition: - """ - This class implements an interface similar to a threading.Condition. It + """This class implements an interface similar to a threading.Condition. It allows a writer to notify readers to wake up and read from the shared memory buffer. This notification is done over a zmq socket. @@ -241,6 +239,7 @@ def check_shm_free_space( Raises: RuntimeError: If the SHM filesystem has insufficient space. + """ if os.path.isdir(shm_path): free_bytes = shutil.disk_usage(shm_path).free @@ -651,7 +650,8 @@ def wait_until_ready(self): def shutdown(self): """If this is an idle reader, wakes it up so it can clean up and shut - down""" + down + """ self.shutting_down = True if self._spin_condition is not None: self._spin_condition.cancel() @@ -761,7 +761,8 @@ def timeout_ms(self) -> int: def should_warn(self) -> bool: """Returns true if it's time to log a warning for a timeout that is not - indefinite""" + indefinite + """ if self._should_warn: elapsed = time.monotonic() - self.started if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning: @@ -941,8 +942,7 @@ def create_from_process_group_single_reader( reader_rank: int = 0, blocking: bool = False, ) -> tuple["MessageQueue", list[Handle]]: - """ - Creates a MessageQueue for a process group with a single reader. + """Creates a MessageQueue for a process group with a single reader. This method is designed for scenarios where only one process (the reader) will consume messages, and all other processes are writers. It sets up @@ -962,6 +962,7 @@ def create_from_process_group_single_reader( tuple[MessageQueue, list[Handle]]: The MessageQueue instance for the calling process, and a list of handles (only non-empty for the reader process). + """ from vllm.platforms.interface import get_assigned_physical_gpu_ids @@ -994,8 +995,7 @@ def create_from_process_group( external_writer_handle=None, blocking: bool = True, ) -> "MessageQueue": - """ - Creates a MessageQueue for a distributed process group with one writer and + """Creates a MessageQueue for a distributed process group with one writer and multiple readers. This method is designed for scenarios where one process (the writer) sends diff --git a/vllm/distributed/device_communicators/shm_object_storage.py b/vllm/distributed/device_communicators/shm_object_storage.py index e2d2b248346b..bc4f1b6e5159 100644 --- a/vllm/distributed/device_communicators/shm_object_storage.py +++ b/vllm/distributed/device_communicators/shm_object_storage.py @@ -20,8 +20,7 @@ class SingleWriterShmRingBuffer: - """ - A single-writer, multiple-reader ring buffer implementation using shared + """A single-writer, multiple-reader ring buffer implementation using shared memory. This class provides a thread-safe ring buffer where one process can write data while multiple processes/threads can read from it. @@ -190,8 +189,7 @@ def byte2int(self, byte_data: bytes) -> int: return int.from_bytes(byte_data, "little", signed=True) def allocate_buf(self, size: int) -> tuple[int, int]: - """ - Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory. + """Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory. Memory layout: `[4-byte monotonic_id][4-byte size][buffer data...]` """ @@ -260,8 +258,7 @@ def free_buf( is_free_fn: Callable[[int, memoryview], bool], nbytes: int | None = None, ) -> Iterable[int]: - """ - Free a buffer of the given size. This is a no-op in shared memory, + """Free a buffer of the given size. This is a no-op in shared memory, but we need to keep track of the metadata. If freed memory spreads across the end and start of the ring buffer, @@ -269,10 +266,13 @@ def free_buf( still might not be a contiguous space of `nbytes` available. Args: + is_free_fn (Callable[[int, memoryview], bool]): Predicate called with + a monotonic id and the buffer, returning True when that buffer + can be reclaimed. nbytes (int, optional): The size of the buffer to free. If None, frees the maximum size of the ring buffer. - """ + """ assert self.is_writer, "Only the writer can free buffers." logger.debug( "Freeing up space in the ring buffer, " @@ -412,8 +412,7 @@ class ShmObjectStorageHandle: class SingleWriterShmObjectStorage: - """ - A single-writer, multiple-reader object storage system built on top of a + """A single-writer, multiple-reader object storage system built on top of a shared memory ring buffer. Provides key-value storage with automatic memory management and cross-process serialization support. @@ -458,8 +457,7 @@ def __init__( serde_class: type[ObjectSerde] = MsgpackSerde, reader_lock: LockType | None = None, ): - """ - Initialize the object storage. + """Initialize the object storage. Args: max_object_size: Maximum size for a single object in bytes. @@ -467,10 +465,11 @@ def __init__( ring_buffer: The shared memory ring buffer for storing objects. serde_class: Serializer/deserializer for objects. reader_lock: Optional lock for synchronizing reader access. + Raises: ValueError: If reader_lock is None for readers. - """ + """ self.max_object_size = max_object_size self.n_readers = n_readers self.serde_class = serde_class @@ -547,22 +546,17 @@ def free_unused(self) -> None: del self.writer_flag[freed_id] def is_cached(self, key: str) -> bool: - """ - Check if the object with the given key is cached. - """ + """Check if the object with the given key is cached.""" return key in self.key_index def get_cached(self, key: str) -> tuple[int, int]: - """ - Get the cached object by key if it exists. - """ + """Get the cached object by key if it exists.""" address, monotonic_id = self.key_index[key] self.increment_writer_flag(monotonic_id) return address, monotonic_id def put(self, key: str, value: Any) -> tuple[int, int]: - """ - Store a key-value pair in the object storage. + """Store a key-value pair in the object storage. Attempts to free max_object_size bytes using FIFO order when the ring buffer runs out of space during a put() operation. @@ -574,6 +568,7 @@ def put(self, key: str, value: Any) -> tuple[int, int]: MemoryError: If there's not enough space in the buffer ValueError: If the serialized object is too large ValueError: If the key already exists in the storage + """ if key in self.key_index: raise ValueError(f"Key '{key}' already exists in the storage.") @@ -639,8 +634,7 @@ def touch( address: int = 0, monotonic_id: int = 0, ) -> None: - """ - Touch an existing cached item to update its eviction status. + """Touch an existing cached item to update its eviction status. For writers (ShmObjectStoreSenderCache): Increment writer_flag For readers (ShmObjectStoreReceiverCache): Increment reader_count @@ -700,8 +694,7 @@ def create_from_handle( ) def default_is_free_check(self, id: int, buf: memoryview) -> bool: - """ - Default is_free function that checks if the first 4 bytes are zero. + """Default is_free function that checks if the first 4 bytes are zero. This indicates that the buffer is free. """ reader_count = int.from_bytes(buf[0:4], "little", signed=True) diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index 7ca132824eca..58c64e79200b 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -205,11 +205,9 @@ def dispatch_router_logits( tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and router logits to the appropriate device. + """Dispatch the hidden states and router logits to the appropriate device. This is a no-op in the base class. """ - assert self.all2all_manager is not None return self.all2all_manager.dispatch_router_logits( hidden_states, @@ -229,8 +227,7 @@ def dispatch( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] ): - """ - Dispatch the hidden states and topk weights/ids to the appropriate device. + """Dispatch the hidden states and topk weights/ids to the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None @@ -245,8 +242,7 @@ def dispatch( def combine( self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False ) -> torch.Tensor: - """ - Combine the hidden states and router logits from the appropriate device. + """Combine the hidden states and router logits from the appropriate device. This is a no-op in the base class. """ assert self.all2all_manager is not None diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index ea44fb6f5beb..ec3166ce3381 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ECConnectorBase Class for Distributed Encoder Cache & +"""ECConnectorBase Class for Distributed Encoder Cache & P2P Encoder cache communication in V1 The class provides the following primitives: @@ -50,8 +49,7 @@ class ECConnectorRole(enum.Enum): class ECConnectorMetadata(ABC): # noqa: B024 - """ - Abstract Metadata used to communicate between the + """Abstract Metadata used to communicate between the Scheduler ECConnector and Worker ECConnector. """ @@ -59,8 +57,7 @@ class ECConnectorMetadata(ABC): # noqa: B024 class ECConnectorWorkerMetadata(ABC): - """ - Abstract Metadata used to communicate back + """Abstract Metadata used to communicate back Worker ECConnector -> Scheduler ECConnector. Each worker can output its own metadata. @@ -73,9 +70,7 @@ class ECConnectorWorkerMetadata(ABC): def aggregate( self, other: "ECConnectorWorkerMetadata" ) -> "ECConnectorWorkerMetadata": - """ - Aggregate metadata with another `ECConnectorWorkerMetadata` object. - """ + """Aggregate metadata with another `ECConnectorWorkerMetadata` object.""" pass @@ -103,8 +98,7 @@ def is_consumer(self) -> bool: return self._is_consumer def shutdown(self) -> None: - """ - Shutdown the connector. This is called when the process + """Shutdown the connector. This is called when the process is shutting down to ensure that all the async operations are completed and the connector is cleaned up properly. """ @@ -123,6 +117,7 @@ def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> No Args: connector_metadata (dict): the connector metadata. + """ self._connector_metadata = connector_metadata @@ -141,8 +136,8 @@ def _get_connector_metadata(self) -> ECConnectorMetadata: Returns: ConnectorMetadata: the connector metadata. - """ + """ # Should only be called while set to valid metadata. assert self._connector_metadata is not None return self._connector_metadata @@ -151,10 +146,11 @@ def register_caches( self, ec_caches: dict[str, torch.Tensor], ): - """ - Initialize with the EC caches. + """Initialize with the EC caches. + Args: ec_caches: dictionary of encoder cache + """ # TODO: Implement this later for P2P feature return @@ -163,8 +159,7 @@ def register_caches( def start_load_caches( self, encoder_cache: dict[str, torch.Tensor], **kwargs ) -> None: - """ - Start loading the cache from the connector into vLLM's encoder cache. + """Start loading the cache from the connector into vLLM's encoder cache. This method loads the encoder cache based on metadata provided by the scheduler. It is called before `_gather_mm_embeddings` for the EC Connector. For EC, @@ -174,6 +169,7 @@ def start_load_caches( encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal data hashes (`mm_hash`) to encoder cache tensors. kwargs (dict): Additional keyword arguments for the connector. + """ pass @@ -181,8 +177,7 @@ def start_load_caches( def save_caches( self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs ) -> None: - """ - Save the encoder cache to the connector. + """Save the encoder cache to the connector. This method saves the encoder cache from the worker's local storage to shared storage or another external connector. @@ -192,14 +187,14 @@ def save_caches( data hashes (`mm_hash`) to encoder cache tensors. mm_hash (str): The hash of the multimodal data whose cache is being saved. kwargs (dict): Additional keyword arguments for the connector. + """ pass def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """ - Notifies worker-side connector ids of requests that have + """Notifies worker-side connector ids of requests that have finished generating tokens on the worker. The scheduler process (via the Executors) will use this output to track which workers are done. @@ -210,16 +205,17 @@ def get_finished( tuple of (sending/saving ids, recving/loading ids). The finished saves/sends req ids must belong to a set provided in a call to this method (this call or a prior one). + """ return None, None def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: - """ - Build the ECConnector worker metadata for this engine step. + """Build the ECConnector worker metadata for this engine step. Returns: ECConnectorWorkerMetadata: the worker metadata. None if no worker metadata is available. + """ return None @@ -232,8 +228,7 @@ def has_cache_item( self, identifier: str, ) -> bool: - """ - Check if a single encoder cache exists + """Check if a single encoder cache exists. Args: identifier (str): the identifier of the media. @@ -241,14 +236,14 @@ def has_cache_item( Returns: A bool where value is True if cache exist for the media + """ pass def ensure_cache_available( self, request: "Request", num_computed_tokens: int ) -> bool: - """ - Ensure encoder cache items are available for the given request. + """Ensure encoder cache items are available for the given request. May initiate asynchronous transfers for items not yet local. Args: @@ -258,6 +253,7 @@ def ensure_cache_available( Returns: True if all items are ready or no transfer is needed. False if any items are still in transit (request should be deferred). + """ return True @@ -282,11 +278,12 @@ def update_state_after_free(self, request: "Request", index: int): @abstractmethod def update_state_after_alloc(self, request: "Request", index: int): - """ - Update ECConnector state to decide allocate cache for requests + """Update ECConnector state to decide allocate cache for requests. Args: request (Request): the request object. + index (int): index of the encoder cache entry being allocated. + """ pass @@ -294,37 +291,37 @@ def update_state_after_alloc(self, request: "Request", index: int): def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> ECConnectorMetadata: - """ - Build the connector metadata for this step. + """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. Also, calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ pass def update_connector_output(self, connector_output: ECConnectorOutput): - """ - Update ECConnector state from worker-side connectors output. + """Update ECConnector state from worker-side connectors output. Args: connector_output (ECConnectorOutput): the worker-side connectors output. + """ return def request_finished( self, request: "Request" ) -> tuple[bool, dict[str, Any] | None]: - """ - Called when a request has finished, before its encoder cache is freed. + """Called when a request has finished, before its encoder cache is freed. Returns: True if the request is being saved/sent asynchronously and cached should not be freed until the request_id is returned from get_finished(). + """ return False, None diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py index 413bf285712a..bfc8aa009e15 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Lightweight mmap-backed shared memory region for encoder cache (EC) data. +"""Lightweight mmap-backed shared memory region for encoder cache (EC) data. Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified for EC: flat shared layout, no multi-tensor cursor, no block_size_factor. diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index cf0a92dcc26c..f1d22a5714a6 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -229,6 +229,7 @@ def _nixl_consumer_admit( Returns: True if no item needs to wait for a remote encoding. False if any item is pending or the request must be aborted. + """ if not request.ec_transfer_params: return True @@ -259,6 +260,7 @@ def _admit_item( Raises: _RemoteUnavailable: The remote source is invalid, the read failed, or the item's total wait budget expired. + """ pos = feature.mm_position mm_hash = feature.identifier @@ -334,6 +336,7 @@ def _fail_or_fallback( item's remote announcement so subsequent steps do not retry it. False if only placeholder metadata (or no input) is available, recording the request for take_unavailable_requests() to drain. + """ mm_hash = feature.identifier self._deferred_since.pop((request.request_id, mm_hash), None) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index 2fd4b01059db..afadcfd42f20 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -15,7 +15,8 @@ class CacheEntry: """A single cache entry. Read `.block_ids` and `.ready` freely; - mutations only through EmbeddingCache methods.""" + mutations only through EmbeddingCache methods. + """ __slots__ = ("block_ids", "_pin_count") diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py index 2db12f903c71..e080f09b6bef 100644 --- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/example_connector.py @@ -67,8 +67,7 @@ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): raise ValueError("ec_transfer_config must be set for ECConnectorBase") def start_load_caches(self, encoder_cache, **kwargs) -> None: - """ - Start loading the cache from the connector into vLLM's encoder cache. + """Start loading the cache from the connector into vLLM's encoder cache. This method loads the encoder cache based on metadata provided by the scheduler. It is called before `_gather_mm_embeddings` for the EC Connector. For EC, @@ -78,6 +77,7 @@ def start_load_caches(self, encoder_cache, **kwargs) -> None: encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal data hashes (`mm_hash`) to encoder cache tensors. kwargs (dict): Additional keyword arguments for the connector. + """ from vllm.platforms import current_platform @@ -108,8 +108,7 @@ def start_load_caches(self, encoder_cache, **kwargs) -> None: logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash) def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None: - """ - Save the encoder cache to the connector. + """Save the encoder cache to the connector. This method saves the encoder cache from the worker's local storage to shared storage or another external connector. @@ -119,6 +118,7 @@ def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None: data hashes (`mm_hash`) to encoder cache tensors. mm_hash (str): The hash of the multimodal data whose cache is being saved. kwargs (dict): Additional keyword arguments for the connector. + """ # Return if it is PD Instance if not self.is_producer: @@ -134,14 +134,14 @@ def has_cache_item( self, identifier: str, ) -> bool: - """ - Check if cache exist externally for the media + """Check if cache exist externally for the media. Args: identifier (str): the identifier of the media. Returns: Bool indicate that media exists in cache or not + """ return self._found_match_for_mm_data(identifier) @@ -150,9 +150,7 @@ def update_state_after_alloc( request: "Request", index: int, ) -> None: - """ - Update ECConnector state after encoder cache allocation. - """ + """Update ECConnector state after encoder cache allocation.""" mm_hash = request.mm_features[index].identifier # Only load cache if it is consumer and cache exists if not self.is_consumer or not self.has_cache_item(mm_hash): @@ -212,8 +210,7 @@ def _generate_foldername_debug( mm_hash: str, create_folder: bool = True, # <- now defaults to True ) -> str: - """ - Return the folder in which the cache for this mm_hash lives. + """Return the folder in which the cache for this mm_hash lives. If `create_folder` is True (default) the directory is created recursively the first time it is needed. """ @@ -223,8 +220,7 @@ def _generate_foldername_debug( return foldername def _generate_filename_debug(self, mm_hash: str) -> str: - """ - Return the full path of the safetensors file for this mm_hash. + """Return the full path of the safetensors file for this mm_hash. Ensures the parent directory exists because `_generate_foldername_debug` is called with its default (`create_folder=True`). diff --git a/vllm/distributed/ec_transfer/ec_transfer_state.py b/vllm/distributed/ec_transfer/ec_transfer_state.py index e58415ef6f4b..798fc42d842d 100644 --- a/vllm/distributed/ec_transfer/ec_transfer_state.py +++ b/vllm/distributed/ec_transfer/ec_transfer_state.py @@ -24,10 +24,7 @@ def has_ec_transfer() -> bool: def ensure_ec_transfer_initialized(vllm_config: "VllmConfig") -> None: - """ - Initialize EC cache connector. - """ - + """Initialize EC cache connector.""" global _EC_CONNECTOR_AGENT if vllm_config.ec_transfer_config is None: diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index ff872c8c2220..5cf63fe5aecd 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -The async worker that transfers experts in the background. -""" +"""The async worker that transfers experts in the background.""" import threading from typing import TYPE_CHECKING diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 162c0f58430b..d05ca8e00cde 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -EPLB communicator implementations and factory. -""" +"""EPLB communicator implementations and factory.""" import contextlib import time @@ -86,7 +84,8 @@ def set_transfer_context( # noqa: B027 @property def needs_profile_buffer_reservation(self) -> bool: """Whether the profile path must run a dummy collective operation to reserve - communication buffers.""" + communication buffers. + """ return True def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: @@ -258,6 +257,7 @@ def __init__( cpu_group: CPU process group for metadata exchange. all_expert_weights: Expert weight tensors for all MoE layers. expert_buffer: Pre-allocated receive buffer tensors. + """ assert all_expert_weights, ( "NixlEplbCommunicator requires non-empty all_expert_weights." @@ -443,7 +443,8 @@ def _init_registered_buffers(self) -> None: def _exchange_remote_send_meta(self) -> None: """Exchange per-layer per-tensor metadata so receivers can compute - remote RDMA addresses at transfer time.""" + remote RDMA addresses at transfer time. + """ local_meta: dict[tuple[int, int], tuple[int, int, int]] = {} for layer_idx, layer_tensors in enumerate(self._all_expert_weights): for t_idx, t in enumerate(layer_tensors): @@ -666,6 +667,7 @@ def create_eplb_communicator( zero-copy RDMA reads. expert_buffer: Pre-allocated receive buffer tensors (one per weight tensor in a single layer). + """ first_layer = expert_weights[0] if expert_weights else [] tensor_device_type = first_layer[0].device.type if first_layer else "cpu" diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index bc3aa36bb2c5..fb5475a607f5 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Expert parallelism load balancer (EPLB) metrics and states. +"""Expert parallelism load balancer (EPLB) metrics and states. # Glossary @@ -74,9 +73,7 @@ def _compute_eplb_load_stats( @dataclass class EplbStats: - """ - Model stats used in EPLB rebalancing algorithm. - """ + """Model stats used in EPLB rebalancing algorithm.""" global_expert_load_window: torch.Tensor """ @@ -233,9 +230,7 @@ class EplbModelState: class EplbState: - """ - EplbState of each expert parallel model. Key is the model config hash. - """ + """EplbState of each expert parallel model. Key is the model config hash.""" def __init__(self, parallel_config: ParallelConfig, device: torch.device): self.parallel_config = parallel_config @@ -305,14 +300,14 @@ def build_initial_global_physical_to_logical_map( num_routed_experts: int, num_redundant_experts: int, ) -> Sequence[int]: - """ - Build an initial expert arrangement using the following structure: + """Build an initial expert arrangement using the following structure: [original routed experts, redundant experts] Returns: physical_to_logical_map (Sequence[int]): A list of integers, where each integer is the index of the logical expert that the corresponding physical expert maps to. + """ global_physical_to_logical_map = list(range(num_routed_experts)) global_physical_to_logical_map += [ @@ -321,8 +316,7 @@ def build_initial_global_physical_to_logical_map( return global_physical_to_logical_map def validate_ep_configuration(self, new_model: MixtureOfExperts): - """ - Validate that the expert parallel configuration of + """Validate that the expert parallel configuration of the new model is the same as the existing models. """ if len(self.model_states) > 0: @@ -361,9 +355,7 @@ def add_model( model: MixtureOfExperts, model_config: ModelConfig, ): - """ - Build the initial EPLB state. - """ + """Build the initial EPLB state.""" self.validate_ep_configuration(model) self.is_async = self.parallel_config.eplb_config.use_async @@ -534,6 +526,7 @@ def prepare_forward( ubatch_slices: When DBO is active, a list of ``UBatchSlice`` objects describing each micro-batch's token range. When ``None``, only ``tensors[0]`` is filled. + """ model_state = self.model_states.get(compute_hash_cached(model_config)) if model_state is None or model_state.num_unpadded_tokens_tensors is None: @@ -556,8 +549,7 @@ def step( is_profile: bool = False, log_stats: bool = False, ) -> None: - """ - Step the EPLB state. + """Step the EPLB state. Args: is_dummy (bool): If `True`, this is a dummy step and the load @@ -574,6 +566,7 @@ def step( - `avg_tokens`: The average load across ranks. - `max_tokens`: The maximum load across ranks. - `balancedness`: The ratio of average load to maximum load. + """ ep_group = get_ep_group().device_group if is_profile: @@ -752,8 +745,7 @@ def rearrange( is_profile: bool = False, rank_mapping: dict[int, int] | None = None, ) -> torch.Tensor | None: - """ - Rearrange the experts according to the current load. + """Rearrange the experts according to the current load. Args: is_profile (bool): If `True`, perform a dummy rearrangement. @@ -761,8 +753,8 @@ def rearrange( no memory movement will be performed. Default is False. rank_mapping (dict[int, int] | None): The rank mapping when scaling is done in EEP. - """ + """ ep_group = get_ep_group().device_group ep_rank = ep_group.rank() @@ -1033,9 +1025,7 @@ def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool: return int(flag.item()) == device_group.size() def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]: - """ - All-reduce a list of tensors. - """ + """All-reduce a list of tensors.""" ep_group = get_ep_group().device_group if len(tensor_list) == 1: all_reduce(tensor_list[0], group=ep_group) @@ -1060,8 +1050,7 @@ def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor] return all_reduce_list def _sync_load_pass(self) -> list[torch.Tensor]: - """ - Sync the expert load pass across all ranks for log stats. + """Sync the expert load pass across all ranks for log stats. Doesn't update the expert load pass in eplb_model_state. """ load_pass_list = [] @@ -1250,8 +1239,7 @@ def compute_logical_maps( physical_to_logical_map: torch.Tensor, num_logical_experts: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Derive logical_to_physical_map and logical_replica_count from + """Derive logical_to_physical_map and logical_replica_count from physical_to_logical_map. Args: @@ -1264,6 +1252,7 @@ def compute_logical_maps( physical slots per logical expert; -1 where unused logical_replica_count: [num_layers, num_logical_experts], number of physical replicas per logical expert + """ device = physical_to_logical_map.device assert physical_to_logical_map.device.type == "cpu" @@ -1339,13 +1328,11 @@ def _commit_eplb_maps_for_layer( new_physical_to_logical_map: torch.Tensor, layer: int, ) -> None: - """ - Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB + """Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB when running async EPLB. Copies all of the new_* maps into model_state. After this function completes, the new mappings will become the current mappings and will be visible to the model. """ - # Commit physical_to_logical_map src = new_physical_to_logical_map dst = model_state.physical_to_logical_map[layer] @@ -1376,12 +1363,10 @@ def _commit_eplb_maps( model_state: EplbModelState, new_physical_to_logical_map: torch.Tensor, ) -> None: - """ - Copies all of the new_* maps into model_state. After this function completes, + """Copies all of the new_* maps into model_state. After this function completes, the new mappings will become the current mappings and will be visible to the model. """ - # Commit physical_to_logical_map src = new_physical_to_logical_map dst = model_state.physical_to_logical_map diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index dee19749745a..3d5be8a8ec4d 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -14,8 +14,7 @@ class CpuGpuEvent: - """ - Combines a CUDA event with a CPU threading event to enforce record->wait + """Combines a CUDA event with a CPU threading event to enforce record->wait ordering across two threads. This class is designed for exactly two threads: one producer that calls @@ -35,8 +34,7 @@ def __init__(self): self._recorded = threading.Event() def wait(self, stream: torch.cuda.Stream | None = None): - """ - Blocks the calling thread until record finishes. Used to guarantee that the + """Blocks the calling thread until record finishes. Used to guarantee that the record kernel is called before wait. Should only be called by the Async Eplb thread. @@ -46,8 +44,7 @@ def wait(self, stream: torch.cuda.Stream | None = None): self._recorded.clear() def record(self, stream: torch.cuda.Stream | None = None): - """ - Unblocks the waiting thread after calling event.record(). + """Unblocks the waiting thread after calling event.record(). Should only be called by the main thread. """ @@ -65,12 +62,12 @@ def override_envs_for_eplb( parallel_config: ParallelConfig, moe_backend: str | None = None, ) -> None: - """ - Override environment variables for EPLB when specific conditions are met. + """Override environment variables for EPLB when specific conditions are met. Args: parallel_config: The parallel configuration object. moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``). + """ is_data_parallel = parallel_config.data_parallel_size > 1 is_eplb_enabled = parallel_config.enable_eplb diff --git a/vllm/distributed/eplb/policy/abstract.py b/vllm/distributed/eplb/policy/abstract.py index d056468b97b2..e0c34b9caf2c 100644 --- a/vllm/distributed/eplb/policy/abstract.py +++ b/vllm/distributed/eplb/policy/abstract.py @@ -18,10 +18,9 @@ def rebalance_experts( num_ranks: int, old_global_expert_indices: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Entry point for expert-parallelism load balancer. + """Entry point for expert-parallelism load balancer. - Parameters: + Args: weight: [layers, num_logical_experts], the load statistics for all logical experts num_replicas: number of physical experts, must be a multiple of @@ -32,8 +31,10 @@ def rebalance_experts( old_global_expert_indices: [layers, num_logical_experts], the old global expert indices. Used to avoid unnecessary weight copying for experts moving within one rank. + Returns: physical_to_logical_map: [layers, num_replicas], the expert index of each replica + """ raise NotImplementedError diff --git a/vllm/distributed/eplb/policy/default.py b/vllm/distributed/eplb/policy/default.py index c2cdc42909fe..822856ff3271 100644 --- a/vllm/distributed/eplb/policy/default.py +++ b/vllm/distributed/eplb/policy/default.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Expert parallelism load balancer (EPLB) for vLLM. +"""Expert parallelism load balancer (EPLB) for vLLM. This module implements the core rearrangement algorithm. @@ -23,17 +22,17 @@ class DefaultEplbPolicy(AbstractEplbPolicy): def balanced_packing( cls, weight: np.ndarray, num_packs: int ) -> tuple[np.ndarray, np.ndarray]: - """ - Pack n weighted objects to m packs, such that each bin contains exactly + """Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs are as balanced as possible. - Parameters: + Args: weight: [X, n], the weight of each item num_packs: number of packs Returns: pack_index: [X, n], the pack index of each item rank_in_pack: [X, n], the rank of the item in the pack + """ num_layers, num_groups = weight.shape assert num_groups % num_packs == 0 @@ -76,17 +75,17 @@ def balanced_packing( def replicate_experts( cls, weight: np.ndarray, num_phy: int ) -> tuple[np.ndarray, np.ndarray]: - """ - Replicate `num_log` experts to `num_phy` replicas, such that the maximum + """Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized. - Parameters: + Args: weight: [X, num_log] num_phy: total number of experts after replication Returns: phy2log: [X, num_phy], logical expert id of each physical expert logcnt: [X, num_log], number of replicas for each logical expert + """ n, num_log = weight.shape num_redundant = num_phy - num_log @@ -109,8 +108,7 @@ def rebalance_experts_hierarchical( num_nodes: int, num_gpus: int, ) -> np.ndarray: - """ - Parameters: + """Parameters weight: [num_moe_layers, num_logical_experts] num_physical_experts: number of physical experts after replication num_groups: number of expert groups @@ -121,6 +119,7 @@ def rebalance_experts_hierarchical( Returns: phy2log: [layers, num_replicas], the expert index of each replica + """ num_layers, num_logical_experts = weight.shape assert num_logical_experts % num_groups == 0 @@ -195,8 +194,7 @@ def preserve_intragpu_slots( num_ranks: int, old_phy2log: np.ndarray, ) -> np.ndarray: - """ - Reorder the new mapping per GPU so that experts that remain on the same GPU + """Reorder the new mapping per GPU so that experts that remain on the same GPU keep their previous slot positions when possible. Incoming experts to that GPU fill any remaining available slots. This is applied only when the number of GPUs is unchanged and the slots per GPU remain the same between @@ -281,10 +279,9 @@ def rebalance_experts( num_ranks: int, old_global_expert_indices: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Entry point for expert-parallelism load balancer. + """Entry point for expert-parallelism load balancer. - Parameters: + Args: weight: [layers, num_logical_experts], the load statistics for all logical experts num_replicas: number of physical experts, must be a multiple of @@ -296,9 +293,12 @@ def rebalance_experts( old_global_expert_indices: [layers, num_logical_experts], the old global expert indices. Used to avoid unnecessary weight copying for experts moving within one rank. + Returns: + ------- phy2log: [layers, num_replicas], the expert index of each replica + """ weight_np = weight.float().cpu().numpy() old_phy2log_np = ( diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index cbd513294565..20ebe4f1f298 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -The actual execution of the rearrangement. +"""The actual execution of the rearrangement. This involves the exchange of expert weights between GPUs. """ @@ -41,9 +40,7 @@ class TransferMetadata: @dataclass class AsyncEplbLayerResult: - """ - The result of one completed async EPLB layer transfer. - """ + """The result of one completed async EPLB layer transfer.""" layer_idx: int """Index of the MoE layer that was transferred.""" @@ -69,8 +66,7 @@ def get_ep_ranks_with_experts_batch( old_indices: np.ndarray, new_indices: np.ndarray, ) -> tuple[dict[int, list[int]], dict[int, list[int]]]: - """ - Get the ranks of the experts that need to be exchanged. + """Get the ranks of the experts that need to be exchanged. Args: expert_ids: 1D array of expert indices to query. @@ -82,6 +78,7 @@ def get_ep_ranks_with_experts_batch( A tuple of two dictionaries mapping expert_id to: - ranks_to_send: The ranks that have this expert and need to send. - ranks_to_recv: The ranks that need to receive this expert. + """ ranks_to_send_map: dict[int, list[int]] = {} ranks_to_recv_map: dict[int, list[int]] = {} @@ -181,8 +178,7 @@ def move_to_buffer( communicator: EplbCommunicator, layer_idx: int = 0, ) -> TransferMetadata: - """ - Rearranges expert weights during EPLB rebalancing. + """Rearranges expert weights during EPLB rebalancing. Args: num_local_experts: Number of local experts. @@ -199,6 +195,7 @@ def move_to_buffer( Returns: TransferMetadata: Metadata needed for completing remote weight transfers. + """ assert old_indices.shape == new_indices.shape recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_) @@ -355,9 +352,8 @@ def move_from_buffer( new_indices: np.ndarray, ep_rank: int, ) -> None: - """ - Copies expert weights from communication buffers back to the target weight tensors - after EPLB rebalancing. + """Copies expert weights from communication buffers back to the target weight + tensors after EPLB rebalancing. Args: expert_weights: List of the actual MoE layer weights used in the execution. @@ -367,6 +363,7 @@ def move_from_buffer( new_indices: (num_experts_total,) mapping from local rows to desired (possibly global) expert id, after rebalance. ep_rank: Rank of the process in the expert parallel group. + """ is_unchanged = transfer_metadata.is_unchanged is_received_locally = transfer_metadata.is_received_locally @@ -437,8 +434,7 @@ def transfer_layer( rank_mapping: dict[int, int] | None = None, layer_idx: int = 0, ) -> TransferMetadata: - """ - Rearranges the expert weights in place according to the new expert indices. + """Rearranges the expert weights in place according to the new expert indices. The value of the indices arguments are logical indices of the experts, while keys are physical. @@ -462,6 +458,7 @@ def transfer_layer( Returns: TransferMetadata: Metadata needed for completing remote weight transfers, including is_unchanged and is_received_locally masks. + """ ep_size = ep_group.size() if rank_mapping is not None: @@ -519,8 +516,7 @@ def rearrange_expert_weights_inplace( is_profile: bool = False, rank_mapping: dict[int, int] | None = None, ) -> None: - """ - Rearranges the expert weights in place according to the new expert indices. + """Rearranges the expert weights in place according to the new expert indices. The value of the indices arguments are logical indices of the experts, while keys are physical. @@ -540,6 +536,7 @@ def rearrange_expert_weights_inplace( This is used during profile run, where we only perform dummy communications to reserve enough memory for the buffers. rank_mapping: A dictionary mapping old rank to new rank. + """ if rank_mapping is not None: if len(rank_mapping) == ep_group.size(): @@ -624,8 +621,7 @@ def _map_old_expert_indices_with_rank_mapping( rank_mapping: dict[int, int], new_ep_size: int, ) -> torch.Tensor: - """ - Map the old global expert indices to the new global expert indices. + """Map the old global expert indices to the new global expert indices. Args: old_global_expert_indices: @@ -636,6 +632,7 @@ def _map_old_expert_indices_with_rank_mapping( Returns: Mapped expert indices with shape (num_layers, new_ep_size * num_local_physical_experts). + """ num_layers, old_num_physical_experts = old_global_expert_indices.shape assert rank_mapping, "Rank mapping is required" diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 96ddd7ce950e..3aa2baa17165 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -39,7 +39,7 @@ class KVCacheEvent( gc=False, # type: ignore[call-arg] tag=True, ): - """Base class for all KV cache-related events""" + """Base class for all KV cache-related events.""" MEDIUM_GPU = "GPU" @@ -135,8 +135,7 @@ class KVEventBatch(EventBatch): class KVEventAggregator: - """ - Aggregates KV events across multiple workers. + """Aggregates KV events across multiple workers. Tracks how many times each event appears and returns only those that were emitted by all workers. """ @@ -150,22 +149,22 @@ def __init__(self, num_workers: int) -> None: self._num_workers: int = num_workers def add_events(self, events: list[KVCacheEvent]) -> None: - """ - Add events from a worker batch. + """Add events from a worker batch. Args: events: List of KVCacheEvent objects. + """ if not isinstance(events, list): raise TypeError("events must be a list of KVCacheEvent.") self._event_counter.update(events) def get_common_events(self) -> list[KVCacheEvent]: - """ - Return events that appeared in all workers. + """Return events that appeared in all workers. Returns: List of events present in all workers. + """ return [ event @@ -174,43 +173,39 @@ def get_common_events(self) -> list[KVCacheEvent]: ] def get_all_events(self) -> list[KVCacheEvent]: - """ - Return all events for all workers. + """Return all events for all workers. Returns: List of events for all workers. + """ return list(self._event_counter.elements()) def clear_events(self) -> None: - """ - Clear all tracked events. - """ + """Clear all tracked events.""" self._event_counter.clear() def increment_workers(self, count: int = 1) -> None: - """ - Increment the number of workers contributing events. + """Increment the number of workers contributing events. Args: count: Number to increment the workers by. + """ if count <= 0: raise ValueError("count must be positive.") self._num_workers += count def reset_workers(self) -> None: - """ - Reset the number of workers to 1. - """ + """Reset the number of workers to 1.""" self._num_workers = 1 def get_number_of_workers(self) -> int: - """ - Return the number of workers. + """Return the number of workers. Returns: int number of workers. + """ return self._num_workers @@ -222,8 +217,7 @@ def __repr__(self) -> str: class KVConnectorKVEvents(ABC): - """ - Abstract base class for KV events. + """Abstract base class for KV events. Acts as a container for KV events from the connector. """ @@ -323,6 +317,7 @@ class ZmqEventPublisher(EventPublisher): Maximum number of events to buffer in memory. topic: Topic to publish events to. + """ SHUTDOWN_TIMEOUT: float = 1.0 @@ -520,6 +515,7 @@ def offset_endpoint_port( Returns: The endpoint with the port offset by data_parallel_rank or suffix appended + """ # Do nothing if input is None or data_parallel_rank is 0 if not endpoint or data_parallel_rank == 0: diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index c08064bd8456..97064e06cdb8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -87,6 +87,7 @@ def get_connector_class_by_name( Returns: The connector class. + """ if connector_name not in cls._registry: raise ValueError(f"Connector '{connector_name}' is not registered.") diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 93552d51b521..9821f912b295 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -KV cache helper for store. -""" +"""KV cache helper for store.""" from collections.abc import Iterator from dataclasses import dataclass @@ -53,7 +51,8 @@ def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None): class KVOutputAggregator: """Utility class to aggregate the output of all workers into a single - output corresponding to Rank 0 for scheduler.""" + output corresponding to Rank 0 for scheduler. + """ def __init__(self, expected_finished_count: int): # Complete transfer tracker. Used to track finished requests @@ -233,11 +232,10 @@ def copy_kv_blocks( def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): - """ - Transforms the layout of received KV cache blocks to the local block_size. + """Transforms the layout of received KV cache blocks to the local block_size. (Only works for local blocksize > remote blocksize) - example: + Example: local blocksize = 16 tokens, remote blocksize = 4 tokens local block[0] = remote block[0, 1, 2, 3] remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|... @@ -246,6 +244,7 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): 1. view => view remote as n_blocks * remote_shape(H,remoteN,D) 2. permute => (H, nblocks, remoteN, D) 3. flatten => (H, localN, D) + """ blocks_to_update = cache.index_select(0, indices) # use physical order @@ -293,8 +292,7 @@ def kv_postprocess_layout_on_receive(cache, indices): def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio): - """ - Transforms the layout of received KV cache to the local block_size and LBHNC. + """Transforms the layout of received KV cache to the local block_size and LBHNC. (Only works for local blocksize > remote blocksize) prefill is LBHNC, smaller block_size @@ -317,9 +315,9 @@ def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_rati def yield_req_data( scheduler_output, ) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]: - """ - Yields: - (req_id, new_block_id_groups, preempted) + """Yields: + (req_id, new_block_id_groups, preempted) + """ # new requests for req_data in scheduler_output.scheduled_new_reqs: @@ -346,6 +344,7 @@ def get_current_attn_backends( Returns: Deduplicated list of attention backend classes. + """ layer_type = cast(type[Any], AttentionLayerBase) layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c97374832ffe..072018c039ed 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State +"""KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State communication in vLLM v1 The class provides the following primitives: @@ -97,8 +96,7 @@ class KVConnectorTransferResults: class SupportsHMA(ABC): - """ - The class that indicates the corresponding connector supports hybrid memory + """The class that indicates the corresponding connector supports hybrid memory allocator (HMA). This is required to use the connector together with hybrid memory allocator. """ @@ -109,8 +107,7 @@ def request_finished_all_groups( request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - """ - Called exactly once when a request has finished for all kv cache groups, + """Called exactly once when a request has finished for all kv cache groups, before its blocks are freed for each group. NOTE(Kuntai): This function is only supported by connectors that support HMA. @@ -124,6 +121,7 @@ def request_finished_all_groups( get_finished(). Optional KVTransferParams to be included in the request outputs returned by the engine. + """ raise NotImplementedError @@ -144,8 +142,7 @@ class KVConnectorRole(enum.Enum): class KVConnectorHandshakeMetadata(ABC): # noqa: B024 - """ - Metadata used for out of band connector handshake between + """Metadata used for out of band connector handshake between P/D workers. This needs to serializable. """ @@ -153,8 +150,7 @@ class KVConnectorHandshakeMetadata(ABC): # noqa: B024 class KVConnectorMetadata(ABC): # noqa: B024 - """ - Abstract Metadata used to communicate + """Abstract Metadata used to communicate Scheduler KVConnector -> Worker KVConnector. """ @@ -162,8 +158,7 @@ class KVConnectorMetadata(ABC): # noqa: B024 class KVConnectorWorkerMetadata(ABC): - """ - Abstract Metadata used to communicate back + """Abstract Metadata used to communicate back Worker KVConnector -> Scheduler KVConnector. Each worker can output its own metadata. @@ -176,16 +171,12 @@ class KVConnectorWorkerMetadata(ABC): def aggregate( self, other: "KVConnectorWorkerMetadata" ) -> "KVConnectorWorkerMetadata": - """ - Aggregate metadata with another `KVConnectorWorkerMetadata` object. - """ + """Aggregate metadata with another `KVConnectorWorkerMetadata` object.""" pass class KVConnectorBase_V1(ABC): - """ - Base class for KV connectors. - """ + """Base class for KV connectors.""" @property def supports_divergent_local_hybrid_hits(self) -> bool: @@ -245,6 +236,7 @@ def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> No Args: connector_metadata (dict): the connector metadata. + """ self._connector_metadata = connector_metadata @@ -263,6 +255,7 @@ def _get_connector_metadata(self) -> KVConnectorMetadata: Returns: ConnectorMetadata: the connector metadata. + """ # Should only be called while set to valid metadata. assert self._connector_metadata is not None @@ -273,22 +266,22 @@ def has_connector_metadata(self) -> bool: Returns: bool: True if connector metadata exists, False otherwise. + """ return self._connector_metadata is not None def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - """ - Initialize with the KV caches. Useful for pre-registering the + """Initialize with the KV caches. Useful for pre-registering the KV Caches in the KVConnector (e.g. for NIXL). Args: kv_caches: dictionary of layer names, kv cache + """ return def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): - """ - Set the xPU-specific ops for copying KV between host and device. + """Set the xPU-specific ops for copying KV between host and device. Needed when host buffer is used for kv transfer (e.g., in NixlConnector) """ return @@ -303,16 +296,14 @@ def reset_capture_state(self) -> None: return def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata): - """ - Handle preempted requests or evicted blocks BEFORE they are overwritten. + """Handle preempted requests or evicted blocks BEFORE they are overwritten. Needed for connectors which use async saves (e.g., OffloadingConnector) """ return @abstractmethod def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - """ - Start loading the KV cache from the connector to vLLM's paged + """Start loading the KV cache from the connector to vLLM's paged KV buffer. Loads required by the current forward start before it; independent asynchronous loads may start after it is submitted. @@ -329,8 +320,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non @abstractmethod def wait_for_layer_load(self, layer_name: str) -> None: - """ - Block until the KV for a specific layer is loaded into vLLM's + """Block until the KV for a specific layer is loaded into vLLM's paged buffer. This is called from within attention layer to ensure async copying from start_load_kv is complete. @@ -338,6 +328,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of that layer + """ pass @@ -349,8 +340,7 @@ def save_kv_layer( attn_metadata: "AttentionMetadata", **kwargs: Any, ) -> None: - """ - Start saving a layer of KV cache from vLLM's paged buffer + """Start saving a layer of KV cache from vLLM's paged buffer to the connector. This is called from within attention layer to enable async copying during execution. @@ -360,13 +350,13 @@ def save_kv_layer( layer in vLLM. attn_metadata (AttentionMetadata): the attention metadata. **kwargs: additional arguments for the save operation. + """ pass @abstractmethod def wait_for_save(self): - """ - Block until all the save operations is done. This is called + """Block until all the save operations is done. This is called as the forward context exits to ensure that the async saving from save_kv_layer is complete before finishing the forward. @@ -377,8 +367,7 @@ def wait_for_save(self): def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """ - Notifies worker-side connector ids of requests that have + """Notifies worker-side connector ids of requests that have finished generating tokens on the worker. The scheduler process (via the Executors) will use this output to track which workers are done. @@ -389,6 +378,7 @@ def get_finished( tuple of (sending/saving ids, recving/loading ids). The finished saves/sends req ids must belong to a set provided in a call to this method (this call or a prior one). + """ return None, None @@ -403,8 +393,7 @@ def get_transfer_results( ) def get_block_ids_with_load_errors(self) -> set[int]: - """ - Get the set of block IDs that failed to load. + """Get the set of block IDs that failed to load. Returns: Set of block IDs that encountered load errors. @@ -419,50 +408,47 @@ def get_block_ids_with_load_errors(self) -> set[int]: block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. + """ return set() def shutdown(self): - """ - Shutdown the connector. This is called when the worker process + """Shutdown the connector. This is called when the worker process is shutting down to ensure that all the async operations are completed and the connector is cleaned up properly. """ return None def get_kv_connector_stats(self) -> "KVConnectorStats | None": - """ - Get the KV connector stats collected during the last interval. - """ + """Get the KV connector stats collected during the last interval.""" return None def get_kv_connector_kv_cache_events(self) -> "KVConnectorKVEvents | None": - """ - Get the KV connector kv cache events collected during the last interval. + """Get the KV connector kv cache events collected during the last interval. This function should be called by the model runner every time after the model execution and before cleanup. """ return None def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: - """ - Get the KVConnector handshake metadata for this connector. + """Get the KVConnector handshake metadata for this connector. This metadata is used for out-of-band connector handshake between P/D workers. Returns: KVConnectorHandshakeMetadata: the handshake metadata. None if no handshake metadata is available. + """ return None def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: - """ - Build the KVConnector worker metadata for this engine step. + """Build the KVConnector worker metadata for this engine step. Returns: KVConnectorWorkerMetadata: the worker metadata. None if no worker metadata is available. + """ return None @@ -476,12 +462,12 @@ def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: self.bind_gpu_block_pool(kv_cache_manager.block_pool) def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: - """ - Bind the GPU block pool to the connector for per-GPU block status tracking. + """Bind the GPU block pool to the connector for per-GPU block status tracking. For example, inc/dec ref counts, or iterate over the prefix cache blocks. Args: gpu_block_pool: the GPU block pool. + """ return @@ -491,8 +477,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded from the + """Get number of new tokens that can be loaded from the external KV cache beyond the num_computed_tokens. Args: @@ -517,6 +502,7 @@ def get_num_new_matched_tokens( call. If the cache cannot be loaded for some tokens (e.g., due to connectivity issues or eviction), those tokens must not be taken into account. + """ pass @@ -524,8 +510,7 @@ def get_num_new_matched_tokens( 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 get_num_new_matched_tokens previously returned True for a request, this function may be called twice for that same request - @@ -543,6 +528,7 @@ def update_state_after_alloc( blocks (KVCacheBlocks): the blocks allocated for the request. num_external_tokens (int): the number of tokens to load from the external KV cache. 0 means nothing should be loaded. + """ pass @@ -550,14 +536,14 @@ def update_state_after_alloc( def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: - """ - Build the connector metadata for this step. + """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. Also, calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ pass @@ -570,12 +556,12 @@ def on_new_request(self, request: "Request") -> None: return def update_connector_output(self, connector_output: KVConnectorOutput): - """ - Update KVConnector state from worker-side connectors output. + """Update KVConnector state from worker-side connectors output. Args: connector_output (KVConnectorOutput): the worker-side connectors output. + """ return @@ -584,8 +570,7 @@ def request_finished( request: "Request", block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: - """ - Called exactly once when a request has finished, before its blocks are + """Called exactly once when a request has finished, before its blocks are freed. The connector may assumes responsibility for freeing the blocks @@ -597,6 +582,7 @@ def request_finished( get_finished(). Optional KVTransferParams to be included in the request outputs returned by the engine. + """ return False, None @@ -614,11 +600,11 @@ def register_finished_partial_tail( return False def take_events(self) -> Iterable["KVCacheEvent"]: - """ - Take the KV cache events from the connector. + """Take the KV cache events from the connector. Yields: New KV cache events since the last call. + """ return () @@ -640,16 +626,16 @@ def has_pending_push_work(self) -> bool: @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: - """ - Get the required KV cache layout for this connector. + """Get the required KV cache layout for this connector. + Args: vllm_config (VllmConfig): the vllm config. Returns: str: the required KV cache layout. e.g. HND, or NHD. None if the connector does not require a specific layout. - """ + """ if cls is KVConnectorBase_V1: raise TypeError( "get_required_kvcache_layout should not be called " @@ -659,8 +645,7 @@ def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: @classmethod def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool: - """ - Check if this connector requires PIECEWISE CUDA graph mode. + """Check if this connector requires PIECEWISE CUDA graph mode. Connectors that use asynchronous layer-by-layer operations (wait_for_layer_load/save_kv_layer) should override this method @@ -676,27 +661,26 @@ def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool: Returns: True if this connector requires PIECEWISE CUDA graph mode, False otherwise. + """ return False def get_finished_count(self) -> int | None: - """ - Get the count of requests expected to complete send/receive operations + """Get the count of requests expected to complete send/receive operations via this connector. This method is used to initialize the KVOutputAggregator, overwriting the default world_size. Returns: int: expected sending or receiving completion count. - """ + """ return None @classmethod def build_kv_connector_stats( cls, data: dict[str, Any] | None = None ) -> "KVConnectorStats | None": - """ - KVConnectorStats resolution method. This method allows dynamically + """KVConnectorStats resolution method. This method allows dynamically registered connectors to return their own KVConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -705,19 +689,18 @@ def build_kv_connector_stats( def set_xfer_handshake_metadata( self, metadata: dict[int, KVConnectorHandshakeMetadata] ) -> None: - """ - Set the KV connector handshake metadata for this connector. + """Set the KV connector handshake metadata for this connector. Args: metadata (KVConnectorHandshakeMetadata): the handshake metadata to set. + """ return None def set_xfer_handshake_metadata_pp_aware( self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] ) -> None: - """ - Set handshake metadata keyed by (pp_rank, tp_rank). + """Set handshake metadata keyed by (pp_rank, tp_rank). - Default implementation assumes pp_rank is always 0 - PP-aware connectors override this to consume all PP producer shards. """ @@ -738,19 +721,18 @@ def build_prom_metrics( labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> "KVConnectorPromMetrics | None": - """ - Create a KVConnectorPromMetrics subclass which should register + """Create a KVConnectorPromMetrics subclass which should register per-connector Prometheus metrics and implement observe() to expose connector transfer stats via Prometheus. """ return None def reset_cache(self) -> bool | None: - """ - Reset the connector's internal cache. + """Reset the connector's internal cache. Returns: bool: True if the cache was successfully reset, False otherwise. + """ logger.debug( "Connector cache reset requested, but %s does not implement reset_cache().", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py index 83b468d26551..cff2367966fc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DecodeBenchConnector: A KV Connector for decode instance performance testing. +"""DecodeBenchConnector: A KV Connector for decode instance performance testing. This connector emulates a prefill-decode disaggregated setting by filling the KV cache with dummy values, allowing measurement of decoder performance @@ -10,7 +9,7 @@ Usage: To use this connector for benchmarking, configure it in the kv_transfer_config: - Example: +Example: vllm serve --kv-transfer-config '{ "kv_connector": "DecodeBenchConnector", "kv_role": "kv_both", @@ -29,6 +28,7 @@ - fill_mean (float): Mean value for random normal fill (default: 0.015) - fill_std (float): Standard deviation for random fill (default: 0.0) Set to 0 for constant values, >0 for random sampling + """ from dataclasses import dataclass @@ -80,8 +80,7 @@ class DecodeBenchConnectorMetadata(KVConnectorMetadata): class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA): - """ - A KV Connector for decode instance performance testing. + """A KV Connector for decode instance performance testing. This connector fills the KV cache with dummy values to emulate a prefill-decode disaggregated setting, enabling performance testing of the @@ -219,8 +218,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int, bool]: - """ - For new requests, return the number of tokens that should be filled + """For new requests, return the number of tokens that should be filled with dummy KV cache values. Returns: @@ -228,6 +226,7 @@ def get_num_new_matched_tokens( - num_tokens_to_fill: number of uncomputed tokens minus 1 (we fill everything except the last token for decode) - is_async: False (synchronous filling) + """ req_id = request.request_id @@ -251,8 +250,7 @@ def get_num_new_matched_tokens( def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): - """ - Called after blocks are allocated. Store the block IDs so we can + """Called after blocks are allocated. Store the block IDs so we can fill them with dummy values. Supports both single- and multi-group KV cache configurations. @@ -334,8 +332,7 @@ def update_state_after_alloc( def build_connector_meta( self, scheduler_output: "SchedulerOutput" ) -> KVConnectorMetadata: - """ - Build metadata containing information about which blocks to fill + """Build metadata containing information about which blocks to fill with dummy KV values. """ meta = DecodeBenchConnectorMetadata(reqs_to_fill=self._pending_fills.copy()) @@ -346,9 +343,7 @@ def build_connector_meta( return meta def request_finished(self, request: "Request"): - """ - Called when a request has finished. Clean up any state. - """ + """Called when a request has finished. Clean up any state.""" self._filled_requests.discard(request.request_id) @@ -389,8 +384,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata): - """ - Fill the allocated KV cache blocks with dummy values. + """Fill the allocated KV cache blocks with dummy values. This simulates having a populated KV cache from a prefill phase, allowing decode performance testing with larger context sizes. @@ -419,13 +413,13 @@ def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata): ) def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int): - """ - Fill specified blocks with dummy values for a specific KV cache group. + """Fill specified blocks with dummy values for a specific KV cache group. Args: group_idx: The KV cache group index to fill block_ids: List of block IDs to fill in this group num_tokens: Total number of tokens to fill across these blocks + """ if not block_ids: return @@ -501,6 +495,7 @@ def _fill_block_tensor( tensor's first dim are ignored. fill_mean: Mean value for the fill. fill_std: Standard deviation for the fill. + """ # Convert block_ids to tensor on device block_ids_tensor = torch.tensor( @@ -551,6 +546,7 @@ def _fill_state_tensor( kv_cache: A state tensor to fill in its entirety. fill_mean: Mean value for the fill. fill_std: Standard deviation for the fill. + """ if fill_std > 0: kv_cache.normal_(mean=fill_mean, std=fill_std) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index 750d3b56d9de..388a8cf3337d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -118,6 +118,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non Note: The number of elements in kv_caches and layer_names should be the same. + """ def inject_kv_into_layer( @@ -134,6 +135,9 @@ def inject_kv_into_layer( src_kv_cache (torch.Tensor): the source KV cache. slot_mapping (torch.Tensor): the slot mapping. In shape [num_tokens]. + attn_metadata (AttentionMetadata): attention metadata for the + current forward pass. + """ # Deliberate bulk transfers in this reference connector. with gpu_sync_allowed(): @@ -201,6 +205,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of that layer + """ return @@ -220,6 +225,7 @@ def save_kv_layer( layer in vLLM. attn_metadata (AttentionMetadata): the attention metadata. **kwargs: additional arguments for the save operation. + """ def extract_kv_from_layer( @@ -260,8 +266,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded from the + """Get number of new tokens that can be loaded from the external KV cache beyond the num_computed_tokens. Args: @@ -272,6 +277,7 @@ def get_num_new_matched_tokens( Returns: the number of tokens that can be loaded from the external KV cache beyond what is already computed. + """ # NOTE: in this debug implementation, we assume that the prompt is # cached_prompt + newly_generated_single_token @@ -295,8 +301,7 @@ def get_num_new_matched_tokens( 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. @@ -315,6 +320,7 @@ def build_connector_meta( Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ meta = ExampleConnectorMetadata() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index d291d4754427..320d9ce8044c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -57,6 +57,7 @@ def load_hidden_states(path: str) -> dict[str, torch.Tensor]: Returns: Dict with "hidden_states" and "token_ids" tensors. + """ lock_path = path + ".lock" with open(lock_path) as lf: @@ -96,8 +97,7 @@ class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata): class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): - """ - Simple debug implementation of a HiddenStatesConnector. + """Simple debug implementation of a HiddenStatesConnector. Simply extracts the hidden states from the kv cache and stores them to disk. Must be used in conjunction with the `extract_hidden_states` spec decoding method. @@ -278,7 +278,8 @@ def wait_for_save(self) -> None: def _on_write_done(self, req_id: str, future: Future) -> None: """Surface any exception from the disk-write thread and drop the - completed future from the in-flight tracking dict.""" + completed future from the in-flight tracking dict. + """ self._req_futures.pop(req_id, None) exc = future.exception() if exc is not None: @@ -436,8 +437,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded from the + """Get number of new tokens that can be loaded from the external KV cache beyond the num_computed_tokens. Args: @@ -448,6 +448,7 @@ def get_num_new_matched_tokens( Returns: the number of tokens that can be loaded from the external KV cache beyond what is already computed. + """ # This connector is store-only, so we don't need to load any tokens return 0, False @@ -471,6 +472,7 @@ def build_connector_meta( Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ meta = ExampleHiddenStatesConnectorMetadata() @@ -509,8 +511,7 @@ def request_finished( request: "Request", block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: - """ - Called exactly once when a request has finished, before its blocks are + """Called exactly once when a request has finished, before its blocks are freed. Returns True to delay block freeing until get_finished extracts @@ -587,16 +588,16 @@ def request_finished_all_groups( @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: - """ - Get the required KV cache layout for this connector. + """Get the required KV cache layout for this connector. + Args: vllm_config (VllmConfig): the vllm config. Returns: str: the required KV cache layout. e.g. HND, or NHD. None if the connector does not require a specific layout. - """ + """ if cls is KVConnectorBase_V1: raise TypeError( "get_required_kvcache_layout should not be called " diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py index 556cba963d5b..bbb64cf24cfe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py @@ -96,6 +96,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: Args: forward_context (ForwardContext): the forward context. **kwargs (Any): additional arguments (unused). + """ self._flexkv_connector.start_load_kv(forward_context, **kwargs) @@ -107,6 +108,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of the layer (unused). + """ self._flexkv_connector.wait_for_layer_load(layer_name) @@ -131,6 +133,7 @@ def save_kv_layer( kv_layer (torch.Tensor): the paged KV buffer (unused). attn_metadata (AttentionMetadata): the attention metadata (unused). **kwargs (Any): additional arguments (unused). + """ self._flexkv_connector.save_kv_layer( layer_name, kv_layer, attn_metadata, **kwargs @@ -160,6 +163,7 @@ def get_finished( that have finished asynchronous transfer. The finished saves/sends req ids must belong to a set provided in a call to this method (this call or a prior one). + """ return self._flexkv_connector.get_finished(finished_req_ids) @@ -169,6 +173,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): Args: kv_caches: dictionary of layer names to kv cache tensors. + """ self._flexkv_connector.register_kv_caches(kv_caches) @@ -192,6 +197,7 @@ def get_num_new_matched_tokens( Tuple of (num_external_tokens, is_ready) where num_external_tokens is the number of additional tokens that can be loaded from the external KV cache. + """ return self._flexkv_connector.get_num_new_matched_tokens( request, num_computed_tokens @@ -215,6 +221,7 @@ def build_connector_meta( Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ return self._flexkv_connector.build_connector_meta(scheduler_output) @@ -224,6 +231,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput): Args: connector_output (KVConnectorOutput): the worker-side connectors output. + """ self._flexkv_connector.update_connector_output(connector_output) @@ -240,6 +248,7 @@ def request_finished( should not be freed until the request_id is returned from :meth:`get_finished`. kv_transfer_params is an optional dict of KVTransferParams to be included in the request outputs. + """ return self._flexkv_connector.request_finished(request, block_ids) @@ -248,6 +257,7 @@ def take_events(self) -> Iterable["KVCacheEvent"]: Returns: New KV cache events since the last call. + """ return self._flexkv_connector.take_events() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py index a54233453bb4..e0bf1dbfc4a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py @@ -67,6 +67,7 @@ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int): size: Total size of the storage file in bytes bytes_per_page: Size of each page in bytes entries: Maximum number of concurrent operations + """ if not HF3FS_AVAILABLE: raise ImportError( @@ -154,7 +155,7 @@ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int): ) def _release_resources(self) -> None: - """Release all acquired resources safely""" + """Release all acquired resources safely.""" # iov must be released before ioring and shm for attr in ("iov_r", "iov_w", "ior_r", "ior_w"): obj = getattr(self, attr, None) @@ -192,6 +193,7 @@ def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[in Returns: List of operation results (0 for success, non-zero for error) + """ self.check(offsets, tensors) assert self.ior_r is not None @@ -228,11 +230,12 @@ def batch_write( Args: offsets: List of byte offsets to write to tensors: List of tensors containing data to write + event: CUDA event to synchronize on before reading the tensors Returns: List of operation results (0 for success, non-zero for error) - """ + """ self.check(offsets, tensors) assert self.ior_w is not None assert self.iov_w is not None @@ -282,6 +285,7 @@ def get_size(self) -> int: Returns: Size of the file in bytes + """ return self.size diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py index 7294d77cef37..b5df658d54d8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -HF3FS KV Connector Implementation for vLLM. +"""HF3FS KV Connector Implementation for vLLM. This module implements a KV connector that uses the 3FS for storing and retrieving KV cache data. @@ -99,9 +98,7 @@ class AsyncOperationManager: - """ - Manages async save/load operations with background threads. - """ + """Manages async save/load operations with background threads.""" def __init__(self, connector: "HF3FSKVConnector"): # Store connector reference and extract commonly used attributes @@ -667,9 +664,7 @@ def get_finished( return self._async_manager.get_finished_operations(finished_req_ids) def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]: - """ - Get the KV connector stats collected during the last interval. - """ + """Get the KV connector stats collected during the last interval.""" # Clear stats for next iteration if ( hasattr(self, "_async_manager") @@ -886,8 +881,7 @@ def _process_cached_requests( def build_kv_connector_stats( cls, data: dict[str, Any] | None = None ) -> Optional["KVConnectorStats"]: - """ - KVConnectorStats resolution method. This method allows dynamically + """KVConnectorStats resolution method. This method allows dynamically registered connectors to return their own KVConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -1021,7 +1015,7 @@ def _align_to_block_size(self, num_tokens: int) -> int: @dataclass class HF3FSKVConnectorStats(KVConnectorStats): - """Container for transfer performance metrics""" + """Container for transfer performance metrics.""" def __post_init__(self): if not self.data: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py index 72792e5eb260..92cf08e5bed4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -HF3FS Metadata Server with key-based organization. -""" +"""HF3FS Metadata Server with key-based organization.""" import argparse import logging @@ -119,6 +117,7 @@ def allocate_pages_for_keys( Returns: Dictionary mapping key -> allocated page index + """ with self.global_lock: if rank not in self.rank_metadata: @@ -167,6 +166,7 @@ def confirm_write_for_keys( rank: Rank ID that confirmed the writes key_confirmations: List of (key, page_index) tuples pages_to_release: List of page indices to release back to free pool + """ with self.global_lock: # Confirm successful writes @@ -197,6 +197,7 @@ def batch_key_exists(self, keys: list[str]) -> list[bool]: Returns: List of boolean values indicating key existence and completion + """ with self.global_lock: results = [] @@ -218,6 +219,7 @@ def get_key_locations(self, rank: int, keys: list[str]) -> list[int | None]: Returns: List of page indices in the same order as input keys (None if key not found) + """ with self.global_lock: if rank not in self.rank_metadata: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py index 4ba3ecee90ee..55c3c8407da0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py @@ -128,6 +128,7 @@ def scatter_kv_caches( num_heads: Size of the H axis content_size: Size of the C axis kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view + """ num_layers = len(kv_caches_ptrs) num_tokens_in_block = len(token_indices) @@ -184,6 +185,7 @@ def gather_kv_caches( num_heads: Size of the H axis content_size: Size of the C axis kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view + """ num_layers = kv_caches_ptrs.shape[0] num_tokens_in_block = len(token_indices) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py index 609d0480e880..96191cd8f7a5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py @@ -32,9 +32,7 @@ class LMCacheKVEvents(KVConnectorKVEvents): - """ - Concrete implementation of KVConnectorKVEvents using KVEventAggregator. - """ + """Concrete implementation of KVConnectorKVEvents using KVEventAggregator.""" def __init__(self, num_workers: int) -> None: self._aggregator = KVEventAggregator(num_workers) @@ -43,9 +41,7 @@ def add_events(self, events: list[KVCacheEvent]) -> None: self._aggregator.add_events(events) def aggregate(self) -> "LMCacheKVEvents": - """ - Aggregate KV events and retain only common events. - """ + """Aggregate KV events and retain only common events.""" common_events = self._aggregator.get_common_events() self._aggregator.clear_events() self._aggregator.add_events(common_events) @@ -72,8 +68,7 @@ def __repr__(self) -> str: class LMCacheConnectorV1(KVConnectorBase_V1): @classmethod def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool: - """ - LMCache requires PIECEWISE CUDA graph mode when layerwise + """LMCache requires PIECEWISE CUDA graph mode when layerwise operations are enabled. The wait_for_layer_load and save_kv_layer methods perform actual async synchronization that cannot be captured in CUDA graphs. @@ -118,12 +113,12 @@ def __init__( # Worker-side methods # ============================== def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - """ - Initialize with the KV caches. Useful for pre-registering the + """Initialize with the KV caches. Useful for pre-registering the KV Caches in the KVConnector (e.g. for NIXL). Args: kv_caches: dictionary of layer names, kv cache + """ if hasattr(self._lmcache_engine, "register_kv_caches"): self._lmcache_engine.register_kv_caches(kv_caches) @@ -149,8 +144,7 @@ def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> No engine.layerwise_retrievers = [] def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - """ - Start loading the KV cache from the connector to vLLM's paged + """Start loading the KV cache from the connector to vLLM's paged KV buffer. This is called from the forward context before the forward pass to enable async loading during model execution. @@ -166,8 +160,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non self._lmcache_engine.start_load_kv(forward_context, **kwargs) def wait_for_layer_load(self, layer_name: str) -> None: - """ - Block until the KV for a specific layer is loaded into vLLM's + """Block until the KV for a specific layer is loaded into vLLM's paged buffer. This is called from within attention layer to ensure async copying from start_load_kv is complete. @@ -175,6 +168,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of that layer + """ self._lmcache_engine.wait_for_layer_load(layer_name) @@ -185,8 +179,7 @@ def save_kv_layer( attn_metadata: AttentionMetadata, **kwargs: Any, ) -> None: - """ - Start saving the a layer of KV cache from vLLM's paged buffer + """Start saving the a layer of KV cache from vLLM's paged buffer to the connector. This is called from within attention layer to enable async copying during execution. @@ -196,14 +189,14 @@ def save_kv_layer( layer in vLLM. attn_metadata (AttentionMetadata): the attention metadata. **kwargs: additional arguments for the save operation. + """ self._lmcache_engine.save_kv_layer( layer_name, kv_layer, attn_metadata, **kwargs ) def wait_for_save(self): - """ - Block until all the save operations is done. This is called + """Block until all the save operations is done. This is called as the forward context exits to ensure that the async saving from save_kv_layer is complete before finishing the forward. @@ -214,8 +207,7 @@ def wait_for_save(self): def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """ - Notifies worker-side connector ids of requests that have + """Notifies worker-side connector ids of requests that have finished generating tokens. Returns: @@ -224,16 +216,17 @@ def get_finished( tuple of (sending/saving ids, recving/loading ids). The finished saves/sends req ids must belong to a set provided in a call to this method (this call or a prior one). + """ return self._lmcache_engine.get_finished(finished_req_ids) def get_block_ids_with_load_errors(self) -> set[int]: - """ - Get the set of block IDs that failed to load. + """Get the set of block IDs that failed to load. Returns: Set of block IDs that encountered load errors. Empty set if no load errors occurred. + """ method = getattr(self._lmcache_engine, "get_block_ids_with_load_errors", None) if callable(method): @@ -243,10 +236,7 @@ def get_block_ids_with_load_errors(self) -> set[int]: return set() def get_kv_connector_kv_cache_events(self) -> LMCacheKVEvents | None: - """ - Get the KV connector kv cache events collected during the last interval. - """ - + """Get the KV connector kv cache events collected during the last interval.""" events = self._lmcache_engine.get_kv_events() # type: ignore [attr-defined] if not events: return None @@ -276,8 +266,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded from the + """Get number of new tokens that can be loaded from the external KV cache beyond the num_computed_tokens. Args: @@ -288,6 +277,7 @@ def get_num_new_matched_tokens( Returns: the number of tokens that can be loaded from the external KV cache beyond what is already computed. + """ return self._lmcache_engine.get_num_new_matched_tokens( request, num_computed_tokens @@ -296,32 +286,30 @@ def get_num_new_matched_tokens( 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.""" self._lmcache_engine.update_state_after_alloc(request, num_external_tokens) def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: - """ - Build the connector metadata for this step. + """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. Also, calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ return self._lmcache_engine.build_connector_meta(scheduler_output) def update_connector_output(self, connector_output: KVConnectorOutput): - """ - Update KVConnector state from worker-side connectors output. + """Update KVConnector state from worker-side connectors output. Args: connector_output (KVConnectorOutput): the worker-side connectors output. + """ # Get the KV events kv_cache_events = connector_output.kv_cache_events @@ -342,8 +330,7 @@ def request_finished( request: "Request", block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: - """ - Called when a request has finished, before its blocks are freed. + """Called when a request has finished, before its blocks are freed. Returns: True if the request is being saved/sent asynchronously and blocks @@ -351,15 +338,16 @@ def request_finished( get_finished(). Optional KVTransferParams to be included in the request outputs returned by the engine. + """ return self._lmcache_engine.request_finished(request, block_ids) def take_events(self) -> Iterable["KVCacheEvent"]: - """ - Take the KV cache events from the connector. + """Take the KV cache events from the connector. Yields: New KV cache events since the last call. + """ if self._kv_cache_events is not None: self._kv_cache_events.aggregate() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py index 2e75519df125..688e1ae18e27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py @@ -44,8 +44,7 @@ def send_lmcache_request( request_type: RequestType, payloads: list[Any], ) -> MessagingFuture[Any]: - """ - Helper function to send the request to the LMCache multiprocess server + """Helper function to send the request to the LMCache multiprocess server. Args: mq_client: The LMCache multiprocess mode message queue client @@ -54,8 +53,8 @@ def send_lmcache_request( Returns: A messaging future for the request - """ + """ future = mq_client.submit_request( request_type, payloads, get_response_class(request_type) ) @@ -65,14 +64,14 @@ def send_lmcache_request( def get_lmcache_chunk_size( mq_client: MessageQueueClient, ) -> int: - """ - Helper function to get the LMCache chunk size from the server + """Helper function to get the LMCache chunk size from the server. Args: mq_client: The LMCache multiprocess mode message queue client Returns: An integer representing the LMCache chunk size + """ future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, []) chunk_size = future.result() @@ -147,16 +146,16 @@ def __init__( vllm_block_size: int, parallel_strategy: ParallelStrategy, ): - """ - Args: - server_url: The server URL for the LMCache message queue - context: The ZMQ context - - model_name: The model name used for LMCache keys - vllm_block_size: The block size used in vLLM - parallel_strategy: - The parallel strategy, which includes `use_mla`, - `world_size`, `worker_id` and so on + """Args: + server_url: The server URL for the LMCache message queue + context: The ZMQ context + + model_name: The model name used for LMCache keys + vllm_block_size: The block size used in vLLM + parallel_strategy: + The parallel strategy, which includes `use_mla`, + `world_size`, `worker_id` and so on + """ self.mq_client = MessageQueueClient(server_url, context) @@ -195,8 +194,7 @@ def maybe_submit_lookup_request( block_hashes: list[bytes] | None = None, token_ids: list[int] | None = None, ) -> None: - """ - Submit a new lookup request to LMCache if there is no ongoing request. + """Submit a new lookup request to LMCache if there is no ongoing request. Supports both token-based and hash-based vLLM: - token_ids: token IDs (token-based vLLM) -> single token-mode key @@ -219,6 +217,7 @@ def maybe_submit_lookup_request( for later retrieve operations. In the meantime, this function will record the lookup request, and the status of the look up request can be checked by `check_lookup_result`. + """ if request_id in self.lookup_futures: # Skip if there is already a lookup request @@ -260,8 +259,7 @@ def maybe_submit_lookup_request( @_lmcache_nvtx_annotate def check_lookup_result(self, request_id: str) -> int | None: - """ - Check the result of a previously submitted lookup request. + """Check the result of a previously submitted lookup request. Args: request_id: The ID of the lookup request submitted in @@ -271,6 +269,7 @@ def check_lookup_result(self, request_id: str) -> int | None: An integer representing the total number of tokens matched in LMCache (prefix matching), or None if the lookup request is not finished yet. + """ assert request_id in self.lookup_futures, ( f"Lookup request for request_id={request_id} has not been submitted" @@ -285,25 +284,27 @@ def check_lookup_result(self, request_id: str) -> int | None: return num_chunks * self.chunk_size def num_blocks_per_chunk(self) -> int: - """ - Returns: - The number of vllm blocks in a LMCache data chunk + """Returns: + The number of vllm blocks in a LMCache data chunk + """ return self.blocks_in_chunk def cleanup_lookup_result(self, request_id: str) -> None: - """ - Clean up lookup future for a finished request to prevent memory leak. + """Clean up lookup future for a finished request to prevent memory leak. + Args: request_id: The ID of the finished request. + """ self.lookup_futures.pop(request_id, None) def end_session(self, request_id: str) -> None: - """ - Notify LMCache server to remove the session for a finished request. + """Notify LMCache server to remove the session for a finished request. + Args: request_id: The ID of the finished request. + """ send_lmcache_request( self.mq_client, @@ -319,7 +320,7 @@ def _create_key( end: int = 0, request_id: str | None = None, ) -> IPCCacheEngineKey: - """Convert token IDs to an IPC cache engine key""" + """Convert token IDs to an IPC cache engine key.""" return IPCCacheEngineKey( model_name=self.model_name, world_size=self.world_size, @@ -334,7 +335,7 @@ def _create_key( def _create_hash_key( self, chunk_hash: bytes, request_id: str | None = None ) -> IPCCacheEngineKey: - """Create a hash-mode IPC cache engine key""" + """Create a hash-mode IPC cache engine key.""" return IPCCacheEngineKey( model_name=self.model_name, world_size=self.world_size, @@ -411,12 +412,12 @@ def is_first_rank_of_pp_group(self) -> bool: ) def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - """ - Register the kv caches with LMCache server + """Register the kv caches with LMCache server. Args: kv_caches: A dict of kv caches to register. The keys are the layer names and the values are the corresponding tensors. + """ # Register kv cache and send the request self.kv_caches = kv_caches @@ -432,14 +433,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): def submit_store_request( self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): - """ - Submit a KV cache store request to LMCache + """Submit a KV cache store request to LMCache. Args: request_id: The ID of the request op: The LoadStoreOp describing the store operation. event: The CUDA event that is recorded after the current model inference step + """ if op.block_hashes is not None: # Hash mode @@ -466,14 +467,14 @@ def submit_store_request( def submit_retrieve_request( self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): - """ - Submit a KV cache retrieve request to LMCache + """Submit a KV cache retrieve request to LMCache. Args: request_id: The ID of the request op: The LoadStoreOp describing the retrieve operation. event: The CUDA event that is recorded after the current model inference step + """ if op.block_hashes is not None: # Hash mode @@ -503,8 +504,7 @@ def batched_submit_store_requests( ops: list[LoadStoreOp], event: torch.cuda.Event, ): - """ - Submit a batched store request to LMCache + """Submit a batched store request to LMCache. Args: request_ids: The IDs of the requests @@ -512,6 +512,7 @@ def batched_submit_store_requests( the same length as request_ids event: The CUDA event that is recorded after the current model inference step + """ all_keys: list[IPCCacheEngineKey] = [] block_ids: list[int] = [] @@ -552,8 +553,7 @@ def batched_submit_retrieve_requests( ops: list[LoadStoreOp], event: torch.cuda.Event, ): - """ - Submit a batched retrieve request to LMCache + """Submit a batched retrieve request to LMCache. Args: request_ids: The IDs of the requests @@ -561,6 +561,7 @@ def batched_submit_retrieve_requests( the same length as request_ids event: The CUDA event that is recorded after the current model inference step + """ all_keys: list[IPCCacheEngineKey] = [] block_ids: list[int] = [] @@ -598,8 +599,7 @@ def batched_submit_retrieve_requests( def get_finished( self, finished_req_ids_from_engine: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """ - Check and get the finished store and retrieve requests. + """Check and get the finished store and retrieve requests. Args: finished_req_ids_from_engine: the set of request ids that are @@ -617,6 +617,7 @@ def get_finished( multiple times in `finished_req_ids_from_engine`. The adapter should take care of deduplicating the request IDs and only return the request IDs that have not been returned before. + """ finished_stores = set() finished_retrieves = set() @@ -675,16 +676,14 @@ def get_finished( return ret_stores, finished_retrieves def num_blocks_per_chunk(self) -> int: - """ - Returns: - The number of vllm blocks in a LMCache data chunk + """Returns: + The number of vllm blocks in a LMCache data chunk + """ return self.blocks_in_chunk def shutdown(self): - """ - Shutdown the LMCache MP worker adapter - """ + """Shutdown the LMCache MP worker adapter.""" logger.info("Unregistering kv caches") send_lmcache_request( self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id] @@ -712,7 +711,7 @@ def _create_key( end: int = 0, request_id: str | None = None, ) -> IPCCacheEngineKey: - """Convert token IDs to an IPC cache engine key""" + """Convert token IDs to an IPC cache engine key.""" return IPCCacheEngineKey( model_name=self.model_name, world_size=self.world_size, @@ -726,7 +725,7 @@ def _create_key( def _create_hash_key( self, chunk_hash: bytes, request_id: str | None = None ) -> IPCCacheEngineKey: - """Create a hash-mode IPC cache engine key""" + """Create a hash-mode IPC cache engine key.""" return IPCCacheEngineKey( model_name=self.model_name, world_size=self.world_size, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py index 64e4893dcbcb..afd9eb862375 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py @@ -59,9 +59,7 @@ def lmcache_get_or_create_config() -> V1Config: def hex_hash_to_int16(s: str) -> int: - """ - Convert a hex hash string to a 16-bit integer. - """ + """Convert a hex hash string to a 16-bit integer.""" return int(s, 16) & 0xFFFF @@ -70,8 +68,7 @@ def apply_mm_hashes_to_token_ids( mm_hashes: list[str], mm_positions: list["PlaceholderRange"], ) -> torch.Tensor: - """ - Overwrite token_ids in-place for multimodal placeholders using + """Overwrite token_ids in-place for multimodal placeholders using efficient slice assignments. """ n = token_ids.size(0) @@ -95,8 +92,7 @@ def mla_enabled(model_config: "ModelConfig") -> bool: def extract_mm_features( request: Union["Request", "NewRequestData"], modify: bool = False ) -> tuple[list[str], list["PlaceholderRange"]]: - """ - Normalize multimodal information from a Request into parallel lists. + """Normalize multimodal information from a Request into parallel lists. This helper reads either: 1) `request.mm_features` (objects each exposing `.identifier` and @@ -119,6 +115,7 @@ def extract_mm_features( Returns: tuple[list[str], list[PlaceholderRange]]: (`mm_hashes`, `mm_positions`). May be `([], [])` when no multimodal data is present. + """ if getattr(request, "mm_features", None): mm_hashes, mm_positions = zip( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py index 6122a20f95b7..a50e603db92b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py @@ -171,6 +171,7 @@ def from_new_request( lmcache_cached_tokens (int): the number of tokens that are cached in LMCache. skip_save (bool): whether the request cache should be saved + """ # vLLM 0.9.0 update: request.block_ids changed from list[int] to # list[list[int]] @@ -219,7 +220,6 @@ def update( """Update the request tracker when a running request is scheduled again """ - self.token_ids.extend(new_token_ids) if new_block_ids is None: @@ -288,6 +288,7 @@ def from_request_tracker( Returns: the request metadata if we need to perform load/save operations, None otherwise. + """ input_token_ids = tracker.token_ids input_token_len = len(input_token_ids) @@ -445,6 +446,7 @@ def _init_lmcache_engine( Returns: The initialized LMCache engine + """ if curr_engine := LMCacheEngineBuilder.get(ENGINE_NAME): return curr_engine @@ -563,6 +565,7 @@ def add_request(self, req_meta: ReqMeta) -> None: Args: req_meta (ReqMeta): the request metadata. + """ self.requests.append(req_meta) @@ -719,6 +722,7 @@ def get_inference_info(self) -> dict: Returns: dict: Dictionary containing inference information + """ # Get vLLM config information vllm_config = self._vllm_config @@ -766,6 +770,7 @@ def get_inference_version(self) -> str: Returns: str: vLLM version string + """ return VLLM_VERSION @@ -816,6 +821,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: Note: The number of elements in kv_caches and layer_names should be the same. + """ if len(self.kv_caches) == 0: self._init_kv_caches_from_forward_context(forward_context) @@ -920,6 +926,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of that layer + """ if self.layerwise_retrievers: logger.debug("Waiting for layer %s to be loaded", self.current_layer) @@ -951,6 +958,7 @@ def save_kv_layer( kv_layer (torch.Tensor): the paged KV buffer of the current layer in vLLM. attn_metadata (AttentionMetadata): the attention metadata. + """ assert self.lmcache_engine is not None @@ -1039,7 +1047,6 @@ def save_kv_layer( @_lmcache_nvtx_annotate def wait_for_save(self): """Blocking until the KV cache is saved to the connector buffer.""" - connector_metadata = self._parent._get_connector_metadata() assert isinstance(connector_metadata, LMCacheConnectorMetadata) @@ -1150,8 +1157,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> int | None: - """ - Check for external KV cache hit. + """Check for external KV cache hit. Args: request (Request): the request object. @@ -1161,6 +1167,7 @@ def get_num_new_matched_tokens( Returns: the number of tokens that can be loaded from the external KV cache beyond what is already computed. + """ if self.kv_role == "kv_producer" and not hasattr( self.lookup_client, "supports_producer_reuse" @@ -1236,13 +1243,11 @@ def get_num_new_matched_tokens( @_lmcache_nvtx_annotate def update_state_after_alloc(self, request: "Request", num_external_tokens: int): - """ - Update KVConnector state after temporary buffer alloc. + """Update KVConnector state after temporary buffer alloc. For SharedStorageConnector, update _request_needs_load if the CacheManager this allocated blocks for us. """ - # Clear local status in lookup client when a new request is # successfully scheduled. self.lookup_client.clear_lookup_status(request.request_id) @@ -1311,8 +1316,8 @@ def build_connector_meta( Args: scheduler_output (SchedulerOutput): the scheduler output object. - """ + """ force_skip_save = self.kv_role == "kv_consumer" or self.force_skip_save meta = LMCacheConnectorMetadata() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index e28076c89b09..31b20f0ed75b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -88,9 +88,7 @@ def extract_world_size_and_kv_rank( rank: int, vllm_config: VllmConfig, ) -> tuple[int, int]: - """ - Convert the rank for the MLA. - """ + """Convert the rank for the MLA.""" use_mla = mla_enabled(vllm_config.model_config) if not use_mla: return world_size, rank @@ -174,8 +172,7 @@ def create_worker_adapter( class LMCacheMPRequestState(enum.Enum): - """ - State machine: + """State machine: PREFETCHING -- update_state_after_alloc --> WAITING_FOR_LOAD WAITING_FOR_LOAD -- process_loading_requests --> READY """ @@ -234,7 +231,8 @@ def __init__(self, request: "Request"): #### def needs_retrieve(self) -> bool: """Check whether the current request needs retrieve, will be used - update_stage_after_alloc""" + update_stage_after_alloc + """ return ( self.num_lmcache_hit_blocks > self.num_vllm_hit_blocks and self.state != LMCacheMPRequestState.READY @@ -242,7 +240,8 @@ def needs_retrieve(self) -> bool: def is_ready_for_retrieving(self) -> bool: """Check whether the current request is ready for retrieving, - will be used in process_loading_requests""" + will be used in process_loading_requests + """ return ( self.state == LMCacheMPRequestState.WAITING_FOR_LOAD and self.needs_retrieve() @@ -301,13 +300,13 @@ def GetStoreMetadata( blocks_in_chunk: int, vllm_block_size: int, ) -> "LMCacheMPRequestMetadata | None": - """ - Generate the store metadata for the current request tracker. + """Generate the store metadata for the current request tracker. Args: tracker: The request tracker to generate the metadata from. blocks_in_chunk: the number of blocks in a LMCache data chunk vllm_block_size: the block size used in vLLM + """ # Store the blocks that has block hashes # NOTE: the invariant here is that `num_stored_blocks` should @@ -375,13 +374,13 @@ def GetRetrieveMetadata( blocks_in_chunk: int, vllm_block_size: int, ) -> "LMCacheMPRequestMetadata | None": - """ - Generate the retrieve metadata for the current request tracker. + """Generate the retrieve metadata for the current request tracker. Args: tracker: The request tracker to generate the metadata from. blocks_in_chunk: the number of blocks in a LMCache data chunk vllm_block_size: the block size used in vLLM + """ if not tracker.is_ready_for_retrieving(): return None @@ -462,8 +461,7 @@ def __repr__(self): class LMCacheMPConnectorUpstream(KVConnectorBase_V1): - """ - The connector for LMCache multi-process mode. + """The connector for LMCache multi-process mode. Extra configs (kv_transfer_config.extra_config): - lmcache.mp.host: the host of the LMCache server. @@ -538,27 +536,26 @@ def _get_connector_metadata(self) -> KVConnectorMetadata: Returns: ConnectorMetadata: the connector metadata. - """ + """ # Should only be called while set to valid metadata. assert self._connector_metadata is not None return self._connector_metadata def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - """ - Initialize with the KV caches. Useful for pre-registering the + """Initialize with the KV caches. Useful for pre-registering the KV Caches in the KVConnector (e.g. for NIXL). Args: kv_caches: dictionary of layer names, kv cache + """ logger.info("Registering kv caches!") self.worker_adapter.register_kv_caches(kv_caches) return def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - """ - Start loading the KV cache from the connector to vLLM's paged + """Start loading the KV cache from the connector to vLLM's paged KV buffer. This is called from the forward context before the forward pass to enable async loading during model execution. @@ -597,8 +594,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non ) def wait_for_layer_load(self, layer_name: str) -> None: - """ - Block until the KV for a specific layer is loaded into vLLM's + """Block until the KV for a specific layer is loaded into vLLM's paged buffer. This is called from within attention layer to ensure async copying from start_load_kv is complete. @@ -606,6 +602,7 @@ def wait_for_layer_load(self, layer_name: str) -> None: Args: layer_name: the name of that layer + """ return @@ -616,8 +613,7 @@ def save_kv_layer( attn_metadata: AttentionMetadata, **kwargs: Any, ) -> None: - """ - Start saving a layer of KV cache from vLLM's paged buffer + """Start saving a layer of KV cache from vLLM's paged buffer to the connector. This is called from within attention layer to enable async copying during execution. @@ -627,12 +623,12 @@ def save_kv_layer( layer in vLLM. attn_metadata (AttentionMetadata): the attention metadata. **kwargs: additional arguments for the save operation. + """ return def wait_for_save(self): - """ - Block until all the save operations is done. This is called + """Block until all the save operations is done. This is called as the forward context exits to ensure that the async saving from save_kv_layer is complete before finishing the forward. @@ -673,8 +669,7 @@ def wait_for_save(self): def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """ - Notifies worker-side connector ids of requests that have + """Notifies worker-side connector ids of requests that have finished generating tokens on the worker. The scheduler process (via the Executors) will use this output to track which workers are done. @@ -685,14 +680,14 @@ def get_finished( tuple of (sending/saving ids, recving/loading ids). The finished saves/sends req ids must belong to a set provided in a call to this method (this call or a prior one). + """ val = self.worker_adapter.get_finished(finished_req_ids) # logger.error("Finished req ids: %s, %s", val[0], val[1]) return val def get_block_ids_with_load_errors(self) -> set[int]: - """ - Get the set of block IDs that failed to load. + """Get the set of block IDs that failed to load. Returns: Set of block IDs that encountered load errors. @@ -707,12 +702,12 @@ def get_block_ids_with_load_errors(self) -> set[int]: block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. + """ return self.worker_adapter.get_block_ids_with_load_errors() def shutdown(self): - """ - Shutdown the connector. This is called when the worker process + """Shutdown the connector. This is called when the worker process is shutting down to ensure that all the async operations are completed and the connector is cleaned up properly. """ @@ -721,9 +716,7 @@ def shutdown(self): return None def get_kv_connector_stats(self) -> "KVConnectorStats | None": - """ - Get the KV connector stats collected during the last interval. - """ + """Get the KV connector stats collected during the last interval.""" return None # ============================== @@ -735,8 +728,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded from the + """Get number of new tokens that can be loaded from the external KV cache beyond the num_computed_tokens. Args: @@ -761,6 +753,7 @@ def get_num_new_matched_tokens( call. If the cache cannot be loaded for some tokens (e.g., due to connectivity issues or eviction), those tokens must not be taken into account. + """ tracker = self._get_or_create_request_tracker(request) # TODO: support loading KV for preempted requests in the future @@ -803,8 +796,7 @@ def get_num_new_matched_tokens( 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 get_num_new_matched_tokens previously returned True for a request, this function may be called twice for that same request - @@ -817,6 +809,7 @@ def update_state_after_alloc( blocks (KVCacheBlocks): the blocks allocated for the request. num_external_tokens (int): the number of tokens that will be loaded from the external KV cache. + """ # NOTE: `blocks` comes from kv_cache_manager.get_blocks(request_id), # which returns ALL blocks for the request (not just newly allocated). @@ -883,14 +876,14 @@ def update_state_after_alloc( def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: - """ - Build the connector metadata for this step. + """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. Also, calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. + """ metadata = LMCacheMPConnectorMetadata() @@ -907,12 +900,12 @@ def build_connector_meta( return metadata def update_connector_output(self, connector_output: KVConnectorOutput): - """ - Update KVConnector state from worker-side connectors output. + """Update KVConnector state from worker-side connectors output. Args: connector_output (KVConnectorOutput): the worker-side connectors output. + """ return @@ -921,8 +914,7 @@ def request_finished( request: "Request", block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: - """ - Called exactly once when a request has finished, before its blocks are + """Called exactly once when a request has finished, before its blocks are freed. The connector may assumes responsibility for freeing the blocks @@ -934,8 +926,8 @@ def request_finished( get_finished(). Optional KVTransferParams to be included in the request outputs returned by the engine. - """ + """ params: dict[str, Any] | None = getattr(request, "kv_transfer_params", None) return_params: dict[str, Any] | None = {} if params is not None else None @@ -962,26 +954,26 @@ def request_finished( return True, return_params def take_events(self) -> Iterable["KVCacheEvent"]: - """ - Take the KV cache events from the connector. + """Take the KV cache events from the connector. Yields: New KV cache events since the last call. + """ return () @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: - """ - Get the required KV cache layout for this connector. + """Get the required KV cache layout for this connector. + Args: vllm_config (VllmConfig): the vllm config. Returns: str: the required KV cache layout. e.g. HND, or NHD. None if the connector does not require a specific layout. - """ + """ if cls is KVConnectorBase_V1: raise TypeError( "get_required_kvcache_layout should not be called " @@ -990,13 +982,13 @@ def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: return None def get_finished_count(self) -> int | None: - """ - Get the count of requests expected to complete send/receive operations + """Get the count of requests expected to complete send/receive operations via this connector. This method is used to initialize the KVOutputAggregator, overwriting the default world_size. Returns: int: expected sending or receiving completion count. + """ return None @@ -1004,8 +996,7 @@ def get_finished_count(self) -> int | None: def build_kv_connector_stats( cls, data: dict[str, Any] | None = None ) -> "KVConnectorStats | None": - """ - KVConnectorStats resolution method. This method allows dynamically + """KVConnectorStats resolution method. This method allows dynamically registered connectors to return their own KVConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -1019,8 +1010,7 @@ def build_prom_metrics( labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> "KVConnectorPromMetrics | None": - """ - Create a KVConnectorPromMetrics subclass which should register + """Create a KVConnectorPromMetrics subclass which should register per-connector Prometheus metrics and implement observe() to expose connector transfer stats via Prometheus. """ @@ -1183,8 +1173,7 @@ def _get_or_create_request_tracker( return self.request_trackers[request_id] def _cleanup_request_tracker(self, request_id: str) -> None: - """ - Clean up request tracker and associated lookup future for a request. + """Clean up request tracker and associated lookup future for a request. This should be called when a request is finished to prevent memory leak. """ # Clean up request tracker diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py index e102272bc78c..770cba8ada9d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py @@ -17,8 +17,7 @@ @dataclass class KVConnectorStats: - """ - Base class for KV Connector Stats, a container for transfer performance + """Base class for KV Connector Stats, a container for transfer performance metrics or otherwise important telemetry from the connector. All sub-classes need to be serializable as stats are sent from worker to logger process. @@ -35,14 +34,11 @@ def reset(self): raise NotImplementedError def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats": - """ - Aggregate stats with another `KVConnectorStats` object. - """ + """Aggregate stats with another `KVConnectorStats` object.""" raise NotImplementedError def reduce(self) -> dict[str, int | float]: - """ - Reduce the observations collected during a time interval to one or + """Reduce the observations collected during a time interval to one or more representative values (eg avg/median/sum of the series). This is meant to be called by the logger to produce a summary of the stats for the last time interval. @@ -95,7 +91,7 @@ def observe(self, transfer_stats_data: dict[str, Any]): ) def log(self, log_fn=logger.info): - """Log transfer metrics periodically, similar to throughput logging""" + """Log transfer metrics periodically, similar to throughput logging.""" if ( self.transfer_stats_accumulator and not self.transfer_stats_accumulator.is_empty() @@ -111,8 +107,7 @@ def log(self, log_fn=logger.info): class KVConnectorPromMetrics: - """ - A base class for per-connector Prometheus metric registration + """A base class for per-connector Prometheus metric registration and recording. """ @@ -131,8 +126,7 @@ def __init__( self.per_engine_labelvalues = per_engine_labelvalues def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): - """ - Record the supplied transfer statistics to Prometheus metrics. These + """Record the supplied transfer statistics to Prometheus metrics. These statistics are engine-specific, and should be recorded to a metric with the appropriate 'engine' label. These metric instances can be created using the create_metric_per_engine() helper method. @@ -141,8 +135,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): class KVConnectorProm: - """ - Support for registering per-connector Prometheus metrics, and + """Support for registering per-connector Prometheus metrics, and recording transfer statistics to those metrics. Uses KVConnectorBase.build_prom_metrics(). """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 5a3f462577c4..4e6844034c7c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -643,7 +643,7 @@ def build_kv_connector_stats( class MooncakeConnectorScheduler: - """Implementation of Scheduler side methods""" + """Implementation of Scheduler side methods.""" def __init__( self, @@ -718,7 +718,8 @@ def get_sw_clipped_blocks( def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" + always recomputes the last token and must start from h(N-1). + """ if self._has_mamba and num_prompt_tokens > 1: return num_prompt_tokens - 1 return num_prompt_tokens @@ -729,7 +730,8 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: derive h(N) correctly. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled.""" + request is preempted and rescheduled. + """ params = request.kv_transfer_params if ( params is not None @@ -756,8 +758,7 @@ def on_new_request(self, request: "Request") -> None: def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int, bool]: - """ - For remote prefill, pull all prompt blocks from remote + """For remote prefill, pull all prompt blocks from remote asynchronously relative to engine execution. Args: @@ -769,8 +770,8 @@ def get_num_new_matched_tokens( external KV cache beyond what is already computed. * true if the external KV cache tokens will be loaded asynchronously (between scheduler steps). - """ + """ params = request.kv_transfer_params logger.debug( "MooncakeConnector get_num_new_matched_tokens: " @@ -881,11 +882,9 @@ def request_finished( request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - """ - Once a request is finished, determine whether request blocks + """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. """ - params = request.kv_transfer_params logger.debug( "MooncakeConnector request_finished, req_id=%s, request_status=%s, " @@ -934,7 +933,7 @@ def request_finished( class MooncakeConnectorWorker: - """Implementation of Worker side methods""" + """Implementation of Worker side methods.""" def __init__( self, @@ -1176,11 +1175,9 @@ async def register_worker_with_bootstrap(self): raise e async def _mooncake_sender_listener(self, ready_event: threading.Event): - """ - Background thread that listens for Mooncake requests, dispatches them + """Background thread that listens for Mooncake requests, dispatches them to a thread pool, and sends acknowledgments upon completion. """ - sock = self.async_zmq_ctx.socket(zmq.ROUTER) self.side_channel_port = sock.bind_to_random_port(f"tcp://{self.hostname}") logger.debug( @@ -1680,7 +1677,8 @@ async def _build_transfer_params( def _bind_sender_thread_device(self) -> None: """ThreadPoolExecutor initializer — binds each pool thread to the correct CUDA device. CUDA device selection is thread-local, so - without this, NVLink transfers fail for TP ranks > 0.""" + without this, NVLink transfers fail for TP ranks > 0. + """ current_platform.set_device(self.device_id) def _send_blocks( @@ -1716,7 +1714,6 @@ def _send_blocks( def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in mooncake.""" - logger.info("Registering KV_Caches. use_mla: %s", self.use_mla) kv_data_ptrs: list[int] = [] @@ -1853,8 +1850,7 @@ async def fetch_finished_sending_reqs(self) -> set[ReqId]: return finished_sending_reqs def get_finished(self) -> tuple[set[str] | None, set[str] | None]: - """ - Get requests that are done sending or recving on this specific worker. + """Get requests that are done sending or recving on this specific worker. The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ @@ -1885,8 +1881,7 @@ def get_finished(self) -> tuple[set[str] | None, set[str] | None]: return finished_sending_reqs or None, finished_recving_reqs or None def get_transfer_results(self) -> KVConnectorTransferResults: - """ - Get transfers that completed on this specific worker, including + """Get transfers that completed on this specific worker, including requests whose remote KV load failed. The scheduler process (via the MultiprocExecutor) will use this output @@ -1909,7 +1904,8 @@ def get_transfer_results(self) -> KVConnectorTransferResults: def get_kv_connector_stats(self) -> KVConnectorStats | None: """Return transfer stats collected since the last call, or None - if nothing has been recorded in this interval.""" + if nothing has been recorded in this interval. + """ if self.xfer_stats.is_empty(): return None return self.xfer_stats.clone_and_reset() @@ -2321,8 +2317,7 @@ def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool: def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]: - """ - Returns the address of the Mooncake bootstrap server. + """Returns the address of the Mooncake bootstrap server. This is only used by prefillers to register workers. Decoders should get addr from kv_transfer_params. """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py index 2d158387f060..df47f9fdb737 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py @@ -32,8 +32,7 @@ class EngineEntry: class MooncakeBootstrapServer: - """ - A centralized server running on the global rank 0 prefiller worker. + """A centralized server running on the global rank 0 prefiller worker. Prefiller workers register their connection info (IP, port, ranks) here. """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index d05b553b64b9..9e75c880d5e6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -73,7 +73,8 @@ def contains(self, group_id: int, block_hash: BlockHash) -> bool: class MooncakeStoreCoordinator: """Mirror of ``HybridKVCacheCoordinator.find_longest_cache_hit`` over an - ``ExternalCachedBlockPool``.""" + ``ExternalCachedBlockPool``. + """ def __init__( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index b9df9648d58e..2510a7ee8e76 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -630,6 +630,7 @@ def process_tokens( ``mask_num``. False entries are skipped before hash access. put_step: Stride for distributing chunks across ranks. put_step_rank: ``chunk_id % put_step`` value this rank stores. + """ assert put_step > 0 if not block_hashes: @@ -662,6 +663,7 @@ class TailKeyBoundary: ``block_hashes[num_tokens // hash_block_size - 1]`` instead of the hash implied by ``MooncakeLookupResult.hit_length``. This changes only the load key, not the reusable prefix. + """ group_id: int @@ -678,6 +680,7 @@ class MooncakeLookupResult: tail_key_boundaries: Hash boundary used to store each cache group's tail block when ``hit_length`` does not identify its store key. There is one entry per group for every nonzero hit. + """ hit_length: int diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 0231ba73948e..fca59a12e0ba 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -815,6 +815,7 @@ def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool: Returns: True when no put is needed or every put succeeds, False otherwise. + """ offloads = req_meta.boundary_state_offloads if not offloads or not req_meta.block_hashes: @@ -2628,7 +2629,8 @@ def lookup( non_block: bool = False, ) -> MooncakeLookupResult | None: """If non_block is True, will return None until the result is ready, - so the caller retries on a later step.""" + so the caller retries on a later step. + """ future = self.futures.get(req_id) if future is None: future = self.executor.submit(self._lookup, num_tokens, list(block_hashes)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 7d06b925ba29..f25e53e3f0f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -588,8 +588,7 @@ def add_new_req( @contextlib.contextmanager def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: - """Context manager for a ZMQ socket""" - + """Context manager for a ZMQ socket.""" if socket_type not in (zmq.ROUTER, zmq.REQ, zmq.DEALER): raise ValueError(f"Unexpected socket type: {socket_type}") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index ecdd59fca3ef..ea6c363df335 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -373,6 +373,7 @@ def has_connector_metadata(self) -> bool: Returns: bool: True if connector metadata exists, False otherwise. + """ try: return self._connector_metadata is not None @@ -381,7 +382,7 @@ def has_connector_metadata(self) -> bool: class MoRIIOConnectorScheduler: - """Implementation of Scheduler side methods""" + """Implementation of Scheduler side methods.""" def __init__( self, @@ -547,7 +548,8 @@ def unmap_request_id( self, request_id: ReqId, transfer_id: TransferId | None = None ): """Unmap request_id/transfer_id. Uses transfer_id for lookup if - exact request_id match fails (handles input_processor mutation).""" + exact request_id match fails (handles input_processor mutation). + """ if request_id in self.request_id_to_transfer_id: tid = self.request_id_to_transfer_id[request_id] del self.request_id_to_transfer_id[request_id] @@ -579,8 +581,7 @@ def get_num_new_matched_tokens( request: "Request", num_computed_tokens: int, ) -> tuple[int, bool]: - """ - For remote prefill, pull all prompt blocks from remote + """For remote prefill, pull all prompt blocks from remote asynchronously relative to engine execution. Args: @@ -592,6 +593,7 @@ def get_num_new_matched_tokens( external KV cache beyond what is already computed. * true if the external KV cache tokens will be loaded asynchronously (between scheduler steps). + """ if self.is_producer: return 0, False @@ -1007,11 +1009,9 @@ def request_finished( request: "Request", block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: - """ - Once a request is finished, determine whether request blocks + """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. """ - request_id = request.request_id params = request.kv_transfer_params # Consumer: can unmap transfer_id<->request_id immediately since done_recving @@ -1177,7 +1177,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None: class MoRIIOConnectorWorker: - """Implementation of Worker side methods""" + """Implementation of Worker side methods.""" def __init__( self, @@ -1435,8 +1435,8 @@ def schedule_write_blocks( kv_layer: KV cache tensor remote_notify_port: Port for completion notification remote_ip: IP address of remote node - """ + """ # synchronization to prevent dirty reads between # transfer and attention operations # we can consider removing this synchronization after ibgda is enabled. @@ -1581,7 +1581,6 @@ def _moriio_handshake_listener( layer_name_to_local_kv_cache_metadata: dict, ): """Background thread for getting new MoRIIO handshakes.""" - encoder = msgspec.msgpack.Encoder() encoded_data = encoder.encode(metadata) size_in_bytes = len(encoded_data) @@ -1637,7 +1636,6 @@ def _moriio_handshake( local-rank mapping _remote_tp_rank -- byte-identical for callers not yet TP-aware. """ - start_time = time.perf_counter() # NOTE(rob): we need each rank to have a unique port. This is @@ -1856,7 +1854,6 @@ def _remote_layer_mr_offset( def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in moriio.""" - self.kv_caches = kv_caches # layer name to kv cache self.kv_cache_shapes = { layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items() @@ -2020,12 +2017,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.moriio_wrapper.async_wait_reqid() def get_finished(self) -> tuple[set[str], set[str]]: - """ - Get requests that are done sending or recving on this specific worker. + """Get requests that are done sending or recving on this specific worker. The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ - done_sending, done_recving = set(), set() if self.is_producer: @@ -2416,8 +2411,7 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No self._eager_handshaked_engines.add(remote_engine_id) def start_load_kv(self, metadata: MoRIIOConnectorMetadata): - """ - Start loading by triggering non-blocking moriio_xfer. + """Start loading by triggering non-blocking moriio_xfer. We check for these trnxs to complete in each step(). """ self.transfer_id_to_request_id = metadata.transfer_id_to_request_id @@ -2686,8 +2680,13 @@ def _compute_block_transfer_offsets( local_block_ids: IDs of local blocks remote_block_ids: IDs of remote blocks remote_moriio_meta: Metadata of the remote MoRIIO agent + remote_tp_size: Tensor parallel size of the remote engine. Falls + back to the locally recorded value when None. + remote_engine_id: Id of the remote engine the blocks live on. + Returns: Tuple of (local_offsets, remote_offsets, transfer_sizes) + """ validate_moriio_heterogeneous_tp_kv_heads( local_tp_size=self.world_size, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index fd8257a19c7b..5c1330521dd6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -87,6 +87,7 @@ def __init__(self, worker: "MoRIIOConnectorWorker"): Args: worker: Reference to the parent worker + """ self._worker_ref: weakref_ref[MoRIIOConnectorWorker] = weakref_ref(worker) self._write_task_q: Queue[WriteTask] = Queue() @@ -108,6 +109,7 @@ def worker(self) -> "MoRIIOConnectorWorker": Raises: RuntimeError: If worker has been garbage collected + """ worker = self._worker_ref() if worker is None: @@ -131,6 +133,7 @@ def schedule_write(self, task: WriteTask) -> bool: Args: task: The write task to schedule + """ self.ensure_worker_started() if self._is_transfer_terminal(task.transfer_id): @@ -178,7 +181,6 @@ def seal_pending_transfers(self) -> None: def _write_worker_loop(self) -> None: """Main loop for the write worker thread.""" - while True: # Process deferred tasks first self._process_deferred_tasks() @@ -275,6 +277,7 @@ def _is_remote_ready(self, task: WriteTask) -> bool: Returns: True if remote blocks are ready + """ return ( task.transfer_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict @@ -291,6 +294,7 @@ def _get_remote_alloc_info(self, transfer_id: str) -> RemoteAllocInfo: Raises: KeyError: If allocation info is missing + """ try: return self.worker.moriio_wrapper.done_remote_allocate_req_dict[transfer_id] @@ -372,9 +376,11 @@ def _prepare_transfer_plan( Args: task: The write task request_info: Remote allocation information + remote_moriio_meta: Metadata of the remote MoRIIO agent Returns: The transfer plan + """ layer_cache = self.worker.kv_caches[task.layer_name] key = (task.layer_name, *_get_write_geometry_key(layer_cache)) @@ -412,6 +418,7 @@ def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]: Args: plan: The transfer plan sessions: List of transfer sessions + """ if plan.use_batch: return [ @@ -510,6 +517,7 @@ class MoRIIOWrapper: moriio_engine: MoRIIO engine instance tp_rank: Tensor parallel rank dp_rank: Data parallel rank + """ def __init__( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index d291d3852804..a6f0d4cafb75 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -71,8 +71,7 @@ def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetada @dataclass class MultiKVConnectorStats(KVConnectorStats): - """ - Maintain a dict of KVConnectorStats objects, one for each connector. + """Maintain a dict of KVConnectorStats objects, one for each connector. This is used to aggregate the stats from all connectors separately. """ @@ -132,8 +131,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): class MultiConnector(KVConnectorBase_V1, SupportsHMA): - """ - A wrapper for using multiple KVConnectors at the same time. + """A wrapper for using multiple KVConnectors at the same time. The current logic is: - Load KV from the first connector that advertises available tokens from @@ -143,8 +141,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): @classmethod def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool: - """ - MultiConnector requires PIECEWISE CUDA graph mode if any of its + """MultiConnector requires PIECEWISE CUDA graph mode if any of its child connectors require it. """ connectors_config = extra_config.get("connectors", []) @@ -480,8 +477,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput): connector_output.kv_connector_worker_meta = multi_connector_worker_meta def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: - """ - Get the KVConnector handshake metadata from sub-connectors. + """Get the KVConnector handshake metadata from sub-connectors. Returns the first non-None metadata from sub-connectors. """ for c in self._connectors: @@ -493,8 +489,7 @@ def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: def set_xfer_handshake_metadata( self, metadata: dict[int, KVConnectorHandshakeMetadata] ) -> None: - """ - Set the KV connector handshake metadata for all sub-connectors. + """Set the KV connector handshake metadata for all sub-connectors. This is needed to start the NIXL listener thread for NixlConnector. """ for c in self._connectors: @@ -590,14 +585,15 @@ def has_pending_push_work(self) -> bool: @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: - """ - Get the required KV cache layout for this connector. + """Get the required KV cache layout for this connector. + Args: vllm_config (VllmConfig): the vllm config. Returns: str: the required KV cache layout. e.g. HND, or NHD. None if the connector does not require a specific layout. + """ assert vllm_config.kv_transfer_config is not None layouts: set[str] = set() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 54184aa33945..ed9123fc6a98 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -299,11 +299,11 @@ def get_exchange_clipped_blocks( def set_xfer_handshake_metadata( self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] ) -> None: - """ - Set the KV connector handshake metadata for this connector. + """Set the KV connector handshake metadata for this connector. Args: metadata (dict): the handshake metadata to set. + """ encoded_data: dict[tuple[int, int], bytes] = {} encoder = msgspec.msgpack.Encoder() @@ -400,7 +400,8 @@ def _prefill_backoff(self) -> int: def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. The number of prompt tokens to load from the prefiller. Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder - will recompute locally.""" + will recompute locally. + """ backoff = self._prefill_backoff() if backoff and num_prompt_tokens > backoff: return num_prompt_tokens - backoff @@ -413,7 +414,8 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: multi-module MTP it is the drafter's whole lookahead window. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled.""" + request is preempted and rescheduled. + """ backoff = self._prefill_backoff() params = request.kv_transfer_params if ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index db392078ca0c..b261a13883c9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1147,8 +1147,7 @@ def _add_notif_only_remote_agent( return self.nixl_wrapper.add_remote_agent(metadata.agent_metadata) def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None: - """ - Initialize transfer buffer in CPU mem for accelerators + """Initialize transfer buffer in CPU mem for accelerators NOT directly supported by NIXL (e.g., tpu) """ xfer_buffers: dict[str, torch.Tensor] = {} @@ -1262,8 +1261,7 @@ def _ensure_handshake( pp_size: int = 1, notif_agents_only: bool = False, ) -> Future[tuple[dict[tuple[int, int], str], float]] | None: - """ - Ensure a handshake is in-flight (or already done) for *engine_id*. + """Ensure a handshake is in-flight (or already done) for *engine_id*. Returns the ``Future`` if a handshake is pending (or was just started), or ``None`` if the handshake already completed @@ -1348,7 +1346,6 @@ def request_ready(f: Future[Any], entry=(req_id, meta)): def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" - self.transfer_topo = TransferTopology( tp_rank=self.transfer_tp_rank, tp_size=self.transfer_tp_size, @@ -1876,7 +1873,8 @@ def _build_mamba_remote( ) -> np.ndarray: """Build remote desc regions (conv sub-projections + ssm) per layer. For hetero-TP, each D rank reads only its sub-projection slice from - the P rank. Returns an Nx3 uint64 array.""" + the P rank. Returns an Nx3 uint64 array. + """ assert nixl_agent_meta.kv_caches_base_addr, ( "Remote KV cache base addresses must not be empty." ) @@ -2034,8 +2032,7 @@ def register_local_xfer_handler( self, block_size: int, ) -> tuple[int, np.ndarray]: - """ - Function used for register local xfer handler with local block_size or + """Function used for register local xfer handler with local block_size or Remote block_size. When local block_size is same as remote block_size, we use local block_size @@ -2124,8 +2121,7 @@ def add_remote_agent( remote_tp_size: int = 1, remote_dcp_size: int = 1, ) -> str: - """ - Add the remote NIXL agent and prepare the descriptors for reading cache + """Add the remote NIXL agent and prepare the descriptors for reading cache blocks from remote. In particular, handle both homogeneous and heterogeneous TP. The former @@ -2395,8 +2391,7 @@ def _validate_remote_agent_handshake( remote_tp_size: int, remote_dcp_size: int = 1, ): - """ - Validate the remote agent handshake metadata ensuring the + """Validate the remote agent handshake metadata ensuring the invariants hold true. """ remote_engine_id = nixl_agent_meta.engine_id @@ -2607,7 +2602,7 @@ def _validate_remote_agent_handshake( assert len(nixl_agent_meta.region_mem_types) == num_remote_regions def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): - """copy recved kv from host buffer to device.""" + """Copy recved kv from host buffer to device.""" assert self.use_host_buffer assert self.copy_blocks is not None @@ -2632,7 +2627,7 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): ) def save_kv_to_host(self, metadata: NixlConnectorMetadata): - """copy kv from device to host buffer.""" + """Copy kv from device to host buffer.""" assert self.use_host_buffer assert self.copy_blocks is not None @@ -2662,7 +2657,8 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): @cached_property def _attention_kv_caches(self) -> list[torch.Tensor]: """Device KV caches of attention layers (mamba states excluded), - as consumed by the receive post-process.""" + as consumed by the receive post-process. + """ assert self.device_kv_caches, ( "_attention_kv_caches accessed before register_kv_caches" ) @@ -2681,8 +2677,7 @@ def post_process_device_kv_on_receive( block_ids_list: list[tuple[list[int], int]], convert: bool = True, ): - """ - Post process device kv cache after receiving from remote. + """Post process device kv cache after receiving from remote. 3 types of conversion supported (``convert``): * kv_cache_postprocess_layout => convert from HND to NHD @@ -2771,8 +2766,7 @@ def post_process_device_kv_on_receive( def post_process_device_kv_on_receive_heterogeneous_attn( self, block_ids: list[int] ): - """ - Post process device kv cache after receiving from remote + """Post process device kv cache after receiving from remote for heterogeneous attention. """ assert self.enable_heterogeneous_attn_post_process @@ -2805,8 +2799,7 @@ def _zero_region_blocks(self, block_ids: BlockIds) -> None: pages[block].zero_() def get_transfer_results(self) -> KVConnectorTransferResults: - """ - Get transfers that completed on this specific worker. + """Get transfers that completed on this specific worker. The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. @@ -2984,6 +2977,7 @@ def _handle_heartbeat(self, payload: str) -> None: Args: payload: comma-separated P-side request IDs, e.g. "req_abc,req_def". + """ new_expiry = time.perf_counter() + self._lease_extension for req_id in payload.split(","): @@ -3010,6 +3004,7 @@ def _pop_done_transfers( Returns: Requests with no outstanding handles and requests with observed failures. Failed requests may still have outstanding handles. + """ done_req_ids: set[str] = set() failed_req_ids: set[str] = set() @@ -3091,9 +3086,7 @@ def _send_pending_recv_notifs(self, req_id: str) -> None: self.xfer_stats.record_failed_notification() def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None: - """ - Send heartbeat notifications to remote engines, extending lease on KV blocks. - """ + """Send heartbeats to remote engines, extending the lease on KV blocks.""" for engine_id, hb_info in metadata.heartbeat_by_engine.items(): # Proactive handshake (this request may still be in waiting queue) so # the **next** heartbeat for this remote can go through. @@ -3126,8 +3119,7 @@ def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None: def get_mapped_blocks( self, block_ids: np.ndarray, block_size_ratio: int ) -> np.ndarray: - """ - Calculates the new set of block IDs by mapping every element + """Calculates the new set of block IDs by mapping every element in the (potentially sparse) input array. Example: block_ids=[0, 2], block_size_ratio=2 get_mapped_blocks 0 1 [2 3] 4 5 @@ -3184,8 +3176,7 @@ def _map_block_ids_for_block_size_ratio( return mapped_local, mapped_remote def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: - """ - Convert block ids to kernel physical block ids. + """Convert block ids to kernel physical block ids. This is required when the logical block size (the one set by the user) does not match the one required by the attn backend. `ratio` is the number of physical blocks per logical block. @@ -3239,6 +3230,7 @@ def _apply_dcp_prefix_caching( An empty result (rank 0 above) takes the notification-only path, allowing the remote to release its blocks without performing a transfer. + """ local_size, remote_size = local_dcp_size, remote_dcp_size @@ -3417,17 +3409,14 @@ def _apply_prefix_caching( return decode_block_ids, prefill_block_ids def get_kv_connector_stats(self) -> KVConnectorStats | None: - """ - Get the KV transfer stats for the connector. - """ + """Get the KV transfer stats for the connector.""" # Clear stats for next iteration if not self.xfer_stats.is_empty(): return self.xfer_stats.clone_and_reset() return None def get_block_ids_with_load_errors(self) -> set[int]: - """ - Return and clear the set of block IDs that failed to load. + """Return and clear the set of block IDs that failed to load. This is called by the scheduler to identify blocks that need to be retried after a NIXL transfer failure. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 99c84765e9db..ac929f7fcd94 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -214,12 +214,12 @@ def request_finished_all_groups( def set_xfer_handshake_metadata_pp_aware( self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] ) -> None: - """ - Set handshake metadata keyed by (pp_rank, tp_rank) so the side + """Set handshake metadata keyed by (pp_rank, tp_rank) so the side channel can serve every PP stage's agent metadata. Args: metadata (dict): the handshake metadata to set. + """ assert self.connector_scheduler is not None self.connector_scheduler.set_xfer_handshake_metadata(metadata) @@ -317,14 +317,14 @@ def shutdown(self): self.connector_scheduler.shutdown() def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: - """ - Get the KVConnector handshake metadata for this connector. + """Get the KVConnector handshake metadata for this connector. This metadata is used for out-of-band connector handshake between P/D workers. Returns: KVConnectorHandshakeMetadata: the handshake metadata. None if no handshake metadata is available. + """ assert self.connector_worker is not None if ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 1ae0ed72cecf..ca6f30d0cfed 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -78,8 +78,7 @@ class NixlAgentMetadata: @dataclass class NixlHandshakePayload(KVConnectorHandshakeMetadata): - """ - Wrapper for NIXL handshake sent over the wire. + """Wrapper for NIXL handshake sent over the wire. Enables two-phase decoding for graceful compatibility checking: 1. Decode NixlHandshakePayload to get compatibility_hash @@ -142,8 +141,7 @@ def compute_nixl_compatibility_hash( attn_backend_name: str, transfer_mode: str = "pull", ) -> str: - """ - Compute compatibility hash for NIXL KV transfer. + """Compute compatibility hash for NIXL KV transfer. Hash only the factors that affect whether two NIXL instances can successfully transfer KV cache data. @@ -169,6 +167,7 @@ def compute_nixl_compatibility_hash( Returns: SHA-256 hex digest + """ from vllm import __version__ as vllm_version from vllm.config.utils import hash_factors diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 9b9012a532ea..f283a40d3c2b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -34,8 +34,7 @@ def __init__( def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int, bool]: - """ - For remote prefill, pull all prompt blocks from remote + """For remote prefill, pull all prompt blocks from remote asynchronously relative to engine execution. Args: @@ -47,8 +46,8 @@ def get_num_new_matched_tokens( external KV cache beyond what is already computed. * true if the external KV cache tokens will be loaded asynchronously (between scheduler steps). - """ + """ params = request.kv_transfer_params logger.debug( "NIXLConnector get_num_new_matched_tokens: " @@ -194,8 +193,7 @@ def request_finished( request: "Request", block_ids: "BlockIds", ) -> tuple[bool, dict[str, Any] | None]: - """ - Once a request is finished, determine whether request blocks + """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. """ from vllm.v1.request import RequestStatus diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 1530b9fb6a69..fea3a00ce247 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -45,8 +45,7 @@ def __init__( super().__init__(vllm_config, engine_id, kv_cache_config) def start_load_kv(self, metadata: NixlConnectorMetadata): - """ - Start loading by triggering non-blocking nixl_xfer. + """Start loading by triggering non-blocking nixl_xfer. We check for these trnxs to complete in each step(). """ for req_id, meta in metadata.reqs_to_recv.items(): @@ -373,8 +372,7 @@ def _read_blocks( expected_consumers: int, awaiting_kvs: bool, ) -> bool: - """ - Post a READ point-to-point xfer request from a single local worker to + """Post a READ point-to-point xfer request from a single local worker to a single remote worker. Returns True when the read was posted (or was unnecessary), False @@ -603,8 +601,7 @@ def _read_blocks_mixed( self._recving_transfers[request_id].append(handle) def _get_new_notifs(self) -> set[str]: - """ - Get req_ids which got a remote xfer message. When multiple consumers + """Get req_ids which got a remote xfer message. When multiple consumers are reading from the same producer (heterogeneous TP or DCP scenario), wait for all consumers to be done pulling. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 01c956a5251d..a61e5f8f7838 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -315,7 +315,8 @@ def _send_registration_to_p( executes on the writer; the handshake runs on the background executor and the request is re-queued onto ``_reg_send_inbox`` once it completes (at which point ``_ensure_handshake`` returns ``None`` and we - send directly).""" + send directly). + """ remote_pp_size = reg_data.get("remote_pp_size", 1) fut = self._ensure_handshake( reg_data["remote_engine_id"], @@ -507,7 +508,8 @@ def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds: ``BlockIds`` is canonically a tuple of per-group lists, but some registration payloads collapse a single-group case to a flat list. Re-wrap that case so downstream group-aware helpers see a - consistent shape.""" + consistent shape. + """ if block_ids and not isinstance(block_ids[0], (list, tuple)): return (list(block_ids),) return block_ids diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index 1e4f5c48e0f7..e66046c0ec7b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -23,7 +23,7 @@ @dataclass class NixlKVConnectorStats(KVConnectorStats): - """Container for transfer performance metrics""" + """Container for transfer performance metrics.""" def __post_init__(self): if not self.data: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py index b86061673487..6ec0e7307810 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py @@ -34,8 +34,7 @@ # TODO: merge with vllm.utils.network_utils.zmq_socket_ctx @contextlib.contextmanager def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: - """Context manager for a ZMQ socket""" - + """Context manager for a ZMQ socket.""" if socket_type not in (zmq.ROUTER, zmq.REQ): raise ValueError(f"Unexpected socket type: {socket_type}") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py index 27f08c1fd834..3c51fa277feb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py @@ -37,7 +37,8 @@ def canonical_format_id(kv_cache_layout: str) -> str: Canonical pages keep the worker's KV layout family, so the id couples the format version with that family; consumers must match it exactly. The family keeps its historical NHD/HND spelling so ids stay stable for KV persisted - before the layout enum existed.""" + before the layout enum existed. + """ layout = KVCacheLayout[kv_cache_layout] legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"} family = legacy.get(layout, layout.name.lower()) @@ -121,7 +122,8 @@ def _interleave_cp_tokens( ctx: _RankContext, ) -> tuple[CopyRun, ...]: """Place each region's num_tokens rows at their canonical token positions, - one run per chunk of interleaved tokens.""" + one run per chunk of interleaved tokens. + """ runs: list[CopyRun] = [] for region in regions: if ctx.cp_size == 1: @@ -256,7 +258,8 @@ def _attention_byte_regions( cp_size: int, ) -> list[ByteRegion] | None: """Byte regions of an attention page, given this rank's head shard. - None when the physical layout is not recognized (fail closed).""" + None when the physical layout is not recognized (fail closed). + """ bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size if tuple(kv_cache.shape) == (num_blocks, heads, bs, 2 * head_size): return _packed_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size) @@ -370,7 +373,8 @@ def _is_exact_partition(intervals: list[tuple[int, int]], size: int) -> bool: def _verify_tiling(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None: """Whichever ranks a block elects as writers must tile the canonical page - exactly once, and each rank's runs must cover exactly its local page.""" + exactly once, and each rank's runs must cover exactly its local page. + """ size = per_rank[0].canonical_page_size_bytes num_writers = per_rank[0].num_writers for mapping in per_rank: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 98720b5b7072..a8789d81c4ef 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -79,7 +79,8 @@ def get_offloading_event_group_spec( class _OffloadEventMetadata: """BlockStored payload snapshot for one OffloadKey, captured while the Request is available and kept until the final matching removal event. - ``medium`` and ``ownership`` are forwarded from the OffloadingEvent.""" + ``medium`` and ``ownership`` are forwarded from the OffloadingEvent. + """ # The chunk's constituent block hashes; the last one is the OffloadKey. block_hashes: tuple[BlockHash, ...] @@ -245,6 +246,7 @@ def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEven Yields: ``BlockStored`` or ``BlockRemoved`` events corresponding to the underlying :class:`OffloadingEvent` stream. + """ for event in events: if event.removed: @@ -254,7 +256,8 @@ def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEven def reset(self) -> None: """Drop all tracked state; pending payloads are stale after a - manager cache reset.""" + manager cache reset. + """ self._pending_event_metadata.clear() def _build_event_metadata( @@ -265,7 +268,8 @@ def _build_event_metadata( ) -> _OffloadEventMetadata: """Build the payload snapshot for one offloaded chunk: its constituent per-block hashes, the whole chunk's tokens, and the - per-block ``block_size``.""" + per-block ``block_size``. + """ hashes_per_chunk = group_config.hashes_per_chunk assert hashes_per_chunk > 0 assert chunk_idx >= 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index 58023c393561..0a4ac96be39f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -170,8 +170,7 @@ class _StatsKey: @dataclass class OffloadingConnectorStats(KVConnectorStats): - """ - Offloading connector stats use flat metric names as keys. + """Offloading connector stats use flat metric names as keys. The ``data`` dict is structured using ``_StatsKey`` / ``_MetricType``:: @@ -243,8 +242,7 @@ def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats": return self def reduce(self) -> dict[str, int | float]: - """ - Reduce the observations collected during a time interval to one or + """Reduce the observations collected during a time interval to one or more representative values (eg avg/median/sum of the series). This is meant to be called by the logger to produce a summary of the stats for the last time interval. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index b253a2399867..961980ee896c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -676,7 +676,8 @@ def _maximal_prefix_lookup( start_chunk_idx: int, ) -> int | None: """Return the number of consecutive offloaded chunks from the start, - or None if the backend deferred a lookup.""" + or None if the backend deferred a lookup. + """ hit_count = 0 defer_lookup = False for local_idx, key in enumerate(keys): @@ -711,7 +712,8 @@ def _sliding_window_lookup( """Return the end index (in `keys`) of the last run of `sliding_window_size` consecutive hits, scanning from the end. The first run may need a larger window for a partial rightmost chunk. - Returns 0 on miss, None if the backend deferred a lookup.""" + Returns 0 on miss, None if the backend deferred a lookup. + """ defer_lookup = False pending_in_window = False consecutive_hits = 0 @@ -750,8 +752,7 @@ def _lookup_complete_chunks( req_status: RequestOffloadState, max_num_new_tokens: int | None = None, ) -> int | None: - """ - Find how many tokens beyond num_locally_computed_tokens can be loaded. + """Find how many tokens beyond num_locally_computed_tokens can be loaded. Iterates full-attention groups first (prefix lookup), then sliding-window groups (suffix lookup). Each group may tighten max_hit_size_tokens, which @@ -1039,14 +1040,15 @@ def get_num_new_matched_tokens( num_computed_tokens: int, max_num_new_tokens: int | None = None, ) -> tuple[int | None, bool]: - """ - Get number of new tokens that can be loaded beyond the + """Get number of new tokens that can be loaded beyond the num_computed_tokens. Args: request (Request): the request object. num_computed_tokens (int): the number of locally computed tokens for this request + max_num_new_tokens (int | None): cap on the number of tokens that + may be loaded beyond `num_computed_tokens`, if any. Returns: A tuple with the following elements: @@ -1057,6 +1059,7 @@ def get_num_new_matched_tokens( should query for this request again later. - `True` if tokens will be loaded asynchronously (between scheduler steps). + """ req_status = self._req_status[request.request_id] for group_state in req_status.group_states: @@ -1208,10 +1211,7 @@ def update_state_after_alloc( req_status.partial_tail_boundary = None def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: - """ - Update request states from the Scheduler's output. - """ - + """Update request states from the Scheduler's output.""" # new_block_ids_end[req_id][i] = end of pre-existing block_ids for # the i-th sliding window group (before this step's extend). # Used to detect sliding window blocks that got re-allocated. @@ -1840,12 +1840,12 @@ def has_pending_push_work(self) -> bool: return bool(self._jobs) or self.manager.has_pending_work() def update_connector_output(self, connector_output: KVConnectorOutput): - """ - Update KVConnector state from worker-side connectors output. + """Update KVConnector state from worker-side connectors output. Args: connector_output (KVConnectorOutput): the worker-side connectors output. + """ meta = connector_output.kv_connector_worker_meta if not isinstance(meta, OffloadingWorkerMetadata): @@ -1938,8 +1938,7 @@ def request_finished( self, request: Request, ) -> tuple[bool, dict[str, Any] | None]: - """ - Called when a request has finished, before its blocks are freed. + """Called when a request has finished, before its blocks are freed. Returns: True if the request is being saved/sent asynchronously and blocks @@ -1947,6 +1946,7 @@ def request_finished( get_finished(). Optional KVTransferParams to be included in the request outputs returned by the engine. + """ req_status = self._req_status.get(request.request_id) @@ -1985,12 +1985,12 @@ def take_events(self) -> Iterable[KVCacheEvent]: Yields: ``BlockStored`` or ``BlockRemoved`` events corresponding to the underlying :class:`OffloadingEvent` stream. + """ yield from self._events_tracker.take_events(self.manager.take_events()) def reset_cache(self) -> None: """Reset the offloading manager cache, evicting all stored chunks.""" - # reset_cache cannot be called in the middle of a schedule step assert not self._current_batch_load_jobs assert not self._current_batch_jobs_to_flush diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 14e826e5525c..94e8a8109d4e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -39,7 +39,7 @@ class OffloadingConnectorWorker: - """Implementation of Worker side methods""" + """Implementation of Worker side methods.""" def __init__( self, @@ -268,14 +268,14 @@ def prepare_store_kv(self, metadata: OffloadingConnectorMetadata): ) def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: - """ - Returns: - tuple of (finished_sending, finished_recving). Stores never - emit finished_sending — the scheduler tracks store completion - via kv_connector_worker_meta.completed_jobs and fences any - block reuse via jobs_to_flush. Loads still emit - finished_recving so the base scheduler can resume requests - blocked on remote KV (and free aborted-during-load reqs). + """Returns: + tuple of (finished_sending, finished_recving). Stores never + emit finished_sending — the scheduler tracks store completion + via kv_connector_worker_meta.completed_jobs and fences any + block reuse via jobs_to_flush. Loads still emit + finished_recving so the base scheduler can resume requests + blocked on remote KV (and free aborted-during-load reqs). + """ assert self.worker is not None finished_recving: set[str] = set() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py index 00dc05bfc4a9..04ce6a0b92c5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -83,6 +83,7 @@ def remote_conv_offsets( < 0: P_TP > D_TP — P pages are smaller, D reads entire P page. Local dims are scaled down by |tp_ratio| to get P-sized offsets. + """ offsets: list[tuple[int, int]] = [] if tp_ratio >= 1: @@ -121,6 +122,7 @@ def derive_mamba_conv_split( Returns: MambaConvSplitInfo with per-rank sub-projection dims, conv_rows, conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes). + """ _supported = ( MambaAttentionBackendEnum.MAMBA1, @@ -228,5 +230,6 @@ def compute_physical_blocks_per_logical( Args: ssm_sizes: (conv_state_bytes, ssm_state_bytes) from NixlAgentMetadata. block_len: the engine's block_len in bytes (from block_lens[0]). + """ return math.ceil((ssm_sizes[0] + ssm_sizes[1]) / block_len) diff --git a/vllm/distributed/kv_transfer/kv_transfer_state.py b/vllm/distributed/kv_transfer/kv_transfer_state.py index 690eb9fa23d9..58d77f827734 100644 --- a/vllm/distributed/kv_transfer/kv_transfer_state.py +++ b/vllm/distributed/kv_transfer/kv_transfer_state.py @@ -38,6 +38,7 @@ def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> boo Note: This function will no-longer be needed after the v1 KV connector becomes the default. + """ if connector is None: connector = _KV_CONNECTOR_AGENT @@ -72,10 +73,7 @@ def _sync_engine_id_across_tp(vllm_config: "VllmConfig") -> None: def ensure_kv_transfer_initialized( vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig" ) -> None: - """ - Initialize KV cache transfer parallel group. - """ - + """Initialize KV cache transfer parallel group.""" global _KV_CONNECTOR_AGENT if _KV_CONNECTOR_AGENT is not None: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index dc3022a16fec..60ed229bc463 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -141,9 +141,11 @@ def _split_tensor_dict( def _get_unique_name(name: str) -> str: """Get a unique name for the group. + Example: _get_unique_name("tp") -> "tp:0" _get_unique_name("tp") -> "tp:1" + """ if name not in _group_name_counter: _group_name_counter[name] = 0 @@ -420,8 +422,7 @@ def patched_fused_scaled_matmul_reduce_scatter( class GroupCoordinator: - """ - PyTorch ProcessGroup wrapper for a group of processes. + """PyTorch ProcessGroup wrapper for a group of processes. PyTorch ProcessGroup is bound to one specific communication backend, e.g. NCCL, Gloo, MPI, etc. GroupCoordinator takes charge of all the communication operations among @@ -634,34 +635,34 @@ def create_single_reader_mq_broadcasters( @property def first_rank(self): - """Return the global rank of the first process in the group""" + """Return the global rank of the first process in the group.""" return self.ranks[0] @property def last_rank(self): - """Return the global rank of the last process in the group""" + """Return the global rank of the last process in the group.""" return self.ranks[-1] @property def is_first_rank(self): - """Return whether the caller is the first process in the group""" + """Return whether the caller is the first process in the group.""" return self.rank == self.first_rank @property def is_last_rank(self): - """Return whether the caller is the last process in the group""" + """Return whether the caller is the last process in the group.""" return self.rank == self.last_rank @property def next_rank(self): - """Return the global rank of the process that follows the caller""" + """Return the global rank of the process that follows the caller.""" rank_in_group = self.rank_in_group world_size = self.world_size return self.ranks[(rank_in_group + 1) % world_size] @property def prev_rank(self): - """Return the global rank of the process that precedes the caller""" + """Return the global rank of the process that precedes the caller.""" rank_in_group = self.rank_in_group world_size = self.world_size return self.ranks[(rank_in_group - 1) % world_size] @@ -719,8 +720,7 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None yield graph_capture_context def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: - """ - User-facing all-reduce function before we actually call the + """User-facing all-reduce function before we actually call the all-reduce operation. We need this because Dynamo does not support passing an arbitrary @@ -809,8 +809,7 @@ def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Ten def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: - """ - NOTE: We assume that the input tensor is on the same device across + """NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ @@ -1363,7 +1362,7 @@ def barrier(self): torch.distributed.barrier(group=self.cpu_group) def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: - """Sends a tensor to the destination rank in a blocking way""" + """Sends a tensor to the destination rank in a blocking way.""" """NOTE: `dst` is the local rank of the destination rank.""" if self.device_communicator is None: raise ValueError("No device communicator found") @@ -1637,8 +1636,7 @@ def graph_capture( device: torch.device, graph_capture_context: GraphCaptureContext | None = None, ): - """ - `graph_capture` is a context manager which should surround the code that + """`graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some operations will be run after the graph is captured, before the graph is replayed. It returns a `GraphCaptureContext` object which contains the @@ -1981,14 +1979,17 @@ def initialize_model_parallel( decode_context_model_parallel_size: int | None = 1, backend: str | None = None, ) -> None: - """ - Initialize model parallel groups. + """Initialize model parallel groups. Arguments: tensor_model_parallel_size: number of GPUs used for tensor model parallelism. pipeline_model_parallel_size: number of GPUs used for pipeline model parallelism. + prefill_context_model_parallel_size: number of GPUs used for context + parallelism during prefill. + decode_context_model_parallel_size: number of GPUs used for context + parallelism during decode. backend: name of torch distributed communication backend. Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we @@ -2003,6 +2004,7 @@ def initialize_model_parallel( are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box. + """ # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() @@ -2476,8 +2478,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): def in_the_same_node_as( pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0 ) -> list[bool]: - """ - This is a collective operation that returns if each rank is in the same node + """This is a collective operation that returns if each rank is in the same node as the source rank. It tests if processes are attached to the same memory system (shared access to shared memory). """ @@ -2568,8 +2569,7 @@ def in_the_same_node_as( def is_global_first_rank() -> bool: - """ - Check if the current process is the first rank globally across all + """Check if the current process is the first rank globally across all parallelism strategies (PP, TP, DP, EP, etc.). Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0` @@ -2579,6 +2579,7 @@ def is_global_first_rank() -> bool: Returns: bool: True if this is the global first rank (rank 0), False otherwise. Returns True if distributed is not initialized (single process). + """ try: # If world group is available, use it for the most accurate check @@ -2599,9 +2600,7 @@ def is_global_first_rank() -> bool: def is_local_first_rank() -> bool: - """ - Check if the current process is the first local rank (rank 0 on its node). - """ + """Check if the current process is the first local rank (rank 0 on its node).""" try: # prefer the initialized world group if available global _WORLD @@ -2622,14 +2621,14 @@ def is_local_first_rank() -> bool: def _node_count(pg: ProcessGroup | StatelessProcessGroup) -> int: - """ - Returns the total number of nodes in the process group. + """Returns the total number of nodes in the process group. Args: pg: The process group to analyze Returns: int: The total number of nodes + """ if isinstance(pg, ProcessGroup): world_size = torch.distributed.get_world_size(group=pg) diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 38c74a97c55c..6ce2aa89ac75 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -59,8 +59,7 @@ def _fetch_group_ports(key: str, coord_store: Store) -> list[int]: class StatelessGroupCoordinator(GroupCoordinator): - """ - A stateless version of the GroupCoordinator class in parallel_state, + """A stateless version of the GroupCoordinator class in parallel_state, It will create CPU, device and TCPStore based communication groups that are independent of PyTorch's WORLD group. Hence, communication groups with a different set of participants GPUs diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index eec9890f0286..f8304200749c 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -59,7 +59,8 @@ def ensure_divisibility(numerator, denominator): def divide(numerator, denominator): """Ensure that numerator is divisible by the denominator and return - the division value.""" + the division value. + """ ensure_divisibility(numerator, denominator) return numerator // denominator @@ -111,6 +112,7 @@ def split_tensor_along_last_dim( Returns: A list of Tensors + """ # Get the size and dimension. last_dim = tensor.dim() - 1 @@ -333,7 +335,6 @@ def all_reduce( def barrier(self, timeout: float = 30.0): """A robust barrier to synchronize all ranks. - Uses a multi-phase approach to ensure all processes reach the barrier before proceeding: @@ -351,6 +352,7 @@ def barrier(self, timeout: float = 30.0): Raises: RuntimeError: If coordination fails or times out + """ # Generate a barrier ID that is globally unique try: @@ -549,8 +551,7 @@ def init_gloo_process_group( group_size: int, timeout: timedelta, ) -> ProcessGroup: - """ - Stateless init ProcessGroup with gloo backend compatible with + """Stateless init ProcessGroup with gloo backend compatible with different torch versions. """ with suppress_stdout(): @@ -583,8 +584,7 @@ def stateless_init_torch_distributed_process_group( return_store: bool = False, listen_socket: socket.socket | None = None, ) -> ProcessGroup | tuple[ProcessGroup, Store]: - """ - A replacement for `torch.distributed.init_process_group` that does not + """A replacement for `torch.distributed.init_process_group` that does not pollute the global state. The created ProcessGroup object can be used for some operations such as `allreduce`, because it does not depend on the global rank. However, some operations such as `broadcast` cannot be used @@ -685,9 +685,8 @@ def stateless_init_torch_distributed_process_group( def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None: - """ - Destroy ProcessGroup returned by - stateless_init_torch_distributed_process_group(). + """Destroy ProcessGroup returned by + stateless_init_torch_distributed_process_group(). """ pg.shutdown() _unregister_process_group(pg.group_name) @@ -706,6 +705,7 @@ def get_worker_rank_suffix(global_rank: int | None = None) -> str: Returns: A string suffix identifying the worker's position in the distributed topology. + """ from vllm.distributed.parallel_state import ( get_dcp_group, diff --git a/vllm/distributed/weight_transfer/__init__.py b/vllm/distributed/weight_transfer/__init__.py index c78fd1f3cc5f..3fcca3282691 100644 --- a/vllm/distributed/weight_transfer/__init__.py +++ b/vllm/distributed/weight_transfer/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Weight transfer engines for syncing model weights from trainers +"""Weight transfer engines for syncing model weights from trainers to inference workers. """ diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2328829fc0bc..2fb60fb71b41 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -193,6 +193,7 @@ def held_names(self) -> "Collection[str] | None": Returns: The held parameter names, or None to hold every one. + """ return None @@ -329,8 +330,7 @@ class WeightTransferUpdateRequest: class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): - """ - Base class for weight transfer engines that handle transport of model weights + """Base class for weight transfer engines that handle transport of model weights from a trainer to inference workers. This abstraction separates weight transfer transport logic from the worker @@ -386,14 +386,14 @@ def __init__( device: torch.device, model: torch.nn.Module, ) -> None: - """ - Initialize the weight transfer engine. + """Initialize the weight transfer engine. Args: config: The configuration for the weight transfer engine vllm_config: The full vLLM config (provides parallel/model config) device: The device this worker's model lives on model: The local model instance which will receive the weights + """ self.config = config self.vllm_config = vllm_config @@ -419,8 +419,7 @@ def reset_weight_update_target(self) -> None: self.model_config = self._default_model_config def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: - """ - Construct typed init info from dict with validation. + """Construct typed init info from dict with validation. Args: init_dict: Dictionary containing backend-specific initialization parameters @@ -430,6 +429,7 @@ def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: Raises: ValueError: If init_dict is invalid for this backend + """ try: return self.init_info_cls(**init_dict) @@ -439,8 +439,7 @@ def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: ) from e def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo: - """ - Construct typed update info from dict with validation. + """Construct typed update info from dict with validation. Args: update_dict: Dictionary containing backend-specific update parameters @@ -450,6 +449,7 @@ def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo: Raises: ValueError: If update_dict is invalid for this backend + """ try: return self.update_info_cls(**update_dict) @@ -460,19 +460,18 @@ def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo: @abstractmethod def init_transfer_engine(self, init_info: TInitInfo) -> None: - """ - Initialize the weight transfer mechanism. + """Initialize the weight transfer mechanism. This is called once at the beginning of training. Args: init_info: Backend-specific initialization info + """ raise NotImplementedError @abstractmethod def start_weight_update(self) -> None: - """ - Prepare the engine for a new weight update. + """Prepare the engine for a new weight update. Engines that receive weights in checkpoint format initialize layerwise reloading here, else this is typically a no-op. @@ -482,8 +481,7 @@ def start_weight_update(self) -> None: @abstractmethod def finish_weight_update(self) -> None: - """ - Finalize the current weight update. + """Finalize the current weight update. Checkpoint-format engines finalize layerwise reloading here; engines that apply weights in place leave this as a no-op. @@ -491,11 +489,11 @@ def finish_weight_update(self) -> None: raise NotImplementedError def update_weights(self, update_info: dict[str, Any]) -> None: - """ - Receive one weight update chunk and load it into the model. + """Receive one weight update chunk and load it into the model. Args: update_info: Dictionary containing backend-specific update info + """ typed_update_info = self.parse_update_info(update_info) self.receive_weights(typed_update_info) @@ -505,19 +503,18 @@ def update_weights(self, update_info: dict[str, Any]) -> None: @abstractmethod def receive_weights(self, update_info: TUpdateInfo) -> None: - """ - Receive weights from the trainer and load them into the model. + """Receive weights from the trainer and load them into the model. Args: update_info: Backend-specific update info containing parameter metadata and any backend-specific data + """ raise NotImplementedError @abstractmethod def shutdown(self) -> None: - """ - Shutdown the weight transfer engine. + """Shutdown the weight transfer engine. This should be called when the worker is shutting down. """ raise NotImplementedError diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index ba486422968f..74a04cc9d8f7 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -58,6 +58,7 @@ def register_engine( Raises: ValueError: If an engine with the same name is already registered + """ if name in cls._registry: raise ValueError(f"Weight transfer engine '{name}' is already registered.") @@ -101,6 +102,7 @@ def create_engine( Raises: ValueError: If the backend is not registered + """ backend = config.backend if backend not in cls._registry: @@ -139,7 +141,8 @@ def register_engine( class_name: str | None = None, ) -> None: """Register a trainer engine. Same conventions as - `WeightTransferEngineFactory.register_engine`.""" + `WeightTransferEngineFactory.register_engine`. + """ if name in cls._registry: raise ValueError( f"Weight transfer trainer engine '{name}' is already registered." @@ -190,6 +193,7 @@ def trainer_init( Raises: ValueError: If `init_info.backend` is not registered. + """ backend = init_info.backend if backend not in cls._registry: diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 7b1563112d19..83b5d15f58af 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -40,7 +40,8 @@ class IPCWeightTransferInitInfo(WeightTransferInitInfo): `packed` is a must-agree wire param: the trainer ships it here at the init handshake so the worker decodes with the same setting the trainer encoded with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does - not need the buffer size (producer-only).""" + not need the buffer size (producer-only). + """ packed: bool = False @@ -53,7 +54,8 @@ class IPCTrainerInitInfo(TrainerInitInfo): `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The trainer propagates `packed` to the worker at `trainer_init` so the two sides - cannot disagree. `backend` is the factory dispatch key.""" + cannot disagree. `backend` is the factory dispatch key. + """ backend: ClassVar[str] = "ipc" packed: bool = False @@ -121,8 +123,7 @@ def __post_init__(self): class IPCWeightTransferEngine( WeightTransferEngine[IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo] ): - """ - Weight transfer engine using CUDA IPC for communication between trainer and workers. + """Weight transfer engine using CUDA IPC between trainer and workers. This implementation uses CUDA IPC to transfer weights from the trainer (rank 0) to all inference workers in a process group. IPC handles are used to share @@ -149,13 +150,13 @@ def __init__( self._packed_importer = PackedBufferImporter() def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None: - """ - Initialize the weight transfer mechanism. No data-plane rendezvous is + """Initialize the weight transfer mechanism. No data-plane rendezvous is needed for IPC; this just records the trainer-supplied wire params so the worker decodes exactly as the trainer encoded. Args: init_info: IPC initialization info (carries `packed`). + """ self.packed = init_info.packed @@ -182,8 +183,7 @@ def finish_weight_update(self) -> None: self._packed_importer.close() def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: - """ - Receive weights from the trainer via CUDA IPC handles and load them. + """Receive weights from the trainer via CUDA IPC handles and load them. Whether the transfer is packed is read from `self.packed`, set at the init handshake from the trainer's init info, so it is guaranteed to @@ -193,6 +193,7 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: update_info: IPC update info containing parameter names, dtypes, shapes, and IPC handles. Each IPC handle is a mapping between physical GPU UUID and the rebuild_cuda_tensor args tuple. + """ # Use the worker's assigned device rather than the ambient current # device: the receive path is no longer wrapped in diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index 67e578640a0d..09e31f0291af 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -118,7 +118,8 @@ def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict: """Serialize a worker init info for `init_weight_transfer_engine`, dropping the unset rendezvous field (the UID in TCP mode) so the wire payload carries only the mode actually in use. Shared by the dense and sparse trainer - engines so the two cannot drift.""" + engines so the two cannot drift. + """ return {key: value for key, value in asdict(init_info).items() if value is not None} @@ -142,8 +143,7 @@ def stateless_init_process_group( world_size: int, device, ) -> "PyNcclCommunicator": - """ - vLLM provides `StatelessProcessGroup` to create a process group + """VLLM provides `StatelessProcessGroup` to create a process group without considering the global process group in torch.distributed. It is recommended to create `StatelessProcessGroup`, and then initialize the data-plane communication (NCCL) between external (train processes) @@ -222,8 +222,7 @@ def worker_init_process_group( def trainer_init( init_info: NCCLRendezvous | dict, ) -> "PyNcclCommunicator": - """ - Initialize NCCL process group for trainer-side weight transfer. + """Initialize NCCL process group for trainer-side weight transfer. The trainer is always rank 0 in the process group. Uses the current CUDA device (torch.accelerator.current_device_index()). @@ -237,6 +236,7 @@ def trainer_init( Returns: PyNcclCommunicator for weight transfer. + """ if isinstance(init_info, dict): master_address = init_info["master_address"] diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 1efd8180dfd5..91298777a918 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -69,7 +69,8 @@ class NCCLTrainerInitInfo(TrainerInitInfo): propagates them to the worker at `trainer_init` so the two sides cannot disagree. Note this defaults to packed, unlike the worker-side `NCCLWeightTransferInitInfo`, whose default only applies when no trainer - ships a value. `backend` is the factory dispatch key.""" + ships a value. `backend` is the factory dispatch key. + """ backend: ClassVar[str] = "nccl" @@ -113,8 +114,7 @@ def __post_init__(self): class NCCLWeightTransferEngine( WeightTransferEngine[NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo] ): - """ - Weight transfer engine using NCCL for communication between trainer and workers. + """Weight transfer engine using NCCL for communication between trainer and workers. This implementation uses NCCL broadcast operations to transfer dense checkpoint-format weights from the trainer (rank 0) to all inference workers @@ -142,14 +142,14 @@ def __init__( self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: - """ - Initialize NCCL process group with the trainer and record the + """Initialize NCCL process group with the trainer and record the trainer-supplied wire params so the worker decodes exactly as the trainer encodes. Args: init_info: NCCL initialization info containing master address, port, rank offset, world size, and the packed wire params + """ self.packed = init_info.packed self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes @@ -175,8 +175,7 @@ def finish_weight_update(self) -> None: finalize_layerwise_reload(self.model, self.model_config) def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None: - """ - Receive weights from trainer via NCCL broadcast. + """Receive weights from trainer via NCCL broadcast. Whether to use packed broadcasting (and the buffer geometry) is read from `self.packed` / `self.packed_*`, set at the init handshake from the @@ -186,6 +185,7 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None: Args: update_info: NCCL update info containing parameter names, dtypes, and shapes + """ if self.model_update_group is None: raise RuntimeError( @@ -365,7 +365,8 @@ def send_weights(self) -> None: def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None: """Iterate the source (materializing each tensor — a collective on all ranks) and, on the sender, broadcast from rank 0, packed or one-by-one. - Non-sender ranks only replay the iteration to stay in the collective.""" + Non-sender ranks only replay the iteration to stay in the collective. + """ if not self.is_sender: for _ in source: pass diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index 0eb9ede2639b..dcd8efc8c4af 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -45,6 +45,7 @@ def unpack_tensor( shapes: List of tensor shapes dtypes: List of tensor dtypes tensor_sizes: List of tensor sizes in bytes + """ unpacked_tensors = packed_tensor.split(tensor_sizes) @@ -85,6 +86,7 @@ def pack_tensors( tensor_list: Pre-existing tensor list to append to (for NCCL multi-buffer reuse). If None, a fresh list is created. current_size: Byte count already accumulated in tensor_list + """ if tensor_list is None: tensor_list = [] @@ -335,6 +337,7 @@ def packed_ipc_producer( buffer_size_bytes: Exact capacity of the reusable IPC buffer. Every chunk is guaranteed to fit within this size. A ``ValueError`` is raised if any single tensor exceeds it. + """ ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="cuda") _, ipc_args = reduce_tensor(ipc_buffer) @@ -463,6 +466,7 @@ def packed_ipc_consumer( export so the buffer mapping is opened and released exactly once (see ``PackedBufferImporter``). ``None`` builds an ephemeral one, which is only safe for single-chunk consumption. + """ props = torch.cuda.get_device_properties(device_index) physical_gpu_id = str(props.uuid) diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py index 70110517d3fc..eeda36596577 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_common.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py @@ -60,6 +60,7 @@ def check_ray_rdt_version() -> None: Raises: ValueError: the installed Ray predates ``RDT_MIN_RAY_VERSION``. + """ import importlib.metadata @@ -162,6 +163,7 @@ def __init__( ValueError: ``workers_per_replica`` does not divide ``num_consumers``, so the deployments are not uniform and two workers of one deployment would share a block index. + """ self.num_producers = max(1, num_producers) self.num_consumers = max(1, num_consumers) @@ -212,7 +214,8 @@ def __init__( def class_of(self, name: str) -> int: """The name's owner class — the planner's bucketing key, since all names - of a chunk must share one producer.""" + of a chunk must share one producer. + """ return self._class_of.get(name, 0) def owners(self, name: str) -> list[int]: @@ -263,6 +266,7 @@ def validate(self) -> None: Raises: ValueError: an owner set is empty or out of range, or a class index does not resolve. + """ for c, row in enumerate(self._owner_sets): if not row: @@ -282,7 +286,8 @@ def validate(self) -> None: def bind(self, actors: list, produce_methods: list, consumer_id: int) -> None: """Attach this consumer's producer handles. Rebuilt wholesale per init, never appended to: every owner index is a position in these lists, so a - rejoining engine that re-inits must not shift them.""" + rejoining engine that re-inits must not shift them. + """ if len(actors) != len(produce_methods): raise ValueError("actors and produce_methods must be parallel") self._actors = list(actors) @@ -345,5 +350,6 @@ def buffer_alloc_bytes(nbytes: int, presize: int = 0) -> int: optional ``presize`` floor, and a coarse 256MB round-up, so the buffer is allocated ONCE and never regrows. Regrowth is a correctness hazard, not just a perf one -- see ``buffer_presize_gb``. Shared by the consumer's receive buffers - and the producer's serve rings.""" + and the producer's serve rings. + """ return max(nbytes, presize, -(-nbytes // (256 << 20)) * (256 << 20)) diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index 96dd53f0fdb9..cfe84d2dd033 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -134,7 +134,8 @@ class _PendingPull: ``targets``/``blob`` must stay strongly referenced until it completes: ``set_target_for_ref`` stores WEAKREFS, so dropping them silently reroutes the - transfer into a fallback buffer.""" + transfer into a fallback buffer. + """ ref: "Any" keys: "list[FetchKey]" @@ -561,7 +562,8 @@ def _resolve_consumer_id(self, init_info: ShardedRDTWeightTransferInitInfo) -> i def _build_static_plan(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: """Build the chunk/free plan once. It never changes across syncs, so - ``update_weights`` needs no per-sync names.""" + ``update_weights`` needs no per-sync names. + """ if not init_info.group_lens: raise ValueError( "Sharded RDT engine requires init_info.group_lens (the gather-" @@ -585,7 +587,8 @@ def _preregister_at_init(self) -> None: Both sides are sized from the static plan: receive buffers are ``ring_depth`` slots at the largest chunk's ``pack_bytes``, and each bound producer is asked to pre-register a serve ring at the max bytes this - consumer will pull from it. A no-op when this worker has no chunks.""" + consumer will pull from it. A no-op when this worker has no chunks. + """ plan = self._cached_plan if plan is None or not plan.chunks: return @@ -630,7 +633,8 @@ def _global_worker_index(self) -> int: ``data_parallel_index``, not ``data_parallel_rank``: vLLM resets the latter to 0 in a dense worker but keeps the former as the distinct global DP rank. Same formula as the sibling ``nccl_engine``, so dense-via-TP and - MoE-via-DP+EP both yield distinct 0..C-1.""" + MoE-via-DP+EP both yield distinct 0..C-1. + """ pc = self.parallel_config return pc.data_parallel_index * pc.world_size + pc.rank # world_size = TP*PP @@ -638,14 +642,16 @@ def _num_consumers(self) -> int: """Total inference-worker count. Prefers the driver-supplied ``init_info.num_consumers`` (authoritative -- the driver knows the whole fleet); else ``world_size_across_dp``, the same stride - ``_global_worker_index`` indexes with, so the two agree at any pp.""" + ``_global_worker_index`` indexes with, so the two agree at any pp. + """ if self._num_consumers_override > 0: return self._num_consumers_override return self.parallel_config.world_size_across_dp def start_weight_update(self) -> None: """Put the model's params on meta so layerwise reload streams them in - as each layer's slices land. Baked replay uses checkpoint format.""" + as each layer's slices land. Baked replay uses checkpoint format. + """ from vllm.model_executor.model_loader.reload import ( initialize_layerwise_reload, ) @@ -654,7 +660,8 @@ def start_weight_update(self) -> None: def finish_weight_update(self) -> None: """Drain the deferred pull/process pipeline (so every layer is fully - loaded) before finalizing the layerwise reload.""" + loaded) before finalizing the layerwise reload. + """ from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, ) @@ -667,7 +674,8 @@ def update_weights(self, update_info: dict[str, Any]) -> None: device sync: post-processing is deferred to background threads and a sync here would block on them and serialize the pull/process pipeline. Completion is guaranteed by ``drain_pending`` in - ``finish_weight_update``.""" + ``finish_weight_update``. + """ self.receive_weights(self.parse_update_info(update_info)) def receive_weights( @@ -709,7 +717,8 @@ def _build_fake_weights( device: torch.device, ) -> list[tuple[str, torch.Tensor]]: """Zero-storage lazies for ``names``, dtype/shape from the init metadata, - all feeding the bake's recording sink.""" + all feeding the bake's recording sink. + """ return [ ( name, @@ -977,7 +986,8 @@ def _restore_after_dry_run(self, model: torch.nn.Module) -> None: def _ensure_proc_worker(self) -> None: """Lazily create the per-slot events, the background CUDA stream, the work - queue, and the single processing thread. Idempotent.""" + queue, and the single processing thread. Idempotent. + """ if self._proc_thread is not None: return import queue @@ -1036,7 +1046,8 @@ def _proc_worker_loop(self) -> None: def _mark_slot_done(self, slot: int) -> None: """Publish that a queued item's read-done event has been recorded (or the item failed) so a pull waiting to reuse ``slot`` can proceed to its - CUDA-event synchronize.""" + CUDA-event synchronize. + """ with self._slot_cv: self._slot_done[slot] += 1 self._slot_cv.notify_all() @@ -1052,7 +1063,8 @@ def drain_pending(self) -> None: """Block until the background thread has processed every queued item and its stream work is complete, then re-raise any error it hit. Called from the worker's ``finish_weight_update`` before ``finalize_layerwise_reload`` - so every baked layer is fully loaded (and ``info.reset()``-ed) first.""" + so every baked layer is fully loaded (and ``info.reset()``-ed) first. + """ if self._proc_queue is not None: self._proc_queue.join() # every put() item task_done()'d # The scatter thread feeds the quant thread, so join it SECOND (all @@ -1104,7 +1116,8 @@ def _chunk_module_scatters( which is exactly what an owner class is. The cut is a pure function of the bake and the ownership table. Copy order within a chunk is bake order, and a module's copies may span chunks (materialize/quant fire on - its first/last chunk; see ``_build_call_plan``).""" + its first/last chunk; see ``_build_call_plan``). + """ assert self._router is not None by_class: dict[int, list[_Scatter]] = {} for copies in modules: @@ -1235,7 +1248,8 @@ def _signal_group_done(self, group_idx: int) -> None: pulls nothing from. Refs are held and drained in ``drain_pending`` so every signal has EXECUTED before the sync ends: ``begin_sync`` resets the counters, and a straggler landing in the next sync would credit a - group it does not belong to.""" + group it does not belong to. + """ assert self._router is not None self._pending_frees.extend(self._router.free_group(group_idx)) @@ -1370,7 +1384,8 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: the modules' scatters via ``ready``; touches only the scattered params (never a receive slot), so it can overlap subsequent chunks' RDMA and scatters. ``info.reset()`` is what makes finalize skip the layer — - drain_pending joins the quant queue before finalize runs.""" + drain_pending joins the quant queue before finalize runs. + """ from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) @@ -1403,7 +1418,8 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: def _quant_worker_loop(self) -> None: """Dedicated quant thread: drains (completed_modules, scatter-done event) - batches. Errors surface via _proc_error like the scatter thread's.""" + batches. Errors surface via _proc_error like the scatter thread's. + """ torch.cuda.set_device(self.device) q = self._quant_queue assert q is not None diff --git a/vllm/distributed/weight_transfer/sharded_rdt_fake.py b/vllm/distributed/weight_transfer/sharded_rdt_fake.py index 5e3613e9d556..536329922ba8 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_fake.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_fake.py @@ -206,7 +206,8 @@ def _make_child( def _meta(self) -> torch.Tensor: """A zero-storage meta tensor of this fake's shape/dtype, so PyTorch - itself computes post-op geometry rather than us reimplementing it.""" + itself computes post-op geometry rather than us reimplementing it. + """ return torch.empty(self.shape, dtype=self.dtype, device="meta") @classmethod @@ -250,7 +251,8 @@ def _intercept( ): """Append the op and return a child, or a tuple of children for multi-return ops. Each child's geometry comes from running the op on a - meta tensor.""" + meta tensor. + """ meta = self_._meta() with torch._C.DisableTorchFunctionSubclass(): meta_result = func(meta, *args, **kwargs) diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 1be17c4fa088..4b16fece8fc9 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -132,7 +132,8 @@ @dataclass class _SharedPack: """One generation of one slot-sharing group: the rendezvous state for a - single chunk, identified by the consumers' issue index.""" + single chunk, identified by the consumers' issue index. + """ group: int slot: int @@ -466,7 +467,8 @@ def wait_freed(self) -> list[int]: def end_sync(self) -> list[int]: """Block until every published group has been freed by its consumers; - return the remaining freed keys so the engine drops its last refs.""" + return the remaining freed keys so the engine drops its last refs. + """ with self._cache_cond: self._wait_for( lambda: bool(self._inflight_groups), @@ -479,7 +481,8 @@ def end_sync(self) -> list[int]: def set_gather_error(self, message: str) -> None: """Record a trainer-side gather failure so blocked serves / publishes - stop waiting and surface it.""" + stop waiting and surface it. + """ with self._cache_cond: self._gather_error = RuntimeError(message) self._cache_cond.notify_all() @@ -508,7 +511,8 @@ def free_group(self, group_idx: int) -> None: def _new_serve_buffer(self, nbytes: int) -> torch.Tensor: """Allocate + NIXL-register one serve slot: the single allocation seam, so registration cannot be skipped on either of the two paths that make - buffers (the init-time reservation and the serve-path backstop).""" + buffers (the init-time reservation and the serve-path backstop). + """ t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) with self._reg_lock: register_nixl_memory(t) @@ -534,6 +538,7 @@ def reserve_serve_buffer( Raises: RuntimeError: two consumers of one sharing group disagree on the chunks they pull from this producer. + """ sg = self._share_group(consumer_id) if plan_digest is not None: @@ -579,13 +584,15 @@ def reserve_serve_buffer( def _share_group(self, consumer_id: int) -> int: """The slot-sharing group ``consumer_id`` belongs to: its index within its own deployment, since that is what fixes the plan. Without a width - the group is the consumer itself and nothing is shared.""" + the group is the consumer itself and nothing is shared. + """ return consumer_id % self._share_width if self._share_width > 0 else consumer_id def _sharers_of(self, sg: int) -> frozenset: """The LIVE consumers of group ``sg``: the rendezvous width. Derived from ``begin_sync``'s live set, so a degraded sync narrows the barrier instead - of waiting forever on a dead deployment. Caller holds ``_cache_cond``.""" + of waiting forever on a dead deployment. Caller holds ``_cache_cond``. + """ cached = self._sharers.get(sg) if cached is None: if not self._sharing_active or self._live_ids is None: @@ -621,7 +628,8 @@ def _join_share_group(self, consumer_id: int, seq: int) -> tuple: def _await_shared_pack(self, gen: _SharedPack) -> torch.Tensor: """Block until this generation's packer published its blob, and return it. A failed pack is re-raised here, so every sharer of a bad pack fails - instead of reading a half-written slot.""" + instead of reading a half-written slot. + """ with self._cache_cond: self._wait_for( lambda: not gen.done and gen.error is None, @@ -642,7 +650,8 @@ def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor: """Group ``sg``'s ring slot ``idx``, grown if this chunk outgrew the reservation. Growing registers memory while the fabric is busy, the hazard ``reserve_serve_buffer`` exists to avoid, so it is a backstop: the - reservation is sized from the same static plan as this pack.""" + reservation is sized from the same static plan as this pack. + """ with self._serve_lock: buffer = self._serve_rings.setdefault(sg, [None] * self._nring)[idx] if buffer is not None and buffer.numel() >= need: @@ -654,7 +663,8 @@ def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor: def _fail_shared_pack(self, gen: _SharedPack, exc: BaseException) -> None: """Publish a pack failure so waiting sharers raise instead of hanging on - a generation that will never complete.""" + a generation that will never complete. + """ with self._cache_cond: gen.error = exc self._cache_cond.notify_all() @@ -681,6 +691,7 @@ def rdt_produce_weights_batched( Raises: ValueError: ``seq`` was not supplied. + """ needed = sorted({n for n, _ in specs}) if self._served_names is not None: @@ -855,7 +866,8 @@ def __init__( def _rpc(self, method: str, *args: Any) -> Any: """Call one of the server actor's methods and block for the result. The single seam through which the engine talks to its server, so tests - can inject a local (non-Ray) fake server.""" + can inject a local (non-Ray) fake server. + """ import ray return ray.get(getattr(self._server, method).remote(*args)) @@ -932,7 +944,8 @@ def _publish_async(self, group_idx: int, entries): """Fire publish_group WITHOUT waiting on the RPC (the gather loop overlaps the publish's server-side rebuild with the next group's gather) and return a handle that ``_await_publish`` resolves. Ray actor - handle in production; a plain (non-Ray) fake server runs inline.""" + handle in production; a plain (non-Ray) fake server runs inline. + """ method = self._server.publish_group remote = getattr(method, "remote", None) if remote is not None: @@ -943,7 +956,8 @@ def _await_publish(self, ref) -> None: """Resolve one async publish so a server-side rebuild error surfaces at window depth instead of at end_sync. Publishes carry nothing back — freed groups flow only through wait_freed/end_sync (one channel; see - publish_group).""" + publish_group). + """ import ray if isinstance(ref, ray.ObjectRef): @@ -1023,6 +1037,7 @@ def get_worker_init_payload(self) -> dict: Raises: RuntimeError: called before ``trainer_init`` cached the server names. + """ if self._server_names is None: raise RuntimeError( @@ -1129,7 +1144,8 @@ def _validate_held_yields(self, gi: int, names, tensors) -> None: by stall watchdog 300s later: consumers route pulls here, the pull passes the served-names guard, and the cache wait never completes. With it, that is an immediate error naming the weight. One set lookup per - name per sync.""" + name per sync. + """ if self._held_names is None: return for name, tensor in zip(names, tensors): @@ -1330,7 +1346,8 @@ def _run_gather_loop( credit releases when every live consumer has signaled the group). So the loop self-paces to the consumers' pull rate with at most `gather_lookahead + 1` groups resident. Runs on every rank; only the - sender has an `update_future` to fail fast on.""" + sender has an `update_future` to fail fast on. + """ assert self.source is not None # guaranteed by trainer_init # The live IDS ride along with the count: the free barrier needs only # the count, but the slot-sharing rendezvous needs to know WHICH diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 9f91da3418da..3441f8febcde 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -68,7 +68,8 @@ class SparseNCCLTrainerInitInfo(TrainerInitInfo): Same rendezvous shape as the dense NCCL backend (the sender opens its endpoint as NCCL rank 0), but with no packed wire params: sparse transfers - are never packed. `backend` is the factory dispatch key.""" + are never packed. `backend` is the factory dispatch key. + """ backend: ClassVar[str] = "sparse_nccl" @@ -113,8 +114,7 @@ def __post_init__(self) -> None: class SparseNCCLWeightTransferEngine( WeightTransferEngine[NCCLWeightTransferInitInfo, SparseNCCLWeightTransferUpdateInfo] ): - """ - Sparse weight transfer engine using NCCL. + """Sparse weight transfer engine using NCCL. Receives checkpoint-coordinate patches broadcast from the trainer and applies them through the model's native weight loader. Sparse updates modify initialized @@ -368,7 +368,8 @@ def _post_send_sync(self) -> None: rebuild or free the patch tensors as soon as a send method returns rather than relying on same-stream ordering. See `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no - cross-rank barrier.""" + cross-rank barrier. + """ if torch.cuda.is_available(): torch.cuda.current_stream().synchronize() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a977b64649eb..e38292c0afe1 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -866,7 +866,6 @@ def __post_init__(self): @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: """Shared CLI arguments for vLLM engine.""" - # Model arguments model_kwargs = get_kwargs(ModelConfig) model_group = parser.add_argument_group( @@ -2042,8 +2041,7 @@ def create_engine_config( usage_context: UsageContext | None = None, headless: bool = False, ) -> VllmConfig: - """ - Create the VllmConfig. + """Create the VllmConfig. NOTE: If VllmConfig is incompatible, we raise an error. """ diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 63f6970e0561..c44fd44d0fb3 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -39,7 +39,7 @@ class StreamingInput: class EngineClient(ABC): - """Protocol class for Clients to Engine""" + """Protocol class for Clients to Engine.""" vllm_config: VllmConfig model_config: ModelConfig @@ -77,6 +77,7 @@ def check_admission( # noqa: B027 Raises: GracefulHTTPError: If the request cannot be admitted. + """ @abstractmethod @@ -124,6 +125,7 @@ async def abort(self, request_id: str | Iterable[str]) -> None: Args: request_id: The unique id of the request, or an iterable of such ids. + """ ... @@ -149,49 +151,49 @@ async def do_log_stats(self) -> None: ... @abstractmethod async def check_health(self) -> None: - """Raise if unhealthy""" + """Raise if unhealthy.""" ... @abstractmethod async def start_profile(self) -> None: - """Start profiling the engine""" + """Start profiling the engine.""" ... @abstractmethod async def stop_profile(self) -> None: - """Stop profiling the engine""" + """Stop profiling the engine.""" ... @abstractmethod async def reset_mm_cache(self) -> None: - """Reset the multi-modal cache""" + """Reset the multi-modal cache.""" ... @abstractmethod async def reset_encoder_cache(self) -> None: - """Reset the encoder cache""" + """Reset the encoder cache.""" ... @abstractmethod async def reset_prefix_cache( self, reset_running_requests: bool = False, reset_connector: bool = False ) -> bool: - """Reset the prefix cache and optionally any configured connector cache""" + """Reset the prefix cache and optionally any configured connector cache.""" ... @abstractmethod async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None: - """Sleep the engine""" + """Sleep the engine.""" ... @abstractmethod async def wake_up(self, tags: list[str] | None = None) -> None: - """Wake up the engine""" + """Wake up the engine.""" ... @abstractmethod async def is_sleeping(self) -> bool: - """Check whether the engine is sleeping""" + """Check whether the engine is sleeping.""" ... @abstractmethod @@ -219,6 +221,7 @@ async def pause_generation( wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV and prefix caches after draining. + """ ... @@ -240,7 +243,7 @@ def shutdown(self, timeout: float | None = None) -> None: async def scale_elastic_ep( self, new_data_parallel_size: int, drain_timeout: int = 300 ) -> None: - """Scale the engine""" + """Scale the engine.""" raise NotImplementedError async def collective_rpc( @@ -256,7 +259,7 @@ async def collective_rpc( async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest ) -> FaultToleranceResult: - """send fault tolerance instruction to the engine""" + """Send fault tolerance instruction to the engine.""" raise NotImplementedError async def get_status(self): @@ -264,7 +267,7 @@ async def get_status(self): raise NotImplementedError async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: - """Get supported tasks""" + """Get supported tasks.""" raise NotImplementedError async def init_weight_transfer_engine( diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 1ef5dbe1f0c4..c879866dfb2c 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Pydantic models for Anthropic API protocol""" +"""Pydantic models for Anthropic API protocol.""" import time from typing import Annotated, Any, Literal @@ -11,21 +11,21 @@ class AnthropicError(BaseModel): - """Error structure for Anthropic API""" + """Error structure for Anthropic API.""" type: str message: str class AnthropicErrorResponse(BaseModel): - """Error response structure for Anthropic API""" + """Error response structure for Anthropic API.""" type: Literal["error"] = "error" error: AnthropicError class AnthropicUsage(BaseModel): - """Token usage information""" + """Token usage information.""" input_tokens: int output_tokens: int @@ -34,7 +34,7 @@ class AnthropicUsage(BaseModel): class AnthropicContentBlock(BaseModel): - """Content block in message""" + """Content block in message.""" type: Literal[ "text", @@ -65,14 +65,14 @@ class AnthropicContentBlock(BaseModel): class AnthropicMessage(BaseModel): - """Message structure""" + """Message structure.""" role: Literal["user", "assistant", "system"] content: str | list[AnthropicContentBlock] class AnthropicTool(BaseModel): - """Tool definition""" + """Tool definition.""" name: str description: str | None = None @@ -91,7 +91,7 @@ def validate_input_schema(cls, v): class AnthropicToolChoice(BaseModel): - """Tool Choice definition""" + """Tool Choice definition.""" type: Literal["auto", "any", "tool", "none"] name: str | None = None @@ -105,7 +105,7 @@ def validate_name_required_for_tool(self) -> "AnthropicToolChoice": class AnthropicJsonOutputFormat(BaseModel): - """JSON output format configuration""" + """JSON output format configuration.""" json_schema: dict[str, Any] | None = Field(default=None, alias="schema") type: Literal["json_schema"] = "json_schema" @@ -119,7 +119,7 @@ class AnthropicOutputConfig(BaseModel): class AnthropicMessagesRequest(BaseModel): - """Anthropic Messages API request""" + """Anthropic Messages API request.""" model: str messages: list[AnthropicMessage] @@ -192,7 +192,7 @@ def validate_max_tokens(cls, v): class AnthropicDelta(BaseModel): - """Delta for streaming responses""" + """Delta for streaming responses.""" type: ( Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"] @@ -211,7 +211,7 @@ class AnthropicDelta(BaseModel): class AnthropicStreamEvent(BaseModel): - """Streaming event""" + """Streaming event.""" type: Literal[ "message_start", @@ -232,7 +232,7 @@ class AnthropicStreamEvent(BaseModel): class AnthropicMessagesResponse(BaseModel): - """Anthropic Messages API response""" + """Anthropic Messages API response.""" id: str type: Literal["message"] = "message" @@ -265,7 +265,7 @@ class AnthropicContextManagement(BaseModel): class AnthropicCountTokensRequest(BaseModel): - """Anthropic messages.count_tokens request""" + """Anthropic messages.count_tokens request.""" model: str messages: list[AnthropicMessage] @@ -291,7 +291,7 @@ def validate_model(cls, v): class AnthropicCountTokensResponse(BaseModel): - """Anthropic messages.count_tokens response""" + """Anthropic messages.count_tokens response.""" input_tokens: int context_management: AnthropicContextManagement | None = None diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index f8784284b49c..e90d70e289a4 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -3,7 +3,7 @@ # Adapted from # https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py -"""Anthropic Messages API serving handler""" +"""Anthropic Messages API serving handler.""" import json import logging @@ -96,7 +96,7 @@ def wrap_data_with_event(data: str, event: str): class AnthropicServingMessages(OpenAIServingChat): - """Handler for Anthropic Messages API requests""" + """Handler for Anthropic Messages API requests.""" def __init__( self, @@ -196,7 +196,7 @@ def to_chat_completion_request( *, merge_inline_system: bool = False, ) -> ChatCompletionRequest: - """Convert Anthropic message format to OpenAI format""" + """Convert Anthropic message format to OpenAI format.""" openai_messages: list[dict[str, Any]] = [] cls._convert_system_message( @@ -224,7 +224,7 @@ def _convert_system_message( *, merge_inline_system: bool = False, ) -> None: - """Convert Anthropic system message to OpenAI format""" + """Convert Anthropic system message to OpenAI format.""" system_parts: list[str] = [] # Top-level system field @@ -278,7 +278,7 @@ def _convert_messages( *, merge_inline_system: bool = False, ) -> None: - """Convert Anthropic messages to OpenAI format""" + """Convert Anthropic messages to OpenAI format.""" for msg in messages: # Handle system messages in-place: extract text, strip billing # headers, and only emit if there is real content. This avoids @@ -310,7 +310,7 @@ def _convert_message_content( openai_msg: dict[str, Any], openai_messages: list[dict[str, Any]], ) -> None: - """Convert complex message content blocks""" + """Convert complex message content blocks.""" content_parts: list[dict[str, Any]] = [] tool_calls: list[dict[str, Any]] = [] reasoning_parts: list[str] = [] @@ -349,7 +349,7 @@ def _convert_block( reasoning_parts: list[str], openai_messages: list[dict[str, Any]], ) -> None: - """Convert individual content block""" + """Convert individual content block.""" if block.type == "text" and block.text: content_parts.append({"type": "text", "text": block.text}) elif block.type == "image" and block.source: @@ -374,7 +374,7 @@ def _convert_block( @classmethod def _convert_tool_use_block(cls, block, tool_calls: list[dict[str, Any]]) -> None: - """Convert tool_use block to OpenAI function call format""" + """Convert tool_use block to OpenAI function call format.""" tool_call = { "id": block.id or f"call_{int(time.time())}", "type": "function", @@ -393,7 +393,7 @@ def _convert_tool_result_block( openai_messages: list[dict[str, Any]], content_parts: list[dict[str, Any]], ) -> None: - """Convert tool_result block to OpenAI format""" + """Convert tool_result block to OpenAI format.""" if role == "user": cls._convert_user_tool_result(block, openai_messages) else: @@ -406,7 +406,7 @@ def _convert_tool_result_block( def _convert_user_tool_result( cls, block, openai_messages: list[dict[str, Any]] ) -> None: - """Convert user tool_result with text and image support""" + """Convert user tool_result with text and image support.""" tool_text = "" tool_image_urls: list[str] = [] tool_reference: list[dict[str, Any]] = [] @@ -468,7 +468,7 @@ def _build_base_request( anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, openai_messages: list[dict[str, Any]], ) -> ChatCompletionRequest: - """Build base ChatCompletionRequest""" + """Build base ChatCompletionRequest.""" if isinstance(anthropic_request, AnthropicCountTokensRequest): return ChatCompletionRequest( model=anthropic_request.model, @@ -498,7 +498,7 @@ def _handle_output_config( req: ChatCompletionRequest, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, ) -> None: - """Handle output configuration such as output format and effort""" + """Handle output configuration such as output format and effort.""" if isinstance(anthropic_request, AnthropicCountTokensRequest): return output_config: AnthropicOutputConfig | None = anthropic_request.output_config @@ -519,7 +519,7 @@ def _handle_streaming_options( req: ChatCompletionRequest, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, ) -> None: - """Handle streaming configuration""" + """Handle streaming configuration.""" if isinstance(anthropic_request, AnthropicCountTokensRequest): return if anthropic_request.stream: @@ -534,7 +534,7 @@ def _convert_tool_choice( anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, req: ChatCompletionRequest, ) -> None: - """Convert Anthropic tool_choice to OpenAI format""" + """Convert Anthropic tool_choice to OpenAI format.""" if anthropic_request.tool_choice is None: req.tool_choice = None return @@ -563,7 +563,7 @@ def _convert_tools( anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, req: ChatCompletionRequest, ) -> None: - """Convert Anthropic tools to OpenAI format""" + """Convert Anthropic tools to OpenAI format.""" if anthropic_request.tools is None: return @@ -593,8 +593,7 @@ async def create_messages( request: AnthropicMessagesRequest, raw_request: Request | None = None, ) -> AsyncGenerator[str, None] | AnthropicMessagesResponse | ErrorResponse: - """ - Messages API similar to Anthropic's API. + """Messages API similar to Anthropic's API. See https://docs.anthropic.com/en/api/messages for the API specification. This API mimics the Anthropic messages API. diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index addcc5f76df8..e3cc744ba70d 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -223,9 +223,7 @@ class ChatCompletionContentPartVideoParam(TypedDict, total=False): class PILImage(BaseModel): - """ - A PIL.Image.Image object. - """ + """A PIL.Image.Image object.""" image_pil: Image.Image model_config = ConfigDict(arbitrary_types_allowed=True) @@ -238,6 +236,7 @@ class CustomChatCompletionContentPILImageParam(TypedDict, total=False): { "image_pil": ImageAsset('cherry_blossom').pil_image } + """ image_pil: PILImage | None @@ -256,6 +255,7 @@ class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False): { "image_url": "https://example.com/image.jpg" } + """ image_url: str | None @@ -273,6 +273,7 @@ class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False): { "audio_url": "https://example.com/audio.mp3" } + """ audio_url: str | None @@ -285,6 +286,7 @@ class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False): { "video_url": "https://example.com/video.mp4" } + """ video_url: str | None @@ -304,6 +306,7 @@ class CustomThinkCompletionContentParam(TypedDict, total=False): "closed": True, "type": "thinking" } + """ thinking: Required[str] @@ -324,6 +327,7 @@ class CustomChatCompletionContentToolReferenceParam(TypedDict, total=False): "name": "get_weather", "type": "tool_reference" } + """ name: str @@ -586,8 +590,7 @@ def _get_embeds_data( class BaseMultiModalItemTracker(ABC, Generic[_T]): - """ - Tracks multi-modal items in a given request and ensures that the number + """Tracks multi-modal items in a given request and ensures that the number of multi-modal items in a given request does not exceed the configured maximum per prompt. """ @@ -651,8 +654,7 @@ def video_processor_name(self) -> str | None: return get_video_processor_cls_name(self.model_config) def add(self, modality: ModalityStr, item: _T) -> str | None: - """ - Add a multi-modal item to the current prompt and returns the + """Add a multi-modal item to the current prompt and returns the placeholder string to use, if any. An optional uuid can be added which serves as a unique identifier of the @@ -663,6 +665,7 @@ def add(self, modality: ModalityStr, item: _T) -> str | None: pre-computed embeddings that do not go through any HF processor, encoder, or model-specific placeholder logic. The corresponding placeholder string is managed by the parser via `_add_placeholder`, so we return None here. + """ add_info = self._validate_add(modality) if add_info is None: @@ -791,13 +794,13 @@ def _resolve_items( mm_processor: BaseMultiModalProcessor | None, modality_order: dict[str, list[str]], ) -> tuple[MultiModalDataDict, MultiModalUUIDDict]: - """ - Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`. + """Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`. Note: `mm_processor` is `None` for text-only models (no registered HF processor) whose only modality is `prompt_embeds`. Every other modality requires a processor, enforced by the guard below. + """ if "image" in items_by_modality and "image_embeds" in items_by_modality: raise VLLMValidationError( @@ -1568,7 +1571,6 @@ def _get_full_multimodal_text_prompt( multimodal_content_part_separator: str = "\n", ) -> str: """Combine multimodal prompts for a multimodal language model.""" - # flatten storage to make it looks like # { # "<|image|>": 2, @@ -1687,8 +1689,7 @@ def _collect_extra_fields(part: dict[str, Any]) -> dict[str, Any]: def _parse_chat_message_content_mm_part( part: ChatCompletionContentPartParam, ) -> tuple[str, _ContentPart]: - """ - Parses a given multi-modal content part based on its type. + """Parses a given multi-modal content part based on its type. Args: part: A dict containing the content part, with a potential 'type' field. @@ -1700,6 +1701,7 @@ def _parse_chat_message_content_mm_part( Raises: ValueError: If the 'type' field is missing and no direct URL is found. + """ assert isinstance( part, dict diff --git a/vllm/entrypoints/cli/benchmark/base.py b/vllm/entrypoints/cli/benchmark/base.py index 60189ab900e8..717081e04cd1 100644 --- a/vllm/entrypoints/cli/benchmark/base.py +++ b/vllm/entrypoints/cli/benchmark/base.py @@ -22,5 +22,6 @@ def cmd(args: argparse.Namespace) -> None: Args: args: The arguments to the command. + """ raise NotImplementedError diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index d173bc392f1d..21a185b3c19d 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The CLI entrypoints of vLLM +"""The CLI entrypoints of vLLM. Note that all future modules must be lazily loaded within main -to avoid certain eager import breakage.""" +to avoid certain eager import breakage. +""" import importlib.metadata import sys diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index dc41d3f067b2..247e968d43ae 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -179,6 +179,7 @@ def _preflight(self, n: int = 1) -> None: Args: n: Number of sequences the request will occupy. + """ if self.engine_client.errored: raise self.engine_client.dead_error @@ -236,7 +237,7 @@ async def _get_trace_headers( @staticmethod def _get_data_parallel_rank(raw_request: Request | None) -> int | None: - """Pulls the data parallel rank from a header, if provided""" + """Pulls the data parallel rank from a header, if provided.""" if raw_request is None: return None @@ -294,7 +295,8 @@ async def _with_kv_transfer_rejection_cleanup( ) -> _T: """Wrap a `create_*` coroutine so that, if it raises or returns an ErrorResponse (i.e. the request never reached the engine), the KV - connector is notified to free any pinned remote-prefill blocks.""" + connector is notified to free any pinned remote-prefill blocks. + """ kv_transfer_params = self.has_kv_connector and request.kv_transfer_params if not kv_transfer_params or not kv_transfer_params.get("do_remote_prefill"): return await awaitable diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 73b9ebe06b8f..b7095131af31 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -53,7 +53,7 @@ def _bitmask_to_token_ids(bitmask_row: torch.Tensor, vocab_size: int) -> list[in class BeamSearchOfflineMixin(OfflineInferenceMixin): - """Offline inference for beam search""" + """Offline inference for beam search.""" def beam_search( self, @@ -63,8 +63,7 @@ def beam_search( use_tqdm: bool = False, concurrency_limit: int | None = None, ) -> list[BeamSearchOutput]: - """ - Generate sequences using beam search. + """Generate sequences using beam search. Args: prompts: A list of prompts. Each prompt can be a string or a list @@ -74,6 +73,7 @@ def beam_search( use_tqdm: Whether to use tqdm to display the progress bar. concurrency_limit: The maximum number of concurrent requests. If None, the number of concurrent requests is unlimited. + """ # TODO: how does beam search work together with length penalty, # frequency, penalty, and stopping criteria, etc.? diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index da02b974c0a5..ac7f0f884f5b 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -20,7 +20,7 @@ class BeamSearchOnlineMixin(ABC): - """online serving for beam search""" + """online serving for beam search.""" renderer: BaseRenderer engine_client: EngineClient diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 753054d1b8a0..cb0c4995bf75 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -60,6 +60,7 @@ class GenerativeScoringRequest(OpenAIBaseModel): the full vocab for those ids (False). item_first: If True, prepend items to query. Otherwise append items to query. add_special_tokens: Whether to add special tokens when tokenizing. + """ model: str | None = None @@ -112,6 +113,7 @@ class GenerativeScoringItemResult(OpenAIBaseModel): index: The index of this item in the input items list. object: Type of object, always "score". score: The probability score for the first label token. + """ index: int @@ -129,6 +131,7 @@ class GenerativeScoringResponse(OpenAIBaseModel): model: The model used for scoring. data: List of scoring results, one per input item. usage: Token usage information. + """ id: str = Field(default="") @@ -188,6 +191,7 @@ async def create_generative_scoring( Returns: GenerativeScoringResponse with probabilities for each item, or ErrorResponse if an error occurred. + """ # Check model error_check_ret = await self._check_model(request) # type: ignore[arg-type] @@ -401,6 +405,7 @@ async def _build_prompts( Returns: Tuple of (list of EngineInput, list of prompt token counts). + """ # Tokenize query if it's a string if isinstance(request.query, str): @@ -455,6 +460,7 @@ def _compute_probabilities( Returns: Dictionary mapping token_id to probability. + """ if apply_softmax: # Normalize over only the label tokens (subset softmax) diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py index 37d7ca381455..f3dcf40362a1 100644 --- a/vllm/entrypoints/launchers/api_server/entry.py +++ b/vllm/entrypoints/launchers/api_server/entry.py @@ -70,14 +70,12 @@ async def build_async_engine_client_from_engine_args( usage_context: UsageContext = UsageContext.OPENAI_API_SERVER, client_config: dict[str, Any] | None = None, ) -> AsyncIterator[EngineClient]: - """ - Create EngineClient, either: + """Create EngineClient, either: - in-process using the AsyncLLMEngine Directly - multiprocess using AsyncLLMEngine RPC Returns the Client or None if the creation failed. """ - # Create the EngineConfig (determines if we can use V1). vllm_config = engine_args.create_engine_config(usage_context=usage_context) @@ -123,7 +121,6 @@ async def build_and_serve( Returns the shutdown task for the caller to await. """ - # Get uvicorn log config (from file or with endpoint filter) log_config = get_uvicorn_log_config(args) if log_config is not None: @@ -162,7 +159,6 @@ async def build_and_serve( async def run_server(args, **uvicorn_kwargs) -> None: """Run a single-worker API server.""" - decorate_logs("APIServer", skip_if_decorated=True) # Interrupt initialization if SIGTERM arrives before uvicorn installs its @@ -180,7 +176,6 @@ async def run_server_worker( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: """Run a single API server worker.""" - if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: ToolParserManager.import_tool_parser(args.tool_parser_plugin) diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 89af22cdc80f..71ae94c10635 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file contains the command line arguments for the vLLM's online server. +"""This file contains the command line arguments for the vLLM's online server. It is kept in a separate file for documentation purposes. """ diff --git a/vllm/entrypoints/launchers/dp_supervisor.py b/vllm/entrypoints/launchers/dp_supervisor.py index 7ebf887ccef1..91754b71e56b 100644 --- a/vllm/entrypoints/launchers/dp_supervisor.py +++ b/vllm/entrypoints/launchers/dp_supervisor.py @@ -182,8 +182,7 @@ async def _probe_endpoint( conn_err_failure_threshold: int = 3, conn_err_retry_delay: float = 5.0, ) -> bool: - """ - Probe /health endpoint for 200 status. + """Probe /health endpoint for 200 status. If there is a connection error, retry every N seconds. """ @@ -251,9 +250,7 @@ def _run_rust_vllm_dp_server(child_args: argparse.Namespace) -> None: def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: - """ - Entrypoint function for the vLLM DP Server. - """ + """Entrypoint function for the vLLM DP Server.""" # Create a fresh process group for the vLLM DP Server, # so that CTRL-C is propagated cleanly. os.setpgrp() @@ -320,8 +317,7 @@ async def run(self) -> None: await supervisor_server_task async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]: - """ - Launch the DPSupervisor HTTP server. + """Launch the DPSupervisor HTTP server. Called only after the vLLM DP Servers are ready so that /health does not return 503 to external probes while the engines are initializing. @@ -366,8 +362,7 @@ async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]: return supervisor_server, supervisor_server_task async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None: - """ - Block until the vLLM DP Servers are ready or shutdown is triggered. + """Block until the vLLM DP Servers are ready or shutdown is triggered. Returns early if monitoring stops (e.g. a DP Server dies during startup), in which case the supervisor server is never started. @@ -379,14 +374,12 @@ async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None: await asyncio.sleep(0.05) def _handle_signal(self, signum: int) -> None: - """ - Signal handler that is added to the event loop. + """Signal handler that is added to the event loop. This catches the SIGTERM from K8s and begins graceful shutdown, by setting the _shutdown_event(), which is watched by the main coroutine monitoring the vLLM DP Servers. """ - if self._shutdown_event.is_set(): return @@ -400,9 +393,7 @@ def _handle_signal(self, signum: int) -> None: self._is_ready = False def _start_children(self) -> None: - """ - Launch vLLM DP Servers on separate GPUs. - """ + """Launch vLLM DP Servers on separate GPUs.""" logger.info("Launching vLLM DP Servers") context = multiprocessing.get_context("spawn") for local_rank in range(self.args.data_parallel_size_local): @@ -416,8 +407,7 @@ def _start_children(self) -> None: self._processes.append(process) async def _probe_all_children(self) -> None: - """ - Background coroutine: probes all child endpoints on each interval. + """Background coroutine: probes all child endpoints on each interval. Exits when any server becomes unhealthy after being ready, signalling _monitor_children to initiate shutdown. @@ -474,8 +464,7 @@ async def _probe_all_children(self) -> None: ) async def _monitor_children(self) -> None: - """ - Main coroutine task that monitors the children vLLM servers. + """Main coroutine task that monitors the children vLLM servers. Before the vLLM servers are /ready: - if the pid is dead, we will shut down diff --git a/vllm/entrypoints/launchers/grpc_server.py b/vllm/entrypoints/launchers/grpc_server.py index 3dc51f2e771d..aab3f57c5ea7 100644 --- a/vllm/entrypoints/launchers/grpc_server.py +++ b/vllm/entrypoints/launchers/grpc_server.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # mypy: ignore-errors -""" -vLLM gRPC Server +"""vLLM gRPC Server Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer from the smg-grpc-servicer package. @@ -16,6 +15,7 @@ --model meta-llama/Llama-2-7b-hf \ --host 0.0.0.0 \ --port 50051 + """ import argparse @@ -54,11 +54,11 @@ async def serve_grpc(args: argparse.Namespace): - """ - Main gRPC serving function. + """Main gRPC serving function. Args: args: Parsed command line arguments + """ log_version_and_model(logger, VLLM_VERSION, args.model) logger.info("vLLM gRPC server args: %s", args) diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 09b0ea97811e..9fabfa3277af 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -53,8 +53,7 @@ async def serve_http( enable_ssl_refresh: bool = False, **uvicorn_kwargs: Any, ): - """ - Start a FastAPI app using Uvicorn, with support for custom Uvicorn config + """Start a FastAPI app using Uvicorn, with support for custom Uvicorn config options. Supports http header limits via h11_max_incomplete_event_size and h11_max_header_count. """ @@ -185,8 +184,7 @@ async def handle_shutdown() -> None: async def watchdog_loop(server: uvicorn.Server, engine: EngineClient): - """ - # Watchdog task that runs in the background, checking + """# Watchdog task that runs in the background, checking # for error state in the engine. Needed to trigger shutdown # if an exception arises is StreamingResponse() generator. """ @@ -197,8 +195,7 @@ async def watchdog_loop(server: uvicorn.Server, engine: EngineClient): def terminate_if_errored(server: uvicorn.Server, engine: EngineClient): - """ - See discussions here on shutting down a uvicorn server + """See discussions here on shutting down a uvicorn server https://github.com/encode/uvicorn/discussions/1103 In this case we cannot await the server shutdown here because handler must first return to close the connection @@ -254,7 +251,6 @@ def validate_api_server_args(args): @instrument(span_name="API server setup") def setup_server(args, *, reuse_port: bool): """Validate API server args and create the server socket.""" - log_version_and_model(logger, VLLM_VERSION, args.model) log_non_default_args(args) diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py index 6fb4123c9b06..33988443d1d9 100644 --- a/vllm/entrypoints/launchers/render/app_state.py +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -31,7 +31,6 @@ async def init_render_app_state( preprocessing pipeline (renderer, input_processor) directly from the :class:`~vllm.config.VllmConfig`. """ - served_model_names = args.served_model_name or [args.model] model_registry = OpenAIModelRegistry( model_config=vllm_config.model_config, diff --git a/vllm/entrypoints/launchers/render/entry.py b/vllm/entrypoints/launchers/render/entry.py index c7465c1a12d3..b36cd6477214 100644 --- a/vllm/entrypoints/launchers/render/entry.py +++ b/vllm/entrypoints/launchers/render/entry.py @@ -30,7 +30,6 @@ async def build_and_serve_renderer( Returns the shutdown task for the caller to await. """ - # Get uvicorn log config (from file or with endpoint filter) log_config = get_uvicorn_log_config(args) if log_config is not None: diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py index 8d1be95b7007..8c05a1201f9c 100644 --- a/vllm/entrypoints/launchers/run_batch.py +++ b/vllm/entrypoints/launchers/run_batch.py @@ -88,8 +88,7 @@ def _raise_audio_filesize_limit(size_bytes: int) -> NoReturn: class BatchTranscriptionRequest(TranscriptionRequest): - """ - Batch transcription request that uses file_url instead of file. + """Batch transcription request that uses file_url instead of file. This class extends TranscriptionRequest but replaces the file field with file_url to support batch processing from audio files written in JSON format. @@ -119,8 +118,7 @@ def validate_no_file(cls, data: Any): class BatchTranslationRequest(TranslationRequest): - """ - Batch translation request that uses file_url instead of file. + """Batch translation request that uses file_url instead of file. This class extends TranslationRequest but replaces the file field with file_url to support batch processing from audio files written in JSON format. @@ -160,8 +158,7 @@ def validate_no_file(cls, data: Any): class BatchRequestInput(OpenAIBaseModel): - """ - The per-line object of the batch input file. + """The per-line object of the batch input file. NOTE: Currently only the `/v1/chat/completions` endpoint is supported. """ @@ -225,9 +222,7 @@ class BatchResponseData(OpenAIBaseModel): class BatchRequestOutput(OpenAIBaseModel): - """ - The per-line object of the batch output and error files - """ + """The per-line object of the batch output and error files""" id: str @@ -358,8 +353,7 @@ async def read_file(path_or_url: str) -> str: async def write_local_file( output_path: str, batch_outputs: list[BatchRequestOutput] ) -> None: - """ - Write the responses to a local file. + """Write the responses to a local file. output_path: The path to write the responses to. batch_outputs: The list of batch outputs to write. """ @@ -371,8 +365,7 @@ async def write_local_file( async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> None: - """ - Upload a local file to a URL. + """Upload a local file to a URL. output_url: The URL to upload the file to. data_or_file: Either the data to upload or the path to the file to upload. from_file: If True, data_or_file is the path to the file to upload. @@ -429,8 +422,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No async def write_file( path_or_url: str, batch_outputs: list[BatchRequestOutput], output_tmp_dir: str ) -> None: - """ - Write batch_outputs to a file or upload to a URL. + """Write batch_outputs to a file or upload to a URL. path_or_url: The path or URL to write batch_outputs to. batch_outputs: The list of batch outputs to write. output_tmp_dir: The directory to store the output file before uploading it @@ -471,8 +463,7 @@ async def download_bytes_from_url( url: str, allowed_media_domains: list[str] | None = None, ) -> bytes: - """ - Download data from a URL or decode from a data URL. + """Download data from a URL or decode from a data URL. Args: url: Either an HTTP/HTTPS URL or a data URL (data:...;base64,...) @@ -482,6 +473,7 @@ async def download_bytes_from_url( Returns: Data as bytes + """ parsed = urlparse(url) @@ -615,8 +607,7 @@ def handle_endpoint_request( handler_getter: Callable[[], Callable | None], wrapper_fn: WrapperFn | None = None, ) -> Awaitable[BatchRequestOutput] | None: - """ - Generic handler for endpoint requests. + """Generic handler for endpoint requests. Args: request: The batch request input @@ -628,6 +619,7 @@ def handle_endpoint_request( Returns: Awaitable[BatchRequestOutput] if the request was handled, None if URL didn't match + """ if not url_matcher(request.url): return None @@ -649,8 +641,7 @@ def make_transcription_wrapper( is_translation: bool, allowed_media_domains: list[str] | None = None, ) -> WrapperFn: - """ - Factory function to create a wrapper for transcription/translation handlers. + """Factory function to create a wrapper for transcription/translation handlers. The wrapper converts BatchTranscriptionRequest or BatchTranslationRequest to TranscriptionRequest or TranslationRequest and calls the appropriate handler. @@ -662,6 +653,7 @@ def make_transcription_wrapper( Returns: A function that takes a handler and returns a wrapped handler + """ def wrapper(handler_fn: Callable): @@ -723,8 +715,7 @@ async def build_endpoint_registry( engine_client: EngineClient, args: Namespace, ) -> dict[str, dict[str, Any]]: - """ - Build the endpoint registry with all serving objects and handler configurations. + """Build the endpoint registry with all serving objects and handler configurations. Args: engine_client: The engine client @@ -732,6 +723,7 @@ async def build_endpoint_registry( Returns: Dictionary mapping endpoint keys to their configurations + """ supported_tasks = await engine_client.get_supported_tasks() logger.info("Supported tasks: %s", supported_tasks) diff --git a/vllm/entrypoints/launchers/utils/server_utils.py b/vllm/entrypoints/launchers/utils/server_utils.py index 082582672381..91f5e620890b 100644 --- a/vllm/entrypoints/launchers/utils/server_utils.py +++ b/vllm/entrypoints/launchers/utils/server_utils.py @@ -29,8 +29,7 @@ def load_log_config(log_config_file: str | None) -> dict | None: def get_uvicorn_log_config(args: Namespace) -> dict | None: - """ - Get the uvicorn log config based on the provided arguments. + """Get the uvicorn log config based on the provided arguments. Priority: 1. If log_config_file is specified, use it diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index e3d0b2f1e520..d9a2a46701c1 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -174,6 +174,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): Note: This class is intended to be used for offline inference. For online serving, use the [AsyncLLMEngine][vllm.AsyncLLMEngine] class instead. + """ def __init__( @@ -224,7 +225,6 @@ def __init__( **kwargs: Any, ) -> None: """LLM constructor.""" - if "disable_log_stats" not in kwargs: kwargs["disable_log_stats"] = True @@ -398,6 +398,7 @@ def get_world_size(self, include_dp: bool = True) -> int: Returns: The world size (tensor_parallel_size * pipeline_parallel_size), optionally multiplied by data_parallel_size if include_dp is True. + """ parallel_config = self.llm_engine.vllm_config.parallel_config if include_dp: @@ -457,6 +458,7 @@ def generate( Returns: A list of `RequestOutput` objects containing the generated completions in the same order as the input prompts. + """ runner_type = self.model_config.runner_type if runner_type != "generate": @@ -507,6 +509,7 @@ def enqueue( Returns: A list of request IDs for the enqueued requests. + """ runner_type = self.model_config.runner_type if runner_type != "generate": @@ -558,6 +561,7 @@ def wait_for_completion( Returns: A list of output objects for all completed requests. + """ if output_type is None: output_type = (RequestOutput, PoolingRequestOutput) @@ -571,8 +575,7 @@ def collective_rpc( args: tuple = (), kwargs: dict[str, Any] | None = None, ) -> list[_R]: - """ - Execute an RPC call on all workers. + """Execute an RPC call on all workers. Args: method: Name of the worker method to execute, or a callable that @@ -592,13 +595,12 @@ def collective_rpc( Note: It is recommended to use this API to only pass control messages, and set up data-plane communication to pass data. - """ + """ return self.llm_engine.collective_rpc(method, timeout, args, kwargs) def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]: - """ - Run a function directly on the model inside each worker, + """Run a function directly on the model inside each worker, returning the result for each of them. !!! warning @@ -625,8 +627,7 @@ def chat( tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[RequestOutput]: - """ - Generate responses for a chat conversation. + """Generate responses for a chat conversation. The chat conversation is converted into a text prompt using the tokenizer and calls the [generate][vllm.LLM.generate] method to generate @@ -670,10 +671,12 @@ def chat( template. tokenization_kwargs: Overrides for `tokenizer.encode`. mm_processor_kwargs: Overrides for `processor.__call__`. + tools: Tools to make available to the model, if any. Returns: A list of `RequestOutput` objects containing the generated responses in the same order as the input messages. + """ model_config = self.model_config runner_type = model_config.runner_type @@ -751,6 +754,7 @@ def enqueue_chat( Returns: A list of request IDs for the enqueued requests. + """ model_config = self.model_config runner_type = model_config.runner_type @@ -787,6 +791,7 @@ def start_profile(self, profile_prefix: str | None = None) -> None: profile_prefix: Optional prefix for the trace file names. If provided, trace files will be named as "_dp_pp_tp". If not provided, default naming will be used. + """ self.llm_engine.start_profile(profile_prefix) @@ -801,8 +806,7 @@ def reset_prefix_cache( ) def sleep(self, level: int = 1, mode: PauseMode = "abort"): - """ - Put the engine to sleep. The engine should not process any requests. + """Put the engine to sleep. The engine should not process any requests. The caller should guarantee that no requests are being processed during the sleep period, before `wake_up` is called. @@ -822,12 +826,12 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"): CPU memory pressure. mode: How to handle any existing requests, can be "abort", "wait", or "keep". + """ self.llm_engine.sleep(level=level, mode=mode) def wake_up(self, tags: list[str] | None = None): - """ - Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep] + """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep] method for more details. Args: @@ -837,6 +841,7 @@ def wake_up(self, tags: list[str] | None = None): is reallocated. wake_up should be called with all tags (or None) before the engine is used again. Use tags=["scheduling"] to resume from level 0 sleep. + """ self.llm_engine.wake_up(tags) @@ -849,17 +854,18 @@ def get_metrics(self) -> list["Metric"]: Note: This method is only available with the V1 LLM engine. + """ return self.llm_engine.get_metrics() def init_weight_transfer_engine( self, request: WeightTransferInitRequest | dict ) -> None: - """ - Initialize weight transfer for RL training. + """Initialize weight transfer for RL training. Args: request: Weight transfer initialization request with backend-specific info + """ init_info_dict = ( request["init_info"] if isinstance(request, dict) else request.init_info @@ -878,11 +884,11 @@ def start_draft_weight_update(self) -> None: self.llm_engine.collective_rpc("start_draft_weight_update") def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: - """ - Update the weights of the model. + """Update the weights of the model. Args: request: Weight update request with backend-specific update info + """ update_info_dict = ( request["update_info"] if isinstance(request, dict) else request.update_info diff --git a/vllm/entrypoints/mcp/tool.py b/vllm/entrypoints/mcp/tool.py index cd25aef087f8..d1f9130bdf22 100644 --- a/vllm/entrypoints/mcp/tool.py +++ b/vllm/entrypoints/mcp/tool.py @@ -23,8 +23,7 @@ def validate_gpt_oss_install(): - """ - Check if the gpt-oss is installed and its version is at least 0.0.7. + """Check if the gpt-oss is installed and its version is at least 0.0.7. If not, raise an ImportError. """ from importlib.metadata import PackageNotFoundError, version @@ -151,8 +150,7 @@ async def get_result(self, context: "ConversationContext") -> Any: return tool_output_msgs async def get_result_parsable_context(self, context: "ConversationContext") -> Any: - """ - This function converts parsable context types to harmony and + """This function converts parsable context types to harmony and back so we can use GPTOSS demo python tool """ from vllm.entrypoints.openai.responses.context import ParsableContext diff --git a/vllm/entrypoints/mcp/tool_server.py b/vllm/entrypoints/mcp/tool_server.py index ed519ce4a51d..a0fe518f2d81 100644 --- a/vllm/entrypoints/mcp/tool_server.py +++ b/vllm/entrypoints/mcp/tool_server.py @@ -76,17 +76,14 @@ def post_process_tools_description( class ToolServer(ABC): @abstractmethod def has_tool(self, tool_name: str) -> bool: - """ - Return True if the tool is supported, False otherwise. - """ + """Return True if the tool is supported, False otherwise.""" pass @abstractmethod def get_tool_description( self, tool_name: str, allowed_tools: list[str] | None = None ) -> ToolNamespaceConfig | None: - """ - Return the tool description for the given tool name. + """Return the tool description for the given tool name. If the tool is not supported, return None. """ pass @@ -95,9 +92,7 @@ def get_tool_description( def new_session( self, tool_name: str, session_id: str, headers: dict[str, str] | None = None ) -> AbstractAsyncContextManager[Any]: - """ - Create a session for the tool. - """ + """Create a session for the tool.""" ... diff --git a/vllm/entrypoints/offline_utils.py b/vllm/entrypoints/offline_utils.py index 650860c3377b..3f75c3033dd6 100644 --- a/vllm/entrypoints/offline_utils.py +++ b/vllm/entrypoints/offline_utils.py @@ -48,7 +48,7 @@ class OfflineInferenceMixin: - """Offline inference utils""" + """Offline inference utils.""" request_counter: Counter renderer: BaseRenderer @@ -113,14 +113,14 @@ def _preprocess_cmpl( tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> Sequence[EngineInput]: - """ - Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into + """Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into a format that can be passed to `_add_request`. Refer to [LLM.generate][] for a complete description of the arguments. Returns: A list of `EngineInput` objects ready to be passed into LLMEngine. + """ renderer = self.renderer model_config = self.model_config @@ -168,14 +168,14 @@ def _preprocess_chat( tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> Sequence[EngineInput]: - """ - Convert a list of conversations into prompts so that they can then + """Convert a list of conversations into prompts so that they can then be used as input for other LLM APIs. Refer to [LLM.chat][] for a complete description of the arguments. Returns: A list of `EngineInput` objects ready to be passed into LLMEngine. + """ renderer = self.renderer diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index 06b3dbfd4a4c..ea6601bff13b 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -50,6 +50,7 @@ async def render_batch_chat_request( Returns: A tuple of (all_conversations, engine_prompts) on success — one entry per conversation — or an ErrorResponse on failure. + """ error_check_ret = await self._check_model(request) if error_check_ret is not None: diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 06e7315c2498..51c0af8e8f15 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -220,8 +220,7 @@ async def render_chat_request( self, request: ChatCompletionRequest, ) -> tuple[list[ConversationMessage], list[EngineInput]] | ErrorResponse: - """ - Validate the model and preprocess a chat completion request. + """Validate the model and preprocess a chat completion request. Delegates preprocessing logic to OnlineRenderer, adding the engine-aware checks (LoRA model validation, engine health). @@ -229,6 +228,7 @@ async def render_chat_request( Returns: A tuple of (conversation, engine_inputs) on success, or an ErrorResponse on failure. + """ error_check_ret = await self._check_model(request) if error_check_ret is not None: @@ -244,8 +244,7 @@ async def create_chat_completion( request: ChatCompletionRequest, raw_request: Request | None = None, ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse: - """ - Chat Completion API similar to OpenAI's API. + """Chat Completion API similar to OpenAI's API. See https://platform.openai.com/docs/api-reference/chat/create for the API specification. This API mimics the OpenAI diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 2fa01478cee6..4d892d89d5f9 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -89,14 +89,14 @@ async def render_completion_request( self, request: CompletionRequest, ) -> list[EngineInput] | ErrorResponse: - """ - Validate the model and preprocess a completion request. + """Validate the model and preprocess a completion request. Delegates preprocessing logic to OnlineRenderer, adding the engine-aware checks (LoRA model validation, engine health). Returns: A list of engine_inputs on success, or an ErrorResponse on failure. + """ error_check_ret = await self._check_model(request) if error_check_ret is not None: diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index 7378d37c6ecf..dc78cc396326 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -123,7 +123,8 @@ def __init__( async def init_static_loras(self): """Loads all static LoRA modules. - Raises if any fail to load""" + Raises if any fail to load + """ if self.static_lora_modules is None: return for lora in self.static_lora_modules: @@ -148,7 +149,8 @@ def model_name(self, lora_request: LoRARequest | None = None) -> str: async def show_available_models(self) -> ModelList: """Show available models. This includes the base model and all - adapters.""" + adapters. + """ model_list = await self.registry.show_available_models() lora_cards = [ ModelCard( @@ -289,6 +291,7 @@ async def resolve_lora(self, lora_name: str) -> LoRARequest | ErrorResponse: LoRARequest if found and loaded successfully. ErrorResponse (404) if no resolver finds the adapter. ErrorResponse (400) if adapter(s) are found but none load. + """ async with self.lora_resolver_lock[lora_name]: # First check if this LoRA is already loaded diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index fca02a91cc49..873b9e92bdab 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -88,8 +88,7 @@ def extract_function_from_recipient(recipient: str) -> str: def has_custom_tools(tool_types: set[str]) -> bool: - """ - Checks if the given tool types are custom tools + """Checks if the given tool types are custom tools (i.e. any tool other than MCP builtin tools) """ return not tool_types.issubset(MCP_BUILTIN_TOOLS) @@ -209,8 +208,7 @@ def get_system_or_developer_message(role: str, instructions: str) -> Message: def parse_chat_inputs_to_harmony_messages(chat_msgs: list) -> list[Message]: - """ - Parse a list of messages from request.messages in the Chat Completion API to + """Parse a list of messages from request.messages in the Chat Completion API to Harmony messages. """ msgs: list[Message] = [] @@ -230,9 +228,8 @@ def parse_chat_inputs_to_harmony_messages(chat_msgs: list) -> list[Message]: def auto_drop_analysis_messages(msgs: list[Message]) -> list[Message]: - """ - Harmony models expect the analysis messages (representing raw chain of thought) to - be dropped after an assistant message to the final channel is produced from the + """Harmony models expect the analysis messages (representing raw chain of thought) + to be dropped after an assistant message to the final channel is produced from the reasoning of those messages. The openai-harmony library does this if the very last assistant message is to the @@ -262,8 +259,7 @@ def auto_drop_analysis_messages(msgs: list[Message]) -> list[Message]: def flatten_input_text_content(content: Any) -> str | None: - """ - Extract text parts from a Chat Completion or Responses API content field and + """Extract text parts from a Chat Completion or Responses API content field and flatten them into a single string. Returns None if no text content is found. """ if content is None or isinstance(content, str): @@ -286,8 +282,7 @@ def flatten_input_text_content(content: Any) -> str | None: def extract_instructions_from_messages( messages: Sequence[Any], ) -> tuple[str | None, list[Any]]: - """ - Peel a leading system/developer Chat Completion or Responses message and + """Peel a leading system/developer Chat Completion or Responses message and flatten its instruction text. """ remaining_messages = list(messages) @@ -327,9 +322,7 @@ def build_harmony_preamble( container_description: str | None = None, with_custom_tools: bool = False, ) -> list[Message]: - """ - Build the standard Harmony system/developer prefix for a request. - """ + """Build the standard Harmony system/developer prefix for a request.""" developer_instructions = system_instructions = None if envs.VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: system_instructions = instructions @@ -359,8 +352,7 @@ def build_harmony_preamble( def parse_chat_input_to_harmony_message( chat_msg, tool_id_names: dict[str, str] | None = None ) -> list[Message]: - """ - Parse a message from request.messages in the Chat Completion API to + """Parse a message from request.messages in the Chat Completion API to Harmony messages. """ tool_id_names = tool_id_names or {} diff --git a/vllm/entrypoints/openai/responses/api_router.py b/vllm/entrypoints/openai/responses/api_router.py index 08b2b5584ec4..ee352dd1933c 100644 --- a/vllm/entrypoints/openai/responses/api_router.py +++ b/vllm/entrypoints/openai/responses/api_router.py @@ -34,7 +34,7 @@ def responses(request: Request) -> OpenAIServingResponses | None: async def _convert_stream_to_sse_events( generator: AsyncGenerator[StreamingResponsesResponse, None], ) -> AsyncGenerator[str, None]: - """Convert the generator to a stream of events in SSE format""" + """Convert the generator to a stream of events in SSE format.""" async for event in generator: event_type = getattr(event, "type", "unknown") # https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 6b93d39b8803..167f006ba0e7 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -141,9 +141,7 @@ async def cleanup_session(self) -> None: def _create_json_parse_error_messages( last_msg: Message, e: json.JSONDecodeError ) -> list[Message]: - """ - Creates an error message when json parse failed. - """ + """Creates an error message when json parse failed.""" error_msg = ( f"Error parsing tool arguments as JSON: {str(e)}. " "Please ensure the tool call arguments are valid JSON and try again." @@ -161,7 +159,7 @@ def _create_json_parse_error_messages( class SimpleContext(ConversationContext): - """This is a context that cannot handle MCP tool calls""" + """This is a context that cannot handle MCP tool calls.""" def __init__( self, @@ -410,7 +408,8 @@ def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None: def need_builtin_tool_call(self) -> bool: """Return true if the last message is a builtin tool call - that the request has enabled.""" + that the request has enabled. + """ last_message = self.response_messages[-1] if last_message.type != "function_call": return False @@ -474,8 +473,7 @@ async def call_search_tool( async def call_container_tool( self, tool_session: Union["ClientSession", Tool], last_msg: Message ) -> list[Message]: - """ - Call container tool. Expect this to be run in a stateful docker + """Call container tool. Expect this to be run in a stateful docker with command line terminal. The official container tool would at least expect the following format: @@ -583,7 +581,7 @@ async def init_tool_sessions( exit_stack.push_async_exit(self.cleanup_session) async def cleanup_session(self, *args, **kwargs) -> None: - """Can be used as coro to used in __aexit__""" + """Can be used as coro to used in __aexit__.""" async def cleanup_tool_session(tool_session): if not isinstance(tool_session, Tool): @@ -693,6 +691,7 @@ def _update_prefill_token_usage(self, output: RequestOutput) -> None: Args: output: The RequestOutput containing prompt token information + """ if output.prompt_token_ids is not None: this_turn_input_tokens = len(output.prompt_token_ids) @@ -757,6 +756,7 @@ def _update_decode_token_usage(self, output: RequestOutput) -> int: Returns: int: Number of output tokens processed in this call + """ updated_output_token_count = 0 if output.outputs: @@ -884,8 +884,7 @@ async def init_tool_sessions( async def call_container_tool( self, tool_session: Union["ClientSession", Tool], last_msg: Message ) -> list[Message]: - """ - Call container tool. Expect this to be run in a stateful docker + """Call container tool. Expect this to be run in a stateful docker with command line terminal. The official container tool would at least expect the following format: @@ -925,7 +924,7 @@ async def call_container_tool( ] async def cleanup_session(self, *args, **kwargs) -> None: - """Can be used as coro to used in __aexit__""" + """Can be used as coro to used in __aexit__.""" async def cleanup_tool_session(tool_session): if not isinstance(tool_session, Tool): diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py index 461ce1f3fa35..c1f3f1b716ac 100644 --- a/vllm/entrypoints/openai/responses/harmony.py +++ b/vllm/entrypoints/openai/responses/harmony.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Harmony ↔ Responses API conversion utilities. +"""Harmony ↔ Responses API conversion utilities. Handles two directions: 1. Response Input → Harmony Messages (input parsing) @@ -53,7 +52,8 @@ def _parse_harmony_format_message(chat_msg: dict) -> Message: """Reconstruct a Message from Harmony-format dict, - preserving channel, recipient, and content_type.""" + preserving channel, recipient, and content_type. + """ author_dict = chat_msg["author"] role = author_dict.get("role") name = author_dict.get("name") diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 27c60967e5b6..f826c932a4ba 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -100,9 +100,7 @@ class ResponseUsage(OpenAIBaseModel): def serialize_message(msg): - """ - Serializes a single message - """ + """Serializes a single message.""" if isinstance(msg, dict): return msg elif hasattr(msg, "to_dict"): @@ -113,15 +111,14 @@ def serialize_message(msg): def serialize_messages(msgs): - """ - Serializes multiple messages - """ + """Serializes multiple messages.""" return [serialize_message(msg) for msg in msgs] if msgs else None class ResponseRawMessageAndToken(OpenAIBaseModel): """Class to show the raw message. - If message / tokens diverge, tokens is the source of truth""" + If message / tokens diverge, tokens is the source of truth + """ message: str tokens: list[int] diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 40abc45b4261..4651bb7205fd 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Streaming SSE event builders for the Responses API. +"""Streaming SSE event builders for the Responses API. Pure functions that translate streaming state + delta data into OpenAI Response API SSE events. Used by the streaming event @@ -125,8 +124,7 @@ def is_mcp_tool_by_namespace( recipient: str | None, allowed_function_tool_names: frozenset[str] | None = None, ) -> bool: - """ - Determine if a tool call is an MCP tool based on recipient prefix. + """Determine if a tool call is an MCP tool based on recipient prefix. Inverse of :func:`is_function_recipient` — everything that is not a function call is an MCP tool. @@ -1145,8 +1143,7 @@ def split_delta(delta: DeltaMessage) -> list[DeltaMessage]: class SimpleStreamingEventProcessor: - """ - State-machine processor for the simple (non-Harmony) streaming path. + """State-machine processor for the simple (non-Harmony) streaming path. Core flow: 1. Resolve the target state from the delta_message @@ -1188,8 +1185,7 @@ def __init__( def resolve_target_state( self, delta_message: DeltaMessage ) -> tuple[_StateType, Any]: - """ - Decide which state the next delta belongs to. + """Decide which state the next delta belongs to. Priority: TOOL_CALL > REASONING > CONTENT, fallback to NONE. For TOOL_CALL the first tool_call object is also returned so @@ -1207,8 +1203,7 @@ def resolve_target_state( return _StateType.NONE, None def needs_transition(self, target_state: _StateType, tool_call: Any) -> bool: - """ - Return True when we must close the current state and open a new one. + """Return True when we must close the current state and open a new one. Two cases trigger a transition: 1. The target state differs from the current state diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 82ab75f756b1..3caff4b056b1 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -115,8 +115,7 @@ def build_response_output_items( def should_continue_final_message( request_input: str | list[ResponseInputOutputItem], ) -> bool: - """ - Determine if the last input message is a partial assistant message + """Determine if the last input message is a partial assistant message that should be continued rather than starting a new generation. This enables partial message completion similar to Anthropic's Messages API, @@ -132,6 +131,7 @@ def should_continue_final_message( Returns: True if the final message should be continued, False otherwise + """ if isinstance(request_input, str): # Simple string input is always a user message @@ -226,8 +226,7 @@ def _construct_message_from_response_item( item: ResponseInputOutputItem, prev_msg: ChatCompletionMessageParam | None = None, ) -> ChatCompletionMessageParam | None: - """ - Returns a new message or None. If `None`, `prev_msg` might be updated. + """Returns a new message or None. If `None`, `prev_msg` might be updated. If `prev_msg` is `None`, a new message is always returned. """ prev_assistant_msg = ( @@ -360,9 +359,7 @@ def extract_function_tool_names(tools: list[Tool]) -> frozenset[str]: def extract_tool_types(tools: list[Tool]) -> set[str]: - """ - Extracts the tool types from the given tools. - """ + """Extracts the tool types from the given tools.""" tool_types: set[str] = set() for tool in tools: if tool.type == "mcp": @@ -379,8 +376,7 @@ def extract_tool_types(tools: list[Tool]) -> set[str]: def convert_tool_responses_to_completions_format( tool: dict, ) -> ChatCompletionToolsParam: - """ - Convert a flat Responses tool schema: + """Convert a flat Responses tool schema: {"type": "function", "name": "...", "description": "...", "parameters": {...}} into a Chat Completions tool param for chat-template rendering. """ diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index f9def808de3a..0a0304e7a4c5 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -37,7 +37,7 @@ class PoolingOfflineMixin(OfflineInferenceMixin): - """Offline inference for pooling models""" + """Offline inference for pooling models.""" runner_type: str chat_template: str | None @@ -93,8 +93,8 @@ def encode( Returns: A list of `PoolingRequestOutput` objects containing the pooled hidden states in the same order as the input prompts. - """ + """ if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin": raise ValueError( "The 'data' field is only supported for the 'plugin' pooling task." @@ -203,8 +203,7 @@ def embed( lora_request: list[LoRARequest] | LoRARequest | None = None, tokenization_kwargs: dict[str, Any] | None = None, ) -> list[EmbeddingRequestOutput]: - """ - Generate an embedding vector for each prompt. + """Generate an embedding vector for each prompt. This class automatically batches the given prompts, considering the memory constraint. For the best performance, put all of your prompts @@ -226,8 +225,8 @@ def embed( Returns: A list of `EmbeddingRequestOutput` objects containing the embedding vectors in the same order as the input prompts. - """ + """ items = self.encode( prompts, use_tqdm=use_tqdm, @@ -248,8 +247,7 @@ def classify( lora_request: list[LoRARequest] | LoRARequest | None = None, tokenization_kwargs: dict[str, Any] | None = None, ) -> list[ClassificationRequestOutput]: - """ - Generate class logits for each prompt. + """Generate class logits for each prompt. This class automatically batches the given prompts, considering the memory constraint. For the best performance, put all of your prompts @@ -271,8 +269,8 @@ def classify( Returns: A list of `ClassificationRequestOutput` objects containing the embedding vectors in the same order as the input prompts. - """ + """ items = self.encode( prompts, use_tqdm=use_tqdm, @@ -330,11 +328,12 @@ def score( chat_template: The chat template to use for the scoring. If None, we use the model's default chat template. tokenization_kwargs: Overrides for `tokenizer.encode`. + Returns: A list of `ScoringRequestOutput` objects containing the generated scores in the same order as the input prompts. - """ + """ if self.runner_type != "pooling": raise ValueError( "LLM.score() is only supported for pooling models. " diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 4b18cfc9e184..51ce9ba0efcf 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -44,7 +44,8 @@ class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): class PluginWithIOProcessorPlugins(PoolingIOProcessor): """IO Processor plugins are a feature that allows pre- and post-processing - of the model input and output for pooling models.""" + of the model input and output for pooling models. + """ name = "plugin" diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py index fd6191b6af8f..a67843e57770 100644 --- a/vllm/entrypoints/pooling/scoring/typing.py +++ b/vllm/entrypoints/pooling/scoring/typing.py @@ -21,8 +21,7 @@ class ScoreMultiModalParam(TypedDict, total=False): - """ - A specialized parameter type for scoring multimodal content + """A specialized parameter type for scoring multimodal content. The reasons why don't reuse `CustomChatCompletionMessageParam` directly: 1. Score tasks don't need the 'role' field (user/assistant/system) that's required in chat completions diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 148b429d3c5a..cff8ff61d914 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -58,8 +58,7 @@ def truncate_text_to_tokens( def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: - """ - Compute ColBERT MaxSim score. + """Compute ColBERT MaxSim score. Args: q_emb: Query token embeddings [query_len, dim] @@ -67,6 +66,7 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens Returns: MaxSim score (sum over query tokens of max similarity to any doc token) + """ # compute in float32 for numerical stability # [query_len, doc_len] @@ -249,8 +249,7 @@ def parse_score_data( def compress_token_type_ids(token_type_ids: list[int]) -> int: - """ - Return position of the first 1 or the length of the list + """Return position of the first 1 or the length of the list if not found. """ first_one = len(token_type_ids) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py index 2fc707cec475..28677b63213a 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py @@ -78,8 +78,7 @@ def attach_router(app: FastAPI): @router.post("/abort_requests") async def abort_requests(raw_request: Request): - """ - Abort one or more requests. To be used in a + """Abort one or more requests. To be used in a Disaggregated Everything setup. """ try: diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py index 60a3560ad737..b754fcc42e96 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Encode/decode utilities for multimodal tensors and field metadata -over JSON/HTTP, used by the disaggregated generate endpoint.""" +over JSON/HTTP, used by the disaggregated generate endpoint. +""" from __future__ import annotations diff --git a/vllm/entrypoints/serve/dev/cache/api_router.py b/vllm/entrypoints/serve/dev/cache/api_router.py index 63737393f29c..81ec91a8666a 100644 --- a/vllm/entrypoints/serve/dev/cache/api_router.py +++ b/vllm/entrypoints/serve/dev/cache/api_router.py @@ -23,8 +23,7 @@ async def reset_prefix_cache( reset_running_requests: bool = Query(default=False), reset_external: bool = Query(default=False), ): - """ - Reset the local prefix cache. + """Reset the local prefix cache. Optionally, if the query parameter `reset_external=true` also resets the external (connector-managed) prefix cache. @@ -35,6 +34,7 @@ async def reset_prefix_cache( Example: POST /reset_prefix_cache?reset_external=true + """ logger.info("Resetting prefix cache...") @@ -46,8 +46,7 @@ async def reset_prefix_cache( @router.post("/reset_mm_cache") async def reset_mm_cache(raw_request: Request): - """ - Reset the multi-modal cache. Note that we currently do not check if the + """Reset the multi-modal cache. Note that we currently do not check if the multi-modal cache is successfully reset in the API server. """ logger.info("Resetting multi-modal cache...") @@ -57,8 +56,7 @@ async def reset_mm_cache(raw_request: Request): @router.post("/reset_encoder_cache") async def reset_encoder_cache(raw_request: Request): - """ - Reset the encoder cache. Note that we currently do not check if the + """Reset the encoder cache. Note that we currently do not check if the encoder cache is successfully reset in the API server. """ logger.info("Resetting encoder cache...") diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py index 392fcf567476..603fdf785c3d 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -36,6 +36,8 @@ async def pause_generation( """Pause generation requests to allow weight updates. Args: + raw_request: The incoming FastAPI request, used to reach the engine + client on the app state. mode: How to handle in-flight requests: - ``"abort"``: Abort all in-flight requests immediately (default). - ``"wait"``: Wait for in-flight requests to complete. @@ -43,8 +45,8 @@ async def pause_generation( wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV/prefix caches after draining. Ignored when mode="keep". - """ + """ engine = engine_client(raw_request) try: @@ -74,7 +76,6 @@ async def pause_generation( @router.post("/resume") async def resume_generation(raw_request: Request) -> JSONResponse: """Resume generation after a pause.""" - engine = engine_client(raw_request) try: @@ -97,7 +98,6 @@ async def abort_requests(raw_request: Request) -> JSONResponse: Empty/missing ``request_ids`` aborts all in-flight requests. """ - engine = engine_client(raw_request) try: @@ -139,7 +139,6 @@ async def abort_requests(raw_request: Request) -> JSONResponse: @router.get("/is_paused") async def is_paused(raw_request: Request) -> JSONResponse: """Return the current pause status.""" - engine = engine_client(raw_request) try: @@ -234,9 +233,12 @@ async def get_world_size( """Get the world size from the parallel config. Args: + raw_request: The incoming FastAPI request, used to reach the engine + client on the app state. include_dp: If True (default), returns the world size including data parallelism (TP * PP * DP). If False, returns the world size without data parallelism (TP * PP). + """ parallel_config = engine_client(raw_request).vllm_config.parallel_config if include_dp: diff --git a/vllm/entrypoints/serve/elastic_ep/middleware.py b/vllm/entrypoints/serve/elastic_ep/middleware.py index 23f45eafeaa0..dcc62aea2c18 100644 --- a/vllm/entrypoints/serve/elastic_ep/middleware.py +++ b/vllm/entrypoints/serve/elastic_ep/middleware.py @@ -20,8 +20,7 @@ def set_scaling_elastic_ep(value): class ScalingMiddleware: - """ - Middleware that checks if the model is currently scaling and + """Middleware that checks if the model is currently scaling and returns a 503 Service Unavailable response if it is. This middleware applies to all HTTP requests and prevents diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 23e392720313..7a1befbe4c92 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/vllm/entrypoints/serve/engine/serving.py @@ -117,7 +117,7 @@ def _log_inputs( def _base_request_id( raw_request: Request | None, default: str | None = None ) -> str | None: - """Pulls the request id to use from a header, if provided""" + """Pulls the request id to use from a header, if provided.""" if raw_request is not None and ( (req_id := raw_request.headers.get("X-Request-Id")) is not None ): diff --git a/vllm/entrypoints/serve/exception_handling/handlers/validation.py b/vllm/entrypoints/serve/exception_handling/handlers/validation.py index 7cdcb3970e0f..acf2e57f3b28 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/validation.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/validation.py @@ -126,7 +126,8 @@ def _format_error(err: object) -> str: def _is_internal_loc_segment(segment: str) -> bool: """True if `segment` is a Pydantic-internal wrapper/union-branch - marker rather than a user-meaningful field name or list index.""" + marker rather than a user-meaningful field name or list index. + """ if _BRACKETED_INTERNAL_RE.search(segment): return True return segment.lower() in _INTERNAL_LOC_MARKERS diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py index 3ee5f0a06514..c0234139374e 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py @@ -28,8 +28,7 @@ async def vllm_error_handler(req: Request, exc: VLLMError): async def engine_error_handler( req: Request, exc: EngineDeadError | EngineGenerateError ): - """ - VLLM V1 AsyncLLM catches exceptions and returns + """VLLM V1 AsyncLLM catches exceptions and returns only two types: EngineGenerateError and EngineDeadError. EngineGenerateError is raised by the per request generate() @@ -51,7 +50,6 @@ async def engine_error_handler( will not automatically shut down. Instead, we use the watchdog background task for check for errored state. """ - if req.app.state.args.log_error_stack: logger.exception( "Engine Exception caught. Request id: %s", diff --git a/vllm/entrypoints/serve/exception_handling/register.py b/vllm/entrypoints/serve/exception_handling/register.py index 9f94873939a0..13c9425511fe 100644 --- a/vllm/entrypoints/serve/exception_handling/register.py +++ b/vllm/entrypoints/serve/exception_handling/register.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -vLLM Exception handlers are registered in four layers: +"""vLLM Exception handlers are registered in four layers: 1. framework errors raised by FastAPI/Starlette 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` diff --git a/vllm/entrypoints/serve/instrumentator/metrics.py b/vllm/entrypoints/serve/instrumentator/metrics.py index 5cba364f5d94..13241d03e884 100644 --- a/vllm/entrypoints/serve/instrumentator/metrics.py +++ b/vllm/entrypoints/serve/instrumentator/metrics.py @@ -55,7 +55,6 @@ class PrometheusResponse(Response): def attach_router(app: FastAPI): """Mount prometheus metrics to a FastAPI app.""" - registry = get_prometheus_registry() # `response_class=PrometheusResponse` is needed to return an HTTP response diff --git a/vllm/entrypoints/serve/lora/api_router.py b/vllm/entrypoints/serve/lora/api_router.py index 61436371c27d..cd7e5bc99ee8 100644 --- a/vllm/entrypoints/serve/lora/api_router.py +++ b/vllm/entrypoints/serve/lora/api_router.py @@ -63,7 +63,8 @@ def _attach_router(app: FastAPI): @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)]) async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request): """Handle POST /v1/load_lora_adapter: load a LoRA adapter into the - serving engine.""" + serving engine. + """ handler: OpenAIServingModels = models(raw_request) response = await handler.load_lora_adapter(request) if isinstance(response, ErrorResponse): @@ -85,7 +86,8 @@ async def unload_lora_adapter( request: UnloadLoRAAdapterRequest, raw_request: Request ): """Handle POST /v1/unload_lora_adapter: unload a LoRA adapter from - the serving engine.""" + the serving engine. + """ handler: OpenAIServingModels = models(raw_request) response = await handler.unload_lora_adapter(request) if isinstance(response, ErrorResponse): diff --git a/vllm/entrypoints/serve/middleware/authenticate.py b/vllm/entrypoints/serve/middleware/authenticate.py index e8a61b21ee46..0a0c1daeec4d 100644 --- a/vllm/entrypoints/serve/middleware/authenticate.py +++ b/vllm/entrypoints/serve/middleware/authenticate.py @@ -12,8 +12,7 @@ class AuthenticationMiddleware: - """ - Pure ASGI middleware that authenticates each request by checking + """Pure ASGI middleware that authenticates each request by checking if the Authorization Bearer token exists and equals anyof "{api_key}". Notes @@ -21,6 +20,7 @@ class AuthenticationMiddleware: There are two cases in which authentication is skipped: 1. The HTTP method is OPTIONS. 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). + """ def __init__(self, app: ASGIApp, tokens: list[str]) -> None: diff --git a/vllm/entrypoints/serve/middleware/log_response.py b/vllm/entrypoints/serve/middleware/log_response.py index e2f5e5bde51e..88f37a75984a 100644 --- a/vllm/entrypoints/serve/middleware/log_response.py +++ b/vllm/entrypoints/serve/middleware/log_response.py @@ -94,7 +94,6 @@ def _extract_content_from_chunk(chunk_data: dict) -> str: def _log_streaming_response(response, response_body: list) -> None: """Log streaming response with robust SSE parsing.""" - sse_decoder = SSEDecoder() chunk_count = 0 diff --git a/vllm/entrypoints/serve/middleware/x_request_id.py b/vllm/entrypoints/serve/middleware/x_request_id.py index 02e85306cd12..1f654caf3c4e 100644 --- a/vllm/entrypoints/serve/middleware/x_request_id.py +++ b/vllm/entrypoints/serve/middleware/x_request_id.py @@ -8,8 +8,7 @@ class XRequestIdMiddleware: - """ - Middleware the set's the X-Request-Id header for each response + """Middleware the set's the X-Request-Id header for each response to a random uuid4 (hex) value if the header isn't already present in the request, otherwise use the provided request id. """ @@ -25,8 +24,7 @@ def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None request_headers = Headers(scope=scope) async def send_with_request_id(message: Message) -> None: - """ - Custom send function to mutate the response headers + """Custom send function to mutate the response headers and append X-Request-Id to it. """ if message["type"] == "http.response.start": diff --git a/vllm/entrypoints/serve/sagemaker/api_router.py b/vllm/entrypoints/serve/sagemaker/api_router.py index 60576f289a0d..5c4e3de74001 100644 --- a/vllm/entrypoints/serve/sagemaker/api_router.py +++ b/vllm/entrypoints/serve/sagemaker/api_router.py @@ -30,7 +30,8 @@ def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]: """Snapshot handler levels for loggers that may be affected by third-party logging configuration side effects (e.g. model_hosting_container_standards - calling configure_root_logger() at import time).""" + calling configure_root_logger() at import time). + """ loggers = [logging.getLogger(), logging.getLogger("vllm")] seen: set[int] = set() snapshot: list[tuple[logging.Handler, int]] = [] @@ -73,7 +74,8 @@ def _attach_router( model_config: ModelConfig | None = None, ): """Register the SageMaker hosting routes (/ping, /invocations) on the - app.""" + app. + """ import model_hosting_container_standards.sagemaker as sagemaker_standards router = APIRouter() @@ -92,7 +94,7 @@ def _attach_router( @router.get("/ping", response_class=Response) @sagemaker_standards.register_ping_handler async def ping(raw_request: Request) -> Response: - """Ping check. Endpoint required for SageMaker""" + """Ping check. Endpoint required for SageMaker.""" return await health(raw_request) @router.post( diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py index 12d487346f76..130a1b4b45c4 100644 --- a/vllm/entrypoints/serve/tokenize/protocol.py +++ b/vllm/entrypoints/serve/tokenize/protocol.py @@ -181,8 +181,7 @@ class DetokenizeResponse(OpenAIBaseModel): class TokenizerInfoResponse(OpenAIBaseModel): - """ - Response containing tokenizer configuration + """Response containing tokenizer configuration equivalent to tokenizer_config.json """ diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index aa089cff44b4..6b82a2afa070 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -34,7 +34,7 @@ async def listen_for_disconnect(request: Request) -> None: - """Returns if a disconnect message is received""" + """Returns if a disconnect message is received.""" while True: message = await request.receive() if message["type"] == "http.disconnect": diff --git a/vllm/entrypoints/serve/utils/fingerprint.py b/vllm/entrypoints/serve/utils/fingerprint.py index e858e667d06c..5203d1003bdf 100644 --- a/vllm/entrypoints/serve/utils/fingerprint.py +++ b/vllm/entrypoints/serve/utils/fingerprint.py @@ -33,7 +33,8 @@ def set_default_fingerprint_mode( custom_value: str | None = None, ) -> None: """Configure the fingerprint mode for subsequent ``get_system_fingerprint`` - calls. Called once at server startup.""" + calls. Called once at server startup. + """ global _DEFAULT_MODE, _CUSTOM_VALUE _DEFAULT_MODE = mode _CUSTOM_VALUE = custom_value @@ -41,7 +42,8 @@ def set_default_fingerprint_mode( def get_system_fingerprint(vllm_config: Any) -> str | None: """Return the fingerprint for ``vllm_config`` using the mode configured by - ``set_default_fingerprint_mode``.""" + ``set_default_fingerprint_mode``. + """ return build_system_fingerprint(vllm_config, _DEFAULT_MODE, _CUSTOM_VALUE) diff --git a/vllm/entrypoints/serve/utils/orca_metrics.py b/vllm/entrypoints/serve/utils/orca_metrics.py index 3808262bf31f..2c0aa4381227 100644 --- a/vllm/entrypoints/serve/utils/orca_metrics.py +++ b/vllm/entrypoints/serve/utils/orca_metrics.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Utility functions that create ORCA endpoint load report response headers. -""" +"""Utility functions that create ORCA endpoint load report response headers.""" import json from collections.abc import Mapping @@ -16,23 +14,22 @@ def create_orca_header( metrics_format: str, named_metrics: list[tuple[str, float]] ) -> Mapping[str, str] | None: - """ - Creates ORCA headers named 'endpoint-load-metrics' in the specified format + """Creates ORCA headers named 'endpoint-load-metrics' in the specified format and adds custom metrics to named_metrics. ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0 ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto - Parameters: - - metrics_format (str): The format of the header ('TEXT', 'JSON'). - - named_metrics (List[Tuple[str, float]]): List of tuples with metric names - and their corresponding double values. + Args: + metrics_format (str): The format of the header ('TEXT', 'JSON'). + named_metrics (List[Tuple[str, float]]): List of tuples with metric names + and their corresponding double values. Returns: - Optional[Mapping[str,str]]: A dictionary with header key as 'endpoint-load-metrics' and values as the ORCA header strings with format prefix and data in with named_metrics in. - """ + """ if metrics_format.lower() not in ["text", "json"]: logger.warning( "Warning: `%s` format is not supported in the ORCA response header", @@ -71,15 +68,17 @@ def create_orca_header( def get_named_metrics_from_prometheus() -> list[tuple[str, float]]: - """ - Collects current metrics from Prometheus and returns some of them + """Collects current metrics from Prometheus and returns some of them in the form of the `named_metrics` list for `create_orca_header()`. - Parameters: + Parameters + ---------- - None - Returns: + Returns + ------- - list[tuple[str, float]]: List of tuples of metric names and their values. + """ named_metrics: list[tuple[str, float]] = [] # Map from prometheus metric names to ORCA named metrics. @@ -98,20 +97,20 @@ def get_named_metrics_from_prometheus() -> list[tuple[str, float]]: def metrics_header(metrics_format: str) -> Mapping[str, str] | None: - """ - Creates ORCA headers named 'endpoint-load-metrics' in the specified format. + """Creates ORCA headers named 'endpoint-load-metrics' in the specified format. Metrics are collected from Prometheus using `get_named_metrics_from_prometheus()`. ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0 ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto - Parameters: - - metrics_format (str): The format of the header ('TEXT', 'JSON'). + Args: + metrics_format (str): The format of the header ('TEXT', 'JSON'). Returns: - Optional[Mapping[str,str]]: A dictionary with header key as 'endpoint-load-metrics' and values as the ORCA header strings with format prefix and data in with named_metrics in. + """ if not metrics_format: return None diff --git a/vllm/entrypoints/serve/utils/tool_calls_utils.py b/vllm/entrypoints/serve/utils/tool_calls_utils.py index 42106f433405..ae2fa24f7f63 100644 --- a/vllm/entrypoints/serve/utils/tool_calls_utils.py +++ b/vllm/entrypoints/serve/utils/tool_calls_utils.py @@ -20,7 +20,6 @@ def maybe_filter_parallel_tool_calls( choice: _ChatCompletionResponseChoiceT, request: ChatCompletionRequest ) -> _ChatCompletionResponseChoiceT: """Filter to first tool call only when parallel_tool_calls is explicitly False.""" - if request.parallel_tool_calls is not False: return choice diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index dc09db6aa319..2e77ceffeffb 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -92,7 +92,8 @@ def asr_inter_chunk_separator( class SpeechToTextBaseServing(GenerateBaseServing): """Base class for speech-to-text operations like transcription and - translation.""" + translation. + """ def __init__( self, @@ -358,8 +359,7 @@ def _get_verbose_segments( segment_class: type[SpeechToTextSegment], start_time: float = 0, ) -> list[SpeechToTextSegment]: - """ - Convert tokens to verbose segments. + """Convert tokens to verbose segments. This method expects the model to produce timestamps as tokens (similar to Whisper). @@ -427,7 +427,8 @@ async def _create_speech_to_text( stream_generator_method: Callable[..., AsyncGenerator[str, None]], ) -> T | V | AsyncGenerator[str, None] | ErrorResponse: """Base method for speech-to-text operations like transcription and - translation.""" + translation. + """ if request.stream and request.use_beam_search: return self.create_error_response( "Streaming is not currently supported with beam search" diff --git a/vllm/entrypoints/speech_to_text/base/utils.py b/vllm/entrypoints/speech_to_text/base/utils.py index bcd29f08e967..8282d6f41dea 100644 --- a/vllm/entrypoints/speech_to_text/base/utils.py +++ b/vllm/entrypoints/speech_to_text/base/utils.py @@ -33,6 +33,7 @@ async def read_upload_with_limit( Raises: VLLMValidationError: If the file exceeds the configured size limit. + """ if max_size_mb is None: max_size_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB diff --git a/vllm/entrypoints/speech_to_text/realtime/protocol.py b/vllm/entrypoints/speech_to_text/realtime/protocol.py index bdcd335f6fe4..7796eee8b848 100644 --- a/vllm/entrypoints/speech_to_text/realtime/protocol.py +++ b/vllm/entrypoints/speech_to_text/realtime/protocol.py @@ -13,14 +13,14 @@ class InputAudioBufferAppend(OpenAIBaseModel): - """Append audio chunk to buffer""" + """Append audio chunk to buffer.""" type: Literal["input_audio_buffer.append"] = "input_audio_buffer.append" audio: str # base64-encoded PCM16 @ 16kHz class InputAudioBufferCommit(OpenAIBaseModel): - """Process accumulated audio buffer""" + """Process accumulated audio buffer.""" type: Literal["input_audio_buffer.commit"] = "input_audio_buffer.commit" final: bool = False @@ -28,14 +28,14 @@ class InputAudioBufferCommit(OpenAIBaseModel): # Server -> Client Events class SessionUpdate(OpenAIBaseModel): - """Configure session parameters""" + """Configure session parameters.""" type: Literal["session.update"] = "session.update" model: str | None = None class SessionCreated(OpenAIBaseModel): - """Connection established notification""" + """Connection established notification.""" type: Literal["session.created"] = "session.created" id: str = Field(default_factory=lambda: f"sess-{random_uuid()}") @@ -43,14 +43,14 @@ class SessionCreated(OpenAIBaseModel): class TranscriptionDelta(OpenAIBaseModel): - """Incremental transcription text""" + """Incremental transcription text.""" type: Literal["transcription.delta"] = "transcription.delta" delta: str # Incremental text class TranscriptionDone(OpenAIBaseModel): - """Final transcription with usage stats""" + """Final transcription with usage stats.""" type: Literal["transcription.done"] = "transcription.done" text: str # Complete transcription @@ -58,7 +58,7 @@ class TranscriptionDone(OpenAIBaseModel): class ErrorEvent(OpenAIBaseModel): - """Error notification""" + """Error notification.""" type: Literal["error"] = "error" error: str diff --git a/vllm/entrypoints/speech_to_text/realtime/serving.py b/vllm/entrypoints/speech_to_text/realtime/serving.py index bec2b95143b6..9745269607ab 100644 --- a/vllm/entrypoints/speech_to_text/realtime/serving.py +++ b/vllm/entrypoints/speech_to_text/realtime/serving.py @@ -68,6 +68,7 @@ async def transcribe_realtime( Yields: StreamingInput objects containing audio prompts for the engine + """ model_config = self.model_config renderer = self.renderer diff --git a/vllm/env_override.py b/vllm/env_override.py index a6b3c373face..35431cccb666 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -261,8 +261,7 @@ def get_output_names(graph_outputs) -> list[str]: def get_graph_partition_signature_patched( self, partitions, skip_cudagraphs: list[bool] ): - """ - Gets signature for each graph partition, including input nodes, output nodes, and + """Gets signature for each graph partition, including input nodes, output nodes, and whether deallocating an input within graph partition. """ from torch._inductor import dependencies @@ -276,8 +275,7 @@ def get_graph_partition_signature_patched( name_to_node = self.get_name_to_nodes() def is_none_layout(buf_name: str) -> bool: - """ - Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated + """Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated so graph partition should not take it as inputs or outputs. """ buf = self.name_to_buf.get(buf_name, None) @@ -425,8 +423,7 @@ def should_partition_patched(self, node, should_log: bool = False) -> bool: # torch._inductor.scheduler.Scheduler.should_partition that modifies # the following piece of code so that we always return True: # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724 - """Return True if we should partition the inductor graph on this node""" - + """Return True if we should partition the inductor graph on this node.""" import torch._inductor.ir as ir from torch._inductor.scheduler import ( BaseSchedulerNode, @@ -507,8 +504,7 @@ def _update_scheduler_patched(self) -> None: # Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches # this method so that we can patch Scheduler.should_partition with the # function above - """ - (Re)initializes the scheduler member. When initializing the scheduler, no CUBIN + """(Re)initializes the scheduler member. When initializing the scheduler, no CUBIN files should be generated (to avoid biasing any benchmarks and pessimizing fusion decisions). """ @@ -585,7 +581,8 @@ def _patch_get_raw_stream_if_needed(): def _apply_constrain_to_fx_strides_patch(): """Patch lowering.constrain_to_fx_strides globally. Safe to call multiple times; only the first call does anything. - Only applies for torch >= 2.11 and < 2.12.""" + Only applies for torch >= 2.11 and < 2.12. + """ global _constrain_to_fx_strides_patched if _constrain_to_fx_strides_patched: return diff --git a/vllm/envs.py b/vllm/envs.py index 333d75b0ff0e..8d836ffc4e71 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -392,8 +392,7 @@ def env_with_choices( choices: list[str] | Callable[[], list[str]], case_sensitive: bool = True, ) -> Callable[[], str | None]: - """ - Create a lambda that validates environment variable against allowed choices + """Create a lambda that validates environment variable against allowed choices. Args: env_name: Name of the environment variable @@ -403,6 +402,7 @@ def env_with_choices( Returns: Lambda function for environment_variables dict + """ def _get_validated_env() -> str | None: @@ -437,8 +437,7 @@ def env_list_with_choices( choices: list[str] | Callable[[], list[str]], case_sensitive: bool = True, ) -> Callable[[], list[str]]: - """ - Create a lambda that validates environment variable + """Create a lambda that validates environment variable containing comma-separated values against allowed choices Args: @@ -450,6 +449,7 @@ def env_list_with_choices( Returns: Lambda function for environment_variables dict that returns list of strings + """ def _get_validated_env_list() -> list[str]: @@ -492,8 +492,7 @@ def env_set_with_choices( choices: list[str] | Callable[[], list[str]], case_sensitive: bool = True, ) -> Callable[[], set[str]]: - """ - Creates a lambda which that validates environment variable + """Creates a lambda which that validates environment variable containing comma-separated values against allowed choices which returns choices as a set. """ @@ -512,6 +511,7 @@ def get_vllm_port() -> int | None: Raises: ValueError: If VLLM_PORT is a URI, suggest k8s service discovery issue. + """ if "VLLM_PORT" not in os.environ: return None @@ -537,8 +537,7 @@ def get_env_or_set_default( env_name: str, default_factory: Callable[[], str], ) -> Callable[[], str]: - """ - Create a lambda that returns an environment variable value if set, + """Create a lambda that returns an environment variable value if set, or generates and sets a default value using the provided factory function. """ @@ -2187,8 +2186,7 @@ def _resolve_rust_cli_path() -> str | None: def __getattr__(name: str): - """ - Gets environment variables lazily. + """Gets environment variables lazily. NOTE: After enable_envs_cache() invocation (which triggered after service initialization), all environment variables will be cached. @@ -2199,14 +2197,13 @@ def __getattr__(name: str): def _is_envs_cache_enabled() -> bool: - """Checked if __getattr__ is wrapped with functools.cache""" + """Checked if __getattr__ is wrapped with functools.cache.""" global __getattr__ return hasattr(__getattr__, "cache_clear") def enable_envs_cache() -> None: - """ - Enables caching of environment variables. This is useful for performance + """Enables caching of environment variables. This is useful for performance reasons, as it avoids the need to re-evaluate environment variables on every call. @@ -2227,8 +2224,7 @@ def enable_envs_cache() -> None: def disable_envs_cache() -> None: - """ - Resets the environment variables cache. It could be used to isolate environments + """Resets the environment variables cache. It could be used to isolate environments between unit tests. """ global __getattr__ @@ -2253,8 +2249,8 @@ def compile_factors() -> dict[str, object]: """Return env vars used for torch.compile cache keys. Start with every known vLLM env var; drop entries in `ignored_factors`; - hash everything else. This keeps the cache key aligned across workers.""" - + hash everything else. This keeps the cache key aligned across workers. + """ ignored_factors: set[str] = { "MAX_JOBS", "VLLM_RPC_BASE_PATH", diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 582e83ec4a6d..c7900a0347e7 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -32,6 +32,7 @@ class VLLMValidationError(VLLMClientError): message: The error message describing the validation failure. parameter: Optional parameter name that failed validation. value: Optional value that was rejected during validation. + """ def __init__( @@ -56,7 +57,7 @@ def __str__(self): class VLLMNotFoundError(VLLMClientError): - """vLLM-specific NotFoundError""" + """vLLM-specific NotFoundError.""" pass @@ -69,6 +70,7 @@ class LoRAAdapterNotFoundError(VLLMNotFoundError): Attributes: message: The error message string describing the exception + """ message: str @@ -96,6 +98,7 @@ class VLLMUnprocessableEntityError(VLLMClientError): message: The error message describing the unprocessable entity. parameter: Optional parameter name that failed validation. value: Optional value that was rejected during validation. + """ def __init__( diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 8d9eb5d7db3b..dc025175a357 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -28,8 +28,7 @@ @dataclass(frozen=True) class BatchDescriptor: - """ - Batch descriptor for cudagraph dispatching. We should keep the num of + """Batch descriptor for cudagraph dispatching. We should keep the num of items as minimal as possible to properly and uniquely describe the padded batch for cudagraph. """ @@ -100,8 +99,7 @@ def make( @contextmanager def sp_local_sizes(self, sequence_parallel_size: int): - """ - Context manager for setting self.local_sizes. Same as self.chunked_sizes + """Context manager for setting self.local_sizes. Same as self.chunked_sizes but without any chunking. """ self.local_sizes = _compute_sp_num_tokens( @@ -211,7 +209,8 @@ def is_forward_context_available() -> bool: def in_piecewise_cudagraph() -> bool: """Whether the current forward runs in piecewise cudagraph mode (graph - segments separated by eager breaks), at capture or replay time.""" + segments separated by eager breaks), at capture or replay time. + """ return ( is_forward_context_available() and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index bda40bebbc86..429e8a829dd0 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -16,8 +16,7 @@ class _InputOptions(TypedDict): - """ - Additional options available to all + """Additional options available to all [`SingletonInput`][vllm.inputs.engine.SingletonInput] types. """ @@ -56,8 +55,7 @@ def tokens_input( prompt: str | None = None, cache_salt: str | None = None, ) -> TokensInput: - """ - Construct [`TokensInput`][vllm.inputs.engine.TokensInput] + """Construct [`TokensInput`][vllm.inputs.engine.TokensInput] from optional values. """ inputs = TokensInput(type="token", prompt_token_ids=prompt_token_ids) @@ -104,8 +102,7 @@ def embeds_input( prompt_token_ids: list[int] | None = None, is_token_ids: list[bool] | None = None, ) -> EmbedsInput: - """ - Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput] + """Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput] from optional values. """ inputs = EmbedsInput(type="embeds", prompt_embeds=prompt_embeds) @@ -190,13 +187,13 @@ def mm_input( class MultiModalEncDecInput(MultiModalInput): - """ - Represents multi-modal input to the engine for encoder-decoder models. + """Represents multi-modal input to the engine for encoder-decoder models. Note: Even text-only encoder-decoder models are currently implemented as multi-modal models for convenience. (Example: https://github.com/vllm-project/bart-plugin) + """ encoder_prompt_token_ids: list[int] @@ -253,8 +250,7 @@ def mm_enc_dec_input( class EncoderDecoderInput(TypedDict): - """ - A rendered [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt] + """A rendered [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt] which can be passed to `LLMEngine.add_request` or `AsyncLLM.add_request`. """ @@ -315,8 +311,7 @@ def _prepare_decoder_input_ids_for_generation( decoder_input_ids: list[int], decoder_start_token_id: int, ) -> list[int]: - """ - Prepare `decoder_input_ids` for generation with encoder-decoder models, + """Prepare `decoder_input_ids` for generation with encoder-decoder models, according to `GenerationMixin._prepare_decoder_input_ids_for_generation()`. Source: diff --git a/vllm/inputs/llm.py b/vllm/inputs/llm.py index a9ac8201fbb1..a95b274c1853 100644 --- a/vllm/inputs/llm.py +++ b/vllm/inputs/llm.py @@ -62,8 +62,7 @@ class MultiModalDataBuiltins(TypedDict, total=False): class _PromptOptions(TypedDict): - """ - Additional options available to all + """Additional options available to all [`SingletonPrompt`][vllm.inputs.llm.SingletonPrompt] types. """ @@ -189,11 +188,11 @@ class EmbedsPrompt(_PromptOptions): class ExplicitEncoderDecoderPrompt(TypedDict): - """ - Schema for a pair of encoder and decoder singleton prompts. + """Schema for a pair of encoder and decoder singleton prompts. Note: This schema is not valid for decoder-only models. + """ encoder_prompt: EncoderPrompt @@ -233,8 +232,7 @@ class ExplicitEncoderDecoderPrompt(TypedDict): class DataPrompt(_PromptOptions): - """ - Represents generic inputs that are converted to + """Represents generic inputs that are converted to [`PromptType`][vllm.inputs.llm.PromptType] by IO processor plugins. """ diff --git a/vllm/ir/op.py b/vllm/ir/op.py index 742d3f33ff85..996bc5ab2ae6 100644 --- a/vllm/ir/op.py +++ b/vllm/ir/op.py @@ -52,23 +52,19 @@ def _validate_name(name: str, entity_type: str) -> None: def set_default_torch_wrap(enable: bool = True) -> None: - """ - Permanently set the torch wrap flag. - """ + """Permanently set the torch wrap flag.""" global _ENABLE_TORCH_WRAP _ENABLE_TORCH_WRAP = enable @contextlib.contextmanager def enable_torch_wrap(enable: bool = True): - """ - Context manager to enable/disable torch custom op wrapping for vLLM IR ops. + """Context manager to enable/disable torch custom op wrapping for vLLM IR ops. When torch wrapping is disabled, the torch custom op layer is skipped and IR ops dispatch directly to the implementation. Helpful for avoiding torch dispatch overhead in eager mode and avoiding the need for lowering for platforms not using Inductor. """ - global _ENABLE_TORCH_WRAP old = _ENABLE_TORCH_WRAP try: @@ -110,8 +106,7 @@ def register_op( activations: list[str] | None = None, allow_inplace: bool = False, ) -> "IrOp | Callable[[Callable], IrOp]": - """ - Register a new vLLM IR op. + """Register a new vLLM IR op. Args: f: the native implementation of the op @@ -131,7 +126,10 @@ def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: @vllm.ir.register_op(name="custom_mul") def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return x * y""" + return x * y + ``` + + """ def decorator(_f: Callable): op_name: str = _f.__name__ if name is None else name @@ -226,16 +224,14 @@ def __init__( self.torch_op: torch._ops.OpOverload = getattr(torch_ops, name).default def register_fake(self, fn: Callable) -> Callable: - """ - Register a fake impl for the torch custom op. If this method is not called, + """Register a fake impl for the torch custom op. If this method is not called, the native implementation is used directly for the fake implementation. """ self._fake_fn = fn return fn def _fake_call(self, *args, **kwargs) -> Any: - """ - Call to the fake implementation of the op. We use indirection because we want + """Call to the fake implementation of the op. We use indirection because we want users to be able to register fake later but also want it to fall back to native directly by default, instead of going through the dispatching mechanism. """ @@ -249,8 +245,8 @@ def register_impl( supports_args: Callable[..., bool] | None = None, inplace: bool = False, ) -> Callable[[Callable[..., Any]], "IrOpImpl"]: - """ - Register an implementation for this custom op. + """Register an implementation for this custom op. + Args: provider: The name of the provider, must be unique. supported: Static support check, use this to check platform support. @@ -302,8 +298,7 @@ def _register_impl(f: Callable): return _register_impl def _inner_call(self, *args, **kwargs) -> Any: - """ - Eager call to torch op lands here. When torch wrapping is disabled, + """Eager call to torch op lands here. When torch wrapping is disabled, __call__ routes straight here instead of going through torch op dispatching. """ impl = self.dispatch(*args, **kwargs) @@ -313,8 +308,7 @@ def _inner_call(self, *args, **kwargs) -> Any: return impl.func_impl_fn(*args, **kwargs) def apply_arg_defaults(self, args) -> tuple: - """ - Return args with default values applied. + """Return args with default values applied. Defaults are taken from the native implementation signature. SHOULD NOT BE USED IN THE DISPATCH PATH (SLOW). @@ -325,8 +319,7 @@ def apply_arg_defaults(self, args) -> tuple: return bound_args.args def dispatch(self, *args, **kwargs) -> "IrOpImpl": - """ - Dispatch to the appropriate implementation based on current priority + """Dispatch to the appropriate implementation based on current priority and argument support checks. Returns the selected IrOpImpl. THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST. @@ -413,8 +406,7 @@ def _filter_priority_impls(self, priority: list[str]) -> list["IrOpImpl"]: return filtered_impls def set_default(self, priority: list[str]) -> None: - """ - Permanently set the dispatch priority for this op. Use this for + """Permanently set the dispatch priority for this op. Use this for process-lifetime setup (e.g., worker startup). For scoped overrides, use ``set_priority`` instead. """ @@ -427,9 +419,7 @@ def set_default(self, priority: list[str]) -> None: @contextlib.contextmanager def set_priority(self, priority: list[str]): - """ - Context manager to set the dispatch priority for implementations for this op. - """ + """Context manager setting the dispatch priority for this op's impls.""" old_priority_impls = self._priority_impls try: self._priority_impls = self._filter_priority_impls(priority) @@ -636,10 +626,10 @@ def supports_args(self, *args, **kwargs) -> bool: @weak_cache def uuid(self): - """ - Compile-time hash to uniquely determine whether the implementation has changed. - Used by vllm-compile hash mechanism and torch.compile lowering pass uuid to - control the vLLM compile cache and AOTAutograd/Inductor caches respectively. + """Compile-time hash to uniquely determine whether the implementation has + changed. Used by vllm-compile hash mechanism and torch.compile lowering pass + uuid to control the vLLM compile cache and AOTAutograd/Inductor caches + respectively. Source file contents do not change so we cache uuid. TODO(luka): Cache the file hash as multiple impls are likely in the same file. @@ -648,8 +638,7 @@ def uuid(self): return hash_source(*sources) def func_impl_fn(self, *args, **kwargs) -> Any: - """ - Copy any inputs in activations if this is an inplace impl, + """Copy any inputs in activations if this is an inplace impl, to ensure functional semantics. """ if not self.inplace: diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index a40028a4d0bc..05bfca73bde0 100644 --- a/vllm/ir/ops/layernorm.py +++ b/vllm/ir/ops/layernorm.py @@ -10,7 +10,7 @@ def rms_norm( x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None ) -> Tensor: - """Weighted root-mean-square layer normalization""" + """Weighted root-mean-square layer normalization.""" orig_dtype = x.dtype x = x.to(torch.float32) x_var = x if variance_size is None else x[..., :variance_size] @@ -48,7 +48,7 @@ def fused_add_rms_norm( epsilon: float, variance_size: int | None = None, ) -> tuple[Tensor, Tensor]: - """Fused add and weighted root-mean-square layer normalization""" + """Fused add and weighted root-mean-square layer normalization.""" orig_dtype = x.dtype x = x.to(torch.float32) x = x + x_residual.to(torch.float32) diff --git a/vllm/ir/util.py b/vllm/ir/util.py index e9240f487aca..150d4219e6cc 100644 --- a/vllm/ir/util.py +++ b/vllm/ir/util.py @@ -10,11 +10,12 @@ def hash_source(*srcs: str | Any) -> str: - """ - Utility method to hash the sources of functions or objects. + """Utility method to hash the sources of functions or objects. + Args: srcs: strings or objects to add to the hash. Objects and functions have their source inspected. + """ hasher = hashlib.sha256() for src in srcs: @@ -34,8 +35,7 @@ def hash_source(*srcs: str | Any) -> str: def weak_lru_cache(maxsize: int | None = 128, typed: bool = False): - """ - LRU Cache decorator that keeps a weak reference to 'self'. + """LRU Cache decorator that keeps a weak reference to 'self'. This avoids memory leakage, which happens when functools.lru_cache stores a reference to self in the global cache. @@ -57,5 +57,5 @@ def inner(self, *args, **kwargs): def weak_cache(user_function, /): - """Simple weak equivalent to functools.cache""" + """Simple weak equivalent to functools.cache.""" return weak_lru_cache(maxsize=None)(user_function) diff --git a/vllm/kernels/helion/case_key.py b/vllm/kernels/helion/case_key.py index 32b544de39cc..0f90dac610aa 100644 --- a/vllm/kernels/helion/case_key.py +++ b/vllm/kernels/helion/case_key.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Structured key for identifying kernel config/autotune/benchmark cases. -""" +"""Structured key for identifying kernel config/autotune/benchmark cases.""" from __future__ import annotations diff --git a/vllm/kernels/helion/config_manager.py b/vllm/kernels/helion/config_manager.py index 052b515c37ff..8887cdaea48a 100644 --- a/vllm/kernels/helion/config_manager.py +++ b/vllm/kernels/helion/config_manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Configuration management for Helion kernels. +"""Configuration management for Helion kernels. This module provides centralized configuration file management for Helion custom operations, including naming conventions, directory resolution, and file I/O. diff --git a/vllm/kernels/helion/ops/__init__.py b/vllm/kernels/helion/ops/__init__.py index 588c070a1245..0fb82f6cfe0d 100644 --- a/vllm/kernels/helion/ops/__init__.py +++ b/vllm/kernels/helion/ops/__init__.py @@ -22,6 +22,7 @@ def import_all_kernels() -> list[str]: Returns: The fully-qualified module names that were imported. + """ imported: list[str] = [] for module_info in pkgutil.iter_modules(__path__): diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index 0ebbf724094e..9a566b0eda41 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -68,7 +68,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/ops/fused_qk_norm_rope.py b/vllm/kernels/helion/ops/fused_qk_norm_rope.py index dbbb9e088641..2fc7cb5dcb06 100644 --- a/vllm/kernels/helion/ops/fused_qk_norm_rope.py +++ b/vllm/kernels/helion/ops/fused_qk_norm_rope.py @@ -118,7 +118,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py index 789fa8dd65fa..cbb19eb800d9 100644 --- a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -94,7 +94,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index 276c7f96701b..90a48abf65b9 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -87,7 +87,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index 8c357386bdd5..adf98601326a 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -112,7 +112,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index 7cc0f99ea5ed..96e24cd988f6 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -96,7 +96,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | pick the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest. """ - if not config_keys: return None diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index 438011b1aa69..e8305c9d3ecf 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -vLLM Helion kernel registration with pre-tuned config selection. +"""vLLM Helion kernel registration with pre-tuned config selection. This module leverages Helion's internal config selection infrastructure to use pre-tuned configs instead of runtime autotuning. @@ -164,8 +163,7 @@ def __call__(self, *args, **kwargs): return self._decorated_kernel(*args, **kwargs) def _create_key_computer(self): - """ - Create a key computer function derived from the config picker. + """Create a key computer function derived from the config picker. The returned function receives kernel arguments unpacked (*args) to match Helion's key signature (called as self._key_fn(*args)). @@ -412,6 +410,10 @@ def register_kernel( """Register a Helion kernel with pre-tuned config selection. Args: + op_name: Name to register the op under. Defaults to the function name. + fake_impl: Optional meta/fake implementation for tracing. + mutates_args: Names of arguments the kernel mutates in place. + helion_settings: Optional Helion settings override. config_picker: Required. Receives ``(args, config_keys)`` where each config key is a ``dict[str, Any]`` mapping parameter names to values. Return the best-matching @@ -435,6 +437,7 @@ def generate_inputs(): "4096": (torch.randn(4096, device="cuda"), 0.5), "8192": (torch.randn(8192, device="cuda"), 0.5), } + """ def decorator(kernel_func: Callable) -> HelionKernelWrapper: diff --git a/vllm/kernels/helion/utils.py b/vllm/kernels/helion/utils.py index f4ace6cff3bf..7887fc9a5996 100644 --- a/vllm/kernels/helion/utils.py +++ b/vllm/kernels/helion/utils.py @@ -60,8 +60,7 @@ def get_gpu_name(device_id: int | None = None) -> str: def canonicalize_gpu_name(name: str) -> str: - """ - Canonicalize GPU name for use as a platform identifier. + """Canonicalize GPU name for use as a platform identifier. Converts to lowercase, replaces separators with underscores, and maps known variant names to their canonical form via _GPU_NAME_ALIASES. diff --git a/vllm/logger.py b/vllm/logger.py index ec67e504bb01..4f528a130193 100644 --- a/vllm/logger.py +++ b/vllm/logger.py @@ -100,7 +100,7 @@ def _print_error_once(logger: Logger, msg: str, *args: Hashable) -> None: def _should_log_with_scope(scope: LogScope) -> bool: - """Decide whether to log based on scope""" + """Decide whether to log based on scope.""" if scope == "global": from vllm.distributed.parallel_state import is_global_first_rank @@ -113,17 +113,16 @@ def _should_log_with_scope(scope: LogScope) -> bool: class _VllmLogger(Logger): - """ - Note: - This class is just to provide type information. - We actually patch the methods directly on the [`logging.Logger`][] - instance to avoid conflicting with other libraries such as - `intel_extension_for_pytorch.utils._logger`. + """Note: + This class is just to provide type information. + We actually patch the methods directly on the [`logging.Logger`][] + instance to avoid conflicting with other libraries such as + `intel_extension_for_pytorch.utils._logger`. + """ def debug_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None: - """ - As [`debug`][logging.Logger.debug], but subsequent calls with + """As [`debug`][logging.Logger.debug], but subsequent calls with the same message are silently dropped. """ if not _should_log_with_scope(scope): @@ -131,8 +130,7 @@ def debug_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> No _print_debug_once(self, msg, *args) def info_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None: - """ - As [`info`][logging.Logger.info], but subsequent calls with + """As [`info`][logging.Logger.info], but subsequent calls with the same message are silently dropped. """ if not _should_log_with_scope(scope): @@ -142,8 +140,7 @@ def info_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> Non def warning_once( self, msg: str, *args: Hashable, scope: LogScope = "local" ) -> None: - """ - As [`warning`][logging.Logger.warning], but subsequent calls with + """As [`warning`][logging.Logger.warning], but subsequent calls with the same message are silently dropped. """ if not _should_log_with_scope(scope): @@ -151,8 +148,7 @@ def warning_once( _print_warning_once(self, msg, *args) def error_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None: - """ - As [`error`][logging.Logger.error], but subsequent calls with + """As [`error`][logging.Logger.error], but subsequent calls with the same message are silently dropped. """ if not _should_log_with_scope(scope): @@ -220,8 +216,8 @@ def _configure_vllm_root_logger() -> None: def init_logger(name: str) -> _VllmLogger: """The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root vllm logger has - already been configured.""" - + already been configured. + """ logger = logging.getLogger(name) for method_name, method in _METHODS_TO_PATCH.items(): @@ -310,8 +306,7 @@ def _trace_calls(log_path, root_dir, frame, event, arg=None): def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): - """ - Enable tracing of every function call in code under `root_dir`. + """Enable tracing of every function call in code under `root_dir`. This is useful for debugging hangs or crashes. `log_file_path` is the path to the log file. `root_dir` is the root directory of the code to trace. If None, it is the diff --git a/vllm/logging_utils/access_log_filter.py b/vllm/logging_utils/access_log_filter.py index 5501bd5bc6e8..086449ead845 100644 --- a/vllm/logging_utils/access_log_filter.py +++ b/vllm/logging_utils/access_log_filter.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Access log filter for uvicorn to exclude specific endpoints from logging. +"""Access log filter for uvicorn to exclude specific endpoints from logging. This module provides a logging filter that can be used to suppress access logs for specific endpoints (e.g., /health, /metrics) to reduce log noise in @@ -13,8 +12,7 @@ class UvicornAccessLogFilter(logging.Filter): - """ - A logging filter that excludes access logs for specified endpoint paths. + """A logging filter that excludes access logs for specified endpoint paths. This filter is designed to work with uvicorn's access logger. It checks the log record's arguments for the request path and filters out records @@ -31,6 +29,7 @@ class UvicornAccessLogFilter(logging.Filter): excluded_paths: A list of URL paths to exclude from logging. Paths are matched exactly. Example: ["/health", "/metrics"] + """ def __init__(self, excluded_paths: list[str] | None = None): @@ -38,14 +37,14 @@ def __init__(self, excluded_paths: list[str] | None = None): self.excluded_paths = set(excluded_paths or []) def filter(self, record: logging.LogRecord) -> bool: - """ - Determine if the log record should be logged. + """Determine if the log record should be logged. Args: record: The log record to evaluate. Returns: True if the record should be logged, False otherwise. + """ if not self.excluded_paths: return True @@ -72,8 +71,7 @@ def create_uvicorn_log_config( excluded_paths: list[str] | None = None, log_level: str = "info", ) -> dict: - """ - Create a uvicorn logging configuration with access log filtering. + """Create a uvicorn logging configuration with access log filtering. This function generates a logging configuration dictionary that can be passed to uvicorn's `log_config` parameter. It sets up the access log @@ -89,6 +87,7 @@ def create_uvicorn_log_config( Example: >>> config = create_uvicorn_log_config(["/health", "/metrics"]) >>> uvicorn.run(app, log_config=config) + """ config = { "version": 1, diff --git a/vllm/logging_utils/formatter.py b/vllm/logging_utils/formatter.py index 22301751a666..4a90009266b5 100644 --- a/vllm/logging_utils/formatter.py +++ b/vllm/logging_utils/formatter.py @@ -19,8 +19,7 @@ def __init__(self, fmt, datefmt=None, style="%"): def format(self, record): def shrink_path(relpath: Path) -> str: - """ - Shortens a file path for logging display: + """Shortens a file path for logging display: - Removes leading 'vllm' folder if present. - If path starts with 'v1', keeps the first two and last two levels, @@ -36,8 +35,10 @@ def shrink_path(relpath: Path) -> str: Args: relpath (Path): The relative path to be shortened. + Returns: str: The shortened path string for display. + """ parts = list(relpath.parts) new_parts = [] diff --git a/vllm/logging_utils/log_time.py b/vllm/logging_utils/log_time.py index 9e94f463711d..c390bc95a3bb 100644 --- a/vllm/logging_utils/log_time.py +++ b/vllm/logging_utils/log_time.py @@ -1,16 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Provides a timeslice logging decorator -""" +"""Provides a timeslice logging decorator.""" import functools import time def logtime(logger, msg=None): - """ - Logs the execution time of the decorated function. + """Logs the execution time of the decorated function. Always place it beneath other decorators. """ diff --git a/vllm/logprobs.py b/vllm/logprobs.py index a82219920675..c4f3340a659a 100644 --- a/vllm/logprobs.py +++ b/vllm/logprobs.py @@ -17,6 +17,7 @@ class Logprob: logprob: The logprob of chosen token rank: The vocab rank of chosen token (>=1) decoded_token: The decoded chosen token index + """ logprob: float @@ -29,8 +30,7 @@ class Logprob: @dataclass class FlatLogprobs(MutableSequence[LogprobsOnePosition | None]): - """ - Flat logprobs of a request into multiple primitive type lists. + """Flat logprobs of a request into multiple primitive type lists. Compared to list[dict[int, Logprob]], this data structure reduced GC overhead significantly. As it flattened logprob information for @@ -61,7 +61,7 @@ class FlatLogprobs(MutableSequence[LogprobsOnePosition | None]): decoded_tokens: list[str | None] = field(default_factory=list) def append(self, logprobs_one_position: LogprobsOnePosition | None) -> None: - """Appends the container with logprobs for the next position""" + """Appends the container with logprobs for the next position.""" self.start_indices.append(len(self.logprobs)) if logprobs_one_position: for token_id, logprob in logprobs_one_position.items(): @@ -78,8 +78,7 @@ def append_fast( ranks: itertools.chain[int], decoded_tokens: Iterable[str | None], ) -> None: - """ - Appends logprobs for the next position without creating + """Appends logprobs for the next position without creating the intermediate logprob dictionary. """ self.start_indices.append(len(self.logprobs)) @@ -93,12 +92,12 @@ def append_fast( self.end_indices.append(len(self.logprobs)) def extend(self, logprobs_multi_positions) -> None: - """Extends the container with logprobs for the next multiple positions""" + """Extends the container with logprobs for the next multiple positions.""" for logprobs_one_position in logprobs_multi_positions: self.append(logprobs_one_position) def __len__(self) -> int: - """Gets number of positions stored in the container""" + """Gets number of positions stored in the container.""" return len(self.start_indices) @overload @@ -108,7 +107,7 @@ def __getitem__(self, position: int) -> LogprobsOnePosition: ... def __getitem__(self, s: slice, /) -> "FlatLogprobs": ... def __getitem__(self, index: int | slice): - """Extracts logprobs of a given position or slice""" + """Extracts logprobs of a given position or slice.""" if isinstance(index, int): return { self.token_ids[i]: Logprob( @@ -149,8 +148,7 @@ def insert(self, index: int, value: dict[int, Logprob] | None) -> None: raise TypeError("Cannot insert logprobs to FlatLogprobs") def __iter__(self) -> Iterator[LogprobsOnePosition]: - """ - Iterates the container and yields LogprobsOnePosition for + """Iterates the container and yields LogprobsOnePosition for each position. """ for i in range(0, len(self.start_indices)): @@ -165,7 +163,7 @@ def __iter__(self) -> Iterator[LogprobsOnePosition]: def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs: - """Creates a container to store prompt logprobs for a request""" + """Creates a container to store prompt logprobs for a request.""" logprobs: PromptLogprobs = FlatLogprobs() if flat_logprobs else [] # NOTE: logprob of first prompt token is None. logprobs.append(None) @@ -173,7 +171,7 @@ def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs: def create_sample_logprobs(flat_logprobs: bool) -> SampleLogprobs: - """Creates a container to store decode logprobs for a request""" + """Creates a container to store decode logprobs for a request.""" return FlatLogprobs() if flat_logprobs else [] @@ -185,7 +183,7 @@ def append_logprobs_for_next_position( rank: int, num_logprobs: int, ) -> None: - """Appends logprobs for the next position""" + """Appends logprobs for the next position.""" if num_logprobs == -1: num_logprobs = len(logprobs) # We do not need a special case for the sampled token diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index df7cfe7122da..36767376c44e 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -244,8 +244,7 @@ def _apply_lora_to_output( def _apply_async_impl( self, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - """ - Forward pass with base linear and LoRA on separate CUDA streams + """Forward pass with base linear and LoRA on separate CUDA streams for overlap, using maybe_execute_in_parallel. Base layer runs on default stream; LoRA runs on aux stream. """ diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index ea776ec2ea4f..6075629c453c 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -83,8 +83,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"): class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA): - """ - LoRA on top of ColumnParallelLinear layer. + """LoRA on top of ColumnParallelLinear layer. LoRA B is sliced for tensor parallelism. There are two types for the `base_layer`: 1. ColumnParallelLinear, e.g.`dense_h_to_4h` in `FalconForCausalLM`. @@ -140,6 +139,7 @@ def forward( Returns: - output - bias + """ bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None @@ -210,8 +210,7 @@ def create_lora_weights( lora_config: LoRAConfig, model_config: PretrainedConfig | None = None, ) -> None: - """ - The main reason for overriding this function is to enhance code + """The main reason for overriding this function is to enhance code maintainability. """ self.lora_config = lora_config @@ -268,8 +267,7 @@ def expand_packed_lora( lora_a: list[torch.Tensor | None], lora_b: list[torch.Tensor | None], ) -> tuple[list[torch.Tensor | None], list[torch.Tensor | None]]: - """ - Expand packed adapter groups when they don't match n_slices. + """Expand packed adapter groups when they don't match n_slices. E.g. in_proj_qkv (covers Q+K+V) + in_proj_z. A None group member means that member was not adapted; the slice(s) @@ -388,8 +386,7 @@ def can_replace_layer( class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): - """ - ColumnParallelLinear layer that is specifically designed for + """ColumnParallelLinear layer that is specifically designed for qkv_proj. Certain models, such as chatglm3 and baichuan-7b, only contains a single LoRA within their qkv_proj layer. @@ -493,8 +490,7 @@ def create_lora_weights( lora_config: LoRAConfig, model_config: PretrainedConfig | None = None, ) -> None: - """ - The main reason for overloading this function is to handle inconsistent + """The main reason for overloading this function is to handle inconsistent weight dimensions in qkv lora. """ super().create_lora_weights(max_loras, lora_config, model_config) @@ -521,8 +517,7 @@ def can_replace_layer( class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA): - """ - Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also. + """Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also. Based on S-LoRA, slicing happens along the rank dim. """ @@ -562,8 +557,7 @@ def can_replace_layer( class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLoRA): - """ - Differs from MergedColumnParallelLinearWithLoRA by slicing the + """Differs from MergedColumnParallelLinearWithLoRA by slicing the LoRA A's also. Based on S-LoRA, slicing happens along the rank dim. @@ -605,8 +599,7 @@ def can_replace_layer( class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA): - """ - Differs from QKVParallelLinearWithLoRA by slicing the + """Differs from QKVParallelLinearWithLoRA by slicing the LoRA A's also. Based on S-LoRA, slicing happens along the rank dim. @@ -641,8 +634,7 @@ def can_replace_layer( class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA): - """ - Differs from MergedQKVParallelLinearWithLoRA by slicing the + """Differs from MergedQKVParallelLinearWithLoRA by slicing the LoRA A's also. Based on S-LoRA, slicing happens along the rank dim. @@ -741,7 +733,8 @@ def set_lora( lora_b: torch.Tensor | list[torch.Tensor], ): """Override to handle single tensor weights - that need to be split into slices.""" + that need to be split into slices. + """ self.reset_lora(index) # Handle case where checkpoint has single tensor weights diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index e47bc79bf469..ed0ccdf7aeff 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -248,7 +248,6 @@ def create_lora_weights( model_config: PretrainedConfig | None = None, ) -> None: """Initializes lora matrices.""" - self._verify_ep_fs(lora_config) self.max_loras = lora_config.max_loras self.fully_sharded = lora_config.fully_sharded_loras @@ -294,9 +293,7 @@ def create_lora_weights( ) def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor: - """ - Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA - """ + """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" if self.tp_size == 1 or not self.fully_sharded: return w13_lora_a @@ -321,9 +318,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b[:, start_idx:end_idx, :] def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: - """ - Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA - """ + """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" if self.tp_size == 1: return w2_lora_a # w2_lora_a shape (num_experts,rank,input_size) @@ -334,9 +329,7 @@ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: return w2_lora_a[:, :, start_idx:end_idx] def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor: - """ - Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA - """ + """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" if self.tp_size == 1 or not self.fully_sharded: return w2_lora_b # Based on S-LoRA, we slice W2 B along the hidden_size dim. @@ -475,7 +468,6 @@ def can_replace_layer( model_config: PretrainedConfig | None = None, ) -> bool: """Returns True if the layer can be replaced by this LoRA layer.""" - # source_layer is MoERunner moe_cls = maybe_get_oot_by_class(MoERunner) return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2 @@ -522,7 +514,6 @@ def create_lora_weights( model_config: PretrainedConfig | None = None, ) -> None: """Initializes lora matrices.""" - if model_config is None: raise ValueError("model_config must be provided for MoE LoRA.") architectures = model_config.architectures @@ -609,30 +600,22 @@ def set_lora( @property def w13_input_size(self): - """ - Full size - """ + """Full size.""" return self.w13_lora_a_stacked[0].shape[-1] @property def w13_output_size(self): - """ - Full size - """ + """Full size.""" return self.w13_lora_b_stacked[0].shape[-2] * self.tp_size @property def w2_input_size(self): - """ - Full size - """ + """Full size.""" return self.w2_lora_a_stacked[0].shape[-1] * self.tp_size @property def w2_output_size(self): - """ - Full size - """ + """Full size.""" return self.hidden_size @classmethod diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 0199437ba22d..7613e9d165b6 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -18,8 +18,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): - """ - LoRA wrapper for LogitsProcessor, with extra logic to handle the + """LoRA wrapper for LogitsProcessor, with extra logic to handle the application of the LoRA adapter. Args: @@ -30,6 +29,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): sharded_to_full_mapping: index mapping from sharded vocab to full vocab received from base_layer.get_sharded_to_full_mapping(). If None, no reindexing will be done. + """ def __init__( diff --git a/vllm/lora/layers/replicated_linear.py b/vllm/lora/layers/replicated_linear.py index 53ae26be4c36..f62dc60b16a4 100644 --- a/vllm/lora/layers/replicated_linear.py +++ b/vllm/lora/layers/replicated_linear.py @@ -25,7 +25,7 @@ def __init__(self, base_layer: ReplicatedLinear) -> None: def forward( self, input_: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: - """Forward of ReplicatedLinearWithLoRA + """Forward of ReplicatedLinearWithLoRA. Args: input_: Tensor whose last dimension is `input_size`. @@ -33,6 +33,7 @@ def forward( Returns: - output - bias + """ bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None diff --git a/vllm/lora/layers/row_parallel_linear.py b/vllm/lora/layers/row_parallel_linear.py index 3d3c7f1d279c..8f607dfbe508 100644 --- a/vllm/lora/layers/row_parallel_linear.py +++ b/vllm/lora/layers/row_parallel_linear.py @@ -42,7 +42,7 @@ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor: def forward( self, input_: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: - """Forward of RowParallelLinear + """Forward of RowParallelLinear. Args: input_: tensor whose last dimension is `input_size`. If @@ -52,6 +52,7 @@ def forward( Returns: - output - bias + """ # set up backprop all-reduce. if self.base_layer.input_is_parallel: @@ -99,8 +100,7 @@ def can_replace_layer( class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA): - """ - Differs from RowParallelLinearWithLoRA by slicing the + """Differs from RowParallelLinearWithLoRA by slicing the LoRA B's also. Based on S-LoRA, slicing happens along the output dim. diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index b4cde2564b39..e725da972349 100644 --- a/vllm/lora/layers/utils.py +++ b/vllm/lora/layers/utils.py @@ -74,8 +74,7 @@ def _get_lora_device(base_layer: nn.Module) -> torch.device: def _not_fully_sharded_can_replace(can_replace): - """ - decorator which adds the condition of not using fully sharded loras + """Decorator which adds the condition of not using fully sharded loras intended to wrap can_replace_layer() """ @@ -88,8 +87,7 @@ def dec(*args, **kwargs): def _fully_sharded_can_replace(can_replace): - """ - decorator which adds the condition of fully sharded loras + """Decorator which adds the condition of fully sharded loras intended to wrap can_replace_layer() """ diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 609f14f85e72..6286533a0fde 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -39,9 +39,7 @@ class MoEEPLoadSpec: def _is_remote_expert_key(raw_name: str, spec: "MoEEPLoadSpec") -> bool: - """ - Decide whether a checkpoint key belongs to a non-local expert. - """ + """Decide whether a checkpoint key belongs to a non-local expert.""" pos = raw_name.find(_EXPERTS_SEPARATOR) if pos < 0: return False @@ -69,15 +67,14 @@ def __init__( *, modules_to_save: dict[str, LoRAFullModuleWeights] | None = None, ) -> None: - """ - Args: - lora_model_id: The integer id for the lora model. - rank: lora rank. - loras: module name -> weights for lora-replaced layers. - is_3d_lora_weight: Whether the on-disk MoE adapter is in the 3D - fused (gate_up_proj / down_proj) layout. Propagated from the - originating LoRARequest. Only consulted by the LoRA model - manager when enable_mixed_moe_lora_format is on. + """Args: + lora_model_id: The integer id for the lora model. + rank: lora rank. + loras: module name -> weights for lora-replaced layers. + is_3d_lora_weight: Whether the on-disk MoE adapter is in the 3D + fused (gate_up_proj / down_proj) layout. Propagated from the + originating LoRARequest. Only consulted by the LoRA model + manager when enable_mixed_moe_lora_format is on. """ self.id = lora_model_id @@ -93,7 +90,8 @@ def __init__( def clone(self, lora_model_id: int) -> "LoRAModel": """Return a copy of the object with different ids. - Will share the underlying tensors.""" + Will share the underlying tensors. + """ return self.__class__( lora_model_id, rank=self.rank, @@ -103,7 +101,7 @@ def clone(self, lora_model_id: int) -> "LoRAModel": ) def get_lora(self, module_name: str) -> LoRALayerWeights | None: - """Get LoRA for a given module by name""" + """Get LoRA for a given module by name.""" return self.loras.get(module_name, None) def get_module_to_save(self, module_name: str) -> LoRAFullModuleWeights | None: @@ -114,7 +112,7 @@ def check_lora_name(self, lora_name: str) -> bool: @staticmethod def _should_skip_module(module_name: str, skip_prefixes: list[str]) -> bool: - """Check if a module should be skipped based on skip prefixes""" + """Check if a module should be skipped based on skip prefixes.""" for prefix in skip_prefixes: if f".{prefix}" in module_name or module_name.startswith(prefix): return True @@ -222,6 +220,12 @@ def from_local_checkpoint( a global counter. device: Device where the lora model is loaded. dtype: dtype of the lora model weights. + model_vocab_size: Vocab size of the base model, used to size the + embedding deltas. + weights_mapper: Optional mapper rewriting checkpoint weight names + to vLLM names. + tensorizer_config_dict: Optional tensorizer config used to load + the checkpoint via tensorizer instead of from disk. skip_prefixes: List of module name prefixes to skip during loading. Models can define this to skip modules not used in inference (e.g., MTP layers). Format: ["mtp."] @@ -233,6 +237,7 @@ def from_local_checkpoint( Returns: Loaded LoRA Model. + """ lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors") lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin") diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index d9b200f0bba5..13a692a62353 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -163,7 +163,6 @@ def pack_moe( If LoRA is None, it signifies that the submodule does not have a LoRA. """ - first_lora = next(lora for lora in loras if lora is not None) assert first_lora is not None rank = first_lora.rank diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 2f3b069f0774..b47c22d6aab9 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -97,6 +97,9 @@ def __init__( in a single batch. vocab_size: the vocab size of the model. lora_config: the LoRA configuration. + device: the device the LoRA tensors are placed on. + vllm_config: the vLLM config for this engine. + """ self.model: SupportsLoRAModel = model self.supported_lora_modules = get_supported_lora_modules(self.model) @@ -785,6 +788,7 @@ def _match_target_modules(self, module_name: str) -> bool: Returns: True if the module passes the filter, False otherwise. + """ return is_in_target_modules( module_name, @@ -793,9 +797,7 @@ def _match_target_modules(self, module_name: str) -> bool: ) def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None: - """ - Determine whether this module supports LoRA and which wrapper to use. - """ + """Determine whether this module supports LoRA and which wrapper to use.""" # For language model (early return) if not self.supports_mm: return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] @@ -1190,9 +1192,7 @@ def _restrict_to_local_experts( return new_module_names[start:end] def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None: - """ - Per-rank slicing metadata for 2D RoutedEXperts LoRA modules. - """ + """Per-rank slicing metadata for 2D RoutedEXperts LoRA modules.""" if not self._use_ep or not self._is_moe: return None module = next( diff --git a/vllm/lora/ops/triton_ops/fp8_kernel_utils.py b/vllm/lora/ops/triton_ops/fp8_kernel_utils.py index 8429562c7621..b9a7acf8627b 100644 --- a/vllm/lora/ops/triton_ops/fp8_kernel_utils.py +++ b/vllm/lora/ops/triton_ops/fp8_kernel_utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Utilities for Punica kernel construction. -""" +"""Utilities for Punica kernel construction.""" from vllm.triton_utils import tl, triton @@ -21,8 +19,7 @@ def _accumulate_mm( group_n: tl.constexpr, use_fp8_w8a8: tl.constexpr, ): - """ - Core matrix multiplication and accumulation logic with quantization support. + """Core matrix multiplication and accumulation logic with quantization support. Args: tiled_a (tl.tensor): Loaded tile from A matrix @@ -36,8 +33,8 @@ def _accumulate_mm( group_k: Block size for K dimension in block-wise quantization group_n: Block size for N dimension in block-wise quantization use_fp8_w8a8: Whether using FP8 W8A8 quantization - """ + """ if use_fp8_w8a8: if group_k > 0 and group_n > 0: # Block-wise quantization: scales are loaded per block @@ -85,8 +82,7 @@ def fp8_mm_k( USE_GDC: tl.constexpr, base_k, ): - """ - FP8-compatible matrix multiplication kernel with quantization support. + """FP8-compatible matrix multiplication kernel with quantization support. Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of B (k x n), iterate through the K dimension to compute the partial/complete matrix block product with proper dequantization. @@ -120,6 +116,7 @@ def fp8_mm_k( b_dtype: datatype of the B matrix USE_GDC: Whether to use PDL. True indicates use. base_k (int): Base offset along K dimension for current SPLIT_K group + """ accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) @@ -238,13 +235,11 @@ def do_shrink_kernel_fp8( per_channel_quant: tl.constexpr, launch_pdl: tl.constexpr, ): - """ - Given an array of integers that identifies the rows of A, ram, + """Given an array of integers that identifies the rows of A, ram, a lora index that identifies which LoRA to use from lora_ptr, lora_index, a slice_id that identifies the input/output slice, compute the matrix product and store in the appropriate output location. """ - # Identify the lora_ptr from slice_id. if SLICE_NUM == 1: cur_lora_ptr = lora_ptr @@ -432,8 +427,7 @@ def do_expand_kernel_fp8( use_fp8_w8a8: tl.constexpr, per_channel_quant: tl.constexpr, ): - """ - FP8-compatible expand kernel for LoRA. + """FP8-compatible expand kernel for LoRA. Given an array of integers that identifies the rows of A, ram, a lora index that identifies which LoRA to use from lora_ptr, lora_index, a slice_id that identifies the input/output slice, @@ -449,7 +443,6 @@ def do_expand_kernel_fp8( - Per-channel quantization - Tensor-wise quantization """ - # ls_d*_ptr can be either an integer or a pointer if SAME_STRIDE: cur_lora_d0_stride = ls_d0_ptr diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py index 033ed1df6ced..780d4a0d70bf 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py +++ b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py @@ -24,7 +24,7 @@ def _get_lora_id( top_k_num, naive_block_assignment: tl.constexpr, ): - """Returns lora_id""" + """Returns lora_id.""" if naive_block_assignment: token_idx = pid_m // top_k_num return tl.load(token_lora_mapping_ptr + token_idx) @@ -41,7 +41,7 @@ def _get_expert_id( max_loras, naive_block_assignment: tl.constexpr, ): - """Returns expert_id""" + """Returns expert_id.""" if naive_block_assignment: return tl.load(expert_ids_ptr + pid_m) else: @@ -61,7 +61,7 @@ def _get_token_offs( naive_block_assignment: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, ): - """Returns token offsets""" + """Returns token offsets.""" if naive_block_assignment: return tl.where(offs == 0, pid_m, num_valid_tokens) else: @@ -76,8 +76,7 @@ def _get_token_offs( def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device): - """ - `_LORA_PTR_DICT` collects the required information during `profile_run`, + """`_LORA_PTR_DICT` collects the required information during `profile_run`, After this, it remains constant and subsequent usage is through LUT. Refer to: https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py @@ -101,9 +100,7 @@ def _adjust_kernel_inputs( sorted_token_ids: torch.Tensor | None, expert_ids: torch.Tensor, ): - """ - helper function to adjust kernel inputs when sorted_token_ids is None - """ + """Helper function to adjust kernel inputs when sorted_token_ids is None.""" if sorted_token_ids is None: stride_tl = 0 stride_el = 0 diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py index 1bde7e084533..f6ef6974dcf7 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py +++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py @@ -25,7 +25,7 @@ def _get_lora_id( top_k_num, naive_block_assignment: tl.constexpr, ): - """Returns lora_id""" + """Returns lora_id.""" if naive_block_assignment: token_idx = pid_m // top_k_num return tl.load(token_lora_mapping_ptr + token_idx) @@ -42,7 +42,7 @@ def _get_expert_id( max_loras, naive_block_assignment: tl.constexpr, ): - """Returns expert_id""" + """Returns expert_id.""" if naive_block_assignment: return tl.load(expert_ids_ptr + pid_m) else: @@ -62,7 +62,7 @@ def _get_token_offs( naive_block_assignment: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, ): - """Returns token offsets""" + """Returns token offsets.""" if naive_block_assignment: return tl.where(offs == 0, pid_m, num_valid_tokens) else: @@ -852,8 +852,7 @@ def _run_fused_moe_lora_small_batch( def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device): - """ - `_LORA_PTR_DICT` collects the required information during `profile_run`, + """`_LORA_PTR_DICT` collects the required information during `profile_run`, After this, it remains constant and subsequent usage is through LUT. Refer to: https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py @@ -877,9 +876,7 @@ def _adjust_kernel_inputs( sorted_token_ids: torch.Tensor | None, expert_ids: torch.Tensor, ): - """ - helper function to adjust kernel inputs when sorted_token_ids is None - """ + """Helper function to adjust kernel inputs when sorted_token_ids is None.""" if sorted_token_ids is None: stride_tl = 0 stride_el = 0 diff --git a/vllm/lora/ops/triton_ops/kernel_utils.py b/vllm/lora/ops/triton_ops/kernel_utils.py index 34a9939ffcd5..fd6ae4442869 100644 --- a/vllm/lora/ops/triton_ops/kernel_utils.py +++ b/vllm/lora/ops/triton_ops/kernel_utils.py @@ -1,19 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Utilities for Punica kernel construction. -""" +"""Utilities for Punica kernel construction.""" from vllm.triton_utils import tl, triton @triton.jit def mm_k( - a_ptr, - b_ptr, - ak_stride, - bk_stride, - offset_k, + a_ptr: tl.tensor, + b_ptr: tl.tensor, + ak_stride: int, + bk_stride: int, + offset_k: tl.tensor, K: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, @@ -23,32 +21,34 @@ def mm_k( CAST_TYPE: tl.constexpr, b_dtype: tl.constexpr, USE_GDC: tl.constexpr, - base_k, + base_k: int, ): - """ - Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of + """Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of B (k x n), iterate, through the K dimension to compute the partial/complete matrix block product. If SPLIT_K == 1, the output m x n product is complete. If SPLIT_K > 1, the thread block computes partial outputs. The partial outputs are then atomically summed in the caller code. + Args: a_ptr: Array of pointers, identifying rows of A b_ptr: Array of pointers, identifying columns of B ak_stride: K dimension stride of the A matrix bk_stride: K dimension stride of the B matrix + offset_k: Offsets within the current BLOCK_K tile along K K: Length of the K dimension BLOCK_M: M dimension of the output block m x n BLOCK_N: N dimension of the output block m x n BLOCK_K: K dimension atom EVEN_K: True if the blocks of A and B can be loaded without any - masking. + masking. SPLIT_K: Parameter signifying parallelism in the K dimension. CAST_TYPE: if True, cast the values from the A matrix to the B - matrix dtype. + matrix dtype. b_dtype: datatype of the B matrix USE_GDC: Whether to use PDL. True indicates use. base_k: Base offset along K dimension for current SPLIT_K group + """ accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) @@ -142,15 +142,13 @@ def do_expand_kernel( ADD_INPUTS: tl.constexpr, USE_GDC: tl.constexpr, ): - """ - Given an array of integers that identifies the rows of A, ram, + """Given an array of integers that identifies the rows of A, ram, a lora index that identifies which LoRA to use from lora_ptr, lora_index, a slice_id that identifies the input/output slice, compute the matrix product and store in the appropriate output location. Given that this is an expand kernel, we don't perform any split-K reduction as the K dimension is assumed to be small. """ - # ls_d*_ptr can be either an integer or a pointer if SAME_STRIDE: # integer @@ -268,13 +266,11 @@ def do_shrink_kernel( SLICE_NUM: tl.constexpr, USE_GDC: tl.constexpr, ): - """ - Given an array of integers that identifies the rows of A, ram, + """Given an array of integers that identifies the rows of A, ram, a lora index that identifies which LoRA to use from lora_ptr, lora_index, a slice_id that identifies the input/output slice, compute the matrix product and store in the appropriate output location. """ - # Identify the lora_ptr from slice_id. if SLICE_NUM == 1: # current lora ptr diff --git a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py index 2dc33965e327..9369fb750b2a 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 @@ -21,8 +20,7 @@ def _get_expand_lora_scale_ptr(lora_weights: list[torch.Tensor], device: torch.device): - """ - `_EXPAND_LORA_SCALE_PTR_DICT` collects the required information during + """`_EXPAND_LORA_SCALE_PTR_DICT` collects the required information during `profile_run`, After this, it remains constant and subsequent usage is through LUT. Refer to: @@ -91,9 +89,7 @@ def _lora_expand_kernel_fp8( per_channel_quant: tl.constexpr, launch_pdl: tl.constexpr, ): - """ - FP8-compatible expand kernel wrapper. - """ + """FP8-compatible expand kernel wrapper.""" cta_n_num = tl.cdiv(N, BLOCK_N) cta_m_num = tl.cdiv(M, BLOCK_M) @@ -192,8 +188,7 @@ def _lora_expand_fp8( use_fp8_w8a8: bool = False, per_channel_quant: bool = False, ) -> None: - """ - FP8-compatible LoRA expand operation. + """FP8-compatible LoRA expand operation. Args: inputs: Input tensor from shrink operation [num_slices, num_tokens, lora_rank] @@ -206,6 +201,8 @@ def _lora_expand_fp8( num_tokens_per_lora: Number of tokens per LoRA lora_token_start_loc: Start location for each LoRA's tokens lora_ids: LoRA IDs to process + num_active_loras: Number of active LoRAs. Accepted for API parity with + the non-FP8 kernel and unused here. no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates if there are any requests that require LoRA. offset_start (int, optional): Offset start for output_tensor. @@ -216,6 +213,7 @@ def _lora_expand_fp8( group_n (int, optional): Block size for N in block-wise quantization. use_fp8_w8a8 (bool, optional): Whether to use FP8 W8A8 quantization. per_channel_quant (bool, optional): Whether to use per-channel quantization. + """ assert no_lora_flag_cpu.numel() == 1 if no_lora_flag_cpu.item(): diff --git a/vllm/lora/ops/triton_ops/lora_expand_op.py b/vllm/lora/ops/triton_ops/lora_expand_op.py index 40f640e2431d..184e58247bb3 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_op.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 @@ -147,32 +146,31 @@ def _lora_expand( offset_start: int = 0, add_inputs: bool = False, ) -> None: - """ - Args: - inputs (torch.Tensor): input tensor - lora_b_weights (list[torch.Tensor]): lora'b weight - output_tensor (torch.Tensor): output tensor - token_lora_mapping (torch.Tensor): A tensor mapping each input token - to the lora-id related to that token. A value of -1 indicates that - LoRA doesn't apply to that token. - token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from - the A matrix grouped by LoRA IDs. - num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number - of tokens that are to be processed by LoRA ID lora_ids[i] - lora_token_start_loc (torch.Tensor): A cumulative sum of - num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that - lora_token_start_loc[i], along with num_tokens_per_lora[i] - identifies the region in token_indices_sorted_by_lora_ids that - LoRA lora_ids[i] should process. - lora_ids (torch.Tensor): LoRA ids to process. - no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates - if there are any requests that require LoRA. - offset_start (int, optional): Offset start for output_tensor. - Defaults to 0. - add_inputs (bool, optional): Whether to add the input tensor to the - output tensor. Defaults to False. - """ + """Args: + inputs (torch.Tensor): input tensor + lora_b_weights (list[torch.Tensor]): lora'b weight + output_tensor (torch.Tensor): output tensor + token_lora_mapping (torch.Tensor): A tensor mapping each input token + to the lora-id related to that token. A value of -1 indicates that + LoRA doesn't apply to that token. + token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from + the A matrix grouped by LoRA IDs. + num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number + of tokens that are to be processed by LoRA ID lora_ids[i] + lora_token_start_loc (torch.Tensor): A cumulative sum of + num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that + lora_token_start_loc[i], along with num_tokens_per_lora[i] + identifies the region in token_indices_sorted_by_lora_ids that + LoRA lora_ids[i] should process. + lora_ids (torch.Tensor): LoRA ids to process. + no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates + if there are any requests that require LoRA. + offset_start (int, optional): Offset start for output_tensor. + Defaults to 0. + add_inputs (bool, optional): Whether to add the input tensor to the + output tensor. Defaults to False. + """ assert no_lora_flag_cpu.numel() == 1 if no_lora_flag_cpu.item(): # None of the inputs require LoRA. diff --git a/vllm/lora/ops/triton_ops/lora_kernel_metadata.py b/vllm/lora/ops/triton_ops/lora_kernel_metadata.py index dd7c2c706a07..e4bd7bfacdf2 100644 --- a/vllm/lora/ops/triton_ops/lora_kernel_metadata.py +++ b/vllm/lora/ops/triton_ops/lora_kernel_metadata.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -LoRA kernels metadata preparation utilities. -""" +"""LoRA kernels metadata preparation utilities.""" import bisect from dataclasses import dataclass, field @@ -107,14 +105,13 @@ def _reset(self): self.num_active_loras_cpu.fill_(0) def prepare_tensors(self, token_lora_mapping: torch.Tensor) -> None: - """ - Prepare kernel metadata tensors for the current forward pass. + """Prepare kernel metadata tensors for the current forward pass. Args: token_lora_mapping (torch.Tensor): Tensor containing lora indices for each input token. - """ + """ self._reset() # Check and record no-lora case. @@ -180,8 +177,7 @@ def meta_args( torch.Tensor, torch.Tensor, ]: - """ - This function returns the kernel metadata required for the current + """This function returns the kernel metadata required for the current forward pass execution of the kernel. The function returns all the metadata required by the kernel, in order, as a tuple, so it can be unpacked directly during the lora_shrink/lora_expand function call. @@ -189,6 +185,10 @@ def meta_args( Args: token_nums (int): Number of input tokens in the current forward pass of the kernel. + specialize_active_lora (bool): If True, report the actual active + LoRA count; otherwise report the padded default, which keeps + the kernel launch shape stable. + """ if specialize_active_lora: num_active_loras = self.num_active_loras_cpu diff --git a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py index 9b4fe75f3c04..c2b852fa40ab 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 @@ -20,8 +19,7 @@ def _get_shrink_lora_scale_ptr( lora_scale_weights: list[torch.Tensor], device: torch.device ): - """ - `_SHRINK_LORA_SCALE_PTR_DICT` collects the required information during + """`_SHRINK_LORA_SCALE_PTR_DICT` collects the required information during `profile_run`. After this, it remains constant and subsequent usage is through LUT. @@ -241,23 +239,23 @@ def _lora_shrink_fp8( use_fp8_w8a8: bool = False, per_channel_quant: bool = False, ) -> None: - """ - Args: - inputs: FP8 or FP16/BF16 input tensor [num_tokens, hidden_size] - lora_a_weights: List of FP8 or FP16/BF16 LoRA A weights per slice - output_tensor: Output tensor (FP16/BF16/FP32) - token_lora_mapping: Token to LoRA ID mapping - token_indices_sorted_by_lora_ids: Sorted token indices - num_tokens_per_lora: Number of tokens per LoRA - lora_token_start_loc: Start location for each LoRA's tokens - lora_ids: LoRA IDs to process - scaling: LoRA scaling factor - a_scale: Activation quantization scales - b_scale: Weight quantization scales per slice - group_k: Block size for K dimension quantization - group_n: Block size for N dimension quantization - use_fp8_w8a8: Whether to use FP8 weights and activations - per_channel_quant: Whether to use per-channel quantization + """Args: + inputs: FP8 or FP16/BF16 input tensor [num_tokens, hidden_size] + lora_a_weights: List of FP8 or FP16/BF16 LoRA A weights per slice + output_tensor: Output tensor (FP16/BF16/FP32) + token_lora_mapping: Token to LoRA ID mapping + token_indices_sorted_by_lora_ids: Sorted token indices + num_tokens_per_lora: Number of tokens per LoRA + lora_token_start_loc: Start location for each LoRA's tokens + lora_ids: LoRA IDs to process + scaling: LoRA scaling factor + a_scale: Activation quantization scales + b_scale: Weight quantization scales per slice + group_k: Block size for K dimension quantization + group_n: Block size for N dimension quantization + use_fp8_w8a8: Whether to use FP8 weights and activations + per_channel_quant: Whether to use per-channel quantization + """ assert no_lora_flag_cpu.numel() == 1 if no_lora_flag_cpu.item(): diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py index d05c8f7c2a52..ec29065b9acf 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 @@ -142,32 +141,31 @@ def _lora_shrink( num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs scaling: float, ) -> None: - """ - Args: - inputs (torch.Tensor): Input tensor - lora_a_weights (list[torch.Tensor]): LoRA weights - output_tensor (torch.Tensor): output tensor - token_lora_mapping (torch.Tensor): A tensor mapping each input token - to the lora-id related to that token. A value of -1 indicates that - LoRA doesn't apply to that token. - token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from - the A matrix grouped by LoRA IDs. - num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number - of tokens that are to be processed by LoRA ID lora_ids[i] - lora_token_start_loc (torch.Tensor): A cumulative sum of - num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that - lora_token_start_loc[i], along with num_tokens_per_lora[i] - identifies the region in token_indices_sorted_by_lora_ids that - LoRA lora_ids[i] should process. - lora_ids (torch.Tensor): LoRA ids to process. - no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates - if there are any requests that require LoRA. - num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the - number of active LoRAs. Stored as a tensor (not int) so - torch.compile treats it as dynamic rather than a constant. - scaling (float): Scaling factor. - """ + """Args: + inputs (torch.Tensor): Input tensor + lora_a_weights (list[torch.Tensor]): LoRA weights + output_tensor (torch.Tensor): output tensor + token_lora_mapping (torch.Tensor): A tensor mapping each input token + to the lora-id related to that token. A value of -1 indicates that + LoRA doesn't apply to that token. + token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from + the A matrix grouped by LoRA IDs. + num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number + of tokens that are to be processed by LoRA ID lora_ids[i] + lora_token_start_loc (torch.Tensor): A cumulative sum of + num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that + lora_token_start_loc[i], along with num_tokens_per_lora[i] + identifies the region in token_indices_sorted_by_lora_ids that + LoRA lora_ids[i] should process. + lora_ids (torch.Tensor): LoRA ids to process. + no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates + if there are any requests that require LoRA. + num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the + number of active LoRAs. Stored as a tensor (not int) so + torch.compile treats it as dynamic rather than a constant. + scaling (float): Scaling factor. + """ assert no_lora_flag_cpu.numel() == 1 if no_lora_flag_cpu.item(): # None of the inputs require LoRA. diff --git a/vllm/lora/ops/triton_ops/utils.py b/vllm/lora/ops/triton_ops/utils.py index dd4c0d2097cf..fbcef91002d2 100644 --- a/vllm/lora/ops/triton_ops/utils.py +++ b/vllm/lora/ops/triton_ops/utils.py @@ -23,8 +23,7 @@ def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device): - """ - `_LORA_A_PTR_DICT` collects the required information during `profile_run`, + """`_LORA_A_PTR_DICT` collects the required information during `profile_run`, After this, it remains constant and subsequent usage is through LUT. Refer to: https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py @@ -75,14 +74,12 @@ def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device): def _get_lora_b_ptr( lora_weights: list[torch.Tensor], offset_start: int, device: torch.device ): - """ - `_LORA_B_PTR_DICT` collects the required information during `profile_run`, + """`_LORA_B_PTR_DICT` collects the required information during `profile_run`, After this, it remains constant and subsequent usage is through LUT. Refer to: https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py """ - key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights) if values := _LORA_B_PTR_DICT.get(key): return values @@ -317,9 +314,7 @@ def get_lora_op_configs( @lru_cache def supports_pdl(device: torch.device | None = None) -> bool: - """ - Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py - """ + """Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py.""" # PDL requires compute capability SM90 or above return ( diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index 8b40af03456e..baffe1698762 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -18,8 +18,7 @@ @dataclass class PEFTHelper: - """ - A helper class for PEFT configurations, specifically designed for LoRA. + """A helper class for PEFT configurations, specifically designed for LoRA. This class handles configuration validation, compatibility checks for various LoRA implementations. """ @@ -40,9 +39,7 @@ class PEFTHelper: vllm_max_position_embeddings: int | None = field(default=False) def _validate_features(self) -> list[str]: - """ - Check if there are any unsupported LoRA features. - """ + """Check if there are any unsupported LoRA features.""" error_msg = [] if self.modules_to_save: unsupported_modules = [ @@ -124,8 +121,7 @@ def from_local_dir( return cls.from_dict(config) def validate_legal(self, lora_config: LoRAConfig) -> None: - """ - Validates the LoRA configuration settings against application + """Validates the LoRA configuration settings against application constraints and requirements. """ error_msg = self._validate_features() diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py index 7ab309515c1a..a0e7751e6ddf 100644 --- a/vllm/lora/punica_wrapper/punica_base.py +++ b/vllm/lora/punica_wrapper/punica_base.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 @@ -20,9 +19,7 @@ class PunicaWrapperABC(ABC): - """ - PunicaWrapper ABC. - """ + """PunicaWrapper ABC.""" @abstractmethod def update_metadata( @@ -33,9 +30,7 @@ def update_metadata( vocab_size: int, **kwargs, ) -> None: - """ - Update the lora-related metadata - """ + """Update the lora-related metadata""" raise NotImplementedError @abstractmethod @@ -47,10 +42,7 @@ def add_shrink( scale: float, **kwargs, ) -> torch.Tensor | None: - """ - Performs GEMM for multiple slices of lora_a. - """ - + """Performs GEMM for multiple slices of lora_a.""" raise NotImplementedError @abstractmethod @@ -64,9 +56,7 @@ def add_expand( add_inputs=True, **kwargs, ) -> torch.Tensor | None: - """ - Performs GEMM for multiple slices of lora_b. - """ + """Performs GEMM for multiple slices of lora_b.""" raise NotImplementedError @abstractmethod @@ -78,8 +68,7 @@ def add_lora_embedding( add_inputs: bool = True, **kwargs, ) -> torch.Tensor | None: - """ - Applies lora specifically for VocabParallelEmbeddingWithLoRA, + """Applies lora specifically for VocabParallelEmbeddingWithLoRA, and this layer only requires the expand operation. """ raise NotImplementedError @@ -97,10 +86,7 @@ def add_lora_linear( buffer: tuple[torch.Tensor, ...] | None = None, **kwargs, ) -> torch.Tensor | None: - """ - Applicable to linear-related lora. - """ - + """Applicable to linear-related lora.""" raise NotImplementedError @abstractmethod @@ -115,9 +101,7 @@ def add_lora_logits( buffer: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | None: - """ - Applies lora specifically for LogitsProcessorWithLoRA. - """ + """Applies lora specifically for LogitsProcessorWithLoRA.""" raise NotImplementedError @abstractmethod @@ -134,8 +118,7 @@ def apply_lora_full_linear( class PunicaWrapperBase(PunicaWrapperABC): - """ - PunicaWrapperBase is designed to manage and provide metadata for the punica + """PunicaWrapperBase is designed to manage and provide metadata for the punica kernel. The main function is to maintain the state information for Multi-LoRA, and to provide the interface for the punica. """ @@ -249,8 +232,7 @@ def _select_full_linear_output( def prefill_metadata( self, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]: - """ - This property provides a convenient way to access the necessary + """This property provides a convenient way to access the necessary metadata for prefill-related kernel computations. 1. seq_start_locs: Tensor of sequence start positions. 2. seq_lengths: Tensor of sequence lengths. @@ -271,8 +253,7 @@ def prefill_metadata( @property def token_lora_indices(self) -> torch.Tensor: - """ - This property provides the lora indices corresponding to each token + """This property provides the lora indices corresponding to each token in the batch. An index of -1 means no lora should be applied. """ token_lora_len = self.indices_len[0] @@ -280,8 +261,7 @@ def token_lora_indices(self) -> torch.Tensor: @property def sampler_indices(self) -> torch.Tensor: - """ - This property is used to access the lora indices specifically for + """This property is used to access the lora indices specifically for LogitsProcessorWithLoRA. """ sampler_indices_len = self.indices_len[1] @@ -289,9 +269,7 @@ def sampler_indices(self) -> torch.Tensor: @property def sampler_indices_padded(self) -> torch.Tensor: - """ - This property provides access to padded sampler indices. - """ + """This property provides access to padded sampler indices.""" indices_padded_len = self.indices_len[2] return self._sampler_indices_padded[:indices_padded_len] @@ -321,8 +299,7 @@ def add_shrink( scale: float, **kwargs, ) -> torch.Tensor | None: - """ - Performs GEMM for multiple slices of lora_a. + """Performs GEMM for multiple slices of lora_a. Semantics: for i in range(len(lora_a_stacked)): @@ -349,8 +326,7 @@ def add_expand( add_inputs=True, **kwargs, ) -> torch.Tensor | None: - """ - Performs GEMM for multiple slices of lora_b. + """Performs GEMM for multiple slices of lora_b. Semantics: offset = offset_start @@ -380,8 +356,7 @@ def add_lora_embedding( add_inputs: bool = True, **kwargs, ) -> torch.Tensor | None: - """ - Applies lora specifically for VocabParallelEmbeddingWithLoRA. + """Applies lora specifically for VocabParallelEmbeddingWithLoRA. and this layer only requires the expand operation. Semantics: y += x @ lora_b_stacked @@ -391,6 +366,7 @@ def add_lora_embedding( x (torch.Tensor): Input tensor. lora_b_stacked (torch.Tensor): lora_b's weights. add_inputs (bool): Default to True. + """ # TODO: implement it based on torch ops raise NotImplementedError @@ -408,8 +384,7 @@ def add_lora_linear( buffer: tuple[torch.Tensor, ...] | None = None, **kwargs, ) -> torch.Tensor | None: - """ - Applicable to linear-related lora. + """Applicable to linear-related lora. Semantics: for i in range(len(lora_a_stacked)): @@ -428,6 +403,7 @@ def add_lora_linear( scale (float): Scaling factor. output_slices (tuple[int, ...]): Every slice's size. buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None. + """ # TODO: implement it based on torch ops raise NotImplementedError @@ -444,8 +420,7 @@ def add_lora_logits( buffer: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | None: - """ - Applies lora specifically for LogitsProcessorWithLoRA. + """Applies lora specifically for LogitsProcessorWithLoRA. Semantics: buffer = (x @ lora_a_stacked) * scale @@ -458,6 +433,7 @@ def add_lora_logits( lora_b_stacked (torch.Tensor):lora_b's weights. scale (float): Scaling factor. buffer (Optional[torch.Tensor]):Default to None. + """ # TODO: implement it based on torch ops raise NotImplementedError @@ -474,8 +450,7 @@ def moe_lora_align_block_size( pad_sorted_ids: bool = False, naive_block_assignment: bool = False, ) -> tuple[torch.Tensor, 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. """ # TODO: implement it based on torch ops @@ -501,8 +476,7 @@ def add_lora_fused_moe( offset: int = 0, token_lora_mapping: torch.Tensor | None = None, ): - """ - Performs a fused forward computation for LoRA of + """Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer. """ # TODO: implement it based on torch ops diff --git a/vllm/lora/punica_wrapper/punica_cpu.py b/vllm/lora/punica_wrapper/punica_cpu.py index d462ce4476d8..436445554fa0 100644 --- a/vllm/lora/punica_wrapper/punica_cpu.py +++ b/vllm/lora/punica_wrapper/punica_cpu.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from typing import Any import torch @@ -20,8 +21,7 @@ # The platforms that are compatible with the PyTorch-native implementation can # inherit this class class PunicaWrapperCPU(PunicaWrapperBase): - """ - PunicaWrapperCPU is designed to manage and provide metadata for the punica + """PunicaWrapperCPU is designed to manage and provide metadata for the punica kernel. The main function is to maintain the state information for Multi-LoRA, and to provide the interface for the pytorch punica ops. """ @@ -133,12 +133,10 @@ def _apply_expand( y_slice_size: int, add_inputs: bool = True, ): - """ - Perform the ` y[:,y_offset:y_offset+y_slice_size]+=x@w_t_all` + """Perform the ` y[:,y_offset:y_offset+y_slice_size]+=x@w_t_all` computation, which is suitable for the GEMM of lora'b. """ - expand_slice_fun: Callable = ( self._expand_slice_prefill if self.is_prefill else self._expand_slice_decode ) @@ -147,8 +145,7 @@ def _apply_expand( def _apply_shrink( self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor, scale: float ): - """ - Perform the ` y+=x@w_t_all` computation, which is suitable for the + """Perform the ` y+=x@w_t_all` computation, which is suitable for the GEMM of lora'a. When `is_prefill is` true, it indicates that it is currently the prefill stage, and the `_shrink_prefill` function should be called. @@ -169,10 +166,9 @@ def add_shrink( x: torch.Tensor, lora_a_stacked: tuple[torch.Tensor, ...], scale: float, - **kwargs, + **kwargs: Any, ): - """ - Performs GEMM for multiple slices of lora_a. + """Performs GEMM for multiple slices of lora_a. When `is_prefill is` true, it indicates that it is currently the prefill stage, and the `_shrink_prefill` function should be called. Otherwise, it is the decode stage, and the _shrink_decode function @@ -187,8 +183,9 @@ def add_shrink( x (torch.Tensor): Input tensor lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights scale (float): Scaling factor for the operation - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ x = x.view(-1, x.shape[-1]) # TODO fuse these kernels for slice_idx in range(len(lora_a_stacked)): @@ -202,10 +199,9 @@ def add_expand( output_slices: tuple[int, ...], offset_start: int = 0, add_inputs=True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Performs GEMM for multiple slices of lora_b. + """Performs GEMM for multiple slices of lora_b. Semantics: for i in range(len(lora_b_stacked)): @@ -218,7 +214,10 @@ def add_expand( x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight output_slices (tuple[int, ...]): Every slice's size + offset_start (int): The starting position of y, defaults to 0. add_inputs (bool): Defaults to True. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ y_org = y y = y.view(-1, y.shape[-1]) @@ -241,10 +240,9 @@ def add_lora_embedding( x: torch.Tensor, lora_b_stacked: torch.Tensor, add_inputs: bool = True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for VocabParallelEmbeddingWithLoRA. + """Applies lora specifically for VocabParallelEmbeddingWithLoRA. Semantics: y += x @ lora_b_stacked @@ -254,8 +252,9 @@ def add_lora_embedding( x (torch.Tensor): Input tensor. lora_b_stacked (torch.Tensor): lora_b's weights. add_inputs (bool): Default to True. - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ # Embedding layer only need expand op expand_fun: Callable = ( self._expand_prefill if self.is_prefill else self._expand_decode @@ -272,10 +271,9 @@ def add_lora_linear( output_slices: tuple[int, ...], *, buffer: tuple[torch.Tensor, ...] | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applicable to linear-related lora. + """Applicable to linear-related lora. Semantics: for i in range(len(lora_a_stacked)): @@ -294,8 +292,9 @@ def add_lora_linear( scale (float): Scaling factor. output_slices (tuple[int, ...]): Every slice's size. buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None. - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices) if buffer is None: @@ -320,10 +319,9 @@ def add_lora_logits( scale, *, buffer: torch.Tensor | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for LogitsProcessorWithLoRA. + """Applies lora specifically for LogitsProcessorWithLoRA. Semantics: buffer = (x @ lora_a_stacked) * scale @@ -336,6 +334,8 @@ def add_lora_logits( lora_b_stacked (torch.Tensor):lora_b's weights. scale (float): Scaling factor. buffer (Optional[torch.Tensor]):Default to None. + **kwargs: Forwarded to `add_shrink` and `add_expand`. + """ y_org = y y = y.view(-1, y.shape[-1]) diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 7deb97e5134e..8604e747f9df 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 """ -from typing import final +from typing import Any, final import torch @@ -32,8 +31,7 @@ @final class PunicaWrapperGPU(PunicaWrapperBase): - """ - PunicaWrapperGPU is designed to manage and provide metadata for the punica + """PunicaWrapperGPU is designed to manage and provide metadata for the punica kernel. The main function is to maintain the state information for Multi-LoRA, and to provide the interface for the punica triton kernel. """ @@ -96,10 +94,9 @@ def add_shrink( x: torch.Tensor, lora_a_stacked: tuple[torch.Tensor, ...], scale: float, - **kwargs, + **kwargs: Any, ): - """ - Performs GEMM for multiple slices of lora_a. + """Performs GEMM for multiple slices of lora_a. Semantics: for i in range(len(lora_a_stacked)): @@ -110,8 +107,9 @@ def add_shrink( x (torch.Tensor): Input tensor lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights scale (float): Scaling factor for the operation - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ x = x.view(-1, x.shape[-1]) lora_shrink( x, @@ -131,10 +129,9 @@ def add_expand( output_slices: tuple[int, ...], offset_start: int = 0, add_inputs=True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Performs GEMM for multiple slices of lora_b. + """Performs GEMM for multiple slices of lora_b. Semantics: for i in range(len(lora_b_stacked)): @@ -147,9 +144,12 @@ def add_expand( x (torch.Tensor): Input tensors lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight output_slices (tuple[int, ...]): Every slice's size + offset_start (int): The starting position of y, defaults to 0. add_inputs (bool): If True, add LoRA output to y; if False, write LoRA-only output to y (used for dual-stream when base and LoRA run on different CUDA streams). Defaults to True. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ y_org = y y = y.view(-1, y.shape[-1]) @@ -177,10 +177,9 @@ def add_lora_embedding( x: torch.Tensor, lora_b_stacked: torch.Tensor, add_inputs: bool = True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for VocabParallelEmbeddingWithLoRA. + """Applies lora specifically for VocabParallelEmbeddingWithLoRA. Semantics: y += x @ lora_b_stacked @@ -190,8 +189,9 @@ def add_lora_embedding( x (torch.Tensor): Input tensor. lora_b_stacked (torch.Tensor): lora_b's weights. add_inputs (bool): Default to True. - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ lora_expand( x.unsqueeze(dim=0), (lora_b_stacked,), @@ -213,10 +213,9 @@ def add_lora_linear( output_slices: tuple[int, ...], *, buffer: torch.Tensor | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applicable to linear-related lora. + """Applicable to linear-related lora. Semantics: for i in range(len(lora_a_stacked)): @@ -226,6 +225,7 @@ def add_lora_linear( @ lora_b_stacked[indices[i], layer_idx, :, :] * scale ).squeeze(0) + Args: y (torch.Tensor): Output tensor. Will be changed in-place. x (torch.Tensor): Input tensor @@ -234,8 +234,10 @@ def add_lora_linear( scale (float): Scaling factor. output_slices (tuple[int, ...]): Every slice's size. buffer (Optional[torch.Tensor]): Defaults to None. - """ + **kwargs: Accepts `add_inputs` (bool, defaults to True); any other + keys are ignored. + """ assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices) assert buffer is None, ( @@ -275,10 +277,9 @@ def add_lora_logits( scale, *, buffer: torch.Tensor | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for LogitsProcessorWithLoRA. + """Applies lora specifically for LogitsProcessorWithLoRA. Semantics: buffer = (x @ lora_a_stacked) * scale @@ -291,6 +292,8 @@ def add_lora_logits( lora_b_stacked (torch.Tensor): lora_b's weights. scale (float): Scaling factor. buffer (Optional[torch.Tensor]): Default to None. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ y_org = y y = y.view(-1, y.shape[-1]) @@ -367,8 +370,7 @@ def moe_lora_align_block_size( naive_block_assignment: bool = False, token_lora_mapping: torch.Tensor | None = None, ) -> tuple[torch.Tensor, 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. When `token_lora_mapping` is provided, it overrides the global mapping @@ -467,9 +469,7 @@ def add_lora_fused_moe( token_lora_mapping: torch.Tensor | None = None, add_inputs: bool = True, ): - """ - Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer. - """ + """Perform a fused forward computation for a LoRA MoE layer.""" ( token_lora_mapping_meta, _, diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py index 3d1cf29683fc..d6af2e36680d 100755 --- a/vllm/lora/punica_wrapper/punica_xpu.py +++ b/vllm/lora/punica_wrapper/punica_xpu.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Based on: +"""Based on: Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). Punica: Multi-Tenant LoRA Serving. https://arxiv.org/abs/2310.18547 """ -from typing import final +from typing import Any, final import torch @@ -29,8 +28,7 @@ @final class PunicaWrapperXPU(PunicaWrapperBase): - """ - PunicaWrapperXPU is designed to manage and provide metadata for the punica + """PunicaWrapperXPU is designed to manage and provide metadata for the punica kernel. The main function is to maintain the state information for Multi-LoRA, and to provide the interface for the punica ipex kernel. """ @@ -135,10 +133,9 @@ def add_shrink( x: torch.Tensor, lora_a_stacked: tuple[torch.Tensor, ...], scale: float, - **kwargs, + **kwargs: Any, ): - """ - Performs GEMM for multiple slices of lora_a. + """Performs GEMM for multiple slices of lora_a. Semantics: for i in range(len(lora_a_stacked)): @@ -149,8 +146,9 @@ def add_shrink( x (torch.Tensor): Input tensor lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights scale (float): Scaling factor for the operation - """ + **kwargs: Unused; accepted for compatibility with the base class signature. + """ x = x.view(-1, x.shape[-1]) for slice_idx in range(len(lora_a_stacked)): self._apply_shrink(y[slice_idx], x, lora_a_stacked[slice_idx], scale) @@ -163,10 +161,9 @@ def add_expand( output_slices: tuple[int, ...], offset_start: int = 0, add_inputs=True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Performs GEMM for multiple slices of lora_b. + """Performs GEMM for multiple slices of lora_b. Semantics: for i in range(len(lora_b_stacked)): @@ -179,7 +176,10 @@ def add_expand( x (torch.Tensor): Input tensors lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight output_slices (tuple[int, ...]): Every slice's size + offset_start (int): The starting position of y, defaults to 0. add_inputs (bool): Defaults to True. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ y_org = y y = y.view(-1, y.shape[-1]) @@ -206,10 +206,9 @@ def add_lora_embedding( x: torch.Tensor, lora_b_stacked: torch.Tensor, add_inputs: bool = True, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for VocabParallelEmbeddingWithLoRA. + """Applies lora specifically for VocabParallelEmbeddingWithLoRA. Semantics: y += x @ lora_b_stacked @@ -219,6 +218,8 @@ def add_lora_embedding( x (torch.Tensor): Input tensor. lora_b_stacked (torch.Tensor): lora_b's weights. add_inputs (bool): Default to True. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ token_lora_indices = self._get_token_lora_indices(x) bgmv_expand(x, lora_b_stacked, y, token_lora_indices, add_inputs) @@ -233,10 +234,9 @@ def add_lora_linear( output_slices: tuple[int, ...], *, buffer: torch.Tensor | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applicable to linear-related lora. + """Applicable to linear-related lora. Semantics: for i in range(len(lora_a_stacked)): @@ -246,6 +246,7 @@ def add_lora_linear( @ lora_b_stacked[indices[i], layer_idx, :, :] * scale ).squeeze(0) + Args: y (torch.Tensor): Output tensor. Will be changed in-place. x (torch.Tensor): Input tensor @@ -254,8 +255,10 @@ def add_lora_linear( scale (float): Scaling factor. output_slices (tuple[int, ...]): Every slice's size. buffer (Optional[torch.Tensor]): Defaults to None. - """ + **kwargs: Accepts `add_inputs` (bool, defaults to True); any other + keys are ignored. + """ assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices) assert buffer is None, ( @@ -287,9 +290,7 @@ def add_lora_linear( @property def sampler_indices_padded(self) -> torch.Tensor: - """ - This property provides access to padded sampler indices. - """ + """This property provides access to padded sampler indices.""" return self._sampler_indices_padded[:] def add_lora_logits( @@ -301,10 +302,9 @@ def add_lora_logits( scale, *, buffer: torch.Tensor | None = None, - **kwargs, + **kwargs: Any, ) -> None: - """ - Applies lora specifically for LogitsProcessorWithLoRA. + """Applies lora specifically for LogitsProcessorWithLoRA. Semantics: buffer = (x @ lora_a_stacked) * scale @@ -317,6 +317,8 @@ def add_lora_logits( lora_b_stacked (torch.Tensor): lora_b's weights. scale (float): Scaling factor. buffer (Optional[torch.Tensor]): Default to None. + **kwargs: Unused; accepted for compatibility with the base class signature. + """ y_org = y y = y.view(-1, y.shape[-1]) @@ -366,8 +368,7 @@ def moe_lora_align_block_size( naive_block_assignment: bool = False, token_lora_mapping: torch.Tensor | None = None, ) -> tuple[torch.Tensor, 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. When `token_lora_mapping` is provided, it overrides the global mapping @@ -466,9 +467,7 @@ def add_lora_fused_moe( token_lora_mapping: torch.Tensor | None = None, add_inputs: bool = True, ): - """ - Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer. - """ + """Perform a fused forward computation for a LoRA MoE layer.""" ( token_lora_mapping_meta, _, diff --git a/vllm/lora/punica_wrapper/utils.py b/vllm/lora/punica_wrapper/utils.py index c29e359f1917..821ba8fc44d9 100644 --- a/vllm/lora/punica_wrapper/utils.py +++ b/vllm/lora/punica_wrapper/utils.py @@ -15,15 +15,13 @@ def compute_meta( token_lora_tensor: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, bool]: - """ - Get the information required for the sgmv kernel. With the features: + """Get the information required for the sgmv kernel. With the features: 1. If consecutive requests in the batch use the same LoRA, this function will combine them into a single request, improving sgmv kernel inference performance. 2. At the beginning of each prefill stage inference, recalculations are needed based on the input, but only once. """ - lora_indices_tensor, seq_length_tensor = torch.unique_consecutive( token_lora_tensor, return_counts=True ) @@ -65,6 +63,7 @@ def convert_mapping( lora_index_to_id: List mapping LoRA ids to LoRA indices. max_loras: Maximum number of LoRAs. vocab_size: Model vocab size. + device: Device the returned index tensors are created on. Returns: A tuple of tensors: @@ -83,6 +82,7 @@ def convert_mapping( indices_len: List of lengths of the above tensors. It contains (base_indices, sampler_indices, sampler_indices_padded, embeddings_indices). + """ index_mapping_indices: list[int] = list(mapping.index_mapping).copy() embedding_indices = index_mapping_indices.copy() diff --git a/vllm/lora/request.py b/vllm/lora/request.py index 7ba1725fa2e8..cd3df54d3b12 100644 --- a/vllm/lora/request.py +++ b/vllm/lora/request.py @@ -10,8 +10,7 @@ class LoRARequest( omit_defaults=True, # type: ignore[call-arg] array_like=True, ): # type: ignore[call-arg] - """ - Request for a LoRA adapter. + """Request for a LoRA adapter. lora_int_id must be globally unique for a given adapter. This is currently not enforced in vLLM. @@ -56,16 +55,14 @@ def path(self): return self.lora_path def __eq__(self, value: object) -> bool: - """ - Overrides the equality method to compare LoRARequest + """Overrides the equality method to compare LoRARequest instances based on lora_name. This allows for identification and comparison lora adapter across engines. """ return isinstance(value, self.__class__) and self.lora_name == value.lora_name def __hash__(self) -> int: - """ - Overrides the hash method to hash LoRARequest instances + """Overrides the hash method to hash LoRARequest instances based on lora_name. This ensures that LoRARequest instances can be used in hash-based collections such as sets and dictionaries, identified by their names across engines. diff --git a/vllm/lora/resolver.py b/vllm/lora/resolver.py index bcfe26467cfb..166cb9ef10fc 100644 --- a/vllm/lora/resolver.py +++ b/vllm/lora/resolver.py @@ -36,6 +36,7 @@ async def resolve_lora( Returns: Optional[LoRARequest]: The resolved LoRA model information, or None if the LoRA model cannot be found. + """ pass @@ -54,9 +55,11 @@ def register_resolver( resolver: LoRAResolver, ) -> None: """Register a LoRA resolver. + Args: resolver_name: Name to register the resolver under. resolver: The LoRA resolver instance to register. + """ if resolver_name in self.resolvers: logger.warning( @@ -70,12 +73,16 @@ def register_resolver( def get_resolver(self, resolver_name: str) -> LoRAResolver: """Get a registered resolver instance by name. + Args: resolver_name: Name of the resolver to get. + Returns: The resolver instance. + Raises: KeyError: If the resolver is not found in the registry. + """ if resolver_name not in self.resolvers: raise KeyError( diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index fba78eddcf25..b1321e9913c2 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -49,8 +49,7 @@ def get_captured_lora_counts(max_loras: int, specialize: bool) -> list[int]: - """ - Returns num_active_loras values for cudagraph capture. + """Returns num_active_loras values for cudagraph capture. When specialize=True: powers of 2 up to max_loras, plus max_loras + 1. When specialize=False: just [max_loras + 1]. @@ -170,7 +169,7 @@ def parse_fine_tuned_lora_name( ) -> tuple[str, bool]: """Parse the name of lora weights. - args: + Args: name: the name of the fine-tuned LoRA, e.g. base_model.model.dense1.weight weights_mapper: maps the name of weight, e.g. @@ -179,8 +178,8 @@ def parse_fine_tuned_lora_name( tuple(module_name, is_lora_a): module_name: the name of the module, e.g. model.dense1, is_lora_a whether the tensor is lora_a or lora_b. - """ + """ # LoRA weight qualified name usually starts with `base_model.model.`, # so we remove the prefix `base_model.model.` to make the following # mapping correctly. @@ -231,10 +230,7 @@ def is_base_embedding_weights(name: str) -> bool: def get_supported_lora_modules(model: nn.Module) -> list[str]: - """ - In vLLM, all linear layers support LoRA. - """ - + """In vLLM, all linear layers support LoRA.""" supported_lora_modules: set[str] = set() for name, module in model.named_modules(): # get the embedding modules if the module's embedding_modules @@ -268,6 +264,7 @@ def is_supported_lora_module( Returns: True if the module is supported, False otherwise. + """ module_suffix = module_name.rsplit(".", 1)[-1] if module_suffix not in supported_lora_modules: @@ -305,6 +302,7 @@ def is_in_target_modules( Returns: True if the module passes the filter, False otherwise. + """ if target_modules is None: return True @@ -329,22 +327,21 @@ def is_in_target_modules( def get_adapter_absolute_path(lora_path: str) -> str: - """ - Resolves the given lora_path to an absolute local path. + """Resolves the given lora_path to an absolute local path. If the lora_path is identified as a Hugging Face model identifier, it will download the model and return the local snapshot path. Otherwise, it treats the lora_path as a local file path and converts it to an absolute path. - Parameters: - lora_path (str): The path to the lora model, which can be an absolute path, - a relative path, or a Hugging Face model identifier. + Args: + lora_path (str): The path to the lora model, which can be an absolute + path, a relative path, or a Hugging Face model identifier. Returns: - str: The resolved absolute local path to the lora model. - """ + str: The resolved absolute local path to the lora model. + """ # Check if the path is an absolute path. Return it no matter exists or not. if os.path.isabs(lora_path): return lora_path diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 93624523e2f8..1fab524e6d08 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -28,7 +28,8 @@ class WorkerLoRAManager: """WorkerLoRAManager that manages LoRA models on the worker side. Every request, the requested LoRAs will be loaded (unless they are already - loaded), and every other LoRA will be unloaded.""" + loaded), and every other LoRA will be unloaded. + """ _manager_cls: type[LoRAModelManager] = LoRAModelManager @@ -74,7 +75,8 @@ def __init__( @contextmanager def dummy_lora_cache(self): """Use this context manager to reuse the dummy lora model - to avoid creating it repeatedly.""" + to avoid creating it repeatedly. + """ self._cached_dummy_lora = None yield self._cached_dummy_lora = False @@ -249,7 +251,8 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager): Uses an LRU Cache. Every request, the requested LoRAs will be loaded (unless they are already loaded) and least recently used LoRAs will - be unloaded if the cache is above capacity.""" + be unloaded if the cache is above capacity. + """ _manager_cls: type[LRUCacheLoRAModelManager] = LRUCacheLoRAModelManager diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index bf98e46e8546..49a71a339b57 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -30,8 +30,7 @@ def maybe_get_oot_by_class(class_type: type) -> type: class PluggableLayer(nn.Module): - """ - Base class for pluggable layers. + """Base class for pluggable layers. A PluggableLayer is a *module-composing* abstraction: it may instantiate other ``torch.nn.Module`` objects as sub-layers, and its functionality depends on @@ -101,8 +100,7 @@ def decorator(layer_cls): class CustomOp(nn.Module): - """ - Base class for custom ops. + """Base class for custom ops. Dispatches the forward method to the appropriate backend. """ @@ -207,8 +205,7 @@ def dispatch_forward(self, compile_native: bool): return self.forward_cuda def maybe_compile(self, fn, *, enable: bool = True): - """ - Compile fn if compilation enabled. + """Compile fn if compilation enabled. Useful for CustomOp instances called from within a torch custom op, meaning the forward call is hidden from the model-level torch.compile. @@ -294,8 +291,7 @@ def enabled(cls) -> bool: @staticmethod def default_on() -> bool: - """ - Behavior controlled by `CompilationConfig.custom_ops`: On by default if + """Behavior controlled by `CompilationConfig.custom_ops`: On by default if 'all', off by default if 'none'. When PyTorch Inductor is used, 'none' is the default value, otherwise 'all'. diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 6c701cfa163c..4d7fdb903dc0 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -231,6 +231,7 @@ def matmul_descriptor_persistent( Returns: Output matrix [M, N] with dtype matching the inputs. + """ assert a.shape[1] == b.shape[0], "Incompatible dimensions" assert a.dtype == b.dtype, "Incompatible dtypes" @@ -378,7 +379,7 @@ def bmm_kernel( B_LARGE: tl.constexpr, C_LARGE: tl.constexpr, ): - """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N) + """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N). Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating along K in a fixed order to preserve batch invariance. @@ -495,8 +496,7 @@ def _log_softmax_kernel( n_cols, BLOCK_SIZE: tl.constexpr, ): - """ - Compute log_softmax along the last dimension of a 2D tensor. + """Compute log_softmax along the last dimension of a 2D tensor. Each block handles one row of the input tensor. """ # Get the row index for this block @@ -550,8 +550,7 @@ def _log_softmax_kernel( def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: - """ - Compute log_softmax using Triton kernel. + """Compute log_softmax using Triton kernel. Args: input: Input tensor @@ -560,6 +559,7 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: Returns: Tensor with log_softmax applied along the specified dimension + """ if dim != -1 and dim != input.ndim - 1: raise ValueError( @@ -607,8 +607,7 @@ def mean_kernel( K, # size after reduction dim BLOCK_SIZE: tl.constexpr, ): - """ - Kernel for computing mean along a single dimension. + """Kernel for computing mean along a single dimension. Input is viewed as (M, N, K) where N is the dimension being reduced. """ # Program ID gives us which output element we're computing @@ -649,8 +648,7 @@ def mean_dim( keepdim: bool = False, dtype: torch.dtype | None = None, ) -> torch.Tensor: - """ - Triton implementation of torch.mean with single dimension reduction. + """Triton implementation of torch.mean with single dimension reduction. Args: input: Input tensor @@ -661,6 +659,7 @@ def mean_dim( Returns: Tensor with mean values along specified dimension + """ # Validate inputs assert -input.ndim <= dim < input.ndim, ( @@ -933,8 +932,7 @@ def _rms_norm_kernel( BLOCK_SIZE: tl.constexpr, HAS_WEIGHT: tl.constexpr, ): - """ - Compute RMS normalization along the last dimension of a 2D tensor. + """Compute RMS normalization along the last dimension of a 2D tensor. RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight Each block handles one row of the input tensor. """ @@ -980,9 +978,7 @@ def rms_norm_batch_invariant( eps: float = 1e-6, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Compute RMS normalization using Triton kernel. - + """Compute RMS normalization using Triton kernel. Args: input: Input tensor of shape (..., hidden_size) @@ -994,6 +990,7 @@ def rms_norm_batch_invariant( Returns: RMS normalized tensor, or ``(output, residual_out)`` when ``residual`` is provided + """ if residual is not None: assert input.shape == residual.shape, ( diff --git a/vllm/model_executor/hw_agnostic/custom_op.py b/vllm/model_executor/hw_agnostic/custom_op.py index 72d3f8522f06..e9c1e2482030 100644 --- a/vllm/model_executor/hw_agnostic/custom_op.py +++ b/vllm/model_executor/hw_agnostic/custom_op.py @@ -30,8 +30,7 @@ def maybe_get_oot_by_class(class_type: type) -> type: class PluggableLayer(nn.Module): - """ - Base class for pluggable layers. + """Base class for pluggable layers. A PluggableLayer is a *module-composing* abstraction: it may instantiate other ``torch.nn.Module`` objects as sub-layers, and its functionality depends on @@ -101,8 +100,7 @@ def decorator(layer_cls): class CustomOp(nn.Module): - """ - Base class for custom ops. + """Base class for custom ops. Dispatches the forward method to the appropriate backend. """ @@ -137,7 +135,8 @@ def forward(self, *args, **kwargs): def forward_native(self, *args, **kwargs): """PyTorch-native implementation; OOT plugins override via - :meth:`forward_oot`.""" + :meth:`forward_oot`. + """ raise NotImplementedError def forward_oot(self, *args, **kwargs): @@ -172,8 +171,7 @@ def dispatch_forward(self, compile_native: bool): return self.forward_native def maybe_compile(self, fn, *, enable: bool = True): - """ - Compile fn if compilation enabled. + """Compile fn if compilation enabled. Useful for CustomOp instances called from within a torch custom op, meaning the forward call is hidden from the model-level torch.compile. @@ -255,8 +253,7 @@ def enabled(cls) -> bool: @staticmethod def default_on() -> bool: - """ - Behavior controlled by `CompilationConfig.custom_ops`: On by default if + """Behavior controlled by `CompilationConfig.custom_ops`: On by default if 'all', off by default if 'none'. When PyTorch Inductor is used, 'none' is the default value, otherwise 'all'. diff --git a/vllm/model_executor/hw_agnostic/layers/layernorm.py b/vllm/model_executor/hw_agnostic/layers/layernorm.py index d2c854afd309..1461e20eb42f 100644 --- a/vllm/model_executor/hw_agnostic/layers/layernorm.py +++ b/vllm/model_executor/hw_agnostic/layers/layernorm.py @@ -10,7 +10,8 @@ @CustomOp.register("rms_norm") class RMSNorm(CustomOp): """``x -> w * x / sqrt(E[x^2] + eps)``. With ``residual``, fuses - ``residual += x`` then RMSNorm and returns ``(normalized, residual)``.""" + ``residual += x`` then RMSNorm and returns ``(normalized, residual)``. + """ def __init__( self, diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index 138833ad149e..0c2ea1aef718 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -209,6 +209,7 @@ def candidate_blocks_to_sparse_indices( ``i * sparse_block_kv + ks % sparse_block_kv``. end: [rows] int32 count of valid sparse-token columns; valid columns always form a prefix because the block ids are sorted. + """ rows, num_candidates = candidate_blocks.shape ratio = candidate_block_size // sparse_block_kv @@ -273,8 +274,17 @@ def sparse_topk_remap( count come back as -1. Args: + logits: [rows, width] bf16 sparse logits, as produced by the sparse + kernels. + sparse_indices: [rows, S] int32 sparse block ids backing ``logits``. + end: [rows] int32 per-row count of valid logit columns. + row_ks: [rows] int32 per-row K range start. + sparse_block_kv: Positions per sparse block (8 or 16). + topk_tokens: Number of tokens to keep per row. + topk_indices: [rows, topk_tokens] int32 output buffer. col_indices: [rows, topk_tokens] int32 scratch for the sparse-column top-k result (row stride 32B-aligned for DeepSelect). + """ rows, width = logits.shape assert logits.dtype == torch.bfloat16 and width >= topk_tokens @@ -329,6 +339,9 @@ def sparse_mqa_logits_prefill_chunk( cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed workspace. candidate_blocks: [rows, K] int32 request-local candidate block ids. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + topk_tokens: Number of tokens to keep per row. topk_indices: [rows, topk_tokens] output buffer. sparse_indices: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`. @@ -341,6 +354,7 @@ def sparse_mqa_logits_prefill_chunk( Returns: The DeepGEMM schedule metadata used, for reuse by later layers. + """ assert weights.dtype == torch.bfloat16, weights.dtype if kernel_metadata is None: @@ -417,6 +431,9 @@ def sparse_mqa_logits_paged_decode( row_indices: [rows] int32 row -> request map (pairing only affects scheduling). candidate_blocks: [rows, K] int32 candidate block ids. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + topk_tokens: Number of tokens to keep per row. topk_indices: [rows, topk_tokens] output buffer. row_ks: [rows] int32 zeros (paged blocks are context-relative). sparse_indices: Caller-owned scratch. @@ -426,6 +443,7 @@ def sparse_mqa_logits_paged_decode( Returns: The DeepGEMM schedule metadata used, for reuse by later layers. + """ rows = q.shape[0] page_kv = kv_cache.shape[1] diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index c610e9b9c8f0..37c73868f7c9 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This module re-exports linear kernel implementations to provide a +"""This module re-exports linear kernel implementations to provide a stable import interface during an ongoing reorganization. Upcoming PRs will remove the scaled_mm and mixed_precision subdirectories and reorganize kernels by provider (aiter, cutlass, flashinfer, etc.) @@ -637,8 +636,7 @@ def choose_scaled_mm_linear_kernel( *, quantization: str, ) -> type[_KernelT]: - """ - Choose a _KernelT that can implement the given config for the + """Choose a _KernelT that can implement the given config for the given compute capability. Attempts to choose the best kernel in terms of performance. @@ -660,8 +658,8 @@ def choose_scaled_mm_linear_kernel( Returns: _KernelT: Chosen kernel. - """ + """ failure_reason_list = [] if force_kernel is not None: @@ -815,8 +813,7 @@ def init_int8_linear_kernel( def choose_mp_linear_kernel( config: MPLinearLayerConfig, compute_capability: int | None = None ) -> type[MPLinearKernel]: - """ - Choose an MPLinearKernel that can implement the given config for the given + """Choose an MPLinearKernel that can implement the given config for the given compute capability. Attempts to choose the best kernel in terms of performance. @@ -832,6 +829,7 @@ def choose_mp_linear_kernel( Returns: type[MPLinearKernel]: Chosen kernel. + """ if compute_capability is None: if current_platform is None: @@ -883,7 +881,8 @@ def choose_mp_linear_kernel( def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8LinearKernel: """Select and instantiate the best MXFP8 linear kernel for the - current platform.""" + current platform. + """ config = Mxfp8LinearLayerConfig(bmm_batch_size=bmm_batch_size) platform = current_platform._enum @@ -935,7 +934,8 @@ def init_mxfp4_linear_kernel( activation_quant_key: QuantKey | None = None, ) -> MxFp4LinearKernel: """Select and instantiate the best MXFP4 linear kernel for the - current platform.""" + current platform. + """ config = MxFp4LinearLayerConfig( activation_quant_key=activation_quant_key, ) @@ -982,7 +982,8 @@ def init_mxfp6_linear_kernel( activation_quant_key: QuantKey | None = None, ) -> MxFp6LinearKernel: """Select and instantiate the best MXFP6 linear kernel for the - current platform.""" + current platform. + """ config = MxFp6LinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key, @@ -1070,7 +1071,8 @@ def init_wfp8_a16_linear_kernel( def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: """Select and instantiate the best NVFP4 linear kernel for the - current platform.""" + current platform. + """ config = NvFp4LinearLayerConfig() a16_kernels = ( FlashInferCuteDslNvFp4W4A16LinearKernel, @@ -1193,8 +1195,7 @@ def register_linear_kernel( platform: PlatformEnum, kernel_type: str = "mp", ) -> None: - """ - Register a new linear kernel class to be considered in kernel selection. + """Register a new linear kernel class to be considered in kernel selection. Args: kernel_class (type): The kernel class to register. @@ -1204,6 +1205,7 @@ def register_linear_kernel( Raises: ValueError: If the kernel_type is not recognized. + """ if kernel_type == "mp": if platform not in _POSSIBLE_KERNELS: diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py index 036b30b7c46d..0db595a3ca42 100644 --- a/vllm/model_executor/kernels/linear/base.py +++ b/vllm/model_executor/kernels/linear/base.py @@ -50,6 +50,7 @@ class Params: # Access typed parameters output = func(input, params.weight, params.weight_scale) ``` + """ weight: torch.Tensor @@ -72,7 +73,7 @@ def from_layer(cls, layer: torch.nn.Module) -> Self: @dataclass class FP8Params(Params): - """FP8 layer parameters with typed fields""" + """FP8 layer parameters with typed fields.""" input_scale_ub: torch.Tensor | None @@ -80,7 +81,7 @@ class FP8Params(Params): @classmethod def from_layer(cls, layer: torch.nn.Module) -> "FP8Params": - """Extract parameters from layer""" + """Extract parameters from layer.""" return cls( weight=getattr(layer, cls.WEIGHT), weight_scale=getattr(layer, cls.WEIGHT_SCALE), @@ -163,6 +164,7 @@ def apply_weights(self, layer, x, bias=None, **kwargs): 2. Initialization: __init__() creates kernel instance with config 3. Weight loading: process_weights_after_loading() preprocesses weights 4. Inference: apply_weights() executes the quantized matmul + """ @classmethod @@ -184,6 +186,7 @@ def is_supported( A tuple of (is_supported, reason): - is_supported: True if the kernel can run on this hardware - reason: If not supported, a string explaining why; otherwise None + """ raise NotImplementedError @@ -204,7 +207,7 @@ def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]: A tuple of (can_implement, reason): - can_implement: True if this kernel supports the config - reason: If not supported, a string explaining why; otherwise None - ``` + """ raise NotImplementedError @@ -214,6 +217,7 @@ def __init__(self, config: _ConfigT) -> None: Args: config: Kernel-specific configuration containing settings like quantization keys, output dtypes, etc. + """ self.config = config @@ -245,6 +249,7 @@ def process_weights_after_loading(self, layer): weight_reordered = reorder_weights(params.weight) replace_parameter(layer, params.WEIGHT, weight_reordered) ``` + """ raise NotImplementedError @@ -270,6 +275,7 @@ def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT: def _get_layer_params(self, layer, **kwargs): return MyKernelParams.from_layer(layer) ``` + """ raise NotImplementedError @@ -281,6 +287,7 @@ def get_output_padding(self) -> int | None: Returns: Number of tokens to pad, or None for no padding (default) + """ return None @@ -306,5 +313,6 @@ def apply_weights( Returns: Output tensor of shape [..., out_features] + """ raise NotImplementedError diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index eb36308bccab..86e029b5f88d 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode. +"""Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode. Routes based on batch size M: M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g) @@ -94,8 +93,7 @@ def _triton_w4a16_skinny_fmt_kernel( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): - """ - Fused W4A16 GEMM reading weights from skinny format [N, K//8]. + """Fused W4A16 GEMM reading weights from skinny format [N, K//8]. B is stored as [N, K//8] int32 using ExLlama shuffle packing: each int32 packs 8 K-values with interleave [0,2,4,6,1,3,5,7]: @@ -196,8 +194,7 @@ def triton_w4a16_skinny_fmt_gemm( zp_bias: int = 8, zp: torch.Tensor | None = None, # [N//8, K//G] int32 zero-points ) -> torch.Tensor: - """ - Fused W4A16 GEMM reading from skinny weight format [N, K//8]. + """Fused W4A16 GEMM reading from skinny weight format [N, K//8]. Args: a: Activation matrix [M, K], float16 or bfloat16. @@ -211,6 +208,7 @@ def triton_w4a16_skinny_fmt_gemm( Returns: Output matrix [M, N], same dtype as a. + """ assert a.is_contiguous(), "Activation matrix must be contiguous" assert b_q.is_contiguous(), "Weight matrix must be contiguous" diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index fd6fb033fe95..ed6c67aa99c5 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Triton-based W4A16 GEMM kernel for ROCm MI300. +"""Triton-based W4A16 GEMM kernel for ROCm MI300. Implements fused int4-weight dequantization + fp16 GEMM in a single kernel, using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]). @@ -67,8 +66,7 @@ def triton_w4a16_gemm_kernel( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): - """ - Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N] + """Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]. B is stored as [K, N//8] int32 using GPTQ sequential packing: each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28]. @@ -170,8 +168,7 @@ def _triton_w4a16_gemm_impl( group_size: int, zp_bias: int = 8, # bias for uint4b8 when qzeros is None ) -> torch.Tensor: - """ - Fused W4A16 GEMM using GPTQ-packed int4 weights. + """Fused W4A16 GEMM using GPTQ-packed int4 weights. Args: a: Activation matrix [M, K], float16 or bfloat16. @@ -184,6 +181,7 @@ def _triton_w4a16_gemm_impl( Returns: Output matrix [M, N], same dtype as a. + """ assert a.is_contiguous(), "Activation matrix must be contiguous" assert b_q.is_contiguous(), "Weight matrix must be contiguous" @@ -302,8 +300,7 @@ def triton_w4a16_gemm( class TritonW4A16LinearKernel(MPLinearKernel): - """ - Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer). + """Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer). Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric) with grouped quantization. Weight tensors are transposed from the @@ -360,8 +357,7 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """ - Convert compressed-tensors checkpoint layout to kernel layout. + """Convert compressed-tensors checkpoint layout to kernel layout. Checkpoint (from compressed_tensors_wNa16.create_weights): weight_packed: [N, K//8] int32 input_dim=1, output_dim=0, packed_dim=1 diff --git a/vllm/model_executor/kernels/linear/mxfp4/base.py b/vllm/model_executor/kernels/linear/mxfp4/base.py index 9c3088df52a1..de0a2a216327 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/base.py +++ b/vllm/model_executor/kernels/linear/mxfp4/base.py @@ -19,6 +19,7 @@ class MxFp4LinearLayerConfig: Attributes: activation_quant_key: Identifies the activation quantization format, or `None` when activations must not be quantized. + """ activation_quant_key: QuantKey | None = None diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py index ef00bc57647d..320693698752 100644 --- a/vllm/model_executor/kernels/linear/mxfp6/base.py +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -21,6 +21,7 @@ class MxFp6LinearLayerConfig: kMxfp6E2M3Static or kMxfp6E3M2Static. activation_quant_key: Identifies the activation quantization format, or `None` when activations must not be quantized. + """ weight_quant_key: QuantKey diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index a555ca17bf89..c5ec08ec8a18 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -177,7 +177,8 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): def input_quant_key(self) -> QuantKey | None: """This kernel supports dynamic quantization of the input. By - convention, pre-quantized blockscales must use the swizzled layout.""" + convention, pre-quantized blockscales must use the swizzled layout. + """ return kNvfp4Dynamic @classmethod diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index 3a053412f3cf..7a38c6964795 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -41,7 +41,8 @@ def from_layer(cls, layer: torch.nn.Module) -> Self: @property def block_scale_attr(self) -> str: """Fp8LinearMethod registers the block scale as ``weight_scale_inv``, - compressed-tensors as ``weight_scale``.""" + compressed-tensors as ``weight_scale``. + """ return ( self.WEIGHT_SCALE if self.weight_scale_inv is None diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index c9046e89f923..2408f4ad6cfc 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -71,8 +71,7 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - """ - `AiterInt8ScaledMMLinearKernel` implements a fused version of + """`AiterInt8ScaledMMLinearKernel` implements a fused version of `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)` where scale_a * a and scale_b * b are implemented using numpy-style broadcasting. diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index 9f69ab0c7377..7df7331f4177 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -174,7 +174,8 @@ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | Non def input_quant_key(self) -> QuantKey | None: """Only static per-tensor activation quantization is supported for external - quantization.""" + quantization. + """ if self.config.activation_quant_key == kFp8StaticTensorSym: return kFp8StaticTensorSym return None @@ -184,7 +185,8 @@ def _pad_to_alignment( x: torch.Tensor, dim: int, alignment: int, value: float = 0.0 ) -> torch.Tensor: """Pad tensor ``x`` along ``dim`` to the next multiple of - ``alignment``.""" + ``alignment``. + """ remainder = x.shape[dim] % alignment if remainder == 0: return x diff --git a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py index 72a3b849840d..1407530e9c78 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py @@ -149,8 +149,7 @@ def apply_block_scaled_mm( class FlashInferFp8DeepGEMMDynamicBlockScaledKernel( Fp8BlockScaledDynamicMMLinearKernel ): - """ - Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM. + """Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM. Dispatches between two kernels based on input batch size: - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation. @@ -214,9 +213,7 @@ def _flashinfer_fp8_blockscale_gemm_fake( weight: torch.Tensor, weight_scale: torch.Tensor, ) -> torch.Tensor: - """ - Required fake/meta implementation for torch.compile graph tracing. - """ + """Required fake/meta implementation for torch.compile graph tracing.""" return torch.empty( input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device ) @@ -236,8 +233,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( group_size: int, use_deep_gemm_e8m0: bool, ) -> torch.Tensor: - """ - Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection. + """Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection. This function switches between two optimized kernels based on the input batch size: - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization. @@ -260,6 +256,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( Returns: Output tensor of shape (batch_size, output_dim) in bfloat16 format + """ def run_flashinfer_deepgemm_swapAB( @@ -323,9 +320,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake( group_size: int, use_deep_gemm_e8m0: bool, ) -> torch.Tensor: - """ - Required fake/meta implementation for torch.compile graph tracing. - """ + """Required fake/meta implementation for torch.compile graph tracing.""" return torch.empty( input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device ) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 9e148ca308cf..3ca89d3cea1f 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -27,8 +27,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): - """ - FP8 Marlin kernel for GPUs that lack FP8 hardware support. + """FP8 Marlin kernel for GPUs that lack FP8 hardware support. Leverages the Marlin kernel for fast weight-only FP8 quantization. """ diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 4fc7ace5a1fe..1f2756cc267e 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -50,8 +50,7 @@ def _supports_torch_fp8_scaled_mm() -> bool: class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): - """ - Base class for FP8 linear kernels using Torch. + """Base class for FP8 linear kernels using Torch. Each subclass represents a kernel variant for specific device capabilities and torch versions. """ diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index d34e29314d75..06b4f1498463 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -19,8 +19,7 @@ def mhc_pre_aiter( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Forward pass for mHC pre block. + """Forward pass for mHC pre block. Args: residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 @@ -40,8 +39,8 @@ def mhc_pre_aiter( post_mix: shape (..., hc_mult), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 - """ + """ hidden_size = residual.shape[-1] assert hidden_size % 256 == 0 from vllm._aiter_ops import rocm_aiter_ops @@ -120,7 +119,7 @@ def mhc_pre_delayed_aiter( comb_res_mix: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """mHC pre with the pre-mix carried in from the previous sublayer. + """MHC pre with the pre-mix carried in from the previous sublayer. Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix* rather than the gate computed here, and that gate is returned as the @@ -151,6 +150,7 @@ def mhc_pre_delayed_aiter( comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 next_pre_mix: shape (..., hc_mult), dtype torch.float32 + """ hidden_size = residual.shape[-1] assert hidden_size % 256 == 0 diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index cea186e49e0d..62741bbda19c 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -109,6 +109,7 @@ def mhc_pre_delayed_tilelang( Post and residual coefficients, optionally normalized BF16 layer input, and the next FP32 pre-mix, with shapes (tokens, hc_mult, 1), (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult). + """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _HC_PRENORM_GEMM_TILELANG_KERNEL, @@ -320,6 +321,7 @@ def mhc_fused_post_pre_delayed_tilelang( (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1), (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult), and (tokens, hidden_size). + """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _HC_PRENORM_GEMM_TILELANG_KERNEL, @@ -551,8 +553,7 @@ def mhc_pre_tilelang( norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Forward pass for mHC pre block. + """Forward pass for mHC pre block. Args: residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 @@ -576,6 +577,7 @@ def mhc_pre_tilelang( post_mix: shape (..., hc_mult), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 + """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, @@ -827,8 +829,7 @@ def mhc_fused_post_pre_tilelang( norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Run one MHC post block followed by the next MHC pre block. + """Run one MHC post block followed by the next MHC pre block. When ``norm_weight`` is provided, the layer_input_cur output is the RMSNorm'd activation (fused into the kernel); otherwise it is the @@ -842,8 +843,8 @@ def mhc_fused_post_pre_tilelang( post_mix_cur: shape (..., hc_mult, 1) comb_mix_cur: shape (..., hc_mult, hc_mult) layer_input_cur: shape (..., hidden_size) - """ + """ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _MHC_FUSED_TILELANG_KERNEL, _MHC_POST_TILELANG_KERNEL, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 9e803d123e29..fa2d5e17004a 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -661,7 +661,7 @@ def mhc_fused_tilelang( tile_n: int = 1, split_k: int = 1, ) -> tilelang.JITKernel: - """Fused mhc post-mapping + pre-norm GEMM FMA""" + """Fused mhc post-mapping + pre-norm GEMM FMA.""" m = T.dynamic("num_tokens") split_k = T.dynamic("split_k") h = hidden diff --git a/vllm/model_executor/kernels/mhc/torch.py b/vllm/model_executor/kernels/mhc/torch.py index 9c57cc169cc6..eb4de35e7c75 100644 --- a/vllm/model_executor/kernels/mhc/torch.py +++ b/vllm/model_executor/kernels/mhc/torch.py @@ -57,8 +57,7 @@ def mhc_pre_torch( sinkhorn_repeat: int, n_splits: int = 1, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Forward pass for mHC pre block. + """Forward pass for mHC pre block. Args: residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 @@ -76,8 +75,8 @@ def mhc_pre_torch( post_mix: shape (..., hc_mult), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 - """ + """ # Validate shapes assert residual.dtype == torch.bfloat16 assert fn.dtype == torch.float32 diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index f159d2dff8b6..f730026d4a90 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -479,7 +479,6 @@ def __init__(self, alpha: float = 1.702, limit: float = 7.0): def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" - gate, up = x[..., ::2], x[..., 1::2] gate = gate.clamp(min=None, max=self.limit) up = up.clamp(min=-self.limit, max=self.limit) @@ -647,9 +646,7 @@ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: # --8<-- [start:relu2] @CustomOp.register("relu2") class ReLUSquaredActivation(CustomOp): - """ - Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2 - """ + """Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2.""" # --8<-- [end:relu2] @@ -676,8 +673,7 @@ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: # --8<-- [start:xielu] @CustomOp.register("xielu") class XIELU(CustomOp): - """ - Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010 + """Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010 If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA Otherwise, we emit a single warning and use xIELU Python """ @@ -747,7 +743,7 @@ def _xielu_python(self, x: torch.Tensor) -> torch.Tensor: ) def _xielu_cuda(self, x: torch.Tensor) -> torch.Tensor: - """Firewall function to prevent torch.compile from seeing .item()""" + """Firewall function to prevent torch.compile from seeing .item().""" assert self._xielu_cuda_obj is not None, "XIELU CUDA object must not be None" original_shape = x.shape # CUDA kernel expects 3D tensors, reshape if needed @@ -808,7 +804,8 @@ def forward_cuda(self, input: torch.Tensor) -> torch.Tensor: def _get_gelu_pytorch_tanh() -> nn.Module: """Get PyTorch GELU with tanh approximation, with ROCm fallback - and fast GELU for ARM.""" + and fast GELU for ARM. + """ if current_platform.is_rocm(): # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile logger.warning_once( diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index a001c8c29c80..79012a4a5cfd 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -166,8 +166,8 @@ def _init_kv_cache_quant( layer: The attention layer instance to initialize. quant_config: Optional quantization configuration. prefix: Layer name prefix for quantization method lookup. - """ + """ # Note [Register q/k/v/prob scales in state dict] # When calling model.to(device), only parameters/buffers in state dict are # moved. If not registering q/k/v/prob scales in state dict, there would @@ -254,8 +254,7 @@ def __init__( head_size_v: int | None = None, **extra_impl_args, ) -> None: - """ - The KV cache is stored inside this class and is accessed via + """The KV cache is stored inside this class and is accessed via `self.kv_cache`. """ super().__init__() @@ -493,8 +492,7 @@ def forward( output_shape: torch.Size | None = None, output_dtype: torch.dtype | None = None, ) -> torch.Tensor: - """ - The KV cache is stored inside this class and is accessed via + """The KV cache is stored inside this class and is accessed via `self.kv_cache`. Attention metadata (`attn_metadata`) is set using a context manager in @@ -691,6 +689,7 @@ def get_attention_context( Note: attn_metadata may be None, but attn_layer and kv_cache are always extracted from the forward context. + """ forward_context: ForwardContext = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -718,8 +717,7 @@ def unified_kv_cache_update( value: torch.Tensor, layer_name: LayerNameType, ) -> torch.Tensor: - """ - Returns a dummy that is passed to unified_attention to signal a side effect and + """Returns a dummy that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ layer_name = _resolve_layer_name(layer_name) diff --git a/vllm/model_executor/layers/attention/cross_attention.py b/vllm/model_executor/layers/attention/cross_attention.py index 9336bc17c241..1422e2778564 100644 --- a/vllm/model_executor/layers/attention/cross_attention.py +++ b/vllm/model_executor/layers/attention/cross_attention.py @@ -34,7 +34,6 @@ def _get_cross_slot_mapping( device: torch.device, ) -> torch.Tensor: """Get cross-attention slot mappings.""" - block_size = kv_cache_spec.block_size slot_mappings = [] @@ -187,8 +186,7 @@ def forward( class CrossAttention(Attention): - """ - Cross-attention for encoder-decoder models. + """Cross-attention for encoder-decoder models. Handles attention between decoder queries and encoder keys/values. """ diff --git a/vllm/model_executor/layers/attention/encoder_only_attention.py b/vllm/model_executor/layers/attention/encoder_only_attention.py index 5805fe2ae1ca..ef73b2798da4 100644 --- a/vllm/model_executor/layers/attention/encoder_only_attention.py +++ b/vllm/model_executor/layers/attention/encoder_only_attention.py @@ -49,9 +49,7 @@ def build( class EncoderOnlyAttention(Attention): - """ - Encoder attention is a special case that doesn't need a KV Cache. - """ + """Encoder attention is a special case that doesn't need a KV Cache.""" def __init__( self, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1e673bd2da14..caa58eefae12 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -# MLA Common Components +"""# MLA Common Components This file implements common components for MLA implementations. @@ -1415,8 +1414,7 @@ def unified_mla_kv_cache_update( kv_cache_dtype: str, k_scale: torch.Tensor, ) -> torch.Tensor: - """ - Returns a dummy that is passed to unified_attention to signal a side effect and + """Returns a dummy that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ layer_name = _resolve_layer_name(layer_name) @@ -1539,8 +1537,7 @@ def dynamic_per_batched_tensor_quant( dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0}, ) class _DecodeConcatQuantFP8(QuantFP8): - """ - QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before + """QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before quantization. When disabled, forward_native is compiled via torch.compile, fusing cat/reshape/quant/view together. """ @@ -1571,7 +1568,8 @@ class MLACommonBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack an inline fp32 scale pair after the - latent data (single-sided: ``head_size_v == 0`` for MLA).""" + latent data (single-sided: ``head_size_v == 0`` for MLA). + """ spec = super().customize_spec(spec) mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: @@ -2030,6 +2028,7 @@ def build_mla_chunked_context_metadata( Returns: The chunked-context metadata, or None when no prefill has any context. + """ # NOTE: it is recommended you read the `Chunked Prefill` section in the # comment at the top of the file before trying to understand this code. @@ -2230,8 +2229,7 @@ def padded_rows(rows: int) -> int: class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): - """ - NOTE: Please read the comment at the top of the file before trying to + """NOTE: Please read the comment at the top of the file before trying to understand this class """ @@ -2316,8 +2314,7 @@ def determine_prefill_query_data_type( vllm_config: VllmConfig, model_dtype: torch.dtype, ) -> torch.dtype: - """ - Determine the query data type for prefill queries. + """Determine the query data type for prefill queries. Return FP8 dtype if cache is FP8 and prefill query quantization is enabled, else model dtype. """ @@ -2484,8 +2481,7 @@ def _build_decode( def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata ) -> M: - """ - This method builds the metadata for full cudagraph capture. + """This method builds the metadata for full cudagraph capture. Currently, only decode is supported for full cudagraphs with MLA. """ m = common_attn_metadata @@ -2671,14 +2667,16 @@ def reorg_kvcache( max_seq_len: int, toks: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - reorg and unpad kvcache after cp local gather to tp layout for attn kernel. + """Reorg and unpad kvcache after cp local gather to tp layout for attn kernel. e.g. allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ..., T0_4, T0_5, pad, pad, T1_2, pad, ...] -> reorganized_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T0_4, T0_5, T1_0, T1_1, T1_2, ...] + Args: + allgatered_kv_c_normed: all-gathered, padded latent KV cache. + allgatered_k_pe: all-gathered, padded RoPE key cache. padded_local_chunk_seq_lens_lst: local chunk context lengths under current CP rank. local_context_lens_allranks: local context lengths on each CP rank. @@ -2687,6 +2685,7 @@ def reorg_kvcache( sum_seq_len: the sum of cp_chunk_seq_lens_lst. max_seq_len: the max value of cp_chunk_seq_lens_lst. toks: the number of tokens for local gather cache. + """ kv_c_segments = [] k_pe_segments = [] @@ -2789,10 +2788,16 @@ def accumulate_mla_context_chunk( remaining token range is initialized. Args: + chunk: The context chunk being folded in. + attn_output: The chunk's attention output. + attn_softmax_lse: The chunk's log-sum-exp values. + output: Running context partial, updated in place. + output_lse: Running log-sum-exp, updated in place. output_written: The chunk's attention output already landed in ``output[chunk.token_slice]`` because the backend was handed it as ``out``, leaving only the lse to fold. Invalid for a continuation chunk, whose leading tokens must be merged rather than overwritten. + """ token_start = chunk.token_slice.start token_end = chunk.token_slice.stop @@ -2819,8 +2824,7 @@ def accumulate_mla_context_chunk( class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): - """ - Shared MLA base providing dense-MHA prefill (via the selected + """Shared MLA base providing dense-MHA prefill (via the selected MLAPrefillBackend) for both dense and sparse impls; subclasses add decode (``forward_mqa``). """ @@ -2856,8 +2860,7 @@ def __init__( def _concat_k_nope_k_pe( self, k_nope: torch.Tensor, k_pe: torch.Tensor ) -> torch.Tensor: - """ - Efficiently concatenate k_nope and k_pe tensors along the last dimension. + """Efficiently concatenate k_nope and k_pe tensors along the last dimension. This function avoids the performance penalty of torch.cat with expanded non-contiguous tensors by pre-allocating the output and using direct copies. @@ -2869,6 +2872,7 @@ def _concat_k_nope_k_pe( Returns: Tensor of shape [..., nope_dim + pe_dim] + """ if k_pe.shape[-1] == 0: # NoPE MLA: nothing to append, so no copy either. @@ -3224,8 +3228,7 @@ def forward_mha( # type: ignore[override] class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]): - """ - NOTE: Please read the comment at the top of the file before trying to + """NOTE: Please read the comment at the top of the file before trying to understand this class """ diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index b8f431ed415f..0af2396ed719 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -322,14 +322,14 @@ def __init__( num_kv_heads: int | None = None, prefix: str = "", ) -> None: - """ - Args: - num_heads: number of attention heads per partition. - head_size: hidden_size per attention head. - scale: scale factor. - num_kv_heads: number of kv heads. - prefix: This has no effect, it is only here to make it easier to - swap between Attention and MultiHeadAttention + """Args: + num_heads: number of attention heads per partition. + head_size: hidden_size per attention head. + scale: scale factor. + num_kv_heads: number of kv heads. + prefix: This has no effect, it is only here to make it easier to + swap between Attention and MultiHeadAttention + """ super().__init__() @@ -517,8 +517,7 @@ def view_qkv_to_4d( q_len: int, kv_len: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Reshape query, key, value to 4D tensors: + """Reshape query, key, value to 4D tensors: (batch_size, seq_len, num_heads, head_size) """ query = query.view(bsz, q_len, self.num_heads, self.head_size) diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index d321d675a5fe..d8917de63a3f 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -114,9 +114,7 @@ def build( @CustomOp.register("static_sink_attention") class StaticSinkAttention(Attention, CustomOp): - """ - Attention with static sink tokens - """ + """Attention with static sink tokens.""" def __init__( self, diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index 197213dc79fe..cee57fdd041e 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -12,8 +12,7 @@ class AttentionLayerBase(ABC): - """ - Base class for attention-like layers (Attention, Mamba, etc.) + """Base class for attention-like layers (Attention, Mamba, etc.) that support the v1 engine. This provides a common interface for getting attention backends @@ -38,8 +37,7 @@ def get_attn_backend(self) -> type[AttentionBackend]: @abstractmethod def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: - """ - Get the KV cache spec for this layer. + """Get the KV cache spec for this layer. May be None if the layer does not need KV cache. """ pass diff --git a/vllm/model_executor/layers/conv.py b/vllm/model_executor/layers/conv.py index 14f60231cdf9..b9719ef58fdd 100644 --- a/vllm/model_executor/layers/conv.py +++ b/vllm/model_executor/layers/conv.py @@ -143,7 +143,7 @@ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: return x def forward_native(self, x: torch.Tensor) -> torch.Tensor: - """Expected input shape: (batch_size, in_channels, height, width)""" + """Expected input shape: (batch_size, in_channels, height, width).""" assert x.dim() == 4 if self.enable_linear: return self._forward_mulmat(x) @@ -193,7 +193,7 @@ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: return x def forward_native(self, x: torch.Tensor) -> torch.Tensor: - """Expected input shape: (batch_size, in_channels, time, height, width)""" + """Expected input shape: (batch_size, in_channels, time, height, width).""" if self.enable_linear: return self._forward_mulmat(x) else: diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index 7046bfc991b9..4f2d5c1783f3 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -55,7 +55,8 @@ def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None: """Workspace token budget for flashinfer fused all-reduce, or None if the - current world size / device is unsupported. Mirrors ``FlashInferAllReduce``.""" + current world size / device is unsupported. Mirrors ``FlashInferAllReduce``. + """ from vllm.config.compilation import PassConfig max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size) diff --git a/vllm/model_executor/layers/fused_embed_norm.py b/vllm/model_executor/layers/fused_embed_norm.py index ddeb5156ae18..876585ad32bd 100644 --- a/vllm/model_executor/layers/fused_embed_norm.py +++ b/vllm/model_executor/layers/fused_embed_norm.py @@ -170,7 +170,8 @@ def _fused_embed_eh_norm_kernel( ``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H]) ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm + - cat. Requires the full table on-rank (replicated embedding).""" + cat. Requires the full table on-rank (replicated embedding). + """ tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py index d83301bed19f..18f97be332ab 100644 --- a/vllm/model_executor/layers/fused_moe/activation.py +++ b/vllm/model_executor/layers/fused_moe/activation.py @@ -56,7 +56,8 @@ def is_gated(self) -> bool: @property def custom_op_name(self) -> str: """Maps to the CustomOp name of activations - in vllm/model_executor/layers/activation.py.""" + in vllm/model_executor/layers/activation.py. + """ return _CUSTOM_OP_NAMES[self] def without_mul(self) -> "MoEActivation": @@ -115,6 +116,7 @@ def activation_without_mul(activation: str) -> str: Returns: The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul") + """ return MoEActivation.from_str(activation).without_mul().value diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 5019867d53ed..75ea91bfca5d 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -122,8 +122,7 @@ def maybe_roundup_layer_hidden_size( act_dtype: torch.dtype, moe_parallel_config: FusedMoEParallelConfig, ) -> int: - """ - Given layer hidden size and MoE configurations, round up hidden_size + """Given layer hidden size and MoE configurations, round up hidden_size if necessary. Args: @@ -135,6 +134,7 @@ def maybe_roundup_layer_hidden_size( Rounded up hidden_size if rounding up is required based on the configs and all2all backend. Original hidden size otherwise. + """ if moe_parallel_config.use_deepep_ht_kernels: hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size( diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 4c7760d00792..ccca24cc6f66 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -41,8 +41,7 @@ def _get_config_dtype_str( use_int4_w4a16: bool = False, ocp_mx_scheme: str | None = None, ) -> str | None: - """ - Return a string used to construct the filename that contains the + """Return a string used to construct the filename that contains the tuning info for a particular quantization scheme. See try_get_optimal_moe_config in fused_moe.py. """ @@ -72,9 +71,7 @@ def _quant_flags_to_group_shape( per_out_ch_quant: bool, block_shape: list[int] | None, ) -> tuple[GroupShape | None, GroupShape | None]: - """ - Convert MoE quantization flags into more generic GroupShapes. - """ + """Convert MoE quantization flags into more generic GroupShapes.""" a_shape: GroupShape | None w_shape: GroupShape | None if block_shape is not None: @@ -173,8 +170,7 @@ def get_routing_method_type( @dataclass class FusedMoEQuantDesc: - """ - A quantization descriptor for fused MoE ops. This class can describe + """A quantization descriptor for fused MoE ops. This class can describe either activations or weights. """ @@ -212,8 +208,7 @@ class FusedMoEQuantDesc: # e.g. for specific arguments bias, precision, etc. @dataclass class FusedMoEQuantConfig: - """ - The FusedMoEQuantConfig contains all the quantization parameters for + """The FusedMoEQuantConfig contains all the quantization parameters for a single FusedMoEMethodBase operation. It consists of four FusedMoEQuantDescs, one for each activation and set of weights. @@ -435,8 +430,7 @@ def use_mxfp4_w4a8(self) -> bool: return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4" def config_name(self, dtype: torch.dtype) -> str | None: - """ - Return a string used to construct the filename that contains the + """Return a string used to construct the filename that contains the tuning info for a particular quantization scheme. See try_get_optimal_moe_config in fused_moe.py. """ @@ -454,8 +448,7 @@ def scale_shape( max_tokens: int, hidden_dim: int, ) -> tuple[int, int] | None: - """ - Construct the proper activation scale shape for this + """Construct the proper activation scale shape for this config. """ if self.is_quantized: @@ -477,8 +470,7 @@ def batched_scale_shape( max_tokens: int, hidden_dim: int, ) -> tuple[int, int, int] | None: - """ - Construct the proper activation batched scale shape for this + """Construct the proper activation batched scale shape for this config, e.g. (num experts, *scale_shape). """ if self.is_quantized: @@ -512,8 +504,7 @@ def make( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> "FusedMoEQuantConfig": - """ - General builder function for a FusedMoEQuantConfig. + """General builder function for a FusedMoEQuantConfig. - quant_dtype: Optional quantization type. None if activations are unquantized or quantized prior to calling. Note: "nvfp4", "mxfp4", "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values @@ -607,9 +598,7 @@ def fp8_w8a8_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for fp8 activations and fp8 weights. - """ + """Construct a quant config for fp8 activations and fp8 weights.""" return FusedMoEQuantConfig.make( current_platform.fp8_dtype(), w1_scale=w1_scale, @@ -640,9 +629,7 @@ def int8_w8a8_moe_quant_config( w2_bias: torch.Tensor | None = None, per_act_token_quant: bool = False, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for int8 activations and int8 weights. - """ + """Construct a quant config for int8 activations and int8 weights.""" return FusedMoEQuantConfig.make( torch.int8, w1_scale=w1_scale, @@ -667,9 +654,7 @@ def gptq_marlin_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ): - """ - Construct a quant config for gptq marlin quantization. - """ + """Construct a quant config for gptq marlin quantization.""" from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) @@ -702,9 +687,7 @@ def mxfp4_w4a16_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for unquantized activations and mxfp4 weights. - """ + """Construct a quant config for unquantized activations and mxfp4 weights.""" return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc(), _a2=FusedMoEQuantDesc(), @@ -730,9 +713,7 @@ def mxfp4_mxfp8_moe_quant_config( mx_alignment: int = 0, is_scale_swizzled: bool = True, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for mxfp4 activations and mxfp4 weights. - """ + """Construct a quant config for mxfp4 activations and mxfp4 weights.""" return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc("mxfp8"), _a2=FusedMoEQuantDesc("mxfp8"), @@ -756,9 +737,7 @@ def mxfp4_w4a8_moe_quant_config( block_shape: list[int] | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for fp8 activations and mxfp4 weights. - """ + """Construct a quant config for fp8 activations and mxfp4 weights.""" return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None), _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None), @@ -782,9 +761,7 @@ def ocp_mx_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for mxfp4 activations and mxfp4 weights. - """ + """Construct a quant config for mxfp4 activations and mxfp4 weights.""" assert quant_dtype in OCP_MX_DTYPES return FusedMoEQuantConfig.make( quant_dtype=quant_dtype, @@ -818,9 +795,7 @@ def nvfp4_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for mxfp4 activations and nvp4 weights. - """ + """Construct a quant config for mxfp4 activations and nvp4 weights.""" return FusedMoEQuantConfig.make( "nvfp4", w1_scale=w1_scale, @@ -845,8 +820,7 @@ def mxfp4_moe_quant_config( w1_scale: torch.Tensor, w2_scale: torch.Tensor, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for MXFP4 x MXFP4 MoE. + """Construct a quant config for MXFP4 x MXFP4 MoE. MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no separate alphas / global activation scales in this config. """ @@ -869,9 +843,7 @@ def nvfp4_w4a16_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for 16-but activations and nvp4 weights. - """ + """Construct a quant config for 16-but activations and nvp4 weights.""" return FusedMoEQuantConfig.make( quant_dtype=None, w1_scale=w1_scale, @@ -899,9 +871,7 @@ def int4_w4a16_moe_quant_config( gemm1_alpha: float | None = None, gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for 16-bit float activations and int4 weights. - """ + """Construct a quant config for 16-bit float activations and int4 weights.""" group_shape = GroupShape(*block_shape) if block_shape is not None else None return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), @@ -924,9 +894,7 @@ def fp8_w8a16_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for 16-bit float activations and fp8 weights. - """ + """Construct a quant config for 16-bit float activations and fp8 weights.""" group_shape = GroupShape(*block_shape) if block_shape is not None else None fp8_dtype = current_platform.fp8_dtype() return FusedMoEQuantConfig( @@ -968,9 +936,7 @@ def int8_w8a16_moe_quant_config( gemm1_alpha: float | None = None, gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for 16-bit float activations and int8 weights. - """ + """Construct a quant config for 16-bit float activations and int8 weights.""" group_shape = GroupShape(*block_shape) if block_shape is not None else None return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), @@ -992,9 +958,7 @@ def int4_w4afp8_moe_quant_config( per_out_ch_quant: bool = False, block_shape: list[int] | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for fp8 activations and int4 weights. - """ + """Construct a quant config for fp8 activations and int4 weights.""" return FusedMoEQuantConfig.make( torch.float8_e4m3fn, # quant dtype for activations w1_scale=w1_scale, @@ -1015,8 +979,7 @@ def biased_moe_quant_config( gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: - """ - Construct a quant config for unquantized activations with biases. + """Construct a quant config for unquantized activations with biases. gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params through to the fused activation kernel (e.g. swigluoai_uninterleave). @@ -1136,8 +1099,7 @@ def make( sp_size_: int, vllm_parallel_config: ParallelConfig, ) -> "FusedMoEParallelConfig": - """ - Determine MoE parallel configuration. Based on the input `tp_size_`, + """Determine MoE parallel configuration. Based on the input `tp_size_`, `dp_size_` and vllm's parallel config, determine what level's of parallelism to use in the fused moe layer. @@ -1145,6 +1107,7 @@ def make( tp_size_ (int): `tp_size` passed into the FusedMoEFactory constructor. pcp_size_ (int): `pcp_size` passed into the FusedMoEFactory constructor. dp_size_ (int): `dp_size` passed into the FusedMoEFactory constructor. + sp_size_ (int): `sp_size` passed into the FusedMoEFactory constructor. vllm_parallel_config (ParallelConfig): vLLM's parallel config object which contains the `enable_expert_parallel` flag. @@ -1208,6 +1171,7 @@ def make( - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3} - Comment: There are 2 engine instances and the experts are split between the 4 devices. + """ use_ep = ( dp_size_ * pcp_size_ * tp_size_ > 1 diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 70158dafa1c7..4ea864eaac3b 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py +"""Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py and updated to fit vllm needs and terminology. """ diff --git a/vllm/model_executor/layers/fused_moe/expert_map_manager.py b/vllm/model_executor/layers/fused_moe/expert_map_manager.py index 71f2186ea4dd..daa2b3ff8e83 100644 --- a/vllm/model_executor/layers/fused_moe/expert_map_manager.py +++ b/vllm/model_executor/layers/fused_moe/expert_map_manager.py @@ -27,8 +27,7 @@ def determine_expert_map( num_fused_shared_experts: int = 0, return_expert_mask: bool = False, ) -> tuple[int, torch.Tensor | None, torch.Tensor | None]: - """ - Calculates how many experts should be assigned to each rank for EP and + """Calculates how many experts should be assigned to each rank for EP and creates a mapping from global to local expert index. Experts are distributed evenly across ranks. Any remaining are assigned to the last rank. @@ -56,6 +55,7 @@ def determine_expert_map( and 0 for sentinel. Returns None if ep_size is 1. Used only when AITER MOE is enabled. + """ from typing import get_args @@ -150,8 +150,7 @@ def determine_expert_placement_strategy( class ExpertMapManager: - """ - Manages expert ID mappings and placement for Expert Parallelism. + """Manages expert ID mappings and placement for Expert Parallelism. Responsibilities: - Calculate local vs global expert counts @@ -193,16 +192,22 @@ def __init__( num_fused_shared_experts: int = 0, rocm_aiter_enabled: bool = False, ): - """ - Initialize expert map manager. + """Initialize expert map manager. Args: + max_num_batched_tokens: Maximum number of batched tokens per step + top_k: Number of experts each token is routed to global_num_experts: Total number of experts across all ranks + num_redundant_experts: Number of redundant expert replicas (EPLB) + num_expert_group: Number of expert groups, when grouped routing is + used moe_parallel_config: MoE parallel configuration (contains ep_size, ep_rank, backend flags) placement_strategy: Strategy for placing experts ('linear' or 'round_robin') num_fused_shared_experts: Number of fused shared experts (for AITER) rocm_aiter_enabled: Whether ROCm AITER fusion is enabled + enable_eplb: Whether expert-parallel load balancing is enabled + """ self.global_num_experts = global_num_experts self.moe_parallel_config = moe_parallel_config @@ -295,8 +300,7 @@ def local_num_experts(self) -> int: @property def expert_map(self) -> torch.Tensor | None: - """ - Mapping from global expert ID to local expert ID. + """Mapping from global expert ID to local expert ID. Returns tensor of shape (global_num_experts,) where: - expert_map[global_id] = local_id if expert is on this rank @@ -308,8 +312,7 @@ def expert_map(self) -> torch.Tensor | None: @property def expert_mask(self) -> torch.Tensor | None: - """ - Expert mask for AITER fusion (ROCm-specific). + """Expert mask for AITER fusion (ROCm-specific). Returns tensor of shape (global_num_experts + num_fused_shared + 1,) where 1 indicates expert is on this rank, 0 otherwise. @@ -325,8 +328,7 @@ def placement_strategy(self) -> ExpertPlacementStrategy: def routing_tables( self, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: - """ - Routing tables for round-robin placement. + """Routing tables for round-robin placement. Returns (global_to_physical, physical_to_global, local_to_global) or None if not using round-robin or tables not needed. @@ -334,8 +336,7 @@ def routing_tables( return self._routing_tables def map_global_to_local(self, global_id: int) -> int: - """ - Map global expert ID to local expert ID. + """Map global expert ID to local expert ID. Args: global_id: Global expert ID (0 to global_num_experts - 1) @@ -345,6 +346,7 @@ def map_global_to_local(self, global_id: int) -> int: Raises: ValueError: If expert is not on this rank + """ if self._expert_map is None: return global_id @@ -369,8 +371,7 @@ def update( moe_parallel_config: FusedMoEParallelConfig, global_num_experts: int, ) -> None: - """ - Update expert mappings for new EP configuration. + """Update expert mappings for new EP configuration. Used during dynamic reconfiguration (e.g., elastic scaling). @@ -378,6 +379,7 @@ def update( global_num_experts: New total number of experts across all ranks moe_parallel_config: New MoE parallel configuration (contains ep_size, ep_rank, backend flags) + """ self.moe_parallel_config = moe_parallel_config self.global_num_experts = global_num_experts @@ -397,8 +399,7 @@ def update( self._init_aiter_shared_experts_topK_buffer() def get_compressed_map_string(self) -> str: - """ - Get compressed string representation of expert map for logging. + """Get compressed string representation of expert map for logging. Returns string mapping local to global expert IDs. """ @@ -457,8 +458,7 @@ def _calculate_expert_maps(self) -> None: def _init_routing_tables( self, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: - """ - Ensure routing tables are initialized if needed for round-robin. + """Ensure routing tables are initialized if needed for round-robin. This is a public method that can be called to explicitly initialize routing tables. It's safe to call multiple times (idempotent). diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index e14e54922a8a..757d1d8295f7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -56,9 +56,7 @@ def _aiter_w4a16_silu_via_a8w4( unpadded_N_w2, unpadded_K_w2, ) -> torch.Tensor: - """ - MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation. - """ + """MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.""" from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul from aiter.ops.triton.quant import dynamic_mxfp8_quant diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 9d35774f4d78..0691d9fed4f2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -29,8 +29,7 @@ # _get_padding_mask and patch_gating_output after # https://github.com/ROCm/aiter/pull/4530 is released and consumed by vLLM. def _get_padding_mask() -> torch.Tensor | None: - """ - Retrieves a boolean mask with non-padding (0) and padding (1) tokens. + """Retrieves a boolean mask with non-padding (0) and padding (1) tokens. `slot_mapping < 0` comes from: @@ -247,8 +246,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): - """ - Monolithic MXFP4 W4A8 expert using AITER triton kernels. + """Monolithic MXFP4 W4A8 expert using AITER triton kernels. This backend uses: - aiter.ops.triton.moe_routing.routing for routing diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 93a4088fc4f7..dcd9b8b543ca 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -36,7 +36,8 @@ def is_aiter_mxfp8_moe_available() -> bool: nightly lacking #3811 would wrongly select FlyDSL instead of falling back to the native Triton dot_scaled path. #3811 added no probe-able public symbol, so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails - closed (returns False -> triton dot_scaled), which is always safe.""" + closed (returns False -> triton dot_scaled), which is always safe. + """ if not (current_platform.is_rocm() and current_platform.supports_mx()): return False try: diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index c578fb641b10..d3ea8f274cff 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -307,8 +307,7 @@ def __init__( max_num_tokens: int, num_dispatchers: int, ): - """ - max_num_tokens: Maximum number of tokens from a DP Rank + """max_num_tokens: Maximum number of tokens from a DP Rank num_dispatchers: The number of DP dispatchers. quant_config: Quantization configuration """ @@ -350,8 +349,7 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return True def supports_packed_ue8m0_act_scales(self) -> bool: - """ - DeepGemm supports packed ue8m0 activation scales on Blackwell-family + """DeepGemm supports packed ue8m0 activation scales on Blackwell-family GPUs (SM100 datacenter and SM120 consumer). """ return is_deep_gemm_e8m0_used() and ( diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py index b1e32462d842..5ee98a8c847f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py @@ -172,6 +172,7 @@ def _activation_kind(activation: MoEActivation) -> int: 0 = SwiGLU_Gu (SiLU(g)*u) 1 = SwiGLU_Ug (SiLU(u)*g) 2 = SiLU + """ if activation == MoEActivation.SWIGLUSTEP: return 0 @@ -217,6 +218,7 @@ def apply( Returns: Output tensor after MoE computation + """ from vllm.model_executor.layers.fused_moe.router.cpu_router import ( select_experts, diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index f7d82c691bf1..b9a96b280411 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -192,7 +192,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def _pad_moe_intermediate(self, layer: torch.nn.Module) -> None: """Zero-pad the per-partition MoE intermediate dim of both weights and - the expert bias, see `_padded_intermediate_size`.""" + the expert bias, see `_padded_intermediate_size`. + """ intermediate_size = self.moe_config.intermediate_size_per_partition padded_size = self._padded_intermediate_size(self.moe_config) if padded_size == intermediate_size: @@ -745,6 +746,7 @@ def prepare_int4_moe_layer_for_cpu( Returns: (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2) + """ E = w13_packed.size(0) @@ -951,7 +953,8 @@ def apply( class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation - modular MoE experts.""" + modular MoE experts. + """ def __init__( self, @@ -1055,7 +1058,6 @@ def supports_expert_map(self) -> bool: def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """VNNI-prepack INT8 MoE weights for CPU kernel.""" - w13 = torch.ops._C.convert_weight_packed(layer.w13_weight) w2 = torch.ops._C.convert_weight_packed(layer.w2_weight) replace_parameter(layer, "w13_weight", w13) @@ -1285,7 +1287,8 @@ def apply( class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): """AMD Zen INT8 MoE with per-token activation and channelwise weight - quantization, dispatched through zentorch.""" + quantization, dispatched through zentorch. + """ def __init__( self, diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index 402077b017f8..cd7da3e1e3a0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -518,8 +518,7 @@ def run_cutlass_moe_fp4( *, activation_config: ApplyMoEActivationConfig | None = None, ) -> None: - """ - MoE implementation for FP4 Inputs + """MoE implementation for FP4 Inputs. # Gemm 1 a: Input tensor: [m, k] (half/bfloat16) diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 211364d24bca..fe95b36858d7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -58,8 +58,7 @@ def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool: def _valid_deep_gemm( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor ) -> bool: - """ - Check if the given problem size is supported by the DeepGemm grouped + """Check if the given problem size is supported by the DeepGemm grouped gemm kernel. All of M, N, K and the quantization block_shape must be aligned by `dg.get_m_alignment_for_contiguous_layout()`. """ diff --git a/vllm/model_executor/layers/fused_moe/experts/fallback.py b/vllm/model_executor/layers/fused_moe/experts/fallback.py index 639b2bf26681..e10d13f0a778 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fallback.py +++ b/vllm/model_executor/layers/fused_moe/experts/fallback.py @@ -30,8 +30,7 @@ def get_clses() -> tuple[ type[mk.FusedMoEExpertsModular], type[mk.FusedMoEExpertsModular], ]: - """ - Get the cls for the experts and fallback experts. + """Get the cls for the experts and fallback experts. Subclasses should implement this method, so that we have a consistent way to call the _supports_* diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py index d269c6f10993..f7e33711667a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py @@ -107,8 +107,7 @@ def workspace_shapes( expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - """ - Compute the shapes for the temporary and final outputs of the two gemms + """Compute the shapes for the temporary and final outputs of the two gemms and activation in the fused expert function. Since the gemms are independent, the workspace for the first gemm can be shared with the workspace for the last gemm. @@ -123,7 +122,6 @@ def workspace_shapes( - Note: in order for activation chunking to work, the first dimension of each tuple must be the number of tokens. """ - # We use global_num_experts due to how moe_align_block_size handles # expert_maps. K_dim = K * 2 if self.use_deep_ep_ll_nvfp4_dispatch else K @@ -212,8 +210,7 @@ def flashinfer_cutedsl_moe_masked( workspace: torch.Tensor, out: torch.Tensor, ): - """ - Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL + """Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL kernels. Args: @@ -229,13 +226,14 @@ def flashinfer_cutedsl_moe_masked( a2_global_scale (torch.Tensor): (l,) w2_blockscale (torch.Tensor): blockscale factors, e4m3, w2_alpha (torch.Tensor): (l,) + out (torch.Tensor): output tensor written in place, [l, m, k], bf16 masked_m (torch.Tensor): Masked dimension indices workspace (torch.Tensor): For gateup_output Notes: - Assumes max(masked_m) <= m. - """ + """ # === Assertions on dtypes === assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}" assert w1_blockscale.dtype == torch.float8_e4m3fn, ( diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index ec0e15bde558..d525822b912b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -29,8 +29,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): - """ - CuteDSL NvFP4 MoE experts using the FlashInfer functional API. + """CuteDSL NvFP4 MoE experts using the FlashInfer functional API. Uses Standard activation format (non-batched). The kernel handles routing, expert computation, and reduction internally. diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index 1614e41b3303..718e018c47e6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -38,8 +38,7 @@ def is_valid_flashinfer_cutlass_fused_moe( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor ) -> bool: - """ - Check if the given problem size is supported by the FlashInfer CUTLASS MoE + """Check if the given problem size is supported by the FlashInfer CUTLASS MoE kernel. """ if not has_flashinfer_cutlass_fused_moe(): @@ -219,8 +218,7 @@ def workspace_shapes( ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: # We use global_num_experts due to how moe_align_block_size handles # expert_maps. - """ - Compute the shapes for the temporary and final outputs of the two gemms + """Compute the shapes for the temporary and final outputs of the two gemms and activation in the fused expert function. Since the gemms are independent, the workspace for the first gemm can be shared with the workspace for the last gemm. diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index f16b351d353d..b8cd58635eb6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -551,8 +551,7 @@ def invoke_moe_batched_triton_kernel( class NaiveBatchedExperts(mk.FusedMoEExpertsModular): - """ - A reference MoE expert class that operates on expert batched format, + """A reference MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ @@ -775,8 +774,7 @@ def batched_moe_kernel_quantize_input( class BatchedTritonExperts(mk.FusedMoEExpertsModular): - """ - A Triton based MoE expert class that operates on expert batched format, + """A Triton based MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 86079605a367..52b6d4e75334 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -352,8 +352,7 @@ def _supports_quant_scheme( ) def _prequantizes_dispatch_activation(self) -> bool: - """ - Whether the prepare/finalize step should quantize activations before + """Whether the prepare/finalize step should quantize activations before the (EP all-to-all) dispatch instead of leaving it to Humming. This is enabled only for block-FP8 (group-128) activations: quantizing @@ -373,8 +372,7 @@ def _prequantizes_dispatch_activation(self) -> bool: @property def expects_unquantized_inputs(self) -> bool: - """ - Whether the prepare/finalize step should defer input quantization to + """Whether the prepare/finalize step should defer input quantization to the experts (by setting defer_input_quant=True and passing unquantized inputs). @@ -832,8 +830,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> None: - """ - Standard apply implementation for Humming indexed experts. + """Standard apply implementation for Humming indexed experts. Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual @@ -967,8 +964,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> None: - """ - Standard apply implementation for Humming grouped experts. + """Standard apply implementation for Humming grouped experts. Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual @@ -1090,8 +1086,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> None: - """ - Standard apply implementation for Humming batched grouped experts. + """Standard apply implementation for Humming batched grouped experts. Note: Humming kernels handle weights internally through the layer object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 8b87597ec268..e2b4987bd050 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -502,8 +502,7 @@ def pack_bitmatrix( BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): - """ - Packs topk_ids into a bitmatrix. + """Packs topk_ids into a bitmatrix. code reference: https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264 """ @@ -949,8 +948,7 @@ def moe_problem_size( w2: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple[int, int, int, int, int]: - """ - Extract the MoE problem size from the given tensor arguments: + """Extract the MoE problem size from the given tensor arguments: - a: The hidden states, input to the MoE layer. - w1: The first set of expert weights. - w2: The second set of expert weights. @@ -1075,8 +1073,7 @@ def apply( class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): - """ - A Triton based MoE expert class that operates on expert standard + """A Triton based MoE expert class that operates on expert standard format and explicitly keeps the activation and reduction (moe_sum) steps unfused from the matmul_ogs kernel. This exposes injection points for activation and moe_sum. diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index 48232913229d..fb41ef544427 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py @@ -10,8 +10,7 @@ @dataclass class MoELoRAContext: - """ - Carries all LoRA state for one MoE forward pass. + """Carries all LoRA state for one MoE forward pass. Built by FusedMoEWithLoRA.forward() and propagated explicitly through the modular kernel path (FusedMoEKernel -> FusedMoEExpertsModular.apply) so diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py index 0aa6f6c30b6d..ac75e7bc798b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py @@ -7,8 +7,7 @@ class LoRAExpertsMixin: - """ - Mixin for FusedMoEExpertsModular subclasses that natively handle + """Mixin for FusedMoEExpertsModular subclasses that natively handle MoELoRAContext inside their apply() implementation. Mixing this class in: diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 378b6a2314d2..a6a70ca95264 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -251,26 +251,56 @@ def fused_marlin_moe( input_dtype: torch.dtype | None = None, activation_config: ApplyMoEActivationConfig | None = None, ) -> torch.Tensor: - """ - This function computes a Mixture of Experts (MoE) layer using two sets of + """This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism. - Parameters: - - hidden_states (torch.Tensor): The input tensor to the MoE layer. - - w1 (torch.Tensor): The first set of expert weights. - - w2 (torch.Tensor): The second set of expert weights. - - w1_scale (torch.Tensor): Scale to be used for w1. - - w2_scale (torch.Tensor): Scale to be used for w2. - - topk_weights (torch.Tensor): Top-k weights. - - topk_ids (torch.Tensor): Indices of topk-k elements. - - w1_zeros (torch.Tensor|None): Optional zero points to be used for w1. - - w2_zeros (torch.Tensor|None): Optional zero points to be used for w2. - - num_bits (bool): The number of bits in expert weights quantization. + Args: + hidden_states (torch.Tensor): The input tensor to the MoE layer. + w1 (torch.Tensor): The first set of expert weights. + w2 (torch.Tensor): The second set of expert weights. + bias1 (torch.Tensor|None): Optional bias for the first GEMM. + bias2 (torch.Tensor|None): Optional bias for the second GEMM. + w1_scale (torch.Tensor): Scale to be used for w1. + w2_scale (torch.Tensor): Scale to be used for w2. + topk_weights (torch.Tensor): Top-k weights. + topk_ids (torch.Tensor): Indices of topk-k elements. + quant_type_id (int): Id of the `ScalarType` the expert weights are + quantized to. Determines the 4- or 8-bit Marlin kernel path. + apply_router_weight_on_input (bool): Apply the topk weights to the + input rather than the output. Only valid when topk is 1. + global_num_experts (int): Total number of experts across all expert + parallel shards. -1 means the local expert count. + activation (MoEActivation): Activation applied between the two GEMMs. + activation_func (Callable|None): Override for the activation + implementation. + moe_sum (Callable|None): Override for the final reduction over the + topk dimension. Defaults to `ops.moe_sum` / `torch.sum`. + expert_map (torch.Tensor|None): Maps global expert indices to the + local expert space of this expert parallel shard. + input_global_scale1 (torch.Tensor|None): NVFP4 global input scale for + the first GEMM. + input_global_scale2 (torch.Tensor|None): NVFP4 global input scale for + the second GEMM. + global_scale1 (torch.Tensor|None): NVFP4 global weight scale for w1. + global_scale2 (torch.Tensor|None): NVFP4 global weight scale for w2. + w1_zeros (torch.Tensor|None): Optional zero points to be used for w1. + w2_zeros (torch.Tensor|None): Optional zero points to be used for w2. + workspace (torch.Tensor|None): Scratch buffer for the Marlin kernels. + intermediate_cache13 (torch.Tensor|None): Scratch buffer for the + first and third intermediate activations. + intermediate_cache2 (torch.Tensor|None): Scratch buffer for the + second intermediate activation. + output (torch.Tensor|None): Tensor to write the result into. A new + tensor is allocated when omitted. + input_dtype (torch.dtype|None): dtype the activations are quantized + to. A 1-byte dtype raises the M block size to at least 16. + activation_config (ApplyMoEActivationConfig|None): Extra + configuration forwarded to the activation. Returns: - - torch.Tensor: The output tensor after applying the MoE layer. - """ + torch.Tensor: The output tensor after applying the MoE layer. + """ quant_type = ScalarType.from_id(quant_type_id) assert quant_type in [ scalar_types.uint4, @@ -399,8 +429,7 @@ def batched_fused_marlin_moe( activation_func: Callable[..., None] | None = None, activation_config: ApplyMoEActivationConfig | None = None, ) -> torch.Tensor: - """ - This function massages the inputs so the batched hidden_states can be + """This function massages the inputs so the batched hidden_states can be presented as a 2D contiguous tensor that could be used with _fused_marlin_moe. @@ -423,7 +452,6 @@ def batched_fused_marlin_moe( expert_ids tensors, so only relevant/valid rows of A (hidden_states) are accessed and are processed with the correct expert_ids. """ - assert hidden_states.ndim == 3, ( f"hidden states must be batched. e.g. [B, MAX_TOKENS, K]." f"But got {hidden_states.size()}" diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index c8063ffb4d63..d5217706e1b5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -34,7 +34,8 @@ def _select_cfg(M, N, K, block_m): """Pick the launch config from host constants only (M=num_valid_tokens, N, K, - block_m) — graph-capture safe (no GPU-scalar branch).""" + block_m) — graph-capture safe (no GPU-scalar branch). + """ # Per-regime winners (measured, isolated cuda-event A/B on gfx950, GPU 3): # BLOCK_K=256 (fewer K-iters + bigger MX scale-load coalesced with the dot), # num_stages=2 (software-pipeline overlaps the E8M0 scale-load with the scaled diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py index cd862cac5958..f52473c68fff 100644 --- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -NVFP4 quantization emulation for MoE. +"""NVFP4 quantization emulation for MoE. This file implements NVFP4 emulation for NVFP4 MOE in case the hardware used does not natively support NVFP4 MOE. @@ -86,8 +85,7 @@ def fused_moe_nvfp4_emulation_kernel( compute_type: tl.constexpr, group_size: tl.constexpr, ): - """ - Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM. + """Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM. Activations A are BF16 (already QDQ'd externally). Weights B are packed uint8 NVFP4 [E, N, K//2] — two FP4 values per byte @@ -334,8 +332,7 @@ def invoke_fused_moe_nvfp4_emulation_kernel( class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): - """ - Extension of TritonExperts to support emulated NVFP4 MoE experts. + """Extension of TritonExperts to support emulated NVFP4 MoE experts. It may be used for NVFP4 models when the device does not have native support for this dtype. @@ -427,8 +424,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - """ - Apply emulated quantized MoE computation. + """Apply emulated quantized MoE computation. This dequantizes the weights on the fly and calls fused_experts_impl with activation quantization support. diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index 2c91e22d9108..da0bca4977b7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -OCP MX quantization emulation for MoE. +"""OCP MX quantization emulation for MoE. This file implements OCP MX (MXFP4/MXFP6) emulation for MoE in case the hardware used does not natively support OCP MX MoE. @@ -50,6 +49,7 @@ def activation_quant_dtype( Raises: NotImplementedError: If the scheme has no emulated activation dtype. + """ if ocp_mx_scheme in { OCP_MX_Scheme.w_mxfp4, @@ -82,8 +82,7 @@ def activation_quant_dtype( class OCP_MXQuantizationEmulationTritonExperts(TritonExperts): - """ - Extension of TritonExperts to support emulated OCP MX MoE experts. + """Extension of TritonExperts to support emulated OCP MX MoE experts. It may be used for OCP MX (MXFP4/MXFP6) models when the device does not have native support for these dtypes. @@ -190,8 +189,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - """ - Apply emulated quantized MoE computation. + """Apply emulated quantized MoE computation. This dequantizes the weights on the fly and calls TritonExperts.apply with activation quantization support. diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 6262234679af..446acf46dd8b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -32,8 +32,7 @@ class TrtLlmBf16ExpertsBase: - """ - BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and + """BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and monolithic interfaces. """ @@ -124,9 +123,7 @@ def expects_unquantized_inputs(self) -> bool: class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular): - """ - BF16 unquantized TRTLLM-Gen MoE kernels. Supports modular interface. - """ + """BF16 unquantized TRTLLM-Gen MoE kernels. Supports modular interface.""" @staticmethod def _supports_parallel_config( @@ -224,9 +221,7 @@ def apply( class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonolithic): - """ - BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface. - """ + """BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.""" @staticmethod def _supports_parallel_config( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 6f256db960e8..e860597bcba3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -58,8 +58,7 @@ def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tens class TrtLlmFp8ExpertsBase: - """ - Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic + """Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic interfaces. """ @@ -194,9 +193,7 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): - """ - Fp8 TRTLLM-Gen MoE kernels. Supports modular interface. - """ + """Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.""" @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -330,9 +327,7 @@ def apply( class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic): - """ - Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. - """ + """Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.""" def supports_routing_replay_capture(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index a3d50011d919..4dc5076cf6db 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -LoRA-aware FlashInfer TRT-LLM MoE experts (BF16). +"""LoRA-aware FlashInfer TRT-LLM MoE experts (BF16). Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153: @@ -109,7 +108,7 @@ def _finalize_lora_kernel( class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular): - """LoRA-aware trtllm MoE experts""" + """LoRA-aware trtllm MoE experts.""" def __init__( self, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 159c77848a56..1c20bc1cbcdf 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -35,9 +35,7 @@ class TrtLlmMxfp4ExpertsBase: - """ - MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic. - """ + """MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic.""" def __init__( self, @@ -164,8 +162,7 @@ def expects_unquantized_inputs(self) -> bool: class TrtLlmMxfp4ExpertsMonolithic( TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsMonolithic ): - """ - Monolithic version of the MXFP4 TRTLLM kernel (router + experts). + """Monolithic version of the MXFP4 TRTLLM kernel (router + experts). Wraps flashinfer.trtllm_fp4_block_scale_moe(). """ @@ -287,8 +284,7 @@ def apply( class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular): - """ - Modular version of the MXFP4 TRTLLM kernel (just the experts). + """Modular version of the MXFP4 TRTLLM kernel (just the experts). Wraps flashinfer.trtllm_fp4_block_scale_routed_moe(). Moved from trtllm_moe.py. """ diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index a9e725f2965c..e3e9bc6f2ee0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -19,8 +19,7 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): - """ - FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface + """FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface (fused router + experts). Wraps flashinfer_trtllm_mxint4_moe(). diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 6a54c0d104ec..cea985830db8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -42,9 +42,7 @@ class TrtLlmNvFp4ExpertsBase: - """ - NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. - """ + """NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface.""" def __init__( self, @@ -300,9 +298,7 @@ def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int: class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular): - """ - Modular version of the implementation (just the experts). - """ + """Modular version of the implementation (just the experts).""" @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -467,9 +463,7 @@ def apply( class TrtLlmNvFp4ExpertsMonolithic( TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic ): - """ - Monolithic version of the kernel (router + experts). - """ + """Monolithic version of the kernel (router + experts).""" def supports_routing_replay_capture(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index dc619b2b12d2..e0610c878d2c 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -110,8 +110,7 @@ def fused_moe_kernel_gptq_awq( use_int4_w4a16: tl.constexpr, use_int8_w8a16: tl.constexpr, ): - """ - Implements the fused computation for a Mixture of Experts (MOE) using + """Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices. Key Parameters: @@ -353,8 +352,7 @@ def fused_moe_kernel( # Tensor-descriptor path for the A gather and B load in the K-loop. USE_TD: tl.constexpr = False, ): - """ - Implements the fused computation for a Mixture of Experts (MOE) using + """Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices. Key Parameters: @@ -1115,15 +1113,13 @@ def get_moe_configs( block_n: int | None = None, block_k: int | None = None, ) -> dict[int, Any] | None: - """ - Return optimized configurations for the fused MoE kernel. + """Return optimized configurations for the fused MoE kernel. The return value will be a dictionary that maps an irregular grid of batch sizes to configurations of the fused_moe kernel. To evaluate the kernel on a given batch size bs, the closest batch size in the grid should be picked and the associated configuration chosen to invoke the kernel. """ - # Avoid optimizing for the batch invariant case. Use default config if envs.VLLM_BATCH_INVARIANT: return None @@ -1187,6 +1183,7 @@ def _ensure_block_size_k_divisible( Returns: A valid BLOCK_SIZE_K that divides size_k and is divisible by group_size. + """ # Fast path: already valid if size_k % block_size_k == 0 and block_size_k % group_size == 0: @@ -1638,8 +1635,7 @@ def _get_config_quant_dtype( use_fp8_w8a8: bool, use_int8_w8a8: bool, ) -> None | torch.dtype | str: - """ - Get the quantization type based on the quantization strategy flags. + """Get the quantization type based on the quantization strategy flags. We don't have a quant_config at this point so we need to work backwards. A return type of None means no quantization is required because the input is unquantized or has been quantized prior to calling diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index 3a7c6eedfba7..c64aeadcc7a5 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -59,8 +59,7 @@ def create_weights( raise NotImplementedError def uses_weight_scale_2_pattern(self) -> bool: - """ - Returns True if this quantization method uses 'weight_scale_2' pattern + """Returns True if this quantization method uses 'weight_scale_2' pattern for per-tensor weight scales (e.g., FP4 variants), False otherwise. This method should be overridden by subclasses that use the @@ -75,8 +74,7 @@ def maybe_roundup_sizes( act_dtype: torch.dtype, moe_parallel_config: FusedMoEParallelConfig, ) -> tuple[int, int]: - """ - Given layer hidden size and intermediate size per partition and MoE + """Given layer hidden size and intermediate size per partition and MoE configurations, round up hidden_size and intermediate_size_per_partition if necessary. @@ -93,6 +91,7 @@ def maybe_roundup_sizes( - rounded_hidden_size is the possibly rounded up hidden size. - rounded_intermediate_size_per_partition is the possibly rounded up intermediate size per partition. + """ from .all2all_utils import maybe_roundup_layer_hidden_size @@ -119,8 +118,7 @@ def skip_forward_padding(self) -> bool: @property def has_unpadded_output(self) -> bool: - """ - Indicates that the hidden_states output might be the unpadded + """Indicates that the hidden_states output might be the unpadded hidden_states shape rather than the full padded shape. """ return False @@ -151,8 +149,7 @@ def apply( shared_experts: "SharedExperts | None", shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - """ - Apply the MoE operation using modular kernels. + """Apply the MoE operation using modular kernels. Args: layer: RoutedExperts instance containing weight parameters @@ -160,9 +157,11 @@ def apply( topk_weights: Expert weights from router topk_ids: Selected expert IDs from router shared_experts_input: Input for shared experts (if any) + shared_experts: The shared experts module, if any Returns: Output tensor from routed experts. + """ raise NotImplementedError @@ -173,15 +172,16 @@ def apply_monolithic( router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: - """ - Apply the MoE operation using monolithic kernels. + """Apply the MoE operation using monolithic kernels. Args: layer: RoutedExperts instance containing weight parameters x: Input tensor router_logits: Router logits (routing done internally) + input_ids: Token ids, for routers that condition on them Returns: Finalized routed states or a deferred-finalize output. + """ raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/hpc_moe.py b/vllm/model_executor/layers/fused_moe/hpc_moe.py index 410b2a733a6f..c7a078e26251 100644 --- a/vllm/model_executor/layers/fused_moe/hpc_moe.py +++ b/vllm/model_executor/layers/fused_moe/hpc_moe.py @@ -122,8 +122,7 @@ def workspace_shapes( ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: # We use global_num_experts due to how moe_align_block_size handles # expert_maps. - """ - Compute the shapes for the temporary and final outputs of the two gemms + """Compute the shapes for the temporary and final outputs of the two gemms and activation in the fused expert function. Since the gemms are independent, the workspace for the first gemm can be shared with the workspace for the last gemm. diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index a9c1a66d87dc..2ee3c5f374a3 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -152,6 +152,11 @@ def FusedMoEFactory( copy that naming convention here and handle any remapping in the load_weights function in each model implementation. + Args: + intermediate_pad: Padding added to the intermediate size, if any. + swiglu_alpha: Optional alpha parameter for the SwiGLU activation. + swiglu_beta: Optional beta parameter for the SwiGLU activation. + Args: num_experts: Number of experts in the model (global count) top_k: Number of experts selected for each token @@ -211,6 +216,7 @@ def FusedMoEFactory( Returns: MoERunner: Configured MoE execution pipeline ready for forward passes + """ vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 120c9e985753..d91743dd951f 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -84,9 +84,7 @@ class FusedMoEActivationFormat(Enum): - """ - The standard activation format (num_tokens, hidden dim). - """ + """The standard activation format (num_tokens, hidden dim).""" Standard = ("standard",) """ @@ -97,9 +95,7 @@ class FusedMoEActivationFormat(Enum): @dataclass class ExpertTokensMetadata: - """ - Metadata regarding expert-token routing. - """ + """Metadata regarding expert-token routing.""" expert_num_tokens: torch.Tensor | None expert_num_tokens_cpu: torch.Tensor | None @@ -122,9 +118,7 @@ def make_from_list( class TopKWeightAndReduce(ABC): - """ - An abstract base class for weight application and reduction implementations. - """ + """An abstract base class for weight application and reduction implementations.""" @abstractmethod def apply( @@ -135,8 +129,7 @@ def apply( topk_ids: torch.Tensor, apply_router_weight_on_input: bool, ) -> torch.Tensor: - """ - Apply topk_weights to the fused_experts_outputs and/or reduce. + """Apply topk_weights to the fused_experts_outputs and/or reduce. If an output tensor is not passed, it will be created in the function. """ @@ -185,8 +178,7 @@ def apply( class FusedMoEPrepareAndFinalize(ABC): - """ - An abstract base class for the [Quantize-Prepare] and [Finalize] steps + """An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above. There are two variants of this class: @@ -195,8 +187,7 @@ class FusedMoEPrepareAndFinalize(ABC): """ def post_init_setup(self, fused_experts: "FusedMoEExperts"): - """ - Initialize FusedMoEPrepareAndFinalizeModular settings that depend on + """Initialize FusedMoEPrepareAndFinalizeModular settings that depend on FusedMoEExpertsModular experts object. The FusedMoEPrepareAndFinalizeModular implementations that have such dependencies may choose to override this function. @@ -206,16 +197,14 @@ def post_init_setup(self, fused_experts: "FusedMoEExperts"): @property @abstractmethod def activation_format(self) -> FusedMoEActivationFormat: - """ - A property indicating the output format of the activations for the + """A property indicating the output format of the activations for the 'prepare' method. """ raise NotImplementedError @abstractmethod def topk_indices_dtype(self) -> torch.dtype | None: - """ - The PrepareFinalize All2All implementations generally constrain the + """The PrepareFinalize All2All implementations generally constrain the dtype of the topk_ids they support. This function returns the required topk indices dtype so it can be respected. Return None if there are no such restrictions. @@ -224,8 +213,7 @@ def topk_indices_dtype(self) -> torch.dtype | None: @abstractmethod def max_num_tokens_per_rank(self) -> int | None: - """ - Some PrepareFinalize All2All implementations are batched. Meaning, + """Some PrepareFinalize All2All implementations are batched. Meaning, they can process only as set of tokens at a time. This function returns the batch size i.e the maximum number of tokens the implementation can process at a time. @@ -239,15 +227,13 @@ def num_dispatchers(self) -> int: @abstractmethod def output_is_reduced(self) -> bool: - """ - Indicates whether or not the output of finalize is reduced across all + """Indicates whether or not the output of finalize is reduced across all ranks. """ raise NotImplementedError def supports_async(self) -> bool: - """ - Indicates whether or not this class implements prepare_async and + """Indicates whether or not this class implements prepare_async and finalize_async. """ return False @@ -255,8 +241,7 @@ def supports_async(self) -> bool: # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): - """ - An abstract base class for the [Quantize-Prepare] and [Finalize] steps + """An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above for the Modular case. """ @@ -272,8 +257,7 @@ def prepare( quant_config: FusedMoEQuantConfig, defer_input_quant: bool, ) -> PrepareResultType: - """ - Perform any quantization (and/or) dispatching needed for this kernel. + """Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - topk_ids: The topk ids. - topk_weights: The topk weights. @@ -309,8 +293,7 @@ def prepare_async( quant_config: FusedMoEQuantConfig, defer_input_quant: bool, ) -> tuple[Callable, ReceiverType] | ReceiverType: - """ - Perform any quantization (and/or) dispatching needed for this kernel + """Perform any quantization (and/or) dispatching needed for this kernel but do not wait for results from other workers. - a1: The (unquantized) input to the MoE layer. - a1_scale: Optional scales for a1 @@ -360,8 +343,7 @@ def finalize( apply_router_weight_on_input: bool, weight_and_reduce_impl: TopKWeightAndReduce, ) -> None: - """ - Perform any combine plus apply weights and perform a reduction on the + """Perform any combine plus apply weights and perform a reduction on the fused experts output. - output: The output tensor, written in place. Must be (M, K) shape. - fused_expert_output: The unweighted, unreduced output of the fused @@ -384,8 +366,7 @@ def finalize_async( apply_router_weight_on_input: bool, weight_and_reduce_impl: TopKWeightAndReduce, ) -> tuple[Callable, Callable] | Callable: - """ - Perform any combine plus apply weights and perform a reduction on the + """Perform any combine plus apply weights and perform a reduction on the fused experts output but do not wait for results from other workers. - output: The output tensor, written in place. Must be (M, K) shape. - fused_expert_output: The unweighted, unreduced output of the fused @@ -419,8 +400,7 @@ def finalize_async( class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): - """ - An abstract base class for the [Quantize-Prepare] and [Finalize] steps + """An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above for the monolithic case. """ @@ -435,8 +415,7 @@ def prepare( quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, ) -> PrepareMonolithicResultType: - """ - Optional method for subclasses compatible with monolithic + """Optional method for subclasses compatible with monolithic FusedMoEExpertsModular kernels. Perform any quantization (and/or) dispatching needed for this kernel. @@ -453,8 +432,7 @@ def prepare( @abstractmethod def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: - """ - Optional method for subclasses compatible with monolithic + """Optional method for subclasses compatible with monolithic FusedMoEExpertsModular kernels. Perform any combine plus apply weights and perform a reduction on the @@ -485,8 +463,7 @@ def __init__( max_num_tokens: int | None = None, num_dispatchers: int | None = None, ): - """ - moe_config: MoE layer configuration. + """moe_config: MoE layer configuration. quant_config: Quantization parameters for this experts instance. """ if self.activation_format() == FusedMoEActivationFormat.Standard and ( @@ -521,8 +498,7 @@ def is_monolithic() -> bool: @property def expects_unquantized_inputs(self) -> bool: - """ - Whether or not the PrepareFinalize should defer input quantization + """Whether or not the PrepareFinalize should defer input quantization in the prepare step. If True, then the Experts kernel will execute the input quantization itself. @@ -533,8 +509,7 @@ def expects_unquantized_inputs(self) -> bool: @staticmethod @abstractmethod def activation_format() -> FusedMoEActivationFormat: - """ - A property which is a tuple of the input and output activation formats + """A property which is a tuple of the input and output activation formats for the 'apply' method. """ raise NotImplementedError @@ -595,8 +570,7 @@ def _make_reason(reason: str) -> str: @staticmethod @abstractmethod def _supports_current_device() -> bool: - """ - Whether the kernel supports the current device type + """Whether the kernel supports the current device type (compute cability and current platform). """ raise NotImplementedError @@ -604,8 +578,7 @@ def _supports_current_device() -> bool: @staticmethod @abstractmethod def _supports_no_act_and_mul() -> bool: - """ - Whether the kernel supports act_and_mul=False, i.e. + """Whether the kernel supports act_and_mul=False, i.e. non-gated MoE models like Nemotron-Nano. """ raise NotImplementedError @@ -621,16 +594,13 @@ def _supports_quant_scheme( @staticmethod @abstractmethod def _supports_activation(activation: MoEActivation) -> bool: - """ - Whether the kernel supports a particular act function. - """ + """Whether the kernel supports a particular act function.""" raise NotImplementedError @staticmethod @abstractmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """ - Whether the kernel supports deployment in particular parallel config. + """Whether the kernel supports deployment in particular parallel config. Can be overridden if a kernel does not support EP, SP or some other configuration. @@ -643,8 +613,7 @@ def _supports_routing_method( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - """ - Whether the kernel supports a routing method (e.g. GroupedTopK). + """Whether the kernel supports a routing method (e.g. GroupedTopK). Can be overridden by monolithic kernels that execute the router in addition to the experts if certain routers are not supported. @@ -656,8 +625,7 @@ def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - """ - Whether a kernel supports a particular dtype for router logits input. + """Whether a kernel supports a particular dtype for router logits input. Can be overridden by monolithic kernels that execute the router in addition to the experts if certain dtypes are not supported. @@ -666,16 +634,14 @@ def _supports_router_logits_dtype( @staticmethod def _supports_shape(hidden_dim: int) -> bool: - """ - Whether a kernel supports a particular shape. Can be overridden if a kernel + """Whether a kernel supports a particular shape. Can be overridden if a kernel has specific shape requirements. """ return True @staticmethod def _supports_batch_invariance() -> bool: - """ - Whether the kernel supports batch invariance, i.e. the output does not + """Whether the kernel supports batch invariance, i.e. the output does not depend on the order of the tokens in the input batch. This is useful for determining if the kernel can used with VLLM_BATCH_INVARIANT=1. """ @@ -764,17 +730,15 @@ def supports_lora() -> bool: return False def supports_packed_ue8m0_act_scales(self) -> bool: - """ - A flag indicating whether or not this class can process packed ue8m0 + """A flag indicating whether or not this class can process packed ue8m0 activation scales. """ return False class FusedMoEExpertsModular(FusedMoEExperts): - """ - An abstract base class for the [Permute-Experts-Unpermute] step described - above. + """An abstract base class for the [Permute-Experts-Unpermute] step described + above. """ @staticmethod @@ -788,8 +752,7 @@ def moe_problem_size( w2: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple[int, int, int, int, int]: - """ - Extract the MoE problem size from the given tensor arguments: + """Extract the MoE problem size from the given tensor arguments: - a: The hidden states, input to the MoE layer. - w1: The first set of expert weights. - w2: The second set of expert weights. @@ -825,9 +788,7 @@ def moe_problem_size( return E, M, N, K, topk def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: - """ - Workspace type: The dtype to use for the workspace tensors. - """ + """Workspace type: The dtype to use for the workspace tensors.""" return act_dtype @abstractmethod @@ -842,8 +803,7 @@ def workspace_shapes( expert_tokens_meta: ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - """ - Compute the shapes for the temporary and final outputs of the two gemms + """Compute the shapes for the temporary and final outputs of the two gemms and activation in the fused expert function. Since the gemms are independent, the workspace for the first gemm can be shared with the workspace for the last gemm. @@ -873,8 +833,7 @@ def workspace_shapes( @staticmethod def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: - """ - Calculate the output dimension for the activation function. + """Calculate the output dimension for the activation function. For *_no_mul activations (e.g. relu2_no_mul), there's no gate/up split, so output size equals input size (N). @@ -888,6 +847,7 @@ def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: Returns: The output dimension after activation. + """ return N if not activation.is_gated else N // 2 @@ -934,49 +894,50 @@ def apply( expert_tokens_meta: ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> None: - """ - This function computes the intermediate result of a Mixture of Experts + """This function computes the intermediate result of a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2. - Parameters: - - output: (torch.Tensor): The unweighted, unreduced output tensor. - - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE - layer. - - w1 (torch.Tensor): The first set of expert weights. - - w2 (torch.Tensor): The second set of expert weights. - - topk_weights: A map of row to expert weights. Some implementations - choose to do weight application. - - topk_ids (torch.Tensor): A map of row to expert id. - - activation (str): The activation function to apply after the first - MoE layer. - - global_num_experts (int): The total number of experts in the global - expert space. - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices - from the global expert space to the local expert space of the expert - parallel shard. - - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be - used for a1. Result of quantization from prepare/finalize and not - from the FusedMoEQuantConfig. - - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs - must be large enough to hold output of either MoE gemm. - - workspace2 (torch.Tensor): A scratch tensor used for the activation - function. - - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional - ExpertTokensMetadata object containing gpu/cpu tensors - as big as the number of local experts with the information about the - number of tokens assigned to each local expert. - - apply_router_weight_on_input: True if router weights are already - applied on the input. This is relevant if the implementation - chooses to do weight application. + Args: + output: (torch.Tensor): The unweighted, unreduced output tensor. + hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE + layer. + w1 (torch.Tensor): The first set of expert weights. + w2 (torch.Tensor): The second set of expert weights. + topk_weights: A map of row to expert weights. Some implementations + choose to do weight application. + topk_ids (torch.Tensor): A map of row to expert id. + activation (str): The activation function to apply after the first + MoE layer. + global_num_experts (int): The total number of experts in the global + expert space. + expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the + expert parallel shard. + a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be + used for a1. Result of quantization from prepare/finalize and not + from the FusedMoEQuantConfig. + a2_scale (Optional[torch.Tensor]): Optional quantized scale to be + used for the second gemm's activations. + workspace13 (torch.Tensor): A scratch tensor used for gemm outputs + must be large enough to hold output of either MoE gemm. + workspace2 (torch.Tensor): A scratch tensor used for the activation + function. + expert_tokens_meta (Optional[ExpertTokensMetadata]): An optional + ExpertTokensMetadata object containing gpu/cpu tensors + as big as the number of local experts with the information about + the number of tokens assigned to each local expert. + apply_router_weight_on_input: True if router weights are already + applied on the input. This is relevant if the implementation + chooses to do weight application. + """ raise NotImplementedError class FusedMoEExpertsMonolithic(FusedMoEExperts): - """ - An abstract base class for the [Permute-Experts-Unpermute] step described - above, but with the monolithic interface (accepts router logits - rather than topk ids and weights). + """An abstract base class for the [Permute-Experts-Unpermute] step described + above, but with the monolithic interface (accepts router logits + rather than topk ids and weights). """ @staticmethod @@ -985,8 +946,7 @@ def _supports_routing_method( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - """ - Whether the kernel supports a routing method (e.g. GroupedTopK). + """Whether the kernel supports a routing method (e.g. GroupedTopK). Monolithic kernels should explicitly opt-in to support. """ @@ -997,8 +957,7 @@ def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - """ - Whether the kernel supports a dtype for router logits. + """Whether the kernel supports a dtype for router logits. Modular kernels should opt-in to support. """ @@ -1083,8 +1042,7 @@ def apply( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: - """ - Same as apply(), except uses router_logits as opposed + """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed to the topk_ids and topk_weights. This is useful for kernels with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). """ @@ -1128,8 +1086,7 @@ def _allocate_buffers( expert_tokens_meta: ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Allocate temporary and output buffers for the fused experts op. + """Allocate temporary and output buffers for the fused experts op. Inputs: - out_dtype: output type of workspace and output tensors. - device: the device of the workspace and output tensors. @@ -1223,11 +1180,9 @@ def _prepare( torch.Tensor, torch.Tensor, ]: - """ - The _prepare method is a wrapper around self.prepare_finalize.prepare + """The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ - if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't # support async prepare/finalize @@ -1392,17 +1347,24 @@ def _finalize( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - """ - The _finalize method is a wrapper around self.prepare_finalize.finalize + """The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap. Args: + output: Tensor the finalized result is written into. + fused_out: The unweighted, unreduced output of the fused experts. + hidden_states: The input tensor to the MoE layer. + topk_weights: A map of row to expert weights. + topk_ids: A map of row to expert id. + apply_router_weight_on_input: True if the router weights were + already applied on the input. shared_experts: SharedExperts | None. The shared experts if any. shared_experts_input: Optional separate input for shared experts. When latent MoE is used, hidden_states is the latent-projected tensor (smaller dimension) used by routed experts, while shared_experts_input is the original hidden_states (full dimension) needed by the shared expert MLP. + """ if not self.prepare_finalize.supports_async(): assert not dbo_enabled() @@ -1463,33 +1425,34 @@ def apply( shared_experts: SharedExperts | None = None, shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: - """ - This function computes a Mixture of Experts (MoE) layer using two sets + """This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism. - Parameters: - - hidden_states: (torch.Tensor): The input tensor to the MoE layer. - - w1 (torch.Tensor): The first set of expert weights. - - w2 (torch.Tensor): The second set of expert weights. - - topk_weights (torch.Tensor): The topk weights applied at the end of the layer. - - topk_ids (torch.Tensor): A map of row to expert id. - - activation (MoEActivation): The activation function to apply after the first - MoE layer. - - global_num_experts (int): The total number of experts in the global - expert space. - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices - from the global expert space to the local expert space of the expert - parallel shard. - - apply_router_weight_on_input (bool): When true, the topk weights are - applied directly on the inputs. This is only applicable when topk is - 1. - - shared_experts: SharedExperts | None. The shared experts if any. - - shared_experts_input (Optional[torch.Tensor]): Optional separate - input for shared experts. For latent MoE, this is the original - hidden_states before latent projection. + Args: + hidden_states: (torch.Tensor): The input tensor to the MoE layer. + w1 (torch.Tensor): The first set of expert weights. + w2 (torch.Tensor): The second set of expert weights. + topk_ids (torch.Tensor): A map of row to expert id. + topk_weights (torch.Tensor): The topk weights applied at the end of + the layer. + activation (MoEActivation): The activation function to apply after + the first MoE layer. + global_num_experts (int): The total number of experts in the global + expert space. + expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the + expert parallel shard. + apply_router_weight_on_input (bool): When true, the topk weights are + applied directly on the inputs. This is only applicable when + topk is 1. + shared_experts: SharedExperts | None. The shared experts if any. + shared_experts_input (Optional[torch.Tensor]): Optional separate + input for shared experts. For latent MoE, this is the original + hidden_states before latent projection. Returns: - - torch.Tensor: The output tensor after applying the MoE layer. + torch.Tensor: The output tensor after applying the MoE layer. + """ output = torch.empty_like(hidden_states) @@ -1571,12 +1534,10 @@ def apply( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: - """ - Same as forward(), except uses router_logits as opposed + """Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( hidden_states, router_logits=router_logits, @@ -1674,8 +1635,7 @@ def supports_lora(self) -> bool: return self.fused_experts.supports_lora() def _post_init_setup(self): - """ - Resolve any leftover setup dependencies between self.prepare_finalize + """Resolve any leftover setup dependencies between self.prepare_finalize and self.fused_experts here. """ self.prepare_finalize.post_init_setup(self.impl.fused_experts) @@ -1685,8 +1645,7 @@ def _post_init_setup(self): ) def output_is_reduced(self) -> bool: - """ - Indicates whether or not the output of fused MoE kernel + """Indicates whether or not the output of fused MoE kernel is reduced across all ranks. """ return self.prepare_finalize.output_is_reduced() diff --git a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py index 7fc8bfcf824d..c8587d3b719b 100644 --- a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py +++ b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py @@ -16,8 +16,7 @@ def moe_align_block_size( pad_sorted_ids: bool = False, ignore_invalid_experts: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Aligns the token distribution across experts to be compatible with block + """Aligns the token distribution across experts to be compatible with block size for matrix multiplication. Note: In the case of expert_parallel, moe_align_block_size initially @@ -26,23 +25,23 @@ def moe_align_block_size( the current GPU rank as -1 so the MoE matmuls could skip those blocks. This requires the num_experts input arg to be the num global experts. - Parameters: - - topk_ids: A tensor of shape [total_tokens, top_k] representing the - top-k expert indices for each token. - - block_size: The block size used in block matrix multiplication. - - num_experts: The total number of experts. - - expert_map: A tensor of shape [num_experts] that maps the expert index - from the global space to the local index space of the current - expert parallel shard. If the expert is not in the current expert - parallel shard, the mapping is set to -1. - - pad_sorted_ids: A flag indicating whether the sorted_token_ids length - should be padded to a multiple of block_size, - - ignore_invalid_experts: A flag indicating whether to ignore invalid - experts. When False, all expert_ids in topk_ids will participate in - counting and ranking, but invalid experts in expert_ids will be marked - as -1. When True, all invalid expert_ids in topk_ids will be ignored - and will not participate in counting or ranking, and there will be no - -1 in expert_ids. + Args: + topk_ids: A tensor of shape [total_tokens, top_k] representing the + top-k expert indices for each token. + block_size: The block size used in block matrix multiplication. + num_experts: The total number of experts. + expert_map: A tensor of shape [num_experts] that maps the expert index + from the global space to the local index space of the current + expert parallel shard. If the expert is not in the current expert + parallel shard, the mapping is set to -1. + pad_sorted_ids: A flag indicating whether the sorted_token_ids length + should be padded to a multiple of block_size, + ignore_invalid_experts: A flag indicating whether to ignore invalid + experts. When False, all expert_ids in topk_ids will participate in + counting and ranking, but invalid experts in expert_ids will be marked + as -1. When True, all invalid expert_ids in topk_ids will be ignored + and will not participate in counting or ranking, and there will be no + -1 in expert_ids. Returns: - sorted_token_ids: A tensor containing the sorted token indices according @@ -70,6 +69,7 @@ def moe_align_block_size( the subsequent matrix multiplication. - The padding ensures that the total number of tokens is now divisible by block_size for proper block matrix operations. + """ max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) if pad_sorted_ids: @@ -106,8 +106,7 @@ def moe_align_block_size( def batched_moe_align_block_size( max_tokens_per_batch: int, block_size: int, expert_num_tokens: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Given num_batches, max_tokens_per_batch, block_size and the number of + """Given num_batches, max_tokens_per_batch, block_size and the number of valid-tokens in each batch, prepare sorted_token_ids, expert_ids and num_tokens_post_pad. sorted_token_ids, expert_ids and num_tokens_post_pad have the same semantics as in moe_align_block_size. @@ -115,12 +114,12 @@ def batched_moe_align_block_size( This function is intended to be a drop in replacement for moe_align_batch_size for the batched case. - Parameters: - - max_tokens_per_batch (int): Number of tokens in each batch (both - valid and invalid). - - block_size (int): block_size to align the data to. - - expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates - the number of valid tokens in batch i. + Args: + max_tokens_per_batch (int): Number of tokens in each batch (both + valid and invalid). + block_size (int): block_size to align the data to. + expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates + the number of valid tokens in batch i. Returns: - sorted_token_ids (torch.Tensor): Torch tensor of size @@ -132,6 +131,7 @@ def batched_moe_align_block_size( - num_tokens_post_pad (torch.Tensor): Torch tensor of size 1 indicating the number of valid blocks with actual data to process. This is represented in terms of num tokens. + Example: Let num_batches=5, max_tokens_per_batch=8, block_size=4, and expert_num_tokens=[2, 3, 0, 6, 8]. This expert_num_tokens tensor @@ -166,8 +166,8 @@ def batched_moe_align_block_size( num_tokens_post_pad will be 24 as sorted_token_ids has valid entries until 24. - """ + """ B = expert_num_tokens.size(0) device = expert_num_tokens.device diff --git a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py index d8ad9b94557d..8d16a419756a 100644 --- a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py +++ b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py @@ -92,8 +92,7 @@ def moe_fused_mul_sum( expert_map: torch.Tensor | None = None, num_valid_tokens: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Fused kernel for MoE (Mixture of Experts) to perform weighted summation + """Fused kernel for MoE (Mixture of Experts) to perform weighted summation of expert outputs. Args: @@ -122,6 +121,7 @@ def moe_fused_mul_sum( Returns: The fused weighted sum of expert outputs. Shape: (num_tokens, hidden_size). + """ assert inputs.ndim == 3 assert topk_weights.ndim == 2 diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index df1a2685ecbe..04081cf6d17b 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -57,8 +57,8 @@ def convert_flashinfer_moe_output( Raises: ValueError: If FlashInfer returns an unexpected layout. - """ + """ if num_tokens < 0: raise ValueError("num_tokens must be non-negative.") if top_k <= 0: diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index ad9fb509a1f7..7a53a960e895 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -112,20 +112,24 @@ def moe_permute( permuted_hidden_states: torch.Tensor | None = None, scratch: MoEPermuteScratch | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - This function expands and permutes activation to gather uncontinuous tokens + """This function expands and permutes activation to gather uncontinuous tokens for each expert. - Parameters: - - hidden_states (torch.Tensor): The input tensor to the MoE layer. - - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states - - topk_ids (torch.Tensor): topk expert route id for each token. - - n_expert (int): The number of expert. - - n_local_expert (int): The number of expert in current EP rank. - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices - from the global expert space to the local expert space of the expert - parallel shard. - - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. - If None, the output tensor will be created in this function. + + Args: + hidden_states (torch.Tensor): The input tensor to the MoE layer. + a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states + topk_ids (torch.Tensor): topk expert route id for each token. + n_expert (int): The number of expert. + n_local_expert (int): The number of expert in current EP rank. + expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. + If None, the output tensor will be created in this function. + scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch + buffers. Validated against hidden_states and topk_ids when given, + otherwise the buffers are allocated in this function. + Returns: - permuted_hidden_states (torch.Tensor): permuted activation. - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states @@ -134,6 +138,7 @@ def moe_permute( of each expert for standard grouped gemm. - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute. - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden. + """ n_token, n_hidden = hidden_states.size() topk = topk_ids.size(1) @@ -249,19 +254,21 @@ def moe_unpermute( inv_permuted_idx: torch.Tensor, expert_first_token_offset: torch.Tensor | None = None, ) -> None: - """ - This function expands and permutes activation to gathering uncontinuous + """This function expands and permutes activation to gathering uncontinuous tokens for each expert. - Parameters: - - out (torch.Tensor): output tensor - - permuted_hidden_states (torch.Tensor): permuted activation. - - topk_weights (torch.Tensor): topk expert route weight for each token. - - inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute. - - expert_first_token_offset (Optional[torch.Tensor]): offset of the first - token of each expert for grouped gemm. + + Args: + out (torch.Tensor): output tensor + permuted_hidden_states (torch.Tensor): permuted activation. + topk_weights (torch.Tensor): topk expert route weight for each token. + inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute. + expert_first_token_offset (Optional[torch.Tensor]): offset of the first + token of each expert for grouped gemm. + Returns: - hidden_states (torch.Tensor): The reduced and unpermuted activation tensor. + """ topk = topk_weights.size(1) n_hidden = permuted_hidden_states.size(-1) diff --git a/vllm/model_executor/layers/fused_moe/oracle/base.py b/vllm/model_executor/layers/fused_moe/oracle/base.py index 06be6becf83f..87c540409ffe 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/base.py +++ b/vllm/model_executor/layers/fused_moe/oracle/base.py @@ -55,24 +55,28 @@ class MoEKernelOracle(ABC, Generic[BackendT]): @abstractmethod def backend_enum_cls(self) -> type[BackendT]: """Return the concrete `Enum` class enumerating this oracle's - backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`).""" + backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`). + """ @abstractmethod def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]: """Return platform-appropriate backends in priority order for - this `moe_config`.""" + this `moe_config`. + """ @abstractmethod def backend_to_kernel_cls( self, backend: BackendT ) -> list[type[mk.FusedMoEExperts]]: """Map a backend enum value to its concrete `FusedMoEExperts` - subclasses, in selection priority order.""" + subclasses, in selection priority order. + """ @abstractmethod def map_backend(self, runner_backend: MoEBackend) -> BackendT: """Map a user-facing `MoEBackend` (from the runner config) to - this oracle's enum.""" + this oracle's enum. + """ @abstractmethod def select_backend( @@ -104,7 +108,8 @@ def make_kernel( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts - combinator) for the chosen backend.""" + combinator) for the chosen backend. + """ def convert_to_kernel_format( self, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index abe47373dea4..72e83cda4245 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -72,12 +72,10 @@ def _get_priority_backends( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> list[Fp8MoeBackend]: - """ - Get available backends in priority order based on platform and config. + """Get available backends in priority order based on platform and config. This function can be extended to become more complex as needed. """ - _AVAILABLE_BACKENDS = [ Fp8MoeBackend.AITER, Fp8MoeBackend.FLASHINFER_TRTLLM, @@ -273,8 +271,7 @@ def refine_fp8_moe_block_shape( config: FusedMoEConfig, weight_block_size: list[int], ) -> list[int] | None: - """ - Compute a refined block shape for block-quantized FP8 MoE weights whose + """Compute a refined block shape for block-quantized FP8 MoE weights whose checkpoint blocks cannot be sharded exactly across TP ranks. TP shards the intermediate dim of the expert weights, so a per-shard size @@ -308,11 +305,9 @@ def select_fp8_moe_backend( activation_key: QuantKey | None, allow_vllm_cutlass: bool = False, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts] | None]: - """ - Select the primary FP8 MoE backend + """Select the primary FP8 MoE backend Note: Shape-specific fallbacks may still occur at runtime. """ - # NOTE: the kernels are selected in the following order. AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key) @@ -456,7 +451,8 @@ def _humming_fp8_weight_schema( layer: RoutedExperts, weight: torch.Tensor, weight_scale: torch.Tensor ) -> dict[str, Any]: """Build the humming weight schema from the canonical on-device fp8/mxfp8 - tensors (scale dtype/shape, block size), not the producing quant method.""" + tensors (scale dtype/shape, block size), not the producing quant method. + """ # mxfp8: e8m0 group-32 scales (stored as uint8 bytes or e8m0). humming has # no compressed-tensors mxfp8 loader; its modelopt schema fits both sources. if weight_scale.dtype in (torch.uint8, torch.float8_e8m0fnu): @@ -614,8 +610,7 @@ def make_fp8_moe_quant_config( gemm1_beta: float | None = None, layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: - """ - Create FusedMoEQuantConfig for the specified FP8 Backend. + """Create FusedMoEQuantConfig for the specified FP8 Backend. The FusedMoEQuantConfig holds the scales that are used at runtime by the Modular Kernel abstraction. @@ -626,7 +621,6 @@ def make_fp8_moe_quant_config( In a future PR, we will have this function should be a method of the modular kernel itself. """ - # MARLIN and CPU are mixed precision W8A16 config. if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU: return fp8_w8a16_moe_quant_config( diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index cd15dceaf5c4..3fdd3ae0a8c4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -37,9 +37,7 @@ class Int8MoeBackend(Enum): def _get_priority_backends( moe_config: FusedMoEConfig, ) -> list[Int8MoeBackend]: - """ - Get available backends in priority order based on platform and config. - """ + """Get available backends in priority order based on platform and config.""" return [ Int8MoeBackend.TRITON, Int8MoeBackend.HUMMING, @@ -100,11 +98,9 @@ def select_int8_moe_backend( weight_key: QuantKey | None = kInt8StaticChannelSym, activation_key: QuantKey | None = kInt8DynamicTokenSym, ) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]: - """ - Select the primary Int8 MoE backend. + """Select the primary Int8 MoE backend. Note: Shape-specific fallbacks may still occur at runtime. """ - AVAILABLE_BACKENDS = _get_priority_backends(config) activation_format = ( @@ -228,7 +224,8 @@ def _humming_int8_weight_schema( weight: torch.Tensor, weight_scale: torch.Tensor ) -> dict[str, Any]: """Build the humming compressed-tensors int8 schema from the canonical - on-device tensors; humming does the signed-int8 -> native conversion.""" + on-device tensors; humming does the signed-int8 -> native conversion. + """ config: dict[str, Any] = { "quant_method": "compressed-tensors", "format": "int-quantized", diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index d47ad9c97e31..f4495137415e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -115,9 +115,7 @@ def backend_to_kernel_cls( def _get_priority_backends() -> list[WNA16MoEBackend]: - """ - Get available backends in priority order based on platform and config. - """ + """Get available backends in priority order based on platform and config.""" if current_platform.is_cpu(): return [WNA16MoEBackend.CPU] if current_platform.is_xpu(): @@ -241,11 +239,13 @@ def select_wna16_moe_backend( quant_config: Quantization structure and checkpoint format description. may_have_zp: Whether the integration can provide weight zero points. may_have_bias: Whether the integration can provide expert bias. + allow_tile_padding: Whether backends that require padding the weights + up to a tile boundary may be selected. Returns: A tuple of (``WNA16MoEBackend``, experts class or ``None``). - """ + """ activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if config.moe_parallel_config.use_batched_activation_format @@ -512,7 +512,8 @@ def _process_weights_flashinfer( def _pad_w13_shard_cols(x: torch.Tensor, unit: int, padded_unit: int) -> torch.Tensor: """Zero-pad each of the two gate/up shards of a ``(E, rows, 2 * unit)`` - tensor along its last dim, from ``unit`` to ``padded_unit`` columns.""" + tensor along its last dim, from ``unit`` to ``padded_unit`` columns. + """ if padded_unit == unit: return x e, rows, _ = x.shape @@ -1080,7 +1081,8 @@ def _humming_wna16_weight_schema( quant_config: QuantizationConfig | QuantizationArgs | None, ) -> dict[str, Any]: """Humming weight schema for a WNA16 checkpoint, derived from the quant - config rather than the running kernel.""" + config rather than the running kernel. + """ from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig @@ -1194,6 +1196,7 @@ def _unpack_and_dequant_int4_gptq( Returns: Dequantized weight tensor in the requested layout. + """ E, K_packed, N = w_int32.shape K = K_packed * 8 @@ -1256,6 +1259,7 @@ def _unpack_and_dequant_int4_awq( Returns: Dequantized weight tensor in the requested layout. + """ E, K, N_packed = w_int32.shape N = N_packed * 8 @@ -1442,6 +1446,15 @@ def convert_to_wna16_moe_kernel_format( layer: the ``MoERunner`` layer whose parameters are being prepared. quant_config: the ``QuantizationConfig`` for this layer. input_dtype: optional activation dtype, usually should be 16 bit. + w13: fused gate/up expert weights. + w2: down-projection expert weights. + w13_scale: quantization scales for ``w13``. + w2_scale: quantization scales for ``w2``. + w13_qzeros: optional zero points for ``w13``. + w2_qzeros: optional zero points for ``w2``. + w13_bias: optional bias for ``w13``. + w2_bias: optional bias for ``w2``. + """ if backend == WNA16MoEBackend.HUMMING: from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 1c38e4378b26..c0ce18194909 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -331,7 +331,8 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: """Available backends in priority order, BF16-act variant before - activation-quantized variant within each vendor family.""" + activation-quantized variant within each vendor family. + """ _AVAILABLE_BACKENDS = [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, @@ -355,8 +356,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: def _get_priority_backends() -> list[Mxfp4MoeBackend]: - """ - Get available backends in priority order. SM100+ prefers DeepGEMM FP4 / + """Get available backends in priority order. SM100+ prefers DeepGEMM FP4 / TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the backend-level ``is_supported_config`` check filters by device capability). """ @@ -411,7 +411,8 @@ def _resolve_activation_key( model_activation_key: QuantKey | None, ) -> QuantKey | None: """Combine the model-supplied activation key with the user override. - Raises on conflict (both set and disagreeing).""" + Raises on conflict (both set and disagreeing). + """ user_override = _user_moe_activation_override() if user_override is None: return model_activation_key @@ -461,7 +462,8 @@ def _filter_by_activation( ) -> list[Mxfp4MoeBackend]: """Pick variants matching ``requested_activation_key``; without one, prefer BF16 if the list has any, else keep the list as-is so explicit - non-BF16 picks (e.g. the ``_afp8`` aliases) still land.""" + non-BF16 picks (e.g. the ``_afp8`` aliases) still land. + """ if requested_activation_key is not None: return [ b @@ -526,8 +528,7 @@ def select_mxfp4_moe_backend( config: FusedMoEConfig, activation_key: QuantKey | None = None, ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: - """ - Select the primary MXFP4 MoE backend. + """Select the primary MXFP4 MoE backend. Args: config: MoE configuration @@ -536,6 +537,7 @@ def select_mxfp4_moe_backend( Use kFp8StaticTensorSym for W4A8 scheme. Note: Shape-specific fallbacks may still occur at runtime. + """ runner_backend = config.moe_backend requested_activation_key = _resolve_activation_key(activation_key) @@ -648,8 +650,7 @@ def select_mxfp4_moe_backend( def select_deepseek_v4_mxfp4_moe_backend( config: FusedMoEConfig, ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: - """ - Select the MXFP4 MoE backend with MXFP8 activation as top priority. + """Select the MXFP4 MoE backend with MXFP8 activation as top priority. Falls back through BF16 and other backends. """ activation_format = ( @@ -799,7 +800,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( torch.Tensor | None, ]: """Convert loaded weights into backend-specific kernel format.""" - if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 2a03eacea8e1..49f61764f0fd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -111,8 +111,8 @@ def select_mxfp8_moe_backend( Returns: A tuple of (fp8_backend, experts_cls). - """ + """ runner_backend = config.moe_backend if runner_backend != "auto": backend = _BACKEND_NAME_MAP.get(runner_backend) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index a860003775d5..ef6a9695f3c3 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -184,11 +184,9 @@ def select_nvfp4_moe_backend( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]: - """ - Select the primary NvFP4 MoE backend + """Select the primary NvFP4 MoE backend Note: Shape-specific fallbacks may still occur at runtime. """ - # NOTE: the kernels are selected in the following order. # FLASHINFER_B12X is intentionally excluded from auto-selection until # the upstream CUTLASS SM121 MMA op guard is resolved; use diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index b24f2a9a797e..910178705933 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -46,8 +46,7 @@ class UnquantizedMoeBackend(Enum): def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]: - """ - Get available backends in priority order based on platform and config. + """Get available backends in priority order based on platform and config. This function can be extended to become more complex as needed. """ @@ -208,11 +207,9 @@ def _trtllm_bf16_lora_supported(moe_config: FusedMoEConfig) -> bool: def select_unquantized_moe_backend( moe_config: FusedMoEConfig, ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]: - """ - Select the primary Unquantized MoE backend. + """Select the primary Unquantized MoE backend. Note: Shape-specific fallbacks may still occur at runtime. """ - if current_platform.is_tpu(): return UnquantizedMoeBackend.TPU, None diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py index c66e43ccba8b..db0a30544373 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py @@ -32,8 +32,7 @@ class W4A8Int8MoeBackend(Enum): def _get_priority_backends( moe_config: FusedMoEConfig, ) -> list[W4A8Int8MoeBackend]: - """ - Get available backends in priority order based on platform and config. + """Get available backends in priority order based on platform and config. Currently only CPU INT4 backend is available for W4A8 INT8 MoE. """ @@ -75,8 +74,7 @@ def select_w4a8_int8_moe_backend( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]: - """ - Select the primary W4A8 Int8 MoE backend. + """Select the primary W4A8 Int8 MoE backend. Args: config: MoE configuration @@ -85,8 +83,8 @@ def select_w4a8_int8_moe_backend( Returns: Tuple of (backend, kernel_class) - """ + """ AVAILABLE_BACKENDS = _get_priority_backends(config) if not AVAILABLE_BACKENDS: @@ -160,8 +158,7 @@ def _return_or_raise( def make_w4a8_int8_moe_quant_config( block_shape: tuple[int, int] | None = None, ) -> FusedMoEQuantConfig: - """ - Create FusedMoEQuantConfig for W4A8 Int8 MoE. + """Create FusedMoEQuantConfig for W4A8 Int8 MoE. Args: block_shape: Quantization block shape (row, col). @@ -170,6 +167,7 @@ def make_w4a8_int8_moe_quant_config( Returns: FusedMoEQuantConfig with appropriate settings for W4A8 Int8 + """ # W4A8 Int8 uses static weight quantization, dynamic activation quantization # Weights are 4-bit (stored as int8, packed to uint8), @@ -196,8 +194,7 @@ def pack_int4_weights_for_kleidi( out_features: int, group_size: int, ) -> torch.Tensor: - """ - Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format. + """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format. Args: int4_as_int8: [out, in] int8 tensor with values in [-8, 7] @@ -209,6 +206,7 @@ def pack_int4_weights_for_kleidi( Returns: Packed weight tensor in KleidiAI format + """ # Shift to unsigned nibble [0, 15] tmp = int4_as_int8.add(8) @@ -250,8 +248,7 @@ def convert_to_w4a8_int8_moe_format( torch.Tensor | None, torch.Tensor | None, ]: - """ - Pack INT4 MoE weights to KleidiAI format. + """Pack INT4 MoE weights to KleidiAI format. This function packs the INT4 weights (stored as int8 values) into the format expected by the KleidiAI dynamic_4bit_int_moe kernel. @@ -267,6 +264,7 @@ def convert_to_w4a8_int8_moe_format( Returns: Tuple of (w13_packed, w2_packed) tensors + """ # Derive dimensions from tensor shapes E = w13_weight.shape[0] # num_experts @@ -315,8 +313,7 @@ def make_w4a8_int8_moe_kernel( experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: - """ - Create FusedMoEKernel for W4A8 Int8 MoE. + """Create FusedMoEKernel for W4A8 Int8 MoE. Args: moe_quant_config: Quantization configuration @@ -326,6 +323,7 @@ def make_w4a8_int8_moe_kernel( Returns: Configured FusedMoEKernel instance + """ # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py index d11beaf98fbd..5a818bfa57c7 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py @@ -16,8 +16,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - A reference prepare/finalize class that reorganizes the tokens into + """A reference prepare/finalize class that reorganizes the tokens into expert batched format, i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index f30e0bea3c39..035d5552693d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py @@ -27,9 +27,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - Prepare/Finalize using DeepEP High-Throughput kernels. - """ + """Prepare/Finalize using DeepEP High-Throughput kernels.""" @staticmethod def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int: diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py index e5d2b601a768..1f357c79a873 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py @@ -34,9 +34,7 @@ def dequant_fp8( expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor ) -> torch.Tensor: - """ - Return dequantized tensor in fp32 - """ + """Return dequantized tensor in fp32.""" # TODO (varun) : Optimize leverage num_tokens_per_expert counts assert expert_x_fp8.is_contiguous() expert_x_scales = expert_x_scales.contiguous() @@ -50,9 +48,7 @@ def dequant_fp8( class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - Prepare/Finalize using DeepEP low-latency kernels. - """ + """Prepare/Finalize using DeepEP low-latency kernels.""" # DeepEP low-latency kernels are compiled only for certain # specific hidden sizes. diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index b3a146274d30..8bd215423da4 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -28,8 +28,7 @@ def _quantize_before_dispatch( quant_config: FusedMoEQuantConfig, defer_input_quant: bool ) -> bool: - """ - Do quantized dispatch for blockfp8 and mxfp8, unless the + """Do quantized dispatch for blockfp8 and mxfp8, unless the subsequent moe kernel requires bf16 inputs. """ if defer_input_quant: @@ -69,8 +68,7 @@ def _unpack_mxfp8_scale( class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API). + """Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API). Supports two modes controlled by the `use_cudagraph` constructor arg: diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index fe3a53941806..35e15325b4db 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py @@ -13,9 +13,7 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - Prepare/Finalize using MoRI kernels. - """ + """Prepare/Finalize using MoRI kernels.""" def __init__( self, @@ -60,8 +58,7 @@ def prepare( quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, ) -> mk.PrepareResultType: - """ - Returns a tuple of: + """Returns a tuple of: - quantized + dispatched a. - Optional quantized + dispatched a1_scales. - Optional ExpertTokensMetadata containing gpu/cpu tensors diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index 89f3843cc504..97d4e0d284c6 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -69,8 +69,7 @@ def _unwrap_scale_and_prepare_for_moe( class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): - """ - Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + """Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. Uses Torch AR/RS or AR for dispatch/combine operations, applied to the topk weights and ids. @@ -121,7 +120,6 @@ def prepare( defer_input_quant: bool = False, ) -> mk.PrepareResultType: """Quantize and Dispatch Topk Weights and Topk Ids.""" - if apply_router_weight_on_input: topk = topk_ids.size(1) assert topk == 1, ( @@ -210,8 +208,7 @@ def finalize( class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic): - """ - Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + """Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. Uses Torch AR/RS or AR for dispatch/combine operations, applied to the router logits (the MoE kernel runs the router internally). @@ -250,7 +247,6 @@ def prepare( defer_input_quant: bool = False, ) -> mk.PrepareMonolithicResultType: """Quantize and Dispatch Router Logits.""" - a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch( a1, quant_config, defer_input_quant ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index b2c85ce0c0ae..6d5dbee3f491 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -33,9 +33,7 @@ def dequant_fp8( expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor ) -> torch.Tensor: - """ - Return dequantized tensor in fp32 - """ + """Return dequantized tensor in fp32.""" assert expert_x_fp8.is_contiguous() expert_x_scales = expert_x_scales.contiguous() num_experts = expert_x_fp8.size(0) @@ -48,9 +46,7 @@ def dequant_fp8( class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): - """ - Prepare/Finalize using NIXL EP kernels. - """ + """Prepare/Finalize using NIXL EP kernels.""" # NIXL EP kernels are compiled only for certain specific hidden sizes. # NOTE: Keep this list sorted, maybe_roundup_layer_hidden_size depends diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 257e19b034a4..88841070f03b 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -44,8 +44,7 @@ class FusedMoeWeightScaleSupported(Enum): @PluggableLayer.register("routed_experts") class RoutedExperts(PluggableLayer): - """ - Container for routed expert weights and execution logic. + """Container for routed expert weights and execution logic. This module owns the expert weight parameters (w13_weight, w2_weight, scales, etc.) and handles: @@ -188,8 +187,7 @@ def _get_quant_method( quant_config: QuantizationConfig | None, moe_config: FusedMoEConfig, ) -> FusedMoEMethodBase: - """ - Helper method to ensure quant_method is never None and + """Helper method to ensure quant_method is never None and of the proper type. """ quant_method = None @@ -318,8 +316,7 @@ def _load_combined_w13_weight_scale( param: torch.Tensor, tp_rank: int, ): - """ - Load w13 weight scales assuming that w1 weight scales and w3 weight + """Load w13 weight scales assuming that w1 weight scales and w3 weight scales are stored in the same loaded_weight tensor. """ shard_size = param.shape[shard_dim] @@ -336,8 +333,7 @@ def _load_model_weight_or_group_weight_scale( loaded_weight: torch.Tensor, tp_rank: int, ): - """ - Load grouped weight scales for group quantization or model weights + """Load grouped weight scales for group quantization or model weights Args: shard_dim: dimension to shard @@ -345,6 +341,7 @@ def _load_model_weight_or_group_weight_scale( shard_id: either w1, w2, or w3 loaded_weight: checkpoint weight to load into the param tp_rank: tensor parallel rank + """ if shard_id == "w2": self._load_w2( @@ -444,6 +441,7 @@ def _narrow_expert_data_for_padding( Must be non-negative. shard_dim: The dimension index corresponding to the shard (intermediate) dimension. Defaults to `None`. + """ dims = (hidden_dim,) if shard_dim is None else (hidden_dim, shard_dim) if loaded_weight.ndim > 0: @@ -990,8 +988,7 @@ def build_expert_params_mapping( lora_base_layer_prefix_on_param_name: str = "", include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: - """ - Create expert parameter mapping for weight loading with redundant experts. + """Create expert parameter mapping for weight loading with redundant experts. This mapping handles the physical-to-logical expert ID conversion needed when loading weights with EPLB redundant experts. @@ -1011,6 +1008,7 @@ def build_expert_params_mapping( ``make_expert_params_mapping`` indexes the model-wide ``params_dict`` (prefix included). include_fused: Prepend the fused pre-fused-checkpoint entries + routed_experts_prefix: Prefix of the routed experts submodule Returns: List of tuples (param_name, weight_name, expert_id, shard_id) @@ -1019,6 +1017,7 @@ def build_expert_params_mapping( - weight_name: Weight name in checkpoint - expert_id: Physical expert ID - shard_id: Shard identifier (w1, w2, w3) + """ num_physical_experts = num_experts + num_redundant_experts @@ -1097,8 +1096,7 @@ def get_expert_weights(self) -> Iterable[torch.Tensor]: def _maybe_make_contiguous( name: str, p: torch.nn.Parameter ) -> torch.nn.Parameter: - """ - In some cases, the last 2 dimensions (the non-expert dimensions) + """In some cases, the last 2 dimensions (the non-expert dimensions) of the weight scale tensor are transposed. This function transforms the tensor (view update) so the tensor is contiguous(). Example: A non-contiguous scale tensor, @@ -1178,8 +1176,7 @@ def forward_modular( shared_experts: "SharedExperts | None" = None, shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Execute routed experts using the quantization method's apply function. + """Execute routed experts using the quantization method's apply function. This is called by the runner after router selection (for modular kernels) quant_method.apply() which accesses the weights on this RoutedExperts @@ -1194,6 +1191,7 @@ def forward_modular( Returns: Output tensor from routed experts. + """ assert not self.quant_method.is_monolithic @@ -1213,8 +1211,7 @@ def forward_monolithic( router_logits: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: - """ - Execute routed experts using the quantization method's apply function. + """Execute routed experts using the quantization method's apply function. This is called by the runner after router selection (for modular kernels) or with router logits (for monolithic kernels). It delegates to @@ -1228,6 +1225,7 @@ def forward_monolithic( Returns: Finalized routed states or a deferred-finalize output. + """ assert self.quant_method.is_monolithic diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index edd3e3beaa19..235808960263 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -136,8 +136,8 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: Args: layer_id: The layer index. topk_ids: Tensor of shape (batch_size, num_routed_experts). - """ + """ ctx = get_forward_context() if ctx.dp_metadata is None: # single dp start_loc = 0 @@ -241,6 +241,7 @@ def get_routed_experts( slot_mappings: Per-KV-cache-group slot mappings for this step, shape ``(num_kv_cache_groups, max_num_batched_tokens)``. num_tokens: Total number of tokens scheduled in this step. + """ return RoutedExpertsTensors( routing_data=self.device_buffer[:num_tokens].clone(), @@ -417,6 +418,7 @@ def get( Returns: Array of shape (num_tokens - token_start, num_layers, num_experts_per_tok). + """ bs = self.block_size block_ids_array = np.array(block_ids, dtype=np.int32) diff --git a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py index 3447f7ce790a..42d4b5d4ad33 100644 --- a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py @@ -16,8 +16,7 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter): - """ - ROCm AITER router for models with fused shared experts (e.g. Qwen3-MoE). + """ROCm AITER router for models with fused shared experts (e.g. Qwen3-MoE). When the AITER topk_softmax kernel supports sigmoid fusion, the routing softmax and shared-expert sigmoid are computed in a single kernel launch. diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 01e674b2b131..0509f02eadf5 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -157,8 +157,7 @@ def eplb_map_to_physical_and_record( class BaseRouter(FusedMoERouter): - """ - Base router class that provides common functionality for all router implementations. + """Base class providing common functionality for all router implementations. This class implements the template method pattern where select_experts() handles common pre-processing and post-processing, delegating the actual routing logic @@ -171,11 +170,11 @@ def __init__( global_num_experts: int, eplb_state: EplbLayerState | None = None, ): - """ - Args: - top_k: Number of experts to select per token - global_num_experts: Total number of experts - eplb_state: Optional EPLBLayerState for load balancing + """Args: + top_k: Number of experts to select per token + global_num_experts: Total number of experts + eplb_state: Optional EPLBLayerState for load balancing + """ super().__init__(eplb_state=eplb_state) self.top_k = top_k @@ -241,8 +240,7 @@ def _compute_routing( *, input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Compute the actual routing logic. + """Compute the actual routing logic. This method must be implemented by subclasses to provide the specific routing algorithm (e.g., grouped_topk, fused_topk, custom routing, etc.). @@ -251,9 +249,11 @@ def _compute_routing( hidden_states: Input hidden states router_logits: Router logits for expert selection indices_type: Desired dtype for expert indices (may be None) + input_ids: Token ids, for routers that condition on them Returns: tuple of (topk_weights, topk_ids) + """ raise NotImplementedError @@ -265,8 +265,7 @@ def _select_experts( *, input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Route the input hidden states to the top-k experts based on the + """Route the input hidden states to the top-k experts based on the router logits. This method implements the template method pattern: @@ -283,6 +282,7 @@ def _select_experts( **Compatibility**: When EPLB is not enabled, the returned ids are equivalent to global logical ids, so should be compatible with plain MoE implementations without redundant experts. + """ # Step 1: Validate EPLB state self._validate_eplb_state() diff --git a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py index f87b7d3ae8d0..def871f428c7 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py @@ -10,8 +10,7 @@ class FusedMoERouter(ABC): - """ - FusedMoERouter is an abstract class that provides a 'select_experts' + """FusedMoERouter is an abstract class that provides a 'select_experts' method that is used for routing hidden states based on router logits. """ @@ -54,8 +53,7 @@ def select_experts( *, input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Route the input hidden states to the top-k experts based on the + """Route the input hidden states to the top-k experts based on the router logits. Returns: @@ -66,8 +64,8 @@ def select_experts( **Compatibility**: When EPLB is not enabled, the returned ids are equivalent to global logical ids, so should be compatible with plain MoE implementations without redundant experts. - """ + """ topk_weights, topk_ids = self._select_experts( hidden_states, router_logits, diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 1f086b448802..f546dfd98c0a 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -217,8 +217,7 @@ def fp32_router_gemm_dispatch_impl( weight: torch.Tensor, allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: - """ - Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, + """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, otherwise optionally run the experimental BF16x3 kernel for medium/large SM100 router batches, then fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index 0f6ec15d4567..203725182491 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -66,8 +66,7 @@ def create_fused_moe_router( bias_vl: torch.Tensor | None = None, image_sentinel_lo: int = 0, ) -> FusedMoERouter: - """ - Factory function to create the appropriate FusedMoERouter subclass based on + """Factory function to create the appropriate FusedMoERouter subclass based on the provided parameters. The selection logic follows this priority order: @@ -114,8 +113,8 @@ def create_fused_moe_router( Returns: An instance of the appropriate FusedMoERouter subclass - """ + """ routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY if routing_strategy != "": return RoutingSimulatorRouter( diff --git a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py index 853e6e67583b..0ff5627d2f04 100644 --- a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py +++ b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py @@ -25,8 +25,7 @@ def route_tokens( top_k: int, indices_type: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Route tokens to experts. + """Route tokens to experts. Args: hidden_states: Input hidden states [num_tokens, hidden_size] @@ -36,13 +35,13 @@ def route_tokens( Returns: tuple of (topk_weights, topk_ids) + """ pass class DistributionBasedRouting(RoutingStrategy): - """ - Distribution-based random routing strategy with configurable distributions. + """Distribution-based random routing strategy with configurable distributions. This routing strategy randomly selects experts for each token based on different probability distributions. Currently supports uniform and normal @@ -50,8 +49,7 @@ class DistributionBasedRouting(RoutingStrategy): """ def __init__(self, distribution: str = "uniform", **distribution_params: Any): - """ - Initialize distribution-based routing. + """Initialize distribution-based routing. Args: distribution: Type of distribution to use for sampling @@ -61,6 +59,7 @@ def __init__(self, distribution: str = "uniform", **distribution_params: Any): chosen distribution For "uniform": No additional parameters needed For "normal": mean (default: 0.0), std (default: 1.0) + """ self.distribution = distribution.lower() self.distribution_params = distribution_params @@ -90,8 +89,7 @@ def route_tokens( top_k: int, indices_type: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Randomly select experts for each token using the specified distribution. + """Randomly select experts for each token using the specified distribution. Args: hidden_states: Input hidden states [num_tokens, hidden_size] @@ -103,6 +101,7 @@ def route_tokens( tuple of (topk_weights, topk_ids) where: - topk_weights: Weights based on distribution sampling - topk_ids: Expert indices sampled from the distribution + """ num_tokens = hidden_states.shape[0] num_experts = router_logits.shape[-1] @@ -129,7 +128,6 @@ def _sample_expert_ids( indices_type: torch.dtype, ) -> torch.Tensor: """Sample expert IDs based on the specified distribution.""" - if self.distribution == "uniform": # Generate random scores, and take the top-k to avoid duplicate topk_ids scores = torch.rand(num_tokens, num_experts, device=device) @@ -218,8 +216,7 @@ def get_distribution_info(self) -> dict: class RoutingSimulator: - """ - Token-to-Expert Routing Simulator. + """Token-to-Expert Routing Simulator. This class provides a framework for testing and comparing different routing strategies for MoE models. It can simulate routing behavior @@ -239,22 +236,22 @@ class RoutingSimulator: @classmethod def register_strategy(cls, name: str, strategy: RoutingStrategy): - """ - Register a custom routing strategy. + """Register a custom routing strategy. Args: name: Name of the strategy strategy: RoutingStrategy instance + """ cls._routing_strategies[name] = strategy @classmethod def get_available_strategies(cls) -> list[str]: - """ - Get list of available routing strategy names. + """Get list of available routing strategy names. Returns: List of available strategy names + """ return list(cls._routing_strategies.keys()) @@ -266,8 +263,7 @@ def simulate_routing( top_k: int, indices_type: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Simulate token-to-expert routing using the specified strategy. + """Simulate token-to-expert routing using the specified strategy. Args: hidden_states: Input hidden states [num_tokens, hidden_size] @@ -278,6 +274,7 @@ def simulate_routing( Returns: tuple of (topk_weights, topk_ids) + """ if strategy_name not in RoutingSimulator._routing_strategies: raise ValueError( diff --git a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py index 0107f8815678..c2dc86dd7cff 100644 --- a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py +++ b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py @@ -71,7 +71,8 @@ def _compute_routing( input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute routing with full bias, compute zero expert output, - mask zero expert IDs.""" + mask zero expert IDs. + """ topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, gating_output=router_logits, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 6e63563ab861..8d3bc7942872 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -225,8 +225,7 @@ def _unpack( class MoERunner(MoERunnerInterface): - """ - Standard MoE runner implementation for executing Mixture of Experts layers. + """Standard MoE runner implementation for executing Mixture of Experts layers. This is the primary concrete implementation of MoE execution logic, providing comprehensive support for standard MoE operations. It handles: @@ -704,7 +703,6 @@ def forward( 1. pytorch cannot handle union types in custom op signatures so _moe_forward and _moe_forward_shared must be split. """ - # Apply transform for routed experts (e.g., latent projection for # latent MoE). When the caller pre-applies the routed input transform # outside the runner (e.g. to overlap it on a separate stream), it @@ -1019,8 +1017,7 @@ def set_eplb_state( logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, ) -> None: - """ - Register the EPLB state in this layer. + """Register the EPLB state in this layer. This is used later in forward pass, where we get the expert mapping and record the load metrics in `expert_load_view`. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py index d579b181562a..4a1270ecbcbe 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py @@ -17,8 +17,7 @@ class MoERunnerInterface(PluggableLayer, ABC): - """ - Abstract base class for Mixture of Experts (MoE) runners. + """Abstract base class for Mixture of Experts (MoE) runners. This class defines the interface that all MoE runner implementations must follow. MoE runners are responsible for executing the forward pass of MoE layers, handling diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index f9543358ecee..585991dfe9df 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -9,8 +9,7 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): - """ - Useful in the case when some FusedMoEExpertsModular + """Useful in the case when some FusedMoEExpertsModular implementation does not perform weight application and reduction but cannot address the needs of all the compatible PrepareAndFinalize implementations. @@ -42,8 +41,7 @@ def apply( class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce): - """ - The fused_experts outputs have already been weight applied and reduced. + """The fused_experts outputs have already been weight applied and reduced. This implementation is a no-op. """ @@ -78,8 +76,7 @@ def apply( class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce): - """ - TopKWeightAndReduce implementation for a fused_experts output + """TopKWeightAndReduce implementation for a fused_experts output of shape (m, topk, K) """ @@ -122,8 +119,7 @@ def apply( class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce): - """ - TopKWeightAndReduce implementation for a fused_experts output + """TopKWeightAndReduce implementation for a fused_experts output of shape (num_experts, batch_size, K) """ diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index cfc3be3e8e12..34e116dfbcbb 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -69,6 +69,7 @@ def resolve_layer_fused_shared_expert( Raises: ValueError: If requested shared-expert fusion is quantization-incompatible. + """ # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported # that returns None if AITER is not available. @@ -164,20 +165,19 @@ def _count_expert_num_tokens( def count_expert_num_tokens( topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None ) -> torch.Tensor: - """ - Count the number to tokens assigned to each expert. + """Count the number to tokens assigned to each expert. - Parameters: - - topk_ids (torch.Tensor): Tensor mapping each token to its - list of experts. - - num_local_experts (int): Number of experts in this rank. - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices - from the global expert space to the local expert space of the expert - parallel shard. + Args: + topk_ids (torch.Tensor): Tensor mapping each token to its list of experts. + num_local_experts (int): Number of experts in this rank. + expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. Returns: A tensor of size num_local_experts, where tensor[i] holds the number of tokens assigned to the ith expert. + """ assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids" expert_num_tokens = torch.empty( @@ -202,8 +202,7 @@ def count_expert_num_tokens( def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor: - """ - Shrink the given tensor and apply the given view to it. This is + """Shrink the given tensor and apply the given view to it. This is used to resize the intermediate fused_moe caches. """ assert prod(v) <= x.numel(), ( @@ -226,8 +225,7 @@ def _fp8_quantize( per_act_token: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Perform fp8 quantization on the inputs. If a block_shape + """Perform fp8 quantization on the inputs. If a block_shape is provided, the output will be blocked. """ if block_shape is None: @@ -252,11 +250,9 @@ def _int8_quantize( per_act_token: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Perform int8 quantization on the inputs. If a block_shape + """Perform int8 quantization on the inputs. If a block_shape is provided, the output will be blocked. """ - # If weights are per-channel (per_channel_quant=True), then # activations apply per-token quantization. Otherwise, assume # activation tensor-wise fp8/int8 quantization, dynamic or static diff --git a/vllm/model_executor/layers/fused_qk_norm_rope.py b/vllm/model_executor/layers/fused_qk_norm_rope.py index d3c645488a83..a62174b338c4 100644 --- a/vllm/model_executor/layers/fused_qk_norm_rope.py +++ b/vllm/model_executor/layers/fused_qk_norm_rope.py @@ -170,6 +170,7 @@ def fused_qk_rmsnorm_rope_gate( Returns: (q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim). ``gate_out`` is the raw (pre-sigmoid) gate. + """ if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0: raise ValueError( diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index e18a35a70da6..a9b14c617000 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Producer side of the QuantizedActivation contract for activation layers. +"""Producer side of the QuantizedActivation contract for activation layers. Given an activation module and the downstream linear it feeds, fuse the activation with that linear's input quantization into a single kernel when the diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index d69415448fa6..c7dbe840db1d 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -A QuantizedActivation is a pre-quantized activation produced by a fused kernel +"""A QuantizedActivation is a pre-quantized activation produced by a fused kernel and consumed directly by a linear layer, letting the layer skip its own input quantization. A linear advertises the key its kernel can consume via expose_input_quant_key; the kernel validates and reads the activation via diff --git a/vllm/model_executor/layers/hpc/gated_mla.py b/vllm/model_executor/layers/hpc/gated_mla.py index f1e495a3a7f2..4edd42f164f5 100644 --- a/vllm/model_executor/layers/hpc/gated_mla.py +++ b/vllm/model_executor/layers/hpc/gated_mla.py @@ -107,6 +107,7 @@ def hpc_gated_mla_gemm( Returns: ``[m, n]`` bfloat16. + """ from hpc.gemm import gated_mla_gemm diff --git a/vllm/model_executor/layers/hpc/hpc_ihc.py b/vllm/model_executor/layers/hpc/hpc_ihc.py index 9a8065a052e3..aaf6ea92f20b 100644 --- a/vllm/model_executor/layers/hpc/hpc_ihc.py +++ b/vllm/model_executor/layers/hpc/hpc_ihc.py @@ -122,6 +122,7 @@ class HpcIHCPre(HpcModule): When given, its weight/eps are folded into the kernel so a single launch covers pre + RMSNorm (the caller then skips the separate layernorm). None keeps the original pre-only behaviour. + """ def __init__( diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index a3b9da4e1ead..c6e027afb5b9 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -87,6 +87,7 @@ def deep_select_topk( Returns: The (num_rows, topk) indices tensor. + """ assert input.dim() == 2 and input.stride(1) == 1 assert ( @@ -162,7 +163,8 @@ def resolve_backend( ) -> str: """Resolve the decode top-k implementation from the configured backend ("auto" = the pre-existing chain, or a validated explicit - value).""" + value). + """ if self._backend == "auto": return self._resolve_auto(logits, topk_tokens, num_rows) @@ -209,7 +211,8 @@ def _resolve_auto( self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: """The pre-existing priority chain: cooperative -> persistent -> - per_row. deep_select/flashinfer/torch are opt-in only.""" + per_row. deep_select/flashinfer/torch are opt-in only. + """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" if self._is_cuda and topk_tokens in (512, 1024, 2048): @@ -271,7 +274,8 @@ def forward( max_seq_len: int, ) -> None: """Run the resolved decode top-k implementation, writing into - topk_indices (int32, -1 fill for rows shorter than topk_tokens).""" + topk_indices (int32, -1 fill for rows shorter than topk_tokens). + """ backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) if backend == "deep_select": row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index cded9daf13f8..a83dd100e693 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -231,6 +231,7 @@ def __init__( device: Device to create parameters on dtype: Data type for parameters activation: Activation function name for gating + """ factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -335,9 +336,7 @@ def forward_xpu( class LayerNorm(nn.Module): - """ - Layer Normalization. - """ + """Layer Normalization.""" def __init__(self, dim: int, eps: float = 1e-6): super().__init__() diff --git a/vllm/model_executor/layers/lightning_attn.py b/vllm/model_executor/layers/lightning_attn.py index d3ea7fb211a0..1df196ca4d05 100644 --- a/vllm/model_executor/layers/lightning_attn.py +++ b/vllm/model_executor/layers/lightning_attn.py @@ -541,8 +541,7 @@ def lightning_attention( block_size: int = 256, kv_history: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Apply lightning attention algorithm + """Apply lightning attention algorithm to compute attention efficiently. Args: @@ -556,6 +555,7 @@ def lightning_attention( Returns: output: Attention output kv: Updated key-value history + """ d = q.shape[-1] e = v.shape[-1] @@ -610,8 +610,7 @@ def _linear_attn_decode_kernel( pad_slot_id: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - """ - Kernel for linear attention decoding with KV cache. + """Kernel for linear attention decoding with KV cache. This kernel computes attention for a single token using the KV cache. """ @@ -683,8 +682,7 @@ def linear_decode_forward_triton( slot_idx: torch.Tensor, BLOCK_SIZE: int = 32, ) -> torch.Tensor: - """ - Perform linear attention decoding using Triton kernels. + """Perform linear attention decoding using Triton kernels. Args: q: Query tensor of shape [B, H, 1, D] @@ -697,6 +695,7 @@ def linear_decode_forward_triton( Returns: output: Attention output tensor + """ B, H, _, D = q.shape assert k.shape == (B, H, 1, D) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 95dc852fb307..c8d52a6a106c 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -146,6 +146,7 @@ def create_weights( input_size: Size of the input dim of the weight across all ranks. output_size: Size of the output dim of the weight across all ranks. params_dtype: Datatype of the parameters. + """ raise NotImplementedError @@ -157,7 +158,8 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: """Apply the weights in layer to the input tensor. - Expects create_weights to have been called before on the layer.""" + Expects create_weights to have been called before on the layer. + """ raise NotImplementedError @@ -250,6 +252,7 @@ class LinearBase(PluggableLayer): prefix: Prefix for parameter names. return_bias: If true, return bias together with outputs in forward pass. disable_tp: If true, tensor parallelism will be disabled for this layer. + """ requires_unquantized_input = False @@ -335,6 +338,7 @@ class ReplicatedLinear(LinearBase): (e.g. model.layers.0.qkv_proj) return_bias: If true, return bias together with outputs in forward pass. disable_tp: Take no effect for replicated linear layers. + """ # --8<-- [end:replicated_linear] @@ -454,6 +458,7 @@ class ColumnParallelLinear(LinearBase): shard per rank (see ``DCPGroupColumnParallelLinear``). tp_size: Override the tensor-parallel world size used for sharding. Defaults to the global TP world size. + """ # --8<-- [end:column_parallel_linear] @@ -690,6 +695,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear): return_bias: If true, return bias together with outputs in forward pass. disable_tp: If true, all weights matrix won't be sharded, this layer will be treated as a "Replicated" MergedLinear. + """ def __init__( @@ -867,8 +873,7 @@ def _load_fused_module_from_checkpoint( loaded_weight: torch.Tensor, output_sizes: list[int] | None = None, ): - """ - Handle special case for models where MLP layers are already + """Handle special case for models where MLP layers are already fused on disk. In this case, we have no shard id. This function determines the shard id by splitting these layers and then calls the weight loader using the shard id. @@ -876,7 +881,6 @@ def _load_fused_module_from_checkpoint( An example of a model with these fused layers: https://huggingface.co/microsoft/Phi-3-mini-4k-instruct """ - current_shard_offset = 0 shard_offsets: list[tuple[int, int, int]] = [] output_sizes = output_sizes or self.output_sizes @@ -1018,6 +1022,7 @@ class QKVParallelLinear(ColumnParallelLinear): (e.g. model.layers.0.qkv_proj) return_bias: If true, return bias together with outputs in forward pass. disable_tp: If true, weights matrix won't be sharded through tp rank. + """ def __init__( @@ -1102,8 +1107,7 @@ def _get_shard_size_mapping(self, loaded_shard_id: str): def _load_fused_module_from_checkpoint( self, param: BasevLLMParameter, loaded_weight: torch.Tensor ): - """ - Handle special case for models where QKV layers are already + """Handle special case for models where QKV layers are already fused on disk. In this case, we have no shard id. This function determines the shard id by splitting these layers and then calls the weight loader using the shard id. @@ -1632,6 +1636,7 @@ class RowParallelLinear(LinearBase): | . | | A_p | - - + Arguments: input_size: first dimension of matrix A. output_size: second dimension of matrix A. @@ -1651,6 +1656,7 @@ class RowParallelLinear(LinearBase): (e.g. model.layers.0.down_proj) return_bias: If true, return bias together with outputs in forward pass. disable_tp: If true, weights matrix won't be sharded through tp rank. + """ # --8<-- [end:row_parallel_linear] diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 694bb6b03960..10f425a50352 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -74,9 +74,9 @@ def __init__( logits_as_input: bool = False, soft_cap: float | None = None, ) -> None: - """ - Args: - scale: A scaling factor to apply to the logits. + """Args: + scale: A scaling factor to apply to the logits. + """ super().__init__() self.scale = scale diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 10ed7d05fc8f..fa0eb257586d 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -16,8 +16,7 @@ class MambaBase(AttentionLayerBase): - """ - Base class for Mamba-like layers which support the v1 engine. + """Base class for Mamba-like layers which support the v1 engine. Inherit from this class if you implement a custom layer. """ @@ -44,8 +43,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: @abstractmethod def get_state_shape(self) -> Iterable[tuple[int, ...]]: - """ - Defines the shape of the state. + """Defines the shape of the state. For mamba layers this is usually a (conv_state, ssm_state) tuple. In this case, returns (conv_state_shape, ssm_state_shape). """ diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index a1a494e121ad..c8f5a9ba64c9 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -45,8 +45,7 @@ @PluggableLayer.register("olmo_hybrid_gated_delta_net_attention") class OlmoHybridGatedDeltaNetAttention(GatedDeltaNetAttention): - """ - Gated DeltaNet linear attention layer for OLMo Hybrid. + """Gated DeltaNet linear attention layer for OLMo Hybrid. This implements the linear attention mechanism that replaces sliding window attention in the hybrid architecture. @@ -294,9 +293,7 @@ def _forward_core( a: torch.Tensor, core_attn_out: torch.Tensor, ): - """ - Core attention computation (called by custom op). - """ + """Core attention computation (called by custom op).""" forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index a57c984db4c7..0cb6ba516068 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -639,9 +639,7 @@ def fix_query_key_value_ordering( mixed_qkvz: torch.Tensor, mixed_ba: torch.Tensor, ): - """ - Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. - """ + """Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.""" new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + ( self.num_k_heads // self.tp_size, ( @@ -692,8 +690,7 @@ def prepare_gdn_attention_core_inputs( mixed_ba: torch.Tensor, num_tokens: int, ): - """ - Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op. + """Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op. For gqa_interleaved_layout (Qwen3-Next): unpack the interleaved [ng, (hk + hk + np/ng*hv + np/ng*hv)] layout into contiguous qkv. @@ -864,7 +861,8 @@ def forward_hip( hidden_states: torch.Tensor, ) -> torch.Tensor: """ROCm forward using AITER Triton fused projection+attention when - available, otherwise falling back to the generic CUDA path.""" + available, otherwise falling back to the generic CUDA path. + """ if GDN_AITER_TRITON_AVAILABLE: num_tokens = hidden_states.size(0) projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) @@ -899,8 +897,7 @@ def forward_cuda( self, hidden_states: torch.Tensor, ) -> torch.Tensor: - """ - Forward pass with three parts: + """Forward pass with three parts: 1. Input projection 2. Core attention (custom op) 3. Output projection @@ -977,8 +974,7 @@ def forward_xpu( self, hidden_states: torch.Tensor, ) -> torch.Tensor: - """ - Forward pass with three parts: + """Forward pass with three parts: 1. Input projection 2. Core attention (custom op) 3. Output projection @@ -1221,6 +1217,7 @@ def _forward_core_rocm( z_out: **output** buffer for z (num_tokens, num_heads, head_dim); mutated in-place. core_attn_out: Pre-allocated output buffer for attention results. + """ forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -1279,6 +1276,7 @@ def _forward_core( b: beta gating vector (num_tokens, num_heads) a: alpha gating vector (num_tokens, num_heads) core_attn_out: Pre-allocated output buffer for attention results. + """ forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -1658,9 +1656,7 @@ def _forward_core_decode_non_spec( core_attn_out: torch.Tensor, attn_metadata: GDNAttentionMetadata, ): - """ - Core attention computation with a packed non-spec decode fast path. - """ + """Core attention computation with a packed non-spec decode fast path.""" non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 self_kv_cache = self.kv_cache # conv_state must be (..., dim, width-1) for the conv kernels. @@ -2030,8 +2026,7 @@ def fused_gdn_gating( beta: float = 1.0, threshold: float = 20.0, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Fused computation of g and beta for Gated Delta Net. + """Fused computation of g and beta for Gated Delta Net. g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) beta_output = b.sigmoid() TODO maybe use torch.compile to replace this triton kernel diff --git a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py index 7f3b3f8dbaab..75a7d097a8b1 100644 --- a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py @@ -523,7 +523,7 @@ def forward( output: torch.Tensor, positions: torch.Tensor, ) -> None: - """Forward method called by torch.ops.vllm.linear_attention""" + """Forward method called by torch.ops.vllm.linear_attention.""" torch.ops.vllm.linear_attention( hidden_states, output, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 5d09da4c3bfe..2b1b1f9c9df2 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -49,8 +49,7 @@ # --8<-- [start:mamba_mixer] @PluggableLayer.register("mamba_mixer") class MambaMixer(MambaBase, PluggableLayer): - """ - Compute ∆, A, B, C, and D the state space parameters and compute + """Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) ∆, B, C are input-dependent @@ -237,8 +236,7 @@ def forward(self, hidden_states: torch.Tensor, output: torch.Tensor): ) def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): - """ - Run the Mamba-1 SSM pipeline. + """Run the Mamba-1 SSM pipeline. Steps ----- @@ -258,7 +256,6 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): decode tokens), both sets of kernels are executed independently and their outputs are concatenated before the final output projection. """ - forward_context: ForwardContext = get_forward_context() attn_metadata_raw = forward_context.attn_metadata diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 1496515e7735..db85ca0dee3f 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -252,8 +252,7 @@ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: # --8<-- [start:mamba_mixer2] @PluggableLayer.register("mamba_mixer2") class MambaMixer2(MambaBase, PluggableLayer): - """ - Compute ∆, A, B, C, and D the state space parameters and compute + """Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) ∆, B, C are input-dependent @@ -1221,7 +1220,6 @@ def share_replayssm_ring_trackers( The final local layer in each group advances its cursors after every layer in that group has consumed the previous values. """ - replayssm_mixers: dict[str, MambaMixer2] = {} for layer_name in ordered_layer_names: layer = forward_context[layer_name] diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 8ca1d839556d..630c27ad2d61 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -261,8 +261,8 @@ def short_conv_state_shape( @classmethod def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int): """Compute the increase in group numbers to account for - replication in order to accompany the head shards.""" - + replication in order to accompany the head shards. + """ # in the case ngoups % tp_size == 0, this will be zero if ngroups % tp_size == 0: return 0 @@ -344,13 +344,13 @@ def append_kda_recoverssm_record( @dataclass class MambaCopySpec: - """ - Data class specifying the memory-copy parameters for Mamba states used for + """Data class specifying the memory-copy parameters for Mamba states used for prefix caching in align mode. Attributes: start_addr (int): Starting address for the memory copy operation. num_elements (int): Number of elements to copy from the starting address. + """ start_addr: int diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 8335c849666d..e8254e4f34eb 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -497,7 +497,7 @@ def causal_conv1d_fn( metadata=None, validate_data=False, ): - """support varlen + continuous batching when x is 2D tensor + """Support varlen + continuous batching when x is 2D tensor. x: (dim,cu_seq_len) cu_seq_len = total tokens of all seqs in that batch @@ -1109,8 +1109,7 @@ def causal_conv1d_update( validate_data=False, out: torch.Tensor | None = None, ): - """ - x: Input tensor which can take the following shapes: + """x: Input tensor which can take the following shapes: - `[batch, dim]` - single token prediction - `[batch, dim, seqlen]` - single or multiple tokens prediction diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index c552d93445e6..ec8511f24bae 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -123,8 +123,7 @@ def causal_conv1d_update_torch( bias: torch.Tensor | None = None, activation: str | None = None, ) -> torch.Tensor: - """ - Pure PyTorch fallback for causal_conv1d_update. + """Pure PyTorch fallback for causal_conv1d_update. Currently used as a fallback for Arm (aarch64) to leverage oneDNN/ACL F.conv1d kernels for batched decoding. """ diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py index 992f4ff68b28..ce1c0c8416eb 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py @@ -173,6 +173,7 @@ def chunk_gated_delta_rule_cutedsl( ``[1, T, Hv, V]`` and ``final_state`` has shape ``[N, Hv, V, K]``. When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of that buffer. + """ q = q.squeeze(0) k = k.squeeze(0) diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py index 7cdd652e7d2b..47458323e1ae 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py @@ -728,7 +728,6 @@ def h_cutedsl( num_stages: int = 2, ) -> None: """Compute H/V_new with the same argument order as the CUDA wrapper.""" - _, H, K_dim = K.shape _, Hv, V_dim = V.shape h_dtype = { diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index af45467886ea..f69368e273e9 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -59,7 +59,8 @@ def get_ssm_device_name() -> str: def _canonical_cache_dtype(cache_dtype: str) -> str: """Canonical key for config lookup. bf16 and fp16 share the same tuned - configs because the kernel only sees bit width when accessing state.""" + configs because the kernel only sees bit width when accessing state. + """ return "float16" if cache_dtype == "bfloat16" else cache_dtype @@ -67,8 +68,7 @@ def _canonical_cache_dtype(cache_dtype: str) -> str: def get_ssm_configs( headdim: int, dstate: int, cache_dtype: str ) -> dict[int, Any] | None: - """ - Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update* + """Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update* keyed by ``effective_batch = batch * nheads``, or ``None`` if no config file is found for the (headdim, dstate, cache_dtype, device) combination. @@ -161,7 +161,8 @@ def _try_get_optimal_ssm_config_cached( @contextmanager def override_ssm_config(config: tuple[int, int]): """Pin ``try_get_optimal_ssm_config`` to ``config`` for the duration of - the context. Used by the tuning benchmark to time specific configs.""" + the context. Used by the tuning benchmark to time specific configs. + """ global _ssm_config_override prev = _ssm_config_override _ssm_config_override = config @@ -515,30 +516,29 @@ def selective_state_update( enable_stochastic_rounding=False, cache_philox_rounds=0, ): - """ - Argument: - state: (batch, dim, dstate) or (batch, nheads, dim, dstate) - x: (batch, dim) or (batch, nheads, dim) - dt: (batch, dim) or (batch, nheads, dim) - A: (dim, dstate) or (nheads, dim, dstate) - B: (batch, dstate) or (batch, ngroups, dstate) - C: (batch, dstate) or (batch, ngroups, dstate) - D: (dim,) or (nheads, dim) - z: (batch, dim) or (batch, nheads, dim) - dt_bias: (dim,) or (nheads, dim) - null_block_id: int - if state_batch_indices is passed, lets the kernel identify - padded entries that will not be processed, - for example: state_batch_indices = [null_block_id, 1, 20, - null_block_id] in this case, the kernel will not process - entries at indices 0 and 3 - out: Preallocated ssm output tensor. Assume same shape as x. - In-place updated. - num_accepted_tokens: (batch,) - number of accepted tokens from previous verification step, - tells the kernel which initial state to use - cu_seqlens: (batch,) - length per sequence, for variable length in speculative decoding cases + """Argument: + state: (batch, dim, dstate) or (batch, nheads, dim, dstate) + x: (batch, dim) or (batch, nheads, dim) + dt: (batch, dim) or (batch, nheads, dim) + A: (dim, dstate) or (nheads, dim, dstate) + B: (batch, dstate) or (batch, ngroups, dstate) + C: (batch, dstate) or (batch, ngroups, dstate) + D: (dim,) or (nheads, dim) + z: (batch, dim) or (batch, nheads, dim) + dt_bias: (dim,) or (nheads, dim) + null_block_id: int + if state_batch_indices is passed, lets the kernel identify + padded entries that will not be processed, + for example: state_batch_indices = [null_block_id, 1, 20, + null_block_id] in this case, the kernel will not process + entries at indices 0 and 3 + out: Preallocated ssm output tensor. Assume same shape as x. + In-place updated. + num_accepted_tokens: (batch,) + number of accepted tokens from previous verification step, + tells the kernel which initial state to use + cu_seqlens: (batch,) + length per sequence, for variable length in speculative decoding cases """ if state.dim() == 3: state = state.unsqueeze(1) @@ -723,8 +723,7 @@ def selective_scan_fn( cu_chunk_seqlen=None, last_chunk_indices=None, ) -> torch.Tensor: - """ - u: (dim, total_length) for varlen or (batch, dim, seqlen) + """u: (dim, total_length) for varlen or (batch, dim, seqlen) applies changes in place. ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate) applies changes in place. @@ -769,9 +768,12 @@ def selective_scan_fn( initial_state_idx: (batch,), dtype int32 The pointer into cache_indices, where the cache block containing the initial state is located. - returns + + Returns + ------- output: (dim, total_length) for varlen or (batch, dim, seqlen) supports inplace replacement + """ if u.stride(-1) != 1: u = u.contiguous() diff --git a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py index 7cf5acf2fab3..7dcb54703fac 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py @@ -146,16 +146,17 @@ def _bmm_chunk_fwd_kernel( def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtype=None): - """ - Argument: + """Argument: a: (seqlen, ngroups, k) b: (seqlen, ngroups, k) chunk_size: int cu_chunk_seq_lens: (nchunks+1,) causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are guaranteed to be correct. + Return: out: (nchunks, ngroups, chunk_size, chunk_size) + """ seqlen, ngroups, k = a.shape assert b.shape == a.shape diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 8c645574b9e7..581cd8bb6d6f 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -175,8 +175,7 @@ def mamba_chunk_scan_combined_varlen( return_intermediate_states=False, state_dtype=None, ): - """ - Argument: + """Argument: x: (seqlen, nheads, headdim) dt: (seqlen, nheads) A: (nheads) @@ -198,7 +197,6 @@ def mamba_chunk_scan_combined_varlen( Return: varlen_states: (batch, nheads, headdim, dstate) """ - assert cu_seqlens is not None, "cu_seqlens must be provided assuming varlen input" assert seq_idx is not None diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 62e4307b37b7..b622924327e8 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Dispatch module for Mamba selective state update (SSU) backends. +"""Dispatch module for Mamba selective state update (SSU) backends. Provides a unified `selective_state_update` function that dispatches to the Triton, FlashInfer, or CPU backend based on the configured diff --git a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py index afae19c75feb..7f8900fa894e 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py +++ b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py @@ -60,8 +60,7 @@ def _cuda_memcpy_d2d(dst: int, src: int, size: int): class IpcBuffer: - """ - Allocates CUDA device memory and exchanges IPC handles with all ranks + """Allocates CUDA device memory and exchanges IPC handles with all ranks so that every rank holds a valid device pointer to every other rank's buffer. """ @@ -133,8 +132,7 @@ def __del__(self): def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int): - """ - Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000). + """Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000). This is the "slot empty" sentinel for the Lamport protocol: the kernel spin-waits until a value is *not* negative zero. """ @@ -153,8 +151,7 @@ def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int): class LamportWorkspace: - """ - Self-contained workspace for Lamport-based cross-GPU AllReduce. + """Self-contained workspace for Lamport-based cross-GPU AllReduce. Parameters ---------- @@ -170,6 +167,7 @@ class LamportWorkspace: process_group : optional ``torch.distributed`` process group for IPC handle exchange. ``None`` uses the default group. + """ def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None): @@ -212,7 +210,8 @@ def __init__(self, rank: int, world_size: int, comm_size: int, process_group=Non @property def workspace(self) -> torch.Tensor: """Device tensor (int64) that can be passed to the kernel - as ``void** workspace``.""" + as ``void** workspace``. + """ return self._workspace # ------------------------------------------------------------------ @@ -225,8 +224,7 @@ def compute_comm_size_for_minimax( world_size: int, fused_qk: bool = True, ) -> int: - """ - Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel. + """Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel. The kernel stores per-token variance scalars in the Lamport buffer: - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot @@ -271,8 +269,7 @@ def get_allreduce_workspace( max_tokens: int = 16384, process_group=None, ) -> torch.Tensor: - """ - Return a cached workspace tensor for the given (rank, world_size) pair. + """Return a cached workspace tensor for the given (rank, world_size) pair. On first call the workspace is allocated and IPC handles are exchanged; subsequent calls with the same arguments return the cached tensor. @@ -288,6 +285,7 @@ def get_allreduce_workspace( Maximum number of tokens per batch (used when ``comm_size is None``). process_group : optional ``torch.distributed`` process group. + """ if comm_size is None: comm_size = LamportWorkspace.compute_comm_size_for_minimax( diff --git a/vllm/model_executor/layers/pooler/abstract.py b/vllm/model_executor/layers/pooler/abstract.py index 82abef4f69ad..a32e38b98638 100644 --- a/vllm/model_executor/layers/pooler/abstract.py +++ b/vllm/model_executor/layers/pooler/abstract.py @@ -22,9 +22,7 @@ def get_supported_tasks(self) -> Set[PoolingTask]: raise NotImplementedError def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: - """ - Construct the updated pooling parameters to use for a supported task. - """ + """Construct the updated pooling parameters to use for a supported task.""" return PoolingParamsUpdate() @abstractmethod diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py index e9790cd2d243..2115e5184ed5 100644 --- a/vllm/model_executor/layers/pooler/seqwise/poolers.py +++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py @@ -42,8 +42,7 @@ class SequencePooler(Pooler): - """ - A layer that pools specific information from hidden states. + """A layer that pools specific information from hidden states. This layer does the following: 1. Extracts specific tokens or aggregates data based on pooling method. diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py index 96eb28cf2667..0dab2d5b1d57 100644 --- a/vllm/model_executor/layers/pooler/tokwise/poolers.py +++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py @@ -46,8 +46,7 @@ class TokenPooler(Pooler): - """ - A layer that pools specific information from hidden states. + """A layer that pools specific information from hidden states. This layer does the following: 1. Extracts specific tokens or aggregates data based on pooling method. diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 57dc2371318d..39e613a61147 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -82,6 +82,7 @@ def register_quantization_config(quantization: str): >>> >>> get_quantization_config("my_quant") + """ # noqa: E501 def _wrapper(quant_config_cls): diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index f8c7269add87..2c123b8768d7 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -393,6 +393,7 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase): Args: quant_config: The AWQ Marlin quantization config. + """ _kernel_backends_being_used: set[str] = set() @@ -701,7 +702,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: def _setup_kernel(self, layer: RoutedExperts) -> None: """Build the FusedMoEKernel for this layer.""" - self.moe_quant_config = self.get_fused_moe_quant_config(layer) self.moe_kernel = make_wna16_moe_kernel( moe_quant_config=self.moe_quant_config, @@ -880,6 +880,7 @@ class AutoAWQLinearMethod(BaseAWQLinearMethod): Args: quant_config: The AWQ quantization config. + """ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 893f85b207f4..e4478e897dae 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -305,6 +305,7 @@ class AutoGPTQLinearMethod(LinearMethodBase): Args: quant_config: The AutoGPTQ quantization config. + """ _kernel_backends_being_used: set[str] = set() @@ -685,7 +686,6 @@ def replace_or_register(name: str, val: torch.Tensor | None): def _setup_kernel(self, layer: RoutedExperts) -> None: """Build the FusedMoEKernel for this layer.""" - self.moe_quant_config = self.get_fused_moe_quant_config(layer) self.moe_kernel = make_wna16_moe_kernel( moe_quant_config=self.moe_quant_config, diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index d1f8df2df7bb..1492997caa9b 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -46,21 +46,24 @@ def create_weights( ): """Create weights for a layer. - The weights will be set as attributes of the layer.""" + The weights will be set as attributes of the layer. + """ raise NotImplementedError @abstractmethod def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Apply the weights in layer to the input tensor. - Expects create_weights to have been called before on the layer.""" + Expects create_weights to have been called before on the layer. + """ raise NotImplementedError # Not required functions def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Gather embeddings in the layer based on indices in the input tensor. - Expects create_weights to have been called before on the layer.""" + Expects create_weights to have been called before on the layer. + """ raise NotImplementedError # Not required functions @@ -73,7 +76,8 @@ def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module): Quantization methods that need special weight handling (e.g. repacked weights) override this. - Expects create_weights to have been called before on the layer.""" + Expects create_weights to have been called before on the layer. + """ layer.weight = embed_tokens.weight return layer @@ -86,8 +90,7 @@ def process_weights_after_loading(self, layer: nn.Module) -> None: def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: - """ - Not all quant methods have embedding implemented, so we need to check that + """Not all quant methods have embedding implemented, so we need to check that it exists for our given method. We check this by making sure the function has been changed from the base implementation. """ @@ -158,8 +161,7 @@ def override_quantization_method( user_quant: str | None, hf_config: Any = None, ) -> QuantizationMethods | None: - """ - Detects if this quantization method can support a given checkpoint + """Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances. @@ -169,6 +171,7 @@ def override_quantization_method( user_quant: The user-specified quantization method string. hf_config: The HuggingFace model config object (e.g. for model_type checks). May be None if not available. + """ return None @@ -203,6 +206,7 @@ def get_quant_method( Returns: The quantize method. None if the given layer doesn't support quant method. + """ raise NotImplementedError @@ -251,13 +255,13 @@ def get_checkpoint_weight_mapper() -> "WeightsMapper": def apply_vllm_mapper( # noqa: B027 self, hf_to_vllm_mapper: "WeightsMapper" ): - """ - Interface for models to update module names referenced in + """Interface for models to update module names referenced in quantization configs in order to reflect the vllm model structure Args: hf_to_vllm_mapper: maps from hf model structure (the assumed structure of the qconfig) to vllm model structure + """ # TODO (@kylesayrs): add implementations for all subclasses pass @@ -268,14 +272,14 @@ def maybe_update_config( # noqa: B027 hf_config: PretrainedConfig | None = None, revision: str | None = None, ): - """ - Interface to update values after config initialization. + """Interface to update values after config initialization. Args: model_name: The name of the model hf_config: The Hugging Face config of the model revision: The revision of the model Returns: + """ # TODO: revision is never passed currently in vllm.py, # but is used in subclasses, should we remove this parameter? diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 306a5400550e..8f134da1dfea 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -121,8 +121,7 @@ def get_name(self) -> QuantizationMethods: return "compressed-tensors" def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): - """ - Transform layer paths in config targets to match vLLM's naming. + """Transform layer paths in config targets to match vLLM's naming. The WeightsMapper is designed for weight paths, but some backends (e.g. transformers) use broad prefix mappings like "" -> "model." @@ -214,8 +213,7 @@ def get_quant_method( return None def _add_fused_moe_to_target_scheme_map(self): # XXXXXXXXXXXXXXXXXXXXXX - """ - Helper function to update target_scheme_map + """Helper function to update target_scheme_map since linear layers get fused into RoutedExperts targeting 'Linear' needs to also match RoutedExperts modules. @@ -269,8 +267,7 @@ def from_config(cls, config: dict[str, Any]) -> "CompressedTensorsConfig": def _parse_sparsity_config( cls, config: dict[str, Any] ) -> tuple[dict[str, SparsityCompressionConfig], list[str]]: - """ - Args: + """Args: config: The `quantization_config` dictionary from config.json Returns: @@ -278,6 +275,7 @@ def _parse_sparsity_config( 1. A dictionary mapping target layer names to their corresponding sparsity_config 2. A list of layer names to ignore for sparsity + """ if not (sparsity_config := config.get(SPARSITY_CONFIG_NAME)): return dict(), [] @@ -301,13 +299,13 @@ def _parse_sparsity_config( def _quantization_scheme_map_from_config( cls, config: dict[str, Any] ) -> QUANTIZATION_SCHEME_MAP_TYPE: - """ - Args: + """Args: config: The `quantization_config` dictionary from config.json Returns: A dictionary mapping target layer names to their corresponding quantization_args for weights and input activations + """ target_scheme_map: dict[str, Any] = dict() quant_format = cast(str, config.get("format")) @@ -667,7 +665,8 @@ def _is_wNa8o8_int( ) -> bool: """Weight N-bit INT (pack-quantized for sub-byte, int-quantized for 8-bit) with static per-tensor INT8 input/output activation quant, applied as a float - fake-quant around a weight-only matmul.""" + fake-quant around a weight-only matmul. + """ is_int_pack_format = format in ( CompressionFormat.pack_quantized.value, CompressionFormat.int_quantized.value, @@ -703,7 +702,8 @@ def _is_wNaM_int( format: str | None, ) -> bool: """Weight N-bit INT with symmetric dynamic INT activation quant - via Humming kernel.""" + via Humming kernel. + """ if input_quant is None: return False is_pack_format = format == CompressionFormat.pack_quantized.value @@ -895,8 +895,7 @@ def _get_scheme_from_parts( def get_scheme( self, layer: torch.nn.Module, layer_name: str | None = None ) -> "CompressedTensorsScheme | None": - """ - compressed-tensors supports non uniform in the following way: + """compressed-tensors supports non uniform in the following way: targets of config_groups: There can be N config_groups which each have a quantization scheme. Each config_group has a list of targets @@ -907,7 +906,6 @@ def get_scheme( use the quantization scheme corresponding to the matched target to select the CompressedTensorsScheme used for inference. """ - # Use the new get_quant_args method to extract QuantizationArgs scheme_dict = self.get_scheme_dict(layer, layer_name) @@ -944,8 +942,7 @@ def get_scheme( def get_scheme_dict( self, layer: torch.nn.Module, layer_name: str | None = None ) -> dict[str, QuantizationArgs | str | None] | None: - """ - Extract the QuantizationArgs for a given layer. + """Extract the QuantizationArgs for a given layer. Returns: dict with { @@ -953,6 +950,7 @@ def get_scheme_dict( "input_activations": QuantizationArgs | None, "format": str | None } | None + """ # TODO (@kylesayrs): support ignore module names with ct matching utils if should_ignore_layer( @@ -1004,8 +1002,7 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): - """ - Use the CompressedTensorsScheme associated with each layer to create + """Use the CompressedTensorsScheme associated with each layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ @@ -1026,8 +1023,7 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ): - """ - Use the output of create_weights and the CompressedTensorsScheme + """Use the output of create_weights and the CompressedTensorsScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details @@ -1039,8 +1035,7 @@ def apply( class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): - """ - Supports loading kv-cache scaling factors from compressed-tensors + """Supports loading kv-cache scaling factors from compressed-tensors checkpoints. """ @@ -1050,12 +1045,12 @@ def __init__(self, quant_config: CompressedTensorsConfig): @staticmethod def validate_kv_cache_scheme(kv_cache_scheme: dict[str, Any] | None): - """ - Validator for the kv cache scheme. Useful for controlling the + """Validator for the kv cache scheme. Useful for controlling the kv cache quantization schemes, that are being supported in vLLM Args: kv_cache_scheme: the compressed-tensors kv cache scheme + """ if kv_cache_scheme is None: return @@ -1091,8 +1086,7 @@ def validate_kv_cache_scheme(kv_cache_scheme: dict[str, Any] | None): ) def create_weights(self, layer: torch.nn.Module): - """ - Initialize placeholder scales and zero points to enable loading of + """Initialize placeholder scales and zero points to enable loading of quantized params from compressed-tensors checkpoints. """ strategy = None # for backward compatibility @@ -1199,8 +1193,7 @@ def _tp_aware_loader( ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """ - Override the default vLLM placeholder scales with the llm-compressor loaded + """Override the default vLLM placeholder scales with the llm-compressor loaded scales. Zero points are not used as only symmetric quantization is supported. """ layer._k_scale = layer.k_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py index 23d25261301c..d8ee32b0fa41 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py @@ -37,7 +37,8 @@ def _dequant_gather_kernel( BLOCK: tl.constexpr, ): """Gather embedding rows by token id, unpack int32-packed INT weights, and - dequantize to ``out`` dtype in one pass (no int8 intermediate).""" + dequantize to ``out`` dtype in one pass (no int8 intermediate). + """ row = tl.program_id(0) col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) col_mask = col < hidden diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index cf98e52bb61a..54f5aa86b34c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -184,9 +184,7 @@ def create_weights( set_weight_attrs(w2_input_scale, extra_weight_attrs) def process_weights_after_loading(self, layer: RoutedExperts) -> None: - """ - Convert NVFP4 MoE weights into kernel format and setup the kernel. - """ + """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod # requires this naming convention. However, the name change breaks # reloading because the state dict no longer matches disk. Once we diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py index 2d2190e614ed..0abae681f083 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py @@ -36,8 +36,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): - """ - CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform + """CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles) - Scales: Fp32 for Channelwise , bf16 for groupwise quantization - Bias: Same data type as original weights diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 7a5b2d949b9f..bcbb687e3bb9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -143,8 +143,7 @@ def get_weight_shape( num_groups_w2: int | None = None, num_groups_w13: int | None = None, ) -> tuple[int, int, int]: - """ - Get the shape of the weight based on the weight name, number of experts + """Get the shape of the weight based on the weight name, number of experts hidden size, intermediate size per partition, number of groups for w2, and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 for weight scales/zero_points. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py index 78419a0dd98d..ff5173b180d2 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py @@ -9,33 +9,26 @@ class CompressedTensorsScheme(ABC): - """ - Abstract class used to describe the weight creation and forward pass + """Abstract class used to describe the weight creation and forward pass of different quantization schemes supported by CompressedTensors. """ @classmethod @abstractmethod def get_min_capability(cls) -> int: - """ - Get minimum device capability. - """ + """Get minimum device capability.""" raise NotImplementedError() @abstractmethod def create_weights(self, *args, **kwargs): - """ - Weight creation for the particular scheme. Inputs to this function - - """ + """Weight creation for the particular scheme. Inputs to this function.""" raise NotImplementedError() @abstractmethod def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None ): - """ - Run the forward pass for the particular scheme. This is where + """Run the forward pass for the particular scheme. This is where scheme-specific dequant/quant steps/kernels should be applied. Args: @@ -43,13 +36,13 @@ def apply_weights( other parameters relevant to the particular scheme. x: input to the layer bias: bias parameter + """ raise NotImplementedError() @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module): - """ - Called after weight loading is complete for any cleanup that + """Called after weight loading is complete for any cleanup that needs to occur. """ raise NotImplementedError() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index 2eb8ba717bde..f3b02b00378a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py @@ -21,8 +21,7 @@ class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme): - """ - Compressed tensors scheme for MXFP4. + """Compressed tensors scheme for MXFP4. Supports models quantized with the compressed-tensors mxfp4-pack-quantized format. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py index 5c511fc98d99..0f72be13a554 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py @@ -22,8 +22,7 @@ class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme): - """ - Compressed tensors scheme for MXFP8 quantization (W8A8). + """Compressed tensors scheme for MXFP8 quantization (W8A8). Loads pre-quantized MXFP8 weights from compressed-tensors checkpoints. Activations are dynamically quantized to MXFP8 at runtime. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py index 187b6438abff..3aa1aac45d7d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py @@ -211,7 +211,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None: """Normalize an int-quantized (plain int8) weight to the canonical - ``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect.""" + ``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect. + """ weight = layer.weight out_features, in_features = weight.shape packed = pack_to_int32(weight.data.contiguous(), self.num_bits) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py index 0cde020627d2..ef4e9d288e7b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py @@ -29,8 +29,7 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase): - """ - Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds + """Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds input and output transforms to either side of the original apply method """ @@ -240,8 +239,7 @@ def get_schemes_args( def get_layer_partition_names( layer_name: str, packed_modules_mapping: dict[str, list[str]] ) -> list[str]: - """ - Get all partition names associated with this layer. + """Get all partition names associated with this layer. Names are returned in order of their partition indices. ```python @@ -251,7 +249,9 @@ def get_layer_partition_names( "gate_proj", "up_proj", ] - assert get_layer_partition_names("mlp.down_proj", mapping) == ["down_proj"]""" + assert get_layer_partition_names("mlp.down_proj", mapping) == ["down_proj"] + ``` + """ for fused_suffix, part_suffixes in packed_modules_mapping.items(): if layer_name.endswith(fused_suffix): return [ diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py index 755059349382..ae8f143a61f2 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py @@ -23,8 +23,7 @@ class HadamardTransform(torch.nn.Module): - """ - Class which handles weight loading, postprocessing, and application of + """Class which handles weight loading, postprocessing, and application of transforms. Meant to be used with `CompressedTensorsLinearTransformMethod` and attention transforms method (not implemented yet) """ diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py index b8b771e3c7a9..70e0c652541f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py @@ -76,8 +76,7 @@ def create_weights( @staticmethod def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str: - """ - Given a kernel, find the string that is needed to be passed into + """Given a kernel, find the string that is needed to be passed into `flashinfer_scaled_fp4_mm`, using vllm.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth """ diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py index 188ee8c4c97d..bc999c7258b5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py @@ -76,8 +76,7 @@ def find_matched_target( targets: Iterable[str], fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), ) -> str | None: - """ - Helper function to look up which "target" in the compressed-tensors + """Helper function to look up which "target" in the compressed-tensors config that a layer corresponds to. Recall that a compressed-tensors configs has a concept of @@ -98,8 +97,8 @@ def find_matched_target( module: torch.nn.Module targets: list of targets to match the layer against fused_mapping: map from fused layer names to its components - """ + """ if layer_name is None: layer_name = "" @@ -115,8 +114,7 @@ def find_matched_target( def _find_first_match( value: str, targets: Iterable[str], check_contains: bool = False ) -> str | None: - """ - Returns first element of target that matches value either + """Returns first element of target that matches value either exactly or as a regex after 're:'. If check_contains is set to True, additionally checks if the target string is contained within the value. @@ -124,8 +122,8 @@ def _find_first_match( value: string to compare the list of targets against targets: list of targets to match the layer against check_contains: whether or not to do a substring match - """ + """ for target in targets: if is_equal_or_regex_match(value, target, check_contains=check_contains): return target @@ -137,8 +135,7 @@ def _match_fused_layer( target_layers: Iterable[str], fused_mapping: Mapping[str, list[str]], ) -> str | None: - """ - Match a fused layer name to its corresponding individual layer in + """Match a fused layer name to its corresponding individual layer in target_layers. Returns first value in fused_mapping which matches targets Implements an "all" matching strategy where a fused layer matches iff @@ -154,6 +151,7 @@ def _match_fused_layer( target_layers = ["model.layers.0.self_attn.q_proj", "model.layers.0.self_attn.k_proj", "model.layers.0.self_attn.v_proj"] + """ # find layer_name in mapping fused = next((key for key in fused_mapping if layer_name.endswith(key)), None) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index d8bef7347df0..5835af111af6 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -251,6 +251,7 @@ class Fp8LinearMethod(LinearMethodBase): Args: quant_config: The quantization config. + """ supports_pre_processed_weights = True @@ -490,6 +491,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): Args: quant_config: The quantization config. + """ supports_pre_processed_weights = True @@ -881,9 +883,7 @@ def apply( class Fp8KVCacheMethod(BaseKVCacheMethod): - """ - Supports loading kv-cache scaling factors from FP8 checkpoints. - """ + """Supports loading kv-cache scaling factors from FP8 checkpoints.""" def __init__(self, quant_config: Fp8Config): super().__init__(quant_config) diff --git a/vllm/model_executor/layers/quantization/fp_quant.py b/vllm/model_executor/layers/quantization/fp_quant.py index c24706252abd..8715db16b74a 100644 --- a/vllm/model_executor/layers/quantization/fp_quant.py +++ b/vllm/model_executor/layers/quantization/fp_quant.py @@ -98,6 +98,7 @@ class FPQuantLinearMethod(LinearMethodBase): Args: quant_config: The FPQuant quantization config. + """ def __init__(self, quant_config: FPQuantConfig): diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index ea00c0f3c408..128a5146f159 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -783,8 +783,7 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - """ - Apply Humming-quantized MoE computation using the standard kernel flow. + """Apply Humming-quantized MoE computation using the standard kernel flow. This method uses FusedMoEKernel.apply() which orchestrates: 1. Preparation (quantization if needed - skipped for Humming via diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index ce3e99bb8da3..e854c582eba3 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -86,7 +86,8 @@ def _resolve_raw( def is_explicitly_configured(name: str) -> bool: """Return True if *name* has an explicit entry in extra_config, - either via exact key match or via a regex pattern key.""" + either via exact key match or via a regex pattern key. + """ if not self._config.extra_config: return False if name in self._config.extra_config: diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py index 8672972aa18a..98c8c982f713 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py @@ -48,7 +48,8 @@ def get_moe_method( layer_config: "INCLayerConfig", ) -> "FusedMoEMethodBase | None": """Optional. Override if this scheme supports MoE. - Default raises NotImplementedError.""" + Default raises NotImplementedError. + """ raise NotImplementedError( f"{type(self).__name__} does not support MoE layers. " f"Layer config: {layer_config}" diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index 2d9cf624aba7..cf455527358e 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -28,8 +28,7 @@ # --8<-- [start:quant_fp8] @CustomOp.register("quant_fp8") class QuantFP8(CustomOp): - """ - Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group). + """Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group). This CustomOp supports both static and dynamic quantization. """ @@ -45,18 +44,18 @@ def __init__( use_ue8m0: bool | None = None, # for Torch compile compile_native: bool = True, ): - """ - Args: - static: static or dynamic quantization - group_shape: quantization group shape (PER_TOKEN, PER_TENSOR, - PER_CHANNEL, or arbitrary block size) - num_token_padding: Pad the token dimension of output to this - size - tma_aligned_scales: For group quantization, output scales in - TMA-aligned layout - column_major_scales: For group quantization, output scales in - column major format - compile_native: Manually compile forward_native if compile mode > None + """Args: + static: static or dynamic quantization + group_shape: quantization group shape (PER_TOKEN, PER_TENSOR, + PER_CHANNEL, or arbitrary block size) + num_token_padding: Pad the token dimension of output to this + size + tma_aligned_scales: For group quantization, output scales in + TMA-aligned layout + column_major_scales: For group quantization, output scales in + column major format + compile_native: Manually compile forward_native if compile mode > None + """ super().__init__(compile_native=compile_native) self.static = static diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py index 447cfb955a83..6d6a8c5ce51b 100644 --- a/vllm/model_executor/layers/quantization/kv_cache.py +++ b/vllm/model_executor/layers/quantization/kv_cache.py @@ -41,8 +41,7 @@ def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> Non class BaseKVCacheMethod(QuantizeMethodBase): - """ - Quant method that adds `_k_scale` and `_v_scale` attributes to the + """Quant method that adds `_k_scale` and `_v_scale` attributes to the Attention layer to support loading those scaling factors from checkpoints. The k/v_scale will be used to: - quantize k/v_cache entries before saving them to the cache @@ -50,6 +49,7 @@ class BaseKVCacheMethod(QuantizeMethodBase): Args: quant_config: the appropriate QuantizationConfig + """ supports_pre_processed_weights = True @@ -58,8 +58,7 @@ def __init__(self, quant_config: QuantizationConfig): self.quant_config = quant_config def create_weights(self, layer: torch.nn.Module): - """ - Create "weight" (aka q_scale, k_scale and v_scale) + """Create "weight" (aka q_scale, k_scale and v_scale) for an attention layer. """ # Initialize the Q and KV cache scales to -1.0, an invalid value. diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e3e8b373696f..f4ea599456f2 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -153,9 +153,7 @@ def algos_owned_by(config_name: str) -> tuple[str, ...]: class ModelOptKVCacheMethod(BaseKVCacheMethod): - """ - Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints. - """ + """Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints.""" def __init__(self, quant_config: "ModelOptQuantConfigBase"): super().__init__(quant_config) @@ -178,8 +176,7 @@ def __init__( self.exclude_modules: list[str] = exclude_modules def is_layer_excluded(self, prefix: str) -> bool: - """ - Check if a layer should be excluded from quantization. + """Check if a layer should be excluded from quantization. Handles both exact matching (for fused layers) and ModelOpt wildcard matching. @@ -472,8 +469,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): """MoE method for ModelOpt FP8. Supports loading FP8 checkpoints with static weight scale and activation scale. + Args: quant_config: The ModelOpt quantization config. + """ def __init__( @@ -818,10 +817,11 @@ def _from_config( class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): - """ - MoE Method for FP4 Quantization. + """MoE Method for FP4 Quantization. + Args: quant_config: NVFP4 Quant Config + """ supports_pre_processed_weights = True @@ -851,9 +851,7 @@ def __init__( ) def uses_weight_scale_2_pattern(self) -> bool: - """ - FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales. - """ + """FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.""" return True def create_weights( @@ -973,9 +971,7 @@ def create_weights( layer.register_parameter("w2_input_scale", w2_input_scale) def process_weights_after_loading(self, layer: RoutedExperts) -> None: - """ - Convert NVFP4 MoE weights into kernel format and setup the kernel. - """ + """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" if is_weights_pre_processed(): if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM: raise RuntimeError( @@ -1992,7 +1988,8 @@ def process(self, layer, role) -> None: class KNvfp4Dynamic(QuantKeyScheme): """NVFP4 activation scheme (W4A4). Has a static global input scale on disk; - the per-group scale is computed at runtime inside the kernel.""" + the per-group scale is computed at runtime inside the kernel. + """ key = kNvfp4Dynamic @@ -2029,7 +2026,8 @@ def process(self, layer, role) -> None: class KFp8StaticTensor(QuantKeyScheme): """Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and - the activation slot (W8A8). One key in both QuantSpec slots.""" + the activation slot (W8A8). One key in both QuantSpec slots. + """ key = kFp8StaticTensorSym @@ -2099,7 +2097,8 @@ def process(self, layer, role) -> None: class KFp8StaticChannel(QuantKeyScheme): """Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only — - there is no static per-channel *activation* today.""" + there is no static per-channel *activation* today. + """ key = kFp8StaticTokenSym @@ -2140,7 +2139,8 @@ def process(self, layer, role) -> None: class KFp8Block128(QuantKeyScheme): """128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D. - No transpose (block kernel keeps [out,in]).""" + No transpose (block kernel keeps [out,in]). + """ key = kFp8Static128BlockSym @@ -2196,7 +2196,8 @@ def process(self, layer, role) -> None: class KMxfp8Static(QuantKeyScheme): """MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale. - Weight-role only. process is validate-only plus an idempotency guard.""" + Weight-role only. process is validate-only plus an idempotency guard. + """ key = kMxfp8Static @@ -2271,7 +2272,8 @@ class KDynamicNoParam(QuantKeyScheme): """Dynamic activation with no stored scale (W8A8): quantized at runtime in the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a non-None activation key. Activation-role only. Serves the fp8 per-token, fp8 - per-block, and mxfp8 dynamic activation keys.""" + per-block, and mxfp8 dynamic activation keys. + """ def create_weights(self, layer, role, ctx, shapes, wl) -> None: if role is not ACT: @@ -2297,7 +2299,7 @@ def process(self, layer, role) -> None: def maybe_fuse_global_scales(layer) -> None: - """alpha = input_global_scale * weight_global_scale, presence-gated. + """Alpha = input_global_scale * weight_global_scale, presence-gated. W4A4 has both -> computed; W4A16 has no input_global_scale -> skipped. """ diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 430abdf70570..fef2c0c7eeae 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -229,6 +229,7 @@ class MoeWNA16Method(FusedMoEMethodBase): Args: quant_config: The MOE WNA16 (W8A16/W4A16) quantization config. + """ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 33f6afac80a9..a4f689f2c11e 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -79,7 +79,8 @@ def get_config_filenames(cls) -> list[str]: def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: """MoE method for RoutedExperts. Subclasses override to pick a - checkpoint-specific kernel family.""" + checkpoint-specific kernel family. + """ return Mxfp4MoEMethod(moe) def get_quant_method( diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index a16a33eaaa33..125ec4f2e32f 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -139,7 +139,8 @@ def _find_matching_targets( class OnlineQuantizationConfig(QuantizationConfig): """Model-level config for online quantization (quantize fp16/bf16 weights - during model loading, without requiring a pre-quantized checkpoint).""" + during model loading, without requiring a pre-quantized checkpoint). + """ def __init__( self, @@ -227,6 +228,7 @@ def _get_method_cls( Returns: The matching method class, or None when ``spec`` has no weight quantization. + """ if spec is None or spec.weight is None: return None @@ -260,6 +262,7 @@ def resolve_quant_method_cls( A tuple of source, quantization key string, target pattern, spec, and method class. Returns None when online quantization does not apply to the layer. + """ quant_spec: QuantSpec | None if self.args.targets is not None: @@ -313,6 +316,7 @@ def _resolve_targets_quant_method_metadata( A tuple of source, quantization key string, target pattern, spec, and dispatch table. Returns None when no pattern applies or the layer is ignored. + """ assert self.args.targets is not None ignored = should_ignore_layer( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 0234702c8bab..c1524e78b6eb 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -114,7 +114,8 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: class OnlineLinearBase(LinearMethodBase): """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint - weights onto meta device and materializes them just-in-time.""" + weights onto meta device and materializes them just-in-time. + """ uses_meta_device: bool = True @@ -158,7 +159,8 @@ def create_weights( class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. - Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + Loads fp16/bf16 weights and quantizes them per-tensor during loading. + """ def __init__(self): super().__init__() @@ -259,7 +261,8 @@ def apply( class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. - Loads fp16/bf16 weights and quantizes them per-block during loading.""" + Loads fp16/bf16 weights and quantizes them per-block during loading. + """ def __init__(self): super().__init__() @@ -430,7 +433,8 @@ def apply( class _Fp8OnlineMoEBase(OnlineMoEMethodBase): """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint - weights onto meta device and materializes them just-in-time.""" + weights onto meta device and materializes them just-in-time. + """ # Declared here for mypy; actual values are set in __init__. fp8_backend: "Fp8MoeBackend" @@ -551,7 +555,8 @@ def get_fused_moe_quant_config( class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): """Online tensorwise FP8 MoE quantization. - Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + Loads fp16/bf16 weights and quantizes them per-tensor during loading. + """ def __init__( self, @@ -608,7 +613,8 @@ def process_weights_after_loading(self, layer: Module) -> None: class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): """Online blockwise FP8 MoE quantization. - Loads fp16/bf16 weights and quantizes them per-block during loading.""" + Loads fp16/bf16 weights and quantizes them per-block during loading. + """ def __init__( self, diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 855a413c45a1..88544fbe994d 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -126,7 +126,6 @@ def maybe_update_config( revision: str | None = None, ): """Enable dynamic MXFP4 only for DeepSeek-V3-family fp4 checkpoints.""" - if hf_config is None: return @@ -162,13 +161,13 @@ def get_name(self) -> QuantizationMethods: def apply_vllm_mapper( # noqa: B027 self, hf_to_vllm_mapper: "WeightsMapper" ): - """ - Interface for models to update module names referenced in + """Interface for models to update module names referenced in quantization configs in order to reflect the vllm model structure Args: hf_to_vllm_mapper: maps from hf model structure (the assumed structure of the qconfig) to vllm model structure + """ quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {} @@ -657,8 +656,7 @@ def _is_w_ocp_mx_a_x( input_quant: dict[str, Any] | None, allow_static_fp8: bool = False, ) -> QuantKeyMatch: - """ - This check returns True only if it is an OCP-MX weight quantization. + """This check returns True only if it is an OCP-MX weight quantization. The activation can be FP16/BF16, OCP MXFP4, MXFP8, FP8. The rationale for checking only the weight type is that the model loading concept and process primarily concerns the weights themselves. @@ -984,8 +982,7 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): - """ - Use the CompressedTensorsScheme associated with each layer to create + """Use the CompressedTensorsScheme associated with each layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ @@ -1006,8 +1003,7 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ): - """ - Use the output of create_weights and the CompressedTensorsScheme + """Use the output of create_weights and the CompressedTensorsScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details @@ -1020,9 +1016,7 @@ def apply( class QuarkKVCacheMethod(BaseKVCacheMethod): - """ - Supports loading kv-cache scaling factors from quark checkpoints. - """ + """Supports loading kv-cache scaling factors from quark checkpoints.""" def __init__(self, quant_config: QuarkConfig): self.validate_kv_cache_config(quant_config.kv_cache_config) @@ -1030,12 +1024,12 @@ def __init__(self, quant_config: QuarkConfig): @staticmethod def validate_kv_cache_config(kv_cache_config: dict[str, Any] | None): - """ - Validator for the kv cache configuration. Useful for controlling the + """Validator for the kv cache configuration. Useful for controlling the kv cache quantization schemes, that are being supported in vLLM Args: kv_cache_config: the quark kv cache scheme + """ if kv_cache_config is None: return diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 735499d18c8b..d2472524a5fe 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1748,10 +1748,7 @@ def create_weights( set_weight_attrs(w2_input_scale_2, extra_weight_attrs) def process_weights_after_loading(self, layer: RoutedExperts) -> None: - """ - Convert NVFP4 MoE weights into kernel format and setup the kernel. - """ - + """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" # Match existing NVFP4 MoE paths: fused w13 uses the w1 global scale. if self.moe.is_act_and_mul and not torch.allclose( layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1] diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py index 7fb0332b574b..777423d78032 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py @@ -31,8 +31,7 @@ class QuarkNVFP4(QuarkScheme): - """ - Quark NVFP4 quantization scheme. + """Quark NVFP4 quantization scheme. Supports loading NVFP4 checkpoints with the following structure: - weight: uint8, shape [out_features, in_features // 2] (packed FP4) diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py index 419055b37217..6e75337e5ac3 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py @@ -11,8 +11,7 @@ class QuarkScheme(ABC): - """ - Abstract class used to describe the weight creation and forward pass + """Abstract class used to describe the weight creation and forward pass of different quantization schemes supported by Quark. """ @@ -36,25 +35,19 @@ def __init__( @classmethod @abstractmethod def get_min_capability(cls) -> int: - """ - Get minimum device capability. - """ + """Get minimum device capability.""" raise NotImplementedError @abstractmethod def create_weights(self, *args, **kwargs): - """ - Weight creation for the particular scheme. Inputs to this function - - """ + """Weight creation for the particular scheme. Inputs to this function.""" raise NotImplementedError @abstractmethod def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None ): - """ - Run the forward pass for the particular scheme. This is where + """Run the forward pass for the particular scheme. This is where scheme-specific dequant/quant steps/kernels should be applied. Args: @@ -62,13 +55,13 @@ def apply_weights( other parameters relevant to the particular scheme. x: input to the layer bias: bias parameter + """ raise NotImplementedError @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module): - """ - Called after weight loading is complete for any cleanup that + """Called after weight loading is complete for any cleanup that needs to occur. """ raise NotImplementedError diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 0f0b2065a8bb..1118db2db137 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -33,8 +33,7 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme): - """ - - Weights: MXFP4 with E8M0 scales per block of 32 + """- Weights: MXFP4 with E8M0 scales per block of 32 - Activations: FP8 E4M3 (static per-tensor quantization) Uses the AITER Triton kernel and falls back to emulation if AITER not available. diff --git a/vllm/model_executor/layers/quantization/qutlass_utils.py b/vllm/model_executor/layers/quantization/qutlass_utils.py index 86b0548307af..a997628a49ef 100644 --- a/vllm/model_executor/layers/quantization/qutlass_utils.py +++ b/vllm/model_executor/layers/quantization/qutlass_utils.py @@ -31,8 +31,7 @@ def triton_scale_swizzle( BLOCK_ROWS: tl.constexpr, BLOCK_COLS: tl.constexpr, ): - """ - Rearranges tensor data from row-major to block-scaled swizzle format. + """Rearranges tensor data from row-major to block-scaled swizzle format. Args: scale_ptr: Pointer to the input scale tensor @@ -43,6 +42,7 @@ def triton_scale_swizzle( output_block_stride: Stride between blocks in the output tensor BLOCK_ROWS: Number of rows in a tile (compile-time constant) BLOCK_COLS: Number of columns in a tile (compile-time constant) + """ pid_row = tl.program_id(0) pid_col = tl.program_id(1) @@ -86,8 +86,7 @@ def triton_scale_swizzle( @torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=()) def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: - """ - Rearranges an E8M0 tensor scale from row-major format to + """Rearranges an E8M0 tensor scale from row-major format to block-scaled swizzle format. This format is suitable for Tmem as described in NVIDIA documentation: @@ -98,6 +97,7 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: Returns: Rearranged tensor in block-scaled swizzle format + """ assert scale_tensor.element_size() == 1, ( "Expected element size to be 1 byte (8 bits)" @@ -154,8 +154,7 @@ def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor: def to_blocked( input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton" ) -> torch.Tensor: - """ - Rearrange a large matrix by breaking it into blocks and applying + """Rearrange a large matrix by breaking it into blocks and applying the rearrangement pattern. See: @@ -167,6 +166,7 @@ def to_blocked( Returns: Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4)) + """ if backend == "triton": return triton_mx_block_rearrange(input_matrix).flatten() diff --git a/vllm/model_executor/layers/quantization/torchao.py b/vllm/model_executor/layers/quantization/torchao.py index 8862578d6f7e..0a24c1e4a9fe 100644 --- a/vllm/model_executor/layers/quantization/torchao.py +++ b/vllm/model_executor/layers/quantization/torchao.py @@ -74,8 +74,7 @@ def torchao_version_at_least(torchao_version: str) -> bool: def should_skip(prefix: str, skip_modules: list[str]) -> bool: - """ - Robust skipping logic: + """Robust skipping logic: should_skip("model.model.layers.1.q_proj", ["model.model.layers.1.q_proj"]) # True should_skip("model.model.layers.10.o_proj", ["o_proj"]) -> True @@ -163,14 +162,14 @@ def get_min_capability(cls) -> int: @staticmethod def get_config_filenames() -> list[str]: - """torchao doesn't require additional config files, we use + """Torchao doesn't require additional config files, we use `config.json` from huggingface: `model_config.hf_config` """ return [] @classmethod def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig": - """Create the quant config from an hf model config""" + """Create the quant config from an hf model config.""" try: from torchao.core.config import config_from_dict except ImportError as err: @@ -287,12 +286,13 @@ def get_quant_method( def torchao_quantize_param_data( param: torch.Tensor, torchao_config: Any ) -> torch.nn.Parameter: - """Quantize a Tensor with torchao quantization specified by torchao_config + """Quantize a Tensor with torchao quantization specified by torchao_config. Args: param: weight parameter of the linear module torchao_config: type of quantization and their arguments we want to use to quantize the Tensor + """ from torchao.core.config import AOBaseConfig from torchao.quantization import quantize_ @@ -322,6 +322,7 @@ class TorchAOLinearMethod(LinearMethodBase): Args: quant_config: The torchao quantization config, a string that encodes the type of quantization and all relevant arguments. + """ def __init__(self, quant_config: TorchAOConfig): diff --git a/vllm/model_executor/layers/quantization/turboquant/centroids.py b/vllm/model_executor/layers/quantization/turboquant/centroids.py index 490265747c5b..65b1c821ccbc 100644 --- a/vllm/model_executor/layers/quantization/turboquant/centroids.py +++ b/vllm/model_executor/layers/quantization/turboquant/centroids.py @@ -45,6 +45,7 @@ def solve_lloyd_max( Returns: centroids: Sorted tensor of 2^bits optimal centroids. boundaries: Sorted tensor of 2^bits - 1 decision boundaries. + """ n_levels = 2**bits sigma2 = 1.0 / d diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py index 84e3940d264d..5b25f32b60e2 100644 --- a/vllm/model_executor/layers/quantization/turboquant/config.py +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -83,6 +83,7 @@ class TurboQuantConfig: norm_correction: Re-normalize centroid vectors to unit norm before inverse rotation during dequant. Fixes quantization-induced norm distortion, improving PPL by ~0.8% at 4-bit. + """ head_dim: int = 128 diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py index cfee80da9627..1beb988ac2ac 100644 --- a/vllm/model_executor/layers/quantization/utils/config_utils.py +++ b/vllm/model_executor/layers/quantization/utils/config_utils.py @@ -63,6 +63,7 @@ def is_shared_expert_quant_fse_compatible( Returns: A compatibility flag and, when incompatible, the reason. + """ if projection_names is None: projection_names = ["gate_up_proj", "down_proj"] @@ -303,13 +304,11 @@ def is_equal_or_regex_match( check_contains: bool = False, use_fnmatch: bool = False, ) -> bool: - """ - Checks whether a value is exactly equal or a regex match for target + """Checks whether a value is exactly equal or a regex match for target if target starts with 're:'. If check_contains is set to True, additionally checks if the target string is contained within the value. If use_fnmatch is set, supports shell-style patterns in target. """ - if target.startswith("re:"): pattern = target[3:] if re.match(pattern, value): diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index efe75fc6b611..c5c01e59ca3a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Utility helpers for NVFP4 + FlashInfer fused-MoE path""" +"""Utility helpers for NVFP4 + FlashInfer fused-MoE path.""" from typing import TYPE_CHECKING diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py index 1b5320615f4a..8df24ae9e6ce 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Utility helpers for MxInt4 + FlashInfer fused-MoE path""" +"""Utility helpers for MxInt4 + FlashInfer fused-MoE path.""" import functools @@ -37,8 +37,7 @@ def prepare_static_weights_for_trtllm_mxint4_moe( gemm2_weights: torch.Tensor, gemm2_scales: torch.Tensor, ) -> dict[str, torch.Tensor]: - """ - Prepare MxInt4 weights for TRT-LLM kernel. + """Prepare MxInt4 weights for TRT-LLM kernel. Input: gemm1_weights: [num_experts, 2*intermediate_size, hidden_size//8] int32 @@ -51,6 +50,7 @@ def prepare_static_weights_for_trtllm_mxint4_moe( Returns: Dict with keys 'gemm1_weights', 'gemm1_scales', 'gemm2_weights', 'gemm2_scales' containing shuffled/packed tensors ready for kernel + """ from flashinfer import block_scale_interleave from flashinfer.fused_moe import ( @@ -192,8 +192,7 @@ def flashinfer_trtllm_mxint4_moe( routing_method_type: int | None = None, routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel. + """Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel. Args: x: Input hidden states. dtype: bfloat16 @@ -209,11 +208,14 @@ def flashinfer_trtllm_mxint4_moe( ep_rank: Expert parallelism rank (default: 0) num_expert_group: Number of expert groups (default: None -> 0) topk_group: Top-k within groups (default: None -> 0) + routing_replay_out: Optional buffer receiving the routing decisions + so they can be replayed later e_score_correction_bias: Optional routing bias. dtype: bfloat16 routing_method_type: FlashInfer RoutingMethodType enum value Returns: Output tensor from MoE layer. dtype: same as x (bfloat16) + """ from flashinfer import RoutingMethodType from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index b31309679479..044285ac8513 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -59,7 +59,7 @@ def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor: def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor, is_gated_activation: bool ): - """Shuffle weights for FI TRT-LLM Format""" + """Shuffle weights for FI TRT-LLM Format.""" from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a epilogue_tile_m = 128 @@ -210,7 +210,6 @@ def align_fp4_moe_weights_for_fi( not satisfied (e.g. with certain tensor-parallel sizes), we pad the gate/up and down projection weights along the intermediate dim. """ - # Current local intermediate size (per partition) is the K dimension of # the down projection. num_experts, hidden_size, intermediate = w2.shape @@ -316,7 +315,6 @@ def align_moe_weights_for_fi( not satisfied (e.g. with certain tensor-parallel sizes), we pad the gate/up and down projection weights along the intermediate dim. """ - # Current local intermediate size (per partition) is the K dimension of # the down projection. num_experts, hidden_size, intermediate = w2.shape @@ -453,8 +451,7 @@ def prepare_fp8_moe_layer_for_fi( w2_input_scale: torch.Tensor | None, is_trtllm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Convert Fp8 MoE weights to flashinfer kernel format + """Convert Fp8 MoE weights to flashinfer kernel format. Note that for trtllm we update the model state dict with the scale format needed for these kernels. @@ -462,7 +459,6 @@ def prepare_fp8_moe_layer_for_fi( Note that for per-tensor, we update the layer's intermediate size if the weights needed padding. """ - assert hasattr(layer.moe_config, "is_act_and_mul") block_quant = ( hasattr(layer, "weight_block_size") and layer.weight_block_size is not None diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index e42f43b39c6c..edf92fb264dd 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -69,7 +69,8 @@ def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: """This function quantizes input values to float8 values " - "with tensor-wise quantization.""" + "with tensor-wise quantization. + """ dtype = current_platform.fp8_dtype() if dtype is None else dtype finfo = torch.finfo(dtype) min_val, max_val = x.aminmax() @@ -347,12 +348,10 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( BLOCK_N: tl.constexpr, ): # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks. - """ - Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then + """Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills the outputs tensors at the right positions. """ - pid_m = tl.program_id(0) pid_n = tl.program_id(1) N_2 = N // 2 @@ -421,8 +420,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor( alpha: float = 1.0, beta: float = 0.0, ): - """ - Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate + """Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate (silu: alpha=1, beta=0; swigluoai: alpha, beta from config). """ GROUP_SIZE = group_size @@ -564,18 +562,23 @@ def per_token_group_quant_fp8( """Function to perform per-token-group quantization on an input tensor `x`. It converts the tensor values into signed float8 values and returns the quantized tensor along with the scaling factor used for quantization. + Args: x: The input tensor with ndim >= 2. group_size: The group size used for quantization. eps: The minimum to avoid dividing zero. dtype: The dtype of output tensor. Note that only `torch.float8_e4m3fn` - is supported for now. + is supported for now. column_major_scales: Outputs scales in column major. tma_aligned_scales: Outputs scales in TMA-aligned layout. out_q: Optional output tensor. If not provided, function will create. + use_ue8m0: If True, round each scale down to a power of two (UE8M0 + format). None defers to the platform default. + Returns: tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor. + """ if use_ue8m0 is None: use_ue8m0 = is_deep_gemm_e8m0_used() @@ -695,6 +698,7 @@ def per_token_group_quant_fp8_packed_for_deepgemm( x_s_packed: Int32 tensor with logical shape [mn, ceil(num_groups_per_row / 4)], laid out with TMA-aligned stride along the packed-K dimension + """ if use_ue8m0 is None: use_ue8m0 = is_deep_gemm_e8m0_used() @@ -789,7 +793,6 @@ def _w8a8_triton_block_scaled_mm( product) on input tensors `A` and `B` with block-wise quantization, and store the result in output tensor `C`. """ - pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) @@ -842,14 +845,12 @@ def _w8a8_triton_block_scaled_mm( def get_w8a8_block_fp8_configs( N: int, K: int, block_n: int, block_k: int ) -> dict[int, Any] | None: - """ - Return optimized configurations for the w8a8 block fp8 kernel. + """Return optimized configurations for the w8a8 block fp8 kernel. The return value will be a dictionary that maps an irregular grid of batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the kernel on a given batch size bs, the closest batch size in the grid should be picked and the associated configuration chosen to invoke the kernel. """ - # First look up if an optimized configuration is available in the configs # directory device_name = get_device_name_as_file_name() @@ -889,18 +890,20 @@ def w8a8_triton_block_scaled_mm( quantization. It takes two input tensors `A` and `B` with scales `As` and `Bs`. The output is returned in the specified `output_dtype`. + Args: A: The input tensor, e.g., activation. B: The input tensor, e.g., weight. As: The per-token-group quantization scale for `A`. 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]. - output_dytpe: The dtype of the returned tensor. + be 2-dim, e.g., [128, 128]. + output_dtype: The dtype of the returned tensor. + Returns: torch.Tensor: The result of matmul. - """ + """ _on_gfx1250 = False if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx1250 @@ -1022,6 +1025,7 @@ def requant_weight_ue8m0_inplace( with shape `(..., M // block_size[0], K // block_size[1])`. block_size: 2-element iterable `[block_m, block_k]` describing the block quantisation granularity. + """ if weight.numel() == 0: return @@ -1191,7 +1195,8 @@ def prepare_fp8_moe_layer_for_deepgemm( def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor: """Pad the weight tensor. This is an optimization on ROCm platform, which - can benefit from tensors located far enough from one another in memory""" + can benefit from tensors located far enough from one another in memory + """ if ( envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm() @@ -1486,7 +1491,6 @@ def process_fp8_input_tensor_strategy_moe( enable_eplb: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Process moe input scales for tensor-wise quantization strategy.""" - if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale): logger.info_once( "Found input_scales that are not equal for " diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 2a3c1bd4f219..f03b1dc7b44f 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -88,8 +88,7 @@ class HummingMoEQuantConfig(FusedMoEQuantConfig): def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: - """ - Map humming group sizes to QuantKey GroupShape. + """Map humming group sizes to QuantKey GroupShape. group_size: elements per group along K (col); 0 means full dimension. group_size_n: elements per group along N (row); 0 means 1 (per-row). @@ -298,7 +297,8 @@ def _humming_input_schema_to_quant_key( schema: "HummingInputSchema", ) -> QuantKey | None: """Convert a HummingInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs.""" + the schema represents unquantized (bf16/fp16) inputs. + """ if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: return None @@ -427,7 +427,8 @@ def convert_linear_layer_to_humming_standard( ): """Rename/reshape a linear layer's quantized params (the canonical MPLinear layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names - and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" + and layout humming's weight schema expects (``weight`` / ``weight_scale``). + """ for name, checkpoint_name in name_map.items(): tensor = getattr(layer, checkpoint_name) delattr(layer, checkpoint_name) @@ -691,11 +692,9 @@ def select_humming_moe_experts( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> type[mk.FusedMoEExperts] | None: - """ - Select the primary Humming MoE Experts class + """Select the primary Humming MoE Experts class Note: Shape-specific fallbacks may still occur at runtime. """ - if not has_humming(): return None @@ -809,14 +808,14 @@ def _replace_layer_parameters( tensors: dict[str, torch.Tensor], preserve_bias: bool = False, ) -> None: - """ - Replace layer parameters for a sublayer with new tensors. + """Replace layer parameters for a sublayer with new tensors. Args: layer: The RoutedExperts layer sublayer_name: Name of the sublayer (e.g., "w13", "w2") tensors: Dict of parameter name to tensor preserve_bias: If True, don't delete bias parameters + """ # Delete old parameters for name, _ in list(layer.named_parameters()): @@ -843,11 +842,11 @@ def _convert_sublayer_to_humming( num_experts: int, param_dtype: torch.dtype, ) -> tuple[Any, Any]: - """ - Convert a sublayer's weights from checkpoint format to Humming format. + """Convert a sublayer's weights from checkpoint format to Humming format. Returns: Tuple of (converted_weight_schema, converted_input_schema) + """ from vllm.utils.humming import HummingWeightSchema @@ -931,8 +930,7 @@ def _process_single_sublayer( param_dtype: torch.dtype, force_weight_schema: Any | None = None, ) -> tuple[Any, Any, "LayerConfig"]: - """ - Process a single sublayer: convert, optionally requant, prepare, and transform. + """Process a single sublayer: convert, optionally requant, prepare, and transform. This combines the common logic from convert_to_humming_moe_kernel_format for processing a single sublayer. @@ -951,6 +949,7 @@ def _process_single_sublayer( Returns: Tuple of the final weight schema, input schema, and Humming layer config. + """ from vllm.utils.humming import HummingWeightSchema @@ -1005,8 +1004,7 @@ def convert_to_humming_moe_kernel_format( input_schema: Any | None = None, force_weight_schema: Any | None = None, ) -> dict[str, "LayerConfig"]: - """ - Convert MoE weights from checkpoint format to Humming kernel format. + """Convert MoE weights from checkpoint format to Humming kernel format. This function processes weights for each sublayer (w13, w2) by: 1. Converting from checkpoint format to humming format if needed @@ -1032,8 +1030,8 @@ def convert_to_humming_moe_kernel_format( - Modifies layer parameters in place - Sets layer.weight_schemas and layer.input_schemas - Sets layer.humming_configs for quant config construction - """ + """ # Build schemas from quant_config if not provided has_bias = layer.moe_config.has_bias num_experts = layer.moe_config.num_local_experts diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index db95d3588a2a..1180e2df3619 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -178,6 +178,7 @@ def per_token_group_quant_int8( Returns: tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ assert x.shape[-1] % group_size == 0, ( "the last dimension of `x` cannot be divisible by `group_size`" diff --git a/vllm/model_executor/layers/quantization/utils/machete_utils.py b/vllm/model_executor/layers/quantization/utils/machete_utils.py index 7bf71ca99c50..b0c3573404c9 100644 --- a/vllm/model_executor/layers/quantization/utils/machete_utils.py +++ b/vllm/model_executor/layers/quantization/utils/machete_utils.py @@ -17,8 +17,7 @@ def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]: def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]: - """ - Queries the supported group sizes for Machete based on the activation type. + """Queries the supported group sizes for Machete based on the activation type. Args: act_type: The activation data type (torch.float16, torch.bfloat16). @@ -27,6 +26,7 @@ def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]: A list of supported group sizes. The group size must be divisible by `TileShapeK = 128 * 8 // num_bits(act_type)`. -1 indicates per-channel quantization. + """ if act_type in [torch.float16, torch.bfloat16]: return [-1, 64, 128] diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 22cf18f4a5a6..8e0e4bba1c2d 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -245,7 +245,8 @@ def marlin_padded_nk(size_n: int, size_k: int, group_size: int = -1) -> tuple[in def marlin_repacked_nk(qweight: torch.Tensor, num_bits: int) -> tuple[int, int]: """Recover the (size_n, size_k) a Marlin weight was repacked with - (including any tile padding) from its packed shape.""" + (including any tile padding) from its packed shape. + """ pack_factor = 32 // num_bits size_k = qweight.size(0) * GPTQ_MARLIN_TILE size_n = qweight.size(1) * pack_factor // GPTQ_MARLIN_TILE @@ -256,7 +257,8 @@ def marlin_pad_qweight( qweight: torch.Tensor, size_n: int, size_k: int, padded_n: int, padded_k: int ) -> torch.Tensor: """Zero-pad a GPTQ-layout packed weight (size_k / pack, size_n) for - gptq_marlin_repack.""" + gptq_marlin_repack. + """ if (padded_n, padded_k) == (size_n, size_k): return qweight pack_factor = size_k // qweight.size(0) @@ -274,7 +276,8 @@ def marlin_pad_scales( group_size: int, ) -> torch.Tensor: """Zero-pad weight scales (num_groups, size_n); call before - marlin_permute_scales and pass the padded extents to it.""" + marlin_permute_scales and pass the padded extents to it. + """ if (padded_n, padded_k) == (size_n, size_k): return scales pad_rows = padded_k // group_size - scales.size(0) if group_size > 0 else 0 @@ -397,8 +400,7 @@ def check_moe_marlin_supports_layer( def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tensor): - """ - Given Marlin packed weight matrices w1_packed, and w2_packed, + """Given Marlin packed weight matrices w1_packed, and w2_packed, return the MoE intermediate size N """ marlin_tile_size = 16 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 3b6990101ed7..1a550576b6fa 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -41,8 +41,8 @@ def _nvfp4_compute_scale_factor( ) -> float: """Compute the power-of-2 scale_factor needed so that all non-zero values in marlin_scales * 2^7 are >= 2 after rescaling. - Returns a Python float (power of 2, >= 1.0).""" - + Returns a Python float (power of 2, >= 1.0). + """ # Since half has a smaller dynamic range compared to bfloat16, # no rescaling is applied here if active dtype is half. if a_dtype is not None and a_dtype == torch.half: @@ -76,9 +76,12 @@ def nvfp4_marlin_process_scales( given value is used directly. The caller is responsible for dividing ``global_scale`` by the returned ``scale_factor`` to preserve numerical correctness. + a_dtype: dtype of the activations the scales will be used with. None + keeps the dtype of ``marlin_scales``. Returns: A tuple of (processed_scales, scale_factor). + """ if not (marlin_scales >= 0).all(): logger.warning_once( @@ -373,7 +376,8 @@ def prepare_nvfp4_moe_layer_for_marlin( def pad_w13(x: torch.Tensor) -> torch.Tensor: """Zero-pad each gate/up shard of a (E, num_shards * N, cols) - tensor to padded_N rows.""" + tensor to padded_N rows. + """ if padded_N == N: return x x = x.view(E, num_shards, N, x.size(-1)) @@ -382,7 +386,8 @@ def pad_w13(x: torch.Tensor) -> torch.Tensor: def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: """Zero-pad the packed N (last) dim of a (E, K, N / packing) - tensor.""" + tensor. + """ if padded_N == N: return x return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 56841cba17c3..251207a9733a 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -217,7 +217,8 @@ def prepare_fp8_layer_for_marlin( def _moe_pad_shard_rows(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor: """Zero-pad each gate/up shard of a ``(E, 2 * n, ...)`` tensor to padded_n - rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing.""" + rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing. + """ if padded_n == n: return x e = x.size(0) @@ -241,14 +242,12 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight_scale: torch.Tensor, w2_weight_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Shuffle weights and scales into marlin format. + """Shuffle weights and scales into marlin format. Note that this function has the side effect of adding a `workspace` attribute to the layer. This `workspace` does not need to be registered as a Parameter as it is not used during weight reloading. """ - logger.warning_once( "Your GPU does not have native support for FP8 computation but " "FP8 quantization is being used. Weight-only FP8 compression will " @@ -372,9 +371,7 @@ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor: def pack_fp8_to_int32( fp8_tensor: torch.Tensor, size_k_first: bool = True ) -> torch.Tensor: - """ - Repack FP8 weights to gptq format (packed int32 elements) - """ + """Repack FP8 weights to gptq format (packed int32 elements).""" assert fp8_tensor.dtype == torch.float8_e4m3fn assert fp8_tensor.ndim == 2 @@ -524,6 +521,7 @@ def prepare_mxfp8_moe_layer_for_marlin( Returns: (w13, w2, w13_scale, w2_scale) in Marlin format. + """ group_size = 32 e = w13.shape[0] diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py index d4a7c563b8f3..5f8cb9f709df 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Utility functions used for tests and benchmarks""" +"""Utility functions used for tests and benchmarks.""" import numpy as np import torch diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 2d0a4fa48c0d..1bd6b3f9314b 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -28,7 +28,7 @@ def should_use_cdna4_mx_scale_swizzle() -> bool: def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): - """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel""" + """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel.""" assert has_triton_kernels() import triton_kernels.matmul_ogs_details.opt_flags as opt_flags from triton_kernels.numerics import InFlexData @@ -374,8 +374,7 @@ def downcast_to_mxfp( BLOCK_OUT_DIM: int = 128, BLOCK_QUANT_DIM: int = 32, ): - """ - Convert the src weights to MXFP4. The src weight is quantized along the + """Convert the src weights to MXFP4. The src weight is quantized along the axis dimension into packed e2m1 values (torch.uint8, two values per byte), so the size of that dimension in the output is half of the logical (unpacked) size. diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index a5272dd9cb19..a66d481159fa 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -657,8 +657,7 @@ def ref_nvfp4_quant(x, global_scale, block_size): def ref_nvfp4_quant_dequant( x: torch.Tensor, global_scale: torch.Tensor, block_size: int ) -> torch.Tensor: - """ - NVFP4 quantize-dequantize operation. + """NVFP4 quantize-dequantize operation. `global_scale` is expected to have a single element. """ diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py index 74e35dce566d..c446f646a18a 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py @@ -11,18 +11,15 @@ def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor: - """ - Pad and block-interleave the FP4 block-scales so that they match the data + """Pad and block-interleave the FP4 block-scales so that they match the data layout expected by the CUTLASS / FlashInfer kernels. - Parameters - ---------- - scale: torch.Tensor + Args: + scale: The FP4 block-scale tensor to swizzle. + + Returns: + The swizzled tensor with the same logical shape as `scale`. - Returns - ------- - torch.Tensor - The swizzled tensor with the same logical shape as *scale*. """ assert scale.dtype == torch.float8_e4m3fn, ( "swizzle_blockscale expects the input tensor to be in " @@ -65,8 +62,7 @@ def pad_nvfp4_weight_for_cutlass( weight: torch.Tensor, alignment: int = 32, ) -> tuple[torch.Tensor, int]: - """ - Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy + """Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy the alignment constraints required by CUTLASS / FlashInfer FP4 kernels. CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible @@ -99,8 +95,7 @@ def pad_nvfp4_weight_for_cutlass( def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int: - """ - Compute K-dim padding (packed bytes) of an NVFP4 weight, + """Compute K-dim padding (packed bytes) of an NVFP4 weight, Aligned with `pad_nvfp4_weight_for_cutlass` above. """ return layer.weight.shape[1] - layer.input_size_per_partition // 2 @@ -110,8 +105,7 @@ def pad_nvfp4_activation_for_cutlass( x_fp4: torch.Tensor, weights_padding_bytes: int, ) -> torch.Tensor: - """ - Pad packed FP4 activations to match the K-dimension padding applied to weights. + """Pad packed FP4 activations to match the K-dimension padding applied to weights. The padding is in bytes (tensor dimension), not FP4 elements. """ if weights_padding_bytes > 0: @@ -123,9 +117,7 @@ def slice_nvfp4_output( out: torch.Tensor, output_size: int, ) -> torch.Tensor: - """ - Slice the output tensor to remove padding in N dimension if weight was padded. - """ + """Slice the output tensor to remove padding in N dimension if weight was padded.""" if out.shape[-1] != output_size: return out[..., :output_size].contiguous() return out diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index f96146e73615..f3f72adc789f 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""This file is used for /tests and /benchmarks""" +"""This file is used for /tests and /benchmarks.""" from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -109,8 +109,7 @@ class _GroupShape(NamedTuple): class GroupShape(_GroupShape): - """ - This class describes the quantization group shape. + """This class describes the quantization group shape. It includes static members for common shapes (per-tensor, per-token). """ @@ -139,8 +138,7 @@ def is_per_group(self) -> bool: @dataclass(frozen=True) class ScaleDesc: - """ - Class for describing a single quantization scaling factor. + """Class for describing a single quantization scaling factor. dtype: data type of the scale static: static scale if True, dynamic if False group_shape: group shape of the scale @@ -166,8 +164,7 @@ def __str__(self): @dataclass(frozen=True) class QuantKey: - """ - Class for identifying the type of quantization. + """Class for identifying the type of quantization. dtype: quantized data type scale: scale descriptor scale2: second-level scale descriptor @@ -378,8 +375,7 @@ def prep_scale_for_group_broadcast( x: torch.Tensor, group_shape: GroupShape | None, ) -> torch.Tensor: - """ - Prepare the input quantization scale for group broadcasting. + """Prepare the input quantization scale for group broadcasting. Args: scale: The scale tensor (scalar or 1D). @@ -388,6 +384,7 @@ def prep_scale_for_group_broadcast( Returns: scale reshaped for correct broadcasting. + """ if scale.numel() == 1: # For per-tensor quant, keep the scale as a scalar (not reshaped to (1, 1)). @@ -432,13 +429,13 @@ def scaled_quantize( quant_dtype: torch.dtype, compute_dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Args: - x: Input tensor to quantize - group_shape: Shape of quantization groups - quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn) - compute_dtype: Optional dtype for intermediate computations. - If None, uses input dtype. Use torch.float32 for higher precision. + """Args: + x: Input tensor to quantize + group_shape: Shape of quantization groups + quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn) + compute_dtype: Optional dtype for intermediate computations. + If None, uses input dtype. Use torch.float32 for higher precision. + """ group_shape = _normalize_quant_group_shape(x, group_shape) assert quant_dtype.is_floating_point, ( @@ -506,7 +503,7 @@ def get_attribute_fallback(obj, attributes: list[str]): def get_and_maybe_dequant_weights( layer: "LinearBase", out_dtype: torch.dtype = torch.float32 ): - """Return layer's unquantized weights in [out, in] layout""" + """Return layer's unquantized weights in [out, in] layout.""" from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod from vllm.model_executor.layers.quantization.online.fp8 import ( @@ -924,8 +921,7 @@ def awq_pack( def convert_bf16_scales_to_fp8( quant_fp8: Callable, scales: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales) + """Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales) expected by W4A8 GEMM kernels. """ assert scales.is_contiguous(), ( @@ -949,9 +945,7 @@ def convert_bf16_scales_to_fp8( def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tensor: - """ - Convert int4b8 (packed to int32) to signed int4 - """ + """Convert int4b8 (packed to int32) to signed int4.""" assert t.is_cuda, "tensor must be on gpu" assert t.dtype == torch.int32, f"expected int32 packed weights but got {t.dtype}" diff --git a/vllm/model_executor/layers/resampler.py b/vllm/model_executor/layers/resampler.py index c9fa8054625e..320508de76de 100644 --- a/vllm/model_executor/layers/resampler.py +++ b/vllm/model_executor/layers/resampler.py @@ -26,8 +26,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -""" -Shared resampler perceiver network used in multimodal models and +"""Shared resampler perceiver network used in multimodal models and related helpers for sincos positional embeddings. Example models: Qwen (Qwen-VL), MiniCPM-V 2.0 @@ -76,8 +75,7 @@ def get_abs_pos(abs_pos: torch.Tensor, tgt_size: torch.Tensor | int) -> torch.Te def get_1d_sincos_pos_embed_from_grid( embed_dim: int, pos: np.ndarray, version: tuple[int, int] = (2, 0) ) -> torch.Tensor: - """ - embed_dim: output dimension for each position + """embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) / (H, W) out: (M, D) / (H, W, D) """ @@ -126,8 +124,7 @@ def get_2d_sincos_pos_embed( cls_token: bool = False, version: tuple[int, int] = (2, 0), ) -> torch.Tensor: - """ - grid_size: int of the grid height and width + """grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) @@ -154,8 +151,7 @@ def get_2d_sincos_pos_embed( class BaseResampler(nn.Module): - """ - A 2D perceiver-resampler network with one cross attention layers by + """A 2D perceiver-resampler network with one cross attention layers by (grid_size**2) learnable queries and 2d sincos pos_emb. Outputs: A tensor with the shape of (grid_size**2, embed_dim) diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py index 006dba17ec48..24c96c252120 100644 --- a/vllm/model_executor/layers/rotary_embedding/common.py +++ b/vllm/model_executor/layers/rotary_embedding/common.py @@ -138,14 +138,14 @@ def forward_static( is_neox_style: bool = True, enable_fp32_compute: bool = False, ) -> torch.Tensor: - """ - Args: - x: [batch_size (optional), seq_len, num_heads, head_size] - cos: [seq_len, head_size // 2] - sin: [seq_len, head_size // 2] - is_neox_style: Whether to use the Neox-style or GPT-J-style. - enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype - for higher accuracy. + """Args: + x: [batch_size (optional), seq_len, num_heads, head_size] + cos: [seq_len, head_size // 2] + sin: [seq_len, head_size // 2] + is_neox_style: Whether to use the Neox-style or GPT-J-style. + enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype + for higher accuracy. + """ origin_dtype = x.dtype if enable_fp32_compute: diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py index f5e6472e235e..46e06371f290 100644 --- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py @@ -263,7 +263,6 @@ def forward_native( inverse: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: """PyTorch-native implementation equivalent to forward().""" - head_size = query.size(-1) query_rot = query[..., -self.rotary_dim :] key_rot = key[..., -self.rotary_dim :] if key is not None else None diff --git a/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py b/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py index ec0c5a3286c7..428901596bb1 100644 --- a/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py @@ -8,7 +8,7 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding): - """3D rotary positional embedding. 3D is t:time h:height w:width""" + """3D rotary positional embedding. 3D is t:time h:height w:width.""" def forward_native( # type: ignore[override] self, diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index 4c8a7bcbfa1d..96ba482184a8 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -185,7 +185,7 @@ def forward_native( return query, key def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): - """load fope weights""" + """Load fope weights.""" world_size = get_tensor_model_parallel_world_size() rank = get_tensor_model_parallel_rank() num_key_value_heads = loaded_weight.size(0) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index a7fe51c46598..185bfd650ca3 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -178,8 +178,12 @@ def triton_mrope( (T/H/W positions with multimodal inputs) mrope_section: [t, h, w] head_size: int + rotary_dim: Number of leading dimensions rotary is applied to. + mrope_interleaved: Whether the T/H/W sections are interleaved rather + than concatenated. is_neox_style: Whether rotary pairs use split-half (NeoX) or adjacent (GPT-J) layout. + """ n_row, n_q_head_head_dim = q.shape n_q_head = n_q_head_head_dim // head_size @@ -320,6 +324,8 @@ def forward_native( [3, num_tokens] (T/H/W positions with multimodal inputs) query: [num_tokens, num_heads * head_size] key: [num_tokens, num_kv_heads * head_size] + offsets: Optional per-token position offsets added to positions. + """ assert positions.ndim == 1 or positions.ndim == 2 assert key is not None diff --git a/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py b/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py index a6f6c276b442..f1cf14b96d1e 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py @@ -85,7 +85,7 @@ def _rebuild_pos_emb( self, positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Interleave the rotary embedding""" + """Interleave the rotary embedding.""" cos_sin = self.cos_sin_cache[positions] mrope_section_3d = [1] * len(self.mrope_dim) mrope_dim = self.mrope_dim @@ -139,8 +139,7 @@ def forward( def get_mrope_interleaved_id_list( a: int, b: int, c: int, force_last: bool = False ) -> list[int]: - """ - Generate an interleaved list of indices for multi-modal rotary embedding. + """Generate an interleaved list of indices for multi-modal rotary embedding. Args: a: Number of indices for first modality @@ -150,6 +149,7 @@ def get_mrope_interleaved_id_list( Returns: List of interleaved indices + """ if force_last: a -= 1 diff --git a/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py index 031a12fceba6..b92dc8e7fd6b 100644 --- a/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py @@ -9,7 +9,8 @@ class NTKScalingRotaryEmbedding(RotaryEmbedding): """RotaryEmbedding extended with fixed and mixed NTK scaling. - https://kexue.fm/archives/9706""" + https://kexue.fm/archives/9706 + """ def __init__( self, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 32902901c72b..e7a938239614 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -140,7 +140,6 @@ def dcp_gather_kv_rows( out_buf: torch.Tensor, ) -> torch.Tensor: """All-gather this rank's KV shard and return it in global token order.""" - gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]] dist.all_gather_into_tensor( gathered, local_padded.contiguous(), group=get_dcp_group().device_group @@ -286,7 +285,8 @@ def _gather_workspace_shapes( """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + - (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales. + """ if use_fp4_cache: return ( ((total_seq_lens, head_dim // 2), torch.uint8), @@ -304,7 +304,8 @@ def kv_cache_as_quant_view( use_fp4_cache: bool, ) -> torch.Tensor: """4D ``[num_blocks, block_size, 1, head_width]`` view expected by - DeepGEMM, from the 3D indexer kv-cache allocation.""" + DeepGEMM, from the 3D indexer kv-cache allocation. + """ if use_fp4_cache: assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 num_blocks, block_size, _ = kv_cache.shape diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py index 78db1b3df5a4..a019f9ec7373 100644 --- a/vllm/model_executor/layers/sparse_mqa_indexer.py +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -41,7 +41,8 @@ def _prefill_k_workspaces( total_seq_lens: int, head_dim: int ) -> tuple[torch.Tensor, torch.Tensor]: """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with - the dense indexer layers of the same model.""" + the dense indexer layers of the same model. + """ values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True ) @@ -119,7 +120,8 @@ def __init__( def _reserve_workspaces(self, device: torch.device) -> None: """Profiling run: claim the K-gather workspace and the peak sparse - logits allocation so the memory estimate covers them.""" + logits allocation so the memory estimate covers them. + """ _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index f483a619b293..a1b79a59a44b 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -48,8 +48,7 @@ def apply_penalties( frequency_penalties: torch.Tensor, repetition_penalties: torch.Tensor, ) -> torch.Tensor: - """ - Applies penalties in place to the logits tensor + """Applies penalties in place to the logits tensor logits : The input logits tensor of shape [num_seqs, vocab_size] prompt_tokens_tensor: A tensor containing the prompt tokens. The prompts are padded to the maximum prompt length within the batch using @@ -259,6 +258,7 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo Returns: The CHUNKK the kernel will dispatch with, and whether the CU budget and split-K workspace admit the shape at all. + """ # Next ^2 of n N_p2 = 1 << (n - 1).bit_length() diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 456a22a7a26a..688b4bf197ef 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -242,6 +242,7 @@ class VocabParallelEmbedding(PluggableLayer): disable_tp: If true, tensor parallelism will be disabled for this layer. quant_method: Preselected quantization method for model-specific layers. parallel_group: Process group used to shard and reduce the embedding. + """ # noqa: E501 # --8<-- [end:vocab_parallel_embedding] @@ -375,7 +376,8 @@ def _get_indices( ) -> VocabParallelEmbeddingShardIndices: """Get start and end indices for vocab parallel embedding, following the layout outlined in the class docstring, based on the given tp_rank and - tp_size.""" + tp_size. + """ num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded padded_org_vocab_start_index, padded_org_vocab_end_index = ( vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, tp_size) @@ -581,6 +583,7 @@ class ParallelLMHead(VocabParallelEmbedding): org_num_embeddings: original vocabulary size (without LoRA). padding_size: padding size for the vocabulary. disable_tp: If true, tensor parallelism will be disabled for this layer. + """ # --8<-- [end:parallel_lm_head] diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index 6305c45f5d0b..226caebe1cd4 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -94,6 +94,7 @@ def register_model_loader(load_format: str): >>> load_config = LoadConfig(load_format="my_loader") >>> type(get_model_loader(load_config)) + """ # noqa: E501 def _wrapper(model_loader_cls): diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 6f356464585a..bbda4a33fdad 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -36,7 +36,8 @@ def download_model(self, model_config: ModelConfig) -> None: @abstractmethod def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: """Load weights into a model. This standalone API allows - inplace weights loading for an already-initialized model""" + inplace weights loading for an already-initialized model + """ raise NotImplementedError def create_model( diff --git a/vllm/model_executor/model_loader/checkpoint_weight_patch.py b/vllm/model_executor/model_loader/checkpoint_weight_patch.py index 0425aabc7dd1..4a26c73ebc86 100644 --- a/vllm/model_executor/model_loader/checkpoint_weight_patch.py +++ b/vllm/model_executor/model_loader/checkpoint_weight_patch.py @@ -36,6 +36,7 @@ class CheckpointWeightPatch(NamedTuple): ``indices`` for a sparse patch. indices: Flat indices into the full checkpoint tensor described by ``shape``. ``None`` makes ``values`` a dense replacement. + """ name: str @@ -176,8 +177,8 @@ def load_checkpoint_weight_patches( Returns: Union of the weight names reported by all ``model.load_weights`` calls. - """ + """ if max_chunk_bytes <= 0: raise ValueError("max_chunk_bytes must be positive") diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 039a577f80d1..e48c5a380ea8 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -135,7 +135,8 @@ def _prepare_weights( ) -> tuple[str, list[str], bool]: """Prepare weights for the model. - If the model is not local, it will be downloaded.""" + If the model is not local, it will be downloaded. + """ model_name_or_path = ( maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path diff --git a/vllm/model_executor/model_loader/ep_weight_filter.py b/vllm/model_executor/model_loader/ep_weight_filter.py index e9152f2cdb2a..4b2f92a18dca 100644 --- a/vllm/model_executor/model_loader/ep_weight_filter.py +++ b/vllm/model_executor/model_loader/ep_weight_filter.py @@ -23,7 +23,8 @@ def parse_expert_id(weight_name: str) -> int | None: Returns ``None`` for dense weights (attention, layernorm, embedding), shared experts, and 3D fused-expert tensors where all experts are stored - in a single tensor without a numeric expert id in the name.""" + in a single tensor without a numeric expert id in the name. + """ m = _EXPERT_ID_RE.search(weight_name) return int(m.group(1)) if m else None @@ -43,8 +44,12 @@ def compute_local_expert_ids( :func:`vllm.model_executor.layers.fused_moe.layer.determine_expert_map`. Args: + num_experts: Total number of experts in the layer. + ep_size: Expert parallel world size. + ep_rank: Rank whose local expert ids are computed. placement: ``"linear"`` for contiguous assignment, ``"round_robin"`` for interleaved assignment. + """ if ep_size <= 1: return None @@ -66,7 +71,8 @@ def should_skip_weight( local_expert_ids: set[int] | None, ) -> bool: """Return ``True`` if *weight_name* is an expert weight that does not - belong to the local rank and should be skipped during loading.""" + belong to the local rank and should be skipped during loading. + """ if local_expert_ids is None: return False eid = parse_expert_id(weight_name) diff --git a/vllm/model_executor/model_loader/reload/__init__.py b/vllm/model_executor/model_loader/reload/__init__.py index 61dc1b66e6f3..a1ac282d290f 100644 --- a/vllm/model_executor/model_loader/reload/__init__.py +++ b/vllm/model_executor/model_loader/reload/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Layerwise weight reloading utilities for vLLM. +"""Layerwise weight reloading utilities for vLLM. This module provides functionality to reload model weights layer-by-layer, which is useful for weight updates without full model reconstruction. diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 5a73738e1511..c52fad985beb 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -54,8 +54,7 @@ def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo: - """ - Get information related to restoring and layerwise processing. If no previous + """Get information related to restoring and layerwise processing. If no previous information existed, a new entry is constructed """ if layer not in LAYERWISE_INFO: @@ -68,8 +67,7 @@ def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo: def record_metadata_for_reloading(model: torch.nn.Module): - """ - Record layer metadata needed for later reloading. + """Record layer metadata needed for later reloading. Stores parameter and buffer metadata as meta tensors for restoration. Must be called before `initialize_layerwise_reload`. @@ -82,8 +80,7 @@ def record_metadata_for_reloading(model: torch.nn.Module): @torch.no_grad() def initialize_layerwise_reload(model: torch.nn.Module): - """ - Set up layerwise weight loading with deferred processing. + """Set up layerwise weight loading with deferred processing. Must be called after `record_metadata_for_reloading`. This function: 1. Saves current kernel tensors for later copying @@ -120,13 +117,13 @@ def initialize_layerwise_reload(model: torch.nn.Module): def initialize_online_processing(layer: torch.nn.Module): - """ - Wrap a layer's weight loaders with online processing loaders. + """Wrap a layer's weight loaders with online processing loaders. Called by either `initialize_layerwise_reload` or an online quantization scheme, prevents double wrapping in the case of online quantization + reloading Args: layer: layer whose parameter weight loaders will be wrapped + """ info = get_layerwise_info(layer) @@ -226,8 +223,7 @@ def online_process_loader(*args, **kwargs): def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig): - """ - Apply processing to any layers which were not layerwise processed during loading. + """Apply processing to any layers which were not layerwise processed during loading. This includes attention layers and layers which have weight elements which are not loaded (due to padding). @@ -237,6 +233,7 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon Args: model: model to finalize processing for model_config: config needed for applying processing to attention layers + """ if hasattr(model, "_original_do_torchao_reload"): model._do_torchao_reload = model._original_do_torchao_reload @@ -311,7 +308,8 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) - Assumes dtype/shapes of attention tensors do not change during processing, since we use .data.copy_() to preserve kernel tensor - references.""" + references. + """ quant_method = getattr(layer, "quant_method", None) if quant_method is not None: # Re-create scale Parameters with sentinel values so unloaded scales @@ -330,8 +328,7 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) - def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): - """ - Finalize layer loading after all weights have been buffered. + """Finalize layer loading after all weights have been buffered. This function: 1. Materializes the layer onto the target device @@ -377,7 +374,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): def _get_original_loader(tensor: torch.Tensor) -> Callable: - """Return the weight loader with any layerwise wrappers removed""" + """Return the weight loader with any layerwise wrappers removed.""" loader = _get_weight_loader(tensor) while loader.__name__ == "online_process_loader": loader = loader.__wrapped__ # type: ignore[union-attr] @@ -391,7 +388,8 @@ def _get_weight_loader(tensor: torch.Tensor): def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): """Copy processed values into original kernel tensor storage and restore - kernel tensor references on the layer. Preserves cudagraph references.""" + kernel tensor references on the layer. Preserves cudagraph references. + """ assert info.kernel_tensors is not None parameters, buffers = info.kernel_tensors non_persistent = info.kernel_non_persistent_buffers diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index a8f8023cf26d..4d08d4916536 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -47,8 +47,7 @@ def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor: def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor: - """ - Materialize a meta tensor into an actual tensor on the current device. + """Materialize a meta tensor into an actual tensor on the current device. Should be called within the torch device context for the given rank. """ tensor = torch.empty_strided( @@ -119,7 +118,7 @@ def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors: def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo): - """Restore a layer to model format with tensors on the meta device""" + """Restore a layer to model format with tensors on the meta device.""" if layer.__class__.__name__ in SKIP_MODULES: return @@ -155,8 +154,7 @@ def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo): class CopyCounter(TorchDispatchMode): - """ - Tracks total number of elements modified with `copy_`. + """Tracks total number of elements modified with `copy_`. Useful for keeping track of weight loading where underlying weights can be arbitrarily transformed (such as with `narrow`) before calling copy. @@ -182,8 +180,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): def get_numel_loaded( weight_loader: Callable, args: inspect.BoundArguments ) -> tuple[int, object]: - """ - Determine how many elements would be loaded by a weight loader call. + """Determine how many elements would be loaded by a weight loader call. Args: weight_loader: used to load weights @@ -192,6 +189,7 @@ def get_numel_loaded( Returns: number of elements loaded by the weight loader, the return value of the weight loader + """ with CopyCounter() as counter: return_value = weight_loader(*args.args, **args.kwargs) diff --git a/vllm/model_executor/model_loader/reload/sanitize.py b/vllm/model_executor/model_loader/reload/sanitize.py index 21c47a2257f5..4e2c780e90f0 100644 --- a/vllm/model_executor/model_loader/reload/sanitize.py +++ b/vllm/model_executor/model_loader/reload/sanitize.py @@ -11,8 +11,7 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor: - """ - Removes references to layer held by tensor attributes. Specifically, removes the + """Removes references to layer held by tensor attributes. Specifically, removes the `__self__` attribute of weight loader methods attached to the tensor. Used by `capture_layer_to_meta` to avoid circular references to layers in @@ -26,6 +25,7 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.T Returns: sanitized tensor + """ for key, value in tensor.__dict__.items(): if isinstance(value, MethodType) and value.__self__ is layer: @@ -35,8 +35,7 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.T def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor: - """ - Restores references to layer held by tensor attributes. + """Restores references to layer held by tensor attributes. Used by `restore_layer_on_meta` to add back layer references, allowing for proper weight loading. @@ -47,6 +46,7 @@ def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Te Returns: sanitized tensor + """ for key, value in tensor.__dict__.items(): if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel: diff --git a/vllm/model_executor/model_loader/reload/torchao_decorator.py b/vllm/model_executor/model_loader/reload/torchao_decorator.py index 7fbc1c32944a..dc8e0acc6eaf 100644 --- a/vllm/model_executor/model_loader/reload/torchao_decorator.py +++ b/vllm/model_executor/model_loader/reload/torchao_decorator.py @@ -27,8 +27,7 @@ def set_torchao_reload_attrs(model: torch.nn.Module, model_config: ModelConfig): def support_quantized_model_reload_from_hp_weights(original_load_weights: FunctionType): - """ - Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support + """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support reloading high precision (bfloat16/float16/float32) weight for an already quantized model, this involves restoring the weights to a high precision weights and then online quantize the weights. diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py index b6d129a79403..b7dc361503d9 100644 --- a/vllm/model_executor/model_loader/reload/utils.py +++ b/vllm/model_executor/model_loader/reload/utils.py @@ -46,14 +46,14 @@ def get_layer_size(layer: torch.nn.Module) -> int: def has_device_tensors(bound_args: BoundArguments) -> bool: - """ - Return True if the loaded weights exist on an accelerator device + """Return True if the loaded weights exist on an accelerator device. Args: bound_args: args to load weights Returns: True if weights are on accelerator device + """ return any( isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu") @@ -62,14 +62,14 @@ def has_device_tensors(bound_args: BoundArguments) -> bool: def get_info_size(info: LayerReloadingInfo) -> int: - """ - Calculate the number of bytes used by loaded weights for a given layer + """Calculate the number of bytes used by loaded weights for a given layer. Args: info: layerwise info to get size of Returns: number of bytes used by loaded weights + """ return sum( value.nbytes diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 3f30a42fe406..b38729d28341 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -19,8 +19,7 @@ class RunaiModelStreamerLoader(BaseModelLoader): - """ - Model loader that can load safetensors + """Model loader that can load safetensors files from local FS, S3, GCS, or Azure Blob Storage. """ @@ -82,8 +81,8 @@ def _prepare_weights( ) -> list[str]: """Prepare weights for the model. - If the model is not local, it will be downloaded.""" - + If the model is not local, it will be downloaded. + """ is_object_storage_path = is_runai_obj_uri(model_name_or_path) is_local = os.path.isdir(model_name_or_path) safetensors_pattern = "*.safetensors" @@ -127,7 +126,7 @@ def _get_weights_iterator( ) def download_model(self, model_config: ModelConfig) -> None: - """Download model if necessary""" + """Download model if necessary.""" self._prepare_weights(model_config.model, model_config.revision) def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index 3f57fe7e0265..d17bfdc6d495 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -27,8 +27,7 @@ class ShardedStateLoader(BaseModelLoader): - """ - Model loader that directly loads each worker's model state dict, which + """Model loader that directly loads each worker's model state dict, 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. See `examples/features/sharded_state/save_sharded_state_offline.py` for creating @@ -57,8 +56,7 @@ def __init__(self, load_config: LoadConfig): def _filter_subtensors( tensors: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]: - """ - Filter out all tensors that share the same memory or a subset of the + """Filter out all tensors that share the same memory or a subset of the memory of another tensor. """ same_storage_groups: dict[Any, list[tuple[str, torch.Tensor]]] = ( diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index 3ef9b6bf2a59..582ee0ea3e05 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -403,8 +403,7 @@ def __init__(self, tensorizer_config: TensorizerConfig): @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: - """Tensorizer CLI arguments""" - + """Tensorizer CLI arguments.""" # Tensorizer options arg group group = parser.add_argument_group( "tensorizer options", @@ -498,7 +497,8 @@ def _check_tensors_on_meta_device(model: nn.Module) -> None: def _resize_lora_embeddings(model: nn.Module): """Modify LoRA embedding layers to use bigger tensors - to allow for adapter added tokens.""" + to allow for adapter added tokens. + """ for child in model.modules(): if ( isinstance(child, VocabParallelEmbedding) @@ -593,8 +593,7 @@ def tensorizer_weights_iterator( def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool: - """ - Infer if the model is a vLLM model by checking the weights for + """Infer if the model is a vLLM model by checking the weights for a vLLM tensorized marker. Args: @@ -603,6 +602,7 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool: Returns: bool: True if the model is a vLLM model, False otherwise. + """ tensorizer_args = tensorizer_config._construct_tensorizer_args() deserializer = TensorDeserializer( @@ -711,7 +711,7 @@ def tensorize_vllm_model( tensorizer_config: TensorizerConfig, generate_keyfile: bool = True, ): - """Utility to load a model and then serialize it with Tensorizer + """Utility to load a model and then serialize it with Tensorizer. Intended to be used separately from running a vLLM server since it creates its own Engine instance. @@ -773,8 +773,7 @@ def shutdown_engine_core() -> None: def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig): - """ - Uses tensorizer to serialize a LoRA adapter. Assumes that the files + """Uses tensorizer to serialize a LoRA adapter. Assumes that the files needed to load a LoRA adapter are a safetensors-format file called adapter_model.safetensors and a json config file called adapter_config.json. diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py index 338f9eac072a..e231eea20bae 100644 --- a/vllm/model_executor/model_loader/tensorizer_loader.py +++ b/vllm/model_executor/model_loader/tensorizer_loader.py @@ -105,7 +105,8 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: Expects a vLLM-tensorized model. See the examples/features/tensorize_vllm_model.py example script - for serializing vLLM models.""" + for serializing vLLM models. + """ if is_vllm_tensorized(self.tensorizer_config): tensorizer_config = self._patch_tensorizer_config(model_config) deserialize_tensorizer_model(model, tensorizer_config) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 11e6be5cf8a1..d80b7d17027d 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -278,8 +278,7 @@ def get_architecture_class_name(model_config: ModelConfig) -> str: def configure_quant_config( quant_config: QuantizationConfig, model_class: type[nn.Module] ): - """ - Pass packed_modules_mapping by reference to quant_config so that + """Pass packed_modules_mapping by reference to quant_config so that quant_config can properly match fused modules Note that model attributes are passed by reference to quant_config, diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index f0af5a284a4d..c009cc6ffa29 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -105,6 +105,7 @@ def export_entries( A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to its ``TensorEntry`` and ``aliases`` maps each duplicate name to its canonical name. + """ entries: dict[str, TensorEntry] = {} aliases: dict[str, str] = {} @@ -206,6 +207,7 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non Args: ready_callback: Invoked once the socket is bound and listening, so the launcher can report overall readiness. + """ socket_path = self._socket_path ensure_private_socket_dir( diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 6c3fb742cdf7..52d770f5d1c9 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """IPC model loader: maps post-quantized weights from a local weight cache -daemon via CUDA IPC instead of loading from disk.""" +daemon via CUDA IPC instead of loading from disk. +""" import dataclasses import socket diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 00d3dbfa8ba9..852856b139f1 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -74,6 +74,7 @@ def check_ipc_platform_support() -> None: Raises: UnsupportedPlatformForIPCError: If the current platform is not CUDA/ROCm. + """ if current_platform.is_cuda_alike(): return @@ -94,6 +95,7 @@ def check_ipc_quant_support(model: torch.nn.Module) -> None: Raises: UnsupportedQuantForIPCError: If any quant method does not declare ``supports_pre_processed_weights``. + """ for name, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) diff --git a/vllm/model_executor/model_loader/weight_tying.py b/vllm/model_executor/model_loader/weight_tying.py index d0c309d2c4ad..bd0bfa4b8038 100644 --- a/vllm/model_executor/model_loader/weight_tying.py +++ b/vllm/model_executor/model_loader/weight_tying.py @@ -78,7 +78,8 @@ def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> Checkpoints produced by quantization or fine-tuning tooling often keep a redundant copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would - otherwise cost.""" + otherwise cost. + """ if not model_config.word_embeddings_untied_by_checkpoint: return if (untied := _get_untied_lm_head(model)) is None: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index d291f1929ca0..2a9937bcd441 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -72,7 +72,7 @@ def enable_xet_high_performance(): - """automatically activates xet high performance mode""" + """Automatically activates xet high performance mode""" if "HF_XET_HIGH_PERFORMANCE" not in os.environ: huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True @@ -103,8 +103,7 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): def atomic_writer( filepath: str | Path, mode: str = "w", encoding: str | None = None ) -> Generator[IO]: - """ - Context manager that provides an atomic file writing routine. + """Context manager that provides an atomic file writing routine. The context manager writes to a temporary file and, if successful, atomically replaces the original file. @@ -116,6 +115,7 @@ def atomic_writer( Yields: file object: A handle to the temporary file. + """ # Create a temporary file in the same directory as the target file # to ensure it's on the same filesystem for an atomic replace. @@ -144,7 +144,8 @@ def atomic_writer( def _natural_sort_key(filepath: str) -> list: """Natural sort key for filenames with numeric components, such as - model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']""" + model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors'] + """ return [ int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath)) @@ -161,7 +162,8 @@ def maybe_download_from_modelscope( """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. Returns the path to the downloaded model, or None if the model is not - downloaded from ModelScope.""" + downloaded from ModelScope. + """ if envs.VLLM_USE_MODELSCOPE: # download model from ModelScope hub, # lazy import so that modelscope is not required for normal use. @@ -482,6 +484,7 @@ def download_weights_from_hf( Returns: str: The path to the downloaded model weights. + """ assert len(allow_patterns) > 0 local_only = huggingface_hub.constants.HF_HUB_OFFLINE @@ -580,6 +583,7 @@ def download_safetensors_index_file_from_hf( subfolder (Optional[str]): The subfolder within the model repository to download weights from. revision (Optional[str]): The revision of the model. + """ # Use file lock to prevent multiple processes from # downloading the same model weights at the same time. @@ -637,8 +641,7 @@ def filter_duplicate_safetensors_files( def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: - """ - Exclude files that are not needed for inference. + """Exclude files that are not needed for inference. See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 """ @@ -1145,7 +1148,8 @@ def instanttensor_weights_iterator( use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files - using instanttensor library.""" + using instanttensor library. + """ try: import instanttensor except ImportError as e: @@ -1407,6 +1411,7 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: str: The remapped parameter name if successful, or the original name if no remapping is needed. None: If the remapped name is not found in params_dict. + """ # Already in vLLM's expected form (e.g. weights pre-renamed by a # `WeightsMapper` from the quant config). Skip the regex remap, which @@ -1506,8 +1511,7 @@ def maybe_remap_moe_expert_param_name( name: str, params_dict: dict[str, torch.nn.Parameter], ) -> str: - """ - Remap MoE expert parameter names to account for routed_experts hierarchy. + """Remap MoE expert parameter names to account for routed_experts hierarchy. This handles the transition from the old FusedMoE structure where weights were directly in the experts module, to the new MoERunner → RoutedExperts @@ -1529,6 +1533,7 @@ def maybe_remap_moe_expert_param_name( Returns: Remapped parameter name if routed_experts hierarchy exists, otherwise the original name + """ # Only remap if this looks like an expert parameter if ".experts." not in name: @@ -1582,8 +1587,7 @@ def remap_moe_expert_weights( weights: Iterable[tuple[str, torch.Tensor]], params_dict: dict[str, torch.nn.Parameter], ) -> Generator[tuple[str, torch.Tensor], None, None]: - """ - Wrapper generator that remaps MoE expert parameter names for backward compatibility. + """Remap MoE expert parameter names for backward compatibility. This allows models with custom weight loading to automatically handle both old and new checkpoint formats without needing model-specific remapping code. @@ -1601,6 +1605,7 @@ def remap_moe_expert_weights( Yields: (remapped_name, tensor) tuples + """ for name, weight in weights: remapped_name = maybe_remap_moe_expert_param_name(name, params_dict) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 031808ee12b8..97f39504286f 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -400,8 +400,7 @@ def forward( class AXK1MLAAttention(nn.Module): - """ - Main reference: DeepseekV2 paper, and FlashInfer Implementation + """Main reference: DeepseekV2 paper, and FlashInfer Implementation (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551). For more info see MLACommonImpl in: diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py index ea75ff8cc0b1..41394607621b 100644 --- a/vllm/model_executor/models/adapters.py +++ b/vllm/model_executor/models/adapters.py @@ -39,7 +39,6 @@ def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None: """Load Sentence-Transformers Dense projection layers.""" - dense_modules = try_get_dense_modules( model_config.model, revision=model_config.revision ) @@ -249,8 +248,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): def as_embedding_model(cls: type[_T]) -> type[_T]: - """ - Subclass an existing vLLM model to support embeddings. + """Subclass an existing vLLM model to support embeddings. By default, the embeddings of the whole prompt are extracted from the normalized hidden state corresponding to the last token. @@ -258,6 +256,7 @@ def as_embedding_model(cls: type[_T]) -> type[_T]: Note: We assume that no extra layers are added to the original model; please implement your own model if this is not the case. + """ # Avoid modifying existing embedding models if is_pooling_model(cls): @@ -306,8 +305,7 @@ def _resolve_num_labels(hf_config: Any, text_config: Any) -> int: def as_seq_cls_model(cls: type[_T]) -> type[_T]: - """ - Subclass an existing vLLM model to support classify and score tasks. + """Subclass an existing vLLM model to support classify and score tasks. By default, the class probabilities are extracted from the softmaxed hidden state corresponding to the last token. @@ -316,6 +314,7 @@ def as_seq_cls_model(cls: type[_T]) -> type[_T]: We assume that the classification head is a single linear layer stored as the attribute `score` of the top-level model; please implement your own model if this is not the case. + """ # Avoid modifying existing classification models if is_pooling_model(cls): @@ -446,8 +445,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module: - """ - Get the language model component for sequence classification conversion. + """Get the language model component for sequence classification conversion. For VLMs, returns the inner language model. For standard LLMs, returns model itself. """ multimodal_model: object = model @@ -481,8 +479,7 @@ def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module: @contextmanager def _disable_seq_cls_loading_on_inner_model(language_model, is_vlm: bool): - """ - Context manager to temporarily disable sequence classification loading + """Context manager to temporarily disable sequence classification loading on inner VLM models to prevent recursive seq_cls_model_loader calls. """ if not is_vlm: diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index 32ad19a8e73b..a75bdaee5abe 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -47,7 +47,8 @@ class ArceeMLP(nn.Module): """Feed-forward layer for Arcee using ReLU^2 activation - (no gating as in LLaMA).""" + (no gating as in LLaMA). + """ def __init__( self, @@ -95,7 +96,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ArceeDecoderLayer(nn.Module): """Transformer decoder block for Arcee, with self-attention and - ReLU^2 MLP.""" + ReLU^2 MLP. + """ def __init__( self, @@ -175,7 +177,8 @@ def forward( @support_torch_compile class ArceeModel(nn.Module, EagleModelMixin): """The transformer model backbone for Arcee (embedding layer + stacked - decoder blocks + final norm).""" + decoder blocks + final norm). + """ def __init__( self, @@ -277,7 +280,8 @@ class ArceeForCausalLM( nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 ): """Arcee Model for causal language modeling, integrated with vLLM - runtime.""" + runtime. + """ hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ @@ -358,7 +362,8 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights into the model (delegates to inner model and handles - tied embeddings).""" + tied embeddings). + """ loader = AutoWeightsLoader(self) # AutoWeightLoader handles weight name remapping, including fusing # separate q_proj, k_proj, v_proj into qkv_proj diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 57eb62d3e26a..4240406f7b55 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -41,13 +41,12 @@ class AriaImagePixelInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - n: Number of images - - c: Number of channels - - h: Height of each image - - w: Width of each image + """Dimensions: + - b: Batch size + - n: Number of images + - c: Number of channels + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -119,8 +118,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class AriaProjector(nn.Module): - """ - A projection module with one cross attention layer and one FFN layer, which + """A projection module with one cross attention layer and one FFN layer, which projects ViT's outputs into MoE's inputs. Args: @@ -129,6 +127,7 @@ class AriaProjector(nn.Module): Outputs: A tensor with the shape of (batch_size, query_number, output_dim) + """ def __init__(self, config: AriaConfig, prefix: str = "") -> None: @@ -187,8 +186,7 @@ def forward( class AriaTextMoELayer(nn.Module): - """ - Mixture of Experts (MoE) Layer for the AriaMoE model. + """Mixture of Experts (MoE) Layer for the AriaMoE model. This layer implements the MoE mechanism, which routes input tokens to different experts based on a routing algorithm, processes them through the @@ -227,8 +225,7 @@ def __init__( ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the MoE Layer. + """Forward pass of the MoE Layer. Args: hidden_states: Input tensor of shape @@ -236,16 +233,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Output tensor after passing through the MoE layer. - """ + """ router_output = torch.nn.functional.linear(hidden_states, self.router_weight) return self.experts(hidden_states, router_output) class AriaTextDecoderLayer(LlamaDecoderLayer): - """ - Custom Decoder Layer for the AriaMoE model which modifies the standard + """Custom Decoder Layer for the AriaMoE model which modifies the standard `LlamaDecoderLayer` by replacing the traditional MLP with a Mixture of Experts (MoE) Layer. """ @@ -262,8 +258,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: class AriaTextModel(LlamaModel, SupportsQuant): - """ - Custom LlamaModel for the AriaMoE model which modifies the standard + """Custom LlamaModel for the AriaMoE model which modifies the standard LlamaModel by replacing the `LlamaDecoderLayer` with `MoEDecoderLayer`. """ @@ -377,8 +372,7 @@ def _get_prompt_updates( dummy_inputs=AriaDummyInputsBuilder, ) class AriaForConditionalGeneration(nn.Module, SupportsMultiModal): - """ - Aria model for conditional generation tasks. + """Aria model for conditional generation tasks. This model combines a vision tower, a multi-modal projector, and a language model to perform tasks that involve both image and text inputs. diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 7d6e51f4d336..4832d445d697 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -71,11 +71,10 @@ class AudioFlamingo3FeatureInputs(TensorSchema): - """ - Dimensions: - - num_chunks: Number of audio chunks (flattened) - - nmb: Number of mel bins - - num_audios: Number of original audio files + """Dimensions: + - num_chunks: Number of audio chunks (flattened) + - nmb: Number of mel bins + - num_audios: Number of original audio files """ type: Literal["audio_features"] @@ -96,12 +95,11 @@ class AudioFlamingo3FeatureInputs(TensorSchema): class AudioFlamingo3EmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size - - naf: Number of audio features - - hs: Hidden size (must match the hidden size of language model - backbone) + """Dimensions: + - bn: Batch size + - naf: Number of audio features + - hs: Hidden size (must match the hidden size of language model + backbone) """ type: Literal["audio_embeds"] = "audio_embeds" @@ -154,8 +152,7 @@ def forward( return hidden_states def _get_feat_extract_output_lengths(self, input_lengths: torch.Tensor): - """ - Computes the output length of the convolutional layers and the output length + """Computes the output length of the convolutional layers and the output length of the audio encoder """ input_lengths = (input_lengths - 1) // 2 + 1 @@ -503,9 +500,7 @@ class AudioFlamingo3ForConditionalGeneration( } def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model.", connector="multi_modal_projector.", diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index 9603dc043cc7..9e0bd4ccbf87 100644 --- a/vllm/model_executor/models/bagel.py +++ b/vllm/model_executor/models/bagel.py @@ -58,12 +58,11 @@ class BagelImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -175,12 +174,12 @@ def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos): return emb def forward(self, position_ids: torch.Tensor) -> torch.Tensor: - """ - Args: + """Args: position_ids: Flattened position IDs, shape (N,) where each ID corresponds to a position in the flattened grid Returns: Position embeddings of shape (N, hidden_size) + """ # Ensure position_ids are on the same device as pos_embed position_ids = position_ids.to(self.pos_embed.device) @@ -322,8 +321,7 @@ def _get_mm_fields_config( class BagelForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP ): - """ - BAGEL: A unified multimodal model for image understanding and generation. + """BAGEL: A unified multimodal model for image understanding and generation. For vLLM, we focus on the image understanding (vision-to-text) capabilities. The image generation part is not supported in vLLM. @@ -507,6 +505,9 @@ def forward( positions: Flattened (concatenated) position ids corresponding to a batch. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. + """ if intermediate_tensors is not None: inputs_embeds = None diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 8731e91c5d6e..58a06c5b2fc1 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -75,9 +75,7 @@ def is_linear_layer(layer_idx, layer_group_size): class BailingMoeV25MLAAttention(nn.Module): - """ - MLA Attention for BailingMoeV2.5 full attention layers. - """ + """MLA Attention for BailingMoeV2.5 full attention layers.""" def __init__( self, @@ -594,7 +592,6 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load checkpoint weights with simplified mapping.""" - params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index cdc7870f3ec5..85ef42675b43 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -458,6 +458,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant): Attributes: model: An instance of BertModel used for forward operations. _pooler: An instance of Pooler used for pooling operations. + """ is_pooling_model = True @@ -575,8 +576,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class SPLADESparsePooler(Pooler): - """ - SPLADE sparse pooling: + """SPLADE sparse pooling: logits = mlm_head(hidden_states) -> log1p(relu(logits)) -> (max|sum over L) @@ -664,8 +664,7 @@ def forward( @default_pooling_type(seq_pooling_type="CLS") class BertSpladeSparseEmbeddingModel(BertEmbeddingModel): - """ - BertEmbeddingModel + SPLADE sparse embedding. + """BertEmbeddingModel + SPLADE sparse embedding. - Make logits by self.mlm_head - pooler: SPLADESparsePooler(mlm_head...) """ @@ -779,6 +778,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu Attributes: model: An instance of BertModel used for forward operations. _pooler: An instance of Pooler used for pooling operations. + """ is_pooling_model = True diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index aefd17003d2f..09f0973a5cca 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Minimal implementation of BlipVisionModel intended to be only used -within a vision language model.""" +within a vision language model. +""" from collections.abc import Iterable @@ -82,7 +83,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: class BlipAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -140,8 +141,7 @@ def forward( self, hidden_states: torch.Tensor, ): - """Input shape: Batch x Time x Channel""" - + """Input shape: Batch x Time x Channel.""" qkv_states, _ = self.qkv(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) out = self.attn(query_states, key_states, value_states) @@ -220,12 +220,12 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class BlipEncoder(nn.Module): - """ - Transformer encoder consisting of `config.num_hidden_layers` self + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`BlipEncoderLayer`]. Args: config: BlipConfig + """ def __init__( diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index de74d1f5ace0..48da4f63e825 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -50,12 +50,11 @@ class Blip2ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -63,11 +62,10 @@ class Blip2ImagePixelInputs(TensorSchema): class Blip2ImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - f: Image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - bn: Batch size * number of images + - f: Image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] @@ -657,11 +655,15 @@ def forward( Args: input_ids: Flattened (concatenated) input_ids corresponding to a batch. + positions: Flattened (concatenated) position ids corresponding to a + batch. + intermediate_tensors: Intermediate tensors from prior forward pass. + inputs_embeds: Optional tensor of input embeddings. Info: [`Blip2ImageInputs`][vllm.model_executor.models.blip2.Blip2ImageInputs] - """ + """ if intermediate_tensors is not None: inputs_embeds = None diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py index f0be4890ace4..6fb311fb14c8 100644 --- a/vllm/model_executor/models/clip.py +++ b/vllm/model_executor/models/clip.py @@ -67,12 +67,11 @@ class CLIPImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -417,8 +416,7 @@ def forward( self, hidden_states: torch.Tensor, ): - """Input shape: Batch x Time x Channel""" - + """Input shape: Batch x Time x Channel.""" qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) out = self.attn(query_states, key_states, value_states) @@ -506,12 +504,12 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class CLIPEncoder(nn.Module): - """ - Transformer encoder consisting of `config.num_hidden_layers` self + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`CLIPEncoderLayer`]. Args: config: CLIPConfig + """ def __init__( diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index d657b78a5da1..ebfd93e857a6 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -96,7 +96,8 @@ def forward(self, hidden_states, residuals=None): def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]: """Returns (norm_class, eps). Uses RMSNorm when config.rms_norm_eps is set, - otherwise falls back to LayerNorm with config.layer_norm_eps.""" + otherwise falls back to LayerNorm with config.layer_norm_eps. + """ rms_eps = getattr(config, "rms_norm_eps", None) if rms_eps is not None: return RMSNorm, rms_eps diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index 42252afd67a2..fbb853beb0a9 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -61,14 +61,13 @@ class Cohere2VisionImagePixelInputs(TensorSchema): - """ - Dimensions: - - np: The total number of patches over each image over each prompt in - the batch - - c: Number of channels - - h: Height of each image patch - - w: Width of each image patch - - bn: Batch size * number of images + """Dimensions: + - np: The total number of patches over each image over each prompt in + the batch + - c: Number of channels + - h: Height of each image patch + - w: Width of each image patch + - bn: Batch size * number of images """ type: Literal["pixel_values"] @@ -134,6 +133,7 @@ def pixel_shuffle(self, image_features: torch.Tensor) -> torch.Tensor: Returns: Downsampled tensor with increased channel dimension + """ height = width = int(image_features.shape[1] ** 0.5) x = image_features.reshape(image_features.shape[0], width, height, -1) @@ -175,8 +175,7 @@ def get_num_patches( processor: Cohere2VisionProcessor, mm_kwargs: Mapping[str, object], ) -> int: - """ - Calculate the number of image patches for a given image. + """Calculate the number of image patches for a given image. Uses the HF processor to determine the actual number of patches. """ image_processor: Cohere2VisionImageProcessorFast = processor.image_processor @@ -372,16 +371,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def _process_image_input( - self, image_input: Cohere2VisionImagePixelInputs, **kwargs + self, image_input: Cohere2VisionImagePixelInputs, **kwargs: object ) -> list[torch.Tensor]: """Process image pixels through vision tower and projector. Args: image_input: Validated image input containing pixel values and patch counts + **kwargs: Unused; accepted for interface compatibility. Returns: List of flattened image embeddings, one per image + """ pixel_values = image_input["pixel_values"] num_patches = image_input["num_patches"] diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index d7f8f87fc6b1..29205e4ed6be 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -300,13 +300,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class FixedPositionalEncoding(nn.Module): - """ - Fixed positional encoding (embedding layer) from sine and cosine functions + """Fixed positional encoding (embedding layer) from sine and cosine functions of different frequencies according to https://arxiv.org/abs/1706.03762 Args: hidden_size: size of the embeddings in the model, also known as d_model max_sequence_length: maximum allowed length of the input sequence + """ def __init__(self, hidden_size: int, max_sequence_length: int = 512) -> None: @@ -664,7 +664,8 @@ def calc_length( repeat_num: int = 1, ) -> torch.Tensor: """Calculates the output length of a Tensor passed - through a convolution or max pooling layer""" + through a convolution or max pooling layer + """ add_pad: float = all_paddings - kernel_size one: float = 1.0 for i in range(repeat_num): @@ -689,10 +690,12 @@ def forward( class PositionalEncoding(torch.nn.Module): """Fixed sinusoidal positional encoding. + Args: d_model (int): embedding dim max_len (int): maximum input length xscale (bool): whether to scale the input by sqrt(d_model) + """ def __init__( @@ -724,12 +727,14 @@ def forward( self, x: torch.Tensor, cache_len: int = 0 ) -> tuple[torch.Tensor, torch.Tensor]: """Adds positional encoding. + Args: x (torch.Tensor): Input. Its shape is (batch, time, feature_size) cache_len (int): the size of the cache which is used to shift positions Returns: x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size) pos_emb (torch.Tensor): Its shape is (1, time, feature_size) + """ input_len = x.size(1) + cache_len if self.xscale: @@ -762,14 +767,15 @@ def forward( self, x: torch.Tensor, cache_len: int = 0 ) -> tuple[torch.Tensor, torch.Tensor]: """Compute positional encoding. + Args: x (torch.Tensor): Input. Its shape is (batch, time, feature_size) cache_len (int): the size of the cache which is used to shift positions Returns: x (torch.Tensor): Its shape is (batch, time, feature_size) pos_emb (torch.Tensor): Its shape is (1, time, feature_size) - """ + """ if self.xscale: x = x * self.xscale @@ -783,15 +789,13 @@ def forward( class Swish(nn.SiLU): - """ - Swish activation function introduced in 'https://arxiv.org/abs/1710.05941' + """Swish activation function introduced in 'https://arxiv.org/abs/1710.05941' Mathematically identical to SiLU. See note in nn.SiLU for references. """ class ConformerFeedForward(nn.Module): - """ - feed-forward module of Conformer model. + """feed-forward module of Conformer model. use_bias (bool): Apply bias to all Linear and Conv1d layers to improve activation flow and stabilize training of huge models. @@ -819,8 +823,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class CausalConv1D(nn.Conv1d): - """ - A causal version of nn.Conv1d where each step would + """A causal version of nn.Conv1d where each step would have limited access to locations on its right or left. All arguments are the same as nn.Conv1d except padding. @@ -893,6 +896,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ConformerConvolution(nn.Module): """The convolution module for the Conformer model. + Args: d_model (int): hidden dimension kernel_size (int): kernel size for depthwise convolution @@ -904,6 +908,7 @@ class ConformerConvolution(nn.Module): use_bias (bool): Use bias in all Linear and Conv1d layers to improve activation flow and stabilize training of huge models. Defaults to True + """ def __init__( @@ -979,10 +984,12 @@ def forward( class CohereASRMultiHeadAttention(nn.Module): """Multi-Head Attention layer of Transformer. + Args: n_head (int): number of heads n_feat (int): size of the features use_bias (bool): whether to remove bias in linear and conv layers + """ def __init__( @@ -1010,6 +1017,7 @@ def forward_qkv( value: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Transforms query, key and value. + Args: query (torch.Tensor): (batch, time1, size) key (torch.Tensor): (batch, time2, size) @@ -1018,6 +1026,7 @@ def forward_qkv( q (torch.Tensor): (batch, head, time1, size) k (torch.Tensor): (batch, head, time2, size) v (torch.Tensor): (batch, head, time2, size) + """ n_batch = query.size(0) q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) @@ -1036,6 +1045,7 @@ def forward_attention( mask: torch.Tensor | None, ) -> torch.Tensor: """Compute attention context vector. + Args: value (torch.Tensor): (batch, time2, size) scores(torch.Tensor): (batch, time1, time2) @@ -1044,6 +1054,7 @@ def forward_attention( value (torch.Tensor): transformed `value` (batch, time2, d_model) weighted by the attention scores + """ n_batch = value.size(0) if mask is not None: @@ -1071,16 +1082,19 @@ def forward( pos_emb: torch.Tensor | None = None, ) -> torch.Tensor: """Compute 'Scaled Dot Product Attention'. + Args: query (torch.Tensor): (batch, time1, size) key (torch.Tensor): (batch, time2, size) value(torch.Tensor): (batch, time2, size) mask (torch.Tensor): (batch, time1, time2) + pos_emb (torch.Tensor): optional relative position embeddings - returns: + Returns: output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention + """ q, k, v = self.forward_qkv(query, key, value) @@ -1130,8 +1144,10 @@ def __init__( def rel_shift(self, x: torch.Tensor) -> torch.Tensor: """Compute relative positional encoding. + Args: x (torch.Tensor): (batch, nheads, time, 2*time-1) + """ b, h, qlen, pos_len = x.size() # (b, h, t1, t2) # need to add a column of zeros on the left side of @@ -1151,17 +1167,19 @@ def forward( pos_emb: torch.Tensor | None = None, ) -> torch.Tensor: """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: query (torch.Tensor): (batch, time1, size) key (torch.Tensor): (batch, time2, size) value(torch.Tensor): (batch, time2, size) mask (torch.Tensor): (batch, time1, time2) - pos_emb (torch.Tensor) : (batch, time1, size) + pos_emb (torch.Tensor): (batch, time1, size) Returns: output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention + """ q, k, v = self.forward_qkv(query, key, value) q = q.transpose(1, 2) # (batch, time1, head, d_k) @@ -1212,6 +1230,7 @@ class ConformerLayer(torch.nn.Module): Conv1d layers from each ConformerLayer to improve activation flow and stabilize training of huge models. Defaults to True. + """ def __init__( @@ -1279,14 +1298,14 @@ def forward( pos_emb: torch.Tensor | None = None, pad_mask: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Args: + """Args: x (torch.Tensor): input signals (B, T, d_model) att_mask (torch.Tensor): attention masks(B, T, T) pos_emb (torch.Tensor): (L, 1, d_model) pad_mask (torch.tensor): padding mask Returns: x (torch.Tensor): (B, T, d_model) + """ residual = x x = self.norm_feed_forward1(x) @@ -1331,8 +1350,7 @@ def forward( class ConformerEncoder(nn.Module): - """ - The encoder for ASR model of Conformer. + """The encoder for ASR model of Conformer. Based on this paper: 'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al. @@ -1473,12 +1491,12 @@ def get_num_encoder_cross_attn_tokens(self, num_encoder_input_tokens: int) -> in return num_encoder_cross_attn_tokens def set_max_audio_length(self, max_audio_length: int) -> None: - """ - Sets maximum input length. + """Sets maximum input length. Pre-calculates internal seq_range mask. Args: max_audio_length (int): New maximum sequence length. + """ device = next(self.parameters()).device dtype = next(self.parameters()).dtype diff --git a/vllm/model_executor/models/cohere_compass.py b/vllm/model_executor/models/cohere_compass.py index 0256f747c90c..1d0dd8c35b05 100644 --- a/vllm/model_executor/models/cohere_compass.py +++ b/vllm/model_executor/models/cohere_compass.py @@ -338,8 +338,7 @@ def pos_embed_interpolate_native( class Cohere_VLImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - np: Number of patches - ni: Number of images - cps: Number of channels * patch_size * patch_size @@ -365,8 +364,7 @@ class Cohere_VLImagePixelInputs(TensorSchema): class Cohere_VLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - nf: Number of image features - hs: Hidden size - ni: Number of images @@ -895,6 +893,7 @@ def prepare_encoder_metadata( instead of computing from cu_seqlens (needed for CUDA graph capture to cover worst-case replay scenarios). device: Device to place tensors on. Defaults to self.device. + """ if device is None: device = self.device @@ -1744,8 +1743,7 @@ def _postprocess_image_embeds_evs( image_embeds_split: tuple[torch.Tensor, ...], image_input: Cohere_VLImageInputs, ) -> tuple[torch.Tensor, ...]: - """ - Append mrope positions for each for images. + """Append mrope positions for each for images. This is necessary to recover correct mrope positions Args: @@ -1757,6 +1755,7 @@ def _postprocess_image_embeds_evs( Tuple of image embeddings for each image item. Resulting embeddings will have extra 5 channels for computed mrope positions, consistent with video embeddings. + """ if self.is_multimodal_pruning_enabled: merge_size = self.visual.spatial_merge_size @@ -1837,8 +1836,7 @@ def recompute_mrope_positions( mrope_positions: torch.LongTensor, num_computed_tokens: int, ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]: - """ - Update part of input mrope positions (starting with + """Update part of input mrope positions (starting with num_computed_tokens index). Original mrope_positions are computed for unpruned sequence and becomes incorrect once pruning occurs, so once we prune media tokens we should reflect this in the @@ -1856,6 +1854,7 @@ def recompute_mrope_positions( Returns: Tuple of (multimodal_embeddings, mrope_positions, mrope_position_delta). + """ return self._recompute_mrope_positions( input_ids=input_ids, @@ -2046,8 +2045,8 @@ def forward( `None` if no images are passed. - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in LLM. `None` if no images are passed. - """ + """ if intermediate_tensors is not None: inputs_embeds = None @@ -2083,9 +2082,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector=["visual.merger", "visual.deepstack_merger_list"], diff --git a/vllm/model_executor/models/colbert.py b/vllm/model_executor/models/colbert.py index 124cfb5fdbec..afdd86072100 100644 --- a/vllm/model_executor/models/colbert.py +++ b/vllm/model_executor/models/colbert.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColBERT late interaction model for retrieval and reranking. +"""ColBERT late interaction model for retrieval and reranking. ColBERT uses per-token embeddings and late interaction (MaxSim) scoring instead of single-vector representations or cross-encoder concatenation. @@ -75,6 +74,7 @@ def _init_colbert_components( ``None``, will be inferred from weights during loading (or auto-loaded from sentence-transformers Dense layers). head_dtype: Data type for the projection layer. + """ self.hidden_size = hidden_size self.colbert_dim = colbert_dim @@ -149,6 +149,7 @@ def _load_colbert_weights( Returns: ``(remaining_weights, loaded_names)`` — the weights that were **not** consumed and the set of names that were loaded. + """ weights_list = list(weights) other_weights: list[tuple[str, torch.Tensor]] = [] diff --git a/vllm/model_executor/models/colpali.py b/vllm/model_executor/models/colpali.py index f31d63c00632..aea809c15ecc 100644 --- a/vllm/model_executor/models/colpali.py +++ b/vllm/model_executor/models/colpali.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColPali late interaction model for multi-modal retrieval and reranking. +"""ColPali late interaction model for multi-modal retrieval and reranking. ColPali extends PaliGemma with a ColBERT-style late interaction head, producing per-token embeddings for both text and image inputs. It uses diff --git a/vllm/model_executor/models/colqwen3.py b/vllm/model_executor/models/colqwen3.py index a2012fffd2ad..3f04c86b9649 100644 --- a/vllm/model_executor/models/colqwen3.py +++ b/vllm/model_executor/models/colqwen3.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColQwen3 late interaction model for multi-modal retrieval and reranking. +"""ColQwen3 late interaction model for multi-modal retrieval and reranking. ColQwen3 extends Qwen3-VL with a ColBERT-style late interaction head, producing per-token embeddings for both text and image inputs. It uses @@ -131,6 +130,7 @@ class ColQwen3Model(Qwen3VLForConditionalGeneration, SupportsLateInteraction): Attributes: custom_text_proj: Linear projection from hidden_size to embed_dim + """ # Mark this as a pooling model so vLLM routes to pooler path diff --git a/vllm/model_executor/models/colqwen3_5.py b/vllm/model_executor/models/colqwen3_5.py index dabb4885ee51..aa8a5d0c81a2 100644 --- a/vllm/model_executor/models/colqwen3_5.py +++ b/vllm/model_executor/models/colqwen3_5.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColQwen3.5 late interaction model for multi-modal retrieval and reranking. +"""ColQwen3.5 late interaction model for multi-modal retrieval and reranking. ColQwen3.5 extends Qwen3.5 with a ColBERT-style late interaction head, producing per-token embeddings for both text and image inputs. It uses @@ -128,6 +127,7 @@ class ColQwen3_5Model( Attributes: custom_text_proj: Linear projection from hidden_size to embed_dim. This is passed to the pooler as pooler.head.projector. + """ # Mark this as a pooling model so vLLM routes to pooler path diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 3fe78f081ea8..cb17cfa32f5a 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -114,8 +114,7 @@ def forward(self, hidden_states, residuals=None): def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]: - """ - Returns the normalization layer class and epsilon value to use. + """Returns the normalization layer class and epsilon value to use. If `config.rms_norm_eps` is present, use RMSNorm. Otherwise default to LayerNorm with `config.layer_norm_eps`. """ diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index bbbe98bc8de0..8288eab236f3 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -445,8 +445,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: - """ - Perform early validation and setup for hybrid attention/mamba models. + """Perform early validation and setup for hybrid attention/mamba models. Block size alignment with mamba page sizes is handled later by Platform.update_block_size_for_backend(), which runs after model @@ -454,6 +453,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: Args: vllm_config: vLLM Config + """ # Enable FULL_AND_PIECEWISE by default MambaModelConfig.verify_and_update_config(vllm_config) @@ -609,12 +609,12 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: class MambaModelConfig(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: - """ - Enable FULL_AND_PIECEWISE cuda graph mode by default (required + """Enable FULL_AND_PIECEWISE cuda graph mode by default (required to get good performance for mamba layers in V1). Args: vllm_config: vLLM Config + """ model_config = vllm_config.model_config cache_config = vllm_config.cache_config diff --git a/vllm/model_executor/models/conformer_encoder.py b/vllm/model_executor/models/conformer_encoder.py index 757dd1ac57c0..f4da532d5c57 100644 --- a/vllm/model_executor/models/conformer_encoder.py +++ b/vllm/model_executor/models/conformer_encoder.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Conformer encoder components for FireRedASR2. +"""Conformer encoder components for FireRedASR2. The audio encoder architecture is (Conv2dSubsampling → RelPositionalEncoding → N × RelPosEmbConformerBlock). @@ -298,9 +297,7 @@ def forward( class ConformerEncoder(nn.Module): - """ - Conformer encoder used by FireRedASR2. - """ + """Conformer encoder used by FireRedASR2.""" def __init__( self, diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py index 38e55405b895..ae65e25c901d 100644 --- a/vllm/model_executor/models/cosmos3_edge.py +++ b/vllm/model_executor/models/cosmos3_edge.py @@ -151,8 +151,7 @@ def patch_merging_by_param( class Cosmos3EdgePatchMerger(nn.Module): - """ - Projector: LayerNorm -> Linear -> GELU -> Linear + """Projector: LayerNorm -> Linear -> GELU -> Linear. Reads config from projector_config (not vision_config). input_hidden_size * spatial_merge_size² -> merger_intermediate_size @@ -490,8 +489,7 @@ class Cosmos3EdgeForConditionalGeneration( SupportsPP, SupportsMRoPE, ): - """ - Cosmos3 Edge model with a SigLIP2 vision encoder. + """Cosmos3 Edge model with a SigLIP2 vision encoder. Architecture: - self.visual: SigLIP2 encoder + patch merger + projector diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py index d75e5bb77d2f..824813eadb15 100644 --- a/vllm/model_executor/models/deepencoder.py +++ b/vllm/model_executor/models/deepencoder.py @@ -248,23 +248,23 @@ def __init__( global_attn_indexes: tuple[int, ...] = (), last_conv_output: int = 1024, ) -> None: - """ - Args: - img_size (int): Input image size. - patch_size (int): Patch size. - in_chans (int): Number of input image channels. - embed_dim (int): Patch embedding dimension. - depth (int): Depth of ViT. - num_heads (int): Number of attention heads in each ViT block. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool): If True, add a learnable bias to query, key, value. - norm_layer (nn.Module): Normalization layer. - act_layer (nn.Module): Activation layer. - use_abs_pos (bool): If True, use absolute positional embeddings. - use_rel_pos (bool): If True, add relative positional embeddings to the attention map. - rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. - window_size (int): Window size for window attention blocks. - global_attn_indexes (list): Indexes for blocks using global attention. + """Args: + img_size (int): Input image size. + patch_size (int): Patch size. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + depth (int): Depth of ViT. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_abs_pos (bool): If True, use absolute positional embeddings. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. + global_attn_indexes (list): Indexes for blocks using global attention. + """ # noqa: E501 super().__init__() self.img_size = img_size @@ -363,7 +363,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Block(nn.Module): """Transformer blocks with support of window attention and residual propagation - blocks""" + blocks + """ def __init__( self, @@ -378,20 +379,20 @@ def __init__( window_size: int = 0, input_size: tuple[int, int] | None = None, ) -> None: - """ - Args: - dim (int): Number of input channels. - num_heads (int): Number of attention heads in each ViT block. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool): If True, add a learnable bias to query, key, value. - norm_layer (nn.Module): Normalization layer. - act_layer (nn.Module): Activation layer. - use_rel_pos (bool): If True, add relative positional embeddings to the attention map. - rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. - window_size (int): Window size for window attention blocks. If it equals 0, then - use global attention. - input_size (tuple(int, int) or None): Input resolution for calculating the relative - positional parameter size. + """Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. If it equals 0, then + use global attention. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ # noqa: E501 super().__init__() self.norm1 = norm_layer(dim) @@ -448,16 +449,16 @@ def __init__( rel_pos_zero_init: bool = True, input_size: tuple[int, int] | None = None, ) -> None: - """ - Args: - dim (int): Number of input channels. - num_heads (int): Number of attention heads. - qkv_bias (bool): If True, add a learnable bias to query, key, value. - use_triton_attention (bool): If True, fuse relative position bias in - a Triton attention kernel on CUDA-alike platforms. - rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. - input_size (tuple(int, int) or None): Input resolution for calculating the relative - positional parameter size. + """Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + use_triton_attention (bool): If True, fuse relative position bias in + a Triton attention kernel on CUDA-alike platforms. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ # noqa: E501 super().__init__() self.num_heads = num_heads @@ -540,8 +541,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def window_partition( x: torch.Tensor, window_size: int ) -> tuple[torch.Tensor, tuple[int, int]]: - """ - Partition into non-overlapping windows with padding if needed. + """Partition into non-overlapping windows with padding if needed. + Args: x (tensor): input tokens with [B, H, W, C]. window_size (int): window size. @@ -549,6 +550,7 @@ def window_partition( Returns: windows: windows after partition with [B * num_windows, window_size, window_size, C]. (Hp, Wp): padded height and width before partition + """ # noqa: E501 B, H, W, C = x.shape @@ -571,8 +573,8 @@ def window_unpartition( pad_hw: tuple[int, int], hw: tuple[int, int], ) -> torch.Tensor: - """ - Window unpartition into original sequences and removing padding. + """Window unpartition into original sequences and removing padding. + Args: windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. window_size (int): window size. @@ -581,6 +583,7 @@ def window_unpartition( Returns: x: unpartitioned sequences with [B, H, W, C]. + """ # noqa: E501 Hp, Wp = pad_hw H, W = hw @@ -596,9 +599,9 @@ def window_unpartition( def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: - """ - Get relative positional embeddings according to the relative positions of + """Get relative positional embeddings according to the relative positions of query and key sizes. + Args: q_size (int): size of query q. k_size (int): size of key k. @@ -606,6 +609,7 @@ def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor Returns: Extracted positional embeddings according to relative positions. + """ max_rel_dist = int(2 * max(q_size, k_size) - 1) # Interpolate rel pos if needed. @@ -641,8 +645,7 @@ def add_decomposed_rel_pos( q_size: tuple[int, int], k_size: tuple[int, int], ) -> torch.Tensor: - """ - Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. + """Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py Args: q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). @@ -653,6 +656,7 @@ def add_decomposed_rel_pos( Returns: attn (Tensor): attention map with added relative positional embeddings. + """ # noqa: E501 q_h, q_w = q_size k_h, k_w = k_size @@ -672,9 +676,7 @@ def add_decomposed_rel_pos( class PatchEmbed(nn.Module): - """ - Image to Patch Embedding. - """ + """Image to Patch Embedding.""" def __init__( self, @@ -684,13 +686,13 @@ def __init__( in_chans: int = 3, embed_dim: int = 768, ) -> None: - """ - Args: - kernel_size (Tuple): kernel size of the projection layer. - stride (Tuple): stride of the projection layer. - padding (Tuple): padding size of the projection layer. - in_chans (int): Number of input image channels. - embed_dim (int): Patch embedding dimension. + """Args: + kernel_size (Tuple): kernel size of the projection layer. + stride (Tuple): stride of the projection layer. + padding (Tuple): padding size of the projection layer. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + """ super().__init__() diff --git a/vllm/model_executor/models/deepencoder2.py b/vllm/model_executor/models/deepencoder2.py index 432b4eb56d1c..3b3860bc2cd4 100644 --- a/vllm/model_executor/models/deepencoder2.py +++ b/vllm/model_executor/models/deepencoder2.py @@ -23,8 +23,7 @@ # --8<-- [start:qwen2_decoder] @PluggableLayer.register("qwen2_decoder") class CustomQwen2Decoder(PluggableLayer): - """ - Qwen2 visual encoder + """Qwen2 visual encoder non-causal attention + causal attention token_type_ids ：0=non-causal, 1=causal """ @@ -75,7 +74,7 @@ def __init__( del self.model.embed_tokens def _create_custom_model(self, config): - """Qwen2Model""" + """Qwen2Model.""" class CustomQwen2ModelInner(Qwen2Model): def forward( @@ -194,11 +193,11 @@ def forward( attention_mask: torch.Tensor = None, **kwargs, ): - """ - Args: - inputs_embeds: [batch_size, seq_len, hidden_dim] - token_type_ids: [batch_size, seq_len], 0=non-causal, 1=causal - attention_mask: [batch_size, seq_len], optional + """Args: + inputs_embeds: [batch_size, seq_len, hidden_dim] + token_type_ids: [batch_size, seq_len], 0=non-causal, 1=causal + attention_mask: [batch_size, seq_len], optional + """ return self.model( inputs_embeds=inputs_embeds, @@ -209,8 +208,7 @@ def forward( class Qwen2Decoder2Encoder(nn.Module): - """ - Decoder based on Multilingual BART + """Decoder based on Multilingual BART Set the initial weights and configuration with a pretrained multilingual BART model, and modify the detailed configurations as a Nougat decoder """ diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 48d019cd357b..9284dd59a09b 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -36,8 +36,7 @@ class DeepseekV2Eagle3DecoderLayer(nn.Module): - """ - Eagle3 decoder layer for Deepseek that: + """Eagle3 decoder layer for Deepseek that: 1. Always uses MLP (not MoE) 2. First layer accepts concatenated embeds + hidden_states """ diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 09f4a847d064..36f0a240bc7d 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -529,8 +529,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 114498764c65..6e438680956f 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -79,13 +79,12 @@ class DeepseekOCRImagePixelInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - n: Number of images - - p: Number of patches - - base_size: Base size of the processor - - image_size: Image size of the processor + """Dimensions: + - b: Batch size + - n: Number of images + - p: Number of patches + - base_size: Base size of the processor + - image_size: Image size of the processor """ type: Literal["pixel_values"] @@ -140,7 +139,8 @@ def __call__( class NGramPerReqLogitsProcessor(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): @@ -639,9 +639,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return autoloaded_weights def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="projector", @@ -676,8 +674,7 @@ def _get_num_input_output_tokens( self, image_spatial_crop: torch.Tensor | None = None, ) -> tuple[int, int, int, int]: - """ - Return (num_input_tokens, num_output_tokens, global_output_token, + """Return (num_input_tokens, num_output_tokens, global_output_token, local_output_token) for a single image described by ``image_spatial_crop``. """ @@ -854,8 +851,7 @@ def _batched_encoder_forward_global_path( self, pixel_values: torch.Tensor, ) -> torch.Tensor: - """ - Encode batched global images with newline tokens inserted. + """Encode batched global images with newline tokens inserted. Output shape: ``[B * 272, n_embed]``. """ bsz = pixel_values.shape[0] @@ -880,8 +876,7 @@ def _batched_encoder_forward_local_path( self, images_crop: torch.Tensor, ) -> torch.Tensor: - """ - Encode local patches without newline insertion (newlines are added later + """Encode local patches without newline insertion (newlines are added later in ``postprocess_encoder_output`` via ``_assemble_patch_grid``). Output shape: ``[P * 100, n_embed]``. """ @@ -951,8 +946,7 @@ def postprocess_encoder_output( clone: bool = False, batch_mm_kwargs: dict[str, Any] | None = None, ) -> None: - """ - Assemble per-image embeddings from global and local encoder outputs. + """Assemble per-image embeddings from global and local encoder outputs. ``output['global']`` contains global-image features with newlines already inserted (from CUDA graph replay or eager fallback): diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index 0e907a344a87..caaf90967183 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -456,9 +456,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return autoloaded_weights def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="projector", diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index bdc634102e2f..d7d4858044ed 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -850,8 +850,7 @@ def forward( def _try_load_fp8_indexer_wk( name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names ): - """ - We fuse the WK and weights_proj projections, but in some checkpoints WK is stored + """We fuse the WK and weights_proj projections, but in some checkpoints WK is stored in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16. Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK weights and scale, and when both are available, dequantizes to BF16 and stores into @@ -911,8 +910,7 @@ def _min_latency_fused_qkv_a_proj_impl( input_: torch.Tensor, weight: torch.Tensor, ) -> torch.Tensor: - """ - Dynamically run min-latency gemm if num_tokens <= 16. + """Dynamically run min-latency gemm if num_tokens <= 16. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ @@ -993,8 +991,7 @@ def forward( class DeepseekV2MLAAttention(nn.Module): - """ - Main reference: DeepseekV2 paper, and FlashInfer Implementation + """Main reference: DeepseekV2 paper, and FlashInfer Implementation (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551). For more info see MLACommonImpl in: diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index e77dd3d13f59..5fb52cff904a 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -64,13 +64,12 @@ class DeepseekVL2ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of images * number of patches - - p: Number of patches - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bnp: Batch size * number of images * number of patches + - p: Number of patches + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -79,11 +78,10 @@ class DeepseekVL2ImagePixelInputs(TensorSchema): class DeepseekVL2VImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - f: Image feature size - - h: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - f: Image feature size + - h: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index c90d78bd6dd1..7df3f744ce38 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -352,7 +352,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): skip duplicate decoder backbone weights, handle vision tower and self-conditioning separately. """ - sc_params = dict( (n, p) for n, p in self.named_parameters() @@ -512,7 +511,8 @@ def _compiled_sample_step( accept/renoise → convergence, all as vectorized PyTorch ops. Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the - caller can compute logprobs outside the compiled region.""" + caller can compute logprobs outside the compiled region. + """ num_decode = decode_slots.shape[0] device = decode_slots.device diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py index c939c2bef5b3..36b1cd530d6b 100644 --- a/vllm/model_executor/models/dots_ocr.py +++ b/vllm/model_executor/models/dots_ocr.py @@ -66,12 +66,11 @@ class DotsOCRImagePixelInputs(TensorSchema): - """ - Dimensions: - - np: The total number of patches over each image over each prompt in - the batch - - ni: Number of images - - cps: Number of channels * patch_size * patch_size + """Dimensions: + - np: The total number of patches over each image over each prompt in + the batch + - ni: Number of images + - cps: Number of channels * patch_size * patch_size """ type: Literal["pixel_values"] @@ -81,11 +80,10 @@ class DotsOCRImagePixelInputs(TensorSchema): class DotsOCRImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nf: Number of image features - - hs: Hidden size - - ni: Number of images + """Dimensions: + - nf: Number of image features + - hs: Hidden size + - ni: Number of images """ type: Literal["image_embeds"] @@ -768,9 +766,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="vision_tower.merger", diff --git a/vllm/model_executor/models/eagle2_5_vl.py b/vllm/model_executor/models/eagle2_5_vl.py index 8fdd7098f224..d4ee31aaf72a 100644 --- a/vllm/model_executor/models/eagle2_5_vl.py +++ b/vllm/model_executor/models/eagle2_5_vl.py @@ -37,13 +37,12 @@ class Eagle2_5_VLImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - bnp: Batch size * number of images * (1 + num_patches) - - c: Number of channels (3) - - h: Height of each image patch - - w: Width of each image patch + """Dimensions: + - bn: Batch size * number of images + - bnp: Batch size * number of images * (1 + num_patches) + - c: Number of channels (3) + - h: Height of each image patch + - w: Width of each image patch """ type: Literal["pixel_values"] @@ -52,11 +51,10 @@ class Eagle2_5_VLImagePixelInputs(TensorSchema): class Eagle2_5_VLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of images - - f: Total image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - n: Number of images + - f: Total image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] @@ -127,8 +125,7 @@ class Eagle2_5_VLMultiModalProcessor( class Eagle2_5_VLForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA ): - """ - Eagle2.5-VL model for conditional generation. + """Eagle2.5-VL model for conditional generation. Architecture: - Vision Encoder: SigLIP @@ -238,8 +235,7 @@ def _init_mlp1(self, config: PretrainedConfig) -> nn.Module: ) def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Tensor: - """ - Pixel shuffle operation for downsampling vision features. + """Pixel shuffle operation for downsampling vision features. Args: x: Input tensor of shape (n, w, h, c) @@ -247,6 +243,7 @@ def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Ten Returns: Downsampled tensor + """ n, w, h, c = x.size() # N, W, H, C --> N, W, H * scale, C // scale @@ -264,14 +261,14 @@ def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Ten return x def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor: - """ - Extract visual features from pixel values. + """Extract visual features from pixel values. Args: pixel_values: Input pixel values of shape (batch, channels, height, width) Returns: Visual embeddings + """ vit_embeds = self.vision_model(pixel_values=pixel_values) diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 716aeaa3b708..81907a517d43 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -119,7 +119,7 @@ def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int): class Ernie4_5_VisionAttention(nn.Module): - """VisionAttention using VLLM framework APIs""" + """VisionAttention using VLLM framework APIs.""" def __init__( self, @@ -497,6 +497,7 @@ def prepare_encoder_metadata( Returns: Dict with ``rotary_pos_emb``, ``cu_seqlens`` and ``max_seqlen``. + """ if device is None: device = self.device @@ -600,12 +601,11 @@ def load_weights(self, weights) -> set[str]: class Ernie4_5_VLImagePixelInputs(TensorSchema): - """ - Dimensions: - - np: The total number of patches over each image over each prompt in - the batch - - ni: Number of images - - cps: Number of channels * patch_size * patch_size + """Dimensions: + - np: The total number of patches over each image over each prompt in + the batch + - ni: Number of images + - cps: Number of channels * patch_size * patch_size """ type: Literal["pixel_values"] @@ -618,13 +618,12 @@ class Ernie4_5_VLImagePixelInputs(TensorSchema): class Ernie4_5_VLVideoPixelInputs(TensorSchema): - """ - Dimensions: - - np: The total number of patches over each image over each prompt in - the batch - - ni: Number of images - - cps: Number of channels * temporal_patch_size * patch_size * - patch_size + """Dimensions: + - np: The total number of patches over each image over each prompt in + the batch + - ni: Number of images + - cps: Number of channels * temporal_patch_size * patch_size * + patch_size """ type: Literal["pixel_values_videos"] diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index c11c64fcba5a..653e7691fba3 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -45,8 +45,7 @@ def unified_kv_cache_update( to_cache: torch.Tensor, layer_name: str, ) -> torch.Tensor: - """ - Returns a dummy that is passed to unified_attention to signal a side effect and + """Returns a dummy that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ forward_context = get_forward_context() @@ -294,6 +293,7 @@ def forward(self, to_cache: torch.Tensor) -> torch.Tensor: Returns: Dummy output tensor (not used) + """ # Note: we set num_heads to num_hidden_layers and # head_size to hidden_size for hidden states storage @@ -373,8 +373,8 @@ def forward(self, hidden_states: torch.Tensor) -> None: Returns: Tuple of (dummy_output, dummy_output) - both unused - """ + """ # Call dummy attention layer to cache hidden states # Output is ignored - we only care about the KV cache side effects _ = self.cache_only_layers[str(self.target_num_hidden_layers)](hidden_states) diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 5b9853ece10b..f18f938774dc 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -140,8 +140,7 @@ def __init__( self._init_mup_vector() def _init_mup_vector(self): - """ - Non learnable per-block scaling vector composed of element-wise + """Non learnable per-block scaling vector composed of element-wise multipliersapplied to each separate contiguous block of the output of the linear projection (in_proj) before further processing (gating, convolution, SSM): @@ -313,8 +312,7 @@ def forward( class FalconH1ParallelHybrid(nn.Module): - """ - A hybrid decoder layer for FalconH1 where the input is processed + """A hybrid decoder layer for FalconH1 where the input is processed in parallel through both the self-attention branch and the SSM (Mamba) branch. Their outputs are then summed to produce the final hidden state. @@ -552,6 +550,7 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index b9d99474597a..847ec89e8f05 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -63,11 +63,10 @@ class FireRedASR2AudioInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - nmb: Number of mel bins - - t: Time frames (M) + """Dimensions: + - b: Batch size + - nmb: Number of mel bins + - t: Time frames (M) """ input_features: Annotated[ diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index b7981f42d364..9a7dc49d7195 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -598,11 +598,10 @@ def forward(self, hidden_states: torch.Tensor, ilens: torch.Tensor): class FunASRAudioInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - nmb: Number of mel bins - - t: Time frames (M) + """Dimensions: + - b: Batch size + - nmb: Number of mel bins + - t: Time frames (M) """ input_features: Annotated[ diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index abe1f5f1bd55..d6aaabe0232f 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -58,14 +58,13 @@ class Gemma3ImagePixelInputs(TensorSchema): - """ - Dimensions: - - p: Number of patches total (over each image over each prompt in the - batch) - - c: Number of channels (3) - - h: Height of each patch - - w: Width of each patch - - bn: Batch size * number of images + """Dimensions: + - p: Number of patches total (over each image over each prompt in the + batch) + - c: Number of channels (3) + - h: Height of each patch + - w: Width of each patch + - bn: Batch size * number of images """ type: Literal["pixel_values"] = "pixel_values" @@ -678,9 +677,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index 56d64735789a..87d72c18c10a 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -185,7 +185,7 @@ def correct( class Gemma3nLaurelBlock(nn.Module): - """Learned Augmented Residual Layer""" + """Learned Augmented Residual Layer.""" def __init__( self, @@ -593,9 +593,7 @@ def forward( enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill ) class Gemma3nSelfDecoder(nn.Module): - """ - Includes altup embedding and self decoder layers - """ + """Includes altup embedding and self decoder layers.""" def __init__( self, @@ -776,9 +774,7 @@ def forward( enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill ) class Gemma3nCrossDecoder(nn.Module): - """ - Cross-decoder layers - """ + """Cross-decoder layers.""" def __init__( self, @@ -818,7 +814,8 @@ def forward( def _kv_sharing_weights_mapper(config: Gemma3nTextConfig) -> WeightsMapper: """KV-shared layers only have q_proj, so qkv_proj packing applies to the other layers. Original checkpoints still ship K/V tensors for the shared - layers (fine-tuned ones omit them); those are dropped.""" + layers (fine-tuned ones omit them); those are dropped. + """ first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers return WeightsMapper( orig_to_new_substr={ diff --git a/vllm/model_executor/models/gemma3n_audio_utils.py b/vllm/model_executor/models/gemma3n_audio_utils.py index bef9bb9a0d4e..427c98e43a17 100644 --- a/vllm/model_executor/models/gemma3n_audio_utils.py +++ b/vllm/model_executor/models/gemma3n_audio_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Lightweight utility functions for Gemma3n audio processing. +"""Lightweight utility functions for Gemma3n audio processing. This module is separate from gemma3n_mm.py to avoid heavy CUDA dependencies, making it testable without a full vLLM build. @@ -37,6 +36,7 @@ def adjust_audio_features_to_expected_length( Tuple of: - adjusted_features: Audio features adjusted to expected_tokens length - tokens_truncated: Number of tokens truncated (0 if padding was applied) + """ audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape tokens_truncated = 0 diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 1c05c165dcfe..769c4c666a5c 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -74,12 +74,11 @@ class Gemma3nImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each patch - - w: Width of each patch + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each patch + - w: Width of each patch """ type: Literal["pixel_values"] = "pixel_values" @@ -113,6 +112,7 @@ def batch_audio_features( Returns: The `(bn, s_max, f)` features and their `(bn, s_max)` mask. + """ if isinstance(input_features, torch.Tensor): assert isinstance(input_features_mask, torch.Tensor) @@ -135,11 +135,10 @@ def batch_audio_features( class Gemma3nAudioInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of audios - - s: seq_length - - f: num_features + """Dimensions: + - bn: Batch size * number of audios + - s: seq_length + - f: num_features """ type: Literal["audio"] = "audio" @@ -187,8 +186,7 @@ def get_image_repl( image_height: int, processor: Gemma3nProcessor, ) -> PromptUpdateDetails: - """ - Get the replacement metadata for image tokens. + """Get the replacement metadata for image tokens. For Gemma3n, this should return the full_image_sequence which includes BOI token, repeated image tokens, and EOI token. @@ -205,8 +203,7 @@ def get_audio_repl( *, processor: Gemma3nProcessor, ) -> PromptUpdateDetails: - """ - Get the replacement metadata for audio tokens. + """Get the replacement metadata for audio tokens. For Gemma3n, this should return the full_audio_sequence which includes BOA token, repeated audio tokens, and EOA token. @@ -316,8 +313,7 @@ def _postprocess_hf_mm_data( return processed_data def _get_num_audio_frames_alone(self, num_samples: int) -> int: - """ - Get the number of mel frames that the feature extractor produces + """Get the number of mel frames that the feature extractor produces for a single audio clip of the given length, independent of batch padding. """ @@ -514,7 +510,8 @@ def get_repl_toks(tok: int) -> list[int]: class Gemma3nMultimodalEmbedder(nn.Module): """Embeds token ids or soft tokens for multimodal content into language - model space.""" + model space. + """ def __init__( self, @@ -571,6 +568,7 @@ def forward( Returns: A torch.Tensor of embeddings with shape `[batch_size, seq_len, self.config.text_config.hidden_size]`. + """ # noqa: E501 if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( @@ -891,9 +889,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", @@ -911,8 +907,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None: @classmethod def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: - """ - Gemma3n supports "free-form" transcription. + """Gemma3n supports "free-form" transcription. We fix its prompt here to standardize transcriptions/translations requests. """ diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 961709a7e799..04e22523c2f2 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -866,6 +866,7 @@ def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: Returns: Per-layer embeddings (num_tokens, num_layers, hidden_size_per_layer_input) + """ if self.embed_tokens_per_layer is None: return None @@ -1199,6 +1200,7 @@ def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: Returns: Per-layer embeddings (num_tokens, num_layers, hidden_size_per_layer_input) + """ return self.self_decoder.get_per_layer_inputs(input_ids) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 2fe2752475a0..752e2243df63 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -132,8 +132,7 @@ def _get_max_soft_tokens( class Gemma4ImagePixelInputs(TensorSchema): - """ - Pre-patchified image inputs from the Gemma4 image processor. + """Pre-patchified image inputs from the Gemma4 image processor. Dimensions: - bn: Batch size * number of images @@ -159,11 +158,10 @@ class Gemma4ImagePixelInputs(TensorSchema): class Gemma4AudioInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of audios - - s: Sequence length (MEL spectrogram frames) - - f: Number of features (MEL bins) + """Dimensions: + - bn: Batch size * number of audios + - s: Sequence length (MEL spectrogram frames) + - f: Number of features (MEL bins) """ type: Literal["audio"] = "audio" @@ -304,9 +302,12 @@ def _compute_num_soft_tokens( for an image of the given dimensions, after padding is stripped. Args: + image_width: Width of the image in pixels. + image_height: Height of the image in pixels. max_soft_tokens: Override for the vision config's ``default_output_length``. When *None*, the value from the model config is used. + """ vision_cfg = self.get_hf_config().vision_config patch_size = vision_cfg.patch_size @@ -345,8 +346,12 @@ def get_image_repl( produce after stripping padding. Args: + image_width: Width of the image in pixels. + image_height: Height of the image in pixels. + processor: The Gemma4 processor, when already constructed. max_soft_tokens: Override for the default token budget. When *None*, falls back to the model config value. + """ if processor is None: processor = self.get_hf_processor() diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 639af0ea2951..01f226ff07a5 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -101,6 +101,7 @@ def _select_and_score( Returns: logits: (num_tokens, num_selected) sparse logits. indices: (num_tokens, num_selected) corresponding vocab indices. + """ num_tokens = hidden_states.shape[0] _, top_k_indices = torch.topk( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 730a015cfe53..87baaa10693a 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -25,7 +25,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.1V & GLM-4.6V-Flash, AutoGLM-Phone-9B model -compatible with HuggingFace weights.""" +compatible with HuggingFace weights. +""" import math from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence @@ -145,12 +146,11 @@ def _to_video_metadata(metadata: Mapping[str, Any]) -> VideoMetadata: class Glm4vImagePixelInputs(TensorSchema): - """ - Dimensions: - - np: Number of patches - - cpp: Number of channels * patch_size * patch_size - - ni: Number of images - - g: Grid dimensions (3 for grid_t, grid_h, grid_w) + """Dimensions: + - np: Number of patches + - cpp: Number of channels * patch_size * patch_size + - ni: Number of images + - g: Grid dimensions (3 for grid_t, grid_h, grid_w) """ type: Literal["pixel_values"] = "pixel_values" @@ -160,12 +160,11 @@ class Glm4vImagePixelInputs(TensorSchema): class Glm4vImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - f: Number of image features (varies based on image resolution) - - h: Hidden size (must match language model backbone) - - n: Number of images - - g: Grid dimensions (3 for grid_t, grid_h, grid_w) + """Dimensions: + - f: Number of image features (varies based on image resolution) + - h: Hidden size (must match language model backbone) + - n: Number of images + - g: Grid dimensions (3 for grid_t, grid_h, grid_w) """ type: Literal["image_embeds"] = "image_embeds" @@ -178,14 +177,13 @@ class Glm4vImageEmbeddingInputs(TensorSchema): class Glm4vVideoPixelInputs(TensorSchema): - """ - Dimensions: - - np: Number of patches - - ctpp: Number of channels * temporal_patch_size * - patch_size * patch_size - - f: Number of frames - - g: Grid dimensions (3 for grid_t which is usually 1 for processed - video, grid_h, grid_w) + """Dimensions: + - np: Number of patches + - ctpp: Number of channels * temporal_patch_size * + patch_size * patch_size + - f: Number of frames + - g: Grid dimensions (3 for grid_t which is usually 1 for processed + video, grid_h, grid_w) """ type: Literal["pixel_values_videos"] = "pixel_values_videos" @@ -195,13 +193,12 @@ class Glm4vVideoPixelInputs(TensorSchema): class Glm4vVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - p: Number of video patches across all frames - - h: Hidden size (must match language model backbone) - - f: Number of frames - - g: Grid dimensions (3 for grid_t which is usually 1 for processed - video, grid_h, grid_w) + """Dimensions: + - p: Number of video patches across all frames + - h: Hidden size (must match language model backbone) + - f: Number of frames + - g: Grid dimensions (3 for grid_t which is usually 1 for processed + video, grid_h, grid_w) """ type: Literal["video_embeds"] = "video_embeds" @@ -771,7 +768,6 @@ def pos_embeds_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor: forward graph compatible with CUDA graph replay. Coordinate generation matches `rot_pos_emb` exactly to guarantee spatial alignment. """ - device = self.embeddings.position_embedding.weight.device dtype = self.dtype all_embeds = [] @@ -850,6 +846,7 @@ def prepare_encoder_metadata( instead of computing from cu_seqlens (needed for CUDA graph capture to cover worst-case replay scenarios). device: Device to place tensors on. Defaults to self.device. + """ if device is None: device = self.device @@ -2347,6 +2344,7 @@ def forward( parallelism. inputs_embeds: Optional pre-computed input embeddings. **kwargs: Additional keyword arguments. + """ if intermediate_tensors is not None: inputs_embeds = None @@ -2370,9 +2368,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model.model", connector="visual.merger.", diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 3c72a148d5ef..18256e9c13bf 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -22,7 +22,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.5, GLM-4.6, GLM-4.7 model -compatible with HuggingFace weights.""" +compatible with HuggingFace weights. +""" from collections.abc import Iterable from itertools import islice diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index 2ef9407fb750..af5a1e113fd8 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -454,8 +454,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index bad4ec8cf333..12592af4a734 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -22,7 +22,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP -model compatible with HuggingFace weights.""" +model compatible with HuggingFace weights. +""" import typing from collections.abc import Callable, Iterable @@ -412,8 +413,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index b00e6873612a..4fc6446ad3c1 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -65,12 +65,11 @@ class GLMVImagePixelInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - c: Number of channels (3) - - h: Height of image - - w: Width of image + """Dimensions: + - b: Batch size + - c: Number of channels (3) + - h: Height of image + - w: Width of image """ type: Literal["pixel_values"] = "pixel_values" @@ -90,14 +89,15 @@ def __init__(self, config): self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size) def forward(self, images: torch.Tensor) -> torch.Tensor: - """ - Parameters: + """Parameters images : torch.Tensor Input image tensor with shape (B, C, H, W) - Returns: + Returns + ------- torch.Tensor Transformed tensor with shape (B, L, D) + """ images = images.to(device=self.proj.weight.device, dtype=self.proj.weight.dtype) x = self.proj(images) @@ -246,8 +246,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", ): - """ - The original implementation is the same as: + """The original implementation is the same as: ```python self.dense_h_to_4h = ColumnParallelLinear( config.hidden_size, @@ -349,14 +348,15 @@ def __init__( self.scaling_factor = vision_config.scaling_factor def forward(self, images: torch.Tensor) -> torch.Tensor: - """ - Parameters: + """Parameters images : torch.Tensor Input image tensor with shape (B, C, H, W) - Returns: + Returns + ------- torch.Tensor Transformed tensor with shape (B, L, D) + """ x = self.patch_embedding(images) x = self.transformer(x) @@ -518,9 +518,7 @@ class GLM4VForCausalLM( } def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="transformer.encoder", connector="transformer.vision.linear_proj", diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py index 05c3ac793578..6bcd33f3dc46 100644 --- a/vllm/model_executor/models/glm_ocr_mtp.py +++ b/vllm/model_executor/models/glm_ocr_mtp.py @@ -246,8 +246,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 219a7ce02de9..710c6378a8cc 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -75,8 +75,7 @@ class GlmAsrEncoderRotaryEmbedding(nn.Module): - """ - Rotary Position Embedding for GLM-ASR encoder. + """Rotary Position Embedding for GLM-ASR encoder. Computes rotary position embeddings on-demand for efficiency. Only caches inv_freq as a buffer; cos/sin are computed during forward @@ -115,8 +114,7 @@ def __init__(self, config) -> None: self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, seq_len: int) -> torch.Tensor: - """ - Compute rotary position frequencies for given sequence length. + """Compute rotary position frequencies for given sequence length. Args: seq_len: The sequence length to compute embeddings for. @@ -124,6 +122,7 @@ def forward(self, seq_len: int) -> torch.Tensor: Returns: Frequency tensor with shape [seq_len, dim/2]. Use .cos() and .sin() to get the rotary embedding components. + """ # Compute on the same device as inv_freq (automatically correct after .to()) seq = torch.arange( @@ -134,8 +133,7 @@ def forward(self, seq_len: int) -> torch.Tensor: class GlmAsrEncoderAttention(nn.Module): - """ - Optimized Multi-headed Grouped Query Attention for GLM-ASR encoder. + """Optimized Multi-headed Grouped Query Attention for GLM-ASR encoder. Uses vLLM's QKVParallelLinear for fused projections, ApplyRotaryEmb for rotary position embeddings, and MMEncoderAttention for hardware-optimized @@ -208,14 +206,14 @@ def forward( rotary_pos_emb_cos: torch.Tensor, rotary_pos_emb_sin: torch.Tensor, ) -> torch.Tensor: - """ - Args: + """Args: hidden_states: [batch_size, seq_len, hidden_size] rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings Returns: [batch_size, seq_len, hidden_size] + """ batch_size, seq_len, _ = hidden_states.shape @@ -255,8 +253,7 @@ def forward( class GlmAsrEncoderMLP(nn.Module): - """ - Optimized MLP for GLM-ASR encoder. + """Optimized MLP for GLM-ASR encoder. Uses vLLM's parallel linear layers for better performance. """ @@ -297,8 +294,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class GlmAsrEncoderLayer(nn.Module): - """ - Optimized Transformer encoder layer for GLM-ASR. + """Optimized Transformer encoder layer for GLM-ASR. Combines attention and MLP with residual connections and layer norms. """ @@ -335,14 +331,14 @@ def forward( rotary_pos_emb_cos: torch.Tensor, rotary_pos_emb_sin: torch.Tensor, ) -> torch.Tensor: - """ - Args: + """Args: hidden_states: [batch_size, seq_len, hidden_size] rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings Returns: [batch_size, seq_len, hidden_size] + """ # Self-attention with residual residual = hidden_states @@ -364,8 +360,7 @@ def forward( class _GlmAsrEncoderOutput: - """ - Simple output container compatible with transformers' BaseModelOutput. + """Simple output container compatible with transformers' BaseModelOutput. This lightweight container holds the encoder output and is compatible with the transformers library's output format while being more efficient @@ -374,6 +369,7 @@ class _GlmAsrEncoderOutput: Attributes: last_hidden_state: Final layer hidden states from the encoder. Shape: [batch_size, seq_len, hidden_size] + """ __slots__ = ("last_hidden_state",) @@ -383,8 +379,7 @@ def __init__(self, last_hidden_state: torch.Tensor): class GlmAsrEncoder(nn.Module): - """ - Optimized GLM-ASR Audio Encoder with vLLM native implementation. + """Optimized GLM-ASR Audio Encoder with vLLM native implementation. This encoder processes audio features through convolutional layers followed by transformer layers with rotary position embeddings. @@ -454,14 +449,14 @@ def __init__( def _get_feat_extract_output_lengths( self, input_lengths: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Compute the output length after convolutions. + """Compute the output length after convolutions. Args: input_lengths: Input sequence lengths [batch_size] Returns: Tuple of (output after conv1, output after conv2) + """ # Conv1: kernel=3, stride=1, padding=1 output_lengths_conv1 = (input_lengths + 2 * 1 - 3) // 1 + 1 @@ -472,8 +467,7 @@ def _get_feat_extract_output_lengths( return output_lengths_conv1, output_lengths_conv2 def forward(self, input_features: torch.Tensor) -> _GlmAsrEncoderOutput: - """ - Forward pass through the encoder. + """Forward pass through the encoder. Args: input_features: [batch_size, num_mel_bins, seq_len] @@ -482,6 +476,7 @@ def forward(self, input_features: torch.Tensor) -> _GlmAsrEncoderOutput: _GlmAsrEncoderOutput: Object with .last_hidden_state attribute \ containing [batch_size, seq_len', hidden_size] where seq_len' \ is the sequence length after convolutions + """ # Apply convolutional layers with GELU activation hidden_states = torch.nn.functional.gelu(self.conv1(input_features)) @@ -515,11 +510,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class GlmAsrFeatureInputs(TensorSchema): - """ - Dimensions: - - num_chunks: Number of audio chunks (flattened) - - nmb: Number of mel bins - - num_audios: Number of original audio files + """Dimensions: + - num_chunks: Number of audio chunks (flattened) + - nmb: Number of mel bins + - num_audios: Number of original audio files """ type: Literal["audio_features"] @@ -538,12 +532,11 @@ class GlmAsrFeatureInputs(TensorSchema): class GlmAsrEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size - - naf: Number of audio features - - hs: Hidden size (must match the hidden size of language model - backbone) + """Dimensions: + - bn: Batch size + - naf: Number of audio features + - hs: Hidden size (must match the hidden size of language model + backbone) """ type: Literal["audio_embeds"] = "audio_embeds" @@ -557,8 +550,7 @@ class GlmAsrEmbeddingInputs(TensorSchema): class GlmAsrMultiModalProjector(nn.Module): - """ - Projects audio encoder outputs to language model hidden space. + """Projects audio encoder outputs to language model hidden space. This projector uses a two-layer MLP to map audio features from the encoder's intermediate size to the language model's hidden size. @@ -601,8 +593,7 @@ def forward(self, audio_features: torch.Tensor) -> torch.Tensor: def _glmasr_field_config( hf_inputs: Mapping[str, torch.Tensor], ) -> dict[str, MultiModalFieldConfig]: - """ - Configure multimodal field batching strategy for GLM-ASR. + """Configure multimodal field batching strategy for GLM-ASR. Determines how to batch audio inputs based on whether chunking is used. When chunk_counts is present, features are flattened across chunks; @@ -614,6 +605,7 @@ def _glmasr_field_config( Returns: Dictionary mapping field names to MultiModalFieldConfig objects \ that specify batching behavior. + """ chunk_counts = hf_inputs.get("chunk_counts") if chunk_counts is not None: @@ -636,8 +628,7 @@ def _glmasr_field_config( class GlmAsrMultiModalDataParser(MultiModalDataParser): - """ - Custom parser for GLM-ASR multimodal data. + """Custom parser for GLM-ASR multimodal data. Extends the base parser to handle GLM-ASR specific audio data formats, including both pre-computed audio embeddings and raw audio features. @@ -658,8 +649,7 @@ def _parse_audio_data( class GlmAsrProcessingInfo(BaseProcessingInfo): - """ - Processing information provider for GLM-ASR model. + """Processing information provider for GLM-ASR model. Provides access to model configuration, processor, and feature extractor needed for audio preprocessing and multimodal integration. @@ -687,8 +677,7 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]: class GlmAsrDummyInputsBuilder(BaseDummyInputsBuilder[GlmAsrProcessingInfo]): - """ - Builder for dummy inputs used in profiling and testing. + """Builder for dummy inputs used in profiling and testing. Generates dummy text prompts and audio data that match the expected format for GLM-ASR model inputs. Used for memory profiling and @@ -727,8 +716,7 @@ def get_dummy_mm_data( class GlmAsrMultiModalProcessor(BaseMultiModalProcessor["GlmAsrProcessingInfo"]): - """ - GLM-ASR processor that inherits directly from BaseMultiModalProcessor + """GLM-ASR processor that inherits directly from BaseMultiModalProcessor for better performance and cleaner implementation. """ diff --git a/vllm/model_executor/models/glmasr_utils.py b/vllm/model_executor/models/glmasr_utils.py index 8dcfcfa89513..7997963ddab8 100644 --- a/vllm/model_executor/models/glmasr_utils.py +++ b/vllm/model_executor/models/glmasr_utils.py @@ -71,8 +71,7 @@ def _get_audio_output_lengths_for_tower( merge_factor: int, conv_params: list[tuple[int, int, int]], ) -> torch.Tensor: - """ - Calculate the output lengths after audio processing. + """Calculate the output lengths after audio processing. The output length accounts for: 1. Convolution layers (downsampling) @@ -86,6 +85,7 @@ def _get_audio_output_lengths_for_tower( Returns: Output lengths after all processing [batch_size] + """ # First, calculate the output length after convolutions if hasattr(audio_tower, "_get_feat_extract_output_lengths"): diff --git a/vllm/model_executor/models/gpt2.py b/vllm/model_executor/models/gpt2.py index 07801cd098b1..c95fe5184898 100644 --- a/vllm/model_executor/models/gpt2.py +++ b/vllm/model_executor/models/gpt2.py @@ -327,6 +327,7 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding): transformer: An instance of GPT2Model used for forward operations. score: A layer for calculating logits. _pooler: An instance of Pooler used for pooling operations. + """ is_pooling_model = True diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 4887007564fc..bd4b38defc53 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -790,6 +790,7 @@ def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: Returns: Weight dtype string (e.g., "mxfp4", "fp8") or None if not available + """ if hasattr(self.layers[layer_id].mlp.experts._quant_method, "weight_dtype"): return self.layers[layer_id].mlp.experts._quant_method.weight_dtype diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 55a90f5ff04e..73d9184f1264 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -84,6 +84,7 @@ def granite_layer_attn_params( Returns: Sliding window size (`None` for full attention), RoPE base theta (`0` for NoPE), and attention sink presence/absence. + """ layer_types = getattr(config, "layer_types", None) sliding_window = ( diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 50b0932bda8e..b62a86543414 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -684,6 +684,7 @@ def _get_all_layer_features( Packing on dim=-1 means the framework's token-level slicing for chunked prefill preserves all levels intact. + """ select_strategy = self._vision_feature_select_strategy diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 0fac643999ff..c8d1f7111bf0 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -91,8 +91,7 @@ ### Audio Input class GraniteSpeechAudioInputs(TensorSchema): - """ - Audio input features for Granite Speech model. + """Audio input features for Granite Speech model. Dimensions: - b: Batch size @@ -477,7 +476,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class GraniteSpeechConformerBlock(nn.Module): """Conformer block, consisting largely of linear layers, - attention, and convolutional layers.""" + attention, and convolutional layers. + """ def __init__(self, config: PretrainedConfig, prefix: str = ""): super().__init__() @@ -732,9 +732,11 @@ def _build_input_features_mask( Args: audio_embed_sizes: torch.Tensor Tensor of num features in each seq in the batch. + Returns: torch.Tensor: Mask of shape (bsz, num_features) to be applied to the audio features prior to splitting the audio embeddings. + """ most_audio_features = int(torch.max(audio_embed_sizes)) mask_indices = torch.arange(most_audio_features).view(1, -1) @@ -763,10 +765,12 @@ def _pad_and_stack_input_features( Args: input_features: list[torch.Tensor] 3D Input features to be coerced into a tensor. + Returns: torch.Tensor: Tensor of shape [bsz, num_features, 160], where num_features is the max number of features of any entry in the batch. + """ feat_lens = [feats.shape[1] for feats in input_features] padding = [max(feat_lens) - length for length in feat_lens] @@ -790,8 +794,10 @@ def _process_audio_input( audio_input: GraniteSpeechAudioInputs Audio inputs object containing Mel features, an input features mask, and the (flattened) number of audio tokens per instance. + Returns: tuple[torch.Tensor]: List of length bsz. + """ # TODO (Alex) - support embedding inputs encoder_embeds = self.encoder(audio_input["input_features"]) diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 8dc47c522654..da7cd0478b55 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -633,6 +633,7 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py index db8d9c20c39d..e5140bc03594 100644 --- a/vllm/model_executor/models/hrm_text.py +++ b/vllm/model_executor/models/hrm_text.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -HRM-Text: Hierarchical Reasoning Model — Text variant. +"""HRM-Text: Hierarchical Reasoning Model — Text variant. Reference Hugging Face implementation: src/transformers/models/hrm_text/modeling_hrm_text.py diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py index 96e718751d59..10c066cb7176 100644 --- a/vllm/model_executor/models/hyperclovax_vision_v2.py +++ b/vllm/model_executor/models/hyperclovax_vision_v2.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -HyperCLOVAX V2 (32B Think Model) Implementation. +"""HyperCLOVAX V2 (32B Think Model) Implementation. This module contains the V2 architecture that uses Qwen2.5 Vision Transformer instead of CLIP/SigLIP used in V1. @@ -57,8 +56,7 @@ class HCXVisionV2ImagePixelInputs(TensorSchema): - """ - V2 Image inputs using Qwen2.5-VL style grid_thw format. + """V2 Image inputs using Qwen2.5-VL style grid_thw format. Dimensions: - np: Number of patches @@ -72,8 +70,7 @@ class HCXVisionV2ImagePixelInputs(TensorSchema): class HCXVisionV2ImageEmbeddingInputs(TensorSchema): - """ - V2 Image embedding inputs. + """V2 Image embedding inputs. Dimensions: - nf: Number of image features @@ -90,8 +87,7 @@ class HCXVisionV2ImageEmbeddingInputs(TensorSchema): class HCXVisionV2VideoPixelInputs(TensorSchema): - """ - V2 Video inputs using Qwen2.5-VL style grid_thw format. + """V2 Video inputs using Qwen2.5-VL style grid_thw format. Dimensions: - np: Number of patches @@ -105,8 +101,7 @@ class HCXVisionV2VideoPixelInputs(TensorSchema): class HCXVisionV2VideoEmbeddingInputs(TensorSchema): - """ - V2 Video embedding inputs. + """V2 Video embedding inputs. Dimensions: - nf: Number of video features @@ -419,8 +414,7 @@ def _get_mm_fields_config( dummy_inputs=HCXVisionV2DummyInputsBuilder, ) class HCXVisionV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): - """ - HyperCLOVAX-SEED Vision-Language Model (V2 architecture). + """HyperCLOVAX-SEED Vision-Language Model (V2 architecture). Supports: - HyperCLOVAX-SEED-Think-32B: Vision + Text diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index 3021e66f759f..09809db7022e 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -45,8 +45,7 @@ class Idefics2VisionEmbeddings(nn.Module): - """ - This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings + """This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings ` to enable images of variable resolution. @@ -153,7 +152,7 @@ def forward( class Idefics2VisionAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -308,10 +307,9 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Args: - hidden_states (`torch.FloatTensor`): - Input to the layer of shape `(batch, seq_len, embed_dim)`. + """Args: + hidden_states (`torch.FloatTensor`): + Input to the layer of shape `(batch, seq_len, embed_dim)`. """ residual = hidden_states @@ -326,13 +324,13 @@ def forward( class Idefics2Encoder(nn.Module): - """ - Transformer encoder consisting of `config.num_hidden_layers` self attention + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Idefics2EncoderLayer`]. Args: config: Idefics2Config + """ def __init__( @@ -368,14 +366,14 @@ def forward( inputs_embeds: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: - r""" - Args: - inputs_embeds (torch.Tensor): - Optionally, instead of passing `input_ids` you can choose to - directly pass an embedded representation. - This is useful if you want more control over how to convert - `input_ids` indices into associated vectorsthan the model's - internal embedding lookup matrix. + r"""Args: + inputs_embeds (torch.Tensor): + Optionally, instead of passing `input_ids` you can choose to + directly pass an embedded representation. + This is useful if you want more control over how to convert + `input_ids` indices into associated vectorsthan the model's + internal embedding lookup matrix. + """ hidden_states = inputs_embeds for encoder_layer in self.layers: diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 4ac8eb5eeb43..014c4ae6b54f 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -70,13 +70,12 @@ class Idefics3ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - bnp: Batch size * number of images * number of patches - - c: Number of channels (3) - - h: Height - - w: Width + """Dimensions: + - bn: Batch size * number of images + - bnp: Batch size * number of images * number of patches + - c: Number of channels (3) + - h: Height + - w: Width """ type: Literal["pixel_values"] @@ -86,11 +85,10 @@ class Idefics3ImagePixelInputs(TensorSchema): class Idefics3ImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - f: Image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - bn: Batch size * number of images + - f: Image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] @@ -829,9 +827,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="model.text_model", connector="model.connector", diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 2e79001add82..e0121f95338d 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -112,8 +112,7 @@ def process_delta(self, text_delta: str, finished: bool) -> str: def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor: - """ - A helper function to be used in the context of + """A helper function to be used in the context of [vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids][] to provide a better error message. """ @@ -225,20 +224,18 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: - """ - Get the placeholder text for the `i`th `modality` item in the prompt. - """ + """Get the placeholder text for the `i`th `modality` item in the prompt.""" ... def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - """ - Returns multimodal embeddings generated from multimodal kwargs + """Returns multimodal embeddings generated from multimodal kwargs to be merged with text embeddings. Note: The returned multimodal embeddings must be in the same order as the appearances of their corresponding multimodal data item in the input prompt. + """ ... @@ -254,14 +251,14 @@ def configure_mm_token_handling(self, vocab_size: int, mm_token_ids: list[int]): ) def get_language_model(self) -> "VllmModel": - """ - Returns the underlying language model used for text generation. + """Returns the underlying language model used for text generation. This is typically the `torch.nn.Module` instance responsible for processing the merged multimodal embeddings and producing hidden states Returns: torch.nn.Module: The core language model component. + """ # Cached assert isinstance(self, nn.Module) @@ -302,8 +299,7 @@ def _mark_language_model( *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, ): - """ - Mark each child module that was assigned to this model during this context + """Mark each child module that was assigned to this model during this context as a language model component. Language model components are automatically skipped in `--mm-encoder-only` @@ -338,8 +334,7 @@ def _mark_tower_model( *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, ): - """ - Mark each child module that was assigned to this model during this context + """Mark each child module that was assigned to this model during this context as a tower model component. Tower model components are automatically skipped when `--limit-mm-per-prompt` @@ -397,8 +392,7 @@ def _mark_composite_model( language_targets: type[nn.Module] | tuple[type[nn.Module], ...], tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]], ): - """ - Composite wrapper over `_mark_language_model` and + """Composite wrapper over `_mark_language_model` and `_mark_tower_model` by modality. """ with ExitStack() as stack: @@ -438,8 +432,7 @@ def get_mm_lora_token_counts( mm_kwargs: "MultiModalKwargsItem | None", num_mm_embeds: int, ) -> tuple[int, int | None]: - """ - Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings. + """Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings. MM LoRA uses these counts to build adapter mappings for the tower and connector forwards. Models with multiple modalities can override this @@ -497,8 +490,7 @@ def embed_input_ids( *, is_multimodal: Tensor | None = None, ) -> Tensor: - """ - Apply token embeddings to `input_ids`. + """Apply token embeddings to `input_ids`. If `multimodal_embeddings` is passed, scatter them into `input_ids` according to the mask `is_multimodal`. @@ -551,8 +543,7 @@ def recompute_mrope_positions( mrope_positions: torch.LongTensor, num_computed_tokens: int, ) -> tuple[Sequence[torch.Tensor], Tensor, int]: - """ - Update part of input mrope positions (starting with + """Update part of input mrope positions (starting with num_computed_tokens index). Original mrope_positions are computed for unpruned sequence and becomes incorrect once pruning occurs, so once we prune media tokens we should reflect this in the @@ -571,6 +562,7 @@ def recompute_mrope_positions( Returns: Tuple of (multimodal_embeddings, mrope_positions, mrope_position_delta). + """ ... @@ -650,16 +642,12 @@ class SupportsScoreTemplate(Protocol): @classmethod def get_score_template(cls, query: str, document: str) -> str | None: - """ - Generate a full prompt by populating the score template with query and document content. - """ # noqa: E501 + """Generate a full prompt by populating the score template with query and document content.""" # noqa: E501 ... @classmethod def post_process_tokens(cls, prompt: "TokensPrompt") -> None: - """ - Perform architecture-specific manipulations on the input tokens. - """ + """Perform architecture-specific manipulations on the input tokens.""" ... @@ -790,8 +778,7 @@ def forward( *, intermediate_tensors: "IntermediateTensors | None", ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": - """ - Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when + """Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when PP rank > 0. Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only @@ -905,7 +892,8 @@ def has_inner_state( @runtime_checkable class IsAttentionFree(Protocol): """The interface required for all models like Mamba that lack attention, - but do have state whose size is constant wrt the number of tokens.""" + but do have state whose size is constant wrt the number of tokens. + """ is_attention_free: ClassVar[Literal[True]] = True """ @@ -933,7 +921,8 @@ def is_attention_free( class IsHybrid(Protocol): """The interface required for all models like Jamba that have both attention and mamba blocks, indicates that - hf_config has 'layers_block_type'""" + hf_config has 'layers_block_type' + """ is_hybrid: ClassVar[Literal[True]] = True """ @@ -953,6 +942,7 @@ def get_mamba_state_shape_from_config( Returns: Shapes for each state cache used by the model. + """ ... @@ -966,6 +956,7 @@ def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: (state, block_ids, cur_block_idx, num_accepted_tokens) and returns a MambaCopySpec describing the memory-copy parameters for prefix caching in align mode. + """ ... @@ -994,9 +985,7 @@ def is_hybrid( @runtime_checkable class MixtureOfExperts(Protocol): - """ - Check if the model is a mixture of experts (MoE) model. - """ + """Check if the model is a mixture of experts (MoE) model.""" expert_weights: MutableSequence[Sequence[Tensor]] """ @@ -1039,8 +1028,7 @@ def set_eplb_state( logical_to_physical_map: Tensor, logical_replica_count: Tensor, ) -> None: - """ - Register the EPLB state in the MoE model. + """Register the EPLB state in the MoE model. Since these are views of the actual EPLB state, any changes made by the EPLB algorithm are automatically reflected in the model's behavior @@ -1054,6 +1042,7 @@ def set_eplb_state( expert_load_view: A view of the expert load metrics tensor. logical_to_physical_map: Mapping from logical to physical experts. logical_replica_count: Count of replicas for each logical expert. + """ self.expert_weights = [] for layer_idx, layer in enumerate(self.moe_layers): @@ -1086,8 +1075,8 @@ def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None: Returns: The MixtureOfExperts instance contained within the model, or None. - """ + """ if is_mixture_of_experts(model): return model @@ -1349,7 +1338,8 @@ def get_generation_prompt( ) -> "PromptType": """Get the prompt for the ASR model. The model has control over the construction, as long as it - returns a valid PromptType.""" + returns a valid PromptType. + """ ... @classmethod @@ -1359,8 +1349,7 @@ def get_other_languages(cls) -> Mapping[str, str]: @classmethod def validate_language(cls, language: str | None) -> str | None: - """ - Ensure the language specified in the transcription request + """Ensure the language specified in the transcription request is a valid ISO 639-1 language code. If the request language is valid, but not natively supported by the model, trigger a warning (but not an exception). @@ -1396,8 +1385,7 @@ def get_num_audio_tokens( stt_config: "SpeechToTextConfig", model_config: "ModelConfig", ) -> int | None: - """ - Map from audio duration to number of audio tokens produced by the ASR + """Map from audio duration to number of audio tokens produced by the ASR model, without running a forward pass. This is used for estimating the amount of processing for this audio. """ @@ -1405,8 +1393,7 @@ def get_num_audio_tokens( @classmethod def post_process_output(cls, text: str) -> str: - """ - Post-process the raw model output text. + """Post-process the raw model output text. Some ASR models output structured formats (e.g., language tags, special tokens) that need to be stripped before returning to the user. @@ -1416,6 +1403,7 @@ def post_process_output(cls, text: str) -> str: Returns: Cleaned transcription text. + """ return text @@ -1432,8 +1420,7 @@ def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegme def get_streaming_post_processor_cls( cls, ) -> type[StreamingTranscriptionPostProcessor]: - """ - Return a stateful post-processor class for streaming output deltas. + """Return a stateful post-processor class for streaming output deltas. Each instance receives the next decoded text delta and whether the request output is final. It returns the cleaned delta that should be @@ -1644,7 +1631,8 @@ def collect_remote_aux_hidden_states( @runtime_checkable class SupportsEagle(SupportsEagleBase, Protocol): """The interface required for models that support - EAGLE-1 and EAGLE-2 speculative decoding.""" + EAGLE-1 and EAGLE-2 speculative decoding. + """ supports_eagle: ClassVar[Literal[True]] = True """ @@ -1674,7 +1662,8 @@ def supports_eagle( @runtime_checkable class SupportsEagle3(SupportsEagleBase, Protocol): """The interface required for models that support - EAGLE-3 speculative decoding.""" + EAGLE-3 speculative decoding. + """ supports_eagle3: ClassVar[Literal[True]] = True """ @@ -1687,12 +1676,12 @@ class SupportsEagle3(SupportsEagleBase, Protocol): """ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: - """ - Set which layers should output auxiliary hidden states for EAGLE-3. + """Set which layers should output auxiliary hidden states for EAGLE-3. Args: layers: Tuple of layer indices that should output auxiliary hidden states. + """ parent_ref = self if hasattr(self, "get_language_model"): @@ -1710,14 +1699,14 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: holder._set_aux_hidden_state_layers(layers) def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: - """ - Get the default layer indices that should output auxiliary hidden states + """Get the default layer indices that should output auxiliary hidden states for EAGLE-3 for this model. Models can override this method to provide different default layers based on their architecture, but it is encouraged to instead include the layer specification in the model's config if possible. Returns: Tuple of layer indices for auxiliary hidden state outputs. + """ parent_ref = self if hasattr(self, "get_language_model"): @@ -1765,8 +1754,7 @@ def get_mrope_input_positions( input_tokens: list[int], mm_features: list["MultiModalFeatureSpec"], ) -> tuple[torch.Tensor, int]: - """ - Get M-RoPE input positions and delta value for this specific model. + """Get M-RoPE input positions and delta value for this specific model. This method should be implemented by each model that supports M-RoPE to provide model-specific logic for computing input positions. @@ -1779,6 +1767,7 @@ def get_mrope_input_positions( llm_positions: Tensor of shape `[num_dims, num_tokens]`, one row per M-RoPE position channel (e.g. T/H/W). mrope_position_delta: Delta for position calculations. + """ ... @@ -1886,8 +1875,7 @@ def postprocess_encoder_output( clone: bool = False, batch_mm_kwargs: dict[str, Any] | None = None, ) -> None: - """ - Post-process encoder output, directly call scatter_output_slices by default. + """Post-process encoder output, directly call scatter_output_slices by default. By default, delegates directly to scatter_output_slices. Override this for models that require additional processing on the raw @@ -1913,10 +1901,17 @@ def prepare_encoder_cudagraph_capture_inputs( """Create dummy inputs and buffers for CUDA graph capture. Args: + token_budget: Token budget the capture is sized for. + max_batch_size: Maximum number of items in a captured batch. + max_frames_per_batch: Maximum number of frames in a captured batch. + device: Device the dummy inputs and buffers are created on. + dtype: Dtype of the dummy inputs and buffers. + path: Configured encoder path. axis_keys: The resolved capture-axis keys (one per axis of ``EncoderCudaGraphConfig.capture_axes``) this capture is for. None or empty when no capture axes are configured; models without capture axes ignore it. + """ ... diff --git a/vllm/model_executor/models/intern_vit.py b/vllm/model_executor/models/intern_vit.py index 44ed3ec0fed8..959f241de9fd 100644 --- a/vllm/model_executor/models/intern_vit.py +++ b/vllm/model_executor/models/intern_vit.py @@ -143,7 +143,7 @@ def forward( class InternParallelAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py index 0ed2723c32b1..057a6ce2479c 100644 --- a/vllm/model_executor/models/interns1.py +++ b/vllm/model_executor/models/interns1.py @@ -89,13 +89,12 @@ def forward(self, image_features): class InternS1ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of images * (1 + num_patches) - - c: Number of channels (3) - - h: Height - - w: Width - - bn: Batch size * number of images + """Dimensions: + - bnp: Batch size * number of images * (1 + num_patches) + - c: Number of channels (3) + - h: Height + - w: Width + - bn: Batch size * number of images """ type: Literal["pixel_values"] = "pixel_values" @@ -104,11 +103,10 @@ class InternS1ImagePixelInputs(TensorSchema): class InternS1ImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - ni: Number of images - - tifs: Total image feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - ni: Number of images + - tifs: Total image feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -119,13 +117,12 @@ class InternS1ImageEmbeddingInputs(TensorSchema): class InternS1VideoPixelInputs(TensorSchema): - """ - Dimensions: - - bnv: Batch size * number of videos * number of frames - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height - - w: Width + """Dimensions: + - bnv: Batch size * number of videos * number of frames + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height + - w: Width """ type: Literal["pixel_values_videos"] = "pixel_values_videos" @@ -134,11 +131,10 @@ class InternS1VideoPixelInputs(TensorSchema): class InternS1VideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nv: Number of videos - - tvfs: Total video feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - nv: Number of videos + - tvfs: Total video feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["video_embeds"] = "video_embeds" @@ -862,9 +858,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 40c5330ad320..8374542cfd6a 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -621,7 +621,7 @@ def get_frope_params_map(self) -> dict[str, str]: return mapper def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - """load weights""" + """Load weights""" orig_to_new_prefix: dict[str, str | None] = { "model.visual.": "visual.", "lm_head.": "language_model.lm_head.", diff --git a/vllm/model_executor/models/interns1_vit.py b/vllm/model_executor/models/interns1_vit.py index a5bc781561ff..e0958a63e3c4 100644 --- a/vllm/model_executor/models/interns1_vit.py +++ b/vllm/model_executor/models/interns1_vit.py @@ -91,15 +91,13 @@ def __init__(self, config: PretrainedConfig): def interpolate_pos_encoding( self, embeddings: torch.Tensor, height: int, width: int ) -> torch.Tensor: - """ - This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + """This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This method is also adapted to support torch.jit tracing. Adapted from: - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 """ # noqa: E501 - num_patches = embeddings.shape[1] - 1 num_positions = self.position_embeddings.shape[1] - 1 @@ -164,7 +162,7 @@ def forward( class InternSdpaAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -225,8 +223,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> torch.Tensor: - """x shape: (B, N, C)""" - + """X shape: (B, N, C).""" q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index c57359b7b5ad..873c594c54b3 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -70,13 +70,12 @@ class InternVLImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - bnp: Batch size * number of images * (1 + num_patches) - - c: Number of channels (3) - - h: Height of each image patch - - w: Width of each image patch + """Dimensions: + - bn: Batch size * number of images + - bnp: Batch size * number of images * (1 + num_patches) + - c: Number of channels (3) + - h: Height of each image patch + - w: Width of each image patch """ type: Literal["pixel_values"] @@ -85,11 +84,10 @@ class InternVLImagePixelInputs(TensorSchema): class InternVLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of images - - f: Total image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - n: Number of images + - f: Total image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] @@ -100,13 +98,12 @@ class InternVLImageEmbeddingInputs(TensorSchema): class InternVLVideoPixelInputs(TensorSchema): - """ - Dimensions: - - bvf: Batch size * number of videos * num_frames - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each video frame - - w: Width of each video frame + """Dimensions: + - bvf: Batch size * number of videos * num_frames + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each video frame + - w: Width of each video frame """ type: Literal["pixel_values_videos"] @@ -115,11 +112,10 @@ class InternVLVideoPixelInputs(TensorSchema): class InternVLVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of videos - - f: Total video feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - n: Number of videos + - f: Total video feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["video_embeds"] @@ -327,7 +323,7 @@ def _get_prompt_updates( class InternVLProcessingInfo(BaseInternVLProcessingInfo): - """InternVL ProcessingInfo extended for video processing""" + """InternVL ProcessingInfo extended for video processing.""" def get_image_processor(self, **kwargs): config = self.get_hf_config() @@ -418,7 +414,7 @@ def get_num_frames_with_most_features( class InternVLDummyInputsBuilder( BaseInternVLDummyInputsBuilder[InternVLProcessingInfo] ): - """InternVL DummyInputsBuilder extended for video support""" + """InternVL DummyInputsBuilder extended for video support.""" def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_videos = mm_counts.get("video", 0) @@ -460,7 +456,7 @@ def get_dummy_mm_data( class InternVLMultiModalProcessor( BaseInternVLMultiModalProcessor[InternVLProcessingInfo] ): - """InternVL MultiModalProcessor extended for video support""" + """InternVL MultiModalProcessor extended for video support.""" def _postprocess_hf_mm_data( self, @@ -906,9 +902,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="mlp1", diff --git a/vllm/model_executor/models/iquest_loopcoder.py b/vllm/model_executor/models/iquest_loopcoder.py index 782bec0ed0df..87783ae9c856 100644 --- a/vllm/model_executor/models/iquest_loopcoder.py +++ b/vllm/model_executor/models/iquest_loopcoder.py @@ -321,6 +321,7 @@ def forward(self, query: torch.Tensor) -> torch.Tensor: Returns: gate: [num_tokens, num_heads * head_dim] (flattened format matching q shape) + """ num_heads, num_tokens, head_dim = query.shape diff --git a/vllm/model_executor/models/isaac.py b/vllm/model_executor/models/isaac.py index a59345dc25f5..f7c1ee2dd131 100644 --- a/vllm/model_executor/models/isaac.py +++ b/vllm/model_executor/models/isaac.py @@ -177,23 +177,23 @@ def create_pixel_shuffle_index_map( scale_factor: int = 1, device: torch.device | None = None, ) -> torch.Tensor: - """ - Build a gather-index map that tells us, for every *output* token after + """Build a gather-index map that tells us, for every *output* token after pixel-shuffle, which `scale_factor**2` *input* tokens are being merged. - Args + Args: ---- seq_sizes : (num_images,) - #patches in each image (row-major order) token_grids : (num_images,2) - (height, width) for every image scale_factor : spatial down-scale factor (≥2) device : (optional) overrides `seq_sizes.device` - Returns + Returns: ------- gather_idx : (new_total_seq_len, scale_factor**2) int64 tensor. gather_idx[i, j] is the *flat* index into the *original* packed sequence for the j-th sub-patch that forms the i-th output token. + """ if device is None: device = seq_sizes.device @@ -266,6 +266,7 @@ def pixel_shuffle_varlen( Raises: ValueError: If more than one batch item is provided. + """ keep_batch_dim = x.dim() == 3 if keep_batch_dim: @@ -403,8 +404,7 @@ def get_dummy_mm_data( class IsaacImagePixelInputs(TensorSchema): - """ - Schema for validating Isaac image inputs. + """Schema for validating Isaac image inputs. Dimensions: - np: Number of patches @@ -687,12 +687,10 @@ def forward( self, packed_seq_patches: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: - r""" - spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): - Tensor containing the spatial dimensions (height, width) - of the input images. + r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): + Tensor containing the spatial dimensions (height, width) + of the input images. """ - seq_patches, token_grids = packed_seq_patches seq_sizes = torch.prod(token_grids, dim=-1) @@ -1015,9 +1013,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="vision_embedding.linear_fc2", # The final linear layer diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index c17228aa50fd..a9a5dfda24aa 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -48,11 +48,10 @@ class KananaVImagePixelInputs(TensorSchema): - """ - Dimensions: - - np: The total number of patches over all images in the batch - - cps: Number of channels * patch_size * patch_size - - ni: Number of images + """Dimensions: + - np: The total number of patches over all images in the batch + - cps: Number of channels * patch_size * patch_size + - ni: Number of images """ type: Literal["pixel_values"] diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index d117b91d53b6..6bbffae0290e 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -125,13 +125,12 @@ def smart_resize( class KeyeImagePixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - bnp: Batch size * Number of patches + - c: Number of channels + - ps: Patch size + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["pixel_values"] @@ -142,13 +141,12 @@ class KeyeImagePixelInputs(TensorSchema): class KeyeImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nf: Number of image features - - hs: Hidden size (must match the hidden size of language model - backbone) - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - nf: Number of image features + - hs: Hidden size (must match the hidden size of language model + backbone) + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["image_embeds"] @@ -160,13 +158,12 @@ class KeyeImageEmbeddingInputs(TensorSchema): class KeyeVideoPixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - bnp: Batch size * Number of patches + - c: Number of channels + - ps: Patch size + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["pixel_values_videos"] @@ -177,13 +174,12 @@ class KeyeVideoPixelInputs(TensorSchema): class KeyeVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nf: Number of video features - - hs: Hidden size (must match the hidden size of language model - backbone) - - nv: Number of videos - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - nf: Number of video features + - hs: Hidden size (must match the hidden size of language model + backbone) + - nv: Number of videos + - g: Grid dimensions (3 for t, h, w) """ type: Literal["video_embeds"] @@ -349,7 +345,8 @@ def apply_rotary_pos_emb_flashatt( class KeyeSiglipAttention(nn.Module): """Multi-headed attention from 'Attention Is All You - Need' paper.""" + Need' paper. + """ def __init__( self, @@ -1465,6 +1462,9 @@ def forward( otherwise it will be `(seq_len,)`. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. + """ if intermediate_tensors is not None: inputs_embeds = None @@ -1588,8 +1588,7 @@ def _process_video_input( def _split_video_grid_thw( grid_thw: torch.Tensor | list[list[int]] | list[int], ) -> list[list[int]]: - """ - Split video grid_thw along the t dimension into per-frame rows. + """Split video grid_thw along the t dimension into per-frame rows. This preserves Keye's current M-RoPE behavior, where a video is emitted as consecutive frame-level multimodal blocks rather than a single block diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index 8625fd9c0188..eec06304551a 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -51,8 +51,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor: - """ - Split grid_thw in t dimension. + """Split grid_thw in t dimension. Args: grid_thw: [N, 3] tensor of [t, h, w] @@ -66,6 +65,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor: tensor([[1, 3, 4], [1, 3, 4], [1, 5, 6]]) + """ t = grid_thw[:, 0] h_w = grid_thw[:, 1:] @@ -76,8 +76,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor: def get_num_patches( grid_thw: torch.Tensor, num_frames: list[int] | torch.Tensor ) -> list[int]: - """ - Return num_patches per video. + """Return num_patches per video. Args: grid_thw: Tensor with shape [N, 3] containing temporal, height, width @@ -101,8 +100,8 @@ def get_num_patches( >>> get_num_patches(grid_thw, num_frames) tensor([16, 1]) # Total patches for first video: 8+8=16, second video: 1. - """ + """ assert len(grid_thw.shape) == 2 if isinstance(num_frames, torch.Tensor): num_frames = num_frames.clone().tolist() @@ -121,13 +120,12 @@ def get_num_patches( class KeyeVL1_5ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - bnp: Batch size * Number of patches + - c: Number of channels + - ps: Patch size + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["pixel_values"] @@ -140,13 +138,12 @@ class KeyeVL1_5ImagePixelInputs(TensorSchema): class KeyeVL1_5ImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nf: Number of image features - - hs: Hidden size (must match the hidden size of language model - backbone) - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - nf: Number of image features + - hs: Hidden size (must match the hidden size of language model + backbone) + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["image_embeds"] @@ -160,13 +157,12 @@ class KeyeVL1_5ImageEmbeddingInputs(TensorSchema): class KeyeVL1_5VideoPixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - bnp: Batch size * Number of patches + - c: Number of channels + - ps: Patch size + - ni: Number of images + - g: Grid dimensions (3 for t, h, w) """ type: Literal["pixel_values_videos"] @@ -179,13 +175,12 @@ class KeyeVL1_5VideoPixelInputs(TensorSchema): class KeyeVL1_5VideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - nf: Number of video features - - hs: Hidden size (must match the hidden size of language model - backbone) - - nv: Number of videos - - g: Grid dimensions (3 for t, h, w) + """Dimensions: + - nf: Number of video features + - hs: Hidden size (must match the hidden size of language model + backbone) + - nv: Number of videos + - g: Grid dimensions (3 for t, h, w) """ type: Literal["video_embeds"] @@ -437,10 +432,10 @@ def _get_prompt_updates( cu_seqlens = torch.cumsum(torch.tensor([0] + num_frames.tolist()), dim=-1) def get_replacement_keye(item_idx: int, modality: str): - """ - Args: - item_idx(int): The item index of modality to replace - modality(str): The modality + """Args: + item_idx(int): The item index of modality to replace + modality(str): The modality + """ if modality == "image": out_item = out_mm_kwargs[modality][item_idx] diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 5b0061e3cb0b..9af46e441ab7 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Kimi-K2.5 Model Implementation for vLLM. +"""Kimi-K2.5 Model Implementation for vLLM. Kimi-K2.5 extends Kimi-K2 with vision support. """ @@ -90,8 +89,7 @@ class MaxImageTokenMeta: class KimiK25MediaPixelInputs(TensorSchema): - """ - Media input schema for K2-VL model. + """Media input schema for K2-VL model. Dimensions: - np: Number of patches (flattened from all media items) diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index eb811c9fee43..03c409a91195 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Vision tower implementation for Kimi-K2.5 model. +"""Vision tower implementation for Kimi-K2.5 model. This module provides the vision encoder components for Kimi-K2.5, including 3D patch embedding, RoPE position embedding, and @@ -284,12 +283,12 @@ def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: def get_freqs_cis( self, grid_thws: torch.Tensor | list[list[int]], device: torch.device ) -> torch.Tensor: - """ - Args: + """Args: grid_thws (torch.Tensor): grid time, height and width Returns: freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ if not hasattr(self, "freqs_cis"): self.register_buffer( @@ -449,6 +448,10 @@ def attention_qkvpacked( Args: x (torch.Tensor): (seqlen, hidden_dim) cu_seqlens (torch.Tensor): cumulative sequence lengths + rope_freqs_cis (torch.Tensor): rotary embedding frequencies + max_seqlen (torch.Tensor | None): longest sequence in the batch + sequence_lengths (torch.Tensor | None): per-sequence lengths + """ seq_length = x.size(0) xqkv, _ = self.wqkv(x) @@ -753,13 +756,13 @@ def forward( *, encoder_metadata: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: - """ - Args: + """Args: pixel_values (torch.Tensor): The input pixel values. grid_thws (torch.Tensor): Temporal, height and width. Returns: torch.Tensor: The output tokens. + """ if encoder_metadata is not None and "pos_embeds" in encoder_metadata: hidden_states = self.patch_embed( diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index 218c257e53c5..9aba4d194524 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/vllm/model_executor/models/kimi_vl.py @@ -135,12 +135,11 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor: class KimiVLImagePixelInputs(TensorSchema): - """ - Dimensions: - - nc: Number of channels - - np: Number of patches - - ps: Patch size - - ni: Number of images + """Dimensions: + - nc: Number of channels + - np: Number of patches + - ps: Patch size + - ni: Number of images """ type: Literal["pixel_values"] = "pixel_values" diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index 6f03150b8461..c43a9dcbcabd 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/vllm/model_executor/models/lfm2.py @@ -440,6 +440,7 @@ def get_mamba_state_shape_from_config( Returns: Tuple containing: - conv_state_shape: Shape for convolutional state cache + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index c2a21e52d18c..92193f0c2a81 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -556,6 +556,7 @@ def get_mamba_state_shape_from_config( Returns: Tuple containing: - conv_state_shape: Shape for convolutional state cache + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/lfm2_siglip2.py b/vllm/model_executor/models/lfm2_siglip2.py index 54fc1c5d2ff0..57a3d4a8d798 100644 --- a/vllm/model_executor/models/lfm2_siglip2.py +++ b/vllm/model_executor/models/lfm2_siglip2.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Implementation of Siglip2VisionModel intended to be only used -within a vision language model.""" +within a vision language model. +""" from collections.abc import Iterable @@ -61,6 +62,7 @@ def forward( Returns: (1, total_tokens, embed_dim) packed embeddings. + """ assert spatial_shapes.device.type == "cpu", ( "Expected `spatial_shapes` on CPU to avoid device-to-host sync in " @@ -113,6 +115,7 @@ def resize_positional_embeddings_packed( Returns: (total_tokens, embed_dim) packed positional embeddings, concatenated in the same order as `lengths_list`. + """ assert spatial_shapes.device.type == "cpu" @@ -155,7 +158,7 @@ def resize_positional_embeddings_packed( class Siglip2Attention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -304,11 +307,11 @@ def forward( cu_seqlens: torch.Tensor, max_seqlen: int | torch.Tensor, ) -> torch.Tensor: - """ - Args: - hidden_states: Input tensor of shape (batch, seq_len, embed_dim). - cu_seqlens: Cumulative sequence lengths tensor. - max_seqlen: Maximum sequence length. + """Args: + hidden_states: Input tensor of shape (batch, seq_len, embed_dim). + cu_seqlens: Cumulative sequence lengths tensor. + max_seqlen: Maximum sequence length. + """ residual = hidden_states @@ -328,12 +331,12 @@ def forward( class Siglip2Encoder(nn.Module): - """ - Transformer encoder consisting of `config.num_hidden_layers` + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Siglip2EncoderLayer`]. Args: config: PretrainedConfig + """ def __init__( @@ -430,8 +433,7 @@ def forward( max_seqlen: torch.Tensor, select_layers: list[int] | None = None, ) -> torch.Tensor: - r""" - spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): + r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. select_layers (`list[int]` or `None`, defaults to `None`): @@ -500,10 +502,15 @@ def forward( """Forward pass through the vision model. Args: + pixel_values_packed: Packed pixel values for all images. + spatial_shapes: Per-image spatial dimensions. + cu_seqlens: Cumulative sequence lengths of the packed images. + max_seqlen: Longest image sequence in the batch. select_layers: Layer indices to select hidden states from. Supports negative indices (e.g., [-2] for second-to-last). If None, returns the last layer output with post_layernorm. Multiple layers can be selected and will be concatenated. + """ return self.vision_model( pixel_values_packed=pixel_values_packed, diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index 15872a1f3f39..9b24095ab6b0 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -81,12 +81,11 @@ def _pad_cumulative_seqlens_buffer( class Lfm2VLImagePixelInputs(TensorSchema): - """ - Dimensions: - - b: Number of images in the prompt - - bn: Batch size * number of images - - d: Number of dimensions - - fd: Number of features per dimension + """Dimensions: + - b: Number of images in the prompt + - bn: Batch size * number of images + - d: Number of dimensions + - fd: Number of features per dimension """ type: Literal["pixel_values"] = "pixel_values" @@ -538,6 +537,7 @@ def forward( Returns: projected_packed: (total_projected_tokens, text_hidden_size) + """ assert spatial_shapes.device.type == "cpu", ( "Expected `spatial_shapes` on CPU to avoid device-to-host sync in " @@ -667,6 +667,7 @@ def get_mamba_state_shape_from_config( Returns: Tuple containing: - conv_state_shape: Shape for convolutional state cache + """ parallel_config = vllm_config.parallel_config hf_language_config = vllm_config.model_config.hf_config.text_config @@ -1276,9 +1277,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 3dc0e9a904df..3ffa47bba36a 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -406,8 +406,7 @@ def load_moe_expert_weights( expert_params_mapping: list[tuple[str, str, int, str]], fused: bool = True, ) -> bool: - """ - Load MoE expert weights. + """Load MoE expert weights. Args: name: The name of the weight to load. @@ -427,8 +426,8 @@ def load_moe_expert_weights( Returns: True if loaded_weight is one of MoE weights and the MoE expert weights are loaded successfully, False otherwise. - """ + """ # Whether the MoE expert weights are loaded successfully. expert_param_loaded = False diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index a32e0e6c63cf..1c4ebff7a04c 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -72,8 +72,7 @@ class LlavaImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of images - c: Number of channels (3) - h: Height @@ -88,8 +87,7 @@ class LlavaImagePixelInputs(TensorSchema): class PixtralHFImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of images - c: Number of channels - h: Height @@ -107,11 +105,10 @@ class PixtralHFImagePixelInputs(TensorSchema): class LlavaImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ifs: Image feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - ifs: Image feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -440,6 +437,7 @@ def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int: Args: hf_config: Model config with vision feature layer(s). + """ feature_layers = hf_config.vision_feature_layer num_hidden_layers = hf_config.vision_config.num_hidden_layers @@ -707,9 +705,12 @@ def forward( positions: Position indices for the input tokens. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. Info: [`LlavaImageInputs`][vllm.model_executor.models.llava.LlavaImageInputs] + """ if intermediate_tensors is not None: inputs_embeds = None @@ -731,9 +732,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py index fee9a4ac59a3..b7b1a8533d6f 100644 --- a/vllm/model_executor/models/llava_next.py +++ b/vllm/model_executor/models/llava_next.py @@ -47,8 +47,7 @@ class LlavaNextImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of images - np: Number of patches + 1 - c: Number of channels (3) @@ -70,11 +69,10 @@ class LlavaNextImagePixelInputs(TensorSchema): class LlavaNextImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ifs: Image feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - ifs: Image feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -578,9 +576,12 @@ def forward( positions: Position indices for the input tokens. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. Info: [`LlavaNextImageInputs`][vllm.model_executor.models.llava_next.LlavaNextImageInputs] + """ if intermediate_tensors is not None: inputs_embeds = None @@ -601,9 +602,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/llava_next_video.py b/vllm/model_executor/models/llava_next_video.py index 96ebdaea8eb1..235d1adbac24 100644 --- a/vllm/model_executor/models/llava_next_video.py +++ b/vllm/model_executor/models/llava_next_video.py @@ -56,8 +56,7 @@ class LlavaNextVideoPixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of videos - f: Number of frames - c: Number of channels (3) @@ -381,8 +380,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def _parse_and_validate_video_input( self, **kwargs: object ) -> LlavaNextVideoPixelInputs | None: - """ - A legal video input should have the following dimensions: + """A legal video input should have the following dimensions: { "pixel_values_videos" : list[b, Tensor(nb_frames, nb_channels, height, width)] @@ -457,10 +455,17 @@ def forward( **kwargs: object, ) -> torch.Tensor | IntermediateTensors: """Run forward pass for LlaVA-NeXT-Video. + Args: input_ids: Flattened (concatenated) input_ids corresponding to a batch. - pixel_values_videos: Pixels in each frames for each input videos. + positions: Flattened (concatenated) position ids corresponding to a + batch. + intermediate_tensors: Intermediate tensors from prior forward pass. + inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. + """ if intermediate_tensors is not None: inputs_embeds = None @@ -486,9 +491,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py index 631f73e64d99..55f63c49baf9 100644 --- a/vllm/model_executor/models/llava_onevision.py +++ b/vllm/model_executor/models/llava_onevision.py @@ -53,17 +53,16 @@ class LlavaOnevisionVideoPixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of videos - - f: Number of frames - - c: Number of channels (3) - - h: Height - - w: Width - - Note that `f` may be different for each batch, and 'num_frames' - may be different for each video, in which case the data is passed as a - list instead of a batched tensor. + """Dimensions: + - bn: Batch size * number of videos + - f: Number of frames + - c: Number of channels (3) + - h: Height + - w: Width + + Note that `f` may be different for each batch, and 'num_frames' + may be different for each video, in which case the data is passed as a + list instead of a batched tensor. """ type: Literal["pixel_values_videos"] = "pixel_values_videos" @@ -75,16 +74,15 @@ class LlavaOnevisionVideoPixelInputs(TensorSchema): class LlavaOnevisionImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - np: Number of patches (1 + num_patches) - - c: Number of channels (3) - - h: Height - - w: Width - - Note that `num_patches` may be different per batch and image, - in which case the data is passed as a list instead of a batched tensor. + """Dimensions: + - bn: Batch size * number of images + - np: Number of patches (1 + num_patches) + - c: Number of channels (3) + - h: Height + - w: Width + + Note that `num_patches` may be different per batch and image, + in which case the data is passed as a list instead of a batched tensor. """ type: Literal["pixel_values"] = "pixel_values" @@ -98,11 +96,10 @@ class LlavaOnevisionImagePixelInputs(TensorSchema): class LlavaOnevisionImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ifs: Image feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - ifs: Image feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -557,8 +554,7 @@ def _parse_and_validate_image_input( def _parse_and_validate_video_input( self, **kwargs: object ) -> LlavaOnevisionVideoPixelInputs | None: - """ - A legal video input should have the following dimensions: + """A legal video input should have the following dimensions: { "pixel_values_videos" : list[b, Tensor(nb_frames, nb_channels, height, width)] @@ -891,10 +887,17 @@ def forward( **kwargs: object, ) -> torch.Tensor | IntermediateTensors: """Run forward pass for LlaVA-Onevision. + Args: input_ids: Flattened (concatenated) input_ids corresponding to a batch. - pixel_values_videos: Pixels in each frames for each input videos. + positions: Flattened (concatenated) position ids corresponding to a + batch. + intermediate_tensors: Intermediate tensors from prior forward pass. + inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. + """ if intermediate_tensors is not None: inputs_embeds = None diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 7eda9dc32025..71c3bc79dc98 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Inference-only LLaVA-OneVision-2 (OV2) model for vLLM. +"""Inference-only LLaVA-OneVision-2 (OV2) model for vLLM. Architecture notes: diff --git a/vllm/model_executor/models/longcat_flash_mtp.py b/vllm/model_executor/models/longcat_flash_mtp.py index 360eb9c76378..851c694f3714 100644 --- a/vllm/model_executor/models/longcat_flash_mtp.py +++ b/vllm/model_executor/models/longcat_flash_mtp.py @@ -313,8 +313,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def _rewrite_spec_layer_name( self, spec_layer: int, name: str, new_to_old_names_mapping: dict ) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/longcat_flash_ngram.py b/vllm/model_executor/models/longcat_flash_ngram.py index 76aa80013eb5..435a11ee17d8 100644 --- a/vllm/model_executor/models/longcat_flash_ngram.py +++ b/vllm/model_executor/models/longcat_flash_ngram.py @@ -156,6 +156,7 @@ def embed_batched( oe_ids: ``[num_tokens, num_embedders]`` global (offset) n-gram ids, as produced by the ``ngram_compute_n_gram_ids`` kernel. Returns: ``[num_tokens, hidden]``. + """ word = self.word_embeddings(input_ids) # [N, H] flat = oe_ids.permute(1, 0).contiguous() # [num_embedders, N] diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py index 343111ee0151..cb8b2b481631 100644 --- a/vllm/model_executor/models/mamba2.py +++ b/vllm/model_executor/models/mamba2.py @@ -197,6 +197,7 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/medusa.py b/vllm/model_executor/models/medusa.py index 6adf3b7669ec..bf43f85d362c 100644 --- a/vllm/model_executor/models/medusa.py +++ b/vllm/model_executor/models/medusa.py @@ -51,7 +51,8 @@ class Medusa(nn.Module): truncated_vocab_size < vocab_size. To use this technique, one has to find the top-k most frequent tokens in target dataset and add that as a tensor in the draft checkpoint (using key token_map). Also, the draft config - needs to have truncated_vocab_size (=k) as an attribute.""" + needs to have truncated_vocab_size (=k) as an attribute. + """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: speculative_config = vllm_config.speculative_config diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py index ebfbaf93fd6f..a75e3b0a8ecc 100644 --- a/vllm/model_executor/models/mellum.py +++ b/vllm/model_executor/models/mellum.py @@ -27,8 +27,7 @@ class MellumAttention(Qwen3MoeAttention): - """ - Differences from `Qwen3MoeAttention`: + """Differences from `Qwen3MoeAttention`: - Supports `per_layer_sliding_window` for `Attention`. """ @@ -122,8 +121,7 @@ def __init__( class MellumDecoderLayer(Qwen3MoeDecoderLayer): - """ - Differences from `Qwen3MoeDecoderLayer`: + """Differences from `Qwen3MoeDecoderLayer`: - Supports interleaved SWA and per-layer RoPE scaling. """ @@ -185,8 +183,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: @support_torch_compile class MellumModel(Qwen3MoeModel): - """ - Differences from `Qwen3MoeModel`: + """Differences from `Qwen3MoeModel`: - Uses `MellumDecoderLayer`. """ @@ -204,8 +201,7 @@ def __init__( class MellumForCausalLM(Qwen3MoeForCausalLM): - """ - Differences from `Qwen3MoeForCausalLM`: + """Differences from `Qwen3MoeForCausalLM`: - Uses `MellumModel`. """ diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index 1129a0e4b284..300f274a187e 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -511,11 +511,9 @@ def forward(self, x, mask=None): # === Audio Inputs === # class MiDashengLMAudioInputs(TensorSchema): - """ - - Dimensions: - - bn: Batch size * number of audios - - p: Number of sampling points + """Dimensions: + - bn: Batch size * number of audios + - p: Number of sampling points """ input_values: Annotated[torch.Tensor, TensorShape("n", "p")] diff --git a/vllm/model_executor/models/mimo_audio.py b/vllm/model_executor/models/mimo_audio.py index f1f901418a78..d0d2513f0424 100644 --- a/vllm/model_executor/models/mimo_audio.py +++ b/vllm/model_executor/models/mimo_audio.py @@ -1315,6 +1315,7 @@ def process_audio(self, audio: torch.Tensor) -> torch.Tensor: Returns: [T//group_size, group_size, audio_channels] + """ T = audio.shape[0] audio = audio[:, : self.audio_channels] @@ -1355,6 +1356,7 @@ def get_audio_feature( Tuple of: - audio_embeds: [total_tokens, out_hidden_size] concatenated embeddings - item_token_lens: list of int, number of tokens per input item + """ if self.audio_tokenizer is None: raise RuntimeError( diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index c7ba39bb0dd9..ec67d72cf2ca 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -776,8 +776,7 @@ def _try_load_fp8_qkv_proj( tp_rank: int, tp_size: int, ) -> bool: - """ - The fused fp8 QKV projection weights and scale are stored separately. + """The fused fp8 QKV projection weights and scale are stored separately. Special care must be taken while sharding these tensors across TP ranks. See _shard_fp8_qkv_proj for more details. @@ -785,6 +784,7 @@ def _try_load_fp8_qkv_proj( True if ``tensor`` was an fp8 qkv_proj weight/scale and was consumed (caller should skip it); False otherwise, so the caller falls through to its normal loading path. + """ is_weight = ( name.endswith("qkv_proj.weight") and tensor.dtype == torch.float8_e4m3fn diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index fa2cabe28860..b08ba229f7a1 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -141,7 +141,7 @@ def forward( class _MiMoV2MTPLayers(nn.Module): - """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*""" + """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*.""" def __init__( self, diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 9281a1689ae1..0edd0736e178 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -275,14 +275,14 @@ def forward( max_seqlen: torch.Tensor, full_attn: bool = True, ) -> torch.Tensor: - """ - Args: - x: [seq_len, batch=1, embed_dim] (seq-first convention) - cu_seqlens: cumulative sequence lengths [num_seqs+1], int32 - rotary_pos_emb_cos: [seq_len, qk_channels // 2] - rotary_pos_emb_sin: [seq_len, qk_channels // 2] - max_seqlen: maximum sequence length - full_attn: if True, full attention; if False, window attention + """Args: + x: [seq_len, batch=1, embed_dim] (seq-first convention) + cu_seqlens: cumulative sequence lengths [num_seqs+1], int32 + rotary_pos_emb_cos: [seq_len, qk_channels // 2] + rotary_pos_emb_sin: [seq_len, qk_channels // 2] + max_seqlen: maximum sequence length + full_attn: if True, full attention; if False, window attention + """ # [seq_len, 1, embed_dim] -> QKV projection qkv, _ = self.qkv(x) # [seq_len, 1, q_size + kv_size + kv_size] @@ -531,6 +531,7 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso Returns: cos: [total_tokens, qk_channels // 2] sin: [total_tokens, qk_channels // 2] + """ cos_list, sin_list = [], [] for i in range(grid_thw.size(0)): @@ -576,12 +577,12 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso return torch.cat(cos_list, dim=0), torch.cat(sin_list, dim=0) def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - """ - Args: + """Args: x: [total_tokens, C] pre-flattened patches grid_thw: [num_images, 3] tensor of (t, h, w) for each image/video Returns: [merged_tokens, out_hidden_size] + """ # Ensure grid_thw is a tensor if not isinstance(grid_thw, torch.Tensor): @@ -1556,8 +1557,10 @@ def forward( batch. **NOTE**: If mrope is enabled (default setting for Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`, otherwise it will be `(seq_len,). - """ + intermediate_tensors: Intermediate tensors from prior forward pass. + inputs_embeds: Optional tensor of input embeddings. + """ if intermediate_tensors is not None: inputs_embeds = None diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index c8b46a45f101..4251514326c6 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -131,13 +131,12 @@ class MiniCPMOAudioFeatureInputs(TensorSchema): - """ - Dimensions: - - bns: Batch size * number of audios * number of slices - - bn: Batch size * number of audios - - c: Number of channels - - l: Length - - s: Number of slices + """Dimensions: + - bns: Batch size * number of audios * number of slices + - bn: Batch size * number of audios + - c: Number of channels + - l: Length + - s: Number of slices """ type: Literal["audio_features"] = "audio_features" @@ -163,8 +162,7 @@ class MiniCPMOAudioFeatureInputs(TensorSchema): class MiniCPMOAudioEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of audios - s: Number of slices - h: Hidden size (must match language model backbone) @@ -1008,8 +1006,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dummy_inputs=MiniCPMODummyInputsBuilder, ) class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6): - """ - MiniCPM-O model with audio support. + """MiniCPM-O model with audio support. Different versions use different LLM backbones: - Version 2.6: Uses Qwen2 - Version 4.5: Uses Qwen3 diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index dac5e8367fbe..cf818bd77e10 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -124,13 +124,12 @@ class MiniCPMVImagePixelInputs(TensorSchema): - """ - Dimensions: - - bns: Batch size * number of images * number of slices - - bn: Batch size * number of images - - c: Number of channels - - h: Height - - w: Width + """Dimensions: + - bns: Batch size * number of images * number of slices + - bn: Batch size * number of images + - c: Number of channels + - h: Height + - w: Width """ type: Literal["pixel_values"] = "pixel_values" @@ -152,11 +151,10 @@ class MiniCPMVImagePixelInputs(TensorSchema): class MiniCPMVImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ns: Number of slices - - hs: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - ns: Number of slices + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] @@ -307,8 +305,7 @@ def __init__( def get_1d_sincos_pos_embed_from_temporal_size( self, embed_dim: int, pos: np.ndarray ): - """ - embed_dim: output dimension for each position + """embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) """ @@ -1219,8 +1216,7 @@ def _get_mm_fields_config( class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): - """ - The abstract class of MiniCPMV can only be inherited, but cannot be + """The abstract class of MiniCPMV can only be inherited, but cannot be instantiated. """ @@ -1419,9 +1415,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="llm", connector="resampler", tower_model="vpm" ) @@ -1494,7 +1488,6 @@ def _mcpmv_normalize_tgt_sizes( slice_counts: list[int], ) -> torch.Tensor: """Normalize tgt_sizes to shape ``(total_slices, 2)``.""" - if isinstance(tgt_sizes, list): if not tgt_sizes: tgt_sizes = torch.zeros((0, 2), dtype=torch.long) @@ -2457,8 +2450,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: dummy_inputs=MiniCPMVDummyInputsBuilder, ) class MiniCPMV(MiniCPMVBaseModel, SupportsMultiModal, SupportsLoRA): - """ - Different versions of MiniCPMV use different visual encoders and LLMs, + """Different versions of MiniCPMV use different visual encoders and LLMs, which is not conducive to the current integration logic of LoRA and bitsandbytes in vLLM. Therefore, it is necessary to separate them. """ diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index d8e10d349884..c8a2fe4bcddb 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -542,19 +542,22 @@ def _get_downsample_mode( def _compute_visual_tokens( self, - image_size, + image_size: ImageSize, max_slice_nums: int | None = None, downsample_mode: str | None = None, ) -> tuple[list[int], int, int]: """Compute grid, source_image_visual_tokens and patch_visual_tokens. Args: + image_size: Size of the source image. + max_slice_nums: Maximum number of slices, or None for the default. downsample_mode: ``"16x"`` (default, full merge) or ``"4x"`` (skip vit_merger, 4x more visual tokens). Returns: (grids, source_image_visual_tokens, patch_visual_tokens) grids is [0, 0] when no slicing occurs. + """ image_processor = self.get_image_processor() if max_slice_nums is None: @@ -855,7 +858,8 @@ def forward( class MiniCPMV4_6DownsampleMLP(nn.Module): """Match HF (transformers v5.7+) parameter naming: pre_norm/linear_1/ - act/linear_2 (instead of pre_norm + Sequential(mlp.0/mlp.2)).""" + act/linear_2 (instead of pre_norm + Sequential(mlp.0/mlp.2)). + """ def __init__( self, @@ -906,10 +910,10 @@ def forward( hidden_states: torch.Tensor, tgt_sizes: torch.Tensor, ) -> list[torch.Tensor]: - """ - Args: - hidden_states: (B, max_patches, D) padded batch. - tgt_sizes: (B, 2) actual (H, W) per sample. + """Args: + hidden_states: (B, max_patches, D) padded batch. + tgt_sizes: (B, 2) actual (H, W) per sample. + """ m1, m2 = self.merge_kernel_size results = [] diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 79c53b1c7e1c..0258fd7bc3a9 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -58,12 +58,11 @@ class Mistral3ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral" @@ -77,9 +76,7 @@ class Mistral3ImagePixelInputs(TensorSchema): class Mistral3PatchMerger(nn.Module): - """ - Learned merging of spatial_merge_size ** 2 patches - """ + """Learned merging of spatial_merge_size ** 2 patches.""" def __init__( self, vision_hidden_size: int, spatial_merge_size: int, patch_size: int @@ -365,6 +362,7 @@ def _get_num_hidden_layers(hf_config: Mistral3Config) -> int: Args: hf_config: Model config with vision feature layer(s). + """ feature_layers = hf_config.vision_feature_layer num_hidden_layers = hf_config.vision_config.num_hidden_layers @@ -588,9 +586,12 @@ def forward( positions: Position indices for the input tokens. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. Info: [`Mistral3ImagePixelInputs`][vllm.model_executor.models.mistral3.Mistral3ImagePixelInputs] + """ if intermediate_tensors is not None: inputs_embeds = None @@ -612,9 +613,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="multi_modal_projector", diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 0096eb42418f..027a8db3b5ce 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -91,12 +91,11 @@ class Llama4ImagePatchInputs(TensorSchema): - """ - Dimensions: - - batch_size: Batch size - - total_num_chunks: Batch size * number of chunks - - num_channels: Number of channels - - image_size: Size of each image + """Dimensions: + - batch_size: Batch size + - total_num_chunks: Batch size * number of chunks + - num_channels: Number of channels + - image_size: Size of each image """ type: Literal["pixel_values"] = "pixel_values" @@ -413,16 +412,15 @@ def forward( self, hidden_states: torch.Tensor, ) -> torch.Tensor: - r""" - Args: - hidden_states: Input tensor of shape - (batch_size, sequence_length, hidden_size). - Hidden states from the model embeddings, representing - the input tokens. - associated vectors than the model's internal embedding - lookup matrix. - """ + r"""Args: + hidden_states: Input tensor of shape + (batch_size, sequence_length, hidden_size). + Hidden states from the model embeddings, representing + the input tokens. + associated vectors than the model's internal embedding + lookup matrix. + """ for encoder_layer in self.layers: layer_outputs = encoder_layer(hidden_states) hidden_states = layer_outputs[0] @@ -1092,7 +1090,8 @@ def _consolidate_qkv_weights( def _rename_weight_for_modelopt_checkpoint(self, name: str) -> str: """Rename weights from ModelOpt llama4 fp8 checkpoints to vLLM - format.""" + format. + """ if name.startswith("model.") or name.startswith("language_model.model."): renamed = ( name.replace("model.", "language_model.model.", 1) @@ -1270,9 +1269,7 @@ def regular_language_model_weights() -> Iterable[tuple[str, torch.Tensor]]: return updated_params def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector=[ diff --git a/vllm/model_executor/models/mlp_speculator.py b/vllm/model_executor/models/mlp_speculator.py index 612baba8eaa6..0e58952102d6 100644 --- a/vllm/model_executor/models/mlp_speculator.py +++ b/vllm/model_executor/models/mlp_speculator.py @@ -19,10 +19,10 @@ class MLPSpeculatorLayerNorm(nn.Module): - """ - A L2 normalization implementation + """A L2 normalization implementation ... - Args + + Args: ---- normalized_shape : int Dimensionality of input data (size of final tensor axis) @@ -32,6 +32,7 @@ class MLPSpeculatorLayerNorm(nn.Module): (i.e. fp16 requires eps >= 6e-8). elementwise_scale_and_shift : bool Include a learned scaling and shift term after normalization. + """ def __init__( @@ -58,8 +59,7 @@ def forward(self, x): class MLPSpeculator(nn.Module): - """ - An implementation of the speculative models introduced in + """An implementation of the speculative models introduced in "Accelerating Production LLMs with Combined Token/Embedding Speculators" https://arxiv.org/pdf/2404.19124 diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index 502b8a5ce6ca..8547b689a03c 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -93,13 +93,12 @@ class MolmoImageInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - bnc: Batch size * number of images * number of crops (dynamic) - - np: Number of patches - - tp: Token sequence positions - - pd: Patch dimension + """Dimensions: + - bn: Batch size * number of images + - bnc: Batch size * number of images * number of crops (dynamic) + - np: Number of patches + - tp: Token sequence positions + - pd: Patch dimension """ images: Annotated[torch.Tensor, TensorShape("bnc", "np", "pd")] @@ -382,9 +381,7 @@ def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tens def forward( self, x: torch.Tensor, patch_num: tuple[int, int] | None = None ) -> list[torch.Tensor]: - """ - : param x: (batch_size, num_patch, n_pixels) - """ + """: param x: (batch_size, num_patch, n_pixels).""" if patch_num is None: patch_num = self.patch_num B, N, D = x.shape @@ -728,9 +725,7 @@ def device(self) -> torch.device: return self.image_vit.patch_embedding.weight.device def encode_image(self, images: torch.Tensor) -> torch.Tensor: - """ - : param images: (batch_size, num_crops, num_patch, n_pixels) - """ + """: param images: (batch_size, num_crops, num_patch, n_pixels).""" B, T, N, D = images.shape mask = ~torch.all(images.view(B * T, N, D) == -1, dim=(1, 2), keepdim=True) @@ -1456,9 +1451,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="model", connector="vision_backbone.image_projector", diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 66de39929424..61a9d2b15385 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -105,15 +105,14 @@ class Molmo2ImageInputs(TensorSchema): - """ - Dimensions: - - nc: The total number of crops (dynamic) - - np: The total number of patches per crop - - cps: Number of channels * patch_size * patch_size - - npp: Number of pooled patches (dynamic) - - pp: pooling_size * pooling_size - - ni: Number of images - - nt: Number of image tokens (dynamic) + """Dimensions: + - nc: The total number of crops (dynamic) + - np: The total number of patches per crop + - cps: Number of channels * patch_size * patch_size + - npp: Number of pooled patches (dynamic) + - pp: pooling_size * pooling_size + - ni: Number of images + - nt: Number of image tokens (dynamic) """ pixel_values: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")] @@ -132,15 +131,14 @@ class Molmo2ImageInputs(TensorSchema): class Molmo2VideoInputs(TensorSchema): - """ - Dimensions: - - nc: The total number of frames (dynamic) - - np: The total number of patches per frame - - cps: Number of channels * patch_size * patch_size - - npp: Number of pooled patches (dynamic) - - pp: pooling_size * pooling_size - - nv: Number of videos - - nt: Number of video tokens (dynamic) + """Dimensions: + - nc: The total number of frames (dynamic) + - np: The total number of patches per frame + - cps: Number of channels * patch_size * patch_size + - npp: Number of pooled patches (dynamic) + - pp: pooling_size * pooling_size + - nv: Number of videos + - nt: Number of video tokens (dynamic) """ pixel_values_videos: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")] @@ -165,7 +163,7 @@ class Molmo2MultiModalInputs(TypedDict, total=False): @dataclass class VitConfig: - """Config for a vision transformer""" + """Config for a vision transformer.""" hidden_size: int = 1152 intermediate_size: int = 4304 @@ -190,7 +188,7 @@ def image_num_patch(self): @dataclass class AdapterConfig: - """Config for a vit-llm adapter""" + """Config for a vit-llm adapter.""" vit_layers: tuple[int, int] = (-3, -9) pooling_attention_mask: bool = False @@ -205,7 +203,7 @@ class AdapterConfig: @dataclass class TextConfig: - """Configuration for a text model transformer""" + """Configuration for a text model transformer.""" hidden_size: int = 3584 """ @@ -530,9 +528,7 @@ def forward( x: torch.Tensor, patch_num: tuple[int, int] | None = None, ) -> list[torch.Tensor]: - """ - : param x: (batch_size, num_patch, n_pixels) - """ + """: param x: (batch_size, num_patch, n_pixels).""" if patch_num is None: patch_num = self.patch_num @@ -545,7 +541,7 @@ def forward( class ImagePoolingAttention(nn.Module): - """Multi-head attention used for image pooling""" + """Multi-head attention used for image pooling.""" def __init__( self, @@ -667,7 +663,7 @@ def forward( class ImageProjectorMLP(nn.Module): - """MLP used for the image projector""" + """MLP used for the image projector.""" def __init__( self, @@ -786,9 +782,7 @@ def device(self) -> torch.device: return self.image_vit.patch_embedding.weight.device def encode_image(self, images: torch.Tensor) -> torch.Tensor: - """ - : param images: (batch_size, num_crops, num_patch, n_pixels) - """ + """: param images: (batch_size, num_crops, num_patch, n_pixels).""" B, T, N, D = images.shape images = images.view(B * T, N, D) image_features = self.image_vit(images) @@ -1444,8 +1438,7 @@ def get_candidate_target_fps( sampling_fps: int | float, max_fps: int | float = _MAX_VIDEO_FPS, ) -> list[float]: - """ - Return the subset of `video_fps` factors that remain multiples + """Return the subset of `video_fps` factors that remain multiples of `sampling_fps`. Examples: @@ -1460,6 +1453,7 @@ def get_candidate_target_fps( ... ValueError: sampling_fps=2 must divide video_fps=5 to produce consistent frame steps. + """ video_fps = int(video_fps) sampling_fps = int(sampling_fps) @@ -1494,9 +1488,7 @@ def get_target_fps( frame_sample_mode: str, candidate_target_fps: list[float], ) -> float | None: - """ - Get the target fps that best spans the video and has the most frames sampled - """ + """Get the target fps that best spans the video and has the most frames sampled.""" num_frames_sampled = 0 selected_target_fps = None for target_fps in candidate_target_fps: @@ -2646,9 +2638,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="model", connector="vision_backbone.image_projector", diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py index 47a240b58953..45b9420184fc 100644 --- a/vllm/model_executor/models/moondream3.py +++ b/vllm/model_executor/models/moondream3.py @@ -295,6 +295,7 @@ def create_patches(self, images: torch.Tensor) -> torch.Tensor: Returns: patches: (batch, num_patches, patch_dim) + """ patch_size = self.config.enc_patch_size batch, channels, height, width = images.shape @@ -321,6 +322,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: Returns: features: (batch, num_patches, hidden_size) + """ # Create patches and embed patches = self.create_patches(pixel_values) @@ -489,7 +491,6 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward pass with expert parallelism and custom GeGLU activation.""" - # Get router logits and compute top-k router_logits, _ = self.gate(x) # [num_tokens, num_experts] topk_logits, topk_ids = torch.topk( @@ -1318,7 +1319,6 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with remapping from HuggingFace format.""" - params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 9269cd29766d..4e43c9c364de 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/vllm/model_executor/models/moonvit.py @@ -80,13 +80,14 @@ def _apply_rope_input_validation(x, freqs_cis): def apply_rope( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Args: (The leading dimensions of all inputs should be the same) + """Args: (The leading dimensions of all inputs should be the same) xq: query, tensor of shape (..., num_heads, head_dim) xk: key, tensor of shape (..., num_heads, head_dim) freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + Returns: xq_out, xk_out: tensors of shape (..., num_heads, head_dim) + """ _apply_rope_input_validation(xq, freqs_cis) _apply_rope_input_validation(xk, freqs_cis) @@ -188,8 +189,7 @@ def forward( *, pos_embeds: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Args: + """Args: x (L, Channels): input tensor grid_hw (N, 2): grid height and width pos_embeds: precomputed positional embeddings of shape @@ -199,6 +199,7 @@ def forward( Returns: (L, Cout) tensor + """ x = self.proj(x).view(x.size(0), -1) if pos_embeds is not None: @@ -226,6 +227,7 @@ class Rope2DPosEmb(nn.Module): max_width (int): the maximum width of the 2D grid theta_base (float): the base of the theta device (str): the device to store the precomputed cis + """ def __init__( @@ -306,24 +308,26 @@ def get_freqs_cis_by_seqlens_list( return freqs_cis def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor: - """ - Args: + """Args: grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples. + Returns: freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ return self.get_freqs_cis_by_seqlens_list(grid_hws.tolist()) def get_freqs_cis_by_idx( self, pos_idx: torch.Tensor, pos_idx_mask: torch.Tensor ) -> torch.Tensor: - """ - Args: + """Args: pos_idx: tensor of shape (..., 2), It contains the (h, w) position indices of each 2D token. pos_idx_mask: a mask of shape (...), the leading dimensions should be the same as pos_idx. Rope will only be applied to the tokens with True mask. `freqs_cis` for the tokens with False mask with be ones. + Return: freqs_cis: tensor of shape (..., dim//2) + """ assert ( pos_idx.shape[:-1] == pos_idx_mask.shape @@ -343,10 +347,10 @@ def get_freqs_cis_by_idx( class MLP2(nn.Module): - """ - Args: - dims: [in_dim, hidden_dim, out_dim] - bias: whether to use bias in linear layer. + """Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ def __init__( @@ -441,13 +445,13 @@ def attention_qkvpacked( rope_freqs_cis: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, ): - """ - Args: - x (torch.Tensor): (seqlen, hidden_dim) - cu_seqlens (torch.Tensor): - max_seqlen: Optional precomputed scalar tensor. When omitted it - is derived from ``cu_seqlens``, which produces a GPU scalar - that breaks CUDA graph capture. + """Args: + x (torch.Tensor): (seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + max_seqlen: Optional precomputed scalar tensor. When omitted it + is derived from ``cu_seqlens``, which produces a GPU scalar + that breaks CUDA graph capture. + """ seq_length = x.size(0) xqkv, _ = self.wqkv(x) @@ -487,13 +491,13 @@ def forward( rope_freqs_cis: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Args: + """Args: hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set max_seqlen: optional precomputed max-sequence-length scalar. Returns: output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input + """ residual = hidden_states hidden_states = self.norm0(hidden_states) @@ -724,6 +728,7 @@ def prepare_encoder_metadata( graph. device: Device for the metadata tensors. Defaults to the model's parameter device. + """ if device is None: device = next(self.parameters()).device @@ -785,18 +790,18 @@ def forward( *, encoder_metadata: dict[str, Any] | None = None, ) -> torch.Tensor | list[torch.Tensor]: - """ - Args: - pixel_values (torch.Tensor): The input pixel values. - grid_hw (torch.Tensor): The grid height and width. - encoder_metadata: Optional precomputed metadata produced by - :meth:`prepare_encoder_metadata`. When provided every - ``.tolist()`` call in the forward path is skipped, the - returned tensor is the packed - ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable - for CUDA graph capture/replay), and ``grid_hw`` is unused. - When ``None`` the legacy path runs and returns a list of - per-image tensors. + """Args: + pixel_values (torch.Tensor): The input pixel values. + grid_hw (torch.Tensor): The grid height and width. + encoder_metadata: Optional precomputed metadata produced by + :meth:`prepare_encoder_metadata`. When provided every + ``.tolist()`` call in the forward path is skipped, the + returned tensor is the packed + ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable + for CUDA graph capture/replay), and ``grid_hw`` is unused. + When ``None`` the legacy path runs and returns a list of + per-image tensors. + """ if encoder_metadata is not None: hidden_states = self.patch_embed( diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index d0bd8510a894..64a31375c862 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -147,11 +147,10 @@ def _ensure_moss_audio_placeholder_tokens(tokenizer: object) -> None: class MossAudioAudioInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - nmb: Number of mel bins - - t: Time frames + """Dimensions: + - b: Batch size + - nmb: Number of mel bins + - t: Time frames """ audio_data: Annotated[torch.Tensor, TensorShape("b", "nmb", "t")] @@ -1091,6 +1090,7 @@ def __call__( -> audio_data has shape [1, mel_dim, max_time] -> mel_dim is the number of mel filter-bank bins, 128 by default -> audio_data_seqlens stores the unpadded mel length + """ del kwargs @@ -1723,6 +1723,7 @@ def _process_audio_input( audio_data=[2, 128, 1200], audio_data_seqlens=[800, 1200] -> returns (audio0_embeds, audio1_embeds), split by token length -> DeepStack packs each item as [main, layer0, ...] on dim -1 + """ audio_data = audio_input["audio_data"] audio_data_seqlens = audio_input["audio_data_seqlens"] diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 4dc3bdb3d743..659b28d36b0f 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -89,12 +89,11 @@ class MossTranscribeDiarizeAudioInputs(TensorSchema): - """ - Dimensions: - - c: Audio chunks - - m: Mel bins - - f: Mel frames - - n: Number of audio items + """Dimensions: + - c: Audio chunks + - m: Mel bins + - f: Mel frames + - n: Number of audio items """ type: Literal["audio_features"] = "audio_features" @@ -114,11 +113,10 @@ class MossTranscribeDiarizeAudioInputs(TensorSchema): class MossTranscribeDiarizeEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of audio items - - t: Number of audio tokens - - h: Hidden size + """Dimensions: + - n: Number of audio items + - t: Number of audio tokens + - h: Hidden size """ type: Literal["audio_embeds"] = "audio_embeds" diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index ade38cc703da..192ac2de9812 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -114,7 +114,8 @@ def _text_config(config): """MuseGlimmer checkpoints may nest the text config under ``text_config`` (multimodal ``MuseGlimmerConfig``) or expose it directly - (``MuseGlimmerTextConfig``).""" + (``MuseGlimmerTextConfig``). + """ return getattr(config, "text_config", config) @@ -485,7 +486,8 @@ def video_replacement(item_idx: int) -> PromptUpdateDetails: def _muse_glimmer_use_qk_norm(config) -> bool: """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``). - Treat a missing/None flag as True — only an explicit ``False`` disables it.""" + Treat a missing/None flag as True — only an explicit ``False`` disables it. + """ val = getattr(config, "use_qk_norm", None) return True if val is None else bool(val) @@ -493,7 +495,8 @@ def _muse_glimmer_use_qk_norm(config) -> bool: def _muse_glimmer_use_attn_output_gate(config) -> bool: """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS applies it; the modular HF ``text_config`` omits ``use_attn_output_gate`` - (reads as ``None``). Missing/None -> True; only explicit ``False`` disables.""" + (reads as ``None``). Missing/None -> True; only explicit ``False`` disables. + """ val = getattr(config, "use_attn_output_gate", None) return True if val is None else bool(val) @@ -1663,9 +1666,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="model", connector=["vision_adapter.", "vision_projection."], diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 6f5e073bd623..2620923a490b 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -103,12 +103,11 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema): - """ - Dimensions: - - c: Number of audio clips (possibly flattened across audio items) - - b: Number of original audio items - - t: Audio feature length - - f: Feature size (mel bins) + """Dimensions: + - c: Number of audio clips (possibly flattened across audio items) + - b: Number of original audio items + - t: Audio feature length + - f: Feature size (mel bins) """ type: Literal["audio_features"] = "audio_features" @@ -118,13 +117,12 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema): class NanoNemotronVLImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - bnp: Batch size * number of images * (1 + num_patches) - - c: Number of channels (3) - - h: Height of each image patch - - w: Width of each image patch + """Dimensions: + - bn: Batch size * number of images + - bnp: Batch size * number of images * (1 + num_patches) + - c: Number of channels (3) + - h: Height of each image patch + - w: Width of each image patch """ type: Literal["pixel_values"] = "pixel_values" @@ -133,8 +131,7 @@ class NanoNemotronVLImagePixelInputs(TensorSchema): class NanoNemotronVLImagePixelInputsDynamic(TensorSchema): - """ - Dynamic-resolution image inputs. + """Dynamic-resolution image inputs. imgs_sizes: per-image (height, width) in pixels. num_tokens_per_image: per-image number of embedding tokens (post downsample). @@ -147,11 +144,10 @@ class NanoNemotronVLImagePixelInputsDynamic(TensorSchema): class NanoNemotronVLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of images - - f: Total image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - n: Number of images + - f: Total image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] @@ -166,14 +162,13 @@ class NanoNemotronVLImageEmbeddingInputs(TensorSchema): class NanoNemotronVLVideoPixelInputs(TensorSchema): - """ - Dimensions: - - bvf: Batch size * number of videos * num_frames - - bn: Batch size * number of videos - - f: Number of frames - - c: Number of channels (3) - - h: Height of each video frame - - w: Width of each video frame + """Dimensions: + - bvf: Batch size * number of videos * num_frames + - bn: Batch size * number of videos + - f: Number of frames + - c: Number of channels (3) + - h: Height of each video frame + - w: Width of each video frame """ type: Literal["pixel_values_videos"] @@ -184,11 +179,10 @@ class NanoNemotronVLVideoPixelInputs(TensorSchema): class NanoNemotronVLVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: - - n: Number of videos - - f: Total video feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - n: Number of videos + - f: Total video feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["video_embeds"] @@ -613,6 +607,7 @@ def _extract_audio_from_videos( Returns: A 3-tuple of (augmented mm_items, extracted audio items, per-video boolean mask indicating which videos have audio). + """ videos = mm_items.get_items("video", VideoProcessorItems) assert isinstance(videos.metadata, list) @@ -1178,7 +1173,8 @@ def _process_video_input( self, video_input: NanoNemotronVLVideoPixelInputs ) -> tuple[torch.Tensor, ...]: """Process video input and create final embeddings with video content - and indicator tokens.""" + and indicator tokens. + """ T = self.video_temporal_patch_size if T > 1: @@ -1513,9 +1509,7 @@ def forward( return hidden_states def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector=["mlp1", "sound_encoder.projection"], diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ea8a28ebd9d7..e71b7d8e6cf0 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -784,6 +784,7 @@ def get_mamba_state_shape_from_config( - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache - x_cache/dt_cache/B_cache ring-buffer shapes (use_replayssm only) + """ parallel_config = vllm_config.parallel_config cache_config = vllm_config.cache_config diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 12c7b580484b..01b557853116 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -62,8 +62,7 @@ class BartScaledWordEmbedding(VocabParallelEmbedding): - """ - This module overrides VocabParallelEmbedding's + """This module overrides VocabParallelEmbedding's forward by multiplying with embeddings scale. """ @@ -138,12 +137,13 @@ def forward( decoder_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, ) -> torch.Tensor: - r""" - Args: + r"""Args: decoder_hidden_states: torch.Tensor of *decoder* input embeddings. encoder_hidden_states: torch.Tensor of *encoder* input embeddings. + Returns: Decoder layer output torch.Tensor + """ residual = decoder_hidden_states @@ -218,12 +218,13 @@ def forward( class MBartDecoderNoPos(nn.Module): - """ - Transformer decoder consisting of *config.decoder_layers* layers. + """Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BartDecoderLayer`] + Args: config: BartConfig embed_tokens (nn.Embedding): output embedding + """ def __init__( @@ -274,13 +275,13 @@ def forward( inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: - r""" - Args: + r"""Args: decoder_input_ids: Indices of *decoder* input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. encoder_hidden_states: Tensor of encoder output embeddings Returns: Decoder output torch.Tensor + """ if inputs_embeds is None: inputs_embeds = self.embed_tokens(decoder_input_ids) @@ -339,12 +340,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class NemotronParsePixelInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - c: Number of channels (3) - - h: Height - - w: Width + """Dimensions: + - b: Batch size + - c: Number of channels (3) + - h: Height + - w: Width """ type: Literal["pixel_values"] @@ -630,14 +630,15 @@ def forward( encoder_outputs: list[torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: - r""" - Args: + r"""Args: input_ids: torch.Tensor of *decoder* input token ids. positions: torch.Tensor of *decoder* position indices. encoder_outputs: List of encoder output tensors (vision embeddings). During profiling, this may be None or empty. + Returns: Output torch.Tensor + """ inputs_embeds = None if encoder_outputs: diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py index 818d8191971c..bf05848858d3 100644 --- a/vllm/model_executor/models/nemotron_vl.py +++ b/vllm/model_executor/models/nemotron_vl.py @@ -387,9 +387,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="mlp1", @@ -467,8 +465,7 @@ def get_hf_processor(self, **kwargs: object) -> LlamaNemotronVLEmbedProcessor: dummy_inputs=BaseInternVLDummyInputsBuilder[LlamaNemotronVLEmbedProcessingInfo], ) class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling): - """ - LlamaNemotronVL model for embeddings. + """LlamaNemotronVL model for embeddings. Inherits from LlamaNemotronVLChatModel and specializes it for embedding tasks: - Uses SigLIP vision encoder instead of C-RADIO diff --git a/vllm/model_executor/models/openpangu_mtp.py b/vllm/model_executor/models/openpangu_mtp.py index 3a04ccdff5be..d25efd8c0369 100644 --- a/vllm/model_executor/models/openpangu_mtp.py +++ b/vllm/model_executor/models/openpangu_mtp.py @@ -235,8 +235,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py index 10b72f929216..8b85d0371b99 100644 --- a/vllm/model_executor/models/openpangu_vl.py +++ b/vllm/model_executor/models/openpangu_vl.py @@ -1135,9 +1135,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="visual.merger.", @@ -1306,9 +1304,7 @@ def rescale_and_normalize( image_std: float | list[float], dtype: torch.dtype = torch.bfloat16, ) -> "torch.Tensor": - """ - Rescale and normalize images. - """ + """Rescale and normalize images.""" image_mean, image_std, do_rescale = _fuse_mean_std_and_rescale_factor( do_normalize=do_normalize, image_mean=image_mean, diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py index ecbfb8a54c27..77099ffb6efc 100644 --- a/vllm/model_executor/models/openvla.py +++ b/vllm/model_executor/models/openvla.py @@ -73,12 +73,11 @@ def _get_num_image_tokens(image_size: int) -> int: class OpenVLAImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (6) - - h: Height - - w: Width + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (6) + - h: Height + - w: Width """ type: Literal["pixel_values"] = "pixel_values" diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index 10b0eb60eea3..f5065659b434 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -200,7 +200,7 @@ def encode(self, pixel_values: torch.Tensor) -> torch.Tensor: return features def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - """[BatchSize, ImageShape] -> [BatchSize, Token, VocabSize]""" + """[BatchSize, ImageShape] -> [BatchSize, Token, VocabSize].""" features = self.encode(pixel_values) logits = self.head(features) tokens = self.tokenize(logits) @@ -217,13 +217,12 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: class OvisImagePatchInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of images * number of patches - - h: Height of each patch - - w: Width of each patch - - patch_indicators: Batch size * (number of patches + 1) - - bn: Batch size * number of images + """Dimensions: + - bnp: Batch size * number of images * number of patches + - h: Height of each patch + - w: Width of each patch + - patch_indicators: Batch size * (number of patches + 1) + - bn: Batch size * number of images """ type: Literal["image_patches"] @@ -331,8 +330,7 @@ def image_indicators_to_visual_tokens( self, image_indicators: list[int], ) -> list[int]: - """ - Filter image indicators placeholders and convert them to corresponding + """Filter image indicators placeholders and convert them to corresponding tokens in visual tokenizer. For example, [-301, -300, -302, -300, -303, -300, -304, -300, -305] should return [vocab_size-1, vocab_size-2, ..., vocab_size-5] diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py index c55f4805377a..607dcff894ce 100644 --- a/vllm/model_executor/models/ovis2_5.py +++ b/vllm/model_executor/models/ovis2_5.py @@ -49,12 +49,11 @@ class Ovis2_5ImagePatchInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of images * number of patches - - patch_size: patch_size_x * patch_size_y * num_channels - - patch_indicators: Batch size * (number of patches + 1) - - bn: Batch size * number of images + """Dimensions: + - bnp: Batch size * number of images * number of patches + - patch_size: patch_size_x * patch_size_y * num_channels + - patch_indicators: Batch size * (number of patches + 1) + - bn: Batch size * number of images """ type: Literal["image_patches"] @@ -66,12 +65,11 @@ class Ovis2_5ImagePatchInputs(TensorSchema): class Ovis2_5VideoPatchInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of videos * number of patches - - patch_size: patch_size_x * patch_size_y * num_channels - - patch_indicators: Batch size * (number of patches + 1) - - bn: Batch size * number of videos + """Dimensions: + - bnp: Batch size * number of videos * number of patches + - patch_size: patch_size_x * patch_size_y * num_channels + - patch_indicators: Batch size * (number of patches + 1) + - bn: Batch size * number of videos """ type: Literal["video_patches"] @@ -83,9 +81,7 @@ class Ovis2_5VideoPatchInputs(TensorSchema): class VisualTokenizer(torch.nn.Module): - """ - VIT - """ + """VIT.""" def __init__( self, @@ -320,8 +316,7 @@ def visual_indicators_to_visual_tokens( self, visual_indicators: list[int], ) -> list[int]: - """ - Filter image indicators placeholders and convert them to corresponding + """Filter image indicators placeholders and convert them to corresponding tokens in visual tokenizer. """ hf_config = self.info.get_hf_config() diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 4a20531b76dd..d2587b020355 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -99,7 +99,6 @@ def smart_resize( 3. The aspect ratio of the image is maintained as closely as possible. """ - if height < factor: width = round((width * factor) / height) height = factor @@ -1015,14 +1014,14 @@ def compute_logits( def iter_mm_grid_thw( self, mm_features: list[MultiModalFeatureSpec] ) -> Iterator[tuple[int, int, int, int, float]]: - """ - Iterate over multimodal features and yield grid information. + """Iterate over multimodal features and yield grid information. Args: mm_features: List of multimodal feature specifications Yields: Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image + """ spatial_merge_size = self.config.vision_config.spatial_merge_size tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0) diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py index 8e3524eeec4f..c1eafa2f09b1 100644 --- a/vllm/model_executor/models/paligemma.py +++ b/vllm/model_executor/models/paligemma.py @@ -57,12 +57,11 @@ class PaliGemmaImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height - - w: Width + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height + - w: Width """ type: Literal["pixel_values"] = "pixel_values" @@ -70,11 +69,10 @@ class PaliGemmaImagePixelInputs(TensorSchema): class PaliGemmaImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ifs: Image feature size - - hs: Hidden size (must match language model backbone) + """Dimensions: + - bn: Batch size * number of images + - ifs: Image feature size + - hs: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" diff --git a/vllm/model_executor/models/parakeet.py b/vllm/model_executor/models/parakeet.py index ba7b96b8244a..2dd527bb65b0 100644 --- a/vllm/model_executor/models/parakeet.py +++ b/vllm/model_executor/models/parakeet.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Modules below used for the audio encoder component in: models/nano_nemotron_vl.py -""" +"""Modules below used for the audio encoder component in: models/nano_nemotron_vl.py.""" from collections.abc import Iterable from functools import cache diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 8b7beb969024..818109734236 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -88,8 +88,7 @@ def _normalize_expert_bias( class Param2MoEAttention(nn.Module): - """ - Grouped-Query Attention (GQA) for Param2MoE. + """Grouped-Query Attention (GQA) for Param2MoE. Notable differences from a vanilla GQA layer: * The checkpoint fuses Q, K, V into a single ``query_key_value`` weight. @@ -252,8 +251,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Param2MoEMoEBlock(nn.Module): - """ - Mixture-of-Experts block for Param2MoE. + """Mixture-of-Experts block for Param2MoE. Routing: * Sigmoid scoring (config.score_function = "sigmoid") @@ -371,8 +369,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class Param2MoEDecoderLayer(nn.Module): - """ - Single transformer decoder block. + """Single transformer decoder block. Dense for the first ``first_k_dense_replace`` layers; MoE thereafter. """ @@ -605,8 +602,7 @@ def update_physical_experts_metadata( class Param2MoEForCausalLM( nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts ): - """ - vLLM-native Param2MoE CausalLM. + """vLLM-native Param2MoE CausalLM. Uses Grouped-Query Attention (GQA) with a Sigmoid-scored, grouped-topk Mixture-of-Experts MLP. diff --git a/vllm/model_executor/models/phi3.py b/vllm/model_executor/models/phi3.py index 56c8755123d3..d641f89ffa73 100644 --- a/vllm/model_executor/models/phi3.py +++ b/vllm/model_executor/models/phi3.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from llama.py -"""Inference-only Phi3 model code inherit from Llama.py""" +"""Inference-only Phi3 model code inherit from Llama.py.""" from vllm.model_executor.models.llama import LlamaForCausalLM diff --git a/vllm/model_executor/models/phi3v.py b/vllm/model_executor/models/phi3v.py index 8cb387347d2f..17477e8de19d 100644 --- a/vllm/model_executor/models/phi3v.py +++ b/vllm/model_executor/models/phi3v.py @@ -119,13 +119,12 @@ def _init_img_processor( class Phi3VImagePixelInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - n: Number of images - - p: Number of patches - - h: Height of each patch - - w: Width of each patch + """Dimensions: + - b: Batch size + - n: Number of images + - p: Number of patches + - h: Height of each patch + - w: Width of each patch """ type: Literal["pixel_values", "image_embeds"] = "pixel_values" @@ -143,12 +142,11 @@ class Phi3VImagePixelInputs(TensorSchema): class Phi3VImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - n: Number of images - - f: Image feature size (e.g., number of tokens per image) - - h: Hidden size (must match language model backbone) + """Dimensions: + - b: Batch size + - n: Number of images + - f: Image feature size (e.g., number of tokens per image) + - h: Hidden size (must match language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -228,8 +226,7 @@ def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor: def forward( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor ) -> torch.FloatTensor: - """ - process image and return vision embeddings. + """Process image and return vision embeddings. pixel_values: (num_images, num_crops, c, h, w) output: (num_images, num_img_tokens, hidden_size) @@ -244,9 +241,7 @@ def forward( return image_features_proj def hd_feature_transform(self, image_features, image_sizes): - """ - image_features: (num_images, num_crops+1, 24*24, 1024) - """ + """image_features: (num_images, num_crops+1, 24*24, 1024).""" assert self.hd_transform_order == "sub_glb", ( f"hd_transform_order `{self.hd_transform_order}` not implemented" ) @@ -304,8 +299,7 @@ def hd_feature_transform(self, image_features, image_sizes): return batch_image_features_proj def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop): - """ - image_features: (num_images*num_crops, 24*24, 1024) + """image_features: (num_images*num_crops, 24*24, 1024) output: (num_images, h_crop*12, w_crop*12, 4096) where h_crop*w_crop == num_crops """ @@ -329,8 +323,7 @@ def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop): return image_features_hd def add_image_newline(self, image_features_hd): - """ - image_features_hd: (num_images, h_crop*12, w_crop*12, 4096) + """image_features_hd: (num_images, h_crop*12, w_crop*12, 4096) output: (num_images, (h_crop*12) * (w_crop*12+1), 4096) """ num_images, h, w, hid_dim = image_features_hd.shape diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index 345f58e5ca5e..a08a17c324ca 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -246,15 +246,13 @@ def forward( image_sizes: torch.Tensor, image_attention_mask: torch.Tensor, ) -> list[torch.FloatTensor]: - """ - process image and return vision embeddings. + """Process image and return vision embeddings. pixel_values: (num_images, num_crops, c, h, w) image_sizes: [[h1, w1], [h2, w2]] image_attention_mask: num_images x num_crops x 32 x 32 output: (num_images, num_img_tokens, hidden_size) """ - # eg # pixel_values: torch.Size([1, 7, 3, 448, 448]) # image_sizes: tensor([[ 896, 1344]], device='cuda:0') @@ -463,16 +461,15 @@ def forward( class Phi4MMImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - p: Number of patches (1 + num_patches) - - c: Number of channels (3) - - h: Height of each image patch - - w: Width of each image patch - - nc: Number of crops - - H_mask: Height of attention mask - - W_mask: Width of attention mask + """Dimensions: + - bn: Batch size * number of images + - p: Number of patches (1 + num_patches) + - c: Number of channels (3) + - h: Height of each image patch + - w: Width of each image patch + - nc: Number of crops + - H_mask: Height of attention mask + - W_mask: Width of attention mask """ type: Literal["pixel_values"] @@ -501,10 +498,9 @@ class Phi4MMImagePixelInputs(TensorSchema): class Phi4MMAudioFeatureInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of audios - - t: Time frames (M) + """Dimensions: + - bn: Batch size * number of audios + - t: Time frames (M) """ type: Literal["audio_features"] @@ -516,12 +512,11 @@ class Phi4MMAudioFeatureInputs(TensorSchema): class Phi4MMAudioEmbeddingInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - n: Number of audios - - f: Audio feature size - - h: Hidden size (must match language model backbone) + """Dimensions: + - b: Batch size + - n: Number of audios + - f: Audio feature size + - h: Hidden size (must match language model backbone) """ type: Literal["audio_embeds"] @@ -543,9 +538,7 @@ def stack_with_pad( tensors: torch.Tensor | list[torch.Tensor], padding_value: int | float = 0, ) -> torch.Tensor: - """ - Stack tensors, padding dimensions that differ across items. - """ + """Stack tensors, padding dimensions that differ across items.""" if isinstance(tensors, torch.Tensor): return tensors @@ -653,8 +646,7 @@ def _compute_num_image_tokens( vit_patch_size: int, token_compression_factor: int = 2, ): - """ - compute the number of tokens an image is expected to take up considering + """Compute the number of tokens an image is expected to take up considering the image encoder architecture and exclude output features containing only padding pixels @@ -776,8 +768,7 @@ def get_image_size_with_most_features(self) -> ImageSize: return ImageSize(height=max_side, width=vit_image_size) def get_audio_num_frames(self, audio_len: int, sr: float) -> int: - """ - Compute the output size of the `extract_features` method. + """Compute the output size of the `extract_features` method. Args: audio_len (int): Length of the input waveform in samples. @@ -787,8 +778,8 @@ def get_audio_num_frames(self, audio_len: int, sr: float) -> int: tuple (int, int): Output size as (T, D), where: T: Number of time frames. D: Number of Mel filterbank bins (80). - """ + """ # Resample to 16000 or 8000 if needed if sr > 16000: resample_ratio = int(sr // 16000) @@ -812,8 +803,7 @@ def get_audio_num_frames(self, audio_len: int, sr: float) -> int: return num_frames def _compute_audio_embed_size(self, audio_frames: int) -> int: - """ - Compute the audio embedding size based on the audio frames and + """Compute the audio embedding size based on the audio frames and compression rate. """ hf_config = self.get_hf_config() @@ -1046,9 +1036,7 @@ def _recompute_cached_prompt_update( dummy_inputs=Phi4MMDummyInputsBuilder, ) class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal): - """ - Implements the Phi-4-multimodal-instruct model in vLLM. - """ + """Implements the Phi-4-multimodal-instruct model in vLLM.""" packed_modules_mapping = { "qkv_proj": [ @@ -1135,8 +1123,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def _parse_and_validate_audio_input( self, **kwargs: object ) -> Phi4MMAudioInputs | None: - """ - Parse and validate the audio input to the model. This handles both + """Parse and validate the audio input to the model. This handles both audio features and audio embeddings, but only the former is used for now. @@ -1145,6 +1132,7 @@ def _parse_and_validate_audio_input( Returns: Optional[Phi4MMAudioInputs]: Parsed and validated audio inputs. + """ audio_features = kwargs.pop("input_audio_embeds", None) audio_embeds = kwargs.pop("audio_embeds", None) @@ -1166,16 +1154,18 @@ def _parse_and_validate_audio_input( def _process_audio_input( self, audio_input: Phi4MMAudioInputs, audio_projection_mode: str ) -> NestedTensors: - """ - Create the audio embeddings from the audio input, where the audio input + """Create the audio embeddings from the audio input, where the audio input is pairs of audio features and audio embed lengths. The audio input is created by `input_mapper_for_phi4mm_audio`. Args: audio_input (Phi4MMAudioInputs): Audio input. + audio_projection_mode (str): Projection mode to use, selecting + which audio projector weights are applied. Returns: NestedTensors: Audio embeddings + """ if audio_input["type"] == "audio_embeds": return audio_input["data"] @@ -1315,9 +1305,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="model.", connector=["audio_projection_for_vision", "audio_projection"], diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 84abee33c0b9..8f6df6a548e1 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -136,6 +136,7 @@ class ConformerEncoderLayer(nn.Module): 1 = typical Multi-Head Attention, 1 < attn_group_sizes < attention_heads = Grouped-Query Attention attn_group_sizes = attention_heads = Multi-Query Attention + """ def __init__( @@ -232,6 +233,7 @@ def forward( mask: mask for x (batch, max_time_in) relative_attention_bias: bias added to attention logits w.r.t. relative positions (1, n_head, time1, time2) + """ x = x + 0.5 * self.feed_forward_in(x) norm_x = self.layer_norm_att(x) @@ -254,7 +256,7 @@ def forward( class TransformerEncoderBase(abc.ABC, nn.Module): - """The Base class for Transformer based encoders + """The Base class for Transformer based encoders. Please set causal = True in streaming model Args: @@ -459,7 +461,6 @@ def _chunk_size_selection( left_chunk: int | list[int] | None = None, ) -> tuple[int, int]: """If chunk size is a list, we will randomly select a chunk size.""" - if chunk_size is None: chunk_size = self.chunk_size if left_chunk is None: @@ -595,7 +596,7 @@ def forward_embeddings( torch.Tensor, ] ): - """Forwarding the inputs through the top embedding layers + """Forwarding the inputs through the top embedding layers. Args: xs_pad: torch.Tensor @@ -606,6 +607,7 @@ def forward_embeddings( non-causal layers left_chunk_nc: (optional, default is None) # of left chunks for non-causal layers + """ # pylint: disable=R0915 # get new lens. @@ -975,13 +977,14 @@ def calculate_hs_mask( def forward( self, xs_pad: torch.Tensor, masks: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Conformer Forward function + """Conformer Forward function. Args: xs_pad: torch.Tensor input tensor masks: torch.Tensor post-embedding input lengths + """ xs_pad = self.encoder_embedding(xs_pad) input_tensor, pos_k, pos_v, hs_mask, masks = self.forward_embeddings( @@ -1065,7 +1068,7 @@ def forward( class WindowQformer(nn.Module): - """Window-level Qformer""" + """Window-level Qformer.""" def __init__( self, @@ -1107,7 +1110,7 @@ def forward( mask: torch.Tensor | None, embed_len: int | None = None, ) -> tuple[torch.Tensor, int | None]: - """forward decoder""" + """Forward decoder.""" # audio_embed: N x T x D => N x D x T audio_embed = audio_embed.transpose(1, 2) @@ -1265,9 +1268,9 @@ def get_audio_features( audio_attention_mask: torch.Tensor | None = None, audio_projection_mode: str = "speech", ) -> torch.Tensor: - """ - arguments: - input_embeds: audio features (B, T, D) B: num audios in a sequence + """arguments: + input_embeds: audio features (B, T, D) B: num audios in a sequence + """ if self.freeze_audio_processor: with torch.no_grad(): @@ -1319,12 +1322,12 @@ def forward( audio_attention_mask: torch.Tensor | None = None, audio_projection_mode: str = "speech", ) -> torch.Tensor: - """ - arguments: + """arguments: audio_features: audio features (T, D) - returns: + Returns: audio_embeds: audio embeddings (num_audio_tokens, hidden_dim) + """ audio_embeds = self.get_audio_features( audio_features.unsqueeze(0), diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py index 08409af4e2f7..bd7ec9c3444d 100644 --- a/vllm/model_executor/models/phi4mm_utils.py +++ b/vllm/model_executor/models/phi4mm_utils.py @@ -13,7 +13,7 @@ class BlockBase(nn.Module): - """Block abstract module""" + """Block abstract module.""" def __init__(self, input_size: int, output_size: int) -> None: super().__init__() @@ -22,13 +22,14 @@ def __init__(self, input_size: int, output_size: int) -> None: def get_activation(name: str = "relu") -> torch.nn.Module: - """Select an activation function by name + """Select an activation function by name. Args: name: str activation function name, one of ["relu", "gelu", "swish", "sigmoid"], default "relu". + """ name = name.lower() if name == "relu": @@ -48,8 +49,7 @@ def get_activation(name: str = "relu") -> torch.nn.Module: def adaptive_enc_mask( x_len: int, chunk_start_idx: list[int], left_window: int = 0, right_window: int = 0 ) -> torch.Tensor: - """ - The function is very important for Transformer Transducer Streaming mode + """The function is very important for Transformer Transducer Streaming mode Args: x_len: sequence length chunk_start_idx: first idx of each chunk, such as [0,18,36,48]. @@ -57,7 +57,8 @@ def adaptive_enc_mask( left_window: how many left chunks can be seen right_window: how many right chunks can be seen. It is used for chunk overlap model. - Returns: + + Returns: mask (torch.Tensor): a mask tensor for streaming model Torch 1.0.1 tensor([[1., 1., 0., 0.], @@ -67,6 +68,7 @@ def adaptive_enc_mask( tensor([[True., True., False., False.], [False., True., True., False.], [False., False., True., True.]]) + """ chunk_start_idx = torch.Tensor( chunk_start_idx @@ -97,7 +99,7 @@ def adaptive_enc_mask( class GLU(nn.Module): - """Implement Gated Linear Unit (GLU) module""" + """Implement Gated Linear Unit (GLU) module.""" def __init__(self, dim: int = -1, act_name: str = "sigmoid") -> None: super().__init__() @@ -184,9 +186,9 @@ def __init__( self.b2 = nn.Parameter(torch.zeros(1, output_dim, 1)) def forward(self, x: Tensor) -> Tensor: - """ - Args: - x: input tensor + """Args: + x: input tensor + """ # to be consistent with GLULinear, we assume the input always has the # #channel (#dim) in the last dimension of the tensor, so need to @@ -273,10 +275,9 @@ def __init__( self.depthwise_seperable_out_channel = depthwise_seperable_out_channel def forward(self, x: Tensor) -> Tensor: - """ + """Args: + x: input tensor - Args: - x: input tensor """ x = self.dw_conv(x) if self.depthwise_seperable_out_channel != 0: @@ -341,6 +342,7 @@ class ConvModule(nn.Module): or onnx export. Typically this is set by the export program or the decoder program, and it isn't present in your config file. default False + """ def __init__( @@ -406,8 +408,7 @@ def __init__( self.ln2 = nn.Linear(input_dim * depthwise_multiplier, input_dim) def _add_ext_pw_layer(self) -> None: - """ - This function is an extension of __init__ function + """This function is an extension of __init__ function and dedicated to the convolution module creation of the conformer. """ @@ -471,6 +472,7 @@ def forward(self, x: Tensor) -> Tensor: Args: x: input tensor. + """ x = self.layer_norm(x) @@ -519,7 +521,7 @@ def forward(self, x: Tensor) -> Tensor: class GLULinear(nn.Module): - """Linear + GLU module + """Linear + GLU module. Args: input_dim: int @@ -531,6 +533,7 @@ class GLULinear(nn.Module): default "sigmoid" (swish function). bias_in_glu: bool, optional If True, the addtive bias is added. Default False. + """ def __init__( @@ -545,10 +548,11 @@ def __init__( self.glu_act = GLU(-1, glu_type) def forward(self, x: Tensor) -> Tensor: - """GLULinear forward + """GLULinear forward. Args: x: input tensor. + """ x = self.linear(x) return self.glu_act(x) @@ -572,6 +576,7 @@ class FeedForward(nn.Module): sigmoid activation is only used with "glu_in_fnn=True", default "sigmoid". bias_in_glu: bool, optional + """ def __init__( @@ -600,6 +605,7 @@ def forward(self, x: Tensor) -> Tensor: Args: x: input tensor. + """ out = self.net(self.layer_norm(x)) @@ -630,8 +636,7 @@ def _pre_hook( class T5RelativeAttentionLogitBias(nn.Module): - """ - This module implements the relative position bias described in Section + """This module implements the relative position bias described in Section 2.1 of the T5 paper: https://arxiv.org/pdf/1910.10683.pdf The Huggingface implementation is used as a reference @@ -671,6 +676,7 @@ class T5RelativeAttentionLogitBias(nn.Module): Whether to use symmetric or asymmetric biases. symmetric=False uses 2x number of bias params to distinguish L->R from R->L. This was found to be better for the encoder. + """ def __init__( @@ -799,6 +805,7 @@ def extend_pe(self, x: torch.Tensor) -> None: Args: x: input tensor + """ if self.pe is not None and self.pe.size(1) >= x.size(1): if self.pe.dtype != x.dtype or self.pe.device != x.device: @@ -841,6 +848,7 @@ class MeanVarianceNormLayer(nn.Module): Args: input_size: int layer input size. + """ def __init__(self, input_size: int) -> None: @@ -850,17 +858,17 @@ def __init__(self, input_size: int) -> None: self.global_invstd = nn.Parameter(torch.ones(input_size)) def forward(self, input_: Tensor) -> Tensor: - """MeanVarianceNormLayer Forward + """MeanVarianceNormLayer Forward. Args: input_: input tensor. + """ return (input_ - self.global_mean) * self.global_invstd class CausalConv1D(nn.Conv1d): - """ - A causal version of nn.Conv1d where each step would have limited access to + """A causal version of nn.Conv1d where each step would have limited access to locations on its right or left All arguments are the same as nn.Conv1d except padding. @@ -956,8 +964,7 @@ def forward( class CausalConv2D(nn.Conv2d): - """ - A causal version of nn.Conv2d where each location in the 2D matrix would + """A causal version of nn.Conv2d where each location in the 2D matrix would have no access to locations on its right or down All arguments are the same as nn.Conv2d except padding which should be set as None @@ -1041,6 +1048,7 @@ class NemoConvSubsampling(torch.nn.Module): activation (Module): activation function, default is nn.ReLU() is_causal (bool): whether to use causal Conv1/2D, where each step will have limited access to locations on its right or left + """ def __init__( @@ -1338,8 +1346,7 @@ def get_streaming_cache_size(self) -> list[int]: return [0, self.subsampling_factor + 1] def forward(self, x: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor | None]: - """ - Forward method for NeMo subsampling. + """Forward method for NeMo subsampling. Args: x: input tensor @@ -1350,6 +1357,7 @@ def forward(self, x: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor | None time_reduction_factor, feat_out) pad_mask: tensor of padded hidden state sequences (B, 1, T // time_reduction_factor) + """ x = x.unsqueeze(1) if self.conv2d_subsampling else x.transpose(1, 2) @@ -1427,7 +1435,7 @@ def reset_parameters(self) -> None: torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale) def conv_split_by_batch(self, x: Tensor) -> tuple[Tensor, bool]: - """Tries to split input by batch, run conv and concat results""" + """Tries to split input by batch, run conv and concat results.""" b, _, _, _ = x.size() if b == 1: # can't split if batch size is 1 return x, False @@ -1454,7 +1462,8 @@ def conv_split_by_batch(self, x: Tensor) -> tuple[Tensor, bool]: def conv_split_by_channel(self, x: Tensor) -> Tensor: """For dw convs, tries to split input by time, run conv and concat - results""" + results + """ x = self.conv[0](x) # full conv2D x = self.conv[1](x) # activation @@ -1493,8 +1502,7 @@ def conv_split_by_channel(self, x: Tensor) -> Tensor: def channel_chunked_conv( self, conv: torch.nn.Module, chunk_size: int, x: Tensor ) -> Tensor: - """Performs channel chunked convolution""" - + """Performs channel chunked convolution.""" ind = 0 out_chunks = [] for chunk in torch.split(x, chunk_size, 1): @@ -1571,14 +1579,14 @@ def calc_length_int( #### multihead attention starts here class AttModule(nn.Module): - """Attention abstraction module""" + """Attention abstraction module.""" def __init__(self) -> None: super().__init__() self.export_mode = False def set_export(self, mode: bool = True) -> None: - """set the export mode""" + """Set the export mode.""" self.export_mode = mode def forward( @@ -1588,13 +1596,14 @@ def forward( pos_emb: Tensor | None = None, att_mask: Tensor | None = None, ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]: - """AttModule forward + """AttModule forward. Args: x: input tensor. memory: memory tensor. pos_emb: positional encoder embedding. att_mask: attention mask tensor. + """ return x, memory, pos_emb, att_mask @@ -1642,6 +1651,7 @@ class MultiHeadedAttention(nn.Module): if group_size > 1: GQA if group_size = 1: MHA if group_size = n_head: MQA + """ inv_sqrt_d_k: torch.jit.Final[float] @@ -1719,6 +1729,7 @@ def forward( relative_attention_bias: bias added to attention logits w.r.t. relative positions (1, n_head, time1, time2) + """ n_batch = query.size(0) @@ -1815,7 +1826,7 @@ def forward( class MultiSequential(torch.nn.Sequential): - """Multi-input multi-output torch.nn.Sequential""" + """Multi-input multi-output torch.nn.Sequential.""" @torch.jit.ignore def forward(self, *args) -> tuple: @@ -1834,6 +1845,7 @@ def get_offset(input_layer: str, time_reduction: int) -> int: time_reduction: time reduction factor for downsampling a feature Returns: int: offset + """ if input_layer in ("conv2d", "nemo_conv") and time_reduction == 4: return 3 @@ -1845,13 +1857,14 @@ def get_offset(input_layer: str, time_reduction: int) -> int: def unfold_tensor(xs_pad: Tensor, max_seq_len: int) -> Tensor: - """ - For a given tensor with shape of (N, T, D), if sequence length T is + """For a given tensor with shape of (N, T, D), if sequence length T is longer than max_seq_len, this function unfold it to a (NT', max_seq_len, D) where T' is T // max_seq_len. + Args: xs_pad: input tensor with shape (N, T, D) max_seq_len: maximum sequence length + """ _, _, D = xs_pad.shape xs_pad = xs_pad.transpose(-1, -2) # convert to N, D, T diff --git a/vllm/model_executor/models/phi4siglip.py b/vllm/model_executor/models/phi4siglip.py index ff880388698c..3c36b1f3eee5 100644 --- a/vllm/model_executor/models/phi4siglip.py +++ b/vllm/model_executor/models/phi4siglip.py @@ -224,11 +224,10 @@ def get_replacement(item_idx: int): class Phi4SiglipImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - d: Max number of patches (padded across images in the batch) - - fd: Features per patch (patch_size * patch_size * channels) + """Dimensions: + - bn: Batch size * number of images + - d: Max number of patches (padded across images in the batch) + - fd: Features per patch (patch_size * patch_size * channels) """ type: Literal["pixel_values"] = "pixel_values" diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index c948fb5ad72e..42812add9b81 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -157,8 +157,7 @@ def _is_layer_none_or_staged(layer: nn.Module) -> bool: class PixtralImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - bn: Batch size * number of images - c: Number of channels (3) - h: Height of each image @@ -688,8 +687,7 @@ class VisionEncoderArgs: def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - """ - freqs_cis: complex - (seq_len, head_dim / 2) + """freqs_cis: complex - (seq_len, head_dim / 2) x: complex - (bsz, seq_len, head_dim / 2) """ ndim = x.ndim @@ -708,9 +706,8 @@ def precompute_freqs_cis_2d( width: int, theta: float, ) -> torch.Tensor: - """ - freqs_cis: 2D complex tensor of shape (height, width, dim // 2) - to be indexed by (height, width) position tuples + """freqs_cis: 2D complex tensor of shape (height, width, dim // 2) + to be indexed by (height, width) position tuples """ # (dim / 2) frequency bases freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) @@ -1018,13 +1015,14 @@ def forward( self, images: list[torch.Tensor], ) -> torch.Tensor: - """ - Args: + """Args: images: list of N_img images of variable sizes, each of shape (C, H, W) + Returns: image_features: tensor of token features for all tokens of all images of shape (N_toks, D) + """ # pass images through initial convolution independently patch_embeds_list = [ @@ -1082,9 +1080,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class PatchMerger(nn.Module): - """ - Learned merging of spatial_merge_size ** 2 patches - """ + """Learned merging of spatial_merge_size ** 2 patches.""" def __init__( self, @@ -1124,8 +1120,7 @@ def permute( x: torch.Tensor, image_sizes: list[tuple[int, int]], ) -> torch.Tensor: - """ - Args: + """Args: x: (N, D) where N is flattened and concatenated patch tokens for all images image_sizes: list of tuple of (height, width) in tokens for @@ -1134,8 +1129,8 @@ def permute( image_features: reorders patch tokens so each grid of (spatial_merge_size, spatial_merge_size) is contiguous. now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2) - """ + """ sub_grids = get_sub_grids( x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size ) # list of [d x sub_grid_size x sub_grid_size x n_patches] @@ -1509,8 +1504,7 @@ def forward( select_layers: list[int] | None = None, feature_select_strategy: VisionFeatureSelectStrategy | None = None, ) -> tuple[torch.Tensor, ...]: - """ - Args: + """Args: pixel_values: Each image to be processed will be a separate tensor in pixel_values. This means it will be a list of tensors because multiple requests batched can have multiple images, @@ -1522,6 +1516,7 @@ def forward( Returns: image_features: tensor of token features for all tokens of all images of shape (N_toks, D) + """ # pass images through initial convolution independently patch_embeds_list = [ diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 26b69bd6dd66..bb6b153eddec 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -139,8 +139,7 @@ def check_interleaved_audio_video( num_video: int, num_audio: int, ) -> bool: - """ - Check if video and audio positions are interleaved in any per-video span. + """Check if video and audio positions are interleaved in any per-video span. For use_audio_in_video=True, each video placeholder is expanded into one local span containing only video/audio pad tokens, bounded by non-pad @@ -187,8 +186,7 @@ def merge_interleaved_embeddings( is_audio: torch.Tensor, is_multimodal: torch.Tensor, ) -> torch.Tensor: - """ - Merge embeddings for interleaved audio-in-video sequences. + """Merge embeddings for interleaved audio-in-video sequences. When use_audio_in_video=True, video and audio tokens are interleaved in the token sequence, but embeddings are provided as separate contiguous @@ -206,6 +204,7 @@ def merge_interleaved_embeddings( Returns: The merged inputs_embeds tensor with multimodal embeddings scattered to their correct positions. + """ from vllm.multimodal.utils import get_mm_embedding_modalities @@ -254,12 +253,11 @@ def _merge_embedding_group( class Qwen2_5OmniAudioFeatureInputs(TensorSchema): - """ - Dimensions: - - na: Number of audios - - nmb: Number of mel bins - - msl: Maximum sequence length - - tsl: Total sequence length + """Dimensions: + - na: Number of audios + - nmb: Number of mel bins + - msl: Maximum sequence length + - tsl: Total sequence length """ type: Literal["audio_features"] @@ -514,8 +512,7 @@ def _derive_audio_from_video_placeholders( placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], mm_prompt_updates: MultiModalPromptUpdates, ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: - """ - Helper to derive audio placeholders from video placeholders when + """Helper to derive audio placeholders from video placeholders when use_audio_in_video=True. """ if "video" not in placeholders: @@ -577,9 +574,7 @@ def _maybe_apply_prompt_updates( mm_kwargs: MultiModalKwargsItems, mm_prompt_updates: MultiModalPromptUpdates, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: - """ - Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. - """ + """Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.""" mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) @@ -644,7 +639,6 @@ def omni_get_updates_use_audio_in_video( <|video_bos|><|VIDEO|><|video_eos|> => <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|> """ - audio_token_id = thinker_config.audio_token_index video_token_id = thinker_config.video_token_index audio_start_token_id = thinker_config.audio_start_token_id @@ -800,9 +794,7 @@ def _apply_hf_processor_main( mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: - """ - Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. - """ + """Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.""" mm_counts = mm_items.get_all_counts() use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) @@ -1171,8 +1163,7 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: def _get_audio_for_video_mapping( self, mm_features: list[MultiModalFeatureSpec] ) -> tuple[dict[int, int], set[int]]: - """ - Map video offset -> paired audio_feature_length for use_audio_in_video. + """Map video offset -> paired audio_feature_length for use_audio_in_video. When use_audio_in_video=True, audio is interleaved within video chunks. The pairing is based on feature order in mm_features. @@ -1180,6 +1171,7 @@ def _get_audio_for_video_mapping( Returns: Tuple of (video_offset -> audio_feature_length mapping, set of paired audio offsets to skip) + """ videos_with_audio = [ f @@ -1207,8 +1199,7 @@ def _compute_audio_token_count(self, audio_feature_length: int) -> int: def iter_mm_features( self, mm_features: list[MultiModalFeatureSpec] ) -> Iterator[tuple[int, str, dict[str, Any]]]: - """ - Iterate over multimodal features sorted by position offset. + """Iterate over multimodal features sorted by position offset. Yields: (offset, modality, feature_data) where feature_data contains: - image: {"grid_t", "grid_h", "grid_w", "t_factor"} @@ -1278,8 +1269,7 @@ def iter_mm_features( def _compute_interleaved_positions( self, start_idx: int, data: dict[str, Any] ) -> tuple[np.ndarray, int]: - """ - Compute positions for interleaved video+audio chunks. + """Compute positions for interleaved video+audio chunks. Returns: (position_ids, total_token_count) """ @@ -1360,8 +1350,7 @@ def get_mrope_input_positions( input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - """ - Compute M-RoPE input positions using mm_features directly. + """Compute M-RoPE input positions using mm_features directly. Example for use_audio_in_video case: (V_i are vision position ids, A_i are audio position ids) @@ -1542,9 +1531,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector="merger.", diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 6cfacc1421fc..a3d0ecfcab97 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -161,8 +161,7 @@ def _pad_flashinfer_cu_seqlens_buffer( class Qwen2_5_VLImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - np: Number of patches - ni: Number of images - cps: Number of channels * patch_size * patch_size @@ -188,8 +187,7 @@ class Qwen2_5_VLImagePixelInputs(TensorSchema): class Qwen2_5_VLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - nf: Number of image features - hs: Hidden size - ni: Number of images @@ -222,8 +220,7 @@ class Qwen2_5_VLImageEmbeddingInputs(TensorSchema): class Qwen2_5_VLVideoPixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - np: Number of patches - nv: Number of videos - ctps: Number of channels * temporal_patch_size * patch_size * @@ -262,8 +259,7 @@ class Qwen2_5_VLVideoPixelInputs(TensorSchema): class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - nf: Number of video features - hs: Hidden size - nv: Number of videos @@ -916,8 +912,8 @@ def prepare_encoder_metadata( cu_window_seqlens (needed for CUDA graph capture to cover worst-case replay scenarios). device: Device to place tensors on. Defaults to self.device. - """ + """ if device is None: device = self.device metadata: dict[str, torch.Tensor] = {} @@ -1261,14 +1257,14 @@ class Qwen2_5_VLForConditionalGeneration( def iter_mm_grid_thw( self, mm_features: list[MultiModalFeatureSpec] ) -> Iterator[tuple[int, int, int, int, float]]: - """ - Iterate over multimodal features and yield grid information. + """Iterate over multimodal features and yield grid information. Args: mm_features: List of multimodal feature specifications Yields: Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image + """ spatial_merge_size = self.config.vision_config.spatial_merge_size tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0) @@ -1462,8 +1458,7 @@ def _postprocess_image_embeds_evs( image_embeds_split: tuple[torch.Tensor, ...], image_input: Qwen2_5_VLImageInputs, ) -> tuple[torch.Tensor, ...]: - """ - Append mrope positions for each for images. + """Append mrope positions for each for images. This is necessary to recover correct mrope positions after video pruning @@ -1476,6 +1471,7 @@ def _postprocess_image_embeds_evs( Tuple of image embeddings for each image item. Resulting embeddings will have extra 4 channels for computed mrope positions. + """ merge_size = self.visual.spatial_merge_size grid_thw = image_input["image_grid_thw"] @@ -1522,8 +1518,7 @@ def _postprocess_video_embeds_evs( video_embeds_split: tuple[torch.Tensor, ...], video_input: Qwen2_5_VLVideoInputs, ) -> tuple[torch.Tensor, ...]: - """ - Prunes video embeddings via Efficient Video Sampling (EVS) + """Prunes video embeddings via Efficient Video Sampling (EVS) and then appends mrope positions for each retained embeddings Args: @@ -1534,6 +1529,7 @@ def _postprocess_video_embeds_evs( Tuple of video embeddings for each video item. Resulting embeddings will have extra 4 channels for computed mrope positions. + """ grid_thw = video_input["video_grid_thw"] assert grid_thw.ndim == 2 @@ -1584,8 +1580,7 @@ def recompute_mrope_positions( mrope_positions: torch.LongTensor, num_computed_tokens: int, ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]: - """ - Update part of input mrope positions (starting with + """Update part of input mrope positions (starting with num_computed_tokens index). Original mrope_positions are computed for unpruned sequence and becomes incorrect once pruning occurs, so once we prune media tokens we should reflect this in the @@ -1602,6 +1597,7 @@ def recompute_mrope_positions( Returns: Tuple of (multimodal_embeddings, mrope_positions, mrope_position_delta). + """ image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id @@ -2018,8 +2014,10 @@ def forward( batch. **NOTE**: If mrope is enabled (default setting for Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`, otherwise it will be `(seq_len,). - """ + intermediate_tensors: Intermediate tensors from prior forward pass. + inputs_embeds: Optional tensor of input embeddings. + """ if intermediate_tensors is not None: inputs_embeds = None @@ -2042,9 +2040,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="visual.merger.", diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 639db2884706..27c603c866b2 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -69,10 +69,9 @@ # # === Audio Inputs === # class Qwen2AudioFeatureInputs(TensorSchema): - """ - Dimensions: - - na: Number of audios - - nmb: Number of mel bins + """Dimensions: + - na: Number of audios + - nmb: Number of mel bins """ type: Literal["audio_features"] @@ -88,12 +87,11 @@ class Qwen2AudioFeatureInputs(TensorSchema): class Qwen2AudioEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size - - naf: Number of audio features - - hs: Hidden size (must match the hidden size of language model - backbone) + """Dimensions: + - bn: Batch size + - naf: Number of audio features + - hs: Hidden size (must match the hidden size of language model + backbone) """ type: Literal["audio_embeds"] = "audio_embeds" diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index c3761f8893f5..65be51e8b2ad 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -120,8 +120,7 @@ class Qwen2VLImagePixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - np: The total number of patches over each image over each prompt in the batch - ni: Number of images @@ -148,8 +147,7 @@ class Qwen2VLImagePixelInputs(TensorSchema): class Qwen2VLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - nf: Number of image features - hs: Hidden size - ni: Number of images @@ -180,8 +178,7 @@ class Qwen2VLImageEmbeddingInputs(TensorSchema): class Qwen2VLVideoPixelInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - np: The total number of patches over each video over each prompt in the batch - ctps: Number of channels * temporal_patch_size * patch_size * @@ -209,8 +206,7 @@ class Qwen2VLVideoPixelInputs(TensorSchema): class Qwen2VLVideoEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - nf: Number of video features - hs: Hidden size - nv: Number of videos @@ -1208,14 +1204,14 @@ class Qwen2VLForConditionalGeneration( def iter_mm_grid_thw( self, mm_features: list[MultiModalFeatureSpec] ) -> Iterator[tuple[int, int, int, int, float]]: - """ - Iterate over multimodal features and yield grid information. + """Iterate over multimodal features and yield grid information. Args: mm_features: List of multimodal feature specifications Yields: Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image + """ spatial_merge_size = self.config.vision_config.spatial_merge_size tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0) @@ -1715,8 +1711,10 @@ def forward( otherwise it will be `(seq_len,)`. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. - """ + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. + """ if intermediate_tensors is not None: inputs_embeds = None @@ -1739,9 +1737,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector="visual.merger.", diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index f991f05d74d3..9204d5305dee 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -583,8 +583,8 @@ def forward( model. `None` if no videos are passed. - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM. `None` if no videos are passed. - """ + """ if intermediate_tensors is not None: inputs_embeds = None diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 3bea939340e5..14925ea55063 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -618,9 +618,7 @@ def get_mrope_input_positions( return llm_positions, mrope_position_delta def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", tower_model=["audio_tower."], @@ -703,8 +701,7 @@ def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: @classmethod def post_process_output(cls, text: str) -> str: - """ - Post-process Qwen3-ASR raw output to extract clean transcription. + """Post-process Qwen3-ASR raw output to extract clean transcription. The model outputs in format: "language {lang}{transcription}" This method strips the language prefix and asr_text tags. diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index ce3884a10eb8..a807f2d4b5aa 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -69,7 +69,8 @@ def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: def dflash_has_any_non_causal(config: Qwen3Config) -> bool: """Whether the draft needs a non-causal-capable backend, resolved from config - (config mirror of the model's ``get_draft_attn_causal``, usable pre-build).""" + (config mirror of the model's ``get_draft_attn_causal``, usable pre-build). + """ return not all( _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers) ) @@ -154,7 +155,8 @@ class DFlashQwen3Attention(nn.Module): Context KVs are pre-inserted into the KV cache before the forward pass. This layer handles only query tokens via standard attention. - Adapted from Qwen3Attention.""" + Adapted from Qwen3Attention. + """ def __init__( self, @@ -248,7 +250,8 @@ def forward( """DFlash attention assumes that the KV cache is already populated with the context K/V from the target model's hidden states. This forward op computes attention for the query tokens only. - See also: precompute_and_store_context_kv""" + See also: precompute_and_store_context_kv + """ qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) @@ -740,7 +743,8 @@ def get_draft_kv_cache_layer_names(self) -> list[str]: def get_draft_attn_causal(self) -> list[bool]: """Per-layer attention causality, aligned with - get_draft_kv_cache_layer_names.""" + get_draft_kv_cache_layer_names. + """ return [layer.self_attn.causal for layer in self.model.layers] def compute_logits( diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index acc09f50382a..ce800e7de668 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -301,11 +301,11 @@ def forward( cu_seqlens: torch.Tensor, max_seqlen: torch.Tensor | None, ) -> torch.Tensor: - """ - Args: - hidden_states: Input tensor of shape (seq_len, hidden_size) - cu_seqlens: Cumulative sequence lengths - max_seqlen: Maximum sequence length in the batch + """Args: + hidden_states: Input tensor of shape (seq_len, hidden_size) + cu_seqlens: Cumulative sequence lengths + max_seqlen: Maximum sequence length in the batch + """ residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) @@ -1316,9 +1316,7 @@ def _maybe_apply_prompt_updates( mm_kwargs: MultiModalKwargsItems, mm_prompt_updates: MultiModalPromptUpdates, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: - """ - Qwen3-Omni reimplements this function to handle `use_audio_in_video`. - """ + """Qwen3-Omni reimplements this function to handle `use_audio_in_video`.""" mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) @@ -1528,8 +1526,7 @@ def _derive_audio_from_video_placeholders( placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], mm_prompt_updates: MultiModalPromptUpdates, ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: - """ - Helper to derive audio placeholders from video placeholders when + """Helper to derive audio placeholders from video placeholders when use_audio_in_video=True. """ if "video" not in placeholders: @@ -1991,8 +1988,7 @@ def _compute_audio_token_count(self, audio_feature_length: int) -> int: def _get_audio_for_video_mapping( self, mm_features: list[MultiModalFeatureSpec] ) -> tuple[dict[int, int], set[int]]: - """ - Map video offset -> paired audio_feature_length for use_audio_in_video. + """Map video offset -> paired audio_feature_length for use_audio_in_video. When use_audio_in_video=True, audio is interleaved within video. The pairing is based on feature order in mm_features. @@ -2000,6 +1996,7 @@ def _get_audio_for_video_mapping( Returns: Tuple of (video_offset -> audio_feature_length mapping, set of paired audio offsets to skip) + """ videos_with_audio = [ f @@ -2022,8 +2019,7 @@ def _get_audio_for_video_mapping( def iter_mm_features( self, mm_features: list[MultiModalFeatureSpec] ) -> Iterator[tuple[int, str, dict[str, Any]]]: - """ - Iterate over multimodal features sorted by position offset. + """Iterate over multimodal features sorted by position offset. Yields: (offset, modality, feature_data) where feature_data contains: - image: {"grid_t", "grid_h", "grid_w", "t_factor"} @@ -2088,8 +2084,7 @@ def iter_mm_features( def _compute_interleaved_positions( self, start_idx: int, data: dict[str, Any] ) -> tuple[np.ndarray, int]: - """ - Compute positions for interleaved video+audio using Qwen3 token-by-token + """Compute positions for interleaved video+audio using Qwen3 token-by-token interleaving logic. Returns: (position_ids [3, N], total_token_count) @@ -2153,9 +2148,7 @@ def get_speech_to_text_config( @classmethod def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: - """ - Construct a transcription/translation prompt for Qwen3-Omni. - """ + """Construct a transcription/translation prompt for Qwen3-Omni.""" audio = stt_params.audio stt_config = stt_params.stt_config model_config = stt_params.model_config @@ -2315,9 +2308,7 @@ def get_mrope_input_positions( return torch.from_numpy(llm_positions), mrope_position_delta def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", connector="visual.merger", diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 644fd134adce..482fd00e5d9b 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -776,6 +776,7 @@ def prepare_encoder_metadata( instead of computing from cu_seqlens (needed for CUDA graph capture to cover worst-case replay scenarios). device: Device to place tensors on. Defaults to self.device. + """ if device is None: device = self.device @@ -1291,6 +1292,7 @@ def _replace_video_token_placeholders( Returns: Token IDs with every placeholder replaced. + """ result: list[int] = [] repl_idx = 0 @@ -1319,7 +1321,8 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: """Transformers>=5.10 processors override `replace_video_token` to expand only the bare video token, keeping the prompt's outer - `<|vision_start|>`/`<|vision_end|>` markers.""" + `<|vision_start|>`/`<|vision_end|>` markers. + """ mixin_impl = getattr(ProcessorMixin, "replace_video_token", None) proc_impl = getattr(type(hf_processor), "replace_video_token", None) return proc_impl is not None and proc_impl is not mixin_impl @@ -1647,9 +1650,11 @@ def get_video_repl( vision_start_token_id: Token ID for vision start marker vision_end_token_id: Token ID for vision end marker video_token_id: Token ID for video content + select_token_id: Whether to return token ids instead of text Returns: PromptUpdateDetails with full token sequence + """ assert len(timestamps) == len(tokens_per_frame), ( "timestamps and tokens_per_frame must have the same length" @@ -2350,8 +2355,7 @@ def _postprocess_image_embeds_evs( image_embeds_split: tuple[torch.Tensor, ...], image_input: Qwen2_5_VLImageInputs, ) -> tuple[torch.Tensor, ...]: - """ - Append mrope positions for each for images. + """Append mrope positions for each for images. This is necessary to recover correct mrope positions after video pruning @@ -2364,6 +2368,7 @@ def _postprocess_image_embeds_evs( Tuple of image embeddings for each image item. Resulting embeddings will have extra 5 channels for computed mrope positions, consistent with video embeddings. + """ if self.is_multimodal_pruning_enabled: merge_size = self.visual.spatial_merge_size @@ -2393,8 +2398,7 @@ def _postprocess_video_embeds_evs( video_embeds_split: tuple[torch.Tensor, ...], video_input: Qwen2_5_VLVideoInputs, ) -> tuple[torch.Tensor, ...]: - """ - Prunes video embeddings via Efficient Video Sampling (EVS) + """Prunes video embeddings via Efficient Video Sampling (EVS) and then appends mrope positions for each retained embeddings Args: @@ -2405,6 +2409,7 @@ def _postprocess_video_embeds_evs( Tuple of video embeddings for each video item. Resulting embeddings will have extra 5 channels for computed mrope positions, and whether the index corresponds to a video embedding. + """ grid_thw = video_input["video_grid_thw"] assert grid_thw.ndim == 2 @@ -2478,7 +2483,6 @@ def _create_final_video_embeddings( These embeddings will replace the placeholder embeddings to create input_embeds for the LLM. """ - device = video_embeddings.device # Generate video replacement token IDs using get_video_repl @@ -2664,6 +2668,7 @@ def _iter_mm_grid_hw( llm_grid_h: Logical grid height (may not match actual token count with EVS). llm_grid_w: Logical grid width (may not match actual token count with EVS). actual_num_tokens: Actual number of video/image tokens in the placeholder. + """ for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): offset = mm_feature.mm_position.offset @@ -2806,8 +2811,7 @@ def recompute_mrope_positions( mrope_positions: torch.LongTensor, num_computed_tokens: int, ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]: - """ - Update part of input mrope positions (starting with + """Update part of input mrope positions (starting with num_computed_tokens index). Original mrope_positions are computed for unpruned sequence and becomes incorrect once pruning occurs, so once we prune media tokens we should reflect this in the @@ -2825,6 +2829,7 @@ def recompute_mrope_positions( Returns: Tuple of (multimodal_embeddings, mrope_positions, mrope_position_delta). + """ return self._recompute_mrope_positions( input_ids=input_ids, @@ -3031,8 +3036,8 @@ def forward( model. `None` if no videos are passed. - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM. `None` if no videos are passed. - """ + """ if intermediate_tensors is not None: inputs_embeds = None @@ -3068,9 +3073,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model", connector=["visual.merger", "visual.deepstack_merger_list"], diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 47dc8f89d326..9f8ff5ed2bf6 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -227,6 +227,7 @@ def forward_video(self, x: torch.Tensor) -> torch.Tensor: Returns: Embedded patches with temporal compression applied. + """ assert self.temporal_patch_size > 1 T = self.temporal_patch_size @@ -598,7 +599,8 @@ def inter_image_mask_metadata( self, imgs_sizes: list[tuple[int, int]], device: torch.device ) -> MaskMetadata: """Build mask metadata from image pixel sizes. Adds num_skip to each - sequence length (cls/register tokens) to match patch generator output.""" + sequence length (cls/register tokens) to match patch generator output. + """ patch_size = self.patch_generator.patch_size num_skip = self.patch_generator.num_skip @@ -611,7 +613,8 @@ def _inter_image_mask_metadata_from_seq_lens( ) -> MaskMetadata: """Build mask metadata from actual sequence lengths (already including cls/register tokens, i.e. patch_count + num_skip per item). - Use inter_image_mask_metadata() when you only have imgs_sizes.""" + Use inter_image_mask_metadata() when you only have imgs_sizes. + """ assert len(seq_lens) > 0 cu_seqlens = torch.tensor( list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 6fc9f9a08037..cdb13df8f5b1 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Whenever you add an architecture to this page, please also update +"""Whenever you add an architecture to this page, please also update `tests/models/registry.py` with example HuggingFace models for it. """ @@ -928,9 +927,7 @@ def load_model_cls(self) -> type[nn.Module]: @dataclass(frozen=True) class _RegisteredModel(_BaseRegisteredModel): - """ - Represents a model that has already been imported in the main process. - """ + """Represents a model that has already been imported in the main process.""" interfaces: _ModelInfo model_cls: type[nn.Module] @@ -951,9 +948,7 @@ def load_model_cls(self) -> type[nn.Module]: @dataclass(frozen=True) class _LazyRegisteredModel(_BaseRegisteredModel): - """ - Represents a model that has not been imported in the main process. - """ + """Represents a model that has not been imported in the main process.""" module_name: str class_name: str @@ -1018,7 +1013,7 @@ def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None: return None def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None: - """save dictionary json file to cache""" + """Save dictionary json file to cache.""" from vllm.model_executor.model_loader.weight_utils import atomic_writer try: @@ -1128,8 +1123,7 @@ def register_model( model_arch: str, model_cls: type[nn.Module] | str, ) -> None: - """ - Register an external model to be used in vLLM. + """Register an external model to be used in vLLM. `model_cls` can be either: diff --git a/vllm/model_executor/models/roberta.py b/vllm/model_executor/models/roberta.py index 59054b0667dc..6a65b3a7e93f 100644 --- a/vllm/model_executor/models/roberta.py +++ b/vllm/model_executor/models/roberta.py @@ -275,6 +275,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding): Attributes: roberta: An instance of BertModel used for forward operations. _pooler: An instance of Pooler used for pooling operations. + """ is_pooling_model = True diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 2cc8a27f4dcf..e7db1d680130 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -537,6 +537,7 @@ def forward( Returns: Intermediate tensors on non-final stages. On the final stage, normalized hidden states, paired with auxiliary states if captured. + """ if get_pp_group().is_first_rank: if inputs_embeds is not None: diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 6d1fd029624c..8a5b226b25d3 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -73,12 +73,11 @@ class SiglipImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height of each image - - w: Width of each image + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image """ type: Literal["pixel_values"] @@ -399,7 +398,7 @@ def forward( self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, None]: - """Input shape: Batch x Time x Channel""" + """Input shape: Batch x Time x Channel.""" qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) out = self.attn(query_states, key_states, value_states) @@ -778,8 +777,9 @@ def maybe_layer_norm_and_apply_head( """Apply the post layer norm and head if they are enabled, given the last hidden states tensor. - args: + Args: encoder_outputs: The last hidden states from the visual encoder. + """ if self.post_layernorm is not None: encoder_outputs = self.post_layernorm(encoder_outputs) diff --git a/vllm/model_executor/models/siglip2navit.py b/vllm/model_executor/models/siglip2navit.py index c16619612fc3..aabb62f4ab59 100644 --- a/vllm/model_executor/models/siglip2navit.py +++ b/vllm/model_executor/models/siglip2navit.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Implementation of SiglipVisionModel intended to be only used -within a vision language model.""" +within a vision language model. +""" from collections.abc import Iterable @@ -86,17 +87,16 @@ def forward( pixel_values: torch.FloatTensor, grid_thws: torch.LongTensor | None = None, ) -> torch.Tensor: - """ - Args: - pixel_values (`torch.FloatTensor`): - Pixel values of shape ( - num_patches, - num_channels * temporal_patch_size * patch_size * patch_size - ) - grid_thws: (`torch.LongTensor`): - grid shape (num_patches, 3) - """ + """Args: + pixel_values (`torch.FloatTensor`): + Pixel values of shape ( + num_patches, + num_channels * temporal_patch_size * patch_size * patch_size + ) + grid_thws: (`torch.LongTensor`): + grid shape (num_patches, 3) + """ # Apply patch embeddings to already patchified pixel values target_dtype = self.patch_embedding.weight.dtype if isinstance(self.patch_embedding, LinearBase): @@ -173,7 +173,7 @@ def apply_rotary_pos_emb( class Siglip2Attention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -236,8 +236,7 @@ def forward( cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Input shape: Batch x Time x Channel""" - + """Input shape: Batch x Time x Channel.""" seq_length, embed_dim = hidden_states.shape qkv_states, _ = self.qkv_proj(hidden_states) @@ -338,11 +337,11 @@ def forward( cu_seqlens: torch.Tensor, position_embeddings: torch.Tensor, ) -> tuple[torch.FloatTensor]: - """ - Args: - hidden_states: Input tensor of shape (batch, seq_len, embed_dim). - cu_seqlens: Cumulative sequence lengths tensor. - position_embeddings: Position embeddings tensor. + """Args: + hidden_states: Input tensor of shape (batch, seq_len, embed_dim). + cu_seqlens: Cumulative sequence lengths tensor. + position_embeddings: Position embeddings tensor. + """ residual = hidden_states @@ -362,12 +361,12 @@ def forward( class Siglip2Encoder(nn.Module): - """ - Transformer encoder consisting of `config.num_hidden_layers` + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Siglip2EncoderLayer`]. Args: config: PretrainedConfig + """ def __init__( @@ -486,15 +485,15 @@ def forward( inputs_embeds: torch.Tensor, grid_thws: torch.Tensor, ) -> torch.Tensor: - r""" - Args: - inputs_embeds: Input tensor of shape - (batch_size, sequence_length, hidden_size). - Embedded representation of the input tokens. - grid_thws: Grid tensor of shape (num_patches, 3) - containing grid dimensions. - Whether or not to return a [`~utils.ModelOutput`] instead of - a plain tuple. + r"""Args: + inputs_embeds: Input tensor of shape + (batch_size, sequence_length, hidden_size). + Embedded representation of the input tokens. + grid_thws: Grid tensor of shape (num_patches, 3) + containing grid dimensions. + Whether or not to return a [`~utils.ModelOutput`] instead of + a plain tuple. + """ rotary_pos_emb = self.rot_pos_emb(grid_thws) window_index, cu_window_seqlens = self.get_window_index(grid_thws) @@ -575,10 +574,9 @@ def forward( pixel_values: torch.FloatTensor, grid_thws: torch.LongTensor, ) -> torch.Tensor: - r""" - spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): - Tensor containing the spatial dimensions (height, width) - of the input images. + r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): + Tensor containing the spatial dimensions (height, width) + of the input images. """ hidden_states = self.embeddings(pixel_values, grid_thws) diff --git a/vllm/model_executor/models/skyworkr1v.py b/vllm/model_executor/models/skyworkr1v.py index ef2626db1f3c..961ee7acb86f 100644 --- a/vllm/model_executor/models/skyworkr1v.py +++ b/vllm/model_executor/models/skyworkr1v.py @@ -49,13 +49,12 @@ class SkyworkR1VImagePixelInputs(TensorSchema): - """ - Dimensions: - - bnp: Batch size * number of images * (1 + num_patches) - - c: Number of channels (3) - - h: Height - - w: Width - - bn: Batch size * number of images + """Dimensions: + - bnp: Batch size * number of images * (1 + num_patches) + - c: Number of channels (3) + - h: Height + - w: Width + - bn: Batch size * number of images """ type: Literal["pixel_values"] = "pixel_values" @@ -72,12 +71,11 @@ class SkyworkR1VImagePixelInputs(TensorSchema): class SkyworkR1VImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - ni: Number of images - - ifs: Image feature size - - hs: Hidden size (must match the hidden size of language model - backbone) + """Dimensions: + - ni: Number of images + - ifs: Image feature size + - hs: Hidden size (must match the hidden size of language model + backbone) """ type: Literal["image_embeds"] = "image_embeds" diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py index 18512528820c..7c4c24813b7d 100644 --- a/vllm/model_executor/models/stablelm.py +++ b/vllm/model_executor/models/stablelm.py @@ -20,7 +20,8 @@ # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/modeling_stablelm_epoch.py # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json """Inference-only StableLM (https://github.com/Stability-AI/StableLM) -model compatible with HuggingFace weights.""" +model compatible with HuggingFace weights. +""" from collections.abc import Iterable from itertools import islice diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index b060ed1d15dd..aed99c84cb57 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -63,15 +63,14 @@ class Step3VLImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - c: Number of channels (3) - - h: Height - - w: Width - - bnp: Batch size * number of images * number of patches - - hp: Height of patch - - wp: Width of patch + """Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height + - w: Width + - bnp: Batch size * number of images * number of patches + - hp: Height of patch + - wp: Width of patch """ type: Literal["pixel_values"] @@ -81,11 +80,10 @@ class Step3VLImagePixelInputs(TensorSchema): class Step3VLImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - f: Image feature size - - h: Hidden size (must match the hidden size of language model backbone) + """Dimensions: + - bn: Batch size * number of images + - f: Image feature size + - h: Hidden size (must match the hidden size of language model backbone) """ type: Literal["image_embeds"] = "image_embeds" @@ -297,7 +295,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: class Step3VisionAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -350,7 +348,7 @@ def forward( self, hidden_states: torch.Tensor, ): - """Input shape: Batch x Time x Channel""" + """Input shape: Batch x Time x Channel.""" bsz, tgt_len, _ = hidden_states.size() # get query proj diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 03343a89a144..9d649ba4c001 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -69,9 +69,7 @@ class FP32ReplicatedLinear(ReplicatedLinear): - """ - Use FP32 for higher precision. - """ + """Use FP32 for higher precision.""" def forward( self, diff --git a/vllm/model_executor/models/step3p5_mtp.py b/vllm/model_executor/models/step3p5_mtp.py index 294b3559c011..ac78384b601e 100644 --- a/vllm/model_executor/models/step3p5_mtp.py +++ b/vllm/model_executor/models/step3p5_mtp.py @@ -306,8 +306,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int | None, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer """ spec_layer_weight_names = [ diff --git a/vllm/model_executor/models/step_vl.py b/vllm/model_executor/models/step_vl.py index 247078e2ec09..3f31547f9824 100644 --- a/vllm/model_executor/models/step_vl.py +++ b/vllm/model_executor/models/step_vl.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""This is basically a copy from perception_models/core/vision_encoder/pe.py""" +"""This is basically a copy from perception_models/core/vision_encoder/pe.py.""" from collections.abc import Callable from functools import partial diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index 7181680fe833..640f94647634 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -15,7 +15,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Wrapper around `Terratorch` models""" +"""Wrapper around `Terratorch` models.""" from collections import OrderedDict from collections.abc import Iterable, Mapping, Sequence diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index cbfe55195233..2d127c430e76 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -14,7 +14,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Wrapper around `transformers` models""" +"""Wrapper around `transformers` models.""" from typing import TYPE_CHECKING diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 89c6486e739e..eeec0a4b919e 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -193,8 +193,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): ) def _patch_config(self): - """ - Patch the config to ensure that the model is created correctly: + """Patch the config to ensure that the model is created correctly: - Sets the attention implementation to "vllm" so the attention instances from `_create_attention_instances` are used @@ -208,7 +207,8 @@ def _patch_config(self): def _mark_model_components(self, vllm_config: "VllmConfig"): """Mark language model and tower submodules as `self.model` is created. - Nothing to do in `Base`, `MultiModalMixin` will override.""" + Nothing to do in `Base`, `MultiModalMixin` will override. + """ yield @cached_property @@ -253,8 +253,7 @@ def _decorate_cls_for_torch_compile( enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, ): - """ - Decorate `cls` to indicate to vLLM that it supports torch compile. + """Decorate `cls` to indicate to vLLM that it supports torch compile. Args: cls: The PreTrainedModel class to decorate. @@ -265,6 +264,7 @@ def _decorate_cls_for_torch_compile( enable_if: A function which takes in the vLLM config and returns whether torch compile should be enabled for this class. is_encoder: Whether the class being decorated is an encoder. + """ logger.debug( "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s", @@ -281,7 +281,8 @@ def _decorate_cls_for_torch_compile( def _decorate_for_torch_compile(self): """Decorate the model's decoder class to indicate to vLLM that it - supports torch compile if `can_enable_torch_compile` is True.""" + supports torch compile if `can_enable_torch_compile` is True. + """ self._decorate_cls_for_torch_compile( cls=self._pre_trained_model_classes.decoder, # Applied to a PreTrainedModel so the batch dimension will exist @@ -295,8 +296,7 @@ def _decorate_for_torch_compile(self): ) def _create_hf_to_vllm_mapper(self): - """ - Create a WeightsMapper to map checkpoint weight names to module qualnames. + """Create a WeightsMapper to map checkpoint weight names to module qualnames. This handles: @@ -347,18 +347,14 @@ def _create_hf_to_vllm_mapper(self): self._maybe_apply_model_mapping() def _get_tie_word_embeddings(self): - """ - Check if the model has tied word embeddings. - """ + """Check if the model has tied word embeddings.""" # Models created with Transformers v4 and v5 will store this in different places tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False) tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False) return tie_word_embeddings_v4 or tie_word_embeddings_v5 def pipeline_parallel(self): - """ - Apply the model's pipeline parallelization plan. - """ + """Apply the model's pipeline parallelization plan.""" if self.pp_group.world_size <= 1: return @@ -697,8 +693,7 @@ def _get_attn_cls(self) -> type[AttentionLayerBase]: return Attention def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None): - """ - If a `parameter` is on the `meta` device, then its parent + """If a `parameter` is on the `meta` device, then its parent `module` is the original module created by: ```python diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index 29e780557cb7..c5231f9e979d 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -65,7 +65,6 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]: """A thin wrapper around `Base.load_weights` to handle the lm_head bias.""" - lm_head_bias = set() def auto_load_lm_head_bias(weights): diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index ad24f2542a47..0b83afe6b727 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -55,17 +55,20 @@ def fuse( self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Apply the fusion to an already-validated `module`, returning the - module to install in its place (mutated in place, or freshly built).""" + module to install in its place (mutated in place, or freshly built). + """ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes - (none unless it stacks weights).""" + (none unless it stacks weights). + """ return {} @property def packed_modules_mapping(self) -> dict[str, list[str]]: """`packed_modules_mapping` entries this fuser contributes (none unless - it stacks weights).""" + it stacks weights). + """ return {} @@ -119,7 +122,8 @@ def fuse( ) -> nn.Module: """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). - Builds the merged submodule and binds the compiled forward.""" + Builds the merged submodule and binds the compiled forward. + """ self.update_attrs(module, prefix, vllm_config) module.forward = types.MethodType(self.fused_forward, module) return module @@ -128,7 +132,8 @@ def fuse( @dataclass class StackedFuser(RewriteFuser): """A fuser that merges sibling projections into one stacked linear and - rewrites the forward to call it.""" + rewrites the forward to call it. + """ merged_name: ClassVar[str] """Attribute name of the merged module created by `update_attrs`.""" @@ -147,14 +152,16 @@ def info(self, name: str) -> str: def shards(self) -> list[tuple[str, ShardId]]: """Each projection's original name and its shard id in the merged module. - Source for both `orig_to_new_stacked` and `packed_modules_mapping`.""" + Source for both `orig_to_new_stacked` and `packed_modules_mapping`. + """ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: """`WeightsMapper.orig_to_new_stacked` entries for one fused instance. Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname so only this exact layer is remapped, never a same-named projection - elsewhere (e.g. an unfused MoE expert's `gate_proj`).""" + elsewhere (e.g. an unfused MoE expert's `gate_proj`). + """ merged = maybe_prefix(prefix, self.merged_name) return { maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards @@ -163,7 +170,8 @@ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: @property def packed_modules_mapping(self) -> dict[str, list[str]]: """`{merged_name: [projection names]}` so quantization can unpack the - fused layer into its per-shard configs.""" + fused layer into its per-shard configs. + """ return {self.merged_name: [name for name, _ in self.shards]} def _unguarded_calls( @@ -174,7 +182,8 @@ def _unguarded_calls( `update_attrs` deletes the projections it merges, so any reference to one beyond its call site must be a guard on its existence; folding those to their constant value keeps the rewritten forward off a name that no longer exists. A reference - that is not such a guard raises, so fusion is skipped.""" + that is not such a guard raises, so fusion is skipped. + """ calls = [] for name in names: call, refs = self_call_and_refs(funcdef, name) diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index fa4b34750c5f..fccb1a8ea966 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -232,7 +232,8 @@ def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: def update_forward(self, module: nn.Module) -> None: """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split. Bypass the KV expansion method so the compressed latent reaches the `vllm_mla` - attention interface unexpanded.""" + attention interface unexpanded. + """ funcdef, fn = recover_forward(type(module)) if self.has_q_lora: # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index af3a101a0ce1..b7f435b4e032 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -80,7 +80,8 @@ def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None: Computations that must run in higher precision say so in their forward (e.g. `hidden_states.type(torch.float32)`), so the dtype is readable from the graph - even when the config does not name it.""" + even when the config does not name it. + """ for node in nodes: if node.op not in ("call_method", "call_function"): continue @@ -125,7 +126,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: """Standard MoE block forward. - Routing and any shared experts are handled inside `self.experts: MoERunner`.""" + Routing and any shared experts are handled inside `self.experts: MoERunner`. + """ orig_shape = hidden_states.shape hidden_states = hidden_states.reshape(-1, orig_shape[-1]) num_tokens = hidden_states.shape[0] @@ -152,7 +154,8 @@ class MoEBlockFuser: def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. - Returns the scoring function and the dtype the router computes in.""" + Returns the scoring function and the dtype the router computes in. + """ state = {name for name, _ in named_state(gate)} if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None @@ -280,7 +283,8 @@ def shared_experts( self, moe_block: nn.Module, prefix: str ) -> SharedExpertMLP | None: """Build the HF shared expert (and its optional gate) - as a `SharedExpertMLP` for vLLM's fused MoE.""" + as a `SharedExpertMLP` for vLLM's fused MoE. + """ if self.shared_name is None: return None shared_experts = getattr(moe_block, self.shared_name) diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index 47d20085a2c7..dd2c5339ce5a 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -51,7 +51,7 @@ def _is_squared(node: object, x: fx.Node) -> bool: def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None: - """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape.""" + """Eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape.""" add = peel(rsqrt.args[0]) if not is_op(add, "add"): return None diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 332bd5414fc3..b0f449a6d212 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -43,7 +43,8 @@ def _reference_weight(module: nn.Module) -> torch.Tensor | None: """A weight whose trailing dim is the module's hidden size. Linears and 2-D gate weights are `[out, hidden]`; norm weights are - `[hidden]`. Used to fabricate a placeholder input of matching size/dtype.""" + `[hidden]`. Used to fabricate a placeholder input of matching size/dtype. + """ for child in module.modules(): if isinstance(child, nn.Linear): return child.weight @@ -58,7 +59,8 @@ class _MetaProxy(fx.Proxy): Shape questions (`len`, iteration, `.shape` unpacks) are answered by executing each op on the meta values, so PyTorch's meta kernels are the - single source of shape inference — no per-op rules.""" + single source of shape inference — no per-op rules. + """ meta: object = _UNKNOWN @@ -82,7 +84,8 @@ class _MetaAttribute(_MetaProxy, fx.proxy.Attribute): """Attribute proxy (e.g. `x.shape`) carrying its meta value. `Proxy.__getattr__` constructs `Attribute` directly, bypassing - `Tracer.proxy`, so the meta value must be grafted on here too.""" + `Tracer.proxy`, so the meta value must be grafted on here too. + """ def __init__(self, root: fx.Proxy, attr: str): super().__init__(root, attr) @@ -218,7 +221,8 @@ def _leaf_attention_interfaces(): """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node. `vllm_attention_function` needs runtime context so it is untraceable. - Every interface returns `(attn_output, attn_weights)`.""" + Every interface returns `(attn_output, attn_weights)`. + """ from transformers.modeling_utils import AttentionInterface original = AttentionInterface.get_interface @@ -296,7 +300,8 @@ def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]: def forward_input_count(cls: type[nn.Module]) -> int: """The number of tensor inputs `cls.forward` declares, excluding `self` and any `*args`/`**kwargs`. Read from the signature, so it is independent of - whether the trace completes (unlike counting placeholders).""" + whether the trace completes (unlike counting placeholders). + """ params = list(forward_parameters(cls).values()) if not params: return 1 # uninspectable: assume a single input and let matching decide @@ -373,7 +378,8 @@ def _in_boolean_context(funcdef: ast.FunctionDef, ref: ast.expr) -> bool: In these positions the object's identity never escapes, so a reference that is always truthy can be replaced by `True`. `and`/`or` are excluded: they - yield an operand, so the module could escape (`x and self.`).""" + yield an operand, so the module could escape (`x and self.`). + """ for node in ast.walk(funcdef): if ( isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert)) @@ -406,7 +412,8 @@ def bypass_existence_guard( `not self.`): the reference itself to `True`. Any other surviving reference escapes the projection's value, which no longer - exists after fusion, so refuse rather than change semantics.""" + exists after fusion, so refuse rather than change semantics. + """ for node in ast.walk(funcdef): # `self. is (not) None`, either operand order. if not (isinstance(node, ast.Compare) and len(node.ops) == 1): @@ -467,7 +474,8 @@ def _base_name(node: ast.expr) -> str | None: def _rebound_names(region: list[ast.stmt]) -> set[str]: """Names rebound outright within `region` (`x = ...`, `del x`). - These are `Name` nodes in a `Store`/`Del` context.""" + These are `Name` nodes in a `Store`/`Del` context. + """ return { node.id for stmt in region @@ -481,7 +489,8 @@ def _inplace_target(node: ast.AST) -> str | None: A write through the name (`x[i] = ...`, `x.attr = ...`) or an in-place method call (`x.mul_(...)`) leaves the base name in a `Load` context, so it is not a - plain `Name` store (see `_rebound_names`).""" + plain `Name` store (see `_rebound_names`). + """ if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance( node.ctx, (ast.Store, ast.Del) ): @@ -702,7 +711,8 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: Non-linear submodules are transparent too (e.g. the dropout GPT-style attentions apply after their output projection). Never walks through a leaf call (e.g. an attention interface): its inputs are what attention consumes, - not what produced the value.""" + not what produced the value. + """ stack = [node] seen: set[fx.Node] = set() while stack: @@ -725,7 +735,8 @@ def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None: """Nearest linear consuming `node`'s output, walking through casts/scalings. Never walks through a leaf call (e.g. an attention interface): what crosses - it is consumed by the attention computation, not projected.""" + it is consumed by the attention computation, not projected. + """ queue = list(node.users) seen: set[fx.Node] = set() while queue: @@ -787,9 +798,7 @@ def is_method(node: object, name: str) -> bool: def is_op(node: object, name: str) -> bool: - """ - Is node `torch.()`, `F.()`, `operator.()`, or `Tensor.()`. - """ + """Is node `.()` for torch, F, operator, or Tensor.""" return any( is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) ) or (hasattr(torch.Tensor, name) and is_method(node, name)) diff --git a/vllm/model_executor/models/transformers/layers.py b/vllm/model_executor/models/transformers/layers.py index 4ff3651dfc37..51c4bfef98cb 100644 --- a/vllm/model_executor/models/transformers/layers.py +++ b/vllm/model_executor/models/transformers/layers.py @@ -22,7 +22,8 @@ def _resolve(module: str, name: str): """Return `name` from the hw-agnostic `module` when enabled and available, - else from vLLM. Logs which source was used.""" + else from vLLM. Logs which source was used. + """ if envs.VLLM_USE_HW_AGNOSTIC: try: obj = getattr(importlib.import_module(f"{_HW_PKG}.{module}"), name) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 480bc7f66e8d..73cf42cc5f80 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -75,7 +75,8 @@ def forward( ) -> torch.Tensor: """In Transformers `experts.forward` will have this signature. - We discard any extra kwargs because we cannot use them here.""" + We discard any extra kwargs because we cannot use them here. + """ # Note: we need to forward through a custom op so the topk_ids # can be transferred without interfering with cudagraphs. return torch.ops.vllm.transformers_moe_forward( @@ -367,7 +368,8 @@ def custom_routing_function( moe_state: TransformersMoEState, ): """Return `topk_weights` from `gating_output` and the - `topk_ids` we stored in the layer earlier.""" + `topk_ids` we stored in the layer earlier. + """ topk_weights = gating_output topk_ids = moe_state.topk_ids assert topk_ids is not None diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 1cdbadb34b2a..ff042e2f92b1 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -224,7 +224,8 @@ def _get_hf_mm_data( mm_items: MultiModalDataItems, ) -> tuple[Mapping[str, object], Mapping[str, object]]: """Rename the parser's `audios` key to the `audio` argument HF audio - processors take.""" + processors take. + """ processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) if self.info._is_audio_model() and "audios" in processor_data: processor_data["audio"] = processor_data.pop("audios") @@ -432,7 +433,8 @@ def _get_prompt_updates( out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: """Empty, because `apply` writes the placeholder ranges itself rather than - deriving them from updates.""" + deriving them from updates. + """ return [] def _get_mm_token_ids(self, modality: str) -> list[int]: @@ -466,7 +468,8 @@ def _apply_audio( num_audios: int, ) -> dict[str, list[PlaceholderRange]]: """Take each contiguous run of the audio token as one item's placeholder, - and record how many tokens the run holds.""" + and record how many tokens the run holds. + """ audio_token_ids = self._get_mm_token_ids("audio") prompt_tensor = torch.tensor(prompt_ids) is_audio = torch.isin(prompt_tensor, torch.tensor(audio_token_ids)) @@ -524,7 +527,8 @@ def _apply_vision( mm_token_type_ids: torch.Tensor | None, ) -> dict[str, list[PlaceholderRange]]: """Split the positions the processor marks as image into one placeholder per - item, sized by the token count it reports for each image.""" + item, sized by the token count it reports for each image. + """ hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) if mm_token_type_ids is None: raise ValueError( @@ -577,7 +581,8 @@ def apply( timing_ctx: TimingContext, ) -> MultiModalInput: """Process the prompt and every multi-modal item in one HF processor call, - then read the placeholder ranges out of the token ids it returns.""" + then read the placeholder ranges out of the token ids it returns. + """ prompt = inputs.prompt mm_items = inputs.mm_data_items hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs @@ -705,7 +710,8 @@ def _get_prompt_updates( out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: """Replace each modality's placeholder token with the token ids that item's - replacement text encodes to, marking which of them hold embeddings.""" + replacement text encodes to, marking which of them hold embeddings. + """ hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) tokenizer = self.info.get_tokenizer() @@ -892,7 +898,8 @@ def _find_encoder_classes( self, model: "PreTrainedModel" ) -> dict[str, type["PreTrainedModel"]]: """Modalities whose encoder cannot be told apart from the model itself are - omitted, as are those `get_encoder` rejects.""" + omitted, as are those `get_encoder` rejects. + """ encoder_classes: dict[str, type[PreTrainedModel]] = {} for modality in _MODALITY_TO_TOKEN_TYPE_ID: try: @@ -953,9 +960,8 @@ def _mark_model_components(self, vllm_config: "VllmConfig"): yield def _decorate_for_torch_compile(self): - """ - Decorate the model's decoder and encoder classes to indicate to vLLM that they - support torch compile if `can_enable_torch_compile` and + """Decorate the model's decoder and encoder classes to indicate to vLLM + that they support torch compile if `can_enable_torch_compile` and `should_torch_compile_mm_encoder` are True respectively. """ super()._decorate_for_torch_compile() @@ -1012,8 +1018,8 @@ def forward( def get_language_model(self) -> torch.nn.Module: """Transformers modeling backend multimodal classes do not contain a separate vLLM language model class. Therefore, in order to return a language model vLLM - class, we use a wrapper to give `self` the same interface as a text model.""" - + class, we use a wrapper to give `self` the same interface as a text model. + """ # Exclude self and object bases = self.__class__.mro()[1:-1] # Keep only classes defined in `vllm.model_executor.models.transformers` @@ -1031,9 +1037,7 @@ def __init__(self, multimodal_model): return LanguageModel(self) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models""" for name in ("language_model", "text_model"): if getattr(self.model, name, None) is not None: return MultiModelKeys.from_string_field(language_model=f"model.{name}") @@ -1169,7 +1173,8 @@ def _select_item_kwargs( ) -> dict[str, Any]: """Narrow the entries of `kwargs` that hold one row per item down to the item at `index`. Length is all there is to match on, so an unrelated entry of the - same length is narrowed too.""" + same length is narrowed too. + """ return { key: value[index : index + 1] if isinstance(value, (torch.Tensor, list)) and len(value) == num_items diff --git a/vllm/model_executor/models/transformers/pooling.py b/vllm/model_executor/models/transformers/pooling.py index f4fa4b496f23..47a9c23d2c3f 100644 --- a/vllm/model_executor/models/transformers/pooling.py +++ b/vllm/model_executor/models/transformers/pooling.py @@ -84,8 +84,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.init_parameters(self.classifier, dtype=self.model_config.head_dtype) class ClassifierWithReshape(self.classifier.__class__): - """ - Token extraction has already been applied in `pooler.pooling`. + """Token extraction has already been applied in `pooler.pooling`. Add dim to match expected input shape of `classifier.forward`. """ diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 116dd7b5166b..47921b988d88 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -49,16 +49,15 @@ # Copied from `accelerate` @contextmanager def init_on_device_without_buffers(device: torch.device): - """ - A context manager under which models are initialized with all + """A context manager under which models are initialized with all parameters on the specified device. However buffers are not initialized on specified device. Args: device (`torch.device`): Device to initialize all parameters on. - """ + """ old_register_parameter = nn.Module.register_parameter def register_empty_parameter(module, name, param): @@ -114,17 +113,18 @@ def replace_linear_class( *, prefix: str = "", ) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear: - """ - Replace nn.Linear with one of vLLM's tensor parallel linear classes. + """Replace nn.Linear with one of vLLM's tensor parallel linear classes. Args: linear: `nn.Linear` to be replaced. style: Tensor parallel style of the new linear, e.g. "colwise". quant_config: Quantization config for the new linear. + prefix: Module prefix of the new linear, used for quantization lookup. + Returns: The new linear. - """ + """ if not isinstance(style, str): raise ValueError(f"Unsupported parallel style type {type(style)}, expected str") @@ -156,9 +156,11 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: Args: conv: `nn.Conv2d` or `nn.Conv3d` to be replaced. + Returns: The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported, returns the original conv module. + """ # vLLM does not handle non-zero padding modes if conv.padding_mode != "zeros": @@ -199,7 +201,8 @@ def setter(obj: object, value: object): class _UninitializedEmbedding(nn.Embedding): """Make `__init__` inert, so that `VocabParallelEmbedding.__init__`' - call to `super().__init__` does not invoke `nn.Embedding.__init__`.""" + call to `super().__init__` does not invoke `nn.Embedding.__init__`. + """ def __init__(self): pass @@ -207,7 +210,8 @@ def __init__(self): class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbedding): """Orders `VocabParallelEmbedding` ahead of `nn.Embedding` in the MRO, so that - `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding.""" + `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding. + """ @lru_cache @@ -217,8 +221,10 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe Args: cls: The `nn.Embedding` subclass to rebase. Cached, so a given `cls` always maps to the same class. + Returns: The new class, to assign to `__class__` of an instance of `cls`. + """ return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) @@ -235,8 +241,10 @@ def replace_embedding_class( embedding: The `nn.Embedding` holding a vocab table. quant_config: Quantization config for the new embedding. prefix: Qualname of `embedding`, used to look up its quantization method. + Returns: The module to install in `embedding`'s place. + """ # Tied tables are reached more than once, so don't rebase an already new embedding if isinstance(embedding, VocabParallelEmbedding): @@ -320,8 +328,7 @@ def get_feature_request_tip( def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool: - """ - Callable to be passed to `@support_torch_compile`'s `enable_if` argument. + """Callable to be passed to `@support_torch_compile`'s `enable_if` argument. Defaults to `True` but is disabled in the following situations: diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index bffe94fdce26..6465972cf161 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -72,8 +72,7 @@ class UltravoxAudioFeatureInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - b: batch size - n: number of chunks - t: Time frames (M) @@ -96,8 +95,7 @@ class UltravoxAudioFeatureInputs(TensorSchema): class UltravoxAudioEmbeddingInputs(TensorSchema): - """ - Dimensions: + """Dimensions: - b: batch size - na: number of audios - afs: audio feature size @@ -340,8 +338,7 @@ def get_replacement_ultravox(item_idx: int): class StackAudioFrames(nn.Module): - """ - Stack the audio embedding frames to reduce the sequence length by a factor + """Stack the audio embedding frames to reduce the sequence length by a factor of `stack_factor`. """ @@ -752,9 +749,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) def get_mm_mapping(self) -> MultiModelKeys: - """ - Get the module prefix in multimodal models - """ + """Get the module prefix in multimodal models.""" return MultiModelKeys.from_string_field( language_model="language_model.", connector="multi_modal_projector.", @@ -939,9 +934,9 @@ def forward( positions: torch.Tensor, intermediate_tensors: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, - **kwargs, + **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - """Run forward pass for Ultravox + """Run forward pass for Ultravox. One key thing to understand is the `input_ids` already accounts for the positions of the to-be-inserted audio embeddings. The to-be-inserted @@ -956,9 +951,10 @@ def forward( positions: Position indices for the input tokens. intermediate_tensors: Intermediate tensors from prior forward pass. inputs_embeds: Optional tensor of input embeddings. + **kwargs: Multimodal inputs for this batch, forwarded to the + multimodal embedding path. """ - if intermediate_tensors is not None: inputs_embeds = None @@ -982,8 +978,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def pad_and_concat_to_dim3( features: torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]], ) -> torch.Tensor: - """ - Pad and concatenate a list of tensors. + """Pad and concatenate a list of tensors. output: Tensor of shape [B, C, M] where M is the maximum length of the input diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 6cd21763a095..5d76066b826c 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -47,7 +47,8 @@ class WeightsMapper: """Maps the name of each weight if they match the following patterns. - If a key maps to a value of `None`, the corresponding weight is ignored.""" + If a key maps to a value of `None`, the corresponding weight is ignored. + """ orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list) orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) @@ -84,6 +85,7 @@ def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None: Returns: (mapped_name, shard_id) if the name should be kept. None if the name should be dropped. + """ # Deprecation warnings if key.endswith(".kv_scale"): @@ -170,7 +172,8 @@ def get_rename_mapper(self) -> "WeightsMapper": rather than being rewritten to the stacked vLLM name (`qkv_proj`). Mappings to `None` are dropped because "do not load this weight" is meaningless to such a consumer, and applying it would silently shrink a quantization config's ignore - list or make LoRA name parsing fail.""" + list or make LoRA name parsing fail. + """ remove_none = lambda d: {k: v for k, v in d.items() if v is not None} return replace( self, @@ -197,8 +200,7 @@ def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]: class AutoWeightsLoader: - """ - Helper class to load weights into a [`torch.nn.Module`][]. It is able + """Helper class to load weights into a [`torch.nn.Module`][]. It is able to automatically detect child modules and parameters while iterating over the weights only once. @@ -315,8 +317,7 @@ def _load_param( def _add_loadable_non_param_tensors( self, module: nn.Module, child_params: dict[str, torch.Tensor] ): - """ - Add tensor names that are not in the model params that may be in the + """Add tensor names that are not in the model params that may be in the safetensors, e.g., batch normalization stats and registered buffers. """ # Add persistent registered buffers. @@ -501,6 +502,7 @@ def maybe_fuse_shared_experts( Yields: `(name, tensor)` pairs with shared experts routed to fused slots. + """ if enabled is None: from vllm._aiter_ops import rocm_aiter_ops @@ -552,6 +554,7 @@ def get_spec_layer_idx_from_weight_name( Returns: The absolute layer index for an MTP-layer weight, else None. + """ if not (n := getattr(config, "num_nextn_predict_layers", 0)): return None @@ -579,6 +582,7 @@ def skip_spec_layers( Yields: `(name, tensor)` pairs whose weight is not an MTP-layer weight. + """ return ( (name, w) @@ -594,8 +598,7 @@ def init_vllm_registered_model( hf_config: "PretrainedConfig | None" = None, architectures: list[str] | None = None, ) -> nn.Module: - """ - Helper function to initialize an inner model registered to vLLM, + """Helper function to initialize an inner model registered to vLLM, based on the arguments passed to the outer vLLM model. """ from vllm.model_executor.model_loader.utils import initialize_model @@ -639,8 +642,7 @@ def flatten_bn( *, concat: bool = False, ) -> list[torch.Tensor] | torch.Tensor: - """ - Flatten the `B` and `N` dimensions of batched multimodal inputs. + """Flatten the `B` and `N` dimensions of batched multimodal inputs. The input tensor should have shape `(B, N, ...)`. """ @@ -654,11 +656,9 @@ def flatten_bn( def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor: - """ - Recursively flattens and concatenates NestedTensors on all but the last + """Recursively flattens and concatenates NestedTensors on all but the last dimension. """ - if isinstance(embeddings, torch.Tensor): # Flatten all but the last dimension. return embeddings.flatten(0, -2) @@ -667,11 +667,9 @@ def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor: def _embedding_count_expression(embeddings: NestedTensors) -> str: - """ - Constructs a debugging representation of the number of embeddings in the + """Constructs a debugging representation of the number of embeddings in the NestedTensors. """ - if isinstance(embeddings, torch.Tensor): return " x ".join([str(dim) for dim in embeddings.shape[:-1]]) @@ -691,13 +689,13 @@ def _merge_multimodal_embeddings( multimodal_embeddings: NestedTensors, is_multimodal: torch.Tensor, ) -> torch.Tensor: - """ - Merge `multimodal_embeddings` into `inputs_embeds` by overwriting the + """Merge `multimodal_embeddings` into `inputs_embeds` by overwriting the positions in `inputs_embeds` corresponding to placeholder tokens in `input_ids`. Note: This updates `inputs_embeds` in place. + """ if len(multimodal_embeddings) == 0: return inputs_embeds @@ -751,8 +749,7 @@ def collect_children( *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, ): - """ - Within this context, collect all direct child assignments to `module`, + """Within this context, collect all direct child assignments to `module`, returning a list of children names that is internally updated until the context is exited. @@ -784,8 +781,7 @@ def no_init_weights( *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, ): - """ - Within this context, prevent weight initialization from using device memory and + """Within this context, prevent weight initialization from using device memory and replace direct child assignments to `module` with the result of `placeholder()`. If `targets` is set, instead prevent weight initialization and @@ -824,9 +820,7 @@ def __call__(self, prefix: str) -> torch.nn.Module: ... class PPMissingLayer(torch.nn.Identity): - """ - A placeholder layer for missing layers in a pipeline parallel model. - """ + """A placeholder layer for missing layers in a pipeline parallel model.""" def __init__(self, *args, **kwargs): super().__init__() @@ -867,6 +861,7 @@ def make_layers( Returns: Tuple of (start_layer, end_layer, modules). + """ from vllm.distributed.parallel_state import get_pp_group from vllm.distributed.utils import get_pp_indices @@ -945,6 +940,7 @@ def maybe_prefix(prefix: str, name: str) -> str: Returns: The string "prefix.name" if prefix was non-empty, otherwise just "name". + """ return name if not prefix else f"{prefix}.{name}" @@ -960,6 +956,7 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non Returns: The draft model's config if available, None otherwise. + """ draft_model_config = vllm_config.speculative_config.draft_model_config draft_load_config = vllm_config.load_config @@ -972,13 +969,14 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int: - """ - Extract the layer index from the module name. + """Extract the layer index from the module name. + Examples: - "encoder.layers.0" -> 0 - "encoder.layers.1.self_attn" -> 1 - "2.self_attn" -> 2 - "model.encoder.layers.0.sub.1" -> ValueError if num_attn_module == 1 + """ subnames = layer_name.split(".") int_vals: list[int] = [] @@ -1013,8 +1011,7 @@ def cast_overflow_tensors(tensors: torch.Tensor, offset: float = 1000) -> torch. def fast_topk( values: torch.Tensor, topk: int, dim: int ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Optimized topk implementation that uses torch.max for k=1 case. + """Optimized topk implementation that uses torch.max for k=1 case. This function provides better performance for the common case of k=1 by using torch.max instead of the more general torch.topk. @@ -1027,6 +1024,7 @@ def fast_topk( Returns: Tuple of (values, indices) where values are the top-k values and indices are their corresponding indices in the input tensor + """ if topk == 1: # Use max along the specified dimension to get both value and index @@ -1086,13 +1084,14 @@ def process_eagle_weight( model: nn.Module, name: str, ) -> None: - """ - Update EAGLE model flags based on loaded weight name. + """Update EAGLE model flags based on loaded weight name. This should be called during weight loading to detect if a model has its own lm_head or embed_tokens weight. + Args: model: The model instance (must support EAGLE) name: The name of the weight to process + """ if not supports_any_eagle(model): return @@ -1111,6 +1110,7 @@ def get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int: Args: feature_layer_index: Index of a required layer in the visual encoder. num_hidden_layers: The total number of hidden layers in the visual encoder. + """ if feature_layer_index < 0: return num_hidden_layers + feature_layer_index + 1 diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index 36008d999654..fb684ddba4d9 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -88,9 +88,7 @@ def _get_vit_attn_backend( *, attn_backend_override: AttentionBackendEnum | None = None, ) -> AttentionBackendEnum: - """ - Get the available attention backend for Vision Transformer. - """ + """Get the available attention backend for Vision Transformer.""" return current_platform.get_vit_attn_backend( head_size, dtype, @@ -102,9 +100,7 @@ def get_vit_attn_backend( head_size: int, dtype: torch.dtype, ) -> AttentionBackendEnum: - """ - Get the attention backend for Vision Transformer. - """ + """Get the attention backend for Vision Transformer.""" mm_cfg = get_multimodal_config() attn_backend_override = ( mm_cfg.mm_encoder_attn_backend if mm_cfg is not None else None @@ -120,7 +116,8 @@ def get_multimodal_config() -> MultiModalConfig | None: """Return the current ``MultiModalConfig``, or ``None`` when no engine config context is active (e.g., during unit tests) or when the current ``model_config`` does not carry a ``multimodal_config`` (e.g., minimal - stubs used in tests).""" + stubs used in tests). + """ vllm_config = get_current_vllm_config_or_none() if vllm_config is None or vllm_config.model_config is None: return None @@ -142,9 +139,7 @@ def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None: def is_vit_use_data_parallel(num_heads: int | None = None) -> bool: - """ - Get the tensor parallel type for Vision Transformer. - """ + """Get the tensor parallel type for Vision Transformer.""" mm_cfg = get_multimodal_config() can_split = ( num_heads % get_tensor_model_parallel_world_size() == 0 @@ -235,6 +230,7 @@ def resolve_visual_encoder_outputs( concatenated with the other select_layers along the last dimension. feature_select_strategy: Defines how to select the hidden states from each layer. + """ if select_layers is None: if not isinstance(encoder_outputs, torch.Tensor): @@ -302,10 +298,11 @@ def run_dp_sharded_vision_model( Args: image_input (torch.Tensor): Image input tensor. vision_model (torch.nn.Module): Vision model. + Returns: torch.Tensor: Output image embeddings - """ + """ num_chunks = image_input.shape[0] mp_world_size = get_tensor_model_parallel_world_size() num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size @@ -329,8 +326,7 @@ def get_load_balance_assignment( sizes: list[int], num_gpus: int = 2, ) -> tuple[list[int], list[int], list[int]]: - """ - Generate load balancing assignment and metadata + """Generate load balancing assignment and metadata for distributing data across GPUs. The load is determined by the total image sizes, not the number of images. @@ -354,7 +350,6 @@ def get_load_balance_assignment( ``` """ - n_samples = len(sizes) # Handle edge cases @@ -414,6 +409,7 @@ def run_dp_sharded_mrope_vision_model( Different rope types have different dimension to do ViT. "rope_3d" for 3D rope (e.g., Qwen2.5-VL) "rope_2d" for 2D rope (e.g., Kimi-VL) + Returns: torch.Tensor: Output image embeddings @@ -585,8 +581,7 @@ def run_dp_sharded_mrope_vision_model( class FusedInputNorm(nn.Module): - """ - Module that applies rescaling and normalization to input images. + """Module that applies rescaling and normalization to input images. Equivalent to: output = (input * rescale_factor - mean) / std """ diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index b41456bcdc43..5559217d3737 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -502,8 +502,7 @@ def get_num_audio_tokens( stt_config: SpeechToTextConfig, model_config: ModelConfig, ) -> int | None: - """ - Map from audio duration to number of audio tokens produced by the ASR + """Map from audio duration to number of audio tokens produced by the ASR model, without running a forward pass. This is used for estimating the amount of processing for this audio. """ @@ -586,8 +585,7 @@ def llm_weights_generator(): def maybe_update_quant_config( self, quant_config: QuantizationConfig ) -> QuantizationConfig: - """ - Update quant config to so that ignored module and target module names + """Update quant config to so that ignored module and target module names match the vLLM model names. Right now this is specific for compressed-tensors format and load_format mistral. diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 961aa3af9c9a..52b636cc2586 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -90,7 +90,7 @@ def _maybe_apply_prompt_updates( class TimeEmbedding(torch.nn.Module): - """Sinusoidal Embedding for encoding time""" + """Sinusoidal Embedding for encoding time.""" def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() @@ -391,7 +391,7 @@ def forward( def embed_multimodal( self, **kwargs ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None: - """Transform audio waveforms -> initial whisper post-conv embeddings""" + """Transform audio waveforms -> initial whisper post-conv embeddings.""" audio_inputs = self._parse_and_validate_audio_arrays(**kwargs) if audio_inputs is None: diff --git a/vllm/model_executor/models/voyage.py b/vllm/model_executor/models/voyage.py index 92d6d7633f02..d131ffff4b0b 100644 --- a/vllm/model_executor/models/voyage.py +++ b/vllm/model_executor/models/voyage.py @@ -16,8 +16,7 @@ class VoyageQwen3BidirectionalEmbedModel(nn.Module): - """ - Qwen3Model + Voyage embedding head + bidirectional attention. + """Qwen3Model + Voyage embedding head + bidirectional attention. Checkpoint conventions (HF): - MLP: gate_proj + up_proj (unfused) diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index d21903bd656c..7fc620e7dde2 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -94,11 +94,10 @@ class WhisperPosEmbedType(enum.Enum): class WhisperAudioInputs(TensorSchema): - """ - Dimensions: - - b: Batch size - - nmb: Number of mel bins - - t: Time frames (M) + """Dimensions: + - b: Batch size + - nmb: Number of mel bins + - t: Time frames (M) """ input_features: Annotated[ @@ -120,9 +119,8 @@ def forward( # Only used for FlashInfer CuDNN backend. sequence_lengths: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Input shape: batch_size x seq_len x hidden_size - or seq_len x hidden_size + """Input shape: batch_size x seq_len x hidden_size + or seq_len x hidden_size """ is_2d = query.dim() == 2 if is_2d: @@ -931,7 +929,6 @@ def parse_language_detection_output( Decodes the first token ID and extracts the language code from the ``<|xx|>`` format. Expects a valid language token from constrained generation. """ - decoded = tokenizer.decode( [token_ids[0]], skip_special_tokens=False, @@ -1058,8 +1055,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def _create_fake_bias_for_k_proj( weights: Iterable[tuple[str, torch.Tensor]], fake_bias_key_name: str ) -> Iterable[tuple[str, torch.Tensor]]: - """ - Create full zeros bias for k_proj weight in self-attn and x-attn layers. + """Create full zeros bias for k_proj weight in self-attn and x-attn layers. So that the bias for k_proj in qkv_proj can be initialized with zeros. """ for name, weight in weights: diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index f80c8a3b0783..9398d997688a 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -71,6 +71,8 @@ def __init__( rank: LoRA rank output_dim: output dimension quant_config: Configuration for model quantization + prefix: Module prefix used for quantization config lookup + """ super().__init__() @@ -130,6 +132,7 @@ def __init__( cache_config: Configuration for key-value caching quant_config: Configuration for model quantization prefix: Optional prefix for parameter names + """ super().__init__() tp_size = get_tensor_model_parallel_world_size() @@ -257,6 +260,7 @@ def forward( Returns: Output tensor [batch_size, seq_len, hidden_size] + """ qkv, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv.split([self.qkv_size] * 3, dim=-1) @@ -311,6 +315,8 @@ def __init__( bare_block_idx: Index of the bare block in the model num_hybrid_layers: Total number of hybrid layers quant_config: Configuration for model quantization + prefix: Module prefix used for quantization config lookup + """ super().__init__() self.config = config @@ -369,6 +375,7 @@ def forward(self, hidden_states: torch.Tensor, block_idx: int) -> torch.Tensor: Returns: Output tensor [batch_size, seq_len, hidden_size] after applying gated feed-forward transformation + """ # Project input to intermediate size with gating gate_up_states, _ = self.gate_up_proj(hidden_states) @@ -415,6 +422,7 @@ def __init__( cache_config: Configuration for key-value caching quant_config: Configuration for model quantization prefix: Optional prefix for parameter names + """ super().__init__() @@ -461,8 +469,8 @@ def forward( Returns: Transformed hidden states after attention and feed-forward - """ + """ # The argument original_hidden_states is concatenated with hidden_states # (which is the output of the previous (mamba) layer). # The concatenated tensor is then used as input of the pre-attention @@ -510,7 +518,11 @@ def __init__( Args: config: The Zamba2 model configuration + model_config: The model config, when available + cache_config: The KV cache config, when available quant_config: Configuration for model quantization + prefix: Module prefix used for quantization config lookup + """ super().__init__() @@ -555,6 +567,7 @@ def forward( Returns: Transformed hidden states with residual connection applied + """ # Store input for residual connection residual = hidden_states @@ -603,6 +616,13 @@ def __init__( Args: shared_transformer: Transformer decoder layer for attention pathway + config: The Zamba2 model configuration + block_idx: Index of this hybrid block in the model + model_config: The model config, when available + cache_config: The KV cache config, when available + quant_config: Configuration for model quantization + prefix: Module prefix used for quantization config lookup + """ super().__init__() self.block_idx = block_idx @@ -644,6 +664,7 @@ def forward( Returns: Output tensor combining transformer and Mamba representations + """ # Process through transformer pathway transformer_hidden_states = self.shared_transformer( @@ -680,6 +701,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: vllm_config: Configuration object containing model, cache, quantization and LoRA settings prefix: Optional prefix for parameter names in state dict + """ super().__init__() @@ -765,6 +787,7 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: Returns: Embedded representation of the input tokens + """ return self.embed_tokens(input_ids) @@ -784,6 +807,7 @@ def forward( Returns: Either final hidden states or intermediate tensors for pipeline parallelism + """ # Handle pipeline parallelism for first rank if inputs_embeds is None: @@ -853,8 +877,8 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache - """ + """ parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_expand * hf_config.hidden_size @@ -884,6 +908,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: Raises: AssertionError: If prefix caching is enabled (not supported by Mamba) + """ config = vllm_config.model_config.hf_config @@ -914,10 +939,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Convert input token IDs to embeddings. + Args: input_ids: Tensor of input token IDs Returns: Embedded representation of the input tokens + """ return self.model.embed_input_ids(input_ids) @@ -938,6 +965,7 @@ def forward( Returns: Output hidden states + """ # Forward pass through model hidden_states = self.model( @@ -959,6 +987,7 @@ def compute_logits( Returns: Logits for next token prediction + """ logits = self.logits_processor(self.lm_head, hidden_states) return logits diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index cb9c35b79c50..2780ecdd08e2 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -77,6 +77,7 @@ def wrap_modules( Returns: List of modules, potentially with offloading hooks installed. + """ pass diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index e682fd0e038b..b87535ac376e 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -151,6 +151,7 @@ class PrefetchOffloader(BaseOffloader): num_in_group: Offload this many layers per group (last N of each group). prefetch_step: Number of layers to prefetch ahead. mode: Offload mode ("cpu" is currently supported). + """ def __init__( diff --git a/vllm/model_executor/offloader/prefetch_ops.py b/vllm/model_executor/offloader/prefetch_ops.py index b7f599210a19..68a78fbd0200 100644 --- a/vllm/model_executor/offloader/prefetch_ops.py +++ b/vllm/model_executor/offloader/prefetch_ops.py @@ -29,6 +29,7 @@ def _wait_prefetch_impl( input_tensor: Input to the layer (e.g., hidden_states) - declared as mutated to create data dependency for torch.compile. layer_idx: Index of the layer to wait for. + """ get_offloader()._wait_for_layer(layer_idx) @@ -49,6 +50,7 @@ def _start_prefetch_impl( prevent torch.compile from reordering this op before the computation that produces output_tensor. layer_idx: Index of the layer to prefetch. + """ get_offloader()._start_prefetch(layer_idx) diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index f04d4d85b1a4..5b76d2a9a1f0 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -34,6 +34,7 @@ class UVAOffloader(BaseOffloader): cpu_offload_max_bytes: Maximum bytes to offload to CPU. cpu_offload_params: Set of parameter name segments to selectively offload. If empty, all parameters are eligible up to the byte limit. + """ supports_tower_offload = True diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py index 7f96ceda09db..de1cc7d21888 100644 --- a/vllm/model_executor/parameter.py +++ b/vllm/model_executor/parameter.py @@ -30,8 +30,7 @@ class BasevLLMParameter(Parameter): - """ - Base parameter for vLLM linear layers. Extends the torch.nn.parameter + """Base parameter for vLLM linear layers. Extends the torch.nn.parameter by taking in a linear weight loader. Will copy the loaded weight into the parameter when the provided weight loader is called. """ @@ -40,14 +39,13 @@ def __new__(cls, data: torch.Tensor | None, **kwargs): return super().__new__(cls, data=data, requires_grad=False) def __init__(self, data: torch.Tensor, weight_loader: Callable): - """ - Initialize the BasevLLMParameter + """Initialize the BasevLLMParameter. Args: data: torch tensor with the parameter data weight_loader: weight loader callable - """ + """ # During weight loading, we often do something like: # narrowed_tensor = param.data.narrow(0, offset, len) # narrowed_tensor.copy_(real_weight) @@ -127,8 +125,7 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): class _ColumnvLLMParameter(BasevLLMParameter): - """ - Private class defining weight loading functionality + """Private class defining weight loading functionality (load_merged_column_weight, load_qkv_weight) for parameters being loaded into linear layers with column parallelism. This includes QKV and MLP layers which are @@ -202,8 +199,7 @@ def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs): class RowvLLMParameter(BasevLLMParameter): - """ - Parameter class defining weight_loading functionality + """Parameter class defining weight_loading functionality (load_row_parallel_weight) for parameters being loaded into linear layers with row parallel functionality. Requires an input_dim to be defined. @@ -231,8 +227,7 @@ def load_row_parallel_weight(self, loaded_weight: torch.Tensor): class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): - """ - Parameter class for linear layer weights. Uses both column and + """Parameter class for linear layer weights. Uses both column and row parallelism. """ @@ -240,8 +235,7 @@ class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): - """ - Parameter class for weight scales loaded for weights with + """Parameter class for weight scales loaded for weights with grouped quantization. Uses both column and row parallelism. """ @@ -249,8 +243,7 @@ class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): class ChannelQuantScaleParameter(_ColumnvLLMParameter): - """ - Parameter class for weight scales loaded for weights with + """Parameter class for weight scales loaded for weights with channel-wise quantization. Equivalent to _ColumnvLLMParameter. """ @@ -258,8 +251,7 @@ class ChannelQuantScaleParameter(_ColumnvLLMParameter): class PerTensorScaleParameter(BasevLLMParameter): - """ - Parameter class for scales where the number of scales is + """Parameter class for scales where the number of scales is equivalent to the number of logical matrices in fused linear layers (e.g. for QKV, there are 3 scales loaded from disk). This is relevant to weights with per-tensor quantization. @@ -291,11 +283,9 @@ def load_column_parallel_weight(self, *args, **kwargs): def _load_into_shard_id( self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs ): - """ - Slice the parameter data based on the shard id for + """Slice the parameter data based on the shard id for loading. """ - param_data = self.data shard_id = self._shard_id_as_int(shard_id) @@ -311,8 +301,7 @@ def _load_into_shard_id( class PackedColumnParameter(_ColumnvLLMParameter): - """ - Parameter for model parameters which are packed on disk + """Parameter for model parameters which are packed on disk and support column parallelism only. See PackedvLLMParameter for more details on the packed properties. """ @@ -351,8 +340,7 @@ def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): class PackedvLLMParameter(ModelWeightParameter): - """ - Parameter for model weights which are packed on disk. + """Parameter for model weights which are packed on disk. Example: GPTQ Marlin weights are int4 or int8, packed into int32. Extends the ModelWeightParameter to take in the packed factor, the packed dimension, and optionally, marlin @@ -395,8 +383,7 @@ def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): - """ - Parameter class for weight scales loaded for weights with + """Parameter class for weight scales loaded for weights with block-wise quantization. Uses both column and row parallelism. """ @@ -404,8 +391,7 @@ class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): class SharedWeightParameter(BasevLLMParameter): - """ - Parameter for weights with many shared tensors across a model + """Parameter for weights with many shared tensors across a model. For example, when applying transforms to the "gate" and "up" partitions of `MergedColumnParallelLinear`, the transform weights must stay separate @@ -446,8 +432,7 @@ def __init__(self, input_dim: int = 1, output_dim: int = 0, **kwargs): ) def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: Any): - """ - Add a partition to the weight parameter. Partitions whose `data_key` + """Add a partition to the weight parameter. Partitions whose `data_key` is the same will share tensor data Args: @@ -455,6 +440,7 @@ def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: An data_key: hashable key used to key shared tensors *args: arguments for `torch.empty` **kwargs: keyword arguments for `torch.empty` + """ # load (shared) tensor using `data_key` if data_key not in self.tensors_registry: @@ -544,8 +530,7 @@ def _fake_weight_loader( def permute_param_layout_( param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs ) -> BasevLLMParameter: - """ - Permute a parameter's layout to the specified input and output dimensions, + """Permute a parameter's layout to the specified input and output dimensions, useful for forcing the parameter into a known layout, for example, if I need a packed (quantized) weight matrix to be in the layout {input_dim = 0, output_dim = 1, packed_dim = 0} @@ -554,7 +539,6 @@ def permute_param_layout_( to ensure x is in the correct layout (permuting it to the correct layout if required, asserting if it cannot get it to the correct layout) """ - curr_input_dim = getattr(param, "input_dim", None) curr_output_dim = getattr(param, "output_dim", None) diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py index 6c5c517d9ce6..33418847da66 100644 --- a/vllm/model_executor/utils.py +++ b/vllm/model_executor/utils.py @@ -45,6 +45,7 @@ def set_weight_attrs( Args: weight: The weight tensor. weight_attrs: A dictionary of attributes to set on the weight tensor. + """ if weight_attrs is None: return @@ -73,9 +74,8 @@ def replace_parameter( new_tensor: torch.Tensor | None, prefer_copy: bool = False, ): - """ - Replace a parameter of a layer while maintaining the ability to reload the weight. - Called within implementations of the `process_weights_after_loading` method. + """Replace a parameter of a layer while maintaining the ability to reload the + weight. Called within implementations of the `process_weights_after_loading` method. Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as ``is_shuffled``) are carried over to the replacement parameter, except @@ -97,6 +97,7 @@ def replace_parameter( a new parameter. This preserves the parameter's storage address (``data_ptr``), which is required for captured CUDA graphs to remain valid across weight updates (e.g. in RL training loops). + """ # should not be used on a tied/shared param @@ -173,7 +174,8 @@ def get_moe_expert_mapping( It will be retrieved from the first module that has a `get_expert_mapping` method. If the model manually implements `get_expert_mapping`, it will be used. - Otherwise, it will use the first RoutedExperts layer.""" + Otherwise, it will use the first RoutedExperts layer. + """ for _, module in model.named_modules(): get_mapping = getattr(module, "get_expert_mapping", None) if get_mapping is not None: diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 35fd6bb8f9a9..6aa5d391e01d 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Warmup deep_gemm kernels. +"""Warmup deep_gemm kernels. DeepGEMM JIT's the kernels. The warmup aims to JIT all the kernels that would be used during model execution beforehand. """ @@ -36,16 +35,15 @@ def _generate_optimal_warmup_m_values( max_tokens: int, n: int, device: torch.device ) -> list[int]: - """ - Generate M values that cover all possible DeepGEMM kernel configurations. + """Generate M values that cover all possible DeepGEMM kernel configurations. Reference: https://github.com/deepseek-ai/DeepGEMM/blob/79f48ee15a82dd5fad5cd9beaa393c1f755e6b55/csrc/jit_kernels/heuristics/common.hpp Args: max_tokens: Maximum number of tokens to warmup for n: The actual N dimension from the weight tensor device: The torch device to get properties from. - """ + """ # DeepGEMM's possible block sizes block_ms = [64, 128, 256] block_ns = list(range(16, min(257, n + 1), 16)) @@ -82,9 +80,7 @@ def _generate_optimal_warmup_m_values( def _extract_data_from_fused_moe_module( m_: torch.nn.Module, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: - """ - Extract weights, weight scales and num_topk from MoERunner module. - """ + """Extract weights, weight scales and num_topk from MoERunner module.""" assert isinstance(m_, MoERunner) m = m_.routed_experts w13 = m.w13_weight @@ -111,8 +107,7 @@ def _extract_data_from_fused_moe_module( def _deep_gemm_linear_data( module: torch.nn.Module, ) -> tuple[torch.Tensor, torch.Tensor] | None: - """ - Return (weight, weight_scale) for a layer whose kernel stamped itself as + """Return (weight, weight_scale) for a layer whose kernel stamped itself as ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None. """ provider = getattr(module, "deep_gemm_warmup_provider", None) diff --git a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py index 056ca5180508..188f26fdcf17 100644 --- a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py @@ -99,7 +99,6 @@ def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None: parsed key as ``__call__`` so runtime does not report a cache miss for an already materialized specialization. """ - with _quiet_tilelang_warmup_logs(): compiled = jit_impl.compile(*args, **kwargs) func = getattr(jit_impl, "func", None) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 80b5160d26c8..bf2c1ad0b693 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Warmup kernels used during model execution. +"""Warmup kernels used during model execution. This is useful specifically for JIT'ed kernels as we don't want JIT'ing to happen during model execution. """ @@ -369,8 +368,7 @@ def _run_flashinfer_autotune_dummy_runs( def flashinfer_autotune(runner: "GPUModelRunner") -> None: - """ - Autotune FlashInfer operations. + """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, autotuning runs benchmarks for each implementation and stores the results. The results are cached transparently and diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py index 3cbc3bf2e086..d467e8eacf33 100644 --- a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py +++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py @@ -18,8 +18,8 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: """Warm every reachable QSA specialization: indexer decode-query-length - profiles plus the sparse attention split-K/merge configs.""" - + profiles plus the sparse attention split-K/merge configs. + """ qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa") attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa") if qsa_module is None or attn_module is None: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index fd02a51076ed..124573b0e397 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -1059,7 +1059,8 @@ def _fused_eh_norm_kernel( ): """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` - ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat. + """ tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index c1a6903937b6..e60083b1184f 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1425,7 +1425,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index d94ff5c2709e..045a731321b9 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -496,8 +496,7 @@ def _resolve_scale_name(name: str) -> str: return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 9c779809650c..8a151dba1d98 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DeepseekV4 MLA Attention Layer -""" +"""DeepseekV4 MLA Attention Layer.""" from abc import ABC, abstractmethod from collections.abc import Callable diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index 919146c551a4..f17555850098 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -91,7 +91,8 @@ def validate_image_sentinel_ids(tokenizer) -> None: def grid_tokens(best_height, best_width, patch_size, downsample_ratio): """Number of LLM tokens the aligner grid occupies (N-layout, including - row/align padding).""" + row/align padding). + """ n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) num_tokens = n_llm_h * (n_llm_w + 1) + 2 @@ -196,7 +197,8 @@ def load_image( def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): """Builds the N-layout token types (final order) and the aligner-row order - for IMAGE slots.""" + for IMAGE slots. + """ compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO pad_h = n_llm_h % 2 rows = n_llm_h + pad_h @@ -239,7 +241,8 @@ def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): class DeepseekV4VLImageProcessor: """Per-image transform (the PIL-input equivalent of the reference - ``load_image``).""" + ``load_image``). + """ def __init__(self, config: DeepseekV4Config) -> None: super().__init__() diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index d98d6a297a43..12fdb0471174 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index +"""Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index preparation. - quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into @@ -61,8 +60,7 @@ def quantize_and_insert_k_kernel( n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) use_fnuz: tl.constexpr = False, ): - """ - Quantize K tensor and insert into paged K cache. + """Quantize K tensor and insert into paged K cache. K Cache block layout (block_size=64 tokens): - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 @@ -174,8 +172,7 @@ def quantize_and_insert_k_cache( is_ue8m0: bool = True, use_fnuz: bool = False, ): - """ - Quantize K tensor and insert into paged K cache. + """Quantize K tensor and insert into paged K cache. K Cache block layout (block_size=64 tokens): - First 64 * 576 = 36864 bytes: Token data diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 326a94f4f43c..50ca86558e05 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels. +"""Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels. Three specialized kernels: - _fused_kv_compress_norm_rope_insert_sparse_attn: @@ -390,7 +389,7 @@ def _compress_gather_split_sparse_attn( NUM_SPLITS: tl.constexpr, HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS ): - """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch + """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch. No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO) """ @@ -639,7 +638,7 @@ def compress_norm_rope_store_two_stage_triton( num_decode_tokens: int, compress_scratch: torch.Tensor, ) -> None: - """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap) + """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap). Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:] to fill the CUs, and use the original single-pass launcher diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 669186edbb1b..d0168ff25e5f 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -312,7 +312,8 @@ def __call__( def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: """Weights dtypes the model's indexer layers ask for: fp32 for the dense scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer - (`SparseMQAIndexer`) is enabled.""" + (`SparseMQAIndexer`) is enabled. + """ if vllm_config.attention_config.indexer_sparse_logits: return (torch.float32, torch.bfloat16) return (torch.float32,) diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py index 77b1aa2236b5..dd25c73af7c2 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention. +"""Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention. Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90, INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout. @@ -380,6 +379,7 @@ def fused_inv_rope_fp8_quant( Returns: Rotated output in [T, G, D] and its FP8 scales. The scale tensor is empty when quantization is disabled. + """ from vllm.utils.deep_gemm import get_tma_aligned_size diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py index 46468897e59d..4c633da6c549 100644 --- a/vllm/models/deepseek_v4/common/vision.py +++ b/vllm/models/deepseek_v4/common/vision.py @@ -263,6 +263,7 @@ def run_dp_sharded_vision_tower( Returns: One ``(n_aligner_rows, hidden_size)`` embedding tensor per image. + """ tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() diff --git a/vllm/models/deepseek_v4/cpu/cpu_compressor.py b/vllm/models/deepseek_v4/cpu/cpu_compressor.py index 6f0b9e06d842..a0c856dc5505 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_compressor.py +++ b/vllm/models/deepseek_v4/cpu/cpu_compressor.py @@ -42,7 +42,8 @@ def cache_norm_weight_fp32(self) -> None: weight loading (see ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU kernel requires fp32/contiguous but the checkpoint loads this weight - in bf16, and it never changes after loading.""" + in bf16, and it never changes after loading. + """ self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous() def forward( diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index 52d4f4c3b017..a7f2af1cbb55 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -118,7 +118,8 @@ def build_tile_scheduler( ) -> dict[str, FlashMLASchedMeta | None]: """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA C++ decode-path concern) -- always return the all-``None`` sentinel - the shared base's own CPU branch returns.""" + the shared base's own CPU branch returns. + """ return { _LAYER_TYPE_SWAONLY: None, _LAYER_TYPE_C4A: None, diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index a2c3b00493e6..efbd6768f62a 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -176,13 +176,15 @@ def _fused_indexer_q_rope_quant_cpu( class _NoOpEvent: """Stand-in for ``torch.cuda.Event`` on CPU: since CPU never sets an aux - stream list, these events are constructed but never actually used.""" + stream list, these events are constructed but never actually used. + """ class DeepseekV4CPUIndexerCache(DeepseekV4IndexerCache): """CPU indexer K-cache descriptor: same fields as the shared base, just pointing ``get_attn_backend()`` at ``DeepseekV4CPUIndexerBackend`` - instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``.""" + instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``. + """ def get_attn_backend(self) -> type[AttentionBackend]: return DeepseekV4CPUIndexerBackend @@ -209,7 +211,8 @@ def forward( compressor run straight-line instead of through ``maybe_execute_in_parallel``; the short-context check drops the CUDA-only ``is_current_stream_capturing()`` guard (always False - here).""" + here). + """ compressor = self.compressor attn_metadata = get_forward_context().attn_metadata @@ -309,7 +312,8 @@ def _wrap_wo_a_process_weights_after_loading(self) -> None: """Snapshot and pack ``wo_a``'s weight into a bf16 copy for ``bmm_cpu`` before the FP8 kernel's own ``process_weights_after_loading`` VNNI-repacks it in place for - row-major reads ``_o_proj`` never does.""" + row-major reads ``_o_proj`` never does. + """ wo_a = self.wo_a quant_method = wo_a.quant_method orig_pwal = quant_method.process_weights_after_loading @@ -400,7 +404,8 @@ def _fused_qnorm_rope_kv_insert( ) -> torch.Tensor: """CPU override: only the fp8_ds_mla (uint8) SWA cache layout is supported here, so the base method's bf16/per-tensor-fp8 branches - are dropped.""" + are dropped. + """ if not isinstance(attn_metadata, dict): # Profile run: kernel doesn't fire; produce a padded tensor so # downstream FlashMLA gets the right shape. @@ -445,7 +450,8 @@ def _split_qkv_and_norm( ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """CPU override: two ``RMSNorm`` calls instead of the shared ``fused_q_kv_rmsnorm``, whose raw ``@triton.jit`` kernel fails to - link under triton-cpu (``undefined symbol: __truncdfbf2``).""" + link under triton-cpu (``undefined symbol: __truncdfbf2``). + """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) return self.q_norm(qr), None, self.kv_norm(kv) @@ -463,7 +469,8 @@ def _prepare_and_attn( ) -> None: """CPU override: no aux streams, so query-projection+KV-insert, the indexer, and the compressor run straight-line instead of through - execute_in_parallel/maybe_execute_in_parallel.""" + execute_in_parallel/maybe_execute_in_parallel. + """ attn_metadata = get_forward_context().attn_metadata indexer = self.indexer compressor = self.compressor @@ -512,7 +519,8 @@ def _sparse_indexer_and_attn( out: torch.Tensor, ) -> None: """CPU override: identical body, minus ``@eager_break_during_capture`` - (this platform never captures a CUDA graph, so it was a no-op).""" + (this platform never captures a CUDA graph, so it was a no-op). + """ if self.indexer is not None and index_q is not None: assert index_weights is not None # index_q_scale is always None on CPU (FP8-only path). diff --git a/vllm/models/deepseek_v4/cpu/cpu_utils.py b/vllm/models/deepseek_v4/cpu/cpu_utils.py index f03bf60b5846..9a5cff9abb89 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_utils.py +++ b/vllm/models/deepseek_v4/cpu/cpu_utils.py @@ -23,6 +23,7 @@ def map_local_to_global_slots_cpu( Returns: [N, K] int64 global slot ids, -1 where ``local_indices`` was invalid. + """ valid = local_indices >= 0 safe_local = local_indices.clamp(min=0).to(torch.int64) diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 6e88d21c2794..a71b6eb63564 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -907,7 +907,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/nvidia/fi_moe.py b/vllm/models/deepseek_v4/nvidia/fi_moe.py index 2e94e75e7dd9..96515374d682 100644 --- a/vllm/models/deepseek_v4/nvidia/fi_moe.py +++ b/vllm/models/deepseek_v4/nvidia/fi_moe.py @@ -159,7 +159,8 @@ def __init__( def _realloc_nvfp4_params(self) -> None: """Swap the mx-recipe scale params for the NVFP4 checkpoint's: fp8-e4m3 per-16 block scales plus the per-tensor second-level - scales (weight_scale_2, input_scale) modelopt exports.""" + scales (weight_scale_2, input_scale) modelopt exports. + """ n_e = self.num_local_experts inter = self.intermediate_size hidden = self.hidden_size diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 883ef5ecafd4..70dc01efabf7 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -56,7 +56,8 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: def _packed_block_span(pool: torch.Tensor) -> int: """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == block_size for unpacked KV, larger when packed (#44577). Raises if not - token-aligned.""" + token-aligned. + """ block_stride = pool.stride(0) token_stride = pool.stride(-2) if block_stride % token_stride != 0: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 4aba14cd03b1..3bc5399f5378 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1985,7 +1985,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 55403559dc65..a19cc707b254 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -518,8 +518,7 @@ def process_weights_after_loading(self) -> None: self.finalize_mega_moe_weights() def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index f3e468ebb1f5..ba47bd9631c2 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1378,7 +1378,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py index 3d1ce97687b9..a19d0d578220 100644 --- a/vllm/models/deepseek_v4/xpu/mtp.py +++ b/vllm/models/deepseek_v4/xpu/mtp.py @@ -493,8 +493,7 @@ def finalize_mega_moe_weights(self) -> None: layer.mtp_block.ffn.finalize_mega_moe_weights() def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: - """ - Rewrite the weight name to match the format of the original model. + """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index b3eddc202452..50bd8e7cb7ec 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1062,7 +1062,8 @@ def get_model_state_cls(): @property def token_lookback_depth(self) -> int: """Tokens before a chunk start the engram hash needs; the model runner - passes them as `lookback_token_ids`.""" + passes them as `lookback_token_ids`. + """ engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 @@ -1086,7 +1087,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-collapse residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 3227b50a4e45..0e379037b490 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -209,7 +209,8 @@ def _build_image_span( self, image_embeds: torch.Tensor, types: torch.Tensor ) -> torch.Tensor: """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END. + """ types = types.to(image_embeds.device) span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) dtype = image_embeds.dtype diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 752b335eafd8..77f5958e5f58 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DeepseekV4 MLA Attention Layer -""" +"""DeepseekV4 MLA Attention Layer""" import math from abc import ABC, abstractmethod @@ -1068,7 +1066,8 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: def _compressed_kv_cache(self) -> torch.Tensor: """The compressed-KV cache tensor of this layer's kv source (own - cache for kv-source layers).""" + cache for kv-source layers). + """ if self.is_kv_source: return self.kv_cache assert self.compressed_cache_prefix is not None diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index c415bbaa6e42..857063ecfb14 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -564,7 +564,8 @@ def _engram_head_shard_weight_loader( param: torch.nn.Parameter, loaded_weight: torch.Tensor ) -> None: """Load this rank's complete head buckets. ue8m0 scales arrive as - float8_e8m0fnu; keep the raw bytes (the param stores uint8).""" + float8_e8m0fnu; keep the raw bytes (the param stores uint8). + """ part_rows = param.shape[0] if loaded_weight.dtype == torch.float8_e8m0fnu: loaded_weight = loaded_weight.view(torch.uint8) @@ -739,7 +740,8 @@ def lookup( def forward(self, indices: torch.Tensor) -> torch.Tensor: """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] - bf16, gathered from all shards for this replica's tokens.""" + bf16, gathered from all shards for this replica's tokens. + """ out = torch.empty( (indices.shape[0], self.part_n_hash_cols, self.dim), dtype=torch.bfloat16, @@ -1000,7 +1002,8 @@ def forward( ) -> torch.Tensor: """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all tokens, pre sequence-parallel shard); token_mask: [T], False shuts - the gate so those positions pass through untouched.""" + the gate so those positions pass through untouched. + """ kv = self.wkv(self.embed(hash_ids).flatten(-2)) num_kv_tokens = hash_ids.shape[0] assert token_mask is None or token_mask.shape == (num_kv_tokens,) diff --git a/vllm/models/deepseek_v41/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py index 4f87b816c9a5..ce56713e8c23 100644 --- a/vllm/models/deepseek_v41/common/mm_preprocess.py +++ b/vllm/models/deepseek_v41/common/mm_preprocess.py @@ -82,7 +82,8 @@ def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int: def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): """Largest aspect-preserving pixel size whose token grid still fits in - max_n_token. Returns (best_height, best_width).""" + max_n_token. Returns (best_height, best_width). + """ r = height / width max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 max_h_float = max_w_float * r @@ -105,7 +106,8 @@ def safe_resize( height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token ): """Shrink the pixel size until the image costs at most max_n_token LLM - tokens.""" + tokens. + """ n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio) if num_image_tokens(n_llm_h, n_llm_w) > max_n_token: best_height, best_width = solve_resize_ratio( @@ -171,7 +173,8 @@ def image_token_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor: class DeepseekV4VLImageProcessor: """Per-image transform (the PIL-input equivalent of the reference - ``load_image``).""" + ``load_image``). + """ def __init__(self, config: DeepseekV41Config) -> None: super().__init__() diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 4a9da21c036b..9a251e0748c9 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index +"""Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index preparation. - quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into @@ -76,8 +75,7 @@ def quantize_and_insert_k_kernel( n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) use_fnuz: tl.constexpr = False, ): - """ - Quantize K tensor and insert into paged K cache. + """Quantize K tensor and insert into paged K cache. K Cache block layout (block_size=64 tokens): - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 @@ -251,8 +249,7 @@ def quantize_and_insert_k_cache( use_fnuz: bool = False, bytes_per_token: int = V4_BYTES_PER_TOKEN, ): - """ - Quantize K tensor and insert into paged K cache. + """Quantize K tensor and insert into paged K cache. ``bytes_per_token`` picks the record (see the module header): the V4 one, or V4.1's all-dims MXFP8 one. diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 0f7830cf58c0..2ba9dbea3de7 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -58,6 +58,7 @@ def fused_save_compress_norm( rms_norm_eps: RMSNorm epsilon. compress_ratio: Group size, either 1 or 2. latent_out: BF16 [tokens, 512], written only at valid group boundaries. + """ assert compress_ratio in (1, 2) assert kv_score.dtype == torch.float32 diff --git a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py index c4bd98b9e99e..ba8ddbd95a9a 100644 --- a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py +++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py @@ -46,11 +46,13 @@ def indexer_k_norm_rope_store( cos_sin_cache: [max_pos, rope_head_dim] GPT-J layout (cos half, then sin half), from the layer's compress-RoPE instance. rms_norm_weight: [128] k_norm weight. + rms_norm_eps: Epsilon of the k_norm RMSNorm. k_cache: uint8 paged indexer cache [num_blocks, block_size, row_bytes]. kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip). compress_ratio: group size; keys are emitted at group boundaries. use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else per-token FP8 with a single fp32 scale. + """ num_tokens = kv_slot_mapping.numel() assert k_pre.ndim == 2 and k_pre.shape[1] == 128 diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0b43cfad190b..650581a7dfd0 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -54,7 +54,8 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: def _packed_block_span(pool: torch.Tensor) -> int: """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == block_size for unpacked KV, larger when packed (#44577). Raises if not - token-aligned.""" + token-aligned. + """ block_stride = pool.stride(0) token_stride = pool.stride(-2) if block_stride % token_stride != 0: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 21fa00e5c963..4d9d9d3108ea 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -1179,7 +1179,8 @@ def get_model_state_cls(): @property def token_lookback_depth(self) -> int: """Tokens before a chunk start the engram hash needs; the model runner - passes them as `lookback_token_ids`.""" + passes them as `lookback_token_ids`. + """ engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 @@ -1203,7 +1204,8 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-collapse residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step.""" + forward(); valid after each target step. + """ return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index bfbcaa89a1a1..21386790cb40 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -209,7 +209,8 @@ def _build_image_span( self, image_embeds: torch.Tensor, types: torch.Tensor ) -> torch.Tensor: """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END. + """ types = types.to(image_embeds.device) span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) dtype = image_embeds.dtype diff --git a/vllm/models/dots3_note/nvidia/audio_encoder.py b/vllm/models/dots3_note/nvidia/audio_encoder.py index dbf7eaded26e..c690c87b7f69 100644 --- a/vllm/models/dots3_note/nvidia/audio_encoder.py +++ b/vllm/models/dots3_note/nvidia/audio_encoder.py @@ -194,7 +194,7 @@ def forward(self, input_ids, past_key_values_length=0): class WhisperAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" + """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__( self, @@ -346,8 +346,7 @@ def forward( rotary_cos: torch.Tensor | None = None, rotary_sin: torch.Tensor | None = None, ) -> tuple[Any, ...]: - """ - Args: + """Args: hidden_states: Input to the layer of shape `(batch, seq_len, embed_dim)`, or `(total_tokens, embed_dim)` when `cu_seqlens_q` is given. @@ -365,6 +364,7 @@ def forward( Returns: A tuple of the output hidden states, followed by the attention weights if `output_attentions` is `True`. + """ residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) @@ -435,12 +435,12 @@ def _init_weights(self, module): class DotsSpeechEncoder(DotsSpeechPreTrainedModel): - """ - Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + """Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`WhisperEncoderLayer`]. Args: config: WhisperConfig + """ def __init__(self, config: WhisperConfig): diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 09e7a40aa5e0..7e78a8fe3294 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -98,7 +98,8 @@ def _kpool_softmax_rotate_write_cache_kernel( BLOCK_D: tl.constexpr, ): """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> - Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" + Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``. + """ row = tl.program_id(0) do_write = True if HAS_WRITE_MASK: @@ -212,6 +213,13 @@ def kpool_compress_and_write_cache( slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score. ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias. loc: ``[n_pools]`` int64 — flat physical slot per pool. + pool_size: Number of tokens compressed into one cache entry. + head_dim: Indexer head dimension. + write_mask: ``[n_pools]`` bool — pools to write, or None for all. + round_scale: Round each fp8 scale down to a power of two. + return_compressed: Also return the compressed K and scales. + write_cache: Write the compressed result into ``kv_cache``. + """ assert slot_k.ndim == 3 assert slot_score.shape == slot_k.shape @@ -592,6 +600,10 @@ def kpool_decode_update_and_maybe_write_cache_batched( ape: ``[pool_size, head_dim]`` fp32. slot_mapping: ``[num_requests, next_n]`` int32. positions: ``[num_requests, next_n]`` int32. + pool_size: Number of tokens compressed into one cache entry. + head_dim: Indexer head dimension. + round_scale: Round each fp8 scale down to a power of two. + """ num_requests, next_n = key.shape[0], key.shape[1] if num_requests == 0 or next_n == 0: diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py index d7de85984f85..d8b25211e0c1 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py @@ -564,8 +564,7 @@ def fused_recurrent_gated_delta_rule( num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: + r"""Args: q (torch.Tensor): queries of shape `[B, T, H, K]`. k (torch.Tensor): @@ -627,6 +626,7 @@ def fused_recurrent_gated_delta_rule( initial_state=h0, cu_seqlens=cu_seqlens ) + """ if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index be417d6b0526..ea86750e7d0a 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -406,19 +406,24 @@ def chunk_kda_scaled_dot_kkt_fwd( chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Compute beta * K * K^T. + r"""Compute beta * K * K^T. Args: + q (torch.Tensor): + The query tensor of shape `[B, T, H, K]`. k (torch.Tensor): The key tensor of shape `[B, T, H, K]`. beta (torch.Tensor): The beta tensor of shape `[B, T, H]`. gk (torch.Tensor): The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + scale (float): + Scale applied to the query-key products. Default: `None`. cu_seqlens (torch.Tensor): The cumulative sequence lengths of the input tensor. Default: None + chunk_indices (torch.Tensor): + Precomputed chunk indices for `cu_seqlens`. Default: `None`. chunk_size (int): The chunk size. Default: 64. output_dtype (torch.dtype): @@ -426,6 +431,7 @@ def chunk_kda_scaled_dot_kkt_fwd( Returns: beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ B, T, H, K = k.shape assert K <= 256 @@ -1318,16 +1324,15 @@ def fused_kda_gate( safe_gate: bool = False, lower_bound: float | None = -5.0, ) -> torch.Tensor: - """ - Forward pass for KDA gate: - input g: [..., H*D] - param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter (softplus branch only) - threshold: softplus threshold parameter (softplus branch only) - safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias); - when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias)) - lower_bound: floor for the safe_gate branch (default -5.0) - return : [..., H, D] + """Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter (softplus branch only) + threshold: softplus threshold parameter (softplus branch only) + safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias); + when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias)) + lower_bound: floor for the safe_gate branch (default -5.0) + return : [..., H, D] """ orig_shape = g.shape[:-1] diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 436b9685cb12..03e423b13cb3 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -130,7 +130,8 @@ def _resolve_kda_prefill_backend( """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton ``chunk_kda_with_fused_gate`` path. ``backend`` comes from - ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda). + """ if backend not in ("auto", "triton", "flashkda"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") capability = current_platform.get_device_capability() @@ -368,7 +369,8 @@ def _flashkda_prefill( bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the attention output into ``out`` (a workspace buffer when ``None``) and - returns ``(out, final_state)``.""" + returns ``(out, final_state)``. + """ assert self._flashkda_buffer_specs is not None final_state, workspace, workspace_out = ( current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py index d2b21345d052..42cabc35b59f 100644 --- a/vllm/models/glm5next/common/multimodal.py +++ b/vllm/models/glm5next/common/multimodal.py @@ -720,7 +720,8 @@ class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor): (image token repeat, video frame/timestamp structure) is owned by vLLM's prompt-update machinery — the inherited ``_get_prompt_updates`` builds the replacement content and the placeholder scan validates against - exactly that.""" + exactly that. + """ def _hf_processor_applies_updates( self, diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py index ad3845457063..703cbb321a82 100644 --- a/vllm/models/glm5next/common/sparse_indexer.py +++ b/vllm/models/glm5next/common/sparse_indexer.py @@ -138,7 +138,8 @@ def _gather_workspace_shapes( """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + - (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales. + """ if use_fp4_cache: return ( ((total_seq_lens, head_dim // 2), torch.uint8), @@ -156,7 +157,8 @@ def kv_cache_as_quant_view( use_fp4_cache: bool, ) -> torch.Tensor: """4D ``[num_blocks, block_size, 1, head_width]`` view expected by - DeepGEMM, from the 3D indexer kv-cache allocation.""" + DeepGEMM, from the 3D indexer kv-cache allocation. + """ if use_fp4_cache: assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 num_blocks, block_size, _ = kv_cache.shape diff --git a/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py index 095db3917b49..f069fdf2c25f 100644 --- a/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py +++ b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py @@ -31,7 +31,8 @@ def _fused_eh_norm_kernel( ): """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` - ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat. + """ tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 70b41021aab0..010b231354ac 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -116,6 +116,7 @@ def fwht128_quant_fp8(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: Returns: (q_fp8 ``[rows, 128]`` float8_e4m3fn, scale ``[rows, 1]`` float32). + """ assert q.ndim == 2 and q.shape[1] == 128, q.shape assert q.dtype == torch.bfloat16 @@ -162,7 +163,8 @@ def _kpool_softmax_rotate_write_cache_kernel( BLOCK_D: tl.constexpr, ): """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> - Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" + Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``. + """ row = tl.program_id(0) do_write = True if HAS_WRITE_MASK: @@ -276,6 +278,13 @@ def kpool_compress_and_write_cache( slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score. ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias. loc: ``[n_pools]`` int64 — flat physical slot per pool. + pool_size: Number of tokens compressed into one cache entry. + head_dim: Indexer head dimension. + write_mask: ``[n_pools]`` bool — pools to write, or None for all. + round_scale: Round each fp8 scale down to a power of two. + return_compressed: Also return the compressed K and scales. + write_cache: Write the compressed result into ``kv_cache``. + """ assert slot_k.ndim == 3 assert slot_score.shape == slot_k.shape @@ -642,6 +651,10 @@ def kpool_decode_update_and_maybe_write_cache_batched( ape: ``[pool_size, head_dim]`` fp32. slot_mapping: ``[num_requests, next_n]`` int32. positions: ``[num_requests, next_n]`` int32. + pool_size: Number of tokens compressed into one cache entry. + head_dim: Indexer head dimension. + round_scale: Round each fp8 scale down to a power of two. + """ num_requests, next_n = key.shape[0], key.shape[1] if num_requests == 0 or next_n == 0: diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py index 152412b8ca8f..a9d0fba476ef 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py @@ -590,8 +590,7 @@ def fused_recurrent_gated_delta_rule( num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: + r"""Args: q (torch.Tensor): queries of shape `[B, T, H, K]`. k (torch.Tensor): @@ -653,6 +652,7 @@ def fused_recurrent_gated_delta_rule( initial_state=h0, cu_seqlens=cu_seqlens ) + """ if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py index b792f606d4e0..9805b3747130 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py @@ -416,19 +416,24 @@ def chunk_kda_scaled_dot_kkt_fwd( chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Compute beta * K * K^T. + r"""Compute beta * K * K^T. Args: + q (torch.Tensor): + The query tensor of shape `[B, T, H, K]`. k (torch.Tensor): The key tensor of shape `[B, T, H, K]`. beta (torch.Tensor): The beta tensor of shape `[B, T, H]`. gk (torch.Tensor): The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + scale (float): + Scale applied to the query-key products. Default: `None`. cu_seqlens (torch.Tensor): The cumulative sequence lengths of the input tensor. Default: None + chunk_indices (torch.Tensor): + Precomputed chunk indices for `cu_seqlens`. Default: `None`. chunk_size (int): The chunk size. Default: 64. output_dtype (torch.dtype): @@ -436,6 +441,7 @@ def chunk_kda_scaled_dot_kkt_fwd( Returns: beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ B, T, H, K = k.shape assert K <= 256 @@ -1328,16 +1334,15 @@ def fused_kda_gate( safe_gate: bool = False, lower_bound: float | None = -5.0, ) -> torch.Tensor: - """ - Forward pass for KDA gate: - input g: [..., H*D] - param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter (softplus branch only) - threshold: softplus threshold parameter (softplus branch only) - safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias); - when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias)) - lower_bound: floor for the safe_gate branch (default -5.0) - return : [..., H, D] + """Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter (softplus branch only) + threshold: softplus threshold parameter (softplus branch only) + safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias); + when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias)) + lower_bound: floor for the safe_gate branch (default -5.0) + return : [..., H, D] """ orig_shape = g.shape[:-1] diff --git a/vllm/models/hy_v4/nvidia/attention.py b/vllm/models/hy_v4/nvidia/attention.py index 40a05f2c6407..5e0a33f18373 100644 --- a/vllm/models/hy_v4/nvidia/attention.py +++ b/vllm/models/hy_v4/nvidia/attention.py @@ -67,6 +67,7 @@ def compute_skip_topk_layers(config: PretrainedConfig) -> set[int]: Raises: ValueError: If ``indexer_types`` has the wrong length or an unknown entry, or if ``index_topk_freq`` is not a positive integer. + """ if not hasattr(config, "index_topk"): return set() @@ -115,6 +116,7 @@ def is_skip_topk_indexer_weight(weight_name: str, skip_topk_layers: set[int]) -> Returns: True when the weight is an indexer weight of a layer that has no indexer module and therefore must be dropped. + """ if ".indexer." not in weight_name or not skip_topk_layers: return False @@ -557,6 +559,7 @@ def _resolve_sink_backend( The backend class to bind, or None when no sink-capable backend is available; the caller then loads the sink weight but disables the bias. + """ head_size = self.kv_lora_rank + self.qk_rope_head_dim dtype = torch.get_default_dtype() diff --git a/vllm/models/hy_v4/nvidia/flashmla_sparse.py b/vllm/models/hy_v4/nvidia/flashmla_sparse.py index 952dd8de5c08..83871755527b 100644 --- a/vllm/models/hy_v4/nvidia/flashmla_sparse.py +++ b/vllm/models/hy_v4/nvidia/flashmla_sparse.py @@ -99,6 +99,7 @@ def _validate_sinks(sinks: torch.Tensor | None, num_heads: int) -> None: Raises: ValueError: If the dtype is not float32 or the shape is not ``(num_heads,)``. + """ if sinks is None: return @@ -134,6 +135,7 @@ def _sinks_for_query( Raises: ValueError: If the sink and query head layouts disagree, or if they live on different devices. + """ sinks = self.sinks if sinks is None: diff --git a/vllm/models/hy_v4/nvidia/hc.py b/vllm/models/hy_v4/nvidia/hc.py index 130d3842ecfc..5c7fcc931ef4 100644 --- a/vllm/models/hy_v4/nvidia/hc.py +++ b/vllm/models/hy_v4/nvidia/hc.py @@ -107,6 +107,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: Returns: A tuple of the pre-gated reduction ``[num_tokens, d]`` and the post gates ``[num_tokens, hc]`` consumed by `HYV4HCPostLayer`. + """ if self.hpc_op is not None: return self.hpc_op(x) @@ -186,6 +187,7 @@ def forward( Returns: The updated residual channels ``[num_tokens, hc, d]``. + """ if self.hpc_op is not None: return self.hpc_op(x, residual, post) @@ -272,6 +274,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: The merged hidden state ``[num_tokens, d]``. + """ if self.hpc_op is not None: return self.hpc_op(x) @@ -367,6 +370,7 @@ def pre( A tuple of the reduced hidden states ``[num_tokens, d]``, the post gates ``[num_tokens, hc]`` (``None`` when iHC is disabled) and the residual (the untouched input). + """ if not self.enable_ihc: return hidden_states, None, hidden_states diff --git a/vllm/models/hy_v4/nvidia/model.py b/vllm/models/hy_v4/nvidia/model.py index 626d704f0e74..b9988da377ef 100644 --- a/vllm/models/hy_v4/nvidia/model.py +++ b/vllm/models/hy_v4/nvidia/model.py @@ -189,7 +189,7 @@ def _forward_ihc( positions: torch.Tensor, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, None]: - """iHC forward: each sub-block reduces and re-scatters the channels.""" + """IHC forward: each sub-block reduces and re-scatters the channels.""" hidden_states = self.hc_attn_layer.prepare_input(hidden_states) hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states) hidden_states = self.input_layernorm(hidden_states) diff --git a/vllm/models/hy_v4/nvidia/mtp.py b/vllm/models/hy_v4/nvidia/mtp.py index 49f1b3d7a508..6c9ea11e81d4 100644 --- a/vllm/models/hy_v4/nvidia/mtp.py +++ b/vllm/models/hy_v4/nvidia/mtp.py @@ -106,6 +106,7 @@ def _resolve_fused_expert_param( Returns: The matching draft parameter name, or None when there is none. + """ if not ckpt_suffix: return param_base if param_base in params_dict else None @@ -164,6 +165,7 @@ def _create_mtp_quant_config( Returns: The quantization config to use for the MTP layers, or None when the MTP layers are unquantized. + """ mtp_quant_algo = getattr(hf_config, "mtp_quant_algo", None) @@ -241,6 +243,7 @@ def _remap_mtp_quant_exclusions( Returns: A shallow copy with the translated exclusions, or the input unchanged. + """ if quant_config is None: return None @@ -671,6 +674,7 @@ def _load_expert_weight( Returns: True when the weight was consumed (even if this rank holds none of the addressed experts). + """ base = name.split(".experts.")[0] for ckpt_proj, tag in ( diff --git a/vllm/models/inkling/amd/logits_processor.py b/vllm/models/inkling/amd/logits_processor.py index 9cf1c663c891..4302c6f674f8 100644 --- a/vllm/models/inkling/amd/logits_processor.py +++ b/vllm/models/inkling/amd/logits_processor.py @@ -37,6 +37,7 @@ class InklingLogitsProcessor(LogitsProcessor): soft_cap: Optional logit soft cap (``None`` for the served checkpoint). logits_mup_width_multiplier: muP width divisor for the final logits; ``None`` or ``0`` disables it. + """ def __init__( diff --git a/vllm/models/inkling/amd/model.py b/vllm/models/inkling/amd/model.py index 711a33855cfc..48ed8a2caab3 100644 --- a/vllm/models/inkling/amd/model.py +++ b/vllm/models/inkling/amd/model.py @@ -72,7 +72,8 @@ def _sconv_add_norm( ROCm uses the portable collective path. The Lamport P2P implementation in the NVIDIA model relies on CUDA GDC and CUDA-specific symmetric-memory - publication semantics which cannot be linked into a gfx950 HSACO.""" + publication semantics which cannot be linked into a gfx950 HSACO. + """ norm_w = norm.weight if norm is not None else None eps = norm.variance_epsilon if norm is not None else 0.0 diff --git a/vllm/models/inkling/amd/moe.py b/vllm/models/inkling/amd/moe.py index d8141eb4cf08..3a71b3ea5ecc 100644 --- a/vllm/models/inkling/amd/moe.py +++ b/vllm/models/inkling/amd/moe.py @@ -192,7 +192,8 @@ def inkling_gate_select( class InklingGate(nn.Module): """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and - global scale (the served checkpoint's only configuration).""" + global scale (the served checkpoint's only configuration). + """ def __init__( self, @@ -246,7 +247,8 @@ def select_experts( self, gating_output: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Full selection: (weights, ids) of [T, K + S]. The first K entries - are the routed top-k; the S trailing entries are the sink gammas.""" + are the routed top-k; the S trailing entries are the sink gammas. + """ return inkling_gate_select( gating_output, self.n_total_experts, @@ -267,7 +269,8 @@ def select_experts( def _inkling_moe_ep_size() -> int: """EP size the MoERunner layer will run with (mirrors FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when - expert parallelism is enabled).""" + expert parallelism is enabled). + """ parallel_config = get_current_vllm_config().parallel_config if not parallel_config.enable_expert_parallel: return 1 diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py index f9cba37f3c71..dc7a8fd1e265 100644 --- a/vllm/models/inkling/amd/mtp.py +++ b/vllm/models/inkling/amd/mtp.py @@ -166,7 +166,8 @@ def embed_input_ids( ) -> torch.Tensor: """Draft-prefill embedding: fused gather + backbone embed_norm, then the target's tower embeddings scattered in unnormed (the backbone - convention — MM embeds are merged after embed_norm).""" + convention — MM embeds are merged after embed_norm). + """ norm = self.backbone_embed_norm embeds = embed_rmsnorm( input_ids, @@ -195,7 +196,8 @@ def fused_input_cat( """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one launch: embedding row gather + the backbone embed_norm + the depth embed_norm chain on one side, hidden_norm on the other, written - straight into the cat buffer.""" + straight into the cat buffer. + """ hidden_w = layer.hidden_norm.weight embed_w = layer.embed_norm.weight eps = layer.hidden_norm.variance_epsilon @@ -327,7 +329,8 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is - skipped entirely.""" + skipped entirely. + """ return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/inkling/amd/ops/norm.py b/vllm/models/inkling/amd/ops/norm.py index d34c1c3259b9..caee8a047e2d 100644 --- a/vllm/models/inkling/amd/ops/norm.py +++ b/vllm/models/inkling/amd/ops/norm.py @@ -238,7 +238,8 @@ def embed_rmsnorm( embedding-table gather. ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the first decoder layer's pre-attention norm) as a second output, still one - launch. Bit-exact vs the unfused module sequence.""" + launch. Bit-exact vs the unfused module sequence. + """ ids = input_ids.view(-1) (T,) = ids.shape n = embed_table.shape[1] @@ -332,7 +333,8 @@ def embed_dual_rmsnorm_cat( (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the backbone embed_norm in front of the depth embed_norm (bit-exact vs the - unfused sequence). The concat copies collapse into direct writes.""" + unfused sequence). The concat copies collapse into direct writes. + """ T, n = hidden.shape if embeds is not None: assert embeds.shape == hidden.shape diff --git a/vllm/models/inkling/nvidia/logits_processor.py b/vllm/models/inkling/nvidia/logits_processor.py index 9cf1c663c891..4302c6f674f8 100644 --- a/vllm/models/inkling/nvidia/logits_processor.py +++ b/vllm/models/inkling/nvidia/logits_processor.py @@ -37,6 +37,7 @@ class InklingLogitsProcessor(LogitsProcessor): soft_cap: Optional logit soft cap (``None`` for the served checkpoint). logits_mup_width_multiplier: muP width divisor for the final logits; ``None`` or ``0`` disables it. + """ def __init__( diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index aeb460d64a21..4363cc883c4f 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -75,7 +75,8 @@ def _sconv_add_norm( """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``. The Lamport path performs reduce-scatter + shard sconv + all-gather + - residual add + norm. The NCCL path handles unsupported configurations.""" + residual add + norm. The NCCL path handles unsupported configurations. + """ attn_metadata = get_forward_context().attn_metadata m = ( attn_metadata.get(sconv.owner.prefix) diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index a8ea968076bb..07c4ce59691c 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -191,7 +191,8 @@ def inkling_gate_select( class InklingGate(nn.Module): """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and - global scale (the served checkpoint's only configuration).""" + global scale (the served checkpoint's only configuration). + """ def __init__( self, @@ -243,7 +244,8 @@ def select_experts( self, gating_output: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Full selection: (weights, ids) of [T, K + S]. The first K entries - are the routed top-k; the S trailing entries are the sink gammas.""" + are the routed top-k; the S trailing entries are the sink gammas. + """ return inkling_gate_select( gating_output, self.n_total_experts, @@ -264,7 +266,8 @@ def select_experts( def _inkling_moe_ep_size() -> int: """EP size the FusedMoE layer will run with (mirrors FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when - expert parallelism is enabled).""" + expert parallelism is enabled). + """ parallel_config = get_current_vllm_config().parallel_config if not parallel_config.enable_expert_parallel: return 1 diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index f9cba37f3c71..dc7a8fd1e265 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -166,7 +166,8 @@ def embed_input_ids( ) -> torch.Tensor: """Draft-prefill embedding: fused gather + backbone embed_norm, then the target's tower embeddings scattered in unnormed (the backbone - convention — MM embeds are merged after embed_norm).""" + convention — MM embeds are merged after embed_norm). + """ norm = self.backbone_embed_norm embeds = embed_rmsnorm( input_ids, @@ -195,7 +196,8 @@ def fused_input_cat( """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one launch: embedding row gather + the backbone embed_norm + the depth embed_norm chain on one side, hidden_norm on the other, written - straight into the cat buffer.""" + straight into the cat buffer. + """ hidden_w = layer.hidden_norm.weight embed_w = layer.embed_norm.weight eps = layer.hidden_norm.variance_epsilon @@ -327,7 +329,8 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is - skipped entirely.""" + skipped entirely. + """ return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/inkling/nvidia/ops/mm_towers.py b/vllm/models/inkling/nvidia/ops/mm_towers.py index 021cf5274bd2..76a339377c52 100644 --- a/vllm/models/inkling/nvidia/ops/mm_towers.py +++ b/vllm/models/inkling/nvidia/ops/mm_towers.py @@ -60,7 +60,8 @@ def dmel_embed_sum_norm( eps: float, ) -> torch.Tensor: """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no - [T, NB, D] intermediate).""" + [T, NB, D] intermediate). + """ T, nb = dmel_idx.shape D = weight.shape[1] assert weight.shape[0] % nb == 0 @@ -151,7 +152,8 @@ def rmsnorm_gelu( ) -> torch.Tensor: """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block when D is small. With ``fold``, x must be [N, T, H, W, D] and the output - comes back as ``fold_timespace_to_depth(result, *fold)``.""" + comes back as ``fold_timespace_to_depth(result, *fold)``. + """ D = x.shape[-1] flat = x.reshape(-1, D) assert flat.stride(1) == 1 and flat.stride(0) == D diff --git a/vllm/models/inkling/nvidia/ops/norm.py b/vllm/models/inkling/nvidia/ops/norm.py index d34c1c3259b9..caee8a047e2d 100644 --- a/vllm/models/inkling/nvidia/ops/norm.py +++ b/vllm/models/inkling/nvidia/ops/norm.py @@ -238,7 +238,8 @@ def embed_rmsnorm( embedding-table gather. ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the first decoder layer's pre-attention norm) as a second output, still one - launch. Bit-exact vs the unfused module sequence.""" + launch. Bit-exact vs the unfused module sequence. + """ ids = input_ids.view(-1) (T,) = ids.shape n = embed_table.shape[1] @@ -332,7 +333,8 @@ def embed_dual_rmsnorm_cat( (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the backbone embed_norm in front of the depth embed_norm (bit-exact vs the - unfused sequence). The concat copies collapse into direct writes.""" + unfused sequence). The concat copies collapse into direct writes. + """ T, n = hidden.shape if embeds is not None: assert embeds.shape == hidden.shape diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index 83c953e1a1d4..96e71700cd7d 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -61,9 +61,7 @@ def _shard_up_proj_tail( shared_output: torch.Tensor, trunc_size: int | None, ) -> torch.Tensor: - """ - Tier 2: column-parallel up-projection folded into the final reduce. - """ + """Tier 2: column-parallel up-projection folded into the final reduce.""" if not self._logged_sharded_tail: self._logged_sharded_tail = True logger.info_once( diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index a74859efb884..2a013ed0fc08 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -317,9 +317,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class KimiMLAAttention(nn.Module): - """ - Main reference: DeepseekV2 vllm Implementation - """ + """Main reference: DeepseekV2 vllm Implementation.""" def __init__( self, diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index 748581771b82..556373c084aa 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -213,9 +213,17 @@ def fused_kda_chunk( Args: qg: ``q * exp2(gk_cumsum)``, ``[1, T, H, 128]``. + w: chunk-local WY representation of the gated keys. + u: chunk-local WY representation of the values. kg_t: chunk-major transposed gated keys, ``[chunks, H, 128, 64]``. + aqk: intra-chunk query-key attention, ``[chunks, H, 64, 64]``. decay: ``exp2`` of each chunk's last gate row, ``[chunks, H, 128]``. out: output buffer, ``[1, T, H, 128]``; may alias ``u``'s source. + scale: scale applied to the query-key products. + cu_seqlens: int32 cumulative sequence lengths. + initial_state: fp32 per-sequence initial recurrent state, or ``None``. + output_final_state: whether to return the final recurrent state. + chunk_offsets: int32 per-sequence first chunk index. checkpoint_state: destination for the mid-prefill state snapshots, fp32 ``[rows, H, 128, 128]``. Without ``checkpoint_state_indices`` it is a staging buffer indexed by @@ -239,6 +247,7 @@ def fused_kda_chunk( has_initial_state: bool ``[N]``; a false entry starts that sequence from a zero state and its cache row is read only, never before the walk writes it. + """ if state_cache is not None: if initial_state is not None or output_final_state: diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py index 4c85ce8624f6..56cb3e148291 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py +++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py @@ -52,6 +52,22 @@ def chunk_kda_prefill( """Run chunk KDA from raw gate and beta projections. Args: + q: query tensor, ``[1, T, H, K]``. + k: key tensor, ``[1, T, H, K]``. + v: value tensor, ``[1, T, H, V]``. + raw_g: raw gate projection, before the activation. + raw_beta: raw beta projection, before the activation. + A_log: log of the per-head gate decay. + g_bias: optional per-head gate bias. + scale: scale applied to the query-key products. Defaults to + ``k.shape[-1] ** -0.5``. + initial_state: fp32 per-sequence initial recurrent state, or ``None``. + output_final_state: whether to return the final recurrent state. + lower_bound: optional floor applied to the gate. + use_qk_l2norm_in_kernel: fold the q/k L2 norm into the kernel. + cu_seqlens: int32 cumulative sequence lengths. + chunk_indices: precomputed chunk indices for ``cu_seqlens``. + chunk_offsets: int32 per-sequence first chunk index. use_fused_chunk: request the two-kernel ROCm path. It is used only when every one of its preconditions holds; otherwise the Triton path runs unchanged. @@ -71,6 +87,7 @@ def chunk_kda_prefill( Returns: The output and, when requested, the final recurrent state. + """ if scale is None: scale = k.shape[-1] ** -0.5 diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py index 0a43249d6e3f..e8b83d23aee1 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py @@ -910,13 +910,12 @@ def fused_kda_gate( threshold: float = 20.0, lower_bound: float | None = None, ) -> torch.Tensor: - """ - Forward pass for KDA gate: - input g: [..., H*D] - param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter - threshold: softplus threshold parameter - return : [..., H, D] + """Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] """ orig_shape = g.shape[:-1] diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py index ca3bc8f75e33..8b9aa86f9ed9 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -59,8 +59,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( USE_SAFE_GATE: tl.constexpr, SOLVE_TRIL_DOT_PRECISION: tl.constexpr, ): - """ - Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + """Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. This kernel: diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py index cc448d95c0eb..86654f264431 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py @@ -151,8 +151,7 @@ def chunk_kda_fwd_intra_token_parallel( chunk_size: int = 64, sub_chunk_size: int = 16, ) -> None: - """ - Token-parallel implementation: each token gets its own thread block. + """Token-parallel implementation: each token gets its own thread block. Supports both fixed-length and variable-length sequences. Reduces wasted computation on padding. @@ -164,10 +163,12 @@ def chunk_kda_fwd_intra_token_parallel( gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) beta: [B, T, HV] Aqk: [B, T, HV, BT] output tensor to write to + cu_seqlens: cumulative sequence lengths for variable-length input Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) scale: attention scale chunk_size: BT (default 64) sub_chunk_size: BC (default 16) + """ B, T, H, K, HV = *q.shape, gk.shape[2] N = len(cu_seqlens) - 1 if cu_seqlens is not None else B diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index 450e46a8e4ec..8eedb5ba89ed 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -275,9 +275,7 @@ def _shard_up_proj_tail( shared_output: torch.Tensor, trunc_size: int | None, ) -> torch.Tensor: - """ - Tier 2: column-parallel up-projection folded into the final reduce. - """ + """Tier 2: column-parallel up-projection folded into the final reduce.""" transform = self.routed_output_transform assert transform is not None diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index e439ec5b6693..cd55362759aa 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -751,6 +751,7 @@ def run_kda_projection_overlap( Returns: The Q/K/V/G projection, F_B output, and beta projection. + """ num_tokens = hidden_states.shape[0] qkvg_weight = packed_weight[:_KDA_QKVG_SIZE] diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 87bf59749c1d..574757844a6e 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -745,7 +745,8 @@ def _decode_concat_cache( slot_mapping: torch.Tensor, ) -> torch.Tensor: """Fused decode query-concat + latent cache insert, dispatched by cache - dtype (same policy as prefill: fp8 cache -> fp8 query).""" + dtype (same policy as prefill: fp8 cache -> fp8 query). + """ if self.kv_cache_dtype == "fp8_ds_mla": cache = self.kv_cache if cache.dtype != torch.uint8: diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 6a26ce78b4c4..7fe612b65466 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -339,6 +339,7 @@ def forward( residual: Optional tensor of the up-projection's output shape to accumulate into. It is consumed in the GEMM's beta-add epilogue, so adding it costs no extra kernel. + """ if self.norm is not None: hidden_states = self.norm(hidden_states) @@ -796,6 +797,7 @@ def _maybe_overlap_router_and_down_proj( ``router_output`` holds the grouped top-k weights and ``topk_ids`` the selected experts; otherwise ``router_output`` holds the raw gate logits and ``topk_ids`` is ``None``. + """ def _router( diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py index b8eed0d7f4a0..67d5cfdf081d 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -59,7 +59,6 @@ def _mapping( def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None: """Validate constraints imposed by the fused collective mapping.""" - if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0: raise ValueError("collective dimensions must be positive") if hidden_dim % tp_size: @@ -86,7 +85,6 @@ def _select_routed_schedule( max_m: int, ) -> tuple[int, int]: """Match upstream MNNVL's one-cluster-per-token occupancy policy.""" - _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) sm_count = torch.cuda.get_device_properties( torch.accelerator.current_device_index() @@ -886,7 +884,6 @@ def compile_kernel( top_k: int = 0, ) -> None: """Compile the rank/M specialization without retaining caller tensors.""" - key = _compile_key( rank, tp_size, diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py index 6a1ba9d77b44..ea9d86d6fa85 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -79,7 +79,6 @@ def validate_configuration( b_prime_stages: int, ) -> None: """Validate constraints imposed by this BF16 SM100 GEMM.""" - if latent_dim <= 0 or shard_dim <= 0: raise ValueError("GEMM K and N must be positive") if latent_dim % 8 or shard_dim % 8: diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py index 7baec0b42b44..3840caf24a27 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -52,7 +52,6 @@ def to_cute_dynamic_m( mode is symbolic, so changing M within an Op's capacity reuses the same compiled kernel. """ - return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( mode=mode, stride_order=tensor.dim_order(), @@ -98,7 +97,6 @@ def load_global_u32x4( side-effecting prevents loop-invariant motion and common-subexpression elimination across polling iterations. """ - address = pointer.toint(loc=loc, ip=ip) opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" out = llvm.inline_asm( @@ -194,7 +192,6 @@ def finalize_top16_bf16( ip=None, ): """Finalize one Kimi K3 top-16 vector with wide metadata loads.""" - addresses = [ pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) for pointer in (gemm2_vector, route_indices, route_weights) @@ -229,7 +226,6 @@ def store_global_u32x4( ip=None, ) -> None: """Store four packed words to an ordinary or NVLS multicast global VA.""" - words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" llvm.inline_asm( @@ -294,7 +290,6 @@ def store_global_u32( @dsl_user_op def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None: """Publish eight BF16 values through an NVLS multicast mapping.""" - words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] llvm.inline_asm( None, @@ -312,7 +307,6 @@ def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None: @dsl_user_op def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: """Reset one Lamport fragment to four FP32 negative-zero bit patterns.""" - address = pointer.toint(loc=loc, ip=ip) value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) llvm.inline_asm( @@ -333,7 +327,6 @@ def red_async_release_gpu_add_u32( pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None ) -> None: """The exact SM100 arrival primitive used by upstream LamportFlags.""" - address = pointer.toint(loc=loc, ip=ip) llvm.inline_asm( None, @@ -378,7 +371,6 @@ def map_shared_to_peer( ip=None, ) -> Int32: """Map a local shared-memory slot to the same slot in a peer CTA.""" - smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value() return Int32( llvm.inline_asm( @@ -422,7 +414,6 @@ def store_shared_cluster_f32( @dsl_user_op def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None): """Load two aligned FP32 DSM partials from local shared memory.""" - address = Int32(pointer.toint(loc=loc, ip=ip)) out = llvm.inline_asm( llvm.StructType.get_literal([T.f32()] * 2), @@ -443,7 +434,6 @@ def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None): @dsl_user_op def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None): """Load four aligned FP32 DSM partials from local shared memory.""" - address = Int32(pointer.toint(loc=loc, ip=ip)) out = llvm.inline_asm( llvm.StructType.get_literal([T.f32()] * 4), @@ -530,7 +520,6 @@ def sanitize_negative_zero_u32x2(packed): @cute.jit def sanitize_negative_zero(packed): """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" - result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) for i in cutlass.range_constexpr(4): word = packed[i] @@ -547,7 +536,6 @@ def sanitize_negative_zero(packed): @cute.jit def fragment_is_dirty(packed): """Bit-exact upstream sentinel check: one comparison per 32-bit word.""" - dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS) for i in cutlass.range_constexpr(1, 4): dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS)) @@ -564,7 +552,6 @@ def warp_sum_specialized( last_warp_mask: cutlass.Constexpr[int], ) -> Float32: """Warp sum supporting a compile-time partial final warp.""" - if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): for offset in cutlass.range_constexpr(16, 0, -1): # range_constexpr does not provide powers-of-two stepping. diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py index 416c316a70e7..3911d20d0358 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -539,9 +539,7 @@ def fused_kda_gate_chunk_cumsum( ) B, T, H, D = raw_g.shape if raw_beta.shape != (B, T, H): - raise ValueError( - f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}" - ) + raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}") if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) @@ -901,13 +899,12 @@ def fused_kda_gate( threshold: float = 20.0, lower_bound: float | None = None, ) -> torch.Tensor: - """ - Forward pass for KDA gate: - input g: [..., H*D] - param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter - threshold: softplus threshold parameter - return : [..., H, D] + """Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] """ orig_shape = g.shape[:-1] diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 087197d10a56..8b9aa86f9ed9 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -23,18 +23,20 @@ ################################################################################ -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BK': BK}, num_warps=num_warps) + triton.Config({"BK": BK}, num_warps=num_warps) for BK in [32, 64] for num_warps in [1, 2, 4] ], key=["H", "HV", "K", "BC"], ) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=["T"]) def chunk_kda_fwd_kernel_inter_solve_fused( q, k, @@ -57,8 +59,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( USE_SAFE_GATE: tl.constexpr, SOLVE_TRIL_DOT_PRECISION: tl.constexpr, ): - """ - Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + """Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. This kernel: @@ -74,8 +75,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( i_h = i_hv // (HV // H) if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -122,21 +129,31 @@ def chunk_kda_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_k0 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_g0 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_q1 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_k1 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_g1 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) # [BC, BK] b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) # [BK] - b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) # [BK, BC] @@ -146,15 +163,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) if i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_q2 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_k2 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_g2 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) # [BC, BK] b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) # [BK] - b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) # [BC, BK] b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) b_qg2 = b_q2 * b_gqn2 @@ -170,15 +195,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk21 += tl.dot(b_kg2, b_kgt) if i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_q3 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_k3 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_g3 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) # [BC, BK] b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) # [BK] - b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) # [BC, BK] b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) b_qg3 = b_q3 * b_gqn3 @@ -203,40 +236,78 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # save off-diagonal Aqk blocks and prepare Akk ################################################################################ if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk10 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) - p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) + p_b1 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) + ) b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) b_Akk10 = b_Akk10 * b_b1[:, None] if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk20 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Aqk21 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) - p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) + p_b2 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) + ) b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) + p_Aqk30 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Aqk31 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Aqk32 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b3 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) + ) b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) b_Akk30 = b_Akk30 * b_b3[:, None] b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk00 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) @@ -256,25 +327,25 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Ai33 = -tl.where(m_A, b_Ai33, 0) for i in range(2, min(BC, T - i_tc0)): - b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.0) b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(BC + 2, min(2*BC, T - i_tc0)): - b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + for i in range(BC + 2, min(2 * BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) - for i in range(2*BC + 2, min(3*BC, T - i_tc0)): - b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0) b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) - for i in range(3*BC + 2, min(4*BC, T - i_tc0)): - b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) + for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0) b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) b_Ai00 += m_I b_Ai11 += m_I @@ -289,53 +360,73 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Ai10 = -tl.dot( tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai00, - input_precision=SOLVE_TRIL_DOT_PRECISION + input_precision=SOLVE_TRIL_DOT_PRECISION, ) b_Ai21 = -tl.dot( tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai11, - input_precision=SOLVE_TRIL_DOT_PRECISION + input_precision=SOLVE_TRIL_DOT_PRECISION, ) b_Ai32 = -tl.dot( tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai22, - input_precision=SOLVE_TRIL_DOT_PRECISION + input_precision=SOLVE_TRIL_DOT_PRECISION, ) b_Ai20 = -tl.dot( b_Ai22, - tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, ) b_Ai31 = -tl.dot( b_Ai33, - tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, ) b_Ai30 = -tl.dot( b_Ai33, - tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, ) ################################################################################ # store full Akk_inv to Akk ################################################################################ - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + p_Akk00 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk10 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Akk20 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk21 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk30 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Akk31 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Akk32 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) @@ -349,18 +440,20 @@ def chunk_kda_fwd_kernel_inter_solve_fused( tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=['BK', 'NC', 'BT', 'HV'], + key=["BK", "NC", "BT", "HV"], ) -@triton.jit(do_not_specialize=['B', 'T']) +@triton.jit(do_not_specialize=["B", "T"]) def chunk_kda_fwd_kernel_intra_sub_chunk( q, k, @@ -386,8 +479,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( i_h = i_hv // (HV // H) if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -406,9 +505,9 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * HV + i_hv) * BT Akk = Akk + (bos * HV + i_hv) * BC - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) @@ -418,10 +517,12 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather( + b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0 + ) else: # caculate offset - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK) b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) b_gn = b_gn[None, :] @@ -429,8 +530,8 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( # less than 85 to avoid overflow in exp2 b_gm = (b_g - b_gn).to(tl.float32) - b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) - b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0) b_kgt = tl.trans(b_k * b_gk) @@ -445,8 +546,10 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) + ) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) @@ -458,8 +561,8 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Ai = -b_Akk for i in range(2, min(BC, T - i_ti)): - b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) - b_a = tl.where(o_i < i, b_a, 0.) + b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.0) b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I @@ -531,8 +634,7 @@ def chunk_kda_fwd_intra( # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) solve_tril_dot_precision = ( "tf32" - if current_platform.is_cuda() - and current_platform.has_device_capability(80) + if current_platform.is_cuda() and current_platform.has_device_capability(80) else "ieee" ) grid = (NT, B * HV) diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py index ecd00a51f9dd..86654f264431 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -16,18 +16,20 @@ from vllm.triton_utils import tl, triton -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) @triton.autotune( configs=[ - triton.Config({'BH': BH}, num_warps=num_warps) + triton.Config({"BH": BH}, num_warps=num_warps) for BH in [1, 2, 4, 8] for num_warps in [1, 2, 4, 8] ], key=["K", "H", "HV"], ) -@triton.jit(do_not_specialize=['T', 'N']) +@triton.jit(do_not_specialize=["T", "N"]) def chunk_kda_fwd_kernel_intra_token_parallel( q, k, @@ -65,7 +67,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel( left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos i_t = i_tg - bos else: @@ -82,11 +87,11 @@ def chunk_kda_fwd_kernel_intra_token_parallel( G: tl.constexpr = HV // H - q += bos * H*K - k += bos * H*K - g += bos * HV*K - Aqk += bos * HV*BT - Akk += bos * HV*BC + q += bos * H * K + k += bos * H * K + g += bos * HV * K + Aqk += bos * HV * BT + Akk += bos * HV * BC beta += bos * HV BK: tl.constexpr = triton.next_power_of_2(K) @@ -103,7 +108,9 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_g = tl.make_block_ptr( + g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) @@ -111,15 +118,25 @@ def chunk_kda_fwd_kernel_intra_token_parallel( for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_gj = tl.make_block_ptr( + g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) - tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv) - tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) + tl.store( + Aqk + i_t * HV * BT + o_hv * BT + j % BT, + b_Aqk.to(Aqk.dtype.element_ty), + mask=m_hv, + ) + tl.store( + Akk + i_t * HV * BC + o_hv * BC + j - i_ts, + b_Akk.to(Akk.dtype.element_ty), + mask=m_hv, + ) def chunk_kda_fwd_intra_token_parallel( @@ -134,8 +151,7 @@ def chunk_kda_fwd_intra_token_parallel( chunk_size: int = 64, sub_chunk_size: int = 16, ) -> None: - """ - Token-parallel implementation: each token gets its own thread block. + """Token-parallel implementation: each token gets its own thread block. Supports both fixed-length and variable-length sequences. Reduces wasted computation on padding. @@ -147,17 +163,21 @@ def chunk_kda_fwd_intra_token_parallel( gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) beta: [B, T, HV] Aqk: [B, T, HV, BT] output tensor to write to + cu_seqlens: cumulative sequence lengths for variable-length input Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) scale: attention scale chunk_size: BT (default 64) sub_chunk_size: BC (default 16) + """ B, T, H, K, HV = *q.shape, gk.shape[2] N = len(cu_seqlens) - 1 if cu_seqlens is not None else B BT = chunk_size BC = sub_chunk_size - def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) + def grid(meta): + return (B * T, triton.cdiv(HV, meta["BH"])) + chunk_kda_fwd_kernel_intra_token_parallel[grid]( q=q, k=k, diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py index 7f5f98400f2d..e65c95670d62 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py @@ -483,9 +483,7 @@ def fused_recurrent_kda( ) -@triton.jit( - do_not_specialize=["stride_beta_token", "stride_state_indices"] -) +@triton.jit(do_not_specialize=["stride_beta_token", "stride_state_indices"]) def fused_recurrent_kda_packed_decode_kernel( mixed_qkv, raw_g, diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 50cdb5518381..1a4ffdbd19ee 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -334,7 +334,8 @@ def _aiter_moe_fused_shared_experts_enabled( class MiniMaxM3MoE(nn.Module): """Sigmoid-routed MoE block with a routing-bias correction and a shared - expert.""" + expert. + """ def __init__( self, diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index ac02bca1fcc8..24a9bd6c09dc 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -173,7 +173,8 @@ class MiniMaxM3IndexerPrefillMetadata: @dataclass class MiniMaxM3IndexerDecodeMetadata: """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform - per-request query length (1, or 1 + num_speculative_tokens).""" + per-request query length (1, or 1 + num_speculative_tokens). + """ seq_lens: torch.Tensor # [num_decodes] int32 block_table: torch.Tensor @@ -206,7 +207,8 @@ class MiniMaxM3IndexerMetadataBuilder( AttentionMetadataBuilder[MiniMaxM3IndexerMetadata] ): """Abstract base: shared setup only. The Triton and MSA builders are - parallel subclasses that each own their full ``build`` (no shared code).""" + parallel subclasses that each own their full ``build`` (no shared code). + """ # Full cudagraphs for uniform decode batches (incl. spec-decode verify). _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index f385f82067fd..5f24645f8827 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -203,7 +203,8 @@ class MiniMaxM3SparsePrefillMetadata: @dataclass class MiniMaxM3SparseDecodeMetadata: """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform - per-request query length (1, or 1 + num_speculative_tokens).""" + per-request query length (1, or 1 + num_speculative_tokens). + """ seq_lens: torch.Tensor # [num_decodes] int32 block_table: torch.Tensor diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py index 8a015d4e6eae..0b3a336cd5c0 100644 --- a/vllm/models/minimax_m3/nvidia/indexer_msa.py +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -85,7 +85,8 @@ class MiniMaxM3IndexerMSAPrefillMetadata: class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): """Decode reuses the inherited base ``decode`` field (the Triton decode metadata); ``prefill_msa`` carries the fmha score plan for the prefill side - (the base ``prefill`` field is unused on this path).""" + (the base ``prefill`` field is unused on this path). + """ prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None # Per-forward view (``[:num_tokens]``) of the builder's persistent unified @@ -111,7 +112,8 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill - fmha plan is built eagerly (prefill batches are not captured).""" + fmha plan is built eagerly (prefill batches are not captured). + """ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 970b06d407ea..734f675fbe62 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -190,7 +190,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class MiniMaxM3MoE(nn.Module): """Sigmoid-routed MoE block with a routing-bias correction and a shared - expert.""" + expert. + """ def __init__( self, @@ -782,7 +783,8 @@ def forward( @property def ffn_all_reduce_deferred(self) -> bool: """This layer's FFN output is left un-reduced; the caller fuses the - all-reduce into the next RMSNorm.""" + all-reduce into the next RMSNorm. + """ if self.is_moe_layer: return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce return not self.mlp.down_proj.reduce_results diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 6b7d8f9f9380..2b8c7b144420 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -33,7 +33,6 @@ def apply_qsa_rope( tensor: torch.Tensor, ) -> torch.Tensor: """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" - num_tokens, _, head_dim = tensor.shape rotary_dim = rotary_emb.rotary_dim cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 @@ -67,7 +66,6 @@ def apply_qsa_rmsnorm( tensor: torch.Tensor, ) -> torch.Tensor: """Use vLLM's portable RMSNorm implementation on ROCm.""" - return cast(torch.Tensor, norm(tensor)) @@ -153,7 +151,6 @@ def project_qk( positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Project replicated Q/K, normalize+rotate Q, and preserve raw K.""" - qk, _ = self.index_qk_proj(hidden_states) q_raw, token_k = qk.split( ( @@ -176,7 +173,6 @@ def normalize_compressed_keys( first_rope_positions: torch.Tensor, ) -> torch.Tensor: """Normalize pooled K and apply the first token's exact group position.""" - keys = compressed_keys.reshape(-1, self.index_head_dim) keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape( -1, 1, self.index_head_dim @@ -284,7 +280,6 @@ def forward( out: torch.Tensor | None = None, ) -> torch.Tensor: """Return fixed-width request-relative token indices padded with ``-1``.""" - metadata = self._metadata() if metadata is None: # Preserve step-0 indices when later MTP steps reuse the buffer. diff --git a/vllm/models/qwen4_exp/amd/low_latency_gemm.py b/vllm/models/qwen4_exp/amd/low_latency_gemm.py index e8c5c2f917ec..c6e54885e2f6 100644 --- a/vllm/models/qwen4_exp/amd/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/amd/low_latency_gemm.py @@ -11,5 +11,4 @@ def enable_qwen4_exp_low_latency_gemm( dtype: torch.dtype, ) -> None: """Keep the standard vLLM linear methods on AMD ROCm.""" - del module, dtype diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index b62179055e83..b3209678922d 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -87,7 +87,6 @@ def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4.""" - if quant_config is not None and quant_config.get_name() == "modelopt_fp4": return None return quant_config @@ -103,7 +102,6 @@ def _remap_qsa_cache_scale_name( that cache directly, so only QSA layers need the final path component moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers. """ - scale_suffixes = { "k_proj.k_scale": "_k_scale", "k_proj.output_scale": "_k_scale", diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index 40ee6a6d4797..a2b420603f06 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -91,7 +91,6 @@ def _remap_quantized_layers( def _remap_mtp_weight_name(name: str) -> str | None: """Map Qwen4Exp checkpoint paths into the standalone draft model.""" - for checkpoint_prefix in ( "model.language_model.", "language_model.", @@ -260,7 +259,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def _iter_qsa_attentions(self): """Yield MTP attention modules that own a QSA indexer.""" - for layer in self.layers: attention = getattr(layer, "self_attn", None) if ( @@ -271,13 +269,11 @@ def _iter_qsa_attentions(self): def set_skip_topk(self, skip: bool) -> None: """Select on MTP step 0 and reuse its QSA indices on later steps.""" - for attention in self._iter_qsa_attentions(): attention.indexer.skip_topk = skip def compact_topk_indices(self, row_indices: torch.Tensor) -> None: """Keep each request's target-aligned step-0 sparse-index row.""" - num_rows = row_indices.numel() for attention in self._iter_qsa_attentions(): buffer = attention.topk_indices_buffer diff --git a/vllm/models/qwen4_exp/amd/ops/qsa.py b/vllm/models/qwen4_exp/amd/ops/qsa.py index 2e45db711d85..6bdb1fe592df 100644 --- a/vllm/models/qwen4_exp/amd/ops/qsa.py +++ b/vllm/models/qwen4_exp/amd/ops/qsa.py @@ -600,7 +600,6 @@ def qsa_mqa_paged( score_scale: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute QSA scores directly from a paged compressed-key cache.""" - _validate_mqa(q) if not q.is_cuda or not HAS_TRITON: raise RuntimeError("paged QSA scoring requires a GPU and Triton") @@ -678,7 +677,6 @@ def expand_qsa_block_indices_cuda( out: torch.Tensor | None = None, ) -> torch.Tensor: """Expand compressed blocks and compact the causal tail of the open group.""" - if not block_indices.is_cuda or not HAS_TRITON: raise RuntimeError("QSA index expansion requires a GPU and Triton") if token_topk % compress_ratio: @@ -738,7 +736,6 @@ def qsa_select_paged_tokens( out: torch.Tensor | None = None, ) -> torch.Tensor: """Score, select, and expand QSA indices without host synchronization.""" - rows = q.shape[0] output_width = token_topk + compress_ratio - 1 if out is None: @@ -829,7 +826,6 @@ def qsa_sparse_paged_attention( out: torch.Tensor | None = None, ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 K/V caches.""" - if not q.is_cuda or not HAS_TRITON: raise RuntimeError("paged QSA sparse attention requires a GPU and Triton") if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: @@ -971,7 +967,6 @@ def qsa_store_cache_rows( rows: torch.Tensor, ) -> None: """Store fixed-width rows in a QSA cache without boolean indexing.""" - if not cache.is_cuda or not HAS_TRITON: raise RuntimeError("QSA cache stores require a GPU and Triton") if cache.ndim != 4 or cache.shape[2] != 1: @@ -1017,7 +1012,6 @@ def qsa_compress_groups_with_ratio( rope_cache: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Pool completed groups from the compressor-state ring and raw token rows.""" - if not raw_keys.is_cuda or not HAS_TRITON: raise RuntimeError("QSA compression requires a GPU and Triton") rows = token_to_req.numel() diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index fc494be055cf..2616a3e9e02d 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -357,7 +357,6 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" - persistent_buffers = { "layer_multipliers": self.layer_multipliers, "ngram_heads_offsets": self.ngram_heads_offsets, diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index c0c766aad141..bed48a4cba39 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -430,7 +430,6 @@ def qwen4_exp_qsa_with_output( layer_name: LayerNameType, ) -> None: """Run the complete QSA state/update/attend transaction.""" - layer_name = _resolve_layer_name(layer_name) layer = get_forward_context().no_compile_layers[layer_name] if not isinstance(layer, Qwen4ExpQSAAttention): diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py index bd132685b528..72fc494b6037 100644 --- a/vllm/models/qwen4_exp/common/ple.py +++ b/vllm/models/qwen4_exp/common/ple.py @@ -28,7 +28,6 @@ def compute_ple_shard_overlap( tp_end: int, ) -> PLEShardOverlap | None: """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" - if checkpoint_start < 0 or checkpoint_rows < 0: raise ValueError("checkpoint shard bounds must be non-negative") if tp_start < 0 or tp_end < tp_start: @@ -54,7 +53,6 @@ def copy_ple_embedding_shard_( tp_end: int, ) -> int: """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" - if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: raise ValueError("destination and loaded weight must have matching ranks") if destination.shape[1:] != loaded_weight.shape[1:]: diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index e806dc0375de..eb0758025fb3 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -47,7 +47,6 @@ def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" - if positions.ndim == 1: positions = positions.unsqueeze(0).expand(3, -1) elif positions.ndim != 2 or positions.shape[0] not in (1, 3): @@ -115,7 +114,6 @@ def circular_qsa_slot_mapping( out: torch.Tensor | None = None, ) -> torch.Tensor: """Map each request to its fixed physical block as a circular token ring.""" - if compressor_state_size <= 0: raise ValueError("QSA circular buffer size must be positive") if block_table.ndim != 2: @@ -167,7 +165,6 @@ def compressed_qsa_slot_mapping( out: torch.Tensor | None = None, ) -> torch.Tensor: """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache.""" - if storage_block_size <= 0 or compress_ratio <= 0: raise ValueError("QSA block size and compression ratio must be positive") compressed_positions = torch.div( diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index bbdc5e4feba9..f4b968f15201 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -32,7 +32,6 @@ def apply_qsa_rope( tensor: torch.Tensor, ) -> torch.Tensor: """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" - num_tokens, _, head_dim = tensor.shape rotary_dim = rotary_emb.rotary_dim cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 @@ -245,7 +244,6 @@ def forward( valid-entry count (the attention kernel's loop bound, never a token index). """ - metadata = self._metadata() if metadata is None: # Preserve step-0 indices when later MTP steps reuse the buffer. diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 386987e18fc8..2e3833965574 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -86,7 +86,6 @@ def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4.""" - if quant_config is not None and quant_config.get_name() == "modelopt_fp4": return None return quant_config @@ -102,7 +101,6 @@ def _remap_qsa_cache_scale_name( that cache directly, so only QSA layers need the final path component moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers. """ - scale_suffixes = { "k_proj.k_scale": "_k_scale", "k_proj.output_scale": "_k_scale", diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index d12b851d688e..af3bcd4659f8 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -90,7 +90,6 @@ def _remap_quantized_layers( def _remap_mtp_weight_name(name: str) -> str | None: """Map Qwen4Exp checkpoint paths into the standalone draft model.""" - for checkpoint_prefix in ( "model.language_model.", "language_model.", @@ -250,7 +249,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def _iter_qsa_attentions(self): """Yield MTP attention modules that own a QSA indexer.""" - for layer in self.layers: attention = getattr(layer, "self_attn", None) if ( @@ -261,13 +259,11 @@ def _iter_qsa_attentions(self): def set_skip_topk(self, skip: bool) -> None: """Select on MTP step 0 and reuse its QSA indices on later steps.""" - for attention in self._iter_qsa_attentions(): attention.indexer.skip_topk = skip def compact_topk_indices(self, row_indices: torch.Tensor) -> None: """Keep each request's target-aligned step-0 sparse-index row.""" - num_rows = row_indices.numel() for attention in self._iter_qsa_attentions(): buffer = attention.topk_indices_buffer diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 84c4e7c5e1bb..c26fd87cff4f 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -869,7 +869,6 @@ def start_prefetch( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" - persistent_buffers = { "layer_multipliers": self.layer_multipliers, "ngram_heads_offsets": self.ngram_heads_offsets, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 6d0df0e741b7..f0aa62a48460 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -727,7 +727,6 @@ def warmup_qsa_sparse_paged_attention( selection_width: int, ) -> tuple[tuple[int, int, int], ...]: """Compile every production-reachable split-K/merge specialization.""" - head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. @@ -870,7 +869,6 @@ def qsa_store_cache_rows( rows: torch.Tensor, ) -> None: """Store fixed-width rows in a QSA cache without boolean indexing.""" - if not cache.is_cuda or not HAS_TRITON: raise RuntimeError("QSA CUDA cache stores require Triton") if cache.ndim != 4 or cache.shape[2] != 1: @@ -916,7 +914,6 @@ def qsa_compress_groups_with_ratio( rope_cache: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Pool completed groups from the compressor-state ring and raw token rows.""" - if not raw_keys.is_cuda or not HAS_TRITON: raise RuntimeError("QSA CUDA compression requires Triton") rows = token_to_req.numel() diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 9ddd65d15935..077dea91f028 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -317,7 +317,6 @@ def warmup_qsa_mqa_paged_decode( max_num_batched_tokens: int, ) -> tuple[tuple[int, int], ...]: """Compile every reachable decode specialization without launching it.""" - page_size = k_cache.shape[1] page_table_width = page_table.shape[1] columns = page_table_width * page_size @@ -444,7 +443,6 @@ def expand_qsa_block_indices( out: torch.Tensor, ) -> None: """Expand compressed blocks and compact the causal tail of the open group.""" - assert token_topk % compress_ratio == 0 block_topk = token_topk // compress_ratio output_width = token_topk + compress_ratio - 1 @@ -524,8 +522,8 @@ def qsa_select_paged_decode( compress_ratio: Number of logical tokens represented by a cache row. decode_query_len: Number of query tokens per request. block_indices: Compressed-index output buffer. - """ + """ assert token_topk % compress_ratio == 0 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio) assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0 @@ -599,8 +597,8 @@ def qsa_select_paged_prefill( max_query_len: Maximum number of query tokens in one request. block_indices: Compressed-index output buffer. max_seq_len: Longest context length in the batch this step. - """ + """ assert token_topk % compress_ratio == 0 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio) assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match" diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 2ec8790e8d3b..b04cdaad2128 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -146,7 +146,6 @@ def _dequantize_embeddings( output_dtype: torch.dtype, ) -> torch.Tensor: """Dequantize PLE lookup output.""" - return self.ple_embedding.ngram_embedding.dequantize( embeddings, output_dtype, diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py index 59f020230372..be026810bf77 100644 --- a/vllm/multimodal/audio.py +++ b/vllm/multimodal/audio.py @@ -45,6 +45,7 @@ def get_audio_duration(*, y: npt.NDArray[np.floating], sr: float = 22050) -> flo Returns: Duration of the audio in seconds. + """ n_samples = y.shape[-1] return float(n_samples) / sr @@ -71,6 +72,7 @@ class AudioSpec: (no normalization). 1 = mono, 2 = stereo, etc. channel_reduction: Method to reduce channels when input has more channels than target. Only used when reducing channels. + """ target_channels: int | None = 1 @@ -119,6 +121,7 @@ def normalize_audio( Raises: ValueError: If audio has unsupported dimensions or channel expansion is requested (e.g., mono to stereo). + """ if not spec.needs_normalization: return audio @@ -195,6 +198,7 @@ def resample_audio_pyav( Returns: Resampled audio with the same shape as the input (1D → 1D, 2D → 2D). + """ orig_sr_int = int(round(orig_sr)) target_sr_int = int(round(target_sr)) @@ -312,6 +316,7 @@ def resample_audio_torchaudio( Returns: Resampled audio with the same shape as the input (1D → 1D, 2D → 2D). + """ orig_sr_int = int(round(orig_sr)) target_sr_int = int(round(target_sr)) @@ -431,6 +436,7 @@ def split_audio( ... ) >>> len(chunks) 3 + """ if audio_data.ndim > 1: raise ValueError( @@ -500,6 +506,7 @@ def find_split_point( ... ) >>> 16000 <= split_idx <= 17600 True + """ segment = wav[start_idx:end_idx] diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache.py index cc48b58298bb..fa6dabce4128 100644 --- a/vllm/multimodal/cache.py +++ b/vllm/multimodal/cache.py @@ -63,12 +63,12 @@ def __init__(self, mm_hashes: list[str]) -> None: class MultiModalProcessorCacheItem: - """ - The data to store inside `MultiModalProcessorOnlyCache`. + """The data to store inside `MultiModalProcessorOnlyCache`. Args: item: The processed tensor data corresponding to a multi-modal item. prompt_updates: The prompt updates corresponding to `item`. + """ def __init__( @@ -83,8 +83,7 @@ def __init__( class MultiModalProcessorCacheItemMetadata: - """ - The metadata to store inside `MultiModalProcessorSenderCache`. + """The metadata to store inside `MultiModalProcessorSenderCache`. Args: item: The processed tensor data corresponding to a multi-modal item. @@ -94,6 +93,7 @@ class MultiModalProcessorCacheItemMetadata: prompt_updates: The prompt updates corresponding to `item`. This needs to stay on P0 because for some models, they are dependent on the processed tensor data (cached on P1). + """ def __init__( @@ -163,8 +163,7 @@ def get_item_size( @classmethod def get_item_complexity(cls, value: MultiModalCacheValue) -> int: - """ - Get the number of leaf elements in a multi-modal cache value. + """Get the number of leaf elements in a multi-modal cache value. This provides a measure of structural complexity that can be useful for debugging cache performance and understanding data patterns. @@ -174,6 +173,7 @@ def get_item_complexity(cls, value: MultiModalCacheValue) -> int: Returns: The number of leaf elements in the nested structure. + """ return json_count_leaves(value) @@ -196,8 +196,7 @@ def get_lru_cache( class BaseMultiModalCache(ABC, Generic[_I, _O]): - """ - Abstract base class to read/write multi-modal items from cache. + """Abstract base class to read/write multi-modal items from cache. The idea of multi-modal caching is based on having a client and server where the client executes in the frontend process (=P0) and @@ -226,8 +225,7 @@ def get_and_update_item( mm_item: _I, mm_hash: str, ) -> _O: - """ - Possibly update a multi-modal item based on whether it is + """Possibly update a multi-modal item based on whether it is in the underlying cache. This update is done out-of-place and updates the cache eviction order. @@ -238,6 +236,7 @@ def get_and_update_item( Returns: The update multi-modal item. + """ raise NotImplementedError @@ -246,8 +245,7 @@ def get_and_update( mm_items: Sequence[_I], mm_hashes: list[str], ) -> list[_O]: - """ - Possibly update a sequence of multi-modal items based on whether they + """Possibly update a sequence of multi-modal items based on whether they are in the underlying cache. This update is done out-of-place and updates the cache eviction order. @@ -258,6 +256,7 @@ def get_and_update( Returns: A new list of updated multi-modal items. + """ assert len(mm_items) == len(mm_hashes) @@ -289,6 +288,7 @@ def cache_if_fits( Returns: `True` if the item was cached, otherwise `False`. + """ if cache.put_if_fits(key, value): return True @@ -325,8 +325,7 @@ class BaseMultiModalProcessorCache( @abstractmethod def is_cached_item(self, mm_hash: str) -> bool: - """ - Check whether a multi-modal item is + """Check whether a multi-modal item is in the underlying cache. This **DOES NOT** update the cache eviction order. @@ -336,12 +335,12 @@ def is_cached_item(self, mm_hash: str) -> bool: Returns: `True` if the item is cached, otherwise `False`. + """ raise NotImplementedError def is_cached(self, mm_hashes: list[str]) -> list[bool]: - """ - Check whether a sequence of multi-modal items are + """Check whether a sequence of multi-modal items are in the underlying cache. This **DOES NOT** update the cache eviction order. @@ -351,6 +350,7 @@ def is_cached(self, mm_hashes: list[str]) -> list[bool]: Returns: For each item, `True` if the item is cached, otherwise `False`. + """ return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes] @@ -367,31 +367,30 @@ def close(self) -> None: @abstractmethod def touch_sender_cache_item(self, mm_hash: str) -> None: - """ - Update the cache eviction order for a multi-modal item. + """Update the cache eviction order for a multi-modal item. This is used to touch the item in the cache without changing its value. Args: mm_hash: The hash of the multi-modal item. + """ raise NotImplementedError @abstractmethod def make_stats(self, *, delta: bool = False) -> CacheInfo: - """ - Get (and reset) the multi-modal cache stats. + """Get (and reset) the multi-modal cache stats. Returns: The current multi-modal caching stats. + """ raise NotImplementedError class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache): - """ - The cache which is used on P0 when IPC caching is disabled. + """The cache which is used on P0 when IPC caching is disabled. How to update each item: @@ -442,8 +441,7 @@ def make_stats(self, *, delta: bool = False) -> CacheInfo: class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache): - """ - The cache which is used on P0 when IPC caching is enabled. + """The cache which is used on P0 when IPC caching is enabled. How to update each item: @@ -507,8 +505,7 @@ def invalidate(self, mm_hash: str) -> None: class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): - """ - The cache which is used on P0 when IPC caching is enabled. + """The cache which is used on P0 when IPC caching is enabled. How to update each item: @@ -609,7 +606,8 @@ def get_and_update_item( @override def touch_sender_cache_item(self, mm_hash: str) -> None: """Touch the item in shared memory cache to prevent eviction. - Increments writer_flag on sender side.""" + Increments writer_flag on sender side. + """ self._shm_cache.touch(mm_hash) @override @@ -662,8 +660,7 @@ def get_and_update_features( self, mm_features: list["MultiModalFeatureSpec"], ) -> list["MultiModalFeatureSpec"]: - """ - Update multimodal features with cached encoder outputs. + """Update multimodal features with cached encoder outputs. Touch all identifier at first before update to avoid item in updated list evict during update. @@ -695,8 +692,7 @@ def touch_receiver_cache_item( mm_hash: str, mm_item: MultiModalKwargsItem | None = None, ) -> None: - """ - Update the cache eviction order for a multi-modal item. + """Update the cache eviction order for a multi-modal item. This is used to touch the item in the cache without changing its value. @@ -705,13 +701,13 @@ def touch_receiver_cache_item( mm_hash: The hash of the multi-modal item. mm_item: The multi-modal item itself. This is optional and may not be needed by some cache implementations. + """ raise NotImplementedError class MultiModalReceiverCache(BaseMultiModalReceiverCache): - """ - The cache which is used on P1 when IPC caching is enabled. + """The cache which is used on P1 when IPC caching is enabled. How to update each item: @@ -762,8 +758,7 @@ def clear_cache(self) -> None: class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): - """ - The cache which is used on P1 Worker Process when IPC caching is enabled. + """The cache which is used on P1 Worker Process when IPC caching is enabled. How to update each item: @@ -827,7 +822,8 @@ def touch_receiver_cache_item( mm_item: MultiModalKwargsItem | None = None, ) -> None: """Touch the item in shared memory cache to prevent eviction. - Increments reader_count on receiver side.""" + Increments reader_count on receiver side. + """ assert mm_item is not None if "address" in mm_item: address = cast(int, mm_item["address"].data) diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index 3042091611f9..ef3da192db21 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/vllm/multimodal/encoder_budget.py @@ -19,8 +19,7 @@ def get_mm_max_toks_per_item( processor: BaseMultiModalProcessor, mm_counts: Mapping[str, int], ) -> Mapping[str, int]: - """ - Get the maximum number of tokens per data item from each modality based + """Get the maximum number of tokens per data item from each modality based on underlying model configuration. """ max_tokens_per_item = processor.info.get_mm_max_tokens_per_item( diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py index 58f807303700..7ccc1bdaa9d0 100644 --- a/vllm/multimodal/gpu_ipc_memory.py +++ b/vllm/multimodal/gpu_ipc_memory.py @@ -183,6 +183,7 @@ def reserve_mm_ipc_gpu_memory( Raises: ValueError: If the reservation leaves no memory for the KV cache. + """ if mm_config is None: return available_kv_cache_memory_bytes diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 3750ae162bb0..eb4ab4d28a40 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -23,8 +23,7 @@ def _get_hasher_factory( algorithm: MMHasherAlgorithm, ) -> Callable[[], "hashlib._Hash"]: - """ - Get the hasher factory based on the configured algorithm. + """Get the hasher factory based on the configured algorithm. Args: algorithm: Hash algorithm name (blake3, sha256, or sha512) @@ -33,8 +32,8 @@ def _get_hasher_factory( Supports blake3 (default), sha256, and sha512 for FIPS compliance. See: https://github.com/vllm-project/vllm/issues/18334 - """ + """ if algorithm == "blake3": from blake3 import blake3 diff --git a/vllm/multimodal/image.py b/vllm/multimodal/image.py index ff2da044a664..4b1d67f2fb86 100644 --- a/vllm/multimodal/image.py +++ b/vllm/multimodal/image.py @@ -38,7 +38,8 @@ def rgba_to_rgb( def _has_transparency(image: Image.Image) -> bool: """Detect whether an image carries transparency data (RGBA, LA, PA, - or tRNS chunk in P/L/RGB PNGs).""" + or tRNS chunk in P/L/RGB PNGs). + """ if image.mode in ("RGBA", "LA", "PA"): return True return "transparency" in getattr(image, "info", {}) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index bd1fc169d144..585acd410d7f 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -120,11 +120,9 @@ class VisionChunkVideo(TypedDict): @dataclass(frozen=True) class PlaceholderRange: - """ - Placeholder location information for multi-modal data. + """Placeholder location information for multi-modal data. Example: - Prompt: `AAAA BBBB What is in these images?` Images A and B will have: @@ -133,6 +131,7 @@ class PlaceholderRange: A: PlaceholderRange(offset=0, length=4) B: PlaceholderRange(offset=5, length=4) ``` + """ offset: int @@ -161,8 +160,7 @@ def get_num_embeds(self) -> int: def get_embeds_indices_in_range( self, start_idx: int, end_idx: int ) -> tuple[int, int]: - """ - Returns the starting and ending indices of the embeddings of encoder outputs + """Returns the starting and ending indices of the embeddings of encoder outputs in the range of [start_idx, end_idx) in the placeholders. For example, given: @@ -190,6 +188,7 @@ def extract_embeds_range(self) -> list[tuple[int, int]]: A tuple `(start, end)` representing the start and end indices (inclusive) of the embedded region. Returns full placeholder range if `is_embed` is `None`. + """ if self.is_embed is None: return [(self.offset, self.offset + self.length - 1)] @@ -234,8 +233,7 @@ def nested_tensors_equal( b: NestedTensors, check_dtype: bool = True, ) -> bool: - """ - Equality check between + """Equality check between [`NestedTensors`][vllm.multimodal.inputs.NestedTensors] objects. If `check_dtype` is `True`, the tensors must have the same dtype. @@ -297,7 +295,8 @@ def _is_dense(t: torch.Tensor) -> bool: """Whether `t`'s elements fill a contiguous storage range, possibly permuted (e.g. a transposed matrix). Such layouts can be copied with pitched cudaMemcpy2D/3DAsync directly from pinned memory, whereas - gapped layouts (e.g. strided slices) stage through a pageable temp.""" + gapped layouts (e.g. strided slices) stage through a pageable temp. + """ if not t.is_contiguous(): expected_stride = 1 for stride, size in sorted(zip(t.stride(), t.shape)): @@ -336,8 +335,7 @@ def _h2d(x: torch.Tensor) -> torch.Tensor: def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> bool: - """ - Equality check between + """Equality check between [`BatchedTensorInputs`][vllm.multimodal.inputs.BatchedTensorInputs] objects. """ return all(k in b and nested_tensors_equal(a[k], b[k]) for k in a) @@ -345,8 +343,7 @@ def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> boo @dataclass class MultiModalFeatureSpec: - """ - Represents a single multimodal input with its processed data and metadata. + """Represents a single multimodal input with its processed data and metadata. Used to track multimodal data through processing and caching. A request containing multiple multimodal items will have one @@ -392,8 +389,7 @@ def gather_kwargs(features: list["MultiModalFeatureSpec"], keys: set[str]): @dataclass class MultiModalFieldElem: - """ - Represents a processed keyword argument to pass to a model for a + """Represents a processed keyword argument to pass to a model for a [`MultiModalKwargsItem`][vllm.multimodal.inputs.MultiModalKwargsItem]. """ @@ -429,8 +425,7 @@ def __eq__(self, other: object) -> bool: @dataclass(frozen=True, kw_only=True) class BaseMultiModalField(ABC): - """ - Defines how to interpret tensor data belonging to a keyword argument for + """Defines how to interpret tensor data belonging to a keyword argument for [`MultiModalKwargsItems`][vllm.multimodal.inputs.MultiModalKwargsItems], and vice versa. """ @@ -458,8 +453,7 @@ def build_elems( key: str, data: NestedTensors, ) -> Sequence[MultiModalFieldElem]: - """ - Construct + """Construct [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem] instances to represent the provided data. @@ -481,8 +475,7 @@ def reduce_data( device: torch.types.Device = None, pin_memory: bool = False, ) -> NestedTensors: - """ - Merge the data from multiple instances of + """Merge the data from multiple instances of [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem]. This is the inverse of @@ -511,9 +504,8 @@ def reduce_data( @dataclass(frozen=True, kw_only=True) class MultiModalBatchedField(BaseMultiModalField): - """ - Info: - [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched] + """Info: + [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched] """ def build_elems( @@ -547,10 +539,9 @@ def _reduce_data( @dataclass(frozen=True, kw_only=True) class MultiModalFlatField(BaseMultiModalField): - """ - Info: - [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat] - [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes] + """Info: + [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat] + [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes] """ slices: Sequence[slice] | Sequence[Sequence[slice]] @@ -637,9 +628,8 @@ def _shape_before_after(tensor: torch.Tensor): @dataclass(frozen=True, kw_only=True) class MultiModalSharedField(BaseMultiModalField): - """ - Info: - [`MultiModalFieldConfig.shared`][vllm.multimodal.inputs.MultiModalFieldConfig.shared] + """Info: + [`MultiModalFieldConfig.shared`][vllm.multimodal.inputs.MultiModalFieldConfig.shared] """ batch_size: int @@ -666,8 +656,7 @@ def _reduce_data( class MultiModalFieldConfig: @staticmethod def batched(modality: str, *, keep_on_cpu: bool = False): - """ - Defines a field where an element in the batch is obtained by + """Defines a field where an element in the batch is obtained by indexing into the first dimension of the underlying data. Args: @@ -676,7 +665,6 @@ def batched(modality: str, *, keep_on_cpu: bool = False): keep_on_cpu: Whether to keep this field on the CPU for the model inputs. Example: - ``` Input: Data: [[AAAA] @@ -688,6 +676,7 @@ def batched(modality: str, *, keep_on_cpu: bool = False): Element 2: [BBBB] Element 3: [CCCC] ``` + """ return MultiModalFieldConfig( field=MultiModalBatchedField(keep_on_cpu=keep_on_cpu), @@ -702,8 +691,7 @@ def flat( *, keep_on_cpu: bool = False, ): - """ - Defines a field where an element in the batch is obtained by + """Defines a field where an element in the batch is obtained by slicing along the first dimension of the underlying data. Args: @@ -716,7 +704,6 @@ def flat( keep_on_cpu: Whether to keep this field on the CPU for the model inputs. Example: - ``` Given: slices: [slice(0, 3), slice(3, 7), slice(7, 9)] @@ -746,6 +733,7 @@ def flat( Element 2: [[B],[B],[B],[B]] Element 3: [[C],[C]] ``` + """ return MultiModalFieldConfig( field=MultiModalFlatField( @@ -764,8 +752,7 @@ def flat_from_sizes( *, keep_on_cpu: bool = False, ): - """ - Defines a field where an element in the batch is obtained by + """Defines a field where an element in the batch is obtained by slicing along the first dimension of the underlying data. Args: @@ -777,7 +764,6 @@ def flat_from_sizes( keep_on_cpu: Whether to keep this field on the CPU for the model inputs. Example: - ``` Given: size_per_item: [3, 4, 2] @@ -807,8 +793,8 @@ def flat_from_sizes( Info: [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat] - """ + """ if size_per_item.ndim != 1: raise ValueError( "size_per_item should be a 1-D tensor, " @@ -836,8 +822,7 @@ def shared( *, keep_on_cpu: bool = False, ): - """ - Defines a field where an element in the batch is obtained by + """Defines a field where an element in the batch is obtained by taking the entirety of the underlying data. This means that the data is the same for each element in the batch. @@ -849,7 +834,6 @@ def shared( keep_on_cpu: Whether to keep this field on the CPU for the model inputs. Example: - ``` Given: batch_size: 4 @@ -863,6 +847,7 @@ def shared( Element 3: [XYZ] Element 4: [XYZ] ``` + """ return MultiModalFieldConfig( field=MultiModalSharedField( @@ -884,8 +869,7 @@ def build_elems( class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]): - """ - A dictionary of processed keyword arguments to pass to the model, + """A dictionary of processed keyword arguments to pass to the model, corresponding to a single item in [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. """ @@ -912,8 +896,7 @@ def get_data(self) -> dict[str, NestedTensors]: class MultiModalKwargsItems(UserDict[str, Sequence[_I]]): - """ - A dictionary of processed multi-modal inputs by modality. + """A dictionary of processed multi-modal inputs by modality. For example, given a processor that processes images into `pixel_values` and `image_grid_thw`, diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index e1e75dd8d744..1f1cda8dccec 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -92,6 +92,10 @@ def load_audio_pyav( Args: path: A :class:`~io.BytesIO` buffer, a filesystem :class:`~pathlib.Path`, or a string path. + sr: Target sample rate, or None to keep the native rate. + mono: Whether to average the channels down to mono. + max_decode_bytes: If set, abort decoding once this many bytes have + been read from the source. max_duration_s: If set, abort decoding once the accumulated sample count exceeds this many seconds of audio. Prevents decompression-bomb attacks where a small compressed file @@ -100,6 +104,7 @@ def load_audio_pyav( Returns: ``(waveform, sample_rate)`` where *waveform* is a 1-D float32 NumPy array and *sample_rate* is the native sample rate in Hz. + """ try: container = av.open(path) @@ -209,7 +214,7 @@ def load_audio_soundfile( max_duration_s: float | None = None, max_decode_bytes: int | None = None, ) -> tuple[np.ndarray, int]: - """Load audio via soundfile""" + """Load audio via soundfile.""" with soundfile.SoundFile(path) as f: native_sr = f.samplerate if max_duration_s is not None: @@ -274,6 +279,7 @@ def load_audio_torchcodec( ``(waveform, sample_rate)`` where *waveform* is a float32 NumPy array (1-D when ``mono=True``) and *sample_rate* is the output sample rate in Hz. + """ if AudioDecoder is None: # Unify "torchcodec not installed" and "system ffmpeg missing" into @@ -412,10 +418,17 @@ def load_audio( """Load audio using the selected decoding backend. Args: + path: Audio file path or in-memory buffer. + sr: Target sample rate, or None to keep the native rate. + mono: Whether to downmix to a single channel. + max_duration_s: Reject audio longer than this many seconds. + max_decode_bytes: Reject audio that would decode to more than this + many bytes. backend: One of ``AUDIO_BACKENDS``. ``None`` (default) selects ``"auto"``, which tries soundfile, then torchcodec, then PyAV; the other values select a single backend with no fallback. + """ backend = backend or "auto" if backend not in AUDIO_BACKENDS: diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index 5a58b9ff72d8..d3e622156fc8 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -14,8 +14,7 @@ @dataclass class MediaWithBytes(Generic[_T]): - """ - Wrapper that couples a media object with its original encoded bytes. + """Wrapper that couples a media object with its original encoded bytes. This ensures the raw bytes and media object remain synchronized, preventing cache corruption from in-place modifications. @@ -89,8 +88,7 @@ def load_bytes(self, data: bytes) -> _T: @abstractmethod def load_base64(self, media_type: str, data: str) -> _T: - """ - List of media types: + """List of media types: https://www.iana.org/assignments/media-types/media-types.xhtml """ raise NotImplementedError diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 8980418b3275..297dc91af232 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -73,6 +73,7 @@ def _wrap_media_fetch_error( 408/429, invalid URL) Original exception for transient errors (5xx, 408, 429, network blips) or other exceptions + """ if isinstance(exc, VLLMValidationError): return exc @@ -162,16 +163,16 @@ def __init__( allowed_local_media_path: str = "", allowed_media_domains: list[str] | None = None, ) -> None: - """ - Args: - media_io_kwargs: Additional args passed to process media - inputs, keyed by modalities. For example, - to set num_frames for video, set - `--media-io-kwargs '{"video":{"num_frames":40}}'` - connection: HTTP connection client to download media contents. - allowed_local_media_path: A local directory to load media files from. - allowed_media_domains: If set, only media URLs that belong to this - domain can be used for multi-modal inputs. + """Args: + media_io_kwargs: Additional args passed to process media + inputs, keyed by modalities. For example, + to set num_frames for video, set + `--media-io-kwargs '{"video":{"num_frames":40}}'` + connection: HTTP connection client to download media contents. + allowed_local_media_path: A local directory to load media files from. + allowed_media_domains: If set, only media URLs that belong to this + domain can be used for multi-modal inputs. + """ super().__init__() @@ -468,9 +469,7 @@ def fetch_audio( self, audio_url: str, ) -> tuple[np.ndarray, int | float]: - """ - Load audio from a URL. - """ + """Load audio from a URL.""" audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) return self.load_from_url( @@ -483,9 +482,7 @@ async def fetch_audio_async( self, audio_url: str, ) -> tuple[np.ndarray, int | float]: - """ - Asynchronously fetch audio from a URL. - """ + """Asynchronously fetch audio from a URL.""" audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) return await self.load_from_url_async( @@ -500,8 +497,7 @@ def fetch_image( *, image_mode: str | None = "RGB", ) -> Image.Image: - """ - Load a PIL image from an HTTP or base64 data URL. + """Load a PIL image from an HTTP or base64 data URL. By default, the image is converted into RGB format. Set `media_io_kwargs={"image": {"image_mode": None}}` to keep the @@ -527,8 +523,7 @@ async def fetch_image_async( *, image_mode: str | None = "RGB", ) -> Image.Image: - """ - Asynchronously load a PIL image from an HTTP or base64 data URL. + """Asynchronously load a PIL image from an HTTP or base64 data URL. By default, the image is converted into RGB format. Set `media_io_kwargs={"image": {"image_mode": None}}` to keep the @@ -555,9 +550,7 @@ def fetch_video( image_mode: str | None = "RGB", video_processor: str | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: - """ - Load video from an HTTP or base64 data URL. - """ + """Load video from an HTTP or base64 data URL.""" image_io = ImageMediaIO( **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) ) @@ -581,8 +574,7 @@ async def fetch_video_async( image_mode: str | None = "RGB", video_processor: str | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: - """ - Asynchronously load video from an HTTP or base64 data URL. + """Asynchronously load video from an HTTP or base64 data URL. By default, the image is converted into RGB format. Set `media_io_kwargs={"image": {"image_mode": None}}` to keep the @@ -608,9 +600,7 @@ def fetch_image_embedding( self, data: str, ) -> torch.Tensor: - """ - Load image embedding from a URL. - """ + """Load image embedding from a URL.""" image_embedding_io = ImageEmbeddingMediaIO() return image_embedding_io.load_base64("", data) @@ -619,9 +609,7 @@ async def fetch_image_embedding_async( self, data: str, ) -> torch.Tensor: - """ - Asynchronously load image embedding from a URL. - """ + """Asynchronously load image embedding from a URL.""" image_embedding_io = ImageEmbeddingMediaIO() loop = asyncio.get_running_loop() @@ -633,9 +621,7 @@ def fetch_audio_embedding( self, data: str, ) -> torch.Tensor: - """ - Load audio embedding from a URL. - """ + """Load audio embedding from a URL.""" audio_embedding_io = AudioEmbeddingMediaIO() return audio_embedding_io.load_base64("", data) @@ -644,9 +630,7 @@ async def fetch_audio_embedding_async( self, data: str, ) -> torch.Tensor: - """ - Asynchronously load audio embedding from a URL. - """ + """Asynchronously load audio embedding from a URL.""" audio_embedding_io = AudioEmbeddingMediaIO() loop = asyncio.get_running_loop() @@ -658,9 +642,7 @@ def fetch_video_embedding( self, data: str, ) -> torch.Tensor: - """ - Load video embedding from a URL. - """ + """Load video embedding from a URL.""" video_embedding_io = VideoEmbeddingMediaIO() return video_embedding_io.load_base64("", data) @@ -669,9 +651,7 @@ async def fetch_video_embedding_async( self, data: str, ) -> torch.Tensor: - """ - Asynchronously load video embedding from a URL. - """ + """Asynchronously load video embedding from a URL.""" video_embedding_io = VideoEmbeddingMediaIO() loop = asyncio.get_running_loop() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index b2f8664246ae..8d7b5aa9b639 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -58,8 +58,7 @@ class ModalityDataItems(ABC, Generic[_T, _I]): - """ - Represents data items for a modality in + """Represents data items for a modality in [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. """ @@ -151,6 +150,7 @@ def validate_embedding_ndim( tensor: The tensor to validate. modality: The modality name for error messages (e.g., "image", "audio"). index: Optional index for list items, included in error messages. + """ if tensor.ndim < 2 or tensor.ndim > 3: idx_str = f" [{index}]" if index is not None else "" @@ -164,8 +164,7 @@ def validate_embedding_ndim( class EmbeddingItems( ModalityDataItems[torch.Tensor | list[torch.Tensor], torch.Tensor] ): - """ - Base class for data items that are expressed as a batched embedding tensor, + """Base class for data items that are expressed as a batched embedding tensor, or a list of embedding tensors (one per item). """ @@ -251,8 +250,7 @@ def get_feature_size(self, item_idx: int) -> int: class DictEmbeddingItems( ModalityDataItems[Mapping[str, torch.Tensor], Mapping[str, torch.Tensor]] ): - """ - Base class for data items that are expressed as a dictionary of tensors. + """Base class for data items that are expressed as a dictionary of tensors. Usually, the dictionary keys correspond to the outputs of HF processor. """ @@ -268,18 +266,18 @@ def __init__( ], optional_fields: Set[str] = frozenset(), ) -> None: - """ - Args: - data: The dictionary of tensors for this modality. - modality: The modality these items belong to. - required_fields: Fields `data` must contain. - fields_factory: Builds the field config from the data. - optional_fields: Fields `data` may omit. Which fields these are is - the caller's decision -- see - `MultiModalDataParser.embedding_field_sets`, where a deployment - that receives embeddings through an EC connector makes the - embeddings optional. They still need a field config, since they - are used whenever they *are* supplied. + """Args: + data: The dictionary of tensors for this modality. + modality: The modality these items belong to. + required_fields: Fields `data` must contain. + fields_factory: Builds the field config from the data. + optional_fields: Fields `data` may omit. Which fields these are is + the caller's decision -- see + `MultiModalDataParser.embedding_field_sets`, where a deployment + that receives embeddings through an EC connector makes the + embeddings optional. They still need a field config, since they + are used whenever they *are* supplied. + """ from transformers.feature_extraction_utils import BatchFeature @@ -471,14 +469,12 @@ def __init__(self, data: Sequence[Any]) -> None: class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]): - """ - A normalized [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] + """A normalized [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] such that each entry corresponds to a list. """ def select(self, modalities: Set[str]): - """ - Construct a new `MultiModalDataItems` instance containing only the + """Construct a new `MultiModalDataItems` instance containing only the selected modalities. """ return MultiModalDataItems( @@ -486,8 +482,7 @@ def select(self, modalities: Set[str]): ) def get_count(self, modality: str, *, strict: bool = True) -> int: - """ - Get the number of data items belonging to a modality. + """Get the number of data items belonging to a modality. If `strict=False`, return `0` instead of raising [`KeyError`][] even if the modality is not found. @@ -513,8 +508,7 @@ def get_items( modality: str, typ: type[_D] | tuple[type[_D], ...], ) -> _D: - """ - Get the data items belonging to a modality, + """Get the data items belonging to a modality, requiring that they belong to a certain type. """ if modality not in self: @@ -541,8 +535,7 @@ def get_items( class MultiModalDataParser: - """ - Parses [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] + """Parses [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] into [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. Args: @@ -561,6 +554,7 @@ class MultiModalDataParser: allow_missing_mm_embeddings (bool): Whether pre-computed embedding tensors may be absent from the request on a disaggregated consumer. Derived by `BaseProcessingInfo.allow_missing_mm_embeddings`. + """ embedding_fields: Mapping[str, Mapping[str, EmbeddingFieldRole]] = {} diff --git a/vllm/multimodal/processing/context.py b/vllm/multimodal/processing/context.py index 07149dbc4a26..c1b81a4d7210 100644 --- a/vllm/multimodal/processing/context.py +++ b/vllm/multimodal/processing/context.py @@ -70,6 +70,7 @@ def overlay_modality_mm_kwargs( Returns: A new dict with the modality-scoped keys overlaid when present. + """ merged = dict(kwargs) if modality is None: @@ -128,8 +129,7 @@ def get_stats_dict(self): @dataclass(frozen=True) class InputProcessingContext: - """ - Contains information about the model which may be used to + """Contains information about the model which may be used to modify the inputs. """ @@ -162,13 +162,13 @@ def get_hf_config( typ: type[Any] | tuple[type[Any], ...] | None = None, /, ) -> Any: - """ - Get the HuggingFace configuration + """Get the HuggingFace configuration (`transformers.PretrainedConfig`) of the model, additionally checking its type. Raises: TypeError: If the configuration is not of the specified type. + """ if typ is None: from transformers.configuration_utils import PretrainedConfig @@ -186,17 +186,15 @@ def get_hf_config( return hf_config def get_hf_image_processor_config(self) -> dict[str, Any]: - """ - Get the HuggingFace image processor configuration of the model. - """ + """Get the HuggingFace image processor configuration of the model.""" return self.model_config.hf_image_processor_config def get_mm_config(self): - """ - Get the multimodal config of the model. + """Get the multimodal config of the model. Raises: RuntimeError: If the model is not a multimodal model. + """ mm_config = self.model_config.multimodal_config if mm_config is None: @@ -221,13 +219,13 @@ def get_hf_processor( /, **kwargs: object, ) -> Any: - """ - Get the HuggingFace processor + """Get the HuggingFace processor (`transformers.ProcessorMixin`) of the model, additionally checking its type. Raises: TypeError: If the processor is not of the specified type. + """ if typ is None: from transformers.processing_utils import ProcessorMixin @@ -254,8 +252,7 @@ def init_processor( /, **kwargs: object, ) -> _T: - """ - Initialize a HuggingFace-like processor class, merging the + """Initialize a HuggingFace-like processor class, merging the keyword arguments with those in the model's configuration. """ merged_kwargs = self.get_merged_mm_kwargs(kwargs) @@ -317,8 +314,7 @@ def call_hf_processor( data: Mapping[str, object], kwargs: Mapping[str, object] = {}, ) -> BatchFeature: - """ - Call `hf_processor` on the prompt `data` + """Call `hf_processor` on the prompt `data` (text, image, audio...) with configurable options `kwargs`. """ assert callable(hf_processor) @@ -385,8 +381,7 @@ def get_hf_config(self) -> PretrainedConfig: return self.ctx.get_hf_config() def get_hf_processor(self, **kwargs: object) -> ProcessorMixin: - """ - Subclasses can override this method to handle + """Subclasses can override this method to handle specific kwargs from model config or user inputs. """ return self.ctx.get_hf_processor(**kwargs) @@ -409,8 +404,7 @@ def default_tok_params(self) -> TokenizeParams: return self.get_default_tok_params() def _get_expected_hidden_size(self) -> int | None: - """ - Get expected hidden size for embedding validation if `mm_embeds` are enabled. + """Get expected hidden size for embedding validation if `mm_embeds` are enabled. This validates hidden dimensions to prevent a vulnerability where embeddings with correct `ndim` but wrong `shape` could cause crashes at inference time. @@ -430,8 +424,7 @@ def allow_missing_mm_embeddings(self) -> bool: return mm_config is not None and mm_config.allow_missing_mm_embeddings def get_data_parser(self) -> MultiModalDataParser: - """ - Constructs a parser to preprocess multi-modal data items + """Constructs a parser to preprocess multi-modal data items before passing them to [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data]. @@ -453,8 +446,7 @@ def skip_prompt_length_check(self) -> bool: return False def get_supported_mm_limits(self) -> Mapping[str, int | None]: - """ - Return the maximum supported number of items for each modality. + """Return the maximum supported number of items for each modality. A value of `None` means unlimited number of items. @@ -486,8 +478,7 @@ def allowed_mm_limits(self) -> Mapping[str, int]: return allowed_limits def validate_num_items(self, modality: str, num_items: int) -> None: - """ - Raise `ValueError` if the number of input items for the given modality + """Raise `ValueError` if the number of input items for the given modality is invalid. """ supported_limit = self.supported_mm_limits.get(modality, 0) @@ -512,8 +503,7 @@ def parse_mm_data( *, validate: bool = True, ) -> MultiModalDataItems: - """ - Normalize [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] + """Normalize [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] to [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems] before passing them to [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data]. @@ -546,8 +536,7 @@ def get_mm_max_tokens_per_item( seq_len: int, mm_counts: Mapping[str, int], ) -> Mapping[str, int] | None: - """ - Return the maximum number of tokens per item of for each modality. + """Return the maximum number of tokens per item of for each modality. When `None` (the default) is returned, vLLM will generate dummy inputs (images/videos) at maximum possible sizes and process them to determine @@ -564,5 +553,6 @@ def get_mm_max_tokens_per_item( length and the maximum number of items of each modality allowed, and agree with dummy inputs (images/videos) at maximum possible sizes. + """ return None diff --git a/vllm/multimodal/processing/dummy_inputs.py b/vllm/multimodal/processing/dummy_inputs.py index 369d4f41a4c0..7ec90d030c20 100644 --- a/vllm/multimodal/processing/dummy_inputs.py +++ b/vllm/multimodal/processing/dummy_inputs.py @@ -26,8 +26,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]): - """ - Abstract base class that constructs the dummy data to profile + """Abstract base class that constructs the dummy data to profile multi-modal models. """ @@ -38,9 +37,7 @@ def __init__(self, info: _I) -> None: @abstractmethod def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - """ - Build the text input corresponding to `mm_counts`. - """ + """Build the text input corresponding to `mm_counts`.""" raise NotImplementedError @abstractmethod @@ -50,8 +47,7 @@ def get_dummy_mm_data( mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions], ) -> MultiModalDataDict: - """ - Build the multimodal input which, after processing, results in + """Build the multimodal input which, after processing, results in the maximum possible number of placeholder tokens. Args: @@ -61,6 +57,7 @@ def get_dummy_mm_data( If None, use model defaults for backward compatibility. If provided, models can use these to customize dummy data generation. + """ raise NotImplementedError @@ -70,14 +67,14 @@ def get_dummy_processor_inputs( mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions], ) -> ProcessorInputs: - """ - Build the input which, after processing, results in + """Build the input which, after processing, results in the maximum possible number of placeholder tokens. Args: seq_len: Sequence length mm_counts: Count of items per modality mm_options: Configurable options per modality (optional) + """ dummy_text = self.get_dummy_text(mm_counts) dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) diff --git a/vllm/multimodal/processing/inputs.py b/vllm/multimodal/processing/inputs.py index 56750c3087cb..639045a40852 100644 --- a/vllm/multimodal/processing/inputs.py +++ b/vllm/multimodal/processing/inputs.py @@ -18,8 +18,7 @@ @dataclass class ProcessorInputs: - """ - Represents the keyword arguments to + """Represents the keyword arguments to [`vllm.multimodal.processing.BaseMultiModalProcessor.apply`][]. """ diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 3db820458880..8e9983798f7f 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -104,8 +104,7 @@ class PromptIndex: class PromptIndexTargets: @staticmethod def start() -> PromptIndex: - """ - Resolves to the start of the prompt (before the first token). + """Resolves to the start of the prompt (before the first token). This results in a match even if the prompt is empty. """ @@ -116,9 +115,7 @@ def start() -> PromptIndex: @staticmethod def prefix(seq: list[int]) -> PromptIndex: - """ - Resolves to a location in the prompt after the given prefix. - """ + """Resolves to a location in the prompt after the given prefix.""" def get_match_index( prompt: list[int], @@ -146,8 +143,7 @@ def get_text_match_index( @staticmethod def end() -> PromptIndex: - """ - Resolves to the end of the prompt (after the last token). + """Resolves to the end of the prompt (after the last token). This results in a match even if the prompt is empty. """ @@ -247,8 +243,7 @@ class UpdateMode(str, Enum): @dataclass class PromptUpdate(ABC): - """ - Defines how to update a prompt with placeholder tokens. + """Defines how to update a prompt with placeholder tokens. Note: The target and content are token sequences. When converting text @@ -256,6 +251,7 @@ class PromptUpdate(ABC): `add_special_tokens=False` to match how the prompt itself is tokenized. Otherwise, the updated prompt may contain duplicated special tokens or fail to match the target. + """ modality: str @@ -294,8 +290,7 @@ def _resolve_content(self, item_idx: int) -> PromptUpdateDetails: return content def resolve(self, item_idx: int) -> "ResolvedPromptUpdate": - """ - Given the index of the processed item within + """Given the index of the processed item within [`modality`][vllm.multimodal.processing.PromptUpdate.modality], output a copy of this object with its lazy attributes resolved. """ @@ -310,11 +305,9 @@ def resolve(self, item_idx: int) -> "ResolvedPromptUpdate": @dataclass class PromptInsertion(PromptUpdate): - """ - Defines how to insert placeholder tokens into a prompt. + """Defines how to insert placeholder tokens into a prompt. Example: - For each image, insert a number of `` feature placeholders equal to the feature size of the vision encoder after the `` token: @@ -355,6 +348,7 @@ class PromptInsertion(PromptUpdate): insertion=[image_token_id] * image_feature_size, ) ``` + """ insertion: PromptUpdateContent = field(repr=False) @@ -379,11 +373,9 @@ def mode(self) -> UpdateMode: @dataclass class PromptReplacement(PromptUpdate): - """ - Defines how to replace portions of an input prompt with placeholder tokens. + """Defines how to replace portions of an input prompt with placeholder tokens. Example: - For each image, replace one `` input placeholder in the prompt with a number of `` feature placeholders equal to the feature size of the vision encoder: @@ -412,6 +404,7 @@ class PromptReplacement(PromptUpdate): ), ) ``` + """ replacement: PromptUpdateContent = field(repr=False) @@ -447,8 +440,7 @@ def modality(self) -> str: ... def full_groupby_modality(values: Iterable[_M]) -> ItemsView[str, list[_M]]: - """ - Convenience function to apply + """Convenience function to apply [`full_groupby`][vllm.utils.collection_utils.full_groupby] based on modality. """ @@ -462,8 +454,7 @@ class PromptTargetMatch(NamedTuple): @dataclass(frozen=True) class ResolvedPromptUpdate: - """ - A [`PromptUpdate`][vllm.multimodal.processing.PromptUpdate] with its + """A [`PromptUpdate`][vllm.multimodal.processing.PromptUpdate] with its lazy attributes resolved. """ @@ -522,8 +513,7 @@ def iter_token_matches( *, start_idx: int = 0, ) -> Generator[_TokenMatch]: - """ - Yield each occurrence of `match_ids` in `token_ids`. + """Yield each occurrence of `match_ids` in `token_ids`. Note that empty matches are ignored. """ @@ -562,8 +552,7 @@ def replace_token_matches( match_ids: list[int], new_ids: list[int], ) -> list[int]: - """ - Replace each occurrence of `match_ids` in `token_ids` + """Replace each occurrence of `match_ids` in `token_ids` with `new_ids`. Note that empty matches are ignored. @@ -846,8 +835,7 @@ def apply_token_matches( prompt: list[int], mm_prompt_updates: "MultiModalPromptUpdates", ) -> tuple[list[int], "MultiModalPromptUpdatesApplyResult"]: - """ - Apply the updates in `mm_prompt_updates` to `prompt`. + """Apply the updates in `mm_prompt_updates` to `prompt`. Matches are exclusive even when multiple modalities share the same placeholder tokens. In that case, the modality that @@ -914,8 +902,7 @@ def _iter_placeholders( prompt: list[int], mm_prompt_updates: "MultiModalPromptUpdates", ) -> Iterable[PlaceholderFeaturesInfo]: - """ - Yield each set of placeholder tokens found in `prompt`. + """Yield each set of placeholder tokens found in `prompt`. Matches are exclusive even when multiple modalities share the same placeholder tokens. In that case, the modality that @@ -1028,8 +1015,7 @@ class MultiModalProcessingInfo(NamedTuple): class BaseMultiModalProcessor(ABC, Generic[_I]): - """ - Abstract base class to process multi-modal inputs to be used in vLLM. + """Abstract base class to process multi-modal inputs to be used in vLLM. Not to be confused with `transformers.ProcessorMixin`. """ @@ -1097,8 +1083,7 @@ def _get_prompt_updates( hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: - """ - Given the original multi-modal items for this modality + """Given the original multi-modal items for this modality and HF-processed data, output the updates to perform. The information returned by this method is used to update token inputs @@ -1167,8 +1152,7 @@ def _get_hf_mm_data( return processor_data, passthrough_data def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str | None: - """ - Get the text to pass to the HF processor alongside the multi-modal + """Get the text to pass to the HF processor alongside the multi-modal data. By default, no text is passed. If the HF processor requires that @@ -1183,8 +1167,7 @@ def _preprocess_hf_mm_data( mm_data: Mapping[str, object], hf_processor_mm_kwargs: Mapping[str, object], ) -> tuple[Mapping[str, object], Mapping[str, object]]: - """ - Pre-process the multi-modal data and HF processor keyword arguments + """Pre-process the multi-modal data and HF processor keyword arguments before they are passed to the HF processor. By default, both are returned as-is. If the HF processor expects the @@ -1201,8 +1184,7 @@ def _postprocess_hf_mm_data( hf_processor_mm_kwargs: Mapping[str, object], processed_data: BatchFeature, ) -> BatchFeature: - """ - Post-process the output of the HF processor. + """Post-process the output of the HF processor. By default, the output is returned as-is. If you need to modify the output of the HF processor before it is converted into multi-modal @@ -1215,9 +1197,7 @@ def _apply_hf_processor_main( mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: - """ - Apply the HF processor on the multi-modal data. - """ + """Apply the HF processor on the multi-modal data.""" valid_mm_items = mm_items.select( {k for k, c in mm_items.get_all_counts().items() if c > 0} ) @@ -1250,8 +1230,7 @@ def _apply_hf_processor_main( ) def _postprocess_prompt(self, prompt: list[int]) -> list[int]: - """ - Post-process the prompt token IDs before locating or applying + """Post-process the prompt token IDs before locating or applying multi-modal prompt updates. By default, the prompt is returned as-is. If the HF processor (or @@ -1303,8 +1282,7 @@ def _recompute_cached_prompt_update( cached_update: ResolvedPromptUpdate, new_item_idx: int, ) -> ResolvedPromptUpdate: - """ - Override this if other attributes of `ResolvedPromptUpdate` + """Override this if other attributes of `ResolvedPromptUpdate` also need to be recomputed after retrieving from the cache. """ return replace(cached_update, item_idx=new_item_idx) @@ -1404,8 +1382,7 @@ def _cached_apply_hf_processor( inputs: ProcessorInputs, timing_ctx: TimingContext, ) -> MultiModalProcessingInfo: - """ - Apply the HF processor on the full prompt text, + """Apply the HF processor on the full prompt text, caching the results and reusing cached results. """ cache = self.cache @@ -1522,8 +1499,7 @@ def _apply_prompt_updates_via_text( *, encode_segments_separately: bool = False, ) -> tuple[list[int], MultiModalPromptUpdatesApplyResult]: - """ - Apply multi-modal prompt updates in text space, then encode the + """Apply multi-modal prompt updates in text space, then encode the result. A non-special-token target may be tokenized differently inside the @@ -1716,8 +1692,7 @@ def apply( inputs: ProcessorInputs, timing_ctx: TimingContext, ) -> MultiModalInput: - """ - Process multi-modal inputs to be used in vLLM. + """Process multi-modal inputs to be used in vLLM. The main steps are: @@ -1763,8 +1738,7 @@ def create_encoder_prompt( prompt: list[int], mm_items: MultiModalDataItems, ) -> list[int]: - """ - Create input prompt for the encoder. HF processor will be applied on + """Create input prompt for the encoder. HF processor will be applied on this prompt during profiling and generation. """ raise NotImplementedError @@ -1805,8 +1779,7 @@ def apply( inputs: ProcessorInputs, timing_ctx: TimingContext, ) -> MultiModalEncDecInput: - """ - Process multi-modal inputs to be used in vLLM. + """Process multi-modal inputs to be used in vLLM. The main processing steps are modified to fit encoder-decoder model: 1. Create encoder prompt from input prompt text. 2. Apply the HF processor on encoder prompt. diff --git a/vllm/multimodal/registry.py b/vllm/multimodal/registry.py index 1d541b975136..f5be2a0ebbf0 100644 --- a/vllm/multimodal/registry.py +++ b/vllm/multimodal/registry.py @@ -45,8 +45,7 @@ class ProcessingInfoFactory(Protocol[_I_co]): - """ - Constructs a + """Constructs a [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] instance from the context. """ @@ -58,8 +57,7 @@ def __call__( class DummyInputsBuilderFactory(Protocol[_I]): # type: ignore[misc] - """ - Constructs a + """Constructs a [`BaseDummyInputsBuilder`][vllm.multimodal.processing.BaseDummyInputsBuilder] instance from the context. """ @@ -68,8 +66,7 @@ def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ... class MultiModalProcessorFactory(Protocol[_I]): # type: ignore[misc] - """ - Constructs a + """Constructs a [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] instance from the context. """ @@ -101,13 +98,10 @@ def build_processor( class MultiModalRegistry: - """ - A registry that dispatches data processing according to the model. - """ + """A registry that dispatches data processing according to the model.""" def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool: - """ - Checks if the model supports multimodal inputs. + """Checks if the model supports multimodal inputs. Returns True if the model is multimodal with any non-zero supported modalities, otherwise returns False, effectively running in text-only mode. @@ -158,8 +152,7 @@ def register_processor( info: ProcessingInfoFactory[_I], dummy_inputs: DummyInputsBuilderFactory[_I], ): - """ - Register a multi-modal processor to a model class. The processor + """Register a multi-modal processor to a model class. The processor is constructed lazily, hence a factory method should be passed. When the model receives multi-modal data, the provided function is @@ -227,9 +220,7 @@ def create_processor( tokenizer: TokenizerLike | None = None, cache: BaseMultiModalProcessorCache | None = None, ) -> BaseMultiModalProcessor[BaseProcessingInfo]: - """ - Create a multi-modal processor for a specific model and tokenizer. - """ + """Create a multi-modal processor for a specific model and tokenizer.""" if not model_config.is_multimodal_model: model_name = model_config.served_model_name or model_config.model raise ValueError(f"{model_name} is not a multimodal model") @@ -250,8 +241,7 @@ def get_dummy_mm_inputs( processor: BaseMultiModalProcessor | None = None, scheduler_config: "SchedulerConfig | None" = None, ) -> MultiModalInput: - """ - Create dummy data for profiling the memory usage of a model. + """Create dummy data for profiling the memory usage of a model. The model is identified by `model_config`. """ diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 4de183e0dcec..c12acfa1e193 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -67,8 +67,7 @@ def encode_image_base64( image_mode: str | None = "RGB", format: str = "PNG", ) -> str: - """ - Encode a pillow image to base64 format. + """Encode a pillow image to base64 format. By default, the image is converted into RGB format before being encoded. Pass `image_mode=None` to keep the original image mode. @@ -83,8 +82,7 @@ def encode_image_url( image_mode: str | None = "RGB", format: str = "PNG", ) -> str: - """ - Encode a pillow image as a data URL. + """Encode a pillow image as a data URL. By default, the image is converted into RGB format before being encoded. Pass `image_mode=None` to keep the original image mode. @@ -147,14 +145,14 @@ def get_mm_features_in_window( def argsort_mm_positions( mm_positions: MultiModalPlaceholders, ) -> list[tuple[str, int]]: - """ - Given a `MultiModalPlaceholders`, output a sequence of keys to + """Given a `MultiModalPlaceholders`, output a sequence of keys to sort the dictionary by `offset` (starting index in the input sequence) in ascending order. Returns: A list of `(modality, idx)`, which can be used to access an item by `mm_positions[modality][idx]`. + """ flat_items = ( (modality, idx, item) @@ -222,7 +220,8 @@ def strip_covered_mm_data( so the workers never consume the payload fields. M-RoPE models keep CPU-side metadata fields used to compute positions. SHM address items are also kept so workers can balance the sender's reference count. The - scheduler-side ``Request`` keeps the full features.""" + scheduler-side ``Request`` keeps the full features. + """ if not mm_features or num_computed_tokens == 0: return mm_features @@ -253,8 +252,7 @@ def group_and_batch_mm_items( device: torch.types.Device = None, pin_memory: bool = False, ) -> Generator[tuple[int, BatchedTensorInputs]]: - """ - Group consecutive items (possibly from different requests) into batches. + """Group consecutive items (possibly from different requests) into batches. Items must be split across groups if any of the following occurs, as the batch would otherwise be invalid: @@ -270,6 +268,7 @@ def group_and_batch_mm_items( A tuple `(num_items, grouped_kwargs)`, where: - `kwargs` is a dictionary of keyword arguments to pass to the model; - `num_items` is the corresponding number of items. + """ start_idx = 0 for end_idx in range(1, len(items) + 1): @@ -296,8 +295,7 @@ def group_and_batch_mm_kwargs( device: torch.types.Device = None, pin_memory: bool = False, ) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]: - """ - Group consecutive items (possibly from different requests) into batches. + """Group consecutive items (possibly from different requests) into batches. Items must be split across groups if any of the following occurs, as the batch would otherwise be invalid: @@ -317,6 +315,7 @@ def group_and_batch_mm_kwargs( - `modality` is the modality of the batch; - `kwargs` is a dictionary of keyword arguments to pass to the model; - `num_items` is the corresponding number of items. + """ for modality, group in groupby(mm_kwargs, key=lambda x: x[0]): items_lst = [item for _, item in group] @@ -333,14 +332,14 @@ def fetch_audio( audio_url: str, audio_io_kwargs: dict[str, Any] | None = None, ) -> tuple[np.ndarray, int | float]: - """ - Args: + """Args: audio_url: URL of the audio file to fetch. audio_io_kwargs: Additional kwargs passed to handle audio IO. Warning: This method has direct access to local files and is only intended to be called by user code. Never call this from the online server! + """ media_io_kwargs = None if not audio_io_kwargs else {"audio": audio_io_kwargs} media_connector = MediaConnector( @@ -354,14 +353,14 @@ def fetch_image( image_url: str, image_io_kwargs: dict[str, Any] | None = None, ) -> Image.Image: - """ - Args: + """Args: image_url: URL of the image file to fetch. image_io_kwargs: Additional kwargs passed to handle image IO. Warning: This method has direct access to local files and is only intended to be called by user code. Never call this from the online server! + """ media_io_kwargs = None if not image_io_kwargs else {"image": image_io_kwargs} media_connector = MediaConnector( @@ -375,14 +374,14 @@ def fetch_video( video_url: str, video_io_kwargs: dict[str, Any] | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: - """ - Args: + """Args: video_url: URL of the video file to fetch. video_io_kwargs: Additional kwargs passed to handle video IO. Warning: This method has direct access to local files and is only intended to be called by user code. Never call this from the online server! + """ media_io_kwargs = None if not video_io_kwargs else {"video": video_io_kwargs} media_connector = MediaConnector( diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index c7391919794f..ff74183a55a7 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -288,6 +288,7 @@ def load_bytes( Returns: Tuple of ``(frames, metadata_dict)``, where ``frames`` is a CPU ``np.ndarray`` unless TorchCodec decodes on ``device="cuda"``. + """ target = VideoTargetMetadata( num_frames=num_frames, fps=fps, max_duration=max_duration @@ -915,8 +916,7 @@ def get_candidate_target_fps( sampling_fps: float, max_fps: float = 8.0, ) -> list[float]: - """ - Return the subset of `video_fps` factors that remain multiples + """Return the subset of `video_fps` factors that remain multiples of `sampling_fps`. Examples: @@ -931,6 +931,7 @@ def get_candidate_target_fps( ... ValueError: sampling_fps=2 must divide video_fps=5 to produce consistent frame steps. + """ if sampling_fps is None: raise ValueError("sampling_fps must be provided") @@ -967,8 +968,8 @@ def get_target_fps( frame_sample_mode: str, candidate_target_fps: list[float], ) -> float | None: - """ - Get the target fps that best spans the videoand has the most frames sampled + """Get the target fps that best spans the video and samples the most + frames. """ num_frames_sampled = 0 selected_target_fps = None @@ -1272,8 +1273,7 @@ def load_bytes( frame_recovery: bool = False, **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: - """ - Load video frames with dynamic sampling based on duration. + """Load video frames with dynamic sampling based on duration. Args: data: Raw video bytes @@ -1284,6 +1284,7 @@ def load_bytes( Returns: Tuple of (frames_array, metadata_dict) + """ # recompute source metadata with adjusted duration to ensure correct # sampling indices computation diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py index 2b86d17a28d9..3ca141c738c4 100644 --- a/vllm/multimodal/video_decoders/opencv.py +++ b/vllm/multimodal/video_decoders/opencv.py @@ -128,8 +128,7 @@ def _read_frames_with_recovery( frame_indices: list[int], total_frames: int, ) -> tuple[npt.NDArray, list[int], dict[int, int]]: - """ - Read frames with dynamic window forward-scan recovery. + """Read frames with dynamic window forward-scan recovery. When a target frame fails to load, the next successfully grabbed frame (before the next target frame) will be used to recover it. @@ -144,6 +143,7 @@ def _read_frames_with_recovery( - frames_array: Array of loaded frames - valid_frame_indices: List of frame indices that were loaded - recovered_map: Dict mapping recovered_idx -> source_idx + """ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) diff --git a/vllm/multimodal/video_prune/evs.py b/vllm/multimodal/video_prune/evs.py index 8584c08299fc..56049b6e9d20 100644 --- a/vllm/multimodal/video_prune/evs.py +++ b/vllm/multimodal/video_prune/evs.py @@ -16,8 +16,7 @@ def compute_retained_tokens_count( tokens_per_frame: int, num_frames: int, q: float ) -> int: - """ - Compute the number of retained tokens for a given video. + """Compute the number of retained tokens for a given video. Method ensures that we retain all the tokens from the first frame regardless of the pruning rate. @@ -28,6 +27,7 @@ def compute_retained_tokens_count( Returns: The number of retained tokens. + """ total_tokens = tokens_per_frame * num_frames evs_num_tokens = int(total_tokens * (1 - q)) @@ -41,8 +41,7 @@ def compute_retention_mask( spatial_merge_size: int, q: float, ) -> torch.Tensor: - """ - Computes the retention mask for input video embeddings. + """Computes the retention mask for input video embeddings. Args: video_embeds (`torch.Tensor`): The input video embeddings @@ -55,6 +54,7 @@ def compute_retention_mask( Returns: `torch.Tensor`: The retention mask for the video embeddings of `(T * H * W // spatial_merge_size ^ 2)` shape. + """ T, H, W = map(int, video_size_thw) @@ -98,8 +98,7 @@ def compute_mrope_for_media( tokens_per_second: float = 1.0, video_second_per_grid: float = 1.0, ) -> torch.Tensor: - """ - Computes the mrope for video embeddings based on the grid dimensions. + """Computes the mrope for video embeddings based on the grid dimensions. Computed mrope positions match original qwen 2.5 implementation, but positions are built for media being the first element in sequence. @@ -113,6 +112,7 @@ def compute_mrope_for_media( Tensor of shape `(T * H * W, 4)` where last dimension represents mrope positions [0:3), while the last channel contains value of llm_grid_w repeated for all positions. + """ llm_grid_t = video_size_thw[0] llm_grid_h = video_size_thw[1] // spatial_merge_size @@ -160,8 +160,7 @@ def recompute_mrope_positions( image_token_id: int, video_token_id: int, ) -> tuple[torch.LongTensor, int]: - """ - Update part of input mrope positions. + """Update part of input mrope positions. Original mrope_positions are computed incorrectly, so once we prune media tokens we should reflect this in the mrope positions for the LLM. @@ -204,8 +203,8 @@ def recompute_mrope_positions( Returns: Tuple of (mrope_positions, mrope_position_delta). - """ + """ # Tensors positions: torch.LongTensor = typing.cast( torch.LongTensor, mrope_positions.clone() diff --git a/vllm/multimodal/video_prune/vidcom2.py b/vllm/multimodal/video_prune/vidcom2.py index 9379f26ca74b..4f583473c9b2 100644 --- a/vllm/multimodal/video_prune/vidcom2.py +++ b/vllm/multimodal/video_prune/vidcom2.py @@ -48,6 +48,7 @@ def compute_retention_mask( Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained tokens. The True count equals `compute_retained_tokens_count` so placeholders sized at prompt-processing time match exactly. + """ T, H, W = map(int, video_size_thw) rows = H // spatial_merge_size diff --git a/vllm/outputs.py b/vllm/outputs.py index 84b9fccdc88d..9d46f327d787 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -53,6 +53,7 @@ class CompletionOutput: ``--per-request-spec-decode-metrics`` is enabled; None otherwise. Surfaced in the response as ``metrics.speculative_decoding`` for single-sequence (``n == 1``) requests. + """ index: int @@ -90,6 +91,7 @@ class PoolingOutput: Args: data: The extracted hidden states. + """ data: torch.Tensor @@ -128,6 +130,7 @@ class RequestOutput: prefix-cache writes for this request. kv_transfer_params: The params for remote K/V transfer. ec_transfer_params: The params for remote encoder-cache transfer. + """ def __init__( @@ -172,7 +175,6 @@ def __init__( def add(self, next_output: "RequestOutput", aggregate: bool) -> None: """Merge subsequent RequestOutput into this one""" - self.finished |= next_output.finished self.kv_transfer_params = next_output.kv_transfer_params self.ec_transfer_params = next_output.ec_transfer_params @@ -232,8 +234,7 @@ def __repr__(self) -> str: class PoolingRequestOutput(Generic[_O]): - """ - The output data of a pooling request to the LLM. + """The output data of a pooling request to the LLM. Args: request_id (str): A unique identifier for the pooling request. @@ -241,6 +242,7 @@ class PoolingRequestOutput(Generic[_O]): prompt_token_ids (list[int]): A list of token IDs used in the prompt. num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. + """ def __init__( @@ -274,6 +276,7 @@ class EmbeddingOutput: Args: embedding: The embedding vector, which is a list of floats. Its length depends on the hidden dimension of the model. + """ embedding: list[float] @@ -315,6 +318,7 @@ class ClassificationOutput: Args: probs: The probability vector, which is a list of floats. Its length depends on the number of classes. + """ probs: list[float] @@ -356,6 +360,7 @@ class ScoringOutput: Args: score: The similarity score, which is a scalar value. + """ score: float diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f099d7752ae5..d51c4455bea2 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -84,8 +84,7 @@ def commit( class Parser: - """ - Abstract Parser class that unifies ReasoningParser and ToolParser into + """Abstract Parser class that unifies ReasoningParser and ToolParser into a single interface for parsing model output. This class provides a unified way to handle both reasoning extraction @@ -191,8 +190,7 @@ def _initialize_history_tool_call_cnt( @abstractmethod def is_reasoning_end(self, input_ids: list[int]) -> bool: - """ - Check if the reasoning content ends in the input_ids. + """Check if the reasoning content ends in the input_ids. Used by structured engines like `xgrammar` to check if the reasoning content ends in the model output. @@ -204,13 +202,13 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: Returns: True if the reasoning content ends in the input_ids. + """ def is_reasoning_end_streaming( self, input_ids: list[int], delta_ids: list[int] ) -> bool: - """ - Check if the reasoning content ends during a decode step. + """Check if the reasoning content ends during a decode step. Must be a pure function of the input_ids and delta_ids. @@ -220,12 +218,12 @@ def is_reasoning_end_streaming( Returns: True if the reasoning content ends in the delta_ids. + """ return self.is_reasoning_end(input_ids) def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: - """ - Locate the token that ends reasoning within one window of tokens. + """Locate the token that ends reasoning within one window of tokens. Args: token_ids: The tokens to examine. @@ -235,13 +233,13 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: ``len(token_ids)`` when reasoning does not end in this window. Parsers that cannot answer from a window alone always return ``None``; callers fall back to :meth:`is_reasoning_end_streaming`. + """ return None @abstractmethod def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """ - Extract content token IDs from the input_ids. + """Extract content token IDs from the input_ids. This extracts the non-reasoning content (e.g., everything after the tag). @@ -251,6 +249,7 @@ def extract_content_ids(self, input_ids: list[int]) -> list[int]: Returns: The extracted content token IDs. + """ @abstractmethod @@ -259,8 +258,7 @@ def extract_reasoning( model_output: str, request: ChatCompletionRequest | ResponsesRequest, ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from a complete model-generated string. + """Extract reasoning content from a complete model-generated string. Used for non-streaming responses where we have the entire model response available before sending to the client. @@ -271,6 +269,7 @@ def extract_reasoning( Returns: A tuple of (reasoning, response_content). + """ @abstractmethod @@ -283,8 +282,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """ - Extract reasoning content from a streaming delta message. + """Extract reasoning content from a streaming delta message. Args: previous_text: Text from all previous tokens. @@ -296,6 +294,7 @@ def extract_reasoning_streaming( Returns: A DeltaMessage with reasoning and/or content fields, or None. + """ # ========== Tool Parser Methods ========== @@ -303,8 +302,7 @@ def extract_reasoning_streaming( def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: - """ - Adjust the request parameters for tool calling. + """Adjust the request parameters for tool calling. Can be overridden by subclasses to modify request parameters (e.g., setting structured output schemas for tool calling). @@ -314,6 +312,7 @@ def adjust_request( Returns: The adjusted request. + """ return request @@ -323,8 +322,7 @@ def extract_tool_calls( model_output: str, request: ChatCompletionRequest | ResponsesRequest, ) -> ExtractedToolCallInformation: - """ - Extract tool calls from a complete model-generated string. + """Extract tool calls from a complete model-generated string. Used for non-streaming responses. @@ -334,6 +332,7 @@ def extract_tool_calls( Returns: ExtractedToolCallInformation containing the tool calls. + """ @abstractmethod @@ -347,8 +346,7 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest | ResponsesRequest, ) -> DeltaMessage | None: - """ - Extract tool calls from a streaming delta message. + """Extract tool calls from a streaming delta message. Args: previous_text: Text from all previous tokens. @@ -361,6 +359,7 @@ def extract_tool_calls_streaming( Returns: A DeltaMessage with tool_calls field, or None. + """ @abstractmethod @@ -381,6 +380,7 @@ def parse( Returns: A tuple of (reasoning, content, tool_calls). + """ @abstractmethod @@ -403,8 +403,7 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: class DelegatingParser(Parser): - """ - A Parser implementation that delegates to separate ReasoningParser and + """A Parser implementation that delegates to separate ReasoningParser and ToolParser instances. This is the recommended base class for creating model-specific parsers diff --git a/vllm/parser/cohere_command.py b/vllm/parser/cohere_command.py index cf8d79b10ec5..e5ed7aed33ae 100644 --- a/vllm/parser/cohere_command.py +++ b/vllm/parser/cohere_command.py @@ -169,8 +169,7 @@ def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: def _tool_definitions_to_schema_list( tools: str | list[Any], ) -> list[CohereNormalizedTool]: - """ - Build the list of ``CohereNormalizedTool`` dicts expected by + """Build the list of ``CohereNormalizedTool`` dicts expected by ``collect_tool_schema``. Accepts: @@ -214,8 +213,7 @@ def _tool_definitions_to_schema_list( def _has_effective_tools( tools: str | list[Any] | None, ) -> TypeGuard[str | list[Any]]: - """ - True when ``tools`` contains at least one tool definition to convert. + """True when ``tools`` contains at least one tool definition to convert. ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest`` uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are @@ -236,8 +234,7 @@ def convert_schema_to_structural_tags( tools: str | list[Any] | None = None, model_architecture: str | None = None, ) -> str | None: - """ - Returns a response_format string accepted by xgrammar's structural tag format. + """Returns a response_format string accepted by xgrammar's structural tag format. Uses the canonical shape: {"type": "structural_tag", "format": {...}} with format.type "triggered_tags" and tag content type "json_schema" or "grammar". @@ -325,8 +322,7 @@ def _unwrap_nested_schema(candidate: Any) -> dict | None: def _schema_from_json_schema_field(js_wr: Any) -> dict | None: - """ - Extract the JSON Schema object from Chat Completions ``json_schema`` payload. + """Extract the JSON Schema object from Chat Completions ``json_schema`` payload. Accepts: - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field diff --git a/vllm/parser/engine/__init__.py b/vllm/parser/engine/__init__.py index 0bd26020bddd..a19f51f8ced4 100644 --- a/vllm/parser/engine/__init__.py +++ b/vllm/parser/engine/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Streaming parser engine framework for tool call and reasoning extraction. +"""Streaming parser engine framework for tool call and reasoning extraction. Instead of hand-rolling a parser for every model's tool-call / reasoning format, each format is declared as a ParserEngineConfig (terminals, diff --git a/vllm/parser/engine/incremental_lexer.py b/vllm/parser/engine/incremental_lexer.py index 8c3820b59ab1..b8f44ad5358c 100644 --- a/vllm/parser/engine/incremental_lexer.py +++ b/vllm/parser/engine/incremental_lexer.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Incremental text lexer that converts text chunks into terminal -tokens, with prefix-match buffering for ambiguous boundaries.""" +tokens, with prefix-match buffering for ambiguous boundaries. +""" from __future__ import annotations diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py index 21d053ef28da..67d52bc34c10 100644 --- a/vllm/parser/engine/streaming_parser_engine.py +++ b/vllm/parser/engine/streaming_parser_engine.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Streaming parser engine that orchestrates token ID scanning, -incremental lexing, and state-machine-driven semantic event emission.""" +incremental lexing, and state-machine-driven semantic event emission. +""" from __future__ import annotations diff --git a/vllm/parser/engine/token_id_scanner.py b/vllm/parser/engine/token_id_scanner.py index 60a64dcae281..6ef0427b5c4e 100644 --- a/vllm/parser/engine/token_id_scanner.py +++ b/vllm/parser/engine/token_id_scanner.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Scan delta token IDs for special tokens and split the stream into -pre-lexed terminals and plain text chunks.""" +pre-lexed terminals and plain text chunks. +""" from __future__ import annotations diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index e3baaf488a8d..185610126f5b 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -83,6 +83,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict: (e.g. partial boolean parsed as bare string). Returns a dict ready for ``json.dumps()``. + """ if not args_str or not args_str.strip(): return {} diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 2081fb278464..6b1bcc7efaa8 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -67,7 +67,8 @@ def _scan_json_value(raw: str, start: int) -> int | None: """Return the end index (exclusive) of the JSON object starting at - ``raw[start]``, or ``None`` when the object is still unterminated.""" + ``raw[start]``, or ``None`` when the object is still unterminated. + """ depth = 0 in_string = False escape = False diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index ce516eea1419..247ef2803fbe 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -136,6 +136,7 @@ def mistral_config( Returns: A frozen :class:`ParserEngineConfig` with ``initial_state=CONTENT``. + """ if reasoning_encoding == "special_token": think_start = _THINK_START_SPECIAL @@ -467,6 +468,7 @@ def _build_guided_schema_pre_v11( Returns: A JSON Schema dict, or ``None`` if the named tool is not found. + """ tool_choice = request.tool_choice tools = request.tools or [] diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index cc9dd3a6c659..a8efd3f0eaa1 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -16,9 +16,7 @@ class ParserManager: - """ - Provides a unified Parser by composing reasoning and tool parser adapters. - """ + """Provides a unified Parser by composing reasoning and tool parser adapters.""" @classmethod def get_tool_parser( @@ -80,8 +78,7 @@ def get_parser( model_name: str | None = None, is_harmony: bool = False, ) -> type[Parser] | None: - """ - Get a Parser that handles both reasoning and tool parsing. + """Get a Parser that handles both reasoning and tool parsing. Composes the individual parsers into a ``DelegatingParser`` subclass. @@ -95,6 +92,7 @@ def get_parser( Returns: A Parser class, or None if neither parser is specified. + """ if not tool_parser_name and not reasoning_parser_name: return None diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index 7cf0843669c3..add01346550b 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -16,9 +16,7 @@ def vllm_version_matches_substr(substr: str) -> bool: - """ - Check to see if the vLLM version matches a substring. - """ + """Check to see if the vLLM version matches a substring.""" from importlib.metadata import PackageNotFoundError, version try: diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 89b68a401e26..6cc1475b441d 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -181,9 +181,7 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: @classmethod def set_device(cls, device: torch.device) -> None: - """ - Set the device for the current platform. - """ + """Set the device for the current platform.""" torch.cpu.set_device(device) @classmethod @@ -508,8 +506,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @classmethod def discover_numa_topology(cls) -> list[list[int]]: - """ - Discover NUMA topology and keep the last physical core of each numa + """Discover NUMA topology and keep the last physical core of each numa into one core group list for nixl start_kv_load() """ SYS_NODE = "/sys/devices/system/node" @@ -570,9 +567,7 @@ def get_punica_wrapper(cls) -> str: @classmethod def get_device_communicator_cls(cls) -> str: - """ - Get device specific communicator class for distributed communication. - """ + """Get device specific communicator class for distributed communication.""" return "vllm.distributed.device_communicators.cpu_communicator.CpuCommunicator" # noqa @classmethod @@ -636,9 +631,7 @@ def pack_kv_cache( kv_cache: torch.Tensor, indices: torch.Tensor, ) -> None: - """ - Rewrite the kv cache shape for the current platform. - """ + """Rewrite the kv cache shape for the current platform.""" # Import lazily: cpu_attn pulls in _custom_ops, which needs a fully # initialized vllm.platforms (avoid circular import while CpuPlatform loads). from vllm._custom_ops import cpu_attn_reshape_and_cache diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 30eed82ccb16..96817b25d65a 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -62,7 +62,8 @@ def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int CUDA_VISIBLE_DEVICES has already been set to the desired value. # This can be removed and simply replaced with torch.cuda.get_device_count - # after https://github.com/pytorch/pytorch/pull/122815 is released.""" + # after https://github.com/pytorch/pytorch/pull/122815 is released. + """ # Note: cuda_visible_devices is not used, but we keep it as an argument for # LRU Cache purposes. @@ -265,9 +266,7 @@ def supported_dtypes(self) -> list[torch.dtype]: @classmethod def set_device(cls, device: torch.device) -> None: - """ - Set the device for the current platform. - """ + """Set the device for the current platform.""" torch.cuda.set_device(device) # With this trick we can force the device to be set eagerly # see https://github.com/pytorch/pytorch/issues/155668 @@ -838,9 +837,7 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: @classmethod @with_nvml_context def is_fully_connected(cls, physical_device_ids: list[int]) -> bool: - """ - query if the set of gpus are fully connected by nvlink (1 hop) - """ + """Query if the set of gpus are fully connected by nvlink (1 hop).""" handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids] for i, handle in enumerate(handles): for j, peer_handle in enumerate(handles): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 98db974a02e4..70937913e2d0 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -42,7 +42,8 @@ def set_assigned_physical_gpu_ids(ids: list[int]) -> None: Idempotent: a second call with the same value is a no-op. Raises RuntimeError if called again with a different value. - This is expected to run during single-threaded worker initialization.""" + This is expected to run during single-threaded worker initialization. + """ global _assigned_physical_gpu_ids if _assigned_physical_gpu_ids is not None: if _assigned_physical_gpu_ids != ids: @@ -122,8 +123,7 @@ def as_version_str(self) -> str: return f"{self.major}.{self.minor}" def to_int(self) -> int: - """ - Express device capability as an integer ``. + """Express device capability as an integer ``. It is assumed that the minor version is always a single digit. """ @@ -175,7 +175,7 @@ class Platform: @property def pass_key(self) -> str: - """Inductor config key for the PassManager custom pass""" + """Inductor config key for the PassManager custom pass.""" return "post_grad_custom_post_pass" @property @@ -238,23 +238,19 @@ def is_cumem_allocator_available(self) -> bool: @classmethod def get_pass_manager_cls(cls) -> str: - """ - Get the pass manager class for this platform. + """Get the pass manager class for this platform. It will be registered as a custom pass under the current_platform.pass_key. """ return "vllm.compilation.passes.pass_manager.PostGradPassManager" @classmethod def get_compile_backend(cls) -> str: - """ - Get the custom compile backend for current platform. - """ + """Get the custom compile backend for current platform.""" return cls.simple_compile_backend @classmethod def import_ir_kernels(cls) -> None: - """ - The default implementation imports ``vllm.kernels``, which registers + """The default implementation imports ``vllm.kernels``, which registers the built-in IR op implementations. Out-of-tree (OOT) platforms should override this method to import their own kernel modules. """ @@ -392,8 +388,7 @@ def get_vit_attn_backend( dtype: torch.dtype, backend: "AttentionBackendEnum | None" = None, ) -> "AttentionBackendEnum": - """ - Get the vision attention backend class of a device. + """Get the vision attention backend class of a device. NOTE: ViT Attention should be checked and override in the platform-specific implementation. we should not override this in any other places, like @@ -426,6 +421,7 @@ def get_device_capability( Args: device_id: Device index in the visible device namespace, matching the argument accepted by torch.cuda. + """ return None @@ -435,8 +431,7 @@ def has_device_capability( capability: tuple[int, int] | int, device_id: int = 0, ) -> bool: - """ - Test whether this platform is compatible with a device capability. + """Test whether this platform is compatible with a device capability. The `capability` argument can either be: @@ -459,8 +454,7 @@ def is_device_capability( capability: tuple[int, int] | int, device_id: int = 0, ) -> bool: - """ - Test whether this platform has exactly the specified device capability. + """Test whether this platform has exactly the specified device capability. The `capability` argument can either be: @@ -483,8 +477,7 @@ def is_device_capability_family( capability: int, device_id: int = 0, ) -> bool: - """ - Returns True if the device capability is any .x. + """Returns True if the device capability is any .x. Mirrors CUDA 13 'family' architecture semantics (e.g. 10.x, 11.x, 12.x). """ current_capability = cls.get_device_capability(device_id=device_id) @@ -532,9 +525,7 @@ def inference_mode(cls): @classmethod def set_device(cls, device: torch.device) -> None: - """ - Set the device for the current platform. - """ + """Set the device for the current platform.""" raise NotImplementedError @classmethod @@ -546,8 +537,7 @@ def manual_seed_all(cls, seed: int) -> None: def pre_register_and_update( cls, parser: FlexibleArgumentParser | None = None ) -> None: - """ - Do some pre-registration or update action for the current platform. + """Do some pre-registration or update action for the current platform. This function is called before global VllmConfig is initialized or cli arguments are parsed. It's used for out-of-tree platforms to register or @@ -560,8 +550,7 @@ def pre_register_and_update( @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: - """ - Apply the platform-specific default values to the config. + """Apply the platform-specific default values to the config. This function is called during the initialization of global VllmConfig, after parsing cli arguments. @@ -574,8 +563,7 @@ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: @classmethod def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: - """ - Check and update the configuration for the current platform. + """Check and update the configuration for the current platform. It can raise an exception if the configuration is not compatible with the current platform, or it can update the configuration to make it @@ -607,8 +595,7 @@ def _find_non_ssm_backend( @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: - """ - Ensure block_size is compatible with the attention backend. + """Ensure block_size is compatible with the attention backend. For hybrid models, also aligns block_size with mamba page sizes. """ from vllm.config.cache import CacheConfig @@ -780,8 +767,7 @@ def _align_hybrid_block_size( vllm_config: "VllmConfig", backend_cls: "type[AttentionBackend]", ) -> None: - """ - For hybrid attention/mamba models, ensure that the attention page + """For hybrid attention/mamba models, ensure that the attention page size is >= the mamba page size, and pad the mamba page size to match. """ from math import lcm @@ -963,15 +949,12 @@ def _align_hybrid_block_size( @classmethod def register_custom_kv_cache_specs(cls, vllm_config: "VllmConfig") -> None: - """ - Register custom KVCacheSpec class on current platform. - """ + """Register custom KVCacheSpec class on current platform.""" pass @classmethod def verify_model_arch(cls, model_arch: str) -> None: - """ - Verify whether the current platform supports the specified model + """Verify whether the current platform supports the specified model architecture. - This will raise an Error or Warning based on the model support on @@ -982,9 +965,7 @@ def verify_model_arch(cls, model_arch: str) -> None: @classmethod def verify_quantization(cls, quant: str) -> None: - """ - Verify whether the quantization is supported by the current platform. - """ + """Verify whether the quantization is supported by the current platform.""" if cls.supported_quantization and quant not in cls.supported_quantization: raise ValueError( f"{quant} quantization is currently not supported in {cls.device_name}." @@ -992,8 +973,7 @@ def verify_quantization(cls, quant: str) -> None: @classmethod def get_cpu_architecture(cls) -> CpuArchEnum: - """ - Determine the CPU architecture of the current system. + """Determine the CPU architecture of the current system. Returns CpuArchEnum indicating the architecture type. """ machine = platform.machine().lower() @@ -1031,50 +1011,37 @@ def is_pin_memory_available(cls) -> bool: def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: - """ - Return the memory usage in bytes. - """ + """Return the memory usage in bytes.""" raise NotImplementedError @classmethod def get_punica_wrapper(cls) -> str: - """ - Return the punica wrapper for current platform. - """ + """Return the punica wrapper for current platform.""" raise NotImplementedError @classmethod def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]: - """ - Return the platform specific values for (-inf, inf) - """ + """Return the platform specific values for (-inf, inf).""" return float("-inf"), float("inf") @classmethod def can_update_inplace(cls) -> bool: - """ - Checks if the platform allows inplace memory updates - """ + """Checks if the platform allows inplace memory updates.""" return True @classmethod def get_lora_vocab_padding_size(cls) -> int: - """ - Returns how much padding the LoRA logits need for kernels - """ + """Returns how much padding the LoRA logits need for kernels.""" return 256 @classmethod def get_device_communicator_cls(cls) -> str: - """ - Get device specific communicator class for distributed communication. - """ + """Get device specific communicator class for distributed communication.""" return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa @classmethod def is_integrated_gpu(cls, device_id: int = 0) -> bool: - """ - Returns whether the GPU is an integrated (UMA) device that shares + """Returns whether the GPU is an integrated (UMA) device that shares system memory with the CPU. On UMA systems (e.g. NVIDIA GH200, DGX Spark, Jetson Orin), @@ -1085,22 +1052,17 @@ def is_integrated_gpu(cls, device_id: int = 0) -> bool: @classmethod def supports_mx(cls) -> bool: - """ - Returns whether the current platform supports MX types. - """ + """Returns whether the current platform supports MX types.""" return False @classmethod def supports_fp8(cls) -> bool: - """ - Returns whether the current platform supports FP8 types. - """ + """Returns whether the current platform supports FP8 types.""" return False @classmethod def is_fp8_fnuz(cls) -> bool: - """ - Returns whether the preferred FP8 type is FNUZ on the current platform. + """Returns whether the preferred FP8 type is FNUZ on the current platform. There are two representations of FP8, OCP FP8 and FNUZ FP8. The OCP specification can be found at https://tinyurl.com/b7jvwpft. @@ -1113,8 +1075,7 @@ def is_fp8_fnuz(cls) -> bool: @classmethod def fp8_dtype(cls) -> torch.dtype: - """ - Returns the preferred FP8 type on the current platform. + """Returns the preferred FP8 type on the current platform. See the documentation for is_fp8_fnuz for details. """ @@ -1122,22 +1083,17 @@ def fp8_dtype(cls) -> torch.dtype: @classmethod def use_all_gather(cls) -> bool: - """ - Whether to use allgather in LogitsProcessor to gather the logits. - """ + """Whether to use allgather in LogitsProcessor to gather the logits.""" return True @classmethod def use_custom_allreduce(cls) -> bool: - """ - Returns if custom allreduce is supported on the current platform - """ + """Returns if custom allreduce is supported on the current platform.""" return False @classmethod def opaque_attention_op(cls) -> bool: - """ - Returns True if we register attention as one giant opaque custom op + """Returns True if we register attention as one giant opaque custom op on the current platform """ return False @@ -1148,7 +1104,7 @@ def validate_request( processed_inputs: "EngineInput", params: "SamplingParams | PoolingParams", ) -> None: - """Raises if this request is unsupported on this platform""" + """Raises if this request is unsupported on this platform.""" def __getattr__(self, key: str): # Pickle checks dunder methods like __getstate__. If we return None @@ -1172,9 +1128,7 @@ def __getattr__(self, key: str): return None def get_global_graph_pool(self) -> Any: - """ - Return the global graph pool for this platform. - """ + """Return the global graph pool for this platform.""" cls = self.__class__ if cls._global_graph_pool is None: cls._global_graph_pool = self.graph_pool_handle() @@ -1182,9 +1136,7 @@ def get_global_graph_pool(self) -> Any: @classmethod def get_static_graph_wrapper_cls(cls) -> str: - """ - Get static graph wrapper class for static graph. - """ + """Get static graph wrapper class for static graph.""" return "vllm.compilation.base_static_graph.AbstractStaticGraphWrapper" @classmethod @@ -1196,36 +1148,27 @@ def stateless_init_device_torch_dist_pg( group_size: int, timeout: timedelta, ) -> "ProcessGroup": - """ - Init platform-specific torch distributed process group. - """ + """Init platform-specific torch distributed process group.""" raise NotImplementedError @classmethod def check_if_supports_dtype(cls, dtype: torch.dtype): - """ - Check if the dtype is supported by the current platform. - """ + """Check if the dtype is supported by the current platform.""" raise NotImplementedError @classmethod def support_hybrid_kv_cache(cls) -> bool: - """ - Returns if the hybrid kv cache is supported by the current platform. - """ + """Returns if the hybrid kv cache is supported by the current platform.""" return False @classmethod def support_static_graph_mode(cls) -> bool: - """ - Returns if the graph mode is supported by the current platform. - """ + """Returns if the graph mode is supported by the current platform.""" return False @classmethod def check_runner_kv_caches_multi_layer(cls) -> None: - """ - Check whether the platform's ModelRunner can handle multiple attention + """Check whether the platform's ModelRunner can handle multiple attention layers that share the same layer index (e.g. cross attention and self attention in the same decoder block of an encoder-decoder model such as BART). @@ -1241,15 +1184,12 @@ def check_runner_kv_caches_multi_layer(cls) -> None: @classmethod def support_deep_gemm(cls) -> bool: - """ - Returns if DeepGEMM is supported by the current platform. - """ + """Returns if DeepGEMM is supported by the current platform.""" return False @classmethod def use_custom_op_collectives(cls) -> bool: - """ - Whether this platform should use torch.ops.vllm.* custom ops for collectives. + """Whether this platform should use torch.ops.vllm.* custom ops for collectives. Returns False by default - platforms must explicitly opt-in. """ @@ -1257,16 +1197,12 @@ def use_custom_op_collectives(cls) -> bool: @classmethod def use_sync_weight_loader(cls) -> bool: - """ - Returns if the current platform needs to sync weight loader. - """ + """Returns if the current platform needs to sync weight loader.""" return False @classmethod def make_synced_weight_loader(cls, original_weight_loader): - """ - Wrap the original weight loader to make it synced. - """ + """Wrap the original weight loader to make it synced.""" if not cls.use_sync_weight_loader(): return original_weight_loader @@ -1280,37 +1216,29 @@ def _synced_weight_loader(param, *args, **kwargs): @classmethod def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]: - """ - Returns a mapping from device_type to a tuple of supported + """Returns a mapping from device_type to a tuple of supported kv_buffer_device for nixl. """ return {} @classmethod def get_nixl_memory_type(cls) -> str | None: - """ - Returns the nixl memory type for the current platform. - """ + """Returns the nixl memory type for the current platform.""" return None @classmethod def check_max_model_len(cls, max_model_len: int) -> int: - """ - Check max_model_len for the current platform. - """ + """Check max_model_len for the current platform.""" return max_model_len @classmethod def set_additional_forward_context(cls, *args, **kwargs) -> dict[str, Any]: - """ - Set some additional forward context for the current platform if needs. - """ + """Set some additional forward context for the current platform if needs.""" return {} @classmethod def num_compute_units(cls, device_id: int = 0) -> int: - """ - Get the number of compute units for the current platform. + """Get the number of compute units for the current platform. (NVIDIA SM / AMD CU / Intel EU) """ raise NotImplementedError( @@ -1329,16 +1257,12 @@ def get_default_ir_op_priority( @classmethod def is_arch_support_pdl(cls) -> bool: - """ - Does the current platform support PDL (Programmatic Dependent Launch)? - """ + """Does the current platform support PDL (Programmatic Dependent Launch)?""" return False @classmethod def validate_environ(cls, hard_fail: bool) -> None: - """ - Validate environment variables for the current platform. - """ + """Validate environment variables for the current platform.""" from vllm import envs for env in os.environ: diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index b3b7fe8b089c..9bf95557ce30 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -91,7 +91,8 @@ def _rocm_device_count_stateless(cuda_visible_devices: str | None = None) -> int CUDA_VISIBLE_DEVICES has already been set to the desired value. # This can be removed and simply replaced with torch.cuda.get_device_count - # after https://github.com/pytorch/pytorch/pull/122815 is released.""" + # after https://github.com/pytorch/pytorch/pull/122815 is released. + """ # Note: cuda_visible_devices is not used, but we keep it as an argument for # LRU Cache purposes. @@ -126,7 +127,8 @@ def _get_wsl_kernel_version() -> tuple[int, ...] | None: def _sync_hip_cuda_env_vars(): """Ensure HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES are consistent. - Treats empty string as unset. Raises on genuine conflicts.""" + Treats empty string as unset. Raises on genuine conflicts. + """ hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None @@ -187,8 +189,7 @@ def _query_total_memory_from_amdsmi(physical_device_id: int) -> int: def _get_gcn_arch() -> str: - """ - Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda. + """Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda. Called once at module level; result stored in _GCN_ARCH. """ try: @@ -223,8 +224,7 @@ def _get_gcn_arch() -> str: def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: - """ - Parse (major, minor) from a GCN arch string, mirroring how + """Parse (major, minor) from a GCN arch string, mirroring how HIP derives hipDeviceProp_t.major / .minor. Format: gfx @@ -238,6 +238,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: Returns None only when the string is not gfx-prefixed at all (i.e. not a ROCm arch string). Raises on any string that looks like a GCN arch but does not match a known layout. + """ m = re.match(r"gfx(\d+)", gcn_arch) if not m: @@ -778,9 +779,7 @@ def get_vit_attn_backend( @classmethod def set_device(cls, device: torch.device) -> None: - """ - Set the device for the current platform. - """ + """Set the device for the current platform.""" torch.cuda.set_device(device) @classmethod @@ -805,9 +804,7 @@ def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: @classmethod @with_amdsmi_context def is_fully_connected(cls, physical_device_ids: list[int]) -> bool: - """ - Query if the set of gpus are fully connected by xgmi (1 hop) - """ + """Query if the set of gpus are fully connected by xgmi (1 hop).""" handles = [amdsmi_get_processor_handles()[i] for i in physical_device_ids] for i, handle in enumerate(handles): for j, peer_handle in enumerate(handles): diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 6920c55e71b7..483c41197e1c 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -31,8 +31,7 @@ def get_mem_info_wrapper( device: int | str | torch.device | None = None, ) -> tuple[int, int]: - """ - Get memory info for a device, compatible with torch.accelerator.get_memory_info API. + """Get memory info for a device, matching `torch.accelerator.get_memory_info`. Args: device: Device specification. Can be: @@ -43,6 +42,7 @@ def get_mem_info_wrapper( Returns: Tuple[int, int]: (free_memory, total_memory) in bytes + """ # Handle None - use current device if device is None: @@ -237,9 +237,7 @@ def get_vit_attn_backend( @classmethod def set_device(cls, device: torch.device) -> None: - """ - Set the device for the current platform. - """ + """Set the device for the current platform.""" torch.xpu.set_device(device) @classmethod diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py index 95e895c279be..7bebfdf03e4d 100644 --- a/vllm/plugins/__init__.py +++ b/vllm/plugins/__init__.py @@ -115,6 +115,7 @@ def load_endpoint_plugins( Returns: Instantiated plugins that passed gating in discovery order. + """ from importlib.metadata import entry_points diff --git a/vllm/plugins/lora_resolvers/filesystem_resolver.py b/vllm/plugins/lora_resolvers/filesystem_resolver.py index 31db0ac4838e..822f410f7dad 100644 --- a/vllm/plugins/lora_resolvers/filesystem_resolver.py +++ b/vllm/plugins/lora_resolvers/filesystem_resolver.py @@ -47,8 +47,7 @@ async def _get_lora_req_from_path( def register_filesystem_resolver(): - """Register the filesystem LoRA Resolver with vLLM""" - + """Register the filesystem LoRA Resolver with vLLM.""" lora_cache_dir = envs.VLLM_LORA_RESOLVER_CACHE_DIR if lora_cache_dir: if not os.path.exists(lora_cache_dir) or not os.path.isdir(lora_cache_dir): diff --git a/vllm/plugins/lora_resolvers/hf_hub_resolver.py b/vllm/plugins/lora_resolvers/hf_hub_resolver.py index 57ccdc2e48d8..9ab249b1987e 100644 --- a/vllm/plugins/lora_resolvers/hf_hub_resolver.py +++ b/vllm/plugins/lora_resolvers/hf_hub_resolver.py @@ -67,6 +67,7 @@ async def _resolve_repo(self, lora_name: str) -> str | None: lora_name: Path to LoRA in HF Hub, e.g., //, match on / (if it contains an adapter directly) or // if it may have one in subdirs. + """ for potential_repo in self.repo_list: if lora_name.startswith(potential_repo) and ( @@ -86,6 +87,7 @@ async def _resolve_repo_subpath( Args: lora_name: Path to LoRA in HF Hub, e.g., // maybe_repo: Path to the repo to match against if one exists. + """ if maybe_repo is None: return None @@ -108,6 +110,7 @@ async def _get_adapter_dirs(self, repo_name: str) -> set[str]: Args: repo_name: Name of the HF hub repo to inspect. + """ repo_files = await asyncio.to_thread( hf_api().list_repo_files, @@ -124,8 +127,7 @@ async def _get_adapter_dirs(self, repo_name: str) -> set[str]: def register_hf_hub_resolver(): - """Register the Hf hub LoRA Resolver with vLLM""" - + """Register the Hf hub LoRA Resolver with vLLM.""" hf_repo_list = envs.VLLM_LORA_RESOLVER_HF_REPO_LIST is_enabled = ( envs.VLLM_PLUGINS is not None and "lora_hf_hub_resolver" in envs.VLLM_PLUGINS diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 272552d61b8d..319b77d5986a 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -28,6 +28,7 @@ class LateInteractionParams( - "score_doc": score a document against a cached query. query_key: stable key used for both DP routing and worker cache lookup. query_uses: expected number of document requests + """ mode: str @@ -47,6 +48,7 @@ class PoolingParams( `None` uses the pooler's default, which is `True` in most cases. dimensions: Reduce the dimensions of embeddings if model support matryoshka representation. + """ # --8<-- [start:common-pooling-params] diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py index a36e4611f3ce..1393e123fbcc 100644 --- a/vllm/profiler/layerwise_profile.py +++ b/vllm/profiler/layerwise_profile.py @@ -224,7 +224,7 @@ def _get_kineto_gpu_event(self, node: _ModuleTreeNode): return next(iterator, None) def _cumulative_cuda_time(self, node: _ModuleTreeNode): - "Return cuda time in microseconds" + """Return cuda time in microseconds.""" def _cumulative_cuda_time_recursive(node: _ModuleTreeNode): if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)): @@ -372,13 +372,13 @@ def df_traversal( class layerwise_profile(profile): def __init__(self, num_running_seqs: int | None = None): - """ - layerwise profile constructor. + """Layerwise profile constructor. Args: num_running_seqs (Optional[int], optional): When given, num_running_seqs will be passed to LayerProfileResults for metadata update. Defaults to None. + """ super().__init__( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index abe8a3431798..ae9da1f4f7c9 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -98,7 +98,8 @@ def start(self) -> None: def step(self) -> None: """Update the profiler state at each worker step, - to handle delayed starts and max iteration limits.""" + to handle delayed starts and max iteration limits. + """ if not self._active: return @@ -138,6 +139,7 @@ def _profiler_step(self) -> bool: Returns: True if the step was an active profiling step (data recorded), False if the step was a warmup step (data discarded). + """ return True @@ -354,6 +356,7 @@ def _profiler_step(self) -> bool: Returns: True if the step was an active profiling step (data recorded), False if the step was a warmup step (data discarded). + """ if self._uses_schedule: self.profiler.step() diff --git a/vllm/ray/lazy_utils.py b/vllm/ray/lazy_utils.py index 06c91cc3943a..674e8429e811 100644 --- a/vllm/ray/lazy_utils.py +++ b/vllm/ray/lazy_utils.py @@ -16,7 +16,6 @@ def is_ray_initialized(): def is_in_ray_actor(): """Check if we are in a Ray actor.""" - try: import ray diff --git a/vllm/ray/ray_env.py b/vllm/ray/ray_env.py index 4103b93172ce..2e2ed48b85c7 100644 --- a/vllm/ray/ray_env.py +++ b/vllm/ray/ray_env.py @@ -76,6 +76,7 @@ def get_env_vars_to_copy( additional_vars: Extra individual env var names to copy. Useful for caller-specific vars (e.g. platform env vars). destination: Label used in log messages only. + """ exclude = (exclude_vars or set()) | RAY_NON_CARRY_OVER_ENV_VARS diff --git a/vllm/reasoning/abs_reasoning_parsers.py b/vllm/reasoning/abs_reasoning_parsers.py index 6fb694545a4a..accc8c9c2727 100644 --- a/vllm/reasoning/abs_reasoning_parsers.py +++ b/vllm/reasoning/abs_reasoning_parsers.py @@ -23,8 +23,7 @@ class ReasoningParser: - """ - Abstract reasoning parser class that should not be used directly. + """Abstract reasoning parser class that should not be used directly. Provided and methods should be used in derived classes. It is used to extract reasoning content from the model output. @@ -71,26 +70,27 @@ def has_engine_confirmed_reasoning_end(self) -> bool: @abstractmethod def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - """ - Check if the reasoning content ends in the input_ids. + """Check if the reasoning content ends in the input_ids. It is used in structured engines like `xgrammar` to check if the reasoning content ends in the model output. - Parameters: + Parameters + ---------- input_ids: list[int] The input_ids of the model output. - Returns: + Returns + ------- bool True if the reasoning content ends in the input_ids. + """ def is_reasoning_end_streaming( self, input_ids: Sequence[int], delta_ids: Iterable[int] ) -> bool: - """ - Check if the reasoning content ends in the input_ids on a + """Check if the reasoning content ends in the input_ids on a decode step. It is used in structured engines like `xgrammar` to check if the @@ -98,29 +98,36 @@ def is_reasoning_end_streaming( `input_ids` the entire model output and `delta_ids` are the last few computed tokens of the model output (like during a decode step). - Parameters: + Parameters + ---------- input_ids: list[int] The entire model output. delta_ids: list[int] The last few computed tokens of the model output at the current decode step. - Returns: + Returns + ------- bool True if the reasoning content ends in the `delta_ids` on a decode step. + """ return self.is_reasoning_end(input_ids) @abstractmethod def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """ - Extract content token ids from the input_ids. - Parameters: + """Extract content token ids from the input_ids. + + Parameters + ---------- input_ids: list[int] The input_ids of the model output. - Returns: + + Returns + ------- list[int] The extracted content from the input_ids. + """ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: @@ -137,8 +144,8 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: Returns: int: Number of tokens that belong to reasoning content. - """ + """ # By default, assume the parser cannot detect reasoning spans. return 0 @@ -148,18 +155,18 @@ def extract_reasoning( model_output: str, request: "ChatCompletionRequest | ResponsesRequest", ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from a complete model-generated string. + """Extract reasoning content from a complete model-generated string. Used for non-streaming responses where we have the entire model response available before sending to the client. - Parameters: + Args: model_output: The model-generated string to extract reasoning content from. request: The request object that was used to generate the model_output. Returns: A tuple containing the reasoning content and the content. + """ @abstractmethod @@ -172,8 +179,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> "DeltaMessage | None": - """ - Instance method that should be implemented for extracting reasoning + """Instance method that should be implemented for extracting reasoning from an incomplete response; for use when handling reasoning calls and streaming. Has to be an instance method because it requires state - the current tokens/diffs, but also the information about what has @@ -203,15 +209,12 @@ def prepare_structured_tag( original_tag: str | None, tool_server: ToolServer | None, ) -> str | None: - """ - Instance method that is implemented for preparing the structured tag - """ + """Instance method that is implemented for preparing the structured tag""" return original_tag class ReasoningParserManager: - """ - Central registry for ReasoningParser implementations. + """Central registry for ReasoningParser implementations. Supports two registration modes: - Eager registration via `register_module` @@ -225,14 +228,14 @@ class ReasoningParserManager: @classmethod def get_reasoning_parser(cls, name: str) -> type[ReasoningParser]: - """ - Retrieve a registered or lazily registered ReasoningParser class. + """Retrieve a registered or lazily registered ReasoningParser class. If the parser is lazily registered, it will be imported and cached on first access. Raises: KeyError: if no parser is found under the given name. + """ if name in cls.reasoning_parsers: return cls.reasoning_parsers[name] @@ -303,8 +306,7 @@ def _register_module( @classmethod def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None: - """ - Register a lazy module mapping for delayed import. + """Register a lazy module mapping for delayed import. Example: ReasoningParserManager.register_lazy_module( @@ -312,6 +314,7 @@ def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> N module_path="vllm.reasoning.qwen3_engine_reasoning_parser", class_name="Qwen3ParserReasoningAdapter", ) + """ cls.lazy_parsers[name] = (module_path, class_name) @@ -324,8 +327,7 @@ def register_module( ) -> ( type[ReasoningParser] | Callable[[type[ReasoningParser]], type[ReasoningParser]] ): - """ - Register module with the given name or name list. it can be used as a + """Register module with the given name or name list. it can be used as a decoder(with module as None) or normal function(with module as not None). """ diff --git a/vllm/reasoning/basic_parsers.py b/vllm/reasoning/basic_parsers.py index f85ccd010a8d..e1d309250648 100644 --- a/vllm/reasoning/basic_parsers.py +++ b/vllm/reasoning/basic_parsers.py @@ -16,8 +16,7 @@ class BaseThinkingReasoningParser(ReasoningParser): - """ - Base class for reasoning parsers that use thinking tokens. + """Base class for reasoning parsers that use thinking tokens. This class provides common functionality for parsers that use start and end tokens to delimit reasoning content ( @@ -87,9 +86,7 @@ def is_reasoning_end_streaming( return end_token_id in delta_ids def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """ - Extract the content after the end tokens - """ + """Extract the content after the end tokens.""" if self.end_token_id not in islice(input_ids, 0, max(0, len(input_ids) - 1)): return [] else: @@ -104,8 +101,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """ - Extract reasoning content from a delta message. + """Extract reasoning content from a delta message. Handles streaming output where previous + delta = current. Uses token IDs for faster processing. """ @@ -157,8 +153,7 @@ def extract_reasoning_streaming( def extract_reasoning( self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from the model output. + """Extract reasoning content from the model output. This is the base implementation that works for most models. Subclasses can override this method for specific behavior. diff --git a/vllm/reasoning/deepseek_r1_reasoning_parser.py b/vllm/reasoning/deepseek_r1_reasoning_parser.py index f16f87f6b0a8..991cf016599d 100644 --- a/vllm/reasoning/deepseek_r1_reasoning_parser.py +++ b/vllm/reasoning/deepseek_r1_reasoning_parser.py @@ -8,8 +8,7 @@ class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser): - """ - Reasoning parser for DeepSeek R1 model. + """Reasoning parser for DeepSeek R1 model. The DeepSeek R1 model uses ... tokens to denote reasoning text. This parser extracts the reasoning content from the model output. diff --git a/vllm/reasoning/deepseek_v3_reasoning_parser.py b/vllm/reasoning/deepseek_v3_reasoning_parser.py index 302024a8e77d..072048af91a4 100644 --- a/vllm/reasoning/deepseek_v3_reasoning_parser.py +++ b/vllm/reasoning/deepseek_v3_reasoning_parser.py @@ -18,8 +18,7 @@ class DeepSeekV3ReasoningParser(ReasoningParser): - """ - V3 parser that delegates to either DeepSeekR1ReasoningParser or + """V3 parser that delegates to either DeepSeekR1ReasoningParser or IdentityReasoningParser based on `thinking` and `separate_reasoning`. """ @@ -81,9 +80,7 @@ def extract_reasoning_streaming( class DeepSeekV3ReasoningWithThinkingParser(DeepSeekV3ReasoningParser): - """ - DeepSeekV3ReasoningParser that defaults to thinking mode. - """ + """DeepSeekV3ReasoningParser that defaults to thinking mode.""" def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} diff --git a/vllm/reasoning/ernie45_reasoning_parser.py b/vllm/reasoning/ernie45_reasoning_parser.py index 0aa926281988..444467b68e66 100644 --- a/vllm/reasoning/ernie45_reasoning_parser.py +++ b/vllm/reasoning/ernie45_reasoning_parser.py @@ -15,8 +15,7 @@ class Ernie45ReasoningParser(BaseThinkingReasoningParser): - """ - Reasoning parser for Ernie45 thinking model. + """Reasoning parser for Ernie45 thinking model. The Ernie45 thinking model output format is abc\n\n\n\ndef\n\n or abc\n\ndef @@ -60,8 +59,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """ - Extract reasoning content from a delta message. + """Extract reasoning content from a delta message. Handles streaming output where previous + delta = current. Uses token IDs for faster processing. The Ernie45 thinking model output format is @@ -138,8 +136,7 @@ def extract_reasoning_streaming( def extract_reasoning( self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" ) -> tuple[str | None, str | None]: - """ - Extract reasoning content from the model output. + """Extract reasoning content from the model output. The Ernie45 thinking model output format is abc\n\n\n\n\ndef\n\n or abc\n\ndef diff --git a/vllm/reasoning/gemma4_utils.py b/vllm/reasoning/gemma4_utils.py index 9cdac72039e8..9566b8af56b8 100644 --- a/vllm/reasoning/gemma4_utils.py +++ b/vllm/reasoning/gemma4_utils.py @@ -75,6 +75,7 @@ def parse_thinking_output(text: str) -> dict[str, str | None]: >>> result = parse_thinking_output(output_text) >>> print(result["thinking"]) # chain-of-thought reasoning or None >>> print(result["answer"]) # final answer + """ if _THINKING_END_TAG in text: parts = text.split(_THINKING_END_TAG, 1) diff --git a/vllm/reasoning/gptoss_reasoning_parser.py b/vllm/reasoning/gptoss_reasoning_parser.py index a6f108015812..5505f1de446b 100644 --- a/vllm/reasoning/gptoss_reasoning_parser.py +++ b/vllm/reasoning/gptoss_reasoning_parser.py @@ -14,8 +14,7 @@ class GptOssReasoningParser(ReasoningParser): - """ - Reasoning parser for GptOss model. + """Reasoning parser for GptOss model. The GptOss model uses harmony to extract reasoning content and this parser is only used for detecting the end of the reasoning content. diff --git a/vllm/reasoning/granite_reasoning_parser.py b/vllm/reasoning/granite_reasoning_parser.py index 463659f8d03d..dc6118224f7e 100644 --- a/vllm/reasoning/granite_reasoning_parser.py +++ b/vllm/reasoning/granite_reasoning_parser.py @@ -16,8 +16,7 @@ class GraniteReasoningParser(ReasoningParser): - """ - Reasoning parser for IBM Granite. + """Reasoning parser for IBM Granite. IBM granite models currently use "Here is my thought process:" and "Here is my response:" to separate its thinking / response outputs. @@ -65,6 +64,7 @@ def extract_reasoning( Returns: tuple[Optional[str], Optional[str]]: Tuple pair containing the reasoning content and non-reasoning content. + """ re_match = self.reasoning_regex.findall(model_output) if not re_match: @@ -108,6 +108,7 @@ def extract_reasoning_streaming( Returns: Union[DeltaMessage, None] DeltaMessage with either reasoning content or content, or None. + """ reasoning, resp_seq_len, content = self._get_content_sections(current_text) # Either we haven't finished the start of the reasoning sequence, @@ -142,6 +143,7 @@ def _is_reasoning_start_substr(self, text: str) -> bool: Returns: bool: True if any of the possible reasoning start seqs match. + """ return any( think_start.startswith(text) for think_start in self.valid_think_starts @@ -155,6 +157,7 @@ def _is_response_start_substr(self, text: str) -> bool: Returns: bool: True if any of the possible response start seqs match. + """ return any( response_start.startswith(text) @@ -175,6 +178,7 @@ def _get_delta_message_with_no_reasoning_bounds( Returns: DeltaMessage: Message containing the parsed content. + """ prev_longest_length = len(current_text) - len(delta_text) is_substr = self._is_reasoning_start_substr(current_text) @@ -213,6 +217,7 @@ def _get_delta_message_with_no_response_bounds( Returns: DeltaMessage: Message containing the parsed content. + """ # If we have no reasoning content or explicitly end with the start of # response sequence, we are in transition to the response; need to be @@ -293,6 +298,7 @@ def _get_delta_message_with_both_bounds( Returns: DeltaMessage: Message containing the parsed content. + """ # Always have content; take length to the end delta_content = delta_text[-len(response_content) :] @@ -329,6 +335,7 @@ def _get_content_sections( tuple[Optional[str], Optional[int], Optional[str]]: Tuple of len 3 containing the reasoning content, the length of the response seq (if there is one) and the non-reasoning content. + """ current_chunk_start = 0 start_reasoning = None diff --git a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py index b15b5608fbe9..0a7e021a464c 100644 --- a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py +++ b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py @@ -16,8 +16,7 @@ class HunyuanA13BReasoningParser(ReasoningParser): - """ - Reasoning parser for Hunyuan A13B Model + """Reasoning parser for Hunyuan A13B Model. HunyuanReasoningParser @@ -102,8 +101,8 @@ def extract_reasoning( Returns: tuple[Optional[str], Optional[str]]: Tuple pair containing the reasoning content and non-reasoning content. - """ + """ re_match = self.full_match_reasoning_regex.findall(model_output) if re_match: reasoning, response_content = re_match[0] @@ -139,7 +138,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """Extract content using token ID sequence state machine""" + """Extract content using token ID sequence state machine.""" # Define sequences think_start_sequence = self.think_start_ids response_start_sequence = self.response_start_ids diff --git a/vllm/reasoning/hy_v3_reasoning_parser.py b/vllm/reasoning/hy_v3_reasoning_parser.py index 0765be74af73..77a1a271d4a8 100644 --- a/vllm/reasoning/hy_v3_reasoning_parser.py +++ b/vllm/reasoning/hy_v3_reasoning_parser.py @@ -18,8 +18,7 @@ class HYV3ReasoningParser(BaseThinkingReasoningParser): - """ - HYV3 parser that delegates to either HYV3ReasoningParser or + """HYV3 parser that delegates to either HYV3ReasoningParser or IdentityReasoningParser based on `reasoning_effort`. The HYV3 model uses ... tokens to denote reasoning text. diff --git a/vllm/reasoning/hy_v4_reasoning_parser.py b/vllm/reasoning/hy_v4_reasoning_parser.py index b65db92b92f5..65035ed53ad4 100644 --- a/vllm/reasoning/hy_v4_reasoning_parser.py +++ b/vllm/reasoning/hy_v4_reasoning_parser.py @@ -46,8 +46,8 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str: RuntimeError: The tokenizer declares the structural tokens through ``model_specific_special_tokens``, which transformers 5 no longer round-trips. - """ + """ import transformers if int(transformers.__version__.split(".")[0]) >= 5: diff --git a/vllm/reasoning/identity_reasoning_parser.py b/vllm/reasoning/identity_reasoning_parser.py index 52fab9914105..f2c9a1e21e33 100644 --- a/vllm/reasoning/identity_reasoning_parser.py +++ b/vllm/reasoning/identity_reasoning_parser.py @@ -15,8 +15,7 @@ class IdentityReasoningParser(ReasoningParser): - """ - Identity reasoning parser. + """Identity reasoning parser. This parser does not attempt to parse or strip out reasoning tokens. It treats the entire model output as content and ignores reasoning. diff --git a/vllm/reasoning/minimax_m2_reasoning_parser.py b/vllm/reasoning/minimax_m2_reasoning_parser.py index b26ea8889ae0..8fa5b45d5853 100644 --- a/vllm/reasoning/minimax_m2_reasoning_parser.py +++ b/vllm/reasoning/minimax_m2_reasoning_parser.py @@ -15,8 +15,7 @@ class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter): # type: ignore[valid-type, misc] - """ - Reasoning parser for MiniMax M2 model. + """Reasoning parser for MiniMax M2 model. MiniMax M2 models don't generate start token, only end token. All content before is reasoning, content after is the @@ -25,9 +24,7 @@ class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter): # type: ignore class MiniMaxM2AppendThinkReasoningParser(ReasoningParser): - """ - Reasoning parser for MiniMax M2 model. - """ + """Reasoning parser for MiniMax M2 model.""" def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__(tokenizer, *args, **kwargs) diff --git a/vllm/reasoning/olmo3_reasoning_parser.py b/vllm/reasoning/olmo3_reasoning_parser.py index b658b92bcbe1..8ac75f91ff76 100644 --- a/vllm/reasoning/olmo3_reasoning_parser.py +++ b/vllm/reasoning/olmo3_reasoning_parser.py @@ -32,8 +32,7 @@ def __len__(self): def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]: - """ - Find the longest overlap where the end of string a matches the start + """Find the longest overlap where the end of string a matches the start of string b. Args: @@ -43,8 +42,8 @@ def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]: Returns: Tuple of IndicesTuples representing the overlapping portions in each string, or a tuple of None if no overlap exists - """ + """ # swap so a is always the shorter string a, b, swap = (a, b, False) if len(a) < len(b) else (b, a, True) @@ -188,8 +187,7 @@ def add_text(self, delta_text: str) -> DeltaMessage | None: class Olmo3ReasoningParser(ReasoningParser): - """ - Reasoning parser for Olmo 3 model + """Reasoning parser for Olmo 3 model. Olmo3ReasoningParser @@ -285,8 +283,8 @@ def extract_reasoning( Returns: tuple[Optional[str], Optional[str]]: Tuple pair containing the reasoning content and non-reasoning content. - """ + """ re_match = self.reasoning_regex.match(model_output) if re_match: reasoning = re_match.group("reasoning") or None @@ -305,8 +303,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """Extract content using token ID sequence state machine""" - + """Extract content using token ID sequence state machine.""" delta_message = self.buffer.add_text(delta_text) if delta_message is None and self.buffer.think_end in self.buffer.buffer: # this is a bit hacky, but, because of how the buffer is diff --git a/vllm/reasoning/poolside_v1_reasoning_parser.py b/vllm/reasoning/poolside_v1_reasoning_parser.py index 30031d8513a9..e854afe4bc46 100644 --- a/vllm/reasoning/poolside_v1_reasoning_parser.py +++ b/vllm/reasoning/poolside_v1_reasoning_parser.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Laguna reasoning parser. +"""Laguna reasoning parser. ``DeepSeekV3ReasoningParser.is_reasoning_end`` walks the entire token sequence backwards and returns ``True`` on the first ```` it diff --git a/vllm/reasoning/step3_reasoning_parser.py b/vllm/reasoning/step3_reasoning_parser.py index d2a575c74ee7..16ab1cf9f400 100644 --- a/vllm/reasoning/step3_reasoning_parser.py +++ b/vllm/reasoning/step3_reasoning_parser.py @@ -17,8 +17,7 @@ class Step3ReasoningParser(ReasoningParser): - """ - Reasoning parser for Step3 model. + """Reasoning parser for Step3 model. The Step3 model uses token to denote the end of reasoning text. This parser extracts all content before as reasoning content. @@ -62,8 +61,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - """ - Extract reasoning content from a delta message. + """Extract reasoning content from a delta message. Handles streaming output where previous + delta = current. Uses token IDs for faster processing. For text "abcxyz": diff --git a/vllm/reasoning/step3p5_reasoning_parser.py b/vllm/reasoning/step3p5_reasoning_parser.py index cc7c227e6d7e..5969ca5228cc 100644 --- a/vllm/reasoning/step3p5_reasoning_parser.py +++ b/vllm/reasoning/step3p5_reasoning_parser.py @@ -14,8 +14,7 @@ class Step3p5ReasoningParser(BaseThinkingReasoningParser): - """ - Reasoning parser for Step3p5 model. + """Reasoning parser for Step3p5 model. Step3p5 uses the ... format, but it tends to emit an extra newline immediately before and/or after the token. This parser diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 4aca17fa4334..36224a62b3c1 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -340,8 +340,7 @@ def _join_mm_warmup(self) -> None: self._mm_warmup_future = None def warmup(self, chat_params: ChatParams) -> None: - """ - Warm up this renderer to avoid first-request latency. + """Warm up this renderer to avoid first-request latency. For chat requests: - Jinja2 template compilation @@ -378,7 +377,8 @@ def warmup(self, chat_params: ChatParams) -> None: async def clear_mm_cache_async(self) -> None: """Serialize clear_mm_cache through the multimodal executor to avoid - races with concurrent process_inputs on the mm_processor_cache.""" + races with concurrent process_inputs on the mm_processor_cache. + """ await self._clear_mm_cache_async() def shutdown(self) -> None: @@ -408,8 +408,7 @@ def get_eos_token_id(self) -> int | None: return self.tokenizer.eos_token_id def get_dec_start_token_id(self) -> int: - """ - Obtain the decoder start token id employed by an encoder/decoder model, + """Obtain the decoder start token id employed by an encoder/decoder model, raising an error if it is not available. """ dec_start_token_id = getattr( diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 268a9aa46367..26df7f20934c 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -103,7 +103,8 @@ def _ensure_prompt_embeds_placeholder_token(tokenizer: HfTokenizer) -> int: """Register `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` as a special token and return - its token ID.""" + its token ID. + """ cached = _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE.get(tokenizer) if cached is not None: return cached @@ -192,7 +193,8 @@ def _build_mixed_prompt_embeds( positions: list[tuple[int, int]], ) -> tuple[torch.Tensor, list[bool]]: """Build the full-length `prompt_embeds` tensor and the `is_token_ids` - mask aligned to `token_ids`.""" + mask aligned to `token_ids`. + """ total_len = len(token_ids) hidden_size = prompt_embeds_tensors[0].shape[1] dtype = prompt_embeds_tensors[0].dtype @@ -900,6 +902,7 @@ def rebuild_mm_uuids_from_mm_data( Returns: Updated UUIDs dictionary with chunk UUIDs + """ vision_chunks = mm_data.get("vision_chunk") if vision_chunks is None: @@ -932,6 +935,7 @@ def build_video_prompts_from_mm_data( Returns: List of video prompts, one per video. + """ vision_chunks = mm_data.get("vision_chunk") if vision_chunks is None: diff --git a/vllm/renderers/inputs/preprocess.py b/vllm/renderers/inputs/preprocess.py index dcf80ac73899..82a49f5638cb 100644 --- a/vllm/renderers/inputs/preprocess.py +++ b/vllm/renderers/inputs/preprocess.py @@ -1,6 +1,4 @@ -""" -Schemas and utilities for preprocessing inputs. -""" +"""Schemas and utilities for preprocessing inputs.""" # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project @@ -90,8 +88,7 @@ def conversation_to_seq( class EncoderDecoderDictPrompt(TypedDict): - """ - A [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt] + """A [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt] that has been standardized into a dictionary. """ @@ -130,9 +127,7 @@ def _validate_prompt_dict(prompt: Mapping[str, object]) -> None: def parse_dec_only_prompt(prompt: PromptType | object) -> DecoderOnlyDictPrompt: - """ - Parse a prompt for a decoder-only model and normalize it to a dictionary. - """ + """Parse a prompt for a decoder-only model and normalize it to a dictionary.""" if isinstance(prompt, str): return TextPrompt(prompt=prompt) @@ -216,9 +211,7 @@ def _parse_dec_prompt(prompt: PromptType | object) -> DecoderDictPrompt: def parse_enc_dec_prompt(prompt: PromptType | object) -> EncoderDecoderDictPrompt: - """ - Parse a prompt for an encoder-decoder model and normalize it to a dictionary. - """ + """Parse a prompt for an encoder-decoder model and normalize it to a dictionary.""" if isinstance(prompt, dict) and "encoder_prompt" in prompt: enc_prompt = prompt["encoder_prompt"] # type: ignore[typeddict-item] dec_prompt = prompt["decoder_prompt"] # type: ignore[typeddict-item] diff --git a/vllm/renderers/inputs/tokenize.py b/vllm/renderers/inputs/tokenize.py index 4168e201203e..bf85051c379d 100644 --- a/vllm/renderers/inputs/tokenize.py +++ b/vllm/renderers/inputs/tokenize.py @@ -1,6 +1,4 @@ -""" -Schemas and utilities for tokenization inputs. -""" +"""Schemas and utilities for tokenization inputs.""" # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project @@ -30,8 +28,7 @@ class EncoderDecoderTokPrompt(TypedDict): - """ - A + """A [`EncoderDecoderDictPrompt`][vllm.renderers.inputs.preprocess.EncoderDecoderDictPrompt] that has been tokenized. """ diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index ca70bb6ace44..09e0451ccc70 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -233,6 +233,7 @@ def _detokenize_delta( Returns: (new_text, updated_state) — the delta text for this chunk and the state to pass to the next call. + """ prev_tokens = list(state.prev_tokens) prefix_offset = state.prefix_offset @@ -296,6 +297,7 @@ async def derender_chat_stream( Returns: (chunk, updated_state) — the derendered SSE chunk and the state the client must pass to the next call. + """ if state is None: state = DerenderStreamState() @@ -461,6 +463,7 @@ async def derender_completion_stream( Returns: (chunk, updated_state) — the derendered chunk and updated state. + """ if state is None: state = DerenderStreamState() @@ -613,7 +616,8 @@ def _convert_chat_logprobs_to_completion_logprobs( logprobs: ChatCompletionLogProbs, ) -> CompletionLogProbs: """Convert ChatCompletionLogProbs (per-token objects) to CompletionLogProbs - (parallel flat lists) as required by the /v1/completions response schema.""" + (parallel flat lists) as required by the /v1/completions response schema. + """ if logprobs.content is None: return CompletionLogProbs() diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index b651a060955f..86470f8db1a4 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -425,8 +425,7 @@ def apply_pre_tokenization( tokenizer: TokenizerLike | None, prompt: TextPrompt, ) -> TextPrompt: - """ - Ensure that the prompt meets the requirements set out by this config. + """Ensure that the prompt meets the requirements set out by this config. If that is not possible, raise a `VLLMValidationError`. This method is run before tokenization occurs. @@ -538,8 +537,7 @@ def apply_post_tokenization( tokenizer: TokenizerLike | None, prompt: TokensPrompt | EmbedsPrompt, ) -> TokensPrompt | EmbedsPrompt: - """ - Ensure that the prompt meets the requirements set out by this config. + """Ensure that the prompt meets the requirements set out by this config. If that is not possible, raise a `VLLMValidationError`. This method is run after tokenization occurs. diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index c0ab45ced702..8072892fe70b 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -127,9 +127,7 @@ def __post_init__(self): ) def all_constraints_none(self) -> bool: - """ - Returns True if all structured-output constraint fields are None. - """ + """Returns True if all structured-output constraint fields are None.""" return all( getattr(self, field) is None for field in ( @@ -143,9 +141,7 @@ def all_constraints_none(self) -> bool: ) def all_non_structural_tag_constraints_none(self) -> bool: - """ - Returns True if all structured-output constraint fields are None. - """ + """Returns True if all structured-output constraint fields are None.""" return all( getattr(self, field) is None for field in ( @@ -692,7 +688,7 @@ def update_from_generation_config( generation_config: dict[str, Any], eos_token_id: int | None = None, ) -> None: - """Update if there are non-default values from generation_config""" + """Update if there are non-default values from generation_config.""" if not self.ignore_eos: self._eos_token_id = eos_token_id @@ -804,7 +800,8 @@ def num_logprobs(self) -> int | None: """Number of sample logprobs to return per output token, or `None` if no sample logprobs were requested. Takes `logprob_token_ids` into account: when `logprobs` is unset but `logprob_token_ids` is set, - returns `len(logprob_token_ids)`.""" + returns `len(logprob_token_ids)`. + """ if self.logprobs is not None: return self.logprobs return len(self.logprob_token_ids) if self.logprob_token_ids else None diff --git a/vllm/scalar_type.py b/vllm/scalar_type.py index db52e93465ca..d73827f499df 100644 --- a/vllm/scalar_type.py +++ b/vllm/scalar_type.py @@ -21,8 +21,7 @@ class NanRepr(Enum): # in sync until the inductor fully supports custom C++ classes. @dataclass(frozen=True) class ScalarType: - """ - ScalarType can represent a wide range of floating point and integer + """ScalarType can represent a wide range of floating point and integer types, in particular it can be used to represent sub-byte data types (something that torch.dtype currently does not support). It is also capable of representing types with a bias, i.e.: @@ -135,8 +134,7 @@ def _raw_min(self) -> int | float: @functools.cached_property def id(self) -> int: - """ - Convert the ScalarType to an int which can be passed to pytorch custom + """Convert the ScalarType to an int which can be passed to pytorch custom ops. This layout of the int must be kept in sync with the C++ ScalarType's from_id method. """ @@ -168,56 +166,51 @@ def size_bits(self) -> int: return self.exponent + self.mantissa + int(self.signed) def min(self) -> int | float: - """ - Min representable value for this scalar type. + """Min representable value for this scalar type. (accounting for bias if there is one) """ return self._raw_min() - self.bias def max(self) -> int | float: - """ - Max representable value for this scalar type. + """Max representable value for this scalar type. (accounting for bias if there is one) """ return self._raw_max() - self.bias def is_signed(self) -> bool: - """ - If the type is signed (i.e. has a sign bit), same as `signed` + """If the type is signed (i.e. has a sign bit), same as `signed` added for consistency with: https://pytorch.org/docs/stable/generated/torch.Tensor.is_signed.html """ return self.signed def is_floating_point(self) -> bool: - "If the type is a floating point type" + """If the type is a floating point type.""" return self.exponent != 0 def is_integer(self) -> bool: - "If the type is an integer type" + """If the type is an integer type.""" return self.exponent == 0 def has_bias(self) -> bool: - "If the type has a non-zero bias" + """If the type has a non-zero bias.""" return self.bias != 0 def has_infs(self) -> bool: - "If the type is floating point and supports infinity" + """If the type is floating point and supports infinity.""" return not self._finite_values_only def has_nans(self) -> bool: return self.nan_repr != NanRepr.NONE.value def is_ieee_754(self) -> bool: - """ - If the type is a floating point type that follows IEEE 754 + """If the type is a floating point type that follows IEEE 754 conventions """ return self.nan_repr == NanRepr.IEEE_754.value and not self._finite_values_only def __str__(self) -> str: - """ - naming generally follows: https://github.com/jax-ml/ml_dtypes + """Naming generally follows: https://github.com/jax-ml/ml_dtypes for floating point types (leading f) the scheme is: `float_em[flags]` flags: @@ -265,7 +258,7 @@ def __len__(self) -> int: @classmethod def int_(cls, size_bits: int, bias: int | None) -> "ScalarType": - "Create a signed integer scalar type (size_bits includes sign-bit)." + """Create a signed integer scalar type (size_bits includes sign-bit).""" ret = cls(0, size_bits - 1, True, bias if bias else 0) ret.id # noqa B018: make sure the id is cached return ret @@ -279,8 +272,7 @@ def uint(cls, size_bits: int, bias: int | None) -> "ScalarType": @classmethod def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType": - """ - Create a standard floating point type + """Create a standard floating point type (i.e. follows IEEE 754 conventions). """ assert mantissa > 0 and exponent > 0 @@ -292,8 +284,7 @@ def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType": def float_( cls, exponent: int, mantissa: int, finite_values_only: bool, nan_repr: NanRepr ) -> "ScalarType": - """ - Create a non-standard floating point type + """Create a non-standard floating point type (i.e. does not follow IEEE 754 conventions). """ assert mantissa > 0 and exponent > 0 diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py index 60fc53350702..8f58ddcf8c7f 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -15,6 +15,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv, next_power_of_2 + @triton.heuristics( { "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 548f600812e2..b5c3a92b65e7 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -1349,6 +1349,7 @@ def grid(meta): ) return y + def _chunk_kda_fwd_with_cumulative_g( q: torch.Tensor, k: torch.Tensor, diff --git a/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py index 79f5abf031fb..997df5a173b1 100644 --- a/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py +++ b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py @@ -335,9 +335,7 @@ def compile(self, compile_key: CompileKey) -> None: assert warmup is not None variant = compile_key.input_variant - def pointer( - name: str, dtype: torch.dtype | None - ) -> TritonWarmupTensor | None: + def pointer(name: str, dtype: torch.dtype | None) -> TritonWarmupTensor | None: return None if dtype is None else variant.pointer(name, dtype) warmup( diff --git a/vllm/third_party/pynvml.py b/vllm/third_party/pynvml.py index 6aabbc217dd0..3bb53e8d9358 100644 --- a/vllm/third_party/pynvml.py +++ b/vllm/third_party/pynvml.py @@ -45,52 +45,54 @@ ## C Type mappings ## ## Enums _nvmlEnableState_t = c_uint -NVML_FEATURE_DISABLED = 0 -NVML_FEATURE_ENABLED = 1 +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 _nvmlBrandType_t = c_uint -NVML_BRAND_UNKNOWN = 0 -NVML_BRAND_QUADRO = 1 -NVML_BRAND_TESLA = 2 -NVML_BRAND_NVS = 3 -NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. -NVML_BRAND_GEFORCE = 5 -NVML_BRAND_TITAN = 6 -NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications -NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC -NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server -NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = ( + 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +) +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming -NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. -NVML_BRAND_QUADRO_RTX = 12 -NVML_BRAND_NVIDIA_RTX = 13 -NVML_BRAND_NVIDIA = 14 -NVML_BRAND_GEFORCE_RTX = 15 # Unused -NVML_BRAND_TITAN_RTX = 16 # Unused -NVML_BRAND_COUNT = 17 +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 17 _nvmlTemperatureThresholds_t = c_uint -NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 -NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 -NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 -NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 -NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 -NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 -NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 -NVML_TEMPERATURE_THRESHOLD_COUNT = 8 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 _nvmlTemperatureSensors_t = c_uint -NVML_TEMPERATURE_GPU = 0 -NVML_TEMPERATURE_COUNT = 1 +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 _nvmlComputeMode_t = c_uint -NVML_COMPUTEMODE_DEFAULT = 0 -NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed -NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 -NVML_COMPUTEMODE_COUNT = 4 +NVML_COMPUTEMODE_COUNT = 4 _nvmlMemoryLocation_t = c_uint NVML_MEMORY_LOCATION_L1_CACHE = 0 @@ -127,158 +129,158 @@ _nvmlNvLinkCapability_t = c_uint NVML_NVLINK_CAP_P2P_SUPPORTED = 0 NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 -NVML_NVLINK_CAP_P2P_ATOMICS = 2 -NVML_NVLINK_CAP_SYSMEM_ATOMICS= 3 -NVML_NVLINK_CAP_SLI_BRIDGE = 4 -NVML_NVLINK_CAP_VALID = 5 -NVML_NVLINK_CAP_COUNT = 6 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 _nvmlNvLinkUtilizationCountPktTypes_t = c_uint -NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 -NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 -NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 -NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 -NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 -NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 -NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 -NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF _nvmlNvLinkUtilizationCountUnits_t = c_uint -NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 -NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 -NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 -NVML_NVLINK_COUNTER_UNIT_COUNT = 4 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 _nvmlNvLinkDeviceType_t = c_uint -NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 -NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 -NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF # These are deprecated, instead use _nvmlMemoryErrorType_t _nvmlEccBitType_t = c_uint -NVML_SINGLE_BIT_ECC = 0 -NVML_DOUBLE_BIT_ECC = 1 +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 NVML_ECC_ERROR_TYPE_COUNT = 2 _nvmlEccCounterType_t = c_uint -NVML_VOLATILE_ECC = 0 -NVML_AGGREGATE_ECC = 1 +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 NVML_ECC_COUNTER_TYPE_COUNT = 2 _nvmlMemoryErrorType_t = c_uint -NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 -NVML_MEMORY_ERROR_TYPE_COUNT = 2 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 _nvmlClockType_t = c_uint -NVML_CLOCK_GRAPHICS = 0 -NVML_CLOCK_SM = 1 -NVML_CLOCK_MEM = 2 -NVML_CLOCK_VIDEO = 3 -NVML_CLOCK_COUNT = 4 +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 _nvmlClockId_t = c_uint -NVML_CLOCK_ID_CURRENT = 0 -NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 -NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 -NVML_CLOCK_ID_COUNT = 4 +NVML_CLOCK_ID_COUNT = 4 _nvmlDriverModel_t = c_uint -NVML_DRIVER_WDDM = 0 -NVML_DRIVER_WDM = 1 -NVML_DRIVER_MCDM = 2 +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 NVML_MAX_GPU_PERF_PSTATES = 16 _nvmlPstates_t = c_uint -NVML_PSTATE_0 = 0 -NVML_PSTATE_1 = 1 -NVML_PSTATE_2 = 2 -NVML_PSTATE_3 = 3 -NVML_PSTATE_4 = 4 -NVML_PSTATE_5 = 5 -NVML_PSTATE_6 = 6 -NVML_PSTATE_7 = 7 -NVML_PSTATE_8 = 8 -NVML_PSTATE_9 = 9 -NVML_PSTATE_10 = 10 -NVML_PSTATE_11 = 11 -NVML_PSTATE_12 = 12 -NVML_PSTATE_13 = 13 -NVML_PSTATE_14 = 14 -NVML_PSTATE_15 = 15 -NVML_PSTATE_UNKNOWN = 32 +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 _nvmlInforomObject_t = c_uint -NVML_INFOROM_OEM = 0 -NVML_INFOROM_ECC = 1 -NVML_INFOROM_POWER = 2 -NVML_INFOROM_DEN = 3 -NVML_INFOROM_COUNT = 4 +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 _nvmlReturn_t = c_uint -NVML_SUCCESS = 0 -NVML_ERROR_UNINITIALIZED = 1 -NVML_ERROR_INVALID_ARGUMENT = 2 -NVML_ERROR_NOT_SUPPORTED = 3 -NVML_ERROR_NO_PERMISSION = 4 -NVML_ERROR_ALREADY_INITIALIZED = 5 -NVML_ERROR_NOT_FOUND = 6 -NVML_ERROR_INSUFFICIENT_SIZE = 7 -NVML_ERROR_INSUFFICIENT_POWER = 8 -NVML_ERROR_DRIVER_NOT_LOADED = 9 -NVML_ERROR_TIMEOUT = 10 -NVML_ERROR_IRQ_ISSUE = 11 -NVML_ERROR_LIBRARY_NOT_FOUND = 12 -NVML_ERROR_FUNCTION_NOT_FOUND = 13 -NVML_ERROR_CORRUPTED_INFOROM = 14 -NVML_ERROR_GPU_IS_LOST = 15 -NVML_ERROR_RESET_REQUIRED = 16 -NVML_ERROR_OPERATING_SYSTEM = 17 -NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 -NVML_ERROR_IN_USE = 19 -NVML_ERROR_MEMORY = 20 -NVML_ERROR_NO_DATA = 21 -NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 -NVML_ERROR_INSUFFICIENT_RESOURCES = 23 -NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 -NVML_ERROR_DEPRECATED = 26 -NVML_ERROR_NOT_READY = 27 -NVML_ERROR_GPU_NOT_FOUND = 28 -NVML_ERROR_INVALID_STATE = 29 -NVML_ERROR_UNKNOWN = 999 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_UNKNOWN = 999 _nvmlFanState_t = c_uint -NVML_FAN_NORMAL = 0 -NVML_FAN_FAILED = 1 +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 _nvmlFanControlPolicy_t = c_uint NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 -NVML_FAN_POLICY_MANUAL = 1 +NVML_FAN_POLICY_MANUAL = 1 _nvmlLedColor_t = c_uint -NVML_LED_COLOR_GREEN = 0 -NVML_LED_COLOR_AMBER = 1 +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 _nvmlGpuOperationMode_t = c_uint -NVML_GOM_ALL_ON = 0 -NVML_GOM_COMPUTE = 1 -NVML_GOM_LOW_DP = 2 +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 _nvmlPageRetirementCause_t = c_uint NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 -NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 -NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 _nvmlRestrictedAPI_t = c_uint -NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 -NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 -NVML_RESTRICTED_API_COUNT = 2 +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 _nvmlBridgeChipType_t = c_uint NVML_BRIDGE_CHIP_PLX = 0 @@ -365,9 +367,9 @@ NVML_TOPOLOGY_SYSTEM = 50 _nvmlGpuP2PCapsIndex_t = c_uint -NVML_P2P_CAPS_INDEX_READ = 0, +NVML_P2P_CAPS_INDEX_READ = (0,) NVML_P2P_CAPS_INDEX_WRITE = 1 -NVML_P2P_CAPS_INDEX_NVLINK =2 +NVML_P2P_CAPS_INDEX_NVLINK = 2 NVML_P2P_CAPS_INDEX_ATOMICS = 3 # # NVML_P2P_CAPS_INDEX_PROP is deprecated. @@ -378,39 +380,39 @@ NVML_P2P_CAPS_INDEX_UNKNOWN = 5 _nvmlGpuP2PStatus_t = c_uint -NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_OK = 0 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 -NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3 -NVML_P2P_STATUS_DISABLED_BY_REGKEY =4 -NVML_P2P_STATUS_NOT_SUPPORTED =5 -NVML_P2P_STATUS_UNKNOWN =6 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4 +NVML_P2P_STATUS_NOT_SUPPORTED = 5 +NVML_P2P_STATUS_UNKNOWN = 6 _nvmlDeviceArchitecture_t = c_uint -NVML_DEVICE_ARCH_KEPLER = 2 -NVML_DEVICE_ARCH_MAXWELL = 3 -NVML_DEVICE_ARCH_PASCAL = 4 -NVML_DEVICE_ARCH_VOLTA = 5 -NVML_DEVICE_ARCH_TURING = 6 -NVML_DEVICE_ARCH_AMPERE = 7 -NVML_DEVICE_ARCH_ADA = 8 -NVML_DEVICE_ARCH_HOPPER = 9 -NVML_DEVICE_ARCH_BLACKWELL = 10 -NVML_DEVICE_ARCH_T23X = 11 -NVML_DEVICE_ARCH_UNKNOWN = 0xffffffff +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_T23X = 11 +NVML_DEVICE_ARCH_UNKNOWN = 0xFFFFFFFF # PCI bus Types _nvmlBusType_t = c_uint NVML_BUS_TYPE_UNKNOWN = 0 -NVML_BUS_TYPE_PCI = 1 -NVML_BUS_TYPE_PCIE = 2 -NVML_BUS_TYPE_FPCI = 3 -NVML_BUS_TYPE_AGP = 4 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 _nvmlPowerSource_t = c_uint -NVML_POWER_SOURCE_AC = 0x00000000 -NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 _nvmlAdaptiveClockInfoStatus_t = c_uint @@ -418,196 +420,290 @@ NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 _nvmlClockLimitId_t = c_uint -NVML_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00 -NVML_CLOCK_LIMIT_ID_TDP = 0xffffff01 -NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xffffff02 +NVML_CLOCK_LIMIT_ID_RANGE_START = 0xFFFFFF00 +NVML_CLOCK_LIMIT_ID_TDP = 0xFFFFFF01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xFFFFFF02 _nvmlPcieLinkMaxSpeed_t = c_uint -NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 -NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 -NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 -NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 _nvmlPcieAtomicsCapability_t = c_uint -NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 -NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 -NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 -NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 -NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 -NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 -NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 -NVML_PCIE_ATOMICS_OPS_MAX = 7 +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 _nvmlAffinityScope_t = c_uint -NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_NODE = 0 NVML_AFFINITY_SCOPE_SOCKET = 1 _nvmlDeviceGpuRecoveryAction_t = c_uint -NVML_GPU_RECOVERY_ACTION_NONE = 0 -NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 -NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 # C preprocessor defined values -nvmlFlagDefault = 0 -nvmlFlagForce = 1 -NVML_INIT_FLAG_NO_GPUS = 1 -NVML_INIT_FLAG_NO_ATTACH = 2 +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 -NVML_MAX_GPC_COUNT = 32 +NVML_MAX_GPC_COUNT = 32 # buffer size -NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 -NVML_DEVICE_UUID_BUFFER_SIZE = 80 -NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 -NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 -NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 -NVML_DEVICE_NAME_BUFFER_SIZE = 64 -NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 -NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 -NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 -NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 -NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 -NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 -NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 -NVML_GRID_LICENSE_BUFFER_SIZE = 128 -NVML_VGPU_NAME_BUFFER_SIZE = 64 -NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 -NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 -NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 -NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH -NVML_PERF_MODES_BUFFER_SIZE = 2048 +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = ( + 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +) +NVML_PERF_MODES_BUFFER_SIZE = 2048 # Format strings -NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" -NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) -''' +""" Field Identifiers. All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. -''' -NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive -NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive - -#ECC Count Totals -NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors -#Individual ECC locations -NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_DEV = 22 # Device memory single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_DEV = 23 # Device memory double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_REG = 24 # Register File single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_REG = 25 # Register File double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_TEX = 26 # Texture memory single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_TEX = 27 # Texture memory double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors +""" +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive + +# ECC Count Totals +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors +# Individual ECC locations +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_DEV = ( + 22 # Device memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_DEV = ( + 23 # Device memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_REG = ( + 24 # Register File single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_REG = ( + 25 # Register File double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_TEX = ( + 26 # Texture memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_TEX = ( + 27 # Texture memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors # Page Retirement -NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors -NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors -NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. +NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors +NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors +NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. # NvLink Flit Error Counters -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = 32 # NVLink flow control CRC Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = 33 # NVLink flow control CRC Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = 34 # NVLink flow control CRC Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = 35 # NVLink flow control CRC Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = 36 # NVLink flow control CRC Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = 37 # NVLink flow control CRC Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = ( + 32 # NVLink flow control CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = ( + 33 # NVLink flow control CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = ( + 34 # NVLink flow control CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = ( + 35 # NVLink flow control CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = ( + 36 # NVLink flow control CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = ( + 37 # NVLink flow control CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = ( + 38 # NVLink flow control CRC Error Counter total for all Lanes +) # NvLink CRC Data Error Counters -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = 39 # NVLink data CRC Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = 40 # NVLink data CRC Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = 41 # NVLink data CRC Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = 42 # NVLink data CRC Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = 43 # NVLink data CRC Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = 44 # NVLink data CRC Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = ( + 39 # NVLink data CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = ( + 40 # NVLink data CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = ( + 41 # NVLink data CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = ( + 42 # NVLink data CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = ( + 43 # NVLink data CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = ( + 44 # NVLink data CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = ( + 45 # NvLink data CRC Error Counter total for all Lanes +) # NvLink Replay Error Counters -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = 52 # NVLink Replay Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = ( + 52 # NVLink Replay Error Counter total for all Lanes +) # NvLink Recovery Error Counters -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = 53 # NVLink Recovery Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = 54 # NVLink Recovery Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = 55 # NVLink Recovery Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = 56 # NVLink Recovery Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = 57 # NVLink Recovery Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = 58 # NVLink Recovery Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = ( + 53 # NVLink Recovery Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = ( + 54 # NVLink Recovery Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = ( + 55 # NVLink Recovery Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = ( + 56 # NVLink Recovery Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = ( + 57 # NVLink Recovery Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = ( + 58 # NVLink Recovery Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = ( + 59 # NVLink Recovery Error Counter total for all Lanes +) # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = ( + 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = ( + 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = ( + 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = ( + 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = ( + 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = ( + 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = ( + 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes +) # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = ( + 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = ( + 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = ( + 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = ( + 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = ( + 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = ( + 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = ( + 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes +) # Perf Policy Counters -NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy -NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy -NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy -NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit -NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = 78 # Perf Policy Counter for Low GPU Utilization Policy -NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy -NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = 80 # Perf Policy Counter for Total App Clock Policy -NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81 # Perf Policy Counter for Total Base Clocks Policy +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = ( + 78 # Perf Policy Counter for Low GPU Utilization Policy +) +NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = ( + 80 # Perf Policy Counter for Total App Clock Policy +) +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = ( + 81 # Perf Policy Counter for Total Base Clocks Policy +) # Memory temperatures -NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device # Energy Counter -NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = ( + 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded +) # NVLink Speed -NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 # NVLink Link Count @@ -622,307 +718,451 @@ NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 # NvLink Flit Error Counters -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = 96 # NVLink flow control CRC Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = 97 # NVLink flow control CRC Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = 98 # NVLink flow control CRC Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = 99 # NVLink flow control CRC Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = 100 # NVLink flow control CRC Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = 101 # NVLink flow control CRC Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = ( + 96 # NVLink flow control CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = ( + 97 # NVLink flow control CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = ( + 98 # NVLink flow control CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = ( + 99 # NVLink flow control CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = ( + 100 # NVLink flow control CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = ( + 101 # NVLink flow control CRC Error Counter for Lane 11 +) # NvLink CRC Data Error Counters -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = 102 # NVLink data CRC Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = 103 # NVLink data CRC Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = 104 # NVLink data CRC Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = 105 # NVLink data CRC Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = 106 # NVLink data CRC Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = 107 # NVLink data CRC Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = ( + 102 # NVLink data CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = ( + 103 # NVLink data CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = ( + 104 # NVLink data CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = ( + 105 # NVLink data CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = ( + 106 # NVLink data CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = ( + 107 # NVLink data CRC Error Counter for Lane 11 +) # NvLink Replay Error Counters -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = 112 # NVLink Replay Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = 113 # NVLink Replay Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = ( + 112 # NVLink Replay Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = ( + 113 # NVLink Replay Error Counter for Lane 11 +) # NvLink Recovery Error Counters -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = 114 # NVLink Recovery Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = 115 # NVLink Recovery Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = 116 # NVLink Recovery Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = 117 # NVLink Recovery Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = 118 # NVLink Recovery Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = 119 # NVLink Recovery Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = ( + 114 # NVLink Recovery Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = ( + 115 # NVLink Recovery Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = ( + 116 # NVLink Recovery Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = ( + 117 # NVLink Recovery Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = ( + 118 # NVLink Recovery Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = ( + 119 # NVLink Recovery Error Counter for Lane 11 +) # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = ( + 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = ( + 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = ( + 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = ( + 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = ( + 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = ( + 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 +) # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = ( + 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = ( + 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = ( + 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = ( + 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = ( + 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = ( + 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 +) # NVLink Speed -NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 # NVLink Throughput Counters -NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB # Row Remapper -NVML_FI_DEV_REMAPPED_COR = 142 -NVML_FI_DEV_REMAPPED_UNC = 143 -NVML_FI_DEV_REMAPPED_PENDING = 144 -NVML_FI_DEV_REMAPPED_FAILURE = 145 +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 -#Remote device NVLink ID +# Remote device NVLink ID NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 # Number of NVLinks connected to NVSwitch NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 # NvLink ECC Data Error Counters -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = 148 #< NVLink data ECC Error Counter for Link 0 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = 149 #< NVLink data ECC Error Counter for Link 1 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = 150 #< NVLink data ECC Error Counter for Link 2 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = 151 #< NVLink data ECC Error Counter for Link 3 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = 152 #< NVLink data ECC Error Counter for Link 4 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = 153 #< NVLink data ECC Error Counter for Link 5 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = 154 #< NVLink data ECC Error Counter for Link 6 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = 155 #< NVLink data ECC Error Counter for Link 7 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = 156 #< NVLink data ECC Error Counter for Link 8 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = 157 #< NVLink data ECC Error Counter for Link 9 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = 158 #< NVLink data ECC Error Counter for Link 10 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = 159 #< NVLink data ECC Error Counter for Link 11 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links - -NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 -NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 -NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 -NVML_FI_DEV_NVLINK_GET_SPEED = 164 -NVML_FI_DEV_NVLINK_GET_STATE = 165 -NVML_FI_DEV_NVLINK_GET_VERSION = 166 - -NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 - -NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 - -NVML_FI_DEV_C2C_LINK_COUNT = 170 -NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 -NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 - -NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 -NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 -NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 -NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 -NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 -NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 -NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 -NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 -NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 -NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 -NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 - -NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 - -NVML_FI_DEV_POWER_AVERAGE = 185 -NVML_FI_DEV_POWER_INSTANT = 186 -NVML_FI_DEV_POWER_MIN_LIMIT = 187 -NVML_FI_DEV_POWER_MAX_LIMIT = 188 -NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 -NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 -NVML_FI_DEV_ENERGY = 191 -NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 - -NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 -NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 -NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 -NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 - -NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 -NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 - -NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 - -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 - -NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 -NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 -NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 -NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 -NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 -NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 -NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 -NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 -NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 -NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 -NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = ( + 148 # < NVLink data ECC Error Counter for Link 0 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = ( + 149 # < NVLink data ECC Error Counter for Link 1 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = ( + 150 # < NVLink data ECC Error Counter for Link 2 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = ( + 151 # < NVLink data ECC Error Counter for Link 3 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = ( + 152 # < NVLink data ECC Error Counter for Link 4 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = ( + 153 # < NVLink data ECC Error Counter for Link 5 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = ( + 154 # < NVLink data ECC Error Counter for Link 6 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = ( + 155 # < NVLink data ECC Error Counter for Link 7 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = ( + 156 # < NVLink data ECC Error Counter for Link 8 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = ( + 157 # < NVLink data ECC Error Counter for Link 9 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = ( + 158 # < NVLink data ECC Error Counter for Link 10 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = ( + 159 # < NVLink data ECC Error Counter for Link 11 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = ( + 160 # < NvLink data ECC Error Counter total for all Links +) + +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 + +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 + +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 + +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 + +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 + +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 + +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 + +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 + +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 + +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 + +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 -NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 -NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 - -NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 -NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 -NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 -NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 - -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 - -NVML_FI_DEV_RESET_STATUS = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead -NVML_FI_DEV_DRAIN_AND_RESET_STATUS = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead -NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 -NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 -NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 - -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 -NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) -NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level -NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) -NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value -NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value -NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value -NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value -NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = 258 # HW Circuitry % lifetime remaining -NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles -NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = 261 # Ramp up rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = 262 # Ramp down rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = 263 # Ramp down hysteresis value in ms for a given profile -NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = 265 # % TMP floor for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = 266 # Ramp up rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = 267 # Ramp down rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = 268 # Ramp down hysteresis value in ms for a given profile - -NVML_FI_MAX = 269 # One greater than the largest field ID defined above +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 + +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = ( + 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +) +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 + +NVML_FI_DEV_RESET_STATUS = ( + 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = ( + 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 + +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = ( + 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = ( + 258 # HW Circuitry % lifetime remaining +) +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = ( + 261 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = ( + 262 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = ( + 263 # Ramp down hysteresis value in ms for a given profile +) +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = ( + 265 # % TMP floor for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = ( + 266 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = ( + 267 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = ( + 268 # Ramp down hysteresis value in ms for a given profile +) + +NVML_FI_MAX = 269 # One greater than the largest field ID defined above # NVML_FI_DEV_NVLINK_GET_STATE state enums NVML_NVLINK_STATE_INACTIVE = 0x0 -NVML_NVLINK_STATE_ACTIVE = 0x1 -NVML_NVLINK_STATE_SLEEP = 0x2 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 -NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS -NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = ( + 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = ( + 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) ## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode -NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU -NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1 # Device is associated with GPU-Passthorugh -NVML_GPU_VIRTUALIZATION_MODE_VGPU = 2 # Device is associated with vGPU inside virtual machine. -NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = 3 # Device is associated with VGX hypervisor in vGPU mode -NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = 4 # Device is associated with VGX hypervisor in vSGA mode +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = ( + 1 # Device is associated with GPU-Passthorugh +) +NVML_GPU_VIRTUALIZATION_MODE_VGPU = ( + 2 # Device is associated with vGPU inside virtual machine. +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = ( + 3 # Device is associated with VGX hypervisor in vGPU mode +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = ( + 4 # Device is associated with VGX hypervisor in vSGA mode +) ## Lib loading ## nvmlLib = None libLoadLock = threading.Lock() -_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown +_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown ## vGPU Management -_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuTypeId_t = c_uint _nvmlVgpuInstance_t = c_uint _nvmlVgpuVmIdType_t = c_uint -NVML_VGPU_VM_ID_DOMAIN_ID = 0 -NVML_VGPU_VM_ID_UUID = 1 +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 _nvmlGridLicenseFeatureCode_t = c_uint -NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 -NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 -NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 -NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. -NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 -NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = ( + 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +) +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 _nvmlGridLicenseExpiryStatus_t = c_uint8 -NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = 0, # Expiry information not available -NVML_GRID_LICENSE_EXPIRY_INVALID = 1, # Invalid expiry or error fetching expiry -NVML_GRID_LICENSE_EXPIRY_VALID = 2, # Valid expiry -NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = 3, # Expiry not applicable -NVML_GRID_LICENSE_EXPIRY_PERMANENT = 4, # Permanent expiry +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = (0,) # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = (1,) # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = (2,) # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = (3,) # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = (4,) # Permanent expiry _nvmlVgpuCapability_t = c_uint -NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported -NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported -NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM -NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type -NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size -NVML_VGPU_CAP_COUNT = 5 +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = ( + 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +) +NVML_VGPU_CAP_EXCLUSIVE_TYPE = ( + 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +) +NVML_VGPU_CAP_EXCLUSIVE_SIZE = ( + 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +) +NVML_VGPU_CAP_COUNT = 5 _nvmlVgpuDriverCapability_t = c_uint -NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = 0 # Supports mixing of different vGPU profiles within one guest VM -NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM -NVML_VGPU_DRIVER_CAP_COUNT = 2 +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = ( + 0 # Supports mixing of different vGPU profiles within one guest VM +) +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 _nvmlDeviceVgpuCapability_t = c_uint -NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations -NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types -NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes -NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second -NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second -NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming -NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = 6 # Set/Get support of mini-quarter vGPU profiles -NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = 7 # Set/Get support for compute media engine vGPU profiles -NVML_DEVICE_VGPU_CAP_WARM_UPDATE = 8 # Query whether the GPU supports FSR and warm update -NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes -NVML_DEVICE_VGPU_CAP_COUNT = 10 +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = ( + 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +) +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = ( + 6 # Set/Get support of mini-quarter vGPU profiles +) +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = ( + 7 # Set/Get support for compute media engine vGPU profiles +) +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = ( + 8 # Query whether the GPU supports FSR and warm update +) +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_COUNT = 10 _nvmlVgpuGuestInfoState_t = c_uint NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 -NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 _nvmlVgpuVmCompatibility_t = c_uint -NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 -NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 -NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 -NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 -NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 _nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint -NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 -NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 -NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 -NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 -NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 _nvmlHostVgpuMode_t = c_uint -NVML_HOST_VGPU_MODE_NON_SRIOV = 0 -NVML_HOST_VGPU_MODE_SRIOV = 1 +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 _nvmlConfComputeGpusReadyState_t = c_uint NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 @@ -945,11 +1185,11 @@ NVML_CC_SYSTEM_MULTIGPU_NONE = 0 NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 - + NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 - + _nvmlConfComputeCcFeature_t = c_uint NVML_CC_SYSTEM_FEATURE_DISABLED = 0 NVML_CC_SYSTEM_FEATURE_ENABLED = 1 @@ -961,62 +1201,71 @@ # GSP firmware NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + class NVMLLibraryMismatchError(Exception): pass + ## Error Checking ## class NVMLError(Exception): _valClassMapping = dict() # List of currently known error codes _errcode_to_string = { - NVML_ERROR_UNINITIALIZED: "Uninitialized", - NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", - NVML_ERROR_NOT_SUPPORTED: "Not Supported", - NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", - NVML_ERROR_NOT_FOUND: "Not Found", - NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", - NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", - NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", - NVML_ERROR_TIMEOUT: "Timeout", - NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", - NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", - NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", - NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", - NVML_ERROR_GPU_IS_LOST: "GPU is lost", - NVML_ERROR_RESET_REQUIRED: "GPU requires restart", - NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", - NVML_ERROR_MEMORY: "Insufficient Memory", - NVML_ERROR_UNKNOWN: "Unknown Error", - } + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + def __new__(typ, value): - ''' + """ Maps value to a proper subclass of NVMLError. See _extractNVMLErrorsAsClasses function for more details - ''' + """ if typ == NVMLError: typ = NVMLError._valClassMapping.get(value, typ) obj = Exception.__new__(typ) obj.value = value return obj + def __str__(self): try: if self.value not in NVMLError._errcode_to_string: - NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value)) + NVMLError._errcode_to_string[self.value] = str( + nvmlErrorString(self.value) + ) return NVMLError._errcode_to_string[self.value] except NVMLError: return "NVML Error with code %d" % self.value + def __eq__(self, other): return self.value == other.value + def nvmlExceptionClass(nvmlErrorCode): if nvmlErrorCode not in NVMLError._valClassMapping: - raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode) + raise ValueError("nvmlErrorCode %s is not valid" % nvmlErrorCode) return NVMLError._valClassMapping[nvmlErrorCode] + def _extractNVMLErrorsAsClasses(): - ''' + """ Generates a hierarchy of classes on top of NVMLError class. Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate @@ -1024,31 +1273,44 @@ def _extractNVMLErrorsAsClasses(): NVMLError is a parent class. Each NVML_ERROR_* gets its own subclass. e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized - ''' + """ this_module = sys.modules[__name__] nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")] for err_name in nvmlErrorsNames: # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized - class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "") + class_name = "NVMLError_" + string.capwords( + err_name.replace("NVML_ERROR_", ""), "_" + ).replace("_", "") err_val = getattr(this_module, err_name) + def gen_new(val): def new(typ): obj = NVMLError.__new__(typ, val) return obj + return new - new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)}) + + new_error_class = type(class_name, (NVMLError,), {"__new__": gen_new(err_val)}) new_error_class.__module__ = __name__ setattr(this_module, class_name, new_error_class) NVMLError._valClassMapping[err_val] = new_error_class + + _extractNVMLErrorsAsClasses() + def _nvmlCheckReturn(ret): - if (ret != NVML_SUCCESS): + if ret != NVML_SUCCESS: raise NVMLError(ret) return ret + ## Function access ## -_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking +_nvmlGetFunctionPointer_cache = ( + dict() +) # function pointers are cached to prevent unnecessary libLoadLock locking + + def _nvmlGetFunctionPointer(name): global nvmlLib @@ -1058,7 +1320,7 @@ def _nvmlGetFunctionPointer(name): libLoadLock.acquire() try: # ensure library was loaded - if (nvmlLib == None): + if nvmlLib == None: raise NVMLError(NVML_ERROR_UNINITIALIZED) try: _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) @@ -1069,6 +1331,7 @@ def _nvmlGetFunctionPointer(name): # lock is always freed libLoadLock.release() + ## Alternative object # Allows the object to be printed # Allows mismatched types to be assigned @@ -1077,9 +1340,11 @@ class nvmlFriendlyObject(object): def __init__(self, dictionary): for x in dictionary: setattr(self, x, dictionary[x]) + def __str__(self): return self.__dict__.__str__() + def nvmlStructToFriendlyObject(struct): d = {} for x in struct._fields_: @@ -1090,6 +1355,7 @@ def nvmlStructToFriendlyObject(struct): obj = nvmlFriendlyObject(d) return obj + # pack the object so it can be passed to the NVML library def nvmlFriendlyObjectToStruct(obj, model): for x in model._fields_: @@ -1102,11 +1368,15 @@ def nvmlFriendlyObjectToStruct(obj, model): setattr(model, key, value) return model + ## Unit structures class struct_c_nvmlUnit_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + class _PrintableStructure(Structure): """ Abstract class that produces nicer __str__ output than ctypes.Structure. @@ -1126,7 +1396,9 @@ class _PrintableStructure(Structure): Exact format of returned str from this class is subject to change in the future. """ + _fmt_ = {} + def __str__(self): result = [] for x in self._fields_: @@ -1138,7 +1410,7 @@ def __str__(self): elif "" in self._fmt_: fmt = self._fmt_[""] result.append(("%s: " + fmt) % (key, value)) - return self.__class__.__name__ + "(" + ", ".join(result) + ")" + return self.__class__.__name__ + "(" + ", ".join(result) + ")" def __getattribute__(self, name): res = super(_PrintableStructure, self).__getattribute__(name) @@ -1159,185 +1431,200 @@ def __setattr__(self, name, value): value = value.encode() super(_PrintableStructure, self).__setattr__(name, value) + class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ - ('name', c_char * 96), - ('id', c_char * 96), - ('serial', c_char * 96), - ('firmwareVersion', c_char * 96), + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), ] + class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): - _fields_ = [ - ('isC2cEnabled', c_uint) - ] + _fields_ = [("isC2cEnabled", c_uint)] + + +nvmlC2cModeInfo_v1 = 0x1000008 -nvmlC2cModeInfo_v1 = 0x1000008; class c_nvmlLedState_t(_PrintableStructure): _fields_ = [ - ('cause', c_char * 256), - ('color', _nvmlLedColor_t), + ("cause", c_char * 256), + ("color", _nvmlLedColor_t), ] + class c_nvmlPSUInfo_t(_PrintableStructure): _fields_ = [ - ('state', c_char * 256), - ('current', c_uint), - ('voltage', c_uint), - ('power', c_uint), + ("state", c_char * 256), + ("current", c_uint), + ("voltage", c_uint), + ("power", c_uint), ] + class c_nvmlUnitFanInfo_t(_PrintableStructure): _fields_ = [ - ('speed', c_uint), - ('state', _nvmlFanState_t), + ("speed", c_uint), + ("state", _nvmlFanState_t), ] + class c_nvmlUnitFanSpeeds_t(_PrintableStructure): - _fields_ = [ - ('fans', c_nvmlUnitFanInfo_t * 24), - ('count', c_uint) - ] + _fields_ = [("fans", c_nvmlUnitFanInfo_t * 24), ("count", c_uint)] + ## Device structures class struct_c_nvmlDevice_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + class nvmlPciInfoExt_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('domain', c_uint), - ('bus', c_uint), - ('device', c_uint), - ('pciDeviceId', c_uint), - ('pciSubSystemId', c_uint), - ('baseClass', c_uint), - ('subClass', c_uint), - ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ("version", c_uint), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + ("pciSubSystemId", c_uint), + ("baseClass", c_uint), + ("subClass", c_uint), + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { - 'version' : "0x%04X", - 'domain' : "0x%04X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - 'baseClass' : "0x%01X", - 'subClass' : "0x%01X", - } + "version": "0x%04X", + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + "baseClass": "0x%01X", + "subClass": "0x%01X", + } + nvmlPciInfoExt_v1 = 0x1000040 + # Legacy pciInfo used for _v1 and _v2 class nvmlPciInfo_v2_t(_PrintableStructure): _fields_ = [ - ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), - ('domain', c_uint), - ('bus', c_uint), - ('device', c_uint), - ('pciDeviceId', c_uint), - + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), # Added in 2.285 - ('pciSubSystemId', c_uint), - ('reserved0', c_uint), - ('reserved1', c_uint), - ('reserved2', c_uint), - ('reserved3', c_uint), + ("pciSubSystemId", c_uint), + ("reserved0", c_uint), + ("reserved1", c_uint), + ("reserved2", c_uint), + ("reserved3", c_uint), ] _fmt_ = { - 'domain' : "0x%04X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - } + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + class nvmlPciInfo_t(_PrintableStructure): _fields_ = [ # Moved to the new busId location below - ('busIdLegacy', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), - ('domain', c_uint), - ('bus', c_uint), - ('device', c_uint), - ('pciDeviceId', c_uint), - + ("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), # Added in 2.285 - ('pciSubSystemId', c_uint), + ("pciSubSystemId", c_uint), # New busId replaced the long deprecated and reserved fields with a # field of the same size in 9.0 - ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { - 'domain' : "0x%08X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - } + "domain": "0x%08X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), + ("version", c_uint), ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ] + SystemDriverBranchInfo_v1 = 0x1000054 + class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): - _fields_ = [ - ('pci', nvmlPciInfo_t), - ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE) - ] + _fields_ = [("pci", nvmlPciInfo_t), ("uuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE)] + class nvmlNvLinkUtilizationControl_t(_PrintableStructure): _fields_ = [ - ('units', _nvmlNvLinkUtilizationCountUnits_t), - ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t), + ("units", _nvmlNvLinkUtilizationCountUnits_t), + ("pktfilter", _nvmlNvLinkUtilizationCountPktTypes_t), ] + class c_nvmlMemory_t(_PrintableStructure): _fields_ = [ - ('total', c_ulonglong), - ('free', c_ulonglong), - ('used', c_ulonglong), + ("total", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), ] - _fmt_ = {'': "%d B"} + _fmt_ = {"": "%d B"} + class c_nvmlMemory_v2_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('total', c_ulonglong), - ('reserved', c_ulonglong), - ('free', c_ulonglong), - ('used', c_ulonglong), + ("version", c_uint), + ("total", c_ulonglong), + ("reserved", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), ] - _fmt_ = {'': "%d B"} + _fmt_ = {"": "%d B"} + nvmlMemory_v2 = 0x02000028 + class c_nvmlBAR1Memory_t(_PrintableStructure): _fields_ = [ - ('bar1Total', c_ulonglong), - ('bar1Free', c_ulonglong), - ('bar1Used', c_ulonglong), + ("bar1Total", c_ulonglong), + ("bar1Free", c_ulonglong), + ("bar1Used", c_ulonglong), ] - _fmt_ = {'': "%d B"} + _fmt_ = {"": "%d B"} + class nvmlClkMonFaultInfo_t(Structure): - _fields_ = [("clkApiDomain", c_uint), - ("clkDomainFaultMask", c_uint) - ] + _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)] + MAX_CLK_DOMAINS = 32 + class nvmlClkMonStatus_t(Structure): - _fields_ = [("bGlobalStatus", c_uint), - ("clkMonListSize", c_uint), - ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS) + _fields_ = [ + ("bGlobalStatus", c_uint), + ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS), ] + # On Windows with the WDDM driver, usedGpuMemory is reported as None # Code that processes this structure should check for None, I.E. # @@ -1351,887 +1638,1031 @@ class nvmlClkMonStatus_t(Structure): # See NVML documentation for more information class c_nvmlProcessInfo_v2_t(_PrintableStructure): _fields_ = [ - ('pid', c_uint), - ('usedGpuMemory', c_ulonglong), - ('gpuInstanceId', c_uint), - ('computeInstanceId', c_uint), + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), ] - _fmt_ = {'usedGpuMemory': "%d B"} + _fmt_ = {"usedGpuMemory": "%d B"} + c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t _nvmlProcessMode_t = c_uint -NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_COMPUTE = 0 NVML_PROCESS_MODE_GRAPHICS = 1 -NVML_PROCESS_MODE_MPS = 2 +NVML_PROCESS_MODE_MPS = 2 + class c_nvmlProcessDetail_v1_t(Structure): _fields_ = [ - ('pid', c_uint), - ('usedGpuMemory', c_ulonglong), - ('gpuInstanceId', c_uint), - ('computeInstanceId', c_uint), - ('usedGpuCcProtectedMemory', c_ulonglong), + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ("usedGpuCcProtectedMemory", c_ulonglong), ] + class c_nvmlProcessDetailList_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('mode', _nvmlProcessMode_t), - ('numProcArrayEntries', c_uint), - ('procArray', POINTER(c_nvmlProcessDetail_v1_t)), + ("version", c_uint), + ("mode", _nvmlProcessMode_t), + ("numProcArrayEntries", c_uint), + ("procArray", POINTER(c_nvmlProcessDetail_v1_t)), ] - _fmt_ = {'numProcArrayEntries': "%d B"} + _fmt_ = {"numProcArrayEntries": "%d B"} + c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t nvmlProcessDetailList_v1 = 0x1000018 + class c_nvmlBridgeChipInfo_t(_PrintableStructure): _fields_ = [ - ('type', _nvmlBridgeChipType_t), - ('fwVersion', c_uint), + ("type", _nvmlBridgeChipType_t), + ("fwVersion", c_uint), ] + class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): _fields_ = [ - ('bridgeCount', c_uint), - ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), + ("bridgeCount", c_uint), + ("bridgeChipInfo", c_nvmlBridgeChipInfo_t * 128), ] + class c_nvmlEccErrorCounts_t(_PrintableStructure): _fields_ = [ - ('l1Cache', c_ulonglong), - ('l2Cache', c_ulonglong), - ('deviceMemory', c_ulonglong), - ('registerFile', c_ulonglong), + ("l1Cache", c_ulonglong), + ("l2Cache", c_ulonglong), + ("deviceMemory", c_ulonglong), + ("registerFile", c_ulonglong), ] + class c_nvmlUtilization_t(_PrintableStructure): _fields_ = [ - ('gpu', c_uint), - ('memory', c_uint), + ("gpu", c_uint), + ("memory", c_uint), ] - _fmt_ = {'': "%d %%"} + _fmt_ = {"": "%d %%"} + # Added in 2.285 class c_nvmlHwbcEntry_t(_PrintableStructure): _fields_ = [ - ('hwbcId', c_uint), - ('firmwareVersion', c_char * 32), + ("hwbcId", c_uint), + ("firmwareVersion", c_char * 32), ] + class c_nvmlValue_t(Union): _fields_ = [ - ('dVal', c_double), - ('uiVal', c_uint), - ('ulVal', c_ulong), - ('ullVal', c_ulonglong), - ('sllVal', c_longlong), - ('siVal', c_int), - ('usVal', c_ushort), + ("dVal", c_double), + ("uiVal", c_uint), + ("ulVal", c_ulong), + ("ullVal", c_ulonglong), + ("sllVal", c_longlong), + ("siVal", c_int), + ("usVal", c_ushort), ] + class c_nvmlSample_t(_PrintableStructure): _fields_ = [ - ('timeStamp', c_ulonglong), - ('sampleValue', c_nvmlValue_t), + ("timeStamp", c_ulonglong), + ("sampleValue", c_nvmlValue_t), ] + class c_nvmlViolationTime_t(_PrintableStructure): _fields_ = [ - ('referenceTime', c_ulonglong), - ('violationTime', c_ulonglong), + ("referenceTime", c_ulonglong), + ("violationTime", c_ulonglong), ] + class c_nvmlFieldValue_t(_PrintableStructure): _fields_ = [ - ('fieldId', c_uint32), - ('scopeId', c_uint32), - ('timestamp', c_int64), - ('latencyUsec', c_int64), - ('valueType', _nvmlValueType_t), - ('nvmlReturn', _nvmlReturn_t), - ('value', c_nvmlValue_t) + ("fieldId", c_uint32), + ("scopeId", c_uint32), + ("timestamp", c_int64), + ("latencyUsec", c_int64), + ("valueType", _nvmlValueType_t), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_nvmlValue_t), ] + NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 -nvmlNvlinkSupportedBwModes_v1 = 0x100001c +nvmlNvlinkSupportedBwModes_v1 = 0x100001C + + class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), - ('totalBwModes', c_uint8) + ("version", c_uint), + ("bwModes", c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ("totalBwModes", c_uint8), ] def __init__(self): - super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(version=nvmlNvlinkSupportedBwModes_v1) + super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__( + version=nvmlNvlinkSupportedBwModes_v1 + ) + + +nvmlNvlinkGetBwMode_v1 = 0x100000C + -nvmlNvlinkGetBwMode_v1 = 0x100000c class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('bIsBest', c_uint), - ('bwMode', c_uint8) - ] + _fields_ = [("version", c_uint), ("bIsBest", c_uint), ("bwMode", c_uint8)] def __init__(self): super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) -nvmlNvlinkSetBwMode_v1 = 0x100000c + +nvmlNvlinkSetBwMode_v1 = 0x100000C + + class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('bSetBest', c_uint), - ('bwMode', c_uint8) - ] + _fields_ = [("version", c_uint), ("bSetBest", c_uint), ("bwMode", c_uint8)] def __init__(self): super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) + class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('mode', c_uint), + ("version", c_uint), + ("mode", c_uint), ] + VgpuHeterogeneousMode_v1 = 0x1000008 + class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('placementId', c_uint), + ("version", c_uint), + ("placementId", c_uint), ] + VgpuPlacementId_v1 = 0x1000008 + class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('count', c_uint), - ('placementSize', c_uint), - ('placementIds', POINTER(c_uint)), + ("version", c_uint), + ("count", c_uint), + ("placementSize", c_uint), + ("placementIds", POINTER(c_uint)), ] + VgpuPlacementList_v1 = 0x1000018 -NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 -NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('placementSize', c_uint), - ('count', c_uint), - ('placementIds', POINTER(c_uint)), - ('mode', c_uint), + ("version", c_uint), + ("placementSize", c_uint), + ("count", c_uint), + ("placementIds", POINTER(c_uint)), + ("mode", c_uint), ] + VgpuPlacementList_v2 = 0x2000020 + class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('bar1Size', c_ulonglong), + ("version", c_uint), + ("bar1Size", c_ulonglong), ] + VgpuTypeBar1Info_v1 = 0x1000010 + class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): _fields_ = [ - ('vgpuInstance', _nvmlVgpuInstance_t), - ('timeStamp', c_ulonglong), - ('smUtil', c_nvmlValue_t), - ('memUtil', c_nvmlValue_t), - ('encUtil', c_nvmlValue_t), - ('decUtil', c_nvmlValue_t), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("timeStamp", c_ulonglong), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), ] + class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('timeStamp', c_ulonglong), - ('vgpuInstance', _nvmlVgpuInstance_t), - ('smUtil', c_nvmlValue_t), - ('memUtil', c_nvmlValue_t), - ('encUtil', c_nvmlValue_t), - ('decUtil', c_nvmlValue_t), - ('jpgUtil', c_nvmlValue_t), - ('ofaUtil', c_nvmlValue_t), + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), + ("jpgUtil", c_nvmlValue_t), + ("ofaUtil", c_nvmlValue_t), ] + class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('sampleValType', _nvmlValueType_t), - ('vgpuInstanceCount', c_uint), - ('lastSeenTimeStamp', c_ulonglong), - ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), + ("version", c_uint), + ("sampleValType", _nvmlValueType_t), + ("vgpuInstanceCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuUtilArray", POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), ] + VgpuInstancesUtilizationInfo_v1 = 0x01000020 + class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ - ('vgpuInstance', _nvmlVgpuInstance_t), - ('pid', c_uint), - ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), - ('timeStamp', c_ulonglong), - ('smUtil', c_uint), - ('memUtil', c_uint), - ('encUtil', c_uint), - ('decUtil', c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), ] + class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), - ('timeStamp', c_ulonglong), - ('vgpuInstance', _nvmlVgpuInstance_t), - ('pid', c_uint), - ('smUtil', c_uint), - ('memUtil', c_uint), - ('encUtil', c_uint), - ('decUtil', c_uint), - ('jpgUtil', c_uint), - ('ofaUtil', c_uint), + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), ] + class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('vgpuProcessCount', c_uint), - ('lastSeenTimeStamp', c_ulonglong), - ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), + ("version", c_uint), + ("vgpuProcessCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuProcUtilArray", POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), ] + VgpuProcessesUtilizationInfo_v1 = 0x01000018 + class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('size', c_ulonglong), + ("version", c_uint), + ("size", c_ulonglong), ] + VgpuRuntimeState_v1 = 0x1000010 + class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): _fields_ = [ - ('year', c_uint32), - ('month', c_uint16), - ('day', c_uint16), - ('hour', c_uint16), - ('min', c_uint16), - ('sec', c_uint16), - ('status', c_uint8), + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), ] -NVML_GRID_LICENSE_STATE_UNKNOWN = 0 -NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 -NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 -NVML_GRID_LICENSE_STATE_UNLICENSED = 4 -NVML_GRID_LICENSE_STATE_LICENSED = 5 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): _fields_ = [ - ('isLicensed', c_uint8), - ('licenseExpiry', c_nvmlVgpuLicenseExpiry_t), - ('currentState', c_uint), + ("isLicensed", c_uint8), + ("licenseExpiry", c_nvmlVgpuLicenseExpiry_t), + ("currentState", c_uint), ] + class c_nvmlEncoderSession_t(_PrintableStructure): _fields_ = [ - ('sessionId', c_uint), - ('pid', c_uint), - ('vgpuInstance', _nvmlVgpuInstance_t), - ('codecType', c_uint), - ('hResolution', c_uint), - ('vResolution', c_uint), - ('averageFps', c_uint), - ('encodeLatency', c_uint), + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("codecType", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFps", c_uint), + ("encodeLatency", c_uint), ] + class c_nvmlProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ - ('pid', c_uint), - ('timeStamp', c_ulonglong), - ('smUtil', c_uint), - ('memUtil', c_uint), - ('encUtil', c_uint), - ('decUtil', c_uint), + ("pid", c_uint), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), ] + class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('timeStamp', c_ulonglong), - ('pid', c_uint), - ('smUtil', c_uint), - ('memUtil', c_uint), - ('encUtil', c_uint), - ('decUtil', c_uint), - ('jpgUtil', c_uint), - ('ofaUtil', c_uint), + ("timeStamp", c_ulonglong), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), ] + class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('processSamplesCount', c_uint), - ('lastSeenTimeStamp', c_ulonglong), - ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)), + ("version", c_uint), + ("processSamplesCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("procUtilArray", POINTER(c_nvmlProcessUtilizationInfo_v1_t)), ] + ProcessesUtilizationInfo_v1 = 0x01000018 + class c_nvmlGridLicenseExpiry_t(_PrintableStructure): _fields_ = [ - ('year', c_uint32), - ('month', c_uint16), - ('day', c_uint16), - ('hour', c_uint16), - ('min', c_uint16), - ('sec', c_uint16), - ('status', c_uint8), + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), ] + class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): _fields_ = [ - ('featureCode', _nvmlGridLicenseFeatureCode_t), - ('featureState', c_uint), - ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('featureEnabled', c_uint), - ('licenseExpiry', c_nvmlGridLicenseExpiry_t), + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), + ("licenseExpiry", c_nvmlGridLicenseExpiry_t), ] + class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): _fields_ = [ - ('isGridLicenseSupported', c_int), - ('licensableFeaturesCount', c_uint), - ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), ] + class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): _fields_ = [ - ('featureCode', _nvmlGridLicenseFeatureCode_t), - ('featureState', c_uint), - ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('featureEnabled', c_uint), + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), ] + class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): _fields_ = [ - ('isGridLicenseSupported', c_int), - ('licensableFeaturesCount', c_uint), - ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), ] + class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): _fields_ = [ - ('featureCode', _nvmlGridLicenseFeatureCode_t), - ('featureState', c_uint), - ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] + class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): _fields_ = [ - ('isGridLicenseSupported', c_int), - ('licensableFeaturesCount', c_uint), - ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), ] + class c_nvmlGridLicensableFeature_t(_PrintableStructure): _fields_ = [ - ('featureCode', _nvmlGridLicenseFeatureCode_t), - ('featureState', c_uint), - ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] + class c_nvmlGridLicensableFeatures_t(_PrintableStructure): _fields_ = [ - ('isGridLicenseSupported', c_int), - ('licensableFeaturesCount', c_uint), - ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), ] + class c_nvmlMarginTemperature_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('marginTemperature', c_int), + ("version", c_uint), + ("marginTemperature", c_int), ] + nvmlMarginTemperature_v1 = 0x1000008 + ## Event structures class struct_c_nvmlEventSet_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) -nvmlEventTypeSingleBitEccError = 0x0000000000000001 -nvmlEventTypeDoubleBitEccError = 0x0000000000000002 -nvmlEventTypePState = 0x0000000000000004 -nvmlEventTypeXidCriticalError = 0x0000000000000008 -nvmlEventTypeClock = 0x0000000000000010 -nvmlEventTypePowerSourceChange = 0x0000000000000080 -nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 -nvmlEventTypeDramRetirementEvent = 0x0000000000000400 -nvmlEventTypeDramRetirementFailure = 0x0000000000000800 -nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 -nvmlEventTypeFatalPoisonError = 0x0000000000002000 -nvmlEventTypeGpuUnavailableError = 0x0000000000004000 -nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 -nvmlEventTypeNone = 0x0000000000000000 -nvmlEventTypeAll = ( - nvmlEventTypeNone - | nvmlEventTypeSingleBitEccError - | nvmlEventTypeDoubleBitEccError - | nvmlEventTypePState - | nvmlEventTypeClock - | nvmlEventTypePowerSourceChange - | nvmlEventTypeXidCriticalError - | nvmlEventMigConfigChange - | nvmlEventTypeSingleBitEccErrorStorm - | nvmlEventTypeDramRetirementEvent - | nvmlEventTypeDramRetirementFailure - | nvmlEventTypeNonFatalPoisonError - | nvmlEventTypeFatalPoisonError - | nvmlEventTypeGpuUnavailableError - | nvmlEventTypeGpuRecoveryAction - ) +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = ( + nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction +) ## Clock Event Reasons defines -nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 -nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting -nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 -nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 -nvmlClocksEventReasonSyncBoost = 0x0000000000000010 -nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 -nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 -nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 -nvmlClocksEventReasonNone = 0x0000000000000000 -nvmlClocksEventReasonAll = ( - nvmlClocksEventReasonNone | - nvmlClocksEventReasonGpuIdle | - nvmlClocksEventReasonApplicationsClocksSetting | - nvmlClocksEventReasonSwPowerCap | - nvmlClocksEventReasonHwSlowdown | - nvmlClocksEventReasonSyncBoost | - nvmlClocksEventReasonSwThermalSlowdown | - nvmlClocksEventReasonHwThermalSlowdown | - nvmlClocksEventReasonHwPowerBrakeSlowdown | - nvmlClocksEventReasonDisplayClockSetting - ) +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = ( + nvmlClocksEventReasonNone + | nvmlClocksEventReasonGpuIdle + | nvmlClocksEventReasonApplicationsClocksSetting + | nvmlClocksEventReasonSwPowerCap + | nvmlClocksEventReasonHwSlowdown + | nvmlClocksEventReasonSyncBoost + | nvmlClocksEventReasonSwThermalSlowdown + | nvmlClocksEventReasonHwThermalSlowdown + | nvmlClocksEventReasonHwPowerBrakeSlowdown + | nvmlClocksEventReasonDisplayClockSetting +) ## Following have been deprecated -nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 -nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting -nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 -nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 -nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 -nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 -nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 -nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 -nvmlClocksThrottleReasonNone = 0x0000000000000000 -nvmlClocksThrottleReasonAll = ( - nvmlClocksThrottleReasonNone | - nvmlClocksThrottleReasonGpuIdle | - nvmlClocksThrottleReasonApplicationsClocksSetting | - nvmlClocksThrottleReasonSwPowerCap | - nvmlClocksThrottleReasonHwSlowdown | - nvmlClocksThrottleReasonSyncBoost | - nvmlClocksThrottleReasonSwThermalSlowdown | - nvmlClocksThrottleReasonHwThermalSlowdown | - nvmlClocksThrottleReasonHwPowerBrakeSlowdown | - nvmlClocksThrottleReasonDisplayClockSetting - ) +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = ( + nvmlClocksThrottleReasonNone + | nvmlClocksThrottleReasonGpuIdle + | nvmlClocksThrottleReasonApplicationsClocksSetting + | nvmlClocksThrottleReasonSwPowerCap + | nvmlClocksThrottleReasonHwSlowdown + | nvmlClocksThrottleReasonSyncBoost + | nvmlClocksThrottleReasonSwThermalSlowdown + | nvmlClocksThrottleReasonHwThermalSlowdown + | nvmlClocksThrottleReasonHwPowerBrakeSlowdown + | nvmlClocksThrottleReasonDisplayClockSetting +) + class c_nvmlEventData_t(_PrintableStructure): _fields_ = [ - ('device', c_nvmlDevice_t), - ('eventType', c_ulonglong), - ('eventData', c_ulonglong), - ('gpuInstanceId', c_uint), - ('computeInstanceId', c_uint) + ("device", c_nvmlDevice_t), + ("eventType", c_ulonglong), + ("eventData", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), ] - _fmt_ = {'eventType': "0x%08X"} + _fmt_ = {"eventType": "0x%08X"} + class c_nvmlAccountingStats_t(_PrintableStructure): _fields_ = [ - ('gpuUtilization', c_uint), - ('memoryUtilization', c_uint), - ('maxMemoryUsage', c_ulonglong), - ('time', c_ulonglong), - ('startTime', c_ulonglong), - ('isRunning', c_uint), - ('reserved', c_uint * 5) + ("gpuUtilization", c_uint), + ("memoryUtilization", c_uint), + ("maxMemoryUsage", c_ulonglong), + ("time", c_ulonglong), + ("startTime", c_ulonglong), + ("isRunning", c_uint), + ("reserved", c_uint * 5), ] + class c_nvmlVgpuVersion_t(Structure): - _fields_ = [("minVersion", c_uint), - ("maxVersion", c_uint) - ] + _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)] + class c_nvmlVgpuMetadata_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("revision", c_uint), - ("guestInfoState", _nvmlVgpuGuestInfoState_t), - ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("reserved", c_uint * 6), - ("vgpuVirtualizationCaps", c_uint), - ("guestVgpuVersion", c_uint), - ("opaqueDataSize", c_uint), - ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE) - ] + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("guestInfoState", _nvmlVgpuGuestInfoState_t), + ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), + ("guestVgpuVersion", c_uint), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE), + ] + class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("revision", c_uint), - ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("pgpuVirtualizationCaps", c_uint), - ("reserved", c_uint * 5), - ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), - ("opaqueDataSize", c_uint), - ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) - ] + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("pgpuVirtualizationCaps", c_uint), + ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE), + ] + class c_nvmlVgpuPgpuCompatibility_t(Structure): - _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), - ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t) - ] + _fields_ = [ + ("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t), + ] + ## vGPU scheduler policy defines -NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 -NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 -NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 -NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 ## Supported vGPU scheduler policy count -NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 -NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 -NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 -NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 -NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): _fields_ = [ - ('avgFactor', c_uint), - ('timeslice', c_uint), + ("avgFactor", c_uint), + ("timeslice", c_uint), ] + class c_nvmlVgpuSchedData_t(_PrintableStructure): _fields_ = [ - ('timeslice', c_uint), + ("timeslice", c_uint), ] + class c_nvmlVgpuSchedulerParams_t(Union): _fields_ = [ - ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t), - ('vgpuSchedData', c_nvmlVgpuSchedData_t), + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedData_t), ] + class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): _fields_ = [ - ('timestamp', c_ulonglong), - ('timeRunTotal', c_ulonglong), - ('timeRun', c_ulonglong), - ('swRunlistId', c_uint), - ('targetTimeSlice', c_ulonglong), - ('cumulativePreemptionTime', c_ulonglong), + ("timestamp", c_ulonglong), + ("timeRunTotal", c_ulonglong), + ("timeRun", c_ulonglong), + ("swRunlistId", c_uint), + ("targetTimeSlice", c_ulonglong), + ("cumulativePreemptionTime", c_ulonglong), ] + class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): _fields_ = [ - ('engineId', c_uint), - ('schedulerPolicy', c_uint), - ('arrMode', c_uint), - ('schedulerParams', c_nvmlVgpuSchedulerParams_t), - ('entriesCount', c_uint), - ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), + ("engineId", c_uint), + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), + ("entriesCount", c_uint), + ( + "logEntries", + c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES, + ), ] + class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): _fields_ = [ - ('schedulerPolicy', c_uint), - ('arrMode', c_uint), - ('schedulerParams', c_nvmlVgpuSchedulerParams_t), + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), ] + class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): _fields_ = [ - ('avgFactor', c_uint), - ('frequency', c_uint), + ("avgFactor", c_uint), + ("frequency", c_uint), ] + class c_nvmlVgpuSchedSetData_t(_PrintableStructure): _fields_ = [ - ('timeslice', c_uint), + ("timeslice", c_uint), ] + class c_nvmlVgpuSchedulerSetParams_t(Union): _fields_ = [ - ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t), - ('vgpuSchedData', c_nvmlVgpuSchedSetData_t), + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedSetDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedSetData_t), ] + class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): _fields_ = [ - ('schedulerPolicy', c_uint), - ('enableARRMode', c_uint), - ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t), + ("schedulerPolicy", c_uint), + ("enableARRMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerSetParams_t), ] + class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): _fields_ = [ - ('supportedSchedulers', c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), - ('maxTimeslice', c_uint), - ('minTimeslice', c_uint), - ('isArrModeSupported', c_uint), - ('maxFrequencyForARR', c_uint), - ('minFrequencyForARR', c_uint), - ('maxAvgFactorForARR', c_uint), - ('minAvgFactorForARR', c_uint), + ("supportedSchedulers", c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), + ("maxTimeslice", c_uint), + ("minTimeslice", c_uint), + ("isArrModeSupported", c_uint), + ("maxFrequencyForARR", c_uint), + ("minFrequencyForARR", c_uint), + ("maxAvgFactorForARR", c_uint), + ("minAvgFactorForARR", c_uint), ] + class c_nvmlFBCStats_t(Structure): - _fields_ = [("sessionsCount", c_uint), - ("averageFPS", c_uint), - ("averageLatency", c_uint) - ] + _fields_ = [ + ("sessionsCount", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), + ] + class c_nvmlFBCSession_t(_PrintableStructure): _fields_ = [ - ('sessionId', c_uint), - ('pid', c_uint), - ('vgpuInstance', _nvmlVgpuInstance_t), - ('displayOrdinal', c_uint), - ('sessionType', c_uint), - ('sessionFlags', c_uint), - ('hMaxResolution', c_uint), - ('vMaxResolution', c_uint), - ('hResolution', c_uint), - ('vResolution', c_uint), - ('averageFPS', c_uint), - ('averageLatency', c_uint), + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("displayOrdinal", c_uint), + ("sessionType", c_uint), + ("sessionFlags", c_uint), + ("hMaxResolution", c_uint), + ("vMaxResolution", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), ] + NVML_DEVICE_MIG_DISABLE = 0x0 -NVML_DEVICE_MIG_ENABLE = 0x1 - -NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 -NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 -NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 -NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 -NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 -NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 -NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 -NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA -NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB -NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC -NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD + class c_nvmlGpuInstancePlacement_t(Structure): - _fields_ = [("start", c_uint), - ("size", c_uint) - ] + _fields_ = [("start", c_uint), ("size", c_uint)] + class c_nvmlGpuInstanceProfileInfo_t(Structure): - _fields_ = [("id", c_uint), - ("isP2pSupported", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("copyEngineCount", c_uint), - ("decoderCount", c_uint), - ("encoderCount", c_uint), - ("jpegCount", c_uint), - ("ofaCount", c_uint), - ("memorySizeMB", c_ulonglong), - ] + _fields_ = [ + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("id", c_uint), - ("isP2pSupported", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("copyEngineCount", c_uint), - ("decoderCount", c_uint), - ("encoderCount", c_uint), - ("jpegCount", c_uint), - ("ofaCount", c_uint), - ("memorySizeMB", c_ulonglong), - ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) - ] - + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] + def __init__(self): - super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2) + super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__( + version=nvmlGpuInstanceProfileInfo_v2 + ) + class c_nvmlGpuInstanceInfo_t(Structure): - _fields_ = [("device", c_nvmlDevice_t), - ("id", c_uint), - ("profileId", c_uint), - ("placement", c_nvmlGpuInstancePlacement_t) - ] + _fields_ = [ + ("device", c_nvmlDevice_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t), + ] + class struct_c_nvmlGpuInstance_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) -NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 -NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 -NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 -NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 -NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 -NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 -NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 -NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + class c_nvmlComputeInstancePlacement_t(Structure): - _fields_ = [("start", c_uint), - ("size", c_uint) - ] + _fields_ = [("start", c_uint), ("size", c_uint)] + class c_nvmlComputeInstanceProfileInfo_t(Structure): - _fields_ = [("id", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint) - ] + _fields_ = [ + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ] + nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("id", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint), - ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) - ] + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] def __init__(self): - super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2) + super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__( + version=nvmlComputeInstanceProfileInfo_v2 + ) + class c_nvmlComputeInstanceInfo_t(Structure): - _fields_ = [("device", c_nvmlDevice_t), - ("gpuInstance", c_nvmlGpuInstance_t), - ("id", c_uint), - ("profileId", c_uint), - ("placement", c_nvmlComputeInstancePlacement_t) - ] + _fields_ = [ + ("device", c_nvmlDevice_t), + ("gpuInstance", c_nvmlGpuInstance_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t), + ] + NVML_MAX_GPU_UTILIZATIONS = 8 -NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 -NVML_GPU_UTILIZATION_DOMAIN_FB = 1 -NVML_GPU_UTILIZATION_DOMAIN_VID = 2 -NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 + + class c_nvmlGpuDynamicPstatesUtilization_t(Structure): - _fields_ = [("bIsPresent", c_uint, 1), - ("percentage", c_uint), - ("incThreshold", c_uint), - ("decThreshold", c_uint)] + _fields_ = [ + ("bIsPresent", c_uint, 1), + ("percentage", c_uint), + ("incThreshold", c_uint), + ("decThreshold", c_uint), + ] + + class c_nvmlGpuDynamicPstatesInfo_t(Structure): - _fields_ = [("flags", c_uint), - ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)] + _fields_ = [ + ("flags", c_uint), + ( + "utilization", + c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS, + ), + ] + NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 -NVML_THERMAL_TARGET_NONE = 0 -NVML_THERMAL_TARGET_GPU = 1 -NVML_THERMAL_TARGET_MEMORY = 2 -NVML_THERMAL_TARGET_POWER_SUPPLY = 4 -NVML_THERMAL_TARGET_BOARD = 8 -NVML_THERMAL_TARGET_VCD_BOARD = 9 -NVML_THERMAL_TARGET_VCD_INLET = 10 -NVML_THERMAL_TARGET_VCD_OUTLET = 11 -NVML_THERMAL_TARGET_ALL = 15 -NVML_THERMAL_TARGET_UNKNOWN = -1 - -NVML_THERMAL_CONTROLLER_NONE = 0 -NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 -NVML_THERMAL_CONTROLLER_ADM1032 = 2 -NVML_THERMAL_CONTROLLER_ADT7461 = 3 -NVML_THERMAL_CONTROLLER_MAX6649 = 4 -NVML_THERMAL_CONTROLLER_MAX1617 = 5 -NVML_THERMAL_CONTROLLER_LM99 = 6 -NVML_THERMAL_CONTROLLER_LM89 = 7 -NVML_THERMAL_CONTROLLER_LM64 = 8 -NVML_THERMAL_CONTROLLER_G781 = 9 -NVML_THERMAL_CONTROLLER_ADT7473 = 10 -NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 -NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 -NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 + +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 -NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 -NVML_THERMAL_CONTROLLER_MAX6649R = 16 -NVML_THERMAL_CONTROLLER_ADT7473S = 17 -NVML_THERMAL_CONTROLLER_UNKNOWN = -1 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + class c_nvmlGpuThermalSensor_t(Structure): - _fields_ = [("controller", c_int), - ("defaultMinTemp", c_int), - ("defaultMaxTemp", c_int), - ("currentTemp", c_int), - ("target", c_int)] + _fields_ = [ + ("controller", c_int), + ("defaultMinTemp", c_int), + ("defaultMaxTemp", c_int), + ("currentTemp", c_int), + ("target", c_int), + ] + + class c_nvmlGpuThermalSettings_t(Structure): - _fields_ = [("count", c_uint), - ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)] + _fields_ = [ + ("count", c_uint), + ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU), + ] + _nvmlCoolerControl_t = c_uint -NVML_THERMAL_COOLER_SIGNAL_NONE = 0 -NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 -NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 -NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 _nvmlCoolerTarget_t = c_uint -NVML_THERMAL_COOLER_TARGET_NONE = (1 << 0) -NVML_THERMAL_COOLER_TARGET_GPU = (1 << 1) -NVML_THERMAL_COOLER_TARGET_MEMORY = (1 << 2) -NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = (1 << 3) -NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY) +NVML_THERMAL_COOLER_TARGET_NONE = 1 << 0 +NVML_THERMAL_COOLER_TARGET_GPU = 1 << 1 +NVML_THERMAL_COOLER_TARGET_MEMORY = 1 << 2 +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = 1 << 3 +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = ( + NVML_THERMAL_COOLER_TARGET_GPU + | NVML_THERMAL_COOLER_TARGET_MEMORY + | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY +) + class c_nvmlCoolerInfo_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("index", c_uint), - ("coolerControlType", _nvmlCoolerControl_t), - ("coolerTarget", _nvmlCoolerTarget_t) - ] + _fields_ = [ + ("version", c_uint), + ("index", c_uint), + ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t), + ] + nvmlCoolerInfo_v1 = 0x1000010 + def nvmlDeviceGetCoolerInfo(handle): c_coolerInfo = c_nvmlCoolerInfo_t() c_coolerInfo.version = nvmlCoolerInfo_v1 @@ -2241,105 +2672,138 @@ def nvmlDeviceGetCoolerInfo(handle): _nvmlCheckReturn(ret) return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + class struct_c_nvmlComputeInstance_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + class c_nvmlDeviceAttributes(Structure): - _fields_ = [("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint), - ("gpuInstanceSliceCount", c_uint), - ("computeInstanceSliceCount", c_uint), - ("memorySizeMB", c_ulonglong), - ] + _fields_ = [ + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + class c_nvmlRowRemapperHistogramValues(Structure): - _fields_ = [("max", c_uint), - ("high", c_uint), - ("partial", c_uint), - ("low", c_uint), - ("none", c_uint) - ] - -NVML_GPU_CERT_CHAIN_SIZE = 0x1000 -NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 -NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 -NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 + _fields_ = [ + ("max", c_uint), + ("high", c_uint), + ("partial", c_uint), + ("low", c_uint), + ("none", c_uint), + ] + + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 -NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + class c_nvmlConfComputeSystemState_t(Structure): - _fields_ = [('environment', c_uint), - ('ccFeature', c_uint), - ('devToolsMode', c_uint), - ] + _fields_ = [ + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ] + nvmlSystemConfComputeSettings_v1 = 0x1000014 + class c_nvmlSystemConfComputeSettings_v1_t(Structure): - _fields_ = [('version', c_uint), - ('environment', c_uint), - ('ccFeature', c_uint), - ('devToolsMode', c_uint), - ('multiGpuMode', c_uint), - ] + _fields_ = [ + ("version", c_uint), + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ("multiGpuMode", c_uint), + ] + def __init__(self): - super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1) + super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__( + version=nvmlSystemConfComputeSettings_v1 + ) + class c_nvmlConfComputeSystemCaps_t(Structure): - _fields_ = [('cpuCaps', c_uint), - ('gpusCaps', c_uint), - ] + _fields_ = [ + ("cpuCaps", c_uint), + ("gpusCaps", c_uint), + ] + class c_nvmlConfComputeMemSizeInfo_t(Structure): - _fields_ = [('protectedMemSizeKib', c_ulonglong), - ('unprotectedMemSizeKib', c_ulonglong), - ] + _fields_ = [ + ("protectedMemSizeKib", c_ulonglong), + ("unprotectedMemSizeKib", c_ulonglong), + ] + class c_nvmlConfComputeGpuCertificate_t(Structure): - _fields_ = [('certChainSize', c_uint), - ('attestationCertChainSize', c_uint), - ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), - ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), - ] + _fields_ = [ + ("certChainSize", c_uint), + ("attestationCertChainSize", c_uint), + ("certChain", c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ("attestationCertChain", c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), + ] + class c_nvmlConfComputeGpuAttestationReport_t(Structure): - _fields_ = [('isCecAttestationReportPresent', c_uint), - ('attestationReportSize', c_uint), - ('cecAttestationReportSize', c_uint), - ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), - ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), - ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), - ] + _fields_ = [ + ("isCecAttestationReportPresent", c_uint), + ("attestationReportSize", c_uint), + ("cecAttestationReportSize", c_uint), + ("nonce", c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), + ("attestationReport", c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), + ("cecAttestationReport", c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), + ] + class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): - _fields_ = [('version', c_uint), - ('maxAttackerAdvantage', c_ulong), - ] + _fields_ = [ + ("version", c_uint), + ("maxAttackerAdvantage", c_ulong), + ] + + ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): - _fields_ = [('version', c_uint), - ('attackerAdvantage', c_ulong), - ] + _fields_ = [ + ("version", c_uint), + ("attackerAdvantage", c_ulong), + ] + + ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 ## string/bytes conversion for ease of use def convertStrBytes(func): - ''' + """ In python 3, strings are unicode instead of bytes, and need to be converted for ctypes Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> ---- Returned from function: b'returned string' Returned to caller: 'returned string' - ''' + """ + @wraps(func) def wrapper(*args, **kwargs): # encoding a str returns bytes in python 2 and 3 @@ -2358,18 +2822,23 @@ def wrapper(*args, **kwargs): return wrapper return func + def throwOnVersionMismatch(func): @wraps(func) def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except NVMLError_FunctionNotFound: - raise NVMLLibraryMismatchError("Unversioned function called and the " - "pyNVML version does not match the NVML lib version. " - "Either use matching pyNVML and NVML lib versions or " - "use a versioned function such as " + func.__name__ + "_v2") + raise NVMLLibraryMismatchError( + "Unversioned function called and the " + "pyNVML version does not match the NVML lib version. " + "Either use matching pyNVML and NVML lib versions or " + "use a versioned function such as " + func.__name__ + "_v2" + ) + return wrapper + ## C function wrappers ## def nvmlInitWithFlags(flags): _LoadNvmlLibrary() @@ -2388,44 +2857,57 @@ def nvmlInitWithFlags(flags): libLoadLock.release() return None + def nvmlInit(): nvmlInitWithFlags(0) return None + def _LoadNvmlLibrary(): - ''' + """ Load the library if it isn't loaded already - ''' + """ global nvmlLib - if (nvmlLib == None): + if nvmlLib == None: # lock to ensure only one caller loads the library libLoadLock.acquire() try: # ensure the library still isn't loaded - if (nvmlLib == None): + if nvmlLib == None: try: - if (sys.platform[:3] == "win"): + if sys.platform[:3] == "win": # cdecl calling convention try: # Check for nvml.dll in System32 first for DCH drivers - nvmlLib = CDLL(os.path.join(os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll")) + nvmlLib = CDLL( + os.path.join( + os.getenv("WINDIR", "C:/Windows"), + "System32/nvml.dll", + ) + ) except OSError as ose: # If nvml.dll is not found in System32, it should be in ProgramFiles # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll - nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll")) + nvmlLib = CDLL( + os.path.join( + os.getenv("ProgramFiles", "C:/Program Files"), + "NVIDIA Corporation/NVSMI/nvml.dll", + ) + ) else: # assume linux nvmlLib = CDLL("libnvidia-ml.so.1") except OSError as ose: _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) - if (nvmlLib == None): + if nvmlLib == None: _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) finally: # lock is always freed libLoadLock.release() + def nvmlShutdown(): # # Leave the library loaded, but shutdown the interface @@ -2437,19 +2919,21 @@ def nvmlShutdown(): # Atomically update refcount global _nvmlLib_refcount libLoadLock.acquire() - if (0 < _nvmlLib_refcount): + if 0 < _nvmlLib_refcount: _nvmlLib_refcount -= 1 libLoadLock.release() return None + # Added in 2.285 @convertStrBytes def nvmlErrorString(result): fn = _nvmlGetFunctionPointer("nvmlErrorString") - fn.restype = c_char_p # otherwise return is an int + fn.restype = c_char_p # otherwise return is an int ret = fn(result) return ret + # Added in 2.285 @convertStrBytes def nvmlSystemGetNVMLVersion(): @@ -2459,6 +2943,7 @@ def nvmlSystemGetNVMLVersion(): _nvmlCheckReturn(ret) return c_version.value + def nvmlSystemGetCudaDriverVersion(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") @@ -2466,6 +2951,7 @@ def nvmlSystemGetCudaDriverVersion(): _nvmlCheckReturn(ret) return c_cuda_version.value + def nvmlSystemGetCudaDriverVersion_v2(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") @@ -2473,6 +2959,7 @@ def nvmlSystemGetCudaDriverVersion_v2(): _nvmlCheckReturn(ret) return c_cuda_version.value + # Added in 2.285 @convertStrBytes def nvmlSystemGetProcessName(pid): @@ -2482,6 +2969,7 @@ def nvmlSystemGetProcessName(pid): _nvmlCheckReturn(ret) return c_name.value + @convertStrBytes def nvmlSystemGetDriverVersion(): c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) @@ -2490,6 +2978,7 @@ def nvmlSystemGetDriverVersion(): _nvmlCheckReturn(ret) return c_version.value + # Added in 2.285 def nvmlSystemGetHicVersion(): c_count = c_uint(0) @@ -2500,12 +2989,11 @@ def nvmlSystemGetHicVersion(): ret = fn(byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # If there are no hics - if (c_count.value == 0): + if c_count.value == 0: return [] hic_array = c_nvmlHwbcEntry_t * c_count.value @@ -2514,14 +3002,16 @@ def nvmlSystemGetHicVersion(): _nvmlCheckReturn(ret) return hics + def nvmlSystemGetDriverBranch(): c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) c_branchInfo.version = SystemDriverBranchInfo_v1 - fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_branchInfo + ## Unit get functions def nvmlUnitGetCount(): c_count = c_uint() @@ -2530,6 +3020,7 @@ def nvmlUnitGetCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlUnitGetHandleByIndex(index): c_index = c_uint(index) unit = c_nvmlUnit_t() @@ -2538,6 +3029,7 @@ def nvmlUnitGetHandleByIndex(index): _nvmlCheckReturn(ret) return unit + def nvmlUnitGetUnitInfo(unit): c_info = c_nvmlUnitInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") @@ -2545,13 +3037,15 @@ def nvmlUnitGetUnitInfo(unit): _nvmlCheckReturn(ret) return c_info + def nvmlUnitGetLedState(unit): - c_state = c_nvmlLedState_t() + c_state = c_nvmlLedState_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") ret = fn(unit, byref(c_state)) _nvmlCheckReturn(ret) return c_state + def nvmlUnitGetPsuInfo(unit): c_info = c_nvmlPSUInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") @@ -2559,6 +3053,7 @@ def nvmlUnitGetPsuInfo(unit): _nvmlCheckReturn(ret) return c_info + def nvmlUnitGetTemperature(unit, type): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") @@ -2566,6 +3061,7 @@ def nvmlUnitGetTemperature(unit, type): _nvmlCheckReturn(ret) return c_temp.value + def nvmlUnitGetFanSpeedInfo(unit): c_speeds = c_nvmlUnitFanSpeeds_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") @@ -2573,17 +3069,19 @@ def nvmlUnitGetFanSpeedInfo(unit): _nvmlCheckReturn(ret) return c_speeds + # added to API def nvmlUnitGetDeviceCount(unit): c_count = c_uint(0) # query the unit to determine device count fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") ret = fn(unit, byref(c_count), None) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = NVML_SUCCESS _nvmlCheckReturn(ret) return c_count.value + def nvmlUnitGetDevices(unit): c_count = c_uint(nvmlUnitGetDeviceCount(unit)) device_array = c_nvmlDevice_t * c_count.value @@ -2593,6 +3091,7 @@ def nvmlUnitGetDevices(unit): _nvmlCheckReturn(ret) return c_devices + ## Device get functions def nvmlDeviceGetCount(): c_count = c_uint() @@ -2601,6 +3100,7 @@ def nvmlDeviceGetCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetHandleByIndex(index): c_index = c_uint(index) device = c_nvmlDevice_t() @@ -2609,6 +3109,7 @@ def nvmlDeviceGetHandleByIndex(index): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleBySerial(serial): c_serial = c_char_p(serial) @@ -2618,6 +3119,7 @@ def nvmlDeviceGetHandleBySerial(serial): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleByUUID(uuid): c_uuid = c_char_p(uuid) @@ -2627,6 +3129,7 @@ def nvmlDeviceGetHandleByUUID(uuid): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleByPciBusId(pciBusId): c_busId = c_char_p(pciBusId) @@ -2636,6 +3139,7 @@ def nvmlDeviceGetHandleByPciBusId(pciBusId): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetName(handle): c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) @@ -2644,14 +3148,17 @@ def nvmlDeviceGetName(handle): _nvmlCheckReturn(ret) return c_name.value + class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), ] + nvmlDevicePerfModes_v1 = 0x1000804 + @convertStrBytes def nvmlDeviceGetPerformanceModes(handle): perfModes = c_nvmlDevicePerfModes_v1_t() @@ -2661,14 +3168,17 @@ def nvmlDeviceGetPerformanceModes(handle): _nvmlCheckReturn(ret) return perfModes.str + class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), ] + nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + @convertStrBytes def nvmlDeviceGetCurrentClockFreqs(handle): currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() @@ -2678,20 +3188,23 @@ def nvmlDeviceGetCurrentClockFreqs(handle): _nvmlCheckReturn(ret) return currentClockFreqs.str + def nvmlDeviceGetBoardId(handle): - c_id = c_uint(); + c_id = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") ret = fn(handle, byref(c_id)) _nvmlCheckReturn(ret) return c_id.value + def nvmlDeviceGetMultiGpuBoard(handle): - c_multiGpu = c_uint(); + c_multiGpu = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") ret = fn(handle, byref(c_multiGpu)) _nvmlCheckReturn(ret) return c_multiGpu.value + def nvmlDeviceGetBrand(handle): c_type = _nvmlBrandType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") @@ -2699,6 +3212,7 @@ def nvmlDeviceGetBrand(handle): _nvmlCheckReturn(ret) return c_type.value + def nvmlDeviceGetC2cModeInfoV1(handle): c_info = c_nvmlC2cModeInfo_v1_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") @@ -2706,9 +3220,11 @@ def nvmlDeviceGetC2cModeInfoV1(handle): _nvmlCheckReturn(ret) return c_info + def nvmlDeviceGetC2cModeInfoV(handle): return nvmlDeviceGetC2cModeInfoV1(handle) + @convertStrBytes def nvmlDeviceGetBoardPartNumber(handle): c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) @@ -2717,6 +3233,7 @@ def nvmlDeviceGetBoardPartNumber(handle): _nvmlCheckReturn(ret) return c_part_number.value + @convertStrBytes def nvmlDeviceGetSerial(handle): c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) @@ -2725,6 +3242,7 @@ def nvmlDeviceGetSerial(handle): _nvmlCheckReturn(ret) return c_serial.value + def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): isReference = type(moduleId) is not c_uint moduleIdRef = moduleId if isReference else byref(moduleId) @@ -2736,6 +3254,7 @@ def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): _nvmlCheckReturn(ret) return moduleId.value + def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): affinity_array = c_ulonglong * nodeSetSize c_affinity = affinity_array() @@ -2744,6 +3263,7 @@ def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() @@ -2752,6 +3272,7 @@ def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() @@ -2760,18 +3281,21 @@ def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceSetCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNumaNodeId(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") node = c_int() @@ -2779,6 +3303,7 @@ def nvmlDeviceGetNumaNodeId(handle): _nvmlCheckReturn(ret) return node.value + def nvmlDeviceGetMinorNumber(handle): c_minor_number = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") @@ -2786,6 +3311,7 @@ def nvmlDeviceGetMinorNumber(handle): _nvmlCheckReturn(ret) return c_minor_number.value + @convertStrBytes def nvmlDeviceGetUUID(handle): c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) @@ -2794,15 +3320,21 @@ def nvmlDeviceGetUUID(handle): _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlDeviceGetInforomVersion(handle, infoRomObject): c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") - ret = fn(handle, _nvmlInforomObject_t(infoRomObject), - c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + ret = fn( + handle, + _nvmlInforomObject_t(infoRomObject), + c_version, + c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE), + ) _nvmlCheckReturn(ret) return c_version.value + # Added in 4.304 @convertStrBytes def nvmlDeviceGetInforomImageVersion(handle): @@ -2812,6 +3344,7 @@ def nvmlDeviceGetInforomImageVersion(handle): _nvmlCheckReturn(ret) return c_version.value + # Added in 4.304 def nvmlDeviceGetInforomConfigurationChecksum(handle): c_checksum = c_uint() @@ -2820,6 +3353,7 @@ def nvmlDeviceGetInforomConfigurationChecksum(handle): _nvmlCheckReturn(ret) return c_checksum.value + # Added in 4.304 def nvmlDeviceValidateInforom(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") @@ -2827,6 +3361,7 @@ def nvmlDeviceValidateInforom(handle): _nvmlCheckReturn(ret) return None + def nvmlDeviceGetLastBBXFlushTime(handle): c_timestamp = c_ulonglong() c_durationUs = c_ulong() @@ -2835,6 +3370,7 @@ def nvmlDeviceGetLastBBXFlushTime(handle): _nvmlCheckReturn(ret) return [c_timestamp.value, c_durationUs.value] + def nvmlDeviceGetDisplayMode(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") @@ -2842,6 +3378,7 @@ def nvmlDeviceGetDisplayMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceGetDisplayActive(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") @@ -2857,12 +3394,14 @@ def nvmlDeviceGetPersistenceMode(handle): _nvmlCheckReturn(ret) return c_state.value + def nvmlDeviceGetPciInfoExt(handle, c_info): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") ret = fn(handle, c_info) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetPciInfo_v3(handle): c_info = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") @@ -2870,9 +3409,11 @@ def nvmlDeviceGetPciInfo_v3(handle): _nvmlCheckReturn(ret) return c_info + def nvmlDeviceGetPciInfo(handle): return nvmlDeviceGetPciInfo_v3(handle) + def nvmlDeviceGetClockInfo(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") @@ -2880,6 +3421,7 @@ def nvmlDeviceGetClockInfo(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 2.285 def nvmlDeviceGetMaxClockInfo(handle, type): c_clock = c_uint() @@ -2888,6 +3430,7 @@ def nvmlDeviceGetMaxClockInfo(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 4.304 def nvmlDeviceGetApplicationsClock(handle, type): c_clock = c_uint() @@ -2896,6 +3439,7 @@ def nvmlDeviceGetApplicationsClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + def nvmlDeviceGetMaxCustomerBoostClock(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") @@ -2903,6 +3447,7 @@ def nvmlDeviceGetMaxCustomerBoostClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + def nvmlDeviceGetClock(handle, type, id): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") @@ -2910,6 +3455,7 @@ def nvmlDeviceGetClock(handle, type, id): _nvmlCheckReturn(ret) return c_clock.value + # Added in 5.319 def nvmlDeviceGetDefaultApplicationsClock(handle, type): c_clock = c_uint() @@ -2918,6 +3464,7 @@ def nvmlDeviceGetDefaultApplicationsClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 4.304 def nvmlDeviceGetSupportedMemoryClocks(handle): # first call to get the size @@ -2925,10 +3472,10 @@ def nvmlDeviceGetSupportedMemoryClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no clocks return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() @@ -2946,6 +3493,7 @@ def nvmlDeviceGetSupportedMemoryClocks(handle): # error case raise NVMLError(ret) + # Added in 4.304 def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # first call to get the size @@ -2953,10 +3501,10 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no clocks return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() @@ -2974,6 +3522,7 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # error case raise NVMLError(ret) + def nvmlDeviceGetFanSpeed(handle): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") @@ -2981,6 +3530,7 @@ def nvmlDeviceGetFanSpeed(handle): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetFanSpeed_v2(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") @@ -2988,15 +3538,18 @@ def nvmlDeviceGetFanSpeed_v2(handle, fan): _nvmlCheckReturn(ret) return c_speed.value + class c_nvmlFanSpeedInfo_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('fan', c_uint), - ('speed', c_uint), + ("version", c_uint), + ("fan", c_uint), + ("speed", c_uint), ] + nvmlFanSpeedInfo_v1 = 0x100000C + def nvmlDeviceGetFanSpeedRPM(handle): c_fanSpeed = c_nvmlFanSpeedInfo_t() c_fanSpeed.fan = 0 @@ -3006,6 +3559,7 @@ def nvmlDeviceGetFanSpeedRPM(handle): _nvmlCheckReturn(ret) return c_fanSpeed.speed + def nvmlDeviceGetTargetFanSpeed(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") @@ -3013,6 +3567,7 @@ def nvmlDeviceGetTargetFanSpeed(handle, fan): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetNumFans(device): c_numFans = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") @@ -3020,12 +3575,14 @@ def nvmlDeviceGetNumFans(device): _nvmlCheckReturn(ret) return c_numFans.value + def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2") ret = fn(handle, index) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) minSpeedRef = minSpeed if isReference else byref(minSpeed) @@ -3035,6 +3592,7 @@ def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): isReference = type(fanControlPolicy) is not c_uint fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) @@ -3043,35 +3601,42 @@ def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else fanControlPolicy.value + def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) _nvmlCheckReturn(ret) return NVML_SUCCESS + class c_nvmlTemperature_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('sensorType', _nvmlTemperatureSensors_t), - ('temperature', c_int), + ("version", c_uint), + ("sensorType", _nvmlTemperatureSensors_t), + ("temperature", c_int), ] + + nvmlTemperature_v1 = 0x100000C + def nvmlDeviceGetTemperatureV1(handle, sensor): c_temp = c_nvmlTemperature_v1_t() c_temp.version = nvmlTemperature_v1 - c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") ret = fn(handle, byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.temperature + def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): if version == nvmlTemperature_v1: return nvmlDeviceGetTemperatureV1(handle, sensor) else: raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + # DEPRECATED use nvmlDeviceGetTemperatureV instead def nvmlDeviceGetTemperature(handle, sensor): c_temp = c_uint() @@ -3080,6 +3645,7 @@ def nvmlDeviceGetTemperature(handle, sensor): _nvmlCheckReturn(ret) return c_temp.value + def nvmlDeviceGetTemperatureThreshold(handle, threshold): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") @@ -3087,6 +3653,7 @@ def nvmlDeviceGetTemperatureThreshold(handle, threshold): _nvmlCheckReturn(ret) return c_temp.value + def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): c_temp = c_uint() c_temp.value = temp @@ -3095,6 +3662,7 @@ def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): _nvmlCheckReturn(ret) return None + def nvmlDeviceGetMarginTemperature(handle): c_marginTempInfo = c_nvmlMarginTemperature_v1_t() c_marginTempInfo.version = nvmlMarginTemperature_v1 @@ -3103,6 +3671,7 @@ def nvmlDeviceGetMarginTemperature(handle): _nvmlCheckReturn(ret) return c_marginTempInfo.marginTemperature + # DEPRECATED use nvmlDeviceGetPerformanceState def nvmlDeviceGetPowerState(handle): c_pstate = _nvmlPstates_t() @@ -3111,6 +3680,7 @@ def nvmlDeviceGetPowerState(handle): _nvmlCheckReturn(ret) return c_pstate.value + def nvmlDeviceGetPerformanceState(handle): c_pstate = _nvmlPstates_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") @@ -3118,6 +3688,7 @@ def nvmlDeviceGetPerformanceState(handle): _nvmlCheckReturn(ret) return c_pstate.value + def nvmlDeviceGetPowerManagementMode(handle): c_pcapMode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") @@ -3125,6 +3696,7 @@ def nvmlDeviceGetPowerManagementMode(handle): _nvmlCheckReturn(ret) return c_pcapMode.value + def nvmlDeviceGetPowerManagementLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") @@ -3132,6 +3704,7 @@ def nvmlDeviceGetPowerManagementLimit(handle): _nvmlCheckReturn(ret) return c_limit.value + # Added in 4.304 def nvmlDeviceGetPowerManagementLimitConstraints(handle): c_minLimit = c_uint() @@ -3141,6 +3714,7 @@ def nvmlDeviceGetPowerManagementLimitConstraints(handle): _nvmlCheckReturn(ret) return [c_minLimit.value, c_maxLimit.value] + # Added in 4.304 def nvmlDeviceGetPowerManagementDefaultLimit(handle): c_limit = c_uint() @@ -3158,6 +3732,7 @@ def nvmlDeviceGetEnforcedPowerLimit(handle): _nvmlCheckReturn(ret) return c_limit.value + def nvmlDeviceGetPowerUsage(handle): c_watts = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") @@ -3165,6 +3740,7 @@ def nvmlDeviceGetPowerUsage(handle): _nvmlCheckReturn(ret) return c_watts.value + def nvmlDeviceGetTotalEnergyConsumption(handle): c_millijoules = c_uint64() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") @@ -3172,6 +3748,7 @@ def nvmlDeviceGetTotalEnergyConsumption(handle): _nvmlCheckReturn(ret) return c_millijoules.value + # Added in 4.304 def nvmlDeviceGetGpuOperationMode(handle): c_currState = _nvmlGpuOperationMode_t() @@ -3181,14 +3758,17 @@ def nvmlDeviceGetGpuOperationMode(handle): _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] + # Added in 4.304 def nvmlDeviceGetCurrentGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[0] + # Added in 4.304 def nvmlDeviceGetPendingGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[1] + def nvmlDeviceGetMemoryInfo(handle, version=None): if not version: c_memory = c_nvmlMemory_t() @@ -3201,6 +3781,7 @@ def nvmlDeviceGetMemoryInfo(handle, version=None): _nvmlCheckReturn(ret) return c_memory + def nvmlDeviceGetBAR1MemoryInfo(handle): c_bar1_memory = c_nvmlBAR1Memory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") @@ -3208,6 +3789,7 @@ def nvmlDeviceGetBAR1MemoryInfo(handle): _nvmlCheckReturn(ret) return c_bar1_memory + def nvmlDeviceGetComputeMode(handle): c_mode = _nvmlComputeMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") @@ -3215,6 +3797,7 @@ def nvmlDeviceGetComputeMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceGetCudaComputeCapability(handle): c_major = c_int() c_minor = c_int() @@ -3223,6 +3806,7 @@ def nvmlDeviceGetCudaComputeCapability(handle): _nvmlCheckReturn(ret) return (c_major.value, c_minor.value) + def nvmlDeviceGetEccMode(handle): c_currState = _nvmlEnableState_t() c_pendingState = _nvmlEnableState_t() @@ -3231,14 +3815,17 @@ def nvmlDeviceGetEccMode(handle): _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] + # added to API def nvmlDeviceGetCurrentEccMode(handle): return nvmlDeviceGetEccMode(handle)[0] + # added to API def nvmlDeviceGetPendingEccMode(handle): return nvmlDeviceGetEccMode(handle)[1] + def nvmlDeviceGetDefaultEccMode(handle): c_defaultState = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") @@ -3246,35 +3833,49 @@ def nvmlDeviceGetDefaultEccMode(handle): _nvmlCheckReturn(ret) return [c_defaultState.value] + def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_count)) + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_count), + ) _nvmlCheckReturn(ret) return c_count.value + # This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): c_counts = c_nvmlEccErrorCounts_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_counts)) + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_counts), + ) _nvmlCheckReturn(ret) return c_counts + # Added in 4.304 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") - ret = fn(handle, - _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), - _nvmlMemoryLocation_t(locationType), - byref(c_count)) + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), + byref(c_count), + ) _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetUtilizationRates(handle): c_util = c_nvmlUtilization_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") @@ -3282,6 +3883,7 @@ def nvmlDeviceGetUtilizationRates(handle): _nvmlCheckReturn(ret) return c_util + def nvmlDeviceGetEncoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3290,6 +3892,7 @@ def nvmlDeviceGetEncoderUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetDecoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3298,6 +3901,7 @@ def nvmlDeviceGetDecoderUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetJpgUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3306,6 +3910,7 @@ def nvmlDeviceGetJpgUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetOfaUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3314,6 +3919,7 @@ def nvmlDeviceGetOfaUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetPcieReplayCounter(handle): c_replay = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") @@ -3321,6 +3927,7 @@ def nvmlDeviceGetPcieReplayCounter(handle): _nvmlCheckReturn(ret) return c_replay.value + def nvmlDeviceGetDriverModel(handle): c_currModel = _nvmlDriverModel_t() c_pendingModel = _nvmlDriverModel_t() @@ -3329,14 +3936,17 @@ def nvmlDeviceGetDriverModel(handle): _nvmlCheckReturn(ret) return [c_currModel.value, c_pendingModel.value] + # added to API def nvmlDeviceGetCurrentDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[0] + # added to API def nvmlDeviceGetPendingDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[1] + # Added in 2.285 @convertStrBytes def nvmlDeviceGetVbiosVersion(handle): @@ -3346,16 +3956,17 @@ def nvmlDeviceGetVbiosVersion(handle): _nvmlCheckReturn(ret) return c_version.value + # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array incase more processes are created c_count.value = c_count.value * 2 + 5 @@ -3368,7 +3979,7 @@ def nvmlDeviceGetComputeRunningProcesses_v2(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3377,6 +3988,7 @@ def nvmlDeviceGetComputeRunningProcesses_v2(handle): # error case raise NVMLError(ret) + # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v3(handle): # first call to get the size @@ -3384,10 +3996,10 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array incase more processes are created c_count.value = c_count.value * 2 + 5 @@ -3402,7 +4014,7 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3412,19 +4024,21 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle): # error case raise NVMLError(ret) + @throwOnVersionMismatch def nvmlDeviceGetComputeRunningProcesses(handle): return nvmlDeviceGetComputeRunningProcesses_v3(handle) + def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array incase more processes are created c_count.value = c_count.value * 2 + 5 @@ -3437,7 +4051,7 @@ def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3446,16 +4060,17 @@ def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array incase more processes are created c_count.value = c_count.value * 2 + 5 @@ -3470,7 +4085,7 @@ def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3480,24 +4095,27 @@ def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): # error case raise NVMLError(ret) + @throwOnVersionMismatch def nvmlDeviceGetGraphicsRunningProcesses(handle): return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + @throwOnVersionMismatch def nvmlDeviceGetMPSComputeRunningProcesses(handle): return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array incase more processes are created c_count.value = c_count.value * 2 + 5 @@ -3512,7 +4130,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3522,16 +4140,17 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") ret = fn(handle, byref(c_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 @@ -3546,7 +4165,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) @@ -3556,6 +4175,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): c_processDetailList = c_nvmlProcessDetailList_t() c_processDetailList.version = version @@ -3565,12 +4185,14 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): # first call to get the size ret = fn(handle, byref(c_processDetailList)) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no running processes return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries - c_processDetailList.procArray = cast((c_procs)(), POINTER(c_nvmlProcessDetail_v1_t)) + c_processDetailList.procArray = cast( + (c_procs)(), POINTER(c_nvmlProcessDetail_v1_t) + ) # make the call again ret = fn(handle, byref(c_processDetailList)) @@ -3580,9 +4202,9 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): for i in range(c_processDetailList.numProcArrayEntries): # use an alternative struct for this object obj = c_processDetailList.procArray[i] - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: obj.usedGpuMemory = None - if (obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: obj.usedGpuCcProtectedMemory = None procs.append(obj) @@ -3591,6 +4213,7 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): # error case raise NVMLError(ret) + def nvmlDeviceGetAutoBoostedClocksEnabled(handle): c_isEnabled = _nvmlEnableState_t() c_defaultIsEnabled = _nvmlEnableState_t() @@ -3598,7 +4221,8 @@ def nvmlDeviceGetAutoBoostedClocksEnabled(handle): ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) _nvmlCheckReturn(ret) return [c_isEnabled.value, c_defaultIsEnabled.value] - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + ## Set functions def nvmlUnitSetLedState(unit, color): @@ -3607,49 +4231,57 @@ def nvmlUnitSetLedState(unit, color): _nvmlCheckReturn(ret) return None + def nvmlDeviceSetPersistenceMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetComputeMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") ret = fn(handle, _nvmlComputeMode_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetEccMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearEccErrorCounts(handle, counterType): fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") ret = fn(handle, _nvmlEccCounterType_t(counterType)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetDriverModel(handle, model): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") ret = fn(handle, _nvmlDriverModel_t(model)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled)) _nvmlCheckReturn(ret) return None - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) _nvmlCheckReturn(ret) return None - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") @@ -3657,24 +4289,28 @@ def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): _nvmlCheckReturn(ret) return None + def nvmlDeviceResetGpuLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) _nvmlCheckReturn(ret) return None + def nvmlDeviceResetMemoryLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) @@ -3683,6 +4319,7 @@ def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_clkMonInfo + # Added in 4.304 def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") @@ -3690,6 +4327,7 @@ def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz) _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceResetApplicationsClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") @@ -3697,6 +4335,7 @@ def nvmlDeviceResetApplicationsClocks(handle): _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceSetPowerManagementLimit(handle, limit): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") @@ -3704,6 +4343,7 @@ def nvmlDeviceSetPowerManagementLimit(handle, limit): _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceSetGpuOperationMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") @@ -3711,6 +4351,7 @@ def nvmlDeviceSetGpuOperationMode(handle, mode): _nvmlCheckReturn(ret) return None + # Added in 2.285 def nvmlEventSetCreate(): fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") @@ -3719,6 +4360,7 @@ def nvmlEventSetCreate(): _nvmlCheckReturn(ret) return eventSet + # Added in 2.285 def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") @@ -3726,6 +4368,7 @@ def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): _nvmlCheckReturn(ret) return None + # Added in 2.285 def nvmlDeviceGetSupportedEventTypes(handle): c_eventTypes = c_ulonglong() @@ -3734,6 +4377,7 @@ def nvmlDeviceGetSupportedEventTypes(handle): _nvmlCheckReturn(ret) return c_eventTypes.value + # raises NVML_ERROR_TIMEOUT exception on timeout def nvmlEventSetWait_v2(eventSet, timeoutms): fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") @@ -3742,9 +4386,11 @@ def nvmlEventSetWait_v2(eventSet, timeoutms): _nvmlCheckReturn(ret) return data + def nvmlEventSetWait(eventSet, timeoutms): return nvmlEventSetWait_v2(eventSet, timeoutms) + # Added in 2.285 def nvmlEventSetFree(eventSet): fn = _nvmlGetFunctionPointer("nvmlEventSetFree") @@ -3752,13 +4398,15 @@ def nvmlEventSetFree(eventSet): _nvmlCheckReturn(ret) return None + # Added in 3.295 def nvmlDeviceOnSameBoard(handle1, handle2): fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") onSameBoard = c_int() ret = fn(handle1, handle2, byref(onSameBoard)) _nvmlCheckReturn(ret) - return (onSameBoard.value != 0) + return onSameBoard.value != 0 + # Added in 3.295 def nvmlDeviceGetCurrPcieLinkGeneration(handle): @@ -3768,6 +4416,7 @@ def nvmlDeviceGetCurrPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 3.295 def nvmlDeviceGetMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") @@ -3776,6 +4425,7 @@ def nvmlDeviceGetMaxPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 3.295 def nvmlDeviceGetCurrPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") @@ -3784,6 +4434,7 @@ def nvmlDeviceGetCurrPcieLinkWidth(handle): _nvmlCheckReturn(ret) return width.value + # Added in 3.295 def nvmlDeviceGetMaxPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") @@ -3792,6 +4443,7 @@ def nvmlDeviceGetMaxPcieLinkWidth(handle): _nvmlCheckReturn(ret) return width.value + def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") gen = c_uint() @@ -3799,36 +4451,41 @@ def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 4.304 def nvmlDeviceGetSupportedClocksThrottleReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + def nvmlDeviceGetSupportedClocksEventReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + # Added in 4.304 def nvmlDeviceGetCurrentClocksThrottleReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + def nvmlDeviceGetCurrentClocksEventReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + # Added in 5.319 def nvmlDeviceGetIndex(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") @@ -3837,6 +4494,7 @@ def nvmlDeviceGetIndex(handle): _nvmlCheckReturn(ret) return c_index.value + # Added in 5.319 def nvmlDeviceGetAccountingMode(handle): c_mode = _nvmlEnableState_t() @@ -3845,35 +4503,40 @@ def nvmlDeviceGetAccountingMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceSetAccountingMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearAccountingPids(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetAccountingStats(handle, pid): stats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") ret = fn(handle, c_uint(pid), byref(stats)) _nvmlCheckReturn(ret) - if (stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + if stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above stats.maxMemoryUsage = None return stats + def nvmlDeviceGetAccountingPids(handle): count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) pids = (c_uint * count.value)() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") ret = fn(handle, byref(count), pids) _nvmlCheckReturn(ret) - return list(map(int, pids[0:count.value])) + return list(map(int, pids[0 : count.value])) + def nvmlDeviceGetAccountingBufferSize(handle): bufferSize = c_uint() @@ -3882,6 +4545,7 @@ def nvmlDeviceGetAccountingBufferSize(handle): _nvmlCheckReturn(ret) return int(bufferSize.value) + def nvmlDeviceGetRetiredPages(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) @@ -3891,8 +4555,7 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter): ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # call again with a buffer @@ -3903,7 +4566,8 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter): c_pages = page_array() ret = fn(device, c_source, byref(c_count), c_pages) _nvmlCheckReturn(ret) - return list(map(int, c_pages[0:c_count.value])) + return list(map(int, c_pages[0 : c_count.value])) + def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) @@ -3914,8 +4578,7 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # call again with a buffer @@ -3928,7 +4591,11 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): c_times = times_array() ret = fn(device, c_source, byref(c_count), c_pages, c_times) _nvmlCheckReturn(ret) - return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ]; + return [ + {"address": int(c_pages[i]), "timestamp": int(c_times[i])} + for i in range(c_count.value) + ] + def nvmlDeviceGetRetiredPagesPendingStatus(device): c_pending = _nvmlEnableState_t() @@ -3937,6 +4604,7 @@ def nvmlDeviceGetRetiredPagesPendingStatus(device): _nvmlCheckReturn(ret) return int(c_pending.value) + def nvmlDeviceGetAPIRestriction(device, apiType): c_permission = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") @@ -3944,12 +4612,14 @@ def nvmlDeviceGetAPIRestriction(device, apiType): _nvmlCheckReturn(ret) return int(c_permission.value) + def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetBridgeChipInfo(handle): bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") @@ -3957,6 +4627,7 @@ def nvmlDeviceGetBridgeChipInfo(handle): _nvmlCheckReturn(ret) return bridgeHierarchy + def nvmlDeviceGetSamples(device, sampling_type, timeStamp): c_sampling_type = _nvmlSamplingType_t(sampling_type) c_time_stamp = c_ulonglong(timeStamp) @@ -3965,17 +4636,32 @@ def nvmlDeviceGetSamples(device, sampling_type, timeStamp): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") ## First Call gets the size - ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None) + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + None, + ) # Stop if this fails - if (ret != NVML_SUCCESS): + if ret != NVML_SUCCESS: raise NVMLError(ret) sampleArray = c_sample_count.value * c_nvmlSample_t c_samples = sampleArray() - ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + c_samples, + ) _nvmlCheckReturn(ret) - return (c_sample_value_type.value, c_samples[0:c_sample_count.value]) + return (c_sample_value_type.value, c_samples[0 : c_sample_count.value]) + def nvmlDeviceGetViolationStatus(device, perfPolicyType): c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) @@ -3987,6 +4673,7 @@ def nvmlDeviceGetViolationStatus(device, perfPolicyType): _nvmlCheckReturn(ret) return c_violTime + def nvmlDeviceGetPcieThroughput(device, counter): c_util = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") @@ -3994,6 +4681,7 @@ def nvmlDeviceGetPcieThroughput(device, counter): _nvmlCheckReturn(ret) return c_util.value + def nvmlSystemGetTopologyGpuSet(cpuNumber): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") @@ -4008,7 +4696,8 @@ def nvmlSystemGetTopologyGpuSet(cpuNumber): c_devices = device_array() ret = fn(cpuNumber, byref(c_count), c_devices) _nvmlCheckReturn(ret) - return list(c_devices[0:c_count.value]) + return list(c_devices[0 : c_count.value]) + def nvmlDeviceGetTopologyNearestGpus(device, level): c_count = c_uint(0) @@ -4025,7 +4714,8 @@ def nvmlDeviceGetTopologyNearestGpus(device, level): c_devices = device_array() ret = fn(device, level, byref(c_count), c_devices) _nvmlCheckReturn(ret) - return list(c_devices[0:c_count.value]) + return list(c_devices[0 : c_count.value]) + def nvmlDeviceGetTopologyCommonAncestor(device1, device2): c_level = _nvmlGpuTopologyLevel_t() @@ -4034,6 +4724,7 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2): _nvmlCheckReturn(ret) return c_level.value + def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): c_rxcounter = c_ulonglong() c_txcounter = c_ulonglong() @@ -4042,24 +4733,28 @@ def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): _nvmlCheckReturn(ret) return (c_rxcounter.value, c_txcounter.value) + def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") ret = fn(device, link, counter, freeze) _nvmlCheckReturn(ret) return None + def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") ret = fn(device, link, counter) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") ret = fn(device, link, counter, byref(control), reset) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): c_control = nvmlNvLinkUtilizationControl_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") @@ -4067,6 +4762,7 @@ def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): _nvmlCheckReturn(ret) return c_control + def nvmlDeviceGetNvLinkCapability(device, link, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") @@ -4074,6 +4770,7 @@ def nvmlDeviceGetNvLinkCapability(device, link, capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): c_result = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") @@ -4081,12 +4778,14 @@ def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): _nvmlCheckReturn(ret) return c_result.value + def nvmlDeviceResetNvLinkErrorCounters(device, link): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") ret = fn(device, link) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNvLinkRemotePciInfo(device, link): c_pci = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") @@ -4094,6 +4793,7 @@ def nvmlDeviceGetNvLinkRemotePciInfo(device, link): _nvmlCheckReturn(ret) return c_pci + def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): c_type = _nvmlNvLinkDeviceType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") @@ -4101,6 +4801,7 @@ def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): _nvmlCheckReturn(ret) return c_type.value + def nvmlDeviceGetNvLinkState(device, link): c_isActive = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") @@ -4108,6 +4809,7 @@ def nvmlDeviceGetNvLinkState(device, link): _nvmlCheckReturn(ret) return c_isActive.value + def nvmlDeviceGetNvLinkVersion(device, link): c_version = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") @@ -4115,12 +4817,14 @@ def nvmlDeviceGetNvLinkVersion(device, link): _nvmlCheckReturn(ret) return c_version.value + def nvmlDeviceModifyDrainState(pciInfo, newState): fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") ret = fn(pointer(pciInfo), newState) _nvmlCheckReturn(ret) return None + def nvmlDeviceQueryDrainState(pciInfo): c_newState = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") @@ -4128,18 +4832,21 @@ def nvmlDeviceQueryDrainState(pciInfo): _nvmlCheckReturn(ret) return c_newState.value + def nvmlDeviceRemoveGpu(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None + def nvmlDeviceDiscoverGpus(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() @@ -4155,6 +4862,7 @@ def nvmlDeviceGetFieldValues(handle, fieldIds): _nvmlCheckReturn(ret) return values + def nvmlDeviceClearFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() @@ -4170,6 +4878,7 @@ def nvmlDeviceClearFieldValues(handle, fieldIds): _nvmlCheckReturn(ret) return values + def nvmlDeviceGetVirtualizationMode(handle): c_virtualization_mode = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") @@ -4177,10 +4886,12 @@ def nvmlDeviceGetVirtualizationMode(handle): _nvmlCheckReturn(ret) return c_virtualization_mode.value + def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") return fn(handle, virtualization_mode) + def nvmlDeviceGetVgpuHeterogeneousMode(handle): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 @@ -4189,6 +4900,7 @@ def nvmlDeviceGetVgpuHeterogeneousMode(handle): _nvmlCheckReturn(ret) return c_vgpuHeterogeneousMode.mode + def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 @@ -4198,6 +4910,7 @@ def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlVgpuInstanceGetPlacementId(vgpuInstance): c_placement = c_nvmlVgpuPlacementId_v1_t(0) c_placement.version = VgpuPlacementId_v1 @@ -4206,9 +4919,10 @@ def nvmlVgpuInstanceGetPlacementId(vgpuInstance): _nvmlCheckReturn(ret) return c_placement.placementId + def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) @@ -4225,14 +4939,15 @@ def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements + def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) @@ -4246,11 +4961,12 @@ def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements + def nvmlGetVgpuDriverCapabilities(capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") @@ -4258,6 +4974,7 @@ def nvmlGetVgpuDriverCapabilities(capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceGetVgpuCapabilities(handle, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") @@ -4265,23 +4982,25 @@ def nvmlDeviceGetVgpuCapabilities(handle, capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceSetVgpuCapabilities(handle, capability, state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetSupportedVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") ret = fn(handle, byref(c_vgpu_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no supported vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value c_vgpu_type_ids = vgpu_type_ids_array() @@ -4297,17 +5016,18 @@ def nvmlDeviceGetSupportedVgpus(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetCreatableVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") ret = fn(handle, byref(c_vgpu_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no supported vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value c_vgpu_type_ids = vgpu_type_ids_array() @@ -4323,132 +5043,148 @@ def nvmlDeviceGetCreatableVgpus(handle): # error case raise NVMLError(ret) + def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): c_profile_id = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") ret = fn(vgpuTypeId, byref(c_profile_id)) _nvmlCheckReturn(ret) - return (c_profile_id.value) + return c_profile_id.value + @convertStrBytes def nvmlVgpuTypeGetClass(vgpuTypeId): c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_class.value + @convertStrBytes def nvmlVgpuTypeGetName(vgpuTypeId): c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_name.value + def nvmlVgpuTypeGetDeviceID(vgpuTypeId): - c_device_id = c_ulonglong(0) + c_device_id = c_ulonglong(0) c_subsystem_id = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) _nvmlCheckReturn(ret) return (c_device_id.value, c_subsystem_id.value) + def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): c_fb_size = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") ret = fn(vgpuTypeId, byref(c_fb_size)) _nvmlCheckReturn(ret) return c_fb_size.value + def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): c_num_heads = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") ret = fn(vgpuTypeId, byref(c_num_heads)) _nvmlCheckReturn(ret) return c_num_heads.value + def nvmlVgpuTypeGetResolution(vgpuTypeId): c_xdim = c_uint(0) c_ydim = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) _nvmlCheckReturn(ret) return (c_xdim.value, c_ydim.value) + @convertStrBytes def nvmlVgpuTypeGetLicense(vgpuTypeId): c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") ret = fn(vgpuTypeId, c_license, c_buffer_size) _nvmlCheckReturn(ret) return c_license.value + def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): c_frl_config = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") ret = fn(vgpuTypeId, byref(c_frl_config)) _nvmlCheckReturn(ret) return c_frl_config.value + def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): c_gsp_heap = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") ret = fn(vgpuTypeId, byref(c_gsp_heap)) _nvmlCheckReturn(ret) return c_gsp_heap.value + def nvmlVgpuTypeGetFbReservation(vgpuTypeId): c_fb_reservation = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") ret = fn(vgpuTypeId, byref(c_fb_reservation)) _nvmlCheckReturn(ret) return c_fb_reservation.value + def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): c_runtime_state = nvmlVgpuRuntimeState_v1_t() c_runtime_state.version = VgpuRuntimeState_v1 - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") ret = fn(vgpuInstance, byref(c_runtime_state)) _nvmlCheckReturn(ret) return c_runtime_state + def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) return c_max_instances.value + def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): c_max_instances_per_vm = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) _nvmlCheckReturn(ret) return c_max_instances_per_vm.value + def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) c_bar1Info.version = VgpuTypeBar1Info_v1 - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") ret = fn(vgpuTypeId, byref(c_bar1Info)) _nvmlCheckReturn(ret) return c_bar1Info + def nvmlDeviceGetActiveVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") ret = fn(handle, byref(c_vgpu_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no active vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value c_vgpu_instances = vgpu_instance_array() @@ -4464,67 +5200,76 @@ def nvmlDeviceGetActiveVgpus(handle): # error case raise NVMLError(ret) + @convertStrBytes def nvmlVgpuInstanceGetVmID(vgpuInstance): c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) - c_vm_id_type = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) _nvmlCheckReturn(ret) return (c_vm_id.value, c_vm_id_type.value) + @convertStrBytes def nvmlVgpuInstanceGetUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) _nvmlCheckReturn(ret) return c_driver_version.value + def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): c_license_status = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") ret = fn(vgpuInstance, byref(c_license_status)) _nvmlCheckReturn(ret) return c_license_status.value + def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") c_license_info = c_nvmlVgpuLicenseInfo_t() ret = fn(vgpuInstance, byref(c_license_info)) _nvmlCheckReturn(ret) return c_license_info + def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): c_frl = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") ret = fn(vgpuInstance, byref(c_frl)) _nvmlCheckReturn(ret) return c_frl.value + def nvmlVgpuInstanceGetEccMode(vgpuInstance): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") @@ -4532,79 +5277,98 @@ def nvmlVgpuInstanceGetEccMode(vgpuInstance): _nvmlCheckReturn(ret) return c_mode.value + def nvmlVgpuInstanceGetType(vgpuInstance): c_vgpu_type = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") ret = fn(vgpuInstance, byref(c_vgpu_type)) _nvmlCheckReturn(ret) return c_vgpu_type.value + def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): c_encoder_capacity = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") ret = fn(vgpuInstance, byref(c_encoder_capacity)) _nvmlCheckReturn(ret) return c_encoder_capacity.value + def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") return fn(vgpuInstance, encoder_capacity) + def nvmlVgpuInstanceGetFbUsage(vgpuInstance): c_fb_usage = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") ret = fn(vgpuInstance, byref(c_fb_usage)) _nvmlCheckReturn(ret) return c_fb_usage.value + def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): c_cap_result = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) _nvmlCheckReturn(ret) - return (c_cap_result.value) + return c_cap_result.value + def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): c_id = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") ret = fn(vgpuInstance, byref(c_id)) _nvmlCheckReturn(ret) - return (c_id.value) + return c_id.value + @convertStrBytes def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") - ret = fn(vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE))) + ret = fn( + vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)) + ) _nvmlCheckReturn(ret) return c_vgpuPciId.value + def nvmlDeviceGetVgpuUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) c_sample_value_type = _nvmlValueType_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") - ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + ret = fn( + handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None + ) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no active vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t c_samples = sampleArray() # make the call again - ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples) + ret = fn( + handle, + c_time_stamp, + byref(c_sample_value_type), + byref(c_vgpu_count), + c_samples, + ) _nvmlCheckReturn(ret) - return c_samples[0:c_vgpu_count.value] + return c_samples[0 : c_vgpu_count.value] else: # error case raise NVMLError(ret) + def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4614,15 +5378,17 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") ret = fn(handle, byref(c_vgpuUtilInfo)) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no active vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case - sampleArray = c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t + sampleArray = ( + c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t + ) c_samples = sampleArray() c_vgpuUtilInfo.vgpuUtilArray = c_samples @@ -4630,29 +5396,33 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): ret = fn(handle, byref(c_vgpuUtilInfo)) _nvmlCheckReturn(ret) - return c_samples[0:c_vgpuUtilInfo.vgpuInstanceCount] + return c_samples[0 : c_vgpuUtilInfo.vgpuInstanceCount] else: # error case raise NVMLError(ret) + def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): c_p2pstatus = _nvmlGpuP2PStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") - ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus)) + ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus)) _nvmlCheckReturn(ret) return c_p2pstatus.value + def nvmlDeviceGetGridLicensableFeatures_v4(handle): c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") ret = fn(handle, byref(c_get_grid_licensable_features)) _nvmlCheckReturn(ret) - return (c_get_grid_licensable_features) + return c_get_grid_licensable_features + def nvmlDeviceGetGridLicensableFeatures(handle): return nvmlDeviceGetGridLicensableFeatures_v4(handle) + def nvmlDeviceGetGspFirmwareVersion(handle, version=None): isUserDefined = version is not None if not isUserDefined: @@ -4662,6 +5432,7 @@ def nvmlDeviceGetGspFirmwareVersion(handle, version=None): _nvmlCheckReturn(ret) return NVML_SUCCESS if isUserDefined else version.value + def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): isReference = type(isEnabled) is not c_uint isEnabledRef = isEnabled if isReference else byref(isEnabled) @@ -4671,6 +5442,7 @@ def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint( _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): c_encoder_capacity = c_ulonglong(0) c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) @@ -4680,18 +5452,19 @@ def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): _nvmlCheckReturn(ret) return c_encoder_capacity.value + def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no active vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t c_samples = sampleArray() @@ -4700,11 +5473,12 @@ def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) _nvmlCheckReturn(ret) - return c_samples[0:c_vgpu_count.value] + return c_samples[0 : c_vgpu_count.value] else: # error case raise NVMLError(ret) + def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4713,15 +5487,17 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") ret = fn(handle, byref(c_vgpuProcUtilInfo)) - if (ret == NVML_SUCCESS): + if ret == NVML_SUCCESS: # special case, no active vGPUs return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case - sampleArray = c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t + sampleArray = ( + c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t + ) c_samples = sampleArray() c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples @@ -4729,11 +5505,12 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): ret = fn(handle, byref(c_vgpuProcUtilInfo)) _nvmlCheckReturn(ret) - return c_samples[0:c_vgpuProcUtilInfo.vgpuProcessCount] + return c_samples[0 : c_vgpuProcUtilInfo.vgpuProcessCount] else: # error case raise NVMLError(ret) + def nvmlDeviceGetEncoderStats(handle): c_encoderCount = c_ulonglong(0) c_encodeFps = c_ulonglong(0) @@ -4743,15 +5520,16 @@ def nvmlDeviceGetEncoderStats(handle): _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + def nvmlDeviceGetEncoderSessions(handle): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") ret = fn(handle, byref(c_session_count), None) - if (ret == NVML_SUCCESS): - if (c_session_count.value != 0): + if ret == NVML_SUCCESS: + if c_session_count.value != 0: # typical case session_array = c_nvmlEncoderSession_t * c_session_count.value c_sessions = session_array() @@ -4769,6 +5547,7 @@ def nvmlDeviceGetEncoderSessions(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetFBCStats(handle): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") @@ -4776,15 +5555,16 @@ def nvmlDeviceGetFBCStats(handle): _nvmlCheckReturn(ret) return c_fbcStats + def nvmlDeviceGetFBCSessions(handle): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") ret = fn(handle, byref(c_session_count), None) - if (ret == NVML_SUCCESS): - if (c_session_count.value != 0): + if ret == NVML_SUCCESS: + if c_session_count.value != 0: # typical case session_array = c_nvmlFBCSession_t * c_session_count.value c_sessions = session_array() @@ -4802,24 +5582,28 @@ def nvmlDeviceGetFBCSessions(handle): # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): - c_encoderCount = c_ulonglong(0) - c_encodeFps = c_ulonglong(0) - c_encoderLatency = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") - ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + ret = fn( + vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency) + ) _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") ret = fn(vgpuInstance, byref(c_session_count), None) - if (ret == NVML_SUCCESS): - if (c_session_count.value != 0): + if ret == NVML_SUCCESS: + if c_session_count.value != 0: # typical case session_array = c_nvmlEncoderSession_t * c_session_count.value c_sessions = session_array() @@ -4837,6 +5621,7 @@ def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetFBCStats(vgpuInstance): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") @@ -4844,15 +5629,16 @@ def nvmlVgpuInstanceGetFBCStats(vgpuInstance): _nvmlCheckReturn(ret) return c_fbcStats + def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") ret = fn(vgpuInstance, byref(c_session_count), None) - if (ret == NVML_SUCCESS): - if (c_session_count.value != 0): + if ret == NVML_SUCCESS: + if c_session_count.value != 0: # typical case session_array = c_nvmlFBCSession_t * c_session_count.value c_sessions = session_array() @@ -4870,15 +5656,16 @@ def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): # error case raise NVMLError(ret) + def nvmlDeviceGetProcessUtilization(handle, timeStamp): # first call to get the size c_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") ret = fn(handle, None, byref(c_count), c_time_stamp) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t c_samples = sampleArray() @@ -4887,11 +5674,12 @@ def nvmlDeviceGetProcessUtilization(handle, timeStamp): ret = fn(handle, c_samples, byref(c_count), c_time_stamp) _nvmlCheckReturn(ret) - return c_samples[0:c_count.value] + return c_samples[0 : c_count.value] else: # error case raise NVMLError(ret) + def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4900,12 +5688,14 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): c_processesUtilInfo.processSamplesCount = c_uint(0) c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") ret = fn(handle, byref(c_processesUtilInfo)) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case - sampleArray = c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t + sampleArray = ( + c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t + ) c_samples = sampleArray() c_processesUtilInfo.procUtilArray = c_samples @@ -4913,11 +5703,12 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): ret = fn(handle, byref(c_processesUtilInfo)) _nvmlCheckReturn(ret) - return c_samples[0:c_processesUtilInfo.processSamplesCount] + return c_samples[0 : c_processesUtilInfo.processSamplesCount] else: # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetMetadata(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") c_vgpuMetadata = c_nvmlVgpuMetadata_t() @@ -4925,13 +5716,14 @@ def nvmlVgpuInstanceGetMetadata(vgpuInstance): # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return c_vgpuMetadata + def nvmlDeviceGetVgpuMetadata(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() @@ -4939,13 +5731,14 @@ def nvmlDeviceGetVgpuMetadata(handle): # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return c_vgpuPgpuMetadata + def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() @@ -4953,6 +5746,7 @@ def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): _nvmlCheckReturn(ret) return c_vgpuPgpuCompatibility + @convertStrBytes def nvmlDeviceGetPgpuMetadataString(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") @@ -4961,13 +5755,14 @@ def nvmlDeviceGetPgpuMetadataString(handle): # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return (c_pgpuMetadata.value, c_bufferSize.value) + def nvmlDeviceGetVgpuSchedulerLog(handle): c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") @@ -4975,6 +5770,7 @@ def nvmlDeviceGetVgpuSchedulerLog(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_log + def nvmlDeviceGetVgpuSchedulerState(handle): c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") @@ -4982,6 +5778,7 @@ def nvmlDeviceGetVgpuSchedulerState(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_state + def nvmlDeviceGetVgpuSchedulerCapabilities(handle): c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") @@ -4989,18 +5786,21 @@ def nvmlDeviceGetVgpuSchedulerCapabilities(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_caps + def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") ret = fn(handle, byref(sched_state)) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSetVgpuVersion(vgpuVersion): fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") ret = fn(byref(vgpuVersion)) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGetVgpuVersion(supported=None, current=None): isUserDefined = (supported is not None) or (current is not None) if not isUserDefined: @@ -5009,10 +5809,15 @@ def nvmlGetVgpuVersion(supported=None, current=None): fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") ret = fn(byref(supported), byref(current)) _nvmlCheckReturn(ret) - return NVML_SUCCESS if isUserDefined else [(supported.minVersion, - supported.maxVersion), - (current.minVersion, - current.maxVersion)] + return ( + NVML_SUCCESS + if isUserDefined + else [ + (supported.minVersion, supported.maxVersion), + (current.minVersion, current.maxVersion), + ] + ) + def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): c_mode = _nvmlEnableState_t() @@ -5021,11 +5826,12 @@ def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): _nvmlCheckReturn(ret) return c_mode.value + def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): c_pidCount = c_uint() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") ret = fn(vgpuInstance, byref(c_pidCount), None) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + if ret == NVML_ERROR_INSUFFICIENT_SIZE: sampleArray = c_pidCount.value * c_uint c_pidArray = sampleArray() ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) @@ -5034,6 +5840,7 @@ def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): raise NVMLError(ret) return (c_pidCount, c_pidArray) + def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): c_accountingStats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") @@ -5041,12 +5848,14 @@ def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): _nvmlCheckReturn(ret) return c_accountingStats + def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") ret = fn(vgpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGetExcludedDeviceCount(): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") @@ -5054,6 +5863,7 @@ def nvmlGetExcludedDeviceCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlGetExcludedDeviceInfoByIndex(index): c_index = c_uint(index) info = c_nvmlExcludedDeviceInfo_t() @@ -5062,6 +5872,7 @@ def nvmlGetExcludedDeviceInfoByIndex(index): _nvmlCheckReturn(ret) return info + def nvmlDeviceGetHostVgpuMode(handle): c_host_vgpu_mode = _nvmlHostVgpuMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") @@ -5069,6 +5880,7 @@ def nvmlDeviceGetHostVgpuMode(handle): _nvmlCheckReturn(ret) return c_host_vgpu_mode.value + def nvmlDeviceSetMigMode(device, mode): c_activationStatus = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") @@ -5076,6 +5888,7 @@ def nvmlDeviceSetMigMode(device, mode): _nvmlCheckReturn(ret) return c_activationStatus.value + def nvmlDeviceGetMigMode(device): c_currentMode = c_uint() c_pendingMode = c_uint() @@ -5084,6 +5897,7 @@ def nvmlDeviceGetMigMode(device): _nvmlCheckReturn(ret) return [c_currentMode.value, c_pendingMode.value] + def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): if version == 2: c_info = c_nvmlGpuInstanceProfileInfo_v2_t() @@ -5097,9 +5911,11 @@ def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): _nvmlCheckReturn(ret) return c_info + # Define function alias for the API exposed by NVML nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo + def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") @@ -5107,12 +5923,16 @@ def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): _nvmlCheckReturn(ret) return c_count.value -def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef): + +def nvmlDeviceGetGpuInstancePossiblePlacements( + device, profileId, placementsRef, countRef +): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") ret = fn(device, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceCreateGpuInstance(device, profileId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") @@ -5120,6 +5940,7 @@ def nvmlDeviceCreateGpuInstance(device, profileId): _nvmlCheckReturn(ret) return c_instance + def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") @@ -5127,18 +5948,21 @@ def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): _nvmlCheckReturn(ret) return c_instance + def nvmlGpuInstanceDestroy(gpuInstance): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") ret = fn(gpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") ret = fn(device, profileId, gpuInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") @@ -5146,6 +5970,7 @@ def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): _nvmlCheckReturn(ret) return c_instance + def nvmlGpuInstanceGetInfo(gpuInstance): c_info = c_nvmlGpuInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") @@ -5153,7 +5978,10 @@ def nvmlGpuInstanceGetInfo(gpuInstance): _nvmlCheckReturn(ret) return c_info -def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2): + +def nvmlGpuInstanceGetComputeInstanceProfileInfo( + device, profile, engProfile, version=2 +): if version == 2: c_info = c_nvmlComputeInstanceProfileInfo_v2_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV") @@ -5161,13 +5989,17 @@ def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, ve c_info = c_nvmlComputeInstanceProfileInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") else: - raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) ret = fn(device, profile, engProfile, byref(c_info)) _nvmlCheckReturn(ret) return c_info + # Define function alias for the API exposed by NVML -nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo +nvmlGpuInstanceGetComputeInstanceProfileInfoV = ( + nvmlGpuInstanceGetComputeInstanceProfileInfo +) + def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): c_count = c_uint() @@ -5176,12 +6008,16 @@ def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): _nvmlCheckReturn(ret) return c_count.value -def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef): + +def nvmlGpuInstanceGetComputeInstancePossiblePlacements( + gpuInstance, profileId, placementsRef, countRef +): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") ret = fn(gpuInstance, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") @@ -5189,25 +6025,33 @@ def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): _nvmlCheckReturn(ret) return c_instance -def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement): + +def nvmlGpuInstanceCreateComputeInstanceWithPlacement( + gpuInstance, profileId, placement +): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") ret = fn(gpuInstance, profileId, placement, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance + def nvmlComputeInstanceDestroy(computeInstance): fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") ret = fn(computeInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS -def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef): + +def nvmlGpuInstanceGetComputeInstances( + gpuInstance, profileId, computeInstancesRef, countRef +): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") @@ -5215,6 +6059,7 @@ def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): _nvmlCheckReturn(ret) return c_instance + def nvmlComputeInstanceGetInfo_v2(computeInstance): c_info = c_nvmlComputeInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") @@ -5222,9 +6067,11 @@ def nvmlComputeInstanceGetInfo_v2(computeInstance): _nvmlCheckReturn(ret) return c_info + def nvmlComputeInstanceGetInfo(computeInstance): return nvmlComputeInstanceGetInfo_v2(computeInstance) + def nvmlDeviceIsMigDeviceHandle(device): c_isMigDevice = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") @@ -5232,6 +6079,7 @@ def nvmlDeviceIsMigDeviceHandle(device): _nvmlCheckReturn(ret) return c_isMigDevice + def nvmlDeviceGetGpuInstanceId(device): c_gpuInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") @@ -5239,6 +6087,7 @@ def nvmlDeviceGetGpuInstanceId(device): _nvmlCheckReturn(ret) return c_gpuInstanceId.value + def nvmlDeviceGetComputeInstanceId(device): c_computeInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") @@ -5246,6 +6095,7 @@ def nvmlDeviceGetComputeInstanceId(device): _nvmlCheckReturn(ret) return c_computeInstanceId.value + def nvmlDeviceGetMaxMigDeviceCount(device): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") @@ -5253,6 +6103,7 @@ def nvmlDeviceGetMaxMigDeviceCount(device): _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetMigDeviceHandleByIndex(device, index): c_index = c_uint(index) migDevice = c_nvmlDevice_t() @@ -5261,6 +6112,7 @@ def nvmlDeviceGetMigDeviceHandleByIndex(device, index): _nvmlCheckReturn(ret) return migDevice + def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") @@ -5268,6 +6120,7 @@ def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): _nvmlCheckReturn(ret) return device + def nvmlDeviceGetAttributes_v2(device): c_attrs = c_nvmlDeviceAttributes() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") @@ -5275,9 +6128,11 @@ def nvmlDeviceGetAttributes_v2(device): _nvmlCheckReturn(ret) return c_attrs + def nvmlDeviceGetAttributes(device): return nvmlDeviceGetAttributes_v2(device) + def nvmlDeviceGetRemappedRows(device): fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") c_corr = c_uint() @@ -5288,6 +6143,7 @@ def nvmlDeviceGetRemappedRows(device): _nvmlCheckReturn(ret) return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + def nvmlDeviceGetRowRemapperHistogram(device): c_vals = c_nvmlRowRemapperHistogramValues() fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") @@ -5295,6 +6151,7 @@ def nvmlDeviceGetRowRemapperHistogram(device): _nvmlCheckReturn(ret) return c_vals + def nvmlDeviceGetArchitecture(device): arch = _nvmlDeviceArchitecture_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") @@ -5302,6 +6159,7 @@ def nvmlDeviceGetArchitecture(device): _nvmlCheckReturn(ret) return arch.value + def nvmlDeviceGetBusType(device): c_busType = _nvmlBusType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") @@ -5309,6 +6167,7 @@ def nvmlDeviceGetBusType(device): _nvmlCheckReturn(ret) return c_busType.value + def nvmlDeviceGetIrqNum(device): c_irqNum = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") @@ -5316,6 +6175,7 @@ def nvmlDeviceGetIrqNum(device): _nvmlCheckReturn(ret) return c_irqNum.value + def nvmlDeviceGetNumGpuCores(device): c_numCores = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") @@ -5323,6 +6183,7 @@ def nvmlDeviceGetNumGpuCores(device): _nvmlCheckReturn(ret) return c_numCores.value + def nvmlDeviceGetPowerSource(device): c_powerSource = _nvmlPowerSource_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") @@ -5330,6 +6191,7 @@ def nvmlDeviceGetPowerSource(device): _nvmlCheckReturn(ret) return c_powerSource.value + def nvmlDeviceGetMemoryBusWidth(device): c_memBusWidth = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") @@ -5337,6 +6199,7 @@ def nvmlDeviceGetMemoryBusWidth(device): _nvmlCheckReturn(ret) return c_memBusWidth.value + def nvmlDeviceGetPcieLinkMaxSpeed(device): c_speed = _nvmlPcieLinkMaxSpeed_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") @@ -5344,6 +6207,7 @@ def nvmlDeviceGetPcieLinkMaxSpeed(device): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetAdaptiveClockInfoStatus(device): c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") @@ -5351,6 +6215,7 @@ def nvmlDeviceGetAdaptiveClockInfoStatus(device): _nvmlCheckReturn(ret) return c_adaptiveClockInfoStatus.value + def nvmlDeviceGetPcieSpeed(device): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") @@ -5358,64 +6223,87 @@ def nvmlDeviceGetPcieSpeed(device): _nvmlCheckReturn(ret) return c_speed.value -def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()): + +def nvmlDeviceGetDynamicPstatesInfo( + device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t() +): isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t - dynamicpstatesinfoRef = c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + dynamicpstatesinfoRef = ( + c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + ) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo") ret = fn(device, dynamicpstatesinfoRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + def nvmlDeviceSetFanSpeed_v2(handle, index, speed): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2") ret = fn(handle, index, speed) _nvmlCheckReturn(ret) return NVML_SUCCESS -def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()): + +def nvmlDeviceGetThermalSettings( + device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t() +): isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings") ret = fn(device, sensorindex, thermalsettingsRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] -def nvmlDeviceGetMinMaxClockOfPState(device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()): + +def nvmlDeviceGetMinMaxClockOfPState( + device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint() +): isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState"); - ret = fn(device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState") + ret = fn( + device, + _nvmlClockType_t(clockType), + _nvmlClockType_t(pstate), + minClockMHzRef, + maxClockMHzRef, + ) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + class c_nvmlClockOffset_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('type', _nvmlClockType_t), - ('pstate', _nvmlPstates_t), - ('clockOffsetMHz', c_int), - ('minClockOffsetMHz', c_int), - ('maxClockOffsetMHz', c_int), + ("version", c_uint), + ("type", _nvmlClockType_t), + ("pstate", _nvmlPstates_t), + ("clockOffsetMHz", c_int), + ("minClockOffsetMHz", c_int), + ("maxClockOffsetMHz", c_int), ] + nvmlClockOffset_v1 = 0x1000018 + def nvmlDeviceGetClockOffsets(device, info): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets") ret = fn(device, info) return NVML_SUCCESS + def nvmlDeviceSetClockOffsets(device, info): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets") ret = fn(device, info) return NVML_SUCCESS + def nvmlDeviceGetSupportedPerformanceStates(device): pstates = [] c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) - c_size = sizeof(c_uint)*c_count.value + c_size = sizeof(c_uint) * c_count.value # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. pstates_array = _nvmlPstates_t * c_count.value @@ -5431,6 +6319,7 @@ def nvmlDeviceGetSupportedPerformanceStates(device): return pstates + def nvmlDeviceGetGpcClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") @@ -5438,6 +6327,7 @@ def nvmlDeviceGetGpcClkVfOffset(device): _nvmlCheckReturn(ret) return offset.value + def nvmlDeviceSetGpcClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") @@ -5445,6 +6335,7 @@ def nvmlDeviceSetGpcClkVfOffset(device, offset): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) @@ -5454,6 +6345,7 @@ def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + def nvmlDeviceGetMemClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") @@ -5461,6 +6353,7 @@ def nvmlDeviceGetMemClkVfOffset(device): _nvmlCheckReturn(ret) return offset.value + def nvmlDeviceSetMemClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") @@ -5468,6 +6361,7 @@ def nvmlDeviceSetMemClkVfOffset(device, offset): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) @@ -5478,6 +6372,7 @@ def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + def nvmlSystemSetConfComputeGpusReadyState(state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") @@ -5485,6 +6380,7 @@ def nvmlSystemSetConfComputeGpusReadyState(state): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetConfComputeGpusReadyState(): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") @@ -5492,6 +6388,7 @@ def nvmlSystemGetConfComputeGpusReadyState(): _nvmlCheckReturn(ret) return c_state.value + def nvmlSystemGetConfComputeCapabilities(): c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") @@ -5499,6 +6396,7 @@ def nvmlSystemGetConfComputeCapabilities(): _nvmlCheckReturn(ret) return c_ccSysCaps + def nvmlSystemGetConfComputeState(): c_state = c_nvmlConfComputeSystemState_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") @@ -5506,16 +6404,19 @@ def nvmlSystemGetConfComputeState(): _nvmlCheckReturn(ret) return c_state + def nvmlSystemGetConfComputeSettings(settings): fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") return fn(settings) + def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") ret = fn(device, c_ccMemSize) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetConfComputeMemSizeInfo(device): c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") @@ -5523,6 +6424,7 @@ def nvmlDeviceGetConfComputeMemSizeInfo(device): _nvmlCheckReturn(ret) return c_ccMemSize + def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): c_memory = c_nvmlMemory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") @@ -5530,6 +6432,7 @@ def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): _nvmlCheckReturn(ret) return c_memory + def nvmlDeviceGetConfComputeGpuCertificate(device): c_cert = c_nvmlConfComputeGpuCertificate_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") @@ -5537,15 +6440,17 @@ def nvmlDeviceGetConfComputeGpuCertificate(device): _nvmlCheckReturn(ret) return c_cert + def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) - setattr(c_attestReport, 'nonce', c_nonce_arr) + setattr(c_attestReport, "nonce", c_nonce_arr) fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") ret = fn(device, byref(c_attestReport)) _nvmlCheckReturn(ret) return c_attestReport + def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 @@ -5555,6 +6460,7 @@ def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 @@ -5563,98 +6469,182 @@ def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): _nvmlCheckReturn(ret) return c_keyRotationThrInfo + ## GPM ## ######### ## Enums/defines #### GPM Metric Identifiers -NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 -NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 -NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 -NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 -NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_DRAM_BW_UTIL = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 -NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 -NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 -NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 -NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec -NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec -NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 -NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec -NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec -NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec -NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = 63 # NvLink write bandwidth for link 0 in MiB/sec -NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec -NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = 65 # NvLink write bandwidth for link 1 in MiB/sec -NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec -NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = 67 # NvLink write bandwidth for link 2 in MiB/sec -NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec -NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = 69 # NvLink write bandwidth for link 3 in MiB/sec -NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec -NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = 71 # NvLink write bandwidth for link 4 in MiB/sec -NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec -NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = 73 # NvLink write bandwidth for link 5 in MiB/sec -NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec -NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = 75 # NvLink write bandwidth for link 6 in MiB/sec -NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec -NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = 77 # NvLink write bandwidth for link 7 in MiB/sec -NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec -NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = 79 # NvLink write bandwidth for link 8 in MiB/sec -NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec -NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = 81 # NvLink write bandwidth for link 9 in MiB/sec -NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = 82 # NvLink read bandwidth for link 10 in MiB/sec -NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = 83 # NvLink write bandwidth for link 10 in MiB/sec -NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = 84 # NvLink read bandwidth for link 11 in MiB/sec -NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = 85 # NvLink write bandwidth for link 11 in MiB/sec -NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = 86 # NvLink read bandwidth for link 12 in MiB/sec -NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = 87 # NvLink write bandwidth for link 12 in MiB/sec -NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = 88 # NvLink read bandwidth for link 13 in MiB/sec -NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = 89 # NvLink write bandwidth for link 13 in MiB/sec -NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = 90 # NvLink read bandwidth for link 14 in MiB/sec -NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = 91 # NvLink write bandwidth for link 14 in MiB/sec -NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = 92 # NvLink read bandwidth for link 15 in MiB/sec -NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = 93 # NvLink write bandwidth for link 15 in MiB/sec -NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = 94 # NvLink read bandwidth for link 16 in MiB/sec -NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = 95 # NvLink write bandwidth for link 16 in MiB/sec -NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = 96 # NvLink read bandwidth for link 17 in MiB/sec -NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = 97 # NvLink write bandwidth for link 17 in MiB/sec -NVML_GPM_METRIC_MAX = 98 +NVML_GPM_METRIC_GRAPHICS_UTIL = ( + 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +) +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = ( + 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +) +NVML_GPM_METRIC_INTEGER_UTIL = ( + 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_ANY_TENSOR_UTIL = ( + 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = ( + 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = ( + 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = ( + 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_DRAM_BW_UTIL = ( + 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP64_UTIL = ( + 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP32_UTIL = ( + 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP16_UTIL = ( + 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = ( + 60 # NvLink read bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = ( + 61 # NvLink write bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = ( + 63 # NvLink write bandwidth for link 0 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = ( + 65 # NvLink write bandwidth for link 1 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = ( + 67 # NvLink write bandwidth for link 2 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = ( + 69 # NvLink write bandwidth for link 3 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = ( + 71 # NvLink write bandwidth for link 4 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = ( + 73 # NvLink write bandwidth for link 5 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = ( + 75 # NvLink write bandwidth for link 6 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = ( + 77 # NvLink write bandwidth for link 7 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = ( + 79 # NvLink write bandwidth for link 8 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = ( + 81 # NvLink write bandwidth for link 9 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = ( + 82 # NvLink read bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = ( + 83 # NvLink write bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = ( + 84 # NvLink read bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = ( + 85 # NvLink write bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = ( + 86 # NvLink read bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = ( + 87 # NvLink write bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = ( + 88 # NvLink read bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = ( + 89 # NvLink write bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = ( + 90 # NvLink read bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = ( + 91 # NvLink write bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = ( + 92 # NvLink read bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = ( + 93 # NvLink write bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = ( + 94 # NvLink read bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = ( + 95 # NvLink write bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = ( + 96 # NvLink read bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = ( + 97 # NvLink write bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_MAX = 98 ## Structs + class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ - ('name', c_char * 96), - ('id', c_char * 96), - ('serial', c_char * 96), - ('firmwareVersion', c_char * 96), + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), ] + class struct_c_nvmlGpmSample_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + class c_metricInfo_t(Structure): _fields_ = [ ("shortName", c_char_p), @@ -5662,47 +6652,55 @@ class c_metricInfo_t(Structure): ("unit", c_char_p), ] + class c_nvmlGpmMetric_t(_PrintableStructure): _fields_ = [ - ('metricId', c_uint), - ('nvmlReturn', _nvmlReturn_t), - ('value', c_double), - ('metricInfo', c_metricInfo_t) + ("metricId", c_uint), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_double), + ("metricInfo", c_metricInfo_t), ] + class c_nvmlGpmMetricsGet_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('numMetrics', c_uint), - ('sample1', c_nvmlGpmSample_t), - ('sample2', c_nvmlGpmSample_t), - ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX) + ("version", c_uint), + ("numMetrics", c_uint), + ("sample1", c_nvmlGpmSample_t), + ("sample2", c_nvmlGpmSample_t), + ("metrics", c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX), ] + NVML_GPM_METRICS_GET_VERSION = 1 + class c_nvmlGpmSupport_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('isSupportedDevice', c_uint), + ("version", c_uint), + ("isSupportedDevice", c_uint), ] + NVML_GPM_SUPPORT_VERSION = 1 ## Functions + def nvmlGpmMetricsGet(metricsGet): fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") ret = fn(byref(metricsGet)) _nvmlCheckReturn(ret) return metricsGet + def nvmlGpmSampleFree(gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") ret = fn(gpmSample) _nvmlCheckReturn(ret) return + def nvmlGpmSampleAlloc(): gpmSample = c_nvmlGpmSample_t() fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") @@ -5710,18 +6708,21 @@ def nvmlGpmSampleAlloc(): _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmSampleGet(device, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") ret = fn(device, gpmSample) _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") ret = fn(device, gpuInstanceId, gpmSample) _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmQueryDeviceSupport(device): gpmSupport = c_nvmlGpmSupport_t() gpmSupport.version = NVML_GPM_SUPPORT_VERSION @@ -5730,6 +6731,7 @@ def nvmlGpmQueryDeviceSupport(device): _nvmlCheckReturn(ret) return gpmSupport + def nvmlGpmSetStreamingEnabled(device, state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") @@ -5737,6 +6739,7 @@ def nvmlGpmSetStreamingEnabled(device, state): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpmQueryIfStreamingEnabled(device): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") @@ -5744,21 +6747,24 @@ def nvmlGpmQueryIfStreamingEnabled(device): _nvmlCheckReturn(ret) return c_state.value + # Low Power Structure and Function -NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 -NVML_NVLINK_POWER_STATE_LOW = 0x1 +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 -NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 -NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF -NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + class c_nvmlNvLinkPowerThres_t(Structure): _fields_ = [ ("lowPwrThreshold", c_uint), ] + def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): c_info = c_nvmlNvLinkPowerThres_t() c_info.lowPwrThreshold = l1threshold @@ -5767,48 +6773,52 @@ def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): _nvmlCheckReturn(ret) return NVML_SUCCESS + NVML_GPU_FABRIC_UUID_LEN = 16 _nvmlGpuFabricState_t = c_uint NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 -NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 -NVML_GPU_FABRIC_STATE_COMPLETED = 3 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + class c_nvmlGpuFabricInfo_t(_PrintableStructure): _fields_ = [ ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), ("status", _nvmlReturn_t), ("cliqueId", c_uint32), - ("state", _nvmlGpuFabricState_t) + ("state", _nvmlGpuFabricState_t), ] + NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 - -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 - -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 nvmlGpuFabricInfo_v2 = 0x02000024 + class c_nvmlGpuFabricInfoV_t(_PrintableStructure): _fields_ = [ ("version", c_uint), @@ -5816,33 +6826,37 @@ class c_nvmlGpuFabricInfoV_t(_PrintableStructure): ("status", _nvmlReturn_t), ("cliqueId", c_uint32), ("state", _nvmlGpuFabricState_t), - ("healthMask", c_uint32) + ("healthMask", c_uint32), ] def __init__(self): super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2) + def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + ###################### ## Enums/defines #### NVML GPU NVLINK BW MODE -NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 -NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 -NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 -NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 -NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 -NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + def nvmlSystemSetNvlinkBwMode(mode): fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") @@ -5850,6 +6864,7 @@ def nvmlSystemSetNvlinkBwMode(mode): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetNvlinkBwMode(): mode = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") @@ -5857,22 +6872,28 @@ def nvmlSystemGetNvlinkBwMode(): _nvmlCheckReturn(ret) return mode.value + _nvmlPowerScopeType_t = c_uint -NVML_POWER_SCOPE_GPU = 0 -NVML_POWER_SCOPE_MODULE = 1 -NVML_POWER_SCOPE_MEMORY = 2 +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + class c_nvmlPowerValue_v2_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('powerScope', _nvmlPowerScopeType_t), - ('powerValueMw', c_uint), + ("version", c_uint), + ("powerScope", _nvmlPowerScopeType_t), + ("powerValueMw", c_uint), ] - _fmt_ = {'': "%d B"} + _fmt_ = {"": "%d B"} + nvmlPowerValue_v2 = 0x0200000C -def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2): + +def nvmlDeviceSetPowerManagementLimit_v2( + device, powerScope, powerLimit, version=nvmlPowerValue_v2 +): c_powerScope = _nvmlPowerScopeType_t(powerScope) c_powerValue = c_nvmlPowerValue_v2_t() c_powerValue.version = c_uint(version) @@ -5882,139 +6903,178 @@ def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version ret = fn(device, byref(c_powerValue)) return NVML_SUCCESS + class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('aggregateUncParity', c_ulonglong), - ('aggregateUncSecDed', c_ulonglong), - ('aggregateCor', c_ulonglong), - ('volatileUncParity', c_ulonglong), - ('volatileUncSecDed', c_ulonglong), - ('volatileCor', c_ulonglong), - ('aggregateUncBucketL2', c_ulonglong), - ('aggregateUncBucketSm', c_ulonglong), - ('aggregateUncBucketPcie', c_ulonglong), - ('aggregateUncBucketMcu', c_ulonglong), - ('aggregateUncBucketOther', c_ulonglong), - ('bThresholdExceeded', c_uint) + ("version", c_uint), + ("aggregateUncParity", c_ulonglong), + ("aggregateUncSecDed", c_ulonglong), + ("aggregateCor", c_ulonglong), + ("volatileUncParity", c_ulonglong), + ("volatileUncSecDed", c_ulonglong), + ("volatileCor", c_ulonglong), + ("aggregateUncBucketL2", c_ulonglong), + ("aggregateUncBucketSm", c_ulonglong), + ("aggregateUncBucketPcie", c_ulonglong), + ("aggregateUncBucketMcu", c_ulonglong), + ("aggregateUncBucketOther", c_ulonglong), + ("bThresholdExceeded", c_uint), ] def __init__(self): - super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1) + super(c_nvmlEccSramErrorStatus_v1_t, self).__init__( + version=nvmlEccSramErrorStatus_v1 + ) + nvmlEccSramErrorStatus_v1 = 0x1000068 + + def nvmlDeviceGetSramEccErrorStatus(device, status): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") ret = fn(device, status) _nvmlCheckReturn(ret) return NVML_SUCCESS -NVML_DEV_CAP_EGM = (1 << 0) + +NVML_DEV_CAP_EGM = 1 << 0 nvmlDeviceCapabilities_v1 = 0x1000008 + class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('capMask', c_uint), + ("version", c_uint), + ("capMask", c_uint), ] def __init__(self): - super(c_nvmlDeviceCapabilities_v1_t, self).__init__(version=nvmlDeviceCapabilities_v1) + super(c_nvmlDeviceCapabilities_v1_t, self).__init__( + version=nvmlDeviceCapabilities_v1 + ) def nvmlDeviceGetCapabilities(device, caps): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") return fn(device, caps) + class c_nvmlPlatformInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('ibGuid', c_char * 16), - ('rackGuid', c_char * 16), - ('chassisPhysicalSlotNumber', c_char), - ('computeSlotIndex', c_char), - ('nodeIndex', c_char), - ('peerType', c_char), - ('moduleId', c_char) + ("version", c_uint), + ("ibGuid", c_char * 16), + ("rackGuid", c_char * 16), + ("chassisPhysicalSlotNumber", c_char), + ("computeSlotIndex", c_char), + ("nodeIndex", c_char), + ("peerType", c_char), + ("moduleId", c_char), ] def __init__(self): super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) -nvmlPlatformInfo_v1 = 0x100002c + +nvmlPlatformInfo_v1 = 0x100002C + + def nvmlDeviceGetPlatformInfo(device, platformInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") ret = fn(device, platformInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + class c_nvmlMask255_t(_PrintableStructure): _fields_ = [ - ('mask', c_uint * 8), + ("mask", c_uint * 8), ] -NVML_WORKLOAD_POWER_MAX_PROFILES = 255 -NVML_POWER_PROFILE_MAX_P = 0 -NVML_POWER_PROFILE_MAX_Q = 1 -NVML_POWER_PROFILE_COMPUTE = 2 -NVML_POWER_PROFILE_MEMORY_BOUND = 3 -NVML_POWER_PROFILE_NETWORK = 4 -NVML_POWER_PROFILE_BALANCED = 5 -NVML_POWER_PROFILE_LLM_INFERENCE = 6 -NVML_POWER_PROFILE_LLM_TRAINING = 7 -NVML_POWER_PROFILE_RBM = 8 -NVML_POWER_PROFILE_DCPCIE = 9 -NVML_POWER_PROFILE_HMMA_SPARSE = 10 -NVML_POWER_PROFILE_HMMA_DENSE = 11 -NVML_POWER_PROFILE_SYNC_BALANCED = 12 -NVML_POWER_PROFILE_HPC = 13 -NVML_POWER_PROFILE_MIG = 14 -NVML_POWER_PROFILE_MAX = 15 - -nvmlWorkloadPowerProfileInfo_v1 = 0x100002c + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 + +nvmlWorkloadPowerProfileInfo_v1 = 0x100002C + + class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('profileId', c_uint), - ('priority', c_uint), - ('conflictingmask', c_nvmlMask255_t) + ("version", c_uint), + ("profileId", c_uint), + ("priority", c_uint), + ("conflictingmask", c_nvmlMask255_t), ] def __init__(self): - super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileInfo_v1) + super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileInfo_v1 + ) + + +nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002BF8 + -nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002bf8 class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('perfProfilesMask', c_nvmlMask255_t), - ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES) + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ( + "perfProfile", + c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES, + ), ] def __init__(self): - super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1) + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileProfilesInfo_v1 + ) + nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 + + class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('perfProfilesMask', c_nvmlMask255_t), - ('requestedProfilesMask', c_nvmlMask255_t), - ('enforcedProfilesMask', c_nvmlMask255_t) + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ("requestedProfilesMask", c_nvmlMask255_t), + ("enforcedProfilesMask", c_nvmlMask255_t), ] def __init__(self): - super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1) + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileCurrentProfiles_v1 + ) + nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 + + class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('requestedProfilesMask', c_nvmlMask255_t), + ("version", c_uint), + ("requestedProfilesMask", c_nvmlMask255_t), ] def __init__(self): - super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1) + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileRequestedProfiles_v1 + ) + def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") @@ -6022,52 +7082,63 @@ def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") ret = fn(device, currentProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") ret = fn(device, supportedBwModes) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetNvlinkBwMode(device, getBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") ret = fn(device, getBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceSetNvlinkBwMode(device, setBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") ret = fn(device, setBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS + nvmlDramEncryptionInfo_v1 = 0x01000008 + class c_nvmlDramEncryptionInfo_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('encryptionState', _nvmlEnableState_t), + ("version", c_uint), + ("encryptionState", _nvmlEnableState_t), ] def __init__(self): - super(c_nvmlDramEncryptionInfo_t, self).__init__(version=nvmlDramEncryptionInfo_v1) + super(c_nvmlDramEncryptionInfo_t, self).__init__( + version=nvmlDramEncryptionInfo_v1 + ) + def nvmlDeviceGetDramEncryptionMode(handle): c_currState = c_nvmlDramEncryptionInfo_t() @@ -6077,64 +7148,79 @@ def nvmlDeviceGetDramEncryptionMode(handle): _nvmlCheckReturn(ret) return [c_currState.encryptionState, c_pendingState.encryptionState] + # added to API def nvmlDeviceGetCurrentDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[0] + # added to API def nvmlDeviceGetPendingDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[1] + def nvmlDeviceSetDramEncryptionMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() - c_dramEncryptionMode.encryptionState = mode; + c_dramEncryptionMode.encryptionState = mode ret = fn(handle, byref(c_dramEncryptionMode)) _nvmlCheckReturn(ret) return None + # Power Smoothing defines -NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 -NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF -NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 -NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 -NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 -nvmlPowerSmoothingState_v1=0x1000008 +nvmlPowerSmoothingState_v1 = 0x1000008 + + class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('state', c_uint), + ("version", c_uint), + ("state", c_uint), ] def __init__(self): - super(c_nvmlPowerSmoothingState_v1_t, self).__init__(version=nvmlPowerSmoothingState_v1) + super(c_nvmlPowerSmoothingState_v1_t, self).__init__( + version=nvmlPowerSmoothingState_v1 + ) + + +nvmlPowerSmoothingProfile_v1 = 0x1000018 + -nvmlPowerSmoothingProfile_v1=0x1000018 class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): _fields_ = [ - ('version', c_uint), - ('profileId', c_uint), - ('paramId', c_uint), - ('value', c_double), + ("version", c_uint), + ("profileId", c_uint), + ("paramId", c_uint), + ("value", c_double), ] def __init__(self): - super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(version=nvmlPowerSmoothingProfile_v1) + super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__( + version=nvmlPowerSmoothingProfile_v1 + ) + def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") ret = fn(device, profile) _nvmlCheckReturn(ret) + def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") ret = fn(device, profile) _nvmlCheckReturn(ret) + def nvmlDevicePowerSmoothingSetState(device, state): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") ret = fn(device, state) _nvmlCheckReturn(ret) - diff --git a/vllm/tilelang_utils/__init__.py b/vllm/tilelang_utils/__init__.py index 803409b1ec65..990c2fca3e77 100644 --- a/vllm/tilelang_utils/__init__.py +++ b/vllm/tilelang_utils/__init__.py @@ -31,6 +31,7 @@ def _ensure_tilelang_imported() -> None: Raises: ImportError: If TileLang is not installed. + """ global T, tilelang diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index 0e9aaa2cae5a..3f0fff938ad9 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -13,9 +13,7 @@ def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: - """ - Wraps a tokenizer to use the custom DeepSeek V3.2 chat template encoding. - """ + """Wraps a tokenizer to use the custom DeepSeek V3.2 chat template encoding.""" dsv32_tokenizer = copy.copy(tokenizer) added_vocab = tokenizer.get_added_vocab() diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index 2a0db215b80f..a5d5020e941b 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -13,9 +13,7 @@ def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: - """ - Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding. - """ + """Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding.""" dsv4_tokenizer = copy.copy(tokenizer) added_vocab = tokenizer.get_added_vocab() diff --git a/vllm/tokenizers/detokenizer_utils.py b/vllm/tokenizers/detokenizer_utils.py index 516f1b70e732..35e73bc4100d 100644 --- a/vllm/tokenizers/detokenizer_utils.py +++ b/vllm/tokenizers/detokenizer_utils.py @@ -205,6 +205,7 @@ def detokenize_incrementally( skip_special_tokens: Whether to skip special tokens. spaces_between_special_tokens: Whether to add spaces between special tokens. + """ new_token_id = all_input_ids[-1] # This is the first iteration for this sequence diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py index 3d35952c5f4c..a091d8e10547 100644 --- a/vllm/tokenizers/hf.py +++ b/vllm/tokenizers/hf.py @@ -23,8 +23,7 @@ class ThreadSafeHFTokenizerMixin: def maybe_make_thread_pool(tokenizer: _T, copies: int = 1): - """ - If `tokenizer` is a `TokenizersBackend`, modify the tokenizer + """If `tokenizer` is a `TokenizersBackend`, modify the tokenizer in-place to make the public interface thread-safe by routing calls through a deep-copied tokenizer pool. @@ -105,8 +104,7 @@ def __reduce__(self): def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: - """ - By default, transformers will recompute multiple tokenizer properties + """By default, transformers will recompute multiple tokenizer properties each time they are called, leading to a significant slowdown. This proxy caches these properties for faster access. """ diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 3ee35365e4a8..3847afa9b227 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -40,8 +40,7 @@ class ToolParser: - """ - Abstract ToolParser class that should not be used directly. Provided + """Abstract ToolParser class that should not be used directly. Provided properties and methods should be used in derived classes. """ @@ -186,8 +185,7 @@ def get_structural_tag( def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Static method that should be implemented for extracting tool calls from + """Static method that should be implemented for extracting tool calls from a complete model-generated string. Used for non-streaming responses where we have the entire model response available before sending to the client. @@ -207,8 +205,7 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - """ - Instance method that should be implemented for extracting tool calls + """Instance method that should be implemented for extracting tool calls from an incomplete response; for use when handling tool calls and streaming. Has to be an instance method because it requires state - the current tokens/diffs, but also the information about what has @@ -220,8 +217,7 @@ def extract_tool_calls_streaming( class ToolParserManager: - """ - Central registry for ToolParser implementations. + """Central registry for ToolParser implementations. Supports two modes: - Eager (immediate) registration via `register_module` @@ -233,8 +229,7 @@ class ToolParserManager: @classmethod def get_tool_parser(cls, name: str) -> type[ToolParser]: - """ - Retrieve a registered or lazily registered ToolParser class. + """Retrieve a registered or lazily registered ToolParser class. If the parser is lazily registered, it will be imported and cached on first access. @@ -301,8 +296,7 @@ def _register_module( @classmethod def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None: - """ - Register a lazy module mapping. + """Register a lazy module mapping. Example: ToolParserManager.register_lazy_module( @@ -310,6 +304,7 @@ def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> N module_path="vllm.tool_parsers.kimi_k2_parser", class_name="KimiK2ToolParser", ) + """ cls.lazy_parsers[name] = (module_path, class_name) @@ -320,8 +315,7 @@ def register_module( force: bool = True, module: type[ToolParser] | None = None, ) -> type[ToolParser] | Callable[[type[ToolParser]], type[ToolParser]]: - """ - Register module immediately or lazily (as a decorator). + """Register module immediately or lazily (as a decorator). Usage: @ToolParserManager.register_module("kimi_k2") diff --git a/vllm/tool_parsers/apertus_tool_parser.py b/vllm/tool_parsers/apertus_tool_parser.py index f23b8601aa61..6825b3ef39e1 100644 --- a/vllm/tool_parsers/apertus_tool_parser.py +++ b/vllm/tool_parsers/apertus_tool_parser.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tool call parser for Apertus models. +"""Tool call parser for Apertus models. Extracts tool calls from the format: <|tools_prefix|>[{"function_name": {"arg1": "value1", ...}}, ...]<|tools_suffix|> @@ -45,8 +44,7 @@ class ApertusToolParser(ToolParser): - """ - Tool call parser for Apertus models. + """Tool call parser for Apertus models. Handles the extraction of tool calls from text in both non-streaming (complete string) and streaming (chunked token) environments. @@ -68,11 +66,11 @@ class ApertusToolParser(ToolParser): "get_weather" >>> info.tool_calls[0].function.arguments '{"city": "Paris"}' + """ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - """ - Initializes the ApertusToolParser. + """Initializes the ApertusToolParser. Args: tokenizer: The model's tokenizer. @@ -82,6 +80,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): Raises: ValueError: If the `model_tokenizer` is not successfully passed to the base class. + """ super().__init__(tokenizer, tools) @@ -101,8 +100,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): self._reset_streaming_state() def _reset_streaming_state(self) -> None: - """ - Resets all streaming state variables for a new completion request. + """Resets all streaming state variables for a new completion request. This clears the delta text buffer and resets the pointers used to track the currently streaming tool index and arguments. Called implicitly @@ -116,8 +114,7 @@ def _reset_streaming_state(self) -> None: def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: - """ - Adjusts the generation request to ensure special tool tokens are not skipped. + """Adjusts the generation request to ensure special tool tokens are not skipped. Forces `skip_special_tokens=False` if tools are actively being evaluated, ensuring the tools special tokens are surfaced to the engine for parsing. @@ -127,6 +124,7 @@ def adjust_request( Returns: The potentially modified chat completion request. + """ request = super().adjust_request(request) if request.tools and request.tool_choice != "none": @@ -134,8 +132,7 @@ def adjust_request( return request def _buffer_delta_text(self, delta_text: str) -> str: - """ - Buffers incoming delta chunks to prevent + """Buffers incoming delta chunks to prevent fragmentation of multi-token special tags. If a chunk ends with a partial match of @@ -155,6 +152,7 @@ def _buffer_delta_text(self, delta_text: str) -> str: "Let me check " # "<|tool" is buffered internally >>> parser._buffer_delta_text("s_prefix|>" \ "<|tools_prefix|>" # Buffer released on completion + """ self.buffered_delta_text += delta_text text = self.buffered_delta_text @@ -178,8 +176,7 @@ def extract_tool_calls( model_output: str, request: ChatCompletionRequest, ) -> ExtractedToolCallInformation: - """ - Extracts tool calls from a completely generated model response (Non-Streaming). + """Extract tool calls from a complete model response (non-streaming). Args: model_output: The full completion string generated by the model. @@ -200,6 +197,7 @@ def extract_tool_calls( 'Let me see.' >>> info.tool_calls[0].function.name 'get_weather' + """ match = self.tool_call_regex.search(model_output) if not match: @@ -280,8 +278,7 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - """ - Handles streaming chunks + """Handles streaming chunks. Args: previous_text: The complete model text generated prior to this chunk. @@ -305,6 +302,7 @@ def extract_tool_calls_streaming( ... ) >>> msg.tool_calls[0].function.arguments 'ation": "Paris"}' + """ delta_text = self._buffer_delta_text(delta_text) if not delta_text: @@ -323,8 +321,7 @@ def extract_tool_calls_streaming( def _extract_streaming( self, current_text: str, delta_text: str ) -> DeltaMessage | None: - """ - Core streaming logic. + """Core streaming logic. Separates visible chat text from JSON blocks and computes diffs. Args: @@ -333,6 +330,7 @@ def _extract_streaming( Returns: A `DeltaMessage` containing the `content` delta and/or `tool_calls` delta. + """ prefix_idx = current_text.rfind(TOOL_CALLS_PREFIX) suffix_idx = current_text.rfind(TOOL_CALLS_SUFFIX) @@ -404,8 +402,7 @@ def _extract_streaming( def _parse_and_diff_json( self, json_str: str, is_final: bool ) -> list[DeltaToolCall]: - """ - Parses an isolated, potentially incomplete streaming JSON array and returns + """Parses an isolated, potentially incomplete streaming JSON array and returns newly accumulated tool call diffs. Args: @@ -417,6 +414,7 @@ def _parse_and_diff_json( A list of `DeltaToolCall` items representing string diffs in function arguments to stream back to the client. + """ try: parsed, _ = partial_json_loads(json_str, Allow.ALL) @@ -460,8 +458,7 @@ def _parse_and_diff_json( def _emit_tool_name( self, parsed: list, index: int, tool_calls: list[DeltaToolCall] ) -> None: - """ - Extracts and emits the function name mapped to a new tool call ID. + """Extracts and emits the function name mapped to a new tool call ID. Args: parsed: The partially parsed JSON list containing tool dictionaries. @@ -472,6 +469,7 @@ def _emit_tool_name( Appends `DeltaToolCall(index=0, function=DeltaFunctionCall(name="get_weather", ...))` to the `tool_calls` list and marks the name as sent. + """ obj = parsed[index] if isinstance(obj, dict) and obj: @@ -491,8 +489,7 @@ def _emit_tool_name( def _get_tool_diff( self, parsed: list, index: int, is_final: bool ) -> DeltaToolCall | None: - """ - Calculates the exact string difference to safely append new tool parameters. + """Calculates the exact string difference to safely append new tool parameters. This ensures characters like `{`, `}`, and `"` don't jump around unevenly in the UI frontend while streaming incomplete JSON arguments. @@ -513,6 +510,7 @@ def _get_tool_diff( >>> # Returns diff (closing bracket suppressed until final): >>> parser._get_tool_diff(parsed, index=0, is_final=False) DeltaToolCall(index=0, function=DeltaFunctionCall(arguments='s')) + """ obj = parsed[index] if not isinstance(obj, dict) or not obj: diff --git a/vllm/tool_parsers/ernie45_tool_parser.py b/vllm/tool_parsers/ernie45_tool_parser.py index 74624b14048d..3c654b44897e 100644 --- a/vllm/tool_parsers/ernie45_tool_parser.py +++ b/vllm/tool_parsers/ernie45_tool_parser.py @@ -29,8 +29,7 @@ class Ernie45ToolParser(ToolParser): def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - """ - Ernie thinking model format: + """Ernie thinking model format: abc\n\n\n\n\ndef\n\n """ super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py index a975f68368a5..708ecb1dd997 100644 --- a/vllm/tool_parsers/functiongemma_tool_parser.py +++ b/vllm/tool_parsers/functiongemma_tool_parser.py @@ -27,8 +27,7 @@ class FunctionGemmaToolParser(ToolParser): - """ - Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it). + """Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it). Handles the FunctionGemma function call format: call:func_name{param:value} diff --git a/vllm/tool_parsers/gemma4_utils.py b/vllm/tool_parsers/gemma4_utils.py index d9aad254c728..33ad20b19396 100644 --- a/vllm/tool_parsers/gemma4_utils.py +++ b/vllm/tool_parsers/gemma4_utils.py @@ -52,6 +52,7 @@ def _parse_tool_arguments(args_str: str) -> dict[str, str]: Returns: Dictionary of argument name → string value. + """ if not args_str or not args_str.strip(): return {} @@ -95,6 +96,7 @@ def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]: >>> tool_calls = parse_tool_calls(output) >>> for tc in tool_calls: ... print(f"Call: {tc['name']}({tc['arguments']})") + """ results = [] @@ -150,6 +152,7 @@ def has_tool_response_tag(text: str) -> bool: >>> if not has_tool_response_tag(model_output): ... # Model used instead — inject <|tool_response> manually ... next_prompt = "<|tool_response>" + tool_result + """ stripped = text.rstrip() return stripped.endswith(_TOOL_RESPONSE_START_TAG) diff --git a/vllm/tool_parsers/gptoss_tool_parser.py b/vllm/tool_parsers/gptoss_tool_parser.py index e147ff1c7418..80881f69a5f5 100644 --- a/vllm/tool_parsers/gptoss_tool_parser.py +++ b/vllm/tool_parsers/gptoss_tool_parser.py @@ -15,8 +15,7 @@ class GptOssToolParser(ToolParser): - """ - Stub tool parser for gpt-oss/harmony models. + """Stub tool parser for gpt-oss/harmony models. All output parsing is handled by HarmonyParser. This stub exists as a capability declaration via HarmonyParser.tool_parser_cls. diff --git a/vllm/tool_parsers/granite_20b_fc_tool_parser.py b/vllm/tool_parsers/granite_20b_fc_tool_parser.py index 432f94b61e5d..4c6c0a1b0efe 100644 --- a/vllm/tool_parsers/granite_20b_fc_tool_parser.py +++ b/vllm/tool_parsers/granite_20b_fc_tool_parser.py @@ -38,8 +38,7 @@ class Granite20bFCToolParser(ToolParser): - """ - Tool call parser for the granite-20b-functioncalling model intended + """Tool call parser for the granite-20b-functioncalling model intended for use with the examples/tool_chat_template_granite20b_fc.jinja template. diff --git a/vllm/tool_parsers/granite_tool_parser.py b/vllm/tool_parsers/granite_tool_parser.py index c1d852dfc352..1e51c3227e89 100644 --- a/vllm/tool_parsers/granite_tool_parser.py +++ b/vllm/tool_parsers/granite_tool_parser.py @@ -36,8 +36,7 @@ class GraniteToolParser(ToolParser): - """ - Tool call parser for the granite 3.0 models. Intended + """Tool call parser for the granite 3.0 models. Intended for use with the examples/tool_chat_template_granite.jinja template. diff --git a/vllm/tool_parsers/hermes_tool_parser.py b/vllm/tool_parsers/hermes_tool_parser.py index c5b318cd19fc..0a300964102c 100644 --- a/vllm/tool_parsers/hermes_tool_parser.py +++ b/vllm/tool_parsers/hermes_tool_parser.py @@ -219,7 +219,8 @@ def extract_tool_calls_streaming( Returns a ``DeltaMessage`` containing either plain content (for text preceding any tool call) or one or more ``DeltaToolCall`` - entries, or ``None`` if there is nothing new to send yet.""" + entries, or ``None`` if there is nothing new to send yet. + """ try: # Extract any content before tool calls. content = self._extract_content(current_text) diff --git a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py index 88b2a20d3a0d..90e291084fd3 100644 --- a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py +++ b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py @@ -100,9 +100,7 @@ def preprocess_model_output( def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract tool calls from a complete model output. - """ + """Extract tool calls from a complete model output.""" try: # Preprocess the model output content, potential_tool_calls = self.preprocess_model_output(model_output) @@ -177,10 +175,7 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - """ - Extract tool calls for streaming mode. - """ - + """Extract tool calls for streaming mode.""" start_idx = consume_space(0, current_text) if current_text[start_idx:].startswith(self.bot_string): start_idx = consume_space(start_idx + len(self.bot_string), current_text) diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py index d73473673fdc..0effb377d0a6 100644 --- a/vllm/tool_parsers/hy_v3_tool_parser.py +++ b/vllm/tool_parsers/hy_v3_tool_parser.py @@ -181,6 +181,7 @@ def _try_parse_wildcard_number(value: str) -> int | float | None: "-3" → int(-3) Return None on failure. + """ try: if "." in value or "e" in value or "E" in value: diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py index 729463162693..c2f82fc4a7be 100644 --- a/vllm/tool_parsers/hy_v4_tool_parser.py +++ b/vllm/tool_parsers/hy_v4_tool_parser.py @@ -101,8 +101,8 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str: RuntimeError: The tokenizer declares the structural tokens through ``model_specific_special_tokens``, which transformers 5 no longer round-trips. - """ + """ import transformers if int(transformers.__version__.split(".")[0]) >= 5: @@ -611,6 +611,7 @@ def extract_tool_calls_streaming( Returns: A streaming delta carrying content and/or the tool calls drained from the buffer, or None when nothing can be emitted yet. + """ content_delta: str | None = None tool_calls: list[StreamToolCall] = [] @@ -1011,6 +1012,7 @@ def get_structural_tag( Returns: The structural tag, or None when structural tagging does not apply. + """ if not envs.VLLM_ENFORCE_STRICT_TOOL_CALLING: return None @@ -1070,6 +1072,7 @@ def _is_guided(request: ChatCompletionRequest) -> bool: Returns: True when the streaming parser must use the string-marker path. + """ structured_outputs = getattr(request, "structured_outputs", None) return ( diff --git a/vllm/tool_parsers/lfm2_tool_parser.py b/vllm/tool_parsers/lfm2_tool_parser.py index 6f4db492be29..7d59005a7a5b 100644 --- a/vllm/tool_parsers/lfm2_tool_parser.py +++ b/vllm/tool_parsers/lfm2_tool_parser.py @@ -42,8 +42,7 @@ class Lfm2ToolParser(ToolParser): - """ - Tool call parser for LiquidAI LFM2/LFM2.5 models that produce pythonic + """Tool call parser for LiquidAI LFM2/LFM2.5 models that produce pythonic tool calls wrapped in <|tool_call_start|> and <|tool_call_end|> tokens. Example model output: @@ -115,7 +114,8 @@ def _strip_echo(raw_after: str) -> str: """Drop any orphan <|tool_call_end|> (and the preceding text) from trailing content. LFM2 occasionally echoes the call body after the first end token and caps it with a second end token; everything - through the last such orphan is model garbage, not user content.""" + through the last such orphan is model garbage, not user content. + """ last_orphan = raw_after.rfind(TOOL_CALL_END) if last_orphan != -1: return raw_after[last_orphan + len(TOOL_CALL_END) :] diff --git a/vllm/tool_parsers/llama4_pythonic_tool_parser.py b/vllm/tool_parsers/llama4_pythonic_tool_parser.py index 18c697f04829..fdffb48ff7fe 100644 --- a/vllm/tool_parsers/llama4_pythonic_tool_parser.py +++ b/vllm/tool_parsers/llama4_pythonic_tool_parser.py @@ -31,8 +31,7 @@ class Llama4PythonicToolParser(ToolParser): - """ - Toolcall parser for Llama4 that produce tool calls in a pythonic style + """Toolcall parser for Llama4 that produce tool calls in a pythonic style Use --enable-auto-tool-choice --tool-call-parser llama4_pythonic """ @@ -67,10 +66,7 @@ def current_tool_index(self, value: int) -> None: def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. - """ - + """Extract the tool calls from a complete model response.""" # remove <|python_start|> and <|python_end|> # as Llama 4 model sometime will output those tokens if model_output.startswith("<|python_start|>"): diff --git a/vllm/tool_parsers/llama_tool_parser.py b/vllm/tool_parsers/llama_tool_parser.py index 5ad7b1cbdd29..93b0980c3775 100644 --- a/vllm/tool_parsers/llama_tool_parser.py +++ b/vllm/tool_parsers/llama_tool_parser.py @@ -37,8 +37,7 @@ class Llama3JsonToolParser(ToolParser): - """ - Tool call parser for Llama 3.x and 4 models intended for use with the + """Tool call parser for Llama 3.x and 4 models intended for use with the examples/tool_chat_template_llama.jinja template. Used when --enable-auto-tool-choice --tool-call-parser llama3_json or @@ -77,8 +76,7 @@ def __init__( def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. + """Extract the tool calls from a complete model response. Only extracts JSON content and ignores any surrounding plain text. Supports both single JSON and multiple JSONs separated by semicolons. """ diff --git a/vllm/tool_parsers/olmo3_tool_parser.py b/vllm/tool_parsers/olmo3_tool_parser.py index 49a871653399..cb7bc632baa8 100644 --- a/vllm/tool_parsers/olmo3_tool_parser.py +++ b/vllm/tool_parsers/olmo3_tool_parser.py @@ -31,8 +31,7 @@ class Olmo3PythonicToolParser(ToolParser): - """ - Tool call parser for Olmo 3 models that produce tool calls as + """Tool call parser for Olmo 3 models that produce tool calls as newline-separated pythonic strings. Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set Code copied from pythonic_tool_parser.py and updated to handle @@ -71,9 +70,7 @@ def current_tool_index(self, value: int) -> None: def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. - """ + """Extract the tool calls from a complete model response.""" original_model_output = model_output # Remove xml tags. match = re.search( diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py index f5b794f0f67b..820e797f4d9f 100644 --- a/vllm/tool_parsers/phi4mini_tool_parser.py +++ b/vllm/tool_parsers/phi4mini_tool_parser.py @@ -28,8 +28,7 @@ class Phi4MiniJsonToolParser(ToolParser): - """ - Tool call parser for phi-4-mini models intended for use with the + """Tool call parser for phi-4-mini models intended for use with the examples/tool_chat_template_llama.jinja template. Used when --enable-auto-tool-choice --tool-call-parser phi4_mini_json @@ -55,9 +54,7 @@ def __init__( def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. - """ + """Extract the tool calls from a complete model response.""" logger.debug("Model output: %s", model_output) pattern = r"functools\[(.*?)\]" diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index 693580e85a2e..178036760634 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GLM-4 Tool Call Parser with incremental string streaming support. +"""GLM-4 Tool Call Parser with incremental string streaming support. This parser fixes the streaming issue reported in Issue #32829 where long string parameters (e.g., file content with 4000+ characters of code) are buffered until diff --git a/vllm/tool_parsers/pythonic_tool_parser.py b/vllm/tool_parsers/pythonic_tool_parser.py index cc584d3d7c95..76b263cfc5bb 100644 --- a/vllm/tool_parsers/pythonic_tool_parser.py +++ b/vllm/tool_parsers/pythonic_tool_parser.py @@ -31,8 +31,7 @@ class PythonicToolParser(ToolParser): - """ - Tool call parser for models that produce tool calls in a pythonic style, + """Tool call parser for models that produce tool calls in a pythonic style, such as Llama 3.2 and Llama 4 models. Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set @@ -69,9 +68,7 @@ def current_tool_index(self, value: int) -> None: def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: - """ - Extract the tool calls from a complete model response. - """ + """Extract the tool calls from a complete model response.""" is_tool_call_pattern = False try: is_tool_call_pattern = ( diff --git a/vllm/tool_parsers/step3_tool_parser.py b/vllm/tool_parsers/step3_tool_parser.py index 929e0c19c600..514945cd0d70 100644 --- a/vllm/tool_parsers/step3_tool_parser.py +++ b/vllm/tool_parsers/step3_tool_parser.py @@ -32,8 +32,7 @@ class Step3ToolParser(ToolParser): - """ - Tool parser for a model that uses a specific XML-like format for tool calls. + """Tool parser for a model that uses a specific XML-like format for tool calls. This version uses a robust, stateful, cursor-based streaming parser and consolidates tool arguments into a single message. """ diff --git a/vllm/tool_parsers/step3p5_tool_parser.py b/vllm/tool_parsers/step3p5_tool_parser.py index 74fcb56ec46f..2639d2f2043a 100644 --- a/vllm/tool_parsers/step3p5_tool_parser.py +++ b/vllm/tool_parsers/step3p5_tool_parser.py @@ -28,8 +28,7 @@ class StreamingXMLToolCallParser: - """ - Simplified streaming XML tool call parser + """Simplified streaming XML tool call parser Supports streaming input, parsing, and output """ @@ -46,8 +45,7 @@ def __init__(self): self.parameter_end_token: str = "" def reset_streaming_state(self): - """Reset streaming parsing state""" - + """Reset streaming parsing state.""" self.deltas = [] # state for streaming self.tool_call_index = 0 @@ -80,8 +78,7 @@ def reset_streaming_state(self): self.setup_parser() def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage: - """ - Parse single streaming XML chunk and return Delta response + """Parse single streaming XML chunk and return Delta response This is the actual streaming interface that receives chunks one by one and maintains internal state @@ -90,6 +87,7 @@ def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage: Returns: DeltaMessage: Contains delta information generated by this chunk, returns empty response if no complete elements + """ # Record delta count before processing initial_delta_count = len(self.deltas) @@ -218,8 +216,7 @@ def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage: return DeltaMessage(content=None) def _escape_xml_special_chars(self, text: str) -> str: - """ - Escape XML special characters + """Escape XML special characters Args: text: Original text Returns: @@ -239,11 +236,11 @@ def _escape_xml_special_chars(self, text: str) -> str: return text def _process_complete_xml_elements(self) -> bool: - """ - Process complete XML elements in buffer + """Process complete XML elements in buffer. Returns: bool: Whether complete elements were found and processed + """ found_any = False @@ -320,8 +317,7 @@ def _process_complete_xml_elements(self) -> bool: return found_any def _fix_incomplete_tag_in_chunk(self, chunk: str) -> str: - """ - Fallback: fix incomplete ) Examples: , @@ -386,8 +382,7 @@ def _fix_incomplete_tag_in_chunk(self, chunk: str) -> str: return chunk def _fix_missing_equals_in_function_tag(self, chunk: str) -> str: - """ - Fix missing = in function tags: or + """Fix missing = in function tags: or Examples: -> -> @@ -425,7 +420,7 @@ def _fix_missing_equals_in_function_tag(self, chunk: str) -> str: return chunk def _validate_function_name(self, func_name: str) -> bool: - """Check if function name exists in tool definitions""" + """Check if function name exists in tool definitions.""" if not self.tools: return False @@ -442,7 +437,7 @@ def _validate_function_name(self, func_name: str) -> bool: return False def _validate_parameter_name(self, param_name: str) -> bool: - """Check if parameter exists in current function's tool definition""" + """Check if parameter exists in current function's tool definition.""" if not self.tools or not self.current_function_name: return True @@ -465,16 +460,15 @@ def _validate_parameter_name(self, param_name: str) -> bool: return True def _should_skip_element(self, element: str) -> bool: - """ - Determine whether an element should be skipped + """Determine whether an element should be skipped. Args: element: Element to evaluate Returns: bool: True means should skip, False means should process - """ + """ # If it's a tool_call XML tag, don't skip if ( element.startswith(self.tool_call_start_token) @@ -501,8 +495,7 @@ def _should_skip_element(self, element: str) -> bool: return not element def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]: - """ - Find next complete XML element from specified position + """Find next complete XML element from specified position. Args: start_pos: Position to start searching @@ -510,6 +503,7 @@ def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]: Returns: (Complete element string, element end position), returns (None, start_pos) if no complete element found + """ buffer = self.streaming_buffer[start_pos:] @@ -588,8 +582,7 @@ def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]: return remaining, start_pos + len(remaining) def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessage: - """ - Merge newly generated deltas from this processing + """Merge newly generated deltas from this processing into a single DeltaMessage Args: @@ -597,6 +590,7 @@ def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessa Returns: Merged DeltaMessage containing all newly generated delta information + """ if len(self.deltas) <= initial_count: return DeltaMessage(content=None) @@ -652,8 +646,7 @@ def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessa ) def _preprocess_xml_chunk(self, chunk: str) -> str: - """ - Preprocess XML chunk, handle non-standard formats, + """Preprocess XML chunk, handle non-standard formats, and escape special characters Args: @@ -661,8 +654,8 @@ def _preprocess_xml_chunk(self, chunk: str) -> str: Returns: Processed XML chunk - """ + """ # Check if this is a tool_call related element is_tool_call = False if chunk.startswith(self.tool_call_start_token) or chunk.startswith( @@ -785,7 +778,7 @@ def _preprocess_xml_chunk(self, chunk: str) -> str: return processed def _emit_delta(self, delta: DeltaMessage): - """Emit Delta response (streaming output)""" + """Emit Delta response (streaming output).""" self.deltas.append(delta) def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None): @@ -814,8 +807,7 @@ def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None): self._end_element("tool_call") def _start_element(self, name: str, attrs: dict[str, str]): - """Handle XML start element events""" - + """Handle XML start element events.""" if name == "root": return @@ -904,7 +896,7 @@ def _start_element(self, name: str, attrs: dict[str, str]): self.current_param_is_first = False def _char_data(self, data: str): - """Handle XML character data events""" + """Handle XML character data events.""" if data and self.current_param_name: # If preprocessing stage determines deferred parsing is needed, # only cache character data, no streaming output @@ -979,8 +971,7 @@ def _char_data(self, data: str): self._emit_delta(delta) def _end_element(self, name: str): - """Handle XML end element events""" - + """Handle XML end element events.""" if name == "root": return @@ -1154,18 +1145,18 @@ def _end_element(self, name: str): self._reset_xml_parser_after_tool_call() def setup_parser(self): - """Set up XML parser event handlers""" + """Set up XML parser event handlers.""" self.parser.buffer_text = True self.parser.StartElementHandler = self._start_element self.parser.EndElementHandler = self._end_element self.parser.CharacterDataHandler = self._char_data def set_tools(self, tools: list[Tool] | None): - """Set tool configuration information""" + """Set tool configuration information.""" self.tools = tools def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None: - """Extract function name from various formats""" + """Extract function name from various formats.""" if attrs and "name" in attrs: return attrs["name"] @@ -1177,7 +1168,7 @@ def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None return None def _extract_parameter_name(self, name: str, attrs: dict[str, str]) -> str | None: - """Extract parameter name from various formats""" + """Extract parameter name from various formats.""" if attrs and "name" in attrs: return attrs["name"] @@ -1195,6 +1186,7 @@ def _get_param_type(self, param_name: str) -> str: Returns: Parameter type + """ if not self.tools or not self.current_function_name: return "string" @@ -1235,6 +1227,7 @@ def repair_param_type(self, param_type: str) -> str: Returns: Repaired parameter type + """ if ( param_type in ["string", "str", "text", "varchar", "char", "enum"] @@ -1264,6 +1257,7 @@ def _convert_param_value(self, param_value: str, param_type: str) -> Any: Returns: Converted value + """ if param_value.lower() == "null": return None @@ -1315,6 +1309,7 @@ def _convert_for_json_streaming(self, converted_value: Any, param_type: str) -> Returns: Converted string for streaming output + """ # Check if value is empty, but exclude numeric 0 if converted_value is None or converted_value == "": @@ -1331,11 +1326,9 @@ def _convert_for_json_streaming(self, converted_value: Any, param_type: str) -> return converted_value def _reset_xml_parser_after_tool_call(self): - """ - Each tool_call is treated as a separate XML document, + """Each tool_call is treated as a separate XML document, so we need to reset the parser after each tool_call. """ - # recreate XML parser self.parser = ParserCreate() self.setup_parser() @@ -1513,7 +1506,5 @@ def extract_tool_calls_streaming( return result def parser_should_check_for_unstreamed_tool_arg_tokens(self) -> bool: - """ - Skip the remaining_call calculation in serving_chat - """ + """Skip the remaining_call calculation in serving_chat.""" return False diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 989250e2ea5b..295d5844b430 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -113,7 +113,6 @@ def get_model_structural_tag( token_suffix: str = "", ) -> StructuralTag | None: """Build a structural tag with xgrammar's builtin model templates.""" - if not tools or tool_choice == "none": return None @@ -158,7 +157,6 @@ def _dump_tool_for_xgrammar( tool: ChatCompletionToolsParam | ResponsesTool, ) -> dict[str, Any]: """Convert tool objects to xgrammar's Chat Completions tool protocol.""" - if isinstance(tool, FunctionTool): function: dict[str, Any] = {"name": tool.name} if tool.description is not None: @@ -178,7 +176,6 @@ def _dump_tool_choice_for_xgrammar( tool_choice: ToolChoice, ) -> dict[str, Any] | str | None: """Convert tool_choice objects to xgrammar's expected protocol.""" - if tool_choice is None: return None @@ -783,6 +780,7 @@ def get_hy_v4_structural_tag( leading colon (e.g. ``":6124c78e"``), or ``""`` when the checkpoint uses unsuffixed tokens. The HYV4 tool parser reads it off the tokenizer vocab and passes it to ``get_model_structural_tag``. + """ del builtin_tools diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index d39c83969f41..1617e25ee174 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -56,8 +56,7 @@ def partial_tag_overlap(text: str, tag: str) -> int: def find_common_prefix(s1: str, s2: str) -> str: - """ - Finds a common prefix that is shared between two strings, if there is one. + """Finds a common prefix that is shared between two strings, if there is one. Order of arguments is NOT important. This function is provided as a UTILITY for extracting information from JSON @@ -79,8 +78,7 @@ def find_common_prefix(s1: str, s2: str) -> str: def find_common_suffix(s1: str, s2: str) -> str: - """ - Finds a common suffix shared between two strings, if there is one. Order of + """Finds a common suffix shared between two strings, if there is one. Order of arguments is NOT important. Stops when the suffix ends OR it hits an alphanumeric character @@ -97,8 +95,7 @@ def find_common_suffix(s1: str, s2: str) -> str: def extract_intermediate_diff(curr: str, old: str) -> str: - """ - Given two strings, extract the difference in the middle between two strings + """Given two strings, extract the difference in the middle between two strings that are known to have a common prefix and/or suffix. This function is provided as a UTILITY for extracting information from JSON @@ -429,7 +426,8 @@ def get_json_schema_from_tools( class UnexpectedAstError(Exception): """Raised when the AST structure does not match the expected - pythonic tool call format.""" + pythonic tool call format. + """ pass @@ -450,6 +448,7 @@ def get_parameter_value(val: ast.expr) -> Any: Raises: UnexpectedAstError: If the AST node is not a supported literal type. + """ if isinstance(val, ast.Constant): if val.value is None or isinstance(val.value, (str, int, float)): @@ -528,6 +527,7 @@ def _ast_callable_dotted_name(node: ast.expr) -> str: Raises: UnexpectedAstError: If the chain does not bottom out in an ``ast.Name`` (e.g. subscript or call expression as receiver). + """ parts: list[str] = [] current: ast.expr = node @@ -550,6 +550,7 @@ def handle_single_tool(call: ast.Call) -> ToolCall: Raises: UnexpectedAstError: If the call target is neither a simple name nor a chain of attribute accesses bottoming out in a name. + """ if not isinstance(call.func, (ast.Name, ast.Attribute)): logger.warning( @@ -921,6 +922,7 @@ def make_valid_python(text: str) -> tuple[str, str] | None: Raises: UnexpectedAstError: If mismatched brackets or parentheses are detected. + """ bracket_stack: list[str] = [] for index, char in enumerate(text): @@ -1088,6 +1090,7 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any: value: The raw string value from the model output. schema_type: One or more JSON Schema type strings (e.g. ``"string"`` or ``["string", "null"]``). + """ if isinstance(schema_type, str): schema_type = [schema_type] @@ -1176,6 +1179,7 @@ def compute_tool_delta( Returns: A DeltaToolCall with only the new argument characters, or None if there is no difference from what was previously sent. + """ new_call_args = new_call.function.arguments if withheld_suffix: diff --git a/vllm/tracing/__init__.py b/vllm/tracing/__init__.py index 4f025970ecbd..9725b57ba979 100644 --- a/vllm/tracing/__init__.py +++ b/vllm/tracing/__init__.py @@ -94,9 +94,7 @@ def instrument( attributes: dict[str, str] | None = None, record_exception: bool = True, ): - """ - Generic decorator to instrument functions. - """ + """Generic decorator to instrument functions.""" if obj is None: return functools.partial( instrument, @@ -135,6 +133,7 @@ def instrument_manual( attributes: Optional dict of span attributes. context: Optional trace context (e.g., from extract_trace_context). kind: Optional SpanKind (e.g., SpanKind.SERVER). + """ is_available, _, _, _, manual_instrument_fn = _REGISTERED_TRACING_BACKENDS["otel"] if is_available(): @@ -146,8 +145,7 @@ def instrument_manual( def is_tracing_available() -> bool: - """ - Returns True if any tracing backend (OTel, Profiler, etc.) is available. + """Returns True if any tracing backend (OTel, Profiler, etc.) is available. Use this to guard expensive tracing logic in the main code. """ check_available = [ diff --git a/vllm/tracing/otel.py b/vllm/tracing/otel.py index ac06ae97255f..36c189b88324 100644 --- a/vllm/tracing/otel.py +++ b/vllm/tracing/otel.py @@ -107,9 +107,7 @@ def init_otel_worker_tracer( process_kind: str, process_name: str, ) -> Tracer: - """ - Backend-specific initialization for OpenTelemetry in a worker process. - """ + """Backend-specific initialization for OpenTelemetry in a worker process.""" # Initialize the tracer if an OTLP endpoint is configured. # The endpoint is propagated via environment variable from the main process. otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") @@ -133,7 +131,6 @@ def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None: def instrument_otel(func, span_name, attributes, record_exception): """Internal wrapper logic for sync and async functions.""" - # Pre-calculate static code attributes once (these don't change) code_attrs = { LoadingSpanAttributes.CODE_FUNCTION: func.__qualname__, @@ -214,8 +211,7 @@ def manual_instrument_otel( def _get_smart_context() -> Context | None: - """ - Determines the parent context. + """Determines the parent context. 1. If a Span is already active in this process, use it. 2. If not, extract from os.environ, handling the case-sensitivity mismatch. """ @@ -239,8 +235,7 @@ def _get_smart_context() -> Context | None: @contextmanager def propagate_trace_to_env(): - """ - Temporarily injects the current OTel context into os.environ. + """Temporarily injects the current OTel context into os.environ. This ensures that any subprocesses (like vLLM workers) spawned within this context inherit the correct traceparent. """ diff --git a/vllm/tracing/utils.py b/vllm/tracing/utils.py index 0e11f850b442..138b801f5125 100644 --- a/vllm/tracing/utils.py +++ b/vllm/tracing/utils.py @@ -13,8 +13,7 @@ class SpanAttributes: - """ - Standard attributes for spans. + """Standard attributes for spans. These are largely based on OpenTelemetry Semantic Conventions but are defined here as constants so they can be used by any backend or logger. @@ -60,8 +59,7 @@ def contains_trace_headers(headers: Mapping[str, str]) -> bool: def extract_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]: - """ - Extract only trace-related headers from a larger header dictionary. + """Extract only trace-related headers from a larger header dictionary. Useful for logging or passing context to a non-OTel client. """ return {h: headers[h] for h in TRACE_HEADERS if h in headers} diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 565caeebc286..24111f52a34d 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -231,7 +231,6 @@ def _patch_hf_transformers_validate_rope(): validate_rope() with the ignore_keys parameter work with newer versions of hf transformers (from v5 onwards) """ - if hasattr(PretrainedConfig.validate_rope, "__vllm_patched__"): return @@ -489,10 +488,11 @@ def register_config_parser(config_format: str): """Register a customized vllm config parser. When a config format is not supported by vllm, you can register a customized config parser to support it. - Args: + + Args: config_format (str): The config parser format name. - Examples: + Examples: >>> from vllm.transformers_utils.config import (get_config_parser, register_config_parser) >>> from vllm.transformers_utils.config_parser_base import ConfigParserBase @@ -511,6 +511,7 @@ def register_config_parser(config_format: str): >>> >>> type(get_config_parser("custom_config_parser")) + """ # noqa: E501 def _wrapper(config_parser_cls): @@ -538,7 +539,8 @@ def _wrapper(config_parser_cls): def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> None: """Some models may have no rope_theta in their config but still use RoPE. - This function sets a default rope_theta if it's missing.""" + This function sets a default rope_theta if it's missing. + """ if getattr(config, "rope_parameters", None) is None: config.rope_parameters = {"rope_type": "default"} if "rope_theta" not in config.rope_parameters: @@ -547,8 +549,8 @@ def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> No def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None: """Patch legacy RoPE type fields for backwards compatibility with - older custom models which would otherwise fail to load.""" - + older custom models which would otherwise fail to load. + """ # No RoPE parameters to patch if rope_parameters is None: return @@ -721,9 +723,7 @@ def _is_encoder_decoder(config: PretrainedConfig) -> bool: def _maybe_update_auto_config_kwargs(kwargs: dict[str, Any], model_type: str): - """ - Update kwargs for AutoConfig initialization based on model_type - """ + """Update kwargs for AutoConfig initialization based on model_type.""" if model_type in _AUTO_CONFIG_KWARGS_OVERRIDES: kwargs.update(_AUTO_CONFIG_KWARGS_OVERRIDES[model_type]) return kwargs @@ -748,8 +748,7 @@ def maybe_override_with_speculators( hf_token: bool | str | None = None, **kwargs, ) -> tuple[str, str | None, dict[str, Any] | None]: - """ - Resolve model configuration when speculators are detected. + """Resolve model configuration when speculators are detected. Checks if the provided model is a speculators model and if so, extracts the target model configuration and builds the speculative config. @@ -764,6 +763,7 @@ def maybe_override_with_speculators( Returns: Tuple of (resolved_model, resolved_tokenizer, speculative_config) + """ kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE config_dict, _ = PretrainedConfig.get_config_dict( @@ -938,8 +938,7 @@ def get_pooling_config( model: str, revision: str | None = "main", ) -> dict[str, Any] | None: - """ - This function gets the pooling and normalize + """This function gets the pooling and normalize config from the model - only applies to sentence-transformers models. @@ -951,6 +950,7 @@ def get_pooling_config( Returns: A dictionary containing the pooling type and whether normalization is used, or None if no pooling configuration is found. + """ modules_file_name = "modules.json" @@ -1035,19 +1035,17 @@ def parse_pooling_type(pooling_name: str): def get_sentence_transformer_tokenizer_config( model: str | Path, revision: str | None = "main" ) -> dict[str, Any] | None: - """ - Returns the tokenization configuration dictionary for a + """Returns the tokenization configuration dictionary for a given Sentence Transformer BERT model. - Parameters: - - model (str|Path): The name of the Sentence Transformer - BERT model. - - revision (str, optional): The revision of the m - odel to use. Defaults to 'main'. + Args: + model (str|Path): The name of the Sentence Transformer BERT model. + revision (str, optional): The revision of the model to use. Defaults to 'main'. Returns: - dict: A dictionary containing the configuration parameters for the Sentence Transformer BERT model. + """ sentence_transformer_config_files = [ "sentence_bert_config.json", @@ -1308,9 +1306,7 @@ def get_safetensors_params_metadata( *, revision: str | None = None, ) -> dict[str, Any]: - """ - Get the safetensors parameters metadata for remote/local model repository. - """ + """Get the safetensors parameters metadata for remote/local model repository.""" if (model_path := Path(model)).exists(): return _read_safetensors_metadata_in_dir(model_path) @@ -1354,6 +1350,7 @@ def checkpoint_has_lm_head(model: str, *, revision: str | None = None) -> bool | Returns: `None` if the checkpoint contents could not be determined, for example because it is not stored as safetensors. + """ metadata = get_safetensors_params_metadata(model, revision=revision) if not metadata: diff --git a/vllm/transformers_utils/configs/AXK1.py b/vllm/transformers_utils/configs/AXK1.py index 09a3a0f46ae6..52af5bec48de 100644 --- a/vllm/transformers_utils/configs/AXK1.py +++ b/vllm/transformers_utils/configs/AXK1.py @@ -6,14 +6,14 @@ class AXK1Config(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`AXK1Model`]. + r"""This is the configuration class to store the configuration of a [`AXK1Model`]. It is used to instantiate an A.X model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the A.X K1. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. + Args: vocab_size (`int`, *optional*, defaults to 163840): Vocabulary size of the A.X K1 model. Defines the number of different @@ -112,6 +112,7 @@ class AXK1Config(PretrainedConfig): layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. + """ model_type = "axk1" diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 3a6874820ea1..a302d5d72928 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Model configs may be defined in this directory for the following reasons: +"""Model configs may be defined in this directory for the following reasons: - There is no configuration file defined by HF Hub or Transformers library. - There is a need to override the existing config to support vLLM. diff --git a/vllm/transformers_utils/configs/colpali.py b/vllm/transformers_utils/configs/colpali.py index c40c58b25ce1..a4e0a83e2929 100644 --- a/vllm/transformers_utils/configs/colpali.py +++ b/vllm/transformers_utils/configs/colpali.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColPali configuration that extends PaliGemmaConfig with embedding projection +"""ColPali configuration that extends PaliGemmaConfig with embedding projection fields. This allows ColPali models to be loaded without trust_remote_code by mapping their custom model_type (colpali) to a standard config class that vLLM understands. diff --git a/vllm/transformers_utils/configs/colqwen3.py b/vllm/transformers_utils/configs/colqwen3.py index 1c09a0a91845..c98995491d8c 100644 --- a/vllm/transformers_utils/configs/colqwen3.py +++ b/vllm/transformers_utils/configs/colqwen3.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ColQwen3 configuration that extends Qwen3VLConfig with embedding projection +"""ColQwen3 configuration that extends Qwen3VLConfig with embedding projection fields. This allows ColQwen3 models to be loaded without trust_remote_code by mapping their custom model_type (colqwen3, ops_colqwen3, etc.) to a standard config class that vLLM understands. diff --git a/vllm/transformers_utils/configs/extract_hidden_states.py b/vllm/transformers_utils/configs/extract_hidden_states.py index 2beec0e3081b..ab705d583462 100644 --- a/vllm/transformers_utils/configs/extract_hidden_states.py +++ b/vllm/transformers_utils/configs/extract_hidden_states.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Config definitions for ExtractHiddenStatesModel, to be used with -the extract_hidden_states spec decoding method.""" +the extract_hidden_states spec decoding method. +""" import os diff --git a/vllm/transformers_utils/configs/falcon.py b/vllm/transformers_utils/configs/falcon.py index c646d241d4eb..f8665a139559 100644 --- a/vllm/transformers_utils/configs/falcon.py +++ b/vllm/transformers_utils/configs/falcon.py @@ -18,7 +18,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Falcon configuration""" +"""Falcon configuration.""" from transformers.configuration_utils import PretrainedConfig diff --git a/vllm/transformers_utils/configs/hy_v3.py b/vllm/transformers_utils/configs/hy_v3.py index 9425caf4e03f..0a67dd8af793 100644 --- a/vllm/transformers_utils/configs/hy_v3.py +++ b/vllm/transformers_utils/configs/hy_v3.py @@ -6,8 +6,7 @@ class HYV3Config(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`HYV3Model`]. + r"""This is the configuration class to store the configuration of a [`HYV3Model`]. It is used to instantiate a HYV3 model (HY V3 MoE language model) according to the specified arguments. @@ -93,6 +92,7 @@ class HYV3Config(PretrainedConfig): >>> config = HYV3Config() >>> model = HYV3Model(config) ``` + """ model_type = "hy_v3" diff --git a/vllm/transformers_utils/configs/hy_v4.py b/vllm/transformers_utils/configs/hy_v4.py index beca23ee7437..687a2206c878 100644 --- a/vllm/transformers_utils/configs/hy_v4.py +++ b/vllm/transformers_utils/configs/hy_v4.py @@ -52,6 +52,7 @@ class HYV4Config(PreTrainedConfig): below zero disable the clamp. num_nextn_predict_layers: Number of MTP layers in the checkpoint. mtp_loss_factor: Training-time MTP loss weight (unused at inference). + """ model_type = "hy_v4" diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py index 9fa823743d66..2033d59c1158 100644 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ b/vllm/transformers_utils/configs/hyperclovax.py @@ -21,8 +21,7 @@ class HyperCLOVAXConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a + r"""This is the configuration class to store the configuration of a [`HyperCLOVAXModel`]. It is used to instantiate a HyperCLOVAX model according to the specified arguments, defining the model architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used @@ -166,6 +165,7 @@ class HyperCLOVAXConfig(PretrainedConfig): `get_rope`. When provided, takes precedence over `rope_theta` and `rope_scaling`. If `None`, it is derived from `rope_theta` and `rope_scaling` automatically. + """ model_type = "hyperclovax" diff --git a/vllm/transformers_utils/configs/kimi_k25.py b/vllm/transformers_utils/configs/kimi_k25.py index 710e9b56367f..095db1d8c46e 100644 --- a/vllm/transformers_utils/configs/kimi_k25.py +++ b/vllm/transformers_utils/configs/kimi_k25.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Kimi-K2.5 Model Configuration. +"""Kimi-K2.5 Model Configuration. This configuration supports video-chunk as an internal modality type. A video-chunk is the smallest independently processable unit of video. @@ -73,6 +72,7 @@ class KimiK25Config(PretrainedConfig): ignore_index: The ignore index for the loss function. media_placeholder_token_id: The token ID for media placeholders. pad_token_id: The token ID for padding. + """ model_type = "kimi_k25" diff --git a/vllm/transformers_utils/configs/lfm2_moe.py b/vllm/transformers_utils/configs/lfm2_moe.py index b399a03c030f..34682a62a1fc 100644 --- a/vllm/transformers_utils/configs/lfm2_moe.py +++ b/vllm/transformers_utils/configs/lfm2_moe.py @@ -6,8 +6,7 @@ class Lfm2MoeConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe + r"""This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the LFM2-8B-A1B model. e.g. [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B) @@ -85,7 +84,9 @@ class Lfm2MoeConfig(PretrainedConfig): >>> # Accessing the model configuration >>> configuration = model.config - ```""" # noqa: E501 + ``` + + """ # noqa: E501 model_type = "lfm2_moe" keys_to_ignore_at_inference = ["past_key_values"] diff --git a/vllm/transformers_utils/configs/mlp_speculator.py b/vllm/transformers_utils/configs/mlp_speculator.py index 75745f227f48..6b5a70e3dbf4 100644 --- a/vllm/transformers_utils/configs/mlp_speculator.py +++ b/vllm/transformers_utils/configs/mlp_speculator.py @@ -24,8 +24,7 @@ def __init__( scale_input: bool = False, **kwargs, ): - """ - Initialize an MLPSpeculatorConfig + """Initialize an MLPSpeculatorConfig. Args: vocab_size: int @@ -52,6 +51,7 @@ def __init__( scale_input: bool if True, will scale the initial hidden states from the base model. + """ if top_k_tokens_per_head is None: top_k_tokens_per_head = [5, 4, 3] diff --git a/vllm/transformers_utils/configs/moondream3.py b/vllm/transformers_utils/configs/moondream3.py index 307bb2977206..8bff5cd10ee7 100644 --- a/vllm/transformers_utils/configs/moondream3.py +++ b/vllm/transformers_utils/configs/moondream3.py @@ -148,5 +148,6 @@ def get_text_config(self, decoder: bool = False) -> "Moondream3TextConfig": Args: decoder: Ignored. Only used for encoder-decoder models. + """ return self.text_config diff --git a/vllm/transformers_utils/configs/nemotron.py b/vllm/transformers_utils/configs/nemotron.py index 62f52703029b..bb15a8b433d9 100644 --- a/vllm/transformers_utils/configs/nemotron.py +++ b/vllm/transformers_utils/configs/nemotron.py @@ -15,7 +15,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Nemotron model configuration""" +"""Nemotron model configuration.""" from transformers import PretrainedConfig from transformers.utils import logging @@ -24,8 +24,7 @@ class NemotronConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a + r"""This is the configuration class to store the configuration of a [`NemotronModel`]. It is used to instantiate a Nemotron model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar @@ -114,7 +113,9 @@ class NemotronConfig(PretrainedConfig): >>> model = NemotronModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config - ```""" + ``` + + """ model_type = "nemotron" keys_to_ignore_at_inference = ["past_key_values"] @@ -193,9 +194,7 @@ def __init__( ) def _rope_parameters_validation(self): - """ - Validate the `rope_parameters` configuration. - """ + """Validate the `rope_parameters` configuration.""" if self.rope_parameters is None: return diff --git a/vllm/transformers_utils/configs/nemotron_h.py b/vllm/transformers_utils/configs/nemotron_h.py index ed62b5d294b3..c83313edfb48 100644 --- a/vllm/transformers_utils/configs/nemotron_h.py +++ b/vllm/transformers_utils/configs/nemotron_h.py @@ -15,7 +15,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""NemotronH model configuration""" +"""NemotronH model configuration.""" import regex as re from transformers.configuration_utils import PretrainedConfig @@ -25,12 +25,12 @@ class NemotronHConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a + r"""This is the configuration class to store the configuration of a [`NemotronHModel`]. It is used to instantiate a NemotronH model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the NemotronH-v0.1 model. + Args: vocab_size (`int`, *optional*, defaults to 131072): Vocabulary size of the NemotronH model. Defines the number of @@ -139,6 +139,7 @@ class NemotronHConfig(PretrainedConfig): Size of chunks for Mamba processing. rescale_prenorm_residual (`bool`, *optional*, defaults to `True`): Whether to rescale the pre-normalization residual connections. + """ model_type = "nemotron_h" diff --git a/vllm/transformers_utils/configs/olmo_hybrid.py b/vllm/transformers_utils/configs/olmo_hybrid.py index cdca81757e70..6e12cdd27402 100644 --- a/vllm/transformers_utils/configs/olmo_hybrid.py +++ b/vllm/transformers_utils/configs/olmo_hybrid.py @@ -6,8 +6,7 @@ class OlmoHybridConfig(PretrainedConfig): - r""" - Configuration class for [`OlmoHybridModel`]. It is used to + r"""Configuration class for [`OlmoHybridModel`]. It is used to instantiate an OLMo Hybrid model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar @@ -19,7 +18,7 @@ class OlmoHybridConfig(PretrainedConfig): can be used to control the model outputs. Read the documentation from [`PreTrainedConfig`] for more information. - Args: + Args: vocab_size (`int`, *optional*, defaults to 100352): Vocabulary size of the OlmoHybrid model. Defines the number of different tokens that can be @@ -159,6 +158,7 @@ class OlmoHybridConfig(PretrainedConfig): >>> model = OlmoHybridModel(configuration) >>> configuration = model.config ``` + """ model_type = "olmo_hybrid" diff --git a/vllm/transformers_utils/configs/ovis.py b/vllm/transformers_utils/configs/ovis.py index 294b4c9037aa..64629566a399 100644 --- a/vllm/transformers_utils/configs/ovis.py +++ b/vllm/transformers_utils/configs/ovis.py @@ -14,6 +14,7 @@ class AIMv2Config(PretrainedConfig): """This is the configuration class to store the configuration of an [`AIMv2Model`]. Instantiating a configuration with the defaults will yield a similar configuration to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224). + Args: hidden_size: Dimension of the hidden representations. intermediate_size: Dimension of the SwiGLU representations. @@ -29,6 +30,7 @@ class AIMv2Config(PretrainedConfig): qkv_bias: Whether to add a bias to the queries, keys and values. use_bias: Whether to add a bias in the feed-forward and projection layers. kwargs: Keyword arguments for the [`PretrainedConfig`]. + """ model_type: str = "aimv2" diff --git a/vllm/transformers_utils/configs/qwen3_5.py b/vllm/transformers_utils/configs/qwen3_5.py index d5820a5783c6..e30cbe6502d4 100644 --- a/vllm/transformers_utils/configs/qwen3_5.py +++ b/vllm/transformers_utils/configs/qwen3_5.py @@ -14,7 +14,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen3.5 model configuration""" +"""Qwen3.5 model configuration.""" from transformers.configuration_utils import PretrainedConfig diff --git a/vllm/transformers_utils/configs/qwen3_5_moe.py b/vllm/transformers_utils/configs/qwen3_5_moe.py index ec229ce81426..938272d0280b 100644 --- a/vllm/transformers_utils/configs/qwen3_5_moe.py +++ b/vllm/transformers_utils/configs/qwen3_5_moe.py @@ -14,7 +14,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen3.5-MoE model configuration""" +"""Qwen3.5-MoE model configuration.""" from transformers.configuration_utils import PretrainedConfig diff --git a/vllm/transformers_utils/configs/qwen3_next.py b/vllm/transformers_utils/configs/qwen3_next.py index de579ed2cf3f..8460725bfd04 100644 --- a/vllm/transformers_utils/configs/qwen3_next.py +++ b/vllm/transformers_utils/configs/qwen3_next.py @@ -14,7 +14,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen3-Next model configuration""" +"""Qwen3-Next model configuration.""" from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging @@ -23,8 +23,7 @@ class Qwen3NextConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a + r"""This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a Qwen3-Next model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct). @@ -157,6 +156,7 @@ class Qwen3NextConfig(PretrainedConfig): >>> # Accessing the model configuration >>> configuration = model.config ``` + """ # noqa: E501 model_type = "qwen3_next" diff --git a/vllm/transformers_utils/configs/radio.py b/vllm/transformers_utils/configs/radio.py index e668c5c5e7f2..4a3e468f84f3 100644 --- a/vllm/transformers_utils/configs/radio.py +++ b/vllm/transformers_utils/configs/radio.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Radio vision model configuration""" +"""Radio vision model configuration.""" from typing import Any @@ -21,8 +21,7 @@ class RadioConfig(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a Radio + r"""This is the configuration class to store the configuration of a Radio vision model. It is used to instantiate a Radio model according to the specified arguments, defining the model architecture. @@ -55,6 +54,7 @@ class RadioConfig(PretrainedConfig): dedicated video patch embedder (3*T*P*P -> hidden) separate from the image embedder (3*P*P -> hidden). When False, a single embedder with input size 3*T*P*P is used for both (images are duplicated T times). + """ model_type = "radio" diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index 9f45183328a9..6407fc4868fa 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -14,8 +14,7 @@ def decorator(fn): @register_speculator("eagle3") def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: - """ - Apply Eagle-3 specific configuration transformations to the `dict` used to + """Apply Eagle-3 specific configuration transformations to the `dict` used to construct the Transformers PreTrainedConfig. Eagle-3 specific fields: @@ -28,7 +27,6 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: provide intermediate hidden states that help the drafter make better predictions. This is the standard field used in Eagle3 checkpoints. """ - pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") if config_dict.get("target_hidden_size") is not None: pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"] @@ -54,8 +52,7 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: @register_speculator("peagle") def update_peagle(config_dict: dict, pre_trained_config: dict) -> None: - """ - Apply PEagle (Parallel Eagle) specific configuration transformations to + """Apply PEagle (Parallel Eagle) specific configuration transformations to the `dict` used to construct the Transformers PreTrainedConfig. PEagle specific fields: @@ -92,8 +89,7 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None: @register_speculator("dflash") def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: - """ - Apply DFlash specific configuration transformations to the `dict` used to + """Apply DFlash specific configuration transformations to the `dict` used to construct the Transformers PreTrainedConfig. DFlash specific fields: @@ -132,8 +128,7 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: @register_speculator("dflash2") def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None: - """ - Apply DFlash2 specific configuration transformations to the `dict` used + """Apply DFlash2 specific configuration transformations to the `dict` used to construct the Transformers PreTrainedConfig. DFlash2 extends DFlash with a grouped convolution layer and a @@ -156,8 +151,7 @@ def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None: @register_speculator("dspark") def update_dspark(config_dict: dict, pre_trained_config: dict) -> None: - """ - Apply DSpark specific configuration transformations to the `dict` used to + """Apply DSpark specific configuration transformations to the `dict` used to construct the Transformers PreTrainedConfig. DSpark extends DFlash with a Markov logit-bias head, reusing the same diff --git a/vllm/transformers_utils/configs/speculators/base.py b/vllm/transformers_utils/configs/speculators/base.py index 1cbaf6e801e1..a867749dbaac 100644 --- a/vllm/transformers_utils/configs/speculators/base.py +++ b/vllm/transformers_utils/configs/speculators/base.py @@ -47,9 +47,7 @@ def from_pretrained( def extract_transformers_pre_trained_config( cls, config_dict: dict[str, Any] ) -> dict[str, Any]: - """ - Extract standard Transformers PreTrainedConfig config from speculators config. - """ + """Extract the Transformers `PreTrainedConfig` from a speculators config.""" speculators_model_type = config_dict.get("speculators_model_type") if speculators_model_type not in SUPPORTED_SPECULATORS_TYPES: raise ValueError( @@ -99,8 +97,7 @@ def validate_speculators_config(cls, config_dict: dict[str, Any]) -> None: def build_vllm_speculative_config( cls, config_dict: dict[str, Any] ) -> dict[str, Any]: - """ - Build vLLM-compatible speculative configuration from speculators format. + """Build vLLM-compatible speculative configuration from speculators format. This method extracts and transforms speculative configuration from the speculators format into the structure expected by vLLM. @@ -110,6 +107,7 @@ def build_vllm_speculative_config( Returns: Dictionary with vLLM-compatible speculative configuration + """ # Extract speculators configuration spec_config = config_dict["speculators_config"] diff --git a/vllm/transformers_utils/configs/ultravox.py b/vllm/transformers_utils/configs/ultravox.py index 31b49b9d993f..b8bd0031994a 100644 --- a/vllm/transformers_utils/configs/ultravox.py +++ b/vllm/transformers_utils/configs/ultravox.py @@ -8,8 +8,7 @@ class UltravoxConfig(transformers.PretrainedConfig): - r""" - This is the configuration class to store the configuration of a + r"""This is the configuration class to store the configuration of a [`UltravoxForConditionalGeneration`]. It is used to instantiate an Ultravox model according to the specified arguments, defining the model architecture. @@ -41,6 +40,7 @@ class UltravoxConfig(transformers.PretrainedConfig): Whether to apply layer normalization at the middle of the projector or at the end. Versions v0.4.1 and below use `False`, but v0.5 and above use `True`. + """ model_type = "ultravox" diff --git a/vllm/transformers_utils/dynamic_module.py b/vllm/transformers_utils/dynamic_module.py index f0702b15bb7e..f42d5752bf57 100644 --- a/vllm/transformers_utils/dynamic_module.py +++ b/vllm/transformers_utils/dynamic_module.py @@ -29,8 +29,7 @@ def try_get_class_from_dynamic_module( warn_on_fail: bool = True, **kwargs, ) -> type | None: - """ - As `transformers.dynamic_module_utils.get_class_from_dynamic_module`, + """As `transformers.dynamic_module_utils.get_class_from_dynamic_module`, but ignoring any errors. """ try: diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 8fd924fcf381..f8cb761848c1 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -37,7 +37,8 @@ def _transformers_v4_compatibility_import(): subset of `ProcessorChatTemplateKwargs` as defined in Transformers v4. In Transformers v5 these were merged into `ProcessorChatTemplateKwargs` and `ChatTemplateLoadKwargs` was removed. For backward compatibility, we add an alias - for `ChatTemplateLoadKwargs` if it doesn't exist.""" + for `ChatTemplateLoadKwargs` if it doesn't exist. + """ old_import = getattr(processing_utils, "ChatTemplateLoadKwargs", None) new_import = getattr(processing_utils, "ProcessorChatTemplateKwargs", None) if old_import is None and new_import is not None: @@ -52,7 +53,8 @@ def _transformers_v4_compatibility_init() -> Any: processor instance before calling the original `ProcessorMixin.__init__`. This can be removed if `Molmo2ForConditionalGeneration` is upstreamed to - Transformers.""" + Transformers. + """ if hasattr(ProcessorMixin.__init__, "_vllm_patched"): return @@ -80,9 +82,7 @@ def __init__(self, *args, **kwargs): class HashableDict(dict): - """ - A dictionary that can be hashed by lru_cache. - """ + """A dictionary that can be hashed by lru_cache.""" # NOTE: pythonic dict is not hashable, # we override on it directly for simplicity @@ -91,9 +91,7 @@ def __hash__(self) -> int: # type: ignore[override] class HashableList(list): - """ - A list that can be hashed by lru_cache. - """ + """A list that can be hashed by lru_cache.""" def __hash__(self) -> int: # type: ignore[override] return hash(tuple(self)) @@ -399,7 +397,8 @@ def get_feature_extractor( **kwargs: Any, ): """Load an audio feature extractor for the given model name - via HuggingFace.""" + via HuggingFace. + """ try: processor_name = convert_model_repo_to_path(processor_name) feature_extractor = AutoFeatureExtractor.from_pretrained( diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index 249a4e58e662..b46183f4e12a 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Multi-modal processors may be defined in this directory for the following +"""Multi-modal processors may be defined in this directory for the following reasons: - There is no processing file defined by HF Hub or Transformers library. diff --git a/vllm/transformers_utils/processors/bagel.py b/vllm/transformers_utils/processors/bagel.py index 3226d7b0c83d..b6fefe032759 100644 --- a/vllm/transformers_utils/processors/bagel.py +++ b/vllm/transformers_utils/processors/bagel.py @@ -18,8 +18,7 @@ class BagelProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call- class BagelProcessor(ProcessorMixin): - """ - Constructs a BAGEL processor which wraps a + """Constructs a BAGEL processor which wraps a SigLIP image processor and a Qwen2 tokenizer. """ @@ -36,9 +35,7 @@ def __call__( images: ImageInput = None, **kwargs: Unpack[BagelProcessorKwargs], ): - """ - Main method to prepare for the model one or several sequences(s) and image(s). - """ + """Prepare one or more sequences and images for the model.""" output_kwargs = self._merge_kwargs( BagelProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, @@ -62,15 +59,11 @@ def __call__( return BatchFeature(data={**pixel_values, **text_inputs}) def batch_decode(self, *args, **kwargs): - """ - This method forwards all its arguments to Qwen2TokenizerFast's batch_decode. - """ + """Forward all arguments to `Qwen2TokenizerFast.batch_decode`.""" return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): - """ - This method forwards all its arguments to Qwen2TokenizerFast's decode. - """ + """This method forwards all its arguments to Qwen2TokenizerFast's decode.""" return self.tokenizer.decode(*args, **kwargs) @property diff --git a/vllm/transformers_utils/processors/cohere_asr.py b/vllm/transformers_utils/processors/cohere_asr.py index e1257de4e735..bfa8fca50dde 100644 --- a/vllm/transformers_utils/processors/cohere_asr.py +++ b/vllm/transformers_utils/processors/cohere_asr.py @@ -228,7 +228,7 @@ def filter_banks(self): return self.fb def splice_frames(self, x, frame_splicing): - """Stacks frames together across feature dim + """Stacks frames together across feature dim. input is batch_size, feature_dim, num_frames output is batch_size, feature_dim*frame_splicing, num_frames @@ -537,7 +537,8 @@ def __call__( class CohereASRProcessor(ProcessorMixin): """HF-compatible processor combining CohereASRFeatureExtractor and a - tokenizer.""" + tokenizer. + """ feature_extractor_class = "CohereASRFeatureExtractor" tokenizer_class = "AutoTokenizer" diff --git a/vllm/transformers_utils/processors/cosmos3_edge.py b/vllm/transformers_utils/processors/cosmos3_edge.py index 95ca2cccd531..a8751cdf6d52 100644 --- a/vllm/transformers_utils/processors/cosmos3_edge.py +++ b/vllm/transformers_utils/processors/cosmos3_edge.py @@ -70,8 +70,7 @@ def _resize_image( resample: PILImageResampling | InterpolationMode | int | None = None, **kwargs, ) -> torch.Tensor: - """ - Rescales the image so that the following conditions are met: + """Rescales the image so that the following conditions are met: 1. Both dimensions (height and width) are divisible by 'factor'. 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. @@ -322,6 +321,7 @@ class Cosmos3EdgeProcessor(Qwen3VLProcessor): normalize, and patch video inputs. chat_template: An optional Jinja template used to render chat messages into a tokenizable string. + """ attributes = ["image_processor", "tokenizer", "video_processor"] diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index 8a9e36448d83..55af43a70957 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -234,9 +234,7 @@ def process_one( images: list[Image.Image], crop_mode: bool = CROP_MODE, ): - """ - - Args: + """Args: prompt (str): the formatted prompt; images (List[ImageType]): the list of images; crop_mode (bool): if True, then crop the image; @@ -248,8 +246,8 @@ def process_one( - pixel_values (torch.FloatTensor): [n_patches, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (List[int]): the number of image tokens - """ + """ if prompt is None or images is None: raise ValueError("prompt and images must be used at the same time.") @@ -309,7 +307,6 @@ def tokenize_with_images( cropping: bool = True, ): """Tokenize text with tags.""" - num_image_tags = conversation.count(self.image_token) if num_image_tags != len(images): raise ValueError( diff --git a/vllm/transformers_utils/processors/deepseek_vl2.py b/vllm/transformers_utils/processors/deepseek_vl2.py index be6d9800a918..6c9966abf8d1 100644 --- a/vllm/transformers_utils/processors/deepseek_vl2.py +++ b/vllm/transformers_utils/processors/deepseek_vl2.py @@ -186,9 +186,7 @@ def process_one( inference_mode: bool = True, **kwargs: Any, ): - """ - - Args: + """Args: prompt (str): the formatted prompt; images (list[ImageType]): the list of images; inference_mode (bool): if True, then remove the last eos token; @@ -201,8 +199,8 @@ def process_one( - pixel_values (torch.FloatTensor): [n_patches, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (list[int]): the number of image tokens - """ + """ if prompt is None or images is None: raise ValueError("prompt and images must be used at the same time.") @@ -276,9 +274,7 @@ def __call__( inference_mode: bool = True, **kwargs: Any, ): - """ - - Args: + """Args: text (str): the formatted prompt; images (list[ImageType]): the list of images; inference_mode (bool): if True, then remove the last eos token; @@ -290,8 +286,8 @@ def __call__( - images (torch.FloatTensor): [n_images, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (list[int]): the number of image tokens - """ + """ prepare = self.process_one( prompt=text, images=images, diff --git a/vllm/transformers_utils/processors/fireredasr2.py b/vllm/transformers_utils/processors/fireredasr2.py index bba7e7ee0495..6225e88e6a70 100644 --- a/vllm/transformers_utils/processors/fireredasr2.py +++ b/vllm/transformers_utils/processors/fireredasr2.py @@ -68,8 +68,7 @@ def __call__(self, sample_rate, wav_np, is_train=False): class FireRedASR2FeatureExtractor(SequenceFeatureExtractor): - r""" - Constructs a FireRedASR2 feature extractor. + r"""Constructs a FireRedASR2 feature extractor. This feature extractor inherits from [`~feature_extraction_sequence_ utils.SequenceFeatureExtractor`] which contains most of the main @@ -99,6 +98,7 @@ class FireRedASR2FeatureExtractor(SequenceFeatureExtractor): Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces the high log_mel_fbank values for signals with hard-zero sections, when VAD cutoff is present in the signal. + """ model_input_names = ["input_features"] @@ -230,9 +230,8 @@ def padding_position_is_0(padded_input, input_lengths): class FireRedASR2Processor(ProcessorMixin): - r""" - Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor and - a FireRedASR2 tokenizer into a single processor. + r"""Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor + and a FireRedASR2 tokenizer into a single processor. [`FireRedASR2Processor`] offers all the functionalities of [`FireRedASR2FeatureExtractor`] and [`Qwen2Tokenizer`]. See the @@ -246,6 +245,7 @@ class FireRedASR2Processor(ProcessorMixin): tokenizer (`Qwen2Tokenizer`): An instance of [`Qwen2Tokenizer`]. The tokenizer is a required input. + """ feature_extractor_class = "FireRedASR2FeatureExtractor" @@ -271,8 +271,7 @@ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True): ) def __call__(self, *args, **kwargs): - """ - Forwards the `audio` argument to FireRedASR2FeatureExtractor's + """Forwards the `audio` argument to FireRedASR2FeatureExtractor's [`~FireRedASR2FeatureExtractor.__call__`] and the `text` argument to [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the above two methods for more information. diff --git a/vllm/transformers_utils/processors/funasr.py b/vllm/transformers_utils/processors/funasr.py index d7a3c4060ceb..aaa0bec48c51 100644 --- a/vllm/transformers_utils/processors/funasr.py +++ b/vllm/transformers_utils/processors/funasr.py @@ -21,10 +21,7 @@ def apply_cmvn(inputs, cmvn): # noqa - """ - Apply CMVN with mvn data - """ - + """Apply CMVN with mvn data.""" device = inputs.device # dtype = inputs.dtype frame, dim = inputs.shape @@ -219,8 +216,7 @@ def forward_lfr_cmvn( class FunASRFeatureExtractor(SequenceFeatureExtractor): - r""" - Constructs a FunASR feature extractor. + r"""Constructs a FunASR feature extractor. This feature extractor inherits from [`~feature_extraction_sequence_ utils.SequenceFeatureExtractor`] which contains most of the main @@ -255,6 +251,7 @@ class FunASRFeatureExtractor(SequenceFeatureExtractor): Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces the high log_mel_fbank values for signals with hard-zero sections, when VAD cutoff is present in the signal. + """ model_input_names = ["input_features"] @@ -369,8 +366,7 @@ def __call__( class FunASRProcessor(ProcessorMixin): - r""" - Constructs a FunASR processor which wraps a FunASR feature extractor and + r"""Constructs a FunASR processor which wraps a FunASR feature extractor and a FunASR tokenizer into a single processor. [`FunASRProcessor`] offers all the functionalities of @@ -385,6 +381,7 @@ class FunASRProcessor(ProcessorMixin): tokenizer (`Qwen2Tokenizer`): An instance of [`Qwen2Tokenizer`]. The tokenizer is a required input. + """ feature_extractor_class = "FunASRFeatureExtractor" @@ -410,8 +407,7 @@ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True): ) def __call__(self, *args, **kwargs): - """ - Forwards the `audio` argument to FunASRFeatureExtractor's + """Forwards the `audio` argument to FunASRFeatureExtractor's [`~FunASRFeatureExtractor.__call__`] and the `text` argument to [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the above two methods for more information. diff --git a/vllm/transformers_utils/processors/glm4v.py b/vllm/transformers_utils/processors/glm4v.py index a8da23955259..e9ee497cea68 100644 --- a/vllm/transformers_utils/processors/glm4v.py +++ b/vllm/transformers_utils/processors/glm4v.py @@ -10,8 +10,7 @@ class GLM4VImageProcessorFast(BaseImageProcessorFast): - """ - Port of https://huggingface.co/zai-org/glm-4v-9b/blob/main/tokenization_chatglm.py#L177 + """Port of https://huggingface.co/zai-org/glm-4v-9b/blob/main/tokenization_chatglm.py#L177 to HF Transformers. """ diff --git a/vllm/transformers_utils/processors/glm5next.py b/vllm/transformers_utils/processors/glm5next.py index b6988d3b8fed..6ad174882392 100644 --- a/vllm/transformers_utils/processors/glm5next.py +++ b/vllm/transformers_utils/processors/glm5next.py @@ -261,7 +261,8 @@ def _resize_or_pad( allow_upscale: bool = False, ) -> torch.Tensor: """Resize onto the aligned canvas, or keep the aspect ratio and - zero-pad the right/bottom sides (``resize_mode="pad"``).""" + zero-pad the right/bottom sides (``resize_mode="pad"``). + """ height, width = stacked_images.shape[-2:] if resize_mode == "resize": @@ -307,7 +308,8 @@ def _pixel_budget( ) -> tuple[int, int]: """(min_pixels, max_pixels) from the token bounds of ``processor_config.json``; one vision token covers - ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels.""" + ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels. + """ if min_image_tokens is None or max_image_tokens is None: raise ValueError( "min_image_tokens and max_image_tokens must be provided by " diff --git a/vllm/transformers_utils/processors/internvl.py b/vllm/transformers_utils/processors/internvl.py index 92c1834c74ee..84976cb178ce 100644 --- a/vllm/transformers_utils/processors/internvl.py +++ b/vllm/transformers_utils/processors/internvl.py @@ -358,8 +358,7 @@ def __call__( class InternVLProcessor(ProcessorMixin): - """ - This model doesn't define its own HF processor, + """This model doesn't define its own HF processor, so we implement our own one here. The code to insert image tokens is based on: diff --git a/vllm/transformers_utils/processors/isaac.py b/vllm/transformers_utils/processors/isaac.py index 2d791df01673..7d10e9be86e9 100644 --- a/vllm/transformers_utils/processors/isaac.py +++ b/vllm/transformers_utils/processors/isaac.py @@ -91,6 +91,7 @@ def get_image_size_for_max_num_patches( `tuple[int, int]`: Height and width (in pixels) that are multiples of `patch_size * pixel_shuffle_scale` and respect both the maximum and optional minimum patch-count constraints. + """ def get_scaled_image_size(scale, original_size, patch_size, pixel_shuffle_scale): @@ -177,9 +178,11 @@ def prepare_image_tensor( The tensor is converted to floating point if needed. scale (`float`, *optional*, defaults to `VISION_SCALE`): Scalar multiplier applied before normalization. + Returns: `torch.Tensor`: Normalized tensor with the same shape as the input and dtype `torch.float32`. + """ if not torch.is_floating_point(image): image = image.float() @@ -209,6 +212,7 @@ def patchify_vision(image: torch.Tensor, patch_size: int) -> torch.Tensor: Raises: ValueError: If `height` or `width` is not divisible by `patch_size`. + """ num_images, height, width, channels = image.shape if height % patch_size or width % patch_size: @@ -265,6 +269,7 @@ def process_vision_for_patches( / patch_size, channels * patch_size**2)` and `dims_virtual` encodes effective `(images, height, width)` dimensions after optional pixel shuffling. + """ # Add batch dim if single image if images.dim() == 3: diff --git a/vllm/transformers_utils/processors/kimi_k25.py b/vllm/transformers_utils/processors/kimi_k25.py index edee9734ce42..13aed85de4ec 100644 --- a/vllm/transformers_utils/processors/kimi_k25.py +++ b/vllm/transformers_utils/processors/kimi_k25.py @@ -28,8 +28,7 @@ def __call__( return_tensors: str | TensorType | None = None, **kwargs, ) -> BatchFeature: - """ - Args: + """Args: text: The text to be field to the model. vision_chunks: List of `VisionChunk` items to be processed. For image: `VisionChunkImage` with @@ -44,6 +43,7 @@ def __call__( Returned when `vision_chunks` is not `None`. - **grid_thws** -- list of image 3D grid in LLM. Returned when `vision_chunks` is not `None`. + """ if vision_chunks is not None: mm_inputs = self.image_processor.preprocess( diff --git a/vllm/transformers_utils/processors/mimo_v2_omni.py b/vllm/transformers_utils/processors/mimo_v2_omni.py index 22349f751364..0555c941b2b4 100644 --- a/vllm/transformers_utils/processors/mimo_v2_omni.py +++ b/vllm/transformers_utils/processors/mimo_v2_omni.py @@ -1071,6 +1071,8 @@ def __call__( tuples (pre-decoded). audio: list of ``str`` (path/url/base64), ``bytes``, ``(waveform_1D, sample_rate)`` tuples, or ``np.ndarray``. + video_audio: audio track extracted from ``videos``, in the same + formats accepted by ``audio``. return_tensors: Passed to :class:`BatchFeature`. Returns: @@ -1079,6 +1081,7 @@ def __call__( - ``pixel_values`` + ``image_grid_thw`` - ``pixel_values_videos`` + ``video_grid_thw`` + ``second_per_grid_ts`` - ``audio_features`` + """ if isinstance(text, list): text = text[0] if len(text) == 1 else "\n".join(text) diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py index d9ac841cd86b..e59a50ccbe8a 100644 --- a/vllm/transformers_utils/processors/minicpmo.py +++ b/vllm/transformers_utils/processors/minicpmo.py @@ -14,9 +14,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -""" -Processor class for MiniCPMO. -""" +"""Processor class for MiniCPMO.""" import math from typing import Literal, TypeAlias @@ -35,8 +33,7 @@ class MiniCPMOProcessor(ProcessorMixin): - r""" - Constructs a MiniCPMV processor which wraps a MiniCPMV image + r"""Constructs a MiniCPMV processor which wraps a MiniCPMV image processor and a MiniCPMV tokenizer into a single processor. [`MiniCPMVProcessor`] offers all the functionalities of @@ -49,6 +46,7 @@ class MiniCPMOProcessor(ProcessorMixin): The image processor is a required input. tokenizer ([`LlamaTokenizerWrapper`], *optional*): The tokenizer is a required input. + """ attributes = ["image_processor", "feature_extractor", "tokenizer"] @@ -253,8 +251,7 @@ def audio_feature_extract( # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode # with CLIP->Llama def batch_decode(self, *args, **kwargs): - """ - This method forwards all its arguments to LlamaTokenizerFast's + """This method forwards all its arguments to LlamaTokenizerFast's [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ @@ -274,8 +271,7 @@ def batch_decode(self, *args, **kwargs): # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode # with CLIP->Llama def decode(self, *args, **kwargs): - """ - This method forwards all its arguments to LlamaTokenizerFast's + """This method forwards all its arguments to LlamaTokenizerFast's [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ @@ -551,15 +547,11 @@ def __init__( ) def __call__(self, audio: torch.Tensor) -> torch.Tensor: - """ - audio: Tensor([num_channels, num_samples]) - """ + """audio: Tensor([num_channels, num_samples]).""" return super().__call__(audio) def forward(self, audio: torch.Tensor) -> torch.Tensor: - """ - audio: Tensor([num_channels, num_samples]) - """ + """audio: Tensor([num_channels, num_samples]).""" mel: torch.Tensor = self.mel_spec(audio) features = torch.log(torch.clip(mel, min=1e-5)) return features diff --git a/vllm/transformers_utils/processors/minicpmv.py b/vllm/transformers_utils/processors/minicpmv.py index 712cbbb47a01..aaedaa376780 100644 --- a/vllm/transformers_utils/processors/minicpmv.py +++ b/vllm/transformers_utils/processors/minicpmv.py @@ -13,9 +13,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -""" -Processor class for MiniCPMV. -""" +"""Processor class for MiniCPMV.""" from typing import TypeAlias @@ -36,8 +34,7 @@ class MiniCPMVProcessor(ProcessorMixin): - r""" - Constructs a MiniCPMV processor which wraps a MiniCPMV image + r"""Constructs a MiniCPMV processor which wraps a MiniCPMV image processor and a MiniCPMV tokenizer into a single processor. [`MiniCPMVProcessor`] offers all the functionalities of @@ -50,6 +47,7 @@ class MiniCPMVProcessor(ProcessorMixin): The image processor is a required input. tokenizer ([`LlamaTokenizerWrapper`], *optional*): The tokenizer is a required input. + """ attributes = ["image_processor", "tokenizer"] @@ -93,8 +91,7 @@ def __call__( # Copied from transformers.models.clip.processing_clip.CLIPProcessor # .batch_decode with CLIP->Llama def batch_decode(self, *args, **kwargs): - """ - This method forwards all its arguments to LlamaTokenizerFast's + """This method forwards all its arguments to LlamaTokenizerFast's [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ @@ -126,8 +123,7 @@ def batch_decode(self, *args, **kwargs): # Copied from transformers.models.clip.processing_clip.CLIPProcessor # .decode with CLIP->Llama def decode(self, *args, **kwargs): - """ - This method forwards all its arguments to LlamaTokenizerFast's + """This method forwards all its arguments to LlamaTokenizerFast's [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ diff --git a/vllm/transformers_utils/processors/moondream3.py b/vllm/transformers_utils/processors/moondream3.py index ae6833b20079..22bdeef21abe 100644 --- a/vllm/transformers_utils/processors/moondream3.py +++ b/vllm/transformers_utils/processors/moondream3.py @@ -67,8 +67,7 @@ def select_tiling( class Moondream3Processor(ProcessorMixin): - """ - Constructs a Moondream3 processor which handles image preprocessing + """Constructs a Moondream3 processor which handles image preprocessing and tokenization for the Moondream3 multimodal model. Args: @@ -78,6 +77,7 @@ class Moondream3Processor(ProcessorMixin): max_crops: Maximum number of crops per image. overlap_margin: Margin for overlapping crops in patches. patch_size: Size of each patch. + """ attributes = ["tokenizer"] @@ -189,8 +189,7 @@ def from_pretrained( pretrained_model_name_or_path, **kwargs, ): - """ - Load the processor, using a separate tokenizer repo. + """Load the processor, using a separate tokenizer repo. The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1' instead of having tokenizer files in the model repo. @@ -294,8 +293,7 @@ def __call__( | list[PreTokenizedInput] = None, **kwargs: Unpack[Moondream3ProcessorKwargs], ) -> BatchFeature: - """ - Process images and text for Moondream3 model. + """Process images and text for Moondream3 model. Args: images: Input images (PIL Image, numpy array, or list thereof). @@ -304,6 +302,7 @@ def __call__( Returns: BatchFeature with processed inputs. + """ output_kwargs = self._merge_kwargs( Moondream3ProcessorKwargs, @@ -428,8 +427,7 @@ def preprocess_image( convert_to_rgb: bool = True, return_tensors: str = "pt", ) -> tuple[torch.Tensor, tuple[int, int]]: - """ - Preprocess an image using overlap-and-resize cropping strategy. + """Preprocess an image using overlap-and-resize cropping strategy. Args: image: Input PIL image, numpy array, or torch tensor. @@ -442,6 +440,7 @@ def preprocess_image( Returns: Tuple of (pixel_values tensor, tiling tuple). + """ image = self._to_pil_image(image) if convert_to_rgb: diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index 3cddbec6c866..3d4d432de2b8 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -292,8 +292,7 @@ def width_and_height_for_max_num_tokens_available( self, target_num_tokens_post_shuffle: int, ) -> tuple[int, int]: - """ - TODO: optimize this so it squeezes closer to target number of tokens. + """TODO: optimize this so it squeezes closer to target number of tokens. Calculate image dimensions that produce approximately `target` tokens after pixel_shuffle. @@ -318,6 +317,7 @@ def width_and_height_for_max_num_tokens_available( ... height // PATCH_SIZE ... ) // 2**2 == 8100 # tokens post-shuffle >>> assert tiler._get_num_embeddings(width=width, height=height) == 8100 + """ side_pixels = ( math.isqrt(target_num_tokens_post_shuffle) @@ -386,6 +386,7 @@ def process_media( num_tokens_available: Number of tokens available for this media Returns: DynamicResolutionParams for the media + """ current_num_tokens_available = num_tokens_available assert isinstance(media, Image.Image), ( @@ -473,6 +474,7 @@ def compute_params( num_tokens_available: Total number of tokens available across all media Returns: List of ImageTilingParams for each media item + """ num_tokens_available = ( num_tokens_available @@ -571,8 +573,7 @@ def rearrange_img(x): class BaseNanoNemotronVLProcessor(ABC): - """ - This model doesn't define its own HF processor, + """This model doesn't define its own HF processor, so we implement our own one here. The code to insert image tokens is based on: @@ -761,8 +762,7 @@ def __call__( class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): - """ - HF Processor with extended video processing logic. + """HF Processor with extended video processing logic. Code for video processing is adapted from video example: https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers """ @@ -1125,8 +1125,7 @@ def get_video_repl( img_context_token_ids: list[int], video_temporal_patch_size: int = 1, ) -> PromptUpdateDetails: - """ - Build prompt replacement for a video. + """Build prompt replacement for a video. The replacement returned is not actually used to replace the placeholder tokens - it's just used to make sure we allocate the correct number of tokens. diff --git a/vllm/transformers_utils/processors/nemotron_vl.py b/vllm/transformers_utils/processors/nemotron_vl.py index 9c5436aa2729..9504f5ee2ba5 100644 --- a/vllm/transformers_utils/processors/nemotron_vl.py +++ b/vllm/transformers_utils/processors/nemotron_vl.py @@ -191,8 +191,7 @@ def _images_to_pixel_values_lst( class LlamaNemotronNanoVLProcessor(InternVLProcessor): - """ - This model doesn't define its own HF processor, + """This model doesn't define its own HF processor, so we implement our own one here. The image processor is given by: @@ -287,8 +286,7 @@ def _images_to_pixel_values_lst( class LlamaNemotronVLEmbedProcessor(InternVLProcessor): - """ - Processor for LlamaNemotronVL embedding model. + """Processor for LlamaNemotronVL embedding model. Inherits from NemotronVLProcessor and specializes it for embedding tasks: - Uses SigLIP transform with normalization instead of base transform diff --git a/vllm/transformers_utils/processors/ovis.py b/vllm/transformers_utils/processors/ovis.py index 907c68035fcf..a77072745f5b 100644 --- a/vllm/transformers_utils/processors/ovis.py +++ b/vllm/transformers_utils/processors/ovis.py @@ -50,10 +50,10 @@ class OvisProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-a class OvisProcessor(ProcessorMixin): - r""" - Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor. + r"""Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor. [`OvisProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`] for more information. + Args: image_processor ([`Qwen2VLImageProcessor`], *optional*): The image processor is a required input. @@ -61,6 +61,7 @@ class OvisProcessor(ProcessorMixin): The tokenizer is a required input. chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string. + """ attributes = ["image_processor", "tokenizer"] @@ -107,12 +108,12 @@ def __call__( | list[PreTokenizedInput] = None, **kwargs: Unpack[OvisProcessorKwargs], ) -> BatchFeature: - """ - Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` + """Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`. - Args: + + Args: images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch tensor. Both channels-first and channels-last formats are supported. @@ -129,7 +130,8 @@ def __call__( - `'pt'`: Return PyTorch `torch.Tensor` objects. - `'np'`: Return NumPy `np.ndarray` objects. - `'jax'`: Return JAX `jnp.ndarray` objects. - Returns: + + Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when @@ -140,8 +142,8 @@ def __call__( - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. - """ + """ max_partition = kwargs.pop("max_partition", 9) covering_threshold = kwargs.pop("covering_threshold", 0.9) @@ -416,28 +418,28 @@ def _get_best_grid(img, side): return torch.tensor(pixel_values), image_placeholders, torch.tensor(grid) def batch_decode(self, *args, **kwargs): - """ - This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please + """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): - """ - This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to + """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.decode(*args, **kwargs) def post_process_image_text_to_text(self, generated_outputs): - """ - Post-process the output of the model to decode the text. + """Post-process the output of the model to decode the text. + Args: generated_outputs (`torch.Tensor` or `np.ndarray`): The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` or `(sequence_length,)`. + Returns: `list[str]`: The decoded text. + """ return self.tokenizer.batch_decode( generated_outputs, diff --git a/vllm/transformers_utils/processors/ovis2_5.py b/vllm/transformers_utils/processors/ovis2_5.py index b9c2d112409c..c31ea915209b 100644 --- a/vllm/transformers_utils/processors/ovis2_5.py +++ b/vllm/transformers_utils/processors/ovis2_5.py @@ -33,13 +33,13 @@ class Ovis2_5ProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[cal class Ovis2_5Processor(ProcessorMixin): - r""" - Constructs an Ovis processor which wraps an Ovis image processor + r"""Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor. [`OvisProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`] for more information. + Args: image_processor ([`Qwen2VLImageProcessor`], *optional*): The image processor is a required input. @@ -48,6 +48,7 @@ class Ovis2_5Processor(ProcessorMixin): chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string. + """ attributes = ["image_processor", "tokenizer"] @@ -110,43 +111,39 @@ def __call__( | list[PreTokenizedInput] = None, **kwargs: Unpack[Ovis2_5ProcessorKwargs], ) -> BatchFeature: - """ - Main method to prepare for the model one or several sequences(s) + """Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`. - Args: - images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, - `list[PIL.Image.Image]`, `list[np.ndarray]`, - `list[torch.Tensor]`): - The image or batch of images to be prepared. - Each image can be a PIL image, NumPy array or PyTorch - tensor. Both channels-first and channels-last formats - are supported. - text (`str`, `list[str]`, `list[list[str]]`): - The sequence or batch of sequences to be encoded. - Each sequence can be a string or a list of strings - (pretokenized string). If the sequences are provided as - list of strings (pretokenized), you must set - `is_split_into_words=True` (to lift the ambiguity with - a batch of sequences). - videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, - `list[torch.Tensor]`): - The image or batch of videos to be prepared. Each video - can be a 4D NumPy array or PyTorch tensor, or a nested - list of 3D frames. Both channels-first and channels-last - formats are supported. - return_tensors (`str` or [`~utils.TensorType`], *optional*): - If set, will return tensors of a particular framework. - Acceptable values are: - - `'tf'`: Return TensorFlow `tf.constant` objects. - - `'pt'`: Return PyTorch `torch.Tensor` objects. - - `'np'`: Return NumPy `np.ndarray` objects. - - `'jax'`: Return JAX `jnp.ndarray` objects. - Returns: + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, or a list): + The image or batch of images to be prepared. Each image can + be a PIL image, NumPy array or PyTorch tensor. Both + channels-first and channels-last formats are supported. + videos (`np.ndarray`, `torch.Tensor`, or a list of either): + The video or batch of videos to be prepared. Each video can + be a 4D NumPy array or PyTorch tensor, or a nested list of + 3D frames. Both channels-first and channels-last formats are + supported. + text (`str`, `list[str]`, `list[list[str]]`): + The sequence or batch of sequences to be encoded. Each + sequence can be a string or a list of strings (pretokenized + string). If the sequences are provided as list of strings + (pretokenized), you must set `is_split_into_words=True` (to + lift the ambiguity with a batch of sequences). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. + Acceptable values are: + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + + Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- list of token ids to be fed to a model. Returned when `text` is not `None`. @@ -164,6 +161,7 @@ def __call__( when `videos` is not `None`. - **second_per_grid_ts** -- list of video seconds per time grid. Returned when `videos` is not `None`. + """ output_kwargs = self._merge_kwargs( Ovis2_5ProcessorKwargs, diff --git a/vllm/transformers_utils/processors/pixtral.py b/vllm/transformers_utils/processors/pixtral.py index 588ad6b7fad8..85604eb5b71d 100644 --- a/vllm/transformers_utils/processors/pixtral.py +++ b/vllm/transformers_utils/processors/pixtral.py @@ -11,8 +11,7 @@ class MistralCommonImageProcessor(ImageProcessingMixin): - """ - Provide a HF-compatible interface for + """Provide a HF-compatible interface for `mistral_common.tokens.tokenizers.multimodal.ImageEncoder`. """ diff --git a/vllm/transformers_utils/processors/voxtral.py b/vllm/transformers_utils/processors/voxtral.py index 5403a7e47f99..4cc23c47754f 100644 --- a/vllm/transformers_utils/processors/voxtral.py +++ b/vllm/transformers_utils/processors/voxtral.py @@ -18,8 +18,7 @@ class MistralCommonFeatureExtractor(SequenceFeatureExtractor): - """ - Provide a HF-compatible interface for + """Provide a HF-compatible interface for `mistral_common.tokens.tokenizers.multimodal.AudioEncoder`. """ diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index ca20390aad51..6b9ff9ec5eb5 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -66,6 +66,7 @@ def resolve_revision( Returns: The resolved revision, or `revision` unchanged if it cannot be resolved (local path, ModelScope, or any Hub error). + """ if Path(repo_id).exists() or envs.VLLM_USE_MODELSCOPE: return revision @@ -106,6 +107,7 @@ def with_retry( retry_delay: Seconds to wait after the first failure, doubled each time. fatal_errors: Exceptions that are a definitive answer rather than a transient failure. These are raised immediately, without logging. + """ for attempt in range(max_retries): try: @@ -353,8 +355,7 @@ def get_hf_file_bytes( def try_get_local_file( model: str | Path, file_name: str, revision: str | None = "main" ) -> Path | Any | None: - """ - Try to get a local file from the HuggingFace repository. + """Try to get a local file from the HuggingFace repository. The possible return values are: @@ -385,20 +386,19 @@ def try_get_local_file( def get_hf_file_to_dict( file_name: str, model: str | Path, revision: str | None = "main" ): - """ - Downloads a file from the Hugging Face Hub and returns + """Downloads a file from the Hugging Face Hub and returns its contents as a dictionary. - Parameters: - - file_name (str): The name of the file to download. - - model (str): The name of the model on the Hugging Face Hub. - - revision (str): The specific version of the model. + Args: + file_name (str): The name of the file to download. + model (str): The name of the model on the Hugging Face Hub. + revision (str): The specific version of the model. Returns: - config_dict (dict): A dictionary containing the contents of the downloaded file. - """ + """ file_path = try_get_local_file(model=model, file_name=file_name, revision=revision) if file_path is None: diff --git a/vllm/transformers_utils/runai_utils.py b/vllm/transformers_utils/runai_utils.py index 248ede6a6f1d..d02992b10b1c 100644 --- a/vllm/transformers_utils/runai_utils.py +++ b/vllm/transformers_utils/runai_utils.py @@ -25,14 +25,14 @@ def list_safetensors(path: str = "") -> list[str]: - """ - List full file names from object path and filter by allow pattern. + """List full file names from object path and filter by allow pattern. Args: path: The object storage path to list from. Returns: list[str]: List of full object storage paths allowed by the pattern + """ return runai_list_safetensors(path) @@ -42,8 +42,7 @@ def is_runai_obj_uri(model_or_path: str) -> bool: class ObjectStorageModel: - """ - A class representing an ObjectStorage model mirrored into a + """A class representing an ObjectStorage model mirrored into a temporary directory. Attributes: @@ -51,6 +50,7 @@ class ObjectStorageModel: Methods: pull_files(): Pull model from object storage to the temporary directory. + """ def __init__(self, url: str) -> None: @@ -86,8 +86,7 @@ def pull_files( allow_pattern: list[str] | None = None, ignore_pattern: list[str] | None = None, ) -> None: - """ - Pull files from object storage into the temporary directory. + """Pull files from object storage into the temporary directory. Args: model_path: The object storage path of the model. diff --git a/vllm/transformers_utils/s3_utils.py b/vllm/transformers_utils/s3_utils.py index 7dfee9fabe36..d0f403c6908d 100644 --- a/vllm/transformers_utils/s3_utils.py +++ b/vllm/transformers_utils/s3_utils.py @@ -36,8 +36,7 @@ def glob( path: str = "", allow_pattern: list[str] | None = None, ) -> list[str]: - """ - List full file names from S3 path and filter by allow pattern. + """List full file names from S3 path and filter by allow pattern. Args: s3: S3 client to use. @@ -46,6 +45,7 @@ def glob( Returns: list[str]: List of full S3 paths allowed by the pattern + """ if s3 is None: s3 = boto3.client("s3") @@ -61,8 +61,7 @@ def list_files( allow_pattern: list[str] | None = None, ignore_pattern: list[str] | None = None, ) -> tuple[str, str, list[str]]: - """ - List files from S3 path and filter by pattern. + """List files from S3 path and filter by pattern. Args: s3: S3 client to use. @@ -77,6 +76,7 @@ def list_files( and the prefix as a dir like string - The third element is a list of files allowed or disallowed by pattern + """ parts = path.removeprefix("s3://").split("/") prefix = "/".join(parts[1:]) diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py index 674f6a05bf98..c5dc1a8ff3cc 100644 --- a/vllm/transformers_utils/utils.py +++ b/vllm/transformers_utils/utils.py @@ -80,16 +80,15 @@ def _maybe_space_split_dict(path: str | PathLike) -> dict[str, str]: @cache def maybe_model_redirect(model: str) -> str: - """ - Use model_redirect to redirect the model name to a local folder. + """Use model_redirect to redirect the model name to a local folder. Args: model: hf model name Returns: maybe redirect to a local folder - """ + """ model_redirect_path = envs.VLLM_MODEL_REDIRECT_PATH if not model_redirect_path: @@ -117,7 +116,8 @@ def parse_safetensors_file_metadata(path: str | PathLike) -> dict[str, Any]: def convert_model_repo_to_path(model_repo: str) -> str: """When VLLM_USE_MODELSCOPE is True convert a model - repository string to a Path str.""" + repository string to a Path str. + """ if not envs.VLLM_USE_MODELSCOPE or Path(model_repo).exists(): return model_repo from modelscope.utils.file_utils import get_model_cache_root diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index fc906cf22727..e573cbc293b3 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -34,6 +34,7 @@ def human_readable_int(value: str) -> int: - '1k' -> 1,000 - '1K' -> 1,024 - '25.6k' -> 25,600 + """ value = value.strip() @@ -82,6 +83,7 @@ def human_readable_int_or_auto(value: str) -> int: - '1K' -> 1,024 - '25.6k' -> 25,600 - '-1' or 'auto' -> -1 (special value for auto-detection) + """ value = value.strip() @@ -95,8 +97,7 @@ class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpForma """SortedHelpFormatter that sorts arguments by their option strings.""" def _split_lines(self, text, width): - """ - 1. Sentences split across lines have their single newlines removed. + """1. Sentences split across lines have their single newlines removed. 2. Paragraphs and lists are split into separate lines. 3. Each line is wrapped to the specified width (width of terminal). """ @@ -466,7 +467,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]: The arguments in config file will be inserted between the argument list. - example: + Example: ```yaml port: 12323 tensor-parallel-size: 4 @@ -492,6 +493,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]: Please note how the config args are inserted after the sub command. this way the order of priorities is maintained when these are args parsed by super(). + """ assert args.count("--config") <= 1, "More than one config file specified!" diff --git a/vllm/utils/async_utils.py b/vllm/utils/async_utils.py index 3faf728384e6..b2c0d57dc921 100644 --- a/vllm/utils/async_utils.py +++ b/vllm/utils/async_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Contains helpers related to asynchronous code. +"""Contains helpers related to asynchronous code. This is similar in concept to the `asyncio` module. """ @@ -29,8 +28,7 @@ def make_async( func: Callable[P, T], executor: Executor | None = None, ) -> Callable[P, Awaitable[T]]: - """ - Take a blocking function, and run it on in an executor thread. + """Take a blocking function, and run it on in an executor thread. This function prevents the blocking function from blocking the asyncio event loop. @@ -49,8 +47,7 @@ def make_async_with_semaphore( func: Callable[P, T], executor: ThreadPoolExecutor, ) -> Callable[P, Awaitable[T]]: - """ - Take a blocking function, and run it on in an executor thread. + """Take a blocking function, and run it on in an executor thread. This function prevents the blocking function from blocking the asyncio event loop. @@ -59,7 +56,6 @@ def make_async_with_semaphore( The function is wrapped in a semaphore to limit the number of concurrent executions making it easier to cancel tasks before they start. """ - semaphore = asyncio.Semaphore(executor._max_workers) async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py index 791b8d848d4f..9fca5ec46541 100644 --- a/vllm/utils/cache.py +++ b/vllm/utils/cache.py @@ -102,8 +102,7 @@ def usage(self) -> float: return self.currsize / self.maxsize def stat(self, *, delta: bool = False) -> CacheInfo: - """ - Gets the cumulative number of hits and queries against this cache. + """Gets the cumulative number of hits and queries against this cache. If `delta=True`, instead gets these statistics since the last call that also passed `delta=True`. @@ -165,6 +164,7 @@ def put_if_fits(self, key: _K, value: _V) -> bool: Returns: `True` if the item was cached, otherwise `False`. + """ try: self[key] = value @@ -175,8 +175,7 @@ def put_if_fits(self, key: _K, value: _V) -> bool: return True def pin(self, key: _K) -> None: - """ - Pins a key in the cache preventing it from being + """Pins a key in the cache preventing it from being evicted in the LRU order. """ if key not in self: @@ -184,8 +183,7 @@ def pin(self, key: _K) -> None: self.pinned_items.add(key) def _unpin(self, key: _K) -> None: - """ - Unpins a key in the cache allowing it to be + """Unpins a key in the cache allowing it to be evicted in the LRU order. """ self.pinned_items.remove(key) diff --git a/vllm/utils/collection_utils.py b/vllm/utils/collection_utils.py index 5852aff549c4..de32ed0c15c2 100644 --- a/vllm/utils/collection_utils.py +++ b/vllm/utils/collection_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Contains helpers that are applied to collections. +"""Contains helpers that are applied to collections. This is similar in concept to the `collections` module. """ @@ -20,8 +19,7 @@ class LazyDict(Mapping[str, _V], Generic[_V]): - """ - Evaluates dictionary items only when they are accessed. + """Evaluates dictionary items only when they are accessed. Adapted from: https://stackoverflow.com/a/47212782/5082708 """ @@ -116,8 +114,7 @@ def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]: def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]): - """ - Unlike [`itertools.groupby`][], groups are not broken by + """Unlike [`itertools.groupby`][], groups are not broken by non-contiguous data. """ groups = defaultdict[_K, list[_V]](list) diff --git a/vllm/utils/counter.py b/vllm/utils/counter.py index c2dce32e97e1..6251cdcfb1b3 100644 --- a/vllm/utils/counter.py +++ b/vllm/utils/counter.py @@ -19,10 +19,10 @@ def reset(self) -> None: class AtomicCounter: - """An atomic, thread-safe counter""" + """An atomic, thread-safe counter.""" def __init__(self, initial: int = 0) -> None: - """Initialize a new atomic counter to given initial value""" + """Initialize a new atomic counter to given initial value.""" super().__init__() self._value = initial @@ -33,13 +33,13 @@ def value(self) -> int: return self._value def inc(self, num: int = 1) -> int: - """Atomically increment the counter by num and return the new value""" + """Atomically increment the counter by num and return the new value.""" with self._lock: self._value += num return self._value def dec(self, num: int = 1) -> int: - """Atomically decrement the counter by num and return the new value""" + """Atomically decrement the counter by num and return the new value.""" with self._lock: self._value -= num return self._value diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index a1e7ab0ee30d..89d3e6ca7c73 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -128,6 +128,7 @@ def check_cgroup_memory_available( Low headroom logs a warning, but does not reject the allocation because cgroup usage can include reclaimable memory. If the cgroup limit or usage cannot be read, the check is skipped. + """ cgroup_limit = get_cgroup_memory_limit() cgroup_usage = get_cgroup_memory_usage() @@ -171,7 +172,7 @@ def get_memory_affinity(pid: int = 0) -> list[int]: def parse_id_list(raw_str: str) -> list[int]: - """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]""" + """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8].""" result: list[int] = [] if not raw_str: return result @@ -273,7 +274,8 @@ def get_visible_memory_node() -> list[int]: def _synthesize_cpu_list() -> list[LogicalCPUInfo]: """Synthesize a flat CPU list: each logical CPU is its own core on NUMA node 0. Used when lscpu output is unavailable or unparsable - (e.g. macOS, RISC-V).""" + (e.g. macOS, RISC-V). + """ cpu_count = os.cpu_count() assert cpu_count return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)] diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index 9ce33b835707..d1696317aab8 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Contains replacement functions to fallback Triton usages in CPU backend -""" +"""Contains replacement functions to fallback Triton usages in CPU backend.""" import ctypes from collections.abc import Callable diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 6fef0072243f..009c7b42582c 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -64,7 +64,7 @@ class DeepGemmQuantScaleFMT(Enum): @classmethod def init_oracle_cache(cls) -> None: - """Initialize the oracle decision and store it in the class cache""" + """Initialize the oracle decision and store it in the class cache.""" cached = getattr(cls, "_oracle_cache", None) if cached is not None: return @@ -426,7 +426,7 @@ def mk_alignment_scope(value: int): def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor: - """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor""" + """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor.""" _lazy_init() if _get_mn_major_tma_aligned_tensor_impl is None: return _missing() @@ -568,6 +568,7 @@ def fp8_fp4_mqa_logits( Returns: Logits tensor of shape [M, N], dtype `torch.float32`. + """ _lazy_init() if _fp8_fp4_mqa_logits_impl is None: @@ -623,6 +624,7 @@ def get_paged_mqa_logits_metadata( Returns: Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits` to schedule work across SMs. + """ _lazy_init() if _get_paged_mqa_logits_metadata_impl is None: @@ -674,6 +676,7 @@ def fp8_fp4_paged_mqa_logits( Returns: Logits tensor of shape [B * next_n, max_model_len], dtype `torch.float32`. + """ _lazy_init() if _fp8_fp4_paged_mqa_logits_impl is None: @@ -701,7 +704,8 @@ def fp8_fp4_paged_mqa_logits( def has_deep_gemm_sparse_mqa() -> bool: """Whether the installed DeepGEMM provides the sparse-indexer kernels - (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only). + """ _lazy_init() return ( _fp8_fp4_sparse_mqa_logits_impl is not None @@ -740,6 +744,7 @@ def get_sparse_mqa_logits_metadata( qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, ``torch.float8_e4m3fn`` for FP8). sparse_block_kv: Tokens per sparse block, 8 or 16. + """ _lazy_init() if _get_sparse_mqa_logits_metadata_impl is None: @@ -777,6 +782,7 @@ def get_paged_sparse_mqa_logits_metadata( block ids, [num_q_tokens, num_max_sparse_blocks], int32. qk_dtype: dtype of the packed Q values. sparse_block_kv: Tokens per sparse block, 8 or 16. + """ _lazy_init() if _get_paged_sparse_mqa_logits_metadata_impl is None: @@ -817,6 +823,7 @@ def fp8_fp4_sparse_mqa_logits( bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; column ``j * sparse_block_kv + o`` scores the token at ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. + """ _lazy_init() if _fp8_fp4_sparse_mqa_logits_impl is None: @@ -848,10 +855,13 @@ def fp8_fp4_paged_sparse_mqa_logits( page stride 512B-aligned. weights: [num_q_tokens, H] ``torch.bfloat16``. metadata: From `get_paged_sparse_mqa_logits_metadata`. + num_max_sparse_blocks: Candidate blocks per row. + sparse_block_kv: Tokens per sparse block, 8 or 16. Returns: bf16 logits of shape [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. + """ _lazy_init() if _fp8_fp4_paged_sparse_mqa_logits_impl is None: @@ -873,8 +883,7 @@ def tf32_hc_prenorm_gemm( sqrsum: torch.Tensor, num_split: int, ) -> torch.Tensor: - """ - Perform the following computation: + """Perform the following computation: out = x.float() @ fn.T sqrsum = x.float().square().sum(-1) @@ -950,7 +959,6 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): and report `1 - sim`. Once kernel accuracy improves this helper can be removed. """ - x, y = x.double(), y.double() denominator = (x * x + y * y).sum() sim = 2 * (x * y).sum() / denominator diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 7c76dd7ba2ff..76471791349d 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -589,8 +589,7 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool: def force_use_trtllm_attention() -> bool | None: - """ - This function should only be called during initialization stage when vllm config + """This function should only be called during initialization stage when vllm config is set. Return `None` if --attention-config.use_trtllm_attention is not set, return `True` if TRTLLM attention is forced to be used, @@ -628,7 +627,6 @@ def use_trtllm_attention( has_spec: bool = False, ) -> bool: """Return `True` if TRTLLM attention is used.""" - # CLI argument is set to 0 - respect it if force_use_trtllm is not None and not force_use_trtllm: return False @@ -736,6 +734,7 @@ def _flashinfer_concat_mla_k( k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim]. k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim]. This is broadcast to all heads. + """ from flashinfer.concat_ops import concat_mla_k diff --git a/vllm/utils/func_utils.py b/vllm/utils/func_utils.py index 5ce23e6a0074..69600decda49 100644 --- a/vllm/utils/func_utils.py +++ b/vllm/utils/func_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Contains helpers that are applied to functions. +"""Contains helpers that are applied to functions. This is similar in concept to the `functools` module. """ @@ -127,8 +126,7 @@ def get_allowed_kwarg_only_overrides( requires_kw_only: bool = True, allow_var_kwargs: bool = False, ) -> dict[str, Any]: - """ - Given a callable which has one or more keyword only params and a dict + """Given a callable which has one or more keyword only params and a dict mapping param names to values, drop values that can be not be kwarg expanded to overwrite one or more keyword-only args. This is used in a few places to handle custom processor overrides for multimodal models, @@ -139,12 +137,14 @@ def get_allowed_kwarg_only_overrides( callable: Callable which takes 0 or more keyword only arguments. If None is provided, all overrides names are allowed. overrides: Potential overrides to be used when invoking the callable. + requires_kw_only: If True, only keyword-only params may be overridden. allow_var_kwargs: Allows overrides that are expandable for var kwargs. Returns: Dictionary containing the kwargs to be leveraged which may be used to overwrite one or more keyword only arguments when invoking the callable. + """ if not overrides: return {} diff --git a/vllm/utils/gc_utils.py b/vllm/utils/gc_utils.py index 32bdfb764861..be6ab4a23fa1 100644 --- a/vllm/utils/gc_utils.py +++ b/vllm/utils/gc_utils.py @@ -14,8 +14,7 @@ class GCDebugConfig: - """ - Config for GC Debugger. + """Config for GC Debugger. - 0: disable GC debugger - 1: enable GC debugger with gc.collect elapsed times - '{"top_objects":5}': enable GC debugger with top 5 collected objects @@ -44,8 +43,7 @@ def __repr__(self) -> str: class GCDebugger: - """ - Debugger for GC which logs helpful information for GC understanding. + """Debugger for GC which logs helpful information for GC understanding. To enable, you should call maybe_attach_gc_debug_callback in the process. """ @@ -59,9 +57,7 @@ def __init__(self, config: GCDebugConfig) -> None: self.gc_top_collected_objects: str = "" def handle(self, phase: str, info: dict[str, int]) -> None: - """ - Handles a GC event (e.g. GC start or GC finish) - """ + """Handles a GC event (e.g. GC start or GC finish).""" generation = info.get("generation") if generation is None: return @@ -122,8 +118,7 @@ def freeze_gc_for_cudagraph_capture(): def freeze_gc_heap() -> None: - """ - Freeze all objects tracked by the garbage collector. It should be invoked + """Freeze all objects tracked by the garbage collector. It should be invoked after server init / warmup, to reduce GC overhead from static objects during serving time. """ @@ -137,9 +132,7 @@ def freeze_gc_heap() -> None: def maybe_attach_gc_debug_callback() -> None: - """ - Attached a callback for GC debug when VLLM_GC_DEBUG is enabled. - """ + """Attached a callback for GC debug when VLLM_GC_DEBUG is enabled.""" config = GCDebugConfig(envs.VLLM_GC_DEBUG) if config.enabled: debugger: GCDebugger = GCDebugger(config) @@ -151,8 +144,7 @@ def gc_callback(phase: str, info: dict[str, int]) -> None: def _compute_detailed_type(o: Any) -> str: - """ - Detailed object type. + """Detailed object type. TODO(Jialin): Further enhance the detailed type with element types for easier debugging. We tried but occasionally it would run into signals @@ -167,9 +159,7 @@ def _compute_detailed_type(o: Any) -> str: def _compute_top_gc_collected_objects(objects: list[Any], top: int) -> str: - """ - Group collected objects by types. - """ + """Group collected objects by types.""" if top <= 0: return "" object_types = [_compute_detailed_type(o) for o in objects] diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 8d0e7ae1871f..096f0d005bee 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -72,7 +72,8 @@ def enable_gpu_sync_check() -> None: def _sync_warning_hook(message, category, filename, lineno, file=None, line=None): """Turn torch's sync warning into an error, but only on a thread that is - being checked and is outside any allow region.""" + being checked and is outside any allow region. + """ if _TORCH_SYNC_WARNING in str(message): mode = _checking.get() if mode is None or _allow_depth.get(): @@ -251,7 +252,8 @@ def _is_dense(t: torch.Tensor) -> bool: """Whether `t`'s elements fill a contiguous storage range, possibly permuted (e.g. a transposed matrix). Such layouts can be copied with pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas - gapped layouts (e.g. strided slices) stage through a pageable temp.""" + gapped layouts (e.g. strided slices) stage through a pageable temp. + """ if t.is_contiguous(): # The common case, answered by a memoized TensorImpl flag. return True diff --git a/vllm/utils/hashing.py b/vllm/utils/hashing.py index f01c6b074ffe..9caee37493ca 100644 --- a/vllm/utils/hashing.py +++ b/vllm/utils/hashing.py @@ -35,6 +35,7 @@ def sha256(input: Any) -> bytes: Returns: Bytes representing the SHA-256 hash of the serialized input. + """ input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL) return hashlib.sha256(input_bytes).digest() @@ -53,6 +54,7 @@ def sha256_cbor(input: Any) -> bytes: Returns: Bytes representing the SHA-256 hash of the CBOR serialized input. + """ input_bytes = cbor2.dumps(input, canonical=True) return hashlib.sha256(input_bytes).digest() @@ -87,6 +89,7 @@ def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]: Returns: A hash function. + """ if hash_fn_name == "sha256": return sha256 @@ -110,6 +113,7 @@ def safe_hash(data: bytes, usedforsecurity: bool = True) -> HASH: Returns: Hash object + """ try: return hashlib.md5(data, usedforsecurity=usedforsecurity) diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 20c510edc025..d6a2aabbb47c 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Contains helpers related to importing modules. +"""Contains helpers related to importing modules. This is similar in concept to the `importlib` module. """ @@ -24,8 +23,7 @@ def import_pynvml(): - """ - Historical comments: + """Historical comments: libnvml.so is the library behind nvidia-smi, and pynvml is a Python wrapper around it. We use it to get GPU @@ -57,8 +55,7 @@ def import_pynvml(): @cache def import_triton_kernels(): - """ - For convenience, prioritize triton_kernels that is available in + """For convenience, prioritize triton_kernels that is available in `site-packages`. Use `vllm.third_party.triton_kernels` as a fall-back. """ if _has_module("triton_kernels"): @@ -83,8 +80,7 @@ def import_triton_kernels(): def import_plugin(plugin_path: str) -> ModuleType | None: - """ - Import a user-defined plugin. + """Import a user-defined plugin. Plugin can be either: * a module in site-packages @@ -104,8 +100,7 @@ def import_plugin(plugin_path: str) -> ModuleType | None: def import_from_path(module_name: str, file_path: str | os.PathLike) -> ModuleType: - """ - Import a Python file according to its file path. + """Import a Python file according to its file path. Based on the official recipe: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly @@ -123,9 +118,7 @@ def import_from_path(module_name: str, file_path: str | os.PathLike) -> ModuleTy def resolve_obj_by_qualname(qualname: str) -> Any: - """ - Resolve an object by its fully-qualified class name. - """ + """Resolve an object by its fully-qualified class name.""" module_name, obj_name = qualname.rsplit(".", 1) module = importlib.import_module(module_name) return getattr(module, obj_name) @@ -148,8 +141,7 @@ def get_vllm_optional_dependencies(): class _PlaceholderBase: - """ - Disallows downstream usage of placeholder modules. + """Disallows downstream usage of placeholder modules. We need to explicitly override each dunder method because [`__getattr__`][vllm.utils.import_utils._PlaceholderBase.__getattr__] @@ -160,8 +152,7 @@ class _PlaceholderBase: """ def __getattr__(self, key: str) -> Never: - """ - The main class should implement this to throw an error + """The main class should implement this to throw an error for attribute accesses representing downstream usage. """ raise NotImplementedError @@ -304,8 +295,7 @@ def __exit__(self, *args: object, **kwargs: object): class PlaceholderModule(_PlaceholderBase): - """ - A placeholder object to use when a module does not exist. + """A placeholder object to use when a module does not exist. This enables more informative errors when trying to access attributes of a module that does not exist. @@ -360,8 +350,7 @@ def __getattr__(self, key: str) -> Never: class LazyLoader(ModuleType): - """ - `LazyLoader` module borrowed from [Tensorflow] + """`LazyLoader` module borrowed from [Tensorflow] (https://github.com/tensorflow/tensorflow/blob/main/tensorflow/python/util/lazy_loader.py) with an addition of "module caching". @@ -562,7 +551,6 @@ def has_tilelang() -> bool: def has_arctic_inference() -> bool: """Whether the optional `arctic_inference` package is available.""" - return _has_module("arctic_inference") diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index bb0bb8df9a1e..ffb63bc493d4 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Monitor unexpected kernel JIT compilation during inference. +"""Monitor unexpected kernel JIT compilation during inference. After server warmup completes, any kernel JIT compilation or autotuning event indicates a cache miss or unexpected input shape that causes a latency spike. diff --git a/vllm/utils/jsontree.py b/vllm/utils/jsontree.py index fe757c2f3374..0f430935acdf 100644 --- a/vllm/utils/jsontree.py +++ b/vllm/utils/jsontree.py @@ -142,8 +142,7 @@ def json_reduce_leaves( initial: _U = ..., # type: ignore[assignment] /, ) -> _T | _U: - """ - Apply a function of two arguments cumulatively to each leaf in a + """Apply a function of two arguments cumulatively to each leaf in a nested JSON structure, from left to right, so as to reduce the sequence to a single value. """ diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py index c681e708b5c1..8fb0009d60bf 100644 --- a/vllm/utils/math_utils.py +++ b/vllm/utils/math_utils.py @@ -13,7 +13,7 @@ def cdiv(a: int, b: int) -> int: def next_power_of_2(n: int) -> int: - """The next power of 2 (inclusive)""" + """The next power of 2 (inclusive).""" return 1 if n < 1 else 1 << (n - 1).bit_length() diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py index c092dff41045..1b449058b27d 100644 --- a/vllm/utils/mem_utils.py +++ b/vllm/utils/mem_utils.py @@ -54,6 +54,7 @@ def release_device_memory_under_pressure(device: torch.device) -> bool: Returns: True if memory was released. + """ if device.type != "cuda" or not current_platform.is_integrated_gpu(device.index): return False @@ -230,8 +231,7 @@ def memory_profiling( baseline_snapshot: MemorySnapshot, weights_memory: int = 0, ) -> Generator[MemoryProfilingResult, None, None]: - """ - Memory profiling context manager. + """Memory profiling context manager. baseline_snapshot: the memory snapshot before the current vLLM instance. weights_memory: memory used by PyTorch when loading the model weights. diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index a8286c84804e..6190b7699178 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -41,6 +41,7 @@ def maybe_execute_in_parallel( Returns: Tuple of (fn0_result, fn1_result). + """ if aux_stream is not None: from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture @@ -100,6 +101,7 @@ def execute_in_parallel( Returns: Tuple of (default_result, aux_results) where aux_results[i] is the result of aux_fns[i] (or None when skipped). + """ aux_results: list[Any] if aux_streams is None or not enable: diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py index 356d913879a1..fa24d4d5f9ab 100644 --- a/vllm/utils/network_utils.py +++ b/vllm/utils/network_utils.py @@ -172,8 +172,7 @@ def _get_reserved_port_range() -> range: def get_open_port() -> int: - """ - Get an open port for the vLLM process to listen on. + """Get an open port for the vLLM process to listen on. An edge case to handle, is when we run data parallel, we need to avoid ports that are potentially used by the data parallel master process. @@ -298,6 +297,7 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: Returns: A properly formatted ZMQ path string. + """ if port is None: return f"{scheme}://{host}" @@ -317,7 +317,6 @@ def make_zmq_socket( router_handover: bool = False, ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined] """Make a ZMQ socket with the proper bind/connect semantics.""" - mem = psutil.virtual_memory() socket = ctx.socket(socket_type) @@ -377,8 +376,7 @@ def zmq_socket_ctx( identity: bytes | None = None, router_handover: bool = False, ) -> Iterator[zmq.Socket]: - """Context manager for a ZMQ socket""" - + """Context manager for a ZMQ socket.""" ctx = zmq.Context() # type: ignore[attr-defined] try: yield make_zmq_socket( diff --git a/vllm/utils/nvtx_pytorch_hooks.py b/vllm/utils/nvtx_pytorch_hooks.py index 25f76edbfaa2..5d979b04cf84 100644 --- a/vllm/utils/nvtx_pytorch_hooks.py +++ b/vllm/utils/nvtx_pytorch_hooks.py @@ -25,7 +25,7 @@ def print_tensor(tensor_obj, prefix, tensor_list=None): def process_layer_params(module_obj): - """Extract the static parameters from LLM and VLM relevant layer types""" + """Extract the static parameters from LLM and VLM relevant layer types.""" param_info = {} # Extract parameters for layers commonly used in LLMs and VLMs if isinstance(module_obj, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)): @@ -186,6 +186,7 @@ def layerwise_nvtx_marker_context(module_name, module_obj, in_tensor=None, kwarg kwargs=kwargs) as ctx: ctx.result = module(*args, **kwargs) return ctx.result + """ holder = ResultHolder() @@ -221,9 +222,9 @@ class PytHooks: a PytHook object. Then call the register_hooks method. Example: - my_hook = PytHook() my_hook.register_hooks(my_network_model) + """ def __init__(self): diff --git a/vllm/utils/ompmultiprocessing.py b/vllm/utils/ompmultiprocessing.py index 0a704e6d2676..57f0016b2260 100644 --- a/vllm/utils/ompmultiprocessing.py +++ b/vllm/utils/ompmultiprocessing.py @@ -187,18 +187,19 @@ def _parse_omp_threads_bind_env(self): def _get_autobind_cpu_ids( self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]] ) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]: - """ - Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for + """Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for other processes. Currently for rank N, only CPU ids on the N-th node in available NUMA node list will be selected. + Args: cpu_selector: a callable object to select CPUs from a CPU list - of a physical core. The input is a LogicalCPUInfo list contains - logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id. - A selected LogicalCPUInfo list should be returned. - """ + of a physical core. The input is a LogicalCPUInfo list contains + logical CPUs of a physical CPU, sorted by the + LogicalCPUInfo.id. A selected LogicalCPUInfo list should be + returned. + """ # this memory node list has been sliced for DP offset allowed_numa_nodes = cr_utils.get_visible_memory_node() logical_cpu_list = cr_utils.get_allowed_cpu_list() diff --git a/vllm/utils/platform_utils.py b/vllm/utils/platform_utils.py index 5d7fed3c9904..975ad696d986 100644 --- a/vllm/utils/platform_utils.py +++ b/vllm/utils/platform_utils.py @@ -29,7 +29,8 @@ def cuda_get_device_properties( device, names: Sequence[str], init_cuda=False ) -> tuple[Any, ...]: """Get specified CUDA device property values without initializing CUDA in - the current process.""" + the current process. + """ if init_cuda or cuda_is_initialized(): props = torch.cuda.get_device_properties(device) return tuple(getattr(props, name) for name in names) diff --git a/vllm/utils/registry.py b/vllm/utils/registry.py index a136d450e7b1..6df0bead0db1 100644 --- a/vllm/utils/registry.py +++ b/vllm/utils/registry.py @@ -6,8 +6,7 @@ class ExtensionManager: - """ - A registry for managing pluggable extension classes. + """A registry for managing pluggable extension classes. This class provides a simple mechanism to register and instantiate extension classes by name. It is commonly used to implement plugin @@ -26,15 +25,11 @@ class ExtensionManager: """ def __init__(self) -> None: - """ - Initialize an empty extension registry. - """ + """Initialize an empty extension registry.""" self.name2class: dict[str, type] = {} def register(self, name: str): - """ - Decorator to register a class with the given name. - """ + """Decorator to register a class with the given name.""" def wrap(cls_to_register: _T) -> _T: self.name2class[name] = cls_to_register @@ -43,9 +38,7 @@ def wrap(cls_to_register: _T) -> _T: return wrap def load(self, cls_name: str, *args, **kwargs) -> Any: - """ - Instantiate and return a registered extension class by name. - """ + """Instantiate and return a registered extension class by name.""" cls = self.name2class.get(cls_name) assert cls is not None, f"Extension class {cls_name} not found" return cls(*args, **kwargs) diff --git a/vllm/utils/sparse_utils.py b/vllm/utils/sparse_utils.py index 79330271dd51..582d91093eae 100644 --- a/vllm/utils/sparse_utils.py +++ b/vllm/utils/sparse_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Safe deserialization helpers for client-supplied embedding tensors. +"""Safe deserialization helpers for client-supplied embedding tensors. `check_sparse_tensor_invariants_threadsafe()` is a thread-safe wrapper for sparse tensor invariant validation. `safe_to_dense()` bounds the memory a diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 6d56df378f13..cc83a62a1690 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -60,8 +60,7 @@ def set_env_var(key: str, value: str) -> Iterator[None]: @contextlib.contextmanager def suppress_stdout(): - """ - Suppress stdout from C libraries at the file descriptor level. + """Suppress stdout from C libraries at the file descriptor level. Only suppresses stdout, not stderr, to preserve error messages. Suppression is disabled when VLLM_LOGGING_LEVEL is set to DEBUG. @@ -70,6 +69,7 @@ def suppress_stdout(): with suppress_stdout(): # C library calls that would normally print to stdout torch.distributed.new_group(ranks, backend="gloo") + """ # Don't suppress if logging level is DEBUG if envs.VLLM_LOGGING_LEVEL == "DEBUG": @@ -114,7 +114,6 @@ def unique_filepath(fn: Callable[[int], Path]) -> Path: def _sync_visible_devices_env_vars(): """Sync HIP/CUDA visibility env vars before spawning (ROCm only).""" - if not current_platform.is_rocm(): return @@ -254,11 +253,11 @@ def decorate_logs( def kill_process_tree(pid: int): - """ - Kills all descendant processes of the given pid by sending SIGKILL. + """Kills all descendant processes of the given pid by sending SIGKILL. Args: pid (int): Process ID of the parent process + """ try: parent = psutil.Process(pid) diff --git a/vllm/utils/tensor_schema.py b/vllm/utils/tensor_schema.py index 526dfd38bac4..02ddda03f1fc 100644 --- a/vllm/utils/tensor_schema.py +++ b/vllm/utils/tensor_schema.py @@ -161,7 +161,6 @@ def _validate_tensor_shape_expected( dynamic_dims: set[str], ) -> None: """Validate that the actual tensor shape matches the expected shape.""" - if len(actual_shape) != len(expected_shape): raise ValueError( f"{field_name} has rank {len(actual_shape)} " diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 0b330edf0038..ab395bd7047f 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -96,8 +96,7 @@ def is_meta_module(module: torch.nn.Module) -> bool: def is_strictly_contiguous(t: torch.Tensor) -> bool: - """ - Check if tensor is contiguous AND has no degenerate strides. + """Check if tensor is contiguous AND has no degenerate strides. A degenerate stride occurs when a dimension has size 1 but the stride doesn't match the canonical contiguous layout. This can cause issues @@ -231,7 +230,8 @@ def read_v1(cg_dir: str) -> float | None: def available_cpu_count() -> int: """CPUs actually usable by this process: scheduling affinity capped by - the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind).""" + the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind). + """ if sys.platform != "linux": return os.cpu_count() or 1 count = len(os.sched_getaffinity(0)) @@ -294,8 +294,7 @@ def set_torch_threads_for_runtime() -> None: @contextlib.contextmanager def set_default_torch_num_threads(num_threads: int | None = None): - """ - Sets the default number of threads for PyTorch to the given value. + """Sets the default number of threads for PyTorch to the given value. `None` means using the value of the environment variable `OMP_NUM_THREADS` (or `1` if that is not available). @@ -363,8 +362,7 @@ def _get_precision_level(dtype: torch.dtype) -> int: def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype): - """ - Test whether it is lossless to cast a tensor from + """Test whether it is lossless to cast a tensor from `src_dtype` to `tgt_dtype`. """ if src_dtype == tgt_dtype: @@ -395,8 +393,7 @@ def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype): def common_broadcastable_dtype(dtypes: Collection[torch.dtype]): - """ - Get the common `dtype` where all of the other `dtypes` can be + """Get the common `dtype` where all of the other `dtypes` can be cast to it without losing any information. """ return max( @@ -567,6 +564,7 @@ def nvfp4_split_data_scale( Returns: ``(data, scale)`` where *data* is uint8 and *scale* is float8_e4m3fn, both views of the same storage. + """ num_pages = kv_side.shape[0] dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2] @@ -744,8 +742,7 @@ def make_ndarray_with_pad( *, max_len: int | None = None, ) -> npt.NDArray: - """ - Make a padded array from 2D inputs. + """Make a padded array from 2D inputs. The padding is applied to the end of each inner list until it reaches `max_len`. @@ -771,8 +768,7 @@ def make_tensor_with_pad( device: str | torch.device | None = None, pin_memory: bool = False, ) -> torch.Tensor: - """ - Make a padded tensor from 2D inputs. + """Make a padded tensor from 2D inputs. The padding is applied to the end of each inner list until it reaches `max_len`. @@ -806,8 +802,7 @@ def __init__(self): def current_stream() -> torch.cuda.Stream: - """ - replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`. + """Replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`. it turns out that `torch.cuda.current_stream()` is quite expensive, as it will construct a new stream object at each call. here we patch `torch.cuda.set_stream` to keep track of the current stream @@ -858,9 +853,7 @@ def current_stream() -> torch.cuda.Stream: def aux_stream() -> torch.cuda.Stream | None: - """ - Ensures aux_stream is initialized only once - """ + """Ensures aux_stream is initialized only once.""" global _aux_stream from vllm.platforms import current_platform @@ -872,8 +865,7 @@ def aux_stream() -> torch.cuda.Stream | None: def weak_ref_tensor(tensor: Any) -> Any: - """ - Create a weak reference to a tensor. + """Create a weak reference to a tensor. The new tensor will share the same data as the original tensor, but will not keep the original tensor alive. This ignores 0-size tensors as those don't allocate any memory. @@ -890,8 +882,7 @@ def weak_ref_tensors( | tuple[torch.Tensor] | IntermediateTensors, ) -> torch.Tensor | list[Any] | tuple[Any] | Any: - """ - Convenience function to create weak references to tensors, + """Convenience function to create weak references to tensors, for single tensor, list of tensors or tuple of tensors. """ if isinstance(tensors, torch.Tensor): @@ -913,9 +904,7 @@ def weak_ref_tensors( def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor: - """ - Get an accelerator view of a CPU tensor using Unified Virtual Addressing (UVA). - """ + """Get an accelerator view of a CPU tensor via Unified Virtual Addressing.""" from vllm.platforms import current_platform if current_platform.is_xpu(): @@ -951,6 +940,7 @@ def is_torch_equal_or_newer(target: str) -> bool: Returns: Whether the condition meets. + """ try: return _is_torch_equal_or_newer(str(torch.__version__), target) @@ -979,6 +969,7 @@ def is_torch_equal(target: str) -> bool: Returns: Whether the condition meets. + """ try: return _is_torch_equal(target) @@ -1064,8 +1055,7 @@ def direct_register_custom_op( dispatch_key: str | None = None, tags: tuple[torch.Tag, ...] = (), ): - """ - `torch.library.custom_op` can have significant overhead because it + """`torch.library.custom_op` can have significant overhead because it needs to consider complicated dispatching logic. This function directly registers a custom op and dispatches it to the CUDA backend. See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 01599e9e9c72..792eeb37b99d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -34,8 +34,7 @@ class AttentionType(str, Enum): - """ - Attention type. + """Attention type. Use string to be compatible with `torch.compile`. """ @@ -366,7 +365,8 @@ def validate_configuration( @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: """Layouts this backend's kernels can consume, most preferred first, or - None when the kernels consume any layout and express no preference.""" + None when the kernels consume any layout and express no preference. + """ return None @classmethod @@ -383,8 +383,7 @@ class AttentionMetadata: @dataclass class CommonAttentionMetadata: - """ - Per-batch attention metadata, shared across layers and backends. + """Per-batch attention metadata, shared across layers and backends. AttentionMetadataBuilder instances use it to construct per-layer metadata. For many of the tensors we keep both GPU and CPU versions. @@ -559,7 +558,8 @@ def unpadded( class AttentionCGSupport(Enum): """Constants for the cudagraph support of the attention backend Here we do not consider the cascade attention, as currently - it is never cudagraph supported.""" + it is never cudagraph supported. + """ ALWAYS = 3 """Cudagraph always supported; supports mixed-prefill-decode""" @@ -655,8 +655,7 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> M: - """ - Central method that builds attention metadata. + """Central method that builds attention metadata. Some builders (MLA) require reorder_batch to be called prior to build. Args: @@ -665,6 +664,7 @@ def build( fast_build: The meta-data will prioritize speed of building over then speed at execution. Can be used for spec-decode where the result of a build call may only be used for few layers/iters. + """ raise NotImplementedError @@ -674,8 +674,7 @@ def update_block_table( blk_table: torch.Tensor, slot_mapping: torch.Tensor, ) -> M: - """ - Update the block table for the attention metadata. + """Update the block table for the attention metadata. Faster when theres multiple kv-cache groups that create virtually the same metadata but just with different block tables. @@ -686,8 +685,7 @@ def update_block_table( def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata ) -> M: - """ - Build attention metadata for CUDA graph capture. Uses build by default. + """Build attention metadata for CUDA graph capture. Uses build by default. Subclasses that override this method should call self.build or super().build_for_cudagraph_capture. """ @@ -700,8 +698,7 @@ def build_for_drafting( common_attn_metadata: CommonAttentionMetadata, draft_index: int, ) -> M: - """ - Build attention metadata for draft model. Uses build by default. + """Build attention metadata for draft model. Uses build by default. Args: common_attn_metadata: The common attention metadata. @@ -710,6 +707,7 @@ def build_for_drafting( draft attempt for the i-th token. For tree-based attention, this index instead refers to the draft attempt for the i-th level in the tree of tokens. + """ return self.build( common_prefix_len=0, @@ -910,8 +908,7 @@ def forward( raise NotImplementedError def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool: - """ - Does this attention implementation support fused output quantization. + """Does this attention implementation support fused output quantization. This is used by the AttnFusionPass to only fuse output quantization onto implementations that support it. @@ -920,20 +917,19 @@ def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool: Returns: is fusion supported for this type of quantization + """ return False def fused_qk_norm_rope_kvcache_supported(self): - """ - Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. + """Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops with the RoPE ops and the KV cache update for implementations that support it. """ return False def fused_rope_kvcache_supported(self): - """ - Does this attention implementation support RoPE+KVCache fusion. + """Does this attention implementation support RoPE+KVCache fusion. This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops with the KV cache update for implementations that support it. """ @@ -954,8 +950,7 @@ def do_qk_norm_rope_kvcache_update( kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - """ - If `fused_qk_norm_rope_kvcache_supported` returns True, this method + """If `fused_qk_norm_rope_kvcache_supported` returns True, this method will be called by the fused custom op. Applies QK-norm + RoPE and writes K/V to the KV cache. Results are written to the pre-allocated q_out and k_out tensors; V is split from QKV at the graph level. @@ -974,8 +969,7 @@ def do_rope_and_kv_cache_update( kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - """ - If `fused_rope_kvcache_supported` returns True, this method will be called + """If `fused_rope_kvcache_supported` returns True, this method will be called by torch.ops.vllm.fused_rope_and_unified_kv_cache_update to perform the inplace RoPE and KV cache update. """ @@ -1040,8 +1034,7 @@ def forward_mqa( raise NotImplementedError def fused_output_quant_supported(self, quant_key: "QuantKey"): - """ - Does this attention implementation support fused output quantization. + """Does this attention implementation support fused output quantization. Since MLA quantization is done manually in forward_impl (common code), all MLA backends support it by default. """ @@ -1080,9 +1073,7 @@ def subclass_attention_backend( attention_backend_cls: type[AttentionBackend], builder_cls: type[AttentionMetadataBuilder[M]], ) -> type[AttentionBackend]: - """ - Return a new subclass where `get_builder_cls` returns `builder_cls`. - """ + """Return a new subclass where `get_builder_cls` returns `builder_cls`.""" name: str = name_prefix + attention_backend_cls.__name__ # type: ignore return type( diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index dd3f34c9167e..d2b9f5a526e8 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -84,7 +84,8 @@ def supports_sliding_window(cls) -> bool: @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """CPU attention supports decoder, - encoder-only and encoder-decoder attention.""" + encoder-only and encoder-decoder attention. + """ return attn_type in ( AttentionType.DECODER, AttentionType.ENCODER, @@ -349,14 +350,22 @@ def forward( """Forward pass for CPU attention backend. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization; not supported + by this backend. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] + """ if output_scale is not None or output_block_scale is not None: raise NotImplementedError( diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 0de6b85217b1..650fc1dea088 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -357,6 +357,7 @@ def is_flash_attn_varlen_func_available() -> bool: Returns: bool: True if a working flash_attn_varlen_func implementation is available. + """ if current_platform.is_cuda() or current_platform.is_xpu(): # CUDA and XPU always have flash_attn_varlen_func available diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index eb2268cc624c..b57af3fd295e 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -766,8 +766,7 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> FlashAttentionMetadata: - """ - fast_build disables AOT scheduling, used when there will be few + """fast_build disables AOT scheduling, used when there will be few iterations i.e. spec-decode """ num_reqs = common_attn_metadata.num_reqs @@ -1178,17 +1177,25 @@ def forward( """Forward pass with FlashAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization; not supported + by this backend. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] NOTE: FP8 quantization, flash-attn expect the size of {q,k,v}_descale to be (num_sequences, num_kv_heads). We use torch's .expand() to avoid duplicating values + """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." @@ -1665,6 +1672,7 @@ def _forward_encoder_attention( output: shape = [num_encoder_tokens, num_heads, head_size] attn_metadata: Encoder attention metadata layer: The attention layer + """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 40846758d833..4bb58ef4ea43 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -131,17 +131,25 @@ def forward( """Forward pass with FlashAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size_v] kv_cache: shape = [num_blocks, block_size, num_kv_heads, head_size + head_size_v] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization; not supported + by this backend. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size_v] NOTE: FP8 quantization, flash-attn expect the size of {q,k,v}_descale to be (num_sequences, num_kv_heads). We use torch's .expand() to avoid duplicating values + """ assert self.vllm_flash_attn_version is not None, ( "FlashAttention version not detected." diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 7ab2b2018e09..694cef4106cf 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -398,7 +398,8 @@ class FlashInferBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """NVFP4 stores K and V as separate per-head slots of packed fp4 data - plus fp8 block scales.""" + plus fp8 block scales. + """ if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4: return spec hs_k = nvfp4_kv_cache_full_dim(spec.head_size) @@ -1253,9 +1254,7 @@ def _compute_flashinfer_kv_metadata( num_reqs: int, page_size: int, ) -> torch.Tensor: - """ - Compute paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len for FlashInfer - attention. + """Compute paged_kv_indptr, paged_kv_indices and paged_kv_last_page_len. Results are stored in self.paged_kv_indptr, self.paged_kv_indices, self.paged_kv_last_page_len buffers. @@ -1964,6 +1963,7 @@ def forward( """Forward pass with FlashInfer. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size], or None for a KV-sharing decoder layer. @@ -1971,8 +1971,15 @@ def forward( KV-sharing decoder layer. kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization. Enables the + attention+quantization fusion path when provided. + output_block_scale: Block scale for fused output quantization, + required for nvfp4 output and rejected for fp8 output. + Returns: shape = [num_tokens, num_heads * head_size] + """ if attn_metadata is None: # Profiling run. @@ -2638,8 +2645,7 @@ def fast_plan_decode( fixed_split_size: int = -1, disable_split_kv: bool = False, ) -> None: - """ - A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for + """A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for cudagraph capture/replay, while the no cudagraph version turns back to the original plan. using original plan after passing host-side buffers: diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 70eded1d093f..9d0ac081725e 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -165,8 +165,7 @@ def physical_to_logical_mapping( block_size: int, total_blocks: int, ) -> torch.Tensor: - """ - Creates an inverse mapping from physical block locations to logical indices. + """Creates an inverse mapping from physical block locations to logical indices. The original block_table maps from logical blocks to physical locations: @@ -239,6 +238,7 @@ def physical_to_logical_mapping( A tensor of shape [max_reqs, total_blocks] where each entry physical_to_logical[req_id, physical_block] contains the logical block index for that physical block, or -1 if unused. + """ max_reqs, max_num_blocks = block_table.shape device = block_table.device @@ -275,8 +275,7 @@ def unique_static_unsorted( ignored_val: int = 0, # value to ignore pad_val: int = -1, # sentinel for unused slots ) -> torch.Tensor: - """ - - Keeps the first occurrence of each non-zero value while preserving order, + """- Keeps the first occurrence of each non-zero value while preserving order, then left-packs those uniques and fills the rest with `pad_val`. - Returns (packed, keep_mask) with the *same shape* as `x`. - Requires that all values be in the range [0, M] @@ -286,6 +285,7 @@ def unique_static_unsorted( Example: x =[3, 1, 0, 1, 2], M=3, ignored_val=0 => [3, 1, 2, -1, -1] + """ if not (-1 <= pad_val <= M): raise ValueError("`pad_val` must lie in [-1, M]") @@ -349,6 +349,7 @@ class BlockSparsityHint(NamedTuple): hint_fn: (q_block_idx [num_tokens, 1], kv_block_idx [1, num_kv_blocks], block_size int) -> bool Tensor [num_tokens, num_kv_blocks]. Returns True for block pairs that may contain non-masked elements. + """ hint_fn: _block_sparsity_hint_signature @@ -435,6 +436,7 @@ def _convert_physical_to_logical( Returns: tuple of (is_valid, logical_q_idx, logical_kv_idx) + """ # Map query indices to corresponding request indices q_req = request_lookup[q_idx] @@ -512,7 +514,6 @@ def get_sliding_window_mask_mod(self) -> _mask_mod_signature: Note that the sliding window mask here is bidirectional, we need to mask it with the bidirectional/causal mask for encoder/decoder. """ - if self.sliding_window is None: raise ValueError("sliding_window must be set for sliding window attention") @@ -540,7 +541,6 @@ def final_mask_mod( def get_prefix_lm_mask_mod(self) -> _mask_mod_signature: """Creates the prefix LM mask_mod function for FlexAttention.""" - assert self.doc_ids is not None request_lookup = self.doc_ids @@ -587,7 +587,6 @@ def get_rswa_mask_mod(self) -> _mask_mod_signature: far-away generated tokens that fall outside the window and outside the prefix. """ - assert self.doc_ids is not None assert self.rswa_prefix_lens is not None assert self.rswa_window is not None @@ -1291,14 +1290,22 @@ def forward( """Forward pass with FLexAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization; not supported + by this backend. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] + """ if output_scale is not None or output_block_scale is not None: raise NotImplementedError( diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 23daff54300b..7a839ff73b4a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -524,8 +524,7 @@ def build( # type: ignore[override] def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata ): - """ - This method builds the metadata for full cudagraph capture. + """This method builds the metadata for full cudagraph capture. Currently, only decode is supported for full cudagraphs with Mamba. """ m = common_attn_metadata diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index 189ef54b236d..0a4584632dca 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -23,8 +23,7 @@ def compute_varlen_chunk_metadata( query_start_loc: torch.Tensor, chunk_size: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels. + """Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels. Given per-sequence cumulative token starts `query_start_loc` of shape [B+1] and a physical `chunk_size`, returns three tensors on the same device: diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index e0b5fab465a1..b5f0b9b182e1 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -227,8 +227,7 @@ def __init__( def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata ) -> M: - """ - This method builds the metadata for full cudagraph capture. + """This method builds the metadata for full cudagraph capture. Currently, only decode is supported for full cudagraphs with Mamba. """ m = common_attn_metadata @@ -276,8 +275,7 @@ def build( num_decode_draft_tokens_cpu: torch.Tensor | None = None, **kwargs: Any, ) -> M: - """ - Default build implementation for Mamba-like attention backends. + """Default build implementation for Mamba-like attention backends. Subclasses (e.g., Mamba2) can override to add additional metadata. """ return self._compute_common_metadata( @@ -294,8 +292,7 @@ def _compute_chunk_metadata( num_computed_tokens_p_cpu: torch.Tensor, query_start_loc_p_cpu: torch.Tensor, ) -> tuple[list[int], list[int], list[int]]: - """ - Compute chunk-specific metadata for Mamba models. + """Compute chunk-specific metadata for Mamba models. The code below carefully constructs the chunks such that: 1. Chunks contain tokens from a *single* sequence only. @@ -382,8 +379,7 @@ def _build_chunk_metadata_tensors( common: M, common_attn_metadata: CommonAttentionMetadata, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Compute chunk metadata and return as device tensors. + """Compute chunk metadata and return as device tensors. Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p). """ num_prefills = common.num_prefills @@ -463,9 +459,7 @@ def _compute_common_metadata( prev_last_scheduled_idx: torch.Tensor | None = None, num_decode_draft_tokens_cpu: torch.Tensor | None = None, ) -> M: - """ - Compute metadata common to both Mamba1 and Mamba2. - """ + """Compute metadata common to both Mamba1 and Mamba2.""" num_reqs = common_attn_metadata.num_reqs # Treat multi-token queries as decode requests when @@ -752,8 +746,7 @@ def _update_metadata_for_cudagraph_capture( self, metadata: M, ) -> M: - """ - Update the metadata for cudagraph capture. + """Update the metadata for cudagraph capture. Currently, only decode is supported for full cudagraphs with Mamba. """ state_indices_tensor_d = metadata.state_indices_tensor_d diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 5f9ac7f17adf..22ee41a4799b 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -63,7 +63,8 @@ def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread - count in place of SM count.""" + count in place of SM count. + """ ideal_splits = 1 while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT): ideal_splits *= 2 diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 3a9ca2ffdb06..6211a4e06c4c 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -437,7 +437,6 @@ def _build_fp8_mixed_decode_prefill( be the full batch or only decodes when prefills use dense MHA. This avoids the BF16 prefill kernel's head-padding overhead at high TP. """ - scheduler_metadata, _ = get_mla_metadata() return FlashMLASparseMetadata.FP8KernelMetadata( scheduler_metadata=scheduler_metadata, @@ -848,7 +847,8 @@ def _gather_prefill_chunk( prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """All-gather this rank's upconverted KV shard so the chunk's rows attend - the whole context, and map their top-k onto the rank-major result.""" + the whole context, and map their top-k onto the rank-major result. + """ shard_rows = int(chunk.chunk_tot_seqlen) assert self.gathered_kv_workspace is not None gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows] diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 53206afd9337..258f20215f13 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1118,7 +1118,8 @@ def _split_pcp_dcp_prefill_chunks( def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: """Per-request KV lengths the prefill chunker budgets logits with; - subclasses whose logits rows are wider than the context override.""" + subclasses whose logits rows are wider than the context override. + """ return seq_lens_cpu @staticmethod diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 3ba5a5a5fad3..1d03b9e77392 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -24,7 +24,7 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend): - """AITER FlashAttention backend for MLA prefill""" + """AITER FlashAttention backend for MLA prefill.""" @staticmethod def get_name() -> str: diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 1b589cb295f4..edea4347271a 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -70,7 +70,8 @@ def is_available(cls) -> bool: def supports_quant_output(self, quant_key: "QuantKey") -> bool: """Whether `run_prefill_new_tokens` can write quantized output directly (fused) for the given quant key, skipping the post-quant - pass. Overridden by backends that support it.""" + pass. Overridden by backends that support it. + """ return False def supports_out(self) -> bool: diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index e7e00a0d0143..1cf1761a7708 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -65,6 +65,7 @@ def get_path(self) -> str: Raises: ValueError: If Backend.CUSTOM is used without being registered + """ path = _MLA_PREFILL_OVERRIDES.get(self, self.value) if not path: @@ -85,6 +86,7 @@ def get_class(self) -> "type[MLAPrefillBackend]": Raises: ImportError: If the backend class cannot be imported ValueError: If CUSTOM is used without being registered + """ return resolve_obj_by_qualname(self.get_path()) @@ -130,6 +132,7 @@ class MyCustomPrefillBackend(MLAPrefillBackend): MLAPrefillBackendEnum.CUSTOM, "my.module.MyCustomPrefillBackend" ) + """ def decorator(cls: type) -> type: diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index c7e935fa75bb..a2b74e4327eb 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -57,6 +57,7 @@ def _get_mla_prefill_backend_priorities( Returns: List of backends in priority order (highest priority first). + """ from vllm.platforms import current_platform @@ -104,6 +105,7 @@ def get_mla_prefill_backend( Returns: The selected prefill backend class. + """ from vllm.platforms import current_platform @@ -173,6 +175,7 @@ def _auto_select_mla_prefill_backend( Returns: The selected prefill backend class. + """ priorities = _get_mla_prefill_backend_priorities( device_capability, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index c9ba2b0a566e..75173e59a3c4 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -710,6 +710,7 @@ def _init_fp8_prefill_ps_buffers( attn_out_dtype: Dtype of the attention output buffer, used to size the padded-head output scratch (small head counts only). device: Target device for the buffers. + """ from aiter import get_ps_metadata_info_v1 @@ -1328,8 +1329,7 @@ def _expand_page_indices_kernel( class AiterMLAHelper: - """ - AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned + """AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned head counts through 128 are padded to the next multiple of 16 by tiling the query heads and slicing to the padded size. Native H24 AITER builds bypass that padding. Small divisors of 16 retain the existing repeat_interleave and diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 86b77db79a81..e1365f4d30c2 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -167,8 +167,7 @@ def triton_convert_req_index_to_global_index( NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 128, # tile width along columns ): - """ - out[token_id, indice_id] = + """out[token_id, indice_id] = block_table[req_id[token_id], token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + token_indices[token_id, indice_id] % BLOCK_SIZE diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 242c8f3b7843..0ac131cee5df 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -426,6 +426,7 @@ def _remap_tiling( Returns: (single_tile, block_n, tiles_per_row, num_warps) + """ single_tile = ( count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS @@ -454,8 +455,7 @@ def triton_convert_req_index_to_global_index( out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - out[token_id, indice_id] = + """out[token_id, indice_id] = block_table[req_id[token_id], token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + token_indices[token_id, indice_id] % BLOCK_SIZE diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index dcaaffaaa8eb..32e2fabeb567 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Attention backend registry""" +"""Attention backend registry.""" from collections.abc import Callable from enum import Enum, EnumMeta @@ -161,6 +161,7 @@ def get_path(self, include_classname: bool = True) -> str: Raises: ValueError: If Backend.CUSTOM is used without being registered + """ path = _ATTN_OVERRIDES.get(self, self.value) if not path: @@ -181,6 +182,7 @@ def get_class(self) -> "type[AttentionBackend]": Raises: ImportError: If the backend class cannot be imported ValueError: If Backend.CUSTOM is used without being registered + """ return resolve_obj_by_qualname(self.get_path()) @@ -189,6 +191,7 @@ def is_overridden(self) -> bool: Returns: True if the backend has a registered override + """ return self in _ATTN_OVERRIDES @@ -224,6 +227,7 @@ def get_path(self, include_classname: bool = True) -> str: Raises: ValueError: If Backend.CUSTOM is used without being registered + """ path = _MAMBA_ATTN_OVERRIDES.get(self, self.value) if not path: @@ -244,6 +248,7 @@ def get_class(self) -> "type[AttentionBackend]": Raises: ImportError: If the backend class cannot be imported ValueError: If Backend.CUSTOM is used without being registered + """ return resolve_obj_by_qualname(self.get_path()) @@ -252,6 +257,7 @@ def is_overridden(self) -> bool: Returns: True if the backend has a registered override + """ return self in _MAMBA_ATTN_OVERRIDES @@ -275,6 +281,7 @@ def register_backend( backend: The AttentionBackendEnum member to register class_path: Optional class path. If not provided and used as decorator, will be auto-generated from the class. + is_mamba: Whether the backend is a Mamba attention backend. Returns: Decorator function if class_path is None, otherwise a no-op @@ -300,6 +307,7 @@ class MyCustomBackend: AttentionBackendEnum.CUSTOM, "my.module.MyCustomBackend" ) + """ def decorator(cls: type) -> type: diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 062aa82285a1..1615f883bc02 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -770,8 +770,7 @@ def build_for_drafting( common_attn_metadata: CommonAttentionMetadata, draft_index: int, ) -> AiterFlashAttentionMetadata: - """ - Build attention metadata for draft model without CPU-GPU sync. + """Build attention metadata for draft model without CPU-GPU sync. During EAGLE drafting all requests are uniform decodes, so we can skip split_decodes_prefills_and_extends() and avoid all .cpu() / @@ -1189,17 +1188,25 @@ def forward( """Forward pass with AiterFlashAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, 2, block_size, num_kv_heads, head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization; not supported + by this backend. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] NOTE: FP8 quantization, flash-attn expect the size of {q,k,v}_descale to be (num_sequences, num_kv_heads). We use torch's .expand() to avoid duplicating values + """ if output_scale is not None or output_block_scale is not None: raise NotImplementedError( diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index e0e4b385509c..b9eacdabca12 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -110,7 +110,8 @@ def get_impl_cls() -> type["RocmAiterUnifiedAttentionImpl"]: @classmethod def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: """Keep K and V packed in the content dim, unlike the native HIP - kernels the base class targets.""" + kernels the base class targets. + """ # block_size == 1 is the per-token page-size probe (see # Platform.get_page_size_bytes); real blocks must be gatherable in # 16-token units by the ROCm kernel. @@ -207,14 +208,21 @@ def forward( """Forward pass with FlashAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, 2, block_size, num_kv_heads, head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] + """ if output_block_scale is not None: raise NotImplementedError( diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index f3bb3b351451..afb545eee8f1 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -244,7 +244,8 @@ def supports_attn_type(cls, attn_type: str) -> bool: @classmethod def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: """K and V as two head groups so the native HIP kernels address each side - as one contiguous region (x-packed interior applied in split_kv_cache).""" + as one contiguous region (x-packed interior applied in split_kv_cache). + """ if spec.state_content_bytes is not None: return spec assert spec.head_size == spec.head_size_v, ( @@ -344,6 +345,7 @@ def _forward_encoder_attention( output: shape = [num_encoder_tokens, num_heads, head_size] attn_metadata: Encoder attention metadata layer: The attention layer + """ # For encoder attention, process FP8 quantization if needed if is_quantized_kv_cache(self.kv_cache_dtype): @@ -390,14 +392,21 @@ def forward( """Forward pass with FlashAttention. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: logical [num_blocks, 2, block_size, num_kv_heads * head_size] under LHBNC (physically K/V-group-first) attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] + """ if output_block_scale is not None: raise NotImplementedError( diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index b7e2a41d00e5..cf7e571599cd 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -275,7 +275,8 @@ class TritonAttentionBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack inline fp32 scales after each head's - data, so the content is (data + one scale) per K/V side.""" + data, so the content is (data + one scale) per K/V side. + """ mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: return spec @@ -557,14 +558,21 @@ def forward( """Forward pass with Paged Attention impl. in Triton. Args: + layer: The attention layer, providing the q/k/v quantization scales. query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. + output: Tensor that the attention result is written into. + output_scale: Scale for fused output quantization. + output_block_scale: Block scale for fused output quantization; + not supported by this backend. + Returns: shape = [num_tokens, num_heads * head_size] + """ if output_block_scale is not None: raise NotImplementedError( @@ -724,6 +732,7 @@ def _forward_encoder_attention( output: shape = [num_encoder_tokens, num_heads, head_size] attn_metadata: Encoder attention metadata layer: The attention layer + """ # Quantized KV cache is not supported for encoder attention. if is_quantized_kv_cache(self.kv_cache_dtype): diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index e656343a9189..3413c5ad00ca 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -324,8 +324,7 @@ def resolve_kv_cache_layout( @dataclass class PerLayerParameters: - """ - Currently, FlashInfer backend only support models in which all layers share + """Currently, FlashInfer backend only support models in which all layers share the same values for the following hyperparameters. Should not be used for trtllm-gen backend since it supports different values for the following hyperparameters. @@ -343,11 +342,9 @@ class PerLayerParameters: def get_per_layer_parameters( vllm_config: VllmConfig, layer_names: list[str], cls_: type["AttentionImpl"] ) -> dict[str, PerLayerParameters]: - """ - Scan layers in `layer_names` and determine some hyperparameters + """Scan layers in `layer_names` and determine some hyperparameters to use during `plan`. """ - layers = get_layers_from_vllm_config( vllm_config, AttentionLayerBase, # type: ignore[type-abstract] @@ -405,8 +402,7 @@ def get_num_attention_heads_from_layers( def infer_global_hyperparameters( per_layer_params: dict[str, PerLayerParameters], ) -> PerLayerParameters: - """ - Currently, FlashInfer backend other than trtllm-gen + """Currently, FlashInfer backend other than trtllm-gen only support models in which all layers share the same values for the following hyperparameters: - `window_left` @@ -416,7 +412,6 @@ def infer_global_hyperparameters( So this function asserts that all layers share the same values for these hyperparameters and returns the global values. """ - assert len(per_layer_params) > 0, "No attention layers found in the model." param_sets = list(per_layer_params.values()) @@ -710,8 +705,7 @@ def split_decodes_prefills_and_extends( common_attn_metadata: CommonAttentionMetadata, decode_threshold: int = 1, ) -> tuple[int, int, int, int, int, int]: - """ - Assuming a reordered batch, finds the boundary between prefill and decode + """Assuming a reordered batch, finds the boundary between prefill and decode requests. Args: @@ -726,6 +720,7 @@ def split_decodes_prefills_and_extends( num_decode_tokens: The number of tokens in the decode requests. num_extend_tokens: The number of tokens in the extend requests. num_prefill_tokens: The number of tokens in the prefill requests. + """ max_query_len = common_attn_metadata.max_query_len num_reqs = common_attn_metadata.num_reqs @@ -784,8 +779,7 @@ def split_decodes_and_prefills( require_uniform: bool = False, treat_short_extends_as_decodes: bool = True, ) -> tuple[int, int, int, int]: - """ - Assuming a reordered batch, finds the boundary between prefill and decode + """Assuming a reordered batch, finds the boundary between prefill and decode requests. The batch is expected to be ordered as: @@ -807,6 +801,7 @@ def split_decodes_and_prefills( num_prefills: The number of prefill requests. num_decode_tokens: The number of tokens in the decode requests. num_prefill_tokens: The number of tokens in the prefill requests. + """ max_query_len = common_attn_metadata.max_query_len num_reqs = common_attn_metadata.num_reqs @@ -855,16 +850,17 @@ def split_decodes_and_prefills( def split_prefill_chunks( seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0 ) -> list[tuple[int, int]]: - """ - Split the prefill requests into chunks such that the total sequence length + """Split the prefill requests into chunks such that the total sequence length of each chunk is less than or equal to the workspace size. Args: seq_lens_cpu: The sequence lengths of the prefill requests on CPU. workspace_size: The maximum workspace size (in tokens) per chunk. request_offset: The offset to add to the request indices. + Returns: A list of tuples of (reqs_start, reqs_end) representing chunk boundaries. + """ chunk_bounds = [] i, n = 0, len(seq_lens_cpu) @@ -884,8 +880,7 @@ def reorder_batch_to_split_decodes_and_prefills( scheduler_output: "SchedulerOutput", decode_threshold: int = 1, ) -> bool: - """ - Reorders the batch to split into prefill and decode requests; places all + """Reorders the batch to split into prefill and decode requests; places all requests with <= decode_threshold tokens at the front of the batch. The batch is reordered into 4 regions: @@ -896,6 +891,7 @@ def reorder_batch_to_split_decodes_and_prefills( Returns: True if the batch was modified, False otherwise. + """ num_reqs = len(input_batch.req_ids) num_scheduled_tokens = [ @@ -960,8 +956,7 @@ def reorder_batch_to_split_decodes_and_prefills( def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor: - """ - Reshapes the query tensor for the specified batch size, so that + """Reshapes the query tensor for the specified batch size, so that it has shape (batch_size, seq_len, num_heads, head_dim). """ assert query.dim() == 3, f"query must be 3D, got {query.dim()}D" @@ -976,8 +971,7 @@ def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tensor: - """ - Reshapes the attention output tensor, so that + """Reshapes the attention output tensor, so that the batch_size and seq_len dimensions are combined. """ if attn_output.dim() == 3: @@ -1143,8 +1137,7 @@ def mamba_get_block_table_tensor( kv_cache_spec: KVCacheSpec, mamba_cache_mode: str, ) -> torch.Tensor: - """ - Get the block table tensor for mamba kernels from the input + """Get the block table tensor for mamba kernels from the input common_attn_metadata.block_table_tensor given different mamba cache modes. - "all": input (#requests, cdiv(max_model_len, block_size) diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index b5618170556b..5c1677cac479 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -178,6 +178,7 @@ def pack_seq_triton( Returns: packed: [B, Lmax, ...] — packed tensor. + """ is_uint8 = x.dtype == torch.uint8 if is_uint8: @@ -334,8 +335,7 @@ def unpack_seq_triton( block_t: int = 64, block_d: int = 64, ) -> torch.Tensor: - """ - Unpack a packed decode query tensor back to the original format. + """Unpack a packed decode query tensor back to the original format. Efficient Triton implementation. Args: @@ -346,8 +346,8 @@ def unpack_seq_triton( Returns: unpacked_tensor: [N, ...] where N = sum(lengths) - """ + """ # Handle multi-dimensional input by reshaping to (B, Lmax, -1) original_shape = packed_tensor.shape if len(original_shape) > 3: diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 58ac1e14a09d..8f844b647641 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -133,8 +133,7 @@ def kernel( N_ROUNDED: tl.constexpr, IS_BASE_E: tl.constexpr, ): - """ - Apply the all-gathered lses to correct each local rank's attention + """Apply the all-gathered lses to correct each local rank's attention output. we still need perform a cross-rank reduction to obtain the final attention output. @@ -147,6 +146,17 @@ def kernel( Pointer to output tensor of shape [ B, H, D ] vlse_ptr (triton.PointerType): Pointer to output tensor of shape [ B, H ] + outputs_stride_B (int): Batch stride of ``outputs_ptr`` + outputs_stride_H (int): Head stride of ``outputs_ptr`` + outputs_stride_D (int): Head-dim stride of ``outputs_ptr`` + lses_stride_N (int): Rank stride of ``lses_ptr`` + lses_stride_B (int): Batch stride of ``lses_ptr`` + lses_stride_H (int): Head stride of ``lses_ptr`` + lse_idx (int): Index of this rank's lse within the all-gathered tensor + HEAD_DIM: Head dimension, as a constexpr + N_ROUNDED: Rank count rounded to a power of two, as a constexpr + IS_BASE_E: Whether the lses are natural-log based, as a constexpr + """ batch_idx = tl.program_id(axis=0).to(tl.int64) head_idx = tl.program_id(axis=1).to(tl.int64) @@ -376,9 +386,11 @@ def correct_attn_out( lses: Tensor of shape [ N, B, H ] cp_rank: Current rank in the context-parallel group ctx: Triton context to avoid recompilation + is_lse_base_on_e: Whether the lses use base e rather than base 2 Returns: Tuple of (out, lse) with corrected attention and final log-sum-exp. + """ if ctx is None: ctx = CPTritonContext() @@ -431,8 +443,7 @@ def _cp_lse_common( seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, ): - """ - cp_attn_out: [ B, H, D ] + """cp_attn_out: [ B, H, D ] cp_attn_lse: [ B, H ] """ if cp_group.world_size == 1: @@ -466,8 +477,7 @@ def cp_lse_ag_out_rs( seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, ): - """ - cp_attn_out: [ B, H, D ] + """cp_attn_out: [ B, H, D ] cp_attn_lse: [ B, H ] """ out, lse = _cp_lse_common( @@ -499,8 +509,7 @@ def cp_lse_ag_out_ar( seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, ): - """ - cp_attn_out: [ B, H, D ] + """cp_attn_out: [ B, H, D ] cp_attn_lse: [ B, H ] """ out, lse = _cp_lse_common( @@ -528,8 +537,7 @@ def _lse_weighted_combine( return_lse: bool = False, is_lse_base_on_e: bool = True, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - CPU reference implementation for LSE-weighted combination. + """CPU reference implementation for LSE-weighted combination. This is a pure PyTorch implementation used for testing and validation. @@ -545,6 +553,7 @@ def _lse_weighted_combine( Returns: Combined output [B, H, D], and optionally global LSE [B, H] + """ N, B, H, D = outputs.shape @@ -936,8 +945,7 @@ def dcp_a2a_lse_reduce( seq_lens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Combine partial attention outputs across DCP ranks using All-to-All. + """Combine partial attention outputs across DCP ranks using All-to-All. The output and LSE are packed into a single output-dtype buffer, sent with one All-to-All, then unpacked and combined with exact LSE weighting. @@ -955,6 +963,7 @@ def dcp_a2a_lse_reduce( Returns: Combined output [B, H/N, D] (head-scattered) If return_lse=True, also returns global_lse [B, H/N] + """ world_size = cp_group.world_size diff --git a/vllm/v1/attention/ops/flashmla.py b/vllm/v1/attention/ops/flashmla.py index c84a495859ba..eb5383e5d13a 100644 --- a/vllm/v1/attention/ops/flashmla.py +++ b/vllm/v1/attention/ops/flashmla.py @@ -49,9 +49,7 @@ def _is_flashmla_available() -> tuple[bool, str | None]: def is_flashmla_dense_supported() -> tuple[bool, str | None]: - """ - Return: is_supported_flag, unsupported_reason (optional). - """ + """Return: is_supported_flag, unsupported_reason (optional).""" is_available, maybe_reason = _is_flashmla_available() if not is_available: return False, maybe_reason @@ -61,9 +59,7 @@ def is_flashmla_dense_supported() -> tuple[bool, str | None]: def is_flashmla_sparse_supported() -> tuple[bool, str | None]: - """ - Return: is_supported_flag, unsupported_reason (optional). - """ + """Return: is_supported_flag, unsupported_reason (optional).""" is_available, maybe_reason = _is_flashmla_available() if not is_available: return False, maybe_reason diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py index 6791e1f5fd1e..d28c78bdb9fa 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -106,8 +106,8 @@ def merge_attn_states( are treated as having context. output_scale: Optional scalar tensor for FP8 static quantization. When provided, output must be FP8 dtype. - """ + """ # Both the CUDA and Triton kernels derive the suffix head stride from # prefix_output, so suffix_output must share the same head stride. assert prefix_output.stride(1) == suffix_output.stride(1), ( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 0e342cb0e064..07cdf181069a 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -654,6 +654,7 @@ def rocm_fp8_paged_mqa_logits( Returns: Logits tensor of shape [B * next_n, max_model_len], dtype `torch.float32`. + """ from vllm._aiter_ops import rocm_aiter_ops @@ -737,6 +738,7 @@ def fp8_mqa_logits_torch( Returns: Logits tensor of shape [M, N], dtype `torch.float32`. + """ k_fp8, scale = kv seq_len_kv = k_fp8.shape[0] @@ -805,8 +807,8 @@ def rocm_fp8_mqa_logits( Returns: Logits tensor of shape [M, N], dtype `torch.float32`. - """ + """ from vllm._aiter_ops import rocm_aiter_ops k_fp8, scale = kv diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py index 6aec1db2e59c..130760d1ded8 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -24,8 +24,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -""" -Memory-efficient attention for decoding. +"""Memory-efficient attention for decoding. It supports page size >= 1. """ diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 0fb31ce7e1dc..00d630edcdd5 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -19,8 +19,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -""" -Memory-efficient attention for prefill. +"""Memory-efficient attention for prefill. It supports page size = 1. """ @@ -229,8 +228,7 @@ def context_attention_fwd( sinks: torch.Tensor | None = None, sinks_bias_key0: bool = False, ): - """ - q, k, v: [b * s, head, head_dim] + """q, k, v: [b * s, head, head_dim] b_start_loc: [b] b_seq_len: [b] out: [b * s, head, head_dim] diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py index 0dc4150cbf15..90b620b345ee 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py @@ -45,7 +45,8 @@ def _store_quantized_value( BLOCK_GRP: tl.constexpr, ): """Uniform quantization of values to VQB bits, pack, and store with - scale/zero in the per-block SoA metadata region (Opt#3 layout).""" + scale/zero in the per-block SoA metadata region (Opt#3 layout). + """ if VQB == 3: val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to( tl.float32 diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index b7aeb8c6a769..b90eb58e9bd4 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file contains ops for ViT attention to be compatible with torch.compile +"""This file contains ops for ViT attention to be compatible with torch.compile as there are operations here not supported by torch.compile (for instance, `.item()` in flash attention) @@ -238,8 +237,7 @@ def apply_sdpa( scale: float | None = None, enable_gqa: bool = False, ) -> torch.Tensor: - """ - Input shape: + """Input shape: (batch_size x seq_len x num_heads x head_size) """ q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v]) diff --git a/vllm/v1/attention/ops/xpu_mla_sparse.py b/vllm/v1/attention/ops/xpu_mla_sparse.py index 9c1a92e10f57..1e7a9ec3fead 100644 --- a/vllm/v1/attention/ops/xpu_mla_sparse.py +++ b/vllm/v1/attention/ops/xpu_mla_sparse.py @@ -185,15 +185,20 @@ def triton_bf16_mla_sparse_interface( d_v: int = 512, block_dpe: int = 64, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - out : [num_tokens, num_heads_q, d_v] + """Out : [num_tokens, num_heads_q, d_v] max_logits : [num_tokens, num_heads_q] lse : logsumexp, [num_tokens, num_heads_q] Args: + q: Query tensor, [num_tokens, num_heads_q, dim_qk]. + kv: Latent KV cache tensor. + indices: Sparse KV indices selected for each query token. + sm_scale: Softmax scale applied to the QK product. + d_v: Size of the value dimension. block_dpe: Size of positional embedding portion of dim_qk. Set to 0 when q/kv contain only the nope latent (e.g. DSv4 prefill where RoPE is not split out). + """ num_tokens, num_heads_q, dim_qk = q.shape _, num_heads_kv, _ = kv.shape diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 286fdcb8dfd2..695916d81947 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -85,6 +85,7 @@ def get_attn_spec_kind( Returns: The ``KVCacheSpecKind`` the layer maps to. + """ from vllm.v1.kv_cache_interface import KVCacheSpecKind @@ -115,7 +116,6 @@ def get_attn_backend( has_sliding_window: bool = False, ) -> type[AttentionBackend]: """Selects which attention backend to use and lazily imports it.""" - if kv_cache_dtype is not None: valid_cache_dtypes = get_args(CacheDType) assert kv_cache_dtype in valid_cache_dtypes, ( diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index e2a06e7e4e16..b4c9dcaa2991 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -31,8 +31,7 @@ class BlockHashToBlockMap: - """ - Cache of blocks that are used for prefix caching. It caches blocks + """Cache of blocks that are used for prefix caching. It caches blocks from hash directly to a block or multiple blocks (i.e. {block_hash: KVCacheBlocks}) - Mostly block_hash maps to a single KVCacheBlock, and KVCacheBlocks @@ -59,9 +58,7 @@ def __init__(self): ] = {} def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None: - """ - Gets any block with the given block hash key. - """ + """Gets any block with the given block hash key.""" blocks = self._cache.get(key) if blocks is not None: if isinstance(blocks, KVCacheBlock): @@ -72,9 +69,7 @@ def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None: return None def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool: - """ - Checks whether the key maps to the given block ID. - """ + """Checks whether the key maps to the given block ID.""" blocks = self._cache.get(key) if blocks is None: return False @@ -86,9 +81,7 @@ def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool: return False def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None: - """ - Inserts the KVCacheBlock to the cache - """ + """Inserts the KVCacheBlock to the cache""" blocks = self._cache.get(key) if blocks is None: # When key is not found, attach a single block to the key @@ -104,9 +97,7 @@ def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None: self._unexpected_blocks_type(blocks) def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None: - """ - Checks if block_hash exists and pop block_id from the cache - """ + """Checks if block_hash exists and pop block_id from the cache""" blocks = self._cache.pop(key, None) if blocks is None: # block_hash not found in the cache @@ -158,6 +149,7 @@ class BlockPool: enable_kv_cache_events: Whether to enable kv cache events. metrics_collector: Optional metrics collector for tracking block residency. medium: Storage medium reported in KV cache events. + """ def __init__( @@ -214,6 +206,7 @@ def get_cached_block( Returns: The cached blocks if exists, or None. + """ cached_blocks = [] for group_id in kv_cache_group_ids: @@ -261,6 +254,7 @@ def cache_full_blocks( consults a subset of blocks (e.g. SWA tail-window), so blocks that can never serve a hit stay out of the prefix-cache hash map. + """ if num_cached_blocks >= num_full_blocks: return @@ -395,6 +389,7 @@ def emit_cached_block_events( num_cached_blocks: Number of blocks that were cache hits. block_size: Number of tokens per block. kv_cache_group_id: The KV cache group ID. + """ if not self.enable_kv_cache_events or num_cached_blocks == 0: return @@ -491,6 +486,7 @@ def cache_partial_block( Returns: The hash key with group ID if a partial entry can be registered; otherwise ``None`` for null blocks. + """ if block.is_null: return None @@ -671,6 +667,7 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: Returns: A list of new block. + """ if num_blocks > self.get_num_free_blocks(): raise ValueError(f"Cannot get {num_blocks} free blocks from the pool") @@ -724,8 +721,7 @@ def unpin_blocks( self.free_block_queue.append_n(released) def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: - """ - If a block is cached in `cached_block_hash_to_block`, we reset its hash + """If a block is cached in `cached_block_hash_to_block`, we reset its hash metadata and evict it from the cache. Args: @@ -733,6 +729,7 @@ def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: Returns: True if the block is evicted, False otherwise. + """ # Clean up metrics tracking first to prevent leaks if self.metrics_collector: @@ -753,6 +750,7 @@ def touch(self, blocks: Sequence[KVCacheBlock]) -> None: Args: blocks: A list of blocks to touch. + """ for block in blocks: # ref_cnt=0 means this block is in the free list (i.e. eviction @@ -774,6 +772,7 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: Args: ordered_blocks: A list of blocks to free ordered by their eviction priority. + """ # Identify blocks with hash (LRU cache) and without it (never match APC) blocks_to_evict_last = [] @@ -800,7 +799,7 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: pool.free_blocks(blocks) def evict_blocks(self, block_ids: set[int]) -> None: - """evict blocks from the prefix cache by their block IDs. + """Evict blocks from the prefix cache by their block IDs. only evicts blocks that are currently cached (have a hash). blocks with ref_cnt > 0 are not freed from the block pool, only evicted @@ -808,6 +807,7 @@ def evict_blocks(self, block_ids: set[int]) -> None: Args: block_ids: Set of block IDs to evict from cache. + """ for block_id in block_ids: assert block_id < len(self.blocks), ( @@ -826,6 +826,7 @@ def reset_prefix_cache(self) -> bool: Returns: bool: True if the prefix cache is successfully reset, False otherwise. + """ num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks() if num_used_blocks != 1: # The null block is always marked as used @@ -863,6 +864,7 @@ def get_num_free_blocks(self) -> int: Returns: The number of free blocks. + """ return self.free_block_queue.num_free_blocks @@ -871,8 +873,8 @@ def get_usage(self) -> float: Returns: The KV cache usage (between 0.0 and 1.0). - """ + """ # Subtract 1 to account for null block. total_gpu_blocks = self.num_gpu_blocks - 1 if not total_gpu_blocks: @@ -884,6 +886,7 @@ def take_events(self) -> list[KVCacheEvent]: Returns: A list of KV cache events. + """ if not self.enable_kv_cache_events: return [] diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index a3ed0c7aa15a..f069ce745318 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -64,6 +64,7 @@ class EncoderCacheManager: make space when needed. freed: List of mm_hash strings that were actually evicted since the last call to get_freed_mm_hashes(). This list is cleared on return. + """ @classmethod @@ -113,6 +114,7 @@ def check_and_update_cache(self, request: Request, input_id: int) -> bool: Returns: True if the encoder output for this input is already cached + """ mm_hash = request.mm_features[input_id].identifier # Not cached at all @@ -162,6 +164,7 @@ def can_allocate( Note: This method does not allocate physical memory for the encoder output but only the state of EncoderCacheManager. + """ num_embeds = request.get_num_encoder_embeds(input_id) @@ -198,8 +201,8 @@ def allocate(self, request: Request, input_id: int) -> None: Note: This method assumes can_allocate() returned True for the same input. - """ + """ mm_hash = request.mm_features[input_id].identifier request_id = request.request_id if mm_hash not in self.cached: @@ -222,7 +225,7 @@ def get_cached_input_ids(self, request: Request) -> set[int]: return self.request_cached_ids.get(request.request_id, set()) def free_encoder_input(self, request: Request, input_id: int) -> None: - """Free the request's reference to the encoder input (`mm_data`) + """Free the request's reference to the encoder input (`mm_data`). When the reference set for the corresponding `mm_hash` becomes empty, the entry is appended to `freeable` and `num_freeable_slots` is @@ -279,6 +282,7 @@ def get_freed_mm_hashes(self) -> list[str]: call to be used by the scheduler to notify workers about which encoder outputs can be removed from their caches. The internal list is cleared after this call. + """ # An entry evicted early in the scheduling pass can be allocated again # later in the same pass. Keep its worker-side tensor in that case. @@ -307,8 +311,8 @@ def compute_mm_encoder_budget( from the input sequence. - Space budget for encoder cache size, measured in number of tokens from the input sequence. - """ + """ if not mm_max_toks_per_item: logger.warning( "All non-text modalities supported by the model have been " diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index afb03ac3eac9..907b9690e022 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -64,9 +64,7 @@ def _validate_prefix_cache_retention_interval( class KVCacheCoordinator(ABC): - """ - Coordinate the KV cache of different KV cache groups. - """ + """Coordinate the KV cache of different KV cache groups.""" enable_partial_hash_hits = False @@ -179,8 +177,7 @@ def get_num_blocks_to_allocate( num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: - """ - Get the number of device blocks needed to be allocated for the request. + """Get the number of device blocks needed to be allocated for the request. Args: request_id: The request ID. @@ -203,6 +200,7 @@ def get_num_blocks_to_allocate( Returns: The number of blocks to allocate. + """ num_blocks_to_allocate = 0 for i, manager in enumerate(self.single_type_managers): @@ -237,8 +235,7 @@ def allocate_new_computed_blocks( num_local_computed_tokens: int, num_external_computed_tokens: int, ) -> None: - """ - Add the new computed blocks to the request. Optionally allocate new + """Add the new computed blocks to the request. Optionally allocate new blocks for external computed tokens (if any). Args: @@ -247,6 +244,7 @@ def allocate_new_computed_blocks( prefix cache. num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. + """ # A running request is already tracked in num_cached_block and won't # have new prefix-cache hits, so this is a no-op for it. @@ -283,8 +281,7 @@ def allocate_new_blocks( num_tokens_main_model: int, num_encoder_tokens: int = 0, ) -> tuple[list[KVCacheBlock], ...]: - """ - Allocate new blocks for the request to give it at least `num_tokens` + """Allocate new blocks for the request to give it at least `num_tokens` token slots. Args: @@ -299,6 +296,7 @@ def allocate_new_blocks( Returns: The new allocated blocks. + """ return tuple( manager.allocate_new_blocks( @@ -333,14 +331,14 @@ def get_replay_boundaries(self, request: Request) -> tuple[int, ...]: return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: - """ - Cache the blocks for the request. + """Cache the blocks for the request. Args: request: The request. num_computed_tokens: The total number of tokens that need to be cached (including tokens that are already cached). + """ boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: @@ -380,18 +378,17 @@ def reset_prefix_cache(self) -> bool: return all([pool.reset_prefix_cache() for pool in pools]) def free(self, request_id: str) -> None: - """ - Free the blocks for the request. + """Free the blocks for the request. Args: request_id: The request ID. + """ for manager in self.single_type_managers: manager.free(request_id) def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: - """ - Pop the request's bookkeeping from all single-type managers and + """Pop the request's bookkeeping from all single-type managers and return its blocks without returning them to the block pool. The caller must eventually pass the returned blocks to `block_pool.free_blocks`, freeing them in reverse order (so that @@ -402,6 +399,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: Returns: The request's blocks in allocation order. + """ blocks: list[KVCacheBlock] = [] for manager in self.single_type_managers: @@ -409,8 +407,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: return blocks def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - """ - Get the number of common prefix blocks for all requests with allocated + """Get the number of common prefix blocks for all requests with allocated KV cache for each kv cache group. Args: @@ -419,6 +416,7 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: Returns: list[int]: The number of common prefix blocks for each kv cache group. + """ return [ manager.get_num_common_prefix_blocks(running_request_id) @@ -431,8 +429,7 @@ def remove_skipped_blocks( processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: - """ - Remove the blocks that are no longer needed from `blocks` and replace + """Remove the blocks that are no longer needed from `blocks` and replace the removed blocks with null_block. Args: @@ -442,6 +439,7 @@ def remove_skipped_blocks( num_prompt_tokens: Optional prompt length. R-SWA managers use this to free gap blocks between the prefill tail and decode window; other manager types ignore it. + """ for manager in self.single_type_managers: manager.remove_skipped_blocks( @@ -449,9 +447,7 @@ def remove_skipped_blocks( ) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: - """ - Get the blocks for the request. - """ + """Get the blocks for the request.""" return tuple( manager.req_to_blocks.get(request_id) or [] for manager in self.single_type_managers @@ -465,7 +461,8 @@ def find_longest_cache_hit( ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]: """Returns the per-group hit blocks, the hit length, and the number of ``num_uncached_common_prefix_tokens`` (a shared prefix that a - sparse-retention group has not cached yet; 0 unless hybrid).""" + sparse-retention group has not cached yet; 0 unless hybrid). + """ pass def new_step_starts(self) -> None: @@ -475,8 +472,7 @@ def new_step_starts(self) -> None: class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): - """ - KV cache coordinator to use if prefix caching is disabled or unsupported. + """KV cache coordinator to use if prefix caching is disabled or unsupported. In contrast to UnitaryKVCacheCoordinator and HybridKVCacheCoordinator, supports arbitrary numbers of KV cache groups (including 0 groups). Does not implement any features related to prefix caching. @@ -527,8 +523,7 @@ def find_longest_cache_hit( class UnitaryKVCacheCoordinator(KVCacheCoordinator): - """ - KV cache coordinator for models with only one KV cache group. This is the + """KV cache coordinator for models with only one KV cache group. This is the case for models with only one KV cache type, e.g., all attention layers use full attention or all attention layers use sliding window attention. """ @@ -613,8 +608,7 @@ class SpecGroup(NamedTuple): class HybridKVCacheCoordinator(KVCacheCoordinator): - """ - KV cache coordinator for hybrid models with multiple KV cache types, and + """KV cache coordinator for hybrid models with multiple KV cache types, and thus multiple kv cache groups. """ @@ -733,8 +727,7 @@ def _cache_hit_alignment_tokens(self) -> int: ) def verify_and_split_kv_cache_groups(self) -> None: - """ - Groups KV cache groups by their spec type for efficient batch processing + """Groups KV cache groups by their spec type for efficient batch processing during cache hit lookup. """ self.attention_groups: list[SpecGroup] = [] @@ -843,8 +836,7 @@ def find_longest_cache_hit( block_hashes: list[BlockHash], max_cache_hit_length: int, ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]: - """ - Find the longest cache hit using an iterative fixed-point algorithm. + """Find the longest cache hit using an iterative fixed-point algorithm. Each attention type either accepts the current candidate length or reduces it. If any type reduces the length, restart checks over all @@ -861,8 +853,8 @@ def find_longest_cache_hit( - The number of tokens of the reconciled (combined) cache hit. - ``num_uncached_common_prefix_tokens``: a shared prefix that a sparse-retention group has not cached yet (0 unless hybrid). - """ + """ num_groups = len(self.kv_cache_config.kv_cache_groups) hit_length = max_cache_hit_length longest_hit_length = 0 @@ -985,8 +977,8 @@ def find_longest_cache_hit_per_group( Returns: (blocks_per_group, hit_lengths_per_group) - """ + """ num_groups = len(self.kv_cache_config.kv_cache_groups) hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] hit_lengths: list[int] = [0] * num_groups diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 226da3a1ed88..9f6997b53d3a 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -33,8 +33,7 @@ @dataclass class KVCacheBlocks: - """ - The allocation result of KVCacheManager, work as the interface between + """The allocation result of KVCacheManager, work as the interface between Scheduler and KVCacheManager, to hide KVCacheManager's internal data structure from the Scheduler. """ @@ -85,8 +84,7 @@ def get_block_ids( *, group_ids: Sequence[int] | None = None, ) -> tuple[list[int], ...] | None: - """ - Converts the KVCacheBlocks instance to block_ids. + """Converts the KVCacheBlocks instance to block_ids. Args: allow_none: Return None when every selected group is empty. @@ -97,6 +95,7 @@ def get_block_ids( - the outer tuple corresponds to KV cache groups - each inner list contains the block_ids of the blocks in that group + """ groups = ( self.blocks @@ -125,9 +124,7 @@ def get_unhashed_block_ids_all_groups(self) -> list[list[int]]: ] def new_empty(self) -> "KVCacheBlocks": - """ - Creates a new KVCacheBlocks instance with no blocks. - """ + """Creates a new KVCacheBlocks instance with no blocks.""" return KVCacheBlocks(tuple(() for _ in range(len(self.blocks)))) @@ -232,6 +229,7 @@ def usage(self) -> float: Returns: The KV cache usage (between 0.0 and 1.0). + """ return self.block_pool.get_usage() @@ -240,6 +238,7 @@ def make_prefix_cache_stats(self) -> PrefixCacheStats | None: Returns: The current prefix caching stats, or None if logging is disabled. + """ if not self.log_stats: return None @@ -278,6 +277,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int window) has not cached yet (Marconi-style APC), or 0 if none. Pinned so sparse prefix-cache retention does not drop the junction and defeat cross-request reuse. + """ # We skip finding the prefix cache hit when prefix caching is # disabled or the request is marked as skipping kv cache read @@ -340,6 +340,7 @@ def get_computed_blocks_for_connector( Returns: The ``get_computed_blocks`` triple (blocks, number of local computed tokens, shared-prefix boundary) plus ``hit_diverged``. + """ coordinator = self.coordinator if not ( @@ -462,6 +463,7 @@ def allocate_slots( Returns: A list of new allocated blocks. + """ # When loading KV data asynchronously, we may have zero new tokens to # compute while still allocating slots for externally computed tokens. @@ -607,6 +609,7 @@ def free(self, request: Request) -> None: Args: request: The request to free the blocks. + """ self.coordinator.free(request.request_id) @@ -624,6 +627,7 @@ def remove_skipped_blocks( processed_computed_tokens: Computed-token prefix length covering fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for R-SWA gap eviction. + """ self.coordinator.remove_skipped_blocks( request_id, processed_computed_tokens, num_prompt_tokens @@ -639,14 +643,16 @@ def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: Returns: The request's blocks in allocation order. + """ return self.coordinator.pop_blocks_for_free(request.request_id) def evict_blocks(self, block_ids: set[int]) -> None: - """evict blocks from the prefix cache by their block IDs. + """Evict blocks from the prefix cache by their block IDs. Args: block_ids: Set of block IDs to evict from cache. + """ self.block_pool.evict_blocks(block_ids) @@ -658,6 +664,7 @@ def reset_prefix_cache(self) -> bool: Returns: bool: True if the prefix cache is successfully reset, False otherwise. + """ if not self.coordinator.reset_prefix_cache(): return False @@ -697,6 +704,7 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: Returns: list[int]: The number of common prefix blocks for each kv cache group. + """ return self.coordinator.get_num_common_prefix_blocks(running_request_id) @@ -705,6 +713,7 @@ def take_events(self) -> list[KVCacheEvent]: Returns: A list of KV cache events. + """ events = self.block_pool.take_events() for event in events: @@ -792,6 +801,7 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: request: The request to cache the blocks. num_computed_tokens: The number of computed tokens, including tokens that are already cached and tokens to be cached. + """ if self.enable_caching: self.coordinator.cache_blocks(request, num_computed_tokens) @@ -846,7 +856,8 @@ def get_zeroing_block_ids_in_range( self, request_id: str, start_token: int, end_token: int ) -> list[int]: """The request's block ids covering [start_token, end_token), from - the groups whose new blocks are zeroed by the worker.""" + the groups whose new blocks are zeroed by the worker. + """ ids: list[int] = [] for mgr in self.coordinator.single_type_managers: if mgr.records_new_block_ids: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index c55c06abf1cc..656765c89604 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -263,6 +263,7 @@ class FreeKVCacheBlockQueue: Args: blocks: A list of KVCacheBlock objects. + """ def __init__(self, blocks: list[KVCacheBlock]) -> None: @@ -300,6 +301,7 @@ def popleft(self) -> KVCacheBlock: Returns: The first free block. + """ if ( self.fake_free_list_head.next_free_block is self.fake_free_list_tail @@ -340,6 +342,7 @@ def popleft_n(self, n: int) -> list[KVCacheBlock]: Returns: A list of n free blocks. + """ if n == 0: return [] @@ -370,6 +373,7 @@ def remove(self, block: KVCacheBlock) -> None: Args: block: The block to remove. + """ if block.prev_free_block is None or block.next_free_block is None: # This should not happen if the block is from the free list. @@ -391,6 +395,7 @@ def append(self, block: KVCacheBlock) -> None: Args: block: The block to append. + """ if self.fake_free_list_tail.prev_free_block is None: raise RuntimeError( @@ -434,6 +439,7 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None: Args: blocks: The blocks to append. + """ if len(blocks) == 0: return @@ -459,6 +465,7 @@ def get_all_free_blocks(self) -> list[KVCacheBlock]: Returns: A list of free blocks. + """ ret = [] if self.fake_free_list_head.next_free_block is None: @@ -505,6 +512,7 @@ def _gen_mm_extra_hash_keys( Returns: A tuple of extra keys and the next multi-modal index. + """ extra_keys: list[Any] = [] @@ -565,6 +573,7 @@ def _gen_lora_extra_hash_keys(request: Request) -> list[str]: Returns: Return LoRA name of the request if it is a LoRA request. Return empty list otherwise. + """ if not request.lora_request: return [] @@ -584,6 +593,7 @@ def _gen_prompt_embeds_extra_hash_keys( Returns: Return a stable hash of the block prompt embeddings if prompt embeds are present. Return empty list otherwise. + """ if request.prompt_embeds is None: return [] @@ -612,6 +622,7 @@ def generate_block_hash_extra_keys( Returns: A tuple of extra keys and the next multi-modal index. + """ mm_extra_keys: list[Any] mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys( @@ -645,6 +656,7 @@ def hash_block_tokens( the contents of the preceding block(s). The hash value is used for prefix caching. We use LRU cache for this function to avoid recomputing hash values for the same block contents. + Args: hash_function: The hash function used to compute block hash. parent_block_hash: The hash of the parent block. None @@ -652,9 +664,11 @@ def hash_block_tokens( curr_block_token_ids: A list of token ids in the current block. The current block is assumed to be full. extra_keys: Extra keys for the block. + Returns: The hash value of the block and the token ids in the block. The entire tuple is used as the hash key of the block. + """ if not parent_block_hash: parent_block_hash = NONE_HASH @@ -813,8 +827,7 @@ def get_request_block_hasher( hash_block_size: int, caching_hash_fn: Callable[[Any], bytes], ) -> Callable[[Request], list[BlockHash]]: - """ - Returns a function which computes the list of un-computed block hashes + """Returns a function which computes the list of un-computed block hashes of a request. Hashes are computed at ``hash_block_size`` granularity and chained over the @@ -912,9 +925,7 @@ def _check_enough_kv_cache_memory( def max_memory_usage_bytes( vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] ) -> int: - """ - Get the maximum memory usage in bytes for the given KV cache specs. - """ + """Get the maximum memory usage in bytes for the given KV cache specs.""" return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) @@ -923,8 +934,7 @@ def estimate_max_model_len( kv_cache_spec: dict[str, KVCacheSpec], available_memory: int, ) -> int: - """ - Estimates the maximum model length that can fit in the available memory + """Estimates the maximum model length that can fit in the available memory using binary search. This function temporarily modifies max_model_len during estimation but @@ -937,6 +947,7 @@ def estimate_max_model_len( Returns: The estimated maximum model length that can fit in the available memory. + """ # Save the original max_model_len to restore after estimation original_max_model_len = vllm_config.model_config.max_model_len @@ -977,8 +988,7 @@ def check_enough_kv_cache_memory( kv_cache_spec: dict[str, KVCacheSpec], available_memory: int, ): - """ - Checks whether `available_memory` is enough for the KV cache to hold at + """Checks whether `available_memory` is enough for the KV cache to hold at least one request with the model's max_model_len. Args: @@ -988,8 +998,8 @@ def check_enough_kv_cache_memory( Raises: ValueError: If there is not enough memory available for the KV cache. - """ + """ # No need to check for available memory if the kv_cache_spec is empty if kv_cache_spec: # Reserve the null block BlockPool permanently holds back, so the check @@ -1012,8 +1022,7 @@ def check_enough_kv_cache_memory( def create_kv_cache_group_specs( kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] ) -> list[KVCacheGroupSpec]: - """ - Create KVCacheGroupSpec object for each kv cache group layer. + """Create KVCacheGroupSpec object for each kv cache group layer. The layers in the same group should share the same KVCacheSpec. @@ -1024,8 +1033,10 @@ def create_kv_cache_group_specs( A list of kv cache groups, where each element is a list of layer names that belong to the same group and should share the same KVCacheSpec. + Returns: A list of KVCacheGroupSpec objects, one for each group. + """ kv_cache_groups = [] for layer_names_one_group in grouped_layer_names: @@ -1040,8 +1051,7 @@ def create_kv_cache_group_specs( def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - """ - Whether all layers in the given KVCacheSpec have the same KV cache spec. + """Whether all layers in the given KVCacheSpec have the same KV cache spec. Note that we regard FullAttentionSpec with and without sliding window as the same type. @@ -1050,8 +1060,8 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: Returns: True if all layers have the same type, False otherwise. - """ + """ if not kv_cache_spec: # Encoder-only models do not have KV cache, kv_cache_type can be # regarded as uniform. @@ -1067,8 +1077,7 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: - """ - Get the maximum concurrency for the given KV cache configuration. + """Get the maximum concurrency for the given KV cache configuration. A request at max_model_len consumes whole blocks from each group's block table — cdiv(per-request bytes, page bytes) of the group's spec — and all @@ -1101,8 +1110,7 @@ def get_max_concurrency_for_kv_cache_config( def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: - """ - Override the number of kv cache blocks if `num_gpu_blocks_override` is set. + """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. The override is logged once, at the call site in `get_kv_cache_configs`. """ if vllm_config.cache_config.num_gpu_blocks_override is not None: @@ -1111,8 +1119,7 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: - """ - Bytes consumed by one block in the worker's shared KV cache pool, mirroring + """Bytes consumed by one block in the worker's shared KV cache pool, mirroring the divisor used by `get_kv_cache_config_from_groups` to convert `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. @@ -1121,9 +1128,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: - """ - Get the page size of the KV cache. - """ + """Get the page size of the KV cache.""" page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} assert len(page_sizes) == 1 return page_sizes.pop() @@ -1132,8 +1137,7 @@ def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: def _get_kv_cache_groups_uniform_spec( kv_cache_specs: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: - """ - Generates the KV cache configuration for a model with the same KV cache + """Generates the KV cache configuration for a model with the same KV cache spec for all layers. Args: @@ -1141,16 +1145,15 @@ def _get_kv_cache_groups_uniform_spec( Returns: The generated KVCacheGroupSpecs - """ + """ return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) def _get_kv_cache_groups_uniform_type( spec: UniformTypeKVCacheSpecs, ) -> list[KVCacheGroupSpec]: - """ - Generates the KV cache configuration for a model with one type of KV cache + """Generates the KV cache configuration for a model with one type of KV cache but different hidden sizes. All layers are merged into one group. Args: @@ -1158,8 +1161,8 @@ def _get_kv_cache_groups_uniform_type( Returns: The generated KVCacheGroupSpecs - """ + """ return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] @@ -1365,8 +1368,7 @@ def _glm5_next_tensor_layout( def unify_kv_cache_spec_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> dict[str, KVCacheSpec]: - """ - Unify the page size of the given KVCacheSpec. If the page size of all layers + """Unify the page size of the given KVCacheSpec. If the page size of all layers are the same, return the original KVCacheSpec. If not same, unify the page size by increasing the block size of layers with smaller page size. Two cases cannot be unified by block size alone and pad their physical page to @@ -1385,6 +1387,7 @@ def unify_kv_cache_spec_page_size( Returns: The updated KVCacheSpec with the same page_size_bytes. + """ page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} if len(page_sizes) <= 1: @@ -1435,8 +1438,7 @@ def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bo def _get_kv_cache_groups_uniform_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: - """ - Generates the KV cache groups for hybrid models with multiple + """Generates the KV cache groups for hybrid models with multiple attention types but still with a uniform page size (physical memory per block per layer) for all layers. @@ -1496,6 +1498,7 @@ def _get_kv_cache_groups_uniform_page_size( kv_cache_spec: The KVCacheSpec of each attention layer in the model Returns: The generated KVCacheGroupSpecs + """ # Group all layers by kv_cache_spec. # E.g., 2 full attention layers and 3 sliding window attention layers, @@ -1659,8 +1662,7 @@ def get_kv_cache_config_from_groups( kv_cache_groups: list[KVCacheGroupSpec], available_memory: int, ) -> KVCacheConfig: - """ - Generate the KV cache configuration from the KV cache groups and spec + """Generate the KV cache configuration from the KV cache groups and spec of each layer. Args: @@ -1669,6 +1671,7 @@ def get_kv_cache_config_from_groups( available_memory: Memory available for KV cache in bytes Returns: The generated KVCacheConfig + """ if len(kv_cache_groups) == 0: # Attention free models do not have KV cache. @@ -1926,15 +1929,14 @@ def _try_get_full_allocation_fallback_groups( def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): - """ - This function tries to convert the KV cache specs to one type if the model + """This function tries to convert the KV cache specs to one type if the model is a hybrid model with multiple type of KV cache. It will convert all SlidingWindowSpec to FullAttentionSpec if both types are present. Args: kv_cache_spec: The kv cache spec of each attention layer in the model - """ + """ if is_kv_cache_spec_uniform( kv_cache_spec ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec): @@ -2166,6 +2168,7 @@ def _annotate_eagle_groups( kv_cache_groups: Groups to annotate in place. use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed group. + """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle_block_drop(): @@ -2196,6 +2199,7 @@ def _warn_if_unannotated_eagle_mamba( Args: vllm_config: Config supplying the speculative method, if any. kv_cache_groups: Groups as they will be handed to consumers. + """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle(): @@ -2262,8 +2266,7 @@ def get_kv_cache_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: - """ - Split the layers in the model into groups with the same KV cache spec. + """Split the layers in the model into groups with the same KV cache spec. Args: vllm_config: The global VllmConfig @@ -2271,6 +2274,7 @@ def get_kv_cache_groups( Returns: The generated KVCacheGroups + """ if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: unify_hybrid_kv_cache_specs(kv_cache_spec) @@ -2356,9 +2360,7 @@ def get_kv_cache_groups( def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: - """ - Generate the KV cache configuration for the scheduler. - """ + """Generate the KV cache configuration for the scheduler.""" assert all( [cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs] ) @@ -2382,9 +2384,7 @@ def generate_scheduler_kv_cache_config( def get_kv_cache_capacity( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> tuple[int, float]: - """ - Get the group-aware KV cache token capacity and max concurrency. - """ + """Get the group-aware KV cache token capacity and max concurrency.""" max_model_len = vllm_config.model_config.max_model_len max_concurrency = get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config @@ -2415,8 +2415,7 @@ def _max_memory_usage_bytes_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> int: - """ - Calculate maximum memory usage in bytes from KV cache groups. + """Calculate maximum memory usage in bytes from KV cache groups. This correctly accounts for padding in hybrid models. For example, if a model has 8 full attention layers and 9 sliding window layers, they will @@ -2475,8 +2474,7 @@ def _estimate_max_model_len_from_groups( kv_cache_groups: list[KVCacheGroupSpec], available_memory: int, ) -> int: - """ - Binary search for the maximum model length that fits in available memory. + """Binary search for the maximum model length that fits in available memory. Returns 0 if even 1 token doesn't fit. """ original_max = vllm_config.model_config.max_model_len @@ -2522,8 +2520,7 @@ def _auto_fit_max_model_len( projected_groups_per_worker: list[list[KVCacheGroupSpec]], available_memory: list[int], ) -> None: - """ - When max_model_len is set to -1, this function estimates the largest + """When max_model_len is set to -1, this function estimates the largest context length that can be supported with the available GPU memory. It uses binary search to find the maximum length that fits across all workers. @@ -2533,6 +2530,7 @@ def _auto_fit_max_model_len( projected_groups_per_worker: KV cache groups projected to each worker. available_memory: Memory available for KV cache in bytes for each worker. + """ original_max = vllm_config.model_config.max_model_len @@ -2585,8 +2583,7 @@ def _project_kv_cache_groups_to_worker( global_kv_cache_groups: list[KVCacheGroupSpec], worker_spec: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: - """ - Projects global KV cache groups onto a single worker's assigned layers. + """Projects global KV cache groups onto a single worker's assigned layers. In pipeline parallelism, each worker only owns a subset of layers. This function filters the global groups to include only layers present on the @@ -2598,6 +2595,7 @@ def _project_kv_cache_groups_to_worker( Returns: The projected KV cache groups containing only this worker's layers. + """ projected_groups: list[KVCacheGroupSpec] = [] for group in global_kv_cache_groups: @@ -2628,8 +2626,7 @@ def get_kv_cache_configs( kv_cache_specs: list[dict[str, KVCacheSpec]], available_memory: list[int], ) -> list[KVCacheConfig]: - """ - Generates the KV cache configurations for a model. + """Generates the KV cache configurations for a model. Since we use a shared centralized controller for all workers, we need the `kv_cache_config` to be consistent across all workers to make sure the KV cache allocation can be applied to all workers. However, different @@ -2656,8 +2653,8 @@ def get_kv_cache_configs( Returns: The generated KVCacheConfigs for each worker. - """ + """ # Merge the KV cache specs of all workers. Different PP stages may have # different layer names, and different TP ranks of the same PP stage should # have the same KV cache spec. @@ -2778,8 +2775,7 @@ def get_kv_cache_configs( class BlockHashListWithBlockSize: - """ - Convert block-hash granularity from `hash_block_size` to `target_block_size`. + """Convert block-hash granularity from `hash_block_size` to `target_block_size`. Used when KV cache groups have different block sizes: `hash_block_size` is the size used to compute the original `block_hashes`; `target_block_size` is the group's actual block size. @@ -2809,6 +2805,7 @@ class BlockHashListWithBlockSize: block_hashes: Block hashes to convert, computed at `hash_block_size`. hash_block_size: Block size at which `block_hashes` were computed. target_block_size: Desired block size; must be a multiple of `hash_block_size`. + """ def __init__( diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index c15296dd0516..d5a059bf4c4c 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -79,6 +79,7 @@ def schedule(self, throttle_prefills: bool = False) -> "SchedulerOutput": Returns: A SchedulerOutput object containing information about the scheduled requests. + """ raise NotImplementedError @@ -105,6 +106,7 @@ def update_from_output( Returns: A dict of client index to EngineCoreOutputs object containing the outputs for each request originating from that client. + """ raise NotImplementedError @@ -115,6 +117,7 @@ def update_draft_token_ids(self, draft_token_ids: "DraftTokenIds") -> None: Args: draft_token_ids: The input draft token ids for each request. + """ raise NotImplementedError @@ -129,6 +132,7 @@ def update_draft_token_ids_in_output( draft_token_ids: The input draft token ids for each request. scheduler_output: Update the given scheduler_output with the corresponding draft token ids. + """ raise NotImplementedError @@ -138,6 +142,7 @@ def add_request(self, request: "Request") -> None: Args: request: The new request being added. + """ raise NotImplementedError @@ -162,6 +167,7 @@ def finish_requests( Returns: List of requests that were aborted. Will not include any that were already finished. + """ raise NotImplementedError @@ -172,7 +178,8 @@ def get_num_unfinished_requests(self) -> int: def has_unfinished_requests(self) -> bool: """Returns True if there are unfinished requests in the scheduler's - internal queue.""" + internal queue. + """ return self.get_num_unfinished_requests() > 0 @abstractmethod @@ -192,7 +199,8 @@ def has_finished_requests(self) -> bool: def has_requests(self) -> bool: """Returns True if there are unfinished requests, or finished requests - not yet returned in SchedulerOutputs.""" + not yet returned in SchedulerOutputs. + """ return self.has_unfinished_requests() or self.has_finished_requests() @property @@ -218,6 +226,8 @@ def reset_prefix_cache( preempted and moved to the waiting queue. Otherwise, this method will only reset the KV prefix cache when there is no running request taking KV cache. + reset_connector: If True, also reset any KV connector state. + """ raise NotImplementedError diff --git a/vllm/v1/core/sched/request_queue.py b/vllm/v1/core/sched/request_queue.py index 38c7db94a45e..ab94e9075294 100644 --- a/vllm/v1/core/sched/request_queue.py +++ b/vllm/v1/core/sched/request_queue.py @@ -43,7 +43,8 @@ def prepend_request(self, request: Request) -> None: @abstractmethod def prepend_requests(self, requests: "RequestQueue") -> None: """Prepend all requests from another queue to the front of this - queue.""" + queue. + """ pass @abstractmethod @@ -129,8 +130,7 @@ def __iter__(self) -> Iterator[Request]: class PriorityRequestQueue(RequestQueue): - """ - A priority queue that supports heap operations. + """A priority queue that supports heap operations. Respects the ordering defined in the Request class, where requests with a smaller value of `priority` are processed first. @@ -161,14 +161,16 @@ def prepend_request(self, request: Request) -> None: """Add a request to the queue according to priority policy. Note: In a priority queue, there is no concept of prepending to the - front. Requests are ordered by (priority, arrival_time).""" + front. Requests are ordered by (priority, arrival_time). + """ self.add_request(request) def prepend_requests(self, requests: RequestQueue) -> None: """Add all requests from another queue according to priority policy. Note: In a priority queue, there is no concept of prepending to the - front. Requests are ordered by (priority, arrival_time).""" + front. Requests are ordered by (priority, arrival_time). + """ for request in requests: self.add_request(request) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b662bcc10b44..08f11911116d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1558,12 +1558,10 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: def _update_request_as_session( self, session: Request, update: StreamingUpdate ) -> None: - """ - Updates the waiting session with the next streaming update. + """Updates the waiting session with the next streaming update. Discards the last sampled output token from the prior input chunk. """ - # Current streaming input behaviour: Keep only computed output tokens # (discard final sampled output token). num_computed_tokens = session.num_computed_tokens @@ -1665,8 +1663,7 @@ def _try_schedule_encoder_inputs( encoder_compute_budget: int, shift_computed_tokens: int = 0, ) -> tuple[list[int], int, int, list[int]]: - """ - Determine which encoder inputs need to be scheduled in the current step, + """Determine which encoder inputs need to be scheduled in the current step, and update `num_new_tokens` and encoder token budget accordingly. An encoder input will be scheduled if: @@ -2507,6 +2504,7 @@ def finish_requests( Returns: List of requests that were aborted. Will not include any that were already finished. + """ assert RequestStatus.is_finished(finished_status) if isinstance(request_ids, str): @@ -2841,8 +2839,7 @@ def get_kv_event_publisher_config(self) -> KVEventsConfig | None: def _connector_finished( self, request: Request ) -> tuple[bool, dict[str, Any] | None]: - """ - Invoke the KV connector request_finished() method if applicable. + """Invoke the KV connector request_finished() method if applicable. Returns optional kv transfer parameters to be included with the request outputs. @@ -2910,14 +2907,12 @@ def _request_remaining_blocks(self, request: Request) -> int: def _inflight_prefill_reserved_blocks(self) -> int: """Num blocks in-flight prefills still need to finish (their reservation).""" - return sum( self._request_remaining_blocks(req) for req in self._inflight_prefills ) def _update_waiting_for_remote_kv(self, request: Request) -> None: - """ - KV Connector: update request state after async recv is finished. + """KV Connector: update request state after async recv is finished. When the kv transfer is ready, we cache the blocks and the request state will be moved back to WAITING from @@ -2959,9 +2954,7 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: self.finished_recving_kv_req_ids.remove(request.request_id) def _try_promote_blocked_waiting_request(self, request: Request) -> bool: - """ - Try to promote a blocked waiting request back to schedulable states. - """ + """Try to promote a blocked waiting request back to schedulable states.""" if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: # finished_recving_kv_req_ids is populated during # update_from_output(), based on worker-side connector signals @@ -2995,8 +2988,7 @@ def _try_promote_blocked_waiting_request(self, request: Request) -> bool: ) def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): - """ - KV Connector: update the scheduler state based on the output. + """KV Connector: update the scheduler state based on the output. The Worker side connectors add finished_recving and finished_sending reqs to the output. @@ -3004,7 +2996,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # if finished_recving: add to state so we can schedule the request during the next step. """ - if self.connector is not None: self.connector.update_connector_output(kv_connector_output) @@ -3030,8 +3021,7 @@ def _update_requests_with_invalid_blocks( num_scheduled_tokens: dict[str, int], evict_blocks: bool = True, ) -> tuple[set[str], int, set[int]]: - """ - Identify and update requests affected by invalid KV cache blocks. + """Identify and update requests affected by invalid KV cache blocks. This method scans the given requests, detects those with invalid blocks and adjusts their `num_computed_tokens` to the longest valid prefix. @@ -3053,6 +3043,7 @@ def _update_requests_with_invalid_blocks( be recomputed across all affected requests. - blocks_to_evict (set[int]): Block IDs to evict from cache, including invalid blocks and downstream dependent blocks. + """ affected_req_ids: set[str] = set() total_affected_tokens = 0 @@ -3145,11 +3136,11 @@ def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]: def _handle_invalid_blocks( self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] ) -> set[str]: - """ - Handle requests affected by invalid KV cache blocks. + """Handle requests affected by invalid KV cache blocks. Returns: Set of affected request IDs to skip in update_from_output main loop. + """ if len(self.kv_cache_config.kv_cache_groups) > 1: # Block IDs are only unique within a group, so a flat set of diff --git a/vllm/v1/core/sched/utils.py b/vllm/v1/core/sched/utils.py index 31fa840e7783..045dacaa9829 100644 --- a/vllm/v1/core/sched/utils.py +++ b/vllm/v1/core/sched/utils.py @@ -30,11 +30,14 @@ def check_sequence_repetition( params: RepetitionDetectionParams, ) -> bool: """Check if a sequence of token IDs has a repetition pattern. + Args: token_ids: List of token IDs params: Repetition detection parameters. + Returns: True if a repetition pattern is found, False otherwise. + """ max_pattern_size = params.max_pattern_size min_pattern_size = params.min_pattern_size @@ -77,6 +80,7 @@ def remove_all(lst: list, items_to_remove: set) -> list: Note: For single item removal, this modifies the original list in-place and returns it. For multiple items, it creates and returns a new list. + """ if not items_to_remove: return lst diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 08eab7628768..9ad0a38ed550 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -49,8 +49,7 @@ class SingleTypeKVCacheManager(ABC): - """ - An abstract base class for a manager that handle the kv cache management + """An abstract base class for a manager that handle the kv cache management logic of one specific type of attention layer. """ @@ -81,14 +80,17 @@ def __init__( needs_kv_cache_zeroing: bool = False, max_admission_blocks_per_request: int | None = None, ) -> None: - """ - Initializes the SingleTypeKVCacheManager. + """Initializes the SingleTypeKVCacheManager. + Args: kv_cache_spec: The kv_cache_spec for this manager. block_pool: The block pool. + enable_caching: Whether prefix caching is enabled. kv_cache_group_id: The id of the kv cache group of this manager. scheduler_block_size: The scheduling granularity (LCM of all group block sizes); a multiple of this manager's ``block_size``. + dcp_world_size: Decode context parallel world size. + pcp_world_size: Prefill context parallel world size. needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs newly allocated block IDs from this manager. max_admission_blocks_per_request: Recycling-aware per-request @@ -97,6 +99,7 @@ def __init__( chunked-local); `None` (the default) means no cap, which is correct for full-attention-style specs that hold every block until the request finishes. + """ self.scheduler_block_size = scheduler_block_size # Hybrid fine-grained lookup may lower this after all participating @@ -184,8 +187,7 @@ def get_num_blocks_to_allocate( num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: - """ - Get the number of blocks needed to be allocated for the request. + """Get the number of blocks needed to be allocated for the request. Args: request_id: The request ID. @@ -206,8 +208,8 @@ def get_num_blocks_to_allocate( Returns: The number of blocks to allocate. - """ + """ num_required_blocks = cdiv(num_tokens, self.block_size) if apply_admission_cap and self._max_admission_blocks_per_request is not None: # Recycling-aware specs (SWA, chunked-local) cap the per-request @@ -269,8 +271,7 @@ def add_local_computed_blocks( num_local_computed_tokens: int, num_external_computed_tokens: int, ) -> None: - """ - Add the locally cached (prefix-hit) blocks to the request: + """Add the locally cached (prefix-hit) blocks to the request: 1. Touch the computed blocks (paired with adding them to `req_blocks`) so their ref_cnt exactly tracks the referencing requests. 1.5. (Optional) For sliding window, skipped blocks are padded with nulls. @@ -282,6 +283,7 @@ def add_local_computed_blocks( prefix cache. num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. + """ # The coordinator only calls this for first-time allocations (running # requests are short-circuited there), so the request has no blocks yet. @@ -327,8 +329,7 @@ def allocate_external_computed_blocks( num_local_computed_tokens: int, num_external_computed_tokens: int, ) -> None: - """ - Allocate new blocks for external (KV-connector) computed tokens. + """Allocate new blocks for external (KV-connector) computed tokens. Must run only after every group's local blocks have been touched via `add_local_computed_blocks`, so this group's `get_new_blocks` cannot @@ -338,6 +339,7 @@ def allocate_external_computed_blocks( request_id: The request ID. num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. + """ num_total_computed_tokens = ( num_local_computed_tokens + num_external_computed_tokens @@ -364,8 +366,7 @@ def allocate_external_computed_blocks( def allocate_new_blocks( self, request_id: str, num_tokens: int, num_tokens_main_model: int ) -> list[KVCacheBlock]: - """ - Allocate new blocks for the request to give it at least `num_tokens` + """Allocate new blocks for the request to give it at least `num_tokens` token slots. Args: @@ -375,8 +376,10 @@ def allocate_new_blocks( num_tokens_main_model: The number of tokens for the main model (aka target model in spec decode). w/o spec decode, it is num_tokens; with spec decode, it is num_tokens - num_lookahead_tokens. + Returns: The new allocated blocks. + """ cow_blocks: list[KVCacheBlock] = [] if request_id in self._partial_hit_reqs: @@ -475,8 +478,7 @@ def cache_blocks( *, replay_boundaries: Sequence[int], ) -> None: - """ - Cache the blocks for the request. + """Cache the blocks for the request. Args: request: The request. @@ -488,6 +490,7 @@ def cache_blocks( a tail once per that-sized segment. Only SWA acts on it. replay_boundaries: Positions a later request replaying this prompt can resume at, from ``get_replay_boundaries``. + """ num_cached_blocks = self.num_cached_block.get(request.request_id, 0) num_full_blocks = num_tokens // self.block_size @@ -552,8 +555,7 @@ def reachable_block_mask( return None def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: - """ - Pop the request's bookkeeping and return its blocks without yet + """Pop the request's bookkeeping and return its blocks without yet returning them to the block pool. The caller is responsible for eventually passing the returned blocks to `block_pool.free_blocks`, freeing them in reverse order (so that tail blocks are evicted first). @@ -563,6 +565,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: Returns: The request's blocks in allocation order. + """ # Default to [] in case a request is freed (aborted) before alloc. req_blocks = self.req_to_blocks.pop(request_id, []) @@ -571,19 +574,18 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: return req_blocks def free(self, request_id: str) -> None: - """ - Free the blocks for the request. + """Free the blocks for the request. Args: request_id: The request ID. + """ # Free blocks in reverse order so that the tail blocks are freed first. self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) @abstractmethod def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """ - Get the number of common prefix blocks for all requests with allocated + """Get the number of common prefix blocks for all requests with allocated KV cache. Args: @@ -592,8 +594,8 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: Returns: The number of common prefix blocks for all requests with allocated KV cache. - """ + """ raise NotImplementedError @classmethod @@ -610,8 +612,7 @@ def find_longest_cache_hit( dcp_world_size: int = 1, pcp_world_size: int = 1, ) -> tuple[tuple[list[KVCacheBlock], ...], int]: - """ - Get the longest cache hit prefix of the blocks that is not longer than + """Get the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix should be a common prefix hit for all the kv cache groups in `kv_cache_group_ids`. If no cache hit is found, return an empty list. @@ -647,8 +648,8 @@ def find_longest_cache_hit( For example, sliding window manager should return a list like ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4 and sliding window 8 and len(kv_cache_group_ids) = 1. - """ + """ raise NotImplementedError def _remove_blocks_in_range( @@ -684,8 +685,8 @@ def remove_skipped_blocks( processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: - """ - Remove and free the blocks that are no longer needed for attention computation. + """Remove and free blocks no longer needed for attention computation. + The removed blocks should be replaced by null_block. This function depends on `get_num_skipped_tokens`, which need to be implemented @@ -698,6 +699,7 @@ def remove_skipped_blocks( num_prompt_tokens: Optional prompt length for attention types (e.g. R-SWA) that evict a middle gap rather than a head prefix. Ignored by the default implementation. + """ del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. @@ -718,14 +720,14 @@ def remove_skipped_blocks( self._remove_blocks_in_range(request_id, 0, num_skipped_blocks) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: - """ - Get the number of tokens that will be skipped for attention computation. + """Get the number of tokens that will be skipped for attention computation. Args: num_computed_tokens: The number of tokens that have been computed. Returns: The number of tokens that will be skipped for attention computation. + """ # The default behavior is to not skip any tokens. return 0 @@ -1144,8 +1146,7 @@ def reachable_block_mask( return mask def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: - """ - Get the number of tokens that will be skipped for attention computation. + """Get the number of tokens that will be skipped for attention computation. For sliding window, this corresponds to the tokens that are prior to the current sliding window. @@ -1174,6 +1175,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: Returns: The number of tokens that will be skipped for attention computation. + """ return max( 0, @@ -1181,8 +1183,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: ) def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """ - NOTE(Chen): The prefix blocks are null blocks for sliding window layers. + """NOTE(Chen): The prefix blocks are null blocks for sliding window layers. So it's not correct to count ref_cnt like FullAttentionManager. Return 0 here for correctness. Need to support cascade attention + sliding window in the future. @@ -1301,8 +1302,7 @@ def find_longest_cache_hit( dcp_world_size: int = 1, pcp_world_size: int = 1, ) -> tuple[tuple[list[KVCacheBlock], ...], int]: - """ - For chunked local attention, we need to find the longest cache hit + """For chunked local attention, we need to find the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix should be a common prefix hit for all the kv cache groups in `kv_cache_group_ids`. If no cache hit is found, return an empty list. @@ -1335,6 +1335,7 @@ def find_longest_cache_hit( Returns: A list of cached blocks + """ assert isinstance(kv_cache_spec, ChunkedLocalAttentionSpec), ( "ChunkedLocalAttentionManager can only be used for " @@ -1389,8 +1390,7 @@ def find_longest_cache_hit( return computed_blocks, hit_length def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: - """ - Get the number of tokens that will be skipped for attention computation. + """Get the number of tokens that will be skipped for attention computation. For chunked local attention, this corresponds to the tokens that are on the left side of the current chunk. @@ -1428,6 +1428,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: Returns: The number of tokens that will be skipped for attention computation. + """ num_skipped_tokens = ( num_computed_tokens // self.attention_chunk_size @@ -1435,9 +1436,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: return num_skipped_tokens def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """ - cascade attention is not supported by chunked local attention. - """ + """Cascade attention is not supported by chunked local attention.""" return 0 @@ -1674,9 +1673,7 @@ def remove_skipped_blocks( blocks[last_state_block_idx] = self._null_block def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """ - cascade attention is not supported by mamba - """ + """Cascade attention is not supported by mamba""" return 0 def _needs_internal_checkpoint( @@ -1979,8 +1976,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: return super().pop_blocks_for_free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: - """ - Get the number of tokens whose mamba state are not needed anymore. Mamba only + """Get the number of tokens whose mamba state are not needed anymore. Mamba only need to keep the state of the last computed token, so we return num_computed_tokens - 1. """ @@ -2600,8 +2596,7 @@ def get_manager_for_kv_cache_spec( role: str | None = None, **kwargs, ) -> SingleTypeKVCacheManager: - """ - Get the appropriate manager for a given KVCacheSpec. + """Get the appropriate manager for a given KVCacheSpec. Uses the KVCacheSpecRegistry to look up the manager class, supporting both built-in and custom specs registered via @register_kv_cache_spec @@ -2612,8 +2607,11 @@ def get_manager_for_kv_cache_spec( max_in_flight_tokens: The max tokens scheduled but not yet settled (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens` max_model_len: The maximum context length the model could serve + role: Optional KV cache group role used to select a role-specific + manager Returns: An instance of the appropriate SingleTypeKVCacheManager subclass + """ manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec, role) assert manager_class is not None, ( diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index 6a48b6282d43..0c1fb0b2c4c9 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -13,8 +13,7 @@ class CudagraphDispatcher: - """ - Runtime cudagraph dispatcher to dispatch keys for multiple set of + """Runtime cudagraph dispatcher to dispatch keys for multiple set of cudagraphs. The dispatcher stores two sets of dispatch keys, one for PIECEWISE and one @@ -109,8 +108,7 @@ def _compute_bs_to_padded_graph_size(self) -> None: ) def _get_lora_cases(self) -> list[int]: - """ - Returns list of has_lora values for CUDA graph capture. + """Returns list of has_lora values for CUDA graph capture. This is the single source of truth for LoRA capture cases. """ lora_config = self.vllm_config.lora_config @@ -241,8 +239,7 @@ def dispatch( valid_modes: AbstractSet[CUDAGraphMode] | None = None, invalid_modes: AbstractSet[CUDAGraphMode] | None = None, ) -> tuple[CUDAGraphMode, BatchDescriptor]: - """ - Given conditions(e.g.,batch descriptor and if using piecewise only), + """Given conditions(e.g.,batch descriptor and if using piecewise only), dispatch to a cudagraph runtime mode and the valid batch descriptor. A new batch descriptor is returned as we might dispatch a uniform batch to a graph that supports a more general batch (uniform to non-uniform). @@ -259,6 +256,7 @@ def dispatch( valid_modes to compute allowed modes. (e.g., {FULL} for features like cascade attention not supported by full cudagraphs). None means no modes are excluded. + """ allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes() @@ -324,13 +322,13 @@ def dispatch( return CUDAGraphMode.NONE, BatchDescriptor(num_tokens) def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]: - """ - Returns capture descriptors for cudagraph capturing. + """Returns capture descriptors for cudagraph capturing. Returns: List of (runtime_mode, batch_descriptors) tuples, ordered PIECEWISE first then FULL. Batch descriptors are sorted largest-first for memory efficiency. + """ if not self.keys_initialized or self.cudagraph_mode == CUDAGraphMode.NONE: return [] diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 7d68b6e7978b..7c424ea0df0c 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -45,8 +45,7 @@ class EEPNotificationType(enum.Enum): class FinishReason(enum.IntEnum): - """ - Reason a request finished - stop, length, abort, error, or repetition. + """Reason a request finished - stop, length, abort, error, or repetition. Int rather than Str for more compact serialization. @@ -284,8 +283,7 @@ def __post_init__(self): class EngineCoreRequestType(enum.Enum): - """ - Request types defined as hex byte strings, so it can be sent over sockets + """Request types defined as hex byte strings, so it can be sent over sockets without separate encoding step. """ @@ -310,9 +308,7 @@ class ReconfigureDistributedRequest(msgspec.Struct): class ReconfigureRankType(enum.IntEnum): - """ - Rank type for reconfiguring distributed request. - """ + """Rank type for reconfiguring distributed request.""" KEEP_CURRENT_RANK = -1 SHUTDOWN_CURRENT_RANK = -2 diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index bf3f5f420ae8..64aa5fa2ce8f 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -94,8 +94,7 @@ def __init__( client_index: int = 0, profiler: TorchProfilerWrapper | None = None, ) -> None: - """ - Create an AsyncLLM. + """Create an AsyncLLM. Args: vllm_config: global configuration. @@ -105,13 +104,21 @@ def __init__( mm_registry: Multi-modal registry. log_requests: Whether to log requests. start_engine_loop: Whether to start the engine loop. + client_addresses: ZMQ addresses of an externally launched engine, + when not launching one in-process. + client_count: Number of API-server clients sharing the engine. + client_index: Index of this client within ``client_count``. + aggregate_engine_logging: Whether to aggregate per-engine logs + into a single stat logger. stat_loggers: customized stat loggers for the engine. If not provided, default stat loggers will be used. PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE. + profiler: Torch profiler wrapper used to trace the frontend. Returns: None + """ # Ensure we can serialize custom transformer configs maybe_register_config_serialize_by_value() @@ -255,7 +262,6 @@ def from_engine_args( stat_loggers: list[StatLoggerFactory] | None = None, ) -> "AsyncLLM": """Create an AsyncLLM from the EngineArgs.""" - # Create the engine configs. vllm_config = engine_args.create_engine_config(usage_context) executor_class = Executor.get_class(vllm_config) @@ -319,6 +325,7 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: Raises: QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded. MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded. + """ max_num_reqs = self.scheduler_config.max_num_queued_reqs if max_num_reqs is not None: @@ -378,7 +385,6 @@ async def add_request( reasoning_parser_kwargs: dict[str, Any] | None = None, ) -> RequestOutputCollector: """Add new request to the AsyncLLM.""" - if self.errored: raise EngineDeadError() @@ -654,8 +660,7 @@ async def generate( reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, ) -> AsyncGenerator[RequestOutput, None]: - """ - Main function called by the API server to kick off a request + """Main function called by the API server to kick off a request * 1) Making an AsyncStream corresponding to the Request. * 2) Processing the Input. * 3) Adding the Request to the Detokenizer. @@ -678,8 +683,8 @@ async def generate( >>> params = self.renderer.default_cmpl_tok_params >>> (engine_input,) = self.renderer.render_cmpl([parsed], params) >>> gen = self.generate(engine_input, sampling_params, request_id) - """ + """ q: RequestOutputCollector | None = None try: q = await self.add_request( @@ -763,7 +768,6 @@ async def generate( def _run_output_handler(self): """Background loop: pulls from EngineCore and pushes to AsyncStreams.""" - if self.output_handler is not None: return @@ -849,7 +853,6 @@ async def abort( self, request_id: str | Iterable[str], internal: bool = False ) -> None: """Abort RequestId in OutputProcessor and EngineCore.""" - request_ids = ( (request_id,) if isinstance(request_id, str) else as_list(request_id) ) @@ -868,7 +871,8 @@ async def notify_kv_transfer_request_rejected( ) -> None: """Submit a pre-aborted request so the connector's request_finished hook runs to free any pre-admission KV-transfer resources (e.g. NIXL - prefill blocks pinned on the P node).""" + prefill blocks pinned on the P node). + """ request = EngineCoreRequest( request_id=request_id, prompt_token_ids=[0], @@ -893,8 +897,7 @@ async def pause_generation( wait_for_inflight_requests: bool | None = None, clear_cache: bool = True, ) -> None: - """ - Pause generation to allow model weight updates. + """Pause generation to allow model weight updates. All mode handling (abort / wait / keep) and cache clearing is done in the engine. New generation/encoding requests will not be scheduled @@ -910,6 +913,7 @@ async def pause_generation( wait_for_inflight_requests: DEPRECATED: use mode argument. clear_cache: Whether to clear KV cache and prefix cache after draining. Set to ``False`` to preserve cache for faster resume. + """ if wait_for_inflight_requests: warnings.warn( @@ -950,8 +954,7 @@ async def encode( tokenization_kwargs: dict[str, Any] | None = None, reasoning_ended: bool | None = None, ) -> AsyncGenerator[PoolingRequestOutput, None]: - """ - Main function called by the API server to kick off a request + """Main function called by the API server to kick off a request * 1) Making an AsyncStream corresponding to the Request. * 2) Processing the Input. * 3) Adding the Request to the EngineCore (separate process). @@ -963,7 +966,6 @@ async def encode( The caller of generate() iterates the returned AsyncGenerator, returning the RequestOutput back to the caller. """ - q: RequestOutputCollector | None = None try: q = await self.add_request( @@ -1116,9 +1118,7 @@ async def collective_rpc( args: tuple = (), kwargs: dict | None = None, ): - """ - Perform a collective RPC call to the given path. - """ + """Perform a collective RPC call to the given path.""" return await self.engine_core.collective_rpc_async( method, timeout, args, kwargs ) @@ -1198,7 +1198,7 @@ async def _scale_elastic_ep( async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest ) -> FaultToleranceResult: - """send fault tolerance instruction to the engine""" + """Send fault tolerance instruction to the engine.""" return await self.engine_core.handle_fault(fault_tolerance_request) async def get_status(self): @@ -1224,11 +1224,11 @@ def dead_error(self) -> BaseException: async def init_weight_transfer_engine( self, request: WeightTransferInitRequest ) -> None: - """ - Initialize weight transfer for RL training. + """Initialize weight transfer for RL training. Args: request: Weight transfer initialization request with backend-specific info + """ await self.collective_rpc( "init_weight_transfer_engine", kwargs={"init_info": request.init_info} @@ -1243,11 +1243,11 @@ async def start_draft_weight_update(self) -> None: await self.collective_rpc("start_draft_weight_update") async def update_weights(self, request: WeightTransferUpdateRequest) -> None: - """ - Batched weight update for RL training. + """Batched weight update for RL training. Args: request: Weight update request with backend-specific update info + """ await self.collective_rpc( "update_weights", kwargs={"update_info": request.update_info} diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b3d94de3e7db..5125954c42e3 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -505,7 +505,6 @@ def add_request(self, request: Request, request_wave: int = 0): def abort_requests(self, request_ids: list[str]): """Abort requests from the scheduler.""" - # TODO: The scheduler doesn't really need to know the # specific finish reason, TBD whether we propagate that # (i.e. client-aborted vs stop criteria met). @@ -599,7 +598,8 @@ def _attach_iteration_details( def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). - Overridden by the DP engine core; never throttles otherwise.""" + Overridden by the DP engine core; never throttles otherwise. + """ return False def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: @@ -608,7 +608,6 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: Returns tuple of outputs and a flag indicating whether the model was executed. """ - # Check for any requests remaining in the scheduler - unfinished, # or finished and not yet removed from the batch. if not self.scheduler.has_requests(): @@ -659,7 +658,6 @@ def step_with_batch_queue( batch in the job queue is finished. 3. Update the scheduler from the output. """ - batch_queue = self.batch_queue assert batch_queue is not None @@ -903,8 +901,8 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None | Future: - Level 2: Discard all GPU memory. mode: Pause mode - how to deal with any existing requests, see documentation of pause_scheduler method. - """ + """ # Pause scheduler before sleeping. clear_prefix_cache = level >= 1 pause_future = self.pause_scheduler(mode=mode, clear_cache=clear_prefix_cache) @@ -935,6 +933,7 @@ def wake_up(self, tags: list[str] | None = None): Args: tags: Tags to wake up. Use ["scheduling"] for level 0 wake up. + """ if tags is not None and "scheduling" in tags: # Remove "scheduling" from tags if there are other tags to process. @@ -1162,8 +1161,7 @@ def _perform_handshakes( vllm_config: VllmConfig, client_handshake_address: str | None, ) -> Generator[EngineZmqAddresses, None, None]: - """ - Perform startup handshakes. + """Perform startup handshakes. For DP=1 or offline mode, this is with the colocated front-end process. @@ -1298,7 +1296,6 @@ def startup_handshake( @staticmethod def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): """Launch EngineCore busy loop in background process.""" - # Ensure we can serialize transformer config after spawning maybe_register_config_serialize_by_value() @@ -1444,7 +1441,6 @@ def _maybe_publish_request_counts(self): def _process_input_queue(self): """Exits when an engine step needs to be performed.""" - waited = False while not self.has_work() and self.is_running(): # Notify callbacks waiting for engine to become idle. @@ -1475,7 +1471,6 @@ def _process_input_queue(self): def _process_engine_step(self) -> bool: """Called only when there are unfinished local requests.""" - # Step the engine core. outputs, model_executed = self.step_fn() # Put EngineCoreOutputs into the output queue. @@ -1549,7 +1544,6 @@ def _handle_client_request( self, request_type: EngineCoreRequestType, request: Any ) -> None: """Dispatch request from client.""" - if request_type == EngineCoreRequestType.WAKEUP: return elif request_type == EngineCoreRequestType.ADD: @@ -1629,7 +1623,8 @@ def _invoke_utility_method( @staticmethod def _convert_msgspec_args(method, args): """If a provided arg type doesn't match corresponding target method - arg type, try converting to msgspec object.""" + arg type, try converting to msgspec object. + """ if not args: return args arg_types = signature(method).parameters.values() @@ -1645,7 +1640,6 @@ def _convert_msgspec_args(method, args): def _send_engine_dead(self): """Send EngineDead status to the EngineCoreClient.""" - # Put ENGINE_CORE_DEAD in the queue. self.output_queue.put_nowait(EngineCoreProc.ENGINE_CORE_DEAD) @@ -1710,7 +1704,6 @@ def process_input_sockets( ready_event: threading.Event, ): """Input socket IO thread.""" - # Msgpack serialization decoding with optional tensor IPC receiver. add_request_decoder = MsgpackDecoder( EngineCoreRequest, oob_tensor_provider=self.tensor_ipc_receiver @@ -1809,7 +1802,6 @@ def process_output_sockets( self, output_paths: list[str], coord_output_path: str | None, engine_index: int ): """Output socket IO thread.""" - # Msgpack serialization encoding. encoder = MsgpackEncoder() # Send buffers to reuse. @@ -2014,7 +2006,8 @@ def _send_abort_outputs(self, aborted_reqs: list[Request]) -> None: class DPEngineCoreProc(EngineCoreProc): """ZMQ-wrapper for running EngineCore in background process - in a data parallel context.""" + in a data parallel context. + """ def __init__( self, @@ -2199,7 +2192,6 @@ def _should_throttle_prefills(self) -> bool: @fault_tolerant_wrapper def run_busy_loop(self): """Core busy loop of the EngineCore for data parallel case.""" - # Loop until process is sent a SIGINT or SIGTERM while self._handle_shutdown(): # 1) Poll the input queue until there is work to do. @@ -2364,8 +2356,7 @@ def commit_prepared_elastic_ep(self) -> None: def _eep_send_engine_core_notification( self, notification_type: EEPNotificationType ): - """ - Send notifications to EngineCoreClient, which can then forward + """Send notifications to EngineCoreClient, which can then forward the notifications to other engine core processes. It is used for: 1) In scale down: removing core engines to notify EngineCoreClient so EngineCoreClient can release their ray placement groups; @@ -2415,9 +2406,7 @@ def _eep_scale_up_before_kv_init(self): class EngineCoreActorMixin: - """ - Ray actor for running EngineCore in a data parallel context - """ + """Ray actor for running EngineCore in a data parallel context.""" def __init__( self, @@ -2510,16 +2499,14 @@ def _perform_handshakes( vllm_config: VllmConfig, client_handshake_address: str | None, ): - """ - For Ray, we don't need to actually perform handshake. + """For Ray, we don't need to actually perform handshake. All addresses information is known before the actor creation. Therefore, we simply yield these addresses. """ yield self.addresses def wait_for_init(self): - """ - Wait until the engine core is initialized. + """Wait until the engine core is initialized. This is just an empty method. When ray.get() on this method (or any other method of the actor) returns, it is guaranteed @@ -2528,9 +2515,7 @@ def wait_for_init(self): pass def run(self): - """ - Run the engine core busy loop. - """ + """Run the engine core busy loop.""" try: self.run_busy_loop() # type: ignore[attr-defined] except SystemExit: diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index b63e8016c470..6a08f2d3afc3 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -77,8 +77,7 @@ class EngineCoreClient(ABC): - """ - EngineCoreClient: subclasses handle different methods for pushing + """EngineCoreClient: subclasses handle different methods for pushing and pulling from the EngineCore for asyncio / multiprocessing. Subclasses: @@ -249,7 +248,8 @@ def collective_rpc( def dp_engines_running(self) -> bool: """Returns True if data parallel engines are collectively in a - running state.""" + running state. + """ raise NotImplementedError async def commit_elastic_ep(self) -> None: @@ -331,8 +331,7 @@ async def get_status(self): class InprocClient(EngineCoreClient): - """ - InprocClient: client for in-process EngineCore. Intended + """InprocClient: client for in-process EngineCore. Intended for use in LLMEngine for V0-style add_request() and step() EngineCore setup in this process (no busy loop). @@ -447,7 +446,8 @@ def dp_engines_running(self) -> bool: @dataclass class BackgroundResources: """Used as a finalizer for clean shutdown, avoiding - circular reference back to the client object.""" + circular reference back to the client object. + """ ctx: zmq.Context # If CoreEngineProcManager, it manages local engines; @@ -469,7 +469,6 @@ class BackgroundResources: def __call__(self): """Clean up background resources.""" - logger.debug_once("[shutdown] MPClient: background resource cleanup start") self.engine_dead = True if self.engine_manager is not None: @@ -543,16 +542,15 @@ class ElasticScalingCache: class MPClient(EngineCoreClient): - """ - MPClient: base client for multi-proc EngineCore. - EngineCore runs in a background process busy loop, getting - new EngineCoreRequests and returning EngineCoreOutputs + """MPClient: base client for multi-proc EngineCore. + EngineCore runs in a background process busy loop, getting + new EngineCoreRequests and returning EngineCoreOutputs - * pushes EngineCoreRequests via input_socket - * pulls EngineCoreOutputs via output_socket + * pushes EngineCoreRequests via input_socket + * pulls EngineCoreOutputs via output_socket - * AsyncMPClient subclass for AsyncLLM usage - * SyncMPClient subclass for LLM usage + * AsyncMPClient subclass for AsyncLLM usage + * SyncMPClient subclass for LLM usage """ def __init__( @@ -803,7 +801,8 @@ def monitor_engine_cores(): def _apply_ready_response(self, payload: bytes) -> None: """Decode an EngineCoreReadyResponse and sync any post-initialization - config changes (e.g. auto-fitted max_model_len) back to the frontend.""" + config changes (e.g. auto-fitted max_model_len) back to the frontend. + """ vllm_config = self.vllm_config response = ( msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) @@ -1334,7 +1333,8 @@ async def get_status(self): class DPAsyncMPClient(AsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) - EngineCore. Assumes external load-balancing by default.""" + EngineCore. Assumes external load-balancing by default. + """ def __init__( self, @@ -1518,7 +1518,8 @@ def get_core_engine_for_request(self, request: EngineCoreRequest): class DPLBAsyncMPClient(DPAsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) - EngineCore. Load-balances between multiple engine processes.""" + EngineCore. Load-balances between multiple engine processes. + """ def __init__( self, @@ -1754,8 +1755,7 @@ async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future: - """ - Wait for core engines to switch to the new setup. + """Wait for core engines to switch to the new setup. In eep_process_engine_core_notification(), a dummy UtilityOutput with EEP_NOTIFICATION_CALL_ID will be set when RECONFIGURE_FINISHED @@ -1827,7 +1827,8 @@ async def _prepare_scale_up_elastic_ep( num_redundant_experts: int, ) -> None: """Prepare scale up by creating new engine cores and reconfiguring - existing ones.""" + existing ones. + """ self._setup_elastic_ep_reconfig_bootstrap() # Phase 1: Send reconfig messages to existing engines @@ -1919,7 +1920,8 @@ async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> N async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: """Scale down the data parallel size by shutting down and - reconfiguring existing engine cores.""" + reconfiguring existing engine cores. + """ cur_data_parallel_size = len(self.core_engines) self.eep_scaling_cache = ElasticScalingCache( diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index f04fa30a1856..c75fdc53cf43 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -94,8 +94,7 @@ def __init__(self, request: EngineCoreRequest): self.output_text = "" def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: - """ - Update RequestState for the request_id by: + """Update RequestState for the request_id by: 1) Detokenize the new token ids incrementally. 2) Evaluate stop criteria. @@ -148,8 +147,8 @@ def decode_next(self, next_token_id: int) -> str: def get_next_output_text(self, finished: bool, delta: bool) -> str: """If delta is True, only new text since the last call to - this method is returned""" - + this method is returned + """ # We return the full output text if the sequence is finished. buffer_length = 0 if finished else self.stop_buffer_length if not delta: diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 3f86f3c8b239..256a72df4254 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -206,8 +206,7 @@ def _get_mm_identifier( mm_hash: str, lora_request: LoRARequest | None, ) -> str: - """ - When enable_tower_connector_lora is True, multi-modal embeddings + """When enable_tower_connector_lora is True, multi-modal embeddings vary depending on the LoRA request. Therefore, the mm_hash must be generated based on the LoRA request to prevent incorrect cache hits. """ diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 32096079e066..c6c77f72e716 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -171,7 +171,6 @@ def from_engine_args( enable_multiprocessing: bool = False, ) -> "LLMEngine": """Creates an LLM engine from the engine arguments.""" - # Create the engine configs. vllm_config = engine_args.create_engine_config(usage_context) executor_class = Executor.get_class(vllm_config) @@ -216,7 +215,6 @@ def get_supported_tasks(self) -> tuple[SupportedTask, ...]: def abort_request(self, request_ids: list[str], internal: bool = False) -> None: """Remove request_ids from EngineCore and Detokenizer.""" - request_ids = self.output_processor.abort_requests(request_ids, internal) self.engine_core.abort_requests(request_ids) diff --git a/vllm/v1/engine/logprobs.py b/vllm/v1/engine/logprobs.py index 8eef567b0715..ff483e128ee2 100644 --- a/vllm/v1/engine/logprobs.py +++ b/vllm/v1/engine/logprobs.py @@ -76,7 +76,6 @@ def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None: logprobs_lists: the lists of logprob tokens, logprobs, and ranks. """ - assert self.num_logprobs is not None assert self.logprobs is not None assert self.cumulative_logprob is not None @@ -129,7 +128,6 @@ def _update_prompt_logprobs( tensors. """ - # Prompt logprobs are enabled. assert self.num_prompt_logprobs is not None assert self.prompt_logprobs is not None @@ -187,7 +185,7 @@ def _update_prompt_logprobs( ) def pop_prompt_logprobs(self) -> PromptLogprobs | None: - """Pop and return all request prompt logprobs + """Pop and return all request prompt logprobs. The logprobs processor aggregates prompt chunk logprobs over one or more prefill chunks. This method returns @@ -199,6 +197,7 @@ def pop_prompt_logprobs(self) -> PromptLogprobs | None: Returns: None if prompt logprobs are disabled for this request. List of all prompt logprobs, otherwise. + """ plp = self.prompt_logprobs if plp: @@ -223,6 +222,7 @@ def _get_sampled_context_ids( Returns: List of sampled token IDs, oldest first, most recent last. + """ if not logprobs_source: return [] @@ -265,6 +265,7 @@ def _correct_decoded_token( Returns: The corrected decoded string, or empty string if the byte sequence is genuinely incomplete at this point. + """ assert self.tokenizer is not None @@ -325,6 +326,7 @@ def _verify_tokens( context_token_ids: Preceding sampled token IDs providing sequential context. If None, extracted from self.logprobs. + """ if context_token_ids is None: context_token_ids = self._get_sampled_context_ids(self.logprobs) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index b8fbddf72ff8..b4f7338048c7 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -49,8 +49,7 @@ class RequestOutputCollector: - """ - Collects streamed RequestOutputs per individual request, + """Collects streamed RequestOutputs per individual request, for hand-off to the consuming asyncio generate task. When streaming deltas, RequestOutputs are merged if the @@ -506,7 +505,6 @@ def has_unfinished_requests(self) -> bool: def propagate_error(self, e: Exception): """Propagate error to all generate() tasks.""" - for _, state in self.request_states.items(): assert state.queue is not None state.queue.put(e) @@ -646,8 +644,7 @@ def process_outputs( engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, ) -> OutputProcessorOutput: - """ - Process the EngineCoreOutputs: + """Process the EngineCoreOutputs: 1) Compute stats for logging 2) Detokenize 3) Create and handle RequestOutput objects: @@ -667,7 +664,6 @@ def process_outputs( If you need to touch every element of the batch, do it from within the loop below. """ - request_outputs: list[RequestOutput | PoolingRequestOutput] = [] reqs_to_abort: list[str] = [] for engine_core_output in engine_core_outputs: diff --git a/vllm/v1/engine/parallel_sampling.py b/vllm/v1/engine/parallel_sampling.py index 8eb6fa057d37..9baf785903a1 100644 --- a/vllm/v1/engine/parallel_sampling.py +++ b/vllm/v1/engine/parallel_sampling.py @@ -53,7 +53,7 @@ def _get_child_sampling_params( self, index: int, ) -> SamplingParams: - """Efficiently obtain child `sampling_params` + """Efficiently obtain child `sampling_params`. If `sampling_params.seed` is not `None` then each child request requires a unique clone of @@ -64,6 +64,7 @@ def _get_child_sampling_params( Returns: Child `sampling_params` instance. + """ seed = self.sampling_params.seed if self.cached_child_sampling_params: @@ -88,6 +89,7 @@ def get_child_info(self, index: int) -> tuple[str, SamplingParams]: Returns: (request ID, sampling_params) tuple + """ child_req_id = f"{index}_{self.request_id}" self.child_requests.add(child_req_id) diff --git a/vllm/v1/engine/tensor_ipc.py b/vllm/v1/engine/tensor_ipc.py index 2a8391b17e6d..2b8fbadd97a5 100644 --- a/vllm/v1/engine/tensor_ipc.py +++ b/vllm/v1/engine/tensor_ipc.py @@ -29,8 +29,7 @@ @dataclasses.dataclass class TensorIpcData: - """ - Data sent via torch.multiprocessing.Queue for zero-copy IPC. + """Data sent via torch.multiprocessing.Queue for zero-copy IPC. Contains the tensor_id and the actual tensor. The tensor is shared in memory (GPU or CPU) for efficient inter-process communication. @@ -130,7 +129,6 @@ def __call__( the queue, buffering them, until the requested tensor is found. Works for CUDA and CPU. """ - # Create lookup key from handle sender_id: str = meta["sender_id"] message_id: int = meta["message_id"] diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 9bce2b7687ef..43f2b80b05e2 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -142,8 +142,7 @@ def _node_ip_from_resources(node_resources: dict) -> str | None: class CoreEngineProcManager: - """ - Utility class to handle creation, readiness, and shutdown + """Utility class to handle creation, readiness, and shutdown of background processes used by the AsyncLLM and LLMEngine. """ @@ -255,7 +254,6 @@ def shutdown(self, timeout: float | None = None) -> None: def monitor_engine_liveness(self) -> None: """Monitor engine core process liveness.""" - sentinel_to_proc = {proc.sentinel: proc for proc in self.processes} sentinels = set(sentinel_to_proc.keys()) @@ -316,8 +314,7 @@ def set_assigned_physical_gpu_ids_for_dp_rank( local_dp_rank: int, user_assigned_gpu_ids: list[int] | None = None, ) -> None: - """ - Populate assigned_physical_gpu_ids on the config for the given DP rank. + """Populate assigned_physical_gpu_ids on the config for the given DP rank. user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the user provided one; this DP rank's shard is sliced from it. It is passed @@ -345,8 +342,7 @@ def get_physical_gpu_ids_for_local_dp_rank( local_world_size: int | None = None, user_assigned_gpu_ids: list[int] | None = None, ) -> list[int]: - """ - Returns list of physical GPU IDs for the specified + """Returns list of physical GPU IDs for the specified data parallel rank. For example, if world_size=2 and local_dp_rank=1, and there are 4 devices, @@ -399,8 +395,7 @@ def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None: class CoreEngineActorManager: - """ - Utility class to handle creation, readiness, and shutdown + """Utility class to handle creation, readiness, and shutdown of core engine Ray actors used by the AsyncLLM and LLMEngine. Different from CoreEngineProcManager, this class manages @@ -552,10 +547,7 @@ def __init__( def create_dp_placement_groups( vllm_config: VllmConfig, ) -> tuple[list["PlacementGroup"], list[int]]: - """ - Create placement groups for data parallel. - """ - + """Create placement groups for data parallel.""" import ray from ray._private.state import available_resources_per_node @@ -769,9 +761,7 @@ def create_dp_placement_groups( def add_dp_placement_groups( old_vllm_config: VllmConfig, new_data_parallel_size: int ) -> tuple[list["PlacementGroup"], list[int]]: - """ - Add placement groups for new data parallel size. - """ + """Add placement groups for new data parallel size.""" import ray from ray._private.state import ( available_resources_per_node, @@ -1051,7 +1041,8 @@ def get_engine_zmq_addresses( Set ``defer_api_server_ports=False`` only when the consumer cannot report a bound port back (e.g. the Rust front-end). IPC paths are - unaffected.""" + unaffected. + """ parallel_config = vllm_config.parallel_config local_engine_count = parallel_config.data_parallel_size_local local_start_index = parallel_config.data_parallel_rank_local @@ -1108,7 +1099,6 @@ def launch_core_engines( addresses: EngineZmqAddresses, ) -> Iterator[CoreEngineLaunch]: """Launch engine and DP coordinator processes as needed.""" - parallel_config = vllm_config.parallel_config dp_size = parallel_config.data_parallel_size local_engine_count = parallel_config.data_parallel_size_local diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 964f16ed2421..ce8a8db826a7 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -36,7 +36,7 @@ class Executor(ABC): - """Abstract base class for vLLM executors." + """Abstract base class for vLLM executors. An executor is responsible for executing the model on one device, or it can be a distributed executor that can execute the model on multiple devices. @@ -139,8 +139,7 @@ def compile_or_warm_up_model(self) -> None: ) def register_failure_callback(self, callback: FailureCallback): # noqa: B027 - """ - Register a function to be called if the executor enters a permanent + """Register a function to be called if the executor enters a permanent failed state. """ pass @@ -168,8 +167,29 @@ def collective_rpc( kwargs: dict | None = None, non_block: Literal[False] = False, ) -> list[_R]: - """ - Execute an RPC call on all workers. + pass + + @overload + def collective_rpc( + self, + method: str | Callable[[WorkerBase], _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None, + non_block: Literal[True] = True, + ) -> Future[list[_R]]: + pass + + @abstractmethod + def collective_rpc( + self, + method: str | Callable[[WorkerBase], _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None, + non_block: bool = False, + ) -> list[_R] | Future[list[_R]]: + """Execute an RPC call on all workers. Args: method: Name of the worker method to execute, or a callable that @@ -191,24 +211,8 @@ def collective_rpc( Note: It is recommended to use this API to only pass control messages, and set up data-plane communication to pass data. - """ - pass - - @overload - def collective_rpc( - self, - method: str | Callable[[WorkerBase], _R], - timeout: float | None = None, - args: tuple = (), - kwargs: dict | None = None, - non_block: Literal[True] = True, - ) -> Future[list[_R]]: - pass - @abstractmethod - def collective_rpc( - self, method, timeout=None, args=(), kwargs=None, non_block: bool = False - ): + """ raise NotImplementedError def get_kv_connector_handshake_metadata( @@ -280,7 +284,8 @@ def save_sharded_state( @abstractmethod def check_health(self) -> None: """Checks if the executor is healthy. If not, it should raise an - exception.""" + exception. + """ raise NotImplementedError def shutdown(self) -> None: @@ -288,7 +293,7 @@ def shutdown(self) -> None: self.collective_rpc("shutdown") def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: - """Init KVOutputAggregator""" + """Init KVOutputAggregator.""" self.kv_output_aggregator = KVOutputAggregator.from_connector( connector, self.parallel_config.world_size ) @@ -381,9 +386,7 @@ def reinitialize_distributed( @classmethod def supports_async_scheduling(cls) -> bool: - """ - Whether the executor supports async scheduling. - """ + """Whether the executor supports async scheduling.""" return False diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 67a631d3ce20..009df6ffba2e 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -384,7 +384,8 @@ def collective_rpc( # type: ignore[override] ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: """Returns single result if unique_reply_rank and/or an output - aggregator is provided, otherwise list.""" + aggregator is provided, otherwise list. + """ assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -451,7 +452,8 @@ def get_response(): def _ensure_worker_termination(worker_procs: list[BaseProcess]): """Ensure that all worker processes are terminated. Assumes workers have received termination requests. Waits for processing, then sends - termination and kill signals if needed.""" + termination and kill signals if needed. + """ def wait_for_termination(procs, timeout): if not time: @@ -500,7 +502,7 @@ def wait_for_termination(procs, timeout): p.kill() def shutdown(self): - """Properly shut down the executor and its workers""" + """Properly shut down the executor and its workers.""" if not getattr(self, "shutting_down", False): worker_count = len(getattr(self, "workers", None) or []) logger.debug( @@ -852,8 +854,8 @@ def death_pipe_monitor(queues_to_shutdown: list[MessageQueue]): @staticmethod def worker_main(*args, **kwargs): """Worker initialization and execution loops. - This runs a background process""" - + This runs a background process + """ # Signal handler used for graceful termination. # SystemExit exception is only raised once to allow this and worker # processes to terminate without error @@ -1010,7 +1012,6 @@ def handle_output(self, output: Any): def async_output_busy_loop(self): """Entrypoint for the thread which handles outputs asynchronously.""" - # set device to the worker device for the thread. # a thread will not inherit the context of the main thread. # when calling any cuda runtime functions, it will implicitly @@ -1027,7 +1028,7 @@ def async_output_busy_loop(self): self.enqueue_output(output) def worker_busy_loop(self): - """Main busy loop for Multiprocessing Workers""" + """Main busy loop for Multiprocessing Workers.""" assert self.rpc_broadcast_mq is not None while True: self._execute_worker_rpc(self.rpc_broadcast_mq.dequeue(indefinite=True)) @@ -1098,8 +1099,8 @@ def setup_proc_title_and_log_prefix(enable_ep: bool) -> None: def set_multiprocessing_worker_envs(local_world_size: int = 1): """Set up environment variables that should be used when there are workers in a multiprocessing environment. This should be called by the parent - process before worker processes are created""" - + process before worker processes are created + """ _maybe_force_spawn() if current_platform.is_cpu() or "OMP_NUM_THREADS" in os.environ: diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index b2c385a862b8..d101885ce2b6 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -52,8 +52,7 @@ @dataclass class RayWorkerMetaData: - """ - Metadata for a Ray worker. + """Metadata for a Ray worker. The order of ray worker creation can be random, and we need to reset the rank after creating all workers. """ @@ -65,7 +64,7 @@ class RayWorkerMetaData: class RayDistributedExecutor(Executor): - """Ray-based distributed executor""" + """Ray-based distributed executor.""" uses_ray: bool = True supports_pp: bool = True @@ -249,8 +248,7 @@ def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwar ip_counts[ip] = ip_counts.get(ip, 0) + 1 def sort_by_driver_then_worker_ip(item: RayWorkerMetaData): - """ - Sort the workers based on 3 properties: + """Sort the workers based on 3 properties: 1. If the worker is on the same node as the driver (vllm engine), it should be placed first. 2. Then, if the worker is on a node with fewer workers, it should @@ -439,6 +437,7 @@ def sample_tokens( # type: ignore[override] Returns: The model runner output. + """ scheduler_output = self.scheduler_output if scheduler_output is None: diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py index 8f4b84ed0f13..b8a2eec73f77 100644 --- a/vllm/v1/executor/ray_executor_v2.py +++ b/vllm/v1/executor/ray_executor_v2.py @@ -72,7 +72,7 @@ class RayWorkerHandle: """run() ObjectRef used as a sentinel for health monitoring""" def run(self): - """Start the worker's busy loop""" + """Start the worker's busy loop.""" self.run_ref = self.actor.run.remote() @@ -190,8 +190,7 @@ def initialize_worker( def _init_message_queues( self, input_shm_handle: Handle, vllm_config: VllmConfig ) -> None: - """ - Workers on the same node as the executor use shared memory for + """Workers on the same node as the executor use shared memory for both the broadcast (input) MQ and the response MQ. Workers on different nodes use TCP (n_local_reader=0). """ diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 71d4cf74717e..0f6d2c9c6416 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -55,7 +55,8 @@ class RayWorkerWrapper(WorkerWrapperBase): """Ray wrapper for vllm.worker.Worker, allowing Worker to be - lazily initialized after Ray sets CUDA_VISIBLE_DEVICES.""" + lazily initialized after Ray sets CUDA_VISIBLE_DEVICES. + """ def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) @@ -68,8 +69,7 @@ def __init__(self, *args, **kwargs) -> None: rpc_rank: int def adjust_rank(self, rank_mapping: dict[int, int]) -> None: - """ - Adjust the rpc_rank based on the given mapping. + """Adjust the rpc_rank based on the given mapping. It is only used during the initialization of the executor, to adjust the rpc_rank of workers after we create all workers. """ @@ -363,8 +363,7 @@ def get_bundles_for_indices( bundle_indices: list[int], world_size: int, ) -> list[tuple[int, str, str]]: - """ - Return GPU bundle indices paired with node IDs and node IPs for + """Return GPU bundle indices paired with node IDs and node IPs for explicit bundle indices specified via VLLM_RAY_BUNDLE_INDICES. """ assert len(bundle_indices) == world_size, ( @@ -391,8 +390,7 @@ def get_bundles_for_indices( def get_bundles_sorted_by_node( placement_group: "PlacementGroup", ) -> list[tuple[int, str, str]]: - """ - Return GPU bundle indices paired with node IDs and node IPs, + """Return GPU bundle indices paired with node IDs and node IPs, sorted driver-first. This utility has to be invoked from the driver node. @@ -533,6 +531,7 @@ def initialize_ray_cluster( on the current (driver) node and pin the first PG bundle to it. Set to False for executors like RayExecutorV2 where all GPU work is delegated to remote Ray actors. + """ assert_ray_available() from vllm.platforms import current_platform diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 7e91300b2fe2..4a2ea1799f68 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -154,9 +154,7 @@ class SparseCacheRole(str, Enum): @dataclass(frozen=True) class KVCacheSpec: - """ - A base class for specifying the KV cache format of one layer. - """ + """A base class for specifying the KV cache format of one layer.""" # number of tokens in a block block_size: int @@ -180,11 +178,11 @@ def state_content_size_bytes(self) -> int: @property def page_size_bytes(self) -> int: - """ - The size of a page with `block_size` tokens in bytes. + """The size of a page with `block_size` tokens in bytes. Returns: The page size + """ raise NotImplementedError @@ -202,37 +200,33 @@ def block_table_token_alignment(self) -> int | None: return 128 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - """ - The maximum possible memory usage of this KV cache in bytes. + """The maximum possible memory usage of this KV cache in bytes. Returns: The KV cache size in bytes + """ raise NotImplementedError def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: - """ - The number of block table entries needed per request, i.e. the row + """The number of block table entries needed per request, i.e. the row length of the worker-side block table for this cache group. Args: vllm_config: The vllm config. max_len: The maximum sequence length to size for, including the encoder length for encoder-decoder models. + """ return cdiv(max_len, self.block_size) def copy_with_new_block_size(self, block_size: int) -> Self: - """ - Create a new KVCacheSpec from self but replacing the block size. - """ + """Create a new KVCacheSpec from self but replacing the block size.""" return replace(self, block_size=block_size) @classmethod def merge(cls, specs: list[Self]) -> Self: - """ - Merge a list of KVCacheSpec objects into a single KVCacheSpec object. - """ + """Merge a list of KVCacheSpec objects into a single KVCacheSpec object.""" if not all(spec == specs[0] for spec in specs[1:]): raise AssertionError( "All layers in the same KV cache group must be the same." @@ -242,9 +236,7 @@ def merge(cls, specs: list[Self]) -> Self: def is_uniform_with_collection( self, kv_cache_specs: dict[str, KVCacheSpec] ) -> bool: - """ - Whether this KVCacheSpec is uniform with all specs of all layers. - """ + """Whether this KVCacheSpec is uniform with all specs of all layers.""" uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self) assert uniform_type_base_spec is not None, ( f"Unsupported KV cache spec type: {type(self)}. " @@ -537,8 +529,7 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: @dataclass(frozen=True, kw_only=True) class FullAttentionSpec(AttentionSpec): - """ - When hybrid allocator is disabled and the model contains both full + """When hybrid allocator is disabled and the model contains both full attention layers and sliding window attention layers, sliding window attention are regarded as full attention in KV cache manager (blocks are allocated for all tokens), while computed as sliding window @@ -582,8 +573,7 @@ def merge_window_sizes(cls, window_sizes: set[int]) -> int | None: @classmethod def merge(cls, specs: list[Self]) -> Self: - """ - Merge a list of FullAttentionSpec objects into a single + """Merge a list of FullAttentionSpec objects into a single FullAttentionSpec object. """ assert all(isinstance(spec, FullAttentionSpec) for spec in specs), ( @@ -1124,9 +1114,7 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: @dataclass(frozen=True) class CrossAttentionSpec(AttentionSpec): - """ - KV cache spec for cross-attention layers in encoder-decoder models. - """ + """KV cache spec for cross-attention layers in encoder-decoder models.""" def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: # For cross-attention, we need to cache encoder states @@ -1141,8 +1129,7 @@ class SinkFullAttentionSpec(FullAttentionSpec): @classmethod def merge(cls, specs: list[Self]) -> Self: - """ - Merge a list of FullAttentionSpec objects into a single + """Merge a list of FullAttentionSpec objects into a single FullAttentionSpec object. """ assert all(isinstance(spec, FullAttentionSpec) for spec in specs), ( @@ -1192,8 +1179,7 @@ def merge(cls, specs: list[Self]) -> Self: @dataclass(frozen=True) class UniformTypeKVCacheSpecs(KVCacheSpec): - """ - A KV cache spec for multiple layers with the same type of attention. Here, + """A KV cache spec for multiple layers with the same type of attention. Here, same types means always need the same number of token slots. For example, sliding window attentions with different window sizes are not the same type and should not be merged into one UniformTypeKVCacheSpecs. @@ -1241,8 +1227,7 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: @classmethod def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: - """ - Whether all layers have the same type of KV cache spec. + """Whether all layers have the same type of KV cache spec. Uses the registry to determine grouping base classes, so custom specs that inherit from FullAttentionSpec are treated as full attention. @@ -1256,8 +1241,7 @@ def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: @classmethod def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None: - """ - Return a SameTypeKVCacheSpecs object if all layers have the same type + """Return a SameTypeKVCacheSpecs object if all layers have the same type of KV cache spec. Return None if not. """ if cls.is_uniform_type(kv_cache_specs): @@ -1268,7 +1252,8 @@ def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None: def get_max_layers_per_page_size(self) -> int: """Max number of layers sharing a page size. For a balanced bucket - this equals the number of repetitions of the layer pattern.""" + this equals the number of repetitions of the layer pattern. + """ return Counter( spec.page_size_bytes for spec in self.kv_cache_specs.values() ).most_common(1)[0][1] @@ -1368,8 +1353,7 @@ def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None: @dataclass class KVCacheTensor: - """ - A class for specifying how the workers should initialize the KV cache. + """A class for specifying how the workers should initialize the KV cache. Placement of a set of same-shaped layers in the KV cache allocation. Layer ``layers[l]``'s page for block ``b`` starts at @@ -1399,8 +1383,7 @@ class KVCacheGroupRole(str, Enum): @dataclass class KVCacheGroupSpec: - """ - Represents a group of model layers that share the same KV cache block table. + """Represents a group of model layers that share the same KV cache block table. These layers are regarded as one layer in the KV cache manager. """ @@ -1421,9 +1404,7 @@ class KVCacheGroupSpec: @dataclass class KVCacheConfig: - """ - The KV cache configuration of a model. - """ + """The KV cache configuration of a model.""" num_blocks: int """The number of KV cache blocks""" diff --git a/vllm/v1/kv_cache_layout.py b/vllm/v1/kv_cache_layout.py index 7c1decd393a9..d784f51b96ac 100644 --- a/vllm/v1/kv_cache_layout.py +++ b/vllm/v1/kv_cache_layout.py @@ -49,7 +49,8 @@ def is_block_contiguous(self) -> bool: @property def is_block_compact(self) -> bool: """True when each page's [H, N, C] bytes form one contiguous run; i.e. - the L and B dimensions are outermost.""" + the L and B dimensions are outermost. + """ return set(self.value[:2]) == {_DIM_L, _DIM_B} @property diff --git a/vllm/v1/kv_cache_spec_registry.py b/vllm/v1/kv_cache_spec_registry.py index 62cb291892a2..3a03fb376ca4 100644 --- a/vllm/v1/kv_cache_spec_registry.py +++ b/vllm/v1/kv_cache_spec_registry.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Registry for KVCacheSpec types and their associated managers. +"""Registry for KVCacheSpec types and their associated managers. This module provides a pluggable architecture for registering custom KVCacheSpec subclasses without modifying vLLM core code. Out-of-tree platforms can define @@ -42,9 +41,7 @@ class KVCacheSpecRegistry: @classmethod def _ensure_registered(cls, vllm_config=None) -> None: - """ - Run full KVCacheSpec registration if the registration is not done. - """ + """Run full KVCacheSpec registration if the registration is not done.""" if _REGISTRY_KVCACHESPEC_LIST: return @@ -67,8 +64,7 @@ def register( manager_class: type["SingleTypeKVCacheManager"] | None = None, uniform_type_base_spec: type["KVCacheSpec"] | None = None, ) -> None: - """ - Register a KVCacheSpec class with its manager and base spec. + """Register a KVCacheSpec class with its manager and base spec. Args: kvcache_spec_cls: The KVCacheSpec subclass to register @@ -77,6 +73,7 @@ def register( instead of being grouped to different kvcache group, `kvcache_spec_cls` and `uniform_type_base_spec` will be trated as uniform type. If None, defaults to kvcache_spec_cls itself (for built-in base specs). + """ assert manager_class is not None, "manager_class is required" if uniform_type_base_spec is None: @@ -118,8 +115,7 @@ def register_role_manager( def get_manager_class( cls, kvcache_spec: "KVCacheSpec", role: str | None = None ) -> type["SingleTypeKVCacheManager"] | None: - """ - Get the single type kvcache manager class for a given kvcache spec instance. + """Get the single type kvcache manager class for a given kvcache spec instance. Args: kvcache_spec: A KVCacheSpec instance @@ -127,6 +123,7 @@ def get_manager_class( Returns: The SingleTypeKVCacheManager class to use for this kvcache_spec + """ cls._ensure_registered() if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)): @@ -144,8 +141,7 @@ def get_manager_class( def get_uniform_type_base_spec( cls, kvcache_spec: "KVCacheSpec" ) -> type["KVCacheSpec"] | None: - """ - Get the base kvcache spec class for grouping compatibility checks. + """Get the base kvcache spec class for grouping compatibility checks. KVCacheSpecs with uniform_type_base_spec will be trated as one group. Args: @@ -153,6 +149,7 @@ def get_uniform_type_base_spec( Returns: The base KVCacheSpec class for checking uniform type kvcache specs + """ cls._ensure_registered() kvcache_spec_cls = type(kvcache_spec) @@ -168,9 +165,7 @@ def get_uniform_type_base_spec( def check_kv_cache_spec_registry( cls, kv_cache_spec: dict[str, "KVCacheSpec"] ) -> None: - """ - Check if the KVCacheSpecs of each layer are registered as expected. - """ + """Check if the KVCacheSpecs of each layer are registered as expected.""" cls._ensure_registered() for layer_name, spec in kv_cache_spec.items(): # use raise instead of assert to make it effective in production environment @@ -192,8 +187,7 @@ def register_kv_cache_spec( manager_class: type["SingleTypeKVCacheManager"] | None = None, uniform_type_base_spec: type["KVCacheSpec"] | None = None, ): - """ - Decorator to register a custom KVCacheSpec class. + """Decorator to register a custom KVCacheSpec class. Args: manager_class: The SingleTypeKVCacheManager to use for this spec. @@ -211,6 +205,7 @@ def register_kv_cache_spec( @dataclass(frozen=True, kw_only=True) class CustomFullAttentionSpec(FullAttentionSpec): pass + """ def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]: diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 988747e06e1e..b22390e444d9 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Core abstractions for KV cache offloading in vLLM v1. -""" +"""Core abstractions for KV cache offloading in vLLM v1.""" from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence @@ -149,8 +147,7 @@ class ScheduleEndContext(NamedTuple): class LoadStoreSpec: - """ - Metadata that encapsulates information allowing a worker + """Metadata that encapsulates information allowing a worker to load, and optionally also to store, blocks of KV data. """ @@ -235,8 +232,7 @@ class OffloadingKVEventsConfig: class OffloadingManager(ABC): @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - """ - Checks whether a single block is offloaded and ready to be read. + """Checks whether a single block is offloaded and ready to be read. Args: key: the key identifying the block to lookup. @@ -246,6 +242,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: HIT if the block is offloaded and ready, MISS if not found, HIT_PENDING if found but not yet readable, or RETRY if the lookup should be retried later. + """ pass @@ -255,8 +252,7 @@ def prepare_load( keys: Collection[OffloadKey], req_context: ReqContext, ) -> LoadStoreSpec: - """ - Prepare the given blocks to be read. + """Prepare the given blocks to be read. The given blocks will be protected from eviction until complete_load is called. It assumes all given blocks are offloaded. @@ -268,27 +264,28 @@ def prepare_load( Returns: A LoadStoreSpec that can be used by a worker to locate and load the actual offloaded KV data. + """ pass def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Mark the given blocks as recently used. + """Mark the given blocks as recently used. This could in practice mean moving them to the end of an LRU list. Args: keys: the keys identifying the blocks. req_context: per-request context (e.g. kv_transfer_params). + """ return def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Marks previous blocks that were prepared to load as done loading. + """Marks previous blocks that were prepared to load as done loading. Args: keys: the keys identifying the blocks. req_context: per-request context (e.g. kv_transfer_params). + """ return @@ -298,8 +295,7 @@ def prepare_store( keys: Collection[OffloadKey], req_context: ReqContext, ) -> PrepareStoreOutput | None: - """ - Prepare the given blocks to be offloaded. + """Prepare the given blocks to be offloaded. The given blocks will be protected from eviction until complete_store is called. @@ -312,6 +308,7 @@ def prepare_store( where to store them (LoadStoreSpec), and list of blocks that were evicted as a result. None is returned if the blocks cannot be stored. + """ pass @@ -321,8 +318,7 @@ def complete_store( req_context: ReqContext, success: bool = True, ): - """ - Marks blocks which were previously prepared to be stored, as stored. + """Marks blocks which were previously prepared to be stored, as stored. Following this call, the blocks become loadable. If success is False, blocks that were not marked as stored will be removed. @@ -331,25 +327,25 @@ def complete_store( keys: the keys identifying the blocks. req_context: per-request context (e.g. kv_transfer_params). success: whether the blocks were stored successfully. + """ return @abstractmethod def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - """ - Called when a new request is first seen by the scheduler. + """Called when a new request is first seen by the scheduler. Returns a RequestOffloadingContext indicating how this request's blocks should be offloaded. Args: req_context: per-request context. + """ pass def on_request_finished(self, req_context: ReqContext) -> None: - """ - Called when a request has finished. + """Called when a request has finished. By the time this is called, the scheduler will issue no more submit-side calls for this request, such as prepare_store() and @@ -364,12 +360,12 @@ def on_request_finished(self, req_context: ReqContext) -> None: Args: req_context: per-request context. + """ return def take_events(self) -> Iterable[OffloadingEvent]: - """ - Take the offloading events from the manager. + """Take the offloading events from the manager. A tier manager emits only events for storage state it owns. A composing manager may aggregate child event streams, but should not @@ -377,6 +373,7 @@ def take_events(self) -> Iterable[OffloadingEvent]: Yields: New OffloadingEvents collected since the last call. + """ return () @@ -410,8 +407,7 @@ def shutdown(self) -> None: class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC): - """ - Spec for loading/storing KV blocks from given block numbers. + """Spec for loading/storing KV blocks from given block numbers. Subclass semantics differ: GPULoadStoreSpec.block_ids are GPU block indices; CPULoadStoreSpec.block_ids are CPU cache chunk indices. @@ -425,8 +421,7 @@ def __repr__(self) -> str: class GPULoadStoreSpec(BlockIDsLoadStoreSpec): - """ - Spec for loading/storing a KV block to GPU memory. + """Spec for loading/storing a KV block to GPU memory. If there are multiple KV groups, the blocks are expected to be ordered by the group index. @@ -460,8 +455,7 @@ def __init__( @dataclass class CanonicalKVCacheTensor: - """ - A canonicalized KV cache tensor whose first dimension is num_blocks. + """A canonicalized KV cache tensor whose first dimension is num_blocks. With standardized layouts (RFC #42082) num_blocks is always the leading logical dimension. @@ -478,7 +472,8 @@ class CopyRun: """A strided byte correspondence between this worker's physical page and a canonical page: for i in range(num_fragments), fragment i spans [local_offset + i * local_stride, +fragment_size) in the worker's page and - [canonical_offset + i * canonical_stride, +fragment_size) canonically.""" + [canonical_offset + i * canonical_stride, +fragment_size) canonically. + """ local_offset: int canonical_offset: int @@ -517,8 +512,7 @@ def is_writer(self, block_id: int) -> bool: @dataclass class CanonicalKVCacheRef: - """ - Per-layer (or group of layers) reference to a specific (by index) + """Per-layer (or group of layers) reference to a specific (by index) CanonicalKVCacheTensor and records the un-padded page size used by that layer. """ @@ -532,8 +526,7 @@ class CanonicalKVCacheRef: @dataclass class CanonicalKVCaches: - """ - Canonicalized block-level representation of the KV caches. + """Canonicalized block-level representation of the KV caches. Composed of: - Unique list of KV cache data tensors, @@ -561,7 +554,8 @@ class TransferResult: class OffloadingWorker(ABC): """Runs in the worker process. Performs async KV transfers for ONE offloaded medium (e.g. CPU). Direction is explicit via submit_store / - submit_load, so there is no (src_medium, dst_medium) routing.""" + submit_load, so there is no (src_medium, dst_medium) routing. + """ @abstractmethod def submit_store( @@ -620,8 +614,7 @@ def __init__(self, config: OffloadingConfig): @abstractmethod def get_manager(self) -> OffloadingManager: - """ - Get an OffloadingManager that will be used + """Get an OffloadingManager that will be used by the scheduler-side offloading connector to track offloaded blocks and manage evictions. """ @@ -629,13 +622,13 @@ def get_manager(self) -> OffloadingManager: @abstractmethod def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: - """ - Get an OffloadingWorker that handles async KV transfers for this spec. + """Get an OffloadingWorker that handles async KV transfers for this spec. Args: kv_caches: Canonicalized KV caches. Returns: An OffloadingWorker instance for this medium. + """ pass diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 8db9e805ff1a..b16fdeef6c36 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -80,8 +80,7 @@ def compute_sub_block_ptrs( tensor: torch.Tensor, skip_count: int = 0, ): - """ - Compute byte pointers for sub-blocks of the given block IDs. + """Compute byte pointers for sub-blocks of the given block IDs. Each block in block_ids contains blocks_per_chunk sub-blocks. The pointer for sub-block j of block b is: @@ -98,6 +97,7 @@ def compute_sub_block_ptrs( output: pre-allocated pointer array to write pointers into. tensor: the source or destination tensor. skip_count: sub-blocks to skip in the first block. + """ assert skip_count < blocks_per_chunk @@ -126,7 +126,8 @@ def compute_sub_block_ptrs( class CopyPlan(NamedTuple): """Precomputed fragment-copy template for one data ref under the canonical CPU layout, unrolled from the ref's mapped runs. Offsets are relative to - the per-block base pointers on each side.""" + the per-block base pointers on each side. + """ frag_offsets_src: np.ndarray frag_offsets_dst: np.ndarray @@ -164,7 +165,8 @@ def _canonical_page_ids( ) -> np.ndarray: """Global canonical page ids matching compute_sub_block_ptrs' enumeration. These identify canonical pages consistently across ranks, so they key - CanonicalPageMapping.is_writer rotation.""" + CanonicalPageMapping.is_writer rotation. + """ if blocks_per_chunk == 1: return block_ids[:count] flat = ( @@ -177,7 +179,8 @@ def _canonical_block_sizes( layer_refs_per_group: list[list[CanonicalKVCacheRef]], num_tensors: int ) -> list[int]: """Canonical CPU bytes per GPU block for each tensor, taken from the refs' - mappings. Requires every ref to carry a mapping.""" + mappings. Requires every ref to carry a mapping. + """ canonical_bytes_per_block = [0] * num_tensors for layer_refs in layer_refs_per_group: for ref in layer_refs: @@ -234,8 +237,7 @@ def _new_descriptor_buffers( class SingleDirectionOffloadingHandler: - """ - Handles transfers for a single direction, either CPU->GPU or GPU->CPU. + """Handles transfers for a single direction, either CPU->GPU or GPU->CPU. Transfers are guaranteed to be executed in order of their submission. Each transfer uses a unique CUDA stream, and its stream will start executing only after the streams of previous transfers have finished. @@ -250,8 +252,7 @@ def __init__( gpu_to_cpu: bool, canonical_layout: bool = False, ): - """ - Initialize a SingleDirectionOffloadingHandler. + """Initialize a SingleDirectionOffloadingHandler. Args: gpu_tensors: list of GPU KV cache tensors. @@ -259,10 +260,12 @@ def __init__( cpu_tensors: list of CPU KV cache tensors. Each of shape (num_cpu_chunks, cpu_page_size_bytes) with dtype int8. Order should match gpu_tensors. + blocks_per_chunk: number of blocks transferred per chunk. layer_refs_per_group: list of CanonicalKVCacheRef per group. gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU. canonical_layout: if True, CPU pages use the canonical layout described by the refs' mappings. + """ assert len(gpu_tensors) == len(cpu_tensors) assert len(gpu_tensors) > 0 @@ -340,7 +343,8 @@ def _estimate_max_copy_ops(self, group_sizes: Sequence[int]) -> int: """Upper bound on the number of copy descriptors for a transfer. Exact for the direct layout. The canonical path may fill fewer: - writer rotation later drops the blocks this rank does not write.""" + writer rotation later drops the blocks this rank does not write. + """ num_copy_ops = 0 for g_idx, (group_size, layer_refs) in enumerate( zip(group_sizes, self.layer_refs_per_group) @@ -369,7 +373,8 @@ def _fill_direct_ops( """Fill one group's copy descriptors for the direct (worker-private) layout: one whole-page copy per (block, ref). - Returns (op_idx past the filled descriptors, bytes added).""" + Returns (op_idx past the filled descriptors, bytes added). + """ num_bytes = 0 for data_ref in self.layer_refs_per_group[g_idx]: t_idx = data_ref.tensor_idx @@ -412,7 +417,8 @@ def _fill_canonical_ops( scatter each block through the ref's precomputed CopyPlan, keeping only the blocks this rank writes. - Returns (op_idx past the filled descriptors, bytes added).""" + Returns (op_idx past the filled descriptors, bytes added). + """ assert self._canonical_copy_plans is not None # Zero-copy reinterpretation for pointer arithmetic: uint64 and the # buffers' int64 are bit-equivalent for addresses @@ -499,7 +505,8 @@ def _filter_writer_blocks( ) -> tuple[np.ndarray, np.ndarray]: """Keep only the blocks this rank writes: replicated ranks take turns writing shared canonical pages, keyed by the rank-consistent CPU-side - canonical page id.""" + canonical page id. + """ cpu_page_ids = _canonical_page_ids( group_dst, self.dst_blocks_per_chunk, diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 9d85090f2d8f..e42cb47f2530 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -44,8 +44,7 @@ class _RequestCacheAccess: class CPUOffloadingManager(OffloadingManager): - """ - An OffloadingManager with a pluggable CachePolicy, resolved by name via + """An OffloadingManager with a pluggable CachePolicy, resolved by name via CachePolicyFactory (built in: "lru", "arc"; external policies can either register their own or be loaded out-of-tree via cache_policy_module_path). diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index 882ee02219c2..c3942e90d850 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -14,8 +14,7 @@ class ARCCachePolicy(CachePolicy): - """ - ARC (Adaptive Replacement Cache) cache policy. + """ARC (Adaptive Replacement Cache) cache policy. Data Structures: T1: Recent cache containing chunks accessed once. diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 3efffbd3961a..d3ef8b859a9b 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -23,8 +23,7 @@ def order_request_keys( class ChunkStatus(ctypes.Structure): - """ - Offloading status for a single chunk of KV data. + """Offloading status for a single chunk of KV data. Holds the following information: ref_cnt - the current number of transfers using this chunk as a source. @@ -42,15 +41,12 @@ def __init__(self, chunk_id: int): @property def is_ready(self) -> bool: - """ - Returns whether the chunk is ready to be read. - """ + """Returns whether the chunk is ready to be read.""" return self.ref_cnt >= 0 class CachePolicy(ABC): - """ - Encapsulates both chunk organization (data structures) and replacement + """Encapsulates both chunk organization (data structures) and replacement decisions (which chunk to evict). LRU and ARC differ in both dimensions — ARC's ghost lists and target_t1_size live at the intersection of storage and eviction, so they cannot be separated cleanly. @@ -73,12 +69,12 @@ def remove(self, key: OffloadKey) -> None: @abstractmethod def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: - """ - Mark chunks as recently used. + """Mark chunks as recently used. Args: keys: Chunks to mark as recently used. req_context: Per-request context for the request touching these chunks. + """ def on_store_miss( @@ -120,8 +116,7 @@ def on_request_finished( def evict( self, n: int, protected: set[OffloadKey] ) -> list[tuple[OffloadKey, ChunkStatus]] | None: - """ - Evict exactly n chunks, skipping any in protected. + """Evict exactly n chunks, skipping any in protected. Returns a list of (key, chunk) for the evicted chunks, or None if n evictions cannot be satisfied. The operation is atomic: @@ -133,8 +128,7 @@ def evict( @abstractmethod def clear(self) -> None: - """ - Remove ALL chunks regardless of ref_cnt. + """Remove ALL chunks regardless of ref_cnt. Ghost lists and adaptive state are also reset. """ diff --git a/vllm/v1/kv_offload/cpu/policies/factory.py b/vllm/v1/kv_offload/cpu/policies/factory.py index eedcc6807697..e4f811d49f0e 100644 --- a/vllm/v1/kv_offload/cpu/policies/factory.py +++ b/vllm/v1/kv_offload/cpu/policies/factory.py @@ -57,6 +57,7 @@ def get_cache_policy_cls( Raises ValueError if the cache policy is neither registered nor resolvable via `module_path`. + """ if name in cls._registry: return cls._registry[name]() diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index 5b5f65dec779..6ca9cb32f42a 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -14,8 +14,7 @@ class LRUCachePolicy(CachePolicy): - """ - LRU caching policy with logical recency independent of transfer pinning. + """LRU caching policy with logical recency independent of transfer pinning. Evictable chunks live in a lazy-invalidating min-heap. A chunk's recency can therefore be updated while it is pinned; when it later becomes diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 284f1de68e8e..af183fbc4f06 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -69,8 +69,7 @@ def _get_populate_write_fn( class SharedOffloadRegion: - """ - Single mmap-backed memory region shared across all workers for a + """Single mmap-backed memory region shared across all workers for a vLLM instance. Workers coordinate via the filesystem: the first worker to open the file with O_EXCL becomes the creator and calls ftruncate; the rest open the existing file and wait until it reaches the expected @@ -264,6 +263,7 @@ def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: Args: tensor_page_size: Bytes per chunk for this tensor. + """ assert self.rank is not None new_offset = self._worker_offset + tensor_page_size @@ -313,6 +313,7 @@ def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: Args: tensor_page_size: Canonical bytes per chunk for this tensor. + """ new_offset = self._canonical_offset + tensor_page_size assert new_offset <= self._row_stride diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 00b36e863a07..9eba094d79f3 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -28,7 +28,8 @@ def _all_workers_barrier() -> None: """Block until every worker rank has reached this point (gloo cpu group). A superset of the node-local mmap openers suffices: once the barrier - releases, every worker sharing the region file has mapped it.""" + releases, every worker sharing the region file has mapped it. + """ from vllm.distributed.parallel_state import ( get_inner_dp_world_group, get_world_group, @@ -161,7 +162,8 @@ def get_manager(self) -> OffloadingManager: def _uses_shared_region(self) -> bool: """Whether the worker CPU buffer is the shared mmap region (vs a private - per-rank tensor); replicated-layout dedup is gated on this being True.""" + per-rank tensor); replicated-layout dedup is gated on this being True. + """ return current_platform.is_cuda_alike() def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index 82ebe31e2d33..1ae0d2d1fb07 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -19,9 +19,7 @@ class FileMapper: - """ - FileMapper maps KV blocks (given by their hash) to file names. - """ + """FileMapper maps KV blocks (given by their hash) to file names.""" def __init__( self, @@ -41,8 +39,7 @@ def __init__( replicated_layout: bool = False, canonical_format: str | None = None, ): - """ - Initialize the file mapper. Each worker constructs its own, but + """Initialize the file mapper. Each worker constructs its own, but `config.json` is shared across workers since rank lives outside the hash. When `parallel_agnostic=True`, tp/pp/pcp/dcp are forced to 1 and rank to 0 so multiple parallelism layouts collapse into the same folder. @@ -136,8 +133,7 @@ def get_config_file_path(self) -> str: @staticmethod def _compute_base_path(root_dir: str, fields: dict) -> str: - """ - Layout: /_/. + """Layout: /_/. safe_model_name replaces '/' with '_' so HuggingFace IDs don't nest. """ canonical = json.dumps(fields, sort_keys=True, separators=(",", ":")) diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index c08c63d01875..93e77591fab2 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -AsyncLookupManager: per-tier async lookup manager for secondary tier +"""AsyncLookupManager: per-tier async lookup manager for secondary tier existence checks. Each secondary tier that wants non-blocking lookups composes its own @@ -65,8 +64,7 @@ class LookupState: class AsyncLookupManager(ABC): - """ - Per-tier async lookup manager for secondary tier existence checks. + """Per-tier async lookup manager for secondary tier existence checks. Each secondary tier that wants non-blocking lookups composes its own AsyncLookupManager instance internally. The manager maintains lookup @@ -126,8 +124,7 @@ def __init__( def batch_lookup( self, keys: list[OffloadKey], req_context: ReqContext ) -> Iterable[bool]: - """ - Check whether a batch of blocks exist in this tier. + """Check whether a batch of blocks exist in this tier. Called from the worker thread — must be synchronous and must not touch the primary tier or scheduler state. @@ -141,13 +138,13 @@ def batch_lookup( # ------------------------------------------------------------------ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - """ - Non-blocking lookup called from the scheduler thread. + """Non-blocking lookup called from the scheduler thread. Returns: True — block is present in this tier. False — block is not present in this tier. None — result not yet available; retry next step. + """ if self._need_to_drain: self.drain_results() @@ -219,7 +216,8 @@ def drain_results(self) -> None: def mark_miss(self, keys: Collection[OffloadKey]) -> None: """Force the cached verdict for ``keys`` to False after a failed load, so the scheduler stops re-issuing the doomed promotion (livelock, #49176). - Keys with no cached entry are skipped.""" + Keys with no cached entry are skipped. + """ for key in keys: state = self._lookup_state.get(key) if state is not None: diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 34aa6ddf78ad..77300fa707b0 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Abstract interfaces and data types for the secondary tiering layer. -""" +"""Abstract interfaces and data types for the secondary tiering layer.""" from abc import ABC, abstractmethod from collections.abc import Collection, Iterable @@ -119,8 +117,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: ... class SecondaryTierManager(ABC): - """ - Abstract interface for managing a single non-primary offloading tier. + """Abstract interface for managing a single non-primary offloading tier. Secondary tiers cannot directly access GPU memory. All data transfers must go through the CPU (primary) tier: @@ -140,12 +137,12 @@ def __init__( primary_kv_view: memoryview, tier_type: str, ) -> None: - """ - Args: - offloading_spec: Offloading configuration. - primary_kv_view: Memoryview of the primary tier's CPU KV cache. - tier_type: Tier type identifier, set by SecondaryTierFactory - from the registered tier type. + """Args: + offloading_spec: Offloading configuration. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory + from the registered tier type. + """ self._offloading_spec = offloading_spec self._primary_kv_view: memoryview = primary_kv_view @@ -154,8 +151,7 @@ def __init__( @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - """ - Check whether a chunk exists in this secondary tier. + """Check whether a chunk exists in this secondary tier. Args: key: Offload key to look up. @@ -165,13 +161,13 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: HIT if the chunk is present and ready, MISS if not found, or RETRY if the chunk is being transferred (retry later). + """ pass @abstractmethod def submit_store(self, job_metadata: TransferJob) -> None: - """ - Submit an async job to store chunks from the primary tier to this + """Submit an async job to store chunks from the primary tier to this secondary tier. This method must be lightweight and non-blocking: allocate metadata @@ -193,13 +189,13 @@ def submit_store(self, job_metadata: TransferJob) -> None: Args: job_metadata: Job metadata including job_id, keys, and chunk_ids identifying the primary-tier slots to read from. + """ pass @abstractmethod def submit_load(self, job_metadata: TransferJob) -> None: - """ - Submit an async job to load chunks from this secondary tier to the + """Submit an async job to load chunks from this secondary tier to the primary tier. This method must be lightweight and non-blocking: mark chunks as @@ -218,13 +214,13 @@ def submit_load(self, job_metadata: TransferJob) -> None: Args: job_metadata: Job metadata including job_id, keys, and chunk_ids identifying the primary-tier slots to write into. + """ pass @abstractmethod def get_finished_jobs(self) -> Iterable[JobResult]: - """ - Return all jobs (loads and stores) that completed since the last call. + """Return all jobs (loads and stores) that completed since the last call. The framework uses these results to release resources and finalize transfers. @@ -232,6 +228,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]: Returns: Iterable of JobResult objects for jobs finished since the last call. + """ pass @@ -248,31 +245,30 @@ def take_events(self) -> Iterable[OffloadingEvent]: return () def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Mark chunks as recently used for eviction policy. + """Mark chunks as recently used for eviction policy. Args: keys: Offload keys to mark as recently used. req_context: Per-request context. + """ return @abstractmethod def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - """ - Called when a new request is first seen by the scheduler. + """Called when a new request is first seen by the scheduler. Returns a RequestOffloadingContext expressing this tier's preference for how chunks should be offloaded for this request. Args: req_context: Per-request context. + """ pass def on_request_finished(self, req_context: ReqContext) -> None: - """ - Called when a request has finished. + """Called when a request has finished. By the time this is called, all per-request calls for this request (submit_store, submit_load, touch) have already been issued, and none @@ -283,6 +279,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: Args: req_context: per-request context. + """ return @@ -300,6 +297,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: Args: context: Per-step context from the scheduler. + """ return diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 967f7cc49826..aef2b7328de3 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ExampleSecondaryTierManager: A simple in-memory secondary tier. +"""ExampleSecondaryTierManager: A simple in-memory secondary tier. This implementation provides a minimal secondary tier that stores chunks in memory (using a dictionary) with immediate completion. It serves as a @@ -35,8 +34,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): - """ - A simple in-memory secondary tier. + """A simple in-memory secondary tier. This implementation: - Stores chunks in a dictionary (key -> True) @@ -52,11 +50,14 @@ def __init__( tier_type: str, custom_param: int = 0, ): - """ - Initialize the example secondary tier. + """Initialize the example secondary tier. Args: + offloading_spec: The offloading spec this tier belongs to. + primary_kv_view: Memoryview over the primary tier's KV buffer. + tier_type: Name identifying this tier type. custom_param: Dummy parameter demonstrating custom args. + """ super().__init__( offloading_spec=offloading_spec, @@ -78,8 +79,7 @@ def __init__( @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - """ - Check whether a chunk exists in this secondary tier. + """Check whether a chunk exists in this secondary tier. Args: key: Offload key to look up. @@ -87,17 +87,18 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: Returns: HIT if the chunk is present, MISS if not found. + """ return LookupResult.HIT if key in self.chunks else LookupResult.MISS @override def submit_store(self, job_metadata: TransferJob) -> None: - """ - Submit a job to store chunks from primary tier to this tier. + """Submit a job to store chunks from primary tier to this tier. Args: job_metadata: Job metadata including job_id, keys, and spec for reading chunks from the primary tier. + """ keys = job_metadata.keys chunk_ids = job_metadata.chunk_ids @@ -118,12 +119,12 @@ def submit_store(self, job_metadata: TransferJob) -> None: @override def submit_load(self, job_metadata: TransferJob) -> None: - """ - Submit a job to load chunks from this tier to primary tier. + """Submit a job to load chunks from this tier to primary tier. Args: job_metadata: Job metadata including job_id, keys, and spec for writing chunks into the primary tier. + """ keys = job_metadata.keys chunk_ids = job_metadata.chunk_ids @@ -149,12 +150,12 @@ def submit_load(self, job_metadata: TransferJob) -> None: @override def get_finished_jobs(self) -> Iterable[JobResult]: - """ - Poll for finished jobs. + """Poll for finished jobs. Returns: Iterable of JobResult objects for all jobs that have finished since the last call. + """ result = self.completed_jobs self.completed_jobs = [] @@ -167,7 +168,8 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @override def drain_jobs(self) -> None: """Synchronous tier — submit_*() returns only after the operation - completes, so there is nothing to wait for.""" + completes, so there is nothing to wait for. + """ return def get_num_chunks(self) -> int: diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 3310167b870b..086392a4065b 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -71,6 +71,7 @@ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]: Raises: ValueError: If 'type' is missing, or is neither registered nor resolvable via module_path. + """ tier_type = tier_config.get("type") if not tier_type: diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py index 9943c90b5472..53249a0ccd9a 100644 --- a/vllm/v1/kv_offload/tiering/fs/io.py +++ b/vllm/v1/kv_offload/tiering/fs/io.py @@ -94,9 +94,7 @@ def _store_block( block_size: int, use_o_direct: bool = True, ) -> None: - """ - Store callback: Writes to a temp file then atomically replaces the destination. - """ + """Store callback: write to a temp file then atomically replace the target.""" # Check if block already exists to avoid redundant writes if os.path.exists(dest_path): return @@ -141,7 +139,8 @@ def _load_block( ) -> None: """Read one KV block from disk; remove the file only on a provable short read (a too-short file is genuine corruption) and leave it untouched on any - other error.""" + other error. + """ fd: int | None = None view_slice = view.cast("B")[offset : offset + block_size] o_direct = O_DIRECT if use_o_direct else 0 @@ -172,8 +171,7 @@ def batch_store_block( block_size: int, use_o_direct: bool = True, ) -> None: - """ - Store a batch of KV blocks from a shared buffer to disk in one call. + """Store a batch of KV blocks from a shared buffer to disk in one call. Each block buffer[offsets[i] : offsets[i]+block_size] is written atomically to dest_paths[i] via a temp-file rename. Raises on first error. @@ -197,8 +195,7 @@ def batch_load_block( block_size: int, use_o_direct: bool = True, ) -> None: - """ - Load a batch of KV blocks from disk into a shared buffer in one call. + """Load a batch of KV blocks from disk into a shared buffer in one call. Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size]. Raises on first error (see _load_block for the delete-on-short-read policy). diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index f73f70e77ffb..ef545f4e3c2c 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -FileSystemTierManager: Pure-Python file system secondary tier for KV cache offloading. +"""FileSystemTierManager: pure-Python filesystem tier for KV cache offloading. Store path: Data is written to a temp file () via os.write, @@ -84,8 +83,7 @@ def batch_lookup( class FileSystemTierManager(SecondaryTierManager): - """ - Pure-Python disk-backed secondary tier. + """Pure-Python disk-backed secondary tier. Read-priority threads service load jobs preferentially; write-priority threads service store jobs preferentially. Both groups can drain either @@ -118,20 +116,20 @@ def __init__( enable_kv_events: bool = False, locality: str | None = None, ): - """ - Args: - offloading_spec: Contains normalized offloading configuration and - blocks_per_chunk. - primary_kv_view: Memoryview of the primary tier's CPU KV cache. - tier_type: Tier type identifier, set by SecondaryTierFactory. - root_dir: Root directory for block files. - n_read_threads: Number of read-priority I/O threads. - n_write_threads: Number of write-priority I/O threads. - enable_kv_events: Emit BlockStored KV events for blocks - successfully stored to this tier. Effective only when KV - cache events are enabled globally (kv_events_config). - locality: Whether this tier's storage is LOCAL or REMOTE relative - to the publishing vLLM instance. + """Args: + offloading_spec: Contains normalized offloading configuration and + blocks_per_chunk. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory. + root_dir: Root directory for block files. + n_read_threads: Number of read-priority I/O threads. + n_write_threads: Number of write-priority I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). + locality: Whether this tier's storage is LOCAL or REMOTE relative + to the publishing vLLM instance. + """ super().__init__(offloading_spec, primary_kv_view, tier_type) self.locality = Locality(locality) if locality is not None else None @@ -271,7 +269,8 @@ def load_task() -> None: @override def get_finished_jobs(self) -> Iterable[JobResult]: """Collect finished jobs; a failed promotion marks only its failed keys - as a miss here (scheduler thread).""" + as a miss here (scheduler thread). + """ results = [] for job_id, success, transfer_time in self._pool.get_finished(): if self.events is not None: @@ -333,8 +332,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: @override def shutdown(self) -> None: - """ - Release resources held by this tier. + """Release resources held by this tier. Shuts down the lookup manager and the thread pool, clearing pending tasks and waiting for active threads to complete. diff --git a/vllm/v1/kv_offload/tiering/fs/thread_pool.py b/vllm/v1/kv_offload/tiering/fs/thread_pool.py index 9db0629192bd..3c246708cffb 100644 --- a/vllm/v1/kv_offload/tiering/fs/thread_pool.py +++ b/vllm/v1/kv_offload/tiering/fs/thread_pool.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Thread pool: - Two queues (load, store) and two sets of threads: - - Load-priority threads: drain the load queue first, then the store queue. - - Store-priority threads: drain the store queue first, then the load queue. - Load jobs are enqueued to the load queue; store jobs to the store queue. +"""Thread pool: +Two queues (load, store) and two sets of threads: +- Load-priority threads: drain the load queue first, then the store queue. +- Store-priority threads: drain the store queue first, then the load queue. +Load jobs are enqueued to the load queue; store jobs to the store queue. """ import threading @@ -20,8 +19,7 @@ class JobState: - """ - Thread-safe completion tracker for a set of per-block I/O tasks. + """Thread-safe completion tracker for a set of per-block I/O tasks. Each task calls task_done(success) when it finishes. """ @@ -50,7 +48,7 @@ def job_id(self) -> JobId: def task_done( self, success: bool, transfer_time: float ) -> tuple[bool, bool, float]: - """Returns if job completed and success flag""" + """Returns if job completed and success flag.""" with self._lock: self._completed += 1 self._transfer_time += transfer_time @@ -60,8 +58,7 @@ def task_done( class DualQueueThreadPool: - """ - Thread pool with two task queues (load and store) and two thread groups. + """Thread pool with two task queues (load and store) and two thread groups. Load-priority threads drain the load queue first, then fall back to the store queue. Store-priority threads do the reverse. Both queues share diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f19616828daa..ba8dee8fadbe 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -TieringOffloadingManager: Multi-tier KV cache offloading orchestrator. +"""TieringOffloadingManager: Multi-tier KV cache offloading orchestrator. This manager coordinates between a CPU primary tier (with direct GPU access) and zero or more secondary tiers (Storage, Network, etc.) to provide @@ -144,7 +143,8 @@ def shutdown(self) -> None: class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the - TieringOffloadingManager with exclude_tier_idx set to the origin tier.""" + TieringOffloadingManager with exclude_tier_idx set to the origin tier. + """ __slots__ = ("_m", "_origin_idx") @@ -174,8 +174,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: class TieringOffloadingManager(OffloadingManager): - """ - Orchestrates multi-tier KV cache offloading. + """Orchestrates multi-tier KV cache offloading. This manager coordinates between a CPU primary tier (with direct GPU access) and zero or more secondary tiers (Storage, Network, etc.) to provide @@ -194,13 +193,13 @@ def __init__( primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, ): - """ - Initialize the TieringOffloadingManager. + """Initialize the TieringOffloadingManager. Args: primary_tier: The primary tier manager (CPU-based). secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. + """ self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] @@ -259,8 +258,7 @@ def _pop_job(self, job_id: JobId) -> JobMetadata | None: return self._jobs.pop(job_id, None) def _maybe_process_finished_jobs(self): - """ - Poll secondary tiers for completed jobs (at most once per step). + """Poll secondary tiers for completed jobs (at most once per step). Guarded by _processed_jobs_this_step: the first call in an engine step does the actual polling; subsequent calls are no-ops. The flag is reset @@ -306,8 +304,7 @@ def _complete_promotion( ) def _process_finished_jobs(self): - """ - Unconditionally poll all secondary tiers for completed jobs. + """Unconditionally poll all secondary tiers for completed jobs. This method: 1. Calls get_finished_jobs() on each secondary tier @@ -350,8 +347,7 @@ def lookup( *, exclude_tier_idx: int | None = None, ) -> LookupResult: - """ - Check whether a single chunk is offloaded and ready. + """Check whether a single chunk is offloaded and ready. Algorithm: 1. Process any completed async jobs first. @@ -362,6 +358,7 @@ def lookup( Args: key: Chunk hash to look up. req_context: Per-request context. + exclude_tier_idx: Skip this tier index during the lookup. Returns: HIT — chunk is ready in the primary tier. @@ -370,6 +367,7 @@ def lookup( RETRY — promotion started or a secondary tier is busy. MISS — chunk not found in any tier, or primary is full and cannot accept a promotion. + """ # Poll first so a promotion that finished since the last call is # already reflected as HIT (not stale HIT_PENDING/MISS) below, and @@ -432,8 +430,7 @@ def _initiate_promotion( key: OffloadKey, req_context: ReqContext, ) -> bool: - """ - Queue a chunk for promotion from a secondary tier to the primary tier. + """Queue a chunk for promotion from a secondary tier to the primary tier. Allocates space in the primary tier immediately (sets ref_cnt=-1 so subsequent lookups within the same step see the slot as in-flight), @@ -448,6 +445,7 @@ def _initiate_promotion( Returns: True if promotion was initiated, False if primary tier is full. + """ # Allocate space in primary tier for promoted chunk. # Must happen immediately so primary.lookup() returns None (in-flight) @@ -505,8 +503,7 @@ def _flush_pending_promotions(self) -> None: def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: - """ - Prepare chunks to be loaded from primary tier to GPU. + """Prepare chunks to be loaded from primary tier to GPU. Callers only pass keys already confirmed HIT by lookup() earlier this step. @@ -520,17 +517,18 @@ def prepare_load( Returns: LoadStoreSpec for reading from primary tier. + """ return self.primary_tier.prepare_load(keys, req_context) @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Mark chunks as recently used in all tiers. + """Mark chunks as recently used in all tiers. Args: keys: Chunks to mark as recently used. req_context: Per-request context. + """ self.primary_tier.touch(keys, req_context) for tier in self.secondary_tiers: @@ -538,8 +536,7 @@ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): @override def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Mark chunks as done loading from primary tier to GPU. + """Mark chunks as done loading from primary tier to GPU. This decrements ref_cnt on the chunks in the primary tier, allowing them to be evicted again. @@ -547,6 +544,7 @@ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): Args: keys: Chunks that finished loading. req_context: Per-request context. + """ self.primary_tier.complete_load(keys, req_context) @@ -554,8 +552,7 @@ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): def prepare_store( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: - """ - Prepare chunks to be stored from GPU to primary tier. + """Prepare chunks to be stored from GPU to primary tier. CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure that any completed async transfers have their ref_cnt decremented @@ -571,6 +568,7 @@ def prepare_store( Returns: PrepareStoreOutput describing where to store chunks and what was evicted, or None if store cannot proceed. + """ # Step 1: Poll for completed async jobs FIRST # _process_finished_jobs() handles two kinds of completions here: @@ -617,8 +615,7 @@ def _cascade_existing_chunks_to_request_level_tiers( req_context: ReqContext, request_level_tiers: set[int], ) -> None: - """ - For tiers that requested request-level policy, submit_store() for + """For tiers that requested request-level policy, submit_store() for chunks that are already present in the primary tier. A key whose primary write is still in flight (HIT_PENDING) cannot be @@ -674,8 +671,7 @@ def complete_store( req_context: ReqContext, success: bool = True, ) -> None: - """ - Mark chunks as done storing from GPU to primary tier. + """Mark chunks as done storing from GPU to primary tier. This is where secondary tier cascading happens — after chunks are confirmed to be in the primary tier, they are cascaded to ALL @@ -691,6 +687,7 @@ def complete_store( keys: Chunks that finished storing. success: Whether the GPU→primary transfer succeeded. req_context: Per-request context forwarded to primary.prepare_read(). + """ # Step 1: Complete store in primary tier (makes chunks loadable) self.primary_tier.complete_store(keys, req_context, success) @@ -748,8 +745,7 @@ def on_new_request( *, exclude_tier_idx: int | None = None, ) -> RequestOffloadingContext: - """ - Query each secondary tier for its offload policy preference. + """Query each secondary tier for its offload policy preference. Returns REQUEST_LEVEL if ANY secondary tier wants request-level. Only stores REQUEST_LEVEL tier decisions for use in prepare_store. @@ -858,6 +854,7 @@ def take_events(self) -> Iterable[OffloadingEvent]: Yields: New OffloadingEvents collected by each tier since the last call. + """ yield from self.primary_tier.take_events() for tier in self.secondary_tiers: diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 50a7e611f674..810a56c7ae85 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -113,19 +113,19 @@ def __init__( enable_kv_events: bool = False, locality: str | None = None, ): - """ - Args: - offloading_spec: Offloading configuration. - primary_kv_view: Memoryview of the primary tier's CPU KV cache. - tier_type: Tier type identifier, set by SecondaryTierFactory. - store_config: Object store connection parameters (see ObjStoreConfig). - prefix: Key prefix prepended to all object keys. - io_threads: Number of NIXL I/O threads. - enable_kv_events: Emit BlockStored KV events for blocks - successfully stored to this tier. Effective only when KV - cache events are enabled globally (kv_events_config). - locality: Whether this tier's storage is LOCAL or REMOTE relative - to the publishing vLLM instance. + """Args: + offloading_spec: Offloading configuration. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory. + store_config: Object store connection parameters (see ObjStoreConfig). + prefix: Key prefix prepended to all object keys. + io_threads: Number of NIXL I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). + locality: Whether this tier's storage is LOCAL or REMOTE relative + to the publishing vLLM instance. + """ super().__init__(offloading_spec, primary_kv_view, tier_type) self.locality = Locality(locality) if locality is not None else None @@ -308,7 +308,8 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: def _poll_active_transfers(self) -> None: """Poll all in-flight transfers once; move newly-completed (success or - failure) into ``_pending_results`` and release their NIXL handles.""" + failure) into ``_pending_results`` and release their NIXL handles. + """ for job_id, entry in list(self._transfers.items()): try: state = self._agent.check_xfer_state(entry.xfer_handle) @@ -363,7 +364,8 @@ def _poll_active_transfers(self) -> None: def get_finished_jobs(self) -> Iterable[JobResult]: """Poll transfers; a failed promotion marks its cached verdicts False - here (scheduler thread).""" + here (scheduler thread). + """ self._poll_active_transfers() results = self._pending_results self._pending_results = [] diff --git a/vllm/v1/kv_offload/tiering/p2p/control/base.py b/vllm/v1/kv_offload/tiering/p2p/control/base.py index 7016cfb162ae..5a7030737bdb 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/base.py @@ -139,6 +139,7 @@ def connect(self, peer_id: str) -> ControlConnection: A connection to peer_id that is registered but no longer alive is retired and replaced; a live one is a duplicate and must not be replaced. + """ ... @@ -158,6 +159,7 @@ def poll(self) -> Sequence[ControlConnection]: The returned sequence is read-only — callers must not mutate it. The caller is responsible for creating sessions for these connections. + """ ... diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index 42dc2c69846b..b28ed76dac97 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ZMQ-based transport layer for P2P KV cache sharing. +"""ZMQ-based transport layer for P2P KV cache sharing. Provides ZmqConnection (per-peer messaging) and ZmqTransport (connection management). Message-content agnostic. diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index ffe609d71c1e..62304b7544a3 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -87,6 +87,7 @@ class PollResult(NamedTuple): Attributes: done: Transfer IDs that completed successfully. failed: Transfer IDs that failed (error, timeout, etc.). + """ done: Sequence[int] @@ -174,6 +175,7 @@ def add_remote_peer( base_addr: Base address of the peer's block memory region. num_blocks: Number of blocks in the peer's region. block_len: Size of each block (must match local block_len). + """ ... @@ -203,6 +205,7 @@ def write_blocks( Returns: A unique transfer_id (int) to track this transfer, or None if the peer is not registered or submission failed. + """ ... @@ -224,6 +227,7 @@ def poll(self, peer_id: str | None = None) -> PollResult: Completed/failed transfers are removed from the inflight set. Must be called periodically to drive progress checking. + """ ... @@ -255,6 +259,7 @@ def cancel( For mode="wait", the subset of *transfer_ids* still tracked as inflight after the cancel attempt. For mode="immediate", always []. + """ ... diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index bf3adf4b604b..e8fb66dd6e30 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL. -""" +"""NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.""" from __future__ import annotations diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index ba1973b76784..17ecaf1a8241 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -P2PSecondaryTierManager: Secondary tier for P2P KV cache sharing. +"""P2PSecondaryTierManager: Secondary tier for P2P KV cache sharing. Owns transports and a single bidirectional P2PSession per remote peer. """ @@ -135,7 +134,8 @@ class P2PDestInfo: def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: """Parse the consumer sub-dict (PD ``remote_prefiller`` or symmetric - ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete.""" + ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete. + """ role = _remote_prefiller_params(kv_params) do_probe = False if role is None: @@ -156,7 +156,8 @@ def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: def _parse_dest(kv_params: dict | None) -> P2PDestInfo | None: """Parse the producer ``remote_decoder`` sub-dict into a ``P2PDestInfo``, - or None if the block is absent (not a remote-decode request).""" + or None if the block is absent (not a remote-decode request). + """ role = _remote_decoder_params(kv_params) if role is None: return None @@ -264,6 +265,7 @@ def __init__( Raises: ValueError: If ``unbound_store_timeout_s`` is not a positive number, or anything convertible to one. + """ super().__init__(offloading_spec, primary_kv_view, tier_type) try: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index ec4153e25b39..8369bbffe83d 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -P2P KV cache sharing protocol constants and documentation. +"""P2P KV cache sharing protocol constants and documentation. Protocol Overview ================= diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index ba2aca6438c6..f43e1a5aab5a 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -P2PSession — bidirectional session combining client + server roles. +"""P2PSession — bidirectional session combining client + server roles. A single P2PSession per remote peer handles BOTH directions of the P2P protocol on one ControlConnection: it can request blocks from the peer diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 9922a4095c4e..5cc438c34ffe 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -TieringOffloadingSpec: Spec for multi-tier KV cache offloading. +"""TieringOffloadingSpec: Spec for multi-tier KV cache offloading. This spec creates a TieringOffloadingManager with a CPU primary tier and configurable secondary tiers (e.g., Storage, Network). @@ -82,8 +81,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec): - """ - Spec for multi-tier KV cache offloading. + """Spec for multi-tier KV cache offloading. Creates a TieringOffloadingManager with: - Primary tier: CPU (LRU or ARC eviction policy) @@ -276,8 +274,7 @@ def __init__(self, config: OffloadingConfig): @override def get_manager(self) -> OffloadingManager: - """ - Get the TieringOffloadingManager. + """Get the TieringOffloadingManager. Creates a TieringOffloadingManager with: - Primary tier: CPU (LRU or ARC) @@ -285,6 +282,7 @@ def get_manager(self) -> OffloadingManager: Returns: TieringOffloadingManager instance + """ if not self._manager: if int(self.extra_config.get("store_threshold", 0)) >= 2: @@ -416,7 +414,8 @@ def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None: """Require a mapping on every ref and record layer portability. Fails loudly rather than persist direct-layout bytes under a - canonical format identity.""" + canonical format identity. + """ all_refs = [ ref for group_refs in kv_caches.group_data_refs for ref in group_refs ] diff --git a/vllm/v1/metrics/buckets.py b/vllm/v1/metrics/buckets.py index 375d73072924..7b9efb8f6dc2 100644 --- a/vllm/v1/metrics/buckets.py +++ b/vllm/v1/metrics/buckets.py @@ -177,6 +177,7 @@ def build_buckets(mantissa_lst: list[int], max_value: int) -> list[float]: Returns: Bucket values in increasing order, capped at `max_value`. + """ exponent = 0 buckets: list[float] = [] @@ -196,6 +197,7 @@ def build_1_2_5_buckets(max_value: int) -> list[float]: Example: >>> build_1_2_5_buckets(100) [1, 2, 5, 10, 20, 50, 100] + """ return build_buckets([1, 2, 5], max_value) @@ -218,6 +220,7 @@ def histogram_buckets( Raises: ValueError: If `family` is `request_tokens` and `max_model_len` is None. + """ if family == "request_tokens": if max_model_len is None: diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 3a8adb6e3244..0ff5876fd1ce 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -91,7 +91,8 @@ def load_stat_logger_plugin_factories() -> list[StatLoggerFactory]: class AggregateStatLoggerBase(StatLoggerBase): """Abstract base class for loggers that - aggregate across multiple DP engines.""" + aggregate across multiple DP engines. + """ @abstractmethod def __init__(self, vllm_config: VllmConfig, engine_indexes: list[int]): ... @@ -1202,16 +1203,15 @@ def log_engine_initialized(self): class StatLoggerManager: - """ - StatLoggerManager: - Logging happens at the level of the EngineCore (per scheduler). - * DP: >1 EngineCore per AsyncLLM - loggers for each EngineCore. - * With Local Logger, just make N copies for N EngineCores. - * With Prometheus, we need a single logger with N "labels" - - This class abstracts away this implementation detail from - the AsyncLLM, allowing the AsyncLLM to just call .record() - and .log() to a simple interface. + """StatLoggerManager: + Logging happens at the level of the EngineCore (per scheduler). + * DP: >1 EngineCore per AsyncLLM - loggers for each EngineCore. + * With Local Logger, just make N copies for N EngineCores. + * With Prometheus, we need a single logger with N "labels" + + This class abstracts away this implementation detail from + the AsyncLLM, allowing the AsyncLLM to just call .record() + and .log() to a simple interface. """ def __init__( diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index ebb676e66dcd..e1e61720b4a5 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Analytic flops/memory estimation module for transformer components, +"""Analytic flops/memory estimation module for transformer components, to help derive MFU (Model Flops Utilization) stats for a running model. """ @@ -33,8 +32,7 @@ class InvalidComponent(Exception): - """ - Custom exception to indicate that a certain ComponentMetric is not + """Custom exception to indicate that a certain ComponentMetric is not applicable to the given VllmConfig. """ @@ -99,8 +97,7 @@ class PerfStats: @dataclass class ExecutionContext: - """ - Represents an execution context for a batch of requests. + """Represents an execution context for a batch of requests. This class aggregates statistics across multiple requests in a batch, separately tracking prefill and decode phases. @@ -202,8 +199,7 @@ def from_single_request( class ParsedArgs: - """ - Syntactic sugar so that Parsers can use dot notations + """Syntactic sugar so that Parsers can use dot notations to access/update the parsed arguments. e.g.) @@ -227,16 +223,14 @@ def model_dump(self) -> dict[str, Any]: class Parser(Protocol): def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: - """ - Parse the vllm config and update the current ParsedArgs and pass it on. + """Parse the vllm config and update the current ParsedArgs and pass it on. If the parser isn't applicable to the vllm_config, it will do nothing. """ ... class ParserChain: - """ - Applies chain of parser in a sequential order. + """Applies chain of parser in a sequential order. Later parsers might overwrite results from previous parsers, so parsers should be chained in the appropriate order if they are not mutually exclusive. @@ -260,8 +254,7 @@ def parse(self, vllm_config: VllmConfig) -> ParsedArgs: class ComponentMetrics(BaseModel, ABC): - """ - Each concrete ComponentMetrics class is associated with: + """Each concrete ComponentMetrics class is associated with: - fields that are required for metric derivation (fields are specified/validated through pydantic model) - parser to parse VllmConfig into fields @@ -277,8 +270,7 @@ def component_type(cls) -> str: @classmethod @abstractmethod def get_parser(cls) -> ParserChain: - """ - Return a ParserChain that provides values for all required fields. + """Return a ParserChain that provides values for all required fields. The returned parser chain must populate ParsedArgs with values for every field defined on this ComponentMetrics class. Missing fields will cause a ValidationError when from_vllm_config() is called. @@ -292,11 +284,9 @@ def __init_subclass__(cls): @classmethod def from_vllm_config(cls, vllm_config: VllmConfig) -> Self: - """ - Instantiate this class from VllmConfig. + """Instantiate this class from VllmConfig. Raises ValidationError if parsing fails. """ - parser = cls.get_parser() parsed_args = parser.parse(vllm_config) try: @@ -343,8 +333,7 @@ def get_write_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int: class BaseConfigParser(Parser): - """ - Parses base model configuration. + """Parses base model configuration. Provides: vocab_size, hidden_size, num_attention_heads, num_hidden_layers, weight_byte_size, activation_byte_size, dp_size, tp_size, pp_size, enable_ep """ @@ -396,8 +385,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class BaseAttentionConfigParser(Parser): - """ - Parses attention-specific configuration. + """Parses attention-specific configuration. Provides: num_key_value_heads, head_dim, cache_byte_size """ @@ -418,8 +406,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class SlidingWindowAttentionParser(Parser): - """ - Parses sliding window attention configuration and layer breakdown. + """Parses sliding window attention configuration and layer breakdown. Provides: sliding_window, num_swa_layers """ @@ -466,8 +453,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class AttentionQuantizationConfigParser(Parser): - """ - Parses quantization configuration for attention layers. + """Parses quantization configuration for attention layers. Overrides: weight_byte_size """ @@ -490,8 +476,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class AttentionDetectionParser(Parser): - """ - Prevents standard AttentionMetrics from being instantiated for MLA models. + """Prevents standard AttentionMetrics from being instantiated for MLA models. MLA models should use MLAAttentionMetrics instead. """ @@ -589,6 +574,7 @@ def get_num_flops_breakdown( Returns: Dict with keys: qkv_proj, out_proj, attn_qk, attn_av. + """ D, q, kv, d = ( self.hidden_size, @@ -642,6 +628,7 @@ def get_read_bytes_breakdown( Returns: Dict with keys: qkv_input, qkv_weight, attn_input (conditional), out_input, out_weight. + """ D, q, kv, d = ( self.hidden_size, @@ -747,8 +734,7 @@ def get_write_bytes_breakdown( class MLADetectionParser(Parser): - """ - Validates that the model uses MLA attention. + """Validates that the model uses MLA attention. Raises InvalidComponent if the model does not use MLA, so MLAAttentionMetrics is silently skipped for non-MLA models. """ @@ -761,8 +747,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class MLAConfigParser(Parser): - """ - Parses MLA-specific configuration fields. + """Parses MLA-specific configuration fields. Provides: kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, q_lora_rank """ @@ -787,8 +772,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class MLAAttentionMetrics(ComponentMetrics): - """ - Performance metrics for Multi-Latent Attention (MLA) layers. + """Performance metrics for Multi-Latent Attention (MLA) layers. MLA uses a compressed latent representation for KV cache: - KV cache stores a single compressed vector of size @@ -1025,8 +1009,7 @@ def get_write_bytes_breakdown( class BaseFfnConfigParser(Parser): - """ - Parses FFN and MoE configuration. + """Parses FFN and MoE configuration. Provides: intermediate_size, num_experts, num_experts_per_tok, moe_intermediate_size, num_shared_experts, num_moe_layers """ @@ -1056,8 +1039,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class FfnParallelParser(Parser): - """ - Parses FFN parallelism configuration. + """Parses FFN parallelism configuration. Provides: ffn_tp_size, ffn_ep_size """ @@ -1077,8 +1059,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class InterleaveMoeLayerStepParser(Parser): - """ - Parses interleave_moe_layer_step field for models like Llama4. + """Parses interleave_moe_layer_step field for models like Llama4. Overrides: num_moe_layers """ @@ -1104,8 +1085,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class MoeLayerFreqParser(Parser): - """ - Parses moe_layer_freq and first_k_dense_replace fields for models like Deepseek. + """Parses moe_layer_freq and first_k_dense_replace fields for models like Deepseek. Overrides: num_moe_layers """ @@ -1129,8 +1109,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: class FfnQuantizationConfigParser(Parser): - """ - Parses quantization configuration for FFN layers. + """Parses quantization configuration for FFN layers. Overrides: weight_byte_size """ @@ -1486,11 +1465,9 @@ def get_write_bytes_breakdown( class ModelMetrics: def __init__(self, vllm_config: VllmConfig) -> None: - """ - Parse vllm_config to instantiate metrics for each component. + """Parse vllm_config to instantiate metrics for each component. is_enabled() will return False if no component metrics could be instantiated. """ - self.vllm_config = vllm_config self.metrics: list[ComponentMetrics] = [] @@ -1558,10 +1535,7 @@ def get_write_bytes_breakdown( def get_step_perf_stats_per_gpu( self, scheduler_output: SchedulerOutput ) -> PerfStats: - """ - Calculate perf stats for the current step based on scheduled tokens. - """ - + """Calculate perf stats for the current step based on scheduled tokens.""" t0 = time.monotonic() # Build a single batch context @@ -1846,7 +1820,8 @@ def get_required(obj: object, attr: str): def getattr_from_list(obj: object, attrs: list[str], default: object = None): """Try to get the first attr that exists in the object - from a list of attrs. Otherwise return None.""" + from a list of attrs. Otherwise return None. + """ for attr in attrs: if hasattr(obj, attr): return getattr(obj, attr) diff --git a/vllm/v1/metrics/prometheus.py b/vllm/v1/metrics/prometheus.py index c8740276713b..91bb88310d69 100644 --- a/vllm/v1/metrics/prometheus.py +++ b/vllm/v1/metrics/prometheus.py @@ -42,6 +42,7 @@ def get_prometheus_registry() -> CollectorRegistry: Returns: Registry: A prometheus registry + """ if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None: logger.debug("Using multiprocess registry for prometheus metrics") @@ -70,7 +71,6 @@ def unregister_vllm_metrics(): def shutdown_prometheus(): """Shutdown prometheus metrics.""" - path = _prometheus_multiproc_dir if path is None: return diff --git a/vllm/v1/metrics/ray_wrappers.py b/vllm/v1/metrics/ray_wrappers.py index 7e2100546e82..2d515d5963ad 100644 --- a/vllm/v1/metrics/ray_wrappers.py +++ b/vllm/v1/metrics/ray_wrappers.py @@ -68,8 +68,7 @@ def labels(self, *labels, **labelskwargs) -> "RayPrometheusMetric": @staticmethod def _get_sanitized_opentelemetry_name(name: str) -> str: - """ - For compatibility with Ray + OpenTelemetry, the metric name must be + """For compatibility with Ray + OpenTelemetry, the metric name must be sanitized. In particular, this replaces disallowed character (e.g., ':') with '_' in the metric name. Allowed characters: a-z, A-Z, 0-9, _ @@ -78,13 +77,13 @@ def _get_sanitized_opentelemetry_name(name: str) -> str: Ref: https://github.com/open-telemetry/opentelemetry-cpp/blob/main/sdk/src/metrics/instrument_metadata_validator.cc#L22-L23 Ref: https://github.com/ray-project/ray/blob/master/src/ray/stats/metric.cc#L107 """ - return re.sub(r"[^a-zA-Z0-9_]", "_", name) class RayGaugeWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Gauge to provide same API as - prometheus_client.Gauge""" + prometheus_client.Gauge + """ def __init__( self, @@ -118,7 +117,8 @@ def set_to_current_time(self): class RayCounterWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Counter to provide same API as - prometheus_client.Counter""" + prometheus_client.Counter + """ def __init__( self, @@ -143,7 +143,8 @@ def inc(self, value: int | float = 1.0): class RayHistogramWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Histogram to provide same API as - prometheus_client.Histogram""" + prometheus_client.Histogram + """ def __init__( self, @@ -169,8 +170,7 @@ def observe(self, value: int | float): class RaySpecDecodingProm(SpecDecodingProm): - """ - RaySpecDecodingProm is used by RayMetrics to log to Ray metrics. + """RaySpecDecodingProm is used by RayMetrics to log to Ray metrics. Provides the same metrics as SpecDecodingProm but uses Ray's util.metrics library. """ @@ -179,8 +179,7 @@ class RaySpecDecodingProm(SpecDecodingProm): class RayKVConnectorProm(KVConnectorProm): - """ - RayKVConnectorProm is used by RayMetrics to log Ray + """RayKVConnectorProm is used by RayMetrics to log Ray metrics. Provides the same metrics as KV connectors but uses Ray's util.metrics library. """ @@ -191,8 +190,7 @@ class RayKVConnectorProm(KVConnectorProm): class RayPerfMetricsProm(PerfMetricsProm): - """ - RayPerfMetricsProm is used by RayMetrics to log Ray + """RayPerfMetricsProm is used by RayMetrics to log Ray metrics. Provides the same MFU metrics as PerfMetricsProm uses Ray's util.metrics library. """ diff --git a/vllm/v1/metrics/reader.py b/vllm/v1/metrics/reader.py index 48c88e5b61cb..e847e862537a 100644 --- a/vllm/v1/metrics/reader.py +++ b/vllm/v1/metrics/reader.py @@ -82,6 +82,7 @@ def get_metrics_snapshot() -> list[Metric]: ... print(f" count = {metric.count}") ... for bucket_le, value in metrics.buckets.items(): ... print(f" {bucket_le} = {value}") + """ collected: list[Metric] = [] for metric in REGISTRY.collect(): diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 24e07e57cc7a..fe17208ce40a 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -34,9 +34,10 @@ class BaseCacheStats: class CachingMetrics: """Metrics for caching with a hit rate of the most recent N requests. + Args: - interval: The number of the most recent requests to aggregate. - Defaults to 1000. + max_recent_requests: The number of the most recent requests to aggregate. + """ def __init__(self, max_recent_requests: int = 1000) -> None: @@ -62,6 +63,7 @@ def observe(self, stats: BaseCacheStats): Args: stats: The prefix cache stats. + """ # reset_prefix_cache was invoked before the current update. # Reset the metrics before aggregating the current stats. @@ -113,8 +115,7 @@ def hit_rate(self) -> float: @dataclass class PrefixCacheStats(BaseCacheStats): - """ - Stores prefix cache hit statistics. + """Stores prefix cache hit statistics. - `reset`: Whether `reset_prefix_cache` was invoked. - `queries`: Refers to the number of tokens that were queried. """ @@ -144,8 +145,7 @@ def record(self, num_tokens: int, num_hits: int, preempted: bool) -> None: @dataclass class MultiModalCacheStats(BaseCacheStats): - """ - Stores multi-modal cache hit statistics. + """Stores multi-modal cache hit statistics. - `reset`: Whether `reset_mm_cache` was invoked. - `queries`: Refers to the number of multi-modal data items that were queried. diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index d8924b273881..35780dce9c6c 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -166,7 +166,6 @@ def empty_cpu( num_positions: int, num_tokens_per_position: int ) -> "LogprobsTensors": """Create empty LogprobsTensors on CPU.""" - logprob_token_ids = torch.empty( (num_positions, num_tokens_per_position), dtype=torch.int32, @@ -437,8 +436,7 @@ class DraftTokenIds: def make_empty_encoder_model_runner_output( scheduler_output: "SchedulerOutput", ) -> ModelRunnerOutput: - """ - Create a ModelRunnerOutput stub that contains the correct + """Create a ModelRunnerOutput stub that contains the correct per-request bookkeeping but no generated data yet. """ if not scheduler_output.num_scheduled_tokens: diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 8b453a09069e..ad875264afb4 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -348,8 +348,7 @@ def take_prefill_stats(self) -> PrefillStats | None: return prefill_stats def __lt__(self, other: "Request") -> bool: - """ - Compare two requests based on priority, arrival time, and request ID. + """Compare two requests based on priority, arrival time, and request ID. Used in priority scheduling. """ if self.priority != other.priority: diff --git a/vllm/v1/sample/logits_processor/__init__.py b/vllm/v1/sample/logits_processor/__init__.py index f319097495d0..ff9a6428c556 100644 --- a/vllm/v1/sample/logits_processor/__init__.py +++ b/vllm/v1/sample/logits_processor/__init__.py @@ -55,8 +55,7 @@ def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]: - """Load all installed logit processor plugins""" - + """Load all installed logit processor plugins.""" from importlib.metadata import entry_points installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP) @@ -171,6 +170,7 @@ def _load_custom_logitsprocs( Returns: A list of all loaded logitproc types + """ from vllm.platforms import current_platform @@ -238,7 +238,7 @@ def validate_logits_processors_parameters( class AdapterLogitsProcessor(LogitsProcessor): - """Wrapper for per-request logits processors + """Wrapper for per-request logits processors. To wrap a specific per-request logits processor, * Subclass `AdapterLogitsProcessor` @@ -264,7 +264,6 @@ def __init__( these arguments are used, the vLLM logits processor interface requires all three arguments to be present. """ - # Map req index -> logits processor state # # State representation is a partial[Tensor] comprising a request-level @@ -301,7 +300,7 @@ def _new_state( prompt_ids: list[int] | None, output_ids: list[int], ) -> partial[torch.Tensor] | None: - """Return state representation for new request + """Return state representation for new request. Returns None if logits processor is not applicable to request diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py index 45b5c47711f2..ea50e56be436 100644 --- a/vllm/v1/sample/logits_processor/builtin.py +++ b/vllm/v1/sample/logits_processor/builtin.py @@ -45,7 +45,7 @@ def __init__( self.min_p: torch.Tensor = self.min_p_device[:0] def is_argmax_invariant(self) -> bool: - """Min-p never impacts greedy sampling""" + """Min-p never impacts greedy sampling.""" return True def get_min_p_by_index(self, index: int) -> float: @@ -129,7 +129,8 @@ def __init__(self, _, device: torch.device, is_pin_memory: bool): def is_argmax_invariant(self) -> bool: """Logit bias can rebalance token probabilities and change the - outcome of argmax in greedy sampling.""" + outcome of argmax in greedy sampling. + """ return False def update_state(self, batch_update: BatchUpdate | None): @@ -188,7 +189,8 @@ def __init__( def is_argmax_invariant(self) -> bool: """By censoring stop tokens, min-tokens can change the outcome - of the argmax operation in greedy sampling.""" + of the argmax operation in greedy sampling. + """ return False @staticmethod @@ -377,7 +379,6 @@ def process_dict_updates( new_state: Callable[[SamplingParams, list[int] | None, list[int]], T | None], ) -> bool: """Utility function to update dict state for sparse LogitsProcessors.""" - if not batch_update: # Nothing to do. return False diff --git a/vllm/v1/sample/logits_processor/interface.py b/vllm/v1/sample/logits_processor/interface.py index dce0706335e0..355ab6a6f88f 100644 --- a/vllm/v1/sample/logits_processor/interface.py +++ b/vllm/v1/sample/logits_processor/interface.py @@ -35,7 +35,7 @@ class MoveDirectionality(Enum): @dataclass(frozen=True) class BatchUpdate: - """Persistent batch state change info for logitsprocs""" + """Persistent batch state change info for logitsprocs.""" batch_size: int # Current num reqs in batch @@ -104,5 +104,6 @@ def update_state( Args: batch_update: Non-None iff there have been changes to the batch makeup. + """ raise NotImplementedError diff --git a/vllm/v1/sample/logits_processor/state.py b/vllm/v1/sample/logits_processor/state.py index 41cbba8dffb3..24360037f7bf 100644 --- a/vllm/v1/sample/logits_processor/state.py +++ b/vllm/v1/sample/logits_processor/state.py @@ -69,7 +69,8 @@ def _ensure_removed_sorted(self) -> None: @property def removed(self) -> list[RemovedRequest]: """Removed request indices sorted in - descending order""" + descending order + """ self._ensure_removed_sorted() return self._removed @@ -81,6 +82,7 @@ def removed_append(self, index: int) -> None: Args: index: request index + """ if self._is_removed_sorted: raise RuntimeError( @@ -93,14 +95,14 @@ def has_removed(self) -> bool: return bool(self._removed) def peek_removed(self) -> int | None: - """Return lowest removed request index""" + """Return lowest removed request index.""" if self.has_removed(): self._ensure_removed_sorted() return self._removed[-1] return None def pop_removed(self) -> int | None: - """Pop lowest removed request index""" + """Pop lowest removed request index.""" if self.has_removed(): self._ensure_removed_sorted() return self._removed.pop() @@ -125,6 +127,7 @@ def get_and_reset(self, batch_size: int) -> BatchUpdate | None: Returns: Frozen logitsprocs batch update instance; `None` if no updates + """ # Reset removal-sorting logic self._is_removed_sorted = False diff --git a/vllm/v1/sample/ops/logprobs.py b/vllm/v1/sample/ops/logprobs.py index 9cdcfa456a80..359a36a81611 100644 --- a/vllm/v1/sample/ops/logprobs.py +++ b/vllm/v1/sample/ops/logprobs.py @@ -9,8 +9,7 @@ @torch.compile(backend=current_platform.simple_compile_backend) def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.Tensor: - """ - Counts elements in each row of x that are greater than the corresponding + """Counts elements in each row of x that are greater than the corresponding value in values. Use torch.compile to generate an optimized kernel for this function. otherwise, it will create additional copies of the input tensors and cause memory issues. @@ -21,6 +20,7 @@ def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.T Returns: torch.Tensor: A 1D tensor of shape (batch_size,) with the counts. + """ torch._check(x.shape[0] >= 1) torch._check(x.shape[0] == values.shape[0]) diff --git a/vllm/v1/sample/ops/penalties.py b/vllm/v1/sample/ops/penalties.py index 7bc6ec7ab893..2e2981bc4c46 100644 --- a/vllm/v1/sample/ops/penalties.py +++ b/vllm/v1/sample/ops/penalties.py @@ -15,9 +15,7 @@ def apply_all_penalties( repetition_penalties: torch.Tensor, output_token_ids: list[list[int]], ) -> torch.Tensor: - """ - Applies presence, frequency and repetition penalties to the logits. - """ + """Applies presence, frequency and repetition penalties to the logits.""" _, vocab_size = logits.shape output_tokens_t = _convert_to_tensors(output_token_ids, vocab_size, logits.device) @@ -41,9 +39,7 @@ def apply_all_penalties( def _convert_to_tensors( output_token_ids: list[list[int]], vocab_size: int, device: torch.device ) -> torch.Tensor: - """ - Convert the different list data structures to tensors. - """ + """Convert the different list data structures to tensors.""" output_tokens_tensor = make_tensor_with_pad( output_token_ids, # Use the value of vocab_size as a pad since we don't have a diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index aeb82dec4232..06b3fc3423f3 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -91,8 +91,7 @@ def flashinfer_sampler_supported() -> bool: class TopKTopPSampler(nn.Module): - """ - Module that performs optional top-k and top-p filtering followed by + """Module that performs optional top-k and top-p filtering followed by weighted random sampling of logits. Implementations may update the logits tensor in-place. @@ -154,8 +153,7 @@ def forward_native( k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - PyTorch-native implementation of top-k and top-p sampling. + """PyTorch-native implementation of top-k and top-p sampling. The logits tensor may be updated in-place. """ @@ -207,8 +205,7 @@ def forward_cpu( k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - PyTorch-native implementation of top-k and top-p sampling for CPU. + """PyTorch-native implementation of top-k and top-p sampling for CPU. The logits tensor may be updated in-place. """ @@ -423,8 +420,7 @@ def apply_top_k_top_p_pytorch( def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor: - """ - Apply top-k mask to the logits. + """Apply top-k mask to the logits. This implementation doesn't involve sorting the entire vocab. Note however that it involves a GPU->CPU sync which can be detrimental for diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 695ba624f77f..20b636e745a1 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Combined Top-K and Top-P Triton kernels. +"""Combined Top-K and Top-P Triton kernels. Based on the paper "Qrita: High-performance Top-k and Top-p Algorithm for GPUs using Pivot-based Truncation and Selection" By Park et al. @@ -1617,8 +1616,7 @@ def apply_top_k_top_p_triton( p: torch.Tensor | None, mask_value: float = float("-inf"), ) -> torch.Tensor: - """ - Apply combined top-k and top-p masking using Triton. + """Apply combined top-k and top-p masking using Triton. Top-k is applied first (by logit value), then top-p is applied to the remaining k values (by probability). @@ -1634,6 +1632,7 @@ def apply_top_k_top_p_triton( Returns: The masked logits tensor. It may or may not be modified in-place. + """ assert logits.ndim == 2 assert logits.dtype == torch.float32 diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 85b623bd84fa..d5f01f0fb40e 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -43,8 +43,7 @@ class RejectionSampler(nn.Module): - """ - The implementation strictly follows the algorithm described in + """The implementation strictly follows the algorithm described in https://arxiv.org/abs/2211.17192. However, we want to clarify the terminology used in the implementation: accepted tokens: tokens that are accepted based on the relationship @@ -113,8 +112,7 @@ def forward( logits: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> SamplerOutput: - """ - Args: + """Args: metadata: Metadata for spec decoding. draft_probs (Optional[torch.Tensor]): @@ -130,10 +128,12 @@ def forward( sampling_metadata (vllm.v1.sample.metadata.SamplingMetadata): Additional metadata needed for sampling, such as temperature, top-k/top-p parameters, or other relevant information. + Returns: SamplerOutput: Contains the final output token IDs and their logprobs if requested. + """ assert metadata.max_spec_len <= MAX_SPEC_LEN @@ -272,6 +272,7 @@ def parse_output( logprobs_tensors: LogprobsTensors | None = None, ) -> tuple[list[list[int]], LogprobsLists | None]: """Parse the output of the rejection sampler. + Args: output_token_ids: The sampled token IDs in shape [batch_size, max_spec_len + 1]. The rejected tokens are @@ -280,8 +281,10 @@ def parse_output( vocab_size: The size of the vocabulary. discard_req_indices: Optional row indices to discard tokens in. logprobs_tensors: Optional logprobs tensors to filter. + Returns: A list of lists of token IDs. + """ output_token_ids_np = output_token_ids.cpu().numpy() # Create mask for valid tokens. @@ -536,6 +539,7 @@ def apply_sampling_constraints( Returns: torch.Tensor: Processed logits if non-greedy sampling is used, otherwise returns the original logits. + """ assert logits.ndim == 2 assert cu_num_draft_tokens.ndim == 1 @@ -597,8 +601,10 @@ def expand_batch_to_tokens( Value to be replaced if it is found in x. replace_to: int = 0 Value to replace with when replace_from is found. + Returns: expanded_x: [num_tokens] tensor. + """ batch_size = x.shape[0] assert cu_num_tokens.shape[0] == batch_size @@ -619,8 +625,7 @@ def generate_uniform_probs( generators: dict[int, torch.Generator], device: torch.device, ) -> torch.Tensor: - """ - Generates a batch of uniform random samples, with optional seeding + """Generates a batch of uniform random samples, with optional seeding if available. This method creates a tensor of shape `(num_tokens, )` filled @@ -639,10 +644,12 @@ def generate_uniform_probs( `torch.Generator` objects. device: torch.device The device on which to allocate the tensor. + Returns: uniform_rand: torch.Tensor A tensor of shape `(num_tokens, )` containing uniform random values in the range [0, 1). + """ # NOTE(woosuk): We deliberately use float64 instead of float32 here # because when using float32, there's a non-negligible chance that diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index 3efaf33d980f..7e8d597d887e 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -19,8 +19,7 @@ class Sampler(nn.Module): - """ - A layer that samples the next tokens from the model's outputs + """A layer that samples the next tokens from the model's outputs with the following steps in order: 1. If logprobs are requested: @@ -170,6 +169,7 @@ def gather_specific_token_logprobs( Returns: LogprobsTensors with logprobs for the specified tokens, or None if no requests have logprob_token_ids. + """ if not logprob_token_ids: return None @@ -252,7 +252,6 @@ def sample( The various logits processing functions called in this method may update the logits tensor in-place. """ - logprobs_mode = logprobs_mode_override or self.logprobs_mode assert not (sampling_metadata.all_greedy and sampling_metadata.all_random) if sampling_metadata.all_random: @@ -312,8 +311,7 @@ def gather_logprobs( num_logprobs: int, token_ids: torch.Tensor, ) -> LogprobsTensors: - """ - Gather logprobs for topk and sampled/prompt token. + """Gather logprobs for topk and sampled/prompt token. Args: logprobs: (num tokens) x (vocab) tensor @@ -329,6 +327,7 @@ def gather_logprobs( Top-k int indices tensor, (num tokens) x (num_logprobs + 1) Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1) Sampled token rank tensor, (num tokens) + """ assert token_ids.dtype == torch.int64 # Find the topK values. diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py index bc4619a7eb3d..8e1e6cf4f66d 100644 --- a/vllm/v1/serial_utils.py +++ b/vllm/v1/serial_utils.py @@ -57,8 +57,7 @@ class OOBTensorConsumer(ABC): @abstractmethod def __call__(self, tensor: torch.Tensor) -> dict | None: - """ - Called with tensors for the current message. + """Called with tensors for the current message. Returns None to reject the tensor (falls back to regular serialization), otherwise a dict with arbitrary placeholder data to be included in the serialized message. @@ -91,7 +90,8 @@ def _typestr(val: Any) -> tuple[str, str] | None: def _encode_type_info_recursive(obj: Any) -> Any: """Recursively encode type information for nested structures of - lists/dicts.""" + lists/dicts. + """ if obj is None: return None if type(obj) is list: @@ -105,7 +105,8 @@ def _decode_type_info_recursive( type_info: Any, data: Any, convert_fn: Callable[[Sequence[str], Any], Any] ) -> Any: """Recursively decode type information for nested structures of - lists/dicts.""" + lists/dicts. + """ if type_info is None: return data if isinstance(type_info, dict): @@ -489,8 +490,7 @@ def run_method( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - """ - Run a method of an object with the given arguments and keyword arguments. + """Run a method of an object with the given arguments and keyword arguments. If the method is string, it will be converted to a method using getattr. If the method is serialized bytes and will be deserialized using cloudpickle. @@ -528,8 +528,7 @@ class PydanticMsgspecMixin: def __get_pydantic_core_schema__( cls, source_type: Any, handler: GetCoreSchemaHandler ) -> core_schema.CoreSchema: - """ - Make msgspec.Struct compatible with Pydantic, respecting defaults. + """Make msgspec.Struct compatible with Pydantic, respecting defaults. Handle JSON=>msgspec.Struct. Used when exposing msgspec.Struct to the API as input or in `/docs`. Note this is cached by Pydantic and not called on every validation. diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index bc7e6c3a13b4..372fc2201db4 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -259,7 +259,8 @@ def _derive_cpu_config( gpu_config: "KVCacheConfig", cpu_capacity_bytes: int ) -> "KVCacheConfig": """Derive a CPU KVCacheConfig from the GPU config. - Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" + Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio. + """ # Import here to avoid potential circular imports from vllm.v1.kv_cache_interface import KVCacheTensor @@ -314,14 +315,14 @@ def _estimate_lazy_target_blocks( def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: """Bind GPU block pool so that we can touch blocks during stores. - Called by Scheduler after kv_cache_manager is ready.""" + Called by Scheduler after kv_cache_manager is ready. + """ self._gpu_block_pool = gpu_block_pool def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int | None, bool]: """Return (num_new_tokens, is_async) from consecutive CPU cache hits.""" - # Pins found CPU blocks so they survive LRU eviction until # update_state_after_alloc() consumes them. Any pin from an earlier # call on the same request (e.g. retry after a failed allocate_slots) @@ -636,8 +637,8 @@ def _prepare_eager_store_specs( Returns: (gpu_block_ids, cpu_block_ids, req_ids, block_meta) for the store event. ``block_meta`` is None when kv cache events are disabled. - """ + """ merged_gpu_block_ids: list[int] = [] merged_cpu_block_ids: list[int] = [] req_ids: list[str] = [] @@ -1130,7 +1131,8 @@ def request_finished( block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: """Always returns (False, None). GPU blocks are protected by ref_cnt, - so the scheduler can free blocks immediately.""" + so the scheduler can free blocks immediately. + """ req_id = request.request_id # Release any temp CPU hit pin from get_num_new_matched_tokens() @@ -1330,7 +1332,6 @@ def reset(self) -> bool: the transfer finished, then release refs without caching abandoned store results. """ - self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks) for transfer in self._pending_finished_stores: self._release_transfer_refs(transfer) diff --git a/vllm/v1/simple_kv_offload/metadata.py b/vllm/v1/simple_kv_offload/metadata.py index 8c8d4511ee60..68ced9b89c73 100644 --- a/vllm/v1/simple_kv_offload/metadata.py +++ b/vllm/v1/simple_kv_offload/metadata.py @@ -14,8 +14,7 @@ @dataclass class SimpleCPUOffloadMetadata(KVConnectorMetadata): - """ - Metadata passed from scheduler to worker for CPU offload operations. + """Metadata passed from scheduler to worker for CPU offload operations. The worker receives flat block lists keyed by a monotonic event_idx. Job->req_id translation is handled by the scheduler-side manager diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index fc950af40fa6..9fc16014c259 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -91,6 +91,7 @@ def register_kv_caches( tensor (attention layers) or a list of tensors (Mamba layers in hybrid models). All values are included for offloading by resolving to their underlying raw storage. + """ if not kv_caches: logger.warning("No KV caches to offload.") @@ -313,6 +314,7 @@ def get_finished( tuple of (finished_sending, finished_recving). - finished_sending: always None (stores use worker metadata). - finished_recving: req_ids whose loads have completed. + """ metadata = self._connector_metadata finished_recving: set[str] = set() diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 284db7dee71f..d7d3e3adb02b 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -219,8 +219,7 @@ def dummy_run( is_graph_capturing: bool = False, slot_mappings: dict[str, torch.Tensor] | None = None, ) -> None: - """ - Key differences to default dummy_run: + """Key differences to default dummy_run: - Only one forward pass due to parallel drafting - DFlash uses context states as unpadded metadata, so hidden_states will use the unpadded num_tokens instead of num_input_tokens diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index 6325c1bf5daa..77700a411fd1 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -106,6 +106,7 @@ def propose( The main purpose is to cache hidden states, not to speculate. Args: + num_speculative_tokens: Number of draft tokens to return per request sampled_token_ids: Sampled token IDs from the target model target_hidden_states: List of hidden state tensors from target model (one per aux hidden state layer) @@ -117,6 +118,7 @@ def propose( Tuple of: - Draft tokens matching sampled tokens, shape [batch_size, 1] - KV connector output (if KV transfer is active), else None + """ assert num_speculative_tokens == self.num_speculative_tokens assert self.model is not None and isinstance(target_hidden_states, list) @@ -318,14 +320,12 @@ def prepare_next_token_ids_padded( gpu_input_batch: InputBatch, discard_request_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Prepare next token IDs for speculative decoding. + """Prepare next token IDs for speculative decoding. Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1). For each request we either use the sampled token (if valid and not discarded) or a backup token from the request state. """ - # Precompute backup token IDs for discarded requests. num_reqs = gpu_input_batch.num_reqs for i in range(num_reqs): @@ -359,6 +359,7 @@ def load_model(self, target_model: nn.Module) -> None: Args: target_model: The target model (passed for compatibility with EagleProposer interface, but not used here) + """ # Get the target model's attention layers before loading draft model target_attn_layer_names = set( @@ -395,7 +396,8 @@ def load_model(self, target_model: nn.Module) -> None: def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: """Validate all drafting layers belong to the same KV cache group - and record the group index for common_attn_metadata selection.""" + and record the group index for common_attn_metadata selection. + """ assert len(self.attn_layer_names) == 1 layer = self.attn_layer_names[0] for gid, group in enumerate(kv_cache_config.kv_cache_groups): diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index bea677dde85a..318b283ceb93 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -221,7 +221,8 @@ def load_model(self, target_model: nn.Module) -> None: def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: """Draft layers span multiple KV cache groups (sliding + full attention with different head dimensions), so skip the base - class single-group assertion.""" + class single-group assertion. + """ def initialize_attn_backend( self, @@ -229,7 +230,8 @@ def initialize_attn_backend( kernel_block_sizes: list[int] | None = None, ) -> None: """Create separate AttentionGroup objects per KV cache spec - so that each head-dim variant gets its own metadata builder.""" + so that each head-dim variant gets its own metadata builder. + """ all_attn_layers = get_layers_from_vllm_config( self.vllm_config, AttentionLayerBase, # type: ignore[type-abstract] diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 1bd1bdc45bbd..dff94b0544dd 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -785,7 +785,8 @@ def _update_positions_dependent_metadata( block_size: int, ) -> torch.Tensor: """Update positions, slot mappings, and sequence metadata for the - next draft step. Returns the updated positions tensor.""" + next draft step. Returns the updated positions tensor. + """ positions_1d = positions[0] if self.uses_mrope else positions if self.uses_mrope: out_pos = self.mrope_positions[0, :batch_size] @@ -1024,8 +1025,7 @@ def prepare_next_token_ids_cpu( gpu_input_batch: InputBatch, num_scheduled_tokens: dict[str, int], ) -> torch.Tensor: - """ - This function is used to prepare the inputs for speculative decoding. + """This function is used to prepare the inputs for speculative decoding. It calculates the next token ids for each request based on the sampled token ids from the CPU. If a request has no sampled token ids (e.g., during the initial decoding steps), it falls back to using the request @@ -1057,8 +1057,7 @@ def prepare_next_token_ids_padded( gpu_input_batch: InputBatch, discard_request_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - This function is used to prepare the inputs for speculative decoding. + """This function is used to prepare the inputs for speculative decoding. It calculates the next token ids and the number of valid sampled tokens for each request, considering the "discarded" requests whose next token is not sampled and comes from `request.get_token_id()` instead. This is denoted @@ -1107,8 +1106,7 @@ def prepare_inputs_padded( spec_decode_metadata: SpecDecodeMetadata, valid_sampled_tokens_count: torch.Tensor, ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]: - """ - This function is used to prepare the inputs for speculative decoding + """This function is used to prepare the inputs for speculative decoding It updates the common_attn_metadata for speculative decoding, but does not consider the rejected tokens. Instead, all tokens are included as inputs to the speculator, with the rejected tokens @@ -1166,8 +1164,7 @@ def prepare_inputs( sampled_token_ids: list[list[int]], num_draft_tokens: list[int], ) -> tuple[CommonAttentionMetadata, torch.Tensor]: - """ - This function is used to prepare the inputs for speculative decoding. + """This function is used to prepare the inputs for speculative decoding. It updates to the common_attn_metadata to account for the rejected tokens (and newly sampled tokens). It also returns the token indices of the tokens that should be fed to the speculator. @@ -1309,8 +1306,7 @@ def _create_draft_vllm_config(self) -> VllmConfig: return base def _get_model(self) -> nn.Module: - """ - Default method to call get_model(). Can be overridden by subclasses which + """Default method to call get_model(). Can be overridden by subclasses which need to customize model loading. """ from vllm.compilation.backends import set_model_tag @@ -1422,8 +1418,7 @@ def load_model(self, target_model: nn.Module) -> None: self.parallel_drafting_hidden_state_tensor.copy_(flat_mask) def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: - """ - Some draft models may not have their own embedding layers, and some may + """Some draft models may not have their own embedding layers, and some may have a duplicate copy of the target model's embedding layers. In these cases, we share the target model's embedding layers with the draft model to save memory. @@ -1515,8 +1510,7 @@ def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: ) def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: - """ - Some draft models may not have their own LM head, and some may have a + """Some draft models may not have their own LM head, and some may have a duplicate copy of the target model's LM head. In these cases, we share the target model's LM head with the draft model to save memory. """ @@ -1680,8 +1674,7 @@ def dummy_run( self.model(**kwargs) def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool: - """ - Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary + """Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary hidden states and directly uses the last layer output just like eagle1. They might indicate this by setting "use_aux_hidden_state" to False inside the "eagle_config" dict of their hf_config. @@ -1696,8 +1689,7 @@ def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool: return use_aux_hidden_state def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: - """ - Validate that all drafting layers belong to the same KVCacheGroup. + """Validate that all drafting layers belong to the same KVCacheGroup. Need this assumption to ensure all drafting layers can use the same AttentionMetadata. May extend to multiple AttentionMetadata in the future. @@ -1723,8 +1715,7 @@ def initialize_attn_backend( kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] | None = None, ) -> None: - """ - Initialize AttentionGroups for draft layers using kv_cache_config. + """Initialize AttentionGroups for draft layers using kv_cache_config. Called from the model runner's initialize_metadata_builders. """ all_attn_layers = get_layers_from_vllm_config( diff --git a/vllm/v1/spec_decode/medusa.py b/vllm/v1/spec_decode/medusa.py index 7adf7cff5f77..7c926ba8372f 100644 --- a/vllm/v1/spec_decode/medusa.py +++ b/vllm/v1/spec_decode/medusa.py @@ -16,9 +16,7 @@ class MedusaProposer: - """ - Medusa proposer class for generating token sequences - """ + """Medusa proposer class for generating token sequences.""" def __init__( self, diff --git a/vllm/v1/spec_decode/ngram_proposer.py b/vllm/v1/spec_decode/ngram_proposer.py index e0240d0e66b8..5e0d2cfc3c43 100644 --- a/vllm/v1/spec_decode/ngram_proposer.py +++ b/vllm/v1/spec_decode/ngram_proposer.py @@ -72,6 +72,8 @@ def batch_propose( """Batch version of ngram proposer using numba for acceleration. Args: + num_requests: + Number of requests in the batch. valid_ngram_requests: Set of indices of requests that need ngram proposals. num_tokens_no_spec: @@ -87,6 +89,7 @@ def batch_propose( list[list[int]]: A list where each element is a list of proposed token IDs for the corresponding request. + """ draft_token_ids: list[list[int]] = [] @@ -211,8 +214,7 @@ def _find_longest_matched_ngram_and_propose_tokens( max_model_len: int, k: int, ) -> np.ndarray: - """ - Find the longest n-gram which matches the suffix of the given tokens + """Find the longest n-gram which matches the suffix of the given tokens whose length is within [min_ngram, max_ngram] (inclusive). If found, we will extract k right after the matched ngram. diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index 2de5c3be007f..061a8f5ea93b 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GPU-accelerated N-gram proposer using fully async PyTorch tensor operations. +"""GPU-accelerated N-gram proposer using fully async PyTorch tensor operations. This version uses a fully vectorized approach with unfold and argmax for finding the first match across all sequences in parallel. @@ -52,8 +51,7 @@ def _find_first_and_extract_all_n_parallel( max_ngram_len: int, num_draft_tokens: int, ) -> torch.Tensor: - """ - Find suffix n-gram matches and extract following tokens. + """Find suffix n-gram matches and extract following tokens. Searches for the earliest prior occurrence of the trailing n-gram, tries multiple lengths, and picks the longest valid match. @@ -66,6 +64,7 @@ def _find_first_and_extract_all_n_parallel( Returns: Draft token predictions; -1 means invalid/no match. + """ batch_size = token_ids.shape[0] max_seq_len = token_ids.shape[1] @@ -165,8 +164,7 @@ def forward( token_ids_gpu: torch.Tensor, combined_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Forward pass for N-gram proposal using GPU tensor operations. + """Forward pass for N-gram proposal using GPU tensor operations. Args: num_tokens_no_spec: Number of tokens for each sequence [batch_size] @@ -177,8 +175,8 @@ def forward( draft_tokens: [batch_size, k] on GPU num_valid_draft_tokens: [batch_size] int32 on GPU, count of leading valid (non -1) tokens per request. - """ + """ device = token_ids_gpu.device # Infer batch size to preserve dynamic shape. @@ -282,8 +280,7 @@ def _generate_dummy_data( pattern_len: int, device: str = "cuda", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Generate random test data with n-gram repetitions. + """Generate random test data with n-gram repetitions. Args: batch_size: Number of sequences in the batch @@ -296,6 +293,7 @@ def _generate_dummy_data( num_tokens: [batch_size] tensor sampled_flags: [batch_size] bool tensor valid_mask: [batch_size] bool tensor + """ token_ids = torch.zeros( batch_size, @@ -321,8 +319,7 @@ def propose( valid_sampled_token_ids_gpu: torch.Tensor, # [batch_size, num_spec_tokens + 1] valid_sampled_tokens_count: torch.Tensor, # [batch_size] ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Propose draft tokens using GPU-accelerated n-gram matching. + """Propose draft tokens using GPU-accelerated n-gram matching. Scatter sampled tokens into `token_ids_gpu`, compute temporary updated lengths, then run the kernel. @@ -338,6 +335,7 @@ def propose( draft_tokens: Proposed draft token IDs [batch_size, k] num_valid_draft_tokens: Count of leading valid draft tokens per request [batch_size] + """ assert num_speculative_tokens == self.k assert token_ids_gpu.device == self.device @@ -398,8 +396,7 @@ def update_token_ids_ngram( num_tokens_no_spec: torch.Tensor, discard_request_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Prepare speculative decoding inputs on device: + """Prepare speculative decoding inputs on device: compute next token ids and valid counts, honoring discarded requests and rejected tokens, without CPU-GPU sync. """ @@ -485,6 +482,7 @@ def update_scheduler_for_invalid_drafts( num_valid_draft_tokens_cpu: CPU buffer of valid draft counts. scheduler_output: Scheduler metadata to update in-place. req_id_to_index: Request-id to batch-index mapping. + """ req_data = scheduler_output.scheduled_cached_reqs num_valid_draft_tokens_event.synchronize() @@ -652,9 +650,7 @@ def copy_num_valid_draft_tokens( num_valid_draft_tokens: torch.Tensor | None, batch_size: int, ) -> None: - """ - Async D2H copy of per-request valid draft counts. - """ + """Async D2H copy of per-request valid draft counts.""" if num_valid_draft_tokens is None: return diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 109d326356a1..7a73473adc2f 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -155,7 +155,6 @@ def build_per_group_and_layer_attn_metadata( def _maybe_share_lm_head(self, target_language_model: torch.nn.Module) -> None: """Step3.5 MTP uses the lm_head stored in each MTP layer.""" - # The base MTP path shares target lm_head into shared_head.head. # Step3.5 checkpoints carry per-MTP-layer shared_head weights. return diff --git a/vllm/v1/spec_decode/suffix_decoding.py b/vllm/v1/spec_decode/suffix_decoding.py index 66137a006316..5d100ce9ed97 100644 --- a/vllm/v1/spec_decode/suffix_decoding.py +++ b/vllm/v1/spec_decode/suffix_decoding.py @@ -7,8 +7,7 @@ class SuffixDecodingProposer: - """ - Speculative decoding proposer for Suffix Decoding (https://arxiv.org/pdf/2411.04975). + """Speculative decoding proposer for Suffix Decoding (https://arxiv.org/pdf/2411.04975). This class imports and uses the official implementation from Arctic Inference (https://github.com/snowflakedb/ArcticInference). """ diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index fbfe4438422a..6db4663660de 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -50,8 +50,7 @@ def eagle_step_slot_mapping_metadata_kernel( PAD_ID: tl.constexpr, batch_size, ): - """ - Fused kernel for EAGLE autoregressive step: updates positions, slot mapping, + """Fused kernel for EAGLE autoregressive step: updates positions, slot mapping, and sequence lengths in a single kernel to reduce launch overhead. Launched with input_batch_size threads. Threads with req_idx >= batch_size @@ -149,8 +148,7 @@ def eagle_step_update_slot_mapping_and_metadata( out_slot_mapping: torch.Tensor, input_batch_size: int | None = None, ) -> None: - """ - Fused update of slot mapping and metadata for one EAGLE autoregressive step. + """Fused update of slot mapping and metadata for one EAGLE autoregressive step. Updates seq_lens in place. Writes to out_clamped_positions and out_slot_mapping. When input_batch_size > batch_size, threads beyond batch_size write @@ -166,6 +164,7 @@ def eagle_step_update_slot_mapping_and_metadata( out_slot_mapping: [input_batch_size] output buffer for slot mapping input_batch_size: total batch size including cudagraph padding; defaults to batch_size (no padding) + """ batch_size = positions_1d.shape[0] if input_batch_size is None: @@ -197,8 +196,7 @@ def eagle_prepare_inputs_padded_kernel( num_rejected_tokens_gpu_ptr, # [num_reqs] (output) num_reqs, # tl.int32 ): - """ - Fused kernel for Eagle prepare_input_padded. This kernel computes the + """Fused kernel for Eagle prepare_input_padded. This kernel computes the token index to sample for each request, taking into account the number of draft tokens and the number of valid sampled tokens (which is one more than the number of accepted tokens). @@ -270,8 +268,7 @@ def eagle_prepare_next_token_padded_kernel( stride_sampled_token_ids, # tl.int32 (stride for dim 0) BLOCK_SIZE_TOKENS: tl.constexpr, # Power-of-2 >= num_sampled_tokens_per_req ): - """ - Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the + """Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the number of valid (1 + accepted) tokens for each request, and the corresponding "next" token id to sample from during speculative decoding. This is the "last accepted token" from the sampled tokens, or the backup token if no @@ -401,8 +398,7 @@ def extend_all_queries_by_N( arange: torch.Tensor, new_slot_mapping: torch.Tensor, ) -> CommonAttentionMetadata: - """ - Creates a new CommonAttentionMetadata with all query lengths increased by N. + """Creates a new CommonAttentionMetadata with all query lengths increased by N. Also all seq lens are increased by N. This is useful e.g. in speculative decoding with parallel drafting, where we extend each sequence by N tokens and predict all tokens in one pass. @@ -453,8 +449,7 @@ def copy_and_expand_eagle_inputs_kernel( shift_input_ids, # tl.bool BLOCK_SIZE_TOKENS: tl.constexpr, # Blocks along token dim to handle prefills ): - """ - Copy and expand inputs from the target model to the drafting buffers for Eagle + """Copy and expand inputs from the target model to the drafting buffers for Eagle speculative decoding. This kernel handles padding slots and parallel drafting tokens, if enabled. """ @@ -668,8 +663,7 @@ def copy_and_expand_dflash_inputs_kernel( BLOCK_SIZE: tl.constexpr, HAS_NUM_REJECTED: tl.constexpr = False, ): - """ - Fused kernel for DFlash first-pass input setup. + """Fused kernel for DFlash first-pass input setup. Per request, this kernel: 1. Copies context positions from target_positions to @@ -853,5 +847,6 @@ def update_num_computed_tokens_for_batch_change( def unconditional_to_conditional_rates(rates: list[float]) -> list[float]: """Convert per-position unconditional rates to per-position conditional - rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}).""" + rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}). + """ return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])] diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 0d8b7377c557..5f43c5a890fc 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -161,7 +161,6 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: Returns True if the parser was advanced successfully. Returns False if the parser failed to advance. """ - if self.ll_tokenizer.eos_token in tokens: if self.ll_matcher.is_stopped() and not self.terminated: self.rollback_lag = 1 diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index dd460f96ec91..bb49fb189b72 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -283,7 +283,7 @@ def subpattern_consumes(parsed) -> bool: def _check_unsupported(parsed) -> None: - """Check for regex features unsupported by regex-automata""" + """Check for regex features unsupported by regex-automata.""" tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # backreference @@ -310,8 +310,7 @@ def _check_unsupported(parsed) -> None: def validate_regex_is_buildable(pattern: str) -> None: - """ - Validates that the input regex is not using unsupported features + """Validates that the input regex is not using unsupported features of the `regex-automata` crate (outlines_core regex engine) and has a universal start state. definition of universal start state used can be found at: diff --git a/vllm/v1/structured_output/backend_types.py b/vllm/v1/structured_output/backend_types.py index d9826c18f534..b3695efd009c 100644 --- a/vllm/v1/structured_output/backend_types.py +++ b/vllm/v1/structured_output/backend_types.py @@ -33,8 +33,7 @@ class StructuredOutputGrammar(ABC): @abstractmethod def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: - """ - Determines whether the provided tokens are accepted for the + """Determines whether the provided tokens are accepted for the given request. Args: @@ -43,12 +42,12 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: Returns: bool: True if the tokens are accepted, False otherwise. + """ @abstractmethod def validate_tokens(self, tokens: list[int]) -> list[int]: - """ - Validates the provided tokens against the grammar. + """Validates the provided tokens against the grammar. Will not advance the FSM. Args: @@ -57,42 +56,41 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns: list[int]: A list of accepted token IDs. Will be a prefix of the input tokens, and empty if none are accepted. + """ @abstractmethod def rollback(self, num_tokens: int) -> None: - """ - Rolls back the state of the grammar by a specified number of tokens. + """Rolls back the state of the grammar by a specified number of tokens. Will also revert counters for the number of processed tokens. Args: num_tokens (int): The number of tokens to roll back. + """ @abstractmethod def fill_bitmask(self, bitmask: "torch.Tensor", batch_index: int) -> None: - """ - Fills the bitmask for a specific batch index. + """Fills the bitmask for a specific batch index. Args: bitmask (torch.Tensor): The bitmask to fill batch_index (int): The index in the bitmask to fill + """ @abstractmethod def is_terminated(self) -> bool: - """ - Checks whether the structured output process has terminated. + """Checks whether the structured output process has terminated. Returns: bool: True if the process is terminated, False otherwise. + """ @abstractmethod def reset(self): - """ - Resets the state of the structured output grammar. - """ + """Resets the state of the structured output grammar.""" @dataclass @@ -110,8 +108,7 @@ def compile_grammar( grammar_spec: str, stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: - """ - Compiles a grammar specification into a structured output grammar. + """Compiles a grammar specification into a structured output grammar. Args: request_type (StructuredOutputOptions): The type of structured @@ -123,20 +120,19 @@ def compile_grammar( Returns: StructuredOutputGrammar: The compiled structured output grammar. + """ @abstractmethod def allocate_token_bitmask(self, max_num_seqs: int) -> "torch.Tensor": - """ - Allocates a token bitmask for the specified maximum number of sequences. + """Allocates a token bitmask for the specified maximum number of sequences. Args: max_num_seqs (int): The maximum number of sequences for which to allocate the bitmask. + """ @abstractmethod def destroy(self): - """ - Backend-specific cleanup. - """ + """Backend-specific cleanup.""" diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 92122f85d0fc..4969789434c2 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -74,6 +74,7 @@ def compile_regex_with_timeout(fn: Callable[[str], _T], pattern: str) -> _T: Raises: ValueError: If compilation exceeds the configured timeout. + """ timeout = envs.VLLM_REGEX_COMPILATION_TIMEOUT_S if timeout <= 0: @@ -103,13 +104,14 @@ def apply_grammar_bitmask( input_batch: InputBatch, logits: torch.Tensor, ) -> None: - """ - Apply grammar bitmask to output logits of the model with xgrammar function. + """Apply grammar bitmask to output logits of the model with xgrammar function. Args: scheduler_output (SchedulerOutput): The result of engine scheduling. + grammar_output (GrammarOutput): The grammar bitmask to apply. input_batch (InputBatch): The input of model runner. logits (torch.Tensor): The output logits of model forward. + """ # Serialization of np.ndarray is much more efficient than a tensor, # so we receive it in that format. @@ -190,8 +192,7 @@ def apply_grammar_bitmask( class OutlinesVocabulary: - """ - Wrapper class for `outlines_core.Vocabulary`, + """Wrapper class for `outlines_core.Vocabulary`, which allows us to store a hash with the vocabulary """ @@ -206,7 +207,7 @@ def __init__(self, vocabulary: oc.Vocabulary) -> None: def get_outlines_cache_path() -> str: - """Get the context object that contains previously-computed return values""" + """Get the context object that contains previously-computed return values.""" outlines_cache_dir = os.getenv("OUTLINES_CACHE_DIR") xdg_cache_home = os.getenv("XDG_CACHE_HOME") home_dir = os.path.expanduser("~") @@ -294,8 +295,7 @@ def clear(self): def get_outlines_cache(): - """Get the Cache instance to be used for index caching""" - + """Get the Cache instance to be used for index caching.""" cache_dir = get_outlines_cache_path() if envs.VLLM_V1_USE_OUTLINES_CACHE: logger.warning( @@ -324,6 +324,7 @@ def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]: Returns: A Dict of token string -> equivalent token ids + """ eos_token_id = tokenizer.eos_token_id @@ -403,8 +404,7 @@ def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary: def grammar_is_likely_lark(grammar_str: str) -> bool: - """ - Check if grammar appears to use Lark syntax. + """Check if grammar appears to use Lark syntax. Args: grammar_str: Input grammar string @@ -417,6 +417,7 @@ def grammar_is_likely_lark(grammar_str: str) -> bool: True >>> grammar_is_likely_lark("rule ::= 'abc'") False + """ if not grammar_str or not isinstance(grammar_str, str): return False @@ -435,8 +436,7 @@ def grammar_is_likely_lark(grammar_str: str) -> bool: def convert_lark_to_ebnf(grammar_str: str) -> str: - """ - Convert a Lark grammar string to EBNF format. + """Convert a Lark grammar string to EBNF format. EBNF reference: https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md @@ -453,6 +453,7 @@ def convert_lark_to_ebnf(grammar_str: str) -> str: >>> print(convert_lark_to_ebnf("rule: 'hello'")) root ::= rule rule ::= "hello" + """ if not isinstance(grammar_str, str): raise ValueError(f"Grammar must be a string, got {type(grammar_str)}") diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index facafd28608f..ef800f247e9c 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -144,7 +144,8 @@ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: """NOTE: Because this method is non-blocking, explicit synchronization - is needed to ensure the data is copied to CPU.""" + is needed to ensure the data is copied to CPU. + """ if n is None: return self.cpu.copy_(self.gpu, non_blocking=True) return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True) @@ -158,7 +159,8 @@ def get_engine_client_zmq_addr( """Return an IPC path (``local_only=True``) or ``tcp://host:port``. ``port=0`` lets the kernel assign the port at ``bind()`` time; the - caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``.""" + caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``. + """ if local_only: return get_open_zmq_ipc_path() return get_tcp_uri(host, port) @@ -201,6 +203,7 @@ def __init__( output_addresses: Output addresses for each API server stats_update_address: Optional stats update address tensor_queue: Optional tensor IPC queue for sharing MM tensors + """ self.listen_address = listen_address self.sock = sock @@ -262,7 +265,8 @@ def gather_actual_addresses( ) -> tuple[list[str], list[str]]: """Return (inputs, outputs) reported by each child, indexed by ``client_index``. Raises ``RuntimeError`` on timeout or premature - child exit.""" + child exit. + """ n = len(self._address_pipes) inputs: list[str | None] = [None] * n outputs: list[str | None] = [None] * n @@ -423,7 +427,8 @@ def shutdown(self, timeout: float | None = None) -> None: class _SubprocessWrapper: """Wraps subprocess.Popen to provide the BaseProcess-like interface - needed by wait_for_completion_or_failure.""" + needed by wait_for_completion_or_failure. + """ def __init__(self, proc, name: str): self._proc = proc @@ -519,7 +524,6 @@ def run_api_server_worker_proc( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: """Entrypoint for individual API server worker processes.""" - from vllm.entrypoints.launchers.api_server.entry import run_server_worker client_config = client_config or {} @@ -550,8 +554,8 @@ def wait_for_completion_or_failure( If CoreEngineProcManager, it manages local engines; if CoreEngineActorManager, it manages all engines. coordinator: The coordinator for data parallel. - """ + """ try: logger.info("Waiting for API servers to complete ...") # Create a mapping of sentinels to their corresponding processes @@ -613,6 +617,7 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: Args: procs: List of processes to shutdown timeout: Maximum time in seconds to wait for graceful shutdown + """ if timeout is None: # Keep a small grace period for best-effort cleanup paths that do not @@ -667,8 +672,7 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: def copy_slice( from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int ) -> torch.Tensor: - """ - Copy the first length elements of a tensor into another tensor in a + """Copy the first length elements of a tensor into another tensor in a non-blocking manner. Used to copy pinned CPU tensor data to pre-allocated GPU tensors. @@ -682,7 +686,6 @@ def report_usage_stats( vllm_config, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT ) -> None: """Report usage statistics if enabled.""" - if not is_usage_stats_enabled(): return @@ -803,6 +806,7 @@ def tensor_data(tensor: torch.Tensor) -> memoryview: Returns: A memoryview of the tensor data as uint8. + """ return tensor.flatten().cpu().contiguous().view(torch.uint8).numpy().data @@ -826,8 +830,7 @@ def __repr__(self) -> str: def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails: - """ - Compute the number of context/generation requests and tokens + """Compute the number of context/generation requests and tokens for the current iteration's scheduler output. A requests is regarded as a context request if its output tokens are still 0, an extended chunk of chunked prefill falls into this category. @@ -838,6 +841,7 @@ def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDet Returns: An IterationDetails object containing the number of context/generation requests and tokens. + """ num_context_requests = 0 num_context_tokens = 0 diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 33b52652a154..c63908dad0d0 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -67,20 +67,20 @@ def __init__( cp_kv_cache_interleave_size: int, slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT, ): - """ - Args: - block_size: Block size used for KV cache memory allocation - max_num_reqs: Maximum number of concurrent requests supported. - max_num_blocks_per_req: Maximum number of blocks per request. - max_num_batched_tokens: Maximum number of tokens in a batch. - pin_memory: Whether to pin memory for faster GPU transfers. - device: Target device for the block table. - kernel_block_size: The block_size of underlying attention kernel. - Will be the same as `block_size` if `block_size` is supported - by the attention kernel. - slot_mapping_mode: How this cache group maps scheduled tokens to - cache slots. Mamba-like state caches do not use token slot - mappings and should use SlotMappingMode.NONE. + """Args: + block_size: Block size used for KV cache memory allocation + max_num_reqs: Maximum number of concurrent requests supported. + max_num_blocks_per_req: Maximum number of blocks per request. + max_num_batched_tokens: Maximum number of tokens in a batch. + pin_memory: Whether to pin memory for faster GPU transfers. + device: Target device for the block table. + kernel_block_size: The block_size of underlying attention kernel. + Will be the same as `block_size` if `block_size` is supported + by the attention kernel. + slot_mapping_mode: How this cache group maps scheduled tokens to + cache slots. Mamba-like state caches do not use token slot + mappings and should use SlotMappingMode.NONE. + """ self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens @@ -254,6 +254,7 @@ def map_to_kernel_blocks( # kv_manager_block_id 0 → kernel block id [0, 1] # kv_manager_block_id 1 → kernel block id [2, 3] # kv_manager_block_id 2 → kernel block id [4, 5] + """ if blocks_per_kv_block == 1: return kv_manager_block_ids diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index fe82ec1427f1..b2f2fe1506a9 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -112,8 +112,7 @@ def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch def _post_process_cudagraph_mode(tensor: torch.Tensor) -> int: - """ - Synchronize cudagraph_mode across DP ranks by taking the minimum. + """Synchronize cudagraph_mode across DP ranks by taking the minimum. If any rank has NONE (0), all ranks use NONE. This ensures all ranks send consistent values (all padded or all unpadded). """ @@ -127,8 +126,7 @@ def _synchronize_dp_ranks( cudagraph_mode: int, parallel_config: ParallelConfig, ) -> tuple[bool, torch.Tensor | None, int]: - """ - 1. Decides if each DP rank is going to microbatch. Either all ranks + """1. Decides if each DP rank is going to microbatch. Either all ranks run with microbatching or none of them do. 2. Determines the total number of tokens that each rank will run. @@ -197,8 +195,7 @@ def coordinate_batch_across_dp( uniform_decode: bool | None = None, cudagraph_mode: int = 0, ) -> tuple[bool, torch.Tensor | None, int]: - """ - Coordinates amongst all DP ranks to determine if and how the full batch + """Coordinates amongst all DP ranks to determine if and how the full batch should be split into microbatches. Args: diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index be676666e84d..8a32f996601a 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Define EC connector functionality mixin for model runners. -""" +"""Define EC connector functionality mixin for model runners.""" from collections.abc import Generator from contextlib import AbstractContextManager, contextmanager, nullcontext diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index a60e0d1892f0..108459642274 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -76,7 +76,8 @@ def __init__( model: SupportsEncoderCudaGraph, ): """Initialize CUDA graph manager with provided token budgets - and max batch size.""" + and max batch size. + """ self.vllm_config = vllm_config self.device = device self.dtype = dtype @@ -351,6 +352,7 @@ def _find_smallest_fitting_budget_given_tokens( Returns: Token budget if found, None if no fitting budget. + """ budgets = budgets if budgets is not None else self.token_budgets for budget in budgets: @@ -411,8 +413,11 @@ def _run_budget_graph( mm_kwargs: Multimodal inputs for the batch. token_budget: Token budget to use. path: Configured encoder path. + axis_keys: Resolved capture-axis keys for this batch. + Returns: Encoder outputs, or None if graph not captured. + """ graph_set = self._get_graph_set(path) num_items = len(self._get_item_specs(mm_kwargs)) @@ -562,6 +567,7 @@ def _dp_shard( images_per_rank: Number of items per rank. max_output_tokens_per_rank: Max output tokens across all ranks (for padding during all_gather). + """ tp_size = get_tensor_model_parallel_world_size() current_rank = get_tensor_model_parallel_rank() @@ -687,6 +693,7 @@ def execute( Returns: List of encoder outputs (one per item). + """ if self.use_dp: per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9657c28fe601..239b47f3d710 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -62,7 +62,8 @@ class AttentionState(NamedTuple): @dataclass(frozen=True) class BatchExecutionDescriptor: """Describes the shape of the batch and CG mode to run; this is used to make shape - matches between the capture and runtime.""" + matches between the capture and runtime. + """ cg_mode: CUDAGraphMode num_tokens: int @@ -90,7 +91,8 @@ def make_cudagraph_stats( class CreateForwardFn(Protocol): """Factory that prepares inputs (OUTSIDE the graph) and returns a forward_fn. Called with warmup=True for the warmup pass and warmup=False - for the captured pass.""" + for the captured pass. + """ def __call__( self, @@ -408,6 +410,8 @@ def capture( it is invoked once with warmup=True and again with warmup=False because attention backends may mutate or lazily initialize metadata during warmup. + progress_bar_desc: Description shown on the capture progress bar. + """ with graph_capture(device=self.device), ExitStack() as stack: if self.ubatch_runner is not None: @@ -503,7 +507,6 @@ def dispatch( num_ubatches: int = 1, ) -> BatchExecutionDescriptor: """Find matching cudagraph descriptor from priority-ordered candidates.""" - effective_loras = self._resolve_effective_loras(num_active_loras) key = (num_tokens, effective_loras) if self._graphs_captured and num_tokens > 0 and key in self._candidates: @@ -952,7 +955,8 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int: """Extrapolate the total FULL capture cost from samples of the largest graphs. The first capture allocates the pool baseline; later graphs mostly - reuse it, so the second sample is taken as the per-graph cost.""" + reuse it, so the second sample is taken as the per-graph cost. + """ if not mem_samples: return 0 first_capture = mem_samples[0] @@ -989,7 +993,8 @@ def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: def _teardown_profiling_state(runner: "GPUModelRunner") -> None: """Release the profiling KV cache and captured graphs while keeping model - weights, so the real ``initialize_kv_cache`` starts from a clean slate.""" + weights, so the real ``initialize_kv_cache`` starts from a clean slate. + """ torch.accelerator.synchronize() if hasattr(runner.model_state, "_mamba_ctx"): runner.model_state._mamba_ctx = None diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index 80309bc009c9..8eabc3adbf1b 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -55,8 +55,7 @@ def sync_cudagraph_and_dp_padding( allow_ubatching: bool = False, uniform_decode: bool = False, ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: - """ - Coordinates the batch descriptor and DP padding across all ranks. + """Coordinates the batch descriptor and DP padding across all ranks. `parallel_config` is only needed to decide whether to microbatch, so callers that never do (`allow_ubatching=False`) can leave it out. @@ -248,6 +247,12 @@ def dispatch_cg_and_sync_dp( need_eager: Force `CUDAGraphMode.NONE` instead of dispatching. num_active_loras: Active LoRA count for this rank. Does not need cross-rank agreement; it never changes a bucket's token count. + parallel_config: Only needed to decide whether to microbatch, so + callers that never do (`allow_ubatching=False`) can leave it out. + allow_ubatching: Whether this rank's batch may be split into + microbatches. Agreed across ranks before it takes effect. + uniform_decode: Whether this rank's batch is a uniform decode, used + to pick the microbatching split. dp_sync: Agreement from a prior dispatch over this same batch, to reuse. Must come from a batch with this same padded `num_tokens` and the same `uniform_token_count`; `num_reqs` may differ, as neither @@ -257,6 +262,7 @@ def dispatch_cg_and_sync_dp( Returns: (batch_desc, sync), where `sync` is this batch's agreement for a later dispatch to reuse. It is None when `dp_size` is 1 or no rank has work. + """ reuse_eager = dp_sync is not None and dp_sync.eager diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py index 6a00b35e4bce..eb81a83662a1 100644 --- a/vllm/v1/worker/gpu/lora_utils.py +++ b/vllm/v1/worker/gpu/lora_utils.py @@ -21,8 +21,7 @@ def get_lora_capture_cases( lora_config: "LoRAConfig | None", compilation_config: "CompilationConfig", ) -> list[int]: - """ - Return num_active_loras values for cudagraph capture. + """Return num_active_loras values for cudagraph capture. When cudagraph_specialize_lora=True: powers of 2 up to max_loras, plus max_loras+1. When False: [0, max_loras+1]. When LoRA disabled: [0]. diff --git a/vllm/v1/worker/gpu/mm/encoder_cache.py b/vllm/v1/worker/gpu/mm/encoder_cache.py index 21d9218182ca..b9aba7b31a58 100644 --- a/vllm/v1/worker/gpu/mm/encoder_cache.py +++ b/vllm/v1/worker/gpu/mm/encoder_cache.py @@ -24,8 +24,7 @@ def remove_request(self, req_id: str) -> None: self.mm_features.pop(req_id, None) def reset_mm_cache(self) -> None: - """ - Clear the multi-modal cache that was used during profiling, + """Clear the multi-modal cache that was used during profiling, but no longer needed during inference. """ # NOTE: v2 encoder cache profiling skips the multi-modal cache diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 82cc697e731c..dca5b260f283 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -NOTE: Coding style guide for this file: +"""NOTE: Coding style guide for this file: This model runner is shared by all models: text and multimodal, generative and embedding, public and private. As a result, this file must only contain code that is common to every model. Model-specific behavior belongs in the @@ -2229,7 +2228,8 @@ def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None: def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that - memory is reclaimable when running in the same process.""" + memory is reclaimable when running in the same process. + """ torch.accelerator.synchronize() self.cudagraph_manager = None self.fast_prefill = None diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index a5c8d6f4cfd7..6f0c9c672c0b 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -144,7 +144,8 @@ def preprocess_state( ) -> None: """Hook run on real batches before the forward pass (after block tables are gathered). Used by mamba "align" prefix caching to pre-copy state - across block boundaries. No-op by default.""" + across block boundaries. No-op by default. + """ return None def postprocess_state( diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index ab1b991d1a8d..f1124c3ee470 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -38,8 +38,8 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: """Return a bool array of shape `[input_batch.num_reqs]` marking requests that produce a sampled token this step, and therefore must have that token (and the draft block proposed from it) propagated to the earlier PP stages. - Returns None if no request in the batch produces a sample.""" - + Returns None if no request in the batch produces a sample. + """ old_computed = input_batch.num_computed_tokens_np prefill_len = input_batch.prefill_len_np # Exclude non-final prefill chunks (they don't produce a sample). @@ -179,7 +179,8 @@ def broadcast_drafts( def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* - requests in the batch.""" + requests in the batch. + """ assert not self.is_last_rank need_sampled_mask = compute_need_sampled_mask(input_batch) if need_sampled_mask is None: diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py index ff3675db130f..01ffa4bdb1e6 100644 --- a/vllm/v1/worker/gpu/sample/output.py +++ b/vllm/v1/worker/gpu/sample/output.py @@ -76,7 +76,8 @@ def _compact_sampling_mask_kernel( class SamplingMaskTensors(NamedTuple): """Device-side masks pending async D2H: compact ids, plus the bitmask as - the exact fallback for rows wider than ``max_num_kept``.""" + the exact fallback for rows wider than ``max_num_kept``. + """ # [num_requests, max_num_kept] token_ids: torch.Tensor diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index b069604c2964..a6199f570ccd 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -107,7 +107,8 @@ def get_logprobs_dims( self, idx_mapping_np: np.ndarray, include_token_ids: bool = True ) -> tuple[int, int] | None: """(num_logprobs, max_per_req_token_ids) for the given requests, or - None when none of them want logprobs.""" + None when none of them want logprobs. + """ max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) max_token_ids = ( self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np) diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index 0a9c64cbf732..b70d8ae271db 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/vllm/v1/worker/gpu/sample/watermark.py @@ -163,6 +163,7 @@ def repeated_context_mask( include_prompt: Search the prompt as well as the generated tokens. skip_partial_context: Mark contexts containing start padding so they use ordinary sampling. + """ if max_history is not None and max_history < 1: raise ValueError("max_history must be positive or None") diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 245f304ac4e3..2d0c25a41e68 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -88,9 +88,14 @@ def build_cost_tables_from_curves( """Build cost tables: graph-padded below the capture limit, smooth above. Args: + draft_curve: (size, cost) samples for the draft model. + verify_curve: (size, cost) samples for the verify model. + max_num_reqs: Largest request count to build a table for. + max_batch_tokens: Largest token count to build a table for. cudagraph_limit: Largest cudagraph-captured size. At or below it, execution pads up to the next captured size, so cost is a step function. Above it there is no padding, so cost is continuous. + """ def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray: @@ -185,7 +190,8 @@ def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int """Dummy-run kwargs whose step timings seed the cost tables. Run these inside StepTimingCollector.collect(), then hand the block's - timings to set_initial_cost_curves.""" + timings to set_initial_cost_curves. + """ max_num_tokens = self.req_states.max_num_batched_tokens size = self._cudagraph_limit = capture_sizes[-1] if capture_sizes else 0 # Also profile beyond the capture limit: real steps run there @@ -255,7 +261,8 @@ def record_confidences( input_batch: "InputBatch", ) -> None: """Publish this step's raw confidences for the ranking kernel and start - copying them to the CPU, where a later step's budget reads them.""" + copying them to the CPU, where a later step's budget reads them. + """ num_reqs = input_batch.num_reqs ready_idx = self._stale_idx ^ 1 with gpu_sync_allowed(): diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 3d243b6ee5da..f166d63e8efb 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -73,8 +73,7 @@ def on_multi_step_decode_end(self, num_reqs: int) -> None: ... @property def advance_draft_positions(self) -> bool: - """ - Whether to increment positions and seq_lens between draft steps. + """Whether to increment positions and seq_lens between draft steps. True for Eagle/standard MTP (each step produces new KV). False for Gemma4 MTP (Q-only, shares target KV, constant positions). diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index 725714fcdbc0..854c958a12d1 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -72,7 +72,8 @@ def _prepare_dflash_inputs_to_capture( class DFlashCudaGraphManager(CudaGraphManager): """DFlash CudaGraphManager for the parallel-drafting query forward, - building its own attention metadata from scratch.""" + building its own attention metadata from scratch. + """ def capture( self, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index b93a79e45513..dd37396b6b4d 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -13,7 +13,6 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: """Share when the draft has no own copy, or its copy matches the target.""" - if not getattr(eagle, flag, False) or draft is None: return True if target is None: @@ -29,7 +28,8 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): """The target's lm_head — from get_language_model() for - *ForConditionalGeneration targets, else the top-level module.""" + *ForConditionalGeneration targets, else the top-level module. + """ return getattr(target_language_model, "lm_head", None) or getattr( target_model, "lm_head", None ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 299cac641620..b09e5196b9fb 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -268,7 +268,8 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: @property def attn_vllm_config(self) -> VllmConfig: """Config for the draft's attention metadata builders. Overridden by - speculators whose attention mode differs from the target's.""" + speculators whose attention mode differs from the target's. + """ return self.vllm_config def set_attn( diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 01b6b5b0cc90..79db440cff5b 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -323,7 +323,6 @@ def _register_add_request(self, request: "CachedRequestState") -> int: """Track add-request operations for logits processors. Not applicable to pooling models. """ - # Fill the next empty index if there is one. if (new_req_index := self.batch_update_builder.pop_removed()) is None: # Append to end otherwise. @@ -531,8 +530,8 @@ def remove_request(self, req_id: str) -> int | None: Returns: Removed request index, or `None` if `req_id` not recognized - """ + """ req_index = self.req_id_to_index.pop(req_id, None) if req_index is None: return None @@ -710,6 +709,7 @@ def condense(self) -> None: Returns: swaps: list of (from,to) swap tuples for moved requests empty_req_indices: indices not filled by condensation + """ num_reqs = self.num_reqs @@ -835,7 +835,6 @@ def condense(self) -> None: def refresh_metadata(self): """Apply any batch updates to sampling metadata.""" - if self.is_pooling_model: batch_changed = self.batch_update_builder.reset() if batch_changed: @@ -1001,18 +1000,19 @@ def _make_prompt_token_ids_cpu_tensor(self) -> torch.Tensor: def make_lora_inputs( self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: - """ - Given the num_scheduled_tokens for each request in the batch, return + """Given the num_scheduled_tokens for each request in the batch, return datastructures used to activate the current LoRAs. + Returns: 1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens) - where, prompt_lora_mapping[i] is the LoRA id to use for the ith - sampled token. + where, prompt_lora_mapping[i] is the LoRA id to use for the ith + sampled token. 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) - where, token_lora_mapping[i] is the LoRA id to use for ith token. + where, token_lora_mapping[i] is the LoRA id to use for ith + token. 3. lora_requests: Set of relevant LoRA requests. - """ + """ req_lora_mapping = self.request_lora_mapping[: self.num_reqs] prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens)) token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) @@ -1028,8 +1028,7 @@ def set_async_sampled_token_ids( sampled_token_ids_cpu: torch.Tensor, async_copy_ready_event: torch.Event, ) -> None: - """ - In async scheduling case, store ref to sampled_token_ids_cpu + """In async scheduling case, store ref to sampled_token_ids_cpu tensor and corresponding copy-ready event. Used to repair output_token_ids prior to sampling, if needed by logits processors. """ @@ -1041,8 +1040,7 @@ def set_async_sampled_token_ids( self.async_copy_ready_event = None def update_async_output_token_ids(self) -> None: - """ - In async scheduling case, update output_token_ids in sampling metadata + """In async scheduling case, update output_token_ids in sampling metadata from prior steps sampled token ids once they've finished copying to CPU. This is called right before they are needed by the logits processors. """ @@ -1088,8 +1086,7 @@ def update_async_output_token_ids(self) -> None: # ^ Implicitly resizes to (first_placeholder + num_to_replace) def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None: - """ - In async scheduling case, update spec_token_ids in sampling metadata with + """In async scheduling case, update spec_token_ids in sampling metadata with real draft token ids from prior step. This is called right before they are needed by the rejection sampler for penalty/bad_words computation. """ diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 97a59bd273c3..bf46c3dbfeff 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -481,7 +481,8 @@ def get_output(self) -> ModelRunnerOutput: class ExecuteModelState(NamedTuple): """Ephemeral cached state transferred between execute_model() and - sample_tokens(), after execute_model() returns None.""" + sample_tokens(), after execute_model() returns None. + """ scheduler_output: "SchedulerOutput" logits: torch.Tensor @@ -1003,8 +1004,7 @@ def update_max_model_len(self, max_model_len: int) -> None: self.effective_drafter_max_model_len = self.max_model_len def reset_mm_cache(self) -> None: - """ - Clear the multi-modal cache that was used during profiling, + """Clear the multi-modal cache that was used during profiling, but no longer needed during inference. """ if self.mm_budget: @@ -1072,7 +1072,8 @@ def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor: first scheduled token (column j is position start - 1 - j); -1 where the position is before the prompt or already past it. Generated positions are left to the model: under async scheduling the CPU token - table holds placeholders for them.""" + table holds placeholders for them. + """ buf = self.lookback_token_ids assert buf is not None buf.np.fill(-1) @@ -1136,14 +1137,14 @@ def _init_model_kwargs(self, num_reqs: int | None = None): return model_kwargs def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None: - """ - Update the order of requests in the batch based on the attention + """Update the order of requests in the batch based on the attention backend's needs. For example, some attention backends (namely MLA) may want to separate requests based on if the attention computation will be compute-bound or memory-bound. Args: scheduler_output: The scheduler output. + """ # Attention free models have zero kv_cache_groups, however models # like Mamba are also attention free but use the kv_cache for @@ -1182,7 +1183,6 @@ def _zero_block_ids(self, block_ids: list[int]) -> None: # Note: used for model runner override. def _init_device_properties(self) -> None: """Initialize attributes from torch.cuda.get_device_properties""" - self.num_sms = num_compute_units(self.device.index) # Note: used for model runner override. @@ -1803,7 +1803,6 @@ def _prepare_input_ids( Uses self.prev_positions[:num_reqs] which maps current pos -> prev pos (-1 for new requests). """ - # Sync the async PP broadcast before reading sampled tokens. if self._pp_recv_work is not None: self._pp_recv_work.wait() @@ -1977,9 +1976,9 @@ def _prepare_inputs( SpecDecodeMetadata | None, int, ]: - """ - Returns: - tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens] + """Returns: + tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens] + """ total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens assert total_num_scheduled_tokens > 0 @@ -2310,9 +2309,9 @@ def _build_attention_metadata( cascade_attn_prefix_lens: list[list[int]] | None = None, slot_mappings: dict[int, torch.Tensor] | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: - """ - Returns: - tuple[attn_metadata, spec_decode_common_attn_metadata] + """Returns: + tuple[attn_metadata, spec_decode_common_attn_metadata] + """ # Attention metadata is not needed for attention free models if len(self.kv_cache_config.kv_cache_groups) == 0: @@ -2651,14 +2650,13 @@ def _compute_cascade_attn_prefix_lens( num_computed_tokens: np.ndarray, num_common_prefix_blocks: list[int], ) -> list[list[int]] | None: - """ - Returns: - Optional[cascade_attn_prefix_lens] - cascade_attn_prefix_lens is 2D: - ``[kv_cache_group_id][attn_group_idx]``, - None if we should not use cascade attention - """ + """Returns: + Optional[cascade_attn_prefix_lens] + cascade_attn_prefix_lens is 2D: + ``[kv_cache_group_id][attn_group_idx]``, + None if we should not use cascade attention + """ use_cascade_attn = False num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups) cascade_attn_prefix_lens: list[list[int]] = [ @@ -2703,12 +2701,16 @@ def _compute_cascade_attn_prefix_len( Args: num_scheduled_tokens: Number of tokens scheduled per request. + num_computed_tokens: Number of tokens already computed per request. num_common_prefix_blocks: Number of shared KV cache blocks. + kv_cache_spec: The KV cache spec of the group being considered. + attn_metadata_builder: Builder used to decide whether cascade + attention can be used. Returns: int: Length of common prefix in tokens. - """ + """ common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size if common_prefix_len == 0: # Common case. @@ -2947,6 +2949,7 @@ def _batch_mm_inputs_from_scheduler( - mm_hashes: List of multimodal hashes for each item - mm_kwargs: List of multimodal kwargs for each item - mm_lora_refs: List of (req_id, placeholder_range) for each item + """ scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs if not scheduled_encoder_inputs: @@ -3203,7 +3206,8 @@ def _execute_mm_encoder( def _get_encoder_output_from_cache(self, mm_hash: str) -> torch.Tensor | None: """Return a cached encoder output for multimodal - embedding gather.""" + embedding gather. + """ return self.encoder_cache.get(mm_hash, None) def _gather_mm_embeddings( @@ -3404,9 +3408,7 @@ def sync_and_gather_intermediate_tensors( ) def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None: - """ - Step for the EPLB (Expert Parallelism Load Balancing) state. - """ + """Step for the EPLB (Expert Parallelism Load Balancing) state.""" if not self.parallel_config.enable_eplb or self.eep_eplb_suppressed: return @@ -3875,6 +3877,7 @@ def _model_forward( Returns: Model output tensor + """ return self.model( input_ids=input_ids, @@ -3892,8 +3895,7 @@ def _is_uniform_decode( num_reqs: int, force_uniform_decode: bool | None = None, ) -> bool: - """ - Checks if it's a decode batch with same amount scheduled tokens + """Checks if it's a decode batch with same amount scheduled tokens across all requests. """ return ( @@ -4065,11 +4067,9 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) def _register_layerwise_nvtx_hooks(self) -> None: - """ - Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled + """Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled to trace detailed information of each layer or module in the model. """ - if ( self.vllm_config.observability_config.enable_layerwise_nvtx_tracing and not self.layerwise_nvtx_hooks_registered @@ -4110,8 +4110,7 @@ def _get_slot_mappings( dict[int, torch.Tensor] | None, dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, ]: - """ - Build slot mappings in both formats needed by the system. + """Build slot mappings in both formats needed by the system. Args: num_tokens_padded: Total number of tokens (padded) @@ -4123,6 +4122,7 @@ def _get_slot_mappings( A tuple of: - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext + """ if not ( hasattr(self, "kv_cache_config") @@ -4179,7 +4179,8 @@ def _is_all_reqs_chunked_prefill(self) -> bool: This is true when `discard_request_mask` is set for every scheduled request (e.g., for chunked prefill requests that are not the last - prefill chunk).""" + prefill chunk). + """ num_reqs = self.input_batch.num_reqs return bool(self.discard_request_mask.np[:num_reqs].all()) @@ -5310,9 +5311,9 @@ def update_config(self, overrides: dict[str, Any]) -> None: @instrument(span_name="Loading (GPU)") def load_model(self, load_dummy_weights: bool = False) -> None: - """ - Args: - load_dummy_weights: load dummy weights instead of real weights. + """Args: + load_dummy_weights: load dummy weights instead of real weights. + """ logger.info_once( "Starting to load model %s...", @@ -5525,6 +5526,7 @@ def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: Returns: Tuple of layer indices if found in draft model config, None otherwise. + """ if not (self.speculative_config and self.speculative_config.draft_model_config): return None @@ -5556,8 +5558,7 @@ def reload_weights( weights_path: str | None = None, is_checkpoint_format: bool = True, ) -> None: - """ - Reload weights from a weights iterator or from disk + """Reload weights from a weights iterator or from disk Args: weights_iterator: weights to load into model @@ -5565,6 +5566,7 @@ def reload_weights( provided. Use path of original model if neither is provided. is_checkpoint_format: set to False if weights have already been processed into kernel format (repacking, renaming, etc.) + """ # TODO(@kylesayrs): generalize to all runners and loaders # argument validation @@ -5774,13 +5776,11 @@ def maybe_randomize_inputs( inputs_embeds: torch.Tensor | None, randomize_inputs: bool = False, ): - """ - Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set. + """Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set. This is to help balance expert-selection - during profile_run - during DP rank dummy run """ - dp_size = self.vllm_config.parallel_config.data_parallel_size randomize_inputs = randomize_inputs or ( envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 @@ -5863,8 +5863,7 @@ def _dummy_run( profile_seq_lens: int | None = None, randomize_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Run a dummy forward pass to warm up/profile run or capture the + """Run a dummy forward pass to warm up/profile run or capture the CUDA graph for the model. Args: @@ -5889,6 +5888,12 @@ def _dummy_run( profile_seq_lens: If provided, use this value for seq_lens instead of max_query_len. Used to profile attention workspace that scales with context length. + allow_microbatching: Whether the dummy batch may be split into + microbatches. + is_graph_capturing: Whether this run is a cudagraph capture. + randomize_inputs: Whether to fill the dummy inputs with random + values rather than zeros. + """ mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: @@ -6568,7 +6573,8 @@ def _init_minimal_kv_cache_for_profiling(self) -> None: def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that - memory is reclaimable when running in the same process.""" + memory is reclaimable when running in the same process. + """ from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT from vllm.v1.worker.workspace import reset_workspace_manager @@ -7038,9 +7044,7 @@ def initialize_attn_backend( kv_cache_config: KVCacheConfig, is_profiling: bool = False, ) -> None: - """ - Initialize the attention backends and attention metadata builders. - """ + """Initialize the attention backends and attention metadata builders.""" assert len(self.attn_groups) == 0, "Attention backends are already initialized" class AttentionGroupKey(NamedTuple): @@ -7143,9 +7147,7 @@ def create_attn_groups( def initialize_metadata_builders( self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] ) -> None: - """ - Create the metadata builders for all KV cache groups and attn groups. - """ + """Create the metadata builders for all KV cache groups and attn groups.""" for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)): for attn_group in self.attn_groups[kv_cache_group_id]: attn_group.create_metadata_builders( @@ -7180,8 +7182,7 @@ def _check_and_update_cudagraph_mode( kv_cache_groups: list[KVCacheGroupSpec], is_profiling: bool = False, ) -> None: - """ - Resolve the cudagraph_mode when there are multiple attention + """Resolve the cudagraph_mode when there are multiple attention groups with potential conflicting CUDA graph support. Then initialize the cudagraph_dispatcher based on the resolved cudagraph_mode. @@ -7234,8 +7235,7 @@ def _check_and_update_cudagraph_mode( self.drafter.initialize_cudagraph_keys(cudagraph_mode) def calculate_reorder_batch_threshold(self) -> None: - """ - Choose the minimum reorder batch threshold from all attention groups. + """Choose the minimum reorder batch threshold from all attention groups. Backends should be able to support lower threshold then what they request just may have a performance penalty due to that backend treating decodes as prefills. @@ -7256,8 +7256,7 @@ def calculate_reorder_batch_threshold(self) -> None: def may_reinitialize_input_batch( self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] ) -> None: - """ - Re-initialize the input batch if the block sizes are different from + """Re-initialize the input batch if the block sizes are different from what it was originally created with. This happens when the final block size (determined after model loading) differs from the placeholder used during __init__, or when there are multiple @@ -7266,6 +7265,7 @@ def may_reinitialize_input_batch( Args: kv_cache_config: The KV cache configuration. kernel_block_sizes: The kernel block sizes for each KV cache group. + """ block_sizes = [] max_num_blocks = [] @@ -7347,18 +7347,19 @@ def initialize_kv_cache_tensors( kernel_block_sizes: list[int], kv_cache_allocation_context: AbstractContextManager | None = None, ) -> dict[str, torch.Tensor]: - """ - Initialize the memory buffer for KV cache. + """Initialize the memory buffer for KV cache. Args: kv_cache_config: The KV cache config kernel_block_sizes: The kernel block sizes for each KV cache group. + kv_cache_allocation_context: Context manager the allocation runs + inside, e.g. to place the KV cache in a specific memory pool. Returns: Dict[str, torch.Tensor]: A map between layer names to their corresponding memory buffer for KV cache. - """ + """ allocation_context = kv_cache_allocation_context or nullcontext() with allocation_context: kv_caches = allocate_kv_cache( @@ -7388,8 +7389,7 @@ def initialize_kv_cache_tensors( def maybe_add_kv_sharing_layers_to_kv_cache_groups( self, kv_cache_config: KVCacheConfig ) -> None: - """ - Add layers that re-use KV cache to KV cache group of its target layer. + """Add layers that re-use KV cache to KV cache group of its target layer. Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()` """ if not self.shared_kv_cache_layers: @@ -7419,11 +7419,15 @@ def initialize_kv_cache( is_profiling: bool = False, kv_cache_allocation_context: AbstractContextManager | None = None, ) -> None: - """ - Initialize KV cache based on `kv_cache_config`. + """Initialize KV cache based on `kv_cache_config`. + Args: kv_cache_config: Configuration for the KV cache, including the KV - cache size of each layer + cache size of each layer + is_profiling: Whether this call is part of a profiling run. + kv_cache_allocation_context: Context manager the allocation runs + inside, e.g. to place the KV cache in a specific memory pool. + """ kv_cache_config = deepcopy(kv_cache_config) self.kv_cache_config = kv_cache_config @@ -7530,9 +7534,7 @@ def init_routed_experts_capturer(self): self.routed_experts_initialized = True def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: - """ - Add encoder-only layers to the KV cache config. - """ + """Add encoder-only layers to the KV cache config.""" block_size = self.vllm_config.cache_config.block_size encoder_only_attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list) attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention) @@ -7556,12 +7558,13 @@ def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: ) def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: - """ - Generates the KVCacheSpec by parsing the kv cache format from each + """Generates the KVCacheSpec by parsing the kv cache format from each Attention module in the static forward context. + Returns: KVCacheSpec: A dictionary mapping layer names to their KV cache format. Layers that do not need KV cache are not included. + """ if has_ec_transfer() and not get_ec_transfer().is_consumer: return {} @@ -7605,11 +7608,11 @@ def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: return pinned.tolist() def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: - """ - Get encoder timing stats for all requests and clear the registry. + """Get encoder timing stats for all requests and clear the registry. Returns: Dictionary mapping request_id to stats dict. + """ with self._encoder_timing_lock: stats = { @@ -7627,14 +7630,14 @@ def timed_encoder_operation( current_item_idx: int, num_items: int, ): - """ - Context manager to time encoder forward operations. + """Context manager to time encoder forward operations. Args: should_time: Whether timing is enabled group_lora_refs: Full list of (request_id, pos_info) tuples current_item_idx: Starting index for this group num_items: Number of items in this group + """ if not should_time: yield diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 0735b39f81c2..a0f2643052ae 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -119,8 +119,7 @@ def unwrap(self) -> Callable: return self.runnable def _capture_ubatches(self, ubatch_metadata, model) -> torch.Tensor: - """ - Capture a cudagraph for a microbatched run. + """Capture a cudagraph for a microbatched run. The logic here is somewhat complicated because we need to make sure that each of the ubatch threads initialize the cuda context before we start diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 2a3531b61b38..44621d08fb20 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -336,7 +336,8 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: def _scoped_allocator_max_split(self, max_split_size_mb: int): """Temporarily set max_split_size_mb to reduce allocator fragmentation at the cost of more cudaMalloc calls (negligible in practice). Restores the original - value on exit.""" + value on exit. + """ if not current_platform.is_cuda(): yield return @@ -536,6 +537,7 @@ def determine_available_memory(self) -> int: Tip: You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. + """ maybe_apply_startup_plan(self) @@ -694,7 +696,6 @@ def get_kv_connector_handshake_metadata( Returned dict is keyed by `(pp_rank, tp_rank)`. """ - if not has_kv_transfer_group(): return None @@ -734,7 +735,6 @@ def update_max_model_len(self, max_model_len: int) -> None: @instrument(span_name="Allocate KV cache") def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: """Allocate GPU KV cache with the specified kv_cache_config.""" - # Update local config with adjusted num blocks after profiling, # so that it's available to the warmup stage. self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks @@ -1364,12 +1364,12 @@ def _check_weight_transfer_engine(self) -> None: ) def init_weight_transfer_engine(self, init_info: dict) -> None: - """ - Initialize weight transfer mechanism. + """Initialize weight transfer mechanism. For NCCL backend, this creates a process group with the trainer. Args: init_info: Dictionary containing backend-specific initialization info + """ self._check_weight_transfer_engine() assert self.weight_transfer_engine is not None @@ -1378,8 +1378,7 @@ def init_weight_transfer_engine(self, init_info: dict) -> None: self.weight_transfer_engine.init_transfer_engine(typed_init_info) def start_weight_update(self) -> None: - """ - Start a new weight update session. + """Start a new weight update session. Delegates engine-specific preparation (e.g. layerwise reload setup) to the configured weight transfer engine. The worker only tracks that a @@ -1389,8 +1388,7 @@ def start_weight_update(self) -> None: self._start_weight_update() def start_draft_weight_update(self) -> None: - """ - Like start_weight_update, but retargets the engine at the speculative + """Like start_weight_update, but retargets the engine at the speculative draft model for this session. """ with set_current_vllm_config(self.vllm_config): @@ -1423,8 +1421,7 @@ def _start_weight_update(self, is_draft: bool = False) -> None: self._weight_update_is_draft = is_draft def update_weights(self, update_info: dict | list[dict]) -> None: - """ - Receive one weight update chunk from the trainer. + """Receive one weight update chunk from the trainer. start_weight_update must be called before update_weights and finish_weight_update must be called after all chunks have been sent. @@ -1434,6 +1431,7 @@ def update_weights(self, update_info: dict | list[dict]) -> None: Args: update_info: Backend-specific update info, or a list indexed by global worker rank across data parallel replicas. + """ self._check_weight_transfer_engine() assert self.weight_transfer_engine is not None diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 6054bb129aba..26ffc159f3c0 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Define KV connector functionality mixin for model runners. -""" +"""Define KV connector functionality mixin for model runners.""" from collections.abc import Generator from contextlib import AbstractContextManager, contextmanager, nullcontext diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 5e7f1f70c56d..9d94e62edb84 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/vllm/v1/worker/lora_model_runner_mixin.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Define LoRA functionality mixin for model runners. -""" +"""Define LoRA functionality mixin for model runners.""" from collections.abc import Callable from contextlib import contextmanager @@ -150,16 +148,17 @@ def maybe_select_dummy_loras( num_sampled_tokens: np.ndarray | None = None, num_active_loras: int = 0, ): - """ - Context manager to select dummy LoRAs for capture/warmup. + """Context manager to select dummy LoRAs for capture/warmup. Args: lora_config: LoRA configuration, or None if LoRA is disabled. num_scheduled_tokens: Array of scheduled token counts per request. + mapping_type: Which LoRA mapping to build (language or encoder). num_sampled_tokens: Array of sampled token counts per request. num_active_loras: Number of distinct active LoRAs to use. - 0: No LoRA active (set up zero mappings). - >0: Use exactly this many distinct LoRAs. + """ if num_sampled_tokens is None: num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32) @@ -255,8 +254,7 @@ def maybe_dummy_run_with_lora( num_active_loras: int = 0, mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE, ): - """ - Context manager for dummy runs with LoRA. + """Context manager for dummy runs with LoRA. Args: lora_config: LoRA configuration. @@ -265,6 +263,8 @@ def maybe_dummy_run_with_lora( remove_lora: Whether to remove LoRAs after the context exits. num_active_loras: Number of distinct active LoRAs to use. LoRA is activated when num_active_loras > 0. + mapping_type: Which LoRA mapping to build (language or encoder). + """ with ( self.maybe_setup_dummy_loras(lora_config, remove_lora), diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 1deb89bda65c..c2e0120e4613 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -415,8 +415,7 @@ def postprocess_mamba_fused_kernel( # the existing 2D-grid contract. TEMPORAL_TILES: tl.constexpr = 1, ): - """ - Fused GPU kernel for postprocess_mamba that computes decisions AND performs + """Fused GPU kernel for postprocess_mamba that computes decisions AND performs mamba state copies without any CPU-GPU synchronization. Grid: (num_reqs, num_states [, TEMPORAL_TILES]) @@ -769,8 +768,7 @@ def create( @dataclasses.dataclass class MambaSpecDecodeGPUContext: - """ - Context for GPU-side Mamba state copy operations during the + """Context for GPU-side Mamba state copy operations during the fused postprocess path. Only used when speculative decoding is enabled on a hybrid model @@ -927,8 +925,7 @@ def initialize_from_forward_context( mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], ) -> None: - """ - Extract and cache memory layout metadata from Mamba state tensors. + """Extract and cache memory layout metadata from Mamba state tensors. This method populates the pre-allocated metadata tensors with information needed by `postprocess_mamba_fused_kernel` to perform state copies entirely @@ -960,6 +957,7 @@ def initialize_from_forward_context( block_tables: per-mamba-group persistent block-table tensors, in the same order as `mamba_group_ids`. Their `data_ptr()` / `stride(0)` are captured once for the kernel to index into. + """ if self.is_initialized: return @@ -1096,7 +1094,7 @@ def compute_aligned_state_indices( seq_lens: torch.Tensor, num_reqs: int, ) -> torch.Tensor: - """compute every Mamba group's aligned physical state IDs in one launch.""" + """Compute every Mamba group's aligned physical state IDs in one launch.""" assert self.is_initialized assert seq_lens.is_cuda assert 0 <= num_reqs <= seq_lens.shape[0] @@ -1137,8 +1135,7 @@ def run_fused_postprocess( num_computed_tokens_gpu: torch.Tensor, num_draft_tokens_gpu: torch.Tensor, ) -> None: - """ - Run the fused postprocess_mamba kernel on GPU. + """Run the fused postprocess_mamba kernel on GPU. This computes decisions and performs mamba state copies entirely on GPU, eliminating the CPU-GPU sync that was previously needed. @@ -1150,6 +1147,7 @@ def run_fused_postprocess( num_scheduled_tokens_gpu: [num_reqs] scheduled token counts num_computed_tokens_gpu: [num_reqs] computed token counts num_draft_tokens_gpu: [num_reqs] draft token counts + """ if num_reqs == 0 or not self.is_initialized: return @@ -1205,6 +1203,7 @@ def run_fused_precopy( token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1). idx_mapping: optional [num_reqs] batch_idx -> req_state_idx. None means V1 batch order already equals request state order. + """ if num_reqs == 0 or not self.is_initialized: return @@ -1442,8 +1441,7 @@ def preprocess_mamba( copy_bufs: MambaCopyBuffers, align_ctx: MambaSpecDecodeGPUContext | None = None, ): - """ - Copy the mamba state of previous step to the last + """Copy the mamba state of previous step to the last (1 + num_speculative_blocks) block. """ fused = _resolve_fused_precopy(align_ctx) diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py index 2c4942070705..5a570b347094 100644 --- a/vllm/v1/worker/startup_plan.py +++ b/vllm/v1/worker/startup_plan.py @@ -77,7 +77,8 @@ def compute_plan_fingerprint( def _plan_path(fingerprint: str) -> str: """Plans are regenerable derived state, so they live under the standard vLLM cache root (like the torch.compile cache) and relocate with - ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own.""" + ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own. + """ # VLLM_CACHE_ROOT is already user-expanded by envs.py. return os.path.join( envs.VLLM_CACHE_ROOT, "startup_plan", f"startup_plan_{fingerprint}.json" @@ -134,7 +135,8 @@ def _applicable_kv_cache_memory_bytes( def maybe_apply_startup_plan(worker: "Worker") -> None: """If enabled and ``--kv-cache-memory`` was not set explicitly, apply a persisted plan by setting ``worker.cache_config.kv_cache_memory_bytes``. - No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``.""" + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``. + """ if ( not envs.VLLM_ENABLE_STARTUP_PLAN or worker.cache_config.kv_cache_memory_bytes is not None @@ -167,7 +169,8 @@ def maybe_apply_startup_plan(worker: "Worker") -> None: def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> None: """Atomically persist this boot's profiling result for future boots. No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``; failures are logged, - never raised.""" + never raised. + """ if not envs.VLLM_ENABLE_STARTUP_PLAN: return fingerprint = compute_plan_fingerprint( diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index 9743322c4165..329c261a0a38 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -43,8 +43,7 @@ def __init__( set_comm_sms: Callable[[int], None], set_compute_sms: Callable[[int], None], ): - """ - Context manager for controlling SM (Streaming Multiprocessor) + """Context manager for controlling SM (Streaming Multiprocessor) allocation. Upon entering the context, it sets the number of SMs allocated for communication and computation to comm_sms and total_sms - comm_sms respectively. Upon exiting, it restores the @@ -57,8 +56,8 @@ def __init__( A function that sets the number of SMs for communication. set_compute_sms (Callable[[int], None]): A function that sets the number of SMs for computation. - """ + """ assert current_platform.is_cuda() or current_platform.is_rocm(), ( "SM/CU control is supported on CUDA and ROCm platforms" ) @@ -221,8 +220,7 @@ def slice_query_start_locs( query_start_loc: torch.Tensor, request_slice: slice, ) -> torch.Tensor: - """ - Creates a new query_start_loc that corresponds to the requests in + """Creates a new query_start_loc that corresponds to the requests in request_slice. Note: This function creates a new tensor to hold the new query_start_locs. @@ -237,11 +235,9 @@ def slice_query_start_locs( def _make_metadata_with_slice( ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata ) -> CommonAttentionMetadata: - """ - This function creates a new CommonAttentionMetadata that corresponds to + """This function creates a new CommonAttentionMetadata that corresponds to the requests included in ubatch_slice """ - assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty" request_slice = ubatch_slice.request_slice @@ -338,8 +334,7 @@ def split_attn_metadata( ubatch_slices: list[UBatchSlice], common_attn_metadata: CommonAttentionMetadata, ) -> list[CommonAttentionMetadata]: - """ - Creates a new CommonAttentionMetadata instance that corresponds to the + """Creates a new CommonAttentionMetadata instance that corresponds to the requests for each UBatchSlice in ubatch_slices. Note: This function does not modify common_attn_metadata diff --git a/vllm/v1/worker/ubatching.py b/vllm/v1/worker/ubatching.py index 6d06ba2135f5..a773ee715d14 100644 --- a/vllm/v1/worker/ubatching.py +++ b/vllm/v1/worker/ubatching.py @@ -18,9 +18,7 @@ class UBatchContext: - """ - Context manager for micro-batching synchronization using threading events. - """ + """Context manager for micro-batching synchronization using threading events.""" def __init__( self, diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index df2ebc7d36ff..6c2a05a6260d 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -331,8 +331,7 @@ def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], ) -> int: - """ - Select a block size that is supported by all backends and is a factor of + """Select a block size that is supported by all backends and is a factor of kv_manager_block_size. If kv_manager_block_size is supported by all backends, return it directly. @@ -347,6 +346,7 @@ def select_common_block_size( Raises: ValueError: If no valid block size found. + """ def block_size_is_supported( @@ -458,8 +458,7 @@ def allocate_kv_cache( def prepare_kernel_block_sizes( kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]] ) -> list[int]: - """ - Generate kernel_block_sizes that matches each block_size. + """Generate kernel_block_sizes that matches each block_size. For attention backends that support virtual block splitting, use the supported block sizes from the backend. @@ -471,6 +470,7 @@ def prepare_kernel_block_sizes( Returns: List of kernel block sizes for each cache group. + """ kernel_block_sizes = [] for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups): @@ -505,8 +505,7 @@ def sanity_check_mm_encoder_outputs( mm_embeddings: MultiModalEmbeddings, expected_num_items: int, ) -> None: - """ - Perform sanity checks for the result of + """Perform sanity checks for the result of [`vllm.model_executor.models.SupportsMultiModal.embed_multimodal`][]. """ assert isinstance(mm_embeddings, (list, tuple, torch.Tensor)), ( @@ -532,8 +531,7 @@ def sanity_check_mm_encoder_outputs( def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int: - """ - Calculate the amount of memory required by vLLM, then validate + """Calculate the amount of memory required by vLLM, then validate that the current amount of free memory is sufficient for that. """ requested_memory = math.ceil( @@ -559,8 +557,7 @@ def add_kv_sharing_layers_to_kv_cache_groups( kv_cache_groups: list[KVCacheGroupSpec], runner_only_attn_layers: set[str] | None = None, ) -> None: - """ - Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches` + """Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches` for layers that do not allocate its own KV cache, based on the mapping in `shared_kv_cache_layers`. Adds these layers to the corresponding KV cache group, which is needed to ensure that attention metadata is assigned later. @@ -571,6 +568,9 @@ def add_kv_sharing_layers_to_kv_cache_groups( means this layer will perform attention using the keys and values from the KV cache of `shared_kv_cache_layers[layer_name]`. kv_cache_groups: The KV cache groups of the model. + runner_only_attn_layers: Attention layers handled by the runner only, + which are excluded from the KV cache groups. + """ if not shared_kv_cache_layers: return @@ -595,8 +595,7 @@ def bind_kv_cache( num_attn_module: int = 1, kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, ) -> None: - """ - Bind the allocated KV cache to both ModelRunner and forward context so + """Bind the allocated KV cache to both ModelRunner and forward context so that the KV cache can be used in the forward pass. This function: @@ -607,9 +606,13 @@ def bind_kv_cache( Args: kv_caches: The allocated kv_caches with layer names as keys. + num_attn_module: Number of attention modules per layer entry. forward_context: The global forward context containing all Attention layers with layer names as keys. runner_kv_caches: The kv_cache declared by ModelRunner. + kv_cache_groups: The KV cache groups of the model, used to resolve + layers that share a KV cache. + """ # Bind kv_caches to ModelRunner assert len(runner_kv_caches) == 0 diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index cf5b0e7ce334..6496a611e6fb 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -55,8 +55,7 @@ def __init__( distributed_init_method: str, is_driver_worker: bool = False, ) -> None: - """ - Initialize common worker components. + """Initialize common worker components. Args: vllm_config: Complete vLLM configuration @@ -65,6 +64,7 @@ def __init__( distributed_init_method: Distributed initialization method is_driver_worker: Whether this worker handles driver responsibilities + """ self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -118,6 +118,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: Returns: Compilation times (language_model, encoder) in seconds. + """ raise NotImplementedError @@ -127,7 +128,8 @@ def check_health(self) -> None: def synchronize_device(self) -> None: """Block until in-flight device work completes; backends outside - ``torch.accelerator`` must override with their own wait.""" + ``torch.accelerator`` must override with their own wait. + """ if torch.accelerator.is_available(): torch.accelerator.synchronize() @@ -209,8 +211,7 @@ def shutdown(self) -> None: class WorkerWrapperBase: - """ - This class represents one process in an executor/engine. It is responsible + """This class represents one process in an executor/engine. It is responsible for lazily initializing the worker and handling the worker's lifecycle. We first instantiate the WorkerWrapper, which remembers the worker module and class name. Then, when we call `update_environment_variables`, and the @@ -222,8 +223,7 @@ def __init__( rpc_rank: int = 0, global_rank: int | None = None, ) -> None: - """ - Initialize the worker wrapper with the given vllm_config and rpc_rank. + """Initialize the worker wrapper with the given vllm_config and rpc_rank. Note: rpc_rank is the rank of the worker in the executor. In most cases, it is also the rank of the worker in the distributed group. However, when multiple executors work together, they can be different. @@ -252,8 +252,7 @@ def update_environment_variables( @instrument(span_name="Worker init") def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: - """ - Here we inject some common logic before initializing the worker. + """Here we inject some common logic before initializing the worker. Arguments are passed to the worker class constructor. """ kwargs = all_kwargs[self.rpc_rank] diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 39b3b10349c3..cddb760acd87 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -123,6 +123,7 @@ def get_simultaneous( Returns: List of tensor views into the workspace buffer, one per shape/dtype pair. + """ actual_bytes = [_compute_bytes(s, d) for s, d in shapes_and_dtypes] aligned_bytes = [round_up(actual, 256) for actual in actual_bytes] @@ -148,6 +149,7 @@ def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: Returns: The current workspace tensor. + """ ubatch_id = dbo_current_ubatch_id() lane = _workspace_lane.get() @@ -228,6 +230,7 @@ def is_workspace_manager_initialized() -> bool: Returns: True if workspace manager is initialized, False otherwise. + """ return _manager is not None @@ -237,6 +240,7 @@ def current_workspace_manager() -> "WorkspaceManager": Raises: AssertionError: If workspace manager has not been initialized. + """ assert _manager is not None, ( "WorkspaceManager not initialized. Call init_workspace_manager() " @@ -259,6 +263,7 @@ def init_workspace_manager( device: The device to allocate workspace on. num_ubatches: Number of workspace ubatch slots. Defaults to 1. num_lanes: Number of independent execution lanes per ubatch. Defaults to 1. + """ global _manager if _manager is not None: @@ -289,6 +294,7 @@ def lock_workspace() -> None: lock_workspace() # Now all get_workspace calls must fit in pre-allocated size + """ current_workspace_manager().lock() diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 4ba4692f6778..c10b381e40f2 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -261,11 +261,41 @@ def flash_attn_varlen_func( return_attn_probs: bool. Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). + seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per + sequence. Mutually exclusive with cu_seqlens_k, and required when + block_table is given. + q_v: optional value tensor fused with q. FA3 only. + block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV + cache block table. Requires seqused_k. + return_softmax_lse: bool. Whether to also return softmax_lse. + out: optional tensor to write the output into. + scheduler_metadata: precomputed scheduling metadata. Not supported by FA2. + q_descale: descale factor for FP8 q. Not supported by FA2, and ignored by + FA4 when the inputs are not FP8. + k_descale: descale factor for FP8 k. Same constraints as q_descale. + v_descale: descale factor for FP8 v. Same constraints as q_descale. + num_splits: int. Number of splits for the split-KV kernel; 0 lets the kernel + choose. FA2 only supports 0 or 1. + output_scale: scale for fused FP8 output quantization. FA4 only. + fa_version: int. FlashAttention major version to dispatch to. + s_aux: learnable per-head attention sink logits. Not supported by FA2. + cp_world_size: int. Context parallel world size. + cp_rank: int. Rank of this shard within the context parallel group. + cp_tot_seqused_k: (batch_size,), dtype torch.int32. Key lengths across the + whole context parallel group. + mask_mod: optional callable applying a custom mask. Not supported by FA2. + block_sparse_tensors: block-sparse index tensors. FA4 only, and must + materialize full_block_cnt and full_block_idx. + aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only. + aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors. + dynamic_causal: optional per-sequence causal offsets. FA4 only. + Return: out: (total, nheads, headdim). softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax normalization factor). + """ assert cu_seqlens_k is not None or seqused_k is not None, ( "cu_seqlens_k or seqused_k must be provided" @@ -577,11 +607,16 @@ def sparse_attn_func( return_attn_probs: bool. Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). + softcap: float. Anything > 0 activates softcapping attention. + return_softmax_lse: bool. Whether to also return softmax_lse. + out: optional tensor to write the output into. + Return: out: (batch_size, seqlen, nheads, headdim). softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax normalization factor). + """ if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -663,11 +698,15 @@ def sparse_attn_varlen_func( return_attn_probs: bool. Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). + return_softmax_lse: bool. Whether to also return softmax_lse. + out: optional tensor to write the output into. + Return: out: (total, nheads, headdim). softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax normalization factor). + """ if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) From 6ca2b23e22aab8534e00a85e8ec5222508a6ecf3 Mon Sep 17 00:00:00 2001 From: Rebecca Lee <110500452+rebklee@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:48:30 -0700 Subject: [PATCH 334/774] [ROCm] Expose kFp8DynamicTokenSym on AITER PTPC linears (#54248) Signed-off-by: Rebecca Lee Co-authored-by: Cursor --- .../fusion/test_quant_activation_contract.py | 8 +++ tests/quantization/test_quark.py | 65 ++++++++++++++++++- .../kernels/linear/scaled_mm/aiter.py | 16 +++++ .../layers/quantization/online/fp8.py | 11 +++- .../quark/schemes/quark_w8a8_fp8.py | 8 ++- 5 files changed, 104 insertions(+), 4 deletions(-) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index ec9bb6110dcd..7e346cc76721 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -24,6 +24,11 @@ FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( CutlassFP8ScaledMMLinearKernel, ) @@ -56,6 +61,9 @@ FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, } diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 9f00adf78314..ced034cb1a94 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -21,9 +21,18 @@ from tests.quantization.utils import load_model_without_vllm_runner from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops -from vllm.config import set_current_vllm_config +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.cache import CacheConfig from vllm.forward_context import set_forward_context +from vllm.model_executor import parameter +from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoeWeightScaleSupported, @@ -113,6 +122,12 @@ def on_gfx950() -> bool: AITER_AVAILABLE = is_aiter_found_and_supported() +AITER_PTPC_KERNELS = ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, +) + DEVICE_TYPE = current_platform.device_type @@ -1095,6 +1110,54 @@ def test_quark_w8a8_fp8_moe_per_block_rejects_misaligned_partition(): ) +def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch): + """QuarkW8A8Fp8 must advertise the key its kernel consumes pre-quantized. + + Only kernel selection runs, no GEMM. The shape matters: AITER PTPC + kernels decline untuned shapes, leaving a torch kernel that has no key. + """ + # Llama-3.1-70B qkv_proj at TP1 (aiter-tuned N, K on gfx950). + N, K = 10240, 8192 + dtype = torch.bfloat16 + kernel_config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticChannelSym, + activation_quant_key=kFp8DynamicTokenSym, + weight_shape=(N, K), + input_dtype=dtype, + out_dtype=dtype, + ) + if not any( + cls.is_supported()[0] and cls.can_implement(kernel_config)[0] + for cls in AITER_PTPC_KERNELS + ): + pytest.skip("no AITER PTPC kernel is usable for this shape and environment") + + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + + scheme = QuarkW8A8Fp8.__new__(QuarkW8A8Fp8) + scheme.weight_qscheme = "per_channel" + scheme.is_static_input_scheme = False + scheme.input_qscheme = "per_channel" + scheme.activation_quant_key = kFp8DynamicTokenSym + scheme.weight_quant_key = kFp8StaticChannelSym + scheme.out_dtype = dtype + scheme.input_dtype = dtype + + layer = torch.nn.Module() + with set_current_vllm_config(VllmConfig()): + scheme.create_weights( + layer, + output_partition_sizes=[N], + input_size_per_partition=K, + params_dtype=dtype, + weight_loader=lambda *args, **kwargs: None, + ) + + assert isinstance(scheme.fp8_linear, AITER_PTPC_KERNELS) + assert layer.input_quant_key == kFp8DynamicTokenSym + + @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) def test_quark_fp8_w_per_tensor_a_per_tensor( kv_cache_dtype: str, monkeypatch, dist_init, workspace_init diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 2408f4ad6cfc..a57e7510de01 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -16,6 +16,8 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, + kFp8DynamicTokenSym, ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -150,6 +152,12 @@ def is_supported( return False, "requires aiter library to be installed." return True, None + def input_quant_key(self) -> QuantKey | None: + # Does not call get_output_padding() - torch fallbacks + # resolve padding from compilation_config, which is unset in + # profile_run. This kernel does not pad activations. + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( @@ -242,6 +250,11 @@ def is_supported( return True, None + def input_quant_key(self) -> QuantKey | None: + # Same per-token FP8 consume path as + # AiterPreshuffledPerTokenFp8ScaledMMLinearKernel. + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( @@ -319,6 +332,9 @@ def is_supported( compute_capability ) + def input_quant_key(self) -> QuantKey | None: + return kFp8DynamicTokenSym + @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: is_ptpc = ( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index c1524e78b6eb..0e12799304fb 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -27,6 +27,10 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( select_fp8_moe_backend, ) +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) from vllm.model_executor.layers.linear import ( LinearMethodBase, ) @@ -405,6 +409,7 @@ def process_weights_after_loading(self, layer: Module) -> None: replace_parameter(layer, "weight_scale", weight_scale) self.fp8_linear.process_weights_after_loading(layer) + expose_input_quant_key(layer, self.fp8_linear) layer._already_called_process_weights_after_loading = True @@ -415,8 +420,10 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: # if batch invariant mode is enabled dequant - if envs.VLLM_BATCH_INVARIANT and not isinstance( - self.fp8_linear, CutlassFP8ScaledMMLinearKernel + if ( + envs.VLLM_BATCH_INVARIANT + and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) + and not isinstance(x, QuantizedActivation) ): weight_dequant = ( layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t() diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py index 57b08361284a..1ef7c1dcc9af 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -9,6 +9,10 @@ from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, @@ -195,10 +199,12 @@ def create_weights( module_name=self.__class__.__name__, ) + expose_input_quant_key(layer, self.fp8_linear) + def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: return self.fp8_linear.apply_weights(layer, x, bias) From 0b9e018fb941652441e5d661ef5b8397e4909ebd Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Thu, 17 Sep 2026 01:21:00 +0800 Subject: [PATCH 335/774] [MM] Further cleanup _apply_hf_processor_main (#53610) Signed-off-by: DarkLight1337 --- docs/contributing/model/multimodal.md | 33 ++- docs/design/mm_processing.md | 11 +- vllm/model_executor/models/audioflamingo3.py | 14 +- .../models/bailing_moe_v3_vl.py | 2 +- vllm/model_executor/models/cohere2_vision.py | 2 +- vllm/model_executor/models/cohere_asr.py | 48 ++--- vllm/model_executor/models/colmodernvbert.py | 14 +- vllm/model_executor/models/colpali.py | 25 +-- vllm/model_executor/models/deepseek_ocr.py | 37 +--- vllm/model_executor/models/deepseek_ocr2.py | 38 +--- vllm/model_executor/models/deepseek_vl2.py | 6 +- vllm/model_executor/models/ernie45_vl.py | 107 +++++----- vllm/model_executor/models/fireredasr2.py | 40 ++-- vllm/model_executor/models/funasr.py | 46 ++-- vllm/model_executor/models/funaudiochat.py | 23 +- vllm/model_executor/models/gemma3_mm.py | 2 +- vllm/model_executor/models/gemma3n_mm.py | 13 +- vllm/model_executor/models/gemma4_mm.py | 43 ++-- vllm/model_executor/models/glm4_1v.py | 53 +++-- vllm/model_executor/models/glmasr.py | 29 ++- vllm/model_executor/models/granite_speech.py | 24 +-- .../models/granite_speech_plus.py | 4 +- vllm/model_executor/models/h2ovl.py | 4 +- .../models/hyperclovax_vision_v2.py | 40 +--- vllm/model_executor/models/idefics3.py | 42 ++-- vllm/model_executor/models/interns1.py | 38 ++-- vllm/model_executor/models/jina_vl.py | 15 +- vllm/model_executor/models/kanana_v.py | 24 ++- vllm/model_executor/models/keye.py | 2 +- vllm/model_executor/models/keye_vl1_5.py | 2 +- vllm/model_executor/models/kimi_audio.py | 32 +-- vllm/model_executor/models/kimi_vl.py | 2 +- vllm/model_executor/models/lfm2_vl.py | 2 +- vllm/model_executor/models/lightonocr.py | 28 +-- vllm/model_executor/models/llava.py | 2 +- vllm/model_executor/models/llava_next.py | 2 +- vllm/model_executor/models/llava_onevision.py | 46 ++-- .../model_executor/models/llava_onevision2.py | 70 +++--- vllm/model_executor/models/midashenglm.py | 49 +++-- vllm/model_executor/models/mimo_v2_omni.py | 38 +--- vllm/model_executor/models/minicpmo.py | 2 +- vllm/model_executor/models/minicpmv.py | 14 +- vllm/model_executor/models/mistral3.py | 4 +- vllm/model_executor/models/mllama4.py | 2 +- vllm/model_executor/models/molmo.py | 41 ++-- vllm/model_executor/models/molmo2.py | 41 ++-- vllm/model_executor/models/moss_audio.py | 27 +-- .../models/moss_transcribe_diarize.py | 18 +- vllm/model_executor/models/muse_glimmer.py | 17 +- .../model_executor/models/nano_nemotron_vl.py | 31 ++- vllm/model_executor/models/ovis.py | 6 +- vllm/model_executor/models/ovis2_5.py | 2 +- vllm/model_executor/models/paddleocr_vl.py | 39 ++-- vllm/model_executor/models/phi3v.py | 21 +- vllm/model_executor/models/phi4mm.py | 66 +++--- vllm/model_executor/models/phi4siglip.py | 35 +-- vllm/model_executor/models/pixtral.py | 53 ++--- .../models/qwen2_5_omni_thinker.py | 166 +++++++------- vllm/model_executor/models/qwen2_audio.py | 33 ++- .../models/qwen3_asr_realtime.py | 11 +- .../models/qwen3_omni_moe_thinker.py | 64 +++--- vllm/model_executor/models/qwen3_vl.py | 38 ++-- vllm/model_executor/models/terratorch.py | 4 +- .../models/transformers/multimodal.py | 79 ++++--- vllm/model_executor/models/ultravox.py | 83 ++----- vllm/model_executor/models/unlimited_ocr.py | 4 +- vllm/model_executor/models/voxtral.py | 57 ++--- .../model_executor/models/voxtral_realtime.py | 11 +- vllm/model_executor/models/whisper.py | 44 ++-- vllm/models/dots3_note/common/processor.py | 44 ++-- vllm/models/inkling/common/mm_preprocess.py | 16 +- .../models/minimax_m3/common/mm_preprocess.py | 50 ++--- vllm/multimodal/processing/context.py | 10 +- vllm/multimodal/processing/processor.py | 202 ++++++++++-------- 74 files changed, 1056 insertions(+), 1331 deletions(-) diff --git a/docs/contributing/model/multimodal.md b/docs/contributing/model/multimodal.md index 951d25587ad9..d374a5ecbc27 100644 --- a/docs/contributing/model/multimodal.md +++ b/docs/contributing/model/multimodal.md @@ -454,17 +454,17 @@ return a schema of the tensors outputted by the HF processor that are related to ??? code ```python - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: # Mistral3Processor requires text corresponding to the images return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], processed_data: BatchFeature, ) -> BatchFeature: - if not mm_data: + if not hf_data: return processed_data pixel_values = processed_data.get("pixel_values") @@ -486,25 +486,36 @@ return a schema of the tensors outputted by the HF processor that are related to calls the HF processor on the multi-modal data without passing any text. If the HF processor instead requires text corresponding to the multi-modal items, you should override - [_get_hf_processor_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_processor_text] + [_get_hf_mm_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_text] to return the dummy text from [BaseDummyInputsBuilder.get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text] like in the example above. If the HF processor expects the multi-modal data - under different keys than those provided by the multi-modal items - (e.g. `audio` instead of `audios`), or requires additional keyword arguments + under different keys than those provided by the multi-modal items, + or requires additional keyword arguments (e.g. `sampling_rate`), you should override - [_preprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._preprocess_hf_mm_data]. + [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs]. If you need to modify the output of the HF processor, you should override [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data]. + Passthrough fields have already been merged into `processed_data` when this + hook is called. + If you need to change how the normalized inputs are passed to the HF + processor (for example, to call a method other than `__call__`, or to use + different kwargs to construct and invoke it), override + [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor]. If the HF processor applies additional transformations to the prompt tokens that are not reflected in its multi-modal outputs (e.g. prepending a BOS token), you should override [_postprocess_prompt][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_prompt] to replicate them. - For even more control over how the HF processor is called, you can override + Only processors that require multiple HF calls or manually construct the + complete output should override [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main] - directly. + directly. Such implementations must obtain their normalized inputs through + [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs] + and return through + [_finalize_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._finalize_hf_mm_data] + so passthrough fields and model-specific postprocessing are applied consistently. Since `pixel_values` is now a list with one tensor per image, we can override [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows: @@ -724,7 +735,7 @@ Examples: ### Custom HF processor -Some models don't define an HF processor class on HF Hub. In that case, you can define a custom HF processor that has the same call signature as HF processors and return it from [BaseProcessingInfo.get_hf_processor][vllm.multimodal.processing.BaseProcessingInfo.get_hf_processor]. It is then applied to the multi-modal data via [InputProcessingContext.call_hf_processor][vllm.multimodal.processing.InputProcessingContext.call_hf_processor] inside [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main]. +Some models don't define an HF processor class on HF Hub. In that case, you can define a custom HF processor that has the same call signature as HF processors and return it from [BaseProcessingInfo.get_hf_processor][vllm.multimodal.processing.BaseProcessingInfo.get_hf_processor]. The default [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor] applies it to the normalized multi-modal data via [InputProcessingContext.call_hf_processor][vllm.multimodal.processing.InputProcessingContext.call_hf_processor]. Override `_call_hf_processor` if the custom processor instead exposes a different callable interface. Examples: diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md index a2112a2922ed..202d8d04c166 100644 --- a/docs/design/mm_processing.md +++ b/docs/design/mm_processing.md @@ -8,9 +8,16 @@ In vLLM's rendering pipeline (see [BaseRenderer][vllm.renderers.base.BaseRendere Since Transformers 5.10, `ProcessorMixin` now allows multi-modal inputs to be passed by themselves. However, certain subclasses (such as `ChameleonProcessor`) and older out-of-tree implementations may still define their own `__call__` method that assumes the presence of text with corresponding placeholder tokens. This causes a problem as we don't have the original text anymore to pass to these processors. -To work around this, each model defines how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text], which its override of [_get_hf_processor_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_processor_text] returns so that [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main] passes it to the HF processor together with the multi-modal inputs to obtain the processed multi-modal data. +To work around this, each model defines how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text], which its override of [_get_hf_mm_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_text] returns. [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs] adds this text and `truncation=False` to the normalized HF inputs. The truncation setting is attached only to calls that use dummy text, so unrelated processor calls do not receive a text-specific option. -Similarly, since the multi-modal data extracted by vLLM may not match what a specific HF processor expects, [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main] allows each model to adapt the inputs via [_preprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._preprocess_hf_mm_data] (e.g. renaming keys like `audios` to `audio` or injecting extra keyword arguments such as `sampling_rate`) and the outputs via [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data], without having to reimplement the entire method. +HF processing is split into four overridable stages: + +1. [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs] extracts and normalizes processor data, kwargs, and passthrough data. It centrally renames vLLM's `audios` field to the standard HF `audio` keyword; processors that require another spelling can adapt it in an override. +2. [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor] invokes the HF processor. Models can override it when they use a different callable or need separate construction and invocation kwargs. +3. [_finalize_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._finalize_hf_mm_data] merges passthrough fields into the processed output. It lazily creates an empty output when there is no processor data. +4. [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data] performs model-specific output transformations after passthrough fields have been merged. + +This division lets most models customize one stage without reimplementing [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main]. Implementations that genuinely require multiple processor calls or manually construct their complete output may still override that method, but must use `_get_hf_mm_inputs` and `_finalize_hf_mm_data` to retain the common input and output behavior. ## Prompt Update Detection diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 4832d445d697..36f256927775 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -374,21 +374,9 @@ def _parse_audio_data( class AudioFlamingo3MultiModalProcessor( BaseMultiModalProcessor[AudioFlamingo3ProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - audios = mm_data.pop("audios", None) - if audios is not None: - mm_data["audio"] = audios - - return mm_data, hf_processor_mm_kwargs - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/bailing_moe_v3_vl.py b/vllm/model_executor/models/bailing_moe_v3_vl.py index c926c49ea8a1..7cd05685309d 100644 --- a/vllm/model_executor/models/bailing_moe_v3_vl.py +++ b/vllm/model_executor/models/bailing_moe_v3_vl.py @@ -119,7 +119,7 @@ def get_dummy_mm_data( class BailingMoeV3VLMultiModalProcessor( BaseMultiModalProcessor[BailingMoeV3VLProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _apply_prompt_updates( diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index fbb853beb0a9..a4748932d67e 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -222,7 +222,7 @@ def get_dummy_mm_data( class Cohere2VisionMultiModalProcessor( BaseMultiModalProcessor[Cohere2VisionProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 29205e4ed6be..5020cadc99ca 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -8,7 +8,7 @@ import torch import torch.nn.functional as F from torch import nn -from transformers import BatchFeature, PretrainedConfig +from transformers import PretrainedConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig @@ -50,7 +50,10 @@ PromptUpdate, TimingContext, ) -from vllm.multimodal.processing.processor import MultiModalProcessingInfo +from vllm.multimodal.processing.processor import ( + HFMultiModalInputs, + MultiModalProcessingResult, +) from vllm.renderers import TokenizeParams from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.processors.cohere_asr import ( @@ -1962,42 +1965,29 @@ def create_encoder_prompt( ) -> list[int]: return [0] - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - - mm_data = dict(mm_data) - mm_data["audio"] = mm_data.pop("audios") - - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) ) - return mm_data, hf_processor_mm_kwargs - - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if "labels" in processed_data: - processed_data["input_ids"] = processed_data.pop("labels") - - return processed_data - def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: # Dithering injects noise into the extracted features, so the # feature extractor is not a pure function of its input. Since the # processing cache assumes that processor outputs are invariant diff --git a/vllm/model_executor/models/colmodernvbert.py b/vllm/model_executor/models/colmodernvbert.py index 89e55d40e275..0c93e7ef52e3 100644 --- a/vllm/model_executor/models/colmodernvbert.py +++ b/vllm/model_executor/models/colmodernvbert.py @@ -157,15 +157,14 @@ class ColModernVBertMultiModalProcessor( def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) if not mm_data: - return BatchFeature(dict(passthrough_data)) + return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) @@ -193,8 +192,9 @@ def _apply_hf_processor_main( result.update(image_outputs) processed_data = result - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/colpali.py b/vllm/model_executor/models/colpali.py index aea809c15ecc..4add3b83a644 100644 --- a/vllm/model_executor/models/colpali.py +++ b/vllm/model_executor/models/colpali.py @@ -20,7 +20,7 @@ import torch import torch.nn as nn -from transformers import BatchFeature, PaliGemmaProcessor +from transformers import PaliGemmaProcessor from vllm.config import VllmConfig from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed @@ -28,6 +28,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing.processor import HFMultiModalInputs from .interfaces import SupportsLateInteraction from .interfaces_base import default_pooling_type @@ -59,31 +60,19 @@ def get_hf_processor(self, **kwargs: object) -> PaliGemmaProcessor: class ColPaliMultiModalProcessor(PaliGemmaMultiModalProcessor): """Multimodal processor for ColPali.""" - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) # The ColPali tokenizer_config.json ships with a small default # max_length (50) that truncates the 1024 image tokens inserted # by PaliGemmaProcessor, causing a token-count mismatch. # vLLM enforces its own max_model_len, so we disable HF # truncation to keep all image + text tokens intact. - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - mm_data, - dict(**hf_processor_mm_kwargs, truncation=False), - ) - processed_data.update(passthrough_data) - return processed_data + return hf_inputs._replace(hf_kwargs=dict(hf_inputs.hf_kwargs, truncation=False)) @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 6e438680956f..3ead25f173dc 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -46,10 +46,10 @@ PromptReplacement, PromptUpdate, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sampling_params import SamplingParams from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config -from vllm.tokenizers.hf import HfTokenizer from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config from vllm.transformers_utils.processors.deepseek_ocr import ( BASE_SIZE, @@ -287,34 +287,19 @@ def get_dummy_mm_data( class DeepseekOCRMultiModalProcessor( BaseMultiModalProcessor[DeepseekOCRProcessingInfo] ): - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - if mm_data: - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(prompt=prompt_text, **mm_data), - hf_processor_mm_kwargs, - ) - - else: - tokenizer = self.info.get_tokenizer() - assert isinstance(tokenizer, HfTokenizer) - processed_data = tokenizer( - prompt_text, add_special_tokens=True, return_tensors="pt" - ) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + if "text" in hf_inputs.hf_data: + hf_inputs.hf_data["prompt"] = hf_inputs.hf_data.pop("text") - processed_data.update(passthrough_data) - return processed_data + return hf_inputs def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index caaf90967183..a64b71b64e49 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -45,9 +45,9 @@ PromptReplacement, PromptUpdate, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config -from vllm.tokenizers.hf import HfTokenizer from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config from vllm.transformers_utils.processors.deepseek_ocr import ( BASE_SIZE, @@ -158,35 +158,19 @@ def get_dummy_mm_data( class DeepseekOCR2MultiModalProcessor( BaseMultiModalProcessor[DeepseekOCR2ProcessingInfo] ): - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - if mm_data: - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(prompt=prompt_text, **mm_data), - hf_processor_mm_kwargs, - ) - - else: - tokenizer = self.info.get_tokenizer() - assert isinstance(tokenizer, HfTokenizer) - processed_data = tokenizer( - prompt_text, add_special_tokens=True, return_tensors="pt" - ) - - processed_data.update(passthrough_data) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + if "text" in hf_inputs.hf_data: + hf_inputs.hf_data["prompt"] = hf_inputs.hf_data.pop("text") - return processed_data + return hf_inputs def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index 5fb52cff904a..6f798a87e647 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -35,7 +35,7 @@ from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, BaseProcessingInfo, - MultiModalProcessingInfo, + MultiModalProcessingResult, ProcessorInputs, PromptReplacement, PromptUpdate, @@ -233,7 +233,7 @@ def get_dummy_mm_data( class DeepseekVL2MultiModalProcessor( BaseMultiModalProcessor[DeepseekVL2ProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( @@ -307,7 +307,7 @@ def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: # The processor logic is different for len(images) <= 2 vs > 2 # Since the processing cache assumes that the processor output is # invariant of how many images are passed per prompt, we only diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 81907a517d43..e1b847c2940f 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -71,6 +71,7 @@ PromptUpdate, cached_encode, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -1112,34 +1113,32 @@ def _pixel_values_norm( pixel_values = pixel_values.to(hf_config.dtype) return pixel_values - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if "images" not in mm_data and "videos" not in mm_data: - return BatchFeature(dict(passthrough_data)) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + hf_data = hf_inputs.hf_data - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + if "images" not in hf_data and "videos" not in hf_data: + return hf_inputs - mm_data = dict(mm_data) - if "images" not in mm_data: - mm_data["images"] = [] - if "videos" not in mm_data: - mm_data["videos"] = [] + prompt_text = hf_data["text"] + hf_data["text"] = [prompt_text] + hf_data.setdefault("images", []) + hf_data.setdefault("videos", []) # Check if HF processor supports video metadata - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + hf_processor = self.info.get_hf_processor(**hf_inputs.hf_kwargs) supports_video_metadata = getattr( hf_processor, "supports_video_metadata", False ) - videos = mm_data["videos"] + videos = hf_data["videos"] assert isinstance(videos, Sequence) if videos and not supports_video_metadata: # Old HF processor, unwrap tuple to pure frames @@ -1147,49 +1146,43 @@ def _apply_hf_processor_main( "HF processor doesn't support video metadata. " "Timestamps will NOT be rendered. Please upgrade the model." ) - mm_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos] + hf_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos] - processor_output = self.info.ctx.call_hf_processor( - hf_processor, - dict( - text=[prompt_text], images=mm_data["images"], videos=mm_data["videos"] - ), - hf_processor_mm_kwargs, - ) + return hf_inputs + + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if "images" not in hf_data and "videos" not in hf_data: + return processed_data # Divide the processor_output into two modalities: image and video. - if processor_output is not None: - pixel_values = processor_output["images"] - if pixel_values is not None: - processor_output["images"] = self._pixel_values_norm( - pixel_values, hf_processor_mm_kwargs - ) - for key in list(processor_output.keys()): - if processor_output[key] is None: - del processor_output[key] - continue - if key == "grid_thw": - grid_thw = processor_output["grid_thw"] - pixel_values_all = processor_output["images"] - # Identify elements where the first - # dimension is greater than 1 and - # treat them as the video modality - mask = grid_thw[:, 0] > 1 - processor_output["video_grid_thw"] = grid_thw[mask] - processor_output["image_grid_thw"] = grid_thw[~mask] - image_patch_num = ( - processor_output["image_grid_thw"].prod(dim=1).sum() - ) - processor_output["pixel_values"] = pixel_values_all[ - :image_patch_num - ] - processor_output["pixel_values_videos"] = pixel_values_all[ - image_patch_num: - ] - del processor_output["images"] + pixel_values = processed_data["images"] + if pixel_values is not None: + processed_data["images"] = self._pixel_values_norm(pixel_values, hf_kwargs) + for key in list(processed_data.keys()): + if processed_data[key] is None: + del processed_data[key] + continue + if key == "grid_thw": + grid_thw = processed_data["grid_thw"] + pixel_values_all = processed_data["images"] + # Identify elements where the first + # dimension is greater than 1 and + # treat them as the video modality + mask = grid_thw[:, 0] > 1 + processed_data["video_grid_thw"] = grid_thw[mask] + processed_data["image_grid_thw"] = grid_thw[~mask] + image_patch_num = processed_data["image_grid_thw"].prod(dim=1).sum() + processed_data["pixel_values"] = pixel_values_all[:image_patch_num] + processed_data["pixel_values_videos"] = pixel_values_all[ + image_patch_num: + ] + del processed_data["images"] - processed_data = processor_output - processed_data.update(passthrough_data) return processed_data def _get_prompt_updates( diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index 847ec89e8f05..82a4a6b50677 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -38,6 +38,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.transformers_utils.processor import cached_processor_from_config from vllm.transformers_utils.processors.fireredasr2 import ( FireRedASR2FeatureExtractor, @@ -228,37 +229,24 @@ def get_dummy_mm_data( class FireRedASR2MultiModalProcessor( BaseMultiModalProcessor[FireRedASR2ProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - - mm_data = dict(mm_data) - mm_data["audio"] = mm_data.pop("audios") - - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) ) - return mm_data, hf_processor_mm_kwargs - - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if "labels" in processed_data: - processed_data["input_ids"] = processed_data.pop("labels") - - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index 9a7dc49d7195..7bf253ab40ee 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -49,7 +49,10 @@ PromptUpdate, TimingContext, ) -from vllm.multimodal.processing.processor import MultiModalProcessingInfo +from vllm.multimodal.processing.processor import ( + HFMultiModalInputs, + MultiModalProcessingResult, +) from vllm.transformers_utils.processor import cached_processor_from_config from vllm.transformers_utils.processors.funasr import FunASRFeatureExtractor from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -752,42 +755,29 @@ def get_dummy_mm_data( class FunASRMultiModalProcessor(BaseMultiModalProcessor[FunASRProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - - mm_data = dict(mm_data) - mm_data["audio"] = mm_data.pop("audios") - - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) ) - return mm_data, hf_processor_mm_kwargs - - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if "labels" in processed_data: - processed_data["input_ids"] = processed_data.pop("labels") - - return processed_data - def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: # Dithering injects noise into the extracted features, so the # feature extractor is not a pure function of its input. Since the # processing cache assumes that processor outputs are invariant diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 607eb4ea5d05..7a093cc2e74c 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -616,29 +616,29 @@ class FunAudioChatMultiModalProcessor( def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) if not mm_data: - return BatchFeature(dict(passthrough_data)) + return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) tokenizer = self.info.get_tokenizer() input_ids = torch.tensor([tokenizer.encode(prompt_text)]) - audios = mm_data.get("audios", []) + audios = mm_data.get("audio", []) if not audios: processed_data = BatchFeature({"input_ids": input_ids}) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) assert isinstance(audios, Sequence) - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) sr = int(feature_extractor.sampling_rate) min_samples = int(getattr(feature_extractor, "n_fft", 400) or 400) @@ -690,8 +690,9 @@ def _apply_hf_processor_main( } processed_data = BatchFeature({"input_ids": input_ids, **mm_inputs}) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index d6aaabe0232f..5bccd9aec0bb 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -263,7 +263,7 @@ def get_dummy_mm_data( class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 769c4c666a5c..dfde046fece6 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -265,20 +265,9 @@ def get_dummy_mm_data( class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - if "audios" in mm_data: - mm_data["audio"] = mm_data.pop("audios") - - return mm_data, hf_processor_mm_kwargs - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 752e2243df63..6ad98cb65086 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -595,22 +595,25 @@ def _get_dummy_videos( class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if not mm_data: - return BatchFeature(dict(passthrough_data)) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) - merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs) + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs) val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs) if val is not None and val not in _SUPPORTED_SOFT_TOKENS: @@ -619,14 +622,12 @@ def _apply_hf_processor_main( f"Valid values are {_SUPPORTED_SOFT_TOKENS}." ) - mm_data = dict(mm_data) - # ---- VIDEO HANDLING ---- # Gemma4 decomposes video into timestamped image frames. # Each frame is processed with max_soft_tokens=70 through the # same vision tower, matching transformers processing_gemma4.py. video_outputs: dict[str, Any] = {} - if videos := mm_data.pop("videos", []): + if videos := hf_data.pop("videos", []): assert isinstance(videos, list) processor = self.info.get_hf_processor() @@ -656,7 +657,7 @@ def _apply_hf_processor_main( timestamps = [idx / fps for idx in frame_indices] # Process frames as images with max_soft_tokens=70 - video_mm_kwargs = dict(hf_processor_mm_kwargs) + video_mm_kwargs = dict(hf_kwargs) # Override configured image options too: these inputs are video # frames, whose cache keys exclude images_kwargs. video_mm_kwargs["images_kwargs"] = {} @@ -726,18 +727,14 @@ def _apply_hf_processor_main( "video_timestamps": video_timestamps_per_video, } - # The processor accepts 'audio' not 'audios'. - if "audios" in mm_data: - mm_data["audio"] = mm_data.pop("audios") - # Warn if any audio waveform exceeds the model's max duration. - if "audio" in mm_data: + if "audio" in hf_data: processor = self.info.get_hf_processor() sr = processor.feature_extractor.sampling_rate max_tokens = processor.audio_seq_length ms_per_tok = processor.audio_ms_per_token max_duration_s = max_tokens * ms_per_tok / 1000.0 - audios = mm_data["audio"] + audios = hf_data["audio"] if not isinstance(audios, (list, tuple)): audios = [audios] for i, waveform in enumerate(audios): @@ -760,13 +757,13 @@ def _apply_hf_processor_main( # NOTE: This requires a corresponding type annotation on the # HF side (Gemma4ProcessorKwargs.images_kwargs) so that # _merge_kwargs routes max_soft_tokens into images_kwargs. - patched_mm_kwargs = dict(hf_processor_mm_kwargs) + patched_mm_kwargs = dict(hf_kwargs) if val is not None and is_top_level_max_soft_tokens: patched_mm_kwargs["max_soft_tokens"] = val processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**patched_mm_kwargs), - dict(text=prompt_text, **mm_data), + dict(text=prompt_text, **hf_data), patched_mm_kwargs, ) @@ -797,9 +794,9 @@ def _apply_hf_processor_main( processed_data["input_features_mask"] = unpadded_masks processed_data.update(video_outputs) - processed_data.update(passthrough_data) - - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 87baaa10693a..c0f77e770da6 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1582,52 +1582,59 @@ def _get_direct_path_inputs( return prepared_data, prepared_kwargs + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - mm_data = dict(mm_data) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + processor = self.info.get_hf_processor(**hf_kwargs) use_direct_path = ( not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA ) if use_direct_path: prepared_data, prepared_kwargs = self._get_direct_path_inputs( - mm_data, hf_processor_mm_kwargs + hf_data, hf_kwargs ) processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**prepared_kwargs), dict(text=prompt_text, **prepared_data), prepared_kwargs, ) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) - videos = mm_data.get("videos") - if isinstance(videos, list) and videos: + if ( + "videos" in hf_data + and isinstance(videos := hf_data["videos"], list) + and len(videos) > 0 + ): + # Videos are processed item-wise below; remove them from hf_data so + # the combined call at the end does not re-sample the raw frames. + del hf_data["videos"] video_grid_thw_lst = [] pixel_values_videos_lst = [] frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor) swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id - mm_data.pop("videos") for item in videos: video_array, metadata = item # don't update mm_kwargs inplace - video_mm_kwargs = dict(**hf_processor_mm_kwargs) + video_mm_kwargs = dict(**hf_kwargs) video_mm_kwargs["do_sample_frames"] = metadata.get( "do_sample_frames", True ) @@ -1667,9 +1674,9 @@ def _apply_hf_processor_main( swap_video_frame_tokens = False processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + self.info.get_hf_processor(**hf_kwargs), + dict(text=prompt_text, **hf_data), + hf_kwargs, ) if swap_video_frame_tokens: input_ids = processed_data["input_ids"] @@ -1677,9 +1684,9 @@ def _apply_hf_processor_main( processed_data["input_ids"] = input_ids processed_data.update(video_outputs) - processed_data.update(passthrough_data) - - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 710c6378a8cc..003b02c28ef9 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -44,6 +44,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.processor import cached_processor_from_config @@ -739,26 +740,24 @@ def _calculate_chunk_counts( chunk_counts.append(min(n_chunks, max_windows)) return chunk_counts - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - if "audios" in mm_data: - mm_data["audio"] = mm_data.pop("audios") - - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) ) - return mm_data, hf_processor_mm_kwargs - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index c8d1f7111bf0..529befd21d24 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -110,7 +110,7 @@ class GraniteSpeechAudioInputs(TensorSchema): """List of audio embedding sizes for each item in batch.""" -class GraniteSpeechMultiModalProcessingInfo(BaseProcessingInfo): +class GraniteSpeechProcessingInfo(BaseProcessingInfo): def get_data_parser(self): feature_extractor = self.get_hf_processor().audio_processor @@ -135,7 +135,7 @@ def get_max_audio_len(self): ### Input Processing & Multimodal utils class GraniteSpeechMultiModalProcessor( - BaseMultiModalProcessor[GraniteSpeechMultiModalProcessingInfo] + BaseMultiModalProcessor[GraniteSpeechProcessingInfo] ): def _get_mm_fields_config( self, @@ -181,23 +181,9 @@ def get_replacement(item_idx: int): ) ] - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - - if audios: - # GraniteSpeechFeatureExtractor accepts "audio" - mm_data["audio"] = audios - - return mm_data, hf_processor_mm_kwargs - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], @@ -216,7 +202,7 @@ def _postprocess_hf_mm_data( class GraniteSpeechDummyInputsBuilder( - BaseDummyInputsBuilder[GraniteSpeechMultiModalProcessingInfo] + BaseDummyInputsBuilder[GraniteSpeechProcessingInfo] ): def get_dummy_mm_data( self, @@ -565,7 +551,7 @@ def forward(self, hidden_states: torch.Tensor): @MULTIMODAL_REGISTRY.register_processor( GraniteSpeechMultiModalProcessor, - info=GraniteSpeechMultiModalProcessingInfo, + info=GraniteSpeechProcessingInfo, dummy_inputs=GraniteSpeechDummyInputsBuilder, ) class GraniteSpeechForConditionalGeneration( diff --git a/vllm/model_executor/models/granite_speech_plus.py b/vllm/model_executor/models/granite_speech_plus.py index ba95cfed131e..7e695b7ed53c 100644 --- a/vllm/model_executor/models/granite_speech_plus.py +++ b/vllm/model_executor/models/granite_speech_plus.py @@ -12,8 +12,8 @@ GraniteSpeechCTCEncoder, GraniteSpeechDummyInputsBuilder, GraniteSpeechForConditionalGeneration, - GraniteSpeechMultiModalProcessingInfo, GraniteSpeechMultiModalProcessor, + GraniteSpeechProcessingInfo, ) ISO639_1_SUPPORTED_LANGS = { @@ -60,7 +60,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @MULTIMODAL_REGISTRY.register_processor( GraniteSpeechMultiModalProcessor, - info=GraniteSpeechMultiModalProcessingInfo, + info=GraniteSpeechProcessingInfo, dummy_inputs=GraniteSpeechDummyInputsBuilder, ) class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGeneration): diff --git a/vllm/model_executor/models/h2ovl.py b/vllm/model_executor/models/h2ovl.py index 5be2d7b46d99..596f24a3eba3 100644 --- a/vllm/model_executor/models/h2ovl.py +++ b/vllm/model_executor/models/h2ovl.py @@ -21,7 +21,7 @@ MultiModalDataItems, ) from vllm.multimodal.processing.processor import ( - MultiModalProcessingInfo, + MultiModalProcessingResult, ProcessorInputs, PromptReplacement, TimingContext, @@ -139,7 +139,7 @@ def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: # The processor logic is different for len(images) <= 1 vs > 1 # Since the processing cache assumes that the processor output is # invariant of how many images are passed per prompt, we only diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py index 10c066cb7176..07106c947134 100644 --- a/vllm/model_executor/models/hyperclovax_vision_v2.py +++ b/vllm/model_executor/models/hyperclovax_vision_v2.py @@ -261,44 +261,8 @@ class HCXVisionV2MultiModalProcessor( ): """Multimodal processor for HyperCLOVAX V2 (32B Think model).""" - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - images = mm_data.get("images") - videos = mm_data.get("videos") - - # Get the HF processor - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - - # Build data dict for HF processor (images/videos only) - # NOTE: We pass the prompt as-is without token normalization. - # Token expansion is handled by vLLM via _get_prompt_updates. - data: dict[str, object] = dict( - text=prompt_text, - images=images, - videos=videos, - ) - - processed_data = self.info.ctx.call_hf_processor( - hf_processor=hf_processor, - data=data, - kwargs=hf_processor_mm_kwargs, - ) - processed_data.update(passthrough_data) - - return processed_data + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 014c4ae6b54f..dd03e17d1ec0 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -52,6 +52,7 @@ PromptUpdateDetails, cached_encode, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -240,31 +241,32 @@ def get_dummy_mm_data( class Idefics3MultiModalProcessor(BaseMultiModalProcessor[Idefics3ProcessingInfo]): - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) image_processor = self.info.get_hf_processor().image_processor if getattr(image_processor, "backend", "pil") == "pil": - hf_processor_mm_kwargs = { - "input_data_format": "channels_last", - **hf_processor_mm_kwargs, - } + hf_inputs = hf_inputs._replace( + hf_kwargs={ + "input_data_format": "channels_last", + **hf_inputs.hf_kwargs, + } + ) - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - mm_data, - hf_processor_mm_kwargs, - ) + return hf_inputs + + def _postprocess_hf_mm_data( + self, + mm_data: Mapping[str, object], + hf_processor_mm_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if not mm_data: + return processed_data images = mm_data.get("images", []) mm_items = self.info.parse_mm_data({"image": images}, validate=False) @@ -289,8 +291,6 @@ def _apply_hf_processor_main( processed_data["pixel_values"].squeeze_(0) processed_data["pixel_attention_mask"].squeeze_(0) - processed_data.update(passthrough_data) - return processed_data def _get_mm_fields_config( diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py index 057a6ce2479c..4a4264660d54 100644 --- a/vllm/model_executor/models/interns1.py +++ b/vllm/model_executor/models/interns1.py @@ -335,28 +335,30 @@ def get_dummy_mm_data( class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo]): """Basic image-only MultiModalProcessor for InternS1-style models.""" + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if not mm_data: - return BatchFeature(dict(passthrough_data)) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) - mm_data = dict(mm_data) - videos = mm_data.pop("videos", []) - images = mm_data.pop("images", []) + videos = hf_data.pop("videos", []) + images = hf_data.pop("images", []) assert isinstance(videos, list) assert isinstance(images, list) - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + hf_processor = self.info.get_hf_processor(**hf_kwargs) tokenizer = hf_processor.tokenizer vocab = tokenizer.get_vocab() @@ -374,9 +376,9 @@ def _apply_hf_processor_main( image_pixel_values = [] for image in images: processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), + self.info.get_hf_processor(**hf_kwargs), dict(text=hf_processor.image_token, **{"images": image}), - hf_processor_mm_kwargs, + hf_kwargs, ) image_pixel_values.append(processed_data.pop("pixel_values")) @@ -398,9 +400,9 @@ def _apply_hf_processor_main( video_pixel_values = [] for video in videos: processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), + self.info.get_hf_processor(**hf_kwargs), dict(text=hf_processor.video_token, **{"videos": video}), - hf_processor_mm_kwargs, + hf_kwargs, ) video_pixel_values.append(processed_data.pop("pixel_values")) @@ -430,9 +432,9 @@ def _apply_hf_processor_main( processed_data.update(text_outputs) processed_data.update(image_outputs) processed_data.update(video_outputs) - processed_data.update(passthrough_data) - - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/jina_vl.py b/vllm/model_executor/models/jina_vl.py index 2895ec13b925..971284561926 100644 --- a/vllm/model_executor/models/jina_vl.py +++ b/vllm/model_executor/models/jina_vl.py @@ -13,7 +13,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalKwargsItems from vllm.multimodal.processing.processor import ( - MultiModalProcessingInfo, + MultiModalProcessingResult, ProcessorInputs, TimingContext, ) @@ -63,19 +63,19 @@ def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: - mm_info = super()._cached_apply_hf_processor(inputs, timing_ctx) + ) -> MultiModalProcessingResult: + mm_res = super()._cached_apply_hf_processor(inputs, timing_ctx) # Score inputs are query-first, while the prompt template is document-first. mm_kwargs = MultiModalKwargsItems( { modality: list(reversed(items)) - for modality, items in mm_info.kwargs.items() + for modality, items in mm_res.kwargs.items() } ) mm_hashes = { modality: list(reversed(hashes)) - for modality, hashes in mm_info.hashes.items() + for modality, hashes in mm_res.hashes.items() } mm_prompt_updates = { modality: [ @@ -85,10 +85,11 @@ def _cached_apply_hf_processor( ] for item_idx, updates in enumerate(reversed(item_updates)) ] - for modality, item_updates in mm_info.prompt_updates.items() + for modality, item_updates in mm_res.prompt_updates.items() } - return MultiModalProcessingInfo( + return MultiModalProcessingResult( + prompt_ids=mm_res.prompt_ids, kwargs=mm_kwargs, hashes=mm_hashes, prompt_updates=mm_prompt_updates, diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index a9a5dfda24aa..e3c27e479b00 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -459,24 +459,27 @@ class KananaVMultiModalProcessor(BaseMultiModalProcessor[KananaVProcessingInfo]) def media_token_id(self) -> int: return self.info.get_hf_config().text_config.eos_token_id + 1 + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Run the underlying HF processor on text and image data.""" - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if not mm_data or not mm_data.get("images", []): - return BatchFeature(dict(passthrough_data)) + if not hf_data or not hf_data.get("images", []): + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) # Images - image_inputs = mm_data.get("images", []) + image_inputs = hf_data.get("images", []) pixel_sizes = [] if not isinstance(image_inputs[0], Image.Image): image_inputs = [Image.fromarray(image) for image in image_inputs] @@ -530,8 +533,9 @@ def _apply_hf_processor_main( pixel_sizes=torch.tensor(pixel_sizes), ) processed_data = BatchFeature(combined_outputs, tensor_type="pt") - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 6bbffae0290e..b11adfa499f8 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -1160,7 +1160,7 @@ class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]): class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _get_prompt_updates( diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index eec06304551a..b52d50296abc 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -373,7 +373,7 @@ def get_supported_mm_limits( class KeyeVL1_5MultiModalProcessor(BaseMultiModalProcessor[KeyeVL1_5ProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _get_prompt_updates( diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index 71cf389f98f7..219c7f0e7b92 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -47,6 +47,7 @@ ) from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, + HFMultiModalInputs, ProcessorInputs, ) from vllm.sequence import IntermediateTensors @@ -227,27 +228,17 @@ def _parse_audio_data( class KimiAudioMultiModalProcessor(BaseMultiModalProcessor[KimiAudioProcessingInfo]): """vLLM multi-modal processor wrapper for Kimi-Audio.""" - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - """Call the HuggingFace processor.""" - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - # Convert mm_data format: {'audios': [...]} -> {'audio': ...} - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) # Convert audio format: [(array, sr), ...] -> [array, ...] # KimiAudioProcessor expects raw numpy arrays - if audios: + mm_data = hf_inputs.hf_data + if audios := mm_data.pop("audio", []): audio_arrays = [] for aud in audios: if isinstance(aud, (tuple, list)) and len(aud) == 2: @@ -259,14 +250,7 @@ def _apply_hf_processor_main( audio_arrays.append(aud) mm_data["audio"] = audio_arrays - # Use the context's call_hf_processor for proper handling - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - mm_data, - hf_processor_mm_kwargs, - ) - processed_data.update(passthrough_data) - return processed_data + return hf_inputs def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index 9aba4d194524..344dfdaaf031 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/vllm/model_executor/models/kimi_vl.py @@ -236,7 +236,7 @@ def get_dummy_mm_data( class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index 9b24095ab6b0..1cd51b93a2ed 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -406,7 +406,7 @@ def get_dummy_mm_data( class Lfm2VLMultiModalProcessor(BaseMultiModalProcessor[Lfm2VLProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index 0ce0d7f3d2b3..d8a124024271 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/vllm/model_executor/models/lightonocr.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Mapping, Sequence -from typing import TypeVar import torch import torch.nn as nn @@ -35,8 +34,6 @@ PromptUpdateDetails, ) -_I = TypeVar("_I", bound=Mistral3ProcessingInfo) - class LightOnOCRProcessingInfo(Mistral3ProcessingInfo): def get_vision_encoder_info( @@ -49,8 +46,8 @@ def get_vision_encoder_info( return PixtralHFEncoderInfo(self.get_hf_config()) -class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: +class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( @@ -62,27 +59,6 @@ def _postprocess_hf_mm_data( if not mm_data: return processed_data - input_ids = processed_data.get("input_ids") - if input_ids is not None: - processor = self.info.get_hf_processor() - tokenizer = self.info.get_tokenizer() - vocab = tokenizer.get_vocab() - - break_id = vocab.get(processor.image_break_token) - end_id = vocab.get(processor.image_end_token) - - # create mask to remove break/end tokens - keep_mask = ~torch.isin( - input_ids, - torch.tensor([break_id, end_id]), - ) - - processed_data["input_ids"] = input_ids[keep_mask].unsqueeze(0) - if "attention_mask" in processed_data: - processed_data["attention_mask"] = processed_data["attention_mask"][ - keep_mask - ].unsqueeze(0) - # un-pad pixel_values per-image so caches remain independent. pixel_values = processed_data.get("pixel_values") if pixel_values is not None: diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index 1c4ebff7a04c..c3a437eb368b 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -320,7 +320,7 @@ def get_hf_processor(self, **kwargs: object): class PixtralHFMultiModalProcessor(BaseMultiModalProcessor[PixtralHFProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py index b7b1a8533d6f..02400a763b22 100644 --- a/vllm/model_executor/models/llava_next.py +++ b/vllm/model_executor/models/llava_next.py @@ -209,7 +209,7 @@ def _get_mm_fields_config( class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py index 55f63c49baf9..66503d774137 100644 --- a/vllm/model_executor/models/llava_onevision.py +++ b/vllm/model_executor/models/llava_onevision.py @@ -324,33 +324,36 @@ def _get_mm_fields_config( pixel_values_videos=MultiModalFieldConfig.batched("video"), ) + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if not mm_data: - return BatchFeature(dict(passthrough_data)) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) - mm_data = dict(mm_data) - videos = mm_data.pop("videos", []) + videos = hf_data.pop("videos", []) assert isinstance(videos, list) if not videos: processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + self.info.get_hf_processor(**hf_kwargs), + dict(text=prompt_text, **hf_data), + hf_kwargs, + ) + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data ) - processed_data.update(passthrough_data) - return processed_data # LLaVA-OneVision processor doesn't support multiple videos # with different sizes when converting back to tensors @@ -360,13 +363,13 @@ def _apply_hf_processor_main( image_token = processor.image_token video_token = processor.video_token - images = mm_data.pop("images", []) + images = hf_data.pop("images", []) assert isinstance(images, list) if images: processor_outputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), + self.info.get_hf_processor(**hf_kwargs), dict(text=image_token * len(images), **{"images": images}), - hf_processor_mm_kwargs, + hf_kwargs, ) image_outputs = { k: v @@ -379,9 +382,9 @@ def _apply_hf_processor_main( pixel_values_videos = [] for video in videos: item_outputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), + self.info.get_hf_processor(**hf_kwargs), dict(text=video_token, **{"videos": video}), - hf_processor_mm_kwargs, + hf_kwargs, ) pixel_values_videos.append(item_outputs["pixel_values_videos"][0]) @@ -394,8 +397,9 @@ def _apply_hf_processor_main( **video_outputs, ) processed_data = BatchFeature(combined_outputs) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 71c3bc79dc98..6988982c2316 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1556,23 +1556,26 @@ def _parse_image_data( class LlavaOnevision2MultiModalProcessor( BaseMultiModalProcessor[LlavaOnevision2ProcessingInfo] ): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if not mm_data: - return BatchFeature(dict(passthrough_data)) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs) + hf_processor = self.info.get_hf_processor(**hf_kwargs) + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs) merged_kwargs.setdefault("return_tensors", "pt") call_kwargs = { k: v @@ -1589,17 +1592,17 @@ def _apply_hf_processor_main( "codec_config", } } - mm_data = dict(mm_data) + # Explicit None + length checks: ``mm_data[...]`` may be a list, numpy # array, or tensor, and ``and `` would raise on the ambiguous # truth value of a multi-element array. - _videos = mm_data.get("videos") + _videos = hf_data.get("videos") if _videos is not None and not isinstance(_videos, Sized): raise TypeError(f"object of type '{type(_videos).__name__}' has no len()") videos_present = _videos is not None and len(_videos) > 0 codec_video_paths = ( - _extract_codec_video_paths(mm_data["videos"]) if videos_present else None + _extract_codec_video_paths(hf_data["videos"]) if videos_present else None ) is_codec_marker = codec_video_paths is not None # Fallback: caller passed video_backend=codec via mm_processor_kwargs @@ -1613,7 +1616,7 @@ def _apply_hf_processor_main( and call_kwargs.get("video_backend") == "codec" ) if is_codec_kwarg: - raw = mm_data["videos"] + raw = hf_data["videos"] if isinstance(raw, str): codec_video_paths = [raw] elif isinstance(raw, (list, tuple)) and all( @@ -1638,7 +1641,7 @@ def _apply_hf_processor_main( # Codec backend: HF processor consumes the path string directly and # performs decode + canvas-packing internally. The dummy ndarray # we attached during prepare_codec_video_input is discarded here. - mm_data["videos"] = ( + hf_data["videos"] = ( codec_video_paths if len(codec_video_paths) > 1 else codec_video_paths[0] @@ -1648,11 +1651,9 @@ def _apply_hf_processor_main( # ``video_backend="codec"`` via mm_kwargs so the wrapped processor # dispatches to its codec branch. output = self.info.ctx.call_hf_processor( - self.info.get_hf_processor( - **{**hf_processor_mm_kwargs, "video_backend": "codec"} - ), - dict(text=prompt_text, **mm_data), - {**hf_processor_mm_kwargs, "video_backend": "codec"}, + self.info.get_hf_processor(**{**hf_kwargs, "video_backend": "codec"}), + dict(text=prompt_text, **hf_data), + {**hf_kwargs, "video_backend": "codec"}, ) data = dict(output) processed_data = BatchFeature( @@ -1660,8 +1661,9 @@ def _apply_hf_processor_main( data, codec_video_paths, hf_processor ) ) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) # ---- Frame backend (registered LlavaOnevision2VideoBackend) ------ # Every non-codec video reaches here as a ``(frames_ndarray, metadata)`` @@ -1676,7 +1678,7 @@ def _apply_hf_processor_main( # provided by the backend's ``compute_frames_index_to_sample``; SSRF / # local-file gating is enforced by the connector before decoding. if videos_present: - timestamp_decimals = hf_processor_mm_kwargs.get( + timestamp_decimals = hf_kwargs.get( "timestamp_decimals", _DEFAULT_TIMESTAMP_DECIMALS ) if not isinstance( @@ -1690,7 +1692,7 @@ def _apply_hf_processor_main( per_video_frames: list[list[Image.Image]] = [] per_video_timestamps: list[list[float]] = [] - videos = mm_data["videos"] + videos = hf_data["videos"] if not isinstance(videos, Iterable): raise TypeError(f"Unsupported videos input: {type(videos)}") for item in videos: @@ -1714,7 +1716,7 @@ def _apply_hf_processor_main( # order) -- otherwise mixed image+video requests bind frames to the # wrong placeholder. ``row_is_video`` labels each grid row so outputs # can be split back into per-modality keys below. - merged_mm_data = dict(mm_data) + merged_mm_data = dict(hf_data) existing_images = merged_mm_data.pop("images", None) caller_images: list[Image.Image] = [] if existing_images: @@ -1760,9 +1762,9 @@ def _apply_hf_processor_main( # **kwargs to the image processor, so passing the full merged kwarg # set here is a no-op beyond return_tensors/padding. output = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), + self.info.get_hf_processor(**hf_kwargs), dict(text=new_prompt, **merged_mm_data), - hf_processor_mm_kwargs, + hf_kwargs, ) data = dict(output) @@ -1820,20 +1822,22 @@ def _gather_rows(tensor, rows): ) processed_data = BatchFeature(data) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) # ---- Image-only / text-only call -------------------------------- # No videos present: delegate to ``ctx.call_hf_processor``, which # runs the wrapped processor over the (possibly empty) image set and # applies float-tensor dtype postprocessing automatically. processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + self.info.get_hf_processor(**hf_kwargs), + dict(text=prompt_text, **hf_data), + hf_kwargs, + ) + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data ) - processed_data.update(passthrough_data) - return processed_data def _rename_codec_outputs_to_video( self, diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index 300f274a187e..7de6ab0e7e2d 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -61,6 +61,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.midashenglm import DashengConfig from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -582,35 +583,33 @@ def get_dummy_mm_data( class MiDashengLMMultiModalProcessor( BaseMultiModalProcessor[MiDashengLMProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - assert isinstance(audios, list) - - min_audio_len = self.info.get_min_audio_len() - processed_audios = [ - np.pad( - audio, - (0, min_audio_len - audio.shape[-1]), - mode="constant", - constant_values=0, - ) - if isinstance(audio, np.ndarray) and audio.shape[-1] < min_audio_len - else audio - for audio in audios - ] - - if processed_audios: - mm_data["audio"] = processed_audios + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + if audios := hf_inputs.hf_data.get("audio"): + assert isinstance(audios, list) + + min_audio_len = self.info.get_min_audio_len() + hf_inputs.hf_data["audio"] = [ + np.pad( + audio, + (0, min_audio_len - audio.shape[-1]), + mode="constant", + constant_values=0, + ) + if isinstance(audio, np.ndarray) and audio.shape[-1] < min_audio_len + else audio + for audio in audios + ] - return mm_data, hf_processor_mm_kwargs + return hf_inputs def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 0edd0736e178..d74edd6d4ad0 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -46,6 +46,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.transformers_utils.configs.mimo_v2_omni import Mimo_VLVisionConfig from vllm.transformers_utils.processors.mimo_v2_omni import ( MiMoOmniProcessor, @@ -930,27 +931,14 @@ def _get_mm_fields_config( fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio") return fields - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo. - Also remap 'audios' → 'audio' since MiMoOmniProcessor.__call__ uses - the singular form. - """ - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - # Remap audios → audio (MiMoOmniProcessor uses singular param name) - if "audios" in mm_data: - mm_data = {**mm_data, "audio": mm_data["audios"]} - mm_data = {k: v for k, v in mm_data.items() if k != "audios"} + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo.""" + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + mm_data = hf_inputs.hf_data # Handle video_audio items: convert video part to (TCHW, timestamps) tuple if "video_audio" in mm_data: @@ -992,7 +980,7 @@ def _apply_hf_processor_main( audio=va_item.audio, ) ) - mm_data = {**mm_data, "video_audio": va_converted} + mm_data["video_audio"] = va_converted if "videos" in mm_data: converted: list[tuple[torch.Tensor, torch.Tensor]] = [] @@ -1026,15 +1014,9 @@ def _apply_hf_processor_main( timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS converted.append((frames, timestamps)) - mm_data = {**mm_data, "videos": converted} + mm_data["videos"] = converted - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - mm_data, - hf_processor_mm_kwargs, - ) - processed_data.update(passthrough_data) - return processed_data + return hf_inputs def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index 4251514326c6..2ec85679212c 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -445,7 +445,7 @@ def process_audios( mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: - if (audios := mm_data.get("audios")) is None: + if (audios := mm_data.get("audio")) is None: return {} mm_items = self.info.parse_mm_data({"audio": audios}, validate=False) diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index cf818bd77e10..14720c5c8630 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -1059,19 +1059,18 @@ def _call_hf_processor_on_prompts( def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) tokenizer = self.info.get_tokenizer() input_ids = torch.tensor([tokenizer.encode(prompt_text)]) - mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs) + mm_inputs = self.process_mm_inputs(mm_data, hf_kwargs) processed_data = BatchFeature( { @@ -1079,8 +1078,9 @@ def _apply_hf_processor_main( **mm_inputs, } ) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 0258fd7bc3a9..f9d7dafe70b3 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -248,7 +248,7 @@ def get_image_size_with_most_features(self) -> ImageSize: return ImageSize(width=width, height=height) -class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder[Mistral3ProcessingInfo]): +class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) @@ -280,7 +280,7 @@ def get_dummy_mm_data( class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 027a8db3b5ce..5c5c2728713c 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -589,7 +589,7 @@ def get_image_size_with_most_features(self) -> ImageSize: class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index 8547b689a03c..9aeef3715c8f 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -1122,6 +1122,9 @@ def get_dummy_mm_data( class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _postprocess_prompt(self, prompt: list[int]) -> list[int]: processor = self.info.get_hf_processor() @@ -1140,28 +1143,28 @@ def _postprocess_prompt(self, prompt: list[int]) -> list[int]: dict(tokens=tokens), )["input_ids"].tolist() - def _apply_hf_processor_main( + def _call_hf_processor( self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - mm_counts = mm_items.get_all_counts() - - valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0}) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not processor_data: - return BatchFeature(dict(passthrough_data)) - - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - prompt_text = self.dummy_inputs.get_dummy_text(mm_counts) - - processed_data = self.info.ctx.call_hf_processor( + hf_processor = self.info.get_hf_processor(**hf_kwargs) + return self.info.ctx.call_hf_processor( hf_processor.process, - dict(text=prompt_text, **processor_data), - hf_processor_mm_kwargs, + hf_data, + hf_kwargs, ) + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if not hf_data: + return processed_data + + hf_processor = self.info.get_hf_processor(**hf_kwargs) tokenizer = hf_processor.tokenizer image_patch_id = tokenizer.vocab[IMAGE_PATCH_TOKEN] @@ -1169,7 +1172,7 @@ def _apply_hf_processor_main( processed_data.pop("input_ids") - if (images := processor_data.get("images")) is not None: + if (images := hf_data.get("images")) is not None: mm_items = self.info.parse_mm_data({"image": images}, validate=False) parsed_images = mm_items.get_items("image", ImageProcessorItems) image_sizes = [ @@ -1193,8 +1196,6 @@ def _apply_hf_processor_main( processed_data["num_crops"] = num_crops processed_data["img_patch_id"] = image_patch_id - processed_data.update(passthrough_data) - return processed_data def _get_mm_fields_config( diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 61a9d2b15385..4c0ce0868e2d 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -1932,6 +1932,9 @@ def _get_dummy_videos( class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _postprocess_prompt(self, prompt: list[int]) -> list[int]: processor = self.info.get_hf_processor() tokenizer = processor.tokenizer @@ -1946,19 +1949,20 @@ def _postprocess_prompt(self, prompt: list[int]) -> list[int]: def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - mm_counts = mm_items.get_all_counts() - - valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0}) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs + ) - prompt_text = self.dummy_inputs.get_dummy_text(mm_counts) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - mm_data = dict(processor_data) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) hf_config = self.info.get_hf_config() - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + hf_processor = self.info.get_hf_processor(**hf_kwargs) def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: res = hf_processor(text=text, images=images, videos=videos, **kwargs) @@ -1973,7 +1977,7 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: tokenizer = hf_processor.tokenizer image_processor = hf_processor.image_processor - if videos := mm_data.pop("videos", []): + if videos := hf_data.pop("videos", []): assert isinstance(videos, Sequence) bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id @@ -1993,16 +1997,16 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: # NOTE: metadata.frames_indices indicates # the sampled frames indices of pre-sampled videos, which is # used to calculate the timestamps. Make sure that - # do_sample_frames in hf_processor_mm_kwargs is false for + # do_sample_frames in hf_kwargs is false for # presampled videos. - # NOTE: a copy of hf_processor_mm_kwargs is created to update + # NOTE: a copy of hf_kwargs is created to update # do_sample_frames, otherwise mm_hash for the object will be # incorrect. - video_mm_kwargs = dict(**hf_processor_mm_kwargs) + video_mm_kwargs = dict(**hf_kwargs) if "do_sample_frames" not in video_mm_kwargs: # molmo_utils already has "do_sample_frames" in - # hf_processor_mm_kwargs, don't overwrite it. + # hf_kwargs, don't overwrite it. video_mm_kwargs["do_sample_frames"] = metadata.get( "do_sample_frames", False ) @@ -2067,11 +2071,11 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: processed_data = self.info.ctx.call_hf_processor( patched_call, - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + dict(text=prompt_text, **hf_data), + hf_kwargs, ) - if (images := mm_data.get("images")) is not None: + if (images := hf_data.get("images")) is not None: mm_items = self.info.parse_mm_data({"image": images}, validate=False) parsed_images = mm_items.get_items("image", ImageProcessorItems) image_sizes = [ @@ -2113,10 +2117,11 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: ) processed_data.update(all_video_outputs) - processed_data.update(passthrough_data) processed_data.pop("input_ids") - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index 64a31375c862..6092812585b3 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -1356,32 +1356,17 @@ def get_dummy_mm_data( class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingInfo]): - def _apply_hf_processor_main( + def _call_hf_processor( self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - if audios: - mm_data["audio"] = audios - hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs) - processor_kwargs = _filter_moss_audio_processor_config(hf_processor_mm_kwargs) - processed_data = self.info.ctx.call_hf_processor( + processor_kwargs = _filter_moss_audio_processor_config(dict(hf_kwargs)) + return self.info.ctx.call_hf_processor( self.info.get_hf_processor(**processor_kwargs), - mm_data, + hf_data, {}, ) - processed_data.update(passthrough_data) - return processed_data def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 659b28d36b0f..32ac69c9afb8 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -66,7 +66,7 @@ PromptUpdateDetails, cached_encode, ) -from vllm.multimodal.processing.processor import ProcessorInputs +from vllm.multimodal.processing.processor import HFMultiModalInputs, ProcessorInputs from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -465,17 +465,19 @@ def get_dummy_processor_inputs( class MossTranscribeDiarizeMultiModalProcessor( BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - audios = _get_audios_from_mm_data(mm_data) + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + if hf_inputs.hf_data: + hf_inputs.hf_data["audio"] = _get_audios_from_mm_data(hf_inputs.hf_data) - return dict(audio=audios), hf_processor_mm_kwargs + return hf_inputs def _postprocess_hf_mm_data( self, diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index 192ac2de9812..38cc475be24f 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -275,16 +275,18 @@ class MuseGlimmerMultiModalProcessor( def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + + if not mm_data: + return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + processor = self.info.get_hf_processor(**hf_kwargs) tokenizer = processor.tokenizer config = self.info.get_hf_config() images = mm_data.get("images", ()) @@ -371,8 +373,9 @@ def _apply_hf_processor_main( video_feature_sizes=torch.tensor(video_sizes), ) processed_data = BatchFeature(data=data, tensor_type=None) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 2620923a490b..eb5e9f06ca4d 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -70,6 +70,7 @@ from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, BaseProcessingInfo, + HFMultiModalInputs, PromptReplacement, PromptUpdate, cached_encode, @@ -196,6 +197,9 @@ class NanoNemotronVLVideoEmbeddingInputs(TensorSchema): class NanoNemotronVLProcessingInfo(BaseProcessingInfo): def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor: + # Not accepted by NanoNemotronVLProcessor.__init__ + kwargs.pop("truncation", None) + return self.ctx.init_processor( NanoNemotronVLProcessor, config=self.get_hf_config(), @@ -358,9 +362,23 @@ def get_mm_max_tokens_per_item( class NanoNemotronVLMultiModalProcessor( BaseMultiModalProcessor[NanoNemotronVLProcessingInfo] ): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) + def _get_hf_mm_inputs( + self, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + # NanoNemotronVLProcessor expects "audios" instead of "audio" + hf_data = hf_inputs.hf_data + if "audio" in hf_data: + hf_data["audios"] = hf_data.pop("audio") + + return hf_inputs + def _get_image_fields_config(self, hf_inputs: BatchFeature): if self.info.is_dynamic_tiler: pixel_values_flat = MultiModalFieldConfig.batched("image") @@ -730,15 +748,12 @@ def apply( # Bypass the cached path: the HF processor must receive the # prompt (with injected ) and the audio data # together so it can perform audio-token replacement natively. - mm_info = self._apply_hf_processor(inputs, timing_ctx) - prompt_ids = self._postprocess_prompt(inputs.prompt) + mm_res = self._apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("apply_prompt_updates"): prompt_ids, mm_placeholders = self._maybe_apply_prompt_updates( mm_items=mm_items, - prompt_ids=prompt_ids, - mm_kwargs=mm_info.kwargs, - mm_prompt_updates=mm_info.prompt_updates, + mm_res=mm_res, ) mm_placeholder_ranges = { @@ -749,8 +764,8 @@ def apply( return MultiModalInput( type="multimodal", prompt_token_ids=prompt_ids, - mm_kwargs=mm_info.kwargs, - mm_hashes=mm_info.hashes, + mm_kwargs=mm_res.kwargs, + mm_hashes=mm_res.hashes, mm_placeholders=mm_placeholder_ranges, ) diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index f5065659b434..5563b9a9cea2 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -326,6 +326,9 @@ def get_dummy_mm_data( class OvisMultiModalProcessor(BaseMultiModalProcessor[OvisProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def image_indicators_to_visual_tokens( self, image_indicators: list[int], @@ -340,9 +343,6 @@ def image_indicators_to_visual_tokens( # -300 is image_atom token, filter them out return [vte_vocab_size + x + 300 for x in image_indicators if x < -300] - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py index 607dcff894ce..0b8a2c5f84bb 100644 --- a/vllm/model_executor/models/ovis2_5.py +++ b/vllm/model_executor/models/ovis2_5.py @@ -327,7 +327,7 @@ def visual_indicators_to_visual_tokens( if x >= INDICATOR_IDS[0] ] - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index d2587b020355..976d82032f20 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -243,36 +243,37 @@ def get_dummy_mm_data( class PaddleOCRVLMultiModalProcessor( BaseMultiModalProcessor[PaddleOCRVLProcessingInfo] ): - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _call_hf_processor( self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - final_mm_kwargs = dict(hf_processor_mm_kwargs or {}) + final_mm_kwargs = dict(hf_kwargs or {}) final_mm_kwargs.setdefault("images_kwargs", {}) # vLLM use PIL.Image, always set channel_last final_mm_kwargs["input_data_format"] = ChannelDimension.LAST - processed_data = self.info.ctx.call_hf_processor( + return self.info.ctx.call_hf_processor( self.info.get_hf_processor(**final_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + hf_data, + hf_kwargs, ) + + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if not hf_data: + return processed_data + num_patches_per_image = processed_data["image_grid_thw"].prod(-1) processed_data["pixel_values"] = processed_data["pixel_values"].split( num_patches_per_image.tolist() ) - processed_data.update(passthrough_data) - return processed_data def _get_mm_fields_config( diff --git a/vllm/model_executor/models/phi3v.py b/vllm/model_executor/models/phi3v.py index 17477e8de19d..a75b4afd2217 100644 --- a/vllm/model_executor/models/phi3v.py +++ b/vllm/model_executor/models/phi3v.py @@ -390,28 +390,9 @@ def get_dummy_mm_data( class Phi3VMultiModalProcessor(BaseMultiModalProcessor[Phi3VProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if not mm_data: - return processed_data - - input_ids = processed_data["input_ids"] - assert isinstance(input_ids, torch.Tensor) - - # Phi3v processor has inserted -1, -2 etc as placeholder in prompt_ids, - # which will cause OverflowError when decoding the prompt_ids. - # Therefore, we need to do an early replacement here - input_ids.masked_fill_(input_ids < 0, _IMAGE_TOKEN_ID) - - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index a08a17c324ca..c08eda612d1e 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -45,6 +45,7 @@ from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, BaseProcessingInfo, + HFMultiModalInputs, PromptReplacement, PromptUpdate, ResolvedPromptUpdate, @@ -867,40 +868,38 @@ def get_dummy_mm_data( class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]): - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) - if not mm_data: - return BatchFeature(dict(passthrough_data)) + # The Phi4MM processor expects "audios" instead of "audio" + hf_data = hf_inputs.hf_data + if "audio" in hf_data: + hf_data["audios"] = hf_data.pop("audio") - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + sr = self.info.get_feature_extractor(**hf_inputs.hf_kwargs).sampling_rate + if audio_data := hf_data.get("audios", []): + assert isinstance(audio_data, Sequence) + hf_data["audios"] = [(data, sr) for data in audio_data] - sr = self.info.get_feature_extractor(**hf_processor_mm_kwargs).sampling_rate - raw_audio_data = mm_data.get("audios", []) - assert isinstance(raw_audio_data, Sequence) - audio_data = [] - for data in raw_audio_data: - assert isinstance(data, (list, np.ndarray, torch.Tensor)) - audio_data.append(data) + return hf_inputs - processor_data = dict(text=prompt_text, **mm_data) - if audio_data: - processor_data["audios"] = [(data, sr) for data in audio_data] - - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - processor_data, - hf_processor_mm_kwargs, - ) + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if not hf_data: + return processed_data - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + hf_processor = self.info.get_hf_processor(**hf_kwargs) num_img_tokens = [ self.info.get_num_image_tokens( image_width=img_size[0], @@ -911,21 +910,26 @@ def _apply_hf_processor_main( ] processed_data["num_img_tokens"] = num_img_tokens - audio_features = processed_data["input_audio_embeds"] + audios = hf_data.get("audios", []) + assert isinstance(audios, list) + audio_features = processed_data.get("input_audio_embeds") + if audio_features is None: + return processed_data + + sr = self.info.get_feature_extractor(**hf_kwargs).sampling_rate feature_sizes = [ - self.info.get_audio_num_frames(len(audio), sr) for audio in audio_data + self.info.get_audio_num_frames(len(audio), sr) for audio, _ in audios ] processed_data["input_audio_embeds"] = [ audio_features[idx, :size] for idx, size in enumerate(feature_sizes) ] - processed_data.update(passthrough_data) return processed_data def _get_mm_fields_config( self, hf_inputs: BatchFeature, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> Mapping[str, MultiModalFieldConfig]: return dict( input_image_embeds=MultiModalFieldConfig.batched("image"), diff --git a/vllm/model_executor/models/phi4siglip.py b/vllm/model_executor/models/phi4siglip.py index 3c36b1f3eee5..fd9fa94c7d8b 100644 --- a/vllm/model_executor/models/phi4siglip.py +++ b/vllm/model_executor/models/phi4siglip.py @@ -143,39 +143,8 @@ def get_dummy_mm_data( class Phi4SiglipMultiModalProcessor( BaseMultiModalProcessor[Phi4SiglipProcessingInfo], ): - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - processed = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, - ) - - # The HF processor's tokenizer_image_token() replaces the "" - # string with IMAGE_TOKEN_INDEX (-200) in input_ids. This breaks - # vLLM's prompt-replacement pipeline which needs to find "" - # as normal sub-tokens. Re-tokenize with the plain tokenizer so - # that "" stays as sub-tokens and can be located by - # PromptReplacement. - # NOTE: tokenizer.__call__() (not .encode()) must be used so that - # added/special tokens like <|user|>, <|end|> are kept as single IDs. - tokenizer = self.info.get_tokenizer() - new_ids = tokenizer(prompt_text).input_ids - processed["input_ids"] = torch.tensor([new_ids]) - - processed_data = processed - processed_data.update(passthrough_data) - return processed_data + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index 42812add9b81..425bba74fdf6 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -14,7 +14,7 @@ from mistral_common.protocol.instruct.messages import UserMessage from mistral_common.protocol.instruct.request import ChatCompletionRequest from packaging.version import Version -from transformers import BatchFeature, PixtralVisionConfig +from transformers import PixtralVisionConfig from transformers.models.pixtral.image_processing_pixtral import ( _num_image_tokens as _get_pixtral_hf_num_image_tokens, ) @@ -41,7 +41,6 @@ from vllm.multimodal.inputs import ( MultiModalFieldConfig, MultiModalKwargsItem, - MultiModalKwargsOptionalItems, NestedTensors, ) from vllm.multimodal.parse import ( @@ -53,7 +52,8 @@ from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, BaseProcessingInfo, - MultiModalPromptUpdates, + HFMultiModalInputs, + MultiModalProcessingResult, PlaceholderFeaturesInfo, ProcessorInputs, PromptReplacement, @@ -67,7 +67,6 @@ MistralCommonImageProcessor, MistralCommonPixtralProcessor, ) -from vllm.utils.collection_utils import is_list_of from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -268,21 +267,19 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]) def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsOptionalItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - self._validate_mm_updates(mm_prompt_updates, mm_item_counts) + self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts) + self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts) mm_placeholders = self._find_mm_placeholders( - prompt_ids, - mm_prompt_updates, + mm_res.prompt_ids, + mm_res.prompt_updates, ) self._validate_mm_placeholders(mm_placeholders, mm_item_counts) - return prompt_ids, mm_placeholders + return mm_res.prompt_ids, mm_placeholders def _get_mm_fields_config( self, @@ -291,36 +288,18 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict(images=MultiModalFieldConfig.batched("image")) - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - - outputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - # Avoid padding issue - dict(**hf_processor_mm_kwargs, return_tensors=None), + # Avoid padding issue + return hf_inputs._replace( + hf_kwargs=dict(hf_inputs.hf_kwargs, return_tensors=None) ) - # Missing batch dimension - if is_list_of(outputs["input_ids"], int): - outputs["input_ids"] = [outputs["input_ids"]] - - processed_data = outputs - processed_data.update(passthrough_data) - return processed_data - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index bb6b153eddec..cecc8345e5ee 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -82,6 +82,8 @@ ) from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, + HFMultiModalInputs, + MultiModalProcessingResult, MultiModalPromptUpdates, PlaceholderFeaturesInfo, PromptReplacement, @@ -498,6 +500,64 @@ def get_dummy_mm_data( class Qwen2_5OmniThinkerMultiModalProcessor( BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo] ): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( + self, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + mm_counts = mm_items.get_all_counts() + if hf_kwargs.get("use_audio_in_video"): + num_videos = mm_counts.get("video", 0) + if mm_counts.get("audio", 0) < num_videos: + raise ValueError( + "Video doesn't have audio track with `audio_in_video=True`" + ) + + if num_videos and "text" in hf_inputs.hf_data: + standalone_mm_counts = dict(mm_counts) + standalone_mm_counts["audio"] -= num_videos + hf_inputs.hf_data["text"] = self._get_hf_mm_text(standalone_mm_counts) + + hf_data = hf_inputs.hf_data + normalized_kwargs = hf_inputs.hf_kwargs + audios = hf_data.pop("audio", []) + + # NOTE: WhisperFeatureExtractor cannot handle empty list of audios + if audios: + hf_data["audio"] = audios + elif ( + normalized_kwargs.get("use_audio_in_video") + and "audio" not in hf_data + and hf_data.get("videos") + ): + raise ValueError( + "Video doesn't have audio track with `audio_in_video=True`" + ) + + merged = self.info.ctx.get_merged_mm_kwargs(normalized_kwargs) + if hf_data.get("videos") and ( + merged.keys() & {"size", "min_pixels", "max_pixels"} + ): + normalized_kwargs = dict(normalized_kwargs) + video_size = dict(self.info.get_hf_processor().video_processor.size) + size_override = merged.get("size") + if size_override is not None: + video_size = video_size | size_override + min_pixels = merged.get("min_pixels") + if min_pixels is not None: + video_size["shortest_edge"] = min_pixels + max_pixels = merged.get("max_pixels") + if max_pixels is not None: + video_size["longest_edge"] = max_pixels + normalized_kwargs["size"] = video_size + + return hf_inputs._replace(hf_kwargs=normalized_kwargs) + def _get_mm_fields_config( self, hf_inputs: BatchFeature, @@ -570,12 +630,14 @@ def _derive_audio_from_video_placeholders( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: """Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.""" mm_item_counts = mm_items.get_all_counts() + prompt_ids = mm_res.prompt_ids + mm_kwargs = mm_res.kwargs + mm_prompt_updates = mm_res.prompt_updates + self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) @@ -789,101 +851,33 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int): ), ] - def _apply_hf_processor_main( + def _postprocess_hf_mm_data( self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, ) -> BatchFeature: - """Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.""" - mm_counts = mm_items.get_all_counts() - - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) - if use_audio_in_video and "video" in mm_counts: - if mm_counts.get("audio", 0) < mm_counts["video"]: - raise ValueError( - "Video doesn't have audio track with `audio_in_video=True`" - ) - mm_counts["audio"] -= mm_counts["video"] - - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_counts) - - mm_data = dict(processor_data) - audios = mm_data.pop("audios", []) - - # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - if audios: - # NOTE: Qwen2.5-Omni processor accept "audio" - mm_data["audio"] = audios - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - ) - elif ( - hf_processor_mm_kwargs.get("use_audio_in_video") - and "audio" not in mm_data - and mm_data.get("videos") - ): - # A subclass (e.g. Qwen3-Omni) may have already popped "audios" - # and populated mm_data["audio"] itself before delegating here, - # in which case this isn't a genuine "no audio" case. Likewise, - # mm_data can be empty on a multimodal-processor-cache hit, where - # there's nothing to (re-)process this call and the real result - # comes from the cache — not a genuine "no audio" case either. - raise ValueError( - "Video doesn't have audio track with `audio_in_video=True`" - ) - - merged = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs) - if mm_data.get("videos") and ( - merged.keys() & {"size", "min_pixels", "max_pixels"} - ): - hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs) - video_size = dict(self.info.get_hf_processor().video_processor.size) - size_override = merged.get("size") - if size_override is not None: - video_size = video_size | size_override - min_pixels = merged.get("min_pixels") - if min_pixels is not None: - video_size["shortest_edge"] = min_pixels - max_pixels = merged.get("max_pixels") - if max_pixels is not None: - video_size["longest_edge"] = max_pixels - hf_processor_mm_kwargs["size"] = video_size - - hf_inputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, - ) - - input_features = hf_inputs.pop("input_features", None) - feature_attention_mask = hf_inputs.get("feature_attention_mask", None) - if "input_audio_features" not in hf_inputs and input_features is not None: + input_features = processed_data.pop("input_features", None) + feature_attention_mask = processed_data.get("feature_attention_mask", None) + if "input_audio_features" not in processed_data and input_features is not None: if feature_attention_mask is not None: input_features = input_features.permute(0, 2, 1)[ feature_attention_mask.bool() ].permute(1, 0) - hf_inputs["input_audio_features"] = input_features + processed_data["input_audio_features"] = input_features if ( - "audio_feature_lengths" not in hf_inputs + "audio_feature_lengths" not in processed_data and feature_attention_mask is not None ): - hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1) + processed_data["audio_feature_lengths"] = feature_attention_mask.sum(-1) - video_second_per_grid = hf_inputs.get("video_second_per_grid", None) + video_second_per_grid = processed_data.get("video_second_per_grid", None) if video_second_per_grid is not None: - hf_inputs["second_per_grid_ts"] = video_second_per_grid - - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) - hf_inputs["use_audio_in_video"] = torch.tensor(use_audio_in_video) + processed_data["second_per_grid_ts"] = video_second_per_grid - processed_data = hf_inputs - processed_data.update(passthrough_data) - processed_data.pop("input_ids") + use_audio_in_video = hf_kwargs.get("use_audio_in_video", False) + processed_data["use_audio_in_video"] = torch.tensor(use_audio_in_video) + processed_data.pop("input_ids", None) return processed_data diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 27c603c866b2..7e7ffd5cfa20 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -59,6 +59,7 @@ PromptReplacement, PromptUpdate, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.sequence import IntermediateTensors from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -233,29 +234,25 @@ def get_dummy_mm_data( class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - if audios: - mm_data["audio"] = audios - - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, - truncation=True, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + truncation=True, + ) ) - return mm_data, hf_processor_mm_kwargs - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/qwen3_asr_realtime.py b/vllm/model_executor/models/qwen3_asr_realtime.py index 15ecd861f14b..bab689c605b3 100644 --- a/vllm/model_executor/models/qwen3_asr_realtime.py +++ b/vllm/model_executor/models/qwen3_asr_realtime.py @@ -37,11 +37,10 @@ ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache -from vllm.multimodal.inputs import MultiModalKwargsOptionalItems from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing.processor import ( BaseDummyInputsBuilder, - MultiModalPromptUpdates, + MultiModalProcessingResult, PlaceholderFeaturesInfo, cached_encode, ) @@ -114,10 +113,10 @@ def __init__( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsOptionalItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: + mm_kwargs = mm_res.kwargs + audios = mm_kwargs.get("audio", []) assert len(audios) == 1, ( f"Expected only one audio input for realtime, got {len(audios)}" @@ -147,7 +146,7 @@ def _maybe_apply_prompt_updates( # Find the audio_pad token position and expand it to audio_len tokens expanded_ids = list[int]() pad_start_idx = -1 - for i, tid in enumerate(prompt_ids): + for i, tid in enumerate(mm_res.prompt_ids): if tid == audio_pad_id and pad_start_idx == -1: pad_start_idx = i expanded_ids.extend([audio_pad_id] * audio_len) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index ce800e7de668..2c3647c81e98 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -73,6 +73,8 @@ ) from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems from vllm.multimodal.processing.processor import ( + HFMultiModalInputs, + MultiModalProcessingResult, MultiModalPromptUpdates, PlaceholderFeaturesInfo, PromptReplacement, @@ -1219,13 +1221,26 @@ def get_mm_max_tokens_per_item( class Qwen3OmniMoeThinkerMultiModalProcessor( Qwen2_5OmniThinkerMultiModalProcessor, ): - def _get_hf_mm_data( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) - audios = processor_data.get("audios") + if mm_items.get_count("audio", strict=False): + # TODO(Isotr0py): Remove this patch after upstream fix PR + # released and Transformers version update: + # https://github.com/huggingface/transformers/pull/41473 + hf_inputs = hf_inputs._replace( + hf_kwargs={ + **hf_inputs.hf_kwargs, + "audio_kwargs": dict(hf_inputs.hf_kwargs.get("audio_kwargs") or {}), + "text_kwargs": dict(hf_inputs.hf_kwargs.get("text_kwargs") or {}), + } + ) + + audios = hf_inputs.hf_data.get("audio") if audios: feature_extractor = self.info.get_feature_extractor() hop_length = feature_extractor.hop_length @@ -1239,47 +1254,30 @@ def pad_to_hop_length(x: np.ndarray) -> np.ndarray: # To make sure the cache works with padding=True, we pre-padded # the audio to a multiple of hop_length. - processor_data = dict(processor_data) - processor_data["audios"] = [ + hf_inputs.hf_data["audio"] = [ pad_to_hop_length(audio) if isinstance(audio, np.ndarray) else (pad_to_hop_length(audio[0]), audio[1]) for audio in audios ] - return processor_data, passthrough_data + return hf_inputs - def _apply_hf_processor_main( + def _postprocess_hf_mm_data( self, - mm_items: MultiModalDataItems, + mm_data: Mapping[str, object], hf_processor_mm_kwargs: Mapping[str, object], + processed_data: BatchFeature, ) -> BatchFeature: - if mm_items.get_all_counts().get("audio", 0): - # TODO(Isotr0py): Remove this patch after upstream fix PR - # released and Transformers version update: - # https://github.com/huggingface/transformers/pull/41473 - hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs) - hf_processor_mm_kwargs["audio_kwargs"] = dict( - hf_processor_mm_kwargs.get("audio_kwargs") or {} - ) - hf_processor_mm_kwargs["text_kwargs"] = dict( - hf_processor_mm_kwargs.get("text_kwargs") or {} - ) - - processed_data = super()._apply_hf_processor_main( - mm_items, - hf_processor_mm_kwargs, + processed_data = super()._postprocess_hf_mm_data( + mm_data, hf_processor_mm_kwargs, processed_data ) if ( "audio_feature_lengths" in processed_data and "feature_attention_mask" in processed_data ): - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - processor_data, _ = self._get_hf_mm_data(valid_mm_items) - audios = processor_data.get("audios", []) + audios = mm_data.get("audio", []) if audios: feature_extractor = self.info.get_feature_extractor( **hf_processor_mm_kwargs @@ -1312,12 +1310,14 @@ def _apply_hf_processor_main( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: """Qwen3-Omni reimplements this function to handle `use_audio_in_video`.""" mm_item_counts = mm_items.get_all_counts() + prompt_ids = mm_res.prompt_ids + mm_kwargs = mm_res.kwargs + mm_prompt_updates = mm_res.prompt_updates + self._validate_mm_kwargs(mm_kwargs, mm_item_counts) use_audio_in_video = False diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 482fd00e5d9b..78dda9c81503 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1327,27 +1327,25 @@ def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: proc_impl = getattr(type(hf_processor), "replace_video_token", None) return proc_impl is not None and proc_impl is not mixin_impl + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - mm_data = dict(mm_data) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) # Separate video processing from image processing. Because the videos # are processed into several image patches video_input_ids_lst: list[list[int]] = [] - if videos := mm_data.pop("videos", []): + if videos := hf_data.pop("videos", []): video_grid_thw_lst = [] pixel_values_videos_lst = [] timestamps_per_video = [] @@ -1370,10 +1368,8 @@ def _apply_hf_processor_main( # NOTE: a copy of is created to update do_sample_frames, # otherwise mm_hash for the object will be incorrect. - video_mm_kwargs = dict(**hf_processor_mm_kwargs) - merged = self.info.ctx.get_merged_mm_kwargs( - hf_processor_mm_kwargs, modality="video" - ) + video_mm_kwargs = dict(**hf_kwargs) + merged = self.info.ctx.get_merged_mm_kwargs(hf_kwargs, modality="video") if merged.keys() & {"size", "min_pixels", "max_pixels"}: video_size = dict(self.info.get_video_processor().size) size_override = merged.get("size") @@ -1485,13 +1481,11 @@ def _apply_hf_processor_main( # fps/num_frames are video-only kwargs already consumed by the loop; # exclude them so the text/image processor call below never gets a list. non_video_mm_kwargs = { - k: v - for k, v in hf_processor_mm_kwargs.items() - if k not in ("fps", "num_frames") + k: v for k, v in hf_kwargs.items() if k not in ("fps", "num_frames") } processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**non_video_mm_kwargs), - dict(text=prompt_text, **mm_data), + hf_data, non_video_mm_kwargs, ) @@ -1516,9 +1510,9 @@ def _apply_hf_processor_main( processed_data["input_ids"] = [expanded_ids] processed_data.update(video_outputs) - processed_data.update(passthrough_data) - - return processed_data + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) def _get_mm_fields_config( self, diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index 640f94647634..ddd1ac81ceb5 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -202,7 +202,9 @@ def apply( hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs with timing_ctx.record("apply_hf_processor"): - _, passthrough_data = self._get_hf_mm_data(mm_items) + passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_processor_mm_kwargs + ).passthrough_data mm_processed_data = BatchFeature( { k: torch.as_tensor(v).unsqueeze(0) diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index ff042e2f92b1..7b435d40ba02 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -219,18 +219,6 @@ class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo] Subclasses add the strategy for locating placeholders in the prompt. """ - def _get_hf_mm_data( - self, - mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - """Rename the parser's `audios` key to the `audio` argument HF audio - processors take. - """ - processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) - if self.info._is_audio_model() and "audios" in processor_data: - processor_data["audio"] = processor_data.pop("audios") - return processor_data, passthrough_data - def _get_modality_field_names(self, modality: str) -> set[str]: """Names of the fields the sub-processor for `modality` produces.""" # TODO: use else branch only once huggingface/transformers#44394 lands. @@ -353,6 +341,9 @@ def _get_mm_fields_config( return mm_fields + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _unpad_audios( self, hf_inputs: "BatchFeature", @@ -597,21 +588,31 @@ def apply( # work for Transformers. The vision path has logic tied to # `mm_tokens_per_modality` in _apply_vision() hf_processor_mm_kwargs = { - **hf_processor_mm_kwargs, + # vLLM needs the untruncated sequence to keep placeholder + # tokens aligned. Note that the text inputs are just dummy + # text, not the original prompt. The original prompt is + # already tokenized by the renderer. + "truncation": False, # HF processors only accept text, and the decoded string already # contains any special tokens, so don't let them be added again # (the HF processor call disables `add_special_tokens`). "add_special_tokens": False, + **hf_processor_mm_kwargs, } - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + processor_data, _, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_processor_mm_kwargs ) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + + # The real prompt is fed to the processor below and the placeholder + # ranges are read back out of its output, so the dummy text meant for + # media-only callers is discarded here. + processor_data.pop("text", None) + has_mm_data = any(processor_data.values()) # Ask which modality owns each token of the expanded prompt, which # is the only marking of the tokens an expansion adds around an item - if any(processor_data.values()): + if has_mm_data: processor_data = {**processor_data, "return_mm_token_type_ids": True} try: @@ -621,7 +622,7 @@ def apply( hf_processor_mm_kwargs, ) except ValueError: - if any(processor_data.values()): + if has_mm_data: raise # Some processors reject a prompt holding placeholders with # no data to go with them, so tokenize it without them @@ -792,31 +793,37 @@ def _get_num_patches_per_image( except (AttributeError, KeyError, TypeError): return None + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> "BatchFeature": - mm_counts = mm_items.get_all_counts() + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs + ) - valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + if not hf_data: + return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) - prompt_text = self.dummy_inputs.get_dummy_text(mm_counts) + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) # Ask for the replacement each placeholder expands to, and record it as # per-item fields: its token ids, and the tokens or patches behind them - if has_mm_data := any(mm_data.values()): - mm_data = {**mm_data, "return_text_replacement_offsets": True} + if has_mm_data := any(hf_data.values()): + hf_data = {**hf_data, "return_text_replacement_offsets": True} try: hf_inputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **mm_data), - hf_processor_mm_kwargs, + self.info.get_hf_processor(**hf_kwargs), + dict(text=prompt_text, **hf_data), + hf_kwargs, ) except ValueError: - if any(mm_data.values()): + if any(hf_data.values()): raise # Some processors reject a prompt holding placeholders with # no data to go with them, so tokenize it without them @@ -825,7 +832,7 @@ def _apply_hf_processor_main( dict(input_ids=[tokenizer.encode(prompt_text)]), tensor_type="pt" ) self._unpad_images(hf_inputs) - self._unpad_audios(hf_inputs, mm_data, hf_processor_mm_kwargs) + self._unpad_audios(hf_inputs, hf_data, hf_kwargs) # Drop the inputs the model would reject hf_inputs.pop("mm_token_type_ids", None) @@ -845,9 +852,10 @@ def _apply_hf_processor_main( "so it can be fixed, and install `transformers<5.15.0` in the " "meantime to locate placeholders in the expanded prompt instead." ) - hf_inputs.update(passthrough_data) hf_inputs.pop("input_ids", None) - return hf_inputs + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, hf_inputs + ) tokenizer = self.info.get_tokenizer() replacements = defaultdict[str, list[list[int]]](list) @@ -865,7 +873,7 @@ def _apply_hf_processor_main( ) if modality == "image": hf_inputs["num_image_patches"] = self._get_num_image_patches( - hf_inputs, mm_data, len(seqs) + hf_inputs, hf_data, len(seqs) ) elif modality == "audio": counts = [] @@ -874,10 +882,11 @@ def _apply_hf_processor_main( counts.append(int(ids.eq(_get_embed_token_id(ids)).sum())) hf_inputs["num_audio_tokens"] = torch.tensor(counts) - hf_inputs.update(passthrough_data) hf_inputs.pop("input_ids") - return hf_inputs + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, hf_inputs + ) # From this version on, a processor reporting no offsets is an error rather than a diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index 6465972cf161..3a1d017f4c46 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -42,6 +42,7 @@ PromptUpdate, cached_encode, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.ultravox import UltravoxConfig @@ -190,72 +191,32 @@ def get_dummy_mm_data( class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo]): - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - audios = mm_data.get("audios", []) - assert isinstance(audios, list) - - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, - include_audio_num_chunks=True, - ) - - return mm_data, hf_processor_mm_kwargs - - def _apply_hf_processor_main( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if processor_data: - processor_data, hf_processor_mm_kwargs = self._preprocess_hf_mm_data( - processor_data, - hf_processor_mm_kwargs, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + # The Ultravox processor accepts "audios" instead of "audio" + hf_data = hf_inputs.hf_data + if "audio" in hf_data: + hf_data["audios"] = hf_data.pop("audio") + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + normalized_kwargs = dict(hf_inputs.hf_kwargs) + # The remote processor already hardcodes truncation=False and forwards + # its kwargs to the audio processor, where another value would collide. + normalized_kwargs.pop("truncation", None) + return hf_inputs._replace( + hf_kwargs=dict( + normalized_kwargs, + sampling_rate=feature_extractor.sampling_rate, + include_audio_num_chunks=True, ) - - prompt_text = self._get_hf_processor_text(mm_items.get_all_counts()) - if prompt_text is not None: - processor_data = dict(text=prompt_text, **processor_data) - - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - - def patched_call(**kwargs) -> BatchFeature: - # The remote UltravoxProcessor does not support the - # `truncation` kwarg injected by `call_hf_processor`: it - # passes `truncation` to its audio processor by itself while - # also forwarding `**kwargs` there, causing a duplicate - # keyword argument error. Its hardcoded `truncation=False` - # already matches vLLM's intended behavior. - kwargs.pop("truncation", None) - - return hf_processor(**kwargs) - - processed_data = self.info.ctx.call_hf_processor( - patched_call, - processor_data, - hf_processor_mm_kwargs, - ) - processed_data.update(passthrough_data) - else: - processed_data = BatchFeature(dict(passthrough_data)) - - return self._postprocess_hf_mm_data( - processor_data, - hf_processor_mm_kwargs, - processed_data, ) def _postprocess_hf_mm_data( diff --git a/vllm/model_executor/models/unlimited_ocr.py b/vllm/model_executor/models/unlimited_ocr.py index c96e5f293d91..d248e5c8e068 100644 --- a/vllm/model_executor/models/unlimited_ocr.py +++ b/vllm/model_executor/models/unlimited_ocr.py @@ -75,7 +75,7 @@ from vllm.multimodal.processing import PromptReplacement, PromptUpdate from vllm.multimodal.processing.context import TimingContext from vllm.multimodal.processing.inputs import ProcessorInputs -from vllm.multimodal.processing.processor import MultiModalProcessingInfo +from vllm.multimodal.processing.processor import MultiModalProcessingResult from vllm.transformers_utils.processors.deepseek_ocr import ( BASE_SIZE, CROP_MODE, @@ -225,7 +225,7 @@ def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: # The processor logic differs for single-image (crop) vs multi-image # (no crop) requests. The processing cache assumes per-item output is # invariant of how many images are passed per prompt, so we only cache diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index 5559217d3737..ad54e8eb20bc 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -16,7 +16,7 @@ from mistral_common.protocol.instruct.request import ChatCompletionRequest from mistral_common.protocol.transcription.request import TranscriptionRequest from mistral_common.tokens.tokenizers.audio import Audio -from transformers import BatchFeature, WhisperConfig +from transformers import WhisperConfig from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.config.multimodal import BaseDummyOptions @@ -36,7 +36,6 @@ from vllm.multimodal.inputs import ( MultiModalFieldConfig, MultiModalKwargsItems, - MultiModalKwargsOptionalItems, NestedTensors, ) from vllm.multimodal.parse import ( @@ -48,7 +47,8 @@ from vllm.multimodal.processing.processor import ( BaseMultiModalProcessor, BaseProcessingInfo, - MultiModalPromptUpdates, + HFMultiModalInputs, + MultiModalProcessingResult, PlaceholderFeaturesInfo, ProcessorInputs, PromptReplacement, @@ -61,7 +61,6 @@ MistralCommonFeatureExtractor, MistralCommonVoxtralProcessor, ) -from vllm.utils.collection_utils import is_list_of from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsTranscription from .utils import init_vllm_registered_model, maybe_prefix @@ -209,21 +208,19 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]) def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsOptionalItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - self._validate_mm_updates(mm_prompt_updates, mm_item_counts) + self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts) + self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts) mm_placeholders = self._find_mm_placeholders( - prompt_ids, - mm_prompt_updates, + mm_res.prompt_ids, + mm_res.prompt_updates, ) self._validate_mm_placeholders(mm_placeholders, mm_item_counts) - return prompt_ids, mm_placeholders + return mm_res.prompt_ids, mm_placeholders def _get_mm_fields_config( self, @@ -241,38 +238,20 @@ def _validate_mm_placeholders( # skip validation here pass - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - - if audios: - # MistralCommonVoxtralProcessor accepts "audio" - mm_data["audio"] = audios + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) # Avoid padding issue - hf_processor_mm_kwargs = dict(**hf_processor_mm_kwargs, return_tensors=None) - - return mm_data, hf_processor_mm_kwargs - - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - # Missing batch dimension - input_ids = processed_data.get("input_ids") - if is_list_of(input_ids, int): - processed_data["input_ids"] = [input_ids] - - return processed_data + return hf_inputs._replace( + hf_kwargs=dict(hf_inputs.hf_kwargs, return_tensors=None) + ) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 52b636cc2586..162fecf03870 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -28,11 +28,10 @@ ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache -from vllm.multimodal.inputs import MultiModalKwargsOptionalItems from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import BaseDummyInputsBuilder from vllm.multimodal.processing.processor import ( - MultiModalPromptUpdates, + MultiModalProcessingResult, PlaceholderFeaturesInfo, ) from vllm.sequence import IntermediateTensors @@ -60,10 +59,10 @@ def __init__( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsOptionalItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: + mm_kwargs = mm_res.kwargs + # there are no placeholder audio tokens for streaming # so we need to build the place placeholder positions manually @@ -86,7 +85,7 @@ def _maybe_apply_prompt_updates( * [0], # only used for length computation, so we can take dummy inputs is_embed=None, ) - return prompt_ids, {"audio": [features_info]} + return mm_res.prompt_ids, {"audio": [features_info]} class TimeEmbedding(torch.nn.Module): diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index 7fc620e7dde2..b50d04d04b59 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -61,6 +61,7 @@ PromptReplacement, PromptUpdate, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.renderers import TokenizeParams from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.jsontree import json_map_leaves @@ -743,6 +744,9 @@ def get_dummy_mm_data( class WhisperMultiModalProcessor(EncDecMultiModalProcessor[WhisperProcessingInfo]): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + def create_encoder_prompt( self, prompt: list[int], @@ -754,37 +758,21 @@ def create_encoder_prompt( # for encoder profiling. return [0] - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) - - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - - mm_data = dict(mm_data) - mm_data["audio"] = mm_data.pop("audios") - - hf_processor_mm_kwargs = dict( - **hf_processor_mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + return hf_inputs._replace( + hf_kwargs=dict( + hf_inputs.hf_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) ) - return mm_data, hf_processor_mm_kwargs - - def _postprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if "labels" in processed_data: - processed_data["input_ids"] = processed_data.pop("labels") - - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/models/dots3_note/common/processor.py b/vllm/models/dots3_note/common/processor.py index e8d365aa380d..addaf77c8141 100644 --- a/vllm/models/dots3_note/common/processor.py +++ b/vllm/models/dots3_note/common/processor.py @@ -34,6 +34,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.transformers_utils.repo_utils import get_hf_file_to_dict from .video import preprocess_dots3_note_video @@ -630,26 +631,33 @@ def get_dummy_mm_data( class Dots3NoteMultiModalProcessor(BaseMultiModalProcessor[Dots3NoteProcessingInfo]): - def _get_hf_mm_data( + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) - if "video" not in mm_items: - return processor_data, passthrough_data - - videos = mm_items.get_items("video", VideoProcessorItems) - raw_videos: list[object] = [] - for index, item in enumerate(videos.data): - if isinstance(item, MediaWithBytes): - raw_videos.append(item.original_bytes) - else: - raw_videos.append(videos.get(index)) - processor_data = dict(processor_data) - processor_data["videos"] = raw_videos - return processor_data, passthrough_data - - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str: + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + + # The Dots3Note processor accepts "audios" instead of "audio" + hf_data = hf_inputs.hf_data + if "audio" in hf_data: + hf_data["audios"] = hf_data.pop("audio") + + if "video" in mm_items: + videos = mm_items.get_items("video", VideoProcessorItems) + + raw_videos: list[object] = [] + for index, item in enumerate(videos.data): + if isinstance(item, MediaWithBytes): + raw_videos.append(item.original_bytes) + else: + raw_videos.append(videos.get(index)) + + hf_data["videos"] = raw_videos + + return hf_inputs + + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) def _get_mm_fields_config( diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py index 249d7edeb43f..cc88c1252951 100644 --- a/vllm/models/inkling/common/mm_preprocess.py +++ b/vllm/models/inkling/common/mm_preprocess.py @@ -182,20 +182,19 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo]) def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + processor = self.info.get_hf_processor(**hf_kwargs) tokenizer = self.info.get_tokenizer() images = mm_data.get("images") or [] - audios = mm_data.get("audios") or [] + audios = mm_data.get("audio") or [] if not isinstance(images, list): images = list(cast(Iterable[Any], images)) if not isinstance(audios, list): @@ -235,8 +234,9 @@ def _apply_hf_processor_main( data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64) processed_data = BatchFeature(data=data, tensor_type=None) - processed_data.update(passthrough_data) - return processed_data + return self._finalize_hf_mm_data( + mm_data, hf_kwargs, passthrough_data, processed_data + ) def _tokenize_with_placeholders( self, diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py index 9f1d08011a03..4e8b791093c4 100644 --- a/vllm/models/minimax_m3/common/mm_preprocess.py +++ b/vllm/models/minimax_m3/common/mm_preprocess.py @@ -32,6 +32,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.processing.processor import HFMultiModalInputs from vllm.multimodal.video import ( VIDEO_LOADER_REGISTRY, VideoBackend, @@ -307,26 +308,21 @@ def get_dummy_mm_data( class MiniMaxM3VLMultiModalProcessor( BaseMultiModalProcessor[MiniMaxM3VLProcessingInfo] ): - def _apply_hf_processor_main( + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} - ) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - - if not mm_data: - return BatchFeature(dict(passthrough_data)) - - prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + hf_data = hf_inputs.hf_data - mm_data = dict(mm_data) # With ``video_needs_metadata=True`` each video arrives as a # ``(frames, metadata)`` tuple. Split the frames back out and forward the # metadata as ``VideoMetadata`` so the processor emits timestamps. - videos = cast(list | None, mm_data.get("videos")) + videos = cast(list | None, hf_data.get("videos")) video_metadata: list[VideoMetadata] | None = None if videos: frames_only = [] @@ -346,21 +342,25 @@ def _apply_hf_processor_main( # stays consistent with _get_prompt_updates. meta.setdefault("total_num_frames", len(frames)) video_metadata.append(VideoMetadata(**meta)) - mm_data["videos"] = frames_only + hf_data["videos"] = frames_only - # Override the video processor's default do_resize=False (set for a - # pre-resized pipeline) to True for vLLM's raw-frame inputs. - merged = dict(do_resize=True, **hf_processor_mm_kwargs) - data = dict(text=prompt_text, **mm_data) if video_metadata is not None: - data["video_metadata"] = video_metadata - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - data, + hf_data["video_metadata"] = video_metadata + + return hf_inputs + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Override the video processor's default for vLLM's raw-frame inputs. + merged = dict(do_resize=True, **hf_kwargs) + return self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**hf_kwargs), + hf_data, merged, ) - processed_data.update(passthrough_data) - return processed_data def _get_mm_fields_config( self, diff --git a/vllm/multimodal/processing/context.py b/vllm/multimodal/processing/context.py index c1b81a4d7210..359c43700d14 100644 --- a/vllm/multimodal/processing/context.py +++ b/vllm/multimodal/processing/context.py @@ -321,12 +321,6 @@ def call_hf_processor( merged_kwargs = self.get_merged_mm_kwargs(kwargs) - # vLLM needs the full untruncated sequence to keep multi-modal - # placeholder tokens aligned; note that the text inputs in - # call_hf_processor are just dummy text, not the original prompt. - # The original prompt is already tokenized by the renderer. - merged_kwargs.setdefault("truncation", False) - allowed_kwargs = get_allowed_kwarg_only_overrides( hf_processor, merged_kwargs, @@ -426,7 +420,7 @@ def allow_missing_mm_embeddings(self) -> bool: def get_data_parser(self) -> MultiModalDataParser: """Constructs a parser to preprocess multi-modal data items before passing them to - [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data]. + [`_get_hf_mm_inputs`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs]. You can support additional modalities by creating a subclass of [`MultiModalDataParser`][vllm.multimodal.parse.MultiModalDataParser] @@ -506,7 +500,7 @@ def parse_mm_data( """Normalize [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] to [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems] before passing them to - [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data]. + [`_get_hf_mm_inputs`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs]. """ mm_items = self.data_parser.parse_mm_data(mm_data) diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 8e9983798f7f..1efe84219ee1 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1008,7 +1008,20 @@ def find_mm_placeholders( _I = TypeVar("_I", bound=BaseProcessingInfo) -class MultiModalProcessingInfo(NamedTuple): +class HFMultiModalInputs(NamedTuple): + hf_data: dict[str, object] + """The data inputs to the HF processor.""" + + hf_kwargs: Mapping[str, object] + """The non-data inputs to the HF processor.""" + + passthrough_data: dict[str, object] + """The inputs to assign to the output of the HF processor, + instead of passing them to the HF processor.""" + + +class MultiModalProcessingResult(NamedTuple): + prompt_ids: list[int] kwargs: MultiModalKwargsOptionalItems hashes: MultiModalHashes prompt_updates: MultiModalPromptUpdates @@ -1137,23 +1150,8 @@ def _find_mm_placeholders( ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: return find_mm_placeholders(new_token_ids, mm_prompt_updates) - def _get_hf_mm_data( - self, - mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - """Extract processor and passthrough data from multi-modal items.""" - processor_data = dict[str, object]() - passthrough_data = dict[str, object]() - - for items in mm_items.values(): - processor_data.update(items.get_processor_data()) - passthrough_data.update(items.get_passthrough_data()) - - return processor_data, passthrough_data - - def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str | None: - """Get the text to pass to the HF processor alongside the multi-modal - data. + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str | None: + """Get the text to pass to the HF processor alongside the multi-modal data. By default, no text is passed. If the HF processor requires that text and multi-modal items correspond to each other, you should @@ -1162,29 +1160,53 @@ def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str | None: """ return None - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - """Pre-process the multi-modal data and HF processor keyword arguments - before they are passed to the HF processor. - - By default, both are returned as-is. If the HF processor expects the - multi-modal data under different keys than those provided by the - multi-modal items (e.g. `audio` instead of `audios`), or requires - additional keyword arguments (e.g. `sampling_rate`), you should - override this method. + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> HFMultiModalInputs: + """Extract the multi-modal data and corresponding keyword arguments + for HF processing. + + If the HF processor expects the multi-modal data under different keys + than those provided by the multi-modal items, or requires additional + keyword arguments, you should override this method. """ - return mm_data, hf_processor_mm_kwargs + hf_data = dict[str, object]() + passthrough_data = dict[str, object]() + + for items in mm_items.values(): + if not items: + continue + + hf_data.update(items.get_processor_data()) + passthrough_data.update(items.get_passthrough_data()) + + if hf_data: + prompt_text = self._get_hf_mm_text(mm_items.get_all_counts()) + if prompt_text is not None: + hf_data = dict(text=prompt_text, **hf_data) + + # vLLM needs the untruncated sequence to keep placeholder + # tokens aligned. Note that the text inputs are just dummy + # text, not the original prompt. The original prompt is + # already tokenized by the renderer. + hf_kwargs = {"truncation": False, **hf_kwargs} + + # HF processors accept (images, text, videos, audio) + # so we need to remap vLLM's "audios" -> HF "audio" + if "audios" in hf_data: + hf_data["audio"] = hf_data.pop("audios") + + return HFMultiModalInputs(hf_data, hf_kwargs, passthrough_data) def _postprocess_hf_mm_data( self, - mm_data: Mapping[str, object], + mm_inputs: Mapping[str, object], hf_processor_mm_kwargs: Mapping[str, object], processed_data: BatchFeature, ) -> BatchFeature: - """Post-process the output of the HF processor. + """Post-process the combined processor and passthrough data. By default, the output is returned as-is. If you need to modify the output of the HF processor before it is converted into multi-modal @@ -1192,41 +1214,51 @@ def _postprocess_hf_mm_data( """ return processed_data + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + """Call the HF processor on the normalized multi-modal inputs.""" + return self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**hf_kwargs), + hf_data, + hf_kwargs, + ) + + def _finalize_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + passthrough_data: Mapping[str, object], + processed_data: BatchFeature | None = None, + ) -> BatchFeature: + """Merge passthrough fields and apply model-specific postprocessing.""" + if processed_data is None: + from transformers.feature_extraction_utils import BatchFeature + + processed_data = BatchFeature() + + processed_data.update(passthrough_data) + return self._postprocess_hf_mm_data(hf_data, hf_kwargs, processed_data) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Apply the HF processor on the multi-modal data.""" - valid_mm_items = mm_items.select( - {k for k, c in mm_items.get_all_counts().items() if c > 0} + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs ) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - if processor_data: - processor_data, hf_processor_mm_kwargs = self._preprocess_hf_mm_data( - processor_data, hf_processor_mm_kwargs - ) - - prompt_text = self._get_hf_processor_text(mm_items.get_all_counts()) - if prompt_text is not None: - processor_data = dict(text=prompt_text, **processor_data) - - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - processor_data, - hf_processor_mm_kwargs, - ) - processed_data.update(passthrough_data) + if hf_data: + processed_data = self._call_hf_processor(hf_data, hf_kwargs) else: - from transformers.feature_extraction_utils import BatchFeature - - processed_data = BatchFeature(dict(passthrough_data)) + processed_data = None - return self._postprocess_hf_mm_data( - processor_data, - hf_processor_mm_kwargs, - processed_data, + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data ) def _postprocess_prompt(self, prompt: list[int]) -> list[int]: @@ -1261,6 +1293,9 @@ def _get_cache_missing_items( mm_missing_data = {} for modality, idxs in mm_missing_idxs.items(): + if not idxs: + continue + missing_modality_data = [] for idx in idxs: data = mm_data_items[modality][idx] @@ -1271,6 +1306,7 @@ def _get_cache_missing_items( ) else: missing_modality_data.append(data) + mm_missing_data[modality] = missing_modality_data mm_missing_items = self.info.parse_mm_data(mm_missing_data, validate=False) @@ -1342,11 +1378,11 @@ def _apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: with timing_ctx.record("apply_hf_processor"): mm_processed_data = self._apply_hf_processor_main( mm_items=inputs.mm_data_items, - hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs, + hf_kwargs=inputs.hf_processor_mm_kwargs, ) mm_kwargs = MultiModalKwargsItems.from_hf_inputs( @@ -1369,26 +1405,28 @@ def _apply_hf_processor( mm_kwargs, ) - mm_info = MultiModalProcessingInfo( + return MultiModalProcessingResult( + prompt_ids=self._postprocess_prompt(inputs.prompt), kwargs=mm_kwargs, hashes=mm_hashes, prompt_updates=mm_prompt_updates, ) - return mm_info - def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: """Apply the HF processor on the full prompt text, caching the results and reusing cached results. """ cache = self.cache + has_passthrough_data = any( + len(items.get_passthrough_data()) > 0 + for items in inputs.mm_data_items.values() + ) - _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items) - if cache is None or passthrough_data: + if cache is None or has_passthrough_data: return self._apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("get_mm_hashes"): @@ -1410,7 +1448,7 @@ def _cached_apply_hf_processor( with timing_ctx.record("apply_hf_processor"): mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, - hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs, + hf_kwargs=inputs.hf_processor_mm_kwargs, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( @@ -1435,14 +1473,13 @@ def _cached_apply_hf_processor( mm_missing_prompt_updates=mm_missing_prompt_updates, ) - mm_info = MultiModalProcessingInfo( + return MultiModalProcessingResult( + prompt_ids=self._postprocess_prompt(inputs.prompt), kwargs=mm_kwargs, hashes=mm_hashes, prompt_updates=mm_prompt_updates, ) - return mm_info - def _apply_token_matches( self, prompt: list[int], @@ -1670,17 +1707,15 @@ def _validate_mm_placeholders( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsOptionalItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - self._validate_mm_updates(mm_prompt_updates, mm_item_counts) + self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts) + self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts) prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, + mm_res.prompt_ids, + mm_res.prompt_updates, ) self._validate_mm_placeholders(mm_placeholders, mm_item_counts) @@ -1705,15 +1740,12 @@ def apply( 4. Extract information about the placeholder tokens from the processed token IDs. """ - prompt_ids = self._postprocess_prompt(inputs.prompt) - mm_info = self._cached_apply_hf_processor(inputs, timing_ctx) + mm_res = self._cached_apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("apply_prompt_updates"): prompt_ids, mm_placeholders = self._maybe_apply_prompt_updates( mm_items=inputs.mm_data_items, - prompt_ids=prompt_ids, - mm_kwargs=mm_info.kwargs, - mm_prompt_updates=mm_info.prompt_updates, + mm_res=mm_res, ) mm_placeholder_ranges = { @@ -1723,8 +1755,8 @@ def apply( return mm_input( prompt_token_ids=prompt_ids, - mm_kwargs=mm_info.kwargs, - mm_hashes=mm_info.hashes, + mm_kwargs=mm_res.kwargs, + mm_hashes=mm_res.hashes, mm_placeholders=mm_placeholder_ranges, ) From 1085b64425a9e6f5ca52876ad32e55fda5665f4e Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 16 Sep 2026 13:21:43 -0400 Subject: [PATCH 336/774] [Benchmark] Retire stale benchmarks and consolidate RMSNorm (#57083) Signed-off-by: Misha Goin --- benchmarks/cutlass_benchmarks/utils.py | 39 -- .../cutlass_benchmarks/w8a8_benchmarks.py | 372 ------------------ .../cutlass_benchmarks/weight_shapes.py | 46 --- benchmarks/kernels/benchmark_rmsnorm.py | 255 ------------ benchmarks/kernels/ir/shapes.py | 5 +- benchmarks/overheads/benchmark_hashing.py | 64 --- tools/pre_commit/check_forbidden_imports.py | 2 - 7 files changed, 3 insertions(+), 780 deletions(-) delete mode 100644 benchmarks/cutlass_benchmarks/utils.py delete mode 100644 benchmarks/cutlass_benchmarks/w8a8_benchmarks.py delete mode 100644 benchmarks/cutlass_benchmarks/weight_shapes.py delete mode 100644 benchmarks/kernels/benchmark_rmsnorm.py delete mode 100644 benchmarks/overheads/benchmark_hashing.py diff --git a/benchmarks/cutlass_benchmarks/utils.py b/benchmarks/cutlass_benchmarks/utils.py deleted file mode 100644 index 659c68bb11d7..000000000000 --- a/benchmarks/cutlass_benchmarks/utils.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Cutlass bench utils - -import torch - - -def to_fp8(tensor: torch.Tensor) -> torch.Tensor: - finfo = torch.finfo(torch.float8_e4m3fn) - return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to( - dtype=torch.float8_e4m3fn - ) - - -def to_int8(tensor: torch.Tensor) -> torch.Tensor: - return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8) - - -def to_bf16(tensor: torch.Tensor) -> torch.Tensor: - return tensor.to(dtype=torch.bfloat16) - - -def to_fp16(tensor: torch.Tensor) -> torch.Tensor: - return tensor.to(dtype=torch.float16) - - -def make_rand_tensors( - dtype: torch.dtype, m: int, n: int, k: int -) -> tuple[torch.Tensor, torch.Tensor]: - a = torch.randn((m, k), device="cuda") * 5 - b = torch.randn((n, k), device="cuda").t() * 5 - - if dtype == torch.int8: - return to_int8(a), to_int8(b) - if dtype == torch.float8_e4m3fn: - return to_fp8(a), to_fp8(b) - - raise ValueError("unsupported dtype") diff --git a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py deleted file mode 100644 index f7325ddd2cbb..000000000000 --- a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py +++ /dev/null @@ -1,372 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import copy -import itertools -import pickle as pkl -import time -from collections.abc import Callable, Iterable - -import torch -import torch.utils.benchmark as TBenchmark -from torch.utils.benchmark import Measurement as TMeasurement -from utils import make_rand_tensors -from weight_shapes import WEIGHT_SHAPES - -from vllm import _custom_ops as ops -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - w8a8_triton_block_scaled_mm, -) -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.utils.math_utils import cdiv - -DEFAULT_MODELS = list(WEIGHT_SHAPES.keys()) -DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512] -DEFAULT_TP_SIZES = [1] - - -# bench -def bench_fn( - label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs -) -> TMeasurement: - min_run_time = 1 - - globals = { - "args": args, - "kwargs": kwargs, - "fn": fn, - } - return TBenchmark.Timer( - stmt="fn(*args, **kwargs)", - globals=globals, - label=label, - sub_label=sub_label, - description=description, - ).blocked_autorange(min_run_time=min_run_time) - - -def bench_int8( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - """Benchmark INT8-based kernels.""" - assert dtype == torch.int8 - a, b = make_rand_tensors(torch.int8, m, n, k) - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32) - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32) - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16) - azp = torch.zeros((m,), device="cuda", dtype=torch.int32) - azp_adj = torch.zeros((n,), device="cuda", dtype=torch.int32) - - bench_fns = { - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16) - ), - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.float16), b.to(dtype=torch.float16) - ), - "cutlass_i8_i8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16 - ), - "cutlass_i8_i8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16, bias - ), - "cutlass_i8_i8_bf16_scaled_mm_azp": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_bias": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, None, bias - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_pt": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp, bias - ), - } - - timers = [] - for name, fn in bench_fns.items(): - # If bench_kernels is None, run all. Otherwise, run only exact matches. - if bench_kernels is None or name in bench_kernels: - print(f"Running {name}") - timers.append(bench_fn(label, sub_label, name, fn)) - - return timers - - -def bench_fp8( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - """Benchmark FP8-based kernels.""" - assert dtype == torch.float8_e4m3fn - a, b = make_rand_tensors(torch.float8_e4m3fn, m, n, k) - a_cont = a.contiguous() - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32) - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32) - - block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32) - block_scale_b = torch.rand( - cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32 - ) - block_scale_a_M_major = block_scale_a.t().contiguous().t() - block_scale_b_K_major = block_scale_b.t().contiguous().t() - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16) - - print(m, k, n) - - bench_fns = { - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16) - ), - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.float16), b.to(dtype=torch.float16) - ), - "pytorch_fp8_fp8_fp16_scaled_mm": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.float16 - ), - "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.float16, use_fast_accum=True - ), - "pytorch_fp8_fp8_bf16_scaled_mm": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.bfloat16 - ), - "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.bfloat16, use_fast_accum=True - ), - "cutlass_fp8_fp8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16 - ), - "cutlass_fp8_fp8_fp16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.float16 - ), - "cutlass_fp8_fp8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16, bias - ), - "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16) - ), - "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm( - a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128) - ), - "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm( - a, b, block_scale_a_M_major, block_scale_b_K_major, torch.float16 - ), - } - - timers = [] - for name, fn in bench_fns.items(): - # If bench_kernels is None, run all. Otherwise, run only exact matches. - if bench_kernels is None or name in bench_kernels: - print(f"Running {name}") - timers.append(bench_fn(label, sub_label, name, fn)) - - return timers - - -def bench( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - if dtype == torch.int8: - return bench_int8(dtype, m, k, n, label, sub_label, bench_kernels) - if dtype == torch.float8_e4m3fn: - return bench_fp8(dtype, m, k, n, label, sub_label, bench_kernels) - raise ValueError("unsupported type") - - -# runner -def print_timers(timers: Iterable[TMeasurement]): - compare = TBenchmark.Compare(timers) - compare.print() - - -def run( - dtype: torch.dtype, - MKNs: Iterable[tuple[int, int, int]], - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - results = [] - for m, k, n in MKNs: - timers = bench( - dtype, - m, - k, - n, - f"scaled-{dtype}-gemm", - f"MKN=({m}x{k}x{n})", - bench_kernels=bench_kernels, - ) - print_timers(timers) - results.extend(timers) - return results - - -def make_output( - data: Iterable[TMeasurement], - MKNs: Iterable[tuple[int, int, int]], - base_description: str, - timestamp=None, -): - print(f"== All Results {base_description} ====") - print_timers(data) - - # pickle all the results - timestamp = int(time.time()) if timestamp is None else timestamp - with open(f"{base_description}-{timestamp}.pkl", "wb") as f: - pkl.dump(data, f) - - -def run_square_bench(args): - dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment)) - MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes)) - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - make_output(data, MKNs, f"square_bench-{args.dtype}") - - -def run_range_bench(args): - dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment)) - n = len(dim_sizes) - Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes - Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes - Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes - MKNs = list(zip(Ms, Ks, Ns)) - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - make_output(data, MKNs, f"range_bench-{args.dtype}") - - -def run_model_bench(args): - print("Benchmarking models:") - for i, model in enumerate(args.models): - print(f"[{i}] {model}") - - def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]: - KNs = [] - for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]): - KN[tp_split_dim] = KN[tp_split_dim] // tp_size - KNs.append(KN) - return KNs - - model_bench_data = [] - models_tps = list(itertools.product(args.models, args.tp_sizes)) - for model, tp_size in models_tps: - Ms = args.batch_sizes - KNs = model_shapes(model, tp_size) - MKNs = [] - for m in Ms: - for k, n in KNs: - MKNs.append((m, k, n)) - - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - model_bench_data.append(data) - - # Print all results - for data, model_tp in zip(model_bench_data, models_tps): - model, tp_size = model_tp - print(f"== Results {args.dtype} {model}-TP{tp_size} ====") - print_timers(data) - - timestamp = int(time.time()) - - all_data = [] - for d in model_bench_data: - all_data.extend(d) - # pickle all data - with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f: - pkl.dump(all_data, f) - - -if __name__ == "__main__": - - def to_torch_dtype(dt): - if dt == "int8": - return torch.int8 - if dt == "fp8": - return torch.float8_e4m3fn - raise ValueError("unsupported dtype") - - parser = FlexibleArgumentParser( - description=""" -Benchmark Cutlass GEMM. - - To run square GEMMs: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64 - - To run constant N and K and sweep M: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384 - - To run dimensions from a model: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1 - - Output: - - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs. - """, # noqa: E501 - formatter_class=argparse.RawTextHelpFormatter, - ) - - parser.add_argument( - "--dtype", - type=to_torch_dtype, - required=True, - help="Available options are ['int8', 'fp8']", - ) - parser.add_argument( - "--kernels", - nargs="+", - type=str, - default=None, - help="Exact names of the kernels to benchmark. If not set, runs all kernels.", - ) - - subparsers = parser.add_subparsers(dest="cmd") - - square_parser = subparsers.add_parser("square_bench") - square_parser.add_argument("--dim-start", type=int, required=True) - square_parser.add_argument("--dim-end", type=int, required=True) - square_parser.add_argument("--dim-increment", type=int, required=True) - square_parser.set_defaults(func=run_square_bench) - - range_parser = subparsers.add_parser("range_bench") - range_parser.add_argument("--dim-start", type=int, required=True) - range_parser.add_argument("--dim-end", type=int, required=True) - range_parser.add_argument("--dim-increment", type=int, required=True) - range_parser.add_argument("--m-constant", type=int, default=None) - range_parser.add_argument("--n-constant", type=int, default=None) - range_parser.add_argument("--k-constant", type=int, default=None) - range_parser.set_defaults(func=run_range_bench) - - model_parser = subparsers.add_parser("model_bench") - model_parser.add_argument( - "--models", - nargs="+", - type=str, - default=DEFAULT_MODELS, - choices=WEIGHT_SHAPES.keys(), - ) - model_parser.add_argument( - "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES - ) - model_parser.add_argument( - "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES - ) - model_parser.set_defaults(func=run_model_bench) - - args = parser.parse_args() - args.func(args) diff --git a/benchmarks/cutlass_benchmarks/weight_shapes.py b/benchmarks/cutlass_benchmarks/weight_shapes.py deleted file mode 100644 index 25b96ef56620..000000000000 --- a/benchmarks/cutlass_benchmarks/weight_shapes.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Weight Shapes are in the format -# ([K, N], TP_SPLIT_DIM) -# Example: -# A shape of ([14336, 4096], 0) indicates the following GEMM shape, -# - TP1 : K = 14336, N = 4096 -# - TP2 : K = 7168, N = 4096 -# A shape of ([4096, 6144], 1) indicates the following GEMM shape, -# - TP1 : K = 4096, N = 6144 -# - TP4 : K = 4096, N = 1536 - -# TP1 shapes -WEIGHT_SHAPES = { - "mistralai/Mistral-7B-v0.1": [ - ([4096, 6144], 1), - ([4096, 4096], 0), - ([4096, 28672], 1), - ([14336, 4096], 0), - ], - "meta-llama/Llama-2-7b-hf": [ - ([4096, 12288], 1), - ([4096, 4096], 0), - ([4096, 22016], 1), - ([11008, 4096], 0), - ], - "meta-llama/Llama-3-8b": [ - ([4096, 6144], 1), - ([4096, 4096], 0), - ([4096, 28672], 1), - ([14336, 4096], 0), - ], - "meta-llama/Llama-2-13b-hf": [ - ([5120, 15360], 1), - ([5120, 5120], 0), - ([5120, 27648], 1), - ([13824, 5120], 0), - ], - "meta-llama/Llama-2-70b-hf": [ - ([8192, 10240], 1), - ([8192, 8192], 0), - ([8192, 57344], 1), - ([28672, 8192], 0), - ], -} diff --git a/benchmarks/kernels/benchmark_rmsnorm.py b/benchmarks/kernels/benchmark_rmsnorm.py deleted file mode 100644 index d8d7f5bcf9da..000000000000 --- a/benchmarks/kernels/benchmark_rmsnorm.py +++ /dev/null @@ -1,255 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import itertools - -import torch -from flashinfer.norm import fused_add_rmsnorm, rmsnorm -from torch import nn - -from vllm import _custom_ops as vllm_ops -from vllm.triton_utils import triton - - -class HuggingFaceRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - orig_dtype = x.dtype - x = x.to(torch.float32) - if residual is not None: - x = x + residual.to(torch.float32) - residual = x.to(orig_dtype) - - variance = x.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + self.variance_epsilon) - x = x.to(orig_dtype) * self.weight - if residual is None: - return x - else: - return x, residual - - -def rmsnorm_naive( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps) - naive_norm.weight = nn.Parameter(weight) - naive_norm = naive_norm.to(x.device) - - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - output = naive_norm(x, residual) - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def rmsnorm_flashinfer( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - if residual is not None: - fused_add_rmsnorm(x, residual, weight, eps) - output = (x, residual) - else: - output = rmsnorm(x, weight, eps) - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def rmsnorm_vllm( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - if residual is not None: - vllm_ops.fused_add_rms_norm(x, residual, weight, eps) - output = (x, residual) - else: - out = torch.empty_like(x) - vllm_ops.rms_norm(out, x, weight, eps) - output = out - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True): - dtype = torch.bfloat16 - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") - weight = torch.ones(hidden_size, dtype=dtype, device="cuda") - residual = torch.randn_like(x) if use_residual else None - - output_naive = rmsnorm_naive( - x.clone(), weight, residual.clone() if residual is not None else None - ) - output_flashinfer = rmsnorm_flashinfer( - x.clone(), weight, residual.clone() if residual is not None else None - ) - output_vllm = rmsnorm_vllm( - x.clone(), weight, residual.clone() if residual is not None else None - ) - - if use_residual: - output_naive = output_naive[0] - output_flashinfer = output_flashinfer[0] - output_vllm = output_vllm[0] - - print(f"Naive output={output_naive}") - print(f"FlashInfer output={output_flashinfer}") - print(f"vLLM output={output_vllm}") - - if torch.allclose( - output_naive, output_flashinfer, atol=1e-2, rtol=1e-2 - ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2): - print("✅ All implementations match") - else: - print("❌ Implementations differ") - - -batch_size_range = [2**i for i in range(0, 7, 2)] -seq_length_range = [2**i for i in range(6, 11, 1)] -head_num_range = [32, 48] -configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range)) - - -def get_benchmark(use_residual): - @triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["head_num", "batch_size", "seq_len"], - x_vals=[list(_) for _ in configs], - line_arg="provider", - line_vals=["huggingface", "flashinfer", "vllm"], - line_names=["HuggingFace", "FlashInfer", "vLLM"], - styles=[("blue", "-"), ("green", "-"), ("red", "-")], - ylabel="us", - plot_name=f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual", - args={}, - ) - ) - def benchmark(head_num, batch_size, seq_len, provider): - dtype = torch.bfloat16 - hidden_size = head_num * 128 # assuming head_dim = 128 - - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") - weight = torch.ones(hidden_size, dtype=dtype, device="cuda") - residual = torch.randn_like(x) if use_residual else None - - quantiles = [0.5, 0.2, 0.8] - - if provider == "huggingface": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_naive( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - elif provider == "flashinfer": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_flashinfer( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - else: - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_vllm( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - - return 1000 * ms, 1000 * max_ms, 1000 * min_ms - - return benchmark - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument( - "--batch-size", - type=int, - default=4, - help="Batch size", - ) - parser.add_argument( - "--seq-len", - type=int, - default=128, - help="Sequence length", - ) - parser.add_argument( - "--hidden-size", - type=int, - default=4096, - help="Hidden size (2nd dimension) of the sequence", - ) - parser.add_argument( - "--use-residual", action="store_true", help="Whether to use residual connection" - ) - parser.add_argument( - "--save-path", - type=str, - default="./configs/rmsnorm/", - help="Path to save rmsnorm benchmark results", - ) - - args = parser.parse_args() - - # Run correctness test - calculate_diff( - batch_size=args.batch_size, - seq_len=args.seq_len, - hidden_size=args.hidden_size, - use_residual=args.use_residual, - ) - - # Get the benchmark function with proper use_residual setting - benchmark = get_benchmark(args.use_residual) - # Run performance benchmark - benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py index bac2920103c2..40f5f8947621 100644 --- a/benchmarks/kernels/ir/shapes.py +++ b/benchmarks/kernels/ir/shapes.py @@ -18,10 +18,11 @@ # Each entry maps an op name to a list of kwarg dicts that will be passed # to that op's registered input generator via op.generate_inputs(**kwargs). SHAPE_CONFIGS: dict[str, list[dict]] = { - "rms_norm": [ + op_name: [ {"num_tokens": n, "hidden_size": d, "dtype": dtype} for dtype in [torch.float16, torch.bfloat16, torch.float32] for d in COMMON_HIDDEN_SIZES for n in NUM_TOKENS - ], + ] + for op_name in ("rms_norm", "fused_add_rms_norm") } diff --git a/benchmarks/overheads/benchmark_hashing.py b/benchmarks/overheads/benchmark_hashing.py deleted file mode 100644 index 178599952d5c..000000000000 --- a/benchmarks/overheads/benchmark_hashing.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import cProfile -import pstats - -from vllm import LLM, SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser - -# A very long prompt, total number of tokens is about 15k. -LONG_PROMPT = ["You are an expert in large language models, aren't you?"] * 1000 -LONG_PROMPT = " ".join(LONG_PROMPT) - - -def main(args): - llm = LLM( - model=args.model, - enforce_eager=True, - enable_prefix_caching=True, - tensor_parallel_size=args.tensor_parallel_size, - ) - - sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len) - profiler = cProfile.Profile() - - print("------warm up------") - for i in range(3): - output = llm.generate(LONG_PROMPT, sampling_params) - print(output[0].outputs[0].text) - - print("------start generating------") - for i in range(3): - profiler.runctx( - "llm.generate(LONG_PROMPT, sampling_params)", globals(), locals() - ) - - # analyze the runtime of hashing function - stats = pstats.Stats(profiler) - stats.sort_stats("cumulative") - total_time = 0 - total_calls = 0 - for func in stats.stats: - if "hash_of_block" in func[2]: - total_time = stats.stats[func][3] - total_calls = stats.stats[func][0] - percentage = (total_time / stats.total_tt) * 100 - print( - f"Hashing took {total_time:.2f} seconds,{percentage:.2f}% of the total runtime." - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser( - description="Benchmark the performance of hashing function in" - "automatic prefix caching." - ) - parser.add_argument("--model", type=str, default="lmsys/longchat-7b-16k") - parser.add_argument("--tensor-parallel-size", "-tp", type=int, default=1) - parser.add_argument("--output-len", type=int, default=10) - parser.add_argument( - "--enable-prefix-caching", action="store_true", help="enable prefix caching" - ) - args = parser.parse_args() - main(args) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 3be7bb243be5..f65b9ec921b7 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -60,8 +60,6 @@ class ForbiddenImport: "benchmarks/kernels/benchmark_lora.py", "benchmarks/kernels/benchmark_machete.py", "benchmarks/fused_kernels/layernorm_rms_benchmarks.py", - "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py", - "benchmarks/cutlass_benchmarks/sparse_benchmarks.py", # cloudpickle "vllm/v1/executor/multiproc_executor.py", "vllm/v1/executor/ray_executor.py", From 651a88c09ac6cbb13c5dd36f2fdbe61a2c68eeef Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:57:24 +0100 Subject: [PATCH 337/774] [CI] Ignore ruff D209, rejoin the docstrings it split, and silence incompatible-rule warnings (#57212) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- .../benchmark_hidden_state_extraction.py | 3 +- .../merge_attn_states_benchmarks.py | 3 +- benchmarks/kernels/benchmark_moe_defaults.py | 3 +- .../benchmark_vocab_parallel_embedding.py | 3 +- docs/mkdocs/gen_files/generate_examples.py | 6 +- .../features/logits_processor/custom_req.py | 6 +- .../logits_processor/custom_req_init.py | 6 +- examples/rl/rlhf_sharded_rdt_small_ep.py | 3 +- pyproject.toml | 4 +- setup.py | 3 +- tests/basic_correctness/test_mem.py | 3 +- tests/benchmarks/test_bfcl_dataset.py | 18 +- ...est_custom_dataset_chat_template_kwargs.py | 3 +- tests/benchmarks/test_sampling_params.py | 9 +- tests/benchmarks/test_skip_tokenizer_init.py | 3 +- tests/compile/passes/test_fusion.py | 3 +- tests/compile/passes/test_split_coalescing.py | 3 +- tests/compile/test_codegen.py | 30 +-- tests/compile/test_compile_ranges.py | 6 +- tests/compile/test_config.py | 9 +- ...test_inductor_fallback_allow_list_patch.py | 3 +- tests/config/test_model_arch_config.py | 6 +- tests/config/test_multimodal_config.py | 15 +- .../test_speculative_draft_hf_overrides.py | 12 +- ...eculative_draft_max_position_embeddings.py | 3 +- tests/distributed/conftest.py | 3 +- .../test_eplb_quant_scale_consistency.py | 12 +- tests/distributed/test_mnnvl_alltoall.py | 6 +- tests/distributed/test_mq_connect_ip.py | 9 +- tests/distributed/test_pynccl.py | 3 +- tests/distributed/test_ray_v2_executor_e2e.py | 3 +- tests/distributed/test_sharded_rdt_plan.py | 183 ++++++------------ .../distributed/test_sharded_rdt_producer.py | 117 ++++------- tests/distributed/test_sharded_rdt_trainer.py | 36 ++-- tests/distributed/test_shm_broadcast.py | 9 +- tests/distributed/test_utils.py | 3 +- tests/distributed/test_weight_transfer.py | 66 +++---- .../test_weight_transfer_nccl_uid.py | 9 +- tests/engine/test_arg_utils.py | 15 +- .../test_anthropic_messages_conversion.py | 30 +-- .../cohere/test_serving_streaming.py | 6 +- .../test_api_server_process_manager.py | 12 +- tests/entrypoints/launchers/test_cli_args.py | 6 +- ...chat_completion_with_mixed_audio_embeds.py | 3 +- ...chat_completion_with_mixed_image_embeds.py | 3 +- .../openai/chat_completion/test_vision.py | 6 +- .../openai/chat_completion/test_chat_error.py | 6 +- .../test_chat_logit_bias_validation.py | 9 +- .../test_extra_content_fields.py | 6 +- .../chat_completion/test_logprob_token_ids.py | 3 +- .../chat_completion/test_serving_chat.py | 12 +- .../test_thinking_token_budget.py | 3 +- .../completion/test_completion_error.py | 9 +- .../parser/test_harmony_render_parity.py | 12 +- .../openai/parser/test_harmony_utils.py | 3 +- .../responses/test_function_call_parsing.py | 9 +- .../test_response_input_to_harmony.py | 12 +- .../openai/responses/test_responses_utils.py | 9 +- .../responses/test_serving_responses.py | 18 +- .../openai/responses/test_simple.py | 3 +- .../openai/responses/test_streaming_events.py | 6 +- .../openai/test_render_token_offsets.py | 9 +- .../pooling/embed/test_cohere_online.py | 6 +- .../embed/test_cohere_openai_parity.py | 3 +- .../pooling/embed/test_io_processor.py | 3 +- .../pooling/embed/test_online_long_text.py | 9 +- .../pooling/embed/test_protocol.py | 3 +- .../scale_out/derender/test_derender.py | 3 +- .../derender/test_derender_parity.py | 9 +- .../scale_out/render/test_render.py | 12 +- .../token_in_token_out/test_protocol.py | 9 +- .../test_error_sanitization.py | 6 +- .../test_validation_exception_handler.py | 15 +- .../serve/lora/test_lora_adapters.py | 3 +- .../test_authentication_middleware.py | 9 +- .../test_sagemaker_handler_overrides.py | 18 +- .../tokenize/test_tokenize_then_chat_vlm.py | 3 +- .../entrypoints/serve/utils/test_api_utils.py | 12 +- .../serve/utils/test_request_logger.py | 3 +- .../serve/utils/test_sse_keep_alive.py | 24 +-- .../realtime/test_realtime_validation.py | 3 +- tests/entrypoints/unit_tests/test_context.py | 21 +- .../test_weight_transfer_llm.py | 3 +- .../fusion/test_quant_activation_contract.py | 3 +- tests/kernels/attention/test_amx_mla.py | 9 +- .../attention/test_attention_selector.py | 6 +- .../attention/test_cutlass_mla_decode.py | 3 +- tests/kernels/attention/test_flashinfer.py | 3 +- .../test_flashinfer_trtllm_attention.py | 3 +- .../test_kimi_k3_mla_fused_epilogue.py | 3 +- .../attention/test_merge_attn_states.py | 3 +- tests/kernels/attention/test_minimax_m3.py | 42 ++-- ...test_mla_cross_layer_kernel_equivalence.py | 24 +-- tests/kernels/attention/test_rocm_aiter_fa.py | 21 +- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 6 +- .../attention/test_triton_decode_attention.py | 6 +- tests/kernels/core/test_fused_embed_norm.py | 6 +- tests/kernels/core/test_fused_q_kv_rmsnorm.py | 6 +- .../kernels/core/test_fused_rms_norm_gated.py | 3 +- tests/kernels/core/test_rocm_aiter_ops.py | 3 +- .../core/test_vocab_parallel_embedding.py | 3 +- tests/kernels/helion/test_autotune.py | 3 +- tests/kernels/helion/test_register.py | 9 +- tests/kernels/ir/test_layernorm.py | 3 +- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 9 +- tests/kernels/mamba/test_mamba_ssm_configs.py | 9 +- ...ayssm_prefill_decode_equivalence_mamba2.py | 3 +- .../test_replayssm_standard_decode_mamba2.py | 6 +- .../moe/modular_kernel_tools/common.py | 3 +- tests/kernels/moe/test_b12x.py | 3 +- tests/kernels/moe/test_batched_moe.py | 9 +- tests/kernels/moe/test_block_int8.py | 6 +- tests/kernels/moe/test_cpu_fused_moe.py | 6 +- tests/kernels/moe/test_flashinfer.py | 3 +- .../moe/test_gpt_oss_triton_kernels.py | 3 +- .../moe/test_marlin_vs_trtllm_mxint4.py | 3 +- .../moe/test_modular_kernel_combinations.py | 3 +- tests/kernels/moe/test_moe_kernel_oracle.py | 3 +- tests/kernels/moe/test_moe_layer.py | 3 +- .../moe/test_mxfp8_aiter_backend_selection.py | 6 +- tests/kernels/moe/test_rocm_aiter_moe.py | 33 ++-- .../test_routed_experts_capture_monolithic.py | 24 +-- tests/kernels/moe/test_routing.py | 3 +- tests/kernels/moe/test_routing_simulator.py | 6 +- ..._mul_per_token_group_quant_fp8_colmajor.py | 3 +- tests/kernels/moe/test_triton_moe_ptpc_fp8.py | 6 +- .../moe/test_unquantized_backend_selection.py | 6 +- tests/kernels/quant_utils.py | 3 +- .../test_aiter_hipb_mm_linear_kernel.py | 3 +- .../kernels/quantization/test_int8_kernel.py | 6 +- .../quantization/test_marlin_tile_padding.py | 15 +- .../test_mxfp4_kernel_selection.py | 18 +- .../quantization/test_mxfp4_triton_ep.py | 3 +- .../test_mxfp6_kernel_selection.py | 9 +- .../test_per_token_group_quant.py | 9 +- .../quantization/test_quantized_embedding.py | 3 +- .../test_rocm_aiter_grouped_quant.py | 6 +- tests/kernels/quantization/test_rocm_mxfp4.py | 3 +- tests/kernels/test_compressor_kv_cache.py | 9 +- tests/kernels/test_cp_gather_fp8.py | 3 +- tests/kernels/test_deepseek_v4_cpu_kernels.py | 57 ++---- tests/kernels/test_engram.py | 9 +- tests/kernels/test_fla_layernorm_guard.py | 3 +- tests/kernels/test_fp32_router_gemm.py | 3 +- .../test_fused_deepseek_v32_norm_rope.py | 3 +- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 6 +- .../test_fused_indexer_q_rope_quant.py | 3 +- .../kernels/test_fused_inv_rope_fp8_quant.py | 9 +- tests/kernels/test_mhc_kernels.py | 6 +- tests/kernels/test_minimax_m3_amd_ops.py | 12 +- tests/kernels/test_shuffle_rows.py | 6 +- tests/kernels/test_top_k_per_row.py | 6 +- tests/lora/test_default_mm_loras.py | 3 +- tests/lora/test_fused_moe_lora_kernel.py | 6 +- tests/lora/test_moe_lora_align_sum.py | 6 +- tests/lora/test_moe_lora_ep_load.py | 6 +- tests/lora/test_punica_ops.py | 6 +- tests/lora/test_punica_ops_fp8.py | 6 +- .../layers/test_fused_shared_expert.py | 6 +- .../layers/test_mla_short_prefill_indexer.py | 36 ++-- .../model_loader/test_ep_weight_filter.py | 9 +- .../model_loader/test_reload.py | 9 +- tests/model_executor/test_deep_gemm_warmup.py | 6 +- .../test_routed_experts_capture.py | 3 +- tests/model_executor/test_weight_utils.py | 24 +-- tests/models/inkling/test_mm_towers.py | 3 +- .../kimi_k3/test_amd_latent_moe_runner.py | 3 +- .../kimi_k3/test_aux_attn_res_stream.py | 6 +- tests/models/kimi_k3/test_kda_metadata.py | 3 +- .../models/language/generation/test_gemma.py | 6 +- .../language/generation/test_mistral.py | 3 +- .../pooling/test_splade_sparse_pooler.py | 3 +- .../generation/test_moondream3_moe.py | 3 +- .../generation/test_transformers_audio.py | 3 +- .../generation/vlm_utils/builders.py | 3 +- .../processing/test_deepseek_ocr.py | 6 +- .../processing/test_llava_next_video.py | 6 +- .../multimodal/processing/test_molmo2.py | 3 +- .../multimodal/processing/test_qwen3_vl.py | 12 +- .../processing/test_transformers_audio.py | 3 +- .../processing/test_transformers_image.py | 24 +-- .../processing/transformers_backend.py | 6 +- tests/models/multimodal/test_cohere_asr.py | 3 +- tests/models/multimodal/test_mapping.py | 3 +- .../multimodal/test_nano_nemotron_vl.py | 9 +- tests/models/qwen4_exp/test_config.py | 3 +- tests/models/test_deepseek_v4_fi_moe_ep.py | 15 +- tests/models/test_deepseek_v4_mega_moe.py | 6 +- tests/models/test_dspark_mla.py | 3 +- tests/models/test_qwen3_5_mtp_config.py | 9 +- tests/models/test_utils.py | 3 +- tests/models/test_vision.py | 3 +- .../models/transformers/fusers/test_linear.py | 69 +++---- tests/models/transformers/fusers/test_mla.py | 15 +- tests/models/transformers/fusers/test_moe.py | 18 +- .../transformers/fusers/test_rms_norm.py | 27 +-- .../transformers/test_layer_registry.py | 6 +- tests/multimodal/media/test_audio.py | 18 +- tests/multimodal/media/test_connector.py | 3 +- tests/multimodal/media/test_video.py | 15 +- tests/multimodal/test_hasher.py | 3 +- tests/multimodal/test_image.py | 3 +- tests/multimodal/test_parse.py | 6 +- tests/multimodal/test_registry.py | 3 +- tests/multimodal/test_utils.py | 3 +- tests/multimodal/test_video.py | 30 +-- tests/parser/cohere/test_reasoning.py | 3 +- tests/parser/cohere/utils.py | 9 +- tests/parser/engine/test_deepseek_v32.py | 3 +- tests/parser/engine/test_deepseek_v4.py | 18 +- tests/parser/engine/test_delegating_replay.py | 3 +- tests/parser/engine/test_engine.py | 57 ++---- .../engine/test_gemma4_streaming_reasoning.py | 27 +-- tests/parser/engine/test_glm47_moe.py | 6 +- tests/parser/engine/test_inkling.py | 30 +-- tests/parser/engine/test_nemotron_v3.py | 18 +- tests/parser/engine/test_parser_engine.py | 69 +++---- tests/parser/engine/test_qwen3.py | 15 +- tests/parser/engine/test_qwen3_reasoning.py | 18 +- tests/parser/engine/test_replay.py | 12 +- tests/parser/engine/test_seed_oss.py | 6 +- .../engine/test_ufffd_reasoning_transition.py | 3 +- tests/parser/engine/trace_builder.py | 3 +- tests/parser/mistral/test_tool_calls.py | 15 +- tests/parser/test_streaming.py | 36 ++-- tests/plugins_tests/test_endpoint_plugins.py | 12 +- .../test_io_processor_plugins.py | 3 +- tests/quantization/test_auto_round.py | 15 +- tests/quantization/test_fp8.py | 3 +- tests/quantization/test_gfx950_moe.py | 3 +- .../test_humming_mxfp4_block_fp8.py | 6 +- tests/quantization/test_int8_moe_oracle.py | 3 +- tests/quantization/test_modelopt.py | 12 +- tests/quantization/test_online_mxfp4.py | 3 +- tests/quantization/test_per_token_kv_cache.py | 3 +- .../test_kimi_k3_reasoning_parser.py | 12 +- .../test_chat_utils_prompt_embeds.py | 9 +- tests/renderers/test_gemma4_chat_template.py | 27 +-- tests/renderers/test_hf.py | 9 +- .../test_process_multi_modal_uuids.py | 3 +- .../test_prompt_embeds_payload_validation.py | 3 +- .../test_sparse_tensor_concurrent_race.py | 3 +- .../test_sparse_tensor_validation.py | 21 +- tests/renderers/test_token_offsets.py | 15 +- tests/test_cmake_utils.py | 3 +- tests/test_config.py | 30 +-- tests/test_logger.py | 24 +-- tests/test_ray_env.py | 6 +- tests/tokenizers_/test_deepseek_v4.py | 3 +- tests/tokenizers_/test_hf.py | 3 +- .../tool_parsers/test_apertus_tool_parser.py | 12 +- .../test_deepseekv32_tool_parser.py | 30 +-- tests/tool_parsers/test_gemma4_tool_parser.py | 3 +- .../test_gigachat3_tool_parser.py | 3 +- .../test_kimi_k3_named_tool_choice.py | 3 +- tests/tool_parsers/test_lfm2_tool_parser.py | 27 +-- tests/tool_use/test_muse_glimmer.py | 3 +- tests/tool_use/test_tool_choice_required.py | 3 +- .../processors/test_glm5next.py | 6 +- tests/transformers_utils/test_config.py | 6 +- .../test_hf_overrides_model_type.py | 3 +- tests/utils.py | 6 +- tests/utils_/test_gpu_sync_debug.py | 15 +- tests/utils_/test_network_utils.py | 15 +- tests/utils_/test_numa_utils.py | 18 +- tests/v1/attention/test_attention_backends.py | 3 +- .../test_attention_backends_selection.py | 3 +- .../v1/attention/test_attention_splitting.py | 3 +- tests/v1/attention/test_dcp_a2a_pack_mask.py | 3 +- .../test_dspark_noncausal_sparse_mla.py | 3 +- .../test_flashinfer_mla_sparse_sm90.py | 3 +- .../v1/attention/test_gdn_metadata_builder.py | 3 +- .../v1/attention/test_indexer_dcp_localize.py | 18 +- .../attention/test_kpool_tail_slot_mapping.py | 21 +- .../test_mamba_update_block_table.py | 21 +- tests/v1/attention/test_mla_backends.py | 3 +- .../test_mla_prefill_quant_output.py | 9 +- .../v1/attention/test_mla_prefill_selector.py | 3 +- .../test_replayssm_metadata_builder.py | 3 +- .../test_rocm_attention_backends_selection.py | 3 +- .../test_sparse_indexer_decode_seq_lens.py | 12 +- .../v1/attention/test_sparse_mla_backends.py | 27 +-- .../test_trtllm_attention_integration.py | 9 +- .../test_partial_prefix_cache_hits.py | 51 ++--- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_contiguous_kv_packing.py | 18 +- tests/v1/core/test_encoder_cache_manager.py | 9 +- tests/v1/core/test_kv_cache_utils.py | 30 +-- tests/v1/core/test_output.py | 9 +- tests/v1/core/test_prefix_caching.py | 75 +++---- tests/v1/core/test_scheduler.py | 84 +++----- .../v1/core/test_swa_inflight_window_free.py | 12 +- .../v1/cudagraph/test_breakable_cudagraph.py | 9 +- .../test_rms_norm_batch_invariant.py | 3 +- tests/v1/distributed/test_async_llm_dp.py | 12 +- tests/v1/distributed/test_external_lb_dp.py | 3 +- tests/v1/distributed/test_hybrid_lb_dp.py | 3 +- tests/v1/distributed/test_internal_lb_dp.py | 9 +- tests/v1/distributed/test_pp_dp_v2.py | 12 +- tests/v1/e2e/general/test_async_scheduling.py | 6 +- .../test_attention_backend_per_kind.py | 3 +- tests/v1/e2e/general/test_context_length.py | 6 +- tests/v1/e2e/general/test_min_tokens.py | 3 +- tests/v1/e2e/general/test_sharded_sampling.py | 6 +- .../mtp_other/test_synthetic.py | 3 +- .../draft_model/test_draft_model.py | 21 +- .../e2e/spec_decode/test_sharded_sampling.py | 3 +- tests/v1/e2e/test_hybrid_chunked_prefill.py | 3 +- .../cpu/scheduler/test_embedding_cache.py | 3 +- .../unit/cpu/scheduler/test_scheduler.py | 36 ++-- .../unit/cpu/test_ec_shared_region.py | 12 +- .../unit/cpu/worker/test_worker.py | 18 +- .../unit/test_epd_proxy_round_robin.py | 3 +- tests/v1/ec_connector/unit/test_metadata.py | 3 +- .../unit/test_scheduler_nixl_producer.py | 3 +- tests/v1/ec_connector/unit/test_session.py | 6 +- .../engine/test_core_engine_actor_manager.py | 3 +- tests/v1/engine/test_engine_core.py | 12 +- tests/v1/engine/test_engine_core_client.py | 9 +- tests/v1/engine/test_logprobs_processor.py | 3 +- tests/v1/engine/test_output_processor.py | 3 +- .../test_canonical_mapping.py | 3 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/offloading_connector/test_metrics.py | 6 +- .../offloading_connector/test_scheduler.py | 48 ++--- .../unit/offloading_connector/test_worker.py | 3 +- .../unit/test_bidirectional_kv_transfer.py | 57 ++---- tests/v1/kv_connector/unit/test_config.py | 18 +- .../unit/test_flexkv_connector.py | 9 +- .../unit/test_handshake_pp_aggregation.py | 9 +- .../v1/kv_connector/unit/test_hf3fs_client.py | 6 +- .../unit/test_lmcache_connector.py | 3 +- .../kv_connector/unit/test_mooncake_stats.py | 3 +- .../unit/test_mooncake_store_connector.py | 9 +- .../unit/test_mooncake_store_coordinator.py | 54 ++---- .../unit/test_mooncake_store_hma_e2e.py | 21 +- .../unit/test_mooncake_store_scheduler.py | 6 +- .../unit/test_mooncake_store_worker.py | 60 ++---- .../unit/test_moriio_connector.py | 21 +- .../unit/test_moriio_routing_fairness.py | 6 +- .../kv_connector/unit/test_multi_connector.py | 3 +- .../kv_connector/unit/test_nixl_connector.py | 6 +- .../unit/test_nixl_connector_hma.py | 51 ++--- .../unit/test_nixl_desc_geometry.py | 30 +-- .../unit/test_nixl_push_connector.py | 105 ++++------ .../unit/test_nixl_rocm_gpu_mem_diag.py | 3 +- .../unit/test_nixl_simple_cpu_offload.py | 9 +- .../unit/test_offloading_connector.py | 6 +- .../unit/test_remote_prefill_lifecycle.py | 3 +- .../unit/test_simple_cpu_offload_connector.py | 3 +- tests/v1/kv_connector/unit/test_tp_mapping.py | 6 +- .../kv_offload/cpu/policies/test_factory.py | 15 +- .../kv_offload/cpu/test_canonical_layout.py | 9 +- .../cpu/test_shared_offload_region.py | 42 ++-- .../v1/kv_offload/tiering/p2p/test_manager.py | 93 +++------ .../kv_offload/tiering/p2p/test_sessions.py | 120 ++++-------- .../kv_offload/tiering/test_async_lookup.py | 6 +- tests/v1/kv_offload/tiering/test_fs_tier.py | 24 +-- tests/v1/kv_offload/tiering/test_obj_tier.py | 6 +- .../tiering/test_tiering_offloading.py | 33 ++-- .../v1/logits_processors/test_correctness.py | 3 +- tests/v1/logits_processors/utils.py | 6 +- tests/v1/metrics/test_perf_metrics.py | 3 +- tests/v1/metrics/test_ray_metrics.py | 21 +- .../sample/test_batched_count_greater_than.py | 3 +- tests/v1/sample/test_logprobs.py | 3 +- tests/v1/sample/test_rejection_sampler.py | 3 +- tests/v1/sample/test_topk_topp_sampler.py | 3 +- tests/v1/sample/utils.py | 3 +- tests/v1/simple_kv_offload/test_kv_events.py | 9 +- tests/v1/simple_kv_offload/test_scheduler.py | 18 +- .../test_backup_token_async_spec.py | 3 +- .../test_eagle_draft_attn_metadata.py | 9 +- .../v1/spec_decode/test_eagle_step_kernel.py | 3 +- .../v1/spec_decode/test_llm_base_proposer.py | 6 +- tests/v1/spec_decode/test_max_len.py | 3 +- tests/v1/structured_output/test_validation.py | 12 +- tests/v1/test_outputs.py | 12 +- tests/v1/test_serial_utils.py | 3 +- .../test_dsv4_packed_zeroer_geometry.py | 3 +- tests/v1/worker/test_encoder_runner.py | 15 +- tests/v1/worker/test_gpu_bad_words.py | 9 +- tests/v1/worker/test_gpu_batch_shard.py | 27 +-- tests/v1/worker/test_gpu_block_table.py | 6 +- tests/v1/worker/test_gpu_logit_bias.py | 3 +- tests/v1/worker/test_gpu_model_runner.py | 6 +- .../worker/test_gpu_model_runner_mm_gather.py | 9 +- tests/v1/worker/test_gpu_model_runner_v2.py | 12 +- ...gpu_model_runner_v2_cudagraph_profiling.py | 3 +- tests/v1/worker/test_gpu_profiler.py | 6 +- .../worker/test_gpu_rejection_sampler_i64.py | 3 +- tests/v1/worker/test_gpu_thinking_budget.py | 3 +- tests/v1/worker/test_gpu_ubatch_slicing.py | 3 +- tests/v1/worker/test_gpu_worker.py | 15 +- tests/v1/worker/test_kv_block_zeroer.py | 3 +- tests/v1/worker/test_mamba_utils.py | 12 +- tests/v1/worker/test_mixed_warmup_gate.py | 3 +- tests/v1/worker/test_prompt_embeds_state.py | 12 +- tests/v1/worker/test_workspace.py | 3 +- tools/report_build_time_ninja.py | 9 +- vllm/_aiter_ops.py | 3 +- .../datasets/create_txt_slices_dataset.py | 3 +- vllm/benchmarks/datasets/datasets.py | 3 +- vllm/benchmarks/lib/endpoint_request_func.py | 6 +- vllm/compilation/backends.py | 3 +- vllm/compilation/compiler_interface.py | 3 +- .../passes/fusion/allreduce_rms_fusion.py | 3 +- .../passes/utility/split_coalescing.py | 3 +- vllm/config/compilation.py | 6 +- vllm/config/model.py | 3 +- vllm/config/multimodal.py | 3 +- vllm/config/observability.py | 3 +- vllm/config/reasoning.py | 6 +- vllm/config/speculative.py | 6 +- vllm/config/utils.py | 6 +- vllm/config/vllm.py | 6 +- vllm/device_allocator/sleep_mode_backend.py | 15 +- .../device_communicators/all_reduce_utils.py | 3 +- .../device_communicators/shm_broadcast.py | 6 +- .../cpu/scheduler/embedding_cache.py | 3 +- vllm/distributed/eplb/eplb_communicator.py | 6 +- .../kv_transfer/kv_connector/utils.py | 3 +- .../v1/example_hidden_states_connector.py | 3 +- .../kv_connector/v1/lmcache_mp_connector.py | 6 +- .../v1/mooncake/mooncake_connector.py | 12 +- .../v1/mooncake/store/coordinator.py | 3 +- .../kv_connector/v1/mooncake/store/worker.py | 3 +- .../v1/moriio/moriio_connector.py | 3 +- .../kv_connector/v1/nixl/base_scheduler.py | 6 +- .../kv_connector/v1/nixl/base_worker.py | 6 +- .../kv_connector/v1/nixl/push_worker.py | 6 +- .../v1/offloading/canonical_mapping.py | 12 +- .../kv_connector/v1/offloading/events.py | 9 +- .../kv_connector/v1/offloading/scheduler.py | 6 +- vllm/distributed/utils.py | 3 +- vllm/distributed/weight_transfer/factory.py | 3 +- .../distributed/weight_transfer/ipc_engine.py | 6 +- .../weight_transfer/nccl_common.py | 3 +- .../weight_transfer/nccl_engine.py | 6 +- .../weight_transfer/sharded_rdt_common.py | 9 +- .../weight_transfer/sharded_rdt_engine.py | 48 ++--- .../weight_transfer/sharded_rdt_fake.py | 6 +- .../weight_transfer/sharded_rdt_trainer.py | 42 ++-- .../weight_transfer/sparse_nccl_engine.py | 6 +- vllm/entrypoints/cli/main.py | 3 +- vllm/entrypoints/generate/base/serving.py | 3 +- vllm/entrypoints/openai/models/serving.py | 6 +- vllm/entrypoints/openai/responses/context.py | 3 +- vllm/entrypoints/openai/responses/harmony.py | 3 +- vllm/entrypoints/openai/responses/protocol.py | 3 +- .../pooling/pooling/io_processor.py | 3 +- .../scale_out/token_in_token_out/mm_serde.py | 3 +- .../exception_handling/handlers/validation.py | 3 +- vllm/entrypoints/serve/lora/api_router.py | 6 +- .../entrypoints/serve/sagemaker/api_router.py | 6 +- vllm/entrypoints/serve/utils/fingerprint.py | 6 +- .../speech_to_text/base/serving.py | 6 +- vllm/env_override.py | 3 +- vllm/envs.py | 3 +- vllm/forward_context.py | 3 +- vllm/logger.py | 3 +- vllm/lora/layers/column_parallel_linear.py | 3 +- vllm/lora/lora_model.py | 3 +- vllm/lora/worker_manager.py | 9 +- vllm/model_executor/hw_agnostic/custom_op.py | 3 +- .../hw_agnostic/layers/layernorm.py | 3 +- .../model_executor/kernels/linear/__init__.py | 12 +- .../kernels/linear/nvfp4/flashinfer.py | 3 +- .../scaled_mm/BlockScaledMMLinearKernel.py | 3 +- .../kernels/linear/scaled_mm/cutlass.py | 6 +- vllm/model_executor/layers/activation.py | 3 +- .../layers/attention/mla_attention.py | 3 +- .../layers/fused_allreduce_gemma_rms_norm.py | 3 +- .../model_executor/layers/fused_embed_norm.py | 3 +- .../layers/fused_moe/activation.py | 3 +- .../fused_moe/experts/aiter_mxfp8_moe.py | 3 +- .../layers/fused_moe/experts/cpu_moe.py | 9 +- .../fused_moe/experts/mxfp8_native_moe.py | 3 +- .../layers/fused_moe/oracle/base.py | 15 +- .../layers/fused_moe/oracle/fp8.py | 3 +- .../layers/fused_moe/oracle/int8.py | 3 +- .../layers/fused_moe/oracle/int_wna16.py | 6 +- .../layers/fused_moe/oracle/mxfp4.py | 9 +- .../fused_moe/router/zero_expert_router.py | 3 +- vllm/model_executor/layers/indexer_topk.py | 9 +- vllm/model_executor/layers/linear.py | 3 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 3 +- .../layers/mamba/mamba_utils.py | 3 +- .../layers/mamba/ops/mamba_ssm.py | 6 +- .../minimax_rms_norm/lamport_workspace.py | 3 +- .../layers/quantization/base_config.py | 12 +- .../compressed_tensors/compressed_tensors.py | 6 +- .../compressed_tensors_embedding.py | 3 +- .../schemes/compressed_tensors_wNa8o8.py | 3 +- .../layers/quantization/inc/config_parser.py | 3 +- .../quantization/inc/schemes/inc_scheme.py | 3 +- .../layers/quantization/modelopt.py | 18 +- .../layers/quantization/mxfp4.py | 3 +- .../layers/quantization/online/base.py | 3 +- .../layers/quantization/online/fp8.py | 18 +- .../layers/quantization/utils/fp8_utils.py | 6 +- .../quantization/utils/humming_utils.py | 6 +- .../layers/quantization/utils/marlin_utils.py | 9 +- .../quantization/utils/marlin_utils_fp4.py | 9 +- .../quantization/utils/marlin_utils_fp8.py | 3 +- .../rotary_embedding/ntk_scaling_rope.py | 3 +- .../layers/sparse_attn_indexer.py | 6 +- .../layers/sparse_mqa_indexer.py | 6 +- .../layers/vocab_parallel_embedding.py | 3 +- .../model_loader/base_loader.py | 3 +- .../model_loader/default_loader.py | 3 +- .../model_loader/ep_weight_filter.py | 6 +- .../model_loader/reload/layerwise.py | 6 +- .../model_loader/runai_streamer_loader.py | 3 +- .../model_executor/model_loader/tensorizer.py | 3 +- .../model_loader/tensorizer_loader.py | 3 +- .../model_loader/weight_cache/ipc_loader.py | 3 +- .../model_loader/weight_tying.py | 3 +- .../model_loader/weight_utils.py | 9 +- vllm/model_executor/models/arcee.py | 15 +- vllm/model_executor/models/blip.py | 3 +- vllm/model_executor/models/cohere2_moe.py | 3 +- vllm/model_executor/models/cohere_asr.py | 3 +- vllm/model_executor/models/deepencoder.py | 3 +- vllm/model_executor/models/deepseek_ocr.py | 3 +- vllm/model_executor/models/diffusion_gemma.py | 3 +- vllm/model_executor/models/gemma3n.py | 3 +- vllm/model_executor/models/gemma3n_mm.py | 3 +- vllm/model_executor/models/glm4_1v.py | 3 +- vllm/model_executor/models/glm4_moe.py | 3 +- vllm/model_executor/models/glm4_moe_mtp.py | 3 +- vllm/model_executor/models/granite_speech.py | 3 +- vllm/model_executor/models/interfaces.py | 15 +- vllm/model_executor/models/keye.py | 3 +- vllm/model_executor/models/lfm2_siglip2.py | 3 +- vllm/model_executor/models/medusa.py | 3 +- vllm/model_executor/models/minicpmv4_6.py | 3 +- vllm/model_executor/models/mllama4.py | 3 +- vllm/model_executor/models/muse_glimmer.py | 9 +- .../model_executor/models/nano_nemotron_vl.py | 3 +- vllm/model_executor/models/phi4mm_utils.py | 3 +- vllm/model_executor/models/qwen3_dflash.py | 12 +- vllm/model_executor/models/qwen3_vl.py | 3 +- vllm/model_executor/models/radio.py | 6 +- vllm/model_executor/models/siglip2navit.py | 3 +- vllm/model_executor/models/stablelm.py | 3 +- .../models/transformers/base.py | 6 +- .../models/transformers/fusers/base.py | 27 +-- .../models/transformers/fusers/mla.py | 3 +- .../models/transformers/fusers/moe.py | 12 +- .../models/transformers/fx_utils.py | 33 ++-- .../models/transformers/layers.py | 3 +- .../model_executor/models/transformers/moe.py | 6 +- .../models/transformers/multimodal.py | 24 +-- .../models/transformers/utils.py | 6 +- vllm/model_executor/models/utils.py | 6 +- vllm/model_executor/models/vision.py | 3 +- vllm/model_executor/utils.py | 3 +- .../warmup/qwen4_exp_qsa_warmup.py | 3 +- vllm/models/deepseek_v32/common/kernels.py | 3 +- vllm/models/deepseek_v4/amd/model.py | 3 +- .../deepseek_v4/common/mm_preprocess.py | 9 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_compressor.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_mla.py | 3 +- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 24 +-- vllm/models/deepseek_v4/cpu/model.py | 3 +- vllm/models/deepseek_v4/nvidia/fi_moe.py | 3 +- .../deepseek_v4/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v4/nvidia/model.py | 3 +- vllm/models/deepseek_v4/xpu/model.py | 3 +- vllm/models/deepseek_v41/amd/model.py | 6 +- vllm/models/deepseek_v41/amd/vl_model.py | 3 +- vllm/models/deepseek_v41/attention.py | 3 +- vllm/models/deepseek_v41/common/engram.py | 9 +- .../deepseek_v41/common/mm_preprocess.py | 9 +- .../deepseek_v41/nvidia/flashinfer_sparse.py | 3 +- vllm/models/deepseek_v41/nvidia/model.py | 6 +- vllm/models/deepseek_v41/nvidia/vl_model.py | 3 +- .../models/glm5next/amd/ops/kpool_compress.py | 3 +- vllm/models/glm5next/common/kda.py | 6 +- vllm/models/glm5next/common/multimodal.py | 3 +- vllm/models/glm5next/common/sparse_indexer.py | 6 +- .../glm5next/nvidia/ops/fused_eh_norm.py | 3 +- .../glm5next/nvidia/ops/kpool_compress.py | 3 +- vllm/models/inkling/amd/model.py | 3 +- vllm/models/inkling/amd/moe.py | 9 +- vllm/models/inkling/amd/mtp.py | 9 +- vllm/models/inkling/amd/ops/norm.py | 6 +- vllm/models/inkling/nvidia/model.py | 3 +- vllm/models/inkling/nvidia/moe.py | 9 +- vllm/models/inkling/nvidia/mtp.py | 9 +- vllm/models/inkling/nvidia/ops/mm_towers.py | 6 +- vllm/models/inkling/nvidia/ops/norm.py | 6 +- vllm/models/kimi_k3/nvidia/mla.py | 3 +- vllm/models/minimax_m3/amd/model.py | 3 +- vllm/models/minimax_m3/common/indexer.py | 6 +- .../minimax_m3/common/sparse_attention.py | 3 +- vllm/models/minimax_m3/nvidia/indexer_msa.py | 6 +- vllm/models/minimax_m3/nvidia/model.py | 6 +- vllm/multimodal/cache.py | 6 +- vllm/multimodal/image.py | 3 +- vllm/multimodal/inputs.py | 3 +- vllm/multimodal/utils.py | 3 +- vllm/parser/engine/incremental_lexer.py | 3 +- vllm/parser/engine/streaming_parser_engine.py | 3 +- vllm/parser/engine/token_id_scanner.py | 3 +- vllm/parser/inkling.py | 3 +- vllm/platforms/cuda.py | 3 +- vllm/platforms/interface.py | 3 +- vllm/platforms/rocm.py | 6 +- vllm/profiler/wrapper.py | 3 +- vllm/renderers/base.py | 3 +- vllm/renderers/hf.py | 6 +- vllm/renderers/online_derenderer.py | 3 +- vllm/sampling_params.py | 3 +- vllm/tool_parsers/hermes_tool_parser.py | 3 +- vllm/tool_parsers/lfm2_tool_parser.py | 3 +- vllm/tool_parsers/utils.py | 3 +- vllm/transformers_utils/config.py | 6 +- .../configs/extract_hidden_states.py | 3 +- vllm/transformers_utils/processor.py | 9 +- .../processors/cohere_asr.py | 3 +- .../transformers_utils/processors/glm5next.py | 6 +- vllm/transformers_utils/utils.py | 3 +- vllm/utils/cpu_resource_utils.py | 3 +- vllm/utils/deep_gemm.py | 3 +- vllm/utils/gpu_sync_debug.py | 6 +- vllm/utils/platform_utils.py | 3 +- vllm/utils/torch_utils.py | 3 +- vllm/v1/attention/backend.py | 6 +- vllm/v1/attention/backends/cpu_attn.py | 3 +- vllm/v1/attention/backends/flashinfer.py | 3 +- vllm/v1/attention/backends/mla/amx_mla.py | 3 +- .../attention/backends/mla/flashmla_sparse.py | 3 +- vllm/v1/attention/backends/mla/indexer.py | 3 +- .../v1/attention/backends/mla/prefill/base.py | 3 +- .../backends/rocm_aiter_unified_attn.py | 3 +- vllm/v1/attention/backends/rocm_attn.py | 3 +- vllm/v1/attention/backends/triton_attn.py | 3 +- .../turboquant_soa/triton_turboquant_store.py | 3 +- vllm/v1/core/kv_cache_coordinator.py | 3 +- vllm/v1/core/kv_cache_manager.py | 3 +- vllm/v1/core/sched/interface.py | 6 +- vllm/v1/core/sched/request_queue.py | 9 +- vllm/v1/engine/async_llm.py | 3 +- vllm/v1/engine/core.py | 9 +- vllm/v1/engine/core_client.py | 21 +- vllm/v1/engine/detokenizer.py | 3 +- vllm/v1/engine/utils.py | 3 +- vllm/v1/executor/abstract.py | 3 +- vllm/v1/executor/multiproc_executor.py | 12 +- vllm/v1/executor/ray_utils.py | 3 +- vllm/v1/kv_cache_interface.py | 3 +- vllm/v1/kv_cache_layout.py | 3 +- vllm/v1/kv_offload/base.py | 6 +- vllm/v1/kv_offload/cpu/gpu_worker.py | 21 +- vllm/v1/kv_offload/cpu/spec.py | 6 +- vllm/v1/kv_offload/tiering/async_lookup.py | 3 +- vllm/v1/kv_offload/tiering/example/manager.py | 3 +- vllm/v1/kv_offload/tiering/fs/io.py | 3 +- vllm/v1/kv_offload/tiering/fs/manager.py | 3 +- vllm/v1/kv_offload/tiering/manager.py | 3 +- vllm/v1/kv_offload/tiering/obj/manager.py | 6 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 6 +- vllm/v1/kv_offload/tiering/spec.py | 3 +- vllm/v1/metrics/loggers.py | 3 +- vllm/v1/metrics/perf.py | 3 +- vllm/v1/metrics/ray_wrappers.py | 9 +- vllm/v1/sample/logits_processor/builtin.py | 6 +- vllm/v1/sample/logits_processor/state.py | 3 +- vllm/v1/serial_utils.py | 6 +- vllm/v1/simple_kv_offload/manager.py | 9 +- vllm/v1/spec_decode/extract_hidden_states.py | 3 +- vllm/v1/spec_decode/gemma4.py | 6 +- vllm/v1/spec_decode/llm_base_proposer.py | 3 +- vllm/v1/spec_decode/utils.py | 3 +- vllm/v1/utils.py | 12 +- vllm/v1/worker/encoder_cudagraph.py | 3 +- vllm/v1/worker/gpu/cudagraph_utils.py | 12 +- vllm/v1/worker/gpu/model_runner.py | 3 +- vllm/v1/worker/gpu/model_states/interface.py | 3 +- vllm/v1/worker/gpu/pp_utils.py | 6 +- vllm/v1/worker/gpu/sample/output.py | 3 +- vllm/v1/worker/gpu/sample/sampler.py | 3 +- .../gpu/spec_decode/adaptive_verification.py | 6 +- .../gpu/spec_decode/dflash/cudagraph.py | 3 +- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 3 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 3 +- vllm/v1/worker/gpu_model_runner.py | 15 +- vllm/v1/worker/gpu_worker.py | 3 +- vllm/v1/worker/startup_plan.py | 9 +- vllm/v1/worker/worker_base.py | 3 +- 693 files changed, 2088 insertions(+), 4171 deletions(-) diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py index 1f5da88c2cf8..878be2c29347 100644 --- a/benchmarks/benchmark_hidden_state_extraction.py +++ b/benchmarks/benchmark_hidden_state_extraction.py @@ -136,8 +136,7 @@ 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() diff --git a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py index a4b4533fa3b3..61dfd7ea127c 100644 --- a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py +++ b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py @@ -84,8 +84,7 @@ 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 diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py index f21fbaa703a9..00c6606db009 100644 --- a/benchmarks/kernels/benchmark_moe_defaults.py +++ b/benchmarks/kernels/benchmark_moe_defaults.py @@ -32,8 +32,7 @@ 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, diff --git a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py index 2922c449eaf0..4925e6c02442 100644 --- a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py +++ b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py @@ -97,8 +97,7 @@ 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( diff --git a/docs/mkdocs/gen_files/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py index c051653cf949..a514dcf879ff 100644 --- a/docs/mkdocs/gen_files/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -75,8 +75,7 @@ def main_file(self) -> Path | None: 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 @@ -91,8 +90,7 @@ def other_files(self) -> list[Path]: """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 [] diff --git a/examples/features/logits_processor/custom_req.py b/examples/features/logits_processor/custom_req.py index 36b83ec95876..eb57029efa0b 100644 --- a/examples/features/logits_processor/custom_req.py +++ b/examples/features/logits_processor/custom_req.py @@ -55,8 +55,7 @@ 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`.""" @@ -75,8 +74,7 @@ def __call__( 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): diff --git a/examples/features/logits_processor/custom_req_init.py b/examples/features/logits_processor/custom_req_init.py index 4f29bd271417..38554f06d8ad 100644 --- a/examples/features/logits_processor/custom_req_init.py +++ b/examples/features/logits_processor/custom_req_init.py @@ -56,8 +56,7 @@ 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`.""" @@ -76,8 +75,7 @@ def __call__( 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): diff --git a/examples/rl/rlhf_sharded_rdt_small_ep.py b/examples/rl/rlhf_sharded_rdt_small_ep.py index ed31a32ad8f1..213eae55ef22 100644 --- a/examples/rl/rlhf_sharded_rdt_small_ep.py +++ b/examples/rl/rlhf_sharded_rdt_small_ep.py @@ -101,8 +101,7 @@ class FSDPTrainWorker: """One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray - actor with no producer mixin. - """ + actor with no producer mixin.""" def __init__(self, model_name, rank, world_size, master_addr, master_port): self.rank = rank diff --git a/pyproject.toml b/pyproject.toml index 83949db8fe89..b5c03b22b7d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,7 +90,9 @@ ignore = [ # don't enforce that everything has a docstring "D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107", # unwanted pydocstyle rules - "D205", "D301", "D400", "D401", "D404", "D415", + "D205", "D209", "D301", "D400", "D401", "D404", "D415", + # incompatible pydocstyle pairs; pin the half `ruff` already picks + "D203", "D213", # star imports "F405", "F403", # lambda expression assignment diff --git a/setup.py b/setup.py index 48cbdcf763b4..b24bcc1a2cc5 100644 --- a/setup.py +++ b/setup.py @@ -724,8 +724,7 @@ def warn_if_rocm_torch_version_mismatch( wheels: list[dict], repo_url: str, arch: str ) -> None: """Warn if installed torch differs from the custom ROCm build on - wheels.vllm.ai and suggest the correct install command. - """ + wheels.vllm.ai and suggest the correct install command.""" try: installed = torch.__version__ except Exception: diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index f04bb397f068..6b4895c64369 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -112,8 +112,7 @@ def test_basic_cumem(): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_discard_tags(): """Test that discard(tags) selectively frees GPU memory for specific - tags while keeping other tags mapped and usable. - """ + tags while keeping other tags mapped and usable.""" allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): diff --git a/tests/benchmarks/test_bfcl_dataset.py b/tests/benchmarks/test_bfcl_dataset.py index 982b93ad8e89..d19192231979 100644 --- a/tests/benchmarks/test_bfcl_dataset.py +++ b/tests/benchmarks/test_bfcl_dataset.py @@ -13,8 +13,7 @@ def _patch_hf_api(side_effect): """Return a patch context that swaps `hf_api()` to a stub whose - `.hf_hub_download` attribute uses `side_effect`. - """ + `.hf_hub_download` attribute uses `side_effect`.""" fake_api = MagicMock() fake_api.hf_hub_download.side_effect = side_effect return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api) @@ -116,8 +115,7 @@ def test_bfcl_dataset_translates_schema_and_attaches_tools( ) -> None: """BFCLDataset should translate schemas to OpenAI tool format, set `messages` directly on SampleRequest, and attach tools/tool_choice via - request_overrides. - """ + request_overrides.""" paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -168,8 +166,7 @@ def test_bfcl_dataset_missing_category_raises_clear_error( hf_tokenizer: PreTrainedTokenizerBase, ) -> None: """A typo'd category should produce an actionable ValueError, not an - opaque huggingface_hub exception. - """ + opaque huggingface_hub exception.""" from huggingface_hub.errors import EntryNotFoundError args = _args_for_bfcl(categories=["simpl"]) # typo @@ -187,8 +184,7 @@ def raise_missing(_repo, filename, **_kwargs): @pytest.mark.benchmark def test_chat_backend_uses_messages_field_when_set() -> None: """When RequestFuncInput.chat_messages is set, the chat backend must use - it verbatim and skip default content construction from `prompt`. - """ + it verbatim and skip default content construction from `prompt`.""" import asyncio from vllm.benchmarks.lib.endpoint_request_func import ( @@ -248,8 +244,7 @@ def post(self, url, json, headers): # noqa: A002 def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None: """prompt_len must reflect tokens from both messages *and* tool schemas, so percentile buckets and input-distribution summaries aren't biased - low for tool-heavy traffic. - """ + low for tool-heavy traffic.""" paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): @@ -296,8 +291,7 @@ def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools( tmp_path: Path, ) -> None: """Older tokenizers don't accept tools=; fallback must still produce a - non-zero prompt_len without crashing. - """ + non-zero prompt_len without crashing.""" paths = _write_fake_files(tmp_path) def fake_download(_repo, filename, **_kwargs): diff --git a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py index 28ecb6e443ec..149ea6625a9a 100644 --- a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py +++ b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py @@ -12,8 +12,7 @@ class _RecordingTokenizer: """Minimal tokenizer stub that records the kwargs forwarded to apply_chat_template, so we can assert chat_template_kwargs propagation - without loading a real model/template. - """ + without loading a real model/template.""" def __init__(self) -> None: self.captured_kwargs: dict | None = None diff --git a/tests/benchmarks/test_sampling_params.py b/tests/benchmarks/test_sampling_params.py index 13dcf8ad208c..3bc34a84b377 100644 --- a/tests/benchmarks/test_sampling_params.py +++ b/tests/benchmarks/test_sampling_params.py @@ -10,8 +10,7 @@ class _FakeTokenizer(TokenizerLike): """Minimal tokenizer implementing the TokenizerLike protocol - for testing get_sampling_params. - """ + for testing get_sampling_params.""" def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None: self._vocab_size = vocab_size @@ -172,8 +171,7 @@ def test_output_range_bounds(self): def test_output_low_clamped_to_one(self): """Even with a high ratio that would push output_low to 0, - the function clamps it to 1. - """ + the function clamps it to 1.""" rng = np.random.default_rng(0) # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1 _, output_lens, _ = get_sampling_params( @@ -235,8 +233,7 @@ def test_invalid_dict_missing_keys(self): def test_input_len_zero_with_special_tokens(self): """input_len < num_special_tokens → real_input_len = 0, which is fine - (range [0, 0]). - """ + (range [0, 0]).""" rng = np.random.default_rng(0) input_lens, _, _ = get_sampling_params( rng, 5, 0.0, 5, 50, self._tok(num_special=10) diff --git a/tests/benchmarks/test_skip_tokenizer_init.py b/tests/benchmarks/test_skip_tokenizer_init.py index 1eb9a2536eb6..741db2e0d626 100644 --- a/tests/benchmarks/test_skip_tokenizer_init.py +++ b/tests/benchmarks/test_skip_tokenizer_init.py @@ -52,8 +52,7 @@ def _write_dataset(path: Path) -> None: def _args(dataset_path: str) -> argparse.Namespace: """Reproduce the argparse.Namespace that serve.py builds from the - failing command, including skip_tokenizer_init=True. - """ + failing command, including skip_tokenizer_init=True.""" return argparse.Namespace( # dataset dataset_name="custom", diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 0945bd622a6c..843052f4afdc 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -732,8 +732,7 @@ def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers( monkeypatch: pytest.MonkeyPatch, ): """Verify that without GDN layers in static_forward_context, - the gated pattern is not registered and no matches occur. - """ + the gated pattern is not registered and no matches occur.""" group_shape = GroupShape(1, 128) vllm_config = VllmConfig( model_config=ModelConfig(dtype=dtype), diff --git a/tests/compile/passes/test_split_coalescing.py b/tests/compile/passes/test_split_coalescing.py index 52e531e12742..ab7a0be1a215 100644 --- a/tests/compile/passes/test_split_coalescing.py +++ b/tests/compile/passes/test_split_coalescing.py @@ -15,8 +15,7 @@ class SplitCoalescingModel(torch.nn.Module): """Model with 3 separate split_with_sizes calls on the same input, - simulating the B200+FP8 graph where CSE fails to merge them. - """ + simulating the B200+FP8 graph where CSE fails to merge them.""" def __init__(self, q_size: int, kv_size: int) -> None: super().__init__() diff --git a/tests/compile/test_codegen.py b/tests/compile/test_codegen.py index a0086cb702e5..21db287a2478 100644 --- a/tests/compile/test_codegen.py +++ b/tests/compile/test_codegen.py @@ -77,8 +77,7 @@ def test_non_primitive_kwargs_lifted_to_consts( into the namespace, so ``device`` is unbound. The fix collects such objects into ``__vllm_consts__`` and references them by index. The unqualified ``device(type=...)`` form must never appear in the - generated source. - """ + generated source.""" split_gm = _trace_and_split(model_fn, (x,), split_ops) code, submod_names, consts = generate_execution_code(split_gm) @@ -99,8 +98,7 @@ def test_dtype_singleton_deduped(x: torch.Tensor) -> None: """``torch.float16`` is a process-wide singleton, so two ops referring to it in the traced graph share a single consts slot via ``id()``-based dedup. Distinct expressions (``x.to(...)`` vs ``(x*2).to(...)``) ensure - the tracer can't CSE the two ops into a single node. - """ + the tracer can't CSE the two ops into a single node.""" def model_fn(x: torch.Tensor) -> torch.Tensor: return x.to(dtype=torch.float16) + (x * 2).to(dtype=torch.float16) @@ -149,8 +147,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_consts_ordering_deterministic(x: torch.Tensor) -> None: """Two independent traces of the same model must produce equal consts lists *in the same order*. Cache artifacts identify const slots by - index, so a non-deterministic order would invalidate cached code. - """ + index, so a non-deterministic order would invalidate cached code.""" def model_fn(x: torch.Tensor) -> torch.Tensor: # Multiple distinct non-primitives encountered in a fixed graph order. @@ -188,8 +185,7 @@ def test_consts_shared_across_split_submods(x: torch.Tensor) -> None: The function below splits into three inlined submods, two of which independently reference ``torch.float16``. The shared ``const_index`` threaded through recursive ``generate_execution_code_with_name`` calls - must collapse the dtype to a single slot used from both submods. - """ + must collapse the dtype to a single slot used from both submods.""" def model_fn(x: torch.Tensor) -> torch.Tensor: a = x.to(dtype=torch.float16) # submod_0: _to_copy(fp16) @@ -234,8 +230,7 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None: """When a child of split_gm is *not* a ``torch.fx.GraphModule`` — as happens in production once ``PiecewiseBackend`` replaces submods — codegen emits ``__vllm_submods__[idx](...)`` instead of inlining, and - the runtime callable is bound from ``submod_callables``. - """ + the runtime callable is bound from ``submod_callables``.""" def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -285,8 +280,7 @@ def forward(self, *args, **kwargs): def test_getitem_in_stitching_graph(x: torch.Tensor) -> None: """``operator.getitem`` on submod tuple returns is the ``call_function`` special case at codegen.py — emitted as ``name = source[index]`` - rather than a function call. - """ + rather than a function call.""" def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid() @@ -304,8 +298,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor: def test_del_emitted_for_intermediate_values(x: torch.Tensor) -> None: """The codegen schedules ``del`` after a value's last use to free memory early. Multi-submod splits naturally have intermediates whose - last use is not the output node. - """ + last use is not the output node.""" def model_fn(x: torch.Tensor) -> torch.Tensor: return x.relu().sigmoid().tanh() @@ -325,8 +318,7 @@ def test_with_submod_false_rejects_call_module() -> None: """``generate_execution_code_with_name(with_submod=False)`` is the recursive entry for inlining a GraphModule into its parent. It must refuse a graph that itself contains ``call_module`` nodes — the parent - is responsible for handling those. - """ + is responsible for handling those.""" g = fx.Graph() x_node = g.placeholder("x") root = torch.nn.Module() @@ -343,8 +335,7 @@ def test_node_ref_recurses_through_containers() -> None: """``_node_ref`` is the recursive walker that lifts non-primitives nested inside list/tuple/dict args. Real aten ops rarely produce such structures, but the path is needed for DTensor placement lists and - other future cases — unit-test the walker directly. - """ + other future cases — unit-test the walker directly.""" consts: list = [] const_index: dict[int, int] = {} cpu = torch.device("cpu") @@ -368,8 +359,7 @@ def test_node_ref_recurses_through_containers() -> None: def test_legacy_code_without_consts() -> None: """``compile_execution_fn(consts=None)`` must still load code that has no ``__vllm_consts__`` reference, so older serialized cache artifacts - keep working. - """ + keep working.""" # Pre-consts codegen: no __vllm_consts__ reference, only torch/operator. legacy_code = ( "import torch\n" diff --git a/tests/compile/test_compile_ranges.py b/tests/compile/test_compile_ranges.py index b3e2ec4731a8..4dfea42a6b4a 100644 --- a/tests/compile/test_compile_ranges.py +++ b/tests/compile/test_compile_ranges.py @@ -129,8 +129,7 @@ def test_compile_config_get_compile_ranges(): class PostGradStaticShapeChecker(InductorPass): """Asserts that compile_sizes entries produce graphs with fully concrete - (non-symbolic) shapes, and compile_ranges entries have symbolic shapes. - """ + (non-symbolic) shapes, and compile_ranges entries have symbolic shapes.""" def __init__(self): self.num_static_calls = 0 @@ -171,8 +170,7 @@ def uuid(self) -> str: def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache): """Verify that compile_sizes entries are compiled with fully concrete - shapes (no SymInts), while compile_ranges entries retain dynamic shapes. - """ + shapes (no SymInts), while compile_ranges entries retain dynamic shapes.""" checker = PostGradStaticShapeChecker() torch.set_default_device("cuda") vllm_config = VllmConfig( diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index fc3295621806..5c9132ccbfe2 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -1229,8 +1229,7 @@ def test_compile_sizes_padding_validation(): def test_inductor_asserts_default_disabled(monkeypatch): """Test that inductor runtime asserts are disabled by default - (INFO logging level) on torch < 2.12. - """ + (INFO logging level) on torch < 2.12.""" monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib @@ -1248,8 +1247,7 @@ def test_inductor_asserts_default_disabled(monkeypatch): def test_inductor_asserts_enabled_in_debug(monkeypatch): """Test that VLLM_LOGGING_LEVEL=DEBUG enables inductor runtime asserts - on torch < 2.12. - """ + on torch < 2.12.""" monkeypatch.setenv("VLLM_LOGGING_LEVEL", "DEBUG") import importlib @@ -1285,8 +1283,7 @@ def test_get_inductor_factors_includes_configs(): def test_inductor_asserts_user_override(monkeypatch): """Test that explicit inductor_compile_config overrides the - debug-logging default. - """ + debug-logging default.""" monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO") import importlib diff --git a/tests/compile/test_inductor_fallback_allow_list_patch.py b/tests/compile/test_inductor_fallback_allow_list_patch.py index c5705da21171..29fe9962e345 100644 --- a/tests/compile/test_inductor_fallback_allow_list_patch.py +++ b/tests/compile/test_inductor_fallback_allow_list_patch.py @@ -236,8 +236,7 @@ def test_kimi_k2_6_style_op_stream(self): def test_inner_set_membership_still_works_for_standard_ops(self): """The patch must not break Inductor's existing fallback decisions - for non-vllm ops such as ``torchvision::roi_align``. - """ + for non-vllm ops such as ``torchvision::roi_align``.""" import torch._inductor.lowering as _lowering allow_list = _lowering.FALLBACK_ALLOW_LIST diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index f71d4875d75c..b68820c0d51c 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -372,8 +372,7 @@ class _HeterogeneousConfig(PretrainedConfig): """A heterogeneous config with no convertor of its own. Mirrors the parts vLLM uses: per-layer configs are shallow copies with the - varying attributes applied and heterogeneity stripped, so they do not recurse. - """ + varying attributes applied and heterogeneity stripped, so they do not recurse.""" is_heterogeneous = True @@ -472,8 +471,7 @@ def test_deepseek_v4_convertor_splits_vision_architecture(): """The vision checkpoint shares model_type/architectures with the text-only DeepSeek-V4; the convertor routes it to the VL wrapper class by rewriting hf_config.architectures (model-class resolution reads the - raw attribute). - """ + raw attribute).""" from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.transformers_utils.model_arch_config_convertor import ( DeepseekV4ModelArchConfigConvertor, diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index abf834608a55..84cf092c2159 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -42,8 +42,7 @@ def test_mm_encoder_attn_backend_hash_updates(): def test_language_model_only_does_not_affect_mm_hash(): """language_model_only does not affect the ViT computation graph, - so it should not change the multimodal config hash. - """ + so it should not change the multimodal config hash.""" base_hash = MultiModalConfig().compute_hash() lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash() assert base_hash == lm_only_hash @@ -51,8 +50,7 @@ def test_language_model_only_does_not_affect_mm_hash(): def test_language_model_only_affects_model_hash(): """language_model_only affects the LM computation graph, - so it should change the model config hash. - """ + so it should change the model config hash.""" model = "llava-hf/llava-1.5-7b-hf" base_hash = ModelConfig(model).compute_hash() lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash() @@ -388,8 +386,7 @@ def _resolve_mm_video_decode_device( torchcodec_available: bool = True, ) -> dict: """Run processor-device then video-decode-device resolution and report - the resulting video media IO kwargs. - """ + the resulting video media IO kwargs.""" mm_config = MultiModalConfig( mm_processor_kwargs={} if device is None else {"device": device}, mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type] @@ -419,8 +416,7 @@ def _resolve_mm_video_decode_device( def test_auto_video_decode_uses_nvdec_on_encoder_instance(): """The processor runs on the accelerator there, so decoded frames should - stay on-device too. - """ + stay on-device too.""" assert _resolve_mm_video_decode_device(ec_role="ec_producer") == { "backend": "torchcodec", "device": "cuda", @@ -432,8 +428,7 @@ def test_auto_video_decode_stays_on_cpu_off_encoder_instance( ec_role: ECRole | None, ): """The processor stays on CPU off encode-only instances, so video - decoding should too. - """ + decoding should too.""" assert _resolve_mm_video_decode_device(ec_role=ec_role) == {} diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 10f8c994f44b..8505172da387 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -34,8 +34,7 @@ def _make_hf_config(**kwargs) -> PretrainedConfig: @pytest.mark.cpu_test def test_dict_overrides_are_not_forwarded_to_draft(): """Dict overrides are target-specific key patches; the draft must get - only the architecture-mapping override. - """ + only the architecture-mapping override.""" composed = SpeculativeConfig.compose_draft_hf_overrides( {"max_position_embeddings": 1234} ) @@ -51,8 +50,7 @@ def test_none_overrides_fall_back_to_arch_mapping(): @pytest.mark.cpu_test def test_callable_overrides_reach_the_draft_config(): """A callable override (config-to-config transform) composes with the - architecture-mapping override and is applied to the draft config. - """ + architecture-mapping override and is applied to the draft config.""" def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.num_hidden_layers = 1 @@ -69,8 +67,7 @@ def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: @pytest.mark.cpu_test def test_arch_mapping_applies_before_callable_override(): """The static arch-mapping override runs first, so the user callable - observes (and may adjust) the post-mapping config. - """ + observes (and may adjust) the post-mapping config.""" seen_architectures: list[str] = [] def record(hf_config: PretrainedConfig) -> PretrainedConfig: @@ -131,8 +128,7 @@ def test_composed_override_is_picklable(): the composed override must be picklable. A nested local closure is not (it raised ``Can't get local object`` on DFlashDraftModel); a ``functools.partial`` over a module-referenceable static method is. - Guard against regressing to a closure. - """ + Guard against regressing to a closure.""" composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) assert isinstance(composed, functools.partial) diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py index 711aaf9b036b..117ae5a66d53 100644 --- a/tests/config/test_speculative_draft_max_position_embeddings.py +++ b/tests/config/test_speculative_draft_max_position_embeddings.py @@ -99,8 +99,7 @@ def test_independent_draft_model_keeps_its_own_limit( vllm_caplog: pytest.LogCaptureFixture, ): """An independent AR draft may genuinely have a smaller context than the - target; its max_position_embeddings must not be resized. - """ + target; its max_position_embeddings must not be resized.""" target_model_config = ModelConfig( AR_MODEL, hf_overrides={"max_position_embeddings": 8192} ) diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index 876ee9dcf518..85a9d73aa5c6 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -179,8 +179,7 @@ def close(self): @pytest.fixture def enable_ray_v2_backend(): """Set env vars for the Ray V2 executor backend and shut down Ray - between tests. - """ + between tests.""" import ray saved = { diff --git a/tests/distributed/test_eplb_quant_scale_consistency.py b/tests/distributed/test_eplb_quant_scale_consistency.py index 49dc8afbf311..3b50b315f0a1 100644 --- a/tests/distributed/test_eplb_quant_scale_consistency.py +++ b/tests/distributed/test_eplb_quant_scale_consistency.py @@ -124,8 +124,7 @@ def _build_processed_layer( backend: str, raw: dict[str, torch.Tensor], device: torch.device ) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]: """Create the method + layer, load ``raw`` and run the real - process_weights_after_loading (which builds the flashinfer kernel). - """ + process_weights_after_loading (which builds the flashinfer kernel).""" moe_config = _make_moe_config(backend) method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False) layer = _RoutedExpertsStub(moe_config).to(device) @@ -158,8 +157,7 @@ def _assert_tensors_equal( def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None: """Permute experts of every registered Parameter in place, the way - EPLB's in-place rearrangement moves expert slices along dim 0. - """ + EPLB's in-place rearrangement moves expert slices along dim 0.""" with torch.no_grad(): for _, param in layer.named_parameters(): param.copy_(param[perm]) @@ -167,8 +165,7 @@ def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> def _ensure_world1_distributed() -> None: """The enable_eplb config makes weight processing all-reduce activation - scale amaxes over the EP group, so a (single-rank) EP group must exist. - """ + scale amaxes over the EP group, so a (single-rank) EP group must exist.""" from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, init_distributed_environment, @@ -256,8 +253,7 @@ def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None: def test_fp8_per_tensor_alphas_registered_and_aliased() -> None: """CPU-only: the fp8 per-tensor oracle must register the fused (w_scale * a_scale) products as layer Parameters aliased by the quant - config, so EPLB rearrangement keeps kernels consistent. - """ + config, so EPLB rearrangement keeps kernels consistent.""" from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, make_fp8_moe_quant_config, diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 0258c781c721..1cd67524602f 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -160,8 +160,7 @@ def _make_forward_context(rank, world_size, num_tokens_per_rank): class _AttnMeta: """Minimal placeholder so set_forward_context's ``attn_metadata is not None`` guard (forward_context.py:334) - is satisfied. The real DPMetadata is built from num_tokens_across_dp. - """ + is satisfied. The real DPMetadata is built from num_tokens_across_dp.""" dp_metadata = None @@ -460,8 +459,7 @@ def _one_sided_workspace_grow_worker(rank, world_size): def test_one_sided_manager_workspace_grow(world_size): """A later initialize() with a larger per-token payload must grow the workspace and rebuild MoeAlltoAll; a later initialize() with a smaller - payload must no-op. - """ + payload must no-op.""" _spawn_workers( _one_sided_workspace_grow_worker, world_size, diff --git a/tests/distributed/test_mq_connect_ip.py b/tests/distributed/test_mq_connect_ip.py index 701dc8f248c7..05617bbb6037 100644 --- a/tests/distributed/test_mq_connect_ip.py +++ b/tests/distributed/test_mq_connect_ip.py @@ -21,8 +21,7 @@ def test_mq_bind_with_local_ip(): """MessageQueue with remote readers should successfully bind - when connect_ip is the local node's IP. - """ + when connect_ip is the local node's IP.""" # n_reader=2, n_local_reader=1 means 1 remote reader, # which triggers the remote ZMQ socket bind. mq = MessageQueue( @@ -44,8 +43,7 @@ def test_mq_bind_with_local_ip(): def test_mq_bind_with_non_local_ip_fails(): """MessageQueue should fail to bind when connect_ip is a non-local IP address (simulating the bug where master_addr - from a different node was used). - """ + from a different node was used).""" # Use a non-local IP that we definitely can't bind to. # 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned. non_local_ip = "198.51.100.1" @@ -59,8 +57,7 @@ def test_mq_bind_with_non_local_ip_fails(): def test_mq_bind_defaults_to_local_ip(): """When connect_ip is None, MessageQueue should auto-detect - the local IP and bind successfully. - """ + the local IP and bind successfully.""" mq = MessageQueue( n_reader=2, n_local_reader=1, diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 0eda9a14e05d..91dccd579869 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -495,8 +495,7 @@ def test_pynccl_suspend_resume_idempotent(): def test_pynccl_suspend_noop_without_symbol(): """RCCL / NCCL < 2.29.7 lack ncclCommSuspend; suspend() must no-op, not - crash, and resume() must then no-op too since nothing was suspended. - """ + crash, and resume() must then no-op too since nothing was suspended.""" from unittest.mock import Mock comm = object.__new__(PyNcclCommunicator) diff --git a/tests/distributed/test_ray_v2_executor_e2e.py b/tests/distributed/test_ray_v2_executor_e2e.py index 27de410e3892..335324c37128 100644 --- a/tests/distributed/test_ray_v2_executor_e2e.py +++ b/tests/distributed/test_ray_v2_executor_e2e.py @@ -23,8 +23,7 @@ def _ray_init(): Without this, workers cannot unpickle actor classes defined in the ``tests`` package, causing FunctionActorManager to fall back to - TemporaryActor which drops async method signatures. - """ + TemporaryActor which drops async method signatures.""" project_root = str(pathlib.Path(__file__).resolve().parents[2]) ray.init( ignore_reinit_error=True, diff --git a/tests/distributed/test_sharded_rdt_plan.py b/tests/distributed/test_sharded_rdt_plan.py index 5cf0a82203b1..7923c413fc2a 100644 --- a/tests/distributed/test_sharded_rdt_plan.py +++ b/tests/distributed/test_sharded_rdt_plan.py @@ -53,8 +53,7 @@ class _FakeLayer: """Stand-in for a leaf module. The planner only stores it and keys on - ``id()``; nothing in the plan path touches its attributes. - """ + ``id()``; nothing in the plan path touches its attributes.""" def __init__(self, tag: str) -> None: self.tag = tag @@ -65,8 +64,7 @@ def __repr__(self) -> str: def _module(layer, copies): """What the bake records for one leaf module: a scatter list whose entries - all carry that module as their ``layer``. - """ + all carry that module as their ``layer``.""" for c in copies: c.layer = layer return copies @@ -83,8 +81,7 @@ def _copy( dtype=torch.bfloat16, ): """A `_Scatter` as the bake would record it: source key, owning layer, and - the meta destination region. - """ + the meta destination region.""" stride = ( (1,) if len(shape) == 1 @@ -143,8 +140,7 @@ def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None): def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): """`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model - where every gather group is one fully-baked leaf module. - """ + where every gather group is one fully-baked leaf module.""" names = ["embed.weight"] names += [f"model.layers.{i}.w" for i in range(n_layers)] names += ["norm.weight"] @@ -165,8 +161,7 @@ def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4): class TestFakeOpChains: """Every allowlisted op must append itself to the chain and hand back a child whose shape/dtype PyTorch itself computed. The chain is the wire - format the producer replays, so its exact contents are load-bearing. - """ + format the producer replays, so its exact contents are load-bearing.""" def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None): return FakeRDTTensor( @@ -227,8 +222,7 @@ def test_kwargs_are_frozen_sorted_for_hashability(self): @pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)]) def test_multi_return_op_emits_one_child_per_output(self, op, n): """chunk/unbind hand back a tuple; each child carries the base op plus a - trailing __getitem__(i) so the producer can index the replayed result. - """ + trailing __getitem__(i) so the producer can index the replayed result.""" parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0) assert isinstance(parts, tuple) assert len(parts) == n @@ -244,8 +238,7 @@ def test_op_chain_is_hashable_as_a_fetch_key(self): class TestFakeUnsupportedOps: """Anything that needs real data must fail loudly at bake time rather than - silently transferring the wrong bytes. - """ + silently transferring the wrong bytes.""" def _fake(self): return FakeRDTTensor( @@ -285,8 +278,7 @@ def test_unsupported_is_a_notimplementederror(self): class TestBakeRecording: """During the dry run the fake's ``copy_`` is the data sink: it records the - source chain plus the meta destination's strided region and moves nothing. - """ + source chain plus the meta destination's strided region and moves nothing.""" def _recorder_and_fake(self, shape=(4, 6)): rec = BakeSink() @@ -322,8 +314,7 @@ def test_copy_marks_the_source_name_live(self): def test_unattributed_copy_is_live_but_unrecorded(self): """A copy_ with no loader stamp cannot be attributed to a param, so its - module must fall back to the plain load — but the name still moved data. - """ + module must fall back to the plain load — but the name still moved data.""" rec, fake = self._recorder_and_fake() rec.current = None torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake) @@ -363,8 +354,7 @@ def test_the_recorded_dtype_is_the_chains_output_not_the_sources(self): producer packs what the replay yields, so the record must carry the POST-chain dtype: taking it from the source name's metadata instead sizes the slice with the wrong itemsize and shifts every later slice in the - chunk, carving the packed blob differently on the two sides. - """ + chunk, carving the packed blob differently on the two sides.""" rec, fake = self._recorder_and_fake(shape=(4,)) layer = _FakeLayer("reinterpreted") rec.current = (layer, "weight") @@ -384,8 +374,7 @@ def test_a_broadcasting_copy_is_refused_at_bake(self): the producer packs what the chain yields, and nothing downstream can catch a mismatch — the consumer's buffer view is exactly prod(dest.shape) elements, so it reshapes cleanly over bytes laid out at other offsets. - So the bake refuses rather than recording a slice it cannot carve. - """ + So the bake refuses rather than recording a slice it cannot carve.""" rec, fake = self._recorder_and_fake(shape=(1, 6)) rec.current = (_FakeLayer("broadcast"), "weight") param = torch.empty((4, 6), dtype=torch.bfloat16, device=META) @@ -401,8 +390,7 @@ def test_a_broadcasting_copy_is_refused_at_bake(self): class TestConsumerIdentity: """Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the worker's producer block and keys the producer's per-consumer serve ring, so - a collision silently serves two workers out of one ring. - """ + a collision silently serves two workers out of one ring.""" def _engine(self, *, dp_index, rank, world_size): eng = object.__new__(ShardedRDTWeightTransferEngine) @@ -439,8 +427,7 @@ def test_one_engine_indexes_by_dp_and_tp(self): def test_independent_engines_offset_into_distinct_ranges(self): """A fleet of separate engines restarts _global_worker_index at 0 in each, so without the replica offset every engine would claim 0..w-1. - The driver sets replica_rank on the payload; these are the ids it buys. - """ + The driver sets replica_rank on the payload; these are the ids it buys.""" ids = [ self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2) for r in (0, 1) @@ -466,8 +453,7 @@ def _router( group_lens=None, ): """A router over ``n`` single-name groups unless told otherwise, so a test - can state expectations in group indices. - """ + can state expectations in group indices.""" if names is None: n = len(name_owner_class or []) or (len(group_lens or []) or 1) names = [f"g{i}" for i in range(n)] @@ -496,8 +482,7 @@ def test_identity_when_fleets_match(self): def test_every_producer_nic_carries_traffic(self): """16 producers / 8 consumers: the block rule spreads each consumer's pulls over its block, alternating by group, so no producer sits idle. - Plain `consumer_id % len(owners)` would idle half of them. - """ + Plain `consumer_id % len(owners)` would idle half of them.""" r = _router(16, 8, group_lens=[1] * 95) names = [f"g{i}" for i in range(95)] for c in range(8): @@ -515,8 +500,7 @@ def test_fan_in_shares_one_producer(self): def test_pipeline_stages_route_to_the_owning_stage(self): """2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner - sets, one per stage; each consumer reaches both. - """ + sets, one per stage; each consumer reaches both.""" r = _router( 16, 8, @@ -532,8 +516,7 @@ def test_pipeline_stages_route_to_the_owning_stage(self): def test_expert_names_route_to_the_holding_ranks(self): """A name held by only some ranks routes to those; a replicated name - keeps the full set. - """ + keeps the full set.""" r = RdtRouter( 8, 4, @@ -552,8 +535,7 @@ def test_expert_names_route_to_the_holding_ranks(self): def test_consumers_spread_over_an_owner_set(self): """Several ranks hold a name (its TP peers); the block rule must spread - consumers across them, not funnel through one. - """ + consumers across them, not funnel through one.""" r = _router( 4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6 ) @@ -562,8 +544,7 @@ def test_consumers_spread_over_an_owner_set(self): def test_a_group_frees_at_every_owner_of_any_of_its_names(self): """The free barrier is per group, so its fan-out is the UNION over the - group's names: a rank holding none of them is not signalled. - """ + group's names: a rank holding none of them is not signalled.""" r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1]) assert r.group_owners(0) == [0, 1, 2, 3] assert r.group_owners(1) == [0, 1] @@ -633,8 +614,7 @@ def test_every_replica_of_a_worker_resolves_the_same_producer(self): def test_adding_a_deployment_does_not_move_the_first_one(self): """The strongest form: an existing consumer's routes are byte-identical - before and after another deployment is provisioned. - """ + before and after another deployment is provisioned.""" one, names = self._fleet(16, 8, 1) two, _ = self._fleet(16, 8, 2) for consumer_id in range(8): @@ -644,8 +624,7 @@ def test_adding_a_deployment_does_not_move_the_first_one(self): def test_one_deployment_is_the_historical_carve(self): """Default (no width) and width == the whole fleet must agree with the - plain block rule, so a single-deployment fleet is untouched. - """ + plain block rule, so a single-deployment fleet is untouched.""" plain = _router(16, 8, group_lens=[1] * 6) explicit, names = self._fleet(16, 8, 1) for consumer_id in range(8): @@ -659,8 +638,7 @@ def test_one_deployment_is_the_historical_carve(self): def test_the_overlay_holds_for_partially_owned_names(self): """Pipeline stages and experts too: two owner sets, one narrow. The narrow one already forced agreement (one owner, no block to carve); the - wide one is what the carve decides. - """ + wide one is what the carve decides.""" r, names = self._fleet( 8, 4, @@ -676,16 +654,14 @@ def test_the_overlay_holds_for_partially_owned_names(self): def test_a_width_that_does_not_divide_the_fleet_raises(self): """A uniform fleet is asserted upstream in ``get_world_size``; if it ever is not, the overlay would map two workers of one deployment onto the same - block index, so refuse rather than serve the wrong bytes. - """ + block index, so refuse rather than serve the wrong bytes.""" with pytest.raises(ValueError, match="does not divide"): RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4) def test_pulls_still_carry_the_global_consumer_id(self): """The carve uses the index WITHIN a deployment, but the wire must keep the fleet-global id: it is what lets the producer tell the sharers of a - slot apart, and count their arrivals separately. - """ + slot apart, and count their arrivals separately.""" class _Method: def __init__(self): @@ -728,8 +704,7 @@ def _producer_pack_offsets(slices): class TestPackedLayout: def test_consumer_layout_matches_the_producer_rule(self): """Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style - group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment. - """ + group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment.""" specs = [ ("w.fp8", "float8_e4m3fn", [17]), ("w.scale", "float32", [3]), @@ -804,8 +779,7 @@ def test_duplicate_source_keys_are_packed_once(self): class TestChunkModuleScatters: """The chunk cut: one chunk per distinct owner class present in the copies, ascending by class index. Derived purely from the bake plus the ownership - table, so vLLM's expert placement needs no cases. - """ + table, so vLLM's expert placement needs no cases.""" def test_unstamped_copies_form_one_chunk(self): layer = _FakeLayer("l") @@ -819,8 +793,7 @@ def test_one_chunk_per_owner_class_ascending(self): """Classes are numbered by first appearance in metadata order, so the replicated class (0) leads and the rest follow — the ordering the old "-1 first, then ascending stamp" rule produced, now falling out of the - table instead of a special case. - """ + table instead of a special case.""" layer = _FakeLayer("moe") names = ["norm", "e5", "e0", "e2"] copies = [_copy(n, shape=(4,)) for n in names] @@ -840,8 +813,7 @@ def test_one_chunk_per_owner_class_ascending(self): def test_copy_order_within_a_chunk_is_bake_order(self): """Scatter order within a chunk must follow the bake, not the stamp - sort: the packed layout and the producer's replay agree on keys order. - """ + sort: the packed layout and the producer's replay agree on keys order.""" layer = _FakeLayer("moe") names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names copies = [_copy(n, shape=(4,)) for n in names] @@ -855,8 +827,7 @@ def test_copy_order_within_a_chunk_is_bake_order(self): def test_stamps_interleave_across_modules_without_reordering_within(self): """Two modules' copies bucket by stamp; within each bucket the order is - module-then-bake order. - """ + module-then-bake order.""" a, b = _FakeLayer("a"), _FakeLayer("b") mods = [ _module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]), @@ -876,8 +847,7 @@ def test_stamps_interleave_across_modules_without_reordering_within(self): def test_scatters_carry_their_own_dtype(self): """Dtype rides the record from the bake, where it is the fake's dtype AFTER its op chain — not a plan-time lookup of the source name, which - would be wrong for any chain that reinterprets dtype. - """ + would be wrong for any chain that reinterprets dtype.""" layer = _FakeLayer("l") eng = _planner({}, name_meta={"w": ("float32", [10])}) ((_er, scatters),) = eng._chunk_module_scatters( @@ -916,8 +886,7 @@ def test_stamps_multiply_chunks_within_a_group(self): def test_materialize_fires_on_a_modules_first_chunk_only(self): """Empty HF params are allocated once per module, by construction — - including a FusedMoE-like module whose copies span owner-class chunks. - """ + including a FusedMoE-like module whose copies span owner-class chunks.""" layer = _FakeLayer("spanning") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -952,8 +921,7 @@ def test_free_signal_fires_on_each_groups_last_chunk(self): def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self): """The load-bearing case: with the group cut across owner-class chunks the signal must hang off the LAST one. Signaling earlier lets the producers - drop the gather buffers while a later chunk's RDMA is still reading. - """ + drop the gather buffers while a later chunk's RDMA is still reading.""" layer = _FakeLayer("l") names = [f"w{i}" for i in range(4)] copies = [_copy(n, shape=(4,)) for n in names] @@ -969,8 +937,7 @@ def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self): def test_free_signal_timing_across_two_split_groups(self): """Two groups x 2 owner-class chunks each: signals land on each group's - last chunk only. - """ + last chunk only.""" names_a = [f"model.layers.0.w{i}" for i in range(4)] names_b = [f"model.layers.1.w{i}" for i in range(4)] baked, name_meta, stamps = {}, {}, {} @@ -988,8 +955,7 @@ def test_free_signal_timing_across_two_split_groups(self): def test_a_group_with_nothing_local_is_signaled_at_sync_start(self): """The owners still published it, so this consumer must signal it — immediately, before the pipeline (signal-before-publish is tolerated), - never hung off another group's chunk. - """ + never hung off another group's chunk.""" baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["model.layers.1.w"] # nothing baked, and not live => no pull eng = _planner( @@ -1010,8 +976,7 @@ def test_leading_groups_with_nothing_local_go_to_pre_free(self): def test_modules_are_deduped_within_a_group(self): """Several names of one fused module map to the same scatter list; the - plan must not chunk it twice. - """ + plan must not chunk it twice.""" layer = _FakeLayer("fused") names = ["qkv.q", "qkv.k", "qkv.v"] copies = [_copy(n, shape=(4,)) for n in names] @@ -1027,8 +992,7 @@ def test_modules_are_deduped_within_a_group(self): def test_plan_is_a_pure_function_of_its_inputs(self): """Called twice with the same arguments it must produce the same plan — - that is what lets the engine build it once and reuse it every sync. - """ + that is what lets the engine build it once and reuse it every sync.""" baked, name_meta, names, group_lens = _one_module_per_layer(2) eng = _planner(baked, name_meta=name_meta) a = eng._build_call_plan(names, group_lens) @@ -1048,8 +1012,7 @@ def test_single_producer_routes_every_chunk_to_local_owner_zero(self): class TestUnbakedNamesGuard: """A live name with no baked plan has no way to load — there is no fallback, and its pull would target groups the pipeline already freed — so the plan - build fails at init, naming the names. - """ + build fails at init, naming the names.""" def test_a_live_unbaked_name_fails_the_plan_build(self): baked, name_meta, names, group_lens = _one_module_per_layer(2) @@ -1060,8 +1023,7 @@ def test_a_live_unbaked_name_fails_the_plan_build(self): def test_never_copied_names_are_dropped_entirely(self): """Experts owned by another EP rank no-op in their loader: not an error, - and not a chunk — their group is pre-freed if nothing else fills it. - """ + and not a chunk — their group is pre-freed if nothing else fills it.""" baked, name_meta, names, group_lens = _one_module_per_layer(2) del baked["norm.weight"] eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"}) @@ -1073,15 +1035,13 @@ def test_never_copied_names_are_dropped_entirely(self): class TestCallPlanRouting: """Under partial ownership each chunk must be pulled from a producer that actually holds every name in it. Consumers bind EVERY producer (signals fan - out to all owners), so local index == trainer rank. - """ + out to all owners), so local index == trainer rank.""" def _routed_planner( self, held_by_group, num_consumers, consumer_id, *, held_by=None ): """``held_by_group[g]`` = the ranks holding group g's names, so a test - can still speak in groups; ``held_by`` overrides individual names. - """ + can still speak in groups; ``held_by`` overrides individual names.""" baked, name_meta, names, group_lens = _one_module_per_layer( len(held_by_group) - 2 ) @@ -1108,16 +1068,14 @@ def _routed_planner( def test_two_stage_ownership_splits_chunks_between_owners(self): """Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the - second. Every chunk goes to a holder of its names. - """ + second. Every chunk goes to a holder of its names.""" eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0) plan = eng._build_call_plan(names, group_lens) assert [c.owner for c in plan.chunks] == [0, 0, 1, 1] def test_a_multi_owner_group_still_resolves_to_one_producer(self): """Every chunk is served by exactly ONE producer: splitting a pull only - multiplies produce calls, since the consumer's NIC bounds it. - """ + multiplies produce calls, since the consumer's NIC bounds it.""" eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0) plan = eng._build_call_plan(names, group_lens) assert all(0 <= c.owner < 2 for c in plan.chunks) @@ -1125,8 +1083,7 @@ def test_a_multi_owner_group_still_resolves_to_one_producer(self): def test_an_expert_name_routes_to_the_rank_holding_it(self): """PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert name inside a stage-1 group is held by rank 3 alone. The name's own owner - set is what routes it. - """ + set is what routes it.""" eng, names, group_lens = self._routed_planner( [[0, 1], [0, 1], [2, 3], [2, 3]], 1, @@ -1141,8 +1098,7 @@ def test_an_expert_name_routes_to_the_rank_holding_it(self): def test_an_unheld_name_is_rejected_at_init(self): """A name no rank holds can never be served. The trainer raises when it transposes the holdings; the router's validate() is the consumer-side - backstop, and either way it fails at init rather than at first pull. - """ + backstop, and either way it fails at init rather than at first pull.""" with pytest.raises(ValueError, match="empty"): RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate() @@ -1151,16 +1107,14 @@ class TestSignalCompleteness: """Every gather group is signaled exactly once — after its last chunk when the worker pulls from it, at sync start otherwise — for ANY expert placement. A group signaled twice over-credits the barrier; a group never - signaled parks a producer credit and hangs end_sync. - """ + signaled parks a producer credit and hangs end_sync.""" def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2): """Pre / n_layers MoE layers / post. Each layer: one norm (its own module, stamp -1) + n_experts expert names stamped ``e // n_local``. The worker's bake covers only ``worker_experts`` (its placement); foreign expert names never copied => dropped entirely, as in the real - bake. - """ + bake.""" n_local = n_experts // ep_size names = ["embed.weight"] group_lens = [1] @@ -1220,8 +1174,7 @@ def test_round_robin_placement_signals_every_group_exactly_once(self): def test_placement_changes_only_the_chunk_count_never_the_signals(self): """Linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6 - hits all 4. More chunks per group, identical signal set. - """ + hits all 4. More chunks per group, identical signal set.""" plans = {} for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])): eng, names, group_lens = self._moe_planner(worker_experts=experts) @@ -1235,8 +1188,7 @@ def test_placement_changes_only_the_chunk_count_never_the_signals(self): def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self): """The FusedMoE shape: one module's copies span every owner-class chunk of - its group; materialize on its first, quant on its last. - """ + its group; materialize on its first, quant on its last.""" eng, names, group_lens = self._moe_planner( worker_experts=[0, 2, 4, 6], n_layers=1 ) @@ -1294,8 +1246,7 @@ def test_partition_is_one_group_per_layer_between_unindexed_blocks(self): def test_flattening_a_partition_returns_the_input_order(self): """Un-indexed names group by POSITION, not by name class — so one that looks like a pre-block name still lands after the layers it follows. - Group index therefore means the same thing on every rank. - """ + Group index therefore means the same thing on every rank.""" names = [ "embed.weight", "model.layers.0.a", @@ -1313,8 +1264,7 @@ def test_layers_are_ordered_by_index_not_appearance(self): def test_stacks_keep_prefix_appearance_order_and_sort_within(self): """Two stacks stay in the order they appear — the vision tower before the - text stack — while each sorts by index internally. - """ + text stack — while each sorts by index internally.""" names = [ "visual.blocks.1.w", "visual.blocks.0.w", @@ -1331,8 +1281,7 @@ def test_stacks_keep_prefix_appearance_order_and_sort_within(self): def test_the_post_block_lands_last_however_early_it_arrives(self): """Megatron-Bridge streams the last pipeline stage's output block before its layers. Sweeping un-indexed names after the first layer into a - trailing group is what lets the gather loop walk groups ascending. - """ + trailing group is what lets the gather loop walk groups ascending.""" names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"] assert layerwise_groups(names) == [ ["model.layers.0.w"], @@ -1342,8 +1291,7 @@ def test_the_post_block_lands_last_however_early_it_arrives(self): def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self): """A raw-checkpoint source may yield in shard-packing order. Coalescing on - the key keeps that from shattering one layer into many groups. - """ + the key keeps that from shattering one layer into many groups.""" names = [ "model.layers.0.a", "model.layers.1.a", @@ -1385,15 +1333,13 @@ def test_a_model_with_no_layers_is_one_pre_group(self): def test_every_supported_naming_convention_partitions_per_layer(self, prefix): """One decoder layer per group under any prefix, with no per-architecture table. A literal `model.layers.` match silently yields one whole-model - group for most of these. - """ + group for most of these.""" names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)] assert layerwise_groups(names) == [[n] for n in names] def test_a_moe_layer_stays_whole(self): """The OUTERMOST index wins, so per-expert names group by layer. Splitting - on the expert would cut FusedMoE — one leaf module — across groups. - """ + on the expert would cut FusedMoE — one leaf module — across groups.""" names = [ "model.layers.0.mlp.experts.0.w1.weight", "model.layers.0.mlp.experts.1.w1.weight", @@ -1410,8 +1356,7 @@ def test_a_moe_layer_stays_whole(self): def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self): """Qwen-VL layout. Today's literal `model.layers.` match puts every one of these in ONE group — the whole model — because the text stack ships as - `model.language_model.layers.`. - """ + `model.language_model.layers.`.""" names = [ "visual.patch_embed.proj.weight", "visual.blocks.0.attn.qkv.weight", @@ -1430,8 +1375,7 @@ def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self): def test_two_stacks_sharing_an_index_do_not_merge(self): """The key carries the prefix, so vision block 0 and text layer 0 are - different groups even though both are index 0. - """ + different groups even though both are index 0.""" names = [ "visual.blocks.0.attn.qkv.weight", "model.language_model.layers.0.self_attn.q_proj.weight", @@ -1468,23 +1412,20 @@ def test_the_two_allowlists_describe_the_same_contract(self): def test_every_bake_recordable_op_is_replayable(self): """Anything the bake can record must be serveable, or the failure lands a - whole sync later than the mistake. - """ + whole sync later than the mistake.""" assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS) def test_the_producer_allows_nothing_the_consumer_cannot_emit(self): """The allowlist is also a guard against a misbehaving or spoofed consumer invoking arbitrary methods on trainer tensors, so it must not be wider than the ops the bake can actually produce. ``to`` in particular - would let a replay change dtype or device. - """ + would let a replay change dtype or device.""" assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values()) assert "to" not in ALLOWED_OPS def test_transpose_via_t_is_serveable(self): """The specific regression: a chain recorded from ``.t()`` must pass the - producer's guard. - """ + producer's guard.""" fake = FakeRDTTensor( name="w", shape=torch.Size((4, 6)), @@ -1529,8 +1470,7 @@ def test_transpose_via_t_is_serveable(self): class TestRequiresTheRayExecutor: """The engine's data plane is Ray's, so a non-Ray executor cannot work. The check is at construction because the alternative is an opaque failure during - the first handshake, long after the misconfiguration. - """ + the first handshake, long after the misconfiguration.""" def _construct(self, backend): cfg = SimpleNamespace( @@ -1557,8 +1497,7 @@ def test_ray_is_accepted(self): def test_a_custom_executor_class_is_left_alone(self): """A `type[Executor]` override is deliberate and unjudgeable here, so it - must not be rejected for merely not being the string "ray". - """ + must not be rejected for merely not being the string "ray".""" class _CustomExecutor: pass diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 7036d102897a..8147841d9aa3 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -236,8 +236,7 @@ def test_rebuilt_views_have_the_requested_dtype_and_geometry(self, server_factor def test_views_of_one_storage_do_not_overlap(self, server_factory): """One IPC export per storage, one as_strided view per name — the whole point of the storage/view split. Overlapping views would serve the same - bytes for two different weights. - """ + bytes for two different weights.""" server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A, nbytes=64) @@ -260,8 +259,7 @@ def test_publish_records_the_group_as_unfreed(self, server_factory): def test_the_default_lookahead_is_one(self): """1 is the sweet spot: group N+1 is gathered AND pullable while N is being pulled, with resident memory at its floor of 2 groups — the bound - the larger-model runs size against. - """ + the larger-model runs size against.""" assert DEFAULT_GATHER_LOOKAHEAD == 1 @@ -276,8 +274,7 @@ def test_single_consumer_signal_releases_the_group(self, server_factory): def test_group_is_held_until_the_last_live_consumer_signals(self, server_factory): """Every live consumer signals every owner; releasing on the first - signal would drop storage other consumers are still reading. - """ + signal would drop storage other consumers are still reading.""" server = server_factory() server.begin_sync(3) _publish(server, GI_A, GROUP_A) @@ -304,8 +301,7 @@ def test_a_signal_that_beats_its_publish_is_completed_at_publish( ): """A consumer with nothing to pull from a group signals it at sync start, which can precede the publish. The publish must then release the - group rather than wait for a signal that will never come again. - """ + group rather than wait for a signal that will never come again.""" server = server_factory() server.begin_sync(1) server.free_group(GI_A) @@ -328,8 +324,7 @@ def test_freed_groups_are_reported_back_to_the_engine(self, server_factory): the server says the group is gone — and it hears that ONLY through ``wait_freed`` / ``end_sync``, never a publish return (a freed notice riding an unharvested async publish while the engine blocks in - ``wait_freed`` would wedge the loop). - """ + ``wait_freed`` would wedge the loop).""" server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -358,8 +353,7 @@ def test_release_drops_cache_and_count_entries(self, server_factory): class _CountingSource: """`iter_groups` yielding one CPU-tensor group per call, built lazily so a - group's tensors exist only once the engine's credit gate lets it gather. - """ + group's tensors exist only once the engine's credit gate lets it gather.""" def __init__(self, groups): self._groups = groups @@ -372,8 +366,7 @@ def iter_groups(self): class _ResidencyDict(dict): """The engine's `_inflight`, instrumented: `max_len` is the high-water mark of groups whose CUDA-IPC export refs were alive at once — the trainer's - resident-memory bound in groups. - """ + resident-memory bound in groups.""" def __init__(self): super().__init__() @@ -387,8 +380,7 @@ def __setitem__(self, key, value): def _loop_engine(server, n_groups, *, lookahead): """A real engine wired straight to a real (non-Ray) server, with only the state `_run_gather_loop` reads. The `_rpc` seam dispatches inline; - `_publish_async` sees no `.remote` and runs inline too. - """ + `_publish_async` sees no `.remote` and runs inline too.""" groups = [(f"model.layers.{gi}.w",) for gi in range(n_groups)] e = ShardedRDTTrainerWeightTransferEngine.__new__( ShardedRDTTrainerWeightTransferEngine @@ -411,16 +403,14 @@ class TestGatherCredit: serveable immediately, and the ENGINE's loop stops gathering while more than `gather_lookahead` groups are unfreed. So at most `lookahead + 1` groups are resident — two at the default — while group N+1 is pulled the instant N's - pulls finish. - """ + pulls finish.""" @pytest.fixture def gather_engine(self, monkeypatch): """Neutralize the two CUDA touches in the export path so the REAL gather loop runs on CPU: `.cuda()` becomes identity, and the storage export ships the `(nbytes, ..., device, ...)` tuple the fixture's - stubbed `rebuild_cuda_tensor` expects. - """ + stubbed `rebuild_cuda_tensor` expects.""" monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self) monkeypatch.setattr( trainer_mod, @@ -432,8 +422,7 @@ def gather_engine(self, monkeypatch): def test_publish_never_blocks(self, server_factory): """Publishing must not park on a credit: the whole point is that a gathered group becomes serveable immediately. Three publishes with no - frees, inline on this thread — a block here hangs the test. - """ + frees, inline on this thread — a block here hangs the test.""" server = server_factory(gather_lookahead=1) server.begin_sync(1) for gi in range(3): @@ -447,8 +436,7 @@ def test_lookahead_one_holds_at_most_two_groups_resident( than 2 groups of gathered weights, no matter how the consumer paces — and the pipeline still overlaps, because the consumer here refuses to free group N until N+1 is already published (pullable). A gate an - off-by-one too tight deadlocks this test; too loose fails max_len. - """ + off-by-one too tight deadlocks this test; too loose fails max_len.""" server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -490,8 +478,7 @@ def test_lookahead_zero_holds_exactly_one_group( ): """0 is the serialized baseline: the loop waits for group N to be freed before gathering N+1, so exactly ONE group is ever resident and nothing - overlaps. A floor of 1 on the bound would silently pipeline instead. - """ + overlaps. A floor of 1 on the bound would silently pipeline instead.""" server = server_factory(gather_lookahead=0) engine = gather_engine(server, 4, lookahead=0) @@ -517,8 +504,7 @@ def test_the_residency_bound_scales_with_the_lookahead( self, server_factory, gather_engine ): """Lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing - consumer sees exactly 3 groups gathered before the loop parks. - """ + consumer sees exactly 3 groups gathered before the loop parks.""" server = server_factory(gather_lookahead=2) engine = gather_engine(server, 8, lookahead=2) @@ -554,8 +540,7 @@ def test_wait_freed_returns_the_freed_backlog(self, server_factory): def test_a_gather_error_releases_a_blocked_wait_freed(self, server_factory): """Otherwise a trainer-side failure on another rank deadlocks this one inside its credit gate. It must RAISE, not return empty — an empty - return would spin the engine straight back into the wait. - """ + return would spin the engine straight back into the wait.""" server = server_factory() server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -636,8 +621,7 @@ def test_begin_sync_sets_the_barrier_target(self, server_factory): def test_the_target_is_floored_at_one(self, server_factory): """A bare/zero call must not make every group free instantly (or divide - the barrier by zero). - """ + the barrier by zero).""" server = server_factory() server.begin_sync(0) assert server._live_count == 1 @@ -645,8 +629,7 @@ def test_the_target_is_floored_at_one(self, server_factory): def test_a_degraded_sync_lowers_the_target(self, server_factory): """FT degraded sync: the live count is the WHOLE degraded-sync mechanism on this side — the group releases after the live consumers - alone. - """ + alone.""" server = server_factory() server.begin_sync(2) # 4 provisioned, 2 alive _publish(server, GI_A, GROUP_A) @@ -670,8 +653,7 @@ def test_begin_sync_clears_per_sync_state(self, server_factory): def test_a_straggler_signal_cannot_credit_the_next_sync(self, server_factory): """begin_sync resets the counts, so a signal that leaked past the sync boundary would over-credit — the reason the consumer drains its fired - signals before finishing. Pin the reset side of that contract. - """ + signals before finishing. Pin the reset side of that contract.""" server = server_factory() server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -693,8 +675,7 @@ def test_the_packed_destination_cache_survives_the_sync_boundary( self, server_factory ): """The packed layout repeats every sync — that is what makes caching the - destination views worth ~5ms per 384-spec group. - """ + destination views worth ~5ms per 384-spec group.""" server = server_factory() server._pack_dsts[("sentinel",)] = (0, []) server.begin_sync(1) @@ -705,8 +686,7 @@ class TestServedNamesGuard: """A producer serves only the names it holds: its stage's groups, and within them its own EP coordinate's experts plus the replicated names. Without this guard a misrouted pull would block forever in the cache wait - for a name this rank never publishes. - """ + for a name this rank never publishes.""" def test_a_pull_for_an_unserved_name_fails_loudly(self, server_factory): server = server_factory(served_names=list(GROUP_A)) @@ -722,8 +702,7 @@ def test_the_error_names_the_unserved_weights(self, server_factory): def test_served_names_none_accepts_anything(self, server_factory): """Gather-to-all: every producer holds every group, so there is nothing - to guard. The call proceeds to the cache wait (not exercised here). - """ + to guard. The call proceeds to the cache wait (not exercised here).""" server = server_factory(served_names=None) assert server._served_names is None @@ -770,8 +749,7 @@ def _wait(): def test_end_sync_fails_when_a_consumer_never_signals(self, server_factory): """The exact mid-sync death signature: the group is published, one of - the two live consumers that owed a signal is gone. - """ + the two live consumers that owed a signal is gone.""" server = server_factory(stall_timeout_s=0.3) server.begin_sync(2) _publish(server, GI_A, GROUP_A) @@ -785,8 +763,7 @@ def test_the_error_reaches_every_blocked_waiter(self, server_factory): """The watchdog fires on the existing `set_gather_error` channel, so one stall unwinds the whole rank through one path rather than each waiter timing out separately — here both a consumer parked in the serve cache - wait and the engine parked in its credit gate. - """ + wait and the engine parked in its credit gate.""" server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -819,8 +796,7 @@ def test_progress_anywhere_keeps_a_slow_sync_alive(self, server_factory): """A consumer that is slow but signaling must never trip it: the stamp is global to the producer, so a steady trickle of signals holds the whole rank open. Nine sequential publishes with a 0.3s timeout each take - longer in total than the timeout, and none may fire. - """ + longer in total than the timeout, and none may fire.""" server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) for gi in range(9): @@ -833,8 +809,7 @@ def test_progress_anywhere_keeps_a_slow_sync_alive(self, server_factory): def test_begin_sync_resets_the_progress_stamp(self, server_factory): """Syncs are minutes apart. Without the reset the first publish of sync N+1 - would measure its stall from somewhere inside sync N and fire immediately. - """ + would measure its stall from somewhere inside sync N and fire immediately.""" server = server_factory(gather_lookahead=1, stall_timeout_s=0.3) server.begin_sync(1) _publish(server, GI_A, GROUP_A) @@ -856,8 +831,7 @@ def test_the_lifecycle_survives_concurrent_publishes_and_signals( self, server_factory ): """Every group published must end up freed and out of the cache, with - end_sync returning cleanly — the property whose violation is a hang. - """ + end_sync returning cleanly — the property whose violation is a hang.""" server = server_factory(gather_lookahead=2) server.begin_sync(1) groups = [(gi, (f"model.layers.{gi}.w",)) for gi in range(12)] @@ -892,8 +866,7 @@ def test_concurrent_signals_of_one_group_release_it_exactly_once( ): """All live consumers signaling at once: the release must happen on the last signal and only once, or the engine drops IPC refs twice (or - never). - """ + never).""" server = server_factory() server.begin_sync(8) _publish(server, GI_A, GROUP_A) @@ -920,8 +893,7 @@ class TestFakeServerAgreesWithTheRealOne: """`test_sharded_rdt_trainer.py`'s `_FakeProducerServer` is a second, independent model of this protocol that the engine-side tests assert on. Pin the two against each other so the fake cannot silently drift from the - semantics the real server enforces. - """ + semantics the real server enforces.""" @staticmethod def _fake(): @@ -952,8 +924,7 @@ class TestExportRing: """The packed export path. Its one hard invariant is that the trainer packs a group exactly the way ``publish_group`` rebuilds it — a layout disagreement would serve the wrong bytes with nothing downstream to catch - it. - """ + it.""" @staticmethod def _engine(): @@ -1045,8 +1016,7 @@ def test_a_slot_is_never_rewritten_while_its_group_is_still_live( ): """Free the later group of each pair first and hold the earlier one live until the NEXT group is gathered -- the window a modular slot counter got - wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's). - """ + wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's).""" server = server_factory(gather_lookahead=1) n_groups = 8 engine = gather_engine(server, n_groups, lookahead=1) @@ -1170,8 +1140,7 @@ class TestSharedServeSlots: def test_one_deployment_never_shares(self, sharing_server): """With the width equal to the live count every group is a singleton, so - each call packs for itself out of its own ring. - """ + each call packs for itself out of its own ring.""" server = sharing_server(workers_per_replica=2) server.begin_sync(2, [0, 1]) _publish(server, GI_A, GROUP_A) @@ -1251,8 +1220,7 @@ def test_the_slot_follows_the_consumers_issue_order(self, sharing_server): def test_a_seqless_call_is_refused(self, sharing_server): """The slot comes from ``seq``, so a caller that sends none cannot be - served safely. - """ + served safely.""" server = sharing_server(workers_per_replica=2) server.begin_sync(1, [0]) _publish(server, GI_A, GROUP_A) @@ -1262,8 +1230,7 @@ def test_a_seqless_call_is_refused(self, sharing_server): def test_a_failed_pack_reaches_every_sharer(self, sharing_server): """Only the packer touches the buffer, so its failure has to be - republished: a waiter must raise rather than read a half-written slot. - """ + republished: a waiter must raise rather than read a half-written slot.""" server = sharing_server(workers_per_replica=2) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1287,8 +1254,7 @@ def test_a_failed_pack_reaches_every_sharer(self, sharing_server): def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server): """A dead deployment must not be waited for: a survivor left alone in - its group packs for itself instead of blocking on the corpse. - """ + its group packs for itself instead of blocking on the corpse.""" server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0) server.begin_sync(3, [0, 1, 3]) # consumer 2 died _publish(server, GI_A, GROUP_A) @@ -1302,8 +1268,7 @@ def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server): def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server): """A consumer that dies INSIDE the sync window has no detector, exactly - as for the free barrier; the rendezvous waits on the same channel. - """ + as for the free barrier; the rendezvous waits on the same channel.""" server = sharing_server(workers_per_replica=2, stall_timeout_s=0.3) server.begin_sync(4, [0, 1, 2, 3]) _publish(server, GI_A, GROUP_A) @@ -1315,8 +1280,7 @@ def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server): class TestSharedReservation: """``reserve_serve_buffer`` under sharing: one ring per GROUP, and the - init-time check that the group's members pull the same chunks. - """ + init-time check that the group's members pull the same chunks.""" def test_matching_plans_reserve_one_ring_between_them(self, sharing_server): server = sharing_server(workers_per_replica=2) @@ -1338,8 +1302,7 @@ def test_sharing_off_reserves_a_ring_per_consumer(self, sharing_server): def test_mismatched_plans_fail_at_reservation(self, sharing_server): """Sharing rests on the sharers pulling the same chunks. Checked here, where it names both consumers, rather than at serve time, where it is a - rendezvous nobody completes. - """ + rendezvous nobody completes.""" server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") @@ -1348,8 +1311,7 @@ def test_mismatched_plans_fail_at_reservation(self, sharing_server): def test_different_groups_may_differ(self, sharing_server): """Only the members of ONE group have to agree; worker 0 and worker 1 of - the same deployment hold different slices and never share a slot. - """ + the same deployment hold different slices and never share a slot.""" server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, "plan-a") server.reserve_serve_buffer(1, 1024, "plan-b") @@ -1358,8 +1320,7 @@ def test_different_groups_may_differ(self, sharing_server): def test_a_missing_digest_skips_the_check(self, sharing_server): """The digest is optional on the wire, so an engine that does not send - one still reserves. - """ + one still reserves.""" server = sharing_server(workers_per_replica=2) server.reserve_serve_buffer(0, 1024, None) server.reserve_serve_buffer(2, 1024, None) diff --git a/tests/distributed/test_sharded_rdt_trainer.py b/tests/distributed/test_sharded_rdt_trainer.py index df9628cc8501..9339e503c45f 100644 --- a/tests/distributed/test_sharded_rdt_trainer.py +++ b/tests/distributed/test_sharded_rdt_trainer.py @@ -34,8 +34,7 @@ class _ListSource(WeightSource): """A WeightSource over an explicit ordered (name, cpu-tensor) list, so the - sharded-RDT group/order logic can be tested without a real model. - """ + sharded-RDT group/order logic can be tested without a real model.""" def __init__(self, pairs): self._pairs = list(pairs) @@ -55,8 +54,7 @@ class _FakeProducerServer: per-group barrier: publish/free are keyed by GROUP INDEX, signals count to the ``begin_sync`` live total, publish_group returns nothing, and freed groups flow back only through wait_freed / end_sync (see - test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne). - """ + test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne).""" def __init__(self, auto_free=True): self.order: list[str] = [] @@ -100,8 +98,7 @@ def wait_freed(self): must already have a banked credit when the gate asks (auto_free, or a test's own free_group) — anything else is the deadlock the real watchdog would kill, so fail loudly. Not appended to ``order``: it is - pacing, not a lifecycle milestone. - """ + pacing, not a lifecycle milestone.""" assert self._pending_freed, ( "wait_freed with nothing freed: the gather loop would deadlock" ) @@ -124,8 +121,7 @@ def _rdt_engine_with_fake_server( ): """Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the - group-major metadata, and route _rpc to the fake. - """ + group-major metadata, and route _rpc to the fake.""" import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod # reduce_tensor needs CUDA; the fake server never rebuilds, so stub it. @@ -325,8 +321,7 @@ def update_weights(self, update_info): class _OwnedSource(_ListSource): """A source holding only some groups' names, like a pipeline-parallel rank. - Takes group indices for convenience and declares the names inside them. - """ + Takes group indices for convenience and declares the names inside them.""" def __init__(self, pairs, owned_group_idx): super().__init__(pairs) @@ -392,8 +387,7 @@ def __iter__(self): ) def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): """The free barrier's target is one integer per sync: the live consumer - count, defaulting to the whole provisioned fleet. - """ + count, defaulting to the whole provisioned fleet.""" server = _FakeProducerServer(auto_free=True) engine = _rdt_engine_with_fake_server( _rdt_source_two_layers(), @@ -411,8 +405,7 @@ def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): class TestLiveCountPlumbing: """``send_weights(live_consumer_ids)`` -> the barrier target. The provisioned geometry is frozen; a degraded sync only lowers the target and - narrows each slot-sharing group's rendezvous. - """ + narrows each slot-sharing group's rendezvous.""" @staticmethod def _engine(num_consumers, world=2, rank=0): @@ -445,8 +438,7 @@ def test_a_full_live_set_matches_the_provisioned_count(self): def test_which_consumers_died_does_not_matter_only_how_many(self): """The whole point of the barrier: no routed per-producer targets, so - the identity of the dead consumer is irrelevant to the producers. - """ + the identity of the dead consumer is irrelevant to the producers.""" counts = [] for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]): engine, got = self._engine(num_consumers=8) @@ -457,8 +449,7 @@ def test_which_consumers_died_does_not_matter_only_how_many(self): def test_the_live_ids_travel_with_the_count(self): """The count sizes the free barrier, the ids size the slot-sharing rendezvous, so a producer that shares slots can tell WHICH consumers it - is still waiting for. They must describe the same set. - """ + is still waiting for. They must describe the same set.""" engine, got = self._engine(num_consumers=8) engine.send_weights([6, 0, 2, 2]) ((count, ids),) = got @@ -526,8 +517,7 @@ def held_names(self): def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): """Every name must be held somewhere or it can never be served — caught - when the holdings are transposed, naming the orphan. - """ + when the holdings are transposed, naming the orphan.""" with pytest.raises(ValueError, match="no trainer rank holds"): _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]), @@ -541,8 +531,7 @@ def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch): """Only the sender's metadata reaches the consumers, so a rank describing - just its own share must fail loudly rather than silently drop the rest. - """ + just its own share must fail loudly rather than silently drop the rest.""" engine = _rdt_engine_with_fake_server( _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]), is_sender=False, @@ -565,8 +554,7 @@ def _serve_ring_server(src_name, src): """A producer server with one cached tensor and a pre-seeded serve ring, so a pull needs no Ray and no NIXL registration. Returns (server, serve) where ``serve(chain)`` packs one spec into the SAME ring slot every time — which is - what puts the destination-view cache, and only it, under test. - """ + what puts the destination-view cache, and only it, under test.""" from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( _RDTProducerServer, ) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 2274f15d97bd..7debed4cbbff 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -419,8 +419,7 @@ def test_tensor_broadcast(): def _dumps_oob(obj) -> tuple[bytes, list]: """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor - dispatch table + out-of-band buffers >= 1MiB. - """ + dispatch table + out-of-band buffers >= 1MiB.""" buffers = [] def callback(buf: pickle.PickleBuffer) -> bool: @@ -491,8 +490,7 @@ def test_tensor_pickle_roundtrip(case: str): @pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) def test_reduce_tensor_fallback(case: str): """Tensors the zero-copy reducer can't safely alias must fall back to - torch's default reduction. - """ + torch's default reduction.""" if case == "cuda": if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -755,8 +753,7 @@ def test_remote_subscribe_addr_unique_concurrent_writers( Pre-fix, the writer probed a port with get_open_port() and bound it afterwards; pinning the probe to one free port makes every writer bind the same port and fail deterministically on that code path, - while the late-binding implementation never consults the probe. - """ + while the late-binding implementation never consults the probe.""" from vllm.distributed.device_communicators import shm_broadcast colliding_port = get_open_port() diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py index fad1d2a00666..031937d7f5e1 100644 --- a/tests/distributed/test_utils.py +++ b/tests/distributed/test_utils.py @@ -31,8 +31,7 @@ def get_cuda_visible_devices(self): def test_cuda_device_count_stateless(): """Test that cuda_device_count_stateless changes return value if - CUDA_VISIBLE_DEVICES is changed. - """ + CUDA_VISIBLE_DEVICES is changed.""" if current_platform.is_rocm(): pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.") actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index ced71ead76c5..1353def73862 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1166,8 +1166,7 @@ def finish_weight_update(self, weight_version: str | None = None) -> None: def _module_with(*pairs): """A tiny nn.Module exposing the given (name, tensor) pairs as parameters, - so trainer tests can build a ModuleSource without a real model. - """ + so trainer tests can build a ModuleSource without a real model.""" module = torch.nn.Module() for name, tensor in pairs: module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) @@ -1288,8 +1287,7 @@ def test_source_is_reiterable(self): def test_metadata_agrees_with_iteration(self): """The two channels must line up element-for-element: engines declare - the round from `metadata()` and then send what iteration yields. - """ + the round from `metadata()` and then send what iteration yields.""" source = ModuleSource( _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3))) ) @@ -1303,13 +1301,11 @@ def test_metadata_agrees_with_iteration(self): class TestWeightSourceGroupContract: """`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what backends gather and free by, so the default must agree with - `layerwise_groups` over `metadata()`, restricted to what this rank holds. - """ + `layerwise_groups` over `metadata()`, restricted to what this rank holds.""" class _Source(WeightSource): """Minimal source over an ordered (name, tensor) list, optionally holding - only some names (in which case it iterates only their groups). - """ + only some names (in which case it iterates only their groups).""" def __init__(self, names, held=None, reverse=False): self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)] @@ -1364,8 +1360,7 @@ def test_groups_keeps_a_partially_held_group_whole(self): def test_groups_order_follows_metadata_not_the_declaration(self): """The declaration is a SET of names; the groups it selects still come out in metadata order, so each pairs with the right ``iter_groups()`` - batch however the source listed them. - """ + batch however the source listed them.""" names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"] source = self._source( names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"] @@ -1399,16 +1394,14 @@ def test_iter_groups_yields_only_held_groups(self): def test_out_of_order_iteration_raises(self): """Materializing is usually a collective, so a rank that iterates out of - order deadlocks its peers -- fail loudly instead. - """ + order deadlocks its peers -- fail loudly instead.""" source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True) with pytest.raises(RuntimeError, match="iteration order must match"): list(source.iter_groups()) def test_a_source_may_override_iter_groups(self): """The extension point: materialize a whole group in one step instead of - one generator resume per tensor. - """ + one generator resume per tensor.""" calls = [] class _Batched(TestWeightSourceGroupContract._Source): @@ -1429,8 +1422,7 @@ class TestDeferredProcessingContract: caller that takes over the update tail (running its own `finalize_layerwise_reload` instead of going through `finish_weight_update`) reads the flag and calls the method. Both must be answerable on any engine, or - that caller ends up reaching through a getattr. - """ + that caller ends up reaching through a getattr.""" def _engines(self): return { @@ -1444,8 +1436,7 @@ def test_every_engine_declares_whether_it_defers(self): def test_every_engine_can_be_drained(self): """The default is a no-op, so a caller never has to check whether the - method exists before calling it. - """ + method exists before calling it.""" for name, cls in self._engines().items(): assert callable(cls.drain_pending), name @@ -1539,8 +1530,7 @@ def test_shutdown_default_is_noop(self): ) def test_ipc_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; - the packed wire param rides the init info, not the per-round update_info. - """ + the packed wire param rides the init info, not the per-round update_info.""" client = RecordingClient() engine = IPCTrainerWeightTransferEngine( client=client, @@ -1559,8 +1549,7 @@ def test_ipc_trainer_send_weights_drives_client_in_order(): def test_ipc_trainer_init_ships_packed_to_worker(): """trainer_init drives the inference-side init handshake and propagates the - must-agree `packed` flag to the worker. - """ + must-agree `packed` flag to the worker.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU (CUDA IPC handles).") @@ -1581,8 +1570,7 @@ def test_ipc_trainer_init_ships_packed_to_worker(): def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the inference-side init handshake with the worker-shaped init info (rank_offset=1) while opening its own endpoint, - and propagates the must-agree wire params to the worker. - """ + and propagates the must-agree wire params to the worker.""" import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod # Bypass the real NCCL rendezvous. @@ -1623,8 +1611,7 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): """The worker engine reads packed + buffer geometry from the trainer-supplied init info at the handshake, not from the config or the - per-round update info. - """ + per-round update info.""" import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod monkeypatch.setattr( @@ -1657,8 +1644,7 @@ def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): """Non-sender trainer ranks build an engine without opening an endpoint or - touching the client; they only join the collectives in send_weights. - """ + touching the client; they only join the collectives in send_weights.""" client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( init_info=NCCLTrainerInitInfo( @@ -1688,8 +1674,7 @@ def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): def test_nccl_trainer_send_weights_drives_client_in_order(): """send_weights issues start -> update -> finish and ships per-round metadata; the packed wire params ride the init handshake, not the - per-round update_info. - """ + per-round update_info.""" client = RecordingClient() engine = NCCLTrainerWeightTransferEngine( client=client, @@ -1710,8 +1695,7 @@ def test_nccl_trainer_send_weights_drives_client_in_order(): class _ScriptedSource(WeightSource): """Declares `meta` but yields whatever `pairs` says — used to drive the - metadata/iteration agreement checks. - """ + metadata/iteration agreement checks.""" def __init__(self, meta, pairs): self._meta = meta @@ -1747,8 +1731,7 @@ def test_nccl_trainer_init_requires_source(): def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): """The worker sizes its buffers (and cuts packed chunks) from metadata(), so - iteration disagreeing with it must raise rather than corrupt the stream. - """ + iteration disagreeing with it must raise rather than corrupt the stream.""" meta = [ ParamMeta("w", torch.float32, (4,)), ParamMeta("b", torch.float32, (2,)), @@ -1762,8 +1745,7 @@ def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch): """A source that declares one wire dtype and materializes another would make - the two sides disagree on every byte offset. - """ + the two sides disagree on every byte offset.""" meta = [ParamMeta("w", torch.float32, (4,))] engine = _mock_group_engine( _ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]), @@ -1790,8 +1772,7 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch): def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch): """NCCL sends numel elements from data_ptr(), so a non-contiguous view must - be linearized first or the worker receives unrelated memory. - """ + be linearized first or the worker receives unrelated memory.""" base = torch.arange(6, dtype=torch.float32).reshape(2, 3) view = base.t() # non-contiguous meta = [ParamMeta("w", torch.float32, tuple(view.shape))] @@ -1861,8 +1842,7 @@ def _sparse_patch(device: str = "cpu") -> SparseWeightPatch: def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): """The sender's trainer_init drives the init handshake with the worker-shaped init info; sparse ships no packed wire params, so the worker - keeps its unpacked defaults. Sparse takes no `source`. - """ + keeps its unpacked defaults. Sparse takes no `source`.""" import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock()) @@ -1915,8 +1895,7 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch): def test_sparse_nccl_trainer_send_weights_empty_round_is_noop(): """A round with no patches must not touch the client (an empty sparse - update info is invalid by construction). - """ + update info is invalid by construction).""" client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() @@ -1952,8 +1931,7 @@ def test_sparse_nccl_trainer_rejects_source(): def test_sparse_nccl_trainer_validates_patch_before_any_rpc(): """Malformed patches must fail on the trainer, before start_weight_update: the worker's own checks only run once the broadcasts are already in flight, - where a size mismatch wedges both sides instead of raising. - """ + where a size mismatch wedges both sides instead of raising.""" client = RecordingClient() engine = SparseNCCLTrainerWeightTransferEngine(client=client) engine.model_update_group = MagicMock() diff --git a/tests/distributed/test_weight_transfer_nccl_uid.py b/tests/distributed/test_weight_transfer_nccl_uid.py index 58201376041a..18a6442a2a94 100644 --- a/tests/distributed/test_weight_transfer_nccl_uid.py +++ b/tests/distributed/test_weight_transfer_nccl_uid.py @@ -403,8 +403,7 @@ def trainer_broadcast_tensor_torch_free( def _run_uid_transfer(trainer_task) -> None: """Mint a unique id, run the worker and the given trainer peer concurrently, and assert the tensor round-trips. Both ranks must enter init together -- - there is no store barrier on the UID path. - """ + there is no store barrier on the UID path.""" from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary _init_ray_for_weight_transfer() @@ -451,8 +450,7 @@ def _run_uid_transfer(trainer_task) -> None: ) def test_nccl_weight_transfer_between_processes_uid(): """Weight transfer over a pre-shared ncclUniqueId (no TCPStore), with both - ranks on vLLM's ``PyNcclCommunicator``. - """ + ranks on vLLM's ``PyNcclCommunicator``.""" _run_uid_transfer(trainer_broadcast_tensor_uid) @@ -462,7 +460,6 @@ def test_nccl_weight_transfer_between_processes_uid(): ) def test_nccl_weight_transfer_torch_free_trainer(): """UID weight transfer against a foreign peer that uses no torch and no vLLM - communicator, only a cupy buffer and raw ``ncclBroadcast``. - """ + communicator, only a cupy buffer and raw ``ncclBroadcast``.""" pytest.importorskip("cupy") _run_uid_transfer(trainer_broadcast_tensor_torch_free) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 9adfedc9f489..68fe25067713 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -570,8 +570,7 @@ def test_attention_config(): def test_multi_node_world_size_includes_pcp(monkeypatch): """PCP expands the process world size, so the --nnodes divisibility check must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world - size of 1 and the launch is rejected before the engine starts. - """ + size of 1 and the launch is rejected before the engine starts.""" import vllm.config.vllm # PCP requires the V2 model runner, which is gated on Triton. @@ -797,8 +796,7 @@ def test_expand_json_human_readable_numbers(input_json, expected_json): ) def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): """Cloud storage URIs should not be passed to get_model_path() - when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs. - """ + when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.""" import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -810,8 +808,7 @@ def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): """Cloud storage tokenizer URI should not be passed to - get_model_path() when HF_HUB_OFFLINE=1. - """ + get_model_path() when HF_HUB_OFFLINE=1.""" import huggingface_hub monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) @@ -973,8 +970,7 @@ class TestDpDeviceIdSharding: def test_dp_supervisor_device_ids_stay_env_relative(self): """Regression test: the DP supervisor must pass env-relative indices, not physical IDs, because each child re-resolves --device-ids - against its inherited device-control env var. - """ + against its inherited device-control env var.""" import argparse from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids @@ -1001,8 +997,7 @@ def test_dp_supervisor_shards_user_device_ids(self): def test_dp_rank_shards_user_assigned_gpu_ids(self): """get_physical_gpu_ids_for_local_dp_rank slices the user-provided - --device-ids list instead of recomputing from the env var. - """ + --device-ids list instead of recomputing from the env var.""" from vllm.platforms import current_platform from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 63e44eacebb4..226908d67368 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -241,8 +241,7 @@ def _make_tool_result_request( self, tool_result_content ) -> AnthropicMessagesRequest: """Build a request with assistant tool_use followed by user - tool_result. - """ + tool_result.""" return _make_request( [ { @@ -327,8 +326,7 @@ def test_tool_result_with_image(self): def test_tool_result_with_text_and_image(self): """Mixed text+image tool_result: text in tool msg, image in user - msg. - """ + msg.""" request = self._make_tool_result_request( [ {"type": "text", "text": "Here is the screenshot"}, @@ -425,8 +423,7 @@ def test_tool_result_no_follow_up_when_no_images(self): class TestAttributionHeaderStripping: def test_billing_header_stripped_from_system(self): """Claude Code's x-anthropic-billing-header block should be - stripped to preserve prefix caching. - """ + stripped to preserve prefix caching.""" request = _make_request( [{"role": "user", "content": "Hello"}], system=[ @@ -1150,8 +1147,7 @@ async def sse_input(): async def test_buffered_content_flushed_on_done_without_usage_chunk(self): """Content buffered during tool_use must be emitted even if the stream jumps straight from finish_reason to [DONE], skipping the - empty-choices usage chunk. - """ + empty-choices usage chunk.""" async def sse_input(): yield _make_stream_chunk( @@ -1255,8 +1251,7 @@ class TestStreamingCacheUsageSemantics: async def test_streaming_cache_fields_absent_then_populated(self): """First chunk lacks prompt_tokens_details (vLLM contract); message_start omits cache fields. The final chunk carries - prompt_tokens_details, so message_delta carries resolved values. - """ + prompt_tokens_details, so message_delta carries resolved values.""" async def sse_input(): yield _make_stream_chunk( @@ -1301,8 +1296,7 @@ async def sse_input(): @pytest.mark.asyncio async def test_streaming_no_cache_hit(self): """When the final chunk reports cached_tokens=0, message_delta carries - cache fields = 0 (cache miss); message_start still omits them. - """ + cache fields = 0 (cache miss); message_start still omits them.""" async def sse_input(): yield _make_stream_chunk( @@ -1343,8 +1337,7 @@ async def sse_input(): @pytest.mark.asyncio async def test_streaming_no_prompt_tokens_details_at_all(self): """If --enable-prompt-tokens-details is off, no chunk carries cache - info; both message_start and message_delta omit cache fields. - """ + info; both message_start and message_delta omit cache fields.""" async def sse_input(): yield _make_stream_chunk( @@ -1465,8 +1458,7 @@ def test_empty_completion_emits_one_text_block(self): class TestCacheSalt: def test_cache_salt_passed_through(self): """cache_salt on the Anthropic request reaches the converted - ChatCompletionRequest so prefix-cache isolation works via /v1/messages. - """ + ChatCompletionRequest so prefix-cache isolation works via /v1/messages.""" request = _make_request( [{"role": "user", "content": "Hello"}], cache_salt="tenant-abc-secret-salt", @@ -1671,8 +1663,7 @@ def _request_body() -> dict: @staticmethod def _conversion_error() -> ValidationError: """A real pydantic ValidationError like the one ChatCompletionRequest - construction raises when Anthropic input violates the OpenAI schema. - """ + construction raises when Anthropic input violates the OpenAI schema.""" class _StubRequest(BaseModel): stop: Annotated[list[str], Field(max_length=4)] | None = None @@ -1683,8 +1674,7 @@ class _StubRequest(BaseModel): def test_validation_error_returns_bad_request(self): """A pydantic ValidationError during Anthropic->OpenAI conversion is - surfaced as a 400 BadRequestError, not a 500. - """ + surfaced as a 400 BadRequestError, not a 500.""" handler = MagicMock(spec=AnthropicServingMessages) handler.create_messages.side_effect = self._conversion_error() diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py index 8723c6ecc792..8c2e8623ad3c 100644 --- a/tests/entrypoints/cohere/test_serving_streaming.py +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -87,8 +87,7 @@ def _make_chunk( citations: list[Any] | None = None, ) -> str: """Build the ``data: {...}\\n\\n`` SSE frame the production code - consumes. - """ + consumes.""" delta: dict[str, Any] = {} if role is not None: delta["role"] = role @@ -136,8 +135,7 @@ async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect - the emitted SSE frames. - """ + the emitted SSE frames.""" request = CohereChatV2Request( model="m", messages=[{"role": "user", "content": "hi"}], stream=True ) diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 7f2f771107c0..668f11544858 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -50,8 +50,7 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit. - """ + endpoints back via the pipe, then exit.""" ctx = zmq.Context() try: in_sock = make_zmq_socket( @@ -358,8 +357,7 @@ def run_with_exception_capture(): def test_gather_actual_addresses_end_to_end(): """Each child binds ROUTER/PULL with a kernel-picked port and reports the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`. - """ + them via :py:meth:`gather_actual_addresses`.""" host = "127.0.0.1" num_servers = 4 @@ -409,8 +407,7 @@ def test_gather_actual_addresses_end_to_end(): @pytest.mark.timeout(30) def test_gather_actual_addresses_child_crash_before_report(): """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots. - """ + clear ``RuntimeError`` rather than hang or return ``None`` slots.""" host = "127.0.0.1" num_servers = 2 placeholder_inputs = [ @@ -447,8 +444,7 @@ def test_gather_actual_addresses_child_crash_before_report(): def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): """The Rust frontend command carries every non-default arg as JSON. - Credentials must not reach the log line. - """ + Credentials must not reach the log line.""" import subprocess as subprocess_mod from vllm.entrypoints.launchers.cli_args import make_arg_parser diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index cf2f678e40da..0bcca8cb41ac 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -186,8 +186,7 @@ def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): def test_passes_with_reasoning_parser(serve_parser): """Ensure validation passes if reasoning is enabled - with a reasoning parser - """ + with a reasoning parser""" args = serve_parser.parse_args( args=[ "--reasoning-parser", @@ -382,8 +381,7 @@ def test_default_chat_template_kwargs_invalid_json(serve_parser): ) def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises): """Ensure that users don't misuse --served-model-name and end up with the default - model tag instead of the one they intended to serve. - """ + model tag instead of the one they intended to serve.""" # Call the serve subparser args.insert(0, "serve") # Create a dummy config file if the test case includes it diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index b12bf3a60b8c..aac68334bf94 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """E2E test for mixing `prompt_embeds` with `audio_embeds` in a single -Chat Completions request. -""" +Chat Completions request.""" import json diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py index b72d94071425..c6ce165cf390 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -68,8 +68,7 @@ async def client(server): def image_url() -> str: """Stable real image as a data URL, kept identical across both the text and prompt_embeds requests so any output difference must come from - how the text content is delivered. - """ + how the text content is delivered.""" return encode_image_url(ImageAsset("stop_sign").pil_image) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py index 6d1c2f24e98f..56a7cf6a1312 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py @@ -145,8 +145,7 @@ def describe_image_messages( ) -> list[dict]: """Build the system + user messages used by the completions-with-image family of tests. *extra_image_fields* is merged into the top-level - image content block (for uuid / bad-key tests). - """ + image content block (for uuid / bad-key tests).""" image_block: dict = { "type": "image_url", "image_url": {"url": image_url}, @@ -176,8 +175,7 @@ async def complete_and_check( temperature: float = 0.0, ) -> str: """Run a chat completion and assert the output is non-empty. - Returns the content string. - """ + Returns the content string.""" chat_completion = await client.chat.completions.create( model=model_name, messages=messages, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index e0f71287441e..b83341cc5014 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -465,8 +465,7 @@ def test_system_message_warns_on_video(video_content): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500. - """ + so the API returns 400 instead of 500.""" with pytest.raises(Exception, match="json_schema.*must be provided"): ChatCompletionRequest( model=MODEL_NAME, @@ -520,8 +519,7 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): @pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500). - """ + TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): ChatCompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py index f1801701e3c5..b415fa116dac 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py @@ -81,8 +81,7 @@ async def test_chat_logit_bias_invalid(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_integer_key(client): """Test that a non-integer logit_bias key is rejected with a clean, - informative error instead of a raw 'invalid literal for int()' message. - """ + informative error instead of a raw 'invalid literal for int()' message.""" with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -102,8 +101,7 @@ async def test_chat_logit_bias_non_integer_key(client): @pytest.mark.asyncio async def test_chat_logit_bias_non_numeric_value(client): """Test that a non-numeric logit_bias value is rejected with a message - that names the specific offending token, not just a generic TypeError. - """ + that names the specific offending token, not just a generic TypeError.""" with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, @@ -122,8 +120,7 @@ async def test_chat_logit_bias_non_numeric_value(client): @pytest.mark.asyncio async def test_chat_logit_bias_multiple_non_integer_keys(client): """Test that ALL invalid logit_bias keys are reported together, - not just the first one encountered. - """ + not just the first one encountered.""" with pytest.raises(openai.BadRequestError) as excinfo: await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py index 9c0a4db68b1e..f4036c4374f4 100644 --- a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py +++ b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py @@ -54,8 +54,7 @@ def stop_sign_image_url(): async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI): """En -> es translation through TranslateGemma's bundled chat template's ``text`` content branch, which depends on ``source_lang_code`` / - ``target_lang_code`` extra fields being preserved on the content part. - """ + ``target_lang_code`` extra fields being preserved on the content part.""" completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ @@ -99,8 +98,7 @@ async def test_translategemma_image_extra_lang_code_fields( ``source_lang_code`` / ``target_lang_code`` extras must survive parsing and end up alongside the ``{"type": "image"}`` placeholder so the template's image branch (``content["type"] == 'image'``) renders the - right OCR-translation prompt. - """ + right OCR-translation prompt.""" completion = await client.chat.completions.create( model=MODEL_NAME, messages=[ diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py index 62f915b054c0..1eea6db48a7c 100644 --- a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -191,8 +191,7 @@ async def test_logprob_token_ids_stream_with_no_top_k(server): async def test_logprob_token_ids_default_behavior_unchanged(server): """Without `logprob_token_ids`, the response carries the natural top-k most-likely tokens. This guards against the new field accidentally - changing the default-path output. - """ + changing the default-path output.""" async with server.get_async_client() as client: resp = await client.chat.completions.create( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 73cffa76cfb5..27573203623e 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1406,8 +1406,7 @@ async def render_chat_request(request): @pytest.mark.asyncio async def test_serving_chat_data_parallel_rank_extraction(): """Test that data_parallel_rank is properly extracted from header and - passed to engine. - """ + passed to engine.""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() @@ -1729,8 +1728,7 @@ async def test_simple_chat(self, serving_chat, stream, include_reasoning): @pytest.mark.asyncio async def test_system_message_without_tools(self, serving_chat, stream): """Leading system message produces a developer message with - DeveloperContent (# Instructions header). - """ + DeveloperContent (# Instructions header).""" messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello"}, @@ -1752,8 +1750,7 @@ async def test_system_message_without_tools(self, serving_chat, stream): @pytest.mark.asyncio async def test_system_message_with_tools(self, serving_chat, stream, weather_tools): """Leading system message is folded into the developer message - alongside tool definitions. - """ + alongside tool definitions.""" messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather?"}, @@ -2168,8 +2165,7 @@ async def test_non_tool_reasoning_empty_content_list(self, serving_chat): @pytest.mark.asyncio async def test_tool_choice_validation_without_parser(): """Test that tool_choice='required' or named tool without tool_parser - returns an appropriate error message. - """ + returns an appropriate error message.""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index b7ed422223bd..ecff3f8ee6be 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -151,8 +151,7 @@ async def client(request, server, server_with_auto_reasoning_config): @pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True) async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): """Test that mixed requests (some with thinking_token_budget, some without) - complete successfully without errors. - """ + complete successfully without errors.""" response_with_budget = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 1fee99be5762..46f4fe452cb9 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -521,8 +521,7 @@ async def mock_generate(*args, **kwargs): def test_json_schema_response_format_missing_schema(): """When response_format type is 'json_schema' but the json_schema field is not provided, request construction should raise a validation error - so the API returns 400 instead of 500. - """ + so the API returns 400 instead of 500.""" with pytest.raises(Exception, match="json_schema.*must be provided"): CompletionRequest( model=MODEL_NAME, @@ -585,8 +584,7 @@ def test_negative_prompt_token_ids_flat(): def test_logprobs_minus_one_allowed(): """logprobs=-1 means "return all logprobs". The sampling layer and the chat top_logprobs / prompt_logprobs validators all accept -1, so the completion - logprobs validator must accept it too instead of rejecting it as negative. - """ + logprobs validator must accept it too instead of rejecting it as negative.""" request = CompletionRequest( model=MODEL_NAME, prompt="Test prompt", @@ -746,8 +744,7 @@ def test_bounded_prompt_embeds_list_allowed(self, monkeypatch): @pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a - TypeError from the mode='before' comparison (which surfaces as HTTP 500). - """ + TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): CompletionRequest( model=MODEL_NAME, diff --git a/tests/entrypoints/openai/parser/test_harmony_render_parity.py b/tests/entrypoints/openai/parser/test_harmony_render_parity.py index ce6e4c80f89d..5cb446122361 100644 --- a/tests/entrypoints/openai/parser/test_harmony_render_parity.py +++ b/tests/entrypoints/openai/parser/test_harmony_render_parity.py @@ -43,8 +43,7 @@ def _system(): class TestResponseInputToHarmonyRenderParity: """Each test drives the same conversation through both APIs and asserts - identical Harmony messages and rendered token sequences. - """ + identical Harmony messages and rendered token sequences.""" # ----------------------------------------------------------------------- # Single-message cases @@ -52,8 +51,7 @@ class TestResponseInputToHarmonyRenderParity: def test_developer_message(self): """Both APIs must render developer messages identically using - DeveloperContent (with the '# Instructions' header). - """ + DeveloperContent (with the '# Instructions' header).""" chat_msgs = parse_chat_input_to_harmony_message( {"role": "developer", "content": "Be concise."} ) @@ -248,8 +246,7 @@ def test_tool_output(self): def test_reasoning_combined_with_function_call(self): """Chat API packs reasoning + tool_calls into one dict; responses API represents them as two separate items. Both must produce the same two - Harmony messages in the same order: analysis then commentary. - """ + Harmony messages in the same order: analysis then commentary.""" chat_msgs = parse_chat_input_to_harmony_message( { "role": "assistant", @@ -532,8 +529,7 @@ def test_multi_turn_two_tool_calls_with_reasoning_between(self): def test_completed_turns_drop_reasoning(self): """Validates that reasoning from completed turns is dropped, while reasoning from the current in-progress tool-call turn is preserved - in both chat completions and responses previous_input_messages. - """ + in both chat completions and responses previous_input_messages.""" first_turn_reasoning = "FIRST_TURN_REASONING" second_turn_reasoning = "SECOND_TURN_REASONING" diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index 816e879342e4..4c1dfe52e033 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -323,8 +323,7 @@ def test_assistant_message_with_empty_tool_call_arguments(self, parse_function): def test_system_message(self, parse_function): """Test parsing system messages, which are parsed into developer messages - with DeveloperContent. - """ + with DeveloperContent.""" chat_msg = { "role": "system", "content": "You are a helpful assistant", diff --git a/tests/entrypoints/openai/responses/test_function_call_parsing.py b/tests/entrypoints/openai/responses/test_function_call_parsing.py index c83fe0f61cce..5b70c7819a0e 100644 --- a/tests/entrypoints/openai/responses/test_function_call_parsing.py +++ b/tests/entrypoints/openai/responses/test_function_call_parsing.py @@ -12,8 +12,7 @@ def test_function_call_dict_converted_to_object(): """Test that function_call dictionaries are correctly parsed into - ResponseFunctionToolCall objects. - """ + ResponseFunctionToolCall objects.""" # Create a request with function_call as dict request_data = { "model": "gpt-oss", @@ -330,8 +329,7 @@ def test_validator_handles_empty_iterator(): def test_assistant_string_content_stays_easyinput(): """EasyInput assistant message with plain string content is not - coerced into a ResponseOutputMessage. - """ + coerced into a ResponseOutputMessage.""" request_data = { "model": "test-model", "input": [ @@ -353,8 +351,7 @@ def test_assistant_string_content_stays_easyinput(): def test_assistant_output_style_content_coerced(): """Assistant message whose content is output-message-shaped (list of - output_text items) should be coerced to ResponseOutputMessage. - """ + output_text items) should be coerced to ResponseOutputMessage.""" request_data = { "model": "test-model", "input": [ diff --git a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py index 2485983bf3ae..e2be8477e435 100644 --- a/tests/entrypoints/openai/responses/test_response_input_to_harmony.py +++ b/tests/entrypoints/openai/responses/test_response_input_to_harmony.py @@ -67,8 +67,7 @@ def test_no_type_key_defaults_to_message_branch(self): def test_system_message(self): """System messages carry developer instructions and must be rendered - as developer messages with DeveloperContent. - """ + as developer messages with DeveloperContent.""" msg = response_input_to_harmony( {"type": "message", "role": "system", "content": "Be helpful."}, prev_responses=[], @@ -90,8 +89,7 @@ def test_assistant_message_gets_final_channel(self): def test_developer_message_gets_instructions_prefix(self): """Developer messages must use DeveloperContent which adds the - '# Instructions' header the model was trained on. - """ + '# Instructions' header the model was trained on.""" msg = response_input_to_harmony( {"type": "message", "role": "developer", "content": "Be concise."}, prev_responses=[], @@ -121,8 +119,7 @@ def test_message_with_array_content(self): def test_developer_message_array_content_concatenated(self): """Array content in developer messages is flattened and rendered - via DeveloperContent with the '# Instructions' header. - """ + via DeveloperContent with the '# Instructions' header.""" msg = response_input_to_harmony( { "type": "message", @@ -213,8 +210,7 @@ def test_function_call_output_channel_recipient_and_author_name(self): def test_function_call_output_uses_most_recent_matching_call(self): """When multiple prev_responses share a call_id, the last one wins - because the search is reversed. - """ + because the search is reversed.""" earlier = ResponseFunctionToolCall( id="fc_old", call_id="call_test", diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 1e59c8224429..967ddacda6de 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -839,13 +839,11 @@ def test_merge_chain_breaks(self, items, num_expected_messages): class TestConstructInputMessagesInstructionsLeak: """Regression tests for #37697: instructions from a prior response - should NOT leak through previous_response_id. - """ + should NOT leak through previous_response_id.""" def test_old_instructions_stripped_from_prev_msg(self): """System message in prev_msg must be dropped so the new request's - instructions are the only system message in the conversation. - """ + instructions are the only system message in the conversation.""" prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "What is 2+2?"}, @@ -862,8 +860,7 @@ def test_old_instructions_stripped_from_prev_msg(self): def test_no_instructions_in_new_request(self): """If the new request has no instructions, old ones should still - be stripped -- they must not carry over. - """ + be stripped -- they must not carry over.""" prev = [ {"role": "system", "content": "old instructions"}, {"role": "user", "content": "Hi"}, diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 1f9d95ea2e1c..a445774fccb7 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -1163,8 +1163,7 @@ def test_commentary_with_function_recipient_not_preamble(self) -> None: def test_preamble_done_emits_text_done_events(self) -> None: """Completed preamble should emit text done + content_part done + - output_item done, same shape as final channel. - """ + output_item done, same shape as final channel.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1185,8 +1184,7 @@ def test_preamble_done_emits_text_done_events(self) -> None: def test_commentary_with_recipient_no_preamble_done(self) -> None: """Commentary + recipient='functions.X' should route to function call - done, not preamble done. - """ + done, not preamble done.""" from vllm.entrypoints.openai.responses.streaming_events import ( emit_previous_item_done_events, ) @@ -1332,8 +1330,7 @@ def mock_parse_delta(**kwargs): class TestStreamingReasoningToContentTransition: """Tests for _process_simple_streaming_events reasoning-to-content transition, specifically the fix for mixed deltas that carry both - reasoning and content simultaneously. - """ + reasoning and content simultaneously.""" @pytest.mark.asyncio async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta( @@ -1343,8 +1340,7 @@ async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta( and content set (e.g. reasoning end and content start in the same chunk), the trailing reasoning text must be emitted as a ResponseReasoningTextDeltaEvent and included in the - ResponseReasoningTextDoneEvent text. - """ + ResponseReasoningTextDoneEvent text.""" monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1413,8 +1409,7 @@ async def test_transition_without_mixed_delta_no_extra_reasoning_event( self, monkeypatch ): """When the transition from reasoning to content is clean (no mixed - delta), no extra reasoning delta event should be emitted. - """ + delta), no extra reasoning delta event should be emitted.""" monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() @@ -1475,8 +1470,7 @@ async def result_generator(): async def test_reasoning_only_stream_no_content(self, monkeypatch): """When the stream has only reasoning deltas and no content, the reasoning done event should be emitted at finalization with the - full accumulated text, and no text delta events should appear. - """ + full accumulated text, and no text delta events should appear.""" monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) serving = _make_serving_instance_with_reasoning() diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index 1880c06392d1..ba55a58e7d3c 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -143,8 +143,7 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_streaming_logprobs(client: OpenAI, model_name: str): """Test that streaming with logprobs returns valid logprob data on - output_text.delta events and that top_logprobs has the requested count. - """ + output_text.delta events and that top_logprobs has the requested count.""" response = await client.responses.create( model=model_name, input="Say hello.", diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index 2e9b47133dbb..e1a8bf46f640 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -124,8 +124,7 @@ def test_parallel_tool_calls(self): def test_split_name_and_args_same_index(self): """Regression: parsers like KimiK2 emit name and args as separate - DeltaToolCalls at the same index within one DeltaMessage. - """ + DeltaToolCalls at the same index within one DeltaMessage.""" tc_name = _make_tool_call(0, name="get_weather") tc_args = _make_tool_call(0, arguments='{"city":"SF"}') delta = DeltaMessage(tool_calls=[tc_name, tc_args]) @@ -141,8 +140,7 @@ def test_split_name_and_args_same_index(self): def test_reasoning_to_content_transition(self): """Regression: the old special case in emit_delta handled this; - now split_delta handles it generically. - """ + now split_delta handles it generically.""" processor = SimpleStreamingEventProcessor() _run_through_processor(processor, DeltaMessage(reasoning="think")) assert processor.state.current_state == _StateType.REASONING diff --git a/tests/entrypoints/openai/test_render_token_offsets.py b/tests/entrypoints/openai/test_render_token_offsets.py index 6ef66ee37bc0..a2e66b7bd6ce 100644 --- a/tests/entrypoints/openai/test_render_token_offsets.py +++ b/tests/entrypoints/openai/test_render_token_offsets.py @@ -25,8 +25,7 @@ def _model_config() -> Mock: def test_completion_flag_forwarded_to_tok_params(): """build_tok_params must forward return_token_offsets, defaulting to False (zero behavioral change for existing callers) and coercing JSON - null to False via the bool() guard. - """ + null to False via the bool() guard.""" cfg = _model_config() default = CompletionRequest(model="m", prompt="hi") @@ -41,8 +40,7 @@ def test_completion_flag_forwarded_to_tok_params(): def test_chat_flag_forwarded_to_tok_params(): """Chat build_tok_params has its own (max_completion_tokens) branch, so - its return_token_offsets forwarding is verified independently. - """ + its return_token_offsets forwarding is verified independently.""" cfg = _model_config() messages = [{"role": "user", "content": "hi"}] @@ -66,8 +64,7 @@ def test_generate_request_token_offsets_default_none(): def test_generate_request_token_offsets_survive_json_round_trip(): """GenerateRequest crosses the disagg serialization boundary; the - tuple[int, int] offsets must survive model_dump and re-validate. - """ + tuple[int, int] offsets must survive model_dump and re-validate.""" req = GenerateRequest( token_ids=[10, 20], sampling_params=SamplingParams(), diff --git a/tests/entrypoints/pooling/embed/test_cohere_online.py b/tests/entrypoints/pooling/embed/test_cohere_online.py index 29276cb0501f..4964d99e0c66 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_online.py +++ b/tests/entrypoints/pooling/embed/test_cohere_online.py @@ -111,8 +111,7 @@ def test_basic_embed(server: RemoteOpenAIServer, model_name: str): def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str): """An input_type not defined in the model's prompt config should be - rejected with a 400 error. - """ + rejected with a 400 error.""" body = { "model": model_name, "input_type": "nonexistent_type", @@ -139,8 +138,7 @@ def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str): """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same - float embeddings for a generic model. - """ + float embeddings for a generic model.""" texts = ["hello world"] v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"]) v1 = _openai_embed(server, model_name, texts) diff --git a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py index 5608cbcb4e06..3e168adbdba9 100644 --- a/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py +++ b/tests/entrypoints/pooling/embed/test_cohere_openai_parity.py @@ -76,8 +76,7 @@ def test_single_text_parity(server: RemoteOpenAIServer): def test_batch_parity(server: RemoteOpenAIServer): """A batch of texts should produce equivalent embeddings via both APIs, - in the same order. - """ + in the same order.""" texts = [ "machine learning", "deep learning", diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index b002aae8db31..f06e43d143ea 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -241,8 +241,7 @@ def test_truncate_none(self): def test_truncate_none_with_max_tokens(self): """truncate=NONE should NOT set truncate_prompt_tokens; the - max_tokens limit is enforced separately via _check_max_tokens. - """ + max_tokens limit is enforced separately via _check_max_tokens.""" req = self._make_request(truncate="NONE", max_tokens=10) tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req) assert tokens is None diff --git a/tests/entrypoints/pooling/embed/test_online_long_text.py b/tests/entrypoints/pooling/embed/test_online_long_text.py index aeb82336fa48..ab3633cd3d0f 100644 --- a/tests/entrypoints/pooling/embed/test_online_long_text.py +++ b/tests/entrypoints/pooling/embed/test_online_long_text.py @@ -240,8 +240,7 @@ async def test_long_text_embedding_1500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~1500 character long text - (~1028 tokens, exceeding 512 token limit). - """ + (~1028 tokens, exceeding 512 token limit).""" # Verify text length # Verify text has sufficient word count (approximately 1500 words) word_count = len(LONG_TEXT_1500_WORDS.split()) @@ -289,8 +288,7 @@ async def test_long_text_embedding_2500_chars( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test embedding processing for ~2500 character long text - (~2048 tokens, requiring multiple chunks). - """ + (~2048 tokens, requiring multiple chunks).""" # Verify text length # Verify text has sufficient word count (approximately 2500 words) word_count = len(LONG_TEXT_2500_WORDS.split()) @@ -382,8 +380,7 @@ async def test_chunked_vs_normal_consistency( client_with_chunked_processing: openai.AsyncOpenAI, model_name: str ): """Test consistency between chunked and - normal processing (using short text). - """ + normal processing (using short text).""" # Use a short text within the 512 token limit short_text = ( "Artificial intelligence technology is changing our world, " diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 24d4f9e6d730..91c853f91098 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for Cohere embed protocol: build_typed_embeddings and its -underlying packing helpers, plus Cohere-specific serving helpers. -""" +underlying packing helpers, plus Cohere-specific serving helpers.""" import struct diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index db702b4b07be..d7e0298ac3ab 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -454,8 +454,7 @@ async def test_derender_completion_empty_generate_responses(client): @pytest.mark.asyncio async def test_derender_completion_logprobs(client): """token_id:N placeholders in logprobs are resolved; CompletionLogProbs - flat-list structure is returned with non-empty tokens and text_offsets. - """ + flat-list structure is returned with non-empty tokens and text_offsets.""" gr1 = await _render_completion(client, "Hello world") ids1 = gr1["token_ids"][:3] token_id = ids1[0] diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 6bcede31992d..adece09c0a1a 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -121,8 +121,7 @@ async def _disagg( def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: """[(name, json normalized args)] so key ordering / whitespace don't - cause false negatives. - """ + cause false negatives.""" return [ (tc["function"]["name"], json.loads(tc["function"]["arguments"])) for tc in (response_choice["message"].get("tool_calls") or []) @@ -144,8 +143,7 @@ async def _run_parity_case( client: httpx.AsyncClient, messages: list[dict], **extra ) -> tuple[dict, dict]: """Run the coupled request then feed its generated tokens into the - disaggregated derender endpoint. Returns (coupled, disagg). - """ + disaggregated derender endpoint. Returns (coupled, disagg).""" coupled = await _coupled(client, messages, **extra) ch = coupled["choices"][0] chat_request = {"model": MODEL, "messages": messages, **extra} @@ -205,8 +203,7 @@ async def test_parity_tool_call(client): @pytest.mark.asyncio async def test_parity_reasoning_and_tool_call(client): """Combined reasoning + tool call parity means the highest drift risk - since it exercises both parser branches on the same output. - """ + since it exercises both parser branches on the same output.""" messages = [{"role": "user", "content": "What's the weather in Paris?"}] coupled, disagg = await _run_parity_case( client, diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index bb6a521ef14b..e21343553dbd 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -455,8 +455,7 @@ async def test_chat_completion_render_with_sampling_params(client): @pytest.mark.asyncio async def test_completion_render_emits_token_offsets(client): """With return_token_offsets, /v1/completions/render returns per-token - (start, end) char offsets aligned with token_ids. - """ + (start, end) char offsets aligned with token_ids.""" prompt = "Hello, world." response = await client.post( "/v1/completions/render", @@ -481,8 +480,7 @@ async def test_completion_render_emits_token_offsets(client): @pytest.mark.asyncio async def test_completion_render_default_no_token_offsets(client): """Without the flag, token_offsets must be null (existing responses - unchanged). - """ + unchanged).""" response = await client.post( "/v1/completions/render", json={ @@ -499,8 +497,7 @@ async def test_completion_render_default_no_token_offsets(client): @pytest.mark.asyncio async def test_chat_render_emits_token_offsets(client): """With return_token_offsets, /v1/chat/completions/render returns - per-token offsets relative to the templated prompt string. - """ + per-token offsets relative to the templated prompt string.""" response = await client.post( "/v1/chat/completions/render", json={ @@ -674,8 +671,7 @@ async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags client, ): """With a ``{% generation %}``-enabled template, the mask marks assistant - tokens and the masked tokens decode to the assistant content. - """ + tokens and the masked tokens decode to the assistant content.""" response = await client.post( "/v1/chat/completions/render", json={ diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index 1dcf93330edf..e51530ce456f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -25,8 +25,7 @@ def _base_payload() -> dict: def test_omitted_max_tokens_is_not_provided(): """Body without ``max_tokens`` must surface as 'not provided' so the - server can apply its own default instead of the dataclass 16. - """ + server can apply its own default instead of the dataclass 16.""" req = GenerateRequest.model_validate(_base_payload()) # SamplingParams' dataclass default leaks through the parsed instance — # this is exactly the bug the server-side defaulting works around. @@ -36,8 +35,7 @@ def test_omitted_max_tokens_is_not_provided(): def test_explicit_max_tokens_is_provided(): """Even when the client picks the same value as the dataclass default, - it must register as explicitly set so the server won't override it. - """ + it must register as explicitly set so the server won't override it.""" payload = _base_payload() payload["sampling_params"] = {"max_tokens": 16} req = GenerateRequest.model_validate(payload) @@ -69,8 +67,7 @@ def test_json_roundtrip_preserves_provided_keys(): def test_internal_instance_construction_treats_all_as_provided(): """When internal callers build ``GenerateRequest`` from a pre-resolved ``SamplingParams`` instance, every field is considered explicitly set - so server-side defaulting can't clobber values resolved upstream. - """ + so server-side defaulting can't clobber values resolved upstream.""" sp = SamplingParams(max_tokens=500, temperature=0.0) req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp) assert req.is_sampling_param_provided("max_tokens") diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py index cb14260ab82c..426d1ede745d 100644 --- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -22,8 +22,7 @@ def test_sanitize_message(): class TestSanitizeMessageFilePaths: """sanitize_message should also strip file paths and traceback - frames, not just memory addresses - see #31683. - """ + frames, not just memory addresses - see #31683.""" def test_strips_traceback_style_frame(self): msg = ( @@ -45,8 +44,7 @@ def test_strips_arbitrary_absolute_path(self): def test_strips_single_parent_container_path(self): """Regression: /app/server.py and /workspace/server.py (common in - container deployments) were missed by the original {2,} quantifier. - """ + container deployments) were missed by the original {2,} quantifier.""" assert "/app/" not in sanitize_message("Error in /app/server.py") assert "/workspace/" not in sanitize_message("Error in /workspace/server.py") diff --git a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py index bd4e87a15543..942dd95c182a 100644 --- a/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py +++ b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py @@ -24,8 +24,7 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: """Minimal stand-in for a FastAPI Request - just enough for the - handler to read req.app.state.args.log_error_stack. - """ + handler to read req.app.state.args.log_error_stack.""" return SimpleNamespace( app=SimpleNamespace( state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack)) @@ -36,8 +35,7 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: class TestValidationErrorParamFallback: """Ensure `param` falls back to the Pydantic error's `loc` when no - custom VLLMValidationError context is present. - """ + custom VLLMValidationError context is present.""" @pytest.mark.parametrize( ("error_type", "msg"), @@ -61,8 +59,7 @@ async def test_param_falls_back_to_loc(self, error_type: str, msg: str): async def test_param_fallback_does_not_crash_on_non_dict_error(self): """Schemathesis fuzzing found that errors[0] isn't always a dict. The fallback must not crash in that case - it should just leave - param as None instead of raising. - """ + param as None instead of raising.""" exc = RequestValidationError(["some unexpected non-dict error"]) response = await validation_exception_handler(_fake_request(), exc) @@ -100,8 +97,7 @@ def test_strips_internal_markers(self, loc, expected): def test_all_internal_falls_back_to_raw_join(self): """If every loc segment looks internal, fall back to the raw - dot-join rather than returning an empty string. - """ + dot-join rather than returning an empty string.""" loc = ("function-wrap[__log_extra_fields__()]",) assert clean_loc_for_param(loc) == "function-wrap[__log_extra_fields__()]" @@ -109,8 +105,7 @@ def test_all_internal_falls_back_to_raw_join(self): class TestValidationErrorDoesNotLeakServerPaths: """FastAPI sets endpoint_file/line/function/path on the exception during routing, not in the constructor - so we set them manually - here to reproduce the real leak. See #31683. - """ + here to reproduce the real leak. See #31683.""" @pytest.mark.asyncio async def test_handler_strips_endpoint_file_context(self): diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index 79356fc951b7..8a860dc6251a 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -245,8 +245,7 @@ async def test_multiple_lora_adapters( client: openai.AsyncOpenAI, tmp_path, qwen3_lora_files ): """Validate that many loras can be dynamically registered and inferenced - with concurrently - """ + with concurrently""" # This test file configures the server with --max-cpu-loras=2 and this test # will concurrently load 10 adapters, so it should flex the LRU cache diff --git a/tests/entrypoints/serve/middleware/test_authentication_middleware.py b/tests/entrypoints/serve/middleware/test_authentication_middleware.py index ee22407fdd0b..528261e644db 100644 --- a/tests/entrypoints/serve/middleware/test_authentication_middleware.py +++ b/tests/entrypoints/serve/middleware/test_authentication_middleware.py @@ -72,8 +72,7 @@ def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]): @pytest.fixture(params=get_args(SupportedTask)) def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: """For each supported task, build an app with only that task's routers, - extract all routes, and return the task name and routes. - """ + extract all routes, and return the task name and routes.""" task = request.param # Enable development mode to register all routes (including dev-only routes). monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") @@ -99,8 +98,7 @@ def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]] def test_auto_discovered_protected_routes_require_auth(task_routes): """For every auto-discovered route that starts with a guarded prefix, - verify that authentication is enforced. - """ + verify that authentication is enforced.""" task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) @@ -134,8 +132,7 @@ def test_auto_discovered_protected_routes_require_auth(task_routes): def test_auto_discovered_unprotected_routes_no_auth(task_routes): """For every auto-discovered route that does NOT start with a guarded - prefix, verify that no authentication is required. - """ + prefix, verify that no authentication is required.""" task, routes = task_routes app = _create_app_with_mock_routes(routes) client = TestClient(app) diff --git a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py index b6137dae93bd..20b917400b3c 100644 --- a/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py +++ b/tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py @@ -114,8 +114,7 @@ def _clear_env_vars(self): def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path): """Test customer scenario: script functions automatically override - framework defaults. - """ + framework defaults.""" try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -176,8 +175,7 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_customer_decorator_usage(self, monkeypatch, tmp_path): """Test customer scenario: using @custom_ping_handler and - @custom_invocation_handler decorators. - """ + @custom_invocation_handler decorators.""" try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -233,8 +231,7 @@ async def my_invoke(request: Request): def test_handler_priority_order(self, monkeypatch, tmp_path): """Test priority: @custom_ping_handler/@custom_invocation_handler - decorators vs script functions. - """ + decorators vs script functions.""" try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -307,8 +304,7 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_environment_variable_script_loading(self, monkeypatch, tmp_path): """Test that environment variables correctly specify script location - and loading. - """ + and loading.""" try: from model_hosting_container_standards.sagemaker.config import ( SageMakerEnvVars, @@ -420,8 +416,7 @@ async def test_framework_default_handlers(self): def test_handler_env_var_override(self, monkeypatch, tmp_path): """Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER - environment variable overrides. - """ + environment variable overrides.""" try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, @@ -515,8 +510,7 @@ async def custom_sagemaker_invocation_handler(request: Request): def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path): """Test that environment variables have highest priority over decorators - and script functions for both ping and invocation handlers. - """ + and script functions for both ping and invocation handlers.""" try: from model_hosting_container_standards.common.fastapi.config import ( FastAPIEnvVars, diff --git a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py index 49e181e36e3b..1d114bc8b81f 100644 --- a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py +++ b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py @@ -51,8 +51,7 @@ async def test_tokenize_then_chat_completion_with_image( local_asset_server, ): """Tokenize a multimodal message, then send the same message to chat - completions. The chat completion must succeed (not 500). - """ + completions. The chat completion must succeed (not 500).""" image_url = local_asset_server.url_for("stop_sign.jpg") messages = [ { diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index 051978b60c47..a41f39187f4f 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -36,13 +36,11 @@ def test_should_include_usage_force_enables_continuous_usage(stream_options, exp class TestGetMaxTokens: """Tests for get_max_tokens() to ensure generation_config's max_tokens acts as a default when from model author, and as a ceiling when - explicitly set by the user. - """ + explicitly set by the user.""" def test_default_sampling_params_used_when_no_request_max_tokens(self): """When user doesn't specify max_tokens, generation_config default - should apply. - """ + should apply.""" result = get_max_tokens( max_model_len=24000, max_tokens=None, @@ -53,8 +51,7 @@ def test_default_sampling_params_used_when_no_request_max_tokens(self): def test_request_max_tokens_not_capped_by_default_sampling_params(self): """When user specifies max_tokens in request, model author's - generation_config max_tokens must NOT cap it (fixes #34005). - """ + generation_config max_tokens must NOT cap it (fixes #34005).""" result = get_max_tokens( max_model_len=24000, max_tokens=5000, @@ -97,8 +94,7 @@ def test_max_model_len_still_caps_output(self): def test_request_max_tokens_smaller_than_default(self): """When user explicitly requests fewer tokens than gen_config default, - that should be respected. - """ + that should be respected.""" result = get_max_tokens( max_model_len=24000, max_tokens=512, diff --git a/tests/entrypoints/serve/utils/test_request_logger.py b/tests/entrypoints/serve/utils/test_request_logger.py index 5df3c3cd8425..717ff190b3a7 100644 --- a/tests/entrypoints/serve/utils/test_request_logger.py +++ b/tests/entrypoints/serve/utils/test_request_logger.py @@ -239,8 +239,7 @@ def test_request_logger_log_outputs_integration(): def test_streaming_complete_logs_full_text_content(): """Test that streaming complete logging includes - full accumulated text, not just token count. - """ + full accumulated text, not just token count.""" mock_logger = MagicMock() with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger): diff --git a/tests/entrypoints/serve/utils/test_sse_keep_alive.py b/tests/entrypoints/serve/utils/test_sse_keep_alive.py index 6c61080f8f51..519b9bb4a129 100644 --- a/tests/entrypoints/serve/utils/test_sse_keep_alive.py +++ b/tests/entrypoints/serve/utils/test_sse_keep_alive.py @@ -17,8 +17,7 @@ def test_disabled_returns_same_generator(): """For 0, negative, and NaN intervals the exact same generator object must be returned (identity), proving no wrapper or per-chunk overhead when - disabled. - """ + disabled.""" async def gen(): yield "x" @@ -34,8 +33,7 @@ async def gen(): @pytest.mark.asyncio async def test_keep_alive_during_silent_start(): """A silent start (queue wait / prefill) emits keep-alive comments before - the first real chunk. - """ + the first real chunk.""" release = asyncio.Event() async def gen(): @@ -57,8 +55,7 @@ async def gen(): @pytest.mark.asyncio async def test_no_keep_alive_when_streaming_is_fast(): """If chunks arrive faster than the interval, no comment is emitted and all - data is forwarded. - """ + data is forwarded.""" async def gen(): for chunk in ["a", "b", "c"]: @@ -72,8 +69,7 @@ async def gen(): @pytest.mark.asyncio async def test_keep_alive_between_chunks_does_not_drop_data(): """Keep-alives interleaved with real chunks must never drop or reorder - data. - """ + data.""" release_b = asyncio.Event() release_c = asyncio.Event() @@ -139,8 +135,7 @@ async def gen(): @pytest.mark.asyncio async def test_aclose_closes_upstream(): """Closing a wrapper that is mid-stream must complete promptly and close - the upstream generator (regression for the earlier leaked-task hang). - """ + the upstream generator (regression for the earlier leaked-task hang).""" finalized = asyncio.Event() async def gen(): @@ -159,8 +154,7 @@ async def gen(): @pytest.mark.asyncio async def test_client_disconnect_cancels_upstream(): """Cancelling the consumer mid-stream must cancel the pending anext and - close the upstream generator so its cleanup runs. - """ + close the upstream generator so its cleanup runs.""" finalized = asyncio.Event() async def gen(): @@ -187,8 +181,7 @@ async def gen(): @pytest.mark.asyncio async def test_streaming_response_emits_comment(): """A keep-alive comment must reach the client through the real - StreamingResponse used by the routers. - """ + StreamingResponse used by the routers.""" release = asyncio.Event() async def gen(): @@ -227,8 +220,7 @@ async def receive(): async def test_streaming_response_disconnect_closes_upstream(): """On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette cancels streaming and the wrapper's finally must close the upstream - generator. - """ + generator.""" finalized = asyncio.Event() async def gen(): diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index b9ea18bbaea9..7fb7a5522d29 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -295,8 +295,7 @@ async def test_session_update_invalid_model_returns_error(model_name): @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_commit_without_session_update_returns_error(model_name): """Test that committing before validating the model returns an error - and does not fall through to processing. - """ + and does not fall through to processing.""" server_args = ["--enforce-eager", "--max-model-len", "2048"] if model_name.startswith("mistralai"): diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index e09687425bb9..c92813ea6612 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -51,8 +51,7 @@ async def generate_mock_outputs( num_turns, prompt_token_counts, output_token_counts, cached_token_counts=None ): """Generate a sequence of mock RequestOutput objects to simulate multiple - turns. - """ + turns.""" if cached_token_counts is None: cached_token_counts = [0] * num_turns @@ -271,8 +270,7 @@ def test_reasoning_tokens_counting(): def test_preamble_tokens_not_counted_as_reasoning(): """Preambles (commentary with no recipient) are visible user text, - not hidden reasoning. They must NOT inflate num_reasoning_tokens. - """ + not hidden reasoning. They must NOT inflate num_reasoning_tokens.""" context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=0) @@ -289,8 +287,7 @@ def test_preamble_tokens_not_counted_as_reasoning(): def test_commentary_with_recipient_counted_as_reasoning(): """Commentary directed at a tool (recipient != None) is hidden from - the user, so it should still count as reasoning tokens. - """ + the user, so it should still count as reasoning tokens.""" context, parser = make_harmony_context() parser.enqueue_chunk_result(reasoning_token_count=3) @@ -349,8 +346,7 @@ async def test_single_turn_no_tool_output(): @pytest.mark.asyncio async def test_negative_tool_tokens_edge_case(): """Test edge case where calculation could result in negative tool - tokens. We should log an error and clamp the value to 0. - """ + tokens. We should log an error and clamp the value to 0.""" # Use patch to check if logger.error was called with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log: context, _ = make_harmony_context(available_tools=["browser"]) @@ -690,8 +686,7 @@ def create_simple_context_output( finished=True, ): """Helper to create a RequestOutput with customizable text for - SimpleContext tests. - """ + SimpleContext tests.""" if token_ids is None: token_ids = [] return RequestOutput( @@ -875,8 +870,7 @@ def test_simple_context_final_output(): def test_simple_context_output_messages_empty_text_with_tokens(): """output_messages should be returned when tokens exist even if text is - empty (e.g. special tokens). - """ + empty (e.g. special tokens).""" context = SimpleContext() context.append_output( create_simple_context_output( @@ -894,8 +888,7 @@ def test_simple_context_output_messages_empty_text_with_tokens(): def test_simple_context_output_messages_no_mutation(): """Each call to output_messages returns a fresh list; callers can't - corrupt internal state. - """ + corrupt internal state.""" context = SimpleContext() context.append_output( create_simple_context_output( diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 9b36ab398d3d..e5c3616d6e65 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -121,8 +121,7 @@ def test_get_world_size_tp1(vllm_runner): @create_new_process_for_each_test() def test_init_weight_transfer_engine_calls_engine(vllm_runner): """Test that init_weight_transfer_engine calls the engine's - init_transfer_engine method. - """ + init_transfer_engine method.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 7e346cc76721..15b9835d6cf5 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -86,8 +86,7 @@ def _all_kernel_classes() -> list[type]: def _probe(cls: type): """A bare kernel instance with a plausible config, so input_quant_key() - can be queried without the hardware-gated constructor. - """ + can be queried without the hardware-gated constructor.""" obj = cls.__new__(cls) # type: ignore[call-overload] if issubclass(cls, NvFp4LinearKernel): obj.config = NvFp4LinearLayerConfig() diff --git a/tests/kernels/attention/test_amx_mla.py b/tests/kernels/attention/test_amx_mla.py index 7dab61f1c4d1..f5fa2fad16c2 100644 --- a/tests/kernels/attention/test_amx_mla.py +++ b/tests/kernels/attention/test_amx_mla.py @@ -247,8 +247,7 @@ def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens): class _FakeLinear: """Minimal duck-typed stand-in for kv_b_proj: only .weight and - .quant_method are read by get_and_maybe_dequant_weights(). - """ + .quant_method are read by get_and_maybe_dequant_weights().""" def __init__(self, weight: torch.Tensor): self.weight = weight @@ -319,8 +318,7 @@ def __init__(self, decode=None, prefill=None, max_seq_len=0): def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl would produce via layer.W_UK_T) and must reproduce plain per-head MLA - decode attention through W_UK/W_UV. - """ + decode attention through W_UK/W_UV.""" set_random_seed(1) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( @@ -386,8 +384,7 @@ def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config): """forward_mha receives raw unabsorbed Q and must attend correctly for a prefill batch with a mix of fresh and continued (cached-prefix) - sequences. - """ + sequences.""" set_random_seed(2) num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 impl, kv_b_weight = _make_amx_mla_impl( diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 0511fe7b07f4..53867a5edf93 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -616,8 +616,7 @@ def test_non_causal_autoselect_backend(): ) def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str): """FlashAttentionBackend must not claim support for kv_cache dtypes - that it cannot handle. - """ + that it cannot handle.""" from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype) @@ -628,8 +627,7 @@ def test_flash_attn_accepts_handled_fp8_variants( kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch ): """FlashAttentionBackend must accept the two fp8 dtypes it can actually - handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'. - """ + handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'.""" import vllm.v1.attention.backends.fa_utils as fa_utils_mod from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend diff --git a/tests/kernels/attention/test_cutlass_mla_decode.py b/tests/kernels/attention/test_cutlass_mla_decode.py index d56d72adba79..7a5011b5caec 100644 --- a/tests/kernels/attention/test_cutlass_mla_decode.py +++ b/tests/kernels/attention/test_cutlass_mla_decode.py @@ -225,8 +225,7 @@ def test_cutlass_mla_decode_cross_layer_view(): """The kernel must read the cache's page-dim stride instead of assuming pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must - match a contiguous cache holding the same data exactly. - """ + match a contiguous cache holding the same data exactly.""" device = torch.device("cuda:0") torch.set_default_dtype(torch.bfloat16) torch.set_default_device(device) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index b786fac095d8..c6cba8d530e4 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -162,8 +162,7 @@ def test_fast_decode_plan_importable() -> None: @torch.inference_mode def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: """On the first call fast_plan_decode must route through self.plan() and - flip vllm_first_call to False on the wrapper object. - """ + flip vllm_first_call to False on the wrapper object.""" from unittest.mock import patch from vllm.v1.attention.backends.flashinfer import fast_plan_decode diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py index b35005d36112..48eb58b5e1d9 100644 --- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py +++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py @@ -134,8 +134,7 @@ def make_quantized_kv_cache( head_size: int, ) -> tuple: """Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf, - kv_scale, ref_kv_cache, is_nvfp4_kv). - """ + kv_scale, ref_kv_cache, is_nvfp4_kv).""" is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE if is_nvfp4_kv: data, scales, kv_scale, ref = make_nvfp4_kv_cache( diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py index e86686ee8668..aa7c836666f7 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -73,8 +73,7 @@ def _strided_context_inputs( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The real prefill-context layouts: ``k_nope``/``v`` as the two strided halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather - workspace (left in the fp8 cache layout for a plain fp8 cache). - """ + workspace (left in the fp8 cache layout for a plain fp8 cache).""" kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype) k_nope, v = kv_nope.split((128, 128), dim=-1) workspace = _randn(num_tokens, 576, dtype=dtype) diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 07295ed0b871..1b89e4db66b3 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -89,8 +89,7 @@ def merge_attn_states_torch( @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16]) def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: """When a token is empty on both sides (both LSE -inf), the 0/0 softmax - scales must not surface as NaN in the merged output. - """ + scales must not surface as NaN in the merged output.""" num_tokens, num_heads, head_size = 6, 8, 128 prefix_output = torch.zeros( num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index c7b763e7c7e1..a7e6d499c176 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -45,8 +45,7 @@ def kv_layout(request) -> KVCacheLayout: def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]: """Per-layer physical stride order for the given layout; the 3-dim - indexer side cache (H=1) is contiguous, so identity. - """ + indexer side cache (H=1) is contiguous, so identity.""" if ndim == 3: return (0, 1, 2) stride_order = layout.layer_view_order @@ -66,8 +65,7 @@ def _main_spec() -> FullAttentionSpec: def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]: """Standardized per-layer logical shape [B, H, N, C] for the main cache, - derived the same way the production allocator does. - """ + derived the same way the production allocator does.""" shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages) return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize) @@ -77,8 +75,7 @@ def _allocate_main_kv_via_contract( ) -> torch.Tensor: """Build the main KV cache exactly as the production allocator does for the given layout: allocate the physical (permuted) tensor, then expose the - inverse-permuted logical [B, H, N, C] view the backend sees. - """ + inverse-permuted logical [B, H, N, C] view the backend sees.""" logical_shape = _main_kv_logical_shape(num_pages) stride_order = _layer_stride_order(layout, len(logical_shape)) physical_shape = tuple(logical_shape[i] for i in stride_order) @@ -1379,8 +1376,7 @@ def test_amd_decode_index_topk_end_to_end( @pytest.mark.parametrize("num_idx_heads", [1, 4]) def test_decode_index_topk_fp8(num_idx_heads: int): """The standalone Triton path must score FP8 inputs in FP32 so its top-k - matches a reference computed from the dequantized FP8 values. - """ + matches a reference computed from the dequantized FP8 values.""" torch.manual_seed(0) topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128 decode_query_len = 1 @@ -1653,8 +1649,7 @@ def test_prefill_sparse_attention_correctness( def test_main_cache_layout_contract(): """The standardized per-layer logical shape is [B, H, N, C] with packed K/V content, and the legacy layout aliases resolve to the expected - per-layer stride orders. - """ + per-layer stride orders.""" nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) assert logical == (nb, h, bs, 2 * d) @@ -1778,8 +1773,7 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): """The indexer side cache is standardized 4D with H=1: under both layouts the allocator's logical view stays contiguous and squeezes (as `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim - [num_blocks, block_size, head_dim] cache the kernels consume. - """ + [num_blocks, block_size, head_dim] cache the kernels consume.""" nb = 5 ispec = MLAAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE @@ -1799,8 +1793,7 @@ def test_indexer_cache_squeezes_to_contiguous_3d(): def test_hnd_allocation_is_packed_head_major(): """Under HND the backend-visible logical view is the packed head-major - physical allocation. - """ + physical allocation.""" nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM logical = _main_kv_logical_shape(nb) stride_order = KVCacheLayout.LBHNC.layer_view_order @@ -1828,8 +1821,7 @@ def test_hnd_allocation_is_packed_head_major(): def test_main_cache_is_block_first_and_unpadded(): """The allocator's contiguous-view branch (not the padded-strided branch) is used for the main GQA cache: its spec is unpadded and the physical - layout keeps num_blocks as the first dimension under both layouts. - """ + layout keeps num_blocks as the first dimension under both layouts.""" spec = _main_spec() # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided(). assert spec.page_size_padded is None @@ -1851,8 +1843,7 @@ def _build_decode_inputs( ): """Shared decode setup: uniform query tokens per request, a non-identity block table, and topk indices selecting the current block plus older causal - blocks for each query token. - """ + blocks for each query token.""" active_batch = len(seq_lens_list) batch = active_batch + num_padded_reqs pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list] @@ -2188,8 +2179,7 @@ def test_decode_sparse_attention_correctness( ): """Decode (split-K) parity under both layouts: this is the only coverage of the decode-site cache feed, and the strided HND case fails if the kernel - ignores the cache strides. - """ + ignores the cache strides.""" torch.manual_seed(0) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs( seq_lens_list, decode_query_len, num_padded_reqs @@ -2239,8 +2229,7 @@ def test_decode_wrong_layout_breaks_parity(): already contiguous-NHD (i.e. skipping the allocator's inverse permute) reorders the K/V content, so the decode output no longer matches the reference computed on the correct logical view. The mislabeled tensor keeps - the same shape as the correct view, so the kernel stays in bounds. - """ + the same shape as the correct view, so the kernel stays in bounds.""" torch.manual_seed(0) seq_lens_list = (130, 257) q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list) @@ -2275,8 +2264,7 @@ def test_decode_wrong_layout_breaks_parity(): def test_main_cache_byte_identical_through_production_allocator(): """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main `FullAttentionSpec` under HND and assert the backend-visible view has the - same shape, stride, and storage offset as the packed-HND allocation. - """ + same shape, stride, and storage offset as the packed-HND allocation.""" nb = 4 spec = _main_spec() raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) @@ -2291,8 +2279,7 @@ def test_main_cache_byte_identical_through_production_allocator(): def test_padded_main_cache_is_flagged(): """AC-2.1 negative: the M3 main cache relies on the allocator's contiguous-view branch (`page_size_padded is None`). A spec that sets - `page_size_padded` is explicitly flagged rather than silently wrong-strided. - """ + `page_size_padded` is explicitly flagged rather than silently wrong-strided.""" def _require_unpadded_block_first(spec, stride_order): inv_order = [stride_order.index(i) for i in range(len(stride_order))] @@ -2329,8 +2316,7 @@ def _require_unpadded_block_first(spec, stride_order): def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout): """AC-5 write path: the `reshape_and_cache_flash` write site now consumes packed-content K/V split views. Writing through those views must persist - into the bound storage under both layouts. - """ + into the bound storage under both layouts.""" torch.manual_seed(0) num_pages = 4 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index 8106ba37a6e9..ffdb76e176f8 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -23,8 +23,7 @@ def test_concat_and_cache_mla_into_unified_slot_view(): """concat_and_cache_mla must write correctly into a per-layer view whose block stride is the full unified slot (block-major), with zero bleed into - the other layers' segments of the same slot. - """ + the other layers' segments of the same slot.""" from vllm import _custom_ops as ops torch.manual_seed(0) @@ -86,8 +85,7 @@ def write(cache): def test_flashmla_dense_decode_unified_slot_view(): """FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA on Hopper) must read a unified-slot block-major view bit-identically to a - contiguous per-layer cache. - """ + contiguous per-layer cache.""" import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -152,8 +150,7 @@ def run(kc): def test_flashinfer_mla_dense_decode_unified_slot_view(): """FlashInfer MLA dense decode must read a unified-slot block-major view (inflated stride(0), non-zero storage offset) bit-identically to a - contiguous per-layer cache. - """ + contiguous per-layer cache.""" try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: @@ -224,8 +221,7 @@ def run(kv): def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a unified-slot block-major view bit-identically to a contiguous fp8_ds_mla - cache, with finite nonzero output. - """ + cache, with finite nonzero output.""" import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_sparse_supported() @@ -314,8 +310,7 @@ def run(kc): def test_indexer_k_quant_and_cache_into_unified_slot_view(): """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must write correctly into a per-layer view whose block stride is the full - unified slot, with zero bleed into the other layers' segments. - """ + unified slot, with zero bleed into the other layers' segments.""" from vllm import _custom_ops as ops torch.manual_seed(0) @@ -360,8 +355,7 @@ def write(cache): def test_flashattn_mla_dense_decode_unified_slot_view(): """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot - block-major view bit-identically to a contiguous per-layer cache. - """ + block-major view bit-identically to a contiguous per-layer cache.""" try: from vllm.vllm_flash_attn import flash_attn_varlen_func except ImportError: @@ -434,8 +428,7 @@ def run(cache): def test_flashmla_dense_fp8_decode_unified_slot_view(): """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache) must read a unified-slot block-major view bit-identically to a contiguous - per-layer fp8 cache. - """ + per-layer fp8 cache.""" import vllm.v1.attention.ops.flashmla as fm ok, reason = fm.is_flashmla_dense_supported() @@ -506,8 +499,7 @@ def run(kc): def test_flashinfer_mla_dense_fp8_decode_unified_slot_view(): """FlashInfer MLA dense decode with an fp8 KV cache must read a unified-slot block-major view bit-identically to a contiguous per-layer - cache. - """ + cache.""" try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla except ImportError: diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index 0fbe7bced2da..dd2870ccf178 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -283,8 +283,7 @@ def _run_direct_flash_attn_case( # Backend contract tests -------------------------------------------------- def test_aiter_mha_backend_contract(): """The ROCm backend advertises the dtypes, shapes, and attention types it - is designed to handle. - """ + is designed to handle.""" from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -312,8 +311,7 @@ def test_aiter_mha_backend_contract(): def test_aiter_mha_backend_validates_kv_cache_block_size(): """The backend should reject KV cache specs that cannot be gathered - correctly by the ROCm kernel. - """ + correctly by the ROCm kernel.""" from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -333,8 +331,7 @@ def test_aiter_mha_backend_validates_kv_cache_block_size(): def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe(): """The backend should trust the ROCm MI3xx probe instead of the raw torch - capability tuple. - """ + capability tuple.""" from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend @@ -359,8 +356,7 @@ def test_aiter_mha_env_flags_control_enablement( use_aiter, use_mha, expected, monkeypatch ): """Both the global AITER flag and the MHA-specific flag must be enabled - before the ROCm MHA path is considered active. - """ + before the ROCm MHA path is considered active.""" from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -378,8 +374,7 @@ def test_aiter_mha_env_flags_control_enablement( def test_aiter_mha_platform_gate_matches_install_and_arch(): """The global AITER availability check should only open on ROCm MI3xx when - the aiter package is installed. - """ + the aiter package is installed.""" from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported assert is_aiter_found_and_supported() is ( @@ -650,8 +645,7 @@ def test_aiter_mha_multi_batch(num_heads, head_size, dtype): @pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS) def test_aiter_fa_large_block_table_matches_reference(num_blocks): """The direct paged-KV path should stay stable for both normal and very - large block tables. - """ + large block tables.""" atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() @@ -673,8 +667,7 @@ def test_aiter_fa_large_block_table_matches_reference(num_blocks): @pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only") def test_aiter_fa_sliding_window_matches_reference(): """The direct kernel should respect the same sliding-window causal mask as - the naive reference implementation. - """ + the naive reference implementation.""" atol = 2e-2 rtol = 2e-2 _assert_aiter_supported() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index 17b59c074332..aabce97fac9e 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -111,8 +111,7 @@ def _make_mixed_common_metadata(): def _patch_build_deps(monkeypatch, events=None): """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs - on CPU. - """ + on CPU.""" def fake_generate_sparse_seqlen_triton( query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len @@ -153,8 +152,7 @@ def test_build_populates_decode_only_split_fields(monkeypatch): def test_build_populates_mixed_split_fields(monkeypatch): """Mixed decode+prefill batch: split is reported, prefill fields stay - default because this impl always runs the MQA path. - """ + default because this impl always runs the MQA path.""" builder = _make_builder() _patch_build_deps(monkeypatch) diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index a46916e14ef3..5ac1d6cdcbf8 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -125,8 +125,7 @@ def _quantize_to_fp8(tensor: torch.Tensor): @pytest.mark.parametrize("PAGE_SIZE", [1, 16]) def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE): """Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales, - and compare against BF16 reference output. - """ + and compare against BF16 reference output.""" assert CACHE_SIZE % PAGE_SIZE == 0 dtype = torch.bfloat16 seq_len = L @@ -249,8 +248,7 @@ def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SI """The kernel must honor the cache's page-dim stride, not assume pages are packed back-to-back. A per-layer view into a cross-layer (block-major) cache has stride(0) inflated by num_layers; outputs must match a - contiguous cache holding the same data exactly. - """ + contiguous cache holding the same data exactly.""" B = 3 seq_len = 1027 CACHE_SIZE = 16384 diff --git a/tests/kernels/core/test_fused_embed_norm.py b/tests/kernels/core/test_fused_embed_norm.py index fc30ab3acaca..9da43dba8361 100644 --- a/tests/kernels/core/test_fused_embed_norm.py +++ b/tests/kernels/core/test_fused_embed_norm.py @@ -41,8 +41,7 @@ def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: def test_fused_embed_norm_matches_reference(): """Main-model fusion: the residual is the exact gather and the second output is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16 - ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel. - """ + ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel.""" set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") @@ -61,8 +60,7 @@ def test_fused_embed_norm_matches_reference(): @torch.inference_mode() def test_fused_embed_eh_norm_matches_reference(): """MTP fusion (folded gather) is bit-exact vs gathering the embeds and - feeding the untouched model-local ``fused_eh_norm``. - """ + feeding the untouched model-local ``fused_eh_norm``.""" set_random_seed(13) table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 81382bdda91d..3d0e2ac11e35 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -62,8 +62,7 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): a [1, q_size] slice of a [1, q_size + kv_size] buffer used to produce a qr_out with row stride q_size + kv_size. Downstream dispatchers that require packed row-major inputs then reject the tensor and silently fall - back to a slower GEMM path on every decode step. - """ + back to a slower GEMM path on every decode step.""" device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 @@ -89,8 +88,7 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int): def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int): """Regression guard: grid used to be (2, num_tokens), hitting CUDA's 65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2) - lifts that bound to 2**31-1. - """ + lifts that bound to 2**31-1.""" device = "cuda" dtype = torch.bfloat16 q_size, kv_size = 192, 576 diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 74fdb00c30bf..86b8e7965673 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests that FusedRMSNormGated decomposes correctly under torch.compile, -matching the eager triton kernel output. -""" +matching the eager triton kernel output.""" import pytest import torch diff --git a/tests/kernels/core/test_rocm_aiter_ops.py b/tests/kernels/core/test_rocm_aiter_ops.py index b2bed636e9f1..c15afd72d7fb 100644 --- a/tests/kernels/core/test_rocm_aiter_ops.py +++ b/tests/kernels/core/test_rocm_aiter_ops.py @@ -579,8 +579,7 @@ def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m): def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view(): """Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match - contiguous inputs. - """ + contiguous inputs.""" require_aiter() require_fp8() from vllm._aiter_ops import rocm_aiter_ops diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index fa50d2b3eb39..b0038c084e3c 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -87,8 +87,7 @@ def _shard(full_table, si, rows_per_rank, dtype): @pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) def test_matches_masked_gather(tp_size, vocab_size, org_vocab_size, hidden, id_dtype): """Every rank's partial output is bit-exact with the eager path, including - the added-vocab (LoRA) rows and the ids owned by another rank. - """ + the added-vocab (LoRA) rows and the ids owned by another rank.""" set_random_seed(7) dtype = torch.bfloat16 rows = pad_vocab_size(vocab_size, 64) // tp_size diff --git a/tests/kernels/helion/test_autotune.py b/tests/kernels/helion/test_autotune.py index f28e8921093b..8b42e145d484 100644 --- a/tests/kernels/helion/test_autotune.py +++ b/tests/kernels/helion/test_autotune.py @@ -59,8 +59,7 @@ def teardown_method(self): def test_autotune_disabled_kernel_produces_valid_config(self): """Register a kernel with no configs (disabled), run autotune, - verify it produces a valid helion.Config. - """ + verify it produces a valid helion.Config.""" with dummy_kernel_registry(configs={}) as register: wrapper = register( "autotune_test_kernel", diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index 572a4126cf24..dd4274bfafd7 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -598,8 +598,7 @@ def default_picker(args, config_keys): ) def test_init_eagerly_initializes_hop_path(self): """Test that register_kernel eagerly builds the configured kernel - on the HOP path (no custom op registration needed). - """ + on the HOP path (no custom op registration needed).""" from vllm.kernels.helion.utils import get_canonical_gpu_name configs: dict[CaseKey, helion.Config] = { @@ -634,8 +633,7 @@ def test_init_eagerly_initializes_hop_path(self): ) def test_init_eagerly_initializes(self): """Test that register_kernel eagerly loads configs and detects GPU - during construction so __call__ needs no further initialization. - """ + during construction so __call__ needs no further initialization.""" from vllm.kernels.helion.utils import get_canonical_gpu_name with ( @@ -963,8 +961,7 @@ class TestTorchCompileHOP: def test_compiled_graph_contains_helion_hop(self): """Verify torch.compile on a HelionKernelWrapper emits a - helion_kernel_wrapper_mutation HOP node in the FX graph. - """ + helion_kernel_wrapper_mutation HOP node in the FX graph.""" configs: dict[CaseKey, helion.Config] = { CaseKey.default(): helion.Config(block_sizes=[4, 4]) } diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index e41b7c09d3fe..c324b189b9df 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -368,8 +368,7 @@ def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): @pytest.mark.parametrize("provider", ["vllm_c"]) def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider): """Test that inplace implementations reuse inputs, - for maybe_inplace overload but not for default overload. - """ + for maybe_inplace overload but not for default overload.""" impl = ir.ops.fused_add_rms_norm.impls[provider] if not impl.supported: pytest.skip(f"{provider} impl not supported on this platform") diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index 9495465b9ec3..c239f831c614 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -1055,8 +1055,7 @@ def test_batch_memcpy_cpu_fallback() -> None: def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"): """High-precision conv reference: everything in float32, no bf16 rounding. - x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim]. - """ + x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim].""" xf = x.float() wf = weight.float().unsqueeze(1) bf = bias.float() @@ -1165,8 +1164,7 @@ def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens): def test_conv_cpp_vnni_packed_matches_torch(seq_lens): """The exact runtime prefill sequence (VNNI-packed weight + SD-layout conv_state view + is_vnni=True) must match the torch fallback. Validates - the packing + layout handoff on any AVX-512BF16 CPU. - """ + the packing + layout handoff on any AVX-512BF16 CPU.""" x, weight, bias = _conv_inputs(sum(seq_lens)) out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True) @@ -1180,8 +1178,7 @@ def test_conv_cpp_vnni_packed_matches_torch(seq_lens): @torch.inference_mode() def test_conv_update_cpp_matches_torch(batch): """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch, - including the in-place conv_state update (the next-step handoff). - """ + including the in-place conv_state update (the next-step handoff).""" from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_update_torch, ) diff --git a/tests/kernels/mamba/test_mamba_ssm_configs.py b/tests/kernels/mamba/test_mamba_ssm_configs.py index 4330caa6115a..7ea1529f2e8b 100644 --- a/tests/kernels/mamba/test_mamba_ssm_configs.py +++ b/tests/kernels/mamba/test_mamba_ssm_configs.py @@ -122,8 +122,7 @@ def test_fallback_when_no_config(monkeypatch, tmp_path): def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): """When effective_batch = batch*nheads is not an exact key, the closest - key should be selected. - """ + key should be selected.""" _write_config( tmp_path, dstate=32, @@ -168,8 +167,7 @@ def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path): def test_non_dict_json_returns_none(monkeypatch, tmp_path): """A valid JSON file that is not a dict (e.g. a list) must be ignored - and return None rather than raising AttributeError. - """ + and return None rather than raising AttributeError.""" device_name = get_ssm_device_name() config_path = tmp_path / get_ssm_config_file_name( _HEADDIM, 16, _CACHE_DTYPE, device_name @@ -191,8 +189,7 @@ def test_non_dict_json_returns_none(monkeypatch, tmp_path): def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path): """An empty JSON object {} must not crash min() — should fall back - to the hard-coded heuristic. - """ + to the hard-coded heuristic.""" _write_config(tmp_path, dstate=64, payload={}) monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path)) diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py index fbc7b691870e..063ff17dd337 100644 --- a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -68,8 +68,7 @@ def _run_prefill_decode_equivalence( the exact fp32 recurrence (and each other). All paths use the production dt flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` - is the recurrent-state precision; ``act_dtype`` the activation/buffer one. - """ + is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" device = current_platform.device_type rtol, atol = _prefill_tolerances(act_dtype) set_random_seed(seed) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py index 6e5ada19bc8f..588ce52dbeb6 100644 --- a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -95,8 +95,7 @@ def _run_standard_decode( ) -> None: """Drive ``num_steps`` decode steps and check every path against the trusted recurrence. ``state_dtype`` is the recurrent-state precision; - ``act_dtype`` is the activation/buffer precision. - """ + ``act_dtype`` is the activation/buffer precision.""" device = "cuda" both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) @@ -540,8 +539,7 @@ def test_replayssm_standard_decode_tp_head_shard_equivalence( separable: one run over all heads must equal concatenating ``tp`` independent per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This guards the per-rank divisors and group->head mapping the state-shape wiring - relies on. The real TP1==TP2 engine check lives in the v1/e2e suite. - """ + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" state_dtype, act_dtype = precision nheads, headdim, dstate, ngroups = geometry assert nheads % tp == 0 and ngroups % tp == 0 diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index dd17298cb367..8adfa6588ca8 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -167,8 +167,7 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: def fe_supports_quant_scheme(self) -> bool: """Check if the fused experts class supports this quant config. - See https://github.com/ROCm/aiter/issues/2419 for AITER gaps. - """ + See https://github.com/ROCm/aiter/issues/2419 for AITER gaps.""" if self.quant_config is None or self.quant_dtype is None: return True if self.quant_dtype != torch.float8_e4m3fn: diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index de8f75d0e9a1..d0b1ce7cae46 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -468,8 +468,7 @@ def test_deepseek_v4_flashinfer_cutlass_falls_through_to_w4a8( monkeypatch: pytest.MonkeyPatch, ) -> None: """Explicit flashinfer_cutlass must try the W4A8 variant when the BF16 - variant is unsupported (the BF16 variant is gated to SM90). - """ + variant is unsupported (the BF16 variant is gated to SM90).""" from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( FlashInferExperts, ) diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index c4ebc5ff4077..9be58e6337bd 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -129,8 +129,7 @@ def test_batched_mm( per_act_token_quant: bool, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware. - """ + and those tests will be skipped on unsupported hardware.""" set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -269,8 +268,7 @@ def test_fused_moe_batched_experts( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware. - """ + and those tests will be skipped on unsupported hardware.""" set_random_seed(7) use_fp8_w8a8 = dtype == torch.float8_e4m3fn @@ -651,8 +649,7 @@ def test_batched_experts_masked_activation_end_to_end( def test_batched_experts_end_to_end(m, n, k, e, topk): """End-to-end BatchedTritonExperts via the reference (no-comms) BatchedPrepareAndFinalize, validated against a torch reference. Exercises - the device-enablement path. - """ + the device-enablement path.""" from vllm.v1.worker.workspace import init_workspace_manager set_random_seed(7) diff --git a/tests/kernels/moe/test_block_int8.py b/tests/kernels/moe/test_block_int8.py index 84a4a0d2d102..415faafb3ab0 100644 --- a/tests/kernels/moe/test_block_int8.py +++ b/tests/kernels/moe/test_block_int8.py @@ -54,8 +54,7 @@ # For test def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape): """This function performs fused moe with block-wise quantization using - native torch. - """ + native torch.""" B, D = a.shape a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) @@ -98,8 +97,7 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed): """Tests the fused_moe kernel with W8A8 INT8 block quantization against a - native torch reference. - """ + native torch reference.""" torch.manual_seed(seed) a = torch.randn((M, K), dtype=dtype) / 10 diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index e7af163b04d9..cf83ef6fc00a 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -397,8 +397,7 @@ def test_cpu_fused_moe_int8( class _StubMoELayer(torch.nn.Module): """Minimal stand-in for the real MoE layer module, exposing just what the unquantized CPU experts read/replace (w13_weight, w2_weight, the - router configuration, and optionally w13_bias/w2_bias). - """ + router configuration, and optionally w13_bias/w2_bias).""" def __init__( self, @@ -501,8 +500,7 @@ def test_cpu_fused_moe_unaligned_intermediate_size( act: MoEActivation, ): """CPU kernels handle unaligned intermediate sizes by zero-padding the - weights before prepacking. - """ + weights before prepacking.""" set_random_seed(0) batch_size = 64 intermediate_size = UNALIGNED_INTERMEDIATE_DIM diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 9627e2c67d4d..2e022ede3627 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -464,8 +464,7 @@ def test_trtllm_fp8_swiglu_clamp_support( ): """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the - per-tensor kernel has no clamp, and Relu2 rejects the parameters. - """ + per-tensor kernel has no clamp, and Relu2 rejects the parameters.""" class _Experts(TrtLlmFp8ExpertsModular): @staticmethod diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 0a16a3179739..4afcba706ccd 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -377,8 +377,7 @@ def test_unit_shuffle(): @pytest.mark.parametrize("n_experts,topk", [(32, 4), (128, 4)]) def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): """routing_data_from_sparse_topk must produce routing structures - identical to make_routing_data for the same topk result. - """ + identical to make_routing_data for the same topk result.""" from vllm.model_executor.layers.fused_moe.experts import ( gpt_oss_triton_kernels_moe as gptoss_moe, ) diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index b647439c49f7..97041e24b9bb 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -289,8 +289,7 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group @torch.inference_mode() def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk): """Test that the flashinfer_trtllm_mxint4_moe wrapper matches the raw - trtllm_mxint4_block_scale_moe kernel call. - """ + trtllm_mxint4_block_scale_moe kernel call.""" pytest.importorskip("flashinfer") from flashinfer import RoutingMethodType from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 6c21eb0d9ade..0c0e1d61f90e 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -352,8 +352,7 @@ def test_modular_kernel_combinations_singlegpu( workspace_init, ): """Note: float8_e4m3fn is not supported on CUDA architecture < 89, - and those tests will be skipped on unsupported hardware. - """ + and those tests will be skipped on unsupported hardware.""" config = Config( Ms=Ms, K=k, diff --git a/tests/kernels/moe/test_moe_kernel_oracle.py b/tests/kernels/moe/test_moe_kernel_oracle.py index 6ba5753f8d19..fbdf804a3b78 100644 --- a/tests/kernels/moe/test_moe_kernel_oracle.py +++ b/tests/kernels/moe/test_moe_kernel_oracle.py @@ -21,8 +21,7 @@ class TestUnquantizedDelegation: """UnquantizedMoEKernelOracle methods must delegate to the existing - module-level functions; behaviour is bit-identical. - """ + module-level functions; behaviour is bit-identical.""" def test_make_kernel_delegates(self) -> None: quant_config = object() diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 38b9dd025228..5b1af59731b7 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -315,8 +315,7 @@ def tp_chunk_gate_up( device: torch.device | int | None = None, ) -> torch.Tensor: """TP-chunk a combined [gate; up] weight, splitting each half separately - so every rank gets a portion of both gate and up. - """ + so every rank gets a portion of both gate and up.""" half = w.shape[dim] // 2 gate = chunk_by_rank( w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index e94882a310a2..51aff4ea3ba0 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -130,8 +130,7 @@ def test_is_supported_config(present, ep_size, supported, reason_substr): def test_explicit_moe_backend_aiter(): """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear - ValueError when the flydsl package is missing. - """ + ValueError when the flydsl package is missing.""" with _gfx950(), _flydsl_installed(True): assert ( _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) @@ -163,8 +162,7 @@ def test_gfx950_picks_aiter(): def test_gfx942_picks_emulation(): """Flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton - dot_scaled backend wins instead. - """ + dot_scaled backend wins instead.""" with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False): backend, experts_cls = select_mxfp8_moe_backend(_config()) assert backend is Fp8MoeBackend.EMULATION diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index de4ae5ae31d8..0967f9499e0f 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -249,8 +249,7 @@ def _make_topk_ids( device: str = "cuda", ) -> torch.Tensor: """Generate distinct expert IDs per token with the same top-k shape as - production routing. - """ + production routing.""" router_logits = torch.randn(num_tokens, num_experts, device=device) _, topk_ids = torch.topk(torch.softmax(router_logits, dim=-1), k=topk, dim=-1) return topk_ids.to(torch.int32) @@ -444,8 +443,7 @@ def test_aiter_fused_moe_custom_op_registered(): def test_aiter_asm_moe_tkw1_custom_op_registered(): """The tkw1 custom op should stay registered for FP8 apply-router-weight - paths. - """ + paths.""" _assert_aiter_supported() import vllm._aiter_ops as aiter_ops # noqa: F401 @@ -455,8 +453,7 @@ def test_aiter_asm_moe_tkw1_custom_op_registered(): def test_aiter_fused_moe_fake_tensor_support(): """The fused-MoE op should preserve fake-tensor compatibility for - torch.compile-style tracing. - """ + torch.compile-style tracing.""" _assert_aiter_supported() import vllm._aiter_ops # noqa: F401 @@ -519,8 +516,7 @@ def test_aiter_moe_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """The fused-MoE gate should depend only on the main AITER toggle and the - MoE-specific toggle. - """ + MoE-specific toggle.""" from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -551,8 +547,7 @@ def test_aiter_moe_shared_experts_enablement_follows_env( monkeypatch: pytest.MonkeyPatch, ): """Shared-expert fusion should only be enabled when the fused-MoE path is - enabled too. - """ + enabled too.""" from vllm._aiter_ops import rocm_aiter_ops _assert_aiter_supported() @@ -693,8 +688,7 @@ def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): """The gfx950 AITER MXFP4 W4A16 MoE path should match the dequantized - MXFP4 reference. - """ + MXFP4 reference.""" from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -743,8 +737,7 @@ def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy(): @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fused_moe_mi350_mxfp4_w4a16_determinism(): """The gfx950 AITER MXFP4 W4A16 MoE path should stay bitwise - deterministic. - """ + deterministic.""" from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, @@ -837,8 +830,7 @@ def test_aiter_fused_moe_bf16_accuracy( intermediate_dim: int, ): """The ROCm AITER fused-MoE BF16 path should match the float32 reference - on representative shapes. - """ + on representative shapes.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -883,8 +875,7 @@ def test_aiter_fused_moe_bf16_accuracy( def test_aiter_fused_moe_gelu_accuracy(): """The GELU activation variant should stay aligned with the float32 - reference. - """ + reference.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -929,8 +920,7 @@ def test_aiter_fused_moe_gelu_accuracy(): def test_aiter_fused_moe_determinism(): """The BF16 fused-MoE kernel should stay bitwise deterministic for the - same inputs. - """ + same inputs.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( ActivationMethod, QuantMethod, @@ -1085,8 +1075,7 @@ def test_aiter_fused_moe_mi3xx_bf16_accuracy(): ) def test_aiter_fused_moe_mi3xx_fp8_accuracy(): """The MI3xx FP8 per-tensor MoE path should stay within the measured FP8 - error budget. - """ + error budget.""" from tests.kernels.moe.utils import make_test_weights from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py index ba5eef662454..eec6180550dc 100644 --- a/tests/kernels/moe/test_routed_experts_capture_monolithic.py +++ b/tests/kernels/moe/test_routed_experts_capture_monolithic.py @@ -185,8 +185,7 @@ def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts( ) -> None: """The capture callback should receive the int16 routed-expert IDs the kernel actually used, the values should be valid expert indices, and - each token should pick ``top_k`` distinct experts. - """ + each token should pick ``top_k`` distinct experts.""" if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -265,8 +264,7 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( routing_method: RoutingMethodType, ) -> None: """Routing replay works for non-DeepSeekV3 routing methods too. - FlashInfer's ``routing_replay_out`` is routing-method-agnostic. - """ + FlashInfer's ``routing_replay_out`` is routing-method-agnostic.""" torch.manual_seed(0) device = torch.device("cuda:0") @@ -319,8 +317,7 @@ def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: """With no callback installed the kernel should not see a - ``routing_replay_out`` tensor — verify the helper short-circuits. - """ + ``routing_replay_out`` tensor — verify the helper short-circuits.""" torch.manual_seed(0) device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( @@ -342,8 +339,7 @@ def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: """FlashInfer's ``routing_replay_out`` is supported by all routing methods, so ``supports_routing_replay_capture`` should be True - regardless of routing method. - """ + regardless of routing method.""" device = torch.device("cuda:0") for routing_method in ( RoutingMethodType.DeepSeekV3, @@ -365,8 +361,7 @@ def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None: """When capture is installed, the allocated buffer is int16 and shaped - ``(num_tokens, experts_per_token)``. - """ + ``(num_tokens, experts_per_token)``.""" device = torch.device("cuda:0") experts, _, _ = _make_bf16_monolithic_experts( num_experts=_DSV3_NUM_EXPERTS, @@ -597,8 +592,7 @@ def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale - path, DSV3 routing) captures valid expert IDs. - """ + path, DSV3 routing) captures valid expert IDs.""" if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " @@ -784,8 +778,7 @@ def _run_nvfp4_monolithic( routing_bias: torch.Tensor, ) -> torch.Tensor: """The monolithic NVFP4 apply expects packed fp4 hidden states + the - matching fp8 per-block scale stored in the ``a1q_scale`` slot. - """ + matching fp8 per-block scale stored in the ``a1q_scale`` slot.""" # Stash the weight tensors on the experts in the locations the apply() # implementation reads from (it pulls them from quant_config / scales # already; w1/w2 come in as args). @@ -813,8 +806,7 @@ def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts( top_k: int, ) -> None: """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs - on the DSV3 routing path. - """ + on the DSV3 routing path.""" if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: pytest.skip( f"DSV3 requires top_k <= n_group * topk_group " diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index f922a3ebd336..174950e06c25 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -250,8 +250,7 @@ def assert_aiter_routing_valid( the Python baseline (different group selection, scoring internals), so numerical comparison is not meaningful. Instead we verify the outputs satisfy the routing contract: correct shapes, valid expert - IDs, non-negative weights, and proper normalization. - """ + IDs, non-negative weights, and proper normalization.""" n_tokens = topk_weights.shape[0] # Shape diff --git a/tests/kernels/moe/test_routing_simulator.py b/tests/kernels/moe/test_routing_simulator.py index b1adcd53d93d..2cb122122fcf 100644 --- a/tests/kernels/moe/test_routing_simulator.py +++ b/tests/kernels/moe/test_routing_simulator.py @@ -76,8 +76,7 @@ def test_basic_functionality( def test_routing_strategy_integration(monkeypatch, device): """Test that the routing strategy environment variable works with - FusedMoEFactory. - """ + FusedMoEFactory.""" pytest.importorskip("vllm.model_executor.layers.fused_moe.layer") import vllm.envs as envs @@ -158,8 +157,7 @@ def test_routing_strategy_integration(monkeypatch, device): def test_distribution_based_routing_with_custom_strategy(): """Test registering and using DistributionBasedRouting with custom - parameters. - """ + parameters.""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Register custom distribution-based strategy diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index 76744da8f4df..204954d1fe46 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -70,8 +70,7 @@ def reference_with_clamp( ) -> tuple[torch.Tensor, torch.Tensor]: """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to match the C++ compute() template, then run the standard silu_and_mul + - quant reference. - """ + quant reference.""" N_2 = x.size(1) // 2 dtype = x.dtype gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype) diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py index 15cef0df38ee..8b10ff75e6db 100644 --- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py +++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py @@ -28,8 +28,7 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization - """ + quantization and per-column weight quantization""" A = A.to(torch.float32) B = B.to(torch.float32) @@ -57,8 +56,7 @@ def fp8_mask(a, mask): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk): """This function performs fused moe with per-column int8 - quantization using native torch. - """ + quantization using native torch.""" B, D = a.shape # Perform per-token quantization a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 48e98f9c1457..d4c037d978a3 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -150,8 +150,7 @@ def test_select_default_backend_by_platform( expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled. - """ + accelerators (FlashInfer, AITER) disabled.""" with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -445,8 +444,7 @@ def test_select_cuda_flashinfer_cutlass_backend( mock_has_flashinfer, ): """Test CUDA backend selection when FlashInfer TRTLLM is not available - and FlashInfer CUTLASS is available. - """ + and FlashInfer CUTLASS is available.""" with ( patch.object(current_platform, "is_cuda", return_value=True), patch.object(current_platform, "is_rocm", return_value=False), diff --git a/tests/kernels/quant_utils.py b/tests/kernels/quant_utils.py index 93a21596ca70..7e14ec95d367 100644 --- a/tests/kernels/quant_utils.py +++ b/tests/kernels/quant_utils.py @@ -159,8 +159,7 @@ def native_per_token_group_quant_fp8( x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn ): """Function to perform per-token-group quantization on an input tensor - `x` using native torch. - """ + `x` using native torch.""" assert x.shape[-1] % group_size == 0, ( "the last dimension of `x` must be divisible by `group_size`" ) diff --git a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py index 70abcbfc8ea4..c855d0e819dc 100644 --- a/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py +++ b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py @@ -304,8 +304,7 @@ def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel) def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): """Kernel output should match a dequantized fp32 reference within - fp8 per-token / per-channel quantization noise. - """ + fp8 per-token / per-channel quantization noise.""" weight_shape = (512, 4096) # (N, K) num_tokens = 32 _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index c01929c5d130..4858ea47adbf 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -30,8 +30,7 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): """Matrix multiplication function that supports per-token input - quantization and per-column weight quantization - """ + quantization and per-column weight quantization""" A = A.to(torch.float32) B = B.to(torch.float32) @@ -54,8 +53,7 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16): def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids): """This function performs fused moe with per-column int8 quantization - using native torch. - """ + using native torch.""" B, D = a.shape # Perform per-token quantization a_q, a_s = per_token_quant_int8(a) diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index f160125809e7..027ffbf10c96 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -360,8 +360,7 @@ def test_gptq_marlin_padded_round_trip(shape, group_size): def test_fp8_block_marlin_padded_round_trip(shape): """Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers): group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the - weight_scale_inv group-wise scale padding. - """ + weight_scale_inv group-wise scale padding.""" size_n, size_k = shape block = 128 dtype = torch.float16 @@ -411,8 +410,7 @@ def test_fp8_block_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 288), (4640, 512)]) def test_mxfp8_marlin_padded_round_trip(shape): """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to - 2^-127 instead of zero and must still contribute nothing. - """ + 2^-127 instead of zero and must still contribute nothing.""" size_n, size_k = shape group_size = 32 # The e8m0-scale Marlin kernels are only instantiated for bf16 activations. @@ -459,8 +457,7 @@ def test_mxfp8_marlin_padded_round_trip(shape): @pytest.mark.parametrize("shape", [(200, 512), (4640, 512)]) def test_awq_zp_marlin_padded_round_trip(shape): """AWQ-style uint4 with runtime zero-points, padded the way - MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0. - """ + MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0.""" size_n, size_k = shape group_size = 128 dtype = torch.float16 @@ -695,8 +692,7 @@ def make_layer(hidden, intermediate): @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_fp8_marlin_moe_padded_round_trip(shape, quant): """FP8 weight-only MoE: pad a tile-misaligned intermediate and check the - real prepare + fused_marlin_moe against the dequantized reference. - """ + real prepare + fused_marlin_moe against the dequantized reference.""" from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk @@ -781,8 +777,7 @@ def q(w): # (out, in) -> fp8 weight, scale, dequant reference @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)]) def test_mxfp8_marlin_moe_padded_round_trip(shape): """MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with - unit e8m0 scales so the reference is the exact fp8 dequant. - """ + unit e8m0 scales so the reference is the exact fp8 dequant.""" from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index de9feb7dfa05..4adfffa7c87b 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -72,8 +72,7 @@ def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an - explicit request for a different activation format must be rejected. - """ + explicit request for a different activation format must be rejected.""" config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) can_implement, reason = kernel_cls.can_implement(config) assert not can_implement @@ -88,8 +87,7 @@ def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( """Marlin/Humming never quantize activations, so an unset activation key, or one that already describes MXFP4-shaped data, is tolerated. When an activation key is explicitly set, a warning must be logged noting that it - is ignored, since these kernels are weight-only (A16). - """ + is ignored, since these kernels are weight-only (A16).""" config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: can_implement, reason = kernel_cls.can_implement(config) @@ -118,8 +116,7 @@ def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): ) def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): """EmulationMxfp4LinearKernel is the universal fallback: it must accept - every supported activation format. - """ + every supported activation format.""" # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -137,8 +134,7 @@ def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): """quant_dequant_func must be derived purely from the config's activation - QuantKey, not set externally. - """ + QuantKey, not set externally.""" # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, # which we are not testing here. monkeypatch.setattr( @@ -162,8 +158,7 @@ def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): def test_aiter_kernel_is_supported_requires_native_mx_support(): """AiterMxfp4LinearKernel must not be selected on platforms without - native MX compute, even if AITER itself is importable. - """ + native MX compute, even if AITER itself is importable.""" with patch( "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", return_value=False, @@ -199,8 +194,7 @@ def apply_weights( @patch("vllm.model_executor.kernels.linear.current_platform") def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp4_linear_kernel should select a registered kernel that - reports itself as supported, and construct it with a fresh config. - """ + reports itself as supported, and construct it with a fresh config.""" platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py index 74996d64e368..1edaa86cac61 100644 --- a/tests/kernels/quantization/test_mxfp4_triton_ep.py +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -16,8 +16,7 @@ class TestTritonMoeForwardExpertMap: """Test that triton_kernel_moe_forward applies expert_map remapping - when expert_map is provided (EP active). - """ + when expert_map is provided (EP active).""" def test_expert_map_remap(self): device = "cuda" if torch.cuda.is_available() else "cpu" diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py index 2c0e94c0f519..5c5b1b3b9e50 100644 --- a/tests/kernels/quantization/test_mxfp6_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py @@ -41,8 +41,7 @@ def test_can_implement_is_abstract(): def test_emulation_kernel_rejects_non_mxfp6_weights(): """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight - format. - """ + format.""" config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static) can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) assert not can_implement @@ -58,8 +57,7 @@ def test_emulation_kernel_accepts_any_supported_config( weight_quant_key, activation_quant_key ): """EmulationMxfp6LinearKernel is the only backend today: it must accept - every supported weight/activation format combination. - """ + every supported weight/activation format combination.""" config = MxFp6LinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key ) @@ -104,8 +102,7 @@ def apply_weights( def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock): """init_mxfp6_linear_kernel should select a registered kernel that reports itself as supported/able to implement the given config, and - construct it with that exact config. - """ + construct it with that exact config.""" platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6") diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index c1226efcc803..6d3c91022ed8 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -118,8 +118,7 @@ def test_per_token_group_quant_fp8_packed( num_tokens, hidden_dim, group_size, poisoned_scales ): """Test the packed DeepGEMM quantization kernel against the Triton - reference (row-major, UE8M0 scales). - """ + reference (row-major, UE8M0 scales).""" device = "cuda" torch.manual_seed(42) @@ -257,8 +256,7 @@ def test_per_token_group_quant_fp8_packed_all_zero(): def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): """Inputs whose absmax/max_8bit produces a non-power-of-2 force the mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior - before optimization. - """ + before optimization.""" device = "cuda" num_tokens, hidden_dim, group_size = 4, 7168, 128 @@ -330,8 +328,7 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( ): """When output_q is allocated with shape (tma_aligned_mn, k) instead of (mn, k), the kernel must overwrite the padded mn rows with zeros so - callers can use ``torch.empty`` instead of ``torch.zeros``. - """ + callers can use ``torch.empty`` instead of ``torch.zeros``.""" device = "cuda" group_size = 128 torch.manual_seed(42) diff --git a/tests/kernels/quantization/test_quantized_embedding.py b/tests/kernels/quantization/test_quantized_embedding.py index 35c0e74713cb..0e4af0a0c1a6 100644 --- a/tests/kernels/quantization/test_quantized_embedding.py +++ b/tests/kernels/quantization/test_quantized_embedding.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the Triton dequant-gather kernel used by -``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup). -""" +``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup).""" import pytest import torch diff --git a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py index 2a5a8081a2f5..c7f0f1eda355 100644 --- a/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py +++ b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py @@ -29,8 +29,7 @@ def test_rocm_aiter_group_fp8_quant_fake_implementation(): """Test that the fake implementation is correctly - defined for torch.ops.vllm.rocm_aiter_group_fp8_quant. - """ + defined for torch.ops.vllm.rocm_aiter_group_fp8_quant.""" # Create test tensors M = 128 N = 4096 @@ -50,8 +49,7 @@ def test_rocm_aiter_group_fp8_quant_fake_implementation(): def test_rocm_aiter_group_fp8_quant_torch_compile_with_cudagraph(): """Test that rocm_aiter_ops.group_fp8_quant with group size 128 can be used with - torch.compile in cudagraph mode. - """ + torch.compile in cudagraph mode.""" # Create test tensors M = 128 N = 4096 diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 7002165af093..0e14e2d2583c 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -209,8 +209,7 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch( @pytest.mark.parametrize("shape", LLAMA_SHAPES) def test_vllm_quant_dequant_mxfp4_matches_reference_on_large_shapes(shape, dtype): """The public vLLM MXFP4 QDQ helper should stay exact against the torch - reference even on larger Llama-like shapes. - """ + reference even on larger Llama-like shapes.""" from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( quant_dequant_mxfp4, ) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 3dc658250fec..95ff6b3c057f 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -670,8 +670,7 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): def test_v41_compressor_metadata_maps_tokens_to_their_ring(): """The ring group's generic slot mapping is disabled (all PAD), so the builder must map every real token to ``ring_block * capacity + pos % - capacity`` and keep padding tokens at PAD. - """ + capacity`` and keep padding tokens at PAD.""" from unittest.mock import MagicMock from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder @@ -1096,8 +1095,7 @@ def test_get_c128_boundary(starts, query_start_loc, expected): @pytest.mark.parametrize("block_size", [16, 64]) def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int): """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache - → compare against original. - """ + → compare against original.""" HEAD_DIM = 512 NOPE_DIM = 448 HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale @@ -1309,8 +1307,7 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("block_size", [16, 64]) def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache - → manual dequant → compare against original. - """ + → manual dequant → compare against original.""" HEAD_DIM = 128 QUANT_BLOCK_SIZE = 128 # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132 diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index d9157c411066..40c5ef5fdbd0 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -432,8 +432,7 @@ def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts): ) def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): """Correctness test with large sequence lengths matching benchmark - scenarios (8K-128K prefill). - """ + scenarios (8K-128K prefill).""" ( cache, block_table, diff --git a/tests/kernels/test_deepseek_v4_cpu_kernels.py b/tests/kernels/test_deepseek_v4_cpu_kernels.py index 9c3fddb6c171..81c03a0bca9f 100644 --- a/tests/kernels/test_deepseek_v4_cpu_kernels.py +++ b/tests/kernels/test_deepseek_v4_cpu_kernels.py @@ -85,8 +85,7 @@ def _decode_token_regions( cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int ): """Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte - regions this layout defines. - """ + regions this layout defines.""" valid = slot_mapping >= 0 idx = valid.nonzero(as_tuple=True)[0] slots = slot_mapping[idx] @@ -111,8 +110,7 @@ def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size): """NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an un-rotated bf16 value, so they match tightly; the RoPE region is rotated in fp32 before its bf16 store and eager/kernel can round a tie to - opposite sides, so allow <=1 ULP there. - """ + opposite sides, so allow <=1 ULP there.""" nope_ref, rope_ref, scale_ref = _decode_token_regions( cache_ref, slot_mapping, block_size ) @@ -192,8 +190,7 @@ def _dup_padded_block_cache( num_blocks: int, block_size: int, width: int, pad_elems: int ) -> tuple[torch.Tensor, torch.Tensor]: """Two independently-stored padded caches with identical values, for ops - that mutate the cache in place. - """ + that mutate the cache in place.""" page_stride = block_size * width + pad_elems template = torch.randn(num_blocks, page_stride, dtype=torch.float32) @@ -210,8 +207,7 @@ def _run_save_partial_states_eager( ): """Reference: state_cache[block_idx, pos_in_block, :W] = kv; state_cache[..., W:] = score + ape[position % compress_ratio]; skip - slot < 0. Mutates state_cache in place, matching the op's contract. - """ + slot < 0. Mutates state_cache in place, matching the op's contract.""" state_width = kv.shape[-1] block_size = state_cache.shape[1] valid = slot_mapping >= 0 @@ -271,8 +267,7 @@ def _run_compress_eager( ): """Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0 quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla - layout. - """ + layout.""" num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -449,8 +444,7 @@ def test_compress_norm_rope_store_cpu_matches_eager( def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride(): """The main (output) kv_cache is a separate physical buffer from the - state cache and can independently be a strided view. - """ + state cache and can independently be a strided view.""" eps = 1e-6 compress_ratio, overlap = 128, False num_tokens = 5 @@ -526,8 +520,7 @@ def _run_compress_indexer_eager( """Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512 path, but the indexer's cache quantizes all 128 post-RoPE dims as one FP8 block with a single raw fp32 scale per token (matching - _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks. - """ + _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks.""" num_tokens, window = gather_slots.shape state_width = state_cache.shape[-1] // 2 state_block_size = state_cache.shape[1] @@ -655,8 +648,7 @@ def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int): """Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and return (cache, dequant_ref) -- the exact post-quantization value the kernel's dequant should reproduce. Slot id == row index, packed - contiguously from block 0. - """ + contiguously from block 0.""" n = kv_f32.shape[0] num_blocks = max(1, -(-n // block_size)) cache = _make_cache(num_blocks, block_size) @@ -714,8 +706,7 @@ def _flash_mla_eager( """MQA attention: gather the (already-dequantized) window+compressed KV rows referenced by *_slots (-1 = invalid/masked), append a learned per-head sink logit with zero value contribution, softmax, weighted - sum. - """ + sum.""" num_tokens, num_heads, _ = q.shape q_f32 = q.float() @@ -814,8 +805,7 @@ def test_flash_mla_cpu_kernel_matches_eager( def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink(): """When every KV slot is masked out, softmax degenerates to the sink - logit alone and the output must be all-zero (no valid value rows). - """ + logit alone and the output must be all-zero (no valid value rows).""" torch.manual_seed(2) dtype = torch.bfloat16 num_tokens, num_heads, num_window = 3, 4, 5 @@ -855,8 +845,7 @@ def _qnorm_kv_insert_eager( q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs ): """Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE - + UE8M0 FP8 quant + paged insert on KV. - """ + + UE8M0 FP8 quant + paged insert on KV.""" num_quant_blocks = NOPE_DIM // QUANT_BLOCK num_tokens_full, num_heads_q, _ = q.shape @@ -982,8 +971,7 @@ def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager( @pytest.mark.parametrize("pad", [1, 5]) def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int): """slot_mapping shorter than q/kv rows: the KV branch must only touch - the first num_tokens rows while Q-norm+RoPE still runs on all rows. - """ + the first num_tokens rows while Q-norm+RoPE still runs on all rows.""" torch.manual_seed(1) dtype = torch.bfloat16 eps = 1e-6 @@ -1044,8 +1032,7 @@ def _inverse_gptj_rope_o_proj_eager( rope_dim: int, ) -> torch.Tensor: """NoPE dims pass through as fp32; the RoPE segment is de-rotated by - R(-pos) in GPT-J (interleaved even/odd) convention. - """ + R(-pos) in GPT-J (interleaved even/odd) convention.""" nope_dim = o.shape[-1] - rope_dim half_dim = rope_dim // 2 @@ -1111,8 +1098,7 @@ def _write_page( ) -> None: """Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size] fp32) into one page row: K-region then scale-region, matching - compress_norm_rope_store_indexer_cpu's layout. - """ + compress_norm_rope_store_indexer_cpu's layout.""" k_region = page_row[: block_size * INDEXER_HEAD_DIM] k_region.copy_(k_fp8.view(torch.uint8).flatten()) scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32) @@ -1148,8 +1134,7 @@ def _build_paged_mqa_logits_batch( ): """Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly non-contiguous in dim 0) with deterministic-per-seed random K/scale - data. - """ + data.""" torch.manual_seed(seed) batch_size = len(seq_lens) max_seq_len = max(seq_lens) @@ -1215,8 +1200,7 @@ def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride(): """The paged K-cache in production is a per-layer view into a shared multi-layer allocation, so its page stride generally exceeds the page byte width -- check the kernel reads a padded/strided buffer identically - to an equivalent tightly-packed one. - """ + to an equivalent tightly-packed one.""" block_size = 64 num_heads = 4 seq_lens = [1, block_size + 7, 3 * block_size] @@ -1342,8 +1326,7 @@ def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel(): CPU op on a CPU tensor, not silently fall through to the Triton kernel (which would also "work" under triton-cpu, masking a dispatch regression). The dispatcher does its CPU import inside the CPU branch, - so the patch target is the source module. - """ + so the patch target is the source module.""" torch.manual_seed(0) num_tokens = 1 @@ -1410,8 +1393,7 @@ def _hc_head_torch( ) -> torch.Tensor: """Eager reference for hc_head_fused_cpu: round-trips the RMSNorm output through bf16 before the linear, matching the kernel's bf16 - arithmetic. - """ + arithmetic.""" hc_mult, hidden_size = hidden_states.shape[-2:] outer_shape = hidden_states.shape[:-2] hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size) @@ -1493,8 +1475,7 @@ def test_mhc_post_cpu(num_tokens: int, hidden_size: int): @pytest.mark.parametrize("hidden_size", [4096, 7168]) def test_hc_head_cpu(num_tokens: int, hidden_size: int): """hc_head_fused_cpu vs. the _hc_head_torch eager reference it - replaces. - """ + replaces.""" torch.set_default_device(DEVICE) set_random_seed(0) hc_mult = 4 diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 8deba1241ae0..c2a7607cbac0 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -383,8 +383,7 @@ def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False): """Run the op on a batch of (token_ids, start, window) requests; `window` is the runner's [depth] lookback for that request (None = all unknown). Token IDs 14 and 22 stand in for image tokens that break n-grams. - Returns one [len(token_ids), ...] hash tensor per request. - """ + Returns one [len(token_ids), ...] hash tensor per request.""" token_map, multipliers, primes, offsets, block_size = tables_meta depth = multipliers.shape[1] - 1 @@ -443,8 +442,7 @@ def test_lookback_window_reproduces_single_instance(runner, capture): V2 supplies every lookback from its device token history and needs no slot cache. V1 supplies prompt positions only and reads generated positions from the slot cache it fills itself. Without any window the - first decode token hashes stale slots (negative control). - """ + first decode token hashes stale slots (negative control).""" block_size = 4 block_table = torch.tensor( [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda" @@ -710,8 +708,7 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): """The fused gather must be bit-exact with the torch dequant path it replaces, from HBM and from pinned host memory alike, and must contribute - zeros for rows another TP rank owns. - """ + zeros for rows another TP rank owns.""" layer = _make_embedding(cpu_offload) cols, rows = 24, layer.part_num_embeddings ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda") diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py index fad8d90805eb..492c099bf1c5 100644 --- a/tests/kernels/test_fla_layernorm_guard.py +++ b/tests/kernels/test_fla_layernorm_guard.py @@ -336,8 +336,7 @@ def test_layer_norm_rows_per_block( @torch.inference_mode() def test_strided_input(dtype: torch.dtype) -> None: """Test that the kernel handles non-contiguous (strided) - inputs correctly. - """ + inputs correctly.""" set_random_seed(42) device = torch.device(DEVICE) num_tokens = 128 diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index b6c2a53ed873..d330abba532a 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -97,8 +97,7 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: """The gate feeds top-k expert selection: the kernel's top-8 must match an fp64 reference's top-8 per token (ties tolerated). This is the business-level correctness of the router — numeric error only matters - if it flips the argsort. - """ + if it flips the argsort.""" _requires_sm90() top_k = 8 device = torch.device("cuda") diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 6691de7ebfff..7f74cde05d5b 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -854,8 +854,7 @@ def test_fused_q_no_indexer(num_tokens: int): @pytest.mark.parametrize("has_indexer", [True, False]) def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool): """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is - produced (bf16, unquantized); ql_nope is consumed directly by the caller. - """ + produced (bf16, unquantized); ql_nope is consumed directly by the caller.""" torch.manual_seed(6) dev = "cuda" max_pos = 8192 diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 3e95913a017a..ced96bc40199 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -176,8 +176,7 @@ def _call_fused( def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor: """Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on - gfx942) encoding the kernel actually wrote, without touching the bytes. - """ + gfx942) encoding the kernel actually wrote, without touching the bytes.""" return t.contiguous().view(torch.uint8).view(FP8_STORE_DTYPE) @@ -713,8 +712,7 @@ def run(q_in, is_q_interleaved): @pytest.mark.parametrize("block_size", [16, 64]) def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int): """slot_mapping.size(0) < q.size(0): the kernel must skip padded - tokens in the KV branch while still processing Q on all rows. - """ + tokens in the KV branch while still processing Q on all rows.""" torch.manual_seed(3) device = "cuda" dtype = torch.bfloat16 diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 20ac8f19f1a6..0ce3a918513f 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -574,8 +574,7 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's sparse MQA-logits kernels take) instead of fp32; the values are the fp32 - result rounded once, so the scoring kernel needs no cast. - """ + result rounded once, so the scoring kernel needs no cast.""" if use_cutedsl and not has_cutedsl(): pytest.skip("cutedsl (cutlass) not installed") diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index 97ea8f7e3443..03d5b7921fa2 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -446,8 +446,7 @@ def test_output_strides(num_tokens, num_heads, n_groups): def test_per_group_contiguity(num_tokens): """FP8 per-group slices must be contiguous. Scale per-group slices are column-major (T-stride=1) — not row-major contiguous, which is - correct for TMA loads. - """ + correct for TMA loads.""" num_heads, n_groups = 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -513,8 +512,7 @@ def test_scales_are_power_of_two(): def test_nope_dims_unchanged(): """Nope dimensions (first 448 per head) should only be quantized, not rotated. Verify by dequantizing and comparing against - quantize-only reference (no RoPE). - """ + quantize-only reference (no RoPE).""" num_tokens, num_heads, n_groups = 16, 64, 8 heads_per_group = num_heads // n_groups max_pos = 4096 @@ -905,8 +903,7 @@ def calc_diff(x, y): @torch.inference_mode() def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config): """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style, - mscale=0, YaRN scaling) matching the production config. - """ + mscale=0, YaRN scaling) matching the production config.""" num_heads = 64 n_groups = 8 heads_per_group = num_heads // n_groups diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 219e42c74294..fafb37d97b3c 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -1372,8 +1372,7 @@ def _patch_first_rank_pp_group(monkeypatch): def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): """First finalize (at the end of load_weights) allocates - hc_attn_fn_broadcast as hc_attn_fn summed over hc streams. - """ + hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) @@ -1388,8 +1387,7 @@ def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): """Re-finalizing after a weight refit must copy into the existing broadcast tensor so its address stays stable for captured CUDA graphs, - while picking up the new hc_attn_fn values. - """ + while picking up the new hc_attn_fn values.""" _patch_first_rank_pp_group(monkeypatch) layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 45fafb0433ca..997c61980913 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -403,8 +403,7 @@ def test_mxfp8_linear_emulation_bf16_at_load( ): """EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the ``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the - same result; the dtype-match (BF16/FP16 activations) must also hold. - """ + same result; the dtype-match (BF16/FP16 activations) must also hold.""" from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) @@ -518,8 +517,7 @@ def _spy(*a, **k): # produced EP garbage). def _capture_expert_mask(expert_mask, *, global_num_experts): """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the - ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``. - """ + ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``.""" from types import SimpleNamespace from unittest import mock @@ -570,8 +568,7 @@ def _fake_fused_moe(hidden_states, w1, w2, tw, ti, *, expert_mask, **kw): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is - (no re-derivation that would collapse an already-0/1 mask to all-ones). - """ + (no re-derivation that would collapse an already-0/1 mask to all-ones).""" # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3). ep_mask = torch.tensor( [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE @@ -585,8 +582,7 @@ def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged(): def test_routed_experts_expert_map_delegates_to_kernel(): """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed - on the resolved kernel, not the global aiter switch. - """ + on the resolved kernel, not the global aiter switch.""" from types import SimpleNamespace from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts diff --git a/tests/kernels/test_shuffle_rows.py b/tests/kernels/test_shuffle_rows.py index 35c2f2644290..4f32949eb765 100644 --- a/tests/kernels/test_shuffle_rows.py +++ b/tests/kernels/test_shuffle_rows.py @@ -17,8 +17,7 @@ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_shuffle_rows_basic(num_tokens: int, hidden_size: int, dtype: torch.dtype): """Test basic functionality of shuffle_rows with various tensor sizes and - dtypes. - """ + dtypes.""" if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") @@ -73,8 +72,7 @@ def test_shuffle_rows_permutation( @pytest.mark.parametrize("hidden_size", [256, 512]) def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int): """Test shuffle_rows with expansion (more output tokens than input - tokens). - """ + tokens).""" if not current_platform.is_cuda(): pytest.skip("shuffle_rows requires CUDA") diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index f249adeaffec..db49b8091b40 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -1740,8 +1740,7 @@ def test_sparse_indexer_decode_topk_explicit_backends( ) -> None: """Every explicit sparse-indexer decode top-k backend must restrict each row to its seq_len (dirty data past the end must never be selected) and - match torch.topk on the valid region. - """ + match torch.topk on the valid region.""" from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk @@ -1812,8 +1811,7 @@ def test_sparse_indexer_decode_topk_short_seq_lens( backend: str, workspace_init ) -> None: """seq_len < next_n: derived per-row ends must clamp at 0 (the reference - kernels do max(0, ...)); negative ends are OOB for the ragged kernels. - """ + kernels do max(0, ...)); negative ends are OOB for the ragged kernels.""" from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py index 3c9505adba7c..e858331c9364 100644 --- a/tests/lora/test_default_mm_loras.py +++ b/tests/lora/test_default_mm_loras.py @@ -132,8 +132,7 @@ def test_default_mm_lora_fails_with_overridden_lora_request( monkeypatch: pytest.MonkeyPatch, ): """Ensure that if the lora_request conflicts with default_mm_loras, - we use the lora_request. - """ + we use the lora_request.""" run_test( vllm_runner, audio_assets, diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index 9ee0ed8584c8..8b5b3d577a32 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -1123,8 +1123,7 @@ def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device): @pytest.mark.parametrize("device", DEVICES) def test_fused_moe_lora_kernel_zero_grid_no_crash(device): """num_active_loras=0 (or num_slices=0) would otherwise launch a grid - with a zero dimension. one-shot wrapper must short-circuit before launch. - """ + with a zero dimension. one-shot wrapper must short-circuit before launch.""" torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024 @@ -1186,8 +1185,7 @@ def test_fused_moe_lora_kernel_zero_grid_no_crash(device): def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device): """one-shot must surface a clear assertion when shrink_block_size_m is not a power of 2 / less than 16, instead of the cryptic Triton compile - failure (`arange's range must be a power of 2`). - """ + failure (`arange's range must be a power of 2`).""" torch.set_default_device(device) set_random_seed(0) num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024 diff --git a/tests/lora/test_moe_lora_align_sum.py b/tests/lora/test_moe_lora_align_sum.py index df27752d2c37..80ba0a56ca40 100644 --- a/tests/lora/test_moe_lora_align_sum.py +++ b/tests/lora/test_moe_lora_align_sum.py @@ -249,8 +249,7 @@ def _build_and_run_align( ) def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras): """Regression test for issue #32235: real LoRA slot must not be skipped - when ``active_lora_ids`` has -1 at position 0. - """ + when ``active_lora_ids`` has -1 at position 0.""" out = _build_and_run_align( num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras ) @@ -289,8 +288,7 @@ def test_moe_lora_align_block_size_disabled_adapter_untouched(): kernels. Pins the invariant protected by the ``adapter_enabled`` guard in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort kernel reads uninitialized ``token_mask`` values for disabled slots and - pollutes ``sorted_token_ids`` / ``cumsum_buffer``. - """ + pollutes ``sorted_token_ids`` / ``cumsum_buffer``.""" max_loras = 1 out = _build_and_run_align( num_lora_tokens=16, diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py index 52cccc6b06ce..60d0c2c9ba7e 100644 --- a/tests/lora/test_moe_lora_ep_load.py +++ b/tests/lora/test_moe_lora_ep_load.py @@ -19,8 +19,7 @@ def _expected_lora_modules() -> set[str]: """Replicate the set ``WorkerLoRAManager._load_adapter`` would build - from this model's ``packed_modules_mapping``. - """ + from this model's ``packed_modules_mapping``.""" expected: set[str] = set(NON_EXPERT_MODULES) for expert in range(GLOBAL_NUM_EXPERTS): for proj in EXPERT_PROJECTIONS: @@ -43,8 +42,7 @@ def _load(lora_dir, peft_helper, *, moe_ep_spec, lora_id): def test_moe_lora_ep2_real_qwen3moe(qwen3moe_lora_files, ep_rank): """ep_size=2 against the real Qwen3-MoE adapter: each rank's loaded LoRA has the right size, the right expert membership, and the - right tensor values. - """ + right tensor values.""" peft_helper = PEFTHelper.from_local_dir( qwen3moe_lora_files, max_position_embeddings=4096 ) diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index f2bb18aac26e..34cbe061dc42 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -45,8 +45,7 @@ def _bgmv_shrink( inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 ): """Memory-efficient shrink reference: per-LoRA matmul loop. - output[mask] = scaling * inputs[mask] @ weight.T - """ + output[mask] = scaling * inputs[mask] @ weight.T""" exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: @@ -67,8 +66,7 @@ def _bgmv_expand( add_inputs=False, ): """Memory-efficient expand reference: per-LoRA matmul loop. - output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T - """ + output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): if lid < 0: diff --git a/tests/lora/test_punica_ops_fp8.py b/tests/lora/test_punica_ops_fp8.py index 18ac6fcd7aad..e447cef9d8a8 100644 --- a/tests/lora/test_punica_ops_fp8.py +++ b/tests/lora/test_punica_ops_fp8.py @@ -877,8 +877,7 @@ def test_lora_shrink_fp8( quant_mode: str, ): """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline. - """ + quantization, comparing against the bf16 baseline.""" torch.set_default_device(device) set_random_seed(seed) @@ -933,8 +932,7 @@ def test_lora_expand_fp8( quant_mode: str, ): """Test FP8 expand kernel with per-tensor, per-channel, and block-wise - quantization, comparing against the bf16 baseline. - """ + quantization, comparing against the bf16 baseline.""" torch.set_default_device(device) set_random_seed(seed) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 6fdd99deaa44..08fb6818e3ab 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -1319,8 +1319,7 @@ def test_fp8_shared_expert_fse_allows_symmetric_ignored_layers() -> None: def test_fp8_shared_expert_fse_expands_packed_projections() -> None: """Both shards of a fused projection ignored is symmetric within that - projection, so the mismatch reported is against the routed experts. - """ + projection, so the mismatch reported is against the routed experts.""" quant_config = _fp8_config( ignored_layers=[ "model.layers.0.mlp.shared_experts.gate_proj", @@ -1344,8 +1343,7 @@ def test_fp8_shared_expert_fse_expands_packed_projections() -> None: def test_fp8_shared_expert_fse_propagates_partial_shard_exclusion() -> None: """Half a fused projection excluded is rejected by `is_layer_skipped` - itself, exactly as it is in `Fp8Config.get_quant_method`. - """ + itself, exactly as it is in `Fp8Config.get_quant_method`.""" quant_config = _fp8_config( ignored_layers=["model.layers.0.mlp.shared_experts.gate_proj"] ) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 91fc039569dc..41a06d0546f0 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -346,8 +346,7 @@ def record_forward_impl(*args): def test_select_candidate_blocks_tolerates_empty_rows(): """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest block pin must not index -1 for them (device-side assert); they select - no candidate blocks and real rows still pin their newest block. - """ + no candidate blocks and real rows still pin their newest block.""" block_size, topk_blocks = 8, 3 logits = torch.zeros(3, 64, device="cuda") logits[0, 3] = 5.0 # block 0 scores highest for row 0 @@ -458,8 +457,7 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): """Production-faithful candidates: top min(avail, K) blocks per row — - full coverage for lens <= K*cbk, a fixed-size subset beyond that. - """ + full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" avail = ((lens.cpu() + cbk - 1) // cbk).tolist() rows = [] for a in avail: @@ -495,8 +493,7 @@ def _candidate_token_mask( def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid - count are -1 (the indexer's "no token" sentinel). - """ + count are -1 (the indexer's "no token" sentinel).""" values, indices = masked.topk(topk, dim=1) return torch.where(values > float("-inf"), indices.int(), -1) @@ -509,8 +506,7 @@ def _assert_same_selection( ) -> None: """Both top-k index sets must pick the same multiset of logit values per row (-1 padding must match too). Indices themselves may differ: the - top-k kernels break ties among equal values in nondeterministic order. - """ + top-k kernels break ties among equal values in nondeterministic order.""" for r in range(logits.shape[0]): base = int(row_starts[r]) if row_starts is not None else 0 g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] @@ -522,8 +518,7 @@ def _assert_same_selection( def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: """Caller-owned buffers the sparse path writes into (the metadata builder - owns these in production). - """ + owns these in production).""" i32 = dict(dtype=torch.int32, device="cuda") return dict( sparse_indices=torch.zeros(rows, num_sparse, **i32), @@ -538,8 +533,7 @@ def _expand_candidates_reference( """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand each candidate block into ``cbk // sbk`` sparse blocks, keep those inside ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by - repeating the last valid block, and count the valid sparse columns. - """ + repeating the last valid block, and count the valid sparse columns.""" ratio = cbk // sbk rows, num_candidates = candidates.shape num_sparse = num_candidates * ratio @@ -565,8 +559,7 @@ def _expand_candidates_reference( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") def test_candidate_blocks_to_sparse_indices_math(): """Candidate->sparse-block expansion: ks anchoring, range filtering, - repeat-last padding, and the per-row valid column count (``end``). - """ + repeat-last padding, and the per-row valid column count (``end``).""" from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -618,8 +611,7 @@ def test_candidate_blocks_to_sparse_indices_math(): def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): """The Triton expansion equals the PyTorch reference on production-shaped random inputs, including strided (column-sliced) candidate rows and - strided output buffers. - """ + strided output buffers.""" from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( candidate_blocks_to_sparse_indices, ) @@ -689,8 +681,7 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): ``torch.topk`` over dense bf16 logits masked to the candidates (the two DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA graph capture, and a later layer reusing the step's expansion/schedule - selects the same values. - """ + selects the same values.""" deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_prefill_chunk, @@ -780,8 +771,7 @@ def run_sparse(kernel_metadata=None): def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' - fp4 data first, then their packed UE8M0 scales. - """ + fp4 data first, then their packed UE8M0 scales.""" device = req_lens.device pages_per = (req_lens + page_kv - 1) // page_kv num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) @@ -817,8 +807,7 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to one query row each), CUDA graph capture/replay and schedule reuse by a - later layer. - """ + later layer.""" deep_gemm = _skip_unless_sm100_sparse_kernels() from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( sparse_mqa_logits_paged_decode, @@ -927,8 +916,7 @@ def run_sparse(o: torch.Tensor, kernel_metadata=None): def test_sparse_mqa_logits_accuracy_vs_torch(): """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized - MXFP4 inputs. - """ + MXFP4 inputs.""" _skip_unless_sm100_sparse_kernels() from deep_gemm.utils import cast_back_from_fp4 diff --git a/tests/model_executor/model_loader/test_ep_weight_filter.py b/tests/model_executor/model_loader/test_ep_weight_filter.py index 9d177f59910f..0ff53d24e71d 100644 --- a/tests/model_executor/model_loader/test_ep_weight_filter.py +++ b/tests/model_executor/model_loader/test_ep_weight_filter.py @@ -232,8 +232,7 @@ def test_fused_3d_expert_not_skipped(self): class TestSafetensorsWeightsIteratorWithEpFilter: """Verify that EP filtering produces a strict subset of unfiltered loading - and that all expected dense + local expert weights are present. - """ + and that all expected dense + local expert weights are present.""" @pytest.fixture(scope="class") def gpt2_files(self): @@ -263,8 +262,7 @@ def test_no_filter_returns_all(self, gpt2_files): def test_empty_filter_skips_experts_only(self, gpt2_files): """GPT-2 has no expert weights, so even an empty local_expert_ids - set should return all weights (all are dense). - """ + set should return all weights (all are dense).""" all_weights = dict(safetensors_weights_iterator(gpt2_files, False)) filtered_weights = dict( safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set()) @@ -275,8 +273,7 @@ def test_empty_filter_skips_experts_only(self, gpt2_files): class TestEpFilterOnSyntheticMoeWeights: """Create synthetic safetensors files with expert-like naming and verify - that the filter correctly skips non-local experts. - """ + that the filter correctly skips non-local experts.""" @pytest.fixture def synthetic_moe_files(self, tmp_path): diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index d123dcb7508b..b4a544d2fc0f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -350,8 +350,7 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold. - """ + reuse the workspace storage whose address captured CUDA graphs hold.""" from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -429,8 +428,7 @@ def load_checkpoint_format_weights(): def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that. - """ + graphs is already unusable; allocating a replacement would hide that.""" from vllm.model_executor.layers.quantization.utils import marlin_utils monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) @@ -611,8 +609,7 @@ def process_weights_after_loading(self, layer): class _LateBiasLayer(torch.nn.Module): """Mimics an online-quantized linear: `weight` is created on meta by `create_weights()`, which wraps the loaders, and the linear base registers - `bias` afterwards. - """ + `bias` afterwards.""" def __init__(self, quant_method): super().__init__() diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index ada073706606..e71061e73317 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -42,8 +42,7 @@ def _run_warmup(monkeypatch, layers) -> list[dict]: @pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"]) def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None: """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the - quantization method that owns it or the name of its scale parameter. - """ + quantization method that owns it or the name of its scale parameter.""" layer = _block_fp8_layer(256, 128, scale_name) calls = _run_warmup(monkeypatch, [layer]) @@ -55,8 +54,7 @@ def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> Non def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None: """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576), - so warmup must not require N % 128 == 0. - """ + so warmup must not require N % 128 == 0.""" assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1 diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index ed10955f963a..b146a235c1c9 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -326,8 +326,7 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises(): def test_routed_experts_attention_group_is_shared_and_fail_closed(): """Both sides key routing data by this gid, so it must skip non-full-attention - groups rather than defaulting to 0, and fail closed when none exists. - """ + groups rather than defaulting to 0, and fail closed when none exists.""" common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) config = SimpleNamespace( kv_cache_groups=[ diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index d19ea0ed05b3..3dea587075f2 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -57,8 +57,7 @@ class TestMaybeRemapKvScaleName: def test_qkv_proj_k_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047 - """ + Regression test for https://github.com/vllm-project/vllm/issues/25047""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -66,8 +65,7 @@ def test_qkv_proj_k_scale(self): def test_qkv_proj_v_scale(self): """Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale - Regression test for https://github.com/vllm-project/vllm/issues/25047 - """ + Regression test for https://github.com/vllm-project/vllm/issues/25047""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) @@ -110,8 +108,7 @@ def test_non_scale_name_unchanged(self): def test_nvfp4_modelopt_k_proj_k_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): k_proj.k_scale -> attn.k_scale. - Validates that NVFP4 checkpoints are not broken by this change. - """ + Validates that NVFP4 checkpoints are not broken by this change.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT ) @@ -120,8 +117,7 @@ def test_nvfp4_modelopt_k_proj_k_scale(self): def test_nvfp4_modelopt_v_proj_v_scale(self): """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4): v_proj.v_scale -> attn.v_scale. - Validates that NVFP4 checkpoints are not broken by this change. - """ + Validates that NVFP4 checkpoints are not broken by this change.""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT ) @@ -129,8 +125,7 @@ def test_nvfp4_modelopt_v_proj_v_scale(self): def test_qwen3_vl_moe_qkv_proj_k_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047). - """ + Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT ) @@ -138,8 +133,7 @@ def test_qwen3_vl_moe_qkv_proj_k_scale(self): def test_qwen3_vl_moe_qkv_proj_v_scale(self): """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE. - Regression test for qwen3_vl_moe.py fix (same bug as #25047). - """ + Regression test for qwen3_vl_moe.py fix (same bug as #25047).""" result = maybe_remap_kv_scale_name( "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT ) @@ -181,8 +175,7 @@ class TestKvCacheScaleMapper: """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent - so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`. - """ + so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`.""" def _mapper(self): # `get_cache_scale_mapper` does not use `self`; call it on the base @@ -272,8 +265,7 @@ def test_idempotent(self, name): def test_composes_with_qkv_mapper(self): """Applied together with a model's qkv/gate_up mapper, the regex scale rules run before the substr rename, so scales are normalized to `.attn.` - and regular projections are still fused correctly. - """ + and regular projections are still fused correctly.""" from vllm.model_executor.models.utils import WeightsMapper model_mapper = WeightsMapper( diff --git a/tests/models/inkling/test_mm_towers.py b/tests/models/inkling/test_mm_towers.py index 51daa506c2b4..d969c8101a29 100644 --- a/tests/models/inkling/test_mm_towers.py +++ b/tests/models/inkling/test_mm_towers.py @@ -42,8 +42,7 @@ def _assert_close_ulps( compare against the reference's local ulp spacing. ``pre_act`` (the reference pre-activation) adds derivative-propagated slack: near ``gelu(x) ~ 0`` an ulp-level input flip legitimately moves the output by - many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps. - """ + many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps.""" g, r = got.float(), ref.float() tol = torch.clamp(max_ulps * _bf16_spacing(ref), min=atol) if pre_act is not None: diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index 4346a7292846..ec270c04fe6c 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -123,8 +123,7 @@ def _check_writes_only_its_own_shard( device: torch.device, tp_size: int, rank: int ) -> None: """Two ranks that swapped both weight rows and write offsets still sum to - the right total, so inspect each slice before the final collective. - """ + the right total, so inspect each slice before the final collective.""" transform = _build_transform(device) runner = _tail_runner(transform, tp_size) group = get_tp_group().device_group diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py index 64cec452d25f..26cb76a55039 100644 --- a/tests/models/kimi_k3/test_aux_attn_res_stream.py +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -130,8 +130,7 @@ def test_disabled_reproduces_the_plain_residual_sum( def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): """The value the next layer reads is the mixture against *its* weights, - so the tap has to reach forward rather than use the current layer's. - """ + so the tap has to reach forward rather than use the current layer's.""" _set_last_rank(monkeypatch, True) _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) @@ -182,8 +181,7 @@ def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( ): """The kernel writes an applied delta back into the prefix in place, which would double-add it into the live residual stream, so the pending output - has to arrive already summed into the prefix with `delta` left None. - """ + has to arrive already summed into the prefix with `delta` left None.""" _set_last_rank(monkeypatch, True) prefix_sum = torch.tensor([1.0, 2.0]) pending = torch.tensor([0.5, 0.25]) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 5dad5754b634..fa810217b188 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -727,8 +727,7 @@ def test_cudagraph_capture_metadata_avoids_device_to_host_copy( """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing scheduled re-stages its FULL-graph metadata through it on every dummy step, so it must not synchronize the device. The host-side draft counts have to - come from `query_start_loc_cpu` and match the device-derived values. - """ + come from `query_start_loc_cpu` and match the device-derived values.""" num_speculative_tokens = 2 batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index 307a572afd1f..c2c8e88e427c 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -60,8 +60,7 @@ def __init__(self, *, vllm_config, prefix): def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: """KV-shared layers have q_proj instead of a packed qkv_proj; their redundant K/V tensors in original checkpoints have no parameter and are - skipped rather than failing the load. - """ + skipped rather than failing the load.""" model = torch.nn.Module() model.config = SimpleNamespace(num_experts=0) model.start_layer, model.end_layer = 0, 2 @@ -93,8 +92,7 @@ def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None: @pytest.mark.cpu_test def test_gemma3n_kv_shared_layer_mapper() -> None: """Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep - q_proj and drop the redundant K/V tensors original checkpoints ship. - """ + q_proj and drop the redundant K/V tensors original checkpoints ship.""" config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2) mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config) weights = [ diff --git a/tests/models/language/generation/test_mistral.py b/tests/models/language/generation/test_mistral.py index a21c5c355a66..fc774bc6e835 100644 --- a/tests/models/language/generation/test_mistral.py +++ b/tests/models/language/generation/test_mistral.py @@ -274,8 +274,7 @@ def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None: def test_mistral_function_call_nested_json(): """Ensure that the function-name regex captures the entire outermost - JSON block, including nested braces. - """ + JSON block, including nested braces.""" # Create a minimal stub tokenizer that provides the few attributes the # parser accesses (`version` and `get_vocab`). diff --git a/tests/models/language/pooling/test_splade_sparse_pooler.py b/tests/models/language/pooling/test_splade_sparse_pooler.py index d470637477c0..0b3561f90fd1 100644 --- a/tests/models/language/pooling/test_splade_sparse_pooler.py +++ b/tests/models/language/pooling/test_splade_sparse_pooler.py @@ -29,8 +29,7 @@ @torch.inference_mode def test_splade_pooler_matches_reference_formula(B, T, H, V): """Ensure SPLADESparsePooler forward() matches the mathematical formula: - log1p(relu(logits)) -> max over sequence length (after masking). - """ + log1p(relu(logits)) -> max over sequence length (after masking).""" torch.manual_seed(0) # Prepare [B] sequences of shape [T, H] diff --git a/tests/models/multimodal/generation/test_moondream3_moe.py b/tests/models/multimodal/generation/test_moondream3_moe.py index 41654d4a7209..79d18627656b 100644 --- a/tests/models/multimodal/generation/test_moondream3_moe.py +++ b/tests/models/multimodal/generation/test_moondream3_moe.py @@ -33,8 +33,7 @@ class _StubGate(torch.nn.Module): """Router stand-in, so the test does not depend on platform-specific - ``ReplicatedLinear`` weight post-processing. - """ + ``ReplicatedLinear`` weight post-processing.""" def __init__(self, hidden_size: int, num_experts: int): super().__init__() diff --git a/tests/models/multimodal/generation/test_transformers_audio.py b/tests/models/multimodal/generation/test_transformers_audio.py index 11ed8e51786a..919557987796 100644 --- a/tests/models/multimodal/generation/test_transformers_audio.py +++ b/tests/models/multimodal/generation/test_transformers_audio.py @@ -81,8 +81,7 @@ def test_transformers_audio_generation( under pytest can cause (possibly flaky) hangs, so they are grouped under the same config. Using VLLM_WORKER_MULTIPROC_METHOD=spawn avoids the deadlock and allows worker processes to terminate cleanly, and release - GPU memory between test runs until the issue is fixed. - """ + GPU memory between test runs until the issue is fixed.""" # TODO: Remove monkeypatch once # https://github.com/vllm-project/vllm/issues/17676 is fixed. disable_envs_cache() diff --git a/tests/models/multimodal/generation/vlm_utils/builders.py b/tests/models/multimodal/generation/vlm_utils/builders.py index 1f4783473d87..1b7e2347be2f 100644 --- a/tests/models/multimodal/generation/vlm_utils/builders.py +++ b/tests/models/multimodal/generation/vlm_utils/builders.py @@ -294,8 +294,7 @@ def sample_frames_with_video_metadata( def apply_image_size_scaling(image, size: float | tuple[int, int], size_type: SizeType): """Applies a size scaler to one image; this can be an image size factor, - which scales the image while maintaining the aspect ratio - """ + which scales the image while maintaining the aspect ratio""" # Special case for embeddings; if it's a tensor, it's only valid if we # are considering size factors at constant scale, i.e., we just clone # the tensor diff --git a/tests/models/multimodal/processing/test_deepseek_ocr.py b/tests/models/multimodal/processing/test_deepseek_ocr.py index 1851244be085..2bcf5e05d37e 100644 --- a/tests/models/multimodal/processing/test_deepseek_ocr.py +++ b/tests/models/multimodal/processing/test_deepseek_ocr.py @@ -108,8 +108,7 @@ def test_populated_images_crop_large_image(self, processor): def test_mismatched_image_size_raises(self, processor): """Deliberately wrong image_size binding should still be caught - by TensorSchema validation. - """ + by TensorSchema validation.""" small_image = Image.new("RGB", (100, 100), color="green") result = processor( @@ -143,8 +142,7 @@ class TestDeepseekOCRInputValidation: def test_mismatched_image_count_raises_value_error(self, processor): """Two ```` tokens with one image must raise ValueError, not AssertionError (which would yield HTTP 500 and, under ``-O``, silently - drop the middle text segment). - """ + drop the middle text segment).""" image = Image.new("RGB", (100, 100), color="red") with pytest.raises(ValueError, match="does not match number of images"): processor(prompt=" and ", images=[image]) diff --git a/tests/models/multimodal/processing/test_llava_next_video.py b/tests/models/multimodal/processing/test_llava_next_video.py index 512440ee8ee6..c4fdecd21ac7 100644 --- a/tests/models/multimodal/processing/test_llava_next_video.py +++ b/tests/models/multimodal/processing/test_llava_next_video.py @@ -16,8 +16,7 @@ class _StubModel: """Carries only the attributes the token-count method reads from `self`, so the real methods can be exercised without constructing - the full `nn.Module` (vision tower, language model, etc.). - """ + the full `nn.Module` (vision tower, language model, etc.).""" patch_grid_length: int pooled_grid_length: int @@ -32,8 +31,7 @@ class _StubModel: def test_num_mm_tokens_match_real_config(model_id): """The stored grid lengths must match what `__init__` derives from the real HF config, and the LoRA token counts must preserve the - frame-level scaling exactly. - """ + frame-level scaling exactly.""" ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1}) config = ctx.model_config.hf_config diff --git a/tests/models/multimodal/processing/test_molmo2.py b/tests/models/multimodal/processing/test_molmo2.py index 98ad241650e0..5633d47a97dc 100644 --- a/tests/models/multimodal/processing/test_molmo2.py +++ b/tests/models/multimodal/processing/test_molmo2.py @@ -70,8 +70,7 @@ def test_dummy_video_num_frames_override_honors_min_of_two( ): """A ``num_frames`` override below 2 must be ignored (the model needs at least 2 frames), matching the "cannot be less than 2, will be ignored" - warning. - """ + warning.""" builder = object.__new__(Molmo2DummyInputsBuilder) builder.info = SimpleNamespace( get_hf_processor=lambda: SimpleNamespace( diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index 47071c2f1bed..c58904a9e2c9 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -31,8 +31,7 @@ def _build_video_mm_data( ``total_num_frames`` is set equal to the ndarray frame count so that HF's ``sample_frames`` indices stay within bounds of the - actual tensor that is passed. - """ + actual tensor that is passed.""" video = np.zeros((num_frames, height, width, 3), dtype=np.uint8) metadata = { "fps": original_fps, @@ -238,8 +237,7 @@ def _build_video_embeds_mm_data( grid_thw: tuple[int, int, int] = (2, 4, 4), ) -> dict[str, Any]: """Create an embeds-only video item as an EPD consumer receives it: - pre-computed embeddings plus the metadata published by the encoder. - """ + pre-computed embeddings plus the metadata published by the encoder.""" import torch t, h, w = grid_thw @@ -258,8 +256,7 @@ def test_processor_video_embeds_with_timestamps(model_id: str) -> None: """Embeds-only video input must size the placeholder range from the grid and the real timestamps published by the encoder (EC consumer path); synthesized or missing timestamps would change the token count - and break embedding merging downstream. - """ + and break embedding merging downstream.""" # `build_model_context` forces enable_mm_embeds off, so build the # config directly; keep the same online-availability skip behavior. HF_EXAMPLE_MODELS.find_hf_info(model_id).check_available_online(on_fail="skip") @@ -304,8 +301,7 @@ def test_processor_video_embeds_with_timestamps(model_id: str) -> None: def test_processor_video_embeds_missing_timestamps(model_id: str) -> None: """Timestamps are required metadata for video embeds: they size the placeholder range, so omitting them must fail loudly at parse time - instead of silently producing a wrong prompt. - """ + instead of silently producing a wrong prompt.""" ctx = build_model_context( model_id, limit_mm_per_prompt={"image": 0, "video": 1}, diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py index 14007c87403d..61517ea0af19 100644 --- a/tests/models/multimodal/processing/test_transformers_audio.py +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -116,8 +116,7 @@ def test_audio_multimodal_processor(model_id, processor_cls): @pytest.mark.parametrize("separator", [" and ", ""]) def test_audio_multiple_inputs(separator, processor_cls): """Multiple audios per prompt are each detected as a separate placeholder - and multi-modal item by the Transformers modelling backend. - """ + and multi-modal item by the Transformers modelling backend.""" model_id = "ibm-granite/granite-speech-3.3-2b" mm_processor = create_processor(model_id, processor_cls) diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index 731794da2667..f5de7d66f4ff 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -83,8 +83,7 @@ def _process_two_images(processor_cls, separator: str): @pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) def test_image_multiple_inputs(processor_cls): """Multiple images per prompt are each detected as a separate placeholder - and multi-modal item by the Transformers modelling backend. - """ + and multi-modal item by the Transformers modelling backend.""" result = _process_two_images(processor_cls, separator="\n and ") assert len(result["mm_placeholders"]["image"]) == 2 @@ -103,8 +102,7 @@ def test_image_adjacent_inputs(processor_cls): @pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) def test_batch_padding_removed_from_image_items(processor_cls): """Emu3 pads every image up to the largest in the batch, which would leave an - item's data dependent on what it was processed with and so uncacheable. - """ + item's data dependent on what it was processed with and so uncacheable.""" mm_processor = create_processor("BAAI/Emu3-Chat-hf", processor_cls) image_token = mm_processor.info.get_hf_processor().image_token @@ -146,8 +144,7 @@ def _process_one_gemma3_image(processor_cls): @offsets_only def test_non_embedding_tokens_excluded_from_placeholders(): """Gemma3 wraps each image in text that carries no embeddings, which must be - inside the placeholder range but masked out of it. - """ + inside the placeholder range but masked out of it.""" _, result = _process_one_gemma3_image(OffsetsMultiModalProcessor) (placeholder,) = result["mm_placeholders"]["image"] @@ -159,8 +156,7 @@ def test_legacy_placeholders_hold_only_image_tokens(): """The legacy path spans whatever `mm_token_type_ids` attributes to the image, which for Gemma3 excludes the text wrapping it, unlike the replacement the offsets path spans. Gemma3 is also the sharp case for the mask: its `image_token_id` is the - marker in the unexpanded prompt, not the token the expansion repeats. - """ + marker in the unexpanded prompt, not the token the expansion repeats.""" hf_processor, result = _process_one_gemma3_image(LegacyMultiModalProcessor) (placeholder,) = result["mm_placeholders"]["image"] @@ -176,8 +172,7 @@ def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): embeddings. Those belong inside the placeholder and masked out, because the token count the processor reports is over the whole span. Idefics3 also refuses a prompt holding `` when no images are passed, which is how the offsets path has to - tokenize it before splicing in the expansion. - """ + tokenize it before splicing in the expansion.""" mm_processor = create_processor( "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls ) @@ -197,8 +192,7 @@ def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): @offsets_only def test_missing_replacement_offsets_names_the_processor(): """A processor that reports no replacement offsets cannot be served, which must - be said plainly rather than surfacing later as a field config mismatch. - """ + be said plainly rather than surfacing later as a field config mismatch.""" model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" mm_processor = create_processor(model_id, OffsetsMultiModalProcessor) hf_processor_cls = type(mm_processor.info.get_hf_processor()) @@ -282,8 +276,7 @@ def process(): def test_spliced_prompt_matches_hf_expansion(model_id, prompt): """The offsets path splices the expansion into a prompt tokenized without any multi-modal data, so its token ids have to come out the same as the ones the HF - processor produces itself, which is what the legacy path returns. - """ + processor produces itself, which is what the legacy path returns.""" prompt_ids = [] for processor_cls in (LegacyMultiModalProcessor, OffsetsMultiModalProcessor): mm_processor = create_processor(model_id, processor_cls) @@ -305,8 +298,7 @@ def test_spliced_prompt_matches_hf_expansion(model_id, prompt): def test_nested_image_fields_split_per_image(processor_cls): """Idefics3 returns image fields with a leading batch dimension, putting the rows belonging to each image one dimension further in. Slicing the batch dimension - instead handed the first image every row and the second an empty tensor. - """ + instead handed the first image every row and the second an empty tensor.""" mm_processor = create_processor( "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls ) diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py index d748e89f20f9..dc49036e8352 100644 --- a/tests/models/multimodal/processing/transformers_backend.py +++ b/tests/models/multimodal/processing/transformers_backend.py @@ -29,8 +29,7 @@ def create_processor(model_id: str, processor_cls): """Build a processor directly, because the registry only ever builds the one the - installed transformers version selects, leaving the other path untested. - """ + installed transformers version selects, leaving the other path untested.""" model_config = ModelConfig(model=model_id, model_impl="transformers") ctx = InputProcessingContext( model_config, cached_tokenizer_from_config(model_config) @@ -41,8 +40,7 @@ def create_processor(model_id: str, processor_cls): def create_cached_processor(model_id: str, processor_cls): """Build a processor backed by a real multi-modal processor cache, and hand the - cache back so a test can check it was actually used. - """ + cache back so a test can check it was actually used.""" model_config = ModelConfig(model=model_id, model_impl="transformers") model_config.multimodal_config.mm_processor_cache_gb = 4 ctx = InputProcessingContext( diff --git a/tests/models/multimodal/test_cohere_asr.py b/tests/models/multimodal/test_cohere_asr.py index a9c45f84055c..85a9b59dd513 100644 --- a/tests/models/multimodal/test_cohere_asr.py +++ b/tests/models/multimodal/test_cohere_asr.py @@ -18,8 +18,7 @@ def test_get_num_audio_tokens_streaming_estimate(audio_duration_s, expected_tokens): """The duration-based estimate must convert ``window_stride`` (seconds) to a sample hop and divide by the encoder subsampling factor. Values are pinned to - concrete numbers so a rounding regression is actually caught. - """ + concrete numbers so a rounding regression is actually caught.""" model_config = SimpleNamespace( hf_config=SimpleNamespace( preprocessor={"window_stride": 0.01, "sample_rate": 16000}, diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py index a167fe205db7..95d60914218a 100644 --- a/tests/models/multimodal/test_mapping.py +++ b/tests/models/multimodal/test_mapping.py @@ -78,8 +78,7 @@ def test_cosmos3_modelopt_quantizer_weights_mapper(): """ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers (``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable ``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former - (it has no parameter for them) while keeping the latter. - """ + (it has no parameter for them) while keeping the latter.""" from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index c64073ffe2be..75900cde53a3 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -81,8 +81,7 @@ class _FakeTensor: """Sentinel stand-in for torch.Tensor in load_weights tests. Supports the .detach().clone() chain used by load_weights for buffered mm weights; both methods return self so identity (and the existing equality - assertions) are preserved through cloning. - """ + assertions) are preserved through cloning.""" def detach(self): return self @@ -166,8 +165,7 @@ def _make_mm_items_with_video_bytes( def test_extract_audio_from_videos_passes_max_duration(): """_extract_audio_from_videos must forward VLLM_MAX_AUDIO_DECODE_DURATION_S - to load_audio_pyav so decompression-bomb audio is rejected. - """ + to load_audio_pyav so decompression-bomb audio is rejected.""" dummy_audio = (np.zeros(16000, dtype=np.float32), 16000.0) mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64) @@ -187,8 +185,7 @@ def test_extract_audio_from_videos_passes_max_duration(): def test_extract_audio_from_videos_rejects_oversized_audio(): """When load_audio_pyav raises due to duration limit the video is - marked as having no audio instead of crashing the server. - """ + marked as having no audio instead of crashing the server.""" mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64) processor = object.__new__(NanoNemotronVLMultiModalProcessor) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index c5ed11d0c693..9d30365b8151 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -144,8 +144,7 @@ def test_qwen4_exp_mtp_override_sets_draft_config( def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> None: """PLE needs raw input_ids, which non-first pipeline ranks never see. The rest of the architecture is PP-capable, so the refusal must be conditional - -- and must land before the engine spends time loading weights. - """ + -- and must land before the engine spends time loading weights.""" vllm_config = SimpleNamespace( model_config=SimpleNamespace( hf_text_config=_text_config(ple_layer_ids=ple_layer_ids), diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 7007eb5aca59..99d661e75598 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -103,8 +103,7 @@ def test_fi_backend_strings_are_registered_mega_moe_backends(): @pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend): """An FI moe_ep backend with a non-DSv4 model must fail at config time - instead of silently falling through to the generic FusedMoE path. - """ + instead of silently falling through to the generic FusedMoE path.""" with pytest.raises(ValueError, match="only supported for DeepSeek-V4"): validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"]) @@ -120,8 +119,7 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): ) def test_native_deep_gemm_mega_moe_not_arch_gated(architectures): """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it); - models validate their own constraints at construction time. - """ + models validate their own constraints at construction time.""" validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures) @@ -134,8 +132,7 @@ def test_all_mega_backends_get_sequence_parallel_moe(moe_backend): """Every mega backend must qualify for sequence-parallel MoE at TP>1/EP: the predicate once matched only the native backend string, which silently ran the fi backends full-batch with an all-reduce on - every rank — 0.42-0.65x native e2e at TP8. - """ + every rank — 0.42-0.65x native e2e at TP8.""" from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel vllm_config = SimpleNamespace( @@ -179,8 +176,7 @@ def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch): """The runtime must not rederive the device from LOCAL_RANK/rank: under a remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the - device it already bound via BootstrapConfig.device. - """ + device it already bound via BootstrapConfig.device.""" import vllm.utils.flashinfer_moe_ep as mod pg = object() @@ -245,8 +241,7 @@ def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo(): def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales(): """One 32-element scale group per row: low nibble is the even element, - high nibble the odd one, ue8m0 scale applies to the whole group. - """ + high nibble the odd one, ue8m0 scale applies to the whole group.""" packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16) sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 31aad5395f4b..9a5723d233a1 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -783,8 +783,7 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): """The loader invokes the model-level PWAL hook for every load format, so it must finalize megamoe + mhc broadcast weights to cover dummy - load, which skips load_weights(). - """ + load, which skips load_weights().""" calls = [] stub = SimpleNamespace( model=SimpleNamespace( @@ -801,8 +800,7 @@ def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): """MTP/DSpark drafters load as their own top-level models, so each needs its own PWAL hook now that the megamoe forward no longer finalizes - weights lazily on first use. - """ + weights lazily on first use.""" calls = [] mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) DeepSeekV4MTP.process_weights_after_loading(mtp) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index ee381ee36b22..d30eb9e7c91d 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -215,8 +215,7 @@ def test_v41_dspark_loads_linear_scales( ): """Checkpoint ``.scale`` maps to the quant method's scale parameter and loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static - loader (see tests/quantization/test_modelopt.py), not in load_weights. - """ + loader (see tests/quantization/test_modelopt.py), not in load_weights.""" from vllm.models.deepseek_v41.nvidia import dspark mxfp8 = scale_dtype != torch.float32 diff --git a/tests/models/test_qwen3_5_mtp_config.py b/tests/models/test_qwen3_5_mtp_config.py index 3c098d3ad55e..6cf54dbc95a8 100644 --- a/tests/models/test_qwen3_5_mtp_config.py +++ b/tests/models/test_qwen3_5_mtp_config.py @@ -73,8 +73,7 @@ def test_mtp_override_extracts_n_predict_from_multimodal_wrapper( model_type: str, expected_arch: str ) -> None: """Verify that multimodal wrapper checkpoints with mtp_num_hidden_layers - in text_config resolve n_predict and architecture correctly. - """ + in text_config resolve n_predict and architecture correctly.""" cfg = SpeculativeConfig.hf_config_override( _multimodal_wrapper_mtp_config(model_type, mtp_layers=2) ) @@ -94,8 +93,7 @@ def test_mtp_override_top_level_precedence_over_nested_text_config( model_type: str, expected_arch: str ) -> None: """Verify that an explicit top-level mtp_num_hidden_layers takes precedence - over a nested text_config value. - """ + over a nested text_config value.""" cfg = _multimodal_wrapper_mtp_config(model_type, mtp_layers=2) cfg.mtp_num_hidden_layers = 3 overridden = SpeculativeConfig.hf_config_override(cfg) @@ -115,8 +113,7 @@ def test_mtp_override_downloads_real_hf_hub_configs( model_id: str, expected_arch: str ) -> None: """Verify that unmodified real-world checkpoints downloaded via - AutoConfig.from_pretrained resolve n_predict=1 from text_config. - """ + AutoConfig.from_pretrained resolve n_predict=1 from text_config.""" hf_config = AutoConfig.from_pretrained(model_id) cfg = SpeculativeConfig.hf_config_override(hf_config) assert cfg.model_type == "qwen3_5_mtp" diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index f9e7d691769b..02903504c684 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -197,8 +197,7 @@ def test_merge_multimodal_embeddings_no_sync(): def test_get_rename_mapper_keeps_only_renames(): """`None` means "do not load", which is meaningless to the consumers of this mapper (LoRA name parsing, quantization config layer lists), and - applying it would silently shrink their lists. - """ + applying it would silently shrink their lists.""" mapper = WeightsMapper( orig_to_new_regex={re.compile(r"^drop_regex\."): None}, orig_to_new_substr={"drop_substr": None, "keep_substr": "kept"}, diff --git a/tests/models/test_vision.py b/tests/models/test_vision.py index b08661abc958..94edce45bf53 100644 --- a/tests/models/test_vision.py +++ b/tests/models/test_vision.py @@ -512,8 +512,7 @@ def _reference_input_norm( def test_fused_input_norm_matches_reference(num_patches: int): """FusedInputNorm must equal the plain affine, including for num_patches above the cuDNN batch-norm grid limit (~65535) that previously raised - CUDNN_STATUS_INTERNAL_ERROR (issue #51717). - """ + CUDNN_STATUS_INTERNAL_ERROR (issue #51717).""" channel = 3 patch_size = 14 * 14 image_mean = [0.48145466, 0.4578275, 0.40821073] diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index d26d742dd3d2..bcad62765a76 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -251,8 +251,7 @@ def forward(self, x): class PlainQKV(nn.Module): """q/k/v returned as a bare tuple: no reshape stands between the calls and - the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it. - """ + the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it.""" def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8): super().__init__() @@ -339,8 +338,7 @@ class GuardedVAttention(FakeAttention): """Gemma 4-style: `v_proj` is called once but also read in a `None` guard. The guard is a fusion invariant (the fuser only binds instances where the - projection exists), so it folds to `True` and the call still fuses. - """ + projection exists), so it folds to `True` and the call still fuses.""" def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -370,8 +368,7 @@ class TruthyGuardAttention(FakeAttention): An `nn.Linear` is always truthy, so the guard folds to `True` and the call still fuses (the projection is a proven invariant, exactly as for a `None` - guard). - """ + guard).""" def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -402,8 +399,7 @@ class StatementGuardVAttention(FakeAttention): The `if`'s test is not an `ast.Compare`, so the identity-check scan skips it; `_in_boolean_context` folds the test to `True` (an `nn.Linear` is always truthy) and the call inside the now-`if True:` body still fuses, with the - live `else` surviving the rewrite. - """ + live `else` surviving the rewrite.""" def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -432,8 +428,7 @@ class BranchedKVAttention(FakeAttention): The three calls do not share a block, so the fused GEMM must be hoisted to the innermost block that dominates all of them (the function body). The - live `if`/`else` must survive the rewrite. - """ + live `if`/`else` must survive the rewrite.""" is_kv_shared_layer = False @@ -464,8 +459,7 @@ class RebindArgAttention(FakeAttention): The branch is dead at trace time so the match succeeds, but the source rewrite must refuse: hoisting the GEMM above the rebind would feed k/v a - different input. Result: no fusion. - """ + different input. Result: no fusion.""" recompute = False @@ -496,8 +490,7 @@ class InPlaceMutatedArgAttention(FakeAttention): Like `RebindArgAttention`, but the mutation keeps the name in a `Load` context (`hidden_states.mul_(2)`), so a bare `Name`-store check would miss it. The rewrite must still refuse: hoisting the GEMM above the mutation changes - the value k/v see. Result: no fusion. - """ + the value k/v see. Result: no fusion.""" recompute = False @@ -527,8 +520,7 @@ class AliasMutatedArgAttention(FakeAttention): `hidden_states.view(-1)` shares storage with its base, so `alias.mul_(2)` changes what k/v read without ever naming `hidden_states` as a target. Only - tracking names that may alias the input catches it. Result: no fusion. - """ + tracking names that may alias the input catches it. Result: no fusion.""" recompute = False @@ -559,8 +551,7 @@ class FunctionalInplaceArgAttention(FakeAttention): `F.relu(hidden_states, inplace=True)` writes through an argument, so the name appears only in a `Load` context inside a call that is not a method on - it -- invisible to a trailing-underscore method check. Result: no fusion. - """ + it -- invisible to a trailing-underscore method check. Result: no fusion.""" recompute = False @@ -590,8 +581,7 @@ class SubscriptMutatedArgAttention(FakeAttention): `hidden_states[..., 0] = 0` leaves the name in a `Load` context on the subscript's value, so it too evades a bare `Name`-store check. The rewrite - must refuse for the same reason. Result: no fusion. - """ + must refuse for the same reason. Result: no fusion.""" recompute = False @@ -621,8 +611,7 @@ class NonFoldableRefAttention(FakeAttention): Deleting `v_proj` would break `isinstance(self.v_proj, nn.Linear)`: it reads the projection's type, not just whether it exists, so the fuser cannot fold - it and must refuse rather than rewrite it. - """ + it and must refuse rather than rewrite it.""" def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs @@ -847,8 +836,7 @@ def test_glu_identifies_down_projection(): It is forced to `RowParallelLinear` in `update_attrs` so its sharded input matches the column-parallel merged gate/up; `None` when there is no such - projection to force (fusion of gate/up still applies). - """ + projection to force (fusion of gate/up still applies).""" with torch.device("meta"): assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj" assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj" @@ -1009,8 +997,7 @@ def test_bypass_existence_guard_folds_identity_none_check(expr, expected): The projection is a proven invariant (it exists as an `nn.Linear`), so an `is (not) None` guard has a constant truth value; folding it lets the call be rewritten away without changing the guard's outcome. `==`/`!=` are excluded - (see the refusal test) because a subclass may override `__eq__`/`__ne__`. - """ + (see the refusal test) because a subclass may override `__eq__`/`__ne__`.""" funcdef, ref = _guard_funcdef(expr) bypass_existence_guard(funcdef, ref, "v_proj") test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.IfExp)).test @@ -1029,8 +1016,7 @@ def test_bypass_existence_guard_folds_bare_truthiness(expr): An `nn.Linear` has no `__bool__`/`__len__`, so it is always truthy; in a pure boolean position (a test, or a `not` operand) the reference can be replaced by - `True`, letting the projection be deleted. - """ + `True`, letting the projection be deleted.""" funcdef, ref = _guard_funcdef(expr) bypass_existence_guard(funcdef, ref, "v_proj") # No `self.v_proj` reference survives; a literal `True` took its place. @@ -1045,8 +1031,7 @@ def test_bypass_existence_guard_folds_bare_if_statement(): The `len(node.ops) == 1` `Compare` scan skips it (it is not a comparison at all); `_in_boolean_context` then matches the `if`'s test and folds it, so the - bare-truthiness statement form is handled, not silently dropped. - """ + bare-truthiness statement form is handled, not silently dropped.""" funcdef = ast.parse( "def f(self, a, b):\n if self.v_proj:\n return a\n return b" ).body[0] @@ -1078,8 +1063,7 @@ def test_bypass_existence_guard_refuses_non_guard_reference(expr): Rewriting it is unsafe (its value depends on more than the projection's existence), so the fuser refuses rather than change semantics. `==`/`!=` are refused too: `is_linear` accepts `nn.Linear` subclasses that could override - `__eq__`/`__ne__`, so equality need not agree with `is (not) None`. - """ + `__eq__`/`__ne__`, so equality need not agree with `is (not) None`.""" funcdef, ref = _guard_funcdef(expr) with pytest.raises(ValueError, match="outside an existence guard"): bypass_existence_guard(funcdef, ref, "v_proj") @@ -1119,8 +1103,7 @@ def test_merged_column_fuser_folds_existence_guard(): Guard folding lives on `StackedFuser`, not just `QKVFuser`, because every stacked fuser deletes the projections it merges -- a surviving reference to - one would fail at runtime. - """ + one would fail at runtime.""" with torch.device("meta"): module = GuardedParallelLinears() fuser = MergedColumnParallelFuser.match(trace(module), module) @@ -1146,8 +1129,7 @@ def test_merged_column_fuser_rejects_input_mutation(cls): The mutation may be direct (`x.add_(1)`), through a view that shares the input's storage, or through a call that writes an argument in place (`F.relu(x, inplace=True)`, `torch.relu_(x)`). Each one silently changes the - fused result, so the rewrite must refuse rather than fuse. - """ + fused result, so the rewrite must refuse rather than fuse.""" with torch.device("meta"): module = cls() fuser = MergedColumnParallelFuser.match(trace(module), module) @@ -1205,8 +1187,7 @@ def test_merged_column_fuser_supports_any_number_of_linears( def test_merged_name_is_unique_per_linear_names(): """The computed name must not collide across different projections, so - stacking two unrelated fusions can never clobber packed_modules_mapping. - """ + stacking two unrelated fusions can never clobber packed_modules_mapping.""" ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c")) assert ab_c.merged_name != a_bc.merged_name @@ -1222,8 +1203,7 @@ def test_heterogeneous_instances_fall_back_to_merged_column_fuser( layer whose out_features aren't a multiple of its own head_dim (as a misconfigured or heterogeneous checkpoint might have) fails QKVFuser's validation and must still fuse via the generic fallback, while a - well-formed sibling of the same class keeps using QKVFuser. - """ + well-formed sibling of the same class keeps using QKVFuser.""" with torch.device("meta"): container = nn.Module() container.compatible = PlainQKV(head_dim=8) @@ -1244,8 +1224,7 @@ def test_detects_and_rewrites_packed_qkv(kv_heads): """A single projection split into q/k/v must be re-sharded, not merged. Only the split sizes change: `QKVParallelLinear` loads the packed - checkpoint weight as-is, and shards q by heads while replicating k/v. - """ + checkpoint weight as-is, and shards q by heads while replicating k/v.""" with torch.device("meta"): meta = PackedQKVAttention(kv_heads=kv_heads) fuser = get_fuser(meta, PackedQKVFuser) @@ -1276,8 +1255,7 @@ def test_detects_and_rewrites_packed_qkv(kv_heads): def test_per_head_split_is_not_packed_qkv(): """The split must consume the whole projection, else its sizes are head - widths and re-sharding by them would be wrong. - """ + widths and re-sharding by them would be wrong.""" with torch.device("meta"): assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None @@ -1397,8 +1375,7 @@ def test_act_and_mul_derived_from_module(default_vllm_config): def _wider_model_config(head_dim: int) -> SimpleNamespace: """A model whose global head size is twice `head_dim`, as a wider layer - elsewhere in a heterogeneous checkpoint would make it. - """ + elsewhere in a heterogeneous checkpoint would make it.""" return SimpleNamespace( model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim), quant_config=None, diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py index 366123389fdf..9f51ab56621b 100644 --- a/tests/models/transformers/fusers/test_mla.py +++ b/tests/models/transformers/fusers/test_mla.py @@ -45,8 +45,7 @@ def _attention(q_lora_rank: int | None): def _match(q_lora_rank: int | None) -> MLAFuser | None: """Match a meta `DeepseekV2Attention` directly (bypassing the per-class - `get_fuser` cache, so both q_lora variants of the same class are seen). - """ + `get_fuser` cache, so both q_lora variants of the same class are seen).""" attn = _attention(q_lora_rank) return MLAFuser.match(trace(attn), attn) @@ -82,8 +81,7 @@ def test_discovers_modules_with_q_lora(): def test_q_lora_stacks_qkv_a_proj(): """The MLA layer reads `q_a_proj` and `kv_a_proj_with_mqa` fused into one - down-projection, so both checkpoint weights must remap into it. - """ + down-projection, so both checkpoint weights must remap into it.""" fuser = _match(q_lora_rank=64) assert isinstance(fuser, MLAFuser) prefix = "model.layers.0.self_attn" @@ -104,8 +102,7 @@ def test_update_forward_rewrites_real_attention(q_lora_rank): The q-LoRA path hoists the merged down-projection to a top-level statement while its calls sit inside the `q_lora_rank` branches, so the region it spans is wide; a check that over-approximates aliasing would silently refuse - here and cost the fusion rather than fail loudly. - """ + here and cost the fusion rather than fail loudly.""" fuser = _match(q_lora_rank) assert isinstance(fuser, MLAFuser) fuser.update_forward(_attention(q_lora_rank)) @@ -129,8 +126,7 @@ def forward(self, x): class RenamedMLA(nn.Module): """An MLA-shaped attention whose children have non-standard names, proving - discovery is by structure and not attribute name. - """ + discovery is by structure and not attribute name.""" def __init__(self): super().__init__() @@ -157,8 +153,7 @@ def forward(self, hidden_states): def test_discovers_modules_under_arbitrary_names(): """Discovery is purely structural: `RenamedMLA` gives its children non-standard - names, and `match` still locates each projection by dataflow. - """ + names, and `match` still locates each projection by dataflow.""" with torch.device("meta"): module = RenamedMLA() fuser = MLAFuser.match(trace(module), module) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index a465116920d3..7de4589eef91 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -277,8 +277,7 @@ def test_moe_fuser_matches_scaled_router(): def test_moe_fuser_matches_grouped_router(): """Group masking between the score and the routing top-k still matches: the - scorer is anchored on the routing (last) top-k, not the group one. - """ + scorer is anchored on the routing (last) top-k, not the group one.""" with torch.device("meta"): block = MoEBlock(GroupedRouter) fuser = MoEBlockFuser.match(block, "experts") @@ -288,8 +287,7 @@ def test_moe_fuser_matches_grouped_router(): def test_moe_fuser_matches_correction_router(): """A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the - rebuilt gate carries it in fp32 for FusedMoE's biased routers. - """ + rebuilt gate carries it in fp32 for FusedMoE's biased routers.""" with torch.device("meta"): block = MoEBlock(CorrectionRouter) fuser = MoEBlockFuser.match(block, "experts") @@ -300,8 +298,7 @@ def test_moe_fuser_matches_correction_router(): def test_moe_fuser_reads_router_dtype_from_the_gate(): """A router that computes its logits in fp32 must keep routing in fp32 when rebuilt, even though no config field names the dtype. The cast back to the - activation dtype after the top-k must not be mistaken for the routing dtype. - """ + activation dtype after the top-k must not be mistaken for the routing dtype.""" with torch.device("meta"): fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts") default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts") @@ -320,8 +317,7 @@ def test_moe_fuser_detects_shared_experts(): def test_moe_fuser_skips_shared_detection_without_extra_children(): """With only experts and gate, shared-expert detection (and its block trace) - is skipped, so a gate-derived add is not misread as a shared expert. - """ + is skipped, so a gate-derived add is not misread as a shared expert.""" with torch.device("meta"): block = MoEBlockNoShared() fuser = MoEBlockFuser.match(block, "experts") @@ -366,8 +362,7 @@ def test_moe_fuser_declines_unsupported(block_cls): def test_moe_fuser_ignores_nested_returns(): """A tuple `return` inside a nested helper must not decline a block whose own - forward returns a single tensor. - """ + forward returns a single tensor.""" with torch.device("meta"): block = MoEBlockNestedTupleReturn() assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) @@ -375,8 +370,7 @@ def test_moe_fuser_ignores_nested_returns(): def test_moe_fuser_router_requires_connected_dataflow(): """A gate with linear + softmax + top-k present but not wired as a router - (top-k selects over the input, not the scored logits) is not detected. - """ + (top-k selects over the input, not the scored logits) is not detected.""" with torch.device("meta"): block = MoEBlock(DisconnectedRouter) assert MoEBlockFuser.match(block, "experts") is None diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index dacf79f77af2..153f7863bde5 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -83,8 +83,7 @@ class ToleranceRMSNorm(NamedEpsRMSNorm): class LiteralEpsRMSNorm(RMSNorm): """eps is a literal in the source, and an unrelated attribute happens to - hold the same value: matching on the value alone would bind to it. - """ + hold the same value: matching on the value alone would bind to it.""" def __init__(self, hidden: int = 16, eps: float = 1e-5): super().__init__(hidden, eps) @@ -96,8 +95,7 @@ def forward(self, x): class AmbiguousEpsRMSNorm(nn.Module): """Two attributes hold the eps value, and the forward reads the second, so - their order cannot pick it. - """ + their order cannot pick it.""" def __init__(self, hidden: int = 16, eps: float = 1e-6): super().__init__() @@ -233,8 +231,7 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config): def test_fused_rms_norm_op_default_eps(default_vllm_config): """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path; - its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`. - """ + its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`.""" from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm with torch.device("meta"): @@ -250,8 +247,7 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): def test_eps_is_derived_per_instance(default_vllm_config): """Two instances of the same norm class with different eps must fuse to their - own eps: the type-cached fuser holds only structure, not this value. - """ + own eps: the type-cached fuser holds only structure, not this value.""" with torch.device("meta"): for eps in (1e-5, 1e-6): module = RMSNorm(16, eps=eps) @@ -263,8 +259,7 @@ def test_eps_is_derived_per_instance(default_vllm_config): @pytest.mark.parametrize("cls", [NamedEpsRMSNorm, ToleranceRMSNorm]) def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): """The eps attribute is identified by holding the traced value, so a norm - stays per-instance correct whatever it names it. - """ + stays per-instance correct whatever it names it.""" with torch.device("meta"): for eps in (1e-5, 1e-6): module = cls(16, eps=eps) @@ -276,8 +271,7 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplog): """A literal eps is recognised as coming from no attribute, even when one - holds the same value, and is taken from the traced source instead. - """ + holds the same value, and is taken from the traced source instead.""" logger = "vllm.model_executor.models.transformers.fusers.rms_norm" with caplog.at_level("DEBUG", logger=logger), torch.device("meta"): module = LiteralEpsRMSNorm() @@ -290,8 +284,7 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): """When several attributes hold the eps value, the one the forward actually - reads is identified, and they are all left as they were found. - """ + reads is identified, and they are all left as they were found.""" with torch.device("meta"): module = AmbiguousEpsRMSNorm(16, eps=1e-6) before = dict(vars(module)) @@ -308,8 +301,7 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): def test_fused_norm_is_gather_capable(default_vllm_config): """Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded projection (OLMoE-style) self-corrects at runtime with no QKV-specific - plumbing. A full-width input skips the gather and equals a plain norm. - """ + plumbing. A full-width input skips the gather and equals a plain norm.""" from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.model_executor.models.transformers.fusers import rms_norm @@ -331,8 +323,7 @@ def test_fused_norm_is_gather_capable(default_vllm_config): def test_gathered_norm_rejects_uneven_sharding(default_vllm_config): """A sharded input (narrower than the full-width weight) that does not tile - the weight evenly across ranks is rejected before any collective. - """ + the weight evenly across ranks is rejected before any collective.""" from vllm.model_executor.models.transformers.fusers import rms_norm norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6) diff --git a/tests/models/transformers/test_layer_registry.py b/tests/models/transformers/test_layer_registry.py index cc809837ae7e..9e20e2ba54ad 100644 --- a/tests/models/transformers/test_layer_registry.py +++ b/tests/models/transformers/test_layer_registry.py @@ -28,8 +28,7 @@ def fake_hw_layernorm(monkeypatch): A `SimpleNamespace` stands in for the module: `importlib.import_module` returns it from `sys.modules` and `getattr` resolves `RMSNorm`, while its attributes are set at construction (no `ModuleType` attribute-set that mypy - rejects, no constant `setattr` that ruff rejects). - """ + rejects, no constant `setattr` that ruff rejects).""" module = types.SimpleNamespace(RMSNorm=type("HwRMSNorm", (), {})) monkeypatch.setitem(sys.modules, HW_MODULE, module) return module @@ -82,8 +81,7 @@ def test_act_and_mul_falls_back_for_unknown_activation( @pytest.fixture(scope="module") def tiny_llama_path(tmp_path_factory): """A randomly-initialized microscopic Llama saved to disk (with an ungated - tokenizer) so vLLM can load it like any local checkpoint. - """ + tokenizer) so vLLM can load it like any local checkpoint.""" from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index 55736fbe2ac8..0404535e1b25 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -54,8 +54,7 @@ def test_audio_media_io_load_base64(dummy_audio_bytes): def test_audio_media_io_load_base64_rejects_malformed(dummy_audio_bytes): """Malformed base64 must surface as a ValueError so the server answers 400. Without strict decoding the bad characters are dropped, the garbage reaches - libsndfile, and the client gets a 500 instead. - """ + libsndfile, and the client gets a 500 instead.""" encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8") malformed = encoded[:8] + "!!!@@@###" + encoded[8:] @@ -134,8 +133,7 @@ def test_small_file_passes_memory_guard(): def test_memory_guard_rejects_large_allocation(): """A file whose frames*channels*4 exceeds the byte limit must be - rejected before allocating the buffer. - """ + rejected before allocating the buffer.""" # 100_000 frames * 8 channels * 4 bytes = 3.2 MB payload = _make_flac_bytes(frames=100_000, channels=8, samplerate=48000) # Set limit to 1 MiB — should reject @@ -150,8 +148,7 @@ def test_memory_guard_rejects_large_allocation(): def test_forged_samplerate_rejected_by_memory_guard(): """The PoC scenario: high sample rate fools the duration guard but - the memory guard catches the large frame*channel allocation. - """ + the memory guard catches the large frame*channel allocation.""" # Forged high sample rate: 655350 Hz, 8 channels, 1M frames # Duration guard sees: 1_000_000 / 655_350 = 1.5s → passes # Memory: 1_000_000 * 8 * 4 = 32 MB @@ -212,8 +209,7 @@ def test_load_audio_unknown_backend_rejected(dummy_audio_bytes): def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes): """`auto` must fall back to the soundfile → PyAV chain when torchcodec - is not importable. - """ + is not importable.""" ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) with patch.object(audio_module, "load_audio_torchcodec", side_effect=ImportError): audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="auto") @@ -223,8 +219,7 @@ def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes): def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes): """Torchcodec installed but system ffmpeg missing (`AudioDecoder is None`) - must surface as ImportError so `auto` falls back to soundfile → PyAV. - """ + must surface as ImportError so `auto` falls back to soundfile → PyAV.""" ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) with patch.object(audio_module, "AudioDecoder", None): with pytest.raises(ImportError, match="torchcodec audio backend"): @@ -239,8 +234,7 @@ def test_load_audio_auto_falls_back_when_libtorchcodec_unloadable( ): """Torchcodec loads its ffmpeg-backed core lazily at decoder construction; a "Could not load libtorchcodec" RuntimeError there (no system ffmpeg) - must also surface as ImportError so `auto` falls back. - """ + must also surface as ImportError so `auto` falls back.""" class _UnloadableDecoder: def __init__(self, *args, **kwargs): diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index 64145e0fea4a..b222d57d3c24 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -228,8 +228,7 @@ async def test_fetch_image_local_files_with_space_in_name(image_url: str): @pytest.mark.asyncio async def test_fetch_image_data_url_with_params(): """RFC 2397 allows parameters between the mediatype and the base64 - marker; they must not be rejected or leak into the media type. - """ + marker; they must not be rejected or leak into the media type.""" connector = MediaConnector() image = Image.new("RGB", (4, 4), color=(255, 0, 0)) diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index 2205b7c7ff3d..90310d6e6d48 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -398,8 +398,7 @@ def raise_unrelated_error(cls, file_path, nvc): class TestMergeKwargsGpuBackendPolicy: """Verify that merge_kwargs blocks request-level GPU backend selection - when the static (engine-level) config did not configure that backend. - """ + when the static (engine-level) config did not configure that backend.""" def test_pynvvideocodec_requires_gpu(self): assert VIDEO_LOADER_REGISTRY.backend_requires_gpu(PYNVVIDEOCODEC_VIDEO_BACKEND) @@ -492,8 +491,7 @@ def test_other_kwargs_preserved_when_gpu_backend_stripped(self): def test_static_pynv_with_different_runtime_gpu_backend(self): """If static sets pynv via video_backend but runtime tries to set it - via the codec-level 'backend' key (without a static match), strip it. - """ + via the codec-level 'backend' key (without a static match), strip it.""" result = VideoMediaIO.merge_kwargs( default_kwargs={"video_backend": "pynvvideocodec"}, runtime_kwargs={"backend": "pynvvideocodec"}, @@ -532,8 +530,7 @@ def test_unknown_backend_not_treated_as_gpu(self): def test_strips_request_level_device(self): """The decode device is a startup-only knob: a request must not move decoding onto the GPU when the startup config did not opt in, nor off - it when it did. - """ + it when it did.""" result = VideoMediaIO.merge_kwargs( default_kwargs={"backend": "torchcodec"}, runtime_kwargs={"device": "cuda"}, @@ -558,8 +555,7 @@ def test_strips_request_level_device(self): def test_switching_backend_drops_stale_codec_options(default_kwargs): """Codec-specific options from the static config must not leak into a different codec backend selected per-request, where they would fail the - new backend's option validation. - """ + new backend's option validation.""" result = VideoMediaIO.merge_kwargs( default_kwargs={**default_kwargs, "num_frames": 8}, runtime_kwargs={"backend": "opencv"}, @@ -581,8 +577,7 @@ def test_pynvvc_frames_normalized_to_nhwc(layout: str): """PyNvVideoCodec frame batches are normalized to NHWC regardless of the per-frame layout the decoder emits (it has varied across versions), so the HF video processors (which materialize a PIL image per frame) receive the - same NHWC shape as every other video backend. - """ + same NHWC shape as every other video backend.""" torch = pytest.importorskip("torch") n, h, w, c = 4, 5, 6, 3 diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 34ecc1c1f657..1b738eeeb89f 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -141,8 +141,7 @@ def item_for_hash(num_frames: int): def test_hash_video_tensor_frames(): """Videos holding tensor frames (e.g. NVDEC-decoded) hash like - array-framed ones, from the original bytes without a D2H copy. - """ + array-framed ones, from the original bytes without a D2H copy.""" source = b"x" * 100 def item_for_hash(frames): diff --git a/tests/multimodal/test_image.py b/tests/multimodal/test_image.py index ad1ce3162866..ea7f18611c3f 100644 --- a/tests/multimodal/test_image.py +++ b/tests/multimodal/test_image.py @@ -80,8 +80,7 @@ def test_l_mode_no_trns_to_rgb(): def test_exif_transpose_normalizes_orientation(): """Image with EXIF orientation 3 (180-degree rotation) should be - normalized so pixel data matches visual display. - """ + normalized so pixel data matches visual display.""" # Built synthetically because the existing assets are PNGs which # don't carry EXIF orientation metadata; we need a JPEG with an # injected EXIF orientation=3 tag. diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index 9c4dbed0f02c..2649f80777c1 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -59,8 +59,7 @@ def test_frame_size_hwc_chw(frame): def test_video_with_metadata_tensor_passthrough(): """Tensor frames pass through unchanged regardless of device: HF video processors accept tensors, and device-resident frames (e.g. NVDEC-decoded) - must not be copied back to host. - """ + must not be copied back to host.""" frames = torch.zeros((4, H, W, 3), dtype=torch.uint8) video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) @@ -71,8 +70,7 @@ def test_video_with_metadata_tensor_passthrough(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_video_with_metadata_keeps_device_tensor(): """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors, - so a device-side HF processor can consume them without a D2H copy. - """ + so a device-side HF processor can consume them without a D2H copy.""" frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda") video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) diff --git a/tests/multimodal/test_registry.py b/tests/multimodal/test_registry.py index 4a08d4fdf66b..6ca887401440 100644 --- a/tests/multimodal/test_registry.py +++ b/tests/multimodal/test_registry.py @@ -57,8 +57,7 @@ def test_dummy_inputs_scheduler_budget( @pytest.mark.core_model def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected): """Test supports_multimodal_inputs returns correct boolean for various - configs. - """ + configs.""" ctx = build_model_context( model_id, limit_mm_per_prompt=limit_mm_per_prompt, diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py index 994e1afc8bc4..97f46de13c18 100644 --- a/tests/multimodal/test_utils.py +++ b/tests/multimodal/test_utils.py @@ -231,8 +231,7 @@ def test_group_and_batch_mm_items_splits_shared_data_by_dtype(): def test_encode_image_url_uncommon_format_has_valid_mimetype(): """A format missing from mimetypes.types_map must still produce a - well-formed type/subtype media type, not a bare 'image'. - """ + well-formed type/subtype media type, not a bare 'image'.""" # ".tga" is not registered in mimetypes.types_map, so it exercises the # fallback path; Pillow can still encode it. url = encode_image_url(Image.new("RGB", (2, 2)), format="TGA") diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 1bc65e86a97f..432521c11f74 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -1214,8 +1214,7 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path): """With device="cuda", torchcodec decodes via NVDEC and keeps the frames - on the GPU instead of returning a host-side numpy array. - """ + on the GPU instead of returning a host-side numpy array.""" pytest.importorskip("torchcodec") with open(dummy_video_path, "rb") as f: @@ -1508,8 +1507,7 @@ def test_glm46v_duration_estimation_from_fps(): class TestGLMGASamplingCaps: """Regression tests for GHSA-58v5-2m8f-94pr: request-controlled fps and max_frames must not create intermediate allocations larger than - the actual frame count. - """ + the actual frame count.""" @staticmethod def _source( @@ -1563,8 +1561,7 @@ def test_glm5next_backend_selected_for_processor(): """Glm5NextVideoProcessor maps to the glm5next loader so only the sampled frames are decoded instead of the whole container. Both the borrowed-config spelling and the dedicated Glm5next class name (landing - with the new checkpoint) must resolve. - """ + with the new checkpoint) must resolve.""" for name in ("Glm5NextVideoProcessor", "Glm5nextVideoProcessor"): assert VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(name) == "glm5next" @@ -1586,8 +1583,7 @@ def test_glm5next_backend_indices_match_sampler( total_frames, original_fps, duration, fps, max_frames ): """The loader must select exactly the frames the processor's sampler - would, with target.fps mapping onto the raw-fps override. - """ + would, with target.fps mapping onto the raw-fps override.""" from vllm.transformers_utils.processors.glm5next import ( glm_sample_frame_indices, ) @@ -1615,8 +1611,7 @@ def test_glm5next_backend_indices_match_sampler( def test_glm5next_backend_metadata_contract(): """create_hf_metadata reports the subset so the processor skips - re-sampling (do_sample_frames=False) and keeps the original totals. - """ + re-sampling (do_sample_frames=False) and keeps the original totals.""" source = VideoSourceMetadata(total_frames_num=900, original_fps=30.0, duration=30.0) target = VideoTargetMetadata(num_frames=-1, fps=-1, max_duration=-1) indices = Glm5NextVideoBackend.compute_frames_index_to_sample(source, target) @@ -1645,8 +1640,7 @@ def test_glm5next_backend_metadata_contract(): def _write_gray_video(tmp_path, total_frames, fps, size=(32, 32)): """Synthetic clip whose frame i is flat gray level i (near-lossless under - mp4v), so a decoded frame's level maps back to its source index. - """ + mp4v), so a decoded frame's level maps back to its source index.""" cv2 = pytest.importorskip("cv2") path = tmp_path / f"gray_{total_frames}_{fps}.mp4" @@ -1687,8 +1681,7 @@ def read(self): def test_glm5next_backend_codec_parity(tmp_path, backend): """Every codec samples the same GLM indices and decodes the same frames; the OpenCV seek reader and torchcodec batched index-exact decode - must agree. - """ + must agree.""" if backend == "torchcodec": pytest.importorskip("torchcodec") @@ -1718,8 +1711,7 @@ def test_glm5next_backend_codec_parity(tmp_path, backend): def test_glm5next_backend_decodes_only_sampled_frames(tmp_path): """End to end over a synthetic clip: load_bytes returns exactly the - sampler's frame count, with the right frame content at each index. - """ + sampler's frame count, with the right frame content at each index.""" pytest.importorskip("cv2") total_frames, fps = 60, 10 @@ -1745,8 +1737,7 @@ def test_glm5next_backend_decodes_only_sampled_frames(tmp_path): def test_glm5next_read_frames_seeks_past_large_gaps(tmp_path): """Sparse targets must not walk the container: the stock reader grabs - every frame up to the last index; the GLM reader seeks instead. - """ + every frame up to the last index; the GLM reader seeks instead.""" cv2 = pytest.importorskip("cv2") total_frames, fps = 200, 10 @@ -1773,8 +1764,7 @@ def test_glm5next_read_frames_seeks_past_large_gaps(tmp_path): def test_glm5next_read_frames_dense_walk_matches_stock(tmp_path): """Dense targets keep the sequential walk (seeking would be slower) and - return the same frames as the stock reader. - """ + return the same frames as the stock reader.""" cv2 = pytest.importorskip("cv2") total_frames, fps = 120, 10 diff --git a/tests/parser/cohere/test_reasoning.py b/tests/parser/cohere/test_reasoning.py index 374c65aaaef1..8c3311e8388c 100644 --- a/tests/parser/cohere/test_reasoning.py +++ b/tests/parser/cohere/test_reasoning.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Reasoning extraction, reasoning-end gating and parser selection for the -unified Cohere Command parser. -""" +unified Cohere Command parser.""" from __future__ import annotations diff --git a/tests/parser/cohere/utils.py b/tests/parser/cohere/utils.py index 0696b9c24a68..f26df5296c6d 100644 --- a/tests/parser/cohere/utils.py +++ b/tests/parser/cohere/utils.py @@ -76,8 +76,7 @@ def token_deltas( ) -> list[tuple[str, list[int]]]: """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing - U+FFFD) into the next step as real streaming does. - """ + U+FFFD) into the next step as real streaming does.""" ids = tokenizer.encode(text) deltas: list[tuple[str, list[int]]] = [] prev = "" @@ -101,8 +100,7 @@ def make_parser( model_config=None, ) -> Parser: """Resolve the unified Cohere parser through ``ParserManager`` by registry - name (``cohere_command3`` / ``cohere_command4``). - """ + name (``cohere_command3`` / ``cohere_command4``).""" cls = ParserManager.get_parser(name, name, enable_auto_tools=True) assert cls is not None return cls( @@ -119,8 +117,7 @@ def drive_parser( deltas: list[tuple[str, list[int]]], ) -> list[DeltaMessage]: """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and - return the non-``None`` messages in order. - """ + return the non-``None`` messages in order.""" out: list[DeltaMessage] = [] for i, (text, ids) in enumerate(deltas): delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1) diff --git a/tests/parser/engine/test_deepseek_v32.py b/tests/parser/engine/test_deepseek_v32.py index aacde642fc89..8db4e0cc76b4 100644 --- a/tests/parser/engine/test_deepseek_v32.py +++ b/tests/parser/engine/test_deepseek_v32.py @@ -289,8 +289,7 @@ def test_missing_invoke_end(self, mock_tokenizer, mock_request): class TestMissingFunctionCallsWrapper: """An invoke without the ``<｜DSML｜function_calls>`` wrapper is still a - tool call (V3.2 counterpart of #48931). - """ + tool call (V3.2 counterpart of #48931).""" def test_non_streaming(self, mock_tokenizer, mock_request): text = _invoke("fn", _param("k", "true", "v")) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 0c8a375ac3b9..7cbf9791a7b0 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -334,8 +334,7 @@ def test_trailing_content_dropped_consistently( ): """Text after a tool block is dropped whether or not the closing wrapper is present, so a missing ```` does not - change what the client sees. - """ + change what the client sees.""" text = "pre\n" + self._ORPHAN if trailing_end: text += DSML_TOOL_END @@ -916,8 +915,7 @@ def test_unwrap_single_tool_still_works(self, mock_tokenizer, mock_request): class TestStreamingWrapperConsistency: """Streamed arg deltas must stay consistent with final extraction - when wrapper params like 'arguments' are unwrapped. - """ + when wrapper params like 'arguments' are unwrapped.""" def test_streaming_wrapper_unwrap_consistency(self, mock_tokenizer, mock_request): tool = _make_tool("get_weather", {"location": {"type": "string"}}) @@ -1198,8 +1196,7 @@ def test_eos_drop_token_does_not_swallow_tool_calls(self): ) def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size): """EOS must not leak into reasoning_content when the model never - emits (generation ends while still in REASONING state). - """ + emits (generation ends while still in REASONING state).""" eos_text = "<｜end▁of▁sentence｜>" eos_id = 128801 vocab = { @@ -1288,8 +1285,7 @@ def test_corrupted_wrapper_ends_reasoning( self, mock_tokenizer, mock_request, wrapper ): """Like the real wrapper, a corrupted one inside ```` closes - the reasoning block and starts the tool call. - """ + the reasoning block and starts the tool call.""" text = f"thinking{wrapper}\n{self._INNER}" parser = self._parser(mock_tokenizer, thinking=True) reasoning, content = parser.extract_reasoning(text, mock_request) @@ -1333,8 +1329,7 @@ def test_prose_mention_matches_canonical_wrapper( prose mention of the real wrapper: the opener starts a tool block and the rest of the message is dropped. That is pre-existing behavior of ``<｜DSML｜tool_calls>``; the variants must not differ from it either - way. - """ + way.""" template = "The opener looks like {} and then params follow. Done." text = template.format(wrapper) reference = template.format(DSML_TOOL_START) @@ -1358,8 +1353,7 @@ def test_prose_mention_matches_canonical_wrapper( @pytest.mark.parametrize("chunk_size", [1, 3, None], ids=lambda c: f"chunk={c}") def test_delegating_parser_corrupted_wrapper(self, wrapper, chunk_size): """The serving-layer shape: reasoning and tool adapters on separate - engines, with the corrupted opener arriving as plain text. - """ + engines, with the corrupted opener arriving as plain text.""" tokens = _dsv4_tokens( reasoning="Checking the weather.", tool_name="get_weather", diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py index 7aaad7c4dfc1..2bc4338785cb 100644 --- a/tests/parser/engine/test_delegating_replay.py +++ b/tests/parser/engine/test_delegating_replay.py @@ -174,8 +174,7 @@ def test_delegating_replay(parser_cls, sample, chunk_size): ) def test_delegating_parse_tool_choice_none(parser_cls, parser_name, sample): """Non-streaming parse() with tool_choice='none' via DelegatingParser - must not leak special tokens into content. - """ + must not leak special tokens into content.""" tokenizer = make_mock_tokenizer(sample) validated_tools = ( _TOOLS_VALIDATOR.validate_python(sample.tools) if sample.tools else None diff --git a/tests/parser/engine/test_engine.py b/tests/parser/engine/test_engine.py index f638cd599f83..36fa1ad8d09a 100644 --- a/tests/parser/engine/test_engine.py +++ b/tests/parser/engine/test_engine.py @@ -275,8 +275,7 @@ def test_text_between_tool_calls(self): def test_unmatched_close_brace_does_not_poison_depth(self): """A stray } in malformed JSON must not kill streaming for - all subsequent content. - """ + all subsequent content.""" engine = StreamingParserEngine(_hermes_config(), tokenizer=None) engine.feed("", []) @@ -292,8 +291,7 @@ def test_unmatched_close_brace_does_not_poison_depth(self): def test_json_args_no_premature_close_brace(self): """Closing braces of the top-level JSON shouldn't be streamed - until confirmed by the end tag. - """ + until confirmed by the end tag.""" engine = StreamingParserEngine(_hermes_config(), tokenizer=None) engine.feed("", []) @@ -348,8 +346,7 @@ class TestLexerBufferFlush: def test_buffered_prefix_emitted_in_current_state(self): """Text buffered by the lexer (e.g. '<') must be emitted as - REASONING_CHUNK before THINK_END transitions to CONTENT. - """ + REASONING_CHUNK before THINK_END transitions to CONTENT.""" engine = StreamingParserEngine(_think_config(), _make_think_tokenizer()) events = engine.feed("", [_START_ID]) @@ -387,13 +384,11 @@ def test_empty_buffer_no_extra_events(self): class TestTokenIdFiltering: """When token IDs are available, lex-matched terminals that also - have token_id_terminal entries should be demoted to content. - """ + have token_id_terminal entries should be demoted to content.""" def test_lex_matched_terminal_demoted_after_token_ids_seen(self): """After receiving token IDs, text that matches a token-ID - terminal should be treated as content, not trigger a transition. - """ + terminal should be treated as content, not trigger a transition.""" engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) # First feed with a non-special token ID to set _ever_had_token_ids @@ -414,8 +409,7 @@ def test_lex_matched_terminal_demoted_after_token_ids_seen(self): def test_scanner_matched_terminal_bypasses_filter(self): """PreLexedTerminals from the scanner bypass the filter and - still trigger state transitions. - """ + still trigger state transitions.""" engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events = engine.feed("", [_TOOL_START_ID]) @@ -428,8 +422,7 @@ def test_scanner_matched_terminal_bypasses_filter(self): def test_no_filtering_without_token_ids(self): """When no token IDs are ever provided (non-streaming), - text matching still triggers transitions. - """ + text matching still triggers transitions.""" engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events = engine.feed('{"name": "f"}', []) @@ -441,8 +434,7 @@ def test_no_filtering_without_token_ids(self): def test_mixed_text_then_real_tool_call(self): """Text mentioning tool syntax followed by a real special-token - tool call. - """ + tool call.""" engine = StreamingParserEngine(_hermes_config(), _make_hermes_tokenizer()) events1 = engine.feed("Mention in text. ", [1, 2, 3, 4]) @@ -502,8 +494,7 @@ def test_deferred_reasoning_tokens_stay_before_next_end_terminal(self): def _func_prefix_config() -> ParserEngineConfig: """Config mixing token-ID terminals (TOOL_START/END) with - text-only terminals (FUNC_PREFIX) and fallback transitions. - """ + text-only terminals (FUNC_PREFIX) and fallback transitions.""" return ParserEngineConfig( name="func_prefix_test", terminals={ @@ -572,13 +563,11 @@ def _make_func_prefix_tokenizer(): class TestTextOnlyFallbackFiltering: """When token IDs are available, transitions marked - skip_in_token_id_mode should be skipped. - """ + skip_in_token_id_mode should be skipped.""" def test_func_prefix_in_prose_demoted_in_strict_mode(self): """ in prose should NOT trigger a tool call - when strict mode is active. - """ + when strict mode is active.""" engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -595,8 +584,7 @@ def test_func_prefix_in_prose_demoted_in_strict_mode(self): def test_normal_flow_after_tool_start_still_works(self): """TOOL_START (special token) -> FUNC_PREFIX (text) should - still parse a tool call normally in strict mode. - """ + still parse a tool call normally in strict mode.""" engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -616,8 +604,7 @@ def test_normal_flow_after_tool_start_still_works(self): def test_fallback_fires_without_token_ids(self): """When no token IDs are provided, fallback transitions should - still fire normally. - """ + still fire normally.""" engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -631,8 +618,7 @@ def test_fallback_fires_without_token_ids(self): def test_tool_between_fallback_blocked_in_strict_mode(self): """The (TOOL_BETWEEN, FUNC_PREFIX) fallback should also be - blocked in strict mode. - """ + blocked in strict mode.""" engine = StreamingParserEngine( _func_prefix_config(), _make_func_prefix_tokenizer() ) @@ -662,8 +648,7 @@ def test_no_tokenizer_attribute(self): class TestArgsResetOnReentry: """When leaving TOOL_ARGS and later re-entering (e.g. two tool calls), the entering-TOOL_ARGS block resets args tracking. The - redundant reset on exit was removed. - """ + redundant reset on exit was removed.""" @staticmethod def _multi_tool_config() -> ParserEngineConfig: @@ -738,8 +723,7 @@ def test_brace_depth_resets_on_reentry(self): class TestToolPreambleFinish: """finish() in TOOL_PREAMBLE state emits TOOL_CALL_END when a tool - call was started (tool_index >= 0), but not when tool_index is -1. - """ + call was started (tool_index >= 0), but not when tool_index is -1.""" @staticmethod def _preamble_with_tool_call_start_config() -> ParserEngineConfig: @@ -798,8 +782,7 @@ def test_finish_no_tool_call_end_without_tool_index(self): class TestRegexTerminalInfraRemoved: """TerminalDef.priority, LexerShape.regex_terminals, and the regex - matching loop were removed. - """ + matching loop were removed.""" def test_terminal_def_no_priority(self): import regex as re @@ -824,8 +807,7 @@ def test_terminals_from_literals_still_works(self): class TestMultiCharTerminalInArgs: """Regression: multi-char terminals falling through in TOOL_ARGS - must be fed char-by-char via _feed_args_text, not _feed_args_char. - """ + must be fed char-by-char via _feed_args_text, not _feed_args_char.""" @staticmethod def _newline_config() -> ParserEngineConfig: @@ -917,8 +899,7 @@ def test_reset_preserves_skip_tool_parsing(self): def _message_header_config(reasoning_end: bool) -> ParserEngineConfig: """MESSAGE_HEADER-based grammar whose tool opener optionally carries - REASONING_END, mirroring the two shapes a transition table can take. - """ + REASONING_END, mirroring the two shapes a transition table can take.""" return ParserEngineConfig( name=f"message_header_test_{reasoning_end}", initial_state=ParserState.MESSAGE_HEADER, diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index ab83bdf5ed4c..ae21d69fc303 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -1097,8 +1097,7 @@ def test_bug_report_scenario(self, tool_call_parser, mock_request): def test_both_extractions_independent(self, parser, request_obj): """Calling extract_reasoning then extract_tool_calls on the same - parser instance should both work (each resets the engine). - """ + parser instance should both work (each resets the engine).""" model_output = FULL_MODEL_OUTPUT reasoning, _ = parser.extract_reasoning(model_output, request_obj) @@ -1112,8 +1111,7 @@ def test_both_extractions_independent(self, parser, request_obj): class TestAdapterExtractReasoning: """The reasoning adapter's extract_reasoning uses skip_tool_parsing - so tool call text is preserved as content for the tool adapter. - """ + so tool call text is preserved as content for the tool adapter.""" @pytest.fixture def adapter(self, mock_tokenizer): @@ -1154,8 +1152,7 @@ def test_no_reasoning_returns_none(self, adapter, request_obj): class TestGemma4SchemaAwareTypeCoercion: """Verify that streaming and non-streaming produce identical type-fixed arguments when tool schemas declare string parameters - but the model outputs bare numbers/booleans. - """ + but the model outputs bare numbers/booleans.""" @pytest.fixture def tools(self): @@ -1185,8 +1182,7 @@ def parser_with_tools(self, tool_call_tokenizer, tools): def test_streaming_string_param_not_coerced(self, parser_with_tools, mock_request): """A numeric value for a string-typed param must remain a string - in the streamed output, matching the non-streaming result. - """ + in the streamed output, matching the non-streaming result.""" chunks = [ "<|tool_call>", "call:update_record{", @@ -1217,8 +1213,7 @@ def test_streaming_mixed_types(self, parser_with_tools, mock_request): def test_streaming_matches_non_streaming(self, parser_with_tools, mock_request): """Concatenated streaming deltas must produce the same arguments - as non-streaming extraction. - """ + as non-streaming extraction.""" text = "<|tool_call>call:update_record{zipcode:12345}" non_streaming = parser_with_tools.extract_tool_calls(text, mock_request) @@ -1239,8 +1234,7 @@ def test_streaming_matches_non_streaming(self, parser_with_tools, mock_request): class TestGemma4SchemaCoercionBoolNumberNull: """Verify that _fix_arg_types coerces string values to non-string - schema types for the Gemma4 parser. - """ + schema types for the Gemma4 parser.""" @pytest.fixture def tools(self): @@ -1301,8 +1295,7 @@ def test_null_stays_string_without_null_schema( def test_streaming_type_stability(self, parser_with_tools, mock_request): """Values streamed incrementally must not cause prefix - incompatibility when types are coerced. - """ + incompatibility when types are coerced.""" text = ( "<|tool_call>call:configure{" "enabled:true," @@ -1402,8 +1395,7 @@ def test_nested_object_coerced(self, parser_with_tools, mock_request): class TestBareThoughtWithoutChannelOpener: """When the model omits <|channel> and starts with bare ``thought``, the parser should auto-inject the channel opener so reasoning is - captured correctly. - """ + captured correctly.""" @pytest.fixture def bare_thought_tokenizer(self): @@ -1474,8 +1466,7 @@ def test_normal_content_not_classified_as_reasoning(self, request_obj): def test_bare_thought_token_at_end_of_stream(self, request_obj): """When the stream ends with just "thought" (no \\n), the parser - should treat it as the thought prefix token, not real reasoning. - """ + should treat it as the thought prefix token, not real reasoning.""" seq: list[tuple[int, str]] = [ (CHANNEL_START_ID, "<|channel>"), (3000, "thought"), diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py index fe6d0ec5c5ad..58fe7aa28cbd 100644 --- a/tests/parser/engine/test_glm47_moe.py +++ b/tests/parser/engine/test_glm47_moe.py @@ -61,8 +61,7 @@ def test_tool_call_without_think_end(self, parser): def test_previous_turn_tool_call_ignored(self, parser): """A finished tool call from an earlier turn says nothing about the - turn currently being generated. - """ + turn currently being generated.""" prompt = [THINK_S, TEXT, TOOL_S, TEXT, TOOL_E, OBSERVATION, TEXT, ASSISTANT] assert not parser.is_reasoning_end(prompt) @@ -89,8 +88,7 @@ def test_every_tool_call_after_think_end_kept(self, parser): def test_falls_back_to_tool_call(self, parser): """Without ````, content starts at the opener itself so the - tool parser still receives a well-formed call. - """ + tool parser still receives a well-formed call.""" ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E] assert parser.extract_content_ids(ids) == [TOOL_S, 30, TOOL_E] diff --git a/tests/parser/engine/test_inkling.py b/tests/parser/engine/test_inkling.py index c38e74544c31..9026d548dfbb 100644 --- a/tests/parser/engine/test_inkling.py +++ b/tests/parser/engine/test_inkling.py @@ -70,8 +70,7 @@ def _tool_block(name: str, args: str) -> str: def _tokenize(text: str) -> list[tuple[int, str]]: """Tokenize like the real stream: markers are atomic special tokens, plain text becomes one token per character (matching the mock - tokenizer's ``chr``-based decode). - """ + tokenizer's ``chr``-based decode).""" tokens: list[tuple[int, str]] = [] i = 0 while i < len(text): @@ -88,8 +87,7 @@ def _tokenize(text: str) -> list[tuple[int, str]]: def _stream(parser, request, text: str, chunk_size: int): """Stream production-shaped deltas: ``chunk_size`` tokens per delta, - with delta_token_ids covering every token (specials and text). - """ + with delta_token_ids covering every token (specials and text).""" tokens = _tokenize(text) results = [] previous_text = "" @@ -121,8 +119,7 @@ def _stream(parser, request, text: str, chunk_size: int): def _stream_text_only(parser, request, text: str, chunk_size: int): """Stream text-only deltas (no token ids), chunked at arbitrary character boundaries — exercises the text-lexing fallback path, - including markers split across chunks. - """ + including markers split across chunks.""" results = [] previous_text = "" for start in range(0, len(text), chunk_size): @@ -173,8 +170,7 @@ def _stream_delegating(parser, request, text, chunk_size, prompt_token_ids): """Stream ``text`` through ``DelegatingParser.parse_delta``, ``chunk_size`` tokens per delta; return ``(content, reasoning, ordered tool names, ordered tool arguments)``. Arguments arrive in fragments, so they are - concatenated per tool index. - """ + concatenated per tool index.""" tokens = _tokenize(text) content, reasoning = "", "" tools: dict[int, str] = {} @@ -526,8 +522,7 @@ class TestToolCallFiltering: """Inkling equivalents of the generic tool-call-filtering replay tests (Inkling is excluded from those in test_replay.py: its structural role/kind tokens and shared block-end token don't fit the generic - reasoning/tool split model). - """ + reasoning/tool split model).""" def test_skip_tool_parsing_round_trip(self, mock_tokenizer, mock_request): # First pass (reasoning adapter, skip_tool_parsing): reasoning is @@ -825,8 +820,7 @@ def test_function_name_header_before_tool_start_streaming( ): """The optional function name between ``<|message_model|>`` and the content-kind marker is metadata: the buffered header must be - discarded on the way out, not flushed into content. - """ + discarded on the way out, not flushed into content.""" tools = [_function_tool()] mock_request.tools = tools content, _, names, args = _stream_delegating( @@ -843,8 +837,7 @@ def test_function_name_header_before_tool_start_streaming( def test_content_state_tool_start_streaming(self, mock_tokenizer, mock_request): """Same opener reached from CONTENT rather than MESSAGE_HEADER: a text block closed with no ``<|message_model|>`` before the tool - block. Preceding text must reach content exactly once, unmarked. - """ + block. Preceding text must reach content exactly once, unmarked.""" tools = [_function_tool()] mock_request.tools = tools content, _, names, args = _stream_delegating( @@ -866,8 +859,7 @@ def test_content_tool_start_emits_reasoning_end_in_reasoning_pass(): boundary (#49876), but the header-flush path (MESSAGE_HEADER --END_MESSAGE--> CONTENT) reaches CONTENT without one, so a tool block opening from there relies on this transition - alone to hand off to the tool pass. - """ + alone to hand off to the tool pass.""" engine = StreamingParserEngine(inkling_config(), tokenizer=None) engine.skip_tool_parsing = True engine.reset(initial_state=ParserState.CONTENT) @@ -884,8 +876,7 @@ def test_skip_reasoning_parsing_inert_for_shared_markers(): markers. Inkling has none — ``<|end_message|>`` is labelled THINK_END yet also closes text, header, and tool blocks — so the flag must be inert: a tool block still closes through its transition and following - text returns to CONTENT instead of leaking into the argument stream. - """ + text returns to CONTENT instead of leaking into the argument stream.""" engine = StreamingParserEngine(inkling_config(), tokenizer=None) engine.skip_reasoning_parsing = True engine.reset(initial_state=ParserState.CONTENT) @@ -905,8 +896,7 @@ class TestToolParserWithoutReasoningParser: labels ``<|end_message|>`` as THINK_END; a bypass keyed on the label would neutralize the closer of every block kind, leaking markers into content and leaving tool calls unterminated. Structure must keep - parsing exactly as with the reasoning parser attached. - """ + parsing exactly as with the reasoning parser attached.""" GEN_PROMPT = [_TML_VOCAB[MSG_MODEL]] diff --git a/tests/parser/engine/test_nemotron_v3.py b/tests/parser/engine/test_nemotron_v3.py index 201ddc4e781c..a7f166f9f3a9 100644 --- a/tests/parser/engine/test_nemotron_v3.py +++ b/tests/parser/engine/test_nemotron_v3.py @@ -56,8 +56,7 @@ def parser(): class TestNemotronSwap: def test_enable_thinking_false_swaps(self, parser): """When enable_thinking=False, model output without think tags - should have reasoning swapped to content. - """ + should have reasoning swapped to content.""" text = "The answer is 42." request = _make_request(enable_thinking=False) reasoning, content = parser.extract_reasoning(text, request) @@ -74,8 +73,7 @@ def test_force_nonempty_content_swaps(self, parser): def test_no_swap_when_content_exists(self, parser): """With enable_thinking=False but real giving content, - no swap occurs. - """ + no swap occurs.""" text = "Some reasoning.Actual content here." request = _make_request(enable_thinking=False) reasoning, content = parser.extract_reasoning(text, request) @@ -182,13 +180,11 @@ def test_streaming_tool_calls(self, parser): class TestParseDeltaTokenIdFiltering: """parse_delta must not trigger tool call parsing when - appears as regular text rather than as a special token ID. - """ + appears as regular text rather than as a special token ID.""" def test_tool_call_text_in_reasoning_is_not_parsed(self, parser): """Literal in model reasoning should be content, - not a tool call. - """ + not a tool call.""" request = _make_request() text = ( @@ -214,8 +210,7 @@ def test_tool_call_text_in_reasoning_is_not_parsed(self, parser): def test_special_token_id_still_triggers_tool_call(self, parser): """When the scanner matches a special token ID, the tool call - must still be parsed correctly. - """ + must still be parsed correctly.""" request = _make_request() parser.parse_delta( @@ -255,8 +250,7 @@ def test_special_token_id_still_triggers_tool_call(self, parser): def test_text_discussion_then_real_tool_call(self, parser): """Model discusses tool syntax in reasoning, then makes a real - tool call via special tokens. - """ + tool call via special tokens.""" request = _make_request() r1 = parser.parse_delta( diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index a64ab2ecd541..8fd7c4d2623f 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -323,8 +323,7 @@ def test_kimi_k2_tool_call_id_includes_func_name(self): def test_multiple_arg_chunks_same_batch_coalesced(self): """Multiple events for the same tool in one batch must produce - at most one DeltaToolCall per index. - """ + at most one DeltaToolCall per index.""" engine = _make_engine() events = [ SemanticEvent(EventType.TOOL_CALL_START, tool_index=0), @@ -552,8 +551,7 @@ class TestPostToolContentDeferral: """Regression: content after TOOL_CALL_END in the same batch must not produce a mixed DeltaMessage(content=..., tool_calls=...) — that causes split_delta to reorder content before tool_calls, breaking the Responses - API state machine. - """ + API state machine.""" def test_text_after_tool_end_deferred(self): engine = _make_engine() @@ -607,8 +605,7 @@ def test_deferred_content_not_flushed_during_arg_continuation(self): """Deferred content from batch N must not mix with arg-continuation tool events in batch N+1 — that creates a DeltaMessage with both content and nameless tool_calls, which crashes the Responses API - state machine (name=None → Pydantic ValidationError). - """ + state machine (name=None → Pydantic ValidationError).""" engine = _make_engine() engine._content_has_nonws = True @@ -837,8 +834,7 @@ def test_content_passthrough(self): class TestEngineBasedPath: """Tests for the _engine_based accumulation behavior in - DelegatingParser.parse_delta. - """ + DelegatingParser.parse_delta.""" def test_engine_based_true_when_both_parsers_engine(self): r = SimpleNamespace(engine_based_streaming=True) @@ -884,8 +880,7 @@ def test_parse_delta_tool_call(self, mock_request): class TestParseTokenIdPassthrough: """parse() must forward model_output_token_ids to _single_pass_parse - so that token-ID-based strict terminal matching is active. - """ + so that token-ID-based strict terminal matching is active.""" def test_literal_tool_tag_in_content_preserved_with_token_ids(self, mock_request): engine = _make_engine(_hermes_config()) @@ -1109,8 +1104,7 @@ class TestAdapterFinishOnStreamEnd: def test_lexer_buffer_flushed_on_finished(self): """Text buffered as a potential terminal prefix must be emitted - as content when the stream ends. - """ + as content when the stream ends.""" tokenizer = make_mock_tokenizer(_VOCAB) parser = _CombinedDelegating(tokenizer) request = _make_delegating_request() @@ -1128,8 +1122,7 @@ def test_lexer_buffer_flushed_on_finished(self): def test_args_buffer_flushed_on_finished(self): """Pending arg buffer text must be emitted when stream ends - mid-tool-call (closing brace held back in buffer). - """ + mid-tool-call (closing brace held back in buffer).""" tokenizer = make_mock_tokenizer(_VOCAB) parser = _CombinedDelegating(tokenizer) request = _make_delegating_request() @@ -1365,8 +1358,7 @@ def test_content_after_transitionless_closer_observable(self): class TestToolAdapterForwardsKwargs: """ParserEngineToolAdapter.__init__ must forward **kwargs to the - parser engine class so chat_template_kwargs reach model parsers. - """ + parser engine class so chat_template_kwargs reach model parsers.""" @pytest.mark.parametrize( "enable_thinking,expected_state", @@ -1396,8 +1388,7 @@ def test_kwargs_forwarded_to_parser_engine(self, enable_thinking, expected_state class TestExtractContentIdsNoEmptyReturn: """extract_content_ids must return input_ids (not []) when there is - no THINK_END token ID and _reasoning_ended is True. - """ + no THINK_END token ID and _reasoning_ended is True.""" _NO_THINK_CONFIG = ParserEngineConfig(name="no_think_end", token_id_terminals={}) @@ -1795,8 +1786,7 @@ class TestDropSpecialTokens: def test_drops_special_token_by_id_from_content(self): """A special token (not a configured terminal) is dropped when - it arrives as its actual token ID. - """ + it arrives as its actual token ID.""" config = ParserEngineConfig( name="drop_content_test", terminals={}, @@ -1831,8 +1821,7 @@ def test_drops_special_token_by_id_from_reasoning(self): def test_drops_via_text_fallback_when_no_token_ids(self): """When no token IDs are provided, text-based lexer catches - drop tokens as a fallback. - """ + drop tokens as a fallback.""" engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1845,8 +1834,7 @@ def test_drops_via_text_fallback_when_no_token_ids(self): def test_regular_tokens_spelling_special_survive(self): """Regular tokens that spell out a drop-token string survive - when token IDs prove they are not the special token. - """ + when token IDs prove they are not the special token.""" vocab = {**_DROP_VOCAB, "h": 72, "<": 73, "bos": 74, ">": 75, "w": 76} engine = _make_engine( vocab=vocab, @@ -1862,8 +1850,7 @@ def test_regular_tokens_spelling_special_survive(self): def test_configured_terminal_not_treated_as_drop(self): """Tokens already in config.terminals (like ) are handled - by the state machine, not the drop mechanism. - """ + by the state machine, not the drop mechanism.""" engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1918,8 +1905,7 @@ def test_no_special_tokens_means_no_drops(self): def test_drops_applied_with_skip_tool_parsing(self): """Drop tokens are always dropped, even with skip_tool_parsing. DROP_TERMINALs have no transitions by construction, so no parser - pass can use them. - """ + pass can use them.""" for initial_state in (ParserState.REASONING, ParserState.CONTENT): engine = _make_engine( vocab=_DROP_VOCAB, @@ -1935,8 +1921,7 @@ def test_drops_applied_with_skip_tool_parsing(self): def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self): """With skip_tool_parsing in REASONING state, drop tokens are - removed but configured terminals still fire their transitions. - """ + removed but configured terminals still fire their transitions.""" engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1972,8 +1957,7 @@ def test_drops_in_tool_args_state(self): def test_mixed_configured_and_drop_terminals(self): """Configured terminals trigger transitions while drop terminals - are silently removed in the same stream. - """ + are silently removed in the same stream.""" engine = _make_engine( vocab=_DROP_VOCAB, special_tokens=list(_DROP_VOCAB.keys()), @@ -1997,8 +1981,7 @@ class TestTruncatedToolOpenerStreamParity: """Regression tests for #47137: when generation terminates inside a ```` opener that has not been promoted to a tool call (via ``max_tokens`` or a ``stop`` string), the non-streaming path - must drop the incomplete markup, matching the streaming path. - """ + must drop the incomplete markup, matching the streaming path.""" _QWEN3_VOCAB = { "": 100, @@ -2044,8 +2027,7 @@ def _stream_content(self, request, chunks: list[str]) -> str: ) def test_truncated_opener_dropped_in_both_paths(self, mock_request, chunks): """A cutoff inside the opener yields no content and no tool calls - in both the non-streaming and streaming paths. - """ + in both the non-streaming and streaming paths.""" text = "".join(chunks) _, content, tool_calls = self._make_parser().parse( @@ -2058,8 +2040,7 @@ def test_truncated_opener_dropped_in_both_paths(self, mock_request, chunks): def test_content_before_truncated_opener_preserved(self, mock_request): """Only the incomplete markup is dropped; content generated before - the opener is returned identically by both paths. - """ + the opener is returned identically by both paths.""" chunks = ["Checking the weather. ", "", "\n", "", "garbage", "", " B"] text = "".join(chunks) @@ -2088,8 +2068,7 @@ def test_content_around_unpromoted_tool_block_stays_ordered(self, mock_request): def test_complete_tool_call_still_promoted(self, mock_request): """Sanity check: a complete tool call still parses in the - non-streaming path after the truncation fix. - """ + non-streaming path after the truncation fix.""" text = ( "\n" "\n" @@ -2109,8 +2088,7 @@ class TestThinkMarkupWithoutReasoningParser: """With no reasoning parser configured, reasoning markup is plain content: engine-backed tool parsers must return it verbatim in both the non-streaming and streaming paths, not consume the markers or - reclassify the block as reasoning and drop it. - """ + reclassify the block as reasoning and drop it.""" _VOCAB = { "": 90, @@ -2176,8 +2154,7 @@ def test_think_block_passes_through_both_paths( def test_think_block_preserved_alongside_tool_call(self, mock_request): """The tools-called branch also returns the tool parser's content; - a think block preceding a promoted tool call must survive in it. - """ + a think block preceding a promoted tool call must survive in it.""" think = "plan" text = think + ( "\n" diff --git a/tests/parser/engine/test_qwen3.py b/tests/parser/engine/test_qwen3.py index b8f96294cb79..07a025876ae7 100644 --- a/tests/parser/engine/test_qwen3.py +++ b/tests/parser/engine/test_qwen3.py @@ -318,8 +318,7 @@ def test_streaming_multi_param(self, parser, mock_request): def test_streaming_args_arrive_incrementally(self, parser, mock_request): """Arguments must stream as intermediate deltas, not batch at - tool-end. - """ + tool-end.""" chunks = [ "\n", "\n", @@ -524,8 +523,7 @@ def test_streaming_value_split_across_chunks(self, parser, mock_request): def test_streaming_split_tool_call_tag(self, parser, mock_request): """ arrives as a single special token; the rest of - the content is split into fine-grained chunks. - """ + the content is split into fine-grained chunks.""" chunks = [ "\n", "\n", @@ -702,8 +700,7 @@ def test_partial_value_with_angle_bracket_and_complete_param(self): class TestSchemaAwareTypeCoercion: """Verify that _fix_arg_types corrects miscoerced values using the - tool schema. - """ + tool schema.""" @pytest.fixture def tools(self): @@ -845,8 +842,7 @@ def parser_with_anyof(self, mock_tokenizer, tools_with_anyof): def test_anyof_string_param_not_coerced(self, parser_with_anyof, mock_request): """A param with anyOf including 'string' must not be coerced - to integer. - """ + to integer.""" text = ( "\n" "\n" @@ -862,8 +858,7 @@ def test_anyof_string_param_not_coerced(self, parser_with_anyof, mock_request): class TestSchemaCoercionBoolNumberNull: """Verify that _fix_arg_types coerces string values to non-string - schema types using coerce_to_schema_type. - """ + schema types using coerce_to_schema_type.""" @pytest.fixture def tools(self): diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index d8e8096fd3b6..3f285e4d0315 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -276,8 +276,7 @@ def test_continued_turn_with_open_thinking_not_end(self, parser): def test_empty_boundary_config_keeps_global_walk(self): """Configs without turn_boundary_tokens keep the pre-existing - behavior: any in the sequence ends reasoning. - """ + behavior: any in the sequence ends reasoning.""" cfg = qwen3_config(turn_boundary_tokens=frozenset()) parser = Qwen3Parser( make_mock_tokenizer(_QWEN3_VOCAB), parser_engine_config=cfg @@ -289,8 +288,7 @@ def test_empty_boundary_config_keeps_global_walk(self): def test_boundary_tokens_missing_from_vocab_keep_global_walk(self): """Boundary tokens absent from the vocabulary resolve to nothing, so subclasses reusing the default on a non-ChatML vocab keep the - pre-existing behavior. - """ + pre-existing behavior.""" vocab = { token: token_id for token, token_id in _QWEN3_VOCAB.items() @@ -301,8 +299,7 @@ def test_boundary_tokens_missing_from_vocab_keep_global_walk(self): def test_boundary_with_thinking_disabled_is_end(self, mock_tokenizer): """With thinking disabled (initial state CONTENT) a walk that stops at a turn - boundary must report reasoning as ended: the model never emits a marker. - """ + boundary must report reasoning as ended: the model never emits a marker.""" parser = Qwen3Parser( mock_tokenizer, chat_template_kwargs={"enable_thinking": False} ) @@ -310,8 +307,7 @@ def test_boundary_with_thinking_disabled_is_end(self, mock_tokenizer): def test_boundary_with_thinking_enabled_not_end(self, parser): """With thinking disabled (initial state REASONING) a walk that stops at a turn - boundary must keep reasoning open - """ + boundary must keep reasoning open""" assert not parser.is_reasoning_end([_IM_START_ID, _TEXT_ID]) @@ -672,8 +668,7 @@ def test_streaming_trailing_newlines_before_tool_call(self, parser_with_strip): class TestWhitespaceStrippingDisabled: """When strip_trailing_reasoning_whitespace is False, - trailing whitespace in reasoning must be preserved. - """ + trailing whitespace in reasoning must be preserved.""" @pytest.fixture def parser_no_strip(self): @@ -731,8 +726,7 @@ def test_default_initial_state_is_reasoning(self, mock_tokenizer): def test_thinking_disabled_streaming_content_only(self, mock_tokenizer): """Plain text with thinking disabled must stream as content, not reasoning. Before the fix, the REASONING initial state caused all - output to be emitted as reasoning chunks. - """ + output to be emitted as reasoning chunks.""" p = Qwen3Parser( mock_tokenizer, chat_template_kwargs={"enable_thinking": False}, diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py index 75fdbfca0fcc..fee74f8d2cca 100644 --- a/tests/parser/engine/test_replay.py +++ b/tests/parser/engine/test_replay.py @@ -440,8 +440,7 @@ def test_replay(self, parser_cls, sample, think_end, tool_start, mode, chunk_siz ) class TestToolCallFilteringNonStreaming: """Non-streaming parse() with tool_choice='none' must suppress tool - calls and not leak special tokens into content. - """ + calls and not leak special tokens into content.""" def test_parse(self, parser_cls, sample, think_end, tool_start): tokenizer = make_mock_tokenizer(sample) @@ -484,8 +483,7 @@ def test_parse(self, parser_cls, sample, think_end, tool_start): ) class TestToolChoiceNoneStreamingParity: """Streaming and non-streaming must return the same content - when tool_choice='none' suppresses tool calls. - """ + when tool_choice='none' suppresses tool calls.""" def test_content_matches(self, parser_cls, sample): tokenizer = make_mock_tokenizer(sample) @@ -543,8 +541,7 @@ def _inject_drop_tokens(sample): class TestDropTokenReplay: """Verify unconfigured special tokens are silently dropped across - all parsers and chunk sizes. - """ + all parsers and chunk sizes.""" @pytest.mark.parametrize( "parser_info", @@ -601,8 +598,7 @@ def test_drop_tokens_removed_from_output(self, parser_info): class TestAdapterReferences: """Verify make_adapters sets reasoning/tool parser class refs on parser engine - parser classes so the serving layer finds them and calls adjust_request. - """ + parser classes so the serving layer finds them and calls adjust_request.""" @pytest.mark.parametrize( "parser_name", diff --git a/tests/parser/engine/test_seed_oss.py b/tests/parser/engine/test_seed_oss.py index 159832135923..e01f2fd9f6d4 100644 --- a/tests/parser/engine/test_seed_oss.py +++ b/tests/parser/engine/test_seed_oss.py @@ -129,8 +129,7 @@ def test_single_tool_call(tool_parser, mock_request): def test_malformed_function_end_does_not_drop_siblings(tool_parser, mock_request): """Regression for #46314: a malformed ```` with no closing ``>`` - on the header must not discard the other, well-formed calls. - """ + on the header must not discard the other, well-formed calls.""" text = ( f"{TOOL_CALL_START}\n\n{TOOL_CALL_END}" f"{TOOL_CALL_START}\n\n" @@ -172,8 +171,7 @@ def test_reasoning_then_tool_call(parser): def test_streaming_think_end_and_tool_call_same_delta(parser): """```` and ```` arriving in one delta must - not leak the terminal tokens into the reasoning text. - """ + not leak the terminal tokens into the reasoning text.""" reasoning, content = simulate_reasoning_streaming( parser, [ diff --git a/tests/parser/engine/test_ufffd_reasoning_transition.py b/tests/parser/engine/test_ufffd_reasoning_transition.py index 6497a6c29f63..ffd2ead95d7b 100644 --- a/tests/parser/engine/test_ufffd_reasoning_transition.py +++ b/tests/parser/engine/test_ufffd_reasoning_transition.py @@ -146,8 +146,7 @@ def test_no_ufffd(self, chunk_size, vocab, delegating_cls): def test_byte_fallback_tokenizer_produces_ufffd(self): """Validate the fixture: decode() returns U+FFFD for isolated - byte-fallback token IDs, proving the old code path would leak. - """ + byte-fallback token IDs, proving the old code path would leak.""" vocab = dict(_MODEL_CONFIGS[0].values[0]) tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS) assert tokenizer.decode([200]) == "�" diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 347df05b4d46..36e347fc572f 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -992,8 +992,7 @@ def _inkling_block( ) -> None: """Append one Inkling content block; blocks after the first start with the ``<|message_model|>`` role token (the first block continues the - generation prompt directly). - """ + generation prompt directly).""" if segs: segs.append(("<|message_model|>", True)) segs.append((kind_token, True)) diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py index 8645a38dd3cb..25b41f44c2d2 100644 --- a/tests/parser/mistral/test_tool_calls.py +++ b/tests/parser/mistral/test_tool_calls.py @@ -492,8 +492,7 @@ def test_extract_tool_calls_pre_v11_regex_fallback( ): """The regex fallback path finds valid JSON via regex when the primary raw_decode fails on leading junk. It should re-serialize arguments - and return a valid tool call. - """ + and return a valid tool call.""" model_output = ( '[TOOL_CALLS] junk [{"name": "add", "arguments":{"a": 1, "b": 2}}] trail' ) @@ -2004,8 +2003,7 @@ def test_streaming_pre_v11_parallel_calls_batched_deltas( mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size, driver ): """A batched delta spanning the boundary between two parallel calls must - keep them on distinct indices (the bug collapsed both onto index 0). - """ + keep them on distinct indices (the bug collapsed both onto index 0).""" model_output = ( '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, ' '{"name": "get_current_weather", "arguments": ' @@ -2123,8 +2121,7 @@ def test_reasoning_active_no_think_block_no_leak( chunk_size, ): """Tool call without a preceding think block must not leak as content - when the reasoning parser is active. - """ + when the reasoning parser is active.""" accumulated_content = "" function_names: list[str] = [] function_args_strs: list[str] = [] @@ -2225,8 +2222,7 @@ def test_adjust_request_pre_v11_required_clears_response_format( response_format so the tool schema is the sole structured-output constraint, matching the base ToolParser (otherwise the request either hits the "multiple constraints" engine error or, with no constraint, - rambles to finish_reason='length'). - """ + rambles to finish_reason='length').""" request = _make_request(tool_choice="required", response_format=response_format) result = mistral_pre_v11_tool_parser.adjust_request(request) @@ -2243,8 +2239,7 @@ def test_adjust_request_non_mistral_tokenizer_required_injects_schema( ) -> None: """The guided-schema injection must also fire for non-Mistral (e.g. HF-mode) tokenizers driving the Mistral tool parser, mirroring the base ToolParser so - required + response_format does not fall back to an unconstrained ramble. - """ + required + response_format does not fall back to an unconstrained ramble.""" request = _make_request( tool_choice="required", response_format={"type": "json_object"} ) diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index 6c4158b81490..cad16746be44 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -229,8 +229,7 @@ def _boundary_chunks(tokenizer, parser, end_token_id=None): def test_parse_delta_reasoning_not_dropped_on_boundary(tokenizer, request_obj): """Regression: reasoning must not be lost when a multi-token delta - spans the reasoning/tool-call boundary. - """ + spans the reasoning/tool-call boundary.""" parser = make_parser(tokenizer, reasoning=True, tool=True) chunks = _boundary_chunks(tokenizer, parser) results = stream_chunks(parser, tokenizer, chunks, request_obj) @@ -248,8 +247,7 @@ def test_parse_delta_reasoning_not_dropped_on_boundary(tokenizer, request_obj): def test_parse_delta_reasoning_boundary_no_tool_parser(tokenizer, request_obj): """When no tool parser is active, boundary-spanning chunks must still - preserve reasoning and pass post- text as content. - """ + preserve reasoning and pass post- text as content.""" parser = make_parser(tokenizer, reasoning=True, tool=False) chunks = _boundary_chunks(tokenizer, parser) results = stream_chunks(parser, tokenizer, chunks, request_obj) @@ -263,8 +261,7 @@ def test_parse_delta_reasoning_boundary_no_tool_parser(tokenizer, request_obj): def test_parse_delta_reasoning_only_no_think_leak(tokenizer, request_obj): """Regression: must not leak into content when streaming - token-by-token with reasoning=True, tool=False. - """ + token-by-token with reasoning=True, tool=False.""" parser = make_parser(tokenizer, reasoning=True, tool=False) results = stream_text( parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] @@ -306,8 +303,7 @@ def test_parse_delta_reasoning_only_thinking_disabled(tokenizer, request_obj): def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, request_obj): """When finished=True but the final parse_delta produces no - tool-call delta, unstreamed args are not flushed. - """ + tool-call delta, unstreamed args are not flushed.""" parser = make_parser(tokenizer, reasoning=False, tool=True) results = stream_text( @@ -330,8 +326,7 @@ def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, reques def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, request_obj): """When all args have been streamed, finished=True must not - produce extra or duplicate arguments. - """ + produce extra or duplicate arguments.""" parser = make_parser(tokenizer, reasoning=False, tool=True) results = stream_text( parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] @@ -351,8 +346,7 @@ def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, reque def test_parse_delta_finished_appends_remaining_args(tokenizer, request_obj): """When finished=True and the tool parser has unstreamed args, - parse_delta appends the remaining arguments to the tool-call delta. - """ + parse_delta appends the remaining arguments to the tool-call delta.""" parser = make_parser(tokenizer, reasoning=False, tool=True) token_ids = tokenizer.encode(MODEL_OUTPUT, add_special_tokens=False) @@ -499,8 +493,7 @@ def test_engine_reasoning_hermes_tool_token_by_token(tokenizer, request_obj): """Qwen3 engine reasoning + Hermes tool parser, token-by-token. Sanity check that the mixed engine/non-engine configuration works - when tokens arrive one at a time (no deferred content). - """ + when tokens arrive one at a time (no deferred content).""" parser = Qwen3ReasoningHermesToolParser(tokenizer) assert parser._reasoning_parser.engine_based_streaming is True @@ -527,8 +520,7 @@ def test_engine_reasoning_hermes_tool_boundary(tokenizer, request_obj): When and are in the same chunk with aligned text and token IDs, the engine processes both terminals and returns - the text as content. - """ + the text as content.""" parser = Qwen3ReasoningHermesToolParser(tokenizer) end_token_id = parser._reasoning_parser._parser_engine._reasoning_end_token_id chunks = _boundary_chunks(tokenizer, parser, end_token_id=end_token_id) @@ -559,8 +551,7 @@ def test_engine_reasoning_hermes_tool_text_holdback(tokenizer, request_obj): Without the fix, finish_streaming() is never called at the reasoning->tool transition when _engine_based is False, so the '<' is lost and the Hermes parser sees 'tool_call>...' instead of - '...'. - """ + '...'.""" parser = Qwen3ReasoningHermesToolParser(tokenizer) vocab = tokenizer.get_vocab() think_end_id = vocab[""] @@ -651,8 +642,7 @@ def test_engine_reasoning_no_tool_batched_content_passthrough(tokenizer, request batched delta carries ```` plus the following content (as happens with stream_interval > 1). The post-```` text must be emitted as content -- not dropped, not reclassified as - reasoning -- and the ```` marker must not leak either way. - """ + reasoning -- and the ```` marker must not leak either way.""" parser = Qwen3ReasoningNoToolParser(tokenizer) assert parser._reasoning_parser.engine_based_streaming is True assert parser._tool_parser is None @@ -680,8 +670,7 @@ def _decode_stream_deltas(tokenizer, groups): This mirrors how vLLM's detokenizer feeds ``parse_delta`` in production: byte-level UTF-8 hold-back means a character whose bytes span multiple tokens is only surfaced once complete (a naive - per-token ``decode`` would instead emit U+FFFD replacement chars). - """ + per-token ``decode`` would instead emit U+FFFD replacement chars).""" from tokenizers.decoders import DecodeStream stream = DecodeStream(skip_special_tokens=False) @@ -713,8 +702,7 @@ def test_engine_reasoning_hermes_tool_multibyte_holdback(tokenizer, request_obj) ``DecodeStream``. Its UTF-8 hold-back yields the correct character round-trip (a naive per-token decode would corrupt it), verifying the boundary stays byte-safe for multi-token - characters. - """ + characters.""" parser = Qwen3ReasoningHermesToolParser(tokenizer) vocab = tokenizer.get_vocab() think_end_id = vocab[""] diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py index 97b74b7bc0ca..2d921539004f 100644 --- a/tests/plugins_tests/test_endpoint_plugins.py +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -81,8 +81,7 @@ def test_no_plugins_loaded_when_allowlist_is_empty_string( ): """`VLLM_PLUGINS=""` parses to `[""]`, not `None` (see `vllm.envs`), so it must be treated as a (non strict) allowlist matching no plugin name, not - as "unset". - """ + as "unset".""" monkeypatch.setenv("VLLM_PLUGINS", "") assert load_endpoint_plugins(("generate",)) == [] @@ -167,8 +166,7 @@ def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch) async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch): """`init_app_state` callers that never ran `build_app` (e.g. `run_batch.py`, which builds a bare `State()`) must not crash just - because `state.endpoint_plugins` was never set. - """ + because `state.endpoint_plugins` was never set.""" from starlette.datastructures import State monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") @@ -186,8 +184,7 @@ async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( """The CPU only render server has no `EngineClient` but a plugin eligible for the `render` task (`required_tasks` is `None` or includes `"render"`) still gets its routes attached at Phase A. Phase B passes `None` for - `engine_client` and it's up to the plugin to handle that. - """ + `engine_client` and it's up to the plugin to handle that.""" monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") args = _build_args() @@ -213,8 +210,7 @@ async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( @pytest.mark.asyncio async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch): """Phase A (attach) + Phase B (init) wired through `build_app` then - exercised with a real HTTP request against the worked example plugin. - """ + exercised with a real HTTP request against the worked example plugin.""" monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") args = _build_args() diff --git a/tests/plugins_tests/test_io_processor_plugins.py b/tests/plugins_tests/test_io_processor_plugins.py index c5c226cb54ef..14e1a2d5769c 100644 --- a/tests/plugins_tests/test_io_processor_plugins.py +++ b/tests/plugins_tests/test_io_processor_plugins.py @@ -38,8 +38,7 @@ def my_plugin_entry_points(): """Patch importlib.metadata.entry_points to expose a single 'my_plugin' entry point backed by DummyIOProcessor, exercising the full plugin-loading code path: entry_points → plugin.load() → func() → - resolve_obj_by_qualname → IOProcessor.__init__. - """ + resolve_obj_by_qualname → IOProcessor.__init__.""" qualname = f"{DummyIOProcessor.__module__}.{DummyIOProcessor.__qualname__}" ep = MagicMock() ep.name = "my_plugin" diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index a2fed9a2d47d..2bc463b593cb 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -382,8 +382,7 @@ def test_inc_get_quant_method_unquantized_layer_with_model_prefix( layer_factory, ) -> None: """``model.``-prefixed extra_config entries trigger early exit for both - LinearBase and ParallelLMHead layers. - """ + LinearBase and ParallelLMHead layers.""" config = make_config(extra_config={"model.layer": {"bits": 16}}) layer = layer_factory() @@ -396,8 +395,7 @@ def test_inc_get_quant_method_unquantized_routed_experts_with_model_prefix( monkeypatch, ) -> None: """``model.``-prefixed extra_config entries trigger early exit for - RoutedExperts layers, returning UnquantizedFusedMoEMethod. - """ + RoutedExperts layers, returning UnquantizedFusedMoEMethod.""" class DummyUnquantizedFusedMoEMethod: def __init__(self, moe_config) -> None: @@ -420,8 +418,7 @@ def __init__(self, moe_config) -> None: def test_inc_get_quant_method_unknown_layer_with_model_prefix_returns_none() -> None: """``model.``-prefixed extra_config entries return None for unhandled - layer types (non-Linear, non-MoE). - """ + layer types (non-Linear, non-MoE).""" config = make_config(extra_config={"model.unknown": {"bits": 16}}) method = config.get_quant_method(DummyLayer(), "unknown") @@ -736,8 +733,7 @@ class DummyMoeConfig: @pytest.mark.parametrize("bits", [4, 8]) def test_wna16_cuda_high_bit_skips_humming(monkeypatch, bits) -> None: """4/8-bit int stays on the Marlin/GPTQ/AWQ path even on CUDA so a single - model can mix high-bit Marlin and low-bit humming layers. - """ + model can mix high-bit Marlin and low-bit humming layers.""" called = {"humming": False} monkeypatch.setattr(current_platform, "is_cuda", lambda: True) @@ -1668,8 +1664,7 @@ def test_inc_get_quant_method_unquantized_moe_returns_unquantized( monkeypatch, ) -> None: """Early-exit returns UnquantizedFusedMoEMethod for FusedMoE layers - when extra_config has bits >= 16. - """ + when extra_config has bits >= 16.""" config = make_config(extra_config={"layer": {"bits": 16}}) layer = object.__new__(RoutedExperts) layer.moe_config = None # UnquantizedFusedMoEMethod accepts moe_config diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 44a508ed57ce..65e119b083fd 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -213,8 +213,7 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales( ): """The VL wrapper must map ``.scale`` keys to the parameter the linear quant method registers, as the text model does. A hardcoded - ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints. - """ + ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" from vllm.models.deepseek_v41.nvidia import model as model_module from vllm.models.deepseek_v41.nvidia import vl_model as vl_module diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index 3adcb020b654..b71e952b769a 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -160,8 +160,7 @@ def is_supported_config(*args, **kwargs): @pytest.fixture def mxfp4_oracle_config(): """Stub the config the oracle reads (``model_config.quantization_config``) - so backend dispatch resolves without a real model / user override. - """ + so backend dispatch resolves without a real model / user override.""" from unittest.mock import patch with patch( diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py index bf508dda8992..dbec989df3f9 100644 --- a/tests/quantization/test_humming_mxfp4_block_fp8.py +++ b/tests/quantization/test_humming_mxfp4_block_fp8.py @@ -85,8 +85,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): (FP8 dtype, [1, 128] block, not per-act-token). This is what makes the DeepEP prepare/finalize step quantize activations to block FP8 *before* the all-to-all dispatch instead of deferring to Humming -- and what makes - HummingExpertsBase.expects_unquantized_inputs return False. - """ + HummingExpertsBase.expects_unquantized_inputs return False.""" from vllm.model_executor.layers.quantization.utils.humming_utils import ( make_humming_moe_quant_config, ) @@ -108,8 +107,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): def test_default_activation_quant_config_defers_to_humming(): """Without a block activation shape the config stays per-token (the deferred path): Humming quantizes internally, preserving pre-existing behavior for - every non-block-FP8 scheme. - """ + every non-block-FP8 scheme.""" from vllm.model_executor.layers.quantization.utils.humming_utils import ( make_humming_moe_quant_config, ) diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py index 048b0e645abb..8c838a785b63 100644 --- a/tests/quantization/test_int8_moe_oracle.py +++ b/tests/quantization/test_int8_moe_oracle.py @@ -70,8 +70,7 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: ) def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor - weights) select the Triton backend. - """ + weights) select the Triton backend.""" config = _make_int8_moe_config() backend, experts_cls = select_int8_moe_backend( config, weight_key=weight_key, activation_key=activation_key diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 8e590d7c3f9e..60b207c1840f 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -587,8 +587,7 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i def test_modelopt_nvfp4_config_dispatches_w4a4_method(): """``quant_method="NVFP4"`` (W4A4) resolves to a ``(kNvfp4Static, kNvfp4Dynamic)`` QuantSpec under the generic - ``ModelOptLinearMethod``. - """ + ``ModelOptLinearMethod``.""" from vllm.model_executor.layers.linear import LinearBase config = ModelOptNvFp4Config( @@ -637,8 +636,7 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): """The bespoke-method escape hatch: a format registered in ``LINEAR_METHOD_BUILDERS`` routes that algo to its own method instead of the generic ``ModelOptLinearMethod``. This is how a format that cannot be a - ``(weight, activation)`` key pair plugs into dispatch. - """ + ``(weight, activation)`` key pair plugs into dispatch.""" from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import modelopt as m @@ -669,8 +667,7 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``: ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The - generic method routes this through ``select_linear_kernel``. - """ + generic method routes this through ``select_linear_kernel``.""" from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.quantization.modelopt import ( RuntimeDtypes, @@ -941,8 +938,7 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch): def test_modelopt_fp8_pb_wo_rejects_non_128_input(): """Input width must still be a multiple of 128 (same as #53132, which only pads the output). A partial input block is refused loudly rather than - silently loading wrong scales. - """ + silently loading wrong scales.""" from vllm.model_executor.layers.quantization import modelopt as mo from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Static128BlockSym, diff --git a/tests/quantization/test_online_mxfp4.py b/tests/quantization/test_online_mxfp4.py index 8d29b4602cd4..9495a15ba681 100644 --- a/tests/quantization/test_online_mxfp4.py +++ b/tests/quantization/test_online_mxfp4.py @@ -411,8 +411,7 @@ def scatter_into_padded( ) -> tuple[torch.Tensor, torch.Tensor]: """Place the source weights into buffers of the padded size, the way the weight loader writes only the unpadded slice of a larger - allocation. - """ + allocation.""" w13 = torch.full( (num_experts, 2 * intermediate_size, hidden_size), padding_value, diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 419b079d21b4..4d6193211829 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -490,8 +490,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig): ) def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str): """Per-token-head should set _k_scale=1.0, _v_scale=1.0 - and delete checkpoint attrs. - """ + and delete checkpoint attrs.""" from vllm.model_executor.layers.quantization.kv_cache import ( BaseKVCacheMethod, ) diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py index dad70ea8d6f7..ea5b374e6e0d 100644 --- a/tests/reasoning/test_kimi_k3_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -276,8 +276,7 @@ def last(needle: list[int]) -> int: def test_is_reasoning_end_streaming_only_scans_the_step_window(): """The decode-step check must not re-derive the answer from the whole sequence: an already-closed think block earlier in the sequence is the - engine's business (it latches `reasoning_ended`), not this call's. - """ + engine's business (it latches `reasoning_ended`), not this call's.""" parser = KimiK3ReasoningParser(DummyTokenizer()) prompt = [*CLOSE_IDS, *OPEN_IDS, 9, 9, 9] @@ -287,8 +286,7 @@ def test_is_reasoning_end_streaming_only_scans_the_step_window(): def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): """A 3-token marker can be split over decode steps; the check carries the - preceding len(marker)-1 tokens so the final token still completes it. - """ + preceding len(marker)-1 tokens so the final token still completes it.""" parser = KimiK3ReasoningParser(DummyTokenizer()) history = [*OPEN_IDS, 5, *CLOSE_IDS[:2]] @@ -299,8 +297,7 @@ def test_is_reasoning_end_streaming_detects_marker_across_step_boundary(): def test_is_reasoning_end_streaming_reopened_block_is_not_ended(): """Kimi K3 may close and immediately reopen the think channel inside one window (speculative decoding lands several tokens per step). The newest - marker wins, matching is_reasoning_end. - """ + marker wins, matching is_reasoning_end.""" parser = KimiK3ReasoningParser(DummyTokenizer()) window = [*CLOSE_IDS, *OPEN_IDS] @@ -327,8 +324,7 @@ def test_is_reasoning_end_streaming_thinking_disabled(): def test_reasoning_end_matches_reference_over_marker_dense_sequences(seed): """Both the full-sequence check and the per-step check must agree with the two-scan reference, on sequences where the markers (which share the suffix - [2, 3]) collide and overlap constantly. - """ + [2, 3]) collide and overlap constantly.""" import random rnd = random.Random(seed) diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index bcaff9b94421..bbd4dc5c5f2e 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -116,8 +116,7 @@ def test_prompt_embeds_keys_registered(): def test_ensure_placeholder_token_is_single_token_and_idempotent(tokenizer): """Ensure the placeholder token is a single token and that multiple calls to - "ensure" are idempotent, across all tokenizer families. - """ + "ensure" are idempotent, across all tokenizer families.""" tid1 = _ensure_prompt_embeds_placeholder_token(tokenizer) tid2 = _ensure_prompt_embeds_placeholder_token(tokenizer) assert tid1 == tid2 @@ -470,8 +469,7 @@ def test_build_mixed_prompt_embeds(stream): @pytest.mark.parametrize("role", ["user", "system"]) async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): """Full renderer pipeline: parse -> chat template -> expand -> locate - -> build mixed prompt, across tokenizers, roles, and sync/async. - """ + -> build mixed prompt, across tokenizers, roles, and sync/async.""" tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE tid = _ensure_prompt_embeds_placeholder_token(tokenizer) @@ -526,8 +524,7 @@ async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): @pytest.mark.asyncio async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn): """Full pipeline with prompt_embeds spread across system + user messages, - verifying ordering and positioning in the final token stream. - """ + verifying ordering and positioning in the final token stream.""" tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE tid = _ensure_prompt_embeds_placeholder_token(tokenizer) diff --git a/tests/renderers/test_gemma4_chat_template.py b/tests/renderers/test_gemma4_chat_template.py index c244b2672312..7d6b0b05f814 100644 --- a/tests/renderers/test_gemma4_chat_template.py +++ b/tests/renderers/test_gemma4_chat_template.py @@ -33,8 +33,7 @@ def _render(template, messages, **kwargs): class TestGemma4ChatTemplate: def test_basic_multiturn_thinking_disabled(self, gemma4_template): """With enable_thinking=False (default), generation prompt ends with - an empty thought channel to suppress thinking. - """ + an empty thought channel to suppress thinking.""" messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, @@ -50,8 +49,7 @@ def test_basic_multiturn_thinking_disabled(self, gemma4_template): def test_basic_multiturn_thinking_enabled(self, gemma4_template): """With enable_thinking=True, generation prompt ends with model - turn opener (no thought suppression). - """ + turn opener (no thought suppression).""" messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, @@ -146,8 +144,7 @@ def test_tool_calls_in_assistant(self, gemma4_template): def test_tool_responses_openai_style(self, gemma4_template): """role='tool' messages are formatted as <|tool_response> blocks - with content dumped as-is. - """ + with content dumped as-is.""" messages = [ {"role": "user", "content": "Weather?"}, { @@ -205,8 +202,7 @@ def test_tool_responses_legacy_style(self, gemma4_template): def test_generation_prompt_not_after_tool_response(self, gemma4_template): """add_generation_prompt=True should NOT add <|turn>model when the - last message type was tool_response (the model turn continues). - """ + last message type was tool_response (the model turn continues).""" messages = [ {"role": "user", "content": "Weather?"}, { @@ -233,8 +229,7 @@ def test_generation_prompt_not_after_tool_response(self, gemma4_template): def test_reasoning_in_tool_chains(self, gemma4_template): """Reasoning field on assistant with tool_calls after last user - message emits <|channel>thought\\n.... - """ + message emits <|channel>thought\\n....""" messages = [ {"role": "user", "content": "Calculate something"}, { @@ -280,8 +275,7 @@ def test_reasoning_not_before_last_user(self, gemma4_template): def test_strip_thinking_in_model_content(self, gemma4_template): """<|channel>... in model content is stripped by the - strip_thinking macro. - """ + strip_thinking macro.""" messages = [ {"role": "user", "content": "Hi"}, { @@ -295,8 +289,7 @@ def test_strip_thinking_in_model_content(self, gemma4_template): def test_multi_turn_tool_chain(self, gemma4_template): """assistant->tool->assistant->tool produces exactly one - <|turn>model (later assistants continue the same turn). - """ + <|turn>model (later assistants continue the same turn).""" messages = [ {"role": "user", "content": "Do two things"}, { @@ -353,8 +346,7 @@ def test_format_argument_types(self, gemma4_template): def test_tool_response_with_multimodal_content(self, gemma4_template): """Multimodal placeholders in tool messages are emitted after the - tool_response block. - """ + tool_response block.""" messages = [ {"role": "user", "content": "Download the image and describe it."}, { @@ -388,8 +380,7 @@ def test_tool_response_with_multimodal_content(self, gemma4_template): def test_tool_response_with_all_modalities(self, gemma4_template): """All multimodal types (image, audio, video) in a single tool - response are rendered. - """ + response are rendered.""" messages = [ {"role": "user", "content": "Process media"}, { diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index 624108a20689..833656aa3c31 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -308,8 +308,7 @@ def apply_chat_template(self, conversation, **kwargs): def test_resolve_chat_template_resolves_name(): """When chat_template is a name, resolve_chat_template should return - the actual Jinja content so that kwargs detection works correctly. - """ + the actual Jinja content so that kwargs detection works correctly.""" from unittest.mock import MagicMock jinja_content = "{{ messages }}{% if tools %}{{ tools }}{% endif %}" @@ -331,8 +330,7 @@ def test_resolve_chat_template_resolves_name(): def test_resolve_chat_template_kwargs_with_template_name(): """Ensures template kwargs are not silently dropped when chat_template - was originally a template name that has been resolved to Jinja content. - """ + was originally a template name that has been resolved to Jinja content.""" from unittest.mock import MagicMock jinja_content = ( @@ -934,8 +932,7 @@ def test_developer_only_no_prior_system(self, model_config, tokenizer): class TestApplyChatTemplateEffortTolerant: """Chat templates that reject unsupported reasoning_effort values should - surface a 400-style client error instead of crashing the request. - """ + surface a 400-style client error instead of crashing the request.""" @pytest.fixture def model_config(self): diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py index 992177a3a633..8987782d71bc 100644 --- a/tests/renderers/test_process_multi_modal_uuids.py +++ b/tests/renderers/test_process_multi_modal_uuids.py @@ -49,8 +49,7 @@ def _build_text_only_renderer() -> HfRenderer: def test_text_only_model_mm_data_maps_to_bad_request(): """Sending multimodal data to a text-only model is a client mistake, so it - must surface as a ValueError and reach the client as HTTP 400, not 500. - """ + must surface as a ValueError and reach the client as HTTP 400, not 500.""" renderer = _build_text_only_renderer() with pytest.raises(ValueError, match="text-only") as exc_info: diff --git a/tests/renderers/test_prompt_embeds_payload_validation.py b/tests/renderers/test_prompt_embeds_payload_validation.py index fec937a4b723..bec28c4c49ac 100644 --- a/tests/renderers/test_prompt_embeds_payload_validation.py +++ b/tests/renderers/test_prompt_embeds_payload_validation.py @@ -127,8 +127,7 @@ def test_a_long_reason_is_truncated_and_says_how_much_was_left_out(model_config) def test_a_short_reason_is_passed_through_whole(model_config): """A reason that fits is not annotated -- the count only appears when it is - telling the reader something. - """ + telling the reader something.""" reason = "invalid magic number; corrupt file?" assert len(reason) <= _MAX_EMBED_ERROR_REASON_CHARS diff --git a/tests/renderers/test_sparse_tensor_concurrent_race.py b/tests/renderers/test_sparse_tensor_concurrent_race.py index c627dc31cbad..a1683eee0bff 100644 --- a/tests/renderers/test_sparse_tensor_concurrent_race.py +++ b/tests/renderers/test_sparse_tensor_concurrent_race.py @@ -218,8 +218,7 @@ class TestCrossLoaderLockSharing: def test_prompt_and_image_share_lock(self, model_config): """Prompt and image loaders cannot run their sparse guards - concurrently. - """ + concurrently.""" valid_prompt = _encode_tensor(_create_valid_dense_tensor()) valid_image = _encode_tensor(_create_valid_dense_tensor(hidden_size=10)) diff --git a/tests/renderers/test_sparse_tensor_validation.py b/tests/renderers/test_sparse_tensor_validation.py index d90ad1e09a18..6cb26a8a1de8 100644 --- a/tests/renderers/test_sparse_tensor_validation.py +++ b/tests/renderers/test_sparse_tensor_validation.py @@ -157,8 +157,7 @@ def test_negative_indices_rejected(self, model_config): def test_hidden_size_mismatch_rejected(self, model_config): """Tensors whose trailing dim doesn't match the model's hidden_size - must be rejected at parse time. - """ + must be rejected at parse time.""" # opt-125m has hidden_size=768, passing 512 triggers the check. wrong_hidden = torch.randn(10, 512, dtype=torch.float32) encoded = _encode_tensor(wrong_hidden) @@ -169,8 +168,7 @@ def test_hidden_size_mismatch_rejected(self, model_config): def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): """Tensors whose dtype doesn't match the model's dtype but are still floating-point are cast, since API clients generally can't know the - server's `--dtype` setting ahead of time. - """ + server's `--dtype` setting ahead of time.""" # Fixture pins model dtype to float32, upload a bfloat16 tensor. mismatched_float = torch.randn(10, 768, dtype=torch.bfloat16) encoded = _encode_tensor(mismatched_float) @@ -183,8 +181,7 @@ def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): def test_non_float_dtype_rejected(self, model_config): """Non-floating-point dtypes cannot be safely cast for embeddings (e.g. integer tensors almost certainly indicate caller confusion), - so they are rejected at parse time. - """ + so they are rejected at parse time.""" non_float = torch.randint(0, 100, (10, 768), dtype=torch.int32) encoded = _encode_tensor(non_float) @@ -193,8 +190,7 @@ def test_non_float_dtype_rejected(self, model_config): def test_non_2d_tensor_rejected(self, model_config): """Tensors that aren't 2D (even after squeezing a leading dim) - must be rejected with a clear error. - """ + must be rejected with a clear error.""" # A 1D tensor cannot be interpreted as (num_tokens, hidden_size). bad = torch.randn(768, dtype=torch.float32) encoded = _encode_tensor(bad) @@ -204,8 +200,7 @@ def test_non_2d_tensor_rejected(self, model_config): def test_non_tensor_payload_rejected(self, model_config): """Deserializing to a non-Tensor object must raise a clear error - instead of propagating an AssertionError. - """ + instead of propagating an AssertionError.""" # `torch.save` will serialize a plain dict; `weights_only=True` allows # loading built-in containers, so this exercises the isinstance check. buffer = io.BytesIO() @@ -463,8 +458,7 @@ def test_oversized_sparse_tensor_rejected(self, small_decode_limit): def test_default_limit_would_reject_the_real_payload(self): """The shipped default has to bound the actual attack, not just the shrunken shape used above. Checked arithmetically so no test ever asks - the allocator for 4 TiB. - """ + the allocator for 4 TiB.""" dense_bytes = BOMB_SHAPE[0] * BOMB_SHAPE[1] * 4 assert dense_bytes > envs.VLLM_MAX_EMBED_DECODE_BYTES > 0 @@ -522,8 +516,7 @@ def test_media_io_rejects_oversized_file( @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO) def test_media_io_rejects_non_tensor_payload(self, io_cls, parameter): """These three had no `isinstance` check, so a non-tensor payload - surfaced as an `AttributeError` (500) rather than a 4xx. - """ + surfaced as an `AttributeError` (500) rather than a 4xx.""" buffer = io.BytesIO() torch.save({"not": "a tensor"}, buffer) encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") diff --git a/tests/renderers/test_token_offsets.py b/tests/renderers/test_token_offsets.py index 973ea175cbb9..fc9192bb43e7 100644 --- a/tests/renderers/test_token_offsets.py +++ b/tests/renderers/test_token_offsets.py @@ -26,8 +26,7 @@ def fast_tokenizer(): def _make_base_renderer_with(tokenizer): """Build a minimal BaseRenderer subclass that exposes the tokenizer so we can call ``_tokenize_prompt`` directly. BaseRenderer is abstract because of - ``render_messages``; we just need a stub. - """ + ``render_messages``; we just need a stub.""" from vllm.renderers.base import BaseRenderer class _StubRenderer(BaseRenderer): @@ -92,8 +91,7 @@ def test_fast_tokenizer_with_flag_returns_offsets(self, fast_tokenizer): def test_base_renderer_without_override_yields_no_offsets(self, fast_tokenizer): """A renderer that does not override ``_can_produce_offsets`` never emits offsets, even with a fast tokenizer and the flag set. This locks - in the base-default-False / subclass-override design. - """ + in the base-default-False / subclass-override design.""" from vllm.renderers.base import BaseRenderer class _BareRenderer(BaseRenderer): @@ -127,8 +125,7 @@ def test_default_flag_no_offsets(self, fast_tokenizer): def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): """Force is_fast=False to simulate a Slow tokenizer: the flag is set - but offsets must not be returned because it cannot produce them. - """ + but offsets must not be returned because it cannot produce them.""" from unittest.mock import PropertyMock, patch renderer = _make_base_renderer_with(fast_tokenizer) @@ -147,8 +144,7 @@ def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): @pytest.mark.parametrize("mm_key", ["multi_modal_data", "multi_modal_uuids"]) def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): """Offsets index the text prompt, which is meaningless once multimodal - data is interleaved, so they are suppressed when MM inputs are present. - """ + data is interleaved, so they are suppressed when MM inputs are present.""" renderer = _make_base_renderer_with(fast_tokenizer) params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) prompt = {"prompt": "Hello.", mm_key: {"image": ["x"]}} @@ -160,8 +156,7 @@ def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): @pytest.mark.asyncio async def test_tokenize_prompt_async_returns_offsets(self, fast_tokenizer): """The async path offloads the sync tokenizer; it must yield the same - offsets as the sync path. - """ + offsets as the sync path.""" renderer = _make_base_renderer_with(fast_tokenizer) params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index 572bc76caa6b..5d217433db71 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -40,8 +40,7 @@ def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( ): """Torch's autodetection can emit a bogus arch=compute_* half (e.g. capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS - target must win, while PTX-only entries keep the virtual arch. - """ + target must win, while PTX-only entries keep the virtual arch.""" repo_root = Path(__file__).parents[1] script = tmp_path / "test_extract_archs.cmake" script.write_text( diff --git a/tests/test_config.py b/tests/test_config.py index 4b37b4f56e03..ce2580934a01 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -441,8 +441,7 @@ def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled - MRV1 paths and off breakable cudagraphs by default. - """ + MRV1 paths and off breakable cudagraphs by default.""" from vllm.config.vllm import ( ROCM_DEFAULT_MRV1_ARCHITECTURES, default_breakable_cudagraph_architectures, @@ -657,8 +656,7 @@ def test_v2_model_runner_supports_extract_hidden_states(): def test_dflash2_draft_forces_v2_model_runner(): """A DFlash2 draft must reach the V2 speculator, the only one that runs its - candidate selector; on V1 it would draft as DFlash1 without raising. - """ + candidate selector; on V1 it would draft as DFlash1 without raising.""" def config(method, architectures): return SimpleNamespace( @@ -771,8 +769,7 @@ def test_late_piecewise_restrictions_without_compilation(monkeypatch, engine_kwa def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): """Cudagraph memory profiling uses a minimal KV cache, so the Mamba - block-count guard must only fire for the real cache sizing. - """ + block-count guard must only fire for the real cache sizing.""" kv_cache_config = SimpleNamespace(has_mamba_layers=True, num_blocks=4) compilation_config = CompilationConfig( @@ -1097,8 +1094,7 @@ def test_v1_model_runner_rejects_v2_only_features(): def test_batch_sharded_sampling_rejects_return_sampling_mask(): """The batch-sharded gather drops sampling masks, so the combination must - fail loudly instead of returning ``sampling_mask=None``. - """ + fail loudly instead of returning ``sampling_mask=None``.""" config = SimpleNamespace( parallel_config=SimpleNamespace( enable_batch_sharded_sampling=True, tensor_parallel_size=2 @@ -2195,8 +2191,7 @@ def __init__(self, model: str, tokenizer: str): @patch("vllm.transformers_utils.runai_utils.ObjectStorageModel.pull_files") def test_s3_url_model_tokenizer_paths(mock_pull_files, s3_url): """Test that S3 URLs create deterministic local directories for model and - tokenizer. - """ + tokenizer.""" # Mock pull_files to avoid actually downloading files during tests mock_pull_files.return_value = None @@ -2294,8 +2289,7 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files): @patch("vllm.transformers_utils.runai_utils.ObjectStorageModel.pull_files") def test_s3_url_different_model_and_tokenizer(mock_pull_files): """Test that when model and tokenizer are different cloud URIs, - pull_files receives the correct URI for each. - """ + pull_files receives the correct URI for each.""" mock_pull_files.return_value = None model_url = "s3://bucket/model/" @@ -2840,8 +2834,7 @@ def test_vllm_config_explicit_overrides(): def test_fusion_pass_op_priority(): """This test checks that custom op enablement & IR op priority - correctly control default fusions - """ + correctly control default fusions""" # Default config, O2, rms_norm+quant fusion disabled cfg1 = VllmConfig() assert not cfg1.compilation_config.pass_config.fuse_norm_quant @@ -2925,8 +2918,7 @@ def test_needs_dp_coordination( def test_fault_tolerance_requires_single_api_server(): """Fault tolerance assumes one AsyncMPClient manages all engines, so it - is incompatible with API server scale-out (_api_process_count > 1). - """ + is incompatible with API server scale-out (_api_process_count > 1).""" with pytest.raises(ValueError, match="single API server"): ParallelConfig(enable_fault_tolerance=True, _api_process_count=2) @@ -2936,8 +2928,7 @@ def test_fault_tolerance_requires_single_api_server(): def test_renderer_num_workers_with_mm_cache(): """Disallow renderer_num_workers > 1 with the mm processor cache only for - pooling models, whose preprocessing runs on the renderer workers. - """ + pooling models, whose preprocessing runs on the renderer workers.""" mm_model = "Qwen/Qwen2-VL-2B-Instruct" # Should raise: pooling + multi-worker + cache enabled (default cache_gb=4) @@ -3263,8 +3254,7 @@ def test_watermarking_forces_model_runner_v2(monkeypatch): def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls): """Regression test: MTP + runai_streamer should use model_weights (original S3 URL) for the draft model, not model (local cache dir set by - pull_runai_model_from_obj_storage). - """ + pull_runai_model_from_obj_storage).""" from unittest.mock import MagicMock s3_url = "s3://my-bucket/Qwen3-35B-A3B-FP8" diff --git a/tests/test_logger.py b/tests/test_logger.py index f051fe205b8c..2ff100151b29 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -51,8 +51,7 @@ def test_trace_function_call(): def test_default_vllm_root_logger_configuration(monkeypatch): """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default - behavior is activated. - """ + behavior is activated.""" monkeypatch.setenv("VLLM_LOGGING_COLOR", "0") _configure_vllm_root_logger() @@ -76,8 +75,7 @@ def test_default_vllm_root_logger_configuration(monkeypatch): def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypatch): """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default - behavior is activated. - """ + behavior is activated.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False) @@ -107,8 +105,7 @@ def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypa def test_logger_configuring_can_be_disabled(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur. - """ + configuration occur.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0") monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False) @@ -120,8 +117,7 @@ def test_logger_configuring_can_be_disabled(monkeypatch): def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs. - """ + configuration occurs.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") monkeypatch.setenv( "VLLM_LOGGING_CONFIG_PATH", @@ -137,8 +133,7 @@ def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist(monke def test_an_error_is_raised_when_custom_logging_config_is_invalid_json(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs. - """ + configuration occurs.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file: @@ -165,8 +160,7 @@ def test_an_error_is_raised_when_custom_logging_config_is_unexpected_json( ): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however it fails before any change in behavior or - configuration occurs. - """ + configuration occurs.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file: @@ -182,8 +176,7 @@ def test_an_error_is_raised_when_custom_logging_config_is_unexpected_json( def test_custom_logging_config_is_parsed_and_used_when_provided(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur. - """ + configuration occur.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1") valid_logging_config = { @@ -207,8 +200,7 @@ def test_custom_logging_config_is_parsed_and_used_when_provided(monkeypatch): def test_custom_logging_config_causes_an_error_if_configure_logging_is_off(monkeypatch): """This test calls _configure_vllm_root_logger again to test custom logging config behavior, however mocks are used to ensure no changes in behavior or - configuration occur. - """ + configuration occur.""" monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0") valid_logging_config = { diff --git a/tests/test_ray_env.py b/tests/test_ray_env.py index 43ec51a8f81b..945b2d80b691 100644 --- a/tests/test_ray_env.py +++ b/tests/test_ray_env.py @@ -15,8 +15,7 @@ class TestDefaultPrefixes: """Built-in prefixes (VLLM_, LMCACHE_, NCCL_, UCX_, HF_, HUGGING_FACE_) - should be forwarded without any extra configuration. - """ + should be forwarded without any extra configuration.""" @patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False) def test_lmcache_prefix(self): @@ -54,8 +53,7 @@ class TestDefaultExtraVars: def test_pythonhashseed_in_result(self): """PYTHONHASHSEED should always be in the result set (as a name to - copy) regardless of whether it is actually set in os.environ. - """ + copy) regardless of whether it is actually set in os.environ.""" result = get_env_vars_to_copy() assert "PYTHONHASHSEED" in result diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index aac88553e6e9..a55aa57666dd 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -452,8 +452,7 @@ def test_deepseek_v4_image_blocks_become_placeholders(): def test_deepseek_v4_image_sentinel_ids_match_tokenizer(): """The borrowed sentinel ids must line up with the reserved - ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer. - """ + ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer.""" from vllm.models.deepseek_v4.common.mm_preprocess import ( IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_TOKEN_NAMES, diff --git a/tests/tokenizers_/test_hf.py b/tests/tokenizers_/test_hf.py index f1dfd046b93e..61c81302f071 100644 --- a/tests/tokenizers_/test_hf.py +++ b/tests/tokenizers_/test_hf.py @@ -51,8 +51,7 @@ def _check_consistency(target: TokenizerLike, expected: TokenizerLike): def test_thread_pool_tokenizer_pickle(model_id: str): """Regression test for issue #45433: the thread-pool tokenizer wrapper reconstructs through maybe_make_thread_pool on unpickling, which used to - fall off the end and return None. - """ + fall off the end and return None.""" reference_tokenizer = AutoTokenizer.from_pretrained(model_id) pooled_tokenizer = maybe_make_thread_pool(deepcopy(reference_tokenizer)) diff --git a/tests/tool_parsers/test_apertus_tool_parser.py b/tests/tool_parsers/test_apertus_tool_parser.py index f3cf5fcaf4ea..5ab0a2bcc732 100644 --- a/tests/tool_parsers/test_apertus_tool_parser.py +++ b/tests/tool_parsers/test_apertus_tool_parser.py @@ -341,8 +341,7 @@ def test_streaming_partial_tag_buffering(self, parser, mock_request): def test_mtp_streaming_massive_chunk(self, parser, mock_request): """Simulates MTP predicting text, tool calls, - and trailing text all in a single chunk. - """ + and trailing text all in a single chunk.""" chunks = [ "Sure! " '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>' @@ -375,8 +374,7 @@ def test_mtp_streaming_multiple_tools_burst(self, parser, mock_request): def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request): """Simulates MTP chunks that jump over entire tools - (e.g., from middle of tool 1 to middle of tool 3). - """ + (e.g., from middle of tool 1 to middle of tool 3).""" chunks = [ '<|tools_prefix|>[{"t1": {"a": 1}', '}, {"t2": {"b": 2}}, {"t3": {"c": 3', @@ -395,8 +393,7 @@ def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request): def test_vllm_streaming_character_by_character(self, parser, mock_request): """Simulates worst-case vLLM fragmentation where - chunks arrive character-by-character. - """ + chunks arrive character-by-character.""" text = ( 'Hi <|tools_prefix|>[{"get_weather": ' '{"location": "London"}}]<|tools_suffix|> ' @@ -414,8 +411,7 @@ def test_vllm_streaming_character_by_character(self, parser, mock_request): def test_vllm_streaming_empty_deltas(self, parser, mock_request): """Simulates vLLM stream producing empty string chunks - (e.g., hidden tokens or artifacts). - """ + (e.g., hidden tokens or artifacts).""" chunks = [ "Wait", "", diff --git a/tests/tool_parsers/test_deepseekv32_tool_parser.py b/tests/tool_parsers/test_deepseekv32_tool_parser.py index d0d032167a37..40ad6033d3bf 100644 --- a/tests/tool_parsers/test_deepseekv32_tool_parser.py +++ b/tests/tool_parsers/test_deepseekv32_tool_parser.py @@ -173,8 +173,7 @@ def test_type_conversion_in_non_streaming(self): def test_string_attr_true_coerced_by_schema(self): """string="true" delivers a string, but the engine's schema-aware - type fixer coerces it to the schema type (integer). - """ + type fixer coerces it to the schema type (integer).""" tool = ChatCompletionToolsParam( function=FunctionDefinition( name="score", @@ -288,8 +287,7 @@ def test_composed_schema_converts_object_and_array_params(self): def test_string_attr_true_coerced_by_composed_schema(self): """string="true" delivers a JSON string, but the engine's schema-aware - type fixer coerces it to the composed schema type (object). - """ + type fixer coerces it to the composed schema type (object).""" tool = ChatCompletionToolsParam( function=FunctionDefinition( name="set_timer", @@ -323,8 +321,7 @@ def test_string_attr_true_coerced_by_composed_schema(self): def test_arguments_wrapper_repaired(self): """A single 'arguments' wrapper parameter must be unwrapped when it - is not part of the tool schema and the inner object matches schema fields. - """ + is not part of the tool schema and the inner object matches schema fields.""" tool = ChatCompletionToolsParam( function=FunctionDefinition( name="get_weather", @@ -495,8 +492,7 @@ def test_multi_typed_null_value(self): def test_null_coerced_back_to_string_by_schema(self): """string="false" with 'null' is json-parsed to None, but the - engine's schema fixer coerces it back to "null" for string schemas. - """ + engine's schema fixer coerces it back to "null" for string schemas.""" tool = ChatCompletionToolsParam( function=FunctionDefinition( name="echo", @@ -659,8 +655,7 @@ def test_type_conversion_in_streaming(self): def test_string_attr_true_coerced_by_schema_streaming(self): """Streaming: string='true' delivers a string but the engine's - schema fixer coerces it to the schema type (integer). - """ + schema fixer coerces it to the schema type (integer).""" tool = ChatCompletionToolsParam( function=FunctionDefinition( name="score", @@ -976,8 +971,7 @@ def test_emits_arguments_before_invoke_completes(self, parser): def test_no_marker_leak_chunked(self, parser): """Chunked streaming must NOT leak DSML start-marker fragments - as content (GitHub #40801). - """ + as content (GitHub #40801).""" full_text = build_tool_call("fn", {"k": "v"}) deltas = self._stream_chunked(parser, full_text, chunk_size=5) content = "".join(d.content for d in deltas if d.content is not None) @@ -987,8 +981,7 @@ def test_no_marker_leak_chunked(self, parser): def test_no_marker_leak_with_prefix_chunked(self, parser): """Content before a tool call must not include start-marker - fragments when chunked (GitHub #40801). - """ + fragments when chunked (GitHub #40801).""" full_text = "Hello!" + build_tool_call("fn", {"a": "b"}) deltas = self._stream_chunked(parser, full_text, chunk_size=5) content = "".join(d.content for d in deltas if d.content is not None) @@ -1000,8 +993,7 @@ def test_no_marker_leak_with_prefix_chunked(self, parser): def test_no_marker_leak_char_by_char(self, parser): """Character-by-character streaming must not leak marker - fragments (GitHub #40801). - """ + fragments (GitHub #40801).""" full_text = build_tool_call("fn", {"k": "v"}) deltas = self._stream_chunked(parser, full_text, chunk_size=1) content = "".join(d.content for d in deltas if d.content is not None) @@ -1011,8 +1003,7 @@ def test_no_marker_leak_char_by_char(self, parser): def test_no_marker_leak_all_split_points(self, parser): """Start token split at every possible boundary must not - leak (GitHub #40801). - """ + leak (GitHub #40801).""" for chunk_size in range(1, len(FC_START) + 2): p = make_parser() full_text = build_tool_call("fn", {"k": "v"}) @@ -1024,8 +1015,7 @@ def test_no_marker_leak_all_split_points(self, parser): def test_false_partial_marker_emitted(self, parser): """Text ending with a prefix of the start token that turns out - NOT to be a marker must still be emitted as content. - """ + NOT to be a marker must still be emitted as content.""" full_text = "<｜DSM some regular text" deltas = self._stream_chunked(parser, full_text, chunk_size=3) content = "".join(d.content for d in deltas if d.content is not None) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index bfb12ff15786..97ed0cd69d4c 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -124,8 +124,7 @@ def mock_request(): class TestParseGemma4Args: """Values are returned as strings; type coercion to proper JSON types - happens at the engine layer. - """ + happens at the engine layer.""" def test_empty_string(self): assert _parse_gemma4_args("") == {} diff --git a/tests/tool_parsers/test_gigachat3_tool_parser.py b/tests/tool_parsers/test_gigachat3_tool_parser.py index be9eb06f8803..e5a1692ffa39 100644 --- a/tests/tool_parsers/test_gigachat3_tool_parser.py +++ b/tests/tool_parsers/test_gigachat3_tool_parser.py @@ -18,8 +18,7 @@ @pytest.fixture(scope="function") def default_tokenizer() -> TokenizerLike: """Override module-scoped default_tokenizer because gigachat tests - mutate the tokenizer via ``add_tokens``. - """ + mutate the tokenizer via ``add_tokens``.""" return AutoTokenizer.from_pretrained("openai-community/gpt2") diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py index 740f4c30130f..36b04b67b195 100644 --- a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Named tool choice for Kimi K3: allowed when the XTML structural tag is -attached (strict tool calling), rejected otherwise. -""" +attached (strict tool calling), rejected otherwise.""" import pytest diff --git a/tests/tool_parsers/test_lfm2_tool_parser.py b/tests/tool_parsers/test_lfm2_tool_parser.py index 599aec7d9236..9ed30e159884 100644 --- a/tests/tool_parsers/test_lfm2_tool_parser.py +++ b/tests/tool_parsers/test_lfm2_tool_parser.py @@ -416,8 +416,7 @@ def test_tool_call( def test_empty_tool_call_block(body: str, lfm2_tokenizer: TokenizerLike): """An empty block ([] or [ ]) means the model called no tools; it must not report tools_called=True with zero tool calls (the invariant - run_tool_extraction asserts). - """ + run_tool_extraction asserts).""" cls = ToolParserManager.get_tool_parser("lfm2") model_output = f"{TOOL_CALL_START}[{body}]{TOOL_CALL_END}" @@ -437,8 +436,7 @@ def test_whitespace_after_start_token(lfm2_tokenizer: TokenizerLike): not break parsing. The streaming path used to feed the indented text to ast.parse verbatim, so every completion candidate raised IndentationError and the call was silently dropped (non-streaming - stripped it and succeeded). - """ + stripped it and succeeded).""" cls = ToolParserManager.get_tool_parser("lfm2") model_output = f"{TOOL_CALL_START} [{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -478,8 +476,7 @@ def test_streaming_full_block_and_trailing_in_single_delta( ): """The entire <|tool_call_start|>[...]<|tool_call_end|> block plus trailing assistant text arrive in one delta. Trailing content must - still be emitted — not silently dropped. - """ + still be emitted — not silently dropped.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = f"{TOOL_CALL_START}[{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}\nDone." @@ -494,8 +491,7 @@ def test_streaming_leading_content_and_full_block_in_single_delta( lfm2_tokenizer: TokenizerLike, ): """Leading assistant text plus the entire tool block arrive in one - delta. Leading content must be emitted — not silently dropped. - """ + delta. Leading content must be emitted — not silently dropped.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( f"Let me check. {TOOL_CALL_START}[{SIMPLE_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -512,8 +508,7 @@ def test_streaming_leading_block_and_trailing_in_single_delta( lfm2_tokenizer: TokenizerLike, ): """Leading text + complete tool block + trailing text in one delta. - Both leading and trailing content must be preserved. - """ + Both leading and trailing content must be preserved.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( "Let me check. " @@ -534,8 +529,7 @@ def test_echoed_tool_call_body_not_leaked_to_content( """LFM2 sometimes emits the tool call body again after the first <|tool_call_end|>, capped with a second <|tool_call_end|>. The echoed body must not surface as assistant content — neither in - streaming nor non-streaming paths. - """ + streaming nor non-streaming paths.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) body = ( "[grocery.orderIngredients(" @@ -570,8 +564,7 @@ def test_streaming_char_by_char_multi_dict_list(lfm2_tokenizer: TokenizerLike): character at a time. Every prefix lands in some partial-parse state (mid-key, mid-value, open quote inside dict, empty dict, etc.). The parser must not raise — incomplete prefixes should silently wait for - more text instead of logging exceptions. - """ + more text instead of logging exceptions.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = ( f"{TOOL_CALL_START}[grocery.orderIngredients(" @@ -602,8 +595,7 @@ def test_streaming_char_by_char_multi_dict_list(lfm2_tokenizer: TokenizerLike): def test_streaming_dotted_name_in_single_delta(lfm2_tokenizer: TokenizerLike): """A pythonic call with a dotted/attribute function name (e.g. ``domain.method(arg=...)``) must be parsed correctly in streaming mode - just as in non-streaming mode. - """ + just as in non-streaming mode.""" tool_parser: ToolParser = ToolParserManager.get_tool_parser("lfm2")(lfm2_tokenizer) full_text = f"{TOOL_CALL_START}[{DOTTED_NAME_FUNCTION_OUTPUT}]{TOOL_CALL_END}" @@ -619,8 +611,7 @@ def test_adjust_request_disables_skip_special_tokens( """When tools are present, the parser must force ``skip_special_tokens=False`` so the engine does not strip the <|tool_call_start|>/<|tool_call_end|> sentinels before they reach the - parser. - """ + parser.""" from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py index 6232f10e4353..77f99a47ec38 100644 --- a/tests/tool_use/test_muse_glimmer.py +++ b/tests/tool_use/test_muse_glimmer.py @@ -220,8 +220,7 @@ def test_reasoning_then_parallel_calls(): def _stream(raw: str, chunk: int): """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming - parsers; return (reasoning, content, tool_calls). - """ + parsers; return (reasoning, content, tool_calls).""" reasoning, content, toolcalls = [], [], [] prev = "" i = 0 diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index bead3c6f3279..f37a3c9681f9 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -376,8 +376,7 @@ def test_streaming_output_valid_with_trailing_extra_data(): class TestNonFunctionToolsSkipped: """Non-function tools (web_search, etc.) must be silently skipped - by the tool-schema utilities instead of raising TypeError. - """ + by the tool-schema utilities instead of raising TypeError.""" def test_find_tool_properties_skips_web_search(self): tools = [WEB_SEARCH_TOOL, FUNCTION_TOOL] diff --git a/tests/transformers_utils/processors/test_glm5next.py b/tests/transformers_utils/processors/test_glm5next.py index 0e24d9853bf4..917a3f627ce6 100644 --- a/tests/transformers_utils/processors/test_glm5next.py +++ b/tests/transformers_utils/processors/test_glm5next.py @@ -330,8 +330,7 @@ def test_sample_frames_fps_interval(video_processor): def test_defaults_mirror_checkpoint_config(): """Bare instantiation matches the checkpoint's ``processor_config.json`` - token-budget style defaults; ``size`` carries no budget. - """ + token-budget style defaults; ``size`` carries no budget.""" image_processor = Glm5NextImageProcessor() assert image_processor.patch_expand_factor == 1 assert image_processor.min_image_tokens == 16 @@ -365,8 +364,7 @@ def test_missing_token_budgets_rejected(): def test_video_config_fields_land(): """fps_interval / max_frame_count_dynamic from the dedicated config - shape sampling without any request overrides. - """ + shape sampling without any request overrides.""" proc = Glm5NextVideoProcessor( fps_interval=4, max_frame_count_dynamic=32, diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 18b33147f69d..fb7a7145ffb0 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -188,8 +188,7 @@ def test_model_config_generation_fallback_forwards_code_revision(): def test_safetensors_metadata_of_repo_without_safetensors(): """A repo storing its weights in another format is an answer, not a failure, - so it must not be retried. - """ + so it must not be retried.""" from huggingface_hub.errors import LocalEntryNotFoundError, NotASafetensorsRepoError get_safetensors_metadata = MagicMock( @@ -228,8 +227,7 @@ def test_mrope_num_dims(section_key, mrope_section, expected_num_dims): @pytest.mark.parametrize("section_name", ["mrope_section", "xdrope_section"]) def test_mrope_num_dims_from_config_attribute(section_name): """Some configs expose the section as an attribute rather than under - `rope_parameters`. - """ + `rope_parameters`.""" config = PretrainedConfig() setattr(config, section_name, [16, 16, 16, 16]) diff --git a/tests/transformers_utils/test_hf_overrides_model_type.py b/tests/transformers_utils/test_hf_overrides_model_type.py index 242594d63196..b739b4294a6e 100644 --- a/tests/transformers_utils/test_hf_overrides_model_type.py +++ b/tests/transformers_utils/test_hf_overrides_model_type.py @@ -22,8 +22,7 @@ def test_hf_overrides_model_type_returns_correct_config_class(): """When hf_overrides sets model_type to a registered custom type whose checkpoint has a *different* model_type on disk, get_config() must return an instance of the registered config class — not the class that matches - the on-disk model_type. - """ + the on-disk model_type.""" # Register the custom config _CONFIG_REGISTRY["test_custom_model"] = _TestCustomConfig diff --git a/tests/utils.py b/tests/utils.py index 5b988d27401e..81f82d98c8a7 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -72,8 +72,7 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None: """Pre-populate the HF cache for (repo_id, filename) pairs that upstream trust_remote_code modules would otherwise fetch from third-party CDNs - (often unreachable from US-based CI). - """ + (often unreachable from US-based CI).""" if HF_HUB_OFFLINE: return for repo_id, filename in assets: @@ -893,8 +892,7 @@ def __init__( max_wait_seconds: float | None = None, ) -> None: """Store custom child process function then invoke superclass - constructor which will indirectly launch it. - """ + constructor which will indirectly launch it.""" self.child_process_fxn = child_process_fxn super().__init__( model=model, diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index efcda2dfd823..10b550c0a174 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -59,8 +59,7 @@ def test_with_env_set(monkeypatch, mode): @create_new_process_for_each_test() def test_other_threads_are_not_policed(monkeypatch): """A background thread that syncs deliberately must not be broken by the - check being armed on the thread running the decorated function. - """ + check being armed on the thread running the decorated function.""" monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -85,8 +84,7 @@ def worker(): def test_allow_on_other_thread_does_not_disarm(monkeypatch): """`gpu_sync_allowed()` on one thread must not suppress the check on another. It is scoped by ContextVar rather than torch's process-global - sync debug mode, which a previous implementation mutated. - """ + sync debug mode, which a previous implementation mutated.""" monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -115,8 +113,7 @@ def worker(): def test_suppressing_works_while_compiling(monkeypatch): """`_suppressing` wraps torch compile entry points, which run with `torch.compiler.is_compiling()` true. `gpu_sync_allowed()` deliberately - no-ops in that state, so `_suppressing` must not route through it. - """ + no-ops in that state, so `_suppressing` must not route through it.""" monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) # Emulate being inside a torch compile, as inductor passes are. @@ -148,8 +145,7 @@ def convert_like(x, like): def test_sync_debug_mode_restored_after_checked_call(monkeypatch): """The mode is armed only for the duration of a checked call. Leaving it on process-wide made every sync outside a checked region emit a - `UserWarning` whenever our handler was not the installed one. - """ + `UserWarning` whenever our handler was not the installed one.""" monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) @@ -294,8 +290,7 @@ def test_genuinely_async_transfers_pass(monkeypatch): asynchronous in both directions -- even with a dtype conversion, which runs GPU-side, and even permuted (e.g. transposed), which uses pitched cudaMemcpy2D/3DAsync; D2H `Tensor.to` allocates a pinned destination. - None of these may be flagged. - """ + None of these may be flagged.""" monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") monkeypatch.setattr(gsd, "_sync_check_enabled", True) gsd._install_copy_checkers() diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index e2935a254954..02889c20f66a 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -29,8 +29,7 @@ def test_get_file_store_init_method_is_unique(): def _call_with_timeout(func, timeout: float = 10.0): """Run func in a daemon thread so a livelock regression fails the test - quickly instead of hanging the suite. - """ + quickly instead of hanging the suite.""" result: dict = {} def target(): @@ -97,8 +96,7 @@ def test_get_open_port_skips_reserved_dp_master_ports( monkeypatch: pytest.MonkeyPatch, ): """VLLM_PORT == VLLM_DP_MASTER_PORT must not livelock by returning the - same free-but-reserved port forever (upstream issue #50024). - """ + same free-but-reserved port forever (upstream issue #50024).""" base = 42123 with monkeypatch.context() as m: m.setenv("VLLM_PORT", str(base)) @@ -111,8 +109,7 @@ def test_get_open_ports_list_skips_reserved_dp_master_ports( monkeypatch: pytest.MonkeyPatch, ): """The VLLM_PORT band scan must not hand out ports inside the window - reserved for the data parallel master process. - """ + reserved for the data parallel master process.""" base = 42223 reserved = range(base, base + 10) with monkeypatch.context() as m: @@ -128,8 +125,7 @@ def test_get_open_port_ephemeral_skips_reserved_range( monkeypatch: pytest.MonkeyPatch, ): """An ephemeral port that happens to fall in the reserved range is - replaced by a rescan starting past the range. - """ + replaced by a rescan starting past the range.""" base = 42323 with monkeypatch.context() as m: m.delenv("VLLM_PORT", raising=False) @@ -148,8 +144,7 @@ def test_get_open_port_ephemeral_without_dp_master_port( monkeypatch: pytest.MonkeyPatch, ): """Without VLLM_PORT and VLLM_DP_MASTER_PORT, an ephemeral port is - returned as before. - """ + returned as before.""" with monkeypatch.context() as m: m.delenv("VLLM_PORT", raising=False) m.delenv("VLLM_DP_MASTER_PORT", raising=False) diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 845fdda2a122..9f718703a7c1 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -157,8 +157,7 @@ def test_pct_binding_filters_cpus(monkeypatch): ) def test_pct_binding_fires_on_every_capable_sku(monkeypatch, sku, expected_cpus): """Each SKU in ``_PCT_CAPABLE_SKUS`` engages the gate at its own - expected ``highest_perf`` and uses its own priority-core stride. - """ + expected ``highest_perf`` and uses its own priority-core stride.""" sku_config = numa_utils._PCT_CAPABLE_SKUS[sku] _patch_pct_gates( monkeypatch, @@ -243,8 +242,7 @@ def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes(): def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch): """PCT: EngineCore for a multi-NUMA shard binds to the union of priority cores across all shard nodes, so worker `--physcpubind` is always a - subset of EngineCore's `cpus_allowed`. - """ + subset of EngineCore's `cpus_allowed`.""" _patch_pct_gates( monkeypatch, model_match=True, @@ -294,8 +292,7 @@ def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes( """external_launcher (or multi-node-within-DP, or Ray) hits the fallback branch. EngineCore must still span every local NUMA node so it can mp-spawn its local workers without ``--physcpubind`` - strict-validation failures. - """ + strict-validation failures.""" _patch_pct_gates( monkeypatch, model_match=True, @@ -319,8 +316,7 @@ def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes( def test_get_numactl_args_engine_core_baseline_multi_node_within_dp_spans_locals(): """Multi-node-within-DP fallback: bind EngineCore to all local NUMA - nodes that the visible ``numa_bind_nodes`` reference. - """ + nodes that the visible ``numa_bind_nodes`` reference.""" vllm_config = _make_config( numa_bind=True, numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1], @@ -343,8 +339,7 @@ def test_get_numactl_args_engine_core_skips_user_cpu_list(monkeypatch): ``--physcpubind`` spawns need. We fall back to ``--cpunodebind`` over the shard's NUMA nodes instead. PCT auto-detect is also bypassed when the user is explicit (its priority-core union may not be a superset - of the user's per-worker cores). - """ + of the user's per-worker cores).""" _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) vllm_config = _make_config( numa_bind=True, @@ -397,8 +392,7 @@ def test_get_numactl_args_requires_detectable_nodes(monkeypatch): def test_configure_subprocess_rejects_unknown_process_kind(): """configure_subprocess only knows 'worker' and 'EngineCore'; anything else must raise ValueError instead of silently routing to the worker - path. - """ + path.""" vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0]) with ( pytest.raises(ValueError, match="process_kind"), diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index af9ed0c6b0d0..66a1c5c08836 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -1392,8 +1392,7 @@ def test_non_causal_backend_correctness( model: str, ): """Test backend's correctness with non-causal (bidirectional) decoder - attention, as used by DFlash speculative decoding. - """ + attention, as used by DFlash speculative decoding.""" def bidirectional_mask_mod( b: torch.Tensor, diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 0481ebc79094..8486d216a125 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -115,8 +115,7 @@ def test_mamba_layers_have_unified_interface( layer_class, expected_backend, expected_mamba_type ): """Test that all Mamba layers have the unified get_attn_backend - interface. - """ + interface.""" assert hasattr(layer_class, "get_attn_backend"), ( f"{layer_class.__name__} should have get_attn_backend method" ) diff --git a/tests/v1/attention/test_attention_splitting.py b/tests/v1/attention/test_attention_splitting.py index ff0e06514e86..bad0f61b7c45 100644 --- a/tests/v1/attention/test_attention_splitting.py +++ b/tests/v1/attention/test_attention_splitting.py @@ -164,8 +164,7 @@ def apply_split_decodes_and_prefills( treat_short_extends_as_decodes: bool = True, ): """Helper function to apply split_decodes_and_prefills and return - the results. - """ + the results.""" device = torch.device("cpu") seq_lens = [10 * (i + 1) for i in range(len(query_lens))] common_metadata = create_common_attn_metadata( diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py index 34621900d5a1..cff1228baf03 100644 --- a/tests/v1/attention/test_dcp_a2a_pack_mask.py +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The empty-shard mask fused into the DCP A2A pack kernel must be bit-exact -with the eager ``mask_dcp_empty_shards_`` pass it replaces. -""" +with the eager ``mask_dcp_empty_shards_`` pass it replaces.""" import pytest import torch diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 6f61a6aeed86..d5c794131474 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -495,8 +495,7 @@ def test_dspark_noncausal_sparse_mla_matches_sdpa( num_heads, ): """Non-causal (window ∪ block, future-pointing) per-token indices must match - a dense SDPA reference over the same indices, for both sparse-MLA backends. - """ + a dense SDPA reference over the same indices, for both sparse-MLA backends.""" _skip_if_backend_unavailable(backend_cls, kv_cache_dtype, block_size) window, block, topk_width = _DSPARK_CONFIGS[config_name] diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index c712cb5844b0..8dc082b08bbc 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -231,8 +231,7 @@ def test_plan_uses_state_params(monkeypatch): def test_kv_lens_host_formula(): """Per-row host lengths: context == position + 1; capped at - index_topk + trailing-pool remainder past the sparse threshold. - """ + index_topk + trailing-pool remainder past the sparse threshold.""" builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index c8032d9eaade..221f933d8942 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -187,8 +187,7 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase): def test_has_initial_state_after_reclassification(): """After reclassification, num_prefills > 0 so the prefill kernel path should compute has_initial_state. For the reclassified request with - context_lens > 0, the corresponding entry must be True. - """ + context_lens > 0, the corresponding entry must be True.""" builder = _create_gdn_builder(num_speculative_tokens=2) batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3]) meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2]) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 4175bb6ff5e1..8bc409adb704 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -56,8 +56,7 @@ def _ref_stable_topk_from_candidates_fp64( ) -> torch.Tensor: """Pure-PyTorch reference for the CuteDSL stable-topk selector order (score desc, then lowest global token id). Selects the same SET as the - kernel; only the set is compared in tests. - """ + kernel; only the set is compared in tests.""" num_rows, num_candidates = candidate_scores.shape device = candidate_scores.device select_k = min(k, num_candidates) @@ -122,8 +121,7 @@ def _dcp_lse_merge( class _FakeDCPGroup: """Single-process stand-in: ``all_gather`` returns the pre-built concatenation of every rank's packed ``(score, global_id)`` candidates, - mirroring the one packed all-gather the merge issues. - """ + mirroring the one packed all-gather the merge issues.""" def __init__(self, gathered_packed: torch.Tensor) -> None: self.gathered_packed = gathered_packed @@ -209,8 +207,7 @@ def _merge_local_topks_global_with_fake_dcp( """Run ``_merge_dcp_topk_global`` for every rank against a faked all-gather and return each rank's global-index result (all ranks should agree). The fake pre-builds the packed candidate concatenation the merge's - single ``all_gather`` would return. - """ + single ``all_gather`` would return.""" # The merge is now CuteDSL-only (no PyTorch fallback), so it runs the real # Triton pack + CuteDSL selector kernels even behind the faked all-gather. if not current_platform.is_cuda() or not has_cutedsl(): @@ -712,8 +709,7 @@ def test_dcp_filter_compaction_matches_reference( """In-kernel compaction must, for every row, produce exactly the rank-owned physical slots packed into [0, valid_count) with -1 in the tail -- the same SET a reference filter + sort/gather produces. Uses wide rows (> BLOCK_N - valid slots), with interior -1 gaps. - """ + valid slots), with interior -1 gaps.""" device = torch.device("cuda") torch.manual_seed(7) dcp_size = 2 @@ -871,8 +867,7 @@ def test_decode_topk_pads_surplus_with_negative_one(): so a non-(-1) pad would be silently attended. This is the common case under DCP, where each rank's local seq_len = global / world is usually << topk. Checked on the *set* of valid indices (the kernel may order them by score, - not position). - """ + not position).""" torch.manual_seed(0) device = torch.device("cuda") topk = 8 @@ -918,8 +913,7 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp(): """End-to-end DCP decode where the global seq_len < topk (so every rank's local top-k is surplus-padded). Exercises the kernel surplus -> merge mask -> global top-k -> physical localize -> LSE merge chain for the common - short-context decode case, vs the non-DCP reference. - """ + short-context decode case, vs the non-DCP reference.""" torch.manual_seed(4) device = torch.device("cuda") world = 2 diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 6377d7df4661..0966f9b20a6b 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -58,8 +58,7 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the request's single KpoolTailManager block, the remaining columns are never - written and stay zero. - """ + written and stay zero.""" bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) return bt @@ -67,8 +66,7 @@ def make_tail_block_table(own_blocks, width=64): def legacy_generic_tail_slots(block_table, query_start_loc, positions): """Reference of the generic ``_compute_slot_mappings_kernel`` arithmetic - (block_table.py:305-313) applied to the tail group's table. - """ + (block_table.py:305-313) applied to the tail group's table.""" slots = [] for req in range(block_table.shape[0]): for i in range(query_start_loc[req], query_start_loc[req + 1]): @@ -109,8 +107,7 @@ def make_batch(per_req_positions, padded_len=None): def test_legacy_generic_mapping_collapses_onto_block_zero(): """The bug: with the manager's 1-column block table, the generic kernel - maps every pos >= kpool onto tail block 0, and distinct requests collide. - """ + maps every pos >= kpool onto tail block 0, and distinct requests collide.""" own_blocks = [5, 9] per_req = [list(range(10)), list(range(12))] # prompts of len 10 and 12 positions, qsl, _, num_actual, num_reqs = make_batch(per_req) @@ -140,8 +137,7 @@ def test_legacy_generic_mapping_collapses_onto_block_zero(): def test_circular_mapping_isolates_requests(): """The fix: every token lands in its own request's block at pos % kpool, and no slot is ever shared by two different requests (slots do recur - within a request every kpool positions -- that is the circular design). - """ + within a request every kpool positions -- that is the circular design).""" own_blocks = [5, 9] per_req = [list(range(10)), list(range(12))] positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) @@ -169,8 +165,7 @@ def test_circular_mapping_isolates_requests(): def test_circular_mapping_matches_generic_for_short_requests(prompt_len): """For pos < kpool the generic kernel already picks the own block, so the two mappings agree while every position fits the request's first block - (single-request behavior is unchanged). - """ + (single-request behavior is unchanged).""" own_blocks = [7] per_req = [list(range(prompt_len))] positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) @@ -280,8 +275,7 @@ class TailRingMirror: """Mirror of _kpool_tail_seed_kernel / _kpool_decode_update_batched_kernel addressing: block = tail_slot // kpool, ring offset = pos % kpool; a pool completing at pos reads ring slots (pool_start + s) % kpool and uses the - current token's own K/score for the last member. - """ + current token's own K/score for the last member.""" def __init__(self, num_blocks, kpool=KPOOL): self.kpool = kpool @@ -324,8 +318,7 @@ def tail_slot_for(mapping, req, pos, own_block): def run_scenario(mapping, interleave): """Requests A (block 5, prompt len 9) and B (block 9, prompt len 11) - decode concurrently; returns A's boundary pool [8, 9, 10, 11]. - """ + decode concurrently; returns A's boundary pool [8, 9, 10, 11].""" ring = TailRingMirror(num_blocks=16) blocks = {"A": 5, "B": 9} prompts = {"A": 9, "B": 11} diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index f281213221bc..f4451fe7af82 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -28,8 +28,7 @@ def _make_vllm_config( block_size: int | None = None, ): """Create a minimal mock VllmConfig with only the fields the builder - accesses, avoiding any model download / HF config inspection. - """ + accesses, avoiding any model download / HF config inspection.""" speculative_config = ( SimpleNamespace( num_speculative_tokens=num_speculative_tokens, @@ -75,8 +74,7 @@ def test_mamba_single_token_prompt_runs_as_prefill(): def test_update_block_table_copies_block_idx_to_persistent_buffers(): """update_block_table() must write block_idx tensors to the current builder's persistent buffers, not leave them pointing to a different - builder's buffers. - """ + builder's buffers.""" block_size = 16 max_model_len = 256 num_reqs = 4 @@ -180,8 +178,7 @@ def test_state_indices_tensor_d_includes_num_speculative_blocks(): bug 1: with mamba_cache_mode='all' and speculative decoding enabled, the cudagraph buffer for state_indices_tensor_d must allocate the same per-request column count as the runtime block table, which includes - num_speculative_blocks trailing scratch columns. - """ + num_speculative_blocks trailing scratch columns.""" block_size = 16 max_model_len = 256 max_num_seqs = 4 @@ -214,8 +211,7 @@ def test_block_idx_cudagraph_capture_padded_by_num_reqs(): bug 2: with mamba_cache_mode='all' and spec decode, _update_metadata_for _cudagraph_capture must slice block_idx_last_{scheduled,computed}_token by the request count (padded_bs == num_reqs), not by num_decode_tokens. - Past num_decodes, the slice must be zero-filled. - """ + Past num_decodes, the slice must be zero-filled.""" block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -295,8 +291,7 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): block_idx_last_scheduled_token_prev_step as a persistent buffer with the same shape as the existing block_idx_last_{scheduled,computed}_token buffers, so cudagraph capture records a stable pointer for the prev-step - input anchor consumed by mamba_mixer2's input gather. - """ + input anchor consumed by mamba_mixer2's input gather.""" block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -324,8 +319,7 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode(): """Without spec decode, the prev-step buffer is unused and must not be - allocated — the input anchor reduces to last_computed_token. - """ + allocated — the input anchor reduces to last_computed_token.""" block_size = 16 max_model_len = 256 max_num_seqs = 8 @@ -349,8 +343,7 @@ def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer(): """_update_metadata_for_cudagraph_capture must copy the prev-step anchor into the builder's persistent buffer (so cudagraph replay reads from the same underlying memory), pad past num_decodes with zero, and return a - slice of the persistent buffer in the metadata. - """ + slice of the persistent buffer in the metadata.""" block_size = 16 max_model_len = 256 max_num_seqs = 8 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 5a8fdea20464..597ffb03fcf7 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -171,8 +171,7 @@ def test_glm5_flashinfer_masked_mha_routing( def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim): """The K concat used by the MLA prefill context loop must equal torch.cat of k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself - instead of allocating and copying. - """ + instead of allocating and copying.""" torch.manual_seed(0) num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256 k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16) diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py index 211cde663bc3..220f6a198dba 100644 --- a/tests/v1/attention/test_mla_prefill_quant_output.py +++ b/tests/v1/attention/test_mla_prefill_quant_output.py @@ -134,8 +134,7 @@ def test_flash_attn_registers_warmup_only_for_fa4( def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): """run_prefill_new_tokens must accept out/output_scale so the direct - (non-**kwargs) call in forward_mha type- and runtime-checks. - """ + (non-**kwargs) call in forward_mha type- and runtime-checks.""" import inspect params = inspect.signature( @@ -151,8 +150,7 @@ def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): def test_mla_impl_forward_mha_accepts_output_scale(): """The abstract MLA impl forward_mha must carry output_scale so every - override (and the unconditional forward_impl call) stays compatible. - """ + override (and the unconditional forward_impl call) stays compatible.""" import inspect from vllm.v1.attention.backend import MLAAttentionImpl @@ -179,8 +177,7 @@ def _fused_fp8_skip_reason() -> str | None: def test_fa4_fused_fp8_output_matches_post_quant(default_vllm_config): """FA4's fused FP8 write (output_scale, flash-attention#135) must match the bf16-attention + standalone static-FP8-quant path it replaces, since - production uses the same output_scale for both. - """ + production uses the same output_scale for both.""" from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index fa52b1484b71..098eddfbf6eb 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -315,8 +315,7 @@ def test_backend_supported_dimension_validation(self): def test_flash_attn_accepts_glm53_flash_nope_dimensions(self): """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a - RoPE-carrying 320-wide query does not. - """ + RoPE-carrying 320-wide query does not.""" try: from vllm.v1.attention.backends.mla.prefill.flash_attn import ( FlashAttnPrefillBackend, diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 351fe0006897..89cb0f566467 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -257,8 +257,7 @@ def test_replayssm_write_pos(case: ReplaySSMBuildCase): def test_resumed_request_differs_from_fresh(): """Same token count, different decode_base: fresh (base 100) -> write_pos 5, - resumed (base 105) -> write_pos 0. - """ + resumed (base 105) -> write_pos 0.""" builder = _create_replayssm_builder(16) batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1]) common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 01018a630c86..2342a279a397 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -517,8 +517,7 @@ def test_auto_selection_for_kv_connector( aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version ): """Auto-selection with a KV connector and AITER enabled resolves to unified attn, - and to triton attn if AITER not enabled. - """ + and to triton attn if AITER not enabled.""" from vllm.platforms.rocm import RocmPlatform with patch( diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py index 9cbb7a6f8f45..d87a82c37b2b 100644 --- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py +++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py @@ -102,8 +102,7 @@ def expected_row_seq_lens(per_req_positions, batch_size, next_n): def test_legacy_flat_layout_misaligns_and_bleeds(): """The bug: ``positions[:n] + 1`` with padded-row coordinates reads other - requests' positions and prefill positions. - """ + requests' positions and prefill positions.""" decode_lens, per_req, positions = make_non_uniform_batch() batch_size = decode_lens.shape[0] next_n = int(decode_lens.max()) @@ -138,8 +137,7 @@ def test_helper_uniform_layout_matches_flat_slice(): def test_helper_padded_layout_per_row(): """Non-uniform batches map every padded row to its own token's position; - pad rows get 0 (empty tail). - """ + pad rows get 0 (empty tail).""" decode_lens, per_req, positions = make_non_uniform_batch() out = _decode_topk_seq_lens(positions, decode_lens, 8, 3, 4, requires_padding=True) expected = expected_row_seq_lens(per_req, 3, 4) @@ -151,8 +149,7 @@ def test_helper_padded_layout_per_row(): def expand_tail_region(dec_seq): """Run the production pure-torch expand + append pair (the fused kernel replicates it exactly on the identity path) and return the tail columns - [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1). - """ + [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1).""" rows = dec_seq.shape[0] pool_ids = torch.arange(SELECT_K, dtype=torch.int64).expand(rows, SELECT_K) valid = torch.ones_like(pool_ids, dtype=torch.bool) @@ -166,8 +163,7 @@ def expand_tail_region(dec_seq): def test_tail_expansion_legacy_vs_fixed(): """End-to-end tail consequence: the fixed mapping appends exactly the request's trailing incomplete pool; the legacy flat slice drops real tail - tokens and emits indices far past the request's own sequence. - """ + tokens and emits indices far past the request's own sequence.""" decode_lens, per_req, positions = make_non_uniform_batch() n = 12 fixed = _decode_topk_seq_lens( diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index ecaf71d55b67..8ff52ecb386c 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -909,8 +909,7 @@ def test_index_group_convert_during_piecewise_capture(): runs in the following eager break. The convert's side stream must not wait on ``logical_topk_ready`` there: an event recorded inside a captured segment is graph-local, and an eager wait on it raises - cudaErrorInvalidValue. - """ + cudaErrorInvalidValue.""" device = torch.device(DEVICE_TYPE) num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16 @@ -1049,8 +1048,7 @@ def test_triton_convert_rejects_req_id_longer_than_token_indices(): req_id but the output is allocated like token_indices, so a full-batch req_id combined with an MQA-subset token_indices wrote past the end of the output buffer. The wrapper must reject the length mismatch instead - of corrupting memory. - """ + of corrupting memory.""" device = torch.device(DEVICE_TYPE) num_topk_tokens = 128 block_size = 64 @@ -1105,8 +1103,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): is active, forward_mqa only receives the leading decode tokens, but _forward_bf16_kv passed the full-batch req_id_per_token to the index conversion, making it write past the end of its output buffer. The call - site must slice req_id_per_token to the MQA tokens. - """ + site must slice req_id_per_token to the MQA tokens.""" device = torch.device(DEVICE_TYPE) num_topk_tokens = 128 block_size = 64 @@ -1204,8 +1201,7 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected): ) def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expected): """A 32K prefill needs the default workspace exactly; shrinking it would - push a supported request onto MQA. - """ + push a supported request onto MQA.""" assert ( _masked_mha_workspace_fits( batch_size=1, @@ -1263,8 +1259,7 @@ def test_is_masked_mha_available_model_dims( ): """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2, GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an - unquantized KV cache, nothing else. - """ + unquantized KV cache, nothing else.""" import vllm.model_executor.layers.attention.sparse_mla_attention as mod monkeypatch.setattr( @@ -3604,8 +3599,7 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of the cross-rank LSE merge, or a NaN would survive the merge even at zero - weight (0 * NaN = NaN). - """ + weight (0 * NaN = NaN).""" num_tokens, num_heads, head_dim = 3, 2, 3 q = torch.empty(num_tokens, num_heads, head_dim, device=DEVICE_TYPE) local_indices = torch.tensor( @@ -3778,8 +3772,7 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): def test_sparse_impl_observes_repointed_indexer_buffer(): """The MTP proposer repoints the draft's indexer at the target model's buffer after the backend impl is built, so the impl must resolve the buffer per read. - Snapshotting it in __init__ leaves the layer reading indices nothing writes. - """ + Snapshotting it in __init__ leaves the layer reading indices nothing writes.""" impl = object.__new__(FlashInferMLASparseImpl) own = torch.zeros(4, 8, dtype=torch.int32) target = torch.ones(4, 8, dtype=torch.int32) @@ -3795,8 +3788,7 @@ def test_sparse_impl_observes_repointed_indexer_buffer(): def test_explicit_topk_buffer_supersedes_indexer(): """Backbone skip-topk layers have no indexer, and the proposer also assigns the - shared buffer directly onto draft submodules. - """ + shared buffer directly onto draft submodules.""" impl = object.__new__(FlashInferMLASparseImpl) indexer = SimpleNamespace( topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32), @@ -3811,8 +3803,7 @@ def test_explicit_topk_buffer_supersedes_indexer(): def test_sparse_mla_common_impl_resolves_buffer_lazily(): """Guards the whole SparseMLACommonImpl family at once: a backend that - snapshots the buffer instead silently loses MTP buffer sharing. - """ + snapshots the buffer instead silently loses MTP buffer sharing.""" assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer) assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( "topk_indices_buffer must stay a lazily-resolved property" diff --git a/tests/v1/attention/test_trtllm_attention_integration.py b/tests/v1/attention/test_trtllm_attention_integration.py index 36c2a710d532..1a49a3895c52 100644 --- a/tests/v1/attention/test_trtllm_attention_integration.py +++ b/tests/v1/attention/test_trtllm_attention_integration.py @@ -198,8 +198,7 @@ def _create_nvfp4_hnd_kv_cache( def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL): """Run TRTLLM attention through the full FlashInfer pipeline - and compare against an SDPA reference. - """ + and compare against an SDPA reference.""" set_random_seed(42) device = torch.device(f"{DEVICE_TYPE}:0") @@ -427,8 +426,7 @@ def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len): def test_trtllm_gen_full_attention_integration(batch_spec_name: str): """Test TRTLLM gen-full attention through the full FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline, - with real TRTLLM kernels on Blackwell. - """ + with real TRTLLM kernels on Blackwell.""" _run_trtllm_integration(BATCH_SPECS[batch_spec_name]) @@ -439,8 +437,7 @@ def test_trtllm_gen_full_attention_integration(batch_spec_name: str): @torch.inference_mode() def test_trtllm_gen_nvfp4_kv_integration(batch_spec_name: str): """Test TRTLLM attention with nvfp4 KV cache through the full - FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline. - """ + FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline.""" _run_trtllm_integration( BATCH_SPECS[batch_spec_name], kv_cache_dtype="nvfp4", diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index eeb977adef65..393ec484651e 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Fine-grained partial prefix-cache hits for hybrid (full attention + mamba "align") models: scheduler chunk splitting, partial tail registration, CoW -on partial hits, and same-step deferral. -""" +on partial hits, and same-step deferral.""" from math import lcm from types import SimpleNamespace @@ -175,8 +174,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: - 0 -> 8192 -> 9728 -> 9984 -> 10000. - """ + 0 -> 8192 -> 9728 -> 9984 -> 10000.""" block_size = 512 scheduler_block_size = block_size * dcp_world_size hash_block_size = 32 @@ -744,8 +742,7 @@ def test_hash_aligned_query_end_uses_regular_partial_tail(): def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): """An external mid-block hit must become a running request even when its - first continuation does not need another Mamba block. - """ + first continuation does not need another Mamba block.""" hash_block_size = 2 mamba_block_size = 4 * hash_block_size kv_cache_config = KVCacheConfig( @@ -817,8 +814,7 @@ def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): def test_boundary_state_offloads_returns_cow_target(): """Boundary hand-offs expose aligned snapshots and the partial-tail CoW - target, never the overwritten CoW source. - """ + target, never the overwritten CoW source.""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1004,8 +1000,7 @@ def request_finished_all_groups(self, request, block_ids): def test_boundary_state_offload_dropped_when_request_freed_before_drain(): """A hand-off recorded in the same scheduling pass as the request's death must not be drained: its release hook has already run, so draining would - leak a pinned block. - """ + leak a pinned block.""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1052,8 +1047,7 @@ def test_boundary_state_offload_dropped_when_request_freed_before_drain(): def test_boundary_state_offloads_block_aligned_prompt(): """A prompt ending on a block boundary registers no CoW partial tail; its - boundary state block is handed off as a snapshot instead (once). - """ + boundary state block is handed off as a snapshot instead (once).""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1107,8 +1101,7 @@ def test_boundary_state_offloads_block_aligned_prompt(): def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): """truncate_computed_blocks slices each group by its own block size, keeps null placeholders in the retained prefix, and leaves the original - lookup result untouched (pure view, no refcount changes). - """ + lookup result untouched (pure view, no refcount changes).""" hash_block_size = 2 kv_cache_config = KVCacheConfig( num_blocks=24, @@ -1163,8 +1156,7 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): def test_truncate_computed_blocks_allows_short_mamba_group_only(): """External state may replace a short Mamba hit, but other groups must - cover the aligned local endpoint. - """ + cover the aligned local endpoint.""" hash_block_size = 2 kv_cache_config = KVCacheConfig( num_blocks=24, @@ -1306,8 +1298,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): """The owner's move re-arms the same-step guard: the moved entry is filled by this step's copy, and chained same-step copies read stale - sources, so a request hitting it in the move step must be deferred. - """ + sources, so a request hitting it in the move step must be deferred.""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1597,8 +1588,7 @@ def test_hybrid_partial_hash_truncates_full_attention_hit_length(): def test_cow_retained_blocks_returned_for_release(): """new_step_starts returns the CoW copy retentions instead of freeing - them; the scheduler owns releasing them once the copy has run. - """ + them; the scheduler owns releasing them once the copy has run.""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -1651,8 +1641,7 @@ def test_cow_retained_blocks_returned_for_release(): def test_free_cow_retained_blocks_defers_until_copy_step_processed(): """Scheduler releases CoW retentions immediately when the copy's step has - been processed (or deferral is off), and defers them otherwise. - """ + been processed (or deferral is off), and defers them otherwise.""" from collections import deque freed: list = [] @@ -1689,8 +1678,7 @@ def test_free_cow_retained_blocks_defers_until_copy_step_processed(): def test_full_attention_eagle_drops_one_hash_unit(): """With fine-grained partial hits, eagle rewinds the hit by one hash unit instead of a whole cache block: the tail block's KV is append-only, so it - still covers the reduced length and stays in the hit as a partial block. - """ + still covers the reduced length and stays in the hit as a partial block.""" from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager @@ -1766,8 +1754,7 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks(): """Regression: with eagle, the mamba group must not receive the eagle lookup margin — its finder never applies the drop, so it could return a hit past the blocks the (dropped) full-attention group covers, crashing - the consumer's CoW with block_idx >= len(req_blocks). - """ + the consumer's CoW with block_idx >= len(req_blocks).""" hash_block_size = 2 block_size = 2 * hash_block_size kv_cache_config = KVCacheConfig( @@ -2369,8 +2356,7 @@ def test_retention_snapshots_handed_off_with_exact_block_ids(): """Under sparse retention, each retained mamba boundary state block is handed to the connector with its exact block id. Connectors must not resolve align-mode state blocks positionally because the block table is not - append-only. - """ + append-only.""" hash_block_size = 2 block_size = 4 manager = make_kv_cache_manager( @@ -2411,8 +2397,7 @@ def test_snapshot_handoff_dense_default_retention(): """With dense (default) retention, every materialized mamba boundary state block is handed off — regular mamba-align + prefix-match-unit deployments get store-able boundary snapshots without setting a - prefix-cache retention interval. - """ + prefix-cache retention interval.""" hash_block_size = 2 block_size = 4 manager = make_kv_cache_manager( @@ -2444,8 +2429,7 @@ def test_snapshot_handoff_dense_default_retention(): def _run_chunked_prefill(manager, req, chunk_size, num_chunks): """Prefill ``req`` one ``chunk_size`` chunk per scheduler step, yielding the - boundary-state hand-offs offered (and claimed) at each step. - """ + boundary-state hand-offs offered (and claimed) at each step.""" computed_blocks, num_computed, _ = manager.get_computed_blocks(req) for step in range(num_chunks): manager.new_step_starts() @@ -2509,8 +2493,7 @@ def test_boundary_states_offered_past_prompt_for_resumed_prefill(): ``num_prompt_tokens``. A resumed request re-prefills its generated tokens and every group re-saves them, so filtering on the original prompt length would silently strip the mamba key for boundaries full attention still - stores; only the connector knows where its save window ends. - """ + stores; only the connector knows where its save window ends.""" block_size = 8 prompt_len = block_size manager = make_kv_cache_manager( diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 097878ff7d50..7e27a4517e03 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -552,8 +552,7 @@ def _create_async_pp_scheduler( def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None: """Delivered tokens must be an order-preserving subset of the emitted - tokens with no duplicates (tokens are globally unique). - """ + tokens with no duplicates (tokens are globally unique).""" it = iter(emitted) for token in delivered: assert token in it, f"token {token} delivered out of order or twice" @@ -563,8 +562,7 @@ def _assert_positions_consistent(req, engine: PipelinedEngine) -> None: """The i-th delivered output token must be one the runner sampled for exactly sequence position prompt_len + i: catches a preempted request's stale output landing on a position the resumed request resampled (or - vice versa), which token-stream equality alone cannot see. - """ + vice versa), which token-stream equality alone cannot see.""" for i, token in enumerate(req.output_token_ids): expected = req.num_prompt_tokens + i actual = engine.emitted_position[token] diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index f3692672e970..2fcf70c6efa7 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -178,8 +178,7 @@ def _shared_layout_config(): class TestCSALinearGrouping: """A CSA + linear-attention model (sparse attention with a compressor ring, plus sharded GDN and one TP-replicated PLE state) goes through the generic - packed-group path; no model-specific branch is involved. - """ + packed-group path; no model-specific branch is involved.""" @staticmethod def _mamba_groups(groups): @@ -194,8 +193,7 @@ def _mamba_groups(groups): def test_replicated_state_gets_its_own_group(self): """The PLE state is TP-replicated and a different size, so it must not share a manager group with the sharded GDN states (the NIXL worker - requires exactly one single-layer replicated group). - """ + requires exactly one single-layer replicated group).""" groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs()) mamba_groups = self._mamba_groups(groups) @@ -248,8 +246,7 @@ def test_every_group_fits_one_packed_block(self): def test_scratch_group_survives_computed_block_truncation(self): """The scratch group contributes no computed blocks, so truncating a lookup result must skip it: its block size is the ring capacity, which - neither divides the hit length nor bounds the (empty) block list. - """ + neither divides the hit length nor bounds the (empty) block list.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True groups = get_kv_cache_groups(config, _make_csa_linear_specs()) @@ -363,8 +360,7 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid per page size, which is emitted whole, or a single-layer state that cannot be split at all -- the mamba layers must be sized against it. Sizing them against a narrower bucket splits them past what the block - already fits, spending a pool block per extra group for no saving. - """ + already fits, spending a pool block per extra group for no saving.""" config = _mock_vllm_config("BLNHC") config.speculative_config = None specs = {} @@ -409,8 +405,7 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would set the block width and the MLA group would fill under a third of every block; capped, it splits into groups no wider than the MLA - page. - """ + page.""" config = _shared_layout_config() specs: dict[str, KVCacheSpec] = {} for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)): @@ -611,8 +606,7 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring claims one block per request, never hashes, and does not disturb the - scheduler/hash block sizes the paged groups imply. - """ + scheduler/hash block sizes the paged groups imply.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True specs: dict[str, KVCacheSpec] = { diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index 7f0747abe0c7..3948f7c6e300 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -367,8 +367,7 @@ def test_free_request_with_duplicate_mm_hashes(): """Freeing a request whose two inputs share the same mm_hash must fully clean up request_cached_ids. After the first free_encoder_input call, cached[mm_hash] becomes empty; the second call must still remove the - remaining input_id from request_cached_ids. - """ + remaining input_id from request_cached_ids.""" manager = EncoderCacheManager(cache_size=20) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) @@ -384,8 +383,7 @@ def test_free_request_with_duplicate_mm_hashes(): def test_duplicate_mm_hash_stays_referenced_until_last_free(): """`cached` holds one reference per request, not per position, so freeing - the first of two occurrences must not release the entry. - """ + the first of two occurrences must not release the entry.""" manager = EncoderCacheManager(cache_size=20) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) @@ -406,8 +404,7 @@ def test_duplicate_mm_hash_stays_referenced_until_last_free(): def test_duplicate_mm_hash_is_not_evicted_before_last_use(): """An item repeated within one request (an image carried across conversation turns) must not become reclaimable capacity while a later - occurrence still has to be spliced in, or the encoder recomputes it. - """ + occurrence still has to be spliced in, or the encoder recomputes it.""" manager = EncoderCacheManager(cache_size=8) req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) other = MockRequest("r2", ["imgB"], [8]) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index a0600abd8510..3c82180474f6 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2593,8 +2593,7 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): """The kpool tail parasitizes the indexer tensors instead of getting its own: sibling idx/tail tensors paired by layer order, zero standalone tail bytes, one shared block per request, and no prefix-caching leakage from - the kpool-sized scratch group. - """ + the kpool-sized scratch group.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2665,8 +2664,7 @@ def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): def test_glm5_kpool_tail_does_not_drag_hash_block_size(): """The tail's kpool-sized scratch block (4 tokens) must not constrain the - prefix-cache hash granularity: participating groups alone decide it. - """ + prefix-cache hash granularity: participating groups alone decide it.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2742,8 +2740,7 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer(): def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba(): """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not - disable slot sharing; the layout re-pads them to the MLA page. - """ + disable slot sharing; the layout re-pads them to the MLA page.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2774,8 +2771,7 @@ def prepadded_mamba_spec(): def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection(): """Round-robin mamba grouping keeps practical PP splits balanced: every stage's largest projected mamba group slice fits its projected MLA - layers, so per-stage slot tensors all have an MLA owner. - """ + layers, so per-stage slot tensors all have an MLA owner.""" model_config = ModelConfig(max_model_len=8192) vllm_config = VllmConfig(model_config=model_config) @@ -2819,8 +2815,7 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatc """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA layers but a round-robin slice of 3 under the minimum 4 mamba groups; grouping bumps to 5 groups so every stage's projection keeps sharing - on instead of silently falling back on stage 0. - """ + on instead of silently falling back on stage 0.""" from vllm.config import ParallelConfig from vllm.distributed.utils import get_pp_indices @@ -2869,8 +2864,7 @@ def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch): def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag(): """Hybrid slot sharing must take precedence over the experimental enable_cross_layers_blocks packed layout: generic packing would give the - MLA slots a strided view, breaking the contiguous virtual split. - """ + MLA slots a strided view, breaking the contiguous virtual split.""" model_config = ModelConfig(max_model_len=8192) kv_transfer_config = KVTransferConfig( kv_connector="NixlConnector", @@ -3172,8 +3166,7 @@ def test_multi_run_layer_compact_strides_place_hoisted_heads(): """A layer-compact run region is its own dense allocation: under LHBNC the head groups sit between the layers and the blocks, so a run's block stride is one head group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models). - """ + on ROCm hybrid models).""" full = new_kv_cache_spec() swa = new_sliding_window_spec(sliding_window=full.block_size * 2) page = full.page_size_bytes @@ -3252,8 +3245,7 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): def test_hidden_states_with_tp_scales_page_size(): """When TP shrinks KV pages below the hidden-state per-token cost, get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states. - """ + common page accommodates the unsharded hidden states.""" # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) # = 16 * 2 * 64 * 2 = 4096 bytes. @@ -4083,8 +4075,7 @@ def test_check_enough_kv_cache_memory_reserves_null_block(): def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers - scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs. - """ + scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) full = FullAttentionSpec(block_size=4, **common) mla = MLAAttentionSpec(block_size=4, **common) @@ -4253,8 +4244,7 @@ def test_no_warning_when_draft_group_is_identified(caplog_vllm): def _deepseek_v4_specs(model_version="deepseek_v4"): """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last. - """ + at differing page sizes, with the MTP draft layer registered last.""" return { "model.layers.0.self_attn.attn": new_mla_spec(), "model.layers.1.self_attn.attn": new_mla_spec(), diff --git a/tests/v1/core/test_output.py b/tests/v1/core/test_output.py index 4e1e22eb8e42..ee8a4a0e78f9 100644 --- a/tests/v1/core/test_output.py +++ b/tests/v1/core/test_output.py @@ -55,8 +55,7 @@ def _mm_feature(offset: int, length: int) -> MultiModalFeatureSpec: def test_strip_covered_mm_data() -> None: """Items fully inside the computed prefix lose their data; items touching - the uncomputed region keep it; already-None data stays None. - """ + the uncomputed region keep it; already-None data stays None.""" from dataclasses import replace covered = _mm_feature(offset=0, length=100) @@ -108,8 +107,7 @@ def _mm_feature_mixed(offset: int, length: int) -> MultiModalFeatureSpec: def test_strip_covered_mm_data_mrope() -> None: """For M-RoPE models, covered items keep their keep_on_cpu metadata fields - (the worker needs them to compute positions); payload fields are dropped. - """ + (the worker needs them to compute positions); payload fields are dropped.""" covered = _mm_feature_mixed(offset=0, length=100) uncovered = _mm_feature_mixed(offset=300, length=100) @@ -129,8 +127,7 @@ def test_strip_covered_mm_data_shm_address_item() -> None: """SHM address descriptors must survive stripping even when covered: the worker needs the address to resolve the payload and to acknowledge the sender's writer reference. Stripping it crashed the SHM receiver cache on - the second identical request (vllm-project/vllm#54994). - """ + the second identical request (vllm-project/vllm#54994).""" address_item = MultiModalKwargsItem( { "address": MultiModalFieldElem(data=4096, field=MultiModalBatchedField()), diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 653d485a82ed..01bf73b6e7fa 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -126,8 +126,7 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the config (LCM of group block sizes) unless explicitly provided. This mirrors ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by - these tests, so callers don't have to pass it at every site. - """ + these tests, so callers don't have to pass it at every site.""" kwargs.setdefault( "scheduler_block_size", lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)), @@ -3330,8 +3329,7 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int): @pytest.mark.parametrize("group_id", [0, 1, 2]) def test_block_stored_event_group_idx(group_id: int): """Test BlockStored events emitted by cache_full_blocks carry the correct - group_idx. - """ + group_idx.""" block_size = 4 num_tokens = block_size * 2 @@ -3624,8 +3622,7 @@ def test_block_removed_event_group_idx(group_id: int): def test_emit_cached_block_events(): """emit_cached_block_events emits one BlockStored for already-cached (reused) prefix blocks, carrying the correct group_idx / - parent_block_hash / token_ids, and without mutating block state. - """ + parent_block_hash / token_ids, and without mutating block state.""" block_size = 4 num_cached_blocks = 3 kv_cache_group_id = 1 @@ -3738,8 +3735,7 @@ def test_emit_cached_block_events_zero_cached(): def test_eagle_enabled_removes_last_block(): """Verify Eagle does NOT remove blocks when request - length is divisible by block size. - """ + length is divisible by block size.""" block_size = 16 manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks=10), @@ -4142,8 +4138,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only(): (its right-to-left scan stops once a contiguous match is found). Blocks earlier in the segment can never serve a hit, so ``HybridKVCacheCoordinator.cache_blocks`` should skip them rather than - polluting the prefix-cache hash map. - """ + polluting the prefix-cache hash map.""" block_size = 8 # Full attn block_size=32, SWA block_size=8, sw=8 -> lcm=32. # tail = ceil(7/8) = 1; per_segment = 32/8 = 4. @@ -4214,8 +4209,7 @@ def test_hybrid_cache_blocks_clamped_to_lcm(): Chunks past the last lcm-aligned boundary can never participate in a cache hit (find_longest_cache_hit always returns lcm-aligned hits), so caching them only pollutes the prefix-cache hash map and keeps blocks - on the LRU list that could otherwise return to the free pool. - """ + on the LRU list that could otherwise return to the free pool.""" block_size = 16 # Full attn block_size=32, SWA block_size=16 -> lcm=32. kv_cache_config = KVCacheConfig( @@ -4350,8 +4344,7 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(): ) def test_hybrid_local_kv_retention_interval_rejects_invalid(interval, expected_match): """A retention interval that is negative or not a multiple of - scheduler_block_size errors out at construction time. - """ + scheduler_block_size errors out at construction time.""" block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -4821,8 +4814,7 @@ def test_block_lookup_cache_multi_blocks_per_key(): def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): """Hybrid full+SWA model with a pool sized at the startup minimum should admit a prompt longer than the SWA cap, because SlidingWindowManager - recycles blocks during chunked prefill (issue #39734). - """ + recycles blocks during chunked prefill (issue #39734).""" block_size = 16 sliding_window = 4 * block_size # 64 tokens max_num_batched_tokens = 8 * block_size # 128 tokens @@ -4917,8 +4909,7 @@ def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk( def test_can_fit_full_sequence_full_attention_still_gates_oversized(): """The cap only loosens the SWA group; a prompt that exceeds the - full-attention pool capacity must still be rejected. - """ + full-attention pool capacity must still be rejected.""" block_size = 16 sliding_window = 4 * block_size max_num_batched_tokens = 8 * block_size @@ -5037,8 +5028,7 @@ def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheB without removing them. These ref_cnt==0 blocks stand in for evictable cache-hit blocks left behind by a previous (e.g. preempted) request, and sitting at the head guarantees a later group's external ``get_new_blocks`` - would contend for them on unpatched code (issue #33775). - """ + would contend for them on unpatched code (issue #33775).""" blocks: list[KVCacheBlock] = [] head = manager.block_pool.free_block_queue.fake_free_list_head for _ in range(num_blocks): @@ -5049,8 +5039,7 @@ def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheB def _assert_no_double_allocation(manager: KVCacheManager, req_id: str) -> None: """No physical block may be handed out twice across groups, and every - block referenced by the request must have a live ref_cnt. - """ + block referenced by the request must have a live ref_cnt.""" block_ids = manager.get_blocks(req_id).get_block_ids() flat = [block_id for group in block_ids for block_id in group] assert len(set(flat)) == len(flat), "Block IDs are not unique across groups" @@ -5077,8 +5066,7 @@ def _cross_group_cache_hit( ) -> Request: """Allocate ``req_id`` with a per-group local prefix hit plus external (connector) computed tokens, driving the coordinator's two-phase path. - Returns the allocated request so callers can free it (e.g. to preempt). - """ + Returns the allocated request so callers can free it (e.g. to preempt).""" block_size = _two_phase_block_size(manager) hit_blocks = _take_free_blocks(manager, num_groups * local_blocks_per_group) cache_hit = KVCacheBlocks( @@ -5173,8 +5161,7 @@ def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(): """Dense retention: freeing an SWA request must place its uncached scratch blocks at the front of the free queue (recycled first) and keep its cached checkpoint blocks at the back (retained for prefix - hits). This split is always-on, independent of the retention interval. - """ + hits). This split is always-on, independent of the retention interval.""" block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -5285,8 +5272,7 @@ def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs) def test_pure_swa_retention_interval_caches_sparse_tails(): """Sparse retention must work for a pure-SWA single-group model, not just hybrid models: only the per-interval tails plus the latest replay tail are - cached, and a replay still hits the latest replayable boundary. - """ + cached, and a replay still hits the latest replayable boundary.""" block_size = 16 manager = _make_pure_swa_manager( block_size, sliding_window=block_size, retention_interval=64 @@ -5358,8 +5344,7 @@ def test_pure_swa_retention_latest_only(): def test_pure_swa_dense_retention_caches_all(): """With retention set to ``None``, a pure-SWA model keeps dense behavior: - every block boundary is a potential hit, so all blocks are cached. - """ + every block boundary is a potential hit, so all blocks are cached.""" block_size = 16 manager = _make_pure_swa_manager(block_size, sliding_window=block_size) @@ -5388,8 +5373,7 @@ def test_mamba_reachable_block_mask_sparsifies_retention(): """Mamba state-snapshot retention: with a configured retention interval, the manager keeps one cached state per interval-sized segment (plus the latest replay boundary) instead of a snapshot per block, which is what - lets a small attention block_size avoid Mamba dominating the KV pool. - """ + lets a small attention block_size avoid Mamba dominating the KV pool.""" from vllm.v1.core.single_type_kv_cache_manager import MambaManager block_size = 16 @@ -5427,8 +5411,7 @@ def test_mamba_reachable_block_mask_pins_shared_prefix(): """A Marconi-detected shared prefix (``shared_prefix_boundary``) lands before ``num_prompt`` so the replay-boundary rule alone would drop it. The mask must pin the single state block ending on that boundary so sparse retention does - not defeat cross-request shared-prefix reuse. - """ + not defeat cross-request shared-prefix reuse.""" from vllm.v1.core.single_type_kv_cache_manager import MambaManager block_size = 16 @@ -5476,8 +5459,7 @@ def test_mamba_shared_prefix_survives_zero_retention(): detection) keeps its Mamba state block cached under ``prefix_cache_retention_interval=0``, which otherwise retains only the end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request. - """ + before ``num_prompt``) would be recomputed by every sharing request.""" block_size = 16 # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. @@ -5519,8 +5501,7 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): detecting request must stay reusable by a later request under ``prefix_cache_retention_interval=0``. Without the pin the junction is masked out and the later request misses; with it (and under dense) the reuse - is preserved. - """ + is preserved.""" block_size = 16 def last_req_hit(retention, pin): @@ -5603,8 +5584,7 @@ def reachable_blocks( def test_swa_reachable_block_mask_pins_shared_prefix(): """SWA analog of the Mamba pin: the shared-prefix junction must keep the ``need``-block sliding-window tail ending on that boundary (not a single - block), so a windowed hit can land there under sparse retention. - """ + block), so a windowed hit can land there under sparse retention.""" from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager block_size = 16 @@ -5649,8 +5629,7 @@ def retained(retention, boundary, window, end_block=16): def test_swa_reachable_block_mask_with_dcp_scaling(): """DCP shards each block's KV across ranks, scaling the effective block size. Verify that dcp_world_size > 1 scales block size in reachability calculations, - producing different masks than dcp_world_size=1. - """ + producing different masks than dcp_world_size=1.""" from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager block_size = 16 @@ -5709,8 +5688,7 @@ def test_swa_reachable_block_mask_sub_block_alignment_is_dense( multiple of the DCP-scaled block size) cannot be represented exactly. This happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the full-attention chunk size and need not divide the SWA block size. The mask - must fall back to dense (None) rather than raise. - """ + must fall back to dense (None) rather than raise.""" from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager spec = SlidingWindowSpec( @@ -5738,8 +5716,7 @@ def test_mamba_reachable_block_mask_ignores_dcp(): """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a state block spans kv_cache_spec.block_size tokens regardless of DCP. The mask must not scale by dcp_world_size, so retention granularity is - identical for any DCP world size. - """ + identical for any DCP world size.""" from vllm.v1.core.single_type_kv_cache_manager import MambaManager spec = MambaSpec( @@ -5773,8 +5750,7 @@ def get_mask(dcp_world_size): def test_mamba_reachable_block_mask_large_dcp_stays_sparse(): """A large DCP world size must not scale the Mamba block size, so it can - never collapse sparse retention into dense caching. - """ + never collapse sparse retention into dense caching.""" from vllm.v1.core.single_type_kv_cache_manager import MambaManager spec = MambaSpec( @@ -5809,8 +5785,7 @@ def test_swa_shared_prefix_reuse_under_zero_retention(): """SWA cross-request analog: a partial shared prefix's sliding-window tail must stay reusable under ``prefix_cache_retention_interval=0``. Without the pin the junction window is masked out and a later request misses; with - it (and under dense) reuse is preserved. - """ + it (and under dense) reuse is preserved.""" block_size = 16 def last_req_hit(retention, pin): diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 5774ec014698..3239ec25a00a 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -538,8 +538,7 @@ def test_throttle_defers_inflight_prefill_chunk(): """DP prefill balancing throttles ALL prefill compute on a throttled step, not just new admissions: an in-progress (chunked) prefill already in the running queue is also deferred, so the step runs decode-only, while a - separate decode keeps being scheduled. - """ + separate decode keeps being scheduled.""" scheduler = create_scheduler( max_num_seqs=16, max_num_batched_tokens=50, enable_chunked_prefill=True ) @@ -593,8 +592,7 @@ def test_throttle_capacity_bound_guard_admits(): """Saturation guard: if a cadence-aligned release step cannot drain the waiting prefill queue (it ran out of token budget), the throttle backs off on the next step so the backlog cannot grow into a TTFT avalanche -- prefills are - admitted even though throttle_prefills is set. - """ + admitted even though throttle_prefills is set.""" scheduler = create_scheduler( max_num_seqs=16, max_num_batched_tokens=200, enable_chunked_prefill=True ) @@ -1293,8 +1291,7 @@ def test_pending_hisparse_reclamation_defers_preemption(monkeypatch): def test_prefix_cache_query_not_inflated_by_connector_defer(): """The GPU prefix-cache query is recorded at admission, so a request the - connector defers several times is counted once, not once per retry. - """ + connector defers several times is counted once, not once per retry.""" num_defers_before_matching = 3 scheduler = create_scheduler( enable_prefix_caching=True, @@ -1322,8 +1319,7 @@ def test_prefix_cache_query_not_inflated_by_connector_defer(): def test_preemption_re_records_prefix_cache_query(): """A preempted request re-enters the lookup on resume, so its recomputation - is counted again into the preempted stats. - """ + is counted again into the preempted stats.""" scheduler = create_scheduler(enable_prefix_caching=True) request = create_requests(num_requests=1)[0] scheduler.add_request(request) @@ -1359,8 +1355,7 @@ def test_preemption_re_records_prefix_cache_query(): def test_prefix_cache_stats_not_recorded_when_caching_disabled(): """With prefix caching off there is no local lookup, so admitting a request - records no phantom miss. - """ + records no phantom miss.""" scheduler = create_scheduler(enable_prefix_caching=False) for request in create_requests(num_requests=2): scheduler.add_request(request) @@ -1374,8 +1369,7 @@ def test_prefix_cache_stats_not_recorded_when_caching_disabled(): def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request(): """A real cache hit rejected once by allocate_slots and admitted on the next - step is counted exactly once, hits included. - """ + step is counted exactly once, hits included.""" block_size = 16 scheduler = create_scheduler( enable_prefix_caching=True, @@ -3094,8 +3088,7 @@ def create_requests_with_priority( def test_priority_scheduling_basic_ordering(): """Test that requests are scheduled in priority order - (lower value = higher priority). - """ + (lower value = higher priority).""" scheduler = create_scheduler_with_priority() # Create requests with different priorities @@ -3124,8 +3117,7 @@ def test_priority_scheduling_basic_ordering(): def test_priority_scheduling_arrival_time_tiebreaker(): """Test that arrival time is used - as tiebreaker when priorities are equal. - """ + as tiebreaker when priorities are equal.""" scheduler = create_scheduler_with_priority() # Create requests with same priority but different arrival times @@ -3282,8 +3274,7 @@ def test_priority_scheduling_preemption(): def test_priority_scheduling_no_preemption_when_space_available(): """Test that preemption doesn't happen - when there's space for new requests. - """ + when there's space for new requests.""" scheduler = create_scheduler_with_priority( max_num_seqs=3, # Allow 3 concurrent requests max_num_batched_tokens=200, # Sufficient token budget @@ -3336,8 +3327,7 @@ def test_priority_scheduling_no_preemption_when_space_available(): def test_priority_scheduling_preemption_victim_selection(): """Test that the correct victim is selected for - preemption based on priority and arrival time. - """ + preemption based on priority and arrival time.""" # This test verifies the priority-based victim selection logic # by checking the waiting queue order after adding requests with different # priorities @@ -3457,8 +3447,7 @@ def test_priority_scheduling_waiting_queue_order(): def test_priority_scheduling_fcfs_fallback(): """Test that FCFS behavior is maintained when all - requests have same priority. - """ + requests have same priority.""" scheduler = create_scheduler_with_priority() # Create requests with same priority but different arrival times @@ -3527,8 +3516,7 @@ def test_priority_scheduling_with_limited_slots(): def test_priority_scheduling_heap_property(): """Test that the waiting queue maintains heap - property for priority scheduling. - """ + property for priority scheduling.""" scheduler = create_scheduler_with_priority( max_num_seqs=1, # Only one request can run at a time ) @@ -3753,8 +3741,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv( use_ec_connector, ec_role, use_v2_model_runner ): """Test that priority scheduling preempts lower priority requests - when out of KV cache space. - """ + when out of KV cache space.""" # Create scheduler with very limited memory to force preemption scheduler = create_scheduler_with_priority( max_num_seqs=2, # Allow multiple requests @@ -3899,8 +3886,7 @@ def test_chunked_prefill_disabled_for_encoder_decoder( enable_chunked_prefill: bool, is_encoder_decoder: bool, expect_enabled: bool ) -> None: """Validate that chunked prefill is appropriately disabled for - encoder-decoder models. - """ + encoder-decoder models.""" scheduler_config = SchedulerConfig( enable_chunked_prefill=enable_chunked_prefill, is_encoder_decoder=is_encoder_decoder, @@ -3929,8 +3915,7 @@ def _validate_chunked_prefill_settings_for_encoder_decoder( scheduler_config: SchedulerConfig, is_encoder_decoder: bool, expect_enabled: bool ) -> None: """Validate chunked prefill settings in the scheduler config for - encoder-decoder models. - """ + encoder-decoder models.""" assert scheduler_config.enable_chunked_prefill is expect_enabled if is_encoder_decoder: # Encoder-decoder models should automatically disable chunked multimodal @@ -4020,8 +4005,7 @@ def _assert_right_encoder_inputs( ): """Verify that requests/mm_hashes should (not) in scheduled encoder input If check_exist is False, this function returns True - if requests are NOT in encoder inputs - """ + if requests are NOT in encoder inputs""" # Get the scheduled encoder inputs # NOTE: scheduled_encoder_inputs is a dictionary with request id as key scheduled_encoder_inputs = output.scheduled_encoder_inputs @@ -4182,8 +4166,7 @@ def test_ec_connector_text_only_request(use_kv_connector): @pytest.mark.parametrize("use_kv_connector", [False, True]) def test_ec_connector_cache_hit_external_load(use_kv_connector): """Test ec_consumer loads from external cache when hit. - A normal basic operation for EPD disaggrgation - """ + A normal basic operation for EPD disaggrgation""" scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", enable_prefix_caching=True, @@ -4630,8 +4613,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption( cache_exist, use_kv_connector ): """Test that priority scheduling preempts lower priority requests - when out of KV cache space. - """ + when out of KV cache space.""" # Create scheduler with very limited memory to force preemption scheduler = create_scheduler_with_priority( model="llava-hf/llava-1.5-7b-hf", @@ -5300,8 +5282,7 @@ def update_connector_output( def test_ec_connector_update_connector_output_called(): """Test that worker-side EC connector output is forwarded to the - EC connector's update_connector_output hook. - """ + EC connector's update_connector_output hook.""" scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", use_ec_connector=True, @@ -5506,8 +5487,7 @@ def test_variable_length_cross_attn_block_allocation(): def test_cross_attn_blocks_not_over_allocated(): """Test that cross-attention blocks are not over-allocated compared to - what each request actually needs. - """ + what each request actually needs.""" from math import ceil block_size = 16 @@ -5603,8 +5583,7 @@ def test_cross_attn_blocks_not_under_allocated(): def test_cross_attn_zero_blocks_without_encoder_inputs(): """Test that requests without encoder inputs get zero cross-attention - blocks, even when the scheduler is configured for encoder-decoder. - """ + blocks, even when the scheduler is configured for encoder-decoder.""" block_size = 16 scheduler = _create_encoder_decoder_scheduler(block_size=block_size) @@ -5709,8 +5688,7 @@ def test_free_encoder_inputs_respects_unconfirmed_placeholders(): evicted entry and crashes the engine with "Encoder cache miss". The scheduler must retain the input until the *confirmed* progress (num_computed_tokens - num_output_placeholders) passes the range end, so - that no pending rejection can rewind into the range. - """ + that no pending rejection can rewind into the range.""" scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -5821,8 +5799,7 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead(): position so the drafter's +1 look-ahead mm-embedding gather (which reads one position past the target's computed range) still finds it cached. This is the primary mechanism that prevents the drafter "Encoder cache miss"; the - worker-side token-embedding fallback is only a backstop. - """ + worker-side token-embedding fallback is only a backstop.""" scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") # create_scheduler only builds ngram spec configs; force the eagle path that # _free_encoder_inputs keys off (its read-ahead deferral). @@ -5852,8 +5829,7 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead(): def test_free_encoder_inputs_unchanged_without_spec_decode(): """Without speculative decoding, encoder inputs are freed as soon as - num_computed_tokens passes the placeholder range, as before. - """ + num_computed_tokens passes the placeholder range, as before.""" scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") mm_positions = [[PlaceholderRange(offset=50, length=100)]] request = create_requests( @@ -5884,8 +5860,7 @@ def test_encoder_cache_retained_across_preemption_and_resume(): encoder input must pull the still-cached entry back out of `freeable` without scheduling a recompute, keeping the scheduler and worker consistent. The spec-rollback retention margin does not gate this path, - so it is covered separately here. - """ + so it is covered separately here.""" scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -5929,8 +5904,7 @@ def test_encoder_cache_recomputed_when_evicted_during_preemption(): pressure before it resumes, the scheduler reports the mm_hash as freed (so the worker drops it) and a resume must schedule a recompute rather than assume the worker still holds it. check_and_update_cache must - return False so the encoder input is re-scheduled. - """ + return False so the encoder input is re-scheduled.""" scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", num_speculative_tokens=3, @@ -6581,8 +6555,7 @@ def _decode_ready_request(scheduler): def test_update_draft_token_ids_strips_ngram_padding(): """ngram_gpu pads unfilled draft slots with -1; manager validation must - strip them before grammar.validate_tokens (which otherwise raises). - """ + strip them before grammar.validate_tokens (which otherwise raises).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) @@ -6601,8 +6574,7 @@ def test_update_draft_token_ids_strips_ngram_padding(): def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to manager validation is stripped). - """ + is preserved (only the input to manager validation is stripped).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py index beaf5203f24c..737796a4bf7d 100644 --- a/tests/v1/core/test_swa_inflight_window_free.py +++ b/tests/v1/core/test_swa_inflight_window_free.py @@ -97,8 +97,7 @@ def test_num_in_flight_tokens_accounting(): def test_swa_free_waits_for_in_flight_step(): """Async: out-of-window blocks stay allocated until the step that still - reads them has been processed. - """ + reads them has been processed.""" scheduler = _create_swa_scheduler(async_scheduling=True) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -133,8 +132,7 @@ def test_swa_free_waits_for_in_flight_step(): def test_swa_free_immediate_when_sync(): """Sync: no in-flight step at schedule time, frees happen at the first - decode allocation as before. - """ + decode allocation as before.""" scheduler = _create_swa_scheduler(async_scheduling=False) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -173,8 +171,7 @@ def test_swa_admission_cap_accounts_for_overlapping_batches(): def test_chunked_local_free_waits_for_in_flight_step(): """Chunked-local attention frees whole chunks left of the current one, and is exposed to the same load-WAR: with async scheduling those chunks must - stay allocated until the in-flight step that still reads them settles. - """ + stay allocated until the in-flight step that still reads them settles.""" scheduler = _create_chunked_scheduler(async_scheduling=True) request = create_requests( num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE @@ -229,8 +226,7 @@ def test_connector_finish_frees_on_settled_basis(): is handed to a KV connector (simple CPU offload / NIXL store), must use the same processed-token basis. Otherwise the connector reads/hands off a block the still-in-flight step is optimistically counted as done with, which is - the load-path WAR this fix closes. - """ + the load-path WAR this fix closes.""" scheduler = create_scheduler( block_size=BLOCK_SIZE, async_scheduling=True, diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 3e87770d4e82..b0b49a9567b1 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -80,8 +80,7 @@ def forward_fn(cg_mode): @pytest.fixture(autouse=True) def _reset_breakable_tls(): """Defensively clear thread-local capture state between tests so a - failure in one test can't leak "nested capture" errors into the next. - """ + failure in one test can't leak "nested capture" errors into the next.""" from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture BreakableCUDAGraphCapture._tls.active = None @@ -291,8 +290,7 @@ def eager_break_op(): def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): """Verify @eager_break_during_capture correctly routes through add_eager when inside a capture context, and runs straight through - when there's no active capture. - """ + when there's no active capture.""" from vllm.compilation.breakable_cudagraph import ( BreakableCUDAGraphCapture, eager_break_during_capture, @@ -554,8 +552,7 @@ def _mock_vllm_config(): def test_wrapper_falls_through_on_runtime_mode_mismatch(): """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing - nothing, so an outer CUDAGraphWrapper can own them. - """ + nothing, so an outer CUDAGraphWrapper can own them.""" from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index d0a4c5c5b2b4..5699a661c17f 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -275,8 +275,7 @@ def test_rms_norm_per_block_quant_batch_invariant( hidden_size: int, dtype: torch.dtype, seed: int ): """C++ ``rms_norm_per_block_quant`` must be batch invariant across the - block threshold (compiled fp8 block-quant path; block pinned to 512). - """ + block threshold (compiled fp8 block-quant path; block pinned to 512).""" import vllm._custom_ops as ops device = torch.device(DEVICE_TYPE) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 4c544fab7b21..cb23910566c7 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -547,8 +547,7 @@ async def controller_task(): async def test_dp_pause_keep_race_staggered_engines(): """Race: send pause(keep) to engine 0, then add two requests, then pause(keep) to engine 1. Ensures no deadlock when pause - requests are staggered and requests arrive in between. - """ + requests are staggered and requests arrive in between.""" if DP_SIZE != 2: pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2") @@ -718,8 +717,7 @@ async def consume_gen(req_id: str) -> None: async def test_dp_pause_wait_mode_drains_in_flight(): """mode="wait" through the DP consensus path: the pause lets the in-flight request run to completion while the consensus is pending, - resolves only once drained, and the engine generates again on resume. - """ + resolves only once drained, and the engine generates again on resume.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) after.callback(engine.shutdown) @@ -750,8 +748,7 @@ async def test_dp_pause_wait_mode_drains_in_flight(): async def test_dp_pause_while_asleep(): """Pausing a sleeping DP engine: sleeping ranks skip dummy batches yet must still reach the pause consensus, and the completion barrier must - be harmless on workers whose memory is unmapped. - """ + be harmless on workers whose memory is unmapped.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) after.callback(engine.shutdown) @@ -787,8 +784,7 @@ async def test_dp_pause_completion_implies_device_idle( """A resolved pause future promises an idle device. Inflate the dummy batches the pause consensus manufactures so any work the pause fails to wait on stays visible, then require a quiet worker stream the moment - pause_generation() returns. - """ + pause_generation() returns.""" with ExitStack() as after: # The probes below ship functions through collective_rpc. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index 95a31f37092c..b53d0d8af09d 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -24,8 +24,7 @@ class ExternalLBServerManager: """Manages data parallel vLLM server instances for external - load balancer testing. - """ + load balancer testing.""" def __init__( self, diff --git a/tests/v1/distributed/test_hybrid_lb_dp.py b/tests/v1/distributed/test_hybrid_lb_dp.py index 748370f37c5c..fcd3c69af54a 100644 --- a/tests/v1/distributed/test_hybrid_lb_dp.py +++ b/tests/v1/distributed/test_hybrid_lb_dp.py @@ -30,8 +30,7 @@ class HybridLBServerManager: """Manages hybrid data parallel vLLM server instances where each node runs a single logical API server that balances requests only to the - DP engines running on that same node. - """ + DP engines running on that same node.""" def __init__( self, diff --git a/tests/v1/distributed/test_internal_lb_dp.py b/tests/v1/distributed/test_internal_lb_dp.py index 9fe8c9f0afd5..b438d898b337 100644 --- a/tests/v1/distributed/test_internal_lb_dp.py +++ b/tests/v1/distributed/test_internal_lb_dp.py @@ -107,8 +107,7 @@ async def _run_request_bursts( class MultinodeInternalLBServerManager: """Manages multi-node data parallel vLLM server instances for internal - load balancer testing using --headless mode. - """ + load balancer testing using --headless mode.""" def __init__( self, @@ -242,8 +241,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): class APIOnlyServerManager: """Manages API-only server (Node 0) and headless engines server (Node 1) - for testing separated API server and engine configuration. - """ + for testing separated API server and engine configuration.""" def __init__( self, @@ -645,8 +643,7 @@ async def test_api_only_multinode_dp_completion_streaming( model_name: str, ) -> None: """Test API-only server streaming with all engines on separate - headless server. - """ + headless server.""" prompt = "What is an LLM?" async def make_streaming_request(): diff --git a/tests/v1/distributed/test_pp_dp_v2.py b/tests/v1/distributed/test_pp_dp_v2.py index f2a274019b12..1e5d4267f474 100644 --- a/tests/v1/distributed/test_pp_dp_v2.py +++ b/tests/v1/distributed/test_pp_dp_v2.py @@ -82,8 +82,7 @@ async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int: @pytest.mark.parametrize("async_scheduling", [True, False]) async def test_pp_dp_v2_low_concurrency(async_scheduling: bool): """A single in-flight request at a time, repeated, to exercise the - PP slot ring under empty batches between decodes. - """ + PP slot ring under empty batches between decodes.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) after.callback(engine.shutdown) @@ -97,8 +96,7 @@ async def test_pp_dp_v2_low_concurrency(async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool): """64 concurrent requests, staggered, to exercise the steady-state - DP all-reduce + PP slot-ring path. - """ + DP all-reduce + PP slot-ring path.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) after.callback(engine.shutdown) @@ -114,8 +112,7 @@ async def _one(i: int) -> int: @pytest.mark.asyncio async def test_pp_dp_v2_abort_mid_decode(): """Cancel half the in-flight requests mid-stream and confirm the - engine survives the abort storm. - """ + engine survives the abort storm.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) after.callback(engine.shutdown) @@ -155,8 +152,7 @@ async def _maybe_cancel(i: int): @pytest.mark.asyncio async def test_pp_dp_v2_pause_resume(): """Pause an engine with a request in flight, then resume and confirm - new requests still work. - """ + new requests still work.""" with ExitStack() as after: engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) after.callback(engine.shutdown) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 0610463484aa..8218f325290d 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -48,8 +48,7 @@ def test_without_spec_decoding( monkeypatch: pytest.MonkeyPatch, ): """Test consistency of combos of async scheduling, preemption, - uni/multiproc executor, prefill chunking. - """ + uni/multiproc executor, prefill chunking.""" struct_outputs = StructuredOutputsParams(json=sample_json_schema) test_sampling_params: list[dict[str, Any]] = [ dict(), @@ -205,8 +204,7 @@ def run_tests( test_sampling_params: list[dict[str, Any]], ): """Test consistency of combos of async scheduling, preemption, - uni/multiproc executor with spec decoding. - """ + uni/multiproc executor with spec decoding.""" # Flex attention supports float32. attention_config = {"backend": "FLEX_ATTENTION"} diff --git a/tests/v1/e2e/general/test_attention_backend_per_kind.py b/tests/v1/e2e/general/test_attention_backend_per_kind.py index 42f2a72554e6..3281247b1c4a 100644 --- a/tests/v1/e2e/general/test_attention_backend_per_kind.py +++ b/tests/v1/e2e/general/test_attention_backend_per_kind.py @@ -72,8 +72,7 @@ def test_backend_per_kind_splits_groups(backend_per_kind, monkeypatch): ) def test_backend_per_kind_overrides_global_backend(monkeypatch): """A per-kind entry wins over the global ``backend`` for its kind; other - kinds fall back to the global backend. - """ + kinds fall back to the global backend.""" monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") llm = LLM( model=MODEL, diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index 48f60747e832..91a450089f6f 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -34,8 +34,7 @@ def test_decoder_max_context_length_validation( max_tokens: int, ) -> None: """Check vLLM decoder model input validation for edge cases where - the prompt length is (almost) equal to the max model length. - """ + the prompt length is (almost) equal to the max model length.""" prompt_ids = [[43] * prompt_len] with vllm_runner( @@ -72,8 +71,7 @@ def test_auto_fit_max_model_len_rejects_oversized_input( """When max_model_len='auto' and KV cache memory is very limited, the engine auto-fits max_model_len to a small value. The frontend must see this reduced value and reject prompts that exceed it, - rather than accepting them and hanging. - """ + rather than accepting them and hanging.""" # Use a small KV cache budget to force auto-fit to a small # max_model_len. Pin block_size=16 so the budget is independent # of the platform's default block size. One block for this model is diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py index 26a1770d39cd..c87339b7685b 100644 --- a/tests/v1/e2e/general/test_min_tokens.py +++ b/tests/v1/e2e/general/test_min_tokens.py @@ -172,8 +172,7 @@ def llm_v1(): def test_min_tokens_with_terminal_structured_output(llm_v1: LLM): """A grammar whose only legal continuation is EOS must stop under - min_tokens instead of sampling outside the grammar. - """ + min_tokens instead of sampling outside the grammar.""" params = SamplingParams( temperature=GREEDY, min_tokens=20, diff --git a/tests/v1/e2e/general/test_sharded_sampling.py b/tests/v1/e2e/general/test_sharded_sampling.py index 4db56337a82a..a580605fb007 100644 --- a/tests/v1/e2e/general/test_sharded_sampling.py +++ b/tests/v1/e2e/general/test_sharded_sampling.py @@ -7,8 +7,7 @@ bitwise identical across the two modes; cross-state boots differ by up to ~0.25 logprob and can flip a near-tie token). A real sharding bug — rows routed to the wrong request, misaligned gathers — produces many-nats -divergence and scrambled top-k sets, far outside these bounds. -""" +divergence and scrambled top-k sets, far outside these bounds.""" import pytest @@ -107,8 +106,7 @@ def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None: @multi_gpu_test(num_gpus=2) def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch): """Generation, logprobs, and prompt logprobs match between batch-sharded - sampling and the replicated fallback, up to measured boot-state noise. - """ + sampling and the replicated fallback, up to measured boot-state noise.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") # Required for the collective_rpc mode check below. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py index 4fcaa2a59c84..8714e4f56a1e 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py @@ -11,8 +11,7 @@ @single_gpu_only def test_synthetic_acceptance_rate(vllm_runner): """Verify that synthetic rejection sampling produces an acceptance - length close to the requested mean acceptance length. - """ + length close to the requested mean acceptance length.""" num_spec_tokens = 3 expected_acceptance_len = 1.875 tolerance = 0.15 diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py index f4cd06ca69d8..3a474ac2039c 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -167,8 +167,7 @@ def test_draft_model_tensor_parallelism(vllm_runner): @multi_gpu_only(num_gpus=2) def test_draft_model_engine_args_tensor_parallelism(): """Ensure the vllm_config for the draft model is created correctly, - and independently of the target model (quantization, TP, etc.) - """ + and independently of the target model (quantization, TP, etc.)""" engine_args = EngineArgs( model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized tensor_parallel_size=2, @@ -199,8 +198,7 @@ def test_draft_model_engine_args_tensor_parallelism(): def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic - so we can test it without instantiating a full proposer. - """ + so we can test it without instantiating a full proposer.""" spec_cfg = vllm_config.speculative_config if spec_cfg.moe_backend is not None: return replace( @@ -225,8 +223,7 @@ def _platform_moe_backend(monkeypatch: pytest.MonkeyPatch) -> MoEBackend: def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): """When moe_backend is set in speculative_config, the draft VllmConfig - should use it while the target keeps its own setting. - """ + should use it while the target keeps its own setting.""" target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", @@ -251,8 +248,7 @@ def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch): """When moe_backend is not set in speculative_config, the draft should - inherit the target's moe_backend. - """ + inherit the target's moe_backend.""" target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", @@ -275,8 +271,7 @@ def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch def test_draft_model_moe_backend_default_auto(): """When neither target nor draft set moe_backend explicitly, both should - default to 'auto'. - """ + default to 'auto'.""" engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, @@ -297,8 +292,7 @@ def test_draft_model_moe_backend_default_auto(): def test_draft_model_engine_args_rejects_invalid_tp_argname(): """The user should pass "draft_tensor_parallel_size" rather than - "tensor_parallel_size". We enforce this with validation. - """ + "tensor_parallel_size". We enforce this with validation.""" engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, @@ -315,8 +309,7 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname(): def assert_draft_model_correctness(args: ArgsTest, vllm_runner): """Compare the outputs using and not using speculative decoding. - In the greedy decoding case, the outputs must match EXACTLY. - """ + In the greedy decoding case, the outputs must match EXACTLY.""" test_prompts: list[Messages] = get_messages( dataset=args.dataset, n=args.num_prompts ) diff --git a/tests/v1/e2e/spec_decode/test_sharded_sampling.py b/tests/v1/e2e/spec_decode/test_sharded_sampling.py index bcb39ad95acd..cc9781ca9fdb 100644 --- a/tests/v1/e2e/spec_decode/test_sharded_sampling.py +++ b/tests/v1/e2e/spec_decode/test_sharded_sampling.py @@ -28,8 +28,7 @@ def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch): sampling under MTP spec decoding: the collectives move the same logits bytes, Gumbel keys derive from (request slot, position, seed), and slot assignment is rank-deterministic. Both runs here are spec decode with - identical math, so outputs must match exactly. - """ + identical math, so outputs must match exactly.""" tp_size = 2 _skip_if_insufficient_gpus_for_tp(tp_size) model_name = "Qwen/Qwen3.5-0.8B-Base" diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index 177c2f4bdbfd..66c8d5b9b76c 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -67,8 +67,7 @@ def test_mtp_speculative_mixed_batch_short_prefill( vllm_runner, model_name, enable_prefix_caching ): """Test to ensure MTP speculative decoding correctly handles - short prefill chunks that fall below the reorder_batch_threshold. - """ + short prefill chunks that fall below the reorder_batch_threshold.""" # Set so large that both prefills will be classified as decodes in a mixed batch # note, with prefix caching we require chunk_size >= mamba_block_size chunk_size = 256 if not enable_prefix_caching else 2048 diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py index df584d3dcb8a..844753290948 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py @@ -381,8 +381,7 @@ def test_pin_if_ready_none_when_absent(): def test_pin_if_ready_returns_unpinned_entry_when_not_ready(): """The producer needs "present but not ready" told apart from "absent": - the first is NACKed as retryable, the second as a miss. - """ + the first is NACKed as retryable, the second as a miss.""" cache = EmbeddingCache(num_blocks=4) cache.alloc("h", 1) # not-ready entry = cache.pin_if_ready("h") diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 9b6338a77383..24ea717cabb7 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -131,8 +131,7 @@ def test_has_cache_item_false_when_not_consumer(monkeypatch): def test_connector_keys_on_identifier_not_mm_hash(monkeypatch): """The connector must key the encoder cache on feature.identifier (what - has_cache_item is called with), NOT feature.mm_hash. - """ + has_cache_item is called with), NOT feature.mm_hash.""" s = _make_scheduler(monkeypatch) req = _Request([_Feature("PROC_KEY", length=1, identifier="ENC_KEY")]) s.update_state_after_alloc(req, 0) @@ -202,8 +201,7 @@ def test_loads_only_serves_hashes_touched_this_step(monkeypatch): def test_repeated_reload_same_step_loads_once(monkeypatch): """The same mm_hash requested twice in one step must appear in - meta.loads exactly once. - """ + meta.loads exactly once.""" s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=2) @@ -226,8 +224,7 @@ def test_load_not_emitted_for_uncached_entry(monkeypatch): def test_load_pin_protects_blocks_until_completion(monkeypatch): """Loaded blocks stay pinned until the worker reports the load memcpy complete. Observable: a save that needs eviction cannot reclaim a - still-pinned loaded entry, but can once the completion report arrives. - """ + still-pinned loaded entry, but can once the completion report arrives.""" # 2 blocks, both occupied by "a" (ready). s = _make_scheduler(monkeypatch, num_blocks=2) _seed_cached(s, "a", n_blocks=2) @@ -254,8 +251,7 @@ def test_load_pin_protects_blocks_until_completion(monkeypatch): def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool: """Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the - blocks currently held by other entries are reclaimable. - """ + blocks currently held by other entries are reclaimable.""" s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0) meta = s.build_connector_meta(scheduler_output=None) return mm_hash in meta.saves @@ -315,8 +311,7 @@ def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch): def test_region_full_skips_save_and_never_blocks(monkeypatch): """When the region is fully occupied by pinned entries, new saves are - silently skipped and ensure_cache_available never blocks. - """ + silently skipped and ensure_cache_available never blocks.""" s = _make_scheduler(monkeypatch, num_blocks=1) _seed_cached(s, "pinned", n_blocks=1) s._cache.pin("pinned") @@ -335,8 +330,7 @@ def test_region_full_skips_save_and_never_blocks(monkeypatch): def test_producer_only_never_emits_loads(monkeypatch): """A producer-only scheduler must never populate meta.loads, even when - entries are ready. - """ + entries are ready.""" s = _make_scheduler(monkeypatch, ec_role="ec_producer") req = _Request([_Feature("h1", length=1)]) @@ -378,8 +372,7 @@ def test_consumer_only_has_cache_item(monkeypatch): def test_save_not_emitted_for_already_cached_entry(monkeypatch): """An entry already in the cache (from a prior save) must not trigger - a second allocation or appear in meta.saves again. - """ + a second allocation or appear in meta.saves again.""" s = _make_scheduler(monkeypatch) req = _Request([_Feature("h1", length=1)]) @@ -400,8 +393,7 @@ def test_save_not_emitted_for_already_cached_entry(monkeypatch): def test_update_connector_output_ignores_foreign_meta(monkeypatch): """A worker output whose payload is not an ECCPUWorkerMetadata is a - no-op — no cache mutation, no crash. - """ + no-op — no cache mutation, no crash.""" s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=1) @@ -415,8 +407,7 @@ class _Foreign: def test_update_connector_output_ignores_unknown_report(monkeypatch): """Reports for a never-seen save hash or a transfer id this scheduler never - dispatched are dropped rather than mutating unrelated state. - """ + dispatched are dropped rather than mutating unrelated state.""" s = _make_scheduler(monkeypatch) s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242])) assert s.has_cache_item("gone") is False @@ -428,8 +419,7 @@ def test_update_connector_output_ignores_unknown_report(monkeypatch): def test_has_pending_push_work_tracks_inflight_save(monkeypatch): """A dispatched-but-unconfirmed save keeps push work pending until the - completion report marks it ready. - """ + completion report marks it ready.""" s = _make_scheduler(monkeypatch) assert s.has_pending_push_work() is False @@ -444,8 +434,7 @@ def test_has_pending_push_work_tracks_inflight_save(monkeypatch): def test_has_pending_push_work_tracks_inflight_load(monkeypatch): """A dispatched-but-unconfirmed load keeps push work pending until the - unpin completion report arrives. - """ + unpin completion report arrives.""" s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=1) assert s.has_pending_push_work() is False # ready + unpinned @@ -464,8 +453,7 @@ def test_has_pending_push_work_tracks_inflight_load(monkeypatch): def test_shutdown_disables_roles_and_cleans_region(monkeypatch): """After shutdown the scheduler serves no items and the region is - cleaned up. - """ + cleaned up.""" s = _make_scheduler(monkeypatch, num_blocks=2) _seed_cached(s, "a", n_blocks=2) assert s.has_cache_item("a") is True diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index fc5a8fc8368a..722c0d89fa21 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -98,8 +98,7 @@ def test_second_instance_opens_existing_file_and_shares_memory(): def test_only_creator_unlinks_file_on_cleanup(): """Critical contract: if the non-creator unlinks, the creator's mmap path becomes a dangling backing file and a third opener would create a new one - out from under the creator. - """ + out from under the creator.""" instance_id = str(uuid.uuid4()) r1 = ECSharedRegion( engine_id=instance_id, @@ -216,8 +215,7 @@ def _spy_fallback(mm, off, ln): def test_non_creator_does_not_probe_or_populate(monkeypatch): """Only the creator pre-faults. A second opener must not probe the advice - nor re-touch pages the creator already populated. - """ + nor re-touch pages the creator already populated.""" from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr instance_id = str(uuid.uuid4()) @@ -314,8 +312,7 @@ def test_madvise_unexpected_oserror_propagates(monkeypatch): def test_pin_memory_success_sets_flag(region): """When cudaHostRegister returns 0, _is_pinned flips to True - and cleanup will correspondingly call cudaHostUnregister. - """ + and cleanup will correspondingly call cudaHostUnregister.""" fake_cudart = MagicMock() success = MagicMock() success.value = 0 @@ -335,8 +332,7 @@ def test_pin_memory_success_sets_flag(region): def test_pin_memory_failure_leaves_flag_false(): """If cudaHostRegister fails (non-zero), don't pretend it succeeded — - cleanup must NOT call cudaHostUnregister on memory we never registered. - """ + cleanup must NOT call cudaHostUnregister on memory we never registered.""" r = _make_region() try: fake_cudart = MagicMock() diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index 317c714ccae0..f4f5f03f6946 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -363,8 +363,7 @@ def test_save_caches_coalesces_only_adjacent_blocks( def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker): """When the scheduler hasn't pre-allocated blocks for ``mm_hash``, - ``save_caches`` + ``flush_saves`` is a pure no-op. - """ + ``save_caches`` + ``flush_saves`` is a pure no-op.""" worker = make_worker() sentinel = 0x42 worker._region.blocks.fill_(sentinel) @@ -382,8 +381,7 @@ def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker): ) def test_save_caches_noop_for_non_save_rank(make_worker, tp_rank, pcp_rank): """Only TP rank 0 + PCP rank 0 writes to mmap. All other combinations - must skip the save entirely. - """ + must skip the save entirely.""" worker = make_worker(tp_rank=tp_rank, pcp_rank=pcp_rank) sentinel = 0x42 worker._region.blocks.fill_(sentinel) @@ -398,8 +396,7 @@ def test_save_caches_noop_for_non_save_rank(make_worker, tp_rank, pcp_rank): def test_save_caches_raises_when_allocated_blocks_too_small(make_worker): """``save_caches`` must raise ``AssertionError`` when the encoder output - is larger than the allocated block space. - """ + is larger than the allocated block space.""" worker = make_worker() src = torch.zeros(3 * _HIDDEN_DIM, dtype=_DTYPE) with pytest.raises(AssertionError, match="exceeds allocated blocks"): @@ -552,8 +549,7 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker): def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): """Every hash in ``meta.loads`` is copied from mmap, including one whose key is already resident in ``encoder_cache`` — there is no worker-side - skip. - """ + skip.""" worker = make_worker() n_blocks = 3 src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape( @@ -782,8 +778,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker): @_requires_accelerator def test_buffer_pool_is_reused_across_load_steps(make_worker): """Once a load copy completes its descriptor buffers return to the pool and - are reused by the next call — no reallocation. - """ + are reused by the next call — no reallocation.""" worker = make_worker() worker._region.blocks[0].fill_(0x01) worker._region.blocks[1].fill_(0x02) @@ -854,8 +849,7 @@ def test_init_pins_memory_when_available(pin_available, expected_pinned): def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm): """``shutdown`` must always call ``region.cleanup`` — and must never - raise. - """ + raise.""" worker = object.__new__(ECCPUWorker) mock_region = Mock(spec=ECSharedRegion) worker._region = mock_region diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py index 1cceddb79c1a..558a80f9cf82 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -43,8 +43,7 @@ def assign(proxy): def _drive(assign, e_urls, counts): """Feed a sequence of request item-counts through the cursor, as - fanout_encoder_primer does one request at a time. - """ + fanout_encoder_primer does one request at a time.""" cursor = 0 all_urls = [] for count in counts: diff --git a/tests/v1/ec_connector/unit/test_metadata.py b/tests/v1/ec_connector/unit/test_metadata.py index e9b78e09f287..65fdc6bbb392 100644 --- a/tests/v1/ec_connector/unit/test_metadata.py +++ b/tests/v1/ec_connector/unit/test_metadata.py @@ -40,8 +40,7 @@ def test_metadata_saves_dict_operations(): def test_metadata_loads_carry_a_transfer_id_and_blocks(): """Each load entry is a (transfer_id, block_ids) pair, so the worker can - report the transfer rather than the mm_hash. - """ + report the transfer rather than the mm_hash.""" meta = ECCPUConnectorMetadata() meta.loads["mm_hash_1"] = (0, [5, 6, 7]) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index 462dcc80a9ce..156d98ae794c 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -79,8 +79,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): """A save whose GPU->mmap copy hasn't been confirmed complete yet is still announced: the entry can't be evicted before it's ready, and a read arriving too early is NACKed NACK_NOT_READY for the consumer to retry — - protecting it is not this method's job. - """ + protecting it is not this method's job.""" s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index fff7ba2ba746..5c25fdf2288d 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -129,8 +129,7 @@ def test_producer_session_grant_nacks_missing_entry(): def test_producer_session_grant_nacks_not_ready_entry_as_retryable(): """An announced entry whose save is still in flight is distinguished from a - miss, so the consumer retries instead of falling back to local compute. - """ + miss, so the consumer retries instead of falling back to local compute.""" cache = EmbeddingCache(num_blocks=8) cache.alloc("h1", 2) # not ready s = _make_producer_session(cache) @@ -438,8 +437,7 @@ def test_consumer_session_poll_nack_goes_to_tombstoned(): def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): """A not-ready NACK must not tombstone: tombstoning admits the request with - no embedding, which is fatal once the media has been rewritten away. - """ + no embedding, which is fatal once the media has been rewritten away.""" import msgspec s = _make_consumer_session() diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 49858f685cd2..69f52e2a7f34 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -68,8 +68,7 @@ def get_addresses(self) -> tuple[list[str], list[str]]: # in ``tests/entrypoints/test_api_server_process_manager.py``. def _bind_and_report_worker(listen_address, sock, args, client_config): """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via ``actual_address_pipe``, then exit. - """ + endpoints back via ``actual_address_pipe``, then exit.""" ctx = zmq.Context() try: in_sock = make_zmq_socket( diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index 5843ce5f55ae..ba1e9a4d200b 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -607,8 +607,7 @@ def _cadenced_dp_engine_core( monkeypatch, results: list[tuple[bool, bool]], dp_sync_interval: int = 32 ): """A bare DPEngineCoreProc whose all-reduce is scripted by `results`; - returns the core and the step numbers at which the all-reduce ran. - """ + returns the core and the step numbers at which the all-reduce ran.""" core = object.__new__(DPEngineCoreProc) core.dp_group = object() core.dp_sync_interval = dp_sync_interval @@ -627,8 +626,7 @@ def scripted_sync(dp_group, has_unfinished, pending_pause): def test_dp_sync_interval_default_is_16(): """Regression pin: lowering the default narrows the mid-wave pause tail - for async RL, where the engine is rarely idle when pause lands. - """ + for async RL, where the engine is rarely idle when pause lands.""" assert ParallelConfig.dp_sync_interval == 16 @@ -646,8 +644,7 @@ def test_dp_sync_interval_normal_wave(monkeypatch, dp_sync_interval: int): def test_dp_sync_interval_idle_pause_consensus_on_first_step(monkeypatch): """A pause of an idle engine arms every rank before its kick-started first step, so the step-1 sync reaches consensus after one dummy batch - regardless of the configured cadence. - """ + regardless of the configured cadence.""" core, synced = _cadenced_dp_engine_core( monkeypatch, [(False, True)], dp_sync_interval=32 ) @@ -673,8 +670,7 @@ def _pausable_engine_core_proc() -> EngineCoreProc: @pytest.mark.parametrize("deferred", [False, True]) def test_pause_synchronizes_device_before_cache_reset(deferred: bool): """A resolved pause promises an idle device: the barrier must run before - caches are cleared and before the caller is unblocked. - """ + caches are cleared and before the caller is unblocked.""" core = _pausable_engine_core_proc() core.engines_running = deferred order: list[str] = [] diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index e6520bfadc8e..a23ae46b8f5d 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -249,8 +249,7 @@ def test_dplb_non_late_interaction_still_uses_lb(): def test_dplb_burst_round_robins_despite_snapshot_rebinds(): """A stats snapshot rebind wipes the optimistic lb_engines increments; - the exact in-flight floor must keep a burst spreading round-robin. - """ + the exact in-flight floor must keep a burst spreading round-robin.""" client = _make_dplb_client(num_engines=4) for _ in range(4): @@ -325,8 +324,7 @@ async def wait_for_snapshot(): def test_dplb_snapshot_backpressure_overrides_inflight(): """An engine reported heavily loaded by the coordinator is avoided even - when this client has routed nothing to it. - """ + when this client has routed nothing to it.""" client = _make_dplb_client(num_engines=2) client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]] @@ -339,8 +337,7 @@ def test_dplb_snapshot_backpressure_overrides_inflight(): def test_dplb_kv_pressure_amplifies_waiting_penalty(): """A waiting queue on a KV-bound engine (slow drain) is penalized, while - the same queue with low KV usage is not (e.g. transient burst). - """ + the same queue with low KV usage is not (e.g. transient burst).""" client = _make_dplb_client(num_engines=2) # Engine 0 has a smaller total but is KV-bound with a queue. client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]] diff --git a/tests/v1/engine/test_logprobs_processor.py b/tests/v1/engine/test_logprobs_processor.py index 3a315c04d501..edb8cef518ca 100644 --- a/tests/v1/engine/test_logprobs_processor.py +++ b/tests/v1/engine/test_logprobs_processor.py @@ -30,8 +30,7 @@ def _make_processor(num_logprobs: int) -> LogprobsProcessor: def test_drops_trailing_sentinel_columns(): """A request that asked for 3 custom token logprobs but ended up in a - batch padded to width 5 must not surface the trailing -inf entries. - """ + batch padded to width 5 must not surface the trailing -inf entries.""" processor = _make_processor(num_logprobs=3) sampled = 42 diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index cdd61445bf6e..9df43088e6c8 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -229,8 +229,7 @@ def test_request_stream_interval_raises_but_not_below_engine_default( ): """A per-request stream_interval can raise the interval above the engine default but not below it (values under the default clamp up), without - altering the generated text. - """ + altering the generated text.""" engine_stream_interval = 5 # Request 0 (below the default) clamps up to 5; request 1 raises it to 10. request_stream_intervals = [1, 10] diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py index 3c2f63c66ab6..9c3ba6d555dd 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py @@ -80,8 +80,7 @@ def _split_hnd_cache(spec) -> torch.Tensor: def _packed_nhd_cache(spec) -> torch.Tensor: """Logical (num_blocks, heads, block_size, 2 * head_size) over an NHD - physical layout — the FlashAttention/FlashInfer/Triton/Flex form. - """ + physical layout — the FlashAttention/FlashInfer/Triton/Flex form.""" return torch.zeros( NUM_BLOCKS, spec.block_size, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 2d87431e319d..f597c8fa27c8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -863,8 +863,7 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): ) def test_canonical_layout_gate(kv_cache_groups, certified): """The canonical layout certifies portability group by group; none of - these shapes are portable in the direct layout. - """ + these shapes are portable in the direct layout.""" assert not _parallelism_agnostic(kv_cache_groups) assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified @@ -872,8 +871,7 @@ def test_canonical_layout_gate(kv_cache_groups, certified): def test_canonical_layout_certifies_v2_model_runner(): """Canonical bytes are certified per layer against live tensor strides at registration, so the static gate must not depend on the model-runner - version — the v2 runner is the case the canonical layout exists for. - """ + version — the v2 runner is the case the canonical layout exists for.""" groups = _groups(_full_attention_spec()) assert _parallelism_agnostic(groups, canonical=True, v2=True) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index 961ce1af12b7..b7d9d6c8efac 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -117,8 +117,7 @@ def _spec_cls_with_metric_definitions( metric_definitions: dict[str, Any], ) -> type: """Build a fake offloading spec class reporting the given metric - definitions, so tests don't need to patch the real CPU spec. - """ + definitions, so tests don't need to patch the real CPU spec.""" class _FakeOffloadingSpec: @staticmethod @@ -188,8 +187,7 @@ def test_build_kv_connector_stats_with_empty_dict(): def test_build_kv_connector_stats_reconstructs_offload_stats(): """Test that OffloadingConnector stats are properly reconstructed with - correct data. - """ + correct data.""" serialized_data = { _StatsKey.TYPES: { LOAD_BYTES: _MetricType.COUNTER, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index df69b9bb66f7..035fb6b0054f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1583,8 +1583,7 @@ def _maximal_lookup(sched, keys, start_chunk_idx: int = 0): @pytest.mark.parametrize("result", list(LookupResult)) def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult): """Whether a result keeps a sliding-window streak alive decides how wide - the demanded chunk set gets, so every member needs deliberate behavior. - """ + the demanded chunk set gets, so every member needs deliberate behavior.""" assert result in _SCAN_BEHAVIOR, f"{result} has no declared scan behavior" expected_lookups, expected_end = _SCAN_BEHAVIOR[result] @@ -2035,8 +2034,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling def test_loads_do_not_populate_fence_index(request_runner): """Loads don't populate _block_id_to_pending_jobs (protected by - delay_free_blocks while in flight). - """ + delay_free_blocks while in flight).""" runner = request_runner( blocks_per_chunk=3, block_size=4, @@ -2106,8 +2104,7 @@ def capture_fence(): def test_fence_at_build_store_jobs(request_runner): """A new prefill (no load -> update_state_after_alloc returns early) reusing a finished request's pending-store block is flushed by - _build_store_jobs's fence. - """ + _build_store_jobs's fence.""" runner = request_runner( blocks_per_chunk=1, block_size=4, @@ -2157,8 +2154,7 @@ def capture_fence(): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_complete_store_called_per_job(request_runner, async_scheduling: bool): """complete_store fires per-job, not deferred to request finish. - Each call carries only that store's keys. - """ + Each call carries only that store's keys.""" tokens_per_block = 4 blocks_per_chunk = 3 tokens_per_chunk = tokens_per_block * blocks_per_chunk @@ -2387,8 +2383,7 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_reset_cache(request_runner, async_scheduling: bool): """reset_cache flushes in-flight loads, calls manager.reset_cache(), resets - next_stored_chunk_idx for active requests and clears job tracking. - """ + next_stored_chunk_idx for active requests and clears job tracking.""" block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -2934,8 +2929,7 @@ def test_stale_sliding_window_block_after_prepare_store_failure( @pytest.mark.parametrize("async_scheduling", [True, False]) def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): """When skip_reading_prefix_cache=True, the offloading connector must not - load any blocks from CPU even if a matching prefix is cached there. - """ + load any blocks from CPU even if a matching prefix is cached there.""" block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -3029,8 +3023,7 @@ def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: class TestEagle: """Tests for Eagle/MTP speculative decoding support in the offloading - connector scheduler — both _lookup() unit tests and integration tests. - """ + connector scheduler — both _lookup() unit tests and integration tests.""" # ------------------------------------------------------------------- # Helpers @@ -3572,8 +3565,7 @@ def test_sw_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): """Eagle sliding-window group stores all prompt blocks but excludes - the trailing chunk while decoding. - """ + the trailing chunk while decoding.""" block_size = 4 sliding_window = 8 num_gpu_blocks = 100 @@ -3625,8 +3617,7 @@ def test_single_block_stored_at_end_of_prefill( self, request_runner, async_scheduling: bool ): """An eagle group with a single-block prompt stores it at the end of - prefill: prompt blocks are stable, so no tail is held back. - """ + prefill: prompt blocks are stable, so no tail is held back.""" block_size = 4 blocks_per_chunk = 1 tokens_per_chunk = block_size * blocks_per_chunk @@ -4263,8 +4254,7 @@ def capture_fence(): def test_chunked_local_attention_reports_its_chunk_window(): """Llama 4 uses chunked local attention, which used to trip the - FullAttentionSpec assert and kill the engine at startup. - """ + FullAttentionSpec assert and kill the engine at startup.""" spec = ChunkedLocalAttentionSpec( block_size=16, num_kv_heads=2, @@ -4503,8 +4493,7 @@ def test_retention_interval_zero_stores_only_replay_boundary( def _shared_kv_mtp_config(): """Speculative config for a shared-group MTP model: eagle-family method whose drafter layer merges into a target KV-cache group, so no group - self-identifies as a drafter group (issue #52735). - """ + self-identifies as a drafter group (issue #52735).""" spec = MagicMock(name="shared_kv_mtp_spec") spec.use_eagle.return_value = True spec.use_eagle_block_drop.return_value = True @@ -4518,14 +4507,12 @@ def _shared_kv_mtp_config(): class TestSharedGroupMTPOffload: """Regression tests for issue #52735: OffloadingConnector must keep serving when speculative decoding is enabled but no KV-cache group is - annotated as a drafter group (shared-group MTP models). - """ + annotated as a drafter group (shared-group MTP models).""" def test_no_annotation_marks_no_groups(self, request_runner): """Spec decode on + zero annotated groups must NOT mark every group as a drafter group; the full store->load roundtrip must match the - non-speculative behavior of test_two_groups_full_and_sliding_window. - """ + non-speculative behavior of test_two_groups_full_and_sliding_window.""" block_size = 4 kv_cache_groups = [ KVCacheGroupSpec( @@ -4582,8 +4569,7 @@ def test_no_annotation_marks_no_groups(self, request_runner): class TestMambaHybridOffloadServing: """Store->finish->lookup flows on a full-attention + mamba-align hybrid, with a manager that only HITs keys that were actually stored. Guards the - two collapse routes of issue #52735. - """ + two collapse routes of issue #52735.""" BLOCK = 4 MAMBA_BLOCK = 16 @@ -4712,8 +4698,7 @@ def test_shared_group_mtp_serves_like_non_speculative(self): """Route 1 of #52735: with spec decode on and no drafter annotation, serving must equal the non-speculative baseline (was 0 before the fix: the all-groups fallback marked the mamba group as a drafter and - the volatile-tail pop consumed its only servable chunk). - """ + the volatile-tail pop consumed its only servable chunk).""" baseline = self._roundtrip_served_tokens(self._make_scheduler(None)) assert baseline == 16 served = self._roundtrip_served_tokens( @@ -4726,8 +4711,7 @@ def test_annotated_eagle_group_does_not_starve_coarse_sibling(self): group must not drag the confirmed boundary below a coarser sibling's chunk granularity. The widened query makes the volatile-tail pop land on an extra queried chunk, holding the boundary at 16 tokens - (was 0 before the fix: 4 hits -> pop -> 12 tokens < mamba chunk). - """ + (was 0 before the fix: 4 hits -> pop -> 12 tokens < mamba chunk).""" scheduler = self._make_scheduler(None, mamba_eagle=True) assert [c.is_eagle_group for c in scheduler.config.kv_group_configs] == [ True, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 9fdac585276d..7a987fef3daa 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -80,8 +80,7 @@ def _allocate_kv_caches( def _single_rank_vllm_config(total_kv_heads: int): """A one-rank (TP=1) parallel config, as canonical mappings are derived - from it. - """ + from it.""" vllm_config = MagicMock() parallel_config = vllm_config.parallel_config parallel_config.tensor_parallel_size = 1 diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 006229b8699f..78a6d930f97a 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -148,8 +148,7 @@ def test_multiturn_lifecycle(): (do_remote_decode=True). Finishes LENGTH_CAPPED with remote_block_ids. Turn 2: P receives remote_block_ids from D. P pulls KV from D because remote_block_ids is not None and external tokens > 0. Computes only - new tokens, finishes LENGTH_CAPPED. - """ + new tokens, finishes LENGTH_CAPPED.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -195,8 +194,7 @@ def test_multiturn_lifecycle(): def test_first_turn_no_remote_blocks(): """First turn: P has no remote_block_ids from D yet. - Standard local prefill, returns kv_transfer_params for future turns. - """ + Standard local prefill, returns kv_transfer_params for future turns.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -276,8 +274,7 @@ def test_abort_p_side_non_length_capped(): def test_remote_blocks_exceed_prompt_tokens(): """D provides more remote tokens than P's prompt needs. - P caps external tokens to prompt length. - """ + P caps external tokens to prompt length.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -313,8 +310,7 @@ def test_remote_blocks_exceed_prompt_tokens(): def test_p_node_pulls_partial_last_block_from_d(): """D sends remote_block_ids with partially filled last block. remote_num_tokens < len(remote_block_ids) * block_size. - P pulls only remote_num_tokens worth of external tokens. - """ + P pulls only remote_num_tokens worth of external tokens.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -359,8 +355,7 @@ def test_p_node_pulls_partial_last_block_from_d(): def test_add_new_req_to_recv_populates_remote_meta(): """add_new_req_to_recv correctly populates RemoteMeta for P-node - bi-directional KV pull from D. - """ + bi-directional KV pull from D.""" meta = NixlConnectorMetadata() kv_params = { "remote_block_ids": [[0, 1, 2]], @@ -388,8 +383,7 @@ def test_add_new_req_to_recv_populates_remote_meta(): def test_build_connector_meta_recv_entries(): """P-node scheduler: do_remote_decode=True + remote_block_ids → - _reqs_need_recv populated, build_connector_meta produces reqs_to_recv. - """ + _reqs_need_recv populated, build_connector_meta produces reqs_to_recv.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -463,8 +457,7 @@ def test_p_node_pull_kv_from_d(dist_init): ) def test_p_node_pull_then_send_kv(dist_init): """Full P-node bi-directional: pull KV from D → prefill → - send KV back to D via notification. - """ + send KV back to D via notification.""" connector, worker = _make_connector_with_fake_worker() meta = _make_p_node_recv_metadata("req-p2", [10, 11], [20, 21]) _do_load_kv(connector, meta) @@ -511,8 +504,7 @@ def test_p_node_deferred_pull_on_no_handshake(dist_init): def test_d_node_request_finished_returns_kv_params(): """D-node request_finished returns kv_transfer_params with do_remote_decode=True, remote_block_ids, remote_num_tokens - for P to pull. These params go directly to P node. - """ + for P to pull. These params go directly to P node.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -593,8 +585,7 @@ def test_d_node_request_finished_remote_num_tokens(): def test_d_node_partial_last_block_remote_num_tokens(): """D-node: remote_num_tokens < len(remote_block_ids) * block_size - when last block is partially filled. - """ + when last block is partially filled.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -628,8 +619,7 @@ def test_no_double_read_blocks_after_reschedule(): bidirectional request (once on initial schedule, once after WAITING_FOR_REMOTE_KVS → reschedule). The _remote_blocks_processed flag must prevent the request from being added to _reqs_need_recv - twice, which would cause P to read D's blocks twice. - """ + twice, which would cause P to read D's blocks twice.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -682,8 +672,7 @@ def test_remote_num_tokens_bounded_by_blocks(): """Edge case 2: D-node request_finished must return remote_num_tokens <= len(remote_block_ids) * block_size. request.num_tokens includes the last sampled token which has no KV - in the cache, so remote_num_tokens must use num_computed_tokens. - """ + in the cache, so remote_num_tokens must use num_computed_tokens.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -719,8 +708,7 @@ def test_remote_num_tokens_bounded_by_blocks(): def test_kv_recompute_threshold_skips_small_transfer(): """Edge case 3: When remote tokens are below kv_recompute_threshold, P should skip the remote pull and compute locally instead of - entering WAITING_FOR_REMOTE_KVS. - """ + entering WAITING_FOR_REMOTE_KVS.""" threshold = 256 vllm_config = create_vllm_config( kv_connector_extra_config={ @@ -762,8 +750,7 @@ def test_p_node_finished_holds_blocks_for_d(): """Edge case 4: P-node finishes with FINISHED_LENGTH_CAPPED and do_remote_decode=True. P must hold blocks (delay_free=True) and return kv_transfer_params with do_remote_prefill=True so D can - read P's blocks. - """ + read P's blocks.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -804,8 +791,7 @@ def test_p_node_finished_holds_blocks_for_d(): def test_cache_miss_first_turn_no_remote_pull(): """Edge case 5: First turn with do_remote_decode=True but no remote_block_ids (cache MISS). P should prefill locally with - num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS. - """ + num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -840,8 +826,7 @@ def test_cache_miss_first_turn_no_remote_pull(): def test_partial_remote_tokens_less_than_prompt(): """Edge case 6: D's remote_num_tokens covers only part of P's prompt. P should pull remote_num_tokens worth of external tokens - and compute the rest locally. - """ + and compute the rest locally.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -886,8 +871,7 @@ def test_partial_remote_tokens_less_than_prompt(): def test_remote_blocks_processed_flag_persists(): """Edge case 7: After recv completes and request is rescheduled, the _remote_blocks_processed flag in kv_transfer_params prevents - the bidirectional path from re-entering _reqs_need_recv. - """ + the bidirectional path from re-entering _reqs_need_recv.""" vllm_config = create_vllm_config( kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, ) @@ -959,8 +943,7 @@ def test_remote_blocks_processed_flag_persists(): ) def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): """P declines on the turn-2 readback if the deadline is known, - near-expiry and a handshake clock offset is known - """ + near-expiry and a handshake clock offset is known""" connector, worker = _make_connector_with_fake_worker() worker._engine_clock_offset[_REMOTE] = offset expiry_time = None if expiry_delta is None else time.perf_counter() + expiry_delta @@ -996,8 +979,7 @@ def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): ) def test_turn2_full_prefix_hit_with_expired_deadline_skips_gate(dist_init): """A full local prefix hit issues no READ. The expiry gate must be skipped - even when D's deadline is expired. - """ + even when D's deadline is expired.""" connector, worker = _make_connector_with_fake_worker() worker._engine_clock_offset[_REMOTE] = 0.0 @@ -1050,8 +1032,7 @@ def test_d_node_request_finished_exports_blocks_expiry_time(): def test_handshake_listener_appends_perf_counter_frame(): """The handshake reply carries a 2nd frame with the listener's live - perf_counter. - """ + perf_counter.""" import threading import msgspec diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index d3dabe4fb468..019b8d1504ac 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -19,15 +19,13 @@ class _StubLMCacheMPConnector: test only asserts on the connector *name* and the ``extra_config`` dict produced by ``VllmConfig``, never instantiates the connector, so a bare placeholder class is sufficient. Not subclassing ``SupportsHMA`` mirrors - the real connector's HMA support (it does not support HMA either). - """ + the real connector's HMA support (it does not support HMA either).""" @pytest.fixture def stub_lmcache_mp_connector(monkeypatch): """Replace the lazy loader so VllmConfig.__post_init__ does not import - ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests. - """ + ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests.""" monkeypatch.setitem( KVConnectorFactory._registry, "LMCacheMPConnector", @@ -107,8 +105,7 @@ def _build_config( enable_cumem_allocator: bool = False, ) -> VllmConfig: """Build a VllmConfig that exercises _verify_kv_transfer_compat without - requiring a real model (avoids HF downloads in CI). - """ + requiring a real model (avoids HF downloads in CI).""" from types import SimpleNamespace kv_transfer_config = ( @@ -134,8 +131,7 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): invalidated when expandable_segments lets the CUDA VMM allocator remap the underlying physical pages. We can't enumerate every connector that does this (especially OOT ones), so reject the combination whenever any - connector is configured. - """ + connector is configured.""" monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") with pytest.raises(ValueError, match="expandable_segments"): _build_config(kv_connector=kv_connector) @@ -143,8 +139,7 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch): """Sleep mode routes KV allocations through CuMemAllocator's pool, which - auto-disables expandable_segments (see #40812). - """ + auto-disables expandable_segments (see #40812).""" monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") _build_config(kv_connector="NixlConnector", enable_sleep_mode=True) @@ -167,8 +162,7 @@ def test_kv_connector_allows_other_alloc_conf(monkeypatch): def test_no_kv_connector_ignores_expandable_segments(monkeypatch): """The expandable_segments check only applies when a KV connector is - configured. - """ + configured.""" monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") _build_config(kv_connector=None) diff --git a/tests/v1/kv_connector/unit/test_flexkv_connector.py b/tests/v1/kv_connector/unit/test_flexkv_connector.py index a40e5f828cab..8cb57366345c 100644 --- a/tests/v1/kv_connector/unit/test_flexkv_connector.py +++ b/tests/v1/kv_connector/unit/test_flexkv_connector.py @@ -50,8 +50,7 @@ def _make_flexkv_module( impl_mock: MagicMock, ) -> tuple[types.ModuleType, types.ModuleType]: """Build a fake ``flexkv`` package hierarchy that returns *impl_mock* - when ``FlexKVConnectorV1Impl`` is instantiated. - """ + when ``FlexKVConnectorV1Impl`` is instantiated.""" flexkv_mod = types.ModuleType("flexkv") integration_mod = types.ModuleType("flexkv.integration") vllm_mod = types.ModuleType("flexkv.integration.vllm") @@ -72,8 +71,7 @@ def _make_flexkv_module( def _install_flexkv_mock(impl_mock: MagicMock): """Insert fake flexkv modules into sys.modules and return a context that - cleans them up afterwards. - """ + cleans them up afterwards.""" flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock) mods = { "flexkv": flexkv_mod, @@ -106,8 +104,7 @@ def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock): class TestFlexKVConnectorImportError: """FlexKVConnectorV1 should fail with a helpful message when flexkv is - absent. - """ + absent.""" def test_import_error_message(self): from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import ( diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index b2cf0f0b64dd..e5cbd977b6d8 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -176,8 +176,7 @@ def test_engine_unwraps_handshake_metadata_for_legacy_connector( ) -> None: """Engine core always asks workers for `(pp_rank, tp_rank)`-keyed metadata, then unwraps to `{tp_rank: metadata}` for a connector that has not opted - into PP-aware handshake (single-PP producer, all `pp_rank == 0`). - """ + into PP-aware handshake (single-PP producer, all `pp_rank == 0`).""" metadata_0 = _Metadata() metadata_1 = _Metadata() connector = _LegacyConnector() @@ -200,8 +199,7 @@ def test_engine_rejects_pp_producer_for_legacy_connector( monkeypatch: pytest.MonkeyPatch, ) -> None: """A connector that has not opted into PP-aware handshake must not silently - drop metadata from `pp_rank > 0`; engine core init raises instead. - """ + drop metadata from `pp_rank > 0`; engine core init raises instead.""" connector = _LegacyConnector() with pytest.raises(ValueError, match="does not support PP-disaggregated"): @@ -216,8 +214,7 @@ def test_engine_passes_handshake_metadata_through_for_pp_aware_connector( monkeypatch: pytest.MonkeyPatch, ) -> None: """A PP-aware connector receives the full `(pp_rank, tp_rank)`-keyed dict - unchanged. - """ + unchanged.""" metadata_0 = _Metadata() metadata_1 = _Metadata() connector = _PPAwareConnector() diff --git a/tests/v1/kv_connector/unit/test_hf3fs_client.py b/tests/v1/kv_connector/unit/test_hf3fs_client.py index d0ae6c24acd0..26a5de11c522 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_client.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_client.py @@ -131,8 +131,7 @@ def _make_client(self, tmp_path): def test_close_idempotent_and_handles_cleared(self, tmp_path): """Multiple close() calls must not raise; deregister_fd called exactly - once, all handles set to None, shm.close() invoked. - """ + once, all handles set to None, shm.close() invoked.""" client, shm_r, shm_w = self._make_client(tmp_path) with ( @@ -238,8 +237,7 @@ def test_constructor_failure_after_shm_alloc_closes_shm(self, tmp_path): def test_constructor_failure_does_not_close_unallocated_shm(self, tmp_path): """Failure before SharedMemory is created must not raise AttributeError - or TypeError from cleanup. - """ + or TypeError from cleanup.""" with ( patch(f"{self._MOD}.HF3FS_AVAILABLE", True), patch(f"{self._MOD}.register_fd"), diff --git a/tests/v1/kv_connector/unit/test_lmcache_connector.py b/tests/v1/kv_connector/unit/test_lmcache_connector.py index ad8ac0566e0d..a64f7f32d5a6 100644 --- a/tests/v1/kv_connector/unit/test_lmcache_connector.py +++ b/tests/v1/kv_connector/unit/test_lmcache_connector.py @@ -212,8 +212,7 @@ def test_does_nothing_when_kv_cache_events_is_not_lmcache_kv_events( self, mock_connector ): """Test that method returns early when kv_cache_events is not - LMCacheKVEvents. - """ + LMCacheKVEvents.""" # Create a mock object that is not LMCacheKVEvents fake_events = MagicMock() connector_output = KVConnectorOutput(kv_cache_events=fake_events) diff --git a/tests/v1/kv_connector/unit/test_mooncake_stats.py b/tests/v1/kv_connector/unit/test_mooncake_stats.py index 6e2bb59fa4f2..a20fcb505330 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_stats.py +++ b/tests/v1/kv_connector/unit/test_mooncake_stats.py @@ -199,8 +199,7 @@ def _bare_worker() -> MooncakeConnectorWorker: """Construct a MooncakeConnectorWorker skipping __init__ (full init requires a live TransferEngine). Only the attributes touched by the methods under test are populated; role flags and async_zmq_ctx keep __del__'s shutdown - path a no-op. - """ + path a no-op.""" worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) worker.xfer_stats = MooncakeKVConnectorStats() worker.engine = MagicMock() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 7e8e532fffb9..f3ffa55a5439 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -77,8 +77,7 @@ def test_store_tp_does_not_override_backend_kv_cache_layout(extra_config): def _make_kv_cache_config() -> KVCacheConfig: """Single-group full-attention KVCacheConfig — enough for the connector - constructor's validate() pass. - """ + constructor's validate() pass.""" spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) return KVCacheConfig( num_blocks=4, @@ -693,8 +692,7 @@ def _poll_lookup(client, req_id, num_tokens=128, block_hashes=(), timeout=5.0): def _gated_recv(gate: threading.Event, value: int): """Mock recv side-effect that blocks until ``gate`` is set, so the - executor's lookup can be held pending deterministically. - """ + executor's lookup can be held pending deterministically.""" def recv(): gate.wait() @@ -705,8 +703,7 @@ def recv(): def test_lookup_key_client_non_block_lookup_async(): """Non-blocking lookup defers to the executor: None first, hit once the - Future resolves. - """ + Future resolves.""" vllm_config = _make_vllm_config() with patch( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 583c24958831..8767b1c0d983 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -33,8 +33,7 @@ def _mamba_align(block_size=32): def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for - the test fixtures. - """ + the test fixtures.""" block_sizes = [g.kv_cache_spec.block_size for g in groups] scheduler_block_size = lcm(*block_sizes) return MooncakeStoreCoordinator( @@ -144,8 +143,7 @@ def test_coordinator_single_full_attention_no_hits(): def test_coordinator_single_swa_tautological_pool_masks_pre_window(): """SWA tautological-pool: hit_length spans full prefix, mask is - tail-window only. - """ + tail-window only.""" groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) # 4 chunks * 16 tokens @@ -192,8 +190,7 @@ def test_coordinator_hybrid_hole_in_full_clips_both(): def test_coordinator_group_block_size_double_hash(): """Group block_size=32 over hash_block_size=16 hashes: adjacent - hashes merge before pool lookup. - """ + hashes merge before pool lookup.""" groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _full(32)), @@ -217,8 +214,7 @@ def test_coordinator_group_block_size_double_hash(): def test_coordinator_fine_grained_partial_tail_hit(): """K3 shape: FA + mamba-align, block_size=32 over hash_block_size=16. When both groups have the sub-block boundary hash, the reconciled hit lands on - the hash boundary (48), not the block boundary (32). - """ + the hash boundary (48), not the block boundary (32).""" groups = [ KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), @@ -238,8 +234,7 @@ def test_coordinator_fine_grained_partial_tail_hit(): def test_coordinator_fine_grained_clips_when_one_group_missing_tail(): """If only one group has the sub-block boundary, min-convergence clips the - reconciled hit back to the block boundary (32). - """ + reconciled hit back to the block boundary (32).""" groups = [ KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), @@ -280,8 +275,7 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): """Hybrid full-attn(block=32) + SWA(block=8, sw=8). lcm=32. With aligned=64 the SWA group should mark exactly the blocks ending at 32 and 64 (i.e. blocks 3 and 7 at block_size=8); the rest can never - participate in any future hit. - """ + participate in any future hit.""" full = _full(32) swa = _swa(block_size=8, sliding_window=8) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -327,8 +321,7 @@ def make(dcp_world_size): def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm - boundary should now span two SWA tail blocks. - """ + boundary should now span two SWA tail blocks.""" full = _full(32) swa = _swa(block_size=8, sliding_window=16) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -358,8 +351,7 @@ def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa(): """DSV4-shaped: full-MLA(B=256) + 4 SWA groups with B in {64, 64, 4, 8} and varied sliding windows. lcm=256, hash_block_size=4. Two lcm segments (aligned_len=512). Validates that the tile-once strategy produces the - expected per-segment tail-window pattern, repeated. - """ + expected per-segment tail-window pattern, repeated.""" full_mla = _full(block_size=256) swa_64_sw128 = _swa(block_size=64, sliding_window=128) swa_64_sw512 = _swa(block_size=64, sliding_window=512) @@ -393,8 +385,7 @@ def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa(): def test_store_mask_fast_path_all_block_sizes_equal_lcm(): """When every non-full-attn group already aligns to lcm_block_size, the - fast path returns all-True without invoking find_longest_cache_hit. - """ + fast path returns all-True without invoking find_longest_cache_hit.""" full = _full(block_size=64) swa = _swa(block_size=64, sliding_window=128) groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -407,8 +398,7 @@ def test_store_mask_fast_path_all_block_sizes_equal_lcm(): def test_store_mask_fast_path_single_attention_group(): """Two groups sharing the same SWA spec collapse to one attention group; - no lcm filter applies, every chunk is True. - """ + no lcm filter applies, every chunk is True.""" swa = _swa(block_size=16, sliding_window=32) groups = [KVCacheGroupSpec(["L0"], swa), KVCacheGroupSpec(["L1"], swa)] coord = _make_coord(groups, hash_block_size=16) @@ -422,8 +412,7 @@ def test_store_mask_fast_path_single_attention_group(): def _retention_groups(): """Hybrid full-attn(block=32) + SWA(block=8, sw=8); lcm=32. The SWA group - densely keeps one tail block per 32-token boundary. - """ + densely keeps one tail block per 32-token boundary.""" full = _full(32) swa = _swa(block_size=8, sliding_window=8) return [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] @@ -431,8 +420,7 @@ def _retention_groups(): def test_store_mask_dense_default_matches_every_lcm_boundary(): """retention_interval=None (default) keeps the SWA tail at every lcm - boundary: tokens 32/64/96/128 -> chunks 3/7/11/15. - """ + boundary: tokens 32/64/96/128 -> chunks 3/7/11/15.""" coord = _make_coord(_retention_groups(), hash_block_size=8) masks = coord.store_mask(128) assert masks[0] is None @@ -442,8 +430,7 @@ def test_store_mask_dense_default_matches_every_lcm_boundary(): def test_store_mask_retention_interval_sparsifies_swa_tails(): """retention_interval=64 keeps an SWA tail once per 64-token segment (chunks 7 and 15) instead of every 32 tokens, dropping the mid-segment - boundaries at 32 and 96. - """ + boundaries at 32 and 96.""" coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64) masks = coord.store_mask(128) assert masks[0] is None # full attn unaffected @@ -452,8 +439,7 @@ def test_store_mask_retention_interval_sparsifies_swa_tails(): def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): """retention_interval=0 drops all segment tails; only the latest replay - boundary (capped at num_prompt-1, aligned down to lcm) is retained. - """ + boundary (capped at num_prompt-1, aligned down to lcm) is retained.""" coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=0) # No replay info -> nothing reachable for the SWA group. assert coord.store_mask(128)[1] == [False] * 16 @@ -464,8 +450,7 @@ def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): def test_store_mask_retention_interval_keeps_segment_and_replay_tails(): """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay - boundary tail (num_prompt=100 -> chunk 11) coexist. - """ + boundary tail (num_prompt=100 -> chunk 11) coexist.""" coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64) masks = coord.store_mask(128, num_prompt_tokens=100) assert masks[1] == [i in (7, 11, 15) for i in range(16)] @@ -503,8 +488,7 @@ def test_store_mask_excludes_mamba_groups_lookup_unaffected(): nulled/freed; speculative blocks relocate), so the positional normal save must never cover mamba chunks — regardless of retention. Lookups still probe mamba boundaries: their keys come from the pinned snapshot/CoW - hand-off path instead. - """ + hand-off path instead.""" groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _mamba_align(16)), @@ -525,8 +509,7 @@ def test_store_mask_excludes_mamba_groups_lookup_unaffected(): def test_lookup_with_eagle_pops_last_full_attention_block(): """Sanity: with use_eagle, find_longest_cache_hit drops the last block. - Pairs with the load_mask test below to lock the round-trip contract. - """ + Pairs with the load_mask test below to lock the round-trip contract.""" groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16, use_eagle=True) hs = _hashes(4) @@ -567,8 +550,7 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention(): def test_load_mask_with_eagle_hybrid_full_plus_swa(): """Hybrid (FullAttn + SWA) with eagle: load_mask must cover every chunk in [0, token_len) for the FullAttn group; SWA group keeps its - tail-window mask. - """ + tail-window mask.""" groups = [ KVCacheGroupSpec(["L0"], _full(16)), KVCacheGroupSpec(["L1"], _swa(16, 32)), diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 85d49700895d..6c2cff0a542c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """End-to-end save->lookup test for MooncakeStoreConnector on a hybrid -(SWA + Full) attention config, using a dict-backed mock store. -""" +(SWA + Full) attention config, using a dict-backed mock store.""" import sys import threading @@ -283,8 +282,7 @@ def _fake_thread_init(*args, **kwargs): def test_recv_skips_swa_blocks_before_window(): """Producer stored every block for both groups; consumer must only fetch - SWA blocks within the sliding window, not the head. - """ + SWA blocks within the sliding window, not the head.""" full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) # sliding_window=32, block_size=16 → 2 contiguous blocks within window. swa = SlidingWindowSpec( @@ -352,8 +350,7 @@ def batch_get_into_multi_buffers(self, keys, addrs, sizes): def test_chunked_token_database_hash_block_size_smaller_than_block_size(): """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens - keys each chunk by its ending fine hash, including a partial tail. - """ + keys each chunk by its ending fine hash, including a partial tail.""" md = KeyMetadata("m", 0, 0, 0, 0, group_id=3) db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4) db.set_kv_caches_base_addr([0]) @@ -386,8 +383,7 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the partial tail is offloaded for both groups under the boundary sub-hash. The full-attention block is read from the request block table; the mamba block - is the core-provided CoW target, not block_ids. - """ + is the core-provided CoW target, not block_ids.""" full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -466,8 +462,7 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): def test_offload_syncs_event_before_put(): """An offload-carrying meta synchronizes its CoW-fence event before the store put reads the blocks, then completes in one pass and drains the - completion counter. - """ + completion counter.""" full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -538,8 +533,7 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): offload must persist every FA block up to the boundary — the normal save floors to the lcm, so those blocks are otherwise never written and the consumer's per-group lookup would miss. The mamba boundary block still - reads the core-provided CoW target. - """ + reads the core-provided CoW target.""" full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -810,8 +804,7 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): """A per-request ring group (CircularBufferSpec, capacity as block_size) sits beside the paged group: worker setup tolerates its block size (the DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged - blocks only, and lookup hits on the paged group alone. - """ + blocks only, and lookup hits on the paged group alone.""" full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) ring = CircularBufferSpec( block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8 diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index af2dfd608313..7ced38d65a41 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -536,8 +536,7 @@ def _make_pending_load_unfinished_request( def _make_pending_load_scheduler_output() -> SimpleNamespace: """scheduler_output for a step where req-0 is parked on a pending load - (not in scheduled_new_reqs or scheduled_cached_reqs). - """ + (not in scheduled_new_reqs or scheduled_cached_reqs).""" return SimpleNamespace( finished_req_ids=set(), preempted_req_ids=set(), @@ -713,8 +712,7 @@ def test_running_request_not_in_resumed_req_ids_appends_blocks(): def test_resumed_request_in_resumed_req_ids_replaces_blocks(): """A request the scheduler marks resumed gets the FULL block table in new_block_ids and must REPLACE the tracker's blocks (not append), even if - a stale tracker from before preemption is still present. - """ + a stale tracker from before preemption is still present.""" scheduler = _make_bare_scheduler() _make_resumed_unfinished_request( scheduler, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 9646331fe014..7dcce073a0ea 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -432,8 +432,7 @@ def _patch_worker_runtime( def test_pool_key_to_string_without_prefix_is_unchanged(): """Default (empty) cache_prefix keeps keys byte-identical to the - historical unprefixed format so existing deployments keep their hits. - """ + historical unprefixed format so existing deployments keep their hits.""" key = PoolKey(KeyMetadata("test-model", 0, 0, 0, 0), "deadbeef") assert ( key.to_string() == "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@deadbeef" @@ -574,8 +573,7 @@ def test_tp_shared_receiving_reads_each_local_store_shard(): def test_pool_key_cache_prefix_namespaces_and_disambiguates(): """A non-empty cache_prefix is prepended, and two instances with - different prefixes never collide on identical block hashes. - """ + different prefixes never collide on identical block hashes.""" md_a = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="depA") md_b = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="depB") @@ -591,8 +589,7 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates(): def test_default_local_buffer_size_matches_pr40900(): """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode - patch preserves it (and the JSON key) so unchanged PR-40900 configs work. - """ + patch preserves it (and the JSON key) so unchanged PR-40900 configs work.""" assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 @@ -1014,8 +1011,7 @@ def test_store_sending_thread_skips_null_sparse_group_blocks(): def test_partial_tail_offload_skips_cap_omitted_mamba_group(): """A Mamba group omitted by the handoff cap must not fall back to the - connector's positional block table. - """ + connector's positional block table.""" from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1143,8 +1139,7 @@ def test_partial_tail_put_failure_activates_pressure_gate(): def test_normal_save_excludes_mamba_group_and_null_blocks(): """The positional normal save must never cover mamba chunks (align-mode block tables are not append-only) and must skip chunks whose block is the - reserved null block, for any group. - """ + reserved null block, for any group.""" from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1205,8 +1200,7 @@ def test_normal_save_excludes_mamba_group_and_null_blocks(): def test_block_aligned_snapshot_offload_uses_provided_block(): """A block-aligned hand-off (sparse-retention mamba checkpoint) uploads exactly the core-pinned block under the boundary-end hash key; the - positional block table entry for that chunk must be ignored. - """ + positional block table entry for that chunk must be ignored.""" store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) @@ -1234,8 +1228,7 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): def test_mixed_snapshot_and_sub_block_offloads(): """A retention snapshot and the prompt-end sub-block CoW tail can arrive in one hand-off; the sub-block path covers FA gap blocks but reads the - mamba boundary only from the CoW block, never positionally. - """ + mamba boundary only from the CoW block, never positionally.""" store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) @@ -1275,8 +1268,7 @@ def test_mixed_snapshot_and_sub_block_offloads(): def test_snapshot_offload_skips_null_handoff_block(): """A hand-off the core could not materialize (null block) carries no - committed state; persisting it would poison the boundary key. - """ + committed state; persisting it would poison the boundary key.""" store = MagicMock() thread = _make_partial_tail_send_thread(store) @@ -1604,8 +1596,7 @@ def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_segment(): # noqa: E501 """Without a preferred_segment the SendingThread still forwards a (default-constructed) ReplicateConfig so the C++ side always sees a - well-defined config object. - """ + well-defined config object.""" store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.return_value = [256, 256] @@ -2115,8 +2106,7 @@ def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch(): def test_recv_thread_skips_split_when_budget_holds_all_keys(): """PR-36 removed the count-based split trigger; with budget for 3 keys, - all three should be requested in a single call. - """ + all three should be requested in a single call.""" store = MagicMock() store.batch_get_into_multi_buffers.return_value = [256, 256, 256] thread = _make_store_recving_thread( @@ -2330,8 +2320,7 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled( monkeypatch, ): """enable_offload=False zeroes out the disk budget so we don't generate - redundant owner GET-RPCs. - """ + redundant owner GET-RPCs.""" store = MagicMock() store.setup.return_value = 0 _install_fake_mooncake(monkeypatch, store) @@ -3006,8 +2995,7 @@ def __init__(self) -> None: def test_store_sending_thread_clamps_token_len_to_lcm(): """Partial chunks past the last lcm boundary aren't stored — cache hits - are always lcm-aligned (mirrors HybridKVCacheCoordinator). - """ + are always lcm-aligned (mirrors HybridKVCacheCoordinator).""" store = MagicMock() store.batch_is_exist.return_value = [0, 0] store.batch_put_from_multi_buffers.return_value = [256, 256] @@ -3032,8 +3020,7 @@ def test_store_sending_thread_clamps_token_len_to_lcm(): def test_store_sending_thread_skips_when_token_len_below_lcm(): """Requests shorter than lcm_block_size cannot produce any aligned chunk, - so neither the existence check nor the put should be issued. - """ + so neither the existence check nor the put should be issued.""" from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec store = MagicMock() @@ -3823,8 +3810,7 @@ def test_lookup_full_hit_with_eagle_pops_once_not_twice(): """Eagle already leaves the last block for the drafter, so a full-prompt re-derivation must never fire for eagle-governed hits: firing would anchor the search one block lower and pop a second - block, regressing the hit by an extra producer boundary. - """ + block, regressing the hit by an extra producer boundary.""" worker = _make_bare_worker(block_size=16) worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( worker._kv_cache_groups, @@ -3981,8 +3967,7 @@ def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable(): the full hit into 48, whose SWA window needs the never-written block 1, so every load failed and the recompute re-entered the same lookup. The re-derivation must report that no stored boundary below the request end - is usable. - """ + is usable.""" from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec worker = _make_bare_worker(block_size=16) @@ -4004,8 +3989,7 @@ def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable(): def test_lookup_swa_single_group_returns_full_when_tail_window_present(): """Single-SWA, sliding_window=32 (= 2 blocks): producer stored only the tail. Coordinator-driven lookup returns full prefix even though the - pre-window blocks are absent. - """ + pre-window blocks are absent.""" from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec worker = _make_bare_worker(block_size=16) @@ -4401,8 +4385,7 @@ def test_config_defaults_to_embedded(): def test_config_pr40900_unchanged(tmp_path): """A literal PR-40900 config (no mode, no enable_offload, no preferred_segment) - parses without raising and resolves to embedded mode. - """ + parses without raising and resolves to embedded mode.""" config_path = _write_mooncake_config( tmp_path, { @@ -4510,8 +4493,7 @@ def test_config_zero_local_buffer(): def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch): """standalone-store + disk: global_segment_size=0, enable_offload=True, preferred_segment set. Assert setup() positional args, ReplicateConfig - wiring, and that the disk-offload buffer budget is allocated. - """ + wiring, and that the disk-offload buffer budget is allocated.""" store = MagicMock() store.setup.return_value = 0 fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store) @@ -4560,8 +4542,7 @@ def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch): def test_topology_embedded_cpu_only(tmp_path, monkeypatch): """Embedded + CPU-only: no mode key (defaults to embedded), global_segment_size>0, enable_offload absent, no preferred_segment. - This is the PR-40900 baseline recipe. - """ + This is the PR-40900 baseline recipe.""" store = MagicMock() store.setup.return_value = 0 fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store) @@ -4750,8 +4731,7 @@ def test_store_worker_close_swallows_store_errors(): def test_blob_block_hashes_wire_roundtrip(): """The lookup wire format sends a ``hash_len`` frame plus the raw hashes concatenated back-to-back; the server rebuilds them through a zero-copy - ``BlobBlockHashes`` view over the frame buffer. - """ + ``BlobBlockHashes`` view over the frame buffer.""" hashes = [BlockHash(bytes([i]) * 16) for i in range(5)] hash_len = len(hashes[0]) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index eae9f71ee3fa..a30b5caf4bbd 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -692,8 +692,7 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups): def test_resolve_host_ip_prefers_extra_config(): """An explicit ``host_ip`` in kv_connector_extra_config overrides get_ip() (so an external router can advertise a routable/internal address); an - absent or empty value falls back to get_ip(). - """ + absent or empty value falls back to get_ip().""" assert resolve_host_ip({"host_ip": "10.0.0.7"}) == "10.0.0.7" fallback = get_ip() @@ -747,8 +746,7 @@ def _read_scheduler( def test_hma_blocks_per_sw_two_groups(): """A Full + SlidingWindow config runs with HMA and computes block budgets - correctly - """ + correctly""" scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) assert scheduler._is_hma_required is True # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention. @@ -766,8 +764,7 @@ def test_hma_blocks_per_sw_two_groups(): def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): """_is_hma_required tracks both the KV cache groups and the --disable-hybrid-kv-cache-manager flag. When HMA is off, - get_exchange_clipped_blocks must be a no-op. - """ + get_exchange_clipped_blocks must be a no-op.""" config = ( _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config() ) @@ -780,8 +777,7 @@ def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): def test_non_sliding_window_hybrid_is_rejected(): """A hybrid group that is not sliding-window (e.g. chunked-local - attention) must fail closed rather than be silently mistransferred. - """ + attention) must fail closed rather than be silently mistransferred.""" full_spec = FullAttentionSpec( block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 ) @@ -815,8 +811,7 @@ def test_non_sliding_window_hybrid_is_rejected(): def test_token_count_basis_uses_full_attention_group(): """Chunked-prefill token counting must use an unclipped full-attention - group (blocks_per_sw == 0), not a clipped sliding-window group. - """ + group (blocks_per_sw == 0), not a clipped sliding-window group.""" scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) # Group 0 is the full-attention group assert scheduler._full_attn_group_idx == 0 @@ -825,8 +820,7 @@ def test_token_count_basis_uses_full_attention_group(): def test_get_exchange_clipped_blocks_clips_only_sw_group(): """get_exchange_clipped_blocks keeps the full attn group intact and clips the - sliding-window group to its window tail. - """ + sliding-window group to its window tail.""" scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) full = [10, 11, 12, 13, 14] sw = [20, 21, 22, 23, 24] @@ -837,8 +831,7 @@ def test_get_exchange_clipped_blocks_clips_only_sw_group(): def test_metadata_hma_block_ids_preserved_per_group(): """add_new_req stores per-group (BlockIds) block lists unchanged for both - read and write, so the hybrid group structure is retained. - """ + read and write, so the hybrid group structure is retained.""" metadata = MoRIIOConnectorMetadata() # Assume: diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py index 76a33cff82c2..e38ec3f7e9be 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -110,8 +110,7 @@ def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0" def build_decode_workers(cfg: PDConfig) -> list: """Decode workers that issue reads for one prefill instance. A TP decode instance reads from every tp rank; a DP decode instance from each dp-rank - worker. Reused across requests so per-worker round-robin state advances. - """ + worker. Reused across requests so per-worker round-robin state advances.""" if cfg.d_tp > 1: return [ make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0) @@ -126,8 +125,7 @@ def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter: """Drive the REAL _resolve_read_source over a fair input stream; tally which prefill GPU each read targets. The only modelled assumption is the proxy contract -- each owner dp-rank delivered equally (``for owner_dp in - range(p_dp)``); no proxy algorithm is reproduced. - """ + range(p_dp)``); no proxy algorithm is reproduced.""" workers = build_decode_workers(cfg) hits: Counter = Counter() for _ in range(rounds): diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index e6a78063be3d..fcda6232c6b2 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -555,8 +555,7 @@ def test_build_kv_connector_stats_with_empty_dict(self): def test_build_kv_connector_stats_reconstructs_nixl_stats(self): """Test that NixlConnector stats are properly reconstructed with - correct data. - """ + correct data.""" serialized_data = { "NixlConnector": { "transfer_duration": [1.5, 2.3], diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d02be5fd564e..c9d2dcf5f176 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -318,8 +318,7 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): """A remote-prefill request added with abort_immediately=True should be added to the scheduler's waiting queue then immediately aborted, so the NIXL connector's request_finished hook enqueues an empty recv to notify - the prefill instance to free its blocks. - """ + the prefill instance to free its blocks.""" from vllm.v1.request import RequestStatus scheduler = create_scheduler(create_vllm_config()) @@ -3146,8 +3145,7 @@ def test_split_read_failure_defers_report_until_last_handle( request — nor invalidate its blocks — while the sibling xfer is still in flight: a posted READ cannot be aborted and would DMA into blocks the scheduler could free and reuse. The report happens exactly once, when the - last handle is terminal. - """ + last handle is terminal.""" request_id = "split_read_partial_failure" err_handle, live_handle = 11, 22 connector, worker, wrapper = _make_split_read_connector( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index d79fcab69819..3f1038ff659c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -281,8 +281,7 @@ def test_update_state_after_alloc_tracks_cached_blocks_per_group(): """Hybrid SWA+FA requests can have different prefix-cache-hit counts per KV cache group. Each group's count must be tracked independently rather than collapsed to a single scalar, or DCP read-phase alignment - (local_num_computed_blocks) would misalign one of the groups. - """ + (local_num_computed_blocks) would misalign one of the groups.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( NixlPullConnectorScheduler, ) @@ -336,8 +335,7 @@ def test_full_local_hit_is_not_awaited_by_the_scheduler(): owed to the producer -- so the worker must not report it in finished_recving. Doing so trips `assert RequestStatus.is_finished` in _update_from_kv_xfer_finished, which only tolerates a finished recv for a - request that is parked or already done. - """ + request that is parked or already done.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( NixlPullConnectorScheduler, ) @@ -442,8 +440,7 @@ def test_apply_dcp_prefix_caching_matches_global_logical_positions( """_apply_dcp_prefix_caching must operate on logical block IDs (DCP interleaves at logical-block granularity) and always return equal-length local/remote slices, so a later, unconditional _apply_prefix_caching - trim is a guaranteed no-op rather than double-processing the read. - """ + trim is a guaranteed no-op rather than double-processing the read.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -842,8 +839,7 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit( def test_apply_prefix_caching_ssm_unpairable_slots_rejected(): """Local SSM slots can only exceed the remote ones by the position D recomputes itself. A larger excess means the lists aren't - position-aligned: fail loudly rather than transfer into wrong slots. - """ + position-aligned: fail loudly rather than transfer into wrong slots.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -1159,8 +1155,7 @@ def test_nixl_region_sort_preserves_pipeline_layer_order(): @pytest.mark.cpu_test def test_get_block_descs_ids_hybrid_ssm(): """Test _compute_desc_ids uses per-group strides for hybrid - FA+SSM when ratio=1 (no kernel block size mismatch). - """ + FA+SSM when ratio=1 (no kernel block size mismatch).""" from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec worker = _make_mock_worker_for_desc_ids( @@ -1244,8 +1239,7 @@ def test_get_block_descs_ids_uses_per_region_pool_capacity(): @pytest.mark.cpu_test def test_get_block_descs_ids_kernel_block_mismatch(): """Test _compute_desc_ids uses different strides for FA - (kernel blocks) vs SSM (logical blocks) when ratio > 1. - """ + (kernel blocks) vs SSM (logical blocks) when ratio > 1.""" from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec ratio = 4 @@ -1276,8 +1270,7 @@ def test_get_block_descs_ids_kernel_block_mismatch(): def test_get_block_descs_ids_hetero_block_size_hybrid(): """With a block-size ratio, FA desc ids are ratio-expanded while SSM desc ids keep the unexpanded logical stride (state blocks are never - sub-split). - """ + sub-split).""" from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec worker = _make_mock_worker_for_desc_ids( @@ -1316,8 +1309,7 @@ def _bind_worker_method(worker, name): @pytest.mark.cpu_test def test_map_block_ids_for_block_size_ratio_hybrid(): """Attention groups expand to remote granularity and clip to the remote - coverage; mamba state blocks pass through 1:1. - """ + coverage; mamba state blocks pass through 1:1.""" from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( @@ -1348,8 +1340,7 @@ def test_map_block_ids_for_block_size_ratio_hybrid(): @pytest.mark.cpu_test def test_post_process_zeroes_untransferred_tail(): """The untransferred sub-blocks of the last local block are zeroed on - receive; mamba state caches are untouched by the attention permute. - """ + receive; mamba state caches are untouched by the attention permute.""" from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( @@ -1392,8 +1383,7 @@ def test_post_process_zeroes_untransferred_tail(): @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM - groups with different block counts (kernel mismatch active). - """ + groups with different block counts (kernel mismatch active).""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -1474,8 +1464,7 @@ def _make_fake_kv_cache_manager(): @pytest.mark.cpu_test def test_zeroing_block_ids_cover_only_loaded_attention_blocks(): """Only zero-recorded (attention) groups contribute, sliced to the - externally-loaded token range; Mamba state blocks are never zeroed. - """ + externally-loaded token range; Mamba state blocks are never zeroed.""" manager = _make_fake_kv_cache_manager() # Tokens [0, 16) are locally cached; the load covers tokens [16, 56). @@ -1503,8 +1492,7 @@ def take_new_block_ids(self): @pytest.mark.cpu_test def test_failed_load_rezeroes_unwritten_skipped_blocks(): """A failed async load leaves zeroing-skipped blocks unwritten beyond - the valid prefix; they must be zeroed before local recompute. - """ + the valid prefix; they must be zeroed before local recompute.""" from unittest.mock import MagicMock from vllm.v1.core.sched.scheduler import Scheduler @@ -1964,8 +1952,7 @@ def test_logical_to_kernel_block_ids_with_remote_ratio( def test_exchange_clipped_blocks_ssm_single_state(): """In single-state cache modes, SSM lists are reduced to the running state slot: speculative scratch slots, null placeholders and the previous - step's state carry nothing. Attention groups pass through untouched. - """ + step's state carry nothing. Attention groups pass through untouched.""" sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) sched.blocks_per_sw = [0, 0] sched._ssm_spec_blocks = [None, 2] @@ -1992,8 +1979,7 @@ def test_exchange_clipped_blocks_ssm_single_state(): @pytest.mark.cpu_test def test_exchange_clipped_blocks_ssm_positional_states(): """In "all" mode every position holds a state, so only the speculative - slots go; placeholders stay to keep the list position-indexed. - """ + slots go; placeholders stay to keep the list position-indexed.""" sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True) sched.blocks_per_sw = [0, 0] sched._ssm_spec_blocks = [None, 2] @@ -2010,8 +1996,7 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed MambaSpec) groups whose layers share the same HMA tensors, with a mamba-aligned unified page and an MLA kernel block smaller than the - logical block. - """ + logical block.""" from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -2151,8 +2136,7 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): """Hybrid MLA+KDA registration: HMA tensors shared by both layer types must be flagged as MLA regions even when a KDA layer registers them first, expose TP-independent kernel-granularity block lens, and build - FA + mamba descriptors for every region. - """ + FA + mamba descriptors for every region.""" from unittest.mock import MagicMock from vllm.config import set_current_vllm_config @@ -2253,8 +2237,7 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): def test_push_write_hybrid_mla_replicates_attention(): """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be written to every covered D rank (replicated MLA latent) while SSM state - is written per-rank through the split handles. - """ + is written per-rank through the split handles.""" import threading from collections import defaultdict from unittest.mock import MagicMock diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index dab6432b5ca3..a4bfa9ab83ae 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -24,8 +24,7 @@ class _RecordingNixl: """Minimal NIXL wrapper stand-in that records descriptor lists and - prepared transfers so tests can resolve desc ids to byte ranges. - """ + prepared transfers so tests can resolve desc ids to byte ranges.""" def __init__(self, *args, **kwargs): self.dlists: dict[int, np.ndarray] = {} @@ -594,8 +593,7 @@ def _register_remote_agents(worker, metadata, tp_size): def _owned_byte_ranges(worker, group_logical_ids): """Byte ranges owned by a request: for each HMA region tensor, every - logical block id of every group maps to one unified page. - """ + logical block id of every group maps to one unified page.""" unified_page = worker._test_unified_page bases = [t.data_ptr() for t in worker._test_tensors] owned = [] @@ -633,8 +631,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and TP-sharded KDA state. Every local descriptor of the request's reads must - stay within its own blocks. - """ + stay within its own blocks.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -701,8 +698,7 @@ def _resolve( """Resolve a desc id to (region, kind, token_start) where kind is 'attn' (desc-page sized, sub-block-aligned, in the request's attention blocks) or 'mamba'. token_start is the request-relative token offset, so local - and remote are comparable even when their kernel blocks differ in size. - """ + and remote are comparable even when their kernel blocks differ in size.""" addr, length, _ = (int(x) for x in desc_arr[int(idx)]) for region, base in enumerate(bases): off = addr - base @@ -896,8 +892,7 @@ def _run_hetero_case( def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): """Sweep prompt lengths across block-boundary residues for several hetero-ppl geometries; assert neighbor-safety, token alignment, and - coverage of every transferred kernel block. - """ + coverage of every transferred kernel block.""" _run_hetero_case( local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens ) @@ -916,8 +911,7 @@ def test_hetero_ppl_with_block_size_ratio(num_tokens): 2). The transfer is clipped at remote sub-block granularity by the pairing and front-trimmed by _apply_prefix_caching, so the untransferred tail can span both a partial block and whole blocks — - the case each of the two former zeroing paths handled only half of. - """ + the case each of the two former zeroing paths handled only half of.""" _run_hetero_case( local_block=24, kernel=8, @@ -966,8 +960,7 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): decode (local) logical block 5760 / kernel 64 (ppl=90), prefill (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with - replicated MLA and 8-way TP-sharded KDA state. - """ + replicated MLA and 8-way TP-sharded KDA state.""" _run_hetero_case( local_block=5760, kernel=64, @@ -981,8 +974,7 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): """The MLA per-token page is TP-independent, so kernel block lengths differing by anything other than the block-size ratio must fail the - handshake loudly rather than transfer at mismatched geometry. - """ + handshake loudly rather than transfer at mismatched geometry.""" worker = _make_mla_hybrid_worker( local_block_size=12, kernel_block_size=4, num_logical_blocks=8 ) @@ -1166,8 +1158,7 @@ def test_csa_linear_registration_discovers_shared_regions(): def test_csa_linear_scratch_descs_follow_the_pages_they_overlay(): """The scratch ring is addressed by the regions it actually registered in, not by inferring that it must sit on the MLA pages: overlaying the main KV - pages instead moves its descriptors with it. - """ + pages instead moves its descriptors with it.""" worker = _make_csa_linear_ple_worker(scratch_aliases="main_kv") # Same MLA regions as before, but the ring now lives in the non-MLA ones. @@ -1301,8 +1292,7 @@ def test_csa_linear_remote_ple_is_copied_whole(): def _make_ring_worker(): """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a - model whose sliding-window KV lives in per-request rings. - """ + model whose sliding-window KV lives in per-request rings.""" from unittest.mock import MagicMock from vllm.config import set_current_vllm_config diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 964cadc3ed76..7ce84bf685b1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -115,8 +115,7 @@ def get_unhashed_block_ids_all_groups(self) -> tuple[list[int], ...]: def _stub_sw_clipping(scheduler) -> None: """Make ``get_exchange_clipped_blocks`` a passthrough so tests don't - need the full sliding-window machinery. - """ + need the full sliding-window machinery.""" scheduler.get_exchange_clipped_blocks = lambda block_ids, clip_ssm=True: block_ids @@ -174,8 +173,7 @@ def test_d_side_update_state_after_alloc_stages_registration(self): def test_p_side_request_finished_stages_blocks(self): """P scheduler pushes blocks into both _finished_request_blocks (lease) - and _newly_finished_push_blocks (metadata for next step). - """ + and _newly_finished_push_blocks (metadata for next step).""" sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -194,8 +192,7 @@ def test_p_side_request_finished_stages_blocks(self): def test_build_connector_meta_drains_both_sides(self): """meta.push_registrations and meta.push_finished_blocks are filled - from the staging dicts and the staging dicts are cleared. - """ + from the staging dicts and the staging dicts are cleared.""" sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -290,8 +287,7 @@ def test_update_connector_output_clears_lease_and_watchdog(self): def test_registration_watchdog_expires(self, caplog): """Stale D registrations whose deadline has passed are dropped at - ``build_connector_meta`` time. - """ + ``build_connector_meta`` time.""" # Watchdog logs a WARNING when it drops the stale entry; that's # what this test is verifying, so silence it in the test report. caplog.set_level( @@ -332,8 +328,7 @@ def test_registration_watchdog_expires(self, caplog): class _StubWriterWorker(NixlPushConnectorWorker): """Construct a worker without invoking ``__init__`` so we can drive - the matching/notif logic without bringing up NIXL or torch. - """ + the matching/notif logic without bringing up NIXL or torch.""" @classmethod def fresh(cls) -> _StubWriterWorker: @@ -439,8 +434,7 @@ class TestPushWriterRegSend: """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of the P engine, and the agent map is read under ``_handshake_lock`` so the send can't race the handshake done-callback's insert or a concurrent - stale-engine eviction (X1). - """ + stale-engine eviction (X1).""" def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]: w = _StubWriterWorker.fresh() @@ -485,8 +479,7 @@ def test_queues_failed_request_when_engine_not_handshaked(self): def test_agent_read_is_serialized_with_handshake_lock(self): """While another thread holds ``_handshake_lock`` (as the handshake done-callback does while writing ``_remote_agents``), the writer's - check-then-send must not run ahead of the critical section. - """ + check-then-send must not run ahead of the critical section.""" w, failures = self._reg_send_worker() w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"} rd = _registration_data("req-1") @@ -596,8 +589,7 @@ def test_handle_push_reg_drops_malformed(self, caplog): class TestPushWriterStartLoadKv: def test_finished_blocks_inbox_matches_stashed_registration(self): """Run the writer-loop's finished-blocks drain against a - pre-populated _pending_d_registrations entry. - """ + pre-populated _pending_d_registrations entry.""" w = _StubWriterWorker.fresh() w._pending_d_registrations["req-C"] = _registration_data("req-C") @@ -622,8 +614,7 @@ def test_finished_blocks_inbox_matches_stashed_registration(self): def test_start_load_kv_enqueues_to_writer(self): """``start_load_kv`` should hand registrations + finished blocks - to the writer queues without doing matching itself. - """ + to the writer queues without doing matching itself.""" w = _StubWriterWorker.fresh() # Stub heartbeats to a no-op; tests don't exercise the heartbeat # path here. @@ -681,8 +672,7 @@ def test_do_start_push_kv_defers_then_writes_when_handshake_ready(): NIXL op from the writer or the executor callback); once it resolves the request is re-queued on ``_deferred_push_inbox`` with the wake set; and on re-drive with the handshake ready the WRITE is issued with a correct - ReqMeta. - """ + ReqMeta.""" w = _StubWriterWorker.fresh() w._logical_to_kernel_block_ids = lambda x, ratio: x xfer_calls: list[dict[str, Any]] = [] @@ -722,8 +712,7 @@ def test_do_start_push_kv_defers_then_writes_when_handshake_ready(): def test_do_start_push_kv_drops_request_on_handshake_failure(): """Handshake raises: the request is dropped (not re-queued, no WRITE) and the failure is logged. Blocks are reclaimed by the lease/watchdog, matching - the old blocking behaviour. - """ + the old blocking behaviour.""" w = _StubWriterWorker.fresh() w._logical_to_kernel_block_ids = lambda x, ratio: x xfer_calls: list[dict[str, Any]] = [] @@ -745,8 +734,7 @@ def test_do_start_push_kv_drops_request_on_handshake_failure(): def test_writer_loop_drains_deferred_push_inbox(): """The writer loop drains ``_deferred_push_inbox`` and re-drives - ``_do_start_push_kv`` for each entry (event-driven, no polling). - """ + ``_do_start_push_kv`` for each entry (event-driven, no polling).""" w = _StubWriterWorker.fresh() w.nixl_wrapper = MagicMock() w.nixl_wrapper.get_new_notifs.return_value = {} @@ -779,8 +767,7 @@ def _tracked(rid, blocks, rd): def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker: """A push worker wired to drive the base ``_ensure_handshake`` eviction - path (``_evict_stale_engines`` + ``_cleanup_remote_engine``). - """ + path (``_evict_stale_engines`` + ``_cleanup_remote_engine``).""" w = _StubWriterWorker.fresh() w._engine_ttl = engine_ttl w._engine_last_active = {} @@ -802,8 +789,7 @@ def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker: def test_active_push_refreshes_engine_last_active(): """The base ``_ensure_handshake`` refreshes liveness only on a new handshake, so an active push to an already-connected engine must refresh - ``_engine_last_active`` itself or it is reaped mid-stream (S1). - """ + ``_engine_last_active`` itself or it is reaped mid-stream (S1).""" w = _eviction_worker(engine_ttl=3600.0) # Already-connected D -> _ensure_handshake returns None, WRITE runs inline. w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"} @@ -819,8 +805,7 @@ def test_active_push_refreshes_engine_last_active(): def test_stale_engine_evicted_on_push(): """A push drives ``_ensure_handshake`` -> ``_evict_stale_engines``, so a D engine silent past its TTL is reaped -- no _remote_agents / NIXL agent leak - across D scale up/down (S1). - """ + across D scale up/down (S1).""" w = _eviction_worker(engine_ttl=30.0) w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} w._engine_last_active["D-old"] = time.perf_counter() - 10_000.0 @@ -837,8 +822,7 @@ def test_stale_engine_evicted_on_push(): def test_cleanup_remote_engine_pops_under_handshake_lock(): """Regression (X1): eviction must not remove ``_remote_agents`` entries while another thread is inside the lock's critical section (e.g. the - push writer reading the map to send a registration). - """ + push writer reading the map to send a registration).""" w = _eviction_worker(engine_ttl=30.0) w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} @@ -863,8 +847,7 @@ def _cleanup(): class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained - on the engine main thread inside ``_get_new_notifs``. - """ + on the engine main thread inside ``_get_new_notifs``.""" w = _StubWriterWorker.fresh() # Pretend the writer thread already forwarded a completion notif # for a request whose KV is being received. @@ -910,8 +893,7 @@ def test_get_transfer_results_evicts_completed_state(self): @staticmethod def _pollable_worker() -> _StubWriterWorker: """Stub worker with the extra base-worker state the real - ``get_transfer_results`` needs to poll ``_sending_transfers``. - """ + ``get_transfer_results`` needs to poll ``_sending_transfers``.""" w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() w.nixl_wrapper = MagicMock() @@ -953,8 +935,7 @@ def test_failed_push_send_is_not_reported_done_sending(self): @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)]) def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma): """P-side sends are completed by consumer notifs / lease expiry, - not by local WRITE-handle completion. - """ + not by local WRITE-handle completion.""" w = self._pollable_worker() w.pp_size = pp_size w._is_hma_required = is_hma @@ -1025,8 +1006,7 @@ def test_update_state_after_alloc_zero_external_tokens_does_not_register(self): def test_request_finished_unfinished_status_does_not_stage(self): """If a request is still RUNNING, request_finished must not stash - blocks for the worker (no push needed). - """ + blocks for the worker (no push needed).""" sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1043,8 +1023,7 @@ def test_request_finished_unfinished_status_does_not_stage(self): def test_request_finished_empty_blocks_does_not_arm_lease(self): """Empty block-id groups should still complete cleanly without - arming the lease/finished maps. - """ + arming the lease/finished maps.""" sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1059,8 +1038,7 @@ def test_request_finished_empty_blocks_does_not_arm_lease(self): def test_update_connector_output_unknown_request_is_noop(self): """Idempotent cleanup: clearing a request that was never staged - must not raise or mutate other state. - """ + must not raise or mutate other state.""" sched = make_nixl_push_scheduler() _stub_sw_clipping(sched) @@ -1093,8 +1071,7 @@ def test_pop_matching_finished_blocks_returns_none_when_empty(self): def test_pop_matching_registration_no_match_when_base_ids_differ(self): """A registration whose base id (after stripping the random suffix) - does NOT match the lookup request_id must not be popped. - """ + does NOT match the lookup request_id must not be popped.""" w = _StubWriterWorker.fresh() # Two unrelated requests: different base UUIDs, so stripping the # trailing ``-<8 hex>`` suffix still yields different base ids. @@ -1110,8 +1087,7 @@ def test_pop_matching_registration_no_match_when_base_ids_differ(self): def test_handle_push_reg_with_non_dict_payload_is_dropped(self, caplog): """msgpack-encoded non-dict payload (e.g. a list) should be - dropped without raising. - """ + dropped without raising.""" caplog.set_level( logging.CRITICAL, logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), @@ -1139,8 +1115,7 @@ def test_handle_push_reg_with_non_string_request_id_is_dropped(self, caplog): def test_handle_push_reg_idempotent_for_same_request_id(self): """Receiving the same PUSH_REG twice (e.g. P retries after a - flake) keeps the entry staged exactly once and never fires. - """ + flake) keeps the entry staged exactly once and never fires.""" w = _StubWriterWorker.fresh() notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode( _registration_data("req-dup") @@ -1155,8 +1130,7 @@ def test_get_transfer_results_enqueues_each_completed_request(self): """``get_transfer_results`` must enqueue every completed request in ``done_sending`` so the writer can drop stale matching state. Unlike the happy-path test, this verifies the *cardinality*: N - completed requests -> N evictions, in order. - """ + completed requests -> N evictions, in order.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], @@ -1179,8 +1153,7 @@ def test_get_transfer_results_enqueues_each_completed_request(self): def test_empty_transfer_results_do_not_enqueue_eviction(self): """If there's nothing newly done, no eviction should be enqueued. The wake event IS still set because ``get_finished`` always wakes - the writer to drain notifs. - """ + the writer to drain notifs.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], @@ -1221,8 +1194,7 @@ def test_failed_push_transfer_does_not_evict_writer_state(self): def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): """A completion notif for a request the worker doesn't know - about should be logged but not crash. - """ + about should be logged but not crash.""" caplog.set_level( logging.CRITICAL, logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), @@ -1254,8 +1226,7 @@ def test_start_load_kv_with_empty_metadata_is_noop(self): def test_get_new_notifs_extends_lease_on_heartbeat(self): """``HB:`` notifs forwarded by the writer thread must extend the leases of tracked P-side requests on the engine main thread, and - ignore request IDs that aren't being tracked. - """ + ignore request IDs that aren't being tracked.""" w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() # _handle_heartbeat reads ``self._lease_extension`` (set in the @@ -1291,13 +1262,11 @@ def test_get_new_notifs_extends_lease_on_heartbeat(self): class TestPushPipelineParallel: """PP-sharded producer: per-stage completion counting, pp_size plumbing, - and remote-region slicing. - """ + and remote-region slicing.""" def test_completion_waits_for_one_notif_per_pp_stage(self): """Each PP stage WRITEs its own layers and sends one notif; D must - collect pp_size notifs before reporting the recv done. - """ + collect pp_size notifs before reporting the recv done.""" w = _StubWriterWorker.fresh() w.transfer_topo = MagicMock() request_id = "req-pp-2" @@ -1318,8 +1287,7 @@ def test_completion_waits_for_one_notif_per_pp_stage(self): def test_req_meta_reads_pp_size_from_kv_transfer_params(self): """D learns the producer's pp_size from kv_transfer_params (forwarded - by the proxy) and defaults to 1 when absent. - """ + by the proxy) and defaults to 1 when absent.""" metadata = NixlConnectorMetadata() params = { "remote_block_ids": ([0],), @@ -1342,8 +1310,7 @@ def test_add_remote_agent_slices_remote_regions_to_local_pp_window(self): this worker holds only a contiguous layer slice; add_remote_agent trims the remote region list to [offset : offset + num_local_regions] before building descriptors. We stop right after the slice via a - sentinel on the next collaborator call. - """ + sentinel on the next collaborator call.""" block_len = 4096 * 16 w = _StubWriterWorker.fresh() # seeds writer-thread state for teardown w.pp_size = 2 @@ -1385,8 +1352,7 @@ class TestPushWriterMlaReplication: """MLA latent KV is replicated across D's TP ranks, so when D_TP > P_TP (``tp_ratio < 0``) the producer must *WRITE* the latent into every D rank it handshook -- not write one and merely notify the rest, which would - leave the un-written ranks decoding against stale KV. - """ + leave the un-written ranks decoding against stale KV.""" @staticmethod def _mla_worker_writing_to(d_ranks): @@ -1611,8 +1577,7 @@ def _written_block_ids(w) -> tuple[list[int], list[int]]: def test_partial_prefix_hit_end_trims_producer_blocks(self): """D registered only its 2 uncomputed suffix blocks; P finished the - full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots. - """ + full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots.""" w, _ = self._worker_driving_xfer() reg = _registration_data("req-pc", local_block_ids=([500, 501],)) diff --git a/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py index bab9ed6e0fff..da8be07157d8 100644 --- a/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py +++ b/tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py @@ -171,8 +171,7 @@ def test_gpu_memory_nixl_hma(model_name, sw_size): @pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"]) def test_gpu_memory_no_nixl_baseline(model_name): """Same workload without NixlConnector. Comparing driver-level memory - between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact. - """ + between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact.""" from vllm import LLM, SamplingParams from vllm.distributed.parallel_state import cleanup_dist_env_and_memory diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index 7044f9f0be52..bdab91a21ffe 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -213,8 +213,7 @@ def _multi_connector_config(swa_enabled: bool = False): def test_nixl_wins_load_over_cpu_offload(): """When NixlConnector (index 0) has matched tokens from a remote prefill, it should win the load: Nixl metadata tracks the recv while CPU offload metadata has no load - scheduled. - """ + scheduled.""" vllm_config, kv_cache_config = _multi_connector_config() scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) mc = scheduler.connector @@ -251,8 +250,7 @@ def test_nixl_wins_load_over_cpu_offload(): @patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper) def test_cpu_offload_wins_when_nixl_has_no_match(): """When NixlConnector returns 0 matched tokens and SimpleCPUOffloadConnector has a - CPU cache hit, the CPU offload connector (index 1) wins the load. - """ + CPU cache hit, the CPU offload connector (index 1) wins the load.""" vllm_config, kv_cache_config = _multi_connector_config() scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) mc = scheduler.connector @@ -311,8 +309,7 @@ def test_cpu_offload_wins_when_nixl_has_no_match(): def test_request_finished_no_async_save(swa_enabled: bool): """A normal request (no P/D) produces no async save from either connector. MultiConnector returns (False, None) via both request_finished and - request_finished_all_groups, and cleans up _requests_to_connector. - """ + request_finished_all_groups, and cleans up _requests_to_connector.""" from vllm.v1.request import RequestStatus vllm_config, kv_cache_config = _multi_connector_config(swa_enabled=swa_enabled) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 1d99ac0d1148..db7ecca62148 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -321,8 +321,7 @@ def test_cpu_offloading( @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_cpu_offloading_metrics() -> None: """Verify that offloading Prometheus metrics (new flat and deprecated - labeled) are emitted after stores and loads. - """ + labeled) are emitted after stores and loads.""" extra_config: dict = { "cpu_bytes_to_use": 500 << 20, "block_size": CPU_BLOCK_SIZES, @@ -511,8 +510,7 @@ def test_tiering_offloading() -> None: def test_fs_tiering_offloading(tmp_path) -> None: """Tests OffloadingConnector with TieringOffloadingSpec - + fs secondary tier. - """ + + fs secondary tier.""" extra_config: dict = { "cpu_bytes_to_use": 1 << 30, "block_size": CPU_BLOCK_SIZES, diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 9ee93c889446..4f82e828d3f9 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -704,8 +704,7 @@ def test_async_load_reserves_blocks_for_inflight(): def test_async_loads_both_admitted_when_pool_fits(): """Sanity: with a pool large enough, the reservation gate admits both async - loads (it is not over-conservative). - """ + loads (it is not over-conservative).""" vllm_config = create_vllm_config() BLOCK_SIZE = vllm_config.cache_config.block_size scheduler = create_scheduler(vllm_config, num_blocks=64) diff --git a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py index f9f790232925..af4c7858fcad 100644 --- a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py +++ b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py @@ -129,8 +129,7 @@ def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig: def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None: """The ring group's capacity is not a token granularity, it holds no - hashed blocks, and it never participates in CPU hits. - """ + hashed blocks, and it never participates in CPU hits.""" vllm_config = create_vllm_config( kv_connector="SimpleCPUOffloadConnector", kv_role="kv_both", diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 201950705a73..3ddc568a524d 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -235,8 +235,7 @@ def test_fa_and_ssm_different_split_factors(self): def test_hetero_block_size_splits(self): """With a block-size ratio, single-source FA sub-block descs pass - through whole; SSM descs are unexpanded and split per source. - """ + through whole; SSM descs are unexpanded and split per source.""" plan = TPMapping( source_ranks_per_group=((0,), (0, 1)), all_source_ranks=(0, 1), @@ -271,8 +270,7 @@ def test_hetero_block_size_splits(self): def test_hetero_block_size_head_sharded_asserts(self): """Head-sharded FA reads (multiple FA sources) are incompatible with - a block-size mismatch and must fail loudly. - """ + a block-size mismatch and must fail loudly.""" plan = TPMapping( source_ranks_per_group=((0, 1), (0, 1)), all_source_ranks=(0, 1), diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py index 35897a9236ae..338ec18d0e67 100644 --- a/tests/v1/kv_offload/cpu/policies/test_factory.py +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -15,8 +15,7 @@ class _DummyCachePolicy(CachePolicy): """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded by module path, so it must be importable at module scope (mirrors - tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec). - """ + tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec).""" def __init__(self, cache_capacity: int) -> None: self.cache_capacity = cache_capacity @@ -55,8 +54,7 @@ class TestCachePolicyFactory: def test_pre_registered_policies_can_be_imported(self): """If someone moves a policy module but forgets to update - factory.py, CI fails. - """ + factory.py, CI fails.""" for name in CachePolicyFactory._registry: cls = CachePolicyFactory._registry[name]() assert issubclass(cls, CachePolicy) @@ -90,8 +88,7 @@ def test_dynamic_load_via_cache_policy_module_path(self): register_cache_policy() call -- this is how external projects integrate a custom CachePolicy without forking/patching vLLM. Mirrors tests/v1/kv_offload/test_factory.py's - test_dynamic_load_via_spec_module_path. - """ + test_dynamic_load_via_spec_module_path.""" policy_cls = CachePolicyFactory.get_cache_policy_cls( "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory" ) @@ -99,8 +96,7 @@ def test_dynamic_load_via_cache_policy_module_path(self): def test_manager_resolves_policy_via_module_path(self): """End-to-end: CPUOffloadingManager resolves an unregistered policy - purely from cache_policy_module_path. - """ + purely from cache_policy_module_path.""" manager = CPUOffloadingManager( num_chunks=4, cache_policy="_DummyCachePolicy", @@ -110,7 +106,6 @@ def test_manager_resolves_policy_via_module_path(self): def test_unregistered_policy_without_module_path_raises(self): """eviction_policy not in registry + no cache_policy_module_path -> - ValueError, same shape as the OffloadingSpecFactory error path. - """ + ValueError, same shape as the OffloadingSpecFactory error path.""" with pytest.raises(ValueError, match="Unknown cache policy"): CachePolicyFactory.get_cache_policy_cls("nonexistent", None) diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py index d4e383801cfd..4594e1eb0924 100644 --- a/tests/v1/kv_offload/cpu/test_canonical_layout.py +++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py @@ -126,8 +126,7 @@ def _canonical_handler(gpu_tensor, cpu_tensor, mapping, gpu_to_cpu): def test_gpu_roundtrip_assembles_canonical_page_across_ranks(): """Two TP2-style rank handlers must scatter into one shared canonical CPU page such that each rank reloads bit-exact and a whole-page (TP1) - reader sees both shards — the cross-topology contract. - """ + reader sees both shards — the cross-topology contract.""" torch.manual_seed(0) num_blocks = 4 gpu_rank = [ @@ -204,8 +203,7 @@ def _nhd_shard_mapping(tp: int, rank: int) -> CanonicalPageMapping: def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: """This rank's local page rows out of the (blocks, 2, tokens, heads, - head_bytes) ground truth. - """ + head_bytes) ground truth.""" local_heads = _TOTAL_HEADS // tp shard = full_kv[:, :, :, rank * local_heads : (rank + 1) * local_heads, :] return shard.reshape(full_kv.shape[0], -1).contiguous() @@ -216,8 +214,7 @@ def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int): """KV written at one tp must be readable at another: writer ranks scatter head shards into a shared canonical region, reader ranks gather their own - shards, and every reader must see the writers' ground-truth bytes. - """ + shards, and every reader must see the writers' ground-truth bytes.""" torch.manual_seed(0) num_blocks = 3 full_kv = torch.randint( diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 1d6a78f9e7ef..4f22d122dbec 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -167,8 +167,7 @@ def _mp_race_construct_and_write( ) -> None: """Race to construct a SharedOffloadRegion, write fill_value, then wait for the parent's cleanup signal before tearing down. The wait gives the - parent a window to read the raw mmap before the creator removes the file. - """ + parent a window to read the raw mmap before the creator removes the file.""" try: region = SharedOffloadRegion( engine_id=engine_id, @@ -197,8 +196,7 @@ def _mp_barrier_construct_and_hold( replicated: bool = False, ) -> None: """Construct with a real cross-process barrier, write, then hold the - mapping (no cleanup) until the parent SIGKILLs this process. - """ + mapping (no cleanup) until the parent SIGKILLs this process.""" try: populate_calls = 0 get_populate_write_fn = region_module._get_populate_write_fn @@ -252,8 +250,7 @@ def iid(): def test_create_next_worker_view_shape_and_stride(iid): """Returned tensor must have shape (num_chunks, tensor_page_size) and - stride (row_stride, 1) where row_stride = cpu_page_size * num_workers. - """ + stride (row_stride, 1) where row_stride = cpu_page_size * num_workers.""" with _region(iid, num_chunks=4, cpu_page_size=2 * PAGE_SIZE) as r: t = r.create_next_worker_view(PAGE_SIZE) assert t.shape == (4, PAGE_SIZE) @@ -293,8 +290,7 @@ def test_create_next_worker_view_row_stride_with_multiple_workers(iid): def test_create_next_worker_view_cursor_advances(iid): """Each create_next_worker_view call must advance _worker_offset by - tensor_page_size. - """ + tensor_page_size.""" with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r: assert r._worker_offset == 0 r.create_next_worker_view(PAGE_SIZE) @@ -351,8 +347,7 @@ def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid): def test_create_next_worker_view_write_visible_in_raw_mmap(iid): """Writes into a create_next_worker_view view must appear at the correct - raw mmap offset - """ + raw mmap offset""" with _region(iid, num_chunks=4) as r: t = r.create_next_worker_view(PAGE_SIZE) t[2, :] = 42 # write to chunk row 2 @@ -385,8 +380,7 @@ def test_create_next_worker_view_multi_tensor_layout(iid): def test_create_next_worker_view_multiprocess_slots(iid): """Each worker process calls create_next_worker_view and writes distinct data; - the parent verifies each slot lands at the correct interleaved offset. - """ + the parent verifies each slot lands at the correct interleaved offset.""" num_workers = 2 num_chunks = 4 @@ -671,8 +665,7 @@ def test_multi_worker_race_shared_memory_visible(iid): def test_multiprocess_race_construct_and_write(iid): """N processes race to construct the same SharedOffloadRegion, each writes - fill_value = rank+1 into their slot; parent verifies interleaved layout. - """ + fill_value = rank+1 into their slot; parent verifies interleaved layout.""" num_workers = 4 num_chunks = 3 @@ -798,8 +791,7 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): def test_cleanup_after_create_next_worker_view_releases_mmap(iid): """cleanup() must close the mmap even after create_next_worker_view was called. create_next_worker_view returns a view that shares storage with _base; both must be - released before mmap.close() can succeed. - """ + released before mmap.close() can succeed.""" r = _make_region(iid) mmap_obj = r.mmap_obj @@ -978,8 +970,7 @@ def test_ftruncate_failure_cleans_up_creator(monkeypatch): def test_backing_file_unlinked_after_barrier(iid): """The file must exist until the barrier releases (late joiners need the - name) and be gone right after, so no exit path can leak it. - """ + name) and be gone right after, so no exit path can leak it.""" path = f"/dev/shm/vllm_offload_{iid}.mmap" seen_at_barrier = [] @@ -1001,8 +992,7 @@ def test_backing_file_unlinked_after_barrier(iid): def test_barrier_failure_unlinks_creator_and_raises(iid): """A failed rendezvous must remove the creator's file and re-raise, not - leave a stub that wedges the next start in _wait_for_file_size. - """ + leave a stub that wedges the next start in _wait_for_file_size.""" path = f"/dev/shm/vllm_offload_{iid}.mmap" with pytest.raises(RuntimeError, match="peer died"): _make_region(iid, barrier=MagicMock(side_effect=RuntimeError("peer died"))) @@ -1012,8 +1002,7 @@ def test_barrier_failure_unlinks_creator_and_raises(iid): @pytest.mark.parametrize("replicated", [False, True]) def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): """With a real cross-process barrier every worker maps one shared inode, - the name is gone while they run, and SIGKILL leaks nothing. - """ + the name is gone while they run, and SIGKILL leaks nothing.""" num_workers = 2 path = f"/dev/shm/vllm_offload_{iid}.mmap" ctx = get_mp_context() @@ -1052,8 +1041,7 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): @pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"]) def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase): """A worker that dies before the rendezvous must still arrive at the - barrier, or its peers block in the collective until it times out. - """ + barrier, or its peers block in the collective until it times out.""" failure = MagicMock(side_effect=ValueError("Insufficient space")) memory_check = failure if failure_phase == "memory" else None if failure_phase == "shm": @@ -1084,8 +1072,7 @@ def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_p def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): """A creator that fails after sizing the file must drop it before arriving - at the barrier, so the next start does not land on a stale file. - """ + at the barrier, so the next start does not land on a stale file.""" path = f"/dev/shm/vllm_offload_{iid}.mmap" monkeypatch.setattr("mmap.mmap", MagicMock(side_effect=OSError("mmap"))) seen_at_barrier = [] @@ -1098,8 +1085,7 @@ def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): def test_barrier_release_failure_keeps_original_error(iid, monkeypatch): """When releasing the peers fails too, the setup error that explains the - failure must be the one that propagates. - """ + failure must be the one that propagates.""" import vllm.v1.kv_offload.cpu.shared_offload_region as region monkeypatch.setattr( diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 2e8e9ff4be69..aceec1cd8fa1 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -48,8 +48,7 @@ def _remote_prefiller_kv_params( kv_request_id: str = "req-1", ) -> dict: """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries - kv_request_id + remote_host + remote_port. - """ + kv_request_id + remote_host + remote_port.""" return { "remote_prefiller": { "kv_request_id": kv_request_id, @@ -65,8 +64,7 @@ def _remote_kv_source_kv_params( kv_request_id: str = "req-1", ) -> dict: """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has - the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port). - """ + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" return { "remote_kv_source": { "kv_request_id": kv_request_id, @@ -78,8 +76,7 @@ def _remote_kv_source_kv_params( def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries - kv_request_id only. - """ + kv_request_id only.""" return {"remote_decoder": {"kv_request_id": kv_request_id}} @@ -244,8 +241,7 @@ def test_lookup_returns_hit_for_different_request_id(self): def test_lookup_returns_miss_without_prefill_key(self): """No ``remote_prefiller`` sub-dict means the request was not routed for remote prefill — local prefill should run instead, so lookup() - returns MISS even when a stale ``remote_decoder`` block is present. - """ + returns MISS even when a stale ``remote_decoder`` block is present.""" mgr = _make_manager() ctx = _req_context(kv_params=_remote_decoder_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.MISS @@ -313,8 +309,7 @@ class TestServeExternalRequests: def test_flushes_failed_serve_ctxs_then_serves_each_session(self): """serve_external_requests releases the failed serves left by reaped sessions via parent.on_request_finished (clearing the - queue), then delegates to every live session with the same parent. - """ + queue), then delegates to every live session with the same parent.""" mgr = _make_manager() ctx = ReqContext(req_id="p2p:peer:req-1:lu1") mgr._failed_serve_ctxs = [ctx] @@ -368,8 +363,7 @@ def test_missing_kv_request_id_fails(self): def test_no_binding_yet_parks_in_unbound_stores(self): """submit_store without a bound session buffers the batch keyed by kv_request_id; no session is pre-created (peer_id is unknown - to the producer at store time). - """ + to the producer at store time).""" mgr = _make_manager() job = _job_metadata( job_id=1, @@ -396,8 +390,7 @@ def test_no_binding_yet_parks_in_unbound_stores(self): def test_routes_to_bound_session(self): """If a session has already received FetchMsg for this kv_request_id (so _kv_to_session is populated), submit_store - forwards directly to that session rather than re-buffering. - """ + forwards directly to that session rather than re-buffering.""" mgr = _make_manager() bound = _FakeSession(peer_id="10.0.0.1:8000", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] @@ -420,8 +413,7 @@ def test_routes_to_bound_session(self): def test_extra_top_level_keys_are_ignored(self): """Producer-side kv_transfer_params should not pre-create a session even when a stale caller still passes a top-level - ``remote_host``/``remote_port`` next to ``remote_decoder``. - """ + ``remote_host``/``remote_port`` next to ``remote_decoder``.""" mgr = _make_manager() params = _remote_decoder_kv_params() params["remote_host"] = "stale" @@ -465,8 +457,7 @@ def test_no_session_fails(self): def test_happy_path_with_active_session(self): """When the peer's session exists, submit_load forwards to - session.request_blocks and does NOT add a finished result yet. - """ + session.request_blocks and does NOT add a finished result yet.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" existing = _FakeSession(peer_id=peer_id, connected=True) @@ -486,8 +477,7 @@ def test_happy_path_with_active_session(self): def test_missing_consumer_flag_fails(self): """Peer fields present but neither do_remote_prefill nor do_p2p_fetch is set — submit_load fails the job rather than - emit a stray FetchMsg. - """ + emit a stray FetchMsg.""" mgr = _make_manager() params = { "remote_host": "10.0.0.1", @@ -531,8 +521,7 @@ def test_no_kv_request_id_does_nothing(self): def test_decoder_side_calls_session_finish_request(self): """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id because the consumer addresses the producer it loaded from. The - session's finish_request cancels the client-role load. - """ + session's finish_request cancels the client-role load.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) @@ -544,8 +533,7 @@ def test_decoder_side_calls_session_finish_request(self): def test_p2p_consumer_side_calls_session_finish_request(self): """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id so the session drops any pending lookups (cancel_lookups) and - cancels any inbound load. - """ + cancels any inbound load.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) @@ -556,8 +544,7 @@ def test_p2p_consumer_side_calls_session_finish_request(self): def test_prefiller_bound_id_routes_via_kv_to_session(self): """Prefiller-side finish for an id whose session is already bound - (FetchMsg received) routes via _kv_to_session and pops the entry. - """ + (FetchMsg received) routes via _kv_to_session and pops the entry.""" mgr = _make_manager() bound = _FakeSession(peer_id="some-peer:1", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] @@ -571,8 +558,7 @@ def test_prefiller_unbound_id_leaves_batches_parked(self): and no session binding is a no-op on `_unbound_stores`. The parked batches survive until a peer fetches them or the `_reap_unbound_stores` timeout fires — `on_request_finished` - must not evict them. - """ + must not evict them.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = _make_manager() @@ -742,8 +728,7 @@ def remove_remote_peer(self, pid): def test_reap_fails_probes(self): """A reaped session's in-flight lookups land in _failed_req_ids so - the consumer's lookup() returns MISS instead of RETRY forever. - """ + the consumer's lookup() returns MISS instead of RETRY forever.""" class FakeData: def remove_remote_peer(self, pid): @@ -778,8 +763,7 @@ def test_unbound_store_kept_within_timeout(self): def test_unbound_store_reaped_after_timeout(self): """Unbound stores past the reap deadline surface as failed and their kv_request_id lands in _reaped_stores so a late FetchMsg is rejected - instead of parking demand nothing can satisfy. - """ + instead of parking demand nothing can satisfy.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = self._make() @@ -801,8 +785,7 @@ def test_unbound_store_reaped_after_timeout(self): def test_submit_store_parks_unbound_batch(self): """submit_store on an unbound id appends a batch with a fresh - submitted_at stamp so the unbound-store sweep can age it out. - """ + submitted_at stamp so the unbound-store sweep can age it out.""" mgr = _make_manager() job = _job_metadata( job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") @@ -824,21 +807,18 @@ class TestHasPendingWork: """has_pending_work() must always return True so the engine keeps ticking the offload pipeline — that's the only thread driving _control.poll() (incoming peer connects) and session.poll() - (incoming fetch messages on existing sessions). - """ + (incoming fetch messages on existing sessions).""" def test_returns_true_unconditionally_to_keep_engine_ticking(self): """Even with no sessions and no jobs, has_pending_work() returns True so the engine keeps calling get_finished_jobs(), which is - what drives _control.poll() for inbound peer connects. - """ + what drives _control.poll() for inbound peer connects.""" mgr = _make_manager() assert mgr.has_pending_work() is True def test_returns_true_even_when_sessions_present(self): """The result is the same regardless of session state — there is - no 'idle' branch. - """ + no 'idle' branch.""" mgr = _make_manager() mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") # type: ignore[assignment] assert mgr.has_pending_work() is True @@ -851,8 +831,7 @@ def test_returns_true_even_when_sessions_present(self): class _ShutdownFakeData: """Fake DataTransport that records cancel/poll/close calls and - drives the wait-cancel loop with a scriptable `still` queue. - """ + drives the wait-cancel loop with a scriptable `still` queue.""" def __init__(self, still_queue: list[list[int]] | None = None) -> None: # Each list in still_queue is the set of ids the next @@ -896,8 +875,7 @@ def close(self) -> None: class TestShutdownDrain: """shutdown() drains inflight transfers via cancel(mode='wait') before calling _data.close(), with a 3s deadline fallback to - cancel(mode='immediate'). - """ + cancel(mode='immediate').""" def _prep( self, @@ -1316,8 +1294,7 @@ def close(self) -> None: class TestAcceptNewPeers: """A second inbound from an already-connected peer is rejected and the - new conn is closed; the existing session is left untouched. - """ + new conn is closed; the existing session is left untouched.""" def test_duplicate_connection_is_closed_and_existing_session_untouched(self): mgr = _make_manager() @@ -1339,8 +1316,7 @@ def test_creates_session_for_new_peer(self): """An inbound conn from a peer with no existing session creates a fresh connected session and registers it under conn.peer_id. The prefiller has no pre-created pending session anymore — the - first signal of a peer's existence is its inbound connection. - """ + first signal of a peer's existence is its inbound connection.""" class FakeData: block_len = 4096 @@ -1387,8 +1363,7 @@ def close(self) -> None: class TestPollOnce: """_poll_once must drain control, accept new peers, poll every session, - surface results, and reap dead sessions — in that order. - """ + surface results, and reap dead sessions — in that order.""" def test_orchestrates_accept_poll_and_reap(self): mgr = _make_manager() @@ -1444,8 +1419,7 @@ def remove_remote_peer(self_inner, pid): def test_new_fetch_id_binds_and_replays_unbound_batches(self): """When session.poll() reports a kv_request_id whose FetchMsg arrived this tick, the manager binds it to that session and - replays every parked submit_store batch via add_stored_blocks. - """ + replays every parked submit_store batch via add_stored_blocks.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = _make_manager() @@ -1479,8 +1453,7 @@ def poll(self_inner): def test_new_fetch_id_with_no_unbound_still_binds(self): """A FetchMsg for a kv_request_id with no parked batches still - records the binding so subsequent submit_stores route fast. - """ + records the binding so subsequent submit_stores route fast.""" mgr = _make_manager() peer = "10.0.0.1:8000" sess = _FakeSession( @@ -1503,8 +1476,7 @@ def poll(self_inner): def test_failed_load_records_kv_request_id(self): """A LoadResult(success=False) from session.poll() must add its - kv_request_id to _failed_req_ids so future lookups return MISS. - """ + kv_request_id to _failed_req_ids so future lookups return MISS.""" mgr = _make_manager() peer = "10.0.0.1:8000" sess = _FakeSession( @@ -1568,8 +1540,7 @@ def test_returns_when_session_has_no_inflight_or_inbound(self): def test_logs_warning_after_5s_then_completes(self, monkeypatch): """A session that stays inflight past 5s triggers the warning, and - once it clears the loop returns. - """ + once it clears the loop returns.""" mgr = _make_manager() mgr._control = _DrainCtrl() # type: ignore[assignment] sess = _FakeSession(peer_id="peer:1") @@ -1627,8 +1598,7 @@ def clearing_sleep(_s): class TestOnScheduleEnd: def test_is_noop(self): """on_schedule_end is a documented no-op; just confirm it doesn't - raise and doesn't mutate state. - """ + raise and doesn't mutate state.""" mgr = _make_manager() before_sessions = dict(mgr._sessions) before_jobs = list(mgr._finished_jobs) @@ -1654,8 +1624,7 @@ class TestConnectionDeathMidTransfer: prefiller-side store no longer travels through the session at store time (it's parked in _unbound_stores keyed by kv_request_id), so its cleanup on connection death is via on_request_finished or the - unbound-store timeout — covered separately below. - """ + unbound-store timeout — covered separately below.""" def test_dead_connection_with_pending_work_surfaces_failures(self): mgr_a, mgr_b = _build_paired_managers() diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index ff5641134ece..a09445668fd3 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -626,8 +626,7 @@ def test_active_loads_work_list_tracks_in_flight(self): """collect_results / has_active_loads use the _active_loads work-list, armed when a fetch is issued and discarded exactly when its load clears — a probe-only request never enters it, and completion empties - it while the entry may briefly linger for GC. - """ + it while the entry may briefly linger for GC.""" session, conn, _ = _make_session() _activate(session, conn) client = session._client @@ -704,8 +703,7 @@ def test_load_abort_ack_timeout_surfaces_failure(self): def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from - _requests — covers the on_abort_ack arrival path. - """ + _requests — covers the on_abort_ack arrival path.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( @@ -742,8 +740,7 @@ class TestLookupFlow: def test_aggregate_flush_resolve_round_trip(self): """register_lookup → flush sends one LookupMsg → response resolves entries → register_lookup returns the cached bool on - every call, and repeat probes never re-issue a LookupMsg. - """ + every call, and repeat probes never re-issue a LookupMsg.""" session, conn, _ = _make_session() _activate(session, conn) @@ -796,8 +793,7 @@ def test_aggregate_flush_resolve_round_trip(self): def test_request_blocks_clears_probe_cache(self): """A resolved HIT probe is popped when its fetch is issued, so a re-scheduled request re-probes instead of trusting the stale True - (the served block is unpinned and may have been evicted). - """ + (the served block is unpinned and may have been evicted).""" session, conn, _ = _make_session() _activate(session, conn) @@ -883,8 +879,7 @@ def test_separate_lookup_msg_per_kv_request_id(self): def test_multiple_lookup_msgs_across_steps(self): """A request's block set may be discovered across scheduler steps: each step that registers new keys flushes its own LookupMsg for - the same kv_request_id, carrying only the newly-probed keys. - """ + the same kv_request_id, carrying only the newly-probed keys.""" session, conn, _ = _make_session() _activate(session, conn) @@ -912,8 +907,7 @@ def test_multiple_lookup_msgs_across_steps(self): def test_split_response_resolves_across_messages(self): """Producer may answer one LookupMsg's keys across multiple - LookupRespMsgs — pairs are self-describing so each lands. - """ + LookupRespMsgs — pairs are self-describing so each lands.""" session, conn, _ = _make_session() _activate(session, conn) @@ -960,8 +954,7 @@ def test_finish_request_cancels_pending_lookups(self): def test_finish_after_flushed_lookup_sends_empty_fetch(self): """LookupMsg flushed but no FetchMsg sent (all-miss case) → finish_request emits an empty FetchMsg so the peer can drop its - lookup state and call parent.on_request_finished. - """ + lookup state and call parent.on_request_finished.""" session, conn, _ = _make_session() _activate(session, conn) @@ -979,8 +972,7 @@ def test_finish_after_flushed_lookup_sends_empty_fetch(self): def test_finish_without_flushed_lookup_sends_no_fetch(self): """No LookupMsg was ever sent → finish_request must not emit an - empty FetchMsg (the peer has no state to release). - """ + empty FetchMsg (the peer has no state to release).""" session, conn, _ = _make_session() _activate(session, conn) @@ -995,8 +987,7 @@ def test_finish_without_flushed_lookup_sends_no_fetch(self): def test_finish_after_real_fetch_sends_no_second_fetch(self): """A real FetchMsg was already sent for the id → finish_request - must not emit a second (empty) FetchMsg. - """ + must not emit a second (empty) FetchMsg.""" session, conn, _ = _make_session() _activate(session, conn) @@ -1027,8 +1018,7 @@ def test_server_lookup_deferred_until_serve_then_all_misses(self): """``poll()`` only enqueues an inbound LookupMsg — no response is sent until ``serve_external_requests``. With an all-miss parent the aggregated LookupRespMsg carries the same keys and - ``hits=[False, ...]``. - """ + ``hits=[False, ...]``.""" session, conn, _ = _make_session() _activate(session, conn) @@ -1088,8 +1078,7 @@ class TestServerLookupHandling: def test_immediate_hits_create_one_store_job(self): """All-HIT batch: one create_store_job call with all keys, one LookupRespMsg with hits=[True]*N, on_request_finished fires at the - end of serve, and `available` is populated for the eventual fetch. - """ + end of serve, and `available` is populated for the eventual fetch.""" cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) session, conn, _ = _make_session() _activate(session, conn) @@ -1119,8 +1108,7 @@ def test_immediate_hits_create_one_store_job(self): def test_all_misses_no_store_job_finish_fires(self): """All-MISS batch: no create_store_job call; one LookupRespMsg - with hits=[False]*N; on_request_finished fires at end of serve. - """ + with hits=[False]*N; on_request_finished fires at end of serve.""" cb = FakeParent() session, conn, _ = _make_session() _activate(session, conn) @@ -1140,8 +1128,7 @@ def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): """HIT/MISS resolutions do not go out on first sight when any key is still HIT_PENDING / RETRY. The lookup parks until every key has settled (or the deadline fires), then one - LookupRespMsg carries all keys in wire order. - """ + LookupRespMsg carries all keys in wire order.""" cb = FakeParent( stored={b"hA": 1}, pending={b"hB"}, @@ -1170,8 +1157,7 @@ def test_pending_resolves_then_aggregate_response_fires(self): the deferred aggregate response: one LookupRespMsg carrying both keys in wire order, and one create_store_job call per HIT (the second HIT is pinned when it resolves, not when the - response goes out). - """ + response goes out).""" cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1207,8 +1193,7 @@ def test_pending_resolves_then_aggregate_response_fires(self): def test_pending_timeout_replies_miss_no_store_job(self): """A HIT_PENDING key that stays pending past the batch ``deadline`` is force-MISS and never pinned; the deferred - aggregate response fires with hits=[False]. - """ + aggregate response fires with hits=[False].""" cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1235,8 +1220,7 @@ def test_pending_timeout_replies_miss_no_store_job(self): def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): """Two LookupMsgs for the same kv_request_id get distinct ctxs - and two on_request_finished calls (one per batch). - """ + and two on_request_finished calls (one per batch).""" cb = FakeParent(stored={b"hA": 1, b"hB": 2}) session, conn, _ = _make_session() _activate(session, conn) @@ -1260,8 +1244,7 @@ def test_close_returns_open_batch_ctxs_as_failed_serves(self): """Tearing the session down with a parked batch returns the synthetic ctx as a failed serve (no parent handle at teardown) so the manager can release the TieringManager's state on its next - serve. - """ + serve.""" cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1282,8 +1265,7 @@ def test_close_returns_open_batch_ctxs_as_failed_serves(self): def test_wire_finish_drops_pending_batches_for_kv_request_id(self): """``ServerRole.finish(kv_request_id)`` drops every parked batch whose kv_request_id matches and queues its ctx for the next - serve's on_request_finished. - """ + serve's on_request_finished.""" cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1314,8 +1296,7 @@ def test_wire_finish_drops_pending_batches_for_kv_request_id(self): def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): """A peer FetchMsg terminates the lookup phase for its id: parked lookups with matching kv_request_id are dropped and their - ctx queued for on_request_finished; other kv_request_ids untouched. - """ + ctx queued for on_request_finished; other kv_request_ids untouched.""" cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) @@ -1357,8 +1338,7 @@ def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): def test_lookup_then_fetch_round_trip_emits_store_result(self): """End-to-end: lookup pins primary slots → fetch matches them → NIXL transfer completes → StoreResult surfaces with the - create_store_job's job_id (the engine releases the pin). - """ + create_store_job's job_id (the engine releases the pin).""" cb = FakeParent(stored={b"hA": 7, b"hB": 8}) session, conn, transport = _make_session() _activate(session, conn) @@ -1480,8 +1460,7 @@ def test_abort_fetch_replies_with_ack(self): def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is - deferred and the abort is parked (abort_started_at set). - """ + deferred and the abort is parked (abort_started_at set).""" session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends @@ -1516,8 +1495,7 @@ def test_abort_fetch_defers_ack_when_cancel_pending(self): def test_abort_fetch_acks_after_drain(self): """Once the transport reports the tid as DONE the parked abort - completes and AbortAckMsg is sent. - """ + completes and AbortAckMsg is sent.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1556,8 +1534,7 @@ def test_abort_fetch_acks_after_drain(self): def test_abort_fetch_force_cancels_after_timeout(self): """If wait-mode never drains, the deadline forces immediate - cancel and an ack is still sent. - """ + cancel and an ack is still sent.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1599,8 +1576,7 @@ def test_abort_fetch_force_cancels_after_timeout(self): def test_abort_fetch_idempotent_while_draining(self): """Receiving AbortFetchMsg twice for the same kv_request_id - keeps a single pending entry and produces a single ack. - """ + keeps a single pending entry and produces a single ack.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1658,8 +1634,7 @@ def test_store_timeout(self): def test_store_timeout_then_late_completion_no_duplicate(self): """A job timed out by _timeout_pending_store_jobs must not also emit a contradictory StoreResult(success=True) when the transport - later reports the same transfer as done. - """ + later reports the same transfer as done.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1692,8 +1667,7 @@ def test_store_timeout_then_late_completion_no_duplicate(self): def test_store_timeout_then_late_failure_no_duplicate(self): """Symmetric guard: a timed-out job must not also emit a second StoreResult(success=False) when the transport later reports the - same transfer as failed. - """ + same transfer as failed.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1736,8 +1710,7 @@ def _last_transfer_done(self, conn: FakeConnection) -> dict | None: def test_no_inflight_unmatched_demand_sends_failure(self): """finish_request with unmatched demand and no inflight -> - immediate TransferDoneMsg(success=False); _outbound cleared. - """ + immediate TransferDoneMsg(success=False); _outbound cleared.""" session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded a block we never stored. @@ -1763,8 +1736,7 @@ def test_no_inflight_unmatched_demand_sends_failure(self): def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the - early-fail message and clears _outbound. - """ + early-fail message and clears _outbound.""" session, conn, transport = _make_session() _activate(session, conn) # Demand 2 blocks; we store 1 (kicks one inflight transfer). @@ -1801,8 +1773,7 @@ def test_with_inflight_defers_then_fires_on_last_transfer(self): def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: - when remaining hits 0, success=True still fires. - """ + when remaining hits 0, success=True still fires.""" session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( @@ -1831,8 +1802,7 @@ def test_prefiller_first_finish_before_fetch(self): """Prefiller-first: finish_request runs before the decoder's fetch arrives. State is held until fetch, then finalized — success=True if all demand was matched against - available blocks, else success=False. - """ + available blocks, else success=False.""" # Case A: all demand satisfied by available blocks. session, conn, transport = _make_session() _activate(session, conn) @@ -1893,8 +1863,7 @@ def test_finish_request_no_inflight_emits_store_failure(self): """finish_request with a stored-but-unmatched job and no inflight -> TransferDoneMsg(success=False) AND deferred StoreResult(success=False) for the submit_store'd job, instead of - the 30s _STORE_TIMEOUT_S path. - """ + the 30s _STORE_TIMEOUT_S path.""" session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded b"demand"; we never stored it. @@ -1930,8 +1899,7 @@ def test_finish_request_no_inflight_emits_store_failure(self): def test_finish_request_remaining_zero_emits_success_via_inflight(self): """Deferred-via-inflight path: finish_request with inflight, last transfer drains remaining to 0 -> TransferDoneMsg(success=True) - AND StoreResult(success=True). - """ + AND StoreResult(success=True).""" session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( @@ -2279,8 +2247,7 @@ def test_close_returns_pending_loads_and_stores(self): def test_send_failure_marks_connection_dead(self): """A raising send must mark the connection dead, not silently drop the message — otherwise the session lingers alive, is never reaped, - and in-flight lookups/loads toward the dead peer hang forever. - """ + and in-flight lookups/loads toward the dead peer hang forever.""" session, conn, _ = _make_session() _activate(session, conn) assert session.alive @@ -2293,8 +2260,7 @@ def test_send_failure_marks_connection_dead(self): def test_close_surfaces_inflight_lookups(self): """close() reports kv_request_ids whose symmetric-P2P probe is still - unresolved; resolved probes are not reported (their answer is in). - """ + unresolved; resolved probes are not reported (their answer is in).""" session, conn, _ = _make_session() _activate(session, conn) @@ -2402,8 +2368,7 @@ class TestDispatchErrorHandling: def test_value_error_disconnects_on_first_occurrence(self): """A FetchMsg that fails validate() raises ValueError and must - terminate the session immediately, with a DisconnectMsg sent. - """ + terminate the session immediately, with a DisconnectMsg sent.""" session, conn, _ = _make_session() _activate(session, conn) # length mismatch → FetchMsg.validate raises ValueError @@ -2430,8 +2395,7 @@ def test_transfer_done_missing_field_disconnects(self): def test_internal_error_does_not_disconnect_once(self): """A non-ValueError raised by a handler is treated as an internal - bug: counter increments, session stays alive on a single hit. - """ + bug: counter increments, session stays alive on a single hit.""" session, conn, _ = _make_session() _activate(session, conn) @@ -2454,8 +2418,7 @@ def _boom(*args, **kwargs): def test_internal_error_threshold_disconnects(self): """Once consecutive non-protocol errors hit the threshold, the - session tears down via _protocol_error. - """ + session tears down via _protocol_error.""" session, conn, _ = _make_session() _activate(session, conn) @@ -2479,8 +2442,7 @@ def _boom(*args, **kwargs): def test_internal_error_counter_resets_on_success(self): """A successful dispatch between errors prevents the threshold - from being reached. - """ + from being reached.""" session, conn, _ = _make_session() _activate(session, conn) @@ -2523,8 +2485,7 @@ class TestInflightPerReqInvariant: def test_invariant_holds_through_lifecycle(self): """Run a full submit→complete sequence for two concurrent kv_request_ids and assert the counter matches `_inflight` at - every observable step, including the empty-after-finish case. - """ + every observable step, including the empty-after-finish case.""" session, conn, transport = _make_session() _activate(session, conn) @@ -2589,8 +2550,7 @@ def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must match the actual presence in `_inflight` for both hits and misses. The whole point of the counter is that this lookup is - constant-time, but we assert correctness, not timing. - """ + constant-time, but we assert correctness, not timing.""" session, _, _ = _make_session() for kv_id_idx in range(100): kv_id = f"req-{kv_id_idx}" diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index 616c6efa9479..3649c8df0488 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -264,8 +264,7 @@ def test_mark_miss_flips_cached_verdict_without_reprobing(self): subsequent lookup then returns False (MISS) directly, WITHOUT enqueuing a fresh batch_lookup — that is what makes the request unable to loop. The entry is retained (as False) not dropped, so cleanup()'s reverse - index stays consistent; an unknown key is a no-op. - """ + index stays consistent; an unknown key is a no-op.""" mgr = InMemoryLookupManager(existing_keys={_key(1), _key(2)}) ctx = _ctx("reqA") assert mgr.lookup(_key(1), ctx) is None @@ -305,8 +304,7 @@ def test_enqueue_once_invariant_enforced(self): receive at most one result per key. Normal operation resolves a key a single time; a second result for an already-decided key trips the assert that guards the invariant (a silent overwrite could flip a corrected - miss back to True and reopen the failed-load livelock). - """ + miss back to True and reopen the failed-load livelock).""" mgr = InMemoryLookupManager(existing_keys={_key(1)}) ctx = _ctx("reqA") diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 13621fb4c432..9c3a64d22a2a 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -369,8 +369,7 @@ def test_shutdown_discards_pending_tasks(fs_tier): @pytest.mark.parametrize("use_c_ext", [True, False]) def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): """Data written by store must be exactly recovered by load, for batches - of any size -- including the empty batch. - """ + of any size -- including the empty batch.""" import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -467,8 +466,7 @@ def test_wait_idle_blocks_until_tasks_complete(): def test_batch_lookup_c_extension(tmp_path): """Validates batch_lookup_C: empty, single, all-existing, all-missing, - mixed ordering, and input type validation. - """ + mixed ordering, and input type validation.""" try: from vllm.fs_io_C import batch_lookup as batch_lookup_C except ImportError: @@ -532,8 +530,7 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): @pytest.mark.parametrize("use_c_ext", [True, False]) def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): """Smoke test: a block id beyond the primary tensor's block count must - fail the job, for both the C extension and the Python fallback. - """ + fail the job, for both the C extension and the Python fallback.""" import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -615,8 +612,7 @@ def test_batched_partial_load_failure_keeps_loaded_blocks( loaded before it (#50321). Corrupt the LAST block: the earlier blocks load fine, so the job reports successful_keys for them and marks only the failed tail a miss. The earlier keys stay HIT — including for the same request — - while the corrupt block stays a MISS (its file was removed). - """ + while the corrupt block stays a MISS (its file was removed).""" import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -668,8 +664,7 @@ def test_batched_load_first_block_fails_marks_whole_batch( ): """When the FIRST block fails, nothing loaded before it: the job reports no successful_keys (None) and the whole batch is marked a miss for the - request. - """ + request.""" import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -700,8 +695,7 @@ def test_transient_load_failure_leaves_file(fs_tier, monkeypatch, use_c_ext): """A transient host error (here ELOOP on open) is NOT a short read: the job fails but the block file must survive untouched, on both the C and Python paths. Deleting on a transient error would turn a passing hiccup into - permanent data loss. - """ + permanent data loss.""" import vllm.v1.kv_offload.tiering.fs.io as io_mod if use_c_ext and not io_mod._HAS_FSIO_C: @@ -805,8 +799,7 @@ def test_mixed_job_results_emit_event_only_for_successful_job( fs_tier_with_events, monkeypatch ): """With a failed and a successful store job in flight, exactly one event - is emitted and its keys belong to the successful job. - """ + is emitted and its keys belong to the successful job.""" import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod tier = fs_tier_with_events @@ -891,8 +884,7 @@ def test_events_require_global_kv_events_flag(tmp_path): def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): """A GPU->CPU->fs cascade surfaces the tier-owned FS stored event via the - TieringOffloadingManager's aggregated take_events(). - """ + TieringOffloadingManager's aggregated take_events().""" from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 0b3168ec3d32..bf8afca42428 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -346,8 +346,7 @@ def test_failed_load_marks_verdict_negative(self): the scheduler would re-issue the same doomed promotion every step for the life of the request. The mark is served from cache with no re-probe, so even though the mock object is still 'present' the SAME - request now resolves to MISS. - """ + request now resolves to MISS.""" ctx = ReqContext(req_id="obj-livelock") self.tier.submit_store(make_job(1, [key(1)], [0])) assert all(r.success for r in drain(self.tier)) @@ -651,8 +650,7 @@ def test_store_event_uses_configured_locality(self, locality, expected): def test_mixed_job_results_emit_event_only_for_successful_job(self): """With a failed and a successful store job resolving in the same poll, exactly one event is emitted and its keys belong to the - successful job. - """ + successful job.""" original = self.agent.check_xfer_state self.agent.check_xfer_state = lambda h: "ERR" if h._id == 0 else original(h) self.tier.submit_store(make_job(1, [key(1)], [0])) # handle 0: fails diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 5cc5f48e3a03..63b4c2c4d0b5 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -273,8 +273,7 @@ def test_basic_store_and_lookup(self): def test_every_lookup_result_has_a_declared_cascade_behavior(): """A new LookupResult must state what the cascade does with it, rather - than falling into whichever branch happens to catch it. - """ + than falling into whichever branch happens to catch it.""" assert set(_CASCADE_SUPPLY_BEHAVIOR) == set(LookupResult) @@ -321,8 +320,7 @@ def _start_request(self, req_context: ReqContext = _CTX): def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): """A failed promotion still finalizes the primary slots with - success=False; the tier corrects its own verdict. - """ + success=False; the tier corrects its own verdict.""" from unittest.mock import patch from vllm.v1.kv_offload.tiering.base import JobResult @@ -1093,8 +1091,7 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( def test_on_new_request_lifecycle(self, manager_setup): """Policy defaults to CHUNK_LEVEL, escalates when a tier requests it, - and is cleaned up on on_request_finished. - """ + and is cleaned up on on_request_finished.""" # Default: all tiers return CHUNK_LEVEL ctx = ReqContext(req_id="req_policy_lifecycle") result = self.manager.on_new_request(ctx) @@ -1171,8 +1168,7 @@ def test_prepare_store_cascades_existing_chunks_to_request_level_tiers( def _make_request_level_request(self, req_id: str) -> ReqContext: """Start a request for which tier1 asks for request-level offloading, - with tier1's submit_store wrapped so the cascade can be observed. - """ + with tier1's submit_store wrapped so the cascade can be observed.""" self.secondary_tier1.on_new_request = lambda req_context: ( RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) ) @@ -1194,8 +1190,7 @@ def _cascaded_keys_for(self, req_id: str) -> list[set[OffloadKey]]: def test_cascade_rejects_retry_from_primary(self, manager_setup): """RETRY would be parked with no guarantee of ever draining, so the cascade refuses it. No real primary tier returns it, so this is the one - disposition that has to be forced. - """ + disposition that has to be forced.""" keys = to_keys(range(3)) self._start_request() assert self.manager.prepare_store(keys, _CTX) is not None @@ -1210,8 +1205,7 @@ def test_cascade_rejects_retry_from_primary(self, manager_setup): def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: """Leave a GPU->primary write for `keys` open, so they look present to - prepare_store but are not yet readable. - """ + prepare_store but are not yet readable.""" writer_ctx = ReqContext(req_id="req_writer") self._start_request(writer_ctx) assert self.manager.prepare_store(keys, writer_ctx) is not None @@ -1220,8 +1214,7 @@ def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setup): """A key another request is still writing must reach the peer once the write lands, not be dropped: prepare_store already counts it as stored - and the scheduler advances past its chunk, so nothing re-offers it. - """ + and the scheduler advances past its chunk, so nothing re-offers it.""" keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1238,8 +1231,7 @@ def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setu def test_deferred_cascade_holds_request_from_finalization(self, manager_setup): """A request that finishes with keys still deferred must stay alive, or - its tiers are torn down before the peer is ever served. - """ + its tiers are torn down before the peer is ever served.""" keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1262,8 +1254,7 @@ def test_cascade_drops_deferred_keys_whose_primary_write_failed( self, manager_setup ): """A failed write frees the chunk, so the deferred key resolves to MISS - and the request finalizes instead of parking forever. - """ + and the request finalizes instead of parking forever.""" keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1283,8 +1274,7 @@ def test_cascade_drops_deferred_keys_whose_primary_write_failed( def test_reset_cache_clears_orchestrator_state(self, manager_setup): """reset_cache wipes every kind of orchestrator state and resets primary tier; pending submissions are dropped without being sent - to the secondary tier. Active request state is retained. - """ + to the secondary tier. Active request state is retained.""" # Cascade — populates primary chunks and leaves cascade jobs # in _jobs (the synchronous example tier has already # queued completions); reset_cache's drain loop will pick them up. @@ -1376,8 +1366,7 @@ def test_tier_filter_skips_filtered_secondary( self, manager_setup, load_tier_filter ): """Filter excluding secondary medium returns MISS from secondaries - even when they hold the chunk; primary is unaffected. - """ + even when they hold the chunk; primary is unaffected.""" chunks = to_keys(range(2)) # Put one chunk in primary, one only in secondary self._start_request() diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index fccdf56409d6..9b12e7412ed6 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -1344,8 +1344,7 @@ def test_thinking_budget_invalid_budget_rejected(invalid_budget): def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index(): """Test thinking budget enforced for a long thinking run, - then a natural end marker. - """ + then a natural end marker.""" h = ThinkingBudgetStateHolder( MockReasoningConfig(), 8, 0, torch.device("cpu"), False ) diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index 85ad2032cde3..76877b655b1a 100644 --- a/tests/v1/logits_processors/utils.py +++ b/tests/v1/logits_processors/utils.py @@ -138,8 +138,7 @@ def __init__(self, group: str): 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`.""" @@ -158,8 +157,7 @@ def __call__( class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): """Example of wrapping a fake request-level logit processor to create a - batch-level logits processor - """ + batch-level logits processor""" def is_argmax_invariant(self) -> bool: return False diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index 32c203705cca..3868cc772e96 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -1241,8 +1241,7 @@ def test_mla_per_gpu_with_pipeline_parallelism(): def test_mla_model_metrics_excludes_standard_attention(): """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics, - for DeepSeek MLA models. - """ + for DeepSeek MLA models.""" hf_config = DeepseekV3Config( hidden_size=7168, num_attention_heads=128, diff --git a/tests/v1/metrics/test_ray_metrics.py b/tests/v1/metrics/test_ray_metrics.py index 7bc49aae6e64..ce4ebcbfdb0e 100644 --- a/tests/v1/metrics/test_ray_metrics.py +++ b/tests/v1/metrics/test_ray_metrics.py @@ -45,8 +45,7 @@ def test_engine_log_metrics_ray( max_tokens: int, ) -> None: """Simple smoke test, verifying this can be used without exceptions. - Need to start a Ray cluster in order to verify outputs. - """ + Need to start a Ray cluster in order to verify outputs.""" @ray.remote(num_gpus=1) class EngineTestActor: @@ -120,8 +119,7 @@ def test_sanitized_opentelemetry_name(): def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock: """Swap the wrapper's underlying Ray metric for a MagicMock while preserving the real metric's ``_tag_keys`` (labels() reads them to - validate arity). - """ + validate arity).""" real_metric = wrapper.metric mock = MagicMock() mock._tag_keys = real_metric._tag_keys @@ -131,8 +129,7 @@ def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock: def test_ray_counter_labels_returns_independent_children(): """RayCounterWrapper.labels() must return distinct labeled children that - each carry their own tag set. - """ + each carry their own tag set.""" base = RayCounterWrapper( name="vllm_test_finish_reason", documentation="", @@ -152,8 +149,7 @@ def test_ray_counter_labels_returns_independent_children(): def test_ray_counter_inc_forwards_per_child_tags(): """.inc() on a labeled counter must forward that child's tags to the - underlying Ray metric (not rely on a shared set_default_tags). - """ + underlying Ray metric (not rely on a shared set_default_tags).""" wrapper = RayCounterWrapper( name="vllm_test_counter_tag_forward", documentation="", @@ -218,8 +214,7 @@ def test_ray_histogram_labels_returns_independent_children_and_forwards_tags(): def test_ray_counter_labels_accepts_non_string_label_values(): """RayPrometheusStatLogger passes ``str(idx)`` for engine indexes; this covers the coercion path for any caller that passes a non-string label - value positionally. - """ + value positionally.""" wrapper = RayCounterWrapper( name="vllm_test_nonstr_label", documentation="", @@ -242,8 +237,7 @@ def test_ray_counter_labels_arity_validation(): def test_unlabeled_inc_carries_replica_id(): """Recording on an unlabeled metric must still pass ReplicaId — it's a - declared tag_key and Ray rejects updates that omit any declared key. - """ + declared tag_key and Ray rejects updates that omit any declared key.""" wrapper = RayCounterWrapper( name="vllm_test_unlabeled_replica_id", documentation="", @@ -256,8 +250,7 @@ def test_unlabeled_inc_carries_replica_id(): def test_double_labels_raises(): """labels() on an already-labeled child should raise, mirroring the - prometheus_client contract. - """ + prometheus_client contract.""" wrapper = RayCounterWrapper( name="vllm_test_double_labels", documentation="", diff --git a/tests/v1/sample/test_batched_count_greater_than.py b/tests/v1/sample/test_batched_count_greater_than.py index 0e2d19288245..c9ace93c6961 100644 --- a/tests/v1/sample/test_batched_count_greater_than.py +++ b/tests/v1/sample/test_batched_count_greater_than.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Test that batched_count_greater_than does not trigger 0/1 specialization -recompiles when batch_size varies. -""" +recompiles when batch_size varies.""" import torch diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 372e1c815643..d2cb99bc26fd 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -573,8 +573,7 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode): def test_prompt_logprobs_mode(): """prompt_logprobs must respect logprobs_mode: *_logits and *_logprobs must return different values. Prompt tokens skip sampling processors, - so processed_* == raw_* on the prompt side. - """ + so processed_* == raw_* on the prompt side.""" from vllm import LLM values: dict[str, float] = {} diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index abb069a26a78..231f6adbb2d7 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/tests/v1/sample/test_rejection_sampler.py @@ -54,8 +54,7 @@ def create_logits_tensor( token_idx_to_override: int | None = None, ) -> torch.Tensor: """Helper function to create logits tensor that - will produce desired token ids on argmax - """ + will produce desired token ids on argmax""" token_ids = [tokens[:-1] for tokens in output_token_ids] num_total_tokens = sum(len(tokens) for tokens in token_ids) logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE_TYPE) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 9f53a8c28760..0b570f49d974 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -884,8 +884,7 @@ def setup(self): def _make_logits(self, pattern: str) -> torch.Tensor: """Build (BATCH, VOCAB) logits with `pattern` applied to row 0 - (rows 1..B-1 stay clean so we can detect cross-row corruption). - """ + (rows 1..B-1 stay clean so we can detect cross-row corruption).""" logits = ( torch.randn( self.BATCH, diff --git a/tests/v1/sample/utils.py b/tests/v1/sample/utils.py index 19c50f79b1cf..e4e9f21092e3 100644 --- a/tests/v1/sample/utils.py +++ b/tests/v1/sample/utils.py @@ -204,8 +204,7 @@ def fake_update_logitsprocs_state( batch_update: BatchUpdate | None, ) -> None: """Imitate logits processors persistent batch state update - in engine core - """ + in engine core""" for logitproc in test_fakes.get_logitsprocs(): logitproc.update_state(batch_update) holder = test_fakes.sampling_metadata.thinking_budget_state_holder diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index 5f5b63c93ecf..f21dae7e0bd8 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -67,8 +67,7 @@ def _make_mixed_kv_cache_config( mamba_block_size: int = BLOCK_SIZE, ) -> KVCacheConfig: """Build a KVCacheConfig with a FullAttention group plus optional - SlidingWindow and Mamba groups for per-group metadata tests. - """ + SlidingWindow and Mamba groups for per-group metadata tests.""" register_all_kvcache_specs(vllm_config=None) groups = [] tensors = [] @@ -239,8 +238,7 @@ def test_events_disabled_no_events() -> None: # --------------------------------------------------------------------------- def test_block_stored_per_group_metadata_full_attention() -> None: """BlockStored carries medium, group_idx, kind, sliding_window, block_size, - and the correct block hash. - """ + and the correct block hash.""" fix = make_events_scheduler() sched = fix.scheduler assert sched.kv_event_medium == MEDIUM_CPU @@ -640,8 +638,7 @@ def test_mamba_block_size_not_scaled() -> None: def test_multi_group_fa_block_stored_group_idx() -> None: """Two distinct groups (FA + SW): each BlockStored carries its group_idx and kind. HybridKVCacheCoordinator requires >=2 distinct spec groups, so - we use FA (group 0) + SlidingWindow (group 1). - """ + we use FA (group 0) + SlidingWindow (group 1).""" kv_cfg = _make_mixed_kv_cache_config(num_blocks=16, sliding_window=32) fix = make_events_scheduler( num_cpu_blocks=8, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index c87a5f7c956a..ffc8b204787f 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1170,8 +1170,7 @@ def test_touched_blocks_survive_eviction() -> None: # --------------------------------------------------------------------------- def test_preemption_no_cpu_block_leak() -> None: """request_finished during in-flight load defers cleanup; - completes after load done. - """ + completes after load done.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1266,8 +1265,7 @@ def test_eager_store_preemption_cleanup() -> None: # --------------------------------------------------------------------------- def test_inflight_finish_deferred_cleanup() -> None: """Store, then start a load, request_finished defers, - load completion fires cleanup. - """ + load completion fires cleanup.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1401,8 +1399,7 @@ def test_chunked_prefill_reads_live_block_ids() -> None: """With chunked prefill, block IDs accumulate across scheduler steps. _prepare_eager_store_specs reads block IDs from scheduler_output via yield_req_data, so the store should reflect the updated (larger) block - list, not a stale snapshot. - """ + list, not a stale snapshot.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler @@ -1973,8 +1970,7 @@ def test_cp_eager_store_and_load_roundtrip( ) -> None: """With CP enabled, store blocks to CPU and reload them for a new request with matching tokens. Verifies that hash matching and transfer-pair - construction work with the virtual block size. - """ + construction work with the virtual block size.""" fix = _make_cp_scheduler( dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, @@ -2118,8 +2114,7 @@ def test_cp_effective_block_size_store_and_load( @pytest.mark.parametrize("cp_world_size", [1, 2, 4]) def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None: """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1 - because each virtual block covers more tokens. - """ + because each virtual block covers more tokens.""" kv_cache_config = _make_kv_cache_config(num_blocks=16) max_batched = 64 @@ -2811,8 +2806,7 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: """A sliding-window group nulls pages that left the window before the connector sees the finished request's block table, while the GPU keeps them hashed in its free queue. The eager store must copy those pages, or - the CPU lookup (which needs the whole window) misses every time. - """ + the CPU lookup (which needs the whole window) misses every time.""" fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False) sched = fix.scheduler gpu_pool = fix.gpu_block_pool diff --git a/tests/v1/spec_decode/test_backup_token_async_spec.py b/tests/v1/spec_decode/test_backup_token_async_spec.py index 3f99a5de0682..9340503ea4f2 100644 --- a/tests/v1/spec_decode/test_backup_token_async_spec.py +++ b/tests/v1/spec_decode/test_backup_token_async_spec.py @@ -131,8 +131,7 @@ def test_various_spec_token_counts(self, num_spec_tokens: int): def test_buggy_code_was_always_off_by_one(self): """The original code used seq_len as index, which is always one past - the end of output_token_ids even without async inflation. - """ + the end of output_token_ids even without async inflation.""" req_ids = ["r0"] requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])} batch = _FakeInputBatch(req_ids, [5]) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index 19166ee193f3..0cc78f3371c2 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -105,8 +105,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus the current draft-step offset. Padded entries are zeroed (matching the - main model runner's convention). - """ + main model runner's convention).""" fake = _make_fake_speculator() base = torch.tensor([100, 200, 300, 0], dtype=torch.int32) @@ -127,8 +126,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): """The padded tail of the upper-bound tensor is zeroed, so it stays a - valid tensor of length ``num_reqs_padded`` regardless of padding. - """ + valid tensor of length ``num_reqs_padded`` regardless of padding.""" fake = _make_fake_speculator() base = torch.zeros(2, dtype=torch.int32) @@ -143,8 +141,7 @@ def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): def test_build_draft_attn_metadata_clamps_to_max_model_len(): """The per-request upper bound (target bound + step) is clamped to the - model length so it never exceeds the allocated KV range. - """ + model length so it never exceeds the allocated KV range.""" fake = _make_fake_speculator(max_model_len=1024) base = torch.tensor([1023, 500], dtype=torch.int32) diff --git a/tests/v1/spec_decode/test_eagle_step_kernel.py b/tests/v1/spec_decode/test_eagle_step_kernel.py index 4d6244eb77fe..83a29bc0469c 100644 --- a/tests/v1/spec_decode/test_eagle_step_kernel.py +++ b/tests/v1/spec_decode/test_eagle_step_kernel.py @@ -132,8 +132,7 @@ def test_eagle_step_slot_mapping_kernel_exceeds_max(): def test_eagle_step_slot_mapping_kernel_cudagraph_padding(): """Test that padding threads write PADDING_SLOT_ID when - input_batch_size > batch_size (cudagraph padding). - """ + input_batch_size > batch_size (cudagraph padding).""" device = torch.device(DEVICE_TYPE) batch_size = 4 input_batch_size = 8 diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py index 622d4250655a..5510c36e60d2 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer.py +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -69,8 +69,7 @@ def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace: def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): """The proposer's slot-mapping math runs against the kernel-granularity - block table, so block_size must come from kernel_block_sizes. - """ + block table, so block_size must come from kernel_block_sizes.""" layer_names = {"draft.0.self_attn.attn"} proposer = _make_proposer(monkeypatch, layer_names) @@ -98,8 +97,7 @@ def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch) def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch): """_draft_attn_layer_names is a set; the attention groups built from it - must not depend on its (process-random) iteration order. - """ + must not depend on its (process-random) iteration order.""" layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"} expected_order = sorted(layer_names) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 3644f0528fe0..81f842419375 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -93,8 +93,7 @@ def test_eagle_max_len( @pytest.mark.parametrize("spec_max_model_len", [80, 150]) def test_mtp_speculative_config_max_model_len(spec_max_model_len: int): """Regression test for #41456: max_model_len in speculative config - should be respected for the draft model. - """ + should be respected for the draft model.""" model_config = ModelConfig( model="XiaomiMiMo/MiMo-7B-Base", runner="generate", diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index ea4429270228..a23189438475 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -29,8 +29,7 @@ def __init__(self, is_diffusion: bool): def test_structured_outputs_rejected_for_diffusion_models(): """Diffusion LLMs denoise the canvas in parallel, which is incompatible with the token-by-token grammar FSM. The request must fail with a clear - validation error instead of an FSM rejection mid-generation (#45436). - """ + validation error instead of an FSM rejection mid-generation (#45436).""" params = SamplingParams( structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) ) @@ -63,8 +62,7 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match): """json_object=False and an empty json schema pass the `is not None` exclusivity check but resolve to no structured-output key, so they must be rejected at request validation (-> 400) instead of reaching and crashing - the engine. - """ + the engine.""" params = SamplingParams(structured_outputs=structured_outputs) with pytest.raises(VLLMValidationError, match=match): params._validate_structured_outputs( @@ -87,8 +85,7 @@ def test_regex_with_nul_byte_rejected(regex): handled by xgrammar's native regex converter. It must be rejected at request validation in every backend mode (a clean 400), instead of reaching that native code or silently falling back to another backend in the default - 'auto' mode. - """ + 'auto' mode.""" params = SamplingParams(structured_outputs=StructuredOutputsParams(regex=regex)) # Rejected before backend selection, so it is a 400 even in 'auto' mode @@ -127,8 +124,7 @@ def test_regex_with_nul_byte_rejected(regex): ) def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): """Only `VLLMClientError` survives `AsyncLLM.generate` untouched; anything else - is wrapped in `EngineGenerateError` and served as a 500 instead of a 400. - """ + is wrapped in `EngineGenerateError` and served as a 500 instead of a 400.""" params = SamplingParams(structured_outputs=structured_outputs) with pytest.raises(VLLMClientError): params._validate_structured_outputs( diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index b5f180710e03..5d52c86374b7 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -46,8 +46,7 @@ def test_logprobs_tensors_cat(): def test_logprobs_tensors_tolists_with_tensor_boundaries(): """Adaptive verification hands over the request boundaries as a tensor (they only exist on device); tolists() must materialize it as a plain - list so slice_request splits requests correctly. - """ + list so slice_request splits requests correctly.""" tensors = LogprobsTensors( torch.tensor([[1, 2], [3, 4], [5, 6]]), torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), @@ -86,8 +85,7 @@ def test_sampling_mask_lists_to_nested_list(): def test_sampling_mask_tensors_match_finite_support(max_num_kept): """Whatever its size (empty, one, around the compact width, beyond the cap, the whole vocab), a sampled row's mask is exactly the ascending set - of its finite logits; rows that sampled nothing are empty. - """ + of its finite logits; rows that sampled nothing are empty.""" from vllm.v1.worker.gpu.sample.output import MAX_COMPACT_SUPPORT vocab_size = 20_001 @@ -114,8 +112,7 @@ def test_sampling_mask_tensors_match_finite_support(max_num_kept): def test_sampling_mask_matches_processed_top_k_top_p_support(): """The mask must exactly mirror whatever support `apply_top_k_top_p` (the real logits-processing function used by the sampler) actually - produces, whatever backend implements it. - """ + produces, whatever backend implements it.""" processed_logits = apply_top_k_top_p( logits=torch.tensor( [[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE @@ -141,8 +138,7 @@ def test_sampling_mask_matches_processed_top_k_top_p_support(): def test_sampling_mask_preserves_top_k_boundary_ties(): """When the kept support is wider than `max_num_kept` (e.g. a top-k boundary tie keeps more than k logits), the mask must fall back to the - exact bitmask instead of silently truncating to `max_num_kept` ids. - """ + exact bitmask instead of silently truncating to `max_num_kept` ids.""" processed_logits = torch.tensor( [[6.0, 5.0, 4.0, 4.0, 4.0, float("-inf"), float("-inf"), float("-inf")]], device=DEVICE_TYPE, diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 8a1af4e468ed..db0d7e68229b 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -427,8 +427,7 @@ def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends it: many requests, each with per-token tensors small enough that pyzmq copies their frames, while the accumulated payload frame is large enough - that pyzmq sends it zero-copy. - """ + that pyzmq sends it zero-copy.""" from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.outputs import LogprobsTensors diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 299ef7695f45..4cc8f9ffe2fd 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -101,8 +101,7 @@ def test_packed_dsv4_zeroer_zeroes_only_each_layers_page(): def test_overlaid_zeroer_dedups_segments_with_max_span(): """Two groups overlay one allocation; the zeroer must emit one segment per distinct byte offset, spanning the widest overlaid page, so a newly allocated block is fully - zeroed no matter which group owns it. - """ + zeroed no matter which group owns it.""" from unittest.mock import MagicMock from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 28df60f49d67..f8bd50c6d9c4 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -97,8 +97,7 @@ def test_draft_lookahead_uses_boundary_feature_when_cached(): """The drafter's +1 look-ahead can reach the feature at offset == processed_end (the next chunk). When its encoder output is already cached (the scheduler encoded it ahead), it is used for the look-ahead position - rather than ignored. - """ + rather than ignored.""" f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end runner = _make_runner([f0, f1], cached=[f0, f1]) @@ -115,8 +114,7 @@ def test_draft_lookahead_uses_boundary_feature_when_cached(): def test_draft_lookahead_tolerates_missing_boundary_feature(): """When the +1 look-ahead feature past the processed boundary is not yet encoded, fall back to the token embedding (the draft token is verified by - the target) instead of raising. - """ + the target) instead of raising.""" f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached runner = _make_runner([f0, f1], cached=[f0]) @@ -133,8 +131,7 @@ def test_draft_lookahead_tolerates_missing_boundary_feature(): def test_draft_lookahead_raises_on_interior_miss(): """A miss for a feature within the processed range (not the look-ahead boundary) is a real invariant violation and must fail loudly, even on the - drafter path. - """ + drafter path.""" f0 = _feature("h0", offset=0, length=8) # interior, within processed range runner = _make_runner([f0], cached=[]) @@ -144,8 +141,7 @@ def test_draft_lookahead_raises_on_interior_miss(): def test_target_path_raises_on_encoder_cache_miss(): """On the target path (no look-ahead) a miss is a real invariant - violation and must fail loudly. - """ + violation and must fail loudly.""" f0 = _feature("h0", offset=0, length=8) runner = _make_runner([f0], cached=[]) @@ -156,8 +152,7 @@ def test_target_path_raises_on_encoder_cache_miss(): @pytest.mark.parametrize("draft_lookahead", [0, 1]) def test_multi_request_batch_gathers_per_request(draft_lookahead): """Two prefilling requests in one batch: per-request query bounds must be - indexed by request, not applied as whole arrays. - """ + indexed by request, not applied as whole arrays.""" a0 = _feature("a0", offset=0, length=8) b0 = _feature("b0", offset=0, length=8) cache = EncoderCache() diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py index 6081163c4b51..150c4ed00711 100644 --- a/tests/v1/worker/test_gpu_bad_words.py +++ b/tests/v1/worker/test_gpu_bad_words.py @@ -73,8 +73,7 @@ def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor: def test_v2_bad_words_prefix_inside_draft_tokens(): """A prefix matching entirely within the draft tokens must mask the bad - word's last token at the draft position that completes the prefix. - """ + word's last token at the draft position that completes the prefix.""" out = _apply([[12, 13, 40]]) expected = torch.zeros_like(out) expected[2, 40] = -float("inf") @@ -83,8 +82,7 @@ def test_v2_bad_words_prefix_inside_draft_tokens(): def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): """A prefix spanning the committed/draft boundary must mask at the row - where the prefix completes, not one draft position later. - """ + where the prefix completes, not one draft position later.""" out = _apply([[11, 12, 30]]) expected = torch.zeros_like(out) expected[1, 30] = -float("inf") @@ -93,8 +91,7 @@ def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): def test_v2_bad_words_no_spurious_match_from_last_committed_token(): """The last committed token must not be double-counted as the first draft - token; [11, 11] never occurs in output [10, 11] + drafts [12, 13]. - """ + token; [11, 11] never occurs in output [10, 11] + drafts [12, 13].""" out = _apply([[11, 11, 50]]) expected = torch.zeros_like(out) torch.testing.assert_close(out, expected) diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 9f0b6da218af..1309686d53ce 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -43,8 +43,7 @@ def _make_batch( slot_pool: np.ndarray | None = None, ) -> InputBatch: """Consistent batch with random unique request slots, 1 + k_i logits - rows per request at the tail of each query segment. - """ + rows per request at the tail of each query segment.""" if slot_pool is None: slot_pool = np.arange(max_num_reqs) slots = rng.choice(slot_pool, size=num_reqs, replace=False) @@ -142,8 +141,7 @@ def _owned_rows(cu: np.ndarray, owned: np.ndarray) -> np.ndarray: def test_local_batch_partition(tp_size: int, seed: int): """Local sub-batches partition the batch by owner, preserve batch order within an owner, and carry exactly the owned requests' entries of every - per-request and per-row array. - """ + per-request and per-row array.""" rng = np.random.default_rng(seed) max_num_reqs = 64 batch = _make_batch(rng, num_reqs=23, max_num_reqs=max_num_reqs, max_spec=3) @@ -225,8 +223,7 @@ def test_local_batch_partition(tp_size: int, seed: int): def test_all_to_all_reassembly_simulated(tp_size: int): """Vocab shards computed from owner-sorted hidden states and routed through the simulated all-to-all reassemble into exactly the full-vocab - logits of each rank's owned rows. - """ + logits of each rank's owned rows.""" rng = np.random.default_rng(7) max_num_reqs = 32 batch = _make_batch(rng, num_reqs=11, max_num_reqs=max_num_reqs, max_spec=4) @@ -264,8 +261,7 @@ def test_all_to_all_reassembly_simulated(tp_size: int): @requires_cuda def test_shard_no_spec(): """Batches without draft tokens (one logits row per request) degenerate - cleanly: rows mirror requests. - """ + cleanly: rows mirror requests.""" rng = np.random.default_rng(11) tp_size = 4 max_num_reqs = 32 @@ -284,8 +280,7 @@ def test_shard_no_spec(): @pytest.mark.parametrize("tp_size", [2, 4]) def test_shard_grammar_output(tp_size: int): """The grammar filter keeps exactly the owned requests' bitmask rows, in - grammar order, and the per-rank pieces partition the global bitmask. - """ + grammar order, and the per-rank pieces partition the global bitmask.""" rng = np.random.default_rng(5) max_num_reqs = 32 num_reqs = 12 @@ -340,8 +335,7 @@ def test_shard_grammar_output(tp_size: int): @pytest.mark.parametrize("tp_size", [2, 4]) def test_gather_src_indices_round_trip(tp_size: int): """Padded per-rank result blocks indexed by gathered_src_indices restore - batch order, including when some rank owns zero requests. - """ + batch order, including when some rank owns zero requests.""" rng = np.random.default_rng(3) max_num_reqs = 40 all_slots = np.arange(max_num_reqs) @@ -376,8 +370,7 @@ def test_gather_src_indices_round_trip(tp_size: int): def test_gather_sampler_output_kernels(tp_size: int): """The fused pack/unpack kernels round-trip per-rank sampler results into batch order: token ids zero-filled past the local width, counts - cast back to int32, identical on every rank. - """ + cast back to int32, identical on every rank.""" from dataclasses import replace as dc_replace from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -472,8 +465,7 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: def test_gather_sampler_output_logprobs_and_nans(): """num_nans is reduced per request and gathered; LogprobsTensors are packed into fixed per-request blocks, width-aligned across ranks (pad - and truncate paths), and restored to global row order. - """ + and truncate paths), and restored to global row order.""" from dataclasses import replace as dc_replace from vllm.v1.outputs import LogprobsTensors @@ -606,8 +598,7 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: def test_finish_requests_frees_slots_in_sorted_order(): """Request slots must be freed in the same order on every TP rank: ownership derives from slot indices, and `finished_req_ids` is a set - whose iteration order is per-process hash-randomized. - """ + whose iteration order is per-process hash-randomized.""" from vllm.v1.worker.gpu.model_runner import GPUModelRunner removed: list[str] = [] diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index f44edabea4cc..e77156eef901 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -217,8 +217,7 @@ def test_v1_block_table_move_row_clears_vacated_row(): not keep stale block ids. Padded dummy-run batches dereference stale rows as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode padding) and write state in place there — corrupting the blocks' new - owner once they are reallocated, e.g. to an in-flight NIXL load. - """ + owner once they are reallocated, e.g. to an in-flight NIXL load.""" from vllm.v1.worker.block_table import BlockTable block_table = BlockTable( @@ -250,8 +249,7 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): PAD-filled, state indices are not), so stale rows would direct dummy state writes at freed — possibly reallocated — blocks. get_dummy_block_tables must hand out zeroed (null block) rows while - preserving the persistent storage address for CUDA graphs. - """ + preserving the persistent storage address for CUDA graphs.""" device = torch.device("cuda") block_tables = BlockTables( block_sizes=[16], diff --git a/tests/v1/worker/test_gpu_logit_bias.py b/tests/v1/worker/test_gpu_logit_bias.py index 8e7148575b7b..3bf62380bbc6 100644 --- a/tests/v1/worker/test_gpu_logit_bias.py +++ b/tests/v1/worker/test_gpu_logit_bias.py @@ -85,8 +85,7 @@ def test_v2_min_tokens_keeps_masking_when_other_candidates_remain(): def test_v2_min_tokens_mixed_batch_gates_restore_per_request(): """CHECK_ALL_MASKED_ROWS is a batch-wide constexpr, so a mixed batch must - still consult the per-request flag. - """ + still consult the per-request flag.""" out = _apply(_only_stop_token_left(2), structured=[True, False]) assert out[0, STOP_TOKEN] == 1.0 diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 2dae08688995..80e8318fa198 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -73,8 +73,7 @@ def _restore_default_dtype(): """Several tests here set the process-wide default dtype to float16 and previously leaked it, corrupting later float-sensitive tests in the same - pytest process (torch.randn silently produced fp16). - """ + pytest process (torch.randn silently produced fp16).""" old = torch.get_default_dtype() yield torch.set_default_dtype(old) @@ -1454,8 +1453,7 @@ class StopInitialization(Exception): def test_hybrid_block_table_initialization(): """Test hybrid block table with different kernel and kvcache_manager block - sizes. - """ + sizes.""" from vllm.v1.worker.block_table import BlockTable # Test configuration: kvcache_manager block size = 32, diff --git a/tests/v1/worker/test_gpu_model_runner_mm_gather.py b/tests/v1/worker/test_gpu_model_runner_mm_gather.py index dcf754fcae28..a86e77c8f52f 100644 --- a/tests/v1/worker/test_gpu_model_runner_mm_gather.py +++ b/tests/v1/worker/test_gpu_model_runner_mm_gather.py @@ -62,8 +62,7 @@ def _gather(features, cached, *, num_scheduled, shift, num_computed=0): def test_draft_shift_uses_boundary_feature_when_cached(): """The drafter's +1 look-ahead reaches the feature at offset == processed_end; when it is already cached it is used for the look-ahead - position rather than ignored. - """ + position rather than ignored.""" f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end mm_embeds, is_mm_embed = _gather([f0, f1], [f0, f1], num_scheduled=8, shift=1) @@ -76,8 +75,7 @@ def test_draft_shift_uses_boundary_feature_when_cached(): def test_draft_shift_tolerates_missing_boundary_feature(): """When the +1 look-ahead feature past the processed boundary is not yet - encoded, fall back to the token embedding instead of raising. - """ + encoded, fall back to the token embedding instead of raising.""" f0 = _feature("h0", offset=0, length=8) f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached mm_embeds, is_mm_embed = _gather([f0, f1], [f0], num_scheduled=8, shift=1) @@ -89,8 +87,7 @@ def test_draft_shift_tolerates_missing_boundary_feature(): def test_draft_shift_raises_on_interior_miss(): """A miss for a feature within the processed range (not the look-ahead - boundary) is a real invariant violation, even on the drafter path. - """ + boundary) is a real invariant violation, even on the drafter path.""" f0 = _feature("h0", offset=0, length=8) # interior, within processed range with pytest.raises(RuntimeError, match="Encoder cache miss"): _gather([f0], [], num_scheduled=8, shift=1) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 891778678b9c..308c9fe9eaab 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -219,8 +219,7 @@ def stage_write(self, *_args): def _make_capture_runner(captured: bool) -> GPUModelRunner: """Minimal V2 runner for capture_model: fakes everything except the - cudagraph_manager's needs_capture decision. - """ + cudagraph_manager's needs_capture decision.""" runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_state = SimpleNamespace(supports_mm_inputs=False) runner.cudagraph_manager = SimpleNamespace( @@ -246,8 +245,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner: def test_capture_model_locks_workspace_after_capture(monkeypatch): """A workspace resize after capture frees the buffer the captured graphs baked in, so capture_model must lock the workspace before returning - (https://github.com/vllm-project/vllm/issues/55336). - """ + (https://github.com/vllm-project/vllm/issues/55336).""" runner = _make_capture_runner(captured=True) monkeypatch.setattr( model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext @@ -268,8 +266,7 @@ def test_capture_model_locks_workspace_after_capture(monkeypatch): def test_capture_model_skips_lock_when_nothing_captured(monkeypatch): """With no graphs to capture (e.g. enforce_eager) there is nothing baked - into the workspace, so the early return must not lock it. - """ + into the workspace, so the early return must not lock it.""" runner = _make_capture_runner(captured=False) lock_calls = [] monkeypatch.setattr( @@ -283,8 +280,7 @@ def test_capture_model_skips_lock_when_nothing_captured(monkeypatch): def test_capture_model_profile_only_skips_lock(monkeypatch): """The memory-profiling capture pass runs before kernel warmup and the real capture; locking there would stop the warmup from growing the - workspace to its scheduler-realistic size. - """ + workspace to its scheduler-realistic size.""" runner = _make_capture_runner(captured=True) monkeypatch.setattr( model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index 0230ff9b83c2..74191206d363 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py @@ -82,8 +82,7 @@ def _capture_model(*, profile_only: bool = False) -> int: class _FakePlatform: """Stands in for current_platform; the global graph pool is a class attr, - matching vllm.platforms.Platform's lazy singleton. - """ + matching vllm.platforms.Platform's lazy singleton.""" _global_graph_pool: Any = GLOBAL_POOL diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index fce0f930771c..a283b6d932bc 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -120,8 +120,7 @@ def test_max_iterations(default_profiler_config): def test_restart_after_max_iterations(default_profiler_config): """A start_profile after an auto-stop must actually restart, not be silently ignored (regression test: auto-stop used to leave _active - True forever, so start() always bailed out early). - """ + True forever, so start() always bailed out early).""" default_profiler_config.max_iterations = 2 profiler = ConcreteWorkerProfiler(default_profiler_config) @@ -332,8 +331,7 @@ def test_skips_annotation_work_after_profiler_stops(self): def test_profiler_entered_during_capture(): """Profiler is used as a context manager in _warmup_and_capture, - confirming it is active during the actual graph capture run. - """ + confirming it is active during the actual graph capture run.""" runner = MagicMock() runner.compilation_config.cudagraph_num_of_warmups = 0 mock_profiler = MagicMock() diff --git a/tests/v1/worker/test_gpu_rejection_sampler_i64.py b/tests/v1/worker/test_gpu_rejection_sampler_i64.py index b69732c8e19c..769444ed64c6 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_i64.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_i64.py @@ -34,8 +34,7 @@ def _run( ) -> tuple[torch.Tensor, torch.Tensor]: """Run rejection_sample on a batch of uniform k=2 requests, with the request under test at request index `roi_req` and draft-logits row - `req_state_idx`. Returns that request's (sampled, num_sampled). - """ + `req_state_idx`. Returns that request's (sampled, num_sampled).""" num_logits = num_reqs * LOGITS_PER_REQ roi_row = roi_req * LOGITS_PER_REQ diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index a7fe0608a7ce..2bc697b08438 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -288,8 +288,7 @@ def test_v2_thinking_budget_clamps_oversized_budget(): def test_v2_thinking_budget_continues_end_prefix_from_prompt(): """A resumed prompt ending with a partial forced-end marker must not - restart the marker sequence and duplicate its first token. - """ + restart the marker sequence and duplicate its first token.""" req_states = _make_req_states([1, START, 10, 11, END_A], prompt_len=5) state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig()) state.add_request(3, SamplingParams(thinking_token_budget=3)) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index 13f1bbe45ecc..e735953d8b15 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -701,8 +701,7 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int): @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA") def test_each_microbatch_owns_its_dcp_buffer(): """Microbatch DCP lengths live in per-microbatch persistent buffers so - CPU-visible shapes stay valid under CUDA graph capture and replay. - """ + CPU-visible shapes stay valid under CUDA graph capture and replay.""" input_batch, _ = _make_cuda_input_batch([4, 4, 4, 4], [8, 9, 10, 11]) ubatches, dcp_buffers = _slice_with_dcp(input_batch, dcp_rank=0) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index f780b2f78f5f..88cbee680240 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -49,8 +49,7 @@ def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path): def test_startup_plan_fingerprint_sensitivity(plan_env): """The fingerprint is the OOM-safety key: stable for identical inputs, - different for anything the profiled value depends on. - """ + different for anything the profiled value depends on.""" fp = startup_plan.compute_plan_fingerprint base = fp(_plan_worker().vllm_config, 0, 1) assert base == fp(_plan_worker().vllm_config, 0, 1) @@ -98,8 +97,7 @@ def _snapshot(free_memory, torch_memory=0): def _profile_result(consumed, reserved_before=0, reserved_after=0): """A result whose free-memory readings agree with `consumed`, which - `memory_profiling` derives as the drop in free memory, negative when it grew. - """ + `memory_profiling` derives as the drop in free memory, negative when it grew.""" return SimpleNamespace( total_consumed=consumed, transient_peak_headroom=0, @@ -127,8 +125,7 @@ def test_profiling_fallback_declines_when_free_memory_dropped(rocm): @pytest.mark.parametrize("rocm", [True], indirect=True) def test_profiling_fallback_replaces_a_released_measurement(rocm): """A negative measurement describes the rest of the device, so it is replaced - by this process's reservation, which the rest of the device cannot move. - """ + by this process's reservation, which the rest of the device cannot move.""" result = _profile_result( consumed=-RELEASED_BY_OTHERS, reserved_after=TORCH_RESERVED, @@ -152,8 +149,7 @@ def test_profiling_fallback_never_returns_a_negative_amount(rocm): @pytest.mark.parametrize("rocm", [False], indirect=True) def test_profiling_fallback_declines_off_rocm(rocm): """Platforms that account frees eagerly keep reporting the error, so the - caller's assertion stays reachable there. - """ + caller's assertion stays reachable there.""" result = _profile_result(consumed=-RELEASED_BY_OTHERS) assert maybe_rocm_profiling_fallback(result) is None @@ -177,8 +173,7 @@ def test_execute_model_waits_previous_pp_send_before_forward( monkeypatch: pytest.MonkeyPatch, ): """Previous device handles are waited before the forward pass; the - metadata handle is left to the GroupCoordinator's reaper. - """ + metadata handle is left to the GroupCoordinator's reaper.""" import torch from vllm.sequence import IntermediateTensors diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index de541c9df670..64a81771ddd3 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -358,8 +358,7 @@ def test_warmup_respects_available_block_count(): def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout): """The zeroer must zero every byte of the target block in every layer and nothing outside it — per head-group region under LHBNC, and never past the target block's - tile under block-major layouts (no out-of-bounds writes, no clobbering). - """ + tile under block-major layouts (no out-of-bounds writes, no clobbering).""" device = torch.device("cuda") num_blocks, num_layers = 4, 2 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index bce4c8c35838..47a27795ca18 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -645,8 +645,7 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: """Build a MambaSpecDecodeGPUContext stand-in exposing only the four - per-request staging buffers touched by stage_postprocess_inputs_to_gpu. - """ + per-request staging buffers touched by stage_postprocess_inputs_to_gpu.""" ctx = MagicMock() ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) @@ -658,8 +657,7 @@ def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): """stage_postprocess_inputs_to_gpu writes each request's state_idx and scheduled/computed/draft counts into slot `i` for `req_ids[i]`, leaves - slots past `num_reqs` untouched, and mirrors the values to the GPU buffers. - """ + slots past `num_reqs` untouched, and mirrors the values to the GPU buffers.""" device = torch.device("cpu") max_num_reqs = 8 ctx = _make_staging_ctx(max_num_reqs, device) @@ -731,8 +729,7 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): """If preprocess_mamba didn't populate mamba_state_idx for a req in the - batch, staging must fail loudly rather than silently writing a stale index. - """ + batch, staging must fail loudly rather than silently writing a stale index.""" device = torch.device("cpu") ctx = _make_staging_ctx(max_num_reqs=4, device=device) scheduler_output = _make_postprocess_scheduler_output( @@ -810,8 +807,7 @@ def _run_gpu_postprocess( device: torch.device, ) -> None: """Initialize the GPU context against `block_table`, run the fused - postprocess kernel for `req_ids`, and synchronize. - """ + postprocess kernel for `req_ids`, and synchronize.""" def t(values): return torch.tensor(values, dtype=torch.int32, device=device) diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index d4123269a952..04f456a7c9cd 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -17,8 +17,7 @@ def _fail(*args, **kwargs): @pytest.mark.parametrize("max_num_reqs", [1, 0]) def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 - the warmup must be skipped without touching the worker callbacks. - """ + the warmup must be skipped without touching the worker callbacks.""" runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) assert ( diff --git a/tests/v1/worker/test_prompt_embeds_state.py b/tests/v1/worker/test_prompt_embeds_state.py index f28c6d2e01ae..1893d0f3470e 100644 --- a/tests/v1/worker/test_prompt_embeds_state.py +++ b/tests/v1/worker/test_prompt_embeds_state.py @@ -78,8 +78,7 @@ def _apply( def test_overlay_chunked_prefill_and_decode(): """Rows within the embeds range come from prompt_embeds at the num_computed offset; requests without embeds and requests past their - embeds length (decode) keep the base embedding. - """ + embeds length (decode) keep the base embedding.""" state = _make_state() embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32) embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32) @@ -98,8 +97,7 @@ def test_overlay_chunked_prefill_and_decode(): def test_overlay_clamps_to_embeds_length(): """A window straddling the end of the prompt embeds writes only the - in-range rows (e.g. final prefill chunk + sampled token). - """ + in-range rows (e.g. final prefill chunk + sampled token).""" state = _make_state() embeds = torch.randn(4, HIDDEN, dtype=torch.float32) state.add_request(3, _NewReqData("a", embeds)) @@ -114,8 +112,7 @@ def test_overlay_clamps_to_embeds_length(): def test_overlay_respects_is_token_ids_mask(): """Mixed mode: positions marked as real token ids keep the base - embedding; only embed positions are overwritten. - """ + embedding; only embed positions are overwritten.""" state = _make_state() embeds = torch.randn(5, HIDDEN, dtype=torch.float32) is_token_ids = [True, False, False, True, False] @@ -133,8 +130,7 @@ def test_overlay_respects_is_token_ids_mask(): def test_index_reuse_clears_stale_entry(): """A request added at a previously-used index without embeds must not - inherit the prior occupant's pointer-table entry. - """ + inherit the prior occupant's pointer-table entry.""" state = _make_state() state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32))) # A second live embeds request so the kernel actually launches (the diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index cc52f28e4b5c..7824c8000d2b 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -84,8 +84,7 @@ def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None: def test_workspace_lock_blocks_growth_and_unlock_restores(monkeypatch) -> None: """Once locked, oversized requests fail loudly instead of reallocating the - buffer that captured CUDA graphs point at; unlock restores growth. - """ + buffer that captured CUDA graphs point at; unlock restores growth.""" monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) diff --git a/tools/report_build_time_ninja.py b/tools/report_build_time_ninja.py index 916e29a35f92..3b43fc988ef8 100644 --- a/tools/report_build_time_ninja.py +++ b/tools/report_build_time_ninja.py @@ -56,8 +56,7 @@ class Target: def __init__(self, start, end): """Creates a target object by passing in the start/end times in seconds - as a float. - """ + as a float.""" self.start = start self.end = end # A list of targets, appended to by the owner of this object. @@ -105,8 +104,7 @@ def DescribeTargets(self): def ReadTargets(log, show_all): """Reads all targets from .ninja_log file |log_file|, sorted by duration. - The result is a list of Target objects. - """ + The result is a list of Target objects.""" header = log.readline() assert header == "# ninja log v5\n", "unrecognized ninja log version {!r}".format( header @@ -155,8 +153,7 @@ def GetExtension(target, extra_patterns): For targets that generate multiple outputs it is important to return a consistent 'canonical' extension. Ultimately the goal is to group build steps - by type. - """ + by type.""" for output in target.targets: if extra_patterns: for fn_pattern in extra_patterns.split(";"): diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index dfcba987d60b..da4bc310d7e4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1974,8 +1974,7 @@ def is_rdna_aiter_enabled(cls) -> bool: enabled aiter. Only aiter's Triton kernels exist on gfx12 (no CK build), so this deliberately stays off the gfx9/CK `@if_aiter_supported` umbrella and gates only the Triton paths rdna4 uses. The gfx12 analog of - `is_enabled()`. - """ + `is_enabled()`.""" if not current_platform.is_rocm() or not IS_AITER_FOUND: return False from vllm.platforms.rocm import on_rdna4 diff --git a/vllm/benchmarks/datasets/create_txt_slices_dataset.py b/vllm/benchmarks/datasets/create_txt_slices_dataset.py index e4555bce9ae3..1bcd93c8ae11 100644 --- a/vllm/benchmarks/datasets/create_txt_slices_dataset.py +++ b/vllm/benchmarks/datasets/create_txt_slices_dataset.py @@ -65,8 +65,7 @@ def create_txt_slices_jsonl( trust_remote_code: bool = False, ) -> None: """Read *input_path*, slice it into prompts, and write JSONL to - *output_path*. - """ + *output_path*.""" tokenizer = AutoTokenizer.from_pretrained( tokenizer_name, trust_remote_code=trust_remote_code ) diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 3456fc633786..8f9e87cf8331 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -4383,8 +4383,7 @@ def sample( def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]: """Generate tokens that decode and re-encode to exactly - target_length. - """ + target_length.""" # Generate random tokens tokens = np.random.randint(0, vocab_size, size=target_length).tolist() diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 5adb54c47f16..2b4027b502c3 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -22,8 +22,7 @@ class StreamedResponseHandler: """Handles streaming HTTP responses by accumulating chunks until complete - messages are available. - """ + messages are available.""" def __init__(self): self.buffer = "" @@ -31,8 +30,7 @@ def __init__(self): def add_chunk(self, chunk_bytes: bytes) -> list[str]: """Add a chunk of bytes to the buffer and return any complete - messages. - """ + messages.""" chunk_str = self._decoder.decode(chunk_bytes) self.buffer += chunk_str diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index fb981a3600f0..4ac694931666 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -150,8 +150,7 @@ def compute_hash(self, vllm_config: VllmConfig) -> str: def compile_context(self, compile_range: Range) -> Generator[None, None, None]: """Provide compilation context for the duration of compilation to set any torch global properties we want to scope to a single Inductor - compilation (e.g. partition rules, pass context). - """ + compilation (e.g. partition rules, pass context).""" with pass_context(compile_range): if self.compilation_config.use_inductor_graph_partition: with inductor_partition_rule_context( diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 0195a0879923..67aff0fe1ea7 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -153,8 +153,7 @@ def _get_vllm_functorch_config() -> dict[str, Any]: """Return the functorch config overrides that vLLM applies at compile time. Used by both set_functorch_config() and get_inductor_factors() to ensure - the compile-time config and cache key are always consistent. - """ + the compile-time config and cache key are always consistent.""" cfg: dict[str, Any] = {} if not envs.VLLM_USE_MEGA_AOT_ARTIFACT: cfg["bundled_autograd_cache"] = False diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index bb41c9550d6f..b2961debcec2 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -66,8 +66,7 @@ def _view_flashinfer_nvfp4_scale_out_as_int32( def _norm_input_weight_dtype_match(match: pm.Match) -> bool: """Prevent fusion when the norm input and weight dtypes differ (e.g. a Gemma - fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm. - """ + fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm.""" for node in match.nodes: if node.target == _IR_RMS_NORM_OP: x, weight = node.args[0], node.args[1] diff --git a/vllm/compilation/passes/utility/split_coalescing.py b/vllm/compilation/passes/utility/split_coalescing.py index 9b60d10fba37..f16623a1aa51 100644 --- a/vllm/compilation/passes/utility/split_coalescing.py +++ b/vllm/compilation/passes/utility/split_coalescing.py @@ -30,8 +30,7 @@ class SplitCoalescingPass(VllmInductorPass): """Replace duplicate ``split_with_sizes`` nodes with a single canonical - node when they share the same input tensor and split sizes. - """ + node when they share the same input tensor and split sizes.""" @VllmInductorPass.time_and_log def __call__(self, graph: fx.Graph) -> None: diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index f1b5e90f2ac7..8ccc7332bd94 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -36,8 +36,7 @@ class CompilationMode(enum.IntEnum): """The compilation approach used for torch.compile-based compilation of the - model. - """ + model.""" NONE = 0 """No torch.compile compilation is applied, model runs in fully eager pytorch mode. @@ -337,8 +336,7 @@ def log_enabled_passes(self) -> None: class DynamicShapesType(str, enum.Enum): """Types of dynamic shapes handling in torch.compile(). see Dynamic shapes and vllm guard dropping in torch_compile.md - for more details. - """ + for more details.""" BACKED = "backed" """Use backed dynamic shapes. torch.compile() guards on backed dynamic diff --git a/vllm/config/model.py b/vllm/config/model.py index a07fbbc7a257..9ee35069dab4 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1005,8 +1005,7 @@ def _resolve_mm_device_do_normalize( def _get_transformers_backend_cls(self) -> str: """Determine which Transformers modeling backend class will be used if - `model_impl` is set to `transformers` or `auto`. - """ + `model_impl` is set to `transformers` or `auto`.""" cls = "Transformers" # If 'hf_config is not hf_text_config' it's a nested config, i.e. multimodal cls += "MultiModal" if self.hf_config is not self.hf_text_config else "" diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 5f93fce0aa58..09e184b8ce2b 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -524,8 +524,7 @@ def is_multimodal_pruning_enabled(self): def get_video_pruning_spec(self) -> tuple[VideoPruningMethod, float] | None: """Return `(method, rate)` when video pruning is enabled, else None. - `rate` is the fraction of video tokens to prune. - """ + `rate` is the fraction of video tokens to prune.""" if self.video_pruning_rate is not None and self.video_pruning_rate > 0: return (self.video_pruning_method, float(self.video_pruning_rate)) return None diff --git a/vllm/config/observability.py b/vllm/config/observability.py index c964fe904db5..e24668b1a6ae 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -156,8 +156,7 @@ def _validate_collect_detailed_traces( cls, value: list[DetailedTraceModules] | None ) -> list[DetailedTraceModules] | None: """Handle the legacy case where users might provide a comma-separated - string instead of a list of strings. - """ + string instead of a list of strings.""" if value is not None and len(value) == 1 and "," in value[0]: value = cast(list[DetailedTraceModules], value[0].split(",")) return value diff --git a/vllm/config/reasoning.py b/vllm/config/reasoning.py index e00d9adc27ab..0ac93be0823a 100644 --- a/vllm/config/reasoning.py +++ b/vllm/config/reasoning.py @@ -47,15 +47,13 @@ class ReasoningConfig: @property def enabled(self) -> bool: """Returns True if reasoning is enabled (i.e. if token IDs have been - initialized), False otherwise. - """ + initialized), False otherwise.""" return self._enabled @property def reasoning_start_token_ids(self) -> list[int] | None: """Token IDs derived from `reasoning_start_str`. Set automatically by - `initialize_token_ids`. Not intended to be configured directly. - """ + `initialize_token_ids`. Not intended to be configured directly.""" return self._reasoning_start_token_ids @property diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 96340df59117..63f729a832c9 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -543,8 +543,7 @@ class SpeculativeConfig: @staticmethod def _acceptance_length_to_rates(length: float, n: int) -> list[float]: """Mean acceptance length to unconditional per-position rates, using - the minimum-variance schedule. - """ + the minimum-variance schedule.""" num_drafts = length - 1 # expected number of accepted draft tokens num_full = int(num_drafts) return ( @@ -558,8 +557,7 @@ def _resolve_synthetic_acceptance_rates( length: float | None, ) -> list[float]: """Return per-position unconditional acceptance rates from exactly one - of `rates` or `length` (validates range, length, and monotonicity). - """ + of `rates` or `length` (validates range, length, and monotonicity).""" if (rates is None) == (length is None): raise ValueError( "rejection_sample_method='synthetic' requires exactly one of " diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 192b8bc8707c..ed724746196a 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -84,8 +84,7 @@ def decorator(cls: type[ConfigT]) -> type[ConfigT]: def get_field(cls: ConfigType, name: str) -> Any: """Get the default factory field of a dataclass by name. Used for getting - default factory fields in `EngineArgs`. - """ + default factory fields in `EngineArgs`.""" if not is_dataclass(cls): raise TypeError("The given class is not a dataclass.") try: @@ -122,8 +121,7 @@ def is_init_field(cls: ConfigType, name: str) -> bool: def replace(dataclass_instance: ConfigT, /, **kwargs) -> ConfigT: """Like [`dataclasses.replace`](https://docs.python.org/3/library/dataclasses.html#dataclasses.replace), but compatible with Pydantic dataclasses which use `pydantic.fields.Field` instead - of `dataclasses.field` - """ + of `dataclasses.field`""" cls = type(dataclass_instance) dataclass_dict = dataclass_instance.__dict__ dataclass_dict = {k: v for k, v in dataclass_dict.items() if is_init_field(cls, k)} diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 174e184c538d..e5710dced396 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -153,8 +153,7 @@ class OptimizationLevel(IntEnum): def enable_norm_fusion(cfg: "VllmConfig") -> bool: """Enable if either RMS norm or quant FP8 custom op is active; - otherwise Inductor handles fusion. - """ + otherwise Inductor handles fusion.""" return ( cfg.compilation_config.is_custom_op_enabled("rms_norm") or cfg.compilation_config.is_custom_op_enabled("quant_fp8") @@ -744,8 +743,7 @@ def use_v2_model_runner(self) -> bool: def _is_dflash2_draft(self) -> bool: """Whether the DFlash draft is a DFlash2 one, by the architecture the - speculator selects on (v1/worker/gpu/spec_decode/__init__.py). - """ + speculator selects on (v1/worker/gpu/spec_decode/__init__.py).""" spec = self.speculative_config if spec is None or spec.method != "dflash": return False diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index bf28278f7cde..45a635d61389 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -70,8 +70,7 @@ def resume(self, tags: list[str] | None = None) -> None: def state(self) -> SleepModeState: """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle - (suspended) engine from a healthy-serving one (see RFC #34303). - """ + (suspended) engine from a healthy-serving one (see RFC #34303).""" return self._state # -- Capability introspection (no instance required) -- @@ -84,30 +83,26 @@ def is_supported(cls) -> bool: @classmethod def preserves_communicators(cls) -> bool: """If False, collective communicators (e.g. NCCL) are destroyed by - ``suspend`` and the executor must re-initialize them on ``resume``. - """ + ``suspend`` and the executor must re-initialize them on ``resume``.""" return False @classmethod def preserves_compiled_artifacts(cls) -> bool: """If True, torch.compile / JIT kernels survive suspend/resume and need - not be recompiled on resume. - """ + not be recompiled on resume.""" return False @classmethod def preserves_graphs_with_communicators(cls) -> bool: """If True, CUDA graphs containing collective communicators (e.g. NCCL) stay valid after resume. False when communicators are rebuilt (embedded - comm handles go stale). - """ + comm handles go stale).""" return False @classmethod def supports_durable_storage(cls) -> bool: """If True, suspended state can be persisted beyond the process - lifetime (disk or object storage) and restored in a new process. - """ + lifetime (disk or object storage) and restored in a new process.""" return False diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index 43e615a3984b..a8298a459fa4 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -156,8 +156,7 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) def should_nccl_symm_mem_ag_rs() -> bool: """Check whether NCCL symmetric memory should be used for - AllGather / ReduceScatter collectives. - """ + AllGather / ReduceScatter collectives.""" from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_enabled, ) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index f457cf61e8a9..bafc5728b85e 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -650,8 +650,7 @@ def wait_until_ready(self): def shutdown(self): """If this is an idle reader, wakes it up so it can clean up and shut - down - """ + down""" self.shutting_down = True if self._spin_condition is not None: self._spin_condition.cancel() @@ -761,8 +760,7 @@ def timeout_ms(self) -> int: def should_warn(self) -> bool: """Returns true if it's time to log a warning for a timeout that is not - indefinite - """ + indefinite""" if self._should_warn: elapsed = time.monotonic() - self.started if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning: diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index afadcfd42f20..2fd4b01059db 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -15,8 +15,7 @@ class CacheEntry: """A single cache entry. Read `.block_ids` and `.ready` freely; - mutations only through EmbeddingCache methods. - """ + mutations only through EmbeddingCache methods.""" __slots__ = ("block_ids", "_pin_count") diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index d05ca8e00cde..61fef468d323 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -84,8 +84,7 @@ def set_transfer_context( # noqa: B027 @property def needs_profile_buffer_reservation(self) -> bool: """Whether the profile path must run a dummy collective operation to reserve - communication buffers. - """ + communication buffers.""" return True def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: @@ -443,8 +442,7 @@ def _init_registered_buffers(self) -> None: def _exchange_remote_send_meta(self) -> None: """Exchange per-layer per-tensor metadata so receivers can compute - remote RDMA addresses at transfer time. - """ + remote RDMA addresses at transfer time.""" local_meta: dict[tuple[int, int], tuple[int, int, int]] = {} for layer_idx, layer_tensors in enumerate(self._all_expert_weights): for t_idx, t in enumerate(layer_tensors): diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 9821f912b295..b645e97c4dd4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -51,8 +51,7 @@ def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None): class KVOutputAggregator: """Utility class to aggregate the output of all workers into a single - output corresponding to Rank 0 for scheduler. - """ + output corresponding to Rank 0 for scheduler.""" def __init__(self, expected_finished_count: int): # Complete transfer tracker. Used to track finished requests diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 320d9ce8044c..e9a3e5d99a3c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -278,8 +278,7 @@ def wait_for_save(self) -> None: def _on_write_done(self, req_id: str, future: Future) -> None: """Surface any exception from the disk-write thread and drop the - completed future from the in-flight tracking dict. - """ + completed future from the in-flight tracking dict.""" self._req_futures.pop(req_id, None) exc = future.exception() if exc is not None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index 31b20f0ed75b..51f4ddf349e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -231,8 +231,7 @@ def __init__(self, request: "Request"): #### def needs_retrieve(self) -> bool: """Check whether the current request needs retrieve, will be used - update_stage_after_alloc - """ + update_stage_after_alloc""" return ( self.num_lmcache_hit_blocks > self.num_vllm_hit_blocks and self.state != LMCacheMPRequestState.READY @@ -240,8 +239,7 @@ def needs_retrieve(self) -> bool: def is_ready_for_retrieving(self) -> bool: """Check whether the current request is ready for retrieving, - will be used in process_loading_requests - """ + will be used in process_loading_requests""" return ( self.state == LMCacheMPRequestState.WAITING_FOR_LOAD and self.needs_retrieve() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 4e6844034c7c..b22387e98b73 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -718,8 +718,7 @@ def get_sw_clipped_blocks( def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1). - """ + always recomputes the last token and must start from h(N-1).""" if self._has_mamba and num_prompt_tokens > 1: return num_prompt_tokens - 1 return num_prompt_tokens @@ -730,8 +729,7 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: derive h(N) correctly. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled. - """ + request is preempted and rescheduled.""" params = request.kv_transfer_params if ( params is not None @@ -1677,8 +1675,7 @@ async def _build_transfer_params( def _bind_sender_thread_device(self) -> None: """ThreadPoolExecutor initializer — binds each pool thread to the correct CUDA device. CUDA device selection is thread-local, so - without this, NVLink transfers fail for TP ranks > 0. - """ + without this, NVLink transfers fail for TP ranks > 0.""" current_platform.set_device(self.device_id) def _send_blocks( @@ -1904,8 +1901,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults: def get_kv_connector_stats(self) -> KVConnectorStats | None: """Return transfer stats collected since the last call, or None - if nothing has been recorded in this interval. - """ + if nothing has been recorded in this interval.""" if self.xfer_stats.is_empty(): return None return self.xfer_stats.clone_and_reset() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 9e75c880d5e6..d05b553b64b9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -73,8 +73,7 @@ def contains(self, group_id: int, block_hash: BlockHash) -> bool: class MooncakeStoreCoordinator: """Mirror of ``HybridKVCacheCoordinator.find_longest_cache_hit`` over an - ``ExternalCachedBlockPool``. - """ + ``ExternalCachedBlockPool``.""" def __init__( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index fca59a12e0ba..762749987b5d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -2629,8 +2629,7 @@ def lookup( non_block: bool = False, ) -> MooncakeLookupResult | None: """If non_block is True, will return None until the result is ready, - so the caller retries on a later step. - """ + so the caller retries on a later step.""" future = self.futures.get(req_id) if future is None: future = self.executor.submit(self._lookup, num_tokens, list(block_hashes)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index ea6c363df335..f15eee9bd57f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -548,8 +548,7 @@ def unmap_request_id( self, request_id: ReqId, transfer_id: TransferId | None = None ): """Unmap request_id/transfer_id. Uses transfer_id for lookup if - exact request_id match fails (handles input_processor mutation). - """ + exact request_id match fails (handles input_processor mutation).""" if request_id in self.request_id_to_transfer_id: tid = self.request_id_to_transfer_id[request_id] del self.request_id_to_transfer_id[request_id] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index ed9123fc6a98..e264a94aafd8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -400,8 +400,7 @@ def _prefill_backoff(self) -> int: def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. The number of prompt tokens to load from the prefiller. Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder - will recompute locally. - """ + will recompute locally.""" backoff = self._prefill_backoff() if backoff and num_prompt_tokens > backoff: return num_prompt_tokens - backoff @@ -414,8 +413,7 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: multi-module MTP it is the drafter's whole lookahead window. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled. - """ + request is preempted and rescheduled.""" backoff = self._prefill_backoff() params = request.kv_transfer_params if ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index b261a13883c9..7bcdd2531bfe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1873,8 +1873,7 @@ def _build_mamba_remote( ) -> np.ndarray: """Build remote desc regions (conv sub-projections + ssm) per layer. For hetero-TP, each D rank reads only its sub-projection slice from - the P rank. Returns an Nx3 uint64 array. - """ + the P rank. Returns an Nx3 uint64 array.""" assert nixl_agent_meta.kv_caches_base_addr, ( "Remote KV cache base addresses must not be empty." ) @@ -2657,8 +2656,7 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): @cached_property def _attention_kv_caches(self) -> list[torch.Tensor]: """Device KV caches of attention layers (mamba states excluded), - as consumed by the receive post-process. - """ + as consumed by the receive post-process.""" assert self.device_kv_caches, ( "_attention_kv_caches accessed before register_kv_caches" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index a61e5f8f7838..01c956a5251d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -315,8 +315,7 @@ def _send_registration_to_p( executes on the writer; the handshake runs on the background executor and the request is re-queued onto ``_reg_send_inbox`` once it completes (at which point ``_ensure_handshake`` returns ``None`` and we - send directly). - """ + send directly).""" remote_pp_size = reg_data.get("remote_pp_size", 1) fut = self._ensure_handshake( reg_data["remote_engine_id"], @@ -508,8 +507,7 @@ def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds: ``BlockIds`` is canonically a tuple of per-group lists, but some registration payloads collapse a single-group case to a flat list. Re-wrap that case so downstream group-aware helpers see a - consistent shape. - """ + consistent shape.""" if block_ids and not isinstance(block_ids[0], (list, tuple)): return (list(block_ids),) return block_ids diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py index 3c51fa277feb..27f08c1fd834 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py @@ -37,8 +37,7 @@ def canonical_format_id(kv_cache_layout: str) -> str: Canonical pages keep the worker's KV layout family, so the id couples the format version with that family; consumers must match it exactly. The family keeps its historical NHD/HND spelling so ids stay stable for KV persisted - before the layout enum existed. - """ + before the layout enum existed.""" layout = KVCacheLayout[kv_cache_layout] legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"} family = legacy.get(layout, layout.name.lower()) @@ -122,8 +121,7 @@ def _interleave_cp_tokens( ctx: _RankContext, ) -> tuple[CopyRun, ...]: """Place each region's num_tokens rows at their canonical token positions, - one run per chunk of interleaved tokens. - """ + one run per chunk of interleaved tokens.""" runs: list[CopyRun] = [] for region in regions: if ctx.cp_size == 1: @@ -258,8 +256,7 @@ def _attention_byte_regions( cp_size: int, ) -> list[ByteRegion] | None: """Byte regions of an attention page, given this rank's head shard. - None when the physical layout is not recognized (fail closed). - """ + None when the physical layout is not recognized (fail closed).""" bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size if tuple(kv_cache.shape) == (num_blocks, heads, bs, 2 * head_size): return _packed_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size) @@ -373,8 +370,7 @@ def _is_exact_partition(intervals: list[tuple[int, int]], size: int) -> bool: def _verify_tiling(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None: """Whichever ranks a block elects as writers must tile the canonical page - exactly once, and each rank's runs must cover exactly its local page. - """ + exactly once, and each rank's runs must cover exactly its local page.""" size = per_rank[0].canonical_page_size_bytes num_writers = per_rank[0].num_writers for mapping in per_rank: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index a8789d81c4ef..6dcfcb5ec098 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -79,8 +79,7 @@ def get_offloading_event_group_spec( class _OffloadEventMetadata: """BlockStored payload snapshot for one OffloadKey, captured while the Request is available and kept until the final matching removal event. - ``medium`` and ``ownership`` are forwarded from the OffloadingEvent. - """ + ``medium`` and ``ownership`` are forwarded from the OffloadingEvent.""" # The chunk's constituent block hashes; the last one is the OffloadKey. block_hashes: tuple[BlockHash, ...] @@ -256,8 +255,7 @@ def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEven def reset(self) -> None: """Drop all tracked state; pending payloads are stale after a - manager cache reset. - """ + manager cache reset.""" self._pending_event_metadata.clear() def _build_event_metadata( @@ -268,8 +266,7 @@ def _build_event_metadata( ) -> _OffloadEventMetadata: """Build the payload snapshot for one offloaded chunk: its constituent per-block hashes, the whole chunk's tokens, and the - per-block ``block_size``. - """ + per-block ``block_size``.""" hashes_per_chunk = group_config.hashes_per_chunk assert hashes_per_chunk > 0 assert chunk_idx >= 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 961980ee896c..d6c210a6f9bf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -676,8 +676,7 @@ def _maximal_prefix_lookup( start_chunk_idx: int, ) -> int | None: """Return the number of consecutive offloaded chunks from the start, - or None if the backend deferred a lookup. - """ + or None if the backend deferred a lookup.""" hit_count = 0 defer_lookup = False for local_idx, key in enumerate(keys): @@ -712,8 +711,7 @@ def _sliding_window_lookup( """Return the end index (in `keys`) of the last run of `sliding_window_size` consecutive hits, scanning from the end. The first run may need a larger window for a partial rightmost chunk. - Returns 0 on miss, None if the backend deferred a lookup. - """ + Returns 0 on miss, None if the backend deferred a lookup.""" defer_lookup = False pending_in_window = False consecutive_hits = 0 diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index f8304200749c..b32269ec2b60 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -59,8 +59,7 @@ def ensure_divisibility(numerator, denominator): def divide(numerator, denominator): """Ensure that numerator is divisible by the denominator and return - the division value. - """ + the division value.""" ensure_divisibility(numerator, denominator) return numerator // denominator diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 74a04cc9d8f7..48f9e4f21281 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -141,8 +141,7 @@ def register_engine( class_name: str | None = None, ) -> None: """Register a trainer engine. Same conventions as - `WeightTransferEngineFactory.register_engine`. - """ + `WeightTransferEngineFactory.register_engine`.""" if name in cls._registry: raise ValueError( f"Weight transfer trainer engine '{name}' is already registered." diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 83b5d15f58af..15864811e96c 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -40,8 +40,7 @@ class IPCWeightTransferInitInfo(WeightTransferInitInfo): `packed` is a must-agree wire param: the trainer ships it here at the init handshake so the worker decodes with the same setting the trainer encoded with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does - not need the buffer size (producer-only). - """ + not need the buffer size (producer-only).""" packed: bool = False @@ -54,8 +53,7 @@ class IPCTrainerInitInfo(TrainerInitInfo): `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The trainer propagates `packed` to the worker at `trainer_init` so the two sides - cannot disagree. `backend` is the factory dispatch key. - """ + cannot disagree. `backend` is the factory dispatch key.""" backend: ClassVar[str] = "ipc" packed: bool = False diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index 09e31f0291af..57a17b9830a4 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -118,8 +118,7 @@ def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict: """Serialize a worker init info for `init_weight_transfer_engine`, dropping the unset rendezvous field (the UID in TCP mode) so the wire payload carries only the mode actually in use. Shared by the dense and sparse trainer - engines so the two cannot drift. - """ + engines so the two cannot drift.""" return {key: value for key, value in asdict(init_info).items() if value is not None} diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 91298777a918..5b98845870fd 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -69,8 +69,7 @@ class NCCLTrainerInitInfo(TrainerInitInfo): propagates them to the worker at `trainer_init` so the two sides cannot disagree. Note this defaults to packed, unlike the worker-side `NCCLWeightTransferInitInfo`, whose default only applies when no trainer - ships a value. `backend` is the factory dispatch key. - """ + ships a value. `backend` is the factory dispatch key.""" backend: ClassVar[str] = "nccl" @@ -365,8 +364,7 @@ def send_weights(self) -> None: def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None: """Iterate the source (materializing each tensor — a collective on all ranks) and, on the sender, broadcast from rank 0, packed or one-by-one. - Non-sender ranks only replay the iteration to stay in the collective. - """ + Non-sender ranks only replay the iteration to stay in the collective.""" if not self.is_sender: for _ in source: pass diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py index eeda36596577..0d913628174a 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_common.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py @@ -214,8 +214,7 @@ def __init__( def class_of(self, name: str) -> int: """The name's owner class — the planner's bucketing key, since all names - of a chunk must share one producer. - """ + of a chunk must share one producer.""" return self._class_of.get(name, 0) def owners(self, name: str) -> list[int]: @@ -286,8 +285,7 @@ def validate(self) -> None: def bind(self, actors: list, produce_methods: list, consumer_id: int) -> None: """Attach this consumer's producer handles. Rebuilt wholesale per init, never appended to: every owner index is a position in these lists, so a - rejoining engine that re-inits must not shift them. - """ + rejoining engine that re-inits must not shift them.""" if len(actors) != len(produce_methods): raise ValueError("actors and produce_methods must be parallel") self._actors = list(actors) @@ -350,6 +348,5 @@ def buffer_alloc_bytes(nbytes: int, presize: int = 0) -> int: optional ``presize`` floor, and a coarse 256MB round-up, so the buffer is allocated ONCE and never regrows. Regrowth is a correctness hazard, not just a perf one -- see ``buffer_presize_gb``. Shared by the consumer's receive buffers - and the producer's serve rings. - """ + and the producer's serve rings.""" return max(nbytes, presize, -(-nbytes // (256 << 20)) * (256 << 20)) diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index cfe84d2dd033..96dd53f0fdb9 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -134,8 +134,7 @@ class _PendingPull: ``targets``/``blob`` must stay strongly referenced until it completes: ``set_target_for_ref`` stores WEAKREFS, so dropping them silently reroutes the - transfer into a fallback buffer. - """ + transfer into a fallback buffer.""" ref: "Any" keys: "list[FetchKey]" @@ -562,8 +561,7 @@ def _resolve_consumer_id(self, init_info: ShardedRDTWeightTransferInitInfo) -> i def _build_static_plan(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: """Build the chunk/free plan once. It never changes across syncs, so - ``update_weights`` needs no per-sync names. - """ + ``update_weights`` needs no per-sync names.""" if not init_info.group_lens: raise ValueError( "Sharded RDT engine requires init_info.group_lens (the gather-" @@ -587,8 +585,7 @@ def _preregister_at_init(self) -> None: Both sides are sized from the static plan: receive buffers are ``ring_depth`` slots at the largest chunk's ``pack_bytes``, and each bound producer is asked to pre-register a serve ring at the max bytes this - consumer will pull from it. A no-op when this worker has no chunks. - """ + consumer will pull from it. A no-op when this worker has no chunks.""" plan = self._cached_plan if plan is None or not plan.chunks: return @@ -633,8 +630,7 @@ def _global_worker_index(self) -> int: ``data_parallel_index``, not ``data_parallel_rank``: vLLM resets the latter to 0 in a dense worker but keeps the former as the distinct global DP rank. Same formula as the sibling ``nccl_engine``, so dense-via-TP and - MoE-via-DP+EP both yield distinct 0..C-1. - """ + MoE-via-DP+EP both yield distinct 0..C-1.""" pc = self.parallel_config return pc.data_parallel_index * pc.world_size + pc.rank # world_size = TP*PP @@ -642,16 +638,14 @@ def _num_consumers(self) -> int: """Total inference-worker count. Prefers the driver-supplied ``init_info.num_consumers`` (authoritative -- the driver knows the whole fleet); else ``world_size_across_dp``, the same stride - ``_global_worker_index`` indexes with, so the two agree at any pp. - """ + ``_global_worker_index`` indexes with, so the two agree at any pp.""" if self._num_consumers_override > 0: return self._num_consumers_override return self.parallel_config.world_size_across_dp def start_weight_update(self) -> None: """Put the model's params on meta so layerwise reload streams them in - as each layer's slices land. Baked replay uses checkpoint format. - """ + as each layer's slices land. Baked replay uses checkpoint format.""" from vllm.model_executor.model_loader.reload import ( initialize_layerwise_reload, ) @@ -660,8 +654,7 @@ def start_weight_update(self) -> None: def finish_weight_update(self) -> None: """Drain the deferred pull/process pipeline (so every layer is fully - loaded) before finalizing the layerwise reload. - """ + loaded) before finalizing the layerwise reload.""" from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, ) @@ -674,8 +667,7 @@ def update_weights(self, update_info: dict[str, Any]) -> None: device sync: post-processing is deferred to background threads and a sync here would block on them and serialize the pull/process pipeline. Completion is guaranteed by ``drain_pending`` in - ``finish_weight_update``. - """ + ``finish_weight_update``.""" self.receive_weights(self.parse_update_info(update_info)) def receive_weights( @@ -717,8 +709,7 @@ def _build_fake_weights( device: torch.device, ) -> list[tuple[str, torch.Tensor]]: """Zero-storage lazies for ``names``, dtype/shape from the init metadata, - all feeding the bake's recording sink. - """ + all feeding the bake's recording sink.""" return [ ( name, @@ -986,8 +977,7 @@ def _restore_after_dry_run(self, model: torch.nn.Module) -> None: def _ensure_proc_worker(self) -> None: """Lazily create the per-slot events, the background CUDA stream, the work - queue, and the single processing thread. Idempotent. - """ + queue, and the single processing thread. Idempotent.""" if self._proc_thread is not None: return import queue @@ -1046,8 +1036,7 @@ def _proc_worker_loop(self) -> None: def _mark_slot_done(self, slot: int) -> None: """Publish that a queued item's read-done event has been recorded (or the item failed) so a pull waiting to reuse ``slot`` can proceed to its - CUDA-event synchronize. - """ + CUDA-event synchronize.""" with self._slot_cv: self._slot_done[slot] += 1 self._slot_cv.notify_all() @@ -1063,8 +1052,7 @@ def drain_pending(self) -> None: """Block until the background thread has processed every queued item and its stream work is complete, then re-raise any error it hit. Called from the worker's ``finish_weight_update`` before ``finalize_layerwise_reload`` - so every baked layer is fully loaded (and ``info.reset()``-ed) first. - """ + so every baked layer is fully loaded (and ``info.reset()``-ed) first.""" if self._proc_queue is not None: self._proc_queue.join() # every put() item task_done()'d # The scatter thread feeds the quant thread, so join it SECOND (all @@ -1116,8 +1104,7 @@ def _chunk_module_scatters( which is exactly what an owner class is. The cut is a pure function of the bake and the ownership table. Copy order within a chunk is bake order, and a module's copies may span chunks (materialize/quant fire on - its first/last chunk; see ``_build_call_plan``). - """ + its first/last chunk; see ``_build_call_plan``).""" assert self._router is not None by_class: dict[int, list[_Scatter]] = {} for copies in modules: @@ -1248,8 +1235,7 @@ def _signal_group_done(self, group_idx: int) -> None: pulls nothing from. Refs are held and drained in ``drain_pending`` so every signal has EXECUTED before the sync ends: ``begin_sync`` resets the counters, and a straggler landing in the next sync would credit a - group it does not belong to. - """ + group it does not belong to.""" assert self._router is not None self._pending_frees.extend(self._router.free_group(group_idx)) @@ -1384,8 +1370,7 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: the modules' scatters via ``ready``; touches only the scattered params (never a receive slot), so it can overlap subsequent chunks' RDMA and scatters. ``info.reset()`` is what makes finalize skip the layer — - drain_pending joins the quant queue before finalize runs. - """ + drain_pending joins the quant queue before finalize runs.""" from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) @@ -1418,8 +1403,7 @@ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: def _quant_worker_loop(self) -> None: """Dedicated quant thread: drains (completed_modules, scatter-done event) - batches. Errors surface via _proc_error like the scatter thread's. - """ + batches. Errors surface via _proc_error like the scatter thread's.""" torch.cuda.set_device(self.device) q = self._quant_queue assert q is not None diff --git a/vllm/distributed/weight_transfer/sharded_rdt_fake.py b/vllm/distributed/weight_transfer/sharded_rdt_fake.py index 536329922ba8..5e3613e9d556 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_fake.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_fake.py @@ -206,8 +206,7 @@ def _make_child( def _meta(self) -> torch.Tensor: """A zero-storage meta tensor of this fake's shape/dtype, so PyTorch - itself computes post-op geometry rather than us reimplementing it. - """ + itself computes post-op geometry rather than us reimplementing it.""" return torch.empty(self.shape, dtype=self.dtype, device="meta") @classmethod @@ -251,8 +250,7 @@ def _intercept( ): """Append the op and return a child, or a tuple of children for multi-return ops. Each child's geometry comes from running the op on a - meta tensor. - """ + meta tensor.""" meta = self_._meta() with torch._C.DisableTorchFunctionSubclass(): meta_result = func(meta, *args, **kwargs) diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 4b16fece8fc9..6b755e772e18 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -132,8 +132,7 @@ @dataclass class _SharedPack: """One generation of one slot-sharing group: the rendezvous state for a - single chunk, identified by the consumers' issue index. - """ + single chunk, identified by the consumers' issue index.""" group: int slot: int @@ -467,8 +466,7 @@ def wait_freed(self) -> list[int]: def end_sync(self) -> list[int]: """Block until every published group has been freed by its consumers; - return the remaining freed keys so the engine drops its last refs. - """ + return the remaining freed keys so the engine drops its last refs.""" with self._cache_cond: self._wait_for( lambda: bool(self._inflight_groups), @@ -481,8 +479,7 @@ def end_sync(self) -> list[int]: def set_gather_error(self, message: str) -> None: """Record a trainer-side gather failure so blocked serves / publishes - stop waiting and surface it. - """ + stop waiting and surface it.""" with self._cache_cond: self._gather_error = RuntimeError(message) self._cache_cond.notify_all() @@ -511,8 +508,7 @@ def free_group(self, group_idx: int) -> None: def _new_serve_buffer(self, nbytes: int) -> torch.Tensor: """Allocate + NIXL-register one serve slot: the single allocation seam, so registration cannot be skipped on either of the two paths that make - buffers (the init-time reservation and the serve-path backstop). - """ + buffers (the init-time reservation and the serve-path backstop).""" t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) with self._reg_lock: register_nixl_memory(t) @@ -584,15 +580,13 @@ def reserve_serve_buffer( def _share_group(self, consumer_id: int) -> int: """The slot-sharing group ``consumer_id`` belongs to: its index within its own deployment, since that is what fixes the plan. Without a width - the group is the consumer itself and nothing is shared. - """ + the group is the consumer itself and nothing is shared.""" return consumer_id % self._share_width if self._share_width > 0 else consumer_id def _sharers_of(self, sg: int) -> frozenset: """The LIVE consumers of group ``sg``: the rendezvous width. Derived from ``begin_sync``'s live set, so a degraded sync narrows the barrier instead - of waiting forever on a dead deployment. Caller holds ``_cache_cond``. - """ + of waiting forever on a dead deployment. Caller holds ``_cache_cond``.""" cached = self._sharers.get(sg) if cached is None: if not self._sharing_active or self._live_ids is None: @@ -628,8 +622,7 @@ def _join_share_group(self, consumer_id: int, seq: int) -> tuple: def _await_shared_pack(self, gen: _SharedPack) -> torch.Tensor: """Block until this generation's packer published its blob, and return it. A failed pack is re-raised here, so every sharer of a bad pack fails - instead of reading a half-written slot. - """ + instead of reading a half-written slot.""" with self._cache_cond: self._wait_for( lambda: not gen.done and gen.error is None, @@ -650,8 +643,7 @@ def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor: """Group ``sg``'s ring slot ``idx``, grown if this chunk outgrew the reservation. Growing registers memory while the fabric is busy, the hazard ``reserve_serve_buffer`` exists to avoid, so it is a backstop: the - reservation is sized from the same static plan as this pack. - """ + reservation is sized from the same static plan as this pack.""" with self._serve_lock: buffer = self._serve_rings.setdefault(sg, [None] * self._nring)[idx] if buffer is not None and buffer.numel() >= need: @@ -663,8 +655,7 @@ def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor: def _fail_shared_pack(self, gen: _SharedPack, exc: BaseException) -> None: """Publish a pack failure so waiting sharers raise instead of hanging on - a generation that will never complete. - """ + a generation that will never complete.""" with self._cache_cond: gen.error = exc self._cache_cond.notify_all() @@ -866,8 +857,7 @@ def __init__( def _rpc(self, method: str, *args: Any) -> Any: """Call one of the server actor's methods and block for the result. The single seam through which the engine talks to its server, so tests - can inject a local (non-Ray) fake server. - """ + can inject a local (non-Ray) fake server.""" import ray return ray.get(getattr(self._server, method).remote(*args)) @@ -944,8 +934,7 @@ def _publish_async(self, group_idx: int, entries): """Fire publish_group WITHOUT waiting on the RPC (the gather loop overlaps the publish's server-side rebuild with the next group's gather) and return a handle that ``_await_publish`` resolves. Ray actor - handle in production; a plain (non-Ray) fake server runs inline. - """ + handle in production; a plain (non-Ray) fake server runs inline.""" method = self._server.publish_group remote = getattr(method, "remote", None) if remote is not None: @@ -956,8 +945,7 @@ def _await_publish(self, ref) -> None: """Resolve one async publish so a server-side rebuild error surfaces at window depth instead of at end_sync. Publishes carry nothing back — freed groups flow only through wait_freed/end_sync (one channel; see - publish_group). - """ + publish_group).""" import ray if isinstance(ref, ray.ObjectRef): @@ -1144,8 +1132,7 @@ def _validate_held_yields(self, gi: int, names, tensors) -> None: by stall watchdog 300s later: consumers route pulls here, the pull passes the served-names guard, and the cache wait never completes. With it, that is an immediate error naming the weight. One set lookup per - name per sync. - """ + name per sync.""" if self._held_names is None: return for name, tensor in zip(names, tensors): @@ -1346,8 +1333,7 @@ def _run_gather_loop( credit releases when every live consumer has signaled the group). So the loop self-paces to the consumers' pull rate with at most `gather_lookahead + 1` groups resident. Runs on every rank; only the - sender has an `update_future` to fail fast on. - """ + sender has an `update_future` to fail fast on.""" assert self.source is not None # guaranteed by trainer_init # The live IDS ride along with the count: the free barrier needs only # the count, but the slot-sharing rendezvous needs to know WHICH diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 3441f8febcde..9c704ab344ea 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -68,8 +68,7 @@ class SparseNCCLTrainerInitInfo(TrainerInitInfo): Same rendezvous shape as the dense NCCL backend (the sender opens its endpoint as NCCL rank 0), but with no packed wire params: sparse transfers - are never packed. `backend` is the factory dispatch key. - """ + are never packed. `backend` is the factory dispatch key.""" backend: ClassVar[str] = "sparse_nccl" @@ -368,8 +367,7 @@ def _post_send_sync(self) -> None: rebuild or free the patch tensors as soon as a send method returns rather than relying on same-stream ordering. See `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no - cross-rank barrier. - """ + cross-rank barrier.""" if torch.cuda.is_available(): torch.cuda.current_stream().synchronize() diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 21a185b3c19d..4ac960d8cb93 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -3,8 +3,7 @@ """The CLI entrypoints of vLLM. Note that all future modules must be lazily loaded within main -to avoid certain eager import breakage. -""" +to avoid certain eager import breakage.""" import importlib.metadata import sys diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 247e968d43ae..5142346df661 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -295,8 +295,7 @@ async def _with_kv_transfer_rejection_cleanup( ) -> _T: """Wrap a `create_*` coroutine so that, if it raises or returns an ErrorResponse (i.e. the request never reached the engine), the KV - connector is notified to free any pinned remote-prefill blocks. - """ + connector is notified to free any pinned remote-prefill blocks.""" kv_transfer_params = self.has_kv_connector and request.kv_transfer_params if not kv_transfer_params or not kv_transfer_params.get("do_remote_prefill"): return await awaitable diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index dc78cc396326..1e13311dce2a 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -123,8 +123,7 @@ def __init__( async def init_static_loras(self): """Loads all static LoRA modules. - Raises if any fail to load - """ + Raises if any fail to load""" if self.static_lora_modules is None: return for lora in self.static_lora_modules: @@ -149,8 +148,7 @@ def model_name(self, lora_request: LoRARequest | None = None) -> str: async def show_available_models(self) -> ModelList: """Show available models. This includes the base model and all - adapters. - """ + adapters.""" model_list = await self.registry.show_available_models() lora_cards = [ ModelCard( diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 167f006ba0e7..077f78c02b25 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -408,8 +408,7 @@ def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None: def need_builtin_tool_call(self) -> bool: """Return true if the last message is a builtin tool call - that the request has enabled. - """ + that the request has enabled.""" last_message = self.response_messages[-1] if last_message.type != "function_call": return False diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py index c1f3f1b716ac..1aae70a381dc 100644 --- a/vllm/entrypoints/openai/responses/harmony.py +++ b/vllm/entrypoints/openai/responses/harmony.py @@ -52,8 +52,7 @@ def _parse_harmony_format_message(chat_msg: dict) -> Message: """Reconstruct a Message from Harmony-format dict, - preserving channel, recipient, and content_type. - """ + preserving channel, recipient, and content_type.""" author_dict = chat_msg["author"] role = author_dict.get("role") name = author_dict.get("name") diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index f826c932a4ba..201e46ed0419 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -117,8 +117,7 @@ def serialize_messages(msgs): class ResponseRawMessageAndToken(OpenAIBaseModel): """Class to show the raw message. - If message / tokens diverge, tokens is the source of truth - """ + If message / tokens diverge, tokens is the source of truth""" message: str tokens: list[int] diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 51ce9ba0efcf..4b18cfc9e184 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -44,8 +44,7 @@ class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): class PluginWithIOProcessorPlugins(PoolingIOProcessor): """IO Processor plugins are a feature that allows pre- and post-processing - of the model input and output for pooling models. - """ + of the model input and output for pooling models.""" name = "plugin" diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py index b754fcc42e96..60a3560ad737 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Encode/decode utilities for multimodal tensors and field metadata -over JSON/HTTP, used by the disaggregated generate endpoint. -""" +over JSON/HTTP, used by the disaggregated generate endpoint.""" from __future__ import annotations diff --git a/vllm/entrypoints/serve/exception_handling/handlers/validation.py b/vllm/entrypoints/serve/exception_handling/handlers/validation.py index acf2e57f3b28..7cdcb3970e0f 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/validation.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/validation.py @@ -126,8 +126,7 @@ def _format_error(err: object) -> str: def _is_internal_loc_segment(segment: str) -> bool: """True if `segment` is a Pydantic-internal wrapper/union-branch - marker rather than a user-meaningful field name or list index. - """ + marker rather than a user-meaningful field name or list index.""" if _BRACKETED_INTERNAL_RE.search(segment): return True return segment.lower() in _INTERNAL_LOC_MARKERS diff --git a/vllm/entrypoints/serve/lora/api_router.py b/vllm/entrypoints/serve/lora/api_router.py index cd7e5bc99ee8..61436371c27d 100644 --- a/vllm/entrypoints/serve/lora/api_router.py +++ b/vllm/entrypoints/serve/lora/api_router.py @@ -63,8 +63,7 @@ def _attach_router(app: FastAPI): @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)]) async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request): """Handle POST /v1/load_lora_adapter: load a LoRA adapter into the - serving engine. - """ + serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.load_lora_adapter(request) if isinstance(response, ErrorResponse): @@ -86,8 +85,7 @@ async def unload_lora_adapter( request: UnloadLoRAAdapterRequest, raw_request: Request ): """Handle POST /v1/unload_lora_adapter: unload a LoRA adapter from - the serving engine. - """ + the serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.unload_lora_adapter(request) if isinstance(response, ErrorResponse): diff --git a/vllm/entrypoints/serve/sagemaker/api_router.py b/vllm/entrypoints/serve/sagemaker/api_router.py index 5c4e3de74001..f2fb255ceefc 100644 --- a/vllm/entrypoints/serve/sagemaker/api_router.py +++ b/vllm/entrypoints/serve/sagemaker/api_router.py @@ -30,8 +30,7 @@ def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]: """Snapshot handler levels for loggers that may be affected by third-party logging configuration side effects (e.g. model_hosting_container_standards - calling configure_root_logger() at import time). - """ + calling configure_root_logger() at import time).""" loggers = [logging.getLogger(), logging.getLogger("vllm")] seen: set[int] = set() snapshot: list[tuple[logging.Handler, int]] = [] @@ -74,8 +73,7 @@ def _attach_router( model_config: ModelConfig | None = None, ): """Register the SageMaker hosting routes (/ping, /invocations) on the - app. - """ + app.""" import model_hosting_container_standards.sagemaker as sagemaker_standards router = APIRouter() diff --git a/vllm/entrypoints/serve/utils/fingerprint.py b/vllm/entrypoints/serve/utils/fingerprint.py index 5203d1003bdf..e858e667d06c 100644 --- a/vllm/entrypoints/serve/utils/fingerprint.py +++ b/vllm/entrypoints/serve/utils/fingerprint.py @@ -33,8 +33,7 @@ def set_default_fingerprint_mode( custom_value: str | None = None, ) -> None: """Configure the fingerprint mode for subsequent ``get_system_fingerprint`` - calls. Called once at server startup. - """ + calls. Called once at server startup.""" global _DEFAULT_MODE, _CUSTOM_VALUE _DEFAULT_MODE = mode _CUSTOM_VALUE = custom_value @@ -42,8 +41,7 @@ def set_default_fingerprint_mode( def get_system_fingerprint(vllm_config: Any) -> str | None: """Return the fingerprint for ``vllm_config`` using the mode configured by - ``set_default_fingerprint_mode``. - """ + ``set_default_fingerprint_mode``.""" return build_system_fingerprint(vllm_config, _DEFAULT_MODE, _CUSTOM_VALUE) diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 2e77ceffeffb..34b7e9c24928 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -92,8 +92,7 @@ def asr_inter_chunk_separator( class SpeechToTextBaseServing(GenerateBaseServing): """Base class for speech-to-text operations like transcription and - translation. - """ + translation.""" def __init__( self, @@ -427,8 +426,7 @@ async def _create_speech_to_text( stream_generator_method: Callable[..., AsyncGenerator[str, None]], ) -> T | V | AsyncGenerator[str, None] | ErrorResponse: """Base method for speech-to-text operations like transcription and - translation. - """ + translation.""" if request.stream and request.use_beam_search: return self.create_error_response( "Streaming is not currently supported with beam search" diff --git a/vllm/env_override.py b/vllm/env_override.py index 35431cccb666..87501eb469f0 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -581,8 +581,7 @@ def _patch_get_raw_stream_if_needed(): def _apply_constrain_to_fx_strides_patch(): """Patch lowering.constrain_to_fx_strides globally. Safe to call multiple times; only the first call does anything. - Only applies for torch >= 2.11 and < 2.12. - """ + Only applies for torch >= 2.11 and < 2.12.""" global _constrain_to_fx_strides_patched if _constrain_to_fx_strides_patched: return diff --git a/vllm/envs.py b/vllm/envs.py index 8d836ffc4e71..0a3f6d3c6751 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2249,8 +2249,7 @@ def compile_factors() -> dict[str, object]: """Return env vars used for torch.compile cache keys. Start with every known vLLM env var; drop entries in `ignored_factors`; - hash everything else. This keeps the cache key aligned across workers. - """ + hash everything else. This keeps the cache key aligned across workers.""" ignored_factors: set[str] = { "MAX_JOBS", "VLLM_RPC_BASE_PATH", diff --git a/vllm/forward_context.py b/vllm/forward_context.py index dc025175a357..6a31b3faff1d 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -209,8 +209,7 @@ def is_forward_context_available() -> bool: def in_piecewise_cudagraph() -> bool: """Whether the current forward runs in piecewise cudagraph mode (graph - segments separated by eager breaks), at capture or replay time. - """ + segments separated by eager breaks), at capture or replay time.""" return ( is_forward_context_available() and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE diff --git a/vllm/logger.py b/vllm/logger.py index 4f528a130193..cdce7748f2cf 100644 --- a/vllm/logger.py +++ b/vllm/logger.py @@ -216,8 +216,7 @@ def _configure_vllm_root_logger() -> None: def init_logger(name: str) -> _VllmLogger: """The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root vllm logger has - already been configured. - """ + already been configured.""" logger = logging.getLogger(name) for method_name, method in _METHODS_TO_PATCH.items(): diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 6075629c453c..d015fd53ce59 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -733,8 +733,7 @@ def set_lora( lora_b: torch.Tensor | list[torch.Tensor], ): """Override to handle single tensor weights - that need to be split into slices. - """ + that need to be split into slices.""" self.reset_lora(index) # Handle case where checkpoint has single tensor weights diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 6286533a0fde..593f85993604 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -90,8 +90,7 @@ def __init__( def clone(self, lora_model_id: int) -> "LoRAModel": """Return a copy of the object with different ids. - Will share the underlying tensors. - """ + Will share the underlying tensors.""" return self.__class__( lora_model_id, rank=self.rank, diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 1fab524e6d08..93624523e2f8 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -28,8 +28,7 @@ class WorkerLoRAManager: """WorkerLoRAManager that manages LoRA models on the worker side. Every request, the requested LoRAs will be loaded (unless they are already - loaded), and every other LoRA will be unloaded. - """ + loaded), and every other LoRA will be unloaded.""" _manager_cls: type[LoRAModelManager] = LoRAModelManager @@ -75,8 +74,7 @@ def __init__( @contextmanager def dummy_lora_cache(self): """Use this context manager to reuse the dummy lora model - to avoid creating it repeatedly. - """ + to avoid creating it repeatedly.""" self._cached_dummy_lora = None yield self._cached_dummy_lora = False @@ -251,8 +249,7 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager): Uses an LRU Cache. Every request, the requested LoRAs will be loaded (unless they are already loaded) and least recently used LoRAs will - be unloaded if the cache is above capacity. - """ + be unloaded if the cache is above capacity.""" _manager_cls: type[LRUCacheLoRAModelManager] = LRUCacheLoRAModelManager diff --git a/vllm/model_executor/hw_agnostic/custom_op.py b/vllm/model_executor/hw_agnostic/custom_op.py index e9c1e2482030..ecb37681f267 100644 --- a/vllm/model_executor/hw_agnostic/custom_op.py +++ b/vllm/model_executor/hw_agnostic/custom_op.py @@ -135,8 +135,7 @@ def forward(self, *args, **kwargs): def forward_native(self, *args, **kwargs): """PyTorch-native implementation; OOT plugins override via - :meth:`forward_oot`. - """ + :meth:`forward_oot`.""" raise NotImplementedError def forward_oot(self, *args, **kwargs): diff --git a/vllm/model_executor/hw_agnostic/layers/layernorm.py b/vllm/model_executor/hw_agnostic/layers/layernorm.py index 1461e20eb42f..d2c854afd309 100644 --- a/vllm/model_executor/hw_agnostic/layers/layernorm.py +++ b/vllm/model_executor/hw_agnostic/layers/layernorm.py @@ -10,8 +10,7 @@ @CustomOp.register("rms_norm") class RMSNorm(CustomOp): """``x -> w * x / sqrt(E[x^2] + eps)``. With ``residual``, fuses - ``residual += x`` then RMSNorm and returns ``(normalized, residual)``. - """ + ``residual += x`` then RMSNorm and returns ``(normalized, residual)``.""" def __init__( self, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 37c73868f7c9..26ad557c67b7 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -881,8 +881,7 @@ def choose_mp_linear_kernel( def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8LinearKernel: """Select and instantiate the best MXFP8 linear kernel for the - current platform. - """ + current platform.""" config = Mxfp8LinearLayerConfig(bmm_batch_size=bmm_batch_size) platform = current_platform._enum @@ -934,8 +933,7 @@ def init_mxfp4_linear_kernel( activation_quant_key: QuantKey | None = None, ) -> MxFp4LinearKernel: """Select and instantiate the best MXFP4 linear kernel for the - current platform. - """ + current platform.""" config = MxFp4LinearLayerConfig( activation_quant_key=activation_quant_key, ) @@ -982,8 +980,7 @@ def init_mxfp6_linear_kernel( activation_quant_key: QuantKey | None = None, ) -> MxFp6LinearKernel: """Select and instantiate the best MXFP6 linear kernel for the - current platform. - """ + current platform.""" config = MxFp6LinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key, @@ -1071,8 +1068,7 @@ def init_wfp8_a16_linear_kernel( def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: """Select and instantiate the best NVFP4 linear kernel for the - current platform. - """ + current platform.""" config = NvFp4LinearLayerConfig() a16_kernels = ( FlashInferCuteDslNvFp4W4A16LinearKernel, diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index c5ec08ec8a18..a555ca17bf89 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -177,8 +177,7 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): def input_quant_key(self) -> QuantKey | None: """This kernel supports dynamic quantization of the input. By - convention, pre-quantized blockscales must use the swizzled layout. - """ + convention, pre-quantized blockscales must use the swizzled layout.""" return kNvfp4Dynamic @classmethod diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index 7a38c6964795..3a053412f3cf 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -41,8 +41,7 @@ def from_layer(cls, layer: torch.nn.Module) -> Self: @property def block_scale_attr(self) -> str: """Fp8LinearMethod registers the block scale as ``weight_scale_inv``, - compressed-tensors as ``weight_scale``. - """ + compressed-tensors as ``weight_scale``.""" return ( self.WEIGHT_SCALE if self.weight_scale_inv is None diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index 7df7331f4177..9f69ab0c7377 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -174,8 +174,7 @@ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | Non def input_quant_key(self) -> QuantKey | None: """Only static per-tensor activation quantization is supported for external - quantization. - """ + quantization.""" if self.config.activation_quant_key == kFp8StaticTensorSym: return kFp8StaticTensorSym return None @@ -185,8 +184,7 @@ def _pad_to_alignment( x: torch.Tensor, dim: int, alignment: int, value: float = 0.0 ) -> torch.Tensor: """Pad tensor ``x`` along ``dim`` to the next multiple of - ``alignment``. - """ + ``alignment``.""" remainder = x.shape[dim] % alignment if remainder == 0: return x diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index f730026d4a90..7166d3f8feea 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -804,8 +804,7 @@ def forward_cuda(self, input: torch.Tensor) -> torch.Tensor: def _get_gelu_pytorch_tanh() -> nn.Module: """Get PyTorch GELU with tanh approximation, with ROCm fallback - and fast GELU for ARM. - """ + and fast GELU for ARM.""" if current_platform.is_rocm(): # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile logger.warning_once( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index caa58eefae12..4cbc380bb151 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1568,8 +1568,7 @@ class MLACommonBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack an inline fp32 scale pair after the - latent data (single-sided: ``head_size_v == 0`` for MLA). - """ + latent data (single-sided: ``head_size_v == 0`` for MLA).""" spec = super().customize_spec(spec) mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index 4f2d5c1783f3..7046bfc991b9 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -55,8 +55,7 @@ def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None: """Workspace token budget for flashinfer fused all-reduce, or None if the - current world size / device is unsupported. Mirrors ``FlashInferAllReduce``. - """ + current world size / device is unsupported. Mirrors ``FlashInferAllReduce``.""" from vllm.config.compilation import PassConfig max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size) diff --git a/vllm/model_executor/layers/fused_embed_norm.py b/vllm/model_executor/layers/fused_embed_norm.py index 876585ad32bd..ddeb5156ae18 100644 --- a/vllm/model_executor/layers/fused_embed_norm.py +++ b/vllm/model_executor/layers/fused_embed_norm.py @@ -170,8 +170,7 @@ def _fused_embed_eh_norm_kernel( ``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H]) ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm + - cat. Requires the full table on-rank (replicated embedding). - """ + cat. Requires the full table on-rank (replicated embedding).""" tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py index 18f97be332ab..46a9bd1b4c58 100644 --- a/vllm/model_executor/layers/fused_moe/activation.py +++ b/vllm/model_executor/layers/fused_moe/activation.py @@ -56,8 +56,7 @@ def is_gated(self) -> bool: @property def custom_op_name(self) -> str: """Maps to the CustomOp name of activations - in vllm/model_executor/layers/activation.py. - """ + in vllm/model_executor/layers/activation.py.""" return _CUSTOM_OP_NAMES[self] def without_mul(self) -> "MoEActivation": diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index dcd9b8b543ca..93a4088fc4f7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -36,8 +36,7 @@ def is_aiter_mxfp8_moe_available() -> bool: nightly lacking #3811 would wrongly select FlyDSL instead of falling back to the native Triton dot_scaled path. #3811 added no probe-able public symbol, so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails - closed (returns False -> triton dot_scaled), which is always safe. - """ + closed (returns False -> triton dot_scaled), which is always safe.""" if not (current_platform.is_rocm() and current_platform.supports_mx()): return False try: diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index b9a96b280411..7b3193fb5f88 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -192,8 +192,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def _pad_moe_intermediate(self, layer: torch.nn.Module) -> None: """Zero-pad the per-partition MoE intermediate dim of both weights and - the expert bias, see `_padded_intermediate_size`. - """ + the expert bias, see `_padded_intermediate_size`.""" intermediate_size = self.moe_config.intermediate_size_per_partition padded_size = self._padded_intermediate_size(self.moe_config) if padded_size == intermediate_size: @@ -953,8 +952,7 @@ def apply( class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation - modular MoE experts. - """ + modular MoE experts.""" def __init__( self, @@ -1287,8 +1285,7 @@ def apply( class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): """AMD Zen INT8 MoE with per-token activation and channelwise weight - quantization, dispatched through zentorch. - """ + quantization, dispatched through zentorch.""" def __init__( self, diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index d5217706e1b5..c8063ffb4d63 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -34,8 +34,7 @@ def _select_cfg(M, N, K, block_m): """Pick the launch config from host constants only (M=num_valid_tokens, N, K, - block_m) — graph-capture safe (no GPU-scalar branch). - """ + block_m) — graph-capture safe (no GPU-scalar branch).""" # Per-regime winners (measured, isolated cuda-event A/B on gfx950, GPU 3): # BLOCK_K=256 (fewer K-iters + bigger MX scale-load coalesced with the dot), # num_stages=2 (software-pipeline overlaps the E8M0 scale-load with the scaled diff --git a/vllm/model_executor/layers/fused_moe/oracle/base.py b/vllm/model_executor/layers/fused_moe/oracle/base.py index 87c540409ffe..06be6becf83f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/base.py +++ b/vllm/model_executor/layers/fused_moe/oracle/base.py @@ -55,28 +55,24 @@ class MoEKernelOracle(ABC, Generic[BackendT]): @abstractmethod def backend_enum_cls(self) -> type[BackendT]: """Return the concrete `Enum` class enumerating this oracle's - backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`). - """ + backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`).""" @abstractmethod def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]: """Return platform-appropriate backends in priority order for - this `moe_config`. - """ + this `moe_config`.""" @abstractmethod def backend_to_kernel_cls( self, backend: BackendT ) -> list[type[mk.FusedMoEExperts]]: """Map a backend enum value to its concrete `FusedMoEExperts` - subclasses, in selection priority order. - """ + subclasses, in selection priority order.""" @abstractmethod def map_backend(self, runner_backend: MoEBackend) -> BackendT: """Map a user-facing `MoEBackend` (from the runner config) to - this oracle's enum. - """ + this oracle's enum.""" @abstractmethod def select_backend( @@ -108,8 +104,7 @@ def make_kernel( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts - combinator) for the chosen backend. - """ + combinator) for the chosen backend.""" def convert_to_kernel_format( self, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 72e83cda4245..7c57e2e5c2bd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -451,8 +451,7 @@ def _humming_fp8_weight_schema( layer: RoutedExperts, weight: torch.Tensor, weight_scale: torch.Tensor ) -> dict[str, Any]: """Build the humming weight schema from the canonical on-device fp8/mxfp8 - tensors (scale dtype/shape, block size), not the producing quant method. - """ + tensors (scale dtype/shape, block size), not the producing quant method.""" # mxfp8: e8m0 group-32 scales (stored as uint8 bytes or e8m0). humming has # no compressed-tensors mxfp8 loader; its modelopt schema fits both sources. if weight_scale.dtype in (torch.uint8, torch.float8_e8m0fnu): diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 3fdd3ae0a8c4..3b8a3c4f878d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -224,8 +224,7 @@ def _humming_int8_weight_schema( weight: torch.Tensor, weight_scale: torch.Tensor ) -> dict[str, Any]: """Build the humming compressed-tensors int8 schema from the canonical - on-device tensors; humming does the signed-int8 -> native conversion. - """ + on-device tensors; humming does the signed-int8 -> native conversion.""" config: dict[str, Any] = { "quant_method": "compressed-tensors", "format": "int-quantized", diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index f4495137415e..c70bea9aa33a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -512,8 +512,7 @@ def _process_weights_flashinfer( def _pad_w13_shard_cols(x: torch.Tensor, unit: int, padded_unit: int) -> torch.Tensor: """Zero-pad each of the two gate/up shards of a ``(E, rows, 2 * unit)`` - tensor along its last dim, from ``unit`` to ``padded_unit`` columns. - """ + tensor along its last dim, from ``unit`` to ``padded_unit`` columns.""" if padded_unit == unit: return x e, rows, _ = x.shape @@ -1081,8 +1080,7 @@ def _humming_wna16_weight_schema( quant_config: QuantizationConfig | QuantizationArgs | None, ) -> dict[str, Any]: """Humming weight schema for a WNA16 checkpoint, derived from the quant - config rather than the running kernel. - """ + config rather than the running kernel.""" from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index c0ce18194909..912b713918ef 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -331,8 +331,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: """Available backends in priority order, BF16-act variant before - activation-quantized variant within each vendor family. - """ + activation-quantized variant within each vendor family.""" _AVAILABLE_BACKENDS = [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, @@ -411,8 +410,7 @@ def _resolve_activation_key( model_activation_key: QuantKey | None, ) -> QuantKey | None: """Combine the model-supplied activation key with the user override. - Raises on conflict (both set and disagreeing). - """ + Raises on conflict (both set and disagreeing).""" user_override = _user_moe_activation_override() if user_override is None: return model_activation_key @@ -462,8 +460,7 @@ def _filter_by_activation( ) -> list[Mxfp4MoeBackend]: """Pick variants matching ``requested_activation_key``; without one, prefer BF16 if the list has any, else keep the list as-is so explicit - non-BF16 picks (e.g. the ``_afp8`` aliases) still land. - """ + non-BF16 picks (e.g. the ``_afp8`` aliases) still land.""" if requested_activation_key is not None: return [ b diff --git a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py index c2dc86dd7cff..0107f8815678 100644 --- a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py +++ b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py @@ -71,8 +71,7 @@ def _compute_routing( input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute routing with full bias, compute zero expert output, - mask zero expert IDs. - """ + mask zero expert IDs.""" topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, gating_output=router_logits, diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index c6e027afb5b9..2ddf71070307 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -163,8 +163,7 @@ def resolve_backend( ) -> str: """Resolve the decode top-k implementation from the configured backend ("auto" = the pre-existing chain, or a validated explicit - value). - """ + value).""" if self._backend == "auto": return self._resolve_auto(logits, topk_tokens, num_rows) @@ -211,8 +210,7 @@ def _resolve_auto( self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: """The pre-existing priority chain: cooperative -> persistent -> - per_row. deep_select/flashinfer/torch are opt-in only. - """ + per_row. deep_select/flashinfer/torch are opt-in only.""" if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" if self._is_cuda and topk_tokens in (512, 1024, 2048): @@ -274,8 +272,7 @@ def forward( max_seq_len: int, ) -> None: """Run the resolved decode top-k implementation, writing into - topk_indices (int32, -1 fill for rows shorter than topk_tokens). - """ + topk_indices (int32, -1 fill for rows shorter than topk_tokens).""" backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) if backend == "deep_select": row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index c8d52a6a106c..9d2e31cfe177 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -158,8 +158,7 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: """Apply the weights in layer to the input tensor. - Expects create_weights to have been called before on the layer. - """ + Expects create_weights to have been called before on the layer.""" raise NotImplementedError diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 0cb6ba516068..807a5b9b1a2a 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -861,8 +861,7 @@ def forward_hip( hidden_states: torch.Tensor, ) -> torch.Tensor: """ROCm forward using AITER Triton fused projection+attention when - available, otherwise falling back to the generic CUDA path. - """ + available, otherwise falling back to the generic CUDA path.""" if GDN_AITER_TRITON_AVAILABLE: num_tokens = hidden_states.size(0) projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 630c27ad2d61..db7e119560c5 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -261,8 +261,7 @@ def short_conv_state_shape( @classmethod def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int): """Compute the increase in group numbers to account for - replication in order to accompany the head shards. - """ + replication in order to accompany the head shards.""" # in the case ngoups % tp_size == 0, this will be zero if ngroups % tp_size == 0: return 0 diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index f69368e273e9..0f2041d2263c 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -59,8 +59,7 @@ def get_ssm_device_name() -> str: def _canonical_cache_dtype(cache_dtype: str) -> str: """Canonical key for config lookup. bf16 and fp16 share the same tuned - configs because the kernel only sees bit width when accessing state. - """ + configs because the kernel only sees bit width when accessing state.""" return "float16" if cache_dtype == "bfloat16" else cache_dtype @@ -161,8 +160,7 @@ def _try_get_optimal_ssm_config_cached( @contextmanager def override_ssm_config(config: tuple[int, int]): """Pin ``try_get_optimal_ssm_config`` to ``config`` for the duration of - the context. Used by the tuning benchmark to time specific configs. - """ + the context. Used by the tuning benchmark to time specific configs.""" global _ssm_config_override prev = _ssm_config_override _ssm_config_override = config diff --git a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py index 7f8900fa894e..0b2ef3a6aabd 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py +++ b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py @@ -210,8 +210,7 @@ def __init__(self, rank: int, world_size: int, comm_size: int, process_group=Non @property def workspace(self) -> torch.Tensor: """Device tensor (int64) that can be passed to the kernel - as ``void** workspace``. - """ + as ``void** workspace``.""" return self._workspace # ------------------------------------------------------------------ diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 1492997caa9b..a2c6eb2a458f 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -46,24 +46,21 @@ def create_weights( ): """Create weights for a layer. - The weights will be set as attributes of the layer. - """ + The weights will be set as attributes of the layer.""" raise NotImplementedError @abstractmethod def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Apply the weights in layer to the input tensor. - Expects create_weights to have been called before on the layer. - """ + Expects create_weights to have been called before on the layer.""" raise NotImplementedError # Not required functions def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Gather embeddings in the layer based on indices in the input tensor. - Expects create_weights to have been called before on the layer. - """ + Expects create_weights to have been called before on the layer.""" raise NotImplementedError # Not required functions @@ -76,8 +73,7 @@ def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module): Quantization methods that need special weight handling (e.g. repacked weights) override this. - Expects create_weights to have been called before on the layer. - """ + Expects create_weights to have been called before on the layer.""" layer.weight = embed_tokens.weight return layer diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 8f134da1dfea..3603fd8bdea9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -665,8 +665,7 @@ def _is_wNa8o8_int( ) -> bool: """Weight N-bit INT (pack-quantized for sub-byte, int-quantized for 8-bit) with static per-tensor INT8 input/output activation quant, applied as a float - fake-quant around a weight-only matmul. - """ + fake-quant around a weight-only matmul.""" is_int_pack_format = format in ( CompressionFormat.pack_quantized.value, CompressionFormat.int_quantized.value, @@ -702,8 +701,7 @@ def _is_wNaM_int( format: str | None, ) -> bool: """Weight N-bit INT with symmetric dynamic INT activation quant - via Humming kernel. - """ + via Humming kernel.""" if input_quant is None: return False is_pack_format = format == CompressionFormat.pack_quantized.value diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py index d8ee32b0fa41..23d25261301c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py @@ -37,8 +37,7 @@ def _dequant_gather_kernel( BLOCK: tl.constexpr, ): """Gather embedding rows by token id, unpack int32-packed INT weights, and - dequantize to ``out`` dtype in one pass (no int8 intermediate). - """ + dequantize to ``out`` dtype in one pass (no int8 intermediate).""" row = tl.program_id(0) col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) col_mask = col < hidden diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py index 3aa1aac45d7d..187b6438abff 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py @@ -211,8 +211,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None: """Normalize an int-quantized (plain int8) weight to the canonical - ``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect. - """ + ``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect.""" weight = layer.weight out_features, in_features = weight.shape packed = pack_to_int32(weight.data.contiguous(), self.num_bits) diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index e854c582eba3..ce3e99bb8da3 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -86,8 +86,7 @@ def _resolve_raw( def is_explicitly_configured(name: str) -> bool: """Return True if *name* has an explicit entry in extra_config, - either via exact key match or via a regex pattern key. - """ + either via exact key match or via a regex pattern key.""" if not self._config.extra_config: return False if name in self._config.extra_config: diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py index 98c8c982f713..8672972aa18a 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_scheme.py @@ -48,8 +48,7 @@ def get_moe_method( layer_config: "INCLayerConfig", ) -> "FusedMoEMethodBase | None": """Optional. Override if this scheme supports MoE. - Default raises NotImplementedError. - """ + Default raises NotImplementedError.""" raise NotImplementedError( f"{type(self).__name__} does not support MoE layers. " f"Layer config: {layer_config}" diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index f4ea599456f2..c649061f339c 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1988,8 +1988,7 @@ def process(self, layer, role) -> None: class KNvfp4Dynamic(QuantKeyScheme): """NVFP4 activation scheme (W4A4). Has a static global input scale on disk; - the per-group scale is computed at runtime inside the kernel. - """ + the per-group scale is computed at runtime inside the kernel.""" key = kNvfp4Dynamic @@ -2026,8 +2025,7 @@ def process(self, layer, role) -> None: class KFp8StaticTensor(QuantKeyScheme): """Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and - the activation slot (W8A8). One key in both QuantSpec slots. - """ + the activation slot (W8A8). One key in both QuantSpec slots.""" key = kFp8StaticTensorSym @@ -2097,8 +2095,7 @@ def process(self, layer, role) -> None: class KFp8StaticChannel(QuantKeyScheme): """Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only — - there is no static per-channel *activation* today. - """ + there is no static per-channel *activation* today.""" key = kFp8StaticTokenSym @@ -2139,8 +2136,7 @@ def process(self, layer, role) -> None: class KFp8Block128(QuantKeyScheme): """128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D. - No transpose (block kernel keeps [out,in]). - """ + No transpose (block kernel keeps [out,in]).""" key = kFp8Static128BlockSym @@ -2196,8 +2192,7 @@ def process(self, layer, role) -> None: class KMxfp8Static(QuantKeyScheme): """MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale. - Weight-role only. process is validate-only plus an idempotency guard. - """ + Weight-role only. process is validate-only plus an idempotency guard.""" key = kMxfp8Static @@ -2272,8 +2267,7 @@ class KDynamicNoParam(QuantKeyScheme): """Dynamic activation with no stored scale (W8A8): quantized at runtime in the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a non-None activation key. Activation-role only. Serves the fp8 per-token, fp8 - per-block, and mxfp8 dynamic activation keys. - """ + per-block, and mxfp8 dynamic activation keys.""" def create_weights(self, layer, role, ctx, shapes, wl) -> None: if role is not ACT: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index a4f689f2c11e..33f6afac80a9 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -79,8 +79,7 @@ def get_config_filenames(cls) -> list[str]: def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: """MoE method for RoutedExperts. Subclasses override to pick a - checkpoint-specific kernel family. - """ + checkpoint-specific kernel family.""" return Mxfp4MoEMethod(moe) def get_quant_method( diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 125ec4f2e32f..41bca8d29d9b 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -139,8 +139,7 @@ def _find_matching_targets( class OnlineQuantizationConfig(QuantizationConfig): """Model-level config for online quantization (quantize fp16/bf16 weights - during model loading, without requiring a pre-quantized checkpoint). - """ + during model loading, without requiring a pre-quantized checkpoint).""" def __init__( self, diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 0e12799304fb..2a8ef0b69f0e 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -118,8 +118,7 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: class OnlineLinearBase(LinearMethodBase): """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint - weights onto meta device and materializes them just-in-time. - """ + weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True @@ -163,8 +162,7 @@ def create_weights( class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. - Loads fp16/bf16 weights and quantizes them per-tensor during loading. - """ + Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" def __init__(self): super().__init__() @@ -265,8 +263,7 @@ def apply( class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. - Loads fp16/bf16 weights and quantizes them per-block during loading. - """ + Loads fp16/bf16 weights and quantizes them per-block during loading.""" def __init__(self): super().__init__() @@ -440,8 +437,7 @@ def apply( class _Fp8OnlineMoEBase(OnlineMoEMethodBase): """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint - weights onto meta device and materializes them just-in-time. - """ + weights onto meta device and materializes them just-in-time.""" # Declared here for mypy; actual values are set in __init__. fp8_backend: "Fp8MoeBackend" @@ -562,8 +558,7 @@ def get_fused_moe_quant_config( class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): """Online tensorwise FP8 MoE quantization. - Loads fp16/bf16 weights and quantizes them per-tensor during loading. - """ + Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" def __init__( self, @@ -620,8 +615,7 @@ def process_weights_after_loading(self, layer: Module) -> None: class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): """Online blockwise FP8 MoE quantization. - Loads fp16/bf16 weights and quantizes them per-block during loading. - """ + Loads fp16/bf16 weights and quantizes them per-block during loading.""" def __init__( self, diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index edf92fb264dd..79859a69d440 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -69,8 +69,7 @@ def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: """This function quantizes input values to float8 values " - "with tensor-wise quantization. - """ + "with tensor-wise quantization.""" dtype = current_platform.fp8_dtype() if dtype is None else dtype finfo = torch.finfo(dtype) min_val, max_val = x.aminmax() @@ -1195,8 +1194,7 @@ def prepare_fp8_moe_layer_for_deepgemm( def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor: """Pad the weight tensor. This is an optimization on ROCm platform, which - can benefit from tensors located far enough from one another in memory - """ + can benefit from tensors located far enough from one another in memory""" if ( envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm() diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index f03b1dc7b44f..f48c33d4cd31 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -297,8 +297,7 @@ def _humming_input_schema_to_quant_key( schema: "HummingInputSchema", ) -> QuantKey | None: """Convert a HummingInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs. - """ + the schema represents unquantized (bf16/fp16) inputs.""" if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: return None @@ -427,8 +426,7 @@ def convert_linear_layer_to_humming_standard( ): """Rename/reshape a linear layer's quantized params (the canonical MPLinear layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names - and layout humming's weight schema expects (``weight`` / ``weight_scale``). - """ + and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" for name, checkpoint_name in name_map.items(): tensor = getattr(layer, checkpoint_name) delattr(layer, checkpoint_name) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 8e0e4bba1c2d..82201dff3b29 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -245,8 +245,7 @@ def marlin_padded_nk(size_n: int, size_k: int, group_size: int = -1) -> tuple[in def marlin_repacked_nk(qweight: torch.Tensor, num_bits: int) -> tuple[int, int]: """Recover the (size_n, size_k) a Marlin weight was repacked with - (including any tile padding) from its packed shape. - """ + (including any tile padding) from its packed shape.""" pack_factor = 32 // num_bits size_k = qweight.size(0) * GPTQ_MARLIN_TILE size_n = qweight.size(1) * pack_factor // GPTQ_MARLIN_TILE @@ -257,8 +256,7 @@ def marlin_pad_qweight( qweight: torch.Tensor, size_n: int, size_k: int, padded_n: int, padded_k: int ) -> torch.Tensor: """Zero-pad a GPTQ-layout packed weight (size_k / pack, size_n) for - gptq_marlin_repack. - """ + gptq_marlin_repack.""" if (padded_n, padded_k) == (size_n, size_k): return qweight pack_factor = size_k // qweight.size(0) @@ -276,8 +274,7 @@ def marlin_pad_scales( group_size: int, ) -> torch.Tensor: """Zero-pad weight scales (num_groups, size_n); call before - marlin_permute_scales and pass the padded extents to it. - """ + marlin_permute_scales and pass the padded extents to it.""" if (padded_n, padded_k) == (size_n, size_k): return scales pad_rows = padded_k // group_size - scales.size(0) if group_size > 0 else 0 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 1a550576b6fa..a12d8b87e98c 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -41,8 +41,7 @@ def _nvfp4_compute_scale_factor( ) -> float: """Compute the power-of-2 scale_factor needed so that all non-zero values in marlin_scales * 2^7 are >= 2 after rescaling. - Returns a Python float (power of 2, >= 1.0). - """ + Returns a Python float (power of 2, >= 1.0).""" # Since half has a smaller dynamic range compared to bfloat16, # no rescaling is applied here if active dtype is half. if a_dtype is not None and a_dtype == torch.half: @@ -376,8 +375,7 @@ def prepare_nvfp4_moe_layer_for_marlin( def pad_w13(x: torch.Tensor) -> torch.Tensor: """Zero-pad each gate/up shard of a (E, num_shards * N, cols) - tensor to padded_N rows. - """ + tensor to padded_N rows.""" if padded_N == N: return x x = x.view(E, num_shards, N, x.size(-1)) @@ -386,8 +384,7 @@ def pad_w13(x: torch.Tensor) -> torch.Tensor: def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: """Zero-pad the packed N (last) dim of a (E, K, N / packing) - tensor. - """ + tensor.""" if padded_N == N: return x return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 251207a9733a..4a5f5ee59a52 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -217,8 +217,7 @@ def prepare_fp8_layer_for_marlin( def _moe_pad_shard_rows(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor: """Zero-pad each gate/up shard of a ``(E, 2 * n, ...)`` tensor to padded_n - rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing. - """ + rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing.""" if padded_n == n: return x e = x.size(0) diff --git a/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py index b92dc8e7fd6b..031a12fceba6 100644 --- a/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py @@ -9,8 +9,7 @@ class NTKScalingRotaryEmbedding(RotaryEmbedding): """RotaryEmbedding extended with fixed and mixed NTK scaling. - https://kexue.fm/archives/9706 - """ + https://kexue.fm/archives/9706""" def __init__( self, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index e7a938239614..7dde4db00fa8 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -285,8 +285,7 @@ def _gather_workspace_shapes( """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + - (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales. - """ + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" if use_fp4_cache: return ( ((total_seq_lens, head_dim // 2), torch.uint8), @@ -304,8 +303,7 @@ def kv_cache_as_quant_view( use_fp4_cache: bool, ) -> torch.Tensor: """4D ``[num_blocks, block_size, 1, head_width]`` view expected by - DeepGEMM, from the 3D indexer kv-cache allocation. - """ + DeepGEMM, from the 3D indexer kv-cache allocation.""" if use_fp4_cache: assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 num_blocks, block_size, _ = kv_cache.shape diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py index a019f9ec7373..78db1b3df5a4 100644 --- a/vllm/model_executor/layers/sparse_mqa_indexer.py +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -41,8 +41,7 @@ def _prefill_k_workspaces( total_seq_lens: int, head_dim: int ) -> tuple[torch.Tensor, torch.Tensor]: """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with - the dense indexer layers of the same model. - """ + the dense indexer layers of the same model.""" values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True ) @@ -120,8 +119,7 @@ def __init__( def _reserve_workspaces(self, device: torch.device) -> None: """Profiling run: claim the K-gather workspace and the peak sparse - logits allocation so the memory estimate covers them. - """ + logits allocation so the memory estimate covers them.""" _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 688b4bf197ef..72cafceb2589 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -376,8 +376,7 @@ def _get_indices( ) -> VocabParallelEmbeddingShardIndices: """Get start and end indices for vocab parallel embedding, following the layout outlined in the class docstring, based on the given tp_rank and - tp_size. - """ + tp_size.""" num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded padded_org_vocab_start_index, padded_org_vocab_end_index = ( vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, tp_size) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index bbda4a33fdad..6f356464585a 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -36,8 +36,7 @@ def download_model(self, model_config: ModelConfig) -> None: @abstractmethod def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: """Load weights into a model. This standalone API allows - inplace weights loading for an already-initialized model - """ + inplace weights loading for an already-initialized model""" raise NotImplementedError def create_model( diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index e48c5a380ea8..039a577f80d1 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -135,8 +135,7 @@ def _prepare_weights( ) -> tuple[str, list[str], bool]: """Prepare weights for the model. - If the model is not local, it will be downloaded. - """ + If the model is not local, it will be downloaded.""" model_name_or_path = ( maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path diff --git a/vllm/model_executor/model_loader/ep_weight_filter.py b/vllm/model_executor/model_loader/ep_weight_filter.py index 4b2f92a18dca..0e23dfc11372 100644 --- a/vllm/model_executor/model_loader/ep_weight_filter.py +++ b/vllm/model_executor/model_loader/ep_weight_filter.py @@ -23,8 +23,7 @@ def parse_expert_id(weight_name: str) -> int | None: Returns ``None`` for dense weights (attention, layernorm, embedding), shared experts, and 3D fused-expert tensors where all experts are stored - in a single tensor without a numeric expert id in the name. - """ + in a single tensor without a numeric expert id in the name.""" m = _EXPERT_ID_RE.search(weight_name) return int(m.group(1)) if m else None @@ -71,8 +70,7 @@ def should_skip_weight( local_expert_ids: set[int] | None, ) -> bool: """Return ``True`` if *weight_name* is an expert weight that does not - belong to the local rank and should be skipped during loading. - """ + belong to the local rank and should be skipped during loading.""" if local_expert_ids is None: return False eid = parse_expert_id(weight_name) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index c52fad985beb..062bfa60906b 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -308,8 +308,7 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) - Assumes dtype/shapes of attention tensors do not change during processing, since we use .data.copy_() to preserve kernel tensor - references. - """ + references.""" quant_method = getattr(layer, "quant_method", None) if quant_method is not None: # Re-create scale Parameters with sentinel values so unloaded scales @@ -388,8 +387,7 @@ def _get_weight_loader(tensor: torch.Tensor): def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): """Copy processed values into original kernel tensor storage and restore - kernel tensor references on the layer. Preserves cudagraph references. - """ + kernel tensor references on the layer. Preserves cudagraph references.""" assert info.kernel_tensors is not None parameters, buffers = info.kernel_tensors non_persistent = info.kernel_non_persistent_buffers diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index b38729d28341..e07330ae57c5 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -81,8 +81,7 @@ def _prepare_weights( ) -> list[str]: """Prepare weights for the model. - If the model is not local, it will be downloaded. - """ + If the model is not local, it will be downloaded.""" is_object_storage_path = is_runai_obj_uri(model_name_or_path) is_local = os.path.isdir(model_name_or_path) safetensors_pattern = "*.safetensors" diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index 582ee0ea3e05..3f865c3f3980 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -497,8 +497,7 @@ def _check_tensors_on_meta_device(model: nn.Module) -> None: def _resize_lora_embeddings(model: nn.Module): """Modify LoRA embedding layers to use bigger tensors - to allow for adapter added tokens. - """ + to allow for adapter added tokens.""" for child in model.modules(): if ( isinstance(child, VocabParallelEmbedding) diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py index e231eea20bae..338f9eac072a 100644 --- a/vllm/model_executor/model_loader/tensorizer_loader.py +++ b/vllm/model_executor/model_loader/tensorizer_loader.py @@ -105,8 +105,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: Expects a vLLM-tensorized model. See the examples/features/tensorize_vllm_model.py example script - for serializing vLLM models. - """ + for serializing vLLM models.""" if is_vllm_tensorized(self.tensorizer_config): tensorizer_config = self._patch_tensorizer_config(model_config) deserialize_tensorizer_model(model, tensorizer_config) diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 52d770f5d1c9..6c3fb742cdf7 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """IPC model loader: maps post-quantized weights from a local weight cache -daemon via CUDA IPC instead of loading from disk. -""" +daemon via CUDA IPC instead of loading from disk.""" import dataclasses import socket diff --git a/vllm/model_executor/model_loader/weight_tying.py b/vllm/model_executor/model_loader/weight_tying.py index bd0bfa4b8038..d0c309d2c4ad 100644 --- a/vllm/model_executor/model_loader/weight_tying.py +++ b/vllm/model_executor/model_loader/weight_tying.py @@ -78,8 +78,7 @@ def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> Checkpoints produced by quantization or fine-tuning tooling often keep a redundant copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would - otherwise cost. - """ + otherwise cost.""" if not model_config.word_embeddings_untied_by_checkpoint: return if (untied := _get_untied_lm_head(model)) is None: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 2a9937bcd441..c781a8bf87c0 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -144,8 +144,7 @@ def atomic_writer( def _natural_sort_key(filepath: str) -> list: """Natural sort key for filenames with numeric components, such as - model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors'] - """ + model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']""" return [ int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath)) @@ -162,8 +161,7 @@ def maybe_download_from_modelscope( """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. Returns the path to the downloaded model, or None if the model is not - downloaded from ModelScope. - """ + downloaded from ModelScope.""" if envs.VLLM_USE_MODELSCOPE: # download model from ModelScope hub, # lazy import so that modelscope is not required for normal use. @@ -1148,8 +1146,7 @@ def instanttensor_weights_iterator( use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files - using instanttensor library. - """ + using instanttensor library.""" try: import instanttensor except ImportError as e: diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index a75bdaee5abe..32ad19a8e73b 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -47,8 +47,7 @@ class ArceeMLP(nn.Module): """Feed-forward layer for Arcee using ReLU^2 activation - (no gating as in LLaMA). - """ + (no gating as in LLaMA).""" def __init__( self, @@ -96,8 +95,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ArceeDecoderLayer(nn.Module): """Transformer decoder block for Arcee, with self-attention and - ReLU^2 MLP. - """ + ReLU^2 MLP.""" def __init__( self, @@ -177,8 +175,7 @@ def forward( @support_torch_compile class ArceeModel(nn.Module, EagleModelMixin): """The transformer model backbone for Arcee (embedding layer + stacked - decoder blocks + final norm). - """ + decoder blocks + final norm).""" def __init__( self, @@ -280,8 +277,7 @@ class ArceeForCausalLM( nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 ): """Arcee Model for causal language modeling, integrated with vLLM - runtime. - """ + runtime.""" hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ @@ -362,8 +358,7 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights into the model (delegates to inner model and handles - tied embeddings). - """ + tied embeddings).""" loader = AutoWeightsLoader(self) # AutoWeightLoader handles weight name remapping, including fusing # separate q_proj, k_proj, v_proj into qkv_proj diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index 09f0973a5cca..a4e821ec6b12 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Minimal implementation of BlipVisionModel intended to be only used -within a vision language model. -""" +within a vision language model.""" from collections.abc import Iterable diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index ebfd93e857a6..d657b78a5da1 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -96,8 +96,7 @@ def forward(self, hidden_states, residuals=None): def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]: """Returns (norm_class, eps). Uses RMSNorm when config.rms_norm_eps is set, - otherwise falls back to LayerNorm with config.layer_norm_eps. - """ + otherwise falls back to LayerNorm with config.layer_norm_eps.""" rms_eps = getattr(config, "rms_norm_eps", None) if rms_eps is not None: return RMSNorm, rms_eps diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 5020cadc99ca..844e8c5e0e7f 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -667,8 +667,7 @@ def calc_length( repeat_num: int = 1, ) -> torch.Tensor: """Calculates the output length of a Tensor passed - through a convolution or max pooling layer - """ + through a convolution or max pooling layer""" add_pad: float = all_paddings - kernel_size one: float = 1.0 for i in range(repeat_num): diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py index 824813eadb15..fa972ca6157c 100644 --- a/vllm/model_executor/models/deepencoder.py +++ b/vllm/model_executor/models/deepencoder.py @@ -363,8 +363,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Block(nn.Module): """Transformer blocks with support of window attention and residual propagation - blocks - """ + blocks""" def __init__( self, diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 3ead25f173dc..3d1921dc2b7f 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -139,8 +139,7 @@ def __call__( class NGramPerReqLogitsProcessor(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): diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7df3f744ce38..bd1d14b6f6dc 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -511,8 +511,7 @@ def _compiled_sample_step( accept/renoise → convergence, all as vectorized PyTorch ops. Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the - caller can compute logprobs outside the compiled region. - """ + caller can compute logprobs outside the compiled region.""" num_decode = decode_slots.shape[0] device = decode_slots.device diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index 87d72c18c10a..3c9df39e674c 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -814,8 +814,7 @@ def forward( def _kv_sharing_weights_mapper(config: Gemma3nTextConfig) -> WeightsMapper: """KV-shared layers only have q_proj, so qkv_proj packing applies to the other layers. Original checkpoints still ship K/V tensors for the shared - layers (fine-tuned ones omit them); those are dropped. - """ + layers (fine-tuned ones omit them); those are dropped.""" first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers return WeightsMapper( orig_to_new_substr={ diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index dfde046fece6..d0f536afe89d 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -499,8 +499,7 @@ def get_repl_toks(tok: int) -> list[int]: class Gemma3nMultimodalEmbedder(nn.Module): """Embeds token ids or soft tokens for multimodal content into language - model space. - """ + model space.""" def __init__( self, diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index c0f77e770da6..8f3db8a3d110 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -25,8 +25,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.1V & GLM-4.6V-Flash, AutoGLM-Phone-9B model -compatible with HuggingFace weights. -""" +compatible with HuggingFace weights.""" import math from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 18256e9c13bf..3c72a148d5ef 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -22,8 +22,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.5, GLM-4.6, GLM-4.7 model -compatible with HuggingFace weights. -""" +compatible with HuggingFace weights.""" from collections.abc import Iterable from itertools import islice diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index 12592af4a734..6d951e04ac7e 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -22,8 +22,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP -model compatible with HuggingFace weights. -""" +model compatible with HuggingFace weights.""" import typing from collections.abc import Callable, Iterable diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 529befd21d24..1c1e98be043c 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -462,8 +462,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class GraniteSpeechConformerBlock(nn.Module): """Conformer block, consisting largely of linear layers, - attention, and convolutional layers. - """ + attention, and convolutional layers.""" def __init__(self, config: PretrainedConfig, prefix: str = ""): super().__init__() diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index e0121f95338d..1cefd63dd63c 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -892,8 +892,7 @@ def has_inner_state( @runtime_checkable class IsAttentionFree(Protocol): """The interface required for all models like Mamba that lack attention, - but do have state whose size is constant wrt the number of tokens. - """ + but do have state whose size is constant wrt the number of tokens.""" is_attention_free: ClassVar[Literal[True]] = True """ @@ -921,8 +920,7 @@ def is_attention_free( class IsHybrid(Protocol): """The interface required for all models like Jamba that have both attention and mamba blocks, indicates that - hf_config has 'layers_block_type' - """ + hf_config has 'layers_block_type'""" is_hybrid: ClassVar[Literal[True]] = True """ @@ -1338,8 +1336,7 @@ def get_generation_prompt( ) -> "PromptType": """Get the prompt for the ASR model. The model has control over the construction, as long as it - returns a valid PromptType. - """ + returns a valid PromptType.""" ... @classmethod @@ -1631,8 +1628,7 @@ def collect_remote_aux_hidden_states( @runtime_checkable class SupportsEagle(SupportsEagleBase, Protocol): """The interface required for models that support - EAGLE-1 and EAGLE-2 speculative decoding. - """ + EAGLE-1 and EAGLE-2 speculative decoding.""" supports_eagle: ClassVar[Literal[True]] = True """ @@ -1662,8 +1658,7 @@ def supports_eagle( @runtime_checkable class SupportsEagle3(SupportsEagleBase, Protocol): """The interface required for models that support - EAGLE-3 speculative decoding. - """ + EAGLE-3 speculative decoding.""" supports_eagle3: ClassVar[Literal[True]] = True """ diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index b11adfa499f8..1322920900bb 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -345,8 +345,7 @@ def apply_rotary_pos_emb_flashatt( class KeyeSiglipAttention(nn.Module): """Multi-headed attention from 'Attention Is All You - Need' paper. - """ + Need' paper.""" def __init__( self, diff --git a/vllm/model_executor/models/lfm2_siglip2.py b/vllm/model_executor/models/lfm2_siglip2.py index 57a3d4a8d798..b3c7ee28ae7a 100644 --- a/vllm/model_executor/models/lfm2_siglip2.py +++ b/vllm/model_executor/models/lfm2_siglip2.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Implementation of Siglip2VisionModel intended to be only used -within a vision language model. -""" +within a vision language model.""" from collections.abc import Iterable diff --git a/vllm/model_executor/models/medusa.py b/vllm/model_executor/models/medusa.py index bf43f85d362c..6adf3b7669ec 100644 --- a/vllm/model_executor/models/medusa.py +++ b/vllm/model_executor/models/medusa.py @@ -51,8 +51,7 @@ class Medusa(nn.Module): truncated_vocab_size < vocab_size. To use this technique, one has to find the top-k most frequent tokens in target dataset and add that as a tensor in the draft checkpoint (using key token_map). Also, the draft config - needs to have truncated_vocab_size (=k) as an attribute. - """ + needs to have truncated_vocab_size (=k) as an attribute.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: speculative_config = vllm_config.speculative_config diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index c8a2fe4bcddb..0a96c2d8479d 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -858,8 +858,7 @@ def forward( class MiniCPMV4_6DownsampleMLP(nn.Module): """Match HF (transformers v5.7+) parameter naming: pre_norm/linear_1/ - act/linear_2 (instead of pre_norm + Sequential(mlp.0/mlp.2)). - """ + act/linear_2 (instead of pre_norm + Sequential(mlp.0/mlp.2)).""" def __init__( self, diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 5c5c2728713c..797c06d18093 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -1090,8 +1090,7 @@ def _consolidate_qkv_weights( def _rename_weight_for_modelopt_checkpoint(self, name: str) -> str: """Rename weights from ModelOpt llama4 fp8 checkpoints to vLLM - format. - """ + format.""" if name.startswith("model.") or name.startswith("language_model.model."): renamed = ( name.replace("model.", "language_model.model.", 1) diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index 38cc475be24f..ed8064542be3 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -114,8 +114,7 @@ def _text_config(config): """MuseGlimmer checkpoints may nest the text config under ``text_config`` (multimodal ``MuseGlimmerConfig``) or expose it directly - (``MuseGlimmerTextConfig``). - """ + (``MuseGlimmerTextConfig``).""" return getattr(config, "text_config", config) @@ -489,8 +488,7 @@ def video_replacement(item_idx: int) -> PromptUpdateDetails: def _muse_glimmer_use_qk_norm(config) -> bool: """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``). - Treat a missing/None flag as True — only an explicit ``False`` disables it. - """ + Treat a missing/None flag as True — only an explicit ``False`` disables it.""" val = getattr(config, "use_qk_norm", None) return True if val is None else bool(val) @@ -498,8 +496,7 @@ def _muse_glimmer_use_qk_norm(config) -> bool: def _muse_glimmer_use_attn_output_gate(config) -> bool: """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS applies it; the modular HF ``text_config`` omits ``use_attn_output_gate`` - (reads as ``None``). Missing/None -> True; only explicit ``False`` disables. - """ + (reads as ``None``). Missing/None -> True; only explicit ``False`` disables.""" val = getattr(config, "use_attn_output_gate", None) return True if val is None else bool(val) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index eb5e9f06ca4d..6e500c0c2c8e 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1188,8 +1188,7 @@ def _process_video_input( self, video_input: NanoNemotronVLVideoPixelInputs ) -> tuple[torch.Tensor, ...]: """Process video input and create final embeddings with video content - and indicator tokens. - """ + and indicator tokens.""" T = self.video_temporal_patch_size if T > 1: diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py index bd7ec9c3444d..af54defd3fb0 100644 --- a/vllm/model_executor/models/phi4mm_utils.py +++ b/vllm/model_executor/models/phi4mm_utils.py @@ -1462,8 +1462,7 @@ def conv_split_by_batch(self, x: Tensor) -> tuple[Tensor, bool]: def conv_split_by_channel(self, x: Tensor) -> Tensor: """For dw convs, tries to split input by time, run conv and concat - results - """ + results""" x = self.conv[0](x) # full conv2D x = self.conv[1](x) # activation diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index a807f2d4b5aa..ce3884a10eb8 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -69,8 +69,7 @@ def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: def dflash_has_any_non_causal(config: Qwen3Config) -> bool: """Whether the draft needs a non-causal-capable backend, resolved from config - (config mirror of the model's ``get_draft_attn_causal``, usable pre-build). - """ + (config mirror of the model's ``get_draft_attn_causal``, usable pre-build).""" return not all( _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers) ) @@ -155,8 +154,7 @@ class DFlashQwen3Attention(nn.Module): Context KVs are pre-inserted into the KV cache before the forward pass. This layer handles only query tokens via standard attention. - Adapted from Qwen3Attention. - """ + Adapted from Qwen3Attention.""" def __init__( self, @@ -250,8 +248,7 @@ def forward( """DFlash attention assumes that the KV cache is already populated with the context K/V from the target model's hidden states. This forward op computes attention for the query tokens only. - See also: precompute_and_store_context_kv - """ + See also: precompute_and_store_context_kv""" qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) @@ -743,8 +740,7 @@ def get_draft_kv_cache_layer_names(self) -> list[str]: def get_draft_attn_causal(self) -> list[bool]: """Per-layer attention causality, aligned with - get_draft_kv_cache_layer_names. - """ + get_draft_kv_cache_layer_names.""" return [layer.self_attn.causal for layer in self.model.layers] def compute_logits( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 78dda9c81503..dc2baa1c5ee4 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1321,8 +1321,7 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: """Transformers>=5.10 processors override `replace_video_token` to expand only the bare video token, keeping the prompt's outer - `<|vision_start|>`/`<|vision_end|>` markers. - """ + `<|vision_start|>`/`<|vision_end|>` markers.""" mixin_impl = getattr(ProcessorMixin, "replace_video_token", None) proc_impl = getattr(type(hf_processor), "replace_video_token", None) return proc_impl is not None and proc_impl is not mixin_impl diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 9f8ff5ed2bf6..bc895913240d 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -599,8 +599,7 @@ def inter_image_mask_metadata( self, imgs_sizes: list[tuple[int, int]], device: torch.device ) -> MaskMetadata: """Build mask metadata from image pixel sizes. Adds num_skip to each - sequence length (cls/register tokens) to match patch generator output. - """ + sequence length (cls/register tokens) to match patch generator output.""" patch_size = self.patch_generator.patch_size num_skip = self.patch_generator.num_skip @@ -613,8 +612,7 @@ def _inter_image_mask_metadata_from_seq_lens( ) -> MaskMetadata: """Build mask metadata from actual sequence lengths (already including cls/register tokens, i.e. patch_count + num_skip per item). - Use inter_image_mask_metadata() when you only have imgs_sizes. - """ + Use inter_image_mask_metadata() when you only have imgs_sizes.""" assert len(seq_lens) > 0 cu_seqlens = torch.tensor( list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device diff --git a/vllm/model_executor/models/siglip2navit.py b/vllm/model_executor/models/siglip2navit.py index aabb62f4ab59..1d0c3c8d96f9 100644 --- a/vllm/model_executor/models/siglip2navit.py +++ b/vllm/model_executor/models/siglip2navit.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Implementation of SiglipVisionModel intended to be only used -within a vision language model. -""" +within a vision language model.""" from collections.abc import Iterable diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py index 7c4c24813b7d..18512528820c 100644 --- a/vllm/model_executor/models/stablelm.py +++ b/vllm/model_executor/models/stablelm.py @@ -20,8 +20,7 @@ # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/modeling_stablelm_epoch.py # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json """Inference-only StableLM (https://github.com/Stability-AI/StableLM) -model compatible with HuggingFace weights. -""" +model compatible with HuggingFace weights.""" from collections.abc import Iterable from itertools import islice diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index eeec0a4b919e..62f6d684e8b6 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -207,8 +207,7 @@ def _patch_config(self): def _mark_model_components(self, vllm_config: "VllmConfig"): """Mark language model and tower submodules as `self.model` is created. - Nothing to do in `Base`, `MultiModalMixin` will override. - """ + Nothing to do in `Base`, `MultiModalMixin` will override.""" yield @cached_property @@ -281,8 +280,7 @@ def _decorate_cls_for_torch_compile( def _decorate_for_torch_compile(self): """Decorate the model's decoder class to indicate to vLLM that it - supports torch compile if `can_enable_torch_compile` is True. - """ + supports torch compile if `can_enable_torch_compile` is True.""" self._decorate_cls_for_torch_compile( cls=self._pre_trained_model_classes.decoder, # Applied to a PreTrainedModel so the batch dimension will exist diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index 0b83afe6b727..ad24f2542a47 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -55,20 +55,17 @@ def fuse( self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Apply the fusion to an already-validated `module`, returning the - module to install in its place (mutated in place, or freshly built). - """ + module to install in its place (mutated in place, or freshly built).""" def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes - (none unless it stacks weights). - """ + (none unless it stacks weights).""" return {} @property def packed_modules_mapping(self) -> dict[str, list[str]]: """`packed_modules_mapping` entries this fuser contributes (none unless - it stacks weights). - """ + it stacks weights).""" return {} @@ -122,8 +119,7 @@ def fuse( ) -> nn.Module: """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). - Builds the merged submodule and binds the compiled forward. - """ + Builds the merged submodule and binds the compiled forward.""" self.update_attrs(module, prefix, vllm_config) module.forward = types.MethodType(self.fused_forward, module) return module @@ -132,8 +128,7 @@ def fuse( @dataclass class StackedFuser(RewriteFuser): """A fuser that merges sibling projections into one stacked linear and - rewrites the forward to call it. - """ + rewrites the forward to call it.""" merged_name: ClassVar[str] """Attribute name of the merged module created by `update_attrs`.""" @@ -152,16 +147,14 @@ def info(self, name: str) -> str: def shards(self) -> list[tuple[str, ShardId]]: """Each projection's original name and its shard id in the merged module. - Source for both `orig_to_new_stacked` and `packed_modules_mapping`. - """ + Source for both `orig_to_new_stacked` and `packed_modules_mapping`.""" def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: """`WeightsMapper.orig_to_new_stacked` entries for one fused instance. Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname so only this exact layer is remapped, never a same-named projection - elsewhere (e.g. an unfused MoE expert's `gate_proj`). - """ + elsewhere (e.g. an unfused MoE expert's `gate_proj`).""" merged = maybe_prefix(prefix, self.merged_name) return { maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards @@ -170,8 +163,7 @@ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: @property def packed_modules_mapping(self) -> dict[str, list[str]]: """`{merged_name: [projection names]}` so quantization can unpack the - fused layer into its per-shard configs. - """ + fused layer into its per-shard configs.""" return {self.merged_name: [name for name, _ in self.shards]} def _unguarded_calls( @@ -182,8 +174,7 @@ def _unguarded_calls( `update_attrs` deletes the projections it merges, so any reference to one beyond its call site must be a guard on its existence; folding those to their constant value keeps the rewritten forward off a name that no longer exists. A reference - that is not such a guard raises, so fusion is skipped. - """ + that is not such a guard raises, so fusion is skipped.""" calls = [] for name in names: call, refs = self_call_and_refs(funcdef, name) diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index fccb1a8ea966..fa4b34750c5f 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -232,8 +232,7 @@ def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: def update_forward(self, module: nn.Module) -> None: """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split. Bypass the KV expansion method so the compressed latent reaches the `vllm_mla` - attention interface unexpanded. - """ + attention interface unexpanded.""" funcdef, fn = recover_forward(type(module)) if self.has_q_lora: # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index b7f435b4e032..af3a101a0ce1 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -80,8 +80,7 @@ def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None: Computations that must run in higher precision say so in their forward (e.g. `hidden_states.type(torch.float32)`), so the dtype is readable from the graph - even when the config does not name it. - """ + even when the config does not name it.""" for node in nodes: if node.op not in ("call_method", "call_function"): continue @@ -126,8 +125,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: """Standard MoE block forward. - Routing and any shared experts are handled inside `self.experts: MoERunner`. - """ + Routing and any shared experts are handled inside `self.experts: MoERunner`.""" orig_shape = hidden_states.shape hidden_states = hidden_states.reshape(-1, orig_shape[-1]) num_tokens = hidden_states.shape[0] @@ -154,8 +152,7 @@ class MoEBlockFuser: def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. - Returns the scoring function and the dtype the router computes in. - """ + Returns the scoring function and the dtype the router computes in.""" state = {name for name, _ in named_state(gate)} if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None @@ -283,8 +280,7 @@ def shared_experts( self, moe_block: nn.Module, prefix: str ) -> SharedExpertMLP | None: """Build the HF shared expert (and its optional gate) - as a `SharedExpertMLP` for vLLM's fused MoE. - """ + as a `SharedExpertMLP` for vLLM's fused MoE.""" if self.shared_name is None: return None shared_experts = getattr(moe_block, self.shared_name) diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index b0f449a6d212..3cdc4e1201ba 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -43,8 +43,7 @@ def _reference_weight(module: nn.Module) -> torch.Tensor | None: """A weight whose trailing dim is the module's hidden size. Linears and 2-D gate weights are `[out, hidden]`; norm weights are - `[hidden]`. Used to fabricate a placeholder input of matching size/dtype. - """ + `[hidden]`. Used to fabricate a placeholder input of matching size/dtype.""" for child in module.modules(): if isinstance(child, nn.Linear): return child.weight @@ -59,8 +58,7 @@ class _MetaProxy(fx.Proxy): Shape questions (`len`, iteration, `.shape` unpacks) are answered by executing each op on the meta values, so PyTorch's meta kernels are the - single source of shape inference — no per-op rules. - """ + single source of shape inference — no per-op rules.""" meta: object = _UNKNOWN @@ -84,8 +82,7 @@ class _MetaAttribute(_MetaProxy, fx.proxy.Attribute): """Attribute proxy (e.g. `x.shape`) carrying its meta value. `Proxy.__getattr__` constructs `Attribute` directly, bypassing - `Tracer.proxy`, so the meta value must be grafted on here too. - """ + `Tracer.proxy`, so the meta value must be grafted on here too.""" def __init__(self, root: fx.Proxy, attr: str): super().__init__(root, attr) @@ -221,8 +218,7 @@ def _leaf_attention_interfaces(): """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node. `vllm_attention_function` needs runtime context so it is untraceable. - Every interface returns `(attn_output, attn_weights)`. - """ + Every interface returns `(attn_output, attn_weights)`.""" from transformers.modeling_utils import AttentionInterface original = AttentionInterface.get_interface @@ -300,8 +296,7 @@ def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]: def forward_input_count(cls: type[nn.Module]) -> int: """The number of tensor inputs `cls.forward` declares, excluding `self` and any `*args`/`**kwargs`. Read from the signature, so it is independent of - whether the trace completes (unlike counting placeholders). - """ + whether the trace completes (unlike counting placeholders).""" params = list(forward_parameters(cls).values()) if not params: return 1 # uninspectable: assume a single input and let matching decide @@ -378,8 +373,7 @@ def _in_boolean_context(funcdef: ast.FunctionDef, ref: ast.expr) -> bool: In these positions the object's identity never escapes, so a reference that is always truthy can be replaced by `True`. `and`/`or` are excluded: they - yield an operand, so the module could escape (`x and self.`). - """ + yield an operand, so the module could escape (`x and self.`).""" for node in ast.walk(funcdef): if ( isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert)) @@ -412,8 +406,7 @@ def bypass_existence_guard( `not self.`): the reference itself to `True`. Any other surviving reference escapes the projection's value, which no longer - exists after fusion, so refuse rather than change semantics. - """ + exists after fusion, so refuse rather than change semantics.""" for node in ast.walk(funcdef): # `self. is (not) None`, either operand order. if not (isinstance(node, ast.Compare) and len(node.ops) == 1): @@ -474,8 +467,7 @@ def _base_name(node: ast.expr) -> str | None: def _rebound_names(region: list[ast.stmt]) -> set[str]: """Names rebound outright within `region` (`x = ...`, `del x`). - These are `Name` nodes in a `Store`/`Del` context. - """ + These are `Name` nodes in a `Store`/`Del` context.""" return { node.id for stmt in region @@ -489,8 +481,7 @@ def _inplace_target(node: ast.AST) -> str | None: A write through the name (`x[i] = ...`, `x.attr = ...`) or an in-place method call (`x.mul_(...)`) leaves the base name in a `Load` context, so it is not a - plain `Name` store (see `_rebound_names`). - """ + plain `Name` store (see `_rebound_names`).""" if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance( node.ctx, (ast.Store, ast.Del) ): @@ -711,8 +702,7 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: Non-linear submodules are transparent too (e.g. the dropout GPT-style attentions apply after their output projection). Never walks through a leaf call (e.g. an attention interface): its inputs are what attention consumes, - not what produced the value. - """ + not what produced the value.""" stack = [node] seen: set[fx.Node] = set() while stack: @@ -735,8 +725,7 @@ def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None: """Nearest linear consuming `node`'s output, walking through casts/scalings. Never walks through a leaf call (e.g. an attention interface): what crosses - it is consumed by the attention computation, not projected. - """ + it is consumed by the attention computation, not projected.""" queue = list(node.users) seen: set[fx.Node] = set() while queue: diff --git a/vllm/model_executor/models/transformers/layers.py b/vllm/model_executor/models/transformers/layers.py index 51c4bfef98cb..4ff3651dfc37 100644 --- a/vllm/model_executor/models/transformers/layers.py +++ b/vllm/model_executor/models/transformers/layers.py @@ -22,8 +22,7 @@ def _resolve(module: str, name: str): """Return `name` from the hw-agnostic `module` when enabled and available, - else from vLLM. Logs which source was used. - """ + else from vLLM. Logs which source was used.""" if envs.VLLM_USE_HW_AGNOSTIC: try: obj = getattr(importlib.import_module(f"{_HW_PKG}.{module}"), name) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 73cf42cc5f80..480bc7f66e8d 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -75,8 +75,7 @@ def forward( ) -> torch.Tensor: """In Transformers `experts.forward` will have this signature. - We discard any extra kwargs because we cannot use them here. - """ + We discard any extra kwargs because we cannot use them here.""" # Note: we need to forward through a custom op so the topk_ids # can be transferred without interfering with cudagraphs. return torch.ops.vllm.transformers_moe_forward( @@ -368,8 +367,7 @@ def custom_routing_function( moe_state: TransformersMoEState, ): """Return `topk_weights` from `gating_output` and the - `topk_ids` we stored in the layer earlier. - """ + `topk_ids` we stored in the layer earlier.""" topk_weights = gating_output topk_ids = moe_state.topk_ids assert topk_ids is not None diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 7b435d40ba02..855b267c02f5 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -424,8 +424,7 @@ def _get_prompt_updates( out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: """Empty, because `apply` writes the placeholder ranges itself rather than - deriving them from updates. - """ + deriving them from updates.""" return [] def _get_mm_token_ids(self, modality: str) -> list[int]: @@ -459,8 +458,7 @@ def _apply_audio( num_audios: int, ) -> dict[str, list[PlaceholderRange]]: """Take each contiguous run of the audio token as one item's placeholder, - and record how many tokens the run holds. - """ + and record how many tokens the run holds.""" audio_token_ids = self._get_mm_token_ids("audio") prompt_tensor = torch.tensor(prompt_ids) is_audio = torch.isin(prompt_tensor, torch.tensor(audio_token_ids)) @@ -518,8 +516,7 @@ def _apply_vision( mm_token_type_ids: torch.Tensor | None, ) -> dict[str, list[PlaceholderRange]]: """Split the positions the processor marks as image into one placeholder per - item, sized by the token count it reports for each image. - """ + item, sized by the token count it reports for each image.""" hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) if mm_token_type_ids is None: raise ValueError( @@ -572,8 +569,7 @@ def apply( timing_ctx: TimingContext, ) -> MultiModalInput: """Process the prompt and every multi-modal item in one HF processor call, - then read the placeholder ranges out of the token ids it returns. - """ + then read the placeholder ranges out of the token ids it returns.""" prompt = inputs.prompt mm_items = inputs.mm_data_items hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs @@ -711,8 +707,7 @@ def _get_prompt_updates( out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: """Replace each modality's placeholder token with the token ids that item's - replacement text encodes to, marking which of them hold embeddings. - """ + replacement text encodes to, marking which of them hold embeddings.""" hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) tokenizer = self.info.get_tokenizer() @@ -907,8 +902,7 @@ def _find_encoder_classes( self, model: "PreTrainedModel" ) -> dict[str, type["PreTrainedModel"]]: """Modalities whose encoder cannot be told apart from the model itself are - omitted, as are those `get_encoder` rejects. - """ + omitted, as are those `get_encoder` rejects.""" encoder_classes: dict[str, type[PreTrainedModel]] = {} for modality in _MODALITY_TO_TOKEN_TYPE_ID: try: @@ -1027,8 +1021,7 @@ def forward( def get_language_model(self) -> torch.nn.Module: """Transformers modeling backend multimodal classes do not contain a separate vLLM language model class. Therefore, in order to return a language model vLLM - class, we use a wrapper to give `self` the same interface as a text model. - """ + class, we use a wrapper to give `self` the same interface as a text model.""" # Exclude self and object bases = self.__class__.mro()[1:-1] # Keep only classes defined in `vllm.model_executor.models.transformers` @@ -1182,8 +1175,7 @@ def _select_item_kwargs( ) -> dict[str, Any]: """Narrow the entries of `kwargs` that hold one row per item down to the item at `index`. Length is all there is to match on, so an unrelated entry of the - same length is narrowed too. - """ + same length is narrowed too.""" return { key: value[index : index + 1] if isinstance(value, (torch.Tensor, list)) and len(value) == num_items diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 47921b988d88..ddb2f4fb7139 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -201,8 +201,7 @@ def setter(obj: object, value: object): class _UninitializedEmbedding(nn.Embedding): """Make `__init__` inert, so that `VocabParallelEmbedding.__init__`' - call to `super().__init__` does not invoke `nn.Embedding.__init__`. - """ + call to `super().__init__` does not invoke `nn.Embedding.__init__`.""" def __init__(self): pass @@ -210,8 +209,7 @@ def __init__(self): class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbedding): """Orders `VocabParallelEmbedding` ahead of `nn.Embedding` in the MRO, so that - `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding. - """ + `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding.""" @lru_cache diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 5d76066b826c..eaa64793de32 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -47,8 +47,7 @@ class WeightsMapper: """Maps the name of each weight if they match the following patterns. - If a key maps to a value of `None`, the corresponding weight is ignored. - """ + If a key maps to a value of `None`, the corresponding weight is ignored.""" orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list) orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) @@ -172,8 +171,7 @@ def get_rename_mapper(self) -> "WeightsMapper": rather than being rewritten to the stacked vLLM name (`qkv_proj`). Mappings to `None` are dropped because "do not load this weight" is meaningless to such a consumer, and applying it would silently shrink a quantization config's ignore - list or make LoRA name parsing fail. - """ + list or make LoRA name parsing fail.""" remove_none = lambda d: {k: v for k, v in d.items() if v is not None} return replace( self, diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index fb684ddba4d9..0c7cfb598e4f 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -116,8 +116,7 @@ def get_multimodal_config() -> MultiModalConfig | None: """Return the current ``MultiModalConfig``, or ``None`` when no engine config context is active (e.g., during unit tests) or when the current ``model_config`` does not carry a ``multimodal_config`` (e.g., minimal - stubs used in tests). - """ + stubs used in tests).""" vllm_config = get_current_vllm_config_or_none() if vllm_config is None or vllm_config.model_config is None: return None diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py index 33418847da66..a31979498890 100644 --- a/vllm/model_executor/utils.py +++ b/vllm/model_executor/utils.py @@ -174,8 +174,7 @@ def get_moe_expert_mapping( It will be retrieved from the first module that has a `get_expert_mapping` method. If the model manually implements `get_expert_mapping`, it will be used. - Otherwise, it will use the first RoutedExperts layer. - """ + Otherwise, it will use the first RoutedExperts layer.""" for _, module in model.named_modules(): get_mapping = getattr(module, "get_expert_mapping", None) if get_mapping is not None: diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py index d467e8eacf33..2065b3833822 100644 --- a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py +++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py @@ -18,8 +18,7 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: """Warm every reachable QSA specialization: indexer decode-query-length - profiles plus the sparse attention split-K/merge configs. - """ + profiles plus the sparse attention split-K/merge configs.""" qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa") attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa") if qsa_module is None or attn_module is None: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 124573b0e397..fd02a51076ed 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -1059,8 +1059,7 @@ def _fused_eh_norm_kernel( ): """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` - ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat. - """ + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index e60083b1184f..c1a6903937b6 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1425,8 +1425,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index f17555850098..919146c551a4 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -91,8 +91,7 @@ def validate_image_sentinel_ids(tokenizer) -> None: def grid_tokens(best_height, best_width, patch_size, downsample_ratio): """Number of LLM tokens the aligner grid occupies (N-layout, including - row/align padding). - """ + row/align padding).""" n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) num_tokens = n_llm_h * (n_llm_w + 1) + 2 @@ -197,8 +196,7 @@ def load_image( def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): """Builds the N-layout token types (final order) and the aligner-row order - for IMAGE slots. - """ + for IMAGE slots.""" compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO pad_h = n_llm_h % 2 rows = n_llm_h + pad_h @@ -241,8 +239,7 @@ def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): class DeepseekV4VLImageProcessor: """Per-image transform (the PIL-input equivalent of the reference - ``load_image``). - """ + ``load_image``).""" def __init__(self, config: DeepseekV4Config) -> None: super().__init__() diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index d0168ff25e5f..669186edbb1b 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -312,8 +312,7 @@ def __call__( def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: """Weights dtypes the model's indexer layers ask for: fp32 for the dense scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer - (`SparseMQAIndexer`) is enabled. - """ + (`SparseMQAIndexer`) is enabled.""" if vllm_config.attention_config.indexer_sparse_logits: return (torch.float32, torch.bfloat16) return (torch.float32,) diff --git a/vllm/models/deepseek_v4/cpu/cpu_compressor.py b/vllm/models/deepseek_v4/cpu/cpu_compressor.py index a0c856dc5505..6f0b9e06d842 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_compressor.py +++ b/vllm/models/deepseek_v4/cpu/cpu_compressor.py @@ -42,8 +42,7 @@ def cache_norm_weight_fp32(self) -> None: weight loading (see ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU kernel requires fp32/contiguous but the checkpoint loads this weight - in bf16, and it never changes after loading. - """ + in bf16, and it never changes after loading.""" self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous() def forward( diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py index a7f2af1cbb55..52d4f4c3b017 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_mla.py +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -118,8 +118,7 @@ def build_tile_scheduler( ) -> dict[str, FlashMLASchedMeta | None]: """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA C++ decode-path concern) -- always return the all-``None`` sentinel - the shared base's own CPU branch returns. - """ + the shared base's own CPU branch returns.""" return { _LAYER_TYPE_SWAONLY: None, _LAYER_TYPE_C4A: None, diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index efbd6768f62a..a2c3b00493e6 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -176,15 +176,13 @@ def _fused_indexer_q_rope_quant_cpu( class _NoOpEvent: """Stand-in for ``torch.cuda.Event`` on CPU: since CPU never sets an aux - stream list, these events are constructed but never actually used. - """ + stream list, these events are constructed but never actually used.""" class DeepseekV4CPUIndexerCache(DeepseekV4IndexerCache): """CPU indexer K-cache descriptor: same fields as the shared base, just pointing ``get_attn_backend()`` at ``DeepseekV4CPUIndexerBackend`` - instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``. - """ + instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``.""" def get_attn_backend(self) -> type[AttentionBackend]: return DeepseekV4CPUIndexerBackend @@ -211,8 +209,7 @@ def forward( compressor run straight-line instead of through ``maybe_execute_in_parallel``; the short-context check drops the CUDA-only ``is_current_stream_capturing()`` guard (always False - here). - """ + here).""" compressor = self.compressor attn_metadata = get_forward_context().attn_metadata @@ -312,8 +309,7 @@ def _wrap_wo_a_process_weights_after_loading(self) -> None: """Snapshot and pack ``wo_a``'s weight into a bf16 copy for ``bmm_cpu`` before the FP8 kernel's own ``process_weights_after_loading`` VNNI-repacks it in place for - row-major reads ``_o_proj`` never does. - """ + row-major reads ``_o_proj`` never does.""" wo_a = self.wo_a quant_method = wo_a.quant_method orig_pwal = quant_method.process_weights_after_loading @@ -404,8 +400,7 @@ def _fused_qnorm_rope_kv_insert( ) -> torch.Tensor: """CPU override: only the fp8_ds_mla (uint8) SWA cache layout is supported here, so the base method's bf16/per-tensor-fp8 branches - are dropped. - """ + are dropped.""" if not isinstance(attn_metadata, dict): # Profile run: kernel doesn't fire; produce a padded tensor so # downstream FlashMLA gets the right shape. @@ -450,8 +445,7 @@ def _split_qkv_and_norm( ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """CPU override: two ``RMSNorm`` calls instead of the shared ``fused_q_kv_rmsnorm``, whose raw ``@triton.jit`` kernel fails to - link under triton-cpu (``undefined symbol: __truncdfbf2``). - """ + link under triton-cpu (``undefined symbol: __truncdfbf2``).""" qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) return self.q_norm(qr), None, self.kv_norm(kv) @@ -469,8 +463,7 @@ def _prepare_and_attn( ) -> None: """CPU override: no aux streams, so query-projection+KV-insert, the indexer, and the compressor run straight-line instead of through - execute_in_parallel/maybe_execute_in_parallel. - """ + execute_in_parallel/maybe_execute_in_parallel.""" attn_metadata = get_forward_context().attn_metadata indexer = self.indexer compressor = self.compressor @@ -519,8 +512,7 @@ def _sparse_indexer_and_attn( out: torch.Tensor, ) -> None: """CPU override: identical body, minus ``@eager_break_during_capture`` - (this platform never captures a CUDA graph, so it was a no-op). - """ + (this platform never captures a CUDA graph, so it was a no-op).""" if self.indexer is not None and index_q is not None: assert index_weights is not None # index_q_scale is always None on CPU (FP8-only path). diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index a71b6eb63564..6e88d21c2794 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -907,8 +907,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/nvidia/fi_moe.py b/vllm/models/deepseek_v4/nvidia/fi_moe.py index 96515374d682..2e94e75e7dd9 100644 --- a/vllm/models/deepseek_v4/nvidia/fi_moe.py +++ b/vllm/models/deepseek_v4/nvidia/fi_moe.py @@ -159,8 +159,7 @@ def __init__( def _realloc_nvfp4_params(self) -> None: """Swap the mx-recipe scale params for the NVFP4 checkpoint's: fp8-e4m3 per-16 block scales plus the per-tensor second-level - scales (weight_scale_2, input_scale) modelopt exports. - """ + scales (weight_scale_2, input_scale) modelopt exports.""" n_e = self.num_local_experts inter = self.intermediate_size hidden = self.hidden_size diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 70dc01efabf7..883ef5ecafd4 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -56,8 +56,7 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: def _packed_block_span(pool: torch.Tensor) -> int: """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == block_size for unpacked KV, larger when packed (#44577). Raises if not - token-aligned. - """ + token-aligned.""" block_stride = pool.stride(0) token_stride = pool.stride(-2) if block_stride % token_stride != 0: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 3bc5399f5378..4aba14cd03b1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1985,8 +1985,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index ba47bd9631c2..f3e468ebb1f5 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1378,8 +1378,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 50bd8e7cb7ec..b3eddc202452 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1062,8 +1062,7 @@ def get_model_state_cls(): @property def token_lookback_depth(self) -> int: """Tokens before a chunk start the engram hash needs; the model runner - passes them as `lookback_token_ids`. - """ + passes them as `lookback_token_ids`.""" engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 @@ -1087,8 +1086,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-collapse residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 0e379037b490..3227b50a4e45 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -209,8 +209,7 @@ def _build_image_span( self, image_embeds: torch.Tensor, types: torch.Tensor ) -> torch.Tensor: """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END. - """ + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" types = types.to(image_embeds.device) span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) dtype = image_embeds.dtype diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 77f5958e5f58..4275d25a76fb 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -1066,8 +1066,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: def _compressed_kv_cache(self) -> torch.Tensor: """The compressed-KV cache tensor of this layer's kv source (own - cache for kv-source layers). - """ + cache for kv-source layers).""" if self.is_kv_source: return self.kv_cache assert self.compressed_cache_prefix is not None diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index 857063ecfb14..c415bbaa6e42 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -564,8 +564,7 @@ def _engram_head_shard_weight_loader( param: torch.nn.Parameter, loaded_weight: torch.Tensor ) -> None: """Load this rank's complete head buckets. ue8m0 scales arrive as - float8_e8m0fnu; keep the raw bytes (the param stores uint8). - """ + float8_e8m0fnu; keep the raw bytes (the param stores uint8).""" part_rows = param.shape[0] if loaded_weight.dtype == torch.float8_e8m0fnu: loaded_weight = loaded_weight.view(torch.uint8) @@ -740,8 +739,7 @@ def lookup( def forward(self, indices: torch.Tensor) -> torch.Tensor: """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] - bf16, gathered from all shards for this replica's tokens. - """ + bf16, gathered from all shards for this replica's tokens.""" out = torch.empty( (indices.shape[0], self.part_n_hash_cols, self.dim), dtype=torch.bfloat16, @@ -1002,8 +1000,7 @@ def forward( ) -> torch.Tensor: """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all tokens, pre sequence-parallel shard); token_mask: [T], False shuts - the gate so those positions pass through untouched. - """ + the gate so those positions pass through untouched.""" kv = self.wkv(self.embed(hash_ids).flatten(-2)) num_kv_tokens = hash_ids.shape[0] assert token_mask is None or token_mask.shape == (num_kv_tokens,) diff --git a/vllm/models/deepseek_v41/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py index ce56713e8c23..4f87b816c9a5 100644 --- a/vllm/models/deepseek_v41/common/mm_preprocess.py +++ b/vllm/models/deepseek_v41/common/mm_preprocess.py @@ -82,8 +82,7 @@ def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int: def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): """Largest aspect-preserving pixel size whose token grid still fits in - max_n_token. Returns (best_height, best_width). - """ + max_n_token. Returns (best_height, best_width).""" r = height / width max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 max_h_float = max_w_float * r @@ -106,8 +105,7 @@ def safe_resize( height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token ): """Shrink the pixel size until the image costs at most max_n_token LLM - tokens. - """ + tokens.""" n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio) if num_image_tokens(n_llm_h, n_llm_w) > max_n_token: best_height, best_width = solve_resize_ratio( @@ -173,8 +171,7 @@ def image_token_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor: class DeepseekV4VLImageProcessor: """Per-image transform (the PIL-input equivalent of the reference - ``load_image``). - """ + ``load_image``).""" def __init__(self, config: DeepseekV41Config) -> None: super().__init__() diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 650581a7dfd0..0b43cfad190b 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -54,8 +54,7 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: def _packed_block_span(pool: torch.Tensor) -> int: """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == block_size for unpacked KV, larger when packed (#44577). Raises if not - token-aligned. - """ + token-aligned.""" block_stride = pool.stride(0) token_stride = pool.stride(-2) if block_stride % token_stride != 0: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 4d9d9d3108ea..21fa00e5c963 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -1179,8 +1179,7 @@ def get_model_state_cls(): @property def token_lookback_depth(self) -> int: """Tokens before a chunk start the engram hash needs; the model runner - passes them as `lookback_token_ids`. - """ + passes them as `lookback_token_ids`.""" engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 @@ -1204,8 +1203,7 @@ def forward( def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-collapse residual stream buffer (max_num_batched_tokens, hc_mult * hidden_size) for the MTP draft model. Populated by - forward(); valid after each target step. - """ + forward(); valid after each target step.""" return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index 21386790cb40..bfbcaa89a1a1 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -209,8 +209,7 @@ def _build_image_span( self, image_embeds: torch.Tensor, types: torch.Tensor ) -> torch.Tensor: """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END. - """ + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" types = types.to(image_embeds.device) span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) dtype = image_embeds.dtype diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 7e78a8fe3294..97d8314c0929 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -98,8 +98,7 @@ def _kpool_softmax_rotate_write_cache_kernel( BLOCK_D: tl.constexpr, ): """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> - Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``. - """ + Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" row = tl.program_id(0) do_write = True if HAS_WRITE_MASK: diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 03e423b13cb3..436b9685cb12 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -130,8 +130,7 @@ def _resolve_kda_prefill_backend( """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton ``chunk_kda_with_fused_gate`` path. ``backend`` comes from - ``additional_config.kda_prefill_backend`` (auto / triton / flashkda). - """ + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" if backend not in ("auto", "triton", "flashkda"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") capability = current_platform.get_device_capability() @@ -369,8 +368,7 @@ def _flashkda_prefill( bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the attention output into ``out`` (a workspace buffer when ``None``) and - returns ``(out, final_state)``. - """ + returns ``(out, final_state)``.""" assert self._flashkda_buffer_specs is not None final_state, workspace, workspace_out = ( current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py index 42cabc35b59f..d2b21345d052 100644 --- a/vllm/models/glm5next/common/multimodal.py +++ b/vllm/models/glm5next/common/multimodal.py @@ -720,8 +720,7 @@ class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor): (image token repeat, video frame/timestamp structure) is owned by vLLM's prompt-update machinery — the inherited ``_get_prompt_updates`` builds the replacement content and the placeholder scan validates against - exactly that. - """ + exactly that.""" def _hf_processor_applies_updates( self, diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py index 703cbb321a82..ad3845457063 100644 --- a/vllm/models/glm5next/common/sparse_indexer.py +++ b/vllm/models/glm5next/common/sparse_indexer.py @@ -138,8 +138,7 @@ def _gather_workspace_shapes( """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + - (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales. - """ + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" if use_fp4_cache: return ( ((total_seq_lens, head_dim // 2), torch.uint8), @@ -157,8 +156,7 @@ def kv_cache_as_quant_view( use_fp4_cache: bool, ) -> torch.Tensor: """4D ``[num_blocks, block_size, 1, head_width]`` view expected by - DeepGEMM, from the 3D indexer kv-cache allocation. - """ + DeepGEMM, from the 3D indexer kv-cache allocation.""" if use_fp4_cache: assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 num_blocks, block_size, _ = kv_cache.shape diff --git a/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py index f069fdf2c25f..095db3917b49 100644 --- a/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py +++ b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py @@ -31,8 +31,7 @@ def _fused_eh_norm_kernel( ): """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` - ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat. - """ + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" tok = tl.program_id(0) off = tl.arange(0, BLOCK) mask = off < H diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 010b231354ac..75ccfa8c2a37 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -163,8 +163,7 @@ def _kpool_softmax_rotate_write_cache_kernel( BLOCK_D: tl.constexpr, ): """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> - Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``. - """ + Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" row = tl.program_id(0) do_write = True if HAS_WRITE_MASK: diff --git a/vllm/models/inkling/amd/model.py b/vllm/models/inkling/amd/model.py index 48ed8a2caab3..711a33855cfc 100644 --- a/vllm/models/inkling/amd/model.py +++ b/vllm/models/inkling/amd/model.py @@ -72,8 +72,7 @@ def _sconv_add_norm( ROCm uses the portable collective path. The Lamport P2P implementation in the NVIDIA model relies on CUDA GDC and CUDA-specific symmetric-memory - publication semantics which cannot be linked into a gfx950 HSACO. - """ + publication semantics which cannot be linked into a gfx950 HSACO.""" norm_w = norm.weight if norm is not None else None eps = norm.variance_epsilon if norm is not None else 0.0 diff --git a/vllm/models/inkling/amd/moe.py b/vllm/models/inkling/amd/moe.py index 3a71b3ea5ecc..d8141eb4cf08 100644 --- a/vllm/models/inkling/amd/moe.py +++ b/vllm/models/inkling/amd/moe.py @@ -192,8 +192,7 @@ def inkling_gate_select( class InklingGate(nn.Module): """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and - global scale (the served checkpoint's only configuration). - """ + global scale (the served checkpoint's only configuration).""" def __init__( self, @@ -247,8 +246,7 @@ def select_experts( self, gating_output: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Full selection: (weights, ids) of [T, K + S]. The first K entries - are the routed top-k; the S trailing entries are the sink gammas. - """ + are the routed top-k; the S trailing entries are the sink gammas.""" return inkling_gate_select( gating_output, self.n_total_experts, @@ -269,8 +267,7 @@ def select_experts( def _inkling_moe_ep_size() -> int: """EP size the MoERunner layer will run with (mirrors FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when - expert parallelism is enabled). - """ + expert parallelism is enabled).""" parallel_config = get_current_vllm_config().parallel_config if not parallel_config.enable_expert_parallel: return 1 diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py index dc7a8fd1e265..f9cba37f3c71 100644 --- a/vllm/models/inkling/amd/mtp.py +++ b/vllm/models/inkling/amd/mtp.py @@ -166,8 +166,7 @@ def embed_input_ids( ) -> torch.Tensor: """Draft-prefill embedding: fused gather + backbone embed_norm, then the target's tower embeddings scattered in unnormed (the backbone - convention — MM embeds are merged after embed_norm). - """ + convention — MM embeds are merged after embed_norm).""" norm = self.backbone_embed_norm embeds = embed_rmsnorm( input_ids, @@ -196,8 +195,7 @@ def fused_input_cat( """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one launch: embedding row gather + the backbone embed_norm + the depth embed_norm chain on one side, hidden_norm on the other, written - straight into the cat buffer. - """ + straight into the cat buffer.""" hidden_w = layer.hidden_norm.weight embed_w = layer.embed_norm.weight eps = layer.hidden_norm.variance_epsilon @@ -329,8 +327,7 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is - skipped entirely. - """ + skipped entirely.""" return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/inkling/amd/ops/norm.py b/vllm/models/inkling/amd/ops/norm.py index caee8a047e2d..d34c1c3259b9 100644 --- a/vllm/models/inkling/amd/ops/norm.py +++ b/vllm/models/inkling/amd/ops/norm.py @@ -238,8 +238,7 @@ def embed_rmsnorm( embedding-table gather. ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the first decoder layer's pre-attention norm) as a second output, still one - launch. Bit-exact vs the unfused module sequence. - """ + launch. Bit-exact vs the unfused module sequence.""" ids = input_ids.view(-1) (T,) = ids.shape n = embed_table.shape[1] @@ -333,8 +332,7 @@ def embed_dual_rmsnorm_cat( (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the backbone embed_norm in front of the depth embed_norm (bit-exact vs the - unfused sequence). The concat copies collapse into direct writes. - """ + unfused sequence). The concat copies collapse into direct writes.""" T, n = hidden.shape if embeds is not None: assert embeds.shape == hidden.shape diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index 4363cc883c4f..aeb460d64a21 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -75,8 +75,7 @@ def _sconv_add_norm( """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``. The Lamport path performs reduce-scatter + shard sconv + all-gather + - residual add + norm. The NCCL path handles unsupported configurations. - """ + residual add + norm. The NCCL path handles unsupported configurations.""" attn_metadata = get_forward_context().attn_metadata m = ( attn_metadata.get(sconv.owner.prefix) diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 07c4ce59691c..a8ea968076bb 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -191,8 +191,7 @@ def inkling_gate_select( class InklingGate(nn.Module): """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and - global scale (the served checkpoint's only configuration). - """ + global scale (the served checkpoint's only configuration).""" def __init__( self, @@ -244,8 +243,7 @@ def select_experts( self, gating_output: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Full selection: (weights, ids) of [T, K + S]. The first K entries - are the routed top-k; the S trailing entries are the sink gammas. - """ + are the routed top-k; the S trailing entries are the sink gammas.""" return inkling_gate_select( gating_output, self.n_total_experts, @@ -266,8 +264,7 @@ def select_experts( def _inkling_moe_ep_size() -> int: """EP size the FusedMoE layer will run with (mirrors FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when - expert parallelism is enabled). - """ + expert parallelism is enabled).""" parallel_config = get_current_vllm_config().parallel_config if not parallel_config.enable_expert_parallel: return 1 diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index dc7a8fd1e265..f9cba37f3c71 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -166,8 +166,7 @@ def embed_input_ids( ) -> torch.Tensor: """Draft-prefill embedding: fused gather + backbone embed_norm, then the target's tower embeddings scattered in unnormed (the backbone - convention — MM embeds are merged after embed_norm). - """ + convention — MM embeds are merged after embed_norm).""" norm = self.backbone_embed_norm embeds = embed_rmsnorm( input_ids, @@ -196,8 +195,7 @@ def fused_input_cat( """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one launch: embedding row gather + the backbone embed_norm + the depth embed_norm chain on one side, hidden_norm on the other, written - straight into the cat buffer. - """ + straight into the cat buffer.""" hidden_w = layer.hidden_norm.weight embed_w = layer.embed_norm.weight eps = layer.hidden_norm.variance_epsilon @@ -329,8 +327,7 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is - skipped entirely. - """ + skipped entirely.""" return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/models/inkling/nvidia/ops/mm_towers.py b/vllm/models/inkling/nvidia/ops/mm_towers.py index 76a339377c52..021cf5274bd2 100644 --- a/vllm/models/inkling/nvidia/ops/mm_towers.py +++ b/vllm/models/inkling/nvidia/ops/mm_towers.py @@ -60,8 +60,7 @@ def dmel_embed_sum_norm( eps: float, ) -> torch.Tensor: """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no - [T, NB, D] intermediate). - """ + [T, NB, D] intermediate).""" T, nb = dmel_idx.shape D = weight.shape[1] assert weight.shape[0] % nb == 0 @@ -152,8 +151,7 @@ def rmsnorm_gelu( ) -> torch.Tensor: """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block when D is small. With ``fold``, x must be [N, T, H, W, D] and the output - comes back as ``fold_timespace_to_depth(result, *fold)``. - """ + comes back as ``fold_timespace_to_depth(result, *fold)``.""" D = x.shape[-1] flat = x.reshape(-1, D) assert flat.stride(1) == 1 and flat.stride(0) == D diff --git a/vllm/models/inkling/nvidia/ops/norm.py b/vllm/models/inkling/nvidia/ops/norm.py index caee8a047e2d..d34c1c3259b9 100644 --- a/vllm/models/inkling/nvidia/ops/norm.py +++ b/vllm/models/inkling/nvidia/ops/norm.py @@ -238,8 +238,7 @@ def embed_rmsnorm( embedding-table gather. ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the first decoder layer's pre-attention norm) as a second output, still one - launch. Bit-exact vs the unfused module sequence. - """ + launch. Bit-exact vs the unfused module sequence.""" ids = input_ids.view(-1) (T,) = ids.shape n = embed_table.shape[1] @@ -333,8 +332,7 @@ def embed_dual_rmsnorm_cat( (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the backbone embed_norm in front of the depth embed_norm (bit-exact vs the - unfused sequence). The concat copies collapse into direct writes. - """ + unfused sequence). The concat copies collapse into direct writes.""" T, n = hidden.shape if embeds is not None: assert embeds.shape == hidden.shape diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 574757844a6e..87bf59749c1d 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -745,8 +745,7 @@ def _decode_concat_cache( slot_mapping: torch.Tensor, ) -> torch.Tensor: """Fused decode query-concat + latent cache insert, dispatched by cache - dtype (same policy as prefill: fp8 cache -> fp8 query). - """ + dtype (same policy as prefill: fp8 cache -> fp8 query).""" if self.kv_cache_dtype == "fp8_ds_mla": cache = self.kv_cache if cache.dtype != torch.uint8: diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 1a4ffdbd19ee..50cdb5518381 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -334,8 +334,7 @@ def _aiter_moe_fused_shared_experts_enabled( class MiniMaxM3MoE(nn.Module): """Sigmoid-routed MoE block with a routing-bias correction and a shared - expert. - """ + expert.""" def __init__( self, diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 24a9bd6c09dc..ac02bca1fcc8 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -173,8 +173,7 @@ class MiniMaxM3IndexerPrefillMetadata: @dataclass class MiniMaxM3IndexerDecodeMetadata: """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform - per-request query length (1, or 1 + num_speculative_tokens). - """ + per-request query length (1, or 1 + num_speculative_tokens).""" seq_lens: torch.Tensor # [num_decodes] int32 block_table: torch.Tensor @@ -207,8 +206,7 @@ class MiniMaxM3IndexerMetadataBuilder( AttentionMetadataBuilder[MiniMaxM3IndexerMetadata] ): """Abstract base: shared setup only. The Triton and MSA builders are - parallel subclasses that each own their full ``build`` (no shared code). - """ + parallel subclasses that each own their full ``build`` (no shared code).""" # Full cudagraphs for uniform decode batches (incl. spec-decode verify). _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 5f24645f8827..f385f82067fd 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -203,8 +203,7 @@ class MiniMaxM3SparsePrefillMetadata: @dataclass class MiniMaxM3SparseDecodeMetadata: """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform - per-request query length (1, or 1 + num_speculative_tokens). - """ + per-request query length (1, or 1 + num_speculative_tokens).""" seq_lens: torch.Tensor # [num_decodes] int32 block_table: torch.Tensor diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py index 0b3a336cd5c0..8a015d4e6eae 100644 --- a/vllm/models/minimax_m3/nvidia/indexer_msa.py +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -85,8 +85,7 @@ class MiniMaxM3IndexerMSAPrefillMetadata: class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): """Decode reuses the inherited base ``decode`` field (the Triton decode metadata); ``prefill_msa`` carries the fmha score plan for the prefill side - (the base ``prefill`` field is unused on this path). - """ + (the base ``prefill`` field is unused on this path).""" prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None # Per-forward view (``[:num_tokens]``) of the builder's persistent unified @@ -112,8 +111,7 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill - fmha plan is built eagerly (prefill batches are not captured). - """ + fmha plan is built eagerly (prefill batches are not captured).""" _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 734f675fbe62..970b06d407ea 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -190,8 +190,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class MiniMaxM3MoE(nn.Module): """Sigmoid-routed MoE block with a routing-bias correction and a shared - expert. - """ + expert.""" def __init__( self, @@ -783,8 +782,7 @@ def forward( @property def ffn_all_reduce_deferred(self) -> bool: """This layer's FFN output is left un-reduced; the caller fuses the - all-reduce into the next RMSNorm. - """ + all-reduce into the next RMSNorm.""" if self.is_moe_layer: return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce return not self.mlp.down_proj.reduce_results diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache.py index fa6dabce4128..b4fabfc21ffb 100644 --- a/vllm/multimodal/cache.py +++ b/vllm/multimodal/cache.py @@ -606,8 +606,7 @@ def get_and_update_item( @override def touch_sender_cache_item(self, mm_hash: str) -> None: """Touch the item in shared memory cache to prevent eviction. - Increments writer_flag on sender side. - """ + Increments writer_flag on sender side.""" self._shm_cache.touch(mm_hash) @override @@ -822,8 +821,7 @@ def touch_receiver_cache_item( mm_item: MultiModalKwargsItem | None = None, ) -> None: """Touch the item in shared memory cache to prevent eviction. - Increments reader_count on receiver side. - """ + Increments reader_count on receiver side.""" assert mm_item is not None if "address" in mm_item: address = cast(int, mm_item["address"].data) diff --git a/vllm/multimodal/image.py b/vllm/multimodal/image.py index 4b1d67f2fb86..ff2da044a664 100644 --- a/vllm/multimodal/image.py +++ b/vllm/multimodal/image.py @@ -38,8 +38,7 @@ def rgba_to_rgb( def _has_transparency(image: Image.Image) -> bool: """Detect whether an image carries transparency data (RGBA, LA, PA, - or tRNS chunk in P/L/RGB PNGs). - """ + or tRNS chunk in P/L/RGB PNGs).""" if image.mode in ("RGBA", "LA", "PA"): return True return "transparency" in getattr(image, "info", {}) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 585acd410d7f..91fa34e86422 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -295,8 +295,7 @@ def _is_dense(t: torch.Tensor) -> bool: """Whether `t`'s elements fill a contiguous storage range, possibly permuted (e.g. a transposed matrix). Such layouts can be copied with pitched cudaMemcpy2D/3DAsync directly from pinned memory, whereas - gapped layouts (e.g. strided slices) stage through a pageable temp. - """ + gapped layouts (e.g. strided slices) stage through a pageable temp.""" if not t.is_contiguous(): expected_stride = 1 for stride, size in sorted(zip(t.stride(), t.shape)): diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index c12acfa1e193..3627dd9355ed 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -220,8 +220,7 @@ def strip_covered_mm_data( so the workers never consume the payload fields. M-RoPE models keep CPU-side metadata fields used to compute positions. SHM address items are also kept so workers can balance the sender's reference count. The - scheduler-side ``Request`` keeps the full features. - """ + scheduler-side ``Request`` keeps the full features.""" if not mm_features or num_computed_tokens == 0: return mm_features diff --git a/vllm/parser/engine/incremental_lexer.py b/vllm/parser/engine/incremental_lexer.py index b8f44ad5358c..8c3820b59ab1 100644 --- a/vllm/parser/engine/incremental_lexer.py +++ b/vllm/parser/engine/incremental_lexer.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Incremental text lexer that converts text chunks into terminal -tokens, with prefix-match buffering for ambiguous boundaries. -""" +tokens, with prefix-match buffering for ambiguous boundaries.""" from __future__ import annotations diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py index 67d52bc34c10..21d053ef28da 100644 --- a/vllm/parser/engine/streaming_parser_engine.py +++ b/vllm/parser/engine/streaming_parser_engine.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Streaming parser engine that orchestrates token ID scanning, -incremental lexing, and state-machine-driven semantic event emission. -""" +incremental lexing, and state-machine-driven semantic event emission.""" from __future__ import annotations diff --git a/vllm/parser/engine/token_id_scanner.py b/vllm/parser/engine/token_id_scanner.py index 6ef0427b5c4e..60a64dcae281 100644 --- a/vllm/parser/engine/token_id_scanner.py +++ b/vllm/parser/engine/token_id_scanner.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Scan delta token IDs for special tokens and split the stream into -pre-lexed terminals and plain text chunks. -""" +pre-lexed terminals and plain text chunks.""" from __future__ import annotations diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 6b1bcc7efaa8..2081fb278464 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -67,8 +67,7 @@ def _scan_json_value(raw: str, start: int) -> int | None: """Return the end index (exclusive) of the JSON object starting at - ``raw[start]``, or ``None`` when the object is still unterminated. - """ + ``raw[start]``, or ``None`` when the object is still unterminated.""" depth = 0 in_string = False escape = False diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 96817b25d65a..a46b607c8e47 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -62,8 +62,7 @@ def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int CUDA_VISIBLE_DEVICES has already been set to the desired value. # This can be removed and simply replaced with torch.cuda.get_device_count - # after https://github.com/pytorch/pytorch/pull/122815 is released. - """ + # after https://github.com/pytorch/pytorch/pull/122815 is released.""" # Note: cuda_visible_devices is not used, but we keep it as an argument for # LRU Cache purposes. diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 70937913e2d0..13e5b55402c6 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -42,8 +42,7 @@ def set_assigned_physical_gpu_ids(ids: list[int]) -> None: Idempotent: a second call with the same value is a no-op. Raises RuntimeError if called again with a different value. - This is expected to run during single-threaded worker initialization. - """ + This is expected to run during single-threaded worker initialization.""" global _assigned_physical_gpu_ids if _assigned_physical_gpu_ids is not None: if _assigned_physical_gpu_ids != ids: diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9bf95557ce30..9474ec4835df 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -91,8 +91,7 @@ def _rocm_device_count_stateless(cuda_visible_devices: str | None = None) -> int CUDA_VISIBLE_DEVICES has already been set to the desired value. # This can be removed and simply replaced with torch.cuda.get_device_count - # after https://github.com/pytorch/pytorch/pull/122815 is released. - """ + # after https://github.com/pytorch/pytorch/pull/122815 is released.""" # Note: cuda_visible_devices is not used, but we keep it as an argument for # LRU Cache purposes. @@ -127,8 +126,7 @@ def _get_wsl_kernel_version() -> tuple[int, ...] | None: def _sync_hip_cuda_env_vars(): """Ensure HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES are consistent. - Treats empty string as unset. Raises on genuine conflicts. - """ + Treats empty string as unset. Raises on genuine conflicts.""" hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index ae9da1f4f7c9..bc679407c39e 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -98,8 +98,7 @@ def start(self) -> None: def step(self) -> None: """Update the profiler state at each worker step, - to handle delayed starts and max iteration limits. - """ + to handle delayed starts and max iteration limits.""" if not self._active: return diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 36224a62b3c1..2cb4dd646cbb 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -377,8 +377,7 @@ def warmup(self, chat_params: ChatParams) -> None: async def clear_mm_cache_async(self) -> None: """Serialize clear_mm_cache through the multimodal executor to avoid - races with concurrent process_inputs on the mm_processor_cache. - """ + races with concurrent process_inputs on the mm_processor_cache.""" await self._clear_mm_cache_async() def shutdown(self) -> None: diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 26df7f20934c..a32824eebda5 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -103,8 +103,7 @@ def _ensure_prompt_embeds_placeholder_token(tokenizer: HfTokenizer) -> int: """Register `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` as a special token and return - its token ID. - """ + its token ID.""" cached = _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE.get(tokenizer) if cached is not None: return cached @@ -193,8 +192,7 @@ def _build_mixed_prompt_embeds( positions: list[tuple[int, int]], ) -> tuple[torch.Tensor, list[bool]]: """Build the full-length `prompt_embeds` tensor and the `is_token_ids` - mask aligned to `token_ids`. - """ + mask aligned to `token_ids`.""" total_len = len(token_ids) hidden_size = prompt_embeds_tensors[0].shape[1] dtype = prompt_embeds_tensors[0].dtype diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 09e0451ccc70..5f30b552ba83 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -616,8 +616,7 @@ def _convert_chat_logprobs_to_completion_logprobs( logprobs: ChatCompletionLogProbs, ) -> CompletionLogProbs: """Convert ChatCompletionLogProbs (per-token objects) to CompletionLogProbs - (parallel flat lists) as required by the /v1/completions response schema. - """ + (parallel flat lists) as required by the /v1/completions response schema.""" if logprobs.content is None: return CompletionLogProbs() diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 8072892fe70b..6cd7b4e81c8a 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -800,8 +800,7 @@ def num_logprobs(self) -> int | None: """Number of sample logprobs to return per output token, or `None` if no sample logprobs were requested. Takes `logprob_token_ids` into account: when `logprobs` is unset but `logprob_token_ids` is set, - returns `len(logprob_token_ids)`. - """ + returns `len(logprob_token_ids)`.""" if self.logprobs is not None: return self.logprobs return len(self.logprob_token_ids) if self.logprob_token_ids else None diff --git a/vllm/tool_parsers/hermes_tool_parser.py b/vllm/tool_parsers/hermes_tool_parser.py index 0a300964102c..c5b318cd19fc 100644 --- a/vllm/tool_parsers/hermes_tool_parser.py +++ b/vllm/tool_parsers/hermes_tool_parser.py @@ -219,8 +219,7 @@ def extract_tool_calls_streaming( Returns a ``DeltaMessage`` containing either plain content (for text preceding any tool call) or one or more ``DeltaToolCall`` - entries, or ``None`` if there is nothing new to send yet. - """ + entries, or ``None`` if there is nothing new to send yet.""" try: # Extract any content before tool calls. content = self._extract_content(current_text) diff --git a/vllm/tool_parsers/lfm2_tool_parser.py b/vllm/tool_parsers/lfm2_tool_parser.py index 7d59005a7a5b..2a0bab07e2a0 100644 --- a/vllm/tool_parsers/lfm2_tool_parser.py +++ b/vllm/tool_parsers/lfm2_tool_parser.py @@ -114,8 +114,7 @@ def _strip_echo(raw_after: str) -> str: """Drop any orphan <|tool_call_end|> (and the preceding text) from trailing content. LFM2 occasionally echoes the call body after the first end token and caps it with a second end token; everything - through the last such orphan is model garbage, not user content. - """ + through the last such orphan is model garbage, not user content.""" last_orphan = raw_after.rfind(TOOL_CALL_END) if last_orphan != -1: return raw_after[last_orphan + len(TOOL_CALL_END) :] diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 1617e25ee174..d2c8c6f41220 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -426,8 +426,7 @@ def get_json_schema_from_tools( class UnexpectedAstError(Exception): """Raised when the AST structure does not match the expected - pythonic tool call format. - """ + pythonic tool call format.""" pass diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 24111f52a34d..54f3086431ac 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -539,8 +539,7 @@ def _wrapper(config_parser_cls): def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> None: """Some models may have no rope_theta in their config but still use RoPE. - This function sets a default rope_theta if it's missing. - """ + This function sets a default rope_theta if it's missing.""" if getattr(config, "rope_parameters", None) is None: config.rope_parameters = {"rope_type": "default"} if "rope_theta" not in config.rope_parameters: @@ -549,8 +548,7 @@ def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> No def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None: """Patch legacy RoPE type fields for backwards compatibility with - older custom models which would otherwise fail to load. - """ + older custom models which would otherwise fail to load.""" # No RoPE parameters to patch if rope_parameters is None: return diff --git a/vllm/transformers_utils/configs/extract_hidden_states.py b/vllm/transformers_utils/configs/extract_hidden_states.py index ab705d583462..2beec0e3081b 100644 --- a/vllm/transformers_utils/configs/extract_hidden_states.py +++ b/vllm/transformers_utils/configs/extract_hidden_states.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Config definitions for ExtractHiddenStatesModel, to be used with -the extract_hidden_states spec decoding method. -""" +the extract_hidden_states spec decoding method.""" import os diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index f8cb761848c1..64af034d9ac4 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -37,8 +37,7 @@ def _transformers_v4_compatibility_import(): subset of `ProcessorChatTemplateKwargs` as defined in Transformers v4. In Transformers v5 these were merged into `ProcessorChatTemplateKwargs` and `ChatTemplateLoadKwargs` was removed. For backward compatibility, we add an alias - for `ChatTemplateLoadKwargs` if it doesn't exist. - """ + for `ChatTemplateLoadKwargs` if it doesn't exist.""" old_import = getattr(processing_utils, "ChatTemplateLoadKwargs", None) new_import = getattr(processing_utils, "ProcessorChatTemplateKwargs", None) if old_import is None and new_import is not None: @@ -53,8 +52,7 @@ def _transformers_v4_compatibility_init() -> Any: processor instance before calling the original `ProcessorMixin.__init__`. This can be removed if `Molmo2ForConditionalGeneration` is upstreamed to - Transformers. - """ + Transformers.""" if hasattr(ProcessorMixin.__init__, "_vllm_patched"): return @@ -397,8 +395,7 @@ def get_feature_extractor( **kwargs: Any, ): """Load an audio feature extractor for the given model name - via HuggingFace. - """ + via HuggingFace.""" try: processor_name = convert_model_repo_to_path(processor_name) feature_extractor = AutoFeatureExtractor.from_pretrained( diff --git a/vllm/transformers_utils/processors/cohere_asr.py b/vllm/transformers_utils/processors/cohere_asr.py index bfa8fca50dde..5665d3cbaeb0 100644 --- a/vllm/transformers_utils/processors/cohere_asr.py +++ b/vllm/transformers_utils/processors/cohere_asr.py @@ -537,8 +537,7 @@ def __call__( class CohereASRProcessor(ProcessorMixin): """HF-compatible processor combining CohereASRFeatureExtractor and a - tokenizer. - """ + tokenizer.""" feature_extractor_class = "CohereASRFeatureExtractor" tokenizer_class = "AutoTokenizer" diff --git a/vllm/transformers_utils/processors/glm5next.py b/vllm/transformers_utils/processors/glm5next.py index 6ad174882392..b6988d3b8fed 100644 --- a/vllm/transformers_utils/processors/glm5next.py +++ b/vllm/transformers_utils/processors/glm5next.py @@ -261,8 +261,7 @@ def _resize_or_pad( allow_upscale: bool = False, ) -> torch.Tensor: """Resize onto the aligned canvas, or keep the aspect ratio and - zero-pad the right/bottom sides (``resize_mode="pad"``). - """ + zero-pad the right/bottom sides (``resize_mode="pad"``).""" height, width = stacked_images.shape[-2:] if resize_mode == "resize": @@ -308,8 +307,7 @@ def _pixel_budget( ) -> tuple[int, int]: """(min_pixels, max_pixels) from the token bounds of ``processor_config.json``; one vision token covers - ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels. - """ + ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels.""" if min_image_tokens is None or max_image_tokens is None: raise ValueError( "min_image_tokens and max_image_tokens must be provided by " diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py index c5dc1a8ff3cc..ad07d901b4fe 100644 --- a/vllm/transformers_utils/utils.py +++ b/vllm/transformers_utils/utils.py @@ -116,8 +116,7 @@ def parse_safetensors_file_metadata(path: str | PathLike) -> dict[str, Any]: def convert_model_repo_to_path(model_repo: str) -> str: """When VLLM_USE_MODELSCOPE is True convert a model - repository string to a Path str. - """ + repository string to a Path str.""" if not envs.VLLM_USE_MODELSCOPE or Path(model_repo).exists(): return model_repo from modelscope.utils.file_utils import get_model_cache_root diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index 89d3e6ca7c73..296fe466460f 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -274,8 +274,7 @@ def get_visible_memory_node() -> list[int]: def _synthesize_cpu_list() -> list[LogicalCPUInfo]: """Synthesize a flat CPU list: each logical CPU is its own core on NUMA node 0. Used when lscpu output is unavailable or unparsable - (e.g. macOS, RISC-V). - """ + (e.g. macOS, RISC-V).""" cpu_count = os.cpu_count() assert cpu_count return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)] diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 009c7b42582c..3d96b496f1b1 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -704,8 +704,7 @@ def fp8_fp4_paged_mqa_logits( def has_deep_gemm_sparse_mqa() -> bool: """Whether the installed DeepGEMM provides the sparse-indexer kernels - (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only). - """ + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" _lazy_init() return ( _fp8_fp4_sparse_mqa_logits_impl is not None diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 096f0d005bee..8d0e7ae1871f 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -72,8 +72,7 @@ def enable_gpu_sync_check() -> None: def _sync_warning_hook(message, category, filename, lineno, file=None, line=None): """Turn torch's sync warning into an error, but only on a thread that is - being checked and is outside any allow region. - """ + being checked and is outside any allow region.""" if _TORCH_SYNC_WARNING in str(message): mode = _checking.get() if mode is None or _allow_depth.get(): @@ -252,8 +251,7 @@ def _is_dense(t: torch.Tensor) -> bool: """Whether `t`'s elements fill a contiguous storage range, possibly permuted (e.g. a transposed matrix). Such layouts can be copied with pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas - gapped layouts (e.g. strided slices) stage through a pageable temp. - """ + gapped layouts (e.g. strided slices) stage through a pageable temp.""" if t.is_contiguous(): # The common case, answered by a memoized TensorImpl flag. return True diff --git a/vllm/utils/platform_utils.py b/vllm/utils/platform_utils.py index 975ad696d986..5d7fed3c9904 100644 --- a/vllm/utils/platform_utils.py +++ b/vllm/utils/platform_utils.py @@ -29,8 +29,7 @@ def cuda_get_device_properties( device, names: Sequence[str], init_cuda=False ) -> tuple[Any, ...]: """Get specified CUDA device property values without initializing CUDA in - the current process. - """ + the current process.""" if init_cuda or cuda_is_initialized(): props = torch.cuda.get_device_properties(device) return tuple(getattr(props, name) for name in names) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index ab395bd7047f..849d0a38f55c 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -230,8 +230,7 @@ def read_v1(cg_dir: str) -> float | None: def available_cpu_count() -> int: """CPUs actually usable by this process: scheduling affinity capped by - the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind). - """ + the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind).""" if sys.platform != "linux": return os.cpu_count() or 1 count = len(os.sched_getaffinity(0)) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 792eeb37b99d..afafa3c92565 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -365,8 +365,7 @@ def validate_configuration( @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: """Layouts this backend's kernels can consume, most preferred first, or - None when the kernels consume any layout and express no preference. - """ + None when the kernels consume any layout and express no preference.""" return None @classmethod @@ -558,8 +557,7 @@ def unpadded( class AttentionCGSupport(Enum): """Constants for the cudagraph support of the attention backend Here we do not consider the cascade attention, as currently - it is never cudagraph supported. - """ + it is never cudagraph supported.""" ALWAYS = 3 """Cudagraph always supported; supports mixed-prefill-decode""" diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index d2b9f5a526e8..b6864897bf9b 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -84,8 +84,7 @@ def supports_sliding_window(cls) -> bool: @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """CPU attention supports decoder, - encoder-only and encoder-decoder attention. - """ + encoder-only and encoder-decoder attention.""" return attn_type in ( AttentionType.DECODER, AttentionType.ENCODER, diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 694cef4106cf..3113c261dd98 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -398,8 +398,7 @@ class FlashInferBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """NVFP4 stores K and V as separate per-head slots of packed fp4 data - plus fp8 block scales. - """ + plus fp8 block scales.""" if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4: return spec hs_k = nvfp4_kv_cache_full_dim(spec.head_size) diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 22ee41a4799b..5f9ac7f17adf 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -63,8 +63,7 @@ def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread - count in place of SM count. - """ + count in place of SM count.""" ideal_splits = 1 while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT): ideal_splits *= 2 diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 6211a4e06c4c..265fb779c3cc 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -847,8 +847,7 @@ def _gather_prefill_chunk( prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """All-gather this rank's upconverted KV shard so the chunk's rows attend - the whole context, and map their top-k onto the rank-major result. - """ + the whole context, and map their top-k onto the rank-major result.""" shard_rows = int(chunk.chunk_tot_seqlen) assert self.gathered_kv_workspace is not None gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows] diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 258f20215f13..53206afd9337 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1118,8 +1118,7 @@ def _split_pcp_dcp_prefill_chunks( def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: """Per-request KV lengths the prefill chunker budgets logits with; - subclasses whose logits rows are wider than the context override. - """ + subclasses whose logits rows are wider than the context override.""" return seq_lens_cpu @staticmethod diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index edea4347271a..1b589cb295f4 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -70,8 +70,7 @@ def is_available(cls) -> bool: def supports_quant_output(self, quant_key: "QuantKey") -> bool: """Whether `run_prefill_new_tokens` can write quantized output directly (fused) for the given quant key, skipping the post-quant - pass. Overridden by backends that support it. - """ + pass. Overridden by backends that support it.""" return False def supports_out(self) -> bool: diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index b9eacdabca12..9e3267c52626 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -110,8 +110,7 @@ def get_impl_cls() -> type["RocmAiterUnifiedAttentionImpl"]: @classmethod def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: """Keep K and V packed in the content dim, unlike the native HIP - kernels the base class targets. - """ + kernels the base class targets.""" # block_size == 1 is the per-token page-size probe (see # Platform.get_page_size_bytes); real blocks must be gatherable in # 16-token units by the ROCm kernel. diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index afb545eee8f1..2671249fb2b9 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -244,8 +244,7 @@ def supports_attn_type(cls, attn_type: str) -> bool: @classmethod def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: """K and V as two head groups so the native HIP kernels address each side - as one contiguous region (x-packed interior applied in split_kv_cache). - """ + as one contiguous region (x-packed interior applied in split_kv_cache).""" if spec.state_content_bytes is not None: return spec assert spec.head_size == spec.head_size_v, ( diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index cf7e571599cd..1b290c0ae9d6 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -275,8 +275,7 @@ class TritonAttentionBackend(AttentionBackend): @classmethod def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack inline fp32 scales after each head's - data, so the content is (data + one scale) per K/V side. - """ + data, so the content is (data + one scale) per K/V side.""" mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: return spec diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py index 90b620b345ee..0dc4150cbf15 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py @@ -45,8 +45,7 @@ def _store_quantized_value( BLOCK_GRP: tl.constexpr, ): """Uniform quantization of values to VQB bits, pack, and store with - scale/zero in the per-block SoA metadata region (Opt#3 layout). - """ + scale/zero in the per-block SoA metadata region (Opt#3 layout).""" if VQB == 3: val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to( tl.float32 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 907b9690e022..8061999409c4 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -461,8 +461,7 @@ def find_longest_cache_hit( ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]: """Returns the per-group hit blocks, the hit length, and the number of ``num_uncached_common_prefix_tokens`` (a shared prefix that a - sparse-retention group has not cached yet; 0 unless hybrid). - """ + sparse-retention group has not cached yet; 0 unless hybrid).""" pass def new_step_starts(self) -> None: diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9f6997b53d3a..49f62a2670a4 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -856,8 +856,7 @@ def get_zeroing_block_ids_in_range( self, request_id: str, start_token: int, end_token: int ) -> list[int]: """The request's block ids covering [start_token, end_token), from - the groups whose new blocks are zeroed by the worker. - """ + the groups whose new blocks are zeroed by the worker.""" ids: list[int] = [] for mgr in self.coordinator.single_type_managers: if mgr.records_new_block_ids: diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index d5a059bf4c4c..c868d249c8a5 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -178,8 +178,7 @@ def get_num_unfinished_requests(self) -> int: def has_unfinished_requests(self) -> bool: """Returns True if there are unfinished requests in the scheduler's - internal queue. - """ + internal queue.""" return self.get_num_unfinished_requests() > 0 @abstractmethod @@ -199,8 +198,7 @@ def has_finished_requests(self) -> bool: def has_requests(self) -> bool: """Returns True if there are unfinished requests, or finished requests - not yet returned in SchedulerOutputs. - """ + not yet returned in SchedulerOutputs.""" return self.has_unfinished_requests() or self.has_finished_requests() @property diff --git a/vllm/v1/core/sched/request_queue.py b/vllm/v1/core/sched/request_queue.py index ab94e9075294..ce58e6e1c7c3 100644 --- a/vllm/v1/core/sched/request_queue.py +++ b/vllm/v1/core/sched/request_queue.py @@ -43,8 +43,7 @@ def prepend_request(self, request: Request) -> None: @abstractmethod def prepend_requests(self, requests: "RequestQueue") -> None: """Prepend all requests from another queue to the front of this - queue. - """ + queue.""" pass @abstractmethod @@ -161,16 +160,14 @@ def prepend_request(self, request: Request) -> None: """Add a request to the queue according to priority policy. Note: In a priority queue, there is no concept of prepending to the - front. Requests are ordered by (priority, arrival_time). - """ + front. Requests are ordered by (priority, arrival_time).""" self.add_request(request) def prepend_requests(self, requests: RequestQueue) -> None: """Add all requests from another queue according to priority policy. Note: In a priority queue, there is no concept of prepending to the - front. Requests are ordered by (priority, arrival_time). - """ + front. Requests are ordered by (priority, arrival_time).""" for request in requests: self.add_request(request) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 64aa5fa2ce8f..3e399a1a816e 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -871,8 +871,7 @@ async def notify_kv_transfer_request_rejected( ) -> None: """Submit a pre-aborted request so the connector's request_finished hook runs to free any pre-admission KV-transfer resources (e.g. NIXL - prefill blocks pinned on the P node). - """ + prefill blocks pinned on the P node).""" request = EngineCoreRequest( request_id=request_id, prompt_token_ids=[0], diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 5125954c42e3..9c84827a63d9 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -598,8 +598,7 @@ def _attach_iteration_details( def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). - Overridden by the DP engine core; never throttles otherwise. - """ + Overridden by the DP engine core; never throttles otherwise.""" return False def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: @@ -1623,8 +1622,7 @@ def _invoke_utility_method( @staticmethod def _convert_msgspec_args(method, args): """If a provided arg type doesn't match corresponding target method - arg type, try converting to msgspec object. - """ + arg type, try converting to msgspec object.""" if not args: return args arg_types = signature(method).parameters.values() @@ -2006,8 +2004,7 @@ def _send_abort_outputs(self, aborted_reqs: list[Request]) -> None: class DPEngineCoreProc(EngineCoreProc): """ZMQ-wrapper for running EngineCore in background process - in a data parallel context. - """ + in a data parallel context.""" def __init__( self, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 6a08f2d3afc3..4ca565d56cd4 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -248,8 +248,7 @@ def collective_rpc( def dp_engines_running(self) -> bool: """Returns True if data parallel engines are collectively in a - running state. - """ + running state.""" raise NotImplementedError async def commit_elastic_ep(self) -> None: @@ -446,8 +445,7 @@ def dp_engines_running(self) -> bool: @dataclass class BackgroundResources: """Used as a finalizer for clean shutdown, avoiding - circular reference back to the client object. - """ + circular reference back to the client object.""" ctx: zmq.Context # If CoreEngineProcManager, it manages local engines; @@ -801,8 +799,7 @@ def monitor_engine_cores(): def _apply_ready_response(self, payload: bytes) -> None: """Decode an EngineCoreReadyResponse and sync any post-initialization - config changes (e.g. auto-fitted max_model_len) back to the frontend. - """ + config changes (e.g. auto-fitted max_model_len) back to the frontend.""" vllm_config = self.vllm_config response = ( msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) @@ -1333,8 +1330,7 @@ async def get_status(self): class DPAsyncMPClient(AsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) - EngineCore. Assumes external load-balancing by default. - """ + EngineCore. Assumes external load-balancing by default.""" def __init__( self, @@ -1518,8 +1514,7 @@ def get_core_engine_for_request(self, request: EngineCoreRequest): class DPLBAsyncMPClient(DPAsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) - EngineCore. Load-balances between multiple engine processes. - """ + EngineCore. Load-balances between multiple engine processes.""" def __init__( self, @@ -1827,8 +1822,7 @@ async def _prepare_scale_up_elastic_ep( num_redundant_experts: int, ) -> None: """Prepare scale up by creating new engine cores and reconfiguring - existing ones. - """ + existing ones.""" self._setup_elastic_ep_reconfig_bootstrap() # Phase 1: Send reconfig messages to existing engines @@ -1920,8 +1914,7 @@ async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> N async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: """Scale down the data parallel size by shutting down and - reconfiguring existing engine cores. - """ + reconfiguring existing engine cores.""" cur_data_parallel_size = len(self.core_engines) self.eep_scaling_cache = ElasticScalingCache( diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index c75fdc53cf43..9658c030ef94 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -147,8 +147,7 @@ def decode_next(self, next_token_id: int) -> str: def get_next_output_text(self, finished: bool, delta: bool) -> str: """If delta is True, only new text since the last call to - this method is returned - """ + this method is returned""" # We return the full output text if the sequence is finished. buffer_length = 0 if finished else self.stop_buffer_length if not delta: diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 43f2b80b05e2..c1002ba1dc5b 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -1041,8 +1041,7 @@ def get_engine_zmq_addresses( Set ``defer_api_server_ports=False`` only when the consumer cannot report a bound port back (e.g. the Rust front-end). IPC paths are - unaffected. - """ + unaffected.""" parallel_config = vllm_config.parallel_config local_engine_count = parallel_config.data_parallel_size_local local_start_index = parallel_config.data_parallel_rank_local diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index ce8a8db826a7..6b07bad17c76 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -284,8 +284,7 @@ def save_sharded_state( @abstractmethod def check_health(self) -> None: """Checks if the executor is healthy. If not, it should raise an - exception. - """ + exception.""" raise NotImplementedError def shutdown(self) -> None: diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 009df6ffba2e..dbaefe218665 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -384,8 +384,7 @@ def collective_rpc( # type: ignore[override] ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: """Returns single result if unique_reply_rank and/or an output - aggregator is provided, otherwise list. - """ + aggregator is provided, otherwise list.""" assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -452,8 +451,7 @@ def get_response(): def _ensure_worker_termination(worker_procs: list[BaseProcess]): """Ensure that all worker processes are terminated. Assumes workers have received termination requests. Waits for processing, then sends - termination and kill signals if needed. - """ + termination and kill signals if needed.""" def wait_for_termination(procs, timeout): if not time: @@ -854,8 +852,7 @@ def death_pipe_monitor(queues_to_shutdown: list[MessageQueue]): @staticmethod def worker_main(*args, **kwargs): """Worker initialization and execution loops. - This runs a background process - """ + This runs a background process""" # Signal handler used for graceful termination. # SystemExit exception is only raised once to allow this and worker # processes to terminate without error @@ -1099,8 +1096,7 @@ def setup_proc_title_and_log_prefix(enable_ep: bool) -> None: def set_multiprocessing_worker_envs(local_world_size: int = 1): """Set up environment variables that should be used when there are workers in a multiprocessing environment. This should be called by the parent - process before worker processes are created - """ + process before worker processes are created""" _maybe_force_spawn() if current_platform.is_cpu() or "OMP_NUM_THREADS" in os.environ: diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 0f6d2c9c6416..e609fb6e7e3d 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -55,8 +55,7 @@ class RayWorkerWrapper(WorkerWrapperBase): """Ray wrapper for vllm.worker.Worker, allowing Worker to be - lazily initialized after Ray sets CUDA_VISIBLE_DEVICES. - """ + lazily initialized after Ray sets CUDA_VISIBLE_DEVICES.""" def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 4a2ea1799f68..17771d9ae67a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1252,8 +1252,7 @@ def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None: def get_max_layers_per_page_size(self) -> int: """Max number of layers sharing a page size. For a balanced bucket - this equals the number of repetitions of the layer pattern. - """ + this equals the number of repetitions of the layer pattern.""" return Counter( spec.page_size_bytes for spec in self.kv_cache_specs.values() ).most_common(1)[0][1] diff --git a/vllm/v1/kv_cache_layout.py b/vllm/v1/kv_cache_layout.py index d784f51b96ac..7c1decd393a9 100644 --- a/vllm/v1/kv_cache_layout.py +++ b/vllm/v1/kv_cache_layout.py @@ -49,8 +49,7 @@ def is_block_contiguous(self) -> bool: @property def is_block_compact(self) -> bool: """True when each page's [H, N, C] bytes form one contiguous run; i.e. - the L and B dimensions are outermost. - """ + the L and B dimensions are outermost.""" return set(self.value[:2]) == {_DIM_L, _DIM_B} @property diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index b22390e444d9..1144477628d2 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -472,8 +472,7 @@ class CopyRun: """A strided byte correspondence between this worker's physical page and a canonical page: for i in range(num_fragments), fragment i spans [local_offset + i * local_stride, +fragment_size) in the worker's page and - [canonical_offset + i * canonical_stride, +fragment_size) canonically. - """ + [canonical_offset + i * canonical_stride, +fragment_size) canonically.""" local_offset: int canonical_offset: int @@ -554,8 +553,7 @@ class TransferResult: class OffloadingWorker(ABC): """Runs in the worker process. Performs async KV transfers for ONE offloaded medium (e.g. CPU). Direction is explicit via submit_store / - submit_load, so there is no (src_medium, dst_medium) routing. - """ + submit_load, so there is no (src_medium, dst_medium) routing.""" @abstractmethod def submit_store( diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index b16fdeef6c36..037e3631e499 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -126,8 +126,7 @@ def compute_sub_block_ptrs( class CopyPlan(NamedTuple): """Precomputed fragment-copy template for one data ref under the canonical CPU layout, unrolled from the ref's mapped runs. Offsets are relative to - the per-block base pointers on each side. - """ + the per-block base pointers on each side.""" frag_offsets_src: np.ndarray frag_offsets_dst: np.ndarray @@ -165,8 +164,7 @@ def _canonical_page_ids( ) -> np.ndarray: """Global canonical page ids matching compute_sub_block_ptrs' enumeration. These identify canonical pages consistently across ranks, so they key - CanonicalPageMapping.is_writer rotation. - """ + CanonicalPageMapping.is_writer rotation.""" if blocks_per_chunk == 1: return block_ids[:count] flat = ( @@ -179,8 +177,7 @@ def _canonical_block_sizes( layer_refs_per_group: list[list[CanonicalKVCacheRef]], num_tensors: int ) -> list[int]: """Canonical CPU bytes per GPU block for each tensor, taken from the refs' - mappings. Requires every ref to carry a mapping. - """ + mappings. Requires every ref to carry a mapping.""" canonical_bytes_per_block = [0] * num_tensors for layer_refs in layer_refs_per_group: for ref in layer_refs: @@ -343,8 +340,7 @@ def _estimate_max_copy_ops(self, group_sizes: Sequence[int]) -> int: """Upper bound on the number of copy descriptors for a transfer. Exact for the direct layout. The canonical path may fill fewer: - writer rotation later drops the blocks this rank does not write. - """ + writer rotation later drops the blocks this rank does not write.""" num_copy_ops = 0 for g_idx, (group_size, layer_refs) in enumerate( zip(group_sizes, self.layer_refs_per_group) @@ -373,8 +369,7 @@ def _fill_direct_ops( """Fill one group's copy descriptors for the direct (worker-private) layout: one whole-page copy per (block, ref). - Returns (op_idx past the filled descriptors, bytes added). - """ + Returns (op_idx past the filled descriptors, bytes added).""" num_bytes = 0 for data_ref in self.layer_refs_per_group[g_idx]: t_idx = data_ref.tensor_idx @@ -417,8 +412,7 @@ def _fill_canonical_ops( scatter each block through the ref's precomputed CopyPlan, keeping only the blocks this rank writes. - Returns (op_idx past the filled descriptors, bytes added). - """ + Returns (op_idx past the filled descriptors, bytes added).""" assert self._canonical_copy_plans is not None # Zero-copy reinterpretation for pointer arithmetic: uint64 and the # buffers' int64 are bit-equivalent for addresses @@ -505,8 +499,7 @@ def _filter_writer_blocks( ) -> tuple[np.ndarray, np.ndarray]: """Keep only the blocks this rank writes: replicated ranks take turns writing shared canonical pages, keyed by the rank-consistent CPU-side - canonical page id. - """ + canonical page id.""" cpu_page_ids = _canonical_page_ids( group_dst, self.dst_blocks_per_chunk, diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 9eba094d79f3..00b36e863a07 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -28,8 +28,7 @@ def _all_workers_barrier() -> None: """Block until every worker rank has reached this point (gloo cpu group). A superset of the node-local mmap openers suffices: once the barrier - releases, every worker sharing the region file has mapped it. - """ + releases, every worker sharing the region file has mapped it.""" from vllm.distributed.parallel_state import ( get_inner_dp_world_group, get_world_group, @@ -162,8 +161,7 @@ def get_manager(self) -> OffloadingManager: def _uses_shared_region(self) -> bool: """Whether the worker CPU buffer is the shared mmap region (vs a private - per-rank tensor); replicated-layout dedup is gated on this being True. - """ + per-rank tensor); replicated-layout dedup is gated on this being True.""" return current_platform.is_cuda_alike() def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index 93e77591fab2..a6879d8c44e5 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -216,8 +216,7 @@ def drain_results(self) -> None: def mark_miss(self, keys: Collection[OffloadKey]) -> None: """Force the cached verdict for ``keys`` to False after a failed load, so the scheduler stops re-issuing the doomed promotion (livelock, #49176). - Keys with no cached entry are skipped. - """ + Keys with no cached entry are skipped.""" for key in keys: state = self._lookup_state.get(key) if state is not None: diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index aef2b7328de3..5ec13f62daac 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -168,8 +168,7 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @override def drain_jobs(self) -> None: """Synchronous tier — submit_*() returns only after the operation - completes, so there is nothing to wait for. - """ + completes, so there is nothing to wait for.""" return def get_num_chunks(self) -> int: diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py index 53249a0ccd9a..cda7e46e5964 100644 --- a/vllm/v1/kv_offload/tiering/fs/io.py +++ b/vllm/v1/kv_offload/tiering/fs/io.py @@ -139,8 +139,7 @@ def _load_block( ) -> None: """Read one KV block from disk; remove the file only on a provable short read (a too-short file is genuine corruption) and leave it untouched on any - other error. - """ + other error.""" fd: int | None = None view_slice = view.cast("B")[offset : offset + block_size] o_direct = O_DIRECT if use_o_direct else 0 diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index ef545f4e3c2c..3d51d95774cd 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -269,8 +269,7 @@ def load_task() -> None: @override def get_finished_jobs(self) -> Iterable[JobResult]: """Collect finished jobs; a failed promotion marks only its failed keys - as a miss here (scheduler thread). - """ + as a miss here (scheduler thread).""" results = [] for job_id, success, transfer_time in self._pool.get_finished(): if self.events is not None: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ba8dee8fadbe..6e0bf90a332a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -143,8 +143,7 @@ def shutdown(self) -> None: class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the - TieringOffloadingManager with exclude_tier_idx set to the origin tier. - """ + TieringOffloadingManager with exclude_tier_idx set to the origin tier.""" __slots__ = ("_m", "_origin_idx") diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 810a56c7ae85..3d6c66c4834a 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -308,8 +308,7 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: def _poll_active_transfers(self) -> None: """Poll all in-flight transfers once; move newly-completed (success or - failure) into ``_pending_results`` and release their NIXL handles. - """ + failure) into ``_pending_results`` and release their NIXL handles.""" for job_id, entry in list(self._transfers.items()): try: state = self._agent.check_xfer_state(entry.xfer_handle) @@ -364,8 +363,7 @@ def _poll_active_transfers(self) -> None: def get_finished_jobs(self) -> Iterable[JobResult]: """Poll transfers; a failed promotion marks its cached verdicts False - here (scheduler thread). - """ + here (scheduler thread).""" self._poll_active_transfers() results = self._pending_results self._pending_results = [] diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 17ecaf1a8241..703c632899d9 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -134,8 +134,7 @@ class P2PDestInfo: def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: """Parse the consumer sub-dict (PD ``remote_prefiller`` or symmetric - ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete. - """ + ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete.""" role = _remote_prefiller_params(kv_params) do_probe = False if role is None: @@ -156,8 +155,7 @@ def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: def _parse_dest(kv_params: dict | None) -> P2PDestInfo | None: """Parse the producer ``remote_decoder`` sub-dict into a ``P2PDestInfo``, - or None if the block is absent (not a remote-decode request). - """ + or None if the block is absent (not a remote-decode request).""" role = _remote_decoder_params(kv_params) if role is None: return None diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 5cc438c34ffe..9adac0f787a6 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -414,8 +414,7 @@ def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None: """Require a mapping on every ref and record layer portability. Fails loudly rather than persist direct-layout bytes under a - canonical format identity. - """ + canonical format identity.""" all_refs = [ ref for group_refs in kv_caches.group_data_refs for ref in group_refs ] diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 0ff5876fd1ce..50e606d3b3f5 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -91,8 +91,7 @@ def load_stat_logger_plugin_factories() -> list[StatLoggerFactory]: class AggregateStatLoggerBase(StatLoggerBase): """Abstract base class for loggers that - aggregate across multiple DP engines. - """ + aggregate across multiple DP engines.""" @abstractmethod def __init__(self, vllm_config: VllmConfig, engine_indexes: list[int]): ... diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index e1e61720b4a5..dd046a8250b6 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -1820,8 +1820,7 @@ def get_required(obj: object, attr: str): def getattr_from_list(obj: object, attrs: list[str], default: object = None): """Try to get the first attr that exists in the object - from a list of attrs. Otherwise return None. - """ + from a list of attrs. Otherwise return None.""" for attr in attrs: if hasattr(obj, attr): return getattr(obj, attr) diff --git a/vllm/v1/metrics/ray_wrappers.py b/vllm/v1/metrics/ray_wrappers.py index 2d515d5963ad..dedc60ba5425 100644 --- a/vllm/v1/metrics/ray_wrappers.py +++ b/vllm/v1/metrics/ray_wrappers.py @@ -82,8 +82,7 @@ def _get_sanitized_opentelemetry_name(name: str) -> str: class RayGaugeWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Gauge to provide same API as - prometheus_client.Gauge - """ + prometheus_client.Gauge""" def __init__( self, @@ -117,8 +116,7 @@ def set_to_current_time(self): class RayCounterWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Counter to provide same API as - prometheus_client.Counter - """ + prometheus_client.Counter""" def __init__( self, @@ -143,8 +141,7 @@ def inc(self, value: int | float = 1.0): class RayHistogramWrapper(RayPrometheusMetric): """Wraps around ray.util.metrics.Histogram to provide same API as - prometheus_client.Histogram - """ + prometheus_client.Histogram""" def __init__( self, diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py index ea50e56be436..99dd157e78e9 100644 --- a/vllm/v1/sample/logits_processor/builtin.py +++ b/vllm/v1/sample/logits_processor/builtin.py @@ -129,8 +129,7 @@ def __init__(self, _, device: torch.device, is_pin_memory: bool): def is_argmax_invariant(self) -> bool: """Logit bias can rebalance token probabilities and change the - outcome of argmax in greedy sampling. - """ + outcome of argmax in greedy sampling.""" return False def update_state(self, batch_update: BatchUpdate | None): @@ -189,8 +188,7 @@ def __init__( def is_argmax_invariant(self) -> bool: """By censoring stop tokens, min-tokens can change the outcome - of the argmax operation in greedy sampling. - """ + of the argmax operation in greedy sampling.""" return False @staticmethod diff --git a/vllm/v1/sample/logits_processor/state.py b/vllm/v1/sample/logits_processor/state.py index 24360037f7bf..7988e8c6e921 100644 --- a/vllm/v1/sample/logits_processor/state.py +++ b/vllm/v1/sample/logits_processor/state.py @@ -69,8 +69,7 @@ def _ensure_removed_sorted(self) -> None: @property def removed(self) -> list[RemovedRequest]: """Removed request indices sorted in - descending order - """ + descending order""" self._ensure_removed_sorted() return self._removed diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py index 8e1e6cf4f66d..882bc44917fd 100644 --- a/vllm/v1/serial_utils.py +++ b/vllm/v1/serial_utils.py @@ -90,8 +90,7 @@ def _typestr(val: Any) -> tuple[str, str] | None: def _encode_type_info_recursive(obj: Any) -> Any: """Recursively encode type information for nested structures of - lists/dicts. - """ + lists/dicts.""" if obj is None: return None if type(obj) is list: @@ -105,8 +104,7 @@ def _decode_type_info_recursive( type_info: Any, data: Any, convert_fn: Callable[[Sequence[str], Any], Any] ) -> Any: """Recursively decode type information for nested structures of - lists/dicts. - """ + lists/dicts.""" if type_info is None: return data if isinstance(type_info, dict): diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 372fc2201db4..85dde88f2d55 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -259,8 +259,7 @@ def _derive_cpu_config( gpu_config: "KVCacheConfig", cpu_capacity_bytes: int ) -> "KVCacheConfig": """Derive a CPU KVCacheConfig from the GPU config. - Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio. - """ + Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" # Import here to avoid potential circular imports from vllm.v1.kv_cache_interface import KVCacheTensor @@ -315,8 +314,7 @@ def _estimate_lazy_target_blocks( def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: """Bind GPU block pool so that we can touch blocks during stores. - Called by Scheduler after kv_cache_manager is ready. - """ + Called by Scheduler after kv_cache_manager is ready.""" self._gpu_block_pool = gpu_block_pool def get_num_new_matched_tokens( @@ -1131,8 +1129,7 @@ def request_finished( block_ids: list[int], ) -> tuple[bool, dict[str, Any] | None]: """Always returns (False, None). GPU blocks are protected by ref_cnt, - so the scheduler can free blocks immediately. - """ + so the scheduler can free blocks immediately.""" req_id = request.request_id # Release any temp CPU hit pin from get_num_new_matched_tokens() diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index 77700a411fd1..18ae2700e937 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -396,8 +396,7 @@ def load_model(self, target_model: nn.Module) -> None: def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: """Validate all drafting layers belong to the same KV cache group - and record the group index for common_attn_metadata selection. - """ + and record the group index for common_attn_metadata selection.""" assert len(self.attn_layer_names) == 1 layer = self.attn_layer_names[0] for gid, group in enumerate(kv_cache_config.kv_cache_groups): diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index 318b283ceb93..bea677dde85a 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -221,8 +221,7 @@ def load_model(self, target_model: nn.Module) -> None: def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: """Draft layers span multiple KV cache groups (sliding + full attention with different head dimensions), so skip the base - class single-group assertion. - """ + class single-group assertion.""" def initialize_attn_backend( self, @@ -230,8 +229,7 @@ def initialize_attn_backend( kernel_block_sizes: list[int] | None = None, ) -> None: """Create separate AttentionGroup objects per KV cache spec - so that each head-dim variant gets its own metadata builder. - """ + so that each head-dim variant gets its own metadata builder.""" all_attn_layers = get_layers_from_vllm_config( self.vllm_config, AttentionLayerBase, # type: ignore[type-abstract] diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index dff94b0544dd..9f7ad68a88a6 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -785,8 +785,7 @@ def _update_positions_dependent_metadata( block_size: int, ) -> torch.Tensor: """Update positions, slot mappings, and sequence metadata for the - next draft step. Returns the updated positions tensor. - """ + next draft step. Returns the updated positions tensor.""" positions_1d = positions[0] if self.uses_mrope else positions if self.uses_mrope: out_pos = self.mrope_positions[0, :batch_size] diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 6db4663660de..471d51d355e5 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -847,6 +847,5 @@ def update_num_computed_tokens_for_batch_change( def unconditional_to_conditional_rates(rates: list[float]) -> list[float]: """Convert per-position unconditional rates to per-position conditional - rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}). - """ + rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}).""" return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])] diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index ef800f247e9c..a89ccbb9ed7c 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -144,8 +144,7 @@ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: """NOTE: Because this method is non-blocking, explicit synchronization - is needed to ensure the data is copied to CPU. - """ + is needed to ensure the data is copied to CPU.""" if n is None: return self.cpu.copy_(self.gpu, non_blocking=True) return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True) @@ -159,8 +158,7 @@ def get_engine_client_zmq_addr( """Return an IPC path (``local_only=True``) or ``tcp://host:port``. ``port=0`` lets the kernel assign the port at ``bind()`` time; the - caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``. - """ + caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``.""" if local_only: return get_open_zmq_ipc_path() return get_tcp_uri(host, port) @@ -265,8 +263,7 @@ def gather_actual_addresses( ) -> tuple[list[str], list[str]]: """Return (inputs, outputs) reported by each child, indexed by ``client_index``. Raises ``RuntimeError`` on timeout or premature - child exit. - """ + child exit.""" n = len(self._address_pipes) inputs: list[str | None] = [None] * n outputs: list[str | None] = [None] * n @@ -427,8 +424,7 @@ def shutdown(self, timeout: float | None = None) -> None: class _SubprocessWrapper: """Wraps subprocess.Popen to provide the BaseProcess-like interface - needed by wait_for_completion_or_failure. - """ + needed by wait_for_completion_or_failure.""" def __init__(self, proc, name: str): self._proc = proc diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 108459642274..3b587ad71f80 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -76,8 +76,7 @@ def __init__( model: SupportsEncoderCudaGraph, ): """Initialize CUDA graph manager with provided token budgets - and max batch size. - """ + and max batch size.""" self.vllm_config = vllm_config self.device = device self.dtype = dtype diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 239b47f3d710..4690c7eae70e 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -62,8 +62,7 @@ class AttentionState(NamedTuple): @dataclass(frozen=True) class BatchExecutionDescriptor: """Describes the shape of the batch and CG mode to run; this is used to make shape - matches between the capture and runtime. - """ + matches between the capture and runtime.""" cg_mode: CUDAGraphMode num_tokens: int @@ -91,8 +90,7 @@ def make_cudagraph_stats( class CreateForwardFn(Protocol): """Factory that prepares inputs (OUTSIDE the graph) and returns a forward_fn. Called with warmup=True for the warmup pass and warmup=False - for the captured pass. - """ + for the captured pass.""" def __call__( self, @@ -955,8 +953,7 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int: """Extrapolate the total FULL capture cost from samples of the largest graphs. The first capture allocates the pool baseline; later graphs mostly - reuse it, so the second sample is taken as the per-graph cost. - """ + reuse it, so the second sample is taken as the per-graph cost.""" if not mem_samples: return 0 first_capture = mem_samples[0] @@ -993,8 +990,7 @@ def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: def _teardown_profiling_state(runner: "GPUModelRunner") -> None: """Release the profiling KV cache and captured graphs while keeping model - weights, so the real ``initialize_kv_cache`` starts from a clean slate. - """ + weights, so the real ``initialize_kv_cache`` starts from a clean slate.""" torch.accelerator.synchronize() if hasattr(runner.model_state, "_mamba_ctx"): runner.model_state._mamba_ctx = None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index dca5b260f283..97571df8a901 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2228,8 +2228,7 @@ def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None: def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that - memory is reclaimable when running in the same process. - """ + memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() self.cudagraph_manager = None self.fast_prefill = None diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 6f0c9c672c0b..a5c8d6f4cfd7 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -144,8 +144,7 @@ def preprocess_state( ) -> None: """Hook run on real batches before the forward pass (after block tables are gathered). Used by mamba "align" prefix caching to pre-copy state - across block boundaries. No-op by default. - """ + across block boundaries. No-op by default.""" return None def postprocess_state( diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index f1124c3ee470..c2aaa1fba25a 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -38,8 +38,7 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: """Return a bool array of shape `[input_batch.num_reqs]` marking requests that produce a sampled token this step, and therefore must have that token (and the draft block proposed from it) propagated to the earlier PP stages. - Returns None if no request in the batch produces a sample. - """ + Returns None if no request in the batch produces a sample.""" old_computed = input_batch.num_computed_tokens_np prefill_len = input_batch.prefill_len_np # Exclude non-final prefill chunks (they don't produce a sample). @@ -179,8 +178,7 @@ def broadcast_drafts( def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* - requests in the batch. - """ + requests in the batch.""" assert not self.is_last_rank need_sampled_mask = compute_need_sampled_mask(input_batch) if need_sampled_mask is None: diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py index 01ffa4bdb1e6..ff3675db130f 100644 --- a/vllm/v1/worker/gpu/sample/output.py +++ b/vllm/v1/worker/gpu/sample/output.py @@ -76,8 +76,7 @@ def _compact_sampling_mask_kernel( class SamplingMaskTensors(NamedTuple): """Device-side masks pending async D2H: compact ids, plus the bitmask as - the exact fallback for rows wider than ``max_num_kept``. - """ + the exact fallback for rows wider than ``max_num_kept``.""" # [num_requests, max_num_kept] token_ids: torch.Tensor diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index a6199f570ccd..b069604c2964 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -107,8 +107,7 @@ def get_logprobs_dims( self, idx_mapping_np: np.ndarray, include_token_ids: bool = True ) -> tuple[int, int] | None: """(num_logprobs, max_per_req_token_ids) for the given requests, or - None when none of them want logprobs. - """ + None when none of them want logprobs.""" max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) max_token_ids = ( self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np) diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 2d0c25a41e68..7e829cf3093c 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -190,8 +190,7 @@ def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int """Dummy-run kwargs whose step timings seed the cost tables. Run these inside StepTimingCollector.collect(), then hand the block's - timings to set_initial_cost_curves. - """ + timings to set_initial_cost_curves.""" max_num_tokens = self.req_states.max_num_batched_tokens size = self._cudagraph_limit = capture_sizes[-1] if capture_sizes else 0 # Also profile beyond the capture limit: real steps run there @@ -261,8 +260,7 @@ def record_confidences( input_batch: "InputBatch", ) -> None: """Publish this step's raw confidences for the ranking kernel and start - copying them to the CPU, where a later step's budget reads them. - """ + copying them to the CPU, where a later step's budget reads them.""" num_reqs = input_batch.num_reqs ready_idx = self._stale_idx ^ 1 with gpu_sync_allowed(): diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index 854c958a12d1..725714fcdbc0 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -72,8 +72,7 @@ def _prepare_dflash_inputs_to_capture( class DFlashCudaGraphManager(CudaGraphManager): """DFlash CudaGraphManager for the parallel-drafting query forward, - building its own attention metadata from scratch. - """ + building its own attention metadata from scratch.""" def capture( self, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index dd37396b6b4d..eb0f4793bf1f 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -28,8 +28,7 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): """The target's lm_head — from get_language_model() for - *ForConditionalGeneration targets, else the top-level module. - """ + *ForConditionalGeneration targets, else the top-level module.""" return getattr(target_language_model, "lm_head", None) or getattr( target_model, "lm_head", None ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index b09e5196b9fb..299cac641620 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -268,8 +268,7 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: @property def attn_vllm_config(self) -> VllmConfig: """Config for the draft's attention metadata builders. Overridden by - speculators whose attention mode differs from the target's. - """ + speculators whose attention mode differs from the target's.""" return self.vllm_config def set_attn( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bf46c3dbfeff..5c85257fa39f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -481,8 +481,7 @@ def get_output(self) -> ModelRunnerOutput: class ExecuteModelState(NamedTuple): """Ephemeral cached state transferred between execute_model() and - sample_tokens(), after execute_model() returns None. - """ + sample_tokens(), after execute_model() returns None.""" scheduler_output: "SchedulerOutput" logits: torch.Tensor @@ -1072,8 +1071,7 @@ def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor: first scheduled token (column j is position start - 1 - j); -1 where the position is before the prompt or already past it. Generated positions are left to the model: under async scheduling the CPU token - table holds placeholders for them. - """ + table holds placeholders for them.""" buf = self.lookback_token_ids assert buf is not None buf.np.fill(-1) @@ -3206,8 +3204,7 @@ def _execute_mm_encoder( def _get_encoder_output_from_cache(self, mm_hash: str) -> torch.Tensor | None: """Return a cached encoder output for multimodal - embedding gather. - """ + embedding gather.""" return self.encoder_cache.get(mm_hash, None) def _gather_mm_embeddings( @@ -4179,8 +4176,7 @@ def _is_all_reqs_chunked_prefill(self) -> bool: This is true when `discard_request_mask` is set for every scheduled request (e.g., for chunked prefill requests that are not the last - prefill chunk). - """ + prefill chunk).""" num_reqs = self.input_batch.num_reqs return bool(self.discard_request_mask.np[:num_reqs].all()) @@ -6573,8 +6569,7 @@ def _init_minimal_kv_cache_for_profiling(self) -> None: def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that - memory is reclaimable when running in the same process. - """ + memory is reclaimable when running in the same process.""" from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT from vllm.v1.worker.workspace import reset_workspace_manager diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 44621d08fb20..acba1d9b67e6 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -336,8 +336,7 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: def _scoped_allocator_max_split(self, max_split_size_mb: int): """Temporarily set max_split_size_mb to reduce allocator fragmentation at the cost of more cudaMalloc calls (negligible in practice). Restores the original - value on exit. - """ + value on exit.""" if not current_platform.is_cuda(): yield return diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py index 5a570b347094..2c4942070705 100644 --- a/vllm/v1/worker/startup_plan.py +++ b/vllm/v1/worker/startup_plan.py @@ -77,8 +77,7 @@ def compute_plan_fingerprint( def _plan_path(fingerprint: str) -> str: """Plans are regenerable derived state, so they live under the standard vLLM cache root (like the torch.compile cache) and relocate with - ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own. - """ + ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own.""" # VLLM_CACHE_ROOT is already user-expanded by envs.py. return os.path.join( envs.VLLM_CACHE_ROOT, "startup_plan", f"startup_plan_{fingerprint}.json" @@ -135,8 +134,7 @@ def _applicable_kv_cache_memory_bytes( def maybe_apply_startup_plan(worker: "Worker") -> None: """If enabled and ``--kv-cache-memory`` was not set explicitly, apply a persisted plan by setting ``worker.cache_config.kv_cache_memory_bytes``. - No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``. - """ + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``.""" if ( not envs.VLLM_ENABLE_STARTUP_PLAN or worker.cache_config.kv_cache_memory_bytes is not None @@ -169,8 +167,7 @@ def maybe_apply_startup_plan(worker: "Worker") -> None: def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> None: """Atomically persist this boot's profiling result for future boots. No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``; failures are logged, - never raised. - """ + never raised.""" if not envs.VLLM_ENABLE_STARTUP_PLAN: return fingerprint = compute_plan_fingerprint( diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 6496a611e6fb..cf869edb26ff 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -128,8 +128,7 @@ def check_health(self) -> None: def synchronize_device(self) -> None: """Block until in-flight device work completes; backends outside - ``torch.accelerator`` must override with their own wait. - """ + ``torch.accelerator`` must override with their own wait.""" if torch.accelerator.is_available(): torch.accelerator.synchronize() From 9f9e1dac26ff0379651dd7e8ca409b573ccfce54 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Wed, 16 Sep 2026 11:09:00 -0700 Subject: [PATCH 338/774] [Bugfix][Model] Restore causal image SWA for DeepSeek V4.1 (#57152) Signed-off-by: Juntian Liu Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../attention/test_deepseek_v4_swa_visible.py | 87 ++++++++++- vllm/models/deepseek_v41/attention.py | 8 - .../deepseek_v41/common/ops/cache_utils.py | 139 ++++-------------- .../nvidia/flash_mla_mega_attn.py | 19 +-- .../deepseek_v41/nvidia/flashinfer_sparse.py | 4 - vllm/models/deepseek_v41/nvidia/flashmla.py | 21 +-- .../configs/deepseek_v41.py | 8 - vllm/v1/attention/backends/mla/sparse_swa.py | 10 +- 8 files changed, 122 insertions(+), 174 deletions(-) diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index d18b7c671a54..1b5097de8e2d 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -462,13 +462,31 @@ def test_flashinfer_mixed_sparse_indices_with_image_spans(): assert actual == exp_rows[1:] -def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder: - overrides: dict = {"sliding_window": WINDOW} +def make_builder( + vision: bool, + window: int = WINDOW, + max_image_tokens: int = MAX_IMG, + max_num_batched_tokens: int = 64, + model_name: str = "meta-llama/Meta-Llama-3-8B", + builder_cls: type[ + DeepseekSparseSWAMetadataBuilder + ] = DeepseekSparseSWAMetadataBuilder, + mm_prefix_clamp_sliding_window: bool = True, +) -> DeepseekSparseSWAMetadataBuilder: + overrides: dict = {"sliding_window": window} if vision: - overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG) + # Emulate the V4 vision config, whose mm_prefix_clamp_sliding_window + # gates the in-kernel SWA widening. V4.1 (causal image tokens) leaves + # it off. + overrides.update( + vision_n_layers=2, + vision_max_n_token=max_image_tokens, + mm_prefix_clamp_sliding_window=mm_prefix_clamp_sliding_window, + ) vllm_config = create_vllm_config( + model_name=model_name, max_model_len=4096, - max_num_batched_tokens=64, + max_num_batched_tokens=max_num_batched_tokens, max_num_seqs=8, hf_config_override=overrides, ) @@ -477,11 +495,11 @@ def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder: num_kv_heads=1, head_size=512, dtype=torch.bfloat16, - sliding_window=WINDOW, + sliding_window=window, cache_dtype_str="auto", model_version="deepseek_v4", ) - return DeepseekSparseSWAMetadataBuilder( + return builder_cls( kv_cache_spec=spec, layer_names=["layer0"], vllm_config=vllm_config, @@ -578,3 +596,60 @@ def test_builder_text_model_unchanged(): ) assert md.prefill_swa_lens.cpu().tolist() == lens assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("compress_ratio", [0, 1, 2]) +@pytest.mark.parametrize("query_len", [100, 4000]) +def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len): + """V4.1 image tokens use the same causal SWA as text, including in chunks.""" + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + combine_topk_swa_indices as combine_v41, + ) + from vllm.models.deepseek_v41.sparse_mla import DeepseekV41SparseSWAMetadataBuilder + + seq_len, window, max_image_tokens = 4000, 128, 2048 + spans = [(1900, 3947)] + builder = make_builder( + True, + window, + max_image_tokens, + query_len, + builder_cls=DeepseekV41SparseSWAMetadataBuilder, + mm_prefix_clamp_sliding_window=False, + ) + assert builder.max_image_tokens == 0 + md = build_metadata(builder, [seq_len], [query_len], {0: spans}) + assert md.prefill_left_visible is None + assert md.prefill_right_visible is None + assert md.prefill_swa_indices.shape[-1] == window + plan = md.get_prefill_chunk_plan( + compress_ratio, prefill_chunk_size=4, has_compressed=compress_ratio > 0 + ) + [(start, end, n, m)] = plan + gather_len = int(md.prefill_gather_lens[0]) + assert (start, end, m - n) == (0, 1, gather_len) + assert gather_len == query_len + min(seq_len - query_len, window - 1) + indices, lens = combine_v41( + torch.empty(query_len, 0, dtype=torch.int32, device="cuda"), + md.query_start_loc, + md.prefill_seq_lens, + md.prefill_gather_lens, + window, + compress_ratio, + 0, + m, + n, + ) + indices, lens = indices.cpu(), lens.cpu().tolist() + paged_indices = md.prefill_swa_indices.cpu() + paged_lens = md.prefill_swa_lens.cpu().tolist() + for token, pos in enumerate(range(seq_len - query_len, seq_len)): + lo, hi = max(0, pos - window + 1), pos + 1 + row = indices[token, : lens[token]] + assert lens[token] == hi - lo + assert torch.all((row >= n) & (row < m)) + absolute_keys = row - n + seq_len - gather_len + assert absolute_keys.tolist() == list(range(lo, hi)) + assert paged_lens[token] == hi - lo + assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi)) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 4275d25a76fb..5b26a98c1c53 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -288,13 +288,6 @@ def __init__( self.n_groups = config.o_groups self.n_local_groups = self.n_groups // tp_size self.window_size = config.sliding_window - # Vision variant: image spans are visible bidirectionally, widening - # prefill SWA index rows by up to max_image_tokens columns. - self.max_image_tokens = ( - getattr(config, "vision_max_n_token", 0) - if getattr(config, "vision_n_layers", 0) > 0 - else 0 - ) # ---- v4.1 sparse-attention topology ---- # compress_ratios has one entry per layer (MTP layers included): # 0 = pure sliding window, 1 = full-length compressed cache, @@ -576,7 +569,6 @@ def __init__( _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup( window_size=self.window_size, block_size=self.swa_cache_layer.block_size, - max_image_tokens=self.max_image_tokens, ) if self.compress_ratio > 1: diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 9a251e0748c9..e7c3650768c2 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -815,16 +815,12 @@ def combine_topk_swa_indices( M: int, N: int, out: tuple[torch.Tensor, torch.Tensor] | None = None, - left_visible: torch.Tensor | None = None, - right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = topk_indices.shape[0] - # max_image_tokens widens the SWA column region for in-image - # bidirectional visibility; the width is fixed per model so the - # caller-provided workspace stays valid on image-free batches. + # The width is fixed per model so the caller-provided workspace stays + # valid on every batch. combined_topk = ( - (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -853,9 +849,6 @@ def combine_topk_swa_indices( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, - left_visible=left_visible, - right_visible=right_visible, - max_image_tokens=max_image_tokens, ) return combined_indices, combined_lens @@ -924,7 +917,6 @@ class CompileKey: TOP_K: int COMPRESS_RATIO: int WINDOW_SIZE: int - IMAGE_WIDTH: int PADDED_TOP_K: int input_variant: TritonPointerInputVariant @@ -946,14 +938,11 @@ def kernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, - left_visible_ptr, - right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, - IMAGE_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, ): batch_idx = tl.program_id(0) @@ -988,19 +977,8 @@ def kernel( # (integer div by constexpr 0 is UB and yields garbage, which # turns the store index below negative). topk_len = 0 - if IMAGE_WIDTH > 0: - # In-image bidirectional visibility: the window starts up to - # max(left - (window - 1), 0) positions earlier and extends - # `right` positions past the query token (both 0 outside - # image spans, reducing to the plain causal window). - left = tl.load(left_visible_ptr + token_idx) - right = tl.load(right_visible_ptr + token_idx) - else: - left = 0 - right = 0 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) - swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) - swa_len = pos + right - swa_start + 1 + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + swa_len = pos - swa_start + 1 offset = tl.arange(0, PADDED_TOP_K) mask = offset < topk_len @@ -1014,18 +992,17 @@ def kernel( mask=mask, ) # Index into gathered buffer: N + (position - gather_start) - # For positions [swa_start, pos + right], the buffer indices are: - # [N + swa_start - gather_start, N + pos + right - gather_start] - for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE): - swa_offset = i + tl.arange(0, WINDOW_SIZE) - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + swa_offset, - M * batch_idx + N + swa_start + swa_offset - gather_start, - mask=swa_offset < swa_len, - ) + # For positions [swa_start, pos], the buffer indices are: + # [N + swa_start - gather_start, N + pos - gather_start] + swa_offset = tl.arange(0, WINDOW_SIZE) + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + swa_offset, + M * batch_idx + N + swa_start + swa_offset - gather_start, + mask=swa_offset < swa_len, + ) combined_len = topk_len + swa_len tl.store(combined_lens_ptr + token_idx, combined_len) @@ -1041,7 +1018,6 @@ def dispatch( # type: ignore[override] topk: int, compress_ratio: int, WINDOW_SIZE: int, - image_width: int = 0, ) -> CompileKey: padded_topk = next_power_of_2(topk_width) input_variant = TritonPointerInputVariant.from_alignment( @@ -1054,7 +1030,6 @@ def dispatch( # type: ignore[override] TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=WINDOW_SIZE, - IMAGE_WIDTH=image_width, PADDED_TOP_K=padded_topk, input_variant=input_variant, ) @@ -1064,19 +1039,10 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: return [] window_size = _hf_config_int(vllm_config, "sliding_window", 128) - image_width = ( - _hf_config_int(vllm_config, "vision_max_n_token", 0) - if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0 - else 0 - ) - # Warm both the plain-window variant (batches without image spans) - # and the in-image bidirectional variant. - image_widths = [0, image_width] if image_width > 0 else [0] return self._trace_dispatch(self.dispatch)( _combine_topk_swa_warmup_inputs(vllm_config), _COMBINE_TOPK_SWA_POINTER_INPUTS, WINDOW_SIZE=window_size, - image_width=image_widths, ) def compile(self, compile_key: CompileKey) -> None: @@ -1093,14 +1059,11 @@ def compile(self, compile_key: CompileKey) -> None: input_variant.pointer("query_start_loc", torch.int32), input_variant.pointer("seq_lens", torch.int32), input_variant.pointer("gather_lens", torch.int32), - int32_ptr, - int32_ptr, 1, # do not specialize M 1, # do not specialize N TOP_K=compile_key.TOP_K, COMPRESS_RATIO=compile_key.COMPRESS_RATIO, WINDOW_SIZE=compile_key.WINDOW_SIZE, - IMAGE_WIDTH=compile_key.IMAGE_WIDTH, PADDED_TOP_K=compile_key.PADDED_TOP_K, grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), ) @@ -1119,13 +1082,8 @@ def __call__( TOP_K: int, COMPRESS_RATIO: int, WINDOW_SIZE: int, - left_visible: torch.Tensor | None = None, - right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> None: num_reqs = seq_lens.shape[0] - has_image = left_visible is not None - image_width = max_image_tokens if has_image else 0 self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( combined_indices, combined_indices.stride(0), @@ -1135,14 +1093,11 @@ def __call__( query_start_loc, seq_lens, gather_lens, - left_visible if has_image else topk_indices, - right_visible if has_image else topk_indices, M, N, TOP_K=TOP_K, COMPRESS_RATIO=COMPRESS_RATIO, WINDOW_SIZE=WINDOW_SIZE, - IMAGE_WIDTH=image_width, PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), ) @@ -1169,9 +1124,6 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token: torch.Tensor | None = None, swa_block_span: int | None = None, compressed_block_span: int | None = None, - prefill_left_visible: torch.Tensor | None = None, - prefill_right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -1181,21 +1133,11 @@ def build_flashinfer_mixed_sparse_indices( per token). Decode tokens read precomputed SWA/compressed indices; prefill tokens derive their SWA window from the position and translate local compressed indices to global slots via the block tables. - - When ``prefill_left_visible``/``prefill_right_visible`` are given (vision - variant), the SWA column region widens by ``max_image_tokens`` and prefill - tokens inside an image span get a bidirectionally widened window; decode - rows are padded with -1 across the extra columns. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 - swa_index_width = decode_swa_indices.shape[-1] - assert swa_index_width >= window_size - has_image = prefill_left_visible is not None - image_width = max_image_tokens if has_image else 0 - swa_total_width = swa_index_width + image_width - if has_image: - assert prefill_right_visible is not None + swa_total_width = decode_swa_indices.shape[-1] + assert swa_total_width >= window_size if decode_compressed_topk_lens is not None: assert decode_compressed_topk_lens.dtype == torch.int32 assert prefill_topk_indices.dtype == torch.int32 @@ -1254,7 +1196,7 @@ def build_flashinfer_mixed_sparse_indices( if num_tokens == 0: return sparse_indices, sparse_topk_lens - window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) + window_block_size = triton.next_power_of_2(max(swa_total_width, 1)) topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 @@ -1278,8 +1220,6 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token, prefill_topk_indices, prefill_topk_indices.stride(0), - prefill_left_visible if has_image else token_to_req_indices, - prefill_right_visible if has_image else token_to_req_indices, query_start_loc, seq_lens, token_to_req_indices, @@ -1293,8 +1233,7 @@ def build_flashinfer_mixed_sparse_indices( compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, - SWA_INDEX_WIDTH=swa_index_width, - IMAGE_WIDTH=image_width, + SWA_INDEX_WIDTH=swa_total_width, COMPRESS_RATIO=compress_ratio, TOP_K=topk, PADDED_TOP_K=padded_topk, @@ -1349,8 +1288,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel( decode_is_valid_token_ptr, prefill_topk_indices_ptr, prefill_topk_stride, - left_visible_ptr, - right_visible_ptr, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, @@ -1365,7 +1302,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel( NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, SWA_INDEX_WIDTH: tl.constexpr, - IMAGE_WIDTH: tl.constexpr, COMPRESS_RATIO: tl.constexpr, TOP_K: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -1377,18 +1313,14 @@ def _build_flashinfer_mixed_sparse_indices_kernel( TOPK_BLOCK_SIZE: tl.constexpr, ): token_idx = tl.program_id(0) - # Total SWA column count; > SWA_INDEX_WIDTH only for in-image - # bidirectional visibility (vision variant), where the extra columns are - # -1-padded for decode rows. - SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH if token_idx < NUM_DECODE_TOKENS: - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_TOTAL_WIDTH + mask = offset < SWA_INDEX_WIDTH values = tl.load( decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, - mask=offset < SWA_INDEX_WIDTH, + mask=mask, other=-1, ) values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) @@ -1431,7 +1363,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_TOTAL_WIDTH + + SWA_INDEX_WIDTH + offset, values, mask=mask, @@ -1445,7 +1377,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1457,26 +1389,17 @@ def _build_flashinfer_mixed_sparse_indices_kernel( start_pos = seq_len - query_len token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query - if IMAGE_WIDTH > 0: - # In-image bidirectional visibility: window starts up to - # max(left - (window - 1), 0) earlier and extends `right` past pos. - left = tl.load(left_visible_ptr + token_idx) - right = tl.load(right_visible_ptr + token_idx) - else: - left = 0 - right = 0 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) - swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) - swa_len = pos + right - swa_start_pos + 1 + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + swa_len = pos - swa_start_pos + 1 if COMPRESS_RATIO > 0: topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) else: # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB. topk_len = 0 - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_TOTAL_WIDTH + mask = offset < SWA_INDEX_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1520,10 +1443,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_TOTAL_WIDTH + + SWA_INDEX_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 233d2f80f903..9722673dd1ee 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -311,7 +311,7 @@ def _reserve_prefill_workspace(self, q: torch.Tensor) -> None: else: assert self.topk_indices_buffer is not None top_k = self.topk_indices_buffer.shape[-1] - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16), ((self.max_num_batched_tokens, combined_topk), torch.int32), @@ -411,7 +411,7 @@ def _forward_prefill_mega( ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: chunk_size = chunk_end - chunk_start kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous( @@ -457,21 +457,6 @@ def _forward_prefill_mega( chunk_m, chunk_n, out=(idx_ws[: qe - qs], lens_ws[: qe - qs]), - left_visible=( - swa_metadata.prefill_left_visible[ - num_decode_tokens + qs : num_decode_tokens + qe - ] - if swa_metadata.prefill_left_visible is not None - else None - ), - right_visible=( - swa_metadata.prefill_right_visible[ - num_decode_tokens + qs : num_decode_tokens + qe - ] - if swa_metadata.prefill_right_visible is not None - else None - ), - max_image_tokens=self.max_image_tokens, ) chunk_out = _token_slice(out, token_base + qs, token_base + qe) # Mega attention over the gathered non-paged bf16 KV (RoPE already diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0b43cfad190b..409725434a17 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -417,10 +417,6 @@ def _build_sparse_index_metadata( decode_is_valid_token=decode_is_valid_token, swa_block_span=swa_block_span, compressed_block_span=compressed_block_span, - prefill_left_visible=swa_metadata.prefill_left_visible, - prefill_right_visible=swa_metadata.prefill_right_visible, - # getattr for tests that bypass __init__ via object.__new__. - max_image_tokens=getattr(self, "max_image_tokens", 0), ) if swa_only: swa_metadata.flashinfer_sparse_index_cache["swa_only"] = ( diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index 47a7b35c76d6..870098975b71 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -120,9 +120,7 @@ def forward_mqa( else: assert self.topk_indices_buffer is not None top_k = self.topk_indices_buffer.shape[-1] - combined_topk = round_up( - top_k + self.window_size + self.max_image_tokens, 128 - ) + combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), ((self.max_num_batched_tokens, combined_topk), torch.int32), @@ -303,7 +301,7 @@ def _forward_prefill( ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: chunk_size = chunk_end - chunk_start workspace = workspace_manager.get_simultaneous( @@ -361,21 +359,6 @@ def _forward_prefill( chunk_M, chunk_N, out=(combined_indices_out, combined_lens_out), - left_visible=( - swa_metadata.prefill_left_visible[ - num_decode_tokens + query_start : num_decode_tokens + query_end - ] - if swa_metadata.prefill_left_visible is not None - else None - ), - right_visible=( - swa_metadata.prefill_right_visible[ - num_decode_tokens + query_start : num_decode_tokens + query_end - ] - if swa_metadata.prefill_right_visible is not None - else None - ), - max_image_tokens=self.max_image_tokens, ) flash_mla_sparse_fwd( q=q[query_start:query_end], diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py index 8f33245cd68f..d8b4335438c6 100644 --- a/vllm/transformers_utils/configs/deepseek_v41.py +++ b/vllm/transformers_utils/configs/deepseek_v41.py @@ -60,11 +60,3 @@ def __init__( self.vision_max_n_token = vision_config.get("max_image_tokens", 1024) self.vision_min_pixels = vision_config.get("min_pixels", 295936) self.vision_max_wh_ratio = vision_config.get("max_wh_ratio") - # The vision variant needs the mm-prefix plumbing (atomic image-span - # prefill + sparse-SWA window widening); the base arch-config - # convertor reads this config-level attribute. - self.is_mm_prefix_lm = vision_n_layers > 0 - # The sparse-SWA index kernels widen the window within image spans - # in-kernel, so mm-prefix ranges longer than sliding_window must be - # kept (they are only consumed by those kernels). - self.mm_prefix_clamp_sliding_window = vision_n_layers > 0 diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 700f328a37c0..e88631c32bb3 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -452,13 +452,15 @@ def __init__(self, *args, **kwargs): assert hasattr(hf_config, "sliding_window") self.window_size = hf_config.sliding_window - # Vision variant: image spans (up to vision_max_n_token tokens) are + # V4 vision variant: image spans (up to vision_max_n_token tokens) are # visible bidirectionally, so prefill index rows widen from - # window_size to window_size + max_image_tokens. Text-only models keep - # max_image_tokens == 0 and take the original code paths everywhere. + # window_size to window_size + max_image_tokens. The V4 config sets + # mm_prefix_clamp_sliding_window exactly for these in-kernel-widened + # ranges; V4.1 image tokens use the plain causal window, and text-only + # models keep max_image_tokens == 0 everywhere. self.max_image_tokens = ( getattr(hf_config, "vision_max_n_token", 0) - if getattr(hf_config, "vision_n_layers", 0) > 0 + if getattr(hf_config, "mm_prefix_clamp_sliding_window", False) else 0 ) self.prefill_index_width = self.window_size + self.max_image_tokens From fc8132a5e523294ea69decb5e7cbf50e9d3e135e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 16 Sep 2026 20:09:05 +0200 Subject: [PATCH 339/774] [Bugfix] Fix np.float64 leaking into the KV transfer metrics log (#57068) Signed-off-by: NickLucche --- .../kv_connector/unit/test_hf3fs_connector.py | 4 +++ .../kv_connector/unit/test_mooncake_stats.py | 5 +++ .../kv_connector/unit/test_multi_connector.py | 35 ++++++++++++++++++- .../kv_connector/unit/test_nixl_connector.py | 3 ++ .../kv_connector/v1/hf3fs/hf3fs_connector.py | 12 ++++--- .../kv_connector/v1/mooncake/stats.py | 8 ++--- .../kv_transfer/kv_connector/v1/nixl/stats.py | 10 +++--- 7 files changed, 63 insertions(+), 14 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_hf3fs_connector.py b/tests/v1/kv_connector/unit/test_hf3fs_connector.py index 8a8f046835fc..de4d53279877 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_connector.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_connector.py @@ -8,6 +8,7 @@ import os from unittest.mock import MagicMock +import numpy as np import pytest import torch @@ -218,6 +219,9 @@ def test_reduce_with_data(self): assert result["Num save task success"] == pytest.approx(2.0, rel=0.01) assert result["Num save task failed"] == pytest.approx(0.0, rel=0.01) assert result["Avg save duration (ms)"] == pytest.approx(2000.0, rel=0.01) + # Reduced values must be plain Python scalars so CLI logging renders + # them without numpy reprs (eg np.float64(...)). + assert all(not isinstance(v, np.generic) for v in result.values()) def test_clone_and_reset(self, hf3fs_stats): """clone_and_reset() returns a copy with data and resets the original.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_stats.py b/tests/v1/kv_connector/unit/test_mooncake_stats.py index a20fcb505330..d52ab5f201a3 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_stats.py +++ b/tests/v1/kv_connector/unit/test_mooncake_stats.py @@ -4,6 +4,8 @@ import threading from unittest.mock import MagicMock +import numpy as np + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import ( MooncakeConnector, MooncakeConnectorWorker, @@ -40,6 +42,9 @@ def test_record_transfer_and_reduce(): assert reduced["Num failed transfers"] == 0 assert reduced["Num failed recvs"] == 0 assert reduced["Num KV expired reqs"] == 0 + # Reduced values must be plain Python scalars so CLI logging renders + # them without numpy reprs (eg np.float64(...)). + assert all(not isinstance(v, np.generic) for v in reduced.values()) def test_record_failures_keeps_stats_non_empty(): diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index fcda6232c6b2..f6d8ed97d88f 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -20,7 +20,10 @@ SupportsHMA, supports_hma, ) -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorLogging, + KVConnectorStats, +) from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiConnector, MultiKVConnectorPromMetrics, @@ -817,6 +820,36 @@ def test_reduce(self): assert "Num successful transfers" in reduced["NixlConnector"] assert reduced["NixlConnector"]["Num successful transfers"] == 2 + def test_log_renders_plain_scalars(self): + """Reduced stats must render as plain numbers in the CLI log, not + numpy reprs (eg np.float64(2.338)).""" + stats = MultiKVConnectorStats( + data={ + "NixlConnector": NixlKVConnectorStats( + data={ + "transfer_duration": [1.0, 2.0], + "post_duration": [0.1, 0.2], + "bytes_transferred": [1024, 2048], + "num_descriptors": [10, 20], + "num_failed_transfers": [], + "num_failed_notifications": [], + } + ) + } + ) + + kv_logging = KVConnectorLogging(kv_transfer_config=None) + kv_logging.transfer_stats_accumulator = stats + log_records: list[tuple] = [] + kv_logging.log(log_fn=lambda *args: log_records.append(args)) + + assert len(log_records) == 1 + msg = log_records[0][1] + assert "np.float" not in msg + assert "'Avg xfer time (ms)': 1500.0" in msg + # The accumulator is reset after logging. + assert kv_logging.transfer_stats_accumulator is None + def test_reset(self): """Test that reset() resets all nested connector stats.""" stats = MultiKVConnectorStats( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index c9d2dcf5f176..2185ba62d27a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1676,6 +1676,9 @@ def test_kv_connector_stats_aggregation(): assert cli_stats["Avg xfer time (ms)"] == 1500.0 assert cli_stats["Avg post time (ms)"] == 1500.0 assert cli_stats["Avg number of descriptors"] == 1.5 + # Reduced values must be plain Python scalars so CLI logging renders + # them without numpy reprs (eg np.float64(...)). + assert all(not isinstance(v, np.generic) for v in cli_stats.values()) def test_multi_kv_connector_stats_aggregation(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py index b5df658d54d8..1b37a8844019 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py @@ -1075,10 +1075,14 @@ def reduce(self) -> dict[str, int | float]: "Num save task failed": num_failed_save, "Num load task success": num_success_load, "Num load task failed": num_failed_load, - "Avg save duration (ms)": round(save_duration.mean() * 1e3, 3), - "P90 save duration (ms)": round(np.percentile(save_duration, 90) * 1e3, 3), - "Avg load duration (ms)": round(load_duration.mean() * 1e3, 3), - "P90 load duration (ms)": round(np.percentile(load_duration, 90) * 1e3, 3), + "Avg save duration (ms)": round(save_duration.mean().item() * 1e3, 3), + "P90 save duration (ms)": round( + np.percentile(save_duration, 90).item() * 1e3, 3 + ), + "Avg load duration (ms)": round(load_duration.mean().item() * 1e3, 3), + "P90 load duration (ms)": round( + np.percentile(load_duration, 90).item() * 1e3, 3 + ), } def is_empty(self) -> bool: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py index d177f55cc720..a9dbd6bdd4ec 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py @@ -126,16 +126,16 @@ def reduce(self) -> dict[str, int | float]: avg_mb = total_mb / n total_time_seconds = xfer_time.sum() throughput_mb_s = ( - total_mb / total_time_seconds if total_time_seconds > 0 else 0.0 + (total_mb / total_time_seconds).item() if total_time_seconds > 0 else 0.0 ) return { "Num successful transfers": n, - "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3), + "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3), "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), - "Avg MB per transfer": round(avg_mb, 3), + "Avg MB per transfer": round(avg_mb.item(), 3), "Throughput (MB/s)": round(throughput_mb_s, 3), - "Avg number of descriptors": round(descs.mean(), 1), + "Avg number of descriptors": round(descs.mean().item(), 1), "Num failed transfers": num_failed_transfers, "Num failed recvs": num_failed_recvs, "Num KV expired reqs": num_kv_expired_reqs, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index e66046c0ec7b..2cff3f843b3f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -115,13 +115,13 @@ def reduce(self) -> dict[str, int | float]: return { "Num successful transfers": n, - "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3), + "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3), "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), - "Avg post time (ms)": round(post_time.mean() * 1e3, 3), + "Avg post time (ms)": round(post_time.mean().item() * 1e3, 3), "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3), - "Avg MB per transfer": round(avg_mb, 3), - "Throughput (MB/s)": round(throughput_mb_s, 3), - "Avg number of descriptors": round(descs.mean(), 1), + "Avg MB per transfer": round(avg_mb.item(), 3), + "Throughput (MB/s)": round(throughput_mb_s.item(), 3), + "Avg number of descriptors": round(descs.mean().item(), 1), } @property From 9639cbde04f58c185ea6e5122dbd2039908abb7a Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:28:08 -0700 Subject: [PATCH 340/774] [Rust Frontend] Expose local DP size in gRPC Control metadata (#57116) Signed-off-by: Alec Flowers Co-authored-by: Codex --- rust/proto/control.proto | 1 + rust/src/server/src/grpc/control.rs | 1 + rust/src/server/src/grpc/tests.rs | 169 +++++++++++++++------------- 3 files changed, 94 insertions(+), 77 deletions(-) diff --git a/rust/proto/control.proto b/rust/proto/control.proto index d2816f68d347..ef3218d35c91 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -68,6 +68,7 @@ message ParallelismInfo { uint32 data_parallel_rank = 4; uint32 decode_context_parallel_size = 5; uint64 world_size = 6; + uint32 data_parallel_size_local = 7; } message GetModelInfoRequest {} diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index de9c5ae65fb9..680ebdfbdaf2 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -50,6 +50,7 @@ impl ControlServiceImpl { data_parallel_rank: ready.data_parallel_rank, decode_context_parallel_size: ready.decode_context_parallel_size, world_size: ready.world_size, + data_parallel_size_local: self.client().engine_count() as u32, } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index a21e013323d4..f4b35c9df4a8 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -31,7 +31,7 @@ use vllm_chat::{ }; use vllm_engine_core_client::mock_engine::{ DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, - default_ready_response, + MockEngineConfig, connect_to_bootstrapped_frontend, default_ready_response, }; use vllm_engine_core_client::protocol::decode_value; use vllm_engine_core_client::protocol::dtype::TensorDtype; @@ -1948,6 +1948,7 @@ async fn control_reports_server_and_model_info() { assert_eq!(parallelism.tensor_parallel_size, 1); assert_eq!(parallelism.pipeline_parallel_size, 1); assert_eq!(parallelism.data_parallel_size, 1); + assert_eq!(parallelism.data_parallel_size_local, 1); assert_eq!(parallelism.data_parallel_rank, 0); assert_eq!(parallelism.decode_context_parallel_size, 1); let rl = server.rl_capabilities.expect("RL capabilities"); @@ -2191,84 +2192,98 @@ async fn control_forwards_weight_update_without_pause_guard() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn control_aggregates_multi_engine_capacity() { - let ipc = IpcNamespace::new().expect("create ipc namespace"); - let handshake_address = ipc.handshake_endpoint(); - - let mut ready_0 = default_ready_response(); - ready_0.max_model_len = 8_192; - ready_0.num_gpu_blocks = 10; - ready_0.effective_data_parallel_size = 2; - ready_0.tensor_parallel_size = 2; - ready_0.pipeline_parallel_size = 3; - ready_0.world_size = 12; - ready_0.weight_transfer_backend = Some("nccl".to_string()); - ready_0.enable_sleep_mode = true; - ready_0.supports_draft_weight_updates = true; - ready_0.effective_attention_block_size = Some(64); - - let mut ready_1 = default_ready_response(); - ready_1.max_model_len = 4_096; - ready_1.num_gpu_blocks = 20; - ready_1.effective_data_parallel_size = 2; - ready_1.tensor_parallel_size = 2; - ready_1.pipeline_parallel_size = 3; - ready_1.world_size = 12; - ready_1.data_parallel_rank = 1; - - let engine_tasks = [ready_0, ready_1].map(|ready| { - let engine_id = EngineId::from_engine_index( - ready.data_parallel_rank.try_into().expect("test rank fits engine identity"), + // Local capacity and ownership must not expand to the deployment-wide DP size. + for (start_rank, global_size) in [(0, 2), (0, 8), (4, 8)] { + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let input_address = ipc.input_endpoint(); + let output_address = ipc.output_endpoint(); + + let mut ready_0 = default_ready_response(); + ready_0.max_model_len = 8_192; + ready_0.num_gpu_blocks = 10; + ready_0.effective_data_parallel_size = global_size as u64; + ready_0.data_parallel_rank = start_rank; + ready_0.tensor_parallel_size = 2; + ready_0.pipeline_parallel_size = 3; + ready_0.world_size = 12; + ready_0.weight_transfer_backend = Some("nccl".to_string()); + ready_0.enable_sleep_mode = true; + ready_0.supports_draft_weight_updates = true; + ready_0.effective_attention_block_size = Some(64); + + let mut ready_1 = default_ready_response(); + ready_1.max_model_len = 4_096; + ready_1.num_gpu_blocks = 20; + ready_1.effective_data_parallel_size = global_size as u64; + ready_1.tensor_parallel_size = 2; + ready_1.pipeline_parallel_size = 3; + ready_1.world_size = 12; + ready_1.data_parallel_rank = start_rank + 1; + + let client_config = EngineCoreClientConfig { + transport_mode: TransportMode::Bootstrapped { + input_address: input_address.clone(), + output_address: output_address.clone(), + engine_start_index: start_rank, + engine_count: 2, + data_parallel_size: global_size, + ready_timeout: Duration::from_secs(2), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + }; + let client_task = tokio::spawn(EngineCoreClient::connect(client_config)); + let mut engine_sockets = Vec::new(); + for ready in [ready_0, ready_1] { + let engine_id = EngineId::from_engine_index( + ready.data_parallel_rank.try_into().expect("test rank fits engine identity"), + ); + engine_sockets.push( + connect_to_bootstrapped_frontend( + &input_address, + &output_address, + engine_id, + MockEngineConfig { + ready_response: ready, + ..Default::default() + }, + ) + .await + .expect("connect mock engine"), + ); + } + let client = client_task.await.unwrap().expect("connect multi-engine client"); + let chat = ChatLlm::from_shared_backend( + Llm::new(client), + Arc::new(FakeTextBackend) as Arc, ); - MockEngineTask::new(spawn_mock_engine_task_with_ready( - handshake_address.clone(), - engine_id, - ready, - |_, _| boxed_test_future(async {}), - )) - }); + let state = AppState::new(vec!["test-model".to_string()], chat); + let service = ControlServiceImpl::new(Arc::new(state)); - let client = EngineCoreClient::connect(EngineCoreClientConfig { - transport_mode: TransportMode::HandshakeOwner { - handshake_address, - advertised_host: "127.0.0.1".to_string(), - engine_count: 2, - ready_timeout: Duration::from_secs(2), - local_input_address: Some(ipc.input_endpoint()), - local_output_address: Some(ipc.output_endpoint()), - }, - coordinator_mode: None, - model_name: "test-model".to_string(), - client_index: 0, - }) - .await - .expect("connect multi-engine client"); - let chat = ChatLlm::from_shared_backend( - Llm::new(client), - Arc::new(FakeTextBackend) as Arc, - ); - let state = AppState::new(vec!["test-model".to_string()], chat); - let service = ControlServiceImpl::new(Arc::new(state)); - - let server = pb::control_server::Control::get_server_info( - &service, - tonic::Request::new(pb::GetServerInfoRequest {}), - ) - .await - .expect("get server info") - .into_inner(); - assert_eq!(server.max_model_len, 4_096); - assert_eq!(server.total_kv_blocks, 30); - assert!(server.effective_attention_block_size.is_none()); - let rl = server.rl_capabilities.expect("RL capabilities"); - assert!(!rl.weight_transfer_enabled); - assert!(rl.weight_transfer_backend.is_empty()); - assert!(!rl.sleep_mode_enabled); - assert!(!rl.draft_weight_updates_enabled); - let parallelism = server.parallelism.unwrap(); - assert_eq!(parallelism.data_parallel_size, 2); - assert_eq!(parallelism.world_size, 12); - - drop(engine_tasks); + let server = pb::control_server::Control::get_server_info( + &service, + tonic::Request::new(pb::GetServerInfoRequest {}), + ) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.max_model_len, 4_096); + assert_eq!(server.total_kv_blocks, 30); + assert!(server.effective_attention_block_size.is_none()); + let rl = server.rl_capabilities.expect("RL capabilities"); + assert!(!rl.weight_transfer_enabled); + assert!(rl.weight_transfer_backend.is_empty()); + assert!(!rl.sleep_mode_enabled); + assert!(!rl.draft_weight_updates_enabled); + let parallelism = server.parallelism.unwrap(); + assert_eq!(parallelism.data_parallel_size, global_size as u32); + assert_eq!(parallelism.data_parallel_rank, start_rank); + assert_eq!(parallelism.data_parallel_size_local, 2); + assert_eq!(parallelism.world_size, 12); + + drop(engine_sockets); + } } #[test] From 6b6bf6b3c35ac72e4876a6fe0925608abcf8316f Mon Sep 17 00:00:00 2001 From: jxj Date: Thu, 17 Sep 2026 02:44:48 +0800 Subject: [PATCH 341/774] [Bugfix] Remove unsupported comma-separated detailed trace values (#55702) Signed-off-by: git-jxj <65210887+git-jxj@users.noreply.github.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/config/observability.py | 13 +------------ vllm/engine/arg_utils.py | 8 -------- 2 files changed, 1 insertion(+), 20 deletions(-) diff --git a/vllm/config/observability.py b/vllm/config/observability.py index e24668b1a6ae..a39962a075ac 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from functools import cached_property -from typing import Any, Literal, cast +from typing import Any, Literal from packaging.version import parse from pydantic import Field, field_validator, model_validator @@ -150,17 +150,6 @@ def _validate_otlp_traces_endpoint(cls, value: str | None) -> str | None: ) return value - @field_validator("collect_detailed_traces") - @classmethod - def _validate_collect_detailed_traces( - cls, value: list[DetailedTraceModules] | None - ) -> list[DetailedTraceModules] | None: - """Handle the legacy case where users might provide a comma-separated - string instead of a list of strings.""" - if value is not None and len(value) == 1 and "," in value[0]: - value = cast(list[DetailedTraceModules], value[0].split(",")) - return value - @model_validator(mode="after") def _validate_tracing_config(self): if self.collect_detailed_traces and not self.otlp_traces_endpoint: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index e38292c0afe1..ccb30a2597ff 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -10,7 +10,6 @@ import sys from collections.abc import Callable from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass -from itertools import permutations from types import UnionType from typing import ( TYPE_CHECKING, @@ -1523,13 +1522,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: observability_group.add_argument( "--otlp-traces-endpoint", **observability_kwargs["otlp_traces_endpoint"] ) - # TODO: generalise this special case - choices = observability_kwargs["collect_detailed_traces"]["choices"] - metavar = f"{{{','.join(str(c) for c in choices)}}}" - observability_kwargs["collect_detailed_traces"]["metavar"] = metavar - observability_kwargs["collect_detailed_traces"]["choices"] += [ - ",".join(p) for p in permutations(get_args(DetailedTraceModules), r=2) - ] observability_group.add_argument( "--collect-detailed-traces", **observability_kwargs["collect_detailed_traces"], From 403f182c02778a447cf72cfd08b72537eeba72f8 Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:47:27 +0900 Subject: [PATCH 342/774] [MM] Keep raw pixels through dp-sharded ViT path (#56872) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> --- vllm/model_executor/models/qwen2_5_vl.py | 4 ++-- vllm/model_executor/models/qwen2_vl.py | 4 ++-- vllm/model_executor/models/vision.py | 5 +++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index a3d0ecfcab97..67f72dfd4796 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -1441,7 +1441,7 @@ def _process_image_input( if self.use_data_parallel: return run_dp_sharded_mrope_vision_model( self.visual, - pixel_values.type(self.visual.dtype), + pixel_values, grid_thw_list, rope_type="rope_3d", ) @@ -1501,7 +1501,7 @@ def _process_video_input( if self.use_data_parallel: return run_dp_sharded_mrope_vision_model( self.visual, - pixel_values_videos.type(self.visual.dtype), + pixel_values_videos, grid_thw_list, rope_type="rope_3d", ) diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 65be51e8b2ad..d5e9d0a28bef 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -1379,7 +1379,7 @@ def _process_image_input( if self.use_data_parallel: return run_dp_sharded_mrope_vision_model( self.visual, - pixel_values.type(self.visual.dtype), + pixel_values, grid_thw.tolist(), rope_type="rope_3d", ) @@ -1404,7 +1404,7 @@ def _process_video_input( if self.use_data_parallel: return run_dp_sharded_mrope_vision_model( self.visual, - pixel_values_videos.type(self.visual.dtype), + pixel_values_videos, grid_thw.tolist(), rope_type="rope_3d", ) diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index 0c7cfb598e4f..f74a66cdfa3b 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -481,6 +481,7 @@ def run_dp_sharded_mrope_vision_model( # to work max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local] + embed_dtype = next(vision_model.parameters()).dtype # Run the vision model on the local pixel_values_local if rope_type == "rope_2d": @@ -495,7 +496,7 @@ def run_dp_sharded_mrope_vision_model( image_embeds_local = torch.empty( (0, embed_dim_reduction_factor, out_dim), device=pixel_values.device, - dtype=pixel_values.dtype, + dtype=embed_dtype, ) else: if pixel_values_local.shape[0] > 0: @@ -505,7 +506,7 @@ def run_dp_sharded_mrope_vision_model( image_embeds_local = torch.empty( (0, vision_model.out_hidden_size), device=pixel_values.device, - dtype=pixel_values.dtype, + dtype=embed_dtype, ) # Pad the output based on max_len_per_rank From bc0f47cd03d6ae99f9f217f096684423bb4ebc2b Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Thu, 17 Sep 2026 02:58:23 +0800 Subject: [PATCH 343/774] [ROCm][Bugfix] Revert #56433 + #51692 to fix accuracy breakdown for DeepSeek-V4 (#57132) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- tests/models/test_deepseek_v4_vl_rocm.py | 125 ------------- vllm/_aiter_ops.py | 32 +--- .../model_executor/kernels/linear/__init__.py | 4 - .../kernels/linear/scaled_mm/aiter.py | 170 ------------------ .../layers/attention/mla_attention.py | 7 +- vllm/model_executor/layers/linear.py | 6 +- vllm/models/deepseek_v4/amd/model.py | 17 +- vllm/models/deepseek_v4/amd/rocm.py | 75 ++------ 8 files changed, 34 insertions(+), 402 deletions(-) diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 750927972ef3..35d5227a6afc 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -15,131 +15,6 @@ ) -def _preshuffled_fp8_linear( - holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128) -) -> nn.Module: - pytest.importorskip("aiter") - from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( - AiterPreshuffledFp8BlockScaledMMKernel, - ) - from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( - FP8ScaledMMLinearLayerConfig, - ) - from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8Dynamic128Sym, - kFp8Static128BlockSym, - ) - - config = FP8ScaledMMLinearLayerConfig( - weight_quant_key=kFp8Static128BlockSym, - activation_quant_key=kFp8Dynamic128Sym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=weight_shape, - ) - layer = nn.Module() - layer.weight = nn.Parameter( - torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()), - requires_grad=False, - ) - layer.register_parameter( - "weight_scale" if holder == "scheme" else "weight_scale_inv", - nn.Parameter( - torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"), - requires_grad=False, - ), - ) - setattr( - layer, - holder, - SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)), - ) - return layer - - -@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) -def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 - - layer = _preshuffled_fp8_linear(holder) - layer.is_bmm = True - original = layer.weight.detach().clone() - kernel = getattr(layer, holder).fp8_linear - kernel.process_weights_after_loading(layer) - - actual = _get_cached_wo_a_bf16(layer, 2, 128, 128) - torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128)) - - -@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) -def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend( - holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config -) -> None: - from vllm._aiter_ops import rocm_aiter_ops - from vllm.models.deepseek_v4.amd import model as rocm_model - - gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder) - original_gate_up = gate_up.weight.detach().clone() - original_down = down.weight.detach().clone() - monkeypatch.setattr( - rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up - ) - monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down) - monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) - model = rocm_model.DeepseekV4MLP(128, 128, "silu") - for linear in (gate_up, down): - getattr(linear, holder).fp8_linear.process_weights_after_loading(linear) - model.prepare_gateup_preshuffle() - - for linear, original in ((gate_up, original_gate_up), (down, original_down)): - expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16)) - torch.testing.assert_close( - linear.weight.view(torch.uint8), expected.view(torch.uint8) - ) - - -def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout( - monkeypatch: pytest.MonkeyPatch, default_vllm_config -) -> None: - from vllm._aiter_ops import rocm_aiter_ops - from vllm.models.deepseek_v4.amd.rocm import ( - DeepseekV4ROCMAiterMLAAttention, - apply_pre_quantized_block_scaled_mm, - ) - - if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256): - pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM") - linear = _preshuffled_fp8_linear(weight_shape=(1024, 256)) - original = linear.weight.detach().clone() - linear.quant_method.fp8_linear.process_weights_after_loading(linear) - attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention) - nn.Module.__init__(attention) - attention.wq_b = linear - attention.indexer = None - attention.q_lora_rank, attention.head_dim = 256, 128 - attention.eps = 1e-5 - attention.q_norm = SimpleNamespace( - weight=torch.ones(256, dtype=torch.bfloat16, device="cuda") - ) - attention.kv_norm = SimpleNamespace( - weight=torch.ones(128, dtype=torch.bfloat16, device="cuda") - ) - monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True) - # Multiple rows and groups with different ranges expose byte-order errors. - inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda") - inputs[:, :128] *= 4 - quantized, scales, _ = attention._split_qkv_and_norm(inputs) - assert scales is not None - output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales) - q = inputs[:, :256].float() - normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps) - reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1) - reference_scales /= torch.finfo(current_platform.fp8_dtype()).max - dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1) - expected = dequantized @ original.float().T - torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01) - - def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None: from vllm.config import CacheConfig from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index da4bc310d7e4..101e1dd10a6d 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -175,26 +175,18 @@ def is_aiter_found_and_supported_on_rdna4() -> bool: @functools.cache def _load_gemm_tuned_configs( - csv_path: str, - filters: tuple[tuple[str, object], ...], - key_cols: tuple[str, ...] = ("N", "K", "M"), -) -> set[tuple[int, ...]]: + q_dtype_w: torch.dtype, csv_path: str +) -> set[tuple[int, int, int]]: try: df = pd.read_csv(csv_path).drop_duplicates() - for col, val in filters: - if col not in df.columns: - continue - if isinstance(val, int): - df = df[df[col].astype(int) == val] - else: - df = df[df[col].astype(str) == str(val)] - return set(zip(*(df[c].astype(int) for c in key_cols))) + df = df[df["q_dtype_w"] == str(q_dtype_w)] + return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int))) except Exception: return set() def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool: - configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),)) + configs = _load_gemm_tuned_configs(q_dtype_w, csv_path) l_m = ( [1, 2, 4] + list(range(8, 513, 8)) @@ -3329,20 +3321,6 @@ def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE return _check_kernel_tuned(N, K, q_dtype_w, csv_path) - @staticmethod - def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool: - """Whether (N, K) has a tuned aiter blockscale bpreshuffle config.""" - if not current_platform.is_rocm(): - return False - import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops - - csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE - gfx = aiter_gemm_a8w8_ops.get_gfx() - cu_num = aiter_gemm_a8w8_ops.get_cu_num() - return (n, k) in _load_gemm_tuned_configs( - csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K") - ) - @staticmethod def shuffle_weight( tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 26ad557c67b7..3c9c616d0c16 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -177,7 +177,6 @@ AiterHipbMMPerTokenFp8ScaledMMLinearKernel, AiterInt8ScaledMMLinearKernel, AiterPerTokenFp8ScaledMMLinearKernel, - AiterPreshuffledFp8BlockScaledMMKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( @@ -322,7 +321,6 @@ def _get_linear_backend(*, quantization: str) -> str: "aiter": { AiterInt8ScaledMMLinearKernel, AiterFp8BlockScaledMMKernel, - AiterPreshuffledFp8BlockScaledMMKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, AiterMxfp4LinearKernel, @@ -466,7 +464,6 @@ def _resolve_backend_kernels( BlockWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ - AiterPreshuffledFp8BlockScaledMMKernel, AiterFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], @@ -1250,7 +1247,6 @@ def register_linear_kernel( "ScaledMMLinearLayerConfig", "AiterHipbMMPerTokenFp8ScaledMMLinearKernel", "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", - "AiterPreshuffledFp8BlockScaledMMKernel", "AiterPerTokenFp8ScaledMMLinearKernel", "NvFp4LinearKernel", "NvFp4LinearLayerConfig", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index a57e7510de01..03e9be51cfa8 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar - import torch from vllm import _custom_ops as ops @@ -474,171 +472,3 @@ def apply_block_scaled_mm( return gemm_a8w8_blockscale_op( A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype ) - - -class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): - """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight.""" - - # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major. - wants_transposed_act_scale: ClassVar[bool] = True - - # process_weights_after_loading shuffles layer.weight to layout (16, 16). - preshuffles_weight: ClassVar[bool] = True - - @classmethod - def is_supported( - cls, compute_capability: int | None = None - ) -> tuple[bool, str | None]: - return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported( - compute_capability - ) - - @classmethod - def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): - can_implement_base, reason = super().can_implement(config) - if not can_implement_base: - return can_implement_base, reason - - act_quant_desc = config.activation_quant_key.scale - if act_quant_desc.group_shape != GroupShape(1, 128): - return ( - False, - ( - "Supports only dynamic per token group activation " - "quantization with group_shape=(1,128)." - ), - ) - - # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8 - # 2D weight with N and K divisible by 128. Require a tuned config as the - # fallback can trigger faults or numerical errors. - if not rocm_aiter_ops.is_linear_fp8_enabled(): - return ( - False, - ( - "requires setting `VLLM_ROCM_USE_AITER=1` " - "and `VLLM_ROCM_USE_AITER_LINEAR=1`." - ), - ) - - n, k = config.weight_shape - if not (n % 128 == 0 and k % 128 == 0): - return ( - False, - ( - f"requires N and K dimensions divisible by 128, received " - f"N={n} and K={k}." - ), - ) - - if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k): - return ( - False, - ( - f"requires a tuned aiter blockscale bpreshuffle config for " - f"N={n} and K={k}." - ), - ) - - return True, None - - @staticmethod - def _reads_weight_directly(layer: torch.nn.Module) -> bool: - """True when something other than apply_weights consumes layer.weight. - - Such a weight must stay in the plain layout: ``is_bmm`` marks a stack - of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj. - Both are stamped after construction, so can_implement cannot see them. - """ - return bool( - getattr(layer, "is_bmm", False) - or getattr(layer, "skip_weight_relayout", False) - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - super().process_weights_after_loading(layer) - - if self._reads_weight_directly(layer): - return - - params = FP8BlockParams.from_layer(layer) - if params.weight_scale_inv is not None: - ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV - else: - ws, attr = params.weight_scale, params.WEIGHT_SCALE - if ws is not None and ws.dtype == torch.float8_e8m0fnu: - replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) - - weight = params.weight - # runtime safety net - assert ( - weight.dim() == 2 - and weight.dtype == current_platform.fp8_dtype() - and weight.shape[0] % 128 == 0 - and weight.shape[1] % 128 == 0 - ), ( - "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight " - "with N and K divisible by 128." - ) - - shuffled_weight = rocm_aiter_ops.shuffle_weight( - weight.contiguous(), layout=(16, 16) - ) - replace_parameter( - layer, - params.WEIGHT, - torch.nn.Parameter(shuffled_weight.data, requires_grad=False), - ) - - def apply_weights( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - params = self._get_layer_params(layer) - Bs = ( - params.weight_scale - if params.weight_scale_inv is None - else params.weight_scale_inv - ) - - # Left unshuffled above; kv_b_proj still reaches here from MLA's - # prefill-context path. - plain = self._reads_weight_directly(layer) - - x_2d = x.view(-1, x.shape[-1]) - A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain) - if plain: - output = rocm_aiter_ops.gemm_a8w8_blockscale( - A, - params.weight, - As, - Bs, - list(self.weight_group_shape), - output_dtype=self.config.out_dtype, - ) - else: - output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, params.weight, As, Bs, output_dtype=self.config.out_dtype - ) - if bias is not None: - output = output + bias - return output.view(*x.shape[:-1], params.weight.shape[0]) - - def apply_block_scaled_mm( - self, - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - ) -> torch.Tensor: - """Block-scaled GEMM for callers that pre-quantize their activations. - - ``As`` must be column-major; a row-major one will not raise, it just - returns wrong numbers for M > 1. - """ - return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, B, As, Bs, output_dtype=self.config.out_dtype - ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4cbc380bb151..26547c73533d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -582,9 +582,10 @@ def __init__( ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) - # MLA reads this weight directly to build W_UK/W_UV, so a backend must - # not relayout it at load time. - kv_b_proj.skip_weight_relayout = True + # AMX reads kv_b_proj's weight directly and never calls it live; the + # reference CPU MLA backend calls it but isn't perf-critical. Skip + # the packed-kernel dispatch either way. + kv_b_proj._cpu_skip_gemm_dispatch = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 9d2e31cfe177..65cda325a221 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -205,9 +205,9 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): - # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical, - # so skip packing and use a plain fallback. - if getattr(layer, "skip_weight_relayout", False): + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not + # perf-critical, so skip packing and use a plain fallback. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): layer.cpu_linear = torch.nn.functional.linear return diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index c1a6903937b6..b23e728b46bb 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -71,10 +71,7 @@ make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.rocm import ( - DeepseekV4ROCMAiterMLAAttention, - weight_already_preshuffled, -) +from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors @@ -152,13 +149,11 @@ def prepare_gateup_preshuffle(self) -> None: return if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Skip if the linear's kernel already shuffled it. - if not weight_already_preshuffled(self.gate_up_proj): - replace_parameter( - self.gate_up_proj, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) self._gateup_scale = ws def forward(self, x): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3cc9827c1fbe..f8ac39cd6241 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -62,21 +62,6 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: return indptr -def weight_already_preshuffled(linear: torch.nn.Module) -> bool: - """True when the linear's kernel already B-preshuffled ``weight``. - - The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped - for those, since shuffle_weight is a permutation rather than an involution. - """ - return any( - getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False) - for method in ( - getattr(linear, "quant_method", None), - getattr(linear, "scheme", None), - ) - ) - - def apply_pre_quantized_block_scaled_mm( linear: torch.nn.Module, x_fp8: torch.Tensor, @@ -792,13 +777,12 @@ def _prep(linear) -> torch.Tensor | None: return None if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Skip if the linear's kernel already shuffled it. - if not weight_already_preshuffled(linear): - replace_parameter( - linear, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Shuffle the weight in place (single weight, no unshuffled copy). + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) return ws self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) @@ -896,26 +880,16 @@ def _run_parallel_input_projections( return qr_kv, kv_score, indexer_kv_score, indexer_weights @functools.cached_property - def _wq_b_act_scale_transpose(self) -> bool | None: - """Activation-scale byte order every wq_b consumer agrees on, else None. - - None means "do not take the fused norm+quant path". Otherwise the value - is the ``transpose_scale`` the producer must pass so that the scale it - writes matches what the selected GEMM reads. + def _wq_b_uses_aiter_block_scaled(self) -> bool: + """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. Cached: the linear kernels and the aiter env gates are fixed once the model is built, so this is evaluated at the first forward only. - Two conditions, both necessary: - - * The fused norm+quant path is only valid if the quant and GEMM it - replaces are exactly the aiter ones; otherwise fall back to the - shared path. - * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have - different (N, K) and so can resolve to kernels wanting opposite byte - orders. A single producer cannot serve both, so refuse the fast path - when they disagree rather than guessing. + The fused norm+quant path is only valid if the quant and GEMM it + replaces are exactly the aiter ones; otherwise fall back to the + shared path. """ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.kernels.linear.scaled_mm import ( @@ -923,31 +897,16 @@ def _wq_b_act_scale_transpose(self) -> bool | None: ) if not rocm_aiter_ops.is_linear_fp8_enabled(): - return None + return False linears = [self.wq_b] if self.indexer is not None: linears.append(self.indexer.wq_b) - layouts: set[bool] = set() for linear in linears: kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) if not isinstance(kernel, Fp8BlockScaledMMLinearKernel): - return None - layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False))) - if len(layouts) != 1: - logger.warning_once( - "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; " - "disabling the fused q/kv norm+quant path.", - scope="global", - ) - return None - transpose_scale = layouts.pop() - logger.debug_once( - "DeepSeek-V4 wq_b: emitting %s-major activation scales", - "column" if transpose_scale else "row", - scope="global", - ) - return transpose_scale + return False + return True def _split_qkv_and_norm( self, qr_kv: torch.Tensor @@ -964,12 +923,11 @@ def _split_qkv_and_norm( when the aiter linear path is not active. """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) - transpose_scale = self._wq_b_act_scale_transpose if not ( qr.dim() == 2 and qr.shape[0] > 0 and self.q_lora_rank % 128 == 0 - and transpose_scale is not None + and self._wq_b_uses_aiter_block_scaled ): return super()._split_qkv_and_norm(qr_kv) @@ -983,8 +941,7 @@ def _split_qkv_and_norm( kv_weight=self.kv_norm.weight.data, kv_epsilon=self.eps, group_size=128, - # Emit the byte order the wq_b GEMMs read. - transpose_scale=transpose_scale, + transpose_scale=False, ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: From 30b847c1fa5461a590606cbc9b8dae4a490c6cdd Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 16 Sep 2026 12:09:30 -0700 Subject: [PATCH 344/774] [Perf][DSV4.1] Remove MegaMoE padding and shared padding workaround (#57204) Signed-off-by: Woosuk Kwon --- tests/models/test_deepseek_v4_mega_moe.py | 129 +++++++--------------- vllm/models/deepseek_v4/nvidia/model.py | 70 +----------- 2 files changed, 42 insertions(+), 157 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 9a5723d233a1..67029b50a911 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -300,8 +300,10 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): @pytest.mark.parametrize("intermediate_size", [512, 2304]) -def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size): - """Pad gate/up separately and zero every added down-projection column.""" +def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( + monkeypatch, intermediate_size +): + """Keep native widths so V4.1 routed and shared experts can fuse.""" experts = DeepseekV4MegaMoEExperts( SimpleNamespace( scheduler_config=SimpleNamespace(max_num_batched_tokens=4), @@ -333,29 +335,15 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat ) experts.finalize_weights() - padded_size = experts.intermediate_size - assert padded_size == (intermediate_size + 511) // 512 * 512 - for actual, original in zip(experts._transformed_l1_weights, originals[:2]): - expected = ( - original - if actual.dtype == torch.int8 - else (experts._ue8m0_uint8_to_float(original)) - ) - actual = actual.view(expected.dtype).unflatten(1, (2, padded_size)) - assert torch.equal( - actual[:, :, :intermediate_size], - expected.unflatten(1, (2, intermediate_size)), - ) - assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0 - for actual, original in zip(experts._transformed_l2_weights, originals[2:]): + assert experts.intermediate_size == intermediate_size + transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) + for actual, original in zip(transformed, originals): expected = ( original if actual.dtype == torch.int8 - else (experts._ue8m0_uint8_to_float(original)) + else experts._ue8m0_uint8_to_float(original) ) - actual = actual.view(expected.dtype) - assert torch.equal(actual[..., : original.shape[-1]], expected) - assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0 + assert torch.equal(actual.view(expected.dtype), expected) transformed_l1 = experts._transformed_l1_weights experts.finalize_weights() @@ -363,20 +351,19 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat @pytest.mark.parametrize( - "hidden_size,intermediate_size,block_size,fused,mxfp8", + "hidden_size,intermediate_size,block_size,mxfp8", [ - pytest.param(128, 512, 128, True, False, id="aligned-block128"), - pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"), - pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash"), - pytest.param(128, 2304, 128, True, False, id="padded-block128"), - pytest.param(5120, 2304, 32, True, True, id="deepseek-v41-flash-mxfp8"), - pytest.param(128, 2304, 128, False, False, id="padded-packed-scale-fallback"), + pytest.param(128, 512, 128, False, id="aligned-block128"), + pytest.param(128, 512, 128, True, id="aligned-block128-mxfp8"), + pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"), + pytest.param(128, 2304, 128, False, id="native-block128"), + pytest.param(5120, 2304, 32, True, id="deepseek-v41-flash-mxfp8"), ], ) def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( - monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8 + monkeypatch, hidden_size, intermediate_size, block_size, mxfp8 ): - """Shared fusion preserves checkpoint channels and zeros padded channels.""" + """Shared fusion preserves checkpoint weights and scales at native widths.""" class FakeDeepGemm: transformed_dims: list[tuple[int, int]] = [] @@ -476,12 +463,6 @@ def scale_parameter(*shape, dtype=torch.int32): originals = [] for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv - if not fused: - # Already packed scales cannot be padded as checkpoint UE8M0 bytes. - linear.weight_scale_inv = scale_parameter( - linear.weight.shape[0], linear.weight.shape[1] // 128 - ) - continue scale.data.view(torch.uint8).random_(124, 130) block_m, block_k = linear.weight_block_size scale_1x32 = ( @@ -493,70 +474,34 @@ def scale_parameter(*shape, dtype=torch.int32): monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() - original_down_ptr = shared_experts.down_proj.weight.data_ptr() experts.finalize_weights(shared_experts) - assert experts.has_fused_shared_experts is fused - if not fused: - assert experts.intermediate_size == 2560 - assert FakeDeepGemm.transformed_dims == [(3, 3)] - assert experts.num_shared_experts == 0 - assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr - assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr - assert shared_experts.gate_up_proj.weight.shape == ( - 2 * intermediate_size, - hidden_size, - ) - assert shared_experts.down_proj.weight.shape == (hidden_size, intermediate_size) - return - + assert experts.has_fused_shared_experts + assert experts.intermediate_size == intermediate_size assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] - padded_size = experts.intermediate_size assert FakeDeepGemm.scale_inputs[-2:] == [ - (1, 2 * padded_size, hidden_size // 32), - (1, hidden_size, padded_size // 32), + (1, 2 * intermediate_size, hidden_size // 32), + (1, hidden_size, intermediate_size // 32), ] - for actual, original, fill in ( + for transformed, actual_scale, (weight, scale) in zip( ( - experts._transformed_shared_l1_weights[0].view(torch.uint8), - originals[0][0], - 0, + experts._transformed_shared_l1_weights, + experts._transformed_shared_l2_weights, ), - (FakeDeepGemm.scale_values[-2][0], originals[0][1], 1), + FakeDeepGemm.scale_values[-2:], + originals, ): - actual = actual.unflatten(0, (2, padded_size)) - assert torch.equal( - actual[:, :intermediate_size], original.unflatten(0, (2, intermediate_size)) - ) - assert torch.all(actual[:, intermediate_size:] == fill) - for actual, original, fill in ( - ( - experts._transformed_shared_l2_weights[0].view(torch.uint8), - originals[1][0], - 0, - ), - (FakeDeepGemm.scale_values[-1][0], originals[1][1], 1), - ): - assert torch.equal(actual[:, : original.shape[1]], original) - assert torch.all(actual[:, original.shape[1] :] == fill) - if padded_size != intermediate_size: - # Generic linear post-load hooks still receive checkpoint-shaped weights. - for linear, (weight, _) in zip( - (shared_experts.gate_up_proj, shared_experts.down_proj), originals - ): - assert torch.equal(linear.weight.view(torch.uint8), weight) - assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr - assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr - else: - assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr - assert ( - experts._transformed_shared_l1_weights[0].data_ptr() - == shared_experts.gate_up_proj.weight.data_ptr() - ) - assert ( - experts._transformed_shared_l2_weights[0].data_ptr() - == shared_experts.down_proj.weight.data_ptr() - ) + assert torch.equal(transformed[0].view(torch.uint8), weight) + assert torch.equal(actual_scale[0], scale) + assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr + assert ( + experts._transformed_shared_l1_weights[0].data_ptr() + == shared_experts.gate_up_proj.weight.data_ptr() + ) + assert ( + experts._transformed_shared_l2_weights[0].data_ptr() + == shared_experts.down_proj.weight.data_ptr() + ) transformed_l1 = experts._transformed_shared_l1_weights experts.finalize_weights(shared_experts) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 4aba14cd03b1..7fef691bd134 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -204,7 +204,6 @@ def __init__( self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size - self.unpadded_intermediate_size = intermediate_size self.num_shared_experts = num_shared_experts self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -387,17 +386,6 @@ def _finalize_shared_expert_weights( # the generic linear post-load hook replaces the raw checkpoint scales # with its 128x128 DeepGEMM layout. checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8) - unpadded_size = self.unpadded_intermediate_size * self.num_shared_experts - padding = self.intermediate_size * self.num_shared_experts - unpadded_size - pad_weights = ( - padding > 0 - and gate_up_weight.dtype == torch.float8_e4m3fn - and down_weight.dtype == torch.float8_e4m3fn - and gate_up_weight.shape == (2 * unpadded_size, self.hidden_size) - and down_weight.shape == (self.hidden_size, unpadded_size) - and gate_up_scale.dtype in checkpoint_scale_dtypes - and down_scale.dtype in checkpoint_scale_dtypes - ) if ( gate_up_scale.dtype in checkpoint_scale_dtypes and down_scale.dtype in checkpoint_scale_dtypes @@ -408,7 +396,6 @@ def _finalize_shared_expert_weights( gate_up_scale, gate_up_weight.shape[0], gate_up_weight.shape[1], - padding=(padding, 0) if pad_weights else (0, 0), ) down_scale = self._prepare_shared_expert_scale( deep_gemm, @@ -416,27 +403,12 @@ def _finalize_shared_expert_weights( down_scale, down_weight.shape[0], down_weight.shape[1], - padding=(0, padding) if pad_weights else (0, 0), ) if gate_up_scale is None or down_scale is None: self.num_shared_experts = 0 return - if pad_weights: - # Pad gate/up separately so the SwiGLU split stays at the midpoint. - gate_up_weight = ( - torch.nn.functional.pad( - gate_up_weight.view(torch.uint8).unflatten(0, (2, unpadded_size)), - (0, 0, 0, padding), - ) - .flatten(0, 1) - .view(gate_up_weight.dtype) - ) - down_weight = torch.nn.functional.pad( - down_weight.view(torch.uint8), (0, padding) - ).view(down_weight.dtype) - shared_intermediate_size = self.intermediate_size * self.num_shared_experts expected_gate_up_shape = ( 2 * shared_intermediate_size, @@ -477,11 +449,11 @@ def _finalize_shared_expert_weights( # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. # The generic linear post-load hook may still repack the serial scales, # but this shared MLP is never called after native fusion is enabled. - # Padded weights need separate storage: the generic linear post-load - # hooks still expect the original checkpoint weight/scale shapes. - if not pad_weights: - gate_up.weight.data = transformed_l1[0] - self._transformed_shared_l1_weights = transformed_l1 + gate_up.weight.data = transformed_l1[0] + self._transformed_shared_l1_weights = ( + gate_up.weight.data, + transformed_l1[1], + ) self._transformed_shared_l2_weights = transformed_l2 def _prepare_shared_expert_scale( @@ -491,8 +463,6 @@ def _prepare_shared_expert_scale( scale: torch.Tensor, mn: int, k: int, - *, - padding: tuple[int, int] = (0, 0), ) -> torch.Tensor | None: block_size = getattr(linear, "weight_block_size", None) if block_size is None or len(block_size) != 2: @@ -527,19 +497,6 @@ def _prepare_shared_expert_scale( .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32] .contiguous() ) - pad_m, pad_k = padding - if pad_m: - scale_1x32 = torch.nn.functional.pad( - scale_1x32.unflatten(0, (2, mn // 2)), - (0, 0, 0, pad_m), - value=1.0, - ).flatten(0, 1) - mn += 2 * pad_m - if pad_k: - scale_1x32 = torch.nn.functional.pad( - scale_1x32, (0, pad_k // 32), value=1.0 - ) - k += pad_k # The grouped API is used with a singleton dimension to request the # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the # 2D layout required for a shared expert. @@ -558,23 +515,6 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: if self._transformed_l1_weights is None: self._check_runtime_supported() - # MegaMoE's 1x32 activation scales need 16-byte TMA rows, so - # pad each gate/up half and the down projection to a multiple of 512. - padded_size = (self.intermediate_size + 511) // 512 * 512 - padding = padded_size - self.intermediate_size - if padding: - for param in (self.w13_weight, self.w13_weight_scale): - gate_up = param.data.unflatten(1, (2, self.intermediate_size)) - param.data = torch.nn.functional.pad( - gate_up, (0, 0, 0, padding) - ).flatten(1, 2) - self.w2_weight.data = torch.nn.functional.pad( - self.w2_weight.data, (0, padding // 2) - ) - self.w2_weight_scale.data = torch.nn.functional.pad( - self.w2_weight_scale.data, (0, padding // 32) - ) - self.intermediate_size = padded_size w13_scale = deep_gemm.transform_sf_into_required_layout( self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), 2 * self.intermediate_size, From 2bdbbc80804b2199cbf39b75e78cf7269d0383a4 Mon Sep 17 00:00:00 2001 From: danisereb Date: Wed, 16 Sep 2026 22:31:43 +0300 Subject: [PATCH 345/774] [BugFix] Fix is_supported of cutlass FP8 linear (selected and fails on A100) (#55884) Signed-off-by: Daniel Serebrenik --- vllm/model_executor/kernels/linear/scaled_mm/cutlass.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index 9f69ab0c7377..da8c69eed0a2 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( CUTLASS_BLOCK_FP8_SUPPORTED, convert_to_channelwise, + cutlass_fp8_supported, ) from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform @@ -164,8 +165,8 @@ def __init__( def is_supported( cls, compute_capability: int | None = None ) -> tuple[bool, str | None]: - if not current_platform.is_cuda(): - return False, "requires CUDA." + if not cutlass_fp8_supported(): + return False, "CUTLASS FP8 kernels not available" return True, None @classmethod From e0705731be04aac90838f318fcdae15c56a53080 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Wed, 16 Sep 2026 17:00:48 -0400 Subject: [PATCH 346/774] [Refactor] Remove unused interface methods (#57194) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../responses/test_parsable_context_unit.py | 9 - .../openai/test_tool_choice_content_none.py | 3 - tests/entrypoints/unit_tests/test_context.py | 4 +- .../test_structural_tag_registry.py | 2 +- vllm/parser/abstract_parser.py | 201 ++---------------- 5 files changed, 19 insertions(+), 200 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py index 01293f62b222..85bf6f2020a1 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py +++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py @@ -36,9 +36,6 @@ class _NoOpParser(DelegatingParser): def is_reasoning_end(self, input_ids: list[int]) -> bool: return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - return input_ids - def extract_reasoning(self, model_output, request): return None, model_output @@ -63,9 +60,6 @@ class _ReasoningOnlyParser(DelegatingParser): def is_reasoning_end(self, input_ids: list[int]) -> bool: return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - return input_ids - def extract_reasoning(self, model_output, request): if "" in model_output and "" in model_output: start = model_output.index("") + len("") @@ -131,9 +125,6 @@ def __init__(self, tokenizer, *args, **kwargs): def is_reasoning_end(self, input_ids: list[int]) -> bool: return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - return input_ids - def extract_reasoning(self, model_output, request): return None, model_output diff --git a/tests/entrypoints/openai/test_tool_choice_content_none.py b/tests/entrypoints/openai/test_tool_choice_content_none.py index 7db815c45f38..a0a8f38e9f7a 100644 --- a/tests/entrypoints/openai/test_tool_choice_content_none.py +++ b/tests/entrypoints/openai/test_tool_choice_content_none.py @@ -31,9 +31,6 @@ class _DummyDelegatingParser(DelegatingParser): def is_reasoning_end(self, input_ids: list[int]) -> bool: return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - return input_ids - def extract_reasoning(self, model_output: str, request): return None, model_output diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index c92813ea6612..8c3d1467edcb 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -71,8 +71,8 @@ async def generate_mock_outputs( class FakeHarmonyParser(HarmonyParser): def __init__(self): # Skip HarmonyParser initialization and script outputs directly. - self.reasoning_parser = None - self.tool_parser = None + self._reasoning_parser = None + self._tool_parser = None self._chunk_results: list[ChunkResult] = [] self._flush_results: list[list[Segment]] = [] self.processed_chunks: list[list[int]] = [] diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 88e1f3f3832c..06c8fb61295c 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -623,7 +623,7 @@ class TestParser(DelegatingParser): tool_choice="auto", ) parser = TestParser(MagicMock(), tools=sample_tools_strict) - parser.reasoning_parser = MagicMock(adjust_request=lambda request: request) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) parser.adjust_request(request) diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index d51c4455bea2..af448e0951c6 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -6,7 +6,6 @@ from abc import abstractmethod from collections.abc import Sequence from dataclasses import dataclass, field -from functools import cached_property from openai.types.responses import ToolChoiceFunction from pydantic import TypeAdapter, ValidationError @@ -84,23 +83,12 @@ def commit( class Parser: - """Abstract Parser class that unifies ReasoningParser and ToolParser into - a single interface for parsing model output. - - This class provides a unified way to handle both reasoning extraction - (e.g., chain-of-thought content in tags) and tool call extraction - (e.g., function calls in XML/JSON format) from model outputs. - - Subclasses can either: - 1. Override the abstract methods directly for custom parsing logic - 2. Set `reasoning_parser` and `tool_parser` properties to delegate to - existing parser implementations - - Class Attributes: - reasoning_parser_cls: The ReasoningParser class to use (for compatibility - with code that needs the class, not instance). - tool_parser_cls: The ToolParser class to use (for compatibility with - code that needs the class, not instance). + """Parse model output into reasoning, content and tool calls. + + The serving layer holds one ``Parser`` per request and calls only the + members defined here. ``ParserEngine`` implements them over a single + declarative engine; ``DelegatingParser`` composes a legacy + ``ReasoningParser`` / ``ToolParser`` pair. """ # Class-level parser classes for compatibility with existing patterns @@ -150,29 +138,16 @@ def __init__( engine_based=self._engine_based, ) - @cached_property - def vocab(self) -> dict[str, int]: - """Get the vocabulary mapping from tokens to IDs.""" - return self.model_tokenizer.get_vocab() - @property def reasoning_parser(self) -> ReasoningParser | None: """The underlying reasoning parser, if any.""" return self._reasoning_parser - @reasoning_parser.setter - def reasoning_parser(self, parser: ReasoningParser | None) -> None: - self._reasoning_parser = parser - @property def tool_parser(self) -> ToolParser | None: """The underlying tool parser, if any.""" return self._tool_parser - @tool_parser.setter - def tool_parser(self, parser: ToolParser | None) -> None: - self._tool_parser = parser - def _initialize_history_tool_call_cnt( self, request: ChatCompletionRequest | ResponsesRequest, @@ -186,119 +161,6 @@ def _initialize_history_tool_call_cnt( state.history_tool_call_cnt = count_history_tool_calls(request) state.history_tool_call_cnt_initialized = True - # ========== Reasoning Parser Methods ========== - - @abstractmethod - def is_reasoning_end(self, input_ids: list[int]) -> bool: - """Check if the reasoning content ends in the input_ids. - - Used by structured engines like `xgrammar` to check if the - reasoning content ends in the model output. - - Must be a pure function of the input_ids. - - Args: - input_ids: The token IDs of the model output. - - Returns: - True if the reasoning content ends in the input_ids. - - """ - - def is_reasoning_end_streaming( - self, input_ids: list[int], delta_ids: list[int] - ) -> bool: - """Check if the reasoning content ends during a decode step. - - Must be a pure function of the input_ids and delta_ids. - - Args: - input_ids: The entire model output token IDs. - delta_ids: The last few computed tokens at the current decode step. - - Returns: - True if the reasoning content ends in the delta_ids. - - """ - return self.is_reasoning_end(input_ids) - - def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: - """Locate the token that ends reasoning within one window of tokens. - - Args: - token_ids: The tokens to examine. - - Returns: - The offset within ``token_ids`` of the last reasoning token, or - ``len(token_ids)`` when reasoning does not end in this window. Parsers - that cannot answer from a window alone always return ``None``; callers - fall back to :meth:`is_reasoning_end_streaming`. - - """ - return None - - @abstractmethod - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - """Extract content token IDs from the input_ids. - - This extracts the non-reasoning content (e.g., everything after - the tag). - - Args: - input_ids: The token IDs of the model output. - - Returns: - The extracted content token IDs. - - """ - - @abstractmethod - def extract_reasoning( - self, - model_output: str, - request: ChatCompletionRequest | ResponsesRequest, - ) -> tuple[str | None, str | None]: - """Extract reasoning content from a complete model-generated string. - - Used for non-streaming responses where we have the entire model - response available before sending to the client. - - Args: - model_output: The complete model-generated string. - request: The request object used to generate the output. - - Returns: - A tuple of (reasoning, response_content). - - """ - - @abstractmethod - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> DeltaMessage | None: - """Extract reasoning content from a streaming delta message. - - Args: - previous_text: Text from all previous tokens. - current_text: Text including the current delta. - delta_text: The new text in this delta. - previous_token_ids: Token IDs from previous generation. - current_token_ids: All token IDs including current. - delta_token_ids: The new token IDs in this delta. - - Returns: - A DeltaMessage with reasoning and/or content fields, or None. - - """ - - # ========== Tool Parser Methods ========== - def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: @@ -317,48 +179,17 @@ def adjust_request( return request @abstractmethod - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest | ResponsesRequest, - ) -> ExtractedToolCallInformation: - """Extract tool calls from a complete model-generated string. + def is_reasoning_end(self, input_ids: list[int]) -> bool: + """Check if the reasoning content ends in the input_ids. - Used for non-streaming responses. + Called with the rendered prompt to decide whether generation starts + after reasoning. Must be a pure function of the input_ids. Args: - model_output: The complete model-generated string. - request: The request object used to generate the output. - - Returns: - ExtractedToolCallInformation containing the tool calls. - - """ - - @abstractmethod - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest | ResponsesRequest, - ) -> DeltaMessage | None: - """Extract tool calls from a streaming delta message. - - Args: - previous_text: Text from all previous tokens. - current_text: Text including the current delta. - delta_text: The new text in this delta. - previous_token_ids: Token IDs from previous generation. - current_token_ids: All token IDs including current. - delta_token_ids: The new token IDs in this delta. - request: The request object. + input_ids: The token IDs of the model output. Returns: - A DeltaMessage with tool_calls field, or None. + True if the reasoning content ends in the input_ids. """ @@ -764,14 +595,14 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: return False return self._reasoning_parser.is_reasoning_end(input_ids) - def is_reasoning_end_streaming( + def _is_reasoning_end_streaming( self, input_ids: list[int], delta_ids: list[int] ) -> bool: if self._reasoning_parser is None: return False return self._reasoning_parser.is_reasoning_end_streaming(input_ids, delta_ids) - def extract_content_ids(self, input_ids: list[int]) -> list[int]: + def _extract_content_ids(self, input_ids: list[int]) -> list[int]: if self._reasoning_parser is None: return input_ids return self._reasoning_parser.extract_content_ids(input_ids) @@ -885,13 +716,13 @@ def parse_delta( reasoning_parser.has_engine_confirmed_reasoning_end() ) else: - should_transition = self.is_reasoning_end_streaming( + should_transition = self._is_reasoning_end_streaming( current_token_ids, delta_token_ids ) if should_transition: state.reasoning_ended = True reasoning_transitioned = True - current_token_ids = self.extract_content_ids(delta_token_ids) + current_token_ids = self._extract_content_ids(delta_token_ids) # Flush whenever the reasoning parser is engine-based (not only # when _engine_based is True): it buffers the post-marker text # (e.g. the "<" of ""), surfaced via finish_streaming(). From 24a63b3340f07952a3623d0f28331492f3dab44e Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Wed, 16 Sep 2026 14:20:28 -0700 Subject: [PATCH 347/774] [CI][Bugfix] Initialize _transfer_layer_group_ids in region_pull_worker fixture (#57211) Signed-off-by: kimi-no-na-wa Co-authored-by: kimi-no-na-wa Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/v1/kv_connector/unit/test_nixl_connector_hma.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 3f1038ff659c..f9cacabdd455 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -85,6 +85,7 @@ def region_pull_worker(): remote_physical_blocks_per_logical=2, ) worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)} + worker._transfer_layer_group_ids = () worker._mixed_mem_types = True worker.src_xfer_handles_by_block_size = {64: 1} worker._dram_src_handles_by_block_size = {64: 2} From 929aa2a7ed3283913b6ca705ac0e63e4a923c6c4 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 16 Sep 2026 14:49:35 -0700 Subject: [PATCH 348/774] [Build] Bump DeepGEMM pin to a6bbb80 (#57218) --- cmake/external_projects/deepgemm.cmake | 2 +- tools/install_deepgemm.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 2ff69bd8b6b3..9ac0d05414c1 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,7 +32,7 @@ else() # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU # alpha/beta, and SiTU for FP8/FP4 Mega MoE. - set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4") + set(_DEEPGEMM_UPSTREAM_TAG "a6bbb8000161c0dc3a85a0300a905f76898a7913") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index db81bf4be310..a3ec840cd897 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -10,7 +10,7 @@ DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 # and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE # activations -DEEPGEMM_GIT_REF="ad1f1726aa540a76c1d26d6a120effb8de21eaa4" +DEEPGEMM_GIT_REF="a6bbb8000161c0dc3a85a0300a905f76898a7913" WHEEL_DIR="" # Parse command line arguments From 5a6ccc589281b4302791cc301e258adcbcc5fead Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Wed, 16 Sep 2026 17:52:34 -0400 Subject: [PATCH 349/774] [Bugfix] Add Responses cache_write_tokens for API compat and CC parity (#57222) Signed-off-by: Yifan Zong Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../entrypoints/openai/responses/test_serving_responses.py | 7 +++++++ tests/entrypoints/unit_tests/test_context.py | 5 +++++ vllm/entrypoints/openai/responses/context.py | 3 +++ vllm/entrypoints/openai/responses/protocol.py | 1 + vllm/entrypoints/openai/responses/serving.py | 1 + 5 files changed, 17 insertions(+) diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index a445774fccb7..0f361c1c6a50 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -18,6 +18,9 @@ ResponseTextConfig, ResponseTextDeltaEvent, ) +from openai.types.responses import ( + ResponseUsage as OpenAIResponseUsage, +) from openai.types.responses.response_format_text_json_schema_config import ( ResponseFormatTextJSONSchemaConfig, ) @@ -975,6 +978,7 @@ def decode(self, token_ids): outputs=[completion], finished=True, num_cached_tokens=0, + num_cache_creation_tokens=3, ) context.append_output(req_output) @@ -995,6 +999,9 @@ async def dummy_result_generator(): ) assert response.usage.output_tokens_details.reasoning_tokens == 1 + assert response.usage.input_tokens_details.cache_write_tokens == 3 + + OpenAIResponseUsage.model_validate(response.usage.model_dump(mode="json")) class TestExtractAllowedToolsFromMcpRequests: diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index 8c3d1467edcb..a5b1afe8c50f 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -682,6 +682,7 @@ def create_simple_context_output( prompt="Test prompt", prompt_token_ids=None, num_cached_tokens=0, + num_cache_creation_tokens=None, logprobs=None, finished=True, ): @@ -707,6 +708,7 @@ def create_simple_context_output( ], finished=finished, num_cached_tokens=num_cached_tokens, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -827,6 +829,7 @@ def test_simple_context_token_counting(): token_ids=[10, 11], prompt_token_ids=[1, 2, 3, 4, 5], num_cached_tokens=2, + num_cache_creation_tokens=3, ) ) context.append_output( @@ -835,12 +838,14 @@ def test_simple_context_token_counting(): token_ids=[12], prompt_token_ids=[1, 2, 3, 4, 5], num_cached_tokens=2, + num_cache_creation_tokens=3, ) ) assert context.num_prompt_tokens == 5 assert context.num_output_tokens == 3 # 2 + 1 assert context.num_cached_tokens == 2 + assert context.num_cache_creation_tokens == 3 def test_simple_context_final_output(): diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 077f78c02b25..3414bd6d82bb 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -189,6 +189,7 @@ def __init__( self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 # todo num_reasoning_tokens is not implemented yet. self.num_reasoning_tokens = 0 # not implemented yet for SimpleContext @@ -204,6 +205,8 @@ def append_output(self, output) -> None: raise ValueError("SimpleContext only supports RequestOutput.") self.num_prompt_tokens = len(output.prompt_token_ids or []) self.num_cached_tokens = output.num_cached_tokens or 0 + if output.num_cache_creation_tokens is not None: + self.num_cache_creation_tokens = output.num_cache_creation_tokens self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 201e46ed0419..ec28d2823ba4 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -80,6 +80,7 @@ class InputTokensDetails(OpenAIBaseModel): cached_tokens: int + cache_write_tokens: int input_tokens_per_turn: list[int] = Field(default_factory=list) cached_tokens_per_turn: list[int] = Field(default_factory=list) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 2c61205e2258..bce54338c4bb 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -837,6 +837,7 @@ async def responses_full_generator( output_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, input_tokens_details=InputTokensDetails( + cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0), cached_tokens=num_cached_tokens, input_tokens_per_turn=[ turn.input_tokens for turn in context.all_turn_metrics From 6c3d9816afbb3e5499a573248d1cf6fcb00c05cd Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:22:07 -0400 Subject: [PATCH 350/774] [BugFix][KV Connector] Fix Deadlock with KVConnector + MTP under KV Pressure (#57104) Signed-off-by: Robert Shaw Co-authored-by: Claude --- .../unit/test_remote_prefill_lifecycle.py | 60 +++++++++++++++++++ tests/v1/kv_connector/unit/utils.py | 8 +++ vllm/v1/core/sched/scheduler.py | 29 +++++++-- 3 files changed, 92 insertions(+), 5 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 4f82e828d3f9..caff9c9a490d 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -731,3 +731,63 @@ def test_async_loads_both_admitted_when_pool_fits(): for req in reqs: assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_async_load_reserves_blocks_for_promotion_margin(): + """An async load is not admitted unless the blocks its own promotion will + need are still free. + + A parked load is allocated without lookahead slots, but promotion pads it + to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without + reserving that margin, two loads can be admitted that together consume the + whole pool; the head then fails ``allocate_slots`` forever while the load + behind it is never reached, and with nothing running no block is ever freed. + + req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so + req_b must be held back in WAITING holding no blocks, leaving req_a room to + be promoted and run. + """ + vllm_config = create_vllm_config(num_speculative_tokens=3) + BLOCK_SIZE = vllm_config.cache_config.block_size + scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7 + + req_a = create_request( + request_id=1, + block_size=BLOCK_SIZE, + num_tokens=BLOCK_SIZE * 4, + do_remote_prefill=True, + max_tokens=1, + ) + req_b = create_request( + request_id=2, + block_size=BLOCK_SIZE, + num_tokens=BLOCK_SIZE * 3, + do_remote_prefill=True, + max_tokens=1, + ) + scheduler.add_request(req_a) + scheduler.add_request(req_b) + + # Both get a full external hit, so each would hold its whole prompt. + with patch.object( + scheduler.connector, + "get_num_new_matched_tokens", + side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)], + ): + scheduler.schedule() + + assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert req_b.status == RequestStatus.WAITING + req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[ + 0 + ].req_to_blocks + assert req_b.request_id not in req_to_blocks + + # req_a's load lands: it must be promotable and actually get scheduled. + scheduler.update_from_output( + scheduler.schedule(), + create_model_runner_output([], finished_recving={req_a.request_id}), + ) + scheduler_output = scheduler.schedule() + assert req_a.status == RequestStatus.RUNNING + assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0 diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 80b97c6a1366..85ada84b0dfe 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -17,6 +17,7 @@ KVTransferConfig, ModelConfig, SchedulerConfig, + SpeculativeConfig, VllmConfig, ) from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory @@ -106,6 +107,7 @@ def create_vllm_config( kv_connector_module_path: str | None = None, kv_role: str = "kv_consumer", disable_hybrid_kv_cache_manager: bool | None = None, + num_speculative_tokens: int | None = None, ) -> VllmConfig: """Initialize VllmConfig For Testing.""" model_config = ModelConfig( @@ -141,6 +143,11 @@ def create_vllm_config( kv_load_failure_policy=kv_load_failure_policy, ) attention_config = AttentionConfig(backend=attention_backend) + speculative_config = ( + SpeculativeConfig(model="ngram", num_speculative_tokens=num_speculative_tokens) + if num_speculative_tokens is not None + else None + ) return VllmConfig( scheduler_config=scheduler_config, model_config=model_config, @@ -148,6 +155,7 @@ def create_vllm_config( kv_transfer_config=kv_transfer_config, device_config=DeviceConfig("cpu"), attention_config=attention_config, + speculative_config=speculative_config, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 08f11911116d..19ea411aa19b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -31,6 +31,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.utils import get_mm_features_in_window +from vllm.utils.math_utils import cdiv from vllm.v1.core.encoder_cache_manager import ( EncoderCacheManager, EncoderDecoderCacheManager, @@ -1140,9 +1141,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if load_kv_async: # An async load holds its blocks for the whole transfer with # no forward progress and isn't preemptible here. Admit it - # only if it fits in (free - other in-flight reservations), to - # avoid deadlock and predictable preemptions. - reserved_blocks = self._inflight_prefill_reserved_blocks() + # only if it fits in (free - other in-flight reservations) + # plus its own spec decode step blocks, to avoid deadlock and + # predictable preemptions. + reserved_blocks = ( + self._inflight_prefill_reserved_blocks() + + self._spec_decode_step_blocks() + ) new_blocks = self.kv_cache_manager.allocate_slots( request, @@ -2892,9 +2897,9 @@ def _connector_finished( return delay_free or partial_tail_delay, kv_xfer_params def _request_remaining_blocks(self, request: Request) -> int: - """Blocks `request` still needs to allocate to hold its full sequence.""" + """Blocks `request` still needs to hold its full sequence and be promoted.""" full_num_tokens = min(request.num_tokens, self.max_model_len) - return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, num_tokens=full_num_tokens, new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, @@ -2904,6 +2909,20 @@ def _request_remaining_blocks(self, request: Request) -> int: num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) + return num_blocks + self._spec_decode_step_blocks() + + def _spec_decode_step_blocks(self) -> int: + """Number of blocks for the extra KV slots a spec decode step needs. + + When using async kv load, scheduler must reserve enough blocks for + full sequence + the spec decode step, otherwise request cannot be + able to run after the async_load if we are out of kv blocks. + """ + if not self.num_spec_tokens: + return 0 + return cdiv( + 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size + ) def _inflight_prefill_reserved_blocks(self) -> int: """Num blocks in-flight prefills still need to finish (their reservation).""" From 0983aef8da9b25ae5be72a5567e571a8f3f82963 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 17 Sep 2026 06:29:07 +0800 Subject: [PATCH 351/774] [Frontend] Return prompt metadata from /inference/v1/generate (#53187) --- .../server/src/routes/inference/generate.rs | 272 ++++++++++++- .../src/routes/inference/generate/convert.rs | 4 + .../src/routes/inference/generate/types.rs | 13 + rust/src/server/src/routes/render.rs | 1 + .../test_generate_stream.py | 357 ++++++++++++++++++ .../test_serving_multimodal_tokens.py | 74 ++-- .../token_in_token_out/mm_features.py | 28 +- .../scale_out/token_in_token_out/protocol.py | 16 + .../scale_out/token_in_token_out/serving.py | 38 +- 9 files changed, 763 insertions(+), 40 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 308ec3ca3f2e..dfb7e6cbb191 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -21,6 +21,7 @@ use thiserror_ext::AsReport as _; use tracing::{error, info, trace}; use tracing_futures::Instrument as _; use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs}; +use vllm_engine_core_client::protocol::multimodal::MmFeatureSpec; use vllm_llm::{ CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage, }; @@ -28,7 +29,7 @@ use vllm_llm::{ use self::convert::{ResponseOptions, prepare_generate_request}; use self::types::{ GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, - GenerateStreamResponse, + GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo, }; pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; @@ -64,6 +65,11 @@ pub async fn generate( } else { None }; + let mm_placeholders = body + .return_token_ids + .unwrap_or(false) + .then(|| extract_mm_placeholders(mm_features.as_deref())) + .flatten(); let prepared = match prepare_generate_request(body, &lora_resolution, request_context, mm_features) { @@ -98,6 +104,7 @@ pub async fn generate( prepared.request_id, api_server_options, prepared.options, + mm_placeholders, ); let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span); @@ -120,6 +127,7 @@ pub async fn generate( prepared.request_id, api_server_options, prepared.options, + mm_placeholders, ) { Ok(response) => response, Err(error) => return error.into_response(), @@ -143,19 +151,26 @@ async fn generate_chunk_stream( include_logprobs, // Ignored: raw generate streaming has no prompt-logprobs wire shape. include_prompt_logprobs: _, + return_token_ids, }: ResponseOptions, + mut mm_placeholders: Option, mut y: TryYielder, ) -> Result<(), ApiError> { pin_mut!(stream); let mut prompt_tokens = None; + let mut prompt_token_ids = None; let mut usage = TokenUsage::default(); while let Some(next) = stream.next().await { match next { Ok(output) => { - if prompt_tokens.is_none() { - prompt_tokens = - output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len()); + if prompt_tokens.is_none() + && let Some(info) = output.prompt_info.as_ref() + { + prompt_tokens = Some(info.prompt_token_ids.len()); + if return_token_ids { + prompt_token_ids = Some(info.prompt_token_ids.to_vec()); + } } usage.prompt_token_count = prompt_tokens.unwrap_or_default(); usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count); @@ -195,6 +210,7 @@ async fn generate_chunk_stream( None }; + let prompt_token_ids = prompt_token_ids.take(); y.yield_ok(GenerateStreamResponse { request_id: request_id.clone(), choices: vec![GenerateResponseStreamChoice { @@ -205,6 +221,8 @@ async fn generate_chunk_stream( }], usage: include_continuous_usage .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), + mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()), + prompt_token_ids, }) .await; } @@ -223,6 +241,8 @@ async fn generate_chunk_stream( request_id, choices: Vec::new(), usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)), + prompt_token_ids: None, + mm_placeholders: None, }) .await; } @@ -244,7 +264,9 @@ fn collect_generate( include_continuous_usage: _, include_logprobs, include_prompt_logprobs, + return_token_ids, }: ResponseOptions, + mm_placeholders: Option, ) -> Result { let logprobs = if include_logprobs { let logprobs = collected.logprobs.as_ref().ok_or_else(|| { @@ -292,11 +314,34 @@ fn collect_generate( token_ids: collected.token_ids, }], prompt_logprobs, + prompt_token_ids: return_token_ids.then_some(collected.prompt_token_ids), + mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(), kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, }) } +fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option { + let features = features?; + if features.is_empty() { + return None; + } + + let mut placeholders = MultiModalPlaceholders::new(); + for feature in features { + placeholders.entry(feature.modality.as_str().to_string()).or_default().push( + PlaceholderRangeInfo { + offset: feature.mm_position.offset, + length: feature.mm_position.length, + }, + ); + } + for ranges in placeholders.values_mut() { + ranges.sort_by_key(|range| range.offset); + } + Some(placeholders) +} + fn raw_logprobs_to_openai_chat(logprobs: &Logprobs) -> Result { let content = logprobs .positions @@ -419,6 +464,7 @@ mod tests { use std::sync::Arc; use futures::{TryStreamExt as _, stream}; + use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; use vllm_llm::GeneratePromptInfo; use super::*; @@ -451,6 +497,13 @@ mod tests { }), ]); + let mm_placeholders = MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 1, + length: 4, + }], + )]); let chunks: Vec<_> = generate_chunk_stream( stream, "raw-stream".to_string(), @@ -461,14 +514,20 @@ mod tests { ResponseOptions { include_usage: true, include_continuous_usage: true, + return_token_ids: true, ..Default::default() }, + Some(mm_placeholders.clone()), ) .try_collect() .await .expect("collect chunks"); assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22])); + assert_eq!(chunks[0].mm_placeholders, Some(mm_placeholders)); + assert!(chunks[1].prompt_token_ids.is_none()); + assert!(chunks[1].mm_placeholders.is_none()); assert_eq!( chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens, 2 @@ -499,6 +558,209 @@ mod tests { ); } + fn stream_output( + prompt_token_ids: Option<&[u32]>, + token_ids: Vec, + finish_reason: Option, + ) -> GenerateOutput { + GenerateOutput { + request_id: String::new(), + prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { + prompt_token_ids: Arc::from(ids), + prompt_logprobs: None, + }), + token_ids, + logprobs: None, + finish_reason, + cached_token_count: 0, + kv_transfer_params: None, + ec_transfer_params: None, + } + } + + async fn collect_chunks( + outputs: Vec, + return_token_ids: bool, + mm_placeholders: Option, + ) -> Vec { + generate_chunk_stream( + stream::iter(outputs.into_iter().map(Ok)), + "raw-stream".to_string(), + ApiServerOptions::default(), + ResponseOptions { + return_token_ids, + ..Default::default() + }, + mm_placeholders, + ) + .try_collect() + .await + .expect("collect chunks") + } + + #[tokio::test] + async fn generate_chunk_stream_omits_prompt_metadata_by_default() { + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), vec![33], None), + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())), + ], + false, + Some(MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 1, + length: 4, + }], + )])), + ) + .await; + + assert_eq!(chunks.len(), 2); + assert!(chunks.iter().all(|c| c.prompt_token_ids.is_none())); + assert!(chunks.iter().all(|c| c.mm_placeholders.is_none())); + } + + #[tokio::test] + async fn generate_chunk_stream_sends_prompt_metadata_on_first_emitted_chunk() { + // An empty first output is skipped and must not consume the metadata. + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), Vec::new(), None), + stream_output(Some(&[11, 22]), vec![33], None), + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())), + ], + true, + None, + ) + .await; + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22])); + assert!(chunks[0].mm_placeholders.is_none()); + assert!(chunks[1].prompt_token_ids.is_none()); + assert!(chunks[1].mm_placeholders.is_none()); + } + + #[test] + fn collect_generate_omits_prompt_metadata_by_default() { + let output = CollectedGenerateOutput { + request_id: "raw-1".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 1, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + }; + + let response = collect_generate( + output, + "raw-1".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + Some(MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 0, + length: 1, + }], + )])), + ) + .expect("response"); + + assert!(response.prompt_token_ids.is_none()); + assert!(response.mm_placeholders.is_none()); + } + + #[test] + fn extract_mm_placeholders_returns_none_without_features() { + assert!(extract_mm_placeholders(None).is_none()); + assert!(extract_mm_placeholders(Some(&[])).is_none()); + } + + #[test] + fn extract_mm_placeholders_groups_ranges_in_prompt_order() { + let feature = |modality: MmModality, offset: usize, length: usize| MmFeatureSpec { + data: None, + modality, + identifier: format!("{}-{offset}", modality.as_str()), + mm_position: PlaceholderRange { + offset, + length, + is_embed: None, + }, + mm_hash: None, + }; + let features = vec![ + feature(MmModality::Image, 8, 2), + feature(MmModality::Audio, 4, 3), + feature(MmModality::Image, 1, 5), + ]; + + let placeholders = extract_mm_placeholders(Some(&features)).expect("metadata"); + + assert_eq!( + placeholders["image"], + vec![ + PlaceholderRangeInfo { + offset: 1, + length: 5, + }, + PlaceholderRangeInfo { + offset: 8, + length: 2, + }, + ] + ); + assert_eq!( + placeholders["audio"], + vec![PlaceholderRangeInfo { + offset: 4, + length: 3, + }] + ); + } + + #[test] + fn collect_generate_returns_prompt_token_ids_when_requested() { + let output = CollectedGenerateOutput { + request_id: "raw-1".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 1, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + }; + + let response = collect_generate( + output, + "raw-1".to_string(), + ApiServerOptions::default(), + ResponseOptions { + return_token_ids: true, + ..Default::default() + }, + None, + ) + .expect("response"); + + assert_eq!(response.prompt_token_ids, Some(vec![10, 20])); + } + #[test] fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() { let output_without_payload = |prompt_token_ids: Vec| CollectedGenerateOutput { @@ -525,6 +787,7 @@ mod tests { include_prompt_logprobs: true, ..Default::default() }, + None, ) .expect("single-token prompt without payload maps to [None]"); let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present"); @@ -539,6 +802,7 @@ mod tests { include_prompt_logprobs: true, ..Default::default() }, + None, ) .expect_err("multi-token prompt without payload is an engine failure"); } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 08c4ce18be7e..8093a34ca29d 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -30,6 +30,8 @@ pub(super) struct ResponseOptions { pub include_logprobs: bool, /// Whether the caller requested top-level prompt logprobs. pub include_prompt_logprobs: bool, + /// Whether the caller requested final prompt token metadata. + pub return_token_ids: bool, } /// Validate and lower one raw generate request into the internal @@ -56,6 +58,7 @@ pub(super) fn prepare_generate_request( .unwrap_or(false); let include_logprobs = request.sampling_params.inner.logprobs.is_some(); let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some(); + let return_token_ids = request.return_token_ids.unwrap_or(false); let mut sampling_params = request.sampling_params.inner; sampling_params.vllm_xargs = merge_kv_transfer_params( sampling_params.vllm_xargs, @@ -93,6 +96,7 @@ pub(super) fn prepare_generate_request( include_continuous_usage, include_logprobs, include_prompt_logprobs, + return_token_ids, }, }) } diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index b9220aa7b767..406ec585692b 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -46,6 +46,7 @@ pub struct GenerateRequest { pub ec_transfer_params: Option>, /// Raw multimodal input; server resolves media. Mutually exclusive with `features`. pub content_parts: Option>, + pub return_token_ids: Option, #[serde(flatten)] pub other: Map, } @@ -81,6 +82,8 @@ pub(super) struct GenerateStreamResponse { pub request_id: String, pub choices: Vec, pub usage: Option, + pub prompt_token_ids: Option>, + pub mm_placeholders: Option, } /// Mirrors the Python vLLM `GenerateResponse` class. @@ -89,10 +92,20 @@ pub(super) struct GenerateResponse { pub request_id: String, pub choices: Vec, pub prompt_logprobs: Option>>>, + pub prompt_token_ids: Option>, + pub mm_placeholders: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, } +pub(super) type MultiModalPlaceholders = HashMap>; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub(super) struct PlaceholderRangeInfo { + pub offset: usize, + pub length: usize, +} + /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads. #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateLogprob { diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index a4ae0298e91d..2d38b63f4991 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -105,6 +105,7 @@ fn lower_render_request( kv_transfer_params: None, ec_transfer_params: None, content_parts: None, + return_token_ids: None, other: Default::default(), }; validate_generate_request(&request, &state.served_model_names)?; diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index ad39958702a8..d9ae424cc85b 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -12,6 +12,9 @@ from vllm.entrypoints.generate.base.protocol import StreamOptions from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.scale_out.token_in_token_out.mm_features import ( + placeholder_ranges_from_engine_input, +) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, GenerateResponse, @@ -19,6 +22,7 @@ from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.exceptions import GenerationError from vllm.logprobs import Logprob +from vllm.multimodal.inputs import PlaceholderRange from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers import renderer_from_config from vllm.renderers.online_renderer import OnlineRenderer @@ -176,6 +180,29 @@ def _parse_sse_chunks(chunks: list[str]) -> list[Any]: return parsed +def test_placeholder_ranges_from_engine_input(): + engine_input: Any = { + "type": "multimodal", + "mm_placeholders": { + "image": [ + PlaceholderRange(offset=1, length=5), + PlaceholderRange(offset=8, length=2), + ] + }, + } + + placeholders = placeholder_ranges_from_engine_input(engine_input) + + assert placeholders is not None + assert [item.model_dump() for item in placeholders["image"]] == [ + {"offset": 1, "length": 5}, + {"offset": 8, "length": 2}, + ] + + text_only: Any = {"type": "token", "prompt_token_ids": [1, 2, 3]} + assert placeholder_ranges_from_engine_input(text_only) is None + + @pytest.mark.asyncio async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution(): engine = _mock_engine() @@ -294,6 +321,336 @@ async def mock_generate(*args, **kwargs): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +async def test_non_stream_returns_final_prompt_token_ids_when_requested(): + engine = _mock_engine() + final_prompt_token_ids = [1, 2, 2, 2, 3] + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + prompt_token_ids=final_prompt_token_ids, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == final_prompt_token_ids + + +@pytest.mark.asyncio +async def test_stream_returns_prompt_token_ids_only_once_when_requested(): + engine = _mock_engine() + final_prompt_token_ids = [1, 2, 2, 2, 3] + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], prompt_token_ids=final_prompt_token_ids + ) + yield _make_request_output( + "req-1", + token_ids=[20], + finish_reason="stop", + finished=True, + prompt_token_ids=final_prompt_token_ids, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks[0]["prompt_token_ids"] == final_prompt_token_ids + assert "prompt_token_ids" not in data_chunks[1] + + +MM_ENGINE_INPUT: Any = { + "type": "multimodal", + "prompt_token_ids": [1, 2, 2, 2, 3], + "mm_placeholders": {"image": [PlaceholderRange(offset=1, length=3)]}, +} +MM_PLACEHOLDERS = {"image": [{"offset": 1, "length": 3}]} + + +def _build_mm_serving_tokens(engine: AsyncLLM) -> ServingTokens: + serving = _build_serving_tokens(engine) + + async def _fake_preprocess(*args, **kwargs): + return [MM_ENGINE_INPUT] + + serving.online_renderer.preprocess_completion = AsyncMock( + side_effect=_fake_preprocess + ) + return serving + + +@pytest.mark.asyncio +@pytest.mark.parametrize("return_token_ids", [False, True]) +async def test_non_stream_prompt_metadata_requires_return_token_ids( + return_token_ids: bool, +): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=return_token_ids, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + if return_token_ids: + assert response.prompt_token_ids == [1, 2, 3] + assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS + else: + assert response.prompt_token_ids is None + assert response.mm_placeholders is None + assert request._response_mm_placeholders is None + + +@pytest.mark.asyncio +async def test_non_stream_text_only_prompt_has_no_mm_placeholders(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == [1, 2, 3] + assert response.mm_placeholders is None + + +@pytest.mark.asyncio +async def test_non_stream_encoder_decoder_omits_mm_placeholders(): + engine = _mock_engine() + engine.model_config.is_encoder_decoder = True + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == [1, 2, 3] + assert response.mm_placeholders is None + + +@pytest.mark.asyncio +async def test_stream_prompt_metadata_lands_on_first_emitted_chunk(): + """An empty first output is skipped; metadata goes on the first real chunk, + once, and never on the sibling chunk of the same iteration (n=2).""" + engine = _mock_engine() + + def _two_outputs(tokens0: list[int], tokens1: list[int], **kwargs): + res = _make_request_output("req-1", token_ids=tokens0, **kwargs) + res.outputs.append( + CompletionOutput( + index=1, + text="", + token_ids=tokens1, + cumulative_logprob=None, + logprobs=None, + finish_reason=kwargs.get("finish_reason"), + ) + ) + return res + + async def mock_generate(*args, **kwargs): + yield _two_outputs([], []) + yield _two_outputs([10], [11]) + yield _two_outputs([20], [21], finish_reason="stop", finished=True) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2, n=2), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert [c["choices"][0]["token_ids"] for c in data_chunks] == [ + [10], + [11], + [20], + [21], + ] + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS + for chunk in data_chunks[1:]: + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + + +@pytest.mark.asyncio +async def test_stream_omits_prompt_metadata_by_default(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks + for chunk in data_chunks: + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + + +@pytest.mark.asyncio +async def test_stream_text_only_prompt_omits_mm_placeholders_key(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert "mm_placeholders" not in data_chunks[0] + + +@pytest.mark.asyncio +async def test_stream_zero_token_completion_still_delivers_prompt_metadata(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert len(data_chunks) == 1 + assert data_chunks[0]["choices"][0]["token_ids"] == [] + assert data_chunks[0]["choices"][0]["finish_reason"] == "stop" + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS + + +@pytest.mark.asyncio +async def test_stream_zero_token_completion_emits_no_chunk_by_default(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + + assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")] + assert parsed[-1] == "[DONE]" + + @pytest.mark.asyncio async def test_stream_error_mid_generation(): """finish_reason='error' mid-stream yields error chunk then [DONE].""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 892d3c8285da..1d017bd3cf1b 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -154,11 +154,28 @@ async def test_render_to_generate_roundtrip(client, test_image): ) -@pytest.mark.asyncio -async def test_content_parts_generates_tokens(client, test_image): - """content_parts with raw media should produce output tokens.""" - data_url = encode_image_url(test_image, format="PNG") +def _collapse_placeholders( + token_ids: list[int], mm_placeholders: dict[str, list[dict[str, int]]] +) -> list[int]: + """Undo placeholder expansion: each expanded range becomes one token. + + ``/generate`` with ``content_parts`` expects unexpanded placeholders and + expands them itself, so the roundtrip must start from a collapsed prompt. + """ + ranges = sorted( + (r["offset"], r["length"]) for rs in mm_placeholders.values() for r in rs + ) + collapsed: list[int] = [] + pos = 0 + for offset, length in ranges: + collapsed.extend(token_ids[pos:offset]) + collapsed.append(token_ids[offset]) + pos = offset + length + collapsed.extend(token_ids[pos:]) + return collapsed + +async def _render_image_prompt(client, data_url: str) -> dict: render_resp = await client.post( RENDER_ENDPOINT, json={ @@ -175,14 +192,26 @@ async def test_content_parts_generates_tokens(client, test_image): }, ) render_resp.raise_for_status() - token_ids = render_resp.json()["token_ids"] + return render_resp.json() + + +@pytest.mark.asyncio +async def test_content_parts_generates_tokens(client, test_image): + """Unexpanded prompt + content_parts must expand exactly like /render.""" + data_url = encode_image_url(test_image, format="PNG") + render_data = await _render_image_prompt(client, data_url) + expanded = render_data["token_ids"] + mm_placeholders = render_data["features"]["mm_placeholders"] + unexpanded = _collapse_placeholders(expanded, mm_placeholders) + assert len(unexpanded) < len(expanded) gen_resp = await client.post( GEN_ENDPOINT, json={ - "token_ids": token_ids, + "token_ids": unexpanded, "content_parts": [{"type": "image_url", "url": data_url}], "sampling_params": {"max_tokens": 10, "temperature": 0.0}, + "return_token_ids": True, }, ) gen_resp.raise_for_status() @@ -192,6 +221,8 @@ async def test_content_parts_generates_tokens(client, test_image): choice = gen_data["choices"][0] assert "token_ids" in choice assert len(choice["token_ids"]) > 0 + assert gen_data["prompt_token_ids"] == expanded + assert gen_data["mm_placeholders"] == mm_placeholders @pytest.mark.asyncio @@ -200,33 +231,20 @@ async def test_content_parts_streaming(client, test_image): import json as json_mod data_url = encode_image_url(test_image, format="PNG") - - render_resp = await client.post( - RENDER_ENDPOINT, - json={ - "model": MODEL_NAME, - "messages": [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": {"url": data_url}}, - {"type": "text", "text": "Describe this."}, - ], - } - ], - }, - ) - render_resp.raise_for_status() - token_ids = render_resp.json()["token_ids"] + render_data = await _render_image_prompt(client, data_url) + expanded = render_data["token_ids"] + mm_placeholders = render_data["features"]["mm_placeholders"] + unexpanded = _collapse_placeholders(expanded, mm_placeholders) async with client.stream( "POST", GEN_ENDPOINT, json={ - "token_ids": token_ids, + "token_ids": unexpanded, "content_parts": [{"type": "image_url", "url": data_url}], "sampling_params": {"max_tokens": 8, "temperature": 0.0}, "stream": True, + "return_token_ids": True, }, ) as resp: resp.raise_for_status() @@ -236,3 +254,9 @@ async def test_content_parts_streaming(client, test_image): chunks.append(json_mod.loads(line[6:])) assert len(chunks) > 0 + prompt_metadata_chunks = [ + chunk for chunk in chunks if chunk.get("prompt_token_ids") is not None + ] + assert len(prompt_metadata_chunks) == 1 + assert prompt_metadata_chunks[0]["prompt_token_ids"] == expanded + assert prompt_metadata_chunks[0]["mm_placeholders"] == mm_placeholders diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py index d7f89e895e5f..62d006783a2e 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py @@ -117,6 +117,24 @@ def mm_kwargs_from_features( return mm_kwargs +def placeholder_ranges_from_engine_input( + engine_input: EngineInput, +) -> dict[str, list[PlaceholderRangeInfo]] | None: + """Return per-modality placeholder ranges, or ``None`` for text-only prompts.""" + if engine_input.get("type") != "multimodal": + return None + + raw_placeholders: MultiModalPlaceholders = cast(MultiModalInput, engine_input)[ + "mm_placeholders" + ] + return { + modality: [ + PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges + ] + for modality, ranges in raw_placeholders.items() + } + + def extract_mm_features( engine_input: EngineInput, *, @@ -134,14 +152,8 @@ def extract_mm_features( mm_engine_input = cast(MultiModalInput, engine_input) mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"] - raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"] - - mm_placeholders = { - modality: [ - PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges - ] - for modality, ranges in raw_placeholders.items() - } + mm_placeholders = placeholder_ranges_from_engine_input(engine_input) + assert mm_placeholders is not None kwargs_data: dict[str, list[str | None]] | None = None mm_metadata: dict[str, list[str | None]] | None = None diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index aec8431b4780..be8130e314cd 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -211,6 +211,15 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest": model: str | None = None + return_token_ids: bool | None = Field( + default=None, + description=( + "If true, return the final prompt token IDs after multimodal " + "placeholder expansion, together with multimodal placeholder ranges. " + "In streaming mode, this metadata is included only in the first chunk." + ), + ) + stream: bool | None = False stream_options: StreamOptions | None = None cache_salt: str | None = Field( @@ -255,6 +264,9 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest": # ``SamplingParams`` instance (e.g. from internal callers that have # already resolved values), in which case all fields are considered set. _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None) + _response_mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = ( + PrivateAttr(default=None) + ) @model_validator(mode="wrap") @classmethod @@ -354,6 +366,8 @@ class GenerateStreamResponse(BaseModel): ) choices: list[GenerateResponseStreamChoice] usage: UsageInfo | None = Field(default=None) + prompt_token_ids: list[int] | None = None + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None class GenerateResponse(BaseModel): @@ -370,6 +384,8 @@ class GenerateResponse(BaseModel): choices: list[GenerateResponseChoice] usage: UsageInfo | None = Field(default=None) prompt_logprobs: list[dict[int, Logprob] | None] | None = None + prompt_token_ids: list[int] | None = None + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None kv_transfer_params: dict[str, Any] | None = Field( default=None, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index d3725e3cf71e..a8a2804a0634 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -44,7 +44,10 @@ from vllm.utils.collection_utils import as_list from vllm.utils.serial_utils import numpy2base64 -from .mm_features import mm_kwargs_from_features +from .mm_features import ( + mm_kwargs_from_features, + placeholder_ranges_from_engine_input, +) from .protocol import ( GenerateRequest, GenerateResponse, @@ -195,6 +198,14 @@ async def serve_tokens( skip_mm_cache=True, ) + # Offsets are relative to the decoder prompt, so they are not + # meaningful for encoder-decoder models. + request._response_mm_placeholders = ( + placeholder_ranges_from_engine_input(engine_input) + if request.return_token_ids and not self.model_config.is_encoder_decoder + else None + ) + # Schedule the request and get the result generator. result_generator: AsyncGenerator[RequestOutput, None] | None = None @@ -359,6 +370,10 @@ async def serve_tokens_full_generator( choices=choices, usage=usage, prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), + prompt_token_ids=( + final_res.prompt_token_ids if request.return_token_ids else None + ), + mm_placeholders=request._response_mm_placeholders, kv_transfer_params=final_res.kv_transfer_params, ec_transfer_params=final_res.ec_transfer_params, ) @@ -395,6 +410,7 @@ async def serve_tokens_stream_generator( num_prompt_tokens = 0 num_generated_tokens: list[int] = [] first_iteration = True + prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params @@ -407,6 +423,8 @@ async def serve_tokens_stream_generator( if first_iteration: if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) + if request.return_token_ids: + prompt_token_ids = res.prompt_token_ids if res.encoder_prompt_token_ids is not None: num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens @@ -421,7 +439,11 @@ async def serve_tokens_stream_generator( finish_reason = output.finish_reason self._raise_if_error(finish_reason, request_id) - if not delta_token_ids: + # Still emit a terminal empty chunk while prompt metadata + # is pending, so zero-token completions deliver it. + if not delta_token_ids and ( + finish_reason is None or prompt_token_ids is None + ): continue if sampling_params.logprobs is not None: @@ -453,6 +475,10 @@ async def serve_tokens_stream_generator( ) ], ) + if prompt_token_ids is not None: + chunk.prompt_token_ids = prompt_token_ids + chunk.mm_placeholders = request._response_mm_placeholders + prompt_token_ids = None if include_continuous_usage: chunk.usage = UsageInfo( prompt_tokens=num_prompt_tokens, @@ -460,7 +486,13 @@ async def serve_tokens_stream_generator( total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - yield f"data: {chunk.model_dump_json()}\n\n" + # Omit absent prompt metadata (Rust skips None fields). + exclude = { + name + for name in ("prompt_token_ids", "mm_placeholders") + if getattr(chunk, name) is None + } + yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo( From fbf2c5e8be9754f31c4e0f549189fc9f3bdd213c Mon Sep 17 00:00:00 2001 From: xinnywinne <36653773+xinnywinne@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:52:05 -0700 Subject: [PATCH 352/774] [Frontend] Add per-request metrics to Responses API (#55084) Signed-off-by: wxin Co-authored-by: Benjamin Chislett --- docs/features/per_request_metrics.md | 12 ++ .../responses/test_serving_responses.py | 147 ++++++++++++++++-- vllm/entrypoints/generate/api_router.py | 1 + vllm/entrypoints/openai/responses/context.py | 7 + vllm/entrypoints/openai/responses/protocol.py | 13 +- vllm/entrypoints/openai/responses/serving.py | 18 ++- 6 files changed, 184 insertions(+), 14 deletions(-) diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md index 25211ecfc39b..da1685f92925 100644 --- a/docs/features/per_request_metrics.md +++ b/docs/features/per_request_metrics.md @@ -120,6 +120,18 @@ the same `metrics` response field. As with `n > 1`, metrics are omitted for requests with multiple prompts, because the timing data cannot be attributed to a single prompt's generation. +## Responses API + +Per-request metrics are available on the `/v1/responses` endpoint using the +common `metrics` object described above. Non-streaming responses include it at +the top level. Streaming responses include it in the final response carried by +the `response.completed` event; intermediate events do not include metrics. + +Metrics are omitted for Responses requests that perform multiple +model-generation turns, such as built-in tool-call workflows, because the +response retains timing data for only one generation turn while token usage is +accumulated across all turns. + ## Relationship to Prometheus Metrics The `metrics` response field provides per-request values for a single request. diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 0f361c1c6a50..1a315902e249 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -48,6 +48,7 @@ SimpleContext, ) from vllm.entrypoints.openai.responses.protocol import ( + ResponseCompletedEvent, ResponseCreatedEvent, ResponseRawMessageAndToken, ResponsesRequest, @@ -71,6 +72,7 @@ _extract_allowed_tools_from_mcp_requests, ) from vllm.sampling_params import SamplingParams +from vllm.v1.metrics.stats import RequestStateStats pytestmark = pytest.mark.skip_global_cleanup @@ -795,6 +797,7 @@ async def generate_output(): pass assert serving_responses_instance.engine_client.generate.call_count == 2 + assert context.request_metrics_cover_all_generation_turns is False followup_engine_input = ( serving_responses_instance.engine_client.generate.call_args_list[1].args[0] ) @@ -1248,9 +1251,19 @@ def test_zero_delta_items_should_preserve_streaming_lifecycle( assert "response.output_item.done" in type_names -def _make_simple_context_with_output(text, token_ids, response_parser=None): - """Create a SimpleContext with a RequestOutput containing the given text.""" - ctx = SimpleContext(response_parser=response_parser) +_PER_REQUEST_STATS = RequestStateStats( + queued_ts=1.0, + scheduled_ts=1.5, + first_token_ts=2.0, + last_token_ts=3.0, +) + + +def _make_request_output( + text, + token_ids, + metrics: RequestStateStats | None = None, +): completion = CompletionOutput( index=0, text=text, @@ -1260,7 +1273,7 @@ def _make_simple_context_with_output(text, token_ids, response_parser=None): finish_reason=None, stop_reason=None, ) - req_output = RequestOutput( + return RequestOutput( request_id="req", prompt="hi", prompt_token_ids=[7, 8], @@ -1268,13 +1281,29 @@ def _make_simple_context_with_output(text, token_ids, response_parser=None): outputs=[completion], finished=False, num_cached_tokens=0, + metrics=metrics, ) + + +def _make_simple_context_with_output( + text, + token_ids, + response_parser=None, + metrics: RequestStateStats | None = None, +): + """Create a SimpleContext with a RequestOutput containing the given text.""" + ctx = SimpleContext(response_parser=response_parser) + req_output = _make_request_output(text, token_ids, metrics) ctx.append_output(req_output) + ctx.request_metrics = metrics return ctx -def _make_serving_instance_with_reasoning(): - """Create an OpenAIServingResponses with a mocked reasoning parser.""" +def _make_serving_instance( + *, + reasoning_parser: str = "", + enable_per_request_metrics: bool = False, +) -> OpenAIServingResponses: engine_client = MagicMock() model_config = MagicMock() model_config.max_model_len = 100 @@ -1285,18 +1314,112 @@ def _make_serving_instance_with_reasoning(): engine_client.input_processor = MagicMock() engine_client.renderer = MagicMock() - models = MagicMock() - - serving = OpenAIServingResponses( + return OpenAIServingResponses( engine_client=engine_client, - models=models, + models=MagicMock(), online_renderer=MagicMock(), request_logger=None, chat_template=None, chat_template_content_format="auto", - reasoning_parser="qwen3", + reasoning_parser=reasoning_parser, + enable_per_request_metrics=enable_per_request_metrics, + ) + + +async def _empty_context_generator(): + if False: + yield + + +async def _make_full_metrics_response( + enable_per_request_metrics: bool, + request_metrics_cover_all_generation_turns: bool = True, +): + serving = _make_serving_instance( + enable_per_request_metrics=enable_per_request_metrics + ) + request = ResponsesRequest(input="hi", tools=[], stream=False, store=False) + sampling_params = SamplingParams(max_tokens=16) + context = SimpleContext() + context.request_metrics_cover_all_generation_turns = ( + request_metrics_cover_all_generation_turns + ) + + async def generate(*args, **kwargs): + yield _make_request_output("hello", [10, 20], _PER_REQUEST_STATS) + + serving.engine_client.generate.side_effect = generate + result_generator = serving._generate_with_builtin_tools( + request_id=request.request_id, + engine_input=tokens_input([7, 8]), + sampling_params=sampling_params, + context=context, ) - return serving + response = await serving.responses_full_generator( + request=request, + sampling_params=sampling_params, + result_generator=result_generator, + context=context, + model_name="test-model", + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="req"), + ) + assert isinstance(response, ResponsesResponse) + return response + + +@pytest.mark.asyncio +async def test_responses_per_request_metrics_follow_server_flag(): + disabled_response = await _make_full_metrics_response(False) + assert disabled_response.metrics is None + assert "metrics" not in disabled_response.model_dump(mode="json") + + enabled_response = await _make_full_metrics_response(True) + assert enabled_response.metrics is not None + assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +@pytest.mark.asyncio +async def test_responses_streaming_metrics_only_on_completed_event(): + serving = _make_serving_instance(enable_per_request_metrics=True) + request = ResponsesRequest(input="hi", tools=[], stream=True, store=False) + context = _make_simple_context_with_output( + "hello", [10, 20], metrics=_PER_REQUEST_STATS + ) + + events = [ + event + async for event in serving.responses_stream_generator( + request=request, + sampling_params=SamplingParams(max_tokens=16), + result_generator=_empty_context_generator(), + context=context, + model_name="test-model", + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="req"), + ) + ] + + for event in events[:-1]: + assert "metrics" not in event.response.model_dump(mode="json") + assert isinstance(events[-1], ResponseCompletedEvent) + assert events[-1].response.metrics is not None + assert events[-1].response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +@pytest.mark.asyncio +async def test_responses_metrics_suppressed_for_multiple_generation_turns(): + response = await _make_full_metrics_response( + True, request_metrics_cover_all_generation_turns=False + ) + + assert response.metrics is None + assert "metrics" not in response.model_dump(mode="json") + + +def _make_serving_instance_with_reasoning(): + """Create an OpenAIServingResponses with a mocked reasoning parser.""" + return _make_serving_instance(reasoning_parser="qwen3") def _identity_increment(event): diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 12c70e37d5c0..ac995a4f031a 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -116,6 +116,7 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_per_request_metrics=args.enable_per_request_metrics, enable_log_outputs=args.enable_log_outputs, default_chat_template_kwargs=default_chat_template_kwargs, ) diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 3414bd6d82bb..0db47cf367e6 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -45,6 +45,8 @@ if TYPE_CHECKING: from mcp.client import ClientSession + from vllm.v1.metrics.stats import RequestStateStats + logger = logging.getLogger(__name__) # This is currently needed as the tool type doesn't 1:1 match the @@ -102,6 +104,11 @@ def copy(self) -> "TurnMetrics": class ConversationContext(ABC): response_parser: Parser | None = None + request_metrics: "RequestStateStats | None" = None + # Built-in tools can trigger additional model generations. In that case, + # the stored engine timestamps cover only one turn, while token usage is + # accumulated across all turns. + request_metrics_cover_all_generation_turns: bool = True @abstractmethod def append_output(self, output: RequestOutput) -> None: diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index ec28d2823ba4..750d1db84c66 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -60,7 +60,11 @@ ChatCompletionMessageParam, ChatTemplateContentFormatOption, ) -from vllm.entrypoints.generate.base.protocol import StopParam, validate_cache_salt +from vllm.entrypoints.generate.base.protocol import ( + PerRequestMetrics, + StopParam, + validate_cache_salt, +) from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger @@ -684,6 +688,11 @@ class ResponsesResponse(OpenAIBaseModel): usage: ResponseUsage | None = None user: str | None = None + # vLLM-specific per-request metrics. Omitted unless enabled server-side. + metrics: PerRequestMetrics | None = Field( + default=None, exclude_if=lambda value: value is None + ) + presence_penalty: float | None = Field( default=None, ge=-2.0, @@ -752,6 +761,7 @@ def from_request( output: list[ResponseOutputItem], status: ResponseStatus, usage: ResponseUsage | None = None, + metrics: PerRequestMetrics | None = None, input_messages: ResponseInputOutputMessage | None = None, output_messages: ResponseInputOutputMessage | None = None, kv_transfer_params: dict[str, Any] | None = None, @@ -795,6 +805,7 @@ def from_request( truncation=request.truncation, user=request.user, usage=usage, + metrics=metrics, kv_transfer_params=kv_transfer_params, ec_transfer_params=ec_transfer_params, ) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index bce54338c4bb..ee4f7d55ca0d 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -30,7 +30,10 @@ DeltaMessage, RequestResponseMetadata, ) -from vllm.entrypoints.generate.base.serving import GenerateBaseServing +from vllm.entrypoints.generate.base.serving import ( + GenerateBaseServing, + build_per_request_timing_metrics, +) from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.context import ( @@ -109,6 +112,7 @@ def __init__( tool_server: ToolServer | None = None, enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_per_request_metrics: bool = False, enable_log_outputs: bool = False, default_chat_template_kwargs: dict[str, Any] | None = None, ) -> None: @@ -136,6 +140,7 @@ def __init__( ) self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage + self.enable_per_request_metrics = enable_per_request_metrics self.default_sampling_params = self.model_config.get_diff_sampling_param() mc = self.model_config @@ -645,6 +650,7 @@ async def _generate_with_builtin_tools( ) async for res in generator: + context.request_metrics = res.metrics context.append_output(res) # NOTE(woosuk): The stop condition is handled by the engine. yield context @@ -654,6 +660,7 @@ async def _generate_with_builtin_tools( break # Call the tool and update the context with the result. + context.request_metrics_cover_all_generation_turns = False tool_output = await context.call_tool() context.append_tool_output(tool_output) @@ -857,6 +864,14 @@ async def responses_full_generator( ], ), ) + per_request_metrics = None + if ( + self.enable_per_request_metrics + and context.request_metrics_cover_all_generation_turns + ): + per_request_metrics = build_per_request_timing_metrics( + context.request_metrics, num_generated_tokens + ) response = ResponsesResponse.from_request( request, sampling_params, @@ -867,6 +882,7 @@ async def responses_full_generator( output=output, status=status, usage=usage, + metrics=per_request_metrics, kv_transfer_params=context.kv_transfer_params, ec_transfer_params=context.ec_transfer_params, ) From 62f948258385418ed248d5cdfceab4d29c1b5bfd Mon Sep 17 00:00:00 2001 From: andrewor14 Date: Wed, 16 Sep 2026 19:11:27 -0400 Subject: [PATCH 353/774] [Tests] Delete deprecate torchao tests for v1 configs (#52956) Signed-off-by: andrewor14 --- tests/quantization/test_torchao.py | 60 ------------------------------ 1 file changed, 60 deletions(-) diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index 50817ce3303d..3c839a687e3a 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -9,7 +9,6 @@ from tests.quantization.utils import load_model_without_vllm_runner from vllm.config import set_current_vllm_config -from vllm.config.load import LoadConfig from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import Attention from vllm.model_executor.model_loader import get_model_loader @@ -67,65 +66,6 @@ def test_pre_quantized_model(monkeypatch, dist_init, workspace_init): assert torch.isfinite(logits).all() -@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") -@pytest.mark.parametrize( - "pt_load_map_location", - [ - f"{DEVICE_TYPE}:0", - # {"": "cuda"}, - ], -) -def test_opt_125m_int8wo_model_loading_with_params( - pt_load_map_location, monkeypatch, dist_init, workspace_init -): - torch._dynamo.reset() - model_name = "jerryzh168/opt-125m-int8wo-partial-quant" - model, vllm_config = load_model_without_vllm_runner( - model_name, - quantization="torchao", - vllm_config_kwargs={ - "load_config": LoadConfig(pt_load_map_location=pt_load_map_location) - }, - ) - - monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) - input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE) - positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE) - with ( - set_current_vllm_config(vllm_config), - set_forward_context(None, vllm_config, num_tokens=input_ids.numel()), - ): - hidden_states = model(input_ids, positions, None) - logits = model.compute_logits(hidden_states) - assert torch.isfinite(logits).all() - - -@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") -def test_qwenvl_int8wo_model_loading_with_params( - monkeypatch, dist_init, workspace_init -): - torch._dynamo.reset() - model_name = "mobicham/Qwen2.5-VL-3B-Instruct_int8wo_ao" - model, vllm_config = load_model_without_vllm_runner( - model_name, - quantization="torchao", - vllm_config_kwargs={ - "load_config": LoadConfig(pt_load_map_location=f"{DEVICE_TYPE}:0") - }, - ) - - monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) - input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE) - positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE).expand(3, -1) - with ( - set_current_vllm_config(vllm_config), - set_forward_context(None, vllm_config, num_tokens=input_ids.numel()), - ): - hidden_states = model(input_ids, positions, None) - logits = model.compute_logits(hidden_states) - assert torch.isfinite(logits).all() - - @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") @pytest.mark.skip( reason="since torchao nightly is only compatible with torch nightly" From 91b96a533cfb06579a360d3aefb2c8341fc89bdb Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:55:14 -0500 Subject: [PATCH 354/774] [ROCm][CI] Shard MI300 Multimodal Processor (#57056) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 78130fa3c719..eadfa36e04fd 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2175,19 +2175,20 @@ steps: - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model -- label: ":amd: (MI300) Multimodal Processor" # 1h 42m +- label: ":amd: (MI300) Multimodal Processor Shard %N" # TBD timeout_in_minutes: 110 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 optional: true + parallelism: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - tests/models/multimodal - tests/models/registry.py commands: - - pytest -v -s models/multimodal/processing/test_tensor_schema.py + - pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI300) Multimodal Models (Extended Pooling)" # TBD timeout_in_minutes: 75 From c58ff535f4a2f30b61a86c846ac4adefe62abb8d Mon Sep 17 00:00:00 2001 From: kyleliang-nv Date: Wed, 16 Sep 2026 17:03:36 -0700 Subject: [PATCH 355/774] [Bugfix][Spec Decode] Fix EAGLE and dense draft startup with EP (#56930) Signed-off-by: Kyle Liang Co-authored-by: Kyle Liang Co-authored-by: OpenAI Codex --- tests/test_config.py | 109 +++++++++++++++++++++++++++++++++++++ vllm/config/speculative.py | 16 +++++- 2 files changed, 122 insertions(+), 3 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index ce2580934a01..3d6a17f7a29b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1520,6 +1520,115 @@ def test_draft_model_enables_async_scheduling_by_default(): assert cfg.scheduler_config.async_scheduling is True +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"]) +@pytest.mark.parametrize( + ("method", "draft_is_moe"), + [ + pytest.param("draft_model", False, id="dense-draft"), + pytest.param("eagle", False, id="eagle"), + pytest.param("eagle3", False, id="eagle3"), + pytest.param("draft_model", True, id="moe-draft"), + pytest.param("mtp", True, id="mtp"), + pytest.param("dspark", True, id="dspark"), + ], +) +def test_draft_inherits_ep_only_for_moe( + tmp_path: Path, + tp_size: int, + target_ep: bool, + method: str, + draft_is_moe: bool, +): + """Validate final draft configs without loading weights or mocking validation.""" + from transformers import LlamaConfig, MixtralConfig + + from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config + + common = dict( + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + max_position_embeddings=2048, + ) + if method in ("mtp", "dspark"): + target_hf_config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + n_routed_experts=4, + num_experts_per_tok=2, + num_nextn_predict_layers=1, + compress_ratios=[1, 1], + head_dim=32, + **common, + ) + else: + target_hf_config = MixtralConfig( + architectures=["MixtralForCausalLM"], num_local_experts=4, **common + ) + target_path = tmp_path / "target" + draft_path = tmp_path / "draft" + _write_json(target_path / "config.json", target_hf_config.to_dict()) + _write_json( + draft_path / "config.json", + LlamaConfig(architectures=["LlamaForCausalLM"], **common).to_dict(), + ) + target_model_config = ModelConfig( + model=str(target_path), tokenizer_mode="skip", max_model_len=2048 + ) + target_parallel_config = ParallelConfig( + tensor_parallel_size=tp_size, + enable_expert_parallel=target_ep, + distributed_executor_backend="mp", + ) + target_model_config.verify_with_parallel_config(target_parallel_config) + speculative_config = SpeculativeConfig( + method=method, + model=str(target_path if draft_is_moe else draft_path), + num_speculative_tokens=1, + target_model_config=target_model_config, + target_parallel_config=target_parallel_config, + ) + + assert speculative_config.method == method + assert speculative_config.draft_model_config.is_moe is draft_is_moe + assert speculative_config.draft_parallel_config.enable_expert_parallel is ( + target_ep and draft_is_moe + ) + assert speculative_config.draft_parallel_config.tensor_parallel_size == tp_size + assert target_parallel_config.enable_expert_parallel is target_ep + assert target_parallel_config.tensor_parallel_size == tp_size + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"]) +@pytest.mark.parametrize("pass_none", [False, True], ids=["omitted", "explicit-none"]) +def test_draft_parallel_config_preserves_ep_without_model( + target_ep: bool, pass_none: bool +): + """Legacy callers without draft model information keep EP inheritance.""" + target_parallel_config = ParallelConfig( + tensor_parallel_size=2, + enable_expert_parallel=target_ep, + distributed_executor_backend="mp", + ) + if pass_none: + draft_parallel_config = SpeculativeConfig.create_draft_parallel_config( + target_parallel_config, 2, draft_model_config=None + ) + else: + draft_parallel_config = SpeculativeConfig.create_draft_parallel_config( + target_parallel_config, 2 + ) + + assert draft_parallel_config.enable_expert_parallel is target_ep + assert draft_parallel_config.tensor_parallel_size == 2 + assert target_parallel_config.enable_expert_parallel is target_ep + + @pytest.mark.parametrize( ("method", "parallel_drafting", "expected_slots"), [ diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 63f729a832c9..ef9a16f903a5 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1545,9 +1545,13 @@ def __post_init__(self): ) ) + # Use the final draft config, after any architecture overrides, + # so EP agrees with ModelConfig's expert-count validation. self.draft_parallel_config = ( SpeculativeConfig.create_draft_parallel_config( - self.target_parallel_config, self.draft_tensor_parallel_size + self.target_parallel_config, + self.draft_tensor_parallel_size, + draft_model_config=self.draft_model_config, ) ) @@ -1743,15 +1747,21 @@ def update_arch_(self): def create_draft_parallel_config( target_parallel_config: ParallelConfig, speculative_draft_tensor_parallel_size: int, + draft_model_config: ModelConfig | None = None, ) -> ParallelConfig: """Create a parallel config for use by the draft worker. - This is mostly a copy of the target parallel config, except the tp_size. + Use the draft TP size and disable inherited EP for known dense drafts. + Without a draft model config, preserve the previous EP inheritance. """ + enable_ep = target_parallel_config.enable_expert_parallel + if draft_model_config is not None: + enable_ep = enable_ep and draft_model_config.is_moe + draft_parallel_config = ParallelConfig( pipeline_parallel_size=1, tensor_parallel_size=speculative_draft_tensor_parallel_size, - enable_expert_parallel=target_parallel_config.enable_expert_parallel, + enable_expert_parallel=enable_ep, distributed_executor_backend=target_parallel_config.distributed_executor_backend, max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce, From a615f53364c109ade2b6bb69248e16ed40405d1f Mon Sep 17 00:00:00 2001 From: Thillai Chithambaram <79466435+thillai-c@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:23:28 -0400 Subject: [PATCH 356/774] [Bugfix][Metrics][MFU] Size activation traffic from the model dtype (#57070) Signed-off-by: Thillai Chithambaram --- tests/v1/metrics/test_perf_metrics.py | 51 ++++++++++++++++++++++++++- vllm/v1/metrics/perf.py | 7 ++-- 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index 3868cc772e96..c0997ec459bf 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -133,7 +133,56 @@ def test_base_config_parser(): assert result.num_attention_heads == 16 assert result.num_hidden_layers == 24 assert result.weight_byte_size == 2 # float16 is 2 bytes - assert result.activation_byte_size == 2 # default activation size + assert result.activation_byte_size == 2 # float16 activations are 2 bytes + + +@pytest.mark.parametrize( + "model_dtype,expected_byte_size", + [("float16", 2), ("bfloat16", 2), ("float32", 4)], +) +def test_base_config_parser_activation_byte_size_follows_dtype( + model_dtype: str, expected_byte_size: int +): + """Activation traffic is sized from the model's compute dtype. + + Activations are produced in the model dtype, so pinning them to 2 bytes + under-reported memory traffic by half for a float32 model while that same + model's weight traffic was sized correctly. + """ + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=24, + ) + vllm_config = create_mock_vllm_config(hf_config, model_dtype=model_dtype) + + result = BaseConfigParser().parse(ParsedArgs(), vllm_config) + + assert result.activation_byte_size == expected_byte_size + assert result.weight_byte_size == expected_byte_size + + +def test_attention_read_bytes_scale_with_activation_dtype(): + """A float32 model reads twice the activation bytes of a bfloat16 one.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=4, + ) + ctx = ExecutionContext.from_single_request( + num_tokens=128, context_len=128, is_prefill=True + ) + + bf16 = AttentionMetrics.from_vllm_config( + create_mock_vllm_config(hf_config, model_dtype="bfloat16") + ).get_read_bytes_breakdown(ctx) + fp32 = AttentionMetrics.from_vllm_config( + create_mock_vllm_config(hf_config, model_dtype="float32") + ).get_read_bytes_breakdown(ctx) + + # qkv_input is pure activation traffic, so it doubles exactly. + assert fp32["qkv_input"] == 2 * bf16["qkv_input"] def test_base_attention_config_parser_with_gqa(): diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index dd046a8250b6..1cd8fa79f1d3 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -369,9 +369,10 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: args.weight_byte_size = get_dtype_size(torch_dtype) - # FIXME: handle this better by parsing whether activations use - # bf16, fp32, etc... - args.activation_byte_size = 2 + # Activations are produced in the model's compute dtype. Quantization + # overrides weight_byte_size below but leaves activations alone, so + # this stays keyed on the model dtype. + args.activation_byte_size = get_dtype_size(torch_dtype) args.dp_size = vllm_config.parallel_config.data_parallel_size args.tp_size = vllm_config.parallel_config.tensor_parallel_size From 95f4925c3a03df8cfcaa21633ccc9dd5b426c7a4 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:53:06 -0400 Subject: [PATCH 357/774] [ROCm][CI] Enable AITER FP8/unquantized cases in modular-kernel sweep + fix MoRI per-tensor FP8 dispatch (#57074) Signed-off-by: Divakar Verma Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 3 +- .../moe/modular_kernel_tools/common.py | 56 ++++++++------ .../moe/modular_kernel_tools/mk_objects.py | 18 +++-- .../moe/test_modular_kernel_combinations.py | 75 +++++++++++-------- .../layers/fused_moe/all2all_utils.py | 15 +++- .../layers/fused_moe/prepare_finalize/mori.py | 9 +++ 6 files changed, 107 insertions(+), 69 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index eadfa36e04fd..4b72256a3fcd 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1704,8 +1704,9 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) Quantization Kernels" # TBD timeout_in_minutes: 55 diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 8adfa6588ca8..3ec522c8ec61 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -15,6 +15,7 @@ dequantize_nvfp4_to_dtype, ) from tests.kernels.utils import torch_experts +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import ( get_dp_group, @@ -33,7 +34,9 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, kFp8Dynamic128Sym, kFp8DynamicTensorSym, kFp8DynamicTokenSym, @@ -165,31 +168,34 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: return vllm_config, env_dict + def fp8_quant_key_pair(self) -> tuple[QuantKey, QuantKey]: + """Derive the (weight_quant_key, activation_quant_key) pair an FP8 + quant config of this shape corresponds to (either OCP or FNUZ FP8, + see ``current_platform.fp8_dtype()``).""" + if self.quant_block_shape is not None: + return kFp8Static128BlockSym, kFp8Dynamic128Sym + if self.is_per_out_ch_quant: + return ( + kFp8StaticChannelSym, + kFp8DynamicTokenSym + if self.is_per_act_token_quant + else kFp8StaticTensorSym, + ) + return ( + kFp8StaticTensorSym, + kFp8DynamicTensorSym + if self.is_per_act_token_quant + else kFp8StaticTensorSym, + ) + def fe_supports_quant_scheme(self) -> bool: """Check if the fused experts class supports this quant config. See https://github.com/ROCm/aiter/issues/2419 for AITER gaps.""" if self.quant_config is None or self.quant_dtype is None: return True - if self.quant_dtype != torch.float8_e4m3fn: + if not is_fp8(self.quant_dtype): return True - # Derive QuantKeys from test config - if self.quant_block_shape is not None: - w_key = kFp8Static128BlockSym - a_key = kFp8Dynamic128Sym - elif self.is_per_out_ch_quant: - w_key = kFp8StaticChannelSym - a_key = ( - kFp8DynamicTokenSym - if self.is_per_act_token_quant - else kFp8StaticTensorSym - ) - else: - w_key = kFp8StaticTensorSym - a_key = ( - kFp8DynamicTensorSym - if self.is_per_act_token_quant - else kFp8StaticTensorSym - ) + w_key, a_key = self.fp8_quant_key_pair() fe_cls = self.fused_experts_type if hasattr(fe_cls, "_supports_quant_scheme"): try: @@ -199,10 +205,7 @@ def fe_supports_quant_scheme(self) -> bool: return True def is_fp8_block_quantized(self): - return ( - self.quant_dtype == torch.float8_e4m3fn - and self.quant_block_shape is not None - ) + return is_fp8(self.quant_dtype) and self.quant_block_shape is not None def is_batched_prepare_finalize(self): info = prepare_finalize_info(self.prepare_finalize_type) @@ -334,6 +337,11 @@ def is_valid(self) -> tuple[bool, str | None]: return False, "Needs Aiter, but Aiter not available." if self.needs_mori() and not has_mori(): # noqa: SIM103 return False, "Needs MoRI, but MoRI not available." + if self.needs_mori() and not rocm_aiter_ops.is_fused_moe_enabled(): + return False, ( + "Mori requires AITER's fused-moe backend to be enabled " + "(VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1)." + ) try: if not self.fused_experts_type._supports_current_device(): @@ -370,7 +378,7 @@ def describe(self): def is_quantized(self) -> bool: # or w1_scale is not None? return ( - self.w1.dtype == torch.float8_e4m3fn + is_fp8(self.w1.dtype) or self.w1.dtype == torch.uint8 or self.w1.dtype == torch.int8 ) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 7f1924bac5ab..e0f37f54360a 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -86,14 +86,14 @@ class ExpertInfo: standard_format = mk.FusedMoEActivationFormat.Standard batched_format = mk.FusedMoEActivationFormat.BatchedExperts common_float_types: list[torch.dtype | str] = [ - torch.float8_e4m3fn, + current_platform.fp8_dtype(), torch.bfloat16, torch.float16, torch.float32, ] common_float_and_int_types = common_float_types + [torch.int8] nvfp4_types = ["nvfp4"] -fp8_types = [torch.float8_e4m3fn] +fp8_types = [current_platform.fp8_dtype()] def register_prepare_and_finalize( @@ -324,7 +324,9 @@ def expert_info(kind) -> ExpertInfo: register_experts( AiterExperts, standard_format, - fp8_types, + # AiterExperts also supports the fully-unquantized (None, None) + # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8. + common_float_types, blocked_quantization_support=True, needs_aiter=True, ) @@ -397,35 +399,35 @@ def expert_info(kind) -> ExpertInfo: None, # per-channel / per-column weights and per-tensor activations TestMoEQuantConfig( - quant_dtype=torch.float8_e4m3fn, + quant_dtype=current_platform.fp8_dtype(), per_out_ch_quant=True, per_act_token_quant=False, block_shape=None, ), # per-channel / per-column weights and per-token activations TestMoEQuantConfig( - quant_dtype=torch.float8_e4m3fn, + quant_dtype=current_platform.fp8_dtype(), per_out_ch_quant=True, per_act_token_quant=True, block_shape=None, ), # per-tensor weights and per-tensor activations TestMoEQuantConfig( - quant_dtype=torch.float8_e4m3fn, + quant_dtype=current_platform.fp8_dtype(), per_out_ch_quant=False, per_act_token_quant=False, block_shape=None, ), # per-tensor weights and per-token activations TestMoEQuantConfig( - quant_dtype=torch.float8_e4m3fn, + quant_dtype=current_platform.fp8_dtype(), per_out_ch_quant=False, per_act_token_quant=True, block_shape=None, ), # block-quantized weights and 128 block per-token activations TestMoEQuantConfig( - quant_dtype=torch.float8_e4m3fn, + quant_dtype=current_platform.fp8_dtype(), per_out_ch_quant=False, per_act_token_quant=False, block_shape=[128, 128], diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 0c0e1d61f90e..fc1d8ac07bcf 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -37,6 +37,7 @@ ProcessGroupInfo, parallel_launch_with_config, ) +from .utils import check_accuracy has_any_multi_gpu_package = ( has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe() @@ -47,12 +48,6 @@ reason="Requires deep_ep or deep_gemm or flashinfer packages", ) -if current_platform.is_fp8_fnuz(): - pytest.skip( - "Tests in this file require float8_e4m3fn and platform does not support", - allow_module_level=True, - ) - def format_result(verbose, msg, ex=None): if ex is not None: @@ -70,6 +65,28 @@ def format_result(verbose, msg, ex=None): print(".", end="") +def assert_aiter_quant_scheme_case(config: Config) -> None: + """Make the AITER (weight_quant_key, activation_quant_key) pair this + config exercises explicit, instead of AiterExperts being reached only + indirectly through the general quant-config sweep. + See https://github.com/vllm-project/vllm/issues/54966.""" + fe_cls = config.fused_experts_type + if getattr(fe_cls, "__name__", "") != "AiterExperts": + return + + if config.quant_config is None: + w_key, a_key = None, None + else: + w_key, a_key = config.fp8_quant_key_pair() + + assert fe_cls._supports_quant_scheme(w_key, a_key), ( + f"AITER case (weight_key={w_key}, activation_key={a_key}) reached " + "the modular-kernel harness, but AiterExperts._supports_quant_scheme " + "does not declare it supported." + ) + print(f"[AITER case] weight_key={w_key}, activation_key={a_key}") + + def rank_worker( pgi: ProcessGroupInfo, vllm_config: VllmConfig, @@ -125,6 +142,25 @@ def rank_worker( count -= 1 continue + # Skip unsupported: AITER x DeepEP-HT/Mori dispatch crashes + # with an illegal memory access at world_size>1 on gfx942. + # https://github.com/vllm-project/vllm/issues/57029 + if ( + config.world_size > 1 + and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts" + and getattr(config.prepare_finalize_type, "__name__", "") + in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize") + ): + print( + f"Skipping[{pgi.rank}]: m={m}, topk={topk}" + " (AITER x DeepEP-HT/Mori illegal memory access," + " https://github.com/vllm-project/vllm/issues/57029)" + ) + count -= 1 + continue + + assert_aiter_quant_scheme_case(config) + # modular kernel out mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors) @@ -152,32 +188,7 @@ def rank_worker( and config.quant_config is not None ) if is_aiter_fp8: - diff = (ref_out - mk_out).abs() - n_total = diff.numel() - max_diff = diff.max().item() - n_exceed = int((diff > atol).sum().item()) - pct_exceed = n_exceed / n_total * 100 - # FP8 hw matmul vs f32 reference: up to ~4% of - # elements may exceed base tolerance, but max - # error should stay within 3x base tolerance. - max_pct_allowed = 5.0 - relaxed_atol = atol * 4 - print( - f"[AITER FP8 precision] " - f"max_diff={max_diff:.6f}, " - f"exceed_atol={n_exceed}/{n_total} " - f"({pct_exceed:.4f}%), " - f"max_pct_allowed={max_pct_allowed}%, " - f"relaxed_limit={relaxed_atol}" - ) - assert pct_exceed <= max_pct_allowed, ( - f"AITER FP8: {pct_exceed:.2f}% elements exceed " - f"atol={atol} (max allowed {max_pct_allowed}%)" - ) - assert max_diff <= relaxed_atol, ( - f"AITER FP8: max_diff={max_diff:.6f} exceeds " - f"relaxed limit {relaxed_atol}" - ) + check_accuracy(ref_out, mk_out, atol=atol, rtol=rtol, percent=0.9) else: torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol) format_result(verbose, config.describe()) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 75ea91bfca5d..6984e34e56fd 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -284,13 +284,20 @@ def maybe_make_prepare_finalize( # Note: We may want to use FP8 dispatch just to reduce # data movement. use_fp8_dispatch = ( - quant_config.is_per_act_token or quant_config.is_block_quantized + quant_config.is_per_act_token + or quant_config.is_block_quantized + or quant_config.is_per_tensor ) if use_fp8_dispatch: - # For PTPC (per token per channel) quant, scale dim is 1 - # For 1x128 quant, scale dim is hidden_dim // 128 + # For PTPC (per token per channel) or per-tensor quant, + # scale dim is 1. For 1x128 quant, scale dim is + # hidden_dim // 128 quant_dtype = quant_config.quant_dtype - scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128 + scale_dim = ( + 1 + if (quant_config.is_per_act_token or quant_config.is_per_tensor) + else moe.hidden_dim // 128 + ) else: # Unquantized dispatch (e.g. AITER with defer_input_quant): # dispatch raw BF16/FP16 data, no scales needed. diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index 35e15325b4db..75b97c69aafb 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py @@ -82,6 +82,15 @@ def prepare( elif quant_config.is_per_act_token: quant_func = get_hip_quant(QuantType.per_Token) a1, scale = quant_func(a1, quant_dtype=current_platform.fp8_dtype()) + elif quant_config.is_per_tensor: + quant_func = get_hip_quant(QuantType.per_Tensor) + a1, scale = quant_func( + a1, + scale=quant_config.a1_scale, + quant_dtype=current_platform.fp8_dtype(), + ) + # mori expects one scale slot per token; broadcast. + scale = scale.expand(a1.shape[0], 1).contiguous() ( dispatch_a1, From 3e267bae70d85785ede4ae724f87f89edd6905a9 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Wed, 16 Sep 2026 20:03:34 -0500 Subject: [PATCH 358/774] [ROCm] triton+triton_kernels 3.8 mxfp4 MoE support (gpt-oss + DeepSeek-V4) (#55934) Signed-off-by: Rohan Potdar Signed-off-by: Rohan138 Co-authored-by: Claude Opus 4.8 Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com> --- docker/Dockerfile.rock | 8 +- .../moe/test_gpt_oss_triton_kernels.py | 3 +- .../model_executor/layers/fused_moe/config.py | 11 +- .../experts/aiter_mxfp4_w4a16_moe.py | 31 +-- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 65 +----- .../experts/gpt_oss_triton_kernels_moe.py | 214 ++++++++++++++---- .../layers/fused_moe/oracle/mxfp4.py | 37 +-- .../layers/quantization/utils/mxfp4_utils.py | 28 ++- vllm/utils/import_utils.py | 20 ++ 9 files changed, 270 insertions(+), 147 deletions(-) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index 14ba523b072f..32623ff1a90b 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -5,7 +5,7 @@ ARG VLLM_BRANCH="main" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base -ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7 +ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 ARG INSTALL_LMCACHE=false # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types @@ -322,8 +322,9 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && export CCACHE_BASEDIR="$PWD" \ && echo "=== ccache stats before ROCm native build ===" \ && (ccache --show-stats || true) \ @@ -669,8 +670,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Install the packaging helpers before modifying requirements/rocm.txt. RUN --mount=type=cache,target=/root/.cache/uv \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ vllm/cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && cd vllm \ && uv pip install --system setuptools_scm regex diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 4afcba706ccd..86028a7feba0 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -384,9 +384,8 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): make_routing_data = gptoss_moe.make_routing_data routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk - use_legacy_triton_kernels = gptoss_moe.use_legacy_triton_kernels - if use_legacy_triton_kernels: + if gptoss_moe.triton_kernels_version == "3.5.1": pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+") from triton_kernels.topk import topk as topk_fn diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index ccca24cc6f66..9a2f8229b1b6 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -17,14 +17,19 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version from vllm.utils.math_utils import cdiv logger = init_logger(__name__) -if has_triton_kernels(): +_triton_kernels_version = get_triton_kernels_version() + +if _triton_kernels_version is not None: try: - from triton_kernels.matmul_ogs import PrecisionConfig + if _triton_kernels_version == "3.8": + from triton_kernels.matmul import PrecisionConfig + else: + from triton_kernels.matmul_ogs import PrecisionConfig except (ImportError, AttributeError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index 757d1d8295f7..f6afd6fbfa34 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -12,9 +12,6 @@ FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( - patch_gating_output, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -26,17 +23,6 @@ ] -def _aiter_raw(t): - """Unwrap the `triton_kernels.tensor.Tensor` that `_swizzle_mxfp4` returns.""" - if t is None or isinstance(t, torch.Tensor): - return t - assert hasattr(t, "storage"), ( - f"expected a triton_kernels wrapped tensor with a .storage attribute, " - f"got {type(t)}" - ) - return t.storage.data - - def _aiter_w4a16_silu_via_a8w4( hidden_states: torch.Tensor, w1_data, @@ -168,11 +154,6 @@ def aiter_triton_kernel_w4a16_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - # See context in #50859. - # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged, - # AITER released, and AITER pin in vLLM increased from 0.1.19. - gating_output = patch_gating_output(gating_output, global_num_experts) - if score_mode is not None: use_grouped_topk = num_expert_group is not None and num_expert_group > 1 routing_data, gather_idx, scatter_idx = aiter_routing( @@ -201,10 +182,14 @@ def aiter_triton_kernel_w4a16_moe_forward( assert quant_config.w1_precision is not None assert quant_config.w2_precision is not None - w1_data = _aiter_raw(w1) - w2_data = _aiter_raw(w2) - w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) - w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + w1_data = w1.storage.data + w2_data = w2.storage.data + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + weight_mx_scale, + ) + + w1_wscale = weight_mx_scale(quant_config.w1_precision).storage.data + w2_wscale = weight_mx_scale(quant_config.w2_precision).storage.data gammas = routing_data.gate_scal if routing_data else None diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 0691d9fed4f2..b327cbeb594c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -5,7 +5,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops -from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -25,60 +24,6 @@ ] -# TODO(akaratza): Remove -# _get_padding_mask and patch_gating_output after -# https://github.com/ROCm/aiter/pull/4530 is released and consumed by vLLM. -def _get_padding_mask() -> torch.Tensor | None: - """Retrieves a boolean mask with non-padding (0) and padding (1) tokens. - - `slot_mapping < 0` comes from: - - slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) - - in gpu_model_runner.py. - - Direct kernel callers do not have model-runner padding metadata, so no - mask is needed when there is no forward context. - """ - if not is_forward_context_available(): - return None - - forward_context = get_forward_context() - - # model runner v2. - if forward_context.is_padding is not None: - return forward_context.is_padding - - # model runner v1. - slot_mapping = forward_context.slot_mapping - - if isinstance(slot_mapping, list): - slot_mapping_dict = slot_mapping[0] - else: - slot_mapping_dict = slot_mapping - - if isinstance(slot_mapping_dict, dict): - slot_mapping_sample = next(iter(slot_mapping_dict.values()), None) - - if isinstance(slot_mapping_sample, torch.Tensor): - return slot_mapping_sample < 0 - else: - return None - - -def patch_gating_output( - gating_output: torch.Tensor, global_num_experts: int -) -> torch.Tensor: - if global_num_experts != 128: - is_padding = _get_padding_mask() - - if is_padding is not None: - assert is_padding.shape[0] == gating_output.shape[0] - gating_output = gating_output.masked_fill(is_padding[:, None], 0.0) - - return gating_output - - def aiter_triton_kernel_w4a8_moe_forward( hidden_states: torch.Tensor, w1, # Tensor or triton_kernels.Tensor @@ -112,11 +57,6 @@ def aiter_triton_kernel_w4a8_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - # See context in #50859. - # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged, - # AITER released, and AITER pin in vLLM increased from 0.1.19. - gating_output = patch_gating_output(gating_output, global_num_experts) - routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) @@ -191,6 +131,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( should_use_cdna4_mx_scale_swizzle, + weight_mx_scale, ) _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None @@ -210,7 +151,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( hidden_states, w1.storage.data, None, - quant_config.w1_precision.weight_scale.storage.data, + weight_mx_scale(quant_config.w1_precision).storage.data, quant_config.w1_precision.flex_ctx.lhs_data.scale, quant_config.w2_precision.flex_ctx.lhs_data.scale, quant_config.w1_bias, @@ -230,7 +171,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( intermediate_cache1, w2.storage.data, None, - quant_config.w2_precision.weight_scale.storage.data, + weight_mx_scale(quant_config.w2_precision).storage.data, quant_config.w2_precision.flex_ctx.lhs_data.scale, None, quant_config.w2_bias, diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index e2b4987bd050..27fdcf8bb2c6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -26,7 +26,7 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels from ..utils import swiglu_limit_func @@ -441,48 +441,59 @@ def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix): _routing.sort_tokens = _sort_tokens_pow2 -# Two API generations of triton_kernels are supported: -# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and -# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` -# constructor takes a `scratchpad` argument. -# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix` -# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only -# when the user has triton_kernels installed system-wide at v3.6.0+. -# -# `use_legacy_triton_kernels` selects between them at import time based on -# whether `SparseMatrix` is importable. -use_legacy_triton_kernels = False +triton_kernels_version = get_triton_kernels_version() -if has_triton_kernels(): +if triton_kernels_version is not None: try: import triton_kernels.swiglu - from triton_kernels.matmul_ogs import ( - FnSpecs, - FusedActivation, - GatherIndx, - RoutingData, - ScatterIndx, - matmul_ogs, - ) - from triton_kernels.tensor import ( - BIT, - Bitmatrix, - ) + if triton_kernels_version == "3.8": + from triton_kernels.matmul import FnSpecs, FusedActivation, matmul + from triton_kernels.matmul_details.opt_flags import ( + update_opt_flags_constraints, + ) + from triton_kernels.tensor import make_ragged_tensor_metadata + + from vllm.platforms.rocm import on_gfx950 + + if on_gfx950(): + # block_n=64/block_k=256 fit gfx950's 160KB LDS + the CDNA4 + # scale swizzle; block_m=32 avoids an autotuner block_m=128 + # padding cliff (~15x) at decode-sized ragged MoE shapes. + update_opt_flags_constraints( + {"block_m": 32, "block_n": 64, "block_k": 256} + ) + else: + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + GatherIndx, + RoutingData, + ScatterIndx, + matmul_ogs, + ) + + if triton_kernels_version == "3.6": + from triton_kernels.tensor import ( + SparseMatrix, + make_ragged_tensor_metadata, + ) + + # BIT/Bitmatrix exist in v3.5.1/v3.6, gone in v3.8. try: from triton_kernels.tensor import ( - SparseMatrix, - make_ragged_tensor_metadata, + BIT, + Bitmatrix, ) except ImportError: - # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once - # the gpt-oss perf regression in v3.6.0+ is resolved upstream. - # Tracking: https://github.com/triton-lang/triton/issues/9969 - use_legacy_triton_kernels = True - if not use_legacy_triton_kernels: + BIT = None + Bitmatrix = None + + # Patch the routing kernels for non-pow2 top_k (DeepSeek-V4 top_k=6); + # v3.8 rebuilds routing from topk ids and needs no patch. + if triton_kernels_version == "3.6": _patch_make_bitmatrix_metadata() - else: - # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4). + elif triton_kernels_version == "3.5.1": _patch_legacy_routing_for_nonpow2_topk() except (AttributeError, ImportError) as e: logger.error( @@ -492,6 +503,21 @@ def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix): ) +if triton_kernels_version == "3.8": + # 3.8's matmul dropped matmul_ogs's in-kernel scatter+topk-reduce, so the + # forward combines externally via index_add; this is the ragged routing + # bundle the (inlined) 3.8 paths pass around, in place of matmul_ogs's. + class RoutingData: # type: ignore[no-redef] + __slots__ = ("ragged", "gate_scal", "gather_tok", "n_expts_act", "n_valid") + + def __init__(self, ragged, gate_scal, gather_tok, n_expts_act, n_valid): + self.ragged = ragged + self.gate_scal = gate_scal + self.gather_tok = gather_tok + self.n_expts_act = n_expts_act + self.n_valid = n_valid + + @triton.jit def pack_bitmatrix( bitmatrix, @@ -562,7 +588,7 @@ def triton_kernel_moe_forward( # three separate kernels used by the generic path below. # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines # equivalent logic via SparseMatrix in `make_routing_data`. - if use_legacy_triton_kernels and expert_map is None: + if triton_kernels_version == "3.5.1" and expert_map is None: from triton_kernels.routing import routing as fused_routing routing_data, gather_idx, scatter_idx = fused_routing( @@ -662,6 +688,52 @@ def triton_kernel_fused_experts( assert activation == MoEActivation.SWIGLUOAI, ( "Only SWIGLUOAI activation is supported" ) + + if triton_kernels_version == "3.8": + # 3.8 has no in-kernel scatter+topk-reduce, so the two matmuls are + # followed by an external index_add combine. + assert quant_config is not None + M, K = hidden_states.shape[-2:] + a_ragged = routing_data.ragged + gather_tok = routing_data.gather_tok + gammas = routing_data.gate_scal + act = FusedActivation( + FnSpecs( + "swiglu", + triton_kernels.swiglu.swiglu_fn, + ("alpha", "limit"), + reduction_n=2, + ), + (swiglu_alpha, swiglu_limit), + ) + inter = matmul( + hidden_states, + w1, + quant_config.w1_bias, + a_ragged, + None, + gather_tok, + None, + quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + ) + y = matmul( + inter, + w2, + quant_config.w2_bias, + a_ragged, + None, + None, + None, + quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + ) + acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device) + acc.index_add_(0, gather_tok.to(torch.int64), y.to(torch.float32)) + out = output_tensor.view(M, K) + out.copy_(acc.to(out.dtype)) + return out assert quant_config is not None # type check, uint8 means mxfp4 @@ -704,7 +776,7 @@ def triton_kernel_fused_experts( ), (swiglu_alpha, swiglu_limit), ) - if not use_legacy_triton_kernels + if triton_kernels_version != "3.5.1" else FusedActivation( FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), (swiglu_alpha, swiglu_limit), @@ -744,6 +816,24 @@ def make_routing_data( topk_weights: torch.Tensor, num_local_experts: int, ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + if triton_kernels_version == "3.8": + _, topk = topk_ids.shape + flat_e = topk_ids.reshape(-1) + valid = flat_e >= 0 + flat_e = flat_e.to(torch.int64).clamp(min=0) + flat_w = topk_weights.reshape(-1).to(torch.float32) * valid.to(torch.float32) + nrows = flat_e.shape[0] + order = torch.argsort(flat_e, stable=True) + col_sum = torch.zeros( + num_local_experts, dtype=torch.int32, device=flat_e.device + ) + col_sum.scatter_add_(0, flat_e, torch.ones_like(flat_e, dtype=torch.int32)) + ragged = make_ragged_tensor_metadata(col_sum, nrows) + gather_tok = (order // topk).to(torch.int32) + gammas = flat_w[order] + routing = RoutingData(ragged, gammas, gather_tok, topk, nrows) + return routing, gather_tok, gather_tok + topk_ids = topk_ids.to(torch.int16) topk_weights = topk_weights.to(torch.bfloat16) @@ -774,7 +864,7 @@ def make_routing_data( Bitmatrix( bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max ) - if not use_legacy_triton_kernels + if triton_kernels_version != "3.5.1" else Bitmatrix( bitmatrix, shape=bitmatrix_shape, @@ -786,7 +876,7 @@ def make_routing_data( # matmul_ogs expects invalid topk_weights to be -1s topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) - if use_legacy_triton_kernels: + if triton_kernels_version == "3.5.1": from triton_kernels.routing import routing_from_bitmatrix return routing_from_bitmatrix( @@ -809,6 +899,10 @@ def routing_data_from_sparse_topk( call already produced; when the topk ids are used as-is (no expert-map remapping) that work is redundant. """ + if triton_kernels_version == "3.8": + # 3.8 routing is rebuilt from raw ids; no SparseMatrix reuse. + return make_routing_data(sparse_topk.indx, sparse_topk.vals, n_expts) + dispatch_indx = sparse_topk.mask_metadata.row_sorted_indx combine_indx = sparse_topk.mask_metadata.col_sorted_indx ragged_batch_metadata = make_ragged_tensor_metadata( @@ -1197,6 +1291,50 @@ def apply( topk = topk_ids.size(1) + if triton_kernels_version == "3.8": + # Activation and topk-reduce kept unfused from the matmuls (3.8 has + # no in-kernel scatter+reduce); combine via external index_add. + assert self._lora_context is None, ( + "tk38 unfused MoE path does not support LoRA yet" + ) + act_out_dim = self.adjust_N_for_activation(w1.shape[2], activation) + M, K = hidden_states.shape[-2:] + a_ragged = routing_data.ragged + gather_tok = routing_data.gather_tok + gammas = routing_data.gate_scal + inter = matmul( + hidden_states, + w1, + quant_config.w1_bias, + a_ragged, + None, + gather_tok, + None, + quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + ) + act_out = torch.empty( + (routing_data.n_valid, act_out_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self.activation(activation, act_out, inter) + down = matmul( + act_out, + w2, + quant_config.w2_bias, + a_ragged, + None, + None, + None, + quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + ) + acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device) + acc.index_add_(0, gather_tok.to(torch.int64), down.to(torch.float32)) + output.view(M, K).copy_(acc.to(output.dtype)) + return + # type check, uint8 means mxfp4 assert hidden_states.dtype == torch.bfloat16 assert ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 912b713918ef..ce19ba39bfc9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -45,7 +45,7 @@ ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d from vllm.platforms import current_platform -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version from vllm.utils.math_utils import round_up if TYPE_CHECKING: @@ -54,9 +54,14 @@ logger = init_logger(__name__) -if has_triton_kernels(): +triton_kernels_version = get_triton_kernels_version() + +if triton_kernels_version is not None: try: - from triton_kernels.matmul_ogs import PrecisionConfig + if triton_kernels_version == "3.8": + from triton_kernels.matmul import FlexCtx, PrecisionConfig + else: + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig except (ImportError, AttributeError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " @@ -65,6 +70,15 @@ ) +def _mx_scale_kwargs(scale): + """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size, + 3.5.1/3.6 use weight_scale. + """ + if triton_kernels_version == "3.8": + return {"b_mx_scale": scale, "b_microblock_size": 32} + return {"weight_scale": scale} + + def _pack_deepgemm_mxfp4_scales( w13_weight: torch.Tensor, w2_weight: torch.Tensor, @@ -1182,7 +1196,6 @@ def swap_every_two_rows(x, axis=-1): elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: # W4A8: MXFP4 weights + static FP8 activations (triton kernel) - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig from triton_kernels.numerics import InFlexData if w13_bias is not None: @@ -1217,11 +1230,11 @@ def swap_every_two_rows(x, axis=-1): # Create PrecisionConfig with both weight and activation info w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), ) @@ -1238,8 +1251,6 @@ def swap_every_two_rows(x, axis=-1): ) elif mxfp4_backend in TRITON_BACKENDS: - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig - if w13_bias is not None: w13_bias = w13_bias.to(torch.float32) if w2_bias is not None: @@ -1255,10 +1266,10 @@ def swap_every_two_rows(x, axis=-1): ) w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex) + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex) + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) # The original mxfp4 block scales have been swizzled into the @@ -1653,8 +1664,6 @@ def convert_weight_to_mxfp4_moe_kernel_format( elif mxfp4_backend in TRITON_BACKENDS or ( mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 ): - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig - if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: # AITER moe_gemm_a16w4 needs gate/up interleaved def interleave_gate_up(w: torch.Tensor) -> torch.Tensor: @@ -1698,10 +1707,10 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor: ) w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex) + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex) + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) # The original mxfp4 block scales have been swizzled into the diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 1bd6b3f9314b..a1d050177507 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -7,7 +7,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer logger = init_logger(__name__) @@ -27,10 +27,23 @@ def should_use_cdna4_mx_scale_swizzle() -> bool: return on_gfx950() and get_tensor_model_parallel_world_size() <= 2 +def weight_mx_scale(precision_config): + """mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it + as b_mx_scale, 3.5.1/3.6 as weight_scale. + """ + if get_triton_kernels_version() == "3.8": + return precision_config.b_mx_scale + return precision_config.weight_scale + + def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel.""" assert has_triton_kernels() - import triton_kernels.matmul_ogs_details.opt_flags as opt_flags + if get_triton_kernels_version() == "3.8": + # 3.8: matmul_ogs_details -> matmul_details + import triton_kernels.matmul_details.opt_flags as opt_flags + else: + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags from triton_kernels.numerics import InFlexData from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout @@ -51,6 +64,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): ) value_layout = StridedLayout scale_layout = StridedLayout + elif current_platform.is_rocm() and get_triton_kernels_version() == "3.8": + # StridedLayout(-2) is triton_kernels' default CDNA mxfp4 weight layout + # (make_default_matmul_mxfp4_w_layout); 3.8 wants layout instances. + from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout + + value_layout = StridedLayout(-2) + scale_layout = ( + CDNA4MXScaleLayout() + if should_use_cdna4_mx_scale_swizzle() + else StridedLayout() + ) elif current_platform.is_rocm(): value_layout = StridedLayout if should_use_cdna4_mx_scale_swizzle(): diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index d6a2aabbb47c..4e90367a3e6c 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -527,6 +527,26 @@ def has_triton_kernels() -> bool: return is_available +@cache +def get_triton_kernels_version() -> str | None: + """The triton_kernels MoE-API generation ("3.5.1"/"3.6"/"3.8"), or None. + + Inferred by capability since the package exposes no usable version: 3.8 + replaced ``matmul_ogs`` with ``matmul``, and 3.5.1 predates ``SparseMatrix``. + """ + if not has_triton_kernels(): + return None + try: + import triton_kernels.matmul_ogs # noqa: F401 + except ImportError: + return "3.8" + try: + from triton_kernels.tensor import SparseMatrix # noqa: F401 + except ImportError: + return "3.5.1" + return "3.6" + + @cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available. From e52da038d19c202ede695a4273adc2316b4ba1a7 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 17 Sep 2026 09:15:05 +0800 Subject: [PATCH 359/774] [Rust Frontend] Normalize native renderer reasoning controls (#56998) Signed-off-by: Bugen Zhao --- rust/src/chat/src/backend/hf.rs | 59 ++- rust/src/chat/src/backend/mod.rs | 4 +- rust/src/chat/src/error.rs | 3 + rust/src/chat/src/lib.rs | 4 +- rust/src/chat/src/reasoning.rs | 397 ++++++++++++++++++ rust/src/chat/src/renderer/deepseek.rs | 142 +++---- .../src/renderer/deepseek_v32/encoding.rs | 8 +- .../src/chat/src/renderer/deepseek_v32/mod.rs | 24 +- .../chat/src/renderer/deepseek_v32/tests.rs | 29 +- rust/src/chat/src/renderer/deepseek_v4/mod.rs | 30 +- .../chat/src/renderer/deepseek_v4/tests.rs | 84 +++- .../src/chat/src/renderer/deepseek_v41/mod.rs | 26 +- .../chat/src/renderer/deepseek_v41/tests.rs | 64 ++- rust/src/chat/src/renderer/harmony/mod.rs | 78 ++-- rust/src/chat/src/renderer/harmony/tests.rs | 48 ++- rust/src/chat/src/renderer/hf/mod.rs | 98 ++++- rust/src/chat/src/renderer/inkling/mod.rs | 62 ++- rust/src/chat/src/renderer/inkling/tests.rs | 80 +++- .../src/chat/src/renderer/kimi_k3/encoding.rs | 69 +-- rust/src/chat/src/renderer/kimi_k3/mod.rs | 32 +- rust/src/chat/src/renderer/kimi_k3/tests.rs | 99 ++++- rust/src/chat/src/renderer/mod.rs | 34 +- rust/src/chat/src/renderer/test_utils.rs | 5 +- rust/src/chat/src/request.rs | 112 +---- rust/src/server/src/error.rs | 26 +- .../routes/openai/chat_completions/convert.rs | 48 ++- .../routes/openai/chat_completions/types.rs | 8 +- .../openai/chat_completions/validate.rs | 2 +- .../server/src/routes/openai/utils/types.rs | 30 ++ 29 files changed, 1252 insertions(+), 453 deletions(-) create mode 100644 rust/src/chat/src/reasoning.rs diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 5b5b8772c55f..af9af8b51ca4 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -70,12 +70,26 @@ impl HfChatBackend { options, multimodal_render_info, )?), - RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()), - RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), - RendererSelection::DeepSeekV41 => Arc::new(DeepSeekV41ChatRenderer::new()), - RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), - RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), - RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())), + RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new( + options.default_chat_template_kwargs, + )), + RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new( + options.default_chat_template_kwargs, + )), + RendererSelection::DeepSeekV41 => Arc::new(DeepSeekV41ChatRenderer::new( + options.default_chat_template_kwargs, + )), + RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new( + options.default_chat_template_kwargs, + )?), + RendererSelection::Inkling => Arc::new(InklingChatRenderer::new( + tokenizer.clone(), + options.default_chat_template_kwargs, + )?), + RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new( + tokenizer.clone(), + options.default_chat_template_kwargs, + )), }; info!( @@ -280,6 +294,39 @@ mod tests { assert_eq!(prompt, "hello"); } + #[test] + fn native_renderer_inherits_deployment_reasoning_below_request_controls() { + let backend = HfChatBackend::from_resolved_model_files( + resolved_files(r#"{"model_type":"deepseek_v4"}"#, "{}"), + "test-model".to_string(), + LoadModelBackendsOptions { + default_chat_template_kwargs: [ + ("thinking".to_string(), serde_json::json!(false)), + ("reasoning_effort".to_string(), serde_json::json!("low")), + ] + .into(), + ..Default::default() + }, + test_tokenizer(), + ) + .unwrap(); + let mut request = request_with_user_text("hello"); + let rendered = backend.chat_renderer().render(&request).unwrap(); + assert!(rendered.prompt.into_text().unwrap().ends_with("")); + assert_eq!(rendered.effective_template_kwargs["enable_thinking"], false); + + request.chat_options.reasoning_effort = Some(crate::EffortValue::from("max")); + let rendered = backend.chat_renderer().render(&request).unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + assert!(prompt.contains("Reasoning Effort: Beyond maximum")); + assert!(prompt.ends_with("")); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + "max" + ); + assert_eq!(rendered.effective_template_kwargs["thinking"], true); + } + #[test] fn auto_uses_harmony_renderer_and_output_processor_for_gpt_oss_model_type() { let backend = backend_for_selection( diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index eed4e95cbe0c..af4579159e82 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -78,8 +78,8 @@ pub struct LoadModelBackendsOptions { /// Optional server-default chat template override, provided either as an /// inline template or as a path to a template file. pub chat_template: Option, - /// Optional server-default keyword arguments merged into every - /// chat-template render before request-level `chat_template_kwargs`. + /// Server-default keyword arguments. HF merges these before request kwargs; + /// native renderers inherit supported reasoning controls below request controls. pub default_chat_template_kwargs: HashMap, /// Maximum number of input items allowed per prompt for each modality. /// Unspecified modalities are unlimited. diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index a8304f17aa18..db985df58d22 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -19,6 +19,8 @@ pub enum Error { ChatTemplate(String), #[error("{0}")] InvalidReasoningEffort(String), + #[error("{message}")] + InvalidReasoningControl { message: String }, #[error("multimodal input is not supported by this chat renderer")] UnsupportedMultimodalRenderer, #[error("unsupported multimodal content: {0}")] @@ -101,6 +103,7 @@ impl Error { match self { Self::PromptTooLong { .. } | Self::InvalidReasoningEffort(_) + | Self::InvalidReasoningControl { .. } | Self::DuplicateToolName { .. } | Self::ToolChoiceRequiresTools | Self::ToolChoiceFunctionNotFound { .. } => true, diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 9c59d1402896..69d39ca78ed3 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -31,6 +31,7 @@ pub use parser::reasoning::{ }; pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use parser::{ParserSelection, validate_parser_overrides}; +pub use reasoning::EffortValue; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer, @@ -39,7 +40,7 @@ pub use renderer::{ }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, - ChatToolChoice, GenerationPromptMode, ReasoningEffort, ResolvedToolContext, SamplingParams, + ChatToolChoice, GenerationPromptMode, ResolvedToolContext, SamplingParams, }; pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage}; pub use vllm_engine_core_client::protocol::multimodal::MmFeatures; @@ -52,6 +53,7 @@ mod event; pub mod multimodal; mod output; mod parser; +mod reasoning; mod renderer; mod request; mod stream; diff --git a/rust/src/chat/src/reasoning.rs b/rust/src/chat/src/reasoning.rs new file mode 100644 index 000000000000..b5a5b17684e6 --- /dev/null +++ b/rust/src/chat/src/reasoning.rs @@ -0,0 +1,397 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Shared reasoning-control resolution for native renderers. +//! +//! Each input source is parsed independently into [`ReasoningControl`]. Sources +//! then fall back in priority order: request, deployment defaults, model default. +//! Within a source, an explicit toggle decides the mode and effort supplies its +//! intensity. Keeping sources separate lets request effort override a deployment +//! toggle, and lets an enabled request inherit effort from enabled defaults. +//! +//! This module owns input types, precedence, and the standard `"none"` sentinel. +//! Model adapters own supported effort names, numeric ranges, and prompt encoding. +//! After model lowering, [`ReasoningControl::template_kwargs`] projects the prompt's +//! decision for downstream consumers. HF prepares raw template inputs separately. + +use std::collections::HashMap; +use std::fmt; + +use serde::{Deserialize, Serialize}; +use serde_json::{Number, Value, json}; + +use crate::error::{Result, invalid_reasoning_control}; +use crate::request::ChatRequest; + +/// A model-specific reasoning effort, before or after renderer lowering. +/// +/// Strings preserve model-specific names; JSON numbers preserve integer and +/// floating-point representations for model-local range validation. Missing/null +/// effort is represented by `Option::None` in [`crate::ChatOptions`]. +/// Supported names and ranges are validated by each renderer or HF template. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum EffortValue { + String(String), + Number(Number), +} + +impl TryFrom for EffortValue { + type Error = crate::error::Error; + + fn try_from(value: Value) -> Result { + match value { + Value::String(value) => Ok(Self::String(value)), + Value::Number(value) => Ok(Self::Number(value)), + _ => Err(invalid_reasoning_control!( + "reasoning effort must be a string or number, got {value}" + )), + } + } +} + +impl TryFrom for EffortValue { + type Error = crate::error::Error; + + /// Construct a numeric effort, rejecting NaN and infinities. + fn try_from(value: f64) -> Result { + Number::from_f64(value).map(Self::Number).ok_or_else(|| { + invalid_reasoning_control!("reasoning effort must be finite, got {value}") + }) + } +} + +impl From<&str> for EffortValue { + fn from(value: &str) -> Self { + Self::String(value.to_owned()) + } +} + +impl fmt::Display for EffortValue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::String(value) => json!(value).fmt(f), + Self::Number(value) => value.fmt(f), + } + } +} + +impl EffortValue { + pub fn as_str(&self) -> Option<&str> { + match self { + Self::String(value) => Some(value), + Self::Number(_) => None, + } + } + + pub fn as_f64(&self) -> Option { + match self { + Self::Number(value) => value.as_f64(), + Self::String(_) => None, + } + } +} + +/// Reasoning intent for one source, or the result of falling back across sources. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum ReasoningControl { + /// This source leaves the mode and effort to a lower-priority source. + Default, + /// Reasoning is explicitly disabled; lower-priority defaults preserve this mode. + Disabled, + /// Reasoning is enabled. Missing effort inherits from an enabled default. + Enabled { effort: Option }, +} + +impl ReasoningControl { + /// Parse request controls; the typed effort field takes precedence over kwargs effort. + pub(crate) fn from_request(request: &ChatRequest) -> Result { + Self::from_kwargs( + &request.chat_options.template_kwargs, + request.chat_options.reasoning_effort.as_ref(), + ) + } + + /// Parse a deployment source, whose controls all come from template kwargs. + pub(crate) fn from_template_kwargs(kwargs: &HashMap) -> Result { + Self::from_kwargs(kwargs, None) + } + + /// Select controls within one source, then derive its mode and optional effort. + fn from_kwargs( + kwargs: &HashMap, + typed_effort: Option<&EffortValue>, + ) -> Result { + // Presence sets alias priority: `thinking` wins, including when its value + // is invalid. Validate the selected spelling as a boolean. + let toggle = ["thinking", "enable_thinking"] + .into_iter() + .find_map(|key| kwargs.get(key).map(|value| (key, value))) + .map(|(key, value)| { + value.as_bool().ok_or_else(|| { + invalid_reasoning_control!( + "template kwarg `{key}` must be a boolean, got {value}" + ) + }) + }) + .transpose()?; + + // The typed request field wins over kwargs. Parse kwargs only when used; + // disabled mode discards effort, and JSON null means omitted effort. + let effort = match (toggle, typed_effort) { + (Some(false), _) => None, + (_, Some(effort)) => Some(effort.clone()), + (_, None) => kwargs + .get("reasoning_effort") + .filter(|effort| !effort.is_null()) + .cloned() + .map(EffortValue::try_from) + .transpose()?, + }; + + match (toggle, effort) { + // 1. Explicit false disables reasoning and discards any effort, + // including values that would fail validation in an enabled mode. + (Some(false), _) => Ok(Self::Disabled), + // 2. `none` disables reasoning when the toggle is absent. Explicit true + // enables reasoning and leaves effort to enabled deployment/model defaults. + (toggle, Some(EffortValue::String(effort))) if effort == "none" => { + Ok(if toggle == Some(true) { + Self::Enabled { effort: None } + } else { + Self::Disabled + }) + } + // 3. Absent toggle and absent/null effort leave the entire decision + // to the next source in the fallback chain. + (None, None) => Ok(Self::Default), + // 4. Explicit true or a concrete effort enables reasoning; + // missing effort inherits later. + (_, effort) => Ok(Self::Enabled { effort }), + } + } + + /// Parse request and deployment independently, then apply request precedence. + /// + /// Each source validates its selected input types. Model adapters subsequently + /// apply their model default and validate/map the resulting effective effort. + pub(crate) fn resolve( + request: &ChatRequest, + defaults: &HashMap, + ) -> Result { + Ok(Self::from_request(request)?.fallback(Self::from_template_kwargs(defaults)?)) + } + + /// Fill this higher-priority source's missing decision or effort from `defaults`. + /// + /// `Default` inherits the whole lower-priority state. `Enabled { effort: None }` + /// inherits effort only from enabled defaults and retains its enabled mode. + /// Explicit disabling and concrete effort are final at this priority level. + pub(crate) fn fallback(self, defaults: Self) -> Self { + match (self, defaults) { + (Self::Default, defaults) => defaults, + (Self::Enabled { effort: None }, Self::Enabled { effort }) => Self::Enabled { effort }, + (request, _) => request, + } + } + + /// Construct an enabled model default or a model-lowered concrete effort. + pub(crate) fn enabled(effort: impl Into) -> Self { + Self::Enabled { + effort: Some(effort.into()), + } + } + + /// Report the resolved mode; callers apply model defaults before rendering. + pub(crate) fn is_enabled(&self) -> bool { + matches!(self, Self::Enabled { .. }) + } + + /// Return the active effort, which may still be missing before model fallback. + pub(crate) fn effort(&self) -> Option<&EffortValue> { + match self { + Self::Enabled { effort } => effort.as_ref(), + Self::Default | Self::Disabled => None, + } + } + + /// Project the same model-lowered decision used to render the native prompt. + /// + /// Preserve request extras and write both toggle aliases plus the effective + /// effort. Disabled mode emits `"none"`; enabled mode with no effective effort + /// clears the original effort key. `Default` preserves request kwargs. + pub(crate) fn template_kwargs(&self, request: &ChatRequest) -> HashMap { + let mut kwargs = request.chat_options.template_kwargs.clone(); + if matches!(self, Self::Default) { + return kwargs; + } + kwargs.insert("thinking".to_string(), json!(self.is_enabled())); + kwargs.insert("enable_thinking".to_string(), json!(self.is_enabled())); + match self { + Self::Disabled => { + kwargs.insert("reasoning_effort".to_string(), json!("none")); + } + Self::Enabled { + effort: Some(effort), + } => { + kwargs.insert("reasoning_effort".to_string(), json!(effort)); + } + Self::Enabled { effort: None } => { + kwargs.remove("reasoning_effort"); + } + Self::Default => unreachable!(), + } + kwargs + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn kwargs(value: Value) -> HashMap { + serde_json::from_value(value).unwrap() + } + + #[test] + fn numeric_effort_rejects_non_finite_values() { + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!(EffortValue::try_from(value).unwrap_err().is_request_validation_error()); + } + } + + #[test] + fn explicit_toggle_decides_mode_and_none_inherits_enabled_effort() { + use ReasoningControl::{Default, Disabled, Enabled}; + + for (source, expected) in [ + (json!({}), Default), + (json!({"reasoning_effort": "none"}), Disabled), + ( + json!({"reasoning_effort": "high"}), + ReasoningControl::enabled("high"), + ), + (json!({"enable_thinking": true}), Enabled { effort: None }), + ( + json!({"enable_thinking": true, "reasoning_effort": "none"}), + Enabled { effort: None }, + ), + ( + json!({"enable_thinking": true, "reasoning_effort": "low"}), + ReasoningControl::enabled("low"), + ), + ( + json!({"enable_thinking": false, "reasoning_effort": "high"}), + Disabled, + ), + ( + json!({"thinking": true, "enable_thinking": false}), + Enabled { effort: None }, + ), + ( + json!({"thinking": false, "enable_thinking": true}), + Disabled, + ), + ( + json!({"reasoning_effort": 37}), + ReasoningControl::enabled(EffortValue::Number(37.into())), + ), + (json!({"reasoning_effort": null}), Default), + ] { + let source = kwargs(source); + assert_eq!( + ReasoningControl::from_template_kwargs(&source).unwrap(), + expected, + "{source:?}" + ); + } + } + + #[test] + fn request_and_deployment_states_fall_back_before_model_defaults() { + for (request, defaults, expected) in [ + ( + json!({}), + json!({"enable_thinking": false}), + ReasoningControl::Disabled, + ), + ( + json!({"reasoning_effort": "high"}), + json!({"thinking": false}), + ReasoningControl::enabled("high"), + ), + ( + json!({"enable_thinking": false}), + json!({"reasoning_effort": "high"}), + ReasoningControl::Disabled, + ), + ( + json!({"thinking": true, "reasoning_effort": "none"}), + json!({"reasoning_effort": "low"}), + ReasoningControl::enabled("low"), + ), + ( + json!({"thinking": true}), + json!({"thinking": false, "reasoning_effort": "low"}), + ReasoningControl::enabled("max"), + ), + (json!({}), json!({}), ReasoningControl::enabled("max")), + ] { + let request = kwargs(request); + let defaults = kwargs(defaults); + let actual = ReasoningControl::from_template_kwargs(&request) + .unwrap() + .fallback(ReasoningControl::from_template_kwargs(&defaults).unwrap()) + .fallback(ReasoningControl::enabled("max")); + assert_eq!( + actual, expected, + "request={request:?}, defaults={defaults:?}" + ); + } + } + + #[test] + fn selected_toggle_requires_a_boolean() { + for value in [ + json!("true"), + json!(1), + json!(null), + json!({"type": "enabled"}), + ] { + assert!( + ReasoningControl::from_template_kwargs(&kwargs(json!({ + "thinking": value, "enable_thinking": true + }))) + .unwrap_err() + .is_request_validation_error() + ); + } + } + + #[test] + fn invalid_effort_types_are_rejected_only_when_the_source_uses_them() { + for value in [json!(true), json!(false), json!([]), json!({})] { + for toggle in [None, Some(true), Some(false)] { + let mut source = kwargs(json!({"reasoning_effort": value})); + if let Some(toggle) = toggle { + source.insert("thinking".to_owned(), json!(toggle)); + } + let control = ReasoningControl::from_template_kwargs(&source); + if toggle == Some(false) { + assert_eq!(control.unwrap(), ReasoningControl::Disabled); + } else { + assert!(control.unwrap_err().is_request_validation_error()); + } + } + + // A typed request effort shadows malformed kwargs effort. + let mut request = ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(EffortValue::from("high")); + request.chat_options.template_kwargs = kwargs(json!({"reasoning_effort": value})); + assert_eq!( + ReasoningControl::from_request(&request).unwrap(), + ReasoningControl::enabled("high") + ); + } + } +} diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 01576a7c4341..7892fd2497c5 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -17,10 +17,10 @@ use serde_json_fmt::JsonFormat; use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER; use super::MediaPartSource; +use crate::EffortValue; use crate::error::{Error, Result}; -use crate::request::{ - ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort, -}; +use crate::reasoning::ReasoningControl; +use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool}; use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall}; const BOS_TOKEN: &str = "<｜begin▁of▁sentence｜>"; @@ -87,21 +87,26 @@ struct RenderedToolSchema<'a> { } /// Render one chat request into the final prompt string. -pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result { - render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt) +pub(super) fn render_request( + request: &ChatRequest, + dialect: DsDialect, + reasoning: &ReasoningControl, +) -> Result { + render_request_with_media_order(request, dialect, reasoning).map(|(prompt, _)| prompt) } /// Render one chat request and record media in placeholder order. pub(super) fn render_request_with_media_order( request: &ChatRequest, dialect: DsDialect, + reasoning: &ReasoningControl, ) -> Result<(String, Vec)> { - let (thinking_mode, reasoning_effort_prompt) = match dialect { - DsDialect::V4 => { - resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))? - } - DsDialect::V41 => resolve_v41_thinking_options(request)?, + let thinking_mode = if reasoning.is_enabled() { + ThinkingMode::Thinking + } else { + ThinkingMode::Chat }; + let reasoning_effort_prompt = reasoning_prompt(reasoning, dialect); let request_tools = request_tools(request); let synthetic_tool_system = needs_synthetic_tool_system(request, request_tools); let drop_thinking = request.parse_template_bool("drop_thinking")?.unwrap_or(true) @@ -254,76 +259,63 @@ pub(super) fn render_request_with_media_order( Ok((out, media_order)) } -/// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer -/// wrapper, the Rust renderer only consumes the typed top-level -/// `reasoning_effort`; the generic template-kwargs map is left for HF -/// templates. -fn resolve_thinking_options(request: &ChatRequest) -> Result<(ThinkingMode, &'static str)> { - let mut thinking_mode = match request.enable_thinking()?.unwrap_or(true) { - true => ThinkingMode::Thinking, - false => ThinkingMode::Chat, - }; - let mut reasoning_effort_prompt = REASONING_EFFORT_HIGH; - - match request.chat_options.reasoning_effort { - Some(ReasoningEffort::None) => thinking_mode = ThinkingMode::Chat, - Some(ReasoningEffort::Max) => { - reasoning_effort_prompt = REASONING_EFFORT_MAX; - } - Some(ReasoningEffort::XHigh | ReasoningEffort::High) => { - reasoning_effort_prompt = REASONING_EFFORT_HIGH; +/// Map the shared reasoning decision to the selected DeepSeek effort dialect. +pub(super) fn resolve_reasoning( + control: ReasoningControl, + dialect: DsDialect, +) -> Result { + let control = control.fallback(ReasoningControl::enabled("high")); + if !control.is_enabled() { + return Ok(control); + } + let effort = control.effort().expect("enabled DeepSeek reasoning has a default effort"); + let effort = match dialect { + DsDialect::V4 => { + let effort = match effort.as_str() { + Some("minimal" | "medium" | "low") => "low", + Some("high" | "xhigh") => "high", + Some("max") => "max", + _ => { + return Err(Error::InvalidReasoningEffort(format!( + "DeepSeek V4 reasoning_effort must be minimal, low, medium, high, xhigh, or max, got {effort}" + ))); + } + }; + EffortValue::from(effort) } - Some(ReasoningEffort::Minimal | ReasoningEffort::Medium | ReasoningEffort::Low) => { - reasoning_effort_prompt = ""; + DsDialect::V41 => { + let budget = match effort { + EffortValue::String(effort) => match effort.as_str() { + "low" => 25, + "high" => 50, + "xhigh" => 75, + "max" => 100, + _ => return Err(invalid_v41_reasoning_effort()), + }, + EffortValue::Number(budget) => budget + .as_u64() + .filter(|budget| (1..=100).contains(budget)) + .ok_or_else(invalid_v41_reasoning_effort)?, + }; + EffortValue::Number(budget.into()) } - None => {} - } - - Ok((thinking_mode, reasoning_effort_prompt)) + }; + Ok(ReasoningControl::enabled(effort)) } -/// Resolve V4.1's numeric reasoning effort using the top-level value before template kwargs. -fn resolve_v41_thinking_options( - request: &ChatRequest, -) -> Result<(ThinkingMode, Cow<'static, str>)> { - let mut thinking = request.enable_thinking()?.unwrap_or(true); - let budget = match request.chat_options.reasoning_effort { - Some(ReasoningEffort::None) => { - thinking = false; - 50 - } - Some(ReasoningEffort::Low) => 25, - Some(ReasoningEffort::High) => 50, - Some(ReasoningEffort::XHigh) => 75, - Some(ReasoningEffort::Max) => 100, - Some(ReasoningEffort::Minimal | ReasoningEffort::Medium) => { - return Err(invalid_v41_reasoning_effort()); - } - None => match request.chat_options.template_kwargs.get("reasoning_effort") { - None | Some(Value::Null) => 50, - Some(Value::String(effort)) => match effort.as_str() { - "low" => 25, - "high" => 50, - "xhigh" => 75, - "max" => 100, - _ => return Err(invalid_v41_reasoning_effort()), - }, - Some(Value::Number(budget)) => budget - .as_u64() - .filter(|budget| (1..=100).contains(budget)) - .ok_or_else(invalid_v41_reasoning_effort)?, - Some(_) => return Err(invalid_v41_reasoning_effort()), - }, +fn reasoning_prompt(control: &ReasoningControl, dialect: DsDialect) -> Cow<'static, str> { + let Some(effort) = control.effort() else { + return Cow::Borrowed(""); }; - if thinking { - Ok(( - ThinkingMode::Thinking, - Cow::Owned(format!( - "Reasoning Effort: {budget} (range 1-100, the higher the value, the more thorough the reasoning)\n\n" - )), - )) - } else { - Ok((ThinkingMode::Chat, Cow::Borrowed(""))) + match dialect { + DsDialect::V4 => Cow::Borrowed(match effort.as_str() { + Some("high") => REASONING_EFFORT_HIGH, + Some("max") => REASONING_EFFORT_MAX, + _ => "", + }), + DsDialect::V41 => Cow::Owned(format!( + "Reasoning Effort: {effort} (range 1-100, the higher the value, the more thorough the reasoning)\n\n" + )), } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 4fc5f26530f8..74c7e4cbf286 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -14,6 +14,7 @@ use serde_json::Value; use serde_json_fmt::JsonFormat; use crate::error::{Error, Result}; +use crate::reasoning::ReasoningControl; use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatRole, ChatTool}; use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall}; @@ -42,8 +43,11 @@ struct RenderedToolSchema<'a> { } /// Render one chat request into the final prompt string. -pub(super) fn render_request(request: &ChatRequest) -> Result { - let thinking_mode = match request.enable_thinking()?.unwrap_or(false) { +pub(super) fn render_request( + request: &ChatRequest, + reasoning: &ReasoningControl, +) -> Result { + let thinking_mode = match reasoning.is_enabled() { true => ThinkingMode::Thinking, false => ThinkingMode::Chat, }; diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs index 03c2765d1913..b4ac1bae1155 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs @@ -3,31 +3,41 @@ mod encoding; +use std::collections::HashMap; + +use serde_json::Value; use vllm_text::Prompt; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; use crate::Result; +use crate::reasoning::ReasoningControl; use crate::request::ChatRequest; /// Dedicated DeepSeek V3.2 renderer. -#[derive(Debug, Clone, Copy, Default)] -pub struct DeepSeekV32ChatRenderer; +#[derive(Debug, Clone, Default)] +pub struct DeepSeekV32ChatRenderer { + default_template_kwargs: HashMap, +} impl DeepSeekV32ChatRenderer { /// Create the dedicated DeepSeek V3.2 renderer. - pub fn new() -> Self { - Self + pub fn new(default_template_kwargs: HashMap) -> Self { + Self { + default_template_kwargs, + } } } impl ChatRenderer for DeepSeekV32ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let reasoning = ReasoningControl::resolve(request, &self.default_template_kwargs)? + .fallback(ReasoningControl::Disabled); Ok(RenderedPrompt { - prompt: Prompt::Text(encoding::render_request(request)?), + prompt: Prompt::Text(encoding::render_request(request, &reasoning)?), media_order: None, - effective_template_kwargs: request_template_kwargs(request), + effective_template_kwargs: reasoning.template_kwargs(request), }) } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 56c03a92cab4..dd9ec13ddc97 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -8,6 +8,7 @@ use serde_json::{Value, json}; use thiserror_ext::AsReport; use super::DeepSeekV32ChatRenderer; +use crate::EffortValue; use crate::error::Error; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; @@ -18,7 +19,7 @@ use crate::request::{ use crate::{ChatRenderer, ChatRole}; fn render_request(request: &ChatRequest) -> String { - DeepSeekV32ChatRenderer::new() + DeepSeekV32ChatRenderer::default() .render(request) .unwrap() .prompt @@ -26,8 +27,30 @@ fn render_request(request: &ChatRequest) -> String { .expect("deepseek renderer should return text prompt") } +#[test] +fn effort_enables_binary_thinking_and_none_disables_it() { + let mut request = ChatRequest::for_test(); + for (effort, enabled) in [ + (EffortValue::from("high"), true), + (EffortValue::from("custom"), true), + (EffortValue::Number(42.into()), true), + (EffortValue::from("none"), false), + ] { + request.chat_options.reasoning_effort = Some(effort); + let rendered = DeepSeekV32ChatRenderer::default().render(&request).unwrap(); + assert_eq!( + rendered.prompt.into_text().unwrap().ends_with(""), + enabled + ); + assert_eq!( + rendered.effective_template_kwargs["enable_thinking"], + enabled + ); + } +} + fn render_result(request: &ChatRequest) -> Result { - DeepSeekV32ChatRenderer::new().render(request).map(|rendered| { + DeepSeekV32ChatRenderer::default().render(request).map(|rendered| { rendered .prompt .into_text() @@ -345,7 +368,7 @@ fn render_rejects_multimodal_input() { ..ChatRequest::for_test() }; - let error = DeepSeekV32ChatRenderer::new().render(&request).unwrap_err(); + let error = DeepSeekV32ChatRenderer::default().render(&request).unwrap_err(); assert!(matches!( error, diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs index 54b111e5050f..136938fbc0c8 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs @@ -1,31 +1,47 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; + +use serde_json::Value; use vllm_text::Prompt; use super::deepseek::{self, DsDialect}; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; use crate::Result; +use crate::reasoning::ReasoningControl; use crate::request::ChatRequest; /// Dedicated DeepSeek V4 renderer. -#[derive(Debug, Clone, Copy, Default)] -pub struct DeepSeekV4ChatRenderer; +#[derive(Debug, Clone, Default)] +pub struct DeepSeekV4ChatRenderer { + default_template_kwargs: HashMap, +} impl DeepSeekV4ChatRenderer { - pub fn new() -> Self { - Self + pub fn new(default_template_kwargs: HashMap) -> Self { + Self { + default_template_kwargs, + } } } impl ChatRenderer for DeepSeekV4ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let reasoning = deepseek::resolve_reasoning( + ReasoningControl::resolve(request, &self.default_template_kwargs)?, + DsDialect::V4, + )?; Ok(RenderedPrompt { - prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?), + prompt: Prompt::Text(deepseek::render_request( + request, + DsDialect::V4, + &reasoning, + )?), media_order: None, - effective_template_kwargs: request_template_kwargs(request), + effective_template_kwargs: reasoning.template_kwargs(request), }) } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 9404b54296a7..21a33ec2d504 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -8,13 +8,14 @@ use serde_json::Value; use super::DeepSeekV4ChatRenderer; use crate::ChatRenderer; +use crate::EffortValue; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::deepseek_v41::DeepSeekV41ChatRenderer; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ChatMessage, ChatRequest, ChatTool, GenerationPromptMode, ReasoningEffort}; +use crate::request::{ChatMessage, ChatRequest, ChatTool, GenerationPromptMode}; fn render_request(request: &ChatRequest) -> String { - DeepSeekV4ChatRenderer::new() + DeepSeekV4ChatRenderer::default() .render(request) .unwrap() .prompt @@ -31,13 +32,13 @@ fn thinking_request(messages: Vec) -> ChatRequest { .chat_options .template_kwargs .insert("thinking".to_string(), Value::Bool(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); request } fn fixture_request(input_name: &str) -> ChatRequest { let mut request = fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options()); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); request } @@ -203,7 +204,7 @@ fn reasoning_effort_max_adds_prefix_when_thinking_is_enabled() { .chat_options .template_kwargs .insert("thinking".to_string(), Value::Bool(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Max); + request.chat_options.reasoning_effort = Some(EffortValue::from("max")); let rendered = render_request(&request); @@ -226,7 +227,7 @@ fn reasoning_effort_high_adds_0731_high_prefix() { .chat_options .template_kwargs .insert("thinking".to_string(), Value::Bool(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::High); + request.chat_options.reasoning_effort = Some(EffortValue::from("high")); let rendered = render_request(&request); @@ -258,7 +259,7 @@ fn reasoning_effort_xhigh_maps_to_high() { messages: vec![ChatMessage::user("solve it")], ..ChatRequest::for_test() }; - request.chat_options.reasoning_effort = Some(ReasoningEffort::XHigh); + request.chat_options.reasoning_effort = Some(EffortValue::from("xhigh")); let rendered = render_request(&request); @@ -268,7 +269,7 @@ fn reasoning_effort_xhigh_maps_to_high() { #[test] fn minimal_and_medium_reasoning_effort_map_to_low() { - for effort in [ReasoningEffort::Minimal, ReasoningEffort::Medium] { + for effort in [EffortValue::from("minimal"), EffortValue::from("medium")] { let mut request = ChatRequest { messages: vec![ChatMessage::user("solve it")], ..ChatRequest::for_test() @@ -292,7 +293,7 @@ fn reasoning_effort_low_keeps_the_default_prompt() { .chat_options .template_kwargs .insert("thinking".to_string(), Value::Bool(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); let rendered = render_request(&request); @@ -305,11 +306,7 @@ fn reasoning_effort_none_disables_thinking() { messages: vec![ChatMessage::user("answer directly")], ..ChatRequest::for_test() }; - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), Value::Bool(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); let rendered = render_request(&request); @@ -317,6 +314,49 @@ fn reasoning_effort_none_disables_thinking() { .assert_eq(&rendered); } +#[test] +fn normalized_reasoning_drives_both_prompt_and_compatibility_kwargs() { + for (kwargs, enabled, effort, prefix) in [ + ( + serde_json::json!({"enable_thinking": true, "reasoning_effort": "none"}), + true, + "high", + "Absolute maximum", + ), + ( + serde_json::json!({"thinking": false, "enable_thinking": true, "reasoning_effort": "max"}), + false, + "none", + "", + ), + ( + serde_json::json!({"reasoning_effort": "xhigh"}), + true, + "high", + "Absolute maximum", + ), + ] { + let mut request = ChatRequest::for_test(); + request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap(); + let rendered = DeepSeekV4ChatRenderer::default().render(&request).unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + assert_eq!(prompt.ends_with(""), enabled); + assert_eq!(prompt.contains("Reasoning Effort:"), enabled); + if enabled { + assert!(prompt.contains(prefix)); + } + assert_eq!(rendered.effective_template_kwargs["thinking"], enabled); + assert_eq!( + rendered.effective_template_kwargs["enable_thinking"], + enabled + ); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + effort + ); + } +} + #[test] fn system_to_assistant_transitions_in_chat_and_thinking_modes() { let rendered = [false, true].map(|enable_thinking| { @@ -338,7 +378,7 @@ fn system_to_assistant_transitions_in_chat_and_thinking_modes() { .chat_options .template_kwargs .insert("enable_thinking".to_string(), Value::Bool(enable_thinking)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); render_request(&request) }); @@ -366,7 +406,7 @@ fn trailing_system_transitions_in_chat_and_thinking_modes() { .chat_options .template_kwargs .insert("enable_thinking".to_string(), Value::Bool(enable_thinking)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); render_request(&request) }); @@ -380,7 +420,7 @@ fn trailing_system_transitions_in_chat_and_thinking_modes() { } #[test] -fn reasoning_effort_template_kwarg_is_ignored() { +fn reasoning_effort_template_kwarg_selects_effort() { let mut request = ChatRequest { messages: vec![ChatMessage::user("solve it")], ..ChatRequest::for_test() @@ -396,15 +436,15 @@ fn reasoning_effort_template_kwarg_is_ignored() { let rendered = render_request(&request); - assert!(rendered.starts_with("<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum")); + assert!(rendered.starts_with("<｜begin▁of▁sentence｜>Reasoning Effort: Beyond maximum")); assert!(rendered.ends_with("<｜Assistant｜>")); } #[test] fn tool_call_arguments_preserve_non_objects_like_recipe() { let renderers: [(&dyn ChatRenderer, &str); 2] = [ - (&DeepSeekV4ChatRenderer::new(), "parameter"), - (&DeepSeekV41ChatRenderer::new(), " parameter"), + (&DeepSeekV4ChatRenderer::default(), "parameter"), + (&DeepSeekV41ChatRenderer::default(), " parameter"), ]; for arguments in [ "not json", @@ -600,7 +640,7 @@ fn drop_thinking_false_keeps_prior_assistant_reasoning() { .chat_options .template_kwargs .insert("drop_thinking".to_string(), Value::Bool(false)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); let rendered = render_request(&request); @@ -693,7 +733,7 @@ fn trailing_system_omits_generation_prompt_when_disabled() { .chat_options .template_kwargs .insert("enable_thinking".to_string(), Value::Bool(false)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + request.chat_options.reasoning_effort = Some(EffortValue::from("low")); let rendered = render_request(&request); diff --git a/rust/src/chat/src/renderer/deepseek_v41/mod.rs b/rust/src/chat/src/renderer/deepseek_v41/mod.rs index 7737f7ebda4f..54544ec3839f 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/mod.rs @@ -1,33 +1,45 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; + +use serde_json::Value; use vllm_text::Prompt; use super::deepseek::{self, DsDialect}; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; use crate::Result; +use crate::reasoning::ReasoningControl; use crate::request::ChatRequest; /// Dedicated DeepSeek V4.1 renderer. -#[derive(Debug, Clone, Copy, Default)] -pub struct DeepSeekV41ChatRenderer; +#[derive(Debug, Clone, Default)] +pub struct DeepSeekV41ChatRenderer { + default_template_kwargs: HashMap, +} impl DeepSeekV41ChatRenderer { - pub fn new() -> Self { - Self + pub fn new(default_template_kwargs: HashMap) -> Self { + Self { + default_template_kwargs, + } } } impl ChatRenderer for DeepSeekV41ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let reasoning = deepseek::resolve_reasoning( + ReasoningControl::resolve(request, &self.default_template_kwargs)?, + DsDialect::V41, + )?; let (prompt, media_order) = - deepseek::render_request_with_media_order(request, DsDialect::V41)?; + deepseek::render_request_with_media_order(request, DsDialect::V41, &reasoning)?; Ok(RenderedPrompt { prompt: Prompt::Text(prompt), media_order: Some(media_order), - effective_template_kwargs: request_template_kwargs(request), + effective_template_kwargs: reasoning.template_kwargs(request), }) } } diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index 3d621c2246fd..62cf6d20247a 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -8,14 +8,15 @@ use serde_json::{Value, json}; use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; +use crate::EffortValue; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ - ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, + ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, }; fn render(request: &ChatRequest) -> String { - DeepSeekV41ChatRenderer::new() + DeepSeekV41ChatRenderer::default() .render(request) .unwrap() .prompt @@ -64,7 +65,7 @@ fn inlines_image_placeholder_at_image_part_positions() { // The placeholder is unconditional (matching the Python encoding's // `IMAGE_PLACEHOLDER`); requests to a multimodal-less backend fail later, // at media preparation. - let renderer = DeepSeekV41ChatRenderer::new(); + let renderer = DeepSeekV41ChatRenderer::default(); let request = ChatRequest { messages: vec![ChatMessage::User { content: ChatContent::Parts(vec![ @@ -87,7 +88,7 @@ fn inlines_image_placeholder_at_image_part_positions() { #[test] fn media_order_follows_reordered_tool_results() { - let renderer = DeepSeekV41ChatRenderer::new(); + let renderer = DeepSeekV41ChatRenderer::default(); let request = ChatRequest { messages: vec![ ChatMessage::assistant_blocks(vec![ @@ -136,7 +137,7 @@ fn media_order_follows_reordered_tool_results() { #[test] fn media_order_excludes_dropped_historical_developer_content() { - let renderer = DeepSeekV41ChatRenderer::new(); + let renderer = DeepSeekV41ChatRenderer::default(); let request = ChatRequest { messages: vec![ ChatMessage::developer(vec![image_part("historical-image")], None), @@ -165,10 +166,11 @@ fn image_part(uuid: &str) -> ChatContentPart { fn maps_reasoning_effort_to_reference_numeric_budget() { for (effort, budget) in [ (None, 50), - (Some(ReasoningEffort::Low), 25), - (Some(ReasoningEffort::High), 50), - (Some(ReasoningEffort::XHigh), 75), - (Some(ReasoningEffort::Max), 100), + (Some(EffortValue::from("low")), 25), + (Some(EffortValue::from("high")), 50), + (Some(EffortValue::from("xhigh")), 75), + (Some(EffortValue::from("max")), 100), + (Some(EffortValue::Number(37.into())), 37), ] { let mut request = request(); request.chat_options.reasoning_effort = effort; @@ -193,11 +195,43 @@ fn accepts_numeric_template_effort_with_top_level_precedence() { request.chat_options.template_kwargs.insert("reasoning_effort".into(), effort); assert!(render(&request).contains(&format!("Reasoning Effort: {budget} (range 1-100"))); - request.chat_options.reasoning_effort = Some(ReasoningEffort::XHigh); + request.chat_options.reasoning_effort = Some(EffortValue::from("xhigh")); assert!(render(&request).contains("Reasoning Effort: 75 (range 1-100")); } } +#[test] +fn enabling_with_none_inherits_deployment_effort_and_reports_numeric_control() { + let renderer = + DeepSeekV41ChatRenderer::new([("reasoning_effort".to_string(), json!(37))].into()); + let mut request = request(); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); + request + .chat_options + .template_kwargs + .insert("enable_thinking".into(), json!(true)); + let rendered = renderer.render(&request).unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + assert!(prompt.contains("Reasoning Effort: 37 (range 1-100")); + assert!(prompt.ends_with("")); + assert_eq!(rendered.effective_template_kwargs["reasoning_effort"], 37); + assert_eq!(rendered.effective_template_kwargs["enable_thinking"], true); + + request.chat_options.template_kwargs.insert("thinking".into(), json!(false)); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".into(), json!(101)); + request.chat_options.reasoning_effort = None; + let rendered = renderer.render(&request).unwrap(); + assert!(rendered.prompt.into_text().unwrap().ends_with("")); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + "none" + ); + assert_eq!(rendered.effective_template_kwargs["enable_thinking"], false); +} + #[test] fn rejects_undefined_effort_names_and_invalid_numeric_budgets() { for effort in [ @@ -211,13 +245,13 @@ fn rejects_undefined_effort_names_and_invalid_numeric_budgets() { ] { let mut request = request(); request.chat_options.template_kwargs.insert("reasoning_effort".into(), effort); - let error = DeepSeekV41ChatRenderer::new().render(&request).unwrap_err(); + let error = DeepSeekV41ChatRenderer::default().render(&request).unwrap_err(); assert!(error.is_request_validation_error()); } - for effort in [ReasoningEffort::Minimal, ReasoningEffort::Medium] { + for effort in [EffortValue::from("minimal"), EffortValue::from("medium")] { let mut request = request(); request.chat_options.reasoning_effort = Some(effort); - let error = DeepSeekV41ChatRenderer::new().render(&request).unwrap_err(); + let error = DeepSeekV41ChatRenderer::default().render(&request).unwrap_err(); assert!(error.is_request_validation_error()); } } @@ -229,8 +263,8 @@ fn chat_mode_and_none_effort_use_closed_thinking_prefix() { let expected = expect!["<｜begin▁of▁sentence｜><｜User｜>question<｜Assistant｜>"]; expected.assert_eq(&render(&request)); - request.chat_options.template_kwargs.insert("thinking".into(), json!(true)); - request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + request.chat_options.template_kwargs.remove("thinking"); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); expected.assert_eq(&render(&request)); } diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs index 71431cc92b46..83b59e91989b 100644 --- a/rust/src/chat/src/renderer/harmony/mod.rs +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -5,21 +5,26 @@ pub(crate) mod encoding; +use std::collections::HashMap; + use openai_harmony::HarmonyEncoding; use openai_harmony::chat::{ Author, Conversation, DeveloperContent, Message, ReasoningEffort as HarmonyReasoningEffort, Role, SystemContent, ToolDescription, }; +use serde_json::Value; use thiserror_ext::AsReport as _; use time::macros::format_description; use vllm_text::Prompt; use self::encoding::harmony_encoding; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; +use crate::AssistantMessageExt as _; +use crate::EffortValue; use crate::error::{Error, Result}; use crate::event::AssistantContentBlock; +use crate::reasoning::ReasoningControl; use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatTool, GenerationPromptMode}; -use crate::{AssistantMessageExt as _, ReasoningEffort}; const SYSTEM_START_DATE_ENV: &str = "VLLM_SYSTEM_START_DATE"; const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS"; @@ -28,6 +33,7 @@ const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTR pub struct HarmonyChatRenderer { encoding: &'static HarmonyEncoding, options: Options, + default_template_kwargs: HashMap, } struct Options { @@ -44,8 +50,9 @@ impl HarmonyChatRenderer { /// unset, the renderer uses the current local date with a UTC fallback. /// - `VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS` moves leading instructions /// into the system model identity when set to a non-zero integer. - pub fn new() -> Result { + pub fn new(default_template_kwargs: HashMap) -> Result { Self::with_options( + default_template_kwargs, env_system_start_date(), env_use_harmony_system_instructions(), ) @@ -57,11 +64,13 @@ impl HarmonyChatRenderer { /// Production code should call [`Self::new`] so the renderer observes the /// same environment contract as the Python Harmony path. pub fn with_options( + default_template_kwargs: HashMap, system_start_date: impl Into, use_system_instructions: bool, ) -> Result { Ok(Self { encoding: harmony_encoding()?, + default_template_kwargs, options: Options { system_start_date: system_start_date.into(), use_system_instructions, @@ -73,7 +82,11 @@ impl HarmonyChatRenderer { /// /// Harmony owns both prompt formatting and tokenization, so the Rust /// frontend bypasses the generic HF tokenizer path for GPT-OSS input. - fn render_token_ids(&self, request: &ChatRequest) -> Result> { + fn render_token_ids( + &self, + request: &ChatRequest, + effort: HarmonyReasoningEffort, + ) -> Result> { if request.has_multimodal() { return Err(Error::UnsupportedMultimodalContent("image_url")); } @@ -86,7 +99,8 @@ impl HarmonyChatRenderer { )); } - let messages = auto_drop_analysis_messages(to_harmony_messages(request, &self.options)?); + let messages = + auto_drop_analysis_messages(to_harmony_messages(request, &self.options, effort)?); let conversation = Conversation::from_messages(messages); // Pass `None` so oss-harmony does not apply its narrower built-in // analysis-drop policy after the Rust-side Python-parity cleanup above. @@ -114,10 +128,13 @@ impl ChatRenderer for HarmonyChatRenderer { /// Render a chat request as [`Prompt::TokenIds`] with template kwargs echoed /// for downstream accounting/debugging. fn render(&self, request: &ChatRequest) -> Result { + let reasoning = ReasoningControl::resolve(request, &self.default_template_kwargs)? + .fallback(ReasoningControl::enabled("medium")); + let effort = to_harmony_reasoning_effort(reasoning.effort())?; Ok(RenderedPrompt { - prompt: Prompt::TokenIds(self.render_token_ids(request)?), + prompt: Prompt::TokenIds(self.render_token_ids(request, effort)?), media_order: None, - effective_template_kwargs: request_template_kwargs(request), + effective_template_kwargs: reasoning.template_kwargs(request), }) } } @@ -127,12 +144,21 @@ impl ChatRenderer for HarmonyChatRenderer { /// This adds the Harmony system/developer preamble, peels at most one leading /// system/developer instruction message, and then lowers the remaining chat /// history message-by-message. -fn to_harmony_messages(request: &ChatRequest, options: &Options) -> Result> { +fn to_harmony_messages( + request: &ChatRequest, + options: &Options, + effort: HarmonyReasoningEffort, +) -> Result> { let (instructions, leading_developer_tools, remaining_messages) = peel_leading_instructions(&request.messages)?; let tool_call_names = tool_call_names(&request.messages); - let mut messages = - build_harmony_preamble(request, instructions, leading_developer_tools, options)?; + let mut messages = build_harmony_preamble( + request, + instructions, + leading_developer_tools, + options, + effort, + )?; for message in remaining_messages { messages.extend(to_harmony_message(message, &tool_call_names, options)?); @@ -177,12 +203,13 @@ fn build_harmony_preamble( instructions: Option, leading_developer_tools: Option<&[ChatTool]>, options: &Options, + effort: HarmonyReasoningEffort, ) -> Result> { let mut messages = vec![Message::from_role_and_content( Role::System, system_content( instructions.as_deref().filter(|_| options.use_system_instructions), - request.chat_options.reasoning_effort, + Some(effort), &options.system_start_date, )?, )]; @@ -227,18 +254,18 @@ fn preamble_tool_descriptions( /// Construct the Harmony system content for the request preamble. /// -/// Harmony defaults the reasoning effort to `medium` when none is provided, so -/// this only sets an explicit effort after validating vLLM's request value. +/// The preamble receives normalized effort; later system messages use Harmony's +/// own default when their effort is omitted. fn system_content( instructions: Option<&str>, - reasoning_effort: Option, + reasoning_effort: Option, system_start_date: &str, ) -> Result { let mut content = SystemContent::new().with_conversation_start_date(system_start_date.to_string()); if let Some(reasoning_effort) = reasoning_effort { - content = content.with_reasoning_effort(to_harmony_reasoning_effort(reasoning_effort)?); + content = content.with_reasoning_effort(reasoning_effort); } if let Some(instructions) = instructions.filter(|text| !text.is_empty()) { @@ -445,20 +472,17 @@ fn to_tool_descriptions(tools: &[ChatTool]) -> Vec { .collect() } -/// Map supported OpenAI reasoning-effort values onto Harmony's enum. +/// Map supported reasoning-effort values onto Harmony's enum. fn to_harmony_reasoning_effort( - reasoning_effort: ReasoningEffort, + reasoning_effort: Option<&EffortValue>, ) -> Result { - match reasoning_effort { - ReasoningEffort::Low => Ok(HarmonyReasoningEffort::Low), - ReasoningEffort::Medium => Ok(HarmonyReasoningEffort::Medium), - ReasoningEffort::High => Ok(HarmonyReasoningEffort::High), - ReasoningEffort::None - | ReasoningEffort::Minimal - | ReasoningEffort::XHigh - | ReasoningEffort::Max => Err(Error::ChatTemplate(format!( - "reasoning_effort={:?} is not supported by Harmony. Supported values are: low, medium, high.", - reasoning_effort.as_str() + match reasoning_effort.and_then(EffortValue::as_str) { + Some("low") => Ok(HarmonyReasoningEffort::Low), + Some("medium") => Ok(HarmonyReasoningEffort::Medium), + Some("high") => Ok(HarmonyReasoningEffort::High), + _ => Err(Error::InvalidReasoningEffort(format!( + "reasoning_effort={} is not supported by Harmony. Supported values are: low, medium, high.", + reasoning_effort.unwrap_or(&EffortValue::from("none")) ))), } } diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs index 4849c932233b..609c1f115920 100644 --- a/rust/src/chat/src/renderer/harmony/tests.rs +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -9,12 +9,11 @@ use thiserror_ext::AsReport as _; use super::HarmonyChatRenderer; use super::encoding::harmony_encoding; use crate::ChatRenderer; +use crate::EffortValue; use crate::error::Error; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ - ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, -}; +use crate::request::{ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode}; const PINNED_DATE: &str = "2025-06-28"; @@ -22,7 +21,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fixture_chat_request( &fixture_path(input_name), FixtureRequestOptions { - enable_thinking: Some(false), + enable_thinking: None, no_generation_prompt_when_last_assistant: false, }, ) @@ -36,7 +35,8 @@ fn fixture_path(name: &str) -> PathBuf { } fn test_renderer(use_system_instructions: bool) -> HarmonyChatRenderer { - HarmonyChatRenderer::with_options(PINNED_DATE, use_system_instructions).unwrap() + HarmonyChatRenderer::with_options(Default::default(), PINNED_DATE, use_system_instructions) + .unwrap() } fn render_token_ids(request: &ChatRequest) -> Vec { @@ -138,14 +138,48 @@ fn drops_stale_analysis_fixture() { #[test] fn rejects_invalid_reasoning_effort() { let mut request = ChatRequest::for_test(); - request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); let error = test_renderer(false).render(&request).unwrap_err(); - expect![[r#"chat template error: reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]] + expect![[r#"reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]] .assert_eq(&error.to_report_string()); } +#[test] +fn normalized_reasoning_respects_harmony_capability_and_deployment_effort() { + let renderer = HarmonyChatRenderer::with_options( + [("reasoning_effort".to_string(), serde_json::json!("low"))].into(), + PINNED_DATE, + false, + ) + .unwrap(); + let mut request = ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); + request + .chat_options + .template_kwargs + .insert("enable_thinking".into(), serde_json::json!(true)); + let rendered = renderer.render(&request).unwrap(); + let prompt = harmony_encoding() + .unwrap() + .tokenizer() + .decode_utf8(rendered.prompt.into_token_ids().unwrap()) + .unwrap(); + assert!(prompt.contains("Reasoning: low")); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + "low" + ); + assert_eq!(rendered.effective_template_kwargs["enable_thinking"], true); + + request + .chat_options + .template_kwargs + .insert("thinking".into(), serde_json::json!(false)); + assert!(renderer.render(&request).is_err()); +} + #[test] fn rejects_unknown_tool_response_id() { let request = ChatRequest { diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index ba2b76d10003..911a73928a33 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -16,7 +16,7 @@ use self::format::{ ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption, }; use self::template::{CompiledChatTemplate, TemplateContext}; -use super::{ChatRenderer, RenderedPrompt, effective_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; use crate::error::Result; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; use crate::{ @@ -33,6 +33,31 @@ pub use template::{load_chat_template, resolve_chat_template}; pub use self::format::ChatTemplateContentFormatOption; +/// Extract the effective chat-template kwargs visible to the renderer from the request, +/// using the provided defaults as the base. +fn effective_template_kwargs( + default_template_kwargs: &HashMap, + request: &ChatRequest, +) -> HashMap { + let mut kwargs = default_template_kwargs.clone(); + kwargs.extend(request.chat_options.template_kwargs.clone()); + + if let Some(reasoning_effort) = &request.chat_options.reasoning_effort { + kwargs.insert( + "reasoning_effort".to_string(), + serde_json::json!(reasoning_effort), + ); + if !request.chat_options.template_kwargs.contains_key("enable_thinking") { + kwargs.insert( + "enable_thinking".to_string(), + serde_json::json!(reasoning_effort.as_str() != Some("none")), + ); + } + } + + kwargs +} + /// Template-visible placeholder tokens per supported modality. /// /// A `None` token means the loaded model does not support that modality, and @@ -580,9 +605,10 @@ mod tests { use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken}; use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo}; + use crate::EffortValue; use crate::request::{ ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, - GenerationPromptMode, ReasoningEffort, ResolvedToolContext, + GenerationPromptMode, ResolvedToolContext, }; use crate::{AssistantContentBlock, ChatRenderer, Error, Result}; @@ -1149,7 +1175,7 @@ mod tests { #[test] fn chat_template_reasoning_effort_overrides_template_kwargs() { let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); - request.chat_options.reasoning_effort = Some(ReasoningEffort::Max); + request.chat_options.reasoning_effort = Some(EffortValue::from("max")); request.chat_options.template_kwargs.insert( "reasoning_effort".to_string(), Value::String("low".to_string()), @@ -1181,7 +1207,7 @@ mod tests { #[test] fn chat_template_reasoning_effort_preserves_request_enable_thinking() { let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); - request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); request .chat_options .template_kwargs @@ -1207,6 +1233,70 @@ mod tests { ); } + #[test] + fn chat_template_preserves_typed_effort_values() { + for effort in [ + serde_json::json!(37), + serde_json::json!(0.37), + serde_json::json!("custom"), + ] { + let mut request = sample_request(vec![ChatMessage::user("hello")]); + request.chat_options.reasoning_effort = + Some(serde_json::from_value(effort.clone()).unwrap()); + request.chat_options.template_kwargs = [ + ("reasoning_effort".to_string(), serde_json::json!("high")), + ("enable_thinking".to_string(), serde_json::json!(true)), + ] + .into(); + let rendered = HfChatRenderer::new( + Some("{{ reasoning_effort|tojson }}|{{ enable_thinking }}".to_string()), + HashMap::new(), + ChatTemplateContentFormatOption::Auto, + ) + .unwrap() + .render(&request) + .unwrap(); + assert_eq!(rendered.prompt, Prompt::Text(format!("{effort}|True"))); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + effort + ); + } + } + + #[test] + fn chat_template_keeps_native_control_conflicts_and_raw_effort_values() { + for effort in [ + serde_json::json!(37), + serde_json::json!(null), + serde_json::json!("none"), + ] { + let mut request = sample_request(vec![ChatMessage::user("hello")]); + request.chat_options.template_kwargs = [ + ("thinking".to_string(), serde_json::json!("custom-mode")), + ("enable_thinking".to_string(), serde_json::json!(true)), + ("reasoning_effort".to_string(), effort), + ] + .into(); + let rendered = HfChatRenderer::new( + Some( + "{{ thinking }}|{{ enable_thinking }}|{{ reasoning_effort is none }}" + .to_string(), + ), + HashMap::new(), + ChatTemplateContentFormatOption::Auto, + ) + .unwrap() + .render(&request) + .unwrap(); + assert!(rendered.prompt.into_text().unwrap().starts_with("custom-mode|True|")); + assert_eq!( + rendered.effective_template_kwargs, + request.chat_options.template_kwargs + ); + } + } + #[test] fn qwen3_template_omits_reasoning_for_historical_assistant_messages() { let request = sample_request(vec![ diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs index e8afe66396dd..9feee371aa92 100644 --- a/rust/src/chat/src/renderer/inkling/mod.rs +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -8,8 +8,10 @@ use thiserror_ext::AsReport as _; use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; +use crate::EffortValue; use crate::error::{Error, Result}; +use crate::reasoning::ReasoningControl; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool}; use crate::{AssistantContentBlock, AssistantToolCall}; @@ -35,6 +37,7 @@ const MAX_REASONING_EFFORT: f64 = 0.99; pub struct InklingChatRenderer { tokenizer: DynTokenizer, special: InklingSpecialTokenIds, + default_template_kwargs: HashMap, } #[derive(Debug, Clone, Copy)] @@ -78,9 +81,16 @@ impl InklingSpecialTokenIds { } impl InklingChatRenderer { - pub fn new(tokenizer: DynTokenizer) -> Result { + pub fn new( + tokenizer: DynTokenizer, + default_template_kwargs: HashMap, + ) -> Result { let special = InklingSpecialTokenIds::resolve(tokenizer.as_ref())?; - Ok(Self { tokenizer, special }) + Ok(Self { + tokenizer, + special, + default_template_kwargs, + }) } fn write_text_tokens(&self, out: &mut Vec, text: &str) -> Result<()> { @@ -150,11 +160,6 @@ impl InklingChatRenderer { } fn write_reasoning_effort(&self, out: &mut Vec, effort: f64) -> Result<()> { - if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) { - return Err(Error::ChatTemplate(format!( - "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}" - ))); - } let formatted = format!("{effort:.2}"); let effort = formatted.trim_end_matches('0').trim_end_matches('.'); let effort = if matches!(effort, "0" | "-0") { @@ -294,11 +299,18 @@ impl ChatRenderer for InklingChatRenderer { let mut out = Vec::new(); let mut tool_call_id_to_name = HashMap::new(); - let effective_template_kwargs = request_template_kwargs(request); + let reasoning = resolve_reasoning(ReasoningControl::resolve( + request, + &self.default_template_kwargs, + )?)?; + let effective_template_kwargs = reasoning.template_kwargs(request); let tools = rendered_tools(request); self.write_tool_declarations(&mut out, &tools)?; - let mut reasoning_effort = - resolve_reasoning_effort(effective_template_kwargs.get("reasoning_effort")); + let mut reasoning_effort = if reasoning.is_enabled() { + reasoning.effort().and_then(EffortValue::as_f64) + } else { + Some(0.0) + }; for message in &request.messages { if !matches!( @@ -352,12 +364,33 @@ impl ChatRenderer for InklingChatRenderer { } } -fn resolve_reasoning_effort(value: Option<&Value>) -> Option { +fn resolve_reasoning(control: ReasoningControl) -> Result { + let control = control.fallback(ReasoningControl::enabled("high")); + if !control.is_enabled() { + return Ok(control); + } + let effort = resolve_reasoning_effort(control.effort()); + if let Some(effort) = effort { + if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) { + return Err(Error::InvalidReasoningEffort(format!( + "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}" + ))); + } + if effort == 0.0 { + return Ok(ReasoningControl::Disabled); + } + } + Ok(ReasoningControl::Enabled { + effort: effort.map(EffortValue::try_from).transpose()?, + }) +} + +fn resolve_reasoning_effort(value: Option<&EffortValue>) -> Option { let Some(value) = value else { return Some(0.9); }; match value { - Value::String(name) => match name.as_str() { + EffortValue::String(name) => match name.as_str() { "none" => Some(0.0), "minimal" => Some(0.1), "low" => Some(0.2), @@ -366,8 +399,7 @@ fn resolve_reasoning_effort(value: Option<&Value>) -> Option { "xhigh" | "max" => Some(0.99), _ => None, }, - Value::Number(number) => number.as_f64(), - _ => None, + EffortValue::Number(number) => number.as_f64(), } } diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index 1c4cd7236dbd..d0dd3055e413 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -9,6 +9,7 @@ use serde_json::json; use thiserror_ext::AsReport; use vllm_text::tokenizer::Tokenizer; +use crate::EffortValue; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use super::{ @@ -17,7 +18,7 @@ use super::{ MESSAGE_TOOL, MESSAGE_USER, }; use crate::event::{AssistantContentBlock, AssistantToolCall}; -use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort}; +use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode}; use crate::{ChatRenderer, Error}; struct FixtureTokenizer; @@ -87,7 +88,7 @@ impl Tokenizer for FixtureTokenizer { } fn renderer() -> InklingChatRenderer { - InklingChatRenderer::new(Arc::new(FixtureTokenizer)).unwrap() + InklingChatRenderer::new(Arc::new(FixtureTokenizer), Default::default()).unwrap() } fn render_token_ids(request: &ChatRequest) -> Vec { @@ -105,7 +106,7 @@ fn fixture_request(name: &str) -> ChatRequest { fn inkling_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: Some(false), + enable_thinking: None, no_generation_prompt_when_last_assistant: false, } } @@ -206,13 +207,13 @@ fn renders_request_and_developer_tools_as_tool_declare() { #[test] fn renders_named_reasoning_effort_after_tool_declarations() { for (effort, expected) in [ - (ReasoningEffort::None, "0.0"), - (ReasoningEffort::Minimal, "0.1"), - (ReasoningEffort::Low, "0.2"), - (ReasoningEffort::Medium, "0.7"), - (ReasoningEffort::High, "0.9"), - (ReasoningEffort::XHigh, "0.99"), - (ReasoningEffort::Max, "0.99"), + (EffortValue::from("none"), "0.0"), + (EffortValue::from("minimal"), "0.1"), + (EffortValue::from("low"), "0.2"), + (EffortValue::from("medium"), "0.7"), + (EffortValue::from("high"), "0.9"), + (EffortValue::from("xhigh"), "0.99"), + (EffortValue::from("max"), "0.99"), ] { let mut request = fixture_request("tool_declare_input.json"); request.chat_options.reasoning_effort = Some(effort); @@ -284,6 +285,43 @@ fn defaults_reasoning_effort_to_high() { ); } +#[test] +fn normalized_mode_and_scalar_are_reflected_in_prompt_and_kwargs() { + for (kwargs, expected, enabled) in [ + ( + json!({"enable_thinking": true, "reasoning_effort": "none"}), + "0.9", + true, + ), + ( + json!({"thinking": false, "reasoning_effort": 10}), + "0.0", + false, + ), + (json!({"reasoning_effort": 0}), "0.0", false), + (json!({"reasoning_effort": null}), "0.9", true), + (json!({"reasoning_effort": 0.37}), "0.37", true), + ] { + let mut request = ChatRequest::for_test(); + request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap(); + let rendered = renderer().render(&request).unwrap(); + let prompt = render_symbolic_tokens(&rendered.prompt.into_token_ids().unwrap()); + assert!(prompt.contains(&format!("Thinking effort level: {expected}<|end_message|>"))); + assert_eq!( + rendered.effective_template_kwargs["enable_thinking"], + enabled + ); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + if enabled { + json!(expected.parse::().unwrap()) + } else { + json!("none") + } + ); + } +} + #[test] fn renders_numeric_reasoning_effort_template_kwarg() { let mut request = ChatRequest::for_test(); @@ -301,19 +339,17 @@ fn renders_numeric_reasoning_effort_template_kwarg() { } #[test] -fn ignores_unsupported_reasoning_effort_values() { - for value in [json!(true), json!("invalid"), json!(null)] { - let mut request = ChatRequest::for_test(); - request - .chat_options - .template_kwargs - .insert("reasoning_effort".to_string(), value); +fn ignores_unsupported_reasoning_effort_names() { + let mut request = ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("invalid")); - assert_eq!( - render_symbolic_tokens(&render_token_ids(&request)), - "<|message_user|><|content_text|>test<|end_message|><|message_model|>" - ); - } + assert_eq!( + render_symbolic_tokens(&render_token_ids(&request)), + "<|message_user|><|content_text|>test<|end_message|><|message_model|>" + ); } #[test] diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index 88b4e3ae6101..b434d72c538d 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -11,7 +11,9 @@ use serde_json::{Map, Value, json}; use vllm_tokenizer::Tokenizer; use super::super::MediaPartSource; +use crate::EffortValue; use crate::error::{Error, Result}; +use crate::reasoning::ReasoningControl; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, }; @@ -76,9 +78,10 @@ impl<'a> K3TokenWriter<'a> { pub(super) fn render_request_with_media_order( request: &ChatRequest, tokenizer: &dyn Tokenizer, + reasoning: &ReasoningControl, ) -> Result<(Vec, Vec)> { - let thinking = thinking_enabled(request)?; - let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; + let thinking = reasoning.is_enabled(); + let thinking_effort = reasoning.effort().and_then(EffortValue::as_str); let tools = request_tools(request); let mut out = K3TokenWriter::new(tokenizer); @@ -244,58 +247,26 @@ fn request_tools(request: &ChatRequest) -> &[ChatTool] { request.initial_tools() } -fn thinking_enabled(request: &ChatRequest) -> Result { - if let Some(thinking) = request.parse_template_bool("thinking")? { - return Ok(thinking); - } - if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { - return Ok(enable_thinking); - } - if let Some(reasoning_effort) = request.chat_options.reasoning_effort { - return Ok(reasoning_effort != crate::request::ReasoningEffort::None); - } - Ok(request - .chat_options - .template_kwargs - .get("reasoning_effort") - .and_then(Value::as_str) - != Some("none")) -} - -fn thinking_effort(request: &ChatRequest) -> Result { - let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") { - value.as_str().ok_or_else(|| { - Error::ChatTemplate(format!( - "template kwarg `thinking_effort` must be a string, got {value}" - )) - })? - } else if let Some(effort) = request.chat_options.reasoning_effort { - if effort == crate::request::ReasoningEffort::None { - DEFAULT_THINKING_EFFORT - } else { - effort.as_str() - } - } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { +/// Resolve standard reasoning controls and validate K3's supported effort grades. +pub(super) fn resolve_reasoning( + request: &ChatRequest, + defaults: &HashMap, +) -> Result { + let reasoning = ReasoningControl::resolve(request, defaults)? + .fallback(ReasoningControl::enabled(DEFAULT_THINKING_EFFORT)); + if let Some(value) = reasoning.effort() { let effort = value.as_str().ok_or_else(|| { - Error::ChatTemplate(format!( - "template kwarg `reasoning_effort` must be a string, got {value}" + Error::InvalidReasoningEffort(format!( + "Kimi K3 reasoning_effort must be a string, got {value}" )) })?; - if effort == "none" { - DEFAULT_THINKING_EFFORT - } else { - effort + if !VALID_THINKING_EFFORTS.contains(&effort) { + return Err(Error::InvalidReasoningEffort(format!( + "unsupported reasoning_effort={effort:?}; supported values are `low`, `high`, and `max`" + ))); } - } else { - DEFAULT_THINKING_EFFORT - }; - - if !VALID_THINKING_EFFORTS.contains(&effort) { - return Err(Error::ChatTemplate(format!( - "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`" - ))); } - Ok(effort.to_string()) + Ok(reasoning) } fn content_is_empty(content: &ChatContent) -> bool { diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs index 90750eeac406..c39ea2985a07 100644 --- a/rust/src/chat/src/renderer/kimi_k3/mod.rs +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -7,10 +7,13 @@ mod encoding; #[cfg(test)] mod tests; +use std::collections::HashMap; + +use serde_json::Value; use vllm_text::Prompt; use vllm_text::tokenizer::DynTokenizer; -use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use super::{ChatRenderer, RenderedPrompt}; use crate::Result; use crate::request::ChatRequest; @@ -18,25 +21,42 @@ use crate::request::ChatRequest; #[derive(Clone)] pub struct KimiK3ChatRenderer { tokenizer: DynTokenizer, + default_template_kwargs: HashMap, } impl KimiK3ChatRenderer { /// Create a Kimi K3 renderer. - pub fn new(tokenizer: DynTokenizer) -> Self { - Self { tokenizer } + pub fn new(tokenizer: DynTokenizer, default_template_kwargs: HashMap) -> Self { + Self { + tokenizer, + default_template_kwargs, + } } } impl ChatRenderer for KimiK3ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; - let (token_ids, media_order) = - encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?; + let reasoning = encoding::resolve_reasoning(request, &self.default_template_kwargs)?; + let mut effective_template_kwargs = reasoning.template_kwargs(request); + // Export K3's native effort key from the resolved standard controls. + // Input `thinking_effort` is ignored; clear it when reasoning is disabled. + if let Some(effort) = reasoning.effort() { + effective_template_kwargs + .insert("thinking_effort".to_string(), serde_json::json!(effort)); + } else { + effective_template_kwargs.remove("thinking_effort"); + } + let (token_ids, media_order) = encoding::render_request_with_media_order( + request, + self.tokenizer.as_ref(), + &reasoning, + )?; Ok(RenderedPrompt { prompt: Prompt::TokenIds(token_ids), media_order: Some(media_order), - effective_template_kwargs: request_template_kwargs(request), + effective_template_kwargs, }) } } diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index d0ecbe44d38b..5aadb3d403c6 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -15,11 +15,10 @@ use vllm_tokenizer::test_utils::TestTokenizer; use super::KimiK3ChatRenderer; use crate::ChatRenderer; +use crate::EffortValue; use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ - ChatContentPart, ChatMessage, ChatTool, GenerationPromptMode, ReasoningEffort, -}; +use crate::request::{ChatContentPart, ChatMessage, ChatTool, GenerationPromptMode}; use crate::{AssistantContentBlock, AssistantToolCall}; const OPEN_ID: u32 = 256; @@ -37,8 +36,70 @@ fn test_tokenizer() -> TestTokenizer { .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID) } +#[test] +fn standard_effort_controls_prompt_and_exported_native_effort() { + let tokenizer = Arc::new(test_tokenizer()); + for (kwargs, typed, defaults, effort) in [ + (json!({}), None, json!({}), "max"), + ( + json!({}), + Some(EffortValue::from("low")), + json!({"reasoning_effort": "high"}), + "low", + ), + ( + json!({"enable_thinking": true, "reasoning_effort": "none"}), + None, + json!({"reasoning_effort": "low"}), + "low", + ), + ( + json!({}), + Some(EffortValue::from("none")), + json!({"reasoning_effort": "low"}), + "none", + ), + ( + json!({"thinking": true}), + Some(EffortValue::from("none")), + json!({}), + "max", + ), + ] { + let renderer = + KimiK3ChatRenderer::new(tokenizer.clone(), serde_json::from_value(defaults).unwrap()); + let mut request = crate::ChatRequest::for_test(); + request.chat_options.reasoning_effort = typed; + request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap(); + let rendered = renderer.render(&request).unwrap(); + let prompt = tokenizer.decode(&rendered.prompt.into_token_ids().unwrap(), false).unwrap(); + let enabled = effort != "none"; + assert_eq!(prompt.ends_with("<|open|>think<|sep|>"), enabled); + assert_eq!( + rendered.effective_template_kwargs["reasoning_effort"], + effort + ); + assert_eq!( + rendered.effective_template_kwargs["enable_thinking"], + enabled + ); + if enabled { + assert!(prompt.contains(&format!("thinking_effort={effort}"))); + assert_eq!( + rendered.effective_template_kwargs["thinking_effort"], + effort + ); + } else { + assert!(!rendered.effective_template_kwargs.contains_key("thinking_effort")); + } + } +} + fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec { - let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt; + let prompt = KimiK3ChatRenderer::new(tokenizer, Default::default()) + .render(request) + .unwrap() + .prompt; let Prompt::TokenIds(token_ids) = prompt else { panic!("kimi k3 renderer should return token IDs") }; @@ -67,7 +128,11 @@ fn kimi_k3_fixture_options() -> FixtureRequestOptions { fn assert_golden(name: &str) { let input_name = format!("{name}_input.json"); - let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + let mut request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + // Reference fixtures use encoder-native keys; Rust accepts standard effort input. + if let Some(effort) = request.chat_options.template_kwargs.remove("thinking_effort") { + request.chat_options.reasoning_effort = Some(EffortValue::try_from(effort).unwrap()); + } let rendered = render_request(&request); expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered); } @@ -310,7 +375,7 @@ fn media_order_follows_reordered_tool_results() { ), ]; request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; - let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())); + let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()), Default::default()); let rendered = renderer.render(&request).unwrap(); let media_order = rendered.media_order.unwrap(); @@ -371,19 +436,18 @@ fn translates_standard_thinking_kwargs() { } #[test] -fn native_k3_kwargs_take_precedence() { +fn thinking_alias_takes_precedence() { let mut request = crate::request::ChatRequest::for_test(); request.chat_options.template_kwargs.extend([ ("thinking".to_string(), json!(true)), ("enable_thinking".to_string(), json!(false)), - ("thinking_effort".to_string(), json!("low")), ("reasoning_effort".to_string(), json!("high")), ]); let rendered = render_request(&request); - assert!(rendered.contains("thinking_effort=low")); - assert!(!rendered.contains("thinking_effort=high")); + assert!(rendered.contains("thinking_effort=high")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); } #[test] @@ -431,7 +495,7 @@ fn enable_thinking_true_overrides_standard_none() { #[test] fn typed_none_disables_thinking() { let mut request = crate::request::ChatRequest::for_test(); - request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + request.chat_options.reasoning_effort = Some(EffortValue::from("none")); let rendered = render_request(&request); @@ -440,20 +504,17 @@ fn typed_none_disables_thinking() { } #[test] -fn rejects_removed_medium_thinking_effort() { +fn rejects_unsupported_standard_reasoning_effort() { let mut request = crate::request::ChatRequest::for_test(); - request - .chat_options - .template_kwargs - .insert("thinking_effort".to_string(), json!("medium")); + request.chat_options.reasoning_effort = Some(EffortValue::from("medium")); - let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())) + let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()), Default::default()) .render(&request) .unwrap_err(); expect![[r#" - ChatTemplate( - "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`", + InvalidReasoningEffort( + "unsupported reasoning_effort=\"medium\"; supported values are `low`, `high`, and `max`", ) "#]] .assert_debug_eq(&error); diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index bc1831c375a9..347ec43d03e1 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -4,11 +4,11 @@ use std::collections::HashMap; use std::sync::Arc; -use serde_json::{Value, json}; +use serde_json::Value; use vllm_text::Prompt; use crate::error::Result; -use crate::request::{ChatRequest, ReasoningEffort}; +use crate::request::ChatRequest; mod deepseek; pub mod deepseek_v32; @@ -62,33 +62,3 @@ pub trait ChatRenderer: Send + Sync { /// Shared trait-object form of [`ChatRenderer`]. pub type DynChatRenderer = Arc; - -/// Extract the effective chat-template kwargs visible to the renderer from the request, -/// using the provided defaults as the base. -pub(crate) fn effective_template_kwargs( - default_template_kwargs: &HashMap, - request: &ChatRequest, -) -> HashMap { - let mut kwargs = default_template_kwargs.clone(); - kwargs.extend(request.chat_options.template_kwargs.clone()); - - if let Some(reasoning_effort) = request.chat_options.reasoning_effort { - kwargs.insert( - "reasoning_effort".to_string(), - Value::String(reasoning_effort.as_str().to_string()), - ); - if !request.chat_options.template_kwargs.contains_key("enable_thinking") { - kwargs.insert( - "enable_thinking".to_string(), - json!(reasoning_effort != ReasoningEffort::None), - ); - } - } - - kwargs -} - -/// Extract the effective chat-template kwargs visible to the renderer from the request. -pub(crate) fn request_template_kwargs(request: &ChatRequest) -> HashMap { - effective_template_kwargs(&HashMap::new(), request) -} diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 1090a4039f5e..2c59ed7a22c4 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -8,10 +8,11 @@ use std::path::Path; use serde::Deserialize; use serde_json::Value; +use crate::EffortValue; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, - GenerationPromptMode, ReasoningEffort, ResolvedToolContext, + GenerationPromptMode, ResolvedToolContext, }; /// Options for constructing a [`ChatRequest`] from a fixture file. @@ -47,7 +48,7 @@ pub(crate) struct FixtureRequest { tools: Vec, messages: Vec, add_generation_prompt: Option, - reasoning_effort: Option, + reasoning_effort: Option, /// Standard response format passed to model-specific renderers. #[serde(default)] response_format: Option, diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index 50e00bbc43a5..e3bbcaa2aa51 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -11,9 +11,9 @@ pub use vllm_parser::tool::Tool as ChatTool; use vllm_text::TextDecodeOptions; pub use vllm_text::{PromptTruncation, SamplingParams}; -use crate::AssistantMessageExt; use crate::error::{Error, Result}; use crate::event::{AssistantContentBlock, AssistantMessage}; +use crate::{AssistantMessageExt, EffortValue}; /// Role label for one text-only chat message. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -362,33 +362,6 @@ pub enum GenerationPromptMode { NoGenerationPrompt, } -/// Effort level for reasoning models. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum ReasoningEffort { - None, - Minimal, - Low, - Medium, - High, - XHigh, - Max, -} - -impl ReasoningEffort { - pub fn as_str(self) -> &'static str { - match self { - Self::None => "none", - Self::Minimal => "minimal", - Self::Low => "low", - Self::Medium => "medium", - Self::High => "high", - Self::XHigh => "xhigh", - Self::Max => "max", - } - } -} - /// Chat-template-related request options. /// /// These are the small subset of chat controls that currently affect prompt @@ -404,8 +377,10 @@ pub struct ChatOptions { /// used instead of the model's default chat template. pub chat_template: Option, - /// Effort level exposed to chat templates for reasoning models. - pub reasoning_effort: Option, + /// Model-specific reasoning effort, typically `none`, `minimal`, `low`, + /// `medium`, `high`, `xhigh`, or `max`. Supported names and numeric ranges + /// are validated by the selected renderer or HF template. + pub reasoning_effort: Option, /// Standard response format available to model-specific renderers. #[serde(default)] @@ -673,29 +648,6 @@ impl ChatRequest { self.tool_context.parsing_enabled() } - /// Return the request-level thinking toggle when explicitly requested. - /// - /// We currently accept the two request kwargs `thinking` and - /// `enable_thinking`. Both must be booleans when present. If both are - /// present, they must have the same value. If neither key is provided, - /// return `None`. - pub(crate) fn enable_thinking(&self) -> Result> { - let thinking = self.parse_template_bool("thinking")?; - let enable_thinking = self.parse_template_bool("enable_thinking")?; - - match (thinking, enable_thinking) { - (None, None) => Ok(None), - (Some(thinking), Some(enable_thinking)) if thinking != enable_thinking => { - Err(Error::ChatTemplate( - "template kwargs `thinking` and `enable_thinking` must match when both are set" - .to_string(), - )) - } - (Some(thinking), _) => Ok(Some(thinking)), - (None, Some(enable_thinking)) => Ok(Some(enable_thinking)), - } - } - pub(crate) fn parse_template_bool(&self, key: &str) -> Result> { match self.chat_options.template_kwargs.get(key) { None => Ok(None), @@ -726,7 +678,7 @@ mod tests { use serde_json::{json, to_value}; use super::{ - ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, + ChatContent, ChatContentPart, ChatMessage, ChatRole, ChatTool, ChatToolChoice, ResolvedToolContext, }; use crate::Error; @@ -919,56 +871,4 @@ mod tests { Error::ToolChoiceFunctionNotFound { name } if name == "missing" )); } - - #[test] - fn enable_thinking_is_none_when_no_kwargs_are_present() { - let request = ChatRequest::for_test(); - assert_eq!(request.enable_thinking().unwrap(), None); - } - - #[test] - fn enable_thinking_accepts_matching_duplicate_kwargs() { - let mut request = ChatRequest::for_test(); - request.chat_options.template_kwargs.insert("thinking".to_string(), json!(true)); - request - .chat_options - .template_kwargs - .insert("enable_thinking".to_string(), json!(true)); - - assert_eq!(request.enable_thinking().unwrap(), Some(true)); - } - - #[test] - fn enable_thinking_rejects_non_boolean_kwargs() { - let mut request = ChatRequest::for_test(); - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), json!("yes")); - - assert!(matches!( - request.enable_thinking(), - Err(Error::ChatTemplate(message)) - if message.contains("`thinking` must be a boolean") - )); - } - - #[test] - fn enable_thinking_rejects_conflicting_duplicate_kwargs() { - let mut request = ChatRequest::for_test(); - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), json!(false)); - request - .chat_options - .template_kwargs - .insert("enable_thinking".to_string(), json!(true)); - - assert!(matches!( - request.enable_thinking(), - Err(Error::ChatTemplate(message)) - if message.contains("`thinking` and `enable_thinking` must match") - )); - } } diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 9f083a1a614e..3280cd1dd503 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -175,16 +175,22 @@ mod tests { } #[test] - fn invalid_reasoning_effort_maps_to_invalid_request() { - let error = vllm_chat::Error::InvalidReasoningEffort( - "DeepSeek V4.1 reasoning_effort must be within [1, 100]".to_string(), - ); - let api_error = chat_submit_error("failed to submit chat request", error); - assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); - assert_eq!( - api_error.to_error_response().error.error_type, - "invalid_request_error" - ); + fn invalid_reasoning_parameters_map_to_invalid_request() { + for error in [ + vllm_chat::Error::InvalidReasoningEffort( + "DeepSeek V4.1 reasoning_effort must be within [1, 100]".to_string(), + ), + vllm_chat::Error::InvalidReasoningControl { + message: "template kwarg `thinking` must be a boolean".to_string(), + }, + ] { + let api_error = chat_submit_error("failed to submit chat request", error); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + assert_eq!( + api_error.to_error_response().error.error_type, + "invalid_request_error" + ); + } } #[test] diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 51b0fb68ad14..18158566360b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -180,7 +180,7 @@ pub(super) fn prepare_chat_request( chat_options: ChatOptions { generation_prompt_mode, chat_template: request.chat_template, - reasoning_effort: request.reasoning_effort, + reasoning_effort: request.reasoning_effort.map(|effort| effort.as_str().into()), response_format, template_kwargs, }, @@ -488,6 +488,50 @@ mod tests { } } + #[test] + fn chat_http_reasoning_effort_preserves_omission_none_and_kwargs() { + for effort in [None, Some(json!(null))].into_iter().chain( + ["none", "minimal", "low", "medium", "high", "xhigh", "max"] + .map(|name| Some(json!(name))), + ) { + let kwargs = json!({"reasoning_effort": 37, "thinking": true}); + let mut value = json!({ + "messages": [{"role": "user", "content": "hello"}], + "chat_template_kwargs": kwargs, + }); + if let Some(effort) = &effort { + value["reasoning_effort"] = effort.clone(); + } + let request: ChatCompletionRequest = serde_json::from_value(value).unwrap(); + let prepared = prepare_chat_request( + request, + &served(&["test-model"]), + request_context(&HeaderMap::new(), None), + ) + .unwrap(); + let options = prepared.chat_request.chat_options; + assert_eq!( + serde_json::to_value(options.reasoning_effort).unwrap(), + effort.unwrap_or(serde_json::Value::Null) + ); + assert_eq!( + serde_json::to_value(options.template_kwargs).unwrap(), + kwargs + ); + } + } + + #[test] + fn chat_http_reasoning_effort_rejects_model_extensions_at_top_level() { + for effort in [json!(37), json!("custom")] { + let request = json!({ + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": effort, + }); + assert!(serde_json::from_value::(request).is_err()); + } + } + #[test] fn chat_http_request_defaults_missing_or_null_model() { for model in [None, Some(serde_json::Value::Null)] { @@ -747,7 +791,7 @@ mod tests { ); let tokenizer = Arc::new(TestTokenizer::new()); - let prompt = KimiK3ChatRenderer::new(tokenizer.clone()) + let prompt = KimiK3ChatRenderer::new(tokenizer.clone(), Default::default()) .render(&prepared.chat_request) .expect("Kimi K3 rendering succeeds") .prompt; diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index a371462694ca..13c7c7233f14 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -9,15 +9,15 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use serde_with::SerializeDisplay; use validator::Validate; -use vllm_chat::ReasoningEffort; use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::TruncationSide; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ - ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope, - StringOrArray, Tool, ToolCall, ToolCallDelta, ToolChoice, Usage, default_true, - deserialize_request_top_k, validate_messages, validate_stop, validate_top_p_value, + ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, ReasoningEffort, StreamOptions, + StreamResponseEnvelope, StringOrArray, Tool, ToolCall, ToolCallDelta, ToolChoice, Usage, + default_true, deserialize_request_top_k, validate_messages, validate_stop, + validate_top_p_value, }; /// vLLM-compatible request type for the Chat Completions API. diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 7260155799b6..1bfc1d9a091d 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -125,8 +125,8 @@ fn validate_function_tools(tools: &[Tool], param: &'static str) -> Result<(), Ap mod tests { use std::collections::HashMap; + use crate::routes::openai::utils::types::ReasoningEffort; use serde_json::json; - use vllm_chat::ReasoningEffort; use super::validate_request_compat; use crate::routes::openai::chat_completions::types::ChatCompletionRequest; diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs index e9aab2d2b297..5554b5d36668 100644 --- a/rust/src/server/src/routes/openai/utils/types.rs +++ b/rust/src/server/src/routes/openai/utils/types.rs @@ -39,6 +39,36 @@ where } } +/// Effort level for reasoning models. +/// +/// Fixed OpenAI HTTP request grades. Request conversion maps these names to +/// the model-independent `vllm_chat::EffortValue` before renderer validation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ReasoningEffort { + None, + Minimal, + Low, + Medium, + High, + XHigh, + Max, +} + +impl ReasoningEffort { + pub fn as_str(self) -> &'static str { + match self { + Self::None => "none", + Self::Minimal => "minimal", + Self::Low => "low", + Self::Medium => "medium", + Self::High => "high", + Self::XHigh => "xhigh", + Self::Max => "max", + } + } +} + // ============================================================================ // String/Array Utilities // ============================================================================ From ed769b4a4520d9ca46657bd9e1ac414a3bbb9cdd Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:29:10 -0400 Subject: [PATCH 360/774] [A2A] Add DeepEPv2 to SP Supported List (#57210) Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw --- vllm/config/parallel.py | 1 + vllm/model_executor/layers/fused_moe/all2all_utils.py | 1 + .../layers/fused_moe/prepare_finalize/deepep_v2.py | 8 ++++++-- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f66af2f30b8d..f0a2b7e45785 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -717,6 +717,7 @@ def use_sequence_parallel_moe(self) -> bool: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", + "deepep_v2", "flashinfer_nvlink_one_sided", "mori_high_throughput", "mori_low_latency", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6984e34e56fd..775f595bfa81 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -276,6 +276,7 @@ def maybe_make_prepare_finalize( num_topk=moe.experts_per_token, use_fp8_dispatch=use_fp8_dispatch, use_cudagraph=use_cudagraph, + sp_size=moe.moe_parallel_config.sp_size, ) elif moe.use_mori_kernels: diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index 8bd215423da4..a584ab5f52ff 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -18,7 +18,7 @@ swizzle_mxfp8_scale, ) from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import round_up +from vllm.utils.math_utils import cdiv, round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, dbo_enabled, @@ -114,6 +114,7 @@ def __init__( num_topk: int, use_fp8_dispatch: bool = False, use_cudagraph: bool = False, + sp_size: int = 1, ): super().__init__() self.buffer = buffer @@ -124,6 +125,7 @@ def __init__( self.num_topk = num_topk self.use_fp8_dispatch = use_fp8_dispatch self.use_cudagraph = use_cudagraph + self.sp_size = sp_size # DBO microbatching: one handle slot per micro-batch. self.handles: list[deep_ep.EPHandle | None] = [None, None] @@ -192,12 +194,14 @@ def _do_dispatch( # high concurrency). Round up to a power of 2 instead: this bounds the # set to ~log2(max_num_batched_tokens) values (compiled once, then # cached) while staying small for decode (e.g. 1 token -> 1) and capped - # at the buffer's init capacity for prefill. + # at the buffer's init capacity for prefill. With sequence parallelism, + # each EP rank holds a ceil(n / sp_size) shard of its DP rank's batch. num_max_tokens_per_rank = None if not do_expand: dp_meta = get_forward_context().dp_metadata if dp_meta is not None: n = int(dp_meta.num_tokens_across_dp_cpu.max()) + n = cdiv(n, self.sp_size) else: n = tokens.shape[0] num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length() From 19b6ff62f24d477a2760602b9e94eb6e4dd7dcdd Mon Sep 17 00:00:00 2001 From: dongluw <108290936+dongluw@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:29:50 -0400 Subject: [PATCH 361/774] [ROCm][Bugfix] Ignore descales for unquantized AITER caches (#56726) Signed-off-by: donglu Co-authored-by: Andreas Karatzas --- tests/v1/attention/test_attention_backends.py | 66 +++++++++++++++++-- vllm/v1/attention/backends/rocm_aiter_fa.py | 41 +++++++----- 2 files changed, 85 insertions(+), 22 deletions(-) diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 66a1c5c08836..b02c85f68202 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -243,14 +243,19 @@ def create_and_prepopulate_kv_cache( class MockAttentionLayer: """A mock attention layer for testing.""" - def __init__(self, device: torch.device): + def __init__( + self, + device: torch.device, + k_scale: float = 1.0, + v_scale: float = 1.0, + ): self._q_scale = torch.tensor(1.0, device=device) - self._k_scale = torch.tensor(1.0, device=device) - self._v_scale = torch.tensor(1.0, device=device) + self._k_scale = torch.tensor(k_scale, device=device) + self._v_scale = torch.tensor(v_scale, device=device) # Add float versions for flashinfer self._q_scale_float = 1.0 - self._k_scale_float = 1.0 - self._v_scale_float = 1.0 + self._k_scale_float = k_scale + self._v_scale_float = v_scale def _clone_kv_cache_in_layout( @@ -285,6 +290,8 @@ def run_attention_backend( kv_cache_dtype: str = "auto", sinks: torch.Tensor | None = None, use_cuda_graph: bool = False, + layer_k_scale: float = 1.0, + layer_v_scale: float = 1.0, ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0") @@ -356,7 +363,7 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): ) # Create mock layer and output buffer - mock_layer = MockAttentionLayer(device) + mock_layer = MockAttentionLayer(device, layer_k_scale, layer_v_scale) output = torch.empty_like(query) if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input: @@ -386,9 +393,11 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): ) graph.replay() else: - output = impl.forward( + backend_output = impl.forward( mock_layer, query, key, value, kv_cache, attn_metadata, output=output ) + if backend_output is not None: + output = backend_output return output @@ -413,6 +422,8 @@ def _test_backend_correctness( model_dtype: torch.dtype | None = None, max_num_seqs: int | None = None, max_num_batched_tokens: int | None = None, + layer_k_scale: float = 1.0, + layer_v_scale: float = 1.0, ): """Test that all backends produce similar outputs to a reference implementation using FlexAttention or an explicit attention-sink reference. @@ -687,6 +698,8 @@ def _test_backend_correctness( kv_cache_dtype=kv_cache_dtype, sinks=sinks, use_cuda_graph=use_cuda_graph, + layer_k_scale=layer_k_scale, + layer_v_scale=layer_v_scale, ) # Check shape and dtype consistency @@ -1317,6 +1330,45 @@ def sliding_window_mask_mod( ) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +def test_rocm_aiter_fa_unquantized_cache_ignores_kv_scales( + default_vllm_config, +): + """Stale scale tensors must not affect an unquantized KV cache.""" + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("AITER is required") + + def sliding_window_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + *, + context_len: int, + sliding_window: int, + ): + causal_mask = q_idx + context_len >= kv_idx + window_mask = q_idx + context_len - kv_idx < sliding_window + return causal_mask & window_mask + + model = "microsoft/Phi-tiny-MoE-instruct" + batch_spec = BATCH_SPECS["small_decode"] + model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens)) + sliding_window = model_config.get_sliding_window() + assert sliding_window is not None + + _test_backend_correctness( + batch_spec, + model, + [AttentionBackendEnum.ROCM_AITER_FA], + partial(sliding_window_mask_mod, sliding_window=sliding_window), + layer_k_scale=0.25, + layer_v_scale=0.5, + ) + + @pytest.mark.parametrize( "batch_spec_name", [ diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 1615f883bc02..4ce8b23d3119 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -989,6 +989,20 @@ def __init__( "cross-attention." ) + def _get_kv_cache_descales( + self, + layer: AttentionLayer, + num_decodes: int, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + # AITER treats non-null descales as an instruction to dequantize. + if not is_quantized_kv_cache(self.kv_cache_dtype): + return None, None + descale_shape = (num_decodes, self.num_kv_heads) + return ( + layer._k_scale.expand(descale_shape), + layer._v_scale.expand(descale_shape), + ) + def extend_for_sliding_window( self, attn_metadata: AiterFlashAttentionMetadata, @@ -1367,6 +1381,9 @@ def forward( or decode_max_query_len > 1 or self.sinks is not None ): + k_descale, v_descale = self._get_kv_cache_descales( + layer, num_decodes + ) assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( "Shuffle KV cache layout is not supported with sliding " "window, sinks, or speculative decoding (multi-token decode)." @@ -1376,7 +1393,6 @@ def forward( flash_attn_with_kvcache, ) - descale_shape = (num_decodes, key_cache.shape[2]) decode_query = query[:num_decode_tokens].reshape( num_decodes, decode_max_query_len, @@ -1393,8 +1409,8 @@ def forward( window_size=self.sliding_window, softcap=self.logits_soft_cap, q_descale=None, - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), + k_descale=k_descale, + v_descale=v_descale, page_table=attn_metadata.block_table[:num_decodes], ) output[:num_decode_tokens].copy_( @@ -1412,10 +1428,6 @@ def forward( unified_attention, ) - descale_shape = ( - num_decodes, - key_cache.shape[2], - ) unified_attention( q=query[:num_decode_tokens], k=key_cache, @@ -1434,8 +1446,8 @@ def forward( block_table=attn_metadata.block_table[:num_decodes], softcap=self.logits_soft_cap, q_descale=None, - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), + k_descale=k_descale, + v_descale=v_descale, sinks=self.sinks, ) return @@ -1449,6 +1461,9 @@ def forward( use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI if use_unified_attention: + k_descale, v_descale = self._get_kv_cache_descales( + layer, num_decodes + ) assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( "unified_attention fallback with shuffle layout " "is not supported yet." @@ -1460,10 +1475,6 @@ def forward( decode_cu_seqlens_q = attn_metadata.query_start_loc[ : num_decodes + 1 ] - descale_shape = ( - num_decodes, - key_cache.shape[2], - ) unified_attention( q=query[:num_decode_tokens], k=key_cache, @@ -1480,8 +1491,8 @@ def forward( block_table=attn_metadata.block_table[:num_decodes], softcap=self.logits_soft_cap, q_descale=None, - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), + k_descale=k_descale, + v_descale=v_descale, ) elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): _, num_heads, head_size = query.shape From 28ec03bc6e51b17cd9a6b48edc3819e542c2ebb1 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:54:43 -0700 Subject: [PATCH 362/774] [Bugfix][Rust Frontend] Accept appended EngineCoreOutput fields (#56533) Co-authored-by: Codex Signed-off-by: Alec Flowers --- .../engine-core-client/src/protocol/mod.rs | 1 + .../engine-core-client/src/protocol/output.rs | 60 +++++++++++++++++++ .../src/protocol/serde_utils.rs | 42 +++++++++++++ .../engine-core-client/src/tests/client.rs | 9 +++ .../src/tests/python_compat.py | 24 ++++++++ 5 files changed, 136 insertions(+) create mode 100644 rust/src/engine-core-client/src/protocol/serde_utils.rs diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 23bd09d8d8a2..8f4ee0612736 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -30,6 +30,7 @@ pub mod multimodal; pub mod output; pub mod request; pub mod sampling; +mod serde_utils; pub mod stats; pub mod structured_outputs; pub mod tensor; diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 85de8c4e3945..c8a8705435c8 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -9,7 +9,9 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_default::DefaultFromSerde; use serde_repr::{Deserialize_repr, Serialize_repr}; use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use serde_with::serde_as; +use super::serde_utils::AllowTrailingFields; use super::utility::UtilityOutput; use crate::error::{Error, Result, ext_value_decode}; use crate::protocol::logprobs::MaybeWireLogprobs; @@ -156,11 +158,14 @@ impl EngineCoreOutput { /// /// Original Python definition: /// +#[serde_as] #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] struct WireEngineCoreOutputs { #[serde(default)] engine_index: u32, /// Outputs grouped for this client in the current engine tick. + // Match msgspec's append-only schema evolution, including Omni extensions. + #[serde_as(deserialize_as = "Vec")] #[serde(default)] outputs: Vec, #[serde(default)] @@ -404,6 +409,61 @@ mod tests { ); } + #[test] + fn engine_core_outputs_ignore_appended_fields() { + let output = EngineCoreOutput { + request_id: "req-1".into(), + new_token_ids: vec![42], + finish_reason: Some(EngineCoreFinishReason::Length), + ..Default::default() + }; + let mut fields = crate::protocol::decode_value(&encode_msgpack(&output).unwrap()) + .unwrap() + .as_array() + .unwrap() + .clone(); + // Omni appends multimodal_output, is_segment_finished, and + // new_prompt_len_snapshot. Unknown payloads may contain nested values. + fields.extend([ + OpaqueValue::Map(vec![( + "audio".into(), + OpaqueValue::Array(vec![OpaqueValue::Ext(1, vec![0, 1])]), + )]), + false.into(), + OpaqueValue::Nil, + ]); + let wire = (0, vec![OpaqueValue::Array(fields); 2]); + let frames = [Bytes::from(encode_msgpack(&wire).unwrap())]; + let decoded = decode_engine_core_outputs(&frames).unwrap(); + // A second output also checks that the first output's tail was consumed. + assert_eq!(decoded.as_request_batch().unwrap().outputs, vec![output; 2]); + } + + #[test] + fn engine_core_output_requires_valid_known_fields() { + for fields in [ + vec![], + vec!["req-1".into()], + vec!["req-1".into(), OpaqueValue::Nil], + vec!["req-1".into(), OpaqueValue::Array(vec![]), false.into()], + ] { + let wire = (0, vec![OpaqueValue::Array(fields)]); + let frames = [Bytes::from(encode_msgpack(&wire).unwrap())]; + assert!(decode_engine_core_outputs(&frames).is_err()); + } + let wire = (0, vec![("req-1", vec![42_u32])]); + let frames = [Bytes::from(encode_msgpack(&wire).unwrap())]; + let decoded = decode_engine_core_outputs(&frames).unwrap(); + assert_eq!( + decoded.as_request_batch().unwrap().outputs, + vec![EngineCoreOutput { + request_id: "req-1".into(), + new_token_ids: vec![42], + ..Default::default() + }] + ); + } + #[test] fn engine_core_outputs_classify_request_batch() { let outputs = WireEngineCoreOutputs { diff --git a/rust/src/engine-core-client/src/protocol/serde_utils.rs b/rust/src/engine-core-client/src/protocol/serde_utils.rs new file mode 100644 index 000000000000..708353e42ea0 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/serde_utils.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::fmt; +use std::marker::PhantomData; + +use serde::de::{IgnoredAny, SeqAccess, Visitor, value::SeqAccessDeserializer}; +use serde::{Deserialize, Deserializer}; +use serde_with::DeserializeAs; + +/// Decode an array-like struct using its derived field types and defaults, +/// skipping appended fields as Python's msgspec does. Opt in only for schemas +/// whose extensions preserve the known field order and semantics. +pub(super) struct AllowTrailingFields; + +impl<'de, T: Deserialize<'de>> DeserializeAs<'de, T> for AllowTrailingFields { + fn deserialize_as(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct ArrayVisitor(PhantomData); + + impl<'de, T: Deserialize<'de>> Visitor<'de> for ArrayVisitor { + type Value = T; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an array-like struct") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let value = T::deserialize(SeqAccessDeserializer::new(&mut seq))?; + while seq.next_element::()?.is_some() {} + Ok(value) + } + } + + deserializer.deserialize_seq(ArrayVisitor(PhantomData)) + } +} diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 3db573d2aadc..739c34d6b751 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2641,6 +2641,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { let multi_connector_stats_hex = lines.next().expect("missing MultiConnector stats fixture line"); let ready_response_hex = lines.next().expect("missing ready response fixture line"); + let extended_outputs_hex = lines.next().expect("missing extended outputs fixture line"); let request_bytes = hex::decode(request_hex).unwrap(); let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap(); @@ -2722,6 +2723,14 @@ fn python_msgpack_fixtures_match_rust_encoding() { )); let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap(); + // Match msgspec's base-schema result for the same extended Python message. + let extended_frames = [bytes::Bytes::from( + hex::decode(extended_outputs_hex).unwrap(), + )]; + assert_eq!( + decode_engine_core_outputs(&extended_frames).unwrap(), + decoded_outputs + ); expect_test::expect![[r#" RequestBatch( RequestBatchOutputs { diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 0cebca972d4d..5013f9e2a57b 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -101,6 +101,14 @@ class EngineCoreOutput( num_nans_in_logits: int = 0 mm_cache_miss_hashes: list[str] | None = None new_sampling_mask: object | None = None + spec_decode_metrics: object | None = None + + +class ExtendedEngineCoreOutput(EngineCoreOutput): + # Match Omni's append-only schema without requiring Omni or torch. + multimodal_output: dict[str, object] | None = None + is_segment_finished: bool = False + new_prompt_len_snapshot: int | None = None class EngineCoreOutputs( @@ -207,6 +215,21 @@ class EngineCoreOutputs( finished_requests={"req-1"}, ) +extended_outputs = EngineCoreOutputs( + outputs=[ + ExtendedEngineCoreOutput( + **msgspec.structs.asdict(outputs.outputs[0]), + multimodal_output={"audio": b"\x00\x01"}, + is_segment_finished=True, + new_prompt_len_snapshot=12, + ) + ], + finished_requests=outputs.finished_requests, +) +extended_outputs_bytes = msgspec.msgpack.encode(extended_outputs) +# The ordinary frontend's schema ignores even non-default extension values. +assert msgspec.msgpack.decode(extended_outputs_bytes, type=EngineCoreOutputs) == outputs + sampling_mask_wire = [ [ " Date: Wed, 16 Sep 2026 22:01:54 -0400 Subject: [PATCH 363/774] Use skip-redaction option for crcr-report (#56532) Signed-off-by: khluu Co-authored-by: khluu Co-authored-by: Codex --- .buildkite/scripts/crcr-report.sh | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/.buildkite/scripts/crcr-report.sh b/.buildkite/scripts/crcr-report.sh index e19146fddbed..3b12d451cdb8 100644 --- a/.buildkite/scripts/crcr-report.sh +++ b/.buildkite/scripts/crcr-report.sh @@ -40,8 +40,7 @@ fi TOKEN_SECRET_KEY="${CRCR_BUILDKITE_TOKEN_SECRET_KEY:-CRCR_BUILDKITE_API_TOKEN}" BK_TOKEN="${BUILDKITE_API_TOKEN:-}" if [[ -z "${BK_TOKEN}" ]]; then - # Not in the job environment, so read it from a Buildkite secret. The agent - # redacts values fetched this way from the log. + # Not in the job environment, so read it from a Buildkite secret. # # Report why a lookup failed. Swallowing stderr made a missing secret, a # denied policy and an unusable agent indistinguishable, all surfacing as the @@ -50,14 +49,9 @@ if [[ -z "${BK_TOKEN}" ]]; then echo "buildkite-agent is not on PATH; cannot read secret '${TOKEN_SECRET_KEY}'" else secret_err="$(mktemp)" - # Deliberately not passing --skip-redaction. On the deployed agent - # (v3.73.1) the flag cannot help: secret_get.go creates the Job API - # client unconditionally and only checks SkipRedaction afterwards, so - # under the docker plugin -- which does not expose the Job API socket to - # the container -- it fails before the flag is read. That ordering was - # only fixed in v3.107.0. Skipping redaction would also stop the token - # being registered with the log redactor, for no gain here. - if BK_TOKEN="$(buildkite-agent secret get "${TOKEN_SECRET_KEY}" 2>"${secret_err}")"; then + # Requires agent >= 3.107.0: the Docker plugin does not mount the Job API + # socket needed for redaction. Capture the value without logging it. + if BK_TOKEN="$(buildkite-agent secret get --skip-redaction "${TOKEN_SECRET_KEY}" 2>"${secret_err}")"; then if [[ -z "${BK_TOKEN}" ]]; then echo "secret '${TOKEN_SECRET_KEY}' resolved but is empty" fi @@ -77,7 +71,8 @@ if [[ -z "${BK_TOKEN}" ]]; then fi AUDIENCE="pytorch-cross-repo-ci-relay" -OIDC_TOKEN="$(buildkite-agent oidc request-token --audience "${AUDIENCE}" 2>/dev/null)" +# OIDC redaction also requires the unavailable Job API socket. +OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)" if [[ -z "${OIDC_TOKEN}" ]]; then echo "could not mint a Buildkite OIDC token -- skipping report" exit 0 From b1f9ebc0c367ccc2aa1e584647e55351b97a5ec3 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Thu, 17 Sep 2026 10:17:45 +0800 Subject: [PATCH 364/774] [Test][Core] Add hybrid model prefix cache hit-rate coverage (#57127) Signed-off-by: zjy0516 Co-authored-by: Codex --- .../test_hybrid_prefix_cache_hit_rate.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py diff --git a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py new file mode 100644 index 000000000000..c75591b69724 --- /dev/null +++ b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Replay fixed chat histories and check per-turn cache hits and metrics.""" + +import asyncio +import json +from uuid import uuid4 + +import httpx +import pytest +from openai.types.chat import ChatCompletionMessageParam +from prometheus_client.parser import text_string_to_metric_families + +from tests.utils import RemoteOpenAIServer + +MODEL = "Qwen/Qwen3.5-0.8B" +NUM_CONVERSATIONS = 8 +MAX_TOKENS = 32 +# Long first turns leave room for hybrid cache alignment and MTP recomputation. +MIN_PROMPT_TOKENS = 2048 +METRICS_TIMEOUT = 30 + + +def _user_turns(conversation: int) -> list[str]: + records = "\n".join( + f"Sensor {conversation}-{i}: temperature {20 + i % 10} C, " + f"humidity {40 + i % 20} percent." + for i in range(160 + conversation) + ) + return [ + f"Session {conversation}. Summarize these sensor readings briefly:\n{records}", + "Which readings should an operator investigate first? Explain briefly.", + "Suggest two checks the operator should perform next.", + ] + + +async def _prefix_cache_counters( + client: httpx.AsyncClient, +) -> tuple[float, float]: + response = await client.get("/metrics") + response.raise_for_status() + names = ("vllm:prefix_cache_queries_total", "vllm:prefix_cache_hits_total") + counters: dict[str, float] = {} + for family in text_string_to_metric_families(response.text): + for sample in family.samples: + if sample.name in names: + counters[sample.name] = counters.get(sample.name, 0) + sample.value + assert all(name in counters for name in names), ( + f"missing prefix-cache metrics: {set(names) - counters.keys()}" + ) + return counters[names[0]], counters[names[1]] + + +async def _wait_for_counters( + client: httpx.AsyncClient, expected: tuple[float, float] +) -> None: + """Allow engine stats to reach the frontend after completions are returned.""" + deadline = asyncio.get_running_loop().time() + METRICS_TIMEOUT + while True: + actual = await _prefix_cache_counters(client) + if actual == expected: + return + assert all(a <= e for a, e in zip(actual, expected)), ( + f"prefix-cache counters exceeded usage: {actual=}, {expected=}" + ) + assert asyncio.get_running_loop().time() < deadline, ( + f"prefix-cache counters did not match usage: {actual=}, {expected=}" + ) + await asyncio.sleep(0.1) + + +@pytest.fixture(scope="module") +def server(request): + args = ["--enable-prompt-tokens-details"] + if request.param: + args += [ + "--speculative-config", + json.dumps({"method": "mtp", "num_speculative_tokens": 1}), + ] + with RemoteOpenAIServer(MODEL, args) as remote_server: + yield remote_server + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server, min_hit_rate", + # MTP re-prefills an extra block; both floors reject one more missed block. + [pytest.param(False, 0.90, id="base"), pytest.param(True, 0.75, id="mtp")], + indirect=["server"], +) +async def test_prefix_cache_hit_rate( + server: RemoteOpenAIServer, min_hit_rate: float +) -> None: + # Isolate sessions so another conversation cannot hide a cache miss. + salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)] + histories: list[list[ChatCompletionMessageParam]] = [ + [] for _ in range(NUM_CONVERSATIONS) + ] + prompt_lengths: list[list[int]] = [[] for _ in range(NUM_CONVERSATIONS)] + + async with ( + server.get_async_client() as client, + httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics, + ): + + async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]: + history = histories[session] + queries = hits = 0 + for turn, user_text in enumerate(_user_turns(session)): + if not replay: + history.append({"role": "user", "content": user_text}) + # Reuse the original replies, even if replay generates different ones. + response = await client.chat.completions.create( + model=MODEL, + messages=history[: 2 * turn + 1], + temperature=0, + max_tokens=MAX_TOKENS, + extra_body={"cache_salt": salts[session]}, + ) + usage = response.usage + assert usage is not None + assert usage.prompt_tokens_details is not None + prompt = usage.prompt_tokens + cached = usage.prompt_tokens_details.cached_tokens + context = f"{session=}, {turn=}, {replay=}: {cached=}, {prompt=}" + assert cached is not None, context + assert 0 <= cached < prompt, context + assert prompt >= MIN_PROMPT_TOKENS, context + if not replay and turn == 0: + assert cached == 0, context + else: + assert cached / prompt >= min_hit_rate, ( + f"{context}, {min_hit_rate=}" + ) + if replay: + assert prompt == prompt_lengths[session][turn], context + else: + prompt_lengths[session].append(prompt) + history.append( + { + "role": "assistant", + "content": response.choices[0].message.content or "", + } + ) + queries += prompt + hits += cached + return queries, hits + + queries, hits = await _prefix_cache_counters(metrics) + for replay in (False, True): + counts = await asyncio.gather( + *(run_conversation(i, replay=replay) for i in range(NUM_CONVERSATIONS)) + ) + queries += sum(prompt for prompt, _ in counts) + hits += sum(cached for _, cached in counts) + await _wait_for_counters(metrics, (queries, hits)) From f730a93d2bf2cad3002578238d5a1035ee1f4177 Mon Sep 17 00:00:00 2001 From: "Chuan (Richard) Li" Date: Wed, 16 Sep 2026 19:22:12 -0700 Subject: [PATCH 365/774] [Scheduler] Add --max-num-active-seqs to cap RUNNING admission (#56758) Signed-off-by: Li Co-authored-by: Cursor --- tests/v1/core/test_scheduler.py | 16 ++++++++++++++++ tests/v1/core/utils.py | 2 ++ vllm/config/scheduler.py | 19 +++++++++++++++++++ vllm/engine/arg_utils.py | 6 ++++++ vllm/v1/core/sched/scheduler.py | 9 ++++++++- 5 files changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3239ec25a00a..3add7beda0e7 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -193,6 +193,22 @@ def test_get_num_unfinished_requests(): assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1 +def test_max_num_active_seqs_caps_admission(): + """Admission is capped by max_num_active_seqs, not max_num_seqs.""" + scheduler = create_scheduler(max_num_seqs=16, max_num_active_seqs=3) + requests = create_requests(num_requests=10) + for request in requests: + scheduler.add_request(request) + + output = scheduler.schedule() + + assert len(output.scheduled_new_reqs) == 3 + assert len(scheduler.running) == 3 + assert len(scheduler.waiting) == 7 + # Execution capacity is untouched. + assert scheduler.max_num_running_reqs == 16 + + def _bind_hisparse_connector(scheduler): connector = object.__new__(HiSparseConnector) connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 28d1f71ac492..f69b10ced346 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -52,6 +52,7 @@ def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int def create_scheduler( model: str = "facebook/opt-125m", max_num_seqs: int = 16, + max_num_active_seqs: int | None = None, max_num_batched_tokens: int = 8192, enable_chunked_prefill: bool = True, enable_prefix_caching: bool = False, @@ -107,6 +108,7 @@ def create_scheduler( max_model_len = max_num_batched_tokens scheduler_config = SchedulerConfig( max_num_seqs=max_num_seqs, + max_num_active_seqs=max_num_active_seqs, max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len, long_prefill_token_threshold=long_prefill_token_threshold, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index d3264b2d8071..321c76fe3058 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -67,6 +67,16 @@ class SchedulerConfig: In real usage, this should be set in `EngineArgs.create_engine_config`. """ + max_num_active_seqs: int | None = Field(default=None, ge=1) + """Maximum number of requests the scheduler admits into RUNNING. + + ``max_num_seqs`` sizes the model runner (per-request buffers and CUDA + graph capture) and is also the default admission limit. Setting this + lowers only the number of requests that may occupy RUNNING, so decode + batches stay smaller without shrinking runner or graph capacity. Must + be ``<= max_num_seqs``. ``None`` (default) keeps current behavior. + """ + long_prefill_token_threshold: int = Field(default=0, ge=0) """For chunked prefill, a request is considered long if the prompt is longer than this number of tokens. 0 disables the cap (default).""" @@ -310,6 +320,15 @@ def verify_max_model_len(self, max_model_len: int) -> Self: f"({self.max_num_seqs})." ) + if ( + self.max_num_active_seqs is not None + and self.max_num_active_seqs > self.max_num_seqs + ): + raise ValueError( + f"max_num_active_seqs ({self.max_num_active_seqs}) cannot be " + f"greater than max_num_seqs ({self.max_num_seqs})." + ) + if self.max_num_batched_tokens > self.max_num_seqs * max_model_len: logger.warning( "max_num_batched_tokens (%d) exceeds max_num_seqs " diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ccb30a2597ff..d637391cd3a4 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -556,6 +556,7 @@ class EngineArgs: max_num_scheduled_tokens: int | None = None long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold max_num_seqs: int | None = None + max_num_active_seqs: int | None = SchedulerConfig.max_num_active_seqs max_num_queued_reqs: int | None = None max_num_queued_tokens: int | None = None max_logprobs: int = ModelConfig.max_logprobs @@ -1589,6 +1590,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "default": None, }, ) + scheduler_group.add_argument( + "--max-num-active-seqs", + **scheduler_kwargs["max_num_active_seqs"], + ) scheduler_group.add_argument( "--max-num-queued-reqs", **scheduler_kwargs["max_num_queued_reqs"] ) @@ -2421,6 +2426,7 @@ def create_engine_config( max_num_batched_tokens=self.max_num_batched_tokens, max_num_scheduled_tokens=self.max_num_scheduled_tokens, max_num_seqs=self.max_num_seqs, + max_num_active_seqs=self.max_num_active_seqs, max_num_queued_reqs=self.max_num_queued_reqs, max_num_queued_tokens=self.max_num_queued_tokens, max_model_len=model_config.max_model_len, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 19ea411aa19b..22c231cc437e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -123,6 +123,13 @@ def __init__( # Scheduling constraints. self.max_num_running_reqs = self.scheduler_config.max_num_seqs + # Admission into RUNNING. Separate from max_num_running_reqs, which + # is the model-runner slot count. + self.max_num_active_reqs = ( + self.scheduler_config.max_num_active_seqs + if self.scheduler_config.max_num_active_seqs is not None + else self.max_num_running_reqs + ) self.max_num_scheduled_tokens = ( self.scheduler_config.max_num_scheduled_tokens if self.scheduler_config.max_num_scheduled_tokens is not None @@ -850,7 +857,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not # in `running` but still hold a model-runner request slot. num_running = len(self.running) + self.num_waiting_for_streaming_input - if num_running >= self.max_num_running_reqs: + if num_running >= self.max_num_active_reqs: break request_queue = self._select_waiting_queue_for_scheduling() From 2307a1d16cdd0077cd882bb9fc41b453f6f1999e Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Wed, 16 Sep 2026 20:30:59 -0600 Subject: [PATCH 366/774] [Core] Don't issue a blocking collective RPC during engine handshake (#57226) Signed-off-by: okorzh Co-authored-by: Andreas Karatzas --- tests/v1/engine/test_engine_core_client.py | 38 ++++++++++++++++++++++ vllm/v1/engine/core.py | 2 ++ 2 files changed, 40 insertions(+) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index a23ae46b8f5d..a153b23b10bb 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -1452,3 +1452,41 @@ def mock_startup_handshake( ) engine_core_proc.shutdown() + + +def _ready_response_with_weight_transfer(weight_transfer_config): + """Build a ready response against a mock executor.""" + from vllm.v1.engine.core import EngineCoreProc + + executor = MagicMock() + executor.supports_draft_weight_updates.return_value = True + + proc = MagicMock() + proc.model_executor = executor + proc.vllm_config = MagicMock() + proc.vllm_config.weight_transfer_config = weight_transfer_config + proc.vllm_config.lora_config = None + proc.vllm_config.model_config.enable_sleep_mode = False + proc.scheduler = MagicMock() + + response = EngineCoreProc._make_ready_response(proc) + return response, executor + + +def test_ready_response_skips_draft_rpc_without_weight_transfer(): + """No weight-transfer config: the executor must not be consulted.""" + response, executor = _ready_response_with_weight_transfer(None) + + executor.supports_draft_weight_updates.assert_not_called() + assert response.supports_draft_weight_updates is False + + +def test_ready_response_queries_executor_with_weight_transfer(): + """Weight transfer configured: the executor is still consulted.""" + weight_transfer_config = MagicMock() + weight_transfer_config.backend = "nixl" + + response, executor = _ready_response_with_weight_transfer(weight_transfer_config) + + executor.supports_draft_weight_updates.assert_called_once() + assert response.supports_draft_weight_updates is True diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9c84827a63d9..7da7f2e0ece1 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1691,6 +1691,8 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode, supports_draft_weight_updates=( self.model_executor.supports_draft_weight_updates() + if self.vllm_config.weight_transfer_config is not None + else False ), ) From fa3622a4a1cd16294a7843895a3246bf29c452c4 Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:36:23 +0300 Subject: [PATCH 367/774] [EC Connector] Add Metrics Collection (#54960) Signed-off-by: omerpaz95 --- tests/v1/core/test_scheduler.py | 82 ++++++++ .../unit/test_ec_connector_metrics.py | 85 +++++++++ .../unit/test_ec_output_aggregator.py | 22 ++- .../unit/test_worker_ec_connector.py | 12 ++ .../ec_transfer/ec_connector/base.py | 43 +++++ .../ec_transfer/ec_connector/metrics.py | 175 ++++++++++++++++++ .../ec_transfer/ec_connector/utils.py | 9 + vllm/v1/core/sched/scheduler.py | 24 +++ vllm/v1/metrics/loggers.py | 18 ++ vllm/v1/metrics/stats.py | 1 + vllm/v1/outputs.py | 3 + vllm/v1/worker/gpu/ec_connector.py | 1 + 12 files changed, 472 insertions(+), 3 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_ec_connector_metrics.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/metrics.py diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3add7beda0e7..ef818f27c772 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -19,6 +19,7 @@ SpeculativeConfig, VllmConfig, ) +from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( HiSparseConnector, HiSparseConnectorScheduler, @@ -5296,6 +5297,87 @@ def update_connector_output( assert final_stats == expected_data +def test_scheduler_ec_connector_stats(): + """Test worker-side, scheduler-side, and combined EC connector stats.""" + + class GenericECConnectorStats(ECConnectorStats): + def reset(self): + self.data = {} + + def aggregate(self, other: ECConnectorStats) -> ECConnectorStats: + self.data.update(other.data) + return self + + def reduce(self) -> dict[str, int | float]: + return {} + + def is_empty(self) -> bool: + return not self.data + + test_cases = ( + ({"worker": 1}, None, {"worker": 1}), + (None, {"scheduler": 2}, {"scheduler": 2}), + ({"worker": 1}, {"scheduler": 2}, {"worker": 1, "scheduler": 2}), + ) + + for worker_data, scheduler_data, expected_data in test_cases: + scheduler = create_scheduler() + worker_stats = ( + GenericECConnectorStats(data=worker_data) if worker_data else None + ) + scheduler_stats = ( + GenericECConnectorStats(data=scheduler_data) if scheduler_data else None + ) + scheduler.ec_connector = Mock() + scheduler.ec_connector.take_unavailable_requests.return_value = set() + scheduler.ec_connector.get_ec_connector_stats.return_value = ( + scheduler_stats if worker_stats is None else None + ) + + def update_connector_output( + ec_connector_output: ECConnectorOutput, + scheduler=scheduler, + scheduler_stats=scheduler_stats, + ): + scheduler.ec_connector.get_ec_connector_stats.return_value = scheduler_stats + + scheduler.ec_connector.update_connector_output.side_effect = ( + update_connector_output + ) + + model_output = ModelRunnerOutput( + req_ids=["req_0"], + req_id_to_index={"req_0": 0}, + sampled_token_ids=[[123]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[None], + ec_connector_output=ECConnectorOutput(ec_connector_stats=worker_stats) + if worker_stats + else None, + ) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=None, + num_scheduled_tokens={"req_0": 1}, + total_num_scheduled_tokens=1, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + engine_core_outputs = scheduler.update_from_output( + scheduler_output, model_output + ) + + final_stats = next( + iter(engine_core_outputs.values()) + ).scheduler_stats.ec_connector_stats + assert final_stats == expected_data + + def test_ec_connector_update_connector_output_called(): """Test that worker-side EC connector output is forwarded to the EC connector's update_connector_output hook.""" diff --git a/tests/v1/ec_connector/unit/test_ec_connector_metrics.py b/tests/v1/ec_connector/unit/test_ec_connector_metrics.py new file mode 100644 index 000000000000..239f0267f060 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_ec_connector_metrics.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the connector-agnostic EC connector metrics plumbing. + +Covers only the base infra: no concrete connector type implements these +hooks yet, so every assertion here is about the documented no-op defaults +and the container contracts, mirroring the equivalent KV connector tests. +""" + +from unittest.mock import Mock + +import pytest + +from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole +from vllm.distributed.ec_transfer.ec_connector.example_connector import ( + ECExampleConnector, +) +from vllm.distributed.ec_transfer.ec_connector.metrics import ( + ECConnectorLogging, + ECConnectorProm, + ECConnectorStats, +) +from vllm.v1.outputs import ECConnectorOutput + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture +def example_connector(tmp_path): + config = Mock(spec=VllmConfig) + config.ec_transfer_config = Mock() + config.ec_transfer_config.get_from_extra_config = Mock(return_value=str(tmp_path)) + config.ec_transfer_config.is_ec_producer = True + config.ec_transfer_config.is_ec_consumer = False + return ECExampleConnector(vllm_config=config, role=ECConnectorRole.SCHEDULER) + + +def test_base_stats_methods_raise_not_implemented(): + """The base container has no aggregation logic of its own.""" + stats = ECConnectorStats() + + with pytest.raises(NotImplementedError): + stats.reset() + with pytest.raises(NotImplementedError): + stats.aggregate(stats) + with pytest.raises(NotImplementedError): + stats.reduce() + with pytest.raises(NotImplementedError): + stats.is_empty() + + +def test_connector_hooks_default_to_none(example_connector): + """No connector type overrides these yet, so every hook is a no-op.""" + assert example_connector.get_ec_connector_stats() is None + assert ECExampleConnector.build_ec_connector_stats() is None + assert ( + ECExampleConnector.build_prom_metrics( + vllm_config=Mock(spec=VllmConfig), + metric_types={}, + labelnames=[], + per_engine_labelvalues={}, + ) + is None + ) + + +def test_ec_connector_output_is_empty_reflects_stats_field(): + output = ECConnectorOutput() + assert output.is_empty() + + output.ec_connector_stats = ECConnectorStats(data={"x": 1}) + assert not output.is_empty() + + +def test_logging_and_prom_are_no_ops_without_ec_transfer_config(): + """A normal run with no EC connector configured must stay silent.""" + logging = ECConnectorLogging(ec_transfer_config=None) + logging.log() # Should not raise, nothing to log. + + vllm_config = Mock(spec=VllmConfig) + vllm_config.ec_transfer_config = None + prom = ECConnectorProm(vllm_config, labelnames=[], per_engine_labelvalues={}) + assert prom.prom_metrics is None + prom.observe({}) # Should not raise. diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py index 4d5b7dae1037..df98ec364f5c 100644 --- a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -5,6 +5,7 @@ import pytest from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata +from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.v1.outputs import ( @@ -29,6 +30,16 @@ def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": return FakeWorkerMeta(self.saves + other.saves) +class FakeConnectorStats(ECConnectorStats): + """Records merge order, mirroring FakeWorkerMeta above.""" + + def __init__(self, saves: list[str]): + self.saves = saves + + def aggregate(self, other: "FakeConnectorStats") -> "FakeConnectorStats": + return FakeConnectorStats(self.saves + other.saves) + + def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: return ModelRunnerOutput( req_ids=[], req_id_to_index={}, ec_connector_output=ec_output @@ -38,19 +49,23 @@ def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: def test_aggregate_folds_every_rank_onto_output_rank(): """EC work done on any rank reaches the scheduler via output_rank's output. - The middle rank reports no worker metadata: it must neither seed nor clobber - the accumulator. + The middle rank reports no worker metadata or stats: it must neither seed + nor clobber either accumulator. """ outputs = [ _worker_output( ECConnectorOutput( finished_sending={"mm0"}, + ec_connector_stats=FakeConnectorStats(["mm0"]), ec_connector_worker_meta=FakeWorkerMeta(["mm0"]), ) ), _worker_output(ECConnectorOutput(finished_recving={"mm1"})), _worker_output( - ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"])) + ECConnectorOutput( + ec_connector_stats=FakeConnectorStats(["mm2"]), + ec_connector_worker_meta=FakeWorkerMeta(["mm2"]), + ) ), ] @@ -59,6 +74,7 @@ def test_aggregate_folds_every_rank_onto_output_rank(): assert result is outputs[2] assert result.ec_connector_output.finished_sending == {"mm0"} assert result.ec_connector_output.finished_recving == {"mm1"} + assert result.ec_connector_output.ec_connector_stats.saves == ["mm0", "mm2"] assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"] diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py index 01dad9d89ebd..911810b53acf 100644 --- a/tests/v1/ec_connector/unit/test_worker_ec_connector.py +++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py @@ -20,6 +20,7 @@ pytestmark = pytest.mark.cpu_test WORKER_META = object() +CONNECTOR_STATS = object() def _scheduler_output() -> SimpleNamespace: @@ -38,6 +39,7 @@ def _connector( fake.is_consumer = is_consumer fake.get_finished.return_value = (None, None) fake.build_connector_worker_meta.return_value = WORKER_META + fake.get_ec_connector_stats.return_value = CONNECTOR_STATS with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake @@ -72,6 +74,16 @@ def test_worker_meta_is_reported_on_context_exit(): assert fake.clear_connector_metadata.called +def test_connector_stats_are_reported_on_context_exit(): + """Reported in the finally block, alongside the worker meta.""" + connector, fake = _connector() + + with connector.maybe_get_output(_scheduler_output()) as output: + assert output.ec_connector_stats is None + + assert output.ec_connector_stats is CONNECTOR_STATS + + def test_v1_producer_receives_existing_encoder_cache(): encoder_cache = {"mm_cached": None} fake = MagicMock(spec=ECConnectorBase) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index ec3166ce3381..716f443af244 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -35,6 +35,12 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.distributed.ec_transfer.ec_connector.metrics import ( + ECConnectorPromMetrics, + ECConnectorStats, + PromMetric, + PromMetricT, + ) from vllm.v1.request import Request logger = init_logger(__name__) @@ -219,6 +225,17 @@ def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: """ return None + def get_ec_connector_stats(self) -> "ECConnectorStats | None": + """ + Get the EC connector stats collected during the last interval. + + Callable on either a worker-side or scheduler-side connector + instance: the worker-side instance reports stats gathered during + cache save/load, while the scheduler-side instance reports stats + gathered during scheduling decisions. + """ + return None + # ============================== # Scheduler-side methods # ============================== @@ -336,3 +353,29 @@ def has_pending_push_work(self) -> bool: leave this as False. """ return False + + @classmethod + def build_ec_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> "ECConnectorStats | None": + """ + ECConnectorStats resolution method. This method allows dynamically + registered connectors to return their own ECConnectorStats object, + which can implement custom aggregation logic on the data dict. + """ + return None + + @classmethod + def build_prom_metrics( + cls, + vllm_config: "VllmConfig", + metric_types: dict[type["PromMetric"], type["PromMetricT"]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> "ECConnectorPromMetrics | None": + """ + Create an ECConnectorPromMetrics subclass which should register + per-connector Prometheus metrics and implement observe() to + expose connector transfer stats via Prometheus. + """ + return None diff --git a/vllm/distributed/ec_transfer/ec_connector/metrics.py b/vllm/distributed/ec_transfer/ec_connector/metrics.py new file mode 100644 index 000000000000..4a329096cee6 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/metrics.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field +from typing import Any, TypeAlias, TypeVar + +from prometheus_client import Counter, Gauge, Histogram + +from vllm.config import ECTransferConfig, VllmConfig +from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory +from vllm.logger import init_logger + +PromMetric: TypeAlias = Gauge | Counter | Histogram +PromMetricT = TypeVar("PromMetricT", bound=PromMetric) + +logger = init_logger(__name__) + + +@dataclass +class ECConnectorStats: + """ + Base class for EC Connector Stats, a container for transfer performance + metrics or otherwise important telemetry from the connector. + All sub-classes need to be serializable as stats are sent from worker to + logger process. + """ + + data: dict[str, Any] = field(default_factory=dict) + + def reset(self): + """Reset the stats, clear the state.""" + raise NotImplementedError + + def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats": + """ + Aggregate stats with another `ECConnectorStats` object. + """ + raise NotImplementedError + + def reduce(self) -> dict[str, int | float]: + """ + Reduce the observations collected during a time interval to one or + more representative values (eg avg/median/sum of the series). + This is meant to be called by the logger to produce a summary of the + stats for the last time interval. + """ + raise NotImplementedError + + def is_empty(self) -> bool: + """Return True if the stats are empty.""" + raise NotImplementedError + + +class ECConnectorLogging: + def __init__(self, ec_transfer_config: ECTransferConfig | None): + # Instantiate the connector's stats class. + if ec_transfer_config and ec_transfer_config.ec_connector: + self.connector_cls = ECConnectorFactory.get_connector_class( + ec_transfer_config + ) + self.reset() + + def reset(self): + self.transfer_stats_accumulator: ECConnectorStats | None = None + + def observe(self, transfer_stats_data: dict[str, Any]): + # Should not be called when an ECConnector is not configured. + assert self.connector_cls is not None + # Called periodically when connector syncs with the scheduler. + # Note that this is not the same as the logging interval. + # We expect transfer_stats_data to be aggregated across all workers and + # consist of observations from a single connector. + transfer_stats = self.connector_cls.build_ec_connector_stats( + transfer_stats_data + ) + if transfer_stats is None: + logger.warning_once( + "The connector %s is collecting stats but " + "does not implement the " + "`build_ec_connector_stats` method. " + "Stats will not be logged.", + self.connector_cls, + ) + return + + if self.transfer_stats_accumulator is None: + self.transfer_stats_accumulator = transfer_stats + else: + # Accumulate last interval stats. + self.transfer_stats_accumulator = self.transfer_stats_accumulator.aggregate( + transfer_stats + ) + + def log(self, log_fn=logger.info): + """Log transfer metrics periodically, similar to throughput logging""" + if ( + self.transfer_stats_accumulator + and not self.transfer_stats_accumulator.is_empty() + ): + # Produce a single cumulative stats object for the last time + # interval from the recorded observations. + xfer_metrics = self.transfer_stats_accumulator.reduce() + xfer_metrics_str = ", ".join(f"{k}={v}" for k, v in xfer_metrics.items()) + log_fn("EC Transfer metrics: %s", xfer_metrics_str) + + # Reset metrics for next interval + self.reset() + + +class ECConnectorPromMetrics: + """ + A base class for per-connector Prometheus metric registration + and recording. + """ + + def __init__( + self, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ): + self._ec_transfer_config = vllm_config.ec_transfer_config + self._gauge_cls = metric_types[Gauge] + self._counter_cls = metric_types[Counter] + self._histogram_cls = metric_types[Histogram] + self._labelnames = labelnames + self.per_engine_labelvalues = per_engine_labelvalues + + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): + """ + Record the supplied transfer statistics to Prometheus metrics. These + statistics are engine-specific, and should be recorded to a metric + with the appropriate 'engine' label. These metric instances can be + created using the create_metric_per_engine() helper method. + """ + raise NotImplementedError + + +class ECConnectorProm: + """ + Support for registering per-connector Prometheus metrics, and + recording transfer statistics to those metrics. Uses + ECConnectorBase.build_prom_metrics(). + """ + + _gauge_cls = Gauge + _counter_cls = Counter + _histogram_cls = Histogram + + def __init__( + self, + vllm_config: VllmConfig, + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ): + self.prom_metrics: ECConnectorPromMetrics | None = None + ec_transfer_config = vllm_config.ec_transfer_config + if ec_transfer_config and ec_transfer_config.ec_connector: + connector_cls = ECConnectorFactory.get_connector_class(ec_transfer_config) + metric_types = { + Gauge: self._gauge_cls, + Counter: self._counter_cls, + Histogram: self._histogram_cls, + } + self.prom_metrics = connector_cls.build_prom_metrics( + vllm_config, + metric_types, + labelnames, + per_engine_labelvalues, + ) + + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): + if self.prom_metrics is None: + return + self.prom_metrics.observe(transfer_stats_data, engine_idx) diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index 46963c2b4113..f54d6877c1b4 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -103,6 +103,7 @@ def aggregate( finished_sending = set[str]() finished_recving = set[str]() worker_meta = None + connector_stats = None for model_runner_output in outputs: assert model_runner_output is not None ec_output = model_runner_output.ec_connector_output @@ -117,9 +118,17 @@ def aggregate( meta if worker_meta is None else worker_meta.aggregate(meta) ) + if stats := ec_output.ec_connector_stats: + connector_stats = ( + stats + if connector_stats is None + else connector_stats.aggregate(stats) + ) + aggregated = ECConnectorOutput( finished_sending=finished_sending or None, finished_recving=finished_recving or None, + ec_connector_stats=connector_stats, ec_connector_worker_meta=worker_meta, ) if aggregated.is_empty(): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 22c231cc437e..ed80daa8c279 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -15,6 +15,7 @@ ECConnectorRole, ) from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory +from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory from vllm.distributed.kv_transfer.kv_connector.v1 import ( @@ -2243,6 +2244,23 @@ def update_from_output( else scheduler_kv_connector_stats ) + # Worker-side EC connector stats from the model runner output. + ec_connector_stats: ECConnectorStats | None = ( + ec_connector_output.ec_connector_stats if ec_connector_output else None + ) + if self.ec_connector: + # Scheduler-side EC connector stats collected after connector update. + scheduler_ec_connector_stats = self.ec_connector.get_ec_connector_stats() + if ( + scheduler_ec_connector_stats is not None + and not scheduler_ec_connector_stats.is_empty() + ): + ec_connector_stats = ( + ec_connector_stats.aggregate(scheduler_ec_connector_stats) + if ec_connector_stats is not None + else scheduler_ec_connector_stats + ) + # collect KV cache events from KV cache manager events = self.kv_cache_manager.take_events() @@ -2287,6 +2305,7 @@ def update_from_output( kv_connector_stats, cudagraph_stats, perf_stats, + ec_connector_stats=ec_connector_stats, ) ) is not None: # Return stats to only one of the front-ends. @@ -2772,6 +2791,7 @@ def make_stats( kv_connector_stats: KVConnectorStats | None = None, cudagraph_stats: CUDAGraphStat | None = None, perf_stats: PerfStats | None = None, + ec_connector_stats: ECConnectorStats | None = None, ) -> SchedulerStats | None: if not self.log_stats: return None @@ -2790,6 +2810,9 @@ def make_stats( connector_stats_payload = ( kv_connector_stats.to_dict() if kv_connector_stats else None ) + ec_connector_stats_payload = ( + ec_connector_stats.data if ec_connector_stats else None + ) return SchedulerStats( num_running_reqs=len(self.running), num_waiting_reqs=len(self.waiting), @@ -2802,6 +2825,7 @@ def make_stats( kv_connector_stats=connector_stats_payload, cudagraph_stats=cudagraph_stats, perf_stats=perf_stats, + ec_connector_stats=ec_connector_stats_payload, ) def make_spec_decoding_stats( diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 50e606d3b3f5..fee6679fe760 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -11,6 +11,10 @@ import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphLogging from vllm.config import SupportsMetricsInfo, VllmConfig +from vllm.distributed.ec_transfer.ec_connector.metrics import ( + ECConnectorLogging, + ECConnectorProm, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( KVConnectorLogging, KVConnectorProm, @@ -118,6 +122,8 @@ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): self.spec_decoding_logging = SpecDecodingLogging(is_diffusion=is_diffusion) kv_transfer_config = self.vllm_config.kv_transfer_config self.kv_connector_logging = KVConnectorLogging(kv_transfer_config) + ec_transfer_config = self.vllm_config.ec_transfer_config + self.ec_connector_logging = ECConnectorLogging(ec_transfer_config) self.cudagraph_logging = None if self.vllm_config.observability_config.cudagraph_metrics: self.cudagraph_logging = CUDAGraphLogging( @@ -224,6 +230,8 @@ def record( self.spec_decoding_logging.observe(scheduler_stats.spec_decoding_stats) if kv_connector_stats := scheduler_stats.kv_connector_stats: self.kv_connector_logging.observe(kv_connector_stats) + if ec_connector_stats := scheduler_stats.ec_connector_stats: + self.ec_connector_logging.observe(ec_connector_stats) if ( self.cudagraph_logging is not None and scheduler_stats.cudagraph_stats is not None @@ -319,6 +327,7 @@ def log(self): self.spec_decoding_logging.log(log_fn=log_fn) self.kv_connector_logging.log(log_fn=log_fn) + self.ec_connector_logging.log(log_fn=log_fn) if self.cudagraph_logging is not None: self.cudagraph_logging.log(log_fn=log_fn) if self._enable_perf_stats(): @@ -451,6 +460,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase): _histogram_cls = Histogram _spec_decoding_cls = SpecDecodingProm _kv_connector_cls = KVConnectorProm + _ec_connector_cls = ECConnectorProm _perf_metrics_cls = PerfMetricsProm def __init__( @@ -488,6 +498,9 @@ def __init__( self.kv_connector_prom = self._kv_connector_cls( vllm_config, labelnames, per_engine_labelvalues ) + self.ec_connector_prom = self._ec_connector_cls( + vllm_config, labelnames, per_engine_labelvalues + ) self.perf_metrics_prom = self._perf_metrics_cls( vllm_config, labelnames, per_engine_labelvalues ) @@ -1064,6 +1077,11 @@ def record( scheduler_stats.kv_connector_stats, engine_idx ) + if scheduler_stats.ec_connector_stats is not None: + self.ec_connector_prom.observe( + scheduler_stats.ec_connector_stats, engine_idx + ) + if scheduler_stats.perf_stats is not None: self.perf_metrics_prom.observe(scheduler_stats.perf_stats, engine_idx) diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index fe17208ce40a..66088c7954aa 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -205,6 +205,7 @@ class SchedulerStats: spec_decoding_stats: SpecDecodingStats | None = None kv_connector_stats: dict[str, Any] | None = None + ec_connector_stats: dict[str, Any] | None = None waiting_lora_adapters: dict[str, int] = field(default_factory=dict) running_lora_adapters: dict[str, int] = field(default_factory=dict) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 35780dce9c6c..c4535f47f227 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata + from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.kv_events import KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorWorkerMetadata, @@ -307,12 +308,14 @@ class ECConnectorOutput: # [mm_hash] finished_sending: set[str] | None = None finished_recving: set[str] | None = None + ec_connector_stats: "ECConnectorStats | None" = None ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None def is_empty(self): return ( not self.finished_sending and not self.finished_recving + and not self.ec_connector_stats and not self.ec_connector_worker_meta ) diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index 1285d9fd14db..c81fd7c11e76 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -80,6 +80,7 @@ def maybe_get_output( output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) ) + output.ec_connector_stats = ec_connector.get_ec_connector_stats() output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() From 5efbf0886a7110ae0801b9540e3072e89a23eb36 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 16 Sep 2026 19:40:21 -0700 Subject: [PATCH 368/774] [CI] Support `FORCE_COLOR` env var; enable colors in CPU CI output (#56913) Signed-off-by: Nick Hill --- .../scripts/hardware_ci/run-cpu-test-arm.sh | 2 +- .../scripts/hardware_ci/run-cpu-test.sh | 2 +- tests/test_logger.py | 23 +++++++++++++++++++ vllm/envs.py | 4 ++++ vllm/logger.py | 2 +- 5 files changed, 30 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 77e0f3e07bad..689ca92d6a05 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -30,7 +30,7 @@ remove_docker_container docker pull "$IMAGE" # Run the image -docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE" +docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE" print_packages() { docker exec "$CONTAINER_NAME" bash -c " diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 43555a19ab58..eecf70ea2bde 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -99,7 +99,7 @@ rm -f "$build_log" OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with' run_test() { local hf_offline=$1 - docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ + docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g "$IMAGE_NAME" \ timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" } diff --git a/tests/test_logger.py b/tests/test_logger.py index 2ff100151b29..9dcc4a1ea825 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import enum +import io import json import logging import os @@ -19,6 +20,7 @@ _DATE_FORMAT, _FORMAT, _configure_vllm_root_logger, + _use_color, enable_trace_function_call, init_logger, ) @@ -72,6 +74,27 @@ def test_default_vllm_root_logger_configuration(monkeypatch): assert formatter.datefmt == _DATE_FORMAT +def test_use_color_force_color(monkeypatch): + """FORCE_COLOR forces colored logs without a TTY, while NO_COLOR and an + explicit VLLM_LOGGING_COLOR=0 take precedence over it.""" + monkeypatch.setattr(sys, "stdout", io.StringIO()) + monkeypatch.setattr(sys, "stderr", io.StringIO()) + for var in ("NO_COLOR", "FORCE_COLOR", "VLLM_LOGGING_COLOR"): + monkeypatch.delenv(var, raising=False) + + assert not _use_color() + + monkeypatch.setenv("FORCE_COLOR", "1") + assert _use_color() + + monkeypatch.setenv("VLLM_LOGGING_COLOR", "0") + assert not _use_color() + monkeypatch.delenv("VLLM_LOGGING_COLOR") + + monkeypatch.setenv("NO_COLOR", "1") + assert not _use_color() + + def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypatch): """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default diff --git a/vllm/envs.py b/vllm/envs.py index 0a3f6d3c6751..c512df457adc 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -44,6 +44,7 @@ VLLM_LOGGING_CONFIG_PATH: str | None = None VLLM_LOGGING_COLOR: str = "auto" NO_COLOR: bool = False + FORCE_COLOR: bool = False VLLM_LOG_STATS_INTERVAL: float = 10.0 VLLM_TRACE_FUNCTION: int = 0 VLLM_USE_FLASHINFER_SAMPLER: bool = True @@ -840,6 +841,8 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_LOGGING_COLOR": lambda: os.getenv("VLLM_LOGGING_COLOR", "auto"), # Standard unix flag for disabling ANSI color codes "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0", + # De-facto standard flag for forcing ANSI color codes (e.g. non-tty case) + "FORCE_COLOR": lambda: os.getenv("FORCE_COLOR", "0") != "0", # If set, vllm will log stats at this interval in seconds # If not set, vllm will log stats every 10 seconds. "VLLM_LOG_STATS_INTERVAL": lambda: ( @@ -2335,6 +2338,7 @@ def compile_factors() -> dict[str, object]: "CUDA_VISIBLE_DEVICES", "NO_COLOR", "VLLM_SKIP_VERSION_SUFFIX", + "FORCE_COLOR", } from vllm.config.utils import normalize_value diff --git a/vllm/logger.py b/vllm/logger.py index cdce7748f2cf..05ee5936d45a 100644 --- a/vllm/logger.py +++ b/vllm/logger.py @@ -29,7 +29,7 @@ def _use_color() -> bool: if envs.NO_COLOR or envs.VLLM_LOGGING_COLOR == "0": return False - if envs.VLLM_LOGGING_COLOR == "1": + if envs.VLLM_LOGGING_COLOR == "1" or envs.FORCE_COLOR: return True if envs.VLLM_LOGGING_STREAM == "ext://sys.stdout": # stdout return hasattr(sys.stdout, "isatty") and sys.stdout.isatty() From 9854b580dfe518171eed3a6ad76cbd402ca2bdfa Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 17 Sep 2026 06:24:06 +0300 Subject: [PATCH 369/774] [Bugfix][KV Offload] Compact canonical MLA rows (#56799) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../unit/offloading_connector/test_config.py | 40 +++++++++++++++++++ .../unit/offloading_connector/test_worker.py | 8 +++- .../kv_connector/v1/offloading/config.py | 14 +++++++ .../kv_connector/v1/offloading/worker.py | 4 +- 4 files changed, 63 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index f597c8fa27c8..acf55972dfa4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,6 +17,7 @@ SchedulerOffloadConfig, ) from vllm.platforms import current_platform +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, @@ -32,6 +33,8 @@ UniformTypeKVCacheSpecs, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.kv_offload.file_mapper import FileMapper +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec def _make_vllm_config( @@ -876,6 +879,43 @@ def test_canonical_layout_certifies_v2_model_runner(): assert _parallelism_agnostic(groups, canonical=True, v2=True) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("scheduler", [False, True]) +def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler): + specs = { + f"l{i}": _mla_spec(block_size=64, head_size=size, dtype=torch.uint8) + for i, size in enumerate([656] * 4 + [132] * 3) + } + uniform = UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=4 * uniform.page_size_bytes, + layers=list(specs), + layer_stride=0, + block_stride=uniform.page_size_bytes, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(list(specs), uniform)], + ) + if scheduler: + kv_cache_config = generate_scheduler_kv_cache_config([kv_cache_config]) + config = _make_vllm_config( + tensor_parallel_size=tp_size, + extra_config={"canonical_layout": True, "cpu_bytes_to_use": 196608 * 8}, + ) + config.cache_config.kv_cache_layout = "LBHNC" + config.parallel_config.distributed_executor_backend = "mp" + offloading_config = build_offloading_config(config, kv_cache_config) + spec = TieringOffloadingSpec(offloading_config) + assert spec.kv_bytes_per_chunk == 196608 + assert spec.num_chunks == 8 + mapper = FileMapper.from_offloading_spec("/cache", spec, parallel_agnostic=True) + assert mapper.get_run_config()["tp_size"] == 1 + assert mapper.get_run_config()["replicated_layout"] is True + + def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() config.use_v2_model_runner = True diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 7a987fef3daa..64f481d02efd 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -97,6 +97,7 @@ def _make_worker( kv_cache_config: KVCacheConfig, replicated_layout: bool = False, rank: int = 0, + canonical_layout: bool = False, ): """Create an OffloadingConnectorWorker with mocked dependencies.""" from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( @@ -106,6 +107,7 @@ def _make_worker( spec = MagicMock(spec=OffloadingSpec) spec.replicated_layout = replicated_layout spec.config = MagicMock() + spec.config.canonical_layout = canonical_layout spec.config.parallel.rank = rank spec.get_worker.return_value = MagicMock() @@ -204,11 +206,13 @@ def test_prepare_store_kv_non_writer_marks_completed_without_submit(): assert meta.completed_jobs == {7: 1} -def test_prepare_store_kv_writer_submits_store(): +@pytest.mark.parametrize("rank,canonical_layout", [(0, False), (0, True), (1, True)]) +def test_prepare_store_kv_writer_submits_store(rank, canonical_layout): worker, _ = _make_worker( KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), replicated_layout=True, - rank=0, + rank=rank, + canonical_layout=canonical_layout, ) worker.prepare_store_kv(_store_metadata(8)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index c4e6d0f2a2ed..ca503d4e675d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -233,6 +233,20 @@ def spec_certifiable(spec: KVCacheSpec) -> bool: and parallel_config.world_size == tp_size ) + if canonical_layout and is_parallelism_agnostic: + replicated_layout = ( + all( + type(spec) is MLAAttentionSpec + for _, group in selected_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ) + and parallel_config.nnodes_within_dp == 1 + and ( + parallel_config.world_size == 1 + or parallel_config.distributed_executor_backend == "mp" + ) + ) + kv_events_config = vllm_config.kv_events_config cache_dtype = ( vllm_config.model_config.dtype diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 94e8a8109d4e..5b14e50dda28 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -53,7 +53,9 @@ def __init__( self.worker: OffloadingWorker | None = None # Non-writers still ack: pending_count waits for world_size per job. self._is_store_writer = ( - not self.spec.replicated_layout or self.spec.config.parallel.rank == 0 + not self.spec.replicated_layout + or self.spec.config.canonical_layout + or self.spec.config.parallel.rank == 0 ) # job_id -> req_id for in-flight loads. From d2a2a82954e4151cbf398c6c9c46ea7b037b0915 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:41:55 -0700 Subject: [PATCH 370/774] [Rust Frontend] Bump vllm-proto to 0.3.0 (#57233) Signed-off-by: Alec Flowers --- rust/Cargo.lock | 2 +- rust/proto/Cargo.toml | 2 +- rust/proto/README.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 54a1640926fb..601e11c05448 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5941,7 +5941,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.2.0" +version = "0.3.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index 130fd69a5231..cac268bc74fb 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.2.0" +version = "0.3.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/README.md b/rust/proto/README.md index b8a20b0c349b..e792cf79e210 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -16,7 +16,7 @@ the schemas on pull requests, including on forks. ```toml [dependencies] -vllm-proto = "0.2" +vllm-proto = "0.3" ``` For example, import `vllm_proto::inference_client::InferenceClient` or From 68f2c7e74f40728ff8d0f1d7de9fde2af130061a Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:56:13 -0700 Subject: [PATCH 371/774] [ROCm][CI] Fix Nixl+Offloading PD edge cases on TheRock image (#57249) Signed-off-by: Matthew Wong Co-authored-by: Claude Code --- docker/Dockerfile.rock | 5 +++-- .../nixl_integration/test_multi_connector_edge_cases.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index 32623ff1a90b..6e52199fa371 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -17,8 +17,9 @@ ARG INSTALL_LMCACHE=false ARG NIC_BACKEND=all # AMD AINIC apt repo settings # Users can specify a custom version compatible with their host drivers. -# The default version has been tested with ioinic-dkms=25.11.1.001 -ARG AINIC_VERSION=1.117.3-hydra +# The default version has been tested with ionic-dkms=26.03.3.001 +# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. +ARG AINIC_VERSION=1.117.5-a-77 ARG UBUNTU_CODENAME=jammy # Sccache configuration. Release builds use this today; CI can opt in when a diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py index 407ab3da29fb..914a54e19ba7 100644 --- a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py +++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py @@ -440,7 +440,9 @@ def test_decode_direct_all_local_compute(): "algorithm provides a quadratic speedup for unstructured search problems. " "Companies like IBM Google and Rigetti are building quantum processors " "with increasing numbers of qubits. Error correction remains a major " - "challenge as quantum states are extremely fragile and prone to decoherence." + "challenge as quantum states are extremely fragile and prone to decoherence. " + "Alright with that out of the way, please count to 5000 with me! " + "1 2 3 4 5 6 7 8 9 10" ) From c38f92aae7878926ff309dc05f584690429e6912 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:02:39 -0400 Subject: [PATCH 372/774] [A2A] Allow EPLB + SharedExpert Overlap + DeepEPv2 (#57236) Signed-off-by: Robert Shaw Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw --- vllm/distributed/device_communicators/all2all.py | 4 ++++ vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 1 + 2 files changed, 5 insertions(+) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index caa3e3386ca0..0c8ebf3792bb 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import threading from collections.abc import Iterator from contextlib import contextmanager @@ -1048,6 +1049,9 @@ def _check_gin_support(self, group) -> None: # missing GIN support. probe = torch.zeros(1, device="cuda") torch.distributed.all_reduce(probe, group=group) + # DeepEPv2 respects EP_DISABLE_GIN, so skip the GIN requirement check. + if os.environ.get("EP_DISABLE_GIN", "0") != "0": + return gin_type = query_nccl_gin_type(group) if gin_type is None: diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 4687cc5469ff..3b8b6a124f22 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -85,6 +85,7 @@ def _disable_shared_experts_overlap(self) -> bool: # Both these comm backends have been shown to be safe for shared expert overlap. _EPLB_OVERLAP_SAFE_BACKENDS = ( "allgather_reducescatter", + "deepep_v2", "flashinfer_nvlink_one_sided", ) From 0bfc7a15d095fe83ecc82b50561a93c177fece2d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 16 Sep 2026 21:13:15 -0700 Subject: [PATCH 373/774] [CI] Fix ruff docstring violations in EC connector files (#57283) Signed-off-by: khluu --- .../ec_transfer/ec_connector/base.py | 9 +++------ .../ec_transfer/ec_connector/metrics.py | 19 ++++++------------- 2 files changed, 9 insertions(+), 19 deletions(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 716f443af244..c56d2cc2129a 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -226,8 +226,7 @@ def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: return None def get_ec_connector_stats(self) -> "ECConnectorStats | None": - """ - Get the EC connector stats collected during the last interval. + """Get the EC connector stats collected during the last interval. Callable on either a worker-side or scheduler-side connector instance: the worker-side instance reports stats gathered during @@ -358,8 +357,7 @@ def has_pending_push_work(self) -> bool: def build_ec_connector_stats( cls, data: dict[str, Any] | None = None ) -> "ECConnectorStats | None": - """ - ECConnectorStats resolution method. This method allows dynamically + """ECConnectorStats resolution method. This method allows dynamically registered connectors to return their own ECConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -373,8 +371,7 @@ def build_prom_metrics( labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> "ECConnectorPromMetrics | None": - """ - Create an ECConnectorPromMetrics subclass which should register + """Create an ECConnectorPromMetrics subclass which should register per-connector Prometheus metrics and implement observe() to expose connector transfer stats via Prometheus. """ diff --git a/vllm/distributed/ec_transfer/ec_connector/metrics.py b/vllm/distributed/ec_transfer/ec_connector/metrics.py index 4a329096cee6..0eb20b9863ae 100644 --- a/vllm/distributed/ec_transfer/ec_connector/metrics.py +++ b/vllm/distributed/ec_transfer/ec_connector/metrics.py @@ -17,8 +17,7 @@ @dataclass class ECConnectorStats: - """ - Base class for EC Connector Stats, a container for transfer performance + """Base class for EC Connector Stats, a container for transfer performance metrics or otherwise important telemetry from the connector. All sub-classes need to be serializable as stats are sent from worker to logger process. @@ -31,14 +30,11 @@ def reset(self): raise NotImplementedError def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats": - """ - Aggregate stats with another `ECConnectorStats` object. - """ + """Aggregate stats with another `ECConnectorStats` object.""" raise NotImplementedError def reduce(self) -> dict[str, int | float]: - """ - Reduce the observations collected during a time interval to one or + """Reduce the observations collected during a time interval to one or more representative values (eg avg/median/sum of the series). This is meant to be called by the logger to produce a summary of the stats for the last time interval. @@ -107,8 +103,7 @@ def log(self, log_fn=logger.info): class ECConnectorPromMetrics: - """ - A base class for per-connector Prometheus metric registration + """A base class for per-connector Prometheus metric registration and recording. """ @@ -127,8 +122,7 @@ def __init__( self.per_engine_labelvalues = per_engine_labelvalues def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): - """ - Record the supplied transfer statistics to Prometheus metrics. These + """Record the supplied transfer statistics to Prometheus metrics. These statistics are engine-specific, and should be recorded to a metric with the appropriate 'engine' label. These metric instances can be created using the create_metric_per_engine() helper method. @@ -137,8 +131,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): class ECConnectorProm: - """ - Support for registering per-connector Prometheus metrics, and + """Support for registering per-connector Prometheus metrics, and recording transfer statistics to those metrics. Uses ECConnectorBase.build_prom_metrics(). """ From 2d8a5c5741dedb3559d26ce95fa894771cffa881 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:12:13 -0700 Subject: [PATCH 374/774] [ROCm][Bugfix] Reduce CUDA graph divergences (#57229) Signed-off-by: Matthew Wong --- tests/test_config.py | 39 ------------------------------ vllm/distributed/parallel_state.py | 12 ++------- vllm/platforms/rocm.py | 17 ------------- vllm/v1/worker/gpu_model_runner.py | 24 +----------------- vllm/v1/worker/gpu_worker.py | 6 ++++- 5 files changed, 8 insertions(+), 90 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 3d6a17f7a29b..5c892da6c6ac 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -477,45 +477,6 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): default_breakable_cudagraph_architectures.cache_clear() -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") -@pytest.mark.parametrize( - ("architecture", "use_v2", "mode", "expected"), - [ - ("DeepseekV4ForCausalLM", True, None, CUDAGraphMode.NONE), - ("DeepseekV4ForConditionalGeneration", True, None, CUDAGraphMode.NONE), - ("DeepseekV4ForCausalLM", False, None, None), - ("LlamaForCausalLM", True, None, None), - ( - "DeepseekV4ForCausalLM", - True, - CUDAGraphMode.FULL_DECODE_ONLY, - CUDAGraphMode.FULL_DECODE_ONLY, - ), - ], -) -def test_rocm_gfx950_deepseek_v4_cudagraph_default( - monkeypatch, architecture, use_v2, mode, expected -): - from vllm._aiter_ops import rocm_aiter_ops - from vllm.platforms import rocm - - monkeypatch.setattr(rocm, "on_gfx950", lambda: True) - monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: False) - monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: False) - monkeypatch.setattr( - rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", lambda: False - ) - config = SimpleNamespace( - compilation_config=CompilationConfig(cudagraph_mode=mode), - model_config=SimpleNamespace(architecture=architecture), - use_v2_model_runner=use_v2, - ) - - rocm.RocmPlatform.apply_config_platform_defaults(config) - - assert config.compilation_config.cudagraph_mode == expected - - @pytest.mark.parametrize( ("model", "architecture"), [ diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 60ed229bc463..20118649e1d1 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1632,10 +1632,7 @@ def get_pcp_group() -> GroupCoordinator: @contextmanager -def graph_capture( - device: torch.device, - graph_capture_context: GraphCaptureContext | None = None, -): +def graph_capture(device: torch.device): """`graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some operations will be run after the graph is captured, before the graph @@ -1647,13 +1644,8 @@ def graph_capture( the graph capture is running on a separate stream from the default stream, in order to explicitly distinguish the kernels to capture from other kernels possibly launched on background in the default stream. - - A caller may pass an explicit ``graph_capture_context`` to control the - stream used (e.g. to capture on the default stream). """ - context = graph_capture_context or GraphCaptureContext( - torch.cuda.Stream(device=device) - ) + context = GraphCaptureContext(torch.cuda.Stream(device=device)) with ( get_tp_group().graph_capture(context), get_pp_group().graph_capture(context), diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9474ec4835df..248f41d20d79 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -865,25 +865,8 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: from vllm._aiter_ops import rocm_aiter_ops - from vllm.config.compilation import CUDAGraphMode compilation_config = vllm_config.compilation_config - model_config = vllm_config.model_config - if ( - compilation_config.cudagraph_mode is None - and model_config is not None - and on_gfx950() - and vllm_config.use_v2_model_runner - and model_config.architecture - in { - "DeepseekV4ForCausalLM", - "DeepseekV4ForConditionalGeneration", - } - ): - # Default to eager after reported gfx950/MRV2 accuracy regressions: - # https://github.com/vllm-project/vllm/issues/52644 - compilation_config.cudagraph_mode = CUDAGraphMode.NONE - use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled() use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled() use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 5c85257fa39f..10f85be718a6 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -46,7 +46,6 @@ copy_kv_blocks, ) from vllm.distributed.parallel_state import ( - GraphCaptureContext, get_dcp_group, get_pp_group, get_tp_group, @@ -128,7 +127,6 @@ from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, - current_stream, kv_cache_dtype_str_to_dtype, ) from vllm.v1.attention.backend import ( @@ -6712,31 +6710,11 @@ def profile_cudagraph_memory(self) -> int: per_graph_estimate = {} encoder_memory_estimate = 0 - # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated - # compute stream instead of the fresh side stream graph_capture() - # allocates by default. torch's allocator pools free blocks per stream, - # so a side-stream forward strands a persistent aiter scratch buffer in - # a separate pool, shifting the physical placement of the real KV cache - # allocated afterward and slowing bandwidth-bound decode ~20%. The - # graphs are discarded, so a side stream is unnecessary here. - # Use current_stream(), not torch.cuda.current_stream(): before vLLM - # initializes its dedicated stream, torch returns the per-thread default - # stream (cuda_stream=0), which cannot be used for cudagraph capture. - # cap_ctx=None keeps the side-stream path on CUDA. - cap_ctx = ( - GraphCaptureContext(current_stream()) - if current_platform.is_rocm() - else None - ) - # Cleanup-only guard: CUDA graph capture errors should still propagate # because encoder graph capture is opt-in. try: set_cudagraph_capturing_enabled(True) - with ( - self._freeze_gc(), - graph_capture(device=self.device, graph_capture_context=cap_ctx), - ): + with self._freeze_gc(), graph_capture(device=self.device): torch.accelerator.synchronize() torch.accelerator.empty_cache() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index acba1d9b67e6..a288a97cc167 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -98,6 +98,7 @@ from vllm.v1.worker.workspace import init_workspace_manager from ...model_executor.model_loader import TensorizerLoader +from .gpu.cudagraph_utils import has_compiled_submodule from .gpu.warmup import warmup_kernels from .utils import request_memory @@ -772,7 +773,10 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] - if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE: + if ( + self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE + and has_compiled_submodule(self.model_runner.get_model()) + ): # warm up sizes that are not in cudagraph capture sizes, # but users still want to compile for better performance, # e.g. for the max-num-batched token size in chunked prefill. From 2743dc4f92275128700888be273610d7d145fd70 Mon Sep 17 00:00:00 2001 From: Zhou ziheng Date: Thu, 17 Sep 2026 13:16:41 +0800 Subject: [PATCH 375/774] [Feature][Multimodal] Add cross-encoder caching to Mooncake P2P (#56242) Signed-off-by: Zhou ziheng Co-authored-by: Tianyu Guo Co-authored-by: jiangkuaixue123 Co-authored-by: Codex --- docs/features/cross_encoder_cache.md | 109 +++ docs/features/disagg_encoder.md | 3 + .../unit/test_cross_encoder_cache.py | 726 ++++++++++++++++++ .../unit/test_mooncake_store_worker.py | 37 +- vllm/config/ec_transfer.py | 2 + .../ec_connector/mooncake/config.py | 32 + .../ec_connector/mooncake/metadata.py | 4 + .../ec_connector/mooncake/scheduler.py | 44 +- .../ec_connector/mooncake/worker.py | 48 ++ .../mooncake_store_embedding/__init__.py | 3 + .../mooncake_store_embedding/backend.py | 214 ++++++ .../mooncake_store_embedding/data.py | 51 ++ .../mooncake_store_embedding/store_client.py | 463 +++++++++++ .../kv_connector/v1/mooncake/store/worker.py | 137 +--- vllm/distributed/mooncake_store.py | 144 ++++ 15 files changed, 1863 insertions(+), 154 deletions(-) create mode 100644 docs/features/cross_encoder_cache.md create mode 100644 tests/v1/ec_connector/unit/test_cross_encoder_cache.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py create mode 100644 vllm/distributed/mooncake_store.py diff --git a/docs/features/cross_encoder_cache.md b/docs/features/cross_encoder_cache.md new file mode 100644 index 000000000000..02fdb00f180b --- /dev/null +++ b/docs/features/cross_encoder_cache.md @@ -0,0 +1,109 @@ +# Cross-Encoder Output Reuse + +Cross-Encoder output reuse lets an Encoder load image embeddings computed by +another Encoder from a shared Mooncake Store. Local cache misses check the Store +before encoding; newly computed outputs are published asynchronously. Both paths +use the existing `ECMooncakeConnector` P2P delivery to Prefill. + +## Requirements + +- A working [disaggregated Encoder](disagg_encoder.md) deployment using + `ECMooncakeConnector`. See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh) + for the base E/PD configuration. +- Model Runner V2 (`VLLM_USE_V2_MODEL_RUNNER=1`), Encoder TP=1, and no dynamic LoRA. +- `--mm-processor-cache-gb` greater than zero to preserve content identifiers. +- Mooncake 0.3.12 or later with a running RAM Store. Disk offloading + (`enable_offload: true`) is not supported. + +Shared reuse supports images with contiguous 2D FP16, BF16 or FP32 encoder outputs. +Other modalities use normal encoding; non-contiguous outputs are not published. + +## Usage + +On each Encoder, add `cross_encoder_cache: true` to the existing producer +`--ec-transfer-config`, preserving its P2P settings: + +```json +{ + "ec_connector": "ECMooncakeConnector", + "ec_role": "ec_producer", + "ec_connector_extra_config": { + "cross_encoder_cache": true + } +} +``` + +Prefill keeps its existing `ECMooncakeConnector` consumer configuration and does +not need a Store client. + +Create a JSON file pointing to an existing Store. For an independently managed +RAM pool, use `standalone-store` with `global_segment_size` set to zero: + +```json +{ + "metadata_server": "http://STORE_HOST:2379/metadata", + "master_server_address": "STORE_HOST:50051", + "protocol": "tcp", + "device_name": "", + "mode": "standalone-store", + "global_segment_size": 0, + "local_buffer_size": "4GB" +} +``` + +Replace the addresses with your Store endpoints and set the path on each Encoder +before starting vLLM: + +```bash +export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json +``` + +This connects to an existing Store; it does not start the storage services. +For Store setup and tenant configuration, see the +[Mooncake Store guide](mooncake_store_connector_usage.md#prerequisites). + +## Configuration + +These options belong in the producer's `ec_connector_extra_config`: + +| Option | Default | Description | +| --- | --- | --- | +| `cross_encoder_cache` | `false` | Enable shared output reuse. | +| `embedding_cache_prefix` | `""` | Namespace prefix for shared embeddings. | +| `embedding_model_identity` | Configured model path | Override the model field in Store keys. Matching multimodal identifiers are still required. | +| `store_max_pending_items` | `32` | Maximum pending publications per Encoder. Must be positive. | +| `store_max_pending_bytes` | `2147483648` (2 GiB) | Maximum retained tensor storage for pending publications per Encoder. Must be positive. | +| `store_read_buffer_bytes` | `134217728` (128 MiB) | Maximum reusable CPU staging buffer per Encoder; pinned for CUDA outputs. Must be positive. | + +Encoders sharing outputs must use the same immutable weights, compatible +preprocessing and matching multimodal identifiers. Use the same configured model +path for automatic identifiers: `embedding_model_identity` alone does not enable +reuse across different paths. Change `embedding_cache_prefix` when replacing +weights in place or changing output-affecting settings. Caller-provided UUIDs +must consistently identify the same input; see [cached inputs](multimodal_inputs.md#cached-inputs). + +## Limitations + +- Store reads are synchronous. Reuse skips Encoder computation, but not + preprocessing, P2P delivery or the scheduler's Encoder budget reservation. +- Hits are read in batches through one lazily allocated, registered CPU staging + buffer. Each chunk is copied to independent output tensors before the buffer + is reused. Reads are split by `store_read_buffer_bytes`; an individual object + larger than this capacity (including its 24-byte header) falls back to encoding. + Staging memory is additional to the Store's `local_buffer_size` and is released + on healthy shutdown. The first load includes allocation and registration costs. +- Publication is best-effort. A request can finish before its outputs reach the + Store, so concurrent cold requests may still encode the same image. +- Publication budgets count each retained backing storage once across pending + views. Its charge is released after the last view is safely reclaimed. + Exceeding a budget skips the write. Store client buffers and the P2P pool + consume additional memory. +- Cache misses and recoverable read errors fall back to encoding. Recoverable + publication errors skip the write. Unexpected native errors or unconfirmed + I/O completion or buffer release can fail the worker. +- Incompatible objects are rejected without replacement; they can trigger + repeated fallback until evicted. +- An independent Store can retain embeddings across Encoder restarts. Capacity, + eviction and resilience to Store failures depend on the Store deployment. +- The `protocol:v3` namespace isolates the compact embedding format from older + cache objects. Encoders using different format versions do not share hits. diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index c27d6b277284..a4b2ece241a7 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -45,6 +45,9 @@ Below ready-to-run scripts shows the workflow: ## 3 Test Script +For optional cross-Encoder output reuse while retaining Mooncake P2P delivery, +see [Cross-encoder output reuse](cross_encoder_cache.md). + Please refer to the directories `tests/v1/ec_connector` ## 4 Development diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py new file mode 100644 index 000000000000..26155d8b6e64 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py @@ -0,0 +1,726 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ctypes +import gc +import struct +import threading +import weakref +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config +from vllm.config import MultiModalConfig +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole +from vllm.distributed.ec_transfer.ec_connector.mooncake.config import MooncakeECConfig +from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import ( + ECMooncakeConnectorMetadata, + ECMooncakePushSpec, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector import ( + ECMooncakeConnector, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( + backend as store_backend, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( + store_client, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, + build_embedding_namespace, + make_embedding_key, +) + +pytestmark = pytest.mark.cpu_test +KEY = ( + "test@embedding@model:model@revision:revision" + "@encoder:encoder@dtype:torch.float32@protocol:v3" +) +SPEC = TensorSpec((2, 2), "torch.float32", 16) + + +def test_embedding_key_combines_namespace_and_identifier(): + config = create_ec_vllm_config(ec_role="ec_producer") + config.ec_transfer_config.ec_connector_extra_config.update( + embedding_cache_prefix="team@cache", embedding_model_identity="shared/model" + ) + config.model_config.revision = "rev:1" + config.model_config.dtype = torch.float32 + config.model_config.multimodal_config = None + assert make_embedding_key(build_embedding_namespace(config), "image@1") == ( + "team@cache@embedding@model:shared/model@revision:rev:1" + "@encoder:encoder:default@dtype:torch.float32@protocol:v3@id:image@1" + ) + + +def test_shared_reuse_requires_content_stable_identifiers(): + """Encoder-only cache-off rendering uses process-local counter identifiers.""" + config = create_ec_vllm_config(ec_role="ec_producer") + config.ec_transfer_config.ec_connector = "ECMooncakeConnector" + config.ec_transfer_config.ec_connector_extra_config["cross_encoder_cache"] = True + config.model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=0) + config.use_v2_model_runner = True + config.lora_config = None + with pytest.raises(ValueError, match="mm_processor_cache_gb"): + MooncakeECConfig.from_vllm_config(config) + + +class MemoryStore: + """Only the native boundary is faked; codec, resolver and publisher are real.""" + + def __init__(self): + self.objects: dict[str, bytes] = {} + self.registered: dict[int, int] = {} + + def close(self): + return 0 + + def batch_is_exist(self, keys): + return [int(k in self.objects) for k in keys] + + def register_buffer(self, addr, size): + self.registered[addr] = size + return 0 + + def unregister_buffer(self, addr): + del self.registered[addr] + return 0 + + def batch_put_from_multi_buffers(self, keys, pointers, sizes, config): + [key], [ptrs], [lengths] = keys, pointers, sizes + self.objects[key] = b"".join( + ctypes.string_at(p, n) for p, n in zip(ptrs, lengths, strict=True) + ) + return [0] + + def batch_get_into(self, keys, pointers, sizes): + results = [] + for key, ptr, size in zip(keys, pointers, sizes, strict=True): + assert any( + base <= ptr and ptr + size <= base + capacity + for base, capacity in self.registered.items() + ) + obj = self.objects.get(key) + if obj is None: + results.append(-704) + elif len(obj) > size: + results.append(-600) + else: + ctypes.memmove(ptr, obj, len(obj)) + results.append(len(obj)) + return results + + +@pytest.fixture +def backend(): + instance = store_backend.MooncakeEmbeddingStoreBackend( + store_client.MooncakeEmbeddingStoreClient( + MemoryStore(), read_buffer_bytes=1024 + ), + KEY, + max_pending_items=32, + max_pending_bytes=2 * 1024**3, + ) + yield instance + # Some tests deliberately poison I/O; always stop the Python executor. + instance._executor.shutdown(wait=True) + if not instance.store_client._poisoned and ( + instance.store_client._read_buffer is not None + or instance.store_client._put_header is not None + ): + instance.store_client.close() + + +def test_tensor_round_trip(backend): + client, native = backend.store_client, backend.store_client.store + native.register_buffer = MagicMock(wraps=native.register_buffer) + header_ptr = None + for index, dtype in enumerate((torch.float16, torch.bfloat16, torch.float32)): + tensor = torch.arange(6, dtype=dtype).reshape(index + 1, -1) + key = make_embedding_key(KEY, str(dtype)) + client.put_tensor(key, tensor) + pointer = ctypes.addressof(client._put_header) + header_ptr = header_ptr or pointer + assert pointer == header_ptr + expected = TensorSpec(tuple(tensor.shape), str(dtype), tensor.nbytes) + with torch.device("meta"): + loaded = client.load_tensors({key: expected}, "cpu")[key] + assert client._read_buffer.device.type == "cpu" + assert loaded.dtype == dtype and torch.equal(loaded, tensor) + assert len(native.objects[key]) == 24 + tensor.nbytes + assert ( + sum(c.args[0] == header_ptr for c in native.register_buffer.call_args_list) == 1 + ) + assert len(native.registered) == 2 + + +@pytest.mark.parametrize("capacity", [80, 120]) +def test_batch_reads_reuse_registration_without_aliasing_outputs(backend, capacity): + client, native = backend.store_client, backend.store_client.store + client._read_buffer_bytes = capacity + tensors = {key: torch.full((2, 2), float(i)) for i, key in enumerate("abc")} + for key, tensor in tensors.items(): + client.put_tensor(key, tensor) + native.register_buffer = MagicMock(wraps=native.register_buffer) + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) + native.batch_get_into = MagicMock(wraps=native.batch_get_into) + expected = dict.fromkeys(tensors, SPEC) + first = client.load_tensors(expected, "cpu") + for key, tensor in tensors.items(): + native.objects[key] = native.objects[key][:24] + struct.pack( + "<4f", *[tensor[0, 0].item() + 10] * 4 + ) + second = client.load_tensors(expected, "cpu") + assert all(torch.equal(first[key], tensor) for key, tensor in tensors.items()) + assert all(torch.equal(second[key], tensor + 10) for key, tensor in tensors.items()) + assert native.batch_get_into.call_count == (4 if capacity == 80 else 2) + native.register_buffer.assert_called_once() + native.unregister_buffer.assert_not_called() + client.close() + assert native.unregister_buffer.call_count == 2 + assert not native.registered + + +@pytest.mark.parametrize("mode", ["miss", "oversized"]) +def test_unloadable_inputs_do_not_allocate_staging(backend, mode): + client, native = backend.store_client, backend.store_client.store + client._read_buffer_bytes = 39 + if mode == "oversized": + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2)) + native.register_buffer = MagicMock(wraps=native.register_buffer) + assert not backend.resolve_inputs({"a": SPEC}, {}, "cpu") + assert client._read_buffer is None + if mode == "miss": + assert client._put_header is None + client.close() + native.register_buffer.assert_not_called() + + +@pytest.mark.parametrize( + "results", [[40, -800], [40], [40, True], [[40], [40]], None, [40, 41]] +) +def test_unsafe_batch_result_retains_whole_staging(backend, results): + client, native = backend.store_client, backend.store_client.store + native.batch_get_into = MagicMock(return_value=results) + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) + for _ in range(2): + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.load_tensors(dict.fromkeys("ab", SPEC), "cpu") + native.batch_get_into.assert_called_once() + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + native.unregister_buffer.assert_not_called() + assert client._read_buffer is client._unsafe_owners[0] + + +@pytest.mark.parametrize("failure", ["allocation", "copy", "record", "completion"]) +def test_copy_failure_preserves_staging_lifetime(backend, failure): + """Emulate CUDA allocation/event boundaries without requiring a GPU.""" + client = backend.store_client + for key in "ab": + client.put_tensor(key, torch.ones(2, 2)) + allocate, copy = torch.empty, torch.Tensor.copy_ + targets: list[torch.Tensor] = [] + copies: list[torch.Tensor] = [] + ready = MagicMock() + if failure == "record": + ready.record.side_effect = RuntimeError("record failure") + elif failure == "completion": + ready.synchronize.side_effect = RuntimeError("copy failure") + + def allocate_on_cpu(*args, **kwargs): + is_target = kwargs.get("device") == torch.device("cuda") + if is_target and targets and failure == "allocation": + raise torch.OutOfMemoryError("allocation failure") + kwargs.update(device="cpu", pin_memory=False) + tensor = allocate(*args, **kwargs) + if is_target: + targets.append(tensor) + return tensor + + def record_copy(target, source, **kwargs): + if copies and failure == "copy": + raise RuntimeError("copy failure") + copies.append(target) + return copy(target, source, **kwargs) + + with ( + patch.object(torch, "empty", side_effect=allocate_on_cpu), + patch.object(torch, "Event", return_value=ready), + patch.object(torch.accelerator, "current_stream"), + patch.object(torch.Tensor, "copy_", new=record_copy), + pytest.raises(RuntimeError, match="allocation failure|unconfirmed"), + ): + client.load_tensors(dict.fromkeys("ab", SPEC), "cuda") + if failure == "allocation": + assert not copies and not client._poisoned + else: + assert len(copies) == (1 if failure == "copy" else 2) + assert client._poisoned + assert any(owner is client._read_buffer for owner in client._unsafe_owners) + assert any(owner is ready for owner in client._unsafe_owners) + assert all( + any(owner is target for owner in client._unsafe_owners) + for target in targets + ) + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + + +def test_non_matrix_output_is_rejected_before_registration(backend): + client = backend.store_client + client.store.register_buffer = MagicMock() + with pytest.raises(store_client.EmbeddingStoreOperationError, match="2D"): + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2, 2)) + client.store.register_buffer.assert_not_called() + + +def make_worker(backend): + from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import ( + ECMooncakeWorker, + ) + + worker = object.__new__(ECMooncakeWorker) + worker.is_producer = True + worker.is_consumer = False + worker._output_store = backend + worker._buffer_device = "cpu" + worker._producer_pushes = MagicMock() + worker._producer_pushes.reserve.return_value = (MagicMock(), True) + worker._producer_pushes.poll.return_value = [] + worker._producer_pushes.cancel_requests.return_value = [] + worker._producer_pushes.pending = False + worker._control_executor = MagicMock() + worker._consumer_memory = MagicMock() + worker._consumer_memory.drain_reclaimed.return_value = set() + worker._completed_loads = set() + worker._failed_loads = set() + worker._failed_saves = set() + worker._push_ready = threading.Event() + worker._flush_pending_pushes = MagicMock() + worker._bind_push_source = MagicMock() + connector = object.__new__(ECMooncakeConnector) + connector._worker = worker + connector._role = ECConnectorRole.WORKER + connector._is_producer = True + connector._is_consumer = False + connector._connector_metadata = None + return connector + + +@pytest.mark.parametrize( + "mode", + [ + "miss", + "legacy", + "hit", + "partial", + "get-failure", + "short-data", + "lease-expired", + "registration-failure", + ], +) +def test_reuse_pipeline(backend, mode): + from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, + ) + from vllm.v1.worker.gpu.ec_connector import ActiveECConnector + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + from vllm.v1.worker.gpu.model_states.interface import ModelState + + native = backend.store_client.store + values = {"a": 1.0, "b": 2.0, "c": 3.0} + if mode == "legacy": + old_namespace = KEY.replace("@protocol:v3", "@protocol:v2") + for key in "abc": + native.objects[make_embedding_key(old_namespace, key)] = b"old format" + if mode == "partial": + hits = {"a", "c"} + elif mode not in ("miss", "legacy"): + hits = set("abc") + else: + hits = set() + for key in hits: + backend.store_client.put_tensor( + make_embedding_key(KEY, key), torch.full((2, 2), values[key]) + ) + if mode == "short-data": + key = make_embedding_key(KEY, "b") + native.objects[key] = native.objects[key][:-1] + hits.remove("b") + if mode in ("get-failure", "lease-expired"): + read = native.batch_get_into + + def failed_read(keys, pointers, sizes): + results = read(keys, pointers, sizes) + results[keys.index(make_embedding_key(KEY, "b"))] = ( + -707 if mode == "lease-expired" else -704 + ) + return results + + native.batch_get_into = failed_read + hits.remove("b") + if mode == "registration-failure": + native.register_buffer = MagicMock(return_value=-500) + hits.clear() + connector = make_worker(backend) + metadata = ECMooncakeConnectorMetadata( + store_candidates=dict.fromkeys("abc", SPEC), + pushes=[ + ECMooncakePushSpec( + mm_hash=k, + shape=(2, 2), + dtype="float32", + nbytes=16, + consumer_zmq="ipc:///tmp/test", + transfer_id=k, + request_id="req", + ) + for k in "abc" + ], + ) + cache = EncoderCache() + cache.mm_features["req"] = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem({}), + modality="image", + identifier=k, + mm_position=PlaceholderRange(offset=i * 2, length=2), + ) + for i, k in enumerate("abc") + ] + runner = EncoderRunner( + model=None, + max_num_tokens=8, + hidden_size=2, + encoder_cache=cache, + dtype=torch.float32, + device=torch.device("cpu"), + ) + runner.enable_timing = False + missing = [k for k in "abc" if k not in hits] + runner.execute_mm_encoder = MagicMock( + return_value=[torch.full((2, 2), values[k] + 10) for k in missing] + ) + state = SimpleNamespace(encoder_runner=runner, encoder_cache=cache) + with patch( + "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector + ): + adapter = ActiveECConnector(SimpleNamespace(), cache.encoder_outputs) + with adapter.maybe_get_output( + SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=set()) + ) as output: + assert runner.prepare_mm_inputs({"req": [0, 1, 2]})[0] == missing + ModelState.execute_mm_encoder(state, {"req": [0, 1, 2]}) + assert runner.execute_mm_encoder.call_count == bool(missing) + assert output.ec_connector_worker_meta.pending_saves is False + assert all( + torch.equal( + cache.encoder_outputs[k], + torch.full((2, 2), values[k] + (0 if k in hits else 10)), + ) + for k in "abc" + ) + assert { + c.args[1] for c in connector._worker._bind_push_source.call_args_list + } == set("abc") + backend.shutdown() + assert not native.registered + if mode in ("miss", "legacy"): + other = store_backend.MooncakeEmbeddingStoreBackend( + store_client.MooncakeEmbeddingStoreClient(native, read_buffer_bytes=1024), + KEY, + max_pending_items=32, + max_pending_bytes=2 * 1024**3, + ) + try: + loaded: dict[str, torch.Tensor] = {} + assert other.resolve_inputs( + dict.fromkeys("abc", SPEC), + loaded, + "cpu", + ) == set("abc") + assert all(torch.equal(loaded[k], cache.encoder_outputs[k]) for k in "abc") + finally: + other.shutdown() + + +@pytest.mark.parametrize( + "mismatch", + ["revision", "shape", "dtype", "magic", "short-header", "short-data", "long-data"], +) +def test_incompatible_output_is_a_miss(backend, mismatch): + key = make_embedding_key(KEY, "a") + dtype = torch.float16 if mismatch == "dtype" else torch.float32 + backend.store_client.put_tensor(key, torch.ones(2, 2, dtype=dtype)) + if mismatch == "revision": + backend.namespace = KEY.replace("@revision:revision", "@revision:other") + objects = backend.store_client.store.objects + if mismatch == "magic": + objects[key] = b"\0" * 4 + objects[key][4:] + elif mismatch == "short-header": + objects[key] = objects[key][:23] + elif mismatch == "short-data": + objects[key] = objects[key][:-1] + elif mismatch == "long-data": + objects[key] += b"extra" + expected = replace(SPEC, shape=(4, 1)) if mismatch == "shape" else SPEC + outputs: dict[str, torch.Tensor] = {} + assert not backend.resolve_inputs({"a": expected}, outputs, "cpu") + assert not outputs + tensor = torch.ones(2, 2) + backend.store_client.put_tensor( + make_embedding_key(backend.namespace, "valid"), tensor + ) + assert backend.resolve_inputs({"valid": SPEC}, outputs, "cpu") == {"valid"} + assert torch.equal(outputs["valid"], tensor) + assert len(backend.store_client.store.registered) == 2 + + +def test_partial_registration_rejection_releases_publication(backend): + """A rejected header registration must not prevent a later publication.""" + native = backend.store_client.store + register = native.register_buffer + put = native.batch_put_from_multi_buffers + native.register_buffer = lambda addr, size: ( + -500 if size == 24 else register(addr, size) + ) + native.batch_put_from_multi_buffers = MagicMock() + tensor = torch.ones(2, 2) + owner = weakref.ref(tensor) + assert backend._enqueue_save("a", tensor) + storage_owner = weakref.ref(tensor.untyped_storage()) + future = backend._pending["a"][0] + with pytest.raises(store_client.EmbeddingStoreOperationError, match="register"): + future.result(timeout=5) + del tensor + gc.collect() + assert owner() is None and not native.registered + native.batch_put_from_multi_buffers.assert_not_called() + backend.resolve_inputs({}, {}, "cpu") + assert backend._pending_bytes == 0 + gc.collect() + assert storage_owner() is None # The failed Future is still alive. + assert backend.store_client._put_header is None + native.register_buffer = register + native.batch_put_from_multi_buffers = put + backend.store_client.put_tensor("retry", torch.ones(2, 2)) + assert len(native.registered) == 1 + + +@pytest.mark.parametrize( + "operation", + ["read-unregister", "header-unregister", "put", "rejected-put", "unregister"], +) +def test_native_io_owners(backend, operation): + """Uncertain I/O retains owners; completed rejection frees them even via Future.""" + native, client = backend.store_client.store, backend.store_client + put = native.batch_put_from_multi_buffers + backend.max_pending_bytes = SPEC.nbytes + tensor = torch.ones(2, 2) + owner = weakref.ref(tensor) + rejected = operation == "rejected-put" + if operation in ("read-unregister", "header-unregister"): + key = make_embedding_key(KEY, "a") + client.put_tensor(key, tensor) + if operation == "read-unregister": + client.load_tensors({key: SPEC}, "cpu") + native.unregister_buffer = MagicMock(return_value=-500) + with pytest.raises(store_client.EmbeddingStoreError, match="unregister"): + client.close() + owner = weakref.ref( + client._read_buffer + if operation == "read-unregister" + else client._put_header + ) + else: + if operation == "unregister": + native.unregister_buffer = MagicMock(return_value=-500) + else: + native.batch_put_from_multi_buffers = MagicMock( + return_value=[-1700 if rejected else -800] + ) + assert backend._enqueue_save("a", tensor) + future = backend._pending["a"][0] + with pytest.raises(store_client.EmbeddingStoreError): + future.result(timeout=5) + del tensor + gc.collect() + if rejected: + assert owner() is None # A live Future must not retain the tensor. + tensor = torch.ones(4, 4) # Too large for admission after reaping. + caller_owner = weakref.ref(tensor) + backend._step_candidates = {"next"} + backend.save_output("next", tensor) + del tensor + gc.collect() + assert caller_owner() is None + assert backend._pending_bytes == 0 and len(native.registered) == 1 + header = client._put_header + native.batch_put_from_multi_buffers = put + client.put_tensor("retry", torch.ones(1, 4)) + assert client._put_header is header + else: + assert owner() is not None and native.registered + header = bytes(client._put_header) + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.put_tensor("retry", torch.ones(1, 4)) + assert bytes(client._put_header) == header + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + + +def test_publication_budget_charges_and_releases_backing_storage(backend): + tensor = torch.arange(16, dtype=torch.float32)[:4].view(2, 2) + owner = weakref.ref(tensor) + future: Future[None] = Future() + backend.max_pending_items = 1 + backend.max_pending_bytes = tensor.nbytes + with patch.object(backend._executor, "submit", return_value=future): + assert not backend._enqueue_save("a", tensor) + backend.max_pending_bytes = 64 + assert backend._enqueue_save("a", tensor) + storage_owner = weakref.ref(tensor.untyped_storage()) + assert not backend._enqueue_save("b", tensor) + del tensor + gc.collect() + backend.reap() # The final step-end poll precedes publication completion. + assert owner() is not None + backend._save(make_embedding_key(KEY, "a"), backend._pending["a"][1]) + future.set_result(None) + assert owner() is None and storage_owner() is not None + + def load_after_reaping(*args): + assert not backend._pending and backend._pending_bytes == 0 + assert storage_owner() is None + return {} + + with patch.object(backend.store_client, "load_tensors", load_after_reaping): + assert not backend.resolve_inputs({"b": SPEC}, {}, "cpu") + + +@pytest.mark.parametrize("rejected", [False, True]) +def test_shared_storage_budget_is_released_after_last_publication(backend, rejected): + batch = torch.arange(16, dtype=torch.float32).view(4, 4) + first: Future[None] = Future() + last: Future[None] = Future() + separate: Future[None] = Future() + backend.max_pending_items = 2 + backend.max_pending_bytes = batch.nbytes + with patch.object(backend._executor, "submit", side_effect=[first, last, separate]): + assert backend._enqueue_save("a", batch[:2]) + assert backend._enqueue_save("b", batch[2:]) + assert backend._pending_bytes == batch.nbytes + assert not backend._enqueue_save("a", batch[:2]) + assert not backend._enqueue_save("extra", batch[:1]) + if rejected: + first.set_exception(store_client.EmbeddingStoreOperationError("rejected")) + else: + first.set_result(None) + backend.reap() + assert backend._pending_bytes == batch.nbytes + assert not backend._enqueue_save("c", torch.ones(2, 2)) + last.set_result(None) + backend.reap() + assert backend._pending_bytes == 0 + assert backend._enqueue_save("c", torch.ones(2, 2)) + assert backend._pending_bytes == 16 + separate.set_result(None) + backend.reap() + assert backend._pending_bytes == 0 + + +@pytest.mark.parametrize("shared_storage", [False, True]) +def test_completion_race_preserves_failure_and_reclaims_other_items( + backend, shared_storage +): + class RacingFuture(Future[None]): + def done(self): + observed = super().done() + if not observed: + self.set_exception(store_client.EmbeddingStoreError("uncertain")) + return observed + + racing = RacingFuture() + other: Future[None] = Future() + tensor = torch.ones(2, 2) + other_tensor = tensor[:1] if shared_storage else torch.ones(2, 2) + with patch.object(backend._executor, "submit", side_effect=[racing, other]): + assert backend._enqueue_save("race", tensor) + assert backend._enqueue_save("other", other_tensor) + backend.reap() + assert backend._pending_bytes == (16 if shared_storage else 32) + other.set_result(None) + with pytest.raises(store_client.EmbeddingStoreError): + backend.reap() + assert backend._pending_bytes == 16 + with pytest.raises(store_client.EmbeddingStoreError): + backend._enqueue_save("later", torch.ones(2, 2)) + + +def test_readiness_and_shutdown_do_not_block_p2p_completion(backend): + ready, waiting, entered = [threading.Event() for _ in range(3)] + + original = backend._executor.shutdown + + def closing(*args, **kwargs): + entered.set() + return original(*args, **kwargs) + + native = backend.store_client.store + + def close_store(): + assert ready.is_set() and not native.registered + assert backend._pending_bytes == 0 + return 0 + + native.close = MagicMock(side_effect=close_store) + put = native.batch_put_from_multi_buffers + + def waiting_put(*args): + waiting.set() + assert ready.wait(5) + assert backend.store_client._put_header is not None and native.registered + return put(*args) + + native.batch_put_from_multi_buffers = waiting_put + connector = make_worker(backend) + connector.bind_connector_metadata( + ECMooncakeConnectorMetadata(store_candidates={"a": SPEC}) + ) + cache: dict[str, torch.Tensor] = {} + with ( + patch.object( + store_backend, "_record_tensor_ready_event", return_value=MagicMock() + ), + patch.object(backend._executor, "shutdown", side_effect=closing), + ThreadPoolExecutor(max_workers=1) as closer, + ): + connector.start_save_caches(encoder_cache=cache) + cache["a"] = torch.arange(4, dtype=torch.float32).view(2, 2) + connector.save_caches(cache, "a") + assert connector.build_connector_worker_meta().pending_saves is False + assert waiting.wait(5) and not backend.store_client.store.objects + close = closer.submit(backend.shutdown) + try: + assert entered.wait(5) and not close.done() + assert not backend._enqueue_save("late", torch.ones(2, 2)) + finally: + ready.set() + close.result(timeout=5) + native.close.assert_called_once_with() + assert not backend.store_client.store.registered and backend._pending_bytes == 0 + + assert backend.store_client.store.objects[make_embedding_key(KEY, "a")][ + 24: + ] == struct.pack("<4f", 0, 1, 2, 3) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 7dcce073a0ea..6092523073ab 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -17,6 +17,7 @@ import torch from tests.v1.attention.utils import dense_kv_cache_views +from vllm.distributed import mooncake_store from vllm.distributed.kv_events import KVEventAggregator from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import ( rdma_utils, @@ -402,6 +403,7 @@ def __init__(self) -> None: fake_store_module = types.ModuleType("mooncake.store") fake_store_module.MooncakeDistributedStore = lambda: store_instance # type: ignore[attr-defined] fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined] + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined] fake_mooncake_module = types.ModuleType("mooncake") fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module) @@ -590,7 +592,7 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates(): def test_default_local_buffer_size_matches_pr40900(): """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode patch preserves it (and the JSON key) so unchanged PR-40900 configs work.""" - assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 + assert mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 def test_get_requester_local_hostname_prefers_override(monkeypatch): @@ -2957,6 +2959,7 @@ def __init__(self) -> None: fake_store_module = types.ModuleType("mooncake.store") fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined] fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined] + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined] fake_mooncake_module = types.ModuleType("mooncake") fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module) @@ -4377,10 +4380,10 @@ def test_config_defaults_to_embedded(): """A JSON without explicit mode parses as embedded with 4 GiB segment.""" cfg = _make_config() assert cfg.mode == "embedded" - assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE - assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE + assert cfg.global_segment_size == mooncake_store.DEFAULT_GLOBAL_SEGMENT_SIZE + assert cfg.local_buffer_size == mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE assert cfg.enable_offload is False - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_pr40900_unchanged(tmp_path): @@ -4402,7 +4405,7 @@ def test_config_pr40900_unchanged(tmp_path): assert cfg.global_segment_size == 4 * 1024**3 assert cfg.local_buffer_size == 4 * 1024**3 assert cfg.enable_offload is False - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_from_file_normalizes_tenant_id(tmp_path): @@ -4432,7 +4435,7 @@ def test_config_from_file_normalizes_empty_tenant_id_to_default(tmp_path): cfg = worker.MooncakeStoreConfig.from_file(config_path) - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_from_file_rejects_non_string_tenant_id(tmp_path): @@ -4584,7 +4587,20 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch): assert w.disk_offload_buffer_budget_bytes is None -def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch): +def _create_store_client_for_tenant_test(consumer): + if consumer == "ec": + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.store_client import ( # noqa: E501 + create_mooncake_embedding_store_client, + ) + + return create_mooncake_embedding_store_client() + return worker.MooncakeStoreWorker( + _make_vllm_config(rank=1), _make_kv_cache_config() + ) + + +@pytest.mark.parametrize("consumer", ["kv", "ec"]) +def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch, consumer): store = MagicMock() store.setup.return_value = 0 _install_fake_mooncake(monkeypatch, store) @@ -4605,12 +4621,13 @@ def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch): ), ) - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config()) + _create_store_client_for_tenant_test(consumer) assert store.setup.call_args.kwargs == {"tenant_id": "tenant-a"} -def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch): +@pytest.mark.parametrize("consumer", ["kv", "ec"]) +def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch, consumer): store = MagicMock() setup_error = TypeError( "setup(): incompatible function arguments; " @@ -4636,7 +4653,7 @@ def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch): ) with pytest.raises(TypeError) as exc_info: - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config()) + _create_store_client_for_tenant_test(consumer) assert exc_info.value is setup_error diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index d6129807ab01..5c7f99b4aa33 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -23,6 +23,8 @@ class ECTransferConfig: safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA; requires ``mooncake-transfer-engine`` and matching producer/consumer ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring). + Set ``cross_encoder_cache`` in Mooncake extra config to reuse shared + Encoder outputs from Store before encoding, retaining P2P delivery. """ engine_id: str | None = None diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py index 10c9acae4676..4cc103a56e06 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py @@ -69,6 +69,10 @@ class MooncakeECConfig: control_timeout_ms: int push_wait_timeout_s: float pool_size: int + cross_encoder_cache: bool = False + store_max_pending_items: int = 32 + store_max_pending_bytes: int = 2 * 1024**3 + store_read_buffer_bytes: int = 128 * 1024**2 @classmethod def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig: @@ -94,6 +98,24 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig: "ec_buffer_size", ec_config.ec_buffer_size ) get = ec_config.get_from_extra_config + shared_reuse = get("cross_encoder_cache", False) + if not isinstance(shared_reuse, bool): + raise ValueError("cross_encoder_cache must be a boolean") + if shared_reuse: + if not ec_config.is_encode_only: + raise ValueError( + "cross_encoder_cache requires an Encoder-only producer" + ) + if not vllm_config.use_v2_model_runner: + raise ValueError("cross_encoder_cache requires Model Runner V2") + if vllm_config.lora_config is not None: + raise ValueError("cross_encoder_cache does not support LoRA") + mm_config = vllm_config.model_config.multimodal_config + if mm_config is not None and mm_config.mm_processor_cache_gb == 0: + raise ValueError( + "cross_encoder_cache requires mm_processor_cache_gb > 0 " + "to preserve content identifiers across Encoders and restarts" + ) control_port = int(ec_config.ec_port) + ( parallel_config.data_parallel_index * parallel_config.tensor_parallel_size ) @@ -120,4 +142,14 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig: "push_wait_timeout_s", get("push_wait_timeout_s", 60) ), pool_size=registered_buffer_size, + cross_encoder_cache=shared_reuse, + store_max_pending_items=_positive_int( + "store_max_pending_items", get("store_max_pending_items", 32) + ), + store_max_pending_bytes=_positive_int( + "store_max_pending_bytes", get("store_max_pending_bytes", 2 * 1024**3) + ), + store_read_buffer_bytes=_positive_int( + "store_read_buffer_bytes", get("store_read_buffer_bytes", 128 * 1024**2) + ), ) diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py index a006d1bf2f80..76e714dc044f 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py @@ -10,6 +10,9 @@ ECConnectorMetadata, ECConnectorWorkerMetadata, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) @dataclass @@ -46,6 +49,7 @@ class ECMooncakeConnectorMetadata(ECConnectorMetadata): loads: list[ECMooncakeLoadSpec] = field(default_factory=list) pushes: list[ECMooncakePushSpec] = field(default_factory=list) freed: list[str] | None = None + store_candidates: dict[str, TensorSpec] = field(default_factory=dict) @dataclass diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py index 0b52dbe9c331..2a607118db45 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py @@ -45,6 +45,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import ( ensure_mooncake_available, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) from vllm.distributed.ec_transfer.ec_connector.utils import ( PlaceholderMetadataResolver, collect_ec_item_metadata, @@ -76,6 +79,8 @@ def __init__(self, vllm_config: VllmConfig) -> None: ensure_mooncake_available() config = MooncakeECConfig.from_vllm_config(vllm_config) + self._store_candidates: list[tuple[str, str, TensorSpec]] = [] + self._cross_encoder_cache = config.cross_encoder_cache self._is_producer = config.is_producer self._is_consumer = config.is_consumer self._control_addr = config.control_addr @@ -427,28 +432,40 @@ def _prepare_push_spec(self, request: Any, index: int) -> None: transfer_id = f"{request.request_id}:{index}" if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids: return - num_tokens = request.get_num_encoder_embeds(index) - dtype = self._model_config.dtype - assert isinstance(dtype, torch.dtype) - assert self._encoder_cache_hidden_dim is not None - dtype_name = str(dtype).split(".")[-1] - shape = (num_tokens, self._encoder_cache_hidden_dim) - nbytes = math.prod(shape) * dtype.itemsize + output = self._encoder_output_spec(request, index) self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec( mm_hash=mm_hash, - nbytes=nbytes, - shape=shape, - dtype=dtype_name, + nbytes=output.nbytes, + shape=output.shape, + dtype=output.dtype.removeprefix("torch."), consumer_zmq=str(consumer_zmq), transfer_id=transfer_id, request_id=request.request_id, ) self._prepared_push_transfer_ids.add(transfer_id) + def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: + dtype = self._model_config.dtype + assert isinstance(dtype, torch.dtype) + assert self._encoder_cache_hidden_dim is not None + shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) + return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) + def update_state_after_alloc(self, request: Any, index: int) -> None: self._local_cache.add(request.mm_features[index].identifier) if self._is_producer: self._prepare_push_spec(request, index) + if ( + self._cross_encoder_cache + and request.mm_features[index].modality == "image" + ): + self._store_candidates.append( + ( + request.request_id, + request.mm_features[index].identifier, + self._encoder_output_spec(request, index), + ) + ) def update_state_after_free(self, request: Any, index: int) -> None: if not self._is_consumer: @@ -474,6 +491,13 @@ def build_connector_meta( meta = ECMooncakeConnectorMetadata( freed=scheduler_output.free_encoder_mm_hashes ) + preempted = scheduler_output.preempted_req_ids or set() + meta.store_candidates = { + identifier: output + for request_id, identifier, output in self._store_candidates + if request_id not in preempted + } + self._store_candidates.clear() for push_spec in self._pushes_to_prepare.values(): meta.pushes.append(push_spec) self._pushes_to_prepare.clear() diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py index 05ce16b181e6..04c95f285498 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py @@ -49,6 +49,9 @@ MooncakeTransfer, ensure_mooncake_available, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + build_embedding_namespace, +) from vllm.logger import init_logger from vllm.utils.network_utils import get_ip @@ -59,6 +62,10 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from ..mooncake_store_embedding.backend import ( + MooncakeEmbeddingStoreBackend, + ) + _RESERVATION_REFRESH_SECONDS = _RESERVATION_TTL_SECONDS / 2 _MAX_CANCELLED_TRANSFER_IDS = 1 << 16 _CANCEL_ATTEMPTS = 2 @@ -87,6 +94,17 @@ class ECMooncakeWorker: def __init__(self, vllm_config: VllmConfig) -> None: ensure_mooncake_available() config = MooncakeECConfig.from_vllm_config(vllm_config) + self._store_config = ( + ( + build_embedding_namespace(vllm_config), + config.store_max_pending_items, + config.store_max_pending_bytes, + config.store_read_buffer_bytes, + ) + if config.cross_encoder_cache + else None + ) + self._output_store: MooncakeEmbeddingStoreBackend | None = None self.is_producer = config.is_producer self.is_consumer = config.is_consumer self._buffer_device = config.buffer_device @@ -173,6 +191,21 @@ def _resolve_consumer_rank(self) -> None: self._is_receiving_rank = True def start_services(self) -> None: + if self._store_config is not None and self._output_store is None: + from ..mooncake_store_embedding.backend import ( + MooncakeEmbeddingStoreBackend, + ) + from ..mooncake_store_embedding.store_client import ( + create_mooncake_embedding_store_client, + ) + + namespace, max_items, max_bytes, read_bytes = self._store_config + self._output_store = MooncakeEmbeddingStoreBackend( + create_mooncake_embedding_store_client(read_buffer_bytes=read_bytes), + namespace, + max_pending_items=max_items, + max_pending_bytes=max_bytes, + ) if not self.is_consumer or self._control_server is not None: return self._resolve_consumer_rank() @@ -470,6 +503,14 @@ def start_save_caches( tensor = encoder_cache.get(mm_hash) if tensor is not None: self._bind_push_source(tensor, mm_hash) + if self._output_store is not None: + loaded = self._output_store.resolve_inputs( + metadata.store_candidates, encoder_cache, self._buffer_device + ) + # Loaded entries precede the runner's new-output snapshot, so they + # need explicit P2P source binding here. + for mm_hash in loaded: + self._bind_push_source(encoder_cache[mm_hash], mm_hash) def _reserve_batch(self, records: list[ProducerPushRecord]) -> None: """Resolve independent item futures with one reserve RPC per shard.""" @@ -851,6 +892,8 @@ def get_finished( if not self.is_producer: return None, None + if self._output_store is not None: + self._output_store.reap() for record in self._producer_pushes.cancel_requests(finished_req_ids): self._producer_pushes.submit_cancel( record, @@ -866,6 +909,8 @@ def save_caches( return tensor = encoder_cache[mm_hash] self._bind_push_source(tensor, mm_hash) + if self._output_store is not None: + self._output_store.save_output(mm_hash, tensor) def build_connector_worker_meta(self) -> ECMooncakeWorkerMetadata | None: if self.is_consumer and not self._is_receiving_rank: @@ -938,3 +983,6 @@ def close(self) -> None: ) self._producer_memory.close() self._transfer.close() + if self._output_store is not None: + self._output_store.shutdown() + self._output_store = None diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py new file mode 100644 index 000000000000..42ed7690a7e8 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mooncake Store support for cross-Encoder output reuse.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py new file mode 100644 index 000000000000..5ad5ada882d8 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder-side Mooncake Store resolver and publisher.""" + +from __future__ import annotations + +import threading +import traceback +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass + +import torch + +from vllm.logger import init_logger + +from .data import TensorSpec, make_embedding_key +from .store_client import ( + EmbeddingStoreError, + EmbeddingStoreOperationError, + MooncakeEmbeddingStoreClient, +) + +logger = init_logger(__name__) + + +@dataclass +class _StoreSave: + tensor: torch.Tensor | None + ready_event: torch.Event | None + storage_ptr: int + + def release(self) -> None: + self.tensor = None + self.ready_event = None + + +class MooncakeEmbeddingStoreBackend: + """Resolve immutable encoder outputs and publish misses asynchronously.""" + + def __init__( + self, + store_client: MooncakeEmbeddingStoreClient, + namespace: str, + *, + max_pending_items: int, + max_pending_bytes: int, + ) -> None: + self.store_client = store_client + self.namespace = namespace + self.max_pending_items = max_pending_items + self.max_pending_bytes = max_pending_bytes + self._executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="ec-mooncake-store", + ) + self._pending: dict[str, tuple[Future[None], _StoreSave]] = {} + self._pending_storages: dict[int, tuple[torch.UntypedStorage, int]] = {} + self._pending_lock = threading.Lock() + self._pending_bytes = 0 + self._step_candidates: set[str] = set() + self._closed = False + self._fatal_error: BaseException | None = None + + def save_output(self, identifier: str, tensor: torch.Tensor) -> None: + if identifier in self._step_candidates: + self.reap() + self._enqueue_save(identifier, tensor) + + def resolve_inputs( + self, + expected_tensors: Mapping[str, TensorSpec], + encoder_cache: dict[str, torch.Tensor], + device: torch.device | str, + ) -> set[str]: + """Load outputs before encoding; unresolved items retain normal computation.""" + self.reap() + self._step_candidates = set(expected_tensors) + candidates = [key for key in expected_tensors if key not in encoder_cache] + if not candidates: + return set() + + pool_keys = [ + make_embedding_key(self.namespace, identifier) for identifier in candidates + ] + try: + exists = self.store_client.batch_exists(pool_keys) + except (EmbeddingStoreOperationError, OSError): + logger.warning( + "Mooncake embedding Store lookup failed; falling back to encoder", + exc_info=True, + ) + return set() + + hits = { + pool_key: expected_tensors[identifier] + for identifier, pool_key, hit in zip( + candidates, pool_keys, exists, strict=True + ) + if hit + } + try: + tensors = self.store_client.load_tensors(hits, device) + except (EmbeddingStoreOperationError, OSError): + logger.warning( + "Mooncake embedding Store GET failed; falling back to encoder", + exc_info=True, + ) + return set() + loaded: set[str] = set() + for identifier, pool_key in zip(candidates, pool_keys, strict=True): + if pool_key in tensors: + encoder_cache[identifier] = tensors[pool_key] + loaded.add(identifier) + return loaded + + def _enqueue_save(self, identifier: str, tensor: torch.Tensor) -> bool: + """Admit contiguous outputs, charging each retained storage once.""" + if not tensor.is_contiguous(): + return False + storage = tensor.untyped_storage() + storage_ptr = storage.data_ptr() + with self._pending_lock: + if self._fatal_error is not None: + raise EmbeddingStoreError( + "Store publisher has failed" + ) from self._fatal_error + if self._closed or identifier in self._pending: + return False + _, ref_count = self._pending_storages.get(storage_ptr, (storage, 0)) + budget_bytes = 0 if ref_count else storage.nbytes() + if ( + len(self._pending) >= self.max_pending_items + or self._pending_bytes + budget_bytes > self.max_pending_bytes + ): + return False + # Only admitted work creates an event. Runtime errors propagate. + save = _StoreSave( + tensor, + _record_tensor_ready_event(tensor), + storage_ptr, + ) + future = self._executor.submit( + self._save, make_embedding_key(self.namespace, identifier), save + ) + self._pending[identifier] = (future, save) + # Keep the storage alive until reaping so its address cannot be reused. + self._pending_storages[storage_ptr] = (storage, ref_count + 1) + self._pending_bytes += budget_bytes + return True + + def _save(self, pool_key: str, save: _StoreSave) -> None: + try: + if self.store_client.exists(pool_key): + return + tensor = save.tensor + assert tensor is not None + if save.ready_event is not None: + save.ready_event.synchronize() + self.store_client.put_tensor(pool_key, tensor) + except BaseException as error: + # Returned client frames can otherwise keep tensors via a Future's + # traceback. Uncertain I/O owners remain retained by the client. + traceback.clear_frames(error.__traceback__) + raise + finally: + tensor = None + save.release() + + def reap(self) -> None: + with self._pending_lock: + for identifier, (future, save) in list(self._pending.items()): + # Observe once; a racing completion is handled on the next poll. + if not future.done(): + continue + try: + future.result() # Already done: never waits on Store I/O. + except (EmbeddingStoreOperationError, OSError) as error: + logger.warning("Store PUT failed for %s: %s", identifier, error) + # This failure is consumed; do not retain publisher/caller frames. + error.__traceback__ = None + except BaseException as error: + # The client may still own unsafe native-I/O buffers. Keep + # their charge and fatal result until worker termination. + if self._fatal_error is None: + self._fatal_error = error + continue + del self._pending[identifier] + storage, ref_count = self._pending_storages[save.storage_ptr] + if ref_count == 1: + del self._pending_storages[save.storage_ptr] + self._pending_bytes -= storage.nbytes() + else: + self._pending_storages[save.storage_ptr] = (storage, ref_count - 1) + fatal = self._fatal_error + if fatal is not None: + raise fatal + + def shutdown(self) -> None: + with self._pending_lock: + if self._closed: + return + self._closed = True + self._executor.shutdown(wait=True, cancel_futures=False) + self.reap() + self.store_client.close() + + +def _record_tensor_ready_event(tensor: torch.Tensor) -> torch.Event | None: + if tensor.device.type != "cuda": + return None + event = torch.Event() + event.record(torch.accelerator.current_stream(tensor.device)) + return event diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py new file mode 100644 index 000000000000..9a7d7dd002ce --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Content identities and output contracts for the shared Encoder Store.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +EMBEDDING_CACHE_KEY_VERSION = "v3" + + +def make_embedding_key(namespace: str, identifier: str) -> str: + return f"{namespace}@id:{identifier}" + + +@dataclass(frozen=True) +class TensorSpec: + """Contiguous output contract used by the Encoder and Store codec.""" + + shape: tuple[int, ...] + dtype: str + nbytes: int + + +def build_embedding_namespace(vllm_config: VllmConfig) -> str: + model_config = vllm_config.model_config + ec_config = vllm_config.ec_transfer_config + assert ec_config is not None + extra_config = ec_config.ec_connector_extra_config + + multimodal_config = model_config.multimodal_config + encoder_hash = ( + multimodal_config.compute_hash() + if multimodal_config is not None + else "encoder:default" + ) + cache_prefix = str(extra_config.get("embedding_cache_prefix", "")) + model_name = str(extra_config.get("embedding_model_identity", model_config.model)) + prefix = f"{cache_prefix}@" if cache_prefix else "" + return ( + f"{prefix}embedding" + f"@model:{model_name}" + f"@revision:{model_config.revision or 'default'}" + f"@encoder:{encoder_hash}" + f"@dtype:{model_config.dtype}" + f"@protocol:{EMBEDDING_CACHE_KEY_VERSION}" + ) diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py new file mode 100644 index 000000000000..4d3ed6739d0e --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Thin Mooncake Store client for embedding objects. + +Originally adapted from vLLM PR #47302 (0d1f71f5d36c). The v3 key namespace +uses a compact header for complete 2D embeddings, not that PR's tensor format. +""" + +from __future__ import annotations + +import ctypes +import struct +import threading +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from enum import IntEnum +from typing import Any + +import torch + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store +from vllm.logger import init_logger +from vllm.utils.network_utils import get_ip + +logger = init_logger(__name__) + +# A failed transfer can still use registered memory after the binding returns. +# Poisoned clients retain their owners until process exit, including if normal +# worker teardown drops the connector. New operations are rejected. +_UNSAFE_CLIENTS: list[Any] = [] + + +# Values from Mooncake v0.3.12 / v0.3.12.post1 types.h (ErrorCode). +# Python exposes integer returns, not this enum. +class _MooncakeErrorCode(IntEnum): + NO_AVAILABLE_HANDLE = -200 + INVALID_PARAMS = -600 + INVALID_REPLICA = -702 + REPLICA_IS_NOT_READY = -703 + OBJECT_NOT_FOUND = -704 + OBJECT_ALREADY_EXISTS = -705 + LEASE_EXPIRED = -707 + RPC_FAIL = -900 + RPC_TIMEOUT = -901 + TENANT_QUOTA_EXCEEDED = -1700 + + +# For RAM batch reads and single-replica writes, these rejections occur +# before submission or after completed I/O. Transfer timeouts can leave I/O +# in flight; all other failures retain registered buffers and poison the client. +_SAFE_IO_REJECTIONS = frozenset( + { + _MooncakeErrorCode.NO_AVAILABLE_HANDLE, + _MooncakeErrorCode.REPLICA_IS_NOT_READY, + _MooncakeErrorCode.OBJECT_NOT_FOUND, + _MooncakeErrorCode.OBJECT_ALREADY_EXISTS, + _MooncakeErrorCode.RPC_FAIL, + _MooncakeErrorCode.RPC_TIMEOUT, + } +) +# Batch GET validates capacity/replicas before I/O and checks leases afterwards. +# PUT rejects tenant quotas before allocating replicas or submitting transfers. +_SAFE_GET_REJECTIONS = _SAFE_IO_REJECTIONS | { + _MooncakeErrorCode.INVALID_PARAMS, + _MooncakeErrorCode.INVALID_REPLICA, + _MooncakeErrorCode.LEASE_EXPIRED, +} +_SAFE_PUT_REJECTIONS = _SAFE_IO_REJECTIONS | {_MooncakeErrorCode.TENANT_QUOTA_EXCEEDED} + + +_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E +# Magic, dtype, token count, hidden dimension. Format versioning is in the key. +_MOONCAKE_TENSOR_HEADER = struct.Struct(" MooncakeEmbeddingStoreClient: + try: + from mooncake.store import ( # type: ignore + MooncakeDistributedStore, + ObjectDataType, + ReplicateConfig, + ) + except ImportError as e: + raise ImportError( + "Install mooncake with Store support to enable cross_encoder_cache." + ) from e + + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils + + config = MooncakeStoreConfig.load_from_config() + if config.enable_offload: + raise ValueError( + "cross_encoder_cache supports a RAM Store; disable enable_offload" + ) + store = MooncakeDistributedStore() + local_ip = get_ip() + local_hostname = rdma_utils.get_requester_local_hostname(local_ip) + setup_mooncake_store(store, config, local_hostname) + + logger.info( + "Initialized embedding Mooncake store mode=%s global_segment_size=%d " + "local_buffer_size=%d", + config.mode, + config.global_segment_size, + config.local_buffer_size, + ) + replicate_config = ReplicateConfig() + replicate_config.data_type = ObjectDataType.TENSOR + return MooncakeEmbeddingStoreClient( + store, replicate_config=replicate_config, read_buffer_bytes=read_buffer_bytes + ) + + +class EmbeddingStoreError(RuntimeError): + """A Store failure that must propagate to the worker.""" + + +class EmbeddingStoreOperationError(EmbeddingStoreError): + """A rejected or completed operation that permits normal Encoder fallback.""" + + +def _batch_get_completed_safely(results: Any, capacities: list[int]) -> bool: + return ( + isinstance(results, list) + and len(results) == len(capacities) + and all( + type(result) is int + and (0 <= result <= capacity or result in _SAFE_GET_REJECTIONS) + for result, capacity in zip(results, capacities, strict=True) + ) + ) + + +class MooncakeEmbeddingStoreClient: + """Wraps Mooncake object and buffer APIs used by embedding transfer.""" + + def __init__( + self, + store: Any, + replicate_config: Any | None = None, + *, + read_buffer_bytes: int = 128 * 1024**2, + ): + self._lifetime_lock = threading.Lock() + self._unsafe_owners: list[Any] = [] + self._poisoned = False + self.store = store + self.replicate_config = replicate_config + self._read_buffer_bytes = read_buffer_bytes + self._read_buffer: torch.Tensor | None = None + self._put_header: ctypes.Array[ctypes.c_char] | None = None + + def _check_healthy(self) -> None: + if self._poisoned: + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; worker must exit" + ) + + def _poison(self, owners: list[Any]) -> None: + with self._lifetime_lock: + self._unsafe_owners.extend(owners) + if not self._poisoned: + self._poisoned = True + _UNSAFE_CLIENTS.append(self) + + @contextmanager + def _registered_io( + self, + buffers: list[tuple[Any, int, int]], + ) -> Iterator[Callable[[int], None]]: + self._check_healthy() + registered = [] + submitted = False + confirmed = False + + def confirm(result: int) -> None: + nonlocal confirmed + if type(result) is not int or ( + result < 0 and result not in _SAFE_PUT_REJECTIONS + ): + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining buffers" + ) + confirmed = True + + try: + for owner, addr, size in buffers: + ret = self.store.register_buffer(addr, size) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding buffer: {ret}" + ) + registered.append(addr) + submitted = True + yield confirm + except BaseException as error: + if submitted and not confirmed: + self._poison([owner for owner, _, _ in buffers]) + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining buffers" + ) from error + raise + finally: + if not submitted or confirmed: + try: + for addr in reversed(registered): + self._unregister_buffer(addr) + except BaseException: + self._poison([owner for owner, _, _ in buffers]) + raise + + def close(self) -> None: + """Close the supported Mooncake binding after all I/O is drained.""" + self._check_healthy() + if self._read_buffer is not None: + try: + self._unregister_buffer(self._read_buffer.data_ptr()) + except BaseException: + self._poison([self._read_buffer]) + raise + self._read_buffer = None + if self._put_header is not None: + try: + self._unregister_buffer(ctypes.addressof(self._put_header)) + except BaseException: + self._poison([self._put_header]) + raise + self._put_header = None + ret = self.store.close() + if ret != 0: + raise EmbeddingStoreError( + f"failed to close embedding Mooncake Store: {ret}" + ) + + def exists(self, pool_key: str) -> bool: + return self.batch_exists([pool_key])[0] + + def batch_exists(self, keys: list[str]) -> list[bool]: + self._check_healthy() + if not keys: + return [] + + states = self.store.batch_is_exist(keys) + if len(states) != len(keys): + raise RuntimeError( + "Mooncake Store returned an unexpected number of lookup results" + ) + return [state == 1 for state in states] + + def load_tensors( + self, expected: Mapping[str, TensorSpec], device: torch.device | str + ) -> dict[str, torch.Tensor]: + """Synchronously load runner inputs into independently owned tensors.""" + self._check_healthy() + device = torch.device(device) + loaded: dict[str, torch.Tensor] = {} + chunk: dict[str, TensorSpec] = {} + chunk_bytes = 0 + for key, spec in expected.items(): + size = _MOONCAKE_TENSOR_HEADER.size + spec.nbytes + if size > self._read_buffer_bytes: + logger.warning( + "Skipping Store GET for %s: object bytes=%d exceed read buffer=%d", + key, + size, + self._read_buffer_bytes, + ) + continue + if chunk_bytes + size > self._read_buffer_bytes: + loaded.update(self._load_chunk(chunk, device)) + chunk = {} + chunk_bytes = 0 + chunk[key] = spec + chunk_bytes += size + if chunk: + loaded.update(self._load_chunk(chunk, device)) + return loaded + + def _load_chunk( + self, expected: dict[str, TensorSpec], device: torch.device + ) -> dict[str, torch.Tensor]: + self._check_healthy() + if self._read_buffer is None: + buffer = torch.empty( + self._read_buffer_bytes, + dtype=torch.uint8, + device="cpu", + pin_memory=device.type == "cuda", + ) + ret = self.store.register_buffer(buffer.data_ptr(), buffer.nbytes) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding read buffer: {ret}" + ) + self._read_buffer = buffer + buffer = self._read_buffer + base = buffer.data_ptr() + sizes = [ + _MOONCAKE_TENSOR_HEADER.size + spec.nbytes for spec in expected.values() + ] + offsets: list[int] = [] + end = 0 + for size in sizes: + offsets.append(end) + end += size + try: + results = self.store.batch_get_into( + list(expected), [base + offset for offset in offsets], sizes + ) + if not _batch_get_completed_safely(results, sizes): + raise EmbeddingStoreError("Unsafe batch GET results") + except BaseException as error: + self._poison([buffer]) + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining read buffer" + ) from error + + valid: list[tuple[str, TensorSpec, int]] = [] + for (key, spec), offset, size, result in zip( + expected.items(), offsets, sizes, results, strict=True + ): + if result != size: + logger.warning( + "Store GET skipped for %s: result=%d, expected bytes=%d", + key, + result, + size, + ) + continue + try: + _validate_mooncake_tensor_metadata( + key, + ctypes.string_at(base + offset, _MOONCAKE_TENSOR_HEADER.size), + spec, + ) + except EmbeddingStoreOperationError as error: + logger.warning("Store GET skipped for %s: %s", key, error) + continue + valid.append((key, spec, offset + _MOONCAKE_TENSOR_HEADER.size)) + if not valid: + return {} + # Allocate every destination before submitting any asynchronous copy. + targets = { + key: torch.empty( + spec.shape, + dtype=getattr(torch, spec.dtype.removeprefix("torch.")), + device=device, + ) + for key, spec, _ in valid + } + ready = torch.Event() if device.type == "cuda" else None + try: + for key, spec, offset in valid: + targets[key].view(torch.uint8).view(-1).copy_( + buffer[offset : offset + spec.nbytes], + non_blocking=ready is not None, + ) + if ready is not None: + ready.record(torch.accelerator.current_stream(device)) + ready.synchronize() + except BaseException as error: + if ready is not None: + self._poison([buffer, *targets.values(), ready]) + raise EmbeddingStoreError( + "Embedding copy completion is unconfirmed; retaining buffers" + ) from error + raise + return targets + + def put_tensor(self, pool_key: str, tensor: torch.Tensor) -> None: + self._check_healthy() + if not tensor.is_contiguous(): + raise EmbeddingStoreOperationError("embedding tensor must be contiguous") + metadata = _encode_mooncake_tensor_metadata(tensor) + # The single publisher reuses this header only after PUT completion. + if self._put_header is None: + header = ctypes.create_string_buffer(len(metadata)) + ret = self.store.register_buffer(ctypes.addressof(header), len(metadata)) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding header: {ret}" + ) + self._put_header = header + header_ptr = ctypes.addressof(self._put_header) + ctypes.memmove(header_ptr, metadata, len(metadata)) + data_size = tensor.numel() * tensor.element_size() + buffers = [ + (tensor, tensor.data_ptr(), data_size), + ] + with self._registered_io(buffers) as confirm: + [result] = self.store.batch_put_from_multi_buffers( + [pool_key], + [[header_ptr, tensor.data_ptr()]], + [[len(metadata), data_size]], + self.replicate_config, + ) + confirm(result) + if result < 0: + raise EmbeddingStoreOperationError( + f"failed to put embedding tensor for {pool_key}" + ) + + def _unregister_buffer(self, addr: int) -> None: + try: + ret = self.store.unregister_buffer(addr) + except Exception as error: + raise EmbeddingStoreError( + f"could not confirm buffer unregistration addr={addr:#x}" + ) from error + if ret != 0: + raise EmbeddingStoreError( + f"failed to unregister embedding buffer addr={addr:#x}: {ret}" + ) + + +def _validate_mooncake_tensor_metadata( + pool_key: str, + metadata: bytes, + expected: TensorSpec, +) -> None: + magic, dtype, num_tokens, hidden_dim = _MOONCAKE_TENSOR_HEADER.unpack(metadata) + if magic != _MOONCAKE_TENSOR_OBJECT_MAGIC: + raise EmbeddingStoreOperationError( + f"invalid Mooncake tensor metadata header for {pool_key}" + ) + if _MOONCAKE_DTYPE_TO_TORCH_DTYPE.get(dtype) != expected.dtype: + raise EmbeddingStoreOperationError( + f"unexpected Mooncake tensor dtype for {pool_key}: {dtype}" + ) + if (num_tokens, hidden_dim) != expected.shape: + raise EmbeddingStoreOperationError( + f"embedding tensor shape mismatch for {pool_key}: expected={expected.shape}" + ) + + +def _encode_mooncake_tensor_metadata(tensor: torch.Tensor) -> bytes: + dtype = str(tensor.dtype) + if dtype not in _TORCH_DTYPE_TO_MOONCAKE_DTYPE: + raise EmbeddingStoreOperationError( + f"unsupported embedding tensor dtype: {dtype}" + ) + if tensor.ndim != 2: + raise EmbeddingStoreOperationError( + f"embedding tensor must be 2D: shape={tuple(tensor.shape)}" + ) + return _MOONCAKE_TENSOR_HEADER.pack( + _MOONCAKE_TENSOR_OBJECT_MAGIC, + _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype], + *tensor.shape, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 762749987b5d..6b6bf18ae48c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -11,7 +11,6 @@ """ import dataclasses -import json import math import os import queue @@ -20,10 +19,8 @@ import time from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from dataclasses import dataclass -from typing import Any, Literal, TypeVar +from typing import Any, TypeVar -import regex as re import torch import zmq @@ -67,6 +64,7 @@ decode_lookup_response, encode_lookup_response, ) +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket @@ -94,10 +92,6 @@ logger = init_logger(__name__) -DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB -DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB -DEFAULT_TENANT_ID = "default" - MOONCAKE_NO_AVAILABLE_HANDLE = -200 _T = TypeVar("_T") @@ -149,115 +143,6 @@ def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str: _DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT -MooncakeMode = Literal["embedded", "standalone-store"] - - -@dataclass -class MooncakeStoreConfig: - """Configuration for MooncakeDistributedStore. - - ``mode`` selects the topology: ``embedded`` (each rank contributes - ``global_segment_size`` in-process) or ``standalone-store`` (rank - contributes 0; an external ``mooncake_client`` process owns the pool - and the SSD tier). - """ - - metadata_server: str - master_server_address: str - protocol: str - device_name: str - mode: MooncakeMode = "embedded" - global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE - local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE - enable_offload: bool = False - tenant_id: str = DEFAULT_TENANT_ID - - def __post_init__(self) -> None: - if self.mode not in ("embedded", "standalone-store"): - raise ValueError(f"unknown Mooncake mode: {self.mode!r}") - if self.local_buffer_size <= 0: - raise ValueError("local_buffer_size must be > 0") - if self.mode == "embedded" and self.global_segment_size == 0: - raise ValueError("embedded mode requires global_segment_size > 0") - if self.mode == "standalone-store" and self.global_segment_size != 0: - raise ValueError("standalone-store mode requires global_segment_size == 0") - - @staticmethod - def from_file(file_path: str) -> "MooncakeStoreConfig": - with open(file_path) as file: - config = json.load(file) - return MooncakeStoreConfig( - metadata_server=config.get("metadata_server", ""), - master_server_address=config.get("master_server_address", ""), - protocol=config.get("protocol", "rdma"), - device_name=config.get("device_name", ""), - mode=config.get("mode", "embedded"), - global_segment_size=_parse_size( - config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) - ), - local_buffer_size=_parse_size( - config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) - ), - enable_offload=bool(config.get("enable_offload", False)), - tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)), - ) - - @staticmethod - def load_from_config() -> "MooncakeStoreConfig": - config_path = os.getenv("MOONCAKE_CONFIG_PATH") - if not config_path: - raise ValueError( - "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." - ) - return MooncakeStoreConfig.from_file(config_path) - - -def _normalize_tenant_id(value: Any) -> str: - if value is None: - return DEFAULT_TENANT_ID - if not isinstance(value, str): - raise TypeError( - f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}" - ) - tenant_id = value.strip() - return tenant_id if tenant_id else DEFAULT_TENANT_ID - - -def _parse_size(value: Any) -> int: - """Parse storage size strings with units: GB, MB, KB, B.""" - if isinstance(value, int): - return value - if not isinstance(value, str): - try: - return int(value) - except (TypeError, ValueError) as e: - raise TypeError(f"Unsupported type for size: {type(value)}") from e - - cleaned = value.strip().lower() - if not cleaned: - raise ValueError("Size cannot be empty.") - - unit_multipliers = { - "gb": 1024**3, - "mb": 1024**2, - "kb": 1024, - "b": 1, - } - match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) - if not match: - raise ValueError(f"Invalid format: '{value}'") - - number_str = match.group(1) - unit = match.group(2) or "b" - multiplier = unit_multipliers[unit] - - try: - numeric_value = float(number_str) - except ValueError as exc: - raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc - return int(numeric_value * multiplier) - - def _align_up(value: int, alignment: int) -> int: return ((value + alignment - 1) // alignment) * alignment @@ -1595,23 +1480,7 @@ def __init__( self.store = MooncakeDistributedStore() local_ip = get_ip() local_hostname = rdma_utils.get_requester_local_hostname(local_ip) - setup_kwargs: dict[str, str] = {} - if store_config.tenant_id != DEFAULT_TENANT_ID: - setup_kwargs["tenant_id"] = store_config.tenant_id - ret = self.store.setup( - local_hostname, - store_config.metadata_server, - store_config.global_segment_size, - store_config.local_buffer_size, - store_config.protocol, - store_config.device_name, - store_config.master_server_address, - **setup_kwargs, - ) - if ret != 0: - msg = "Initialize MooncakeDistributedStore failed." - logger.error(msg) - raise RuntimeError(msg) + setup_mooncake_store(self.store, store_config, local_hostname) preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config) self.preferred_segment = preferred_segment diff --git a/vllm/distributed/mooncake_store.py b/vllm/distributed/mooncake_store.py new file mode 100644 index 000000000000..f76bf0276618 --- /dev/null +++ b/vllm/distributed/mooncake_store.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared configuration and setup for EC and KV Mooncake Store clients.""" + +import json +import os +from dataclasses import dataclass +from typing import Any, Literal + +import regex as re + +DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB +DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB +DEFAULT_TENANT_ID = "default" + + +MooncakeMode = Literal["embedded", "standalone-store"] + + +@dataclass +class MooncakeStoreConfig: + """Configuration for MooncakeDistributedStore. + + ``mode`` selects the topology: ``embedded`` (each rank contributes + ``global_segment_size`` in-process) or ``standalone-store`` (rank + contributes 0; an external ``mooncake_client`` process owns the pool + and the SSD tier). + """ + + metadata_server: str + master_server_address: str + protocol: str + device_name: str + mode: MooncakeMode = "embedded" + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE + enable_offload: bool = False + tenant_id: str = DEFAULT_TENANT_ID + + def __post_init__(self) -> None: + if self.mode not in ("embedded", "standalone-store"): + raise ValueError(f"unknown Mooncake mode: {self.mode!r}") + if self.local_buffer_size <= 0: + raise ValueError("local_buffer_size must be > 0") + if self.mode == "embedded" and self.global_segment_size == 0: + raise ValueError("embedded mode requires global_segment_size > 0") + if self.mode == "standalone-store" and self.global_segment_size != 0: + raise ValueError("standalone-store mode requires global_segment_size == 0") + + @staticmethod + def from_file(file_path: str) -> "MooncakeStoreConfig": + with open(file_path) as file: + config = json.load(file) + return MooncakeStoreConfig( + metadata_server=config.get("metadata_server", ""), + master_server_address=config.get("master_server_address", ""), + protocol=config.get("protocol", "rdma"), + device_name=config.get("device_name", ""), + mode=config.get("mode", "embedded"), + global_segment_size=_parse_size( + config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) + ), + local_buffer_size=_parse_size( + config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) + ), + enable_offload=bool(config.get("enable_offload", False)), + tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)), + ) + + @staticmethod + def load_from_config() -> "MooncakeStoreConfig": + config_path = os.getenv("MOONCAKE_CONFIG_PATH") + if not config_path: + raise ValueError( + "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." + ) + return MooncakeStoreConfig.from_file(config_path) + + +def _normalize_tenant_id(value: Any) -> str: + if value is None: + return DEFAULT_TENANT_ID + if not isinstance(value, str): + raise TypeError( + f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}" + ) + tenant_id = value.strip() + return tenant_id if tenant_id else DEFAULT_TENANT_ID + + +def _parse_size(value: Any) -> int: + """Parse storage size strings with units: GB, MB, KB, B.""" + if isinstance(value, int): + return value + if not isinstance(value, str): + try: + return int(value) + except (TypeError, ValueError) as e: + raise TypeError(f"Unsupported type for size: {type(value)}") from e + + cleaned = value.strip().lower() + if not cleaned: + raise ValueError("Size cannot be empty.") + + unit_multipliers = { + "gb": 1024**3, + "mb": 1024**2, + "kb": 1024, + "b": 1, + } + match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) + if not match: + raise ValueError(f"Invalid format: '{value}'") + + number_str = match.group(1) + unit = match.group(2) or "b" + multiplier = unit_multipliers[unit] + + try: + numeric_value = float(number_str) + except ValueError as exc: + raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc + return int(numeric_value * multiplier) + + +def setup_mooncake_store( + store: Any, config: MooncakeStoreConfig, local_hostname: str +) -> None: + """Initialize a native Store, preserving explicit tenant configuration.""" + setup_kwargs: dict[str, str] = {} + if config.tenant_id != DEFAULT_TENANT_ID: + setup_kwargs["tenant_id"] = config.tenant_id + ret = store.setup( + local_hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name, + config.master_server_address, + **setup_kwargs, + ) + if ret != 0: + raise RuntimeError("Initialize MooncakeDistributedStore failed.") From e6b1a5e5e3fce6b777c5f63398f87c7c911a40b3 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Wed, 16 Sep 2026 22:19:13 -0700 Subject: [PATCH 376/774] [XPU] Enable XPU eplb (#44987) Signed-off-by: mayuyuace Signed-off-by: Qiming Zhang --- tests/conftest.py | 2 +- tests/distributed/eplb_utils.py | 8 +- tests/distributed/test_elastic_ep.py | 4 + tests/distributed/test_eplb_algo.py | 5 +- tests/distributed/test_eplb_events.py | 29 +++-- tests/distributed/test_eplb_execute.py | 107 +++++++++++++----- .../distributed/test_eplb_fused_moe_layer.py | 9 +- .../test_eplb_fused_moe_layer_dep_nvfp4.py | 3 + vllm/config/parallel.py | 17 ++- vllm/distributed/eplb/async_worker.py | 28 ++--- vllm/distributed/eplb/eplb_communicator.py | 95 +++++++++++++--- vllm/distributed/eplb/eplb_state.py | 22 ++-- vllm/distributed/eplb/eplb_utils.py | 25 +++- vllm/distributed/eplb/rebalance_execute.py | 19 ++-- .../layers/fused_moe/experts/xpu_moe.py | 1 + .../layers/fused_moe/router/base_router.py | 2 +- 16 files changed, 267 insertions(+), 109 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 5fd594a04a2e..a52a38f35775 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -935,7 +935,7 @@ def hf_runner(): def _default_block_size() -> int: - if torch.xpu.is_available(): + if current_platform.is_xpu(): return 64 if current_platform.is_cpu(): return 128 diff --git a/tests/distributed/eplb_utils.py b/tests/distributed/eplb_utils.py index e6a41842e19d..d868069e0a75 100644 --- a/tests/distributed/eplb_utils.py +++ b/tests/distributed/eplb_utils.py @@ -13,6 +13,7 @@ from vllm.distributed.parallel_state import ( init_distributed_environment, ) +from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables mp.set_start_method("spawn", force=True) @@ -73,13 +74,16 @@ def distributed_run(fn, world_size, *args): def set_env_vars_and_device(env: dict[str, str]) -> None: update_environment_variables(env) local_rank = os.environ["LOCAL_RANK"] - device = torch.device(f"cuda:{local_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{local_rank}") torch.accelerator.set_device_index(device) + backend = "xccl" if current_platform.is_xpu() else "nccl" + # Create a minimal vllm config for init_distributed_environment vllm_config = VllmConfig() with set_current_vllm_config(vllm_config): - init_distributed_environment() + init_distributed_environment(backend=backend) atexit.register(_destroy_process_group_if_initialized) # Ensure each worker process has the same random seed random.seed(42) diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 8ccea8628389..1dc6b5d3bbdd 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -10,6 +10,8 @@ import pytest import requests +from vllm.platforms import current_platform + from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k from ..utils import RemoteOpenAIServer, multi_gpu_test @@ -233,6 +235,8 @@ def test_elastic_ep_scaling( if not has_nixl(): pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if current_platform.is_xpu(): + pytest.skip("Elastic EP with NIXL is not supported on XPU") if all2all_backend == "nixl_ep" and not has_nixl_ep(): pytest.skip("NIXL EP is not installed") diff --git a/tests/distributed/test_eplb_algo.py b/tests/distributed/test_eplb_algo.py index aaa987b9fc42..8cd5cf9e3c05 100644 --- a/tests/distributed/test_eplb_algo.py +++ b/tests/distributed/test_eplb_algo.py @@ -10,6 +10,7 @@ compute_logical_maps, ) from vllm.distributed.eplb.policy.default import DefaultEplbPolicy +from vllm.platforms import current_platform def test_eplb_load_stats_reduce_across_ranks(): @@ -263,11 +264,13 @@ def test_global_load_balance_fallback(): assert torch.sum(logcnt) == num_replicas -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cpu", "cuda", "xpu"]) def test_device_compatibility(device): """Test device compatibility.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") + if device == "xpu" and not current_platform.is_xpu(): + pytest.skip("XPU not available") weight = torch.tensor([[10, 20, 30, 40]], device=device) num_replicas = 6 diff --git a/tests/distributed/test_eplb_events.py b/tests/distributed/test_eplb_events.py index 963e3cbb77f8..f924b107d354 100644 --- a/tests/distributed/test_eplb_events.py +++ b/tests/distributed/test_eplb_events.py @@ -4,15 +4,17 @@ import threading import time +import pytest import torch -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream +from vllm.platforms import current_platform def test_wait_blocks_until_record(): event = CpuGpuEvent() - record_stream = torch.cuda.Stream() - wait_stream = torch.cuda.Stream() + record_stream = torch.Stream() + wait_stream = torch.Stream() wait_returned = threading.Event() def waiter(): @@ -33,8 +35,8 @@ def waiter(): def test_reuse_across_multiple_cycles(): wrapper = CpuGpuEvent() - record_stream = torch.cuda.Stream() - wait_stream = torch.cuda.Stream() + record_stream = torch.Stream() + wait_stream = torch.Stream() NUM_CYCLES = 8 completed_cycles = [] barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)] @@ -56,13 +58,18 @@ def waiter(): assert len(completed_cycles) == NUM_CYCLES -def test_producer_consumer(): - """This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared - GPU tensor on multiple CPU threads. +@pytest.mark.parametrize("device", ["cuda", "xpu"]) +def test_producer_consumer(device): + """This test uses the CpuGpuEvent to synchronize reads and writes to/from + a shared GPU tensor on multiple CPU threads. """ - worker_stream = torch.cuda.Stream() + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + if device == "xpu" and not current_platform.is_xpu(): + pytest.skip("XPU not available") + worker_stream = torch.Stream() # Create a single element counter that will be shared between two threads - buf = torch.zeros(1, device="cuda") + buf = torch.zeros(1, device=device) NUM_ROUNDS = 5 ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)] @@ -76,7 +83,7 @@ def worker(): if i > 0: events[i - 1].wait(stream=worker_stream) - with torch.cuda.stream(worker_stream): + with device_stream(worker_stream): buf.fill_(float(i + 1)) worker_stream.synchronize() diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index aed0decc6d35..2037535e3ddf 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -23,6 +23,7 @@ ensure_model_parallel_initialized, get_tp_group, ) +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -208,10 +209,14 @@ def verify_redundant_experts_have_same_weights( ep_rank: int, world_size: int, num_local_experts: int, + cpu_group, ) -> bool: """Verify that all replicas of the same logical expert have the same weights.""" num_layers = len(expert_weights) total_physical_experts = world_size * num_local_experts + # XPU: gloo does not support XPU tensors, so gather on CPU then move back. + # CUDA: gloo supports CUDA tensors directly, keep original device. + use_cpu_gather = current_platform.is_xpu() ok = True for layer in range(num_layers): @@ -219,31 +224,35 @@ def verify_redundant_experts_have_same_weights( all_weights: list[torch.Tensor] = [] for weight_idx, hidden_size in enumerate(hidden_sizes): + orig_device = expert_weights[layer][weight_idx].device + + if use_cpu_gather: + local_weights = expert_weights[layer][weight_idx].cpu() + gather_device = "cpu" + else: + local_weights = expert_weights[layer][weight_idx] + gather_device = orig_device + # Create tensor to store all expert weights # Shape: [total_physical_experts, hidden_size] gathered_weights = torch.zeros( total_physical_experts, hidden_size, - device=expert_weights[layer][weight_idx].device, - dtype=expert_weights[layer][weight_idx].dtype, + device=gather_device, + dtype=local_weights.dtype, ) - # Use all_gather to collect expert weights from current node - # expert_weights[layer][weight_idx] shape: - # [num_local_experts, hidden_size] - local_weights = expert_weights[layer][ - weight_idx - ] # [num_local_experts, hidden_size] - # Split tensor along dim 0 into a list for all_gather gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0) torch.distributed.all_gather( - # Output list: each element corresponds to one rank's weights list(gathered_weights_list), - local_weights, # Input: current rank's local weights + local_weights, + group=cpu_group, ) + if use_cpu_gather: + gathered_weights = gathered_weights.to(orig_device) all_weights.append(gathered_weights) # Verify that all replicas of the same logical expert have the same @@ -294,9 +303,23 @@ def verify_redundant_experts_have_same_weights( return ok -def assert_verification_synced(local_ok: bool, msg: str) -> None: - ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32) - torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN) +def assert_verification_synced(local_ok: bool, msg: str, cpu_group) -> None: + if current_platform.is_xpu(): + # XPU: default group backend (gloo) does not support XPU tensors; + # use CPU tensor with cpu_group instead. + ok_tensor = torch.tensor( + [1 if local_ok else 0], device="cpu", dtype=torch.int32 + ) + torch.distributed.all_reduce( + ok_tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group + ) + else: + # CUDA: nccl (default group) supports device tensors directly. + device = torch.accelerator.current_accelerator() + ok_tensor = torch.tensor( + [1 if local_ok else 0], device=device, dtype=torch.int32 + ) + torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN) assert bool(ok_tensor.item()), msg @@ -336,8 +359,10 @@ def _test_async_transfer_layer_without_mtp_worker( ep_group_coordinator = get_tp_group() ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") total_physical_experts = world_size * num_local_experts hidden_sizes = [16, 32] @@ -376,7 +401,7 @@ def _test_async_transfer_layer_without_mtp_worker( new_indices_cpu = new_indices.cpu() expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] - cuda_stream = torch.cuda.Stream(device=device) + stream = torch.Stream(device=device) communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, @@ -384,7 +409,7 @@ def _test_async_transfer_layer_without_mtp_worker( expert_weights=expert_weights, expert_buffer=expert_buffer, ) - communicator.set_stream(cuda_stream) + communicator.set_stream(stream) for layer_idx in range(num_layers): transfer_metadata = transfer_layer( @@ -394,10 +419,10 @@ def _test_async_transfer_layer_without_mtp_worker( expert_weights_buffer=expert_buffer, ep_group=ep_group, communicator=communicator, - cuda_stream=cuda_stream, + stream=stream, layer_idx=layer_idx, ) - cuda_stream.synchronize() + stream.synchronize() move_from_buffer( expert_weights=expert_weights[layer_idx], expert_weights_buffers=expert_buffer, @@ -421,6 +446,7 @@ def _test_async_transfer_layer_without_mtp_worker( ep_rank, world_size, num_local_experts, + cpu_group=cpu_group, ) and local_ok ) @@ -428,6 +454,7 @@ def _test_async_transfer_layer_without_mtp_worker( local_ok, "Async transfer verification failed on at least one rank. " "See logs for details.", + cpu_group=cpu_group, ) @@ -452,9 +479,11 @@ def _test_rearrange_expert_weights_with_redundancy( ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") # Test parameters total_physical_experts = world_size * num_local_experts @@ -523,12 +552,14 @@ def _test_rearrange_expert_weights_with_redundancy( ep_rank, world_size, num_local_experts, + cpu_group=cpu_group, ) and local_ok ) assert_verification_synced( local_ok, "Rearrange verification failed on at least one rank. See logs for details.", + cpu_group=cpu_group, ) @@ -556,7 +587,7 @@ def _test_rearrange_expert_weights_with_redundancy( ], ) @pytest.mark.parametrize( - "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] + "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"] ) def test_rearrange_expert_weights_with_redundancy( world_size, @@ -568,6 +599,12 @@ def test_rearrange_expert_weights_with_redundancy( """Test the functionality of rearranging expert weights with redundancy.""" if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") + if eplb_communicator == "nixl" and current_platform.is_xpu(): + pytest.skip("NIXL does not support XPU") + if eplb_communicator in ("torch_nccl", "pynccl") and not torch.cuda.is_available(): + pytest.skip(f"{eplb_communicator} requires CUDA") + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): + pytest.skip("torch_xccl requires XPU") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -592,9 +629,11 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") num_layers = 2 num_local_experts = 2 @@ -623,9 +662,10 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: original_weights.append(layer_copy) expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, - backend="torch_nccl", + backend=default_backend, expert_weights=expert_weights, expert_buffer=expert_buffer, ) @@ -657,6 +697,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: assert_verification_synced( local_ok, "No-change EPLB verification failed on at least one rank.", + cpu_group=cpu_group, ) @@ -666,7 +707,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: (2, 2, 2, 3), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"]) +@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"]) def test_async_transfer_layer_without_mtp( world_size: int, num_layers: int, @@ -677,6 +718,10 @@ def test_async_transfer_layer_without_mtp( """Exercise async EPLB transfer path without MTP/spec decode.""" if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") + if eplb_communicator == "nixl" and current_platform.is_xpu(): + pytest.skip("NIXL does not support XPU") + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): + pytest.skip("torch_xccl requires XPU") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") @@ -715,9 +760,11 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") num_layers = 1 num_local_experts = 2 @@ -753,9 +800,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: original_weights.append(layer_copy) expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, - backend="torch_nccl", + backend=default_backend, expert_weights=expert_weights, expert_buffer=expert_buffer, ) @@ -788,6 +836,7 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: assert_verification_synced( local_ok, "Profile-mode EPLB verification failed on at least one rank.", + cpu_group=cpu_group, ) diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py index d8d870bc98c8..c6ae9e60d6e5 100644 --- a/tests/distributed/test_eplb_fused_moe_layer.py +++ b/tests/distributed/test_eplb_fused_moe_layer.py @@ -17,6 +17,7 @@ get_tp_group, ) from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -81,7 +82,7 @@ def make_fused_moe_layer( ) re = fml.routed_experts - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{torch.accelerator.current_accelerator().type}:{rank}") from functools import partial @@ -190,13 +191,14 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): vllm_config = VllmConfig() vllm_config.parallel_config.tensor_parallel_size = world_size vllm_config.parallel_config.enable_expert_parallel = True + vllm_config.parallel_config.enable_eplb = True with set_current_vllm_config(vllm_config): ensure_model_parallel_initialized( tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 ) - ep_group = get_tp_group().cpu_group + ep_group = get_tp_group().device_group ep_rank = torch.distributed.get_rank() fml_layers = [ @@ -216,9 +218,10 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): shuffled_indices[lidx] = torch.randperm(test_config.num_experts) expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), - backend="torch_nccl", + backend=default_backend, expert_weights=rank_expert_weights, expert_buffer=expert_buffer, ) diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index fe11e13cd60b..979742d5014e 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -24,6 +24,7 @@ ModelOptNvFp4Config, ModelOptNvFp4FusedMoE, ) +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -270,6 +271,8 @@ def test_eplb_fml( ): if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") + if current_platform.is_xpu(): + pytest.skip("NVFP4 quantization and flashinfer backends are CUDA-only") num_local_experts = num_experts // world_size num_topk = 4 diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f0a2b7e45785..ba744e049fbc 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -38,7 +38,9 @@ DataParallelBackend = Literal["ray", "mp"] EPLBPolicyOption = Literal["default"] DCPCommBackend = Literal["ag_rs", "a2a"] -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"] +EPLBCommunicatorBackend = Literal[ + "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl" +] All2AllBackend = Literal[ "naive", "pplx", @@ -94,9 +96,11 @@ class EPLBConfig: Backend for EPLB expert weight communication: - "torch_nccl": Use torch.distributed on the device process group - "torch_gloo": Use torch.distributed gloo with CPU staging + - "torch_xccl": Use torch.distributed XCCL device P2P on XPU - "nixl": Use NIXL with staged send/recv buffers - "pynccl": Use PyNccl send/recv - - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo") + - None: Auto-select backend ("torch_xccl" on XPU, prefers "nixl" + on CUDA, falls back to "torch_gloo") """ @model_validator(mode="after") @@ -525,10 +529,10 @@ def _validate_parallel_config(self) -> Self: ) if self.enable_eplb: - if not current_platform.is_cuda_alike(): + if not current_platform.is_cuda_alike() and not current_platform.is_xpu(): raise ValueError( "Expert parallelism load balancing is only supported on " - "CUDA devices or ROCm devices now." + "CUDA devices or ROCm devices or XPU devices now." ) if not self.enable_expert_parallel: raise ValueError("enable_expert_parallel must be True to use EPLB.") @@ -1034,7 +1038,10 @@ def __post_init__(self) -> None: # See https://github.com/pytorch/pytorch/issues/174288 from vllm.distributed.nixl_utils import is_nixl_available - if is_nixl_available(): + if current_platform.is_xpu(): + # On XPU, use the device-native XCCL P2P backend. + self.eplb_config.communicator = "torch_xccl" + elif is_nixl_available(): self.eplb_config.communicator = "nixl" elif self.enable_elastic_ep: self.eplb_config.communicator = "pynccl" diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index 5cf63fe5aecd..3b7bc908eca7 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -10,7 +10,7 @@ from vllm.distributed.parallel_state import get_eplb_group from vllm.logger import init_logger -from .eplb_utils import CpuGpuEvent +from .eplb_utils import CpuGpuEvent, device_stream from .rebalance_execute import AsyncEplbLayerResult, transfer_layer if TYPE_CHECKING: @@ -24,17 +24,17 @@ def start_async_worker( is_profile: bool = False, ) -> threading.Thread: rank = get_eplb_group().device_group.rank() - device_index = state.cuda_device_index + device_index = state.device_index assert state.is_async def thread_target() -> None: assert device_index is not None torch.accelerator.set_device_index(device_index) - cuda_stream = torch.cuda.Stream(device=device_index) + stream = torch.Stream(device=torch.accelerator.current_accelerator()) try: transfer_run_periodically( state=state, - cuda_stream=cuda_stream, + stream=stream, is_profile=is_profile, ) except Exception as exc: # pragma: no cover - diagnostic path @@ -49,13 +49,13 @@ def run_rebalance_experts( model_state: "EplbModelState", eplb_state: "EplbState", physical_to_logical_map_cpu: torch.Tensor, - cuda_stream: torch.cuda.Stream, + stream: torch.Stream, ) -> torch.Tensor: assert model_state.eplb_stats is not None eplb_stats = model_state.eplb_stats # Move the global expert load window to CPU for computation. - with torch.cuda.stream(cuda_stream): + with device_stream(stream): global_expert_load_window = eplb_stats.global_expert_load_window.cpu() # Compute new expert mappings for the model new_physical_to_logical_map = eplb_state.policy.rebalance_experts( @@ -73,11 +73,11 @@ def run_rebalance_experts( def transfer_run_periodically( state: "EplbState", - cuda_stream: torch.cuda.Stream, + stream: torch.Stream, is_profile: bool = False, ) -> None: while True: - state.rearrange_event.wait(stream=cuda_stream) + state.rearrange_event.wait(stream=stream) eplb_group = get_eplb_group().device_group eplb_cpu_group = get_eplb_group().cpu_group @@ -87,16 +87,16 @@ def transfer_run_periodically( for model_state in state.model_states.values(): layer_idx = 0 # Set the async worker's CUDA stream on the communicator - model_state.communicator.set_stream(cuda_stream) + model_state.communicator.set_stream(stream) num_layers = model_state.model.num_moe_layers # Snapshot the physical_to_logical_map (synchronized with # rearrange_event) and copy it to CPU - with torch.cuda.stream(cuda_stream): + with device_stream(stream): physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu() new_physical_to_logical_map = run_rebalance_experts( - model_state, state, physical_to_logical_map_cpu, cuda_stream + model_state, state, physical_to_logical_map_cpu, stream ) # Execute one EPLB layer transfer per model forward pass. Each iteration @@ -132,13 +132,13 @@ def transfer_run_periodically( communicator=model_state.communicator, ep_group=eplb_group, is_profile=is_profile, - cuda_stream=cuda_stream, + stream=stream, layer_idx=layer_idx, ) # Wait until all writes to expert_buffer have finished before making the # AsyncEplbLayerResult visible to the main thread. - cuda_stream.synchronize() + stream.synchronize() # This event guarantees that expert_buffer will not be overwritten by # subsequent iterations of this loop until the main thread has consumed @@ -155,6 +155,6 @@ def transfer_run_periodically( # Block this thread until the main thread and main stream # finish copying model_state.expert_buffer into # model_state.model.expert_weights[layer_idx] - consumed_event.wait(stream=cuda_stream) + consumed_event.wait(stream=stream) assert model_state.pending_result is None layer_idx += 1 diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 61fef468d323..730c01a3882e 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -22,6 +22,7 @@ from vllm.distributed.device_communicators.pynccl_wrapper import ( ncclDataTypeEnum, ) +from vllm.distributed.eplb.eplb_utils import device_stream from vllm.distributed.parallel_state import ( GroupCoordinator, get_pp_group, @@ -87,8 +88,8 @@ def needs_profile_buffer_reservation(self) -> bool: communication buffers.""" return True - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: - self._cuda_stream = cuda_stream + def set_stream(self, stream: torch.Stream | None) -> None: + self._stream = stream def _log_initialized(self) -> None: if is_local_first_rank(): @@ -101,10 +102,10 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator): def __init__( self, ep_group: ProcessGroup, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._ep_group = ep_group - self._cuda_stream = cuda_stream + self._stream = stream self._p2p_ops: list[P2POp] = [] self._log_initialized() @@ -144,7 +145,7 @@ def execute(self) -> None: if not self._p2p_ops: return try: - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): reqs = batch_isend_irecv(self._p2p_ops) for req in reqs: req.wait() @@ -158,10 +159,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): def __init__( self, cpu_group: ProcessGroup, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._cpu_group = cpu_group - self._cuda_stream = cuda_stream + self._stream = stream self._ops: list[tuple[str, torch.Tensor, int]] = [] self._log_initialized() @@ -217,7 +218,7 @@ def build_ops() -> None: recv_staging.append((tensor, cpu_tensor)) try: - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): build_ops() finally: self._ops.clear() @@ -225,10 +226,10 @@ def build_ops() -> None: # Wait for all D2H copies to finish # before issuing gloo batch_isend_irecv operations. with gpu_sync_allowed(): - if self._cuda_stream is not None: - self._cuda_stream.synchronize() + if self._stream is not None: + self._stream.synchronize() else: - torch.cuda.current_stream().synchronize() + torch.accelerator.current_stream().synchronize() reqs = batch_isend_irecv(p2p_ops) for req in reqs: @@ -236,11 +237,69 @@ def build_ops() -> None: if not recv_staging: return - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): for dst_tensor, cpu_tensor in recv_staging: dst_tensor.copy_(cpu_tensor, non_blocking=True) +class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator): + """EPLB communicator using XCCL device-to-device P2P on XPU.""" + + def __init__( + self, + ep_group: ProcessGroup, + stream: torch.Stream | None = None, + ) -> None: + self._ep_group = ep_group + self._stream = stream + self._ops: list[tuple[str, torch.Tensor, int]] = [] + self._log_initialized() + + def add_send( + self, + tensors: list[torch.Tensor], + dst_rank: int, + expert_id: int, # unused by this backend + ) -> None: + for tensor in tensors: + self._ops.append(("send", tensor, dst_rank)) + + def add_recv( + self, + tensors: list[torch.Tensor], + src_rank: int, + expert_id: int, # unused by this backend + ) -> None: + for tensor in tensors: + self._ops.append(("recv", tensor, src_rank)) + + def execute(self) -> None: + if not self._ops: + return + + p2p_ops: list[P2POp] = [] + try: + for op, tensor, peer_rank in self._ops: + send_or_recv = ( + torch.distributed.isend if op == "send" else torch.distributed.irecv + ) + p2p_ops.append( + P2POp( + send_or_recv, + tensor, + peer_rank, + self._ep_group, + ) + ) + finally: + self._ops.clear() + + with device_stream(self._stream): + reqs = batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + class NixlEplbCommunicator(EplbCommunicator): """EPLB communicator backed by NIXL READ transfers.""" @@ -343,7 +402,7 @@ def _make_agent_name(self) -> str: uid = uuid.uuid4().hex[:8] return f"eplb-{self._rank}{pp_suffix}-{uid}" - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: + def set_stream(self, stream: torch.Stream | None) -> None: pass def add_send( @@ -602,10 +661,10 @@ class PyNcclEplbCommunicator(EplbCommunicator): def __init__( self, pynccl_comm: PyNcclCommunicator, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._pynccl_comm = pynccl_comm - self._cuda_stream = cuda_stream + self._stream = stream self._group_started = False self._log_initialized() @@ -622,7 +681,7 @@ def add_send( ) -> None: self._ensure_group_started() for tensor in tensors: - self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream) + self._pynccl_comm.send(tensor, dst_rank, stream=self._stream) def add_recv( self, @@ -632,7 +691,7 @@ def add_recv( ) -> None: self._ensure_group_started() for tensor in tensors: - self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream) + self._pynccl_comm.recv(tensor, src_rank, stream=self._stream) def execute(self) -> None: if self._group_started: @@ -750,6 +809,8 @@ def _create_pynccl() -> EplbCommunicator: return TorchDistGlooStagedEplbCommunicator( cpu_group=group_coordinator.cpu_group, ) + elif backend == "torch_xccl": + return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group) elif backend == "torch_nccl": return TorchDistNcclEplbCommunicator(ep_group=torch_group) elif backend == "pynccl": diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index fb5475a607f5..e4d805b9fead 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -203,9 +203,9 @@ class EplbModelState: """ EPLB stats for the model. """ - cuda_device_index: int | None + device_index: int | None """ - CUDA device index for the async EPLB worker thread. + Device index for the async EPLB worker thread. """ communicator: EplbCommunicator """ @@ -286,14 +286,14 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device): """ Background thread handling async transfers. """ - self.cuda_device_index: int | None = None + self.device_index: int | None = None """ - CUDA device index for the async EPLB worker thread. + Device index for the async EPLB worker thread. """ - if self.device.type == "cuda": - self.cuda_device_index = self.device.index - if self.cuda_device_index is None and torch.cuda.is_available(): - self.cuda_device_index = torch.accelerator.current_device_index() + if self.device.type in ("cuda", "xpu"): + self.device_index = self.device.index + if self.device_index is None and torch.accelerator.is_available(): + self.device_index = torch.accelerator.current_device_index() @staticmethod def build_initial_global_physical_to_logical_map( @@ -504,7 +504,7 @@ def add_model( expert_buffer=expert_buffer, rebalanced=False, eplb_stats=None, - cuda_device_index=self.cuda_device_index, + device_index=self.device_index, communicator=communicator, num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) @@ -763,8 +763,8 @@ def rearrange( is_main_rank = ep_rank == 0 if is_main_rank: if not self.is_async or is_profile: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) start_event.record() logger.info( "Rearranging experts %s %s...", diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 3d5be8a8ec4d..943f6d9bef06 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for EPLB (Expert Parallel Load Balancing).""" +import contextlib import os import threading @@ -13,6 +14,22 @@ logger = init_logger(__name__) +@contextlib.contextmanager +def device_stream(stream: torch.Stream | None): + """Platform-agnostic context manager that activates *stream* as the + current accelerator stream for the duration of the ``with`` block. + A no-op when *stream* is ``None``.""" + if stream is None: + yield + return + prev = torch.accelerator.current_stream() + torch.accelerator.set_stream(stream) + try: + yield + finally: + torch.accelerator.set_stream(prev) + + class CpuGpuEvent: """Combines a CUDA event with a CPU threading event to enforce record->wait ordering across two threads. @@ -30,10 +47,10 @@ class CpuGpuEvent: """ def __init__(self): - self._event = torch.cuda.Event() + self._event = torch.Event() self._recorded = threading.Event() - def wait(self, stream: torch.cuda.Stream | None = None): + def wait(self, stream: torch.Stream | None = None): """Blocks the calling thread until record finishes. Used to guarantee that the record kernel is called before wait. @@ -43,7 +60,7 @@ def wait(self, stream: torch.cuda.Stream | None = None): self._event.wait(stream) self._recorded.clear() - def record(self, stream: torch.cuda.Stream | None = None): + def record(self, stream: torch.Stream | None = None): """Unblocks the waiting thread after calling event.record(). Should only be called by the main thread. @@ -53,7 +70,7 @@ def record(self, stream: torch.cuda.Stream | None = None): "CpuGpuEvent.record() called before the previous event was " "consumed by wait()" ) - self._event = torch.cuda.Event() + self._event = torch.Event() self._event.record(stream) self._recorded.set() diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 20ebe4f1f298..7a52dd5f56c0 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -13,7 +13,7 @@ from torch.distributed import ProcessGroup, all_gather from vllm.distributed.eplb.eplb_communicator import EplbCommunicator -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -173,7 +173,7 @@ def move_to_buffer( new_indices: np.ndarray, expert_weights: Sequence[torch.Tensor], expert_weights_buffers: Sequence[torch.Tensor], - cuda_stream: torch.cuda.Stream | None, + stream: torch.Stream | None, ep_rank: int, communicator: EplbCommunicator, layer_idx: int = 0, @@ -183,12 +183,12 @@ def move_to_buffer( Args: num_local_experts: Number of local experts. old_indices: (num_experts_total,) ndarray of current (old) - global-to-local expert assignments. + global-to-locals expert assignments. new_indices: (num_experts_total,) ndarray of desired (new) global-to-local assignments after rebalance. expert_weights: Original expert weights for the layer. expert_weights_buffers: Intermediate buffers (one per tensor). - cuda_stream: CUDA stream for async copies (can be None for sync mode). + stream: CUDA/XPU stream for async copies (can be None for sync mode). ep_rank: Rank of this process in expert parallel group. communicator: EplbCommunicator instance for P2P communication. layer_idx: Index of the MoE layer being transferred. @@ -261,10 +261,9 @@ def move_to_buffer( expert = new_local_expert_ids[dst] src_local = expert_to_src_map.get(expert, -1) if src_local != -1: - with torch.cuda.stream(cuda_stream): + with device_stream(stream): for w, b in zip(expert_weights, expert_weights_buffers): b[dst].copy_(w[src_local], non_blocking=True) - communicator.set_transfer_context(old_indices, layer_idx) # 2. Post sends @@ -430,7 +429,7 @@ def transfer_layer( ep_group: ProcessGroup, communicator: EplbCommunicator, is_profile: bool = False, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, rank_mapping: dict[int, int] | None = None, layer_idx: int = 0, ) -> TransferMetadata: @@ -451,7 +450,7 @@ def transfer_layer( is_profile (bool): If `True`, do not perform any actual weight copy. This is used during profile run, where we only perform dummy communications to reserve enough memory for the buffers. - cuda_stream: CUDA stream for async copies (can be None for sync mode). + stream: CUDA stream for async copies (can be None for sync mode). rank_mapping: Optional rank mapping for elastic expert parallelism. layer_idx: Index of the MoE layer being transferred. @@ -499,7 +498,7 @@ def transfer_layer( new_indices=new_layer_indices_np, expert_weights=expert_weights, expert_weights_buffers=expert_weights_buffer, - cuda_stream=cuda_stream, + stream=stream, ep_rank=ep_group.rank(), communicator=communicator, layer_idx=layer_idx, @@ -601,7 +600,7 @@ def rearrange_expert_weights_inplace( new_indices=new_global_expert_indices_cpu[layer_idx], expert_weights=expert_weights[layer_idx], expert_weights_buffers=weights_buffer, - cuda_stream=None, + stream=None, ep_rank=ep_rank, communicator=communicator, layer_idx=layer_idx, diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index dfc578b44170..ddd891b1c090 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -171,6 +171,7 @@ def apply( num_experts=self.moe_config.num_local_experts, ep_rank=self.moe_config.ep_rank, ep_size=self.moe_config.ep_size, + expert_map=expert_map, gemm1_clamp_limit=self.gemm1_clamp_limit, ) assert self.fused_moe_impl is not None diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 0509f02eadf5..b51dbf90dc7e 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -13,7 +13,7 @@ from vllm.triton_utils import tl, triton from vllm.v1.worker.ubatching import dbo_current_ubatch_id -if current_platform.is_cuda_alike(): +if current_platform.is_cuda_alike() or current_platform.is_xpu(): @triton.jit def _eplb_map_and_record_i32_kernel( From 8538017f4b1d01b543345f5dbe1993cb66c80fa9 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 17 Sep 2026 13:58:56 +0800 Subject: [PATCH 377/774] [Frontend] Omit absent fields from /inference/v1/generate stream chunks (#57264) Signed-off-by: aoshen02 Co-authored-by: Claude Fable 5.1 --- .../server/src/routes/inference/generate.rs | 30 ++++++++++++++++ .../test_generate_stream.py | 34 +++++++++++++++++++ .../scale_out/token_in_token_out/serving.py | 10 ++---- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index dfb7e6cbb191..2ce42d8c33f2 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -679,6 +679,36 @@ mod tests { assert!(response.mm_placeholders.is_none()); } + #[test] + fn stream_chunk_omits_absent_fields() { + // Contract shared with the Python frontend: absent optional fields + // are omitted from stream chunks, not serialized as null. + let chunk = GenerateStreamResponse { + request_id: "raw-1".to_string(), + choices: vec![GenerateResponseStreamChoice { + index: 0, + logprobs: None, + finish_reason: None, + token_ids: vec![1], + }], + usage: None, + prompt_token_ids: None, + mm_placeholders: None, + }; + + let json = serde_json::to_value(&chunk).expect("serialize chunk"); + + for key in ["usage", "prompt_token_ids", "mm_placeholders"] { + assert!(json.get(key).is_none(), "{key} should be omitted"); + } + for key in ["logprobs", "finish_reason"] { + assert!( + json["choices"][0].get(key).is_none(), + "{key} should be omitted" + ); + } + } + #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index d9ae424cc85b..0cb8a6d397be 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -569,6 +569,40 @@ async def mock_generate(*args, **kwargs): assert "mm_placeholders" not in chunk +@pytest.mark.asyncio +async def test_stream_chunks_omit_absent_fields(): + """Absent optional fields are omitted from stream chunks, like Rust.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert len(data_chunks) == 2 + for chunk in data_chunks: + assert "usage" not in chunk + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + assert "logprobs" not in chunk["choices"][0] + assert "finish_reason" not in data_chunks[0]["choices"][0] + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + + @pytest.mark.asyncio async def test_stream_text_only_prompt_omits_mm_placeholders_key(): engine = _mock_engine() diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a8a2804a0634..838fe90c5caf 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -486,13 +486,9 @@ async def serve_tokens_stream_generator( total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent prompt metadata (Rust skips None fields). - exclude = { - name - for name in ("prompt_token_ids", "mm_placeholders") - if getattr(chunk, name) is None - } - yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" + # Omit absent fields, like the Rust frontend and the + # final usage chunk below. + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo( From b978df3299518ed7b02584f69d6daac55eab0b71 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 17 Sep 2026 14:31:59 +0800 Subject: [PATCH 378/774] Revert "[Frontend] Omit absent fields from /inference/v1/generate stream chunks" (#57302) --- .../server/src/routes/inference/generate.rs | 30 ---------------- .../test_generate_stream.py | 34 ------------------- .../scale_out/token_in_token_out/serving.py | 10 ++++-- 3 files changed, 7 insertions(+), 67 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 2ce42d8c33f2..dfb7e6cbb191 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -679,36 +679,6 @@ mod tests { assert!(response.mm_placeholders.is_none()); } - #[test] - fn stream_chunk_omits_absent_fields() { - // Contract shared with the Python frontend: absent optional fields - // are omitted from stream chunks, not serialized as null. - let chunk = GenerateStreamResponse { - request_id: "raw-1".to_string(), - choices: vec![GenerateResponseStreamChoice { - index: 0, - logprobs: None, - finish_reason: None, - token_ids: vec![1], - }], - usage: None, - prompt_token_ids: None, - mm_placeholders: None, - }; - - let json = serde_json::to_value(&chunk).expect("serialize chunk"); - - for key in ["usage", "prompt_token_ids", "mm_placeholders"] { - assert!(json.get(key).is_none(), "{key} should be omitted"); - } - for key in ["logprobs", "finish_reason"] { - assert!( - json["choices"][0].get(key).is_none(), - "{key} should be omitted" - ); - } - } - #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 0cb8a6d397be..d9ae424cc85b 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -569,40 +569,6 @@ async def mock_generate(*args, **kwargs): assert "mm_placeholders" not in chunk -@pytest.mark.asyncio -async def test_stream_chunks_omit_absent_fields(): - """Absent optional fields are omitted from stream chunks, like Rust.""" - engine = _mock_engine() - - async def mock_generate(*args, **kwargs): - yield _make_request_output("req-1", token_ids=[10]) - yield _make_request_output( - "req-1", token_ids=[20], finish_reason="stop", finished=True - ) - - engine.generate = MagicMock(side_effect=mock_generate) - serving = _build_serving_tokens(engine) - request = GenerateRequest( - token_ids=[1, 2, 3], - sampling_params=SamplingParams(max_tokens=2), - model=MODEL_NAME, - stream=True, - ) - - response = await serving.serve_tokens(request) - parsed = _parse_sse_chunks([chunk async for chunk in response]) - data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] - - assert len(data_chunks) == 2 - for chunk in data_chunks: - assert "usage" not in chunk - assert "prompt_token_ids" not in chunk - assert "mm_placeholders" not in chunk - assert "logprobs" not in chunk["choices"][0] - assert "finish_reason" not in data_chunks[0]["choices"][0] - assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" - - @pytest.mark.asyncio async def test_stream_text_only_prompt_omits_mm_placeholders_key(): engine = _mock_engine() diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 838fe90c5caf..a8a2804a0634 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -486,9 +486,13 @@ async def serve_tokens_stream_generator( total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent fields, like the Rust frontend and the - # final usage chunk below. - yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + # Omit absent prompt metadata (Rust skips None fields). + exclude = { + name + for name in ("prompt_token_ids", "mm_placeholders") + if getattr(chunk, name) is None + } + yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo( From fdcd42350bd2475abffc01bb73eec5d7c7a1e549 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Thu, 17 Sep 2026 14:46:00 +0800 Subject: [PATCH 379/774] [ROCm][Perf] Enable HCA dual-stream overlap for DeepSeek-V4 (#56853) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- vllm/models/deepseek_v4/amd/rocm.py | 44 +++++++++++++++++++++++------ 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index f8ac39cd6241..545f29bd356b 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -558,14 +558,17 @@ def __init__(self, *args, **kwargs): self._fused_compressor_split_sizes: tuple[int, int] | None = None if self.indexer is None: - # Only enable multi-stream overlap for CSA layer now. - self.aux_stream_list = None + # Dense layers have no compressor work to overlap; HCA layers + # (compressor, no indexer) keep the streams for the dual-stream + # fork below. + if self.compressor is None: + self.aux_stream_list = None else: # Disable indexer inner overlap. self.indexer.aux_stream = None - def _enable_csa_multi_stream(self) -> bool: - """All CSA multi-stream gates: env var, streams, and capture region. + def _enable_multi_stream_overlap(self) -> bool: + """ROCm multi-stream gates: streams and capture region. Dict metadata marks piecewise cudagraph, whose eager breaks rebuild the attention inputs on the owning stream. Forking side streams @@ -573,7 +576,7 @@ def _enable_csa_multi_stream(self) -> bool: this overlap on ROCm (event waits can hang), so multi-stream only runs where the fork/join becomes static graph edges: inside capture, or with non-dict metadata (full cudagraph or the profile run), which - has no eager breaks. + has no eager breaks. Covers both the HCA and CSA forks. """ attn_metadata = get_forward_context().attn_metadata return self.aux_stream_list is not None and ( @@ -624,7 +627,7 @@ def forward( device=hidden_states.device, ) - if self._enable_csa_multi_stream(): + if self._enable_multi_stream_overlap(): # The ROCm override consumes these sentinels inside the capture # boundary, moving the stream fan-out ahead of the projections. self._prepare_and_attn_fn( @@ -658,7 +661,7 @@ def _prepare_and_attn( positions: torch.Tensor, o_padded: torch.Tensor, ) -> None: - """Run the ROCm CSA fork/join inside the graph capture boundary.""" + """Run the ROCm fork/join (HCA or CSA) inside the capture boundary.""" aux_streams = self.aux_stream_list # The sequential pipeline disables aux_stream_list before calling # back with real projection inputs; aux_streams is None ends that @@ -684,13 +687,13 @@ def _prepare_and_attn( # Re-check: forward's gate ran inside a captured segment that # _prepare_and_attn_eager (MRV1) then broke, making this region eager. - if not self._enable_csa_multi_stream(): + if not self._enable_multi_stream_overlap(): self._run_sequential_pipeline(hidden_states, positions, o_padded) return indexer = self.indexer compressor = self.compressor - assert indexer is not None and compressor is not None + assert compressor is not None def default_chain(): qr_kv = self._fused_wqa_wkv_gemm(hidden_states) @@ -710,6 +713,24 @@ def main_compressor_chain() -> None: ) compressor(score, positions, self.rotary_emb) + if indexer is None: + # HCA dual-stream: the main compressor runs on aux stream 0 while + # the default stream produces q and inserts KV into the SWA cache. + # Both branches only read hidden_states, so the join merely has to + # precede the sparse attention that consumes the compressed KV. + (q, _qr_out, _qr_scale_out, kv_out), _ = execute_in_parallel( + default_chain, + [main_compressor_chain], + self.ln_events[0], + [self.ln_events[1]], + aux_streams[:1], + enable=True, + ) + self._sparse_indexer_and_attn( + hidden_states, None, None, None, q, kv_out, positions, o_padded + ) + return + def indexer_compressor_chain() -> None: score = torch.mm( hidden_states, @@ -718,6 +739,11 @@ def indexer_compressor_chain() -> None: ) indexer.compressor(score, positions, self.indexer_rotary_emb) + # CSA three-stream: the main and indexer compressors run on aux + # streams 0 and 1 while the default stream produces q and inserts KV + # into the SWA cache. Every branch only reads hidden_states plus its + # own state, so the join merely has to precede the indexer op and the + # sparse attention, which consume the compressed KV caches. (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel( default_chain, [main_compressor_chain, indexer_compressor_chain], From bd2d7ae7c197ba995800bb03a23b4d79c38627a4 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 17 Sep 2026 03:01:19 -0400 Subject: [PATCH 380/774] [Perf] Reuse MoE workspace for DeepGEMM warmup (#57268) Signed-off-by: Lucas Wilkinson --- vllm/model_executor/warmup/deep_gemm_warmup.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 6aa5d391e01d..1266f4ccb265 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -30,6 +30,7 @@ ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units +from vllm.v1.worker.workspace import current_workspace_manager def _generate_optimal_warmup_m_values( @@ -275,15 +276,15 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ) if not warmup_cases: return - device = w1.device def _warmup(w: torch.Tensor, w_scale: torch.Tensor): _, n, k = w.size() - a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.zeros( - (MAX_M, k // block_m), device=device, dtype=torch.float32 + a1q, a1q_scales, out = current_workspace_manager().get_simultaneous( + ((MAX_M, k), torch.float8_e4m3fn), + ((MAX_M, k // block_m), torch.float32), + ((MAX_M, n), torch.bfloat16), ) - out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16) + a1q_scales.zero_() for num_tokens, align_used, expert_ids in warmup_cases: with mk_alignment_scope(align_used): From 7bbce752b85f183eb34e27ea6df9647f80986060 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:02:24 +0800 Subject: [PATCH 381/774] [Perf][Model] Qwen4Exp QSA: sm_90 tuning table for _select_config (#57273) Signed-off-by: Shuolei Wang Signed-off-by: Shuolei Wang <948904026@qq.com> Co-authored-by: Roger Wang --- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 40 ++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index f0aa62a48460..816ccdd498d9 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -22,6 +22,12 @@ def _is_sm120() -> bool: return current_platform.get_device_capability() == (12, 0) +@lru_cache(maxsize=1) +def _is_sm90() -> bool: + """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" + return current_platform.get_device_capability() == (9, 0) + + @triton.jit(do_not_specialize=["num_rows", "num_requests"]) def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, @@ -509,6 +515,33 @@ def _select_sm120_config( return 32, 4, 1 +def _select_sm90_config( + base_programs: int, use_prefill_config: bool, is_fp8: bool +) -> tuple[int, int, int]: + """(block_n, target_splits, num_warps) tuned on sm_90 (H20).""" + if base_programs > 2048: + return (32, 1, 1) if use_prefill_config else (64, 1, 2) + if is_fp8: + if base_programs <= 24: + return 64, 64, 2 + if base_programs <= 32: + return 32, 16, 1 + if base_programs <= 64: + return 32, 8, 1 + if base_programs <= 128: + return 32, 4, 1 + if base_programs <= 256: + return 32, 8, 1 + return 32, 4, 1 + if base_programs <= 32: + return 64, 64, 2 + if base_programs <= 64: + return 32, 16, 1 + if base_programs <= 256: + return 32, 8, 1 + return 32, 4, 1 + + def _select_config( num_rows: int, num_kv_heads: int, @@ -521,13 +554,18 @@ def _select_config( Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the uniform decode/verify length). This default table was tuned on GB300; - sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead. + sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead, + and sm_90 (Hopper) to _select_sm90_config. """ base_programs = num_rows * num_kv_heads if _is_sm120(): BLOCK_N, target_splits, num_warps = _select_sm120_config( base_programs, use_prefill_config, is_fp8 ) + elif _is_sm90(): + BLOCK_N, target_splits, num_warps = _select_sm90_config( + base_programs, use_prefill_config, is_fp8 + ) elif base_programs > 2048: BLOCK_N, target_splits, num_warps = ( (32, 1, 1) if use_prefill_config else (64, 1, 2) From b3079e6e46a21109a8641c6d48641453a9c7cb86 Mon Sep 17 00:00:00 2001 From: leo <92263211+liuyao0322@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:04:53 +0800 Subject: [PATCH 382/774] [Perf][DSpark] Stack DeepSeek V4 context WKV projections (#54674) Signed-off-by: liuyao0322 Co-authored-by: OpenAI Codex --- tests/models/test_dspark_mla.py | 73 ++++++++++++++++++++++++ vllm/models/deepseek_v4/nvidia/dspark.py | 52 +++++++++++++++-- 2 files changed, 119 insertions(+), 6 deletions(-) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index d30eb9e7c91d..3e70450f2ce0 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead from vllm.model_executor.models.registry import ModelRegistry +from vllm.models.deepseek_v4.nvidia import dspark as dsv4_dspark from vllm.models.kimi_k3.nvidia import dspark_mla from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel @@ -262,3 +263,75 @@ def load_scale(param, weight, *args): assert loaded == {runtime_name} assert shards == [() if shard_id is None else (shard_id,)] torch.testing.assert_close(param, checkpoint_scale) + + +def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer(): + weights = [ + ("mtp.0.attn.wkv.weight", torch.arange(4)), + ("mtp.1.attn.wq_a.weight", torch.arange(3)), + ("mtp.2.attn.wkv.scale", torch.tensor(0.5)), + ("mtp.3.attn.wkv.weight", torch.arange(2)), + ] + + duplicated = list(dsv4_dspark._duplicate_context_wkv_weights(weights, 3)) + + assert [name for name, _ in duplicated] == [ + "mtp.0.attn.wkv.weight", + "context_wkv_proj.weight", + "mtp.1.attn.wq_a.weight", + "mtp.2.attn.wkv.scale", + "context_wkv_proj.scale", + "mtp.3.attn.wkv.weight", + ] + assert duplicated[1][1].shard_id == 0 + assert duplicated[4][1].shard_id == 2 + assert duplicated[0][1].data_ptr() == duplicated[1][1].data_ptr() + assert duplicated[3][1].data_ptr() == duplicated[4][1].data_ptr() + + +def test_dsv4_context_kv_uses_one_stacked_wkv_projection(monkeypatch): + calls = [] + stacked_output = torch.arange(24, dtype=torch.float32).view(2, 12) + + class StackedProjection: + def __init__(self): + self.calls = 0 + + def __call__(self, main_x): + self.calls += 1 + assert main_x.shape == (2, 5) + return stacked_output + + projection = StackedProjection() + layers = [ + SimpleNamespace(attn=SimpleNamespace(kv_norm=lambda kv, offset=i: kv + offset)) + for i in range(3) + ] + model = SimpleNamespace( + config=SimpleNamespace(head_dim=4), + context_wkv_proj=projection, + layers=layers, + num_dspark_layers=3, + ) + slot_mappings = [torch.tensor([0, 1]), None, torch.tensor([4, 5])] + monkeypatch.setattr( + dsv4_dspark, + "_insert_context_kv", + lambda attn, kv, positions, slots: calls.append( + (attn, kv.clone(), positions, slots) + ), + ) + + dsv4_dspark.DSparkDeepseekV4Model.precompute_and_store_context_kv( + model, + torch.zeros(2, 5), + torch.tensor([7, 8]), + slot_mappings, + ) + + assert projection.calls == 1 + assert len(calls) == 2 + assert torch.equal(calls[0][1], stacked_output.view(2, 3, 4)[:, 0]) + assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2) + assert calls[0][3] is slot_mappings[0] + assert calls[1][3] is slot_mappings[2] diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 1a3c2c780d28..52783889d6ca 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -31,7 +31,10 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, +) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -61,6 +64,24 @@ # MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders): # fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``. _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") +_CONTEXT_WKV_RE = re.compile(r"^mtp\.(\d+)\.attn\.wkv\.(.+)$") + + +def _duplicate_context_wkv_weights( + weights: Iterable[tuple[str, torch.Tensor]], num_layers: int +) -> Iterable[tuple[str, torch.Tensor]]: + """Load every draft layer's WKV into the cross-layer projection.""" + for name, weight in weights: + yield name, weight + match = _CONTEXT_WKV_RE.fullmatch(name) + if match is None: + continue + layer_idx = int(match.group(1)) + if layer_idx >= num_layers: + continue + stacked_weight = weight.detach() + stacked_weight.shard_id = layer_idx + yield f"context_wkv_proj.{match.group(2)}", stacked_weight class DSparkDeepseekV4Model(nn.Module): @@ -113,6 +134,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: for i in range(self.num_dspark_layers) ] ) + self.context_wkv_proj = MergedColumnParallelLinear( + config.hidden_size, + [config.head_dim] * self.num_dspark_layers, + bias=False, + return_bias=False, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "context_wkv_proj"), + disable_tp=True, + ) # Heads: final norm + hc_head, and the Markov + confidence heads # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint @@ -173,15 +203,16 @@ def precompute_and_store_context_kv( place draft layers in different groups). ``None`` (or a ``None`` entry) runs the projection to reserve workspace but writes nothing (profiling). """ - for i, layer in enumerate(self.layers): + all_kv = self.context_wkv_proj(main_x).view( + main_x.shape[0], self.num_dspark_layers, self.config.head_dim + ) + for i, (layer, kv) in enumerate( + zip(self.layers, all_kv.unbind(1), strict=True) + ): slot_mapping = ( None if context_slot_mappings is None else context_slot_mappings[i] ) attn = layer.attn - # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection - # (q_lora part discarded), then RoPE/quant/insert via the fused op. - qr_kv, _ = attn.fused_wqa_wkv(main_x) - kv = qr_kv[..., attn.q_lora_rank :] kv = attn.kv_norm(kv) if slot_mapping is None: continue @@ -429,6 +460,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: head_start = n_local_head * tp_rank head_end = n_local_head * (tp_rank + 1) + weights = _duplicate_context_wkv_weights(weights, len(self.model.layers)) for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) if mapped is None: @@ -452,6 +484,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: self.quant_config, name, loaded_weight ) + if name.startswith("model.context_wkv_proj."): + param = params_dict[name] + param.weight_loader(param, loaded_weight, loaded_weight.shard_id) + loaded_params.add(name) + continue + # E8M0 expert scales: keep raw exponent bytes. if ".experts." in name: if ( @@ -522,6 +560,8 @@ def _remap_dspark_name(self, name: str) -> str | None: Returns None for non-mtp weights (owned by the target model). """ + if name.startswith("context_wkv_proj."): + return f"model.{name}" m = re.match(r"mtp\.(\d+)\.(.*)", name) if m is None: return None From 528fa835fd37d0805a824d516e22f88df61a771a Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Thu, 17 Sep 2026 00:14:43 -0700 Subject: [PATCH 383/774] [ROCm] Stage large pageable H2D copies instead of registering them (#56343) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Shanshan Shen <467638484@qq.com> --- vllm/platforms/rocm.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 248f41d20d79..9089d025e0e5 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.config.kernel import IrOpPriorityConfig + from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.selector import AttentionSelectorConfig logger = init_logger(__name__) @@ -493,6 +494,9 @@ class RocmPlatform(Platform): dist_backend: str = "nccl" # rocm shares the same device control env var as CUDA device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + # Set in pre_register_and_update, so it exists only on the driver; Ray + # workers are separate processes and copy env vars by allowlist. + additional_env_vars: list[str] = ["GPU_PINNED_MIN_XFER_SIZE"] ray_noset_device_env_vars: list[str] = [ "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", @@ -862,6 +866,16 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: ) return torch.cuda.get_device_properties(device_id).total_memory + @classmethod + def pre_register_and_update( + cls, parser: "FlexibleArgumentParser | None" = None + ) -> None: + # Keep mmap'd weight pages on the HIP staging path: above this + # threshold the runtime registers the pageable source instead, and each + # registration's MMU notifier makes KFD suspend our queues. In KB, so + # 4 GiB. + os.environ.setdefault("GPU_PINNED_MIN_XFER_SIZE", str(4 * 1024 * 1024)) + @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: from vllm._aiter_ops import rocm_aiter_ops From 4e5ffda19d897f3c34455167dc0503620f1b3e2a Mon Sep 17 00:00:00 2001 From: "Jiahan Chang (Cyrus)" <173873397+jiahanc@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:38:58 +0800 Subject: [PATCH 384/774] [Bugfix] Isolate supplemental FlashInfer BF16 autotuning (#57285) Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: OpenAI Codex --- vllm/model_executor/warmup/kernel_warmup.py | 43 +++++++++++++++++++-- vllm/utils/flashinfer.py | 43 +++++++++++++++++---- 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index bf2c1ad0b693..81eca7ac50e5 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -330,13 +330,16 @@ def _flashinfer_deferred_moe_token_counts( return tuple(dict.fromkeys(token_counts)) -def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]: +def _flashinfer_autotune_token_counts( + runner: "GPUModelRunner", *, include_bf16: bool = True +) -> tuple[int, ...]: max_tokens = runner.scheduler_config.max_num_batched_tokens # Tune the widest bucket set first so bounded passes reuse its configs. token_counts = [max_tokens] linear_backend = runner.vllm_config.kernel_config.linear_backend if ( - linear_backend == "flashinfer_cutedsl" + include_bf16 + and linear_backend == "flashinfer_cutedsl" and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS ): token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) @@ -350,7 +353,7 @@ def _run_flashinfer_autotune_dummy_runs( import vllm.utils.flashinfer as fi_utils dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} - for num_tokens in _flashinfer_autotune_token_counts(runner): + for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False): tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) logger.info( "Running FlashInfer autotune with %d tokens and token buckets %s.", @@ -367,6 +370,36 @@ def _run_flashinfer_autotune_dummy_runs( ) +def _run_flashinfer_bf16_autotune_dummy_run( + runner: "GPUModelRunner", + *, + skip_ops: set[str] | None = None, + skip_attn: bool = False, +) -> None: + import vllm.utils.flashinfer as fi_utils + + if ( + runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl" + or runner.scheduler_config.max_num_batched_tokens + <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS + ): + return + + num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS + tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) + logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens) + # Other ops execute from their existing cache or fallback, without the + # BF16 bucket cap. No full-model autotune context may enclose this pass. + with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops): + runner._dummy_run( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + randomize_inputs=True, + **({"skip_attn": True} if skip_attn else {}), + ) + + def flashinfer_autotune(runner: "GPUModelRunner") -> None: """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, @@ -435,6 +468,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) + with torch.inference_mode(): + _run_flashinfer_bf16_autotune_dummy_run( + runner, skip_ops=skip_ops, skip_attn=hisparse_enabled + ) finally: set_autotune_process_group(None) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 76471791349d..abc4d75500c1 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -11,7 +11,8 @@ import importlib.util import os import shutil -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextvars import ContextVar from typing import Any, NoReturn import requests @@ -26,6 +27,24 @@ logger = init_logger(__name__) +_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar( + "flashinfer_bf16_autotune_buckets", default=None +) + + +@contextlib.contextmanager +def autotune_bf16_only( + tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None +) -> Iterator[None]: + """Tune BF16 calls with bounded buckets, outside full-model autotuning.""" + token = _bf16_autotune_buckets.set(tuning_buckets) + try: + with autotune(tune_mode=False, skip_ops=skip_ops): + yield + finally: + _bf16_autotune_buckets.reset(token) + + def flashinfer_bf16_mm( a: torch.Tensor, b: torch.Tensor, @@ -35,14 +54,21 @@ def flashinfer_bf16_mm( ) -> torch.Tensor: from flashinfer import mm_bf16 - return mm_bf16( - a, - b, - bias=bias, - pdl=pdl, - out_dtype=torch.bfloat16, - backend=backend, + tuning_buckets = _bf16_autotune_buckets.get() + tuning = ( + autotune(tune_mode=True, tuning_buckets=tuning_buckets) + if tuning_buckets is not None + else contextlib.nullcontext() ) + with tuning: + return mm_bf16( + a, + b, + bias=bias, + pdl=pdl, + out_dtype=torch.bfloat16, + backend=backend, + ) # This is the storage path for the cubins, it can be replaced @@ -1237,6 +1263,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: __all__ = [ "has_flashinfer", "flashinfer_bf16_mm", + "autotune_bf16_only", "has_flashinfer_bf16_gemm", "is_flashinfer_bf16_gemm_supported", "is_flashinfer_cutedsl_bf16_gemm_supported", From de24e5190820cc6640f2d55295405f228fffbc41 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Thu, 17 Sep 2026 16:12:17 +0800 Subject: [PATCH 385/774] [MM][V2] Enable encoder-only ViT CUDA graph capture (#56922) Signed-off-by: jiangkuaixue123 Co-authored-by: Codex --- .../test_areas/disaggregated_mooncake.yaml | 10 ++ .../generation/test_vit_cudagraph.py | 115 ++++++++++++++++++ tests/v1/cudagraph/test_encoder_cudagraph.py | 95 +++++++++++++++ .../run_epd_mooncake_ec_full_pipeline.sh | 36 +++++- .../integration/test_epd_correctness.py | 21 ++++ tests/v1/worker/test_encoder_runner.py | 55 ++++++++- vllm/v1/worker/mm_encoder_model_runner.py | 35 +++++- 7 files changed, 362 insertions(+), 5 deletions(-) diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index 8c061a95ad8f..f6e944d70581 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.buildkite/test_areas/disaggregated_mooncake.yaml @@ -40,7 +40,14 @@ steps: - vllm/v1/worker/gpu_worker.py - vllm/v1/worker/gpu/ec_connector.py - vllm/v1/worker/gpu/model_runner.py + - vllm/v1/worker/mm_encoder_model_runner.py + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py + - vllm/v1/attention/ops/vit_attn_wrappers.py + - vllm/v1/worker/gpu/model_states/ - vllm/v1/worker/gpu/mm/ + - tests/v1/cudagraph/test_encoder_cudagraph.py + - tests/models/multimodal/generation/test_vit_cudagraph.py - examples/disaggregated/disaggregated_encoder/ - tests/v1/ec_connector/ - requirements/kv_connectors.txt @@ -49,6 +56,9 @@ steps: PYTHONUNBUFFERED: "1" VLLM_HOST_IP: "127.0.0.1" VLLM_USE_V2_MODEL_RUNNER: "1" + E_CUDAGRAPH_MM_ENCODER: "1" + # FA padding NaNs are tracked separately in #57136. + E_MM_ENCODER_ATTN_BACKEND: "TRITON_ATTN" MOONCAKE_EC_PROTOCOL: "tcp" USE_MM_PROMPTS: "1" SKIP_BASELINE: "0" diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index f8136ae8f3a7..1ed84ebc5275 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -499,3 +499,118 @@ def test_vit_cudagraph_video(model_id, vllm_runner, video_assets): # Ensure the output is a string assert isinstance(output_text, str) + + +def _check_eonly_encoder_outputs(worker, batches): + """Compare real E-only replay with eager, retaining outputs across replay.""" + import torch + + from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner + + runner = worker.model_runner + assert isinstance(runner, MMEncoderModelRunner) + manager = runner.model_state.encoder_runner.cudagraph_manager + assert manager is not None and manager.is_captured() + retained: list[tuple[torch.Tensor, torch.Tensor]] = [] + max_error = 0.0 + stats = [] + with torch.inference_mode(): + for batch in batches: + kwargs = {key: value.to(runner.device) for key, value in batch.items()} + actual = manager.execute(kwargs) + expected = runner.model.embed_multimodal(**kwargs) + assert len(actual) == len(expected) + for output, reference in zip(actual, expected): + assert torch.isfinite(output).all() + # Two BF16 ulps at unit scale, fixed before model validation. + torch.testing.assert_close(output, reference, rtol=0.016, atol=0.016) + max_error = max(max_error, (output - reference).abs().max().item()) + for output, snapshot in retained: + torch.testing.assert_close(output, snapshot, rtol=0, atol=0) + retained.extend((output, output.clone()) for output in actual) + stats.append(manager.get_cumulative_stats()) + assert stats[0]["graph_hits"] > 0 + assert stats[1]["graph_hits"] > stats[0]["graph_hits"] + assert stats[2]["graph_misses"] > stats[1]["graph_misses"] + assert stats[3]["graph_hits"] > stats[2]["graph_hits"] + assert stats[4]["graph_hits"] > stats[3]["graph_hits"] + assert stats[5]["graph_hits"] > stats[4]["graph_hits"] + assert stats[6]["graph_hits"] > stats[5]["graph_hits"] + return {"max_abs_error": max_error, "stats": stats} + + +@pytest.mark.parametrize( + "model_id", + params_with_marks({key: MODEL_CONFIGS[key] for key in ("qwen3_vl", "qwen3_5")}), +) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA graphs") +def test_eonly_vit_cudagraph_outputs( + model_id, vllm_runner, image_assets, video_assets, monkeypatch +): + """Exercise startup capture, image/video replay, fallback and output ownership.""" + import numpy as np + from PIL import Image + from transformers import AutoProcessor + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + config = MODEL_CONFIGS[model_id] + token_budgets = [64, 256, 2048] + processor = AutoProcessor.from_pretrained(config.model) + first, second = [asset.pil_image for asset in image_assets] + batches = [ + [first.resize((224, 224))], + [second.resize((448, 224)), first.resize((224, 224))], + [second.resize((1792, 1792))], + [second.resize((224, 224))], + [first.resize((1280, 720))], + [first.resize((224, 224))], + ] + frames = sample_frames_from_video(video_assets[0].np_ndarrays, 2) + video = np.stack( + [np.asarray(Image.fromarray(frame).resize((224, 224))) for frame in frames] + ) + with vllm_runner( + config.model, + dtype=config.dtype, + mm_encoder_only=True, + # FA padding NaNs are tracked separately in #57136. + mm_encoder_attn_backend="FLASHINFER", + enable_prefix_caching=False, + max_model_len=4096, + max_num_seqs=2, + limit_mm_per_prompt={"image": 2, "video": 1}, + compilation_config={ + "cudagraph_mode": "NONE", + "cudagraph_mm_encoder": True, + "encoder_cudagraph_token_budgets": token_budgets, + "encoder_cudagraph_max_vision_items_per_batch": 2, + "encoder_cudagraph_max_frames_per_batch": 2, + }, + ) as model: + mm_config = model.llm.llm_engine.vllm_config.model_config.multimodal_config + processor_kwargs = mm_config.merge_mm_processor_kwargs({}) + inputs = [ + dict( + processor.image_processor( + images=images, return_tensors="pt", **processor_kwargs + ) + ) + for images in batches + ] + inputs.append( + dict( + processor.video_processor( + videos=[video], + do_sample_frames=False, + return_tensors="pt", + **processor_kwargs, + ) + ) + ) + results = model.llm.collective_rpc( + partial(_check_eonly_encoder_outputs, batches=inputs) + ) + for result in results: + assert result["stats"][0]["token_budgets"] == token_budgets + print(f"E-only {model_id}: {results}") diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 096f0e40750f..1f145bdc8962 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -13,7 +13,9 @@ """ from collections.abc import Hashable +from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock, patch import pytest import torch @@ -575,6 +577,99 @@ def test_chunking_when_images_exceed_max_batch(self): assert out.shape == (4, _HIDDEN) +# --------------------------------------------------------------------------- +# E-only capture and output lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) +@pytest.mark.usefixtures("dist_init", "workspace_init") +@pytest.mark.parametrize("profile_only", [False, True]) +@torch.inference_mode() +def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only): + """The E-only entry captures only the encoder and preserves cached outputs.""" + from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, + ) + from vllm.v1.worker.gpu.ec_connector import ActiveECConnector + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner + from vllm.v1.worker.workspace import lock_workspace + + device = torch.device("cuda:0") + dtype = torch.float16 + model = SimpleMockViTModel().to(device).half() + manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype) + encoder = object.__new__(EncoderRunner) + encoder.device = device + encoder.cudagraph_manager = manager + runner = object.__new__(MMEncoderModelRunner) + runner.model_state = SimpleNamespace(encoder_runner=encoder) + # No decoder manager is installed: capture must be encoder-only. + with patch( + "vllm.v1.worker.mm_encoder_model_runner.lock_workspace", wraps=lock_workspace + ) as lock: + runner.capture_model(profile_only=profile_only) + assert lock.call_count == int(not profile_only) + assert len(manager.budget_graphs["default"]) == len(_BUDGETS) + + cache: dict[str, torch.Tensor] = {} + pending_sends: dict[str, torch.Tensor] = {} + connector = MagicMock(spec=ECConnectorBase) + connector.is_producer = True + connector.is_consumer = False + connector.get_finished.return_value = (None, None) + connector.save_caches.side_effect = lambda *, encoder_cache, mm_hash: ( + pending_sends.update({mm_hash: encoder_cache[mm_hash]}) + ) + with patch( + "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector + ): + ec = ActiveECConnector(SimpleNamespace(), cache) + scheduled = SimpleNamespace( + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset() + ) + saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = [] + # Mixed sizes exercise packing order; 64 is the boundary, 81 falls back. + for grids in ([[1, 8, 8], [1, 4, 4]], [[1, 16, 16]], [[1, 18, 18]], [[1, 4, 4]]): + inputs = _make_mm_kwargs(grids, device, dtype) + expected = model.encoder_eager_forward(inputs).split( + [t * (h // 2) * (w // 2) for t, h, w in grids] + ) + with ec.maybe_get_output(scheduled) as ec_output: + outputs = manager.execute(inputs) + assert outputs is not None + cache[str(len(cache))] = outputs[0] + assert ec_output.finished_sending is None + assert outputs is not None + for actual, eager in zip(outputs, expected): + torch.testing.assert_close(actual, eager) + for previous, snapshot in saved_outputs: + torch.testing.assert_close(previous, snapshot, rtol=0, atol=0) + saved_outputs.extend((output, output.clone()) for output in outputs) + assert manager.graph_hits == 4 + assert manager.graph_misses == 1 + assert pending_sends.keys() == cache.keys() + connector.get_finished.return_value = (set(cache), None) + with ec.maybe_get_output(scheduled) as ec_output: + pass + assert ec_output.finished_sending == set(cache) + + +def test_eonly_without_encoder_graph_skips_capture(): + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner + + encoder = object.__new__(EncoderRunner) + encoder.cudagraph_manager = None + runner = object.__new__(MMEncoderModelRunner) + runner.model_state = SimpleNamespace(encoder_runner=encoder) + assert runner.capture_model() == 0 + + # --------------------------------------------------------------------------- # SimpleMockViTVideoModel — extends SimpleMockViTModel with video support # --------------------------------------------------------------------------- diff --git a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh index 05af1bb6fdbc..c4e9b0d06975 100755 --- a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh +++ b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh @@ -20,6 +20,8 @@ # SKIP_BASELINE set to 1 to reuse existing BASELINE_FILE # CONCURRENCY / REPEAT concurrent requests / rounds (defaults 3 / 2) # MAX_MODEL_LEN context length (default 16384) +# E_CUDAGRAPH_MM_ENCODER enable encoder CUDA graphs (default 0) +# E_MM_ENCODER_ATTN_BACKEND optional encoder attention backend override set -euo pipefail @@ -35,6 +37,15 @@ if [[ "$USE_MM_PROMPTS" == "1" ]]; then TEST_ARGS+=(--use_mm_prompts --mm_smoke_test) fi MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}" +E_ATTN_ARGS=() +if [[ -n "${E_MM_ENCODER_ATTN_BACKEND:-}" ]]; then + E_ATTN_ARGS+=(--mm-encoder-attn-backend "$E_MM_ENCODER_ATTN_BACKEND") +fi +case "${E_CUDAGRAPH_MM_ENCODER:-0}" in + 0) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":false}' ;; + 1) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":true}' ;; + *) echo "E_CUDAGRAPH_MM_ENCODER must be 0 or 1"; exit 1 ;; +esac GPU_SINGLE="${GPU_SINGLE:-0}" GPU_E="${GPU_E:-0}" @@ -203,7 +214,7 @@ run_epd_mooncake() { cleanup_instances echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \ + VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \ --port "$ENCODE_PORT" \ --gpu-memory-utilization 0.35 \ --mm-tensor-ipc torch_shm \ @@ -215,6 +226,9 @@ run_epd_mooncake() { --limit-mm-per-prompt '{"image":2,"video":0}' \ --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \ --ec-transfer-config "$ENC_EC_JSON" \ + --compilation-config "$E_COMPILATION_CONFIG" \ + "${E_ATTN_ARGS[@]}" \ + --worker-extension-cls tests.v1.ec_connector.integration.test_epd_correctness.EncoderGraphTestWorkerExtension \ >"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 & local ENCODER_PID=$! PIDS+=("$ENCODER_PID") @@ -266,6 +280,26 @@ run_epd_mooncake() { --baseline_file "$BASELINE_FILE" \ "${TEST_ARGS[@]}" + curl -fsS "http://localhost:${ENCODE_PORT}/collective_rpc" \ + -H 'Content-Type: application/json' \ + -d '{"method":"encoder_graph_stats"}' \ + >"${LOG_PATH}/encoder_graph_stats.json" + "$PYTHON_BIN" - "${LOG_PATH}/encoder_graph_stats.json" "${E_CUDAGRAPH_MM_ENCODER:-0}" <<'PY' +import json +import sys + +with open(sys.argv[1]) as f: + results = json.load(f)["results"] +assert results, "No encoder workers returned graph statistics" +for stats in results: + if sys.argv[2] == "1": + assert stats["is_captured"], stats + assert stats["graph_hits"] > 0, stats + else: + assert stats["graph_hits"] == 0, stats +print("Encoder graph statistics:", results) +PY + cleanup_instances } diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index b8565a0acd8c..4e37637ea452 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/tests/v1/ec_connector/integration/test_epd_correctness.py @@ -26,6 +26,7 @@ import os import time from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING import openai import requests @@ -33,6 +34,26 @@ from vllm.assets.image import ImageAsset from vllm.multimodal.utils import encode_image_url +if TYPE_CHECKING: + from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner + + +class EncoderGraphTestWorkerExtension: + """Expose encoder replay evidence through the existing test RPC endpoint.""" + + model_runner: "MMEncoderModelRunner" + + def encoder_graph_stats(self): + from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner + + runner = self.model_runner + assert isinstance(runner, MMEncoderModelRunner) + manager = runner.model_state.encoder_runner.cudagraph_manager + if manager is None: + return {"is_captured": False, "graph_hits": 0, "graph_misses": 0} + return {"is_captured": manager.is_captured(), **manager.get_cumulative_stats()} + + MAX_OUTPUT_LEN = 256 # Sample prompts with multimodal content diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index f8bd50c6d9c4..98aaa568f0b6 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -9,12 +9,15 @@ the processed range still fails loudly. """ -from unittest.mock import MagicMock +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import numpy as np import pytest import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldElem, @@ -31,6 +34,56 @@ HIDDEN = 4 +@pytest.mark.parametrize( + "enabled,enforce_eager,supported,expected", + [ + (True, False, True, True), + (False, False, True, False), + (True, True, True, False), + (True, False, False, False), + ], +) +def test_encoder_graph_configuration(enabled, enforce_eager, supported, expected): + """Encoder graph eligibility is independent of decoder graph mode.""" + config = MagicMock() + config.model_config.enforce_eager = enforce_eager + config.model_config.dtype = torch.float32 + config.model_config.get_inputs_embeds_size.return_value = HIDDEN + config.scheduler_config.max_num_batched_tokens = 8 + config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + config.compilation_config.cudagraph_mm_encoder = enabled + model = MagicMock(spec=SupportsEncoderCudaGraph) if supported else torch.nn.Module() + state = _model_state(EncoderCache()) + with patch("vllm.v1.worker.gpu.model_states.interface.EncoderCudaGraphManager"): + ModelState.__init__(state, config, model, state.encoder_cache, state.device) + assert state.encoder_runner.has_cudagraph() is expected + + +def test_encoder_graph_unsupported_modality_uses_eager_output(): + """An image graph must not intercept another modality's encoder output.""" + from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager + + runner = _make_runner([], []) + output = torch.ones(2, HIDDEN) + runner.model = MagicMock() + runner.model.embed_multimodal.return_value = [output] + manager = MagicMock(spec=EncoderCudaGraphManager) + manager.config = SimpleNamespace(modalities=["image"]) + manager.is_captured.return_value = True + manager.supports_modality.side_effect = lambda modality: ( + EncoderCudaGraphManager.supports_modality(manager, modality) + ) + runner.cudagraph_manager = manager + with patch( + "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs", + return_value=[("audio", 1, {"audio_values": torch.zeros(2)})], + ): + result = runner.execute_mm_encoder([]) + assert len(result) == 1 + assert result[0] is output + manager.execute.assert_not_called() + + def _model_state(cache: EncoderCache) -> MagicMock: """A mock ModelState backed by a real EncoderCache.""" state = MagicMock() diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index e19f1b332a68..4dc0bf77dd01 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/vllm/v1/worker/mm_encoder_model_runner.py @@ -4,9 +4,10 @@ An encoder-only instance -- `--mm-encoder-only`, or the producer side of encoder-cache disaggregation -- encodes the multi-modal items and publishes the -embeddings. It runs no language model: no KV cache, no sampler, no CUDA graphs. +embeddings. It runs no language model: no KV cache, sampler, or decoder graphs. """ +import time from collections.abc import Iterator from contextlib import contextmanager from typing import TYPE_CHECKING, Any @@ -15,7 +16,9 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.sequence import IntermediateTensors +from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.v1.kv_cache_interface import KVCacheSpec from vllm.v1.outputs import ( ModelRunnerOutput, @@ -24,10 +27,13 @@ from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu.model_runner import GPUModelRunner +from vllm.v1.worker.workspace import lock_workspace if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput +logger = init_logger(__name__) + class MMEncoderModelRunner(GPUModelRunner): """Encoder-only variant of the V2 GPU model runner.""" @@ -67,8 +73,31 @@ def input_tensor_semaphore(self) -> Iterator[None]: def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: return {} + @torch.inference_mode() def capture_model(self, *, profile_only: bool = False) -> int: - return 0 + encoder_runner = self.model_state.encoder_runner + if not encoder_runner.has_cudagraph(): + return 0 + + start_time = time.perf_counter() + with freeze_gc_for_cudagraph_capture(): + torch.accelerator.empty_cache() + start_free_memory = torch.accelerator.get_memory_info()[0] + encoder_runner.capture() + graph_memory = start_free_memory - torch.accelerator.get_memory_info()[0] + + if not profile_only: + lock_workspace() + + manager = encoder_runner.cudagraph_manager + assert manager is not None + logger.info( + "Encoder graph capturing finished in %.0f secs, took %.2f GiB (%d graphs)", + time.perf_counter() - start_time, + graph_memory / (1 << 30), + manager.get_num_graphs_to_capture(), + ) + return graph_memory def needs_cudagraph_capture(self) -> bool: return False @@ -114,7 +143,7 @@ def execute_model( batch_req_state, _ = self.gather_batch_req_state(scheduler_output, False) assert batch_req_state is not None - # No CUDA graph, and no DP peer to agree a padded shape with. + # No decoder graph or DP padding; encoder graphs run independently. batch_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=batch_req_state.num_tokens, From 88afb777008bab6d33a6f5ea965e2a0d83a0a7ec Mon Sep 17 00:00:00 2001 From: coderfornow <39123181+coderfornow@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:05:17 +0530 Subject: [PATCH 386/774] [CPU][s390x] Pin protobuf to 7.36.1 and drop C++ extension removal workaround (#54978) Signed-off-by: coderfornow Co-authored-by: Li, Jiang --- docker/Dockerfile.s390x | 6 ------ docs/getting_started/installation/cpu.s390x.inc.md | 13 ------------- requirements/common.txt | 3 ++- 3 files changed, 2 insertions(+), 20 deletions(-) diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index 9517ff371850..6faba540a92b 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -169,8 +169,6 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/ ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH" ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -# Force pure Python protobuf to avoid s390x C++ extension crashes -ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python COPY . /workspace/vllm WORKDIR /workspace/vllm @@ -214,10 +212,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ uv pip install "$(echo dist/*.whl)[tensorizer]" -# Remove protobuf C++ extension that crashes on s390x -RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \ - /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true - # setup non-root user for vllm RUN umask 002 && \ /usr/sbin/useradd --uid 2000 --gid 0 vllm && \ diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 88767e4c6ed2..e1ff4ee3cb0d 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -138,19 +138,6 @@ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ The Docker image (`Dockerfile.s390x`) already includes TCMalloc and sets `LD_PRELOAD` automatically. -!!! warning "Protobuf workaround for s390x" - The C++ protobuf extension crashes on s390x. After installation, set the - following environment variable and remove the C++ extensions: - - ```bash - export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python - - # Remove C++ protobuf extensions that crash on s390x - SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])") - rm -rf "$SITE_PKGS/google/_upb/"*.so \ - "$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true - ``` - --8<-- [end:build-wheel-from-source] --8<-- [start:pre-built-images] diff --git a/requirements/common.txt b/requirements/common.txt index acfc986c7533..ff691ebe47ac 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -11,7 +11,8 @@ transformers >= 5.10.4 huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export tokenizers >= 0.21.1 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 -protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 +protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x" +protobuf == 7.36.1; platform_machine == "s390x" fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides. starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1 aiohttp >= 3.13.3 From 23cd0346cec1cb98862d93e0c1bfa490f827e6a3 Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Thu, 17 Sep 2026 01:43:08 -0700 Subject: [PATCH 387/774] [CPU] Align scheduler and NIXL CPU affinity per local rank (#53636) Signed-off-by: Li, Tianmu Co-authored-by: OpenAI Codex Co-authored-by: Li, Jiang --- docs/getting_started/installation/cpu.md | 4 ++-- vllm/platforms/cpu.py | 3 ++- vllm/utils/ompmultiprocessing.py | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 539c28d7097b..53b8eb596d28 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -150,7 +150,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu - `VLLM_CPU_KVCACHE_SPACE`: specify the KV Cache size (e.g, `VLLM_CPU_KVCACHE_SPACE=40` means 40 GiB space for KV cache), larger setting will allow vLLM to run more requests in parallel. This parameter should be set based on the hardware configuration and memory management pattern of users. Default value is `0`. - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable. -- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`. +- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the total number of CPU cores which are not dedicated to OpenMP threads. The variable only takes effect when `VLLM_CPU_OMP_THREADS_BIND` is set to `auto`. By default, x86, ARM, and RISC-V reserve `local_world_size` CPUs in total, while PowerPC and S390X retain a one-CPU default. When KV transfer is enabled, one additional CPU is reserved per local rank. Reserved CPUs are distributed across the local ranks. - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence. - `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations). @@ -253,7 +253,7 @@ See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, support ### How to decide `VLLM_CPU_OMP_THREADS_BIND`? -- Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and 1 CPU per rank will be reserved for other vLLM components when `world_size > 1`. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following. +- Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and the default reserved CPUs will be kept for other vLLM components. On x86, ARM, and RISC-V, `local_world_size` CPUs are reserved in total, plus one additional CPU per local rank when KV transfer is enabled. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following. - On a hyper-threading enabled platform with 16 logical CPU cores / 8 physical CPU cores: diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 6cc1475b441d..390c0dbfede7 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -515,6 +515,7 @@ def discover_numa_topology(cls) -> list[list[int]]: if not (os.path.exists(SYS_NODE) and os.path.exists(SYS_CPU)): return [] + use_highest_sibling = cls.get_cpu_architecture() == CpuArchEnum.X86 core_rsv_for_kv = [] for node in os.listdir(SYS_NODE): if not node.startswith("node") or not node[4:].isdigit(): @@ -547,7 +548,7 @@ def discover_numa_topology(cls) -> list[list[int]]: else: siblings = [cpu_id] - phys = min(siblings) + phys = max(siblings) if use_highest_sibling else min(siblings) if phys not in seen_phys: seen_phys.add(phys) diff --git a/vllm/utils/ompmultiprocessing.py b/vllm/utils/ompmultiprocessing.py index 57f0016b2260..34395489fb44 100644 --- a/vllm/utils/ompmultiprocessing.py +++ b/vllm/utils/ompmultiprocessing.py @@ -46,12 +46,12 @@ def __init__(self, config: "VllmConfig"): self.reserve_cpu_num = ( self.local_world_size if current_platform.get_cpu_architecture() - in (CpuArchEnum.ARM, CpuArchEnum.RISCV) + in (CpuArchEnum.X86, CpuArchEnum.ARM, CpuArchEnum.RISCV) else 1 ) # reserve at one more core for nixl_connector under p/d case if config.kv_transfer_config: - self.reserve_cpu_num += 1 + self.reserve_cpu_num += self.local_world_size if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None: if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU: From f1c2f6ada80f7c0d896b7870f53db97103ff339d Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 17 Sep 2026 03:44:11 -0500 Subject: [PATCH 388/774] [ROCm][CI] Stage H gating and MI355 test reallocation (#57080) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex --- .buildkite/test-amd.yaml | 1714 ++++++++++++----- .buildkite/test_areas/compile.yaml | 4 +- .buildkite/test_areas/cuda.yaml | 4 +- .buildkite/test_areas/disaggregated.yaml | 14 +- .buildkite/test_areas/distributed.yaml | 58 +- .buildkite/test_areas/e2e_integration.yaml | 2 +- .buildkite/test_areas/engine.yaml | 4 + .buildkite/test_areas/entrypoints.yaml | 24 +- .buildkite/test_areas/kernels.yaml | 189 +- .buildkite/test_areas/lm_eval.yaml | 28 +- .buildkite/test_areas/lora.yaml | 18 +- .buildkite/test_areas/misc.yaml | 43 +- .buildkite/test_areas/models_basic.yaml | 21 +- .buildkite/test_areas/models_distributed.yaml | 8 +- .buildkite/test_areas/models_language.yaml | 13 +- .buildkite/test_areas/models_multimodal.yaml | 7 +- .buildkite/test_areas/plugins.yaml | 14 + .buildkite/test_areas/pytorch.yaml | 16 + .buildkite/test_areas/quantization.yaml | 44 + .buildkite/test_areas/ray_compat.yaml | 16 + .buildkite/test_areas/rust_frontend.yaml | 4 +- .buildkite/test_areas/spec_decode.yaml | 42 +- .buildkite/test_areas/weight_loading.yaml | 2 +- 23 files changed, 1719 insertions(+), 570 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 4b72256a3fcd..d038d661769f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -124,8 +124,9 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/transformers/fusers/ + - tests/models/test_qwen3_5_mtp_config.py commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ - label: ":computer: (CPU) Multimodal Processor Shard %N" # TBD timeout_in_minutes: 90 @@ -214,7 +215,7 @@ steps: - .buildkite/test_areas/rust_frontend_cargo.yaml - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy - label: ":computer: (CPU) Rust Frontend Cargo" # TBD timeout_in_minutes: 35 @@ -229,7 +230,7 @@ steps: - .buildkite/test_areas/rust_frontend_cargo.yaml - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test - label: ":computer: (CPU) Docker Build Metadata" # TBD timeout_in_minutes: 35 @@ -281,7 +282,7 @@ steps: #-------------------------------------------------------- mi250 · distributed --------------------------------------------------------# - label: ":amd: (MI250) Pipeline + Context Parallelism" # TBD - timeout_in_minutes: 80 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_4 num_gpus: 4 @@ -297,14 +298,18 @@ steps: - tests/distributed/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/v1/distributed/test_pp_dp_v2.py + optional: true commands: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py + - pytest -v -s v1/distributed/test_pp_dp_v2.py #---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------# -- label: ":amd: (MI250) Helion Kernels" # TBD - timeout_in_minutes: 40 +- label: ":amd: (MI250) Helion Kernels Shard %N" + timeout_in_minutes: 45 + parallelism: 2 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true @@ -316,25 +321,43 @@ steps: - vllm/platforms/rocm.py commands: - pip install helion==1.4.0 - - pytest -v -s kernels/helion/ + - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT #----------------------------------------------------- mi250 · models / language -----------------------------------------------------# - label: ":amd: (MI250) Language Models (PPL)" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - tests/models/language/generation_ppl_test + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/transformers_utils/configs/qwen3_5_moe.py + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/model_executor/layers/fla/ops/ + - vllm/_aiter_ops.py + - vllm/v1/attention/backends/triton_attn.py + - vllm/v1/attention/backends/rocm_attn.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/v1/attention/backends/flex_attention.py + - vllm/v1/attention/ops/ + - vllm/platforms/rocm.py + optional: true commands: - pytest -v -s models/language/generation_ppl_test #---------------------------------------------------- mi250 · models / multimodal ----------------------------------------------------# - label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true @@ -367,7 +390,7 @@ steps: - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - label: ":amd: (MI250) CUDAGraph" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true @@ -378,9 +401,13 @@ steps: - vllm/config/compilation.py - vllm/compilation - vllm/platforms/rocm.py + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py commands: - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - pytest -v -s v1/cudagraph/test_cudagraph_mode.py + - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py + - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py - label: ":amd: (MI250) E2E Core" # TBD timeout_in_minutes: 50 @@ -423,6 +450,7 @@ steps: - vllm/engine/arg_utils.py - vllm/transformers_utils/config.py - vllm/platforms/rocm.py + - vllm/v1/engine/ commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py @@ -447,7 +475,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/draft_model/ - label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD - timeout_in_minutes: 65 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true @@ -461,6 +489,7 @@ steps: - vllm/transformers_utils/configs/speculators/ - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py + - vllm/v1/attention/backends/ commands: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ @@ -468,7 +497,7 @@ steps: #------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# - label: ":amd: (MI300) Python-only Installation" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -501,14 +530,16 @@ steps: - tests/basic_correctness/test_basic_correctness - tests/basic_correctness/test_cpu_offload - tests/basic_correctness/test_mem.py + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s basic_correctness/test_mem.py - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py - pytest -v -s basic_correctness/test_cpu_offload.py -- label: ":amd: (MI300) Distributed Models" # TBD +- label: ":amd: (MI300) Distributed Models Shard %N" timeout_in_minutes: 80 + parallelism: 3 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -531,21 +562,17 @@ steps: - tests/basic_correctness/ - tests/model_executor/model_loader/test_sharded_state_loader.py - tests/models/ + - vllm/v1/worker/ + - tests/conftest.py + - tests/utils.py + - tests/models/test_vision.py commands: - - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' - - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' + - case "$$BUILDKITE_PARALLEL_JOB" in 0) TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; 1) pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && pytest models/language -v -s -m 'distributed(num_gpus=2)' && VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; 2) pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py && pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' && pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' ;; *) exit 2 ;; esac #-------------------------------------------------------- mi300 · benchmarks ---------------------------------------------------------# - label: ":amd: (MI300) Benchmarks CLI" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -554,13 +581,14 @@ steps: source_file_dependencies: - vllm/ - tests/benchmarks/ + - rust/src/bench/tests/python_serve_flags.txt commands: - pytest -v -s benchmarks/ #---------------------------------------------------------- mi300 · compile ----------------------------------------------------------# - label: ":amd: (MI300) PyTorch Compilation" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -576,6 +604,29 @@ steps: - csrc/ - tests/compile - vllm/platforms/rocm.py + - vllm/__init__.py + - vllm/_aiter_ops.py + - vllm/_custom_ops.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/env_override.py + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/ir/ + - vllm/kernels/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/plugins/ + - vllm/sampling_params.py + - vllm/sequence.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ commands: - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" @@ -596,6 +647,7 @@ steps: - tests/compile/fusions_e2e/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + optional: true commands: - rocm-smi - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" @@ -615,6 +667,7 @@ steps: - tests/compile/fusions_e2e/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + optional: true commands: - rocm-smi # Run all models and attn backends but only Inductor partition and native custom ops @@ -631,11 +684,37 @@ steps: source_file_dependencies: - vllm/ - tests/compile/passes + - vllm/platforms/rocm.py + - vllm/__init__.py + - vllm/_aiter_ops.py + - vllm/_custom_ops.py + - vllm/compilation/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/env_override.py + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/ir/ + - vllm/kernels/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/plugins/ + - vllm/sampling_params.py + - vllm/sequence.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ + optional: true commands: - pytest -s -v compile/passes --ignore compile/passes/distributed - label: ":amd: (MI300) PyTorch Fullgraph" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -649,11 +728,33 @@ steps: - csrc/ - tests/compile - vllm/platforms/rocm.py + - vllm/__init__.py + - vllm/_aiter_ops.py + - vllm/_custom_ops.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/env_override.py + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/ir/ + - vllm/kernels/ + - vllm/logger.py + - vllm/multimodal/ + - vllm/platforms/ + - vllm/plugins/ + - vllm/sampling_params.py + - vllm/sequence.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ commands: - - pytest -v -s compile/fullgraph/test_full_graph.py + - find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}' - label: ":amd: (MI300) PyTorch Nightly Dependency Override Check" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -704,7 +805,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition" - label: ":amd: (MI300) Distributed Compile + RPC" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -722,9 +823,9 @@ steps: - tests/compile/test_wrapper.py - tests/entrypoints/llm/test_collective_rpc.py - vllm/platforms/rocm.py + optional: true commands: - pytest -v -s entrypoints/llm/test_collective_rpc.py - - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - pytest -v -s ./compile/test_wrapper.py #----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------# @@ -738,9 +839,16 @@ steps: source_file_dependencies: - vllm/ - tests/cuda + - vllm/envs.py + - vllm/logger.py + - vllm/platforms/ + - vllm/plugins/ + - vllm/utils/ + optional: true commands: - pytest -v -s cuda/test_cuda_context.py - pytest -v -s cuda/test_platform_no_cuda_init.py + - pytest -v -s cuda/test_cuda_compatibility_path.py #-------------------------------------------------------- mi300 · detokenizer --------------------------------------------------------# @@ -756,6 +864,19 @@ steps: - tests/detokenizer - tests/multimodal - tests/utils_ + - vllm/assets/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ commands: - pytest -v -s detokenizer - pytest -v -s -m 'not cpu_test' multimodal @@ -782,7 +903,7 @@ steps: - pytest -v -s distributed/test_shm_storage.py - label: ":amd: (MI300) EPLB Algorithm" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -814,7 +935,7 @@ steps: - pytest -v -s distributed/test_eplb_spec_decode.py - label: ":amd: (MI300) Distributed Features" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -831,6 +952,26 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - csrc/custom_quickreduce.cu + - csrc/ops.h + - csrc/torch_bindings.cpp + - vllm/model_executor/layers/ + - vllm/entrypoints/llm.py + - vllm/config/parallel.py + - vllm/v1/engine/ + - vllm/v1/executor/ + - vllm/v1/worker/ + - vllm/_custom_ops.py + - vllm/envs.py + - tests/distributed/test_rocm_aiter_custom_ar.py + - tests/distributed/test_rocm_quick_reduce.py + - tests/distributed/test_quick_all_reduce.py + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py + - tests/utils.py + - examples/rl/rlhf_async_new_apis.py + - tests/distributed/test_weight_transfer.py + - tests/distributed/test_packed_tensor.py + optional: true commands: - pytest -v -s tests/distributed/test_context_parallel.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py @@ -839,6 +980,10 @@ steps: - pytest -v -s tests/v1/distributed/test_dbo.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py - pytest -v -s tests/distributed/test_packed_tensor.py + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/distributed/test_rocm_quick_reduce.py + - pytest -v -s tests/distributed/test_quick_all_reduce.py - label: ":amd: (MI300) Distributed" # TBD timeout_in_minutes: 40 @@ -993,7 +1138,7 @@ steps: #---------------------------------------------------------- mi300 · engine -----------------------------------------------------------# - label: ":amd: (MI300) Engine Core" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1008,6 +1153,17 @@ steps: - tests/test_vllm_port - tests/jit_monitor/test_hooks.py - tests/jit_monitor/test_hooks_gpu.py + - vllm/compilation/ + - vllm/config/ + - vllm/engine/ + - vllm/entrypoints/logger.py + - vllm/envs.py + - vllm/logger.py + - vllm/logging_utils/ + - vllm/platforms/ + - vllm/sequence.py + - vllm/triton_utils/ + - vllm/utils/ commands: - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py @@ -1049,8 +1205,9 @@ steps: - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests -- label: ":amd: (MI300) Entrypoints Integration (API Server)" # TBD +- label: ":amd: (MI300) Entrypoints Integration (API Server) %N" timeout_in_minutes: 65 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1063,12 +1220,12 @@ steps: - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - - pytest -v -s entrypoints/scale_out + - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi + - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1081,10 +1238,10 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness + - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness - label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" # TBD - timeout_in_minutes: 55 + timeout_in_minutes: 70 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1114,7 +1271,6 @@ steps: - tests/entrypoints/cohere - tests/entrypoints/generate commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tool_use - pytest -v -s entrypoints/tool_parsers - pytest -v -s entrypoints/generate @@ -1122,7 +1278,7 @@ steps: - pytest -v -s entrypoints/cohere - label: ":amd: (MI300) Entrypoints Integration (Responses API)" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1133,7 +1289,6 @@ steps: - vllm/ - tests/entrypoints/openai/responses commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/responses - label: ":amd: (MI300) Entrypoints Integration (Speech to Text)" # TBD @@ -1146,6 +1301,7 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/speech_to_text + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text @@ -1181,7 +1337,7 @@ steps: - pytest -v -s entrypoints/pooling - label: ":amd: (MI300) OpenAI API Correctness" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1220,11 +1376,12 @@ steps: - .buildkite/scripts/scheduled_integration_test/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + optional: true commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 - label: ":amd: (MI300) LM Eval Small Models" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1238,11 +1395,12 @@ steps: - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + optional: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt - label: ":amd: (MI300) MRCR Eval Small Models" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1275,11 +1433,12 @@ steps: - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - .buildkite/lm-eval-harness/ commands: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt - label: ":amd: (MI300) Multimodal Accuracy Eval (Small Models)" # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1295,7 +1454,7 @@ steps: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 - label: ":amd: (MI300) GPQA Eval (GPT-OSS)" # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -1317,7 +1476,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt -- label: ":amd: (MI300) LM Eval Small Models FP8 + Mixed" # TBD +- label: ":amd: (MI300) LM Eval Small Models Distributed" timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1333,6 +1492,8 @@ steps: - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt + optional: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt @@ -1517,7 +1678,7 @@ steps: #--------------------------------------------------------- mi300 · examples ----------------------------------------------------------# - label: ":amd: (MI300) Examples" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1529,27 +1690,27 @@ steps: - examples/ - vllm/platforms/rocm.py commands: - - pip install tensorizer # Basic - - python3 basic/offline_inference/chat.py - - python3 basic/offline_inference/generate.py --model facebook/opt-125m - - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 - - python3 basic/offline_inference/classify.py - - python3 basic/offline_inference/embed.py - - python3 basic/offline_inference/score.py # Multi-modal models - - python3 generate/multimodal/audio_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 # Pooling models - - python3 pooling/embed/vision_embedding_offline.py --seed 0 # Features demo - - python3 features/automatic_prefix_caching/prefix_caching_offline.py - - python3 deployment/llm_engine_example.py - - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors - - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 + - pip install --no-deps tensorizer + - python3 basic/offline_inference/chat.py + - python3 basic/offline_inference/generate.py --model facebook/opt-125m + - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 + - python3 basic/offline_inference/classify.py + - python3 basic/offline_inference/embed.py + - python3 basic/offline_inference/score.py + - python3 generate/multimodal/audio_language_offline.py --seed 0 + - python3 generate/multimodal/vision_language_offline.py --seed 0 + - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 + - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 + - python3 pooling/embed/vision_embedding_offline.py --seed 0 + - python3 features/automatic_prefix_caching/prefix_caching_offline.py + - python3 deployment/llm_engine_example.py + - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors + - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 + - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 #---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# @@ -1565,6 +1726,8 @@ steps: - vllm/kernels - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/ir/ + - tests/kernels/ir/ commands: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir @@ -1584,15 +1747,16 @@ steps: - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py + optional: true commands: - - rocm-smi + - amd-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: ":amd: (MI300) Attention Kernels Shard %N" # TBD - timeout_in_minutes: 70 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1607,27 +1771,10 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/utils/flashinfer.py commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI300) Core Operation Kernels Shard %N" # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - parallelism: 3 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - tests/kernels/core - - tests/kernels/test_top_k_per_row.py - - tests/kernels/test_concat_mla_q.py - - vllm/model_executor/layers/rotary_embedding/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1643,6 +1790,12 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py + - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu + - csrc/libtorch_stable/minimax_reduce_rms_kernel.h + - csrc/libtorch_stable/ops.h + - csrc/libtorch_stable/torch_bindings.cpp + - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py + optional: true commands: - pytest -v -s kernels/core/test_minimax_reduce_rms.py @@ -1667,7 +1820,7 @@ steps: - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py - label: ":amd: (MI300) Mamba Kernels" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1682,7 +1835,7 @@ steps: - pytest -v -s kernels/mamba - label: ":amd: (MI300) MoE Kernels Shard %N" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1708,8 +1861,9 @@ steps: - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI300) Quantization Kernels" # TBD - timeout_in_minutes: 55 +- label: ":amd: (MI300) Quantization Kernels Shard %N" + timeout_in_minutes: 120 + parallelism: 6 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1729,8 +1883,10 @@ steps: - vllm/platforms/rocm.py - vllm/model_executor/kernels/ - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/config/ + - tests/kernels/quantization/test_rocm_skinny_gemms.py commands: - - pytest -v -s kernels/quantization + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) DeepEP FP8 MoE Kernels" # TBD timeout_in_minutes: 55 @@ -1754,7 +1910,7 @@ steps: #----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# - label: ":amd: (MI300) LoRA Shard %N" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 85 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -1768,8 +1924,9 @@ steps: commands: - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py -- label: ":amd: (MI300) LoRA TP (Distributed)" # TBD +- label: ":amd: (MI300) LoRA TP (Distributed) %N" timeout_in_minutes: 65 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_4 @@ -1780,14 +1937,17 @@ steps: - vllm/model_executor/layers/fused_moe/ - tests/lora - vllm/platforms/rocm.py + optional: true commands: - - pytest -v -s -x lora/test_chatglm3_tp.py - - pytest -v -s -x lora/test_llama_tp.py - - pytest -v -s -x lora/test_qwen3_with_multi_loras.py - - pytest -v -s -x lora/test_olmoe_tp.py - - pytest -v -s -x lora/test_gptoss_tp.py - - pytest -v -s -x lora/test_qwen35_densemodel_lora.py - - pytest -v -s -x lora/test_gemma4_tp.py + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac + - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_chatglm3_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_llama_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_qwen3_with_multi_loras.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_olmoe_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_gptoss_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_qwen35_densemodel_lora.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_gemma4_tp.py; fi #------------------------------------------------------ mi300 · model_executor -------------------------------------------------------# @@ -1806,189 +1966,50 @@ steps: - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/model_executor/warmup + - tests/model_executor/test_jit_warmup.py + - tests/model_executor/test_jit_warmup_cutedsl_launcher.py + - tests/model_executor/test_jit_warmup_triton_launcher.py commands: - apt-get update && apt-get install -y curl libsodium23 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s model_executor -m '(not slow_test)' - - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py + - export PYTHONFAULTHANDLER=1 + - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread + - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread -#---------------------------------------------------- mi300 · model_runner_v2 -------------------------------------------------------# +#------------------------------------------------------ mi300 · models / basic -------------------------------------------------------# -- label: ":amd: (MI300) Model Runner V2 Core" # TBD - timeout_in_minutes: 55 +- label: ":amd: (MI300) Basic Models (Initialization)" # TBD + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/worker/gpu/ - - vllm/v1/worker/gpu_worker.py - - vllm/v1/core/sched/ - - vllm/v1/attention/ - - tests/v1/engine/test_llm_engine.py - - tests/v1/e2e/ - - tests/entrypoints/llm/test_struct_output_generate.py - - vllm/platforms/rocm.py + - vllm/ + - tests/models/test_initialization.py + - tests/models/registry.py + optional: true commands: - - set -x - - export VLLM_USE_V2_MODEL_RUNNER=1 - - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" - - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" - - pytest -v -s v1/e2e/general/test_context_length.py - - pytest -v -s v1/e2e/general/test_min_tokens.py - - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" + - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset -- label: ":amd: (MI300) Model Runner V2 Examples" # TBD - timeout_in_minutes: 40 +- label: ":amd: (MI300) Basic Models (Other)" # TBD + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 optional: true - working_dir: "/vllm-workspace/examples" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/worker/gpu/ - - vllm/v1/core/sched/ - - vllm/v1/worker/gpu_worker.py - - examples/basic/offline_inference/ - - examples/generate/multimodal/ - - examples/features/ - - examples/pooling/embed/vision_embedding_offline.py - - examples/features/tensorize_vllm_model.py - - examples/deployment/llm_engine_example.py - - vllm/platforms/rocm.py + - vllm/ + - tests/models/test_terratorch.py + - tests/models/transformers/test_backend.py + - tests/models/test_registry.py + - tests/models/test_hyv4_rocm.py + - tests/models/test_vision.py + - tests/models/test_deepseek_v4_vl_rocm.py commands: - - set -x - - export VLLM_USE_V2_MODEL_RUNNER=1 - - pip install tensorizer - - python3 basic/offline_inference/chat.py - - python3 basic/offline_inference/generate.py --model facebook/opt-125m - - python3 generate/multimodal/audio_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 - - python3 pooling/embed/vision_embedding_offline.py --seed 0 - - python3 features/automatic_prefix_caching/prefix_caching_offline.py - - python3 deployment/llm_engine_example.py - - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors - - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 - -- label: ":amd: (MI300) Model Runner V2 Distributed" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/worker/gpu/ - - vllm/v1/worker/gpu_worker.py - - tests/basic_correctness/test_basic_correctness.py - - tests/v1/distributed/test_async_llm_dp.py - - tests/v1/distributed/test_eagle_dp.py - - vllm/platforms/rocm.py - commands: - - set -x - - export VLLM_USE_V2_MODEL_RUNNER=1 - - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True" - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray" - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - -- label: ":amd: (MI300) Model Runner V2 Pipeline Parallelism" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/worker/gpu/ - - vllm/v1/worker/gpu_worker.py - - tests/distributed/test_pipeline_parallel.py - - tests/distributed/test_pp_cudagraph.py - - vllm/platforms/rocm.py - commands: - - set -x - - export VLLM_USE_V2_MODEL_RUNNER=1 - - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray" - - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray" - -- label: ":amd: (MI300) Model Runner V2 Spec Decode" # TBD - timeout_in_minutes: 90 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/worker/gpu/ - - vllm/v1/worker/gpu_worker.py - - tests/v1/spec_decode/test_max_len.py - - tests/v1/spec_decode/test_rejection_sampler_utils.py - - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - set -x - - export VLLM_USE_V2_MODEL_RUNNER=1 - - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp" - - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py - - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - pytest -v -s v1/e2e/spec_decode/eagle/ - - pytest -v -s v1/e2e/spec_decode/speculators/ - - pytest -v -s v1/e2e/spec_decode/mtp/ - -#------------------------------------------------------ mi300 · models / basic -------------------------------------------------------# - -- label: ":amd: (MI300) Basic Models (Extra Initialization) Shard %N" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - parallelism: 6 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - tests/models/test_initialization.py - - tests/models/registry.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - -- label: ":amd: (MI300) Basic Models (Initialization)" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/test_initialization.py - - tests/models/registry.py - commands: - - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - -- label: ":amd: (MI300) Basic Models (Other)" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/test_terratorch.py - - tests/models/transformers/test_backend.py - - tests/models/test_registry.py - - tests/models/test_hyv4_rocm.py - - tests/models/test_vision.py - commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge @@ -2007,8 +2028,9 @@ steps: commands: - pytest -v -s models/language/pooling_mteb_test -- label: ":amd: (MI300) Language Models (Extended Pooling)" # TBD - timeout_in_minutes: 80 +- label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N" + timeout_in_minutes: 95 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2018,10 +2040,10 @@ steps: - vllm/ - tests/models/language/pooling commands: - - pytest -v -s models/language/pooling -m 'not core_model' + - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI300) Language Models (Standard)" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2072,22 +2094,6 @@ steps: - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: ":amd: (MI300) Language Models (Hybrid) Shard %N" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - parallelism: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/generation - commands: - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - #---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------# - label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" # TBD @@ -2105,8 +2111,9 @@ steps: - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py -- label: ":amd: (MI300) Multimodal Models (Extended Generation 2)" # TBD - timeout_in_minutes: 65 +- label: ":amd: (MI300) Multimodal Models (Extended Generation 2) Shard %N" + timeout_in_minutes: 100 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2116,11 +2123,11 @@ steps: - vllm/ - tests/models/multimodal/generation commands: - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI300) Multimodal Models (Extended Generation 3)" # TBD - timeout_in_minutes: 55 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2133,7 +2140,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - label: ":amd: (MI300) Multimodal Models (Standard) 1: qwen2" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2147,7 +2154,7 @@ steps: - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - label: ":amd: (MI300) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2157,27 +2164,13 @@ steps: - vllm/ - tests/models/multimodal/generation - tests/models/multimodal/test_mapping.py + - tests/models/multimodal commands: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model -- label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" # TBD - timeout_in_minutes: 70 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - - label: ":amd: (MI300) Multimodal Processor Shard %N" # TBD - timeout_in_minutes: 110 + timeout_in_minutes: 115 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2205,7 +2198,7 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' - label: ":amd: (MI300) Multimodal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2233,7 +2226,9 @@ steps: - vllm/platforms/rocm.py - tests/models/quantization - vllm/model_executor/model_loader/ + optional: true commands: + - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------# @@ -2311,7 +2306,7 @@ steps: #---------------------------------------------------------- mi300 · plugins ----------------------------------------------------------# - label: ":amd: (MI300) Plugin Integration" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -2324,43 +2319,43 @@ steps: - vllm/platforms/rocm.py commands: # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform + # END: platform plugin tests + # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin + # END: `io_processor` plugins test + # BEGIN: `bge_m3_sparse io_processor` test + # END: `bge_m3_sparse io_processor` test + # BEGIN: `colbert_query io_processor` test + # END: `colbert_query io_processor` test + # BEGIN: `stat_logger` plugins test + # END: `stat_logger` plugins test + # BEGIN: `endpoint` plugins test + # END: `endpoint` plugins test + # BEGIN: other tests - pip install -e ./plugins/vllm_add_dummy_platform - pytest -v -s plugins_tests/test_platform_plugins.py - pip uninstall vllm_add_dummy_platform -y - # END: platform plugin tests - # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin + - pytest -v -s ./plugins_tests/test_io_processor_plugins.py - pip install -e ./plugins/prithvi_io_processor_plugin - - pytest -v -s plugins_tests/test_io_processor_plugins.py - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py - pip uninstall prithvi_io_processor_plugin -y - # END: `io_processor` plugins test - # BEGIN: `bge_m3_sparse io_processor` test - pip install -e ./plugins/bge_m3_sparse_plugin - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py - pip uninstall bge_m3_sparse_plugin -y - # END: `bge_m3_sparse io_processor` test - # BEGIN: `colbert_query io_processor` test - pip install -e ./plugins/colbert_query_plugin - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py - pip uninstall colbert_query_plugin -y - # END: `colbert_query io_processor` test - # BEGIN: `stat_logger` plugins test - pip install -e ./plugins/vllm_add_dummy_stat_logger - pytest -v -s plugins_tests/test_stats_logger_plugins.py - pip uninstall dummy_stat_logger -y - # END: `stat_logger` plugins test - # BEGIN: `endpoint` plugins test - pip install -e ./plugins/vllm_add_dummy_endpoint_plugin - pytest -v -s plugins_tests/test_endpoint_plugins.py - pip uninstall vllm_add_dummy_endpoint_plugin -y - # END: `endpoint` plugins test - # BEGIN: other tests - pytest -v -s plugins_tests/test_scheduler_plugins.py - pip install -e ./plugins/vllm_add_dummy_model - pytest -v -s distributed/test_distributed_oot.py - - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process - - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process - - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins + - pytest -v -s plugins_tests/test_oot_registration_online.py + - pytest -v -s plugins_tests/test_oot_registration_offline.py + - pytest -v -s plugins_tests/lora_resolvers - label: ":amd: (MI300) GGUF Plugin" # TBD timeout_in_minutes: 45 @@ -2402,18 +2397,24 @@ steps: - tests/entrypoints/serve/instrumentator/test_uds.py - tests/v1/sample/test_logprobs_e2e.py - vllm/platforms/rocm.py + - tests/entrypoints/openai/chat_completion/test_include_reasoning.py + - tests/entrypoints/openai/chat_completion/test_serving_chat.py + - tests/entrypoints/openai/completion/test_shutdown.py + - tests/entrypoints/openai/test_uds.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming" + - pytest -v -s entrypoints/openai/chat_completion/test_include_reasoning.py - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" + - pytest -v -s entrypoints/openai/chat_completion/test_serving_chat.py -k "test_chat_per_request_metrics_follow_server_flag or test_streaming_reasoning_usage_counts_across_deltas or test_completion_tokens_details" - pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" -- label: ":amd: (MI300) Rust Frontend Serve + Admin Coverage" # TBD +- label: ":amd: (MI300) Rust Frontend Serve/Admin Coverage" timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2509,8 +2510,9 @@ steps: #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# -- label: ":amd: (MI300) Quantization" # TBD +- label: ":amd: (MI300) Quantization Shard %N" timeout_in_minutes: 115 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2524,19 +2526,19 @@ steps: - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - tests/quantization - tests/rocm/test_moe_weight_replay.py + optional: true commands: - # temporary install here since we need nightly, will move to requirements/test.in # after torchao 0.12 release, and pin a working version of torchao nightly here - # since torchao nightly is only compatible with torch nightly currently # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now # we can only upgrade after this is resolved # TODO(jerryzh168): resolve the above comment + - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py - - pytest -v -s rocm/test_moe_weight_replay.py + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi #--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# @@ -2556,6 +2558,7 @@ steps: - tests/conftest.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/entrypoints/generate/beam_search commands: - pytest -v -s samplers @@ -2570,6 +2573,18 @@ steps: source_file_dependencies: - vllm/ - tests/test_regression + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + optional: true commands: - pip install 'modelscope<1.38' - pytest -v -s test_regression.py @@ -2586,6 +2601,8 @@ steps: - requirements/ - setup.py - vllm/platforms/rocm.py + - .buildkite/scripts/check-ray-compatibility.sh + optional: true commands: - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh @@ -2608,7 +2625,7 @@ steps: - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test - label: ":amd: (MI300) E2E Scheduling" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2622,7 +2639,7 @@ steps: - pytest -v -s v1/e2e/general/test_async_scheduling.py - label: ":amd: (MI300) V1 Engine" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2632,12 +2649,14 @@ steps: - vllm/v1/engine/ - tests/v1/engine/ - vllm/platforms/rocm.py + - tests/v1/test_tensor_ipc_queue.py commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/test_tensor_ipc_queue.py - label: ":amd: (MI300) Spec Decode Draft Model" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2654,7 +2673,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: ":amd: (MI300) Spec Decode Eagle" # TBD +- label: ":amd: (MI300) Spec Decode Eagle 1: DeepSeek + Qwen" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2670,7 +2689,26 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode/eagle/ + - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3" + +- label: ":amd: (MI300) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3" + - label: ":amd: (MI300) Spec Decode N-Gram + Suffix" # TBD timeout_in_minutes: 35 @@ -2693,7 +2731,7 @@ steps: - python3 spec_decode/test_custom_proposer.py - label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2708,12 +2746,13 @@ steps: - vllm/transformers_utils/configs/speculators/ - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py + - vllm/v1/attention/backends/ commands: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ -- label: ":amd: (MI300) Speculators Correctness" # TBD - timeout_in_minutes: 35 +- label: ":amd: (MI300) Speculators Correctness Nightly" + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2740,6 +2779,7 @@ steps: - tests/evals/gsm8k/ - tests/v1/spec_decode/test_speculators_correctness.py - vllm/platforms/rocm.py + - vllm/v1/worker/gpu/spec_decode/dflash/ commands: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test @@ -2753,6 +2793,17 @@ steps: source_file_dependencies: - vllm/ - tests/v1/spec_decode + - vllm/config/ + - vllm/distributed/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/models/qwen4_exp/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not slow_test' v1/spec_decode @@ -2763,7 +2814,7 @@ steps: dind: false agent_pool: mi300_2 num_gpus: 2 - working_dir: "/vllm-workspace/" + working_dir: /vllm-workspace source_file_dependencies: - vllm/config/speculative.py - vllm/distributed/kv_transfer/kv_connector/ @@ -2785,12 +2836,15 @@ steps: - vllm/_aiter_ops.py - tests/v1/kv_connector/extract_hidden_states_integration/ - vllm/platforms/rocm.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py + - tests/v1/kv_connector/extract_hidden_states_integration + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration + - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration - label: ":amd: (MI300) V1 Attention Shard %N" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 125 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 @@ -2808,7 +2862,7 @@ steps: commands: - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI300) V1 Core + KV + Metrics" # TBD +- label: ":amd: (MI300) V1 Core" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2826,71 +2880,267 @@ steps: - tests/v1/ec_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + optional: true commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + # - export HSA_NO_SCRATCH_RECLAIM=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/simple_kv_offload - - pytest -v -s v1/worker - - pytest -v -s v1/streaming_input - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - # - export HSA_NO_SCRATCH_RECLAIM=1 - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: ":amd: (MI300) V1 Sample + Logits" # TBD - timeout_in_minutes: 70 +- label: ":amd: (MI300) V1 Executor + Worker" + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" + working_dir: /vllm-workspace/tests source_file_dependencies: - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - -- label: ":amd: (MI300) Distributed DP Basic" # TBD - timeout_in_minutes: 45 + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/simple_kv_offload + - tests/v1/worker + - tests/v1/streaming_input + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/executor + - pytest -v -s v1/worker + - pytest -v -s v1/streaming_input + optional: true + +- label: ":amd: (MI300) V1 KV Offload" + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false - agent_pool: mi300_2 + agent_pool: mi300_1 + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/simple_kv_offload + - tests/v1/worker + - tests/v1/streaming_input + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh + - pytest -v -s v1/kv_offload + - pytest -v -s v1/simple_kv_offload + optional: true + +- label: ":amd: (MI300) V1 KV Connectors Shard %N" + timeout_in_minutes: 60 + parallelism: 4 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/simple_kv_offload + - tests/v1/worker + - tests/v1/streaming_input + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + optional: true + +- label: ":amd: (MI300) V1 Metrics + LM Eval" + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/simple_kv_offload + - tests/v1/worker + - tests/v1/streaming_input + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - export GIT_CONFIG_COUNT=1 + - export GIT_CONFIG_KEY_0=http.version + - export GIT_CONFIG_VALUE_0=HTTP/1.1 + - export GIT_TERMINAL_PROMPT=0 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + optional: true + + +- label: ":amd: (MI300) V1 Sample" + timeout_in_minutes: 70 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 optional: true - num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + - vllm/config/ - vllm/distributed/ - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/v1/distributed - - tests/entrypoints/launchers/api_server/test_multi_api_servers.py - - vllm/platforms/rocm.py + - vllm/inputs/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ commands: - - export NCCL_CUMEM_HOST_ENABLE=0 - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/sample + +- label: ":amd: (MI300) V1 Logits + Oracle" + timeout_in_minutes: 70 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py + - label: ":amd: (MI300) Metrics, Tracing" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -2901,6 +3151,18 @@ steps: - vllm/ - tests/v1/tracing - tests/tracing/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tracing/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ commands: - "pip install \ 'opentelemetry-sdk>=1.26.0' \ @@ -2911,7 +3173,7 @@ steps: - pytest -v -s tracing - label: ":amd: (MI300) V1 E2E" # TBD - timeout_in_minutes: 25 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_2 @@ -2921,31 +3183,32 @@ steps: source_file_dependencies: - vllm/ - tests/v1/e2e + - vllm/compilation/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/logger.py + - vllm/logging_utils/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/e2e/spec_decode/ commands: - >- pytest -v -s v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism -- label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh - - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_4 @@ -2984,38 +3247,6 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py -- label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3089,16 +3320,18 @@ steps: - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh - label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_4 num_gpus: 4 optional: true source_file_dependencies: - - vllm/v1/attention/backends/utils.py - - vllm/v1/worker/gpu_model_runner.py - - tests/v1/e2e/test_hybrid_chunked_prefill.py + - vllm/v1/attention/backends/utils.py + - vllm/v1/worker/gpu_model_runner.py + - tests/v1/e2e/test_hybrid_chunked_prefill.py + - vllm/v1/core/ + working_dir: /vllm-workspace/tests commands: - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py @@ -3114,6 +3347,7 @@ steps: source_file_dependencies: - vllm/ - tests/weight_loading + optional: true commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt @@ -3156,11 +3390,90 @@ steps: #-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# -- label: ":amd: (MI355) Distributed Features" # TBD +- label: ":amd: (MI355) Distributed DP Basic" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 + optional: true + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/v1/distributed + - tests/entrypoints/launchers/api_server/test_multi_api_servers.py + - vllm/platforms/rocm.py + - tests/v1/e2e/general/test_sharded_sampling.py + - tests/v1/e2e/spec_decode/test_sharded_sampling.py + - tests/v1/e2e/spec_decode/utils.py + commands: + - export NCCL_CUMEM_HOST_ENABLE=0 + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py + - pytest -v -s v1/e2e/general/test_sharded_sampling.py + - pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py + +- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh + +- label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 40 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: ":amd: (MI355) Distributed Features" # TBD + timeout_in_minutes: 90 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 num_gpus: 2 optional: true working_dir: "/vllm-workspace/" @@ -3191,18 +3504,28 @@ steps: - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/v1/distributed/test_dbo.py - tests/utils.py + - examples/features/data_parallel/data_parallel_offline.py + - examples/rl/rlhf_async_new_apis.py + - tests/distributed/test_weight_transfer.py + - tests/distributed/test_packed_tensor.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization + - pytest -v -s tests/v1/distributed/test_dbo.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - pytest -v -s tests/distributed/test_packed_tensor.py - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py - - pytest -v -s tests/v1/distributed/test_dbo.py #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# -- label: ":amd: (MI355) Entrypoints Integration (API Server)" # TBD - timeout_in_minutes: 55 +- label: ":amd: (MI355) Entrypoints Integration (API Server) %N" + timeout_in_minutes: 65 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3215,12 +3538,12 @@ steps: - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - - pytest -v -s entrypoints/scale_out + - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi + - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355) Entrypoints Integration (OpenAI API completion)" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3236,7 +3559,7 @@ steps: - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness - label: ":amd: (MI355) Entrypoints Integration (OpenAI API chat_completion)" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 70 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3251,7 +3574,7 @@ steps: - pytest -v -s entrypoints/openai/chat_completion - label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3264,15 +3587,16 @@ steps: - tests/entrypoints/tool_parsers - tests/entrypoints/anthropic - tests/entrypoints/generate + - tests/entrypoints/cohere commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tool_use - pytest -v -s entrypoints/tool_parsers - pytest -v -s entrypoints/generate - pytest -v -s entrypoints/anthropic + - pytest -v -s entrypoints/cohere - label: ":amd: (MI355) Entrypoints Integration (Speech to Text)" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3281,12 +3605,13 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/speech_to_text + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text - label: ":amd: (MI355) Entrypoints Integration (Multimodal)" - timeout_in_minutes: 40 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3295,12 +3620,13 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/multimodal + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal - label: ":amd: (MI355) Entrypoints Integration (Pooling)" # TBD - timeout_in_minutes: 55 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3309,6 +3635,7 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/pooling + optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling @@ -3361,8 +3688,8 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt -- label: ":amd: (MI355) LM Eval Small Models" # TBD - timeout_in_minutes: 45 +- label: ":amd: (MI355) LM Eval Small Models Distributed" + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 @@ -3378,6 +3705,7 @@ steps: - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt @@ -3424,7 +3752,7 @@ steps: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-tp4.txt --tp-size=4 - label: ":amd: (MI355) LM Eval Large Models" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_8 @@ -3449,7 +3777,7 @@ steps: #--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# - label: ":amd: (MI355) Examples" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3459,23 +3787,24 @@ steps: - vllm/multimodal - examples/ - vllm/platforms/rocm.py + optional: true commands: - - pip install tensorizer # Basic + # Multi-modal models + # Pooling models + # Features demo + - pip install --no-deps tensorizer - python3 basic/offline_inference/chat.py - python3 basic/offline_inference/generate.py --model facebook/opt-125m - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 - python3 basic/offline_inference/classify.py - python3 basic/offline_inference/embed.py - python3 basic/offline_inference/score.py - # Multi-modal models - python3 generate/multimodal/audio_language_offline.py --seed 0 - python3 generate/multimodal/vision_language_offline.py --seed 0 - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 - # Pooling models - python3 pooling/embed/vision_embedding_offline.py --seed 0 - # Features demo - python3 features/automatic_prefix_caching/prefix_caching_offline.py - python3 deployment/llm_engine_example.py - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors @@ -3502,8 +3831,10 @@ steps: - vllm/model_executor/layers/rotary_embedding/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/kernels/test_fused_qk_norm_rope_gate.py + optional: true commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355) MiniMax Reduce RMS Kernels" # TBD timeout_in_minutes: 45 @@ -3520,6 +3851,12 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py + - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu + - csrc/libtorch_stable/minimax_reduce_rms_kernel.h + - csrc/libtorch_stable/ops.h + - csrc/libtorch_stable/torch_bindings.cpp + - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py + optional: true commands: - pytest -v -s kernels/core/test_minimax_reduce_rms.py @@ -3543,8 +3880,12 @@ steps: - vllm/v1/attention/selector.py - vllm/platforms/rocm.py - vllm/_aiter_ops.py + - examples/basic/offline_inference/chat.py + - tests/kernels/attention/test_attention_selector.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + optional: true commands: - - rocm-smi + - amd-smi - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN - pytest -v -s tests/kernels/attention/test_attention_selector.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -3564,15 +3905,16 @@ steps: - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py + optional: true commands: - - rocm-smi + - amd-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: ":amd: (MI355) Attention Kernels Shard %N" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3586,11 +3928,13 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/utils/flashinfer.py + optional: true commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355) MoE Kernels Shard %N" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3610,6 +3954,7 @@ steps: - vllm/config - vllm/_aiter_ops.py - vllm/platforms/rocm.py + optional: true commands: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py @@ -3646,11 +3991,15 @@ steps: - vllm/utils/torch_utils.py - vllm/platforms/ - vllm/_aiter_ops.py + - csrc/quantization/cutlass_w8a8/moe/ + - vllm/distributed/device_communicators/ + - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py -- label: ":amd: (MI355) Quantization Kernels" # TBD - timeout_in_minutes: 50 +- label: ":amd: (MI355) Quantization Kernels Shard %N" + timeout_in_minutes: 120 + parallelism: 6 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3669,8 +4018,11 @@ steps: - vllm/platforms/rocm.py - vllm/model_executor/kernels/ - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/config/ + - tests/kernels/quantization/test_rocm_skinny_gemms.py + optional: true commands: - - pytest -v -s kernels/quantization + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355) FP8 MoE Kernels" # TBD timeout_in_minutes: 40 @@ -3685,6 +4037,13 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py - vllm/envs.py + - tests/kernels/moe/test_gpt_oss_triton_kernels.py + - tests/kernels/moe/test_modular_oai_triton_moe.py + - tests/kernels/moe/test_moe.py + - tests/kernels/moe/test_block_int8.py + - tests/kernels/moe/test_triton_moe_no_act_mul.py + - tests/kernels/moe/test_triton_moe_ptpc_fp8.py + optional: true commands: - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py @@ -3694,7 +4053,7 @@ steps: - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py - label: ":amd: (MI355) DeepEP FP8 MoE Kernels" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 @@ -3710,12 +4069,51 @@ steps: - vllm/platforms/rocm.py - vllm/envs.py commands: - - pytest -v -s kernels/moe/test_deepep_moe.py - -#----------------------------------------------------- mi355 · models / language -----------------------------------------------------# + - pytest -v -s kernels/moe/test_deepep_moe.py + +#------------------------------------------------------ mi355 · models / basic -------------------------------------------------------# + +- label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + parallelism: 6 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - tests/models/test_initialization.py + - tests/models/registry.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + optional: true + commands: + - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +#----------------------------------------------------- mi355 · models / language -----------------------------------------------------# + +- label: ":amd: (MI355) Language Models (Hybrid) Shard %N" # TBD + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + parallelism: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation + commands: + - export GIT_CONFIG_COUNT=1 + - export GIT_CONFIG_KEY_0=http.version + - export GIT_CONFIG_VALUE_0=HTTP/1.1 + - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355) Language Models (Extended Generation)" # TBD - timeout_in_minutes: 65 + timeout_in_minutes: 95 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3723,13 +4121,15 @@ steps: source_file_dependencies: - vllm/ - tests/models/language/generation + optional: true commands: - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0' - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: ":amd: (MI355) Language Models (Extended Pooling)" # TBD - timeout_in_minutes: 65 +- label: ":amd: (MI355) Language Models (Extended Pooling) Shard %N" + timeout_in_minutes: 95 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3739,10 +4139,10 @@ steps: - vllm/ - tests/models/language/pooling commands: - - pytest -v -s models/language/pooling -m 'not core_model' + - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355) Language Models (PPL)" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3767,6 +4167,8 @@ steps: - vllm/v1/attention/ops/ - vllm/platforms/rocm.py - tests/models/language/generation_ppl_test + - vllm/model_executor/layers/fla/ops/ + - vllm/ commands: - pytest -v -s models/language/generation_ppl_test @@ -3779,6 +4181,7 @@ steps: source_file_dependencies: - vllm/ - tests/models/language + optional: true commands: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' @@ -3786,7 +4189,7 @@ steps: #---------------------------------------------------- mi355 · models / multimodal ----------------------------------------------------# - label: ":amd: (MI355) Multimodal Models (Extended Generation 1)" # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3801,7 +4204,7 @@ steps: - pytest -v -s models/multimodal/test_mapping.py - label: ":amd: (MI355) Multimodal Models (Extended Generation 3)" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3814,7 +4217,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - label: ":amd: (MI355) Multimodal Models (Extended Pooling)" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3827,7 +4230,7 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' - label: ":amd: (MI355) Multimodal Models (Standard) 1: qwen2" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3841,7 +4244,7 @@ steps: - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - label: ":amd: (MI355) Multimodal Models (Standard) 4: other + whisper" # TBD - timeout_in_minutes: 55 + timeout_in_minutes: 70 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3853,13 +4256,13 @@ steps: commands: - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model + - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model #----------------------------------------------------- mi355 · models / quantized -----------------------------------------------------# - label: ":amd: (MI355) Quantized Models" # TBD - timeout_in_minutes: 50 + timeout_in_minutes: 80 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3870,13 +4273,16 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py - vllm/model_executor/model_loader/ + optional: true commands: + - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# -- label: ":amd: (MI355) Quantization" # TBD - timeout_in_minutes: 80 +- label: ":amd: (MI355) Quantization Shard %N" + timeout_in_minutes: 115 + parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3887,10 +4293,16 @@ steps: - tests/quantization - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/model_executor/layers/fused_moe/oracle/unquantized.py + - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py + - tests/rocm/test_moe_weight_replay.py + optional: true commands: + - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi # - label: Quantized MoE Test (B200-MI355) # TBD # timeout_in_minutes: 180 @@ -3922,7 +4334,7 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# - label: ":amd: (MI355) V1 Attention Shard %N" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 125 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3936,11 +4348,12 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + optional: true commands: - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI355) V1 Core + KV + Metrics" # TBD - timeout_in_minutes: 55 +- label: ":amd: (MI355) V1 Core" + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3956,24 +4369,243 @@ steps: - tests/v1/kv_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/simple_kv_offload + - tests/v1/streaming_input + - tests/v1/ec_connector/unit commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not cpu_test' v1/core + +- label: ":amd: (MI355) V1 Executor + Worker" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/cudagraph + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/simple_kv_offload + - tests/v1/streaming_input + - tests/v1/ec_connector/unit + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - pytest -v -s v1/worker - - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s v1/streaming_input + +- label: ":amd: (MI355) V1 KV Offload" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/cudagraph + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/simple_kv_offload + - tests/v1/streaming_input + - tests/v1/ec_connector/unit + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh + - pytest -v -s v1/kv_offload + - pytest -v -s v1/simple_kv_offload + +- label: ":amd: (MI355) V1 KV Connectors Shard %N" + timeout_in_minutes: 60 + parallelism: 4 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/cudagraph + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/simple_kv_offload + - tests/v1/streaming_input + - tests/v1/ec_connector/unit + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +- label: ":amd: (MI355) V1 Metrics + LM Eval" + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/cudagraph + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/entrypoints/pooling/ + - vllm/inputs/ + - vllm/lora/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/outputs.py + - vllm/platforms/ + - vllm/pooling_params.py + - vllm/profiler/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/simple_kv_offload + - tests/v1/streaming_input + - tests/v1/ec_connector/unit + commands: + - export GIT_CONFIG_COUNT=1 + - export GIT_CONFIG_KEY_0=http.version + - export GIT_CONFIG_VALUE_0=HTTP/1.1 + - export GIT_TERMINAL_PROMPT=0 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not cpu_test' v1/metrics - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: ":amd: (MI355) V1 Sample + Logits" # TBD +- label: ":amd: (MI355) CUDAGraph" timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/cudagraph + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/v1/cudagraph_dispatcher.py + - vllm/config/compilation.py + - vllm/compilation + - vllm/platforms/rocm.py + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py + commands: + - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py + - pytest -v -s v1/cudagraph/test_cudagraph_mode.py + - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py + - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py + + +- label: ":amd: (MI355) V1 Sample" + timeout_in_minutes: 70 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -3982,15 +4614,56 @@ steps: - tests/v1/test_oracle.py - tests/v1/test_request.py - tests/v1/test_outputs.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/sample + +- label: ":amd: (MI355) V1 Logits + Oracle" + timeout_in_minutes: 70 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/logits_processors - pytest -v -s v1/test_oracle.py - pytest -v -s v1/test_request.py - pytest -v -s v1/test_outputs.py + - label: ":amd: (MI355) V1 Spec Decode" # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -3998,11 +4671,24 @@ steps: source_file_dependencies: - vllm/ - tests/v1/spec_decode + - vllm/config/ + - vllm/distributed/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/models/qwen4_exp/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + optional: true commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not slow_test' v1/spec_decode -- label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance" # TBD - timeout_in_minutes: 180 +- label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance Nightly" + timeout_in_minutes: 55 + parallelism: 3 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 @@ -4026,13 +4712,17 @@ steps: - tests/utils.py - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py + - vllm/v1/worker/gpu/spec_decode/ + - tests/v1/e2e/spec_decode/conftest.py commands: - - pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi #------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# - label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 70 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 @@ -4041,6 +4731,7 @@ steps: source_file_dependencies: - vllm/ - tests/weight_loading + optional: true commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt @@ -4070,6 +4761,17 @@ steps: source_file_dependencies: - vllm/ - tests/test_regression + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ commands: - pip install 'modelscope<1.38' - pytest -v -s test_regression.py diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 8eea80068646..219a4ca018e5 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -92,7 +92,7 @@ steps: dind: false device: mi355_1 working_dir: "/vllm-workspace/" - timeout_in_minutes: 25 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: @@ -133,7 +133,7 @@ steps: dind: false device: mi355_1 working_dir: "/vllm-workspace/" - timeout_in_minutes: 30 + timeout_in_minutes: 40 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 230781cd5da8..a031c37ba371 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -23,7 +23,7 @@ steps: dind: false device: mi355_1 working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 20 + timeout_in_minutes: 25 depends_on: - image-build-amd source_file_dependencies: @@ -50,7 +50,7 @@ steps: amd: label: ":amd: (MI250) CUDAGraph" device: mi250_1 - timeout_in_minutes: 45 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 226d7ac94ae7..0fbb32f8f9ff 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -16,9 +16,9 @@ steps: - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: - label: ":amd: (MI300) Distributed NixlConnector PD accuracy" + label: ":amd: (MI355) Distributed NixlConnector PD accuracy" dind: false - device: mi300_4 + device: mi355_4 timeout_in_minutes: 60 depends_on: - image-build-amd @@ -71,9 +71,9 @@ steps: - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: - label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" + label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" dind: false - device: mi300_4 + device: mi355_4 timeout_in_minutes: 40 depends_on: - image-build-amd @@ -130,7 +130,7 @@ steps: label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" dind: false device: mi300_4 - timeout_in_minutes: 55 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -221,9 +221,9 @@ steps: - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh mirror: amd: - label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance" + label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" dind: false - device: mi300_2 + device: mi355_2 timeout_in_minutes: 45 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 3ea26ce89c1d..b83d2893bdae 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -65,9 +65,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py mirror: amd: - label: ":amd: (MI300) Distributed DP Basic" + label: ":amd: (MI355) Distributed DP Basic" dind: false - device: mi300_2 + device: mi355_2 timeout_in_minutes: 45 depends_on: - image-build-amd @@ -125,6 +125,9 @@ steps: - tests/compile/test_wrapper.py - tests/entrypoints/llm/test_collective_rpc.py - vllm/platforms/rocm.py + commands: + - pytest -v -s entrypoints/llm/test_collective_rpc.py + - pytest -v -s ./compile/test_wrapper.py - label: ":nvidia: (L4) Distributed Torchrun + Shutdown" key: distributed-torchrun-shutdown-tests-2-gpus @@ -296,6 +299,57 @@ steps: - pytest -v -s tests/v1/distributed/test_dbo.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py - pytest -v -s tests/distributed/test_packed_tensor.py + mirror: + amd: + label: ':amd: (MI355) Distributed Features' + dind: false + device: mi355_2 + timeout_in_minutes: 90 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/v1/distributed/test_dbo.py + - tests/distributed/test_context_parallel.py + - examples/features/data_parallel/data_parallel_offline.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - csrc/custom_quickreduce.cu + - csrc/ops.h + - csrc/torch_bindings.cpp + - vllm/model_executor/layers/ + - vllm/entrypoints/llm.py + - vllm/config/parallel.py + - vllm/v1/engine/ + - vllm/v1/executor/ + - vllm/v1/worker/ + - vllm/_custom_ops.py + - vllm/envs.py + - tests/distributed/test_rocm_aiter_custom_ar.py + - tests/distributed/test_rocm_quick_reduce.py + - tests/distributed/test_quick_all_reduce.py + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py + - tests/utils.py + - examples/rl/rlhf_async_new_apis.py + - tests/distributed/test_weight_transfer.py + - tests/distributed/test_packed_tensor.py + commands: + - pytest -v -s tests/distributed/test_context_parallel.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization + - pytest -v -s tests/v1/distributed/test_dbo.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - pytest -v -s tests/distributed/test_packed_tensor.py + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/distributed/test_rocm_quick_reduce.py + - pytest -v -s tests/distributed/test_quick_all_reduce.py - label: ":nvidia: (B200) Distributed" key: distributed-tests-2xb200 diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 02ba409e3433..20bec50b5e8c 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -57,7 +57,7 @@ steps: dind: false device: mi355_1 working_dir: "/vllm-workspace" - timeout_in_minutes: 20 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index bb848cd9e5e4..179954740a3f 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -56,6 +56,8 @@ steps: timeout_in_minutes: 45 depends_on: - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) E2E Scheduling" key: e2e-scheduling-1-gpu @@ -73,6 +75,8 @@ steps: timeout_in_minutes: 55 depends_on: - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) E2E Core" device: h200_18gb diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 00091707d5b4..60ac5d932c3a 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -21,9 +21,11 @@ steps: label: ":amd: (MI355) Entrypoints Unit" dind: false device: mi355_1 - timeout_in_minutes: 35 + timeout_in_minutes: 50 depends_on: - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (LLM)" device: h200_35gb @@ -44,7 +46,7 @@ steps: label: ":amd: (MI300) Entrypoints Integration (LLM)" dind: false device: mi300_1 - timeout_in_minutes: 55 + timeout_in_minutes: 90 depends_on: - image-build-amd @@ -176,6 +178,15 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text + mirror: + amd: + label: ":amd: (MI355) Entrypoints Integration (Speech to Text)" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Multimodal)" device: h200_35gb @@ -210,6 +221,15 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling + mirror: + amd: + label: ":amd: (MI355) Entrypoints Integration (Pooling)" + dind: false + device: mi355_1 + timeout_in_minutes: 65 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd - label: ":nvidia: (H200 MIG 18GB) OpenAI API Correctness" key: openai-api-correctness diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index ed4ea67f933f..5a7832d945d2 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -13,6 +13,22 @@ steps: commands: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir + mirror: + amd: + label: ":amd: (MI355) vLLM IR" + dind: false + device: mi355_1 + timeout_in_minutes: 50 + working_dir: "/vllm-workspace/" + depends_on: + - image-build-amd + source_file_dependencies: + - tests/ir/ + - tests/kernels/ir/ + - vllm/ir + - vllm/kernels + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) Core Operation Kernels Shard %N" device: h200_35gb @@ -26,6 +42,29 @@ steps: commands: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 3 + mirror: + amd: + label: ":amd: (MI355) Core Operation Kernels Shard %N" + dind: false + device: mi355_1 + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - tests/kernels/core + - tests/kernels/test_top_k_per_row.py + - tests/kernels/test_concat_mla_q.py + - tests/kernels/test_rocm_fp32_router_gemm.py + - vllm/model_executor/layers/fused_moe/router/gate_linear.py + - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py + - vllm/model_executor/layers/rotary_embedding/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/kernels/test_fused_qk_norm_rope_gate.py + commands: + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (H100) MiniMax Reduce RMS Kernels" key: kernels-minimax-reduce-rms-test-2-gpus @@ -145,6 +184,7 @@ steps: dind: false device: mi300_1 timeout_in_minutes: 90 + parallelism: 6 depends_on: - image-build-amd source_file_dependencies: @@ -208,9 +248,6 @@ steps: label: ":amd: (MI300) Quantization Kernels Shard %N" dind: false device: mi300_1 - # Pin the AMD mirror at its current 2 shards so raising NVIDIA - # parallelism does not multiply AMD copies (needs ci-infra #473). - parallelism: 2 timeout_in_minutes: 120 source_file_dependencies: - csrc/quantization/ @@ -222,6 +259,12 @@ steps: - vllm/kernels/aiter_ops.py - vllm/platforms/rocm.py - vllm/model_executor/kernels/ + - csrc/rocm/ + - vllm/v1/attention/backends/rocm_aiter_fa.py + - tests/kernels/quant_utils.py + - tests/kernels/utils.py + - vllm/_custom_ops.py + - vllm/envs.py depends_on: - image-build-amd @@ -258,6 +301,10 @@ steps: - vllm/config - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - csrc/rocm/ + - tests/kernels/utils.py + - vllm/model_executor/layers/quantization/mxfp4.py + - vllm/model_executor/layers/quantization/utils/quant_utils.py depends_on: - image-build-amd @@ -278,6 +325,8 @@ steps: timeout_in_minutes: 55 depends_on: - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (H100) DeepGEMM Kernels" key: kernels-deepgemm-test-h100 @@ -463,6 +512,9 @@ steps: timeout_in_minutes: 45 depends_on: - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py + - scripts/benchmark_helion_kernels.py - label: ":nvidia: (H100) FP8 MoE Kernels" @@ -491,6 +543,25 @@ steps: commands: - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py - pytest -v -s kernels/moe/test_deepep_moe.py + mirror: + amd: + label: ":amd: (MI355) DeepEP FP8 MoE Kernels" + dind: false + device: mi355_2 + timeout_in_minutes: 55 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/moe/ + - csrc/quantization/w8a8/cutlass/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py + commands: + - pytest -v -s kernels/moe/test_deepep_moe.py - label: ":nvidia: (B200) FP4 MoE Kernels" key: kernels-fp4-moe-test-b200 @@ -520,7 +591,30 @@ steps: - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py - + mirror: + amd: + label: ":amd: (MI355) FusedMoE Layer Kernels" + dind: false + device: mi355_2 + timeout_in_minutes: 35 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/ + - vllm/config/ + - vllm/forward_context.py + - vllm/v1/worker/workspace.py + - vllm/utils/import_utils.py + - vllm/utils/math_utils.py + - vllm/utils/torch_utils.py + - vllm/platforms/ + - vllm/_aiter_ops.py - label: ":nvidia: (B200) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-b200s @@ -538,3 +632,90 @@ steps: commands: - pytest -v -s kernels/moe/test_moe_layer.py - pytest -v -s kernels/moe/test_deepep_v2_moe.py + +- label: ":amd: (MI355) MLA Kernels" + key: amd-mla-kernels-mi355 + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 25 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + - tests/kernels/attention/test_rocm_aiter_mla_decode.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py + - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py + commands: + - amd-smi + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py + +- label: ":amd: (MI355) FP8 MoE Kernels" + key: amd-fp8-moe-kernels-mi355 + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py + - tests/kernels/moe/test_gpt_oss_triton_kernels.py + - tests/kernels/moe/test_modular_oai_triton_moe.py + - tests/kernels/moe/test_moe.py + - tests/kernels/moe/test_block_int8.py + - tests/kernels/moe/test_triton_moe_no_act_mul.py + - tests/kernels/moe/test_triton_moe_ptpc_fp8.py + commands: + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py + - pytest -v -s kernels/moe/test_moe.py + - pytest -v -s kernels/moe/test_block_int8.py + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py + +- label: ":amd: (MI355) Kernels" + key: amd-kernels-mi355 + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 25 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/quantization/fp4/ + - csrc/attention/mla/ + - csrc/quantization/cutlass_w8a8/moe/ + - vllm/model_executor/layers/fused_moe/cutlass_moe.py + - vllm/v1/attention/backends/triton_attn.py + - vllm/v1/attention/backends/rocm_attn.py + - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/v1/attention/backends/mla/aiter_triton_mla.py + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + - examples/basic/offline_inference/chat.py + - tests/kernels/attention/test_attention_selector.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + commands: + - amd-smi + - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN + - pytest -v -s tests/kernels/attention/test_attention_selector.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index ec45f7e69e73..e94cbe5dfc15 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -30,6 +30,27 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py +- label: ":amd: (MI300) LM Eval Small Models Harness" + key: amd-lm-eval-small-models-harness + dind: false + device: mi300_1 + timeout_in_minutes: 35 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/.buildkite/lm-eval-harness + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - .buildkite/lm-eval-harness/ + commands: + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt + - label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking" device: h200_35gb key: lm-eval-watermarking @@ -100,7 +121,7 @@ steps: label: ":amd: (MI355) LM Eval Small Models Distributed" dind: false device: mi355_2 - timeout_in_minutes: 45 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: @@ -286,13 +307,16 @@ steps: label: ":amd: (MI300) LM Eval Large Models" dind: false device: mi300_8 - timeout_in_minutes: 40 + timeout_in_minutes: 65 depends_on: - image-build-amd commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt + source_file_dependencies: + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) MoE Refactor Integration TEMPORARY" key: moe-refactor-integration-test-h100-temporary diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 93fb5f5b81f0..28f695f391a9 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -61,7 +61,7 @@ steps: dind: false device: mi300_4 working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 25 + timeout_in_minutes: 65 depends_on: - image-build-amd source_file_dependencies: @@ -73,11 +73,11 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # expandable_segments breaks custom all-reduce IPC on ROCm. - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac - - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_chatglm3_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_llama_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_qwen3_with_multi_loras.py - - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_olmoe_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp + - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_chatglm3_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_llama_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_qwen3_with_multi_loras.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_olmoe_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_gptoss_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_qwen35_densemodel_lora.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_gemma4_tp.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp; fi diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 940447bf1b52..f17d154fbe98 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -99,7 +99,7 @@ steps: label: ":amd: (MI300) V1 Core" dind: false device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd @@ -119,7 +119,7 @@ steps: label: ":amd: (MI300) V1 Executor + Worker" dind: false device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd @@ -138,7 +138,7 @@ steps: label: ":amd: (MI300) V1 KV Offload" dind: false device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd @@ -158,7 +158,7 @@ steps: label: ":amd: (MI300) V1 KV Connectors Shard %N" dind: false device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd @@ -258,7 +258,7 @@ steps: label: ":amd: (MI355) Extract Hidden States Integration" dind: false device: mi355_2 - timeout_in_minutes: 20 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: @@ -305,6 +305,15 @@ steps: - pip install 'modelscope<1.38' - pytest -v -s test_regression.py working_dir: "/vllm-workspace/tests" # optional + mirror: + amd: + label: ":amd: (MI355) Regression" + dind: false + device: mi355_1 + timeout_in_minutes: 25 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/tests - label: ":nvidia: (H200 MIG 35GB) Examples" device: h200_35gb @@ -446,9 +455,15 @@ steps: label: ":amd: (MI355) Async Engine, Inputs, Utils, Worker" dind: false device: mi355_1 - timeout_in_minutes: 40 + timeout_in_minutes: 50 depends_on: - image-build-amd + commands: + - pytest -v -s detokenizer + - pytest -v -s -m 'not cpu_test' multimodal + # test_cache.py marks every test cpu_test, including this GPU regression. + - pytest -v -s multimodal/test_cache.py::test_sleep_wake_preserves_mm_cache_consistency + - pytest -v -s utils_ - label: ":computer: (CPU) Params, Env, Tokenizers, Parser" key: cpu-params-env-tokenizers-parser @@ -572,7 +587,7 @@ steps: amd: label: ":amd: (MI250) Batch Invariance" device: mi250_1 - timeout_in_minutes: 30 + timeout_in_minutes: 35 depends_on: - image-build-amd source_file_dependencies: @@ -623,3 +638,17 @@ steps: commands: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test + mirror: + amd: + label: ":amd: (MI355) Acceptance Length (Large Models)" + dind: false + device: mi355_1 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/model_executor/models/mlp_speculator.py + - tests/v1/spec_decode/test_acceptance_length.py + - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 870cc9b393de..94b0723f8aab 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -19,7 +19,7 @@ steps: label: ":amd: (MI355) Basic Models (Initialization)" dind: false device: mi355_1 - timeout_in_minutes: 25 + timeout_in_minutes: 30 depends_on: - image-build-amd @@ -37,6 +37,23 @@ steps: # test.) Also run if model initialization test file is modified - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 14 + mirror: + amd: + label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + parallelism: 10 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - tests/models/test_initialization.py + - tests/models/registry.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) Basic Models (Other)" device: h200_35gb @@ -162,6 +179,6 @@ steps: device: mi250_1 no_gpu: true working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 25 + timeout_in_minutes: 35 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index f735a9b01871..a054c00db860 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -38,7 +38,7 @@ steps: label: ":amd: (MI300) Distributed Models Shard %N" dind: false device: mi300_2 - timeout_in_minutes: 35 + timeout_in_minutes: 80 depends_on: - image-build-amd source_file_dependencies: @@ -63,7 +63,7 @@ steps: - | case "$$BUILDKITE_PARALLEL_JOB" in 0) - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; 1) @@ -72,8 +72,8 @@ steps: VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; 2) - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' + pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py && + pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' ;; *) exit 2 ;; diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 68f1cdd429ce..3e5024fe8741 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -19,7 +19,7 @@ steps: label: ":amd: (MI300) Language Models (Standard)" dind: false device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd @@ -82,10 +82,10 @@ steps: parallelism: 2 mirror: amd: - label: ":amd: (MI300) Language Models (Hybrid) Shard %N" + label: ":amd: (MI355) Language Models (Hybrid) Shard %N" dind: false - device: mi300_1 - timeout_in_minutes: 60 + device: mi355_1 + timeout_in_minutes: 65 depends_on: - image-build-amd commands: @@ -133,7 +133,7 @@ steps: label: ":amd: (MI355) Language Models (Extended Generation)" dind: false device: mi355_1 - timeout_in_minutes: 70 + timeout_in_minutes: 95 depends_on: - image-build-amd commands: @@ -197,11 +197,8 @@ steps: dind: false device: mi300_1 timeout_in_minutes: 95 - parallelism: 2 depends_on: - image-build-amd - commands: - - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":nvidia: (H200 MIG 18GB) Language Models (MTEB)" key: language-models-test-mteb diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index b09e8cc2d907..e8c5ddd2cb27 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -77,9 +77,9 @@ steps: - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: - label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" + label: ":amd: (MI355) Multimodal Models (Standard) 4: other + whisper" dind: false - device: mi300_1 + device: mi355_1 timeout_in_minutes: 70 depends_on: - image-build-amd @@ -104,7 +104,8 @@ steps: device: mi250_1 no_gpu: true working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 80 + timeout_in_minutes: 90 + parallelism: 6 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 65e7a0b258a0..11f63ff52dfd 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -75,6 +75,20 @@ steps: commands: - pip install "vllm-gguf-plugin >= 0.0.2" - pytest -v -s plugins_tests/gguf + mirror: + amd: + label: ":amd: (MI355) GGUF Plugin" + dind: false + device: mi355_1 + timeout_in_minutes: 45 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/layers/quantization + - tests/plugins_tests/test_gguf_plugin.py + - tests/plugins_tests/gguf + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 18GB) BitsAndBytes Plugin" key: bitsandbytes-plugin diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index f34ff50e672d..459f7faeb4a1 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -122,6 +122,21 @@ steps: - tests/compile/passes commands: - pytest -s -v compile/passes --ignore compile/passes/distributed + mirror: + amd: + label: ':amd: (MI300) PyTorch Compilation Passes' + dind: false + device: mi300_1 + timeout_in_minutes: 80 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/ + - tests/compile/passes + - vllm/platforms/rocm.py + commands: + - pytest -s -v compile/passes --ignore compile/passes/distributed - label: ":nvidia: (H200 MIG 35GB) PyTorch Fullgraph" device: h200_35gb @@ -154,6 +169,7 @@ steps: - vllm/utils/ - vllm/v1/ - tests/compile + - csrc/ commands: # Run smoke tests under fullgraph directory, except test_full_cudagraph.py # as it is a heavy test that is covered in other steps. diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 1d895e1f2f90..702bbe702452 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -20,6 +20,31 @@ steps: # The SM90-only checkpoint currently contains a removed weight_chan_scale # parameter. It was not exercised by the previous L4 job. - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + mirror: + amd: + label: ":amd: (MI355) Quantization Shard %N" + dind: false + device: mi355_1 + timeout_in_minutes: 115 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/quantization + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/layers/fused_moe/oracle/unquantized.py + - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py + - tests/rocm/test_moe_weight_replay.py + commands: + # Use the ROCm default instead of the NVIDIA runner override. + - unset VLLM_USE_V2_MODEL_RUNNER + - uv pip install --system torchao==0.17.0 + - uv pip install --system conch-triton-kernels + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions" device: h200_35gb @@ -63,3 +88,22 @@ steps: - tests/models/quantization commands: - pytest -v -s models/quantization + mirror: + amd: + label: ":amd: (MI355) Quantized Models" + dind: false + device: mi355_1 + timeout_in_minutes: 80 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/layers/quantization + - tests/models/quantization + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ + commands: + # Use the ROCm default instead of the NVIDIA runner override. + - unset VLLM_USE_V2_MODEL_RUNNER + - pytest -v -s models/quantization diff --git a/.buildkite/test_areas/ray_compat.yaml b/.buildkite/test_areas/ray_compat.yaml index d24c38430b8d..dc8571b26045 100644 --- a/.buildkite/test_areas/ray_compat.yaml +++ b/.buildkite/test_areas/ray_compat.yaml @@ -16,3 +16,19 @@ steps: - setup.py commands: - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh + mirror: + amd: + label: ":amd: (MI355) Ray Dependency Compatibility Check" + dind: false + device: mi355_1 + timeout_in_minutes: 25 + working_dir: "/" + depends_on: + - image-build-amd + source_file_dependencies: + - .buildkite/scripts/check-ray-compatibility.sh + - requirements/ + - setup.py + - vllm/platforms/rocm.py + commands: + - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 770d6b64df22..a679b405ed66 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -136,7 +136,7 @@ steps: label: ":amd: (MI355) Rust Frontend Core Correctness" dind: false device: mi355_1 - timeout_in_minutes: 20 + timeout_in_minutes: 25 depends_on: - image-build-amd source_file_dependencies: @@ -166,7 +166,7 @@ steps: label: ":amd: (MI355) Rust Frontend Tool Use" dind: false device: mi355_1 - timeout_in_minutes: 25 + timeout_in_minutes: 30 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 838727035980..c06642ea1a13 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -47,7 +47,7 @@ steps: label: ":amd: (MI300) Spec Decode Eagle 1: DeepSeek + Qwen" dind: false device: mi300_1 - timeout_in_minutes: 35 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -74,7 +74,7 @@ steps: label: ":amd: (MI300) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" dind: false device: mi300_1 - timeout_in_minutes: 25 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -187,7 +187,7 @@ steps: label: ":amd: (MI300) Spec Decode Draft Model" dind: false device: mi300_1 - timeout_in_minutes: 55 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -233,6 +233,9 @@ steps: timeout_in_minutes: 40 depends_on: - image-build-amd + source_file_dependencies: + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (B200) Spec Decode DeepSeek MTP Parallel Load" key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300 @@ -362,6 +365,33 @@ steps: # Each Gemma4 parametrization is an indivisible ~23m floor. Keep them in # separate shards and put the remaining tests in shard 0, where new tests # land by default. - - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]" - - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" - - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]" + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi + mirror: + amd: + label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance Nightly" + dind: false + device: mi355_1 + timeout_in_minutes: 55 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/tests + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/sample/ + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/worker/gpu_worker.py + - vllm/model_executor/model_loader/ + - vllm/model_executor/models/gemma4.py + - vllm/model_executor/models/gemma4_mm.py + - vllm/model_executor/models/gemma4_mtp.py + - vllm/model_executor/models/llama.py + - vllm/model_executor/models/llama_eagle3.py + - vllm/model_executor/models/medusa.py + - vllm/model_executor/models/registry.py + - tests/evals/gsm8k/ + - tests/utils.py + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 154b355778b2..b8afa261e592 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -20,7 +20,7 @@ steps: label: ":amd: (MI355) Weight Loading Multi-GPU" dind: false device: mi355_2 - timeout_in_minutes: 40 + timeout_in_minutes: 70 depends_on: - image-build-amd commands: From 08633cb5cd77ea315351ae531945a9a253016788 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 17 Sep 2026 16:56:28 +0800 Subject: [PATCH 389/774] [Qwen3.8-Flash-Next] Enable FP8 TP with FlashInfer TRTLLM MoE (#55867) Signed-off-by: Canlin Signed-off-by: Isotr0py Co-authored-by: OpenAI Codex Co-authored-by: Isotr0py --- tests/kernels/moe/test_fp8_tp_loading.py | 311 ++++++++++++++++++ .../model_executor/layers/fused_moe/config.py | 2 + .../layers/fused_moe/oracle/fp8.py | 77 +++++ .../layers/fused_moe/routed_experts.py | 35 +- .../model_executor/layers/quantization/fp8.py | 43 +-- 5 files changed, 434 insertions(+), 34 deletions(-) create mode 100644 tests/kernels/moe/test_fp8_tp_loading.py diff --git a/tests/kernels/moe/test_fp8_tp_loading.py b/tests/kernels/moe/test_fp8_tp_loading.py new file mode 100644 index 000000000000..680ccb812cca --- /dev/null +++ b/tests/kernels/moe/test_fp8_tp_loading.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch + +import vllm.distributed as distributed +from tests.kernels.moe.utils import make_dummy_moe_config +from vllm.model_executor.layers.fused_moe import ( + unquantized_fused_moe_method as unquantized_module, +) +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.quantization import fp8 as fp8_module +from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + +def _make_fp8_tp_experts( + monkeypatch, + tp_size, + tp_rank, + backend="flashinfer_trtllm", + *, + mock_backend=True, + num_experts=2, + hidden_dim=256, + intermediate_size=640, + quant_config=None, + config_overrides=None, +): + # Exercise allocation and the public weight loader without selecting a GPU kernel. + if mock_backend: + monkeypatch.setattr( + fp8_module, "select_fp8_moe_backend", lambda **kwargs: (None, None) + ) + monkeypatch.setattr( + unquantized_module, + "select_unquantized_moe_backend", + lambda **kwargs: (None, None), + ) + monkeypatch.setattr( + fp8_module, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + # The block-shape validator imports this function from vllm.distributed. + monkeypatch.setattr( + distributed, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + config = make_dummy_moe_config( + num_experts=num_experts, + hidden_dim=hidden_dim, + intermediate_size=intermediate_size, + experts_per_token=min(num_experts, 10), + ) + config.routing_method = RoutingMethodType.RenormalizeNaive + config.moe_parallel_config.tp_size = tp_size + config.moe_parallel_config.tp_rank = tp_rank + config.intermediate_size_per_partition = intermediate_size // tp_size + config.intermediate_size_per_partition_unpadded = intermediate_size // tp_size + config.moe_backend = backend + for field, value in (config_overrides or {}).items(): + setattr(config, field, value) + return RoutedExperts( + "model.layers.0.mlp.experts", + torch.bfloat16, + config, + quant_config + or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]), + expert_map_manager=SimpleNamespace( + local_num_experts=num_experts, + placement_strategy="linear", + expert_map=None, + expert_mask=None, + routing_tables=None, + map_global_to_local=lambda index: index, + ), + ) + + +@pytest.mark.parametrize("tp_size", [2, 4, 8]) +@pytest.mark.parametrize("batched", [False, True]) +@pytest.mark.parametrize("intermediate_size", [640, 896]) +def test_fp8_block_aligned_tp_preserves_checkpoint( + monkeypatch, tp_size, batched, intermediate_size +): + """All ranks reconstruct the original dequantized projections, including + padding-only ranks. Reloading must clear stale weights and scales. + """ + generator = torch.Generator().manual_seed(42) + weights = { + "w1": torch.randn(2, intermediate_size, 256, generator=generator).to( + torch.float8_e4m3fn + ), + "w3": torch.randn(2, intermediate_size, 256, generator=generator).to( + torch.float8_e4m3fn + ), + "w2": torch.randn(2, 256, intermediate_size, generator=generator).to( + torch.float8_e4m3fn + ), + } + scales = { + name: torch.rand(2, w.shape[1] // 128, w.shape[2] // 128, generator=generator) + + 0.1 + for name, w in weights.items() + } + reconstructed: dict[str, list[torch.Tensor]] = {name: [] for name in weights} + num_blocks = intermediate_size // 128 + for rank in range(tp_size): + layer = _make_fp8_tp_experts( + monkeypatch, tp_size, rank, intermediate_size=intermediate_size + ) + assert layer.quant_method.weight_scale_refine is None + assert layer.quant_method.moe_block_shape == [128, 128] + width = layer.moe_config.intermediate_size_per_partition + assert width == ((num_blocks + tp_size - 1) // tp_size) * 128 + for name in weights: + prefix = "w2" if name == "w2" else "w13" + for suffix, checkpoint in ( + ("weight", weights[name]), + ("weight_scale_inv", scales[name]), + ): + param_name = f"{prefix}_{suffix}" + param = getattr(layer, param_name) + # Also reload after dirtying the destination, to catch padding leaks. + for _ in range(2): + if name != "w3": + param.data.fill_(7) + if batched: + assert param.weight_loader( + param, checkpoint, param_name, name, 0, return_success=True + ) + else: + for expert in range(2): + assert param.weight_loader( + param, + checkpoint[expert], + param_name, + name, + expert, + return_success=True, + ) + w = getattr(layer, f"{prefix}_weight").float() + s = getattr(layer, f"{prefix}_weight_scale_inv") + if name in ("w1", "w3"): + half = 0 if name == "w1" else 1 + w = w.chunk(2, dim=1)[half] + s = s.chunk(2, dim=1)[half] + dim = 2 if name == "w2" else 1 + valid = max(0, min(width, intermediate_size - rank * width)) + assert torch.count_nonzero(w.narrow(dim, valid, width - valid)) == 0 + assert torch.all(s.narrow(dim, valid // 128, (width - valid) // 128) == 1) + dequant = w * s.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2) + reconstructed[name].append(dequant.narrow(dim, 0, valid)) + for name, parts in reconstructed.items(): + expected = weights[name].float() * scales[name].repeat_interleave( + 128, dim=1 + ).repeat_interleave(128, dim=2) + assert torch.equal(torch.cat(parts, dim=2 if name == "w2" else 1), expected) + + +@pytest.mark.parametrize("backend", ["auto", "triton"]) +@pytest.mark.parametrize("tp_size,block", [(2, 64), (4, 32)]) +def test_fp8_tp_default_keeps_refined_layout(monkeypatch, backend, tp_size, block): + layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0, backend) + assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size + assert layer.quant_method.moe_block_shape == [block, block] + + +@pytest.mark.parametrize( + "field,value", + [ + ("hidden_dim", 192), + ("intermediate_size", 672), + ("is_lora_enabled", True), + ("has_bias", True), + ], +) +def test_fp8_block_aligned_tp_rejects_unsupported_layout(monkeypatch, field, value): + with pytest.raises(ValueError, match="Block-aligned FP8 TP sharding requires"): + _make_fp8_tp_experts(monkeypatch, 4, 0, config_overrides={field: value}) + + +@pytest.mark.parametrize("tp_size", [1, 5]) +def test_fp8_aligned_tp_keeps_original_layout(monkeypatch, tp_size): + layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0) + assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size + assert not layer.moe_config.tp_shard_with_padding + assert layer.quant_method.weight_scale_refine is None + + +def test_fp8_skipped_layer_keeps_original_tp_layout(monkeypatch, default_vllm_config): + layer = _make_fp8_tp_experts( + monkeypatch, + 4, + 0, + quant_config=Fp8Config( + is_checkpoint_fp8_serialized=True, + weight_block_size=[128, 128], + ignored_layers=["model.layers.0.mlp.experts"], + ), + ) + assert isinstance(layer.quant_method, unquantized_module.UnquantizedFusedMoEMethod) + assert not layer.moe_config.tp_shard_with_padding + assert layer.moe_config.intermediate_size_per_partition == 160 + assert layer.w13_weight.shape == (2, 320, 256) + assert layer.w2_weight.shape == (2, 256, 160) + + +def test_fp8_block_aligned_tp_rejects_presharded_weights(monkeypatch): + layer = _make_fp8_tp_experts(monkeypatch, 4, 0) + with pytest.raises(ValueError, match="unsharded checkpoint"): + layer.w13_weight.weight_loader( + layer.w13_weight, + torch.zeros(160, 256, dtype=torch.float8_e4m3fn), + "w13_weight", + "w1", + 0, + ) + + +@pytest.mark.parametrize("tp_size", [2, 4, 8]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +@torch.inference_mode() +def test_fp8_block_aligned_tp_flashinfer_matches_unsharded( + monkeypatch, + tp_size, + num_tokens, + dist_init, + workspace_init, +): + """Compare sharded FI with unsharded FI, reporting the cross-backend + difference separately because activation quantization can differ. + """ + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config + from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + from vllm.platforms import current_platform + + if not current_platform.is_device_capability_family(100): + pytest.skip("Requires datacenter Blackwell and FlashInfer TRTLLM") + torch.manual_seed(42) + e, h, n, topk = 16, 2560, 640, 10 + with torch.device("cuda"), set_current_vllm_config(VllmConfig()): + x = torch.randn(num_tokens, h, dtype=torch.bfloat16) / 10 + logits = torch.randn(num_tokens, e, dtype=torch.float32) + probabilities = logits.softmax(dim=-1) + topk_weights, topk_ids = probabilities.topk(topk, dim=-1) + topk_weights /= topk_weights.sum(dim=-1, keepdim=True) + weights = { + "w1": torch.randn(e, n, h).to(torch.float8_e4m3fn), + "w3": torch.randn(e, n, h).to(torch.float8_e4m3fn), + "w2": torch.randn(e, h, n).to(torch.float8_e4m3fn), + } + scales = { + name: (torch.rand(e, w.shape[1] // 128, w.shape[2] // 128) + 0.5) / 32 + for name, w in weights.items() + } + ref = fused_experts( + x, + torch.cat([weights["w1"], weights["w3"]], dim=1), + weights["w2"], + topk_weights, + topk_ids, + quant_config=fp8_w8a8_moe_quant_config( + w1_scale=torch.cat([scales["w1"], scales["w3"]], dim=1), + w2_scale=scales["w2"], + block_shape=[128, 128], + ), + ) + total = torch.zeros_like(ref, dtype=torch.float32) + fi_reference = None + for rank in range(-1, tp_size): + layer = _make_fp8_tp_experts( + monkeypatch, + 1 if rank == -1 else tp_size, + max(rank, 0), + mock_backend=False, + num_experts=e, + hidden_dim=h, + ) + for name in weights: + prefix = "w2" if name == "w2" else "w13" + for suffix, checkpoint in ( + ("weight", weights[name]), + ("weight_scale_inv", scales[name]), + ): + param_name = f"{prefix}_{suffix}" + param = getattr(layer, param_name) + param.weight_loader(param, checkpoint, param_name, name, 0) + method = layer.quant_method + method.process_weights_after_loading(layer) + assert method.is_monolithic + out = method.apply_monolithic(layer, x, logits) + assert torch.isfinite(out).all() + if rank == -1: + fi_reference = out.float().clone() + else: + total += out.float() + # Different FP8 GEMM reduction orders need not be bitwise identical. + assert fi_reference is not None + relative_l2 = (total - fi_reference).norm() / fi_reference.norm().clamp_min( + 1e-8 + ) + cross_backend_l2 = ( + fi_reference - ref.float() + ).norm() / ref.float().norm().clamp_min(1e-8) + print( + f"TP{tp_size} M={num_tokens}: shard_l2={relative_l2.item():.6f}, " + f"unsharded_fi_vs_triton_l2={cross_backend_l2.item():.6f}" + ) + assert relative_l2 < 0.01, relative_l2.item() diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 9a2f8229b1b6..35973e635cc0 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1309,6 +1309,8 @@ class FusedMoEConfig: # Set by __post_init__ intermediate_size_per_partition: int = -1 + # Use the allocated width as the checkpoint TP stride, including for scales. + tp_shard_with_padding: bool = False rocm_aiter_fmoe_enabled: bool = False aiter_fmoe_shared_expert_enabled: bool = False diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 7c57e2e5c2bd..6a6ce34d160d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -31,11 +31,14 @@ prepare_fp8_moe_layer_for_marlin, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, QuantKey, + create_fp8_quant_key, kFp8Dynamic128Sym, kFp8Static128BlockSym, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up logger = init_logger(__name__) @@ -299,6 +302,80 @@ def refine_fp8_moe_block_shape( return [refine, refine] +def pad_tp_shard_to_weight_blocks( + config: FusedMoEConfig, + weight_block_size: list[int], +) -> bool: + """Pad the TP shard to whole checkpoint blocks, keeping scales rank-local.""" + block_n, block_k = weight_block_size + if ( + block_n != block_k + or config.tp_size == 1 + or config.intermediate_size_per_partition % block_n == 0 + ): + return False + if ( + config.intermediate_size % block_n != 0 + or config.hidden_dim % block_n != 0 + or config.ep_size != 1 + or config.is_lora_enabled + or config.has_bias + ): + raise ValueError( + f"Block-aligned FP8 TP sharding requires {block_n}-aligned " + "global expert dimensions, pure TP, and no LoRA or " + "expert bias." + ) + config.tp_shard_with_padding = True + config.intermediate_size_per_partition = round_up( + config.intermediate_size_per_partition, block_n + ) + return True + + +def resolve_fp8_moe_weight_block_shape( + config: FusedMoEConfig, + weight_block_size: list[int], + activation_key: QuantKey, + is_checkpoint_fp8_serialized: bool, +) -> tuple[list[int], tuple[int, int] | None]: + """Return the TP-adapted block shape and refine factor: + refine if kernels allow, else pad to the TP shard.""" + refined_shape = refine_fp8_moe_block_shape(config, weight_block_size) + if is_checkpoint_fp8_serialized and config.moe_backend != "auto": + kernel_classes = backend_to_kernel_cls(map_fp8_backend(config.moe_backend)) + can_refine = refined_shape is not None and any( + k_cls._supports_quant_scheme( + create_fp8_quant_key( + static=True, group_shape=GroupShape(*refined_shape) + ), + activation_key, + ) + for k_cls in kernel_classes + ) + if not can_refine and pad_tp_shard_to_weight_blocks(config, weight_block_size): + logger.info_once( + "FP8 %s TP loading uses complete checkpoint blocks: " + "local allocation %d, without weight requantization.", + config.moe_backend, + config.intermediate_size_per_partition, + ) + return weight_block_size, None + if refined_shape is None: + return weight_block_size, None + logger.info_once( + "FP8 MoE block scales refined from %s to %s to fit " + "the TP-sharded intermediate size %d.", + str(weight_block_size), + str(refined_shape), + config.intermediate_size_per_partition, + ) + return refined_shape, ( + weight_block_size[0] // refined_shape[0], + weight_block_size[1] // refined_shape[1], + ) + + def select_fp8_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 88841070f03b..42ea15dcda7d 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -332,6 +332,7 @@ def _load_model_weight_or_group_weight_scale( shard_id: str, loaded_weight: torch.Tensor, tp_rank: int, + is_scale: bool = False, ): """Load grouped weight scales for group quantization or model weights @@ -341,14 +342,41 @@ def _load_model_weight_or_group_weight_scale( shard_id: either w1, w2, or w3 loaded_weight: checkpoint weight to load into the param tp_rank: tensor parallel rank + is_scale: whether padding should use unit scales instead of zero weights. """ + padded_tp = self.moe_config.tp_shard_with_padding + if padded_tp: + destination = expert_data + if shard_id in ("w1", "w3") and self.moe_config.is_act_and_mul: + half = destination.shape[shard_dim] // 2 + destination = destination.narrow( + shard_dim, 0 if shard_id == "w1" else half, half + ) + shard_size = destination.shape[shard_dim] + expected_size = ( + self.moe_config.intermediate_size + * shard_size + // self.moe_config.intermediate_size_per_partition + ) + if loaded_weight.shape[shard_dim] != expected_size: + raise ValueError( + "Block-aligned TP loading expects an unsharded checkpoint " + f"projection of size {expected_size}, got " + f"{loaded_weight.shape[shard_dim]}." + ) + start = min(tp_rank * shard_size, expected_size) + size = min(shard_size, expected_size - start) + loaded_weight = loaded_weight.narrow(shard_dim, start, size) + destination.fill_(1 if is_scale else 0) + if shard_id == "w2": self._load_w2( shard_dim=shard_dim, loaded_weight=loaded_weight, expert_data=expert_data, tp_rank=tp_rank, + load_full=padded_tp, ) elif shard_id in ("w1", "w3"): self._load_w13( @@ -357,6 +385,7 @@ def _load_model_weight_or_group_weight_scale( loaded_weight=loaded_weight, expert_data=expert_data, tp_rank=tp_rank, + load_full=padded_tp, ) def _load_per_channel_weight_scale( @@ -511,11 +540,12 @@ def _load_w2( shard_dim: int, loaded_weight: torch.Tensor, tp_rank: int, + load_full: bool = False, ): # Index the loaded weight for tp sharding. # down_proj: "RowParallel" so tp sharding on input_dim - # Only narrow if the loaded_weight is not a scalar (0-dim tensor). - if loaded_weight.ndim > 0: + # Padded TP weights have already been sliced by the grouped loader. + if not load_full and loaded_weight.ndim > 0: # Same padding fix as _load_w13: use unpadded per-rank size. tp_size = self.moe_config.moe_parallel_config.tp_size loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size @@ -800,6 +830,7 @@ def weight_loader( loaded_weight=loaded_weight, expert_data=expert_data, tp_rank=self.moe_config.tp_rank, + is_scale=True, ) elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value: self._load_per_tensor_weight_scale( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 5835af111af6..c85a2c288e76 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -8,7 +8,6 @@ import vllm.envs as envs from vllm.config import get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, ) @@ -24,14 +23,12 @@ SharedExperts, UnquantizedFusedMoEMethod, ) -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEQuantConfig, -) +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, - refine_fp8_moe_block_shape, + resolve_fp8_moe_weight_block_shape, select_fp8_moe_backend, ) from vllm.model_executor.layers.linear import ( @@ -90,8 +87,6 @@ ACTIVATION_SCHEMES = ["static", "dynamic"] -logger = init_logger(__name__) - class Fp8Config(QuantizationConfig): """Config class for FP8.""" @@ -511,32 +506,16 @@ def __init__(self, quant_config: Fp8Config, layer: RoutedExperts): # Set weight key and activation key for kernel compatibility if self.block_quant: assert self.weight_block_size is not None - # TP shards the intermediate dim of the expert weights, so a - # per-shard size that is not a multiple of the checkpoint's block - # size makes the checkpoint's block scales impossible to shard - # exactly. When a finer block size (>= 32) divides both the - # checkpoint blocks and all involved dims, the weight scales are - # refined to that granularity at load time (a lossless upsampling, - # since the refined block divides the checkpoint block). The - # refined block shape is encoded in the weight key, so the oracle - # only selects kernels that support it (e.g. Triton, which takes - # the block shape as a runtime argument). - refined_shape = refine_fp8_moe_block_shape(self.moe, self.weight_block_size) - if refined_shape is not None: - block_n, block_k = self.weight_block_size - self.weight_scale_refine = ( - block_n // refined_shape[0], - block_k // refined_shape[1], - ) - self.moe_block_shape = refined_shape - logger.info_once( - "FP8 MoE block scales refined from %s to %s to fit " - "the TP-sharded intermediate size %d.", - str(self.weight_block_size), - str(refined_shape), - self.moe.intermediate_size_per_partition, + # Adapt the checkpoint block shape to TP sharding before + # building the weight key. + self.moe_block_shape, self.weight_scale_refine = ( + resolve_fp8_moe_weight_block_shape( + self.moe, + self.weight_block_size, + kFp8Dynamic128Sym, + self.quant_config.is_checkpoint_fp8_serialized, ) - assert self.moe_block_shape is not None + ) weight_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.moe_block_shape) ) From 40b40d1d39bb82563dd699be3456abaa95d0722a Mon Sep 17 00:00:00 2001 From: drakosha Date: Thu, 17 Sep 2026 13:19:56 +0300 Subject: [PATCH 390/774] [Bugfix][KV Offload] Register the offload region in chunks (#51081) Signed-off-by: Mikhail Kostryukov Co-authored-by: Claude Co-authored-by: Colton Ottley Co-authored-by: Or Ozeri --- .../kv_offload/cpu/test_canonical_layout.py | 1 + tests/v1/kv_offload/cpu/test_gpu_worker.py | 16 ++++ .../cpu/test_shared_offload_region.py | 82 +++++++++++++++++ .../device_communicators/cuda_wrapper.py | 26 ++++++ vllm/v1/kv_offload/cpu/gpu_worker.py | 88 +++++++++++++++---- 5 files changed, 198 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py index 4594e1eb0924..749bd6fac501 100644 --- a/tests/v1/kv_offload/cpu/test_canonical_layout.py +++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py @@ -241,6 +241,7 @@ def canonical_view(rank: int, world_size: int) -> torch.Tensor: # The Triton load path dereferences CPU pointers on the GPU, which is # only legal on pinned memory; production pins via CPUOffloadingWorker pin_mmap_region(region) + assert region.is_pinned return region.create_next_canonical_view(_CANONICAL_PAGE) try: diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 5306d0557102..731d61a7e924 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -49,6 +49,22 @@ def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None: ) +@pytest.mark.skipif(not gpu_worker.HAS_TRITON, reason="Requires Triton") +def test_unpinned_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None: + """The Triton load path dereferences CPU pointers on the GPU, so pageable + host memory takes the DMA path even for pages where Triton would win.""" + monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: False) + + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]] + assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is not ( + ops.swap_blocks_batch + ) + assert gpu_worker._select_swap_blocks_fn( + refs, gpu_to_cpu=False, host_memory_is_pinned=False + ) is (ops.swap_blocks_batch) + + def test_worker_shutdown_releases_region_and_runs_both_handlers() -> None: """Both directions drain before the worker releases its shared region.""" worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker) diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 4f22d122dbec..f84e1bd7296c 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -15,6 +15,7 @@ import pytest from vllm.utils.system_utils import get_mp_context +from vllm.v1.kv_offload.cpu import gpu_worker from vllm.v1.kv_offload.cpu import shared_offload_region as region_module from vllm.v1.kv_offload.cpu.shared_offload_region import ( SharedOffloadRegion, @@ -788,6 +789,87 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): ] +@pytest.fixture +def host_register(monkeypatch): + """Pretend to run on CUDA with host registrations capped at seven pages. + + Registration goes to a mock CudaRTLibrary and the region's cleanup to a + separate mock torch runtime. Both stay installed through the cleanup that + ``_region`` runs on exit. + """ + cudart = MagicMock() + cudart.cudaHostRegister.return_value = 0 + cudart.cudaHostUnregister.return_value = 0 + cudart.cudaGetLastError.return_value = 2 + torch_cudart = MagicMock() + torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0) + monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE) + monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart) + monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart) + return cudart, torch_cudart + + +def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register): + """Chunks end on row boundaries: 3-page rows under a 7-page cap register + as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1.""" + cudart, torch_cudart = host_register + with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: + gpu_worker.pin_mmap_region(region) + base = region._base.data_ptr() + assert cudart.mock_calls == [ + call.cudaHostRegister(base, 6 * PAGE_SIZE), + call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE), + call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE), + ] + assert region.is_pinned + + assert torch_cudart.cudaHostUnregister.call_args_list == [ + call(base + 12 * PAGE_SIZE), + call(base + 6 * PAGE_SIZE), + call(base), + ] + + +@pytest.mark.parametrize("fail_at", [0, 1, 2]) +def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at): + """A failed chunk drains its pending error and unregisters the chunks + before it on the same handle, so the region is pinned whole or not at all.""" + cudart, torch_cudart = host_register + cudart.cudaHostRegister.side_effect = [0] * fail_at + [2] + with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: + gpu_worker.pin_mmap_region(region) + base = region._base.data_ptr() + assert not region.is_pinned + assert region.pinned_addresses == [] + + registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)] + failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE + assert cudart.mock_calls == [ + *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered), + call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size), + call.cudaGetLastError(), + *(call.cudaHostUnregister(address) for address in reversed(registered)), + ] + torch_cudart.cudaHostUnregister.assert_not_called() + + +def test_pin_mmap_region_without_cuda_runtime_stays_pageable( + iid, host_register, monkeypatch +): + """A CUDA runtime that cannot be loaded leaves the region pageable instead + of failing worker startup.""" + + def missing_runtime(): + raise AssertionError("libcudart is not loaded in the current process") + + monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime) + with _region(iid) as region: + gpu_worker.pin_mmap_region(region) + assert not region.is_pinned + assert region.pinned_addresses == [] + + def test_cleanup_after_create_next_worker_view_releases_mmap(iid): """cleanup() must close the mmap even after create_next_worker_view was called. create_next_worker_view returns a view that shares storage with _base; both must be diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py index a5026163e934..d27d84c03045 100644 --- a/vllm/distributed/device_communicators/cuda_wrapper.py +++ b/vllm/distributed/device_communicators/cuda_wrapper.py @@ -78,6 +78,16 @@ class CudaRTLibrary: cudaError_t, [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint], ), + # cudaError_t cudaHostRegister ( void* ptr, size_t size, unsigned int flags ) + Function( + "cudaHostRegister", + cudaError_t, + [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint], + ), + # cudaError_t cudaHostUnregister ( void* ptr ) + Function("cudaHostUnregister", cudaError_t, [ctypes.c_void_p]), + # cudaError_t cudaGetLastError ( void ) + Function("cudaGetLastError", cudaError_t, []), ] # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa @@ -92,6 +102,9 @@ class CudaRTLibrary: "cudaMemcpy": "hipMemcpy", "cudaIpcGetMemHandle": "hipIpcGetMemHandle", "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle", + "cudaHostRegister": "hipHostRegister", + "cudaHostUnregister": "hipHostUnregister", + "cudaGetLastError": "hipGetLastError", } # class attribute to store the mapping from the path to the library @@ -185,3 +198,16 @@ def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p: ) ) return devPtr + + def cudaHostRegister(self, ptr: int, size: int, flags: int = 0) -> int: + """Return the raw error code instead of raising, since callers can + fall back to pageable memory. Drain a failure with cudaGetLastError.""" + return self.funcs["cudaHostRegister"](ctypes.c_void_p(ptr), size, flags) + + def cudaHostUnregister(self, ptr: int) -> int: + """Return the raw error code instead of raising.""" + return self.funcs["cudaHostUnregister"](ctypes.c_void_p(ptr)) + + def cudaGetLastError(self) -> int: + """Return and clear the error pending on this thread.""" + return self.funcs["cudaGetLastError"]() diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 037e3631e499..797abed5eaa4 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -11,6 +11,7 @@ import torch from vllm import _custom_ops as ops +from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, triton @@ -38,10 +39,13 @@ def _select_swap_blocks_fn( layer_refs_per_group: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, + host_memory_is_pinned: bool = True, ): """Resolve the swap_blocks function for a handler at init time.""" # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton. - if gpu_to_cpu: + # The Triton kernel dereferences CPU pointers on the GPU, which is only + # valid for pinned host memory. + if gpu_to_cpu or not host_memory_is_pinned: return ops.swap_blocks_batch # Fall back to the C++ DMA path on platforms where Triton isn't usable # (e.g. ROCm host mappings) or where GPU kernels cannot directly @@ -190,8 +194,12 @@ def _canonical_block_sizes( return canonical_bytes_per_block +# Bound registration size to avoid driver limits on large host allocations. +MAX_HOST_REGISTER_CHUNK_BYTES = 64 * 1024**3 + + def pin_mmap_region(region: SharedOffloadRegion) -> None: - """Register the entire mmap as CUDA pinned memory via cudaHostRegister.""" + """Register row-aligned chunks, rolling back on failure.""" if not current_platform.is_cuda_alike(): logger.info( "Skipping mmap host registration on %s; cudaHostRegister is only " @@ -201,23 +209,65 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None: return rank = region.rank + try: + cudart = CudaRTLibrary() + except (AssertionError, AttributeError, OSError): + logger.warning( + "Could not load the CUDA runtime for host registration on rank=%d; " + "the offload region stays pageable", + rank, + exc_info=True, + ) + return base_ptr = region._base.data_ptr() - result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0) - if result.value != 0: + total_size = region.total_size_bytes + # Chunks end on block-row boundaries, which are page aligned, so neither the + # driver's page rounding nor any single block transfer straddles two + # registrations. + rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1) + chunk_size = rows_per_chunk * region._row_stride + + # Register, drain and roll back through the same runtime handle, so a + # failed chunk leaves neither a pending error nor a partly pinned region. + addresses: list[int] = [] + for offset in range(0, total_size, chunk_size): + address = base_ptr + offset + size = min(chunk_size, total_size - offset) + result = cudart.cudaHostRegister(address, size) + if result == 0: + addresses.append(address) + continue + cudart.cudaGetLastError() logger.warning( - "cudaHostRegister failed for rank=%d (code=%d) — " - "transfers will still work but may be slower (unpinned DMA)", + "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); " + "the offload region stays pageable", rank, + offset / 1e9, + total_size / 1e9, result, ) - else: - logger.debug( - "cudaHostRegister rank=%d %.2f GB", - rank, - region.total_size_bytes / 1e9, - ) - region.is_pinned = True + for registered in reversed(addresses): + unregister_result = cudart.cudaHostUnregister(registered) + if unregister_result != 0: + cudart.cudaGetLastError() + logger.warning( + "cudaHostUnregister failed for rank=%d at %#x (code=%d); " + "that chunk stays registered until the process exits", + rank, + registered, + unregister_result, + ) + return + + region.pinned_addresses.extend(addresses) + region.is_pinned = True + logger.debug( + "cudaHostRegister rank=%d %.2f GB in %d chunk(s)", + rank, + total_size / 1e9, + len(addresses), + ) def _new_descriptor_buffers( @@ -248,6 +298,7 @@ def __init__( layer_refs_per_group: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, canonical_layout: bool = False, + host_memory_is_pinned: bool = True, ): """Initialize a SingleDirectionOffloadingHandler. @@ -262,6 +313,8 @@ def __init__( gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU. canonical_layout: if True, CPU pages use the canonical layout described by the refs' mappings. + host_memory_is_pinned: whether the CPU tensors are pinned, so GPU + kernels may dereference them directly. """ assert len(gpu_tensors) == len(cpu_tensors) @@ -299,7 +352,7 @@ def __init__( self.gpu_to_cpu: bool = gpu_to_cpu self.layer_refs_per_group = layer_refs_per_group self._swap_blocks_batch = _select_swap_blocks_fn( - layer_refs_per_group, gpu_to_cpu + layer_refs_per_group, gpu_to_cpu, host_memory_is_pinned ) # GPU blocks may be smaller @@ -646,7 +699,7 @@ def transfer_async( last_event = last_transfer.end_event # assure job will start only after the previous one completes stream.wait_event(last_event) - # CPU->GPU reads from host pinned memory, which is never written + # CPU->GPU reads from host memory, which is never written # by a concurrent GPU stream, so CU_MEMCPY_SRC_ACCESS_ORDER_ANY is # safe and lets the driver pipeline source reads. GPU->CPU reads # from the live GPU KV cache, which the compute stream keeps @@ -764,6 +817,9 @@ def __init__( logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) if mmap_region is not None and pin_memory: pin_mmap_region(mmap_region) + host_memory_is_pinned = pin_memory and ( + mmap_region is None or mmap_region.is_pinned + ) canonical_bytes_per_block = ( _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors)) @@ -813,6 +869,7 @@ def __init__( layer_refs_per_group=kv_caches.group_data_refs, gpu_to_cpu=True, canonical_layout=canonical_layout, + host_memory_is_pinned=host_memory_is_pinned, ) self._load_handler = SingleDirectionOffloadingHandler( @@ -822,6 +879,7 @@ def __init__( layer_refs_per_group=kv_caches.group_data_refs, gpu_to_cpu=False, canonical_layout=canonical_layout, + host_memory_is_pinned=host_memory_is_pinned, ) def submit_store( From b6e7c1f1f0430b5d4784aea391c42067581b7f76 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:22:49 +0300 Subject: [PATCH 391/774] [kv_offload] Skip scratch groups (#57145) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../unit/offloading_connector/test_config.py | 33 +++++++++++++++ .../unit/offloading_connector/test_worker.py | 42 +++++++++++++++++++ .../kv_connector/v1/offloading/config.py | 6 +-- .../kv_connector/v1/offloading/worker.py | 5 ++- 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index acf55972dfa4..fe0d258462b8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -19,6 +19,7 @@ from vllm.platforms import current_platform from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KVCacheConfig, @@ -340,6 +341,38 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): assert offloading_config.cache.blocks_per_chunk == 2 +def test_offloading_skips_scratch_group(): + kv_cache_config = _make_mamba_hybrid_kv_cache_config() + kv_cache_config.kv_cache_groups.insert( + 1, + KVCacheGroupSpec( + ["scratch"], + CircularBufferSpec( + block_size=4, num_kv_heads=1, head_size=128, dtype=torch.float32 + ), + ), + ) + page_size = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes + kv_cache_config.kv_cache_tensors = [ + KVCacheTensor( + size=page_size * kv_cache_config.num_blocks, + layers=["full_layer", "scratch", "mamba_layer"], + layer_stride=0, + block_stride=page_size, + ) + ] + config = _make_vllm_config() + + offloading_config = build_offloading_config(config, kv_cache_config) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), config, kv_cache_config + ) + + assert [group.group_id for group in offloading_config.groups] == [0, 2] + assert [group.group_idx for group in scheduler_config.kv_group_configs] == [0, 2] + assert offloading_config.worker_kv_bytes_per_block == page_size + + def test_hisparse_offloads_only_indexer_group(): source = KVCacheGroupSpec( ["source"], diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 64f481d02efd..32bd4f3eca75 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -13,6 +13,7 @@ from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -520,6 +521,47 @@ def test_register_kv_caches(backend): assert actual.mapping is not None +def test_register_packed_kv_caches_skips_scratch_group(): + attn_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + page = attn_spec.page_size_bytes + layers = ["layer0", "layer1"] + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor( + size=2 * page * NUM_BLOCKS, + layers=layers, + layer_stride=page, + block_stride=2 * page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layers, attn_spec), + KVCacheGroupSpec( + ["scratch"], + CircularBufferSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=DTYPE + ), + ), + ], + ) + kv_caches = _allocate_kv_caches( + kv_cache_config, [], device=torch.device(f"{DEVICE_TYPE}:0") + ) + worker, spec = _make_worker(kv_cache_config) + worker.register_kv_caches(kv_caches) + + canonical = spec.get_worker.call_args[0][0] + assert len(canonical.tensors) == 1 + assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, 2 * page) + assert canonical.group_data_refs == [[CanonicalKVCacheRef(0, 2 * page)]] + + @pytest.mark.parametrize("backend", ATTN_BACKENDS) def test_register_kv_caches_uniform_type(backend): """Test register_kv_caches with UniformTypeKVCacheSpecs. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index ca503d4e675d..e20c1189af1e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -35,7 +35,7 @@ def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: if kv_cache_config.hisparse_host_num_blocks is None: - return tuple(range(len(kv_cache_config.kv_cache_groups))) + return kv_cache_config.prefix_cacheable_group_ids return tuple( group_id for group_id, group in enumerate(kv_cache_config.kv_cache_groups) @@ -135,12 +135,12 @@ def build_offloading_config( ) worker_kv_bytes_per_block = 0 - all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups) if ( - all_groups_selected + kv_cache_config.hisparse_host_num_blocks is None and kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors ): + # Scratch filtering must preserve the scheduler/worker allocation stride. # Every KVCacheTensor describes placement within the same backing allocation, # so its size is the total, not a per-tensor share. total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 5b14e50dda28..2bb00d88f089 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -141,8 +141,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ), None, ) - if packed_layer_name is not None and len(selected_groups) == len( - kv_cache_config.kv_cache_groups + if ( + packed_layer_name is not None + and kv_cache_config.hisparse_host_num_blocks is None ): (tensor,) = tensors_per_block[packed_layer_name] num_blocks = tensor.shape[0] From 683f6c8edf9656da99926bbe5aae5ffffaa850b3 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 17 Sep 2026 03:33:44 -0700 Subject: [PATCH 392/774] [CI] Raise H200 LM Eval Large Models timeout to 120 min (#57335) Signed-off-by: khluu Co-authored-by: Kimi --- .buildkite/test_areas/lm_eval.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index e94cbe5dfc15..3eb7f0a44430 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -296,7 +296,7 @@ steps: - label: ":nvidia: (H200) LM Eval Large Models" key: lm-eval-large-models-8xh200 - timeout_in_minutes: 50 + timeout_in_minutes: 120 device: h200 optional: true num_devices: 8 From 99ea5f525da6b5428f2c5740f3cd15377f284c1d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 17 Sep 2026 03:34:00 -0700 Subject: [PATCH 393/774] [CI] Raise DSv4-Flash disaggregated engine readiness timeout to 1800s (#57334) Signed-off-by: khluu Co-authored-by: Kimi --- .buildkite/test_areas/disaggregated.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 0fbb32f8f9ff..d217090bc5e7 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -287,7 +287,7 @@ steps: PREFILL_BLOCK_SIZE: "256" DECODE_BLOCK_SIZE: "256" MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" - VLLM_ENGINE_READY_TIMEOUT_S: "900" + VLLM_ENGINE_READY_TIMEOUT_S: "1800" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" # The V4 indexer only supports block-outermost layouts; let it pick. PREFILLER_KV_LAYOUT: "" From 0d726275653800e82c324e1d32b074f782e89315 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Thu, 17 Sep 2026 18:45:52 +0800 Subject: [PATCH 394/774] [Perf][EPD] Batch image requests per encoder (#57095) Signed-off-by: Tianyu Guo Signed-off-by: Zhou ziheng Co-authored-by: Zhou ziheng Co-authored-by: OpenAI Codex --- .../disaggregated_encoder/README.md | 16 ++ .../disaggregated_encoder/disagg_epd_proxy.py | 126 +++++++--- .../unit/test_ec_example_connector.py | 15 ++ .../ec_connector/unit/test_epd_proxy_retry.py | 216 ++++++++++++++++-- .../unit/test_scheduler_nixl_producer.py | 4 +- .../ec_transfer/ec_connector/utils.py | 11 +- 6 files changed, 329 insertions(+), 59 deletions(-) diff --git a/examples/disaggregated/disaggregated_encoder/README.md b/examples/disaggregated/disaggregated_encoder/README.md index bd2a46abbe52..71a0d0e72ea8 100644 --- a/examples/disaggregated/disaggregated_encoder/README.md +++ b/examples/disaggregated/disaggregated_encoder/README.md @@ -106,6 +106,22 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. | | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). | +The proxy batches images from the same user request assigned to the same encoder. +Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number +of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each +image separately. Audio and video remain separate subrequests. + +For encoders with a smaller `--limit-mm-per-prompt` image limit than P/PD, set +`ENCODER_MAX_BATCH_SIZE` no higher than the smallest encoder image limit. +For example, for encoders configured with `--limit-mm-per-prompt '{"image": 2}'`: + +```bash +ENCODER_MAX_BATCH_SIZE=2 python disagg_epd_proxy.py \ + --encode-servers-urls "http://e1:8001,http://e2:8002" \ + --prefill-servers-urls disable \ + --decode-servers-urls "http://pd1:8003" +``` + Example usage: For E + PD setup: diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index c44c9026906b..3317b091fb52 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -10,8 +10,8 @@ For MM input we: 1. Extract *every* image/audio/video item. - 2. Fire N concurrent requests to the encoder cluster - (one request per item, with **all text removed**). + 2. Send concurrent encoder requests with all text removed, + grouping images assigned to the same encoder. 3. Wait for all of them to succeed. 4. Forward the *original* request to a decode server. """ @@ -99,6 +99,11 @@ def validate_ec_consumer_routing( # Diagnostic switch: forward the original request to the decoder so the # only difference from the rewrite path is the rewrite itself. NO_REWRITE = False +# Maximum images per encoder subrequest; 0 leaves batches unlimited. +ENCODER_MAX_BATCH_SIZE = int(os.getenv("ENCODER_MAX_BATCH_SIZE", "0")) +if ENCODER_MAX_BATCH_SIZE < 0: + raise ValueError("ENCODER_MAX_BATCH_SIZE must be non-negative") + # Decode-side retries for a retryable internal error (`finish_reason="error"`, # e.g. an encoder embedding the connector could not deliver). Re-issuing runs # the encode again, which produces a fresh transfer. @@ -222,7 +227,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. Group images by encoder, retaining per-item round-robin assignment. 2. Send them concurrently to the encode cluster. 3. Raise if any of them fails. @@ -248,6 +253,7 @@ async def fanout_encoder_primer( item_uuids: dict[int, str] = {} item_transfer_ids: dict[int, str] = {} item_meta: dict[int, dict] = {} + transfer_items: dict[int, dict[str, str]] = {} ec_params: dict[str, Any] = {} # Round-robin over encode servers to distribute load a bit. The cursor @@ -259,29 +265,45 @@ async def fanout_encoder_primer( e_urls, encoder_rr_idx, len(mm_items) ) + groups: list[tuple[str, list[int]]] = [] + image_groups: dict[str, list[int]] = {} for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)): + if item["type"] == "image_url": + indices = image_groups.get(target_url) + if indices is None or ( + ENCODER_MAX_BATCH_SIZE and len(indices) >= ENCODER_MAX_BATCH_SIZE + ): + indices = [] + image_groups[target_url] = indices + groups.append((target_url, indices)) + indices.append(idx) + else: + groups.append((target_url, [idx])) + + for target_url, indices in groups: # Derive a *child* request id: :: - child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}" + child_req_id = f"{req_id}:{indices[0]}:{uuid.uuid4().hex[:6]}" headers = {"x-request-id": child_req_id, "Content-Type": "application/json"} # With --no-rewrite the decoder still receives the raw image and derives # the cache key by hashing it, so the encoder must do the same -- passing # a uuid here would make the two disagree and silently defeat the EC # transfer, leaving the decoder to encode the image itself. - item_uuid = None if NO_REWRITE else content_uuid(item) - if item_uuid is not None: - item_uuids[idx] = item_uuid - transfer_id = uuid.uuid4().hex - item_transfer_ids[idx] = transfer_id + content = [] + for idx in indices: + item = mm_items[idx] + item_uuid = None if NO_REWRITE else content_uuid(item) + if item_uuid is not None: + item_uuids[idx] = item_uuid + item_transfer_ids[idx] = uuid.uuid4().hex + content.append(item if item_uuid is None else {**item, "uuid": item_uuid}) encoder_req = { "model": orig_request.get("model"), "messages": [ { "role": "user", - "content": [ - item if item_uuid is None else {**item, "uuid": item_uuid} - ], + "content": content, }, ], # No max_tokens cap: the encoder instance never samples, it finishes @@ -297,17 +319,13 @@ async def fanout_encoder_primer( if key in orig_request: encoder_req[key] = orig_request[key] if consumer_zmq is not None: - # No mm_hash here on purpose. The encoder's own - # `mm_features[i].identifier` is derived from the uuid *and* the - # engine's media_io_kwargs / mm_processor_kwargs, so this proxy - # cannot know it before the encoder runs. Sending the bare uuid - # would fail the connector's hash match and make the producer - # invent its own transfer id, which the consumer could then never - # cancel. Omitting it lets the connector match by position, which - # is exact: one encoder request carries exactly one item. + # The engine may rehash UUIDs with processing options. Match by + # position instead; batches contain only images in input order. encoder_req["ec_transfer_params"] = { "consumer_zmq": consumer_zmq, - "ec_items": [{"transfer_id": transfer_id}], + "ec_items": [ + {"transfer_id": item_transfer_ids[idx]} for idx in indices + ], } tasks.append( encode_session.post( @@ -320,7 +338,8 @@ async def fanout_encoder_primer( results = await asyncio.gather(*tasks, return_exceptions=True) # Fail fast if any sub-request failed - for idx, r in enumerate(results): + for (_, indices), r in zip(groups, results): + idx = indices[0] if isinstance(r, Exception): logger.error( "[%s] Encoder request #%d raised exception: %s", @@ -357,18 +376,45 @@ async def fanout_encoder_primer( logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx) params = {} if params: - # One encoder request carries exactly one item, so there is a - # single reported entry. Do not key it by this proxy's uuid: when - # media_io_kwargs or mm_processor_kwargs are set the engine - # re-hashes the uuid together with them, so the encoder's own - # `mm_features[i].identifier` is a derived value this proxy cannot - # predict. Fall back to the sole entry, and carry the key the - # encoder actually used through as `ec_mm_hash`. - ec_mm_hash = item_uuids.get(idx) - reported = params.get(ec_mm_hash) - if reported is None and len(params) == 1: - ((ec_mm_hash, reported),) = params.items() - if reported: + by_index = {} + for mm_hash, reported in params.items(): + if len(indices) == 1: + break + for local_index in reported.get("item_indices", []): + if type(local_index) is not int or not 0 <= local_index < len( + indices + ): + raise HTTPException(502, "Invalid encoder item index") + if local_index in by_index: + raise HTTPException(502, "Duplicate encoder item index") + by_index[local_index] = (mm_hash, reported) + for local_index, idx in enumerate(indices): + matched = by_index.get(local_index) + if matched is None: + # Compatibility with encoders without position metadata. + mm_hash = item_uuids.get(idx) + if mm_hash in params: + matched = (mm_hash, params[mm_hash]) + elif len(indices) == 1 and len(params) == 1: + matched = next(iter(params.items())) + elif ( + mm_items[idx]["type"] == "video_url" + and (orig_request.get("mm_processor_kwargs") or {}).get( + "use_audio_in_video" + ) + and len(params) == 2 + and consumer_zmq is None + and all( + entry.keys() <= {"metadata", "item_indices"} + for entry in params.values() + ) + ): + # Keep raw video for its audio/video features when no transfer + # handles need to be forwarded. + continue + else: + raise HTTPException(502, "Encoder metadata cannot be matched") + ec_mm_hash, reported = matched metadata = reported.get("metadata") or {} if metadata: item_meta[idx] = { @@ -384,12 +430,16 @@ async def fanout_encoder_primer( # looks it up by mm_hash on the request, so carry it through. ec_params[ec_mm_hash] = reported if NO_REWRITE and consumer_zmq is not None: - ec_params.setdefault("ec_items", []).append( - {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]} - ) + transfer_items[idx] = { + "mm_hash": ec_mm_hash, + "transfer_id": item_transfer_ids[idx], + } + + if transfer_items: + ec_params["ec_items"] = [transfer_items[idx] for idx in sorted(transfer_items)] logger.info( - "[%s] All %d encoder requests completed successfully", req_id, len(mm_items) + "[%s] All %d encoder requests completed successfully", req_id, len(groups) ) return item_meta, ec_params diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index dac809d1a42c..a6645fccf4f3 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -16,6 +16,7 @@ ECExampleConnectorMetadata, MMMeta, ) +from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange from vllm.v1.core.sched.output import SchedulerOutput @@ -41,6 +42,20 @@ def get_num_encoder_embeds(self, input_id: int) -> int: return self._token_counts[input_id] +def test_metadata_preserves_repeated_item_positions(): + request = MockRequest("request", ["A", "B", "A"], [4, 4, 4]) + resolver = Mock() + resolver.fields_for.return_value = {"image_grid_thw"} + for feature in request.mm_features: + feature.data = Mock() + feature.data.get_data.return_value = {"image_grid_thw": torch.tensor([1, 2, 2])} + items = collect_ec_item_metadata(request.mm_features, resolver) + assert items == { + "A": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [0, 2]}, + "B": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [1]}, + } + + @pytest.fixture def temp_storage(tmp_path): """Fixture providing temporary storage path.""" diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 1a8c877c72e2..0c48dc84d6f8 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -9,6 +9,8 @@ import asyncio import importlib.util from pathlib import Path +from types import SimpleNamespace +from typing import Any import httpx import msgspec @@ -83,6 +85,128 @@ async def post(self, url, data=None, headers=None): return _EncoderResponse(self._replies.pop(0)) +@pytest.mark.parametrize("no_rewrite", [False, True]) +@pytest.mark.parametrize( + "batch_size, expected_sizes", [(0, [3, 3]), (1, [1] * 6), (2, [2, 2, 1, 1])] +) +def test_image_batches_preserve_item_identity( + proxy, monkeypatch, no_rewrite, batch_size, expected_sizes +): + """Rehashed and repeated images retain metadata and per-occurrence transfers.""" + from vllm.distributed.ec_transfer.ec_connector.mooncake.scheduler import ( + ECMooncakeScheduler, + ) + + seen = [] + + class Session: + async def post(self, url, data=None, headers=None): + body = msgspec.json.decode(data) + seen.append((url, body)) + params: dict[str, dict[str, Any]] = {} + for index, item in enumerate(body["messages"][0]["content"]): + image = item["image_url"]["url"] + entry = params.setdefault( + image + "-processed", + { + "metadata": {"image_grid_thw": [1, 2, ord(image)]}, + "item_indices": [], + "peer_port": 1234, + }, + ) + entry["item_indices"].append(index) + # Response ordering must not affect the image-to-metadata mapping. + return _EncoderResponse(dict(reversed(list(params.items())))) + + monkeypatch.setattr(proxy, "encode_session", Session()) + monkeypatch.setattr(proxy, "NO_REWRITE", no_rewrite) + monkeypatch.setattr(proxy, "ENCODER_MAX_BATCH_SIZE", batch_size) + monkeypatch.setattr(proxy, "encoder_rr_idx", 0) + images = ["A", "A", "C", "D", "A", "F"] + body = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image}} + for image in images + ], + } + ], + "mm_processor_kwargs": {"max_pixels": 262144}, + } + meta, handles = asyncio.run( + proxy.fanout_encoder_primer( + body, ["http://e0", "http://e1"], "r", "tcp://consumer:1" + ) + ) + assert [len(b["messages"][0]["content"]) for _, b in seen] == expected_sizes + assert ( + len( + { + x["transfer_id"] + for _, b in seen + for x in b["ec_transfer_params"]["ec_items"] + } + ) + == 6 + ) + for url, batch in seen: + assert batch["mm_processor_kwargs"] == body["mm_processor_kwargs"] + for item in batch["messages"][0]["content"]: + assert item["image_url"]["url"] in ( + {"A", "C"} if url.startswith("http://e0/") else {"A", "D", "F"} + ) + assert [meta[i]["image_grid_thw"][-1] for i in range(6)] == list(map(ord, images)) + assert [meta[i]["ec_mm_hash"] for i in range(6)] == [ + x + "-processed" for x in images + ] + assert handles["A-processed"]["peer_port"] == 1234 + assert "ec_transfer_params" not in body + + transfers_by_encoder: list[list[dict[str, str]]] = [[], []] + for url, batch in seen: + rank = 0 if url.startswith("http://e0/") else 1 + transfers_by_encoder[rank].extend(batch["ec_transfer_params"]["ec_items"]) + expected_transfers = [ + transfers_by_encoder[i % 2][i // 2]["transfer_id"] for i in range(len(images)) + ] + params = ( + handles + if no_rewrite + else proxy.rewrite_for_decode(body, meta)["ec_transfer_params"] + ) + request = SimpleNamespace( + ec_transfer_params=params, + mm_features=[ + SimpleNamespace(identifier=image + "-processed") for image in images + ], + ) + assert [ + ECMooncakeScheduler._request_transfer_id(request, i) for i in range(len(images)) + ] == expected_transfers + + +def test_batch_rejects_ambiguous_metadata(proxy, monkeypatch): + monkeypatch.setattr( + proxy, + "encode_session", + _EncoderSession([{"unknown": {"metadata": {"image_grid_thw": [1, 2, 2]}}}]), + ) + body = { + "messages": [ + { + "content": [ + {"type": "image_url", "image_url": {"url": image}} + for image in ["A", "B"] + ] + } + ] + } + with pytest.raises(proxy.HTTPException, match="cannot be matched"): + asyncio.run(proxy.fanout_encoder_primer(body, ["http://e0"], "r")) + + def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch): """The retry loop re-enters `prepare_for_decode` with the same body. @@ -187,20 +311,21 @@ async def backend(request): body = await request.json() seen[stage].append(body) if stage == "encode": - item = body["messages"][0]["content"][0] - mm_hash = item["uuid"] - if body.get("mm_processor_kwargs"): - mm_hash += "-processed" - return web.json_response( - { - "ec_transfer_params": { - mm_hash: { - "metadata": {"image_grid_thw": [[1, 2, 2]]}, - "peer_port": len(seen[stage]), - } - } - } - ) + params: dict[str, dict[str, Any]] = {} + for index, item in enumerate(body["messages"][0]["content"]): + mm_hash = item["uuid"] + if body.get("mm_processor_kwargs"): + mm_hash += "-processed" + entry = params.setdefault( + mm_hash, + { + "metadata": {"image_grid_thw": [[1, 2, 2]]}, + "peer_port": len(seen[stage]), + "item_indices": [], + }, + ) + entry["item_indices"].append(index) + return web.json_response({"ec_transfer_params": params}) if stage == "prefill": assert body["max_tokens"] == 1 and body["stream"] is False return web.json_response( @@ -263,7 +388,7 @@ async def backend(request): finally: await proxy.on_shutdown() - assert len(seen["encode"]) == 4 + assert len(seen["encode"]) == 2 assert len(seen["decode"]) == 2 for requests in seen.values(): for forwarded in requests: @@ -296,7 +421,7 @@ async def backend(request): if request_options: mm_hash += "-processed" assert set(final["ec_transfer_params"]) - {"ec_items"} == {mm_hash} - assert final["ec_transfer_params"][mm_hash]["peer_port"] in (3, 4) + assert final["ec_transfer_params"][mm_hash]["peer_port"] == 2 if prefill: assert final["kv_transfer_params"] == {"remote_block_ids": [2]} @@ -324,3 +449,62 @@ def test_pooled_connections_are_retired_before_the_server_closes_them( assert proxy.decode_session.connector._keepalive_timeout == pooled_for finally: asyncio.run(proxy.on_shutdown()) + + +@pytest.mark.parametrize( + "no_rewrite, transfer", + [(False, None), (True, None), (False, "push"), (False, "handle")], +) +def test_video_audio_fallback(proxy, monkeypatch, no_rewrite, transfer): + """Keep raw video and neighboring images, but never discard transfer handles.""" + import copy + + from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata + + video = {"type": "video_url", "video_url": {"url": "video-with-audio"}} + images = [{"type": "image_url", "image_url": {"url": key}} for key in ("A", "B")] + body: dict[str, Any] = { + "messages": [{"content": [images[0], video, images[1]]}], + "mm_processor_kwargs": {"use_audio_in_video": True}, + } + original = copy.deepcopy(body) + # Collector indices describe two processed features from one video item. + video_metadata = collect_ec_item_metadata( + [SimpleNamespace(identifier=key, data=None) for key in ("video", "audio")], None + ) + if transfer == "handle": + video_metadata["audio"]["transfer_id"] = "reservation" + image_metadata = { + key: {"metadata": {"image_grid_thw": [1, 2, size]}, "item_indices": [i]} + for i, (key, size) in enumerate((("A", 2), ("B", 4))) + } + # Images form the first group, the singleton video the second. + monkeypatch.setattr(proxy, "ENCODER_MAX_BATCH_SIZE", 0) + monkeypatch.setattr(proxy, "NO_REWRITE", no_rewrite) + replies = _EncoderSession([image_metadata, video_metadata]) + monkeypatch.setattr(proxy, "encode_session", replies) + consumer = "tcp://consumer:1" if transfer == "push" else None + preparation = proxy.prepare_for_decode(body, "r", ["http://e0"], "", consumer) + if transfer: + with pytest.raises(proxy.HTTPException, match="cannot be matched"): + asyncio.run(preparation) + return + + prepared, _, _ = asyncio.run(preparation) + assert body == original + assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"] + final = prepared["messages"][0]["content"] + assert final[1] == video + handles = prepared["ec_transfer_params"] + assert not video_metadata.keys() & handles.keys() + if no_rewrite: + assert final == original["messages"][0]["content"] + assert "ec_items" not in handles + else: + for position, image, key in ((0, images[0], "A"), (2, images[1], "B")): + assert final[position] == { + "type": "image_embeds", + "image_embeds": image_metadata[key]["metadata"], + "uuid": proxy.content_uuid(image), + } + assert [item["mm_hash"] for item in handles["ec_items"]] == ["A", "B"] diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index 156d98ae794c..d7ef19b5f4f4 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -67,6 +67,7 @@ def test_request_finished_producer_emits_params(monkeypatch): assert params == { "h1": { "metadata": {}, + "item_indices": [0], "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, @@ -91,6 +92,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): assert params == { "h1": { "metadata": {}, + "item_indices": [0], "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, @@ -110,5 +112,5 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): assert delay is False # The item's placeholder metadata is still reported even though there's # no cache entry to transfer (empty "metadata": no fields, no transfer). - assert params == {"h1": {"metadata": {}}} + assert params == {"h1": {"metadata": {}, "item_indices": [0]}} s.shutdown() diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index f54d6877c1b4..35d26bb44bf9 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -60,15 +60,16 @@ def collect_ec_item_metadata( Keyed by mm_hash, each entry carries a `metadata` dict with whatever placeholder fields `resolver` says this model needs published for its - modality, so a consumer can skip the image transform. `data` is None for - items served from the processor cache, in which case the metadata is + modality, so a consumer can skip the image transform. `item_indices` + identifies every occurrence in `mm_features`, including repeated hashes. + `data` is None for items served from the processor cache, so the metadata is unavailable here and the consumer has to fall back to processing the media itself. A connector that also has transfer coordinates to report (e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside `metadata`, not into it. """ items: dict[str, dict[str, Any]] = {} - for feature in mm_features: + for index, feature in enumerate(mm_features): metadata: dict[str, Any] = {} if feature.data is not None: wanted = resolver.fields_for(feature.modality) @@ -81,7 +82,9 @@ def collect_ec_item_metadata( # Some metadata (e.g. Qwen3-VL video timestamps) is # produced as a plain list rather than a tensor. metadata[key] = value - items[feature.identifier] = {"metadata": metadata} + entry = items.setdefault(feature.identifier, {"item_indices": []}) + entry["metadata"] = metadata + entry["item_indices"].append(index) return items From 9ff08f9493aa745b0559c02cf6148f71d9817169 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 17 Sep 2026 18:46:28 +0800 Subject: [PATCH 395/774] [Bugfix] Honor skip_reading_prefix_cache for KV connector hits (#57269) Signed-off-by: aoshen02 Co-authored-by: aoshen02 --- tests/v1/core/test_scheduler.py | 19 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 2 ++ 2 files changed, 21 insertions(+) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index ef818f27c772..298c8c618576 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2302,6 +2302,25 @@ def test_has_sync_kv_loads( assert output.has_sync_kv_loads is expected_has_sync_loads +def test_kv_connector_honors_skip_reading_prefix_cache(): + """A request that must score every prompt row takes no external hit.""" + BLOCK_SIZE = 16 + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv(matched_tokens=BLOCK_SIZE * 2, is_async=False), + block_size=BLOCK_SIZE, + ) + plain, scoring = create_requests(num_requests=2, num_tokens=BLOCK_SIZE * 4) + scoring.sampling_params.prompt_logprobs = 1 + scoring.skip_reading_prefix_cache = True + scheduler.add_request(plain) + scheduler.add_request(scoring) + + output = scheduler.schedule() + assert output.num_scheduled_tokens[plain.request_id] == BLOCK_SIZE * 2 + assert output.num_scheduled_tokens[scoring.request_id] == BLOCK_SIZE * 4 + + @pytest.mark.parametrize("is_async", [False, True]) def test_kv_connector_basic(is_async: bool): """Test whether Scheduler with KVConnector schedules tokens, allocates diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ed80daa8c279..e0c4997b63ff 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -935,6 +935,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: request, block_aligned_local ) ) + if request.skip_reading_prefix_cache: + ext_tokens, load_kv_async = 0, False if ext_tokens is None: # The request cannot be scheduled because From da3c07bf7dda0c5cbe94d06854426907cd30c7ed Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Thu, 17 Sep 2026 19:57:22 +0800 Subject: [PATCH 396/774] [XPU][CI] skip test_hybrid_prefix_cache_hit_rate (#57301) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index ef2b4ba07396..458bef9518f2 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -98,7 +98,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pip install dpctl && - pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && + pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py --ignore=v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && From 667b26e50bc3287793efd92672c520fa101ecf80 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:02:39 +0200 Subject: [PATCH 397/774] [Bugfix][ROCm][GLM-5.3-Flash] Apply deferred tilelang.jit already on attribute access (#57192) Signed-off-by: simondanielsson --- tests/kernels/test_mhc_tilelang_jit.py | 36 ++++++++++++++---- vllm/tilelang_utils/__init__.py | 51 ++++++++++++++++++-------- 2 files changed, 64 insertions(+), 23 deletions(-) diff --git a/tests/kernels/test_mhc_tilelang_jit.py b/tests/kernels/test_mhc_tilelang_jit.py index 4f39fae02058..1c430e7ae338 100644 --- a/tests/kernels/test_mhc_tilelang_jit.py +++ b/tests/kernels/test_mhc_tilelang_jit.py @@ -24,19 +24,26 @@ class _PassConfigKey: def _install_tilelang_stub( monkeypatch: pytest.MonkeyPatch, ) -> dict[str, int]: - calls = {"jit_decorate": 0, "compiled_call": 0} + calls = {"jit_decorate": 0, "compiled_call": 0, "compiled_compile": 0} tilelang: Any = ModuleType("tilelang") + class _JitImpl: + def __init__(self, func: Any) -> None: + self.func = func + + def __call__(self, *args: Any, **kw: Any) -> Any: + calls["compiled_call"] += 1 + return self.func.__name__ + + def compile(self, *args: Any, **kw: Any) -> Any: + calls["compiled_compile"] += 1 + return self.func.__name__ + def jit(**kwargs: Any) -> Any: def decorate(func: Any) -> Any: calls["jit_decorate"] += 1 - - def compiled(*args: Any, **kw: Any) -> Any: - calls["compiled_call"] += 1 - return func.__name__ - - return compiled + return _JitImpl(func) return decorate @@ -78,6 +85,21 @@ def test_tilelang_jit_decorator_is_lazy_only_on_rocm( assert calls["compiled_call"] == 1 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm") +def test_tilelang_jit_proxies_compile_only_warmup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _install_tilelang_stub(monkeypatch) + module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels" + monkeypatch.delitem(sys.modules, module_name, raising=False) + module = importlib.import_module(module_name) + + assert module.mhc_post_tilelang.compile() == "mhc_post_tilelang" + assert calls["jit_decorate"] == 1 + assert calls["compiled_compile"] == 1 + assert calls["compiled_call"] == 0 + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm") def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols() -> None: if importlib.util.find_spec("tilelang") is None: diff --git a/vllm/tilelang_utils/__init__.py b/vllm/tilelang_utils/__init__.py index 990c2fca3e77..5ac03c858ed0 100644 --- a/vllm/tilelang_utils/__init__.py +++ b/vllm/tilelang_utils/__init__.py @@ -61,8 +61,41 @@ def _get_pass_configs() -> dict[Any, Any]: return pass_configs +class _DeferredTileLangJitKernel: + """Stand-in for a `tilelang.jit` kernel that decorates on first use. + + Both attribute access and calling apply the decoration. Required for + compile-only JIT warmup on platforms that defer import of tilelang. + """ + + _kernel_function: Callable[..., Any] | None = None + _jit_kernel: Any = None + + def __init__(self, kernel_function: Callable[..., Any]) -> None: + self._kernel_function = kernel_function + functools.update_wrapper(self, kernel_function) + + def _ensure_jit_kernel(self) -> Any: + if self._jit_kernel is None: + _ensure_tilelang_imported() + kernel_function = self._kernel_function + assert kernel_function is not None + kernel_function.__globals__["tilelang"] = tilelang + kernel_function.__globals__["T"] = T + self._jit_kernel = tilelang.jit(pass_configs=_get_pass_configs())( + kernel_function + ) + return self._jit_kernel + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self._ensure_jit_kernel()(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._ensure_jit_kernel(), name) + + def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]: - """Apply `tilelang.jit`, deferring until first call on ROCm. + """Apply `tilelang.jit`, deferring until first use on ROCm. ROCm defers JIT decoration so importing the caller's module does not require TileLang immediately. CUDA keeps the eager decoration behavior. @@ -75,18 +108,4 @@ def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]: _ensure_tilelang_imported() return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function) - compiled_kernel: Callable[..., Any] | None = None - - @functools.wraps(kernel_function) - def wrapper(*args: Any, **kwargs: Any) -> Any: - nonlocal compiled_kernel - if compiled_kernel is None: - _ensure_tilelang_imported() - kernel_function.__globals__["tilelang"] = tilelang - kernel_function.__globals__["T"] = T - compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())( - kernel_function - ) - return compiled_kernel(*args, **kwargs) - - return wrapper + return _DeferredTileLangJitKernel(kernel_function) From 91a4c40b45696c4dd31281877ad46af0910efcfa Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:04:56 -0400 Subject: [PATCH 398/774] [Compile] Fix compile warning #177-D (#56609) Signed-off-by: yewentao256 --- csrc/libtorch_stable/cache_kernels.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 887c83f3af87..c5e35c52b82f 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -84,7 +84,9 @@ void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, namespace { // ROCm hipMemcpyBatchAsync faults for count > 8192 (MI350X/gfx950), so chunk // at that ceiling. +#if defined(USE_ROCM) constexpr int64_t kRocmDefaultMaxBatchDescriptors = 8192; +#endif constexpr const char* kMaxBatchDescriptorsEnv = "VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS"; From a85dd55177d59b99903432c71c7376309d7e7757 Mon Sep 17 00:00:00 2001 From: Laura Lopez-Real <113060513+laulopezreal@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:22:57 +0200 Subject: [PATCH 399/774] Fix Laguna patch mutating flat RoPE parameters (#57189) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Laura López Real Co-authored-by: Laura López Real Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_config.py | 37 +++++++++++++++++++++++++++++++ vllm/transformers_utils/config.py | 2 +- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/test_config.py b/tests/test_config.py index 5c892da6c6ac..a07929abfcbf 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -46,6 +46,7 @@ from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel from vllm.platforms import current_platform from vllm.transformers_utils.config import ( + _patch_hf_transformers_nested_rope_validation, get_pooling_config, try_get_dense_modules, ) @@ -54,6 +55,42 @@ DEVICE_TYPE = current_platform.device_type +def test_nested_rope_validation_patch_preserves_flat_rope_parameters(monkeypatch): + calls = [] + + def original_validate_rope(config, *args, **kwargs): + calls.append(config) + + from transformers import PretrainedConfig + + monkeypatch.setattr(PretrainedConfig, "validate_rope", original_validate_rope) + _patch_hf_transformers_nested_rope_validation() + + nested_rope_parameters = { + "full_attention": {"rope_type": "default"}, + "original_max_position_embeddings": 32768, + } + PretrainedConfig.validate_rope( + SimpleNamespace(rope_parameters=nested_rope_parameters) + ) + assert nested_rope_parameters == {"full_attention": {"rope_type": "default"}} + + flat_rope_parameters = { + "rope_type": "linear", + "factor": 8.0, + "rope_theta": 500000.0, + } + PretrainedConfig.validate_rope( + SimpleNamespace(rope_parameters=flat_rope_parameters) + ) + assert flat_rope_parameters == { + "rope_type": "linear", + "factor": 8.0, + "rope_theta": 500000.0, + } + assert len(calls) == 2 + + def _write_json(path: Path, value: object) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value), encoding="utf-8") diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 54f3086431ac..20d0d73afb92 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -273,7 +273,7 @@ def patched_validate_rope(self, *args, **kwargs): layer_types = set(rope_parameters) & set( hf_configuration_utils.ALLOWED_LAYER_TYPES ) - if shared_keys := set(rope_parameters) - layer_types: + if layer_types and (shared_keys := set(rope_parameters) - layer_types): for key in shared_keys: del rope_parameters[key] logger.warning( From 2e50824766ead3e30f53cba711e8c7ffb2bb6b2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mart=C3=ADn=20el=20Cheikh?= <157884230+melcheikh@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:27:54 -0300 Subject: [PATCH 400/774] [Bugfix] Fall back to native sampling when FlashInfer cannot target the GPU (#48956) Signed-off-by: martin el cheikh Co-authored-by: Claude --- tests/v1/sample/test_topk_topp_sampler.py | 38 ++++++++++++++++++++++ vllm/v1/sample/ops/topk_topp_sampler.py | 39 ++++++++++++++++++++++- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 0b570f49d974..8ff16e559917 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -237,6 +237,44 @@ def test_flashinfer_sampler(): ) +@pytest.mark.skipif( + not FLASHINFER_TOPK_TOPP_SUPPORTED, + reason="Requires CUDA and flashinfer", +) +def test_flashinfer_sampler_fallback_when_jit_cannot_target_gpu(monkeypatch): + """FlashInfer swallows arch-detection errors (e.g. SM 12.x with a CUDA + toolkit older than 12.9), leaving an empty target-arch set that makes + every JIT spec fail with a misleading "requires sm75 or higher" error at + first use, killing the engine during startup profiling (issue #42393). + The sampler gate must detect this and fall back to native sampling, and + surface the real reason when the user explicitly opted in. + """ + import flashinfer.jit.core + from flashinfer.compilation_context import CompilationContext + + from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sampler_supported + + def fake_check_cuda_arch(): + raise RuntimeError("FlashInfer requires GPUs with sm75 or higher") + + def fake_normalize_cuda_arch(major, minor): + raise RuntimeError("SM 12.x requires CUDA >= 12.9") + + monkeypatch.setattr(flashinfer.jit.core, "check_cuda_arch", fake_check_cuda_arch) + monkeypatch.setattr( + CompilationContext, + "_normalize_cuda_arch", + staticmethod(fake_normalize_cuda_arch), + ) + + monkeypatch.delenv("VLLM_USE_FLASHINFER_SAMPLER", raising=False) + assert flashinfer_sampler_supported() is False + + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "1") + with pytest.raises(RuntimeError, match="SM 12.x requires CUDA >= 12.9"): + flashinfer_sampler_supported() + + # ============================================================================= # Triton kernel tests # ============================================================================= diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 06b3fc3423f3..477e99c2fabf 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -10,6 +10,7 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform +from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import HAS_TRITON if HAS_TRITON: @@ -41,12 +42,46 @@ def _skip_aiter_sampler_on_gfx1250() -> bool: return on_gfx1250() +def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None: + """Return why FlashInfer JIT codegen cannot target the current GPU, or + None if it can. + + FlashInfer swallows arch-detection errors when building its compilation + context (e.g. SM 12.x with a CUDA toolkit older than 12.9), leaving an + empty target-arch set that makes every JIT spec fail with a misleading + "requires sm75 or higher" error at first use — killing the engine during + startup profiling (https://github.com/vllm-project/vllm/issues/42393). + """ + try: + from flashinfer.jit.core import check_cuda_arch + except ImportError: + return None + try: + check_cuda_arch() + return None + except RuntimeError as e: + reason = str(e) + try: + # Re-derive the real error that FlashInfer swallowed during arch + # detection, e.g. "SM 12.x requires CUDA >= 12.9". + from flashinfer.compilation_context import CompilationContext + + CompilationContext._normalize_cuda_arch(capability.major, capability.minor) + except RuntimeError as e: + reason = str(e) + except Exception: + # FlashInfer internals changed; keep the original error message. + pass + return reason + + def flashinfer_sampler_supported() -> bool: """Decide whether FlashInfer's top-p/top-k sampler can be used. Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER`` is 0, when the platform isn't CUDA, when the GPU's compute capability is - unsupported. Raises ``RuntimeError`` if the user explicitly opted in + unsupported, or when FlashInfer cannot JIT-compile for the current + GPU/CUDA toolchain. Raises ``RuntimeError`` if the user explicitly opted in via the env var but FlashInfer is unavailable. Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``; @@ -73,6 +108,8 @@ def flashinfer_sampler_supported() -> bool: unsupported_reason = ( f"unsupported compute capability {capability.as_version_str()}" ) + else: + unsupported_reason = _flashinfer_jit_unsupported_reason(capability) if unsupported_reason is None: logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global") From f46968cb6e28ff89182bfa6fc291b4f851ab07d1 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 17 Sep 2026 05:42:59 -0700 Subject: [PATCH 401/774] [CI] Make ci-clean-log.sh portable to macOS/BSD sed (#57367) Signed-off-by: khluu --- .buildkite/scripts/ci-clean-log.sh | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/ci-clean-log.sh b/.buildkite/scripts/ci-clean-log.sh index e2e21483d546..2fe45f3b7745 100644 --- a/.buildkite/scripts/ci-clean-log.sh +++ b/.buildkite/scripts/ci-clean-log.sh @@ -10,11 +10,20 @@ fi INPUT_FILE="$1" +# GNU sed uses "-i [-r]" while BSD/macOS sed needs "-i '' [-E]"; detect once. +if sed --version >/dev/null 2>&1; then + sedi() { sed -i "$@"; } + SED_ERE=-r +else + sedi() { sed -i '' "$@"; } + SED_ERE=-E +fi + # Strip timestamps -sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE" +sedi 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE" # Strip Buildkite inline timestamp markers (ESC _bk;t= BEL) -sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE" +sedi 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE" # Strip colorization -sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE" +sedi "$SED_ERE" 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE" From 8be3ca35f311fc205262122aedcf775d56ebc3ac Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:52:02 -0700 Subject: [PATCH 402/774] [ROCm][CI] Fix AMD CI pipeline upload rejected by an invalid block-step key (#57375) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5 (1M context) --- .buildkite/test-amd.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d038d661769f..2f5a84aaf1ea 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2414,7 +2414,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" -- label: ":amd: (MI300) Rust Frontend Serve/Admin Coverage" +- label: ":amd: (MI300) Rust Frontend Serve-Admin Coverage" timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false From 0ef8366ba39fc38b88c60b17ee2ddf05aa8a7596 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:53:11 +0100 Subject: [PATCH 403/774] [Frontend] Only show the summary line of config docstrings in `--help` (#57357) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/utils_/test_argparse_utils.py | 17 +++++++++++++++++ vllm/utils/argparse_utils.py | 25 ++++++++++++++++++------- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/utils_/test_argparse_utils.py b/tests/utils_/test_argparse_utils.py index baf57d994ff9..89e0f9977d0c 100644 --- a/tests/utils_/test_argparse_utils.py +++ b/tests/utils_/test_argparse_utils.py @@ -524,3 +524,20 @@ def test_flat_product(): (3, 4, "a", 5, 6), (3, 4, "b", 5, 6), ] + + +def test_group_description_is_summary_only(): + """Group descriptions are config docstrings, too long for the terminal.""" + parser = FlexibleArgumentParser() + group = parser.add_argument_group( + title="MyConfig", + description="Summary line.\n\nDetails which only belong in the docs.", + ) + group.add_argument("--my-arg") + + assert "Summary line." in parser.format_help() + assert "only belong in the docs" not in parser.format_help() + + parser._search_keyword = "myconfig" + assert "Summary line." in parser.format_help() + assert "only belong in the docs" not in parser.format_help() diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index e573cbc293b3..d214da3402cd 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -93,6 +93,11 @@ def human_readable_int_or_auto(value: str) -> int: return human_readable_int(value) +def summary_line(description: str | None) -> str: + """Return the first paragraph of a config docstring as a single line.""" + return " ".join((description or "").split("\n\n", 1)[0].split()) + + class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter): """SortedHelpFormatter that sorts arguments by their option strings.""" @@ -203,7 +208,7 @@ def format_help(self): for group in self._action_groups: if group.title and group.title.lower() == search_keyword: formatter.start_section(group.title) - formatter.add_text(group.description) + formatter.add_text(summary_line(group.description)) formatter.add_arguments(group._group_actions) formatter.end_section() formatter.add_text(self._json_tip) @@ -242,12 +247,18 @@ def format_help(self): # positionals, optionals and user-defined groups formatter.start_section("Config Groups") config_groups = "" - for group in self._action_groups: - if not group._group_actions: - continue - title = group.title - description = group.description or "" - config_groups += f"{title: <24}{description}\n" + groups = [g for g in self._action_groups if g._group_actions] + titles = [g.title or "" for g in groups] + title_width = max((len(title) for title in titles), default=0) + 1 + width = max(formatter._width - formatter._current_indent, title_width + 20) + for title, group in zip(titles, groups): + lines = textwrap.wrap( + summary_line(group.description), + width, + initial_indent=f"{title: <{title_width}}", + subsequent_indent=" " * title_width, + ) + config_groups += ("\n".join(lines) or title) + "\n" formatter.add_text(config_groups) formatter.end_section() From 75c71390d5b399f5397a9166920fc45902f99f14 Mon Sep 17 00:00:00 2001 From: CaiJohn Date: Thu, 17 Sep 2026 20:57:47 +0800 Subject: [PATCH 404/774] [Bugfix] Escape control characters in xgrammar choice grammar (#48115) Signed-off-by: john <6712432+CaiJohn@users.noreply.github.com> Co-authored-by: john <6712432+CaiJohn@users.noreply.github.com> --- tests/v1/structured_output/test_utils.py | 46 ++++++++++++++++++++++++ vllm/v1/structured_output/utils.py | 15 ++++++-- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index f5f80202205b..3e4a12f7e5c1 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -5,9 +5,12 @@ from xgrammar import Grammar from xgrammar.testing import _is_grammar_accept_string +from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, + validate_xgrammar_grammar, ) +from vllm.v1.structured_output.utils import choice_as_grammar pytestmark = pytest.mark.cpu_test @@ -232,3 +235,46 @@ def test_pattern_properties_constrains_keys_and_values(pattern_properties_schema assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') + + +class TestIsGrammarAcceptString: + class TestPR48115Regressions: + @pytest.mark.parametrize("codepoint", [*range(0x20), 0x7F]) + def test_choice_as_grammar_preserves_control_characters(self, codepoint): + choice = f"a{chr(codepoint)}f" + grammar = Grammar.from_ebnf(choice_as_grammar([choice])) + + assert _is_grammar_accept_string(grammar, choice) + assert not _is_grammar_accept_string(grammar, "a") + assert not _is_grammar_accept_string(grammar, "af") + assert not _is_grammar_accept_string(grammar, f"a\\u{codepoint:04x}f") + + @pytest.mark.parametrize( + ("choice", "other"), + [ + (r"a\nf", "a\nf"), + ('a quote " and a backslash \\', 'a quote " and a backslash '), + ("café", "cafe"), + ("日本語", "日本"), + ("😀", "😁"), + ], + ) + def test_choice_as_grammar_preserves_literal_choices(self, choice, other): + grammar = Grammar.from_ebnf(choice_as_grammar(["yes", choice])) + + assert _is_grammar_accept_string(grammar, "yes") + assert _is_grammar_accept_string(grammar, choice) + assert not _is_grammar_accept_string(grammar, other) + assert not _is_grammar_accept_string(grammar, choice + "extra") + + def test_validate_xgrammar_preserves_multiline_choice(self): + structured_outputs = StructuredOutputsParams(choice=["yes", "no\nplease"]) + validate_xgrammar_grammar( + SamplingParams(structured_outputs=structured_outputs) + ) + + assert structured_outputs.choice is None + grammar = Grammar.from_ebnf(structured_outputs.grammar) + assert _is_grammar_accept_string(grammar, "yes") + assert _is_grammar_accept_string(grammar, "no\nplease") + assert not _is_grammar_accept_string(grammar, r"no\nplease") diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 4969789434c2..4bb438d6f380 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -567,9 +567,18 @@ def extract_references(text: str) -> set[str]: def choice_as_grammar(choice: list[str]) -> str: def escape_ebnf_string(s: str) -> str: - """Escape special characters in a EBNF string.""" - # Escape double quotes and backslashes - return re.sub(r'(["\\])', r"\\\1", s) + """Escape EBNF literals, including raw LF, CR, and NUL terminators.""" + escapes = {"\\": r"\\", '"': r"\"", "\n": r"\n", "\r": r"\r", "\t": r"\t"} + + def escape_char(ch: str) -> str: + if ch in escapes: + return escapes[ch] + # Escape remaining C0 controls (U+0000-U+001F) and DEL (U+007F). + if ord(ch) < 0x20 or ord(ch) == 0x7F: + return f"\\u{ord(ch):04x}" + return ch + + return "".join(escape_char(ch) for ch in s) escaped_choices = (escape_ebnf_string(c) for c in choice) grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices) From fe284ea17d606ceba4a5d3d55d3e3d9a34f478d9 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 17 Sep 2026 09:35:34 -0400 Subject: [PATCH 405/774] [BugFxi] Fix DeepGEMM FP8 workspace over allocation (#53914) Signed-off-by: Lucas Wilkinson Signed-off-by: Lucas Wilkinson Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/kernels/moe/test_deepgemm.py | 40 +++++++++++++++++++ .../layers/fused_moe/deep_gemm_utils.py | 3 +- .../layers/fused_moe/experts/deep_gemm_moe.py | 33 ++++++++++++++- 3 files changed, 73 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index f6df64c51bf7..cee0faa59229 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -26,6 +26,10 @@ ) from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( deepgemm_moe_permute, + ep_gather, +) +from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( + _fp8_workspace_shape, ) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -43,6 +47,42 @@ BLOCK_SIZE = [128, 128] +def test_ep_gather_uses_64_bit_row_offsets(): + hidden_size = 4096 + source_row = 1 << 18 + input_tensor = torch.empty( + (source_row + 1, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + input_tensor[source_row].fill_(1) + output = torch.empty((1, hidden_size), device="cuda", dtype=torch.bfloat16) + + ep_gather( + input_tensor=input_tensor, + recv_topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64), + recv_topk_weight=torch.ones((1, 1), device="cuda", dtype=torch.float32), + input_index=torch.full((1, 1), source_row, device="cuda", dtype=torch.int32), + expert_map=None, + output_tensor=output, + ) + + torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) + + +@pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("num_columns", [2048, 6144, 6145]) +def test_fp8_workspace_shape(workspace_dtype, num_columns): + num_rows = 17 + shape = _fp8_workspace_shape(num_rows, num_columns, workspace_dtype) + + allocated_bytes = math.prod(shape) * workspace_dtype.itemsize + required_bytes = num_rows * num_columns * torch.float8_e4m3fn.itemsize + + assert allocated_bytes >= required_bytes + assert allocated_bytes - required_bytes < num_rows * workspace_dtype.itemsize + + @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init): hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 4ea864eaac3b..33a229695561 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -393,12 +393,13 @@ def _fwd_kernel_ep_gather( source_token_index = tl.load( input_index + cur_token * input_index_stride0 + topk_index ) + source_token_index_i64 = source_token_index.to(tl.int64) acc_weight = tl.load( recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index ) tmp = tl.load( input_tensor - + source_token_index * input_tensor_stride0 + + source_token_index_i64 * input_tensor_stride0 + cur_block * BLOCK_D + off_d ) diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index fe95b36858d7..699b029bad86 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -50,6 +50,21 @@ logger = init_logger(__name__) +def _fp8_workspace_shape( + num_rows: int, num_columns: int, workspace_dtype: torch.dtype +) -> tuple[int, int]: + """Size an FP8 byte buffer stored in a workspace of another dtype.""" + bytes_per_workspace_element = workspace_dtype.itemsize + fp8_columns_per_workspace_element = ( + bytes_per_workspace_element // torch.float8_e4m3fn.itemsize + ) + assert bytes_per_workspace_element % torch.float8_e4m3fn.itemsize == 0 + return ( + num_rows, + -(-num_columns // fp8_columns_per_workspace_element), + ) + + def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool: align = get_mk_alignment_for_contiguous_layout()[0] return align <= M and N % align == 0 and K % align == 0 @@ -216,7 +231,14 @@ def workspace_shapes( assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (M_sum, max(activation_out_dim, K)) + # workspace1 is allocated in the activation dtype by the workspace + # manager, but is only ever viewed and used as FP8 in apply(). Size it + # by bytes instead of reserving one BF16/FP16 element per FP8 element. + workspace1 = _fp8_workspace_shape( + M_sum, + max(activation_out_dim, K), + self.workspace_dtype(self.moe_config.in_dtype), + ) workspace2 = (M_sum, max(N, K)) output = (M, K) return (workspace1, workspace2, output) @@ -480,7 +502,14 @@ def workspace_shapes( assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (M_sum, max(activation_out_dim, K)) + # workspace1 holds the permuted and requantized FP8 activations. The + # workspace manager allocates it in the model activation dtype, so + # account for the dtype sizes rather than overallocating BF16/FP16. + workspace1 = _fp8_workspace_shape( + M_sum, + max(activation_out_dim, K), + self.workspace_dtype(self.moe_config.in_dtype), + ) workspace2 = (M_sum, max(N, K)) output = (M, K) return (workspace1, workspace2, output) From 0eae9acd4d01574e12d4ecf6a0229813f7fdb799 Mon Sep 17 00:00:00 2001 From: Shantipriya Parida Date: Thu, 17 Sep 2026 16:39:26 +0300 Subject: [PATCH 406/774] [Bugfix][ROCm][MoE] Fall back instead of crashing when AITER MoE is requested for a non-gated (is_act_and_mul=False) model (#56590) Signed-off-by: Shantipriya Parida --- .../moe/test_unquantized_backend_selection.py | 60 +++++++++++++++++++ .../layers/fused_moe/oracle/unquantized.py | 28 ++++++++- 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index d4c037d978a3..9d0f4d6900c4 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -6,6 +6,7 @@ import torch from tests.kernels.moe.utils import make_dummy_moe_config +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, @@ -228,6 +229,65 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer): assert expert_cls is not None +@patch( + "vllm.utils.flashinfer.has_flashinfer", + return_value=False, +) +@patch( + "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops." + "is_fused_moe_enabled", + return_value=True, +) +@patch( + "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops." + "is_rdna_aiter_enabled", + return_value=False, +) +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific backend selection test" +) +def test_select_rocm_aiter_backend_non_gated_activation_falls_back( + mock_rdna_disabled, mock_aiter_enabled, mock_has_flashinfer, monkeypatch +): + """Test ROCm backend selection falls back (not raises) for non-gated MoE.""" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + + with patch( + "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" + ) as mock_platform: + mock_platform.is_cuda.return_value = False + mock_platform.is_rocm.return_value = True + mock_platform.is_cpu.return_value = False + mock_platform.is_xpu.return_value = False + mock_platform.is_tpu.return_value = False + mock_platform.is_out_of_tree.return_value = False + + moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL) + assert moe_config.is_act_and_mul is False + + selected_backend, expert_cls = select_unquantized_moe_backend( + moe_config=moe_config, + ) + + assert selected_backend != UnquantizedMoeBackend.AITER + assert expert_cls is not None + + +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific backend selection test" +) +def test_explicit_aiter_backend_non_gated_activation_still_raises(): + """Explicit `--moe-backend aiter` still raises for non-gated MoE; + only the env-var opt-in path falls back.""" + moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL) + moe_config.moe_backend = "aiter" # explicit pin, not "auto" + assert moe_config.is_act_and_mul is False + + with pytest.raises(ValueError): + select_unquantized_moe_backend(moe_config=moe_config) + + @patch( "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config", return_value=(True, None), diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 910178705933..d1437933c3ef 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -305,7 +305,33 @@ def _return_or_raise( AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER) else: backend = UnquantizedMoeBackend.AITER - return _return_or_raise(backend, moe_config, activation_format) + reason = None + for k_cls in backend_to_kernel_cls(backend): + supported, reason = k_cls.is_supported_config( + k_cls, moe_config, None, None, activation_format + ) + if supported: + logger.info_once(_make_log_backend(backend)) + return backend, k_cls + # AITER was explicitly requested but does not support this + # deployment configuration (e.g. a non-gated MoE activation, + # which AiterExperts._supports_no_act_and_mul() rejects + # unconditionally). Rather than hard-crashing at engine init, + # warn — explicitly stating that this is falling back, and + # which backends will be tried next — and fall back to the + # remaining backends in priority order below. + logger.warning_once( + "VLLM_ROCM_USE_AITER_MOE=1 was requested, but %s " + "Falling back to try the remaining available MoE backends: %s.", + _make_log_unsupported(backend, reason), + ", ".join( + b.value + for b in AVAILABLE_BACKENDS + if b != UnquantizedMoeBackend.AITER + ), + ) + if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS: + AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER) for backend in AVAILABLE_BACKENDS: for k_cls in backend_to_kernel_cls(backend): From 117cf43fb65713a0a7a3c35b23be078d5fb319d4 Mon Sep 17 00:00:00 2001 From: Rukhaiya2004 <162583766+Rukhaiya2004@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:15:13 +0530 Subject: [PATCH 407/774] [HARDWARE][POWER] Enable W8A8 INT8 MoE on POWER (#55316) Signed-off-by: Rukhaiya Signed-off-by: Akash kaothalkar Co-authored-by: Akash kaothalkar Co-authored-by: Antigravity Co-authored-by: Li, Jiang --- cmake/cpu_extension.cmake | 6 + csrc/cpu/cpu_fused_moe_int8.cpp | 19 +- csrc/cpu/cpu_types_vsx.hpp | 31 +- .../micro_gemm/cpu_micro_gemm_int8_vsx.hpp | 287 ++++++++++++++++++ csrc/cpu/torch_bindings.cpp | 9 +- tests/kernels/moe/test_cpu_fused_moe.py | 22 +- .../layers/fused_moe/experts/cpu_moe.py | 177 +++++++++++ .../layers/fused_moe/oracle/int8.py | 8 +- 8 files changed, 537 insertions(+), 22 deletions(-) create mode 100644 csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index ff2911dc11c8..6359a96887a5 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -522,6 +522,12 @@ if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) ${VLLM_EXT_SRC}) endif() +if (POWER10_FOUND OR POWER11_FOUND) + set(VLLM_EXT_SRC + "csrc/cpu/cpu_fused_moe_int8.cpp" + ${VLLM_EXT_SRC}) +endif() + if(USE_ONEDNN) set(VLLM_EXT_SRC "csrc/cpu/dnnl_kernels.cpp" diff --git a/csrc/cpu/cpu_fused_moe_int8.cpp b/csrc/cpu/cpu_fused_moe_int8.cpp index 649317b4786a..27d0f6a7d539 100644 --- a/csrc/cpu/cpu_fused_moe_int8.cpp +++ b/csrc/cpu/cpu_fused_moe_int8.cpp @@ -26,10 +26,23 @@ #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON: #endif +#if defined(__powerpc64__) + #include "cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp" + #define VSX_DISPATCH(SCALAR_TYPE, ...) \ + case cpu_utils::ISA::VSX: { \ + using gemm_t = \ + cpu_micro_gemm::MicroGemmINT8; \ + return __VA_ARGS__(); \ + } +#else + #define VSX_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::VSX: +#endif + #define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \ [&] { \ switch (ISA_TYPE) { \ NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ + VSX_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ default: { \ TORCH_CHECK(false, "Invalid CPU ISA type."); \ } \ @@ -81,9 +94,9 @@ void fused_moe_int8_impl( const int32_t thread_num = cpu_utils::get_max_threads(); const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( - gemm_m_tile_size * input_size_13 * sizeof(int8_t)); - const int32_t w2_input_buffer_size = - cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t)); + (gemm_m_tile_size * input_size_13 + 64) * sizeof(int8_t)); + const int32_t w2_input_buffer_size = cpu_utils::round_up<64>( + (gemm_m_tile_size * input_size_2 + 64) * sizeof(int8_t)); const int32_t w13_n_tile_size = [&]() { const int64_t cache_size = cpu_utils::get_available_l2_size(); diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index d46906130e92..eaea921d98b8 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -13,10 +13,10 @@ namespace vec_op { struct fp8_e4m3_tag {}; struct fp8_e5m2_tag {}; -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +// Note: FP16 (Half) is not supported on POWER VSX — no FP16Vec16 type. #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) @@ -416,6 +416,7 @@ struct FP32Vec8 : public Vec { tmp.val[0] = two_x; tmp.val[1] = two_x; FP32Vec8 temp_vec(tmp); + FP32Vec8 exp_vec = temp_vec.exp(); vector float e = temp_vec.exp().reg.val[0]; vector float num = vec_sub(e, one); @@ -634,6 +635,16 @@ struct FP32Vec16 : public Vec { reg.val[3] = vec_ctf(v.reg.val[3], 0); } + explicit FP32Vec16(const float* __restrict__ base, const INT32Vec16& index) { + INT32Vec16::AliasReg idx; + idx.reg = index.reg; + AliasReg ar; + for (int i = 0; i < VEC_ELEM_NUM; ++i) { + ar.values[i] = base[idx.values[i]]; + } + reg = ar.reg; + } + FP32Vec16 operator*(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]), vec_mul(reg.val[1], b.reg.val[1]), @@ -648,6 +659,13 @@ struct FP32Vec16 : public Vec { vec_add(reg.val[3], b.reg.val[3])})); } + FP32Vec16 operator-() const { + const __vector float zero = vec_splats(0.0f); + return FP32Vec16( + f32x4x4_t({vec_sub(zero, reg.val[0]), vec_sub(zero, reg.val[1]), + vec_sub(zero, reg.val[2]), vec_sub(zero, reg.val[3])})); + } + FP32Vec16 operator-(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_sub(reg.val[0], b.reg.val[0]), vec_sub(reg.val[1], b.reg.val[1]), @@ -655,11 +673,6 @@ struct FP32Vec16 : public Vec { vec_sub(reg.val[3], b.reg.val[3])})); } - FP32Vec16 operator-() const { - return FP32Vec16( - f32x4x4_t({-reg.val[0], -reg.val[1], -reg.val[2], -reg.val[3]})); - } - FP32Vec16 operator/(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]), vec_div(reg.val[1], b.reg.val[1]), diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp new file mode 100644 index 000000000000..60333537281d --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp @@ -0,0 +1,287 @@ +#ifndef CPU_MICRO_GEMM_INT8_VSX_HPP +#define CPU_MICRO_GEMM_INT8_VSX_HPP + +#include +#include +#include +#include +#include + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#include +#include +#include + +namespace cpu_micro_gemm { + +namespace vsx_int8 { + +constexpr int32_t K = 4; +constexpr int32_t Mr = 4; +constexpr int32_t Nr = 32; + +// --------------------------------------------------------------------------- +// 8-Accumulator MMA 4x32 micro-GEMM (POWER10 xvi8ger4pp). +// +// Computes 4 rows x 32 columns (128 outputs) using all 8 hardware accumulators. +// Operates on pre-biased uint8 weights (B + 128) and subtracts 128 * a_sums. +// Works for all m in [1, 4] with zero scalar loop overhead. +// --------------------------------------------------------------------------- +__attribute__((target("cpu=power10"))) FORCE_INLINE void gemm_mma_4x32( + const int8_t* __restrict__ a_packed, const uint8_t* __restrict__ b_panel, + const int32_t* __restrict__ a_sums, int32_t* __restrict__ c_ptr, + const int32_t m, const int32_t k_size, const int64_t ldc) { + TORCH_CHECK(m > 0 && m <= Mr); + TORCH_CHECK(k_size % K == 0); + + __vector_quad acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7; + __builtin_mma_xxsetaccz(&acc0); + __builtin_mma_xxsetaccz(&acc1); + __builtin_mma_xxsetaccz(&acc2); + __builtin_mma_xxsetaccz(&acc3); + __builtin_mma_xxsetaccz(&acc4); + __builtin_mma_xxsetaccz(&acc5); + __builtin_mma_xxsetaccz(&acc6); + __builtin_mma_xxsetaccz(&acc7); + + const int8_t* __restrict__ ap = a_packed; + const uint8_t* __restrict__ bp = b_panel; + +#pragma GCC unroll 4 + for (int32_t k = 0; k < k_size; k += 4) { + const __vector unsigned char avec = + vec_xl(0, reinterpret_cast(ap)); + __builtin_mma_xvi8ger4pp(&acc0, avec, vec_xl(0, bp)); + __builtin_mma_xvi8ger4pp(&acc1, avec, vec_xl(16, bp)); + __builtin_mma_xvi8ger4pp(&acc2, avec, vec_xl(32, bp)); + __builtin_mma_xvi8ger4pp(&acc3, avec, vec_xl(48, bp)); + __builtin_mma_xvi8ger4pp(&acc4, avec, vec_xl(64, bp)); + __builtin_mma_xvi8ger4pp(&acc5, avec, vec_xl(80, bp)); + __builtin_mma_xvi8ger4pp(&acc6, avec, vec_xl(96, bp)); + __builtin_mma_xvi8ger4pp(&acc7, avec, vec_xl(112, bp)); + ap += 16; + bp += 128; + } + + __vector signed int res[4]; +#define STORE_ACC(ACC, CG_IDX) \ + { \ + __builtin_mma_disassemble_acc(res, &ACC); \ + for (int r = 0; r < m; ++r) { \ + const __vector signed int vcorr = vec_splats(128 * a_sums[r]); \ + const __vector signed int val = vec_sub(res[r], vcorr); \ + vec_xst(val, 0, \ + reinterpret_cast(c_ptr + r * ldc + (CG_IDX) * 4)); \ + } \ + } + STORE_ACC(acc0, 0); + STORE_ACC(acc1, 1); + STORE_ACC(acc2, 2); + STORE_ACC(acc3, 3); + STORE_ACC(acc4, 4); + STORE_ACC(acc5, 5); + STORE_ACC(acc6, 6); + STORE_ACC(acc7, 7); +#undef STORE_ACC +} + +} // namespace vsx_int8 + +// ============================================================================= +// MicroGemmINT8 specialisation for VSX (POWER10+) +// ============================================================================= +template +class MicroGemmINT8 { + public: + static constexpr int32_t K = vsx_int8::K; // 4 + static constexpr int32_t Mr = vsx_int8::Mr; // 4 + static constexpr int32_t Nr = vsx_int8::Nr; // 32 + static constexpr int32_t NrGemv = Nr; + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = Nr; + static constexpr int32_t WeightOCGroupSize = Nr; + static_assert(MaxMSize % Mr == 0); + + // ------------------------------------------------------------------------- + // quantize_row: scalar_t -> INT8 + per-token scale + // Single-pass 16-element vector conversion with aligned 16-byte stores. + // ------------------------------------------------------------------------- + static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output, + float& scale, const int32_t size) { + TORCH_CHECK_EQ(size % 16, 0); + + const __vector signed short z = vec_splats((signed short)0); + + auto load_lo_hi = + [&](int32_t i) -> std::pair<__vector float, __vector float> { + if constexpr (std::is_same_v) { + const float* p = reinterpret_cast(input) + i; + return {vec_xl(0, p), vec_xl(16, p)}; + } else { + const signed short* p = + reinterpret_cast(input + i); + const __vector signed short raw = (__vector signed short)vec_xl(0, p); + return {(__vector float)vec_mergeh(z, raw), + (__vector float)vec_mergel(z, raw)}; + } + }; + + __vector float max_v = vec_splats(0.0f); + for (int32_t i = 0; i < size; i += 8) { + auto [lo, hi] = load_lo_hi(i); + max_v = vec_max(max_v, vec_abs(lo)); + max_v = vec_max(max_v, vec_abs(hi)); + } + __vector float t = vec_max(max_v, vec_sld(max_v, max_v, 8)); + t = vec_max(t, vec_sld(t, t, 4)); + const float abs_max = std::max(vec_extract(t, 0), 1.0e-7f); + scale = abs_max / 127.0f; + const __vector float inv_scale = vec_splats(127.0f / abs_max); + + for (int32_t i = 0; i < size; i += 16) { + auto [lo0, hi0] = load_lo_hi(i); + auto [lo1, hi1] = load_lo_hi(i + 8); + + const __vector signed int i0 = + vec_cts(vec_round(vec_mul(lo0, inv_scale)), 0); + const __vector signed int i1 = + vec_cts(vec_round(vec_mul(hi0, inv_scale)), 0); + const __vector signed int i2 = + vec_cts(vec_round(vec_mul(lo1, inv_scale)), 0); + const __vector signed int i3 = + vec_cts(vec_round(vec_mul(hi1, inv_scale)), 0); + + const __vector signed short s0 = vec_packs(i0, i1); + const __vector signed short s1 = vec_packs(i2, i3); + const __vector signed char s8 = vec_packs(s0, s1); + + vec_xst(s8, 0, reinterpret_cast(output + i)); + } + } + + // ------------------------------------------------------------------------- + // dequantize_tile: INT32 accumulator -> FP32 scaled output + // ------------------------------------------------------------------------- + static FORCE_INLINE void dequantize_tile( + int32_t* input, float* output, const float* __restrict__ input_scales, + const float* __restrict__ weight_scales, const int32_t m, const int32_t n, + const int32_t stride) { + TORCH_CHECK_EQ(n % 4, 0); + for (int32_t mi = 0; mi < m; ++mi) { + const __vector float a_scale = vec_splats(input_scales[mi]); + for (int32_t ni = 0; ni < n; ni += 4) { + const __vector signed int i32 = vec_xl(0, input + mi * stride + ni); + const __vector float wscale = vec_xl(0, weight_scales + ni); + vec_xst(vec_mul(vec_ctf(i32, 0), vec_mul(a_scale, wscale)), 0, + output + mi * stride + ni); + } + } + } + + // ------------------------------------------------------------------------- + // pack_input_from_rows: pack rows into 4-row MMA tiles and compute row sums + // Header: 64 bytes (contains int32 a_sums[0..7]) + // Data: a_packed + 64 (contains 4x4 tiles of A) + // ------------------------------------------------------------------------- + static void pack_input_from_rows(const int8_t* const* __restrict__ rows, + int8_t* __restrict__ a_packed, + const int32_t m, const int32_t k) { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK_EQ(k % K, 0); + + int32_t* a_sums = reinterpret_cast(a_packed); + for (int32_t r = 0; r < 8; ++r) a_sums[r] = 0; + + int8_t* a_data = a_packed + 64; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* const* panel_rows = rows + row_base; + int32_t* panel_sums = a_sums + row_base; + int8_t* dst_panel = a_data + (row_base / Mr) * (k * Mr); + + for (int32_t ki = 0; ki < k; ki += 4) { + int8_t* dst = dst_panel + (ki / 4) * 16; + for (int32_t r = 0; r < 4; ++r) { + if (r < panel_m) { + const int8_t* src = panel_rows[r] + ki; + dst[r * 4 + 0] = src[0]; + dst[r * 4 + 1] = src[1]; + dst[r * 4 + 2] = src[2]; + dst[r * 4 + 3] = src[3]; + panel_sums[r] += + static_cast(src[0]) + src[1] + src[2] + src[3]; + } else { + dst[r * 4 + 0] = 0; + dst[r * 4 + 1] = 0; + dst[r * 4 + 2] = 0; + dst[r * 4 + 3] = 0; + } + } + } + } + } + + // ------------------------------------------------------------------------- + // pack_weight: [output_size, input_size] INT8 -> packed layout + // Pre-biases with +128 in [N/32][K/4][8 col-groups][4 cols x 4 k] layout + // ------------------------------------------------------------------------- + static void pack_weight(const int8_t* __restrict__ weight, + int8_t* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK_EQ(output_size % NSize, 0); + TORCH_CHECK_EQ(input_size % K, 0); + + uint8_t* dst_base = reinterpret_cast(packed_weight); + + for (int32_t o = 0; o < output_size; o += Nr) { + uint8_t* dst_panel = dst_base + o * input_size; + for (int32_t k = 0; k < input_size; k += 4) { + for (int32_t cg = 0; cg < 8; ++cg) { + uint8_t* dst = dst_panel + (k / 4) * 128 + cg * 16; + for (int32_t col = 0; col < 4; ++col) { + for (int32_t ki = 0; ki < 4; ++ki) { + const int8_t w = weight[(o + cg * 4 + col) * input_size + k + ki]; + dst[col * 4 + ki] = + static_cast(static_cast(w) + 128); + } + } + } + } + } + } + + // ------------------------------------------------------------------------- + // gemm: 8-accumulator MMA GEMM over 32-column tiles + // ------------------------------------------------------------------------- + void gemm(const int8_t* __restrict__ a_packed, + const int8_t* __restrict__ b_packed, int32_t* __restrict__ c, + const int32_t m, const int32_t k, const int64_t b_n_group_stride, + const int64_t ldc) const { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK_EQ(k % K, 0); + (void)b_n_group_stride; + + const int32_t* a_sums = reinterpret_cast(a_packed); + const int8_t* a_data = a_packed + 64; + const uint8_t* b_u8 = reinterpret_cast(b_packed); + + for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr) { + const uint8_t* b_panel = b_u8 + n_idx * k; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* a_panel = a_data + (row_base / Mr) * (k * Mr); + const int32_t* panel_sums = a_sums + row_base; + + vsx_int8::gemm_mma_4x32(a_panel, b_panel, panel_sums, + c + row_base * ldc + n_idx, panel_m, k, ldc); + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif // CPU_MICRO_GEMM_INT8_VSX_HPP diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 8f663a93c36d..88170068dea3 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -883,8 +883,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ - !defined(__APPLE__) + +#if (defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ + !defined(__APPLE__)) || \ + defined(__powerpc64__) ops.def( "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, " "str isa) -> ()"); @@ -895,8 +897,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool " "skip_weighted, str act, str isa) -> ()"); ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8); -#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && - // !defined(__APPLE__) +#endif // ARM_I8MM+BF16 or __powerpc64__ ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index cf83ef6fc00a..d4c543722131 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -306,17 +306,28 @@ def test_cpu_fused_moe( @pytest.mark.skipif( - current_platform.get_cpu_architecture() != CpuArchEnum.ARM, - reason="Requires Arm CPU", + current_platform.get_cpu_architecture() + not in (CpuArchEnum.ARM, CpuArchEnum.POWERPC), + reason="Requires Arm or POWER CPU", ) @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("expert_num", EXPERT_NUM) @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) @pytest.mark.parametrize("use_bias", USE_BIAS) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "dtype", + ( + [torch.float32, torch.float16, torch.bfloat16] + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM + else [torch.float32, torch.bfloat16] + ), +) @pytest.mark.parametrize("act", ACT) -@pytest.mark.parametrize("isa", ["neon"]) +@pytest.mark.parametrize( + "isa", + ["neon"] if current_platform.get_cpu_architecture() == CpuArchEnum.ARM else ["vsx"], +) def test_cpu_fused_moe_int8( batch_size: int, expert_num: int, @@ -386,7 +397,8 @@ def test_cpu_fused_moe_int8( isa, ) - torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) + atol = rtol = 1e-1 if isa == "vsx" else 2e-2 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) # moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 7b3193fb5f88..858b35186448 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -950,6 +950,29 @@ def apply( # =========================================================================== +def prepare_int8_moe_layer_for_cpu( + w13: torch.Tensor, + w2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepack INT8 MoE weights for the current CPU architecture.""" + # SMMLA packing for AArch64 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return ( + cpu_prepack_moe_weight_int8(w13, "neon"), + cpu_prepack_moe_weight_int8(w2, "neon"), + ) + # VSX packing for POWER + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return ( + cpu_prepack_moe_weight_int8(w13, "vsx"), + cpu_prepack_moe_weight_int8(w2, "vsx"), + ) + # VNNI packing for x86 + packed_w13 = torch.ops._C.convert_weight_packed(w13) + packed_w2 = torch.ops._C.convert_weight_packed(w2) + return packed_w13, packed_w2 + + class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation modular MoE experts.""" @@ -1283,6 +1306,160 @@ def apply( ) +class PowerCPUExpertsInt8(mk.FusedMoEExpertsModular): + """POWER VSX INT8 MoE with per-token activation and channelwise weight quant.""" + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExpertsModular.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + return supported, reason + if moe_config.in_dtype not in (torch.float32, torch.bfloat16): + return False, "kernel requires float32 or bfloat16 activations" + if moe_config.hidden_dim % 16 != 0: + return False, "kernel requires hidden dim divisible by 16" + if moe_config.intermediate_size_per_partition % 16 != 0: + return False, "kernel requires intermediate dim divisible by 16" + return True, None + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC + and hasattr(torch.ops._C, "cpu_fused_moe_int8") + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + ) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return not moe_parallel_config.use_ep + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == ( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "vsx") + w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "vsx") + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # cpu_fused_moe_int8 manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + # The modular prepare step applies router weights to the input when + # apply_router_weight_on_input is enabled. + assert self.w1_scale is not None + assert self.w2_scale is not None + cpu_fused_moe_int8( + output, + hidden_states, + w1, + w2, + self.w1_scale, + self.w2_scale, + self.w1_bias, + self.w2_bias, + topk_weights, + topk_ids, + activation.value, + "vsx", + skip_weighted=apply_router_weight_on_input, + ) + + class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): """AMD Zen INT8 MoE with per-token activation and channelwise weight quantization, dispatched through zentorch.""" diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 3b8a3c4f878d..1a23b0173663 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -71,10 +71,16 @@ def backend_to_kernel_cls( from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( ArmCPUExpertsInt8, CPUExpertsInt8, + PowerCPUExpertsInt8, ZenCPUExpertsInt8, ) - return [ZenCPUExpertsInt8, ArmCPUExpertsInt8, CPUExpertsInt8] + return [ + ZenCPUExpertsInt8, + PowerCPUExpertsInt8, + ArmCPUExpertsInt8, + CPUExpertsInt8, + ] else: raise ValueError(f"Unknown Int8 MoE backend: {backend.value}") From 78fdf4efb86dffde1b0f5733207aed038c3530f8 Mon Sep 17 00:00:00 2001 From: Yuchen Wang <93700456+yuchenwang3@users.noreply.github.com> Date: Thu, 17 Sep 2026 07:00:04 -0700 Subject: [PATCH 408/774] [Bugfix][MoE] Convert FlashInfer BF16 weights in place (#54699) Signed-off-by: Yuchen Wang <93700456+yuchenwang3@users.noreply.github.com> Signed-off-by: Dakai An Co-authored-by: Dakai An Co-authored-by: OpenAI Codex Co-authored-by: aoshen02 --- tests/kernels/moe/test_flashinfer.py | 258 ++++++++++++++++++ .../moe/test_moe_weight_loading_padded.py | 126 +++++++++ tests/kernels/moe/test_trtllm_bf16_moe.py | 8 +- .../model_loader/test_reload.py | 105 +++++++ .../layers/fused_moe/routed_experts.py | 18 ++ .../fused_moe/unquantized_fused_moe_method.py | 122 +++++++-- .../quantization/utils/flashinfer_utils.py | 86 +++--- .../model_loader/reload/meta.py | 4 +- .../model_loader/reload/utils.py | 8 +- 9 files changed, 667 insertions(+), 68 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2e022ede3627..53d760c339c8 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import multiprocessing from dataclasses import dataclass import pytest @@ -433,6 +434,263 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout( assert w13_converted.shape[0] == num_experts assert w2_converted.shape[0] == num_experts + assert w13_converted.data_ptr() == w13.data_ptr() + assert w2_converted.data_ptr() == w2.data_ptr() + + +@pytest.mark.parametrize("is_gated_act_gemm", [True, False]) +def test_convert_moe_weights_to_flashinfer_trtllm_block_layout_values( + is_gated_act_gemm, +): + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) + + num_experts, intermediate, hidden = 2, 256, 256 + w13_multiplier = 2 if is_gated_act_gemm else 1 + w13 = torch.randn( + (num_experts, w13_multiplier * intermediate, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + w2 = torch.randn( + (num_experts, hidden, intermediate), + dtype=torch.bfloat16, + device="cuda", + ) + + def _reference_block_layout( + weight: torch.Tensor, + is_w13: bool, + cache: dict[torch.Size, torch.Tensor], + ) -> torch.Tensor: + outputs = [] + for expert in weight: + expert_uint8 = expert.view(torch.uint8) + if is_w13: + indices = _maybe_get_cached_w3_w1_permute_indices( + cache, + expert_uint8, + 128, + is_gated_act_gemm=is_gated_act_gemm, + ) + if is_gated_act_gemm: + indices = (indices + expert_uint8.shape[0] // 2) % ( + expert_uint8.shape[0] + ) + else: + indices = get_w2_permute_indices_with_cache( + cache, + expert_uint8, + 128, + ) + rows, cols = expert_uint8.shape + blocks = expert_uint8.view(rows, cols // 128, 128).permute(1, 0, 2) + outputs.append(torch.index_select(blocks, 1, indices.to(weight.device))) + return torch.stack(outputs).view(torch.bfloat16) + + reference_cache: dict[torch.Size, torch.Tensor] = {} + expected_w13 = _reference_block_layout(w13, is_w13=True, cache=reference_cache) + expected_w2 = _reference_block_layout(w2, is_w13=False, cache=reference_cache) + w13_ptr = w13.data_ptr() + w2_ptr = w2.data_ptr() + + actual_w13, actual_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout( + {}, + w13, + w2, + is_gated_act_gemm=is_gated_act_gemm, + ) + + assert actual_w13.data_ptr() == w13_ptr + assert actual_w2.data_ptr() == w2_ptr + assert torch.equal(actual_w13, expected_w13) + assert torch.equal(actual_w2, expected_w2) + + +def _make_unquantized_flashinfer_test_layer( + monkeypatch, intermediate, is_gated, *, device="cuda" +): + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=256, + intermediate_size=intermediate, + activation=MoEActivation.SILU if is_gated else MoEActivation.RELU2_NO_MUL, + ) + moe_config.intermediate_size_per_partition = (intermediate + 127) // 128 * 128 + method = object.__new__(UnquantizedFusedMoEMethod) + method.moe = moe_config + method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM + method.moe_kernel = None + monkeypatch.setattr( + method, + "_init_moe_kernel", + lambda _: setattr(method, "moe_kernel", object()), + ) + + layer = torch.nn.Module() + layer.moe_config = moe_config + with torch.device(device): + method.create_weights( + layer, + num_experts=2, + hidden_size=256, + intermediate_size_per_partition=moe_config.intermediate_size_per_partition, + params_dtype=torch.bfloat16, + ) + return method, layer + + +@pytest.mark.parametrize("intermediate", [192, 256]) +@pytest.mark.parametrize("is_gated", [True, False]) +def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( + monkeypatch, intermediate, is_gated +): + """Each raw reload clears padding left dirty by prior in-place conversion.""" + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) + + method, layer = _make_unquantized_flashinfer_test_layer( + monkeypatch, intermediate, is_gated + ) + w13_shape, w2_shape = layer.w13_weight.shape, layer.w2_weight.shape + padded = w2_shape[-1] + w13_ptr = layer.w13_weight.data_ptr() + w2_ptr = layer.w2_weight.data_ptr() + + for _ in range(2): + reloaded_w13 = torch.randn_like(layer.w13_weight) + reloaded_w2 = torch.randn_like(layer.w2_weight) + reloaded_w13[:, intermediate:padded].zero_() + if is_gated: + reloaded_w13[:, padded + intermediate :].zero_() + reloaded_w2[:, :, intermediate:].zero_() + expected_w13, expected_w2 = ( + convert_moe_weights_to_flashinfer_trtllm_block_layout( + {}, + reloaded_w13.clone(), + reloaded_w2.clone(), + is_gated_act_gemm=is_gated, + ) + ) + + # The loader writes only logical checkpoint slices, not padding. + layer.w13_weight.fill_(float("nan")) + layer.w2_weight.fill_(float("nan")) + layer.w13_weight[:, :intermediate].copy_(reloaded_w13[:, :intermediate]) + if is_gated: + layer.w13_weight[:, padded : padded + intermediate].copy_( + reloaded_w13[:, padded : padded + intermediate] + ) + layer.w2_weight[:, :, :intermediate].copy_(reloaded_w2[:, :, :intermediate]) + method.process_weights_after_loading(layer) + + assert layer.w13_weight.shape == w13_shape + assert layer.w2_weight.shape == w2_shape + assert layer.w13_weight.data_ptr() == w13_ptr + assert layer.w2_weight.data_ptr() == w2_ptr + kernel_w13, kernel_w2 = method._kernel_weights(layer) + assert torch.equal(kernel_w13, expected_w13) + assert torch.equal(kernel_w2, expected_w2) + + +def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index): + from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader + from vllm.model_executor.utils import weights_already_processed + + torch.accelerator.set_device_index(device_index) + with pytest.MonkeyPatch.context() as monkeypatch: + method, layer = _make_unquantized_flashinfer_test_layer( + monkeypatch, 192, is_gated, device="meta" + ) + loader = object.__new__(IpcModelLoader) + loader.mode = mode + loader._apply_entries(layer, entries, {}, device_index) + pointers = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr()) + + # Like ipc_cache: a fresh method, only tensor metadata, no _setup_kernel. + with weights_already_processed(): + method.process_weights_after_loading(layer) + assert method.moe_kernel is not None + actual = method._kernel_weights(layer) + for weight, reference, pointer in zip(actual, expected, pointers): + assert weight.data_ptr() == pointer + assert torch.equal(weight.cpu(), reference) + # Cross-process addresses need not match; writes prove sharing/isolation. + with torch.no_grad(): + for weight in actual: + weight.fill_(1.0) + torch.accelerator.synchronize() + + +@pytest.mark.parametrize("is_gated", [True, False]) +@pytest.mark.parametrize("cache_ndim", [3, 4]) +@pytest.mark.parametrize("mode", ["copy", "zero_copy"]) +def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( + monkeypatch, is_gated, cache_ndim, mode +): + """A spawned IPC consumer restores views without transient method state.""" + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) + from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry + + method, layer = _make_unquantized_flashinfer_test_layer( + monkeypatch, 192, is_gated, device="meta" + ) + packed = convert_moe_weights_to_flashinfer_trtllm_block_layout( + {}, + torch.randn_like(layer.w13_weight, device="cuda"), + torch.randn_like(layer.w2_weight, device="cuda"), + is_gated_act_gemm=is_gated, + ) + expected = [weight.cpu() for weight in packed] + entries = {} + for name, weight in zip(("w13_weight", "w2_weight"), packed): + cached = weight.view_as(getattr(layer, name)) if cache_ndim == 3 else weight + entries[name] = TensorEntry.from_tensor(cached, kind="param") + consumer = multiprocessing.get_context("spawn").Process( + target=_check_flashinfer_ipc_weights, + args=( + entries, + expected, + is_gated, + mode, + torch.accelerator.current_device_index(), + ), + ) + consumer.start() + try: + # Keep producer allocations alive until the consumer releases its views. + consumer.join(timeout=180) + assert not consumer.is_alive(), "IPC consumer timed out" + assert consumer.exitcode == 0, "IPC consumer failed; see child traceback" + for source, reference in zip(packed, expected): + expected_source = ( + torch.ones_like(reference) if mode == "zero_copy" else reference + ) + assert torch.equal(source.cpu(), expected_source) + finally: + if consumer.is_alive(): + consumer.terminate() + consumer.join(timeout=10) + if consumer.is_alive(): + consumer.kill() + consumer.join(timeout=10) + consumer.close() @pytest.mark.parametrize( diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index d3716a196b02..423f18113254 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -12,7 +12,14 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, +) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) from .utils import make_dummy_moe_config @@ -249,6 +256,57 @@ def test_narrowed_data_shares_storage(self): class TestWeightLoadingWithPaddedHiddenSize: """Integration-style tests that simulate padded weight loading.""" + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("tp_rank", [0, 1]) + def test_load_w2_chunks_noncontiguous_cpu_source_to_padded_cuda( + self, dtype, tp_rank + ): + hidden = 4096 + intermediate = 1024 + loaded_weight = ( + torch.arange(hidden * intermediate * 2) + .remainder(251) + .to(dtype) + .reshape(hidden, intermediate * 2) + ) + tp_source = loaded_weight[ + :, intermediate * tp_rank : intermediate * (tp_rank + 1) + ] + expert_data_full = torch.zeros( + hidden + 8, intermediate + 8, device="cuda", dtype=dtype + ) + destination = expert_data_full[:hidden, :intermediate] + + assert not tp_source.is_contiguous() + assert tp_source.nbytes > 1 << 20 + assert not destination.is_contiguous() + + experts = object.__new__(RoutedExperts) + torch.nn.Module.__init__(experts) + experts.moe_config = make_dummy_moe_config() + experts.moe_config.moe_parallel_config.tp_size = 2 + + torch.accelerator.synchronize() + allocated_before = torch.accelerator.memory_allocated() + torch.accelerator.reset_peak_memory_stats() + experts._load_w2( + expert_data=expert_data_full, + shard_dim=1, + loaded_weight=loaded_weight, + tp_rank=tp_rank, + ) + torch.accelerator.synchronize() + peak_extra = torch.accelerator.max_memory_allocated() - allocated_before + + # A single copy into the strided CUDA view allocates a full-shard + # temporary. Correct values alone would not catch that regression. + assert peak_extra < tp_source.nbytes + + torch.testing.assert_close(destination.cpu(), tp_source) + assert torch.count_nonzero(expert_data_full[hidden:, :]) == 0 + assert torch.count_nonzero(expert_data_full[:hidden, intermediate:]) == 0 + def test_load_w2_with_padding(self): """Simulate loading w2 weights when hidden_size is padded.""" padded_hidden = 3072 @@ -364,6 +422,74 @@ def test_narrow_shard_dim(self): ) +class TestUnquantizedTrtLlmPrePadding: + @staticmethod + def _make_method( + intermediate: int, + backend: UnquantizedMoeBackend = UnquantizedMoeBackend.FLASHINFER_TRTLLM, + ) -> UnquantizedFusedMoEMethod: + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=64, + intermediate_size=intermediate, + ) + method = object.__new__(UnquantizedFusedMoEMethod) + method.moe = moe_config + method.unquantized_backend = backend + method.moe_kernel = None + return method + + @pytest.mark.parametrize( + "backend,original,expected", + [ + (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1344, 1408), + (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1408, 1408), + (UnquantizedMoeBackend.TRITON, 1344, 1344), + ], + ) + def test_rounds_intermediate_before_weight_allocation( + self, + backend: UnquantizedMoeBackend, + original: int, + expected: int, + ): + method = self._make_method(original, backend) + + hidden, intermediate = method.maybe_roundup_sizes( + hidden_size=64, + intermediate_size_per_partition=original, + act_dtype=torch.bfloat16, + moe_parallel_config=method.moe.moe_parallel_config, + ) + + assert hidden == 64 + assert intermediate == expected + + @pytest.mark.parametrize("is_act_and_mul", [True, False]) + def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool): + method = self._make_method(1344) + if not is_act_and_mul: + method.moe.activation = MoEActivation.RELU2_NO_MUL + method.moe.intermediate_size_per_partition = 1408 + layer = torch.nn.Module() + + method.create_weights( + layer=layer, + num_experts=2, + hidden_size=64, + intermediate_size_per_partition=1408, + params_dtype=torch.bfloat16, + ) + + up_mult = 2 if is_act_and_mul else 1 + assert layer.w13_weight.shape == (2, up_mult * 1408, 64) + assert layer.w2_weight.shape == (2, 64, 1408) + assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0 + if is_act_and_mul: + assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0 + assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0 + + class TestLoadWeightsExpertBias: """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize all-zero per-expert `.bias` tensors for models whose experts have no bias diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index 9e50b4fc3bab..f7466c63c895 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/tests/kernels/moe/test_trtllm_bf16_moe.py @@ -75,6 +75,10 @@ def test_trtllm_bf16_moe_modular_no_graph( a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 + # The FlashInfer conversion may rewrite unpadded input storage in place. + # Preserve the original layout for the independent torch reference. + reference_w1 = w1.clone() + reference_w2 = w2.clone() score = torch.randn((m, e), device="cuda", dtype=dtype) scores = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(scores, topk) @@ -133,8 +137,8 @@ def test_trtllm_bf16_moe_modular_no_graph( torch_output = torch_moe( a, - w1, - w2, + reference_w1, + reference_w2, score, topk, activation=MoEActivation.SILU, diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b4a544d2fc0f..844149ce1b31 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -5,6 +5,7 @@ import inspect import sys import types +from types import SimpleNamespace from unittest.mock import Mock from weakref import WeakKeyDictionary, ref @@ -467,6 +468,110 @@ def test_model_cleanup(dist_init, default_vllm_config): assert len(mock_info_dict) == 0 +@pytest.mark.parametrize("is_gated", [False, True]) +@pytest.mark.parametrize("has_bias", [False, True]) +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("padded", [False, True]) +def test_padded_moe_reload_releases_each_layer( + monkeypatch, is_gated, has_bias, tp_rank, padded +): + """Checkpoint-sized copies finish each layer without global finalization.""" + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + hidden, intermediate, experts = 4, 3, 2 + stored_hidden, stored_intermediate = (8, 8) if padded else (hidden, intermediate) + config = SimpleNamespace( + hidden_dim_unpadded=hidden, + intermediate_size_per_partition_unpadded=intermediate, + is_act_and_mul=is_gated, + has_bias=has_bias, + tp_rank=tp_rank, + moe_parallel_config=SimpleNamespace(tp_size=2), + ) + model = torch.nn.ModuleList() + processed: list[torch.nn.Module] = [] + for _ in range(2): + method = object.__new__(UnquantizedFusedMoEMethod) + torch.nn.Module.__init__(method) + method.moe = config + # The regression concerns streaming reload, not kernel conversion. + monkeypatch.setattr(method, "process_weights_after_loading", processed.append) + layer = object.__new__(RoutedExperts) + torch.nn.Module.__init__(layer) + layer.moe_config = config + layer.quant_config = None + layer.quant_method = method + layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) + layer._loaded_expert_biases = set() + method.create_weights( + layer, + experts, + stored_hidden, + stored_intermediate, + torch.float32, + weight_loader=layer.weight_loader, + ) + model.append(layer) + + record_metadata_for_reloading(model) + original_params = [dict(layer.named_parameters()) for layer in model] + shards = ["w1", "w3", "w2"] if is_gated else ["w1", "w2"] + for cycle in range(2): + initialize_layerwise_reload(model) + for layer_index, layer in enumerate(model): + info = reload_layerwise.get_layerwise_info(layer) + inputs = [] + expected = { + name: torch.full_like(p, float("nan")) + for name, p in original_params[layer_index].items() + } + params = dict(layer.named_parameters()) + calls = [ + (e, s, b) + for e in range(experts) + for s in shards + for b in ([False, True] if has_bias else [False]) + ] + for call_index, (expert, shard, bias) in enumerate(calls): + name = ("w2" if shard == "w2" else "w13") + ( + "_bias" if bias else "_weight" + ) + shape = ( + ((hidden,) if bias else (hidden, 2 * intermediate)) + if shard == "w2" + else ((2 * intermediate,) if bias else (2 * intermediate, hidden)) + ) + weight = torch.arange( + torch.Size(shape).numel(), dtype=torch.float32 + ).reshape(shape) + weight = weight + 100 * (1 + call_index + cycle) + inputs.append(ref(weight)) + # Direct checkpoint loading is the reference for deferred reload. + layer.weight_loader(expected[name], weight, name, shard, expert) + param = params[name] + param.weight_loader(param, weight, name, shard, expert) + del weight + if call_index != len(calls) - 1: + assert info.can_load(), ( + "Layer processed before its final checkpoint shard" + ) + + assert not info.can_load(), ( + "Padding must not defer the layer until finalization" + ) + assert not info.loaded_weights + assert len(processed) == cycle * len(model) + layer_index + 1 + assert all(source() is None for source in inputs) + for name, original in original_params[layer_index].items(): + assert getattr(layer, name) is original + # Kernel-specific tests cover padding, which is not checkpoint data. + mask = torch.isfinite(expected[name]) + assert torch.equal(original[mask], expected[name][mask]) + + def test_get_numel_loaded(): param = torch.empty(10, device="meta") loaded_weight = torch.empty(10) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 42ea15dcda7d..e3a287ddc2da 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -27,6 +27,7 @@ QuantizationConfig, resolve_quant_method, ) +from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts @@ -566,6 +567,23 @@ def _load_w2( hidden_dim=hidden_dim, shard_dim=shard_dim, ) + if ( + loaded_weight.device.type == "cpu" + and expert_data.device.type == "cuda" + and not expert_data.is_contiguous() + and expert_data.ndim == 2 + ): + # A strided CPU-to-CUDA copy can allocate a full-shard CUDA + # temporary. Make the TP slice contiguous on CPU and copy rows + # in chunks to limit the temporary used for the padded view. + loaded_weight = loaded_weight.contiguous() + num_chunks = max(cdiv(loaded_weight.nbytes, 1 << 20), 1) + for dst, src in zip( + expert_data.chunk(num_chunks, dim=0), + loaded_weight.chunk(num_chunks, dim=0), + ): + dst.copy_(src) + return expert_data.copy_(loaded_weight) def _load_single_value( diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 523d1df0d80e..a1b2471f5d41 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -11,6 +11,7 @@ from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, + FusedMoEParallelConfig, FusedMoEQuantConfig, biased_moe_quant_config, ) @@ -33,6 +34,7 @@ set_weight_attrs, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -68,22 +70,38 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): + unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded + assert unpadded_intermediate is not None + has_intermediate_padding = ( + intermediate_size_per_partition > unpadded_intermediate + ) if self.moe.is_act_and_mul: w13_up_dim = 2 * intermediate_size_per_partition else: w13_up_dim = intermediate_size_per_partition + unpadded_hidden = self.moe.hidden_dim_unpadded + assert unpadded_hidden is not None + unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1) # Fused gate_up_proj (column parallel) - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - w13_up_dim, - hidden_size, - dtype=params_dtype, - ), - requires_grad=False, + w13_data = torch.empty( + num_experts, + w13_up_dim, + hidden_size, + dtype=params_dtype, ) + if has_intermediate_padding: + # The checkpoint loader only writes the unpadded slices. Zero the + # tails now so no full-size padding copy is needed after loading. + w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_() + if self.moe.is_act_and_mul: + w13_data[ + :, + intermediate_size_per_partition + unpadded_intermediate :, + ].zero_() + w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) + w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden if self.moe.has_bias: w13_bias = torch.nn.Parameter( torch.zeros(num_experts, w13_up_dim, dtype=params_dtype), @@ -91,18 +109,22 @@ def create_weights( ) layer.register_parameter("w13_bias", w13_bias) set_weight_attrs(w13_bias, extra_weight_attrs) + w13_bias.weight_loader_numel = num_experts * unpadded_up # down_proj (row parallel) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=params_dtype, - ), - requires_grad=False, + w2_data = torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, ) + if has_intermediate_padding: + w2_data[:, :, unpadded_intermediate:].zero_() + w2_weight = torch.nn.Parameter(w2_data, requires_grad=False) layer.register_parameter("w2_weight", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) + w2_weight.weight_loader_numel = ( + num_experts * unpadded_hidden * unpadded_intermediate + ) if self.moe.has_bias: w2_bias = torch.nn.Parameter( torch.zeros(num_experts, hidden_size, dtype=params_dtype), @@ -110,6 +132,26 @@ def create_weights( ) layer.register_parameter("w2_bias", w2_bias) set_weight_attrs(w2_bias, extra_weight_attrs) + w2_bias.weight_loader_numel = num_experts * unpadded_hidden + + def maybe_roundup_sizes( + self, + hidden_size: int, + intermediate_size_per_partition: int, + act_dtype: torch.dtype, + moe_parallel_config: FusedMoEParallelConfig, + ) -> tuple[int, int]: + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes( + hidden_size, + intermediate_size_per_partition, + act_dtype, + moe_parallel_config, + ) + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + intermediate_size_per_partition = round_up( + intermediate_size_per_partition, 128 + ) + return hidden_size, intermediate_size_per_partition def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor: # Pad the weight tensor. This is an optimization on ROCm platform, which @@ -135,6 +177,21 @@ def _setup_kernel( w13: torch.Tensor, w2: torch.Tensor, ) -> None: + w13_parameter_shape = w13.shape + w2_parameter_shape = w2.shape + + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + unpadded = self.moe.intermediate_size_per_partition_unpadded + assert unpadded is not None + intermediate = w2.shape[-1] + # Reloads overwrite only checkpoint slices; the previous in-place + # permutation can leave nonzero values in the raw padding slots. + if intermediate > unpadded: + w13[:, unpadded:intermediate].zero_() + if self.moe.is_act_and_mul: + w13[:, intermediate + unpadded :].zero_() + w2[:, :, unpadded:].zero_() + # Shuffle weights to runtime format. w13_new, w2_new = convert_to_unquantized_kernel_format( self.unquantized_backend, @@ -142,6 +199,10 @@ def _setup_kernel( w13_weight=w13, w2_weight=w2, ) + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + w13_new = w13_new.view(w13_parameter_shape) + w2_new = w2_new.view(w2_parameter_shape) + # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__; # On the first call we replace the parameter normally. On subsequent # calls (e.g. RL weight updates that re-trigger @@ -170,6 +231,25 @@ def _setup_kernel( assert self.moe_kernel is not None self.moe_kernel.fused_experts.process_weights_after_loading(layer) + def _kernel_weights( + self, layer: "RoutedExperts" + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM: + return layer.w13_weight, layer.w2_weight + + # IPC caches restore tensor shapes/storage, not method attributes. + # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements. + def block_view(weight: torch.Tensor) -> torch.Tensor: + if weight.ndim == 4: + return weight + experts, rows, cols = weight.shape + return weight.view(experts, cols // 64, rows, 64) + + return ( + block_view(layer.w13_weight), + block_view(layer.w2_weight), + ) + def _init_moe_kernel(self, layer: "RoutedExperts") -> None: """Build the MoE kernel from the layer's current (shuffled) weights.""" self.moe_quant_config = self.get_fused_moe_quant_config(layer) @@ -274,10 +354,11 @@ def forward_native( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert self.moe_kernel is not None + w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=w13, + w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -315,10 +396,11 @@ def apply_monolithic( ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None + w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply_monolithic( x, - layer.w13_weight, - layer.w2_weight, + w13, + w2, router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 044285ac8513..1bfd7e61f90a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -112,8 +112,9 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout( ) -> tuple[torch.Tensor, torch.Tensor]: """Convert expert weights to FlashInfer's block layout. - This reorders W13 and W2 into the expected epilogue-tiled block layout and - returns the shuffled weight tensors. + This reorders W13 and W2 in place into the expected epilogue-tiled block + layout and returns views of the shuffled weight tensors. Using one expert + as scratch space avoids allocating another full copy of both weights. """ if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16: raise ValueError( @@ -147,51 +148,50 @@ def _copy_permuted_expert_to_block_layout( out=out, ) - w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape - w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape - w13_weights_shuffled_tensor = torch.empty( - (num_experts, w13_cols // block_k, w13_rows, block_k), - dtype=torch.uint8, - device=w13_weight.device, - ) - w2_weights_shuffled_tensor = torch.empty( - (num_experts, w2_cols // block_k, w2_rows, block_k), - dtype=torch.uint8, - device=w2_weight.device, - ) - - for i in range(num_experts): - w13_expert_uint8 = w13_weight[i].view(torch.uint8) - - permute_indices = _maybe_get_cached_w3_w1_permute_indices( - cache_permute_indices, - w13_expert_uint8, - epilogue_tile_m, - is_gated_act_gemm=is_gated_act_gemm, - ) - if is_gated_act_gemm: - rows = w13_expert_uint8.shape[0] - permute_indices = (permute_indices + rows // 2) % rows - _copy_permuted_expert_to_block_layout( - w13_weights_shuffled_tensor[i], - w13_expert_uint8, - permute_indices, + def _convert_weight_in_place( + weight: torch.Tensor, + is_w13: bool, + ) -> torch.Tensor: + rows, cols = weight[0].view(torch.uint8).shape + block_layout_shape = (num_experts, cols // block_k, rows, block_k) + expert_scratch = torch.empty( + block_layout_shape[1:], + dtype=torch.uint8, + device=weight.device, ) - permute_indices = get_w2_permute_indices_with_cache( - cache_permute_indices, - w2_weight[i].view(torch.uint8), - epilogue_tile_m, - ) - _copy_permuted_expert_to_block_layout( - w2_weights_shuffled_tensor[i], - w2_weight[i].view(torch.uint8), - permute_indices, - ) + for i in range(num_experts): + expert_uint8 = weight[i].view(torch.uint8) + if is_w13: + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + cache_permute_indices, + expert_uint8, + epilogue_tile_m, + is_gated_act_gemm=is_gated_act_gemm, + ) + if is_gated_act_gemm: + permute_indices = ( + permute_indices + expert_uint8.shape[0] // 2 + ) % expert_uint8.shape[0] + else: + permute_indices = get_w2_permute_indices_with_cache( + cache_permute_indices, + expert_uint8, + epilogue_tile_m, + ) + + _copy_permuted_expert_to_block_layout( + expert_scratch, + expert_uint8, + permute_indices, + ) + expert_uint8.view(-1).copy_(expert_scratch.view(-1)) + + return weight.view(torch.uint8).view(block_layout_shape).view(torch.bfloat16) return ( - w13_weights_shuffled_tensor.view(torch.bfloat16), - w2_weights_shuffled_tensor.view(torch.bfloat16), + _convert_weight_in_place(w13_weight, is_w13=True), + _convert_weight_in_place(w2_weight, is_w13=False), ) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 4d08d4916536..9874ebd67be9 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -9,7 +9,7 @@ from .sanitize import restore_layer_refs, sanitize_layer_refs from .types import LayerReloadingInfo, LayerTensors -from .utils import get_layer_params_buffers, get_layer_tensors +from .utils import get_layer_params_buffers, get_layer_tensors, get_tensor_load_numel __all__ = [ "to_meta_tensor", @@ -206,5 +206,5 @@ def get_numel_loaded( numel = counter.copied_numel param = args.arguments.get("param", None) if isinstance(param, torch.Tensor): - numel = min(numel, param.numel()) + numel = min(numel, get_tensor_load_numel(param)) return numel, return_value diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py index b7dc361503d9..203aa82af303 100644 --- a/vllm/model_executor/model_loader/reload/utils.py +++ b/vllm/model_executor/model_loader/reload/utils.py @@ -11,6 +11,7 @@ "get_layer_tensors", "get_layer_params_buffers", "get_layer_size", + "get_tensor_load_numel", "has_device_tensors", "get_info_size", ] @@ -30,6 +31,11 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors: ) +def get_tensor_load_numel(tensor: torch.Tensor) -> int: + """Count checkpoint elements, excluding padding declared by the weight creator.""" + return getattr(tensor, "weight_loader_numel", tensor.numel()) + + def get_layer_size(layer: torch.nn.Module) -> int: """Calculate total number of elements across loadable tensors in a layer. @@ -39,7 +45,7 @@ def get_layer_size(layer: torch.nn.Module) -> int: from .meta import SKIP_LOAD_TENSORS return sum( - tensor.numel() + get_tensor_load_numel(tensor) for name, tensor in get_layer_tensors(layer).items() if name not in SKIP_LOAD_TENSORS ) From 861a299adc72bbbb1b60853b6173075465a6b425 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:11:14 -0400 Subject: [PATCH 409/774] [Bugfix][Model Runner V2] Route dummy tokens to MoE experts during profiling (#57270) Signed-off-by: Robert Shaw Co-authored-by: Claude Opus 5 (1M context) --- tests/v1/worker/test_gpu_input_batch_v2.py | 25 ++++++++++++++++++++++ vllm/v1/worker/gpu/input_batch.py | 3 ++- vllm/v1/worker/gpu/model_runner.py | 3 +++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py index 8937f2dc18ef..13a92d2c0f05 100644 --- a/tests/v1/worker/test_gpu_input_batch_v2.py +++ b/tests/v1/worker/test_gpu_input_batch_v2.py @@ -6,6 +6,9 @@ import pytest import torch +from vllm.config import VllmConfig +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.platforms import current_platform from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.gpu import cp_utils @@ -56,6 +59,28 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA top-k.") +@pytest.mark.parametrize("is_padding", [True, False]) +def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool): + """Dummy tokens marked as padding are dropped by the MoE router (top-k id + -1), which is wanted for idle DP ranks. Profile runs must not mark them, or + no token reaches the experts and MoE memory is never profiled.""" + monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") + num_tokens = 16 + buffers = InputBuffers( + max_num_reqs=4, max_num_tokens=num_tokens, device=torch.device(DEVICE) + ) + batch = InputBatch.make_dummy(4, num_tokens, buffers, is_padding=is_padding) + hidden_states = torch.randn(num_tokens, 4, device=DEVICE) + router_logits = torch.randn(num_tokens, 8, device=DEVICE) + + with set_forward_context(None, VllmConfig(), is_padding=batch.is_padding): + _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False) + + assert bool((topk_ids == -1).all()) is is_padding + assert bool((topk_ids >= 0).all()) is not is_padding + + def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch): """The batch must view the caller-owned buffer, sliced to padded length. diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index d3c65bd06f11..32c083e45974 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -125,6 +125,7 @@ def make_dummy( num_tokens: int, input_buffers: InputBuffers, max_query_len: int | None = None, + is_padding: bool = True, ) -> "InputBatch": assert 0 < num_reqs <= num_tokens device = input_buffers.device @@ -168,7 +169,7 @@ def make_dummy( input_ids = input_buffers.input_ids[:num_tokens].zero_() positions = input_buffers.positions[:num_tokens].zero_() - input_buffers.is_padding[:num_tokens].fill_(True) + input_buffers.is_padding[:num_tokens].fill_(is_padding) is_padding = input_buffers.is_padding[:num_tokens] logits_indices = query_start_loc[1:] - 1 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 97571df8a901..631e6b04d426 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1733,6 +1733,9 @@ def execute_model( batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + # Profiling and warmup must route the dummy tokens to experts + # so MoE memory is measured and MoE kernels are exercised. + is_padding=not is_profile, ) if self.pcp_manager is not None: input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) From 9e1cd262cbf43b3f3839a28cfa151e2e4b27e2d8 Mon Sep 17 00:00:00 2001 From: Alexander Lee Date: Thu, 17 Sep 2026 10:44:36 -0400 Subject: [PATCH 410/774] [Bugfix] Use DP index for dense DP weight updates and EC CPU region (#56950) Signed-off-by: lxy-alexander Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5 Co-authored-by: aoshen02 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/cpu/test_ec_shared_region.py | 28 +++++++++++++++++++ .../worker/test_gpu_worker_weight_transfer.py | 18 +++++++++--- .../ec_transfer/ec_connector/cpu/common.py | 4 +-- vllm/v1/worker/gpu_worker.py | 2 +- 4 files changed, 45 insertions(+), 7 deletions(-) diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index 722c0d89fa21..c168970147b0 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -13,6 +13,8 @@ import pytest import torch +from vllm.config import ParallelConfig +from vllm.distributed.ec_transfer.ec_connector.cpu import common from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( ECSharedRegion, _wait_for_file_size, @@ -95,6 +97,32 @@ def test_second_instance_opens_existing_file_and_shares_memory(): r1.cleanup() +def test_dense_dp_engines_get_distinct_regions(monkeypatch): + """Dense DP engines reset data_parallel_rank to 0 but share instance_id, + so the region must be keyed by data_parallel_index.""" + monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8) + instance_id = str(uuid.uuid4()) + regions = [] + try: + for dp_index in range(2): + parallel_config = ParallelConfig( + data_parallel_size=2, data_parallel_rank=dp_index + ) + parallel_config.reconfigure_for_independent_dp_rank() + cfg = MagicMock() + cfg.instance_id = instance_id + cfg.parallel_config = parallel_config + cfg.model_config.dtype = torch.float16 + cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} + regions.append(common.create_ec_shared_region(cfg)) + + assert all(r._is_creator for r in regions) + assert regions[0]._mmap_path != regions[1]._mmap_path + finally: + for r in regions: + r.cleanup() + + def test_only_creator_unlinks_file_on_cleanup(): """Critical contract: if the non-creator unlinks, the creator's mmap path becomes a dangling backing file and a third opener would create a new one diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index e576ebd03177..537653bda38d 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -11,7 +11,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config from vllm.lora.layers import BaseLayerWithLoRA from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload from vllm.v1.worker.gpu_worker import Worker @@ -136,12 +136,22 @@ def test_rank_local_update_selects_worker_payload(rank, expected): assert worker._weight_update_active is True -def test_rank_local_update_includes_data_parallel_rank(): +def test_rank_local_update_uses_data_parallel_index_after_reconfigure(): engine = _RecordingEngine() worker = _make_worker(engine) worker.rank = 0 - worker.vllm_config.parallel_config.data_parallel_size = 4 - worker.vllm_config.parallel_config.data_parallel_rank = 2 + parallel_config = ParallelConfig( + data_parallel_size=4, + data_parallel_rank=2, + ) + assert parallel_config.data_parallel_rank == 2 + assert parallel_config.data_parallel_index == 2 + + parallel_config.reconfigure_for_independent_dp_rank() + assert parallel_config.data_parallel_rank == 0 + assert parallel_config.data_parallel_index == 2 + + worker.vllm_config.parallel_config = parallel_config Worker.start_weight_update(worker) Worker.update_weights( diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index e2c37b72a6be..eff16e08b0ae 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -99,8 +99,8 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: ec_config = vllm_config.ec_transfer_config assert ec_config is not None, "ec_transfer_config required to build region" - dp_rank = vllm_config.parallel_config.data_parallel_rank - engine_id = f"{vllm_config.instance_id}_dp{dp_rank}" + dp_index = vllm_config.parallel_config.data_parallel_index + engine_id = f"{vllm_config.instance_id}_dp{dp_index}" dtype = vllm_config.model_config.dtype hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a288a97cc167..02f599e57e8c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1449,7 +1449,7 @@ def update_weights(self, update_info: dict | list[dict]) -> None: if isinstance(update_info, list): parallel_config = self.vllm_config.parallel_config local_update_info = update_info[ - parallel_config.data_parallel_rank * parallel_config.world_size + parallel_config.data_parallel_index * parallel_config.world_size + self.rank ] else: From 438434b5b5f8936a4dc7d651c708928c90f085d6 Mon Sep 17 00:00:00 2001 From: akii96 Date: Thu, 17 Sep 2026 17:44:54 +0300 Subject: [PATCH 411/774] [ROCm][Perf] Insert MiniMax-M3 sparse-PA K/V without a contiguous copy (#56849) Signed-off-by: Aakif Nawaz --- tests/kernels/test_minimax_m3_amd_ops.py | 76 ++++++++++++++++++++++++ vllm/models/minimax_m3/amd/model.py | 30 +++++++++- 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 997c61980913..27bd74ce41dc 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -10,6 +10,7 @@ * Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch * Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul * Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math + * Strided K/V sparse-PA cache insert -> contiguous-copy insert The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T`` @@ -39,6 +40,7 @@ _mxfp8_e4m3_quantize_triton, dequant_mxfp8_to_bf16, ) +from vllm.models.minimax_m3.amd.model import _kv_insert_operand # noqa: E402 from vllm.models.minimax_m3.amd.ops import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, @@ -608,3 +610,77 @@ def resolve(consumes_mask, *, has_moe_kernel=True): assert torch.equal(resolve(True), mask) # AITER kernel -> 0/1 mask assert torch.equal(resolve(False), canonical) # non-AITER -> canonical map assert torch.equal(resolve(False, has_moe_kernel=False), canonical) # non-modular + + +# --------------------------------------------------------------------------- # +# Strided K/V operand for the AITER sparse-PA cache insert +# --------------------------------------------------------------------------- # +def _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim, num_q_heads=16): + """K/V as the sparse-PA insert sees them: column slices of a fused qkv.""" + kv = num_kv_heads * head_dim + row = num_q_heads * head_dim + 2 * kv + 2 * head_dim + qkv = torch.randn((num_tokens, row), dtype=torch.bfloat16, device=DEVICE) + k_start = num_q_heads * head_dim + k = qkv[:, k_start : k_start + kv].view(num_tokens, num_kv_heads, head_dim) + v = qkv[:, k_start + kv : k_start + 2 * kv].view(num_tokens, num_kv_heads, head_dim) + return k, v + + +def _asm_kv_cache(pages, num_kv_heads, head_dim): + """Page-16 K/V cache views matching ``_ensure_aiter_sparse_pa_kv_cache``.""" + x = 16 // torch.tensor([], dtype=torch.bfloat16).element_size() + key = torch.zeros( + (pages, num_kv_heads, head_dim // x, 16, x), + dtype=torch.bfloat16, + device=DEVICE, + ) + value = torch.zeros( + (pages, num_kv_heads, 16 // x, head_dim, x), + dtype=torch.bfloat16, + device=DEVICE, + ) + return key, value + + +@pytest.mark.parametrize("num_tokens", [1, 8, 129, 1024]) +@pytest.mark.parametrize("num_kv_heads", [1, 2]) +def test_kv_insert_operand_matches_contiguous(num_tokens, num_kv_heads): + """A strided K/V slice must insert bit-identically to a contiguous copy.""" + reshape_and_cache = pytest.importorskip("aiter").reshape_and_cache + head_dim = 128 + pages = num_tokens // 16 + 4 + slots = torch.randperm(pages * 16, device=DEVICE)[:num_tokens].to(torch.int32) + + torch.manual_seed(num_tokens) + k, v = _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim) + assert num_tokens == 1 or not k.is_contiguous() + + caches = [] + for operand in (lambda t: t.contiguous(), _kv_insert_operand): + key_cache, value_cache = _asm_kv_cache(pages, num_kv_heads, head_dim) + reshape_and_cache( + operand(k), + operand(v), + key_cache, + value_cache, + slots, + kv_cache_dtype="auto", + asm_layout=True, + ) + caches.append((key_cache, value_cache)) + + (ref_k, ref_v), (got_k, got_v) = caches + assert torch.equal(ref_k, got_k) + assert torch.equal(ref_v, got_v) + + +def test_kv_insert_operand_copies_unsupported_layouts(): + """Only a row-strided slice passes through; anything else is copied.""" + k, _ = _fused_qkv_kv_slices(32, 2, 128) + assert _kv_insert_operand(k) is k + + inner_strided = k[:, :, ::2] + assert _kv_insert_operand(inner_strided).is_contiguous() + + flat = k.reshape(32, -1) + assert _kv_insert_operand(flat).is_contiguous() diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 50cdb5518381..a0a58253d1a8 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -579,6 +579,32 @@ def forward( return output +# Widest row pitch aiter's reshape_and_cache can carry in its `int` stride args. +_MAX_KV_INSERT_ROW_STRIDE = 2**31 - 1 + + +def _kv_insert_operand(t: torch.Tensor) -> torch.Tensor: + """K or V operand for ``aiter.reshape_and_cache``, without a needless copy. + + The AITER sparse-PA insert takes K and V as column slices of the fused + ``[q | k | v | index_q | index_k]`` projection, so they are row-strided + views and never contiguous. ``reshape_and_cache`` does not need them to be: + it passes ``key.stride(0)`` to the kernel as the row pitch and reads + ``key[token_idx * key_stride + i]`` for ``i`` over ``num_heads * + head_size``, so the only requirement is that the trailing + ``(num_heads, head_size)`` block be contiguous within a row, which such a + slice always satisfies. Any other layout falls back to a copy. + """ + if ( + t.dim() == 3 + and t.stride(2) == 1 + and t.stride(1) == t.shape[2] + and t.stride(0) <= _MAX_KV_INSERT_ROW_STRIDE + ): + return t + return t.contiguous() + + class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): """Block-sparse attention layer with the lightning-indexer branch. @@ -902,8 +928,8 @@ def _insert_aiter_sparse_pa_kv( else "auto" ) reshape_and_cache( - k.contiguous(), - v.contiguous(), + _kv_insert_operand(k), + _kv_insert_operand(v), key_cache, value_cache, slot_mapping, From 3affd35f3164ba17503ba63cd8733becd1b03bf1 Mon Sep 17 00:00:00 2001 From: vllmellm Date: Thu, 17 Sep 2026 22:52:12 +0800 Subject: [PATCH 412/774] [Fix][ROCm] MXFP4 MoE round-up inflates TP-sharded expert weights on CDNA3, starving KV cache (#56359) Signed-off-by: vllmellm Co-authored-by: TJian --- .../layers/fused_moe/oracle/mxfp4.py | 37 +++++++++++++------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index ce19ba39bfc9..d1c151aca247 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -769,18 +769,31 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 128) elif current_platform.is_rocm(): - if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and ( - activation == MoEActivation.SITU or activation == MoEActivation.SILU - ): - # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size - # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a - # no-op for K3's shapes) rather than the generic ROCm 256 round-up, - # which would inflate weights and OOM. - intermediate_size = round_up(intermediate_size, 128) - hidden_size = round_up(hidden_size, 128) - else: - intermediate_size = round_up(intermediate_size, 256) - hidden_size = round_up(hidden_size, 256) + from vllm.platforms.rocm import get_cdna_version + + is_situ_or_silu = activation in ( + MoEActivation.SITU, + MoEActivation.SILU, + ) + + # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size + # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 + # rather than the generic ROCm 256 round-up, which would inflate + # weights and OOM. + aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 + + # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs. + # CDNA4's F16xMXFP4 configuration uses block_k=256. + triton_uses_128 = ( + backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4 + ) + + alignment = ( + 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256 + ) + + intermediate_size = round_up(intermediate_size, alignment) + hidden_size = round_up(hidden_size, alignment) elif backend == Mxfp4MoeBackend.CPU: # CPU AMX kernel uses BLOCK_N=32, align to 32 intermediate_size = round_up(intermediate_size, 32) From 4f9a41ac94beb940249d9712a6dd2996cf37da61 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Thu, 17 Sep 2026 11:05:31 -0400 Subject: [PATCH 413/774] [CI/Build][Hardware][NVIDIA] Add Rubin CUDA 13.4 nightly images (#55953) Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Misha Goin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/release-pipeline.yaml | 92 +++++++++++++++++++ .buildkite/scripts/cleanup-nightly-builds.sh | 4 +- docker/Dockerfile | 5 + tests/v1/engine/test_output_processor.py | 3 +- .../kv_connector/unit/test_nixl_connector.py | 3 +- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 3 +- 6 files changed, 105 insertions(+), 5 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 4af4eaf56ebc..2b90fef8201c 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -316,6 +316,66 @@ steps: - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" release-images' + - label: "Build release image - x86_64 - Rubin CUDA 13.4" + depends_on: ~ + id: build-release-image-x86-rubin-cuda-13-4 + if: build.env("NIGHTLY") == "1" + agents: + queue: cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker buildx build --load \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \ + --platform linux/amd64 \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=13.4 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg INSTALL_RUBIN_PRERELEASE=true \ + --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \ + --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ + --build-arg RUN_WHEEL_CHECK=false \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \ + --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images' + + - label: "Build release image - aarch64 - Rubin CUDA 13.4" + depends_on: ~ + id: build-release-image-arm64-rubin-cuda-13-4 + if: build.env("NIGHTLY") == "1" + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker buildx build --load \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \ + --platform linux/arm64 \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=13.4 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg INSTALL_RUBIN_PRERELEASE=true \ + --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \ + --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ + --build-arg RUN_WHEEL_CHECK=false \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \ + --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images' + - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04" depends_on: ~ id: build-release-image-x86-ubuntu2404 @@ -513,6 +573,20 @@ steps: - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" release-manifests' + - label: "Create multi-arch manifest - Rubin CUDA 13.4" + depends_on: + - build-release-image-x86-rubin-cuda-13-4 + - build-release-image-arm64-rubin-cuda-13-4 + id: create-multi-arch-manifest-rubin-cuda-13-4 + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu134 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu134 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" release-manifests' + - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" depends_on: - build-release-image-x86-ubuntu2404 @@ -585,6 +659,24 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" + - label: "Publish nightly multi-arch Rubin image to DockerHub - CUDA 13.4" + depends_on: + - create-multi-arch-manifest-rubin-cuda-13-4 + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/scripts/push-nightly-builds.sh cu134" + # Clean up old nightly builds (keep only last 14) + - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu134-nightly-" + plugins: + - docker-login#v3.0.0: + username: vllmbot + password-env: DOCKERHUB_TOKEN + env: + DOCKER_BUILDKIT: "1" + DOCKERHUB_USERNAME: "vllmbot" + - label: "Publish nightly CPU multi-arch image to DockerHub" depends_on: - build-cpu-release-image-x86 diff --git a/.buildkite/scripts/cleanup-nightly-builds.sh b/.buildkite/scripts/cleanup-nightly-builds.sh index b324778b1f79..41b04996e455 100755 --- a/.buildkite/scripts/cleanup-nightly-builds.sh +++ b/.buildkite/scripts/cleanup-nightly-builds.sh @@ -56,8 +56,8 @@ get_all_tags() { set -x # Get both last_updated timestamp and tag name, separated by | - # Exclude nightly-dev tags from cleanup - local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"') + # Exclude mutable architecture tags and nightly-dev tags from cleanup + local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | select(.name != ($prefix + "x86_64") and .name != ($prefix + "aarch64")) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"') if [ -z "$tags" ]; then break diff --git a/docker/Dockerfile b/docker/Dockerfile index d1a7afcbfb70..ad010285b181 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -231,10 +231,12 @@ RUN --mount=type=cache,target=/opt/uv/cache \ echo "Installing Rubin prerelease requirements..." \ && /opt/venv/bin/python3 use_existing_torch.py --prefix \ && uv pip install --python /opt/venv/bin/python3 \ + --prerelease=allow --index-strategy unsafe-first-match \ -r /tmp/vllm-requirements/rubin-prerelease.txt \ && uv pip install --python /opt/venv/bin/python3 \ -r requirements/cuda.txt \ && uv pip install --python /opt/venv/bin/python3 \ + --prerelease=allow --index-strategy unsafe-first-match \ -r /tmp/vllm-requirements/rubin-prerelease.txt; \ elif [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing torch nightly..." \ @@ -941,11 +943,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ uv pip install --system \ + --prerelease=allow --index-strategy unsafe-first-match \ -r /tmp/vllm-requirements/rubin-prerelease.txt \ && sed -Ei '/^[[:space:]]*(torch|torchvision|torchaudio)[[:space:]]*==/Id' \ /tmp/requirements-cuda.txt \ && uv pip install --system -r /tmp/requirements-cuda.txt \ && uv pip install --system \ + --prerelease=allow --index-strategy unsafe-first-match \ -r /tmp/vllm-requirements/rubin-prerelease.txt \ && printf '%s\n' \ "/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \ @@ -1262,6 +1266,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ echo "Installing Rubin prerelease requirements..." \ && uv pip install --system \ + --prerelease=allow --index-strategy unsafe-first-match \ -r /tmp/vllm-requirements/rubin-prerelease.txt; \ fi diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 9df43088e6c8..69128f55b704 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -177,7 +177,8 @@ def test_incremental_detokenization( def test_remote_prefill_cached_tokens_override(do_remote_prefill: bool): """P/D disaggregation: num_cached_tokens should report the P worker's cache hits (passed via kv_transfer_params) instead of the local count, - which sees the KVs pulled from the remote prefill as a ~100% hit.""" + which sees the KVs pulled from the remote prefill as a ~100% hit. + """ output_processor = OutputProcessor(tokenizer=None, log_stats=False) prompt_tokens = [1, 2, 3, 4, 5, 6, 7, 8] diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 2185ba62d27a..9ab4b0ae95f6 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -350,7 +350,8 @@ def test_prefill_exports_cached_tokens_in_kv_transfer_params(): """The P worker reports its own prefix-cache hits in the returned kv_transfer_params so the D worker can surface them in prompt_tokens_details instead of the ~100% local hit it measures - when pulling the KVs from the remote.""" + when pulling the KVs from the remote. + """ vllm_config = create_vllm_config() scheduler = create_scheduler(vllm_config) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 816ccdd498d9..401048e02c55 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -482,7 +482,8 @@ def _select_sm120_config( """(block_n, target_splits, num_warps) retuned on sm_120 (RTX PRO 6000 Blackwell), split by cache dtype: bf16 and fp8 favour different configs on sm_120, most visibly on the large-prefill region. Each entry is the fastest - config that stays correct on its own path; the main win is more warps.""" + config that stays correct on its own path; the main win is more warps. + """ if is_fp8: if base_programs > 2048: return (32, 2, 1) if use_prefill_config else (64, 1, 2) From 1df336c3bab8e2ac2eb3d325bd90e62b94656e9c Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Thu, 17 Sep 2026 18:05:48 +0300 Subject: [PATCH 414/774] [Perf][ROCm] Enable AITER QuickReduce + RMSNorm fusion (#48249) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_quick_all_reduce.py | 172 +++++++++++++++++++++ vllm/_aiter_ops.py | 107 ++++++++++++- 2 files changed, 278 insertions(+), 1 deletion(-) diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index f9228fbd6832..107d95fa613a 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -3,12 +3,17 @@ import multiprocessing import random +import sys +import types +from types import SimpleNamespace import pytest import ray import torch import torch.distributed as dist +import vllm._aiter_ops as aiter_ops +import vllm.distributed as vllm_distributed import vllm.envs as envs from vllm import _custom_ops as ops from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa @@ -233,6 +238,173 @@ def fake_qr_all_reduce( assert called_quant_level == QuickReduceRegime.FP.value +def test_rocm_aiter_fused_rmsnorm_uses_aiter_qr_rmsnorm_for_prefill( + monkeypatch: pytest.MonkeyPatch, +): + class FakeAiterAllReduce: + aiter_ca = SimpleNamespace(world_size=2, fully_connected=True) + + def use_1stage_fused_ar_rms(self, inp): + return False + + def custom_fused_ar_rms(self, *args, **kwargs): + raise AssertionError("expected fused QR+RMSNorm dispatch") + + class FakeQuickReduce: + disabled = False + + def __init__(self): + self.checked_input = None + + def should_quick_allreduce(self, inp): + self.checked_input = inp + return True + + def _get_qr_quant_level(self, inp): + return QuickReduceRegime.INT4.value + + qr_comm = FakeQuickReduce() + device_comm = SimpleNamespace(cpu_group=object(), qr_comm=qr_comm) + monkeypatch.setattr( + vllm_distributed, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_comm), + ) + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr(aiter_ops.dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(aiter_ops.dist, "get_rank", lambda group: 0) + monkeypatch.setattr( + aiter_ops.dist, + "all_gather_object", + lambda handles, handle, group: handles.__setitem__(0, handle), + ) + + calls = [] + fake_aiter = types.SimpleNamespace( + init_custom_qr=lambda rank, world_size, qr_max_size: 123, + qr_get_handle=lambda ptr: b"handle", + qr_open_handles=lambda ptr, handles: None, + ) + + def fake_qr_all_reduce_rmsnorm( + ptr, + inp, + residual, + residual_out, + out, + weight, + eps, + hidden_dim, + quant_level, + cast_bf2half, + ): + calls.append( + { + "ptr": ptr, + "inp": inp, + "residual": residual, + "residual_out": residual_out, + "out": out, + "weight": weight, + "eps": eps, + "hidden_dim": hidden_dim, + "quant_level": quant_level, + "cast_bf2half": cast_bf2half, + } + ) + out.copy_(inp) + residual_out.copy_(residual) + + fake_aiter.qr_all_reduce_rmsnorm = fake_qr_all_reduce_rmsnorm + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None) + monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", True) + + # token_num > 80 forces vLLM's AITER fused AR+RMS wrapper onto its + # non-1stage prefill path. hidden_dim=4096 bf16 gives an 8 KiB row, so four + # independent RMSNorm rows fit in one 32 KiB QR tile. + inp = torch.randn(81, 4096, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(4096, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6 + ) + + assert qr_comm.checked_input is inp + assert len(calls) == 1 + assert calls[0]["ptr"] == 123 + assert calls[0]["inp"] is inp + assert calls[0]["residual"] is residual + assert calls[0]["out"] is out + assert calls[0]["residual_out"] is residual_out + assert calls[0]["weight"] is weight + assert calls[0]["hidden_dim"] == 4096 + assert calls[0]["quant_level"] == QuickReduceRegime.INT4.value + assert calls[0]["cast_bf2half"] is True + torch.testing.assert_close(out, inp) + torch.testing.assert_close(residual_out, residual) + + +def test_rocm_aiter_fused_rmsnorm_keeps_1stage_decode_path( + monkeypatch: pytest.MonkeyPatch, +): + calls = [] + + class FakeAiterCA: + world_size = 2 + fully_connected = True + + def custom_fused_ar_rms( + self, inp, residual, weight, epsilon, *, use_1stage, gemma_norm + ): + calls.append( + { + "inp": inp, + "residual": residual, + "weight": weight, + "eps": epsilon, + "use_1stage": use_1stage, + "gemma_norm": gemma_norm, + } + ) + return inp + 1, residual + 1 + + class FakeAiterAllReduce: + aiter_ca = FakeAiterCA() + + def use_1stage_fused_ar_rms(self, inp): + return True + + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + inp = torch.randn(4, 4096, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(4096, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6 + ) + + assert len(calls) == 1 + assert calls[0]["inp"] is inp + assert calls[0]["residual"] is residual + assert calls[0]["weight"] is weight + assert calls[0]["use_1stage"] is True + assert calls[0]["gemma_norm"] is False + torch.testing.assert_close(out, inp + 1) + torch.testing.assert_close(residual_out, residual + 1) + + @ray.remote(num_gpus=1, max_calls=1) def graph_quickreduce( monkeypatch: pytest.MonkeyPatch, diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 101e1dd10a6d..4b9a145f4830 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -4,9 +4,10 @@ import functools import os from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch +import torch.distributed as dist from torch._ops import OpOverload import vllm.envs as envs @@ -34,6 +35,63 @@ # which is a host op, so we cache it once here. FP8_DTYPE = current_platform.fp8_dtype() _HIPB_MM_INITIALIZED_DEVICES: set[int] = set() +KB = 1024 +MB = 1024 * KB + + +def _get_or_create_aiter_qr_rmsnorm_comm( + device_comm: Any, +) -> tuple[int, int, bool] | None: + """Create the AITER QR communicator used by fused QR+RMSNorm. + + vLLM's normal QuickReduce communicator owns vLLM C++ QR buffers. The fused + QR+RMSNorm kernel added in AITER needs an AITER DeviceComms instance, so we + lazily create one on the same TP CPU group and keep it on the vLLM device + communicator for reuse. + """ + existing = getattr(device_comm, "_aiter_qr_rmsnorm_comm", None) + if existing is False: + return None + if existing is not None: + return existing + + try: + import aiter + except Exception: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + if not hasattr(aiter, "qr_all_reduce_rmsnorm"): + device_comm._aiter_qr_rmsnorm_comm = False + return None + + group = getattr(device_comm, "cpu_group", None) + if group is None: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + world_size = dist.get_world_size(group=group) + rank = dist.get_rank(group=group) + if world_size not in (2, 4, 8): + device_comm._aiter_qr_rmsnorm_comm = False + return None + + qr_max_size_mb = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB + qr_max_size = None if qr_max_size_mb is None else qr_max_size_mb * MB + + try: + ptr = aiter.init_custom_qr(rank, world_size, qr_max_size) + handle = aiter.qr_get_handle(ptr) + handles = [None] * world_size + dist.all_gather_object(handles, handle, group=group) + aiter.qr_open_handles(ptr, handles) + except Exception: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + comm = (ptr, world_size, envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16) + device_comm._aiter_qr_rmsnorm_comm = comm + return comm def _ensure_hipb_mm_extension_initialized() -> None: @@ -991,6 +1049,53 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( ca = aiter_ar.aiter_ca use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_) + if not use_1stage: + from vllm.distributed import get_tp_group + + device_comm = get_tp_group().device_communicator + qr_comm = getattr(device_comm, "qr_comm", None) + hidden_dim = input_.shape[-1] + row_size = hidden_dim * input_.element_size() + fused_qr_rmsnorm_ok = ( + qr_comm is not None + and not getattr(qr_comm, "disabled", True) + and hasattr(qr_comm, "should_quick_allreduce") + and qr_comm.should_quick_allreduce(input_) + and input_.shape == residual.shape + and input_.dtype in (torch.bfloat16, torch.float16) + and input_.dtype == residual.dtype + and input_.dtype == weight.dtype + and weight.dim() == 1 + and weight.numel() == hidden_dim + and input_.numel() % hidden_dim == 0 + and row_size > 0 + and row_size <= 32 * KB + and (32 * KB) % row_size == 0 + ) + if fused_qr_rmsnorm_ok: + assert qr_comm is not None + aiter_qr_comm = _get_or_create_aiter_qr_rmsnorm_comm(device_comm) + if aiter_qr_comm is not None: + import aiter + + ptr, _, cast_bf2half = aiter_qr_comm + quant_level = qr_comm._get_qr_quant_level(input_) + out = torch.empty_like(input_) + residual_out = torch.empty_like(residual) + aiter.qr_all_reduce_rmsnorm( + ptr, + input_, + residual, + residual_out, + out, + weight, + epsilon, + hidden_dim, + quant_level, + cast_bf2half, + ) + return out, residual_out + result = ca.custom_fused_ar_rms( input_, residual, From 493d4b3c04a439c926df2c4d2fa3cdfabb1a8605 Mon Sep 17 00:00:00 2001 From: HieDean <34408026+HieDean@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:20:37 +0800 Subject: [PATCH 415/774] [bugfix] Mark draft tokens to rebuilt their embeddings. (#57356) Signed-off-by: HieDean <799287043@qq.com> Co-authored-by: hiedean --- vllm/v1/worker/gpu_model_runner.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 10f85be718a6..77d11621c2cb 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1858,8 +1858,12 @@ def _prepare_input_ids( total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens if self.enable_prompt_embeds: # The multimodal embed path reads is_token_ids.gpu; its .cpu copy is - # refreshed every step but the async fast paths below only scatter - # input_ids.gpu, so refresh is_token_ids.gpu here too. + # refreshed every step but the async fast path below replaces draft + # tokens directly on the GPU. Mark those positions as token IDs before + # refreshing the GPU mask so their embeddings are rebuilt. Sampled + # tokens are already marked by bookkeeping. + if spec_flattened_indices: + self.is_token_ids.np[spec_flattened_indices] = True self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens) if num_common_tokens < total_without_spec: # If not all requests are decodes from the last iteration, From 6a7de3ba3ad95965e41618568ff803fa01ab168d Mon Sep 17 00:00:00 2001 From: Zoe923 <52874215+Zoe923@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:49:56 +0800 Subject: [PATCH 416/774] [Bugfix][GDN] Fix CuteDSL BF16 KKT inversion divergence (#53864) Signed-off-by: Zoe923 <871844693@qq.com> --- .../kernels/mamba/test_gdn_prefill_cutedsl.py | 83 +++++++++++++++++++ .../gdn_chunk_cutedsl/kernel_kkt_inv_uw.py | 4 +- 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py index 861b686e32da..d86cfde2c7c7 100644 --- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py +++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py @@ -192,3 +192,86 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): assert no_buffer_state_error.mean().item() < 6e-4 assert buffer_o_error.max().item() == 0 assert buffer_state_error.max().item() == 0 + + +def test_gdn_chunk_cutedsl_correlated_keys_stability(): + """Cover a KKT block where three BF16 Newton rounds diverge.""" + total_tokens = 64 + num_k_heads = 4 + num_v_heads = 8 + head_dim = 128 + rng = torch.Generator("cuda").manual_seed(3456) + + q = torch.zeros( + 1, + total_tokens, + num_k_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.zeros_like(q) + q[..., 0] = 1 + k[..., 0] = 1 + v = torch.randn( + 1, + total_tokens, + num_v_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + generator=rng, + ) + g = torch.zeros( + 1, + total_tokens, + num_v_heads, + device="cuda", + dtype=torch.float32, + ) + beta = torch.full_like(g, 0.9) + initial_state = torch.zeros( + 1, + num_v_heads, + head_dim, + head_dim, + device="cuda", + dtype=torch.float32, + ) + cu_seqlens = torch.tensor([0, total_tokens], device="cuda", dtype=torch.int32) + chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens) + + ref_o, ref_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=False, + ) + actual_o, actual_state = chunk_gated_delta_rule_cutedsl( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + ) + torch.accelerator.synchronize() + + assert torch.isfinite(actual_o).all() + assert torch.isfinite(actual_state).all() + output_error = (actual_o.float() - ref_o.float()).abs() + state_error = ( + actual_state.float() - ref_state.to(actual_state.dtype).float() + ).abs() + assert output_error.max().item() < 1e-1 + assert output_error.mean().item() < 2e-3 + assert state_error.max().item() < 2e-1 + assert state_error.mean().item() < 6e-4 diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py index ba61ec68b5cc..4b4046a09327 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py @@ -573,8 +573,8 @@ def store_ab_abg( for i in cutlass.range_constexpr(4): M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i]) - # 3 rounds of Newton-Schulz - for _ in cutlass.range_constexpr(3): + # The fourth round keeps correlated-key KKT blocks stable in BF16. + for _ in cutlass.range_constexpr(4): # First MMA: -AiM = Ai @ (-M) cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)]) cute.arch.sync_warp() From d7e755c23080988f7450f107466d6f6a74b042ff Mon Sep 17 00:00:00 2001 From: Ziyang Ma Date: Thu, 17 Sep 2026 23:50:14 +0800 Subject: [PATCH 417/774] [Bugfix] Fix wrong vLLM version reported by pip install (proto-v* tag collision) (#57295) Signed-off-by: RyanMa29 Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Codex --- .buildkite/scripts/ci-bake-rocm.sh | 2 +- setup.py | 22 ++++++++++++++++++++-- tools/build_rust.py | 17 ++++++++++++++++- 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 6c5f3308b785..0a9e6e7589fb 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -418,7 +418,7 @@ validate_ci_build_context_source() { describe_ci_revision() { git -C "$1" describe --tags --long --abbrev=10 \ - --match '*[0-9]*' "$2" 2>/dev/null + --match 'v[0-9]*' "$2" 2>/dev/null } write_ci_git_archival_metadata() { diff --git a/setup.py b/setup.py index b24bcc1a2cc5..ed9616195ff0 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,18 @@ from setuptools_scm import get_version from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME +# Select vLLM release tags, excluding crate tags such as "proto-v0.3.0". +VLLM_GIT_DESCRIBE_COMMAND = [ + "git", + "describe", + "--dirty", + "--tags", + "--long", + "--abbrev=40", + "--match", + "v[0-9]*", +] + def load_module_from_path(module_name, path): spec = importlib.util.spec_from_file_location(module_name, path) @@ -1249,9 +1261,15 @@ def get_vllm_version() -> str: if env_version := os.getenv("VLLM_VERSION_OVERRIDE"): print(f"Overriding VLLM version with {env_version} from VLLM_VERSION_OVERRIDE") os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version - return get_version(write_to="vllm/_version.py") + return get_version( + write_to="vllm/_version.py", + git_describe_command=VLLM_GIT_DESCRIBE_COMMAND, + ) - version = get_version(write_to="vllm/_version.py") + version = get_version( + write_to="vllm/_version.py", + git_describe_command=VLLM_GIT_DESCRIBE_COMMAND, + ) sep = "+" if "+" not in version else "." # dev versions might contain + if not envs.VLLM_SKIP_VERSION_SUFFIX: diff --git a/tools/build_rust.py b/tools/build_rust.py index a000ec8169fd..320cf75e3db4 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -16,13 +16,28 @@ ROOT_DIR = Path(__file__).resolve().parents[1] VLLM_RS_BUILD_VERSION = "VLLM_RS_BUILD_VERSION" +# Select vLLM release tags, excluding crate tags such as "proto-v0.3.0". +VLLM_GIT_DESCRIBE_COMMAND = [ + "git", + "describe", + "--dirty", + "--tags", + "--long", + "--abbrev=40", + "--match", + "v[0-9]*", +] + def prepare_build_environment() -> str | None: """Set the device-independent vLLM source version for Rust artifacts.""" version = os.getenv(VLLM_RS_BUILD_VERSION) or None if version is None: try: - version = get_version(root=ROOT_DIR) + version = get_version( + root=ROOT_DIR, + git_describe_command=VLLM_GIT_DESCRIBE_COMMAND, + ) except LookupError: return None From b78fbab4be552650286778a9e489980d850f76c3 Mon Sep 17 00:00:00 2001 From: Young <33064860+YCH188@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:52:10 +0800 Subject: [PATCH 418/774] [Docs] Fix the typos in the document (#57190) Co-authored-by: Young --- docs/benchmarking/cli.md | 4 ++-- docs/benchmarking/dashboard.md | 2 +- docs/contributing/README.md | 2 +- docs/design/metrics.md | 2 +- docs/features/quantization/quark.md | 2 +- docs/features/speculative_decoding/draft_model.md | 2 +- docs/getting_started/installation/cpu.x86.inc.md | 2 +- docs/governance/committers.md | 2 +- docs/training/weight_transfer/base.md | 2 +- 9 files changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index d46ab04b0e2a..e7ee794da6f1 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -147,7 +147,7 @@ With standard decoding, each streamed output usually contains one token, so ITL and TPOT are typically similar. With speculative decoding, one streamed output can contain multiple tokens, -such as several accepted draft tokens within a single engine tstep. ITL records +such as several accepted draft tokens within a single engine step. ITL records only the gaps between streamed outputs; it does not add zero-duration gaps for tokens in the same output. TPOT instead amortizes the request's decoding time over every output token. @@ -451,7 +451,7 @@ vllm bench serve \ --num-prompts -1 ``` -Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card. +Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructured data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card. #### BFCL (Tool-Calling) Benchmark diff --git a/docs/benchmarking/dashboard.md b/docs/benchmarking/dashboard.md index 44effc078e35..4c1132c17fd1 100644 --- a/docs/benchmarking/dashboard.md +++ b/docs/benchmarking/dashboard.md @@ -36,7 +36,7 @@ When run, benchmark script generates results under **benchmark/results** folder, - `ON_CPU`: set the value to '1' on Intel® Xeon® and Arm® Neoverse™ Processors. Default value is 0. - `SERVING_JSON`: JSON file to use for the serving tests. Default value is empty string (use default file). - `LATENCY_JSON`: JSON file to use for the latency tests. Default value is empty string (use default file). -- `THROUGHPUT_JSON`: JSON file to use for the throughout tests. Default value is empty string (use default file). +- `THROUGHPUT_JSON`: JSON file to use for the throughput tests. Default value is empty string (use default file). - `REMOTE_HOST`: IP for the remote vLLM service to benchmark. Default value is empty string. - `REMOTE_PORT`: Port for the remote vLLM service to benchmark. Default value is empty string. - `PROMPTS_PER_CONCURRENCY`: Multiplier to compute `num_prompts` for serving tests (`num_prompts = max_concurrency × value`). Overrides JSON `num_prompts`. Default is NULL. diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 7a1af9c1a7b2..e4b86f99aed5 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -50,7 +50,7 @@ To rebuild only the Rust frontend binary: ./build_rust.sh --debug # faster build for development ``` -If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first: +If you are developing vLLM's Python and CUDA/C++ code, install PyTorch first: ```bash uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 3d108890d702..3e070f588890 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -670,7 +670,7 @@ fall under the more general heading of "Observability". vLLM has support for OpenTelemetry tracing: - Added by and reinstated by -- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces` +- Configured with `--otlp-traces-endpoint` and `--collect-detailed-traces` - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/) - [User-facing docs](../../examples/observability/opentelemetry/README.md) - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f) diff --git a/docs/features/quantization/quark.md b/docs/features/quantization/quark.md index 8445288ff97f..e5cfb8653bd6 100644 --- a/docs/features/quantization/quark.md +++ b/docs/features/quantization/quark.md @@ -100,7 +100,7 @@ kv-cache and the quantization algorithm is AutoSmoothQuant. !!! note Note the quantization algorithm needs a JSON config file and the config file is located in - [Quark Pytorch examples](https://quark.docs.amd.com/latest/pytorch/pytorch_examples.html), + [Quark PyTorch examples](https://quark.docs.amd.com/latest/pytorch/pytorch_examples.html), under the directory `examples/torch/language_modeling/llm_ptq/models`. For example, AutoSmoothQuant config file for Llama is `examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json`. diff --git a/docs/features/speculative_decoding/draft_model.md b/docs/features/speculative_decoding/draft_model.md index 5ab4629409f4..880ea1df5989 100644 --- a/docs/features/speculative_decoding/draft_model.md +++ b/docs/features/speculative_decoding/draft_model.md @@ -96,7 +96,7 @@ The code used to request completions as a client remains unchanged: }, gpu_memory_utilization=0.5, ) -outputs = llm.generate(prompts，sampling_params) +outputs = llm.generate(prompts, sampling_params) for output in outputs: prompt = output.prompt diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index a64d0e570333..260ffe286362 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -33,7 +33,7 @@ uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_V pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl" --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" - Before use vLLM CPU installed via wheels, make Intel OpenMP is added to `LD_PRELOAD`: + Before use vLLM CPU installed via wheels, make sure Intel OpenMP is added to `LD_PRELOAD`: ```bash # manually find the path sudo find / -iname *libiomp5.so diff --git a/docs/governance/committers.md b/docs/governance/committers.md index bd3675912217..48ed4b26073f 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -2,7 +2,7 @@ This document lists the current committers of the vLLM project and the core areas they maintain. Committers have write access to the vLLM repository and are responsible for reviewing and merging PRs. -You can also refer to the [CODEOWNERS](https://github.com/vllm-project/vllm/blob/main/.github/CODEOWNERS) file for concrete file-level ownership and reviewers. Both this documents and the CODEOWNERS file are living documents and they complement each other. +You can also refer to the [CODEOWNERS](https://github.com/vllm-project/vllm/blob/main/.github/CODEOWNERS) file for concrete file-level ownership and reviewers. Both these documents and the CODEOWNERS file are living documents and they complement each other. ## Active Committers diff --git a/docs/training/weight_transfer/base.md b/docs/training/weight_transfer/base.md index d73c88436ae0..759c55ef6fe9 100644 --- a/docs/training/weight_transfer/base.md +++ b/docs/training/weight_transfer/base.md @@ -233,7 +233,7 @@ Two implementations ship with vLLM: | `RayVLLMWeightSyncClient(handle)` | One or more `AsyncLLM`/`LLM` Ray actors. Accepts a list and fans each call out to every handle, blocking on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit | | `HTTPVLLMWeightSyncClient(base_url, timeout=300)` | A vLLM server over the RLHF HTTP routes | -Custom weight sync clients can be implement like so: +Custom weight sync clients can be implemented like so: ```python class MyFrameworkWeightSyncClient: From e960ead0aa70362a57f205be831083b68c7a5bdb Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Thu, 17 Sep 2026 11:11:31 -0500 Subject: [PATCH 419/774] [Mypy] Fix mypy typing for Transformers models (#54320) Signed-off-by: Taneem Ibrahim Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/models/transformers/fusers/test_mla.py | 6 +-- tools/pre_commit/mypy.py | 1 - vllm/compilation/decorators.py | 13 +++-- .../models/transformers/__init__.py | 24 ++++----- .../models/transformers/base.py | 44 +++++++++------ .../models/transformers/causal.py | 4 +- .../models/transformers/fusers/attention.py | 12 +++-- .../models/transformers/fusers/base.py | 7 ++- .../models/transformers/fusers/mla.py | 36 ++++++------- .../models/transformers/fusers/rms_norm.py | 12 ++--- .../models/transformers/fx_utils.py | 25 +++++---- .../models/transformers/legacy.py | 31 ++++++----- .../model_executor/models/transformers/moe.py | 42 +++++++++++---- .../models/transformers/multimodal.py | 54 ++++++++++++------- .../models/transformers/pooling.py | 37 ++++++++----- .../models/transformers/utils.py | 21 ++++---- 16 files changed, 224 insertions(+), 145 deletions(-) diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py index 9f51ab56621b..625fc4dc0815 100644 --- a/tests/models/transformers/fusers/test_mla.py +++ b/tests/models/transformers/fusers/test_mla.py @@ -53,7 +53,6 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None: def test_discovers_modules_without_q_lora(): fuser = _match(q_lora_rank=None) assert isinstance(fuser, MLAFuser) - assert not fuser.has_q_lora assert fuser.q_proj_name == "q_proj" assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa" assert fuser.kv_a_layernorm_name == "kv_a_layernorm" @@ -68,7 +67,6 @@ def test_discovers_modules_without_q_lora(): def test_discovers_modules_with_q_lora(): fuser = _match(q_lora_rank=64) assert isinstance(fuser, MLAFuser) - assert fuser.has_q_lora assert fuser.q_a_proj_name == "q_a_proj" assert fuser.q_a_layernorm_name == "q_a_layernorm" assert fuser.q_b_proj_name == "q_b_proj" @@ -107,7 +105,7 @@ def test_update_forward_rewrites_real_attention(q_lora_rank): assert isinstance(fuser, MLAFuser) fuser.update_forward(_attention(q_lora_rank)) names = set(fuser.fused_forward.__code__.co_names) - if fuser.has_q_lora: + if fuser.q_a_proj_name is not None: assert _FUSED_QKV_A_PROJ in names assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names @@ -158,7 +156,7 @@ def test_discovers_modules_under_arbitrary_names(): module = RenamedMLA() fuser = MLAFuser.match(trace(module), module) assert isinstance(fuser, MLAFuser) - assert not fuser.has_q_lora + assert fuser.q_a_proj_name is None assert fuser.q_proj_name == "alpha" assert fuser.kv_a_proj_name == "beta" assert fuser.kv_a_layernorm_name == "gamma" diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index b7a278b9bd6b..234acc029858 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -116,7 +116,6 @@ r"vllm/model_executor/models/[nN]", r"vllm/model_executor/models/[oO]", r"vllm/model_executor/models/[qQ]", - r"vllm/model_executor/models/transformers", r"vllm/model_executor/models/[uU]", r"vllm/model_executor/models/[vV]", r"vllm/model_executor/models/[wW]", diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index 242e7227e922..35c0e6dee5ee 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -7,7 +7,7 @@ import os import sys from collections.abc import Callable, Generator -from typing import TYPE_CHECKING, Any, TypeVar, overload +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload from unittest.mock import patch import torch @@ -32,6 +32,9 @@ from .monitor import monitor_profiling_run, monitor_torch_compile +DynamicArgDims: TypeAlias = dict[str, int | list[int] | dict[int, str]] +"""Argument name -> the dynamic dimension(s) of that argument.""" + # shape_id parameter was added to mark_unbacked in PyTorch 2.11.0 _SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0") @@ -89,7 +92,7 @@ def support_torch_compile( @overload def support_torch_compile( *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None, + dynamic_arg_dims: DynamicArgDims | None, ) -> Callable[[type[_T]], type[_T]]: ... @@ -103,7 +106,7 @@ def support_torch_compile( @overload def support_torch_compile( *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None, + dynamic_arg_dims: DynamicArgDims | None = None, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, @@ -117,7 +120,7 @@ def support_torch_compile(cls: type[_T]) -> type[_T]: ... def support_torch_compile( cls: type[_T] | None = None, *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None, + dynamic_arg_dims: DynamicArgDims | None = None, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, @@ -328,7 +331,7 @@ def _try_load_aot_compiled_fn( def _support_torch_compile( cls: type[_T], - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]], + dynamic_arg_dims: DynamicArgDims, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 2d127c430e76..e19de6d4e891 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -21,7 +21,7 @@ import torch.nn.functional as F from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR, Base +from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin from vllm.model_executor.models.transformers.moe import MoEMixin @@ -126,10 +126,10 @@ def vllm_mla_attention_forward( # Text only models -class TransformersForCausalLM(CausalMixin, Base): ... +class TransformersForCausalLM(CausalMixin): ... -class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... +class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ... # Multimodal models @@ -138,7 +138,7 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... +class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin): ... @MULTIMODAL_REGISTRY.register_processor( @@ -146,16 +146,14 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalMoEForCausalLM( - MoEMixin, MultiModalMixin, CausalMixin, Base -): ... +class TransformersMultiModalMoEForCausalLM(MoEMixin, MultiModalMixin, CausalMixin): ... # Embedding models -class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ... +class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin): ... -class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... +class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin): ... @MULTIMODAL_REGISTRY.register_processor( @@ -163,17 +161,17 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ... +class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin): ... # Sequence classification models class TransformersForSequenceClassification( - SequenceClassificationMixin, LegacyMixin, Base + SequenceClassificationMixin, LegacyMixin ): ... class TransformersMoEForSequenceClassification( - SequenceClassificationMixin, MoEMixin, Base + SequenceClassificationMixin, MoEMixin ): ... @@ -183,7 +181,7 @@ class TransformersMoEForSequenceClassification( dummy_inputs=MultiModalDummyInputsBuilder, ) class TransformersMultiModalForSequenceClassification( - SequenceClassificationMixin, MultiModalMixin, Base + SequenceClassificationMixin, MultiModalMixin ): ... diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 62f6d684e8b6..ff8fc98dbc53 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -35,7 +35,7 @@ get_model_conversion_mapping, ) -from vllm.compilation.decorators import support_torch_compile +from vllm.compilation.decorators import DynamicArgDims, support_torch_compile from vllm.config.utils import getattr_iter from vllm.distributed import get_pp_group, get_tp_group from vllm.distributed.utils import get_pp_indices @@ -76,6 +76,7 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + ShardId, WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix, @@ -110,10 +111,13 @@ class Base( SupportsEagle, SupportsEagle3, ): - embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it + hf_to_vllm_mapper: WeightsMapper + + # TODO transformers will have a util to get these + embedding_modules = {"embed_tokens": "input_embeddings"} def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): - super().__init__() + nn.Module.__init__(self) logger.info("Using Transformers modeling backend.") self.vllm_config = vllm_config @@ -248,7 +252,7 @@ def _pre_trained_model_classes(self) -> PreTrainedModelClasses: def _decorate_cls_for_torch_compile( self, cls: type["PreTrainedModel"], - dynamic_arg_dims: dict[str, int] | None, + dynamic_arg_dims: DynamicArgDims | None, enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, ): @@ -284,7 +288,7 @@ def _decorate_for_torch_compile(self): self._decorate_cls_for_torch_compile( cls=self._pre_trained_model_classes.decoder, # Applied to a PreTrainedModel so the batch dimension will exist - dynamic_arg_dims=dict[str, int]( + dynamic_arg_dims=DynamicArgDims( input_ids=1, # shape: [1, seq_len] inputs_embeds=1, # shape: [1, seq_len, hidden_size] position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope @@ -303,9 +307,8 @@ def _create_hf_to_vllm_mapper(self): - Checkpoints saved with no base model prefix - Any quantization config specific mappings """ - self.hf_to_vllm_mapper = WeightsMapper() - orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming - orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex + orig_to_new_renaming: list[WeightRenaming] = [] + orig_to_new_regex: dict[re.Pattern, str | None] = {} for mapping in get_model_conversion_mapping(self.model): # Handle weights which have been renamed in Transformers @@ -341,6 +344,11 @@ def _create_hf_to_vllm_mapper(self): nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)") orig_to_new_regex[nested_lm_head_pattern] = r"\2" + self.hf_to_vllm_mapper = WeightsMapper( + orig_to_new_renaming=orig_to_new_renaming, + orig_to_new_regex=orig_to_new_regex, + ) + # Apply mapping to quantization config if needed self._maybe_apply_model_mapping() @@ -371,8 +379,8 @@ def pipeline_parallel(self): "modeling backend will infer the split from the layers of %s in order " "of declaration and keep parameter-free modules on every rank. This " "may fail if the model's structure is non-standard. %s", - type(self.model), - type(module), + type(self.model).__name__, + type(module).__name__, tip, ) @@ -470,7 +478,7 @@ def recursive_replace(self): "backend will shard the model the best it can during graph fusion and " "replicate the rest. This may be suboptimal or fail if the model does " "not fuse cleanly. %s", - type(self.model), + type(self.model).__name__, tip, ) @@ -481,6 +489,8 @@ def recursive_replace(self): vocab_embeddings = self._vocab_embeddings() + orig_to_new_stacked: dict[str, tuple[str, ShardId]] = {} + def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module): """Register a fused layer's mappings just before it is built.""" self.fusers.setdefault(prefix, []).append(fuser) @@ -498,8 +508,7 @@ def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module): ) self.attention_fusers[index] = (prefix, fuser) - orig_to_new_stacked = fuser.orig_to_new_stacked(prefix) - self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked) + orig_to_new_stacked.update(fuser.orig_to_new_stacked(prefix)) packed_modules_mapping = fuser.packed_modules_mapping self.packed_modules_mapping.update(packed_modules_mapping) @@ -564,6 +573,8 @@ def _recursive_replace(module: nn.Module, prefix: str): _recursive_replace(self.model, prefix="model") + self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked) + def _create_attention_instances(self): """Create `Attention` instances to inform KV cache allocation.""" text_config = self.text_config @@ -589,8 +600,11 @@ def _create_attention_instances(self): for i in range(start, end): if i not in self.attention_fusers: - in_range = layer_types and i < len(layer_types) - layer = f"{i} ({layer_types[i]})" if in_range else str(i) + layer = ( + f"{i} ({layer_types[i]})" + if layer_types and i < len(layer_types) + else str(i) + ) raise ValueError( f"Layer {layer} does not dispatch through the Transformers " "attention interface and vLLM has no other way to handle it." diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index c5231f9e979d..e7dba82b0f80 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -27,13 +27,15 @@ from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from .base import Base + if TYPE_CHECKING: import torch from vllm.config import VllmConfig -class CausalMixin(VllmModelForTextGeneration): +class CausalMixin(VllmModelForTextGeneration, Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO super(VllmModelForTextGeneration, self).__init__( diff --git a/vllm/model_executor/models/transformers/fusers/attention.py b/vllm/model_executor/models/transformers/fusers/attention.py index fc16527f15e3..d8d324536d27 100644 --- a/vllm/model_executor/models/transformers/fusers/attention.py +++ b/vllm/model_executor/models/transformers/fusers/attention.py @@ -5,6 +5,7 @@ import ast import inspect import textwrap +from collections.abc import Callable from dataclasses import dataclass from functools import cache from typing import TYPE_CHECKING, ClassVar @@ -36,16 +37,16 @@ def _is_interface_lookup(node: ast.expr | None) -> bool: @cache -def interface_call(cls: type[nn.Module]) -> ast.Call | None: - """The attention interface call in `cls.forward`, if it makes exactly one.""" +def interface_call(forward: Callable) -> ast.Call | None: + """The attention interface call in `forward`, if it makes exactly one.""" try: - source = inspect.getsource(inspect.unwrap(cls.forward)) + source = inspect.getsource(inspect.unwrap(forward)) tree = ast.parse(textwrap.dedent(source)) except (AttributeError, OSError, SyntaxError, TypeError): return None # Find the names of the local variables that read from the interface lookup - names = set() + names: set[str] = set() for node in ast.walk(tree): if not isinstance(node, (ast.Assign, ast.AnnAssign)): continue @@ -97,7 +98,7 @@ def info(self, name: str) -> str: def match( cls, graph: fx.Graph | None, module: nn.Module ) -> "AttentionFuser | None": - if (call := interface_call(type(module))) is None: + if (call := interface_call(type(module).forward)) is None: return None scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"] scale_expr = scaling[0] if len(scaling) == 1 else None @@ -120,6 +121,7 @@ def fuse( self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: if (sinks := self.sinks(module)) is not None: + assert isinstance(self.s_aux_expr, ast.Attribute) size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size device = vllm_config.device_config.device data = torch.empty(size, dtype=sinks.dtype, device=device) diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index ad24f2542a47..937b70c5ed45 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -130,11 +130,14 @@ class StackedFuser(RewriteFuser): """A fuser that merges sibling projections into one stacked linear and rewrites the forward to call it.""" - merged_name: ClassVar[str] - """Attribute name of the merged module created by `update_attrs`.""" merged_cls_name: ClassVar[str] """Name of the vLLM class the merged projection becomes (for logging).""" + @property + @abstractmethod + def merged_name(self) -> str: + """Attribute name of the merged module created by `update_attrs`.""" + def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index fa4b34750c5f..96a4a0b87b3d 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -24,6 +24,7 @@ upstream_linear, ) from vllm.model_executor.models.transformers.utils import ( + Style, log_replacement, replace_linear_class, ) @@ -128,30 +129,26 @@ class MLAFuser(StackedFuser): merged_name: ClassVar[str] = "fused_qkv_a_proj" merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" - @property - def has_q_lora(self) -> bool: - return self.q_a_proj_name is not None - def info(self, name: str) -> str: info_str = ( f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)" ) - if self.has_q_lora: + if self.q_a_proj_name is not None: info_str += "; " + super().info(name).removeprefix("Fused: ") return info_str @property def shards(self) -> list[tuple[str, ShardId]]: """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection.""" - if self.has_q_lora: - return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)] - return [] + if self.q_a_proj_name is None: + return [] + return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)] @property def packed_modules_mapping(self) -> dict[str, list[str]]: - if self.has_q_lora: - return super().packed_modules_mapping - return {} + if self.q_a_proj_name is None: + return {} + return super().packed_modules_mapping @classmethod def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None": @@ -234,10 +231,10 @@ def update_forward(self, module: nn.Module) -> None: Bypass the KV expansion method so the compressed latent reaches the `vllm_mla` attention interface unexpanded.""" funcdef, fn = recover_forward(type(module)) - if self.has_q_lora: + if (q_a_proj_name := self.q_a_proj_name) is not None: # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. # The fused call is inserted at the top-level statement preceding both. - names = [self.q_a_proj_name, self.kv_a_proj_name] + names = [q_a_proj_name, self.kv_a_proj_name] calls = self._unguarded_calls(funcdef, names) if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]): raise ValueError("down-projections read different inputs") @@ -258,7 +255,7 @@ def update_forward(self, module: nn.Module) -> None: def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"): quant_config = vllm_config.quant_config - def replace_linear_by_name(name: str, style: str): + def replace_linear_by_name(name: str, style: Style): linear = module.get_submodule(name) replacement = replace_linear_class( linear, style, quant_config, prefix=maybe_prefix(prefix, name) @@ -266,8 +263,10 @@ def replace_linear_by_name(name: str, style: str): setattr(module, name, replacement) log_replacement(maybe_prefix(prefix, name), linear, replacement) - if self.has_q_lora: - q_a = module.get_submodule(self.q_a_proj_name) + if (q_a_proj_name := self.q_a_proj_name) is not None: + # `match` sets the q-LoRA names together, or none of them + assert self.q_b_proj_name is not None + q_a = module.get_submodule(q_a_proj_name) kv_a = module.get_submodule(self.kv_a_proj_name) merged = MergedColumnParallelLinear( input_size=q_a.in_features, @@ -280,7 +279,7 @@ def replace_linear_by_name(name: str, style: str): ) logger.debug( "%s: %s, %s: %s -> %s: %s", - self.q_a_proj_name, + q_a_proj_name, q_a, self.kv_a_proj_name, kv_a, @@ -289,10 +288,11 @@ def replace_linear_by_name(name: str, style: str): ) setattr(module, self.merged_name, merged) # The rewritten forward calls the merged projection instead. - delattr(module, self.q_a_proj_name) + delattr(module, q_a_proj_name) delattr(module, self.kv_a_proj_name) replace_linear_by_name(self.q_b_proj_name, "colwise") else: + assert self.q_proj_name is not None replace_linear_by_name(self.kv_a_proj_name, "replicate") replace_linear_by_name(self.q_proj_name, "colwise") diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index dd2c5339ce5a..f7972f03a6e7 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -216,7 +216,7 @@ def _eps_source( finally: for name, value in candidates.items(): setattr(module, name, value) - if (name := markers.get(marked)) is not None: + if marked is not None and (name := markers.get(marked)) is not None: return name, None logger.debug_once( "%s does not hold its eps (%s) in an attribute. Every instance in this " @@ -251,18 +251,18 @@ def fuse( ) -> nn.Module: """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" weight = getattr(module, "weight", None) - has_weight = weight is not None - hidden_size = weight.size(0) if has_weight else 0 + hidden_size = weight.size(0) if weight is not None else 0 eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps if not isinstance(eps, (int, float)): # If eps was not detected, match torch behaviour. - dtype = weight.dtype if has_weight else vllm_config.model_config.dtype + config_dtype = vllm_config.model_config.dtype + dtype = weight.dtype if weight is not None else config_dtype eps = torch.finfo(dtype).eps if self.zero_centered: return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps) return TPAwareRMSNorm( hidden_size=hidden_size, eps=eps, - has_weight=has_weight, - dtype=weight.dtype if has_weight else None, + has_weight=weight is not None, + dtype=weight.dtype if weight is not None else None, ) diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 3cdc4e1201ba..cb630866a390 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -15,17 +15,19 @@ import textwrap from collections.abc import Callable from itertools import chain +from typing import Any from unittest import mock import torch from torch import fx, nn from torch.nn import functional as F +from typing_extensions import TypeIs from vllm.logger import init_logger logger = init_logger(__name__) -_UNKNOWN = object() +_UNKNOWN: Any = object() """Sentinel meta value for proxies whose concrete value could not be inferred. Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`).""" @@ -34,7 +36,7 @@ `call_module` nodes; meta execution must call modules for real.""" -def is_leaf_call(node: object) -> bool: +def is_leaf_call(node: object) -> TypeIs[fx.Node]: """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" return isinstance(node, fx.Node) and node.meta.get("leaf_call", False) @@ -60,7 +62,7 @@ class _MetaProxy(fx.Proxy): executing each op on the meta values, so PyTorch's meta kernels are the single source of shape inference — no per-op rules.""" - meta: object = _UNKNOWN + meta: Any = _UNKNOWN def __len__(self) -> int: if self.meta is not _UNKNOWN: @@ -149,6 +151,7 @@ def meta_of(arg: object) -> object: if unknown: return _UNKNOWN if kind == "call_function": + assert callable(target) return target(*meta_args, **meta_kwargs) if kind == "call_method": receiver, *rest = meta_args @@ -576,7 +579,7 @@ def _aliasing_reads(node: ast.expr) -> set[str]: would taint nearly every later name. """ names: set[str] = set() - stack = [node] + stack: list[ast.AST] = [node] while stack: current = stack.pop() if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS: @@ -646,7 +649,7 @@ def aliasing_names( """ names = set(seed) assigns = [ - node + (node, node.value) for node in ast.walk(funcdef) if isinstance(node, (ast.Assign, ast.AnnAssign)) and node.value is not None @@ -655,10 +658,10 @@ def aliasing_names( changed = True while changed: changed = False - for node in assigns: - if not _aliasing_reads(node.value) & names: + for node, value in assigns: + if not _aliasing_reads(value) & names: continue - targets = getattr(node, "targets", None) or [node.target] + targets = node.targets if isinstance(node, ast.Assign) else [node.target] stores = { child.id for target in targets @@ -770,7 +773,7 @@ def peel(node: object) -> object: return node -def is_fn(node: object, target: Callable) -> bool: +def is_fn(node: object, target: object) -> TypeIs[fx.Node]: """Is node `()`.""" return ( isinstance(node, fx.Node) @@ -779,14 +782,14 @@ def is_fn(node: object, target: Callable) -> bool: ) -def is_method(node: object, name: str) -> bool: +def is_method(node: object, name: str) -> TypeIs[fx.Node]: """Is node `.()`.""" return ( isinstance(node, fx.Node) and node.op == "call_method" and node.target == name ) -def is_op(node: object, name: str) -> bool: +def is_op(node: object, name: str) -> TypeIs[fx.Node]: """Is node `.()` for torch, F, operator, or Tensor.""" return any( is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) diff --git a/vllm/model_executor/models/transformers/legacy.py b/vllm/model_executor/models/transformers/legacy.py index 9c8a5c45b6f1..7a256f29d5be 100644 --- a/vllm/model_executor/models/transformers/legacy.py +++ b/vllm/model_executor/models/transformers/legacy.py @@ -20,35 +20,38 @@ import torch +from vllm.model_executor.models.utils import WeightsMapper from vllm.sequence import IntermediateTensors +from .base import Base + if TYPE_CHECKING: from vllm.config import VllmConfig -class LegacyMixin: +class LegacyMixin(Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - # Drop unsupported/unwanted output embeddings layers. - self.hf_to_vllm_mapper.orig_to_new_prefix.update( - { + self.hf_to_vllm_mapper |= WeightsMapper( + # Drop unsupported/unwanted output embeddings layers. + orig_to_new_prefix={ "model.lm_head.": None, "model.predictions.": None, "model.qa_outputs.": None, "model.embeddings_project.": None, "model.discriminator_predictions.": None, - } + }, + orig_to_new_substr={ + # Some encoder models have the position_ids buffer in the checkpoint. + # vLLM always passes position_ids as an argument, so drop the buffer. + "position_ids": None, + # Some encoder models have the bias of the final classifier layer in + # the checkpoint. vLLM does not use this bias, so drop it. + "score.bias": None, + }, ) - # Some encoder models have the position_ids buffer in the checkpoint. vLLM will - # always pass position_ids as an argument, so we drop the buffer if it exists. - self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None - - # Some encoder models have the bias of the final classifier layer in the - # checkpoint. vLLM does not use this bias, so we drop it if it exists. - self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None - # roberta-like models an extra padding in positions. # FIXME(Isotr0py): This is quite hacky for roberta edge case, # we should find a better way to handle this. @@ -61,6 +64,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + **kwargs, ) -> torch.Tensor | IntermediateTensors: if self.is_roberta: # RoBERTa positions start at padding_idx + 1. @@ -72,4 +76,5 @@ def forward( positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, + **kwargs, ) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 480bc7f66e8d..9f63e7333670 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -41,6 +41,7 @@ from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op +from .base import Base from .utils import log_replacement, maybe_per_layer if TYPE_CHECKING: @@ -66,14 +67,15 @@ def __init__(self, *args, moe_state: TransformersMoEState, **kwargs): self._moe_state = moe_state self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel - def forward( + def forward( # type: ignore[override] self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, **kwargs: Any, ) -> torch.Tensor: - """In Transformers `experts.forward` will have this signature. + """In Transformers `experts.forward` will have this signature, which is why + it deliberately does not match `MoERunner.forward`. We discard any extra kwargs because we cannot use them here.""" # Note: we need to forward through a custom op so the topk_ids @@ -126,10 +128,18 @@ def _transformers_moe_forward_fake( class TransformersRoutedExperts(RoutedExperts): def get_expert_mapping( - self, include_fused: bool = False + self, + ckpt_gate_proj_name: str | None = None, + ckpt_down_proj_name: str | None = None, + ckpt_up_proj_name: str | None = None, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: - common_names = ("gate_proj", "down_proj", "up_proj") - common_map = super().get_expert_mapping(*common_names, include_fused) + common_map = super().get_expert_mapping( + ckpt_gate_proj_name or "gate_proj", + ckpt_down_proj_name or "down_proj", + ckpt_up_proj_name or "up_proj", + include_fused, + ) mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused) if not include_fused: return common_map + mixtral_map @@ -138,7 +148,13 @@ def get_expert_mapping( return common_fused + mixtral_fused + common_unfused + mixtral_unfused -class MoEMixin(MixtureOfExperts): +class MoEMixin(MixtureOfExperts, Base): + mlp_layers: list[nn.Module] + """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods.""" + + # Redeclared because `recursive_replace` assigns it below the first read + num_local_physical_experts: int + def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.check_version("5.0.0", "MoE models support") # Skip MixtureOfExperts.__init__ and call the next class in MRO @@ -197,7 +213,8 @@ def recursive_replace(self): # Dtype the router computes in, if it is not the activation dtype. config_router_dtype = getattr(text_config, "moe_router_dtype", None) - config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) + if config_router_dtype is not None: + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) # Grouped topk routing kwargs num_expert_group = getattr(text_config, "n_group", None) @@ -222,7 +239,7 @@ def recursive_replace(self): # MixtureOfExperts mixin settings ep_size = get_ep_group().world_size - self.mlp_layers = [] # Used for MixtureOfExperts methods + self.mlp_layers = [] self.moe_layers = [] self.num_expert_groups = 1 if num_expert_group is None else num_expert_group self.num_logical_experts = num_experts @@ -321,9 +338,10 @@ def _recursive_replace(module: nn.Module, prefix: str): if shared_experts is not None: hf_shared = shared_experts.shared_experts glu_fuser = get_fuser(hf_shared, GLUFuser) - down_name = glu_fuser and glu_fuser.down_name - if down_name is not None: - shared_down_projs.append((hf_shared, down_name)) + if glu_fuser is not None: + down_name = glu_fuser.down_name + if down_name is not None: + shared_down_projs.append((hf_shared, down_name)) # Prefer config, otherwise read it from fuser. router_dtype = config_router_dtype or fuser.router_dtype gate = fuser.gate(moe_block, prefix, router_dtype) @@ -374,7 +392,9 @@ def custom_routing_function( # Handle all gather in expert parallel if topk_ids.size(0) != hidden_states.size(0): dp_metadata = get_forward_context().dp_metadata + assert dp_metadata is not None sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + assert sizes is not None is_sp = moe_state.is_sequence_parallel group = get_ep_group() if is_sp else get_dp_group() assert sizes[group.rank_in_group] == topk_ids.shape[0] diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 855b267c02f5..164a4dbd8396 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -17,17 +17,24 @@ """Transformers modeling backend mixin for multi-modal models.""" from collections import defaultdict -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import ExitStack, contextmanager from typing import TYPE_CHECKING, Any import torch import transformers from packaging.version import Version +from transformers.utils.generic import ModelOutput from vllm.compilation.decorators import should_torch_compile_mm_encoder +from vllm.config.multimodal import AudioDummyOptions from vllm.config.utils import getattr_iter -from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input +from vllm.inputs import ( + MultiModalDataBuiltins, + MultiModalDataDict, + MultiModalInput, + mm_input, +) from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -62,6 +69,8 @@ from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d +from .base import Base + if TYPE_CHECKING: from transformers import BatchFeature, PreTrainedModel @@ -148,7 +157,7 @@ def get_max_audio_tokens(self) -> int: def get_max_image_tokens(self) -> int: width, height = self.get_image_size_with_most_features() processor = self.get_hf_processor() - multimodal_config = self.ctx.model_config.multimodal_config + multimodal_config = self.ctx.model_config.get_multimodal_config() mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {} mm_tokens = processor._get_num_multimodal_tokens( image_sizes=([height, width],), **mm_processor_kwargs @@ -188,7 +197,7 @@ def get_dummy_mm_data( mm_counts: Mapping[str, int], mm_options: Mapping[str, "BaseDummyOptions"], ) -> MultiModalDataDict: - data: MultiModalDataDict = {} + data = MultiModalDataBuiltins() if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): sampling_rate = self.info._get_audio_sampling_rate() sub = self.info._get_audio_processor() @@ -196,10 +205,14 @@ def get_dummy_mm_data( if chunk_length is None: chunk_length = 30 audio_len = int(chunk_length * sampling_rate) + audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance( + audio_overrides, AudioDummyOptions + ) data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, - overrides=mm_options.get("audio"), + overrides=audio_overrides, ) if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): target_width, target_height = self.info.get_image_size_with_most_features() @@ -359,7 +372,9 @@ def _unpad_audios( itself is to process it by itself. """ audios = mm_data.get("audio") - if not audios or len({len(audio) for audio in audios}) == 1: + if not isinstance(audios, Iterable) or not audios: + return + if len({len(audio) for audio in audios}) == 1: return alone = [ @@ -444,7 +459,8 @@ def _get_mm_token_ids(self, modality: str) -> list[int]: token_id = getattr(processor, names[0], getattr_iter(config, names)) if token_id is None: token = getattr(processor, f"{modality}_token", None) - token_id = info.get_tokenizer().get_vocab().get(token) + if token is not None: + token_id = info.get_tokenizer().get_vocab().get(token) if token_id is None: raise ValueError( f"Cannot find {modality}_token_id on processor or model config" @@ -722,13 +738,15 @@ def get_target_token_ids(token: str | int | Sequence[int]) -> list[int]: for modality, items in out_mm_kwargs.items(): # Popped so they are neither cached nor sent to the model; the updates # they produce are cached alongside the item instead - replacements = [ - PromptUpdateDetails.select_token_id( - (ids := item.pop(f"{modality}_replacement_ids").data).tolist(), - _get_embed_token_id(ids), + replacements = [] + for item in items: + ids = item.pop(f"{modality}_replacement_ids").data + assert isinstance(ids, torch.Tensor) + replacements.append( + PromptUpdateDetails.select_token_id( + ids.tolist(), _get_embed_token_id(ids) + ) ) - for item in items - ] token = getattr(hf_processor, f"{modality}_token") updates.append( PromptReplacement( @@ -777,7 +795,7 @@ def _get_num_patches_per_image( ) -> list[int] | None: """Ask the HF processor how many rows of image data each image produces.""" images = mm_data.get("images") - if not images: + if not isinstance(images, Iterable) or not images: return None try: sizes = [(image.height, image.width) for image in images] @@ -893,7 +911,7 @@ def _apply_hf_processor_main( ) -class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): +class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip SupportsMRoPE.__init__ and call the next class in MRO super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix) @@ -925,7 +943,7 @@ def _mark_model_components(self, vllm_config: "VllmConfig"): if model_config.skip_tokenizer_init: # Determining the supported modalities needs the HF processor, which in # turn needs a tokenizer - mm_config = model_config.multimodal_config + mm_config = model_config.get_multimodal_config() if mm_config.mm_encoder_only or any( mm_config.get_limit_per_prompt(modality) == 0 for modality in encoder_classes @@ -1029,7 +1047,7 @@ def get_language_model(self) -> torch.nn.Module: # Exclude MultiModalMixin itself bases = [b for b in bases if b is not MultiModalMixin] - class LanguageModel(*bases): + class LanguageModel(*bases): # type: ignore[misc] def __init__(self, multimodal_model): # Don't call super().__init__() to avoid re-initialization self.__dict__.update(multimodal_model.__dict__) @@ -1203,7 +1221,7 @@ def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> Any: # from `self.get_image_features` if isinstance(features, tuple): return features[0] - if isinstance(features, dict): + if isinstance(features, ModelOutput): return features.pooler_output return features diff --git a/vllm/model_executor/models/transformers/pooling.py b/vllm/model_executor/models/transformers/pooling.py index 47a9c23d2c3f..79594387b740 100644 --- a/vllm/model_executor/models/transformers/pooling.py +++ b/vllm/model_executor/models/transformers/pooling.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING import torch +import torch.nn as nn from transformers import AutoModelForSequenceClassification from vllm.config.utils import getattr_iter @@ -26,11 +27,25 @@ from vllm.model_executor.models.interfaces import SupportsCrossEncoding from vllm.model_executor.models.interfaces_base import VllmModelForPooling +from .base import Base + if TYPE_CHECKING: from vllm.config import VllmConfig -class EmbeddingMixin(VllmModelForPooling): +class ClassifierWithReshape(nn.Module): + """Token extraction has already been applied in `pooler.pooling`. + + Add dim to match expected input shape of `classifier.forward`. + """ + + def forward(self, *args, **kwargs): + if len(args) > 0: + args = (args[0].unsqueeze(1), *args[1:]) + return super().forward(*args, **kwargs) + + +class EmbeddingMixin(VllmModelForPooling, Base): default_seq_pooling_type = "CLS" def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): @@ -45,7 +60,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.pooler = DispatchPooler.for_embedding(pooler_config) -class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling): +class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling, Base): default_seq_pooling_type = "CLS" def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): @@ -83,17 +98,13 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): ) self.init_parameters(self.classifier, dtype=self.model_config.head_dtype) - class ClassifierWithReshape(self.classifier.__class__): - """Token extraction has already been applied in `pooler.pooling`. - Add dim to match expected input shape of `classifier.forward`. - """ - - def forward(self, *args, **kwargs): - if len(args) > 0: - args = (args[0].unsqueeze(1), *args[1:]) - return super().forward(*args, **kwargs) - - self.classifier.__class__ = ClassifierWithReshape + # Order `ClassifierWithReshape` ahead of the classifier's own class so that + # its `super().forward(...)` reaches the original implementation. + self.classifier.__class__ = type( + "ClassifierWithReshape", + (ClassifierWithReshape, type(self.classifier)), + {}, + ) self.pooler = DispatchPooler.for_seq_cls( pooler_config, diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index ddb2f4fb7139..0443f390c465 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -18,11 +18,10 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager -from functools import lru_cache from itertools import chain from operator import attrgetter from pathlib import Path -from typing import TYPE_CHECKING, Literal, TypeVar +from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar import torch from torch import nn @@ -70,7 +69,7 @@ def register_empty_parameter(module, name, param): module._parameters[name].to(device), **kwargs ) - tensor_constructors_to_patch = {} + tensor_constructors_to_patch: dict[str, Callable] = {} def patch_tensor_constructor(fn): def wrapper(*args, **kwargs): @@ -147,8 +146,8 @@ def replace_linear_class( ) -TorchConv = nn.Conv2d | nn.Conv3d -VllmConv = Conv2dLayer | Conv3dLayer +TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d +VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: @@ -212,7 +211,9 @@ class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbeddin `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding.""" -@lru_cache +_rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {} + + def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]: """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`. @@ -224,7 +225,10 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe The new class, to assign to `__class__` of an instance of `cls`. """ - return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) + if cls not in _rebased_embedding_classes: + rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) + _rebased_embedding_classes[cls] = rebased_cls + return _rebased_embedding_classes[cls] def replace_embedding_class( @@ -278,10 +282,9 @@ def _recursive_replace(module: nn.Module, prefix: str): qual_name = maybe_prefix(prefix, child_name) # Replace modules as needed if isinstance(child_module, nn.Linear): - style = "replicate" new_module = replace_linear_class( child_module, - style, + "replicate", quant_config, prefix=qual_name, ) From fef55fd26dc633fcc87978f12f1a97e08e351c02 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:21:22 -0700 Subject: [PATCH 420/774] [ROCm][CI] Adapt MoE tests to the triton_kernels 3.8 API (#57385) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5 (1M context) --- .../moe/test_gpt_oss_triton_kernels.py | 38 ++++++++++++------- .../moe/test_modular_oai_triton_moe.py | 26 ++++++++----- tests/kernels/moe/utils.py | 24 ++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 16 +++----- .../layers/quantization/utils/mxfp4_utils.py | 9 +++++ 5 files changed, 79 insertions(+), 34 deletions(-) diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 86028a7feba0..c1d83d8349bf 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from vllm.platforms import current_platform -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels if not has_triton_kernels(): pytest.skip( @@ -15,9 +15,15 @@ allow_module_level=True, ) -import triton_kernels.matmul_ogs_details.opt_flags as opt_flags import triton_kernels.swiglu -from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig + +if get_triton_kernels_version() == "3.8": + # 3.8: matmul_ogs -> matmul, matmul_ogs_details -> matmul_details + import triton_kernels.matmul_details.opt_flags as opt_flags + from triton_kernels.matmul import FlexCtx, PrecisionConfig +else: + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig from triton_kernels.numerics import InFlexData from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor @@ -28,9 +34,10 @@ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( triton_kernel_moe_forward, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs from vllm.utils.math_utils import round_up -from .utils import shuffle_weight +from .utils import mxfp4_w_layouts, shuffle_weight def deshuffle(w: torch.Tensor): @@ -133,12 +140,12 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int): x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0) - w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1) - w_scale_layout, w_scale_layout_opts = ( - layout.make_default_matmul_mxfp4_w_scale_layout( - mx_axis=1, num_warps=num_warps - ) - ) + ( + w_layout, + w_layout_opts, + w_scale_layout, + w_scale_layout_opts, + ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps) w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1) w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, torch.bfloat16, axis=1) @@ -165,10 +172,10 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int): ) pc1 = PrecisionConfig( - weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData()) + **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) ) pc2 = PrecisionConfig( - weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData()) + **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) ) # tucuate so the rest can run properly @@ -281,8 +288,6 @@ class Case: @pytest.mark.parametrize("num_token", [2]) @pytest.mark.parametrize("tp", [1, 2, 4, 8]) def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init): - from triton_kernels.tensor_details import layout - if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"): pytest.skip("make_default_matmul_mxfp4_w_layout not available") @@ -387,6 +392,11 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): if gptoss_moe.triton_kernels_version == "3.5.1": pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+") + if gptoss_moe.triton_kernels_version == "3.8": + pytest.skip( + "3.8 rebuilds routing from raw ids, so routing_data_from_sparse_topk " + "just forwards to make_routing_data and the comparison is vacuous" + ) from triton_kernels.topk import topk as topk_fn diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 47be82dd2ab1..9665079e373d 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -10,7 +10,7 @@ from tests.utils import wait_for_gpu_memory_to_clear from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels if not has_triton_kernels(): pytest.skip( @@ -18,11 +18,14 @@ allow_module_level=True, ) -from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig +if get_triton_kernels_version() == "3.8": + # 3.8: matmul_ogs -> matmul + from triton_kernels.matmul import FlexCtx, PrecisionConfig +else: + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig from triton_kernels.numerics import InFlexData from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor -from triton_kernels.tensor_details import layout from triton_kernels.testing import assert_close from vllm.config import VllmConfig, set_current_vllm_config @@ -35,11 +38,12 @@ UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import set_random_seed -from .utils import make_dummy_moe_config, shuffle_weight +from .utils import make_dummy_moe_config, mxfp4_w_layouts, shuffle_weight MNK = [ (1, 512, 384), @@ -130,10 +134,12 @@ def make_weights(dtype, k, n, e): w2 = w2[..., :n, :k] num_warps = 8 - w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1) - w_scale_layout, w_scale_layout_opts = ( - layout.make_default_matmul_mxfp4_w_scale_layout(mx_axis=1, num_warps=num_warps) - ) + ( + w_layout, + w_layout_opts, + w_scale_layout, + w_scale_layout_opts, + ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps) w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts) w1_scale_tri = convert_layout( @@ -150,10 +156,10 @@ def make_weights(dtype, k, n, e): ) w1_precision_config = PrecisionConfig( - weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData()) + **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData()) + **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) ) return ( diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index e7624729502d..d92e4d6129ba 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -705,3 +705,27 @@ def check_accuracy(a, b, atol, rtol, percent): f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " f"(threshold: {1 - percent:.4f})" ) + + +def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): + """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs. + + triton_kernels 3.8 returns layout instances; earlier versions return a + (layout, opts) tuple. + """ + from triton_kernels.tensor_details import layout + + from vllm.utils.import_utils import get_triton_kernels_version + + if get_triton_kernels_version() == "3.8": + w = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) + s = layout.make_default_matmul_mxfp4_w_scale_layout( + mx_axis=mx_axis, num_warps=num_warps + ) + return w, {}, s, {} + + w, w_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) + s, s_opts = layout.make_default_matmul_mxfp4_w_scale_layout( + mx_axis=mx_axis, num_warps=num_warps + ) + return w, w_opts, s, s_opts diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index d1c151aca247..7fed4124f569 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -31,7 +31,12 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + _swizzle_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mx_scale_kwargs as _mx_scale_kwargs, +) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, ) @@ -70,15 +75,6 @@ ) -def _mx_scale_kwargs(scale): - """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size, - 3.5.1/3.6 use weight_scale. - """ - if triton_kernels_version == "3.8": - return {"b_mx_scale": scale, "b_microblock_size": 32} - return {"weight_scale": scale} - - def _pack_deepgemm_mxfp4_scales( w13_weight: torch.Tensor, w2_weight: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index a1d050177507..8cc626f3866a 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -36,6 +36,15 @@ def weight_mx_scale(precision_config): return precision_config.weight_scale +def mx_scale_kwargs(scale): + """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size, + 3.5.1/3.6 use weight_scale. + """ + if get_triton_kernels_version() == "3.8": + return {"b_mx_scale": scale, "b_microblock_size": 32} + return {"weight_scale": scale} + + def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel.""" assert has_triton_kernels() From e30bf70c5dfb90fd9b31a8f4692b49df9b223704 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:22:39 -0700 Subject: [PATCH 421/774] [ROCm][CI] Fix Entrypoints Integration (Pooling) tests on TheRock image (#57380) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5 (1M context) --- .../test_cross_encoder_online_vision.py | 58 ++++++------------- 1 file changed, 17 insertions(+), 41 deletions(-) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index 9f2d7da80116..bf3c771fdd1b 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -27,31 +27,10 @@ ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else ["auto"] -# Per-backend tolerance with explicit entries; "default" is the fallback -BACKEND_TOL: dict[str, float] = { - "default": 0.05, # 5% tolerance for other backends (e.g. FLASH_ATTN) - # Relaxed tolerances for ROCm attn - # See: https://github.com/vllm-project/vllm/issues/35569 - "ROCM_ATTN": 0.09, # gfx950:~8.45%, gfx942:~3.70% - "ROCM_AITER_FA": 0.045, # gfx950:~2.00%, gfx942:~0.80% - "TRITON_ATTN": 0.045, # gfx950:~3.00%, gfx942:~2.20% - "FLEX_ATTENTION": 0.045, # gfx950:~3.25%, gfx942:~1.10% -} - -# Some ROCm attention backends show small absolute drift on the low -# text-vs-text probability even though larger scores remain well inside the -# relative tolerance. The absolute drift is uniform across score magnitudes -# (~0.005-0.010), so it only exceeds the relative tolerance for the small -# ~0.10 text-vs-text value. Keep the relative tolerances tight and add only a -# small absolute floor for the affected backends. -# TRITON_ATTN: gfx942/ROCm 7.2 drifts ~0.008 abs on text-vs-text (~7.9% rel). -BACKEND_ABS_TOL: dict[str, float] = { - "default": 0.0, - "auto": 0.007, - "ROCM_AITER_FA": 0.005, - "TRITON_ATTN": 0.009, - "FLEX_ATTENTION": 0.00875, -} +# The absolute floor matters only for the low ~0.10 text-vs-text score, where a +# few thousandths of drift would otherwise exceed the relative bound. +REL_TOL = 0.05 +ABS_TOL = 0.007 # ROCm: disable prefix caching and eliminate batch variance to reduce # test flakiness. @@ -62,17 +41,9 @@ ) -def get_tol(backend: str) -> float: - return BACKEND_TOL.get(backend, BACKEND_TOL["default"]) - - -def get_abs_tol(backend: str) -> float: - return BACKEND_ABS_TOL.get(backend, BACKEND_ABS_TOL["default"]) - - def assert_score(actual: float, expected: float, backend: str, label: str): - tol = get_tol(backend) - abs_tol = get_abs_tol(backend) + tol = REL_TOL + abs_tol = ABS_TOL diff = abs(actual - expected) rel_diff = diff / abs(expected) if expected != 0 else diff print( @@ -104,9 +75,16 @@ def assert_score(actual: float, expected: float, backend: str, label: str): }, ] -TEXT_VS_TEXT = 0.10040374100208282 -TEXT_VS_IMAGE = 0.7423753142356873 -TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838 +# One bf16 ULP in attention compounds to ~0.1% here, so ROCm backends each +# settle differently; these are their midpoint across backends and archs. +if current_platform.is_rocm(): + TEXT_VS_TEXT = 0.105243 + TEXT_VS_IMAGE = 0.742472 + TEXT_VS_TEXT_PLUS_IMAGE = 0.544568 +else: + TEXT_VS_TEXT = 0.10040374100208282 + TEXT_VS_IMAGE = 0.7423753142356873 + TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838 @pytest.fixture(scope="module", params=ATTN_BACKENDS) @@ -507,6 +485,4 @@ async def test_rerank_api_instruction_field_matches_chat_template_kwargs( kwargs_scores = [ r.relevance_score for r in sorted(kwargs_rerank.results, key=lambda x: x.index) ] - assert field_scores == pytest.approx( - kwargs_scores, rel=get_tol(backend), abs=get_abs_tol(backend) - ) + assert field_scores == pytest.approx(kwargs_scores, rel=REL_TOL, abs=ABS_TOL) From b2f2bd71ba6cff6d0678e2f83d45c1c08425144c Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Thu, 17 Sep 2026 10:02:24 -0700 Subject: [PATCH 422/774] [CI][Bugfix] Add tp_shard_with_padding to padded MoE reload test mock (#57402) Signed-off-by: kimi-no-na-wa Co-authored-by: kimi-no-na-wa Co-authored-by: Misha Goin --- tests/model_executor/model_loader/test_reload.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 844149ce1b31..fde8b86cdaef 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -489,6 +489,7 @@ def test_padded_moe_reload_releases_each_layer( is_act_and_mul=is_gated, has_bias=has_bias, tp_rank=tp_rank, + tp_shard_with_padding=False, moe_parallel_config=SimpleNamespace(tp_size=2), ) model = torch.nn.ModuleList() From f3aa88d2309523f8438a239a55f59a28179ce5ef Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:05:52 -0400 Subject: [PATCH 423/774] [Kimi K3 Bug] Fix kimi k3 reasoning parser (#57098) Signed-off-by: yewentao256 Co-authored-by: Flora Feng <4florafeng@gmail.com> --- vllm/reasoning/kimi_k3_reasoning_parser.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py index d3ae2b3abe38..e5618d47be4a 100644 --- a/vllm/reasoning/kimi_k3_reasoning_parser.py +++ b/vllm/reasoning/kimi_k3_reasoning_parser.py @@ -256,11 +256,13 @@ def extract_reasoning( ) -> tuple[str | None, str | None]: """Split full text into ``(reasoning, rest)`` for the non-streaming path. - Handles three shapes: - * no think channel at all -> ``(None, model_output)`` (all content) + Handles four shapes: + * response opener without think markers -> no think channel; all content * open marker present -> reasoning starts after ``<|open|>think<|sep|>`` * open marker absent but a close marker exists (gen-prefix consumed the open) -> reasoning starts at offset 0 + * neither marker present -> truncated reasoning after a consumed + generation prefix ``rest`` is whatever follows the close marker, fed on to the tool parser. """ if not self._thinking_enabled: @@ -270,11 +272,14 @@ def extract_reasoning( # reasoning content begins right after think-open (or at start if the # open marker was already consumed as a generation prefix) content_start = m_open.end() if m_open is not None else 0 - # if there is no think channel at all, everything is content - if m_open is None and self._think_close_re.search(model_output) is None: - return None, self._content_after_reasoning(model_output, request) m_close = self._think_close_re.search(model_output, content_start) + if ( + m_open is None + and m_close is None + and self._response_open_re.search(model_output) is not None + ): + return None, self._content_after_reasoning(model_output, request) if m_close is not None: reasoning = model_output[content_start : m_close.start()] rest = model_output[m_close.end() :] From ff6b5808c4571a37ba3d27f74d74ef43c3a14245 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Thu, 17 Sep 2026 20:51:35 +0300 Subject: [PATCH 424/774] [ROCm] Resolve the indexer fp8 cache dtype once at import (#53792) Signed-off-by: Sriram Kumar Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 07cdf181069a..3d891c8ce0f7 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -29,6 +29,8 @@ logger = init_logger(__name__) +FP8_DTYPE = current_platform.fp8_dtype() + @functools.cache def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None: @@ -260,8 +262,7 @@ def indexer_k_quant_and_cache_triton( # In real layout, we store the first portion as kv cache value # and second portion as kv cache scale kv_cache = kv_cache.view(num_blocks, -1) - fp8_dtype = current_platform.fp8_dtype() - kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype) + kv_cache_value = kv_cache[:, : block_size * head_dim].view(FP8_DTYPE) kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32) head_tile_size = head_tile_size // kv_cache.element_size() layout = "NORMAL" if block_size == 1 else "SHUFFLE" @@ -279,7 +280,7 @@ def indexer_k_quant_and_cache_triton( layout, block_tile_size, head_tile_size, - IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz, + IS_FNUZ=torch.float8_e4m3fnuz == FP8_DTYPE, USE_UE8M0=scale_fmt == "ue8m0", ) @@ -465,8 +466,7 @@ def cp_gather_indexer_k_quant_cache_triton( num_blocks = k_cache.shape[0] # we assume the kv cache already been split to 2 portion k_cache = k_cache.view(num_blocks, -1) - fp8_dtype = current_platform.fp8_dtype() - k_cache_value = k_cache[:, : block_size * head_dim].view(fp8_dtype) + k_cache_value = k_cache[:, : block_size * head_dim].view(FP8_DTYPE) k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32) grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) @@ -516,7 +516,6 @@ def fp8_paged_mqa_logits_torch( ): from vllm.utils.math_utils import cdiv - fp8_dtype = current_platform.fp8_dtype() batch_size, next_n, _, dim = q.size() if next_n == 1: block_size = kv_cache.shape[1] @@ -540,7 +539,7 @@ def fp8_paged_mqa_logits_torch( cache = kv_cache_flat[pages] scale_offset = block_size * dim cache_value = ( - cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32) + cache[..., :scale_offset].view(dtype=FP8_DTYPE).to(torch.float32) ) cache_scale = ( cache[..., scale_offset:].view(dtype=torch.float32).contiguous() @@ -558,7 +557,7 @@ def fp8_paged_mqa_logits_torch( kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:] scale = scale.contiguous().view(torch.float) q = q.float() - kv_cache = kv_cache.view(fp8_dtype).float() * scale + kv_cache = kv_cache.view(FP8_DTYPE).float() * scale num_block, block_size, _, dim = kv_cache.size() logits = torch.full( [batch_size * next_n, max_model_len], @@ -1021,7 +1020,6 @@ def rocm_aiter_sparse_attn_indexer( # careful! this will be None in dummy run forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata - fp8_dtype = current_platform.fp8_dtype() from vllm.utils.torch_utils import _resolve_layer_name k_cache_prefix = _resolve_layer_name(k_cache_prefix) @@ -1037,7 +1035,7 @@ def rocm_aiter_sparse_attn_indexer( # Prefill k_fp8 and k_scale buffers, used by # rocm_aiter_sparse_attn_indexer's prefill path workspace_manager.get_simultaneous( - ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, head_dim), FP8_DTYPE), ((total_seq_lens, 4), torch.uint8), ) @@ -1117,7 +1115,7 @@ def rocm_aiter_sparse_attn_indexer( workspace_manager = current_workspace_manager() k_fp8_full, k_scale_full = workspace_manager.get_simultaneous( - ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, head_dim), FP8_DTYPE), ((total_seq_lens, 4), torch.uint8), ) for chunk in prefill_metadata.chunks: From 09c379ab6183085fd90e8195792cc63f5247cdf3 Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Thu, 17 Sep 2026 19:09:18 +0100 Subject: [PATCH 425/774] [Bugfix] Max-load throughput cliff when `max_num_seqs` is not a multiple of 8 (#57355) Signed-off-by: Andy Lo Co-authored-by: OpenAI Codex --- tests/compile/test_config.py | 19 +++++++++++++++++++ vllm/config/vllm.py | 8 +++++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5c9132ccbfe2..7465ea442c89 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -541,6 +541,25 @@ def _mock_config_for_cudagraph_sizes( return config +@pytest.mark.parametrize("max_num_seqs", [100, 101]) +def test_default_cudagraph_capture_sizes_cover_off_stride_max_num_seqs( + max_num_seqs: int, +) -> None: + compilation_config = CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE + ) + config = _mock_config_for_cudagraph_sizes( + max_num_seqs=max_num_seqs, + num_speculative_tokens=0, + max_num_batched_tokens=32768, + compilation_config=compilation_config, + ) + + VllmConfig._set_cudagraph_sizes(config) + + assert max_num_seqs in compilation_config.cudagraph_capture_sizes + + @pytest.mark.parametrize( ("max_num_seqs", "num_speculative_tokens", "widest_is_captured"), [ diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e5710dced396..4431743a3d57 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2223,9 +2223,9 @@ def _set_cudagraph_sizes(self): max_cudagraph_capture_size = ( self.compilation_config.max_cudagraph_capture_size ) - # Decode sizes to cover, in tokens. Populated only when a request - # is more than one token wide and only when the default is computed - # here, so an explicit capture range is left exactly as configured. + # Decode sizes to cover, in tokens. Populated only when the default + # is computed here, so an explicit capture range is left exactly as + # configured. uniform_decode_sizes: list[int] = [] if max_cudagraph_capture_size is None: from vllm.platforms import current_platform @@ -2316,6 +2316,8 @@ def request_counts(max_reqs: int) -> list[int]: if n * query_len <= max_cudagraph_capture_size } ) + elif max_num_seqs <= max_cudagraph_capture_size: + uniform_decode_sizes = [max_num_seqs] max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) From 6a344d8c8595fc7ca32a633b6f146d36dd227430 Mon Sep 17 00:00:00 2001 From: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:19:13 +0200 Subject: [PATCH 426/774] [Bugfix][Metrics] Do not log a 0.0% prefix cache hit rate before any query (#54990) Signed-off-by: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 16 ++++++++++++++++ vllm/v1/metrics/loggers.py | 11 ++++------- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 3c82180474f6..07c222214586 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1068,6 +1068,22 @@ def _stats(requests: int, queries: int, hits: int) -> PrefixCacheStats: return PrefixCacheStats(requests=requests, queries=queries, hits=hits) +def test_metrics_empty_distinguishes_no_queries_from_no_hits(): + """`hit_rate` alone cannot tell the two apart; `empty` can. + + Both an unobserved window and a genuine all-miss window report a hit + rate of 0.0, so anything surfacing that number to a human has to check + `empty` first - which is what the prefix-cache log line does. + """ + metrics = CachingMetrics(max_recent_requests=5) + assert metrics.empty + assert metrics.hit_rate == 0.0 + + metrics.observe(_stats(1, 20, 0)) + assert not metrics.empty + assert metrics.hit_rate == 0.0 + + def test_metrics(): """Test the prefix caching metrics.""" metrics = CachingMetrics(max_recent_requests=5) diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index fee6679fe760..26a4142df9f1 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -296,19 +296,16 @@ def log(self): log_parts.append("Preemptions: %d") log_args.append(self.num_preemptions) - log_parts.extend( - [ - "%s KV cache usage: %.1f%%", - "Prefix cache hit rate: %.1f%%", - ] - ) + log_parts.append("%s KV cache usage: %.1f%%") log_args.extend( [ self.kv_cache_device, self.last_scheduler_stats.kv_cache_usage * 100, - self.prefix_caching_metrics.hit_rate * 100, ] ) + if not self.prefix_caching_metrics.empty: + log_parts.append("Prefix cache hit rate: %.1f%%") + log_args.append(self.prefix_caching_metrics.hit_rate * 100) if envs.VLLM_COMPUTE_NANS_IN_LOGITS: log_parts.append("Corrupted: %d reqs") From a9a7e45f3174781f2b20c03c7ac1beee2b98f9d5 Mon Sep 17 00:00:00 2001 From: afriedri Date: Thu, 17 Sep 2026 13:39:52 -0500 Subject: [PATCH 427/774] [ROCm] Restore `VLLM_ROCM_USE_AITER_FP4_ASM_GEMM` and default w4a4 ASM GEMM back to off (#57055) Signed-off-by: Andy Friedrich --- tests/kernels/quantization/test_rocm_mxfp4.py | 16 ++++++++++++---- vllm/_aiter_ops.py | 6 +++++- vllm/envs.py | 6 ++++++ 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 0e14e2d2583c..8d1a4bb0a734 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -156,22 +156,26 @@ def test_fp4_env_defaults(): """ROCm FP4 env defaults should stay stable for the AITER gates.""" import vllm.envs as envs + assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True @pytest.mark.parametrize( ( "use_aiter", + "use_fp4_asm_gemm", "use_fp4bmm", ), [ - (True, True), - (True, False), - (False, True), + (True, True, True), + (True, True, False), + (True, False, True), + (False, True, True), ], ) def test_rocm_aiter_fp4_enablement_follows_env_and_arch( use_aiter, + use_fp4_asm_gemm, use_fp4bmm, monkeypatch, ): @@ -187,11 +191,15 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch( _assert_aiter_supported() on_gfx950_value = on_gfx950() - expected_asm_gemm = use_aiter and on_gfx950_value + expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value with monkeypatch.context() as mp: mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") + mp.setenv( + "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", + "1" if use_fp4_asm_gemm else "0", + ) mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0") _reload_envs() rocm_aiter_ops.refresh_env_variables() diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4b9a145f4830..c514b7c0bfef 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1893,6 +1893,7 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen. VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention. VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. + VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4). @@ -1960,6 +1961,8 @@ def get_moe_dispatch_policy(cls) -> int: _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM _LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM + # TODO: Consolidate under _LINEAR_ENABLED + _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS @@ -1989,6 +1992,7 @@ def refresh_env_variables(cls): cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM + cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 @@ -2230,7 +2234,7 @@ def is_linear_hipbmm_enabled(cls) -> bool: def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool: from vllm.platforms.rocm import on_gfx950 - return cls._AITER_ENABLED and on_gfx950() + return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950() @classmethod @if_aiter_supported diff --git a/vllm/envs.py b/vllm/envs.py index c512df457adc..82d6e771b528 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -143,6 +143,7 @@ VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True + VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False VLLM_ROCM_USE_AITER_FP8BMM: bool = True VLLM_ROCM_USE_AITER_FP4BMM: bool = True @@ -1305,6 +1306,11 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_USE_AITER_MHA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1") ), + # Whether to use aiter fp4 gemm asm. + # By default is disabled. + "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1") + ), # Whether to use aiter rope. # By default is disabled. "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: ( From 41f9104fa649aed2c26af8a459210619ad9514fe Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 18 Sep 2026 03:06:36 +0800 Subject: [PATCH 428/774] [DSv4.1] Integrate Mega-Gate from DeepGEMM (#56266) Signed-off-by: zjy0516 Signed-off-by: Thien Tran Co-authored-by: OpenAI Codex Co-authored-by: zjy0516 --- .buildkite/test_areas/kernels.yaml | 1 + tests/kernels/test_mhc_kernels.py | 5 +- tests/models/test_deepseek_v4_mega_moe.py | 150 ++++++++++++++++++++-- vllm/models/deepseek_v4/nvidia/dspark.py | 14 ++ vllm/models/deepseek_v4/nvidia/model.py | 129 ++++++++++++++++--- vllm/models/deepseek_v4/nvidia/mtp.py | 21 ++- vllm/models/deepseek_v41/nvidia/dspark.py | 14 ++ vllm/models/deepseek_v41/nvidia/model.py | 16 ++- vllm/utils/deep_gemm.py | 54 ++++++++ 9 files changed, 369 insertions(+), 35 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 5a7832d945d2..583378fdb49b 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -120,6 +120,7 @@ steps: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - vllm/models/deepseek_v4/nvidia/ + - vllm/utils/deep_gemm.py - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - tests/models/test_deepseek_v4_mega_moe.py commands: diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index fafb37d97b3c..08b1a24c85e7 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -132,6 +132,7 @@ def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) draft = SimpleNamespace( + use_mega_moe=False, use_sequence_parallel=False, hc_mult=hc_mult, layers=[make_layer(mix) for mix in mixes], @@ -452,7 +453,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( device=DEVICE, dtype=torch.bfloat16 ) decoder.attn = lambda positions, x, _: x * 0.5 - decoder.ffn = lambda x, input_ids: x * 0.25 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 with torch.device(DEVICE): decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 @@ -536,7 +537,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf device=DEVICE, dtype=torch.bfloat16 ) decoder.attn = lambda positions, x, _: x * 0.5 - decoder.ffn = lambda x, input_ids: x * 0.25 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 with torch.device(DEVICE): decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 67029b50a911..ad029160c5af 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -15,6 +15,7 @@ DeepseekV4MegaMoEExperts, DeepseekV4MoE, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs @@ -22,6 +23,7 @@ from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config +from vllm.utils.torch_utils import set_default_torch_dtype pytestmark = pytest.mark.skipif( not current_platform.is_cuda(), @@ -33,6 +35,7 @@ def v41_moe_config(dist_init): return SimpleNamespace( model_config=SimpleNamespace( + dtype=torch.bfloat16, hf_config=DeepseekV41Config( text_config=dict( hidden_size=128, @@ -63,28 +66,52 @@ def v41_moe_config(dist_init): ) +@pytest.mark.parametrize("use_cudagraph", [False, True]) +@pytest.mark.parametrize("above_threshold", [False, True]) @pytest.mark.parametrize("vision", [False, True]) -@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)]) +@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 384, 6), (2, 128, 3)]) def test_deepseek_v41_moe_routes_without_hash_table( - v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k + v41_moe_config, + monkeypatch, + vision, + layer_id, + num_experts, + top_k, + above_threshold, + use_cudagraph, ): """Main and draft layers select experts by score and preserve image routing.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM Mega Gate requires SM100") + config = v41_moe_config.model_config.hf_config + config.hidden_size = 256 + config.n_routed_experts = 384 + config.num_experts_per_tok = 6 + config.dspark_n_routed_experts = 128 + num_tokens = (1 if num_experts == 128 else 16) + int(above_threshold) config.vision_n_layers = int(vision) - with torch.device("cuda"): + with ( + set_default_torch_dtype(v41_moe_config.model_config.dtype), + torch.device("cuda"), + ): moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") - hidden_states = torch.randn(4, config.hidden_size) + hidden_states = torch.randn( + num_tokens, config.hidden_size, dtype=torch.bfloat16 + ) input_ids = ( - torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]) + torch.tensor( + [42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257] + ).repeat((num_tokens + 3) // 4)[:num_tokens] if vision else None ) assert moe.gate.tid2eid is None assert moe.gate.weight.shape == (num_experts, config.hidden_size) assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) - assert moe.experts.w13_weight.shape == (num_experts, 256, 64) + assert moe.experts.w13_weight.shape == (num_experts, 256, 128) assert moe.experts.top_k == top_k - assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2) + assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6) with torch.no_grad(): moe.gate.weight.normal_(std=0.01) @@ -93,7 +120,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias) scores = torch.nn.functional.softplus( - torch.nn.functional.linear(hidden_states, moe.gate.weight) + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) ).sqrt() bias = moe.gate.e_score_correction_bias if vision: @@ -105,15 +132,49 @@ def test_deepseek_v41_moe_routes_without_hash_table( dim=-1, keepdim=True ) + routed = {} + def check_routing(x, weights, ids, *, activation_clamp): - torch.testing.assert_close(ids, expected_ids) - torch.testing.assert_close(weights, expected_weights) + routed["ids"], routed["weights"] = ids, weights assert activation_clamp == config.swiglu_limit return x.clone() monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) - torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states) + routing_input_ids = ( + input_ids + if input_ids is not None + else torch.zeros(num_tokens, dtype=torch.int64, device="cuda") + ) + metadata = prepare_mega_gate_routing_metadata( + routing_input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID if vision else None, + ) + if use_cudagraph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + moe(hidden_states, input_ids, metadata) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = moe(hidden_states, input_ids, metadata) + graph.replay() + torch.accelerator.synchronize() + else: + output = moe(hidden_states, input_ids, metadata) + + ids, weights = routed["ids"], routed["weights"] + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close( + weights, + expected_weights, + rtol=1e-3, + atol=1e-4, + ) + torch.testing.assert_close(output, 3 * hidden_states) @pytest.mark.parametrize( @@ -156,6 +217,67 @@ def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config): moe(torch.zeros(1, 128)) +@pytest.mark.parametrize("num_tokens", [16, 17]) +def test_deepseek_v4_mega_gate_hash_routing_correctness( + v41_moe_config, monkeypatch, num_tokens +): + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM Mega Gate requires SM100") + + config = v41_moe_config.model_config.hf_config + config.hidden_size = 256 + config.num_hash_layers = 1 + config.vocab_size = 32 + with ( + set_default_torch_dtype(v41_moe_config.model_config.dtype), + torch.device("cuda"), + ): + moe = DeepseekV4MoE( + v41_moe_config, + prefix="model.layers.0.ffn", + num_hash_layers=config.num_hash_layers, + ) + hidden_states = torch.randn( + num_tokens, config.hidden_size, dtype=torch.bfloat16 + ) + input_ids = torch.arange(num_tokens) + + token_ids = torch.arange(config.vocab_size, device="cuda") + fixed_ids = torch.stack( + ( + token_ids % config.n_routed_experts, + (token_ids + 3) % config.n_routed_experts, + ), + dim=1, + ) + with torch.no_grad(): + moe.gate.weight.normal_(std=0.01) + moe.gate.tid2eid.copy_(fixed_ids) + + expected_ids = fixed_ids[input_ids] + scores = torch.nn.functional.softplus( + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) + ).sqrt() + expected_weights = scores.gather(1, expected_ids) + expected_weights *= config.routed_scaling_factor / expected_weights.sum( + dim=-1, keepdim=True + ) + + def check_routing(x, weights, ids, *, activation_clamp): + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4) + return x.clone() + + monkeypatch.setattr(moe.experts, "forward", check_routing) + monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like) + metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=True, + image_sentinel_base_id=None, + ) + torch.testing.assert_close(moe(hidden_states, input_ids, metadata), hidden_states) + + def test_deepseek_v4_mega_moe_expert_mapping(): mapping = make_deepseek_v4_expert_params_mapping(2) @@ -513,6 +635,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( monkeypatch, fused ): class FakeGate(torch.nn.Module): + weight = torch.empty(2, 128) tid2eid = None e_score_correction_bias = None @@ -760,7 +883,8 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", ) -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(): +@pytest.mark.parametrize("nonfinite_padding", [False, True]) +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_padding): from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 device = torch.device("cuda") @@ -796,6 +920,8 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(): [False, True, False, False, True, False, True], device=device, ) + if nonfinite_padding: + topk_weights[is_padding] = float("nan") ref_x, ref_x_sf = per_token_cast_to_fp8( hidden_states, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 52783889d6ca..987909702cf4 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -17,6 +17,7 @@ import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -51,12 +52,14 @@ sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, DeepseekV4Model, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -96,6 +99,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 @@ -235,6 +239,15 @@ def forward( ) inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) @@ -247,6 +260,7 @@ def forward( post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 7fef691bd134..59156365c519 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -98,6 +98,37 @@ logger = init_logger(__name__) +class MegaGateRoutingMetadata(typing.NamedTuple): + safe_hash_input_ids: torch.Tensor | None + image_token_mask: torch.Tensor | None + hash_token_mask: torch.Tensor | None + + +def prepare_mega_gate_routing_metadata( + input_ids: torch.Tensor, + *, + has_hash_routing: bool, + image_sentinel_base_id: int | None, +) -> MegaGateRoutingMetadata: + if image_sentinel_base_id is None: + return MegaGateRoutingMetadata( + input_ids if has_hash_routing else None, + None, + torch.ones_like(input_ids, dtype=torch.bool) if has_hash_routing else None, + ) + + image_token_mask = (input_ids >= image_sentinel_base_id) & ( + input_ids < image_sentinel_base_id + 5 + ) + if not has_hash_routing: + return MegaGateRoutingMetadata(None, image_token_mask, None) + return MegaGateRoutingMetadata( + torch.where(image_token_mask, 0, input_ids), + image_token_mask, + ~image_token_mask, + ) + + class DeepseekV4MLP(nn.Module): def __init__( self, @@ -803,6 +834,10 @@ def __init__( raise NotImplementedError( "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." ) + if self.use_mega_moe and not self.renormalize: + raise NotImplementedError( + "DeepSeek V4 MegaMoE requires normalized top-k probabilities." + ) if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": raise NotImplementedError( "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" @@ -1000,7 +1035,10 @@ def _init_fused_moe_experts( ) def forward( - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None = None, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") @@ -1013,23 +1051,61 @@ def forward( return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape - router_logits, _ = self.gate(hidden_states) - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=router_logits, - scoring_func=self.scoring_func, - e_score_correction_bias=self.gate.e_score_correction_bias.data - if self.gate.e_score_correction_bias is not None - else None, - topk=self.n_activated_experts, - renormalize=self.renormalize, - indices_type=self.hash_indices_dtype, - input_tokens=input_ids, - hash_indices_table=self.gate.tid2eid, - routed_scaling_factor=self.routed_scaling_factor, - bias_vl=bias_vl.data if bias_vl is not None else None, - image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, - ) + # Small local padded batches favor GateLinear; 128-expert gates cross earlier. + gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 + if hidden_states.shape[0] <= gate_threshold: + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + bias_vl=bias_vl.data if bias_vl is not None else None, + image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, + ) + else: + from vllm.utils.deep_gemm import bf16_mega_gate + + if mega_gate_metadata is None: + raise RuntimeError( + "MegaMoE routing metadata must be prepared once by the model." + ) + + unmapped_topk_idx = None + fix_routing_mask = None + if self.gate.tid2eid is not None: + if mega_gate_metadata.safe_hash_input_ids is None: + raise RuntimeError("Hash routing requires prepared input IDs.") + if mega_gate_metadata.hash_token_mask is None: + raise RuntimeError("Hash routing requires a prepared routing mask.") + unmapped_topk_idx = self.gate.tid2eid[ + mega_gate_metadata.safe_hash_input_ids + ] + fix_routing_mask = mega_gate_metadata.hash_token_mask + + topk_weights, topk_ids = bf16_mega_gate( + hidden_states, + self.gate.weight, + self.n_activated_experts, + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + ep_rank=self.ep_rank, + bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + image_bias=bias_vl.data if bias_vl is not None else None, + image_token_mask=mega_gate_metadata.image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + ) activation_clamp = ( float(self.swiglu_limit) if self.swiglu_limit is not None else None ) @@ -1248,6 +1324,7 @@ def forward( post_mix: torch.Tensor | None = None, res_mix: torch.Tensor | None = None, residual: torch.Tensor | None = None, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: attn_norm_weight = self.attn_norm.weight.data attn_norm_eps = self.attn_norm.variance_epsilon @@ -1332,7 +1409,7 @@ def forward( norm_eps=ffn_norm_eps, ) - x = self.ffn(x, input_ids) + x = self.ffn(x, input_ids, mega_gate_metadata) return x, residual, post_mix, res_mix @@ -1394,6 +1471,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + self.has_local_hash_moe = self.start_layer < min( + self.end_layer, config.num_hash_layers + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) @@ -1496,6 +1576,16 @@ def forward( hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=self.has_local_hash_moe, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) + residual, post_mix, res_mix = None, None, None remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) aux_hidden_states: list[torch.Tensor] = [] @@ -1511,6 +1601,7 @@ def forward( post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) if idx + 1 in self.aux_hidden_state_layers: # Reconstruct the aux hidden state for draft models diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index a19cc707b254..50886b0afab7 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -20,6 +20,7 @@ import vllm.envs as envs from vllm.config import VllmConfig +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -51,6 +52,7 @@ sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import ( _FUSED_MTP_INPUT_RMSNORM_KERNEL, _MTP_SHARED_HEAD_RMSNORM_KERNEL, @@ -62,6 +64,7 @@ DeepseekV4Model, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -92,6 +95,7 @@ def __init__( self.config = config quant_config = vllm_config.quant_config self.rms_norm_eps = config.rms_norm_eps + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -181,8 +185,23 @@ def forward( hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) + mega_gate_metadata = None + if self.use_mega_moe: + routing_input_ids = input_ids + if self.mtp_block.use_sequence_parallel: + routing_input_ids = sp_shard(routing_input_ids) + mega_gate_metadata = prepare_mega_gate_routing_metadata( + routing_input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=input_ids + positions=positions, + x=hidden_states, + input_ids=input_ids, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) if self.mtp_block.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 9a07189bc3bb..c08de12a12f5 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -18,6 +18,7 @@ import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -50,12 +51,14 @@ sp_shard, ) +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, DeepseekV4Model, _linear_scale_param_name, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -77,6 +80,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = ( @@ -202,6 +206,15 @@ def forward( ) inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) @@ -215,6 +228,7 @@ def forward( post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) # Collapse the hc copies with the pre-mix from the last layer's FFN diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 21fa00e5c963..6d1faace6510 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -61,7 +61,9 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v4.nvidia.model import ( + MegaGateRoutingMetadata, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( @@ -339,6 +341,7 @@ def forward( engram_mask: torch.Tensor | None = None, *, capture_previous_aux: bool = False, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -462,7 +465,7 @@ def forward( norm_weight=self.ffn_norm.weight, norm_eps=self.ffn_norm.variance_epsilon, ) - x = self.ffn(x, input_ids) + x = self.ffn(x, input_ids, mega_gate_metadata) return x, residual, post_mix, res_mix, ffn_pre, previous_aux @@ -705,6 +708,16 @@ def forward( hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) + residual, post_mix, res_mix = None, None, None pre_mix: torch.Tensor | None = None if not get_pp_group().is_first_rank: @@ -728,6 +741,7 @@ def forward( engram_hashes, engram_mask, capture_previous_aux=idx in self.aux_hidden_state_layers, + mega_gate_metadata=mega_gate_metadata, ) if previous_aux is not None: # idx is the one-based id of the layer whose post this is. diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3d96b496f1b1..2a0d4eee4aba 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -164,6 +164,7 @@ def _missing(*_: Any, **__: Any) -> NoReturn: _get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _mega_mhc_impl: Callable[..., Any] | None = None +_bf16_mega_gate_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -240,6 +241,7 @@ def _lazy_init() -> None: global _get_sparse_mqa_logits_metadata_impl global _get_paged_sparse_mqa_logits_metadata_impl global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl + global _bf16_mega_gate_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -260,6 +262,7 @@ def _lazy_init() -> None: or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None or _mega_mhc_impl is not None + or _bf16_mega_gate_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -315,6 +318,7 @@ def _lazy_init() -> None: ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) _mega_mhc_impl = getattr(_dg, "mega_mhc", None) + _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -482,6 +486,55 @@ def cublaslt_gemm_nt(*args, **kwargs): return _cublaslt_gemm_nt_impl(*args, **kwargs) +def bf16_mega_gate( + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + *, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DeepGEMM Mega Gate. + + The output order follows vLLM's router contract: weights, then indices. + Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA + graph capture. + """ + _lazy_init() + if _bf16_mega_gate_impl is None: + raise RuntimeError( + "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate." + ) + + topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device) + topk_weights = torch.empty( + (x.shape[0], num_topk), dtype=torch.float32, device=x.device + ) + _bf16_mega_gate_impl( + x, + weight, + num_topk, + use_shared_as_routed=False, + num_shared_experts=0, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + scoring_func=scoring_func, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + out=(topk_idx, topk_weights), + ) + return topk_weights, topk_idx + + def fp8_gemm_nt(*args, **kwargs): _lazy_init() if _fp8_gemm_nt_impl is None: @@ -988,6 +1041,7 @@ def should_use_deepgemm_for_fp8_linear( __all__ = [ "calc_diff", + "bf16_mega_gate", "DeepGemmQuantScaleFMT", "fp8_gemm_nt", "fp8_einsum", From 67e5b0acc9988afc50d019db64ad9da0dadaa15e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 17 Sep 2026 21:26:10 +0200 Subject: [PATCH 429/774] [Bugfix][HiSparse][NIXL] Import full blocks without tail prefill on D (#57049) Signed-off-by: NickLucche Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson --- tests/v1/core/test_prefix_caching.py | 104 ++++++++++++++++-- tests/v1/worker/test_utils.py | 65 +++++++++++ .../kv_connector/v1/hisparse/worker.py | 55 +++++++-- .../kv_connector/v1/nixl/base_worker.py | 3 +- .../kv_connector/v1/nixl/pull_worker.py | 10 +- vllm/v1/core/single_type_kv_cache_manager.py | 20 ++-- vllm/v1/hisparse/coordinator.py | 68 ++++++++---- vllm/v1/hisparse/types.py | 5 +- 8 files changed, 277 insertions(+), 53 deletions(-) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 01bf73b6e7fa..229f9c17f9ee 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -325,7 +325,7 @@ def test_hisparse_reports_when_context_is_fully_resident(): def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): - """Keep indexer-only imports host-backed in the resident block table.""" + """Restore the final host page when completing an indexer-only import.""" manager = make_hisparse_kv_cache_manager( 32, 16, @@ -360,11 +360,15 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks assert len(source) == len(indexer) == len(resident) == 4 assert len(hot) == 2 - # The local prefix is adopted from shadow pages (GPU-resident), while the - # externally imported page stays host-backed until its tail allocation. assert not any(block.is_null for block in resident[:2]) - assert resident[2].is_null + assert not resident[2].is_null assert not resident[3].is_null + coordinator = get_hisparse_coordinator(manager) + assert not coordinator.build_offload_command().page_transfers + coordinator.finish_host_import(resumed.request_id, failed=False) + transfers = coordinator.build_offload_command().page_transfers + assert len(transfers) == 1 and transfers[0].restore + assert transfers[0].resident_block_ids == (resident[2].block_id,) def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix(): @@ -465,7 +469,8 @@ def lookup(request, num_computed_tokens, max_num_new_tokens=None): ) allocated = manager.get_blocks(resumed.request_id).blocks assert len(allocated[0]) == host_blocks - assert allocated[2][1].is_null == (external > 0) + assert allocated[2][1].is_null == (host_blocks > 2) + assert not allocated[2][-1].is_null manager.free(resumed) # A surviving indexer offload must not make absent host KV look computed. @@ -490,6 +495,85 @@ def allocate_external_prefix( ) +@pytest.mark.parametrize("num_tokens", [1, 15, 16, 17, 31, 32, 33, 127]) +def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens): + """Import every real token and restore the page used by last-token replay.""" + from tests.v1.kv_connector.unit.utils import create_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + config = create_vllm_config(block_size=HISPARSE_BLOCK_SIZE) + connector = NixlPullConnectorScheduler(config, "test", manager.kv_cache_config) + request = make_request( + "import", list(range(num_tokens)), HISPARSE_BLOCK_SIZE, sha256 + ) + request.kv_transfer_params = {"do_remote_prefill": True} + count, is_async = connector.get_num_new_matched_tokens(request, 0) + assert count == num_tokens and is_async + assert allocate_external_prefix(manager, request, count) is not None + source, _, resident, _ = manager.get_blocks(request.request_id).blocks + assert len(source) == len(resident) == (num_tokens + 15) // 16 + assert all(block.is_null for block in resident[:-1]) + assert not resident[-1].is_null + coordinator = get_hisparse_coordinator(manager) + assert not coordinator.build_offload_command().page_transfers + + coordinator.finish_host_import(request.request_id, failed=False) + transfers = coordinator.build_offload_command().page_transfers + assert len(transfers) == 1 + restore = transfers[0] + assert restore.restore and not restore.after_forward + assert restore.host_block_id == source[-1].block_id + assert restore.resident_block_ids == (resident[-1].block_id,) + request.num_computed_tokens = count + scheduler = SimpleNamespace( + connector=connector, + kv_cache_manager=manager, + failed_recving_kv_req_ids=set(), + finished_recving_kv_req_ids={request.request_id}, + ) + Scheduler._update_waiting_for_remote_kv(scheduler, request) + assert request.num_tokens - request.num_computed_tokens == 1 + assert manager.allocate_slots(request, num_new_tokens=1) is not None + assert coordinator.build_row_mirrors([(request.request_id, num_tokens - 1, 1)]) + # A pending restore must not expose uninitialized GPU copies to prefix hits. + assert all(resident[-1] not in copies for copies in coordinator.copies.values()) + + coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1}) + assert resident[-1].ref_cnt == 2 + assert all(resident[-1] not in copies for copies in coordinator.copies.values()) + coordinator.update_spills({}, {restore.transfer_id: 1}) + assert resident[-1].ref_cnt == 1 + assert coordinator.request_states[request.request_id].valid_pages == set( + range(num_tokens // HISPARSE_BLOCK_SIZE) + ) + if num_tokens % HISPARSE_BLOCK_SIZE: + assert source[-1].block_hash is None + else: + assert coordinator.copies[source[-1].block_hash] == (resident[-1],) + + +def test_hisparse_aborted_tail_restore_retains_both_endpoints(): + """An abort cannot recycle storage still used by a queued tail restore.""" + manager = make_hisparse_kv_cache_manager(16, 16) + request = make_request("import", list(range(17)), HISPARSE_BLOCK_SIZE, sha256) + assert allocate_external_prefix(manager, request, 17) is not None + source, _, resident, _ = manager.get_blocks(request.request_id).blocks + host_tail, gpu_tail = source[-1], resident[-1] + coordinator = get_hisparse_coordinator(manager) + coordinator.finish_host_import(request.request_id, failed=False) + restore = coordinator.build_offload_command().page_transfers[0] + manager.free(request) + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1 + coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1}) + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1 + coordinator.update_spills({}, {restore.transfer_id: 1}) + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 0 + assert not coordinator.has_pending_work() + + @pytest.mark.parametrize("enable_caching", [False, True]) def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching): """Later admissions must not remove a scheduled request's only readable tier.""" @@ -669,7 +753,7 @@ def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool(): def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages(): """An in-flight host import must reserve its unwritten resident pages.""" manager = make_hisparse_kv_cache_manager( - 11, + 12, 16, transfer_device_cache=True, ) @@ -1127,9 +1211,10 @@ def test_hisparse_external_import_uses_hard_gpu_footprint(): assert len(source) == num_prompt_blocks assert len(indexer) == num_prompt_blocks assert len(resident) == num_prompt_blocks - assert all(block.is_null for block in resident) + assert all(block.is_null for block in resident[:-1]) + assert not resident[-1].is_null assert len(hot) == 2 - assert manager.block_pool.get_num_free_blocks() == 1 + assert manager.block_pool.get_num_free_blocks() == 0 def test_hisparse_external_import_survives_capacity_retry(): @@ -1149,7 +1234,8 @@ def test_hisparse_external_import_survives_capacity_retry(): manager.free(first) assert allocate_external_prefix(manager, second, len(tokens)) is not None _, _, resident, hot = manager.get_blocks(second.request_id).blocks - assert all(block.is_null for block in resident) + assert all(block.is_null for block in resident[:-1]) + assert not resident[-1].is_null assert len(hot) == 2 diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 55cab287fd26..9a345e926c72 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -531,6 +531,71 @@ def test_hisparse_eager_mirror_records_transfer_without_page_copy(): worker._enqueue_transfers.assert_not_called() +@pytest.mark.parametrize("is_host_writer", [False, True]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hisparse_tail_restore_preserves_imported_rows( + monkeypatch, request, is_host_writer, device +): + """Restore externally pinned host rows on every rank under sync checking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required for the device copy check") + worker = _make_hisparse_worker() + worker.is_host_writer = is_host_writer + worker.kernel_block_size = 4 + worker.host_caches = tuple( + torch.arange(24, dtype=torch.float32).reshape(12, 2) + layer * 100 + for layer in range(2) + ) + if device == "cuda": + # Match production's cudaHostRegister pool, which torch.is_pinned() + # does not recognize, rather than torch's caching pinned allocator. + host_pool, registered = hisparse_runtime_module.allocate_pinned_host_pool( + sum(cache.nbytes for cache in worker.host_caches) + ) + request.addfinalizer( + lambda: hisparse_runtime_module.release_pinned_state([], [registered]) + ) + host_caches = host_pool.view(torch.float32).view(2, 12, 2).unbind() + for destination, source in zip(host_caches, worker.host_caches): + destination.copy_(source) + worker.host_caches = host_caches + worker.resident_caches = tuple( + torch.full((4, 4, 2), -1.0, device=device) for _ in range(2) + ) + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(resident_source_index=i, eager_host_mirror=True) + ) + for i in range(2) + ] + worker._record_transfer_completion = MagicMock() + stream = torch.cuda.current_stream() if device == "cuda" else MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True) + + restore_pages = worker._restore_pages + if device == "cuda": + import vllm.utils.gpu_sync_debug as gpu_sync_debug + + monkeypatch.setattr(gpu_sync_debug, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", False) + gpu_sync_debug.enable_gpu_sync_check() + restore_pages = gpu_sync_debug.with_gpu_sync_check(restore_pages) + restore_pages([transfer]) + + # Appending into unused capacity must preserve every imported prompt row. + for cache, host, block in zip( + worker.resident_caches, worker.host_caches, transfer.resident_block_ids + ): + assert torch.equal(cache[block].cpu(), host[8:12]) + cache[block, 2:] = 999 + assert torch.equal(cache[block, :2].cpu(), host[8:10]) + assert torch.all(cache[0] == -1) + worker._record_transfer_completion.assert_called_once_with( + [transfer], stream=stream + ) + + def test_hisparse_lazy_mirror_copies_transfer_pages(): worker = _make_hisparse_worker() worker.cache_handles = [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index b7ffe9ec97f9..b223616401e5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -22,6 +22,7 @@ get_tensor_model_parallel_rank, get_tp_group, ) +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import current_stream from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX @@ -414,11 +415,18 @@ def start_step( transfers = ( metadata.command.page_transfers if metadata.command is not None else [] ) + self._restore_pages([transfer for transfer in transfers if transfer.restore]) self._post_forward_transfers = [ - transfer for transfer in transfers if transfer.after_forward + transfer + for transfer in transfers + if not transfer.restore and transfer.after_forward ] self._submit_transfers( - [transfer for transfer in transfers if not transfer.after_forward] + [ + transfer + for transfer in transfers + if not transfer.restore and not transfer.after_forward + ] ) self._pending_invalid_block_ids.extend(metadata.source_block_ids) if request_state_indices is not None: @@ -708,11 +716,17 @@ def _enqueue_layer_mirror(self, layer_index: int) -> None: self._submitted_mirror_layers.update(pending_layers) def _record_transfer_completion( - self, transfers: list[SparseKVPageTransfer] + self, + transfers: list[SparseKVPageTransfer], + *, + stream: torch.Stream | None = None, ) -> None: - if not transfers or not self.is_host_writer: + if not transfers: return - stream = self.dma_stream + if stream is None: + if not self.is_host_writer: + return + stream = self.dma_stream assert stream is not None completion_event = torch.Event() completion_event.record(stream) @@ -720,6 +734,30 @@ def _record_transfer_completion( self._pending_transfer_events.append((completion_event, transfer_ids)) self._enqueued_transfer_ids.extend(transfer_ids) + def _restore_pages(self, transfers: list[SparseKVPageTransfer]) -> None: + """Restore imported tails on every rank before its forward or graph replay.""" + if not transfers: + return + for layer_index, cache in enumerate(self.cache_handles): + source_index = cache.runtime.resident_source_index + source = self.host_caches[layer_index] + destination = self.resident_caches[layer_index] + for transfer in transfers: + host_start = transfer.host_block_id * self.kernel_block_size + resident_block = transfer.resident_block_ids[source_index] + # Each import pays one H2D page copy and enqueue per layer/rank + # instead of tail prefill on D. Restore the prompt rows before + # decode appends to this page; unused rows do not extend the + # logical sequence. + # cudaHostRegister pins this pool outside PyTorch's allocator, + # so the sync checker cannot recognize it via is_pinned(). + with gpu_sync_allowed(): + destination[resident_block].copy_( + source[host_start : host_start + self.kernel_block_size], + non_blocking=True, + ) + self._record_transfer_completion(transfers, stream=current_stream()) + def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: if self.cache_handles[0].runtime.eager_host_mirror: self._record_transfer_completion(transfers) @@ -733,15 +771,12 @@ def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: descriptor_count = len(transfers) * num_layers descriptors = self._acquire_dma_descriptors(descriptor_count) destination_rows = np.fromiter( - ( - transfer.destination_block_id * self.kernel_block_size - for transfer in transfers - ), + (transfer.host_block_id * self.kernel_block_size for transfer in transfers), dtype=np.int64, count=len(transfers), ) source_blocks_by_transfer = np.asarray( - [transfer.source_block_ids for transfer in transfers], dtype=np.int64 + [transfer.resident_block_ids for transfer in transfers], dtype=np.int64 ) if source_blocks_by_transfer.ndim != 2: raise RuntimeError( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7bcdd2531bfe..60cf20303c3e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -3278,7 +3278,8 @@ def _apply_prefix_caching_by_region( """Pair an uncached decode suffix with the same prefill regions.""" assert len(decode_block_ids) == len(prefill_block_ids) if not any(decode_block_ids): - return [], prefill_block_ids + empty_regions: list[list[int]] = [[] for _ in decode_block_ids] + return empty_regions, empty_regions.copy() if num_computed_blocks is not None: assert num_remote_blocks is not None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index fea3a00ce247..203c6bcfa668 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -168,8 +168,14 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): local_block_ids = meta.local_physical_block_ids remote_region_groups = self.dst_region_group_ids[engine_id] local_region_groups = self.region_group_ids or remote_region_groups - groups_differ = local_region_groups != remote_region_groups - if groups_differ: + if not local_block_ids: + # Region expansion cannot index empty groups. Pass empty specs to + # _read_blocks so its existing cache-hit notification path runs. + read_specs = [ + ReadSpec(remote_rank=rank, local_block_ids=[], remote_block_ids=[]) + for rank in plan.all_source_ranks + ] + elif local_region_groups != remote_region_groups: if not self.use_mla or self._has_mamba: raise NotImplementedError( "Different NIXL cache-group layouts are only supported for " diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 9ad0a38ed550..835a434c1e7c 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2498,12 +2498,14 @@ def get_num_blocks_to_allocate( del num_tokens_main_model assert not new_computed_blocks if total_computed_tokens > num_local_computed_tokens: - if total_computed_tokens % self.block_size != 0: - raise ValueError( - "A host-only HiSparse import must end on a cache-block boundary." - ) - imported_pages = total_computed_tokens // self.block_size - return max(cdiv(num_tokens, self.block_size) - imported_pages, 0) + # Keep the last imported page writable, including the last-token + # replay on a full prompt hit. Only earlier pages stay host-only. + imported_pages = cdiv(total_computed_tokens, self.block_size) + existing = len(self.req_to_blocks.get(request_id, ())) + needs_tail = self.get_resident_page(request_id, imported_pages - 1) is None + return int(needs_tail) + max( + cdiv(num_tokens, self.block_size) - max(imported_pages, existing), 0 + ) existing = len(self.req_to_blocks.get(request_id, ())) host_pages = cdiv(num_local_computed_tokens, self.block_size) required = cdiv(num_tokens, self.block_size) @@ -2518,13 +2520,15 @@ def allocate_external_computed_blocks( num_local_computed_tokens: int, num_external_computed_tokens: int, ) -> None: - """Represent imported host history with null resident pages.""" + """Reserve a writable final page; earlier imported pages stay on host.""" assert num_external_computed_tokens > 0 num_tokens = num_local_computed_tokens + num_external_computed_tokens blocks = self.req_to_blocks[request_id] - tail_page = cdiv(num_tokens, self.block_size) - 1 + tail_page = (num_tokens - 1) // self.block_size if len(blocks) <= tail_page: blocks.extend([self._null_block] * (tail_page + 1 - len(blocks))) + if blocks[tail_page].is_null: + blocks[tail_page] = self.block_pool.get_new_blocks(1)[0] def add_local_computed_blocks( self, diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index ecbc0a270ac1..056e19479a33 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -73,6 +73,7 @@ class _PendingSpill: request_state: _HiSparseRequestState host_block: KVCacheBlock resident_blocks: tuple[KVCacheBlock, ...] + restore: bool = False expected_worker_completions: int = 0 worker_completions: int = 0 enqueue_applied: bool = False @@ -271,25 +272,29 @@ def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: num_computed_tokens // self.host_manager.block_size, len(host_blocks) ) for host_idx in range(state.copies_recorded_blocks, num_host_blocks): - host_block = host_blocks[host_idx] - if host_block.is_null or host_block.block_hash is None: - continue - blocks: list[KVCacheBlock] = [] - for manager in self.resident_managers: - block = manager.get_resident_page(request_id, host_idx) - if block is None: - break - blocks.append(block) - if len(blocks) != len(self.resident_managers): + if host_idx in state.pending_pages: continue - self._drop_copy(host_block.block_hash) - self.copies[host_block.block_hash] = tuple(blocks) - for block in blocks: - self._copy_by_block[block.block_id] = host_block.block_hash + self._record_copy(request_id, host_idx, host_blocks[host_idx]) state.copies_recorded_blocks = max( state.copies_recorded_blocks, num_host_blocks ) + def _record_copy( + self, request_id: str, page_idx: int, host_block: KVCacheBlock + ) -> None: + if host_block.is_null or host_block.block_hash is None: + return + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, page_idx) + if block is None: + return + blocks.append(block) + self._drop_copy(host_block.block_hash) + self.copies[host_block.block_hash] = tuple(blocks) + for block in blocks: + self._copy_by_block[block.block_id] = host_block.block_hash + def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None: blocks = self.copies.pop(host_hash, None) if blocks is None: @@ -406,9 +411,14 @@ def plan_prefix_materialization( assert self.host_manager is not None host_block_size = self.host_manager.block_size num_pages = num_computed_tokens // host_block_size + importing_pages = cdiv( + self._pending_imports.get(request_id, 0), host_block_size + ) state = self._get_request_state(request_id) budget = max(self.max_spill_pages - len(self.spills_to_send), 0) for page_idx in range(num_pages): + if page_idx < importing_pages: + continue if budget == 0: break if page_idx in state.pending_pages: @@ -420,7 +430,7 @@ def plan_prefix_materialization( for manager in self.resident_managers ): continue - if self._plan_spill(request_id, page_idx, after_forward=True): + if self._plan_page_transfer(request_id, page_idx, after_forward=True): budget -= 1 def publish_when_ready( @@ -491,17 +501,27 @@ def complete_host_import(self, request_id: str, num_computed_tokens: int) -> Non state = self._get_request_state(request_id) state.valid_pages = set(range(num_pages)) state.ready_prefix_pages = num_pages + if num_computed_tokens: + self._plan_page_transfer( + request_id, + (num_computed_tokens - 1) // block_size, + after_forward=False, + restore=True, + ) self._publish_host_blocks_if_ready(request_id) - def _plan_spill( + def _plan_page_transfer( self, request_id: str, page_idx: int, *, after_forward: bool, + restore: bool = False, ) -> bool: assert self.host_manager is not None state = self._get_request_state(request_id) + if page_idx in state.pending_pages: + return False host_blocks = self.host_manager.req_to_blocks.get(request_id) if host_blocks is None or page_idx >= len(host_blocks): return False @@ -523,9 +543,10 @@ def _plan_spill( self.next_spill_id += 1 plan = SparseKVPageTransfer( transfer_id=spill_id, - destination_block_id=host_block.block_id, - source_block_ids=tuple(block.block_id for block in blocks), + host_block_id=host_block.block_id, + resident_block_ids=tuple(block.block_id for block in blocks), after_forward=after_forward, + restore=restore, ) self.pending_spills[spill_id] = _PendingSpill( transfer_id=spill_id, @@ -533,6 +554,7 @@ def _plan_spill( request_state=state, host_block=host_block, resident_blocks=tuple(blocks), + restore=restore, ) state.pending_pages[page_idx] = spill_id self.spills_to_send.append(plan) @@ -666,7 +688,7 @@ def update_spills( self._complete_host_writes() def _apply_enqueued_spills(self) -> None: - """Drop the DMA pins once the worker has enqueued a spill.""" + """Drop the GPU pins once every worker has completed the transfer.""" for pending in self.pending_spills.values(): if ( pending.enqueue_applied @@ -692,8 +714,12 @@ def _complete_host_writes(self) -> None: request_id, page_idx = pending.page pending.request_state.pending_pages.pop(page_idx, None) if self.request_states.get(request_id) is pending.request_state: - pending.request_state.valid_pages.add(page_idx) - self._page_became_clean(request_id, pending.request_state, page_idx) + if not pending.restore: + pending.request_state.valid_pages.add(page_idx) + if page_idx in pending.request_state.valid_pages: + self._page_became_clean(request_id, pending.request_state, page_idx) + if pending.restore: + self._record_copy(request_id, page_idx, pending.host_block) completed_request_ids.add(request_id) assert self.host_manager is not None self.host_manager.block_pool.free_blocks([pending.host_block]) diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py index 8618f3cb20a4..269793ea8521 100644 --- a/vllm/v1/hisparse/types.py +++ b/vllm/v1/hisparse/types.py @@ -9,9 +9,10 @@ class SparseKVPageTransfer: """Copy one logical KV page between cache-manager and worker-owned tiers.""" transfer_id: int - destination_block_id: int - source_block_ids: tuple[int, ...] + host_block_id: int + resident_block_ids: tuple[int, ...] after_forward: bool + restore: bool = False @dataclass(frozen=True) From ac2f0ea82c0d3005eb1de081cda715f0ee524c16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Thu, 17 Sep 2026 21:36:19 +0200 Subject: [PATCH 430/774] [MoE][Bugfix] Skip SP padded rows in grouped MoE routing (#56079) Signed-off-by: Elvir Crncevic Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Tyler Michael Smith --- tests/kernels/moe/test_grouped_topk.py | 32 +++++++ tests/kernels/moe/test_routing.py | 15 ++++ tests/v1/worker/test_gpu_model_runner.py | 14 +++ tests/v1/worker/test_gpu_model_runner_v2.py | 14 +++ tests/v1/worker/test_gpu_ubatch_slicing.py | 90 +++++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 3 + .../fused_moe/router/grouped_topk_router.py | 14 +++ .../layers/fused_moe/router/router_factory.py | 3 + vllm/models/kimi_k3/nvidia/model.py | 1 + vllm/v1/worker/gpu/model_runner.py | 19 ++-- vllm/v1/worker/gpu_model_runner.py | 16 ++++ vllm/v1/worker/gpu_ubatch_wrapper.py | 9 ++ 12 files changed, 224 insertions(+), 6 deletions(-) diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index c65cd381db86..2004feae6595 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -15,8 +15,10 @@ get_cached_compilation_config, set_current_vllm_config, ) +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( GroupedTopk, + GroupedTopKRouter, fused_grouped_topk, ) from vllm.platforms import current_platform @@ -70,6 +72,36 @@ def _single_group_reference( return values, indices.to(torch.int32) +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize("skip_padding", [False, True]) +def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool): + hidden_states = torch.randn((2, 16), device="cuda") + gating_output = torch.randn((2, 896), device="cuda") + is_padding = torch.tensor([False, True], device="cuda") + router = GroupedTopKRouter( + top_k=16, + global_num_experts=896, + num_expert_group=1, + topk_group=1, + skip_padding=skip_padding, + ) + + with set_forward_context(None, VllmConfig(), is_padding=is_padding): + topk_weights, topk_ids = router._compute_routing( + hidden_states, gating_output, torch.int32 + ) + + assert torch.all(topk_ids[0] >= 0) + if skip_padding: + assert torch.all(topk_ids[1] == -1) + assert torch.all(topk_weights[1] == 0) + else: + assert torch.all(topk_ids[1] >= 0) + assert torch.all(topk_weights[1] > 0) + + @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 174950e06c25..4baad24cc3cd 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -87,6 +87,21 @@ def test_multiple_expert_groups_use_grouped_topk() -> None: ) assert isinstance(router, GroupedTopKRouter) + assert not router.skip_padding + + +def test_grouped_topk_padding_skip_must_be_enabled() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=8, + topk_group=4, + skip_padding=True, + ) + + assert isinstance(router, GroupedTopKRouter) + assert router.skip_padding def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None: diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 80e8318fa198..32341ce944db 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -174,6 +174,20 @@ def test_freeze_gc_disables_and_restores_automatic_gc( gc.disable() +def test_prepare_padding_mask_marks_sequence_parallel_padding(): + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.is_padding = torch.empty(8, dtype=torch.bool) + + mask = runner._prepare_padding_mask(1, 8) + + assert mask.tolist() == [False, True, True, True, True, True, True, True] + assert mask.data_ptr() == runner.is_padding.data_ptr() + + mask = runner._prepare_padding_mask(0, 8) + + assert mask.all() + + @pytest.fixture def model_runner(): vllm_config = get_vllm_config() diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 308c9fe9eaab..83cd1e9d07b6 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -21,6 +21,20 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner +def test_prepare_padding_mask_marks_sequence_parallel_padding(): + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.input_buffers = SimpleNamespace(is_padding=torch.empty(8, dtype=torch.bool)) + + mask = runner._prepare_padding_mask(1, 8) + + assert mask.tolist() == [False, True, True, True, True, True, True, True] + assert mask.data_ptr() == runner.input_buffers.is_padding.data_ptr() + + mask = runner._prepare_padding_mask(0, 8) + + assert mask.all() + + def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): runner = GPUModelRunner.__new__(GPUModelRunner) runner.max_model_len = 262144 diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index e735953d8b15..9d3a92e623ec 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -10,6 +10,7 @@ """ import threading +from contextlib import nullcontext from dataclasses import replace from types import SimpleNamespace from typing import Any, cast @@ -30,6 +31,7 @@ from vllm.forward_context import create_forward_context from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker import gpu_ubatch_wrapper as legacy_gpu_ubatch_wrapper from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import cudagraph_utils, dp_utils from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -426,6 +428,94 @@ def test_slice_model_inputs_handles_mrope_positions(): assert sliced["intermediate_tensors"] is None +@pytest.mark.parametrize( + ("cudagraph_runtime_mode", "execution_method"), + [ + (CUDAGraphMode.FULL, "_capture_ubatches"), + (CUDAGraphMode.NONE, "_run_ubatches"), + ], +) +def test_legacy_ubatch_contexts_slice_padding_mask( + cudagraph_runtime_mode: CUDAGraphMode, execution_method: str +): + """Replacement contexts preserve each microbatch's padding rows.""" + ubatch_slices = [ + UBatchSlice(slice(0, 2), slice(0, 3)), + UBatchSlice(slice(2, 4), slice(3, 6)), + ] + is_padding = torch.tensor([False, True, False, True, False, True]) + forward_context = SimpleNamespace( + batch_descriptor=object(), + ubatch_slices=ubatch_slices, + cudagraph_runtime_mode=cudagraph_runtime_mode, + attn_metadata=[{}, {}], + slot_mapping={}, + dp_metadata=object(), + is_padding=is_padding, + ) + + wrapper = object.__new__(legacy_gpu_ubatch_wrapper.UBatchWrapper) + wrapper.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(data_parallel_size=2) + ) + wrapper.comm_stream = object() + wrapper.ready_barrier = object() + wrapper.cudagraphs = {} + wrapper.sm_control = nullcontext() + wrapper.runnable = object() + + def make_contexts(*, forward_contexts, **kwargs): + return [ + SimpleNamespace(id=i, forward_context=context) + for i, context in enumerate(forward_contexts) + ] + + def create_context(attn_metadata, vllm_config, **kwargs): + return SimpleNamespace(attn_metadata=attn_metadata, **kwargs) + + model_inputs = { + "input_ids": torch.arange(6), + "positions": torch.arange(6), + "inputs_embeds": None, + "intermediate_tensors": None, + } + expected = object() + with ( + patch.object( + legacy_gpu_ubatch_wrapper, + "get_forward_context", + return_value=forward_context, + ), + patch.object( + legacy_gpu_ubatch_wrapper.torch.cuda, + "current_stream", + return_value=object(), + ), + patch.object( + legacy_gpu_ubatch_wrapper.DPMetadata, "make", return_value=object() + ), + patch.object( + legacy_gpu_ubatch_wrapper, + "create_forward_context", + side_effect=create_context, + ), + patch.object( + legacy_gpu_ubatch_wrapper, + "make_ubatch_contexts", + side_effect=make_contexts, + ), + patch.object(wrapper, execution_method, return_value=expected) as execute, + ): + result = wrapper(**model_inputs) + + assert result is expected + contexts = [ + metadata.context.forward_context for metadata in execute.call_args.args[0] + ] + torch.testing.assert_close(contexts[0].is_padding, is_padding[:3]) + torch.testing.assert_close(contexts[1].is_padding, is_padding[3:]) + + def _make_dbo_config() -> VllmConfig: return VllmConfig( model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 2ee3c5f374a3..49557bb3f9a0 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -137,6 +137,7 @@ def FusedMoEFactory( runner_args: dict[str, Any] | None = None, routed_experts_cls: type[RoutedExperts] | None = None, routed_experts_args: dict[str, Any] | None = None, + skip_padding: bool = False, ) -> MoERunner: """Factory function for creating MoE execution pipeline. @@ -213,6 +214,7 @@ def FusedMoEFactory( runner_args: Additional arguments for runner constructor routed_experts_cls: Custom RoutedExperts class (None = use default) routed_experts_args: Additional arguments for routed_experts constructor + skip_padding: Whether grouped routing should invalidate padding rows. Returns: MoERunner: Configured MoE execution pipeline ready for forward passes @@ -308,6 +310,7 @@ def FusedMoEFactory( else 1.0, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=num_fused_shared_experts, + skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels), # Fused shared-expert slot weight. With apply_routed_scale_to_output # the runner scales the combined output by routed_scaling_factor, so # the shared slot weight must be 1/routed_scaling_factor for its net diff --git a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py index 2df3489358b1..9b134d385a63 100644 --- a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py @@ -8,6 +8,7 @@ from vllm import envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, @@ -258,6 +259,7 @@ def __init__( e_score_correction_bias: torch.Tensor | None = None, num_fused_shared_experts: int = 0, eplb_state: EplbLayerState | None = None, + skip_padding: bool = False, ): super().__init__( top_k=top_k, @@ -271,6 +273,7 @@ def __init__( self.routed_scaling_factor = routed_scaling_factor self.e_score_correction_bias = e_score_correction_bias self.num_fused_shared_experts = num_fused_shared_experts + self.skip_padding = skip_padding @property def routing_method_type(self) -> RoutingMethodType: @@ -346,4 +349,15 @@ def valid_grouping() -> bool: e_score_correction_bias=self.e_score_correction_bias, ) + if ( + self.skip_padding + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[: topk_ids.shape[0]].unsqueeze(1) + topk_weights = topk_weights.masked_fill(is_padding, 0) + topk_ids = topk_ids.masked_fill(is_padding, -1) + return topk_weights, topk_ids diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index 203725182491..f03fb106f795 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -65,6 +65,7 @@ def create_fused_moe_router( # Deepseek V4 vision routing bias parameters bias_vl: torch.Tensor | None = None, image_sentinel_lo: int = 0, + skip_padding: bool = False, ) -> FusedMoERouter: """Factory function to create the appropriate FusedMoERouter subclass based on the provided parameters. @@ -91,6 +92,7 @@ def create_fused_moe_router( topk_group: Top-k within each group (for grouped routing) scoring_func: Scoring function to use ("softmax" or "sigmoid") num_fused_shared_experts: Number of fused shared experts (for ROCm AITER) + skip_padding: Whether to invalidate routes for padding rows Grouped topk and fused topk bias arguments: routed_scaling_factor: Scaling factor for routed weights @@ -209,6 +211,7 @@ def advertised_routing_method(groups: int | None) -> RoutingMethodType: routed_scaling_factor=routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=num_fused_shared_experts, + skip_padding=skip_padding, ) # Otherwise fall through to the non-grouped chain below. diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 7fe612b65466..a970479da81d 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -761,6 +761,7 @@ def __init__( routed_input_transform=None, routed_output_transform=self.routed_output_transform, is_sequence_parallel=use_sequence_parallel, + skip_padding=True, runner_cls=LatentMoERunner if self.use_latent_moe else None, ) if self.padded_moe_intermediate_size != moe_intermediate_size: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 631e6b04d426..903e5abd6771 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1258,6 +1258,14 @@ def gather_batch_req_state( num_reqs, num_toks, max_query_len, batch_state.has_prefill ) + def _prepare_padding_mask( + self, num_tokens: int, num_tokens_after_padding: int + ) -> torch.Tensor: + is_padding = self.input_buffers.is_padding[:num_tokens_after_padding] + is_padding[:num_tokens].fill_(False) + is_padding[num_tokens:].fill_(True) + return is_padding + def prepare_inputs( self, scheduler_output: SchedulerOutput, @@ -1267,12 +1275,11 @@ def prepare_inputs( num_tokens = batch_req_state.num_tokens num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) assert num_tokens > 0 + is_padding = self.input_buffers.is_padding[:num_tokens_after_padding] if envs.VLLM_MOE_SKIP_PADDING: - # Mark trailing cudagraph-padding rows so kernels can skip work for - # them when supported. - is_padding = self.input_buffers.is_padding - is_padding[:num_tokens].fill_(False) - is_padding[num_tokens:num_tokens_after_padding].fill_(True) + is_padding = self._prepare_padding_mask( + num_tokens, num_tokens_after_padding + ) req_ids = batch_req_state.req_ids num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens @@ -1440,7 +1447,7 @@ def prepare_inputs( has_prefill=batch_req_state.has_prefill, input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], positions=self.input_buffers.positions[:num_tokens_after_padding], - is_padding=self.input_buffers.is_padding[:num_tokens_after_padding], + is_padding=is_padding, logits_indices=logits_indices, cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 77d11621c2cb..c943f02212bb 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -816,6 +816,9 @@ def __init__( self.positions = torch.zeros( self.max_num_tokens, dtype=torch.int64, device=self.device ) + self.is_padding = torch.zeros( + self.max_num_tokens, dtype=torch.bool, device=self.device + ) self.query_start_loc = self._make_buffer( self.max_num_reqs + 1, dtype=torch.int32 ) @@ -3500,6 +3503,14 @@ def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: return round_up(num_scheduled_tokens, tp_size) return num_scheduled_tokens + def _prepare_padding_mask( + self, num_tokens_unpadded: int, num_tokens_padded: int + ) -> torch.Tensor: + padding_mask = self.is_padding[:num_tokens_padded] + padding_mask[:num_tokens_unpadded].fill_(False) + padding_mask[num_tokens_unpadded:].fill_(True) + return padding_mask + def _prepare_mm_inputs( self, num_tokens: int ) -> tuple[torch.Tensor | None, torch.Tensor]: @@ -4441,6 +4452,7 @@ def execute_model( num_tokens_unpadded, ubatch_slices_padded, ) + is_padding = self._prepare_padding_mask(num_tokens_unpadded, num_tokens_padded) with ( set_forward_context( attn_metadata, @@ -4452,6 +4464,7 @@ def execute_model( ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, skip_compiled=has_encoder_input, + is_padding=is_padding, ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output( @@ -6159,6 +6172,8 @@ def _dummy_run( if num_tokens_across_dp is not None: num_tokens_across_dp[:] = num_tokens_padded + is_padding = self._prepare_padding_mask(0, num_tokens_padded) + with ( self.maybe_randomize_inputs( input_ids, inputs_embeds, randomize_inputs=randomize_inputs @@ -6172,6 +6187,7 @@ def _dummy_run( batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, + is_padding=is_padding, ), ): outputs = self.model( diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index a0f2643052ae..f2cab7898be5 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -261,6 +261,7 @@ def _make_ubatch_metadata( dp_metadata, batch_descriptor, cudagraph_runtime_mode, + is_padding, ) -> list[UbatchMetadata]: # Create one forward context per ubatch forward_contexts = [] @@ -276,6 +277,11 @@ def _make_ubatch_metadata( batch_descriptor=batch_descriptor, cudagraph_runtime_mode=cudagraph_runtime_mode, slot_mapping=slot_mapping[i] if has_slot_mapping else None, + is_padding=( + is_padding[ubatch_slice.token_slice] + if is_padding is not None + else None + ), ) ) @@ -373,6 +379,7 @@ def __call__(self, *args, **kwargs): attn_metadata = forward_context.attn_metadata slot_mapping = forward_context.slot_mapping + is_padding = forward_context.is_padding num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) input_ids = kwargs["input_ids"] positions = kwargs["positions"] @@ -414,6 +421,7 @@ def __call__(self, *args, **kwargs): dp_metadata=ubatch_dp_metadata, batch_descriptor=batch_descriptor, cudagraph_runtime_mode=CUDAGraphMode.NONE, + is_padding=is_padding, ) with self.sm_control: return self._capture_ubatches(ubatch_metadata, self.runnable) @@ -440,6 +448,7 @@ def __call__(self, *args, **kwargs): dp_metadata=ubatch_dp_metadata, batch_descriptor=batch_descriptor, cudagraph_runtime_mode=CUDAGraphMode.NONE, + is_padding=is_padding, ) with self.sm_control: return self._run_ubatches(ubatch_metadata, self.runnable) From 9a5bd373cfe59e236f2b32c30530f4dbbf24120d Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 17 Sep 2026 16:04:03 -0400 Subject: [PATCH 431/774] [CI] Retire Weight Loading smoke tests (#57398) Signed-off-by: mgoin --- .buildkite/test-amd.yaml | 60 ------------------- .buildkite/test_areas/weight_loading.yaml | 38 ------------ .github/CODEOWNERS | 1 - .../test_rocm_compressed_tensors_w4a16.py | 1 - tests/weight_loading/models-amd.txt | 18 ------ tests/weight_loading/models-large-amd.txt | 4 -- tests/weight_loading/models-large.txt | 6 -- tests/weight_loading/models.txt | 18 ------ .../run_model_weight_loading_test.sh | 58 ------------------ tests/weight_loading/test_weight_loading.py | 49 --------------- 10 files changed, 253 deletions(-) delete mode 100644 .buildkite/test_areas/weight_loading.yaml delete mode 100644 tests/weight_loading/models-amd.txt delete mode 100644 tests/weight_loading/models-large-amd.txt delete mode 100644 tests/weight_loading/models-large.txt delete mode 100644 tests/weight_loading/models.txt delete mode 100755 tests/weight_loading/run_model_weight_loading_test.sh delete mode 100644 tests/weight_loading/test_weight_loading.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 2f5a84aaf1ea..ba8f382755e9 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3335,36 +3335,6 @@ steps: commands: - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py -#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------# - -- label: ":amd: (MI300) Weight Loading Multi-GPU" # TBD - timeout_in_minutes: 70 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - optional: true - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - -- label: ":amd: (MI300) Weight Loading Multi-GPU (Large Models)" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt - ######################################################################################################################################### # # # MI355 (gfx950) tests # @@ -4719,36 +4689,6 @@ steps: - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi -#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# - -- label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD - timeout_in_minutes: 70 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - optional: true - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - -- label: ":amd: (MI355) Weight Loading Multi-GPU (Large Models)" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - working_dir: "/vllm-workspace/tests" - num_gpus: 2 - optional: true - source_file_dependencies: - - vllm/ - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt - #----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# - label: ":amd: (MI355) Regression" # TBD diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml deleted file mode 100644 index b8afa261e592..000000000000 --- a/.buildkite/test_areas/weight_loading.yaml +++ /dev/null @@ -1,38 +0,0 @@ -group: Weight Loading -depends_on: - - image-build -steps: -- label: ":nvidia: (L4) Weight Loading Multi-GPU" - key: weight-loading-multiple-gpu - timeout_in_minutes: 50 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - optional: true - source_file_dependencies: - - vllm/ - - "!vllm/distributed/kv_transfer/" - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt - mirror: - amd: - label: ":amd: (MI355) Weight Loading Multi-GPU" - dind: false - device: mi355_2 - timeout_in_minutes: 70 - depends_on: - - image-build-amd - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - -# - label: Weight Loading Multiple GPU - Large Models # optional -# working_dir: "/vllm-workspace/tests" -# num_devices: 2 -# device: a100 -# optional: true -# source_file_dependencies: -# - vllm/ -# - tests/weight_loading -# commands: -# - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index aa5b2a9fc9dd..f64cf5f56107 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -110,7 +110,6 @@ /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera /tests/v1/structured_output @mgoin @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium -/tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee /tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 /tests/v1/kv_connector/nixl_integration @NickLucche diff --git a/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py index a9b35a4ea642..8e4563ab8004 100644 --- a/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py +++ b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py @@ -17,7 +17,6 @@ @pytest.mark.parametrize( "model_path", [ - # Listed in tests/weight_loading/models.txt "nm-testing/tinyllama-oneshot-w4a16-group128-v2", ], ) diff --git a/tests/weight_loading/models-amd.txt b/tests/weight_loading/models-amd.txt deleted file mode 100644 index a6a85736cb16..000000000000 --- a/tests/weight_loading/models-amd.txt +++ /dev/null @@ -1,18 +0,0 @@ -# quantization, model, revision[, minimum capability[, safetensors load strategy]] -fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main -None, amd/Llama-3.2-1B-Instruct-FP8-KV, main -fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main, 80, prefetch -gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main -gptq, TheBloke/Llama-2-7B-GPTQ, main -gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True -compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main -compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main -compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main -compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main -awq, casperhansen/mixtral-instruct-awq, main -fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main diff --git a/tests/weight_loading/models-large-amd.txt b/tests/weight_loading/models-large-amd.txt deleted file mode 100644 index 41e8ee5200af..000000000000 --- a/tests/weight_loading/models-large-amd.txt +++ /dev/null @@ -1,4 +0,0 @@ -# quantization, model, revision[, minimum capability[, safetensors load strategy]] -fp8, amd/Meta-Llama-3.1-70B-Instruct-FP8-KV, main, 80, prefetch -None, microsoft/phi-4, main, 80, prefetch -fp8, amd/Mixtral-8x22B-Instruct-v0.1-FP8-KV, main, 80, prefetch diff --git a/tests/weight_loading/models-large.txt b/tests/weight_loading/models-large.txt deleted file mode 100644 index 8b3913cc2d50..000000000000 --- a/tests/weight_loading/models-large.txt +++ /dev/null @@ -1,6 +0,0 @@ -compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-quantized, main -compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-channel-quantized, main -compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W8A16-quantized, main -gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, main -awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main -compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main diff --git a/tests/weight_loading/models.txt b/tests/weight_loading/models.txt deleted file mode 100644 index 87b965087092..000000000000 --- a/tests/weight_loading/models.txt +++ /dev/null @@ -1,18 +0,0 @@ -gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main -gptq, TheBloke/Llama-2-7B-GPTQ, main -gptq, LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit, main -gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True -gptq, TechxGenus/gemma-2b-it-GPTQ, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main -compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main -compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main -compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main -compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main -#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main -awq, casperhansen/mixtral-instruct-awq, main -awq_marlin, casperhansen/mixtral-instruct-awq, main -fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main diff --git a/tests/weight_loading/run_model_weight_loading_test.sh b/tests/weight_loading/run_model_weight_loading_test.sh deleted file mode 100755 index 598f1620c7da..000000000000 --- a/tests/weight_loading/run_model_weight_loading_test.sh +++ /dev/null @@ -1,58 +0,0 @@ -#!/bin/bash -SUCCESS=0 - -while getopts "c:" OPT; do - case ${OPT} in - c ) - CONFIG="$OPTARG" - ;; - \? ) - usage - exit 1 - ;; - esac -done - - -IFS=$'\n' read -d '' -r -a MODEL_CONFIGS < "$CONFIG" - -for MODEL_CONFIG in "${MODEL_CONFIGS[@]}" -do - if [[ $MODEL_CONFIG == \#* ]]; then - echo "=== SKIPPING MODEL: $MODEL_CONFIG ===" - continue - fi - - LOCAL_SUCCESS=0 - IFS=', ' read -r -a array <<< "$MODEL_CONFIG" - - echo "=== RUNNING MODEL: $MODEL_CONFIG ===" - - unset WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY - export QUANTIZATION=${array[0]} - export MODEL_NAME=${array[1]} - export REVISION=${array[2]} - # If array length is larger than 3, then MIN_CAPABILITY is provided - if [ ${#array[@]} -gt 3 ]; then - export MIN_CAPABILITY=${array[3]} - fi - if [ ${#array[@]} -gt 4 ]; then - export WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY=${array[4]} - fi - pytest -s weight_loading/test_weight_loading.py || LOCAL_SUCCESS=$? - - if [[ $LOCAL_SUCCESS == 0 ]]; then - echo "=== PASSED MODEL: ${MODEL_CONFIG} ===" - else - echo "=== FAILED MODEL: ${MODEL_CONFIG} ===" - fi - - SUCCESS=$((SUCCESS + LOCAL_SUCCESS)) - -done - -if [ "${SUCCESS}" -eq "0" ]; then - exit 0 -else - exit 1 -fi diff --git a/tests/weight_loading/test_weight_loading.py b/tests/weight_loading/test_weight_loading.py deleted file mode 100644 index c32ce71acae2..000000000000 --- a/tests/weight_loading/test_weight_loading.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import os - -import pytest -import torch - -from vllm.platforms import current_platform -from vllm.utils.mem_constants import GiB_bytes - -MAX_MODEL_LEN = 1024 -MODEL_NAME = os.environ.get( - "MODEL_NAME", "robertgshaw2/zephyr-7b-beta-channelwise-gptq" -) -REVISION = os.environ.get("REVISION", "main") -QUANTIZATION = os.environ.get("QUANTIZATION", "gptq_marlin") -MIN_CAPABILITY = os.environ.get("MIN_CAPABILITY", "80") -SAFETENSORS_LOAD_STRATEGY = os.environ.get( - "WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY" -) - - -@pytest.mark.skipif( - MODEL_NAME == "casperhansen/deepseek-coder-v2-instruct-awq", reason="OOM in the CI" -) -@pytest.mark.skipif( - not current_platform.has_device_capability(int(MIN_CAPABILITY)), - reason="Current system does not have minimum capability.", -) -def test_weight_loading(vllm_runner): - """Test parameter weight loading with tp>1.""" - # MoE models need fp16. - NEEDS_FP16 = QUANTIZATION == "gptq" - with vllm_runner( - model_name=MODEL_NAME, - revision=REVISION, - dtype=torch.half if NEEDS_FP16 else "auto", - quantization=None if QUANTIZATION == "None" else QUANTIZATION, - max_model_len=MAX_MODEL_LEN, - tensor_parallel_size=2, - safetensors_load_strategy=SAFETENSORS_LOAD_STRATEGY, - # Generating 20 tokens needs a tiny KV cache, while sizing and - # allocating a full-size one dominates this test on large devices. - kv_cache_memory_bytes=2 * GiB_bytes, - ) as model: - output = model.generate_greedy("Hello world!", max_tokens=20) - print(output) - assert output From 2b02c6c29b72147de18905ba4e7dc8f284655b40 Mon Sep 17 00:00:00 2001 From: Doug Brown Date: Thu, 17 Sep 2026 16:09:36 -0400 Subject: [PATCH 432/774] [Model] Add Cohere2MoE Eagle3 auxiliary hidden states (#49819) Signed-off-by: Doug Brown --- tests/v1/spec_decode/test_dflash_causality.py | 12 ++++++++++ vllm/model_executor/models/cohere2_moe.py | 23 ++++++++++++++----- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py index 0f1de86cf95b..dc13e5570090 100644 --- a/tests/v1/spec_decode/test_dflash_causality.py +++ b/tests/v1/spec_decode/test_dflash_causality.py @@ -103,6 +103,18 @@ def test_dflash_fc_uses_aux_layer_count(): assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096 +def test_dflash_target_layers_map_to_post_layer_boundaries(): + vllm_config = _vllm_config( + dflash_config={"target_layer_ids": [0, 16, 31]}, + ) + + assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == ( + 1, + 17, + 32, + ) + + @pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"]) def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name): layer_ids = [1, 17, 32] diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index d657b78a5da1..98e649dec62d 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -35,7 +35,7 @@ from vllm.sequence import IntermediateTensors from .commandr import LayerNorm -from .interfaces import SupportsPP, SupportsQuant +from .interfaces import EagleModelMixin, SupportsEagle3, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, WeightsMapper, @@ -385,7 +385,7 @@ def forward( @support_torch_compile -class Cohere2MoeModel(nn.Module): +class Cohere2MoeModel(nn.Module, EagleModelMixin): """Transformer decoder for Cohere2Moe.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -439,7 +439,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -450,17 +450,28 @@ def forward( assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - for layer in islice(self.layers, self.start_layer, self.end_layer): + + aux_hidden_states = [] + if self.start_layer in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states) + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual = layer(positions, hidden_states, residual) + if layer_idx + 1 in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states) if not get_pp_group().is_last_rank: return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) hidden_states, _ = self.norm(hidden_states, residual) + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states -class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant): +class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant, SupportsEagle3): is_text_generation_model = True hf_to_vllm_mapper = WeightsMapper( @@ -520,7 +531,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) def compute_logits( From acc2ed2a5f1fbe79575a220fc743a80c2d068ce8 Mon Sep 17 00:00:00 2001 From: rasmith Date: Thu, 17 Sep 2026 15:20:35 -0500 Subject: [PATCH 433/774] [ROCm][CI][The Rock 10] Fix (MI355) Quantized Models failure on The Rock 10 with Triton 3.8.x (#54849) Signed-off-by: Randall Smith Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../kernels/linear/mxfp8/rocm_native.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 29739afae3f2..1d9789ccedcf 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -119,6 +119,18 @@ def _mxfp8_dot_scaled_linear( return out +# Triton 3.8 enables TRITON_HIP_USE_ASYNC_COPY by default on gfx950; its extra LDS +# buffer leaves no room for num_stages=3 at BLOCK_K=256. +# TODO(rasmith)(Rohan138): Remove the 3.8 check once +# https://github.com/vllm-project/vllm/pull/50605 merges. +_BK256_STAGES = 3 +if triton.__version__.startswith("3.8") and current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 + + if on_gfx950(): + _BK256_STAGES = 2 + + def _select_cfg(M, N, K): """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950. @@ -166,7 +178,7 @@ def _select_cfg(M, N, K): # 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits # "out of resources", so use 128x128x256 -- within the known-good footprint.) if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256: - return 128, 128, 256, 8, 3 + return 128, 128, 256, 8, _BK256_STAGES return 128, 128, 128, 8, 3 # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked): # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since @@ -186,7 +198,7 @@ def _select_cfg(M, N, K): # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K # / very-large-M shapes. if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))): - return 128, 128, 256, 8, 3 + return 128, 128, 256, 8, _BK256_STAGES # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there). if N <= 1024 and K % 256 == 0: From 9612f77077e09acbae9cc1d1b2ddf14a23e91597 Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Thu, 17 Sep 2026 23:23:15 +0300 Subject: [PATCH 434/774] [Bugfix] Release stale FlashMLA workspace views after growth (#56902) Signed-off-by: Maroon Ayoub Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../attention/backends/mla/flashmla_sparse.py | 45 +++++++++++-------- 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 265fb779c3cc..b3332eb86292 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -720,7 +720,9 @@ def __init__( ) self.pcp_dcp_kv_gather = False - self.gathered_kv_workspace: torch.Tensor | None = None + self.workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ + (q_concat_shape, torch.bfloat16) + ] if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None @@ -738,20 +740,14 @@ def __init__( # shards into a workspace of the full prefill size. shard_rows //= parallel_config.decode_context_parallel_size self.prefill_workspace_shape = (shard_rows, head_size) - specs = [ - (q_concat_shape, torch.bfloat16), - (self.prefill_workspace_shape, torch.bfloat16), - ] + self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16)) if self.pcp_dcp_kv_gather: - specs.append(((prefill_workspace_size, head_size), torch.bfloat16)) - buffers = current_workspace_manager().get_simultaneous(*specs) - self.q_concat_buffer, self.prefill_bf16_workspace = buffers[:2] - if self.pcp_dcp_kv_gather: - self.gathered_kv_workspace = buffers[2] - else: - (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( - (q_concat_shape, torch.bfloat16), - ) + self.workspace_specs.append( + ((prefill_workspace_size, head_size), torch.bfloat16) + ) + # Reserve capacity without retaining views that prevent old storage + # from being released when another layer grows the shared workspace. + current_workspace_manager().get_simultaneous(*self.workspace_specs) def _forward_bf16_kv( self, @@ -849,8 +845,10 @@ def _gather_prefill_chunk( """All-gather this rank's upconverted KV shard so the chunk's rows attend the whole context, and map their top-k onto the rank-major result.""" shard_rows = int(chunk.chunk_tot_seqlen) - assert self.gathered_kv_workspace is not None - gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows] + _, _, gathered_kv_workspace = current_workspace_manager().get_simultaneous( + *self.workspace_specs + ) + gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows] dist.all_gather_into_tensor( gathered_kv, shard, group=get_dcp_group().device_group ) @@ -902,10 +900,13 @@ def _forward_fp8_kv_separate_prefill_decode( assert fp8_metadata.prefill is not None first_chunk = fp8_metadata.prefill.chunks[0] assert isinstance(index_group, HiSparseMLAIndexGroup) + _, prefill_bf16_workspace, *_ = ( + current_workspace_manager().get_simultaneous(*self.workspace_specs) + ) prefill_ready = index_group.gather_fp8_prefill( self.index_group_index, kv_c_and_k_pe_cache, - self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen], + prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen], first_chunk.block_table, first_chunk.workspace_starts, len(first_chunk.block_table), @@ -1022,8 +1023,11 @@ def _fp8_decode( ) assert fp8_metadata.prefill is not None + _, prefill_bf16_workspace, *_ = ( + current_workspace_manager().get_simultaneous(*self.workspace_specs) + ) for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks): - chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen] + chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen] if uses_host_cache and chunk_index > 0: assert isinstance(index_group, HiSparseMLAIndexGroup) prefill_ready = index_group.gather_fp8_prefill( @@ -1302,7 +1306,10 @@ def forward_mqa( actual_num_heads = self.num_heads if isinstance(q, tuple): ql_nope, q_pe = q - q = self.q_concat_buffer[: ql_nope.shape[0]] + q_concat_buffer, *_ = current_workspace_manager().get_simultaneous( + *self.workspace_specs + ) + q = q_concat_buffer[: ql_nope.shape[0]] ops.concat_mla_q(ql_nope, q_pe, q) else: actual_num_heads = q.shape[1] From 668d6c3a778b9ac525f78d39b0c413c8604235d1 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:11:40 -0700 Subject: [PATCH 435/774] [CI] Shard multimodal Processor 1->4 (#56316) Signed-off-by: Thang Nguyen Co-authored-by: Thang Nguyen Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Kimi Code --- .buildkite/test_areas/models_multimodal.yaml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index e8c5ddd2cb27..a324da687f03 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -109,22 +109,24 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 18GB) Multimodal Processor" +- label: ":nvidia: (H200 MIG 18GB) Multimodal Processor Shard %N" key: multi-modal-processor timeout_in_minutes: 98 device: h200_18gb + parallelism: 4 source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" - tests/models/multimodal - tests/models/registry.py commands: - - pytest -v -s models/multimodal/processing/test_tensor_schema.py + - pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: - label: ":amd: (MI355) Multimodal Processor" + label: ":amd: (MI355) Multimodal Processor Shard %N" dind: false device: mi355_1 + parallelism: 4 timeout_in_minutes: 115 depends_on: - image-build-amd From db7a24c230a4db6ae3a568ed01b01f3f7172069d Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 17 Sep 2026 17:38:59 -0400 Subject: [PATCH 436/774] [Perf] Add fused DFlash2 grouped convolution (#55960) Signed-off-by: mgoin --- .buildkite/test_areas/spec_decode.yaml | 4 + .../acceptance_rates/dflash/test_dflash.py | 21 +++ tests/v1/spec_decode/test_dflash2.py | 50 +++++++ vllm/model_executor/models/qwen3_dflash2.py | 124 ++++++++++++++++++ 4 files changed, 199 insertions(+) diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index c06642ea1a13..fdb12efe1158 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -283,6 +283,8 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/qwen3_dflash2.py + - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/laguna_dflash.py - tests/v1/e2e/spec_decode/ commands: @@ -293,6 +295,8 @@ steps: dind: false device: mi300_1 timeout_in_minutes: 120 + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" depends_on: - image-build-amd source_file_dependencies: diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py index 2b1800f78234..48dcde2bbced 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py @@ -26,6 +26,8 @@ class DFlashCorrectnessConfig: use_chat_completions: bool = False enforce_eager: bool = False disable_flashinfer_sampler: bool = False + language_model_only: bool = False + chat_template_kwargs: dict[str, object] | None = None QWEN3_DFLASH = DFlashCorrectnessConfig( @@ -35,6 +37,18 @@ class DFlashCorrectnessConfig: expected_acceptance_len=3.5, ) +QWEN3_8_DFLASH2_NVFP4 = DFlashCorrectnessConfig( + model="nvidia/Qwen3.8-27B-NVFP4", + draft_model="z-lab/Qwen3.8-27B-DFlash2", + expected_accuracy=0.90, + expected_acceptance_len=6.14 * 0.95, + num_speculative_tokens=7, + max_num_seqs=32, + use_chat_completions=True, + language_model_only=True, + chat_template_kwargs={"enable_thinking": False}, +) + LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig( model="poolside/Laguna-XS-2.1-NVFP4", draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", @@ -104,6 +118,11 @@ def test_dflash_reference_acceptance_lengths( True, id="laguna-nvfp4-mrv2", ), + pytest.param( + QWEN3_8_DFLASH2_NVFP4, + True, + id="qwen3.8-dflash2-nvfp4-mrv2", + ), ], ) def test_dflash_correctness( @@ -134,12 +153,14 @@ def test_dflash_correctness( disable_log_stats=False, enable_chunked_prefill=None, compilation_config=CompilationConfig(), + language_model_only=config.language_model_only, ) as spec_runner: spec_llm = spec_runner.llm results = evaluate_gsm8k_offline( spec_llm, num_questions=config.num_questions, use_chat_completions=config.use_chat_completions, + chat_template_kwargs=config.chat_template_kwargs, ) accuracy = results["accuracy"] acceptance_len = compute_acceptance_len(spec_llm.get_metrics()) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index e7ff2b6d5d0b..72379fa7d7fc 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -7,6 +7,7 @@ import torch from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges +from vllm.platforms import current_platform from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator @@ -35,6 +36,55 @@ def test_grouped_conv_matches_reference(block_size: int): torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "batch,num_groups,group_size,block_size,taps", + [ + (7, 13, 11, 5, 1), + (7, 13, 11, 5, 3), + (7, 13, 11, 8, 2), + (16, 64, 16, 8, 2), + (16, 160, 16, 8, 2), + ], +) +def test_grouped_conv_triton_matches_reference( + dtype: torch.dtype, + batch: int, + num_groups: int, + group_size: int, + block_size: int, + taps: int, +): + torch.manual_seed(0) + rows = batch * block_size + hidden = torch.randn(rows, num_groups * group_size, device="cuda", dtype=dtype) + base = torch.randn(taps, num_groups * group_size, device="cuda", dtype=dtype) + projected = torch.randn(rows, 2, taps, num_groups, device="cuda", dtype=dtype) + delta = projected[:, 1] + + actual = _grouped_conv( + hidden, delta, base, block_size, num_groups, group_size, taps + ) + + hidden_blocks = hidden.float().view(batch, block_size, num_groups, group_size) + expected = torch.zeros_like(hidden_blocks) + base_blocks = base.float().view(taps, num_groups, group_size) + delta_blocks = delta.float().view(batch, block_size, taps, num_groups) + for position in range(block_size): + for tap in range(min(taps, position + 1)): + expected[:, position] += ( + base_blocks[tap] + delta_blocks[:, position, tap, :, None] + ) * hidden_blocks[:, position - tap] + + torch.testing.assert_close( + actual, + expected.flatten(0, 1).flatten(-2).to(dtype), + rtol=1e-2 if dtype is torch.bfloat16 else 1e-5, + atol=1e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py index cb027293f981..2043307d0250 100644 --- a/vllm/model_executor/models/qwen3_dflash2.py +++ b/vllm/model_executor/models/qwen3_dflash2.py @@ -11,6 +11,8 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op from .qwen3_dflash import ( DFlashQwen3DecoderLayer, @@ -20,6 +22,125 @@ from .utils import maybe_prefix +@triton.jit +def _dflash2_grouped_conv_kernel( + x_ptr, + delta_ptr, + base_ptr, + output_ptr, + x_stride_row, + delta_stride_row, + delta_stride_tap, + base_stride_tap, + output_stride_row, + NUM_CHANNELS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + GROUP_SIZE: tl.constexpr, + TAPS: tl.constexpr, + ELEMENT_BLOCK: tl.constexpr, +) -> None: + row = tl.program_id(0) // triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK) + col_block = tl.program_id(0) % triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK) + channels = col_block * ELEMENT_BLOCK + tl.arange(0, ELEMENT_BLOCK) + mask = channels < NUM_CHANNELS + groups = channels // GROUP_SIZE + position = row % BLOCK_SIZE + + delta_row = delta_ptr + row * delta_stride_row + x_row = x_ptr + row * x_stride_row + accumulator = ( + tl.load(base_ptr + channels, mask=mask, other=0.0).to(tl.float32) + + tl.load(delta_row + groups, mask=mask, other=0.0).to(tl.float32) + ) * tl.load(x_row + channels, mask=mask, other=0.0).to(tl.float32) + + for tap in tl.static_range(1, TAPS): + tap_mask = mask & (position >= tap) + coefficient = tl.load( + base_ptr + tap * base_stride_tap + channels, + mask=tap_mask, + other=0.0, + ).to(tl.float32) + tl.load( + delta_row + tap * delta_stride_tap + groups, + mask=tap_mask, + other=0.0, + ).to(tl.float32) + x = tl.load( + x_ptr + (row - tap) * x_stride_row + channels, + mask=tap_mask, + other=0.0, + ).to(tl.float32) + accumulator += coefficient * x + + tl.store( + output_ptr + row * output_stride_row + channels, + accumulator, + mask=mask, + ) + + +def dflash2_grouped_conv_impl( + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + block_size: int, + group_size: int, +) -> torch.Tensor: + """DFlash2 grouped convolution adapted from LightSeek TokenSpeed PR #1399.""" + num_rows, num_channels = x.shape + output = torch.empty_like(x) + if num_rows == 0: + return output + + element_block = 1024 if num_rows >= 128 and num_channels % 1024 == 0 else 512 + grid = (num_rows * triton.cdiv(num_channels, element_block),) + _dflash2_grouped_conv_kernel[grid]( + x, + delta, + base, + output, + x.stride(0), + delta.stride(0), + delta.stride(1), + base.stride(0), + output.stride(0), + NUM_CHANNELS=num_channels, + BLOCK_SIZE=block_size, + GROUP_SIZE=group_size, + TAPS=base.shape[0], + ELEMENT_BLOCK=element_block, + num_warps=4, + ) + return output + + +def dflash2_grouped_conv_fake( + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + block_size: int, + group_size: int, +) -> torch.Tensor: + return torch.empty_like(x) + + +direct_register_custom_op( + op_name="dflash2_grouped_conv", + op_func=dflash2_grouped_conv_impl, + fake_impl=dflash2_grouped_conv_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +def dflash2_grouped_conv( + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + block_size: int, + group_size: int, +) -> torch.Tensor: + return torch.ops.vllm.dflash2_grouped_conv(x, delta, base, block_size, group_size) + + def _grouped_conv( hidden_states: torch.Tensor, delta: torch.Tensor, @@ -29,6 +150,9 @@ def _grouped_conv( group_size: int, taps: int, ) -> torch.Tensor: + if hidden_states.is_cuda: + return dflash2_grouped_conv(hidden_states, delta, base, block_size, group_size) + blocks = hidden_states.unflatten(-1, (num_groups, group_size)) coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1) output = coefficients[:, 0] * blocks From ac78445bc9ee3bb615b8b21db22fb7dac6698ece Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 17 Sep 2026 17:47:30 -0400 Subject: [PATCH 437/774] [MoE] Encapsulate TRT-LLM BF16 weight layout handling (#57405) Signed-off-by: mgoin --- tests/kernels/moe/test_flashinfer.py | 30 ++++-- .../moe/test_moe_weight_loading_padded.py | 25 ----- tests/kernels/moe/test_trtllm_bf16_moe.py | 49 +++++++++ .../fused_moe/experts/trtllm_bf16_moe.py | 17 ++- .../fused_moe/experts/trtllm_lora_moe.py | 7 +- .../layers/fused_moe/oracle/unquantized.py | 27 ++++- .../fused_moe/unquantized_fused_moe_method.py | 100 +++++------------- 7 files changed, 138 insertions(+), 117 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 53d760c339c8..d17416de99b8 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( FlashInferExperts, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + view_as_block_major_k, +) from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( TrtLlmFp8ExpertsModular, TrtLlmFp8ExpertsMonolithic, @@ -570,6 +573,10 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( padded = w2_shape[-1] w13_ptr = layer.w13_weight.data_ptr() w2_ptr = layer.w2_weight.data_ptr() + load_numel = ( + layer.w13_weight.weight_loader_numel, + layer.w2_weight.weight_loader_numel, + ) for _ in range(2): reloaded_w13 = torch.randn_like(layer.w13_weight) @@ -602,19 +609,24 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( assert layer.w2_weight.shape == w2_shape assert layer.w13_weight.data_ptr() == w13_ptr assert layer.w2_weight.data_ptr() == w2_ptr - kernel_w13, kernel_w2 = method._kernel_weights(layer) + assert ( + layer.w13_weight.weight_loader_numel, + layer.w2_weight.weight_loader_numel, + ) == load_numel + kernel_w13 = view_as_block_major_k(layer.w13_weight) + kernel_w2 = view_as_block_major_k(layer.w2_weight) assert torch.equal(kernel_w13, expected_w13) assert torch.equal(kernel_w2, expected_w2) -def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index): +def _check_flashinfer_ipc_weights(entries, expected, mode, device_index): from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader from vllm.model_executor.utils import weights_already_processed torch.accelerator.set_device_index(device_index) with pytest.MonkeyPatch.context() as monkeypatch: method, layer = _make_unquantized_flashinfer_test_layer( - monkeypatch, 192, is_gated, device="meta" + monkeypatch, 192, is_gated=True, device="meta" ) loader = object.__new__(IpcModelLoader) loader.mode = mode @@ -625,7 +637,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde with weights_already_processed(): method.process_weights_after_loading(layer) assert method.moe_kernel is not None - actual = method._kernel_weights(layer) + actual = ( + view_as_block_major_k(layer.w13_weight), + view_as_block_major_k(layer.w2_weight), + ) for weight, reference, pointer in zip(actual, expected, pointers): assert weight.data_ptr() == pointer assert torch.equal(weight.cpu(), reference) @@ -636,11 +651,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde torch.accelerator.synchronize() -@pytest.mark.parametrize("is_gated", [True, False]) @pytest.mark.parametrize("cache_ndim", [3, 4]) @pytest.mark.parametrize("mode", ["copy", "zero_copy"]) def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( - monkeypatch, is_gated, cache_ndim, mode + monkeypatch, cache_ndim, mode ): """A spawned IPC consumer restores views without transient method state.""" from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( @@ -649,13 +663,12 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry method, layer = _make_unquantized_flashinfer_test_layer( - monkeypatch, 192, is_gated, device="meta" + monkeypatch, 192, is_gated=True, device="meta" ) packed = convert_moe_weights_to_flashinfer_trtllm_block_layout( {}, torch.randn_like(layer.w13_weight, device="cuda"), torch.randn_like(layer.w2_weight, device="cuda"), - is_gated_act_gemm=is_gated, ) expected = [weight.cpu() for weight in packed] entries = {} @@ -667,7 +680,6 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( args=( entries, expected, - is_gated, mode, torch.accelerator.current_device_index(), ), diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index 423f18113254..ef31de06d14a 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -12,7 +12,6 @@ import pytest import torch -from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, ) @@ -465,30 +464,6 @@ def test_rounds_intermediate_before_weight_allocation( assert hidden == 64 assert intermediate == expected - @pytest.mark.parametrize("is_act_and_mul", [True, False]) - def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool): - method = self._make_method(1344) - if not is_act_and_mul: - method.moe.activation = MoEActivation.RELU2_NO_MUL - method.moe.intermediate_size_per_partition = 1408 - layer = torch.nn.Module() - - method.create_weights( - layer=layer, - num_experts=2, - hidden_size=64, - intermediate_size_per_partition=1408, - params_dtype=torch.bfloat16, - ) - - up_mult = 2 if is_act_and_mul else 1 - assert layer.w13_weight.shape == (2, up_mult * 1408, 64) - assert layer.w2_weight.shape == (2, 64, 1408) - assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0 - if is_act_and_mul: - assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0 - assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0 - class TestLoadWeightsExpertBias: """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index f7466c63c895..5e569e153a7c 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/tests/kernels/moe/test_trtllm_bf16_moe.py @@ -12,6 +12,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.moe.utils import make_dummy_moe_config from tests.kernels.utils import torch_moe from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -27,6 +28,9 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( TrtLlmBf16ExpertsModular, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import ( + TrtLlmBf16LoRAExperts, +) from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, convert_to_unquantized_kernel_format, @@ -150,3 +154,48 @@ def test_trtllm_bf16_moe_modular_no_graph( atol=1e-1, rtol=2e-1, ) + + +@pytest.mark.parametrize("has_lora_delta", [False, True]) +@torch.inference_mode() +def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta): + """LoRA dispatch accepts both 3D parameters and legacy 4D packed views.""" + config = make_dummy_moe_config( + num_experts=128, hidden_dim=256, intermediate_size=128 + ) + experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG) + x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10 + w1, w2 = convert_to_unquantized_kernel_format( + UnquantizedMoeBackend.FLASHINFER_TRTLLM, + config, + torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10, + torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10, + ) + assert w1.ndim == w2.ndim == 3 + topk = ( + torch.tensor([[0], [1]], device="cuda", dtype=torch.int32), + torch.ones(2, 1, device="cuda", dtype=torch.float32), + ) + delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16) + + def invoke(w1, w2): + result = experts.invoke_routed_moe( + hidden_states=x, + w1=w1, + w2=w2, + topk_ids_and_weights=topk, + gemm1_lora_delta=delta if has_lora_delta else None, + global_num_experts=128, + a1q_scale=None, + output=torch.empty_like(x), + ) + if has_lora_delta: + # Only compare rows belonging to real tokens, excluding padding. + indices = result[2].flatten().long() + return result[0][indices], result[3][indices] + return result + + actual = invoke(w1, w2) + expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64)) + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 446acf46dd8b..54227d3cd1eb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -31,6 +31,15 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +def view_as_block_major_k(weight: torch.Tensor) -> torch.Tensor: + """View packed storage, including legacy 4D IPC cache entries.""" + if weight.ndim == 4: + return weight + experts, rows, cols = weight.shape + block_k = 128 // weight.element_size() + return weight.view(experts, cols // block_k, rows, block_k) + + class TrtLlmBf16ExpertsBase: """BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and monolithic interfaces. @@ -200,8 +209,8 @@ def apply( result = flashinfer.fused_moe.trtllm_bf16_routed_moe( topk_ids=(topk_ids, topk_weights), hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), num_experts=global_num_experts, top_k=topk_ids.size(1), n_group=None, @@ -281,8 +290,8 @@ def apply( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), num_experts=global_num_experts, top_k=self.topk, n_group=num_expert_group, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 4dc5076cf6db..1c5698932a4e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -34,6 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( LoRAExpertsMixin, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + view_as_block_major_k, +) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -527,8 +530,8 @@ def invoke_routed_moe( ret = flashinfer.fused_moe.trtllm_bf16_routed_moe( topk_ids=topk_ids_and_weights, hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), gemm1_lora_delta=gemm1_lora_delta, num_experts=global_num_experts, top_k=self.topk, diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index d1437933c3ef..bcb00eecbbb9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -26,6 +26,7 @@ swap_w13_to_w31, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey @@ -349,6 +350,17 @@ def _return_or_raise( ) +def unquantized_round_up_hidden_size_and_intermediate_size( + backend: UnquantizedMoeBackend, + hidden_size: int, + intermediate_size: int, +) -> tuple[int, int]: + """Round up dimensions before allocation to satisfy the selected kernel.""" + if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + intermediate_size = round_up(intermediate_size, 128) + return hidden_size, intermediate_size + + def convert_to_unquantized_kernel_format( unquantized_backend: UnquantizedMoeBackend, moe_config: FusedMoEConfig, @@ -377,13 +389,26 @@ def convert_to_unquantized_kernel_format( ) moe_config.intermediate_size_per_partition = padded_intermediate + # Reloads only overwrite checkpoint slices. An earlier in-place + # permutation can leave nonzero values in the raw padding slots. + unpadded = moe_config.intermediate_size_per_partition_unpadded + assert unpadded is not None + if padded_intermediate > unpadded: + w13_weight[:, unpadded:padded_intermediate].zero_() + if is_act_and_mul: + w13_weight[:, padded_intermediate + unpadded :].zero_() + w2_weight[:, :, unpadded:].zero_() + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} - w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout( + convert_moe_weights_to_flashinfer_trtllm_block_layout( _cache_permute_indices, w13_weight, w2_weight, is_gated_act_gemm=is_act_and_mul, ) + # Keep checkpoint-shaped parameters for reload/IPC. The experts create + # BlockMajorK views at dispatch without changing the underlying storage. + return w13_weight, w2_weight if ( unquantized_backend == UnquantizedMoeBackend.TRITON diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index a1b2471f5d41..e0835f538c32 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -24,6 +24,7 @@ convert_to_unquantized_kernel_format, make_unquantized_moe_kernel, select_unquantized_moe_backend, + unquantized_round_up_hidden_size_and_intermediate_size, ) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -34,7 +35,6 @@ set_weight_attrs, ) from vllm.platforms import current_platform -from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -72,9 +72,6 @@ def create_weights( ): unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded assert unpadded_intermediate is not None - has_intermediate_padding = ( - intermediate_size_per_partition > unpadded_intermediate - ) if self.moe.is_act_and_mul: w13_up_dim = 2 * intermediate_size_per_partition else: @@ -83,22 +80,15 @@ def create_weights( assert unpadded_hidden is not None unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1) # Fused gate_up_proj (column parallel) - w13_data = torch.empty( - num_experts, - w13_up_dim, - hidden_size, - dtype=params_dtype, + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_up_dim, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, ) - if has_intermediate_padding: - # The checkpoint loader only writes the unpadded slices. Zero the - # tails now so no full-size padding copy is needed after loading. - w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_() - if self.moe.is_act_and_mul: - w13_data[ - :, - intermediate_size_per_partition + unpadded_intermediate :, - ].zero_() - w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden @@ -111,15 +101,15 @@ def create_weights( set_weight_attrs(w13_bias, extra_weight_attrs) w13_bias.weight_loader_numel = num_experts * unpadded_up # down_proj (row parallel) - w2_data = torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=params_dtype, + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, ) - if has_intermediate_padding: - w2_data[:, :, unpadded_intermediate:].zero_() - w2_weight = torch.nn.Parameter(w2_data, requires_grad=False) layer.register_parameter("w2_weight", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) w2_weight.weight_loader_numel = ( @@ -147,11 +137,9 @@ def maybe_roundup_sizes( act_dtype, moe_parallel_config, ) - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - intermediate_size_per_partition = round_up( - intermediate_size_per_partition, 128 - ) - return hidden_size, intermediate_size_per_partition + return unquantized_round_up_hidden_size_and_intermediate_size( + self.unquantized_backend, hidden_size, intermediate_size_per_partition + ) def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor: # Pad the weight tensor. This is an optimization on ROCm platform, which @@ -177,21 +165,6 @@ def _setup_kernel( w13: torch.Tensor, w2: torch.Tensor, ) -> None: - w13_parameter_shape = w13.shape - w2_parameter_shape = w2.shape - - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - unpadded = self.moe.intermediate_size_per_partition_unpadded - assert unpadded is not None - intermediate = w2.shape[-1] - # Reloads overwrite only checkpoint slices; the previous in-place - # permutation can leave nonzero values in the raw padding slots. - if intermediate > unpadded: - w13[:, unpadded:intermediate].zero_() - if self.moe.is_act_and_mul: - w13[:, intermediate + unpadded :].zero_() - w2[:, :, unpadded:].zero_() - # Shuffle weights to runtime format. w13_new, w2_new = convert_to_unquantized_kernel_format( self.unquantized_backend, @@ -199,10 +172,6 @@ def _setup_kernel( w13_weight=w13, w2_weight=w2, ) - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - w13_new = w13_new.view(w13_parameter_shape) - w2_new = w2_new.view(w2_parameter_shape) - # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__; # On the first call we replace the parameter normally. On subsequent # calls (e.g. RL weight updates that re-trigger @@ -231,25 +200,6 @@ def _setup_kernel( assert self.moe_kernel is not None self.moe_kernel.fused_experts.process_weights_after_loading(layer) - def _kernel_weights( - self, layer: "RoutedExperts" - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM: - return layer.w13_weight, layer.w2_weight - - # IPC caches restore tensor shapes/storage, not method attributes. - # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements. - def block_view(weight: torch.Tensor) -> torch.Tensor: - if weight.ndim == 4: - return weight - experts, rows, cols = weight.shape - return weight.view(experts, cols // 64, rows, 64) - - return ( - block_view(layer.w13_weight), - block_view(layer.w2_weight), - ) - def _init_moe_kernel(self, layer: "RoutedExperts") -> None: """Build the MoE kernel from the layer's current (shuffled) weights.""" self.moe_quant_config = self.get_fused_moe_quant_config(layer) @@ -354,11 +304,10 @@ def forward_native( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert self.moe_kernel is not None - w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply( hidden_states=x, - w1=w13, - w2=w2, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -396,11 +345,10 @@ def apply_monolithic( ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None - w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply_monolithic( x, - w13, - w2, + layer.w13_weight, + layer.w2_weight, router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, From e8a2a0a8bbe675885ced6efddd49c1da0ff0a218 Mon Sep 17 00:00:00 2001 From: Shimi Bandiel Date: Thu, 17 Sep 2026 15:02:00 -0700 Subject: [PATCH 438/774] [Bugfix][Frontend] Reject stop strings on --tokens-only servers instead of silently ignoring them (#57058) Signed-off-by: Shimi Bandiel Co-authored-by: Claude Co-authored-by: Cyrus Leung Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../token_in_token_out/test_serving_tokens.py | 25 +++++++++++++++++++ .../scale_out/token_in_token_out/serving.py | 10 ++++++++ 2 files changed, 35 insertions(+) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 1c349e8d6724..aa90618a9068 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -422,6 +422,31 @@ async def test_stop_string_workflow(client, tokenizer, messages): assert generate_res == completions_res +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="--tokens-only is not supported by the Rust frontend", +) +@pytest.mark.parametrize("server", [["--tokens-only"]], indirect=True) +async def test_stop_strings_rejected_when_tokens_only(client): + """--tokens-only forces detokenize=False after request validation, so + stop strings must be rejected explicitly rather than silently ignored.""" + sampling_params = {"max_tokens": 5, "stop": ["never"]} + payload = { + "model": MODEL_NAME, + "token_ids": [1, 2, 3], + "sampling_params": sampling_params, + "stream": False, + } + resp = await client.post(GEN_ENDPOINT, json=payload) + assert resp.status_code == 400 + assert "stop strings" in resp.json()["error"]["message"] + + sampling_params["stop"] = None + resp = await client.post(GEN_ENDPOINT, json=payload) + resp.raise_for_status() + + @pytest.mark.asyncio @pytest.mark.parametrize( "server", diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a8a2804a0634..7813c7facc60 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -134,6 +134,16 @@ async def serve_tokens( f"sampling_params.n must be at most the server's max_num_seqs " f"({max_num_seqs}), got {sampling_params.n}." ) + if self.force_no_detokenize and sampling_params.stop: + # SamplingParams rejects stop with detokenize=False at request + # validation, but this server forces detokenize=False afterwards, + # so the combination must be rejected here or stop strings are + # silently never applied. + return self.create_error_response( + "stop strings are not supported on a --tokens-only server " + "because detokenization is disabled. Check stop strings on " + "the coordinator, or use stop_token_ids." + ) try: msgspec.msgpack.encode( ( From d3074acddade913b530e192143b3611f58938ce3 Mon Sep 17 00:00:00 2001 From: rasmith Date: Thu, 17 Sep 2026 17:13:13 -0500 Subject: [PATCH 439/774] [AMD][CI][The Rock] Fix language models standard for The Rock on mi355 (#53837) Signed-off-by: Randall Smith Co-authored-by: Andreas Karatzas --- tests/models/language/generation/test_common.py | 10 ++-------- tests/models/registry.py | 5 ++++- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 33ff4e7ccb38..50a0a5ed8ce8 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -26,7 +26,7 @@ "meta-llama/Llama-3.2-1B-Instruct", "openbmb/MiniCPM3-4B", "Qwen/Qwen2.5-0.5B-Instruct", - "TitanML/tiny-mixtral", + "axolotl-ai-co/tiny-mixtral-30m", "Qwen/Qwen3-8B", ] @@ -95,7 +95,7 @@ pytest.param("stabilityai/stablelm-3b-4e1t"), # stablelm pytest.param("bigcode/starcoder2-3b"), # starcoder2 pytest.param( - "TitanML/tiny-mixtral", # mixtral + "axolotl-ai-co/tiny-mixtral-30m", # mixtral marks=[pytest.mark.core_model], ), pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus @@ -126,12 +126,6 @@ def test_models( model_info.check_available_online(on_fail="skip") model_info.check_transformers_version(on_fail="skip") - if current_platform.is_rocm() and model == "TitanML/tiny-mixtral": - # Its single-token router selects LLMM1, whose low-precision - # accumulation can change the top-2 experts. Keep the optimized kernel - # enabled generally, but use the reference GEMM for this accuracy test. - monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0") - if use_rocm_aiter and (model in AITER_MODEL_LIST): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") elif use_rocm_aiter and model not in AITER_MODEL_LIST: diff --git a/tests/models/registry.py b/tests/models/registry.py index 8adb9ed15241..6a0fe8318e7e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -447,7 +447,10 @@ def check_available_online( ), "MixtralForCausalLM": _HfExamplesInfo( "mistralai/Mixtral-8x7B-Instruct-v0.1", - {"tiny": "TitanML/tiny-mixtral"}, + { + "tiny": "axolotl-ai-co/tiny-mixtral-30m", + "tiny-random": "TitanML/tiny-mixtral", + }, ), "NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"), "NemotronHForCausalLM": _HfExamplesInfo( From 80447d27655918da6bfbccd0d3a40e975bda220a Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 17 Sep 2026 15:14:37 -0700 Subject: [PATCH 440/774] [Bugfix][DSv4.1] Fix FlashInfer DSpark non-causal attention (#57432) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- .../test_dspark_noncausal_sparse_mla.py | 150 ++++++++++++++++++ .../deepseek_v41/nvidia/flashinfer_sparse.py | 68 +++++++- vllm/v1/attention/backends/mla/sparse_swa.py | 2 + 3 files changed, 214 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index d5c794131474..30c35a730efc 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -604,3 +604,153 @@ def test_dspark_noncausal_differs_from_causal( f"non-causal backend output matches the causal reference " f"(max abs diff={causal_err}); future-pointing indices are not attended to" ) + + +@pytest.mark.parametrize("context_len", [20, 128, 900]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_dsv41_flashinfer_dspark_window_matches_reference( + context_len, dtype, num_heads, monkeypatch +): + """Draft queries see the full block without treating padded slots as keys.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DSV4 TRTLLM sparse attention requires SM100") + from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( + DeepseekSparseSWAFlashInferMetadataBuilder, + DeepseekV4FlashInferMLAAttention, + ) + from vllm.models.deepseek_v41.sparse_mla import ( + DeepseekV41SparseSWAMetadataBuilder, + ) + + torch.manual_seed(123) + device = "cuda" + query_lens = [5, 3] + context_lens = [context_len, context_len + 137] + num_real_tokens, num_tokens, head_dim = 8, 9, 512 + cache = torch.randn(32, 128, head_dim, device=device, dtype=torch.bfloat16).to( + dtype + ) + query = torch.randn( + num_tokens, num_heads, head_dim, device=device, dtype=torch.bfloat16 + ).to(dtype) + indices = torch.full((num_tokens, 256), -1, device=device, dtype=torch.int32) + visible_indices = [] + visible_lens = [] + for req, (context, query_len) in enumerate(zip(context_lens, query_lens)): + visible = ( + torch.arange(max(context - 128, 0), context + query_len, device=device) + + req * 16 * 128 + ) + visible_indices.extend([visible] * query_len) + visible_lens.extend([visible.numel()] * query_len) + for token, visible in enumerate(visible_indices): + indices[token, : visible.numel()] = visible.to(torch.int32) + query_start_loc = torch.tensor([0, 5, 8, 9], dtype=torch.int32) + metadata = SimpleNamespace( + num_decodes=3, + num_prefills=0, + num_decode_tokens=num_tokens, + num_prefill_tokens=0, + seq_lens=torch.tensor( + [context_lens[0] + 5, context_lens[1] + 3, 1], + device=device, + dtype=torch.int32, + ), + query_start_loc=query_start_loc.to(device), + query_start_loc_cpu=query_start_loc, + token_to_req_indices=torch.tensor( + [0] * 5 + [1] * 3 + [2], device=device, dtype=torch.int32 + ), + decode_swa_indices=indices, + decode_swa_width=256, + decode_swa_lens=torch.tensor( + visible_lens + [0], device=device, dtype=torch.int32 + ), + block_table=torch.arange(48, device=device, dtype=torch.int32).view(3, -1), + block_size=128, + flashinfer_sparse_index_cache={}, + max_decode_query_len=5, + ) + + # Exercise FlashInfer preparation without constructing a model/config. + def init_parent(builder): + builder._max_tokens = num_tokens + builder.device = device + builder.window_size = 128 + + monkeypatch.setattr(DeepseekV41SparseSWAMetadataBuilder, "__init__", init_parent) + monkeypatch.setattr( + DeepseekV41SparseSWAMetadataBuilder, + "build", + lambda *args: metadata, + ) + builder = DeepseekSparseSWAFlashInferMetadataBuilder() + common_metadata = SimpleNamespace(causal=False) + builder.build(0, common_metadata) + prepared = ( + metadata.flashinfer_decode_topk_lens, + metadata.flashinfer_decode_seq_lens, + ) + attention = SimpleNamespace( + kv_cache_torch_dtype=dtype, + window_size=128, + compress_ratio=0, + topk_indices_buffer=torch.empty( + num_tokens, 0, device=device, dtype=torch.int32 + ), + scale=1 / math.sqrt(head_dim), + _flashinfer_fp8_bmm1_scale=1 / math.sqrt(head_dim), + _flashinfer_fp8_bmm2_scale=1.0, + attn_sink=None, + ) + attention._build_sparse_index_metadata = MethodType( + DeepseekV4FlashInferMLAAttention._build_sparse_index_metadata, attention + ) + output = torch.empty_like(query, dtype=torch.bfloat16) + + def forward(): + metadata.flashinfer_sparse_index_cache.clear() + DeepseekV4FlashInferMLAAttention._forward( + attention, query, None, cache, metadata, None, True, output + ) + + def check_output(): + references = [] + for token, visible in enumerate(visible_indices): + keys = cache.flatten(0, 1)[visible].float() + weights = torch.softmax( + query[token].float() @ keys.T / math.sqrt(head_dim), -1 + ) + references.append(weights @ keys) + atol = 0.01 if dtype == torch.bfloat16 else 0.05 + torch.testing.assert_close( + output[:num_real_tokens].float(), + torch.stack(references), + atol=atol, + rtol=0.05, + ) + + forward() + check_output() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + # Replay must consume the current inputs, including KV written since capture. + query.copy_(torch.randn_like(query, dtype=torch.bfloat16).to(dtype)) + cache.copy_(torch.randn_like(cache, dtype=torch.bfloat16).to(dtype)) + # A new step changes visibility as well as Q/KV, without recapturing. + metadata.seq_lens.sub_(1) + metadata.decode_swa_lens[:num_real_tokens].sub_(1) + for token, visible in enumerate(visible_indices): + indices[token, visible.numel() - 1] = -1 + visible_indices[token] = visible[:-1] + builder.build(0, common_metadata) + assert metadata.flashinfer_decode_topk_lens.data_ptr() == prepared[0].data_ptr() + assert metadata.flashinfer_decode_seq_lens.data_ptr() == prepared[1].data_ptr() + torch.testing.assert_close(prepared[0], metadata.decode_swa_lens.clamp_min(128)) + torch.testing.assert_close( + prepared[1], metadata.seq_lens[metadata.token_to_req_indices.long()] + ) + graph.replay() + check_output() diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 409725434a17..a247ae7d7409 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -26,7 +26,11 @@ ) from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 -from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf +from vllm.v1.attention.backend import ( + AttentionCGSupport, + CommonAttentionMetadata, + MultipleOf, +) from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) @@ -183,6 +187,39 @@ class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBui _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Graphs retain these addresses while each build refreshes their contents. + self._decode_topk_lens = torch.empty( + self._max_tokens, dtype=torch.int32, device=self.device + ) + self._decode_seq_lens = torch.empty_like(self._decode_topk_lens) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> "DeepseekSparseSWAMetadata": + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) + num_tokens = metadata.num_decode_tokens + if not common_attn_metadata.causal and num_tokens > 0: + assert metadata.decode_swa_lens is not None + assert metadata.seq_lens is not None + assert metadata.token_to_req_indices is not None + topk_lens = self._decode_topk_lens[:num_tokens] + seq_lens = self._decode_seq_lens[:num_tokens] + torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens) + torch.index_select( + metadata.seq_lens, + 0, + metadata.token_to_req_indices[:num_tokens], + out=seq_lens, + ) + metadata.flashinfer_decode_topk_lens = topk_lens + metadata.flashinfer_decode_seq_lens = seq_lens + return metadata + class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): @staticmethod @@ -493,21 +530,40 @@ def _forward( # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for # uniform-q batches, and this avoids flattening mixed batches into one call. if num_decode_tokens > 0: + decode_query = query[:num_decode_tokens] + decode_output = output[:num_decode_tokens] decode_cu = query_start_loc[: num_decodes + 1] + decode_seq_lens = seq_lens[:num_decodes] + decode_topk_lens = sparse_topk_lens[:num_decode_tokens] + max_decode_query_len = swa_metadata.max_decode_query_len + if swa_metadata.decode_swa_width > self.window_size: + # DSpark's non-causal window extends past the fixed 128 SWA + # columns into the aliased compressed pool. Exclude padding, + # and expose the full block to each query instead of letting + # TRTLLM derive a causal SWA length from its query position. + assert swa_only + assert swa_metadata.flashinfer_decode_topk_lens is not None + assert swa_metadata.flashinfer_decode_seq_lens is not None + decode_topk_lens = swa_metadata.flashinfer_decode_topk_lens + decode_seq_lens = swa_metadata.flashinfer_decode_seq_lens + decode_query = decode_query.unsqueeze(1) + decode_output = decode_output.unsqueeze(1) + decode_cu = None + max_decode_query_len = 1 flashinfer_trtllm_batch_decode_sparse_mla_dsv4( - query=query[:num_decode_tokens], + query=decode_query, swa_kv_cache=swa_k_cache, workspace_buffer=workspace, sparse_indices=sparse_indices[:num_decode_tokens], compressed_kv_cache=compressed_kv_cache, - sparse_topk_lens=sparse_topk_lens[:num_decode_tokens], - seq_lens=seq_lens[:num_decodes], - out=output[:num_decode_tokens], + sparse_topk_lens=decode_topk_lens, + seq_lens=decode_seq_lens, + out=decode_output, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, sinks=self.attn_sink, cum_seq_lens_q=decode_cu, - max_q_len=swa_metadata.max_decode_query_len, + max_q_len=max_decode_query_len, ) if num_prefill_tokens > 0: diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index e88631c32bb3..0fa5ae414ce5 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -205,6 +205,8 @@ class DeepseekSparseSWAMetadata: num_decode_tokens: int = 0 num_prefill_tokens: int = 0 max_decode_query_len: int = 1 + flashinfer_decode_topk_lens: torch.Tensor | None = None + flashinfer_decode_seq_lens: torch.Tensor | None = None # Pre-computed prefill metadata shared across all DeepseekV4 attention layers. prefill_seq_lens: torch.Tensor | None = None From e0050f287aae8b3ddbe2947a5af1c5dd23db2781 Mon Sep 17 00:00:00 2001 From: Mustafa YILDIRIM <2341474+mustafayildirim@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:34:56 -0700 Subject: [PATCH 441/774] [Bugfix][ROCm] Add record_logical_topk_ready to ROCMAiterMLASparseImpl (GLM-5.3-Flash boot crash) (#57252) Signed-off-by: Mustafa YILDIRIM Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Andreas Karatzas --- vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e1365f4d30c2..69b6e342cd13 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -808,6 +808,11 @@ def __init__( (q_concat_shape, vllm_config.model_config.dtype), ) + def record_logical_topk_ready(self) -> None: + # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer + # but does not participate in sparse-MLA index groups. + pass + def _forward_mla( self, layer: AttentionLayer, From ccde9c61e958d9dbfc21af6c62df6ff15881c811 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:01:35 -0400 Subject: [PATCH 442/774] [Bugfix] Restore KV cache metadata GET method for external event consumer (#56925) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../v1/engine/test_kv_cache_group_metadata.py | 70 +++++++++++++++++++ vllm/v1/engine/core.py | 32 ++++++++- 2 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 tests/v1/engine/test_kv_cache_group_metadata.py diff --git a/tests/v1/engine/test_kv_cache_group_metadata.py b/tests/v1/engine/test_kv_cache_group_metadata.py new file mode 100644 index 000000000000..1b2dedec4dca --- /dev/null +++ b/tests/v1/engine/test_kv_cache_group_metadata.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.engine.core import EngineCore +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, + MLAAttentionSpec, +) + +pytestmark = pytest.mark.cpu_test + +BASE_BLOCK_SIZE = 1536 +DCP_WORLD_SIZE = 8 + + +def _make_engine_core_with_dcp_manager() -> EngineCore: + mla_spec = MLAAttentionSpec( + block_size=BASE_BLOCK_SIZE, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + block_size=BASE_BLOCK_SIZE, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["mla"], mla_spec), + KVCacheGroupSpec(["mamba"], mamba_spec), + ], + ) + manager = KVCacheManager( + config, + max_model_len=BASE_BLOCK_SIZE * DCP_WORLD_SIZE * 2, + scheduler_block_size=BASE_BLOCK_SIZE * DCP_WORLD_SIZE, + hash_block_size=BASE_BLOCK_SIZE, + dcp_world_size=DCP_WORLD_SIZE, + ) + scheduler = SimpleNamespace(kv_cache_config=config, kv_cache_manager=manager) + engine_core = object.__new__(EngineCore) + engine_core.scheduler = scheduler + return engine_core + + +def test_kv_cache_group_metadata_uses_effective_dcp_block_sizes() -> None: + engine_core = _make_engine_core_with_dcp_manager() + + metadata = engine_core.get_kv_cache_group_metadata() + + # DCP scales attention blocks but does not scale replicated Mamba state. + assert [item["block_size"] for item in metadata] == [12288, 1536] + assert [item["kind"] for item in metadata] == ["mla_attention", "mamba"] + + managers = engine_core.scheduler.kv_cache_manager.coordinator.single_type_managers + assert [item["block_size"] for item in metadata] == [ + manager.block_size for manager in managers + ] diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 7da7f2e0ece1..14224f95e0bc 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -85,7 +85,11 @@ EngineCoreSentinel, fault_tolerant_wrapper, ) -from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + get_kv_cache_spec_kind, + is_full_attention_spec, +) from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus @@ -457,6 +461,32 @@ def log_field(name: str, field: str) -> str: supported_pooling_tasks, ) + def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]: + """Return serializable KV cache metadata for external event consumers + (e.g., Dynamo). This function has no call sites within vLLM but needs to + be kept. + """ + kv_cache_config = getattr(self.scheduler, "kv_cache_config", None) + if kv_cache_config is None: + return [] + + kv_cache_manager = cast(Any, self.scheduler).kv_cache_manager + managers = kv_cache_manager.coordinator.single_type_managers + metadata: list[dict[str, int | str | None]] = [] + for group_idx, (group, manager) in enumerate( + zip(kv_cache_config.kv_cache_groups, managers, strict=True) + ): + spec = group.kv_cache_spec + metadata.append( + { + "group_idx": group_idx, + "kind": get_kv_cache_spec_kind(spec).value, + "block_size": manager.block_size, + "sliding_window": getattr(spec, "sliding_window", None), + } + ) + return metadata + def add_request(self, request: Request, request_wave: int = 0): """Add request to the scheduler. From a524f80b1f0dd85889bf227310670ee8fdcd5fd2 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Thu, 17 Sep 2026 17:09:06 -0600 Subject: [PATCH 443/774] [AMD][Bugfix] Make the nested-RoPE patch reach automatic validation (#57289) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Andreas Karatzas --- vllm/transformers_utils/config.py | 48 ++++++++++++++++++++++++++----- 1 file changed, 41 insertions(+), 7 deletions(-) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 20d0d73afb92..e78a718e2844 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -162,11 +162,12 @@ def __getitem__(self, key): _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"} -# Model types whose checkpoints store shared RoPE parameters alongside the -# per-layer-type dicts (e.g. Laguna's `original_max_position_embeddings`). -# Since transformers 5.17, `validate_rope` treats every top-level value of such -# a dict as a layer's parameters and raises on the shared ones. -_PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna"} +# Model types whose checkpoints carry shared RoPE parameters alongside the +# per-layer-type dicts. Since transformers 5.17, `validate_rope` treats every +# top-level value of such a dict as a layer's parameters and raises on the +# shared ones. `laguna` gets them injected by `convert_rope_params_to_dict`; +# `gemma4_text` ships them in the checkpoint itself. +_PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna", "gemma4_text"} # Model types whose checkpoints declare `layer_types` entries that upstream # transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config @@ -224,6 +225,35 @@ def _mistral_patch_hf_hub_constants() -> Iterator[None]: constants.SAFETENSORS_INDEX_FILE = hf_safetensors_index_file +def _install_hf_config_validator( + name: str, validator: Callable, supersedes: Callable +) -> None: + """Replace a ``PretrainedConfig`` validator on every ``@strict`` snapshot. + + ``@strict`` snapshots each ``validate_*`` method into ``__class_validators__`` + at class creation and automatic post-``__init__`` validation dispatches off + that frozen list, so assigning the class attribute alone only affects + explicit ``config.validate_*()`` calls. Every ``@strict``-decorated config + class owns a snapshot, so rewrite them all, matching ``supersedes`` by + identity to leave a genuine per-model override in place. + """ + setattr(PretrainedConfig, name, validator) + + seen: set[int] = set() + stack = [PretrainedConfig] + while stack: + cls = stack.pop() + if id(cls) in seen: + continue + seen.add(id(cls)) + validators = cls.__dict__.get("__class_validators__") + if isinstance(validators, list) and any(v is supersedes for v in validators): + cls.__class_validators__ = [ + validator if v is supersedes else v for v in validators + ] + stack.extend(cls.__subclasses__()) + + def _patch_hf_transformers_validate_rope(): """Transformers v5 moved the ignore_keys option from the method signature of validate_rope and replaced it with the ignore_keys_at_rope_validation parameter @@ -250,7 +280,9 @@ def patched_validate_rope(self, *args, **kwargs): return result patched_validate_rope.__vllm_patched__ = True # type: ignore[attr-defined] - PretrainedConfig.validate_rope = patched_validate_rope + _install_hf_config_validator( + "validate_rope", patched_validate_rope, _original_validate_rope + ) def _patch_hf_transformers_nested_rope_validation() -> None: @@ -284,7 +316,9 @@ def patched_validate_rope(self, *args, **kwargs): return _original_validate_rope(self, *args, **kwargs) patched_validate_rope.__vllm_nested_rope_patched__ = True # type: ignore[attr-defined] - PretrainedConfig.validate_rope = patched_validate_rope + _install_hf_config_validator( + "validate_rope", patched_validate_rope, _original_validate_rope + ) def _patch_hf_transformers_allowed_layer_types( From 092bdd6d57ac7c1cd5272339c80372053fd51bbe Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:12:12 +0800 Subject: [PATCH 444/774] [Frontend][Core] Add release_kv_cache_memory() API (#44890) Signed-off-by: Dakai An Signed-off-by: AlanFokCo Signed-off-by: aoshen02 Signed-off-by: Dakai An <77474977+andakai@users.noreply.github.com> Co-authored-by: AlanFokCo Co-authored-by: aoshen02 Co-authored-by: OpenAI Codex --- docs/features/sleep_mode.md | 21 +++++ rust/src/engine-core-client/src/client.rs | 17 +++- rust/src/mock-engine/src/engine.rs | 10 +-- rust/src/mock-engine/src/tests.rs | 1 + rust/src/server/src/middleware/metrics.rs | 1 + rust/src/server/src/routes.rs | 4 + rust/src/server/src/routes/sleep.rs | 33 +++++-- rust/src/server/src/routes/tests.rs | 89 ++++++++++++++++++- tests/basic_correctness/test_mem.py | 85 ++++++++++++++++++ tests/entrypoints/serve/dev/test_sleep.py | 25 +++++- tests/multimodal/test_cache.py | 75 ++++++++++++++++ tests/v1/core/test_scheduler.py | 76 ++++++++++++++++ tests/v1/engine/test_engine_core.py | 24 +++++ vllm/device_allocator/sleep_mode_backend.py | 14 +++ vllm/engine/protocol.py | 5 ++ vllm/entrypoints/llm.py | 8 ++ .../entrypoints/serve/dev/sleep/api_router.py | 6 ++ vllm/v1/engine/async_llm.py | 11 ++- vllm/v1/engine/core.py | 29 +++++- vllm/v1/engine/core_client.py | 31 +++++-- vllm/v1/engine/llm_engine.py | 11 ++- vllm/v1/executor/abstract.py | 32 +++++-- vllm/v1/worker/cpu_worker.py | 3 + vllm/v1/worker/gpu_worker.py | 3 + 24 files changed, 572 insertions(+), 42 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index dcb6ad8c940c..2616f374ea5a 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -74,6 +74,18 @@ llm.wake_up(tags=["weights"]) llm.wake_up(tags=["kv_cache"]) ``` +#### Release only KV cache memory + +`LLM.release_kv_cache_memory()` discards KV cache physical memory while keeping model weights resident. It requires a completed pause and all executor memory to be resident: full sleep, partial wake-up, and repeated release without restoring memory are rejected. Requests retained with `mode="keep"` are recomputed after wake-up. + +Use the default `cumem` backend with managed allocations: `enable_cumem_allocator=True` on CUDA/ROCm (also enabled by `enable_sleep_mode=True`), or `enable_sleep_mode=True` on XPU. Other backends must implement selective discard; CPU and unmanaged allocations are unsupported. + +```python +llm.sleep(level=0, mode="keep") # Wait for the pause to complete. +llm.release_kv_cache_memory() +llm.wake_up(tags=["kv_cache"]) # Reallocate KV cache and resume scheduling. +``` + ### Online Serving To enable sleep mode in a vLLM server you need to initialize it with the flag `VLLM_SERVER_DEV_MODE=1` and pass `--enable-sleep-mode` to the vLLM server. @@ -107,9 +119,18 @@ curl -X POST 'http://localhost:8000/collective_rpc' -H 'Content-Type: applicatio curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache' ``` +To release only KV cache memory, wait for the level 0 pause to complete before calling the release endpoint. The same resident-memory and backend requirements as the Python API apply. + +```bash +curl -X POST 'http://localhost:8000/sleep?level=0&mode=keep' +curl -X POST 'http://localhost:8000/release_kv_cache_memory' +curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache' +``` + #### HTTP endpoints - `POST /sleep?level=1` — Put the model to sleep (`level=1`). +- `POST /release_kv_cache_memory` — Discard KV cache memory after a completed pause, while all executor memory is resident. - `POST /wake_up` — Wake up the model. Supports optional `tags` query parameters for partial wake-up (e.g., `?tags=weights`). - `POST /collective_rpc` — Perform a collective remote procedure call (RPC). - `GET /is_sleeping` — Check if the model is sleeping. diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 4cd7420e561e..aa6aacf0e017 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -911,13 +911,22 @@ impl EngineCoreClient { Ok(()) } - /// Wake the engine from sleep, optionally limiting the wake-up to specific - /// tags. - pub async fn wake_up(&self, tags: Option>) -> Result<()> { - self.call_utility::<(), _>("wake_up", (tags,)).await?; + /// Release KV cache memory while keeping model weights resident. + pub async fn release_kv_cache_memory(&self) -> Result<()> { + self.call_utility::<(), _>("release_kv_cache_memory", ()).await?; Ok(()) } + /// Wake the engine from sleep, optionally limiting the wake-up to specific + /// tags, and return whether every engine is fully awake. + pub async fn wake_up(&self, tags: Option>) -> Result { + Ok(self + .call_utility::("wake_up", (tags,)) + .await? + .into_iter() + .all(|fully_awake| fully_awake)) + } + /// Pause the scheduler so generation can be halted pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> { self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?; diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 38f9b21be1d7..4844aba12018 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -95,12 +95,10 @@ fn utility_response( "get_supported_tasks" => utility_envelope(vec!["generate"]), "is_sleeping" => utility_envelope(false), "reset_prefix_cache" => utility_envelope(true), - "reset_mm_cache" - | "reset_encoder_cache" - | "profile" - | "sleep" - | "wake_up" - | "execute_dummy_batch" => utility_envelope(()), + "wake_up" => utility_envelope(true), + "reset_mm_cache" | "reset_encoder_cache" | "profile" | "sleep" | "execute_dummy_batch" => { + utility_envelope(()) + } _ => utility_envelope(Value::Nil), }?; diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index bee710fd0c15..c1cfbb49cb2d 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -195,6 +195,7 @@ async fn utility_requests_return_minimal_success_responses() { assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache")); client.reset_mm_cache().await.expect("reset mm cache"); client.reset_encoder_cache().await.expect("reset encoder cache"); + assert!(client.wake_up(Some(vec!["kv_cache".to_string()])).await.expect("wake up")); shutdown_mock(client, shutdown, task).await; } diff --git a/rust/src/server/src/middleware/metrics.rs b/rust/src/server/src/middleware/metrics.rs index 1e6a79a556e1..de07620cf75b 100644 --- a/rust/src/server/src/middleware/metrics.rs +++ b/rust/src/server/src/middleware/metrics.rs @@ -26,6 +26,7 @@ const EXCLUDED_HANDLERS: &[&str] = &[ "/reset_encoder_cache", "/collective_rpc", "/sleep", + "/release_kv_cache_memory", "/wake_up", "/is_sleeping", ]; diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 41e49878820b..1cd6af3d4fa2 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -142,6 +142,10 @@ fn build_router_with_options( .route("/weight_info", get(weight_transfer::weight_info)) .route("/abort_requests", post(abort_requests::abort_requests)) .route("/sleep", post(sleep::sleep)) + .route( + "/release_kv_cache_memory", + post(sleep::release_kv_cache_memory), + ) .route("/wake_up", post(sleep::wake_up)) .route("/is_sleeping", get(sleep::is_sleeping)) .route("/pause", post(pause::pause)) diff --git a/rust/src/server/src/routes/sleep.rs b/rust/src/server/src/routes/sleep.rs index 4bf9981402ca..d64a0850b6e4 100644 --- a/rust/src/server/src/routes/sleep.rs +++ b/rust/src/server/src/routes/sleep.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use axum::Json; use axum::extract::rejection::QueryRejection; -use axum::extract::{Query, State}; +use axum::extract::{Query, RawQuery, State}; use axum::http::StatusCode; use serde::{Deserialize, Serialize}; use vllm_engine_core_client::protocol::utility::PauseMode; @@ -27,12 +27,6 @@ pub(crate) struct SleepParams { mode: PauseMode, } -#[derive(Debug, Default, Deserialize)] -pub(crate) struct WakeUpParams { - #[serde(default)] - tags: Option>, -} - const fn default_sleep_level() -> u32 { 1 } @@ -57,14 +51,35 @@ pub async fn sleep( Ok(StatusCode::OK) } +/// Release KV cache memory while keeping model weights resident. +pub async fn release_kv_cache_memory( + State(state): State>, +) -> Result { + state + .engine_core_client() + .release_kv_cache_memory() + .await + .map_err(|error| utility_call_error("release_kv_cache_memory", error))?; + + Ok(StatusCode::OK) +} + /// Wake the engine from sleep mode. pub async fn wake_up( State(state): State>, - Query(params): Query, + RawQuery(query): RawQuery, ) -> Result { + let tags = query.and_then(|query| { + let tags = url::form_urlencoded::parse(query.as_bytes()) + .filter(|(key, _)| key == "tags") + .map(|(_, value)| value.into_owned()) + .collect::>(); + (!tags.is_empty()).then_some(tags) + }); + state .engine_core_client() - .wake_up(params.tags) + .wake_up(tags) .await .map_err(|error| utility_call_error("wake_up", error))?; diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 4731958e2f87..c27ac20fad32 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -5805,7 +5805,46 @@ async fn sleep_route_uses_python_compatible_default_query_values() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn wake_up_route_without_tags_sends_none() { +async fn release_kv_cache_memory_route_sends_expected_utility_call() { + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("release_kv_cache_memory")); + assert_eq!(array[3], Value::Array(Vec::new())); + + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/release_kv_cache_memory") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + let status = response.status(); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body)); + assert!(body.is_empty()); + engine_task.await.expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn wake_up_route_without_tags_accepts_bool_result() { let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { boxed_test_future(async move { let utility = recv_engine_message(dealer).await; @@ -5818,7 +5857,7 @@ async fn wake_up_route_without_tags_sends_none() { assert_eq!(array[2], Value::from("wake_up")); assert_eq!(array[3], Value::Array(vec![Value::Nil])); - send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; }) }) .await; @@ -5842,6 +5881,51 @@ async fn wake_up_route_without_tags_sends_none() { engine_task.await.expect("mock engine task"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn wake_up_route_accepts_repeated_tags() { + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("wake_up")); + assert_eq!( + array[3], + Value::Array(vec![Value::Array(vec![ + Value::from("weights"), + Value::from("kv_cache"), + ])]) + ); + + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/wake_up?tags=weights&tags%5B%5D=ignored&tags=kv_cache") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + let status = response.status(); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body)); + assert!(body.is_empty()); + engine_task.await.expect("mock engine task"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn is_sleeping_route_returns_json_payload() { @@ -6160,6 +6244,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { for (method, uri) in [ ("GET", "/is_sleeping"), ("POST", "/sleep"), + ("POST", "/release_kv_cache_memory"), ("POST", "/wake_up"), ("GET", "/is_paused"), ("POST", "/pause"), diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 6b4895c64369..746838dc40b6 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -109,6 +109,91 @@ def test_basic_cumem(): assert torch.allclose(output, torch.ones_like(output) * 3) +@pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"]) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch): + """Rejected release and KV-only/full sleep cycles preserve greedy tokens.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + kv_cache_memory_bytes = 256 * 1024 * 1024 + llm = LLM( + "Qwen/Qwen3-0.6B", + enable_sleep_mode=True, + enforce_eager=True, + kv_cache_memory_bytes=kv_cache_memory_bytes, + max_model_len=1024, + max_num_seqs=4, + gpu_memory_utilization=0.05, + ) + prompt = "How are you?" + sampling_params = SamplingParams(temperature=0, max_tokens=10) + expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids + + def get_mapped_bytes(worker): + return mapped_usage(get_mem_allocator_instance()) + + mapped_before = llm.collective_rpc(get_mapped_bytes)[0] + # Utility RPC transports the engine error as a plain Exception. + with pytest.raises(Exception, match="requires a completed pause"): + llm.release_kv_cache_memory() + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before + assert not llm.llm_engine.is_sleeping() + assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected + + llm.sleep(level=0) + llm.release_kv_cache_memory() + mapped_after = llm.collective_rpc(get_mapped_bytes)[0] + assert mapped_before - mapped_after >= kv_cache_memory_bytes * 0.99 + assert mapped_after > 0 + assert llm.llm_engine.is_sleeping() + + if full_sleep: + llm.sleep(level=full_sleep) + assert llm.collective_rpc(get_mapped_bytes)[0] == 0 + llm.wake_up(tags=None if full_sleep else ["kv_cache"]) + if full_sleep == 2: + llm.collective_rpc("reload_weights") + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before + assert not llm.llm_engine.is_sleeping() + actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids + assert actual == expected + + +@pytest.mark.parametrize("first_level", [1, 2]) +@pytest.mark.parametrize("second_level", [1, 2]) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch): + """Repeated sleep after KV-only wake preserves mappings and recoverability.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + llm = LLM( + "Qwen/Qwen3-0.6B", + enable_sleep_mode=True, + enforce_eager=True, + kv_cache_memory_bytes=256 * 1024 * 1024, + max_model_len=1024, + max_num_seqs=4, + gpu_memory_utilization=0.05, + ) + prompt = "How are you?" + sampling_params = SamplingParams(temperature=0, max_tokens=10) + expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids + + def get_mapped_bytes(worker): + return mapped_usage(get_mem_allocator_instance()) + + llm.sleep(level=first_level) + llm.wake_up(tags=["kv_cache"]) + mapped_before = llm.collective_rpc(get_mapped_bytes)[0] + assert mapped_before > 0 + llm.sleep(level=second_level) + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before + llm.wake_up(tags=["weights"]) + if first_level == 2: + llm.collective_rpc("reload_weights") + assert not llm.llm_engine.is_sleeping() + actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids + assert actual == expected + + @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") def test_discard_tags(): """Test that discard(tags) selectively frees GPU memory for specific diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py index 5ebee506002a..34e69c9ecef4 100644 --- a/tests/entrypoints/serve/dev/test_sleep.py +++ b/tests/entrypoints/serve/dev/test_sleep.py @@ -1,14 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import AsyncMock + +import pytest import requests +from fastapi import FastAPI +from fastapi.testclient import TestClient from prometheus_client.parser import text_string_to_metric_families from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.serve.dev.sleep.api_router import attach_router MODEL_NAME = "meta-llama/Llama-3.2-1B" +@pytest.mark.cpu_test +@pytest.mark.parametrize("fails", [False, True]) +def test_release_kv_cache_memory_route(fails): + app = FastAPI() + app.state.engine_client = AsyncMock() + release = app.state.engine_client.release_kv_cache_memory + if fails: + release.side_effect = RuntimeError("requires a completed pause first") + attach_router(app) + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/release_kv_cache_memory") + + assert response.status_code == (500 if fails else 200) + release.assert_awaited_once_with() + + def test_sleep_mode(): # dtype, max-len etc set so that this can run in CI args = [ @@ -63,7 +86,7 @@ def test_sleep_mode(): ) assert response.status_code == 200 - # is sleeping should be false after waking up any part of the engine + # Partial wake keeps the engine sleeping. response = requests.get(remote_server.url_for("is_sleeping")) assert response.status_code == 200 assert response.json().get("is_sleeping") is True diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index f77dacd341b8..0c68d2bac0fd 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor +from functools import partial from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest @@ -34,6 +37,8 @@ PlaceholderRange, ) from vllm.multimodal.processing import PromptInsertion +from vllm.renderers.base import BaseRenderer +from vllm.utils.async_utils import make_async from vllm.utils.mem_constants import GiB_bytes, MiB_bytes from ..utils import create_new_process_for_each_test @@ -755,6 +760,76 @@ def test_processor_cache_shared_across_loras(): assert feature_lora_b.data == item_data +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "release_error", + [ + None, + "requires a completed pause first", + "requires all executor memory to be resident", + ], + ids=["released", "not-paused", "nonresident-memory"], +) +@pytest.mark.asyncio +async def test_release_kv_cache_resends_mm_payload(use_async, release_error): + """Release must not leave a sender hit pointing at a cleared receiver.""" + from vllm.v1.engine.async_llm import AsyncLLM + from vllm.v1.engine.llm_engine import LLMEngine + + model_config = SimpleNamespace( + get_multimodal_config=lambda: MultiModalConfig(mm_processor_cache_gb=1) + ) + sender = MultiModalProcessorSenderCache(model_config) + receiver = MultiModalReceiverCache(model_config) + item = _dummy_item({"pixel_values": 16}) + mm_hash = "image_A" + payload, _ = sender.get_and_update_item((item, []), mm_hash) + assert receiver.get_and_update_item(payload, mm_hash) is item + assert sender.get_and_update_item(None, mm_hash)[0] is None + + renderer = SimpleNamespace(mm_processor_cache=sender, _mm_cache_stats=None) + renderer.clear_mm_cache = partial(BaseRenderer.clear_mm_cache, renderer) + + def release(): + if release_error: + raise RuntimeError(release_error) + receiver.clear_cache() + + with ThreadPoolExecutor(max_workers=1) as executor: + renderer._clear_mm_cache_async = make_async( + renderer.clear_mm_cache, executor=executor + ) + renderer.clear_mm_cache_async = partial( + BaseRenderer.clear_mm_cache_async, renderer + ) + engine = SimpleNamespace( + renderer=renderer, + engine_core=SimpleNamespace( + release_kv_cache_memory=release, + release_kv_cache_memory_async=make_async(release, executor=executor), + ), + logger_manager=Mock(), + ) + + async def call_release(): + if use_async: + await AsyncLLM.release_kv_cache_memory(engine) + else: + LLMEngine.release_kv_cache_memory(engine) + + if release_error: + with pytest.raises(RuntimeError, match=release_error): + await call_release() + engine.logger_manager.record_sleep_state.assert_not_called() + else: + await call_release() + engine.logger_manager.record_sleep_state.assert_called_once_with(1, 0) + + payload, _ = sender.get_and_update_item((item, []), mm_hash) + assert payload is item + assert receiver.get_and_update_item(payload, mm_hash) is item + + _SLEEP_VISION_PROMPT = ( "<|im_start|>system\nYou are a helpful assistant.<|im_end|>" "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 298c8c618576..4951014d1ba1 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -36,10 +36,13 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash +from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.engine import FinishReason +from vllm.v1.engine.core import EngineCore +from vllm.v1.executor.uniproc_executor import UniProcExecutor from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -1487,6 +1490,79 @@ def test_scheduler_reset_prefix_cache(): assert scheduler.waiting[i] == request +def test_kv_cache_release_after_keep_pause_preserves_requests(): + """A completed keep pause permits release and recomputation of live requests.""" + scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True) + running, waiting = create_requests(num_requests=2) + for request in (running, waiting): + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[running.request_id], + req_id_to_index={running.request_id: 0}, + sampled_token_ids=[[1]], + ), + ) + assert scheduler.kv_cache_manager.get_block_ids(running.request_id)[0] + + core = object.__new__(EngineCore) + core.scheduler = scheduler + core.model_executor = Mock(is_sleeping=False) + core.mm_receiver_cache = None + core.batch_queue = None + + assert core.pause_scheduler(mode="keep", clear_cache=False) is None + assert scheduler.running == [running] + assert list(scheduler.waiting) == [waiting] + assert not scheduler.has_requests() + + core.release_kv_cache_memory() + + core.model_executor.discard.assert_called_once_with(("kv_cache",)) + assert scheduler.requests == { + request.request_id: request for request in (running, waiting) + } + assert not scheduler.running + assert list(scheduler.waiting) == [running, waiting] + assert running.status == RequestStatus.PREEMPTED + assert running.num_computed_tokens == 0 + assert list(running.output_token_ids) == [1] + assert not scheduler.kv_cache_manager.get_block_ids(running.request_id)[0] + assert scheduler.schedule().total_num_scheduled_tokens == 0 + + assert core.wake_up(tags=["kv_cache"]) + resumed = scheduler.schedule() + assert resumed.num_scheduled_tokens == {running.request_id: running.num_tokens} + + +@pytest.mark.parametrize( + "sleeping_tags", + [{"weights", "kv_cache"}, {"weights"}, {"kv_cache"}], + ids=["full-sleep", "weights-only", "kv-only-or-repeated-release"], +) +def test_kv_cache_release_rejects_nonresident_memory(sleeping_tags): + """Reject release without resetting caches or discarding more memory.""" + core = object.__new__(EngineCore) + core.scheduler = create_scheduler() + core.scheduler.set_pause_state(PauseState.PAUSED_ALL) + core.batch_queue = None + core._reset_caches = Mock() + core.model_executor = object.__new__(UniProcExecutor) + core.model_executor.sleeping_tags = sleeping_tags.copy() + core.model_executor.collective_rpc = Mock() + + with pytest.raises( + RuntimeError, match="requires all executor memory to be resident" + ): + core.release_kv_cache_memory() + + core._reset_caches.assert_not_called() + core.model_executor.collective_rpc.assert_not_called() + assert core.model_executor.sleeping_tags == sleeping_tags + + def test_reset_connector_cache_no_connector_is_no_op_success(): """``reset_connector_cache`` must return True when no connector is configured. diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index ba1e9a4d200b..e16e0550c21d 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -23,6 +23,7 @@ from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.utils.torch_utils import set_default_torch_num_threads +from vllm.v1.core.sched.interface import PauseState from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc from vllm.v1.executor.abstract import Executor @@ -667,6 +668,29 @@ def _pausable_engine_core_proc() -> EngineCoreProc: return core +@pytest.mark.parametrize( + "pause_state,has_requests,has_batches", + [ + pytest.param(PauseState.UNPAUSED, False, False, id="not-paused"), + pytest.param(PauseState.PAUSED_ALL, True, False, id="pending-requests"), + pytest.param(PauseState.PAUSED_ALL, False, True, id="pending-batches"), + ], +) +def test_kv_cache_release_rejects_unsafe_state(pause_state, has_requests, has_batches): + """Reject release before touching caches or memory if work can still use KV.""" + core = _pausable_engine_core_proc() + core.scheduler.pause_state = pause_state + core.scheduler.has_requests.return_value = has_requests + core.batch_queue = [object()] if has_batches else None + core._reset_caches = MagicMock() + + with pytest.raises(RuntimeError, match="requires a completed pause"): + core.release_kv_cache_memory() + + core._reset_caches.assert_not_called() + core.model_executor.discard.assert_not_called() + + @pytest.mark.parametrize("deferred", [False, True]) def test_pause_synchronizes_device_before_cache_reset(deferred: bool): """A resolved pause promises an idle device: the barrier must run before diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 45a635d61389..023af04275c4 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -73,6 +73,13 @@ def state(self) -> SleepModeState: (suspended) engine from a healthy-serving one (see RFC #34303).""" return self._state + def discard(self, tags: tuple[str, ...]) -> None: + """Discard selected tagged allocations without preserving contents.""" + raise RuntimeError( + f"Sleep-mode backend {type(self).__name__!r} does not support " + "selective discard." + ) + # -- Capability introspection (no instance required) -- @classmethod @@ -132,6 +139,13 @@ def resume(self, tags: list[str] | None = None) -> None: allocator.wake_up(tags) self._state = "RUNNING" + def discard(self, tags: tuple[str, ...]) -> None: + from vllm.device_allocator import get_mem_allocator_instance + + allocator = get_mem_allocator_instance() + allocator.discard(tags) + self._state = "SUSPENDED" + @classmethod def preserves_communicators(cls) -> bool: # Comm identity/topology survive memory suspension; no reinit needed. diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index c44fd44d0fb3..602e2570a25f 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -186,6 +186,11 @@ async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None: """Sleep the engine.""" ... + @abstractmethod + async def release_kv_cache_memory(self) -> None: + """Discard KV cache physical GPU memory. Requires a completed pause.""" + ... + @abstractmethod async def wake_up(self, tags: list[str] | None = None) -> None: """Wake up the engine.""" diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index d9a2a46701c1..999310aa6169 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -830,6 +830,14 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"): """ self.llm_engine.sleep(level=level, mode=mode) + def release_kv_cache_memory(self) -> None: + """Release the GPU physical memory backing the KV cache. + + Requires a completed ``sleep(level=0)`` and resident executor memory. + Restore with ``wake_up(tags=["kv_cache"])``; kept requests are recomputed. + """ + self.llm_engine.release_kv_cache_memory() + def wake_up(self, tags: list[str] | None = None): """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep] method for more details. diff --git a/vllm/entrypoints/serve/dev/sleep/api_router.py b/vllm/entrypoints/serve/dev/sleep/api_router.py index fedff8da749b..846a487706b7 100644 --- a/vllm/entrypoints/serve/dev/sleep/api_router.py +++ b/vllm/entrypoints/serve/dev/sleep/api_router.py @@ -27,6 +27,12 @@ async def sleep(raw_request: Request): return Response(status_code=200) +@router.post("/release_kv_cache_memory") +async def release_kv_cache_memory(raw_request: Request): + await engine_client(raw_request).release_kv_cache_memory() + return Response(status_code=200) + + @router.post("/wake_up") async def wake_up(raw_request: Request): tags = raw_request.query_params.getlist("tags") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 3e399a1a816e..ad9d9bf6c776 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1079,10 +1079,17 @@ async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: if self.logger_manager is not None: self.logger_manager.record_sleep_state(1, level) - async def wake_up(self, tags: list[str] | None = None) -> None: - await self.engine_core.wake_up_async(tags) + async def release_kv_cache_memory(self) -> None: + await self.renderer.clear_mm_cache_async() + await self.engine_core.release_kv_cache_memory_async() if self.logger_manager is not None: + self.logger_manager.record_sleep_state(1, 0) + + async def wake_up(self, tags: list[str] | None = None) -> None: + fully_awake = await self.engine_core.wake_up_async(tags) + + if self.logger_manager is not None and fully_awake: self.logger_manager.record_sleep_state(0, 0) async def checkpoint_prepare(self) -> None: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 14224f95e0bc..5ee2c6f8c02c 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -957,12 +957,15 @@ def pause_complete(f: Future): pause_future.add_done_callback(pause_complete) return future - def wake_up(self, tags: list[str] | None = None): + def wake_up(self, tags: list[str] | None = None) -> bool: """Wake up the engine from sleep. Args: tags: Tags to wake up. Use ["scheduling"] for level 0 wake up. + Returns: + Whether all executor memory is resident again (fully awake). + """ if tags is not None and "scheduling" in tags: # Remove "scheduling" from tags if there are other tags to process. @@ -973,8 +976,30 @@ def wake_up(self, tags: list[str] | None = None): # Partial wakes intentionally keep the remaining allocations asleep. # Resume scheduling only once all executor memory is resident again. - if not self.model_executor.is_sleeping: + fully_awake = not self.model_executor.is_sleeping + if fully_awake: self.resume_scheduler() + return fully_awake + + def release_kv_cache_memory(self) -> None: + """Discard KV cache physical memory. Requires a completed pause + and all executor memory to be resident. Kept requests are recomputed + after wake-up. + """ + if not ( + self.is_scheduler_paused() + and not self.scheduler.has_requests() + and not self.batch_queue + ): + raise RuntimeError( + "release_kv_cache_memory() requires a completed pause first" + ) + if self.model_executor.is_sleeping: + raise RuntimeError( + "release_kv_cache_memory() requires all executor memory to be resident" + ) + self._reset_caches() + self.model_executor.discard(("kv_cache",)) def is_sleeping(self) -> bool: """Check if engine is sleeping at any level.""" diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 4ca565d56cd4..399af7c0b9d8 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -193,7 +193,10 @@ def reset_encoder_cache(self) -> None: def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: raise NotImplementedError - def wake_up(self, tags: list[str] | None = None) -> None: + def release_kv_cache_memory(self) -> None: + raise NotImplementedError + + def wake_up(self, tags: list[str] | None = None) -> bool: raise NotImplementedError def is_sleeping(self) -> bool: @@ -285,7 +288,10 @@ async def reset_encoder_cache_async(self) -> None: async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None: raise NotImplementedError - async def wake_up_async(self, tags: list[str] | None = None) -> None: + async def release_kv_cache_memory_async(self) -> None: + raise NotImplementedError + + async def wake_up_async(self, tags: list[str] | None = None) -> bool: raise NotImplementedError async def is_sleeping_async(self) -> bool: @@ -397,8 +403,11 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: result = self.engine_core.sleep(level, mode) assert result is None - def wake_up(self, tags: list[str] | None = None) -> None: - self.engine_core.wake_up(tags) + def release_kv_cache_memory(self) -> None: + self.engine_core.release_kv_cache_memory() + + def wake_up(self, tags: list[str] | None = None) -> bool: + return self.engine_core.wake_up(tags) def is_sleeping(self) -> bool: return self.engine_core.is_sleeping() @@ -1021,8 +1030,11 @@ def pin_lora(self, lora_id: int) -> bool: def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: self.call_utility("sleep", level, mode) - def wake_up(self, tags: list[str] | None = None) -> None: - self.call_utility("wake_up", tags) + def release_kv_cache_memory(self) -> None: + self.call_utility("release_kv_cache_memory") + + def wake_up(self, tags: list[str] | None = None) -> bool: + return self.call_utility("wake_up", tags) def is_sleeping(self) -> bool: return self.call_utility("is_sleeping") @@ -1265,8 +1277,11 @@ async def reset_encoder_cache_async(self) -> None: async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None: await self.call_utility_async("sleep", level, mode) - async def wake_up_async(self, tags: list[str] | None = None) -> None: - await self.call_utility_async("wake_up", tags) + async def release_kv_cache_memory_async(self) -> None: + await self.call_utility_async("release_kv_cache_memory") + + async def wake_up_async(self, tags: list[str] | None = None) -> bool: + return await self.call_utility_async("wake_up", tags) async def is_sleeping_async(self) -> bool: return await self.call_utility_async("is_sleeping") diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index c6c77f72e716..a2d97b5544fb 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -375,10 +375,17 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"): if self.logger_manager is not None: self.logger_manager.record_sleep_state(1, level) - def wake_up(self, tags: list[str] | None = None): - self.engine_core.wake_up(tags) + def release_kv_cache_memory(self) -> None: + self.renderer.clear_mm_cache() + self.engine_core.release_kv_cache_memory() if self.logger_manager is not None: + self.logger_manager.record_sleep_state(1, 0) + + def wake_up(self, tags: list[str] | None = None): + fully_awake = self.engine_core.wake_up(tags) + + if self.logger_manager is not None and fully_awake: self.logger_manager.record_sleep_state(0, 0) def is_sleeping(self) -> bool: diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 6b07bad17c76..f234e943b4fd 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -32,6 +32,9 @@ _R = TypeVar("_R") +# Memory-pool tags managed by sleep/wake_up/discard. +SLEEP_TAGS = frozenset(("weights", "kv_cache")) + FailureCallback = Callable[[], None] @@ -108,7 +111,6 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config self._init_executor() - self.is_sleeping = False self.sleeping_tags: set[str] = set() self.kv_output_aggregator: KVOutputAggregator | None = None self.ec_output_aggregator: ECOutputAggregator | None = None @@ -338,15 +340,18 @@ def reset_encoder_cache(self) -> None: """Reset the encoder cache in each worker to clear cached encoder outputs.""" self.collective_rpc("reset_encoder_cache") + @property + def is_sleeping(self) -> bool: + return bool(self.sleeping_tags) + def sleep(self, level: int = 1): - if self.is_sleeping: + if "weights" in self.sleeping_tags: logger.warning("Executor is already sleeping.") return time_before_sleep = time.perf_counter() self.collective_rpc("sleep", kwargs=dict(level=level)) time_after_sleep = time.perf_counter() - self.sleeping_tags = {"weights", "kv_cache"} - self.is_sleeping = True + self.sleeping_tags |= SLEEP_TAGS logger.info( "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep ) @@ -375,8 +380,23 @@ def wake_up(self, tags: list[str] | None = None): self.sleeping_tags.remove(tag) else: self.sleeping_tags.clear() - if not self.sleeping_tags: - self.is_sleeping = False + + def discard(self, tags: tuple[str, ...]) -> None: + tags_to_discard = set(tags) - self.sleeping_tags + if not tags_to_discard: + logger.warning("Tags %s are already sleeping.", tags) + return + time_before_discard = time.perf_counter() + try: + self.collective_rpc("discard", args=(tuple(tags_to_discard),)) + finally: + self.sleeping_tags |= tags_to_discard + time_after_discard = time.perf_counter() + logger.info( + "It took %.6f seconds to discard tags %s.", + time_after_discard - time_before_discard, + tags_to_discard, + ) def reinitialize_distributed( self, reconfig_request: ReconfigureDistributedRequest diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index 2ee2164dc31e..1badcfe6e82e 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -198,6 +198,9 @@ def wake_up(self, tags: list[str] | None = None) -> None: logger.warning("sleep mode is not supported on CPU, ignore it.") pass + def discard(self, tags: tuple[str, ...]) -> None: + raise RuntimeError("Selective discard is not supported on CPU.") + def _should_warm_up_model(self) -> bool: # VLLM_CPU_CI_ENV always skips warmup to save CI time. if envs.VLLM_CPU_CI_ENV: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 02f599e57e8c..8ced3bdd4307 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -304,6 +304,9 @@ def wake_up(self, tags: list[str] | None = None) -> None: self.synchronize_device() + def discard(self, tags: tuple[str, ...]) -> None: + self.sleep_mode_backend.discard(tags) + def checkpoint_prepare(self) -> None: checkpoint_prepare_distributed_state() From 48f663c37f8c77803d9281778d49ec0727abcacf Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 18 Sep 2026 07:52:01 +0800 Subject: [PATCH 445/774] [XPU] Fix incorrect context-key normalization for Qwen DFlash-based models (#56431) Signed-off-by: Yan Ma --- vllm/model_executor/models/qwen3_dflash.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index ce3884a10eb8..5504304c733a 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -32,6 +32,7 @@ VocabParallelEmbedding, ) from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType @@ -561,6 +562,15 @@ def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor: # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) --- # The weight is selected per layer by the outermost (layer) index. all_k_normed = torch.empty_like(all_k) + if current_platform.is_xpu(): + for layer_idx in range(all_k.shape[0]): + ops.rms_norm( + all_k_normed[layer_idx], + all_k[layer_idx], + self._k_norm_weights[layer_idx], + self._rms_norm_eps, + ) + return all_k_normed ops.rms_norm( all_k_normed, all_k, From 213c379fcafe5ac3bdc235a480e13140ee13c7b5 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Thu, 17 Sep 2026 19:06:05 -0500 Subject: [PATCH 446/774] [ROCm][Bugfix] Gate AITER MXFP8 MoE on the aiter enable flag (#57426) Signed-off-by: Rohan Potdar Co-authored-by: Claude --- .../moe/test_mxfp8_aiter_backend_selection.py | 37 +++++++++------ tests/models/quantization/test_mxfp8.py | 15 ------ .../fused_moe/experts/aiter_mxfp8_moe.py | 47 +------------------ 3 files changed, 24 insertions(+), 75 deletions(-) diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 51aff4ea3ba0..61b26af62c09 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -2,10 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MXFP8 MoE backend selection for the AITER FlyDSL kernel (gfx950). -GPU-free: mocks the platform (gfx950) and the ``flydsl`` package check, then +GPU-free: mocks the platform (gfx950) and the aiter fused-MoE enable flag, then exercises the oracle so the FlyDSL backend is auto-picked when usable (including under expert parallelism, since apply() forwards the expert_map as aiter's -expert_mask) and skipped (native fallback) when the device/package is missing. +expert_mask) and skipped (native fallback) when the device is unsupported or the +aiter runtime is disabled. """ import dataclasses @@ -48,6 +49,8 @@ kMxfp8Static, ) +from vllm._aiter_ops import rocm_aiter_ops # noqa: E402 # isort: skip + _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe" @@ -82,8 +85,8 @@ def _gfx950(): ) -def _flydsl_installed(present: bool): - return patch(f"{_AITER_MOD}.is_aiter_mxfp8_moe_available", return_value=present) +def _aiter_moe_enabled(present: bool): + return patch.object(rocm_aiter_ops, "is_fused_moe_enabled", return_value=present) def test_aiter_mxfp8_registered(): @@ -109,13 +112,18 @@ def test_ep_supported(ep_size): @pytest.mark.parametrize( "present,ep_size,supported,reason_substr", [ - (True, 1, True, None), # gfx950 + flydsl + TP -> selectable - (True, 2, True, None), # gfx950 + flydsl + EP -> selectable (expert_mask) - (False, 1, False, "flydsl package"), # package missing -> clear reason + (True, 1, True, None), # gfx950 + aiter MoE + TP -> selectable + (True, 2, True, None), # gfx950 + aiter MoE + EP -> selectable (expert_mask) + ( + False, + 1, + False, + "does not support current device", + ), # disabled -> not selected ], ) def test_is_supported_config(present, ep_size, supported, reason_substr): - with _gfx950(), _flydsl_installed(present): + with _gfx950(), _aiter_moe_enabled(present): ok, reason = AiterMxfp8Experts.is_supported_config( AiterMxfp8Experts, _config(ep_size), @@ -130,8 +138,8 @@ def test_is_supported_config(present, ep_size, supported, reason_substr): def test_explicit_moe_backend_aiter(): """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear - ValueError when the flydsl package is missing.""" - with _gfx950(), _flydsl_installed(True): + ValueError when the aiter MoE runtime is disabled.""" + with _gfx950(), _aiter_moe_enabled(True): assert ( _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) is AiterMxfp8Experts @@ -142,18 +150,17 @@ def test_explicit_moe_backend_aiter(): ) with ( _gfx950(), - _flydsl_installed(False), - pytest.raises(ValueError, match="flydsl package"), + _aiter_moe_enabled(False), + pytest.raises(ValueError, match="does not support current device"), ): _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) def test_gfx950_picks_aiter(): - """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins.""" - # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require VLLM_ROCM_USE_AITER=1 + """Auto-select on real ROCm hardware with aiter MoE enabled -> FlyDSL wins.""" with ( patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True), - _flydsl_installed(True), + _aiter_moe_enabled(True), ): backend, experts_cls = select_mxfp8_moe_backend(_config()) assert backend is Fp8MoeBackend.AITER_MXFP8 diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index c12a72a09c0a..0463468eae96 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/tests/models/quantization/test_mxfp8.py @@ -111,11 +111,6 @@ def test_mxfp8_aiter_requires_swigluoai_activation( "_supports_current_device", staticmethod(lambda: True), ) - monkeypatch.setattr( - aiter_mxfp8_moe, - "is_aiter_mxfp8_moe_available", - lambda: True, - ) config = FusedMoEConfig( num_experts=8, @@ -164,11 +159,6 @@ def test_mxfp8_aiter_requires_swigluoai_params( "_supports_current_device", staticmethod(lambda: True), ) - monkeypatch.setattr( - aiter_mxfp8_moe, - "is_aiter_mxfp8_moe_available", - lambda: True, - ) config = FusedMoEConfig( num_experts=8, @@ -218,11 +208,6 @@ def test_mxfp8_aiter_accepts_swigluoai_params( "_supports_current_device", staticmethod(lambda: True), ) - monkeypatch.setattr( - aiter_mxfp8_moe, - "is_aiter_mxfp8_moe_available", - lambda: True, - ) config = FusedMoEConfig( num_experts=8, diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 93a4088fc4f7..c15a1a13d42e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -11,6 +11,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( @@ -24,41 +25,6 @@ _AITER_SWIGLU_BETA = 1.0 -def is_aiter_mxfp8_moe_available() -> bool: - """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` - package is importable, AND the installed aiter carries the mxfp8 FlyDSL - 2-stage support from ROCm/aiter#3811. - - ``flydsl`` and ``aiter`` are separate packages, so ``is_flydsl_available()`` - (flydsl pkg + arch) is necessary but not sufficient: an older aiter without - #3811 still ships the flydsl pkg and the ``aiter.ops.flydsl`` module but a - broken/missing ``per_1x32 + fp8`` 2-stage path. Without this extra gate a - nightly lacking #3811 would wrongly select FlyDSL instead of falling back to - the native Triton dot_scaled path. #3811 added no probe-able public symbol, - so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails - closed (returns False -> triton dot_scaled), which is always safe.""" - if not (current_platform.is_rocm() and current_platform.supports_mx()): - return False - try: - import os - - import aiter - from aiter.ops.flydsl.utils import is_flydsl_available - - if not is_flydsl_available(): - return False - return os.path.exists( - os.path.join( - os.path.dirname(aiter.__file__), - "configs", - "model_configs", - "minimax_m3_mxfp8_tuned_fmoe.csv", - ) - ) - except Exception: - return False - - class AiterMxfp8Experts(Mxfp8TritonExpertsBase): """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950).""" @@ -79,10 +45,7 @@ def expects_unquantized_inputs(self) -> bool: @staticmethod def _supports_current_device() -> bool: - # Device capability only (gfx950 / MX-capable ROCm). The flydsl package - # check lives in is_supported_config so a missing package is reported - # distinctly from an unsupported device. - return current_platform.is_rocm() and current_platform.supports_mx() + return current_platform.supports_mx() and rocm_aiter_ops.is_fused_moe_enabled() @staticmethod def _supports_parallel_config(moe_parallel_config) -> bool: @@ -95,12 +58,6 @@ def is_supported_config( is_supported, reason = super().is_supported_config( cls, moe_config, weight_key, activation_key, activation_format ) - # _supports_current_device() only gates on the device; surface a clear - # reason when the device is fine but the flydsl package is missing. - if is_supported and not is_aiter_mxfp8_moe_available(): - return False, ( - "kernel requires the aiter flydsl package, which is not installed" - ) if ( is_supported and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE From 2c88fb131c7ae0be01907cd8c276911db5e7aad4 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 17 Sep 2026 18:24:58 -0600 Subject: [PATCH 447/774] [Bugfix] DiffusionGemma: hand out stashed logprobs only on the committing step (#57414) --- vllm/model_executor/models/diffusion_gemma.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index bd1d14b6f6dc..af0cb35168dc 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -1312,7 +1312,6 @@ def __call__( is_committing = states.is_encoder_phase[decode_slots].clone() slots_np = input_batch.idx_mapping_np[:num_reqs] - is_decode_np = per_req_nlogits_np > 0 max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in @@ -1396,13 +1395,16 @@ def __call__( # stashed logprobs and attach to SamplerOutput. logprobs_tensors = None if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs: + committing_slots = set( + decode_slots_np[is_committing.cpu().numpy()].tolist() + ) parts_ids, parts_lp, parts_ranks = [], [], [] cu_gen: list[int] = [] flat_offset = 0 for i in range(num_reqs): cu_gen.append(flat_offset) slot = int(slots_np[i]) - if is_decode_np[i] and slot in self._pending_logprobs: + if slot in committing_slots and slot in self._pending_logprobs: lp = self._pending_logprobs.pop(slot) parts_ids.append(lp.logprob_token_ids) parts_lp.append(lp.logprobs) From 64563d0ec4b761745d5babbc5750b263f0783926 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 17 Sep 2026 19:00:10 -0700 Subject: [PATCH 448/774] [Bugfix][DSv4.1] Preserve NaN-scored candidate block indices (#57454) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- .../layers/test_mla_short_prefill_indexer.py | 22 ++++++++++++++++++- .../kernels/attention/dsa/candidate_blocks.py | 3 ++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 41a06d0546f0..882b745ece79 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -365,6 +365,26 @@ def test_select_candidate_blocks_tolerates_empty_rows(): assert (out[0] >= 0).all() and (out[2, :2] >= 0).all() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("score", [0.0, float("nan")]) +@pytest.mark.parametrize("start", [0, 3]) +def test_candidate_selection_keeps_valid_blocks_with_nan_scores(score, start): + """NaN warmup scores must not turn real blocks into repeated sparse padding.""" + logits = torch.full((3, 64), score, device="cuda") + starts = torch.full((3,), start, device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 17, 35], device="cuda", dtype=torch.int32) + out = torch.empty(3, 8, device="cuda", dtype=torch.int32) + + sparse_indexer._select_candidate_blocks(logits, starts, starts + lengths, 8, 8, out) + + for row, count in enumerate([0, 3, 5]): + valid = out[row][out[row] >= 0].sort().values + torch.testing.assert_close( + valid, torch.arange(count, device="cuda", dtype=torch.int32) + ) + assert (out[row] == -1).sum().item() == 8 - count + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( "width,block_size,k,decode", @@ -413,7 +433,7 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( top = reduced.topk(min(k, nblocks), dim=-1) expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32) expected[:, : top.indices.shape[1]] = torch.where( - top.values > -torch.inf, top.indices, -1 + top.values != -torch.inf, top.indices, -1 ).int() actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2] select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat) diff --git a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py index efc7c418cad6..b9cdc923b14a 100644 --- a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py +++ b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py @@ -66,9 +66,10 @@ def _store_candidates_kernel( cols = tl.program_id(1) * TILE + tl.arange(0, TILE) value = tl.load(values + row * k + cols, cols < k, other=-float("inf")) index = tl.load(indices + row * k + cols, cols < k, other=-1) + # NaN scores can occur during warmup; only -inf denotes padding. tl.store( output + row * out_stride_row + cols * out_stride_col, - tl.where(value > -float("inf"), index, -1), + tl.where(value != -float("inf"), index, -1), cols < OUT_K, ) From 2bbdfcfcfea79db5a3a4a32a573e0998fd6a58eb Mon Sep 17 00:00:00 2001 From: Chauncey Date: Fri, 18 Sep 2026 10:31:35 +0800 Subject: [PATCH 449/774] [Perf][GLM5.3-Flash] Use cooperative top-k for small GLM decode batches (#57327) Signed-off-by: chaunceyjiang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_sparse_indexer_topk_dispatch.py | 41 +++++++++++++++++++ vllm/models/glm5next/nvidia/sparse_indexer.py | 23 ++++++++++- 2 files changed, 63 insertions(+), 1 deletion(-) create mode 100644 tests/models/glm5next/test_sparse_indexer_topk_dispatch.py diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py new file mode 100644 index 000000000000..83dfdc9f6ffd --- /dev/null +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Focused unit tests for the GLM decode top-k dispatch.""" + +import pytest +import torch + +from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA-only dispatch" +) + + +@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"]) +@torch.inference_mode() +def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int): + logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda") + assert _use_cooperative_topk(logits, 512, num_rows) + + +@pytest.mark.parametrize( + ("num_rows", "select_k", "stride_padding"), + [ + (65, 512, 0), + (1, 256, 0), + (1, 512, 1), + ], + ids=["too_many_rows", "unsupported_k", "unaligned_stride"], +) +@torch.inference_mode() +def test_cooperative_topk_falls_back_on_unsupported_batch( + num_rows: int, select_k: int, stride_padding: int +): + logits = torch.empty( + num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda" + ) + if stride_padding: + logits = logits[:, :8192] + assert not _use_cooperative_topk(logits, select_k, num_rows) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 1fc42c9e4114..9dab659d4cc0 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -38,6 +38,22 @@ logger = init_logger(__name__) + +def _use_cooperative_topk( + logits: torch.Tensor, + select_k: int, + num_rows: int, +) -> bool: + return ( + current_platform.is_cuda() + and select_k in (512, 1024, 2048) + and num_rows <= 64 + and logits.stride(0) % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + + # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -573,7 +589,12 @@ def sparse_attn_indexer_kpool( (topk_workspace,) = workspace_manager.get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) - torch.ops._C.persistent_topk( + topk_op = ( + torch.ops._C.cooperative_topk + if _use_cooperative_topk(logits, select_k, num_rows) + else torch.ops._C.persistent_topk + ) + topk_op( logits, seq_lens, topk_dst, From 7b942936276d59cc1912185a442b97dea8d2386c Mon Sep 17 00:00:00 2001 From: Ren Yuzhou <54501155+yuzhouo7@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:35:23 +0800 Subject: [PATCH 450/774] [Bugfix][ROCm][KV Offload] Use private pinned tensors for CPU KV offload (#57160) Signed-off-by: Yuzhou Ren --- tests/v1/kv_offload/test_factory.py | 49 +++++++++++++++++++---------- vllm/v1/kv_offload/cpu/spec.py | 2 +- 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 16539677c490..01ef756cae53 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -275,11 +275,12 @@ def fake_region_ctor(**kwargs): @pytest.mark.parametrize("world_size", [2, 4, 8]) def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: int): - # On shared-region (CUDA-alike) platforms the default spec now honors + # On CUDA the default spec honors # replicated layout: a single MLA copy (num_copies=1), matching tiering. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * 8, @@ -295,19 +296,19 @@ def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: in assert spec.num_chunks == 8 +@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"]) @pytest.mark.parametrize("world_size", [2, 4, 8]) def test_cpu_spec_replicated_disabled_without_shared_region( - monkeypatch, world_size: int + monkeypatch, world_size: int, rocm: bool ): - # Data-loss guard: non-CUDA-alike platforms keep a per-rank private pinned + # Data-loss guard: platforms without shared regions keep a private pinned # tensor (no shared medium), so replicated layout MUST stay off. Otherwise # the rank-0 writer gate would ack rank>0 stores without writing, leaving # those private buffers empty and corrupting subsequent loads. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module - monkeypatch.setattr( - cpu_spec_module.current_platform, "is_cuda_alike", lambda: False - ) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2, @@ -324,17 +325,26 @@ def test_cpu_spec_replicated_disabled_without_shared_region( @pytest.mark.parametrize("config_replicated", [True, False]) -@pytest.mark.parametrize("cuda_alike", [True, False]) +@pytest.mark.parametrize( + ("cuda_alike", "rocm", "shared_region"), + [(True, False, True), (True, True, False), (False, False, False)], + ids=["cuda", "rocm", "non-cuda-alike"], +) def test_cpu_spec_replicated_layout_truth_matrix( - monkeypatch, cuda_alike: bool, config_replicated: bool + monkeypatch, + cuda_alike: bool, + rocm: bool, + shared_region: bool, + config_replicated: bool, ): # replicated_layout is enabled iff the config gate passes AND the deployment - # actually allocates on the shared region (CUDA-alike). + # actually allocates on the shared region. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module monkeypatch.setattr( cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike ) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * 8, @@ -344,10 +354,10 @@ def test_cpu_spec_replicated_layout_truth_matrix( ) assert isinstance(spec, CPUOffloadingSpec) - assert spec.replicated_layout is (cuda_alike and config_replicated) + assert spec.replicated_layout is (shared_region and config_replicated) -def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch): +def test_cpu_spec_create_worker_uses_mmap_on_cuda(monkeypatch): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT @@ -371,6 +381,7 @@ def fake_worker_ctor(**kwargs): return MagicMock() monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) monkeypatch.setattr( @@ -386,7 +397,10 @@ def fake_worker_ctor(**kwargs): assert worker_calls[0]["mmap_region"] is region -def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch): +@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"]) +def test_cpu_spec_create_worker_uses_tensor_path_without_shared_region( + monkeypatch, rocm +): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4) @@ -403,15 +417,14 @@ def fake_worker_ctor(**kwargs): worker_calls.append(kwargs) return MagicMock() - monkeypatch.setattr( - cpu_spec_module.current_platform, "is_cuda_alike", lambda: False - ) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) spec.create_worker(MagicMock()) - # Non-CUDA-alike platforms keep the per-rank pinned-tensor path. + # ROCm and non-CUDA-alike platforms use the per-rank pinned-tensor path. assert region_calls == [] assert worker_calls[0]["mmap_region"] is None @@ -420,7 +433,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module # worker_kv_bytes_per_block=0 yields num_chunks=0; a zero-byte region cannot - # be mmap'd, so even on CUDA-alike this must fall back to the tensor path. + # be mmap'd, so even on CUDA this must fall back to the tensor path. spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4) assert isinstance(spec, CPUOffloadingSpec) assert spec.num_chunks == 0 @@ -429,6 +442,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): worker_calls: list[dict[str, Any]] = [] monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) monkeypatch.setattr( cpu_spec_module, "SharedOffloadRegion", @@ -461,6 +475,7 @@ def test_cpu_spec_create_worker_rank_assignment( import vllm.v1.kv_offload.cpu.spec as cpu_spec_module monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * 8, diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 00b36e863a07..fddf56608590 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -162,7 +162,7 @@ def get_manager(self) -> OffloadingManager: def _uses_shared_region(self) -> bool: """Whether the worker CPU buffer is the shared mmap region (vs a private per-rank tensor); replicated-layout dedup is gated on this being True.""" - return current_platform.is_cuda_alike() + return current_platform.is_cuda_alike() and not current_platform.is_rocm() def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: mmap_region: SharedOffloadRegion | None = None From 70df48dc3d01fe1bb3206f7d8a832c7e8de348e8 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:42:54 +0200 Subject: [PATCH 451/774] [Bugfix][KV Cache][GLM-5.3-Flash] Disable slot mapping kernel for the kpool tail buffer (#57317) Signed-off-by: simondanielsson --- .../attention/test_kpool_tail_slot_mapping.py | 24 +++++ .../test_gpu_kpool_tail_slot_mapping.py | 98 +++++++++++++++++++ vllm/v1/kv_cache_interface.py | 4 + 3 files changed, 126 insertions(+) create mode 100644 tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 0966f9b20a6b..cb39e5403719 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -32,6 +32,7 @@ ) from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides from vllm.v1.kv_cache_layout import KVCacheLayout +from vllm.v1.worker.block_table import get_block_table_width KPOOL = 4 @@ -55,6 +56,29 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): assert content_stride == 1 +def test_tail_spec_opts_out_of_generic_slot_mapping(): + """The tail row is one block wide (padded to the block-table alignment), so + the generic kernel's ``pos // kpool`` column index runs off the end of the + allocation for long prompts. The spec must opt out of it entirely.""" + spec = KpoolTailSpec( + block_size=KPOOL, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=KPOOL, + ) + max_len = 1 << 20 + width = get_block_table_width( + spec.max_num_blocks_per_req(None, max_len), + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) + + assert width * KPOOL < max_len + assert spec.uses_slot_mapping is False + + def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the request's single KpoolTailManager block, the remaining columns are never diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py new file mode 100644 index 000000000000..8a7445f7d49a --- /dev/null +++ b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU regression test for the kpool tail group's slot mapping. + +The kpool tail cache holds one circular block per request, so its block table +row is orders of magnitude shorter than the sequence it serves. +`_compute_slot_mappings_kernel` indexes block tables by +`position // kernel_block_size` with an effectively unmasked load, so +this leads to a read of neighboring requests' rows (and for the later requests, outside +the buffer entirely). +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import KpoolTailSpec +from vllm.v1.worker.block_table import get_block_table_width +from vllm.v1.worker.gpu.block_table import BlockTables + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="requires CUDA or ROCm", +) + +KPOOL = 4 +MAX_MODEL_LEN = 1 << 20 +MLA_BLOCK_SIZE = 640 +MAX_NUM_REQS = 4 +POISON_BLOCK_ID = 1 << 20 + + +def make_tail_spec(): + return KpoolTailSpec( + block_size=KPOOL, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=KPOOL, + ) + + +def test_kpool_tail_group_never_position_indexes_its_block_table(): + device = torch.device("cuda") + spec = make_tail_spec() + tail_width = get_block_table_width( + spec.max_num_blocks_per_req(None, MAX_MODEL_LEN), + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) + mla_width = get_block_table_width(MAX_MODEL_LEN // MLA_BLOCK_SIZE, MLA_BLOCK_SIZE) + + # Positions chosen against the tail row: `in_row` still lands inside + # request 0's own row, `next_row` lands on request 1's poisoned row, and + # `past_table` is past the whole [MAX_NUM_REQS, tail_width] allocation. + in_row = KPOOL - 1 + next_row = tail_width * KPOOL + past_table = MAX_NUM_REQS * tail_width * KPOOL + KPOOL + assert past_table < MAX_MODEL_LEN + positions = [0, in_row, next_row, past_table] + + block_tables = BlockTables( + block_sizes=[MLA_BLOCK_SIZE, spec.block_size], + max_num_reqs=MAX_NUM_REQS, + max_num_batched_tokens=len(positions), + max_num_blocks_per_group=[mla_width, tail_width], + device=device, + kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size], + slot_mapping_enabled=[True, spec.uses_slot_mapping], + ) + mla_blocks = list(range(1, mla_width + 1)) + block_tables.append_block_ids(0, (mla_blocks, [7]), overwrite=True) + block_tables.append_block_ids(1, ([2], [POISON_BLOCK_ID]), overwrite=True) + block_tables.apply_staged_writes() + + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor( + [0, len(positions)], dtype=torch.int32, device=device + ) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + torch.tensor(positions, dtype=torch.int64, device=device), + num_tokens_padded=len(positions), + ) + torch.accelerator.synchronize() + + tail_slots = slot_mappings[1].tolist() + assert tail_slots == [PAD_SLOT_ID] * len(positions) + assert POISON_BLOCK_ID * KPOOL not in tail_slots + + mla_slots = slot_mappings[0].tolist() + assert mla_slots == [ + mla_blocks[p // MLA_BLOCK_SIZE] * MLA_BLOCK_SIZE + p % MLA_BLOCK_SIZE + for p in positions + ] diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 17771d9ae67a..6595be14d21d 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -987,6 +987,10 @@ def is_uniform_with_collection( def prefix_cacheable(self) -> bool: return False + @property + def uses_slot_mapping(self) -> bool: + return False + @dataclass(frozen=True) class MambaSpec(KVCacheSpec): From 2909ad8fa472c04f7db8330f0e751ff43ae76965 Mon Sep 17 00:00:00 2001 From: Jason <31175216+thegoldenflow@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:07:25 -0400 Subject: [PATCH 452/774] [Bugfix] Attach request-level tools to existing system message in DeepSeek V4 Python renderer (#51856) Signed-off-by: Jason Yao Co-authored-by: Chauncey --- .../fixtures/deepseek_v4/test_output_1.txt | 4 +- tests/tokenizers_/test_deepseek_v4.py | 143 ++++++++++++++++++ vllm/tokenizers/deepseek_v4.py | 14 +- 3 files changed, 157 insertions(+), 4 deletions(-) diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt b/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt index dbd823476c1c..7e3c9bd5a394 100644 --- a/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt +++ b/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt @@ -1,4 +1,4 @@ -<｜begin▁of▁sentence｜> +<｜begin▁of▁sentence｜>You are a helpful assistant. ## Tools @@ -26,7 +26,7 @@ Otherwise, output directly after with tool calls or final response. {"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}} You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. -You are a helpful assistant.<｜User｜>What's the weather in Beijing?<｜Assistant｜>The user wants to know the weather in Beijing. I should use the get_weather tool. +<｜User｜>What's the weather in Beijing?<｜Assistant｜>The user wants to know the weather in Beijing. I should use the get_weather tool. <｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index a55aa57666dd..f7bc5909ba31 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -485,3 +485,146 @@ def convert_tokens_to_ids(self, token: str) -> int: ] ) assert image_sentinel_mask(ids).tolist() == [False, True, True, False, False] + + +def _request_tools(): + return [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + + +def _encode_reference(messages): + from vllm.tokenizers.deepseek_v4_encoding import encode_messages + + return encode_messages(messages, thinking_mode="thinking", reasoning_effort="low") + + +def test_deepseek_v4_attaches_request_tools_to_existing_system_message(): + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Weather in Paris?"}, + ] + + prompt = _tokenizer().apply_chat_template( + messages, + tools=_request_tools(), + tokenize=False, + thinking=True, + reasoning_effort="low", + ) + + expected = _encode_reference( + [{**messages[0], "tools": _request_tools()}, messages[1]] + ) + assert prompt == expected + assert prompt.startswith("<｜begin▁of▁sentence｜>You are helpful.\n\n## Tools") + assert prompt.count("## Tools") == 1 + + +def test_deepseek_v4_synthetic_system_only_when_no_system_message(): + """Without a system message, request tools still get a synthetic leading + system entry.""" + messages = [{"role": "user", "content": "Weather in Paris?"}] + + prompt = _tokenizer().apply_chat_template( + messages, + tools=_request_tools(), + tokenize=False, + thinking=True, + reasoning_effort="low", + ) + + expected = _encode_reference( + [{"role": "system", "tools": _request_tools()}, messages[0]] + ) + assert prompt == expected + assert prompt.startswith("<｜begin▁of▁sentence｜>\n\n## Tools") + + +def test_deepseek_v4_request_tools_override_system_message_tools(): + message_tools = [ + { + "type": "function", + "function": { + "name": "message_level_tool", + "description": "Should not be rendered", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + system_message = { + "role": "system", + "content": "You are helpful.", + "tools": message_tools, + } + user_message = {"role": "user", "content": "Weather in Paris?"} + + prompt = _tokenizer().apply_chat_template( + [system_message, user_message], + tools=_request_tools(), + tokenize=False, + thinking=True, + reasoning_effort="low", + ) + + expected = _encode_reference( + [{**system_message, "tools": _request_tools()}, user_message] + ) + assert prompt == expected + assert '"name": "get_weather"' in prompt + assert "message_level_tool" not in prompt + + +def test_deepseek_v4_attaches_request_tools_to_first_system_message_only(): + """Only the first system message, at any index, gets the request tools.""" + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "system", "content": "First system."}, + {"role": "system", "content": "Second system."}, + {"role": "user", "content": "Weather in Paris?"}, + ] + + prompt = _tokenizer().apply_chat_template( + messages, + tools=_request_tools(), + tokenize=False, + thinking=True, + reasoning_effort="low", + ) + + expected = _encode_reference( + [ + messages[0], + {**messages[1], "tools": _request_tools()}, + messages[2], + messages[3], + ] + ) + assert prompt == expected + assert "First system.\n\n## Tools" in prompt + assert prompt.count("## Tools") == 1 + + +def test_deepseek_v4_request_tools_do_not_mutate_caller_messages(): + import copy + + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Weather in Paris?"}, + ] + snapshot = copy.deepcopy(messages) + + _tokenizer().apply_chat_template(messages, tools=_request_tools(), tokenize=False) + + assert messages == snapshot diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index a5d5020e941b..532ca0df7c9b 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -35,8 +35,18 @@ def apply_chat_template( conversation = kwargs.get("conversation", messages) messages = conversation.copy() if tools is not None and len(tools) > 0: - messages.insert(0, {"role": "system"}) - messages[0]["tools"] = tools # type: ignore[typeddict-unknown-key] + # Match the Rust renderer: request tools attach to the first + # system message; synthesize one only when none exists. + system_idx = next( + (i for i, m in enumerate(messages) if m.get("role") == "system"), + None, + ) + if system_idx is None: + messages.insert(0, {"role": "system"}) + system_idx = 0 + else: + messages[system_idx] = copy.copy(messages[system_idx]) + messages[system_idx]["tools"] = tools # type: ignore[typeddict-unknown-key] reasoning_effort = kwargs.get("reasoning_effort") if not isinstance(reasoning_effort, str): From bf13ecc2f7c4d69b59854f7cef37547be58805e5 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Thu, 17 Sep 2026 20:14:36 -0700 Subject: [PATCH 453/774] [Perf][MRV2] Share token-to-request mappings across KV cache groups (#57102) Signed-off-by: Juntian Liu Co-authored-by: Woosuk Kwon --- vllm/v1/worker/gpu/attn_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index eccfdbabcdfd..89f54d4d1140 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -426,6 +426,7 @@ def build_attn_metadata( seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] attn_metadata: dict[str, Any] = {} + token_to_req_indices: torch.Tensor | None = None num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): if not attn_groups[i]: @@ -472,6 +473,7 @@ def build_attn_metadata( mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, req_idx=req_idx, + _token_to_req_indices_cache=token_to_req_indices, **common_attn_metadata_extra_kwargs, ) @@ -497,6 +499,7 @@ def build_attn_metadata( ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata + token_to_req_indices = common_attn_metadata._token_to_req_indices_cache return attn_metadata From d12c2768530a4da8d7c7d4a4f5d6adfa351675cd Mon Sep 17 00:00:00 2001 From: Mustafa YILDIRIM <2341474+mustafayildirim@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:36:55 -0700 Subject: [PATCH 454/774] [Bugfix][ROCm] Alias SparseAttnIndexerKpool.forward_cuda to forward_native (GLM-5.3-Flash boot crash) (#57425) Signed-off-by: Mustafa YILDIRIM Co-authored-by: Hermes Agent Co-authored-by: Tyler Michael Smith --- vllm/models/glm5next/amd/sparse_indexer.py | 25 +++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index f47c9164bc8b..c9b68e973f84 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -665,7 +665,7 @@ def __init__( self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache - def forward_native( + def forward_hip( self, hidden_states: torch.Tensor, q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -727,3 +727,26 @@ def forward_native( self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, ) + + def forward_native( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + *, + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.forward_hip( + hidden_states, + q_quant, + k, + weights, + gate_score=gate_score, + compress_ape=compress_ape, + index_kpool=index_kpool, + positions=positions, + ) From dee37d89115db4c94a820a79a78a7828e141c910 Mon Sep 17 00:00:00 2001 From: lazymio Date: Fri, 18 Sep 2026 12:32:39 +0800 Subject: [PATCH 455/774] [Frontend] Add `--tool-strict-level` for server-side control for structural tag activation (#56268) Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Bugen Zhao Signed-off-by: lazymio Signed-off-by: Bugen Zhao --- docs/features/tool_calling.md | 40 ++- rust/src/chat/src/backend/hf.rs | 2 + rust/src/chat/src/backend/mod.rs | 3 +- rust/src/chat/src/lib.rs | 19 +- rust/src/chat/src/output/default/mod.rs | 15 +- .../chat/src/output/default/structural_tag.rs | 220 +++++++++++++-- rust/src/chat/src/parser/mod.rs | 60 +++++ rust/src/chat/tests/chat.rs | 1 + rust/src/chat/tests/roundtrip.rs | 1 + rust/src/cmd/src/cli.rs | 16 ++ rust/src/cmd/src/cli/tests.rs | 6 + .../examples/external_engine_openai_qwen.rs | 2 + rust/src/server/src/config.rs | 3 + rust/src/server/src/grpc/tests.rs | 2 + rust/src/server/src/lib.rs | 3 +- rust/src/server/src/render.rs | 14 +- .../server/src/routes/http_client_tests.rs | 1 + rust/src/server/src/routes/tests.rs | 1 + .../test_structural_tag_registry.py | 254 +++++++++++++++++- vllm/entrypoints/anthropic/serving.py | 2 + vllm/entrypoints/cohere/serving.py | 2 + vllm/entrypoints/generate/api_router.py | 4 + .../launchers/api_server/app_state.py | 2 + vllm/entrypoints/launchers/cli_args.py | 8 + .../entrypoints/launchers/render/app_state.py | 2 + .../openai/chat_completion/serving.py | 2 + vllm/entrypoints/openai/responses/serving.py | 2 + vllm/parser/abstract_parser.py | 4 + vllm/parser/parser_manager.py | 10 + vllm/renderers/online_derenderer.py | 2 + vllm/renderers/online_renderer.py | 2 + vllm/tool_parsers/abstract_tool_parser.py | 3 + vllm/tool_parsers/hy_v4_tool_parser.py | 4 + vllm/tool_parsers/structural_tag_registry.py | 59 +++- vllm/tool_parsers/tool_strict_level.py | 32 +++ 35 files changed, 742 insertions(+), 61 deletions(-) create mode 100644 vllm/tool_parsers/tool_strict_level.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 9628bc431a1e..fb70c5024835 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -109,18 +109,22 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t ## Constrained Decoding Behavior -Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode and the per-tool `strict` field: +Structural-tag parsers resolve call obligation, grammar activation, and argument-schema enforcement separately. With structural-tag enforcement enabled: -| `tool_choice` value | Schema-constrained decoding | Behavior | +| `tool_choice` value | Call obligation | Structural-tag activation | | --- | --- | --- | -| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. | -| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. | -| `"auto"` | Only when `strict: true` is set on at least one tool | Structural-tag parsers constrain tool-call arguments when a tool opts in with `strict: true`. Without it, the model generates freely and tool calls are extracted from raw text. | -| `"none"` | N/A | No tool calls are produced. | +| Named function | Call the selected function | Always | +| `"required"` | Produce at least one tool call | Always | +| `"auto"` | Tool calls are optional | When at least one tool sets `strict: true`, or `--tool-strict-level` is `function` or `parameter` | +| `"none"` | Tool calling is disabled | Disabled | + +When a structural tag applies, each tool's declared parameter schema is enforced only when that tool sets `strict: true` or the server uses `--tool-strict-level parameter`. Tools with omitted or false `strict` receive broad argument-syntax constraints at levels `auto` and `function`, including for required and named calls. + +For parsers using schema-derived JSON constraints, required and named calls continue to enforce the declared parameter schemas. ### Strict Mode -For `tool_choice="required"` or named function calling, structural-tag constraints are always applied regardless of the `strict` field. For `tool_choice="auto"`, setting `strict: true` on at least one tool opts in to structural-tag constraints; without it, the model generates freely and tool calls are extracted from raw text. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages. +Structural-tag constraints have two layers. The tool-call *envelope* (markup and function name) is constrained whenever a structural tag applies: always for `tool_choice="required"` and named function calling, and for `tool_choice="auto"` once at least one tool sets `strict: true` or the server raises the floor. The *argument schema* of an individual tool is pinned only when that tool sets `strict: true` (or under `--tool-strict-level parameter`); a tool that omits `strict` keeps its arguments unconstrained, even when another tool in the same request is strict. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages. For best compatibility with strict schema enforcement, define tool parameter schemas in the OpenAI strict-schema style: @@ -134,6 +138,22 @@ vLLM also provides a global toggle via the `VLLM_ENFORCE_STRICT_TOOL_CALLING` en VLLM_ENFORCE_STRICT_TOOL_CALLING=false vllm serve ... ``` +### Server-Side Strictness Floor + +Most OpenAI-compatible clients and agent frameworks never set `strict` on their tools, so with `tool_choice="auto"` the model generates tool calls without any grammar and malformed markup can leak into the response. The `--tool-strict-level` option lets the server operator raise the floor for every request that carries tools, independently of what the client declares: + +| Value | Behavior | +| --- | --- | +| `auto` (default) | Follow the request's tool choice and per-tool strictness. Required/named choices activate structural tags; `tool_choice="auto"` activates them when at least one tool sets `strict: true`. | +| `function` | Constrain the tool-call envelope (markup and the function name) for every request with tools, leaving argument contents free unless the client marked the tool `strict: true`. | +| `parameter` | Additionally pin argument schemas for every tool, as if every tool had `strict: true`. | + +```bash +vllm serve ... --tool-strict-level function +``` + +The floor never relaxes a constraint the request would already receive: tools the client marked `strict: true` keep their schemas at every level. With `tool_choice="auto"`, the grammar does not force a tool call; a plain text response stays valid. `VLLM_ENFORCE_STRICT_TOOL_CALLING=false` disables structural tags entirely and takes precedence over this option. + ## Automatic Function Calling To enable this feature, you should set the following flags: @@ -152,7 +172,7 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template! !!! note - With `tool_choice="auto"`, schema-level constraint requires both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`. When these conditions are met and the selected parser supports structural tags, vLLM constrains tool-call arguments. Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema. + With `tool_choice="auto"`, structural-tag constraints require both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`, or a server-side floor set via `--tool-strict-level`. When these conditions are met and the selected parser supports structural tags, vLLM constrains the tool-call envelope and pins the argument schema of each tool that sets `strict: true` (or of every tool under `--tool-strict-level parameter`). Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema. ### Hermes Models (`hermes`) @@ -162,8 +182,8 @@ All Nous Research Hermes-series models newer than Hermes 2 Pro should be support * `NousResearch/Hermes-2-Theta-*` * `NousResearch/Hermes-3-*` -_Note that the Hermes 2 **Theta** models are known to have degraded tool call quality and capabilities due to the merge -step in their creation_. +*Note that the Hermes 2 **Theta** models are known to have degraded tool call quality and capabilities due to the merge +step in their creation*. Flags: `--tool-call-parser hermes` diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index af9af8b51ca4..07f0b2dce815 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -134,6 +134,7 @@ impl ChatBackend for HfChatBackend { self.tokenizer.clone(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } @@ -343,6 +344,7 @@ mod tests { let error = match backend.new_chat_output_processor( &mut request, NewChatOutputProcessorOptions { + tool_strict_level: crate::ToolStrictLevel::Auto, tool_call_parser: &ParserSelection::Explicit("json".to_string()), reasoning_parser: &ParserSelection::Auto, }, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index af4579159e82..a0343e53f327 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -13,7 +13,7 @@ use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo}; use crate::output::DynChatOutputProcessor; use crate::renderer::DynChatRenderer; use crate::request::ChatRequest; -use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; +use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection, ToolStrictLevel}; pub mod hf; @@ -21,6 +21,7 @@ pub mod hf; pub struct NewChatOutputProcessorOptions<'a> { pub tool_call_parser: &'a ParserSelection, pub reasoning_parser: &'a ParserSelection, + pub tool_strict_level: ToolStrictLevel, } /// Minimal prompt-processing backend needed by `vllm-chat`. diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 69d39ca78ed3..0273396ebd85 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -30,7 +30,7 @@ pub use parser::reasoning::{ ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory, }; pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; -pub use parser::{ParserSelection, validate_parser_overrides}; +pub use parser::{ParserSelection, ToolStrictLevel, validate_parser_overrides}; pub use reasoning::EffortValue; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ @@ -74,6 +74,8 @@ pub struct ChatRequestProcessor { tool_call_parser: ParserSelection, /// Reasoning parser selection used when preparing generation requests. reasoning_parser: ParserSelection, + /// Server-side floor for tool-call structural tags. + tool_strict_level: ToolStrictLevel, } impl ChatRequestProcessor { @@ -85,6 +87,7 @@ impl ChatRequestProcessor { model_dtype: Some(model_dtype), tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + tool_strict_level: ToolStrictLevel::Auto, } } @@ -95,6 +98,7 @@ impl ChatRequestProcessor { model_dtype: None, tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + tool_strict_level: ToolStrictLevel::Auto, } } @@ -109,6 +113,12 @@ impl ChatRequestProcessor { self } + /// Configure the server-side floor for tool-call structural tags. + pub fn with_tool_strict_level(mut self, tool_strict_level: ToolStrictLevel) -> Self { + self.tool_strict_level = tool_strict_level; + self + } + async fn finalize_rendered_prompt( &self, request: &ChatRequest, @@ -207,6 +217,7 @@ impl ChatRequestProcessor { NewChatOutputProcessorOptions { tool_call_parser: &self.tool_call_parser, reasoning_parser: &self.reasoning_parser, + tool_strict_level: self.tool_strict_level, }, )?; let text_request = self.prepare_text_request(request).await?; @@ -255,6 +266,12 @@ impl ChatLlm { self } + /// Set the server-side floor for tool-call structural tags. + pub fn with_tool_strict_level(mut self, tool_strict_level: ToolStrictLevel) -> Self { + self.processor.tool_strict_level = tool_strict_level; + self + } + /// Override the effective model dtype used for multimodal tensor encoding. pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self { self.processor.model_dtype = Some(model_dtype); diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index a73115306342..9fb3d02b114f 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -18,10 +18,10 @@ use self::unified::unified_event_stream; use super::structured::structured_chat_event_stream; use crate::error::Result; use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream}; -use crate::parser::ParserSelection; use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; use crate::parser::unified::UnifiedParserFactory; +use crate::parser::{ParserSelection, ToolStrictLevel}; use crate::request::{ChatRequest, ChatTool}; use crate::{Error, Result as ChatResult}; @@ -49,6 +49,7 @@ impl DefaultChatOutputProcessor { tokenizer: DynTokenizer, tool_call_parser: &ParserSelection, reasoning_parser: &ParserSelection, + tool_strict_level: ToolStrictLevel, ) -> ChatResult { let parser = if let Some(parser) = Self::resolve_optional_unified_parser( request.tools(), @@ -73,7 +74,11 @@ impl DefaultChatOutputProcessor { Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box }; - apply_structural_tag_constraint(request, parser.structural_tag_builder())?; + apply_structural_tag_constraint( + request, + parser.structural_tag_builder(), + tool_strict_level, + )?; if parser.preserve_special_tokens() { request.decode_options.skip_special_tokens = false; @@ -191,7 +196,7 @@ mod tests { use vllm_tokenizer::test_utils::TestTokenizer; use super::DefaultChatOutputProcessor; - use crate::parser::ParserSelection; + use crate::parser::{ParserSelection, ToolStrictLevel}; use crate::request::ChatRequest; fn tokenizer() -> Arc { @@ -213,6 +218,7 @@ mod tests { tokenizer(), &selection, &selection, + ToolStrictLevel::Auto, ) .unwrap(); } @@ -227,6 +233,7 @@ mod tests { tokenizer(), &ParserSelection::Auto, &ParserSelection::Auto, + ToolStrictLevel::Auto, ) .unwrap(); } @@ -244,6 +251,7 @@ mod tests { tokenizer(), tool, reasoning, + ToolStrictLevel::Auto, ) .unwrap(); } @@ -258,6 +266,7 @@ mod tests { tokenizer(), &ParserSelection::Auto, &ParserSelection::Explicit("gemma4".to_string()), + ToolStrictLevel::Auto, ) { Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"), Err(error) => error, diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index 2ef6f0868593..116cd53a4abe 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -14,19 +14,21 @@ use xgrammar_structural_tag::{ build_structural_tag, }; +use crate::parser::ToolStrictLevel; use crate::request::{ChatRequest, ChatToolChoice}; use crate::{Error, Result as ChatResult}; /// Apply structural tag constraints to the request based on the tool parser's structural tag -/// support and the request's tool choice. +/// support, the request's tool choice, and the server-side strictness floor. pub(super) fn apply_structural_tag_constraint( request: &mut ChatRequest, builder: Option<&dyn StructuralTagBuilder>, + strict_level: ToolStrictLevel, ) -> ChatResult<()> { let Some(builder) = builder else { return Ok(()); }; - let Some(tool_choice) = structural_tag_tool_choice(request) else { + let Some(tool_choice) = structural_tag_tool_choice(request, strict_level) else { return Ok(()); }; @@ -34,11 +36,14 @@ pub(super) fn apply_structural_tag_constraint( .tools() .iter() .map(|tool| { + // A tool without an explicit `strict` is non-strict: its call envelope is + // constrained, but its arguments stay free unless the server level is `parameter`. + let strict = strict_level >= ToolStrictLevel::Parameter || tool.strict == Some(true); ToolParam::Function(FunctionToolParam::new(FunctionDefinition { name: tool.name.clone(), description: tool.description.clone(), parameters: Some(tool.parameters.clone()), - strict: tool.strict, + strict: Some(strict), })) }) .collect::>(); @@ -66,17 +71,21 @@ pub(super) fn apply_structural_tag_constraint( /// Resolve the tool choice used for [`xgrammar_structural_tag`] based on the request. /// /// Returns `None` if no structural tag constraints should be applied. -fn structural_tag_tool_choice(request: &ChatRequest) -> Option { +fn structural_tag_tool_choice( + request: &ChatRequest, + strict_level: ToolStrictLevel, +) -> Option { if request.tools().is_empty() { return None; } match request.tool_choice() { - // For `Auto`, only apply the structural tag if there's at least one strict tool. - ChatToolChoice::Auto if request.tools().iter().any(|tool| tool.strict == Some(true)) => { - Some(StructuralTagToolChoice::auto()) - } - ChatToolChoice::Auto | ChatToolChoice::None => None, + ChatToolChoice::None => None, + // For `Auto`, apply the structural tag only when the server raises the floor or at + // least one tool opts in with `strict: true`. + ChatToolChoice::Auto => (strict_level >= ToolStrictLevel::Function + || request.tools().iter().any(|tool| tool.strict == Some(true))) + .then(StructuralTagToolChoice::auto), ChatToolChoice::Required => Some(StructuralTagToolChoice::required()), ChatToolChoice::Function { name } => Some(StructuralTagToolChoice::function(name.clone())), @@ -148,8 +157,12 @@ mod tests { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let tag = structural_tag_value(&request); assert_eq!(tag["type"], "structural_tag"); @@ -172,8 +185,12 @@ mod tests { }; let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let tag = structural_tag_value(&request); assert!(tag.to_string().contains("lookup")); @@ -199,8 +216,12 @@ mod tests { }; let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let tag = structural_tag_value(&request).to_string(); assert!(tag.contains("search")); @@ -212,8 +233,12 @@ mod tests { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag decision should succeed"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag decision should succeed"); assert!(request.sampling_params.structured_outputs.is_none()); } @@ -227,8 +252,12 @@ mod tests { }); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let params = structured_outputs(&request); assert!(params.constraint.is_structural_tag()); @@ -242,8 +271,12 @@ mod tests { let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let tag = structural_tag_value(&request); assert_eq!(tag["type"], "structural_tag"); @@ -259,8 +292,12 @@ mod tests { }); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let params = structured_outputs(&request); assert!(params.constraint.is_structural_tag()); @@ -279,8 +316,12 @@ mod tests { ); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag should build"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag should build"); let tag = structural_tag_value(&request).to_string(); assert!(tag.contains("lookup")); @@ -292,8 +333,12 @@ mod tests { let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag decision should succeed"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag decision should succeed"); assert!(request.sampling_params.structured_outputs.is_none()); } @@ -307,10 +352,127 @@ mod tests { }); let parser = qwen3_coder_parser(request.tools()); - apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) - .expect("structural tag decision should succeed"); + apply_structural_tag_constraint( + &mut request, + parser.structural_tag_builder(), + ToolStrictLevel::Auto, + ) + .expect("structural tag decision should succeed"); let params = structured_outputs(&request); assert!(params.constraint.is_json_object()); } + + fn build(tool_choice: ChatToolChoice, tools: Vec, level: ToolStrictLevel) -> ChatRequest { + let mut request = request(tool_choice, tools); + let parser = qwen3_coder_parser(request.tools()); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder(), level) + .expect("structural tag decision should succeed"); + request + } + + fn tag_string(request: &ChatRequest) -> String { + structural_tag_value(request).to_string() + } + + #[test] + fn function_level_constrains_auto_without_strict_tools() { + let request = build( + ChatToolChoice::Auto, + vec![chat_tool("search", None)], + ToolStrictLevel::Function, + ); + + let tag = structural_tag_value(&request); + assert_eq!(tag["type"], "structural_tag"); + assert!(tag.to_string().contains("search")); + } + + #[test] + fn absent_strict_is_non_strict() { + let unset = build( + ChatToolChoice::Required, + vec![chat_tool("search", None)], + ToolStrictLevel::Auto, + ); + let non_strict = build( + ChatToolChoice::Required, + vec![chat_tool("search", Some(false))], + ToolStrictLevel::Auto, + ); + let strict = build( + ChatToolChoice::Required, + vec![chat_tool("search", Some(true))], + ToolStrictLevel::Auto, + ); + + assert_eq!(tag_string(&unset), tag_string(&non_strict)); + assert_ne!(tag_string(&unset), tag_string(&strict)); + } + + #[test] + fn strict_tool_does_not_pin_its_neighbour() { + let mixed = build( + ChatToolChoice::Auto, + vec![chat_tool("weather", Some(true)), chat_tool("search", None)], + ToolStrictLevel::Auto, + ); + let explicit = build( + ChatToolChoice::Auto, + vec![ + chat_tool("weather", Some(true)), + chat_tool("search", Some(false)), + ], + ToolStrictLevel::Auto, + ); + + assert_eq!(tag_string(&mixed), tag_string(&explicit)); + } + + #[test] + fn parameter_level_pins_argument_schemas() { + let pinned = build( + ChatToolChoice::Auto, + vec![chat_tool("search", None)], + ToolStrictLevel::Parameter, + ); + let strict = build( + ChatToolChoice::Auto, + vec![chat_tool("search", Some(true))], + ToolStrictLevel::Auto, + ); + let envelope_only = build( + ChatToolChoice::Auto, + vec![chat_tool("search", None)], + ToolStrictLevel::Function, + ); + + assert_eq!(tag_string(&pinned), tag_string(&strict)); + assert_ne!(tag_string(&pinned), tag_string(&envelope_only)); + } + + #[test] + fn tool_strict_level_parses_case_insensitively() { + assert_eq!("AUTO".parse::(), Ok(ToolStrictLevel::Auto)); + assert_eq!( + "PARAMETER".parse::(), + Ok(ToolStrictLevel::Parameter) + ); + assert_eq!( + "function".parse::(), + Ok(ToolStrictLevel::Function) + ); + assert_eq!(ToolStrictLevel::default(), ToolStrictLevel::Auto); + assert_eq!( + serde_json::to_value(ToolStrictLevel::Auto).unwrap(), + json!("auto") + ); + assert_eq!( + serde_json::from_value::(json!("auto")).unwrap(), + ToolStrictLevel::Auto + ); + assert!("off".parse::().is_err()); + assert!("strict".parse::().is_err()); + } } diff --git a/rust/src/chat/src/parser/mod.rs b/rust/src/chat/src/parser/mod.rs index c10a4f11bf3e..5a86b24934be 100644 --- a/rust/src/chat/src/parser/mod.rs +++ b/rust/src/chat/src/parser/mod.rs @@ -57,6 +57,66 @@ impl ParserSelection { } } +/// Server-side floor for tool-call structural tags (`--tool-strict-level`). +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Default, + DeserializeFromStr, + SerializeDisplay, +)] +pub enum ToolStrictLevel { + /// Derive tool constraints from the request's tool choice and per-tool strictness. + #[default] + Auto, + /// Constrain the tool-call envelope for every request with tools. + Function, + /// Additionally pin argument schemas, as if every tool were `strict: true`. + Parameter, +} + +impl ToolStrictLevel { + pub const AUTO_LITERAL: &str = "auto"; + pub const FUNCTION_LITERAL: &str = "function"; + pub const PARAMETER_LITERAL: &str = "parameter"; +} + +impl FromStr for ToolStrictLevel { + type Err = String; + + fn from_str(value: &str) -> Result { + if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) { + Ok(Self::Auto) + } else if value.eq_ignore_ascii_case(Self::FUNCTION_LITERAL) { + Ok(Self::Function) + } else if value.eq_ignore_ascii_case(Self::PARAMETER_LITERAL) { + Ok(Self::Parameter) + } else { + Err(format!( + "unknown tool strict level {value:?}; expected one of {}, {}, {}", + Self::AUTO_LITERAL, + Self::FUNCTION_LITERAL, + Self::PARAMETER_LITERAL + )) + } + } +} + +impl fmt::Display for ToolStrictLevel { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Auto => Self::AUTO_LITERAL, + Self::Function => Self::FUNCTION_LITERAL, + Self::Parameter => Self::PARAMETER_LITERAL, + }) + } +} + /// Validate explicit parser override names without starting request processing. pub fn validate_parser_overrides( tool_call_parser: &ParserSelection, diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ed96134cdc23..ab366b8827a4 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -227,6 +227,7 @@ impl ChatBackend for FakeChatBackend { self.tokenizer(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 5669e255e0d5..db3559c44f4c 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -717,6 +717,7 @@ async fn parse_completion( let processor = backends.chat_backend.new_chat_output_processor( &mut request, NewChatOutputProcessorOptions { + tool_strict_level: vllm_chat::ToolStrictLevel::Auto, tool_call_parser: &case.tool_call_parser, reasoning_parser: &case.reasoning_parser, }, diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index cda2e6c1a8b9..5bdf66da57f8 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -24,6 +24,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; use uuid::Uuid; use vllm_chat::GenerationConfigMode; +use vllm_chat::ToolStrictLevel; use vllm_chat::multimodal::MmLimitPerPrompt; use vllm_engine_core_client::TransportMode; use vllm_managed_engine::ManagedEngineConfig; @@ -134,6 +135,10 @@ pub struct RenderArgs { /// `none` to disable parsing. #[arg(long, default_value_t)] reasoning_parser: ParserSelection, + /// Server-side floor for structural-tag based tool calling: `auto`, + /// `function`, or `parameter`. + #[arg(long, default_value_t)] + tool_strict_level: ToolStrictLevel, /// Select the native chat renderer implementation. #[arg(long = "tokenizer-mode", default_value_t)] renderer: RendererSelection, @@ -171,6 +176,7 @@ impl RenderArgs { port: self.port, tool_call_parser: self.tool_call_parser, reasoning_parser: self.reasoning_parser, + tool_strict_level: self.tool_strict_level, renderer: self.renderer, chat_template: self.chat_template, default_chat_template_kwargs: self.default_chat_template_kwargs.unwrap_or_default(), @@ -241,6 +247,14 @@ pub struct SharedRuntimeArgs { #[arg(long, default_value_t)] #[serde(default = "default_py_bootstrap_parser_selection")] pub reasoning_parser: ParserSelection, + /// Server-side floor for structural-tag based tool calling, applied on + /// top of the per-tool `strict` field: `auto` follows the request's + /// tool choice and per-tool strictness, + /// `function` constrains the tool-call envelope for every request with + /// tools, `parameter` additionally pins argument schemas. + #[arg(long, default_value_t)] + #[serde(default)] + pub tool_strict_level: ToolStrictLevel, /// Select the chat renderer implementation. #[arg(long = "tokenizer-mode", default_value_t)] #[serde(default, rename = "tokenizer_mode")] @@ -517,6 +531,7 @@ impl SharedRuntimeArgs { listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, tool_call_parser: self.tool_call_parser, reasoning_parser: self.reasoning_parser, + tool_strict_level: self.tool_strict_level, renderer: self.renderer, language_model_only: self.language_model_only, chat_template: self.chat_template, @@ -574,6 +589,7 @@ impl SharedRuntimeArgs { listener_mode, tool_call_parser: self.tool_call_parser, reasoning_parser: self.reasoning_parser, + tool_strict_level: self.tool_strict_level, renderer: self.renderer, language_model_only: self.language_model_only, chat_template: self.chat_template, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 5cc6b1921450..f7f6d0a65815 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -184,6 +184,7 @@ fn serve_args_forward_python_flags_with_separator() { engine_ready_timeout_secs: 600, tool_call_parser: Auto, reasoning_parser: Auto, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, max_logprobs: None, @@ -1023,6 +1024,7 @@ fn frontend_args_accept_json() { engine_ready_timeout_secs: 600, tool_call_parser: None, reasoning_parser: None, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, max_logprobs: None, @@ -1693,6 +1695,7 @@ fn serve_args_accept_handshake_aliases() { engine_ready_timeout_secs: 600, tool_call_parser: Auto, reasoning_parser: Auto, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, max_logprobs: None, @@ -1850,6 +1853,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { }, tool_call_parser: Auto, reasoning_parser: Auto, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, chat_template: None, @@ -1942,6 +1946,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { }, tool_call_parser: Auto, reasoning_parser: Auto, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, chat_template: None, @@ -2055,6 +2060,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present }, tool_call_parser: None, reasoning_parser: None, + tool_strict_level: Auto, renderer: Auto, language_model_only: false, chat_template: None, diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index ad3872843a1f..3de227babdbe 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -16,6 +16,7 @@ use clap::Parser; use futures::StreamExt as _; use tokio_util::sync::CancellationToken; use tracing_subscriber::EnvFilter; +use vllm_chat::ToolStrictLevel; use vllm_engine_core_client::TransportMode; use vllm_server::{ ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig, @@ -70,6 +71,7 @@ async fn main() -> Result<()> { }, tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + tool_strict_level: ToolStrictLevel::Auto, renderer: RendererSelection::Auto, language_model_only: false, chat_template: None, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 4c337adbef8f..afcfaeccfb0a 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -14,6 +14,7 @@ use serde_json::Value; use vllm_chat::multimodal::MmLimitPerPrompt; use vllm_chat::{ ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, + ToolStrictLevel, }; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; use vllm_text::backend::hf::HfOverrides; @@ -229,6 +230,8 @@ pub struct Config { pub tool_call_parser: ParserSelection, /// Reasoning parser selection. pub reasoning_parser: ParserSelection, + /// Server-side floor for tool-call structural tags. + pub tool_strict_level: ToolStrictLevel, /// Chat renderer selection. pub renderer: RendererSelection, /// Disable frontend-side multimodal preprocessing and render the model as diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f4b35c9df4a8..18b6e6eb2990 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -244,6 +244,7 @@ impl ChatBackend for FakeTextBackend { self.tokenizer(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } @@ -292,6 +293,7 @@ impl ChatBackend for FakeMultimodalBackend { self.tokenizer(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 49d3d57f3264..91564db73c5b 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -139,7 +139,8 @@ async fn build_state(config: &Config) -> Result> { let chat = ChatLlm::new(text, chat_backend) .with_tool_call_parser(config.tool_call_parser.clone()) - .with_reasoning_parser(config.reasoning_parser.clone()); + .with_reasoning_parser(config.reasoning_parser.clone()) + .with_tool_strict_level(config.tool_strict_level); let state = Arc::new( AppState::new(served_model_names, chat) diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index 4fc55fe60be9..50485e26a408 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -10,7 +10,7 @@ use tokio_util::sync::CancellationToken; use tracing::info; use vllm_chat::{ ChatRequestProcessor, ChatTemplateContentFormatOption, LoadModelBackendsOptions, - ParserSelection, RendererSelection, load_model_backends, + ParserSelection, RendererSelection, ToolStrictLevel, load_model_backends, }; use vllm_text::TextRequestProcessor; use vllm_text::backend::hf::HfOverrides; @@ -32,6 +32,7 @@ pub struct RenderConfig { pub port: u16, pub tool_call_parser: ParserSelection, pub reasoning_parser: ParserSelection, + pub tool_strict_level: ToolStrictLevel, pub renderer: RendererSelection, pub chat_template: Option, pub default_chat_template_kwargs: HashMap, @@ -88,10 +89,12 @@ async fn build_state(config: &RenderConfig) -> Result> { let max_model_len = config.max_model_len.unwrap_or(u32::MAX); let text = TextRequestProcessor::new(loaded.text_backend, max_model_len) .with_max_logprobs(config.max_logprobs); - let chat = ChatRequestProcessor::render_only(loaded.chat_backend).with_parser_selections( - config.tool_call_parser.clone(), - config.reasoning_parser.clone(), - ); + let chat = ChatRequestProcessor::render_only(loaded.chat_backend) + .with_parser_selections( + config.tool_call_parser.clone(), + config.reasoning_parser.clone(), + ) + .with_tool_strict_level(config.tool_strict_level); Ok(Arc::new(RenderState { model: config.model.clone(), served_model_names, @@ -154,6 +157,7 @@ mod tests { shutdown.cancel(); let error = serve_render( RenderConfig { + tool_strict_level: ToolStrictLevel::Auto, model: "test-model".to_string(), revision: None, hf_overrides: Default::default(), diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index a0e0872e64ce..762d974b9b6b 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -168,6 +168,7 @@ impl ChatBackend for FakeChatBackend { self.tokenizer(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index c27ac20fad32..76d9ab862731 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -517,6 +517,7 @@ impl ChatBackend for FakeChatBackend { self.tokenizer(), options.tool_call_parser, options.reasoning_parser, + options.tool_strict_level, )?)) } } diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 06c8fb61295c..c216b398cc5b 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from types import SimpleNamespace from unittest.mock import MagicMock @@ -33,9 +34,11 @@ SUPPORTED_STRUCTURAL_TAG_MODELS, VLLM_BUILTIN_STRUCTURAL_TAG_MODELS, XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS, + ToolChoice, get_function_parameters, get_model_structural_tag, ) +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel @pytest.fixture @@ -231,11 +234,11 @@ def test_get_model_structural_tag_supports_all_xgrammar_builtins( def test_get_model_structural_tag_supports_vllm_hermes( - sample_tools: list[ChatCompletionToolsParam], + sample_tools_strict: list[ChatCompletionToolsParam], ): tag = get_model_structural_tag( model="hermes", - tools=sample_tools, + tools=sample_tools_strict, tool_choice="required", reasoning=False, ) @@ -273,6 +276,22 @@ def test_get_model_structural_tag_supports_vllm_hermes( assert content["json_schema"] == expected_schema +def test_hermes_required_without_strict_leaves_arguments_free( + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model="hermes", + tools=sample_tools, + tool_choice="required", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + tags = tag.model_dump()["format"]["tags"] + assert tags + assert all(tag_dump["content"]["json_schema"] is True for tag_dump in tags) + + def test_hermes_required_tool_calls_use_empty_separator(): tools = [ ChatCompletionToolsParam( @@ -361,10 +380,19 @@ def _k3_tools(*calls: str) -> str: return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}" -def _k3_grammar(tool_choice, tools=None): +def _as_strict(tools: list[ChatCompletionToolsParam]) -> list[ChatCompletionToolsParam]: + for tool in tools: + tool.function.strict = True + return tools + + +def _k3_grammar(tool_choice, tools=None, *, strict: bool = True): + """Build the K3 grammar; tools are strict by default so the typed + argument channel is exercised (an unset ``strict`` leaves arguments free).""" + tools = tools if tools is not None else _k3_tools_by_name() tag = get_model_structural_tag( model="kimi_k3", - tools=tools if tools is not None else _k3_tools_by_name(), + tools=_as_strict(tools) if strict else tools, tool_choice=tool_choice, reasoning=False, ) @@ -467,6 +495,18 @@ def test_kimi_k3_required_rejects_invalid(body: str): assert not _is_grammar_accept_string(_k3_grammar("required"), body) +def test_kimi_k3_required_without_strict_leaves_arguments_free(): + # An unset ``strict`` pins the call envelope only: a value the schema + # would reject is accepted, and the declared-required argument may be + # omitted. + grammar = _k3_grammar("required", strict=False) + body = _k3_response() + _k3_tools( + _k3_call("get_weather", _k3_arg("days", "number", "abc")) + ) + assert _is_grammar_accept_string(grammar, body) + assert not _is_grammar_accept_string(_k3_grammar("required"), body) + + def test_kimi_k3_schema_without_required_accepts_empty_call(): tools = [ ChatCompletionToolsParam( @@ -829,3 +869,209 @@ def test_kimi_k3_forced_tool_choice_builds_single_mandatory_call(): response_only = _k3_response("no call here") assert _is_grammar_accept_string(grammar, ok) assert not _is_grammar_accept_string(grammar, response_only) + + +def _pins_argument_schema(tag: StructuralTag) -> bool: + return '"json_schema": {' in json.dumps(tag.model_dump(), ensure_ascii=False) + + +def _dumped(tag: StructuralTag) -> str: + return json.dumps(tag.model_dump(), ensure_ascii=False) + + +def test_tool_strict_level_auto_is_the_default( + sample_tools: list[ChatCompletionToolsParam], +): + """Auto + no strict tool gets no tag unless the operator raises the floor.""" + assert ( + get_model_structural_tag( + model="deepseek_v4", + tools=sample_tools, + tool_choice="auto", + reasoning=False, + ) + is None + ) + + +@pytest.mark.parametrize("model", sorted(SUPPORTED_STRUCTURAL_TAG_MODELS)) +def test_tool_strict_level_function_lifts_auto_gate( + model: str, + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice="auto", + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + + assert tag is not None + + +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_tool_strict_level_function_pins_envelope_only( + model: str, + sample_tools: list[ChatCompletionToolsParam], +): + """The request's own tools stay untouched; unset ``strict`` stays free.""" + tag = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice="auto", + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + + assert tag is not None + assert not _pins_argument_schema(tag) + assert all(tool.function.strict is None for tool in sample_tools) + + +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_tool_strict_level_parameter_pins_argument_schemas( + model: str, + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice="auto", + reasoning=False, + strict_level=ToolStrictLevel.PARAMETER, + ) + + assert tag is not None + assert _pins_argument_schema(tag) + + +def test_tool_strict_level_function_keeps_client_strict_tools_strict( + sample_tools_strict: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model="deepseek_v4", + tools=sample_tools_strict, + tool_choice="auto", + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + + assert tag is not None + assert _pins_argument_schema(tag) + + +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_absent_strict_is_not_pinned_next_to_a_strict_tool(model: str): + """A strict tool must not drag a neighbour with unset ``strict`` into + schema enforcement.""" + weather = ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "strict": True, + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + ) + search = ChatCompletionToolsParam( + type="function", + function={ + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + ) + tag = get_model_structural_tag( + model=model, + tools=[weather, search], + tool_choice="auto", + reasoning=False, + ) + + assert tag is not None + dumped = _dumped(tag) + assert '"city"' in dumped + assert '"query"' not in dumped + + +@pytest.mark.parametrize( + "tool_choice", + [ + "required", + ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name="get_weather") + ), + ], +) +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_forced_tool_choice_pins_schema_only_for_strict_tools( + model: str, + tool_choice: ToolChoice, + sample_tools: list[ChatCompletionToolsParam], + sample_tools_strict: list[ChatCompletionToolsParam], +): + """Required / named always constrain the call; the argument schema + follows the per-tool ``strict`` (or the parameter level).""" + free = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice=tool_choice, + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + pinned = get_model_structural_tag( + model=model, + tools=sample_tools_strict, + tool_choice=tool_choice, + reasoning=False, + ) + floored = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice=tool_choice, + reasoning=False, + strict_level=ToolStrictLevel.PARAMETER, + ) + + assert free is not None and not _pins_argument_schema(free) + assert pinned is not None and _pins_argument_schema(pinned) + assert floored is not None and _pins_argument_schema(floored) + + +def test_tool_strict_level_parameter_overrides_client_strict_false(): + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + "strict": False, + }, + ) + ] + tag = get_model_structural_tag( + model="deepseek_v4", + tools=tools, + tool_choice="required", + reasoning=False, + strict_level=ToolStrictLevel.PARAMETER, + ) + + assert tag is not None + assert _pins_argument_schema(tag) + + +def test_tool_strict_level_from_name(): + assert ToolStrictLevel.from_name("AUTO") is ToolStrictLevel.AUTO + assert ToolStrictLevel.from_name("Parameter") is ToolStrictLevel.PARAMETER + with pytest.raises(ValueError, match="expected one of auto, function, parameter"): + ToolStrictLevel.from_name("strict") + with pytest.raises(ValueError, match="expected one of auto, function, parameter"): + ToolStrictLevel.from_name("off") diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index e90d70e289a4..6cc1271add1c 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -112,6 +112,7 @@ def __init__( reasoning_parser: str = "", enable_auto_tools: bool = False, tool_parser: str | None = None, + tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, default_chat_template_kwargs: dict[str, Any] | None = None, @@ -127,6 +128,7 @@ def __init__( return_tokens_as_token_ids=return_tokens_as_token_ids, reasoning_parser=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 788558645243..63d00f0c576a 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -249,6 +249,7 @@ def __init__( reasoning_parser: str = "", enable_auto_tools: bool = False, tool_parser: str | None = None, + tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, default_chat_template_kwargs: dict[str, Any] | None = None, @@ -265,6 +266,7 @@ def __init__( return_tokens_as_token_ids=return_tokens_as_token_ids, reasoning_parser=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index ac995a4f031a..189c4e6779b0 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -112,6 +112,7 @@ async def init_generate_state( return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, tool_server=state.tool_server, reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, @@ -137,6 +138,7 @@ async def init_generate_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, @@ -178,6 +180,7 @@ async def init_generate_state( return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, @@ -198,6 +201,7 @@ async def init_generate_state( return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py index 09bb0ec63020..d33be659494c 100644 --- a/vllm/entrypoints/launchers/api_server/app_state.py +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -95,6 +95,7 @@ async def init_app_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.structured_outputs_config.reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, @@ -111,6 +112,7 @@ async def init_app_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.structured_outputs_config.reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 71ae94c10635..488a7f2c43c8 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -20,6 +20,7 @@ ) from vllm.entrypoints.openai.models.protocol import LoRAModulePath from vllm.tool_parsers import ToolParserManager +from vllm.tool_parsers.tool_strict_level import ToolStrictLevelName from vllm.utils.argparse_utils import FlexibleArgumentParser from .utils.constants import ( @@ -116,6 +117,13 @@ class BaseFrontendArgs: """Special the tool parser plugin write to parse the model-generated tool into OpenAI API format, the name register in this plugin can be used in `--tool-call-parser`.""" + tool_strict_level: ToolStrictLevelName = "auto" + """Server-side floor for structural-tag based tool calling, applied on top + of the per-tool `strict` field. `auto` follows the request's tool choice + and per-tool strictness; `function` constrains the + tool-call envelope (markup and function name) for every request with + tools; `parameter` additionally pins argument schemas, as if every tool + were `strict: true`.""" tool_server: str | None = None """Comma-separated list of host:port pairs (IPv4, IPv6, or hostname). Examples: 127.0.0.1:8000, [::1]:8000, localhost:1234. Or `demo` for diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py index 33988443d1d9..1d242977be29 100644 --- a/vllm/entrypoints/launchers/render/app_state.py +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -60,6 +60,7 @@ async def init_render_app_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, @@ -76,6 +77,7 @@ async def init_render_app_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, + tool_strict_level=args.tool_strict_level, reasoning_parser=args.reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 51c0af8e8f15..3e2d8f7369ac 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -132,6 +132,7 @@ def __init__( enable_auto_tools: bool = False, exclude_tools_when_tool_choice_none: bool = False, tool_parser: str | None = None, + tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, enable_log_outputs: bool = False, @@ -161,6 +162,7 @@ def __init__( tool_parser_name=tool_parser, reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, model_name=self.model_config.model, is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index ee4f7d55ca0d..8108c5f77537 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -109,6 +109,7 @@ def __init__( reasoning_parser: str = "", enable_auto_tools: bool = False, tool_parser: str | None = None, + tool_strict_level: str = "auto", tool_server: ToolServer | None = None, enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, @@ -135,6 +136,7 @@ def __init__( tool_parser_name=tool_parser, reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, model_name=self.model_config.model, is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index af448e0951c6..0e84e710ed79 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -36,6 +36,7 @@ extract_named_tool_call_streaming, extract_required_tool_call_streaming, ) +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel logger = init_logger(__name__) @@ -95,6 +96,8 @@ class Parser: # Subclasses should override these if they use specific parser classes reasoning_parser_cls: type[ReasoningParser] | None = None tool_parser_cls: type[ToolParser] | None = None + # Server-side floor for tool-call structural tags (--tool-strict-level). + tool_strict_level: ToolStrictLevel = ToolStrictLevel.AUTO def __init__( self, @@ -415,6 +418,7 @@ def _apply_structural_tag( structure_tag = self._tool_parser.get_structural_tag( request, reasoning=False, + strict_level=self.tool_strict_level, ) if structure_tag is None: return request diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index a8efd3f0eaa1..41b27967a6cb 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING from vllm.logger import init_logger +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel if TYPE_CHECKING: from vllm.parser.abstract_parser import Parser @@ -77,6 +78,7 @@ def get_parser( enable_auto_tools: bool = False, model_name: str | None = None, is_harmony: bool = False, + tool_strict_level: str = "auto", ) -> type[Parser] | None: """Get a Parser that handles both reasoning and tool parsing. @@ -89,6 +91,8 @@ def get_parser( model_name: The model name for parser-specific warnings. is_harmony: Whether the selected model uses the Harmony format. If True, HarmonyParser is always returned. + tool_strict_level: Server-side floor for tool-call structural + tags (``--tool-strict-level``). Returns: A Parser class, or None if neither parser is specified. @@ -105,11 +109,14 @@ def get_parser( if reasoning_parser_cls is None and tool_parser_cls is None: return None + strict_level = ToolStrictLevel.from_name(tool_strict_level) + if is_harmony: from vllm.parser.harmony import HarmonyParser HarmonyParser.reasoning_parser_cls = reasoning_parser_cls HarmonyParser.tool_parser_cls = tool_parser_cls + HarmonyParser.tool_strict_level = strict_level return HarmonyParser if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": @@ -121,6 +128,7 @@ def get_parser( class _KimiK3Parser(KimiK3Parser): reasoning_parser_cls = r_cls tool_parser_cls = t_cls + tool_strict_level = strict_level return _KimiK3Parser @@ -136,6 +144,7 @@ class _KimiK3Parser(KimiK3Parser): class _CohereCommandParser(CohereCommandParser): reasoning_parser_cls = r_cls tool_parser_cls = t_cls + tool_strict_level = strict_level return _CohereCommandParser @@ -147,5 +156,6 @@ class _CohereCommandParser(CohereCommandParser): class _Parser(DelegatingParser): reasoning_parser_cls = r_cls tool_parser_cls = t_cls + tool_strict_level = strict_level return _Parser diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 5f30b552ba83..05f2710afb95 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -54,6 +54,7 @@ def __init__( exclude_tools_when_tool_choice_none: bool = False, tool_parser: str | None = None, reasoning_parser: str | None = None, + tool_strict_level: str = "auto", default_chat_template_kwargs: dict[str, Any] | None = None, log_error_stack: bool = False, ) -> None: @@ -68,6 +69,7 @@ def __init__( tool_parser_name=tool_parser, reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, model_name=model_config.model, is_harmony=self.use_harmony, ) diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index 2ab5725b73fc..f4d21dedad06 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -127,6 +127,7 @@ def __init__( exclude_tools_when_tool_choice_none: bool = False, tool_parser: str | None = None, reasoning_parser: str | None = None, + tool_strict_level: str = "auto", default_chat_template_kwargs: dict[str, Any] | None = None, log_error_stack: bool = False, ) -> None: @@ -141,6 +142,7 @@ def __init__( tool_parser_name=tool_parser, reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, + tool_strict_level=tool_strict_level, model_name=model_config.model, is_harmony=self.use_harmony, ) diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 3847afa9b227..953782daa96f 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -30,6 +30,7 @@ StructuredOutputsParams, ) from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel from vllm.tool_parsers.utils import Tool, get_json_schema_from_tools from vllm.utils.collection_utils import is_list_of from vllm.utils.import_utils import import_plugin @@ -168,6 +169,7 @@ def get_structural_tag( request: ChatCompletionRequest | ResponsesRequest, *, reasoning: bool = False, + strict_level: ToolStrictLevel = ToolStrictLevel.AUTO, ): if self.structural_tag_model is None: return None @@ -180,6 +182,7 @@ def get_structural_tag( tools=request.tools, tool_choice=request.tool_choice, reasoning=reasoning, + strict_level=strict_level, ) def extract_tool_calls( diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py index c2f82fc4a7be..b9c857dbd9e6 100644 --- a/vllm/tool_parsers/hy_v4_tool_parser.py +++ b/vllm/tool_parsers/hy_v4_tool_parser.py @@ -25,6 +25,7 @@ from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel from vllm.tool_parsers.utils import partial_tag_overlap if TYPE_CHECKING: @@ -994,6 +995,7 @@ def get_structural_tag( request: ChatCompletionRequest | ResponsesRequest, *, reasoning: bool = False, + strict_level: ToolStrictLevel = ToolStrictLevel.AUTO, ) -> StructuralTag | None: """Build a structural tag matching HYV4's tool tokens. @@ -1009,6 +1011,7 @@ def get_structural_tag( Args: request: The request being adjusted. reasoning: Whether the grammar also covers the reasoning phase. + strict_level: Server-side floor from ``--tool-strict-level``. Returns: The structural tag, or None when structural tagging does not apply. @@ -1041,6 +1044,7 @@ def get_structural_tag( tool_choice=request.tool_choice, reasoning=reasoning, token_suffix=self._extractor.token_suffix, + strict_level=strict_level, ) except Exception: logger.warning( diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 295d5844b430..5d65b186fe30 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -35,6 +35,7 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, ) +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel ToolChoice: TypeAlias = ( Literal["none", "auto", "required"] @@ -94,15 +95,57 @@ def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder: return decorator +def _tool_is_strict(tool: ChatCompletionToolsParam | ResponsesTool) -> bool: + if isinstance(tool, FunctionTool): + return tool.strict is True + if isinstance(tool, ChatCompletionToolsParam): + return tool.function.strict is True + return False + + def _any_tool_strict( tools: Sequence[ChatCompletionToolsParam | ResponsesTool], ) -> bool: - for tool in tools: - if isinstance(tool, FunctionTool) and tool.strict is True: - return True - if isinstance(tool, ChatCompletionToolsParam) and tool.function.strict is True: - return True - return False + return any(_tool_is_strict(tool) for tool in tools) + + +def _with_tool_strict( + tool: ChatCompletionToolsParam | ResponsesTool, + strict: bool, +) -> ChatCompletionToolsParam | ResponsesTool: + """Return a copy of ``tool`` with ``strict`` set, leaving the request's alone.""" + if isinstance(tool, FunctionTool): + return tool.model_copy(update={"strict": strict}) + if isinstance(tool, ChatCompletionToolsParam): + return tool.model_copy( + update={"function": tool.function.model_copy(update={"strict": strict})} + ) + return tool + + +def _resolve_tool_strictness( + tools: Sequence[ChatCompletionToolsParam | ResponsesTool], + tool_choice: ToolChoice, + strict_level: ToolStrictLevel, +) -> Sequence[ChatCompletionToolsParam | ResponsesTool] | None: + """Decide whether a structural tag applies and pin each tool's ``strict``. + + A tool without an explicit ``strict`` is treated as non-strict: its call + envelope is still constrained, but its arguments stay free unless the + server level is PARAMETER. ``None`` means no structural tag. + """ + if ( + tool_choice == "auto" + and strict_level == ToolStrictLevel.AUTO + and not _any_tool_strict(tools) + ): + return None + return [ + _with_tool_strict( + tool, strict_level >= ToolStrictLevel.PARAMETER or _tool_is_strict(tool) + ) + for tool in tools + ] def get_model_structural_tag( @@ -111,12 +154,14 @@ def get_model_structural_tag( tool_choice: ToolChoice, reasoning: bool, token_suffix: str = "", + strict_level: ToolStrictLevel = ToolStrictLevel.AUTO, ) -> StructuralTag | None: """Build a structural tag with xgrammar's builtin model templates.""" if not tools or tool_choice == "none": return None - if tool_choice == "auto" and not _any_tool_strict(tools): + tools = _resolve_tool_strictness(tools, tool_choice, strict_level) + if tools is None: return None dumped_tools = [_dump_tool_for_xgrammar(tool) for tool in tools] diff --git a/vllm/tool_parsers/tool_strict_level.py b/vllm/tool_parsers/tool_strict_level.py new file mode 100644 index 000000000000..03a7d8b170a1 --- /dev/null +++ b/vllm/tool_parsers/tool_strict_level.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import enum +from typing import Literal + +ToolStrictLevelName = Literal["auto", "function", "parameter"] + + +class ToolStrictLevel(enum.IntEnum): + """Server-side floor for tool-call structural tags (``--tool-strict-level``). + + AUTO: derive call constraints from ``tool_choice`` and argument-schema + enforcement from each tool's explicit ``strict`` field. + FUNCTION: constrain the tool-call envelope for every request with tools. + PARAMETER: additionally pin argument schemas, as if every tool were + ``strict: true``. + """ + + AUTO = 0 + FUNCTION = 1 + PARAMETER = 2 + + @classmethod + def from_name(cls, name: str) -> "ToolStrictLevel": + try: + return cls[name.strip().upper()] + except KeyError: + expected = ", ".join(level.name.lower() for level in cls) + raise ValueError( + f"Unknown tool strict level {name!r}; expected one of {expected}." + ) from None From 39e33db7f3b10ebd3aca6a0008f2990978f118e7 Mon Sep 17 00:00:00 2001 From: lazymio Date: Fri, 18 Sep 2026 13:11:24 +0800 Subject: [PATCH 456/774] [Frontend] Fix the parsing of missing `string=` in DeepSeek V4 (#56271) Co-authored-by: Claude Opus 5 Co-authored-by: Bugen Zhao Signed-off-by: lazymio --- .../src/tool/deepseek_dsml/deepseek_v32.rs | 28 ++++++++ .../src/tool/deepseek_dsml/deepseek_v4.rs | 23 ++++++ .../tool/deepseek_dsml/deepseek_v41/tests.rs | 22 ++++++ rust/src/parser/src/tool/deepseek_dsml/mod.rs | 11 +-- tests/parser/engine/test_deepseek_v4.py | 70 +++++++++++++++++++ tests/parser/engine/test_deepseek_v41.py | 2 + vllm/parser/deepseek_v4.py | 32 +++++---- vllm/parser/deepseek_v41.py | 15 +--- 8 files changed, 174 insertions(+), 29 deletions(-) diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs index 632eda600f77..9bbe39367b5e 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs @@ -193,6 +193,34 @@ mod tests { ); } + #[test] + fn deepseek_v32_parse_complete_keeps_parameter_without_string_attr() { + // The model sometimes omits `string="..."`; the parameter must still + // be kept and converted through the schema, like `string="false"`. + let mut parser = DeepSeekV32ToolParser::new(&test_tools()); + let output = parser + .parse_complete( + "<｜DSML｜function_calls>\n\ + <｜DSML｜invoke name=\"convert\">\n\ + <｜DSML｜parameter name=\"whole\">5.0\n\ + <｜DSML｜parameter name=\"flag\" >true\n\ + <｜DSML｜parameter name=\"payload\" string=\"true\">{\"nested\":true}\n\ + \n\ + ", + ) + .unwrap(); + + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ + "whole": 5.0, + "flag": true, + "payload": "{\"nested\":true}", + }) + ); + } + #[test] fn deepseek_v32_parse_complete_preserves_raw_closing_tag_text_in_parameter_value() { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs index b65df235ca20..032cfbe5e284 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs @@ -147,6 +147,29 @@ mod tests { ); } + #[test] + fn deepseek_v4_streaming_keeps_parameter_without_string_attr() { + let mut parser = DeepSeekV4ToolParser::new(&test_tools()); + let output = collect_stream( + &mut parser, + &[ + "<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n", + "<｜DSML｜parameter name=\"loc", + "ation\">Par", + "is\n", + "<｜DSML｜parameter name=\"date\" string=\"true\">tomorrow\n", + "\n", + ], + ); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "location": "Paris", "date": "tomorrow" }) + ); + } + #[test] fn tool_framing_preserves_body_whitespace_across_chunk_boundaries() { assert_tool_framing_preserves_body_whitespace::( diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs index d40f82d533cd..5a6758f5c935 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs @@ -71,6 +71,28 @@ fn spaced_dsml_streaming_preserves_shared_schema_coercion() { } } +#[test] +fn spaced_dsml_keeps_parameter_without_string_attr() { + let wire = concat!( + "<｜DSML｜ calls>\n", + "<｜DSML｜ invoke name=\"lookup\">\n", + "<｜DSML｜ parameter name=\"query\">value\n", + "<｜DSML｜ parameter name=\"limit\">2\n", + "\n", + "", + ); + let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap(); + + let output = collect_stream(parser.as_mut(), &[wire]); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "query": "value", "limit": 2 }) + ); +} + #[test] fn v4_tags_remain_text_in_v41_dialect() { let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap(); diff --git a/rust/src/parser/src/tool/deepseek_dsml/mod.rs b/rust/src/parser/src/tool/deepseek_dsml/mod.rs index 01a3fbe282ea..9c7dd203312d 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/mod.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/mod.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; -use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated}; +use winnow::combinator::{alt, delimited, eof, opt, preceded, repeat, seq, terminated}; use winnow::prelude::*; use winnow::stream::Partial; use winnow::token::{literal, rest, take_until}; @@ -293,13 +293,16 @@ fn parse_invoke_params(invoke_body: &str, tokens: DsmlTokens) -> ModalResult ModalResult { seq! {DsmlParameter { _: literal(tokens.parameter_start), _: ws1, name: name_attr.map(|name: &str| name.to_string()), - _: ws1, - is_string: string_attr.map(|value| value == "true"), + is_string: opt(preceded(ws1, string_attr)).map(|value| value == Some("true")), _: ws0, _: ">", value: take_until(0.., tokens.parameter_end).map(str::to_string), @@ -315,7 +318,7 @@ fn name_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> { /// Parse a string attribute. fn string_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> { - delimited("string=\"", alt(("true", "false")), "\"").parse_next(input) + delimited("string=\"", take_until(0.., "\""), "\"").parse_next(input) } /// Parse a DSML name attribute. diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 7cbf9791a7b0..2cf3a8bb85f3 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -119,6 +119,35 @@ def test_empty_args(self): result = json.loads(_dsml_arg_converter("", partial=False)) assert result == {} + def _bare(self, name: str, value: str) -> str: + return f'<｜DSML｜parameter name="{name}">{value}{_PARAM_CLOSE}' + + def test_missing_string_attribute_keeps_literal(self): + """The model sometimes omits ``string``; the parameter must survive.""" + raw = self._bare("name", "alpha") + assert json.loads(_dsml_arg_converter(raw, partial=False)) == {"name": "alpha"} + + def test_missing_string_attribute_parses_json(self): + raw = self._bare("count", "7") + self._bare("opts", '{"k": [1]}') + assert json.loads(_dsml_arg_converter(raw, partial=False)) == { + "count": 7, + "opts": {"k": [1]}, + } + + def test_missing_string_attribute_alongside_annotated(self): + raw = self._bare("a", "x") + _param("b", "true", "42") + assert json.loads(_dsml_arg_converter(raw, partial=False)) == { + "a": "x", + "b": "42", + } + + def test_missing_string_attribute_partial_value(self): + raw = _param("a", "true", "x") + '<｜DSML｜parameter name="n">12' + assert json.loads(_dsml_arg_converter(raw, partial=True)) == { + "a": "x", + "n": 12, + } + def test_invalid_json_fallback(self): raw = self._raw(("data", "false", "[broken")) result = json.loads(_dsml_arg_converter(raw, partial=False)) @@ -232,6 +261,47 @@ def test_streaming_split_next_parameter_tag_is_buffered( } +class TestMissingStringAttribute: + def test_non_streaming_keeps_parameter(self, mock_tokenizer, mock_request): + text = ( + f"{DSML_TOOL_START}" + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n" + f'<｜DSML｜parameter name="location">Paris{_PARAM_CLOSE}\n' + f"{_param('date', 'true', 'tomorrow')}\n" + f"{DSML_INVOKE_END}{DSML_TOOL_END}" + ) + + result = DeepSeekV4Parser(mock_tokenizer).extract_tool_calls(text, mock_request) + + assert result.tools_called is True + assert json.loads(result.tool_calls[0].function.arguments) == { + "location": "Paris", + "date": "tomorrow", + } + + def test_streaming_keeps_parameter(self, mock_tokenizer, mock_request): + chunks = [ + DSML_TOOL_START, + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n", + '<｜DSML｜parameter name="loc', + 'ation">Par', + "is", + _PARAM_CLOSE, + f"\n{_param('date', 'true', 'tomorrow')}\n", + DSML_INVOKE_END, + DSML_TOOL_END, + ] + + results = simulate_tool_streaming( + DeepSeekV4Parser(mock_tokenizer), mock_request, chunks + ) + + assert json.loads(collect_tool_arguments(results)) == { + "location": "Paris", + "date": "tomorrow", + } + + # ── Bare absorption and duplicate absorption ───────── diff --git a/tests/parser/engine/test_deepseek_v41.py b/tests/parser/engine/test_deepseek_v41.py index cc7d6e0aedf4..d99ca6b2a17b 100644 --- a/tests/parser/engine/test_deepseek_v41.py +++ b/tests/parser/engine/test_deepseek_v41.py @@ -137,10 +137,12 @@ def test_python_argument_conversion_and_partial_values(): '<｜DSML｜ parameter name="object" string="false">' '{"a": [true, null]}' '<｜DSML｜ parameter name="bad" string="false">[broken' + '<｜DSML｜ parameter name="bare">7' '<｜DSML｜ parameter name="text" string="true"> a", ) -_ESCAPED_DSML = re.escape(_DSML) -_PARAM_RE = re.compile( - rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">' - rf"(.*?)" - rf"(?:|(?=<{_ESCAPED_DSML}parameter\s+name=))", - re.DOTALL, -) -_PARTIAL_PARAM_RE = re.compile( - rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">' - rf"(.*)$", - re.DOTALL, -) + +def _param_patterns( + param_start: str, param_close: str +) -> tuple[re.Pattern, re.Pattern]: + """Complete and trailing-partial parameter regexes for one DSML dialect. + + The ``string`` attribute is optional: the model sometimes omits it, and + dropping such a parameter hands the client a tool call with no arguments. + """ + start, close = re.escape(param_start), re.escape(param_close) + head = rf'{start}\s+name="([^"]+)"(?:\s+string="([^"]*)")?\s*>' + return ( + re.compile(rf"{head}(.*?)(?:{close}|(?={start}\s+name=))", re.DOTALL), + re.compile(rf"{head}(.*)$", re.DOTALL), + ) + + +_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE) def _dsml_arg_converter( @@ -80,6 +86,8 @@ def _dsml_arg_converter( params: dict[str, object] = {} last_end = 0 + # A missing ``string`` attribute is treated like ``string="false"``: try + # JSON, fall back to the literal text. for m in param_re.finditer(raw_args): name, is_str, value = m.group(1), m.group(2), m.group(3) if is_str == "true": diff --git a/vllm/parser/deepseek_v41.py b/vllm/parser/deepseek_v41.py index 668bdd4a6bed..84216005102a 100644 --- a/vllm/parser/deepseek_v41.py +++ b/vllm/parser/deepseek_v41.py @@ -5,11 +5,10 @@ import functools from dataclasses import replace -import regex as re - from vllm.parser.deepseek_v4 import ( DeepSeekV4Parser, _dsml_arg_converter, + _param_patterns, deepseek_v4_config, ) from vllm.parser.engine.parser_engine_config import ParserEngineConfig @@ -21,17 +20,7 @@ DSML_PARAM_START = "<｜DSML｜ parameter" DSML_PARAM_CLOSE = "" -_PARAM_RE = re.compile( - r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">' - r"(.*?)" - r"(?:|(?=<｜DSML｜ parameter\s+name=))", - re.DOTALL, -) -_PARTIAL_PARAM_RE = re.compile( - r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">' - r"(.*)$", - re.DOTALL, -) +_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE) @functools.cache From 42a33039eb43c2adc885329a1c26ee1a5d8cc0fb Mon Sep 17 00:00:00 2001 From: Sun <2114377220@qq.com> Date: Fri, 18 Sep 2026 14:18:41 +0800 Subject: [PATCH 457/774] [Perf] Fuse CohereASR relative attention score accumulation (#55190) Signed-off-by: levius <2114377220@qq.com> Co-authored-by: Codex --- tests/models/multimodal/test_cohere_asr.py | 58 +++++++++++++++++++++- vllm/model_executor/models/cohere_asr.py | 16 ++++-- 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/tests/models/multimodal/test_cohere_asr.py b/tests/models/multimodal/test_cohere_asr.py index 85a9b59dd513..950542f90019 100644 --- a/tests/models/multimodal/test_cohere_asr.py +++ b/tests/models/multimodal/test_cohere_asr.py @@ -4,9 +4,65 @@ from types import SimpleNamespace import pytest +import torch from vllm.config import SpeechToTextConfig -from vllm.model_executor.models.cohere_asr import CohereAsrForConditionalGeneration +from vllm.model_executor.models.cohere_asr import ( + CohereAsrForConditionalGeneration, + RelPositionMultiHeadAttention, +) + + +class _ReferenceRelPositionMultiHeadAttention(RelPositionMultiHeadAttention): + """Relative attention implementation before score accumulation is fused.""" + + def forward(self, query, key, value, mask, pos_emb=None): + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) + + assert pos_emb is not None + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) + + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)] + scores = (matrix_ac + matrix_bd) / self.s_d_k + return self.forward_attention(v, scores, mask) + + +@pytest.mark.parametrize(("query_len", "key_len"), [(2, 2), (3, 5)]) +def test_relative_attention_fuses_score_accumulation(monkeypatch, query_len, key_len): + torch.manual_seed(7) + batch, heads, model_dim = 2, 4, 32 + optimized = RelPositionMultiHeadAttention(heads, model_dim, None, None) + reference = _ReferenceRelPositionMultiHeadAttention(heads, model_dim, None, None) + reference.load_state_dict(optimized.state_dict()) + + query = torch.randn(batch, query_len, model_dim) + key = torch.randn(batch, key_len, model_dim) + value = torch.randn(batch, key_len, model_dim) + pos_emb = torch.randn(1, query_len + key_len - 1, model_dim) + expected = reference(query, key, value, None, pos_emb) + + original_baddbmm = torch.baddbmm + calls = 0 + + def tracked_baddbmm(*args, **kwargs): + nonlocal calls + calls += 1 + return original_baddbmm(*args, **kwargs) + + monkeypatch.setattr(torch, "baddbmm", tracked_baddbmm) + actual = optimized(query, key, value, None, pos_emb) + + torch.testing.assert_close(actual, expected) + assert calls == 1 @pytest.mark.parametrize( diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 844e8c5e0e7f..33b4aa7b0d70 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -1206,10 +1206,18 @@ def forward( matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - # drops extra elements in the matrix_bd to match the matrix_ac's size - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)] - scores = (matrix_ac + matrix_bd) / self.s_d_k # (batch, head, time1, time2) + # drops extra elements in matrix_bd to match the key sequence length + matrix_bd = matrix_bd[:, :, :, : k.size(-2)] + batch, heads, query_len, key_len = matrix_bd.shape + scale = 1.0 / self.s_d_k + scores = torch.baddbmm( + matrix_bd.flatten(0, 1), + q_with_bias_u.flatten(0, 1), + k.transpose(-2, -1).flatten(0, 1), + beta=scale, + alpha=scale, + ).view(batch, heads, query_len, key_len) + return self.forward_attention(v, scores, mask) From d5f0a6e829faa69d1db289bf62b14dae136c02b2 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 18 Sep 2026 01:33:41 -0500 Subject: [PATCH 458/774] [Bugfix][DBO] Fix DeepEP low-latency profiling crash with DP+EP+DBO (#57502) Signed-off-by: Micah Williamson Co-authored-by: Andreas Karatzas --- vllm/v1/worker/gpu_model_runner.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c943f02212bb..cedf95e319a2 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6164,6 +6164,11 @@ def _dummy_run( num_tokens_padded, None, False ) + num_tokens_unpadded = num_tokens_padded if is_profile else 0 + is_padding = self._prepare_padding_mask( + num_tokens_unpadded, num_tokens_padded + ) + if ubatch_slices_padded is not None: # Adjust values to reflect a single ubatch. # TODO(sage,lucas): this is cruft that should be addressed in @@ -6172,8 +6177,6 @@ def _dummy_run( if num_tokens_across_dp is not None: num_tokens_across_dp[:] = num_tokens_padded - is_padding = self._prepare_padding_mask(0, num_tokens_padded) - with ( self.maybe_randomize_inputs( input_ids, inputs_embeds, randomize_inputs=randomize_inputs From 95d723700b1df07b20263c696c0342b92c188bfe Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 18 Sep 2026 01:26:17 -0700 Subject: [PATCH 459/774] [Rust Frontend] Return sampling masks over gRPC (#56777) Co-authored-by: Bugen Zhao Signed-off-by: Biswa Panda --- rust/Cargo.lock | 2 +- rust/proto/Cargo.toml | 2 +- rust/proto/inference.proto | 2 + rust/src/chat/src/output/default/unified.rs | 1 + rust/src/chat/src/output/harmony/tests.rs | 1 + rust/src/chat/tests/roundtrip.rs | 2 + .../src/protocol/logprobs.rs | 2 +- .../src/protocol/logprobs/array.rs | 2 +- .../engine-core-client/src/protocol/mod.rs | 1 + .../engine-core-client/src/protocol/output.rs | 6 +- .../src/protocol/sampling_mask.rs | 237 ++++++++++++++++++ .../engine-core-client/src/tests/client.rs | 12 +- .../src/tests/python_compat.py | 8 +- rust/src/llm/src/error.rs | 8 + rust/src/llm/src/output.rs | 31 +++ rust/src/llm/tests/generate.rs | 69 ++++- rust/src/server/Cargo.toml | 1 + rust/src/server/src/grpc/convert.rs | 44 ++++ rust/src/server/src/grpc/inference.rs | 1 + .../server/src/routes/inference/generate.rs | 6 + .../server/src/routes/openai/completions.rs | 6 + rust/src/text/src/output/decoded.rs | 146 +++++++++++ rust/src/text/src/output/mod.rs | 7 + 23 files changed, 570 insertions(+), 27 deletions(-) create mode 100644 rust/src/engine-core-client/src/protocol/sampling_mask.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 601e11c05448..dee3f3078966 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5941,7 +5941,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.3.0" +version = "0.4.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index cac268bc74fb..492cb5a06d67 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.3.0" +version = "0.4.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 46313f915a40..4eaf6f92f4b4 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -160,6 +160,8 @@ message SequenceOutput { // Only present in final output for this sequence optional FinishInfo finish_info = 8; + // Only present in final output when sampling-mask capture is enabled. + repeated TokenIds sampling_mask = 10; } // Prompt info, returned in the first response diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index 964f0bf35a3a..023e3c3ffa3a 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -438,6 +438,7 @@ mod tests { finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), } } diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 8dcf84dd9866..035aef3cd4aa 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -58,6 +58,7 @@ fn finished() -> Finished { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index db3559c44f4c..9cba9f0aa228 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -806,6 +806,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), } }); @@ -817,6 +818,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }); events.push(DecodedTextEvent::TextDelta { decoded: vllm_text::DecodedText::unattributed(chunk.delta), diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index 646c106a04ed..72d852ba2ced 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -mod array; +pub(crate) mod array; #[cfg(test)] mod tests; mod wire; diff --git a/rust/src/engine-core-client/src/protocol/logprobs/array.rs b/rust/src/engine-core-client/src/protocol/logprobs/array.rs index 4adf1ae9755f..919ef80f46a3 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/array.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/array.rs @@ -50,7 +50,7 @@ pub(super) fn decode_array2_u32( }) } -pub(super) fn decode_array1_u32( +pub(crate) fn decode_array1_u32( value: WireNdArray, field: &str, frames: &[Bytes], diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 8f4ee0612736..742fed59bdeb 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -30,6 +30,7 @@ pub mod multimodal; pub mod output; pub mod request; pub mod sampling; +pub mod sampling_mask; mod serde_utils; pub mod stats; pub mod structured_outputs; diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c8a8705435c8..a45cbe37315e 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -15,6 +15,7 @@ use super::serde_utils::AllowTrailingFields; use super::utility::UtilityOutput; use crate::error::{Error, Result, ext_value_decode}; use crate::protocol::logprobs::MaybeWireLogprobs; +use crate::protocol::sampling_mask::MaybeWireSamplingMask; use crate::protocol::stats::{PrefillStats, SchedulerStats}; use crate::protocol::{OpaqueValue, decode_msgpack}; @@ -127,7 +128,7 @@ pub struct EngineCoreOutput { #[serde(default)] pub mm_cache_miss_hashes: Option>, #[serde(default)] - pub new_sampling_mask: Option, + pub new_sampling_mask: Option, /// Per-request speculative-decoding acceptance metrics, set on the final /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here; /// the Rust frontend does not yet surface it in responses. @@ -150,6 +151,9 @@ impl EngineCoreOutput { self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take()) .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors")) .transpose()?; + self.new_sampling_mask = (self.new_sampling_mask.take()) + .map(|value| value.resolve(frames, "new_sampling_mask")) + .transpose()?; Ok(()) } } diff --git a/rust/src/engine-core-client/src/protocol/sampling_mask.rs b/rust/src/engine-core-client/src/protocol/sampling_mask.rs new file mode 100644 index 000000000000..7f5c4592a7d3 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/sampling_mask.rs @@ -0,0 +1,237 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::ops::{Deref, DerefMut}; + +use bytes::Bytes; +use enum_as_inner::EnumAsInner; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::error::{Error, Result, bail_ext_value_decode}; +use crate::protocol::logprobs::array::decode_array1_u32; +use crate::protocol::tensor::WireNdArray; + +/// Decoded sampling-mask support sets, one row per generated token position. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct SamplingMask { + pub rows: Vec>, +} + +/// Python `SamplingMaskLists` tuple before ndarray raw views are resolved. +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)] +pub struct WireSamplingMask { + token_ids: WireNdArray, + #[serde(default)] + offsets: Option, + #[serde(default)] + cu_num_generated_tokens: Option>, +} + +/// Sampling-mask field while it transitions from Python wire data to rows. +#[derive(Debug, Clone, PartialEq, EnumAsInner)] +pub enum MaybeWireSamplingMask { + Wire(Box), + Direct(SamplingMask), +} + +impl Deref for MaybeWireSamplingMask { + type Target = SamplingMask; + + fn deref(&self) -> &Self::Target { + match self { + Self::Wire(_) => panic!("sampling mask is still in wire format"), + Self::Direct(value) => value, + } + } +} + +impl DerefMut for MaybeWireSamplingMask { + fn deref_mut(&mut self) -> &mut Self::Target { + match self { + Self::Wire(_) => panic!("sampling mask is still in wire format"), + Self::Direct(value) => value, + } + } +} + +impl<'de> Deserialize<'de> for MaybeWireSamplingMask { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + WireSamplingMask::deserialize(deserializer).map(|value| Self::Wire(Box::new(value))) + } +} + +impl Serialize for MaybeWireSamplingMask { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + match self { + Self::Wire(value) => value.serialize(serializer), + Self::Direct(value) => WireSamplingMask::from_direct(value) + .map_err(serde::ser::Error::custom)? + .serialize(serializer), + } + } +} + +impl MaybeWireSamplingMask { + pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result { + match self { + Self::Direct(value) => Ok(Self::Direct(value)), + Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct), + } + } +} + +impl WireSamplingMask { + fn from_direct(value: &SamplingMask) -> std::result::Result { + let [row] = value.rows.as_slice() else { + return Err(format!( + "per-request engine-core sampling mask must contain exactly one row, got {}", + value.rows.len() + )); + }; + let token_ids = row.iter().map(|&token_id| i64::from(token_id)).collect::>(); + Ok(Self { + token_ids: WireNdArray::from_i64(vec![token_ids.len()], token_ids)?, + offsets: None, + cu_num_generated_tokens: None, + }) + } + + fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result { + if let Some(indices) = self.cu_num_generated_tokens { + bail_ext_value_decode!( + "{field_prefix}.cu_num_generated_tokens: expected None for per-request engine-core sampling-mask payload, got {indices:?}" + ); + } + + let token_ids = + decode_array1_u32(self.token_ids, &format!("{field_prefix}.token_ids"), frames)?; + let Some(offsets) = self.offsets else { + return Ok(SamplingMask { + rows: vec![token_ids], + }); + }; + let offsets = decode_array1_u32(offsets, &format!("{field_prefix}.offsets"), frames)?; + if offsets.first().copied() != Some(0) { + bail_ext_value_decode!("{field_prefix}.offsets: first offset must be zero"); + } + if offsets.windows(2).any(|pair| pair[0] > pair[1]) { + bail_ext_value_decode!("{field_prefix}.offsets: offsets must be nondecreasing"); + } + if offsets.last().copied().map(|value| value as usize) != Some(token_ids.len()) { + bail_ext_value_decode!( + "{field_prefix}.offsets: final offset {:?} does not match token_ids length {}", + offsets.last(), + token_ids.len() + ); + } + + let rows = offsets + .windows(2) + .map(|pair| token_ids[pair[0] as usize..pair[1] as usize].to_vec()) + .collect(); + Ok(SamplingMask { rows }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::dtype::{NumpyDtype, TensorDtype}; + use crate::protocol::output::{ + EngineCoreOutput, RequestBatchOutputs, decode_engine_core_outputs, + }; + use crate::protocol::tensor::WireArrayData; + use crate::protocol::{decode_msgpack, encode_msgpack}; + + #[test] + fn decodes_single_position_without_offsets() { + let wire = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![3], vec![2, 12, 16]).unwrap(), + offsets: None, + cu_num_generated_tokens: None, + }; + let encoded = encode_msgpack(&wire).unwrap(); + let decoded: MaybeWireSamplingMask = decode_msgpack(&encoded).unwrap(); + let decoded = decoded.resolve(&[], "new_sampling_mask").unwrap(); + + assert_eq!(decoded.into_direct().unwrap().rows, vec![vec![2, 12, 16]]); + } + + #[test] + fn decodes_csr_rows_and_rejects_bad_terminal_offset() { + let wire = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![4], vec![2, 12, 16, 18]).unwrap(), + offsets: Some(WireNdArray::from_i64(vec![3], vec![0, 2, 4]).unwrap()), + cu_num_generated_tokens: None, + }; + let decoded = MaybeWireSamplingMask::Wire(Box::new(wire)) + .resolve(&[], "new_sampling_mask") + .unwrap(); + assert_eq!( + decoded.into_direct().unwrap().rows, + vec![vec![2, 12], vec![16, 18]] + ); + + let malformed = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![2], vec![2, 12]).unwrap(), + offsets: Some(WireNdArray::from_i64(vec![2], vec![0, 1]).unwrap()), + cu_num_generated_tokens: None, + }; + let error = MaybeWireSamplingMask::Wire(Box::new(malformed)) + .resolve(&[], "new_sampling_mask") + .unwrap_err(); + assert!(error.to_string().contains("final offset")); + } + + #[test] + fn decodes_sampling_mask_from_multipart_aux_frames() { + let output = EngineCoreOutput { + request_id: "req-multipart-mask".to_string(), + new_token_ids: vec![16, 18], + new_sampling_mask: Some(MaybeWireSamplingMask::Wire(Box::new(WireSamplingMask { + token_ids: WireNdArray { + dtype: NumpyDtype::little(TensorDtype::I32), + shape: vec![4], + data: WireArrayData::AuxIndex(1), + }, + offsets: Some(WireNdArray { + dtype: NumpyDtype::little(TensorDtype::I64), + shape: vec![3], + data: WireArrayData::AuxIndex(2), + }), + cu_num_generated_tokens: None, + }))), + ..Default::default() + }; + let primary = encode_msgpack(&crate::protocol::output::EngineCoreOutputs::from( + RequestBatchOutputs { + outputs: vec![output], + ..Default::default() + }, + )) + .unwrap(); + let token_ids = + [2_i32, 12, 16, 18].into_iter().flat_map(i32::to_le_bytes).collect::>(); + let offsets = [0_i64, 2, 4].into_iter().flat_map(i64::to_le_bytes).collect::>(); + + let decoded = decode_engine_core_outputs(&[ + Bytes::from(primary), + Bytes::from(token_ids), + Bytes::from(offsets), + ]) + .unwrap(); + let mask = decoded.as_request_batch().unwrap().outputs[0] + .new_sampling_mask + .as_ref() + .unwrap(); + + assert_eq!(mask.rows, vec![vec![2, 12], vec![16, 18]]); + } +} diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 739c34d6b751..eb8e713fb402 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2712,15 +2712,15 @@ fn python_msgpack_fixtures_match_rust_encoding() { decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap()); assert_eq!(python_mm_features, rust_mm_features); - let decoded_sampling_mask_outputs: EngineCoreOutputs = - rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap(); + let decoded_sampling_mask_outputs = + decode_engine_core_outputs(&[bytes::Bytes::from(sampling_mask_outputs_bytes)]).unwrap(); let sampling_mask_output = &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0]; assert!(sampling_mask_output.mm_cache_miss_hashes.is_none()); - assert!(matches!( - sampling_mask_output.new_sampling_mask.as_ref(), - Some(Value::Array(fields)) if fields.len() == 3 - )); + assert_eq!( + sampling_mask_output.new_sampling_mask.as_ref().unwrap().rows, + vec![vec![2, 12, 16, 17, 18]] + ); let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap(); // Match msgspec's base-schema result for the same extended Python message. diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 5013f9e2a57b..a19dc9fdaa45 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -234,14 +234,10 @@ class EngineCoreOutputs( [ ", + /// Sampling support sets aligned one-to-one with generated token positions. + pub sampling_mask: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -155,6 +158,8 @@ pub struct GenerateOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned one-to-one with `token_ids`. + pub sampling_mask: Option, } impl GenerateOutput { @@ -202,6 +207,7 @@ impl GenerateOutput { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } } @@ -271,6 +277,16 @@ impl Stream for GenerateOutputStream { } let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap()); + let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap()); + if let Some(mask) = sampling_mask.as_ref() + && mask.rows.len() != raw.new_token_ids.len() + { + return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { + request_id: raw.request_id, + token_count: raw.new_token_ids.len(), + row_count: mask.rows.len(), + }))); + } let cached_token_count = raw .prefill_stats .as_ref() @@ -291,6 +307,7 @@ impl Stream for GenerateOutputStream { cached_token_count, kv_transfer_params: raw.kv_transfer_params, ec_transfer_params: raw.ec_transfer_params, + sampling_mask, }; Poll::Ready(Some(Ok(output))) @@ -341,6 +358,7 @@ impl> + Send> T { while let Some(output) = stream.next().await.transpose()? { cached_token_count = cached_token_count.max(output.cached_token_count); + let sampling_mask = output.sampling_mask; if let Some(info) = output.prompt_info { if prompt_token_ids.is_none() { prompt_token_ids = Some(info.prompt_token_ids.to_vec()); @@ -359,6 +377,9 @@ impl> + Send> T { existing.logprobs = Some(step_logprobs); } } + if let Some(mut mask) = sampling_mask { + existing.sampling_mask.get_or_insert_default().rows.append(&mut mask.rows); + } } else { collected = Some(CollectedGenerateOutput { request_id: output.request_id, @@ -374,6 +395,7 @@ impl> + Send> T { }, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask, }); } @@ -387,6 +409,15 @@ impl> + Send> T { }; collected.kv_transfer_params = output.kv_transfer_params; collected.ec_transfer_params = output.ec_transfer_params; + if let Some(mask) = collected.sampling_mask.as_ref() + && mask.rows.len() != collected.token_ids.len() + { + return Err(crate::Error::SamplingMaskTokenCountMismatch { + request_id: collected.request_id, + token_count: collected.token_ids.len(), + row_count: mask.rows.len(), + }); + } return Ok(collected); } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 2d09aff889a8..09b519976772 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; -use std::sync::Once; +use std::sync::{Arc, Once}; use std::time::Duration; use futures::StreamExt as _; @@ -18,6 +18,7 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::sampling_mask::{MaybeWireSamplingMask, SamplingMask}; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; @@ -326,6 +327,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { num_local_cached_tokens: 1, ..Default::default() }), + new_sampling_mask: Some(MaybeWireSamplingMask::Direct( + SamplingMask { + rows: vec![vec![1, 33, 99]], + }, + )), ..request_output_with_logprobs( &request.request_id, vec![33], @@ -334,15 +340,22 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { Some(prompt_logprobs()), ) }, - request_output_with_logprobs_and_kv( - &request.request_id, - vec![44], - Some(EngineCoreFinishReason::Stop), - Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), - None, - Some(serde_json::json!({"connector": "x"})), - None, - ), + EngineCoreOutput { + new_sampling_mask: Some(MaybeWireSamplingMask::Direct( + SamplingMask { + rows: vec![vec![2, 44, 88]], + }, + )), + ..request_output_with_logprobs_and_kv( + &request.request_id, + vec![44], + Some(EngineCoreFinishReason::Stop), + Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), + None, + Some(serde_json::json!({"connector": "x"})), + None, + ) + }, ], ..Default::default() } @@ -375,6 +388,42 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { collected.kv_transfer_params, Some(serde_json::json!({"connector": "x"})) ); + assert_eq!( + collected.sampling_mask, + Some(SamplingMask { + rows: vec![vec![1, 33, 99], vec![2, 44, 88]], + }) + ); +} + +#[tokio::test] +async fn collect_output_rejects_partial_sampling_mask() { + let output = vllm_llm::GenerateOutput { + request_id: "req-partial-mask".to_string(), + prompt_info: Some(GeneratePromptInfo { + prompt_token_ids: Arc::from([11_u32, 22]), + prompt_logprobs: None, + }), + token_ids: vec![33, 44], + logprobs: None, + finish_reason: Some(FinishReason::Length), + cached_token_count: 0, + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: Some(SamplingMask { + rows: vec![vec![1, 33, 99]], + }), + }; + + let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); + assert!(matches!( + error, + Error::SamplingMaskTokenCountMismatch { + token_count: 2, + row_count: 1, + .. + } + )); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index caccbdefc5ac..8990ba2f3bc3 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -19,6 +19,7 @@ itertools.workspace = true libc.workspace = true llm-multimodal.workspace = true openssl.workspace = true +prost.workspace = true prost-types.workspace = true rmp-serde.workspace = true rmpv.workspace = true diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 5f6ce5d42bc7..0f28e147ef07 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -316,6 +316,11 @@ pub fn to_sequence_output( _ => (vec![], vec![], vec![]), }; + let sampling_mask = finished + .and_then(|finished| finished.sampling_mask.as_ref()) + .map(|mask| mask.rows.iter().map(|row| pb::TokenIds { ids: row.clone() }).collect()) + .unwrap_or_default(); + Ok(pb::SequenceOutput { index: 0, // TODO: multi-sequence (n > 1) not supported text: if opts.output_text { @@ -333,6 +338,7 @@ pub fn to_sequence_output( ranks: rank_values, candidate_tokens: candidates, finish_info, + sampling_mask, }) } @@ -524,6 +530,7 @@ impl ResponseOpts { mod tests { use prost::Message as _; use vllm_engine_core_client::protocol::output::StopReason; + use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_text::{ FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, }; @@ -720,6 +727,7 @@ mod tests { finish_reason: reason, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } @@ -804,4 +812,40 @@ mod tests { assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32); assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30))); } + + #[test] + fn sequence_output_only_carries_sampling_mask_on_terminal_output() { + let mut fin = finished(FinishReason::Length); + fin.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, 10], vec![2, 20]], + }); + + let terminal = + to_sequence_output("", &[10, 20], None, Some(&fin), &ResponseOpts::default()) + .expect("terminal output"); + let intermediate = to_sequence_output("", &[9], None, None, &ResponseOpts::default()) + .expect("intermediate output"); + + assert_eq!( + terminal.sampling_mask, + vec![ + pb::TokenIds { ids: vec![1, 10] }, + pb::TokenIds { ids: vec![2, 20] } + ] + ); + assert!(intermediate.sampling_mask.is_empty()); + } + + #[test] + fn sampling_mask_uses_additive_proto_tag_ten() { + let response = pb::SequenceOutput { + sampling_mask: vec![pb::TokenIds { ids: vec![7, 8] }], + ..Default::default() + }; + + assert_eq!( + response.encode_to_vec(), + vec![0x52, 0x04, 0x0a, 0x02, 0x07, 0x08] + ); + } } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 72943b6af0bb..d56415c7475e 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -257,6 +257,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl { finish_reason: collected.finish_reason, kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, + sampling_mask: collected.sampling_mask, }; let outputs = convert::to_sequence_output( diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index dfb7e6cbb191..f82059dffda4 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -481,6 +481,7 @@ mod tests { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -494,6 +495,7 @@ mod tests { cached_token_count: 2, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ]); @@ -575,6 +577,7 @@ mod tests { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } @@ -658,6 +661,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, prompt_token_ids: vec![10, 20], + sampling_mask: None, }; let response = collect_generate( @@ -744,6 +748,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, prompt_token_ids: vec![10, 20], + sampling_mask: None, }; let response = collect_generate( @@ -776,6 +781,7 @@ mod tests { }, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index a122afbcd97a..06874a43bd31 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -694,6 +694,7 @@ mod tests { ))), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -806,6 +807,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -858,6 +860,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -913,6 +916,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -985,6 +989,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -1059,6 +1064,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index a438c8e267d6..d3dc42265641 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize}; use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder}; @@ -50,6 +51,8 @@ pub struct Finished { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned with all generated token positions. + pub sampling_mask: Option, } /// Sample metadata emitted by one engine output update. @@ -120,6 +123,7 @@ pub async fn decoded_text_event_stream( let mut sampled_token_ids = Vec::new(); let mut output_token_count: usize = 0; let mut sampled_logprobs: Option = None; + let mut sampling_mask: Option = None; while let Some(next) = raw_stream.next().await { let output = next?; @@ -213,6 +217,7 @@ pub async fn decoded_text_event_stream( let mut new_token_ids = output.token_ids; let mut new_logprobs = output.logprobs; + let mut new_sampling_mask = output.sampling_mask; // Trim tokens and logprobs if we matched stop string. if let Some(num_tokens) = truncate_tokens_to { @@ -220,9 +225,15 @@ pub async fn decoded_text_event_stream( if let Some(logprobs) = &mut new_logprobs { logprobs.positions.truncate(num_tokens); } + if let Some(mask) = &mut new_sampling_mask { + mask.rows.truncate(num_tokens); + } } output_token_count += new_token_ids.len(); + if let Some(mut mask) = new_sampling_mask { + sampling_mask.get_or_insert_default().rows.append(&mut mask.rows); + } let decoded_logprobs = new_logprobs .as_ref() @@ -246,6 +257,16 @@ pub async fn decoded_text_event_stream( } if let Some(reason) = finish_reason { + if let Some(mask) = sampling_mask.as_ref() + && mask.rows.len() != output_token_count + { + return Err(vllm_llm::Error::SamplingMaskTokenCountMismatch { + request_id, + token_count: output_token_count, + row_count: mask.rows.len(), + } + .into()); + } // Flush any remaining buffered text. let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?; let text_len = full_decoded.text.len(); @@ -293,6 +314,7 @@ pub async fn decoded_text_event_stream( finish_reason: reason, kv_transfer_params, ec_transfer_params, + sampling_mask, })), }) .await; @@ -354,6 +376,7 @@ mod tests { use futures::{Stream, StreamExt as _, stream}; use vllm_engine_core_client::AbortCause; + use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use vllm_llm::GenerateOutput; use vllm_tokenizer::{TokenAnchor, TokenAttribution, test_utils::TestTokenizer}; @@ -384,6 +407,73 @@ mod tests { s.bytes().map(u32::from).collect() } + #[tokio::test] + async fn sampling_masks_survive_text_decode_and_collection() { + let prompt: Arc<[u32]> = Arc::from([]); + let mut first = GenerateOutput::for_test(Some(prompt), vec![b'a' as u32], None); + first.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, b'a' as u32]], + }); + let mut terminal = + GenerateOutput::for_test(None, vec![b'b' as u32], Some(FinishReason::Length)); + terminal.sampling_mask = Some(SamplingMask { + rows: vec![vec![2, b'b' as u32]], + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let collected = decoded_text_event_stream( + "test-mask".into(), + tokenizer, + stream::iter([Ok(first), Ok(terminal)]), + TextDecodeOptions::default(), + false, + ) + .collect_output() + .await + .unwrap(); + + assert_eq!( + collected.sampling_mask, + Some(SamplingMask { + rows: vec![vec![1, b'a' as u32], vec![2, b'b' as u32]], + }) + ); + } + + #[tokio::test] + async fn sampling_mask_mismatch_reuses_llm_error() { + let prompt: Arc<[u32]> = Arc::from([]); + let mut terminal = GenerateOutput::for_test( + Some(prompt), + vec![b'a' as u32, b'b' as u32], + Some(FinishReason::Length), + ); + terminal.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, b'a' as u32]], + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let error = decoded_text_event_stream( + "test-mask-mismatch".into(), + tokenizer, + stream::iter([Ok(terminal)]), + TextDecodeOptions::default(), + false, + ) + .collect_output() + .await + .unwrap_err(); + + assert!(matches!( + error, + crate::Error::Llm(vllm_llm::Error::SamplingMaskTokenCountMismatch { + token_count: 2, + row_count: 1, + .. + }) + )); + } + fn opts(stop: &[&str], min_tokens: u32) -> TextDecodeOptions { TextDecodeOptions { stop_strings: Some(stop.iter().map(|s| s.to_string()).collect()), @@ -472,6 +562,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ), ], @@ -503,6 +594,7 @@ mod tests { ))), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ), ], @@ -576,6 +668,60 @@ mod tests { assert!(output.finish_reason.is_stop()); } + #[tokio::test] + async fn stream_stop_string_mid_chunk_keeps_token_logprob_mask_alignment() { + let token_ids = ascii_tokens("abcxyz"); + let mut raw = GenerateOutput::for_test( + Some(Arc::from([])), + token_ids.clone(), + Some(FinishReason::Length), + ); + raw.logprobs = Some(Logprobs { + positions: token_ids + .iter() + .map(|&token_id| PositionLogprobs { + entries: vec![TokenLogprob { + token_id, + logprob: -0.25, + rank: 1, + }], + }) + .collect(), + }); + raw.sampling_mask = Some(SamplingMask { + rows: token_ids.iter().map(|&token_id| vec![token_id, token_id + 100]).collect(), + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let output = decoded_text_event_stream( + "test-mid-chunk-stop".into(), + tokenizer, + stream::iter([Ok(raw)]), + opts(&["cx"], 0), + false, + ) + .collect_output() + .await + .unwrap(); + + assert_eq!(output.text, "ab"); + assert!(output.finish_reason.is_stop()); + assert_eq!(output.token_ids, ascii_tokens("abcx")); + assert_eq!(output.logprobs.as_ref().unwrap().positions.len(), 4); + assert_eq!( + output.sampling_mask, + Some(SamplingMask { + rows: vec![ + vec![b'a' as u32, b'a' as u32 + 100], + vec![b'b' as u32, b'b' as u32 + 100], + vec![b'c' as u32, b'c' as u32 + 100], + vec![b'x' as u32, b'x' as u32 + 100], + ], + }) + ); + assert_eq!(output.usage.output_token_count, 4); + } + #[tokio::test] async fn stream_stop_string_at_end() { let output = run_to_completion(ascii_tokens("abcxyz"), opts(&["xyz"], 0)).await; diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index a415bb8900c2..1293ac6025ff 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -17,6 +17,7 @@ mod logprobs; use std::sync::Arc; use futures::{StreamExt as _, pin_mut}; +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use crate::{Error, FinishReason, Result, TextOutputStream}; @@ -35,6 +36,8 @@ pub struct CollectedTextOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned with generated token positions. + pub sampling_mask: Option, } #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")] @@ -90,6 +93,7 @@ impl T { usage: vllm_llm::TokenUsage::default(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }) }; @@ -99,6 +103,7 @@ impl T { collected.usage = finished.usage; collected.kv_transfer_params = finished.kv_transfer_params; collected.ec_transfer_params = finished.ec_transfer_params; + collected.sampling_mask = finished.sampling_mask; return Ok(collected); } } @@ -173,6 +178,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), }), ]); @@ -295,6 +301,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), }), ]); From 8da75d61904cd0b2c5c322746ea15251773ec51f Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 18 Sep 2026 16:39:27 +0800 Subject: [PATCH 460/774] C3x SM100 FP8 blockwise - Pad activation scales to multiple of 4 (#57377) Signed-off-by: Thien Tran Co-authored-by: Codex --- .buildkite/test_areas/kernels.yaml | 3 ++ ...scaled_mm_blockwise_sm100_fp8_dispatch.cuh | 33 ++++++++++++++++--- tests/kernels/quantization/test_block_fp8.py | 9 +++-- 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 583378fdb49b..207d827f3600 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -367,6 +367,7 @@ steps: - csrc/quantization/fp4/ - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu - csrc/libtorch_stable/quantization/fp4/ + - csrc/libtorch_stable/quantization/w8a8/cutlass/ - csrc/attention/mla/ - csrc/quantization/cutlass_w8a8/moe/ - vllm/model_executor/layers/fusion/ @@ -408,6 +409,7 @@ steps: - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py + - tests/kernels/quantization/test_block_fp8.py - tests/models/quantization/test_nvfp4.py commands: - nvidia-smi @@ -442,6 +444,7 @@ steps: pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2' pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8' + pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh index 8f74f2991469..fa2add918213 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh @@ -152,6 +152,29 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te int32_t m = a.size(0), n = b.size(1), k = a.size(1); + // The SM100 blockwise SF copy requires the SFA M extent to be a multiple + // of 4 (4-row copy atom). For misaligned M on the non-swapAB path, pad only + // the per-token activation scales (column-major (m, k/128), 1/64 the size + // of the bf16 activation) and build layout_SFA from the padded extent. + int32_t m_sf = m; + std::optional a_scales_pad; + void const* a_scales_data = a_scales.data_ptr(); + if constexpr (!swap_ab) { + if (m % 4 != 0) { + m_sf = (m + 3) & ~3; + int64_t num_groups = a_scales.numel() / m; + a_scales_pad = torch::stable::new_empty( + a_scales, {num_groups * m_sf}, + torch::headeronly::ScalarType::Float); + cudaMemcpy2DAsync(a_scales_pad->data_ptr(), m_sf * sizeof(float), + a_scales.data_ptr(), m * sizeof(float), + m * sizeof(float), num_groups, + cudaMemcpyDeviceToDevice, + get_current_cuda_stream(a.get_device())); + a_scales_data = a_scales_pad->data_ptr(); + } + } + StrideA a_stride; StrideB b_stride; StrideC c_stride; @@ -162,16 +185,16 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te c_stride = cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1)); - LayoutSFA layout_SFA = swap_ab ? + LayoutSFA layout_SFA = swap_ab ? ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) : - ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + ScaleConfig::tile_atom_to_shape_SFA(make_shape(m_sf, n, k, 1)); LayoutSFB layout_SFB = swap_ab ? ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) : ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); auto a_ptr = static_cast(a.data_ptr()); auto b_ptr = static_cast(b.data_ptr()); - auto a_scales_ptr = static_cast(a_scales.data_ptr()); + auto a_scales_ptr = static_cast(a_scales_data); auto b_scales_ptr = static_cast(b_scales.data_ptr()); typename GemmKernel::MainloopArguments mainloop_args{}; @@ -212,7 +235,9 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out, constexpr int TILE_K = 128; // TODO: better heuristics - bool swap_ab = (m < 16) || (m % 4 != 0); + // swapAB only wins for small M; misaligned M >= 65 is faster on the + // non-swapAB path with padded activation scales (measured on GB300). + bool swap_ab = (m <= 64); bool use_tma_epilogue = (m * n) % 4 == 0; if (!swap_ab) { constexpr int TILE_N = 128; diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index d02527aba048..038b94bd2d74 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -162,11 +162,16 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed): @pytest.mark.skipif( not current_platform.is_cuda(), reason="CUTLASS only supported on CUDA platform." ) +@pytest.mark.parametrize( + # 65/66/67 cover all M%4 residue classes above the SM100 swapAB + # threshold (m <= 64); 1026 crosses multiple 128-row SF atoms. + "M", + [32, 65, 66, 67, 1026], +) @torch.inference_mode() -def test_w8a8_block_fp8_cutlass_matmul(): +def test_w8a8_block_fp8_cutlass_matmul(M): # Test simple case where weight.shape % 128 != 0, # like in DSV3 kv_a_proj_with_mqa - M = 32 N = 576 K = 7168 block_size = [128, 128] From 32636580a6f1c3bc41deefc4bf7800f850c24034 Mon Sep 17 00:00:00 2001 From: olka-amd Date: Fri, 18 Sep 2026 11:49:53 +0300 Subject: [PATCH 461/774] [Bugfix][MLA] TritonMLA: fix illegal memory access on causal multi-token decode (#51065) Signed-off-by: Oleksandr Kachur Co-authored-by: Claude Opus 5 --- .../test_triton_mla_causal_verify_flatten.py | 215 ++++++++++++++++++ vllm/v1/attention/backends/mla/triton_mla.py | 91 +++----- 2 files changed, 252 insertions(+), 54 deletions(-) create mode 100644 tests/kernels/attention/test_triton_mla_causal_verify_flatten.py diff --git a/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py b/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py new file mode 100644 index 000000000000..a27552c90e3e --- /dev/null +++ b/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the multi-token decode flatten in ``TritonMLAImpl.forward_mqa``. + +``decode_attention_fwd`` has no causal flag: it launches one program per row of +``q`` and reads that row's KV extent from ``B_Seqlen``. So a ``query_len``-token +decode block is flattened to one row per query token, and intra-block causality +has to be expressed through the per-row sequence lengths -- row ``t`` gets +``seq_len - (query_len - 1) + t``, i.e. the committed prefix plus block tokens +``0..t``. A non-causal draft block instead gives every row the same full extent. + +The kernel is replaced by a spy, so these tests assert on the ``block_table`` +and ``seq_lens`` the backend actually submits and need no GPU or Triton runtime. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.mla.triton_mla import ( + TritonMLAImpl, + TritonMLAMetadataBuilder, +) + +NUM_HEADS = 16 +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM +PAGE_SIZE = 16 + +# Committed context per request, including a fresh request at zero context +# where the whole KV extent is the block itself. A decode block's seq_len +# already counts the block, so seq_len = context + query_len. +CONTEXT_LENS = [0, 1, 37, 512] + + +def _seq_lens(query_len: int) -> list[int]: + return [c + query_len for c in CONTEXT_LENS] + + +def _make_impl(dcp_world_size: int = 1) -> TritonMLAImpl: + """Only the attributes ``forward_mqa`` reads; __init__ wants a full config.""" + impl = object.__new__(TritonMLAImpl) + impl.kv_lora_rank = KV_LORA_RANK + impl.scale = HEAD_SIZE**-0.5 + impl._sm_count = 304 + impl.dcp_world_size = dcp_world_size + return impl + + +def _make_metadata(seq_lens: list[int], query_len: int, causal: bool): + num_decodes = len(seq_lens) + return SimpleNamespace( + num_decodes=num_decodes, + num_decode_tokens=num_decodes * query_len, + max_seq_len=max(seq_lens), + causal=causal, + decode=SimpleNamespace( + # arange page ids so each request's window is an identifiable slice. + block_table=torch.arange(num_decodes * 64, dtype=torch.int32).reshape( + num_decodes, 64 + ), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + ), + ) + + +def _run_forward_mqa( + seq_lens: list[int], + query_len: int, + causal: bool, + q_rows=None, + dcp_world_size: int = 1, +): + """Drive forward_mqa with the decode kernel spied out. + + Returns the ``(block_table, seq_lens)`` handed to ``decode_attention_fwd``. + """ + captured: dict = {} + + def spy(q, kv_cache, kv_c_cache, o, lse, block_table, b_seq_len, *args, **kwargs): + captured["block_table"] = block_table.detach().clone() + captured["seq_lens"] = b_seq_len.detach().clone() + + metadata = _make_metadata(seq_lens, query_len, causal) + rows = metadata.num_decode_tokens if q_rows is None else q_rows + q = torch.zeros(rows, NUM_HEADS, HEAD_SIZE, dtype=torch.bfloat16) + kv_cache = torch.zeros( + len(seq_lens) * 64, PAGE_SIZE, HEAD_SIZE, dtype=torch.bfloat16 + ) + layer = SimpleNamespace(_k_scale=torch.ones(1), layer_name="test") + + with patch("vllm.v1.attention.backends.mla.triton_mla.decode_attention_fwd", spy): + _make_impl(dcp_world_size).forward_mqa(q, kv_cache, metadata, layer) + + assert captured, "forward_mqa did not reach the decode kernel" + return captured + + +def test_multi_token_decode_flags_are_declared_together(): + """Capture probes max_query_len > 1, which only clears the assert in + build_for_cudagraph_capture once query_len_support raised the threshold.""" + assert ( + TritonMLAMetadataBuilder._cudagraph_support == AttentionCGSupport.UNIFORM_BATCH + ) + assert TritonMLAMetadataBuilder.query_len_support == QueryLenSupport.UNIFORM + + +@pytest.mark.parametrize("query_len", [2, 3, 4, 5, 8]) +def test_causal_block_rows_see_prefix_plus_own_position(query_len): + """Causal verify: row t sees seq_len - (query_len - 1) + t KV entries. + + Regression guard: the causal branch used to be skipped entirely, so a verify + token could attend the draft siblings it is supposed to be checking. Fails + on unmodified upstream, where no flatten happens at all for causal blocks. + """ + seq_lens = _seq_lens(query_len) + captured = _run_forward_mqa(seq_lens, query_len, causal=True) + + got = captured["seq_lens"].tolist() + want = [s - (query_len - 1) + t for s in seq_lens for t in range(query_len)] + assert got == want, ( + "causal flatten is not causal: verify row r*query_len+t must get " + f"seq_len_r - {query_len - 1} + t entries.\n" + f" seq_lens {seq_lens}\n got {got}\n expected {want}" + ) + + # Per-request invariants the arithmetic must hold regardless of query_len. + for r, seq_len in enumerate(seq_lens): + rows = got[r * query_len : (r + 1) * query_len] + assert rows == sorted(rows) and len(set(rows)) == query_len, ( + f"request {r} rows must strictly increase, got {rows}" + ) + assert rows[-1] == seq_len, "last row must see the full sequence" + assert rows[0] == seq_len - (query_len - 1) + assert all(n > 0 for n in rows), f"non-positive KV extent in {rows}" + + +def test_cudagraph_padding_rows_present_no_kv_extent(): + """Causal flattening pins cudagraph padding rows (seq_len 0) to exactly + zero KV extent, keeping them in the kernel's skipped branch.""" + query_len = 4 + seq_lens = _seq_lens(query_len) + [0] + captured = _run_forward_mqa(seq_lens, query_len, causal=True) + + padding = captured["seq_lens"].tolist()[-query_len:] + assert padding == [0] * query_len, ( + f"padding rows must present exactly zero KV extent, got {padding}" + ) + + +@pytest.mark.parametrize("query_len", [2, 4]) +def test_non_causal_block_rows_all_see_the_full_prefix(query_len): + """The DSpark draft block is generated in one pass, so no row is masked.""" + seq_lens = _seq_lens(query_len) + captured = _run_forward_mqa(seq_lens, query_len, causal=False) + + want = [s for s in seq_lens for _ in range(query_len)] + assert captured["seq_lens"].tolist() == want + + +@pytest.mark.parametrize("causal", [True, False]) +def test_block_table_is_expanded_to_match_q_rows(causal): + """The kernel indexes block_table by program_id, so it needs one row per + query token -- a short block_table is read out of bounds.""" + query_len = 4 + captured = _run_forward_mqa(_seq_lens(query_len), query_len, causal=causal) + + num_rows = len(CONTEXT_LENS) * query_len + assert captured["block_table"].shape[0] == num_rows + assert captured["seq_lens"].shape[0] == num_rows + # repeat_interleave, not repeat: a request's rows must stay adjacent. + for r in range(len(CONTEXT_LENS)): + block = captured["block_table"][r * query_len : (r + 1) * query_len] + assert torch.equal(block, block[:1].expand_as(block)) + + +@pytest.mark.parametrize("causal", [True, False]) +def test_single_token_decode_is_untouched(causal): + """query_len == 1 must short-circuit: ordinary decode sees no expansion.""" + seq_lens = _seq_lens(1) + captured = _run_forward_mqa(seq_lens, query_len=1, causal=causal) + + assert captured["seq_lens"].tolist() == seq_lens + assert captured["block_table"].shape[0] == len(seq_lens) + + +def test_expansion_factor_follows_the_q_row_count(): + """The kernel indexes block_table and seq_lens by program_id, so the + expansion must follow q rows, not the token/request counts.""" + num_rows = len(CONTEXT_LENS) * 4 + captured = _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=num_rows) + + assert captured["block_table"].shape[0] == num_rows + assert captured["seq_lens"].shape[0] == num_rows + + +def test_non_uniform_block_is_rejected(): + """Q rows must divide evenly across the decode requests.""" + with pytest.raises(AssertionError, match="non-uniform decode block"): + _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=6) + + +def test_causal_multi_token_decode_is_rejected_under_dcp(): + """Per-row extents offset the global seq_len, but DCP passes the rank-local + slice, so the causal arithmetic would silently address the wrong KV.""" + with pytest.raises(AssertionError, match="not supported with DCP"): + _run_forward_mqa(_seq_lens(4), query_len=4, causal=True, dcp_world_size=2) + + # Non-causal rows all take the same extent, which stays correct when local. + _run_forward_mqa(_seq_lens(4), query_len=4, causal=False, dcp_world_size=2) diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 78787d5c264d..3deb996a5b3e 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -6,7 +6,6 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( @@ -14,6 +13,7 @@ MLACommonImpl, MLACommonMetadata, MLACommonMetadataBuilder, + QueryLenSupport, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability @@ -26,7 +26,6 @@ MultipleOf, ) from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd -from vllm.v1.kv_cache_interface import KVCacheSpec from vllm.v1.worker.workspace import ( current_workspace_manager, is_workspace_manager_initialized, @@ -50,50 +49,16 @@ def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int: class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): - _cudagraph_support: ClassVar[AttentionCGSupport] = ( - AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - ) - # Non-causal DSpark block is flattened to one decode row per query token in - # forward_mqa, so no intra-block causal masking is required. + # forward_mqa flattens a uniform multi-token block to one decode row per + # query token, so causal and non-causal blocks both take the decode path. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM supports_non_causal_multi_token_decode: ClassVar[bool] = True - @classmethod - def get_cudagraph_support( - cls, - vllm_config: VllmConfig, - kv_cache_spec: KVCacheSpec, - ) -> AttentionCGSupport: - """Report UNIFORM_BATCH where a non-causal multi-token block is served. - - ``_cudagraph_support`` is a class constant, so serving the DSpark - draft's (1 + num_spec) block through the decode path reports - UNIFORM_SINGLE_TOKEN_DECODE and, because the engine takes the minimum - over all attention groups, downgrades the *whole* engine off full - cudagraphs. ``forward_mqa`` flattens that block with - ``repeat_interleave`` on a Python int and performs no device->host - sync, so it does satisfy the UNIFORM_BATCH contract. - - ``non_causal_multi_token_decode`` is a KV-cache-group property, not a - per-layer one: ``MLAAttentionSpec.merge`` ORs it over every layer in - the group, so a group holding both a draft and its target reports it - for both. That is the same predicate ``__init__`` below already uses to - raise ``reorder_batch_threshold``, so the two stay consistent, but it - does mean this lifts a causal target sharing the draft's KV cache group - as well. - """ - if getattr(kv_cache_spec, "non_causal_multi_token_decode", False): - return AttentionCGSupport.UNIFORM_BATCH - return cls._cudagraph_support - def __init__(self, kv_cache_spec, layer_names, vllm_config, device): super().__init__(kv_cache_spec, layer_names, vllm_config, device) # DCP local sequence lengths are not advanced between draft steps. self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 - # Only the non-causal DSpark draft group serves multi-token blocks via - # the decode path; raise its reorder threshold to the spec block length - # so full-cudagraph capture admits it. Causal usage stays single-token. - if getattr(self, "non_causal_multi_token_decode", False): - self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self._reserve_attn_logits_workspace() def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None: @@ -109,12 +74,12 @@ def _reserve_attn_logits_workspace(self) -> None: """ if not is_workspace_manager_initialized(): return - # Decode reorder threshold is 1, so decode tokens <= max_num_seqs. - B = self.vllm_config.scheduler_config.max_num_seqs - # Non-causal DSpark draft flattens each request's block to query_len - # decode rows; cover max_num_seqs * block_len rows. - if getattr(self, "non_causal_multi_token_decode", False): - B *= self.reorder_batch_threshold + # forward_mqa flattens each request's block to query_len decode rows, + # and query_len is bounded by the reorder threshold. + B = ( + self.vllm_config.scheduler_config.max_num_seqs + * self.reorder_batch_threshold + ) # DCP all-gathers the query heads before forward_mqa. q_num_heads = self.num_heads * self.dcp_world_size max_splits = _compute_num_kv_splits( @@ -313,14 +278,32 @@ def forward_mqa( block_table = attn_metadata.decode.block_table seq_lens = attn_metadata.decode.seq_lens - if not attn_metadata.causal: - # Non-causal DSpark block: flatten to one decode row per query token. - # Each row attends to the same committed KV prefix (per-row seq_lens) - # and never to sibling block tokens = non-causal block semantics. - # Mirrors FlashInferMLA's non-causal path. - query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes - if query_len > 1: - block_table = block_table.repeat_interleave(query_len, dim=0) + # decode_attention_fwd has no causal flag: it launches one program per + # row of q and reads that row's KV extent from seq_lens, so intra-block + # causality is expressed as per-row extents. Deriving query_len from the + # tensors the kernel indexes keeps the three row counts in step. + num_decodes = seq_lens.shape[0] + query_len, remainder = divmod(B, num_decodes) if num_decodes else (1, 0) + assert remainder == 0, ( + f"non-uniform decode block: {B} query rows over {num_decodes} requests" + ) + if query_len > 1: + block_table = block_table.repeat_interleave(query_len, dim=0) + if attn_metadata.causal: + # Per-row extents are offsets off the global sequence length; + # under DCP seq_lens is this rank's local slice instead. + assert self.dcp_world_size == 1, ( + "causal multi-token decode is not supported with DCP" + ) + # Row t attends the prefix plus block tokens 0..t. Clamp holds + # padding rows at extent 0; the kernel skips them either way + # (split_kv_end == split_kv_start). + offsets = torch.arange( + 1 - query_len, 1, device=seq_lens.device, dtype=seq_lens.dtype + ) + seq_lens = (seq_lens.unsqueeze(1) + offsets).flatten().clamp(min=0) + else: + # Non-causal draft block: every row sees the same prefix. seq_lens = seq_lens.repeat_interleave(query_len) # Run MQA — always pass layer scales. When KV cache is From 67a8a3f9269795d683868254d134a199c29aff4f Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Fri, 18 Sep 2026 02:05:13 -0700 Subject: [PATCH 462/774] [Bugfix][CPU] Fix macOS multimodal SHM cache initialization (#57142) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- tests/test_envs.py | 5 +++++ vllm/envs.py | 10 +++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/test_envs.py b/tests/test_envs.py index 1ef40cde0cc0..dfd08a9ec287 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -18,6 +18,11 @@ from vllm.exceptions import VLLMValidationError +def test_object_storage_shm_default_name(): + """The generated name must fit macOS's shared-memory name limit.""" + assert len(envs._generate_shm_name()) <= 30 + + def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch): assert envs.VLLM_HOST_IP == "" assert envs.VLLM_PORT is None diff --git a/vllm/envs.py b/vllm/envs.py index 82d6e771b528..26c00c6da623 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -535,6 +535,14 @@ def get_vllm_port() -> int | None: raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err +def _generate_shm_name() -> str: + import pybase64 + + # Fit macOS's 30-character limit without truncating the UUID. + encoded_uuid = pybase64.urlsafe_b64encode(uuid.uuid4().bytes).decode("ascii") + return "vllm_mm_" + encoded_uuid.rstrip("=") + + def get_env_or_set_default( env_name: str, default_factory: Callable[[], str], @@ -1943,7 +1951,7 @@ def _resolve_rust_cli_path() -> str | None: # if not explicitly set. "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME": get_env_or_set_default( "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", - lambda: f"VLLM_OBJECT_STORAGE_SHM_BUFFER_{uuid.uuid4().hex}", + _generate_shm_name, ), # The size in MB of the buffers (NVL and RDMA) used by DeepEP "VLLM_DEEPEP_BUFFER_SIZE_MB": lambda: int( From 6df2b1a8c64a7956c51410cbfe1b786a69574263 Mon Sep 17 00:00:00 2001 From: meng-amd Date: Fri, 18 Sep 2026 17:52:46 +0800 Subject: [PATCH 463/774] [Quantization] Support native Quark W4A16 INT4/UINT4 exports in vLLM (#48606) Signed-off-by: HongWei Meng --- tests/quantization/test_quark.py | 440 +++++++++++++++++- .../model_executor/kernels/linear/__init__.py | 1 + .../layers/fused_moe/experts/triton_moe.py | 2 +- .../layers/quantization/auto_awq.py | 6 +- .../layers/quantization/quark/quark.py | 56 ++- .../layers/quantization/quark/quark_moe.py | 354 +++++++++++++- .../quantization/quark/schemes/__init__.py | 2 + .../quark/schemes/quark_w4a16_int4.py | 212 +++++++++ .../layers/quantization/quark/utils.py | 57 +++ 9 files changed, 1124 insertions(+), 6 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index ced034cb1a94..39481456fd6d 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -58,12 +58,17 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501 QuarkMoEMethod, QuarkW4A8Fp8MoEMethod, + QuarkW4A16Int4MoEMethod, QuarkW8A8Fp8MoEMethod, QuarkW8A8Int8MoEMethod, ) -from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme +from vllm.model_executor.layers.quantization.quark.schemes import ( + QuarkScheme, + QuarkW4A16Int4, +) from vllm.model_executor.layers.quantization.quark.utils import ( QuarkQTensorHint, + canonicalize_quark_packed_int4, should_ignore_layer, ) from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( @@ -97,6 +102,7 @@ kNvfp4Static, ) from vllm.model_executor.models.llama import LlamaForCausalLM +from vllm.model_executor.models.utils import WeightsMapper from vllm.platforms import current_platform from vllm.transformers_utils.repo_utils import hf_api @@ -1876,3 +1882,435 @@ def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv(): ) assert target.loaded_shard is None # o_proj is a direct (non-fused) proj assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) + + +_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7] + + +def _quark_int4_config( + *, + pack_method: str = "reorder", + symmetric: bool = True, + exclude: list[str] | None = None, +) -> dict: + return { + "quant_method": "quark", + "export": {"pack_method": pack_method, "kv_cache_group": []}, + "global_quant_config": { + "weight": { + "dtype": "int4", + "group_size": 128, + "symmetric": symmetric, + } + }, + "exclude": exclude or [], + } + + +def _sign_extend_int4_nibbles(t: torch.Tensor) -> torch.Tensor: + mask = (t & 0x8).bool() + t = t.clone() + t[mask] = t[mask] | 0xF0 + return t + + +def _dequantize_quark_signed_awq_torch( + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + group_size: int, + *, + pack_reorder: bool = True, +) -> torch.Tensor: + bits = 4 + shifts = torch.arange(0, 32, bits, device=qweight.device) + iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8) + iweights = iweights.view(qweight.shape[0], -1) + zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8) + zeros = zeros.view(qzeros.shape[0], -1) + + if pack_reorder: + order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device) + else: + order = torch.arange(8, device=qweight.device) + iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape( + qweight.shape[0], -1 + ) + zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1) + iweights = _sign_extend_int4_nibbles(iweights & 0xF) + zeros = _sign_extend_int4_nibbles(zeros & 0xF) + + scales = scales.repeat_interleave(group_size, dim=0) + zeros = zeros.repeat_interleave(group_size, dim=0) + return (iweights - zeros) * scales + + +def _dequantize_awq_unsigned_torch( + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + group_size: int, + *, + pack_reorder: bool = True, +) -> torch.Tensor: + bits = 4 + shifts = torch.arange(0, 32, bits, device=qweight.device) + iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8) + iweights = iweights.view(qweight.shape[0], -1) + zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8) + zeros = zeros.view(qzeros.shape[0], -1) + + if pack_reorder: + order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device) + else: + order = torch.arange(8, device=qweight.device) + iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape( + qweight.shape[0], -1 + ) + zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1) + + scales = scales.repeat_interleave(group_size, dim=0) + zeros = zeros.repeat_interleave(group_size, dim=0) + return (iweights - zeros) * scales + + +def _pack_int4_nibbles(nibbles: torch.Tensor, *, pack_reorder: bool) -> torch.Tensor: + pack_order = ( + torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=nibbles.device) + if pack_reorder + else torch.arange(8, device=nibbles.device) + ) + shifts = pack_order * 4 + return ((nibbles.to(torch.int64) & 0xF) << shifts).sum(dim=-1).to(torch.int32) + + +class TestQuarkInt4Format: + """Tests for Quark INT4 export format compatibility.""" + + def test_quark_order_pack_method_uses_native_int4_scheme(self): + quant_config = QuarkConfig.from_config(_quark_int4_config(pack_method="order")) + weight_config = quant_config.quant_config["global_quant_config"] + weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config( + weight_config + ) + assert scheme_cls is QuarkW4A16Int4 + + scheme = quant_config.init_scheme( + scheme_cls, + weight_quant_key=weight_key, + activation_quant_key=act_key, + weight_config=weight_config["weight"], + ) + assert isinstance(scheme, QuarkW4A16Int4) + assert not scheme.pack_reorder + + def test_quark_int4_scheme_supports_asymmetric_weights(self): + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False)) + weight_config = quant_config.quant_config["global_quant_config"] + weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config( + weight_config + ) + assert scheme_cls is QuarkW4A16Int4 + + scheme = quant_config.init_scheme( + scheme_cls, + weight_quant_key=weight_key, + activation_quant_key=act_key, + weight_config=weight_config["weight"], + ) + assert isinstance(scheme, QuarkW4A16Int4) + assert not scheme.is_symmetric + + @pytest.mark.parametrize("missing_field", ["group_size", "symmetric"]) + def test_quark_int4_scheme_requires_weight_config_fields(self, missing_field): + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt4Static, + ) + + weight_config = { + "dtype": "int4", + "group_size": 128, + "symmetric": True, + } + weight_config.pop(missing_field) + quant_config = QuarkConfig.from_config( + { + "quant_method": "quark", + "export": {"pack_method": "reorder", "kv_cache_group": []}, + "global_quant_config": {"weight": weight_config}, + "exclude": [], + } + ) + + with pytest.raises(ValueError, match=missing_field): + quant_config.init_scheme( + QuarkW4A16Int4, + weight_quant_key=kInt4Static, + activation_quant_key=None, + weight_config=weight_config, + ) + + def test_quark_int4_moe_uses_native_moe_method(self): + from unittest.mock import MagicMock, patch + + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt4StaticAsym, + ) + + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False)) + moe_config = type("MoeConfig", (), {})() + moe_config.has_bias = False + + mock_backend = MagicMock() + mock_experts_cls = MagicMock() + with patch( + "vllm.model_executor.layers.quantization.quark.quark_moe" + ".select_wna16_moe_backend", + return_value=(mock_backend, mock_experts_cls), + ): + method = QuarkW4A16Int4MoEMethod( + kInt4StaticAsym, + None, + quant_config.quant_config["global_quant_config"]["weight"], + quant_config.pack_method, + moe_config, + ) + + assert method.group_size == 128 + assert method.pack_reorder + assert method.use_wna16_backend + + def test_triton_wna16_experts_supports_asymmetric_int4(self): + from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( + TritonWNA16Experts, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, + ) + + for key in [kInt4Static, kInt4Static32, kInt4StaticAsym, kInt4Static32Asym]: + assert TritonWNA16Experts._supports_quant_scheme(key, None), ( + f"{key} should be supported by TritonWNA16Experts" + ) + + def test_quark_moe_loader_shards_w2_zero_point_across_tp_ranks(self): + """w2 zero-points are sharded along the intermediate dim per TP rank. + + Loading the same Quark checkpoint tensor at tensor_parallel_size=2 for + both ranks must jointly reconstruct the tensor_parallel_size=1 (full) + result, with no overlap or gap. This guards the custom loader's TP + split for the packed zero-point layout. + """ + + class _FakeMoEConfig: + has_bias = False + + def __init__(self, tp_size, tp_rank): + self.tp_size = tp_size + self.tp_rank = tp_rank + + class _FakeLayer: + def __init__(self, moe_config): + self.moe_config = moe_config + self.intermediate_size_per_partition = 64 + self.group_size_div_factor = 1 + + def load_w2_zero_point(raw_zp, tp_size, tp_rank): + from unittest.mock import MagicMock, patch + + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt4StaticAsym, + ) + + quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False)) + moe_config = _FakeMoEConfig(tp_size, tp_rank) + with patch( + "vllm.model_executor.layers.quantization.quark.quark_moe" + ".select_wna16_moe_backend", + return_value=(MagicMock(), MagicMock()), + ): + method = QuarkW4A16Int4MoEMethod( + kInt4StaticAsym, + None, + quant_config.quant_config["global_quant_config"]["weight"], + quant_config.pack_method, + moe_config, + ) + + layer = _FakeLayer(moe_config) + # The loader converts the packed zero-point, then writes the + # per-rank slice into param.data[expert_id]. A single expert dim is + # enough; the last dim shrinks by tp_size after sharding. + param = torch.nn.Parameter( + torch.zeros(1, 16, 8 // tp_size, dtype=torch.uint8), + requires_grad=False, + ) + loader = method.get_weight_loader(layer, weight_loader=None) + loader( + param, + raw_zp.clone(), + weight_name="w2_weight_zero_point", + shard_id="w2", + expert_id=0, + ) + return param.data[0] + + # Quark-packed zero-point for one expert: rows = hidden // pack_factor, + # cols carry the (packed) intermediate dim, split evenly across ranks. + raw_zp = torch.randint(0, 256, (8, 16), dtype=torch.uint8) + + full = load_w2_zero_point(raw_zp, tp_size=1, tp_rank=0) + shard0 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=0) + shard1 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=1) + + expected = full.view(full.size(0), 2, -1) + torch.testing.assert_close(shard0, expected[:, 0]) + torch.testing.assert_close(shard1, expected[:, 1]) + + def test_quark_shared_expert_gate_keeps_quantized_tensors(self): + quant_config = QuarkConfig.from_config(_quark_int4_config()) + mapper = quant_config.get_cache_scale_mapper() + + output_names = { + name + for name, _ in mapper.apply( + [ + ( + "model.language_model.layers.0.mlp.shared_expert_gate.weight", + torch.zeros(1), + ), + ( + "model.language_model.layers.0.mlp.shared_expert_gate" + ".weight_scale", + torch.zeros(1), + ), + ( + "model.language_model.layers.0.mlp.shared_expert_gate" + ".weight_zero_point", + torch.zeros(1), + ), + ] + ) + } + + assert ( + "model.language_model.layers.0.mlp.shared_expert_gate.weight" + in output_names + ) + assert ( + "model.language_model.layers.0.mlp.shared_expert_gate.weight_scale" + in output_names + ) + assert ( + "model.language_model.layers.0.mlp.shared_expert_gate.weight_zero_point" + in output_names + ) + + def test_quark_apply_mapper_updates_exclude_and_layer_quant_config(self): + quant_config = QuarkConfig.from_config( + { + **_quark_int4_config(), + "exclude": ["lm_head"], + "layer_quant_config": { + "model.language_model.layers.0.mlp.gate": { + "weight": {"dtype": "float16"}, + }, + }, + } + ) + quant_config.apply_vllm_mapper( + WeightsMapper( + orig_to_new_prefix={ + "lm_head": "language_model.lm_head", + "model.language_model.": "language_model.model.", + }, + ) + ) + + layer_quant_config = quant_config.quant_config["layer_quant_config"] + assert "language_model.model.layers.0.mlp.gate" in layer_quant_config + assert "model.language_model.layers.0.mlp.gate" not in layer_quant_config + assert quant_config.quant_config["exclude"] == [ + "language_model.lm_head", + ] + + def test_quark_apply_mapper_ignores_non_string_list_entries(self): + quant_config = QuarkConfig.from_config( + { + **_quark_int4_config(), + "algo_config": [ + {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]} + ], + } + ) + quant_config.apply_vllm_mapper(WeightsMapper()) + + assert quant_config.quant_config["algo_config"] == [ + {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]} + ] + + def test_quark_exclude_matches_via_mapper_prefix(self): + quant_config = QuarkConfig.from_config(_quark_int4_config(exclude=["lm_head"])) + quant_config.apply_vllm_mapper( + WeightsMapper(orig_to_new_prefix={"lm_head": "language_model.lm_head"}) + ) + + exclude_layers = quant_config.quant_config["exclude"] + assert should_ignore_layer("language_model.lm_head", ignore=exclude_layers) + assert not should_ignore_layer( + "language_model.model.layers.0.mlp.gate", ignore=exclude_layers + ) + + +@pytest.mark.parametrize("symmetric", [False, True]) +@pytest.mark.parametrize("pack_method", ["order", "reorder"]) +def test_quark_int4_canonicalizes_pack_for_kernel_layout(pack_method, symmetric): + """Quark pack order is normalized to the layout expected by awq_* ops. + + This reuses the existing AWQ dequant/gemm kernels for compute only; loading + still goes through the native Quark quantization path, not AutoAWQ. + """ + pack_reorder = pack_method == "reorder" + group_size = 2 + packed_values = torch.tensor( + [ + [0, 1, 7, 8, 9, 15, 2, 14], + [15, 8, 0, 3, 12, 7, 1, 9], + ], + dtype=torch.int32, + ) + packed_zeros = torch.zeros((1, 8), dtype=torch.int32) + qweight = _pack_int4_nibbles(packed_values, pack_reorder=pack_reorder).view(2, 1) + qzeros = _pack_int4_nibbles(packed_zeros, pack_reorder=pack_reorder).view(1, 1) + scales = torch.ones((1, 8), dtype=torch.float16) + + dequantize = ( + _dequantize_quark_signed_awq_torch + if symmetric + else _dequantize_awq_unsigned_torch + ) + expected = dequantize( + qweight, + scales, + qzeros, + group_size, + pack_reorder=pack_reorder, + ) + canonical_weight = canonicalize_quark_packed_int4( + qweight, + pack_reorder=pack_reorder, + is_symmetric=symmetric, + ) + canonical_zero = canonicalize_quark_packed_int4( + qzeros, + pack_reorder=pack_reorder, + is_symmetric=symmetric, + ) + actual = _dequantize_awq_unsigned_torch( + canonical_weight, scales, canonical_zero, group_size + ) + + assert torch.equal(actual, expected) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 3c9c616d0c16..8c57059e5727 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -864,6 +864,7 @@ def choose_mp_linear_kernel( can_implement, failure_reason = kernel.can_implement(config) if can_implement: + logger.info_once("Using %s for mixed-precision linear", kernel.__name__) return kernel else: failure_reasons.append( diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 9f982644b569..da9666a66b8f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -596,9 +596,9 @@ def _supports_quant_scheme( ) -> bool: SUPPORTED_W = [ kInt4Static, + kInt4StaticAsym, kInt8Static, kInt4Static32, - kInt4StaticAsym, kInt4Static32Asym, # other group sizes? ] diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 2c123b8768d7..71b858e91969 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -93,7 +93,7 @@ def _replace_or_register_parameter( def _convert_awq_to_standard_format( layer: torch.nn.Module, w_q_name: str, - w_zp_name: str, + w_zp_name: str | None, size_bits: int, ) -> None: """Convert AWQ weight and zero-point tensors to standard GPTQ-like format. @@ -101,6 +101,7 @@ def _convert_awq_to_standard_format( AWQ packs qweight along the output dim with a non-standard bit order. This converts to standard bit order and repacks qweight along the input dim, matching the format expected by the MPLinearKernel framework. + If w_zp_name is None (symmetric quantization), only the weight is converted. """ pack_factor = 32 // size_bits mask = (1 << size_bits) - 1 @@ -139,6 +140,9 @@ def _noop_loader(*args, **kwargs): ) setattr(layer, w_q_name, new_param) + if w_zp_name is None: + return + # --- Convert qzeros: fix AWQ bit ordering and repack # AWQ qzeros: (G, N // pack) packed along dim 1, AWQ bit order # Target: (N // pack, G) packed along dim 0, standard bit order diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 88544fbe994d..a3a3a36f0047 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -32,6 +32,7 @@ QuarkOCP_MX, QuarkScheme, QuarkW4A8_MXFP4_FP8, + QuarkW4A16Int4, QuarkW8A8Fp8, QuarkW8A8Fp8PerBlock, QuarkW8A8Int8, @@ -39,6 +40,7 @@ from vllm.model_executor.layers.quantization.quark.utils import ( QuarkQTensorHint, deep_compare, + parse_w4a16_int4_weight_config, should_ignore_layer, ) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( @@ -54,6 +56,10 @@ kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt4W4A8StaticChannelSym, kInt8DynamicTensorAsym, kInt8DynamicTensorSym, @@ -172,8 +178,10 @@ def apply_vllm_mapper( # noqa: B027 quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {} for k, v in self.quant_config.items(): - if isinstance(v, list): + if isinstance(v, list) and all(isinstance(item, str) for item in v): quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_list(v) + elif isinstance(v, list): + quant_config_with_hf_to_vllm_mapper[k] = v elif isinstance(v, dict): quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_dict(v) else: @@ -215,6 +223,9 @@ def get_quant_method( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key, dynamic_mxfp4_quant=dynamic_mxfp4_quant, + weight_config=self._find_matched_config(prefix, type(layer)).get( + "weight" + ), ) layer.scheme = scheme return QuarkLinearMethod(self) @@ -227,6 +238,7 @@ def get_quant_method( method_cls=method_cls, weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key, + layer_name=prefix, ) return None @@ -572,6 +584,26 @@ def _is_w8a8_int8( ) return QuantKeyMatch(True, activation_quant_key, weight_quant_key) + def _is_w4a16_int4( + self, + weight_quant: dict[str, Any] | None, + input_quant: dict[str, Any] | None, + ) -> QuantKeyMatch: + if weight_quant is None or input_quant is not None: + return QuantKeyMatch(False, None, None) + + is_int4 = weight_quant.get("dtype") in ("int4", "uint4") + is_packed = self.pack_method in ("order", "reorder") + if not (is_int4 and is_packed): + return QuantKeyMatch(False, None, None) + + group_size, is_symmetric = parse_w4a16_int4_weight_config(weight_quant) + if is_symmetric: + weight_key = kInt4Static32 if group_size == 32 else kInt4Static + else: + weight_key = kInt4Static32Asym if group_size == 32 else kInt4StaticAsym + return QuantKeyMatch(True, None, weight_key) + def _is_w4a8_mxfp4_fp8( self, weight_quant: dict[str, Any] | None, @@ -890,6 +922,12 @@ def _get_scheme_cls_from_config( match.activation_quant_key, QuarkW8A8Int8, ) + elif match := self._is_w4a16_int4(weight_config, input_config): + return ( + match.weight_quant_key, + match.activation_quant_key, + QuarkW4A16Int4, + ) elif match := self._is_w4a8_mxfp4_fp8(weight_config, input_config): return ( match.weight_quant_key, @@ -919,6 +957,7 @@ def init_scheme( weight_quant_key: QuantKey | None, activation_quant_key: QuantKey | None, dynamic_mxfp4_quant: bool = False, + weight_config: dict[str, Any] | None = None, ) -> "QuarkScheme": """Construct a Quark scheme selected by get_scheme_cls.""" if scheme_cls not in ( @@ -928,9 +967,24 @@ def init_scheme( QuarkOCP_MX, QuarkNVFP4, QuarkW4A8_MXFP4_FP8, + QuarkW4A16Int4, ): raise AssertionError(f"Unsupported Quark scheme class: {scheme_cls}") + scheme: QuarkScheme + if scheme_cls is QuarkW4A16Int4: + normalized_weight_config = self._unwrap_single_quant_config(weight_config) + assert normalized_weight_config is not None + assert weight_quant_key is not None + scheme = QuarkW4A16Int4( + weight_quant_key=weight_quant_key, + activation_quant_key=activation_quant_key, + pack_method=self.pack_method, + weight_config=normalized_weight_config, + ) + self._check_scheme_supported(scheme.get_min_capability()) + return scheme + kwargs: dict[str, Any] = {"activation_quant_key": activation_quant_key} # These classes handle different possible `weight_quant_key`. diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index d2472524a5fe..3a9f6d63662b 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -40,6 +40,12 @@ make_int8_moe_quant_config, select_int8_moe_backend, ) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + convert_to_wna16_moe_kernel_format, + make_wna16_moe_kernel, + make_wna16_moe_quant_config, + select_wna16_moe_backend, +) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, @@ -56,7 +62,10 @@ make_nvfp4_moe_quant_config, select_nvfp4_moe_backend, ) -from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint +from vllm.model_executor.layers.quantization.quark.utils import ( + QuarkQTensorHint, + parse_w4a16_int4_weight_config, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( validate_fp8_block_shape_moe, ) @@ -75,6 +84,10 @@ kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt4W4A8StaticChannelSym, kInt8DynamicTensorAsym, kInt8DynamicTensorSym, @@ -102,6 +115,7 @@ __all__ = [ "QuarkMoEMethod", + "QuarkW4A16Int4MoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkOCP_MX_MoEMethod", "QuarkNvfp4MoEMethod", @@ -179,6 +193,12 @@ def get_moe_method_target( match.activation_quant_key, QuarkW8A8Fp8MoEMethod, ) + if match := quant_config._is_w4a16_int4(weight_config, input_config): + return ( + match.weight_quant_key, + match.activation_quant_key, + QuarkW4A16Int4MoEMethod, + ) if match := quant_config._is_w_ocp_mx_a_x( weight_config, input_config, allow_static_fp8=True ): @@ -202,6 +222,7 @@ def get_moe_method( method_cls: type["QuarkMoEMethod"], weight_quant_key: QuantKey | None, activation_quant_key: QuantKey | None, + layer_name: str | None = None, ) -> "QuarkMoEMethod": if weight_quant_key is None: raise RuntimeError("Unsupported FusedMoe scheme") @@ -226,10 +247,339 @@ def get_moe_method( return QuarkW8A8Int8MoEMethod( module.moe_config, weight_quant_key, activation_quant_key ) + elif method_cls is QuarkW4A16Int4MoEMethod: + assert layer_name is not None + layer_quant_config = quant_config._find_matched_config( + layer_name, type(module) + ) + weight_config = quant_config._unwrap_single_quant_config( + layer_quant_config.get("weight") + ) + if weight_config is None: + raise RuntimeError( + f"Missing weight quantization config for layer {layer_name}" + ) + return QuarkW4A16Int4MoEMethod( + weight_quant_key, + activation_quant_key, + weight_config, + quant_config.pack_method, + module.moe_config, + quant_config=quant_config, + ) else: raise RuntimeError("Unsupported FusedMoe scheme") +class QuarkW4A16Int4MoEMethod(QuarkMoEMethod): + """Quark packed INT4 weight-only MoE method.""" + + supported_activation_quant_keys: list[QuantKey | None] = [None] + supported_weight_quant_keys: list[QuantKey] = [ + kInt4Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, + ] + + def __init__( + self, + weight_quant_key: QuantKey, + activation_quant_key: QuantKey | None, + weight_config: dict[str, Any], + pack_method: str, + moe: FusedMoEConfig, + quant_config: "QuarkConfig | None" = None, # type: ignore # noqa F821 + ): + super().__init__(moe, weight_quant_key, activation_quant_key) + + self.quant_config = quant_config + # group_size is read from the Quark config rather than weight_quant_key: + # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4 + # checkpoints commonly use group_size 128, which the QuantKey cannot + # represent. Symmetry is unambiguous, so it comes from the key. + self.group_size = parse_w4a16_int4_weight_config(weight_config)[0] + self.is_symmetric = weight_quant_key.symmetric + self.bit8_pack_factor = 2 + self.pack_reorder = pack_method == "reorder" + + self.use_wna16_backend = True + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + weight_key=weight_quant_key, + quant_config=self.quant_config, + may_have_zp=not self.is_symmetric, + may_have_bias=False, + ) + + def create_weights( + self, + layer: RoutedExperts, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + group_size = self.group_size + group_size_div_factor = 1 + + while intermediate_size_per_partition % group_size or hidden_size % group_size: + group_size = group_size // 2 + group_size_div_factor *= 2 + assert group_size >= 32 + layer.group_size = group_size + layer.group_size_div_factor = group_size_div_factor + + strategy = FusedMoeWeightScaleSupported.GROUP.value + extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False}) + + assert "weight_loader" in extra_weight_attrs + weight_loader = extra_weight_attrs["weight_loader"] + extra_weight_attrs["weight_loader"] = self.get_weight_loader( + layer, weight_loader + ) + + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w13_weight_zero_point = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition // self.bit8_pack_factor, + hidden_size // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point) + set_weight_attrs(w13_weight_zero_point, extra_weight_attrs) + + w2_weight_zero_point = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size // self.bit8_pack_factor, + intermediate_size_per_partition // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point) + set_weight_attrs(w2_weight_zero_point, extra_weight_attrs) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + w13_qzeros = layer.w13_weight_zero_point if not self.is_symmetric else None + w2_qzeros = layer.w2_weight_zero_point if not self.is_symmetric else None + + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=None, + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_qzeros=w13_qzeros, + w2_qzeros=w2_qzeros, + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + + ( + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, + _, # w13_qzeros + _, # w2_qzeros + _, # w13_input_global_scale + _, # w2_input_global_scale + _, # w13_bias + _, # w2_bias + ) = converted + + replace_parameter(layer, "w13_weight", w13_qweight) + replace_parameter(layer, "w2_weight", w2_qweight) + replace_parameter(layer, "w13_weight_scale", w13_scales) + replace_parameter(layer, "w2_weight_scale", w2_scales) + + self._setup_kernel(layer) + + def _setup_kernel(self, layer: RoutedExperts) -> None: + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + backend=self.wna16_backend, + routing_tables=layer._expert_routing_tables(), + ) + + def get_fused_moe_quant_config( + self, layer: RoutedExperts + ) -> FusedMoEQuantConfig | None: + w1_zp = layer.w13_weight_zero_point if not self.is_symmetric else None + w2_zp = layer.w2_weight_zero_point if not self.is_symmetric else None + return make_wna16_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + group_size=layer.group_size, + num_bits=4, + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + def get_weight_loader(self, layer: RoutedExperts, weight_loader): + def convert_quark_tensor(tensor: torch.Tensor, tensor_type: str): + size0 = tensor.size(0) + tensor = tensor.view(torch.uint8) + + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) + tensor = (tensor[:, :, None] >> shifter) & 0xF + + if self.pack_reorder: + reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7] + tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order] + else: + tensor = tensor.view(-1, 8) + if self.is_symmetric: + tensor = tensor ^ 0x8 + tensor = tensor.view(size0, -1) + + tensor = tensor.T.contiguous() + if tensor_type == "weight": + tensor = tensor[:, 1::2] * 16 + tensor[:, ::2] + elif tensor_type == "zero_point": + tensor = tensor[1::2, :] * 16 + tensor[::2, :] + return tensor + + def quark_int4_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ): + loaded_weight = loaded_weight.to(param.data.device) + shard_size = layer.intermediate_size_per_partition + + if weight_name.endswith("_weight"): + loaded_weight = convert_quark_tensor(loaded_weight, "weight") + elif weight_name.endswith("_weight_zero_point"): + loaded_weight = convert_quark_tensor(loaded_weight, "zero_point") + elif weight_name.endswith("_weight_scale"): + loaded_weight = loaded_weight.T + + if layer.group_size_div_factor > 1 and ( + "weight_zero_point" in weight_name or "weight_scale" in weight_name + ): + loaded_weight = loaded_weight.repeat_interleave( + layer.group_size_div_factor, 1 + ) + + if "w13_weight_zero_point" in weight_name: + tensor = loaded_weight.view( + layer.moe_config.tp_size, -1, loaded_weight.size(1) + )[layer.moe_config.tp_rank] + if shard_id == "w1": + param.data[expert_id, : shard_size // 2] = tensor + else: + param.data[expert_id, shard_size // 2 :] = tensor + return True if return_success else None + elif "w2_weight_zero_point" in weight_name: + param.data[expert_id] = loaded_weight.view( + loaded_weight.size(0), layer.moe_config.tp_size, -1 + )[:, layer.moe_config.tp_rank] + return True if return_success else None + + return weight_loader( + param, + loaded_weight, + weight_name, + shard_id, + expert_id, + return_success=return_success, + ) + + return quark_int4_weight_loader + + class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): supported_activation_quant_keys = [ kFp8StaticTensorSym, diff --git a/vllm/model_executor/layers/quantization/quark/schemes/__init__.py b/vllm/model_executor/layers/quantization/quark/schemes/__init__.py index 4db1ba384555..a5857a776f5e 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/__init__.py @@ -5,6 +5,7 @@ from .quark_ocp_mx import QuarkOCP_MX from .quark_scheme import QuarkScheme from .quark_w4a8_mxfp4_fp8 import QuarkW4A8_MXFP4_FP8 +from .quark_w4a16_int4 import QuarkW4A16Int4 from .quark_w8a8_fp8 import QuarkW8A8Fp8, QuarkW8A8Fp8PerBlock from .quark_w8a8_int8 import QuarkW8A8Int8 @@ -14,6 +15,7 @@ "QuarkW8A8Fp8PerBlock", "QuarkW8A8Int8", "QuarkOCP_MX", + "QuarkW4A16Int4", "QuarkW4A8_MXFP4_FP8", "QuarkNVFP4", ] diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py new file mode 100644 index 000000000000..cded0f80a9eb --- /dev/null +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Callable +from typing import Any + +import torch + +from vllm.model_executor.kernels.linear import ( + MPLinearKernel, + MPLinearLayerConfig, + choose_mp_linear_kernel, +) +from vllm.model_executor.layers.quantization.auto_awq import ( + _convert_awq_to_standard_format, +) +from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import ( + QuarkScheme, +) +from vllm.model_executor.layers.quantization.quark.utils import ( + canonicalize_quark_packed_int4, + parse_w4a16_int4_weight_config, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.scalar_type import scalar_types + + +class QuarkW4A16Int4(QuarkScheme): + """Quark packed INT4 weight-only linear scheme via MPLinearKernel.""" + + supported_activation_quant_keys: list[QuantKey | None] = [None] + supported_weight_quant_keys: list[QuantKey] = [ + kInt4Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, + ] + + def __init__( + self, + weight_quant_key: QuantKey, + activation_quant_key: QuantKey | None, + pack_method: str, + weight_config: dict[str, Any], + ): + super().__init__(weight_quant_key, activation_quant_key) + + # group_size is read from the Quark config rather than weight_quant_key: + # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4 + # checkpoints commonly use group_size 128, which the QuantKey cannot + # represent. Symmetry is unambiguous, so it comes from the key. + self.group_size = parse_w4a16_int4_weight_config(weight_config)[0] + self.pack_factor = 8 + self.pack_reorder = pack_method == "reorder" + self.is_symmetric = weight_quant_key.symmetric + self.quant_type = ( + scalar_types.uint4b8 if self.is_symmetric else scalar_types.uint4 + ) + self.kernel: MPLinearKernel | None = None + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def create_weights( + self, + layer: torch.nn.Module, + output_partition_sizes: list[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + input_size = kwargs["input_size"] + output_size = kwargs["output_size"] + group_size = ( + self.group_size if self.group_size != -1 else input_size_per_partition + ) + if input_size_per_partition % group_size != 0: + raise ValueError( + "The input size is not aligned with the quantized weight shape. " + "This can be caused by too large tensor parallel size. " + f"input_size_per_partition={input_size_per_partition}, " + f"group_size={group_size}." + ) + + output_size_per_partition = sum(output_partition_sizes) + packed_output_size_per_partition = math.ceil( + output_size_per_partition / self.pack_factor + ) + layer.output_size_per_partition = output_size_per_partition + layer.packed_output_size_per_partition = packed_output_size_per_partition + + mp_linear_kernel_config = MPLinearLayerConfig( + full_weight_shape=(input_size, output_size), + partition_weight_shape=( + input_size_per_partition, + output_size_per_partition, + ), + weight_type=self.quant_type, + act_type=params_dtype, + group_size=group_size, + zero_points=not self.is_symmetric, + ) + kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) + + def weight_scale_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + *args, + **loader_kwargs, + ) -> None: + if loaded_weight.shape[1] < param.data.shape[1]: + padded_weight = loaded_weight.new_zeros(param.data.shape) + padded_weight[:, : loaded_weight.shape[1]] = loaded_weight + loaded_weight = padded_weight + weight_loader(param, loaded_weight, *args, **loader_kwargs) + + weight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + packed_output_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.pack_factor, + weight_loader=weight_loader, + ) + num_groups = input_size_per_partition // group_size + weight_zero_point = PackedvLLMParameter( + data=torch.zeros( + num_groups, + packed_output_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.pack_factor, + weight_loader=weight_loader, + ) + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + num_groups, + packed_output_size_per_partition * self.pack_factor, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_scale_loader, + ) + + layer.register_parameter("weight", weight) + layer.register_parameter("weight_zero_point", weight_zero_point) + layer.register_parameter("weight_scale", weight_scale) + + self.kernel = kernel_type( + mp_linear_kernel_config, + w_q_param_name="weight", + w_s_param_name="weight_scale", + w_zp_param_name="weight_zero_point" if not self.is_symmetric else None, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight.data = canonicalize_quark_packed_int4( + layer.weight.data, + pack_reorder=self.pack_reorder, + is_symmetric=self.is_symmetric, + pack_factor=self.pack_factor, + ) + if not self.is_symmetric: + layer.weight_zero_point.data = canonicalize_quark_packed_int4( + layer.weight_zero_point.data, + pack_reorder=self.pack_reorder, + is_symmetric=self.is_symmetric, + pack_factor=self.pack_factor, + ) + output_size = layer.output_size_per_partition + packed_output_size = layer.packed_output_size_per_partition * self.pack_factor + if output_size < packed_output_size: + layer.weight_scale.data[:, output_size:].zero_() + + _convert_awq_to_standard_format( + layer, + "weight", + "weight_zero_point" if not self.is_symmetric else None, + self.quant_type.size_bits, + ) + assert self.kernel is not None + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ): + assert self.kernel is not None + return self.kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/quark/utils.py b/vllm/model_executor/layers/quantization/quark/utils.py index 48c978263ee7..ac455e38b83f 100644 --- a/vllm/model_executor/layers/quantization/quark/utils.py +++ b/vllm/model_executor/layers/quantization/quark/utils.py @@ -66,6 +66,63 @@ def should_ignore_layer( return all(shards_ignored) +def parse_w4a16_int4_weight_config( + weight_config: Mapping[str, Any], +) -> tuple[int, bool]: + """Parse required W4A16 INT4/UINT4 weight fields from Quark config.""" + if "group_size" not in weight_config: + raise ValueError( + "Quark W4A16 INT4/UINT4 configs must specify weight.group_size" + ) + if "symmetric" not in weight_config: + raise ValueError("Quark W4A16 INT4/UINT4 configs must specify weight.symmetric") + + group_size = weight_config["group_size"] + is_symmetric = weight_config["symmetric"] + if not isinstance(group_size, int) or group_size <= 0: + raise ValueError( + f"Quark W4A16 weight.group_size must be a positive int, got {group_size!r}" + ) + if not isinstance(is_symmetric, bool): + raise ValueError( + f"Quark W4A16 weight.symmetric must be a bool, got {is_symmetric!r}" + ) + return group_size, is_symmetric + + +def canonicalize_quark_packed_int4( + packed_weight: torch.Tensor, + *, + pack_reorder: bool, + is_symmetric: bool, + pack_factor: int = 8, +) -> torch.Tensor: + """Convert Quark export nibble layout to AWQ checkpoint layout.""" + from vllm.model_executor.layers.quantization.auto_awq import ( + _REVERSE_AWQ_PACK_ORDER, + ) + + if pack_reorder: + source_order = torch.tensor( + _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32 + ) + else: + source_order = torch.arange( + pack_factor, device=packed_weight.device, dtype=torch.int32 + ) + target_order = torch.tensor( + _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32 + ) + source_shifts = source_order * 4 + target_shifts = target_order * 4 + + values = (packed_weight.to(torch.int32)[..., None] >> source_shifts) & 0xF + if is_symmetric: + values = values ^ 0x8 + packed = (values.to(torch.int64) << target_shifts.to(torch.int64)).sum(dim=-1) + return packed.to(torch.int32) + + # utility for tensor dims > 2 cases def quark_quantize_weight_to_mxfp4(w: torch.Tensor): assert w.dtype == torch.bfloat16, ( From b34647906550f6e9dc65c55ecfb36b697d5e1f43 Mon Sep 17 00:00:00 2001 From: Yuwen Zhou Date: Fri, 18 Sep 2026 17:59:42 +0800 Subject: [PATCH 464/774] [CPU] Add CPU FP8 W8A8 linear/MoE support (#49942) Signed-off-by: yuwenzho Co-authored-by: Li, Jiang --- cmake/cpu_extension.cmake | 5 + csrc/cpu/cpu_attn.cpp | 13 +- csrc/cpu/cpu_isa.cpp | 20 + csrc/cpu/sgl-kernels/common.h | 83 ++ csrc/cpu/sgl-kernels/gemm.h | 54 +- csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp | 1093 +++++++++++++++++ csrc/cpu/sgl-kernels/moe.cpp | 100 +- csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp | 432 +++++++ csrc/cpu/torch_bindings.cpp | 37 +- tests/kernels/moe/test_cpu_quant_fused_moe.py | 198 +++ .../quantization/test_cpu_fp8_scaled_mm.py | 229 +++- vllm/_custom_ops.py | 44 + .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/cpu.py | 310 ++++- .../layers/attention/mla_attention.py | 4 +- .../layers/fused_moe/experts/cpu_moe.py | 217 ++++ .../layers/fused_moe/oracle/fp8.py | 32 +- .../layers/quantization/online/fp8.py | 6 +- vllm/v1/attention/backends/mla/amx_mla.py | 87 +- 20 files changed, 2924 insertions(+), 46 deletions(-) create mode 100644 csrc/cpu/cpu_isa.cpp create mode 100644 csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp create mode 100644 csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 6359a96887a5..cca2a98cee57 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -488,6 +488,7 @@ set(VLLM_EXT_SRC "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/torch_bindings.cpp") if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND @@ -546,11 +547,13 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/gemm.cpp" "csrc/cpu/sgl-kernels/gemm_int8.cpp" "csrc/cpu/sgl-kernels/gemm_fp8.cpp" + "csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp" "csrc/cpu/sgl-kernels/gemm_int4.cpp" "csrc/cpu/sgl-kernels/moe.cpp" "csrc/cpu/sgl-kernels/moe_int8.cpp" "csrc/cpu/sgl-kernels/moe_int4.cpp" "csrc/cpu/sgl-kernels/moe_fp8.cpp" + "csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp" "csrc/cpu/sgl-kernels/bmm.cpp" "csrc/cpu/sgl-kernels/decode.cpp" "csrc/cpu/sgl-kernels/extend.cpp" @@ -571,6 +574,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/torch_bindings.cpp" @@ -587,6 +591,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/cpu_isa.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/torch_bindings.cpp" diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index b31a6faa0104..ea0b4e49b273 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -1,16 +1,7 @@ #include "cpu_attn_dispatch_generated.h" -#if defined(__x86_64__) || defined(_M_X64) - #include -// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21]. -static bool runtime_has_amx_fp8() { - unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; - if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) { - return (eax >> 21) & 1u; - } - return false; -} -#endif // x86_64 +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp. +extern bool runtime_has_amx_fp8(); // Maps kv_cache_dtype string to Fp8KVCacheDataType enum. // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2. diff --git a/csrc/cpu/cpu_isa.cpp b/csrc/cpu/cpu_isa.cpp new file mode 100644 index 000000000000..1c4a8faf1918 --- /dev/null +++ b/csrc/cpu/cpu_isa.cpp @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// AMX-FP8 is reported at CPUID leaf 0x1E, subleaf 1, EAX[4]. +#if defined(__x86_64__) || defined(_M_X64) + #include + +bool runtime_has_amx_fp8() { + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; + if (__get_cpuid_max(0, nullptr) < 0x1e) return false; + if (__get_cpuid_count(0x1e, 1, &eax, &ebx, &ecx, &edx)) { + return (eax >> 4) & 1u; + } + return false; +} + +bool cpu_has_amx_fp8() { return runtime_has_amx_fp8(); } +#else +bool cpu_has_amx_fp8() { return false; } +#endif diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index 6f27ba2d5c10..de1260bb3321 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -12,11 +12,45 @@ #include #include #include +#include #if defined(_OPENMP) #include #endif +// Adapted from sglang: dispatch utilities for FP8 W8A8 kernels +#define AT_DISPATCH_OUT_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, \ + NAME, \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Float, out_t, __VA_ARGS__) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::BFloat16, out_t, __VA_ARGS__) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Half, out_t, __VA_ARGS__)) + +#define AT_DISPATCH_CASE_ENUM_NO_RETURN(VALUE, TYPE, HINT, ...) \ + case VALUE: { \ + constexpr TYPE HINT = VALUE; \ + __VA_ARGS__; \ + break; \ + } + +#define AT_DISPATCH_BOOL_NO_RETURN(VALUE, NAME, HINT, ...) \ + [&]() { \ + switch (VALUE) { \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(true, bool, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(false, bool, HINT, __VA_ARGS__) \ + } \ + }() + +#define AT_DISPATCH_QUANT_MODE_NO_RETURN(MODE, NAME, HINT, ...) \ + [&]() { \ + switch (MODE) { \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_TENSOR, int, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_ROW, int, HINT, __VA_ARGS__) \ + AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_GROUP, int, HINT, __VA_ARGS__) \ + } \ + }() + namespace { // dispatch bool @@ -472,4 +506,53 @@ inline T* conditional_data_ptr(const std::optional& opt) { return opt.has_value() ? opt.value().data_ptr() : nullptr; } +// Adapted from sglang: M blocking utilities for FP8 W8A8 kernels +inline int64_t get_m_block(int64_t M) { + if (M <= 48) { + return M; + } else if (M < 64) { + return 32; + } else if (M < 96) { + return 64; + } else { + return 128; + } +} + +inline std::tuple get_m_blocking(int64_t M) { + bool parallel_on_M = M > 128; + int64_t block_m = get_m_block(M); + int64_t Mc = (M + block_m - 1) / block_m; + int64_t Mc_parallel = parallel_on_M ? Mc : 1; + return std::make_tuple(parallel_on_M, block_m, Mc, Mc_parallel); +} + +#if defined(CPU_CAPABILITY_AVX512) +template +void zero_buffer(T* data, int64_t size) { + const int32_t vec_size = at::vec::Vectorized::size(); + auto zero_vec = at::vec::Vectorized(0); + int64_t d = 0; + for (; d < size - (size % vec_size); d += vec_size) { + zero_vec.store(data + d); + } + if (d < size) { + zero_vec.store(data + d, size - d); + } +} +#else +template +void zero_buffer(T* data, int64_t size) { + memset(data, 0, sizeof(T) * size); +} +#endif + +#if defined(__x86_64__) +bool avx10_2_available() { + // __builtin_cpu_supports returns the masked feature bit rather than 0/1. + static const bool available = __builtin_cpu_supports("avx10.2"); + return available; +} +#endif + } // anonymous namespace diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h index db7ad884d09a..dc5f9c585d5d 100644 --- a/csrc/cpu/sgl-kernels/gemm.h +++ b/csrc/cpu/sgl-kernels/gemm.h @@ -94,7 +94,7 @@ enum class CPUActMethod : int { gelu_and_mul = 2, }; -enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 }; +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 }; constexpr bool operator==(CPUQuantMethod a, int64_t b) { return static_cast(a) == b; @@ -395,3 +395,55 @@ void tinygemm_kernel( bool brg, int64_t block_size_K, bool do_unpack = true); + +// FP8 W8A8 tinygemm kernel (defined in gemm_fp8_w8a8.cpp) +void tinygemm_kernel( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf); + +// FP8 W8A8 quantize activation (defined in gemm_fp8_w8a8.cpp) +std::tuple quantize_fp8e4m3_vec( + const at::Tensor& t, + bool channelwise, + c10::optional scale_opt); + +// FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp) +template +void fused_experts_fp8_a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + at::Float8_e4m3fn* __restrict__ A_tmp, + scalar_t* __restrict__ B_tmp, + float* __restrict__ C_tmp, + float* __restrict__ Ukernel_tmp, + const at::Float8_e4m3fn* __restrict__ input, + const at::Float8_e4m3fn* __restrict__ packed_w1, + const at::Float8_e4m3fn* __restrict__ packed_w2, + const float* __restrict__ As, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int64_t block_size_N, + int64_t block_size_K, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad); diff --git a/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp new file mode 100644 index 000000000000..854566d682bd --- /dev/null +++ b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp @@ -0,0 +1,1093 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include +#include +#include + +#include "common.h" +#include "vec.h" + +// Runtime check for AMX-FP8, implemented in cpu_isa.cpp. +extern bool cpu_has_amx_fp8(); + +namespace { + +#define BLOCK_N 32 + +#define PER_TENSOR 1 +#define PER_ROW 2 +#define PER_GROUP 3 + +static bool cpublas_checked = false; +static bool cpublas_can_pack = false; + +bool cpublas_could_pack() { + // the could_pack check requires AMX support implicitly + if (cpublas_checked) { + return cpublas_can_pack; + } + cpublas_can_pack = cpu_has_amx_fp8() && at::native::cpublas::could_pack(at::kFloat8_e4m3fn); + cpublas_checked = true; + return cpublas_can_pack; +} + +#if defined(CPU_CAPABILITY_AVX512) +static void cvt_f8e4m3_to_bf16( + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) { + if (stride == cols) { + // A contiguous buffer + size_t len = rows * cols; + size_t i = 0; + for (; i < len; i += 32) { + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[i]); + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec); + _mm512_storeu_si512((__m512i*)(out + i), (__m512i)bf16_vec); + } + for (; i < len; ++i) { + out[i] = (at::BFloat16)in[i]; + } + } else { + // Non-contiguous. Access each row with stride + TORCH_CHECK(stride > cols); + for (int r = 0; r < rows; ++r) { + size_t i = 0; + size_t vec_len = cols / 32 * 32; + for (; i < vec_len; i += 32) { + __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]); + __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec); + _mm512_storeu_si512((__m512i*)(out + r * cols + i), (__m512i)bf16_vec); + } + for (; i < cols; ++i) { + out[r * cols + i] = (at::BFloat16)in[r * stride + i]; + } + } + } +} + +// accumulate and store result to buffer +// if act/wei are per_group quantized, apply scales +template +static void _accumulate_result( + float* __restrict__ output, + const float* __restrict__ input, + const float* __restrict__ scale_a, + const float* __restrict__ scale_b, + int M, + int ldi, + int ldo, + int ldsa = 1) { + float a_scale, b_scale; + __m512 va_scale; + __m512 vb_scale; + for (int m = 0; m < M; ++m) { + if constexpr (act_quant_mode == PER_GROUP) { + a_scale = *(scale_a + m * ldsa); + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto i) { + constexpr int n = i * 16; + __m512 vc_f = _mm512_loadu_ps(input + m * ldi + n); + if constexpr (act_quant_mode == PER_GROUP) { + vc_f = _mm512_mul_ps(vc_f, va_scale); + } + if constexpr (wei_quant_mode == PER_GROUP) { + vb_scale = _mm512_loadu_ps(scale_b + n); + vc_f = _mm512_mul_ps(vc_f, vb_scale); + } + if constexpr (accum) { + __m512 vo = _mm512_loadu_ps(output + m * ldo + n); + _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f)); + } else { + _mm512_storeu_ps(output + m * ldo + n, vc_f); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int n = tail_start; n < N; ++n) { + float dq_val = input[m * ldi + n]; + if constexpr (act_quant_mode == PER_GROUP) { + dq_val = dq_val * a_scale; + } + if constexpr (wei_quant_mode == PER_GROUP) { + b_scale = scale_b[n]; + dq_val = dq_val * b_scale; + } + if constexpr (accum) { + output[m * ldo + n] += dq_val; + } else { + output[m * ldo + n] = dq_val; + } + } + } +} + +// Store result to output buffer with dtype conversion +// If act/wei are per_row or per_tensor quantized, apply scales +// If bias is not null, add bias +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + __m512 va_scale, vb_scale; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + vb_scale = _mm512_set1_ps(b_scale); + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + if constexpr (act_quant_mode != PER_GROUP) { + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto idx) { + constexpr int j = idx * 16; + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps(); + if constexpr (act_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, va_scale); + } + if constexpr (wei_quant_mode == PER_ROW) { + vb_scale = _mm512_loadu_ps(scales_b + j); + } + if constexpr (wei_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, vb_scale); + } + y_vec = _mm512_add_ps(y_vec, bias_vec); + if constexpr (std::is_same::value) { + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); + } else if constexpr (std::is_same::value) { + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); + } else if constexpr (std::is_same::value) { + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); + } else { + TORCH_CHECK(false, "Unsupported output dtype"); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int j = tail_start; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#else // no AVX512 + +static void cvt_f8e4m3_to_bf16( + const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) { + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + out[r * cols + c] = (at::BFloat16)in[r * stride + c]; + } + } +} + +// Store result to output buffer with dtype conversion +// If act/wei are per_row or per_tensor quantized, apply scales +// If bias is not null, add bias +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + for (int j = 0; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#endif // CPU_CAPABILITY_AVX512 + +template +void _micro_gemm( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf) { + // If FP8 brgemm is not available, convert A/B to bf16 for computation + // Compute GEMM fp8 * fp8 -> fp32 (or bf16 * bf16 -> fp32) + // If per_group quant, apply scales. Otherwise, don't apply scales here + // Finally accumulate and store results +#if defined(CPU_CAPABILITY_AVX512) + if constexpr (cpublas_can_pack) { +#ifdef CPUBLAS_BRGEMM_F8F8F32 + at::native::cpublas::brgemm( + M, N, K, lda /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, A, B, ukernel_buf, true /* is_vnni */); +#else + cvt_f8e4m3_to_bf16(A, dqA_buf, M, K, lda); + cvt_f8e4m3_to_bf16(B, dqB_buf, K, N, N); + at::native::cpublas::brgemm( + M, N, K, K /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, dqA_buf, dqB_buf, ukernel_buf, true /* is_vnni */); +#endif + _mm_prefetch(B + N * (K + 128), _MM_HINT_T0); + _mm_prefetch(A + K + 128, _MM_HINT_T0); + _accumulate_result( + C, ukernel_buf, scales_a, scales_b, M, N /*ldi*/, ldc, ldsa); + } else +#endif + { + for (int64_t i = 0; i < M; ++i) { + for (int64_t j = 0; j < N; ++j) { + float sum = 0; + for (int64_t k = 0; k < K; ++k) { + sum += ((float)A[i * lda + k] * (float)B[k * N + j]); + } + if constexpr (act_quant_mode == PER_GROUP) { + sum *= scales_a[i * ldsa]; + } + if constexpr (wei_quant_mode == PER_GROUP) { + sum *= scales_b[j]; + } + C[i * ldc + j] += sum; + } + } + } +} + +template +void _float8_linear_impl( + const at::Tensor& input, + const at::Tensor& input_scales, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::Tensor& output) { + // input shape = [..., K] + // input is per token quantized + int64_t K = input.size(-1); + auto input_view = input.view({-1, K}); + int64_t M = input_view.size(0); + + // weight shape = [Nc, Kc, block_k, block_n] + // scales shape = [Nc, G, block_n] + int64_t Nc = weight.size(0); + int64_t Kc = wei_quant_mode != PER_GROUP ? 1 : weight.size(1); + int64_t block_k = wei_quant_mode != PER_GROUP ? weight.size(1) * weight.size(2) : weight.size(2); + constexpr int64_t block_n = BLOCK_N; + TORCH_CHECK(weight.size(3) == block_n, "Float8 linear: unexpected weight shape"); + int64_t N = Nc * block_n; + TORCH_CHECK(K == Kc * block_k, "Float8 linear: weight and input shapes mismatch"); + auto [parallel_on_M, block_m, Mc, Mc_parallel] = get_m_blocking(M); + + // scales shape = [Nc, G, block_n] + int64_t num_groups = wei_quant_mode == PER_TENSOR ? 1 : weight_scales.size(1); + TORCH_CHECK(K % num_groups == 0, "K should be divisible by num_groups"); + int64_t group_size = K / num_groups; + TORCH_CHECK(group_size % block_k == 0, "Float8 linear: group_size should be divisible by block_k"); + int64_t block_per_group = group_size / block_k; + TORCH_CHECK( + input_scales.numel() == 1 || input_scales.numel() == M || input_scales.numel() == M * num_groups, + "Float8 linear: unexpected input scales shape"); + auto ldsa = act_quant_mode == PER_TENSOR ? 0 : act_quant_mode == PER_ROW ? 1 : num_groups; + + const at::Float8_e4m3fn* a_ptr = input_view.data_ptr(); + const float* a_scales_ptr = input_scales.data_ptr(); + const at::Float8_e4m3fn* b_ptr = weight.data_ptr(); + const float* b_scales_ptr = weight_scales.data_ptr(); + out_dtype* c_ptr = output.data_ptr(); + const float* bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr; + + int64_t block_size = block_m * block_n; + int64_t num_thread = at::get_num_threads(); + at::Tensor y_buffer = at::empty({num_thread, block_size}, output.options().dtype(at::kFloat)); + // Create buffer for brgemm output and dqA/dqB (optional) +#if defined(CPU_CAPABILITY_AVX512) + // buffer for brgemm output in float32 + int64_t buffer_size = block_size * 2; // float32 = bfloat16 * 2 +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // buffers for dqA & dqB in bf16 + buffer_size += (block_k * block_n + block_m * block_k); +#endif + at::Tensor micro_gemm_buffer = at::empty({num_thread, buffer_size}, output.options().dtype(at::kBFloat16)); +#endif + + parallel_2d(Mc, Nc, [&](int64_t mc0, int64_t mc1, int64_t nc0, int64_t nc1) { + int tid = get_thread_num(); + float* y_buf = y_buffer.data_ptr() + tid * block_size; + at::BFloat16 *dqA_buffer = nullptr, *dqB_buffer = nullptr; + float* ukernel_buf = nullptr; +#if defined(CPU_CAPABILITY_AVX512) + at::BFloat16* micro_gemm_buf = micro_gemm_buffer.data_ptr() + tid * buffer_size; + ukernel_buf = reinterpret_cast(micro_gemm_buf); +#ifndef CPUBLAS_BRGEMM_F8F8F32 + dqA_buffer = micro_gemm_buf; + dqB_buffer = micro_gemm_buf + block_m * block_k; + ukernel_buf = reinterpret_cast(micro_gemm_buf + block_m * block_k + block_k * block_n); +#endif +#endif + + loop_2d(mc0, mc1, nc0, nc1, block_n * K, [&](int64_t mci, int64_t nc, int64_t) { + int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m; + zero_buffer(y_buf, m_size * block_n); + for (int kci = 0; kci < Kc; ++kci) { + auto scales_a = a_scales_ptr + mci * block_m * num_groups + kci / block_per_group; + auto scales_b = b_scales_ptr + nc * block_n * num_groups + kci / block_per_group * block_n; + _micro_gemm( + /* C */ y_buf, + /* A */ a_ptr + mci * block_m * K + kci * block_k, + /* scales_a */ scales_a, + /* B */ b_ptr + (nc * Kc + kci) * block_n * block_k, + /* scales_b */ scales_b, + /* M */ m_size, + /* K */ block_k, + /* lda */ K, + /* ldc */ block_n, + /* ldsa */ ldsa, + /* ukernel_buf */ ukernel_buf, + /* dqA_buf */ dqA_buffer, + /* dqB_buf */ dqB_buffer); + } + auto scales_a = act_quant_mode == PER_TENSOR ? a_scales_ptr + : act_quant_mode == PER_ROW ? a_scales_ptr + mci * block_m + : nullptr; + auto scales_b = wei_quant_mode == PER_TENSOR ? b_scales_ptr + : wei_quant_mode == PER_ROW ? b_scales_ptr + nc * block_n + : nullptr; + auto bias_data = bias_ptr ? bias_ptr + nc * block_n : nullptr; + store_out( + y_buf, c_ptr + mci * block_m * N + nc * block_n, m_size, N /*lda*/, scales_a, scales_b, bias_data); + }); + + if constexpr (cpublas_can_pack) { + at::native::cpublas::brgemm_release(); + } + }); +} + +} // anonymous namespace + +void tinygemm_kernel( + float* C, + const at::Float8_e4m3fn* A, + const float* scales_a, + const at::Float8_e4m3fn* B, + const float* scales_b, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + int64_t ldsa, + float* ukernel_buf, + at::BFloat16* dqA_buf, + at::BFloat16* dqB_buf) { + // act_quant_mode = per row (2), wei_quant_mode = per group (3) + if (cpublas_could_pack()) { + _micro_gemm(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf); + } else { + _micro_gemm(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf); + } +} + +#define INSTANTIATE_TINYGEMM_TEMPLATE() \ + void tinygemm_kernel( \ + float* C, \ + const at::Float8_e4m3fn* A, \ + const float* scales_a, \ + const at::Float8_e4m3fn* B, \ + const float* scales_b, \ + int64_t M, \ + int64_t K, \ + int64_t lda, \ + int64_t ldc, \ + int64_t ldsa, \ + float* ukernel_buf, \ + at::BFloat16* dqA_buf, \ + at::BFloat16* dqB_buf) + +INSTANTIATE_TINYGEMM_TEMPLATE(); +/* +return: packed_weight, packed_scales +*/ +std::tuple float8_linear_prepack_impl(const at::Tensor& weight, const at::Tensor& scales) { + // weight shape = [N, K] + // scales shape = [N, G] + TORCH_CHECK(weight.dim() == 2, "Float8 linear CPU: Weight should be a 2D tensor for packing"); + TORCH_CHECK(weight.size(1) % 2 == 0, "Float8 linear CPU: Weight should have even number of columns for packing"); + + auto new_scales = scales; + if (new_scales.dim() == 1) { + new_scales.unsqueeze_(1); + } + new_scales = new_scales.to(at::kFloat); + int N = weight.size(0); + int K = weight.size(1); + int G = scales.size(1); + int group_size = K / G; + int block_k = group_size > 128 ? 128 : group_size; + while (K % block_k != 0) { + block_k /= 2; + } + TORCH_CHECK( + block_k > 0 && block_k <= group_size, "Float8 linear CPU: Invalid block_k size, should be in (0, group_size]"); + constexpr int block_n = BLOCK_N; + int Nc = N / block_n; + int Kc = K / block_k; + + // Reorder weight to [N/block_n, K/block_k, block_k, block_n] + // Reorder scales to [N/block_n, G, block_n] + auto weight_view = weight.view({Nc, block_n, Kc, block_k}); + at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous(); + at::Tensor blocked_weight; + at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous(); + +#if defined(CPU_CAPABILITY_AVX512) + if (cpublas_could_pack()) { +#ifdef CPUBLAS_BRGEMM_F8F8F32 + constexpr int vnni_size = 4; // for fp8 +#else + constexpr int vnni_size = 2; // for float16 +#endif + blocked_weight = at::empty({Nc, Kc, block_k, block_n}, weight.options()); + auto weight_ptr = reinterpret_cast(weight_reordered.data_ptr()); + auto blocked_weight_ptr = reinterpret_cast(blocked_weight.data_ptr()); + int64_t num_blocks = Nc * Kc; + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) { + for (const auto i : c10::irange(begin, end)) { + auto in_ptr = weight_ptr + i * block_k * block_n; + auto out_ptr = blocked_weight_ptr + i * block_k * block_n; + + // Reorder weight block to VNNI + // plain shape = [block_k, block_n] + // packed shape = [block_k / VNNI_SIZE, block_n, VNNI_SIZE] viewed as [block_k, block_n] + constexpr int n_group_size = 8; + constexpr int n_group = block_n / n_group_size; // 4 + for (int nb = 0; nb < n_group; ++nb) { + for (int k = 0; k < block_k; k += vnni_size) { + for (int ni = 0; ni < n_group_size; ++ni) { + for (int ki = 0; ki < vnni_size; ++ki) { + int src_idx = nb * n_group_size + ni + (k + ki) * block_n; + int dst_idx = (nb * n_group_size + ni) * vnni_size + k * block_n + ki; + *(out_ptr + dst_idx) = *(in_ptr + src_idx); + } + } + } + } + } + }); + } else +#endif + { + blocked_weight = weight_reordered; + } + + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales)); +} + +// AVX512 optimized channel-wise max computation +inline void +compute_channel_max_avx512(const float* data, float* max_vals, int64_t num_channels, int64_t elements_per_channel) { + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const float* channel_data = data + c * elements_per_channel; + float max_val = 0.0f; + + int64_t i = 0; + // Process 16 elements at a time + __m512 max_vec = _mm512_setzero_ps(); + + for (; i <= elements_per_channel - 16; i += 16) { + __m512 data_vec = _mm512_loadu_ps(&channel_data[i]); + __m512 abs_vec = _mm512_abs_ps(data_vec); + max_vec = _mm512_max_ps(max_vec, abs_vec); + } + + // Horizontal max of the vector + float hmax = _mm512_reduce_max_ps(max_vec); + max_val = std::max(max_val, hmax); + + // Handle remaining elements + for (; i < elements_per_channel; ++i) { + max_val = std::max(max_val, std::abs(channel_data[i])); + } + + max_vals[c] = max_val; + } + }); +} + +// AVX512 optimized scaling and clamping for channel-wise quantization +inline void scale_clamp_channelwise_avx512( + const float* src, + float* dst, + const float* scales, + int64_t num_channels, + int64_t elements_per_channel, + float quant_max, + float neg_quant_max) { + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + float scale_val = scales[c]; + float scale_reciprocal = 1.0f / scale_val; + + const float* channel_src = src + c * elements_per_channel; + float* channel_dst = dst + c * elements_per_channel; + + int64_t i = 0; + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max); + + for (; i <= elements_per_channel - 16; i += 16) { + __m512 src_vec = _mm512_loadu_ps(&channel_src[i]); + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec); + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec); + _mm512_storeu_ps(&channel_dst[i], clamped_vec); + } + + // Handle remaining elements + for (; i < elements_per_channel; ++i) { + float scaled = channel_src[i] * scale_reciprocal; + channel_dst[i] = std::clamp(scaled, neg_quant_max, quant_max); + } + } + }); +} + +// AVX512 optimized scaling and clamping for global quantization +inline void scale_clamp_global_avx512( + const float* src, float* dst, int64_t size, float scale_reciprocal, float quant_max, float neg_quant_max) { + if (size <= 0) return; + + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max); + + at::parallel_for(0, size, 4096, [&](int64_t start, int64_t end) { + int64_t block_size = end - start; + const float* block_src = src + start; + float* block_dst = dst + start; + + int64_t i = 0; + for (; i <= block_size - 16; i += 16) { + __m512 src_vec = _mm512_loadu_ps(&block_src[i]); + __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec); + __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec); + _mm512_storeu_ps(&block_dst[i], clamped_vec); + } + + // Handle remaining elements in block + for (; i < block_size; ++i) { + float scaled = block_src[i] * scale_reciprocal; + block_dst[i] = std::clamp(scaled, neg_quant_max, quant_max); + } + }); +} + +std::tuple +_quantize_fp8e4m3(const at::Tensor& t, bool channelwise, c10::optional scale_opt = c10::nullopt) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Ensure input is contiguous and in float32 + auto t_float = t.to(at::ScalarType::Float).contiguous(); + at::Tensor qt; + at::Tensor scale_tensor; + + if (channelwise) { + // Channel-wise quantization with AVX512 optimization + int64_t num_channels = t_float.size(0); + int64_t elements_per_channel = t_float.numel() / num_channels; + + if (elements_per_channel * num_channels != t_float.numel()) { + throw std::runtime_error("Tensor must be divisible by number of channels for channel-wise quantization"); + } + + // Allocate scale tensor + scale_tensor = at::empty({num_channels}, t_float.options()); + float* scale_data = scale_tensor.data_ptr(); + + // Compute channel-wise max using AVX512 + compute_channel_max_avx512(t_float.data_ptr(), scale_data, num_channels, elements_per_channel); + + // Apply quant_max and EPS + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + scale_data[c] = std::max(scale_data[c] / quant_max, eps); + } + }); + + // Create output tensor for quantized values + qt = at::empty_like(t_float); + + // Scale and clamp using AVX512 + scale_clamp_channelwise_avx512( + t_float.data_ptr(), + qt.data_ptr(), + scale_data, + num_channels, + elements_per_channel, + quant_max, + -quant_max); + } else { + // Global quantization with AVX512 optimization + if (!scale_opt.has_value()) { + throw std::runtime_error("Scale must be provided for non-channelwise quantization in AVX512 version"); + } + + scale_tensor = scale_opt.value().to(at::ScalarType::Float).contiguous(); + + // Handle scalar scale case + float scale_val; + if (scale_tensor.numel() == 1) { + scale_val = std::max(scale_tensor.item(), eps); + } else { + // Compute max of scale tensor if it's not scalar + float* scale_data = scale_tensor.data_ptr(); + __m512 max_vec = _mm512_set1_ps(eps); + + int64_t i = 0; + for (; i <= scale_tensor.numel() - 16; i += 16) { + __m512 scale_vec = _mm512_loadu_ps(&scale_data[i]); + max_vec = _mm512_max_ps(max_vec, scale_vec); + } + + float hmax = _mm512_reduce_max_ps(max_vec); + for (; i < scale_tensor.numel(); ++i) { + hmax = std::max(hmax, scale_data[i]); + } + scale_val = std::max(hmax, eps); + } + + scale_tensor = at::tensor({scale_val}, scale_tensor.options()); + float scale_reciprocal = 1.0f / scale_val; + + // Create output tensor + qt = at::empty_like(t_float); + + // Scale and clamp using AVX512 + scale_clamp_global_avx512( + t_float.data_ptr(), qt.data_ptr(), t_float.numel(), scale_reciprocal, quant_max, -quant_max); + } + + // Final conversion to FP8 E4M3FN using ATen's native conversion + qt = at::_to_copy(qt, at::ScalarType::Float8_e4m3fn); + + return std::make_tuple(qt, scale_tensor); +} + +inline __m128i cvtfp32_fp8e4m3(__m512& src) { + // cvt 16x32 from fp32 to fp8 e4m3 + const __m512i sign_mask = _mm512_set1_epi32(0x80000000); + const __m512i fp8_max = _mm512_set1_epi32(UINT32_C(1087) << 20); + const __m512i denorm_thresh = _mm512_set1_epi32(UINT32_C(121) << 23); + const __m512i denorm_mask = _mm512_set1_epi32(UINT32_C(141) << 23); + const __m512i bias_part1 = _mm512_set1_epi32((uint32_t)(7 - 127) << 23); + const __m512i rounding_bias = _mm512_set1_epi32(0x7FFFF); + __m512i f_bits = _mm512_castps_si512(src); + // Extract and save sign + __m512i sign = _mm512_and_epi32(f_bits, sign_mask); + f_bits = _mm512_xor_epi32(f_bits, sign); + + // Prepare result containers + __m512i result = _mm512_setzero_si512(); + + // Step 1: Handle case of overflow + // (f_bits >= fp8_max): set result = 0x7f + __mmask16 overflow_mask = _mm512_cmpge_epu32_mask(f_bits, fp8_max); + if (overflow_mask) { + result = _mm512_mask_set1_epi32(result, overflow_mask, 0x7f); + } + + // Step 2: Handle small numbers (denormals) + // Small numbers (f_bits < denorm_thresh) + __mmask16 denorm_thresh_mask = _mm512_cmplt_epu32_mask(f_bits, denorm_thresh); + + if (denorm_thresh_mask) { + __m512 small_input = _mm512_castsi512_ps(f_bits); + __m512 small_denorm = _mm512_add_ps(small_input, _mm512_castsi512_ps(denorm_mask)); + __m512i small_denorm_bits = _mm512_castps_si512(small_denorm); + __m512i small_result = _mm512_sub_epi32(small_denorm_bits, denorm_mask); + result = _mm512_mask_mov_epi32(result, denorm_thresh_mask, small_result); + } + + // Step 3: Handle normal numbers + __mmask16 normal_mask = ~(overflow_mask | denorm_thresh_mask); + + if (normal_mask) { + // mant_odd = (f_bits >> 20) & 1 + __m512i mant_odd = _mm512_and_epi32(_mm512_srli_epi32(f_bits, 20), _mm512_set1_epi32(1)); + // f_bits += bias_part1 + rounding_bias + __m512i rounded = _mm512_add_epi32(f_bits, bias_part1); + rounded = _mm512_add_epi32(rounded, rounding_bias); + // Add mant_odd + rounded = _mm512_add_epi32(rounded, mant_odd); + // Shift right by 20 bits + __m512i normal_result = _mm512_srli_epi32(rounded, 20); + result = _mm512_mask_mov_epi32(result, normal_mask, normal_result); + } + + // Merge back the sign + __m512i sign_shifted = _mm512_srli_epi32(sign, 24); + result = _mm512_or_epi32(result, sign_shifted); + + // Now result is 16 x 32-bit integers, but we only need 8-bit for each + __m512i packed = _mm512_and_si512(result, _mm512_set1_epi32(0xFF)); + + // Narrow 32-bit integers to 8-bit + return _mm512_cvtepi32_epi8(packed); +} + +__attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) { + __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec)); +} + +// Carries the whole loop rather than just the conversion, as GCC cannot inline an avx10.2 +// function into a caller lacking that target, so branching per vector is costly. +__attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2( + const at::BFloat16* __restrict__ src, at::Float8_e4m3fn* __restrict__ dst, int64_t count, float scale_reciprocal) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + + for (int64_t i = 0; i <= count - 32; i += 32) { + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&src[i + 16])); + + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3_avx10_2(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3_avx10_2(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2); + } +} + +std::tuple _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Input validation and preparation + assert(t.scalar_type() == at::ScalarType::BFloat16); + auto t_bf16 = t.contiguous(); + int64_t num_channels = t_bf16.size(0); + int64_t elements_per_channel = t_bf16.numel() / num_channels; + assert(elements_per_channel % 32 == 0); // do not consider tile currently + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); + + // Allocate output tensors + at::Tensor scale_tensor = at::empty({num_channels}, t_bf16.options().dtype(at::ScalarType::Float)); + float* scale_data = scale_tensor.data_ptr(); + + // Unified processing: compute max, apply scale, and quantize in single pass + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const at::BFloat16* channel_src = t_bf16.data_ptr() + c * elements_per_channel; + at::Float8_e4m3fn* quant_dst = quant_t.data_ptr() + c * elements_per_channel; + + // Step 1: Compute channel-wise max using AVX512 + float channel_max = 0.0f; + int64_t i = 0; + + // Process 32 elements at a time for max computation + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values (2x 256-bit vectors) + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i + 16])); + + // Convert BF16 to FP32 + __m512i fp32_int1 = _mm512_cvtepu16_epi32(src_vec1); + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int1, 16)); + + __m512i fp32_int2 = _mm512_cvtepu16_epi32(src_vec2); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int2, 16)); + + // Compute absolute values + __m512 abs_vec1 = _mm512_abs_ps(fp32_vec1); + __m512 abs_vec2 = _mm512_abs_ps(fp32_vec2); + + // Find max in each vector + float max1 = _mm512_reduce_max_ps(abs_vec1); + float max2 = _mm512_reduce_max_ps(abs_vec2); + channel_max = std::max(channel_max, std::max(max1, max2)); + } + + // Step 2: Apply quant_max and EPS to compute scale + float scale_val = std::max(channel_max / quant_max, eps); + float scale_reciprocal = 1.0f / scale_val; + scale_data[c] = scale_val; + + // Step 3: Scale and clamp using AVX512 (reuse the same loop structure) + if (avx10_2_available()) { + quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal); + continue; + } + i = 0; + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&channel_src[i + 16])); + + // Convert BF16 to FP32 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + // Scale and clamp + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i + 16), fp8_vec2); + } + } + }); + + return std::make_tuple(quant_t, scale_tensor); +} + +std::tuple +_quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& scale_tensor) { + constexpr float quant_max = 448.0f; // torch.finfo(torch.float8_e4m3fn).max + constexpr float eps = std::numeric_limits::epsilon(); + + // Input validation + assert(t.scalar_type() == at::ScalarType::BFloat16); + + auto t_bf16 = t.contiguous(); + auto scale_tensor_contig = scale_tensor.contiguous().to(at::ScalarType::Float); + + // Get scale value (handle scalar or tensor) + float scale_val; + if (scale_tensor_contig.numel() == 1) { + scale_val = scale_tensor_contig.item(); + } else { + // Take max if scale is a tensor (though should be scalar for global quantization) + scale_val = scale_tensor_contig.max().item(); + } + + // Apply EPS to scale + scale_val = std::max(scale_val, eps); + float scale_reciprocal = 1.0f / scale_val; + + at::Tensor scale_output = at::tensor({scale_val}, scale_tensor_contig.options()); + + // Apply scale and clamp using AVX512 + const at::BFloat16* src_data = t_bf16.data_ptr(); + at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn)); + int64_t total_elements = t_bf16.numel(); + const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal); + const __m512 quant_max_vec = _mm512_set1_ps(quant_max); + const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max); + int64_t num_channels = t_bf16.size(0); + int64_t elements_per_channel = total_elements / num_channels; + assert(elements_per_channel % 32 == 0); // do not consider tile currently + // Process in parallel blocks + at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) { + for (int64_t c = start; c < end; ++c) { + const at::BFloat16* src_data = t_bf16.data_ptr() + c * elements_per_channel; + at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr() + c * elements_per_channel; + + if (avx10_2_available()) { + quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal); + continue; + } + + int64_t i = 0; + // Process 32 elements at a time using AVX512 + for (; i <= elements_per_channel - 32; i += 32) { + // Load 32 BF16 values (2x 256-bit vectors) + __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast(&src_data[i])); + __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast(&src_data[i + 16])); + + // Convert BF16 to FP32 + __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16)); + __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16)); + + // Scale and clamp + __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec); + __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec); + + __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec); + __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec); + + __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1); + __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i), fp8_vec1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i + 16), fp8_vec2); + } + } + }); + + return std::make_tuple(quant_t, scale_output); +} + +std::tuple +_quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional scale_opt) { + if (channelwise) { + return _quantize_fp8e4m3_bf16_per_tensor_no_scale(t); + } else { + assert(scale_opt.has_value()); + return _quantize_fp8e4m3_bf16_per_tensor_with_scale(t, scale_opt.value()); + } +} + +// Public wrapper for torch op registration +std::tuple +quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional scale_opt) { + return _quantize_fp8e4m3_vec(t, channelwise, scale_opt); +} + +at::Tensor fp8_scaled_mm_with_quant( + const at::Tensor& act, + const std::optional& act_scales, + bool channelwise, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype) { + // TODO: refine here to use tensor ptr and in below dispatch api. + std::tuple quant_act = act.scalar_type() == at::ScalarType::BFloat16 + ? _quantize_fp8e4m3_vec(act, channelwise, act_scales) + : _quantize_fp8e4m3(act, channelwise, act_scales); + auto input = std::get<0>(quant_act); + auto input_scales = std::get<1>(quant_act); + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW + : PER_GROUP; + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; + // Case to fall back + if (weight.dim() == 2) { + TORCH_CHECK( + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, + "FP8 linear: Per-group quantization is not supported in the fallback path"); + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); + return y_fp32.to(output_dtype); + } + + static bool cpublas_can_pack = cpublas_could_pack(); + auto out_sizes = input.sizes().vec(); + out_sizes.back() = N; + auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); + +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ + AT_DISPATCH_BOOL_NO_RETURN( \ + CAN_PACK, \ + "cpublas_can_pack", \ + can_pack, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + A_QUANT_MODE, \ + "act_quant_mode", \ + a_quant_mode, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + B_QUANT_MODE, \ + "wei_quant_mode", \ + b_quant_mode, \ + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) + + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { + _float8_linear_impl( + input, input_scales, weight, weight_scales, bias, output); + }); + return output; +} + +at::Tensor float8_linear_impl( + const at::Tensor& input, + const at::Tensor& input_scales, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype) { + int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0); + int act_quant_mode = input_scales.numel() == 1 ? PER_TENSOR + : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW + : PER_GROUP; + int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP; + // Case to fall back + if (weight.dim() == 2) { + TORCH_CHECK( + act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP, + "FP8 linear: Per-group quantization is not supported in the fallback path"); + auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias); + return y_fp32.to(output_dtype); + } + + static bool cpublas_can_pack = cpublas_could_pack(); + auto out_sizes = input.sizes().vec(); + out_sizes.back() = N; + auto output = at::empty(out_sizes, input.options().dtype(output_dtype)); + +#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \ + AT_DISPATCH_BOOL_NO_RETURN( \ + CAN_PACK, \ + "cpublas_can_pack", \ + can_pack, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + A_QUANT_MODE, \ + "act_quant_mode", \ + a_quant_mode, \ + AT_DISPATCH_QUANT_MODE_NO_RETURN( \ + B_QUANT_MODE, \ + "wei_quant_mode", \ + b_quant_mode, \ + AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__)))) + + AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() { + _float8_linear_impl( + input, input_scales, weight, weight_scales, bias, output); + }); + return output; +} \ No newline at end of file diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp index 3afdc1bd6791..09a170eb6d4c 100644 --- a/csrc/cpu/sgl-kernels/moe.cpp +++ b/csrc/cpu/sgl-kernels/moe.cpp @@ -836,6 +836,11 @@ static inline void check_moe_scales( TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4."); TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8."); TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8."); + } else if constexpr (quant == CPUQuantMethod::FP8_W8A8) { + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a8."); + TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a8."); + TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a8."); + TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2."); } } @@ -850,9 +855,14 @@ static inline void check_moe_scales( check_moe_scales(w1_scale, w2_scale, block_size); } else if (moe_comp_method == CPUQuantMethod::MXFP4) { check_moe_scales(w1_scale, w2_scale, block_size); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + check_moe_scales(w1_scale, w2_scale, block_size); } } +// CHECK_MOE_SCALES_FP8: validate scale shapes for FP8_W8A16. +// w1s shape [E, 2N/block_N, K/block_K, block_N]: DIM0=1, DIM1=2 +// w2s shape [E, K/block_N, N/block_K]: DIM0=1, DIM1=2 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \ auto w1s = w1_scale.value(); \ auto w2s = w2_scale.value(); \ @@ -864,6 +874,17 @@ static inline void check_moe_scales( TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N)); \ TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K)) +// CHECK_MOE_SCALES_FP8_W8A8: extract scale info for FP8_W8A8 without shape +// validation, because w1s is pre-packed by float8_linear_prepack_cpu with +// shape [E, 2N/BLOCK_N, G, BLOCK_N] where BLOCK_N=32 != block_size_N(128). +// This matches the sglang upstream behavior (checks are commented out there). +#define CHECK_MOE_SCALES_FP8_W8A8(DIM0, DIM1) \ + auto w1s = w1_scale.value(); \ + auto w2s = w2_scale.value(); \ + auto block_size_val = block_size.value(); \ + int64_t block_size_N = block_size_val[0]; \ + int64_t block_size_K = block_size_val[1] + // hidden_states: [M, K] // w1: [E, 2N, K] or [E, 2N, K / 2] for uint8 // w2: [E, K, N] or [E, K, N / 2] for uint8 @@ -883,6 +904,7 @@ void fused_experts_cpu( const std::optional& w2_scale, const std::optional& w1_zero, const std::optional& w2_zero, + const std::optional& a1_scale, const std::optional> block_size, const std::optional& w1_bias, const std::optional& w2_bias, @@ -895,7 +917,10 @@ void fused_experts_cpu( constexpr int64_t BLOCK_M = block_size_m(); constexpr int64_t BLOCK_N = block_size_n(); - const auto st = hidden_states.scalar_type(); + // FP8_W8A8 hidden_states input is FP8; scalar_t should be BF16 for output + const auto st = (moe_comp_method == CPUQuantMethod::FP8_W8A8) + ? at::kBFloat16 + : hidden_states.scalar_type(); // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax) // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt); @@ -908,6 +933,9 @@ void fused_experts_cpu( if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) { CHECK_DIM(4, w1); CHECK_DIM(4, w2); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // w1 is prepackaged by float8_linear_prepack_cpu (5D), skip w1 dim check + CHECK_DIM(3, w2); } else { CHECK_DIM(3, w1); CHECK_DIM(3, w2); @@ -924,8 +952,9 @@ void fused_experts_cpu( int64_t M = hidden_states.size(0); int64_t K = hidden_states.size(1); - int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2 - : w1.size(1) / 2; + int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 + ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2 + : (moe_comp_method == CPUQuantMethod::FP8_W8A8 ? w2.size(2) : w1.size(1) / 2); int64_t E = w1.size(0); int64_t topk = topk_weights_.size(1); @@ -935,7 +964,10 @@ void fused_experts_cpu( // check weight shapes CHECK_EQ(w2.size(0), E); - if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) { + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + CHECK_EQ(w2.size(1), K); + CHECK_EQ(packed_w2.size(2), packed_N); + } else if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) { CHECK_EQ(w2.size(1), K); CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); @@ -993,22 +1025,35 @@ void fused_experts_cpu( // 5. Aq_tmp : [M, K] or [M * topk, N] // 6. As_tmp : [M * topk] // - // for fp8 w8a16 and mxfp4: + // for fp8 w8a16, mxfp4, and fp8 w8a8: // 7. intermediate_cache0 : [M * topk, 2N] // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)] // + // for fp8 w8a8 additionally: + // 9. Ukernel_tmp : [T, 2 * BLOCK_M * BLOCK_N] (float32 for FP8×FP8 brgemm) + // int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 + num_threads * BLOCK_M * K * - (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) + + (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 + | moe_comp_method == CPUQuantMethod::FP8_W8A8 ? 1 : 2) + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); if (moe_comp_method == CPUQuantMethod::INT8_W8A8) { buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float); } - if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) { + if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4 + || moe_comp_method == CPUQuantMethod::FP8_W8A8) { buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2; } + if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // Ukernel buffer: float32 workspace per thread for FP8×FP8 brgemm + buffer_size_nbytes += num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // BF16 fallback: also need dqA [T, BLOCK_M, BLOCK_K] + dqB [T, BLOCK_K, BLOCK_N] + buffer_size_nbytes += num_threads * (BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N) * 2; +#endif + } if (moe_comp_method == CPUQuantMethod::INT4_W4A8) { buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) + num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N; @@ -1181,6 +1226,47 @@ void fused_experts_cpu( E, topk, num_tokens_post_pad); + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) { + // Stage 1: FP8×FP8 GEMM (hidden_states is FP8 input) + // Stage 2: BF16×FP8 GEMM (W8A16 path for down-projection) + at::Float8_e4m3fn* __restrict__ A_tmp = + (at::Float8_e4m3fn*)((void*)(intermediate_cache2 + M * topk * K)); + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); + scalar_t* __restrict__ intermediate_cache0 = + (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); + scalar_t* __restrict__ B_tmp = + (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); + float* __restrict__ Ukernel_tmp = + (float*)((void*)(B_tmp + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N))); + + CHECK_MOE_SCALES_FP8_W8A8(1, 2); + fused_experts_fp8_a8_kernel_impl( + out_hidden_states.data_ptr(), + intermediate_cache0, + intermediate_cache1, + intermediate_cache2, + A_tmp, + B_tmp, + C_tmp, + Ukernel_tmp, + hidden_states.data_ptr(), + packed_w1.data_ptr(), + packed_w2.data_ptr(), + a1_scale.has_value() ? a1_scale.value().data_ptr() : nullptr, + w1s.data_ptr(), + w2s.data_ptr(), + block_size_N, + block_size_K, + topk_weights_.data_ptr(), + sorted_ids, + expert_ids, + offsets, + M, + N, + K, + E, + topk, + num_tokens_post_pad); } else { scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K; float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); diff --git a/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp new file mode 100644 index 000000000000..0f393a34a853 --- /dev/null +++ b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp @@ -0,0 +1,432 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include + +#include "common.h" +#include "gemm.h" +#include "moe.h" + +namespace { + +// Byte-wise copy for FP8 elements (no remainder assumption) +inline void copy_stub2( + at::Float8_e4m3fn* __restrict__ out, + const at::Float8_e4m3fn* __restrict__ input, + int64_t size) { + for (int64_t d = 0; d < size; d++) { + out[d] = input[d]; + } +} + +#define PER_TENSOR 1 +#define PER_ROW 2 +#define PER_GROUP 3 + +// Store float32 GEMM result to output buffer with dequantization and dtype +// conversion. Activations are PER_ROW quantized; weights are PER_GROUP quantized. +// Uses AVX512 on capable platforms. +#if defined(CPU_CAPABILITY_AVX512) + +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + __m512 va_scale, vb_scale; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + vb_scale = _mm512_set1_ps(b_scale); + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + if constexpr (act_quant_mode != PER_GROUP) { + va_scale = _mm512_set1_ps(a_scale); + } + constexpr int N_UNROLL = N / 16; + c10::ForcedUnroll{}([&](auto idx) { + constexpr int j = idx * 16; + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps(); + if constexpr (act_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, va_scale); + } + if constexpr (wei_quant_mode == PER_ROW) { + vb_scale = _mm512_loadu_ps(scales_b + j); + } + if constexpr (wei_quant_mode != PER_GROUP) { + y_vec = _mm512_mul_ps(y_vec, vb_scale); + } + y_vec = _mm512_add_ps(y_vec, bias_vec); + if constexpr (std::is_same::value) { + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); + } else if constexpr (std::is_same::value) { + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); + } else if constexpr (std::is_same::value) { + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); + } else { + TORCH_CHECK(false, "Unsupported output dtype"); + } + }); + constexpr int tail_start = N / 16 * 16; + for (int j = tail_start; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + c_ptr[i * lda + j] = static_cast(y_buf[i * N + j] * a_scale * b_scale); + } + } // for M +} + +#else // no AVX512 + +template +inline void store_out( + const float* y_buf, + out_dtype* c_ptr, + int64_t M, + int64_t lda, + const float* scales_a, + const float* scales_b, + const float* bias) { + float a_scale = 1.0, b_scale = 1.0; + if constexpr (act_quant_mode == PER_TENSOR) { + a_scale = *scales_a; + } + if constexpr (wei_quant_mode == PER_TENSOR) { + b_scale = *scales_b; + } + for (int i = 0; i < M; ++i) { + if constexpr (act_quant_mode == PER_ROW) { + a_scale = *(scales_a + i); + } + for (int j = 0; j < N; ++j) { + if constexpr (wei_quant_mode == PER_ROW) { + b_scale = scales_b[j]; + } + float val = y_buf[i * N + j] * a_scale * b_scale; + if (bias) val += bias[j]; + c_ptr[i * lda + j] = static_cast(val); + } + } +} + +#endif // CPU_CAPABILITY_AVX512 + +} // anonymous namespace + +// FP8 W8A8 fused MoE kernel. +// +// hidden_states: [M, K] FP8 E4M3, pre-quantized activation +// packed_w1: [E, 2*Nc, Kc, BLOCK_K, BLOCK_N] FP8 W8A8 prepack (float8_linear_prepack_cpu) +// packed_w2: [E, K, N_packed] FP8 W8A16 VNNI pack (convert_weight_packed) +// As: [M] float32 per-token activation scales +// w1s: [E, num_groups, 2N] float32 per-group weight scales for w1 +// w2s: [E, ...] float32 per-group weight scales for w2 +// +// Stage 1: hidden_states(FP8) @ w1(FP8) → ic0(BF16) [M*topk, 2N] +// Uses FP8×FP8 GEMM (tinygemm_kernel with FP8 inputs) +// Stage 1.5: ic1 = silu(ic0[:,:N]) * ic0[:,N:] [M*topk, N] +// Stage 2: ic1(BF16) @ w2(FP8) → ic2(BF16) [M*topk, K] +// Uses BF16×FP8 GEMM (W8A16 path, tinygemm_kernel) +// Stage 3: out = weighted_sum(ic2, topk_weights) [M, K] + +template +void fused_experts_fp8_a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + at::Float8_e4m3fn* __restrict__ A_tmp, + scalar_t* __restrict__ B_tmp, + float* __restrict__ C_tmp, + float* __restrict__ Ukernel_tmp, + const at::Float8_e4m3fn* __restrict__ input, + const at::Float8_e4m3fn* __restrict__ packed_w1, + const at::Float8_e4m3fn* __restrict__ packed_w2, + const float* __restrict__ As, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int64_t block_size_N, + int64_t block_size_K, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad) { + // 1. intermediate_cache1 : [M * topk, N] + // 2. intermediate_cache2 : [M * topk, K] + // 3. A_tmp : [T, BLOCK_M * K] FP8 + // 4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N] float32 + // 5. intermediate_cache0 : [M * topk, 2N] + // 6. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, max(K, N)] scalar_t + // 7. Ukernel_tmp: [T, 2 * BLOCK_M * BLOCK_N] float32 + // 8. (non-FP8-brgemm only) dqA [T, BLOCK_M * BLOCK_K] + dqB [T, BLOCK_K * BLOCK_N] in BF16 + // appended after Ukernel_tmp; see moe.cpp FP8_W8A8 buffer allocation + + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N_VAL = block_size_n(); // = 32 + const int64_t KB = K / BLOCK_K; // K blocks + + int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N_VAL * std::max(K, N); + + // stage 1: intermediate_cache0 = hidden_states @ w1 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N_VAL); + + int64_t scale_size_N = div_up(2 * N, block_size_N); + int64_t scale_size_K = div_up(K, block_size_K); + int64_t blocks_n_per_group = block_size_N / BLOCK_N_VAL; + int64_t num_groups = div_up(K, block_size_K); // G + int64_t blocks_k_per_group = block_size_K / BLOCK_K; + const int64_t stride_e = 2 * N * K; // per-expert stride for w1 (total elements) + bool use_brgemm = true; + int64_t num_thread = at::get_num_threads(); + + // For the BF16 fallback path (no native FP8 brgemm), tinygemm_kernel needs temporary + // BF16 buffers to dequantize FP8 inputs before running AMX-BF16 brgemm. + // These are placed after the Ukernel_tmp region (see moe.cpp FP8_W8A8 buffer allocation). +#ifndef CPUBLAS_BRGEMM_F8F8F32 + // dqA: [T, BLOCK_M, BLOCK_K] BF16 + // dqB: [T, BLOCK_K, BLOCK_N] BF16 + // Layout: immediately after Ukernel_tmp[num_thread * 2 * BLOCK_M * BLOCK_N_VAL] + at::BFloat16* __restrict__ dq_base = reinterpret_cast( + Ukernel_tmp + num_thread * 2 * BLOCK_M * BLOCK_N_VAL); + const int64_t dq_per_thread = BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N_VAL; +#endif + + at::parallel_for(0, MB * NB, 1, [&](int64_t begin, int64_t end) { + int tid = get_thread_num(); + // A_tmp: FP8 activation buffer for this thread [BLOCK_M, K] + at::Float8_e4m3fn* __restrict__ A = A_tmp + tid * BLOCK_M * K; + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N_VAL; + alignas(64) float As_[BLOCK_M]; + float* __restrict__ ukernel_buf_1 = Ukernel_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; +#ifndef CPUBLAS_BRGEMM_F8F8F32 + at::BFloat16* __restrict__ dqA_buf = dq_base + tid * dq_per_thread; + at::BFloat16* __restrict__ dqB_buf = dqA_buf + BLOCK_M * BLOCK_K; +#else + at::BFloat16* __restrict__ dqA_buf = nullptr; + at::BFloat16* __restrict__ dqB_buf = nullptr; +#endif + float* __restrict__ ukernel_buf_2 = ukernel_buf_1 + BLOCK_M * BLOCK_N_VAL; + auto ldsa = 1; // activation per-row quantized + + for (int64_t i = begin; i < end; ++i) { + int64_t mb = i / NB; + int64_t nb = i % NB; + int64_t nb1 = nb + NB; // gate is first N columns, up is second N columns + + int64_t n_size = std::min(N - nb * BLOCK_N_VAL, BLOCK_N_VAL); + + // B (weight) for expert expert_id: gate+up fused [2N, K] prepackaged as FP8 + int32_t expert_id = expert_ids[mb]; + const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e; + const float* __restrict__ Bs = w1s + expert_id * num_groups * (2 * N); + + // Load activation tile for this M-block + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + int64_t m_size = offsets[mb + 1] - offsets[mb]; + + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m] / topk; + copy_stub2(A + m * K, input + index * K, K); + As_[m] = As[index]; + } + + const int64_t offset = offsets[mb]; + + // Gate path (nb): w1[:N, :] + zero_buffer(C0, BLOCK_M * BLOCK_N_VAL); + for (int kci = 0; kci < KB; ++kci) { + // FP8×FP8 tinygemm: act_quant=PER_ROW(2), wei_quant=PER_GROUP(3) + tinygemm_kernel( + /* C */ C0, + /* A */ A + kci * BLOCK_K, + /* scales_a*/ As_ + kci / blocks_k_per_group, + /* B */ B + (nb * KB + kci) * BLOCK_K * BLOCK_N_VAL, + /* scales_b*/ Bs + nb * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, + /* M */ m_size, + /* K */ BLOCK_K, + /* lda */ K, + /* ldc */ BLOCK_N_VAL, + /* ldsa */ ldsa, + /* ukernel */ ukernel_buf_1, + /* dqA_buf */ dqA_buf, + /* dqB_buf */ dqB_buf); + } + store_out( + C0, + ic0 + offset * 2 * N + nb * BLOCK_N_VAL, + m_size, + 2 * N, // lda + As_, + nullptr, + nullptr); + + // Up path (nb1): w1[N:, :] + zero_buffer(C1, BLOCK_M * BLOCK_N_VAL); + for (int kci = 0; kci < KB; ++kci) { + tinygemm_kernel( + /* C */ C1, + /* A */ A + kci * BLOCK_K, + /* scales_a*/ As_ + kci / blocks_k_per_group, + /* B */ B + (nb1 * KB + kci) * BLOCK_K * BLOCK_N_VAL, + /* scales_b*/ Bs + nb1 * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, + /* M */ m_size, + /* K */ BLOCK_K, + /* lda */ K, + /* ldc */ BLOCK_N_VAL, + /* ldsa */ ldsa, + /* ukernel */ ukernel_buf_2, + /* dqA_buf */ dqA_buf, + /* dqB_buf */ dqB_buf); + } + store_out( + C1, + ic0 + offset * 2 * N + nb1 * BLOCK_N_VAL, + m_size, + 2 * N, // lda + As_, + nullptr, + nullptr); + } + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:] + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); + } + }); + + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2 + // w2: [E, K, N] FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight) + const int64_t OC = K; // output channels (hidden dim) + const int64_t IC = N; // input channels (expert intermediate dim) + const int64_t MB2 = MB; + const int64_t NB2 = div_up(OC, BLOCK_N_VAL); + scale_size_N = div_up(K, block_size_N); + scale_size_K = div_up(N, block_size_K); + const int64_t stride_e2 = OC * IC; // per-expert stride for w2 + const int64_t stride_oc = IC; + + int64_t avg_M = std::max(int64_t(1), M * topk / E); + bool use_brgemm2 = can_use_brgemm(avg_M); + + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + int tid = get_thread_num(); + alignas(64) scalar_t C[BLOCK_M * BLOCK_K]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N_VAL * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t m_size = offsets[mb + 1] - offsets[mb]; + int64_t n_size = std::min(OC - nb * BLOCK_N_VAL, BLOCK_N_VAL); + + // A from ic1: [M*topk, N] in sorted order (contiguous after stage 1.5) + const scalar_t* __restrict__ A2 = ic1 + offsets[mb] * N; + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + + // B from w2: [E, K, N] VNNI packed, offset to expert + output block + int32_t expert_id = expert_ids[mb]; + const at::Float8_e4m3fn* __restrict__ B2 = + packed_w2 + expert_id * stride_e2 + nb * BLOCK_N_VAL * stride_oc; + const float* __restrict__ Bs2 = + w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K; + + // Only unpack B on first visit or new expert + int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1]; + bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id); + + tinygemm_kernel( + /* A */ A2, + /* B */ B2, + /* C */ C, + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N_VAL * IC, + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL, + /* Bbias */ nullptr, + /* scale */ Bs2, + /* M */ m_size, + /* N */ n_size, + /* K */ IC, + /* lda */ IC, + /* ldb */ n_size, + /* ldc */ BLOCK_N_VAL, + /* brg */ use_brgemm2, + /* block_sz_K */ block_size_K, + /* do_unpack */ do_unpack); + + // Scatter ic2 in original token order, weighting by topk_weights + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m]; + float weight = topk_weights[index]; + copy_mul_stub(ic2 + index * K + nb * BLOCK_N_VAL, C + m * BLOCK_N_VAL, weight, n_size); + } + }); + + if (use_brgemm2) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 3: output = sum(ic2, dim=topk) [M, K] + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + sum_stub(output + m * K, ic2 + m * topk * K, topk, K); + } + }); +} + +#define INSTANTIATE_MOE_FP8_A8_TEMPLATE(TYPE) \ + template void fused_experts_fp8_a8_kernel_impl( \ + TYPE* __restrict__ output, \ + TYPE* __restrict__ ic0, \ + TYPE* __restrict__ ic1, \ + TYPE* __restrict__ ic2, \ + at::Float8_e4m3fn* __restrict__ A_tmp, \ + TYPE* __restrict__ B_tmp, \ + float* __restrict__ C_tmp, \ + float* __restrict__ Ukernel_tmp, \ + const at::Float8_e4m3fn* __restrict__ input, \ + const at::Float8_e4m3fn* __restrict__ packed_w1, \ + const at::Float8_e4m3fn* __restrict__ packed_w2, \ + const float* __restrict__ As, \ + const float* __restrict__ w1s, \ + const float* __restrict__ w2s, \ + int64_t block_size_N, \ + int64_t block_size_K, \ + const float* __restrict__ topk_weights, \ + const int32_t* __restrict__ sorted_ids, \ + const int32_t* __restrict__ expert_ids, \ + const int32_t* __restrict__ offsets, \ + int64_t M, \ + int64_t N, \ + int64_t K, \ + int64_t E, \ + int64_t topk, \ + int64_t num_tokens_post_pad) + +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16); +INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 88170068dea3..accd5fcec354 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -73,6 +73,7 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states, const std::optional& w2_scale, const std::optional& w1_zero, const std::optional& w2_zero, + const std::optional& a1_scale, const std::optional> block_size, const std::optional& w1_bias, const std::optional& w2_bias, @@ -91,6 +92,18 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni); +// Adapted from sglang: FP8 W8A8 kernels +std::tuple float8_linear_prepack_impl( + const at::Tensor& weight, const at::Tensor& scales); +at::Tensor fp8_scaled_mm_with_quant(const at::Tensor& act, + const std::optional& act_scales, + bool channelwise, const at::Tensor& weight, + const at::Tensor& weight_scales, + const std::optional& bias, + at::ScalarType output_dtype); +std::tuple quantize_fp8e4m3_vec( + const at::Tensor& t, bool channelwise, c10::optional scale_opt); + // Adapted from sglang: MLA CPU kernels (AMX-only) void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& output, @@ -281,6 +294,7 @@ void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, const std::string& activation); bool cpu_attn_has_isa(const std::string& isa); +bool cpu_has_amx_fp8(); torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, @@ -611,7 +625,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, " "Tensor w2, Tensor topk_weights, Tensor topk_ids, " "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " - "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, " + "Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, " "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, " "bool is_vnni) -> " "()"); @@ -629,7 +643,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool is_vnni) -> Tensor"); ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); - // Adapted from sglang: casual_conv1d kernels + // Adapted from sglang: FP8 W8A8 kernels + ops.def( + "float8_linear_prepack_cpu(Tensor weight, Tensor scales) -> (Tensor, " + "Tensor)"); + ops.impl("float8_linear_prepack_cpu", torch::kCPU, + &float8_linear_prepack_impl); + ops.def( + "fp8_scaled_mm_with_quant(Tensor act, Tensor? act_scales, bool " + "channelwise, " + "Tensor weight, Tensor weight_scales, Tensor? bias, ScalarType " + "output_dtype) -> Tensor"); + ops.impl("fp8_scaled_mm_with_quant", torch::kCPU, &fp8_scaled_mm_with_quant); + ops.def( + "quantize_fp8e4m3_vec(Tensor input, bool channelwise, Tensor? scale_opt) " + "-> (Tensor, Tensor)"); + ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec); + ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor"); ops.impl("causal_conv1d_weight_pack", torch::kCPU, &causal_conv1d_weight_pack); @@ -832,6 +862,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "-> (Tensor, Tensor)"); ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); + // Pure runtime CPU ISA capability check. + ops.def("cpu_has_amx_fp8() -> bool", &cpu_has_amx_fp8); + // CPU attention kernels ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa); ops.def( diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 99cae4534be7..e89fbb55a089 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -1114,5 +1114,203 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace): ) +# =========================================================================== +# FP8 W8A8 MoE +# =========================================================================== + + +requires_cpu_fp8_w8a8 = pytest.mark.skipif( + not ops.cpu_has_amx_fp8(), + reason="requires native AMX-FP8 hardware", +) + +FP8_W8A8_MAX = torch.finfo(torch.float8_e4m3fn).max +FP8_W8A8_QUANT_GROUP = 128 # quantization group size for both K and N dimensions + + +def _make_fp8_w8a8_weight_w13( + E: int, two_n: int, K: int, group_K: int = FP8_W8A8_QUANT_GROUP +) -> tuple[torch.Tensor, torch.Tensor]: + """Create w13 [E, 2N, K] FP8 with group scales [E, 2N, G] where G = K//group_K.""" + G = K // group_K + w_list, s_list = [], [] + for _ in range(E): + w_f32 = torch.randn(two_n, K) + # Per-row, per-K-group quantization: scale [2N, G] + w_re = w_f32.view(two_n, G, group_K) + abs_max = w_re.abs().amax(dim=2, keepdim=True).clamp(min=1e-7) # [2N, G, 1] + scale = (abs_max / FP8_W8A8_MAX).squeeze(2) # [2N, G] + w_q = ( + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn) + ) + w_list.append(w_q.view(two_n, K).contiguous()) + s_list.append(scale.float()) + return torch.stack(w_list), torch.stack(s_list) # [E, 2N, K], [E, 2N, G] + + +def _make_fp8_w8a8_weight_w2( + E: int, + K: int, + N: int, + group_K: int = FP8_W8A8_QUANT_GROUP, + group_N: int = FP8_W8A8_QUANT_GROUP, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create w2 [E, K, N] FP8 with block scales [E, K//gK, N//gN].""" + nK = K // group_K + nN = N // group_N + w_list, s_list = [], [] + for _ in range(E): + w_f32 = torch.randn(K, N) + w_re = w_f32.view(nK, group_K, nN, group_N) + abs_max = w_re.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-7) + scale = (abs_max / FP8_W8A8_MAX).squeeze(1).squeeze(2) # [nK, nN] + w_q = ( + (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn) + ) + w_list.append(w_q.view(K, N).contiguous()) + s_list.append(scale.float()) + return torch.stack(w_list), torch.stack(s_list) # [E, K, N], [E, K//gK, N//gN] + + +def _pack_fp8_w8a8_w13( + w13: torch.Tensor, # [E, 2N, K] FP8 + w13_scale: torch.Tensor, # [E, 2N, G] float32 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack w13 for CPU FP8 W8A8 MoE kernel using float8_linear_prepack_cpu.""" + E = w13.size(0) + packed_list, scale_list = [], [] + for i in range(E): + pw, ps = torch.ops._C.float8_linear_prepack_cpu( + w13[i].contiguous(), w13_scale[i].contiguous() + ) + packed_list.append(pw) + scale_list.append(ps) + return torch.stack(packed_list), torch.stack(scale_list) + + +def _ref_fp8_w8a8_moe( + hidden_states: torch.Tensor, # BF16 [M, K_hidden] (after FP8 round-trip dequant) + w13: torch.Tensor, # FP8 [E, 2N, K_hidden] + w13_scale: torch.Tensor, # float32 [E, 2N, G] (G = K_hidden//group_K) + w2: torch.Tensor, # FP8 [E, K_hidden, N] + w2_scale: torch.Tensor, # float32 [E, K_hidden//gK, N//gN] + topk_weights: torch.Tensor, # float32 [M, top_k] + topk_ids: torch.Tensor, # int64 [M, top_k] + N: int, + group_K: int = FP8_W8A8_QUANT_GROUP, + group_N: int = FP8_W8A8_QUANT_GROUP, +) -> torch.Tensor: + """Reference implementation: dequantize weights then compute MoE.""" + M, K_hidden = hidden_states.shape + top_k = topk_ids.shape[1] + output = torch.zeros(M, K_hidden, dtype=torch.bfloat16) + + for tok in range(M): + x = hidden_states[tok].float() # [K_hidden] + for k_idx in range(top_k): + expert_id = topk_ids[tok, k_idx].item() + w = topk_weights[tok, k_idx].item() + + # Stage 1: x @ w13.T — gate and up projections + w13_e = w13[expert_id].float() # [2N, K_hidden] + ws_e = w13_scale[expert_id].float() # [2N, G] + G = ws_e.shape[1] + w13_dq = torch.zeros_like(w13_e) + for g in range(G): + c0, c1 = g * group_K, (g + 1) * group_K + w13_dq[:, c0:c1] = w13_e[:, c0:c1] * ws_e[:, g : g + 1] + gate_up = (x.unsqueeze(0) @ w13_dq.T).squeeze(0) # [2N] + + gate = gate_up[:N] + up = gate_up[N:] + act = F.silu(gate) * up # [N] + + # Stage 2: act @ w2.T — down projection + w2_e = w2[expert_id].float() # [K_hidden, N] + ws2_e = w2_scale[expert_id].float() # [K_hidden//gK, N//gN] + nK2 = K_hidden // group_K + nN2 = N // group_N + w2_dq = torch.zeros_like(w2_e) + for gi in range(nK2): + for gj in range(nN2): + r0, r1 = gi * group_K, (gi + 1) * group_K + c0, c1 = gj * group_N, (gj + 1) * group_N + w2_dq[r0:r1, c0:c1] = w2_e[r0:r1, c0:c1] * ws2_e[gi, gj] + out_e = (act.unsqueeze(0) @ w2_dq.T).squeeze(0) # [K_hidden] + + output[tok] += (w * out_e).to(torch.bfloat16) + + return output + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize( + "E,N,K,M,top_k", + [ + # N and K must be multiples of FP8_W8A8_QUANT_GROUP=128 + (4, 128, 256, 1, 2), # single token (decode phase) + (4, 128, 256, 8, 2), + (8, 256, 512, 16, 2), + (4, 256, 512, 4, 1), + ], +) +def test_fp8_w8a8_cpu_fused_moe(E: int, N: int, K: int, M: int, top_k: int): + """Test fused_experts_cpu FP8_W8A8 shape/dtype and accuracy vs a + dequantized BF16 reference.""" + set_random_seed(42) + + w13, w13_scale = _make_fp8_w8a8_weight_w13(E, 2 * N, K) + w2, w2_scale = _make_fp8_w8a8_weight_w2(E, K, N) + + packed_w13, packed_w13_scale = _pack_fp8_w8a8_w13(w13, w13_scale) + packed_w2 = _prepack_experts(w2) + + hidden_states = torch.randn(M, K, dtype=torch.bfloat16) * 0.1 + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(hidden_states, True, None) + assert x_fp8.dtype == torch.float8_e4m3fn + assert x_scales.shape == (M,) + + gen = torch.Generator().manual_seed(0) + topk_weights = torch.softmax(torch.randn(M, top_k, generator=gen), dim=-1).float() + topk_ids = torch.zeros(M, top_k, dtype=torch.int32) + for i in range(M): + topk_ids[i] = torch.randperm(E, generator=torch.Generator().manual_seed(i))[ + :top_k + ].int() + + output = torch.empty(M, K, dtype=torch.bfloat16) + ops.fused_experts_cpu( + output, + x_fp8, + packed_w13, + packed_w2, + topk_weights, + topk_ids, + ops.CPUQuantMethod.FP8_W8A8, + packed_w13_scale, # w1_scale: [E, Nc, G, BLOCK_N] + w2_scale, # w2_scale: [E, K//128, N//128] + None, # w1_zero + None, # w2_zero + [FP8_W8A8_QUANT_GROUP, FP8_W8A8_QUANT_GROUP], # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni (weights already packed) + x_scales, # a1_scale: per-token FP8 activation scales + ) + + assert output.shape == (M, K) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + + x_dequant = (x_fp8.float() * x_scales.unsqueeze(1)).bfloat16() + ref = _ref_fp8_w8a8_moe( + x_dequant, w13, w13_scale, w2, w2_scale, topk_weights, topk_ids.long(), N + ) + torch.testing.assert_close(output.float(), ref.float(), atol=1e-2, rtol=1e-2) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index e4dadb53bd66..0b0d12a66c85 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu). +"""Tests for CPU FP8 scaled-mm GEMM kernels. + +Covers: + * fp8_scaled_mm_cpu (W8A16) + * fp8_scaled_mm_with_quant (W8A8) Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`. """ @@ -23,19 +27,31 @@ if not current_platform.is_cpu(): pytest.skip("skipping CPU-only tests", allow_module_level=True) -if not ops._supports_cpu_fp8_w8a16: - pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True) +requires_cpu_fp8_w8a16 = pytest.mark.skipif( + not ops._supports_cpu_fp8_w8a16, reason="fp8_scaled_mm_cpu op not available" +) +requires_cpu_fp8_w8a8 = pytest.mark.skipif( + not ops.cpu_has_amx_fp8(), + reason="requires native AMX-FP8 hardware", +) if not torch.cpu._is_amx_tile_supported(): pytest.skip("requires AMX tile support", allow_module_level=True) +FP8_MAX = torch.finfo(torch.float8_e4m3fn).max BLOCK_SIZE = [128, 128] +GROUP_SIZE = 128 def cdiv(a: int, b: int) -> int: return -(a // -b) +# --------------------------------------------------------------------------- +# W8A16 helpers +# --------------------------------------------------------------------------- + + def quantize_weight_block_fp8( weight: torch.Tensor, block_size: list[int], @@ -126,7 +142,7 @@ def ref_fp8_block_scaled_mm( # --------------------------------------------------------------------------- -# Test parameters +# W8A16 test: fp8_scaled_mm_cpu # --------------------------------------------------------------------------- M_SIZES = [1, 4, 16, 64, 128] # (N, K) — weight shape is [N, K], output has N columns. @@ -141,6 +157,7 @@ def ref_fp8_block_scaled_mm( ] +@requires_cpu_fp8_w8a16 @pytest.mark.parametrize("M", M_SIZES) @pytest.mark.parametrize("N,K", NK_SIZES) @pytest.mark.parametrize("use_bias", [False, True]) @@ -200,6 +217,77 @@ def ref_fp8_per_tensor_scaled_mm( return out.to(out_dtype) +# --------------------------------------------------------------------------- +# W8A8 helpers +# --------------------------------------------------------------------------- + + +def quantize_weight_per_tensor( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with a single per-tensor scale, broadcast + to [N, 1] (the shape float8_linear_prepack_cpu expects).""" + N = weight.shape[0] + abs_max = weight.abs().max() + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q, scale.expand(N).contiguous().view(N, 1) + + +def quantize_weight_per_channel( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with per-channel (per-row) scales [N, 1].""" + abs_max = weight.abs().amax(dim=1, keepdim=True) + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q, scale.float() + + +def quantize_weight_per_group( + weight: torch.Tensor, + group_size: int = GROUP_SIZE, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with per-group (128-block) scales [N, G].""" + N, K = weight.shape + G = K // group_size + w_groups = weight.view(N, G, group_size) + abs_max = w_groups.abs().amax(dim=2, keepdim=True) + scale = (abs_max / FP8_MAX).clamp(min=1e-7) + q = (w_groups.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q.view(N, K), scale.squeeze(2).float() + + +def quantize_act_per_token(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Reference per-token (row) dynamic FP8 quantization of activations.""" + x_f = x.float() + scale = (x_f.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-7) + scaled = (x_f / scale).clamp(-FP8_MAX, FP8_MAX) + # Round via fp16 to match the AVX10.2 kernel path. + q = scaled.half().to(torch.float8_e4m3fn) + return q, scale + + +def reference_static_fp8_linear( + x: torch.Tensor, # BF16 [M, K] + weight_fp8: torch.Tensor, # FP8 [N, K] + weight_scale: torch.Tensor, # float32 [N, 1] + act_scale: torch.Tensor, # float32 scalar + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Static per-tensor: quantize/dequant x with act_scale, dequant weight, matmul.""" + x_scaled = (x.float() / act_scale.item()).clamp(-FP8_MAX, FP8_MAX) + # Round via fp16 to match the AVX10.2 kernel path. + x_q = x_scaled.half().to(torch.float8_e4m3fn) + x_dq = x_q.float() * act_scale.item() + w_dq = weight_fp8.float() * weight_scale.view(-1, 1) + out = torch.mm(x_dq, w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + NK_SIZES_PER_TENSOR = [ (32, 64), (5120, 5120), @@ -274,3 +362,136 @@ def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int): ) supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config) assert not supported + + +def reference_dynamic_fp8_linear( + x: torch.Tensor, # BF16 [M, K] + weight_fp8: torch.Tensor, # FP8 [N, K] + weight_scale: torch.Tensor, # float32 [N, G] + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Dynamic per-token act x per-group weight: quantize x per-token, + dequant weight per-group, matmul.""" + N, K = weight_fp8.shape + G = weight_scale.shape[1] + group_size = K // G + x_q, x_scale = quantize_act_per_token(x) + x_dq = x_q.float() * x_scale + w_dq = ( + weight_fp8.float().view(N, G, group_size) * weight_scale.view(N, G, 1) + ).view(N, K) + out = torch.mm(x_dq, w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +# --------------------------------------------------------------------------- +# W8A8 tests: fp8_scaled_mm_with_quant +# --------------------------------------------------------------------------- +W8A8_M_SIZES = [1, 4, 16, 64] +W8A8_NK_SIZES = [(128, 256), (256, 512), (64, 128)] + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("static_activation", [True, False]) +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_tensor( + M: int, N: int, K: int, use_bias: bool, static_activation: bool +): + """CPUFP8W8A8ScaledMMLinearKernel with per-tensor weight scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_tensor(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + + if static_activation: + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_static_fp8_linear( + x, w_fp8, w_scale, act_scale, bias, out_dtype + ) + atol = 0.05 + else: + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + atol = 0.05 + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol) + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("static_activation", [True, False]) +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_channel( + M: int, N: int, K: int, use_bias: bool, static_activation: bool +): + """CPUFP8W8A8ScaledMMLinearKernel with per-channel weight scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_channel(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + + if static_activation: + act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, act_scale, False, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_static_fp8_linear( + x, w_fp8, w_scale, act_scale, bias, out_dtype + ) + atol = 0.05 + else: + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + atol = 0.05 + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol) + + +@requires_cpu_fp8_w8a8 +@pytest.mark.parametrize("M", W8A8_M_SIZES) +@pytest.mark.parametrize("N,K", W8A8_NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_w8a8_per_group(M: int, N: int, K: int, use_bias: bool): + """CPUFp8W8A8BlockScaledMMKernel: per-group (128-block) weight scale + + dynamic per-token activation scale.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=torch.bfloat16) + weight = torch.randn(N, K, dtype=torch.bfloat16) + bias = torch.randn(N, dtype=torch.float32) if use_bias else None + + w_fp8, w_scale = quantize_weight_per_group(weight) + packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale) + kernel_out = ops.fp8_scaled_mm_with_quant( + x, None, True, packed_weight, packed_scale, bias, out_dtype + ) + ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype) + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=0.05) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 96f47e2069f1..2a1654231fc2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import IntEnum +from functools import cache from typing import TYPE_CHECKING, Literal import torch @@ -3384,6 +3385,7 @@ class CPUQuantMethod(IntEnum): FP8_W8A16 = 2 INT4_W4A8 = 3 MXFP4 = 4 + FP8_W8A8 = 5 if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): @@ -3422,6 +3424,7 @@ def fused_experts_cpu( alpha: float | None = None, limit: float | None = None, is_vnni: bool = True, + a1_scale: torch.Tensor | None = None, ) -> None: torch.ops._C.fused_experts_cpu( out, @@ -3435,6 +3438,7 @@ def fused_experts_cpu( w2_scale, w1_zero, w2_zero, + a1_scale, block_size, w1_bias, w2_bias, @@ -3564,6 +3568,46 @@ def fp8_scaled_mm_cpu( ) +# FP8 W8A8 CPU kernels +@cache +def cpu_has_amx_fp8() -> bool: + """Whether this CPU has native AMX-FP8 MMA support.""" + if not hasattr(torch.ops._C, "cpu_has_amx_fp8"): + return False + return bool(torch.ops._C.cpu_has_amx_fp8()) + + +if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"): + + @register_fake("_C::fp8_scaled_mm_with_quant") + def fp8_scaled_mm_with_quant_fake( + act: torch.Tensor, + act_scales: torch.Tensor | None, + channelwise: bool, + weight: torch.Tensor, + weight_scales: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, + ) -> torch.Tensor: + M = act.reshape(-1, act.size(-1)).size(0) + N = weight.size(0) * weight.size(-1) + return torch.empty((M, N), dtype=output_dtype, device=act.device) + + +def fp8_scaled_mm_with_quant( + act: torch.Tensor, + act_scales: torch.Tensor | None, + channelwise: bool, + weight: torch.Tensor, + weight_scales: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, +) -> torch.Tensor: + return torch.ops._C.fp8_scaled_mm_with_quant( + act, act_scales, channelwise, weight, weight_scales, bias, output_dtype + ) + + def chunk_gated_delta_rule_cpu( query: torch.Tensor, key: torch.Tensor, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 8c57059e5727..fb7a170b82fb 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -186,6 +186,8 @@ from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUFp8PerTensorScaledMMLinearKernel, + CPUFp8W8A8BlockScaledMMKernel, + CPUFP8W8A8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -436,6 +438,7 @@ def _resolve_backend_kernels( ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.CPU: [ + CPUFP8W8A8ScaledMMLinearKernel, CPUFp8PerTensorScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, @@ -468,6 +471,7 @@ def _resolve_backend_kernels( TritonFp8BlockScaledMMKernel, ], PlatformEnum.CPU: [ + CPUFp8W8A8BlockScaledMMKernel, # W8A8 preferred; falls back to W8A16 below CPUFp8BlockScaledMMKernel, ], PlatformEnum.XPU: [ diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index f4a062b7f2b3..75db4bcaa9ed 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -10,6 +10,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUFp8PerTensorScaledMMLinearKernel, + CPUFP8W8A8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -69,5 +70,6 @@ "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", "CPUFp8PerTensorScaledMMLinearKernel", + "CPUFP8W8A8ScaledMMLinearKernel", "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index e72dffcbbc50..1d2ede1853c8 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -270,12 +270,17 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Skip the base class process (FP8 padding / fnuz normalization) # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. params = self._get_layer_params(layer) - packed_weight = torch.ops._C.convert_weight_packed(params.weight) - replace_parameter( - layer, - params.WEIGHT, - torch.nn.Parameter(packed_weight, requires_grad=False), - ) + + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and + # let the kernel pack on the fly instead. + if not getattr(layer, "_cpu_skip_gemm_dispatch", False): + packed_weight = torch.ops._C.convert_weight_packed(params.weight) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) # Re-wrap scale as a plain Parameter so the kernel can read it # without weight-loader metadata interfering. @@ -320,7 +325,7 @@ def apply_weights( list(self.weight_group_shape), bias, x.dtype, - True, # is_vnni (weight already prepacked) + not getattr(layer, "_cpu_skip_gemm_dispatch", False), # is_vnni ) return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out @@ -430,3 +435,294 @@ def apply_scaled_mm( raise NotImplementedError( "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly." ) + + +class CPUFp8W8A8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """FP8 W8A8 block-quantized GEMM with dynamic per-token activation quant on CPU.""" + + # Activation quant is done inside fp8_scaled_mm_with_quant; + # the base class must not apply a second quant step. + apply_input_quant = False + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops.cpu_has_amx_fp8(): + return False, ( + "requires native AMX-FP8 hardware; falls back " + "to CPUFp8BlockScaledMMKernel (W8A16) otherwise." + ) + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + # Same block-shape requirements as the W8A16 sibling kernel. + weight_gs = config.weight_quant_key.scale.group_shape + if weight_gs.col <= 0 or weight_gs.col != 128: + return False, ( + "CPU FP8 W8A8 block kernel requires K-dimension block size of 128, " + f"got {weight_gs.col}." + ) + if weight_gs.row <= 0 or weight_gs.row % 32 != 0: + return False, ( + "CPU FP8 W8A8 block kernel requires N-dimension block size to be " + f"a positive multiple of 32, got {weight_gs.row}." + ) + if config.out_dtype not in (torch.bfloat16, torch.float32): + return False, "Only bfloat16/float32 output dtype supported." + n = config.weight_shape[0] + if n % 32 != 0: + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepack block-quantized FP8 weight for W8A8 AMX BRGEMM. + + For block-quant FP8 models weight is stored as [N, K] (NOT transposed, + unlike non-block FP8 where vLLM stores [K, N]). + + float8_linear_prepack_cpu expects: + weight: [N, K] + scales: [N, G] where G = K / block_k + + weight_scale_inv from the checkpoint has shape [ceil(N/block_n), G], + so we expand rows from ceil(N/block_n) to N via repeat_interleave. + """ + params = self._get_layer_params(layer) + weight = params.weight # [N, K] + + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this + # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and + # its block scales; apply_weights falls back to the W8A16 op. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + return + + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) # [ceil(N/block_n), ceil(K/block_k)] + assert weight_scale is not None, "FP8 W8A8 block kernel requires weight scale." + + N = weight.shape[0] + block_n = self.weight_group_shape.row # e.g. 128 + + # Expand scale rows: [ceil(N/block_n), G] → [N, G] + if weight_scale.size(0) != N: + weight_scale_expanded = weight_scale.repeat_interleave(block_n, dim=0)[ + :N, : + ].contiguous() + else: + weight_scale_expanded = weight_scale.contiguous() + + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu( + weight, weight_scale_expanded.to(torch.float32) + ) + + scale_attr = ( + params.WEIGHT_SCALE_INV + if params.weight_scale_inv is not None + else params.WEIGHT_SCALE + ) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) + replace_parameter( + layer, + scale_attr, + torch.nn.Parameter(packed_scales, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + params = self._get_layer_params(layer) + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + # Unpacked weight: fp8_scaled_mm_with_quant has no is_vnni option, + # so fall back to the W8A16 op (BF16 activation, block-scaled FP8 + # weight) for this layer. + out = torch.ops._C.fp8_scaled_mm_cpu( + x_2d, + params.weight, + weight_scale, + list(self.weight_group_shape), + bias, + x.dtype, + False, # is_vnni + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + # act_scales=None + channelwise=True → dynamic per-token activation quant + out = ops.fp8_scaled_mm_with_quant( + x_2d, + None, + True, + params.weight, + weight_scale, + bias, + x.dtype, + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "CPUFp8W8A8BlockScaledMMKernel overrides apply_weights directly." + ) + + +class CPUFP8W8A8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """FP8 W8A8 GEMM with dynamic per-token activation quantization on CPU.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops.cpu_has_amx_fp8(): + return False, ( + "requires native AMX-FP8 hardware (Diamond Rapids); falls back " + "to the Torch FP8 kernels otherwise." + ) + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + if config.out_dtype not in (torch.bfloat16, torch.float16, torch.float32): + return False, "Only bfloat16/float16/float32 output dtype supported." + # Reject block-quantized weights — those go to CPUFp8BlockScaledMMKernel + weight_gs = config.weight_quant_key.scale.group_shape + if weight_gs.col > 0 and weight_gs.row > 0: + return False, "Block-quantized weights not handled by FP8 W8A8 kernel." + n = config.weight_shape[0] + if n % 32 != 0: + # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepack weights with float8_linear_prepack_cpu (VNNI + block layout).""" + w_name, ws_name, xs_name, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, ws_name) + + # vLLM stores linear weight as [K, N], while the CPU prepack op expects [N, K]. + N = weight.shape[-1] + + # Ensure scale is 2D [N, G] + ws = weight_scale + if ws is None: + raise RuntimeError("FP8 W8A8 kernel requires weight scale.") + ws = ws.to(torch.float32) + + # Fused modules may carry per-shard scales (e.g. QKV/MLP shards). + # Expand those to per-output-channel scales before prepack. + if ( + len(getattr(layer, "logical_widths", [])) > 1 + and ws.dim() <= 1 + and ws.numel() == len(layer.logical_widths) + ): + ws = convert_to_channelwise(ws.reshape(-1), layer.logical_widths) + + if ws.dim() == 0: + ws = ws.view(1, 1).expand(N, 1) + elif ws.dim() == 1: + if ws.numel() == 1: + ws = ws.view(1, 1).expand(N, 1) + elif ws.numel() == N: + ws = ws.unsqueeze(1) # [N] -> [N, 1] + else: + raise RuntimeError( + "FP8 W8A8 kernel got unsupported 1D weight scale shape: " + f"{tuple(ws.shape)} for N={N}." + ) + + # Align scale rows to output channels N. + if ws.size(0) != N: + if ws.size(0) == 1: + ws = ws.expand(N, ws.size(1)).contiguous() + else: + raise RuntimeError( + "FP8 W8A8 kernel got mismatched weight scale rows: " + f"ws_rows={ws.size(0)}, expected N={N}." + ) + + w_2d = weight.t().contiguous() # [N, K] + packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(w_2d, ws) + + replace_parameter( + layer, + w_name, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) + replace_parameter( + layer, + ws_name, + torch.nn.Parameter(packed_scales, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + w_name, ws_name, xs_name, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, ws_name) + # For static quantization, use pre-computed activation scale + act_scales = getattr(layer, xs_name, None) if xs_name else None + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + # channelwise=True: per-token (per-row) quantization when act_scales is None + # channelwise=False: per-tensor quantization when act_scales is provided + channelwise = act_scales is None + out = ops.fp8_scaled_mm_with_quant( + x_2d, + act_scales, + channelwise, + weight, + weight_scale, + bias, + x.dtype, + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_scaled_mm(self, A, B, out_dtype, As, Bs, bias=None, **kwargs): + raise NotImplementedError( + "CPUFP8W8A8ScaledMMLinearKernel uses apply_weights directly." + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 26547c73533d..3104ac25c6ef 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1055,7 +1055,7 @@ def forward_impl( mqa_q_nope, self.impl._w_uk_packed, # type: ignore[attr-defined] True, - None, + self.impl._w_scale, # type: ignore[attr-defined] ) mqa_ql_nope = mqa_ql_nope.transpose(0, 1) else: @@ -1401,7 +1401,7 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): x, self.impl._w_uv_packed, # type: ignore[attr-defined] True, - None, + self.impl._w_scale, # type: ignore[attr-defined] ) else: # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 858b35186448..3e7376fe9015 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -15,6 +15,7 @@ convert_weight_packed_scale_zp, cpu_fused_moe, cpu_fused_moe_int8, + cpu_has_amx_fp8, cpu_prepack_moe_weight, cpu_prepack_moe_weight_int8, fused_experts_cpu, @@ -467,6 +468,7 @@ def _supports_routing_method( RoutingMethodType.Default, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepSeekV3, RoutingMethodType.DeepseekV4, RoutingMethodType.DeepSeekV3, ] @@ -1612,3 +1614,218 @@ def apply( self._w1_scale_bf16, self._w2_scale_bf16, ) + + +# =========================================================================== +# FP8 W8A8 MoE +# =========================================================================== + + +def prepare_fp8_w8a8_moe_layer_for_cpu( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepack FP8 W8A8 MoE weights for CPU kernel.""" + E = w13.size(0) + + packed_w13_list = [] + packed_w13_scale_list = [] + + for i in range(E): + w1_exp = w13[i] # [2N, K] FP8 + ws_exp = w13_scale[i] # [2N, K/128] or [2N_blocks, K/128] + + # Align scale rows to weight rows if necessary (for block-128 checkpoints + # where scale may have fewer rows: [2N/128, K/128] → [2N, K/128]) + if ws_exp.size(0) < w1_exp.size(0): + repeat_factor = w1_exp.size(0) // ws_exp.size(0) + ws_exp = torch.repeat_interleave(ws_exp, repeat_factor, dim=0) + ws_exp = ws_exp[: w1_exp.size(0), :].contiguous() + + pw, ps = torch.ops._C.float8_linear_prepack_cpu(w1_exp, ws_exp) + packed_w13_list.append(pw) + packed_w13_scale_list.append(ps) + + packed_w13 = torch.stack(packed_w13_list) + packed_w13_scale = torch.stack(packed_w13_scale_list) + + # w2 uses the W8A16 path (BF16 activation × FP8 weight), VNNI packed + packed_w2 = torch.ops._C.convert_weight_packed(w2) + return packed_w13, packed_w13_scale, packed_w2, w2_scale + + +class CPUExpertsFp8W8A8(mk.FusedMoEExpertsModular): + """CPU FP8 W8A8 block-quantized modular MoE experts.""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + @property + def expects_unquantized_inputs(self) -> bool: + # Receives BF16 hidden states; quantization to FP8 is done inside apply() + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + # Requires native AMX-FP8 hardware. + return current_platform.is_cpu() and cpu_has_amx_fp8() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # Block-quantized FP8. + return (weight_key, activation_key) == ( + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepSeekV3, + RoutingMethodType.DeepseekV4, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # float8_linear_prepack_cpu blocks w1 into [E, Nc, Kc, block_k, block_n], + # not the (E, N, K) layout the base implementation assumes -- N isn't + # recoverable from that shape, so read it from moe_config instead. + E = w1.shape[0] + K = a1.size(-1) + N = ( + self.moe_config.intermediate_size_per_partition + * self.moe_config.w13_num_shards + ) + M = a1.size(0) if a1.dim() == 2 else a1.size(1) + topk = topk_ids.size(1) + return E, M, N, K, topk + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + if apply_router_weight_on_input: + # fused_experts_cpu always applies topk_weights internally on + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would + # also pre-apply it to hidden_states, double-weighting the + # output. Not needed by any CPU FP8 model today. + raise NotImplementedError( + "CPUExpertsFp8W8A8 does not support apply_router_weight_on_input=True." + ) + + block_shape = ( + list(self.quant_config.block_shape) + if self.quant_config.block_shape + else ( + [ + self.quant_config._w1.shape.row, + self.quant_config._w1.shape.col, + ] + if self.quant_config._w1.shape is not None + else None + ) + ) + + # Quantize hidden_states (BF16) → FP8 with per-token scales + hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1]) + x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec( + hidden_states_2d, True, None + ) + + fused_experts_cpu( + output, + x_fp8, + w1, + w2, + topk_weights, + topk_ids, + CPUQuantMethod.FP8_W8A8, + self.w1_scale, # w1_scale + self.w2_scale, # w2_scale + None, # w1_zero + None, # w2_zero + block_shape, # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni + x_scales, # a1_scale: per-token FP8 activation scales + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 6a6ce34d160d..abf089620408 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -57,7 +57,9 @@ class Fp8MoeBackend(Enum): VLLM_CUTLASS = "VLLM_CUTLASS" BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS" XPU = "XPU" + # CPU FP8 W8A16 (BF16 activations); CPU_W8A8 needs native AMX-FP8. CPU = "CPU" + CPU_W8A8 = "CPU_W8A8" HPC = "HPC" # Dequantize-to-BF16 emulation for MXFP8 on devices without a native # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here. @@ -92,6 +94,7 @@ def _get_priority_backends( Fp8MoeBackend.BATCHED_VLLM_CUTLASS, Fp8MoeBackend.BATCHED_TRITON, Fp8MoeBackend.XPU, + Fp8MoeBackend.CPU_W8A8, Fp8MoeBackend.CPU, Fp8MoeBackend.HPC, ] @@ -129,8 +132,10 @@ def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> Non _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU) if current_platform.is_cpu(): - # CPU platform uses FP8 W8A16 fused MoE kernel. + # W8A8 first: it falls through to the W8A16 backend whenever the + # hardware (AMX-FP8) or the config isn't supported. _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU) + _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU_W8A8) return _AVAILABLE_BACKENDS @@ -238,6 +243,13 @@ def backend_to_kernel_cls( return [CPUExpertsFp8] + elif backend == Fp8MoeBackend.CPU_W8A8: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsFp8W8A8, + ) + + return [CPUExpertsFp8W8A8] + elif backend == Fp8MoeBackend.HPC: from vllm.model_executor.layers.fused_moe.hpc_moe import ( HPCExperts, @@ -652,6 +664,14 @@ def convert_to_fp8_moe_kernel_format( ) w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2) + elif fp8_backend == Fp8MoeBackend.CPU_W8A8: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + prepare_fp8_w8a8_moe_layer_for_cpu, + ) + + w13, w13_scale, w2, w2_scale = prepare_fp8_w8a8_moe_layer_for_cpu( + w13, w2, w13_scale, w2_scale + ) else: if fp8_backend not in [ Fp8MoeBackend.TRITON, @@ -697,7 +717,15 @@ def make_fp8_moe_quant_config( In a future PR, we will have this function should be a method of the modular kernel itself. """ - # MARLIN and CPU are mixed precision W8A16 config. + if fp8_backend == Fp8MoeBackend.CPU_W8A8: + return fp8_w8a8_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + block_shape=block_shape, + ) + + # MARLIN and CPU (W8A16) are mixed precision W8A16 configs. if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU: return fp8_w8a16_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 2a8ef0b69f0e..ed018c81bf0a 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -737,11 +737,13 @@ def __init__( ) # Reject backends whose make_fp8_moe_quant_config branch silently # drops per_act_token_quant / per_out_ch_quant or collapses scales: - # MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_* - # fold scales into a per-tensor alpha (oracle/fp8.py). + # MARLIN / CPU route through fp8_w8a16_moe_quant_config, CPU_W8A8 + # ignores both flags, and FLASHINFER_* fold scales into a per-tensor + # alpha (oracle/fp8.py). if self.fp8_backend in ( Fp8MoeBackend.MARLIN, Fp8MoeBackend.CPU, + Fp8MoeBackend.CPU_W8A8, Fp8MoeBackend.FLASHINFER_CUTLASS, Fp8MoeBackend.FLASHINFER_TRTLLM, ): diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 5f9ac7f17adf..38c8b68af9f7 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -61,6 +61,44 @@ _SPLIT_OCCUPANCY_MULTIPLIER = 2 +def _block_quant_to_tensor_quant( + weight: torch.Tensor, # FP8 [N, K] block-quantized weight + weight_scale: torch.Tensor, # [ceil(N/bn), ceil(K/bk)] block scales + block_size: list[int], # [block_n, block_k] +) -> tuple[torch.Tensor, torch.Tensor]: + """Re-quantize FP8 block-quant weight to FP8 per-tensor quant. + + Returns (fp8_weight, dequant_scale) where fp8_weight * dequant_scale + approximates the original float weight. dequant_scale is a 0-dim + scalar tensor as expected by bmm_cpu's ``scale`` argument. + + NOTE: block->tensor re-quantization may slightly affect accuracy + (equivalent to SGLang's block_quant_to_tensor_quant for CPU). + """ + block_n, block_k = block_size[0], block_size[1] + n, k = weight.shape + + # Expand block scales [ceil(N/bn), ceil(K/bk)] -> [N, K] + scale = weight_scale.to(torch.float32) + if scale.shape[0] < n: + scale = scale.repeat_interleave(block_n, dim=0)[:n, :] + if scale.shape[1] < k: + scale = scale.repeat_interleave(block_k, dim=1)[:, :k] + + # Dequantize to float32 + x_f32 = weight.to(torch.float32) * scale + + # Re-quantize per-tensor back to FP8 + fp8_max = torch.finfo(weight.dtype).max # 448.0 for e4m3fn + amax = x_f32.abs().amax().clamp(min=1e-12) + quant_scale = fp8_max / amax + x_fp8 = (x_f32 * quant_scale).clamp(-fp8_max, fp8_max).to(weight.dtype) + + # dequant_scale = 1 / quant_scale (0-dim tensor for bmm_cpu) + dequant_scale = (amax / fp8_max).to(torch.float32) + return x_fp8, dequant_scale + + def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread count in place of SM count.""" @@ -208,6 +246,9 @@ def __init__( ) self._w_uk_packed: torch.Tensor | None = None self._w_uv_packed: torch.Tensor | None = None + self._w_scale: torch.Tensor | None = ( + None # per-tensor dequant scale (FP8 W8A16) + ) def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: # forward_mha (unlike forward_mqa) receives no `layer` argument, so it @@ -216,12 +257,40 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: # impl derives its own copies from self.kv_b_proj (already available # via the constructor) and VNNI-packs them once, ahead of time, for # use with bmm_cpu in forward_mha. - kv_b_proj_weight = get_and_maybe_dequant_weights( - self.kv_b_proj, out_dtype=act_dtype - ).T - kv_b_proj_weight = kv_b_proj_weight.view( - self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) + # + # kv_b_proj carries `_cpu_skip_gemm_dispatch`, so the CPU FP8 kernels + # leave its weight in the checkpoint [N, K] layout. For block-quant + # FP8 checkpoints, re-quantize per-tensor and keep the BMM in FP8 + # (W8A16); everything else goes through the BF16 path. + weight = getattr(self.kv_b_proj, "weight", None) + block_size = getattr(self.kv_b_proj, "weight_block_size", None) + if ( + weight is not None + and weight.dtype == torch.float8_e4m3fn + and block_size is not None + ): + block_scale = getattr(self.kv_b_proj, "weight_scale_inv", None) + if block_scale is None: + block_scale = self.kv_b_proj.weight_scale + fp8_weight, self._w_scale = _block_quant_to_tensor_quant( + weight, block_scale, block_size + ) + # [N, K] -> [K, N] = [kv_lora_rank, heads * (nope + v)] + kv_b_proj_weight = fp8_weight.T.contiguous().view( + self.kv_lora_rank, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + else: + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + self._w_scale = None w_uk, w_uv = kv_b_proj_weight.split( [self.qk_nope_head_dim, self.v_head_dim], dim=-1 ) @@ -348,7 +417,9 @@ def forward_mha( ql_nope = torch.empty( num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device ) - ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None) + ops.bmm_cpu( + ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, self._w_scale + ) mqa_q = torch.cat([ql_nope, q_pe], dim=-1) kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) @@ -402,4 +473,4 @@ def forward_mha( output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose( 0, 1 ) - ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None) + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, self._w_scale) From 5bb596201a19ba73e08e05e062a7ad64e96ecda2 Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Fri, 18 Sep 2026 03:10:38 -0700 Subject: [PATCH 465/774] [CI] Raise Elastic EP Scaling step timeout 30m -> 40m (#57287) Signed-off-by: kimi-no-na-wa Co-authored-by: kimi-no-na-wa --- .buildkite/test_areas/expert_parallelism.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index f74b1adf2e2a..485d4f3d3124 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -43,7 +43,7 @@ steps: - label: ":nvidia: (H100) Elastic EP Scaling" key: elastic-ep-scaling-test - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: h100 working_dir: "/vllm-workspace/tests" num_devices: 4 From 5397f967b9ee282354f9803af3e758c9270588e9 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 03:11:07 -0700 Subject: [PATCH 466/774] [CI] Add residual timeout headroom after JIT rollback (#56649) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/engine.yaml | 2 +- .buildkite/test_areas/spec_decode.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 179954740a3f..f6e243a1b726 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -81,7 +81,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) E2E Core" device: h200_18gb key: e2e-core-1-gpu - timeout_in_minutes: 50 + timeout_in_minutes: 65 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index fdb12efe1158..f60d745df334 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -174,7 +174,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Spec Decode Draft Model" key: spec-decode-draft-model - timeout_in_minutes: 45 + timeout_in_minutes: 60 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ From 4c6c1a40b3cb4f9018e2a6c79fdbe11621acbda1 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Fri, 18 Sep 2026 03:13:45 -0700 Subject: [PATCH 467/774] [ROCm][Bugfix] Fix intermittent ROCR host segfault (#57328) Signed-off-by: Matthew Wong --- docker/Dockerfile.rocm_base | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index d842a7b07e68..fa0a88fd206e 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -15,7 +15,7 @@ ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350" -ARG ROCM_RUNTIME_COMMIT="b539bf7eebfd99ad0a69668caa1f4037034d501f" +ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521" # Sccache configuration (only used in release pipeline) ARG USE_SCCACHE From f2b6152ec2448b4ea877153d6a613c381e7fd8bb Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 03:40:32 -0700 Subject: [PATCH 468/774] [CI] Add flaky rerun markers / timeout skips to sibling tests lacking them (#57359) Signed-off-by: khluu Co-authored-by: Kimi --- .../pooling/test_token_classification.py | 2 ++ tests/multimodal/media/test_connector.py | 17 +++++++++++++---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index d8d813f404fb..c205a3c42a8a 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -115,6 +115,7 @@ def test_modernbert_models( @pytest.mark.parametrize("model", ["Davlan/xlm-roberta-base-ner-hrl"]) @pytest.mark.parametrize("dtype", ["float"]) +@pytest.mark.flaky(reruns=3) @torch.inference_mode def test_xlm_roberta_models( hf_runner, @@ -197,6 +198,7 @@ def test_openai_privacy_filter( @pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"]) @pytest.mark.parametrize("dtype", ["float"]) +@pytest.mark.flaky(reruns=3) @torch.inference_mode def test_auto_conversion( hf_runner, diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index b222d57d3c24..5949c66b562a 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -299,6 +299,7 @@ async def test_fetch_video_http(video_url: str, num_frames: int): assert metadata_sync == metadata_async +@pytest.mark.flaky(reruns=3, reruns_delay=5) @pytest.mark.asyncio @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) @pytest.mark.parametrize("max_duration", [1, 60, 1800]) @@ -320,8 +321,11 @@ async def test_fetch_video_http_with_dynamic_loader( } ) - video_sync, metadata_sync = connector.fetch_video(video_url) - video_async, metadata_async = await connector.fetch_video_async(video_url) + try: + video_sync, metadata_sync = connector.fetch_video(video_url) + video_async, metadata_async = await connector.fetch_video_async(video_url) + except (TimeoutError, asyncio.TimeoutError) as e: + pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) assert metadata_sync == metadata_async @@ -379,6 +383,7 @@ def test_placeholder_range_extract_embeds_range(offset, is_embed, expected): assert pr.extract_embeds_range() == expected +@pytest.mark.flaky(reruns=3, reruns_delay=5) @pytest.mark.asyncio @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) @pytest.mark.parametrize("num_frames", [-1, 32, 1800]) @@ -394,8 +399,12 @@ async def test_allowed_media_domains(video_url: str, num_frames: int): ], ) - video_sync, metadata_sync = connector.fetch_video(video_url) - video_async, metadata_async = await connector.fetch_video_async(video_url) + try: + video_sync, metadata_sync = connector.fetch_video(video_url) + video_async, metadata_async = await connector.fetch_video_async(video_url) + except (TimeoutError, asyncio.TimeoutError) as e: + pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") + assert np.array_equal(video_sync, video_async) assert metadata_sync == metadata_async From 04a3a00eb79118dda521577caaaba2ba6b65bb33 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 03:40:54 -0700 Subject: [PATCH 469/774] [CI] Raise GSM8K startup max wait to 2400s for flaky 30B+ MoE eval configs (#57361) Signed-off-by: khluu --- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 2 ++ .../Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml | 2 ++ .../Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml | 2 ++ .../Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml | 2 ++ .../Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml | 2 ++ 5 files changed, 10 insertions(+) diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index bc5c11fa53e2..141bf68162a3 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml @@ -3,6 +3,8 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# 30B+ MoE with EP/PCP: download + kernel autotuning can exceed the 1200s default. +startup_max_wait_seconds: 2400 # Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation. server_args: >- --compilation-config '{"cudagraph_mode":"NONE"}' diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml index 5f10684d2e7c..f81eb1c18824 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml @@ -2,6 +2,8 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 +# 30B MoE + DeepGEMM/DeepEP JIT compile can exceed the 1200s default startup wait. +startup_max_wait_seconds: 2400 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log" env: VLLM_USE_DEEP_GEMM: "1" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml index 37e6039e9694..3987065f6ffd 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml @@ -2,6 +2,8 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block" accuracy_threshold: 0.85 num_questions: 1319 num_fewshot: 5 +# 30B MoE + DeepGEMM/DeepEP JIT compile can exceed the 1200s default startup wait. +startup_max_wait_seconds: 2400 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log" env: VLLM_USE_DEEP_GEMM: "1" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml index 6624cea1ef23..004e3e15e302 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml @@ -2,4 +2,6 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 +# 30B MoE + CuTeDSL/DeepEP JIT compile can exceed the 1200s default startup wait. +startup_max_wait_seconds: 2400 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml index f2d4588e3aee..55892f21969b 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml @@ -2,4 +2,6 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 +# 30B MoE + CuTeDSL/DeepEP JIT compile can exceed the 1200s default startup wait. +startup_max_wait_seconds: 2400 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl" From cc09352e5c50ca6b7d2c6d79b29d6d3b18df3046 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 18 Sep 2026 18:44:39 +0800 Subject: [PATCH 470/774] [Bugfix][Multimodal] Preserve DeepSeek V4 image block spacing (#56882) Signed-off-by: Xucheng Zhou Signed-off-by: Isotr0py Co-authored-by: OpenAI Codex Co-authored-by: Isotr0py Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../fixtures/deepseek_v4/test_input_5.json | 39 +++++++++++++++++++ .../fixtures/deepseek_v4/test_output_5.txt | 9 +++++ tests/tokenizers_/test_deepseek_v4.py | 14 +++---- vllm/renderers/deepseek_v4.py | 4 +- vllm/tokenizers/deepseek_v4_encoding.py | 6 +-- 5 files changed, 60 insertions(+), 12 deletions(-) create mode 100644 tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json create mode 100644 tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json b/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json new file mode 100644 index 000000000000..2a2fa6a89e81 --- /dev/null +++ b/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json @@ -0,0 +1,39 @@ +{ + "thinking_mode": "thinking", + "reasoning_effort": "max", + "messages": [ + { + "role": "system", + "content": "You are a helpful vision assistant." + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "请按“第一张、第二张”的顺序回答：第一张图" + }, + { + "type": "image_url", + "image_url": { + "url": "examples/images/carrots.jpeg" + } + }, + { + "type": "text", + "text": "和第二张图" + }, + { + "type": "image_url", + "image_url": { + "url": "examples/images/corn.jpeg" + } + }, + { + "type": "text", + "text": "中分别是什么食材？它们通常食用的部位分别是什么？" + } + ] + } + ] +} diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt b/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt new file mode 100644 index 000000000000..1203dc8db114 --- /dev/null +++ b/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt @@ -0,0 +1,9 @@ +<｜begin▁of▁sentence｜>You are a helpful vision assistant.<｜User｜>请按“第一张、第二张”的顺序回答：第一张图 + +<｜deepseek_image｜> + +和第二张图 + +<｜deepseek_image｜> + +中分别是什么食材？它们通常食用的部位分别是什么？<｜Assistant｜> \ No newline at end of file diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index f7bc5909ba31..4403c2a877d0 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -54,13 +54,9 @@ def _load_reference_case(case_id: int): def _render_reference_case(case_id: int, **kwargs): messages, tools = _load_reference_case(case_id) - conversation, _, _ = parse_chat_messages( - messages, - _model_config(), - content_format="string", - ) + # Preserve content blocks for encoding without fetching fixture media. return _tokenizer().apply_chat_template( - conversation=conversation, + conversation=messages, messages=messages, tools=tools, tokenize=False, @@ -397,6 +393,7 @@ def test_deepseek_v4_maps_xhigh_to_high_reasoning_effort(): (2, {"thinking": True, "reasoning_effort": "low"}), (3, {"thinking": True, "reasoning_effort": "low"}), (4, {"thinking": False}), + (5, {"thinking": False}), ], ) def test_deepseek_v4_matches_reference_golden_fixtures(case_id, kwargs): @@ -447,7 +444,10 @@ def test_deepseek_v4_image_blocks_become_placeholders(): thinking=False, ) - assert "<｜User｜>first:<｜deepseek_image｜>second:<｜deepseek_image｜>" in prompt + assert ( + "<｜User｜>first:\n\n<｜deepseek_image｜>\n\n" + "second:\n\n<｜deepseek_image｜>" in prompt + ) def test_deepseek_v4_image_sentinel_ids_match_tokenizer(): diff --git a/vllm/renderers/deepseek_v4.py b/vllm/renderers/deepseek_v4.py index e93069209a20..ede2e9530c00 100644 --- a/vllm/renderers/deepseek_v4.py +++ b/vllm/renderers/deepseek_v4.py @@ -40,7 +40,7 @@ def render_messages( conversation, mm_data, mm_uuids = parse_chat_messages( messages, self.model_config, - content_format="string", + content_format="openai", media_io_kwargs=params.media_io_kwargs, mm_processor_kwargs=params.mm_processor_kwargs, ) @@ -67,7 +67,7 @@ async def render_messages_async( conversation, mm_data, mm_uuids = await parse_chat_messages_async( messages, self.model_config, - content_format="string", + content_format="openai", media_io_kwargs=params.media_io_kwargs, mm_processor_kwargs=params.mm_processor_kwargs, ) diff --git a/vllm/tokenizers/deepseek_v4_encoding.py b/vllm/tokenizers/deepseek_v4_encoding.py index fb5bc44e854b..37dc55936688 100644 --- a/vllm/tokenizers/deepseek_v4_encoding.py +++ b/vllm/tokenizers/deepseek_v4_encoding.py @@ -224,7 +224,7 @@ def flatten_content_blocks(content: Any) -> Any: raise ValueError( f"Unsupported content block type: {block.get('type')!r}" ) - return "".join(parts) + return "\n\n".join(parts) def render_message( @@ -267,7 +267,7 @@ def render_message( ) role = msg.get("role") - content = msg.get("content") + content = flatten_content_blocks(msg.get("content")) tools = msg.get("tools") response_format = msg.get("response_format") tool_calls = msg.get("tool_calls") @@ -337,7 +337,7 @@ def render_message( parts.append(f"[Unsupported {block_type}]") prompt += "\n\n".join(parts) else: - prompt += flatten_content_blocks(content) or "" + prompt += content or "" elif role == "latest_reminder": prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content) From c58532c86b8ab821fc1d819636480b7588bce7b5 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 18 Sep 2026 19:44:05 +0800 Subject: [PATCH 471/774] [Frontend] Upgrade XGrammar to 0.2.7 and Rust structural tags to 0.3.0 (#57272) Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 4 +- requirements/test/cuda.txt | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.txt | 2 +- rust/Cargo.lock | 4 +- rust/Cargo.toml | 2 +- .../src/tool/deepseek_dsml/deepseek_v41.rs | 4 +- .../deepseek_v41/structural_tag.rs | 108 ----------------- .../tool/deepseek_dsml/deepseek_v41/tests.rs | 10 +- rust/src/parser/src/tool/hy/structural_tag.rs | 12 +- .../src/unified/kimi_k3/structural_tag.rs | 25 ++-- .../test_structural_tag_registry.py | 71 +++-------- .../deepseekv41_engine_tool_parser.py | 2 +- vllm/tool_parsers/structural_tag_registry.py | 110 +----------------- 15 files changed, 57 insertions(+), 303 deletions(-) delete mode 100644 rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs diff --git a/requirements/common.txt b/requirements/common.txt index ff691ebe47ac..b55f680b769a 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -26,7 +26,7 @@ lm-format-enforcer == 0.11.3 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x" outlines_core == 0.2.14 lark == 1.2.2 -xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" +xgrammar == 0.2.7; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" typing_extensions >= 4.10 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317 partial-json-parser # used for parsing partial JSON outputs diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 996caab96a93..c1d49139538e 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -43,7 +43,7 @@ anyio==4.14.1 # sse-starlette # starlette # watchfiles -apache-tvm-ffi==0.1.9 +apache-tvm-ffi==0.1.14.post0 # via xgrammar arctic-inference==0.1.1 # via -r requirements/test/cuda.in @@ -1298,7 +1298,7 @@ word2number==1.1 # via lm-eval wrapt==1.17.2 # via smart-open -xgrammar==0.2.3 +xgrammar==0.2.7 # via -r requirements/test/../common.txt xxhash==3.5.0 # via diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index ea94558000ec..3207d395aa24 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1409,7 +1409,7 @@ word2number==1.1 # via lm-eval wrapt==1.17.2 # via smart-open -xgrammar==0.2.3 +xgrammar==0.2.7 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 631c3952fd95..c961b6294aa7 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1360,7 +1360,7 @@ word2number==1.1 # via lm-eval wrapt==2.1.2 # via smart-open -xgrammar==0.2.1 +xgrammar==0.2.7 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 4e4fd628b656..d9572389eb71 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1197,7 +1197,7 @@ werkzeug==3.1.8 # via schemathesis word2number==1.1 # via lm-eval -xgrammar==0.2.3 +xgrammar==0.2.7 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/rust/Cargo.lock b/rust/Cargo.lock index dee3f3078966..b4917a303b8c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -6596,9 +6596,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" [[package]] name = "xgrammar-structural-tag" -version = "0.2.0+xgrammar.0.2.4.dd729e7" +version = "0.3.0+xgrammar.0.2.7.82505d0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a" +checksum = "3864005575c2ab76057f789898f59fb9fa1378d07c5ceabc24764e012d0e90bb" dependencies = [ "auto_impl", "serde", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index dde065e79044..ea8b813619c6 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -153,7 +153,7 @@ vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } -xgrammar-structural-tag = "0.2.0" +xgrammar-structural-tag = "0.3.0" zeromq = { version = "0.6.0", default-features = false, features = [ "tokio-runtime", "all-transport", diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs index 44658b1d44f8..be68b6d2e5f3 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs @@ -4,8 +4,6 @@ use super::{DeepSeekDsmlToolParser, DsmlTokens}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; -mod structural_tag; - /// Tool parser for DeepSeek V4.1's spaced DSML tags. /// /// Arguments are emitted only after a full `invoke` block is parsed. @@ -27,7 +25,7 @@ impl ToolParser for DeepSeekV41ToolParser { } fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { - Some(&structural_tag::DeepSeekV41StructuralTagBuilder) + Some(xgrammar_structural_tag::Model::DeepSeekV41.builder()) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs deleted file mode 100644 index d5fb6118fb62..000000000000 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs +++ /dev/null @@ -1,108 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -//! Visible-text grammar for DeepSeek V4.1 tool calls. -//! -//! Adapted from xgrammar-structural-tag 0.2.0's DeepSeek V4 builder, with -//! V4.1's spaced DSML tags and schema-unconstrained parameter values. - -use serde_json::Value; -use xgrammar_structural_tag::Result; -use xgrammar_structural_tag::builders::{ - StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, -}; -use xgrammar_structural_tag::format::{ - Format, JsonSchemaFormat, StructuralTag, TagFormat, TriggeredTagsFormat, -}; -use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam}; - -use super::DsmlTokens; - -const TOKENS: DsmlTokens = DsmlTokens::V41; - -pub(super) struct DeepSeekV41StructuralTagBuilder; - -impl StructuralTagBuilder for DeepSeekV41StructuralTagBuilder { - fn build(&self, ctx: StructuralTagContext<'_>) -> Result { - let tags = ctx - .function_tools - .iter() - .map(|tool| call_tag(tool, ctx.options)) - .collect::>(); - let excludes = if ctx.options.exclude_special_tokens { - &["", ""][..] - } else { - &[] - }; - let format = match ctx.tool_choice { - BuilderToolChoice::Auto if tags.is_empty() => Format::any_text_excluding(excludes), - BuilderToolChoice::Auto => Format::TriggeredTags( - TriggeredTagsFormat::new( - &[TOKENS.tool_calls_start], - vec![TagFormat::new( - format!("{}\n", TOKENS.tool_calls_start), - Format::tags_with_separator(tags, "", true, false), - TOKENS.tool_calls_end, - )], - ) - .with_excludes(excludes), - ), - BuilderToolChoice::Forced | BuilderToolChoice::Required => Format::sequence(vec![ - Format::const_string(format!("\n\n{}\n", TOKENS.tool_calls_start)), - Format::tags_with_separator( - tags, - "", - true, - ctx.tool_choice == BuilderToolChoice::Forced, - ), - Format::const_string(TOKENS.tool_calls_end), - ]), - }; - // Reasoning boundaries remain owned by the engine reasoner; serving - // calls this visible-text builder with reasoning=false. - Ok(StructuralTag::new(format)) - } -} - -fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { - // TODO: enforce tool.parameters and strict schema semantics once V4.1 - // argument conversion is available. This builder intentionally ignores - // parameters and strict: it constrains tool names and DSML/value syntax - // only. Parameter names, presence, uniqueness and value schemas remain - // unconstrained, including for strict=true. Serving still uses strict to - // decide whether auto tool choice activates a grammar. - TagFormat::new( - format!("{} name=\"{}\">\n", TOKENS.invoke_start, tool.function.name), - Format::star(parameter(options)), - format!("{}\n", TOKENS.invoke_end), - ) -} - -fn parameter(options: StructuralTagOptions) -> Format { - Format::tag( - format!("{} name=\"", TOKENS.parameter_start), - Format::sequence(vec![ - Format::regex(r#"[^"]+"#), - Format::const_string("\" string=\""), - Format::or(vec![ - Format::sequence(vec![ - Format::const_string("true\">"), - Format::any_text_excluding(&[ - TOKENS.parameter_end, - TOKENS.invoke_end, - TOKENS.tool_calls_end, - ]), - ]), - Format::sequence(vec![ - Format::const_string("false\">"), - Format::JsonSchema( - JsonSchemaFormat::new(Value::Bool(true)) - .with_any_order(options.any_order) - .with_max_whitespace_cnt(options.max_whitespace_cnt), - ), - ]), - ]), - ]), - format!("{}\n", TOKENS.parameter_end), - ) -} diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs index 5a6758f5c935..208d88ea33a8 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs @@ -102,7 +102,7 @@ fn v4_tags_remain_text_in_v41_dialect() { } #[test] -fn structural_tag_ignores_parameter_schemas_and_strict() { +fn structural_tag_respects_parameter_schemas_and_strict() { let parser = DeepSeekV41ToolParser::create(&tools()).unwrap(); let builder = parser.structural_tag_builder().unwrap(); let options = StructuralTagOptions::default().with_reasoning(false); @@ -128,7 +128,13 @@ fn structural_tag_ignores_parameter_schemas_and_strict() { ] { let mut function = FunctionDefinition::new("lookup").with_parameters(parameters); function.strict = strict; - assert_eq!(build(function), expected); + let tag = build(function); + if strict == Some(false) { + assert_eq!(tag, expected); + } else { + assert_ne!(tag, expected); + } + assert!(tag.to_json_string().unwrap().contains("deepseek_v4_1_xml")); } } } diff --git a/rust/src/parser/src/tool/hy/structural_tag.rs b/rust/src/parser/src/tool/hy/structural_tag.rs index cce98d5b0940..e883acbafd50 100644 --- a/rust/src/parser/src/tool/hy/structural_tag.rs +++ b/rust/src/parser/src/tool/hy/structural_tag.rs @@ -219,7 +219,7 @@ mod tests { ToolChoice::required(), ); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tags_with_separator","tags":[{"begin":"get_weather\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"days"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}}]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"get_weather\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"days"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}}]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); } #[test] @@ -238,7 +238,7 @@ mod tests { ToolChoice::required(), ); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"tags_with_separator","tags":[{"begin":"get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","",""]},{"type":"const_string","value":""}]}]},"end":""}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":""}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""}]}]},"end":""}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":""}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); } #[test] @@ -259,7 +259,7 @@ mod tests { ); let value = serde_json::to_value(tag).unwrap(); - expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"limit"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq( + expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"limit"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq( &serde_json::to_string( &value["format"]["elements"][1]["tags"][0]["content"], ) @@ -276,8 +276,8 @@ mod tests { let auto = build("", &tools, ToolChoice::auto()).to_json_string().unwrap(); let forced = build("", &tools, ToolChoice::function("lookup")).to_json_string().unwrap(); - expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":[""],"tags":[{"begin":"\n","content":{"type":"tags_with_separator","tags":[{"begin":"search\n","content":{"type":"any_text","excludes":[]},"end":""},{"begin":"lookup\n","content":{"type":"any_text","excludes":[]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto); - expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[]},"end":""},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&forced); + expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":[""],"tags":[{"type":"tag","begin":"\n","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"search\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto); + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":""},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&forced); } #[test] @@ -314,7 +314,7 @@ mod tests { ); let value = serde_json::to_value(tag).unwrap(); - expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"tenant"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq( + expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"tenant"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""],"max_tokens":null,"max_chars":null},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq( &serde_json::to_string( &value["format"]["elements"][1]["tags"][0]["content"], ) diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs index 0a86a274e1ff..fc30f455e9cc 100644 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -6,14 +6,14 @@ use serde_json::{Map, Value}; use xgrammar_structural_tag::Result; use xgrammar_structural_tag::builders::{ - StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, + ReasoningMode, StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, }; use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; use super::{ ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, - RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, + RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, }; pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = @@ -50,14 +50,23 @@ impl StructuralTagBuilder for KimiK3StructuralTagBuilder { } } -fn response_prefix(reasoning: bool) -> Vec { +fn response_prefix(reasoning: ReasoningMode) -> Vec { let mut elements = Vec::new(); - if reasoning { - elements.push(Format::tag( - "", + if reasoning != ReasoningMode::Disabled { + let prefix = Format::tag( + if reasoning == ReasoningMode::Auto { + THINK_OPEN + } else { + "" + }, Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), THINK_CLOSE, - )); + ); + elements.push(if reasoning == ReasoningMode::Auto { + Format::optional(prefix) + } else { + prefix + }); elements.push(Format::const_string(RESPONSE_OPEN)); } else { elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); @@ -386,7 +395,7 @@ mod tests { ) .unwrap(); - expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"],"max_tokens":null,"max_chars":null},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); } #[test] diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index c216b398cc5b..3221d1fb2c41 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock import pytest -from xgrammar import Grammar, StructuralTag, builtin_structural_tag +from xgrammar import Grammar, StructuralTag from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -23,6 +23,7 @@ from vllm.tool_parsers.deepseekv32_engine_tool_parser import ( DeepSeekV32EngineToolParser, ) +from vllm.tool_parsers.deepseekv41_engine_tool_parser import DeepSeekV41EngineToolParser from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser @@ -83,51 +84,9 @@ def test_supported_structural_tag_models_include_vllm_builtins(): assert "hermes" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS -@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) -def test_deepseek_v41_older_xgrammar_keeps_format_constraints( - choice, - sample_tools_strict, - monkeypatch, -): - monkeypatch.setattr( - builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None, raising=False - ) - tool_choice = ( - ChatCompletionNamedToolChoiceParam( - function=ChatCompletionNamedFunction(name=choice) - ) - if choice == "get_weather" - else choice - ) - tag = get_model_structural_tag( - "deepseek_v41", sample_tools_strict, tool_choice, reasoning=False - ) - grammar = Grammar.from_structural_tag(tag) - begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n' - end = "\n" - parameter = ( - '<｜DSML｜ parameter name="undeclared" string="false">' - '{"nested": [true, null, 1.5]}\n' - ) - # Required city is omitted; extra and repeated fields remain unconstrained. - assert _is_grammar_accept_string(grammar, begin + end) - assert _is_grammar_accept_string(grammar, begin + parameter * 2 + end) - assert not _is_grammar_accept_string( - grammar, - begin + parameter.replace('{"nested": [true, null, 1.5]}', "invalid") + end, - ) - assert not _is_grammar_accept_string( - grammar, (begin + end).replace('name="get_weather"', 'name="unknown"') - ) - assert not _is_grammar_accept_string( - grammar, begin + parameter.replace("｜ parameter", "｜parameter") + end - ) - assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto") - - def test_deepseek_v41_named_choice_emits_one_call(sample_tools): tag = get_model_structural_tag( - "deepseek_v41", + "deepseek_v4_1", sample_tools, ChatCompletionNamedToolChoiceParam( function=ChatCompletionNamedFunction(name="get_weather") @@ -147,19 +106,14 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): grammar, "\n\n<｜DSML｜ calls>\n" + call * 2 + "" ) assert ( - get_model_structural_tag("deepseek_v41", sample_tools, "auto", reasoning=False) + get_model_structural_tag("deepseek_v4_1", sample_tools, "auto", reasoning=False) is None ) -@pytest.mark.skipif( - not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), - reason="Requires XGrammar's DeepSeek V4.1 builtin", -) @pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) -@pytest.mark.parametrize("reasoning", [False, True]) def test_deepseek_v41_constrains_parameters_after_reasoning( - sample_tools_strict, choice, reasoning + sample_tools_strict, choice ): """Strict calls must enforce the schema without requiring another think close.""" tool_choice = ( @@ -171,7 +125,7 @@ def test_deepseek_v41_constrains_parameters_after_reasoning( ) sample_tools_strict[0].function.parameters["additionalProperties"] = False tag = get_model_structural_tag( - "deepseek_v41", sample_tools_strict, tool_choice, reasoning=reasoning + "deepseek_v4_1", sample_tools_strict, tool_choice, reasoning=False ) grammar = Grammar.from_structural_tag(tag) begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n' @@ -195,14 +149,10 @@ def test_deepseek_v41_constrains_parameters_after_reasoning( assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto") -@pytest.mark.skipif( - not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), - reason="Requires XGrammar's DeepSeek V4.1 builtin", -) def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools): sample_tools[0].function.strict = False tag = get_model_structural_tag( - "deepseek_v41", sample_tools, "required", reasoning=False + "deepseek_v4_1", sample_tools, "required", reasoning=False ) grammar = Grammar.from_structural_tag(tag) call = ( @@ -571,6 +521,7 @@ def test_get_model_structural_tag_supports_named_tool_choice( (DeepSeekV31ToolParser, "deepseek_v3_1"), (DeepSeekV32EngineToolParser, "deepseek_v3_2"), (DeepSeekV4EngineToolParser, "deepseek_v4"), + (DeepSeekV41EngineToolParser, "deepseek_v4_1"), (Glm47MoeModelToolParser, "glm_4_7"), (Hermes2ProToolParser, "hermes"), (KimiK2ToolParser, "kimi"), @@ -638,7 +589,11 @@ def fake_get_model_structural_tag(*, reasoning: bool, **kwargs): assert captured == [False] +@pytest.mark.parametrize( + "parser_cls", [Qwen3EngineToolParser, DeepSeekV41EngineToolParser] +) def test_unified_parser_get_structural_tag_disables_reasoning( + parser_cls, monkeypatch: pytest.MonkeyPatch, sample_tools_strict: list[ChatCompletionToolsParam], ): @@ -654,7 +609,7 @@ def fake_get_model_structural_tag(*, reasoning: bool, **kwargs): ) class TestParser(DelegatingParser): - tool_parser_cls = Qwen3EngineToolParser + tool_parser_cls = parser_cls request = ChatCompletionRequest( messages=[], diff --git a/vllm/tool_parsers/deepseekv41_engine_tool_parser.py b/vllm/tool_parsers/deepseekv41_engine_tool_parser.py index a4c699aeabc0..a94883d06a52 100644 --- a/vllm/tool_parsers/deepseekv41_engine_tool_parser.py +++ b/vllm/tool_parsers/deepseekv41_engine_tool_parser.py @@ -5,4 +5,4 @@ class DeepSeekV41EngineToolParser(DeepSeekV41ParserToolAdapter): # type: ignore[valid-type, misc] - structural_tag_model = "deepseek_v41" + structural_tag_model = "deepseek_v4_1" diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 5d65b186fe30..ba86395b7a38 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -10,7 +10,6 @@ from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from openai.types.responses.tool_choice_function import ToolChoiceFunction from xgrammar import StructuralTag, normalize_tool_choice -from xgrammar import builtin_structural_tag as xgrammar_builtin_structural_tag from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag from xgrammar.openai_tool_call_schema import ( BuiltinToolParam, @@ -71,11 +70,10 @@ "deepseek_v3_2", "glm_4_7", "deepseek_v4", + "deepseek_v4_1", } ) -VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( - {"deepseek_v41", "hermes", "hy_v4", "kimi_k3"} -) +VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "hy_v4", "kimi_k3"}) SUPPORTED_STRUCTURAL_TAG_MODELS = ( XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS ) @@ -270,110 +268,6 @@ def get_function_parameters(function) -> dict[str, Any] | bool: return function.parameters if function.parameters is not None else True -_V41_CALLS_START = "<｜DSML｜ calls>" -_V41_CALLS_END = "" -_V41_INVOKE_END = "" -_V41_PARAMETER_END = "" - - -@register_vllm_structural_tag("deepseek_v41") -def get_deepseek_v41_structural_tag( - tools: list[FunctionToolParam], - builtin_tools: list[BuiltinToolParam], - tool_choice: SimplifiedToolChoice, - reasoning: bool, - token_suffix: str = "", -) -> StructuralTag: - # Serving enables this visible-text grammar after the reasoning boundary. - builder = getattr( - xgrammar_builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None - ) - if builder is not None: - return builder( - tools=tools, - builtin_tools=builtin_tools, - tool_choice=tool_choice, - reasoning="disabled", - ) - - del builtin_tools, reasoning, token_suffix - - # Compatibility with xgrammar releases without the V4.1 builtin. This constrains - # tool names and DSML/value syntax only. Parameter names, presence, uniqueness - # and value schemas remain unconstrained, including for strict=true. The request - # layer still uses strict to decide whether auto tool choice activates a grammar. - parameter = TagFormat( - begin='<｜DSML｜ parameter name="', - content=SequenceFormat( - elements=[ - RegexFormat(pattern=r'[^"]+'), - ConstStringFormat(value='" string="'), - OrFormat( - elements=[ - SequenceFormat( - elements=[ - ConstStringFormat(value='true">'), - AnyTextFormat( - excludes=[ - _V41_PARAMETER_END, - _V41_INVOKE_END, - _V41_CALLS_END, - ] - ), - ] - ), - SequenceFormat( - elements=[ - ConstStringFormat(value='false">'), - JSONSchemaFormat(json_schema=True), - ] - ), - ] - ), - ] - ), - end=f"{_V41_PARAMETER_END}\n", - ) - calls = TagsWithSeparatorFormat( - tags=[ - TagFormat( - begin=f'<｜DSML｜ invoke name="{tool.function.name}">\n', - content=StarFormat(content=parameter), - end=f"{_V41_INVOKE_END}\n", - ) - for tool in tools - ], - separator="", - at_least_one=True, - stop_after_first=tool_choice == "forced", - ) - if tool_choice == "auto": - return StructuralTag( - format=TriggeredTagsFormat( - triggers=[_V41_CALLS_START], - tags=[ - TagFormat( - begin=f"{_V41_CALLS_START}\n", - content=calls, - end=_V41_CALLS_END, - ) - ], - excludes=["", ""], - ) - if tools - else AnyTextFormat(excludes=["", ""]) - ) - return StructuralTag( - format=SequenceFormat( - elements=[ - ConstStringFormat(value=f"\n\n{_V41_CALLS_START}\n"), - calls, - ConstStringFormat(value=_V41_CALLS_END), - ] - ) - ) - - def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: arguments_field_prefix = '", "arguments": ' formats = [ From e562003adbb9c4eb177194d1318eb2e6862a8e15 Mon Sep 17 00:00:00 2001 From: Aditya Ganesh Kumar Date: Fri, 18 Sep 2026 19:46:41 +0800 Subject: [PATCH 472/774] [Tokenizer] Drop dead Mistral tokenizer shims for transformers#41962 (#56325) Signed-off-by: Aditya Ganesh Kumar Co-authored-by: Claude Opus 5 Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/tokenizers/mistral.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 66b69089f0e6..4ffefe0db6d9 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -354,22 +354,13 @@ def __call__( "`text_pair` is not supported by `MistralTokenizer.__call__`." ) - encoded = self.transformers_tokenizer( + return self.transformers_tokenizer( text=text, text_pair=text_pair, add_special_tokens=add_special_tokens, truncation=truncation, max_length=max_length, ) - # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962 - # is in, revert to only call self.transformers_tokenizer(...). - # Hack to fix wrongly added eos token, when fix will be supported the condition - # below will be False even before the revert is done. - if encoded["input_ids"] and encoded["input_ids"][-1] == self.eos_token_id: - encoded["input_ids"].pop(-1) - if attention_mask := encoded.get("attention_mask"): - attention_mask.pop(-1) - return encoded @property def vocab(self) -> list[str]: @@ -389,8 +380,10 @@ def encode( max_length: int | None = None, add_special_tokens: bool = True, ) -> list[int]: - # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962 - # is in, directly call self.transformers_tokenizer.encode(...). + # NOTE: transformers guards truncation with `if max_length and ...`, so + # `max_length=0` returns the full sequence instead of an empty one. vLLM + # treats `truncate_prompt_tokens=0` as an empty prompt, so keep slicing + # here until that is fixed upstream. encoded = self.tokenizer.encode(text, bos=add_special_tokens, eos=False) if truncation is not False and max_length is not None: @@ -437,11 +430,6 @@ def apply_chat_template( def decode( self, ids: Sequence[int] | int, skip_special_tokens: bool = False ) -> str: - # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962 - # is in, directly call self.transformers_tokenizer.decode(...). - if isinstance(ids, int): - ids = [ids] - return self.transformers_tokenizer.decode( ids, skip_special_tokens=skip_special_tokens ) From 2b8a5495a88a9a44a9d92e8c4f4aae3b55b6b74b Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 18 Sep 2026 06:50:01 -0500 Subject: [PATCH 473/774] [Mypy] Fix mypy typing for N/O models (#54142) Signed-off-by: Taneem Ibrahim --- tools/pre_commit/mypy.py | 2 - .../layers/mamba/mamba_utils.py | 10 +- vllm/model_executor/models/aimv2.py | 19 +- vllm/model_executor/models/interfaces.py | 7 + .../model_executor/models/nano_nemotron_vl.py | 112 ++++++++--- vllm/model_executor/models/nemotron_h.py | 9 +- vllm/model_executor/models/nemotron_h_mtp.py | 24 ++- vllm/model_executor/models/nemotron_vl.py | 19 +- vllm/model_executor/models/nvlm_d.py | 1 + vllm/model_executor/models/olmo_hybrid.py | 2 + vllm/model_executor/models/opencua.py | 12 +- vllm/model_executor/models/openpangu.py | 6 +- vllm/model_executor/models/openpangu_mtp.py | 17 +- vllm/model_executor/models/openpangu_vl.py | 175 ++++++++++-------- vllm/model_executor/models/openvla.py | 1 + vllm/model_executor/models/opt.py | 2 + vllm/model_executor/models/ovis.py | 13 +- vllm/model_executor/models/ovis2_5.py | 29 ++- .../processors/nano_nemotron_vl.py | 4 +- 19 files changed, 321 insertions(+), 143 deletions(-) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 234acc029858..71dd929d378e 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -113,8 +113,6 @@ EXCLUDE = [ r"vllm/model_executor/models/[kK]", - r"vllm/model_executor/models/[nN]", - r"vllm/model_executor/models/[oO]", r"vllm/model_executor/models/[qQ]", r"vllm/model_executor/models/[uU]", r"vllm/model_executor/models/[vV]", diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index db7e119560c5..8c89488d2d58 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -222,12 +222,18 @@ def mamba2_state_shape( @classmethod def append_replayssm_ring( cls, - base_shapes: tuple[tuple[int, ...], ...], + base_shapes: tuple[tuple[int, int], tuple[int, int, int]], n_groups: int, tp_world_size: int, logical_window: int, backend: MambaBackendEnum, - ) -> tuple[tuple[int, ...], ...]: + ) -> tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int], + tuple[int, int, int], + ]: """Append the physical ReplaySSM ring shapes. ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index 3461b5713734..4e31b99944f1 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -26,7 +26,10 @@ class AIMv2SwiGLUFFN(nn.Module): def __init__( - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str + self, + config: AIMv2Config, + quant_config: QuantizationConfig | None, + prefix: str, ): super().__init__() hidden_features = config.intermediate_size @@ -91,7 +94,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Attention(nn.Module): def __init__( - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str + self, + config: AIMv2Config, + quant_config: QuantizationConfig | None, + prefix: str, ): super().__init__() self.config = config @@ -144,7 +150,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Block(nn.Module): def __init__( - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str + self, + config: AIMv2Config, + quant_config: QuantizationConfig | None, + prefix: str, ): super().__init__() self.attn = AIMv2Attention( @@ -166,7 +175,7 @@ class AIMv2Transformer(nn.Module): def __init__( self, config: AIMv2Config, - quant_config: QuantizationConfig, + quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, prefix: str = "", @@ -206,7 +215,7 @@ class AIMv2Model(torch.nn.Module): def __init__( self, config: AIMv2Config, - quant_config: QuantizationConfig, + quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, prefix: str = "", diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 1cefd63dd63c..21a09cb7ba2e 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -91,6 +91,13 @@ tuple[int, int, int], tuple[int, int, int], ] + | tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int], + tuple[int, int, int], + ] ) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 6e500c0c2c8e..fee431617c8b 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -12,7 +12,7 @@ from collections.abc import Iterable, Mapping, Sequence from functools import cached_property from io import BytesIO -from typing import Annotated, Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias, TypedDict, cast import torch import torch.nn as nn @@ -20,7 +20,11 @@ from vllm import envs from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import ( + AudioDummyOptions, + BaseDummyOptions, + VideoDummyOptions, +) from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.logger import init_logger from vllm.lora.layers.base import BaseLayerWithLoRA @@ -114,7 +118,7 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema): type: Literal["audio_features"] = "audio_features" input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")] feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")] - audio_num_clips: list[int] + audio_num_clips: list[int] | torch.Tensor class NanoNemotronVLImagePixelInputs(TensorSchema): @@ -195,6 +199,12 @@ class NanoNemotronVLVideoEmbeddingInputs(TensorSchema): ) +class NanoNemotronVLMultiModalInputs(TypedDict, total=False): + images: NanoNemotronVLImageInputs | None + videos: NanoNemotronVLVideoInputs | None + audios: NanoNemotronVLAudioFeatureInputs | None + + class NanoNemotronVLProcessingInfo(BaseProcessingInfo): def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor: # Not accepted by NanoNemotronVLProcessor.__init__ @@ -465,8 +475,11 @@ def get_image_replacement(item_idx: int): if isinstance(images, ImageEmbeddingItems): feature_size = images.get_feature_size(item_idx) elif self.info.is_dynamic_tiler: - feature_size = out_mm_data["num_tokens_per_image"][item_idx] + num_tokens_per_image = out_mm_data["num_tokens_per_image"] + assert isinstance(num_tokens_per_image, list) + feature_size = num_tokens_per_image[item_idx] else: + assert isinstance(images, ImageProcessorItems) image_size = images.get_image_size(item_idx) max_num_tiles = hf_processor.max_num_tiles feature_size = hf_processor.get_num_image_tokens( @@ -530,10 +543,12 @@ def get_video_replacement(item_idx: int): assert isinstance(num_patches, int) T = hf_processor.video_temporal_patch_size + num_tubelets: int | None if T > 1 and num_patches is not None: num_tubelets = math.ceil(num_patches / T) else: num_tubelets = num_patches + assert num_tubelets is not None video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate if video_pruning_rate is not None and video_pruning_rate > 0.0: @@ -635,6 +650,7 @@ def _extract_audio_from_videos( audio_items: list[AudioItem] = [] has_audio: list[bool] = [] for idx, metadata in enumerate(metadata_list): + assert metadata is not None video_bytes = metadata.get("original_video_bytes") if video_bytes is None or len(video_bytes) == 0: raise ValueError( @@ -661,21 +677,25 @@ def _extract_audio_from_videos( # Create a new VideoProcessorItems with metadata that does not contain # the large video bytes, to avoid modifying the input `mm_items`. - new_metadata_list = [ - {k: v for k, v in meta.items() if k != "original_video_bytes"} - for meta in metadata_list - ] + new_metadata_list = metadata_list.copy() + for idx, metadata in enumerate(metadata_list): + assert metadata is not None + new_metadata_list[idx] = { + key: value + for key, value in metadata.items() + if key != "original_video_bytes" + } new_videos = VideoProcessorItems(data=videos.data, metadata=new_metadata_list) - audio_parsed = {} + audio_parsed = MultiModalDataItems() if audio_items: audio_parsed = self.data_parser.parse_mm_data({"audio": audio_items}) # Create a new MultiModalDataItems with the new video and audio items. new_mm_items_dict = {**mm_items, **audio_parsed, "video": new_videos} - mm_items = MultiModalDataItems(new_mm_items_dict) + new_mm_items = MultiModalDataItems(new_mm_items_dict) - return mm_items, audio_items, has_audio + return new_mm_items, audio_items, has_audio def apply( self, @@ -800,7 +820,7 @@ def _get_dummy_videos( ) videos = [v.copy() for v in videos] - video_items = [] + video_items: list[VideoItem] = [] for video in videos: video_num_frames = video.shape[0] video_metadata = { @@ -869,6 +889,9 @@ def get_dummy_mm_data( target_num_frames = num_frames num_videos = mm_counts.get("video", 0) video_overrides = mm_options.get("video") + assert video_overrides is None or isinstance( + video_overrides, VideoDummyOptions + ) dummy_video = { "video": self._get_dummy_videos( width=video_width, @@ -884,6 +907,9 @@ def get_dummy_mm_data( if sound_config := self.info.sound_config: num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") if mm_options else None + assert audio_overrides is None or isinstance( + audio_overrides, AudioDummyOptions + ) tokens_per_audio = max(1, seq_len // max(num_audios, 1)) max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate calculated_max_audio_num_samples = ParakeetExtractor.audio_length( @@ -946,6 +972,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): model_config = vllm_config.model_config config = model_config.hf_config multimodal_config = model_config.multimodal_config + assert multimodal_config is not None image_size = config.force_image_size patch_size = config.patch_size self.patch_size = patch_size @@ -1138,17 +1165,26 @@ def _parse_and_validate_image_input( return None if self.dynamic_resolution: + assert isinstance(pixel_values_flat, (list, torch.Tensor)) pixel_values_flat = DynamicResolutionImageTiler.stack( pixel_values_flat, self.patch_size ) + imgs_sizes = kwargs.pop("imgs_sizes") + num_tokens_per_image = kwargs.pop("num_tokens_per_image") + assert isinstance(imgs_sizes, list) + assert isinstance(num_tokens_per_image, list) return NanoNemotronVLImagePixelInputsDynamic( - pixel_values_flat=pixel_values_flat, **kwargs + pixel_values_flat=pixel_values_flat, + imgs_sizes=imgs_sizes, + num_tokens_per_image=num_tokens_per_image, ) else: + assert isinstance(pixel_values_flat, torch.Tensor) + image_num_patches = kwargs.pop("image_num_patches") + assert isinstance(image_num_patches, torch.Tensor) return NanoNemotronVLImagePixelInputs( pixel_values_flat=pixel_values_flat, - num_patches=kwargs.pop("image_num_patches"), - **kwargs, + num_patches=image_num_patches, ) def _process_image_input_dynamic( @@ -1166,7 +1202,8 @@ def _process_image_input_dynamic( return image_embeds.split(num_tokens_per_image) def _process_image_input( - self, image_input: NanoNemotronVLImagePixelInputs + self, + image_input: NanoNemotronVLImagePixelInputs | NanoNemotronVLVideoPixelInputs, ) -> tuple[torch.Tensor, ...]: image_embeds = self.extract_feature(image_input["pixel_values_flat"]) num_patches = image_input["num_patches"] @@ -1362,7 +1399,8 @@ def _create_final_video_embeddings( # of them than max_num_batched_tokens. # Embed them with the base weights - # the LoRA delta is undefined for tokens with no mapping entry. - embed_tokens = self.get_language_model().model.embed_tokens + language_model = cast(NemotronHForCausalLM, self.get_language_model()) + embed_tokens = language_model.model.embed_tokens if isinstance(embed_tokens, BaseLayerWithLoRA): embed_tokens = embed_tokens.base_layer text_embeddings = embed_tokens(repl_token_ids) @@ -1388,26 +1426,29 @@ def _parse_and_validate_video_input( return None if video_embeds is not None: + assert isinstance(video_embeds, (torch.Tensor, list)) return NanoNemotronVLVideoEmbeddingInputs( type="video_embeds", data=video_embeds, ) if pixel_values_flat_video is not None: - if torch.is_tensor(frames_indices): + if isinstance(frames_indices, torch.Tensor): frames_indices = frames_indices.flatten() else: + assert isinstance(frames_indices, (list, tuple)) frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0) - if torch.is_tensor(frame_duration_ms): + if isinstance(frame_duration_ms, torch.Tensor): frame_duration_ms = frame_duration_ms.flatten() else: + assert isinstance(frame_duration_ms, (list, tuple)) frame_duration_ms = torch.cat( [f.flatten() for f in frame_duration_ms], dim=0 ) if ( - torch.is_tensor(pixel_values_flat_video) + isinstance(pixel_values_flat_video, torch.Tensor) and pixel_values_flat_video.ndim == 5 ): # batched._reduce_data stacked same-shape videos into @@ -1415,9 +1456,11 @@ def _parse_and_validate_video_input( # same-H,W cat path below handles it uniformly. pixel_values_flat_video = list(pixel_values_flat_video) - if not torch.is_tensor(pixel_values_flat_video): + if not isinstance(pixel_values_flat_video, torch.Tensor): + assert isinstance(pixel_values_flat_video, (list, tuple)) pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0) + assert isinstance(video_num_patches, torch.Tensor) expected_h = pixel_values_flat_video.shape[-2] expected_w = pixel_values_flat_video.shape[-1] num_frames = video_num_patches[0].item() @@ -1434,8 +1477,10 @@ def _parse_and_validate_video_input( raise AssertionError("This line should be unreachable.") - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - modalities = {} + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> NanoNemotronVLMultiModalInputs: + modalities: NanoNemotronVLMultiModalInputs = {} # Preserve the order of modalities if there are multiple of them # from the order of kwargs. for input_key in kwargs: @@ -1458,8 +1503,17 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: ) and "audios" not in modalities ): + input_audio_features = kwargs.get("input_audio_features") + feature_attention_mask = kwargs.get("feature_attention_mask") + audio_num_clips = kwargs.get("audio_num_clips") + assert isinstance(input_audio_features, torch.Tensor) + assert isinstance(feature_attention_mask, torch.Tensor) + assert isinstance(audio_num_clips, (list, torch.Tensor)) modalities["audios"] = NanoNemotronVLAudioFeatureInputs( - **kwargs, validate=False + input_audio_features=input_audio_features, + feature_attention_mask=feature_attention_mask, + audio_num_clips=audio_num_clips, + validate=False, ) return modalities @@ -1479,23 +1533,26 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: for modality in modalities: if modality == "images": image_input = modalities["images"] - if image_input["type"] == "image_embeds": + assert image_input is not None + if isinstance(image_input, NanoNemotronVLImageEmbeddingInputs): image_embeddings = image_input["data"] - elif self.dynamic_resolution: - assert image_input["type"] == "pixel_values_dynamic" + elif isinstance(image_input, NanoNemotronVLImagePixelInputsDynamic): image_embeddings = self._process_image_input_dynamic(image_input) else: image_embeddings = self._process_image_input(image_input) multimodal_embeddings += tuple(image_embeddings) if modality == "videos": video_input = modalities["videos"] + assert video_input is not None if video_input["type"] == "video_embeds": video_embeddings = video_input["data"] else: + assert isinstance(video_input, NanoNemotronVLVideoPixelInputs) video_embeddings = self._process_video_input(video_input) multimodal_embeddings += tuple(video_embeddings) if modality == "audios": audio_input = modalities["audios"] + assert audio_input is not None audio_embeddings = self._process_audio_input(audio_input) multimodal_embeddings += tuple(audio_embeddings) @@ -1538,6 +1595,7 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): mm_config = self.model_config.multimodal_config + assert mm_config is not None load_multimodal_weights = not all( mm_config.get_limit_per_prompt(modality) == 0 for modality in ("image", "video", "audio") diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e71b7d8e6cf0..e5451d16ba02 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -68,6 +68,7 @@ EagleModelMixin, HasInnerState, IsHybrid, + MambaStateShapes, MixtureOfExperts, SupportsEagle, SupportsEagle3, @@ -138,6 +139,7 @@ def __init__( prefix: str = "", ): super().__init__() + assert parallel_config is not None self.tp_size = get_tensor_model_parallel_world_size() self.routed_scaling_factor = config.routed_scaling_factor @@ -188,6 +190,8 @@ def __init__( prefix=f"{prefix}.shared_experts", ) + self.fc1_latent_proj: ReplicatedLinear | None + self.fc2_latent_proj: ReplicatedLinear | None if self.use_latent_moe: self.fc1_latent_proj = ReplicatedLinear( input_size=config.hidden_size, @@ -722,7 +726,7 @@ class NemotronHForCausalLM( SupportsReplaySSM, ): # Relevant only if self.has_moe is True - is_non_gated_moe: bool = True + is_non_gated_moe = True hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={"backbone": "model", "mtp": None}, @@ -773,7 +777,7 @@ def get_mamba_state_dtype_from_config( def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[tuple[int, ...], ...]: + ) -> MambaStateShapes: """Calculate shapes for Mamba's convolutional and state caches. Args: @@ -858,6 +862,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_layers.append(layer.mixer.experts) self.num_moe_layers = len(self.moe_layers) + assert example_moe is not None self.num_logical_experts = example_moe.n_logical_experts self.num_physical_experts = example_moe.n_physical_experts self.num_local_physical_experts = example_moe.n_local_physical_experts # noqa: E501 diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index 96f1b6685ed9..6b9de6300918 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -92,10 +92,12 @@ def __init__( def forward( self, - inputs_embeds: torch.Tensor, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None = None, + *, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, ) -> tuple[torch.Tensor, torch.Tensor | None]: # Start projections (Fusion) if self.has_start_projections: @@ -179,10 +181,11 @@ def __init__( def forward( self, - inputs_embeds: torch.Tensor, - positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None = None, + *, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, ) -> tuple[torch.Tensor, torch.Tensor | None]: # Start projections (Fusion) if self.has_start_projections: @@ -305,13 +308,14 @@ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: def forward( self, - input_ids: torch.Tensor, + input_ids: torch.Tensor | None, positions: torch.Tensor, hidden_states: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: if inputs_embeds is None: + assert input_ids is not None inputs_embeds = self.get_input_embeddings(input_ids) residual = None @@ -386,6 +390,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @staticmethod def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None: vllm_config = kwargs.get("vllm_config") + assert isinstance(vllm_config, VllmConfig) return get_draft_quant_config(vllm_config) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -393,14 +398,15 @@ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: def forward( self, - input_ids: torch.Tensor, + input_ids: torch.Tensor | None, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | None = None, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor: """Forward - applies attention-based MTP.""" + assert hidden_states is not None hidden_states = self.model( input_ids, positions, @@ -473,9 +479,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = name.replace("backbone.", "model.") if "scale" in name or "zero_point" in name: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: continue + name = remapped_name # Handle stacked parameters (qkv_proj) for attention layers is_stacked = False @@ -547,6 +554,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) + assert weight_loader is not None weight_loader(param, loaded_weight) loaded_params.add(name) diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py index bf05848858d3..7b2777a71087 100644 --- a/vllm/model_executor/models/nemotron_vl.py +++ b/vllm/model_executor/models/nemotron_vl.py @@ -138,7 +138,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "language_model"), ) - self.img_context_token_id = None + self.img_context_token_id: int | None = None self.visual_token_mask = None self.make_empty_intermediate_tensors = ( @@ -146,7 +146,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) def _patch_quant_config( - self, config: PretrainedConfig, quant_config: QuantizationConfig + self, config: PretrainedConfig, quant_config: QuantizationConfig | None ): # the awq models from OpenGVLab missing `modules_to_not_convert` # patch the quant_config to add `modules_to_not_convert` back @@ -518,12 +518,21 @@ def _init_vision_model( use_head=False, ) - def _init_mlp1(self, config: PretrainedConfig) -> nn.Module: + def _init_mlp1( + self, + config: PretrainedConfig, + vit_hidden_size: int | None = None, + vision_projection_hidden_size: int | None = None, + ) -> nn.Module: """Override to use different MLP structure for embedding model.""" + if vit_hidden_size is None: + vit_hidden_size = config.vision_config.hidden_size + if vision_projection_hidden_size is None: + vision_projection_hidden_size = config.get_text_config().hidden_size return super()._init_mlp1( config, - vit_hidden_size=config.vision_config.hidden_size, - vision_projection_hidden_size=config.get_text_config().hidden_size, + vit_hidden_size=vit_hidden_size, + vision_projection_hidden_size=vision_projection_hidden_size, ) def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/nvlm_d.py b/vllm/model_executor/models/nvlm_d.py index d937fade1d95..a69f024a0158 100644 --- a/vllm/model_executor/models/nvlm_d.py +++ b/vllm/model_executor/models/nvlm_d.py @@ -128,6 +128,7 @@ def get_replacement_nvlm(item_idx: int): if isinstance(images, ImageEmbeddingItems): feature_size = images.get_feature_size(item_idx) else: + assert isinstance(images, ImageProcessorItems) image_size = images.get_image_size(item_idx) feature_size = self.info.get_num_image_tokens( image_width=image_size.width, diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py index 9a24df5ec9a6..e32087ea2134 100644 --- a/vllm/model_executor/models/olmo_hybrid.py +++ b/vllm/model_executor/models/olmo_hybrid.py @@ -143,6 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): rope_parameters["rope_theta"] is not None ) + self.rotary_emb: nn.Module | None if self._use_rope: self.rotary_emb = get_rope( self.head_dim, @@ -183,6 +184,7 @@ def forward( q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = self._apply_qk_norm(q, k) if self._use_rope: + assert self.rotary_emb is not None q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/opencua.py b/vllm/model_executor/models/opencua.py index 6a875c813cb9..6b12d1d3fbe7 100644 --- a/vllm/model_executor/models/opencua.py +++ b/vllm/model_executor/models/opencua.py @@ -7,6 +7,7 @@ """Inference-only OpenCUA-7B model compatible with HuggingFace weights.""" +import typing from collections.abc import Mapping, Sequence from typing import Any @@ -27,6 +28,7 @@ ) from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, BaseMultiModalProcessor, PromptReplacement, PromptUpdate, @@ -178,7 +180,13 @@ def get_replacement_opencua(item_idx: int): ] -class OpenCUADummyInputsBuilder(Qwen2VLDummyInputsBuilder): +if typing.TYPE_CHECKING: + _OpenCUADummyInputsBuilderBase = BaseDummyInputsBuilder[OpenCUAProcessingInfo] +else: + _OpenCUADummyInputsBuilderBase = Qwen2VLDummyInputsBuilder + + +class OpenCUADummyInputsBuilder(_OpenCUADummyInputsBuilderBase): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) @@ -220,7 +228,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config - multimodal_config = vllm_config.model_config.multimodal_config + multimodal_config = vllm_config.model_config.get_multimodal_config() self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.config = config diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index b8937f2e09fb..4aa603b458fa 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -175,6 +175,7 @@ def __init__( self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.shared_experts: OpenPanguMLP | None if config.n_shared_experts is not None: intermediate_size = config.moe_intermediate_size * config.n_shared_experts self.shared_experts = OpenPanguMLP( @@ -283,7 +284,7 @@ def __init__( ) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) self.q_b_proj = ColumnParallelLinear( - q_lora_rank, + self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, quant_config=quant_config, @@ -568,6 +569,9 @@ def __init__( self.qk_nope_dim = getattr(config, "qk_nope_dim", None) self.qk_rope_dim = getattr(config, "qk_rope_dim", None) self.v_channels = getattr(config, "v_channels", None) + assert self.qk_nope_dim is not None + assert self.qk_rope_dim is not None + assert self.v_channels is not None self.head_dim = self.qk_rope_dim + self.qk_nope_dim self.q_size = self.num_heads * self.head_dim self.k_size = self.num_kv_heads * self.head_dim diff --git a/vllm/model_executor/models/openpangu_mtp.py b/vllm/model_executor/models/openpangu_mtp.py index d25efd8c0369..c91453342a40 100644 --- a/vllm/model_executor/models/openpangu_mtp.py +++ b/vllm/model_executor/models/openpangu_mtp.py @@ -52,7 +52,9 @@ class OpenPanguMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer): def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: nn.Module.__init__(self) - config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + config = speculative_config.draft_model_config.hf_config self.config = config quant_config = vllm_config.quant_config @@ -200,10 +202,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: break else: for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: + ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) = mapping + if expert_weight_name not in name: continue - name = name.replace(weight_name, param_name) + name = name.replace(expert_weight_name, expert_param_name) param = params_dict[name] weight_loader = param.weight_loader @@ -211,7 +218,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: param, loaded_weight, name, - shard_id=shard_id, + shard_id=expert_shard_id, expert_id=expert_id, ) break diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py index 8b85d0371b99..741ec2155fe7 100644 --- a/vllm/model_executor/models/openpangu_vl.py +++ b/vllm/model_executor/models/openpangu_vl.py @@ -22,7 +22,7 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import lru_cache, partial -from typing import Annotated, Literal, Optional +from typing import Annotated, Literal, Optional, TypeAlias import torch import torch.nn as nn @@ -627,7 +627,7 @@ def _padding_weight(name: str, w: torch.Tensor) -> torch.Tensor: pad[-(dim + 1) * 2 + 1] = pad_len return F.pad(w, pad, mode="constant", value=0) - stacked_params_mapping = [ + stacked_params_mapping: list[tuple[str, str, str | int]] = [ ("attn.qkv.", "attn.q.", "q"), ("attn.qkv.", "attn.k.", "k"), ("attn.qkv.", "attn.v.", "v"), @@ -677,15 +677,11 @@ class OpenPanguVLProcessingInfo(Qwen2_5_VLProcessingInfo): def get_hf_config(self): return self.ctx.model_config.hf_config - def get_hf_processor( - self, - *, - min_pixels: int | None = None, - max_pixels: int | None = None, - size: dict[str, int] | None = None, - fps: float | list[float] | None = None, - **kwargs: object, - ): + def get_hf_processor(self, **kwargs: object): + kwargs.pop("min_pixels", None) + kwargs.pop("max_pixels", None) + kwargs.pop("size", None) + fps = kwargs.pop("fps", None) if fps is not None: kwargs["fps"] = fps @@ -747,11 +743,19 @@ class OpenPanguVLVideoEmbeddingInputs(TensorSchema): ] +OpenPanguVLImageInputs: TypeAlias = ( + OpenPanguVLImagePixelInputs | OpenPanguVLImageEmbeddingInputs +) +OpenPanguVLVideoInputs: TypeAlias = ( + OpenPanguVLVideoPixelInputs | OpenPanguVLVideoEmbeddingInputs +) + + class OpenPanguVLMultiModalProcessor(Qwen2_5_VLMultiModalProcessor): def _get_prompt_updates( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, any], + hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) @@ -878,10 +882,12 @@ def _parse_preprocess_params(self, vision_config): self.do_rescale = image_processor.do_rescale self.rescale_factor = image_processor.rescale_factor self.do_normalize = image_processor.do_normalize - self.image_mean = tuple(image_processor.image_mean) - self.image_std = tuple(image_processor.image_std) + self.image_mean: tuple[float, ...] = tuple(image_processor.image_mean) + self.image_std: tuple[float, ...] = tuple(image_processor.image_std) - def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig): + def _maybe_ignore_quant_config( + self, quant_config: QuantizationConfig | None + ) -> QuantizationConfig | None: if isinstance(quant_config, AutoGPTQConfig): return None return quant_config @@ -904,7 +910,9 @@ def _validate_and_reshape_mm_tensor( else: return torch.concat(mm_input) - def _parse_and_validate_image_input(self, **kwargs: object): + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> OpenPanguVLImageInputs | None: pixel_values = kwargs.pop("pixel_values", None) image_embeds = kwargs.pop("image_embeds", None) image_grid_thw = kwargs.pop("image_grid_thw", None) @@ -932,26 +940,27 @@ def _parse_and_validate_image_input(self, **kwargs: object): image_grid_thw=image_grid_thw, ) - if image_embeds is not None: - image_embeds = self._validate_and_reshape_mm_tensor( - image_embeds, "image embeds" - ) - image_grid_thw = self._validate_and_reshape_mm_tensor( - image_grid_thw, "image grid_thw" - ) + assert image_embeds is not None + image_embeds = self._validate_and_reshape_mm_tensor( + image_embeds, "image embeds" + ) + image_grid_thw = self._validate_and_reshape_mm_tensor( + image_grid_thw, "image grid_thw" + ) - if not isinstance(image_embeds, torch.Tensor): - raise ValueError( - "Incorrect type of image embeddings. " - f"Got type: {type(image_embeds)}" - ) - return OpenPanguVLImageEmbeddingInputs( - type="image_embeds", - image_embeds=image_embeds, - image_grid_thw=image_grid_thw, + if not isinstance(image_embeds, torch.Tensor): + raise ValueError( + f"Incorrect type of image embeddings. Got type: {type(image_embeds)}" ) + return OpenPanguVLImageEmbeddingInputs( + type="image_embeds", + image_embeds=image_embeds, + image_grid_thw=image_grid_thw, + ) - def _parse_and_validate_video_input(self, **kwargs: object): + def _parse_and_validate_video_input( + self, **kwargs: object + ) -> OpenPanguVLVideoInputs | None: pixel_values_videos = kwargs.pop("pixel_values_videos", None) video_embeds = kwargs.pop("video_embeds", None) video_grid_thw = kwargs.pop("video_grid_thw", None) @@ -973,27 +982,30 @@ def _parse_and_validate_video_input(self, **kwargs: object): video_grid_thw=video_grid_thw, ) - if video_embeds is not None: - video_embeds = self._validate_and_reshape_mm_tensor( - video_embeds, "video embeds" - ) - video_grid_thw = self._validate_and_reshape_mm_tensor( - video_grid_thw, "video grid_thw" - ) + assert video_embeds is not None + video_embeds = self._validate_and_reshape_mm_tensor( + video_embeds, "video embeds" + ) + video_grid_thw = self._validate_and_reshape_mm_tensor( + video_grid_thw, "video grid_thw" + ) - if not isinstance(video_embeds, torch.Tensor): - raise ValueError( - "Incorrect type of video embeddings. " - f"Got type: {type(video_embeds)}" - ) - return OpenPanguVLVideoEmbeddingInputs( - type="video_embeds", - video_embeds=video_embeds, - video_grid_thw=video_grid_thw, + if not isinstance(video_embeds, torch.Tensor): + raise ValueError( + f"Incorrect type of video embeddings. Got type: {type(video_embeds)}" ) + return OpenPanguVLVideoEmbeddingInputs( + type="video_embeds", + video_embeds=video_embeds, + video_grid_thw=video_grid_thw, + ) - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - mm_input_by_modality = {} + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> dict[str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None]: + mm_input_by_modality: dict[ + str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None + ] = {} for input_key in kwargs: if ( input_key in ("pixel_values", "image_embeds") @@ -1019,16 +1031,19 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: multimodal_embeddings: tuple[torch.Tensor, ...] = () for modality in mm_input_by_modality: - multimodal_input = mm_input_by_modality[modality] if modality == "image": - vision_embeddings = self._process_image_input(multimodal_input) + image_input = mm_input_by_modality["image"] + assert isinstance(image_input, OpenPanguVLImageInputs) + vision_embeddings = self._process_image_input(image_input) multimodal_embeddings = ( multimodal_embeddings if not vision_embeddings else (multimodal_embeddings + vision_embeddings) ) if modality == "video": - video_embeddings = self._process_video_input(multimodal_input) + video_input = mm_input_by_modality["video"] + assert isinstance(video_input, OpenPanguVLVideoInputs) + video_embeddings = self._process_video_input(video_input) multimodal_embeddings = ( multimodal_embeddings if not video_embeddings @@ -1041,20 +1056,19 @@ def get_input_embeddings( input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None, ) -> torch.Tensor: - if multimodal_embeddings is None: - return self.language_model.embed_input_ids(input_ids) - - mm_token_ids = input_ids.new_tensor( - [self.config.image_token_id, self.config.video_token_id] - ) - is_multimodal = torch.isin(input_ids, mm_token_ids) - return self.embed_input_ids( - input_ids, - multimodal_embeddings, - is_multimodal=is_multimodal, - ) + inputs_embeds = self.language_model.embed_input_ids(input_ids) + if multimodal_embeddings is not None: + inputs_embeds = self.embed_input_ids( # type: ignore[call-overload] + input_ids, + inputs_embeds, + multimodal_embeddings, + [self.config.image_token_id, self.config.video_token_id], + ) + return inputs_embeds - def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]: + def _process_image_input( + self, image_input: OpenPanguVLImageInputs + ) -> tuple[torch.Tensor, ...]: grid_thw = image_input["image_grid_thw"] if grid_thw.ndim != 2: raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}") @@ -1085,7 +1099,9 @@ def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]: sizes = grid_thw.prod(-1) // merge_size // merge_size return image_embeds.split(sizes.tolist()) - def _process_video_input(self, video_input) -> torch.Tensor: + def _process_video_input( + self, video_input: OpenPanguVLVideoInputs + ) -> tuple[torch.Tensor, ...]: grid_thw = video_input["video_grid_thw"] if grid_thw.ndim != 2: raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}") @@ -1158,8 +1174,13 @@ def iter_mm_grid_thw( for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): offset = mm_feature.mm_position.offset modality = mm_feature.modality + feature_data = mm_feature.data + assert feature_data is not None if modality == "image": - t, h, w = mm_feature.data["image_grid_thw"].data.tolist() + grid_thw_item = feature_data["image_grid_thw"] + grid_thw = grid_thw_item.data + assert isinstance(grid_thw, torch.Tensor) + t, h, w = grid_thw.tolist() assert t == 1, f"Image must have 1 frame, got {t}" yield ( modality, @@ -1169,7 +1190,10 @@ def iter_mm_grid_thw( w // spatial_merge_size, ) elif modality == "video": - t, h, w = mm_feature.data["video_grid_thw"].data.tolist() + grid_thw_item = feature_data["video_grid_thw"] + grid_thw = grid_thw_item.data + assert isinstance(grid_thw, torch.Tensor) + t, h, w = grid_thw.tolist() yield ( modality, offset, @@ -1281,14 +1305,17 @@ def normalize(image, mean, std): @lru_cache(maxsize=10) def _fuse_mean_std_and_rescale_factor( do_normalize: bool | None = None, - image_mean: float | list[float] | None = None, - image_std: float | list[float] | None = None, + image_mean: float | tuple[float, ...] | None = None, + image_std: float | tuple[float, ...] | None = None, do_rescale: bool | None = None, rescale_factor: float | None = None, device: Optional["torch.device"] = None, ) -> tuple: if do_rescale and do_normalize: # Fused rescale and normalize + assert rescale_factor is not None + assert image_mean is not None + assert image_std is not None image_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor) image_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor) do_rescale = False @@ -1300,8 +1327,8 @@ def rescale_and_normalize( do_rescale: bool, rescale_factor: float, do_normalize: bool, - image_mean: float | list[float], - image_std: float | list[float], + image_mean: float | tuple[float, ...], + image_std: float | tuple[float, ...], dtype: torch.dtype = torch.bfloat16, ) -> "torch.Tensor": """Rescale and normalize images.""" diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py index 77099ffb6efc..add7eac41f56 100644 --- a/vllm/model_executor/models/openvla.py +++ b/vllm/model_executor/models/openvla.py @@ -349,6 +349,7 @@ def get_insertion(item_idx: int) -> PromptUpdateDetails: if isinstance(images, ImageEmbeddingItems): num_image_tokens = images.get_feature_size(item_idx) else: + assert isinstance(images, ImageProcessorItems) image_size = images.get_image_size(item_idx) num_image_tokens = self.info.get_num_image_tokens( image_width=image_size.width, diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py index 1efe9eb8a8ca..4726c72b417b 100644 --- a/vllm/model_executor/models/opt.py +++ b/vllm/model_executor/models/opt.py @@ -218,6 +218,7 @@ def __init__( ) # Project out & in will be replicated if they exist. + self.project_out: ReplicatedLinear | None if config.word_embed_proj_dim != config.hidden_size: self.project_out = ReplicatedLinear( config.hidden_size, @@ -229,6 +230,7 @@ def __init__( else: self.project_out = None + self.project_in: ReplicatedLinear | None if config.word_embed_proj_dim != config.hidden_size: self.project_in = ReplicatedLinear( config.word_embed_proj_dim, diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index 5563b9a9cea2..bd00686b41d7 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -59,7 +59,12 @@ from vllm.transformers_utils.processors.ovis import OvisProcessor from vllm.utils.tensor_schema import TensorSchema, TensorShape -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + supports_pp, +) # Cannot find the following number from hf config. IMAGE_TOKEN = "" @@ -281,7 +286,7 @@ def get_image_segment_len(self) -> int: # minus 1 for presented image token return (patch_grid_length // hidden_stride) ** 2 - 1 - def get_image_pad_token(self) -> str: + def get_image_pad_token(self) -> str | None: hf_text_config = self.get_hf_config().get_text_config() text_model_type = hf_text_config.model_type return IMAGE_PAD_TOKEN_MAP.get(text_model_type) @@ -439,8 +444,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): text_model_type = self.config.get_text_config().model_type self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type] + language_model = self.get_language_model() + assert supports_pp(language_model) self.make_empty_intermediate_tensors = ( - self.get_language_model().make_empty_intermediate_tensors + language_model.make_empty_intermediate_tensors ) def _parse_and_validate_image_input( diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py index 0b8a2c5f84bb..588c06b8dea6 100644 --- a/vllm/model_executor/models/ovis2_5.py +++ b/vllm/model_executor/models/ovis2_5.py @@ -4,14 +4,14 @@ from collections.abc import Iterable, Mapping from functools import partial -from typing import Annotated, Literal +from typing import Annotated, Literal, TypedDict import torch import torch.nn as nn from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -40,7 +40,12 @@ from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor from vllm.utils.tensor_schema import TensorSchema, TensorShape -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + supports_pp, +) IMAGE_TOKEN = "" VIDEO_TOKEN = "``, which deepseek_r1 classifies + wholly as reasoning, so the text lands there rather than in content. + What this pins is detokenization fidelity through the parser path. + """ messages = [{"role": "user", "content": "What is 2+2?"}] gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) @@ -814,16 +819,22 @@ async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer): "model": PARSER_MODEL, "generate_response": _e2e_generate_response(output_ids), "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": {"model": PARSER_MODEL, "messages": messages}, }, ) assert resp.status_code == 200, resp.text - content = resp.json()["choices"][0]["message"]["content"] - assert content == expected + message = resp.json()["choices"][0]["message"] + assert message["reasoning"] == expected + assert message["content"] is None @pytest.mark.asyncio async def test_e2e_token_identity(parser_client, parser_tokenizer): - """encode(derender(token_ids)) == token_ids (RL invariant).""" + """encode(derender(token_ids)) == token_ids (RL invariant). + + Markerless output comes back as reasoning (see + ``test_e2e_plain_roundtrip``), so re-encode that. + """ messages = [{"role": "user", "content": "Hi"}] gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) @@ -836,17 +847,22 @@ async def test_e2e_token_identity(parser_client, parser_tokenizer): "model": PARSER_MODEL, "generate_response": _e2e_generate_response(output_ids), "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": {"model": PARSER_MODEL, "messages": messages}, }, ) assert resp.status_code == 200 - content = resp.json()["choices"][0]["message"]["content"] - re_encoded = _encode(parser_tokenizer, content) + reasoning = resp.json()["choices"][0]["message"]["reasoning"] + re_encoded = _encode(parser_tokenizer, reasoning) assert output_ids == re_encoded @pytest.mark.asyncio async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer): - """CJK + emoji roundtrip without U+FFFD.""" + """CJK + emoji roundtrip without U+FFFD. + + Markerless output comes back as reasoning (see + ``test_e2e_plain_roundtrip``). + """ messages = [{"role": "user", "content": "Reply in Chinese"}] gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) @@ -859,11 +875,12 @@ async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer): "model": PARSER_MODEL, "generate_response": _e2e_generate_response(output_ids), "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": {"model": PARSER_MODEL, "messages": messages}, }, ) assert resp.status_code == 200 - content = resp.json()["choices"][0]["message"]["content"] - assert "�" not in content + reasoning = resp.json()["choices"][0]["message"]["reasoning"] + assert "�" not in reasoning @pytest.mark.asyncio @@ -976,8 +993,10 @@ async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenize @pytest.mark.asyncio -async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer): - """Without chat_request, derender falls back to plain detokenization.""" +async def test_e2e_no_chat_request_rejected(parser_client, parser_tokenizer): + """Without chat_request a parser configured model rejects with 400 + rather than silently falling back to plain detokenization. This is to + prevent the leak of raw reasoning/tool markup into content.""" messages = [{"role": "user", "content": "Hello"}] gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) @@ -992,9 +1011,8 @@ async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer): "prompt_tokens": len(gen_req["token_ids"]), }, ) - assert resp.status_code == 200 - content = resp.json()["choices"][0]["message"]["content"] - assert "Hi" in content + assert resp.status_code == 400 + assert "chat_request" in resp.json()["error"]["message"] # --------------------------------------------------------------------------- diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index df10b31d2847..4032d9cbf308 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -8,6 +8,8 @@ 1. Unit tests (no server): covers ``_detokenize_delta`` correctness (chunked == one-shot) and ``derender_completion_stream`` / ``derender_chat_stream`` logic via a real tokenizer on a tiny model. + The parser path is covered both with a deterministic stub parser and + with the real ``HarmonyParser`` (skipped without ``openai_harmony``). 2. Integration tests (require a running render server): covers the full HTTP round-trip through the streaming endpoint. Marked with @@ -15,19 +17,142 @@ fixtures from the sibling ``test_derender.py``. """ +import json +from collections.abc import Callable + import pytest import pytest_asyncio +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, +) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderStreamState, + GenerateResponse, + GenerateResponseChoice, GenerateResponseStreamChoice, GenerateStreamResponse, ) from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo +from vllm.parser import Parser +from vllm.utils import random_uuid MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" +def _merge_fake_deltas( + accumulated: DeltaMessage | None, new: DeltaMessage | None +) -> DeltaMessage | None: + """Fold one token's worth of `_FakeParser` output into a per call + accumulator, the same way a real parser merges multiple tokens' + effects into the single `DeltaMessage` it returns per `parse_delta` + call. + """ + if new is None: + return accumulated + if accumulated is None: + accumulated = DeltaMessage() + + if new.content: + accumulated.content = (accumulated.content or "") + new.content + if new.reasoning: + accumulated.reasoning = (accumulated.reasoning or "") + new.reasoning + + for tc in new.tool_calls: + starts_new_call = tc.id is not None or ( + tc.function is not None and tc.function.name is not None + ) + existing = ( + None + if starts_new_call + else next((t for t in accumulated.tool_calls if t.index == tc.index), None) + ) + if existing is None: + accumulated.tool_calls.append(tc) + continue + if existing.function is None: + existing.function = DeltaFunctionCall() + if tc.function is not None and tc.function.arguments: + existing.function.arguments = ( + existing.function.arguments or "" + ) + tc.function.arguments + + return accumulated + + +class _FakeParser(Parser): + """Deterministic `Parser` stub for unit testing the replay/merge/pin + mechanics of `OnlineDerenderer._derender_chat_stream_parsed` in + isolation from any real reasoning/tool parser's markup grammar. + + Processes every id in `delta_token_ids` in order and folds their + effects into one `DeltaMessage`, like a real parser handling a + multi-token delta in a single call. Ignores `delta_text`: + + - `TOOL_START` opens a new tool call at index 0 (id + name). + - `TOOL_ARG` appends one `"a"` to that tool call's arguments. + - `REASON` emits one `"r"` of reasoning (suppressed when the request + has `include_reasoning=False`, mirroring the real parsers). + - `CONTENT` emits one `"c"` of content. + - An empty `delta_token_ids` with `finished=True` (the finish only + flush call) emits a `"FLUSH"` content sentinel so tests can confirm + it happened. + - Anything else emits nothing. + """ + + TOOL_START = 9 + TOOL_ARG = 10 + REASON = 11 + CONTENT = 12 + + def parse_delta( + self, + delta_text, + delta_token_ids, + request, + prompt_token_ids=None, + *, + finished, + ): + if not delta_token_ids: + return DeltaMessage(content="FLUSH") if finished else None + + accumulated: DeltaMessage | None = None + for tok in delta_token_ids: + accumulated = _merge_fake_deltas( + accumulated, self._parse_one_token(tok, request) + ) + return accumulated + + def _parse_one_token(self, tok: int, request) -> DeltaMessage | None: + if tok == self.TOOL_START: + return DeltaMessage( + tool_calls=[ + DeltaToolCall( + id=f"call-{random_uuid()}", + type="function", + function=DeltaFunctionCall(name="get_weather", arguments=""), + index=0, + ) + ] + ) + if tok == self.TOOL_ARG: + return DeltaMessage( + tool_calls=[ + DeltaToolCall(index=0, function=DeltaFunctionCall(arguments="a")) + ] + ) + if tok == self.REASON: + if not request.include_reasoning: + return None + return DeltaMessage(reasoning="r") + if tok == self.CONTENT: + return DeltaMessage(content="c") + return None + + # --------------------------------------------------------------------------- # Helpers shared across tests # --------------------------------------------------------------------------- @@ -111,6 +236,40 @@ def derenderer(tokenizer): ) +@pytest.fixture(scope="module") +def parsed_derenderer(tokenizer): + """OnlineDerenderer with `_FakeParser` wired in as `self.parser`. + + The parser configured path runs `_derender_chat_stream_parsed` via + `make_async(..., executor=renderer._executor)`, so unlike + `derenderer` above (whose streaming paths never touch the executor), + `renderer._executor` must be a real `ThreadPoolExecutor` here. + `loop.run_in_executor` cannot submit work to a `MagicMock`. + """ + from concurrent.futures import ThreadPoolExecutor + from unittest.mock import MagicMock + + from vllm.renderers.online_derenderer import OnlineDerenderer + + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + renderer._executor = ThreadPoolExecutor(max_workers=2) + + model_config = MagicMock() + model_config.hf_config.model_type = "llama" + model_config.model = MODEL_NAME + + dr = OnlineDerenderer( + model_config=model_config, + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + ) + dr.parser = _FakeParser + return dr + + class TestDetokenizeDelta: """_detokenize_delta: chunked decode must equal one shot decode.""" @@ -378,47 +537,454 @@ async def test_chunked_equals_oneshot(self, derenderer, tokenizer): assert streamed == one_shot @pytest.mark.asyncio - @pytest.mark.parametrize("with_chat_request", [True, False]) - async def test_parser_raises_not_implemented(self, tokenizer, with_chat_request): - """Stream chat with a parser active must fail closed (NotImplementedError). + async def test_parser_configured_missing_chat_request_raises( + self, parsed_derenderer + ): + """A parser configured model must never fall through to plain detok + even when `chat_request` is omitted. If allowed this would leak raw + reasoning/tool markup into `delta.content`. `ServingDerender` + pre-checks this too (400 before touching the tokenizer). This pins + the `OnlineDerenderer` level backstop.""" + with pytest.raises(ValueError, match="chat_request"): + await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]), + state=None, + chat_request=None, + ) - Checks that a parser configured model must 501 even when ``chat_request`` is - omitted, otherwise reasoning/tool markup would leak into ``delta.content` - via the plain detok fallback. - """ - from unittest.mock import MagicMock - from vllm.renderers.online_derenderer import OnlineDerenderer +def _chat_request(**kwargs): + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + + kwargs.setdefault("messages", [{"role": "user", "content": "hi"}]) + kwargs.setdefault("model", MODEL_NAME) + return ChatCompletionRequest(**kwargs) - renderer = MagicMock() - renderer.get_tokenizer.return_value = tokenizer - model_config = MagicMock() - model_config.hf_config.model_type = "llama" - model_config.model = MODEL_NAME +class TestDerenderChatStreamParsed: + """derender_chat_stream: parser branch (replay + parse_delta) exercised + against the deterministic `_FakeParser` so these pin OnlineDerenderer's + own replay/merge/pin/finish_reason logic independent of any real + parser's markup grammar (covered separately by the parser_server backed + integration tests below).""" - # Construct a derenderer WITH a tool/reasoning parser active. - # ParserManager.get_parser returns None when no parser is named, so - # we inject a mock parser class directly. - dr = OnlineDerenderer( - model_config=model_config, - renderer=renderer, - request_logger=None, - chat_template=None, - chat_template_content_format="string", + @pytest.mark.asyncio + async def test_dispatches_and_emits_content(self, parsed_derenderer): + chunk, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.CONTENT], finish_reason="stop" + ), + chat_request=_chat_request(), ) - dr.parser = MagicMock() # simulate a parser being active + assert chunk.choices[0].delta.content == "c" + assert chunk.choices[0].delta.role == "assistant" + assert chunk.choices[0].finish_reason == "stop" + assert state.output_token_ids == [_FakeParser.CONTENT] - chat_request = MagicMock() if with_chat_request else None - token_ids = tokenizer.encode("hello")[:3] + @pytest.mark.asyncio + async def test_role_sent_once(self, parsed_derenderer): + chat_request = _chat_request() + chunk1, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]), + chat_request=chat_request, + ) + chunk2, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.CONTENT], finish_reason="stop" + ), + state=state, + chat_request=chat_request, + ) + assert chunk1.choices[0].delta.role == "assistant" + assert chunk2.choices[0].delta.role is None - with pytest.raises(NotImplementedError): - await dr.derender_chat_stream( - model=MODEL_NAME, - generate_chunk=_make_stream_chunk(token_ids), - state=None, - chat_request=chat_request, + @pytest.mark.asyncio + async def test_finish_only_chunk_flushes(self, parsed_derenderer): + """A finish only chunk (no new tokens) still calls `parse_delta` + once with `finished=True` to flush buffered state.""" + chat_request = _chat_request() + _, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]), + chat_request=chat_request, + ) + chunk2, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([], finish_reason="stop"), + state=state, + chat_request=chat_request, + ) + assert chunk2.choices[0].delta.content == "FLUSH" + assert chunk2.choices[0].finish_reason == "stop" + + @pytest.mark.asyncio + async def test_chunking_invariance(self, parsed_derenderer): + """1 token per chunk vs a single whole chunk assemble identically. + + Every chunk (live or replayed) is fed to `parse_delta` at its own + producer granularity (see `_derender_chat_stream_parsed`'s + docstring). This proves that choice doesn't change the assembled + output for a well behaved incremental parser. Each call's effects + fold together the same way regardless of how many tokens land in + one call vs. many.""" + token_ids = [ + _FakeParser.REASON, + _FakeParser.REASON, + _FakeParser.TOOL_START, + _FakeParser.TOOL_ARG, + _FakeParser.TOOL_ARG, + _FakeParser.CONTENT, + ] + chat_request = _chat_request( + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tool_choice="auto", + include_reasoning=True, + ) + + async def _assemble(chunks: list[list[int]]) -> dict: + state = None + content = "" + reasoning = "" + tool_args = "" + for i, tids in enumerate(chunks): + finish = "stop" if i == len(chunks) - 1 else None + chunk, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(tids, finish_reason=finish), + state=state, + chat_request=chat_request, + ) + delta = chunk.choices[0].delta + content += delta.content or "" + reasoning += delta.reasoning or "" + for tc in delta.tool_calls: + if tc.function and tc.function.arguments: + tool_args += tc.function.arguments + return {"content": content, "reasoning": reasoning, "tool_args": tool_args} + + whole = await _assemble([token_ids]) + one_at_a_time = await _assemble([[t] for t in token_ids]) + + assert whole == one_at_a_time + assert whole == {"content": "c", "reasoning": "rr", "tool_args": "aa"} + + @pytest.mark.asyncio + async def test_replay_preserves_chunk_boundaries( + self, parsed_derenderer, monkeypatch + ): + """A multi-token chunk (e.g. a speculative decoding step) must be + replayed as one `parse_delta` call, not one per token. Parsers + with per-call state such as `history_tool_call_cnt` would + otherwise diverge from standard serving.""" + + class _CallCountingParser(_FakeParser): + """Emits how many `parse_delta` calls it has seen.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.num_calls = 0 + + def parse_delta( + self, + delta_text, + delta_token_ids, + request, + prompt_token_ids=None, + *, + finished, + ): + self.num_calls += 1 + return DeltaMessage(content=str(self.num_calls)) + + monkeypatch.setattr(parsed_derenderer, "parser", _CallCountingParser) + chat_request = _chat_request() + + _, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([_FakeParser.CONTENT] * 3), + chat_request=chat_request, + ) + assert state.output_chunk_lens == [3] + + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.CONTENT], finish_reason="stop" + ), + state=state, + chat_request=chat_request, + ) + assert chunk.choices[0].delta.content == "2" + + @pytest.mark.asyncio + async def test_tool_call_id_pinned_across_chunks(self, parsed_derenderer): + """Once an index's ID is recorded in `last_tool_call_ids`, a later + id bearing delta for that same index is pinned to the recorded + value rather than a freshly (re-)generated one.""" + chat_request = _chat_request( + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tool_choice="auto", + ) + chunk1, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([_FakeParser.TOOL_START]), + chat_request=chat_request, + ) + first_id = chunk1.choices[0].delta.tool_calls[0].id + assert first_id is not None + assert state.last_tool_call_ids == [first_id] + + chunk2, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.TOOL_START], finish_reason="stop" + ), + state=state, + chat_request=chat_request, + ) + assert chunk2.choices[0].delta.tool_calls[0].id == first_id + assert state.last_tool_call_ids == [first_id] + + @pytest.mark.asyncio + async def test_finish_reason_rewritten_to_tool_calls(self, parsed_derenderer): + chat_request = _chat_request( + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tool_choice="auto", + ) + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.TOOL_START, _FakeParser.TOOL_ARG], finish_reason="stop" + ), + chat_request=chat_request, + ) + assert chunk.choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_finish_reason_stop_for_named_tool_choice(self, parsed_derenderer): + # ChatCompletionRequest's `check_tool_usage` is a mode="before" + # validator, so it sees the raw value and only accepts "auto", + # "required" or a dict, never an already built + # ChatCompletionNamedToolChoiceParam. + chat_request = _chat_request( + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + ) + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.TOOL_START, _FakeParser.TOOL_ARG], finish_reason="stop" + ), + chat_request=chat_request, + ) + assert chunk.choices[0].finish_reason == "stop" + + @pytest.mark.asyncio + async def test_include_reasoning_false_suppresses_reasoning( + self, parsed_derenderer + ): + chat_request = _chat_request(include_reasoning=False) + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [_FakeParser.REASON, _FakeParser.CONTENT], finish_reason="stop" + ), + chat_request=chat_request, + ) + delta = chunk.choices[0].delta + assert delta.reasoning is None + assert delta.content == "c" + + +# --------------------------------------------------------------------------- +# Harmony / GPT-OSS replay — unit, no server +# --------------------------------------------------------------------------- + +HARMONY_MODEL = "openai/gpt-oss-20b" +HARMONY_REASONING = "The user wants 2 plus 3." +HARMONY_ANSWER = "The answer is 5." +HARMONY_PROMPT = "<|start|>user<|message|>Add 2 and 3.<|end|><|start|>assistant" +HARMONY_OUTPUT = ( + f"<|channel|>analysis<|message|>{HARMONY_REASONING}<|end|>" + f"<|start|>assistant<|channel|>final<|message|>{HARMONY_ANSWER}<|return|>" +) + + +@pytest.fixture(scope="module") +def harmony_encode(): + """Encoder for canned GPT-OSS harmony token sequences.""" + pytest.importorskip("openai_harmony") + + # Pre-caches the o200k_base BPE file that openai-harmony's Rust backend + # downloads on first use, same as the sibling suite's E2E harmony tests. + from tests.entrypoints.scale_out.derender.test_derender import ( + _ensure_harmony_vocab, + ) + from vllm.entrypoints.openai.parser.harmony_utils import get_encoding + + _ensure_harmony_vocab() + encoding = get_encoding() + + def _encode(harmony_str: str) -> list[int]: + return encoding.encode(harmony_str, allowed_special="all") + + return _encode + + +@pytest.fixture(scope="module") +def harmony_tokenizer(): + pytest.importorskip("openai_harmony") + from vllm.tokenizers import get_tokenizer + + return get_tokenizer(HARMONY_MODEL, trust_remote_code=True) + + +@pytest.fixture(scope="module") +def harmony_derenderer(harmony_tokenizer): + """OnlineDerenderer whose parser resolves to the real `HarmonyParser`. + + Same shape as `parsed_derenderer` (mocked renderer, real executor) but + with a real tokenizer and parser, so the replay path is exercised + against Harmony's actual channel grammar rather than a stub. + """ + from concurrent.futures import ThreadPoolExecutor + from unittest.mock import MagicMock + + from vllm.parser.harmony import HarmonyParser + from vllm.renderers.online_derenderer import OnlineDerenderer + + renderer = MagicMock() + renderer.get_tokenizer.return_value = harmony_tokenizer + renderer._executor = ThreadPoolExecutor(max_workers=2) + + model_config = MagicMock() + model_config.model = HARMONY_MODEL + model_config.hf_config.model_type = "gpt_oss" + model_config.hf_text_config.model_type = "gpt_oss" + model_config.hf_overrides = None + + dr = OnlineDerenderer( + model_config=model_config, + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + enable_auto_tools=True, + tool_parser="openai", + reasoning_parser="openai_gptoss", + ) + assert dr.use_harmony + assert dr.parser is HarmonyParser + return dr + + +async def _stream_harmony_deltas( + derenderer, + chat_request, + output_ids: list[int], + prompt_ids: list[int], + chunk_size: int, +) -> list[DeltaMessage]: + """Drive `output_ids` through the parser path in fixed size chunks. + + Returns the per chunk `DeltaMessage`s in order so callers can assert on + intermediate emissions, not just the assembled result. + """ + deltas: list[DeltaMessage] = [] + state = None + for start in range(0, len(output_ids), chunk_size): + tids = output_ids[start : start + chunk_size] + is_last = start + chunk_size >= len(output_ids) + chunk, state = await derenderer.derender_chat_stream( + model=HARMONY_MODEL, + generate_chunk=_make_stream_chunk( + tids, finish_reason="stop" if is_last else None + ), + state=state, + chat_request=chat_request, + prompt_token_ids=prompt_ids, + ) + deltas.append(chunk.choices[0].delta) + return deltas + + +class TestDerenderChatStreamHarmony: + """derender_chat_stream: replay + `parse_delta` against real HarmonyParser. + + Skipped where `openai_harmony` is not installed. `prompt_token_ids` is + threaded through for call shape fidelity only, since HarmonyParser + derives its state from the output tokens alone. + """ + + @pytest.mark.asyncio + async def test_reasoning_never_leaks_as_content_midstream( + self, harmony_derenderer, harmony_encode + ): + """Regression guard for the RFC's original replay + diff design. + + `HarmonyParser.parse()` always flushes to EOS. Mid-stream that + raises `HarmonyError` and the recovery branch re-emits the in + flight message on the `final` channel, so partial analysis surfaces + as content. Replay + `parse_delta` must never do that. Driven one + token per chunk, the finest granularity a client can produce. + """ + output_ids = harmony_encode(HARMONY_OUTPUT) + prompt_ids = harmony_encode(HARMONY_PROMPT) + chat_request = _chat_request(model=HARMONY_MODEL, include_reasoning=True) + + deltas = await _stream_harmony_deltas( + harmony_derenderer, chat_request, output_ids, prompt_ids, chunk_size=1 + ) + + reasoning = "" + content = "" + for i, delta in enumerate(deltas): + reasoning += delta.reasoning or "" + content += delta.content or "" + # Any analysis text emitted as content breaks the prefix + # property, since the two channels share no prefix here. + assert HARMONY_ANSWER.startswith(content), ( + f"content after chunk {i} is not a prefix of the final " + f"channel text: {content!r}" ) + assert HARMONY_REASONING.startswith(reasoning) + assert not delta.tool_calls + + assert reasoning == HARMONY_REASONING + assert content == HARMONY_ANSWER + + @pytest.mark.asyncio + async def test_stream_matches_batch(self, harmony_derenderer, harmony_encode): + """Streamed assembly equals one shot `/derender` over the same IDs.""" + output_ids = harmony_encode(HARMONY_OUTPUT) + prompt_ids = harmony_encode(HARMONY_PROMPT) + chat_request = _chat_request(model=HARMONY_MODEL, include_reasoning=True) + + deltas = await _stream_harmony_deltas( + harmony_derenderer, chat_request, output_ids, prompt_ids, chunk_size=3 + ) + reasoning = "".join(d.reasoning or "" for d in deltas) + content = "".join(d.content or "" for d in deltas) + + batch_choices = await harmony_derenderer.derender_chat( + GenerateResponse( + request_id="test-harmony-batch", + choices=[ + GenerateResponseChoice( + index=0, token_ids=output_ids, finish_reason="stop" + ) + ], + ), + chat_request, + ) + message = batch_choices[0].message + + assert reasoning == message.reasoning + assert content == message.content class TestDerenderStreamStateValidation: @@ -446,6 +1012,24 @@ def test_prev_tokens_at_cap_accepted(self): state = DerenderStreamState(prev_tokens=["a"] * 1024) assert len(state.prev_tokens) == 1024 + def test_output_chunk_lens_mismatch_rejected(self): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + DerenderStreamState(output_token_ids=[1, 2, 3], output_chunk_lens=[1, 1]) + + def test_output_chunk_lens_zero_entry_rejected(self): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + DerenderStreamState(output_token_ids=[1, 2], output_chunk_lens=[0, 2]) + + def test_output_chunk_lens_matching_accepted(self): + state = DerenderStreamState( + output_token_ids=[1, 2, 3], output_chunk_lens=[1, 2] + ) + assert state.output_chunk_lens == [1, 2] + class TestServingDerenderStreamErrorHandling: """Malformed stream_state must surface as 400 and not an unhandled 500.""" @@ -458,8 +1042,10 @@ def _make_serving(self, side_effect: Exception): models = MagicMock() models.is_base_model.return_value = True models.model_config = MagicMock() + models.model_config.max_model_len = 100_000 online_derenderer = MagicMock() + online_derenderer.parser = None online_derenderer.derender_completion_stream = AsyncMock( side_effect=side_effect ) @@ -504,6 +1090,179 @@ async def test_chat_stream_bad_state_returns_400(self, exc): assert result.error.code == 400 +class TestServingDerenderStreamValidation: + """If checks cannot run in `derender_chat_stream_response`: all must + reject with 400 before ever calling into `online_derenderer` (i.e. + before touching the tokenizer).""" + + def _make_serving(self, *, parser_configured: bool, max_model_len: int = 100_000): + from unittest.mock import AsyncMock, MagicMock + + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, + ) + from vllm.entrypoints.scale_out.derender.serving import ServingDerender + + models = MagicMock() + models.is_base_model.return_value = True + models.model_config = MagicMock() + models.model_config.max_model_len = max_model_len + + online_derenderer = MagicMock() + online_derenderer.parser = MagicMock() if parser_configured else None + online_derenderer.derender_chat_stream = AsyncMock( + return_value=( + ChatCompletionStreamResponse(id="t", model=MODEL_NAME, choices=[]), + DerenderStreamState(), + ) + ) + return ServingDerender(models=models, online_derenderer=online_derenderer) + + @pytest.mark.asyncio + async def test_missing_chat_request_with_parser_rejected(self): + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=True) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + assert "chat_request" in result.error.message + serving.online_derenderer.derender_chat_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_missing_chat_request_without_parser_ok(self): + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=False) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + ) + result = await serving.derender_chat_stream_response(request) + assert not isinstance(result, ErrorResponse) + + @pytest.mark.asyncio + async def test_missing_prompt_token_ids_with_parser_rejected(self): + """A parser configured model must reject a missing prompt_token_ids + the same way it rejects a missing chat_request. Without it, + parse_delta cannot tell whether the prompt left reasoning open and + would misclassify reasoning content as plain content.""" + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=True) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + chat_request=_chat_request(), + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + assert "prompt_token_ids" in result.error.message + serving.online_derenderer.derender_chat_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_prompt_token_ids_present_with_parser_ok(self): + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=True) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2]), + chat_request=_chat_request(), + prompt_token_ids=[1, 2, 3], + ) + result = await serving.derender_chat_stream_response(request) + assert not isinstance(result, ErrorResponse) + + @pytest.mark.asyncio + async def test_oversized_output_token_ids_rejected(self): + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=False, max_model_len=4) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1, 2, 3]), + stream_state=DerenderStreamState( + output_token_ids=[1, 2], output_chunk_lens=[2] + ), + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + assert "max_model_len" in result.error.message + serving.online_derenderer.derender_chat_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_oversized_prompt_token_ids_rejected(self): + """prompt_token_ids is caller supplied and otherwise unbounded. A + parser configured deployment rescans it in full on every chunk + (is_reasoning_end / adjust_initial_state_from_prompt), so it must be + bounded the same way output_token_ids is.""" + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=True, max_model_len=4) + request = DerenderChatStreamRequest( + stream=True, + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([1]), + chat_request=_chat_request(), + prompt_token_ids=[1, 2, 3, 4, 5], + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + assert "max_model_len" in result.error.message + serving.online_derenderer.derender_chat_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_too_many_choices_rejected(self): + """Each streamed chunk contains at most one choice (a single + DerenderStreamState is threaded through every choice). The check + could never fire since derender_chat_stream itself rejects anything + above 1 first.""" + from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatStreamRequest, + ) + + serving = self._make_serving(parser_configured=False) + two_choices = GenerateStreamResponse( + request_id="t", + choices=[ + GenerateResponseStreamChoice(index=i, token_ids=[1]) for i in range(2) + ], + ) + request = DerenderChatStreamRequest( + stream=True, model=MODEL_NAME, generate_chunk=two_choices + ) + result = await serving.derender_chat_stream_response(request) + assert isinstance(result, ErrorResponse) + assert result.error.code == 400 + assert "at most one choice" in result.error.message + serving.online_derenderer.derender_chat_stream.assert_not_called() + + # --------------------------------------------------------------------------- # Integration tests — require a live render server # --------------------------------------------------------------------------- @@ -757,3 +1516,416 @@ async def test_streaming_usage_chunk(client): assert d2["chunk"]["choices"] == [] assert d2["chunk"]["usage"]["prompt_tokens"] == 10 assert d2["chunk"]["usage"]["completion_tokens"] == len(token_ids) + + +# --------------------------------------------------------------------------- +# Integration tests for parser configured (reasoning + tool calls) require a +# live render server. Mirrors the parser_server / parser_tokenizer pattern +# from test_derender.py. +# --------------------------------------------------------------------------- + +PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +_PARSER_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } +] + + +@pytest.fixture(scope="module") +def parser_server(): + from tests.utils import RemoteLaunchRenderServer + + args = [ + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + "--reasoning-parser", + "deepseek_r1", + ] + with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def parser_client(parser_server): + import httpx + + async with httpx.AsyncClient( + base_url=parser_server.url_for(""), timeout=60.0 + ) as http_client: + yield http_client + + +@pytest.fixture(scope="module") +def parser_tokenizer(): + from vllm.tokenizers import get_tokenizer + + return get_tokenizer(PARSER_MODEL) + + +def _require_parser_markers(tokenizer, text: str, *markers: str) -> list[int]: + """Encode text and skip the test if any marker is lost in roundtrip.""" + ids = tokenizer.encode(text, add_special_tokens=False) + decoded = tokenizer.decode(ids, skip_special_tokens=False) + for m in markers: + if m not in decoded: + pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip") + return ids + + +async def _render_parser_chat(client, messages: list[dict]) -> dict: + resp = await client.post( + "/v1/chat/completions/render", + json={"model": PARSER_MODEL, "messages": messages}, + ) + assert resp.status_code == 200, resp.text + return resp.json() + + +async def _stream_chat_derender( + client, + output_ids: list[int], + chunk_sizes: list[int], + chat_request: dict, + prompt_tokens: int, + prompt_token_ids: list[int], + on_chunk: Callable[[list[dict]], None] | None = None, +) -> dict: + """Feed `output_ids` through the streaming chat derender endpoint in + the given `chunk_sizes`, threading `stream_state` across calls and + return the assembled message. + + If `on_chunk` is given, it is called after every chunk with a + snapshot (deep copy) of the `tool_calls` accumulator so far, letting + callers assert properties of the intermediate deltas (e.g. monotonic + argument growth) rather than only the final assembled result. + """ + state = None + content = "" + reasoning = "" + tool_calls: list[dict] = [] + finish_reason = None + + pos = 0 + for i, size in enumerate(chunk_sizes): + tids = output_ids[pos : pos + size] + pos += size + is_last = i == len(chunk_sizes) - 1 + resp = await client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": PARSER_MODEL, + "generate_chunk": { + "request_id": "stream-test", + "choices": [ + { + "index": 0, + "token_ids": tids, + "finish_reason": "stop" if is_last else None, + } + ], + }, + "stream_state": state, + "prompt_tokens": prompt_tokens, + "prompt_token_ids": prompt_token_ids, + "chat_request": chat_request, + }, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + state = data["stream_state"] + delta = data["chunk"]["choices"][0]["delta"] + content += delta.get("content") or "" + reasoning += delta.get("reasoning") or "" + for tc in delta.get("tool_calls") or []: + idx = tc["index"] + while len(tool_calls) <= idx: + tool_calls.append({"id": None, "name": None, "arguments": ""}) + if tc.get("id"): + tool_calls[idx]["id"] = tc["id"] + fn = tc.get("function") or {} + if fn.get("name"): + tool_calls[idx]["name"] = fn["name"] + if fn.get("arguments"): + tool_calls[idx]["arguments"] += fn["arguments"] + if is_last: + finish_reason = data["chunk"]["choices"][0]["finish_reason"] + if on_chunk is not None: + on_chunk([dict(tc) for tc in tool_calls]) + + return { + "content": content or None, + "reasoning": reasoning or None, + "tool_calls": tool_calls, + "finish_reason": finish_reason, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunking", + [ + "whole", + "per_token", + pytest.param( + "triples", + marks=pytest.mark.xfail( + reason=" leaks into reasoning when it shares a " + "multi-token delta with reasoning text (vllm#55195)" + ), + ), + ], +) +async def test_stream_parsed_matches_batch_reasoning( + parser_client, parser_tokenizer, chunking +): + """Streamed == batch: assembled reasoning/content equal the non + streaming `/derender` result over the same token IDs, for every + chunking of the same output (chunking invariance).""" + messages = [{"role": "user", "content": "What is 2+3?"}] + gen_req = await _render_parser_chat(parser_client, messages) + + reasoning_text = "The user wants 2 plus 3. That is 5." + answer_text = "The answer is 5." + output_text = f"{reasoning_text}{answer_text}" + output_ids = _require_parser_markers(parser_tokenizer, output_text, "") + + chat_request = { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": True, + } + + batch_resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": { + "request_id": "batch", + "choices": [ + {"index": 0, "token_ids": output_ids, "finish_reason": "stop"} + ], + }, + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": chat_request, + }, + ) + assert batch_resp.status_code == 200, batch_resp.text + batch_msg = batch_resp.json()["choices"][0]["message"] + + n = len(output_ids) + chunk_sizes = { + "whole": [n], + "per_token": [1] * n, + "triples": [3] * (n // 3) + [n % 3], + }[chunking] + chunk_sizes = [c for c in chunk_sizes if c > 0] + streamed = await _stream_chat_derender( + parser_client, + output_ids, + chunk_sizes, + chat_request, + len(gen_req["token_ids"]), + gen_req["token_ids"], + ) + assert streamed["content"] == batch_msg["content"] + assert streamed["reasoning"] == batch_msg.get("reasoning") + + +@pytest.mark.asyncio +async def test_stream_parsed_matches_batch_tool_call(parser_client, parser_tokenizer): + """Tool call name+id+arguments parity between streamed and batch, plus + the finish_reason -> "tool_calls" rewrite for auto tool choice.""" + messages = [{"role": "user", "content": "Weather in Paris?"}] + gen_req = await _render_parser_chat(parser_client, messages) + + output_text = ( + "Let me check." + '\n{"name": "get_weather", ' + '"arguments": {"city": "Paris"}}\n' + ) + output_ids = _require_parser_markers( + parser_tokenizer, output_text, "", "", "" + ) + chat_request = { + "model": PARSER_MODEL, + "messages": messages, + "tools": _PARSER_TOOLS, + "tool_choice": "auto", + } + + batch_resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": { + "request_id": "batch", + "choices": [ + {"index": 0, "token_ids": output_ids, "finish_reason": "stop"} + ], + }, + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": chat_request, + }, + ) + assert batch_resp.status_code == 200, batch_resp.text + batch_choice = batch_resp.json()["choices"][0] + if not batch_choice["message"]["tool_calls"]: + pytest.skip("Model did not emit a tool call") + + for chunk_sizes in ([len(output_ids)], [1] * len(output_ids)): + # Snapshot of tool_calls[0]["arguments"] after every chunk, to check + # it only ever grows by appending. It never retracts or duplicates + # already streamed text. That's the exact failure mode the + # replay + parse_delta design (over diff-based streaming) exists to + # avoid. Parity with the batch result alone wouldn't catch a + # transient mid-stream regression. + arg_snapshots: list[str] = [] + + def _record( + tool_calls: list[dict], _snapshots: list[str] = arg_snapshots + ) -> None: + if tool_calls and tool_calls[0]["arguments"]: + _snapshots.append(tool_calls[0]["arguments"]) + + streamed = await _stream_chat_derender( + parser_client, + output_ids, + chunk_sizes, + chat_request, + len(gen_req["token_ids"]), + gen_req["token_ids"], + on_chunk=_record, + ) + assert streamed["tool_calls"] + assert streamed["tool_calls"][0]["name"] == "get_weather" + assert streamed["tool_calls"][0]["id"] is not None + assert json.loads(streamed["tool_calls"][0]["arguments"]) == json.loads( + batch_choice["message"]["tool_calls"][0]["function"]["arguments"] + ) + assert streamed["finish_reason"] == "tool_calls" + + assert arg_snapshots, "expected at least one tool-call argument delta" + for prev, curr in zip(arg_snapshots, arg_snapshots[1:]): + assert curr.startswith(prev), ( + f"tool-call arguments regressed: {prev!r} -> {curr!r}" + ) + assert arg_snapshots[-1] == streamed["tool_calls"][0]["arguments"] + + +@pytest.mark.asyncio +async def test_stream_parsed_cjk_across_chunk_boundaries( + parser_client, parser_tokenizer +): + """CJK/emoji tokens split across chunk boundaries with a parser active + across the reasoning -> content transition (parser path regression + guard analogous to #46159 in the plain detok path).""" + messages = [{"role": "user", "content": "Reply in Chinese"}] + gen_req = await _render_parser_chat(parser_client, messages) + + reasoning_text = "思考中" + answer_text = "你好世界 😀" + output_text = f"{reasoning_text}{answer_text}" + output_ids = _require_parser_markers(parser_tokenizer, output_text, "") + + chat_request = { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": True, + } + streamed = await _stream_chat_derender( + parser_client, + output_ids, + [1] * len(output_ids), + chat_request, + len(gen_req["token_ids"]), + gen_req["token_ids"], + ) + assert "�" not in (streamed["content"] or "") + assert "�" not in (streamed["reasoning"] or "") + assert answer_text in (streamed["content"] or "") + + +@pytest.mark.asyncio +async def test_stream_parsed_include_reasoning_false(parser_client, parser_tokenizer): + """include_reasoning=False emits no reasoning deltas on the streaming + parser path.""" + messages = [{"role": "user", "content": "What is 2+3?"}] + gen_req = await _render_parser_chat(parser_client, messages) + + output_text = "reasoning hereThe answer is 5." + output_ids = _require_parser_markers(parser_tokenizer, output_text, "") + + chat_request = { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": False, + } + streamed = await _stream_chat_derender( + parser_client, + output_ids, + [1] * len(output_ids), + chat_request, + len(gen_req["token_ids"]), + gen_req["token_ids"], + ) + assert streamed["reasoning"] is None + + +@pytest.mark.asyncio +async def test_stream_parsed_missing_chat_request_rejected(parser_client): + """Parser configured + no chat_request on the streaming endpoint -> 400 + (the live-server counterpart to the mocked + `TestServingDerenderStreamValidation` checks above).""" + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": PARSER_MODEL, + "generate_chunk": { + "request_id": "reject-test", + "choices": [{"index": 0, "token_ids": [1, 2], "finish_reason": None}], + }, + }, + ) + assert resp.status_code == 400 + assert "chat_request" in resp.json()["error"]["message"] + + +@pytest.mark.asyncio +async def test_stream_parsed_missing_prompt_token_ids_rejected(parser_client): + """Parser configured + chat_request given but no prompt_token_ids on the + streaming endpoint -> 400 (the live-server counterpart to the mocked + `test_missing_prompt_token_ids_with_parser_rejected` above). Without + prompt_token_ids, parse_delta cannot tell whether the prompt left + reasoning open and would silently misclassify reasoning as content.""" + messages = [{"role": "user", "content": "Hello"}] + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": PARSER_MODEL, + "generate_chunk": { + "request_id": "reject-test", + "choices": [{"index": 0, "token_ids": [1, 2], "finish_reason": None}], + }, + "chat_request": { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": True, + }, + }, + ) + assert resp.status_code == 400 + assert "prompt_token_ids" in resp.json()["error"]["message"] diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index 81e919d840b7..d3c5f3aac1cb 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -138,6 +138,13 @@ async def derender_chat_response( if bounds_error is not None: return bounds_error + if self.online_derenderer.parser is not None and request.chat_request is None: + return self.create_error_response( + "chat_request is required when a tool or reasoning parser is " + "configured because plain detokenization would leak raw parser " + "markup into content." + ) + try: choices = await self.online_derenderer.derender_chat( request.generate_response, request.chat_request @@ -248,14 +255,67 @@ async def derender_chat_stream_response( Processes one ``GenerateStreamResponse`` chunk and returns the derendered chunk together with the updated client carried state. - - ``parser is None`` or no ``chat_request`` until reasoning/tool call - functionality added in future PR. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: return error_check_ret + if self.online_derenderer.parser is not None and request.chat_request is None: + return self.create_error_response( + "chat_request is required when a tool or reasoning parser is " + "configured because plain detokenization would leak raw parser " + "markup into content." + ) + + if ( + self.online_derenderer.parser is not None + and request.prompt_token_ids is None + ): + return self.create_error_response( + "prompt_token_ids is required when a tool or reasoning " + "parser is configured. Without it parse_delta cannot tell " + "whether the prompt left reasoning open (e.g. a chat " + "template that pre-opens ) and would misclassify " + "reasoning content as plain content." + ) + + # Each streamed chunk contains at most one choice per SSE event. + # A single DerenderStreamState is threaded + # through every choice in the chunk, so >1 would corrupt detok/parser + # state across choices. See the matching check in + # OnlineDerenderer.derender_chat_stream which this fails ahead of to + # avoid touching the tokenizer at all on a malformed chunk. + num_choices = len(request.generate_chunk.choices) + if num_choices > 1: + return self.create_error_response( + f"derender_chat_stream expects at most one choice per chunk " + f"(got {num_choices})." + ) + + # prompt_token_ids is folded in here too. It is caller supplied and + # otherwise unbounded. A parser configured deployment rescans it + # in full on every chunk (is_reasoning_end / + # adjust_initial_state_from_prompt, see abstract_parser.parse_delta) + # since parser state cannot be carried across calls. + prompt_len = ( + len(request.prompt_token_ids) if request.prompt_token_ids is not None else 0 + ) + prior_len = ( + len(request.stream_state.output_token_ids) + if request.stream_state is not None + else 0 + ) + delta_len = sum( + len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids + ) + max_model_len = self.model_config.max_model_len + if prompt_len + prior_len + delta_len > max_model_len: + return self.create_error_response( + f"prompt_token_ids length ({prompt_len}) plus " + f"output_token_ids length ({prior_len}) plus delta " + f"({delta_len}) exceeds max_model_len ({max_model_len})." + ) + model_name = request.model or self.models.model_name() try: chunk, updated_state = await self.online_derenderer.derender_chat_stream( @@ -264,15 +324,21 @@ async def derender_chat_stream_response( state=request.stream_state, chat_request=request.chat_request, prompt_tokens=request.prompt_tokens, + prompt_token_ids=request.prompt_token_ids, ) - except NotImplementedError as exc: - return self.create_error_response(exc) - except ValueError as exc: - return self.create_error_response(str(exc)) - except (KeyError, IndexError) as exc: - return self.create_error_response( - f"invalid stream_state: detokenization failed ({exc!r})" + except Exception as exc: + # The two ValueErrors derender_chat_stream can raise directly + # (missing chat_request, >1 choice per chunk) are already + # pre-checked above, so anything reaching here comes from the + # parser replaying entirely client controlled state. Hermes + # swallows its own streaming exceptions but finalize_generation + # and engine based parsers don't, so an unexpected parser + # failure on malformed (but well typed) state must surface as + # a 400 here rather than an unhandled 500. + logger.debug( + "derender stream replay failed: %s: %s", type(exc).__name__, exc ) + return self.create_error_response("invalid stream_state: derender failed") logger.debug( "derender_chat_stream request_id=%s model=%s delta_tokens=%d", @@ -313,7 +379,7 @@ async def derender_completion_stream_response( return self.create_error_response(str(exc)) except (KeyError, IndexError) as exc: return self.create_error_response( - f"invalid stream_state: detokenization failed ({exc!r})" + f"invalid stream_state: detokenization failed ({exc})" ) logger.debug( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index be8130e314cd..2bd47b105978 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any, Literal, TypeAlias +from typing import Annotated, Any, Literal, TypeAlias from pydantic import ( BaseModel, @@ -483,16 +483,32 @@ class DerenderStreamState(BaseModel): streaming derender endpoint. All fields are plain JSON serializable data. No opaque tokenizer or parser internals are stored here. + Two separate sets of fields support two different streaming modes: + + - For plain streaming (no parser configured including the completions path): + `prev_tokens`, `prefix_offset` and `read_offset` maintain a bounded incremental + decoding window. This requires O(window) transport and O(delta) computation + per chunk. + - For parser enabled chat streaming: `output_token_ids`, `output_chunk_lens`, + `tools_streamed` and `last_tool_call_ids` are used to replay `parse_delta()` + from scratch on every chunk because parser state cannot be serialized. This + incurs O(n) transport per chunk (O(n²) per generation) and O(n²) + `parse_delta()` calls per generation. Since many parsers re-scan the + entire accumulated text on each invocation, `parse_delta()` itself is + O(n) yielding a true worst case compute cost of O(n³) per generation. + No caching is performed. Work is bounded by `max_model_len`. + See `OnlineDerenderer._derender_chat_stream_parsed`. + The detokenization strategy carries the incremental decode offsets directly rather than re-sending the whole token history each chunk. - ``detokenize_incrementally`` only ever reads the trailing token window - ``prev_tokens[prefix_offset:]``, so we carry just that tail plus the two + `detokenize_incrementally` only ever reads the trailing token window + `prev_tokens[prefix_offset:]`, so we carry just that tail plus the two offsets. Each chunk resumes exactly where the last one stopped, including - any partially processed multi-byte character (tracked by ``read_offset``), + any partially processed multi-byte character (tracked by `read_offset`), then trims and rebases the window so it never grows with generation length. Performance: - - Compute per chunk is O(delta). One ``detokenize_incrementally`` call per + - Compute per chunk is O(delta). One `detokenize_incrementally` call per new token, independent of how many tokens preceded it. - Transport per chunk is O(window). The carried tail is bounded by the incremental detokenization offset, so cumulative bytes over the wire are @@ -500,18 +516,18 @@ class DerenderStreamState(BaseModel): """ prev_tokens: list[str] = Field(default_factory=list) - """Trailing decode window. Token strings from ``prefix_offset`` onward. + """Trailing decode window. Token strings from `prefix_offset` onward. Bounded, trimmed and rebased each chunk to the tail - ``detokenize_incrementally`` still reads, so it does not grow with the + `detokenize_incrementally` still reads, so it does not grow with the number of chunks. """ prefix_offset: int = Field(default=0, ge=0) - """Prefix offset into ``prev_tokens`` for incremental detokenization.""" + """Prefix offset into `prev_tokens` for incremental detokenization.""" read_offset: int = Field(default=0, ge=0) - """Read offset into ``prev_tokens`` for incremental detokenization.""" + """Read offset into `prev_tokens` for incremental detokenization.""" @field_validator("prev_tokens") @classmethod @@ -525,23 +541,53 @@ def _bound_prev_tokens(cls, v: list[str]) -> list[str]: return v role_sent: bool = False - """True once the initial ``role: "assistant"`` delta has been emitted. + """True once the initial `role: "assistant"` delta has been emitted. Prevents re-emitting the role on subsequent chunks even when the detok window is transiently empty (e.g. usage only final chunk). """ - # TODO: Properties used in follow on PR for tool call parsing - last_content: str | None = None - """Last emitted cumulative assistant content text.""" + output_token_ids: list[int] = Field(default_factory=list) + """All output tokens seen so far. Parser path only. + + Replay buffer: each chunk rebuilds a fresh parser and replays every + token in here through `parse_delta` (discarding the result) before + processing the current chunk's tokens since parser internal state + cannot be serialized into this stateless model. Unavoidably O(n) + bounded by ``max_model_len`` (enforced server side, not by a field + validator here since the bound is model dependent). + """ + + output_chunk_lens: list[Annotated[int, Field(gt=0)]] = Field(default_factory=list) + """Token count of each chunk in `output_token_ids`. Parser path only. + + Replay uses these to reproduce the original `parse_delta` call + boundaries. Must sum to `len(output_token_ids)`. + """ + + @model_validator(mode="after") + def _validate_output_chunk_lens(self) -> "DerenderStreamState": + total = sum(self.output_chunk_lens) + if total != len(self.output_token_ids): + raise ValueError( + f"output_chunk_lens must sum to len(output_token_ids) " + f"(got sum={total}, len(output_token_ids)=" + f"{len(self.output_token_ids)})" + ) + return self + + tools_streamed: bool = False + """True once a tool call delta has been emitted. Parser path only. - last_reasoning: str | None = None - """Last emitted cumulative reasoning text.""" + Drives the `finish_reason` -> `"tool_calls"` rewrite on the final + chunk mirroring the generate streaming path. + """ last_tool_call_ids: list[str] = Field(default_factory=list) - """Stable tool-call IDs, assigned once when each call first appears. + """Stable tool call IDs, assigned once when each call first appears. - Prevents ID regeneration across re-parsing. + Indexed by tool call index. Parser path only. Prevents ID regeneration + when replay reprocesses a tool call that already has a pinned ID. """ @@ -570,6 +616,19 @@ class DerenderChatStreamRequest(BaseModel): prompt_tokens: int | None = None """Prompt token count for usage. Forwarded from the render step.""" + prompt_token_ids: list[int] | None = None + """Prompt token IDs. Required by the parser path's `parse_delta` to + settle its initial reasoning state (e.g. chat templates that pre-open + ````). `prompt_tokens` is a usage count and cannot serve this + purpose. Sourced from `GenerateRequest.token_ids` at the render step. + + Rejected with a 400 (by `ServingDerender`) when a tool or reasoning + parser is configured and this is omitted. Without it, `parse_delta` + cannot tell whether the prompt left reasoning open and would silently + misclassify reasoning content as plain content. Unused on the plain + detokenization path. + """ + chat_request: ChatCompletionRequest | None = None """The original (post adjust_request) ChatCompletionRequest from /render.""" diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 05f2710afb95..a89bc843b81e 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -4,7 +4,10 @@ from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption -from vllm.entrypoints.generate.base.protocol import DeltaMessage, ToolCall +from vllm.entrypoints.generate.base.protocol import ( + DeltaMessage, + ToolCall, +) from vllm.entrypoints.generate.base.serving import resolve_token_id_placeholder from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProbs, @@ -29,6 +32,9 @@ ) from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.entrypoints.serve.utils.tool_calls_utils import ( + maybe_filter_parallel_tool_calls, +) from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer @@ -95,6 +101,11 @@ def __init__( self._derender_completion_async = make_async( self._derender_completion, executor=renderer._executor ) + # Replay is O(n) per chunk (unlike the O(delta) detok-only paths), + # so it must not run on the event loop either. + self._derender_chat_stream_parsed_async = make_async( + self._derender_chat_stream_parsed, executor=renderer._executor + ) async def derender_chat( self, @@ -144,6 +155,7 @@ def _derender_chat( tokenizer, chat_request.tools, chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, ) reasoning, content, tool_calls = parser.parse( decoded_text, @@ -278,23 +290,30 @@ async def derender_chat_stream( state: DerenderStreamState | None = None, chat_request: ChatCompletionRequest | None = None, prompt_tokens: int | None = None, + prompt_token_ids: list[int] | None = None, ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState]: """Process one GenerateStreamResponse chunk for streaming chat derender. - TODO: parse path for reasoning and tool calls is implemented in future PR. - - Unlike OpenAI's API, which always emits ``role: "assistant"`` on the + Unlike OpenAI's API, which always emits `role: "assistant"` on the very first chunk, this emits it on the first chunk with a non empty - ``choices`` list. A leading usage only chunk therefore defers the + `choices` list. A leading usage only chunk therefore defers the role to the following content chunk instead of sending an empty role only delta. Args: model: Model name for the response object. - generate_chunk: One SSE chunk from ``/inference/v1/generate``. - state: Client carried detok state (``None`` for first call). - chat_request: Original ChatCompletionRequest from ``/render``. + generate_chunk: One SSE chunk from `/inference/v1/generate`. + state: Client carried detok state (`None` for first call). + chat_request: Original ChatCompletionRequest from `/render`. + Required when a reasoning or tool parser is configured + (validated by the caller — see `ServingDerender`) because + plain detokenization would leak raw parser markup into `content`. prompt_tokens: Prompt token count for the usage chunk. + prompt_token_ids: Prompt token IDs. Parser path only. Lets + `parse_delta` settle its initial reasoning state. Required + when a reasoning or tool parser is configured (validated by + the caller — see ``ServingDerender``). Without it reasoning + left open by the prompt would be misclassified as content. Returns: (chunk, updated_state) — the derendered SSE chunk and the state @@ -304,28 +323,36 @@ async def derender_chat_stream( if state is None: state = DerenderStreamState() - if self.parser is not None: - # TODO: Follow on PR will implement the parse path. Check on the - # parser alone (fail closed). A parser configured model must never - # fall through to plain detok on the streaming path, even when - # ``chat_request`` is omitted or reasoning/tool markup would leak - # into ``delta.content``. - raise NotImplementedError( - "Streaming chat derender is not yet supported for models with " - "a reasoning or tool parser configured. Use the non-streaming " - "derender endpoint (stream=false) for parsed output." - ) - # A single DerenderStreamState is threaded through every choice in # this chunk. Correct only when there is at most one choice per SSE # event (n=1, one call per index), as the streaming derender # protocol assumes. Multiple choices sharing one chunk would corrupt - # each other's detok window. + # each other's detok/parser state. if len(generate_chunk.choices) > 1: raise ValueError( "derender_chat_stream expects at most one choice per chunk" ) + parser_cls = self.parser + if parser_cls is not None: + # Fail (mirrors ServingDerender's pre-check) because a parser + # configured model must never fall through to plain detok or + # reasoning/tool markup would leak into `delta.content`. + if chat_request is None: + raise ValueError( + "chat_request is required for streaming chat derender " + "when a tool or reasoning parser is configured" + ) + return await self._derender_chat_stream_parsed_async( + parser_cls, + model, + generate_chunk, + state, + chat_request, + prompt_tokens, + prompt_token_ids, + ) + tokenizer = self.renderer.get_tokenizer() skip_special = ( chat_request.skip_special_tokens if chat_request is not None else True @@ -374,6 +401,203 @@ async def derender_chat_stream( ) return chunk, updated_state + def _derender_chat_stream_parsed( + self, + parser_cls: type[Parser], + model: str, + generate_chunk: GenerateStreamResponse, + state: DerenderStreamState, + chat_request: ChatCompletionRequest, + prompt_tokens: int | None, + prompt_token_ids: list[int] | None, + ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState]: + """Parser path for streaming chat derender: replay + `parse_delta`. + + Parser internal state (buffered markup, reasoning/tool phase, etc.) + cannot be serialized into `DerenderStreamState`, so each call + builds a fresh parser and replays every prior output token through + `parse_delta` (discarding the result) before processing this + chunk's tokens for real. + + Every chunk, replayed or live, goes through one `parse_delta` call + with the tokens it arrived with (more than one under e.g. + speculative decoding). Replay recovers those boundaries from + `state.output_chunk_lens`, so the rebuilt parser state matches + standard serving, which calls `parse_delta` once per engine step. + Text comes from a fresh incremental detokenizer with special tokens + preserved (``skip_special_tokens=False``). + """ + tokenizer = self.renderer.get_tokenizer() + + chat_template_kwargs: dict[str, Any] = {} + if not self.use_harmony: + chat_template_kwargs = ( + chat_request.build_chat_params( + self.chat_template, + self.chat_template_content_format, + ) + .with_defaults(self.default_chat_template_kwargs) + .chat_template_kwargs + ) + + parser = parser_cls( + tokenizer, + chat_request.tools, + chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, + ) + + # Ephemeral incremental detok window, local to this call. Threaded + # across both the replay and current chunk phases (via `nonlocal`) + # so multi-byte characters split across that boundary still decode + # correctly. Discarded once the call returns. + detok_state = DerenderStreamState() + + def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: + """Replay prior chunks through `parse_delta` to rebuild parser + state, discarding the result. Never `finished` since that only + applies to the current chunk. + """ + nonlocal detok_state + start = 0 + for chunk_len in chunk_lens: + chunk = token_ids[start : start + chunk_len] + start += chunk_len + text, detok_state = self._detokenize_delta( + tokenizer, chunk, detok_state, skip_special_tokens=False + ) + parser.parse_delta( + text, + chunk, + chat_request, + prompt_token_ids=prompt_token_ids, + finished=False, + ) + + # Replay history to reconstruct parser state. The result is thrown + # away and only the current chunk's emission goes to the client. + _replay(state.output_token_ids, state.output_chunk_lens) + + stream_choices: list[ChatCompletionResponseStreamChoice] = [] + output_token_ids = list(state.output_token_ids) + output_chunk_lens = list(state.output_chunk_lens) + role_sent = state.role_sent + tools_streamed = state.tools_streamed + last_tool_call_ids = list(state.last_tool_call_ids) + + # At most one choice: the caller (derender_chat_stream) already + # rejects >1 before dispatching here. role_sent/tools_streamed/ + # output_token_ids below are updated for that single choice, not + # accumulated across choices. Looping over generate_chunk.choices + # could silently corrupt output if chunks were ever allowed to + # contain multiple choices. + if generate_chunk.choices: + choice = generate_chunk.choices[0] + delta_tids = choice.token_ids or [] + is_finished = choice.finish_reason is not None + + if delta_tids: + # One parse_delta call for the whole chunk (producer + # granularity), not one per token. See the granularity note + # in this method's docstring. + text, detok_state = self._detokenize_delta( + tokenizer, delta_tids, detok_state, skip_special_tokens=False + ) + delta_message = parser.parse_delta( + text, + delta_tids, + chat_request, + prompt_token_ids=prompt_token_ids, + finished=is_finished, + ) + elif is_finished: + # Finish only chunk (no new tokens). Still flush any + # buffered tool call arguments. + delta_message = parser.parse_delta( + "", + [], + chat_request, + prompt_token_ids=prompt_token_ids, + finished=True, + ) + else: + delta_message = None + + if delta_tids: + output_token_ids.extend(delta_tids) + output_chunk_lens.append(len(delta_tids)) + + if delta_message is None: + delta_message = DeltaMessage() + + if delta_message.tool_calls: + tools_streamed = True + for tc in delta_message.tool_calls: + if tc.id is None: + continue + if tc.index < len(last_tool_call_ids): + # Pin: reuse the ID already recorded for this index + # rather than one a from scratch replay regenerated. + # Real trigger not just defensive with + # tool_choice="required", + # extract_required_tool_call_streaming resets + # function_name_returned to False whenever the + # partial JSON transiently fails to parse which + # re-emits id+name for the same index on replay. + tc.id = last_tool_call_ids[tc.index] + else: + last_tool_call_ids.append(tc.id) + + if not role_sent: + delta_message.role = "assistant" + role_sent = True + + finish_reason = choice.finish_reason + if finish_reason is not None: + is_named_tool_choice = ( + type(chat_request.tool_choice) is ChatCompletionNamedToolChoiceParam + ) + if tools_streamed and not is_named_tool_choice: + finish_reason = "tool_calls" + + stream_choice = ChatCompletionResponseStreamChoice( + index=choice.index, + delta=delta_message, + finish_reason=finish_reason, + ) + stream_choices.append( + maybe_filter_parallel_tool_calls(stream_choice, chat_request) + ) + + updated_state = state.model_copy( + update={ + "output_token_ids": output_token_ids, + "output_chunk_lens": output_chunk_lens, + "role_sent": role_sent, + "tools_streamed": tools_streamed, + "last_tool_call_ids": last_tool_call_ids, + } + ) + + usage: UsageInfo | None = None + if generate_chunk.usage is not None: + u = generate_chunk.usage + pt = prompt_tokens if prompt_tokens is not None else (u.prompt_tokens or 0) + ct = u.completion_tokens or 0 + usage = UsageInfo( + prompt_tokens=pt, + completion_tokens=ct, + total_tokens=pt + ct, + ) + + chunk = ChatCompletionStreamResponse( + id=generate_chunk.request_id, + model=model, + choices=stream_choices, + usage=usage, + ) + return chunk, updated_state + async def derender_completion( self, generate_responses: list[GenerateResponse], From 1cdf1689e55c483a754c6c24baac3e47d39d4ece Mon Sep 17 00:00:00 2001 From: Artur Fierka Date: Fri, 18 Sep 2026 16:38:15 +0200 Subject: [PATCH 479/774] [Quantization][XPU] Enable int8_w8a8 MoE on the Triton backend for XPU (#53162) Signed-off-by: Artur Fierka Co-authored-by: Claude Co-authored-by: Kunshang Ji --- tests/quantization/test_int8_moe_oracle.py | 62 ++++++++++++++----- .../layers/fused_moe/experts/triton_moe.py | 14 ++++- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py index 8c838a785b63..9580eb9ab328 100644 --- a/tests/quantization/test_int8_moe_oracle.py +++ b/tests/quantization/test_int8_moe_oracle.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for INT8 (W8A8) fused-MoE oracle backend selection. -These exercise ``select_int8_moe_backend`` only (no kernels are launched), so -they run on any platform where the Triton INT8 MoE kernel is available — CUDA -(SM >= 7.5) or ROCm — not just gfx950. +These exercise ``select_int8_moe_backend`` only (no MoE kernels are launched), +so they run on any platform where the Triton INT8 MoE kernel is available — +CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950. """ import pytest @@ -27,17 +27,22 @@ ) from vllm.platforms import current_platform -# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on -# ROCm CDNA GPUs. Gate on that rather than on a specific arch. +# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+), on +# ROCm CDNA GPUs and on XPU. Gate on that rather than on a specific arch. INT8_MOE_SUPPORTED = ( - current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) -) or current_platform.is_rocm() + (current_platform.is_cuda() and current_platform.has_device_capability((7, 5))) + or current_platform.is_rocm() + or current_platform.is_xpu() +) requires_int8_moe = pytest.mark.skipif( not INT8_MOE_SUPPORTED, - reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)", + reason="Requires Triton INT8 MoE support (CUDA SM>=7.5, ROCm, or XPU)", ) +# So FusedMoEConfig.device names the device the test actually runs on. +DEVICE = current_platform.device_type + def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -52,7 +57,7 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), activation=MoEActivation.SILU, in_dtype=torch.bfloat16, - device="cuda", + device=DEVICE, routing_method=RoutingMethodType.Renormalize, moe_backend=moe_backend, ) @@ -62,15 +67,20 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: @pytest.mark.parametrize( "weight_key,activation_key", [ - # per-channel weight + dynamic per-token activation - (kInt8StaticChannelSym, kInt8DynamicTokenSym), - # per-tensor weight + dynamic per-tensor activation - (kInt8StaticTensorSym, kInt8DynamicTensorSym), + pytest.param( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + id="per_channel_weight-per_token_act", + ), + pytest.param( + kInt8StaticTensorSym, + kInt8DynamicTensorSym, + id="per_tensor_weight-per_tensor_act", + ), ], ) def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): - """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor - weights) select the Triton backend.""" + """Both dynamic-activation INT8 MoE schemes select the Triton backend.""" config = _make_int8_moe_config() backend, experts_cls = select_int8_moe_backend( config, weight_key=weight_key, activation_key=activation_key @@ -79,6 +89,28 @@ def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): assert experts_cls is not None +@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU-only behaviour") +def test_scaled_int8_quant_is_available_on_xpu(): + """The per-tensor scheme's activation quantization works on XPU. + + ``_int8_quantize`` sends per-tensor activations through + ``ops.scaled_int8_quant``, which has an XPU branch. This is what lets the + per-tensor scheme be enabled here alongside CUDA and ROCm; pin it so a + future change cannot silently remove that branch. + """ + from vllm import _custom_ops as ops + + x = torch.randn(4, 16, device=DEVICE, dtype=torch.bfloat16) + scale = torch.full((1,), 0.05, device=DEVICE, dtype=torch.float32) + + q, returned_scale, azp = ops.scaled_int8_quant(x, scale=scale) + + assert q.dtype == torch.int8 + assert q.shape == x.shape + assert azp is None + torch.testing.assert_close(returned_scale, scale) + + @requires_int8_moe def test_int8_explicit_moe_backend_triton(): """An explicit --moe-backend triton selects the Triton INT8 backend.""" diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index da9666a66b8f..700d113921e8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -130,10 +130,18 @@ def _supports_quant_scheme( # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and # the Triton int8_w8a8 fused MoE kernel handles them. + # XPU reaches the same Triton kernel through the same launcher, and both + # int8 activation-quant paths work there: per-token uses the Triton + # `per_token_quant_int8`, per-tensor resolves to the XPU branch of + # `scaled_int8_quant`, which is plain elementwise arithmetic. device_supports_int8 = ( - current_platform.is_cuda() - and current_platform.has_device_capability((7, 5)) - ) or current_platform.is_rocm() + ( + current_platform.is_cuda() + and current_platform.has_device_capability((7, 5)) + ) + or current_platform.is_rocm() + or current_platform.is_xpu() + ) supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] if device_supports_int8: From dbf4b89cf229fa52869aa36744d2e0ec0c8997ed Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Fri, 18 Sep 2026 23:01:35 +0800 Subject: [PATCH 480/774] [Misc] Remove unnecessary Transformers version guards (#57556) Signed-off-by: DarkLight1337 --- .../test_basic_correctness.py | 11 ----- tests/lora/test_qwenvl.py | 20 +++------ .../models/language/generation/test_common.py | 12 ------ .../multimodal/generation/test_common.py | 6 --- .../generation/vlm_utils/model_utils.py | 10 ----- .../pooling/test_jinavl_reranker.py | 22 ---------- .../multimodal/processing/test_idefics3.py | 6 --- .../multimodal/processing/test_qwen2_vl.py | 14 ------- .../multimodal/processing/test_smolvlm.py | 6 --- vllm/model_executor/models/glm4_1v.py | 42 +++++-------------- .../models/qwen3_omni_moe_thinker.py | 12 ------ 11 files changed, 15 insertions(+), 146 deletions(-) diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 1b3c92ee5b60..688d355c318d 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/tests/basic_correctness/test_basic_correctness.py @@ -11,8 +11,6 @@ import pytest import torch -from packaging.version import Version -from transformers import __version__ as TRANSFORMERS_VERSION import vllm.envs as envs from vllm import LLM @@ -149,15 +147,6 @@ def test_models( if enable_prompt_embeds: with torch.no_grad(): prompt_embeds = hf_model.get_prompt_embeddings(example_prompts) - if model == "hmellor/tiny-random-Gemma2ForCausalLM" and ( - Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0") - ): - # For Gemma 1/2 models with Transformers 5.4.0+, the prompt embeddings - # are normalised in `get_prompt_embeddings`, like Gemma 3. - # For older versions, we need to manually normalise. - embed_scale = hf_model.config.hidden_size**0.5 - normalizer = torch.tensor(embed_scale, dtype=prompt_embeds[0].dtype) - prompt_embeds = [p_e * normalizer for p_e in prompt_embeds] with VllmRunner( model, diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index 063de2b5875d..360b9aff05b0 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -3,8 +3,6 @@ from dataclasses import dataclass import pytest -from packaging.version import Version -from transformers import __version__ as TRANSFORMERS_VERSION import vllm from tests.conftest import VllmRunner @@ -29,20 +27,12 @@ class TestConfig: def __post_init__(self): if self.mm_processor_kwargs is None: - # There is a bug in transformers v4 where size is ignored by - # `Qwen2VLProcessor.__call__` - if Version(TRANSFORMERS_VERSION) < Version("5.2.0"): - self.mm_processor_kwargs = { - "min_pixels": 28 * 28, - "max_pixels": 1280 * 28 * 28, - } - else: - self.mm_processor_kwargs = { - "size": { - "shortest_edge": 28 * 28, - "longest_edge": 1280 * 28 * 28, - } + self.mm_processor_kwargs = { + "size": { + "shortest_edge": 28 * 28, + "longest_edge": 1280 * 28 * 28, } + } class Qwen2VLTester: diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 50a0a5ed8ce8..f70d5a0c6892 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -3,8 +3,6 @@ import pytest import torch -from packaging.version import Version -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.platforms import current_platform @@ -162,16 +160,6 @@ def test_models( if prompt_embeds is not None: embed = hf_model.model.get_input_embeddings()(token_ids) - if "gemma" in model.lower() and ( - Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0") - ): - # For Gemma 1/2 models with Transformers 5.4.0+, the prompt - # embeddings are normalised in `get_prompt_embeddings`, - # like Gemma 3. For older versions, we need to manually normalise. - embed_scale = hf_model.config.hidden_size**0.5 - normalizer = torch.tensor(embed_scale, dtype=embed.dtype) - embed *= normalizer - # MiniCPM models apply scale_emb to embeddings internally. # vLLM expects pre-scaled embeddings when using inputs_embeds. if model in EMBED_SCALING_MODELS: diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index c364bab3828f..4f68887d1d15 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -317,8 +317,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): img_idx_to_prompt=lambda idx: "", auto_cls=AutoModelForImageTextToText, vllm_output_post_proc=model_utils.blip2_vllm_to_hf_output, - # FIXME: https://github.com/huggingface/transformers/pull/38510 - marks=[pytest.mark.skip("Model is broken")], ), "cosmos3": VLMTestInfo( models=["nvidia/Cosmos3-Nano"], @@ -819,10 +817,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): patch_hf_runner=model_utils.paddleocr_vl_patch_hf_runner, image_size_factors=[(0.25,)], marks=[ - pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) == Version("4.57.3"), - reason="This model is broken in Transformers v4.57.3", - ), pytest.mark.skipif( Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), reason="Model's custom code uses ROPE_INIT_FUNCTIONS" diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py index 77beb64601af..ae3f376e1f7f 100644 --- a/tests/models/multimodal/generation/vlm_utils/model_utils.py +++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py @@ -308,16 +308,6 @@ def processor(*args, **kwargs): hf_model.processor = processor - orig_generate = hf_model.model.generate - - def _generate(self, *args, **kwargs): - # FIXME: https://github.com/huggingface/transformers/issues/38333 - kwargs["disable_compile"] = True - - return orig_generate(*args, **kwargs) - - hf_model.model.generate = types.MethodType(_generate, hf_model.model) - return hf_model diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index c71a2f08cbf1..05433e7269d5 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -import transformers -from packaging import version from transformers import AutoModel from vllm.assets.base import VLLM_S3_BUCKET_URL @@ -246,10 +244,6 @@ def _run_test( @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.skipif( - version.parse(transformers.__version__) == version.parse("4.57.5"), - reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295", -) def test_model_text_image( hf_runner, vllm_runner, @@ -269,10 +263,6 @@ def test_model_text_image( @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.skipif( - version.parse(transformers.__version__) == version.parse("4.57.5"), - reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295", -) def test_model_text_text( hf_runner, vllm_runner, @@ -292,10 +282,6 @@ def test_model_text_text( @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.skipif( - version.parse(transformers.__version__) == version.parse("4.57.5"), - reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295", -) def test_model_image_text( hf_runner, vllm_runner, @@ -315,10 +301,6 @@ def test_model_image_text( @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.skipif( - version.parse(transformers.__version__) == version.parse("4.57.5"), - reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295", -) def test_model_image_image( hf_runner, vllm_runner, @@ -338,10 +320,6 @@ def test_model_image_image( @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -@pytest.mark.skipif( - version.parse(transformers.__version__) == version.parse("4.57.5"), - reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295", -) def test_model_text_mixed_documents( hf_runner, vllm_runner, diff --git a/tests/models/multimodal/processing/test_idefics3.py b/tests/models/multimodal/processing/test_idefics3.py index 7365db59f2bc..feffbf6568a1 100644 --- a/tests/models/multimodal/processing/test_idefics3.py +++ b/tests/models/multimodal/processing/test_idefics3.py @@ -3,9 +3,7 @@ """Tests for Idefics3's multimodal preprocessing kwargs.""" import pytest -from packaging.version import Version from transformers import Idefics3Config -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.multimodal import MULTIMODAL_REGISTRY @@ -13,10 +11,6 @@ from ...utils import build_model_context -@pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) < Version("5.2.0"), - reason="See https://github.com/huggingface/transformers/pull/43948", -) @pytest.mark.parametrize("model_id", ["HuggingFaceM4/Idefics3-8B-Llama3"]) @pytest.mark.parametrize( ("mm_processor_kwargs", "expected_toks_per_img"), diff --git a/tests/models/multimodal/processing/test_qwen2_vl.py b/tests/models/multimodal/processing/test_qwen2_vl.py index 4976b05c3e0a..57b8e7fd391f 100644 --- a/tests/models/multimodal/processing/test_qwen2_vl.py +++ b/tests/models/multimodal/processing/test_qwen2_vl.py @@ -3,9 +3,7 @@ import pytest import torch -from packaging.version import Version from PIL import Image -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.model_executor.models.vision import FusedInputNorm from vllm.multimodal import MULTIMODAL_REGISTRY @@ -133,12 +131,6 @@ def test_processor_override( kwargs_on_init: bool, ): """Ensure Qwen2VLMultiModalProcessor handles min/max pixels properly.""" - if ( - Version(TRANSFORMERS_VERSION) < Version("5.2.0") - and "size" in mm_processor_kwargs - ): - pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`") - ctx = build_model_context( model_id, mm_processor_kwargs=mm_processor_kwargs if kwargs_on_init else None, @@ -184,12 +176,6 @@ def test_get_image_size_with_most_features( model_id: str, mm_processor_kwargs: dict[str, object], ): - if ( - Version(TRANSFORMERS_VERSION) < Version("5.2.0") - and "size" in mm_processor_kwargs - ): - pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`") - ctx = build_model_context( model_id, mm_processor_kwargs=mm_processor_kwargs, diff --git a/tests/models/multimodal/processing/test_smolvlm.py b/tests/models/multimodal/processing/test_smolvlm.py index 678b3fd39db1..f27c4eb7b8f5 100644 --- a/tests/models/multimodal/processing/test_smolvlm.py +++ b/tests/models/multimodal/processing/test_smolvlm.py @@ -3,9 +3,7 @@ """Tests for smolvlm's multimodal preprocessing kwargs.""" import pytest -from packaging.version import Version from transformers import SmolVLMConfig -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.multimodal import MULTIMODAL_REGISTRY @@ -13,10 +11,6 @@ from ...utils import build_model_context -@pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) < Version("5.2.0"), - reason="See https://github.com/huggingface/transformers/pull/43948", -) @pytest.mark.parametrize("model_id", ["HuggingFaceTB/SmolVLM2-2.2B-Instruct"]) @pytest.mark.parametrize( ("mm_processor_kwargs", "expected_toks_per_img"), diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 8f3db8a3d110..1537ae7d779f 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -36,9 +36,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -import transformers from einops import rearrange -from packaging.version import Version from transformers import BatchFeature, Glm4vProcessor from transformers.image_processing_base import ImageProcessingMixin from transformers.models.glm4v.configuration_glm4v import ( @@ -132,8 +130,6 @@ # For profile run _MAX_FRAMES_PER_VIDEO = 600 -TRANSFORMERS_WITH_GA = Version(transformers.__version__) >= Version("5.10.0.dev0") - def _to_video_metadata(metadata: Mapping[str, Any]) -> VideoMetadata: return VideoMetadata( @@ -1389,17 +1385,6 @@ def _get_video_second_idx_glmga( timestamps_list = full_second_idxs[::2] return list(timestamps_list) - def _get_video_frame_embed_token_id(self, hf_processor: object) -> int: - attr = ( - "image_token_id" - if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA - else "video_token_id" - ) - token_id = getattr(hf_processor, attr, None) - if not isinstance(token_id, int): - raise ValueError(f"Processor has no valid {attr}") - return token_id - def _construct_video_placeholder( self, video_array: np.ndarray, @@ -1437,7 +1422,7 @@ def _construct_video_placeholder( num_tokens_per_frame = int(H * W) // merge_length placeholder = [] placeholder.append(bov_token_id) - frame_embed_token_id = self._get_video_frame_embed_token_id(hf_processor) + frame_embed_token_id = hf_processor.image_token_id for frame_idx in frames_idx_token: placeholder.append(boi_token_id) placeholder.extend([frame_embed_token_id] * num_tokens_per_frame) @@ -1601,10 +1586,7 @@ def _apply_hf_processor_main( processor = self.info.get_hf_processor(**hf_kwargs) - use_direct_path = ( - not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA - ) - if use_direct_path: + if not isinstance(processor, Glm4vProcessor): prepared_data, prepared_kwargs = self._get_direct_path_inputs( hf_data, hf_kwargs ) @@ -1627,8 +1609,6 @@ def _apply_hf_processor_main( del hf_data["videos"] video_grid_thw_lst = [] pixel_values_videos_lst = [] - frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor) - swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id for item in videos: video_array, metadata = item @@ -1651,10 +1631,10 @@ def _apply_hf_processor_main( video_mm_kwargs, ) input_ids = video_outputs.pop("input_ids") - if swap_video_frame_tokens: - input_ids[input_ids == processor.image_token_id] = ( - processor.video_token_id - ) + + input_ids[input_ids == processor.image_token_id] = ( + processor.video_token_id + ) video_placeholder = processor.tokenizer.batch_decode(input_ids)[0] prompt_text = prompt_text.replace( "<|begin_of_video|><|video|><|end_of_video|>", @@ -1670,17 +1650,15 @@ def _apply_hf_processor_main( ) else: video_outputs = dict() - swap_video_frame_tokens = False processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**hf_kwargs), dict(text=prompt_text, **hf_data), hf_kwargs, ) - if swap_video_frame_tokens: - input_ids = processed_data["input_ids"] - input_ids[input_ids == processor.video_token_id] = processor.image_token_id - processed_data["input_ids"] = input_ids + input_ids = processed_data["input_ids"] + input_ids[input_ids == processor.video_token_id] = processor.image_token_id + processed_data["input_ids"] = input_ids processed_data.update(video_outputs) return self._finalize_hf_mm_data( @@ -1728,7 +1706,7 @@ def get_video_replacement(item_idx: int): ) return PromptUpdateDetails.select_token_id( placeholder, - embed_token_id=self.info._get_video_frame_embed_token_id(hf_processor), + embed_token_id=hf_processor.image_token_id, ) return [ diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 2c3647c81e98..85d0bcaad08e 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1228,18 +1228,6 @@ def _get_hf_mm_inputs( ) -> HFMultiModalInputs: hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) - if mm_items.get_count("audio", strict=False): - # TODO(Isotr0py): Remove this patch after upstream fix PR - # released and Transformers version update: - # https://github.com/huggingface/transformers/pull/41473 - hf_inputs = hf_inputs._replace( - hf_kwargs={ - **hf_inputs.hf_kwargs, - "audio_kwargs": dict(hf_inputs.hf_kwargs.get("audio_kwargs") or {}), - "text_kwargs": dict(hf_inputs.hf_kwargs.get("text_kwargs") or {}), - } - ) - audios = hf_inputs.hf_data.get("audio") if audios: feature_extractor = self.info.get_feature_extractor() From 017dced6a6fd3cf430e4686a47b354da1cadfbf5 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:17:07 +0300 Subject: [PATCH 481/774] [DeepSeek V4] Fix fused MoE expert distribution (#57465) Signed-off-by: Itay Alroy Co-authored-by: Tyler Michael Smith --- vllm/models/deepseek_v4/nvidia/model.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 59156365c519..39d3c87bc64f 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -993,20 +993,18 @@ def _init_fused_moe_experts( prefix: str, ) -> None: parallel_config = vllm_config.parallel_config - self.tp_rank = get_tensor_model_parallel_rank() + ep_group = get_ep_group() + ep_size = ep_group.world_size + ep_rank = ep_group.rank_in_group eplb_config = parallel_config.eplb_config self.n_redundant_experts = eplb_config.num_redundant_experts self.n_shared_experts = config.n_shared_experts or 0 self.n_logical_experts = self.n_routed_experts self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts - assert self.n_physical_experts % self.tp_size == 0, ( - f"n_physical_experts={self.n_physical_experts} must be divisible by " - f"tp_size={self.tp_size}. Adjust num_redundant_experts." - ) - self.n_local_physical_experts = self.n_physical_experts // self.tp_size + self.n_local_physical_experts = self.n_physical_experts // ep_size self.n_local_experts = self.n_local_physical_experts - self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_start_idx = ep_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts self.physical_expert_start = self.experts_start_idx self.physical_expert_end = self.experts_end_idx From 4991f976698d8ebf1a3b6d743c959b52bac689cc Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 18 Sep 2026 23:49:17 +0800 Subject: [PATCH 482/774] [XPU] sleep mode: fix KV cache release test (#57485) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji --- tests/basic_correctness/test_mem.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 746838dc40b6..6a483d26aa9f 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -128,31 +128,40 @@ def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch): sampling_params = SamplingParams(temperature=0, max_tokens=10) expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids - def get_mapped_bytes(worker): - return mapped_usage(get_mem_allocator_instance()) + def get_mapped_bytes_by_tag(worker): + allocator = get_mem_allocator_instance() + mapped_bytes = mapped_usage(allocator) + kv_cache_bytes = sum( + data.handle[1] + for data in allocator.pointer_to_data.values() + if data.tag == "kv_cache" and not data.is_asleep + ) + return mapped_bytes, kv_cache_bytes - mapped_before = llm.collective_rpc(get_mapped_bytes)[0] + mapped_before, kv_cache_before = llm.collective_rpc(get_mapped_bytes_by_tag)[0] # Utility RPC transports the engine error as a plain Exception. with pytest.raises(Exception, match="requires a completed pause"): llm.release_kv_cache_memory() - assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before + assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before assert not llm.llm_engine.is_sleeping() assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected llm.sleep(level=0) llm.release_kv_cache_memory() - mapped_after = llm.collective_rpc(get_mapped_bytes)[0] - assert mapped_before - mapped_after >= kv_cache_memory_bytes * 0.99 + mapped_after, kv_cache_after = llm.collective_rpc(get_mapped_bytes_by_tag)[0] + assert kv_cache_before > 0 + assert kv_cache_after == 0 + assert mapped_before - mapped_after >= kv_cache_before assert mapped_after > 0 assert llm.llm_engine.is_sleeping() if full_sleep: llm.sleep(level=full_sleep) - assert llm.collective_rpc(get_mapped_bytes)[0] == 0 + assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == 0 llm.wake_up(tags=None if full_sleep else ["kv_cache"]) if full_sleep == 2: llm.collective_rpc("reload_weights") - assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before + assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before assert not llm.llm_engine.is_sleeping() actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids assert actual == expected From a7156060c1fc17c5b31d46d3760829bed8ac1a58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 18 Sep 2026 18:08:44 +0200 Subject: [PATCH 483/774] [Core] Make parallel sampling (n>1) reqs admission atomic (#53936) Signed-off-by: NickLucche --- tests/v1/engine/test_admission_control.py | 90 +++++++++++++++++++++++ vllm/v1/engine/async_llm.py | 24 ++++-- 2 files changed, 107 insertions(+), 7 deletions(-) diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index 219850f5d489..15be166c8c0e 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -15,6 +15,7 @@ import argparse import asyncio import multiprocessing +from collections.abc import Awaitable, Callable from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -35,6 +36,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import human_readable_int +from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor @@ -78,6 +80,45 @@ def _make_output_processor(**request_states) -> OutputProcessor: return op +def _make_engine_request(request_id: str, n: int) -> EngineCoreRequest: + return EngineCoreRequest( + request_id=request_id, + prompt_token_ids=[1], + mm_features=None, + sampling_params=SamplingParams(n=n), + pooling_params=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + ) + + +def _make_request_test_llm( + max_num_queued_reqs: int, + add_request_async: Callable[[EngineCoreRequest], Awaitable[None]], +) -> AsyncLLM: + llm = _make_async_llm(max_num_queued_reqs=max_num_queued_reqs) + llm.output_processor = OutputProcessor(None, log_stats=False) + llm.engine_core = SimpleNamespace( + resources=SimpleNamespace(engine_dead=False), + add_request_async=AsyncMock(side_effect=add_request_async), + abort_requests_async=AsyncMock(), + shutdown=MagicMock(), + ) + llm.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(kv_sharing_fast_prefill=False) + ) + llm.output_handler = None + llm.log_requests = False + llm._run_output_handler = MagicMock() + llm.input_processor = MagicMock() + llm.input_processor.assign_request_id.side_effect = lambda request: setattr( + request, "external_req_id", request.request_id + ) + return llm + + def _make_scheduler_config(**kwargs) -> SchedulerConfig: return SchedulerConfig( runner_type="generate", @@ -340,6 +381,55 @@ async def add_request_async(_request): assert len(request_states) == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("first_n", [1, 3]) +async def test_parallel_admission_is_all_or_nothing(first_n: int): + """An n>1 request reserves either all its capacity or none of it.""" + + async def add_request_async(_request): + await asyncio.sleep(0) + + llm = _make_request_test_llm(3, add_request_async) + second_n = 3 if first_n == 1 else 1 + first = _make_engine_request("first", first_n) + second = _make_engine_request("second", second_n) + results = await asyncio.gather( + llm.add_request("first", first, first.params), + llm.add_request("second", second, second.params), + return_exceptions=True, + ) + + assert not isinstance(results[0], Exception) + assert isinstance(results[1], QueueOverflowError) + assert llm.output_processor.get_num_unfinished_requests() == first_n + assert llm.engine_core.add_request_async.await_count == first_n + + +@pytest.mark.asyncio +async def test_parallel_admission_cancellation_cleans_up_all_children(): + """Cancellation during core submission must release every reserved slot.""" + first_submission_started = asyncio.Event() + + async def add_request_async(_request): + first_submission_started.set() + await asyncio.Event().wait() + + llm = _make_request_test_llm(3, add_request_async) + request = _make_engine_request("parallel", 3) + output = llm.generate(request, request.params, request.request_id) + generate_task = asyncio.create_task(anext(output)) + + await first_submission_started.wait() + generate_task.cancel() + with pytest.raises(asyncio.CancelledError): + await generate_task + + assert llm.engine_core.add_request_async.await_count == 1 + assert llm.output_processor.get_num_unfinished_requests() == 0 + aborted_ids = llm.engine_core.abort_requests_async.await_args.args[0] + assert set(aborted_ids) == {"0_parallel", "1_parallel", "2_parallel"} + + # --------------------------------------------------------------------------- # create_error_response integration # --------------------------------------------------------------------------- diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index ad9d9bf6c776..7a57d0ff1857 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -401,10 +401,6 @@ async def add_request( "prompt logprobs" ) - if isinstance(params, SamplingParams) and params.n > 1: - # TODO (NickLucche) Batch check admission check for all n requests - self.check_admission(params.n, request_id) - if isinstance(prompt, AsyncGenerator): if reasoning_ended is not None or reasoning_parser_kwargs is not None: raise NotImplementedError @@ -500,14 +496,28 @@ async def add_request( # Fan out child requests (for n>1). parent_request = ParentRequest(request) + child_requests: list[EngineCoreRequest] = [] for idx in range(parent_params.n): request_id, child_params = parent_request.get_child_info(idx) child_request = request if idx == parent_params.n - 1 else copy(request) child_request.request_id = request_id child_request.sampling_params = child_params - await self._add_request( - child_request, prompt_text, parent_request, idx, queue - ) + child_requests.append(child_request) + + self.check_admission(parent_params.n, parent_request.request_id) + try: + for idx, child_request in enumerate(child_requests): + self.output_processor.add_request( + child_request, prompt_text, parent_request, idx, queue + ) + + for child_request in child_requests: + await self.engine_core.add_request_async(child_request) + if self.log_requests: + logger.info("Added request %s.", child_request.request_id) + except BaseException: + await self.abort(parent_request.request_id, internal=True) + raise return queue async def _add_request( From 164dca8d1e65b8f10346321f01b25e23caa642a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 18 Sep 2026 18:08:59 +0200 Subject: [PATCH 484/774] [Nixl] Separate transport-failure metrics from KV expiry (#55854) Signed-off-by: NickLucche --- docs/features/nixl_connector_usage.md | 15 ++-- .../unit/test_bidirectional_kv_transfer.py | 3 + .../kv_connector/unit/test_multi_connector.py | 24 ++++++ .../kv_connector/unit/test_nixl_connector.py | 82 +++++++++++++++++++ .../unit/test_nixl_push_connector.py | 8 ++ .../kv_connector/v1/nixl/base_worker.py | 33 ++++++-- .../kv_connector/v1/nixl/pull_worker.py | 21 ++++- .../kv_connector/v1/nixl/push_worker.py | 4 + .../kv_transfer/kv_connector/v1/nixl/stats.py | 56 +++++++++++-- 9 files changed, 224 insertions(+), 22 deletions(-) diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 9799eb9cf1b8..0d1f16e2aa45 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -418,12 +418,13 @@ activity for the last reporting interval. Example output: ```text KV Transfer metrics: Num successful transfers=4, Avg xfer time (ms)=1.381, P90 xfer time (ms)=2.601, Avg post time (ms)=0.672, P90 post time (ms)=0.801, -Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0 +Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0, +Num failed transfers=0, Num KV expired reqs=0 ``` The table below describes each field. All timing values cover only the -successful transfers recorded in the current interval; failed transfers are -counted separately via Prometheus (see +successful transfers recorded in the current interval; failures are reported +as interval counts in the same line and via Prometheus (see [Prometheus metrics](#prometheus-metrics) below). | Metric | Unit | Description | @@ -436,6 +437,8 @@ counted separately via Prometheus (see | `Avg MB per transfer` | MB | Mean payload size per transfer, computed as `total bytes transferred / number of transfers`. Reflects the average KV cache footprint of a single request (sequence length × layers × head dimension × dtype bytes). | | `Throughput (MB/s)` | MB/s | Effective bandwidth over the interval: `total MB transferred / total xfer time (s)` across all successful transfers. This is aggregate throughput, not per-request bandwidth. | | `Avg number of descriptors` | count | Mean number of NIXL memory descriptors (scatter-gather segments) submitted per transfer. More descriptors indicate more fragmented or larger KV cache allocations; very high counts can increase descriptor-registration overhead. | +| `Num failed transfers` | count | Failed NIXL transfers, handshakes, and completion notifications (`send_notif`) in the interval, grouped because all are sporadic lower-transport-layer events. | +| `Num KV expired reqs` | count | Requests whose KV blocks expired before being read. Reported separately from the failure counts above because it is an autoscaler/lease-tuning signal rather than a transport issue. | ### Prometheus metrics @@ -448,9 +451,9 @@ exported when NixlConnector is active: | `vllm:nixl_post_time_seconds` | Histogram | Time to submit the transfer request to the RDMA backend (seconds). | | `vllm:nixl_bytes_transferred` | Histogram | Bytes moved per transfer. | | `vllm:nixl_num_descriptors` | Histogram | Descriptor count per transfer. | -| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers. | -| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed completion notifications (`send_notif`). | -| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). | +| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers, including handshake and notification (`send_notif`) failures. These are grouped because all are sporadic lower-transport-layer events. | +| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed notifications (`send_notif`), retained for compatibility. These failures are also included in `vllm:nixl_num_failed_transfers`; do not sum the two counters. | +| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). Kept separate from the failure counters above: KV expiry is an autoscaler/lease-tuning signal, not a transport issue. | !!! tip High `vllm:nixl_num_kv_expired_reqs` indicates that the prefiller's lease diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 78a6d930f97a..3d539f97c8cb 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -969,6 +969,9 @@ def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): assert worker.xfer_stats.data["num_kv_expired_reqs"] == ( [1] if expect_declined else [] ) + # KV expiry must not pollute the transport failure counters. + assert worker.xfer_stats.data["num_failed_transfers"] == [] + assert worker.xfer_stats.data["num_failed_handshakes"] == [] _, done_recving = connector.get_finished(finished_req_ids=set()) assert "req" in done_recving diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index f6d8ed97d88f..1a85865560ab 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -567,6 +567,8 @@ def test_build_kv_connector_stats_reconstructs_nixl_stats(self): "num_descriptors": [10, 20], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } } @@ -590,6 +592,8 @@ def test_build_kv_connector_stats_with_multiple_connectors(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], }, "MockConnector": {"mock_field": [1, 2, 3]}, } @@ -619,6 +623,8 @@ def test_build_kv_connector_stats_raises_error_for_unknown_connector(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], }, } @@ -638,6 +644,8 @@ def test_build_kv_connector_stats_with_already_instantiated_objects(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) mock_stats = MockConnectorStats(data={"mock_field": [1, 2, 3]}) @@ -667,6 +675,8 @@ def test_build_kv_connector_stats_with_mixed_objects_and_dicts(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) @@ -698,6 +708,8 @@ def test_build_kv_connector_stats_skips_connectors_without_custom_stats(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], }, "ExampleConnector": {"some_field": [1, 2, 3]}, } @@ -735,6 +747,8 @@ def test_aggregate_same_connector(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } @@ -750,6 +764,8 @@ def test_aggregate_same_connector(self): "num_descriptors": [20], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } @@ -781,6 +797,8 @@ def test_aggregate_new_connector(self): "num_descriptors": [10], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } @@ -807,6 +825,8 @@ def test_reduce(self): "num_descriptors": [10, 20], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } @@ -833,6 +853,8 @@ def test_log_renders_plain_scalars(self): "num_descriptors": [10, 20], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } @@ -862,6 +884,8 @@ def test_reset(self): "num_descriptors": [10, 20], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], + "num_kv_expired_reqs": [], } ) } diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 9ab4b0ae95f6..3b9f6078652c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1682,6 +1682,83 @@ def test_kv_connector_stats_aggregation(): assert all(not isinstance(v, np.generic) for v in cli_stats.values()) +def test_kv_connector_stats_failure_grouping(): + """Transfer, handshake and notification failures are reported as one + transport-failure count, while KV expiry is reported separately: the + former are sporadic lower-transport-layer events, the latter an + autoscaler signal.""" + stats = NixlKVConnectorStats() + assert stats.is_empty() + + stats.record_failed_transfer() + stats.record_failed_handshake() + stats.record_failed_notification() + stats.record_kv_expired_req() + assert not stats.is_empty() + + # No successful transfers: latency stats are zero but the failure + # counts still surface. + reduced = stats.reduce() + assert reduced["Num successful transfers"] == 0 + assert reduced["Num failed transfers"] == 3 + assert reduced["Num KV expired reqs"] == 1 + + +def test_nixl_prom_metrics_group_handshake_with_transfer_failures(): + """vllm:nixl_num_failed_transfers counts handshake and notification + failures too, while vllm:nixl_num_kv_expired_reqs stays a separate + counter.""" + from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlPromMetrics, + ) + + registry = CollectorRegistry() + + class RegistryGauge(Gauge): + def __init__(self, *args, **kwargs): + super().__init__(*args, registry=registry, **kwargs) + + class RegistryCounter(Counter): + def __init__(self, *args, **kwargs): + super().__init__(*args, registry=registry, **kwargs) + + class RegistryHistogram(Histogram): + def __init__(self, *args, **kwargs): + super().__init__(*args, registry=registry, **kwargs) + + vllm_config = create_vllm_config() + metric_types = { + Gauge: RegistryGauge, + Counter: RegistryCounter, + Histogram: RegistryHistogram, + } + prom = NixlPromMetrics( + vllm_config, + metric_types, + labelnames=["engine"], + per_engine_labelvalues={0: ["engine-0"]}, + ) + + stats = NixlKVConnectorStats() + stats.record_failed_transfer() + stats.record_failed_handshake() + stats.record_failed_notification() + stats.record_kv_expired_req() + prom.observe(stats.data, engine_idx=0) + + def counter_value(name: str) -> float: + for metric in registry.collect(): + for sample in metric.samples: + if sample.name == name: + return sample.value + raise AssertionError(f"metric {name} not found in registry") + + assert counter_value("vllm:nixl_num_failed_transfers_total") == 3.0 + assert counter_value("vllm:nixl_num_kv_expired_reqs_total") == 1.0 + + def test_multi_kv_connector_stats_aggregation(): """Test MultiKVConnectorStats aggregation across TP ranks using KVOutputAggregator (used by MultiprocExecutor). @@ -3029,6 +3106,11 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): assert request_id in done_recving assert connector.get_block_ids_with_load_errors() == {1, 2, 3} + # Handshake failures are recorded as transport failures, separately + # from KV expiry. + assert connector.connector_worker.xfer_stats.data["num_failed_handshakes"] + assert connector.connector_worker.xfer_stats.data["num_failed_transfers"] == [] + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 7ce84bf685b1..e35ead001575 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -47,6 +47,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( NixlPushConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlKVConnectorStats, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, @@ -360,6 +363,7 @@ def fresh(cls) -> _StubWriterWorker: w._recv_failures = set() w._recving_transfers = defaultdict(list) w._is_hma_required = False + w.xfer_stats = NixlKVConnectorStats() w._reqs_to_process = set() w._reqs_to_send = {} w.consumer_notification_counts_by_req = defaultdict(int) @@ -730,6 +734,10 @@ def test_do_start_push_kv_drops_request_on_handshake_failure(): assert xfer_calls == [] assert len(failures) == 1 assert failures[0]["failure_type"] == "push_handshake_failed" + # The handshake metric is counted once by _ensure_handshake's own + # callback (stubbed out here); the push side must not double-record it. + assert w.xfer_stats.data["num_failed_handshakes"] == [] + assert w.xfer_stats.data["num_failed_transfers"] == [] def test_writer_loop_drains_deferred_push_inbox(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 60cf20303c3e..c2b20ffc597e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1306,6 +1306,8 @@ def done_callback( error=e, remote_engine_id=eid, ) + # Count once per handshake, regardless of waiting requests. + self.xfer_stats.record_failed_handshake() fut.add_done_callback(done_callback) return fut @@ -2811,13 +2813,17 @@ def get_transfer_results(self) -> KVConnectorTransferResults: done_sending.update(self._replicated_pcp_done_sending) self._replicated_pcp_done_sending.clear() - # Process receive failures reported by background threads. + # Process receive failures reported by background threads. Each + # producer already recorded a specific metric (handshake or + # notification), so only the failure bookkeeping runs here. while not self._failed_recv_reqs.empty(): try: req_id = self._failed_recv_reqs.get_nowait() except queue.Empty: break - self._handle_failed_transfer(req_id, None, self._recv_failures) + self._handle_failed_transfer( + req_id, None, self._recv_failures, record_failed_transfer=False + ) failed_recv_reqs: set[ReqId] = set() if self._recv_failures: @@ -3059,10 +3065,27 @@ def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool: return True def _handle_failed_transfer( - self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None + self, + req_id: str, + handle: int | None, + failed_req_ids: set[str] | None = None, + record_failed_transfer: bool = True, ) -> bool: - """Record a failure and release its handle, returning False to retain it.""" - self.xfer_stats.record_failed_transfer() + """Record a failure and release its handle, returning False to retain it. + + Args: + req_id: The request ID. + handle: The transfer handle. + failed_req_ids: Requests observed as failed; each stays listed + until it has no outstanding handles. + record_failed_transfer: Whether to count the failure toward the + transport-failure metric. Callers that already recorded a more + specific metric (eg KV expiry, reported separately from + transport failures) pass False. + + """ + if record_failed_transfer: + self.xfer_stats.record_failed_transfer() if failed_req_ids is not None: failed_req_ids.add(req_id) return handle is None or self._try_release_xfer_handle(req_id, handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 203c6bcfa668..7711a35954ca 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -149,7 +149,11 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): engine_id, ) self.xfer_stats.record_kv_expired_req() - self._handle_failed_transfer(req_id, None, self._recv_failures) + # KV expiry is reported separately from transport failures, so only + # the state cleanup side of _handle_failed_transfer runs here. + self._handle_failed_transfer( + req_id, None, self._recv_failures, record_failed_transfer=False + ) return if any(len(group) > 0 for group in meta.local_block_ids): @@ -364,7 +368,20 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]: remote_agents = self._remote_agents[meta.remote.engine_id] for rank_to_notify, agent in remote_agents.items(): if rank_to_notify != (0, read_specs[0].remote_rank): - self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) + try: + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) + except Exception as e: + self._log_failure( + failure_type="notification_failed", + msg="Remote rank will not update request state. " + "This may indicate network issues.", + req_id=req_id, + error=e, + dst_engine_id=meta.remote.engine_id, + remote_rank=rank_to_notify[1], + remote_agent_name=agent, + ) + self.xfer_stats.record_failed_notification() def _read_blocks( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 01c956a5251d..fbe85f46a327 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -365,6 +365,7 @@ def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: engine_id, req_id, ) + self.xfer_stats.record_failed_notification() self._failed_recv_reqs.put(req_id) return for rank, agent_name in agents.items(): @@ -377,6 +378,9 @@ def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: error=e, remote_rank=rank, ) + self.xfer_stats.record_failed_notification() + # Earlier registrations may still trigger WRITEs into D's blocks. + # Keep the receive pending until those writes are finished. logger.debug( "Sent PUSH_REG for %s to engine %s (%dB)", req_id, engine_id, len(notif_msg) ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index 2cff3f843b3f..f6d0684b6b3e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -39,6 +39,7 @@ def reset(self): "num_descriptors": [], "num_failed_transfers": [], "num_failed_notifications": [], + "num_failed_handshakes": [], "num_kv_expired_reqs": [], } @@ -57,6 +58,10 @@ def record_failed_notification(self): """Record a failed NIXL notification (send_notif).""" self.data["num_failed_notifications"].append(1) + def record_failed_handshake(self): + """Record a failed NIXL handshake.""" + self.data["num_failed_handshakes"].append(1) + def record_kv_expired_req(self): """Record a request that had its KV blocks expire.""" self.data["num_kv_expired_reqs"].append(1) @@ -72,6 +77,7 @@ def is_empty(self) -> bool: self.num_successful_transfers == 0 and len(self.data["num_failed_transfers"]) == 0 and len(self.data["num_failed_notifications"]) == 0 + and len(self.data["num_failed_handshakes"]) == 0 and len(self.data["num_kv_expired_reqs"]) == 0 ) @@ -84,10 +90,22 @@ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: return self def reduce(self) -> dict[str, int | float]: - # Compute compact representative stats suitable for CLI logging + # Compute compact representative stats suitable for CLI logging. + # Failure counts are reported on every interval: transfer, handshake + # and notification failures are grouped as sporadic + # lower-transport-layer events, while KV expiry is reported separately + # as it is an actionable autoscaler signal rather than a transport + # issue. + failure_counts = { + "Num failed transfers": len(self.data["num_failed_transfers"]) + + len(self.data["num_failed_handshakes"]) + + len(self.data["num_failed_notifications"]), + "Num KV expired reqs": len(self.data["num_kv_expired_reqs"]), + } if self.num_successful_transfers == 0: - # CLI logging only reports successful transfers stats. If all requests in - # the interval were unsuccessful, Prom will report failures stats instead. + # Timing / throughput stats only cover successful transfers. If + # all requests in the interval were unsuccessful, the failure + # counts above still surface what happened. return { "Num successful transfers": 0, "Avg xfer time (ms)": 0, @@ -97,6 +115,7 @@ def reduce(self) -> dict[str, int | float]: "Avg MB per transfer": 0, "Throughput (MB/s)": 0, "Avg number of descriptors": 0, + **failure_counts, } xfer_time = np.asarray(self.data["transfer_duration"]) @@ -122,6 +141,7 @@ def reduce(self) -> dict[str, int | float]: "Avg MB per transfer": round(avg_mb.item(), 3), "Throughput (MB/s)": round(throughput_mb_s.item(), 3), "Avg number of descriptors": round(descs.mean().item(), 1), + **failure_counts, } @property @@ -212,15 +232,20 @@ def __init__( ) counter_nixl_num_failed_transfers = self._counter_cls( name="vllm:nixl_num_failed_transfers", - documentation="Number of failed NIXL KV Cache transfers.", + documentation="Number of failed NIXL KV Cache transfers, including" + " handshake and notification failures. NOTE: KV expiry is tracked" + " separately in vllm:nixl_num_kv_expired_reqs.", labelnames=labelnames, ) self.counter_nixl_num_failed_transfers = create_metric_per_engine( counter_nixl_num_failed_transfers, self.per_engine_labelvalues ) + counter_nixl_num_failed_notifications = self._counter_cls( name="vllm:nixl_num_failed_notifications", - documentation="Number of failed NIXL KV Cache notifications.", + documentation="Number of failed NIXL KV Cache notifications. " + "Retained for compatibility; these failures are also included in " + "vllm:nixl_num_failed_transfers.", labelnames=labelnames, ) self.counter_nixl_num_failed_notifications = create_metric_per_engine( @@ -254,13 +279,26 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): ): for list_item in transfer_stats_data[list_item_key]: prom_obj[engine_idx].observe(list_item) - for counter_obj, counter_item_key in zip( + for counter_obj, counter_item_keys in zip( [ self.counter_nixl_num_failed_transfers, self.counter_nixl_num_failed_notifications, self.counter_nixl_num_kv_expired_reqs, ], - ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"], + [ + # Transfer, handshake and notification failures are grouped: + # all are sporadic lower-transport-layer events. KV expiry is + # reported separately as it signals autoscaler behavior, not + # transport health. + ( + "num_failed_transfers", + "num_failed_handshakes", + "num_failed_notifications", + ), + ("num_failed_notifications",), + ("num_kv_expired_reqs",), + ], ): - for list_item in transfer_stats_data[counter_item_key]: - counter_obj[engine_idx].inc(list_item) + for counter_item_key in counter_item_keys: + for list_item in transfer_stats_data[counter_item_key]: + counter_obj[engine_idx].inc(list_item) From 56cac80b145ed1e9814bba519735ad4c71c530f5 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Sat, 19 Sep 2026 00:17:51 +0800 Subject: [PATCH 485/774] [Deprecation] Remove assistant_token_mask support (#57520) Signed-off-by: DarkLight1337 --- .../scale_out/render/test_render.py | 170 ------------------ .../openai/chat_completion/protocol.py | 13 -- vllm/entrypoints/scale_out/render/serving.py | 22 --- .../scale_out/token_in_token_out/protocol.py | 8 - vllm/inputs/engine.py | 10 -- vllm/renderers/hf.py | 123 +++---------- vllm/renderers/params.py | 8 - 7 files changed, 20 insertions(+), 334 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index e21343553dbd..e65a26a7d140 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -594,176 +594,6 @@ async def test_completion_render_multiple_prompts_token_offsets(client): assert 0 <= start <= end <= len(prompt) -@pytest.mark.asyncio -async def test_chat_completion_render_assistant_tokens_mask_default(client): - """Without return_assistant_tokens_mask, assistant_tokens_mask should be null.""" - response = await client.post( - "/v1/chat/completions/render", - json={ - "model": MODEL_NAME, - "messages": [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - {"role": "user", "content": "How are you?"}, - ], - }, - ) - - assert response.status_code == 200 - data = response.json() - assert data.get("assistant_tokens_mask") is None - - -@pytest.mark.asyncio -async def test_chat_completion_render_assistant_tokens_mask_false(client): - """Explicitly setting return_assistant_tokens_mask=false gives null.""" - response = await client.post( - "/v1/chat/completions/render", - json={ - "model": MODEL_NAME, - "messages": [ - {"role": "user", "content": "Hello"}, - ], - "return_assistant_tokens_mask": False, - }, - ) - - assert response.status_code == 200 - data = response.json() - assert data.get("assistant_tokens_mask") is None - - -@pytest.mark.asyncio -async def test_chat_render_assistant_tokens_mask_null_without_gen_tags( - client, -): - """The tiny test model lacks ``{% generation %}`` tags, so the mask is null.""" - response = await client.post( - "/v1/chat/completions/render", - json={ - "model": MODEL_NAME, - "messages": [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - ], - "return_assistant_tokens_mask": True, - }, - ) - - assert response.status_code == 200 - assert response.json().get("assistant_tokens_mask") is None - - -# A minimal chat template with {% generation %} tags so we can test that -# the mask correctly marks assistant tokens. -_TEMPLATE_WITH_GENERATION = ( - "{% for m in messages %}" - "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n" - "{% elif m['role'] == 'assistant' %}" - "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}" - "{% endif %}" - "{% endfor %}" -) - - -@pytest.mark.asyncio -async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags( - client, -): - """With a ``{% generation %}``-enabled template, the mask marks assistant - tokens and the masked tokens decode to the assistant content.""" - response = await client.post( - "/v1/chat/completions/render", - json={ - "model": MODEL_NAME, - "messages": [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - {"role": "user", "content": "Bye"}, - ], - "chat_template": _TEMPLATE_WITH_GENERATION, - "return_assistant_tokens_mask": True, - }, - ) - - assert response.status_code == 200 - data = response.json() - - mask = data["assistant_tokens_mask"] - token_ids = data["token_ids"] - assert mask is not None - assert isinstance(mask, list) - assert len(mask) == len(token_ids) - assert all(v in (0, 1) for v in mask) - assert sum(mask) > 0, "mask should mark at least one assistant token" - - # Detokenize masked (assistant) and unmasked (non-assistant) tokens - # separately to verify the mask is correct, not just non-empty. - masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m] - unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m] - - detok = await client.post( - "/detokenize", - json={"model": MODEL_NAME, "tokens": masked_ids}, - ) - assert detok.status_code == 200 - assert "Hi!" in detok.json()["prompt"] - - detok_rest = await client.post( - "/detokenize", - json={"model": MODEL_NAME, "tokens": unmasked_ids}, - ) - assert detok_rest.status_code == 200 - assert "Hi!" not in detok_rest.json()["prompt"] - assert "Bye" in detok_rest.json()["prompt"] - - -@pytest.mark.asyncio -async def test_chat_render_assistant_tokens_mask_follows_truncation(client): - """The assistant mask must be truncated with the prompt it describes. - - `assistant_tokens_mask` is positional: entry i labels token i. Truncating - `token_ids` from the left without truncating the mask leaves the two - describing different positions, and the mask ends up marking whichever - tokens happen to sit at the old offsets. - """ - messages = [ - # Deliberately lopsided: a long leading user turn and a short trailing - # one, so keeping the head of the mask is distinguishable from keeping - # its tail. - {"role": "user", "content": "Hello hello hello hello hello hello"}, - {"role": "assistant", "content": "Hi!"}, - {"role": "user", "content": "Bye"}, - ] - body = { - "model": MODEL_NAME, - "messages": messages, - "chat_template": _TEMPLATE_WITH_GENERATION, - "return_assistant_tokens_mask": True, - } - - full = await client.post("/v1/chat/completions/render", json=body) - assert full.status_code == 200 - full_token_ids = full.json()["token_ids"] - full_mask = full.json()["assistant_tokens_mask"] - assert sum(full_mask) > 0 - - keep = len(full_token_ids) - 4 - # Precondition: with this prompt the head and tail slices of the mask - # really do differ, so the assertion below can tell them apart. - assert full_mask[-keep:] != full_mask[:keep] - - truncated = await client.post( - "/v1/chat/completions/render", - json={**body, "truncate_prompt_tokens": keep, "truncation_side": "left"}, - ) - assert truncated.status_code == 200 - data = truncated.json() - - assert data["token_ids"] == full_token_ids[-keep:] - assert data["assistant_tokens_mask"] == full_mask[-keep:] - - @pytest.mark.asyncio async def test_messages_render_basic(client): """Test basic Anthropic Messages render endpoint.""" diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index badd5a0344c6..225e99244141 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -453,18 +453,6 @@ class ChatCompletionRequest(OpenAIBaseModel): ), ) - return_assistant_tokens_mask: bool = Field( - default=False, - description=( - "If true, the /render response will include an " - "``assistant_tokens_mask`` field — a per-token list of 0/1 " - "values indicating which tokens were assistant-generated. " - "Requires the chat template to use ``{% generation %}`` " - "tags. When the template does not support it, " - "``assistant_tokens_mask`` will be ``null``." - ), - ) - cache_salt: str | None = Field( default=None, min_length=1, @@ -605,7 +593,6 @@ def build_chat_params( extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, - return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask), # No-tools requests default to tool_choice="none" at the API # layer. Collapse that default before rendering, so K3 emits a # model-visible tool-choice instruction only for requests with a diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index b79138664f56..7973ac3c3d88 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -126,33 +126,11 @@ async def render_chat_request( ) params = request.to_sampling_params(max_tokens, self.default_sampling_params) - assistant_tokens_mask: list[int] | None = engine_input.get( # type: ignore[assignment] - "assistant_tokens_mask" - ) - if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len( - token_ids - ): - logger.warning( - "assistant_tokens_mask length (%d) != token_ids length (%d); " - "this can happen with multimodal inputs where " - "placeholder expansion changes the token count. " - "The mask may be positionally misaligned.", - len(assistant_tokens_mask), - len(token_ids), - ) - if len(assistant_tokens_mask) < len(token_ids): - assistant_tokens_mask.extend( - [0] * (len(token_ids) - len(assistant_tokens_mask)) - ) - else: - assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)] - request_id = f"chatcmpl-{random_uuid()}" return GenerateRequest( request_id=request_id, token_ids=token_ids, - assistant_tokens_mask=assistant_tokens_mask, features=self._extract_mm_features(engine_input), sampling_params=params, model=request.model, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 2bd47b105978..ae0fd62b6efb 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -150,14 +150,6 @@ class GenerateRequest(BaseModel): token_ids: list[int] = Field(min_length=1) """The token ids to generate text from.""" - assistant_tokens_mask: list[int] | None = None - """Per-token mask (1 = assistant-generated, 0 = not). - - Only populated when the render request sets ``return_assistant_tokens_mask=True`` - and the chat template supports ``{% generation %}``. - ``None`` when the mask was not requested or could not be computed. - """ - @field_validator("token_ids") @classmethod def validate_token_ids(cls, v: list[int]) -> list[int]: diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index 429e8a829dd0..26b840f76f3f 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -43,11 +43,6 @@ class TokensInput(_InputOptions): """Char-level (start, end) offsets per token, propagated from the renderer's TokensPrompt when offsets were computed.""" - assistant_tokens_mask: NotRequired[list[int] | None] - """Per-token 0/1 mask marking assistant-generated tokens. - Populated when ``return_assistant_tokens_mask=True`` is set on the - render request and the chat template supports ``{% generation %}``.""" - def tokens_input( prompt_token_ids: list[int], @@ -155,11 +150,6 @@ class MultiModalInput(_InputOptions): `prompt_token_ids`. """ - assistant_tokens_mask: NotRequired[list[int] | None] - """Per-token 0/1 mask marking assistant-generated tokens. - Populated when ``return_assistant_tokens_mask=True`` is set on the - render request and the chat template supports ``{% generation %}``.""" - def mm_input( prompt_token_ids: list[int], diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index a32824eebda5..1c35802d0d16 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -7,7 +7,7 @@ import itertools import weakref from collections import defaultdict, deque -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -758,7 +758,6 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = ..., chat_template: str | None = ..., tokenize: Literal[True] = ..., - return_assistant_tokens_mask: Literal[False] = ..., **kwargs, ) -> list[int]: ... @overload @@ -770,20 +769,8 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = ..., chat_template: str | None = ..., tokenize: Literal[False] = ..., - return_assistant_tokens_mask: Literal[False] = ..., **kwargs, ) -> str: ... -@overload -def safe_apply_chat_template( - model_config: ModelConfig, - tokenizer: HfTokenizer, - conversation: list[ConversationMessage], - *, - tools: list[dict[str, Any]] | None = ..., - chat_template: str | None = ..., - return_assistant_tokens_mask: Literal[True], - **kwargs, -) -> tuple[list[int], list[int] | None]: ... def safe_apply_chat_template( model_config: ModelConfig, tokenizer: HfTokenizer, @@ -792,9 +779,8 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = None, chat_template: str | None = None, tokenize: bool = True, - return_assistant_tokens_mask: bool = False, **kwargs, -) -> str | list[int] | tuple[list[int], list[int] | None]: +) -> str | list[int]: chat_template = resolve_chat_template( tokenizer, chat_template=chat_template, @@ -822,38 +808,6 @@ def safe_apply_chat_template( chat_template_kwargs=kwargs, ) - # assistant_tokens_mask requires tokenized output — force tokenize=True. - if return_assistant_tokens_mask: - tokenize = True - - # When return_assistant_tokens_mask is requested and the template supports it, - # request assistant_tokens_mask via return_dict. - # Check for the actual Jinja tag, not just the word "generation" - # (which also appears in add_generation_prompt). - if return_assistant_tokens_mask and "{% generation %}" in chat_template: - resolved_kwargs["return_assistant_tokens_mask"] = True - resolved_kwargs["return_dict"] = True - resolved_kwargs.pop("tokenize", None) - try: - result = tokenizer.apply_chat_template( - conversation=conversation, # type: ignore[arg-type] - tools=tools, # type: ignore[arg-type] - chat_template=chat_template, - tokenize=True, - **resolved_kwargs, - ) - except (TypeError, ValueError) as exc: - logger.warning( - "apply_chat_template failed for assistant_tokens_mask: %s", exc - ) - else: - if isinstance(result, Mapping): - token_ids = list(result.get("input_ids", [])) - mask_raw = result.get("assistant_masks") - mask = list(mask_raw) if mask_raw is not None else None - return token_ids, mask - return list(result), None - # transformers v5 changed the default of `return_dict` to True, which # makes `apply_chat_template(tokenize=True)` return a `BatchEncoding` # instead of `list[int]`. Force `return_dict=False` so downstream code @@ -879,9 +833,6 @@ def safe_apply_chat_template( logger.warning("Chat template rejected the request: %s", e) raise VLLMValidationError(_template_error_reason(e)) from e - if return_assistant_tokens_mask: - assert isinstance(plain, list), f"Expected list[int], got {type(plain)}" - return plain, None return plain @@ -1066,22 +1017,12 @@ def render_messages( logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) chat_template_kwargs["tokenize"] = True - assistant_tokens_mask: list[int] | None = None - if params.return_assistant_tokens_mask: - prompt_raw, assistant_tokens_mask = safe_apply_chat_template( - model_config, - tokenizer, - conversation, - return_assistant_tokens_mask=True, - **chat_template_kwargs, - ) - else: - prompt_raw = safe_apply_chat_template( - model_config, - tokenizer, - conversation, - **chat_template_kwargs, - ) + prompt_raw = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + **chat_template_kwargs, + ) # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 # model which uses unified vision chunks for both images and videos. @@ -1107,9 +1048,6 @@ def render_messages( prompt = parse_dec_only_prompt(prompt_raw) - if assistant_tokens_mask is not None: - cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask - # When `prompt_embeds` is mixed with other modality data, # `_process_tokens` runs `_process_multimodal` first (expanding # `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then @@ -1183,30 +1121,12 @@ async def render_messages_async( logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) chat_template_kwargs["tokenize"] = True - assistant_tokens_mask: list[int] | None = None - if params.return_assistant_tokens_mask: - result_with_mask = cast( - tuple[list[int], list[int] | None], - await make_async( - safe_apply_chat_template, - executor=self._executor, - )( - model_config, - tokenizer, - conversation, - return_assistant_tokens_mask=True, # type: ignore[arg-type] - **chat_template_kwargs, - ), - ) - prompt_raw: str | list[int] = result_with_mask[0] - assistant_tokens_mask = result_with_mask[1] - else: - prompt_raw = await self._apply_chat_template_async( - model_config, - tokenizer, - conversation, - **chat_template_kwargs, - ) + prompt_raw = await self._apply_chat_template_async( + model_config, + tokenizer, + conversation, + **chat_template_kwargs, + ) # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 # model which uses unified vision chunks for both images and videos. @@ -1232,9 +1152,6 @@ async def render_messages_async( prompt = parse_dec_only_prompt(prompt_raw) - if assistant_tokens_mask is not None: - cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask - # See `render_messages` for the rationale. if prompt_embeds_tensors and mm_data: assert prompt_embeds_placeholder_token_id is not None @@ -1276,7 +1193,6 @@ def _process_tokens( processor records all placeholder offsets in the final (post-expansion) coordinate space, no offset shifting needed afterwards. """ - assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None) prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) if prompt_embeds_info is not None: tensors, placeholder_token_id = prompt_embeds_info @@ -1284,7 +1200,9 @@ def _process_tokens( cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( list(prompt["prompt_token_ids"]), mm_updates ) + engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) + if prompt_embeds_info is not None: tensors, _ = prompt_embeds_info self._apply_prompt_embeds_to_engine_input( @@ -1292,8 +1210,7 @@ def _process_tokens( tensors, mm_updates, ) - if assistant_tokens_mask is not None: - engine_input["assistant_tokens_mask"] = assistant_tokens_mask + return engine_input @override @@ -1304,7 +1221,6 @@ async def _process_tokens_async( skip_mm_cache: bool = False, ) -> TokensInput | MultiModalInput: """Async equivalent of `_process_tokens`.""" - assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None) prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) if prompt_embeds_info is not None: tensors, placeholder_token_id = prompt_embeds_info @@ -1312,9 +1228,11 @@ async def _process_tokens_async( cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( list(prompt["prompt_token_ids"]), mm_updates ) + engine_input = await super()._process_tokens_async( prompt, skip_mm_cache=skip_mm_cache ) + if prompt_embeds_info is not None: tensors, _ = prompt_embeds_info self._apply_prompt_embeds_to_engine_input( @@ -1322,8 +1240,7 @@ async def _process_tokens_async( tensors, mm_updates, ) - if assistant_tokens_mask is not None: - engine_input["assistant_tokens_mask"] = assistant_tokens_mask + return engine_input @staticmethod diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 86470f8db1a4..9f76a2ed49a4 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -27,13 +27,9 @@ # Prompt keys whose entry i describes token i of the prompt, so they must be # reduced with the same slice as the prompt tokens themselves. Anything added # here is truncated by `TokenizeParams.apply_post_tokenization`. -# `_assistant_tokens_mask` is renderer-internal: `HfRenderer.render_messages` -# stashes it on the prompt so `_process_tokens` can move it onto the engine -# input. _PARALLEL_TO_PROMPT_TOKENS = ( "prompt_token_offsets", "prompt_is_token_ids", - "_assistant_tokens_mask", ) @@ -99,9 +95,6 @@ class ChatParams: mm_processor_kwargs: dict[str, Any] | None = None """The kwargs to pass to the multi-modal processor.""" - return_assistant_tokens_mask: bool = False - """Request a per-token assistant mask from apply_chat_template.""" - tool_choice: Any | None = None """Request-level tool choice for renderers that need API metadata.""" @@ -136,7 +129,6 @@ def with_defaults( default_mm_processor_kwargs, self.mm_processor_kwargs, ), - return_assistant_tokens_mask=self.return_assistant_tokens_mask, tool_choice=self.tool_choice, response_format=self.response_format, ) From 8a5cf5438728180210aea43897d6a717de4b46ec Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:29:11 +0100 Subject: [PATCH 486/774] [Multimodal] Type dummy options per modality (#57576) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- .../processing/test_audioflamingo3.py | 3 +- .../multimodal/processing/test_common.py | 15 ++-- .../multimodal/processing/test_molmo2.py | 6 +- .../multimodal/processing/test_qwen3_vl.py | 9 ++- .../processing/test_tensor_schema.py | 13 +-- vllm/config/multimodal.py | 79 +++++++++++++------ vllm/engine/arg_utils.py | 15 ++++ vllm/model_executor/models/aria.py | 12 +-- vllm/model_executor/models/audioflamingo3.py | 11 +-- vllm/model_executor/models/bagel.py | 11 +-- .../models/bailing_moe_v3_vl.py | 4 +- vllm/model_executor/models/bee.py | 13 +-- vllm/model_executor/models/blip2.py | 12 +-- vllm/model_executor/models/clip.py | 12 +-- vllm/model_executor/models/cohere2_vision.py | 11 +-- vllm/model_executor/models/cohere_asr.py | 4 +- vllm/model_executor/models/cohere_compass.py | 11 +-- vllm/model_executor/models/colmodernvbert.py | 10 +-- vllm/model_executor/models/deepseek_ocr.py | 8 +- vllm/model_executor/models/deepseek_ocr2.py | 8 +- vllm/model_executor/models/deepseek_vl2.py | 12 +-- vllm/model_executor/models/dots_ocr.py | 12 +-- vllm/model_executor/models/ernie45_vl.py | 19 ++--- vllm/model_executor/models/fireredasr2.py | 12 ++- vllm/model_executor/models/funasr.py | 12 +-- vllm/model_executor/models/funaudiochat.py | 11 +-- vllm/model_executor/models/gemma3_mm.py | 12 +-- vllm/model_executor/models/gemma3n_mm.py | 18 ++--- vllm/model_executor/models/gemma4_mm.py | 28 ++----- vllm/model_executor/models/gemma4_unified.py | 6 +- vllm/model_executor/models/glm4_1v.py | 19 ++--- vllm/model_executor/models/glm4v.py | 11 +-- vllm/model_executor/models/glmasr.py | 11 +-- vllm/model_executor/models/granite_speech.py | 12 +-- .../models/hyperclovax_vision_v2.py | 20 ++--- vllm/model_executor/models/idefics3.py | 11 +-- vllm/model_executor/models/interns1.py | 18 ++--- vllm/model_executor/models/internvl.py | 16 ++-- vllm/model_executor/models/isaac.py | 11 +-- vllm/model_executor/models/kanana_v.py | 4 +- vllm/model_executor/models/keye.py | 18 ++--- vllm/model_executor/models/kimi_audio.py | 4 +- vllm/model_executor/models/kimi_k25.py | 4 +- vllm/model_executor/models/kimi_vl.py | 12 +-- vllm/model_executor/models/lfm2_vl.py | 12 +-- vllm/model_executor/models/llava.py | 12 +-- .../model_executor/models/llava_next_video.py | 13 +-- vllm/model_executor/models/llava_onevision.py | 19 ++--- .../model_executor/models/llava_onevision2.py | 4 +- vllm/model_executor/models/midashenglm.py | 13 +-- vllm/model_executor/models/mimo_v2_omni.py | 23 ++---- vllm/model_executor/models/minicpmo.py | 12 +-- vllm/model_executor/models/minicpmv.py | 14 ++-- vllm/model_executor/models/mistral3.py | 12 +-- vllm/model_executor/models/mllama4.py | 12 +-- vllm/model_executor/models/molmo.py | 11 +-- vllm/model_executor/models/molmo2.py | 9 +-- vllm/model_executor/models/moondream3.py | 7 +- vllm/model_executor/models/moss_audio.py | 11 +-- .../models/moss_transcribe_diarize.py | 10 +-- vllm/model_executor/models/muse_glimmer.py | 16 +--- .../model_executor/models/nano_nemotron_vl.py | 18 +---- vllm/model_executor/models/nemotron_parse.py | 4 +- vllm/model_executor/models/nvlm_d.py | 11 +-- vllm/model_executor/models/openvla.py | 10 +-- vllm/model_executor/models/ovis.py | 12 +-- vllm/model_executor/models/ovis2_5.py | 19 ++--- vllm/model_executor/models/paddleocr_vl.py | 11 +-- vllm/model_executor/models/paligemma.py | 12 +-- vllm/model_executor/models/phi3v.py | 12 +-- vllm/model_executor/models/phi4mm.py | 19 ++--- vllm/model_executor/models/phi4siglip.py | 7 +- vllm/model_executor/models/pixtral.py | 14 ++-- .../models/qwen2_5_omni_thinker.py | 24 ++---- vllm/model_executor/models/qwen2_audio.py | 11 +-- vllm/model_executor/models/qwen2_vl.py | 18 ++--- vllm/model_executor/models/qwen3_asr.py | 12 +-- vllm/model_executor/models/qwen3_vl.py | 15 ++-- vllm/model_executor/models/rvl.py | 12 +-- vllm/model_executor/models/siglip.py | 12 +-- vllm/model_executor/models/step3_vl.py | 11 +-- vllm/model_executor/models/terratorch.py | 4 +- .../models/transformers/multimodal.py | 20 ++--- vllm/model_executor/models/ultravox.py | 11 +-- vllm/model_executor/models/voxtral.py | 14 ++-- vllm/model_executor/models/whisper.py | 11 +-- .../deepseek_v4/common/mm_preprocess.py | 8 +- .../deepseek_v41/common/mm_preprocess.py | 8 +- vllm/models/dots3_note/common/processor.py | 13 ++- vllm/models/inkling/common/mm_preprocess.py | 12 +-- vllm/models/kimi_k3/common/mm_preprocess.py | 7 +- .../models/minimax_m3/common/mm_preprocess.py | 10 +-- vllm/multimodal/processing/dummy_inputs.py | 9 +-- 93 files changed, 461 insertions(+), 720 deletions(-) diff --git a/tests/models/multimodal/processing/test_audioflamingo3.py b/tests/models/multimodal/processing/test_audioflamingo3.py index f5920a3be4b8..f223145d9895 100644 --- a/tests/models/multimodal/processing/test_audioflamingo3.py +++ b/tests/models/multimodal/processing/test_audioflamingo3.py @@ -25,6 +25,7 @@ from transformers import PretrainedConfig from tests.models.registry import HF_EXAMPLE_MODELS +from vllm.config.multimodal import MultiModalDummyOptions class MockAudioFlamingo3Config(PretrainedConfig): @@ -127,7 +128,7 @@ def test_dummy_data_generation(mock_ctx): builder = AudioFlamingo3DummyInputsBuilder(info) mm_counts = {"audio": 2} - dummy_data = builder.get_dummy_mm_data(100, mm_counts, {}) + dummy_data = builder.get_dummy_mm_data(100, mm_counts, MultiModalDummyOptions()) assert "audio" in dummy_data assert len(dummy_data["audio"]) == 2 diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 1fbb2149a60d..ea4305976cde 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -13,6 +13,7 @@ AudioDummyOptions, BaseDummyOptions, ImageDummyOptions, + MultiModalDummyOptions, VideoDummyOptions, ) from vllm.inputs import MultiModalDataDict, MultiModalInput @@ -174,7 +175,7 @@ def get_token_prompt( inputs = dummy_inputs.get_dummy_processor_inputs( model_config.max_model_len, mm_counts, - mm_options={}, + mm_options=MultiModalDummyOptions(), # Assume all Mistral models define this extra argument mm_data=mm_data, # type: ignore[call-arg] ) @@ -182,7 +183,7 @@ def get_token_prompt( inputs = dummy_inputs.get_dummy_processor_inputs( model_config.max_model_len, mm_counts, - mm_options={}, + mm_options=MultiModalDummyOptions(), ) if not isinstance(inputs.prompt, list): @@ -278,10 +279,12 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: return BaseDummyOptions(count=count) # Assign normalized DummyOptions to the model config - model_config.get_multimodal_config().limit_per_prompt = { - modality: _to_dummy_options(modality, count) - for modality, count in limit_mm_per_prompt_ints.items() - } + model_config.get_multimodal_config().limit_per_prompt = MultiModalDummyOptions( + { + modality: _to_dummy_options(modality, count) + for modality, count in limit_mm_per_prompt_ints.items() + } + ) baseline_processor = factories.build_processor(ctx, cache=None) cached_processor = factories.build_processor(ctx, cache=cache) diff --git a/tests/models/multimodal/processing/test_molmo2.py b/tests/models/multimodal/processing/test_molmo2.py index 5633d47a97dc..68d3c8163cb7 100644 --- a/tests/models/multimodal/processing/test_molmo2.py +++ b/tests/models/multimodal/processing/test_molmo2.py @@ -6,7 +6,7 @@ import pytest import torch -from vllm.config.multimodal import VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions from vllm.model_executor.models.molmo2 import ( Molmo2DummyInputsBuilder, build_flat_image_bool_length, @@ -83,7 +83,9 @@ def test_dummy_video_num_frames_override_honors_min_of_two( data = builder.get_dummy_mm_data( seq_len=128, mm_counts={"image": 0, "video": 1}, - mm_options={"video": VideoDummyOptions(num_frames=num_frames_override)}, + mm_options=MultiModalDummyOptions( + video=VideoDummyOptions(num_frames=num_frames_override) + ), ) video, _metadata = data["video"][0] diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index c58904a9e2c9..04f87049c961 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -12,6 +12,7 @@ import pytest from vllm.config import ModelConfig +from vllm.config.multimodal import MultiModalDummyOptions from vllm.multimodal import MULTIMODAL_REGISTRY from ...registry import HF_EXAMPLE_MODELS @@ -339,7 +340,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non limit_mm_per_prompt={"image": 0, "video": 1}, ) capped = MULTIMODAL_REGISTRY.create_processor(capped_ctx.model_config) - capped_dummy = capped.dummy_inputs.get_dummy_mm_data(1024, {"video": 1}, {}) + capped_dummy = capped.dummy_inputs.get_dummy_mm_data( + 1024, {"video": 1}, MultiModalDummyOptions() + ) capped_frames = capped_dummy["video"][0][0].shape[0] assert capped_frames == 16, ( f"Expected the dummy to spread the budget over 16 frames, got {capped_frames}" @@ -351,7 +354,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non limit_mm_per_prompt={"image": 0, "video": 1}, ) uncapped = MULTIMODAL_REGISTRY.create_processor(uncapped_ctx.model_config) - uncapped_dummy = uncapped.dummy_inputs.get_dummy_mm_data(1024, {"video": 1}, {}) + uncapped_dummy = uncapped.dummy_inputs.get_dummy_mm_data( + 1024, {"video": 1}, MultiModalDummyOptions() + ) uncapped_frames = uncapped_dummy["video"][0][0].shape[0] assert uncapped_frames == 2, ( f"Expected the uncapped dummy to keep 2 frames, got {uncapped_frames}" diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index e1fb0cef7439..f249a25269c4 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -18,6 +18,7 @@ AudioDummyOptions, BaseDummyOptions, ImageDummyOptions, + MultiModalDummyOptions, VideoDummyOptions, ) from vllm.distributed import ( @@ -98,7 +99,7 @@ def create_batched_mm_kwargs( processor_inputs = dummy_inputs.get_dummy_processor_inputs( seq_len=model_config.max_model_len, mm_counts=mm_counts, - mm_options={}, + mm_options=MultiModalDummyOptions(), ) mm_items = processor_inputs.mm_data_items resized_mm_data = { @@ -248,10 +249,12 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: return AudioDummyOptions(count=count) return BaseDummyOptions(count=count) - model_config.get_multimodal_config().limit_per_prompt = { - modality: _to_dummy_options(modality, count) - for modality, count in limit_mm_per_prompt.items() - } + model_config.get_multimodal_config().limit_per_prompt = MultiModalDummyOptions( + { + modality: _to_dummy_options(modality, count) + for modality, count in limit_mm_per_prompt.items() + } + ) processor = factories.build_processor(ctx, cache=None) with initialize_dummy_model(model_cls, model_config) as model: diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 09e184b8ce2b..4eca42430c21 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -3,11 +3,18 @@ from collections.abc import Mapping from pathlib import Path -from typing import Any, Literal, TypeAlias, TypedDict, final +from typing import Any, Literal, TypeAlias, TypeVar, final, overload import torch -from pydantic import ConfigDict, Field, field_validator, model_validator +from pydantic import ( + ConfigDict, + Field, + GetCoreSchemaHandler, + field_validator, + model_validator, +) from pydantic.dataclasses import dataclass +from pydantic_core import core_schema import vllm.envs as envs from vllm.config.ec_transfer import ECTransferConfig @@ -18,6 +25,8 @@ logger = init_logger(__name__) +_T = TypeVar("_T") + @dataclass class BaseDummyOptions: @@ -51,17 +60,50 @@ class AudioDummyOptions(BaseDummyOptions): @final -class MultiModalDummyOptionsBuiltins(TypedDict, total=False): - """Type annotations for modality types predefined by vLLM.""" +class MultiModalDummyOptions(dict[str, BaseDummyOptions]): + """Dummy data options for each modality. + + Lookups of the modalities predefined by vLLM return their own options + class, while any other modality returns + [`BaseDummyOptions`][vllm.config.multimodal.BaseDummyOptions]. + """ + + @overload # type: ignore[override] + def get(self, key: Literal["image"], /) -> ImageDummyOptions | None: ... + + @overload + def get(self, key: Literal["image"], default: _T, /) -> ImageDummyOptions | _T: ... + + @overload + def get(self, key: Literal["video"], /) -> VideoDummyOptions | None: ... + + @overload + def get(self, key: Literal["video"], default: _T, /) -> VideoDummyOptions | _T: ... + + @overload + def get(self, key: Literal["audio"], /) -> AudioDummyOptions | None: ... - image: ImageDummyOptions - """Options for dummy images.""" + @overload + def get(self, key: Literal["audio"], default: _T, /) -> AudioDummyOptions | _T: ... - video: VideoDummyOptions - """Options for dummy videos.""" + @overload + def get(self, key: str, /) -> BaseDummyOptions | None: ... - audio: AudioDummyOptions - """Options for dummy audios.""" + @overload + def get(self, key: str, default: _T, /) -> BaseDummyOptions | _T: ... + + def get(self, key: str, default: object = None, /) -> object: + return super().get(key, default) + + @classmethod + def __get_pydantic_core_schema__( + cls, + source_type: Any, + handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + return core_schema.no_info_after_validator_function( + cls, handler.generate_schema(dict[str, BaseDummyOptions]) + ) MMEncoderTPMode = Literal["weights", "data"] @@ -75,15 +117,6 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False): `"xpu"` on XPU). Validated against that set by the CLI.""" -MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions] -""" -A dictionary containing an entry for each modality type of dummy data. - -The built-in modalities are defined by -[`MultiModalDummyOptionsBuiltins`][vllm.config.multimodal.MultiModalDummyOptionsBuiltins]. -""" - - @config class MultiModalConfig: """Controls the behavior of multimodal models.""" @@ -91,7 +124,9 @@ class MultiModalConfig: language_model_only: bool = False """If True, disables all multimodal inputs by setting all modality limits to 0. Equivalent to setting `--limit-mm-per-prompt` to 0 for every modality.""" - limit_per_prompt: MMDummyOptions = Field(default_factory=dict) + limit_per_prompt: MultiModalDummyOptions = Field( + default_factory=MultiModalDummyOptions + ) """The maximum number of input items and options allowed per prompt for each modality. @@ -257,8 +292,8 @@ class MultiModalConfig: def _validate_limit_per_prompt( cls, value: dict[str, int | dict[str, int]], - ) -> MMDummyOptions: - out: MMDummyOptions = {} + ) -> MultiModalDummyOptions: + out = MultiModalDummyOptions() for k, v in value.items(): # Handle legacy format where only count is specified diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index d637391cd3a4..2dd4b04f2b56 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -22,6 +22,7 @@ cast, get_args, get_origin, + is_typeddict, ) import huggingface_hub @@ -183,6 +184,18 @@ def is_type(type_hint: TypeHint, type: TypeHintT) -> TypeIs[TypeHintT]: return type_hint is type or get_origin(type_hint) is type +def is_dict_subclass(type_hint: TypeHint) -> bool: + """Check if the type hint is a subclass of `dict`. + + `TypedDict`s are excluded because they do not support class checks. + """ + return ( + isinstance(type_hint, type) + and not is_typeddict(type_hint) + and issubclass(type_hint, dict) + ) + + def contains_type(type_hints: set[TypeHint], type: TypeHintT) -> bool: """Check if the type hints contain a specific type.""" return any(is_type(type_hint, type) for type_hint in type_hints) @@ -305,6 +318,8 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: for field in fields(cls): # Get the set of possible types for the field type_hints: set[TypeHint] = get_type_hints(field.type) + # Subclasses of dict (e.g. MultiModalDummyOptions) are CLI dicts + type_hints = {dict if is_dict_subclass(th) else th for th in type_hints} # If the field is a dataclass, we can use the model_validate_json generator = (th for th in type_hints if is_dataclass(th)) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 4240406f7b55..01119b4a00bd 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -10,7 +10,7 @@ from transformers.models.aria.processing_aria import AriaProcessor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.fused_moe import FusedMoEFactory @@ -315,22 +315,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: vision_config = self.info.get_vision_config() max_image_size = vision_config.image_size - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) return { "image": self._get_dummy_images( width=max_image_size, height=max_image_size, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 36f256927775..4be996d2784c 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -30,7 +30,7 @@ from transformers.models.qwen2_audio import Qwen2AudioEncoder from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import ModalityData, MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.models.module_mapping import MultiModelKeys @@ -216,21 +216,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: hf_processor = self.info.get_hf_processor() feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = int(hf_processor.max_audio_len * sampling_rate) - num_audios = mm_counts.get("audio", 0) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index 9e0bd4ccbf87..e9ebb0e5edc9 100644 --- a/vllm/model_executor/models/bagel.py +++ b/vllm/model_executor/models/bagel.py @@ -14,7 +14,7 @@ import torch.nn as nn from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.activation import get_act_fn @@ -248,23 +248,20 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) hf_config = self.info.get_hf_config() vit_config = hf_config.vit_config # Use the configured image size image_size = vit_config.image_size - image_overrides = mm_options.get("image") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) return { "image": self._get_dummy_images( width=image_size, height=image_size, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/bailing_moe_v3_vl.py b/vllm/model_executor/models/bailing_moe_v3_vl.py index 7cd05685309d..a236d0722b4c 100644 --- a/vllm/model_executor/models/bailing_moe_v3_vl.py +++ b/vllm/model_executor/models/bailing_moe_v3_vl.py @@ -13,7 +13,7 @@ from transformers import AutoProcessor, BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -103,7 +103,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() return { diff --git a/vllm/model_executor/models/bee.py b/vllm/model_executor/models/bee.py index 4a57413b6e83..cc6c0f09abc4 100644 --- a/vllm/model_executor/models/bee.py +++ b/vllm/model_executor/models/bee.py @@ -8,7 +8,7 @@ from transformers.activations import GELUActivation from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY @@ -90,21 +90,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index 48da4f63e825..cd019e41ec48 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -15,7 +15,7 @@ ) from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.quantization import QuantizationConfig @@ -444,23 +444,19 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: hf_config = self.info.get_hf_config() vision_config = hf_config.vision_config max_image_size = vision_config.image_size - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) return { "image": self._get_dummy_images( width=max_image_size, height=max_image_size, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py index 6fb311fb14c8..7c5ba496ea8d 100644 --- a/vllm/model_executor/models/clip.py +++ b/vllm/model_executor/models/clip.py @@ -15,7 +15,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import divide, get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.model_executor.layers.activation import get_act_fn @@ -205,20 +205,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index a4748932d67e..4fdaff44442e 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -18,7 +18,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import MulAndSilu from vllm.model_executor.layers.linear import ( @@ -202,19 +202,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) image_size = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=image_size.width, height=image_size.height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 33b4aa7b0d70..c8a196ffdc58 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -12,7 +12,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt @@ -1944,7 +1944,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, mm_processor_kwargs=None, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() diff --git a/vllm/model_executor/models/cohere_compass.py b/vllm/model_executor/models/cohere_compass.py index 1d0dd8c35b05..9e8adabb808a 100644 --- a/vllm/model_executor/models/cohere_compass.py +++ b/vllm/model_executor/models/cohere_compass.py @@ -51,7 +51,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, @@ -1238,11 +1238,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - image_overrides = mm_options.get("image") - target_image_width, target_image_height = ( self.info.get_image_size_with_most_features() ) @@ -1251,8 +1248,8 @@ def get_dummy_mm_data( "image": self._get_dummy_images( width=target_image_width, height=target_image_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/colmodernvbert.py b/vllm/model_executor/models/colmodernvbert.py index 0c93e7ef52e3..a14476f7747b 100644 --- a/vllm/model_executor/models/colmodernvbert.py +++ b/vllm/model_executor/models/colmodernvbert.py @@ -15,7 +15,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed from vllm.multimodal import MULTIMODAL_REGISTRY @@ -135,18 +135,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 3d1921dc2b7f..62a46fc49852 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -11,7 +11,7 @@ from transformers import BatchFeature, CLIPVisionConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -268,17 +268,15 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - max_image_size = self.info.get_image_size_with_most_features() return { "image": self._get_dummy_images( width=max_image_size.width, height=max_image_size.height, - num_images=num_images, + num_images=mm_counts.get("image", 0), ) } diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index a64b71b64e49..938598a68bc6 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -11,7 +11,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -140,17 +140,15 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - max_image_size = self.info.get_image_size_with_most_features() return { "image": self._get_dummy_images( width=max_image_size.width, height=max_image_size.height, - num_images=num_images, + num_images=mm_counts.get("image", 0), ) } diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index 6f798a87e647..c720f623351b 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -15,7 +15,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig @@ -212,20 +212,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - max_image_size = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=max_image_size.width, height=max_image_size.height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py index 36b1cd530d6b..adcc560ba5ea 100644 --- a/vllm/model_executor/models/dots_ocr.py +++ b/vllm/model_executor/models/dots_ocr.py @@ -9,7 +9,7 @@ from transformers.models.qwen2_vl import Qwen2VLProcessor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import utils as dist_utils from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, @@ -104,20 +104,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index e1b847c2940f..a800458bbf81 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -36,7 +36,7 @@ from transformers import BaseImageProcessor, BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -1276,33 +1276,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index 82a4a6b50677..be5f5efbff55 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -12,7 +12,7 @@ ) from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import MultiModalDataDict, PromptType from vllm.logger import init_logger @@ -207,20 +207,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = feature_extractor.chunk_length * sampling_rate - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) ret = { "audio": self._get_dummy_audios( - length=audio_len, num_audios=num_audios, overrides=audio_overrides + length=audio_len, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } return ret diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index 7bf253ab40ee..f8db52378570 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -14,7 +14,7 @@ ) from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict, PromptType @@ -734,22 +734,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = feature_extractor.chunk_length * sampling_rate - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ), } diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 7a093cc2e74c..8b40ddf0cd43 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -26,7 +26,7 @@ from transformers.modeling_outputs import BaseModelOutput from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear @@ -580,7 +580,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = int(feature_extractor.sampling_rate) @@ -597,15 +597,12 @@ def get_dummy_mm_data( 1, (target_num_frames * sampling_rate + token_fps - 1) // token_fps, ) - num_audios = int(mm_counts.get("audio", 0)) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=int(mm_counts.get("audio", 0)), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 5bccd9aec0bb..c6166bc2f423 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -11,7 +11,7 @@ from transformers.models.gemma3.processing_gemma3 import Gemma3ProcessorKwargs from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import GemmaRMSNorm @@ -244,20 +244,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index d0f536afe89d..be2e8edeb2a7 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -18,7 +18,7 @@ from transformers.models.siglip import SiglipImageProcessorFast from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import MultiModalDataDict, PromptType, TextPrompt from vllm.logger import init_logger @@ -232,10 +232,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_audios = mm_counts.get("audio", 0) processor = self.info.get_hf_processor() audio_feature_extractor: Gemma3nAudioFeatureExtractor = ( processor.feature_extractor @@ -245,21 +243,17 @@ def get_dummy_mm_data( img_width = image_processor.size.get("width", 224) img_height = image_processor.size.get("height", 224) - image_overrides = mm_options.get("image") - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - return { "image": self._get_dummy_images( width=img_width, height=img_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ), } diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index f8cb9752c861..796dfbb6e8fa 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -36,11 +36,7 @@ from vllm.config import VllmConfig from vllm.config.model import get_served_model_name -from vllm.config.multimodal import ( - AudioDummyOptions, - BaseDummyOptions, - VideoDummyOptions, -) +from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm @@ -273,10 +269,7 @@ def get_mm_max_tokens_per_item( num_frames = _VIDEO_MAX_FRAMES mm_config = self.ctx.model_config.get_multimodal_config() video_opts = mm_config.limit_per_prompt.get("video") - if ( - isinstance(video_opts, VideoDummyOptions) - and video_opts.num_frames is not None - ): + if video_opts is not None and video_opts.num_frames is not None: num_frames = min(num_frames, video_opts.num_frames) tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6) return tokens @@ -502,9 +495,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) num_audios = mm_counts.get("audio", 0) num_videos = mm_counts.get("video", 0) processor = self.info.get_hf_processor() @@ -517,18 +509,12 @@ def get_dummy_mm_data( img_width = size.get("width", 224) img_height = size.get("height", 224) - image_overrides = mm_options.get("image") if mm_options else None - audio_overrides = mm_options.get("audio") if mm_options else None - video_overrides = mm_options.get("video") if mm_options else None - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - data: dict[str, Any] = { "image": self._get_dummy_images( width=img_width, height=img_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image") if mm_options else None, ), } @@ -537,7 +523,7 @@ def get_dummy_mm_data( data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, - overrides=audio_overrides, + overrides=mm_options.get("audio") if mm_options else None, ) if num_videos > 0: @@ -546,7 +532,7 @@ def get_dummy_mm_data( height=img_height, num_frames=_VIDEO_MAX_FRAMES, num_videos=num_videos, - overrides=video_overrides, + overrides=mm_options.get("video") if mm_options else None, ) return data diff --git a/vllm/model_executor/models/gemma4_unified.py b/vllm/model_executor/models/gemma4_unified.py index 7c9aaa8499ec..1a2624651528 100644 --- a/vllm/model_executor/models/gemma4_unified.py +++ b/vllm/model_executor/models/gemma4_unified.py @@ -28,7 +28,6 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import VideoDummyOptions from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.models.gemma3n_mm import batch_audio_features from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM @@ -177,10 +176,7 @@ def get_mm_max_tokens_per_item( num_frames = _VIDEO_MAX_FRAMES mm_config = self.ctx.model_config.get_multimodal_config() video_opts = mm_config.limit_per_prompt.get("video") - if ( - isinstance(video_opts, VideoDummyOptions) - and video_opts.num_frames is not None - ): + if video_opts is not None and video_opts.num_frames is not None: num_frames = min(num_frames, video_opts.num_frames) tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6) return tokens diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 1537ae7d779f..95cdcc5a2063 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -52,7 +52,7 @@ from transformers.video_utils import VideoMetadata from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -1455,33 +1455,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index 4fc6446ad3c1..7b42d5734f3a 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -16,7 +16,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import SiluAndMul, get_act_fn @@ -445,22 +445,19 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: hf_config = self.info.get_hf_config() vision_config = hf_config.vision_config target_width = target_height = vision_config["image_size"] - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 003b02c28ef9..cf12dedbe5fb 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -11,7 +11,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensPrompt @@ -694,13 +694,10 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate - num_audios = mm_counts.get("audio", 0) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) max_audio_len = getattr( self.info.get_hf_processor(), "max_audio_len", DEFAULT_MAX_AUDIO_LEN_S @@ -710,8 +707,8 @@ def get_dummy_mm_data( return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 1c1e98be043c..53cbe7df43dd 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -34,7 +34,7 @@ from transformers import BatchFeature, PretrainedConfig from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear @@ -208,17 +208,13 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - return { "audio": self._get_dummy_audios( length=self.info.get_max_audio_len(), - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py index 07106c947134..e269e2a05993 100644 --- a/vllm/model_executor/models/hyperclovax_vision_v2.py +++ b/vllm/model_executor/models/hyperclovax_vision_v2.py @@ -18,7 +18,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.forward_context import set_forward_context from vllm.inputs import MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY @@ -193,7 +193,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, mm_processor_kwargs: Mapping[str, object] | None = None, ) -> ProcessorInputs: """Build dummy processor inputs for memory profiling.""" @@ -225,31 +225,25 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, mm_processor_kwargs: Mapping[str, object] | None = None, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = 16 # Default for video - image_overrides = mm_options.get("image") if mm_options else None - video_overrides = mm_options.get("video") if mm_options else None - result: MultiModalDataDict = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, # type: ignore + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image") if mm_options else None, ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, # type: ignore + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video") if mm_options else None, ), } diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index dd03e17d1ec0..80c71cfe7992 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -29,7 +29,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -221,21 +221,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) hf_processor = self.info.get_hf_processor() image_processor: Idefics3ImageProcessor = hf_processor.image_processor longest_edge = image_processor.max_image_size["longest_edge"] - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=longest_edge, height=longest_edge, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py index 4a4264660d54..8d818e9a78fb 100644 --- a/vllm/model_executor/models/interns1.py +++ b/vllm/model_executor/models/interns1.py @@ -22,7 +22,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.interns1_vit import InternS1VisionModel @@ -299,35 +299,29 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) config = self.info.get_hf_config() image_size_h, image_size_w = config.vision_config.image_size - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=image_size_w, height=image_size_h, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index 873c594c54b3..3640f8d6a31b 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -17,7 +17,7 @@ from transformers import BatchFeature, PretrainedConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig @@ -203,19 +203,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } @@ -425,7 +422,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: dummy_image = super().get_dummy_mm_data(seq_len, mm_counts, mm_options) if self.info.ctx_video_token: @@ -436,9 +433,6 @@ def get_dummy_mm_data( ) num_videos = mm_counts.get("video", 0) video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance( - video_overrides, VideoDummyOptions - ) dummy_video = { "video": self._get_dummy_videos( width=image_size, diff --git a/vllm/model_executor/models/isaac.py b/vllm/model_executor/models/isaac.py index f7c1ee2dd131..f8f6dee32c96 100644 --- a/vllm/model_executor/models/isaac.py +++ b/vllm/model_executor/models/isaac.py @@ -13,7 +13,7 @@ from transformers.image_processing_utils import BatchFeature from vllm.config import ModelConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -386,19 +386,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index e3c27e479b00..fd076d66e7b8 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -18,7 +18,7 @@ from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY @@ -442,7 +442,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: num_images = mm_counts.get("image", 0) return { diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 1322920900bb..bfd2af714b82 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -17,7 +17,7 @@ from transformers.utils import torch_int from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import ModalityData, MultiModalDataDict from vllm.logger import init_logger @@ -1124,30 +1124,24 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features(seq_len) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - mm_data = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index 219c7f0e7b92..3541fcf371aa 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -13,7 +13,7 @@ from transformers import WhisperConfig as HFWhisperConfig from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import PromptType, TokensPrompt from vllm.model_executor.model_loader import DefaultModelLoader @@ -180,7 +180,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> ProcessorInputs: dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 9af46e441ab7..81b5ec164629 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -14,7 +14,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig @@ -216,7 +216,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: # TODO: Support mm_options for vision_chunk to allow user configuration dummy_items = self.get_dummy_mm_items() diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index 344dfdaaf031..13779042b62e 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/vllm/model_executor/models/kimi_vl.py @@ -53,7 +53,7 @@ from transformers.activations import GELUActivation from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.interfaces import ( @@ -219,18 +219,14 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=MaxImageTokenMeta.width, height=MaxImageTokenMeta.height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index 1cd51b93a2ed..644da00c81c2 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -21,7 +21,7 @@ from typing_extensions import Buffer from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.forward_context import set_forward_context from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear @@ -387,20 +387,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index c3a437eb368b..3ea0df9d163c 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -19,7 +19,7 @@ from transformers.models.pixtral import PixtralProcessor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear @@ -228,20 +228,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/llava_next_video.py b/vllm/model_executor/models/llava_next_video.py index 235d1adbac24..a33480ef08eb 100644 --- a/vllm/model_executor/models/llava_next_video.py +++ b/vllm/model_executor/models/llava_next_video.py @@ -10,7 +10,7 @@ from transformers import BatchFeature, LlavaNextVideoConfig, LlavaNextVideoProcessor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.models.clip import CLIPVisionModel @@ -171,25 +171,20 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - return { "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ) } diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py index 66503d774137..6fc96c8e6774 100644 --- a/vllm/model_executor/models/llava_onevision.py +++ b/vllm/model_executor/models/llava_onevision.py @@ -14,7 +14,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.multimodal import MULTIMODAL_REGISTRY @@ -278,33 +278,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 6988982c2316..d8d514be2381 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -54,7 +54,7 @@ support_torch_compile, ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import ModalityData, MultiModalDataDict @@ -1471,7 +1471,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, ) -> MultiModalDataDict: n_img = mm_counts.get("image", 0) n_vid = mm_counts.get("video", 0) diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index 7de6ab0e7e2d..6dae25ffb286 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -36,7 +36,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn @@ -564,18 +564,13 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - return { "audio": self._get_dummy_audios( length=self.info.get_max_audio_len(), - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index d74edd6d4ad0..91a3177b7f50 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -15,11 +15,7 @@ from typing_extensions import TypedDict from vllm.config import VllmConfig -from vllm.config.multimodal import ( - BaseDummyOptions, - ImageDummyOptions, - VideoDummyOptions, -) +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -1184,33 +1180,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index 2ec85679212c..cc8a1066f2d9 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -40,7 +40,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import ModalityData, MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems from vllm.multimodal.inputs import ( @@ -402,22 +402,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) audio_len = ( self.info.get_max_audio_chunks_with_most_features() * self.info.get_default_audio_sampling_rate() ) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - audio_mm_data = { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 14720c5c8630..8027dc1027d4 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -45,9 +45,8 @@ from vllm.config import VllmConfig from vllm.config.multimodal import ( - BaseDummyOptions, ImageDummyOptions, - VideoDummyOptions, + MultiModalDummyOptions, ) from vllm.inputs import ModalityData, MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig @@ -848,9 +847,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) num_videos = mm_counts.get("video", 0) image_width, image_height = self.info.get_image_size_with_most_features() @@ -859,14 +857,12 @@ def get_dummy_mm_data( seq_len, mm_counts ) - image_overrides = mm_options.get("image") video_overrides = mm_options.get("video") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) # Convert video overrides to image overrides for per-frame image generation, # and apply num_frames override to num_video_frames. video_frame_overrides: ImageDummyOptions | None = None - if isinstance(video_overrides, VideoDummyOptions): + if video_overrides is not None: if video_overrides.num_frames: num_video_frames = min(num_video_frames, video_overrides.num_frames) if video_overrides.width or video_overrides.height: @@ -880,8 +876,8 @@ def get_dummy_mm_data( "image": self._get_dummy_images( width=image_width, height=image_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": [ self._get_dummy_images( diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index f9d7dafe70b3..41352dc42766 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -11,7 +11,7 @@ from transformers.models.pixtral import PixtralProcessor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.layernorm import RMSNorm @@ -261,20 +261,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 797c06d18093..89d4edbe8666 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -36,7 +36,7 @@ support_torch_compile, ) from vllm.config import VllmConfig, set_current_vllm_config -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.attention import MMEncoderAttention @@ -709,20 +709,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - (target_width, target_height) = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index 9aeef3715c8f..17ad7f0be1b0 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -21,7 +21,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -1104,19 +1104,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 4c0ce0868e2d..881be146e790 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -26,7 +26,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -1834,7 +1834,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: num_images = mm_counts.get("image", 0) num_videos = mm_counts.get("video", 0) @@ -1845,13 +1845,11 @@ def get_dummy_mm_data( if num_images > 0: target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - dummy_images = self._get_dummy_images( width=target_width, height=target_height, num_images=num_images, - overrides=image_overrides, + overrides=mm_options.get("image"), ) if num_videos > 0: @@ -1864,7 +1862,6 @@ def get_dummy_mm_data( video_overrides = mm_options.get("video") if video_overrides: - assert isinstance(video_overrides, VideoDummyOptions) num_frames_override = video_overrides.num_frames if num_frames_override: if num_frames_override > target_num_frames: diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py index 45b9420184fc..1e750fc9d0fb 100644 --- a/vllm/model_executor/models/moondream3.py +++ b/vllm/model_executor/models/moondream3.py @@ -13,7 +13,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -943,15 +943,14 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, mm_processor_kwargs: Mapping[str, object] | None = None, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) return { "image": self._get_dummy_images( width=378, height=378, - num_images=num_images, + num_images=mm_counts.get("image", 0), ) } diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index 6092812585b3..02551d5663bc 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -19,7 +19,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, @@ -1341,16 +1341,13 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=16000, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 32ac69c9afb8..208d5d4bc1ab 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -20,7 +20,7 @@ from transformers import BatchFeature from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt from vllm.model_executor.models.interfaces import ( @@ -423,20 +423,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: num_audios = mm_counts.get("audio", 0) if num_audios == 0: return {} feature_extractor = self.info.get_feature_extractor() - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=_get_max_audio_samples(feature_extractor), num_audios=num_audios, - overrides=audio_overrides, + overrides=mm_options.get("audio"), ) } @@ -444,7 +442,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> ProcessorInputs: dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index ed8064542be3..a7938bb56b84 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -45,11 +45,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import ( - BaseDummyOptions, - ImageDummyOptions, - VideoDummyOptions, -) +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import ( divide, get_pp_group, @@ -234,7 +230,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: processor = self.info.get_hf_processor() video_processor = processor.video_processor @@ -245,16 +241,12 @@ def get_dummy_mm_data( * int(video_processor.downsample_factor) * video_grid ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) return { "image": self._get_dummy_images( width=image_width, height=image_height, num_images=mm_counts.get("image", 0), - overrides=image_overrides, + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=video_size, @@ -263,7 +255,7 @@ def get_dummy_mm_data( seq_len, mm_counts ), num_videos=mm_counts.get("video", 0), - overrides=video_overrides, + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index fee431617c8b..de61880b5f65 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -21,8 +21,7 @@ from vllm import envs from vllm.config import VllmConfig from vllm.config.multimodal import ( - AudioDummyOptions, - BaseDummyOptions, + MultiModalDummyOptions, VideoDummyOptions, ) from vllm.inputs import MultiModalDataDict, MultiModalInput @@ -839,22 +838,19 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) (target_width, target_height), _ = ( self.info.get_dummy_image_size_and_max_tokens(mm_counts) ) processor = self.info.get_hf_processor() - image_overrides = mm_options.get("image") - dummy_image = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } @@ -889,9 +885,6 @@ def get_dummy_mm_data( target_num_frames = num_frames num_videos = mm_counts.get("video", 0) video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance( - video_overrides, VideoDummyOptions - ) dummy_video = { "video": self._get_dummy_videos( width=video_width, @@ -907,9 +900,6 @@ def get_dummy_mm_data( if sound_config := self.info.sound_config: num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") if mm_options else None - assert audio_overrides is None or isinstance( - audio_overrides, AudioDummyOptions - ) tokens_per_audio = max(1, seq_len // max(num_audios, 1)) max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate calculated_max_audio_num_samples = ParakeetExtractor.audio_length( diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 01b557853116..5c4375b60058 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -22,7 +22,7 @@ from vllm.config import CacheConfig, VllmConfig from vllm.config.lora import LoRAConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.activation import get_act_fn @@ -388,7 +388,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: num_images = mm_counts.get("image", 0) diff --git a/vllm/model_executor/models/nvlm_d.py b/vllm/model_executor/models/nvlm_d.py index a69f024a0158..5efcdb6f5360 100644 --- a/vllm/model_executor/models/nvlm_d.py +++ b/vllm/model_executor/models/nvlm_d.py @@ -13,7 +13,7 @@ import torch.nn as nn from transformers import PretrainedConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.multimodal import MULTIMODAL_REGISTRY @@ -83,19 +83,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py index add7eac41f56..9242f7e1f51f 100644 --- a/vllm/model_executor/models/openvla.py +++ b/vllm/model_executor/models/openvla.py @@ -9,7 +9,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.linear import ( @@ -299,18 +299,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - image_overrides = mm_options.get("image") image_size = self.info.get_image_size_with_most_features() return { "image": self._get_dummy_images( width=image_size.width, height=image_size.height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index bd00686b41d7..ed0ea7ab7baa 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -29,7 +29,7 @@ from transformers import BatchFeature, PretrainedConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -311,20 +311,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - mm_data = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } return mm_data diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py index 588c06b8dea6..2520b38ae396 100644 --- a/vllm/model_executor/models/ovis2_5.py +++ b/vllm/model_executor/models/ovis2_5.py @@ -11,7 +11,7 @@ from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -290,33 +290,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) - mm_data = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } return mm_data diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 976d82032f20..69dae97d8d79 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -32,7 +32,7 @@ from transformers.utils import torch_int from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -223,19 +223,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - max_image_size = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=max_image_size.width, height=max_image_size.height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py index c1eafa2f09b1..d0173e4c116e 100644 --- a/vllm/model_executor/models/paligemma.py +++ b/vllm/model_executor/models/paligemma.py @@ -8,7 +8,7 @@ from transformers import BatchFeature, PaliGemmaConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY @@ -130,22 +130,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: hf_config = self.info.get_hf_config() vision_config = hf_config.vision_config max_image_size = vision_config.image_size - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=max_image_size, height=max_image_size, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/phi3v.py b/vllm/model_executor/models/phi3v.py index a75b4afd2217..2486dc3e8c79 100644 --- a/vllm/model_executor/models/phi3v.py +++ b/vllm/model_executor/models/phi3v.py @@ -29,7 +29,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig @@ -371,20 +371,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index c08eda612d1e..86cb4859678f 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -17,7 +17,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_pp_group from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -839,28 +839,21 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) - mm_data = { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "audio": self._get_dummy_audios( length=_AUDIO_MAX_SOUNDFILE_SIZE, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ), } diff --git a/vllm/model_executor/models/phi4siglip.py b/vllm/model_executor/models/phi4siglip.py index fd9fa94c7d8b..078f3155c227 100644 --- a/vllm/model_executor/models/phi4siglip.py +++ b/vllm/model_executor/models/phi4siglip.py @@ -14,7 +14,7 @@ from transformers import BatchFeature, PretrainedConfig, Siglip2VisionConfig from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY @@ -126,15 +126,14 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) size = self.info.get_image_size_with_most_features() return { "image": self._get_dummy_images( width=size.width, height=size.height, - num_images=num_images, + num_images=mm_counts.get("image", 0), overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index 425bba74fdf6..f5c167aec3da 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -21,7 +21,7 @@ from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import divide, get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn @@ -208,20 +208,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } @@ -229,7 +225,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, mm_data: MultiModalDataDict | None = None, ) -> ProcessorInputs: tokenizer = self.info.get_tokenizer() diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index cecc8345e5ee..ce7cb058b77a 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -44,7 +44,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.forward_context import set_forward_context from vllm.inputs import ModalityData, MultiModalDataDict from vllm.logger import init_logger @@ -448,12 +448,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - feature_extractor = self.info.get_feature_extractor() target_audio_length = ( @@ -469,28 +465,24 @@ def get_dummy_mm_data( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - audio_overrides = mm_options.get("audio") - mm_data = { "audio": self._get_dummy_audios( length=target_audio_length, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ), "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 7e7ffd5cfa20..b643fe422727 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -37,7 +37,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import ModalityData, MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( @@ -214,21 +214,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = feature_extractor.chunk_length * sampling_rate - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index d5e9d0a28bef..4bbddb0167c7 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -43,7 +43,7 @@ from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import parallel_state, tensor_model_parallel_all_gather from vllm.distributed import utils as dist_utils from vllm.inputs import ModalityData, MultiModalDataDict @@ -1096,32 +1096,26 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - target_width, target_height = self.info.get_image_size_with_most_features() target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) - image_overrides = mm_options.get("image") - video_overrides = mm_options.get("video") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, - num_videos=num_videos, - overrides=video_overrides, + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), ), } diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 14925ea55063..241e26fa4e72 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -32,7 +32,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensPrompt from vllm.logger import init_logger @@ -226,10 +226,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - feature_extractor = self.info.get_feature_extractor() target_audio_length = ( @@ -240,13 +238,11 @@ def get_dummy_mm_data( * feature_extractor.sampling_rate ) - audio_overrides = mm_options.get("audio") - return { "audio": self._get_dummy_audios( length=target_audio_length, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ), } diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index dc2baa1c5ee4..f7b553926cc8 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -54,8 +54,8 @@ ) from vllm.config import VllmConfig from vllm.config.multimodal import ( - BaseDummyOptions, MultiModalConfig, + MultiModalDummyOptions, VideoDummyOptions, VideoPruningMethod, ) @@ -1119,11 +1119,8 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - image_overrides = mm_options.get("image") video_overrides = mm_options.get("video") target_image_width, target_image_height = ( @@ -1133,7 +1130,6 @@ def get_dummy_mm_data( # treat videos as special images target_num_frames = 2 if video_overrides: - assert isinstance(video_overrides, VideoDummyOptions) num_frames_override = video_overrides.num_frames if num_frames_override: if num_frames_override > target_num_frames: @@ -1204,7 +1200,6 @@ def get_dummy_mm_data( target_video_size.height, ) if video_overrides: - assert isinstance(video_overrides, VideoDummyOptions) width_override = video_overrides.width if width_override: if width_override > target_video_width: @@ -1230,14 +1225,14 @@ def get_dummy_mm_data( "image": self._get_dummy_images( width=target_image_width, height=target_image_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=target_video_width, height=target_video_height, num_frames=target_num_frames, - num_videos=num_videos, + num_videos=mm_counts.get("video", 0), ), } diff --git a/vllm/model_executor/models/rvl.py b/vllm/model_executor/models/rvl.py index bf9661bcaf8e..5e822edae1db 100644 --- a/vllm/model_executor/models/rvl.py +++ b/vllm/model_executor/models/rvl.py @@ -8,7 +8,7 @@ from transformers.activations import GELUActivation from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY @@ -40,20 +40,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ), } diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 8a5b226b25d3..ab7146c5981d 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -16,7 +16,7 @@ ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import divide, get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.model_executor.layers.activation import get_act_fn @@ -158,20 +158,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - target_width, target_height = self.info.get_image_size_with_most_features() - image_overrides = mm_options.get("image") - return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index aed99c84cb57..ac747a7ecb39 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -11,7 +11,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn @@ -139,19 +139,16 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: target_width, target_height = self.info.get_image_size_with_most_features() - num_images = mm_counts.get("image", 0) - - image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( width=target_width, height=target_height, - num_images=num_images, - overrides=image_overrides, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), ) } diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index ddd1ac81ceb5..ce04acc72d37 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -33,7 +33,7 @@ from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalInput, mm_input from vllm.logger import init_logger from vllm.model_executor.layers.pooler import IdentityPooler @@ -153,7 +153,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: # Dummy data is generated based on the 'input' section # defined in the HF configuration file diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 164a4dbd8396..cecb4f564a98 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -27,7 +27,6 @@ from transformers.utils.generic import ModelOutput from vllm.compilation.decorators import should_torch_compile_mm_encoder -from vllm.config.multimodal import AudioDummyOptions from vllm.config.utils import getattr_iter from vllm.inputs import ( MultiModalDataBuiltins, @@ -75,7 +74,7 @@ from transformers import BatchFeature, PreTrainedModel from vllm.config import VllmConfig - from vllm.config.multimodal import BaseDummyOptions + from vllm.config.multimodal import MultiModalDummyOptions logger = init_logger(__name__) @@ -195,7 +194,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, "BaseDummyOptions"], + mm_options: "MultiModalDummyOptions", ) -> MultiModalDataDict: data = MultiModalDataBuiltins() if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): @@ -204,21 +203,16 @@ def get_dummy_mm_data( chunk_length = getattr(sub, "chunk_length", None) if sub else None if chunk_length is None: chunk_length = 30 - audio_len = int(chunk_length * sampling_rate) - audio_overrides = mm_options.get("audio") - assert audio_overrides is None or isinstance( - audio_overrides, AudioDummyOptions - ) data["audio"] = self._get_dummy_audios( - length=audio_len, + length=int(chunk_length * sampling_rate), num_audios=num_audios, - overrides=audio_overrides, + overrides=mm_options.get("audio"), ) if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): - target_width, target_height = self.info.get_image_size_with_most_features() + width, height = self.info.get_image_size_with_most_features() data["image"] = self._get_dummy_images( - width=target_width, - height=target_height, + width=width, + height=height, num_images=num_images, overrides=mm_options.get("image"), ) diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index 3a1d017f4c46..bb92833f7799 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -16,7 +16,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn @@ -169,7 +169,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() @@ -177,15 +177,12 @@ def get_dummy_mm_data( audio_len = ( feature_extractor.chunk_length * sampling_rate * _MAX_ENCODER_BATCH_SIZE ) - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index ad54e8eb20bc..0b8d889bbc69 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -19,7 +19,7 @@ from transformers import WhisperConfig from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt from vllm.logger import init_logger @@ -137,19 +137,15 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - target_length = self.info.get_max_audio_array_len() - audio_overrides = mm_options.get("audio") - return { "audio": self._get_dummy_audios( length=target_length, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } @@ -157,7 +153,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, mm_data: MultiModalDataDict | None = None, ) -> ProcessorInputs: tokenizer = self.info.get_tokenizer() diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index b50d04d04b59..59268619b0a4 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -19,7 +19,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import ( @@ -724,21 +724,18 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: feature_extractor = self.info.get_feature_extractor() sampling_rate = feature_extractor.sampling_rate audio_len = feature_extractor.chunk_length * sampling_rate - num_audios = mm_counts.get("audio", 0) - - audio_overrides = mm_options.get("audio") return { "audio": self._get_dummy_audios( length=audio_len, - num_audios=num_audios, - overrides=audio_overrides, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), ) } diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index 919146c551a4..d28848bb8f8c 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -22,7 +22,7 @@ import math from collections.abc import Mapping, Sequence -from typing import Any, cast +from typing import Any import numpy as np import torch @@ -30,7 +30,7 @@ from transformers import BatchFeature from typing_extensions import assert_never -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import ImageSize, MultiModalDataItems @@ -375,7 +375,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: size = self.info.get_image_size_with_most_features() return { @@ -383,7 +383,7 @@ def get_dummy_mm_data( width=size.width, height=size.height, num_images=mm_counts.get("image", 0), - overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + overrides=mm_options.get("image"), ), } diff --git a/vllm/models/deepseek_v41/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py index 4f87b816c9a5..7710edb43f85 100644 --- a/vllm/models/deepseek_v41/common/mm_preprocess.py +++ b/vllm/models/deepseek_v41/common/mm_preprocess.py @@ -19,14 +19,14 @@ import math from collections.abc import Mapping, Sequence -from typing import Any, cast +from typing import Any import numpy as np import torch from PIL import Image, ImageOps from transformers import BatchFeature -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import ImageSize, MultiModalDataItems @@ -296,7 +296,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: size = self.info.get_image_size_with_most_features() return { @@ -304,7 +304,7 @@ def get_dummy_mm_data( width=size.width, height=size.height, num_images=mm_counts.get("image", 0), - overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + overrides=mm_options.get("image"), ), } diff --git a/vllm/models/dots3_note/common/processor.py b/vllm/models/dots3_note/common/processor.py index addaf77c8141..4589874ec0a1 100644 --- a/vllm/models/dots3_note/common/processor.py +++ b/vllm/models/dots3_note/common/processor.py @@ -13,10 +13,7 @@ from transformers import BatchFeature from vllm.config.multimodal import ( - AudioDummyOptions, - BaseDummyOptions, - ImageDummyOptions, - VideoDummyOptions, + MultiModalDummyOptions, ) from vllm.inputs import MultiModalDataDict from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems @@ -584,7 +581,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: data: dict[str, Any] = {} num_images = mm_counts.get("image", 0) @@ -594,7 +591,7 @@ def get_dummy_mm_data( width=width, height=height, num_images=num_images, - overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + overrides=mm_options.get("image"), ) num_audios = mm_counts.get("audio", 0) if num_audios: @@ -611,7 +608,7 @@ def get_dummy_mm_data( data["audio"] = self._get_dummy_audios( length=min(chunk_samples, max(1, seq_len) * stride), num_audios=num_audios, - overrides=cast(AudioDummyOptions | None, mm_options.get("audio")), + overrides=mm_options.get("audio"), ) num_videos = mm_counts.get("video", 0) if num_videos: @@ -625,7 +622,7 @@ def get_dummy_mm_data( height=height, num_frames=num_frames, num_videos=num_videos, - overrides=cast(VideoDummyOptions | None, mm_options.get("video")), + overrides=mm_options.get("video"), ) return data diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py index cc88c1252951..a949e4ec7403 100644 --- a/vllm/models/inkling/common/mm_preprocess.py +++ b/vllm/models/inkling/common/mm_preprocess.py @@ -13,9 +13,7 @@ from transformers.feature_extraction_utils import BatchFeature from vllm.config.multimodal import ( - AudioDummyOptions, - BaseDummyOptions, - ImageDummyOptions, + MultiModalDummyOptions, ) from vllm.inputs import MultiModalDataDict from vllm.multimodal.inputs import ( @@ -145,7 +143,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: config = self.info.get_hf_config() num_images = mm_counts.get("image", 0) @@ -156,12 +154,11 @@ def get_dummy_mm_data( patch_size = getattr(config.vision_config, "patch_size", 40) # A square image ~4 patches wide so the dummy emits several patches. side = patch_size * 4 - image_overrides = mm_options.get("image") mm_data["image"] = self._get_dummy_images( width=side, height=side, num_images=num_images, - overrides=cast(ImageDummyOptions | None, image_overrides), + overrides=mm_options.get("image"), ) if num_audios: # Size the dummy at the maximum allowed audio so memory/encoder @@ -169,11 +166,10 @@ def get_dummy_mm_data( params = self.info.get_hf_processor().audio_feature_extractor.params hop = int(round(params.audio_token_duration_s * params.sample_rate)) audio_len = MAX_AUDIO_TOKENS * hop - audio_overrides = mm_options.get("audio") mm_data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, - overrides=cast(AudioDummyOptions | None, audio_overrides), + overrides=mm_options.get("audio"), ) return mm_data diff --git a/vllm/models/kimi_k3/common/mm_preprocess.py b/vllm/models/kimi_k3/common/mm_preprocess.py index cfe603b47b41..dfd90fe9a969 100644 --- a/vllm/models/kimi_k3/common/mm_preprocess.py +++ b/vllm/models/kimi_k3/common/mm_preprocess.py @@ -9,7 +9,7 @@ import torch from transformers import BatchFeature -from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.config.multimodal import ImageDummyOptions, MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.multimodal.inputs import ( @@ -204,7 +204,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: MultiModalDummyOptions | None = None, ) -> MultiModalDataDict: media_proc_cfg = self.info.image_processor.media_proc_cfg max_size = self.info.get_max_image_size( @@ -214,7 +214,6 @@ def get_dummy_mm_data( media_proc_cfg["patch_limit_on_one_side"], media_proc_cfg["fixed_output_tokens"], ) - num_images = mm_counts.get("image", 0) image_overrides = cast( ImageDummyOptions | None, mm_options.get("image") if mm_options else None, @@ -223,7 +222,7 @@ def get_dummy_mm_data( "image": self._get_dummy_images( width=max_size.width, height=max_size.height, - num_images=num_images, + num_images=mm_counts.get("image", 0), overrides=image_overrides, ) } diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py index 4e8b791093c4..958e4f189b78 100644 --- a/vllm/models/minimax_m3/common/mm_preprocess.py +++ b/vllm/models/minimax_m3/common/mm_preprocess.py @@ -11,9 +11,7 @@ from transformers.video_utils import VideoMetadata from vllm.config.multimodal import ( - BaseDummyOptions, - ImageDummyOptions, - VideoDummyOptions, + MultiModalDummyOptions, ) from vllm.inputs import MultiModalDataDict from vllm.multimodal.inputs import ( @@ -283,7 +281,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: size = self.info.get_image_size_with_most_features() video_size = self.info.get_video_size_with_most_features() @@ -293,14 +291,14 @@ def get_dummy_mm_data( width=size.width, height=size.height, num_images=mm_counts.get("image", 0), - overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + overrides=mm_options.get("image"), ), "video": self._get_dummy_videos( width=video_size.width, height=video_size.height, num_frames=num_frames, num_videos=mm_counts.get("video", 0), - overrides=cast(VideoDummyOptions | None, mm_options.get("video")), + overrides=mm_options.get("video"), ), } diff --git a/vllm/multimodal/processing/dummy_inputs.py b/vllm/multimodal/processing/dummy_inputs.py index 7ec90d030c20..53e9f4f5b69d 100644 --- a/vllm/multimodal/processing/dummy_inputs.py +++ b/vllm/multimodal/processing/dummy_inputs.py @@ -10,8 +10,8 @@ from vllm.config.multimodal import ( AudioDummyOptions, - BaseDummyOptions, ImageDummyOptions, + MultiModalDummyOptions, VideoDummyOptions, ) from vllm.inputs import MultiModalDataDict @@ -45,7 +45,7 @@ def get_dummy_mm_data( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> MultiModalDataDict: """Build the multimodal input which, after processing, results in the maximum possible number of placeholder tokens. @@ -65,7 +65,7 @@ def get_dummy_processor_inputs( self, seq_len: int, mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions], + mm_options: MultiModalDummyOptions, ) -> ProcessorInputs: """Build the input which, after processing, results in the maximum possible number of placeholder tokens. @@ -124,9 +124,8 @@ def _get_dummy_images( width: int, height: int, num_images: int, - overrides: BaseDummyOptions | None = None, + overrides: ImageDummyOptions | None = None, ) -> list[Image.Image]: - assert overrides is None or isinstance(overrides, ImageDummyOptions) if num_images == 0: return [] if overrides: From 71fc70d3ae1df53300a23ec69f6d97b7207a109f Mon Sep 17 00:00:00 2001 From: frida-andersson Date: Fri, 18 Sep 2026 18:57:44 +0200 Subject: [PATCH 487/774] [ROCm][DSv4] Fix sparse-indexer logits collapse on gfx950/gfx942 (#50455) Signed-off-by: Frida Andersson Signed-off-by: frida-andersson --- .../attention/test_rocm_triton_attn_dsv4.py | 186 +++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 192 +++++++++++++++++- 2 files changed, 373 insertions(+), 5 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index f0ac4d838511..256b50d49a3b 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1430,3 +1430,189 @@ def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None: # Second call returns the same cached object. assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out + + +@requires_split_decode_arch +@torch.inference_mode() +def test_fp8_paged_mqa_logits_triton_matches_torch_ref() -> None: + """Block-flat Triton decode logits vs the eager torch reference. + + Compare only valid key positions: the kernel skips the -inf pad that + the torch ref writes past seq_len. + """ + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + from vllm.v1.worker.workspace import init_workspace_manager + + device = torch.device("cuda") + init_workspace_manager(device) + torch.manual_seed(0) + + fp8_dtype = current_platform.fp8_dtype() + batch_size = 2 + next_n = 1 + num_heads = 4 + head_size = 128 + block_size = 64 + max_model_len = 256 + seq_lens = [180, 64] + num_pages = [(seq_len + block_size - 1) // block_size for seq_len in seq_lens] + num_blocks = sum(num_pages) + + q = torch.randn( + batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16 + ).to(fp8_dtype) + weights = torch.rand( + batch_size * next_n, num_heads, device=device, dtype=torch.float32 + ) + context_lens = torch.tensor(seq_lens, device=device, dtype=torch.int32) + block_tables = torch.full( + (batch_size, max(num_pages)), 0, device=device, dtype=torch.int32 + ) + page = 0 + for i, n_pages in enumerate(num_pages): + block_tables[i, :n_pages] = torch.arange( + page, page + n_pages, device=device, dtype=torch.int32 + ) + page += n_pages + + kv_cache = torch.empty( + num_blocks, block_size, 1, head_size + 4, device=device, dtype=torch.uint8 + ) + values = torch.randn( + num_blocks, block_size, head_size, device=device, dtype=torch.bfloat16 + ).to(fp8_dtype) + scales = torch.rand( + num_blocks, block_size, device=device, dtype=torch.float32 + ).clamp(min=1e-3) + kv_flat = kv_cache.view(num_blocks, -1) + scale_off = block_size * head_size + kv_flat[:, :scale_off] = values.reshape(num_blocks, -1).view(torch.uint8) + kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_blocks, -1) + + got = mod.rocm_fp8_paged_mqa_logits_triton( + q, kv_cache, weights, context_lens, block_tables, max_model_len + ) + ref = mod.fp8_paged_mqa_logits_torch( + q, kv_cache, weights, context_lens, block_tables, max_model_len + ) + for i, seq_len in enumerate(seq_lens): + torch.testing.assert_close( + got[i, :seq_len], ref[i, :seq_len], atol=2e-3, rtol=2e-3 + ) + + +def test_indexer_k_is_c4a_block_flat_gate() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _indexer_k_is_c4a_block_flat, + ) + + assert not _indexer_k_is_c4a_block_flat(1) + assert not _indexer_k_is_c4a_block_flat(2) + assert _indexer_k_is_c4a_block_flat(4) + assert not _indexer_k_is_c4a_block_flat(128) + + +@requires_split_decode_arch +@torch.inference_mode() +def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None: + """Ratio 1/2 stay on AITER; ratio 4 takes Triton.""" + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + fp8_paged_mqa_logits_torch, + indexer_k_quant_and_cache_triton, + rocm_fp8_paged_mqa_logits, + rocm_fp8_paged_mqa_logits_triton, + ) + from vllm.v1.worker.workspace import init_workspace_manager + + device = torch.device("cuda") + init_workspace_manager(device) + torch.manual_seed(1) + + fp8_dtype = current_platform.fp8_dtype() + batch_size = 1 + next_n = 1 + num_heads = 32 + head_size = 128 + block_size = 64 + seq_len = 180 + max_model_len = 256 + num_pages = (seq_len + block_size - 1) // block_size + num_tokens = num_pages * block_size + + q = torch.randn( + batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16 + ).to(fp8_dtype) + weights = torch.rand( + batch_size * next_n, num_heads, device=device, dtype=torch.float32 + ) + context_lens = torch.tensor([seq_len], device=device, dtype=torch.int32) + block_tables = torch.arange(num_pages, device=device, dtype=torch.int32).view(1, -1) + dummy_sched = torch.zeros(8, 2, device=device, dtype=torch.int32) + + k_bf16 = torch.randn(num_tokens, head_size, device=device, dtype=torch.bfloat16) + slot = torch.arange(num_tokens, device=device, dtype=torch.int32) + cache_shuffle = torch.zeros( + num_pages, block_size, head_size + 4, device=device, dtype=torch.uint8 + ) + indexer_k_quant_and_cache_triton(k_bf16, cache_shuffle, slot, 128, "ue8m0") + cache_shuffle = cache_shuffle.unsqueeze(2) + + aiter = rocm_fp8_paged_mqa_logits( + q, + cache_shuffle, + weights, + context_lens, + block_tables, + dummy_sched, + max_model_len, + compress_ratio=1, + ) + for ratio in (1, 2): + got = rocm_fp8_paged_mqa_logits( + q, + cache_shuffle, + weights, + context_lens, + block_tables, + dummy_sched, + max_model_len, + compress_ratio=ratio, + ) + torch.testing.assert_close(got[:, :seq_len], aiter[:, :seq_len], atol=0, rtol=0) + + cache_flat = torch.empty( + num_pages, block_size, 1, head_size + 4, device=device, dtype=torch.uint8 + ) + values = torch.randn( + num_pages, block_size, head_size, device=device, dtype=torch.bfloat16 + ).to(fp8_dtype) + scales = torch.rand( + num_pages, block_size, device=device, dtype=torch.float32 + ).clamp(min=1e-3) + kv_flat = cache_flat.view(num_pages, -1) + scale_off = block_size * head_size + kv_flat[:, :scale_off] = values.reshape(num_pages, -1).view(torch.uint8) + kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_pages, -1) + + gated_c4a = rocm_fp8_paged_mqa_logits( + q, + cache_flat, + weights, + context_lens, + block_tables, + dummy_sched, + max_model_len, + compress_ratio=4, + ) + triton = rocm_fp8_paged_mqa_logits_triton( + q, cache_flat, weights, context_lens, block_tables, max_model_len + ) + ref = fp8_paged_mqa_logits_torch( + q, cache_flat, weights, context_lens, block_tables, max_model_len + ) + torch.testing.assert_close( + gated_c4a[:, :seq_len], triton[:, :seq_len], atol=0, rtol=0 + ) + torch.testing.assert_close( + gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3 + ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 3d891c8ce0f7..6b1530d74e65 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -66,6 +66,11 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 +def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool: + """V4.0 C4A is block-flat (NORMAL). Ratio 1 and 2 are 16×16 SHUFFLE.""" + return compress_ratio == 4 + + def _get_aiter_top_k_kernel( *, is_prefill: bool, @@ -458,6 +463,7 @@ def cp_gather_indexer_k_quant_cache_triton( token_to_seq: torch.Tensor, block_tile_size: int = 16, head_tile_size: int = 16, + cache_layout: str | None = None, ): num_tokens = k_fp8.size(0) block_size = k_cache.size(1) @@ -470,7 +476,10 @@ def cp_gather_indexer_k_quant_cache_triton( k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32) grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) - layout = "NORMAL" if block_size == 1 else "SHUFFLE" + if cache_layout is None: + layout = "NORMAL" if block_size == 1 else "SHUFFLE" + else: + layout = cache_layout kernel_args = ( k_cache_value, k_cache_scale, @@ -607,6 +616,158 @@ def fp8_paged_mqa_logits_torch( return logits +@triton.jit +def _fp8_paged_mqa_logits_decode_kernel( + q_ptr, # fp8 [B, NEXT_N, H, D] + kv_val_ptr, # fp8, block-flat: [num_blocks, block_size*(D+4)] + kv_scale_ptr, # fp32, block-flat: [num_blocks, block_size*(D+4)//4] + weights_ptr, # fp32 [B*NEXT_N, H] + ctx_lens_ptr, # int32 [B*NEXT_N] if CTX_PER_ROW else [B] + block_tables_ptr, # int32 [B, max_blocks] + logits_ptr, # fp32 [B*NEXT_N, max_model_len] + stride_q_b, + stride_q_n, + stride_q_h, + stride_w_row, + stride_kvblk_fp8, + stride_kvblk_f32, + scale_region_off, # block_size*D // 4 (fp32 offset of scale region within a block) + stride_bt_b, + stride_logits_row, + max_blocks, # block_tables width; guards the block-table gather + max_model_len, + NUM_HEADS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, # paged-cache page size + BLOCK_KV: tl.constexpr, # positions per tile; must divide BLOCK_SIZE + N_SPLITS: tl.constexpr, # KV-tile parallelism factor (grid dim 1) + NEXT_N: tl.constexpr, # query positions per batch (1 = decode; >1 = MTP verify) + CTX_PER_ROW: tl.constexpr, # ctx_lens is already per-(b, n) +): + # scale_region_off = block_size*D//4 needs D % 4 == 0 (D=128). + tl.static_assert(HEAD_SIZE % 4 == 0) + # Grid (B*NEXT_N, N_SPLITS): disjoint KV tiles per program, no host sync. + row = tl.program_id(0) + split = tl.program_id(1) + b = row // NEXT_N + n = row % NEXT_N + if CTX_PER_ROW: + seq_len = tl.load(ctx_lens_ptr + row) + else: + seq_len = tl.load(ctx_lens_ptr + b) - NEXT_N + n + 1 + seq_len = tl.minimum(tl.maximum(seq_len, 0), max_model_len) + + h = tl.arange(0, NUM_HEADS) + d = tl.arange(0, HEAD_SIZE) + # keep q/kv fp8 -> fp8 MFMA (f32 accum), not the slow f32 path + q = tl.load( + q_ptr + b * stride_q_b + n * stride_q_n + h[:, None] * stride_q_h + d[None, :] + ) + w = tl.load(weights_ptr + row * stride_w_row + h).to(tl.float32) # [H] + + kv_col = tl.arange(0, BLOCK_KV) + for kv_start in tl.range(split * BLOCK_KV, seq_len, N_SPLITS * BLOCK_KV): + pos = kv_start + kv_col + # Tiles are page-aligned (BLOCK_KV | BLOCK_SIZE), so a tile is one page. + logical_blk = kv_start // BLOCK_SIZE + blk_ok = logical_blk < max_blocks # guard against a wild page -> fault + mask_pos = (pos < seq_len) & blk_ok + page = tl.load( + block_tables_ptr + b * stride_bt_b + logical_blk, mask=blk_ok, other=0 + ).to(tl.int64) + pos_in_blk = pos - logical_blk * BLOCK_SIZE # [BLOCK_KV] == pos % BLOCK_SIZE + + # block-flat layout: values region, then scales region. + val_off = page * stride_kvblk_fp8 + pos_in_blk[:, None] * HEAD_SIZE + d[None, :] + kv = tl.load( + kv_val_ptr + val_off, mask=mask_pos[:, None], other=0.0 + ) # [BLOCK_KV, D] fp8 + sc = tl.load( + kv_scale_ptr + page * stride_kvblk_f32 + scale_region_off + pos_in_blk, + mask=mask_pos, + other=0.0, + ) # [BLOCK_KV] + + # fp8 upcasts exactly; bit-identical on gfx942, ~1e-5 rel on gfx950 + # (accum order), top-k unchanged. + s = tl.dot(kv, tl.trans(q)) # [BLOCK_KV, H] + s = tl.maximum(s, 0.0) + s = s * w[None, :] + s = tl.sum(s, axis=1) # [BLOCK_KV] + s = s * sc + tl.store(logits_ptr + row * stride_logits_row + pos, s, mask=mask_pos) + + +def rocm_fp8_paged_mqa_logits_triton( + q_fp8: torch.Tensor, + kv_cache_fp8: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + max_model_len: int, +) -> torch.Tensor: + """Triton paged MQA-logits for decode and MTP; matches the torch ref but + has no host sync, so it is safe to capture under a full CUDA graph.""" + batch_size, next_n, num_heads, head_size = q_fp8.shape + block_size = kv_cache_fp8.shape[1] + BLOCK_KV = 64 + assert block_size % BLOCK_KV == 0 + + fp8_dtype = current_platform.fp8_dtype() + num_blocks = kv_cache_fp8.shape[0] + kv_flat = kv_cache_fp8.reshape( + num_blocks, -1 + ) # uint8 [num_blocks, block_size*(D+4)] + kv_val = kv_flat.view(fp8_dtype) # [num_blocks, block_size*(D+4)] fp8 + kv_scale = kv_flat.view(torch.float32) # [num_blocks, block_size*(D+4)//4] fp32 + + cl = context_lens.reshape(-1) + ctx_per_row = not (next_n > 1 and cl.numel() == batch_size) + + max_blocks = block_tables.shape[1] + (out_logits,) = current_workspace_manager().get_simultaneous( + ((batch_size * next_n, max_model_len), torch.float32), + ) + + # Memory-bound over the KV range: split each row's keys across programs so + # few-row / long-context launches still fill the GPU. Cap splits at the + # device CU count (304 on gfx942, 256 on gfx950) rather than a gfx950-sized + # constant. All terms are static at launch, so the grid stays CUDA-graph-safe. + rows = batch_size * next_n + tiles_cap = (max_model_len + BLOCK_KV - 1) // BLOCK_KV + N_SPLITS = max(1, min(max(1, _decode_cu_count()), tiles_cap, 1024 // rows)) + _fp8_paged_mqa_logits_decode_kernel[(rows, N_SPLITS)]( + q_fp8, + kv_val, + kv_scale, + weights, + cl, + block_tables, + out_logits, + q_fp8.stride(0), + q_fp8.stride(1), + q_fp8.stride(2), + weights.stride(0), + kv_val.stride(0), + kv_scale.stride(0), + (block_size * head_size) // 4, + block_tables.stride(0), + out_logits.stride(0), + max_blocks, + max_model_len, + NUM_HEADS=num_heads, + HEAD_SIZE=head_size, + BLOCK_SIZE=block_size, + BLOCK_KV=BLOCK_KV, + N_SPLITS=N_SPLITS, + NEXT_N=next_n, + CTX_PER_ROW=ctx_per_row, + num_warps=4, + num_stages=2, + ) + return out_logits + + @functools.lru_cache def paged_mqa_logits_module(): paged_mqa_logits_module_path = None @@ -632,6 +793,8 @@ def rocm_fp8_paged_mqa_logits( block_tables: torch.Tensor, schedule_metadata: torch.Tensor, max_model_len: int, + *, + compress_ratio: int = 1, ) -> torch.Tensor: """Compute FP8 MQA logits using paged KV-cache. @@ -639,8 +802,7 @@ def rocm_fp8_paged_mqa_logits( q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to `torch.float8_e4m3fn` by caller. kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape - [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last - 4 bytes per (block,pos) store the `float` dequant scale. + [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. context_lens: Tensor of shape [B], dtype int32; effective context length for each batch element. @@ -649,6 +811,7 @@ def rocm_fp8_paged_mqa_logits( schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. + compress_ratio: C4A (4) takes block-flat Triton; 1 and 2 stay on AITER. Returns: Logits tensor of shape [B * next_n, max_model_len], dtype @@ -657,11 +820,26 @@ def rocm_fp8_paged_mqa_logits( """ from vllm._aiter_ops import rocm_aiter_ops - aiter_paged_mqa_logits_module = None - # if rocm_aiter_ops.is_enabled(): batch_size, next_n = q_fp8.shape[:2] block_size = kv_cache_fp8.shape[1] + # C4A only: Flash/DSv3.2 also skip insert but still write SHUFFLE. + if ( + (_ON_GFX950 or _ON_GFX942) + and _indexer_k_is_c4a_block_flat(compress_ratio) + and block_size > 1 + ): + if block_size % 64 == 0: + return rocm_fp8_paged_mqa_logits_triton( + q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len + ) + # Non 64-aligned page size (not used in prod): eager torch ref. + return fp8_paged_mqa_logits_torch( + q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len + ) + + aiter_paged_mqa_logits_module = None + if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled(): aiter_paged_mqa_logits_module = paged_mqa_logits_module() @@ -1128,6 +1306,9 @@ def rocm_aiter_sparse_attn_indexer( chunk.block_table, chunk.cu_seq_lens, token_to_seq=chunk.token_to_seq, + cache_layout=( + "NORMAL" if _indexer_k_is_c4a_block_flat(compress_ratio) else None + ), ) logits = rocm_fp8_mqa_logits( q_fp8[chunk.token_start : chunk.token_end], @@ -1225,6 +1406,7 @@ def rocm_aiter_sparse_attn_indexer( decode_metadata.block_table, decode_metadata.schedule_metadata, max_model_len=max_model_len, + compress_ratio=compress_ratio, ) if candidate_blocks is not None: From 23e26e058839fb2a3e77c78fbf2184f563593227 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:47:48 -0400 Subject: [PATCH 488/774] [KV Offloading] Back-pressure detection and remediation (#50045) Signed-off-by: Bill Nell Signed-off-by: bnellnm <49004751+bnellnm@users.noreply.github.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- .../unit/offloading_connector/test_events.py | 89 +++ .../v1/kv_offload/tiering/p2p/test_manager.py | 34 + .../kv_offload/tiering/test_backpressure.py | 706 ++++++++++++++++++ vllm/v1/kv_offload/tiering/backpressure.py | 384 ++++++++++ vllm/v1/kv_offload/tiering/base.py | 28 +- vllm/v1/kv_offload/tiering/example/manager.py | 4 + vllm/v1/kv_offload/tiering/factory.py | 42 ++ vllm/v1/kv_offload/tiering/fs/manager.py | 15 +- vllm/v1/kv_offload/tiering/manager.py | 52 +- vllm/v1/kv_offload/tiering/metrics.py | 31 + vllm/v1/kv_offload/tiering/obj/manager.py | 10 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 20 +- vllm/v1/kv_offload/tiering/spec.py | 54 ++ 13 files changed, 1464 insertions(+), 5 deletions(-) create mode 100644 tests/v1/kv_offload/tiering/test_backpressure.py create mode 100644 vllm/v1/kv_offload/tiering/backpressure.py diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index da0ccb02b11c..c409fd125d70 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -682,3 +682,92 @@ def test_tiering_accepts_self_describing_kv_events(): assert spec.kv_events_config.enable_kv_cache_events assert spec.kv_events_config.self_describing_kv_events assert tracker.self_describing_enabled + + +def _build_tiering_spec(secondary_tiers, *, top_level_backpressure=None): + vllm_config = create_vllm_config( + block_size=4, + max_num_batched_tokens=16, + disable_hybrid_kv_cache_manager=False, + ) + extra_config = { + "spec_name": "TieringOffloadingSpec", + "cpu_bytes_to_use": 1 << 20, + "secondary_tiers": secondary_tiers, + } + if top_level_backpressure is not None: + extra_config["backpressure"] = top_level_backpressure + vllm_config.kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config=extra_config, + ) + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + return TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + + +def test_partial_tier_backpressure_inherits_top_level_defaults(): + """A partial tier override still inherits missing fields. + + Field-by-field merge means an override that sets only ``high_water_s`` + keeps ``backpressure_cls`` (and any other field) from the top-level + default, so the resolved dict reaching the factory is complete. + """ + spec = _build_tiering_spec( + [{"type": "example", "backpressure": {"high_water_s": 0.1}}], + top_level_backpressure={ + "backpressure_cls": "EMABackpressureDetector", + "high_water_s": 0.5, + "low_water_s": 0.05, + }, + ) + + resolved = spec.secondary_tier_configs[0]["backpressure"] + # Tier override wins for the field it sets. + assert resolved["high_water_s"] == 0.1 + # Missing fields fall back to the top-level default. + assert resolved["backpressure_cls"] == "EMABackpressureDetector" + assert resolved["low_water_s"] == 0.05 + + +def test_tier_override_wins_over_top_level(): + """Precedence is tier override > top-level default, field-by-field.""" + spec = _build_tiering_spec( + [ + { + "type": "example", + "backpressure": {"high_water_s": 0.1, "low_water_s": 0.02}, + } + ], + top_level_backpressure={ + "backpressure_cls": "EMABackpressureDetector", + "high_water_s": 0.5, + "low_water_s": 0.05, + }, + ) + + resolved = spec.secondary_tier_configs[0]["backpressure"] + # Tier override wins for the fields it sets. + assert resolved["high_water_s"] == 0.1 + assert resolved["low_water_s"] == 0.02 + # Fields only in top-level still fill in. + assert resolved["backpressure_cls"] == "EMABackpressureDetector" + + +def test_tier_without_backpressure_stays_unconfigured(): + spec = _build_tiering_spec([{"type": "example"}]) + assert "backpressure" not in spec.secondary_tier_configs[0] diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index aceec1cd8fa1..8e806582fd9f 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -180,6 +180,40 @@ def test_seed_resolved_after_init_none_hash(self, monkeypatch): assert mgr._get_hash_seed() == "random-seed-abc" +class TestBackpressureRejected: + def _patch_transports(self, monkeypatch) -> None: + monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) + monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) + monkeypatch.setattr( + manager_module.FileMapper, + "from_offloading_spec", + lambda **k: SimpleNamespace(get_run_config=lambda: {}), + ) + + def test_explicit_backpressure_detector_rejected(self, monkeypatch): + """An explicitly configured detector must fail fast for P2P. + + The generic store-latency detector cannot fail fast in PD mode and + rendezvous time makes store latency an unreliable pressure signal, + so P2P rejects it rather than degrading silently. + """ + self._patch_transports(monkeypatch) + with pytest.raises(ValueError, match="not supported for the P2P"): + P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + backpressure_detector=object(), + ) + + def test_no_detector_constructs_normally(self, monkeypatch): + self._patch_transports(monkeypatch) + mgr = P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + assert mgr.bp_detector is None + + # --------------------------------------------------------------------------- # Tests for _peer_id_from_params # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_offload/tiering/test_backpressure.py b/tests/v1/kv_offload/tiering/test_backpressure.py new file mode 100644 index 000000000000..b809bc4dc490 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_backpressure.py @@ -0,0 +1,706 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for back-pressure detection in TieringOffloadingManager. + +These tests use a DelayedSecondaryTierManager that holds completed jobs +until explicitly released, allowing precise control over when the manager +observes store completions and their apparent latency. +""" + +import time +from collections.abc import Iterable +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + OffloadPolicy, + ReqContext, + RequestOffloadingContext, + ScheduleEndContext, + make_offload_key, +) +from vllm.v1.kv_offload.tiering.backpressure import ( + DropAccountingPolicy, + DropStorePolicy, + EMABackpressureDetector, + ThrottledDropPolicy, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + SecondaryTierManager, + TieringOffloadingMetrics, + TransferJob, +) +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) + +_BP_EMA_ALPHA = EMABackpressureDetector.DEFAULT_ALPHA +_BP_WARMUP = EMABackpressureDetector.DEFAULT_WARMUP_COMPLETIONS +# Water marks in s/MiB, scaled for the test's tiny 16-byte blocks. +_BP_HIGH_WATER_S = 1000.0 +_BP_LOW_WATER_S = 500.0 + +_CTX = ReqContext(req_id="test") +_MOCK_OFFLOADING_SPEC = MagicMock() + + +def _mock_mmap_region(num_blocks: int, row_bytes: int = 16): + mock = MagicMock() + view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy()) + mock.create_kv_memoryview.return_value = view + return mock + + +def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: + return [make_offload_key(str(i).encode(), 0) for i in int_ids] + + +class DelayedSecondaryTierManager(SecondaryTierManager): + """Secondary tier that holds completed store jobs until released. + + Jobs are stored immediately but their completion is not reported + by get_finished_jobs() until release_jobs() is called. This lets + tests control the apparent latency of store operations by adjusting + submit_time on the JobMetadata before triggering the completion poll. + """ + + def __init__( + self, offloading_spec, primary_kv_view, tier_type, backpressure_detector=None + ): + super().__init__( + offloading_spec, + primary_kv_view, + tier_type, + backpressure_detector=backpressure_detector, + ) + self.blocks: dict[OffloadKey, bool] = {} + self._held_jobs: list[JobResult] = [] + self._released_jobs: list[JobResult] = [] + self._request_policy: OffloadPolicy | None = None + + def lookup(self, key, req_context): + return LookupResult.HIT if key in self.blocks else LookupResult.MISS + + def submit_store(self, job_metadata: TransferJob) -> None: + for key in job_metadata.keys: + self.blocks[key] = True + self._held_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + + def submit_load(self, job_metadata: TransferJob) -> None: + self._released_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + + def get_finished_jobs(self) -> Iterable[JobResult]: + result = self._released_jobs + self._released_jobs = [] + return result + + def release_jobs(self): + self._released_jobs.extend(self._held_jobs) + self._held_jobs.clear() + + def on_new_request(self, req_context): + if self._request_policy is not None: + return RequestOffloadingContext(policy=self._request_policy) + return RequestOffloadingContext() + + def drain_jobs(self): + self.release_jobs() + + def has_pending_work(self): + return bool(self._held_jobs) + + def get_num_blocks(self): + return len(self.blocks) + + +class TestDefaultConfig: + def test_fs_gets_local_watermarks(self): + cfg = EMABackpressureDetector.default_config("fs") + assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S + assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S + + def test_obj_gets_network_watermarks(self): + cfg = EMABackpressureDetector.default_config("obj") + assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S + assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S + + def test_p2p_gets_local_watermarks(self): + cfg = EMABackpressureDetector.default_config("p2p") + assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S + assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S + + def test_fs_with_remote_locality_gets_network_watermarks(self): + cfg = EMABackpressureDetector.default_config("fs", locality="REMOTE") + assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S + assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S + + def test_fs_with_local_locality_gets_local_watermarks(self): + cfg = EMABackpressureDetector.default_config("fs", locality="LOCAL") + assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S + assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S + + +class TestIdleDecay: + """Tests for EMA idle decay — recovery from pressure without completions.""" + + def _make_detector(self, **kwargs): + defaults = dict( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + decay_half_life_s=2.0, + ) + defaults.update(kwargs) + return EMABackpressureDetector(**defaults) + + def test_ema_decays_while_under_pressure(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._ema = _BP_HIGH_WATER_S * 2 + bp._under_pressure = True + bp._last_update = time.monotonic() + + with patch("time.monotonic", return_value=bp._last_update + 4.1): + # ~2.05 half-lives → EMA * 0.24, well below low water (500) + assert bp.is_under_pressure() is False + assert bp.store_latency_ema < _BP_LOW_WATER_S + + def test_ema_stays_pressured_if_not_enough_decay(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._ema = _BP_HIGH_WATER_S * 4 + bp._under_pressure = True + bp._last_update = time.monotonic() + + with patch("time.monotonic", return_value=bp._last_update + 2.0): + # 2s = 1 half-life → EMA * 0.5 = HIGH * 2, still above LOW + assert bp.is_under_pressure() is True + assert bp.store_latency_ema == pytest.approx(_BP_HIGH_WATER_S * 2, rel=0.01) + + def test_no_decay_before_warmup(self): + bp = self._make_detector() + bp._ema = _BP_HIGH_WATER_S * 2 + bp._under_pressure = True + bp._last_update = time.monotonic() + + with patch("time.monotonic", return_value=bp._last_update + 10.0): + # Still in warmup (_completions=0 < _warmup_completions), + # so decay should not be applied in is_under_pressure. + assert bp.is_under_pressure() is True + assert bp.store_latency_ema == _BP_HIGH_WATER_S * 2 + + def test_no_decay_applied_during_new_sample(self): + """on_store_completed applies the new sample directly without + idle decay — decay only runs in is_under_pressure().""" + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._ema = _BP_HIGH_WATER_S * 2 + bp._under_pressure = True + t0 = time.monotonic() + bp._last_update = t0 + + block_bytes = 16 + with patch("time.monotonic", return_value=t0 + 4.0): + bp.on_store_completed(0.0, block_bytes) + # No decay in on_store_completed: EMA = alpha*0 + (1-alpha)*2000 + expected = (1 - _BP_EMA_ALPHA) * _BP_HIGH_WATER_S * 2 + assert bp.store_latency_ema == pytest.approx(expected, rel=0.01) + + def test_reset_clears_last_update(self): + bp = self._make_detector() + bp._last_update = time.monotonic() + bp.reset() + assert bp._last_update == 0.0 + + +class TestBackpressure: + @pytest.fixture + def setup(self): + mock_region = _mock_mmap_region(20) + self.primary = CPUPrimaryTierOffloadingManager( + num_chunks=20, mmap_region=mock_region + ) + mock_view = mock_region.create_kv_memoryview() + self.tier = DelayedSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_view, + tier_type="delayed", + backpressure_detector=EMABackpressureDetector( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + ), + ) + self.manager = TieringOffloadingManager( + primary_tier=self.primary, + secondary_tiers=[self.tier], + ) + + def _start_request(self, ctx=_CTX): + if ctx.req_id not in self.manager._req_state: + self.manager.on_new_request(ctx) + + def _store_blocks(self, keys, ctx=_CTX): + self._start_request(ctx) + result = self.manager.prepare_store(keys, ctx) + assert result is not None + self.manager.complete_store(keys, ctx, success=True) + + def _simulate_on_schedule_end(self): + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + self.manager.on_schedule_end(ctx) + + def _backdate_held_jobs(self, age_s: float): + """Set submit_time on all in-flight store jobs so they appear + to have taken ``age_s`` seconds when completed.""" + self.manager._processed_jobs_this_step = False + now = time.monotonic() + for job_id, meta in self.manager._jobs.items(): + tj = meta.transfer_job + if not tj.is_promotion and tj.submit_time > 0: + tj.submit_time = now - age_s + + def test_ema_updates_on_store_completion(self, setup): + bp = self.tier.bp_detector + fast_latency = 0.01 + # EMA is in s/MiB; with 16-byte blocks, scale = MiB / block_bytes. + scale = EMABackpressureDetector._MIB / self.tier.block_size_bytes + + # During warmup, EMA stays at 0. + for i in range(_BP_WARMUP - 1): + keys = to_keys([300 + i]) + self._store_blocks(keys) + self._backdate_held_jobs(fast_latency) + self.tier.release_jobs() + self._simulate_on_schedule_end() + assert bp.store_latency_ema == 0.0 + + # Final warmup sample seeds EMA with the mean. + keys = to_keys([400]) + self._store_blocks(keys) + self._backdate_held_jobs(fast_latency) + self.tier.release_jobs() + self._simulate_on_schedule_end() + fast_s_per_mib = fast_latency * scale + assert bp.store_latency_ema == pytest.approx(fast_s_per_mib, rel=0.1) + + # After warmup, EMA updates normally. + new_latency = 0.05 + keys = to_keys([401]) + self._store_blocks(keys) + self._backdate_held_jobs(new_latency) + self.tier.release_jobs() + self._simulate_on_schedule_end() + new_s_per_mib = new_latency * scale + expected = _BP_EMA_ALPHA * new_s_per_mib + (1 - _BP_EMA_ALPHA) * fast_s_per_mib + assert bp.store_latency_ema == pytest.approx(expected, rel=0.1) + + def test_pressure_activates_above_high_water(self, setup): + bp = self.tier.bp_detector + + # Warm up the detector with fast completions first. + for i in range(_BP_WARMUP): + keys = to_keys([200 + i]) + self._store_blocks(keys) + self._backdate_held_jobs(0.001) + self.tier.release_jobs() + self._simulate_on_schedule_end() + + keys = to_keys([100]) + self._store_blocks(keys) + self._simulate_on_schedule_end() + + # 0.1s for a single 16-byte block → ~6553 s/MiB, well above high water. + self._backdate_held_jobs(0.1) + self.tier.release_jobs() + self._simulate_on_schedule_end() + + assert bp.store_latency_ema > _BP_HIGH_WATER_S + assert bp.is_under_pressure() is True + + def test_pressure_clears_below_low_water(self, setup): + bp = self.tier.bp_detector + bp.store_latency_ema = _BP_HIGH_WATER_S * 2 + bp._under_pressure = True + bp._completions = _BP_WARMUP + + # Drive several fast completions to bring EMA below low water. + block_id = 100 + for _ in range(50): + if bp.store_latency_ema < _BP_LOW_WATER_S: + break + keys = to_keys([block_id]) + block_id += 1 + self._store_blocks(keys) + # Pressure is on, so the cascade was skipped. Force-submit + # a store job manually to get a completion to feed the EMA. + job_meta = self.manager.create_store_job(keys, _CTX) + job_meta.submit_time = time.monotonic() - 0.001 + self.tier.submit_store(job_meta) + self.tier.release_jobs() + self._simulate_on_schedule_end() + else: + pytest.fail(f"EMA did not recover below low water: {bp.store_latency_ema}") + + assert bp.is_under_pressure() is False + + def test_hysteresis_prevents_oscillation(self, setup): + bp = self.tier.bp_detector + bp._completions = _BP_WARMUP + block_bytes = self.tier.block_size_bytes + + # Set EMA between low and high water marks (in s/MiB). + mid = (_BP_LOW_WATER_S + _BP_HIGH_WATER_S) / 2 + bp.store_latency_ema = mid + bp._last_update = time.monotonic() + + # Raw seconds that produce `mid` s/MiB for a given number of blocks. + def mid_latency_s(num_blocks: int) -> float: + return mid * (num_blocks * block_bytes / EMABackpressureDetector._MIB) + + # When not under pressure, mid-range EMA should not activate. + bp._under_pressure = False + keys = to_keys(range(2)) + self._store_blocks(keys) + self._backdate_held_jobs(mid_latency_s(len(keys))) + self.tier.release_jobs() + self._simulate_on_schedule_end() + assert bp.is_under_pressure() is False + + # When under pressure, mid-range EMA should not deactivate. + bp._under_pressure = True + bp._last_update = time.monotonic() + keys2 = to_keys(range(10, 12)) + # Force-submit since pressure is on (cascade would skip). + for k in keys2: + self.primary.prepare_store([k], _CTX) + self.primary.complete_store([k], _CTX) + job_meta = self.manager.create_store_job(keys2, _CTX) + job_meta.submit_time = time.monotonic() - mid_latency_s(len(keys2)) + self.tier.submit_store(job_meta) + self.tier.release_jobs() + self._simulate_on_schedule_end() + assert bp.is_under_pressure() is True + + def test_stores_skipped_under_pressure(self, setup): + bp = self.tier.bp_detector + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + initial_blocks = self.tier.get_num_blocks() + + keys = to_keys(range(50, 53)) + self._store_blocks(keys) + + # Blocks should be in primary but NOT in secondary. + for k in keys: + assert self.primary.lookup(k, _CTX) is LookupResult.HIT + assert self.tier.get_num_blocks() == initial_blocks + + def test_stores_resume_after_pressure_clears(self, setup): + bp = self.tier.bp_detector + + # Start under pressure — stores skipped. + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + keys1 = to_keys(range(60, 62)) + self._store_blocks(keys1) + assert all(k not in self.tier.blocks for k in keys1) + + # Clear pressure — next store should cascade. + bp._under_pressure = False + keys2 = to_keys(range(70, 72)) + self._store_blocks(keys2) + assert all(k in self.tier.blocks for k in keys2) + + def test_dropped_store_count_tracked(self, setup): + bp = self.tier.bp_detector + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + policy = bp.policy + assert policy.pop_stores_dropped() == (0, 0) + + self._store_blocks(to_keys([80])) + assert policy._stores_dropped == 1 + assert policy._blocks_dropped == 1 + + self._store_blocks(to_keys([81, 82])) + assert policy._stores_dropped == 2 + assert policy._blocks_dropped == 3 + + def test_metrics_reported_via_get_stats(self, setup): + bp = self.tier.bp_detector + bp._under_pressure = True + bp.store_latency_ema = 2.5 + policy = bp.policy + policy._stores_dropped = 5 + policy._blocks_dropped = 12 + + stats = self.manager.get_stats() + assert stats is not None + reduced = stats.reduce() + + ema_key = TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA + stores_key = TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED + blocks_key = TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED + assert reduced[f"{ema_key}:('1:delayed',)"] == pytest.approx(2.5) + assert reduced[f"{stores_key}:('1:delayed',)"] == 5 + assert reduced[f"{blocks_key}:('1:delayed',)"] == 12 + + # Dropped counts reset after get_stats. + assert policy._stores_dropped == 0 + assert policy._blocks_dropped == 0 + + def test_reset_cache_clears_backpressure(self, setup): + bp = self.tier.bp_detector + bp.store_latency_ema = 5.0 + bp._under_pressure = True + + self.manager.reset_cache() + + assert bp.store_latency_ema == 0.0 + assert bp.is_under_pressure() is False + + def test_request_level_tier_respects_backpressure(self, setup): + """Request-level cascade also skips pressured tiers.""" + # Store blocks in primary first. + existing = to_keys(range(3)) + self._store_blocks(existing) + self.tier.release_jobs() + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + + # Make tier request-level for a new request. + self.tier._request_policy = OffloadPolicy.REQUEST_LEVEL + ctx = ReqContext(req_id="req_rl_bp") + self.manager.on_new_request(ctx) + + # Activate pressure. + bp = self.tier.bp_detector + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + initial_blocks = self.tier.get_num_blocks() + + # prepare_store with existing + new blocks. + new = to_keys(range(10, 12)) + all_keys = existing + new + result = self.manager.prepare_store(all_keys, ctx) + assert result is not None + + # Existing blocks would normally cascade to request-level tier, + # but should be skipped under pressure. + assert self.tier.get_num_blocks() == initial_blocks + + def test_loads_continue_under_pressure(self, setup): + """Loads from a pressured tier still work.""" + blocks = to_keys(range(3)) + for b in blocks: + self.tier.blocks[b] = True + + bp = self.tier.bp_detector + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + + # Lookups should still initiate promotion. + for b in blocks: + result = self.manager.lookup(b, _CTX) + assert result is LookupResult.HIT_PENDING + + +class TestThrottledDropPolicy: + """Tests for proportional throttling under pressure.""" + + def _make_detector(self, **kwargs): + defaults = dict( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + cooldown_s=0.0, + ) + defaults.update(kwargs) + return EMABackpressureDetector(**defaults) + + def test_no_drop_when_not_under_pressure(self): + bp = self._make_detector() + for _ in range(20): + assert bp.should_store(1) is True + + def test_no_drop_at_exactly_high_water(self): + bp = self._make_detector() + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S + for _ in range(20): + assert bp.should_store(1) is True + + def test_partial_drop_at_moderate_pressure(self): + bp = self._make_detector() + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 2 + results = [bp.should_store(1) for _ in range(20)] + assert any(results), "some stores should be allowed" + assert not all(results), "some stores should be dropped" + + def test_full_drop_at_severe_pressure(self): + bp = self._make_detector() + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 4 + for _ in range(20): + assert bp.should_store(1) is False + + def test_drop_rate_increases_with_severity(self): + """Higher EMA → higher drop rate.""" + drop_counts = [] + for multiplier in [1.5, 2.0, 2.5, 3.5]: + bp = self._make_detector() + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * multiplier + drops = sum(1 for _ in range(100) if not bp.should_store(1)) + drop_counts.append(drops) + for i in range(len(drop_counts) - 1): + assert drop_counts[i] <= drop_counts[i + 1] + + def test_explicit_drop_store_policy_still_drops_all(self): + """DropStorePolicy (explicit) still drops everything under pressure.""" + bp = EMABackpressureDetector( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + policy=DropStorePolicy(), + cooldown_s=0.0, + ) + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 1.5 + for _ in range(20): + assert bp.should_store(1) is False + + +class TestSharedDropAccounting: + """Both policies share drop accounting via DropAccountingPolicy.""" + + @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) + def test_shared_base(self, policy_cls): + assert issubclass(policy_cls, DropAccountingPolicy) + + @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) + def test_pop_returns_and_resets(self, policy_cls): + policy = policy_cls() + policy.on_store_skipped(3) + policy.on_store_skipped(2) + assert policy.pop_stores_dropped() == (2, 5) + assert policy.pop_stores_dropped() == (0, 0) + + @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) + def test_reset_clears_counters(self, policy_cls): + policy = policy_cls() + policy.on_store_skipped(4) + policy.reset() + assert policy.pop_stores_dropped() == (0, 0) + + def test_throttled_reset_also_clears_call_count(self): + policy = ThrottledDropPolicy() + policy._call_count = 7 + policy.reset() + assert policy._call_count == 0 + + +class TestHealthyBypass: + """Tests for the fast-path bypass when a tier is consistently healthy.""" + + def _make_detector(self, **kwargs): + defaults = dict( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + cooldown_s=0.0, + ) + defaults.update(kwargs) + return EMABackpressureDetector(**defaults) + + def test_not_healthy_initially(self): + bp = self._make_detector() + assert bp.is_healthy() is False + + def test_healthy_after_streak(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._last_update = time.monotonic() + block_bytes = 16 + fast_s_per_mib = _BP_LOW_WATER_S * 0.5 + fast_elapsed = fast_s_per_mib * (block_bytes / EMABackpressureDetector._MIB) + for _ in range(EMABackpressureDetector._HEALTHY_THRESHOLD + _BP_WARMUP): + bp.on_store_completed(fast_elapsed, block_bytes) + assert bp.is_healthy() is True + + def test_healthy_resets_on_high_latency(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + bp._ema = _BP_LOW_WATER_S + bp._last_update = time.monotonic() + block_bytes = 16 + slow_s_per_mib = _BP_HIGH_WATER_S * 2 + slow_elapsed = slow_s_per_mib * (block_bytes / EMABackpressureDetector._MIB) + bp.on_store_completed(slow_elapsed, block_bytes) + assert bp.is_healthy() is False + + def test_healthy_bypasses_policy(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + assert bp.should_store(1) is True + + def test_reset_clears_healthy_streak(self): + bp = self._make_detector() + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + bp.reset() + assert bp._healthy_streak == 0 + assert bp.is_healthy() is False + + +class TestCooldown: + """Tests for the post-pressure cooldown period.""" + + def _make_detector(self, **kwargs): + defaults = dict( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + cooldown_s=1.0, + ) + defaults.update(kwargs) + return EMABackpressureDetector(**defaults) + + def test_cooldown_throttles_stores(self): + bp = self._make_detector() + bp._completions = _BP_WARMUP + bp._under_pressure = True + bp._ema = _BP_HIGH_WATER_S * 2 + bp._last_update = time.monotonic() + + with patch("time.monotonic") as mock_time: + t0 = bp._last_update + 10.0 + mock_time.return_value = t0 + assert bp.is_under_pressure() is False + assert bp._pressure_cleared_at > 0 + + mock_time.return_value = t0 + 0.1 + results = [bp.should_store(1) for _ in range(10)] + assert sum(results) == 5, "cooldown should allow ~50%" + + def test_no_cooldown_when_never_pressured(self): + bp = self._make_detector() + for _ in range(20): + assert bp.should_store(1) is True + + def test_no_cooldown_after_window_expires(self): + bp = self._make_detector() + bp._pressure_cleared_at = time.monotonic() - 2.0 + for _ in range(20): + assert bp.should_store(1) is True + + def test_reset_clears_cooldown(self): + bp = self._make_detector() + bp._pressure_cleared_at = time.monotonic() + bp._cooldown_count = 5 + bp.reset() + assert bp._pressure_cleared_at == 0.0 + assert bp._cooldown_count == 0 diff --git a/vllm/v1/kv_offload/tiering/backpressure.py b/vllm/v1/kv_offload/tiering/backpressure.py new file mode 100644 index 000000000000..deba37c94b2d --- /dev/null +++ b/vllm/v1/kv_offload/tiering/backpressure.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import time +from abc import ABC, abstractmethod + + +class BackpressurePolicy(ABC): + """Decides what action to take when a detector signals pressure.""" + + @abstractmethod + def should_store(self, detector: BackpressureDetector) -> bool: ... + + @abstractmethod + def on_store_skipped(self, num_blocks: int) -> None: ... + + @abstractmethod + def pop_stores_dropped(self) -> tuple[int, int]: + """Return and reset (stores_dropped, blocks_dropped).""" + ... + + @abstractmethod + def reset(self) -> None: ... + + +class DropAccountingPolicy(BackpressurePolicy): + """Base policy that tracks dropped store/block counters. + + Subclasses implement ``should_store``; the shared drop accounting + (``on_store_skipped``/``pop_stores_dropped``/``reset``) lives here so + every policy reports these counters identically. Subclasses that add + their own state override ``reset`` and call ``super().reset()``. + """ + + def __init__(self): + self._stores_dropped: int = 0 + self._blocks_dropped: int = 0 + + def on_store_skipped(self, num_blocks: int) -> None: + self._stores_dropped += 1 + self._blocks_dropped += num_blocks + + def pop_stores_dropped(self) -> tuple[int, int]: + stores, blocks = self._stores_dropped, self._blocks_dropped + self._stores_dropped = 0 + self._blocks_dropped = 0 + return stores, blocks + + def reset(self) -> None: + self._stores_dropped = 0 + self._blocks_dropped = 0 + + +class DropStorePolicy(DropAccountingPolicy): + """Silently drop stores to pressured tiers.""" + + def should_store(self, detector) -> bool: + return not detector.is_under_pressure() + + +class ThrottledDropPolicy(DropAccountingPolicy): + """Drop stores proportionally to pressure severity. + + Instead of dropping all stores when pressure is detected, the drop + rate scales with how far the EMA exceeds the high watermark:: + + overshoot = (ema - high_water) / high_water + drop_rate = min(1.0, overshoot / ramp_factor) + + With the default ``ramp_factor=2.0``, stores ramp from 0% drop at + the high watermark to 100% drop at 3x the high watermark. This + keeps secondary tiers populated under moderate pressure, reducing + cache misses (and therefore recomputation / ITL regression) while + still protecting TTFT under severe congestion. + """ + + def __init__(self, ramp_factor: float = 2.0): + super().__init__() + self._call_count: int = 0 + self._ramp_factor = ramp_factor + + def should_store(self, detector: BackpressureDetector) -> bool: + # Proportional throttling needs the EMA and high watermark, which + # only the EMA detector exposes. + assert isinstance(detector, EMABackpressureDetector), ( + "ThrottledDropPolicy requires an EMABackpressureDetector" + ) + if not detector.is_under_pressure(): + return True + ema = detector.store_latency_ema + high = detector._high + if high <= 0: + return False + overshoot = max(0.0, ema - high) / high + drop_rate = min(1.0, overshoot / self._ramp_factor) + if drop_rate <= 0: + return True + if drop_rate >= 1.0: + return False + self._call_count += 1 + period = max(2, round(1.0 / drop_rate)) + return (self._call_count % period) != 0 + + def reset(self) -> None: + super().reset() + self._call_count = 0 + + +class BackpressureDetector(ABC): + """Observes store completion signals and determines pressure state.""" + + def __init__( + self, + policy: BackpressurePolicy | None = None, + ): + self._policy = policy or ThrottledDropPolicy() + + @property + def policy(self) -> BackpressurePolicy: + return self._policy + + @abstractmethod + def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: ... + + @abstractmethod + def is_under_pressure(self) -> bool: ... + + @abstractmethod + def reset(self) -> None: ... + + def update(self, submit_time: float, num_bytes: int) -> None: + """Update pressure state from a completed store job. + + Args: + submit_time: ``time.monotonic()`` when the job was submitted. + num_bytes: Total bytes written (num_blocks * block_size_bytes). + + """ + if num_bytes <= 0: + return + elapsed = time.monotonic() - submit_time + self.on_store_completed(elapsed, num_bytes) + + def should_store(self, num_blocks: int) -> bool: + """Check policy and record skip if rejected.""" + if not self._policy.should_store(self): + self._policy.on_store_skipped(num_blocks) + return False + return True + + @classmethod + def default_config( + cls, tier_type: str, *, locality: str | None = None + ) -> dict | None: + """Return default constructor kwargs for ``tier_type``, or None. + + Args: + tier_type: The tier type string (e.g. ``"fs"``, ``"obj"``). + locality: Optional locality hint from the tier config + (``"LOCAL"`` or ``"REMOTE"``). When ``"REMOTE"``, tiers + that would normally get local-storage watermarks (e.g. + ``"fs"``) receive network watermarks instead. + + """ + return None + + @property + def stats(self) -> dict[str, float]: + return {} + + +class EMABackpressureDetector(BackpressureDetector): + """EMA of store latency normalized by transfer size. + + The EMA tracks seconds per megabyte (s/MiB) so that the metric is + comparable regardless of how many blocks are in a job or how large + each block is. Water marks are in the same unit. + + Default water marks are derived from fio benchmarks on the WDC H100 + cluster. Two presets are provided: + + LOCAL (NVMe/SSD): NVMe sustains ~5 GB/s writes. + high=0.005 s/MiB (~200 MB/s) catches severe congestion; + low=0.001 s/MiB (~1 GB/s) requires meaningful recovery. + + NETWORK (CephFS, object store, ``obj``tiers, or any tier + with ``"locality": "REMOTE"``): CephFS sustains ~1.5 GB/s. + high=0.020 s/MiB (~50 MB/s); low=0.005 s/MiB (~200 MB/s). + + ``obj`` tiers get NETWORK defaults automatically. An + ``fs`` tier defaults to LOCAL (appropriate for NVMe/SSD); set + ``"locality": "REMOTE"`` in the tier config for network-backed + filesystems like CephFS. + """ + + _MIB = 1 << 20 + + # EMA smoothing factor: higher values (→1) react faster to latency + # spikes but are noisier; lower values (→0) smooth more but lag. + DEFAULT_ALPHA = 0.15 + + # Number of store completions to collect before seeding the EMA. + # During warmup, pressure is never signalled; the EMA is initialized + # to the mean of the warmup samples to avoid cold-start false positives. + DEFAULT_WARMUP_COMPLETIONS = 10 + + # When under pressure and no completions arrive, the EMA decays + # toward zero so the detector can eventually recover. The half-life + # controls how quickly: after this many seconds of silence the EMA + # halves. With a 2 s half-life and high_water=0.005, an EMA frozen + # at 0.010 reaches low_water=0.001 in ~6.6 s. + DEFAULT_DECAY_HALF_LIFE_S = 2.0 + + # After pressure clears, stores are throttled to 50% for this many + # seconds to prevent a burst of stores from re-overwhelming the tier. + DEFAULT_COOLDOWN_S = 1.0 + + # Number of consecutive sub-low-watermark completions before the + # detector enters a "healthy" fast-path that bypasses policy checks. + _HEALTHY_THRESHOLD = 10 + + LOCAL_HIGH_WATER_S = 0.010 + LOCAL_LOW_WATER_S = 0.002 + + NETWORK_HIGH_WATER_S = 0.050 + NETWORK_LOW_WATER_S = 0.010 + + _NETWORK_TIER_TYPES = frozenset({"obj"}) + + @classmethod + def default_config( + cls, tier_type: str, *, locality: str | None = None + ) -> dict | None: + is_remote = locality is not None and locality.upper() == "REMOTE" + if is_remote or tier_type in cls._NETWORK_TIER_TYPES: + return { + "high_water_s": cls.NETWORK_HIGH_WATER_S, + "low_water_s": cls.NETWORK_LOW_WATER_S, + } + return { + "high_water_s": cls.LOCAL_HIGH_WATER_S, + "low_water_s": cls.LOCAL_LOW_WATER_S, + } + + def __init__( + self, + high_water_s: float, + low_water_s: float, + alpha: float = DEFAULT_ALPHA, + warmup_completions: int = DEFAULT_WARMUP_COMPLETIONS, + decay_half_life_s: float = DEFAULT_DECAY_HALF_LIFE_S, + cooldown_s: float = DEFAULT_COOLDOWN_S, + policy: BackpressurePolicy | None = None, + ): + super().__init__(policy=policy) + if not (0 < alpha <= 1): + raise ValueError(f"alpha must be in (0, 1], got {alpha}") + if low_water_s > high_water_s: + raise ValueError( + f"low_water_s ({low_water_s}) must be <= high_water_s ({high_water_s})" + ) + + self._alpha = alpha + self._high = high_water_s + self._low = low_water_s + self._warmup_completions = warmup_completions + self._decay_half_life_s = decay_half_life_s + self._cooldown_s = cooldown_s + # Exponential moving average of store latency in seconds. + self._ema: float = 0.0 + # Whether the tier is currently considered under pressure. + self._under_pressure: bool = False + # Number of store completions observed (for warmup gating). + self._completions: int = 0 + # Samples collected during warmup; used to seed the EMA. + self._warmup_samples: list[float] = [] + # Monotonic timestamp of the last EMA update; used for idle decay. + self._last_update: float = 0.0 + # Consecutive completions with EMA below low watermark. + self._healthy_streak: int = 0 + # Monotonic timestamp when pressure was last cleared (for cooldown). + self._pressure_cleared_at: float = 0.0 + self._cooldown_count: int = 0 + + def _apply_idle_decay(self) -> None: + """Decay the EMA toward zero based on elapsed idle time. + + Uses the formula ``ema * 0.5^(dt / half_life)`` so the EMA + halves every ``decay_half_life_s`` seconds of silence. This + lets the detector recover from pressure when no stores are + attempted (and therefore no completions arrive to update the + EMA organically). + """ + now = time.monotonic() + dt = now - self._last_update + if dt <= 0 or self._decay_half_life_s <= 0: + return + self._ema *= 0.5 ** (dt / self._decay_half_life_s) + self._last_update = now + + def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: + now = time.monotonic() + s_per_mib = elapsed_s / (num_bytes / self._MIB) + self._completions += 1 + if self._completions <= self._warmup_completions: + self._warmup_samples.append(s_per_mib) + if self._completions == self._warmup_completions: + self._ema = sum(self._warmup_samples) / len(self._warmup_samples) + self._last_update = now + return + self._ema = self._alpha * s_per_mib + (1 - self._alpha) * self._ema + self._last_update = now + if self._ema > self._high: + self._under_pressure = True + self._pressure_cleared_at = 0.0 + self._healthy_streak = 0 + elif self._ema < self._low: + if self._under_pressure: + self._pressure_cleared_at = now + self._cooldown_count = 0 + self._under_pressure = False + self._healthy_streak += 1 + else: + self._healthy_streak = 0 + + def is_under_pressure(self) -> bool: + if self._under_pressure and self._completions >= self._warmup_completions: + self._apply_idle_decay() + if self._ema < self._low: + self._under_pressure = False + self._pressure_cleared_at = time.monotonic() + self._cooldown_count = 0 + return self._under_pressure + + def is_healthy(self) -> bool: + """True when the tier has been consistently fast. + + After ``_HEALTHY_THRESHOLD`` consecutive store completions with + EMA below the low watermark (and no active pressure), the + detector enters a healthy fast-path: ``should_store()`` returns + ``True`` immediately, skipping the policy check entirely. + """ + return ( + not self._under_pressure and self._healthy_streak >= self._HEALTHY_THRESHOLD + ) + + def should_store(self, num_blocks: int) -> bool: + under_pressure = self.is_under_pressure() + if self.is_healthy(): + return True + if ( + not under_pressure + and self._pressure_cleared_at > 0 + and time.monotonic() - self._pressure_cleared_at < self._cooldown_s + ): + self._cooldown_count += 1 + if self._cooldown_count % 2 == 0: + self._policy.on_store_skipped(num_blocks) + return False + return super().should_store(num_blocks) + + def reset(self) -> None: + self._ema = 0.0 + self._under_pressure = False + self._completions = 0 + self._warmup_samples.clear() + self._last_update = 0.0 + self._healthy_streak = 0 + self._pressure_cleared_at = 0.0 + self._cooldown_count = 0 + self._policy.reset() + + @property + def store_latency_ema(self) -> float: + return self._ema + + @store_latency_ema.setter + def store_latency_ema(self, value: float) -> None: + self._ema = value + + @property + def stats(self) -> dict[str, float]: + return {"store_latency_ema": self._ema} diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 77300fa707b0..a90d33eb1136 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Abstract interfaces and data types for the secondary tiering layer.""" +import time from abc import ABC, abstractmethod from collections.abc import Collection, Iterable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -20,6 +21,9 @@ RequestOffloadingContext, ScheduleEndContext, ) +from vllm.v1.kv_offload.tiering.backpressure import ( + BackpressureDetector, +) if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( @@ -52,6 +56,11 @@ class TieringOffloadingMetrics: ) ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs" ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs" + BACKPRESSURE_STORE_LATENCY_EMA = ( + "vllm:kv_offload_tiering_backpressure_store_latency_ema" + ) + BACKPRESSURE_STORES_DROPPED = "vllm:kv_offload_tiering_backpressure_stores_dropped" + BACKPRESSURE_BLOCKS_DROPPED = "vllm:kv_offload_tiering_backpressure_blocks_dropped" @dataclass @@ -63,6 +72,7 @@ class TransferJob: chunk_ids: np.ndarray is_promotion: bool req_context: ReqContext + submit_time: float = field(default_factory=time.monotonic) @dataclass @@ -77,6 +87,7 @@ class JobResult: # indicated by `success`. Must be a subset of the job's original keys. successful_keys: Collection[OffloadKey] | None = None transfer_time: float | None = None + transfer_bytes: int | None = None class ParentManager(ABC): @@ -136,18 +147,33 @@ def __init__( offloading_spec: "OffloadingSpec", primary_kv_view: memoryview, tier_type: str, + backpressure_detector: BackpressureDetector | None = None, ) -> None: """Args: offloading_spec: Offloading configuration. primary_kv_view: Memoryview of the primary tier's CPU KV cache. tier_type: Tier type identifier, set by SecondaryTierFactory from the registered tier type. + backpressure_detector: Optional `BackpressureDetector`. """ self._offloading_spec = offloading_spec self._primary_kv_view: memoryview = primary_kv_view + assert primary_kv_view.strides is not None, ( + "primary_kv_view.strides cannot be None" + ) + self._block_size_bytes: int = primary_kv_view.strides[0] self.tier_type = tier_type self.locality: Locality | None = None + self._bp_detector = backpressure_detector + + @property + def block_size_bytes(self) -> int: + return self._block_size_bytes + + @property + def bp_detector(self) -> BackpressureDetector | None: + return self._bp_detector @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 5ec13f62daac..e3f7c93d7adb 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector class ExampleSecondaryTierManager(SecondaryTierManager): @@ -49,6 +50,7 @@ def __init__( primary_kv_view: memoryview, tier_type: str, custom_param: int = 0, + backpressure_detector: "BackpressureDetector | None" = None, ): """Initialize the example secondary tier. @@ -57,12 +59,14 @@ def __init__( primary_kv_view: Memoryview over the primary tier's KV buffer. tier_type: Name identifying this tier type. custom_param: Dummy parameter demonstrating custom args. + backpressure_detector: Optional backpressure detector. """ super().__init__( offloading_spec=offloading_spec, primary_kv_view=primary_kv_view, tier_type=tier_type, + backpressure_detector=backpressure_detector, ) logger.info( diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 086392a4065b..6a460724de94 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -8,11 +8,33 @@ from vllm.v1.kv_offload.tiering.base import SecondaryTierManager if TYPE_CHECKING: + from vllm.v1.kv_offload.backpressure import BackpressureDetector from vllm.v1.kv_offload.base import OffloadingSpec logger = init_logger(__name__) +_DEFAULT_PACKAGE = "vllm.v1.kv_offload.tiering.backpressure" + + +def _import_class(name: str) -> type: + """Import a class by name. + + If ``name`` contains a dot it is treated as a fully qualified path + (e.g. ``'my_package.MyDetector'``). Otherwise it is looked up in + the default backpressure module + (``vllm.v1.kv_offload.tiering.backpressure``), so short names like + ``'DropStorePolicy'`` or ``'EMABackpressureDetector'`` work out of + the box. + """ + if "." in name: + module_path, _, class_name = name.rpartition(".") + else: + module_path, class_name = _DEFAULT_PACKAGE, name + module = importlib.import_module(module_path) + return getattr(module, class_name) + + class SecondaryTierFactory: """Registry for SecondaryTierManager implementations, resolved by type. @@ -47,10 +69,30 @@ def create_secondary_tier( config = tier_config.copy() tier_type = config.pop("type") config.pop("module_path", None) + locality = config.get("locality") + bp_config = config.pop("backpressure", None) + bp_detector = None + if bp_config is not None: + bp_config = bp_config.copy() + bp_cls_name = bp_config.pop("backpressure_cls", None) + if bp_cls_name is None: + raise ValueError( + "backpressure config must include 'backpressure_cls' " + "(e.g. 'EMABackpressureDetector')" + ) + policy_cls_name = bp_config.pop("policy_cls", None) + policy = _import_class(policy_cls_name)() if policy_cls_name else None + detector_cls: type[BackpressureDetector] = _import_class(bp_cls_name) + defaults = detector_cls.default_config(tier_type, locality=locality) + if defaults: + for k, v in defaults.items(): + bp_config.setdefault(k, v) + bp_detector = detector_cls(policy=policy, **bp_config) return tier_cls( offloading_spec=offloading_spec, primary_kv_view=primary_kv_view, tier_type=tier_type, + backpressure_detector=bp_detector, **config, ) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 3d51d95774cd..6522cfd23646 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -40,6 +40,7 @@ ) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager +from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector from vllm.v1.kv_offload.tiering.base import ( JobId, JobResult, @@ -115,6 +116,7 @@ def __init__( n_write_threads: int = 16, enable_kv_events: bool = False, locality: str | None = None, + backpressure_detector: BackpressureDetector | None = None, ): """Args: offloading_spec: Contains normalized offloading configuration and @@ -129,9 +131,12 @@ def __init__( cache events are enabled globally (kv_events_config). locality: Whether this tier's storage is LOCAL or REMOTE relative to the publishing vLLM instance. + backpressure_detector: Optional backpressure detector. """ - super().__init__(offloading_spec, primary_kv_view, tier_type) + super().__init__( + offloading_spec, primary_kv_view, tier_type, backpressure_detector + ) self.locality = Locality(locality) if locality is not None else None self.events: list[OffloadingEvent] | None = None @@ -150,6 +155,8 @@ def __init__( # Keys of in-flight load (promotion) jobs, so a failed load can mark # its own cached lookup verdicts False (see get_finished_jobs). self._load_job_keys: dict[JobId, list[OffloadKey]] = {} + # Block count per in-flight job, used to report transfer_bytes. + self._job_block_counts: dict[JobId, int] = {} # Per load job: how many blocks loaded before a failure (partial keep). # Written by the pool worker inside the load task before it raises (so # before task_done publishes the job); read on the scheduler thread in @@ -225,6 +232,7 @@ def submit_store(self, job_metadata: TransferJob) -> None: self._block_size, self._use_o_direct, ) + self._job_block_counts[job_metadata.job_id] = len(keys) self._pool.enqueue_store(job_metadata.job_id, 1, [task]) @override @@ -234,6 +242,7 @@ def submit_load(self, job_metadata: TransferJob) -> None: # keys as a miss (see get_finished_jobs). keys = list(job_metadata.keys) self._load_job_keys[job_id] = keys + self._job_block_counts[job_id] = len(keys) paths = [self.file_mapper.get_file_name(key) for key in keys] offsets = [int(cid) * self._block_size for cid in job_metadata.chunk_ids] @@ -272,6 +281,8 @@ def get_finished_jobs(self) -> Iterable[JobResult]: as a miss here (scheduler thread).""" results = [] for job_id, success, transfer_time in self._pool.get_finished(): + block_count = self._job_block_counts.pop(job_id, 0) + transfer_bytes = block_count * self._block_size if block_count else None if self.events is not None: keys = self._store_job_keys.pop(job_id, None) if success and keys: @@ -299,6 +310,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]: success=False, successful_keys=tuple(successful) if successful else None, transfer_time=transfer_time, + transfer_bytes=transfer_bytes, ) ) continue @@ -307,6 +319,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]: job_id=job_id, success=success, transfer_time=transfer_time, + transfer_bytes=transfer_bytes, ) ) return results diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 6e0bf90a332a..ce518042b656 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -238,6 +238,10 @@ def __init__( for tier_idx, tier in enumerate(self.secondary_tiers) } + self._tier_index: dict[SecondaryTierManager, int] = { + tier: i for i, tier in enumerate(self.secondary_tiers) + } + @property def _transfer_jobs(self) -> dict[JobId, JobMetadata]: return self._jobs @@ -337,6 +341,43 @@ def _process_finished_jobs(self): self.primary_tier.complete_read( transfer_job.keys, transfer_job.req_context ) + if completed_job.success: + self._update_backpressure(tier, job_metadata, completed_job) + + def _should_store_to_tier( + self, tier: SecondaryTierManager, num_blocks: int + ) -> bool: + detector = tier.bp_detector + if detector is None: + return True + return detector.should_store(num_blocks) + + def _update_backpressure( + self, + tier: SecondaryTierManager, + job_metadata: JobMetadata, + completed_job: JobResult, + ) -> None: + detector = tier.bp_detector + if detector is None: + return + was_under_pressure = detector.is_under_pressure() + tj = job_metadata.transfer_job + num_bytes = ( + completed_job.transfer_bytes + if completed_job.transfer_bytes is not None + else len(tj.keys) * tier.block_size_bytes + ) + detector.update(tj.submit_time, num_bytes) + if detector.is_under_pressure() != was_under_pressure: + tier_idx = self._tier_index[tier] + logger.info( + "Tier #%d (%s) back-pressure %s (stats=%s)", + tier_idx, + tier.tier_type, + "activated" if detector.is_under_pressure() else "cleared", + detector.stats, + ) @override def lookup( @@ -643,8 +684,10 @@ def _cascade_existing_chunks_to_request_level_tiers( return for tier_idx in request_level_tiers: - job_metadata = self.create_store_job(ready_keys, req_context, tier_idx) tier = self.secondary_tiers[tier_idx] + if not self._should_store_to_tier(tier, len(ready_keys)): + continue + job_metadata = self.create_store_job(ready_keys, req_context, tier_idx) tier.submit_store(job_metadata) def _flush_pending_cascades(self) -> None: @@ -698,6 +741,8 @@ def complete_store( # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier_idx, tier in enumerate(self.secondary_tiers): + if not self._should_store_to_tier(tier, len(keys)): + continue job_metadata = self.create_store_job(keys, req_context, tier_idx) tier.submit_store(job_metadata) @@ -904,6 +949,10 @@ def reset_cache(self) -> None: del self._req_state[req_id] self._processed_jobs_this_step = False + for tier in self.secondary_tiers: + if tier.bp_detector is not None: + tier.bp_detector.reset() + @override def get_stats(self) -> OffloadingConnectorStats | None: stats = self.primary_tier.get_stats() @@ -911,6 +960,7 @@ def get_stats(self) -> OffloadingConnectorStats | None: if stats is not None and stats.is_empty(): stats = None + self._metrics.record_backpressure(self.secondary_tiers) metrics_stats = self._metrics.take_stats() if metrics_stats is not None: if stats is None: diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 55d56999398f..240ea57dd70a 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -12,6 +12,7 @@ from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.tiering.base import ( JobResult, + SecondaryTierManager, TieringOffloadingMetrics, TransferJob, ) @@ -130,6 +131,36 @@ def on_promotion_allocation_failure(self) -> None: TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES ) + def record_backpressure( + self, + tiers: list[SecondaryTierManager], + ) -> None: + for i, tier in enumerate(tiers): + detector = tier.bp_detector + if detector is None: + continue + label = self.tier_label(i) + ema = detector.stats.get("store_latency_ema") + if ema is not None: + self._stats.set_gauge( + TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA, + ema, + labelvalues=label, + ) + stores_dropped, blocks_dropped = detector.policy.pop_stores_dropped() + if stores_dropped > 0: + self._stats.increase_counter( + TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED, + stores_dropped, + labelvalues=label, + ) + if blocks_dropped > 0: + self._stats.increase_counter( + TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED, + blocks_dropped, + labelvalues=label, + ) + def take_stats(self) -> OffloadingConnectorStats | None: active_transfer_stats = OffloadingConnectorStats() self._observe_active_transfer_stats(active_transfer_stats) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 3d6c66c4834a..1886e3b81168 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -22,6 +22,7 @@ ) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager +from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector from vllm.v1.kv_offload.tiering.base import ( JobId, JobResult, @@ -112,6 +113,7 @@ def __init__( io_threads: int = 4, enable_kv_events: bool = False, locality: str | None = None, + backpressure_detector: BackpressureDetector | None = None, ): """Args: offloading_spec: Offloading configuration. @@ -125,9 +127,15 @@ def __init__( cache events are enabled globally (kv_events_config). locality: Whether this tier's storage is LOCAL or REMOTE relative to the publishing vLLM instance. + backpressure_detector: Optional backpressure detector. """ - super().__init__(offloading_spec, primary_kv_view, tier_type) + super().__init__( + offloading_spec, + primary_kv_view, + tier_type, + backpressure_detector, + ) self.locality = Locality(locality) if locality is not None else None self.events: list[OffloadingEvent] | None = None diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 703c632899d9..59b9c320695f 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -265,7 +265,25 @@ def __init__( number, or anything convertible to one. """ - super().__init__(offloading_spec, primary_kv_view, tier_type) + backpressure_detector = kwargs.pop("backpressure_detector", None) + if backpressure_detector is not None: + # The generic (store-latency) detector is unreliable for P2P: in + # PD mode dropping a store can leave the decoder waiting until the + # load timeout instead of failing fast, and rendezvous time makes + # store latency a poor pressure signal. Reject it until a + # P2P-specific fail-fast path exists. + raise ValueError( + "Backpressure is not supported for the P2P secondary tier. " + "The generic store-latency detector cannot fail fast in PD " + "mode and rendezvous time makes store latency an unreliable " + "pressure signal. Remove the 'backpressure' config from the " + "p2p tier." + ) + super().__init__( + offloading_spec, + primary_kv_view, + tier_type, + ) try: timeout_s = float(unbound_store_timeout_s) except (TypeError, ValueError): diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 9adac0f787a6..93c6cbabcb85 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -247,6 +247,34 @@ def build_metric_definitions( assert isinstance(tier_config, dict) tier_cls = SecondaryTierFactory.get_tier_class(tier_config) metrics.update(tier_cls.build_metric_definitions(tier_config)) + + metrics[TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA] = ( + OffloadingGaugeMetadata( + documentation=( + "Exponential moving average of store latency " + "for back-pressure detection, in s/MiB." + ), + labelnames=("tier",), + ) + ) + metrics[TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED] = ( + OffloadingCounterMetadata( + documentation=( + "Number of store operations dropped due to " + "back-pressure on a secondary tier." + ), + labelnames=("tier",), + ) + ) + metrics[TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED] = ( + OffloadingCounterMetadata( + documentation=( + "Number of blocks dropped due to back-pressure on a secondary tier." + ), + labelnames=("tier",), + ) + ) + return metrics def __init__(self, config: OffloadingConfig): @@ -259,6 +287,32 @@ def __init__(self, config: OffloadingConfig): if not isinstance(self.secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") + # Backpressure config is merged field-by-field in priority order + # (highest first): + # 1. Per-tier ``backpressure`` dict in the tier config + # 2. Top-level ``backpressure`` in kv_connector_extra_config + # Merging per field (rather than per whole dict) means a partial + # tier override still inherits missing fields from the top-level + # default, so the resolved dict reaching the factory is complete. + # Within each tier's resolved dict, tier-type-aware water marks + # are filled in last so a bare ``"backpressure": {}`` picks up + # sensible thresholds for the storage medium. + bp_defaults = self.extra_config.get("backpressure") + + for tier_cfg in self.secondary_tier_configs: + tier_override = tier_cfg.get("backpressure") + # Overlay from lowest to highest precedence so higher-precedence + # fields win while lower-precedence ones fill in the gaps. + merged: dict[str, Any] = {} + for source in (bp_defaults, tier_override): + if source: + merged.update(source) + # Only set a resolved dict when at least one source contributed + # (or an explicit ``backpressure`` key was present, e.g. ``{}``), + # so tiers without any backpressure config stay unconfigured. + if merged or "backpressure" in tier_cfg: + tier_cfg["backpressure"] = merged + # Scheduler-side mmap (rank=None); kept for cleanup self._scheduler_mmap: SharedOffloadRegion | None = None From 50d812b66a697c464a9964b3a061b47583b83e34 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 18 Sep 2026 12:57:19 -0700 Subject: [PATCH 489/774] [Perf][DSV4.1] Optimize MegaMoE staging and NVFP4 cache gathers (#57604) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- tests/kernels/test_compressor_kv_cache.py | 95 +++++++++++++------ tests/models/test_deepseek_v4_mega_moe.py | 22 +++-- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 67 ++++++++----- .../deepseek_v41/common/ops/cache_utils.py | 6 +- .../nvidia/flash_mla_mega_attn.py | 2 +- 5 files changed, 130 insertions(+), 62 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 95ff6b3c057f..06897f04c0b7 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -473,15 +473,18 @@ def test_v41_mxfp8_cache_round_trip(): def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - """Unpack one 256-byte e2m1 row and apply its 32 e4m3 tile scales.""" + """Unpack e2m1 rows and apply their e4m3 tile scales.""" mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device) - codes = torch.empty(512, dtype=torch.uint8, device=packed.device) - codes[0::2] = packed & 0xF # even element in the low nibble - codes[1::2] = packed >> 4 + codes = torch.empty( + (*packed.shape[:-1], 512), dtype=torch.uint8, device=packed.device + ) + codes[..., 0::2] = packed & 0xF # even element in the low nibble + codes[..., 1::2] = packed >> 4 vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0) return ( - vals.view(32, 16) * scales.view(torch.float8_e4m3fn).float().view(32, 1) - ).flatten() + vals.reshape(*packed.shape[:-1], 32, 16) + * scales.view(torch.float8_e4m3fn).float().reshape(*packed.shape[:-1], 32, 1) + ).flatten(-2) @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @@ -543,7 +546,10 @@ def test_v41_rope_insert_nvfp4_record(compress_ratio: int): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") -def test_v41_nvfp4_gather_matches_insert(): +@pytest.mark.parametrize("num_reqs", [1, 8, 32]) +@pytest.mark.parametrize("num_tokens", [70, 128, 129, 2051]) +@pytest.mark.parametrize("partial", [False, True]) +def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial): """The NVFP4 gather dequantizes exactly what the insert kernel wrote.""" from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( @@ -553,38 +559,73 @@ def test_v41_nvfp4_gather_matches_insert(): torch.manual_seed(17) device = "cuda" block_size = 64 - num_tokens = 70 - num_blocks = 4 + num_blocks = math.ceil(num_tokens / block_size) page_bytes = math.ceil(block_size * 288 / 512) * 512 positions = torch.arange(num_tokens, dtype=torch.int64, device=device) latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) - angles = torch.randn(128, 32, device=device) + angles = torch.randn(num_tokens, 32, device=device) cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1)) slots = torch.arange(num_tokens, dtype=torch.int64, device=device) rope_quant_insert(latent, positions, cos_sin, cache, slots, 1) - out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) - dequantize_and_gather_k_cache( - out, - cache, - seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), - gather_lens=None, - block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( - 1, -1 - ), - block_size=block_size, - offset=0, + offset = 7 if partial else 0 + seq_lens = num_tokens - torch.arange(num_reqs, dtype=torch.int32, device=device) + gather_lens = seq_lens - 3 if partial else None + storage = torch.full( + (num_reqs, num_tokens + offset + 16, 512), + -123.0, + dtype=torch.bfloat16, + device=device, ) - for t in (0, 1, block_size, num_tokens - 1): - page, row = divmod(t, block_size) - expected = _decode_nvfp4_row( - backing[page, row * 256 : (row + 1) * 256], - backing[page, block_size * 256 + row * 32 :][:32], + out = storage[:, : num_tokens + offset] + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).repeat( + num_reqs, 1 + ) + + def gather(): + dequantize_and_gather_k_cache( + out, + cache, + seq_lens=seq_lens, + gather_lens=gather_lens, + block_table=block_table, + block_size=block_size, + offset=offset, ) - torch.testing.assert_close(out[0, t].float(), expected, rtol=0, atol=0) + + expected = _decode_nvfp4_row( + backing[:, : block_size * 256].reshape(-1, 256), + backing[:, block_size * 256 : block_size * 288].reshape(-1, 32), + ) + + def check(shorter_by): + assert (storage[:, num_tokens + offset :] == -123).all() + for req in range(num_reqs): + start = 3 if partial else 0 + length = num_tokens - req - start - shorter_by + assert (out[req, :offset] == -123).all() + assert (out[req, offset + length :] == -123).all() + torch.testing.assert_close( + out[req, offset : offset + length].float(), + expected[start : start + length], + rtol=0, + atol=0, + ) + + gather() + check(0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + gather() + storage.fill_(-123) + seq_lens.sub_(1) + if gather_lens is not None: + gather_lens.sub_(1) + graph.replay() + check(1) @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index ad029160c5af..f87dd75fb220 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -686,12 +686,14 @@ def forward(self, hidden_states): not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", ) -def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(): +@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) +@pytest.mark.parametrize("hidden_size", [256, 5120]) +def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact( + num_tokens, hidden_size +): from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 device = torch.device("cuda") - num_tokens = 7 - hidden_size = 256 top_k = 8 generator = torch.Generator(device=device) @@ -846,6 +848,7 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) assert torch.equal(fused_x_sf, ref_x_sf) assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) + assert torch.all(fused_shared_x_sf[~populated] == -1) def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): @@ -884,11 +887,13 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", ) @pytest.mark.parametrize("nonfinite_padding", [False, True]) -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_padding): +@pytest.mark.parametrize("num_tokens", [7, 65, 135]) +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding( + nonfinite_padding, num_tokens +): from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 device = torch.device("cuda") - num_tokens = 7 hidden_size = 256 top_k = 8 @@ -916,10 +921,9 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_paddin dtype=torch.float32, generator=generator, ) - is_padding = torch.tensor( - [False, True, False, False, True, False, True], - device=device, - ) + is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) + is_padding[1::3] = True + is_padding[-1] = True if nonfinite_padding: topk_weights[is_padding] = float("nan") diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index d254dad9d6a1..aa9482bed643 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -12,7 +12,7 @@ from vllm.triton_utils import tl, triton -@triton.jit +@triton.jit(do_not_specialize=["num_tokens"]) def _prepare_megamoe_inputs_kernel( hidden_states, x_fp8, @@ -40,27 +40,32 @@ def _prepare_megamoe_inputs_kernel( topk_idx_stride_k: tl.constexpr, topk_weights_out_stride_m: tl.constexpr, topk_weights_out_stride_k: tl.constexpr, + num_tokens, hidden_size: tl.constexpr, top_k: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_K: tl.constexpr, BLOCK_TOPK: tl.constexpr, SHARED_BLOCK_M: tl.constexpr, ) -> None: - token_id = tl.program_id(0) + token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + token_mask = token_id < num_tokens k_block_id = tl.program_id(1) k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = k_offsets < hidden_size + k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size) hidden = tl.load( - hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k, + hidden_states + + token_id[:, None] * hidden_stride_m + + k_offsets[None, :] * hidden_stride_k, mask=k_mask, other=0.0, ).to(tl.float32) num_groups: tl.constexpr = BLOCK_K // GROUP_K - hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K]) - amax = tl.max(hidden_groups, axis=1) + hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K]) + amax = tl.max(hidden_groups, axis=2) amax = tl.maximum(amax, 1.0e-4) scale = amax / 448.0 @@ -71,21 +76,24 @@ def _prepare_megamoe_inputs_kernel( scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) - hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / rounded_scale)[:, None] - scaled = tl.reshape(scaled, [BLOCK_K]) + hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) + scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None] + scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) fp8 = scaled.to(tl.float8e4nv) tl.store( - x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k, + x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, fp8, mask=k_mask, ) scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32) + packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( + tl.int32 + ) tl.store( x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, packed_scale, + mask=token_mask, ) # DeepGEMM's SM100 shared-expert TMA loads require the activation scales @@ -105,41 +113,48 @@ def _prepare_megamoe_inputs_kernel( + shared_row * shared_x_sf_stride_m + k_block_id * shared_x_sf_stride_k, packed_scale, + mask=token_mask, ) if k_block_id == 0: topk_offsets = tl.arange(0, BLOCK_TOPK) - topk_mask = topk_offsets < top_k - token_is_padding = False + topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k) + token_is_padding = tl.full((BLOCK_M,), False, tl.int1) if is_padding is not None: - token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m) + token_is_padding = tl.load( + is_padding + token_id * is_padding_stride_m, + mask=token_mask, + other=True, + ) ids = tl.load( - topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, + topk_ids + + token_id[:, None] * topk_ids_stride_m + + topk_offsets[None, :] * topk_ids_stride_k, mask=topk_mask, other=0, ).to(tl.int64) - ids = tl.where(token_is_padding, -1, ids) + ids = tl.where(token_is_padding[:, None], -1, ids) tl.store( topk_idx_out - + token_id * topk_idx_stride_m - + topk_offsets * topk_idx_stride_k, + + token_id[:, None] * topk_idx_stride_m + + topk_offsets[None, :] * topk_idx_stride_k, ids, mask=topk_mask, ) weights = tl.load( topk_weights - + token_id * topk_weights_stride_m - + topk_offsets * topk_weights_stride_k, + + token_id[:, None] * topk_weights_stride_m + + topk_offsets[None, :] * topk_weights_stride_k, mask=topk_mask, other=0.0, ) - weights = tl.where(token_is_padding, 0.0, weights) + weights = tl.where(token_is_padding[:, None], 0.0, weights) tl.store( topk_weights_out - + token_id * topk_weights_out_stride_m - + topk_offsets * topk_weights_out_stride_k, + + token_id[:, None] * topk_weights_out_stride_m + + topk_offsets[None, :] * topk_weights_out_stride_k, weights, mask=topk_mask, ) @@ -195,7 +210,9 @@ def prepare_megamoe_inputs( ) block_k = 128 - grid = (num_tokens, triton.cdiv(hidden_size, block_k)) + # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. + block_m = 8 if num_tokens >= 64 else 1 + grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) block_topk = triton.next_power_of_2(top_k) padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 _prepare_megamoe_inputs_kernel[grid]( @@ -225,8 +242,10 @@ def prepare_megamoe_inputs( topk_idx_out.stride(1), topk_weights_out.stride(0), topk_weights_out.stride(1), + num_tokens, hidden_size, top_k, + BLOCK_M=block_m, BLOCK_K=block_k, GROUP_K=32, BLOCK_TOPK=block_topk, diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index e7c3650768c2..04b2a9bac1aa 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -584,7 +584,11 @@ def dequantize_and_gather_k_cache_triton( num_reqs = seq_lens.shape[0] NUM_WORKERS = 128 if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN: - _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, NUM_WORKERS)]( + # Keep short gathers unchanged; cap total workers for multi-request chunks. + num_workers = min( + max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset) + ) + _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)]( out, out.stride(0), out.stride(1), diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 9722673dd1ee..6f4dc7e96c4a 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -422,7 +422,7 @@ def _forward_prefill_mega( if not swa_only: assert flashmla_metadata is not None dequantize_and_gather_k_cache( - kv_ws[:chunk_size], + kv_ws[:chunk_size, :chunk_n], self._compressed_kv_cache(), seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, gather_lens=None, From 63d9ad0a3a435cdf3a44495028b10f390a38f960 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 18 Sep 2026 14:11:27 -0600 Subject: [PATCH 490/774] [Model] DiffusionGemma: honor logprob_token_ids on the converging step (#57417) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 --- vllm/model_executor/models/diffusion_gemma.py | 30 +++++++++++++++++-- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 558bc82bcb10..401d99087066 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -967,6 +967,7 @@ def __init__( tp_group_name: str = "", ): self.sampling_states = sampler.sampling_states + self.logprob_token_ids_state = sampler.logprob_token_ids_state self.req_states = sampler.req_states self.logits_mode = sampler.logprobs_mode in ("raw_logits", "processed_logits") # Self-conditioning soft embed = probs @ embed_weight * normalizer, @@ -1023,9 +1024,11 @@ def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> No # that was aborted between its converging denoise and commit steps. self._pending_logprobs.pop(req_idx, None) self.sampling_states.add_request(req_idx, sampling_params) + self.logprob_token_ids_state.add_request(req_idx, sampling_params) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() + self.logprob_token_ids_state.apply_staged_writes() @property def penalties_state(self): @@ -1214,6 +1217,11 @@ def __call__( slots_np = input_batch.idx_mapping_np[:num_reqs] max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) + # Requests may ask for specific token ids' logprobs instead of, or as + # well as, a top-k. + max_token_ids = self.logprob_token_ids_state.max_num_token_ids(slots_np) + want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0 + num_logprobs = max(max_num_logprobs, 0) # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in # _compiled_sample_step keeps several live [group * CL, vocab] copies, so @@ -1271,7 +1279,7 @@ def __call__( # Logprobs for denoise steps that just converged (is_encoder_phase # flipped False→True), stashed per tile so `scaled` is freed each tile. - if max_num_logprobs >= 0: + if want_logprobs: converged_mask = states.is_encoder_phase[tile_slots] just_converged = converged_mask & ~is_committing[tile] if just_converged.any(): @@ -1285,17 +1293,33 @@ def __call__( # positions are never emitted. k_i = int(valid_canvas_len_np[start_req + li]) pos = li * CL + per_req_ids = max_token_ids > 0 self._pending_logprobs[slot.item()] = compute_topk_scores( flat_logits[pos : pos + k_i], - max_num_logprobs, + num_logprobs, argmax_tokens[local_idx][:k_i], + logprob_token_ids_state=( + self.logprob_token_ids_state if per_req_ids else None + ), + # every row of this stash belongs to one slot + expanded_idx_mapping=( + torch.full( + (k_i,), + slot.item(), + dtype=torch.int32, + device=flat_logits.device, + ) + if per_req_ids + else None + ), + max_per_req_token_ids=max_token_ids, logits_mode=self.logits_mode, ) # Commit steps: is_committing was True at entry. Reassemble previously # stashed logprobs and attach to SamplerOutput. logprobs_tensors = None - if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs: + if want_logprobs and is_committing.any() and self._pending_logprobs: committing_slots = set( decode_slots_np[is_committing.cpu().numpy()].tolist() ) From 49e0427ca5183a56ccb112e5be1300feb2873ca9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 18 Sep 2026 22:33:02 +0200 Subject: [PATCH 491/774] [Bugfix][NIXL] Avoid receive reports for notification-only requests (#57570) Signed-off-by: NickLucche Co-authored-by: OpenAI Codex --- .buildkite/test_areas/disaggregated.yaml | 3 ++ .../unit/test_kv_load_failure_recovery.py | 33 +++++++++++++ .../kv_connector/unit/test_nixl_connector.py | 47 +++++++++++++++++++ .../kv_connector/v1/nixl/pull_worker.py | 10 +++- 4 files changed, 91 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index d217090bc5e7..b187169e7675 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -311,6 +311,9 @@ steps: PREFILLER_TP_SIZE: "2" DECODER_TP_SIZE: "2" MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct" + # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. + PREFILL_BLOCK_SIZE: "2048" + DECODE_BLOCK_SIZE: "2048" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index a54a023a13c4..a8bc2d0e12f0 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -35,6 +35,39 @@ def scheduler(): return create_scheduler(vllm_config) +@pytest.mark.parametrize("policy", ["fail", "recompute"]) +def test_failed_receive_completion_honors_failure_policy(policy): + """Failed receives still need completion to free blocks or resume computation.""" + scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy)) + request = create_request(num_tokens=3 * scheduler.block_size) + scheduler.add_request(request) + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = [ + (2 * scheduler.block_size, True), + (0, False), + ] + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.take_events.return_value = () + block_pool = scheduler.kv_cache_manager.block_pool + free_blocks_before = block_pool.get_num_free_blocks() + scheduler_output = scheduler.schedule() + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + output = create_model_runner_output([], finished_recving={request.request_id}) + assert output.kv_connector_output is not None + output.kv_connector_output.failed_recving = {request.request_id} + scheduler.update_from_output(scheduler_output, output) + + if policy == "fail": + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert block_pool.get_num_free_blocks() == free_blocks_before + else: + assert request.num_computed_tokens == 0 + next_output = scheduler.schedule() + assert request.request_id in next_output.num_scheduled_tokens + + @pytest.mark.parametrize( "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs", [ diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 3b9f6078652c..39f8f88d5c70 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -10,6 +10,7 @@ import time import uuid from collections import defaultdict +from concurrent.futures import Future from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, patch @@ -2899,6 +2900,52 @@ def test_empty_recv_is_reported_only_when_awaited( assert (request_id in done_recving) is awaiting_kvs +@pytest.mark.parametrize("is_hma", [False, True]) +@pytest.mark.parametrize( + "local_block_ids,awaiting_kvs", + [((), False), (([],), False), (([],), True), (([1, 2, 3],), True)], +) +def test_handshake_failure_reports_only_awaited_recvs( + recv_worker, is_hma, local_block_ids, awaiting_kvs +): + """A failed cleanup handshake must not complete an already-aborted request.""" + worker = recv_worker + worker._is_hma_required = is_hma + worker._recving_metadata.clear() + worker._remote_agents = {} + worker._handshake_lock = contextlib.nullcontext() + worker._ready_requests = queue.Queue() + worker._reqs_to_process = set() + worker.pcp_rank = 0 + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="request", + local_block_ids=local_block_ids, + kv_transfer_params={ + "remote_block_ids": ([4, 5, 6],), + "remote_engine_id": "prefill", + "remote_request_id": "prefill-request", + "remote_host": "localhost", + "remote_port": 1234, + }, + awaiting_kvs=awaiting_kvs, + ) + handshake = Future[None]() + handshake.set_exception(RuntimeError("handshake failed")) + with patch.object(worker, "_ensure_handshake", return_value=handshake): + worker.start_load_kv(metadata) + + results = worker.get_transfer_results() + expected = {"request"} if awaiting_kvs else set() + assert results.finished_recving == results.failed_recving == expected + assert worker.get_block_ids_with_load_errors() == ( + set(local_block_ids[0]) if awaiting_kvs and not is_hma else set() + ) + assert not worker._recving_metadata + assert not worker._recv_failures + assert not worker._pending_recv_notifs + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 7711a35954ca..43fbbe0c3c25 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -64,8 +64,14 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): len(meta.local_physical_block_ids), len(meta.remote.block_ids), ) - # always store metadata for failure recovery - self._recving_metadata[req_id] = meta + # Full local hits and aborted cleanup only notify P; no recv is awaited. + # On a full local hit: + # - Notification failure must not fail the request: its KV is local. + # - Receive completion must not be reported: num_external_tokens == 0, + # so the scheduler never entered WAITING_FOR_REMOTE_KVS to begin with. + # Aborted cleanup requests have already been removed from the scheduler. + if meta.awaiting_kvs or any(meta.local_block_ids): + self._recving_metadata[req_id] = meta if remote_engine_id not in self._remote_agents: # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: From 2c3fb4c545005db887788747b48d040d930604df Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 18 Sep 2026 16:13:35 -0500 Subject: [PATCH 492/774] [CI][ROCm] Deprecate DinD for MI250 test groups (#56162) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .../scripts/hardware_ci/run-amd-test.sh | 8 ++++ .buildkite/test-amd.yaml | 22 ++++++++++- .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/distributed.yaml | 1 + .buildkite/test_areas/docker.yaml | 1 + .buildkite/test_areas/engine.yaml | 3 ++ .buildkite/test_areas/kernels.yaml | 2 + .buildkite/test_areas/misc.yaml | 1 + .buildkite/test_areas/models_basic.yaml | 4 +- .buildkite/test_areas/models_language.yaml | 2 + .buildkite/test_areas/models_multimodal.yaml | 2 + .buildkite/test_areas/plugins.yaml | 1 + .../test_areas/rust_frontend_cargo.yaml | 2 + .buildkite/test_areas/samplers.yaml | 1 + tests/conftest.py | 2 +- .../model_executor/layers/test_activation.py | 38 +++++++++++++++++++ .../tools/test_docker_build_metadata_args.py | 7 +++- vllm/model_executor/layers/activation.py | 10 ++++- 18 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 tests/model_executor/layers/test_activation.py diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 541a8f55258c..2705c3b87e7b 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -1612,6 +1612,13 @@ fi echo "Final commands: $commands" +# Match native CPU jobs even when the container can see AMD devices. +cpu_platform_env=() +if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" \ + && "$commands" != *python_only_compile.sh* ]]; then + cpu_platform_env=(-e "VLLM_TARGET_DEVICE=cpu") +fi + standalone_merge_base_env=() if [[ "$commands" == *python_only_compile.sh* ]]; then # The ROCm test image often ships /vllm-workspace without .git. Resolve the @@ -1769,6 +1776,7 @@ else -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \ -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \ -e "PYTORCH_ROCM_ARCH=" \ + "${cpu_platform_env[@]}" \ "${standalone_merge_base_env[@]}" \ --name "${container_name}" \ "${image_name}" \ diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ba8f382755e9..789ca5426034 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -114,6 +114,7 @@ steps: - label: ":computer: (CPU) Basic Models Other" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -124,13 +125,15 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/transformers/fusers/ + - tests/model_executor/layers/test_activation.py - tests/models/test_qwen3_5_mtp_config.py commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py - label: ":computer: (CPU) Multimodal Processor Shard %N" # TBD timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true parallelism: 6 @@ -146,6 +149,7 @@ steps: - label: ":computer: (CPU) V1 Others" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -166,6 +170,7 @@ steps: - label: ":computer: (CPU) Async Engine, Inputs, Utils, Worker, Config" # TBD timeout_in_minutes: 80 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -205,6 +210,7 @@ steps: - label: ":computer: (CPU) Rust Frontend Cargo Style + Clippy" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -220,6 +226,7 @@ steps: - label: ":computer: (CPU) Rust Frontend Cargo" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -235,6 +242,7 @@ steps: - label: ":computer: (CPU) Docker Build Metadata" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 no_gpu: true optional: true @@ -265,6 +273,7 @@ steps: - label: ":amd: (MI250) PyTorch Fullgraph Smoke" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -284,6 +293,7 @@ steps: - label: ":amd: (MI250) Pipeline + Context Parallelism" # TBD timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_4 num_gpus: 4 working_dir: "/vllm-workspace/tests" @@ -311,6 +321,7 @@ steps: timeout_in_minutes: 45 parallelism: 2 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -328,6 +339,7 @@ steps: - label: ":amd: (MI250) Language Models (PPL)" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -359,6 +371,7 @@ steps: - label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -374,6 +387,7 @@ steps: - label: ":amd: (MI250) Batch Invariance" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -392,6 +406,7 @@ steps: - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -412,6 +427,7 @@ steps: - label: ":amd: (MI250) E2E Core" # TBD timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -425,6 +441,7 @@ steps: - label: ":amd: (MI250) E2E Scheduling" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -438,6 +455,7 @@ steps: - label: ":amd: (MI250) V1 Engine" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -460,6 +478,7 @@ steps: - label: ":amd: (MI250) Spec Decode Draft Model" # TBD timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" @@ -477,6 +496,7 @@ steps: - label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index a031c37ba371..8ac8e1960ade 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -49,6 +49,7 @@ steps: mirror: amd: label: ":amd: (MI250) CUDAGraph" + dind: false device: mi250_1 timeout_in_minutes: 55 depends_on: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b83d2893bdae..a10824656395 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -416,6 +416,7 @@ steps: mirror: amd: label: ":amd: (MI250) Pipeline + Context Parallelism" + dind: false device: mi250_4 timeout_in_minutes: 90 depends_on: diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index ef50062a7f1c..4c80f26ca52e 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -18,6 +18,7 @@ steps: mirror: amd: label: ":amd: (MI250) Docker Build Metadata" + dind: false device: mi250_1 no_gpu: true timeout_in_minutes: 35 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index f6e243a1b726..8e47a84e18cd 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -52,6 +52,7 @@ steps: mirror: amd: label: ":amd: (MI250) V1 Engine" + dind: false device: mi250_1 timeout_in_minutes: 45 depends_on: @@ -71,6 +72,7 @@ steps: mirror: amd: label: ":amd: (MI250) E2E Scheduling" + dind: false device: mi250_1 timeout_in_minutes: 55 depends_on: @@ -90,6 +92,7 @@ steps: mirror: amd: label: ":amd: (MI250) E2E Core" + dind: false commands: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py device: mi250_1 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 207d827f3600..7e43b123f12e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -322,6 +322,7 @@ steps: mirror: amd: label: ":amd: (MI250) Mamba Kernels" + dind: false device: mi250_1 timeout_in_minutes: 55 depends_on: @@ -512,6 +513,7 @@ steps: mirror: amd: label: ":amd: (MI250) Helion Kernels Shard %N" + dind: false device: mi250_1 timeout_in_minutes: 45 depends_on: diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index f17d154fbe98..1bff8e18500b 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -586,6 +586,7 @@ steps: mirror: amd: label: ":amd: (MI250) Batch Invariance" + dind: false device: mi250_1 timeout_in_minutes: 35 depends_on: diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 94b0723f8aab..6702eb6dd3fe 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -170,12 +170,14 @@ steps: - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py - tests/models/transformers/fusers/ + - tests/model_executor/layers/test_activation.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py mirror: amd: label: ":amd: (MI250) Basic Models Other" + dind: false device: mi250_1 no_gpu: true working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 3e5024fe8741..3f80dfdfa3c4 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -155,6 +155,7 @@ steps: mirror: amd: label: ":amd: (MI250) Language Models (PPL)" + dind: false device: mi250_1 timeout_in_minutes: 65 depends_on: @@ -214,6 +215,7 @@ steps: mirror: amd: label: ":amd: (MI250) Language Models (MTEB)" + dind: false device: mi250_1 timeout_in_minutes: 60 depends_on: diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index a324da687f03..779801a2ea2b 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -57,6 +57,7 @@ steps: mirror: amd: label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" + dind: false device: mi250_1 timeout_in_minutes: 55 depends_on: @@ -101,6 +102,7 @@ steps: mirror: amd: label: ":amd: (MI250) Multimodal Processor Shard %N" + dind: false device: mi250_1 no_gpu: true working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 11f63ff52dfd..6b901d8f490b 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -53,6 +53,7 @@ steps: mirror: amd: label: ":amd: (MI250) Plugin Integration" + dind: false device: mi250_2 timeout_in_minutes: 50 depends_on: diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 0b49787d5d76..53d1da7edf41 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -17,6 +17,7 @@ steps: mirror: amd: label: ":amd: (MI250) Rust Frontend Cargo Style + Clippy" + dind: false device: mi250_1 no_gpu: true working_dir: "/vllm-workspace" @@ -40,6 +41,7 @@ steps: mirror: amd: label: ":amd: (MI250) Rust Frontend Cargo" + dind: false device: mi250_1 no_gpu: true working_dir: "/vllm-workspace" diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 18f857f65a5f..e1bc4c08d8ad 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -20,6 +20,7 @@ steps: mirror: amd: label: ":amd: (MI250) Samplers" + dind: false device: mi250_1 timeout_in_minutes: 40 depends_on: diff --git a/tests/conftest.py b/tests/conftest.py index a52a38f35775..0d70be62eb07 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -241,7 +241,7 @@ def dist_init(): rank=0, distributed_init_method=f"file://{temp_file}", local_rank=0, - backend="nccl", + backend="gloo" if current_platform.is_cpu() else "nccl", ) initialize_model_parallel(1, 1) yield diff --git a/tests/model_executor/layers/test_activation.py b/tests/model_executor/layers/test_activation.py new file mode 100644 index 000000000000..d0ef42c99afe --- /dev/null +++ b/tests/model_executor/layers/test_activation.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from functools import partial + +import pytest +import torch +import torch.nn as nn + +from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul +from vllm.platforms import CpuArchEnum, current_platform + + +@pytest.mark.skipif( + not current_platform.is_cpu() + or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC, + reason="Native CPU activations require a non-PowerPC CPU platform", +) +@pytest.mark.parametrize( + "activation_factory,reference", + [ + (SiluAndMul, nn.SiLU()), + (GeluAndMul, nn.GELU()), + (partial(GeluAndMul, approximate="tanh"), nn.GELU(approximate="tanh")), + ], + ids=["silu", "gelu", "gelu_tanh"], +) +def test_cpu_act_and_mul_without_compiled_ops( + default_vllm_config, monkeypatch, activation_factory, reference +): + """CPU activations must construct and run without compiled wheel operators.""" + default_vllm_config.compilation_config.custom_ops = ["all"] + monkeypatch.setattr(torch.ops, "_C", object()) + x = torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0]], device="cpu") + + layer = activation_factory() + + torch.testing.assert_close(layer(x), reference(x[:, :4]) * x[:, 4:]) diff --git a/tests/tools/test_docker_build_metadata_args.py b/tests/tools/test_docker_build_metadata_args.py index 8095ad34aa0e..127768afc733 100644 --- a/tests/tools/test_docker_build_metadata_args.py +++ b/tests/tools/test_docker_build_metadata_args.py @@ -155,10 +155,15 @@ def test_vllm_openai_image_embeds_metadata_contract() -> None: def test_rust_build_cache_excludes_git_metadata() -> None: + import torch + from vllm.platforms import current_platform dockerfile_names = ["Dockerfile", "Dockerfile.cpu"] - if not current_platform.is_rocm(): + # CPU jobs can reuse ROCm artifacts, which omit the XPU Dockerfile. + if not current_platform.is_rocm() and ( + not current_platform.is_cpu() or torch.version.hip is None + ): dockerfile_names.append("Dockerfile.xpu") for name in dockerfile_names: dockerfile = (REPO_ROOT / "docker" / name).read_text() diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 7166d3f8feea..35a6fdf9b010 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -129,8 +129,11 @@ def __init__(self, *, compile_native: bool = True): super().__init__(compile_native=compile_native) if ( current_platform.is_cuda_alike() - or current_platform.is_cpu() or current_platform.is_xpu() + or ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC + ) ): self.op = torch.ops._C.silu_and_mul @@ -424,8 +427,11 @@ def __init__(self, approximate: str = "none"): raise ValueError(f"Unknown approximate mode: {approximate}") if ( current_platform.is_cuda_alike() - or current_platform.is_cpu() or current_platform.is_xpu() + or ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC + ) ): if approximate == "none": self.op = torch.ops._C.gelu_and_mul From 0390299309c5bc11b8aa78ea0f1d10993273dbb4 Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:16:29 -0500 Subject: [PATCH 493/774] [ROCm][CI] Shard MI300 Entrypoints Integration (Pooling) (#57583) Signed-off-by: aarushjain29 --- .buildkite/test-amd.yaml | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 789ca5426034..17b0fa408e8b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1341,21 +1341,6 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal -- label: ":amd: (MI300) Entrypoints Integration (Pooling)" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - fast_check: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/pooling - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/pooling - - label: ":amd: (MI300) OpenAI API Correctness" # TBD timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3615,12 +3600,13 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal -- label: ":amd: (MI355) Entrypoints Integration (Pooling)" # TBD +- label: ":amd: (MI355) Entrypoints Integration (Pooling) Shard %N" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 fast_check: true + parallelism: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -3628,7 +3614,7 @@ steps: optional: true commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/pooling + - pytest -v -s entrypoints/pooling --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB #----------------------------------------------------------- mi355 · evals -----------------------------------------------------------# From 62af3df7347cf890b4b925a18ab6a2afd00c70b9 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Fri, 18 Sep 2026 17:22:18 -0400 Subject: [PATCH 494/774] [Bugfix][MLA] Reserve sparse prefill buffers before KV cache sizing (#57575) Signed-off-by: Lucas Wilkinson Co-authored-by: Codex Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../v1/attention/test_sparse_mla_backends.py | 15 +++++++ vllm/models/hy_v4/nvidia/flashmla_sparse.py | 12 ++++- .../attention/backends/mla/flashmla_sparse.py | 44 ++++++++++++++----- 3 files changed, 59 insertions(+), 12 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 8ff52ecb386c..7704600d09d4 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -421,7 +421,18 @@ def test_sparse_backend_decode_correctness( workspace_init, q_scale: float, k_scale: float, + monkeypatch, ): + if ( + batch_name == "large_q_pure_prefill" + and backend_cls == FlashMLASparseBackend + and kv_cache_dtype == "fp8_ds_mla" + and tensor_parallel_size == 4 + ): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse.split_prefill_chunks", + lambda rows, capacity: [(i, i + 1) for i in range(len(rows))], + ) if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes: pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}") @@ -774,6 +785,10 @@ def test_sparse_backend_decode_correctness( ) with torch.inference_mode(): + if backend_cls == FlashMLASparseBackend and kv_cache_dtype == "fp8_ds_mla": + from vllm.v1.worker.workspace import current_workspace_manager + + current_workspace_manager().lock() backend_output = mock_layer.forward_impl( query_vllm, kv_c_vllm, diff --git a/vllm/models/hy_v4/nvidia/flashmla_sparse.py b/vllm/models/hy_v4/nvidia/flashmla_sparse.py index 83871755527b..1d62f2f556de 100644 --- a/vllm/models/hy_v4/nvidia/flashmla_sparse.py +++ b/vllm/models/hy_v4/nvidia/flashmla_sparse.py @@ -34,6 +34,7 @@ flash_mla_sparse_fwd, flash_mla_with_kvcache, ) +from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: from vllm.model_executor.models.deepseek_v2 import Indexer @@ -220,6 +221,7 @@ def _bf16_flash_mla_kernel( topk_indices: torch.Tensor, topk_length: torch.Tensor | None = None, actual_num_heads: int | None = None, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = q.shape[0] kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( @@ -250,7 +252,14 @@ def _bf16_flash_mla_kernel( # Zero (not new_empty) the padded lanes: topk_indices is shared by # all heads, so the kernel reduces across the head group and NaNs # from uninitialized memory would leak into the real heads. - q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2])) + if out is None: + q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2])) + else: + *_, q_padded, _, _ = current_workspace_manager().get_simultaneous( + *self.workspace_specs + ) + q_padded = q_padded[:num_tokens] + q_padded[:, actual_num_heads:].zero_() q_padded[:, :actual_num_heads, :] = q q = q_padded @@ -262,6 +271,7 @@ def _bf16_flash_mla_kernel( self.softmax_scale, attn_sink=attn_sink, topk_length=topk_length, + out=out, ) output = output[:, :actual_num_heads, :] diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index b3332eb86292..aa022066f045 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -19,6 +19,7 @@ ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability +from vllm.utils.math_utils import round_up from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import ( async_tensor_h2d, @@ -745,6 +746,20 @@ def __init__( self.workspace_specs.append( ((prefill_workspace_size, head_size), torch.bfloat16) ) + prefill_query_heads = num_heads + if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size: + prefill_query_heads *= parallel_config.tensor_parallel_size + padded_prefill_query_heads = round_up( + prefill_query_heads, self.prefill_padding + ) + self.workspace_specs.extend( + (shape, torch.bfloat16) + for shape in ( + (max_tokens, padded_prefill_query_heads, head_size), + (max_tokens, padded_prefill_query_heads, self.kv_lora_rank), + (max_tokens, prefill_query_heads, self.kv_lora_rank), + ) + ) # Reserve capacity without retaining views that prevent old storage # from being released when another layer grows the shared workspace. current_workspace_manager().get_simultaneous(*self.workspace_specs) @@ -845,7 +860,7 @@ def _gather_prefill_chunk( """All-gather this rank's upconverted KV shard so the chunk's rows attend the whole context, and map their top-k onto the rank-major result.""" shard_rows = int(chunk.chunk_tot_seqlen) - _, _, gathered_kv_workspace = current_workspace_manager().get_simultaneous( + _, _, gathered_kv_workspace, *_ = current_workspace_manager().get_simultaneous( *self.workspace_specs ) gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows] @@ -1007,12 +1022,10 @@ def _fp8_decode( q, decode_topk if decode_topk is not None else topk_indices ) else: - # Mixed or pure prefill: allocate output tensor - attn_out = q.new_empty( - (num_mqa_tokens, self.num_heads, self.kv_lora_rank), - dtype=q.dtype, - device=q.device, + _, prefill_bf16_workspace, *_, padded_out, compact_out = ( + current_workspace_manager().get_simultaneous(*self.workspace_specs) ) + attn_out = compact_out[:num_mqa_tokens] if num_decode_tokens > 0: attn_out[:num_decode_tokens] = _fp8_decode( @@ -1023,9 +1036,6 @@ def _fp8_decode( ) assert fp8_metadata.prefill is not None - _, prefill_bf16_workspace, *_ = ( - current_workspace_manager().get_simultaneous(*self.workspace_specs) - ) for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks): chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen] if uses_host_cache and chunk_index > 0: @@ -1076,12 +1086,15 @@ def _fp8_decode( assert topk_length is not None chunk_topk_length = topk_length[chunk.tokens_slice] - attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( + chunk_out, chunk_lse = self._bf16_flash_mla_kernel( chunk_q, chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length, + out=padded_out[chunk.tokens_slice], ) + attn_out[chunk.tokens_slice].copy_(chunk_out) + del chunk_out, chunk_lse if self.pcp_dcp_kv_gather and lse is None: # No decode rows: the DCP merge still expects an LSE, of no rows. @@ -1254,6 +1267,7 @@ def _bf16_flash_mla_kernel( topk_indices: torch.Tensor, topk_length: torch.Tensor | None = None, actual_num_heads: int | None = None, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = q.shape[0] kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( @@ -1275,7 +1289,14 @@ def _bf16_flash_mla_kernel( f"Padding num_heads from {actual_num_heads} to " f"{padded_num_heads} for BF16 sparse prefill kernel" ) - q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2])) + if out is None: + q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2])) + else: + *_, q_padded, _, _ = current_workspace_manager().get_simultaneous( + *self.workspace_specs + ) + q_padded = q_padded[:num_tokens] + q_padded[:, actual_num_heads:].zero_() q_padded[:, :actual_num_heads, :] = q q = q_padded @@ -1286,6 +1307,7 @@ def _bf16_flash_mla_kernel( topk_indices, self.softmax_scale, topk_length=topk_length, + out=out, ) output = output[:, :actual_num_heads, :] From 6aab78f4777c945b02123af330ff351f72785c0f Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Sat, 19 Sep 2026 05:53:02 +0800 Subject: [PATCH 495/774] [Bugfix][KVConnector] Make ExampleHiddenStatesConnector abort-safe (#56841) Signed-off-by: ys2025-AI --- .../unit/test_hidden_states_connector.py | 59 +++++++++++++++++++ .../v1/example_hidden_states_connector.py | 43 +++++++++++--- 2 files changed, 93 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py index 8427833b121a..735b042259b5 100644 --- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -169,3 +169,62 @@ def test_hidden_state_group_isolated_from_packed_mixed_page_groups(): assert ( ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == hidden_group_ids[0] ) + + +# ---- abort-path robustness -------------------------------------------------- + + +def _bare_connector() -> ExampleHiddenStatesConnector: + """Instance with scheduler-side state but bypassing ``__init__`` (no engine).""" + conn = ExampleHiddenStatesConnector.__new__(ExampleHiddenStatesConnector) + conn._request_filenames = {} + conn._pending_saves = {} + conn._lock_fds = {} + conn._cache_kv_group_id = 1 + conn._connector_metadata = None + conn._req_copy_events = {} + conn._accumulated_finished_req_ids = set() + return conn + + +def test_get_finished_count_is_one(): + # Only TP rank 0 writes, so KVOutputAggregator must expect a single + # finished_sending notification per request (not the TP world size). + assert _bare_connector().get_finished_count() == 1 + + +def test_request_finished_is_noop_for_never_scheduled_request(): + # A request aborted while still queued never reaches build_connector_meta, + # so no filename was recorded. request_finished must not raise KeyError. + conn = _bare_connector() + request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None) + assert conn.request_finished(request, []) == (False, None) + + +def test_request_finished_all_groups_handles_missing_group(): + # Guard against indexing a nonexistent per-group block table. + conn = _bare_connector() + conn._cache_kv_group_id = 2 + request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None) + assert conn.request_finished_all_groups(request, ([], [])) == (False, None) + + +def test_get_finished_does_not_report_untracked_request(): + # A never-scheduled aborted request has no copy event. get_finished must not + # report it as done_sending, or the scheduler asserts it is still tracked. + conn = _bare_connector() + assert conn.get_finished({"cmpl-aborted"}) == (None, None) + assert conn._accumulated_finished_req_ids == set() + + +def test_get_finished_reports_tracked_completed_request(): + conn = _bare_connector() + + class _DoneEvent: + def query(self) -> bool: + return True + + conn._req_copy_events["cmpl-done"] = _DoneEvent() + done_sending, done_recving = conn.get_finished({"cmpl-done"}) + assert done_sending == {"cmpl-done"} + assert done_recving is None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index e9a3e5d99a3c..47990d22a5a5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -103,6 +103,18 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): Must be used in conjunction with the `extract_hidden_states` spec decoding method. """ + def get_finished_count(self) -> int: + """Number of ``finished_sending`` notifications expected per request. + + Only TP rank 0 writes hidden states to disk, so each request yields a + single notification. The default (``None``) makes ``KVOutputAggregator`` + wait for every worker in the world size; since ranks 1..N-1 never submit + a copy, the request would never be reported as finished and its blocks + would be leaked forever (the scheduler then spins on + ``has_finished_requests()``). + """ + return 1 + @classmethod def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int: """Index of the KV cache group holding the extracted hidden states. @@ -517,7 +529,13 @@ def request_finished( the hidden states from the KV cache. """ req_id = request.request_id - filename = self._request_filenames.pop(req_id) + filename = self._request_filenames.pop(req_id, None) + if filename is None: + # Request was aborted before it was ever scheduled, so + # build_connector_meta never recorded a filename for it. There are + # no hidden states to save and no blocks worth delaying. + self._pending_saves.pop(req_id, None) + return False, None kv_params = request.kv_transfer_params or {} if kv_params.get("include_output_tokens", False): # Exclude the final token — it was the model's output, never an @@ -566,15 +584,18 @@ def get_finished( done_sending: set[str] = set() for req_id in list(self._accumulated_finished_req_ids): event = self._req_copy_events.get(req_id) - if event is None or event.query(): - self._req_copy_events.pop(req_id, None) + if event is not None and not event.query(): + continue # DtoH copy still in flight + self._req_copy_events.pop(req_id, None) + self._accumulated_finished_req_ids.discard(req_id) + # Only report requests we actually tracked for sending. + if event is not None: done_sending.add(req_id) - self._accumulated_finished_req_ids.discard(req_id) - # Clean up any leftover lock fds (e.g. aborted requests - # that never went through _submit_async_write). - lock_fd = self._lock_fds.pop(req_id, None) - if lock_fd is not None: - os.close(lock_fd) + # Clean up any leftover lock fds (e.g. aborted requests that never + # went through _submit_async_write). + lock_fd = self._lock_fds.pop(req_id, None) + if lock_fd is not None: + os.close(lock_fd) return done_sending or None, None @@ -583,6 +604,10 @@ def request_finished_all_groups( request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: + if not 0 <= self._cache_kv_group_id < len(block_ids): + # No per-group block table (e.g. a request that was never + # scheduled); treat as a no-op instead of indexing out of bounds. + return False, None return self.request_finished(request, block_ids[self._cache_kv_group_id]) @classmethod From 1dc2d854c12099f1f08233c886368a8c7baf430f Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Fri, 18 Sep 2026 15:01:24 -0700 Subject: [PATCH 496/774] [Feat][Model] Support encoder-side SWA-bounded replay for DeepSeek-V4.1-Flash (#56227) Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 Co-authored-by: Woosuk Kwon --- .../models/test_deepseek_v41_replay_start.py | 112 +++++++ .../attention/test_deepseek_v4_swa_visible.py | 184 +++++++++++- .../test_dspark_noncausal_sparse_mla.py | 1 + .../test_indexer_deepseek_v4_slot_mapping.py | 23 ++ tests/v1/core/test_contiguous_kv_packing.py | 64 ++++ tests/v1/core/test_prefix_caching.py | 4 +- tests/v1/core/test_prefix_replay.py | 279 ++++++++++++++++++ .../unit/test_nixl_connector_hma.py | 1 + vllm/config/cache.py | 8 +- vllm/engine/arg_utils.py | 5 + vllm/models/deepseek_v4/amd/rocm.py | 2 + vllm/models/deepseek_v4/sparse_mla.py | 1 + vllm/models/deepseek_v41/amd/rocm.py | 2 + vllm/models/deepseek_v41/attention.py | 20 ++ .../deepseek_v41/common/ops/cache_utils.py | 13 + .../deepseek_v41/nvidia/flashinfer_sparse.py | 7 +- .../models/deepseek_v41/nvidia/model_state.py | 145 ++++++++- vllm/models/deepseek_v41/sparse_mla.py | 1 + .../backends/mla/compressor_utils.py | 31 +- vllm/v1/attention/backends/mla/indexer.py | 11 +- vllm/v1/attention/backends/mla/sparse_swa.py | 55 +++- vllm/v1/core/kv_cache_manager.py | 13 +- vllm/v1/core/sched/output.py | 3 + vllm/v1/core/sched/scheduler.py | 79 ++++- vllm/v1/core/single_type_kv_cache_manager.py | 2 + vllm/v1/kv_cache_interface.py | 24 +- vllm/v1/request.py | 3 + vllm/v1/worker/gpu/model_states/default.py | 7 +- .../worker/gpu/model_states/encoder_only.py | 3 + .../worker/gpu/model_states/mamba_hybrid.py | 2 + 30 files changed, 1067 insertions(+), 38 deletions(-) create mode 100644 tests/models/test_deepseek_v41_replay_start.py create mode 100644 tests/v1/core/test_prefix_replay.py diff --git a/tests/models/test_deepseek_v41_replay_start.py b/tests/models/test_deepseek_v41_replay_start.py new file mode 100644 index 000000000000..7d4e45fe8170 --- /dev/null +++ b/tests/models/test_deepseek_v41_replay_start.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SWA bounded replay in DeepseekV41ModelState: the batch's replay starts reach +the sliding-window builders, and the replayed tokens' slots are padded in the +prefix-cacheable groups only.""" + +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.model_states.default import DefaultModelState + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +DEVICE = torch.device("cuda") +WINDOW = 4 +# Group 0 is prefix-cacheable (compressed KV), group 1 is the replayed window. +KV_CACHE_CONFIG = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + kv_cache_spec=SimpleNamespace(prefix_cacheable=True, prefix_replay_tokens=0) + ), + SimpleNamespace( + kv_cache_spec=SimpleNamespace( + prefix_cacheable=False, prefix_replay_tokens=WINDOW + ) + ), + ] +) + + +@pytest.fixture +def state(monkeypatch): + cfg = MagicMock() + cfg.model_config.enable_prompt_embeds = False + cfg.model_config.uses_mrope = False + cfg.model_config.is_multimodal_model = False + cfg.scheduler_config.max_num_seqs = 8 + cfg.scheduler_config.max_num_batched_tokens = 64 + seen: dict = {} + + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw): + seen["slot_mappings"] = slot_mappings.clone() + seen["extra"] = kw["model_specific_attn_metadata"] + return {} + + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) + state = DeepseekV41ModelState( + cfg, SimpleNamespace(token_lookback_depth=0), None, DEVICE + ) + state.seen = seen # type: ignore[attr-defined] + return state + + +def _batch(query_lens, seq_lens, idx_mapping, is_prefilling): + num_reqs, num_tokens = len(query_lens), sum(query_lens) + batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(8, 64, DEVICE)) + return replace( + batch, + query_start_loc=torch.tensor( + [0, *np.cumsum(query_lens)], dtype=torch.int32, device=DEVICE + ), + idx_mapping_np=np.array(idx_mapping, dtype=np.intp), + is_prefilling_np=np.array(is_prefilling, dtype=np.bool_), + positions=torch.cat( + [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)] + ).to(DEVICE), + ) + + +def _prepare(state, batch): + slot_mappings = torch.arange(2 * batch.num_tokens, device=DEVICE).view(2, -1) + state.prepare_attn( + batch, CUDAGraphMode.NONE, (), slot_mappings, [], KV_CACHE_CONFIG + ) + swa_builder = MagicMock(spec=DeepseekSparseSWAMetadataBuilder) + extra = state.seen["extra"] + replay_start = extra.get_extra_attn_kwargs(swa_builder, batch.num_reqs)[ + "replay_start" + ] + assert extra.get_extra_attn_kwargs(MagicMock(), batch.num_reqs) == {} + return replay_start.tolist(), state.seen["slot_mappings"] + + +def test_replayed_tokens_write_only_the_window_group(state): + # Request state 3 resumes a hit at 16: positions 16..19 are replayed. + state.add_request(3, SimpleNamespace(replay_start=16)) + state.add_request(1, SimpleNamespace(replay_start=0)) + batch = _batch([6, 2], [22, 2], idx_mapping=[3, 1], is_prefilling=[True, True]) + replay_start, slots = _prepare(state, batch) + assert replay_start == [16, 0] + padded = [PAD_SLOT_ID] * WINDOW + [4, 5, 6, 7] + assert slots[0].tolist() == padded # compressed KV keeps its cached rows + assert slots[1].tolist() == list(range(8, 16)) # the window is rebuilt + + +def test_only_prefills_carry_a_replay_start(state): + state.add_request(0, SimpleNamespace(replay_start=16)) + # The same request, decoding: its rows sit above the hit. + batch = _batch([1], [40], idx_mapping=[0], is_prefilling=[False]) + replay_start, slots = _prepare(state, batch) + assert replay_start == [0] + assert slots.tolist() == [[0], [1]] diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 1b5097de8e2d..7a74ce8c81be 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -11,6 +11,8 @@ ``pos + right`` (inclusive). """ +import inspect + import pytest import torch from typing_extensions import TypedDict @@ -57,10 +59,13 @@ def ref_left_right( return lefts, rights -def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]: - """Reference [start, end) window bounds for one query token.""" +def ref_swa_bounds( + pos: int, window: int, left: int, right: int, replay_start: int = 0 +) -> tuple[int, int]: + """Reference [start, end) window bounds for one query token. Under SWA + bounded replay no window KV exists below ``replay_start``.""" left_add = max(left - (window - 1), 0) - start = max(pos - (window - 1) - left_add, 0) + start = max(pos - (window - 1) - left_add, 0, replay_start) return start, pos + right + 1 @@ -97,12 +102,14 @@ def ref_swa_slot_rows( window: int, max_image_tokens: int, width: int, + replay_starts: list[int] | None = None, ) -> tuple[list[list[int]], list[int]]: """Reference paged slot-id rows and lens for every token in the batch.""" block_table_cpu = block_table.cpu() lefts, rights = ref_left_right( seq_lens, query_lens, spans_per_req, max_image_tokens ) + replay_starts = replay_starts or [0] * len(seq_lens) rows: list[list[int]] = [] lens: list[int] = [] token = 0 @@ -110,7 +117,9 @@ def ref_swa_slot_rows( prefix_len = seq_len - query_len for i in range(query_len): pos = prefix_len + i - start, end = ref_swa_bounds(pos, window, lefts[token], rights[token]) + start, end = ref_swa_bounds( + pos, window, lefts[token], rights[token], replay_starts[req] + ) row = [] for p in range(start, end): blk = int(block_table_cpu[req, p // BLOCK_SIZE]) @@ -139,6 +148,7 @@ def run_swa_kernel( swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device) swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device) is_valid = slot_mapping >= 0 + replay_start_t = torch.zeros(len(seq_lens), dtype=torch.int32, device=device) if with_image: lefts, rights = ref_left_right( @@ -164,6 +174,7 @@ def run_swa_kernel( block_table, block_table.stride(0), BLOCK_SIZE, + replay_start_t, token_offset=0, HAS_IMAGE=with_image, TRITON_BLOCK_SIZE=1024, @@ -289,18 +300,25 @@ def combine_case( query_lens: list[int], spans: list[list[tuple[int, int]]], with_image: bool, + replay_starts: list[int] | None = None, + combine_fn=combine_topk_swa_indices, ): """Run combine_topk_swa_indices and return (indices, lens, expected).""" device = torch.device("cuda") num_reqs = len(seq_lens) + replay_starts = replay_starts or [0] * num_reqs query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) query_start_loc[1:] = torch.tensor( query_lens, dtype=torch.int32, device=device ).cumsum(0) num_tokens = int(query_start_loc[-1]) seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + # The builder's gather covers only the context above the replay start. gather_lens = torch.tensor( - [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)], + [ + q + min(max(s - q - r, 0), WINDOW - 1) + for s, q, r in zip(seq_lens, query_lens, replay_starts) + ], dtype=torch.int32, device=device, ) @@ -319,7 +337,14 @@ def combine_case( else: left_t = right_t = None - combined_indices, combined_lens = combine_topk_swa_indices( + # V4.1's combine has no image windows (causal SWA only). + has_image = "left_visible" in inspect.signature(combine_fn).parameters + image_kwargs = ( + dict(left_visible=left_t, right_visible=right_t, max_image_tokens=MAX_IMG) + if has_image + else {} + ) + combined_indices, combined_lens = combine_fn( topk_indices, query_start_loc, seq_lens_t, @@ -329,9 +354,7 @@ def combine_case( topk, M, N, - left_visible=left_t, - right_visible=right_t, - max_image_tokens=MAX_IMG, + **image_kwargs, ) # Reference rows. @@ -342,7 +365,7 @@ def combine_case( MAX_IMG, ) topk_cpu = topk_indices.cpu() - width = WINDOW + MAX_IMG + width = WINDOW + (MAX_IMG if has_image else 0) combined_topk = (topk + width + 127) // 128 * 128 rows = [] lens = [] @@ -354,6 +377,8 @@ def combine_case( pos = prefix_len + i topk_len = min((pos + 1) // compress_ratio, topk) start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token]) + # The window never reaches below the gathered buffer. + start = max(start, gather_start) swa_len = end - start row = [-1] * combined_topk for j in range(topk_len): @@ -388,6 +413,31 @@ def test_combine_topk_swa_with_image_spans(cfg): assert indices.cpu().tolist() == rows +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("cfg", COMBINE_CASES) +def test_v41_combine_topk_swa_stops_at_replay_start(cfg): + """SWA bounded replay: the gathered buffer starts at replay_start, so the + window never indexes below it.""" + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + combine_topk_swa_indices as combine_v41, + ) + + # Request 0 replays [16, 40): the windows of its first rows would + # otherwise reach below 16. + indices, lens, rows, exp_lens = combine_case( + cfg["compress_ratio"], + cfg["topk"], + seq_lens=[40, 12], + query_lens=[24, 12], + spans=[[], []], + with_image=False, + replay_starts=[16, 0], + combine_fn=combine_v41, + ) + assert lens.cpu().tolist() == exp_lens + assert indices.cpu().tolist() == rows + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("cfg", COMBINE_CASES) def test_combine_topk_swa_without_image_unchanged(cfg): @@ -472,6 +522,7 @@ def make_builder( DeepseekSparseSWAMetadataBuilder ] = DeepseekSparseSWAMetadataBuilder, mm_prefix_clamp_sliding_window: bool = True, + bounded_replay: bool = False, ) -> DeepseekSparseSWAMetadataBuilder: overrides: dict = {"sliding_window": window} if vision: @@ -498,6 +549,7 @@ def make_builder( sliding_window=window, cache_dtype_str="auto", model_version="deepseek_v4", + bounded_replay=bounded_replay, ) return builder_cls( kv_cache_spec=spec, @@ -512,11 +564,17 @@ def build_metadata( seq_lens: list[int], query_lens: list[int], mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None, + replay_starts: list[int] | None = None, ): device = torch.device("cuda") query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch( seq_lens, query_lens, device ) + replay_start = ( + None + if replay_starts is None + else torch.tensor(replay_starts, dtype=torch.int32, device=device) + ) return builder.build( 0, CommonAttentionMetadata( @@ -533,6 +591,7 @@ def build_metadata( causal=True, mm_req_doc_ranges=mm_req_doc_ranges, ), + replay_start=replay_start, ) @@ -653,3 +712,108 @@ def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len): assert absolute_keys.tolist() == list(range(lo, hi)) assert paged_lens[token] == hi - lo assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi)) + + +# SWA bounded replay: every prefill index path clamps the window at replay_start. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_builder_replay_start_bounds_prefill_window_and_gather(): + """Paged-direct prefill indices stop at replay_start and the FlashMLA + gather only covers the context the request may see.""" + seq_lens = [40, 12, 30] + query_lens = [24, 12, 10] + replay_starts = [12, 0, 0] + builder = make_builder(vision=False) + md = build_metadata(builder, seq_lens, query_lens, None, replay_starts) + assert md.replay_start is not None + assert md.replay_start.cpu().tolist() == replay_starts + + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) + rows, lens = ref_swa_slot_rows( + seq_lens, + query_lens, + [[], [], []], + block_table, + WINDOW, + MAX_IMG, + WINDOW, + replay_starts=replay_starts, + ) + assert md.prefill_swa_lens.cpu().tolist() == lens + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows + # Request 0: the first replayed row (pos 16) sees [12, 16]; its last row + # (pos 39) is a full window above the replay start and sees all of it. + assert md.prefill_swa_lens[0].item() == 16 - 12 + 1 + assert md.prefill_swa_lens[23].item() == WINDOW + # gather_len = query_len + min(prefix_len - replay_start, WINDOW - 1). + assert md.prefill_gather_lens.cpu().tolist() == [ + 24 + min(16 - 12, WINDOW - 1), + 12, + 10 + min(20, WINDOW - 1), + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_builder_replaying_group_builds_without_replay_start(): + """Graph captures build without model-specific kwargs: a replaying group's + metadata then bounds nothing.""" + builder = make_builder(vision=False, bounded_replay=True) + md = build_metadata(builder, [40, 12], [24, 12], None) + assert not md.replay_start.any() + assert md.prefill_swa_lens[0].item() == WINDOW + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start(): + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + build_flashinfer_mixed_sparse_indices as build_v41, + ) + + device = torch.device("cuda") + # 1 decode token (req 0) + two prefill requests; req 1 resumes at 16. + seq_lens = [20, 40, 9] + query_lens = [1, 24, 9] + replay_starts = [0, 16, 0] + query_start_loc = torch.tensor([0, 1, 25, 34], dtype=torch.int32, device=device) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + token_to_req = torch.tensor( + [0] + [1] * 24 + [2] * 9, dtype=torch.int32, device=device + ) + block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1) + kwargs = dict( + decode_swa_indices=torch.zeros((1, WINDOW), dtype=torch.int32, device=device), + decode_compressed_indices=None, + decode_compressed_topk_lens=None, + prefill_topk_indices=torch.empty((33, 0), dtype=torch.int32, device=device), + query_start_loc=query_start_loc, + seq_lens=seq_lens_t, + token_to_req_indices=token_to_req, + swa_block_table=block_table, + swa_block_size=BLOCK_SIZE, + compressed_block_table=None, + compressed_block_size=BLOCK_SIZE, + window_size=WINDOW, + compress_ratio=1, + topk=0, + ) + plain, _ = build_v41( + replay_start=torch.zeros(3, dtype=torch.int32, device=device), **kwargs + ) + bounded, _ = build_v41( + replay_start=torch.tensor(replay_starts, dtype=torch.int32, device=device), + **kwargs, + ) + rows, _ = ref_swa_slot_rows( + seq_lens, + query_lens, + [[], [], []], + block_table, + WINDOW, + MAX_IMG, + WINDOW, + replay_starts=replay_starts, + ) + # Decode row is copied through; prefill rows follow the bounded window. + assert bounded[0].cpu().tolist() == plain[0].cpu().tolist() + assert bounded[1:].cpu().tolist() == rows[1:] diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 30c35a730efc..3d4af1741362 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -669,6 +669,7 @@ def test_dsv41_flashinfer_dspark_window_matches_reference( ), block_table=torch.arange(48, device=device, dtype=torch.int32).view(3, -1), block_size=128, + replay_start=torch.zeros(3, device=device, dtype=torch.int32), flashinfer_sparse_index_cache={}, max_decode_query_len=5, ) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index f3b4b7f5ba24..83fceef0181e 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -14,6 +14,7 @@ from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla.compressor_utils import ( CompressedSlotMappingKernel, + get_compressed_slot_mapping, ) from vllm.v1.attention.backends.mla.indexer import ( BuildPrefillChunkMetadataKernel, @@ -235,6 +236,28 @@ def test_compressed_slot_mapping_warmup_includes_index_kpool(): assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)} +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compressed_slot_mapping_inherits_padded_token_slots(): + """A token whose own slot is padded (SWA bounded replay) closes no + compressed state either.""" + device = torch.device("cuda") + query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device) + seq_lens = torch.tensor([8], dtype=torch.int32, device=device) + block_table = torch.tensor([[3]], dtype=torch.int32, device=device) + slot_mapping = torch.arange(8, dtype=torch.int64, device=device) + slot_mapping[:4] = -1 + compressed = get_compressed_slot_mapping( + 8, + slot_mapping, + query_start_loc, + seq_lens, + block_table, + block_size=4, + compress_ratio=2, + ) + assert compressed.tolist() == [-1, -1, -1, -1, -1, 3 * 4 + 2, -1, 3 * 4 + 3] + + def test_index_conversion_warmup_uses_physical_block_stride(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=64), diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 2fcf70c6efa7..58d4e20bf70c 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -651,3 +651,67 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): scheduler_block_size=128, ) assert len(manager.coordinator.single_type_managers) == len(groups) + + +class TestSWABoundedReplayGrouping: + def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self): + """Under SWA bounded replay only the paged MLA group hashes (hash + block = its block) and the coordinator probes one spec group.""" + config = _shared_layout_config() + config.cache_config.enable_prefix_caching = True + specs: dict[str, KVCacheSpec] = { + "layers.2.attn": MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=2, + alignment=576, + ), + } + swa = SlidingWindowMLASpec( + block_size=32, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=128, + alignment=576, + bounded_replay=True, + ) + for layer in range(4): + specs[f"layers.{layer}.attn.swa_cache"] = swa + + groups = get_kv_cache_groups(config, specs) + # The worker reads the replay window off the (packed) group specs to + # arm the window clamp; the wrapper must forward it like cacheability. + assert sorted( + (g.kv_cache_spec.prefix_cacheable, g.kv_cache_spec.prefix_replay_tokens) + for g in groups + ) == [(False, 128)] * 4 + [(True, 0)] + kv_cache_config = get_kv_cache_config_from_groups( + config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + ) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([kv_cache_config]), + max_model_len=8192, + enable_caching=True, + hash_block_size=128, + scheduler_block_size=128, + ) + assert len(manager.coordinator.attention_groups) == 1 + assert isinstance( + manager.coordinator.attention_groups[0].spec, MLAAttentionSpec + ) + + # Allocating and caching a request goes through every group's manager; + # the 32-token SWA group must skip hashing rather than assert on the + # 128-token hash block. + from tests.v1.core.utils import create_requests + + request = create_requests(num_requests=1, num_tokens=256, block_size=128)[0] + assert manager.allocate_slots(request, 256) is not None + request.num_computed_tokens = 256 + manager.cache_blocks(request, 256) + for single in manager.coordinator.single_type_managers: + cached = single.num_cached_block.get(request.request_id, 0) + assert cached == (2 if single.kv_cache_spec.prefix_cacheable else 0) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 229f9c17f9ee..40e810c58094 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -6,7 +6,7 @@ from collections.abc import Callable from dataclasses import replace from math import lcm -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import MagicMock import pytest @@ -533,7 +533,9 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens): kv_cache_manager=manager, failed_recving_kv_req_ids=set(), finished_recving_kv_req_ids={request.request_id}, + prefix_replay_tokens=0, ) + scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler) Scheduler._update_waiting_for_remote_kv(scheduler, request) assert request.num_tokens - request.num_computed_tokens == 1 assert manager.allocate_slots(request, num_new_tokens=1) is not None diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py new file mode 100644 index 000000000000..b2e78e1f925e --- /dev/null +++ b/tests/v1/core/test_prefix_replay.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SWA bounded replay: after a prefix hit the scheduler recomputes the tail of +the hit to rebuild the non-cacheable sliding-window group, keeps the hit's +blocks, and hands the worker the replay start.""" + +import pytest +import torch + +from tests.v1.kv_connector.unit.utils import MockKVConfig, create_model_runner_output +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + SlidingWindowMLASpec, +) +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import RequestStatus +from vllm.v1.structured_output import StructuredOutputManager + +from .utils import create_requests, create_scheduler + +BLOCK_SIZE = 16 +WINDOW = 32 +NUM_PROMPT_TOKENS = 100 +# 100 tokens -> 6 full blocks cached -> 96-token hit. +HIT_TOKENS = NUM_PROMPT_TOKENS // BLOCK_SIZE * BLOCK_SIZE +FULL, SWA = 0, 1 + + +def _replay_scheduler( + *, + long_prefill_token_threshold: int = 0, + use_kv_connector: MockKVConfig | None = None, +) -> Scheduler: + """A hybrid layout: one prefix-cacheable full-attention group and one + replayed sliding-window group.""" + base = create_scheduler( + block_size=BLOCK_SIZE, + enable_prefix_caching=True, + long_prefill_token_threshold=long_prefill_token_threshold, + use_kv_connector=use_kv_connector, + ) + vllm_config = base.vllm_config + kv_cache_config = KVCacheConfig( + num_blocks=10000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["swa"], + SlidingWindowMLASpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=64, + dtype=torch.bfloat16, + sliding_window=WINDOW, + bounded_replay=True, + ), + ), + ], + ) + scheduler = Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + block_size=BLOCK_SIZE, + log_stats=True, + structured_output_manager=StructuredOutputManager(vllm_config), + ) + scheduler.use_v2_model_runner = True + assert scheduler.prefix_replay_tokens == WINDOW + return scheduler + + +def _prefill(scheduler: Scheduler, request) -> None: + scheduler.add_request(request) + out = scheduler.schedule() + assert out.num_scheduled_tokens[request.request_id] == request.num_prompt_tokens + scheduler.update_from_output(out, _step_output(out, [request])) + + +def _step_output(out, requests) -> ModelRunnerOutput: + """Model output for every scheduled request: a token once its prefill is + done (non-zero, so outputs never extend the all-zero shared prompt in the + cache), nothing while it is still prefilling. schedule() has already + advanced num_computed_tokens past this step's chunk.""" + scheduled = [r for r in requests if r.request_id in out.num_scheduled_tokens] + return ModelRunnerOutput( + req_ids=[r.request_id for r in scheduled], + req_id_to_index={r.request_id: i for i, r in enumerate(scheduled)}, + sampled_token_ids=[ + [1000] if r.num_computed_tokens >= r.num_prompt_tokens else [] + for r in scheduled + ], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=None, + ) + + +def _new_req_data(out, request): + return next(r for r in out.scheduled_new_reqs if r.req_id == request.request_id) + + +def test_hit_replays_window_without_reallocating(): + scheduler = _replay_scheduler() + expected_replay = WINDOW + first, second = create_requests( + num_requests=2, + num_tokens=NUM_PROMPT_TOKENS, + same_prompt=True, + block_size=BLOCK_SIZE, + ) + _prefill(scheduler, first) + manager = scheduler.kv_cache_manager + + scheduler.add_request(second) + out = scheduler.schedule() + + replay_start = HIT_TOKENS - expected_replay + new_req = _new_req_data(out, second) + assert new_req.num_computed_tokens == replay_start + assert new_req.replay_start == replay_start + assert ( + out.num_scheduled_tokens[second.request_id] == NUM_PROMPT_TOKENS - replay_start + ) + + first_blocks = manager.get_blocks(first.request_id).blocks + second_blocks = manager.get_blocks(second.request_id).blocks + # The full-attention hit is shared (replayed tokens included); only the + # tail block is new. + num_hit_blocks = HIT_TOKENS // BLOCK_SIZE + assert second_blocks[FULL][:num_hit_blocks] == first_blocks[FULL][:num_hit_blocks] + assert all(b.ref_cnt == 2 for b in second_blocks[FULL][:num_hit_blocks]) + assert len(second_blocks[FULL]) == num_hit_blocks + 1 + # The window group never hits: blocks below the window are null, the + # window (replayed and new tokens) gets fresh blocks. + num_skipped_blocks = (HIT_TOKENS - WINDOW + 1) // BLOCK_SIZE + assert all(b.is_null for b in second_blocks[SWA][:num_skipped_blocks]) + num_window_blocks = len(second_blocks[SWA]) - num_skipped_blocks + assert num_window_blocks == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) - num_skipped_blocks + assert not any(b.is_null for b in second_blocks[SWA][num_skipped_blocks:]) + + scheduler.update_from_output(out, _step_output(out, [first, second])) + assert second.num_computed_tokens == NUM_PROMPT_TOKENS + + +def test_replay_only_chunks_make_progress(): + """A chunk cap smaller than the replay window schedules replay-only + chunks; they advance the request instead of stalling it.""" + scheduler = _replay_scheduler(long_prefill_token_threshold=BLOCK_SIZE) + first, second = create_requests( + num_requests=2, + num_tokens=NUM_PROMPT_TOKENS, + same_prompt=True, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(first) + while first.num_computed_tokens < NUM_PROMPT_TOKENS: + out = scheduler.schedule() + scheduler.update_from_output(out, _step_output(out, [first])) + + scheduler.add_request(second) + + def full_blocks(): + return scheduler.kv_cache_manager.get_blocks(second.request_id).blocks[FULL] + + scheduled: list[int] = [] + while second.num_computed_tokens < NUM_PROMPT_TOKENS: + out = scheduler.schedule() + if not scheduled: + new_req = _new_req_data(out, second) + assert new_req.num_computed_tokens == HIT_TOKENS - WINDOW + assert new_req.replay_start == HIT_TOKENS - WINDOW + # A replay-only chunk adopts the hit and allocates nothing new. + assert len(full_blocks()) == HIT_TOKENS // BLOCK_SIZE + scheduled.append(out.num_scheduled_tokens[second.request_id]) + scheduler.update_from_output(out, _step_output(out, [first, second])) + # 32 replayed + 4 new tokens in 16-token chunks: two replay-only chunks. + assert scheduled == [BLOCK_SIZE, BLOCK_SIZE, NUM_PROMPT_TOKENS - HIT_TOKENS] + assert len(full_blocks()) == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) + assert second.status == RequestStatus.RUNNING + + +def test_async_remote_kv_hit_replays_after_load(): + """A KV-connector hit loaded asynchronously carries no window state + either; the replay is applied when the load completes.""" + matched = 64 + scheduler = _replay_scheduler( + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True) + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + out = scheduler.schedule() + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert request.request_id not in out.num_scheduled_tokens + scheduler.update_from_output( + out, create_model_runner_output([], finished_recving={request.request_id}) + ) + + out = scheduler.schedule() + new_req = _new_req_data(out, request) + assert new_req.num_computed_tokens == matched - WINDOW + assert new_req.replay_start == matched - WINDOW + assert ( + out.num_scheduled_tokens[request.request_id] + == NUM_PROMPT_TOKENS - matched + WINDOW + ) + + +def test_remote_kv_hit_is_taken_in_whole_blocks(): + """A hit ending one token short of a block boundary would put the replay + window's first token in a block the sliding-window group retires, so a + connector hit is cut back to whole blocks.""" + matched = 4 * BLOCK_SIZE - 1 + scheduler = _replay_scheduler( + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True) + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + out = scheduler.schedule() + scheduler.update_from_output( + out, create_model_runner_output([], finished_recving={request.request_id}) + ) + out = scheduler.schedule() + new_req = _new_req_data(out, request) + hit = 3 * BLOCK_SIZE + assert new_req.replay_start == hit - WINDOW + assert new_req.num_computed_tokens == hit - WINDOW + swa_manager = scheduler.kv_cache_manager.coordinator.single_type_managers[SWA] + swa_blocks = scheduler.kv_cache_manager.get_blocks(request.request_id).blocks[SWA] + assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block + + +@pytest.mark.parametrize("via_connector", [False, True]) +def test_hit_no_longer_than_window_is_ignored(via_connector): + """Such a hit would be recomputed in full anyway. It is not adopted (an + async load of it is not even started), so a request has computed tokens + iff it replays.""" + if via_connector: + scheduler = _replay_scheduler( + use_kv_connector=MockKVConfig(matched_tokens=WINDOW, is_async=True) + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + else: + scheduler = _replay_scheduler() + first = create_requests( + num_requests=1, num_tokens=WINDOW, same_prompt=True, block_size=BLOCK_SIZE + )[0] + _prefill(scheduler, first) + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + same_prompt=True, + block_size=BLOCK_SIZE, + req_ids=["long"], + )[0] + scheduler.add_request(request) + out = scheduler.schedule() + new_req = _new_req_data(out, request) + assert request.status == RequestStatus.RUNNING + assert new_req.num_computed_tokens == 0 + assert new_req.replay_start == 0 + assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index f9cacabdd455..a4dfc168f31f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1505,6 +1505,7 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks(): scheduler.kv_cache_manager.cache_blocks = MagicMock() scheduler.failed_recving_kv_req_ids = {"req-1"} scheduler.finished_recving_kv_req_ids = {"req-1"} + scheduler.prefix_replay_tokens = 0 request = MagicMock() request.request_id = "req-1" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 37ad5b2b34fe..c12433a367a8 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -227,6 +227,11 @@ class CacheConfig: necessary for implementing this optimization in some models (e.g. Gemma3n) """ + swa_bounded_replay: bool = True + """Keep the sliding-window KV of models that support it (DeepSeek-V4.1) + out of prefix caching and rebuild it after a prefix hit by recomputing the + hit's last window. Requires model runner V2.""" + kv_cache_memory_bytes: int | None = None """Size of KV Cache per GPU in bytes. By default, this is set to None and vllm can automatically infer the kv cache size based on @@ -281,8 +286,9 @@ def compute_hash(self) -> str: "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", - # WIP feature toggle not impacting compiled graph shape + # Feature toggles not impacting compiled graph shape "kv_sharing_fast_prefill", + "swa_bounded_replay", } from vllm.config.utils import get_hash_factors, hash_factors diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 2dd4b04f2b56..67d8a2f2a0ba 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -781,6 +781,7 @@ class EngineArgs: stream_interval: int = SchedulerConfig.stream_interval kv_sharing_fast_prefill: bool = CacheConfig.kv_sharing_fast_prefill + swa_bounded_replay: bool = CacheConfig.swa_bounded_replay optimization_level: OptimizationLevel = VllmConfig.optimization_level performance_mode: PerformanceMode = VllmConfig.performance_mode @@ -1310,6 +1311,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--kv-sharing-fast-prefill", **cache_kwargs["kv_sharing_fast_prefill"] ) + cache_group.add_argument( + "--swa-bounded-replay", **cache_kwargs["swa_bounded_replay"] + ) cache_group.add_argument( "--mamba-cache-dtype", **cache_kwargs["mamba_cache_dtype"] ) @@ -2118,6 +2122,7 @@ def create_engine_config( prefix_cache_retention_interval=self.prefix_cache_retention_interval, kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, + swa_bounded_replay=self.swa_bounded_replay, mamba_cache_dtype=self.mamba_cache_dtype, mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 545f29bd356b..3cf26f8767ce 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -495,11 +495,13 @@ def build( common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, + replay_start: torch.Tensor | None = None, ) -> DeepseekV4ROCMAiterSparseSWAMetadata: base = super().build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, fast_build=fast_build, + replay_start=replay_start, ) ragged_indices = None diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 4953bdc9256f..a4000183a0e3 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -195,6 +195,7 @@ def build( if self.compress_ratio > 1: slot_mapping = get_compressed_slot_mapping( cm.num_actual_tokens, + cm.slot_mapping, cm.query_start_loc, cm.seq_lens, cm.block_table_tensor.clamp_(min=0), diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index b3d80f8d2591..4d838e121e69 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -425,11 +425,13 @@ def build( common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, + replay_start: torch.Tensor | None = None, ) -> DeepseekV4ROCMAiterSparseSWAMetadata: base = super().build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, fast_build=fast_build, + replay_start=replay_start, ) ragged_indices = None diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 5b26a98c1c53..bf645a8f156b 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -512,6 +512,25 @@ def __init__( "decodes only on SM100." ) + swa_bounded_replay = cache_config.swa_bounded_replay + if swa_bounded_replay and not vllm_config.use_v2_model_runner: + logger.warning_once( + "SWA bounded replay needs model runner V2 (only it skips the " + "paged-KV writes of replayed tokens); the sliding-window cache " + "takes part in prefix caching instead." + ) + swa_bounded_replay = False + if ( + swa_bounded_replay + and vllm_config.parallel_config.prefill_context_parallel_size > 1 + ): + logger.warning_once( + "SWA bounded replay is off under prefill context parallelism " + "(the replayed tokens' slot padding knows the rank-local batch " + "only); the sliding-window cache takes part in prefix caching " + "instead." + ) + swa_bounded_replay = False self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, window_size=self.window_size, @@ -522,6 +541,7 @@ def __init__( block_size=32, packed_bytes_per_token=self.swa_bytes_per_token, packed_page_alignment=self.kv_page_alignment, + bounded_replay=swa_bounded_replay, ) # The attention layer itself was already registered with the diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 04b2a9bac1aa..02fb0d55396c 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -982,6 +982,10 @@ def kernel( # turns the store index below negative). topk_len = 0 swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + # gather_len already excludes context below the request's replay + # start (SWA bounded replay), so the window cannot start before + # the gathered buffer does. + swa_start = tl.maximum(swa_start, gather_start) swa_len = pos - swa_start + 1 offset = tl.arange(0, PADDED_TOP_K) @@ -1128,6 +1132,8 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token: torch.Tensor | None = None, swa_block_span: int | None = None, compressed_block_span: int | None = None, + *, + replay_start: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -1137,6 +1143,9 @@ def build_flashinfer_mixed_sparse_indices( per token). Decode tokens read precomputed SWA/compressed indices; prefill tokens derive their SWA window from the position and translate local compressed indices to global slots via the block tables. + + ``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every + prefill token's window: positions below it hold no window KV. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 @@ -1227,6 +1236,7 @@ def build_flashinfer_mixed_sparse_indices( query_start_loc, seq_lens, token_to_req_indices, + replay_start, swa_block_table, swa_block_table.stride(0), swa_block_size, @@ -1295,6 +1305,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, + replay_start_ptr, swa_block_table_ptr, swa_block_table_stride, swa_block_size, @@ -1394,6 +1405,8 @@ def _build_flashinfer_mixed_sparse_indices_kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + # SWA bounded replay: no window KV exists below the request's replay start. + swa_start_pos = tl.maximum(swa_start_pos, tl.load(replay_start_ptr + req_idx)) swa_len = pos - swa_start_pos + 1 if COMPRESS_RATIO > 0: topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index a247ae7d7409..4dcc972e2c3f 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -200,8 +200,11 @@ def build( common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, + replay_start: torch.Tensor | None = None, ) -> "DeepseekSparseSWAMetadata": - metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = super().build( + common_prefix_len, common_attn_metadata, fast_build, replay_start + ) num_tokens = metadata.num_decode_tokens if not common_attn_metadata.causal and num_tokens > 0: assert metadata.decode_swa_lens is not None @@ -375,6 +378,7 @@ def _build_sparse_index_metadata( assert swa_metadata.token_to_req_indices is not None assert swa_metadata.decode_swa_indices is not None assert swa_metadata.block_table is not None + assert swa_metadata.replay_start is not None decode_swa_indices = swa_metadata.decode_swa_indices.reshape( num_decode_tokens, swa_metadata.decode_swa_width @@ -454,6 +458,7 @@ def _build_sparse_index_metadata( decode_is_valid_token=decode_is_valid_token, swa_block_span=swa_block_span, compressed_block_span=compressed_block_span, + replay_start=swa_metadata.replay_start[:num_reqs], ) if swa_only: swa_metadata.flashinfer_sparse_index_cache["swa_only"] = ( diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 1ba0f4160333..af20976dedbf 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -2,15 +2,23 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +import numpy as np import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.utils import AttentionGroup @triton.jit @@ -41,14 +49,69 @@ def _gather_lookback_kernel( tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH) +@triton.jit +def _pad_replayed_slots_kernel( + slot_mappings_ptr, # [num_groups, num_tokens_padded] + group_stride, + cacheable_groups_ptr, # [NUM_GROUPS] indices of the prefix-cacheable groups + query_start_loc_ptr, # [num_reqs + 1] + positions_ptr, # [num_tokens] + replay_start_ptr, # [num_reqs] + window, + pad_slot_id, + NUM_GROUPS: tl.constexpr, + BLOCK: tl.constexpr, +): + req = tl.program_id(0) + start = tl.load(replay_start_ptr + req) + if start == 0: + return + begin = tl.load(query_start_loc_ptr + req) + end = tl.load(query_start_loc_ptr + req + 1) + for tok in range(begin, end, BLOCK): + offs = tok + tl.arange(0, BLOCK) + pos = tl.load(positions_ptr + offs, mask=offs < end, other=0) + replayed = (offs < end) & (pos >= start) & (pos < start + window) + for i in tl.static_range(NUM_GROUPS): + group = tl.load(cacheable_groups_ptr + i) + tl.store( + slot_mappings_ptr + group * group_stride + offs, + pad_slot_id, + mask=replayed, + ) + + +class ReplayAttnMetadata(ModelSpecificAttnMetadata): + """Hands the batch's replay starts to the sliding-window builders.""" + + def __init__(self, replay_start: torch.Tensor) -> None: + self.replay_start = replay_start + + def get_extra_attn_kwargs( + self, attn_metadata_builder: Any, num_reqs: int + ) -> dict[str, Any]: + if isinstance(attn_metadata_builder, DeepseekSparseSWAMetadataBuilder): + return {"replay_start": self.replay_start[:num_reqs]} + return {} + + class DeepseekV41ModelState(DefaultModelState): - """DefaultModelState plus the engram lookback window. + """DefaultModelState plus the engram lookback window and SWA bounded replay. The engram n-gram hash needs the ids of the ``depth`` tokens preceding each request's chunk start (see ``common/engram.py``). The runner keeps the full token history on device, so the window is gathered there every step: exact for prompt and generated tokens alike, whatever instance produced their KV. + + SWA bounded replay (``CacheConfig.swa_bounded_replay``) keeps the + sliding-window KV out of prefix caching and rebuilds it after a prefix hit + by recomputing the hit's last window; the scheduler tells each request the + position from which it holds window KV (``NewRequestData.replay_start``). + ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots + in the prefix-cacheable groups so the cached KV stays as is, and hands the + starts to the sliding-window metadata builders, whose kernels read no + window KV below them. """ def __init__( @@ -67,6 +130,20 @@ def __init__( (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device ) + # Per request state index; batches gather from it (see prepare_attn). + self._replay_start_np = np.zeros(self.max_num_reqs, dtype=np.int32) + self._replay_start = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self._replay_start_staging = UvaBufferPool(self.max_num_reqs, torch.int32) + # The replay window and the indices of the prefix-cacheable groups, from + # the KV cache config on first use; None until then. + self._replay: tuple[int, torch.Tensor] | None = None + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + self._replay_start_np[req_index] = new_req_data.replay_start + def prepare_inputs( self, input_batch: InputBatch, req_states: RequestState ) -> dict[str, torch.Tensor | None]: @@ -96,3 +173,67 @@ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any] self.lookback_token_ids.fill_(-1) model_inputs["lookback_token_ids"] = self.lookback_token_ids return model_inputs + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, + ) -> dict[str, Any]: + if self._replay is None: + specs = [group.kv_cache_spec for group in kv_cache_config.kv_cache_groups] + self._replay = ( + max(spec.prefix_replay_tokens for spec in specs), + torch.tensor( + [i for i, spec in enumerate(specs) if spec.prefix_cacheable], + dtype=torch.int32, + device=self.device, + ), + ) + window, cacheable_groups = self._replay + if window: + num_reqs = input_batch.num_reqs + # Decode rows sit above the hit, so only prefills carry a replay + # start; dummy batches (captures, profiling) carry none. + replay_start_np = np.where( + input_batch.is_prefilling_np[:num_reqs], + self._replay_start_np[input_batch.idx_mapping_np[:num_reqs]], + 0, + ).astype(np.int32) + replay_start = self._replay_start_staging.copy_to_gpu( + replay_start_np, out=self._replay_start[:num_reqs] + ) + if replay_start_np.any(): + # The replayed tokens rebuild window KV only: their slots in the + # prefix-cacheable groups are padded so the cached KV stays as is. + _pad_replayed_slots_kernel[(num_reqs,)]( + slot_mappings, + slot_mappings.stride(0), + cacheable_groups, + input_batch.query_start_loc, + input_batch.positions, + replay_start, + window, + PAD_SLOT_ID, + NUM_GROUPS=cacheable_groups.numel(), + BLOCK=1024, + ) + assert model_specific_attn_metadata is None + model_specific_attn_metadata = ReplayAttnMetadata(replay_start) + return super().prepare_attn( + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=for_capture, + ubatch_idx=ubatch_idx, + model_specific_attn_metadata=model_specific_attn_metadata, + ) diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 3d69499d2a80..8db6d4ed2b12 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -195,6 +195,7 @@ def build( if self.compress_ratio > 1: slot_mapping = get_compressed_slot_mapping( cm.num_actual_tokens, + cm.slot_mapping, cm.query_start_loc, cm.seq_lens, cm.block_table_tensor.clamp_(min=0), diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index b32af7234b0c..25e24f347c0e 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -42,6 +42,8 @@ class CompileKey: @triton.jit(do_not_specialize=["block_table_stride"]) def kernel( # [num_tokens] + compressed_slot_mapping_ptr, + # [num_tokens] the tokens' own slots slot_mapping_ptr, # [num_reqs + 1] query_start_loc_ptr, @@ -69,7 +71,12 @@ def kernel( mask = offset < query_len pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE) - is_valid = (pos + 1) % COMPRESS_RATIO == 0 + # A replayed token has a PAD slot (SWA bounded replay: its KV is + # cached already); its compressed KV is cached too, so PAD here. + slot = tl.load( + slot_mapping_ptr + query_start + offset, mask=mask, other=PAD_ID + ) + is_valid = ((pos + 1) % COMPRESS_RATIO == 0) & (slot != PAD_ID) pos_after_compress = pos // COMPRESS_RATIO block_ids = pos_after_compress // block_size @@ -81,7 +88,9 @@ def kernel( # NOTE slot_ids = tl.where(is_valid, slot_ids, PAD_ID) - tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask) + tl.store( + compressed_slot_mapping_ptr + query_start + offset, slot_ids, mask=mask + ) def dispatch( # type: ignore[override] self, @@ -121,6 +130,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: int32_ptr = TritonWarmupTensor(torch.int32) return dict( + compressed_slot_mapping=TritonWarmupTensor(torch.int64), slot_mapping=TritonWarmupTensor(torch.int64), query_start_loc=int32_ptr, seq_lens=int32_ptr, @@ -132,6 +142,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: @kernel_launcher def __call__( self, + compressed_slot_mapping: torch.Tensor, slot_mapping: torch.Tensor, query_start_loc: torch.Tensor, seq_lens: torch.Tensor, @@ -149,6 +160,7 @@ def __call__( def get_compressed_slot_mapping( num_tokens: int, + slot_mapping: torch.Tensor, query_start_loc: torch.Tensor, seq_lens: torch.Tensor, block_table: torch.Tensor, @@ -156,19 +168,28 @@ def get_compressed_slot_mapping( compress_ratio: int, out: torch.Tensor | None = None, ) -> torch.Tensor: + """Slot mapping for writing the compressed states of ``num_tokens`` tokens. + + Every ``compress_ratio`` tokens share one compressed state, written by the + last of them: that token maps to the state's slot, the others to PAD. A + token whose own ``slot_mapping`` entry is PAD maps to PAD too: SWA bounded + replay recomputes tokens whose KV is cached already, and their compressed + states must not be rewritten either. + """ if out is not None: # Guard: for padded / invalid sequences. # Negative positions produce bogus block indices that lead to illegal memory # accesses inside the block_table load. # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`. out.fill_(-1) - slot_mapping = out[:num_tokens] + compressed_slot_mapping = out[:num_tokens] else: - slot_mapping = torch.full( + compressed_slot_mapping = torch.full( (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device ) _COMPRESSED_SLOT_MAPPING_KERNEL( + compressed_slot_mapping, slot_mapping, query_start_loc, seq_lens, @@ -176,7 +197,7 @@ def get_compressed_slot_mapping( block_size, compress_ratio, ) - return slot_mapping + return compressed_slot_mapping _COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel() diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 53206afd9337..430bd10c990c 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -795,6 +795,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.pcp_world_size = parallel_config.prefill_context_parallel_size self.use_pcp = self.pcp_world_size > 1 + self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size # NOTE(Chen):an estimated max size of flattened_kv. Need to double check. self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config) @@ -1217,10 +1218,18 @@ def build( factor = self.kv_cache_spec.block_size // kernel_block_size indexer_block_table = (block_table[:, ::factor] // factor).contiguous() padded_num_tokens = num_tokens - if self.pcp_world_size > 1: + local_slot_mapping = slot_mapping + if self.use_pcp: + # The gathered layout holds each rank's local tokens, padded, in + # rank order, so this rank's segment lines up with query_start_loc. padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size + local_slot_mapping = slot_mapping[ + self.pcp_rank * padded_num_tokens : (self.pcp_rank + 1) + * padded_num_tokens + ] compressed_slot_mapping = get_compressed_slot_mapping( num_tokens, + local_slot_mapping, query_start_loc, seq_lens, indexer_block_table, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 0fa5ae414ce5..6fe20e0e6934 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -81,9 +81,12 @@ def __init__( block_size: int = 64, packed_bytes_per_token: int = 584, packed_page_alignment: int = 576, + bounded_replay: bool = False, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend + # DeepseekV4.1 SWA bounded replay. + self.bounded_replay = bounded_replay self.kv_cache = torch.tensor([]) self.head_dim = head_dim self.window_size = window_size @@ -117,6 +120,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: "nvfp4_ds_mla", ) return SlidingWindowMLASpec( + bounded_replay=self.bounded_replay, block_size=self.block_size, num_kv_heads=1, head_size=self.head_dim, @@ -198,6 +202,10 @@ class DeepseekSparseSWAMetadata: # None when the model is text-only or the batch has no image spans. prefill_left_visible: torch.Tensor | None = None prefill_right_visible: torch.Tensor | None = None + # SWA bounded replay: [num_reqs] lower bound of the positions a request's + # window attention may read (SWA bounded replay); + # zeros when nothing replays. + replay_start: torch.Tensor | None = None # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0 @@ -256,10 +264,12 @@ def get_prefill_chunk_plan( has_compressed = compress_ratio > 1 # query_len <= max_num_batched_tokens and - # gather_len = query_len + min(prefix_len, window_size - 1), so the - # worst-case gathered width is bounded by - # max_num_batched_tokens + window_size - 1. The compressed prefix pool - # is bounded by ceil(max_model_len / compress_ratio). + # gather_len <= query_len + min(prefix_len, window_size - 1) (a replay + # boundary only shrinks it), so the worst-case gathered width is + # bounded by max_num_batched_tokens + window_size - 1. The compressed + # prefix pool is bounded by ceil(max_model_len / compress_ratio). The + # CPU gather lengths below are this upper bound, for chunk planning + # only; the kernels use the exact per-request values. max_workspace_area = prefill_chunk_size * ( (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0) + self.prefill_window_size @@ -331,6 +341,7 @@ def kernel( # Inputs seq_lens_ptr, query_start_loc_ptr, + replay_start_ptr, num_prefills, num_decodes, window_size, @@ -352,7 +363,10 @@ def kernel( query_len = qsl_end - qsl_start prefix_len = seq_len - query_len - gather_len = query_len + tl.minimum(prefix_len, window_size - 1) + # Context below replay_start has no window KV (SWA bounded replay). + replay_start = tl.load(replay_start_ptr + num_decodes + safe_offset, mask=mask) + visible_prefix = tl.maximum(prefix_len - replay_start, 0) + gather_len = query_len + tl.minimum(visible_prefix, window_size - 1) tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) @@ -384,6 +398,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: prefill_gather_lens=int32_ptr, seq_lens=int32_ptr, query_start_loc=int32_ptr, + replay_start=int32_ptr, num_prefills=compile_key.block_size, num_decodes=0, window_size=1, @@ -395,6 +410,7 @@ def __call__( prefill_gather_lens: torch.Tensor, seq_lens: torch.Tensor, query_start_loc: torch.Tensor, + replay_start: torch.Tensor, num_prefills: int, num_decodes: int, window_size: int, @@ -550,12 +566,19 @@ def __init__(self, *args, **kwargs): ) self.decode_swa_indices_noncausal: torch.Tensor | None = None self._max_tokens = max_tokens + # replay_start when no group replays: no lower bound. + self.no_replay_start = torch.zeros( + self.vllm_config.scheduler_config.max_num_seqs + 1, + dtype=torch.int32, + device=self.device, + ) def build( self, common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, + replay_start: torch.Tensor | None = None, ) -> DeepseekSparseSWAMetadata: """Build SWA metadata for mixed decode/prefill batches. @@ -564,6 +587,10 @@ def build( separate window_topk_idxs for each portion. For prefill, we use chunked prefill to align with the indexer's chunking. + + ``replay_start`` ([num_reqs], SWA bounded replay): the position from + which each request holds window KV; the model state passes it for every + batch of a replaying group, graph captures pass nothing. """ seq_lens = common_attn_metadata.seq_lens seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound @@ -587,6 +614,9 @@ def build( is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] is_valid_token.copy_(slot_mapping >= 0) + if replay_start is None: + # Graph captures build without it; their dummy batches replay nothing. + replay_start = self.no_replay_start non_causal = not common_attn_metadata.causal decode_swa_width = ( @@ -609,6 +639,8 @@ def build( device=self.device, ) decode_swa_indices = self.decode_swa_indices_noncausal + # No replay_start: decode rows sit above the hit, so a + # request's replayed window never reaches them. _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( decode_swa_indices, self.decode_swa_lens, @@ -637,6 +669,7 @@ def build( is_valid_token, block_table, self.block_size, + replay_start, num_tokens=num_decode_tokens, token_offset=0, ) @@ -683,6 +716,7 @@ def build( is_valid_token, block_table, self.block_size, + replay_start, num_tokens=num_prefill_tokens, token_offset=num_decode_tokens, has_image=has_image, @@ -696,6 +730,7 @@ def build( seq_lens_cpu, query_start_loc, query_start_loc_cpu, + replay_start, ) # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta @@ -727,6 +762,7 @@ def build( ), prefill_left_visible=prefill_left_visible, prefill_right_visible=prefill_right_visible, + replay_start=replay_start, block_size=self.block_size, num_decodes=num_decodes, num_prefills=num_prefills, @@ -803,6 +839,7 @@ def update_draft_decode_metadata( assert metadata.is_valid_token is not None assert metadata.decode_swa_indices is not None assert metadata.decode_swa_lens is not None + assert metadata.replay_start is not None _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( metadata.decode_swa_indices, @@ -817,6 +854,7 @@ def update_draft_decode_metadata( metadata.is_valid_token, metadata.block_table, self.block_size, + metadata.replay_start, num_tokens=metadata.num_decode_tokens, token_offset=0, ) @@ -872,6 +910,7 @@ def _build_deepseek_v4_metadata( seq_lens_cpu: torch.Tensor | None, query_start_loc: torch.Tensor, query_start_loc_cpu: torch.Tensor, + replay_start: torch.Tensor, ) -> dict[str, torch.Tensor | int | None]: """Pre-compute DeepseekV4 prefill metadata during the metadata build phase. @@ -893,6 +932,7 @@ def _build_deepseek_v4_metadata( pfx_gather_lens, seq_lens, query_start_loc, + replay_start, num_prefills, num_decodes, self.window_size, @@ -982,6 +1022,7 @@ def _compute_swa_indices_and_lens_kernel( block_table_ptr, block_table_stride, block_size, + replay_start_ptr, token_offset, HAS_IMAGE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -1022,6 +1063,8 @@ def _compute_swa_indices_and_lens_kernel( right = 0 left_add = tl.maximum(left - (window_size - 1), 0) start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0) + # SWA bounded replay: no window KV exists below the request's replay start. + start_pos = tl.maximum(start_pos, tl.load(replay_start_ptr + req_idx)) end_pos = pos + right + 1 swa_len = end_pos - start_pos @@ -1111,6 +1154,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: is_valid_token=TritonWarmupTensor(torch.bool), block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), block_size=compile_key.block_size, + replay_start=int32_ptr, num_tokens=1, token_offset=0, has_image=bool(compile_key.has_image), @@ -1131,6 +1175,7 @@ def __call__( is_valid_token: torch.Tensor, block_table: torch.Tensor, block_size: int, + replay_start: torch.Tensor, *, num_tokens: int, token_offset: int, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 49f62a2670a4..1adf386cab2e 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -465,12 +465,17 @@ def allocate_slots( A list of new allocated blocks. """ - # When loading KV data asynchronously, we may have zero new tokens to - # compute while still allocating slots for externally computed tokens. - if num_new_tokens == 0 and num_external_computed_tokens == 0: + # A step may need no slots of its own while still adopting computed + # tokens: an async KV load, or a chunk that ends inside the replayed + # range of a hit (SWA bounded replay). + if ( + num_new_tokens == 0 + and num_external_computed_tokens == 0 + and num_new_computed_tokens == 0 + ): raise ValueError( "num_new_tokens must be greater than 0 when there are no " - "external computed tokens" + "computed tokens to adopt" ) if new_computed_blocks is not None: diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 3a46fa40a71e..7501c51b2ccd 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -48,6 +48,8 @@ class NewRequestData: # Only used for v2 model runner. prefill_token_ids: list[int] | None = None + # DeepSeek-V4.1 only: SWA bounded replay; see Request.replay_start. + replay_start: int = 0 @classmethod def from_request( @@ -73,6 +75,7 @@ def from_request( prompt_embeds=request.prompt_embeds, prompt_is_token_ids=request.prompt_is_token_ids, prefill_token_ids=prefill_token_ids, + replay_start=request.replay_start, ) @property diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index e0c4997b63ff..19ac16284cd7 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -271,6 +271,18 @@ def __init__( self.use_eagle_block_drop = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = vllm_config.num_lookahead_tokens + # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt + # after a prefix hit by recomputing its trailing tokens. One window + # for all such groups, so the rewind matches every group's allocation. + replay_windows = { + group.kv_cache_spec.prefix_replay_tokens + for group in kv_cache_config.kv_cache_groups + if group.kv_cache_spec.prefix_replay_tokens > 0 + } + assert len(replay_windows) <= 1, ( + f"Prefix replay windows should agree: {sorted(replay_windows)}" + ) + self.prefix_replay_tokens = replay_windows.pop() if replay_windows else 0 self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: @@ -946,6 +958,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: step_skipped_waiting.prepend_request(request) continue + if self.prefix_replay_tokens: + # SWA bounded replay recomputes the hit's last + # window, from the block holding its first + # token; the sliding-window groups retire + # whole blocks below the window of the hit's + # next token. Hits end on a block boundary, or + # a hit ending one token short of one would + # retire the block the replay starts in. + ext_tokens -= ext_tokens % self.block_size + load_kv_async = load_kv_async and ext_tokens > 0 + if partial_tail and ext_tokens > partial_tail: # Remote strictly exceeds the full local hit: drop the # sub-block tail so no CoW is needed, and let the load @@ -988,6 +1011,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_new_local_computed_tokens + num_external_computed_tokens ) assert num_computed_tokens <= request.num_tokens + if 0 < num_computed_tokens <= self.prefix_replay_tokens: + # SWA bounded replay: a hit no longer than the replayed + # window would be recomputed in full and save nothing. + new_computed_blocks = ( + self.kv_cache_manager.empty_kv_cache_blocks + ) + num_new_local_computed_tokens = 0 + num_external_computed_tokens = 0 + connector_prefix_cache_hits = 0 + num_computed_tokens = 0 + load_kv_async = False # Skip request with pending mm encoding prefetches if self._ec_transfer_pending(request, num_computed_tokens): @@ -1022,6 +1056,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: new_encoder_compute_budget = encoder_compute_budget pad_spec_decode = False + # SWA bounded replay: recompute the tail of the hit without + # rewriting its cached KV. An async load replays once the KV + # has arrived (_update_waiting_for_remote_kv). + num_replay_tokens = 0 + if load_kv_async: # KVTransfer: loading remote KV, do not allocate for new work. assert num_external_computed_tokens > 0 @@ -1032,6 +1071,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: break else: request_token_budget = min(token_budget, input_budget - draft_slots) + if did_prefix_cache_lookup: + # Fresh admission: local and/or sync external hit. + num_replay_tokens = self._mark_prefix_replay( + request, num_computed_tokens + ) + num_computed_tokens -= num_replay_tokens # Number of tokens to be scheduled. # We use `request.num_tokens` instead of # `request.num_prompt_tokens` to consider the resumed @@ -1159,9 +1204,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + self._spec_decode_step_blocks() ) + # Replayed tokens are already counted in the adopted hit; a + # chunk that ends inside the replayed range needs no new slots. + num_tokens_past_hit = max(num_new_tokens - num_replay_tokens, 0) new_blocks = self.kv_cache_manager.allocate_slots( request, - num_new_tokens, + num_tokens_past_hit, num_new_computed_tokens=num_new_local_computed_tokens, new_computed_blocks=new_computed_blocks, num_lookahead_tokens=effective_lookahead_tokens, @@ -2963,6 +3011,20 @@ def _inflight_prefill_reserved_blocks(self) -> int: self._request_remaining_blocks(req) for req in self._inflight_prefills ) + def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int: + """Record where a prefix hit's replay starts and return the number of + hit tokens to recompute; the caller rewinds the computed count by it. + + The replayed tokens are the hit's last window (see + ``Request.replay_start``): the worker rebuilds their sliding-window KV + and leaves their cached KV alone. + """ + if not self.prefix_replay_tokens: + return 0 + num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens) + request.replay_start = num_hit_tokens - num_replay_tokens + return num_replay_tokens + def _update_waiting_for_remote_kv(self, request: Request) -> None: """KV Connector: update request state after async recv is finished. @@ -2998,11 +3060,16 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: # This will cache the blocks iff caching is enabled. self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) - # on a full prompt hit, we need to re-compute the last token - # in order to be able to sample the next token - if request.num_computed_tokens == request.num_tokens: - request.num_computed_tokens = request.num_tokens - 1 - + # SWA bounded replay recomputes the tail of the hit, which covers the + # last token; otherwise a full prompt hit re-computes that token so the + # next one can be sampled. + num_replay_tokens = self._mark_prefix_replay( + request, request.num_computed_tokens + ) + if num_replay_tokens > 0: + request.num_computed_tokens -= num_replay_tokens + elif request.num_computed_tokens == request.num_tokens: + request.num_computed_tokens = request.num_tokens - 1 self.finished_recving_kv_req_ids.remove(request.request_id) def _try_promote_blocked_waiting_request(self, request: Request) -> bool: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 835a434c1e7c..9fe59dfdbe4d 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -492,6 +492,8 @@ def cache_blocks( can resume at, from ``get_replay_boundaries``. """ + if not self.kv_cache_spec.prefix_cacheable: + return num_cached_blocks = self.num_cached_block.get(request.request_id, 0) num_full_blocks = num_tokens // self.block_size diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 6595be14d21d..fc8b88921a39 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -164,6 +164,11 @@ def prefix_cacheable(self) -> bool: """Whether this spec's group participates in prefix caching.""" return True + @property + def prefix_replay_tokens(self) -> int: + """DeepSeek-V4.1 only: bounded replay. Currently only for DSV41 SWA.""" + return 0 + @property def num_heads(self) -> int: raise NotImplementedError @@ -908,6 +913,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec): # DeepseekV4-only: see MLAAttentionSpec.model_version. alignment: int | None = None # Default to None for no padding. model_version: str | None = None + bounded_replay: bool = False # MLA stores a single latent vector per state; there is no separate V. head_size_v: int = 0 @@ -919,6 +925,14 @@ def __post_init__(self): super().__post_init__() _apply_alignment_padding(self) + @property + def prefix_cacheable(self) -> bool: + return not self.bounded_replay + + @property + def prefix_replay_tokens(self) -> int: + return self.sliding_window if self.bounded_replay else 0 + @classmethod def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), ( @@ -930,16 +944,18 @@ def merge(cls, specs: list[Self]) -> Self: model_version_set = set(spec.model_version for spec in specs) sliding_window_set = set(spec.sliding_window for spec in specs) extra_retained_set = set(spec.extra_retained_tokens for spec in specs) + bounded_replay_set = set(spec.bounded_replay for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 and len(sliding_window_set) == 1 and len(extra_retained_set) == 1 + and len(bounded_replay_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, sliding " - "window size, and retained token count." + "window size, retained token count, and replay policy." ) return cls( block_size=specs[0].block_size, @@ -954,6 +970,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), + bounded_replay=bounded_replay_set.pop(), ) def is_uniform_with_collection( @@ -962,6 +979,7 @@ def is_uniform_with_collection( return all( isinstance(spec, SlidingWindowMLASpec) and spec.sliding_window == self.sliding_window + and spec.bounded_replay == self.bounded_replay for spec in kv_cache_specs.values() ) @@ -1195,6 +1213,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): def prefix_cacheable(self) -> bool: return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values()) + @property + def prefix_replay_tokens(self) -> int: + return max(spec.prefix_replay_tokens for spec in self.kv_cache_specs.values()) + @property def first_spec(self) -> KVCacheSpec: """Return the first spec in the group.""" diff --git a/vllm/v1/request.py b/vllm/v1/request.py index ad875264afb4..dcf679093b66 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -201,6 +201,9 @@ def __init__( # in the (sparse) prefix cache; 0 means none. Set at admission for # hybrid/Mamba models when a shared prefix is detected (Marconi-style). self.shared_prefix_boundary = 0 + # DeepSeek-V4.1 only: SWA bounded replay. The request holds no + # sliding-window KV below this position; 0 when nothing replays. + self.replay_start = 0 # The number of NaNs in logits. A value greater than 0 # indicates that the output is corrupted diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 27a9591e4be0..35ca9ccba102 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -16,7 +16,10 @@ from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.rope import get_rope_state -from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.gpu.model_states.interface import ( + ModelSpecificAttnMetadata, + ModelState, +) from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState from vllm.v1.worker.gpu.states import RequestState @@ -173,6 +176,7 @@ def prepare_attn( kv_cache_config: KVCacheConfig, for_capture: bool = False, ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, ) -> dict[str, Any]: if cudagraph_mode == CUDAGraphMode.FULL: # Use padded sizes - padding is handled by model_runner.prepare_attn. @@ -229,5 +233,6 @@ def prepare_attn( req_idx=input_batch.idx_mapping_np, ubatch_idx=ubatch_idx, fast_prefill=input_batch.fast_prefill, + model_specific_attn_metadata=model_specific_attn_metadata, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py index 32a70e7331ee..19c97e265812 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_only.py +++ b/vllm/v1/worker/gpu/model_states/encoder_only.py @@ -19,6 +19,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup @@ -168,6 +169,7 @@ def prepare_attn( kv_cache_config: KVCacheConfig, for_capture: bool = False, ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, ) -> dict[str, Any]: assert ubatch_idx == 0, "DBO is not supported" attn_metadata = super().prepare_attn( @@ -178,6 +180,7 @@ def prepare_attn( attn_groups, kv_cache_config, for_capture, + model_specific_attn_metadata=model_specific_attn_metadata, ) attn_metadata.update( self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 0ac3f2f63b96..bfad8a1b738f 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -239,8 +239,10 @@ def prepare_attn( kv_cache_config: KVCacheConfig, for_capture: bool = False, ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, ) -> dict[str, Any]: assert ubatch_idx == 0, "DBO is not supported" + assert model_specific_attn_metadata is None if cudagraph_mode == CUDAGraphMode.FULL: num_reqs = input_batch.num_reqs_after_padding num_tokens = input_batch.num_tokens_after_padding From a1bf8ac12d9f1537ff2d233f5ab3d1346fd8bd44 Mon Sep 17 00:00:00 2001 From: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Date: Sat, 19 Sep 2026 06:13:23 +0800 Subject: [PATCH 497/774] [Bugfix][MRV2] Match fast-prefill padding to active LoRA batches (#56456) Signed-off-by: waizuichougou <2082431897@qq.com> Signed-off-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> --- tests/v1/worker/test_attn_utils.py | 61 +++++++++++++++++++++++ vllm/v1/worker/gpu/attn_utils.py | 3 +- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/mm_encoder_model_runner.py | 4 +- 4 files changed, 69 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 2dadc8c980e6..99481a1293f9 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -9,11 +9,13 @@ from types import SimpleNamespace +import numpy as np import pytest import torch import vllm.v1.hisparse.binding as attn_utils_module from tests.v1.attention.utils import dense_kv_cache_views +from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches @@ -30,9 +32,11 @@ ) from vllm.v1.worker.gpu import attn_utils from vllm.v1.worker.gpu.attn_utils import ( + FastPrefillHelper, get_attn_cg_support, get_query_lens_mismatch_unsupported_backend, ) +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.utils import ( AttentionGroup, allocate_kv_cache, @@ -240,6 +244,63 @@ def check( assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set() +@pytest.mark.parametrize( + "num_tokens", [1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128] +) +@pytest.mark.parametrize("num_active_loras", [1, 2, 4]) +def test_fast_prefill_dispatch_preserves_active_lora_count( + num_tokens: int, num_active_loras: int +): + """Fast-prefill padding must match the main dispatch's LoRA variant.""" + + class FakeCudaGraphManager: + device = "cpu" + + def __init__(self): + self.dispatch_calls = [] + + def dispatch(self, **kwargs): + self.dispatch_calls.append(kwargs) + # A captured no-LoRA graph is padded to 8 tokens, while the + # active-LoRA path stays eager at the unpadded token count. + if kwargs["num_active_loras"] == 0: + num_tokens = 8 + mode = CUDAGraphMode.PIECEWISE + else: + num_tokens = kwargs["num_tokens"] + mode = CUDAGraphMode.NONE + return BatchExecutionDescriptor( + cg_mode=mode, + num_tokens=num_tokens, + num_reqs=kwargs["num_reqs"], + num_active_loras=kwargs["num_active_loras"], + num_ubatches=1, + ) + + manager = FakeCudaGraphManager() + helper = FastPrefillHelper(manager, max_num_tokens=max(32, num_tokens)) + metadata = helper.prepare( + torch.arange(num_tokens, dtype=torch.int32), + num_reqs=1, + cu_num_logits_np=np.array([0, num_tokens], dtype=np.int32), + has_prefill=True, + batch_desc=BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=1, + num_active_loras=num_active_loras, + num_ubatches=1, + ), + num_active_loras=num_active_loras, + ) + + assert metadata is not None + assert metadata.num_logits_indices == num_tokens + assert metadata.logits_indices_padded.shape[0] == num_tokens + assert metadata.max_logits_per_req == num_tokens + assert manager.dispatch_calls[-1]["num_active_loras"] == num_active_loras + + class _FakeSharedHostRegion: def __init__(self) -> None: self.cleanup_calls = 0 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 89f54d4d1140..e8b4da8f3629 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -95,6 +95,7 @@ def prepare( cu_num_logits_np: np.ndarray, has_prefill: bool, batch_desc: "BatchExecutionDescriptor", + num_active_loras: int = 0, ) -> FastPrefillBatchMetadata | None: if ( not has_prefill @@ -116,7 +117,7 @@ def prepare( num_reqs=num_reqs, num_tokens=num_logits, uniform_token_count=None, - num_active_loras=0, + num_active_loras=num_active_loras, ) num_logits_padded = min(desc.num_tokens, self.max_num_tokens) return FastPrefillBatchMetadata( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 903e5abd6771..a05658e7ce56 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1271,6 +1271,7 @@ def prepare_inputs( scheduler_output: SchedulerOutput, batch_req_state: "BatchReqState", batch_desc: BatchExecutionDescriptor, + num_active_loras: int, ) -> InputBatch: num_tokens = batch_req_state.num_tokens num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) @@ -1405,6 +1406,7 @@ def prepare_inputs( cu_num_logits_np, batch_req_state.has_prefill, batch_desc, + num_active_loras=num_active_loras, ) # CPU upper bound on seq_lens; padded entries left at zero. @@ -1710,7 +1712,7 @@ def execute_model( cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks) assert batch_req_state is not None input_batch = self.prepare_inputs( - scheduler_output, batch_req_state, batch_desc + scheduler_output, batch_req_state, batch_desc, num_active_loras ) block_tables, slot_mappings = self.prepare_attn(input_batch) # Mamba "align" pre-copy: migrate recurrent state across block diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index 4dc0bf77dd01..f0d19fa12ef2 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/vllm/v1/worker/mm_encoder_model_runner.py @@ -149,7 +149,9 @@ def execute_model( num_tokens=batch_req_state.num_tokens, num_reqs=None, ) - self.prepare_inputs(scheduler_output, batch_req_state, batch_desc) + self.prepare_inputs( + scheduler_output, batch_req_state, batch_desc, num_active_loras=0 + ) scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs if self.lora_config is not None: From 729ebac4983e1510e035bc579142b0fc210a49a3 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 17:04:14 -0700 Subject: [PATCH 498/774] [CI] Reclaim GPU memory between model initialization tests (#57362) Signed-off-by: khluu --- tests/models/test_initialization.py | 45 ++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index e3a08980e90c..073f616af875 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc +from collections.abc import Generator from functools import partial from unittest.mock import patch import pytest +import torch from vllm import LLM +from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.attention.backends.utils import resolve_kv_cache_layout from vllm.v1.core.kv_cache_utils import ( @@ -15,7 +20,12 @@ ) from vllm.v1.engine.core import EngineCore as V1EngineCore -from ..utils import create_new_process_for_each_test, requires_spawn_multiprocessing +from ..utils import ( + create_new_process_for_each_test, + requires_spawn_multiprocessing, + wait_for_gpu_memory_to_clear, + wait_for_memory_to_settle, +) from .registry import ( _TRANSFORMERS_BACKEND_MODELS, AUTO_EXAMPLE_MODELS, @@ -24,6 +34,8 @@ ) from .utils import dummy_hf_overrides +logger = init_logger(__name__) + # This minimal list of model architectures is smaller than the total list of # supported models. The intention is that in the "typical" regression testing # scenario, we only test initializing these models. This subset was chosen @@ -53,6 +65,37 @@ ) +@pytest.fixture(autouse=True) +def cleanup_gpu_memory_between_tests() -> Generator[None, None, None]: + """Reclaim GPU memory left over by each test's forked subprocess. + + These tests are known to leave behind GPU memory (see the comment at + ``gpu_memory_utilization=0.80`` below). Each parametrized case builds an + LLM in a forked subprocess whose VRAM is only reclaimed asynchronously + when the subprocess exits, so the leftovers can OOM a later test in this + module (observed with the heavier multimodal archs, e.g. + InternVLChatModel). Mirrors ``clean_gpu_memory_between_tests`` in the + top-level conftest, but unconditionally for this module. + """ + yield + gc.collect() + if not torch.cuda.is_available(): + return + torch.accelerator.empty_cache() + if current_platform.is_rocm() or current_platform.is_xpu(): + wait_for_memory_to_settle() + return + try: + wait_for_gpu_memory_to_clear( + devices=list(range(torch.accelerator.device_count())), + threshold_ratio=0.1, + ) + except Exception as e: + # NVML is not usable on MIG slices (raises NVMLError_NoPermission); + # cleanup is best-effort, so never fail a test over it. + logger.info("Failed to clean GPU memory: %s", e) + + @create_new_process_for_each_test() def can_initialize( model_arch: str, monkeypatch: pytest.MonkeyPatch, EXAMPLE_MODELS: HfExampleModels From 90531daf6dde1c961d326d43c493477feca510c7 Mon Sep 17 00:00:00 2001 From: hangangqiang <36526001+touch869@users.noreply.github.com> Date: Sat, 19 Sep 2026 09:11:05 +0800 Subject: [PATCH 499/774] [Core][Frontend] Bind KV-event publishers at port 0 and expose the bound endpoints (#55844) Signed-off-by: aoshen02 Co-authored-by: mukeshbhandarkar Co-authored-by: aoshen02 Co-authored-by: Claude Fable 5.1 Co-authored-by: aoshen02 --- rust/src/server/src/routes.rs | 2 + .../src/server/src/routes/kv_event_sources.rs | 25 +++++++ rust/src/server/src/routes/tests.rs | 50 +++++++++++++ tests/distributed/conftest.py | 4 + tests/distributed/test_events.py | 74 +++++++++++++++++++ tests/v1/engine/test_engine_core_client.py | 55 ++++++++++++++ vllm/distributed/kv_events.py | 54 ++++++++++---- vllm/engine/protocol.py | 5 ++ .../entrypoints/serve/instrumentator/basic.py | 7 ++ vllm/v1/engine/async_llm.py | 4 + vllm/v1/engine/core_client.py | 16 ++++ 11 files changed, 282 insertions(+), 14 deletions(-) create mode 100644 rust/src/server/src/routes/kv_event_sources.rs diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 1cd6af3d4fa2..4611d0c6d6ce 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -6,6 +6,7 @@ mod cache; mod collective_rpc; mod health; mod inference; +mod kv_event_sources; mod load; mod lora; mod metrics; @@ -91,6 +92,7 @@ fn build_router_with_options( .route("/metrics", get(metrics::scrape)) .route("/load", get(load::load)) .route("/version", get(version::version)) + .route("/kv_event_sources", get(kv_event_sources::kv_event_sources)) // OpenAI-compatible endpoints .route("/v1/models", get(openai::list_models)) .route("/v1/completions", post(openai::completions)) diff --git a/rust/src/server/src/routes/kv_event_sources.rs b/rust/src/server/src/routes/kv_event_sources.rs new file mode 100644 index 000000000000..c3d36493a9a0 --- /dev/null +++ b/rust/src/server/src/routes/kv_event_sources.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::BTreeMap; +use std::sync::Arc; + +use axum::Json; +use axum::extract::State; +use vllm_engine_core_client::protocol::handshake::KvEventsConfig; + +use crate::state::AppState; + +/// KV-event publisher config of each data-parallel rank, keyed by rank. +pub async fn kv_event_sources( + State(state): State>, +) -> Json> { + Json( + state + .engine_core_client() + .ready_responses() + .into_iter() + .filter_map(|ready| Some((ready.data_parallel_rank, ready.kv_events_config.clone()?))) + .collect(), + ) +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 76d9ab862731..89ce35ce6845 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -2054,6 +2054,56 @@ async fn version_returns_engine_vllm_version() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn kv_event_sources_reports_each_ready_engine() { + let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse { + data_parallel_rank: 1, + kv_events_config: Some( + vllm_engine_core_client::protocol::handshake::KvEventsConfig { + enable_kv_cache_events: true, + publisher: "zmq".to_string(), + endpoint: "tcp://10.0.0.2:41233".to_string(), + replay_endpoint: None, + buffer_steps: 10_000, + hwm: 100_000, + max_queue_size: 100_000, + topic: "kv".to_string(), + }, + ), + ..default_ready_response() + }; + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await; + let response = app + .call( + Request::builder() + .uri("/kv_event_sources") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!( + json, + json!({ + "1": { + "enable_kv_cache_events": true, + "publisher": "zmq", + "endpoint": "tcp://10.0.0.2:41233", + "replay_endpoint": null, + "buffer_steps": 10000, + "hwm": 100000, + "max_queue_size": 100000, + "topic": "kv", + } + }) + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn server_info_endpoint_is_dev_mode_only() { diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index 85a9d73aa5c6..1837d3706e9a 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -91,8 +91,11 @@ def __init__( if isinstance(replay_endpoints, str): replay_endpoints = [replay_endpoints] + ipv6 = any("[" in ep for ep in pub_endpoints + (replay_endpoints or [])) + # Set up subscriber socket - connect to all endpoints self.sub = self.ctx.socket(zmq.SUB) + self.sub.setsockopt(zmq.IPV6, ipv6) self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8")) for endpoint in pub_endpoints: self.sub.connect(endpoint) @@ -103,6 +106,7 @@ def __init__( if replay_endpoints: for replay_endpoint in replay_endpoints: replay = self.ctx.socket(zmq.DEALER) + replay.setsockopt(zmq.IPV6, ipv6) replay.connect(replay_endpoint) self.replay_sockets.append(replay) diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index aae80d9cce31..dee0b99bb0ec 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -332,3 +332,77 @@ def test_event_publisher_factory(random_port): publisher = EventPublisherFactory.create(config, DP_RANK) assert isinstance(publisher, ZmqEventPublisher) publisher.shutdown() + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + [("tcp://*:5557", "tcp://*:5558"), ("tcp://*:0", "tcp://*:0")], +) +def test_offset_endpoint_port_keeps_port_zero(endpoint, expected): + """Port 0 is not offset: every rank asks the OS for its own port.""" + from vllm.distributed.kv_events import ZmqEventPublisher + + assert ZmqEventPublisher.offset_endpoint_port(endpoint, 1) == expected + + +@pytest.mark.parametrize("node_ip", ["127.0.0.1", "::1"]) +def test_ephemeral_publishers_report_bound_endpoints(monkeypatch, node_ip): + """Two ranks binding tcp://*:0 get distinct OS-assigned ports on this + node's IP and report endpoints a subscriber can use for events and replay.""" + from vllm.config.kv_events import KVEventsConfig + from vllm.distributed import kv_events + from vllm.utils.network_utils import split_zmq_path + + from .conftest import MockSubscriber + + monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip) + config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint="tcp://*:0", + replay_endpoint="tcp://*:0", + topic="kv", + ) + pubs = [EventPublisherFactory.create(config, rank) for rank in range(2)] + try: + resolved = [pub.get_publisher_config() for pub in pubs] + endpoints = [e for c in resolved for e in (c.endpoint, c.replay_endpoint)] + hosts_ports = {split_zmq_path(e)[1:] for e in endpoints} + assert {h for h, _ in hosts_ports} == {node_ip} + assert len(hosts_ports) == 4 and "0" not in {p for _, p in hosts_ports} + assert config.endpoint == "tcp://*:0" + + sub = MockSubscriber(endpoints[::2], endpoints[1::2], topic="kv") + try: + # PUB/SUB drops messages sent before the subscription lands, so + # keep publishing until both ranks are heard. + ranks: set[int] = set() + deadline = time.time() + 10 + while ranks != {0, 1} and time.time() < deadline: + for pub in pubs: + pub.publish(create_test_events(1)) + while (result := sub.receive_one(timeout=100)) is not None: + ranks.add(result[1].data_parallel_rank) + assert ranks == {0, 1} + sub.request_replay(0, socket_idx=1) + assert sub.receive_replay(socket_idx=1) + finally: + sub.close() + finally: + for pub in pubs: + pub.shutdown() + + +@pytest.mark.parametrize("node_ip", ["0.0.0.0", "::"]) +def test_ephemeral_bind_needs_a_node_ip(monkeypatch, node_ip): + """tcp://*:0 is advertised as this node's IP; with no usable IP the + publisher refuses instead of reporting a wildcard.""" + from vllm.config.kv_events import KVEventsConfig + from vllm.distributed import kv_events + + monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip) + config = KVEventsConfig( + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://*:0" + ) + with pytest.raises(ValueError, match="VLLM_HOST_IP"): + EventPublisherFactory.create(config, DP_RANK) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index a153b23b10bb..15dec72a4195 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -214,6 +214,7 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM client.core_engines = [bytes([i, 0]) for i in range(num_engines)] client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)] client.eng_start_index = 0 + client._kv_event_sources = {} return client @@ -1490,3 +1491,57 @@ def test_ready_response_queries_executor_with_weight_transfer(): executor.supports_draft_weight_updates.assert_called_once() assert response.supports_draft_weight_updates is True + + +def test_apply_ready_response_retains_kv_event_sources(): + """Each engine's resolved publisher config is kept by DP rank; engines + without KV events contribute nothing.""" + import msgspec + + from vllm.config.kv_events import KVEventsConfig + + client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() + client._kv_event_sources = {} + client.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), + model_config=SimpleNamespace(max_model_len=8192), + ) + + def ready(rank: int, config: KVEventsConfig | None) -> bytes: + return msgspec.msgpack.encode( + EngineCoreReadyResponse( + max_model_len=8192, + num_gpu_blocks=100, + block_size=16, + dp_stats_address=None, + dtype="bfloat16", + vllm_version="test", + world_size=2, + data_parallel_size=2, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=rank, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", + supports_lora=False, + max_loras=0, + kv_events_config=config, + ) + ) + + rank1 = KVEventsConfig( + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:41233" + ) + client._apply_ready_response(ready(1, rank1)) + client._apply_ready_response(ready(0, None)) + assert client.get_kv_event_sources() == {1: rank1} + + # A rank that comes back (elastic re-add) reports its new port. + rank1_again = KVEventsConfig( + enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:47001" + ) + client._apply_ready_response(ready(1, rank1_again)) + assert client.get_kv_event_sources() == {1: rank1_again} diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 3aa2baa17165..9d7ec2f827f4 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -17,6 +17,12 @@ from vllm.config.kv_events import KVEventsConfig from vllm.logger import init_logger +from vllm.utils.network_utils import ( + get_ip, + get_tcp_uri, + is_valid_ipv6_address, + split_zmq_path, +) from vllm.v1.core.kv_cache_utils import ExternalBlockHash logger = init_logger(__name__) @@ -349,6 +355,8 @@ def __init__( replay_endpoint, self._dp_rank ) assert self._endpoint is not None + self._hwm = hwm + self._socket_setup() self._publisher_config = KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", @@ -359,8 +367,6 @@ def __init__( max_queue_size=max_queue_size, topic=topic, ) - self._hwm = hwm - self._socket_setup() # Payload self._seq_gen = count() @@ -423,16 +429,10 @@ def _socket_setup(self) -> None: if self._pub is None: self._pub = self._ctx.socket(zmq.PUB) self._pub.set_hwm(self._hwm) - # Heuristic: bind if wildcard / * present, else connect. - # bind stable, connect volatile convention - if self._endpoint is not None and ( - "*" in self._endpoint - or "::" in self._endpoint - or self._endpoint.startswith("ipc://") - or self._endpoint.startswith("inproc://") - ): - self._pub.bind(self._endpoint) - elif self._endpoint is not None: + assert self._endpoint is not None + if self._is_bind_endpoint(self._endpoint): + self._endpoint = self._bind(self._pub, self._endpoint) + else: self._pub.connect(self._endpoint) # Set up replay socket: use ROUTER @@ -441,7 +441,31 @@ def _socket_setup(self) -> None: # 3) works in our non‑blocking poll loop alongside PUB if self._replay_endpoint is not None: self._replay = self._ctx.socket(zmq.ROUTER) - self._replay.bind(self._replay_endpoint) + self._replay_endpoint = self._bind(self._replay, self._replay_endpoint) + + @staticmethod + def _is_bind_endpoint(endpoint: str) -> bool: + """Bind wildcard and ipc/inproc endpoints; connect to the rest.""" + return ( + "*" in endpoint + or "::" in endpoint + or endpoint.startswith(("ipc://", "inproc://")) + ) + + @staticmethod + def _bind(sock: zmq.Socket, endpoint: str) -> str: + """Bind; tcp://*:0 binds this node's IP and returns the bound endpoint.""" + if endpoint != "tcp://*:0": + sock.bind(endpoint) + return endpoint + ip = get_ip() + if ip in ("0.0.0.0", "::"): + raise ValueError("tcp://*:0 needs a node IP; set VLLM_HOST_IP") + if is_valid_ipv6_address(ip): + sock.setsockopt(zmq.IPV6, 1) + sock.bind(get_tcp_uri(ip, 0)) + _, _, port = split_zmq_path(sock.getsockopt_string(zmq.LAST_ENDPOINT)) + return get_tcp_uri(ip, int(port)) def _publisher_thread(self) -> None: """Background thread that processes the event queue.""" @@ -514,7 +538,7 @@ def offset_endpoint_port( Returns: The endpoint with the port offset by data_parallel_rank - or suffix appended + or suffix appended; a tcp port of 0 is left for the OS. """ # Do nothing if input is None or data_parallel_rank is 0 @@ -529,6 +553,8 @@ def offset_endpoint_port( last_colon_idx = endpoint.rfind(":") base_addr = endpoint[:last_colon_idx] base_port = int(endpoint[last_colon_idx + 1 :]) + if base_port == 0: + return endpoint new_port = base_port + data_parallel_rank return f"{base_addr}:{new_port}" return endpoint diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 602e2570a25f..1e6011a2300f 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any from vllm.config import ModelConfig, VllmConfig +from vllm.config.kv_events import KVEventsConfig from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, WeightTransferUpdateRequest, @@ -62,6 +63,10 @@ def errored(self) -> bool: ... @abstractmethod def dead_error(self) -> BaseException: ... + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: + """KV-event publisher config of each engine, keyed by DP rank.""" + return {} + def check_admission( # noqa: B027 self, n: int = 1, request_id: str | None = None ) -> None: diff --git a/vllm/entrypoints/serve/instrumentator/basic.py b/vllm/entrypoints/serve/instrumentator/basic.py index be091a1f4330..64201db3f3d9 100644 --- a/vllm/entrypoints/serve/instrumentator/basic.py +++ b/vllm/entrypoints/serve/instrumentator/basic.py @@ -54,3 +54,10 @@ async def get_server_load_metrics(request: Request): async def show_version(): ver = {"version": VLLM_VERSION} return JSONResponse(content=ver) + + +@router.get("/kv_event_sources") +async def kv_event_sources(raw_request: Request): + """KV-event publisher config per DP rank; empty on render-only servers.""" + client: EngineClient | None = raw_request.app.state.engine_client + return client.get_kv_event_sources() if client is not None else {} diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7a57d0ff1857..67b8ac560d52 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -12,6 +12,7 @@ import vllm.envs as envs from vllm import TokensPrompt from vllm.config import VllmConfig +from vllm.config.kv_events import KVEventsConfig from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, WeightTransferUpdateRequest, @@ -1282,3 +1283,6 @@ async def update_weight_version(self, new_version: str) -> None: async def get_weight_version(self) -> str: """Return the latest committed weight version.""" return await self.engine_core.get_weight_version_async() + + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: + return self.engine_core.get_kv_event_sources() diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 399af7c0b9d8..6ad05e1f8ced 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -23,6 +23,7 @@ from vllm import envs from vllm.config import VllmConfig +from vllm.config.kv_events import KVEventsConfig from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S from vllm.logger import init_logger from vllm.lora.request import LoRARequest @@ -86,6 +87,10 @@ class EngineCoreClient(ABC): * AsyncMPClient: ZMQ + background proc EngineCore w/ asyncio (for AsyncLLM) """ + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: + """KV-event publisher config of each ready engine, keyed by DP rank.""" + return {} + @staticmethod def make_client( multiprocess_mode: bool, @@ -572,6 +577,7 @@ def __init__( self.vllm_config = vllm_config self._renderer: BaseRenderer | None = renderer self._effective_attention_block_sizes: set[int | None] = set() + self._kv_event_sources: dict[int, KVEventsConfig] = {} # ZMQ setup. sync_ctx = zmq.Context(io_threads=2) @@ -860,6 +866,14 @@ def _apply_ready_response(self, payload: bytes) -> None: else: assert response.dp_stats_address == self.stats_update_address + if response.kv_events_config is not None: + self._kv_event_sources[response.data_parallel_rank] = ( + response.kv_events_config + ) + + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: + return dict(self._kv_event_sources) + def _process_utility_output( output: UtilityOutput, utility_results: dict[int, AnyFuture] @@ -1944,6 +1958,8 @@ async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> No self.lb_engines = self.lb_engines[:new_data_parallel_size] # Pending coordinator snapshots must use the surviving ranks too. self.engine_ranks_managed = self.engine_ranks_managed[:new_data_parallel_size] + for rank in range(new_data_parallel_size, cur_data_parallel_size): + self._kv_event_sources.pop(rank, None) removed_dp_size = cur_data_parallel_size - new_data_parallel_size pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size pause_futures = [ From 468663a5e56eba13367ce94f1b005bea4e183793 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 18 Sep 2026 20:36:45 -0500 Subject: [PATCH 500/774] [ROCm] Bump AITER to v0.1.22.post1 (#56885) Signed-off-by: Micah Williamson --- docker/Dockerfile.rock_base | 2 +- docker/Dockerfile.rocm_base | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index 80b6aaca2964..adde16af3c0f 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -6,7 +6,7 @@ ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0 ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.0.0 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.21.post1" +ARG AITER_BRANCH="v0.1.22.post1" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0 ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index fa0a88fd206e..eb4bb5cc9750 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.11.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.21.post2" +ARG AITER_BRANCH="v0.1.22.post1" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" From 038cb14bd7874030e98bb6edc4b6f4efe649fff7 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 18 Sep 2026 19:44:48 -0600 Subject: [PATCH 501/774] [CI][Bugfix] Correct the Laguna DFlash acceptance-length reference (#57647) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Andreas Karatzas --- tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py index 48dcde2bbced..0dff04333602 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py @@ -53,7 +53,7 @@ class DFlashCorrectnessConfig: model="poolside/Laguna-XS-2.1-NVFP4", draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", expected_accuracy=0.7, # Standard GSM8K sanity floor. - expected_acceptance_len=3.55 * 0.9, + expected_acceptance_len=3.05 * 0.9, num_speculative_tokens=15, max_model_len=8192, max_num_seqs=32, From 2fd6268e6d75f16f05c048cfb26d300ac49f548e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 20:29:00 -0700 Subject: [PATCH 502/774] [CI] Deflake MTEB score tests with per-model mteb_tol and reruns (#57364) Signed-off-by: khluu --- .../language/pooling_mteb_test/mteb_embed_utils.py | 12 ++++++++++-- .../language/pooling_mteb_test/mteb_score_utils.py | 12 ++++++++++-- tests/models/language/pooling_mteb_test/test_baai.py | 5 +++++ .../pooling_mteb_test/test_bge_reranker_v2_gemma.py | 2 ++ .../language/pooling_mteb_test/test_cross_encoder.py | 3 +++ tests/models/language/pooling_mteb_test/test_gte.py | 8 ++++++++ .../language/pooling_mteb_test/test_intfloat.py | 3 +++ tests/models/language/pooling_mteb_test/test_jina.py | 6 ++++++ .../pooling_mteb_test/test_modernbert_fp8.py | 1 + .../language/pooling_mteb_test/test_mxbai_rerank.py | 2 ++ .../language/pooling_mteb_test/test_nemotron.py | 4 ++++ .../models/language/pooling_mteb_test/test_nomic.py | 3 +++ .../pooling_mteb_test/test_qwen3_reranker.py | 3 +++ .../pooling_mteb_test/test_snowflake_arctic_embed.py | 6 ++++++ .../language/pooling_mteb_test/test_st_projector.py | 3 +++ .../models/language/pooling_mteb_test/test_voyage.py | 2 ++ tests/models/utils.py | 4 ++++ 17 files changed, 75 insertions(+), 4 deletions(-) diff --git a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py index b9deecff8d1b..0423e94d2165 100644 --- a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py +++ b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py @@ -161,10 +161,14 @@ def mteb_test_embed_models( model_info: EmbedModelInfo, vllm_extra_kwargs=None, hf_model_callback=None, - atol=MTEB_EMBED_TOL, + atol: float | None = None, prompt_prefix: str | None = None, vllm_model_callback=None, ): + if atol is None: + atol = ( + model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_EMBED_TOL + ) vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs) # Test embed_dims, isnan and whether to use normalize @@ -268,4 +272,8 @@ def mteb_test_embed_models( # We are not concerned that the vllm mteb results are better # than SentenceTransformers, so we only perform one-sided testing. - assert st_main_score - vllm_main_score < atol + diff = st_main_score - vllm_main_score + assert diff < atol, ( + f"diff={diff:.6g} tol={atol} model={model_info.name} " + f"(st={st_main_score}, vllm={vllm_main_score})" + ) diff --git a/tests/models/language/pooling_mteb_test/mteb_score_utils.py b/tests/models/language/pooling_mteb_test/mteb_score_utils.py index 16081cbe238b..2d83a2ebaab6 100644 --- a/tests/models/language/pooling_mteb_test/mteb_score_utils.py +++ b/tests/models/language/pooling_mteb_test/mteb_score_utils.py @@ -230,8 +230,12 @@ def mteb_test_rerank_models( hf_runner=HFMtebCrossEncoder, vllm_extra_kwargs=None, vllm_mteb_encoder=VllmMtebCrossEncoder, - atol=MTEB_RERANK_TOL, + atol: float | None = None, ): + if atol is None: + atol = ( + model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_RERANK_TOL + ) vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs) # Maybe load chat_template. @@ -308,4 +312,8 @@ def mteb_test_rerank_models( # We are not concerned that the vllm mteb results are better # than SentenceTransformers, so we only perform one-sided testing. - assert st_main_score - vllm_main_score < atol + diff = st_main_score - vllm_main_score + assert diff < atol, ( + f"diff={diff:.6g} tol={atol} model={model_info.name} " + f"(st={st_main_score}, vllm={vllm_main_score})" + ) diff --git a/tests/models/language/pooling_mteb_test/test_baai.py b/tests/models/language/pooling_mteb_test/test_baai.py index ec11960fda07..a3a4ffaef196 100644 --- a/tests/models/language/pooling_mteb_test/test_baai.py +++ b/tests/models/language/pooling_mteb_test/test_baai.py @@ -17,6 +17,7 @@ "BAAI/bge-base-en", architecture="BertModel", mteb_score=0.779336792, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -54,6 +55,7 @@ "BAAI/bge-m3", architecture="XLMRobertaModel", mteb_score=0.787343078, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -82,6 +84,7 @@ "BAAI/bge-reranker-base", architecture="XLMRobertaForSequenceClassification", mteb_score=0.32398, + mteb_tol=1e-2, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -101,6 +104,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) @@ -113,6 +117,7 @@ def test_embed_models_correctness( correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py index 23bc95548bcd..b750b622b9fc 100644 --- a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py +++ b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py @@ -26,6 +26,7 @@ "method": "no_post_processing", }, mteb_score=0.33757, + mteb_tol=1e-2, seq_pooling_type="LAST", attn_type="decoder", is_prefix_caching_supported=True, @@ -127,6 +128,7 @@ def get_inputs(pairs, tokenizer, prompt=None): return torch.Tensor(scores) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models( diff --git a/tests/models/language/pooling_mteb_test/test_cross_encoder.py b/tests/models/language/pooling_mteb_test/test_cross_encoder.py index 0d1067d5e2c5..8f5d86fa8d1f 100644 --- a/tests/models/language/pooling_mteb_test/test_cross_encoder.py +++ b/tests/models/language/pooling_mteb_test/test_cross_encoder.py @@ -17,6 +17,7 @@ is_prefix_caching_supported=False, is_chunked_prefill_supported=False, mteb_score=0.32898, + mteb_tol=1e-2, ), RerankModelInfo( "tomaarsen/Qwen3-Reranker-0.6B-seq-cls", @@ -27,10 +28,12 @@ is_chunked_prefill_supported=True, chat_template_name="qwen3_reranker.jinja", mteb_score=0.33459, + mteb_tol=1e-2, ), ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_gte.py b/tests/models/language/pooling_mteb_test/test_gte.py index 0a54262e124f..c6b9ebc274fc 100644 --- a/tests/models/language/pooling_mteb_test/test_gte.py +++ b/tests/models/language/pooling_mteb_test/test_gte.py @@ -18,6 +18,7 @@ EmbedModelInfo( "thenlper/gte-large", mteb_score=0.76807651, + mteb_tol=2e-3, architecture="BertModel", seq_pooling_type="MEAN", attn_type="encoder_only", @@ -44,6 +45,7 @@ "Alibaba-NLP/gte-multilingual-base", architecture="GteNewModel", mteb_score=0.775074696, + mteb_tol=2e-3, hf_overrides={"architectures": ["GteNewModel"]}, seq_pooling_type="CLS", attn_type="encoder_only", @@ -79,6 +81,7 @@ EmbedModelInfo( "Alibaba-NLP/gte-modernbert-base", mteb_score=0.748193353, + mteb_tol=2e-3, architecture="ModernBertModel", seq_pooling_type="CLS", attn_type="encoder_only", @@ -90,6 +93,7 @@ EmbedModelInfo( "Qwen/Qwen3-Embedding-0.6B", mteb_score=0.771163695, + mteb_tol=2e-3, architecture="Qwen3ForCausalLM", seq_pooling_type="LAST", attn_type="decoder", @@ -109,6 +113,7 @@ # classifier_pooling: mean "Alibaba-NLP/gte-reranker-modernbert-base", mteb_score=0.33386, + mteb_tol=1e-2, architecture="ModernBertForSequenceClassification", seq_pooling_type="CLS", attn_type="encoder_only", @@ -119,6 +124,7 @@ RerankModelInfo( "Alibaba-NLP/gte-multilingual-reranker-base", mteb_score=0.33062, + mteb_tol=1e-2, architecture="GteNewForSequenceClassification", hf_overrides={"architectures": ["GteNewForSequenceClassification"]}, seq_pooling_type="CLS", @@ -130,6 +136,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) @@ -142,6 +149,7 @@ def test_embed_models_correctness( correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: vllm_extra_kwargs = {} diff --git a/tests/models/language/pooling_mteb_test/test_intfloat.py b/tests/models/language/pooling_mteb_test/test_intfloat.py index adadb60eeefe..805a3f91edb1 100644 --- a/tests/models/language/pooling_mteb_test/test_intfloat.py +++ b/tests/models/language/pooling_mteb_test/test_intfloat.py @@ -13,6 +13,7 @@ "intfloat/e5-small", architecture="BertModel", mteb_score=0.742285423, + mteb_tol=2e-3, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -29,6 +30,7 @@ "intfloat/multilingual-e5-base", architecture="XLMRobertaModel", mteb_score=0.779325955, + mteb_tol=2e-3, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -48,6 +50,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_jina.py b/tests/models/language/pooling_mteb_test/test_jina.py index 6c9308745e3f..753e6125101f 100644 --- a/tests/models/language/pooling_mteb_test/test_jina.py +++ b/tests/models/language/pooling_mteb_test/test_jina.py @@ -23,6 +23,7 @@ EmbedModelInfo( "jinaai/jina-embeddings-v3", mteb_score=0.824413164, + mteb_tol=2e-3, architecture="XLMRobertaModel", is_matryoshka=True, seq_pooling_type="MEAN", @@ -33,6 +34,7 @@ EmbedModelInfo( "jinaai/jina-embeddings-v5-text-small", mteb_score=0.794535707854956, + mteb_tol=2e-3, architecture="JinaEmbeddingsV5Model", seq_pooling_type="LAST", attn_type="decoder", @@ -42,6 +44,7 @@ EmbedModelInfo( "jinaai/jina-embeddings-v5-text-nano", architecture="JinaEmbeddingsV5Model", + mteb_tol=2e-3, dtype="bfloat16" if current_platform.is_rocm() else "auto", seq_pooling_type="LAST", attn_type="encoder_only", @@ -54,6 +57,7 @@ RerankModelInfo( "jinaai/jina-reranker-v2-base-multilingual", mteb_score=0.33643, + mteb_tol=1e-2, architecture="XLMRobertaForSequenceClassification", seq_pooling_type="CLS", attn_type="encoder_only", @@ -63,6 +67,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: task = "retrieval" if "v5" in model_info.name else "text-matching" @@ -98,6 +103,7 @@ def hf_model_callback(model): ) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py b/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py index 8e062f0d76db..f97057f11fc1 100644 --- a/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py +++ b/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py @@ -56,6 +56,7 @@ def _assert_modernbert_online_fp8(model) -> None: ) +@pytest.mark.flaky(reruns=2) @pytest.mark.skipif( _fp8_scaled_mm_unsupported(), reason="No FP8 ScaledMM kernel is available on this GPU type.", diff --git a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py index 74fe760e7839..1ad1f68b9adf 100644 --- a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py +++ b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py @@ -30,6 +30,7 @@ is_chunked_prefill_supported=True, chat_template_name="mxbai_rerank_v2.jinja", mteb_score=0.33651, + mteb_tol=1e-2, enable_test=True, ), RerankModelInfo( @@ -104,6 +105,7 @@ def compute_logits(inputs): return torch.Tensor(scores) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner) diff --git a/tests/models/language/pooling_mteb_test/test_nemotron.py b/tests/models/language/pooling_mteb_test/test_nemotron.py index 79fae2833990..711224a1a29f 100644 --- a/tests/models/language/pooling_mteb_test/test_nemotron.py +++ b/tests/models/language/pooling_mteb_test/test_nemotron.py @@ -19,6 +19,7 @@ "nvidia/llama-nemotron-embed-1b-v2", architecture="LlamaBidirectionalModel", mteb_score=0.689164662128673, + mteb_tol=2e-3, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -32,6 +33,7 @@ architecture="LlamaBidirectionalForSequenceClassification", chat_template_name="nemotron-rerank.jinja", mteb_score=0.33994, + mteb_tol=1e-2, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -40,11 +42,13 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_nomic.py b/tests/models/language/pooling_mteb_test/test_nomic.py index d7947f73f95e..38e34b78b024 100644 --- a/tests/models/language/pooling_mteb_test/test_nomic.py +++ b/tests/models/language/pooling_mteb_test/test_nomic.py @@ -17,6 +17,7 @@ # HF version and remove revision set. revision="720244025c1a7e15661a174c63cce63c8218e52b", mteb_score=0.737568559, + mteb_tol=2e-3, enable_test=True, seq_pooling_type="MEAN", attn_type="encoder_only", @@ -35,6 +36,7 @@ "nomic-ai/nomic-embed-text-v2-moe", architecture="NomicBertModel", mteb_score=0.715488912, + mteb_tol=2e-3, enable_test=True, seq_pooling_type="MEAN", attn_type="encoder_only", @@ -44,6 +46,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py index 3c182cb046b5..10c90e733e6a 100644 --- a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py +++ b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py @@ -32,6 +32,7 @@ is_prefix_caching_supported=True, is_chunked_prefill_supported=True, mteb_score=0.33459, + mteb_tol=1e-2, enable_test=True, ), RerankModelInfo( @@ -108,11 +109,13 @@ def compute_logits(inputs): return torch.Tensor(scores) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner) +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", RERANK_MODELS) @multi_gpu_test(num_gpus=2) def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None: diff --git a/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py b/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py index f3afbe84fa93..b68b2778b687 100644 --- a/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py +++ b/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py @@ -14,6 +14,7 @@ is_matryoshka=False, architecture="BertModel", mteb_score=0.714927797, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -37,6 +38,7 @@ is_matryoshka=False, architecture="NomicBertModel", mteb_score=0.681146831, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -54,6 +56,7 @@ is_matryoshka=True, architecture="BertModel", mteb_score=0.649088363, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -65,6 +68,7 @@ is_matryoshka=True, architecture="XLMRobertaModel", mteb_score=0.712258299, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -76,6 +80,7 @@ is_matryoshka=True, architecture="GteModel", mteb_score=0.706622444, + mteb_tol=2e-3, seq_pooling_type="CLS", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -85,6 +90,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_st_projector.py b/tests/models/language/pooling_mteb_test/test_st_projector.py index 395846347fb3..7397db8535ff 100644 --- a/tests/models/language/pooling_mteb_test/test_st_projector.py +++ b/tests/models/language/pooling_mteb_test/test_st_projector.py @@ -14,6 +14,7 @@ "TencentBAC/Conan-embedding-v1", architecture="BertModel", mteb_score=0.688611955, + mteb_tol=2e-3, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -24,6 +25,7 @@ "google/embeddinggemma-300m", architecture="Gemma3TextModel", mteb_score=0.7473819294684156, + mteb_tol=2e-3, seq_pooling_type="MEAN", attn_type="encoder_only", is_prefix_caching_supported=False, @@ -33,6 +35,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: mteb_test_embed_models(hf_runner, vllm_runner, model_info) diff --git a/tests/models/language/pooling_mteb_test/test_voyage.py b/tests/models/language/pooling_mteb_test/test_voyage.py index 99ef1de9adfd..628ea31ca0e9 100644 --- a/tests/models/language/pooling_mteb_test/test_voyage.py +++ b/tests/models/language/pooling_mteb_test/test_voyage.py @@ -22,6 +22,7 @@ "num_labels": 2048, }, mteb_score=0.7054, + mteb_tol=2e-3, # === MTEB Results === # STS12: 0.6613 # STS13: 0.6906 @@ -34,6 +35,7 @@ ] +@pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("model_info", MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: # Encoder-only attention models need enforce_eager=True to avoid diff --git a/tests/models/utils.py b/tests/models/utils.py index b52855bcdd30..e4cc5e673b39 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -396,6 +396,8 @@ class ModelInfo: @dataclass class EmbedModelInfo(ModelInfo): mteb_score: float | None = None + # Per-model MTEB score tolerance; falls back to MTEB_EMBED_TOL. + mteb_tol: float | None = None is_matryoshka: bool = False matryoshka_dimensions: list[int] | None = None @@ -403,6 +405,8 @@ class EmbedModelInfo(ModelInfo): @dataclass class RerankModelInfo(ModelInfo): mteb_score: float | None = None + # Per-model MTEB score tolerance; falls back to MTEB_RERANK_TOL. + mteb_tol: float | None = None chat_template_name: str | None = None From 9d75dd454431ccbb3c249b172cf525a5ddc745ee Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 20:30:02 -0700 Subject: [PATCH 503/774] [CI] Deflake pooling shards with GPU teardown fixtures between tests (#57371) Signed-off-by: khluu Co-authored-by: Kimi --- tests/models/language/pooling/conftest.py | 28 +++++++++++++++ tests/models/language/pooling/test_colbert.py | 7 ++-- .../pooling/test_nomic_max_model_len.py | 18 ++++++++++ .../language/pooling_mteb_test/conftest.py | 34 +++++++++++++++++++ tests/models/multimodal/pooling/conftest.py | 34 +++++++++++++++++++ 5 files changed, 119 insertions(+), 2 deletions(-) create mode 100644 tests/models/language/pooling_mteb_test/conftest.py create mode 100644 tests/models/multimodal/pooling/conftest.py diff --git a/tests/models/language/pooling/conftest.py b/tests/models/language/pooling/conftest.py index 0571f1cd6d75..2da4c65133f5 100644 --- a/tests/models/language/pooling/conftest.py +++ b/tests/models/language/pooling/conftest.py @@ -2,12 +2,40 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Pytest configuration for vLLM language generation tests.""" +import gc + +import pytest import torch +from tests.utils import wait_for_memory_to_settle +from vllm.logger import init_logger from vllm.platforms import current_platform +logger = init_logger(__name__) + def pytest_sessionstart(session): """Configure ROCm-specific settings before test session starts.""" if current_platform.is_rocm(): torch.set_float32_matmul_precision("high") + + +@pytest.fixture(autouse=True) +def release_gpu_memory_between_tests(): + """Reclaim cached GPU memory between tests in this shard. + + These tests start many vLLM/HF engines sequentially in one pytest + process; without explicit reclamation between tests, VRAM fragmentation + accumulates and later tests fail at engine startup ("Engine core + initialization failed"). + """ + yield + gc.collect() + if torch.accelerator.is_available(): + torch.accelerator.empty_cache() + try: + wait_for_memory_to_settle() + except ValueError as e: + # Longer-lived (class/module-scoped) engine fixtures may legitimately + # still hold VRAM; nothing further to reclaim in that case. + logger.info("Failed to clean GPU memory: %s", e) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 75323d54bf8d..8829558a065d 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -128,9 +128,12 @@ def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device """Download and return the ColBERT linear projection weight.""" from safetensors.torch import load_file - from vllm.transformers_utils.repo_utils import hf_api + from vllm.transformers_utils.repo_utils import hf_api, with_retry - path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]) + path = with_retry( + lambda: hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]), + f"Error downloading {hf_spec['weights_file']} from {model_name}", + ) weights = load_file(path) return weights[hf_spec["weights_key"]].to(device) diff --git a/tests/models/language/pooling/test_nomic_max_model_len.py b/tests/models/language/pooling/test_nomic_max_model_len.py index 6ea29f6d0677..4ecbe77bc634 100644 --- a/tests/models/language/pooling/test_nomic_max_model_len.py +++ b/tests/models/language/pooling/test_nomic_max_model_len.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: SIM117 +from functools import partial + import pytest from ...utils import EmbedModelInfo @@ -26,6 +28,22 @@ max_model_len = int(original_max_position_embeddings * factor) +@pytest.fixture(scope="module", autouse=True) +def predownload_models(): + """Warm the HF cache so engine startup never hits transient hub errors.""" + from vllm.transformers_utils.repo_utils import hf_api, with_retry + + for model_info in MODELS: + with_retry( + partial( + hf_api().snapshot_download, + model_info.name, + revision=model_info.revision, + ), + f"Error downloading {model_info.name}", + ) + + @pytest.mark.parametrize("model_info", MODELS) def test_default(model_info, vllm_runner): with vllm_runner( diff --git a/tests/models/language/pooling_mteb_test/conftest.py b/tests/models/language/pooling_mteb_test/conftest.py new file mode 100644 index 000000000000..5a31a5f41580 --- /dev/null +++ b/tests/models/language/pooling_mteb_test/conftest.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pytest configuration for vLLM pooling MTEB tests.""" + +import gc + +import pytest +import torch + +from tests.utils import wait_for_memory_to_settle +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@pytest.fixture(autouse=True) +def release_gpu_memory_between_tests(): + """Reclaim cached GPU memory between tests in this shard. + + These tests start many vLLM/HF engines sequentially in one pytest + process; without explicit reclamation between tests, VRAM fragmentation + accumulates and later tests fail at engine startup ("Engine core + initialization failed"). + """ + yield + gc.collect() + if torch.accelerator.is_available(): + torch.accelerator.empty_cache() + try: + wait_for_memory_to_settle() + except ValueError as e: + # Longer-lived (class/module-scoped) engine fixtures may legitimately + # still hold VRAM; nothing further to reclaim in that case. + logger.info("Failed to clean GPU memory: %s", e) diff --git a/tests/models/multimodal/pooling/conftest.py b/tests/models/multimodal/pooling/conftest.py new file mode 100644 index 000000000000..32f598b30a8b --- /dev/null +++ b/tests/models/multimodal/pooling/conftest.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pytest configuration for vLLM multimodal pooling tests.""" + +import gc + +import pytest +import torch + +from tests.utils import wait_for_memory_to_settle +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@pytest.fixture(autouse=True) +def release_gpu_memory_between_tests(): + """Reclaim cached GPU memory between tests in this shard. + + These tests start many vLLM/HF engines sequentially in one pytest + process; without explicit reclamation between tests, VRAM fragmentation + accumulates and later tests fail at engine startup ("Engine core + initialization failed"). + """ + yield + gc.collect() + if torch.accelerator.is_available(): + torch.accelerator.empty_cache() + try: + wait_for_memory_to_settle() + except ValueError as e: + # Longer-lived (class/module-scoped) engine fixtures may legitimately + # still hold VRAM; nothing further to reclaim in that case. + logger.info("Failed to clean GPU memory: %s", e) From 01c7bf88135458940d89584c05ae47404aac77cb Mon Sep 17 00:00:00 2001 From: Jiang Wu Date: Fri, 18 Sep 2026 21:06:30 -0700 Subject: [PATCH 504/774] [Bugfix][Model] Fix Aria expert weight names and layout (#57487) Signed-off-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com> Co-authored-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com> --- tests/models/multimodal/test_aria.py | 100 +++++++++++++++++++++++++++ vllm/model_executor/models/aria.py | 13 ++-- 2 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 tests/models/multimodal/test_aria.py diff --git a/tests/models/multimodal/test_aria.py b/tests/models/multimodal/test_aria.py new file mode 100644 index 000000000000..fed7b2064660 --- /dev/null +++ b/tests/models/multimodal/test_aria.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +from transformers import AriaTextConfig +from transformers.models.aria.modeling_aria import AriaTextMoELayer as HFMoE + +from vllm.config import VllmConfig +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, +) +from vllm.model_executor.models.aria import AriaTextModel, AriaTextMoELayer +from vllm.model_executor.models.utils import AutoWeightsLoader + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +@pytest.mark.parametrize( + "checkpoint_name,loaded_name,param_name", + [ + ( + "experts.fc1.weight", + "experts.w13_weight", + "experts.routed_experts.w13_weight", + ), + ( + "experts.fc2.weight", + "experts.w2_weight", + "experts.routed_experts.w2_weight", + ), + ( + "shared_experts.down_proj.weight", + "shared_experts.down_proj.weight", + "shared_experts.down_proj.weight", + ), + ], +) +def test_aria_expert_weights_load_with_checkpoint_layout( + checkpoint_name: str, loaded_name: str, param_name: str +): + """Real Aria checkpoint weights must reach the right parameter and layout.""" + config = AriaTextConfig( + hidden_size=32, + intermediate_size=64, + moe_num_experts=2, + moe_topk=1, + moe_num_shared_experts=1, + ) + checkpoint_weight = HFMoE(config).state_dict()[checkpoint_name] + checkpoint_weight.copy_( + torch.arange( + checkpoint_weight.numel(), dtype=checkpoint_weight.dtype + ).reshape_as(checkpoint_weight) + ) + expected = checkpoint_weight.clone() + if checkpoint_name.startswith("experts."): + expected = expected.transpose(-1, -2) + + layer = AriaTextMoELayer(config, quant_config=None, prefix="mlp") + loaded = AutoWeightsLoader(layer).load_weights( + [(checkpoint_name, checkpoint_weight)], mapper=AriaTextModel.hf_to_vllm_mapper + ) + + assert loaded == {loaded_name} + torch.testing.assert_close( + dict(layer.named_parameters())[param_name], expected, rtol=0, atol=0 + ) + + +@pytest.mark.cpu_test +def test_aria_quant_config_renames_expert_modules(): + quant_config = CompressedTensorsConfig( + target_scheme_map={}, + ignore=[ + "model.layers.0.mlp.experts.fc1", + "model.layers.0.mlp.experts.fc2", + ], + quant_format="pack-quantized", + ) + AriaTextModel.__new__( + AriaTextModel, vllm_config=VllmConfig(quant_config=quant_config) + ) + + assert quant_config.ignore == [ + "model.layers.0.mlp.experts.gate_up_proj", + "model.layers.0.mlp.experts.down_proj", + ] + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "name,expected", [("fc1", "gate_up_proj"), ("fc2", "down_proj")] +) +def test_aria_expert_scale_names(name: str, expected: str): + suffixes = ["weight_scale", "input_scale"] + names = [f"model.layers.0.mlp.experts.{name}.{suffix}" for suffix in suffixes] + assert AriaTextModel.hf_to_vllm_mapper.apply_list(names) == [ + f"model.layers.0.mlp.experts.{expected}.{suffix}" for suffix in suffixes + ] diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 01119b4a00bd..da91c9d9bf97 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -3,6 +3,7 @@ from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Literal +import regex as re import torch import torch.nn as nn from transformers import AriaConfig, AriaTextConfig, BatchFeature @@ -222,6 +223,7 @@ def __init__( intermediate_size=config.intermediate_size, quant_config=quant_config, prefix=f"{prefix}.experts", + is_fused_checkpoint_transposed=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -267,14 +269,17 @@ class AriaTextModel(LlamaModel, SupportsQuant): "gate_up_proj": ["gate_proj", "up_proj"], } - # Aria packs all experts into single (transposed) fc1/fc2 tensors, which is - # exactly the pre-fused checkpoint layout FusedMoE self-loads once fc1/fc2 - # are renamed to the fused gate_up_proj/down_proj names. + # The fused expert loader expects names without the .weight suffix. hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_regex={ + re.compile(r"experts\.fc1\.weight$"): "experts.gate_up_proj", + re.compile(r"experts\.fc2\.weight$"): "experts.down_proj", + }, + # Quantization configs also use the expert module names. orig_to_new_substr={ "experts.fc1": "experts.gate_up_proj", "experts.fc2": "experts.down_proj", - } + }, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): From a8d1aa9c99b8698a2a78b611b7a10c30e6b3995b Mon Sep 17 00:00:00 2001 From: Hotragn Pettugani <103170876+Hotragn@users.noreply.github.com> Date: Sat, 19 Sep 2026 00:06:48 -0400 Subject: [PATCH 505/774] [Bugfix][Multimodal] Frame the multi-modal hash digest input (#54283) Signed-off-by: hotragn Co-authored-by: Claude Opus 5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/multimodal/test_hasher.py | 40 +++++++++++++++ tests/multimodal/test_processing.py | 37 ++++++++++++++ vllm/multimodal/hasher.py | 75 +++++++++++++++++++++++------ 3 files changed, 138 insertions(+), 14 deletions(-) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 1b738eeeb89f..971d97a2ca94 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -268,3 +268,43 @@ def test_hash_media_io_noop_config_preserves_hash(): assert hasher.hash_kwargs("blake3", image=loaded) == hasher.hash_kwargs( "blake3", image=plain ) + + +# The digest input is a concatenation of byte chunks, so it has to be uniquely +# decodable. Each case below is a pair of distinct processor kwargs that used to +# serialize to the same bytes, which made two requests share an mm hash -- and +# therefore share both the processor cache entry and the prefix-cache block key. +IMAGE = b"\x89PNG\r\n\x1a\n" + + +def _hash(**mm_processor_kwargs: object) -> str: + return MultiModalHasher.hash_kwargs( + "blake3", model_id="m", image=IMAGE, **mm_processor_kwargs + ) + + +def test_hash_collision_kwargs_key_value_boundary(): + # Both used to flatten to b"ab" + b"c". + assert _hash(**{"ab": "c"}) != _hash(**{"a": "bc"}) + + +def test_hash_collision_nested_vs_flattened_key(): + nested = _hash(size={"shortest_edge": 224}) + flattened = _hash(**{"size.shortest_edge": 224}) + assert nested != flattened + + +def test_hash_collision_sequence_vs_mapping(): + assert _hash(fps=[2, 4]) != _hash(fps={"0": 2, "1": 4}) + + +@pytest.mark.parametrize("empty", ["", b"", []]) +def test_hash_collision_none_vs_empty(empty): + assert _hash(video_pruning_rate=None) != _hash(video_pruning_rate=empty) + + +def test_hash_collision_empty_container_vs_omitted(): + omitted = _hash() + assert _hash(size={}) != omitted + assert _hash(size=[]) != omitted + assert _hash(size={}) != _hash(size=[]) diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index 7d6a7fd7e168..de3e333886b8 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -1324,3 +1324,40 @@ def test_processor_inputs_hashes_ignore_unrelated_kwargs(): ) assert inputs.get_mm_hashes("test-model", "blake3") == {"image": ["image-uuid"]} + + +@pytest.mark.parametrize( + ("left", "right"), + [ + # Shifting the key/value boundary: both flatten to the dotted key + # "mm_processor_kwargs.abc" followed by no value bytes. + ({"ab": "c"}, {"a": "bc"}), + # A nested mapping and a caller-supplied dotted key flatten alike. + ({"size": {"shortest_edge": 224}}, {"size.shortest_edge": 224}), + # A sequence and a mapping keyed by stringified indices flatten alike. + ({"fps": [2, 4]}, {"fps": {"0": 2, "1": 4}}), + # None contributes the key alone, which a zero-byte value also does. + ({"video_pruning_rate": None}, {"video_pruning_rate": ""}), + # An empty container contributes nothing, as does omitting the key. + ({"size": {}}, {}), + ], +) +def test_processor_inputs_hashes_distinguish_kwargs_shapes(left, right): + """Distinct processor kwargs must not share a multi-modal hash. + + ``hf_processor_mm_kwargs`` is per-request input, so both the keys and the + values here are caller-controlled. The hash is the identity of the + processor cache entry and is mixed into the prefix-cache block key, so two + requests sharing one is a cross-request cache hit. + """ + image = random_image(np.random.RandomState(0), min_wh=8, max_wh=9) + mm_data_items = MultiModalDataParser().parse_mm_data({"image": [image]}) + + def hash_with(hf_processor_mm_kwargs): + return ProcessorInputs( + prompt=[], + mm_data_items=mm_data_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + ).get_mm_hashes("test-model", "blake3")["image"][0] + + assert hash_with(left) != hash_with(right) diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index eb4ab4d28a40..ae52ceefc409 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -18,6 +18,28 @@ logger = init_logger(__name__) +# Framing for the digest input. The hash is built by feeding a stream of byte +# chunks to `hasher.update`, so the stream must be uniquely decodable: without +# an explicit length in front of every chunk and a tag in front of every +# container, distinct inputs serialize to identical bytes and share a cache +# entry. See `test_hash_collision_*` in tests/multimodal/test_hasher.py. +_LENGTH_BYTES = 8 +_TAG_NONE = b"\x00" +_TAG_SEQUENCE = b"\x01" +_TAG_MAPPING = b"\x02" +_TAG_LEAF = b"\x03" + + +def _encode_length(value: int) -> bytes: + return value.to_bytes(_LENGTH_BYTES, "little") + + +def _framed(chunk: bytes | memoryview) -> Iterable[bytes | memoryview]: + """Yield *chunk* preceded by its size in bytes.""" + size = chunk.nbytes if isinstance(chunk, memoryview) else len(chunk) + yield _encode_length(size) + yield chunk + @functools.lru_cache(maxsize=3) def _get_hasher_factory( @@ -61,20 +83,27 @@ def _get_image_id_bytes(image: Image.Image) -> bytes | None: class MultiModalHasher: + """Derives multi-modal cache keys. + + Every method here yields *framed* chunks: each chunk is preceded by its + length and each container by its kind, so that the concatenation fed to the + digest is uniquely decodable and distinct inputs cannot share a key. + """ + @classmethod def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: # Simple cases if isinstance(obj, (bytes, memoryview)): - return (obj,) + return _framed(obj) if isinstance(obj, str): - return (obj.encode("utf-8"),) + return _framed(obj.encode("utf-8")) if isinstance(obj, (int, float)): - return (np.array(obj).tobytes(),) + return _framed(np.array(obj).tobytes()) if isinstance(obj, Image.Image): image_id = _get_image_id_bytes(obj) if image_id is not None: - return (image_id,) + return _framed(image_id) data = {"mode": obj.mode, "data": np.asarray(obj)} palette = obj.palette @@ -88,7 +117,7 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image): image_id = _get_image_id_bytes(obj.media) if image_id is not None: - return (image_id,) + return _framed(image_id) if obj.io_config: return cls.iter_item_to_bytes( @@ -151,7 +180,7 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: "No serialization method found for %s. Falling back to pickle.", type(obj) ) - return (pickle.dumps(obj),) + return _framed(pickle.dumps(obj)) @classmethod def iter_item_to_bytes( @@ -159,18 +188,36 @@ def iter_item_to_bytes( key: str, obj: object, ) -> Iterable[bytes | memoryview]: + """Yield the digest input for a single ``key``/``obj`` pair.""" + yield from _framed(key.encode("utf-8")) + yield from cls.iter_value_to_bytes(obj) + + @classmethod + def iter_value_to_bytes( + cls, + obj: object, + ) -> Iterable[bytes | memoryview]: + """Yield the digest input for a value, tagged by its container kind. + + Containers carry their kind and length so that a nested structure can + never serialize to the same bytes as a differently shaped one (a list + and a mapping keyed by stringified indices, for example). + """ if obj is None: - yield key.encode("utf-8") - return - # Recursive cases - if isinstance(obj, (list, tuple)): - for i, elem in enumerate(obj): - yield from cls.iter_item_to_bytes(f"{key}.{i}", elem) + yield _TAG_NONE + elif isinstance(obj, (list, tuple)): + yield _TAG_SEQUENCE + yield _encode_length(len(obj)) + for elem in obj: + yield from cls.iter_value_to_bytes(elem) elif isinstance(obj, dict): + yield _TAG_MAPPING + yield _encode_length(len(obj)) for k, v in obj.items(): - yield from cls.iter_item_to_bytes(f"{key}.{k}", v) + yield from _framed(str(k).encode("utf-8")) + yield from cls.iter_value_to_bytes(v) else: - yield key.encode("utf-8") + yield _TAG_LEAF yield from cls.serialize_item(obj) @classmethod From 6af44e2887e3dbfb1c59807d1315af88f232df11 Mon Sep 17 00:00:00 2001 From: Jiang Wu Date: Fri, 18 Sep 2026 22:46:13 -0700 Subject: [PATCH 506/774] [Bugfix][CPU] Fix DeepSeek V4.1 import without Triton (#57654) Signed-off-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com> Co-authored-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com> --- tests/test_triton_utils.py | 4 +++- vllm/triton_utils/importing.py | 2 +- vllm/v1/worker/gpu/sample/gumbel.py | 10 ++++------ 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_triton_utils.py b/tests/test_triton_utils.py index 0ae4c0d6d69d..7f6422466f5e 100644 --- a/tests/test_triton_utils.py +++ b/tests/test_triton_utils.py @@ -97,7 +97,8 @@ def test_triton_placeholder_language(): lang = TritonLanguagePlaceholder() assert isinstance(lang, types.ModuleType) assert lang.__name__ == "triton.language" - assert lang.constexpr is None + assert lang.constexpr(2**31 - 1) == 2**31 - 1 + assert lang.constexpr(1.5) == 1.5 assert lang.dtype is None assert lang.int64 is None assert lang.int32 is None @@ -143,3 +144,4 @@ def test_no_triton_fallback(): assert triton.__class__.__name__ == "TritonPlaceholder" assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder" assert tl.__class__.__name__ == "TritonLanguagePlaceholder" + assert tl.constexpr(2**31 - 1) == 2**31 - 1 diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index 966ab55f326b..e312564c0bd4 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -124,7 +124,7 @@ def decorator(*args, **kwargs): class TritonLanguagePlaceholder(types.ModuleType): def __init__(self): super().__init__("triton.language") - self.constexpr = None + self.constexpr = lambda value: value self.dtype = None self.int64 = None self.int32 = None diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 722021236926..44cbadb9e09c 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,22 +2,20 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.triton_utils import tl, triton # Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the # retained Philox helper for rejection sampling. # # Triton requires globals accessed from `@triton.jit` functions to be wrapped -# in `tl.constexpr(...)`. We can only do that when Triton is actually -# available — on the CPU worker path `tl` is a placeholder whose `constexpr` -# attribute is `None`, and `tl.constexpr(...)` would crash at import time. -_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 +# in `tl.constexpr(...)`. +_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) # Offset salt keeping the draft's Gumbel noise disjoint from the target's. # Verification is a probability-ratio test, not a Gumbel coupling, so a proposal # and the residual it is resampled from must not share a noise vector. # Positions are int64 and never approach 2**30, so the streams cannot collide. -_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) if HAS_TRITON else (1 << 30) +_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) @triton.jit From 59e6682c21595b3f71c4a569318d1a1ac328add3 Mon Sep 17 00:00:00 2001 From: Ryan Rock Date: Sat, 19 Sep 2026 02:13:16 -0500 Subject: [PATCH 507/774] [CI][AMD] Bump torchao to v18 for Python 3.14 (#53009) Signed-off-by: Ryan Rock Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 17b0fa408e8b..5046506e0542 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2540,7 +2540,7 @@ steps: # we can only upgrade after this is resolved # TODO(jerryzh168): resolve the above comment - unset VLLM_USE_V2_MODEL_RUNNER - - uv pip install --system torchao==0.17.0 + - uv pip install --system torchao==0.18.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi From 38537331d3e8b37e480f1bbc174f78888e6c5252 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 19 Sep 2026 00:27:11 -0700 Subject: [PATCH 508/774] [Bugfix][DSA] Avoid runtime JIT for offset candidate end buffers (#57667) Signed-off-by: Woosuk Kwon --- .../layers/test_mla_short_prefill_indexer.py | 9 ++++++--- .../kernels/attention/dsa/sparse_mqa_logits.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 882b745ece79..340d59a39286 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -628,7 +628,10 @@ def test_candidate_blocks_to_sparse_indices_math(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") @pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)]) -def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): +@pytest.mark.parametrize("end_offset", [0, 1]) +def test_candidate_blocks_to_sparse_indices_matches_reference( + cbk: int, sbk: int, end_offset: int +): """The Triton expansion equals the PyTorch reference on production-shaped random inputs, including strided (column-sliced) candidate rows and strided output buffers.""" @@ -663,7 +666,7 @@ def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int ke, cbk, sbk, - out=(si_buf[:, : num_candidates * ratio], end_buf[::2]), + out=(si_buf[:, : num_candidates * ratio], end_buf[end_offset::2]), ) assert torch.equal(got_indices, ref_indices) assert torch.equal(got_end, ref_end) @@ -671,7 +674,7 @@ def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7) ) assert not si_buf[:, num_candidates * ratio :].any() - assert not end_buf[1::2].any() + assert not end_buf[1 - end_offset :: 2].any() def _skip_unless_sm100_sparse_kernels(): diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index 0c2ea1aef718..345c711ae645 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -73,7 +73,7 @@ def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: _INT32_BIG = tl.constexpr(2**31 - 1) -@triton.jit +@triton.jit(do_not_specialize_on_alignment=["end_ptr"]) def _expand_candidates_kernel( cand_ptr, cand_stride, From c3b48446349569512749db7f6e2164aa8a33437d Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 19 Sep 2026 00:53:42 -0700 Subject: [PATCH 509/774] [Docs] Explain pointer-alignment JIT specialization in Triton skill (#57669) Signed-off-by: Woosuk Kwon --- .agents/skills/triton-kernel-writing/SKILL.md | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/.agents/skills/triton-kernel-writing/SKILL.md b/.agents/skills/triton-kernel-writing/SKILL.md index 6dd6a8520be0..a41de016b78d 100644 --- a/.agents/skills/triton-kernel-writing/SKILL.md +++ b/.agents/skills/triton-kernel-writing/SKILL.md @@ -22,10 +22,14 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo a small, legible heuristic when workloads need different choices. Use `triton.autotune` only when tuning is critical to performance, such as for a matrix multiplication. Otherwise prioritize simple code and fast startup. -- Be careful to avoid unintended runtime JIT compilation. For example, put - unimportant runtime integer scalars in `do_not_specialize`, especially those - that may alternate between values such as 0 and 1, which can produce - different specialization keys. +- Avoid unintended runtime JIT compilation. Put unimportant runtime integer + scalars in `do_not_specialize`, especially values that alternate between 0 + and 1. Pointer alignment also affects specialization: an offset tensor view + can trigger a new variant after warmup used an aligned allocation. For + scalar-only pointers whose alignment provides no useful optimization, use + `@triton.jit(do_not_specialize_on_alignment=["end_ptr"])`. Retain alignment + specialization when vectorized accesses benefit from it. See the + [Triton JIT options](https://triton-lang.org/main/python-api/generated/triton.jit.html). - The Triton compiler does not guarantee safe ordering when a kernel writes to a pointer and subsequently reads from the same pointer. This pattern must have a `tl.debug_barrier()` between the write and read. The barrier @@ -49,7 +53,9 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo ## Validation - Check correctness at boundary shapes and at sizes that exercise masks and - large offsets. + large offsets. Include offset tensor views, not only fresh allocations. + Warm aligned buffers, then exercise offset views with vLLM's + `--jit-monitor-mode error --jit-monitor-verbose` to catch missing variants. - Choose accumulation and intermediate dtypes explicitly. Test numerically difficult inputs, not only random, well-scaled tensors. - Use `$kernel-microbenchmark` for benchmark construction, measurement, and From a5a30471ff2bb7f0824f2da10e358af98d304472 Mon Sep 17 00:00:00 2001 From: Aditya Ganesh Kumar Date: Sat, 19 Sep 2026 17:13:54 +0800 Subject: [PATCH 510/774] [Tests] Cover get_unhashed_block_ids_all_groups (#55928) Signed-off-by: Aditya Ganesh Kumar --- tests/v1/core/test_prefix_caching.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 40e810c58094..d6848a6a7f3d 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -5937,3 +5937,29 @@ def test_device_eviction_keeps_host_prefix(): assert host_block.block_hash is not None, ( "Device eviction invalidated unrelated host KV" ) + + +def test_get_unhashed_block_ids_all_groups(): + """Unhashed, non-null block ids are reported per KV cache group. + A group with no unhashed blocks reports an empty list, not a missing entry. + """ + + def hashed(block_id: int, group_id: int) -> KVCacheBlock: + block = KVCacheBlock(block_id=block_id) + block.set_block_hash(make_block_hash_with_group_id(BlockHash(b"h"), group_id)) + return block + + blocks = KVCacheBlocks( + ( + [ + KVCacheBlock(block_id=1), # unhashed + hashed(2, group_id=0), # cached + KVCacheBlock(block_id=3, is_null=True), # null padding + KVCacheBlock(block_id=4), # unhashed + ], + # No unhashed blocks in this group. + [hashed(5, group_id=1), KVCacheBlock(block_id=6, is_null=True)], + ) + ) + + assert blocks.get_unhashed_block_ids_all_groups() == [[1, 4], []] From 4cc15f2121b3421478f3186aac38de901b4a44ef Mon Sep 17 00:00:00 2001 From: siyu Date: Sat, 19 Sep 2026 17:16:39 +0800 Subject: [PATCH 511/774] [Quantization] Let ModelOpt MXFP8 layers load pre-processed weights Purpose (#57316) Signed-off-by: liusy58 Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com> --- tests/quantization/test_modelopt.py | 30 +++++++++++++++++++ .../kernels/linear/mxfp8/Mxfp8LinearKernel.py | 4 +++ .../kernels/linear/mxfp8/emulation.py | 2 ++ .../kernels/linear/mxfp8/flashinfer.py | 4 +++ .../kernels/linear/mxfp8/rocm_native.py | 2 ++ .../layers/quantization/modelopt.py | 19 +++++++++--- 6 files changed, 57 insertions(+), 4 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 4ca1af53a2ae..ae8548c37058 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -144,6 +144,36 @@ def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypat assert torch.equal(linear.weight_scale, scales) +def test_modelopt_mxfp8_pre_processed_weights_follow_kernel(monkeypatch): + """The weight cache daemon can only serve MXFP8 layers whose kernel needs + no post-load state beyond the parameters it exports.""" + from vllm.config.quantization import QuantSpec + from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearKernel + + method = ModelOptLinearMethod.__new__(ModelOptLinearMethod) + method.spec = QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic) + method.kernel = None + assert not method.supports_pre_processed_weights + + method.kernel = Mock(spec=Mxfp8LinearKernel) + method.kernel.supports_pre_processed_weights = False + assert not method.supports_pre_processed_weights + method.kernel.supports_pre_processed_weights = True + assert method.supports_pre_processed_weights + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt.is_weights_pre_processed", + lambda: True, + ) + layer = torch.nn.Module() + method.process_weights_after_loading(layer) + method.kernel.process_weights_after_loading.assert_not_called() + + method.kernel.supports_pre_processed_weights = False + with pytest.raises(RuntimeError, match="pre-processed"): + method.process_weights_after_loading(layer) + + def test_modelopt_fp8_updates_weight_dims_after_transpose(): """Humming reads weight.input_dim/output_dim. Swapping the ModelWeightParameter for a plain Parameter drops them, so the per-tensor diff --git a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py index 3fe34029e5cb..fdcfeb264b29 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py +++ b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py @@ -27,6 +27,10 @@ class Mxfp8LinearKernel(ABC): Marlin, emulation). """ + supports_pre_processed_weights: bool = False + """True if ``process_weights_after_loading`` only rewrites parameters, so + weights exported by the weight cache daemon can be used as-is.""" + def __init__(self, c: Mxfp8LinearLayerConfig) -> None: assert self.can_implement(c)[0] assert self.is_supported()[0] diff --git a/vllm/model_executor/kernels/linear/mxfp8/emulation.py b/vllm/model_executor/kernels/linear/mxfp8/emulation.py index 79b2fba3889b..da48655d37db 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/emulation.py +++ b/vllm/model_executor/kernels/linear/mxfp8/emulation.py @@ -16,6 +16,8 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel): """Software emulation fallback for MXFP8 (dequant to BF16).""" + supports_pre_processed_weights = True + @classmethod def is_supported( cls, compute_capability: int | None = None diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index 02700a0ced0c..cc9b35d34bdd 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -27,6 +27,8 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+).""" + supports_pre_processed_weights = True + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -117,6 +119,8 @@ def apply_weights( class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel): """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103).""" + supports_pre_processed_weights = True + @classmethod def is_supported( cls, compute_capability: int | None = None diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 1d9789ccedcf..c7e6537f4948 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -212,6 +212,8 @@ def _select_cfg(M, N, K): class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" + supports_pre_processed_weights = True + @classmethod def is_supported( cls, compute_capability: int | None = None diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index c649061f339c..364c52b4c879 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2482,9 +2482,16 @@ def __init__( @property def supports_pre_processed_weights(self) -> bool: # type: ignore[override] - # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack. + # TODO(Isotr0py): support fp8 ModelOpt kernels transpose/repack. w = self.spec.weight - return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE + return isinstance(w, QuantKey) and ( + w.dtype == FP4_DTYPE + or ( + w == kMxfp8Static + and self.kernel is not None + and self.kernel.supports_pre_processed_weights + ) + ) def create_weights( self, @@ -2524,7 +2531,13 @@ def create_weights( expose_input_quant_key(layer, self.kernel) def process_weights_after_loading(self, layer) -> None: + if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False): + self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size) if is_weights_pre_processed(): + if not self.supports_pre_processed_weights: + raise RuntimeError( + f"{type(self.kernel).__name__} cannot use pre-processed weights" + ) return self.fmt.pre_process(layer) self.wkey.process(layer, WEIGHT) @@ -2556,8 +2569,6 @@ def process_weights_after_loading(self, layer) -> None: persistent=False, ) layer._nvfp4_group_size_for_gather = self.ctx.group_size - if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False): - self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size) self.kernel.process_weights_after_loading(layer) def apply(self, layer, x, bias=None): From 211e252d0b4f8429f9b15fc52bdfed07782c7f70 Mon Sep 17 00:00:00 2001 From: vllmellm Date: Sat, 19 Sep 2026 21:05:00 +0800 Subject: [PATCH 512/774] [Bugfix][Model] Fix MiMo-V2.5 fused fp8 qkv_proj sharding (pre-shard count is num_key_value_heads; MTP path too) (#57508) Signed-off-by: vllmellm --- .../quantization/test_mimo_v2_qkv_shard.py | 225 ++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 196 ++++++++++----- vllm/model_executor/models/mimo_v2_mtp.py | 57 ++++- 3 files changed, 415 insertions(+), 63 deletions(-) create mode 100644 tests/models/quantization/test_mimo_v2_qkv_shard.py diff --git a/tests/models/quantization/test_mimo_v2_qkv_shard.py b/tests/models/quantization/test_mimo_v2_qkv_shard.py new file mode 100644 index 000000000000..91eb036bc7c5 --- /dev/null +++ b/tests/models/quantization/test_mimo_v2_qkv_shard.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit test for MiMo-V2 fused fp8 ``qkv_proj`` sharding. + +The checkpoint stores the fused QKV pre-sharded at ``num_key_value_heads`` +chunks (``ckpt_tp``), each holding that slice's ``[Q_c | K_c | V_c]`` rows with +its own block scale grid. ``ckpt_tp`` is not the layer's KV-head count: a +MiMo-V2.5 SWA layer has 8 KV heads but 4 chunks of 3712 rows (= 29 blocks), so +each chunk carries two KV heads. Reading the scales as one slice per KV head +expands 14 x 128 = 1792 rows for an 1856-row slice and fails. + +These tests build such checkpoints for the GA, SWA and Pro geometries and check +that every rank gets exactly the weights it should own, for TP sizes below, +equal to and above ``ckpt_tp`` (the last case replicating KV heads). +""" + +import pytest +import torch + +from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj +from vllm.utils.math_utils import cdiv + +pytestmark = pytest.mark.cpu_test + +BLOCK = 128 +FP8_DTYPE = torch.float8_e4m3fn +FP8_MAX = torch.finfo(FP8_DTYPE).max +COLS = 128 # one column block is enough to exercise the row mapping +# fp8 e4m3 has 3 mantissa bits, so a single block-quantization round trip costs +# ~2-3% relative L2; a wrong row -> scale pairing costs >= 10%. +TOL = 0.06 + +# (num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): MiMo-V2.5 GA, its +# SWA/MTP layers, and MiMo-V2.5-Pro. +GEOMETRIES = { + "ga": (64, 4, 192, 128, 4), + "swa": (64, 8, 192, 128, 4), + "pro": (128, 8, 192, 128, 8), +} + + +def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): + return ( + num_heads // ckpt_tp * head_dim + + num_kv_heads // ckpt_tp * head_dim + + num_kv_heads // ckpt_tp * v_head_dim + ) + + +def _quantize_chunks(truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): + """Block-quantize each checkpoint chunk on its own grid, like the export.""" + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp) + weights, scales = [], [] + for chunk in range(ckpt_tp): + rows = truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] + quantized = torch.zeros_like(rows) + scale = torch.zeros( + cdiv(rows_per_chunk, BLOCK), rows.shape[1] // BLOCK, dtype=torch.float32 + ) + for row in range(0, rows_per_chunk, BLOCK): + for col in range(0, rows.shape[1], BLOCK): + block = rows[row : row + BLOCK, col : col + BLOCK] + amax = block.abs().max().clamp(min=1e-12) / FP8_MAX + scale[row // BLOCK, col // BLOCK] = amax + quantized[row : row + BLOCK, col : col + BLOCK] = block / amax + weights.append(quantized.to(FP8_DTYPE)) + scales.append(scale) + return torch.cat(weights, dim=0), torch.cat(scales, dim=0) + + +def _owned_rows( + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, tp_rank, tp_size +): + """Rows the rank must own, in the forward's `[Q | K | V]` order.""" + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp) + q_per_chunk = num_heads // ckpt_tp * head_dim + k_per_chunk = num_kv_heads // ckpt_tp * head_dim + kv_heads: list[int] = [] + if tp_size <= num_kv_heads: + kv_heads = list( + range( + tp_rank * (num_kv_heads // tp_size), + (tp_rank + 1) * (num_kv_heads // tp_size), + ) + ) + else: + kv_heads = [tp_rank // (tp_size // num_kv_heads)] + index = [] + for head in range( + tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size) + ): + chunk, offset = divmod(head, num_heads // ckpt_tp) + start = chunk * rows_per_chunk + offset * head_dim + index += list(range(start, start + head_dim)) + for head in kv_heads: + chunk, offset = divmod(head, num_kv_heads // ckpt_tp) + start = chunk * rows_per_chunk + q_per_chunk + offset * head_dim + index += list(range(start, start + head_dim)) + for head in kv_heads: + chunk, offset = divmod(head, num_kv_heads // ckpt_tp) + start = chunk * rows_per_chunk + q_per_chunk + k_per_chunk + offset * v_head_dim + index += list(range(start, start + v_head_dim)) + return truth[torch.tensor(index)] + + +def _dequantize(weight, scale): + """What the kernel computes: scale[row // block] applied per weight row.""" + per_row = scale.repeat_interleave(BLOCK, dim=1) + return weight.to(torch.float32) * per_row[torch.arange(weight.shape[0]) // BLOCK] + + +def _relative_l2(got, expected): + return ((got - expected).norm() / expected.norm()).item() + + +@pytest.mark.parametrize("geometry", list(GEOMETRIES)) +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +def test_fused_qkv_proj_sharding(geometry, tp_size): + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry] + if num_heads % tp_size or (tp_size > num_kv_heads and tp_size % num_kv_heads): + pytest.skip("TP size does not divide the head counts") + + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp) + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS) + for chunk in range(ckpt_tp): # per-chunk magnitude spread, as in the export + truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk + weight, scale = _quantize_chunks( + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp + ) + + kv_per_rank = (num_kv_heads // tp_size) if tp_size <= num_kv_heads else 1 + rows_rank = (num_heads // tp_size) * head_dim + kv_per_rank * ( + head_dim + v_head_dim + ) + for tp_rank in range(tp_size): + w_rank, s_rank = _shard_fp8_qkv_proj( + weight, + scale, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + v_head_dim=v_head_dim, + tp_rank=tp_rank, + tp_size=tp_size, + ckpt_tp=ckpt_tp, + ) + assert w_rank.shape == (rows_rank, COLS) + assert s_rank.shape == (cdiv(rows_rank, BLOCK), scale.shape[1]) + expected = _owned_rows( + truth, + num_heads, + num_kv_heads, + head_dim, + v_head_dim, + ckpt_tp, + tp_rank, + tp_size, + ) + assert _relative_l2(_dequantize(w_rank, s_rank), expected) <= TOL + + +@pytest.mark.parametrize("geometry", ["ga", "swa"]) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_exact_chunk_matches_checkpoint(geometry, tp_size): + """At tp_size <= ckpt_tp the chunk count is the only thing that matters.""" + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry] + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp) + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS) + weight, scale = _quantize_chunks( + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp + ) + if tp_size == ckpt_tp: + w_rank, s_rank = _shard_fp8_qkv_proj( + weight, + scale, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + v_head_dim=v_head_dim, + tp_rank=0, + tp_size=tp_size, + ckpt_tp=ckpt_tp, + ) + # The checkpoint chunk *is* the rank's slice: no re-quantization. + # (torch.equal has no fp8 kernel, so compare the underlying bytes.) + expected_w = weight.chunk(tp_size, dim=0)[0] + assert torch.equal(w_rank.view(torch.uint8), expected_w.view(torch.uint8)) + assert torch.equal(s_rank, scale.chunk(tp_size, dim=0)[0]) + + +def test_wrong_chunk_count_is_detected(): + """Reading SWA scales as one slice per KV head must not look correct. + + This is the regression the fix addresses: 116 scale rows are 4 x 29 (29 = + ceil(3712 / 128)), not 8 x 15, so ``ckpt_tp = num_kv_heads`` pairs every + weight row with the wrong scale. + """ + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES["swa"] + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp) + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS) + for chunk in range(ckpt_tp): + truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk + weight, scale = _quantize_chunks( + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp + ) + assert scale.shape[0] == ckpt_tp * cdiv(rows_per_chunk, BLOCK) + assert scale.shape[0] != num_kv_heads * cdiv(rows_per_chunk, BLOCK) + + kw = dict( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + v_head_dim=v_head_dim, + tp_rank=0, + tp_size=1, + ) + correct = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=ckpt_tp, **kw) + wrong = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=num_kv_heads, **kw) + assert correct[0].shape == wrong[0].shape + + expected = _owned_rows( + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, 0, 1 + ) + assert _relative_l2(_dequantize(*correct), expected) <= TOL + assert _relative_l2(_dequantize(*wrong), expected) > 5 * TOL diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index ec67d72cf2ca..be665e7e064a 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -50,6 +50,7 @@ ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -460,6 +461,28 @@ def is_compressed_softmax_layer(self) -> bool: return self.config.hybrid_layer_pattern[self.layer_id] == 1 +def _requantize_fp8( + grouped: torch.Tensor, rows_rank: int, block: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Block-quantize a rank's ``[Q | K | V]`` rows back to fp8. + + A rank's rows need not end on a block boundary (a single 1856-row slice is + 14.5 blocks) while ``scaled_quantize`` requires both dims to be multiples of + the block size, so pad the tail with zeros: zero rows cannot raise a block's + amax, so every scale -- and the number of scale rows -- is unchanged. The + padding is dropped again here. + """ + padded = cdiv(rows_rank, block) * block + if padded != rows_rank: + grouped = torch.cat( + [grouped, grouped.new_zeros(padded - rows_rank, grouped.shape[1])], dim=0 + ) + w_rank, s_rank = scaled_quantize( + grouped, GroupShape(block, block), dtype, compute_dtype=torch.float32 + ) + return w_rank[:rows_rank], s_rank + + def _shard_fp8_qkv_proj( w_full: torch.Tensor, s_full: torch.Tensor, @@ -469,75 +492,136 @@ def _shard_fp8_qkv_proj( v_head_dim: int, tp_rank: int, tp_size: int, + ckpt_tp: int, block: int = 128, ) -> tuple[torch.Tensor, torch.Tensor]: """Shard the fp8 qkv_proj weights for ``tp_rank``. - The checkpoint stores the fused QKV as ``num_kv_heads`` contiguous groups - (one per KV head; ``n`` below), each ordered ``[Q | K | V]``: + The checkpoint stores the fused QKV pre-sharded for ``ckpt_tp`` ranks (the + model config's ``num_key_value_heads``), each chunk holding that slice's + Q, K and V rows: - [Q_1 | K_1 | V_1 | Q_2 | K_2 | V_2 | ... | Q_n | K_n | V_n] + [Q_0 | K_0 | V_0 | Q_1 | K_1 | V_1 | ... | Q_n | K_n | V_n] (n = ckpt_tp) - Per group, Q has ``(num_heads / num_kv_heads) * head_dim`` rows, K has - ``head_dim`` rows, and V has ``v_head_dim`` rows. + Per chunk, Q has ``(num_heads / ckpt_tp) * head_dim`` rows, K has + ``(num_kv_heads / ckpt_tp) * head_dim`` rows, and V has + ``(num_kv_heads / ckpt_tp) * v_head_dim`` rows, and the fp8 block scales + are tiled per chunk too (``ceil(rows_per_chunk / block)`` rows each). - Each TP rank owns ``g = num_kv_heads / tp_size`` of these groups, and the - forward expects them de-interleaved into a single Q, K, and V block: + ``ckpt_tp`` is not the layer's KV-head count: a MiMo-V2.5 SWA layer has 8 + KV heads over 4 chunks, so each chunk carries two KV heads (3712 rows = 29 + blocks) while a GA layer has 4 KV heads over 4 chunks (3392 rows each). + + The forward expects each rank's slice de-interleaved: [Q_1 | Q_2 | ... | Q_g | K_1 | K_2 | ... | K_g | V_1 | V_2 | ... | V_g] - When ``g == 1`` the rank's slice is already ``[Q | K | V]``, so a plain - chunk suffices. When ``g > 1`` we cannot reach the de-interleaved layout by - re-permuting the fp8 block scales: each scale covers a 128-row block, and - since K is 192 rows (1.5 blocks) a block straddles the K/V boundary, so no - whole-block permutation produces it. Instead we dequantize this rank's - groups to float (dropping the block constraint), reorder the rows into the - layout above (Q, K, and V then each span a whole number of blocks), and - re-quantize to fp8. + When ``tp_size == ckpt_tp`` the checkpoint chunk *is* that layout, so a + plain chunk of both weight and scale suffices. Otherwise each rank's Q, K + and V rows are gathered from the chunks that hold them, dequantized with + the chunk's own scales, reordered, and re-quantized to fp8. """ - assert tp_size <= num_kv_heads and num_kv_heads % tp_size == 0, ( - "TP size must evenly split the number of KV heads." + assert num_heads % tp_size == 0, ( + f"num_heads={num_heads} must be divisible by tp_size={tp_size}." ) - - kv_heads_per_rank = num_kv_heads // tp_size - if kv_heads_per_rank == 1: - # One KV head per rank. The weights and scale can be trivially sharded - # without re-quantization. - w = w_full.chunk(tp_size, dim=0)[tp_rank] - s = s_full.chunk(tp_size, dim=0)[tp_rank] - return w, s - - q_rows_per_group = (num_heads // num_kv_heads) * head_dim - k_rows_per_group = head_dim - v_rows_per_group = v_head_dim - rows_per_group = q_rows_per_group + k_rows_per_group + v_rows_per_group - scale_rows_per_group = s_full.shape[0] // num_kv_heads - qs, ks, vs = [], [], [] - for g_idx in range(tp_rank * kv_heads_per_rank, (tp_rank + 1) * kv_heads_per_rank): - row_start = g_idx * rows_per_group - scale_row_start = g_idx * scale_rows_per_group - # Dequantize this group's weights. - w_g = w_full[row_start : row_start + rows_per_group].to(torch.float32) - s_g = s_full[scale_row_start : scale_row_start + scale_rows_per_group].to( - torch.float32 - ) - s_g_expanded = s_g.repeat_interleave(block, dim=0).repeat_interleave( - block, dim=1 - )[:rows_per_group] - w_g_dequant = w_g * s_g_expanded - # Track the dequantized q, k, and v weights separately. - qs.append(w_g_dequant[:q_rows_per_group]) - ks.append(w_g_dequant[q_rows_per_group : q_rows_per_group + k_rows_per_group]) - vs.append(w_g_dequant[q_rows_per_group + k_rows_per_group :]) - - # Combine the q, k, and v weights into the following layout: - # [Q_1, Q_2, .., Q_g, K_1, K_2, ..., K_g, V_1, V_2, ..., V_g] - grouped = torch.cat([torch.cat(qs), torch.cat(ks), torch.cat(vs)], dim=0) - # Quantize back to fp8. - return scaled_quantize( - grouped, GroupShape(block, block), w_full.dtype, compute_dtype=torch.float32 + if ckpt_tp <= 0 or num_heads % ckpt_tp or num_kv_heads % ckpt_tp: + raise ValueError( + f"fused qkv_proj is pre-sharded at {ckpt_tp} chunks, which do not " + f"divide num_heads={num_heads} / num_kv_heads={num_kv_heads}." + ) + # When there are fewer KV heads than ranks, vLLM replicates them + # (`num_kv_head_replicas`) and rank r owns KV head r // replicas, which + # keeps every Q head grouped with the KV head it attends to. + if tp_size <= num_kv_heads: + assert num_kv_heads % tp_size == 0, ( + f"num_kv_heads={num_kv_heads} must be divisible by tp_size={tp_size}." + ) + kv_head_ids = list( + range( + tp_rank * (num_kv_heads // tp_size), + (tp_rank + 1) * (num_kv_heads // tp_size), + ) + ) + else: + assert tp_size % num_kv_heads == 0, ( + f"tp_size={tp_size} must be divisible by num_kv_heads={num_kv_heads}." + ) + kv_head_ids = [tp_rank // (tp_size // num_kv_heads)] + q_head_ids = list( + range(tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size)) ) + rows_per_chunk = w_full.shape[0] // ckpt_tp + q_per_chunk = (num_heads // ckpt_tp) * head_dim + k_per_chunk = (num_kv_heads // ckpt_tp) * head_dim + v_per_chunk = (num_kv_heads // ckpt_tp) * v_head_dim + if q_per_chunk + k_per_chunk + v_per_chunk != rows_per_chunk: + raise ValueError( + f"fused qkv_proj has {w_full.shape[0]} rows, not {ckpt_tp} chunks of " + f"{q_per_chunk} Q + {k_per_chunk} K + {v_per_chunk} V rows." + ) + + # The scales are tiled per chunk; they collapse to one continuous grid when + # a chunk is a whole number of blocks (the SWA chunks are: 3712 = 29 * 128). + chunk_scale_rows = cdiv(rows_per_chunk, block) + rows = torch.arange(w_full.shape[0]) + per_chunk_scales = s_full.shape[0] == ckpt_tp * chunk_scale_rows + if per_chunk_scales: + scale_index = (rows // rows_per_chunk) * chunk_scale_rows + ( + rows % rows_per_chunk + ) // block + elif s_full.shape[0] == cdiv(w_full.shape[0], block): + scale_index = rows // block + else: + raise ValueError( + f"fused qkv_proj scale has {s_full.shape[0]} rows, expected either " + f"{ckpt_tp * chunk_scale_rows} ({ckpt_tp} chunks of " + f"{chunk_scale_rows} rows) or {cdiv(w_full.shape[0], block)} " + f"(one continuous grid)" + ) + + if tp_size == ckpt_tp and per_chunk_scales: + # One checkpoint chunk per rank: already [Q | K | V] for that rank, and + # its scale rows line up with ceil(rows_per_chunk / block). + return ( + w_full.chunk(ckpt_tp, dim=0)[tp_rank], + s_full.chunk(ckpt_tp, dim=0)[tp_rank], + ) + + # Gather this rank's Q, K and V rows from the chunks that hold them. + q_heads_per_chunk = num_heads // ckpt_tp + kv_heads_per_chunk = num_kv_heads // ckpt_tp + head_rows = torch.arange(head_dim) + v_head_rows = torch.arange(v_head_dim) + row_index: list[torch.Tensor] = [] + for head in q_head_ids: + chunk = head // q_heads_per_chunk + row_index.append( + chunk * rows_per_chunk + (head % q_heads_per_chunk) * head_dim + head_rows + ) + for head in kv_head_ids: + chunk = head // kv_heads_per_chunk + row_index.append( + chunk * rows_per_chunk + + q_per_chunk + + (head % kv_heads_per_chunk) * head_dim + + head_rows + ) + for head in kv_head_ids: + chunk = head // kv_heads_per_chunk + row_index.append( + chunk * rows_per_chunk + + q_per_chunk + + k_per_chunk + + (head % kv_heads_per_chunk) * v_head_dim + + v_head_rows + ) + index = torch.cat(row_index) + grouped = w_full[index].to(torch.float32) * s_full[ + scale_index[index] + ].repeat_interleave(block, dim=1) + return _requantize_fp8(grouped, index.numel(), block, w_full.dtype) + @support_torch_compile class MiMoV2Model(nn.Module, EagleModelMixin): @@ -819,6 +903,8 @@ def _try_load_fp8_qkv_proj( v_head_dim=attn.v_head_dim, tp_rank=tp_rank, tp_size=tp_size, + # The fused qkv_proj is pre-sharded for this many ranks. + ckpt_tp=self.config.num_key_value_heads, ) sharded = {"weight": w_rank, "weight_scale_inv": s_rank} for kind, tensor in sharded.items(): diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index b08ba229f7a1..6fd8309350aa 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -45,7 +45,7 @@ SupportsMultiModal, _require_is_multimodal, ) -from .mimo_v2 import MiMoV2Attention, MiMoV2MLP +from .mimo_v2 import MiMoV2Attention, MiMoV2MLP, _shard_fp8_qkv_proj from .utils import _merge_multimodal_embeddings, maybe_prefix # MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports @@ -266,6 +266,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + pending_qkv_proj: dict[str, dict[str, torch.Tensor]] = {} for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: @@ -279,14 +280,54 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ): continue - # Support fused qkv_proj checkpoint (Pro format). - # The checkpoint is stored pre-sharded for TP=8 as - # [Q_rank0, K_rank0, V_rank0, Q_rank1, ...], so splitting along - # dim 0 with chunk(tp_size) gives each rank its Q+K+V slice for - # both the FP8 weight and the block weight_scale_inv. This matches - # how the main model loads the same layout. + # Fused qkv_proj: the fp8 weight and its block scales arrive as + # separate tensors and are pre-sharded at `config.num_key_value_heads` + # chunks of [Q_c | K_c | V_c] (see _shard_fp8_qkv_proj for the + # layout). Shard them together: a plain chunk(tp_size) only happens + # to be right when tp_size is that chunk count *and* the scales are + # tiled per chunk (they are 4 x 29 = 116 rows for MiMo-V2.5, not + # tp_size slices of an even split). if "qkv_proj" in name: - if name in params_dict: + is_fp8_weight = ( + name.endswith("qkv_proj.weight") + and loaded_weight.dtype == torch.float8_e4m3fn + ) + is_fp8_scale = name.endswith("qkv_proj.weight_scale_inv") + if is_fp8_weight or is_fp8_scale: + prefix, kind = name.rsplit(".", 1) + if ( + f"{prefix}.weight" not in params_dict + or f"{prefix}.weight_scale_inv" not in params_dict + ): + continue + entry = pending_qkv_proj.setdefault(prefix, {}) + entry[kind] = loaded_weight + if "weight" not in entry or "weight_scale_inv" not in entry: + # Waiting for the other half of the fused projection. + continue + del pending_qkv_proj[prefix] + + # Geometry comes from the module the weights belong to. + attn = self.get_submodule(prefix.rsplit(".", 1)[0]) + w_rank, s_rank = _shard_fp8_qkv_proj( + entry["weight"], + entry["weight_scale_inv"], + num_heads=attn.total_num_heads, + num_kv_heads=attn.total_num_kv_heads, + head_dim=attn.head_dim, + v_head_dim=attn.v_head_dim, + tp_rank=tp_rank, + tp_size=tp_size, + ckpt_tp=self.config.num_key_value_heads, + ) + for k, tensor in (("weight", w_rank), ("weight_scale_inv", s_rank)): + param = params_dict[f"{prefix}.{k}"] + if tensor.shape[0] > param.shape[0]: + tensor = tensor[: param.shape[0]] + default_weight_loader(param, tensor) + loaded_params.add(f"{prefix}.{k}") + elif name in params_dict: + # Non-fp8 fused checkpoint: unchanged behaviour. param = params_dict[name] loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank] default_weight_loader(param, loaded_weight) From fbe8a157fbdd4b2f813ae6395397232de97d5c5d Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Sat, 19 Sep 2026 16:05:04 +0300 Subject: [PATCH 513/774] [ROCm][Perf] Enable the AITER GDN decode fast path for flat qkvz layouts (#53623) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor Co-authored-by: TJian --- .buildkite/test-amd.yaml | 1 + .buildkite/test_areas/kernels.yaml | 1 + .../mamba/test_gdn_rocm_layout_dispatch.py | 124 ++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 14 +- 4 files changed, 132 insertions(+), 8 deletions(-) create mode 100644 tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5046506e0542..fe2cfd18c89e 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1834,6 +1834,7 @@ steps: source_file_dependencies: - csrc/mamba/ - tests/kernels/mamba + - vllm/model_executor/layers/mamba/gdn - vllm/model_executor/layers/mamba/ops - vllm/platforms/rocm.py commands: diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 7e43b123f12e..aab6e52b8eb9 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -316,6 +316,7 @@ steps: source_file_dependencies: - csrc/mamba/ - tests/kernels/mamba + - vllm/model_executor/layers/mamba/gdn - vllm/model_executor/layers/mamba/ops commands: - pytest -v -s kernels/mamba diff --git a/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py b/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py new file mode 100644 index 000000000000..ae34ffcb8058 --- /dev/null +++ b/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Layout dispatch tests for ``QwenGatedDeltaNetAttention._forward_core_rocm``. + +The AITER fused reshape+conv kernel learned Qwen3.5's flat ``[q|k|v|z]`` +packing in https://github.com/ROCm/aiter/pull/3251, so flat-layout models take +the decode fast path too instead of falling back to the generic path. + +These tests run host-side on CPU: ``_forward_core_rocm`` is bound to a stub +layer whose ``_forward_core_decode_aiter``/``_forward_core`` record which +branch ran, so no GPU or AITER install is needed. +""" + +from __future__ import annotations + +import types +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + +PREFIX = "model.layers.0.linear_attn" +H = 2 # num key heads +HV = 4 # num value heads +K = 8 # head_k_dim +V = 8 # head_v_dim + + +def _make_metadata( + *, + num_prefills: int = 0, + num_decodes: int = 2, + spec_sequence_masks: torch.Tensor | None = None, +): + """Only the fields the dispatch condition reads carry meaningful values.""" + return GDNAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefills, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_actual_tokens=num_prefills + num_decodes, + spec_sequence_masks=spec_sequence_masks, + ) + + +def _make_layer(gqa_interleaved_layout: bool): + """Stub layer running the real ``_forward_core_rocm``, recording dispatch.""" + layer = types.SimpleNamespace() + layer.prefix = PREFIX + layer.gqa_interleaved_layout = gqa_interleaved_layout + layer.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat" + + layer.calls = [] + layer._forward_core_decode_aiter = lambda **kw: layer.calls.append("aiter") + layer._forward_core = lambda **kw: layer.calls.append("generic") + layer.prepare_gdn_attention_core_inputs = lambda qkvz, ba, n: ( + torch.zeros(n, H * K * 2 + HV * V), + torch.zeros(n, HV, V), + torch.zeros(n, HV), + torch.zeros(n, HV), + ) + layer._forward_core_rocm = types.MethodType( + QwenGatedDeltaNetAttention._forward_core_rocm, layer + ) + return layer + + +def _run(layer, meta) -> str: + num_tokens = max(meta.num_actual_tokens, 1) + ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta}) + with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx): + layer._forward_core_rocm( + qkvz=torch.zeros(num_tokens, 2 * H * K + 2 * HV * V), + ba=torch.zeros(num_tokens, 2 * HV), + z_out=torch.zeros(num_tokens, HV, V), + core_attn_out=torch.zeros(num_tokens, HV, V), + ) + assert len(layer.calls) == 1 + return layer.calls[0] + + +@pytest.mark.parametrize("gqa_interleaved_layout", [True, False]) +def test_pure_decode_takes_the_fast_path(gqa_interleaved_layout: bool) -> None: + """Both packings reach the fast path; flat used to be guarded out.""" + layer = _make_layer(gqa_interleaved_layout) + assert _run(layer, _make_metadata()) == "aiter" + + +@pytest.mark.parametrize("gqa_interleaved_layout", [True, False]) +def test_layout_string_matches_the_packing(gqa_interleaved_layout: bool) -> None: + """The value handed to the kernel's ``qkvz_layout`` parameter.""" + layer = _make_layer(gqa_interleaved_layout) + expected = "interleaved" if gqa_interleaved_layout else "flat" + assert layer.qkvz_layout == expected + + +@pytest.mark.parametrize("gqa_interleaved_layout", [True, False]) +@pytest.mark.parametrize( + "meta_kwargs", + [ + {"num_prefills": 1}, + {"num_decodes": 0}, + {"spec_sequence_masks": torch.zeros(1, dtype=torch.bool)}, + ], + ids=["has_prefill", "no_decode", "spec_decode"], +) +def test_non_pure_decode_batches_use_the_generic_path( + gqa_interleaved_layout: bool, meta_kwargs: dict +) -> None: + layer = _make_layer(gqa_interleaved_layout) + assert _run(layer, _make_metadata(**meta_kwargs)) == "generic" diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 807a5b9b1a2a..ccf1c60f0fe2 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -399,6 +399,7 @@ def __init__( self.key_dim = self.head_k_dim * self.num_k_heads self.value_dim = self.head_v_dim * self.num_v_heads self.gqa_interleaved_layout = gqa_interleaved_layout + self.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat" if current_platform.is_xpu(): self._forward_method = self.forward_xpu elif current_platform.is_cpu(): @@ -1206,9 +1207,9 @@ def _forward_core_rocm( """ROCm AITER fast path: conv1d + recurrent attention from packed qkvz/ba layout. - For decode-only (no spec, no prefill) interleaved-GQA layouts, - dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks - the packed layout and falls through to ``_forward_core``. + For decode-only (no spec, no prefill) batches, dispatches directly to + ``_forward_core_decode_aiter``. Otherwise unpacks the packed layout and + falls through to ``_forward_core``. Args: qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim) @@ -1231,12 +1232,8 @@ def _forward_core_rocm( assert isinstance(attn_metadata, GDNAttentionMetadata) - # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved - # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use - # the generic path below to split/rearrange inputs correctly. if ( - self.gqa_interleaved_layout - and attn_metadata.spec_sequence_masks is None + attn_metadata.spec_sequence_masks is None and attn_metadata.num_prefills == 0 and attn_metadata.num_decodes > 0 ): @@ -1625,6 +1622,7 @@ def _forward_core_decode_aiter( : attn_metadata.num_actual_tokens ], validate_data=True, + qkvz_layout=self.qkvz_layout, ) ) From 41c4a3ed4e45792739bd8a5bbcd01269b06eae2c Mon Sep 17 00:00:00 2001 From: Lio Einaudi Date: Sat, 19 Sep 2026 21:30:55 +0800 Subject: [PATCH 514/774] [Kernel][Perf] Add sm_120 (RTX PRO 6000 / RTX 50) tuned configs for batch-invariant persistent matmul (#57456) Signed-off-by: LioEinaudi --- .../determinism/batch_invariant_configs.py | 229 ++++++++++++++++++ 1 file changed, 229 insertions(+) diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py index 39cda984f7f6..874447bb8295 100644 --- a/vllm/model_executor/determinism/batch_invariant_configs.py +++ b/vllm/model_executor/determinism/batch_invariant_configs.py @@ -254,6 +254,233 @@ class _MatmulShapeConfig(NamedTuple): ), ), }, + "sm120": { + (12288, 2048): _MatmulShapeConfig( # Qwen3-1.7B gate_up_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 128, 8, 5)), + (4, _MatmulMConfig(16, 128, 8, 4)), + (8, _MatmulMConfig(16, 128, 8, 4)), + (16, _MatmulMConfig(16, 128, 4, 4)), + (32, _MatmulMConfig(32, 128, 8, 5)), + (64, _MatmulMConfig(64, 128, 4, 3)), + (256, _MatmulMConfig(64, 64, 4, 4)), + (512, _MatmulMConfig(64, 128, 4, 5)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2048, 6144): _MatmulShapeConfig( # Qwen3-1.7B down_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 32, 8, 5)), + (4, _MatmulMConfig(16, 32, 8, 3)), + (8, _MatmulMConfig(16, 32, 8, 4)), + (16, _MatmulMConfig(16, 32, 4, 5)), + (32, _MatmulMConfig(32, 32, 4, 5)), + (64, _MatmulMConfig(32, 32, 4, 4)), + (256, _MatmulMConfig(64, 64, 4, 5)), + (512, _MatmulMConfig(128, 64, 4, 3)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 64, 8, 4)), + ), + ), + (4096, 2048): _MatmulShapeConfig( # Qwen3-1.7B qkv_proj + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 4, 4)), + (4, _MatmulMConfig(16, 64, 8, 5)), + (8, _MatmulMConfig(16, 64, 8, 5)), + (16, _MatmulMConfig(16, 64, 8, 3)), + (32, _MatmulMConfig(16, 64, 8, 5)), + (64, _MatmulMConfig(32, 64, 4, 2)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(64, 64, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + (151936, 2048): _MatmulShapeConfig( # Qwen3-1.7B lm_head + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 5)), + (4, _MatmulMConfig(16, 128, 4, 4)), + (8, _MatmulMConfig(16, 128, 4, 5)), + (16, _MatmulMConfig(16, 128, 8, 3)), + (32, _MatmulMConfig(32, 128, 8, 3)), + (64, _MatmulMConfig(32, 128, 8, 5)), + (256, _MatmulMConfig(128, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 8, 3)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2048, 2048): _MatmulShapeConfig( # Qwen3-1.7B o_proj + block_k=32, + m_buckets=( + (1, _MatmulMConfig(16, 128, 4, 3)), + (4, _MatmulMConfig(16, 128, 4, 4)), + (8, _MatmulMConfig(16, 64, 4, 5)), + (16, _MatmulMConfig(32, 32, 8, 4)), + (32, _MatmulMConfig(16, 64, 4, 5)), + (64, _MatmulMConfig(32, 32, 4, 5)), + (256, _MatmulMConfig(128, 32, 4, 5)), + (512, _MatmulMConfig(128, 64, 4, 4)), + (1024, _MatmulMConfig(128, 128, 8, 5)), + (2048, _MatmulMConfig(128, 256, 8, 4)), + ), + ), + (19456, 2560): _MatmulShapeConfig( # Qwen3-4B gate_up_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 128, 8, 4)), + (4, _MatmulMConfig(16, 128, 8, 3)), + (8, _MatmulMConfig(16, 128, 8, 5)), + (16, _MatmulMConfig(16, 128, 8, 5)), + (32, _MatmulMConfig(32, 128, 4, 4)), + (64, _MatmulMConfig(64, 128, 8, 4)), + (256, _MatmulMConfig(128, 128, 4, 3)), + (512, _MatmulMConfig(64, 128, 4, 5)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2560, 9728): _MatmulShapeConfig( # Qwen3-4B down_proj + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 4, 5)), + (4, _MatmulMConfig(16, 64, 4, 4)), + (8, _MatmulMConfig(16, 64, 4, 4)), + (16, _MatmulMConfig(16, 32, 4, 5)), + (32, _MatmulMConfig(32, 32, 4, 3)), + (64, _MatmulMConfig(32, 64, 8, 4)), + (256, _MatmulMConfig(64, 64, 4, 3)), + (512, _MatmulMConfig(64, 128, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + (6144, 2560): _MatmulShapeConfig( # Qwen3-4B qkv_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 128, 8, 4)), + (4, _MatmulMConfig(16, 256, 4, 3)), + (8, _MatmulMConfig(16, 128, 4, 5)), + (16, _MatmulMConfig(16, 64, 8, 3)), + (32, _MatmulMConfig(16, 128, 8, 4)), + (64, _MatmulMConfig(64, 64, 4, 5)), + (256, _MatmulMConfig(128, 128, 8, 3)), + (512, _MatmulMConfig(64, 64, 4, 5)), + (1024, _MatmulMConfig(64, 128, 8, 4)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (151936, 2560): _MatmulShapeConfig( # Qwen3-4B lm_head + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 5)), + (4, _MatmulMConfig(16, 64, 8, 5)), + (8, _MatmulMConfig(16, 64, 4, 5)), + (16, _MatmulMConfig(16, 64, 4, 5)), + (32, _MatmulMConfig(32, 64, 4, 5)), + (64, _MatmulMConfig(64, 64, 4, 5)), + (256, _MatmulMConfig(128, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 4, 3)), + (1024, _MatmulMConfig(128, 128, 8, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2560, 4096): _MatmulShapeConfig( # Qwen3-4B o_proj + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 3)), + (4, _MatmulMConfig(16, 64, 8, 4)), + (8, _MatmulMConfig(16, 64, 8, 5)), + (16, _MatmulMConfig(16, 64, 4, 4)), + (32, _MatmulMConfig(16, 32, 4, 4)), + (64, _MatmulMConfig(16, 64, 8, 3)), + (256, _MatmulMConfig(64, 64, 8, 3)), + (512, _MatmulMConfig(64, 128, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + (24576, 4096): _MatmulShapeConfig( # Qwen3-8B gate_up_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 256, 8, 3)), + (4, _MatmulMConfig(16, 256, 8, 3)), + (8, _MatmulMConfig(16, 256, 8, 3)), + (16, _MatmulMConfig(16, 256, 8, 3)), + (32, _MatmulMConfig(32, 256, 8, 3)), + (64, _MatmulMConfig(64, 256, 4, 3)), + (256, _MatmulMConfig(64, 128, 4, 5)), + (512, _MatmulMConfig(128, 128, 4, 3)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (4096, 12288): _MatmulShapeConfig( # Qwen3-8B down_proj + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 4, 5)), + (4, _MatmulMConfig(16, 64, 4, 5)), + (8, _MatmulMConfig(16, 64, 4, 4)), + (16, _MatmulMConfig(16, 32, 4, 4)), + (32, _MatmulMConfig(32, 32, 8, 5)), + (64, _MatmulMConfig(32, 64, 4, 5)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(64, 64, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + (6144, 4096): _MatmulShapeConfig( # Qwen3-8B qkv_proj + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 4)), + (4, _MatmulMConfig(16, 128, 8, 4)), + (8, _MatmulMConfig(16, 128, 4, 4)), + (16, _MatmulMConfig(16, 64, 4, 5)), + (32, _MatmulMConfig(16, 128, 8, 5)), + (64, _MatmulMConfig(64, 64, 4, 4)), + (256, _MatmulMConfig(128, 128, 8, 3)), + (512, _MatmulMConfig(64, 64, 4, 5)), + (1024, _MatmulMConfig(128, 64, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (151936, 4096): _MatmulShapeConfig( # Qwen3-8B lm_head + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 5)), + (4, _MatmulMConfig(16, 64, 8, 5)), + (8, _MatmulMConfig(16, 64, 4, 5)), + (16, _MatmulMConfig(16, 64, 4, 5)), + (32, _MatmulMConfig(32, 64, 4, 5)), + (64, _MatmulMConfig(32, 128, 8, 5)), + (256, _MatmulMConfig(128, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 8, 3)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (4096, 4096): _MatmulShapeConfig( # Qwen3-8B o_proj + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 4)), + (4, _MatmulMConfig(16, 64, 4, 3)), + (8, _MatmulMConfig(16, 32, 4, 4)), + (16, _MatmulMConfig(16, 32, 4, 3)), + (32, _MatmulMConfig(16, 64, 4, 5)), + (64, _MatmulMConfig(64, 32, 8, 4)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(64, 64, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + }, } _TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = ( @@ -265,6 +492,8 @@ class _MatmulShapeConfig(NamedTuple): def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None: if capability is None: return None + if capability.major == 12: + return "sm120" if capability.major == 10: return "blackwell" if capability.major == 9: From 751f6807d9cb3de50c27a5f27188c4fb04fe0e2b Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Sat, 19 Sep 2026 23:38:00 +0800 Subject: [PATCH 515/774] [Refactor] Use dynamic MM cache in processor (#57674) Signed-off-by: DarkLight1337 --- .../processing/test_audio_in_video.py | 24 +-- .../multimodal/processing/test_common.py | 22 ++- .../multimodal/processing/test_gemma4.py | 15 +- .../multimodal/processing/test_mllama4.py | 6 +- .../multimodal/processing/test_moss_audio.py | 6 +- .../multimodal/processing/test_qwen2_vl.py | 2 +- .../processing/test_tensor_schema.py | 2 +- .../processing/test_transformers_image.py | 1 + .../processing/transformers_backend.py | 2 +- tests/multimodal/test_cache.py | 69 ++++--- tests/multimodal/test_processing.py | 37 +++- tests/multimodal/test_registry.py | 30 --- tests/renderers/test_warmup.py | 172 +++++++++--------- vllm/model_executor/models/llava.py | 5 - .../models/qwen3_asr_realtime.py | 35 ++-- .../model_executor/models/voxtral_realtime.py | 19 +- vllm/multimodal/encoder_budget.py | 65 +++---- vllm/multimodal/processing/inputs.py | 3 + vllm/multimodal/processing/processor.py | 43 ++++- vllm/multimodal/registry.py | 56 +----- vllm/renderers/base.py | 104 +++++------ vllm/v1/worker/gpu/model_runner.py | 6 +- vllm/v1/worker/gpu_model_runner.py | 14 +- 23 files changed, 340 insertions(+), 398 deletions(-) diff --git a/tests/models/multimodal/processing/test_audio_in_video.py b/tests/models/multimodal/processing/test_audio_in_video.py index f696d7a561a7..acb102f60fb0 100644 --- a/tests/models/multimodal/processing/test_audio_in_video.py +++ b/tests/models/multimodal/processing/test_audio_in_video.py @@ -63,36 +63,32 @@ def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> Non mm_processor_cache_gb=1, ) - # Baseline: no cache, always processes from scratch. - baseline_processor = MULTIMODAL_REGISTRY.create_processor( - ctx.model_config, cache=None - ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + # Sender cache: on a cache hit returns (None, prompt_updates) for each # item, setting mm_kwargs["video"] = [None] – the exact condition that # triggered the original bug. sender_cache = MultiModalProcessorSenderCache(ctx.model_config) - cached_processor = MULTIMODAL_REGISTRY.create_processor( - ctx.model_config, cache=sender_cache - ) - video_token_id = baseline_processor.info.get_hf_config().video_token_id + video_token_id = processor.info.get_hf_config().video_token_id mm_data = create_mm_data(num_videos) hf_processor_mm_kwargs = {"use_audio_in_video": True} - def run(processor): + def run(processor, cache): return processor( [video_token_id] * num_videos, - mm_items=baseline_processor.info.parse_mm_data(mm_data), + mm_items=processor.info.parse_mm_data(mm_data), hf_processor_mm_kwargs=hf_processor_mm_kwargs, + cache=cache, )["prompt_token_ids"] - baseline_ids = run(baseline_processor) + baseline_ids = run(processor, cache=None) # First call on the sender-cache processor: cache miss. # mm_kwargs["video"] items are real tensors; use_audio_in_video is # detected normally from the item data. - first_ids = run(cached_processor) + first_ids = run(processor, cache=sender_cache) assert first_ids == baseline_ids, ( "Cache-miss call produced different prompt_token_ids than baseline.\n" f" baseline : {baseline_ids}\n" @@ -103,7 +99,7 @@ def run(processor): # MultiModalProcessorSenderCache.get_and_update_item returns (None, …), # so mm_kwargs["video"] = [None]. Before the fix, use_audio_in_video was # not detected, yielding wrong token ids. - second_ids = run(cached_processor) + second_ids = run(processor, cache=sender_cache) assert second_ids == baseline_ids, ( "Cache-hit call produced different prompt_token_ids than baseline.\n" "This is the regression introduced when use_audio_in_video detection\n" @@ -123,7 +119,7 @@ def test_use_audio_in_video_without_audio_track(model_id: str) -> None: model_id, limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1}, ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=None) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) video_token_id = processor.info.get_hf_config().video_token_id rng = np.random.RandomState(0) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index ea4305976cde..db6f37dc3d5f 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -72,6 +72,8 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict: _IGNORE_MM_KEYS = { + # Dithering causes minor divergence + "cohere_asr": {"input_features"}, # In Ultravox, the audio_features can be different depending on padding # The slight difference should not be a problem though, since # attention_mask lets us ignore the difference. @@ -286,8 +288,7 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: } ) - baseline_processor = factories.build_processor(ctx, cache=None) - cached_processor = factories.build_processor(ctx, cache=cache) + processor = factories.build_processor(ctx) rng = np.random.RandomState(0) @@ -336,41 +337,42 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: _test_processing_correctness_one( model_config, mm_data, - baseline_processor, - cached_processor, + processor, batch_idx, hit_rate, num_batches, simplify_rate, + cache=cache, ) def _test_processing_correctness_one( model_config: ModelConfig, mm_data: MultiModalDataDict, - baseline_processor: BaseMultiModalProcessor, - cached_processor: BaseMultiModalProcessor, + processor: BaseMultiModalProcessor, batch_idx: int, hit_rate: float, num_batches: int, simplify_rate: float, + cache: MultiModalProcessorOnlyCache, ): model_type = model_config.hf_config.model_type - token_prompt = get_token_prompt(baseline_processor, mm_data) - mm_items = baseline_processor.info.parse_mm_data(mm_data) + token_prompt = get_token_prompt(processor, mm_data) + mm_items = processor.info.parse_mm_data(mm_data) ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]()) - baseline_tokenized_result = baseline_processor( + baseline_tokenized_result = processor( token_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, ) - cached_tokenized_result = cached_processor( + cached_tokenized_result = processor( token_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) _assert_inputs_equal( diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index a28cc50be5fb..6430f40d8ba7 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -221,7 +221,7 @@ def test_video_cache_is_independent_of_image_kwargs( mm_processor_cache_gb=1, ) cache = MultiModalProcessorOnlyCache(ctx.model_config) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=cache) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) hf_processor = processor.info.get_hf_processor() image = PILImage.new("RGB", (48, 48), color=(128, 128, 128)) frames = np.stack([np.asarray(image)] * 2) @@ -230,15 +230,16 @@ def test_video_cache_is_independent_of_image_kwargs( {"image": image, "video": [(frames, metadata)]} ) - def process(mm_kwargs): + def process(mm_kwargs, cache): return processor( hf_processor.image_token + hf_processor.video_token, mm_items, mm_uuid_items={"video": [video_uuid]}, hf_processor_mm_kwargs=mm_kwargs, + cache=cache, ) - baseline = process({}) + baseline = process({}, cache=None) if kwargs_on_init: ctx = build_model_context( model_id, @@ -247,12 +248,12 @@ def process(mm_kwargs): mm_processor_cache_gb=1, ) cache = MultiModalProcessorOnlyCache(ctx.model_config) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=cache) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) request_kwargs = {} if kwargs_on_init else kwargs - process(request_kwargs) - cached = process(request_kwargs) + process(request_kwargs, cache=cache) + cached = process(request_kwargs, cache=cache) cache.clear_cache() - fresh = process(request_kwargs) + fresh = process(request_kwargs, cache=cache) def pixels(result, modality, field): return result["mm_kwargs"][modality][0][field].data diff --git a/tests/models/multimodal/processing/test_mllama4.py b/tests/models/multimodal/processing/test_mllama4.py index e5ba6ae24107..6cf946eff279 100644 --- a/tests/models/multimodal/processing/test_mllama4.py +++ b/tests/models/multimodal/processing/test_mllama4.py @@ -23,11 +23,9 @@ def test_profiling(model_id: str, max_model_len: int): model_config_kwargs=model_config_kwargs, limit_mm_per_prompt=mm_counts, ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - mm_inputs = MULTIMODAL_REGISTRY.get_dummy_mm_inputs( - ctx.model_config, - mm_counts=mm_counts, - ) + mm_inputs = processor.get_dummy_mm_inputs(mm_counts) hf_config = ctx.get_hf_config(Llama4Config) image_size = hf_config.vision_config.image_size diff --git a/tests/models/multimodal/processing/test_moss_audio.py b/tests/models/multimodal/processing/test_moss_audio.py index 87e7079a810b..17863b791ca0 100644 --- a/tests/models/multimodal/processing/test_moss_audio.py +++ b/tests/models/multimodal/processing/test_moss_audio.py @@ -168,7 +168,6 @@ def _build_moss_audio_processor(cache=None): MossAudioMultiModalProcessor( info, MossAudioDummyInputsBuilder(info), - cache=cache, ), ctx, ) @@ -259,21 +258,24 @@ def test_moss_audio_multimodal_processor_handles_token_and_cache_paths(): ) cache = MultiModalProcessorOnlyCache(ctx.model_config) - cached_processor, _ = _build_moss_audio_processor(cache=cache) + cached_processor, _ = _build_moss_audio_processor() cached_text_miss = cached_processor( prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) cached_text_hit = cached_processor( prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) cached_token_hit = cached_processor( token_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len) diff --git a/tests/models/multimodal/processing/test_qwen2_vl.py b/tests/models/multimodal/processing/test_qwen2_vl.py index 57b8e7fd391f..963b5959ca9d 100644 --- a/tests/models/multimodal/processing/test_qwen2_vl.py +++ b/tests/models/multimodal/processing/test_qwen2_vl.py @@ -26,7 +26,6 @@ def test_jina_vl_processing_order() -> None: processor = MULTIMODAL_REGISTRY.create_processor( ctx.model_config, tokenizer=ctx.tokenizer, - cache=cache, ) placeholder = "<|vision_start|><|image_pad|><|vision_end|>" @@ -37,6 +36,7 @@ def process(images: list[Image.Image]): return processor( placeholder * len(images), mm_items=processor.info.parse_mm_data({"image": images}), + cache=cache, ) query = process([query_image]) diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index f249a25269c4..403a01a9e2bb 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -255,7 +255,7 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: for modality, count in limit_mm_per_prompt.items() } ) - processor = factories.build_processor(ctx, cache=None) + processor = factories.build_processor(ctx) with initialize_dummy_model(model_cls, model_config) as model: for modality, _, mm_kwargs in create_batched_mm_kwargs(model_config, processor): diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index f5de7d66f4ff..8f5a47be7b94 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -245,6 +245,7 @@ def process(): prompt="\nWhat is this?", mm_items=mm_processor.info.parse_mm_data({"image": image}), hf_processor_mm_kwargs={}, + cache=cache, ) first, second = process(), process() diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py index dc49036e8352..869afe945653 100644 --- a/tests/models/multimodal/processing/transformers_backend.py +++ b/tests/models/multimodal/processing/transformers_backend.py @@ -48,4 +48,4 @@ def create_cached_processor(model_id: str, processor_cls): ) info = MultiModalProcessingInfo(ctx) cache = MultiModalProcessorOnlyCache(model_config) - return processor_cls(info, MultiModalDummyInputsBuilder(info), cache=cache), cache + return processor_cls(info, MultiModalDummyInputsBuilder(info)), cache diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index 0c68d2bac0fd..02bd5119700b 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing as mp -from concurrent.futures import ThreadPoolExecutor -from functools import partial from types import SimpleNamespace from unittest.mock import Mock @@ -37,10 +35,13 @@ PlaceholderRange, ) from vllm.multimodal.processing import PromptInsertion -from vllm.renderers.base import BaseRenderer +from vllm.renderers import renderer_from_config from vllm.utils.async_utils import make_async from vllm.utils.mem_constants import GiB_bytes, MiB_bytes +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.engine.llm_engine import LLMEngine +from ..models.utils import build_model_context from ..utils import create_new_process_for_each_test pytestmark = pytest.mark.cpu_test @@ -773,12 +774,14 @@ def test_processor_cache_shared_across_loras(): @pytest.mark.asyncio async def test_release_kv_cache_resends_mm_payload(use_async, release_error): """Release must not leave a sender hit pointing at a cleared receiver.""" - from vllm.v1.engine.async_llm import AsyncLLM - from vllm.v1.engine.llm_engine import LLMEngine - - model_config = SimpleNamespace( - get_multimodal_config=lambda: MultiModalConfig(mm_processor_cache_gb=1) + ctx = build_model_context( + "llava-hf/llava-v1.6-mistral-7b-hf", + mm_processor_kwargs=None, + limit_mm_per_prompt={"image": 1}, + mm_processor_cache_gb=1, ) + model_config = ctx.model_config + sender = MultiModalProcessorSenderCache(model_config) receiver = MultiModalReceiverCache(model_config) item = _dummy_item({"pixel_values": 16}) @@ -787,43 +790,39 @@ async def test_release_kv_cache_resends_mm_payload(use_async, release_error): assert receiver.get_and_update_item(payload, mm_hash) is item assert sender.get_and_update_item(None, mm_hash)[0] is None - renderer = SimpleNamespace(mm_processor_cache=sender, _mm_cache_stats=None) - renderer.clear_mm_cache = partial(BaseRenderer.clear_mm_cache, renderer) + renderer = renderer_from_config(VllmConfig(model_config=model_config)) + renderer._mm_processor_cache = sender def release(): if release_error: raise RuntimeError(release_error) receiver.clear_cache() - with ThreadPoolExecutor(max_workers=1) as executor: - renderer._clear_mm_cache_async = make_async( - renderer.clear_mm_cache, executor=executor - ) - renderer.clear_mm_cache_async = partial( - BaseRenderer.clear_mm_cache_async, renderer - ) - engine = SimpleNamespace( - renderer=renderer, - engine_core=SimpleNamespace( - release_kv_cache_memory=release, - release_kv_cache_memory_async=make_async(release, executor=executor), + engine = SimpleNamespace( + renderer=renderer, + engine_core=SimpleNamespace( + release_kv_cache_memory=release, + release_kv_cache_memory_async=make_async( + release, + executor=renderer._executor, ), - logger_manager=Mock(), - ) - - async def call_release(): - if use_async: - await AsyncLLM.release_kv_cache_memory(engine) - else: - LLMEngine.release_kv_cache_memory(engine) + ), + logger_manager=Mock(), + ) - if release_error: - with pytest.raises(RuntimeError, match=release_error): - await call_release() - engine.logger_manager.record_sleep_state.assert_not_called() + async def call_release(): + if use_async: + await AsyncLLM.release_kv_cache_memory(engine) else: + LLMEngine.release_kv_cache_memory(engine) + + if release_error: + with pytest.raises(RuntimeError, match=release_error): await call_release() - engine.logger_manager.record_sleep_state.assert_called_once_with(1, 0) + engine.logger_manager.record_sleep_state.assert_not_called() + else: + await call_release() + engine.logger_manager.record_sleep_state.assert_called_once_with(1, 0) payload, _ = sender.get_and_update_item((item, []), mm_hash) assert payload is item diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index de3e333886b8..0bd19a0ea0dc 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -7,7 +7,7 @@ import numpy as np import pytest -from vllm.config import ModelConfig +from vllm.config import ModelConfig, SchedulerConfig from vllm.exceptions import VLLMValidationError from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.hasher import MultiModalHasher @@ -32,6 +32,7 @@ ) from vllm.utils.collection_utils import flatten_2d_lists +from ..models.utils import build_model_context from .utils import random_image pytestmark = pytest.mark.cpu_test @@ -1361,3 +1362,37 @@ def hash_with(hf_processor_mm_kwargs): ).get_mm_hashes("test-model", "blake3")["image"][0] assert hash_with(left) != hash_with(right) + + +@pytest.mark.parametrize( + ("chunked_prefill", "max_model_len", "expected_seq_len"), + [(None, 491520, 491520), (True, 491520, 8192), (True, 128, 128), (False, 128, 128)], +) +def test_dummy_inputs_scheduler_budget( + chunked_prefill, max_model_len, expected_seq_len +): + ctx = build_model_context( + "llava-hf/llava-v1.6-mistral-7b-hf", + mm_processor_kwargs=None, + limit_mm_per_prompt={"image": 1}, + ) + ctx.model_config.max_model_len = max_model_len + + processor = MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ) + processor.apply = lambda *args, **kwargs: {"prompt_token_ids": [7]} + + kwargs = {} + if chunked_prefill is not None: + kwargs["scheduler_config"] = SchedulerConfig( + max_model_len=max_model_len, + is_encoder_decoder=False, + max_num_batched_tokens=8192, + max_num_seqs=1, + enable_chunked_prefill=chunked_prefill, + ) + + result = processor.get_dummy_mm_inputs({"image": 1}, **kwargs) + assert len(result["prompt_token_ids"]) == expected_seq_len diff --git a/tests/multimodal/test_registry.py b/tests/multimodal/test_registry.py index 6ca887401440..d7947845afe3 100644 --- a/tests/multimodal/test_registry.py +++ b/tests/multimodal/test_registry.py @@ -5,11 +5,9 @@ """ from types import SimpleNamespace -from unittest.mock import MagicMock import pytest -from vllm.config import SchedulerConfig from vllm.multimodal import MULTIMODAL_REGISTRY from ..models.utils import build_model_context @@ -17,34 +15,6 @@ pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize( - ("chunked_prefill", "max_model_len", "expected_seq_len"), - [(None, 491520, 491520), (True, 491520, 8192), (True, 128, 128), (False, 128, 128)], -) -def test_dummy_inputs_scheduler_budget( - chunked_prefill, max_model_len, expected_seq_len -): - model_config = MagicMock() - model_config.max_model_len = max_model_len - processor = MagicMock() - processor.apply.return_value = {"prompt_token_ids": [7]} - kwargs = {} - if chunked_prefill is not None: - kwargs["scheduler_config"] = SchedulerConfig( - max_model_len=max_model_len, - is_encoder_decoder=False, - max_num_batched_tokens=8192, - max_num_seqs=1, - enable_chunked_prefill=chunked_prefill, - ) - result = MULTIMODAL_REGISTRY.get_dummy_mm_inputs( - model_config, {"image": 1}, processor=processor, **kwargs - ) - get_inputs = processor.dummy_inputs.get_dummy_processor_inputs - assert get_inputs.call_args.kwargs["seq_len"] == expected_seq_len - assert len(result["prompt_token_ids"]) == expected_seq_len - - @pytest.mark.parametrize( "model_id,limit_mm_per_prompt,expected", [ diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index f6f9f7a4db6c..2627cdd98501 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -5,6 +5,8 @@ These tests exercise: - Zero-limit modalities are filtered from mm_counts passed to get_dummy_processor_inputs (e.g. --limit-mm-per-prompt image=0 ...) + - The warmup runs through the processor-only cache and clears it afterwards, + so its dummy inputs never land in the sender cache used by the serving path - MM warmup is skipped entirely when mm_processor is None - The multimodal warmup is launched as a task on the single-worker _mm_executor to overlap engine-core init (future lifecycle, join by @@ -21,10 +23,13 @@ """ from concurrent.futures import Future, ThreadPoolExecutor +from types import MethodType, SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from vllm.entrypoints.chat_utils import ChatTemplateResolutionError +from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.base import BaseRenderer from vllm.renderers.params import ChatParams @@ -34,9 +39,11 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock: render_chat is mocked to raise ChatTemplateResolutionError so the chat warmup block is skipped cleanly, keeping the test focused on MM warmup. - """ - from vllm.entrypoints.chat_utils import ChatTemplateResolutionError + The processor is a mock at its I/O boundaries (``dummy_inputs`` and + ``apply``) but keeps the real ``get_dummy_mm_inputs``, since that is where + the warmup's seq_len capping and cache plumbing actually live. + """ renderer = MagicMock() # chat warmup: make render_chat raise so we skip past it cleanly @@ -45,13 +52,17 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock: # MM processor with configurable limits mm_processor = MagicMock() mm_processor.info.allowed_mm_limits = mm_limits + mm_processor.info.ctx.model_config = renderer.model_config mm_processor.apply.return_value = {"prompt_token_ids": [1]} + mm_processor.get_dummy_mm_inputs = MethodType( + BaseMultiModalProcessor.get_dummy_mm_inputs, mm_processor + ) renderer.mm_processor = mm_processor - renderer._readonly_mm_processor = None + renderer._mm_processor_cache = MagicMock(name="sender_cache") + renderer._mm_processor_only_cache = MagicMock(name="processor_only_cache") renderer._warmup_mm_processor = BaseRenderer._warmup_mm_processor.__get__( renderer, BaseRenderer ) - renderer._clear_processor_cache = BaseRenderer._clear_processor_cache renderer.warmup_mm = BaseRenderer.warmup_mm.__get__(renderer, BaseRenderer) renderer.start_mm_warmup_in_background = ( BaseRenderer.start_mm_warmup_in_background.__get__(renderer, BaseRenderer) @@ -64,7 +75,6 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock: renderer._mm_warmup_future = None # MM warmup has not run yet; warmup_mm must actually execute on the mock. renderer._mm_warmup_done = False - renderer.clear_mm_cache = MagicMock() renderer.model_config.max_model_len = 128 renderer.config.scheduler_config.max_num_batched_tokens = 8192 renderer.config.scheduler_config.enable_chunked_prefill = True @@ -80,8 +90,7 @@ def test_zero_limit_modality_excluded_from_mm_counts(self): """A modality with limit=0 must not appear in mm_counts.""" renderer = _make_renderer_mock({"image": 1, "video": 0}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs get_inputs.assert_called_once() @@ -93,8 +102,7 @@ def test_all_zero_limits_passes_empty_mm_counts(self): """When all limits are 0, mm_counts must be empty.""" renderer = _make_renderer_mock({"image": 0, "video": 0}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs get_inputs.assert_called_once() @@ -105,8 +113,7 @@ def test_positive_limits_all_included_in_mm_counts(self): """All modalities with limit > 0 must be present in mm_counts.""" renderer = _make_renderer_mock({"image": 2, "video": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs get_inputs.assert_called_once() @@ -125,8 +132,7 @@ def test_warmup_sequence_length_is_bounded(self, max_model_len, expected_seq_len renderer = _make_renderer_mock({"image": 1, "video": 1}) renderer.model_config.max_model_len = max_model_len - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs assert get_inputs.call_args.kwargs["seq_len"] == expected_seq_len @@ -134,26 +140,44 @@ def test_warmup_sequence_length_is_bounded(self, max_model_len, expected_seq_len def test_processor_apply_called(self): renderer = _make_renderer_mock({"image": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) renderer.mm_processor.apply.assert_called_once() - def test_mm_cache_cleared_after_warmup(self): + def test_render_chat_called_with_warmup_message(self): renderer = _make_renderer_mock({"image": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) + BaseRenderer.warmup(renderer, ChatParams()) - renderer.clear_mm_cache.assert_called_once() + renderer.render_chat.assert_called_once() - def test_render_chat_called_with_warmup_message(self): + +class TestMmWarmupCacheRouting: + # The warmup runs through the processor-only cache and drops it afterwards, + # so the dummy inputs never reach the sender cache used by the serving path. + + def test_warmup_passes_processor_only_cache(self): renderer = _make_renderer_mock({"image": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): + with patch.object( + renderer.mm_processor, + "get_dummy_mm_inputs", + wraps=renderer.mm_processor.get_dummy_mm_inputs, + ) as spy: BaseRenderer.warmup(renderer, ChatParams()) - renderer.render_chat.assert_called_once() + assert spy.call_args.kwargs["cache"] is renderer._mm_processor_only_cache + assert ( + spy.call_args.kwargs["scheduler_config"] is renderer.config.scheduler_config + ) + + def test_only_cache_cleared_and_sender_cache_untouched(self): + renderer = _make_renderer_mock({"image": 1}) + + BaseRenderer.warmup(renderer, ChatParams()) + + renderer._mm_processor_only_cache.clear_cache.assert_called_once() + renderer._mm_processor_cache.clear_cache.assert_not_called() class TestMmWarmupSkippedWhenNoProcessor: @@ -168,51 +192,30 @@ def test_no_warmup_without_processor(self): renderer.model_config.get_multimodal_config.assert_not_called() -class TestReadonlyMmWarmup: - """Readonly MM processor warmup must mirror the render path behavior.""" - - def test_readonly_processor_apply_called_and_cache_cleared(self): - renderer = _make_renderer_mock({"image": 1}) - readonly_mm_processor = MagicMock() - readonly_mm_processor.info.allowed_mm_limits = {"image": 1} - renderer._readonly_mm_processor = readonly_mm_processor - - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) - - readonly_mm_processor.apply.assert_called_once() - readonly_mm_processor.cache.clear_cache.assert_called_once() - - class TestWarmupFaultIsolation: - # A failure during a multimodal processor warmup is caught so it does not - # abort the remaining warmup steps; warmup itself must not raise. + # A failure during a warmup step is caught so it does not abort the other + # step; warmup itself must not raise. def test_chat_failure_does_not_abort_mm_warmup(self): renderer = _make_renderer_mock({"image": 1}) renderer.render_chat.side_effect = RuntimeError("chat boom") - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) # must not raise + BaseRenderer.warmup(renderer, ChatParams()) # must not raise renderer.mm_processor.apply.assert_called_once() - def test_mm_failure_does_not_abort_readonly_warmup(self): + def test_mm_failure_is_swallowed_and_cache_cleared(self): renderer = _make_renderer_mock({"image": 1}) - readonly_mm_processor = MagicMock() - readonly_mm_processor.info.allowed_mm_limits = {"image": 1} - renderer._readonly_mm_processor = readonly_mm_processor - # main processor warmup blows up before apply() + # the warmup blows up before apply() renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs.side_effect = ( RuntimeError("mm boom") ) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - BaseRenderer.warmup(renderer, ChatParams()) # must not raise + BaseRenderer.warmup(renderer, ChatParams()) # must not raise - readonly_mm_processor.apply.assert_called_once() - # cache is still cleared in the failed task's finally - renderer.clear_mm_cache.assert_called_once() + renderer.mm_processor.apply.assert_not_called() + # the cache is still cleared in the failed warmup's finally + renderer._mm_processor_only_cache.clear_cache.assert_called_once() class TestBackgroundMmWarmup: @@ -234,16 +237,15 @@ def _make_bg_renderer(self, mm_limits: dict[str, int]): def test_start_submits_future_to_mm_executor(self): renderer = self._make_bg_renderer({"image": 1}) try: - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - # Spy on the real executor's submit to confirm the warmup is - # dispatched through _mm_executor (not a separate Thread). - with patch.object( - renderer._mm_executor, "submit", wraps=renderer._mm_executor.submit - ) as spy: - renderer.start_mm_warmup_in_background() - assert isinstance(renderer._mm_warmup_future, Future) - assert spy.called - renderer._join_mm_warmup() + # Spy on the real executor's submit to confirm the warmup is + # dispatched through _mm_executor (not a separate Thread). + with patch.object( + renderer._mm_executor, "submit", wraps=renderer._mm_executor.submit + ) as spy: + renderer.start_mm_warmup_in_background() + assert isinstance(renderer._mm_warmup_future, Future) + assert spy.called + renderer._join_mm_warmup() assert renderer._mm_warmup_future is None renderer.mm_processor.apply.assert_called_once() finally: @@ -252,7 +254,6 @@ def test_start_submits_future_to_mm_executor(self): def test_start_noop_for_text_only_model(self): renderer = _make_renderer_mock({}) renderer.mm_processor = None - # _readonly_mm_processor is already None renderer.start_mm_warmup_in_background() @@ -261,12 +262,11 @@ def test_start_noop_for_text_only_model(self): def test_start_is_run_at_most_once(self): renderer = self._make_bg_renderer({"image": 1}) try: - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - renderer.start_mm_warmup_in_background() - first = renderer._mm_warmup_future - renderer.start_mm_warmup_in_background() # must not spawn a second - assert renderer._mm_warmup_future is first - renderer._join_mm_warmup() + renderer.start_mm_warmup_in_background() + first = renderer._mm_warmup_future + renderer.start_mm_warmup_in_background() # must not spawn a second + assert renderer._mm_warmup_future is first + renderer._join_mm_warmup() finally: renderer._mm_executor.shutdown(wait=True) @@ -275,9 +275,8 @@ def test_warmup_joins_background_and_does_not_rerun_mm(self): # run only the chat warmup — the MM warmup must not run twice. renderer = self._make_bg_renderer({"image": 1}) try: - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - renderer.start_mm_warmup_in_background() - BaseRenderer.warmup(renderer, ChatParams()) + renderer.start_mm_warmup_in_background() + BaseRenderer.warmup(renderer, ChatParams()) # MM apply called exactly once (by the background warmup task). renderer.mm_processor.apply.assert_called_once() @@ -294,13 +293,12 @@ def test_warmup_does_not_rerun_mm_after_reset_joins_background(self): # re-run the MM warmup — the _mm_warmup_done flag survives the join. renderer = self._make_bg_renderer({"image": 1}) try: - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - renderer.start_mm_warmup_in_background() - # Simulate reset_mm_cache joining the background warmup first - # (clears _mm_warmup_future, sets _mm_warmup_done via the join). - renderer._join_mm_warmup() - assert renderer._mm_warmup_future is None # joined & cleared - BaseRenderer.warmup(renderer, ChatParams()) + renderer.start_mm_warmup_in_background() + # Simulate reset_mm_cache joining the background warmup first + # (clears _mm_warmup_future, sets _mm_warmup_done via the join). + renderer._join_mm_warmup() + assert renderer._mm_warmup_future is None # joined & cleared + BaseRenderer.warmup(renderer, ChatParams()) # MM apply called exactly once (by the background warmup task); # warmup() did not re-run warmup_mm despite the future being cleared. @@ -313,9 +311,8 @@ def test_warmup_mm_runs_at_most_once(self): # Direct repeated calls to warmup_mm run the MM warmup only once. renderer = _make_renderer_mock({"image": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - renderer.warmup_mm() - renderer.warmup_mm() # second call must be a no-op + renderer.warmup_mm() + renderer.warmup_mm() # second call must be a no-op renderer.mm_processor.apply.assert_called_once() @@ -324,10 +321,9 @@ def test_shutdown_joins_background_warmup(self): # so the mm_processor_cache is never touched concurrently. renderer = self._make_bg_renderer({"image": 1}) - with patch("vllm.multimodal.processing.TimingContext", autospec=True): - renderer.start_mm_warmup_in_background() - future = renderer._mm_warmup_future - renderer.shutdown() + renderer.start_mm_warmup_in_background() + future = renderer._mm_warmup_future + renderer.shutdown() assert renderer._mm_warmup_future is None # The background warmup was allowed to complete (apply ran) and the @@ -344,10 +340,6 @@ class TestEngineStartWarmupHook: # These tests pin the renderer plumbing contract on EngineCoreClient. def _mock_config(self): - from types import SimpleNamespace - - from vllm.v1.engine.core_client import EngineCoreClient # noqa: F401 - cfg = SimpleNamespace( parallel_config=SimpleNamespace( data_parallel_size=1, diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index 3ea0df9d163c..22cd34fc8557 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -25,7 +25,6 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.cache import BaseMultiModalProcessorCache from vllm.multimodal.inputs import ( MultiModalFieldConfig, MultiModalKwargsItem, @@ -407,21 +406,17 @@ def _build_llava_or_pixtral_hf_info( def _build_llava_or_pixtral_hf_processor( info: _I, dummy_inputs: BaseDummyInputsBuilder[_I], - *, - cache: BaseMultiModalProcessorCache | None = None, ) -> BaseMultiModalProcessor: if isinstance(info, PixtralHFProcessingInfo): return PixtralHFMultiModalProcessor( info, dummy_inputs, # type: ignore - cache=cache, ) if isinstance(info, LlavaProcessingInfo): return LlavaMultiModalProcessor( info, dummy_inputs, # type: ignore - cache=cache, ) raise NotImplementedError(type(info)) diff --git a/vllm/model_executor/models/qwen3_asr_realtime.py b/vllm/model_executor/models/qwen3_asr_realtime.py index bab689c605b3..6eb762b6281e 100644 --- a/vllm/model_executor/models/qwen3_asr_realtime.py +++ b/vllm/model_executor/models/qwen3_asr_realtime.py @@ -25,21 +25,10 @@ from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import ( - SupportsRealtime, -) -from vllm.model_executor.models.qwen3_asr import ( - Qwen3ASRDummyInputsBuilder, - Qwen3ASRForConditionalGeneration, - Qwen3ASRMultiModalProcessor, - Qwen3ASRProcessingInfo, - _get_feat_extract_output_lengths, -) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ProcessorInputs, TimingContext from vllm.multimodal.processing.processor import ( - BaseDummyInputsBuilder, MultiModalProcessingResult, PlaceholderFeaturesInfo, cached_encode, @@ -47,6 +36,15 @@ from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.processor import cached_processor_from_config +from .interfaces import SupportsRealtime +from .qwen3_asr import ( + Qwen3ASRDummyInputsBuilder, + Qwen3ASRForConditionalGeneration, + Qwen3ASRMultiModalProcessor, + Qwen3ASRProcessingInfo, + _get_feat_extract_output_lengths, +) + logger = init_logger(__name__) _PRE_ALLOCATE_BUFFER_SIZE_IN_S = 60 @@ -101,14 +99,13 @@ def flush(self) -> np.ndarray | None: class Qwen3ASRRealtimeMultiModalProcessor(Qwen3ASRMultiModalProcessor): - def __init__( + def _cached_apply_hf_processor( self, - info: _I, - dummy_inputs: BaseDummyInputsBuilder[_I], - *, - cache: BaseMultiModalProcessorCache | None = None, - ) -> None: - super().__init__(info, dummy_inputs, cache=None) + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> MultiModalProcessingResult: + # realtime can't make use of a cache yet + return self._apply_hf_processor(inputs, timing_ctx) def _maybe_apply_prompt_updates( self, diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 162fecf03870..9fdc946df990 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -27,9 +27,8 @@ VoxtralProcessingInfo, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache from vllm.multimodal.parse import MultiModalDataItems -from vllm.multimodal.processing import BaseDummyInputsBuilder +from vllm.multimodal.processing import ProcessorInputs, TimingContext from vllm.multimodal.processing.processor import ( MultiModalProcessingResult, PlaceholderFeaturesInfo, @@ -38,23 +37,19 @@ from vllm.tokenizers import cached_tokenizer_from_config from vllm.utils.torch_utils import is_torch_equal_or_newer -from .utils import ( - _flatten_embeddings, -) +from .utils import _flatten_embeddings logger = init_logger(__name__) class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor): - def __init__( + def _cached_apply_hf_processor( self, - info: _I, - dummy_inputs: BaseDummyInputsBuilder[_I], - *, - cache: BaseMultiModalProcessorCache | None = None, - ) -> None: + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> MultiModalProcessingResult: # realtime can't make use of a cache yet - super().__init__(info, dummy_inputs, cache=None) + return self._apply_hf_processor(inputs, timing_ctx) def _maybe_apply_prompt_updates( self, diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index ef3da192db21..13448b9096dc 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/vllm/multimodal/encoder_budget.py @@ -2,8 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping -from vllm.config import ModelConfig, VllmConfig +from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.multimodal.cache import BaseMultiModalProcessorCache from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.multimodal.registry import MultiModalRegistry @@ -14,26 +15,21 @@ def get_mm_max_toks_per_item( - model_config: ModelConfig, - mm_registry: MultiModalRegistry, processor: BaseMultiModalProcessor, mm_counts: Mapping[str, int], + cache: BaseMultiModalProcessorCache | None, ) -> Mapping[str, int]: """Get the maximum number of tokens per data item from each modality based on underlying model configuration. """ max_tokens_per_item = processor.info.get_mm_max_tokens_per_item( - seq_len=model_config.max_model_len, + seq_len=processor.info.ctx.model_config.max_model_len, mm_counts=mm_counts, ) if max_tokens_per_item is not None: return max_tokens_per_item - mm_inputs = mm_registry.get_dummy_mm_inputs( - model_config, - mm_counts=mm_counts, - processor=processor, - ) + mm_inputs = processor.get_dummy_mm_inputs(mm_counts, cache=cache) return { modality: sum(item.get_num_embeds() for item in placeholders) @@ -65,7 +61,7 @@ def __init__( if enable_cache else None ) - processor = mm_registry.create_processor(model_config, cache=cache) + processor = mm_registry.create_processor(model_config) self.cache = cache self.processor = processor @@ -91,9 +87,8 @@ def __init__( active_modalities = tower_modalities | embed_only_modalities all_mm_max_toks_per_item = get_mm_max_toks_per_item( - model_config, - mm_registry, processor, + cache=cache, mm_counts=dict.fromkeys(active_modalities, 1), ) @@ -194,26 +189,32 @@ def get_modality_with_max_tokens(self) -> str: def get_encoder_budget(self) -> int: return min(self.encoder_compute_budget, self.encoder_cache_size) - def reset_cache(self) -> None: - if self.cache is not None: - self.cache.clear_cache() - + def get_dummy_encoder_profile_inputs( + self, + modality: str | None = None, + max_items_per_batch: int | None = None, + ) -> list[tuple[str, MultiModalKwargsItem]]: + if self.get_encoder_budget() <= 0 or not self.mm_max_toks_per_item: + return [] + + if modality is None: + modality = self.get_modality_with_max_tokens() + if max_items_per_batch is None: + max_items_per_batch = self.mm_max_items_per_batch[modality] + + # Don't use `max_items_per_batch` here to avoid redundant computation + dummy_mm_inputs = self.processor.get_dummy_mm_inputs( + mm_counts={modality: 1}, + cache=self.cache, + ) + dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] -def get_dummy_encoder_profile_inputs( - mm_registry: MultiModalRegistry, - budget: MultiModalBudget, -) -> list[tuple[str, MultiModalKwargsItem]]: - if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item: - return [] + # We use the cache so that the item is saved to the cache, + # but not read from the cache + assert dummy_mm_item is not None, "Item should not already be cached" - modality = budget.get_modality_with_max_tokens() - max_items_per_batch = budget.mm_max_items_per_batch[modality] - dummy_mm_inputs = mm_registry.get_dummy_mm_inputs( - budget.model_config, - mm_counts={modality: 1}, - processor=budget.processor, - ) - dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] - assert dummy_mm_item is not None, "Dummy item should be generated" + return [(modality, dummy_mm_item)] * max_items_per_batch - return [(modality, dummy_mm_item)] * max_items_per_batch + def reset_cache(self) -> None: + if self.cache is not None: + self.cache.clear_cache() diff --git a/vllm/multimodal/processing/inputs.py b/vllm/multimodal/processing/inputs.py index 639045a40852..6e547c000434 100644 --- a/vllm/multimodal/processing/inputs.py +++ b/vllm/multimodal/processing/inputs.py @@ -6,6 +6,7 @@ from vllm.config.multimodal import MMHasherAlgorithm from vllm.inputs import MultiModalHashes +from ..cache import BaseMultiModalProcessorCache from ..hasher import MultiModalHasher from ..parse import MultiModalDataItems, MultiModalUUIDItems @@ -28,6 +29,8 @@ class ProcessorInputs: hf_processor_mm_kwargs: Mapping[str, object] = field(default_factory=dict) media_io_kwargs: Mapping[str, Mapping[str, object]] = field(default_factory=dict) + cache: BaseMultiModalProcessorCache | None = None + def get_mm_hashes( self, model_id: str, diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 1efe84219ee1..b8e51b205301 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -18,6 +18,7 @@ import torch from typing_extensions import TypeVar, assert_never +from vllm.config import SchedulerConfig from vllm.inputs import ( MultiModalEncDecInput, MultiModalHashes, @@ -1039,8 +1040,6 @@ def __init__( self, info: _I, dummy_inputs: "BaseDummyInputsBuilder[_I]", - *, - cache: BaseMultiModalProcessorCache | None = None, ) -> None: super().__init__() @@ -1053,16 +1052,50 @@ def __init__( self.info = info self.dummy_inputs = dummy_inputs - self.cache = cache self.data_parser = self.info.get_data_parser() + def get_dummy_mm_inputs( + self, + mm_counts: Mapping[str, int], + *, + cache: BaseMultiModalProcessorCache | None = None, + scheduler_config: "SchedulerConfig | None" = None, + ) -> MultiModalInput: + """Create dummy data for profiling the memory usage of a model.""" + model_config = self.info.ctx.model_config + seq_len = model_config.max_model_len + if scheduler_config is not None and scheduler_config.enable_chunked_prefill: + seq_len = min(seq_len, scheduler_config.max_num_batched_tokens) + + mm_config = model_config.get_multimodal_config() + + processor_inputs = self.dummy_inputs.get_dummy_processor_inputs( + seq_len=seq_len, + mm_counts=mm_counts, + mm_options=mm_config.limit_per_prompt, + ) + processor_inputs.cache = cache + + mm_inputs = self.apply( + processor_inputs, + timing_ctx=TimingContext(enabled=False), + ) + + prompt_token_ids = mm_inputs["prompt_token_ids"] + total_len = len(prompt_token_ids) + if total_len < seq_len: + prompt_token_ids.extend([0] * (seq_len - total_len)) + + return mm_inputs + def __call__( self, prompt: str | list[int], mm_items: MultiModalDataItems, mm_uuid_items: MultiModalUUIDItems | None = None, hf_processor_mm_kwargs: Mapping[str, object] | None = None, + cache: BaseMultiModalProcessorCache | None = None, ) -> MultiModalInput: if isinstance(prompt, str): tokenizer = self.info.get_tokenizer() @@ -1076,6 +1109,7 @@ def __call__( mm_items, mm_uuid_items, hf_processor_mm_kwargs=hf_processor_mm_kwargs or {}, + cache=cache, ) return self.apply(processor_inputs, TimingContext(enabled=False)) @@ -1420,7 +1454,7 @@ def _cached_apply_hf_processor( """Apply the HF processor on the full prompt text, caching the results and reusing cached results. """ - cache = self.cache + cache = inputs.cache has_passthrough_data = any( len(items.get_passthrough_data()) > 0 for items in inputs.mm_data_items.values() @@ -1827,6 +1861,7 @@ def apply( inputs.mm_uuid_items, hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs, media_io_kwargs=inputs.media_io_kwargs, + cache=inputs.cache, ) encoder_inputs = super().apply(encoder_processor_inputs, timing_ctx) diff --git a/vllm/multimodal/registry.py b/vllm/multimodal/registry.py index f5be2a0ebbf0..24a7a696b4dd 100644 --- a/vllm/multimodal/registry.py +++ b/vllm/multimodal/registry.py @@ -2,12 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading from collections import defaultdict -from collections.abc import Mapping from dataclasses import dataclass from multiprocessing.synchronize import Lock as LockType from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast -from vllm.inputs import MultiModalInput from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config @@ -29,12 +27,7 @@ ) if TYPE_CHECKING: - from vllm.config import ( - ModelConfig, - ObservabilityConfig, - SchedulerConfig, - VllmConfig, - ) + from vllm.config import ModelConfig, ObservabilityConfig, VllmConfig from vllm.model_executor.models.interfaces import SupportsMultiModal logger = init_logger(__name__) @@ -75,8 +68,6 @@ def __call__( self, info: _I, dummy_inputs: BaseDummyInputsBuilder[_I], - *, - cache: BaseMultiModalProcessorCache | None = None, ) -> BaseMultiModalProcessor[_I]: ... @@ -89,12 +80,10 @@ class _ProcessorFactories(Generic[_I]): def build_processor( self, ctx: InputProcessingContext, - *, - cache: BaseMultiModalProcessorCache | None = None, ): info = self.info(ctx) dummy_inputs_builder = self.dummy_inputs(info) - return self.processor(info, dummy_inputs_builder, cache=cache) + return self.processor(info, dummy_inputs_builder) class MultiModalRegistry: @@ -218,7 +207,6 @@ def create_processor( model_config: "ModelConfig", *, tokenizer: TokenizerLike | None = None, - cache: BaseMultiModalProcessorCache | None = None, ) -> BaseMultiModalProcessor[BaseProcessingInfo]: """Create a multi-modal processor for a specific model and tokenizer.""" if not model_config.is_multimodal_model: @@ -230,45 +218,7 @@ def create_processor( ctx = self._create_processing_ctx(model_config, tokenizer) - return factories.build_processor(ctx, cache=cache) - - def get_dummy_mm_inputs( - self, - model_config: "ModelConfig", - mm_counts: Mapping[str, int], - *, - cache: BaseMultiModalProcessorCache | None = None, - processor: BaseMultiModalProcessor | None = None, - scheduler_config: "SchedulerConfig | None" = None, - ) -> MultiModalInput: - """Create dummy data for profiling the memory usage of a model. - - The model is identified by `model_config`. - """ - seq_len = model_config.max_model_len - if scheduler_config is not None and scheduler_config.enable_chunked_prefill: - seq_len = min(seq_len, scheduler_config.max_num_batched_tokens) - - if processor is None: - processor = self.create_processor(model_config, cache=cache) - - mm_config = model_config.get_multimodal_config() - processor_inputs = processor.dummy_inputs.get_dummy_processor_inputs( - seq_len=seq_len, - mm_counts=mm_counts, - mm_options=mm_config.limit_per_prompt, - ) - mm_inputs = processor.apply( - processor_inputs, - timing_ctx=TimingContext(enabled=False), - ) - - prompt_token_ids = mm_inputs["prompt_token_ids"] - total_len = len(prompt_token_ids) - if total_len < seq_len: - prompt_token_ids.extend([0] * (seq_len - total_len)) - - return mm_inputs + return factories.build_processor(ctx) def _get_cache_type( self, diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 2cb4dd646cbb..f4916faf6bb4 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -127,8 +127,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: self._async_tokenizer_decode = make_async(self._decode, executor=self._executor) self.mm_processor: BaseMultiModalProcessor | None = None - self._readonly_mm_processor: BaseMultiModalProcessor | None = None - self._mm_cache_stats: MultiModalCacheStats | None = None + self._clear_mm_cache_async = make_async( self.clear_mm_cache, executor=self._mm_executor ) @@ -138,6 +137,16 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: self._safe_load_prompt_embeds_async = make_async( safe_load_prompt_embeds, executor=self._executor ) + + self._mm_processor_cache: BaseMultiModalProcessorCache | None = None + + # A second processor with its own processor-only cache. + # Used by the tokenize endpoint so that tokenize-only + # requests don't pollute the sender cache. + self._mm_processor_only_cache: BaseMultiModalProcessorCache | None = None + + self._mm_cache_stats: MultiModalCacheStats | None = None + if mm_registry.supports_multimodal_inputs(config.model_config): # Install the process-global GPU memory pool used to gate # frontend GPU-side multimodal decoding (no-op when the budget @@ -149,30 +158,16 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: config.parallel_config._api_process_count, ) - mm_processor_cache = mm_registry.processor_cache_from_config(config) - with set_default_torch_num_threads(): self.mm_processor = mm_registry.create_processor( config.model_config, tokenizer=self.tokenizer, - cache=mm_processor_cache, ) - if mm_processor_cache: - self._mm_cache_stats = MultiModalCacheStats() - self._resources.callback(mm_processor_cache.close) - - # A second processor with its own processor-only cache. - # Used by the tokenize endpoint so that tokenize-only - # requests don't pollute the sender cache. - ro_cache = mm_registry.processor_only_cache_from_config(config) - if ro_cache is not None: - with set_default_torch_num_threads(): - self._readonly_mm_processor = mm_registry.create_processor( - config.model_config, - tokenizer=self.tokenizer, - cache=ro_cache, - ) + self._mm_processor_cache = mm_registry.processor_cache_from_config(config) + self._mm_processor_only_cache = ( + mm_registry.processor_only_cache_from_config(config) + ) # This is used to generate internal request ID for MM processing # It has no relation to the request ID for engine core @@ -181,6 +176,12 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: config.observability_config ) + if self._mm_processor_cache: + self._mm_cache_stats = MultiModalCacheStats() + self._resources.callback(self._mm_processor_cache.close) + if self._mm_processor_only_cache: + self._resources.callback(self._mm_processor_only_cache.close) + def get_tokenizer(self) -> _T: tokenizer = self.tokenizer if tokenizer is None: @@ -198,11 +199,8 @@ def get_mm_processor(self) -> "BaseMultiModalProcessor": return self.mm_processor @property - def mm_processor_cache(self) -> "BaseMultiModalProcessorCache | None": - if self.mm_processor is None: - return None - - return self.mm_processor.cache + def mm_processor_cache(self) -> BaseMultiModalProcessorCache | None: + return self._mm_processor_cache def stat_mm_cache(self) -> MultiModalCacheStats | None: mm_cache_stats = self._mm_cache_stats @@ -214,7 +212,7 @@ def stat_mm_cache(self) -> MultiModalCacheStats | None: return mm_cache_stats def update_mm_cache_stats(self) -> None: - mm_processor_cache = self.mm_processor_cache + mm_processor_cache = self._mm_processor_cache mm_cache_stats = self._mm_cache_stats if mm_processor_cache and mm_cache_stats: @@ -222,38 +220,26 @@ def update_mm_cache_stats(self) -> None: mm_cache_stats.record(delta.total, delta.hits) def clear_mm_cache(self) -> None: - mm_processor_cache = self.mm_processor_cache + mm_processor_cache = self._mm_processor_cache if mm_processor_cache is not None: mm_processor_cache.clear_cache() if self._mm_cache_stats is not None: self._mm_cache_stats.reset = True - @staticmethod - def _clear_processor_cache( - processor: "BaseMultiModalProcessor | None", - ) -> None: - if processor is None: - return - - processor_cache = processor.cache - if processor_cache is not None: - processor_cache.clear_cache() - def _warmup_mm_processor( self, - processor: "BaseMultiModalProcessor", + processor: BaseMultiModalProcessor, + cache: BaseMultiModalProcessorCache | None, *, log_prefix: str, ) -> None: - model_config = self.model_config mm_limits = {k: v for k, v in processor.info.allowed_mm_limits.items() if v > 0} start_time = time.perf_counter() - _ = mm_registry.get_dummy_mm_inputs( - model_config, - mm_counts=dict.fromkeys(mm_limits, 1), - processor=processor, + _ = processor.get_dummy_mm_inputs( + dict.fromkeys(mm_limits, 1), + cache=cache, scheduler_config=self.config.scheduler_config, ) @@ -274,28 +260,20 @@ def warmup_mm(self) -> None: # prevent MM processor hangs with set_default_torch_num_threads(1): if self.mm_processor: + mm_processor_only_cache = self._mm_processor_only_cache + try: logger.debug("Warming up multi-modal processing...") self._warmup_mm_processor( self.mm_processor, + mm_processor_only_cache, log_prefix="Multi-modal", ) except Exception: logger.warning("Multi-modal warmup failed") finally: - self.clear_mm_cache() - - if self._readonly_mm_processor is not None: - try: - logger.debug("Warming up readonly multi-modal processing...") - self._warmup_mm_processor( - self._readonly_mm_processor, - log_prefix="Readonly multi-modal", - ) - except Exception: - logger.warning("Readonly multi-modal warmup failed") - finally: - self._clear_processor_cache(self._readonly_mm_processor) + if mm_processor_only_cache: + mm_processor_only_cache.clear_cache() finally: self._mm_warmup_done = True @@ -323,7 +301,7 @@ def start_mm_warmup_in_background(self) -> None: # which is not designed for concurrent access (its get_and_update must # stay ordered with the engine-core cache) — is never touched # concurrently by the warmup and the serving path. - if self.mm_processor is None and self._readonly_mm_processor is None: + if self.mm_processor is None: return # text-only model: nothing to warm up if self._mm_warmup_future is not None: return # already launched @@ -853,10 +831,7 @@ def _process_multimodal( *, skip_mm_cache: bool = False, ) -> "MultiModalInput": - if skip_mm_cache and self._readonly_mm_processor is not None: - mm_processor = self._readonly_mm_processor - else: - mm_processor = self.get_mm_processor() + mm_processor = self.get_mm_processor() mm_req_id = f"renderer{self.api_process_rank}-mm-{self._mm_req_counter.inc(1)}" @@ -873,6 +848,11 @@ def _process_multimodal( mm_uuid_items, hf_processor_mm_kwargs=mm_processor_kwargs or {}, media_io_kwargs=media_io_kwargs or {}, + cache=( + self._mm_processor_only_cache + if skip_mm_cache + else self._mm_processor_cache + ), ) mm_timing_ctx = self._mm_timing_registry.get(mm_req_id) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a05658e7ce56..4dd8865cbd22 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -57,7 +57,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import ( MultiModalBudget, - get_dummy_encoder_profile_inputs, ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask @@ -935,10 +934,7 @@ def profile_run(self) -> None: self.mm_registry, enable_cache=False, ) - dummy_mm_inputs = get_dummy_encoder_profile_inputs( - self.mm_registry, - mm_budget, - ) + dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs() self.model_state.encoder_runner.profile_encoder_cache( dummy_mm_inputs, mm_budget ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index cedf95e319a2..7d933917bc53 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5837,21 +5837,15 @@ def _get_mm_dummy_batch( assert self.mm_budget is not None # Don't use `max_items_per_batch` here to avoid redundant computation - dummy_mm_inputs = self.mm_registry.get_dummy_mm_inputs( - self.model_config, - mm_counts={modality: 1}, - cache=self.mm_budget.cache, + dummy_mm_inputs = self.mm_budget.get_dummy_encoder_profile_inputs( + modality, + max_items_per_batch, ) - dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] - - # We use the cache so that the item is saved to the cache, - # but not read from the cache - assert dummy_mm_item is not None, "Item should not already be cached" return next( mm_kwargs_batch for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( - [(modality, dummy_mm_item)] * max_items_per_batch, + dummy_mm_inputs, device=self.device, pin_memory=PIN_MEMORY, ) From 5302d1fe40e6969b854e0ab288d34825653767c6 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Sat, 19 Sep 2026 10:46:49 -0700 Subject: [PATCH 516/774] [Model Runner V2] Support custom logits processors (#56497) Signed-off-by: David Cheung Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- docs/features/custom_logitsprocs.md | 343 ++++-------------- tests/test_config.py | 9 + tests/test_sampling_params.py | 8 + .../test_input_processor_trace_replay.py | 1 + .../logits_processors/test_custom_offline.py | 9 + .../logits_processors/test_custom_online.py | 3 + tests/v1/worker/test_gpu_bad_words.py | 55 ++- tests/v1/worker/test_gpu_batch_shard.py | 6 + tests/v1/worker/test_gpu_logit_bias.py | 35 +- tests/v1/worker/test_gpu_logits_processors.py | 114 ++++++ .../worker/test_gpu_model_runner_v2_eplb.py | 11 +- .../test_gpu_rejection_sampler_chunking.py | 3 + tests/v1/worker/test_gpu_sampler_flags.py | 45 ++- tests/v1/worker/test_gpu_thinking_budget.py | 24 +- .../v1/worker/test_logits_processor_loader.py | 228 ++++++++++++ vllm/config/vllm.py | 12 - vllm/model_executor/models/diffusion_gemma.py | 2 +- vllm/sampling_params.py | 8 - vllm/v1/engine/input_processor.py | 29 +- vllm/v1/watermarking/gpu_sampler.py | 6 +- vllm/v1/worker/gpu/model_runner.py | 12 +- vllm/v1/worker/gpu/sample/bad_words.py | 55 ++- vllm/v1/worker/gpu/sample/batch_shard.py | 4 + vllm/v1/worker/gpu/sample/logit_bias.py | 71 ++-- .../gpu/sample/logits_processor/__init__.py | 26 ++ .../gpu/sample/logits_processor/interface.py | 158 ++++++++ .../gpu/sample/logits_processor/loader.py | 211 +++++++++++ vllm/v1/worker/gpu/sample/penalties.py | 54 ++- vllm/v1/worker/gpu/sample/sampler.py | 123 +++---- vllm/v1/worker/gpu/sample/states.py | 19 +- vllm/v1/worker/gpu/sample/thinking_budget.py | 28 +- .../gpu/spec_decode/rejection_sampler.py | 3 + 32 files changed, 1210 insertions(+), 505 deletions(-) create mode 100644 tests/v1/worker/test_gpu_logits_processors.py create mode 100644 tests/v1/worker/test_logits_processor_loader.py create mode 100644 vllm/v1/worker/gpu/sample/logits_processor/__init__.py create mode 100644 vllm/v1/worker/gpu/sample/logits_processor/interface.py create mode 100644 vllm/v1/worker/gpu/sample/logits_processor/loader.py diff --git a/docs/features/custom_logitsprocs.md b/docs/features/custom_logitsprocs.md index 232f4363efee..8c1c25939928 100644 --- a/docs/features/custom_logitsprocs.md +++ b/docs/features/custom_logitsprocs.md @@ -8,84 +8,40 @@ A "custom" logits processor is written by a user of vLLM and is loaded into vLLM This document shows how to write, load and use a custom logits processor. +!!! note + This document covers the Model Runner V2 (MRV2) interface. The legacy V1 + model runner (`VLLM_USE_V2_MODEL_RUNNER=0`) uses a different interface, + `vllm.v1.sample.logits_processor.LogitsProcessor`, which is not covered + here. + ## Logits Processors Background A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior. -In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_requests) x (vocab_size)` tensor of raw logits output by the model. For all requests which enable the logits processor, the logits processor applies a transformation to the corresponding row of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then passed to softmax. +In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_logits_rows) x (vocab_size)` tensor of raw logits output by the model. Note that a logits row is not a request: rows are reordered every step, and under speculative decoding a request owns one row per draft token. The logits processor applies a transformation to the rows of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then used for sampling. ## Creating a Custom Logits Processor -Custom logits processors must subclass `vllm.v1.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods: +Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods: -* `validate_params(cls, sampling_params: SamplingParams)`: - * Raise `ValueError` if `SamplingParams` has invalid arguments (especially custom arguments) used by logits processor. - * When request is sent to entrypoint, `validate_params()` will validate `SamplingParams` and refuse request with invalid arguments. - * **Note:** it's important to implement `validate_params()` to prevent invalid parameters for custom logits processor. Otherwise requests with invalid parameters can cause unexpected behaviour in custom logits processor. - -* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)` +* `__init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState)`: * `vllm_config`: engine configuration data structure - * `device`: hardware accelerator device info - * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation - -* `apply(self, logits: torch.Tensor) -> torch.Tensor`: - * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`) - * Apply logits processor transformation at batch granularity - * Return a transformed `(num_requests) x (vocab_size)` logits tensor - * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient - -* `is_argmax_invariant(self) -> bool`: - * Return `True` if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), `False` if the logits processor may modify argmax - * `is_argmax_invariant()` is evaluated once at startup; if `True`, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling - -* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`: - * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step - * Use the `BatchUpdate` members to update logits processor internal state - * **Note:** batch update data structure may be `None`, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated `output_token_ids` lists that it could have retained when they were added. - -### How the vLLM engine builds the `BatchUpdate` data structure - -!!! important - Some logits processors design changes are still in progress. We expect - that in the future you will not need to account for batch state changes - when implementing a logits processor, and the information in this section - will become irrelevant. - -Logits processor `update_state()` implementations should assume the following model for how the model runner updates persistent batch state (expressed here in terms of the `BatchUpdate` abstraction): - -1. Identify indices of requests which finished in the current engine step - -2. Identify new requests introduced in the current step - -3. Use Add operations to replace as many finished requests with new requests, in order of increasing index of the replaced request starting with the lowest index - -4. Based on the relative number of new and finished requests: - - 1. If the numbers of new and finished requests are the same, proceed to next step + * `req_states`: a narrow, read-only view of the persistent batch, exposing the on-device token history (`all_token_ids`, `prompt_len`, `prefill_len`, `total_len`) plus `device`, `max_num_reqs` and `vocab_size` - 2. *If there are more new requests than finished requests:* apply Add operations to extend the batch with the remaining new requests which did not replace finished requests. Assign consecutive indices to these new requests, starting with `current_max_batch_index + 1` +* `add_request(self, req_idx, sampling_params) -> bool`: + * Initialize per-slot state for a request entering the batch. Slots are recycled through a free list, so per-slot state must be fully (re)initialized here; there is no removal hook, since freed slots are never read + * Return whether this processor modifies logits for the request. The sampler skips the logits-processing pipeline for batches in which no request needs it - 3. *If there are fewer new requests than finished requests:* +* `apply_staged_writes(self) -> None` (optional): + * Flush host-side writes staged by `add_request()` to the device; called once per step before the forward pass - * Apply Remove operations to finished requests which were not replaced with new requests. These removed request indices will necessarily be greater than the greatest index of the finished requests which were replaced in the previous step. The Removes may leave the batch in a non-contiguous state +* `validate_params(cls, sampling_params) -> None` (optional classmethod): + * Raise `ValueError` for invalid per-request arguments (especially custom arguments); runs at request admission, so invalid arguments fail the request instead of reaching the sampler - * **"Condense" the batch to be contiguous:** starting with the lowest-index empty slot (which was caused by a Remove), apply a Unidirectional Move from the current highest non-empty slot in the batch to fill the empty slot. Proceed with additional Unidirectional Move operations in order of increasing empty slot destination index and decreasing non-empty slot source index until the batch is contiguous - - * **Shrink the batch:** a side effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots - -5. Reorder the batch for improved efficiency. Depending on the attention backend implementation and the current characteristics of the batch, zero or more Swap Move operations may be applied to reorder the batch - -Notes: - -* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves - -* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations - * Example: if a request is Added at index 5 and then swapped with index 3, the Add operation in `BatchUpdate.added` will be associated with index 5 not 3 - * In other words Move operations can be assumed to be applied after Adds and Removes - -* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved` - -* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None` +* `apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor`: + * Consume a `(num_logits_rows) x (vocab_size)` logits tensor and the step's batch layout (`ctx`: row-to-slot mappings, `input_ids`, positions) + * Apply logits processor transformation at batch granularity + * Return a transformed logits tensor. You can modify the input logits in-place or out-of-place; in-place is more memory-efficient ### Passing Custom Argument to a Custom Logits Processor @@ -93,7 +49,7 @@ Unlike built-in logits processors, custom logits processors may require configur ### Example Custom Logits Processor Implementation -The contrived example below implements a custom logits processor which consumes a `(num\_requests) \times (vocab\_size)` logits tensor and masks out all tokens except for one (`target_token`) with `float(-inf)`. The logits processor is disabled for any request that does not specify `target_token`. To determine whether the logits processor is enabled and which token to leave unmasked, the logits processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request: +The contrived example below implements a custom logits processor which masks out all tokens except for one (`target_token`) for the requests that enable it. The processor is disabled for any request that does not specify `target_token`; `add_request()` reports this via its return value, so the sampler can skip the logits-processing pipeline entirely when no request enables the processor. To determine whether the logits processor is enabled and which token to leave unmasked, the processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request: ??? code "Example custom logits processor definition" @@ -101,193 +57,55 @@ The contrived example below implements a custom logits processor which consumes import torch from vllm.config import VllmConfig from vllm.sampling_params import SamplingParams - from vllm.v1.sample.logits_processor import (BatchUpdate, - LogitsProcessor, - MoveDirectionality) + from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsProcRequestState, + LogitsContext, + LogitsProcessor, + ) - class DummyLogitsProcessor(LogitsProcessor): - """Fake logit processor to support unit testing and examples""" - @classmethod - def validate_params(cls, params: SamplingParams): - target_token: int | None = params.extra_args and params.extra_args.get( - "target_token" - ) - if target_token is not None and not isinstance(target_token, int): - raise ValueError(f"target_token value {target_token} is not int") - - def __init__(self, vllm_config: "VllmConfig", device: torch.device, - is_pin_memory: bool): - self.req_info: dict[int, int] = {} - - def is_argmax_invariant(self) -> bool: - """Never impacts greedy sampling""" - return False - - def update_state(self, batch_update: BatchUpdate | None): - if not batch_update: - return - - # Process added requests. - for index, params, _, _ in batch_update.added: - assert params is not None - self.validate_params(params) - if params.extra_args and (target_token := - params.extra_args.get("target_token")): - self.req_info[index] = target_token - else: - self.req_info.pop(index, None) - - if self.req_info: - # Process removed requests. - for index in batch_update.removed: - self.req_info.pop(index, None) - - # Process moved requests, unidirectional move (a->b) and swap - # (a<->b) - for adx, bdx, direct in batch_update.moved: - a_val = self.req_info.pop(adx, None) - b_val = self.req_info.pop(bdx, None) - if a_val is not None: - self.req_info[bdx] = a_val - if direct == MoveDirectionality.SWAP and b_val is not None: - self.req_info[adx] = b_val - - def apply(self, logits: torch.Tensor) -> torch.Tensor: - if not self.req_info: - return logits + class TargetTokenLogitsProcessor(LogitsProcessor): + """Masks out all tokens except `target_token` (a per-request custom + argument); requests without it are left alone.""" - # Save target values before modification - cols = torch.tensor( - list(self.req_info.values()), dtype=torch.long, device=logits.device + def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): + # Per-slot target; -1 means disabled. Staged on the host and + # flushed to the device once per step in apply_staged_writes(). + self.target_token = torch.full( + (req_states.max_num_reqs,), -1, dtype=torch.int64 ) - rows = torch.tensor( - list(self.req_info.keys()), dtype=torch.long, device=logits.device + self.target_token_dev = torch.full( + (req_states.max_num_reqs,), -1, dtype=torch.int64, device=req_states.device ) - values_to_keep = logits[rows, cols].clone() - - # Mask all but target tokens - logits[rows] = float('-inf') - logits[rows, cols] = values_to_keep - - return logits - - ``` -In the rest of this document, we will use `DummyLogitsProcessor` as an example of a custom logits processor. + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + target = (sampling_params.extra_args or {}).get("target_token") + self.target_token[req_idx] = target if target is not None else -1 + return target is not None -The `DummyLogitsProcessor.update_state()` implementation maintains a "sparse" representation of the batched requests in the `self.req_info` dictionary: only those requests which specify a `target_token` value have a key in the dictionary. `update_state()` adjusts the stored request indices and `target_token` values (keys and values respectively in `self.req_info`) in response to Add, Remove and Move operations against the persistent batch. + def apply_staged_writes(self) -> None: + self.target_token_dev.copy_(self.target_token, non_blocking=True) -### Wrapping an Existing Request-Level Logits Processor - -Although the vLLM engine applies logits processors at batch granularity, some users may want to use vLLM with a "request-level" logits processor implementation - an implementation which operates on individual requests. This will be especially true if your logits processor was developed for vLLM version 0, which required it to be a `Callable` (as described [here][vllm.logits_process]) conforming to the following type annotation: - -``` python -RequestLogitsProcessor = Union[ - - # (output token ids, logits tensor) -> logits tensor - Callable[[list[int], Tensor], Tensor], - - # (prompt token ids, output token ids, logits tensor) -> logits tensor - Callable[[list[int], list[int], Tensor], Tensor], -] -``` - -While request-level logits processors are explicitly *not* supported in the vLLM engine, vLLM *does* provide a convenient process to wrap an existing `Callable` request-level logits processor and create a batch-level logits processor that is compatible with vLLM. The `Callable` must conform to the type annotation above; if your request-level logits processor has a different interface, then in order to wrap it, you may need to modify it or implement an additional wrapper layer to comply with the interface specification above. - -You can wrap the request-level logits processor by subclassing `AdapterLogitsProcessor` as shown in the example below (in this example, `DummyPerReqLogitsProcessor` is a stand-in for your request-level logits processor which needs to be wrapped.): - -* Override `AdapterLogitsProcessor.validate_params(cls,params)` to validate request's sampling parameters. - -* Override `AdapterLogitsProcessor.is_argmax_invariant(self)` to accurately reflect whether your request-level logits processor may impact which token has the highest-value logit. - -* Override `AdapterLogitsProcessor.new_req_logits_processor(self,params)` to create a new request-level logits processor instance from a `SamplingParams` instance: - -??? code "Example of Wrapping a Request-Level Logits Processor" - - ``` python - ... - - from vllm.v1.sample.logits_processor import ( - AdapterLogitsProcessor, # Wrapper base-class - RequestLogitsProcessor, # Request-level logitsproc type annotation - ) - - ... - - # Stand-in for your request-level logits processor: - class DummyPerReqLogitsProcessor: - """The request-level logits processor masks out all logits except the - token id identified by `target_token`""" - - def __init__(self, target_token: int) -> None: - """Specify `target_token`""" - self.target_token = target_token - - def __call__( - self, - output_ids: list[int], - logits: torch.Tensor, - ) -> torch.Tensor: - val_to_keep = logits[self.target_token].item() - logits[:] = float("-inf") - logits[self.target_token] = val_to_keep + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + # ctx.expanded_idx_mapping maps each logits row to its request slot. + cols = self.target_token_dev[ctx.expanded_idx_mapping.long()] + rows = torch.nonzero(cols >= 0).squeeze(1) + if rows.numel() == 0: + return logits + kept = logits[rows, cols[rows]].clone() + logits[rows] = float("-inf") + logits[rows, cols[rows]] = kept return logits - - ... - - # Example of wrapping the request-level logits processor: - class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): - """Example of wrapping a fake request-level logit processor to create a - batch-level logits processor""" - - @classmethod - def validate_params(cls, params: SamplingParams): - target_token: Any | None = params.extra_args and params.extra_args.get( - "target_token" - ) - if target_token is not None and not isinstance(target_token, int): - raise ValueError( - f"target_token value {target_token} is not int" - ) - - def is_argmax_invariant(self) -> bool: - return False - - def new_req_logits_processor( - self, - params: SamplingParams, - ) -> Optional[RequestLogitsProcessor]: - """This method returns a new request-level logits processor, customized - to the `target_token` value associated with a particular request. - - Returns None if the logits processor should not be applied to the - particular request. To use the logits processor the request must have - a "target_token" custom argument with an integer value. - - Args: - params: per-request sampling params - - Returns: - `Callable` request logits processor, or None - """ - target_token: Any | None = params.extra_args and params.extra_args.get( - "target_token" - ) - if target_token is None: - return None - return DummyPerReqLogitsProcessor(target_token) ``` -!!! note - Your `new_req_logits_processor()` override can return `None` to signal that the wrapped logits processor should not be applied to the request in question. - -Once you have created a custom subclass (like `WrappedPerReqLogitsProcessor`) which wraps your request level logits processor, you can pass the custom subclass to vLLM via any of the methods described in the following section. +Per-request state is keyed by the request slot index, and slots are recycled through a free list. The example keeps a per-slot state where `-1` means disabled; note that `add_request()` overwrites the slot's entry unconditionally, which is what keeps recycled slots from leaking state between requests. ## Ways to Load Your Custom Logits Processor in vLLM Logits processors are loaded at initialization. Critically, the set of loaded logits processors cannot be modified after the vLLM engine finishes loading, and new logits processors cannot be loaded on-demand for individual requests. +Loaded classes are validated against the Model Runner V2 interface; a class that does not subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` is rejected with a clear error. + This section details different ways of making your logits processor visible to vLLM and triggering vLLM to load your logits processor. ### Method 1: Pass the Custom Logits Processor Fully-Qualified Class Name (FQCN) to vLLM at Initialization Time @@ -314,7 +132,7 @@ See examples below: # Pass in FQCN llm = LLM( model="facebook/opt-125m", - logits_processors=["your.module.path:DummyLogitsProcessor"], + logits_processors=["your.module.path:TargetTokenLogitsProcessor"], ) ``` @@ -323,14 +141,14 @@ See examples below: ``` python # Pass in FQCN engine_args = AsyncEngineArgs(model="facebook/opt-125m", - logits_processors=["your.module.path:DummyLogitsProcessor"]) + logits_processors=["your.module.path:TargetTokenLogitsProcessor"]) async_llm = AsyncLLM.from_engine_args(engine_args) ``` ??? code "Passing custom logits processor FQCN to vLLM server via CLI" ```bash - vllm serve facebook/opt-125m --logits_processors your.module.path:DummyLogitsProcessor + vllm serve facebook/opt-125m --logits_processors your.module.path:TargetTokenLogitsProcessor ``` ### Method 2: Automatically Detect Custom Logits Processors Installed in Your Python Environment As Entry Points @@ -345,7 +163,7 @@ Suppose that you have developed a Python package that holds your custom logits p ``` toml [project.entry-points."vllm.logits_processors"] - dummy_logits_processor = "your.module.path:DummyLogitsProcessor" + target_token_logits_processor = "your.module.path:TargetTokenLogitsProcessor" ``` Once your package is installed, your custom logits processor will be loaded automatically whenever vLLM is initialized. You do *not* need to pass the custom logits processor to the `LLM` or `AsyncLLM` constructors or to the vLLM server explicitly at initialization time if your logits processor is exposed as an entry point. @@ -361,26 +179,26 @@ You can pass one or more custom logits processor class objects to the `LLM` and ``` python # Import custom logits processor - from some.module import DummyLogitsProcessor + from some.module import TargetTokenLogitsProcessor # ...or... # Define custom logits processor locally - from vllm.v1.sample.logits_processor import LogitsProcessor + from vllm.v1.worker.gpu.sample.logits_processor import LogitsProcessor - class DummyLogitsProcessor(LogitsProcessor): - # See DummyLogitsProcessor implementation above + class TargetTokenLogitsProcessor(LogitsProcessor): + # See TargetTokenLogitsProcessor implementation above ... # Pass class object to LLM constructor llm = LLM( model="facebook/opt-125m", - logits_processors=[DummyLogitsProcessor], + logits_processors=[TargetTokenLogitsProcessor], ) # Pass class object to AsyncLLM constructor engine_args = AsyncEngineArgs(model="facebook/opt-125m", - logits_processors=[DummyLogitsProcessor]) + logits_processors=[TargetTokenLogitsProcessor]) async_llm = AsyncLLM.from_engine_args(engine_args) ``` @@ -388,7 +206,7 @@ You can pass one or more custom logits processor class objects to the `LLM` and The design of the custom logits processor determines whether the logits processor must be enabled/disabled for a given request, and what arguments must be provided to configure the logits processor. -The examples below show how a user would pass a custom argument (`target_token`) to `DummyLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior. +The examples below show how a user would pass a custom argument (`target_token`) to `TargetTokenLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior. ??? code "vLLM REST API: configure custom logits processor for a request" @@ -419,7 +237,7 @@ The examples below show how a user would pass a custom argument (`target_token`) ??? code "Offline: configure custom logits processor for an `LLM` request" ``` python - outputs_logitproc = llm.generate("your prompt", + outputs_logitproc = llm.generate("your prompt", SamplingParams(..., extra_args={"target_token": 67})) ``` @@ -438,31 +256,14 @@ The examples below show how a user would pass a custom argument (`target_token`) ## Best Practices for Writing Custom Logits Processors -Once vLLM loads a logits processor during initialization, then vLLM will invoke `update_state()` and `apply()` against that logits processor in every engine step. Both methods operate on all requests which currently reside in the vLLM persistent batch. Thus, it is important to implement these methods efficiently. - -* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity - * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()` - * However, if you think that a logits processor may be used infrequently, it may be appropriate to use a "sparse" representation of request state i.e. the class can represent request configuration using a dictionary which only stores metadata about requests that enable the logits processor - * **Note:** wrapped request-level logits processors do not need to implement `apply()` and `update_state()`; the default `AdapterLogitsProcessor.update_state()` implementation maintains a sparse representation of request state, wherein requests for which `new_req_logits_processor()` returns `None` are not represented in the base-class state dictionary. The default implementation of `AdapterLogitsProcessor.apply()` applies the request-level logits processor to each row of input logits sequentially and assembles the output logits tensor. If the performance of this `AdapterLogitsProcessor` default implementation is insufficient, then avoid wrapping your request-level logits processor and instead re-implement it as a `LogitsProcessor` subclass with optimized `apply()` and `update_state()` implementations that operate at batch granularity +* Write an efficient `apply()` implementation in light of the fact that logits processors operate at batch granularity. For example, you may be able to use efficient vectorized operations to implement `apply()`, and to stage per-request state into contiguous tensors in `add_request()` instead of per-row Python loops * It is up to the logits processor author to determine: - 1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `update_state()` override determines how `SamplingParams` fields are mapped into logits processor state - - * **Note:** for wrapped request-level logits processors, `new_req_logits_processor()` determines how `SamplingParams` fields are used to initialize a request-level logits processor instance. - - 2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Try to save compute and memory for requests which disable the logits processor - - * **Note:** for wrapped per-request logits processors, the default `AdapterLogitsProcessor.update_state()` implementation ensures that the request-level logits processor is disabled when `new_req_logits_processor()` returns `None` for that request - - 3. **The conditions under which the logits processor is short-circuited at the batch level.** Even if you have defined a way to disable the custom logits processor at the request level, it may be difficult to translate this into compute savings i.e. if your `update_state()` and `apply()` implementations use efficient vectorized implementations that operate on the whole persistent batch in a single command. For example, you cannot skip an entire vectorized operation in `apply()` just because one request disabled the logits processor. To save compute in the edge-case where no running requests utilize the custom logits processor, we recommend designing `apply()` to return the unmodified input tensor if all requests have the logits processor disabled. Similarly, consider whether steps can be skipped in `update_state()` if no requests enable the logits processor - - * Additionally, an easy way to save compute in `update_state()` is to exit early when the `batch_update` is `None` - - * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class implements the above optimizations by default + 1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `add_request()` override determines how `SamplingParams` fields are mapped into logits processor state -* Ensure that the logits processor `update_state` method discards information about finished requests (i.e. requests which are replaced by an Add or which are subject to a Remove) + 2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Return `False` from `add_request()` for such requests so that the sampler can skip the logits-processing pipeline for batches in which no request needs it - * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class handles this by default +* Stage host-side writes in `add_request()` and flush them once per step in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows -* `is_argmax_invariant()` can be hard-coded to `True` or `False` if the logits processor has consistent behavior. However, the argmax invariance may also be determined programmatically (i.e. if your logits processor is user-customizable in some way that impacts whether the logits processor is argmax invariant). For this reason, `is_argmax_invariant()` is not a class method +* Since there is no removal hook, `add_request()` must fully (re)initialize the slot's state; do not rely on state left behind by the slot's previous occupant diff --git a/tests/test_config.py b/tests/test_config.py index a07929abfcbf..9da6eced9429 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -652,6 +652,15 @@ def test_v2_model_runner_supports_extract_hidden_states(): assert config._get_v2_model_runner_unsupported_features() == [] +def test_v2_model_runner_supports_custom_logits_processors(): + config = VllmConfig() + config.model_config = cast( + ModelConfig, SimpleNamespace(logits_processors=["a.b:C"]) + ) + + assert config._get_v2_model_runner_unsupported_features() == [] + + def test_dflash2_draft_forces_v2_model_runner(): """A DFlash2 draft must reach the V2 speculator, the only one that runs its candidate selector; on V1 it would draft as DFlash1 without raising.""" diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index cd23ea993816..3c909f4969df 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -51,6 +51,14 @@ def test_non_diffusion_models_unaffected(): params.verify(MockModelConfig(), None, None, None) +def test_verify_leaves_logits_processors_to_admission(): + """verify() is runner-agnostic; LP validation lives in the admission + layer, so an unimportable FQCN must not fail verify().""" + SamplingParams().verify( + MockModelConfig(logits_processors=["no.such:Cls"]), None, None, None + ) + + @pytest.mark.parametrize("value", [-(2**63) - 1, 2**64]) def test_extra_args_rejects_nested_integer_overflow(value): """Reject extension values before they reach the engine transport.""" diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py index 894293bc3db1..24302f6de616 100644 --- a/tests/v1/engine/test_input_processor_trace_replay.py +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -85,6 +85,7 @@ def _validate(enable_trace_replay: bool) -> None: speculative_config=None, structured_outputs_config=None, tokenizer=None, + validate_logits_processors_params=lambda params: None, ) params = SamplingParams(trace_decode_token_ids=[1, 2, 3]) with patch.object(SamplingParams, "verify"): diff --git a/tests/v1/logits_processors/test_custom_offline.py b/tests/v1/logits_processors/test_custom_offline.py index 0d0b9531dcf3..f04f5cfa47c8 100644 --- a/tests/v1/logits_processors/test_custom_offline.py +++ b/tests/v1/logits_processors/test_custom_offline.py @@ -162,6 +162,9 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource """ # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + # These tests exercise the V1-interface logits processor; Model Runner V2 + # rejects V1-interface processors at load time by design. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") set_random_seed(40) # Choose LLM args based on logitproc source @@ -217,6 +220,9 @@ def test_custom_logitsprocs_req(monkeypatch): """ # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + # These tests exercise the V1-interface logits processor; Model Runner V2 + # rejects V1-interface processors at load time by design. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") set_random_seed(40) _run_test( {"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True @@ -261,6 +267,9 @@ def test_rejects_custom_logitsprocs( """ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + # These tests exercise V1-interface rejection and V1 runner internals; + # Model Runner V2 rejects V1-interface processors at load time by design. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") set_random_seed(40) test_params: dict[str, dict[str, Any]] = { diff --git a/tests/v1/logits_processors/test_custom_online.py b/tests/v1/logits_processors/test_custom_online.py index 4a7ad8cc0069..b117875625ab 100644 --- a/tests/v1/logits_processors/test_custom_online.py +++ b/tests/v1/logits_processors/test_custom_online.py @@ -79,6 +79,9 @@ def server(default_server_args, request, monkeypatch): """ # Test that logitproc info is passed to workers monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + # These tests exercise the V1-interface logits processor; Model Runner V2 + # rejects V1-interface processors at load time by design. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") if request.param: # Launch server, append FQCN argument, inject dummy logitproc module diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py index 150c4ed00711..3cce549536a4 100644 --- a/tests/v1/worker/test_gpu_bad_words.py +++ b/tests/v1/worker/test_gpu_bad_words.py @@ -14,6 +14,10 @@ from vllm.sampling_params import SamplingParams from vllm.v1.worker.gpu.sample.bad_words import BadWordsState +from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsContext, + LogitsProcRequestState, +) from vllm.v1.worker.gpu.states import RequestState DEVICE = torch.device("cuda") @@ -47,7 +51,7 @@ def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, in req_states.apply_staged_writes() req_idx = req_states.req_id_to_index["req"] - state = BadWordsState(req_states) + state = BadWordsState(None, LogitsProcRequestState.from_request_state(req_states)) state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids)) state.apply_staged_writes() return state, req_idx @@ -61,16 +65,55 @@ def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor: expanded_idx_mapping = torch.tensor( [req_idx] * num_logits, dtype=torch.int32, device=DEVICE ) - state.apply_bad_words( + state.apply( logits, - expanded_idx_mapping, - idx_mapping_np, - torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE), - torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE), + LogitsContext( + expanded_idx_mapping=expanded_idx_mapping, + idx_mapping=torch.tensor([req_idx], dtype=torch.int32, device=DEVICE), + idx_mapping_np=idx_mapping_np, + expanded_local_pos=torch.tensor( + LOCAL_POS, dtype=torch.int32, device=DEVICE + ), + input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE), + pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE), + seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64), + ), ) return logits.cpu() +def test_v2_bad_words_slot_reuse_clears_state(): + """A slot recycled from a bad-words request must not keep masking. + + There is no removal hook by design; safety comes from add_request() + overwriting the slot's state unconditionally. + """ + state, req_idx = _make_state([[10, 11, 60]]) + state.add_request(req_idx, SamplingParams()) + state.apply_staged_writes() + assert int(state.num_bad_words.np[req_idx]) == 0 + + num_logits = len(INPUT_IDS) + idx_mapping = torch.tensor([req_idx], dtype=torch.int32, device=DEVICE) + out = state.apply( + torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE), + LogitsContext( + expanded_idx_mapping=torch.tensor( + [req_idx] * num_logits, dtype=torch.int32, device=DEVICE + ), + idx_mapping=idx_mapping, + idx_mapping_np=np.array([req_idx], dtype=np.intp), + expanded_local_pos=torch.tensor( + LOCAL_POS, dtype=torch.int32, device=DEVICE + ), + input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE), + pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE), + seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64), + ), + ) + assert not torch.isinf(out).any() + + def test_v2_bad_words_prefix_inside_draft_tokens(): """A prefix matching entirely within the draft tokens must mask the bad word's last token at the draft position that completes the prefix.""" diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 1309686d53ce..ffb092209859 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -183,6 +183,12 @@ def test_local_batch_partition(tp_size: int, seed: int): ).all() assert (_np(local.logits_indices) == _np(batch.logits_indices)[rows]).all() assert (_np(local.seq_lens) == _np(batch.seq_lens)[owned]).all() + # Indexed by batch position, like seq_lens: a consumer holding local row + # indices reads another request's length if this is not localised too. + assert ( + _np(local.seq_lens_cpu_upper_bound) + == _np(batch.seq_lens_cpu_upper_bound)[owned] + ).all() assert batch.num_draft_tokens_per_req is not None assert local.num_draft_tokens_per_req is not None assert ( diff --git a/tests/v1/worker/test_gpu_logit_bias.py b/tests/v1/worker/test_gpu_logit_bias.py index 3bf62380bbc6..617a36120476 100644 --- a/tests/v1/worker/test_gpu_logit_bias.py +++ b/tests/v1/worker/test_gpu_logit_bias.py @@ -14,6 +14,11 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState +from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsContext, + LogitsProcRequestState, +) +from vllm.v1.worker.gpu.states import RequestState DEVICE = torch.device("cuda") VOCAB_SIZE = 128 @@ -45,17 +50,35 @@ def _only_stop_token_left(num_rows: int) -> torch.Tensor: def _apply(logits: torch.Tensor, structured: list[bool]) -> torch.Tensor: """Run the fused bias kernel with one logits row per request.""" - state = LogitBiasState(max_num_reqs=4, device=DEVICE) + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + # add_request() reads prompt_len per slot; every slot here uses the same one. + req_states.prompt_len.np[:] = PROMPT_LEN + + state = LogitBiasState(None, LogitsProcRequestState.from_request_state(req_states)) for req_idx, is_structured in enumerate(structured): - state.add_request(req_idx, PROMPT_LEN, _params(is_structured)) + state.add_request(req_idx, _params(is_structured)) state.apply_staged_writes() num_reqs = len(structured) - state.apply_logit_bias( + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) + state.apply( logits, - torch.arange(num_reqs, dtype=torch.int32, device=DEVICE), - np.arange(num_reqs, dtype=np.intp), - torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE), + LogitsContext( + expanded_idx_mapping=idx_mapping, + idx_mapping=idx_mapping, + idx_mapping_np=np.arange(num_reqs, dtype=np.intp), + expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE), + input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE), + pos=torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE), + seq_lens_upper_bound_np=np.full(num_reqs, POS + 1, dtype=np.int64), + ), ) return logits.cpu() diff --git a/tests/v1/worker/test_gpu_logits_processors.py b/tests/v1/worker/test_gpu_logits_processors.py new file mode 100644 index 000000000000..f0e5bd43a54b --- /dev/null +++ b/tests/v1/worker/test_gpu_logits_processors.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for Model Runner V2 logits processor tests", + allow_module_level=True, + ) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsContext, + LogitsProcessor, +) +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 +TARGET_TOKEN = 7 +BIAS = 5.0 +BOOST = 100.0 +TEMPERATURE = 2.0 + + +class MockReasoningConfig: + reasoning_start_token_ids = [90] + reasoning_end_token_ids = [91] + natural_reasoning_end_token_ids = [91] + + +class RecordingProcessor(LogitsProcessor): + """Adds BOOST to TARGET_TOKEN and records every lifecycle callback.""" + + def __init__(self, vllm_config: Any, req_state: Any): + self.events: list[str] = [] + self.seen_logits: torch.Tensor | None = None + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + self.events.append(f"add_request:{req_idx}") + return True + + def apply_staged_writes(self) -> None: + self.events.append("apply_staged_writes") + + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + self.events.append("apply") + self.seen_logits = logits.clone() + logits[:, TARGET_TOKEN] += BOOST + return logits + + +def _make_sampler(processor: RecordingProcessor) -> Sampler: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + return Sampler( + vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()), + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + custom_logits_processors=[processor], + ) + + +def test_custom_processor_lifecycle_and_pipeline_position(): + """A custom processor's hooks fire in order, it runs after the built-in + bias stage, and before temperature scaling.""" + proc = RecordingProcessor(None, None) + sampler = _make_sampler(proc) + sampler.add_request( + 0, + SamplingParams(logit_bias={TARGET_TOKEN: BIAS}, temperature=TEMPERATURE), + ) + sampler.apply_staged_writes() + + num_reqs = 1 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) + raw = torch.zeros(num_reqs, VOCAB_SIZE, device=DEVICE) + out = sampler.apply_sampling_params( + raw, + expanded_idx_mapping=idx_mapping, + idx_mapping=idx_mapping, + idx_mapping_np=np.arange(num_reqs, dtype=np.intp), + pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE), + input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE), + expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE), + seq_lens_upper_bound_np=np.ones(num_reqs, dtype=np.int64), + ) + + assert proc.events == [ + "add_request:0", + "apply_staged_writes", + "apply", + ] + # The processor sees the built-in bias applied but no temperature scaling. + assert proc.seen_logits is not None + assert proc.seen_logits[0, TARGET_TOKEN].item() == BIAS + # Its own modification is then temperature-scaled by the sampler. + assert out[0, TARGET_TOKEN].item() == (BIAS + BOOST) / TEMPERATURE diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 1176857dffa4..b9cee16ac4d3 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -56,7 +56,16 @@ def from_mapping(cls, **kwargs: Any) -> "FakeEplbState": def _make_runner(**overrides: Any) -> Any: runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) runner.device = torch.device("cpu") - runner.model_config = SimpleNamespace(model="test-model") + runner.model_config = SimpleNamespace(model="test-model", logits_processors=None) + runner.req_states = SimpleNamespace( + device=runner.device, + max_num_reqs=8, + vocab_size=0, + all_token_ids=None, + prompt_len=None, + prefill_len=None, + total_len=None, + ) runner.load_config = SimpleNamespace(load_format="hf") runner.parallel_config = SimpleNamespace( enable_eplb=True, diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py index 2e5adfabd13c..efd2b59d75ed 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -47,6 +47,9 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str): [np.arange(count, dtype=np.int32) for count in num_logits_per_req] ) ).to(device), + seq_lens_cpu_upper_bound=torch.from_numpy( + np.array([10, 20, 30, 40], dtype=np.int32) + ), ) rejection_sampler = object.__new__(RejectionSampler) rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode) diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py index 3fd1cb3b3c68..3f05a2cfb59f 100644 --- a/tests/v1/worker/test_gpu_sampler_flags.py +++ b/tests/v1/worker/test_gpu_sampler_flags.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence +from types import SimpleNamespace + import numpy as np import pytest import torch @@ -10,6 +13,10 @@ pytest.skip("CUDA required for sampler flag tests", allow_module_level=True) from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsContext, + LogitsProcessor, +) from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.states import RequestState @@ -23,7 +30,7 @@ class MockReasoningConfig: natural_reasoning_end_token_ids = [91] -def _make_sampler() -> Sampler: +def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> Sampler: req_states = RequestState( max_num_reqs=4, max_model_len=64, @@ -33,11 +40,12 @@ def _make_sampler() -> Sampler: device=DEVICE, ) return Sampler( + vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()), max_num_reqs=4, vocab_size=VOCAB_SIZE, device=DEVICE, req_states=req_states, - reasoning_config=MockReasoningConfig(), + custom_logits_processors=custom_logits_processors, ) @@ -65,23 +73,46 @@ def test_logits_processing_cache_matches_request_features( sampling_params: SamplingParams, expected: bool ): sampler = _make_sampler() - sampler.add_request(3, prompt_len=1, sampling_params=sampling_params) + sampler.add_request(3, sampling_params=sampling_params) assert sampler.needs_logits_processing[3] == expected def test_logits_processing_cache_is_overwritten_when_slot_is_reused(): sampler = _make_sampler() - sampler.add_request(3, 1, SamplingParams.for_sampler_warmup()) - sampler.add_request(3, 1, SamplingParams()) + sampler.add_request(3, SamplingParams.for_sampler_warmup()) + sampler.add_request(3, SamplingParams()) assert not sampler.needs_logits_processing[3] +class _GateProcessor(LogitsProcessor): + """Reports a fixed admission decision from add_request().""" + + def __init__(self, admitted: bool): + self.admitted = admitted + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + return self.admitted + + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + raise AssertionError("the pipeline never runs in this test") + + +@pytest.mark.parametrize("admitted", [True, False]) +def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool): + """A custom processor's add_request() return value is OR-ed into the + per-request needs_logits_processing flag.""" + sampler = _make_sampler(custom_logits_processors=[_GateProcessor(admitted)]) + sampler.add_request(3, SamplingParams()) + + assert sampler.needs_logits_processing[3] == admitted + + def test_logits_processing_cache_only_checks_active_requests(): sampler = _make_sampler() - sampler.add_request(0, 1, SamplingParams(temperature=0.0)) - sampler.add_request(2, 1, SamplingParams.for_sampler_warmup()) + sampler.add_request(0, SamplingParams(temperature=0.0)) + sampler.add_request(2, SamplingParams.for_sampler_warmup()) sampling_only = np.array([0], dtype=np.int32) with_processing = np.array([0, 2], dtype=np.int32) diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index 2bc697b08438..245602733535 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import numpy as np import pytest import torch @@ -12,6 +15,7 @@ ) from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState from vllm.v1.worker.gpu.states import RequestState @@ -75,11 +79,17 @@ def _apply( idx_mapping_np = idx_mapping.cpu().numpy() state.apply( logits, - expanded_idx_mapping, - idx_mapping, - idx_mapping_np, - torch.tensor(input_ids, dtype=torch.int32, device=DEVICE), - torch.tensor(local_pos, dtype=torch.int32, device=DEVICE), + LogitsContext( + expanded_idx_mapping=expanded_idx_mapping, + idx_mapping=idx_mapping, + idx_mapping_np=idx_mapping_np, + expanded_local_pos=torch.tensor( + local_pos, dtype=torch.int32, device=DEVICE + ), + input_ids=torch.tensor(input_ids, dtype=torch.int32, device=DEVICE), + pos=torch.zeros(len(input_ids), dtype=torch.int32, device=DEVICE), + seq_lens_upper_bound_np=np.full(1, len(input_ids), dtype=np.int64), + ), ) return logits.cpu() @@ -190,15 +200,14 @@ def test_v2_greedy_sampling_applies_thinking_budget(): """Greedy-only requests must not bypass thinking-budget processing.""" req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) sampler = Sampler( + vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()), max_num_reqs=4, vocab_size=VOCAB_SIZE, device=DEVICE, req_states=req_states, - reasoning_config=MockReasoningConfig(), ) sampler.add_request( req_idx=3, - prompt_len=1, sampling_params=SamplingParams( temperature=0.0, thinking_token_budget=3, @@ -219,6 +228,7 @@ def test_v2_greedy_sampling_applies_thinking_budget(): torch.tensor([4], dtype=torch.int32, device=DEVICE), input_ids, torch.tensor([0], dtype=torch.int32, device=DEVICE), + np.full(1, 4, dtype=np.int64), ) assert out[0, END].item() == pytest.approx(1.0e9) diff --git a/tests/v1/worker/test_logits_processor_loader.py b/tests/v1/worker/test_logits_processor_loader.py new file mode 100644 index 000000000000..7e851e5237d1 --- /dev/null +++ b/tests/v1/worker/test_logits_processor_loader.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for the V2 custom logits processor loader. + +Load-time validation is pure Python (importlib + issubclass), so it runs +without CUDA. Pipeline behavior is covered by test_gpu_logits_processors.py. +""" + +import subprocess +import sys +from enum import Enum, auto +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +import vllm.v1.worker.gpu.sample.logits_processor.loader as loader +from vllm.exceptions import VLLMValidationError +from vllm.sampling_params import SamplingParams +from vllm.v1.sample.logits_processor import ( + LogitsProcessor as V1LogitsProcessor, +) +from vllm.v1.worker.gpu.sample.logits_processor import ( + LogitsContext, + LogitsProcessor, +) + + +class DummyV2Processor(LogitsProcessor): + """Records its constructor args so tests can assert how it was built.""" + + def __init__(self, vllm_config: Any, req_state: Any): + self.ctor_args = (vllm_config, req_state) + + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + return logits + + +class AltV2Processor(DummyV2Processor): + """Second V2 processor, used to assert plugin/FQCN load ordering.""" + + +class DummyV1Processor(V1LogitsProcessor): + """A V1-interface processor. Never instantiated: the loader must reject + it at the issubclass check, before any constructor runs.""" + + +def not_a_processor() -> None: + """FQCN target that resolves but is not a class.""" + + +def _fake_entrypoint(cls: type) -> SimpleNamespace: + return SimpleNamespace( + name=cls.__name__.lower(), + value=f"fake:{cls.__name__}", + load=lambda: cls, + ) + + +@pytest.fixture(autouse=True) +def no_installed_plugins(monkeypatch: pytest.MonkeyPatch): + """Isolate tests from entry points registered in the environment.""" + monkeypatch.setattr(loader, "entry_points", lambda group: []) + + +class ProcSource(Enum): + CLASS = auto() + FQCN = auto() + ENTRYPOINT = auto() + + +@pytest.mark.parametrize("source", list(ProcSource)) +def test_loads_v2_processors(monkeypatch: pytest.MonkeyPatch, source: ProcSource): + """A V2 subclass loads via class object, FQCN string, or entrypoint.""" + fake_req_states = SimpleNamespace( + device="cpu", + max_num_reqs=4, + vocab_size=128, + all_token_ids=None, + prompt_len=None, + prefill_len=None, + total_len=None, + ) + custom: list[Any] + if source is ProcSource.CLASS: + custom = [DummyV2Processor] + elif source is ProcSource.FQCN: + custom = [f"{__name__}:DummyV2Processor"] + else: + monkeypatch.setattr( + loader, "entry_points", lambda group: [_fake_entrypoint(DummyV2Processor)] + ) + # Plugins load before user-specified processors. + custom = [f"{__name__}:AltV2Processor"] + + procs = loader.build_custom_logits_processors(None, fake_req_states, False, custom) + + if source is ProcSource.ENTRYPOINT: + assert [type(p) for p in procs] == [DummyV2Processor, AltV2Processor] + else: + assert [type(p) for p in procs] == [DummyV2Processor] + vllm_config, req_state = procs[0].ctor_args + assert vllm_config is None + assert isinstance(req_state, loader.LogitsProcRequestState) + assert req_state.max_num_reqs == 4 + + +@pytest.mark.parametrize( + ("custom", "plugin_cls", "exc_type", "msg_fragment"), + [ + pytest.param( + [DummyV1Processor], + None, + ValueError, + "not a subclass", + id="v1-class-object", + ), + pytest.param( + [f"{__name__}:DummyV1Processor"], + None, + ValueError, + "not a subclass", + id="v1-fqcn", + ), + pytest.param( + [], + DummyV1Processor, + ValueError, + "V1-interface plugins are not supported", + id="v1-entrypoint", + ), + pytest.param( + [f"{__name__}:not_a_processor"], + None, + ValueError, + "must be a type", + id="non-type-fqcn", + ), + pytest.param( + ["no.such.module:Nope"], + None, + RuntimeError, + "no.such.module:Nope", + id="missing-module", + ), + pytest.param( + ["no_colon_here"], + None, + ValueError, + "Expected format", + id="fqcn-missing-colon", + ), + pytest.param( + ["too:many:colons"], + None, + ValueError, + "Expected format", + id="fqcn-extra-colons", + ), + ], +) +def test_rejects_invalid( + monkeypatch: pytest.MonkeyPatch, + custom: list, + plugin_cls: type | None, + exc_type: type[Exception], + msg_fragment: str, +): + """V1 processors and malformed references fail at load time, not at + sampling time, with a message naming the offending reference.""" + if plugin_cls is not None: + monkeypatch.setattr( + loader, "entry_points", lambda group: [_fake_entrypoint(plugin_cls)] + ) + + with pytest.raises(exc_type, match=msg_fragment): + loader.build_custom_logits_processors(None, None, False, custom) + + +def test_pooling_model_rejects_custom_logitsprocs(): + """Pooling models reject custom processors instead of ignoring them.""" + with pytest.raises(ValueError, match="Pooling models do not support"): + loader.build_custom_logits_processors(None, None, True, [DummyV2Processor]) + assert loader.build_custom_logits_processors(None, None, True, []) == [] + + +class ValidatingProcessor(DummyV2Processor): + """Rejects target_token=-1 via the optional validate_params hook.""" + + @classmethod + def validate_params(cls, sampling_params: SamplingParams) -> None: + if (sampling_params.extra_args or {}).get("target_token") == -1: + raise ValueError("target_token must not be -1") + + +def test_validate_params_runs_at_admission(): + """The factory loads classes up front; the returned validator runs + validate_params per request and wraps rejections as VLLMValidationError.""" + loader._cached_load_v2_logitsprocs.cache_clear() + validate = loader.build_custom_logits_processors_params_validator( + [ValidatingProcessor] + ) + validate(SamplingParams(extra_args={"target_token": 1})) + with pytest.raises(VLLMValidationError, match="target_token"): + validate(SamplingParams(extra_args={"target_token": -1})) + + +def test_validator_factory_loads_classes_eagerly(): + """A bad reference fails when the validator is built (at startup), not on + the first request.""" + loader._cached_load_v2_logitsprocs.cache_clear() + with pytest.raises(RuntimeError, match="no.such.module:Nope"): + loader.build_custom_logits_processors_params_validator(["no.such.module:Nope"]) + + +def test_loader_import_stays_frontend_safe(): + """The frontend process imports the loader to validate params; importing + it must not pull in model-runner side modules.""" + code = ( + "import sys\n" + "import vllm.v1.worker.gpu.sample.logits_processor.loader\n" + "assert 'vllm.sampling_params' not in sys.modules\n" + "assert 'vllm.v1.sample.logits_processor' not in sys.modules\n" + "assert 'vllm.v1.worker.gpu.states' not in sys.modules\n" + "assert 'vllm.v1.worker.gpu.buffer_utils' not in sys.modules\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2ee458e7ff65..78e30cf9d058 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2880,7 +2880,6 @@ def _validate_mm_processor_device(self) -> None: def _get_v2_model_runner_unsupported_features(self) -> list[str]: """Collect features not yet supported by the V2 model runner.""" unsupported: list[str] = [] - model_config = self.model_config speculative_config = self.speculative_config if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: @@ -2929,17 +2928,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: if self.parallel_config.enable_elastic_ep: unsupported.append("elastic expert parallelism") - has_logitsproc_plugins = False - if model_config is not None: - from importlib.metadata import entry_points - - has_logitsproc_plugins = bool(entry_points(group="vllm.logits_processors")) - - if model_config is not None and ( - model_config.logits_processors or has_logitsproc_plugins - ): - unsupported.append("custom logits processors") - if self.cache_config.mamba_cache_mode == "all": unsupported.append("mamba cache mode 'all'") diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 401d99087066..0c1cc8c55e1a 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -1014,7 +1014,7 @@ def __init__( # on the subsequent commit step when num_sampled=CANVAS_LEN. self._pending_logprobs: dict[int, LogprobsTensors] = {} - def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None: + def add_request(self, req_idx: int, sampling_params: Any) -> None: if use_penalty(sampling_params): logger.warning_once( "DiffusionGemma does not support repetition/frequency/presence " diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 6cd7b4e81c8a..120dbaddfd80 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -823,7 +823,6 @@ def verify( self._validate_logit_bias(model_config) self._validate_trace_replay(model_config, speculative_config) self._validate_stop_token_ids(model_config) - self._validate_logits_processors(model_config) self._validate_allowed_token_ids(model_config) self._validate_spec_decode(speculative_config) self._validate_diffusion(model_config) @@ -993,13 +992,6 @@ def _validate_trace_replay( value=invalid_token_ids, ) - def _validate_logits_processors(self, model_config: ModelConfig) -> None: - from vllm.v1.sample.logits_processor import ( - validate_logits_processors_parameters, - ) - - validate_logits_processors_parameters(model_config.logits_processors, self) - def _validate_allowed_token_ids(self, model_config: ModelConfig) -> None: allowed_token_ids = self.allowed_token_ids if allowed_token_ids is None: diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 256a72df4254..06b2b32efd2a 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from functools import partial from typing import Any, Literal import vllm.envs as envs @@ -51,6 +52,11 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.structured_outputs_config = vllm_config.structured_outputs_config self.observability_config = vllm_config.observability_config + # Load the custom logits processor classes once; the returned callable + # runs their validate_params hooks per request at admission. + self.validate_logits_processors_params = ( + self._build_logits_processors_params_validator() + ) self.generation_config_fields = model_config.try_get_generation_config() @@ -81,6 +87,25 @@ def tokenizer(self) -> TokenizerLike | None: def get_tokenizer(self) -> TokenizerLike: return self.renderer.get_tokenizer() + def _build_logits_processors_params_validator( + self, + ) -> Callable[[SamplingParams], None]: + """Load the custom logits processor classes and return the per-request + params validator for the active model runner.""" + custom_logitsprocs = self.model_config.logits_processors + if self.vllm_config.use_v2_model_runner: + from vllm.v1.worker.gpu.sample.logits_processor import ( + build_custom_logits_processors_params_validator, + ) + + return build_custom_logits_processors_params_validator(custom_logitsprocs) + + from vllm.v1.sample.logits_processor import ( + validate_logits_processors_parameters, + ) + + return partial(validate_logits_processors_parameters, custom_logitsprocs) + def _validate_params( self, params: SamplingParams | PoolingParams, @@ -101,6 +126,8 @@ def _validate_params( self.tokenizer, ) + self.validate_logits_processors_params(params) + if self.model_config.return_sampling_mask: if params.temperature <= 0: raise ValueError( diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index afa95ee8781b..c9624b2b25ce 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -36,10 +36,8 @@ def __init__( self.watermarking.np.fill(True) self.watermarking.copy_to_uva() - def add_request( - self, req_idx: int, prompt_len: int, sampling_params: SamplingParams - ) -> None: - super().add_request(req_idx, prompt_len, sampling_params) + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + super().add_request(req_idx, sampling_params) self.watermarking.np[req_idx] = sampling_params.watermarking if sampling_params.watermarking and sampling_params.temperature == 0: logger.warning_once( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4dd8865cbd22..e76c447ea5cc 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -143,6 +143,7 @@ all_to_all_logits, gather_sampler_output, ) +from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors from vllm.v1.worker.gpu.sample.output import SamplerOutput from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker from vllm.v1.worker.gpu.sample.sampler import Sampler @@ -447,8 +448,13 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: ) # Initialize samplers. Model states may override via custom_sampler(). + config_processors = self.model_config.logits_processors or () + custom_logits_processors = build_custom_logits_processors( + self.vllm_config, self.req_states, self.is_pooling_model, config_processors + ) if self.is_last_pp_rank and not self.is_pooling_model: sampler_kwargs: dict[str, Any] = { + "vllm_config": self.vllm_config, "max_num_reqs": self.max_num_reqs, "vocab_size": self.vocab_size, "device": self.device, @@ -457,8 +463,8 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: "num_speculative_tokens": self.decode_query_len, "use_fp64_gumbel": self.model_config.use_fp64_gumbel, "enable_trace_replay": self.model_config.enable_trace_replay, - "reasoning_config": self.vllm_config.reasoning_config, "return_sampling_mask": self.model_config.return_sampling_mask, + "custom_logits_processors": custom_logits_processors, } if self.vllm_config.watermark_config is None: self.sampler = Sampler(**sampler_kwargs) @@ -1147,9 +1153,7 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: if self.is_last_pp_rank and new_req_data.sampling_params is not None: assert self.sampler is not None - self.sampler.add_request( - req_index, prompt_len, new_req_data.sampling_params - ) + self.sampler.add_request(req_index, new_req_data.sampling_params) assert self.prompt_logprobs_worker is not None self.prompt_logprobs_worker.add_request( req_id, req_index, new_req_data.sampling_params diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py index 3fb9d148ecac..974b24359889 100644 --- a/vllm/v1/worker/gpu/sample/bad_words.py +++ b/vllm/v1/worker/gpu/sample/bad_words.py @@ -1,43 +1,47 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import numpy as np +from typing import TYPE_CHECKING + import torch import vllm.envs as envs from vllm.sampling_params import SamplingParams from vllm.triton_utils import tl, triton from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor -from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsContext, + LogitsProcessor, + LogitsProcRequestState, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig -class BadWordsState: - def __init__(self, req_states: RequestState): +class BadWordsState(LogitsProcessor): + def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): self.req_states = req_states - self.max_num_reqs = req_states.max_num_reqs - self.device = req_states.device + max_num_reqs = req_states.max_num_reqs + device = req_states.device max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS # flattened bad word tokens: [max_num_reqs, VLLM_MAX_BAD_WORDS_TOTAL_TOKENS] self.bad_word_token_ids = StagedWriteTensor( - (self.max_num_reqs, max_total_tokens), - dtype=torch.int32, - device=self.device, + (max_num_reqs, max_total_tokens), dtype=torch.int32, device=device ) # cumulative offsets of bad words: [max_num_reqs, VLLM_MAX_NUM_BAD_WORDS + 1] self.bad_word_offsets = StagedWriteTensor( - (self.max_num_reqs, max_num_bad_words + 1), - dtype=torch.int32, - device=self.device, + (max_num_reqs, max_num_bad_words + 1), dtype=torch.int32, device=device ) # number of bad words per request - self.num_bad_words = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) + self.num_bad_words = UvaBackedTensor(max_num_reqs, dtype=torch.int32) - def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: bad_words_token_ids = sampling_params.bad_words_token_ids if not bad_words_token_ids: self.num_bad_words.np[req_idx] = 0 - return + return False num_bad_words = len(bad_words_token_ids) max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS @@ -65,38 +69,33 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens) self.bad_word_offsets.stage_write(req_idx, 0, offsets) self.num_bad_words.np[req_idx] = num_bad_words + return True def apply_staged_writes(self) -> None: self.num_bad_words.copy_to_uva() self.bad_word_token_ids.apply_write() self.bad_word_offsets.apply_write() - def apply_bad_words( - self, - logits: torch.Tensor, - expanded_idx_mapping: torch.Tensor, - idx_mapping_np: np.ndarray, - input_ids: torch.Tensor, - expanded_local_pos: torch.Tensor, - ) -> None: - max_num_bad_words = int(self.num_bad_words.np[idx_mapping_np].max()) + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + max_num_bad_words = int(self.num_bad_words.np[ctx.idx_mapping_np].max()) if max_num_bad_words == 0: # No request uses bad words. Skip the kernel launch. - return + return logits apply_bad_words( logits, - expanded_idx_mapping, + ctx.expanded_idx_mapping, self.bad_word_token_ids.gpu, self.bad_word_offsets.gpu, self.num_bad_words.gpu, self.req_states.all_token_ids.gpu, self.req_states.prompt_len.gpu, self.req_states.total_len.gpu, - input_ids, - expanded_local_pos, + ctx.input_ids, + ctx.expanded_local_pos, max_num_bad_words, ) + return logits @triton.jit diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index 0f484f1a1710..b2803de9daeb 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -200,6 +200,9 @@ def shard_sampler_inputs( num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:] ) max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1 + local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[ + torch.from_numpy(local_req_indices_np) + ] # Shard the input batch GPU tensors. sorted_logits_indices = torch.empty( @@ -269,6 +272,7 @@ def shard_sampler_inputs( expanded_idx_mapping=local_expanded_idx_mapping, expanded_local_pos=local_expanded_local_pos, seq_lens=local_seq_lens, + seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound, logits_indices=local_logits_indices, cu_num_logits=local_cu_num_logits, cu_num_logits_np=local_cu_num_logits_np, diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index 846edbe1e74d..53e20802b69f 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -1,60 +1,61 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import numpy as np import torch from vllm.sampling_params import SamplingParams from vllm.triton_utils import tl, triton from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsContext, + LogitsProcessor, + LogitsProcRequestState, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig MAX_NUM_ALLOWED_TOKEN_IDS = 1024 MAX_NUM_LOGIT_BIAS_TOKENS = 1024 MAX_NUM_STOP_TOKEN_IDS = 128 -class LogitBiasState: - def __init__(self, max_num_reqs: int, device: torch.device): - self.max_num_reqs = max_num_reqs +class LogitBiasState(LogitsProcessor): + def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): + self.prompt_len_np = req_states.prompt_len.np + max_num_reqs = req_states.max_num_reqs + device = req_states.device # Allowed token IDs. - self.num_allowed_token_ids = UvaBackedTensor( - self.max_num_reqs, dtype=torch.int32 - ) + self.num_allowed_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32) self.allowed_token_ids = StagedWriteTensor( - (self.max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS), - dtype=torch.int32, - device=device, + (max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS), dtype=torch.int32, device=device ) # Logit bias. - self.num_logit_bias = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) + self.num_logit_bias = UvaBackedTensor(max_num_reqs, dtype=torch.int32) self.logit_bias_token_ids = StagedWriteTensor( - (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), - dtype=torch.int32, - device=device, + (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), dtype=torch.int32, device=device ) self.logit_bias = StagedWriteTensor( - (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), + (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), dtype=torch.float32, device=device, ) # Min tokens. - self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) - self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) - self.restore_when_all_masked = UvaBackedTensor( - self.max_num_reqs, dtype=torch.int32 - ) + self.min_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + self.num_stop_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + self.restore_when_all_masked = UvaBackedTensor(max_num_reqs, dtype=torch.int32) self.stop_token_ids = StagedWriteTensor( - (self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), - dtype=torch.int32, - device=device, + (max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), dtype=torch.int32, device=device ) # Using any of the above. self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool) - def add_request( - self, req_idx: int, prompt_len: int, sampling_params: SamplingParams - ) -> None: + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + prompt_len = int(self.prompt_len_np[req_idx]) # Using any logit bias. use_logit_bias = False @@ -112,6 +113,7 @@ def add_request( self.restore_when_all_masked.np[req_idx] = 0 self.use_logit_bias[req_idx] = use_logit_bias + return use_logit_bias def apply_staged_writes(self) -> None: self.num_allowed_token_ids.copy_to_uva() @@ -126,25 +128,19 @@ def apply_staged_writes(self) -> None: self.restore_when_all_masked.copy_to_uva() self.stop_token_ids.apply_write() - def apply_logit_bias( - self, - logits: torch.Tensor, - expanded_idx_mapping: torch.Tensor, - idx_mapping_np: np.ndarray, - pos: torch.Tensor, - ) -> None: - if not np.any(self.use_logit_bias[idx_mapping_np]): + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + if not np.any(self.use_logit_bias[ctx.idx_mapping_np]): # No request uses logit bias. Skip the kernel launch. - return + return logits enable_stop_token_restore = bool( - np.any(self.restore_when_all_masked.np[idx_mapping_np]) + np.any(self.restore_when_all_masked.np[ctx.idx_mapping_np]) ) apply_logit_bias( logits, - expanded_idx_mapping, - pos, + ctx.expanded_idx_mapping, + ctx.pos, self.num_allowed_token_ids.gpu, self.allowed_token_ids.gpu, self.num_logit_bias.gpu, @@ -156,6 +152,7 @@ def apply_logit_bias( self.stop_token_ids.gpu, enable_stop_token_restore, ) + return logits @triton.jit diff --git a/vllm/v1/worker/gpu/sample/logits_processor/__init__.py b/vllm/v1/worker/gpu/sample/logits_processor/__init__.py new file mode 100644 index 000000000000..f76eed35ceb3 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/logits_processor/__init__.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom logits processors for the V2 model runner. + +Kept import-light: the frontend process imports this package to validate +per-request params, so nothing here may pull in model-runner side modules +(torch, triton, worker state) at import time. +""" + +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsContext, + LogitsProcessor, + LogitsProcRequestState, +) +from vllm.v1.worker.gpu.sample.logits_processor.loader import ( + build_custom_logits_processors, + build_custom_logits_processors_params_validator, +) + +__all__ = [ + "LogitsProcRequestState", + "LogitsContext", + "LogitsProcessor", + "build_custom_logits_processors", + "build_custom_logits_processors_params_validator", +] diff --git a/vllm/v1/worker/gpu/sample/logits_processor/interface.py b/vllm/v1/worker/gpu/sample/logits_processor/interface.py new file mode 100644 index 000000000000..9934492ec227 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/logits_processor/interface.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""V2 logits processor interface. + +Import-light by design: custom processor classes are also loaded in the +frontend process (to validate per-request params), so this module must not +pull in model-runner side modules (torch, triton, worker state) at import +time. Heavy imports live under TYPE_CHECKING; annotations are deferred. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + import torch + + from vllm.config import VllmConfig + from vllm.sampling_params import SamplingParams + from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor + from vllm.v1.worker.gpu.states import RequestState + + +@dataclass(frozen=True) +class LogitsProcRequestState: + """State associated with active requests, shared with logits processors. + + Wraps the model runner's per-slot buffers, which are mutated in place, + so reads always see current values. Processors must treat every field + as read-only. + """ + + device: torch.device + max_num_reqs: int + vocab_size: int + + # [max_num_reqs, max_model_len] committed token ids per request slot. + all_token_ids: StagedWriteTensor + # [max_num_reqs] tokens in the user-provided prompt. + prompt_len: UvaBackedTensor + # [max_num_reqs] tokens fed at the latest (re)fill: the prompt plus any + # partial output on resumption after preemption. + prefill_len: UvaBackedTensor + # [max_num_reqs] prompt_len + output_len; grows as the request progresses. + total_len: StagedWriteTensor + + @classmethod + def from_request_state(cls, req_states: RequestState) -> LogitsProcRequestState: + return cls( + device=req_states.device, + max_num_reqs=req_states.max_num_reqs, + vocab_size=req_states.vocab_size, + all_token_ids=req_states.all_token_ids, + prompt_len=req_states.prompt_len, + prefill_len=req_states.prefill_len, + total_len=req_states.total_len, + ) + + +@dataclass(frozen=True) +class LogitsContext: + """The current step's batch layout, passed to every ``apply()`` call. + + A row is a logits row, not a request: rows are reordered every step, and + under speculative decoding a request owns one row per draft token. + Committed tokens live in ``req_states.all_token_ids`` (valid up to + ``total_len``); this step's draft tokens are only in ``input_ids``. + """ + + # [num_logits_rows] row -> persistent request slot. + expanded_idx_mapping: torch.Tensor + # [num_reqs] batch position -> persistent request slot. + idx_mapping: torch.Tensor + # [num_reqs] batch position -> persistent request slot, on the host, for + # skipping work without a device sync. + idx_mapping_np: np.ndarray + # [num_logits_rows] row -> its offset among the rows of its own request. + expanded_local_pos: torch.Tensor + # [num_logits_rows] token fed to the model at each row's input position. + input_ids: torch.Tensor + # [num_logits_rows] position of each row within its sequence. + pos: torch.Tensor + # [num_reqs] batch position -> upper bound of tokens visible to the token + # being sampled this step, on the host. Exact when spec decoding isn't in use. + # Exact per-row lengths on device are `pos + 1`. + seq_lens_upper_bound_np: np.ndarray + + +class LogitsProcessor(ABC): + """Custom logits processor for Model Runner V2. + + Per-request state is keyed by the request slot index; slots are recycled + through a free list, so per-slot state must be fully (re)initialized in + ``add_request()``. + + ``apply()`` runs after the built-in bias, penalty, bad-words and grammar + stages and before temperature, min_p and top-k/top-p, so it sees unscaled + logits and must not re-inflate grammar-masked tokens. Thinking-budget + forcing runs after ``apply()`` and wins over its edits for requests whose + budget is exhausted. + + State that is constant for a request belongs in ``__init__()`` or + ``add_request()``; ``apply()`` receives only what changes per step. + """ + + def __init__( # noqa: B027 + self, vllm_config: VllmConfig, req_states: LogitsProcRequestState + ): + """Capture what stays constant for the processor's lifetime. + + ``req_states`` exposes the on-device token history and batch constants a + processor may read. Treat it as read-only. + """ + + @classmethod # noqa: B027 + def validate_params(cls, sampling_params: SamplingParams) -> None: + """Raise ``ValueError`` for invalid per-request arguments. + + Runs at request admission, so invalid arguments fail the request + with an error instead of reaching the sampler. + """ + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + """Initialize per-slot state for a request entering the batch. + + The slot may hold a previous occupant's state; overwrite or + neutralize all of it here. + + Returns whether this processor modifies logits for the request. + """ + return True + + def apply_staged_writes(self) -> None: # noqa: B027 + """Flush any host-side writes staged by ``add_request()`` to the device. + + Called once per step before the forward pass, after the model runner + has flushed ``req_states``, so a processor that stages writes here can + read the request's tokens back on device. + """ + + @abstractmethod + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + """Modify logits in place or return a new tensor. In-place modification + is preferred for efficiency. + + ``apply()`` is called once for the whole batch, including rows of + requests this processor declined in ``add_request()``, so filter rows + via ``ctx.expanded_idx_mapping``. + + Args: + logits: [num_logits_rows, vocab_size] float32 tensor. + ctx: this step's batch layout. + + """ + raise NotImplementedError diff --git a/vllm/v1/worker/gpu/sample/logits_processor/loader.py b/vllm/v1/worker/gpu/sample/logits_processor/loader.py new file mode 100644 index 000000000000..e3b94d539485 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/logits_processor/loader.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Loading of custom logits processor classes for the V2 model runner. + +This module is deliberately import-light: the frontend process uses it to +validate per-request params, so importing it must not pull in model-runner +side modules (torch, triton, worker state). Imports needed to instantiate +processors are deferred to ``build_custom_logits_processors``. +""" + +import importlib +from collections.abc import Callable, Sequence +from functools import lru_cache +from importlib.metadata import entry_points +from typing import TYPE_CHECKING + +from vllm.exceptions import VLLMValidationError +from vllm.logger import init_logger +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsProcessor, + LogitsProcRequestState, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.sampling_params import SamplingParams + from vllm.v1.worker.gpu.states import RequestState + +logger = init_logger(__name__) + +# Same entry-points group as +# vllm.v1.sample.logits_processor.LOGITSPROCS_GROUP; duplicated here so this +# module stays importable without the V1 sampler machinery. +LOGITSPROCS_GROUP = "vllm.logits_processors" + + +def _load_v2_logitsprocs_plugins() -> list[type[LogitsProcessor]]: + """Load installed logits processor plugins. + + Plugins registered under ``LOGITSPROCS_GROUP`` ship V1-interface + processors today; those crash under the V2 sampler, so reject them at + load time. + """ + from vllm.utils.torch_utils import guard_cuda_initialization + + installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP) + if len(installed_logitsprocs_plugins) == 0: + logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP) + return [] + + logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP) + classes: list[type[LogitsProcessor]] = [] + for entrypoint in installed_logitsprocs_plugins: + try: + logger.debug( + "- Loading logitproc plugin entrypoint=%s target=%s", + entrypoint.name, + entrypoint.value, + ) + with guard_cuda_initialization(): + cls = entrypoint.load() + except Exception as e: + logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e) + raise RuntimeError( + f"Failed to load LogitsProcessor plugin {entrypoint}" + ) from e + if not (isinstance(cls, type) and issubclass(cls, LogitsProcessor)): + raise ValueError( + f"LogitsProcessor plugin {entrypoint.name} is not a subclass of " + "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor; " + "V1-interface plugins are not supported by the V2 model runner." + ) + classes.append(cls) + return classes + + +def _load_v2_logitsprocs_by_fqcns( + logits_processors: Sequence[str | type], +) -> list[type[LogitsProcessor]]: + """Resolve a mixed list of processor types and FQCN strings into types. + + FQCN syntax is : i.e. x.y.z:CustomLogitProc. + """ + from vllm.utils.torch_utils import guard_cuda_initialization + + classes: list[type[LogitsProcessor]] = [] + for ldx, logitproc in enumerate(logits_processors): + if isinstance(logitproc, type): + logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__) + if not issubclass(logitproc, LogitsProcessor): + raise ValueError( + f"{logitproc.__name__} is not a subclass of " + "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor" + ) + classes.append(logitproc) + continue + + logger.debug("- Loading logits processor %s", logitproc) + parts = logitproc.split(":") + if len(parts) != 2: + raise ValueError( + f"Invalid logits processor FQCN {logitproc!r}. " + "Expected format: ':'" + ) + module_path, qualname = parts + + try: + with guard_cuda_initialization(): + module = importlib.import_module(module_path) + except Exception as e: + logger.error( + "Failed to load %sth LogitsProcessor plugin %s: %s", + ldx, + logitproc, + e, + ) + raise RuntimeError( + f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}" + ) from e + + obj = module + for attr in qualname.split("."): + obj = getattr(obj, attr) + if not isinstance(obj, type): + raise ValueError("Loaded logit processor must be a type.") + if not issubclass(obj, LogitsProcessor): + raise ValueError( + f"{obj.__name__} is not a subclass of " + "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor" + ) + classes.append(obj) + + return classes + + +def _load_v2_logitsprocs( + logits_processors: Sequence[str | type] | None, +) -> list[type[LogitsProcessor]]: + """Load all custom logits processors for the V2 model runner. + + Combines installed plugins (validated as V2 processors) with the + user-specified list of processor types and FQCN strings. + """ + return _load_v2_logitsprocs_plugins() + _load_v2_logitsprocs_by_fqcns( + logits_processors or () + ) + + +@lru_cache +def _cached_load_v2_logitsprocs( + custom_logitsprocs: tuple[str | type, ...], +) -> list[type[LogitsProcessor]]: + return _load_v2_logitsprocs(list(custom_logitsprocs)) + + +def build_custom_logits_processors( + vllm_config: "VllmConfig", + req_states: "RequestState", + is_pooling_model: bool, + custom_logitsprocs: Sequence[str | type] = (), +) -> list[LogitsProcessor]: + """Load and instantiate custom logits processors, entrypoint plugins first. + + Raises: + ValueError: if a pooling model specifies custom processors, or a + loaded class does not implement the V2 interface. + RuntimeError: if an FQCN fails to import. + + """ + from vllm.v1.sample.logits_processor import STR_POOLING_REJECTS_LOGITSPROCS + + if is_pooling_model: + if custom_logitsprocs: + raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS) + logger.debug( + "Skipping logits processor loading because pooling models" + " do not support logits processors." + ) + return [] + + custom_logitsprocs_classes = _load_v2_logitsprocs(custom_logitsprocs) + lp_req_state = LogitsProcRequestState.from_request_state(req_states) + return [ctor(vllm_config, lp_req_state) for ctor in custom_logitsprocs_classes] + + +def build_custom_logits_processors_params_validator( + custom_logitsprocs: Sequence[str | type] | None, +) -> "Callable[[SamplingParams], None]": + """Load custom processor classes once and return a params validator. + + Called from the frontend at startup. The returned callable runs each + processor's ``validate_params`` at request admission. + + Raises: + ValueError: if a loaded class does not implement the V2 interface. + RuntimeError: if an FQCN fails to import. + + """ + if not custom_logitsprocs: + return lambda _: None + + classes = _cached_load_v2_logitsprocs(tuple(custom_logitsprocs)) + + def validate_params(sampling_params: "SamplingParams") -> None: + for cls in classes: + try: + cls.validate_params(sampling_params) + except ValueError as e: + raise VLLMValidationError(str(e)) from e + + return validate_params diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index 25cb2f211d94..eb751e8be355 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import numpy as np import torch @@ -8,16 +10,22 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor -from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsContext, + LogitsProcessor, + LogitsProcRequestState, +) +if TYPE_CHECKING: + from vllm.config import VllmConfig -class PenaltiesState: - def __init__(self, req_states: RequestState): - self.req_states = req_states - max_num_reqs = req_states.max_num_reqs - self.vocab_size = req_states.vocab_size +class PenaltiesState(LogitsProcessor): + def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): + self.req_states = req_states self.device = req_states.device + max_num_reqs = req_states.max_num_reqs + vocab_size = req_states.vocab_size self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32) @@ -30,20 +38,17 @@ def __init__(self, req_states: RequestState): # Statistics for penalties. self.prompt_bin_mask = torch.zeros( - max_num_reqs, - cdiv(self.vocab_size, 32), - dtype=torch.int32, - device=self.device, + max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32, device=self.device ) # TODO(woosuk): This tensor is rarely used but can be very large, taking up # GBs of GPU memory. Optimize the memory usage. self.output_bin_counts = torch.zeros( - max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device + max_num_reqs, vocab_size, dtype=torch.int32, device=self.device ) self._new_penalties_reqs: list[int] = [] - def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty self.presence_penalty.np[req_idx] = sampling_params.presence_penalty @@ -52,13 +57,12 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: self.use_penalty[req_idx] = do_penalty if do_penalty: self._new_penalties_reqs.append(req_idx) + return do_penalty def apply_staged_writes(self) -> None: if self._new_penalties_reqs: idx_mapping = async_tensor_h2d( - self._new_penalties_reqs, - dtype=torch.int32, - device=self.device, + self._new_penalties_reqs, dtype=torch.int32, device=self.device ) prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs] @@ -78,29 +82,23 @@ def apply_staged_writes(self) -> None: self.frequency_penalty.copy_to_uva() self.presence_penalty.copy_to_uva() - def apply_penalties( - self, - logits: torch.Tensor, - expanded_idx_mapping: torch.Tensor, - idx_mapping_np: np.ndarray, - input_ids: torch.Tensor, - expanded_local_pos: torch.Tensor, - ) -> None: - if not np.any(self.use_penalty[idx_mapping_np]): + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: + if not np.any(self.use_penalty[ctx.idx_mapping_np]): # No request uses penalties. Skip the kernel launch. - return + return logits apply_penalties( logits, - expanded_idx_mapping, - input_ids, - expanded_local_pos, + ctx.expanded_idx_mapping, + ctx.input_ids, + ctx.expanded_local_pos, self.repetition_penalty.gpu, self.frequency_penalty.gpu, self.presence_penalty.gpu, self.prompt_bin_mask, self.output_bin_counts, ) + return logits @triton.jit diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index b069604c2964..92e9ecee4d66 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence + import numpy as np import torch import vllm.envs as envs +from vllm.config import VllmConfig from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode -from vllm.config.reasoning import ReasoningConfig from vllm.sampling_params import SamplingParams from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p, @@ -18,6 +20,11 @@ from vllm.v1.worker.gpu.sample.bad_words import BadWordsState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState +from vllm.v1.worker.gpu.sample.logits_processor.interface import ( + LogitsContext, + LogitsProcessor, + LogitsProcRequestState, +) from vllm.v1.worker.gpu.sample.logprob import ( LogprobTokenIdsState, compute_topk_scores, @@ -33,6 +40,7 @@ class Sampler: def __init__( self, + vllm_config: VllmConfig, max_num_reqs: int, vocab_size: int, device: torch.device, @@ -41,8 +49,8 @@ def __init__( num_speculative_tokens: int = 1, use_fp64_gumbel: bool = False, enable_trace_replay: bool = False, - reasoning_config: ReasoningConfig | None = None, return_sampling_mask: bool = False, + custom_logits_processors: Sequence[LogitsProcessor] = (), ): self.logprobs_mode = logprobs_mode self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default. @@ -50,11 +58,25 @@ def __init__( self.req_states = req_states self.sampling_states = SamplingStates(max_num_reqs, vocab_size) - self.penalties_state = PenaltiesState(req_states) - self.logit_bias_state = LogitBiasState(max_num_reqs, device) - self.bad_words_state = BadWordsState(req_states) + + lp_req_state = LogitsProcRequestState.from_request_state(req_states) + self.penalties_state = PenaltiesState(vllm_config, lp_req_state) + logit_bias_state = LogitBiasState(vllm_config, lp_req_state) + bad_words_state = BadWordsState(vllm_config, lp_req_state) + + # List order is pipeline order: bias adds, penalties scale, so the + # two do not commute. + self.logits_processors: list[LogitsProcessor] = [ + logit_bias_state, + self.penalties_state, + bad_words_state, + *custom_logits_processors, + ] + self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) - self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config) + self.thinking_budget_state = ThinkingBudgetState( + req_states, vllm_config.reasoning_config + ) self.trace_replay_state = ( TraceReplayState(req_states) if enable_trace_replay else None ) @@ -65,41 +87,25 @@ def __init__( not return_sampling_mask and flashinfer_sampler_supported() ) - def add_request( - self, req_idx: int, prompt_len: int, sampling_params: SamplingParams - ) -> None: - self.sampling_states.add_request(req_idx, sampling_params) - self.penalties_state.add_request(req_idx, sampling_params) - self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params) - self.bad_words_state.add_request(req_idx, sampling_params) + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + needs_processing = self.sampling_states.add_request(req_idx, sampling_params) + needs_processing |= self.thinking_budget_state.add_request( + req_idx, sampling_params + ) + for processor in self.logits_processors: + needs_processing |= processor.add_request(req_idx, sampling_params) + self.needs_logits_processing[req_idx] = needs_processing + self.logprob_token_ids_state.add_request(req_idx, sampling_params) - self.thinking_budget_state.add_request(req_idx, sampling_params) if self.trace_replay_state is not None: self.trace_replay_state.add_request(req_idx, sampling_params) - states = self.sampling_states - temperature = states.temperature.np[req_idx] - self.needs_logits_processing[req_idx] = ( - self.logit_bias_state.use_logit_bias[req_idx] - or self.penalties_state.use_penalty[req_idx] - or self.bad_words_state.num_bad_words.np[req_idx] > 0 - or ( - self.thinking_budget_state.enabled - and self.thinking_budget_state.use_thinking_budget[req_idx] - ) - or (temperature != 0.0 and temperature != 1.0) - or states.min_p.np[req_idx] != 0.0 - or states.top_k.np[req_idx] != states.vocab_size - or states.top_p.np[req_idx] != 1.0 - ) - def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() - self.penalties_state.apply_staged_writes() - self.logit_bias_state.apply_staged_writes() - self.bad_words_state.apply_staged_writes() - self.logprob_token_ids_state.apply_staged_writes() + for processor in self.logits_processors: + processor.apply_staged_writes() self.thinking_budget_state.apply_staged_writes() + self.logprob_token_ids_state.apply_staged_writes() if self.trace_replay_state is not None: self.trace_replay_state.apply_staged_writes() @@ -131,6 +137,7 @@ def __call__( expanded_local_pos = input_batch.expanded_local_pos pos = input_batch.positions[input_batch.logits_indices] input_ids = input_batch.input_ids[input_batch.logits_indices] + seq_lens_upper_bound_np = input_batch.seq_lens_cpu_upper_bound.numpy() # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear # that num_nans is computed before applying penalties and temperature. @@ -146,6 +153,7 @@ def __call__( pos, input_ids, expanded_local_pos, + seq_lens_upper_bound_np, return_logprobs=logprobs_dims is not None, ) @@ -215,6 +223,7 @@ def apply_sampling_params( pos: torch.Tensor, input_ids: torch.Tensor, expanded_local_pos: torch.Tensor, + seq_lens_upper_bound_np: np.ndarray, skip_top_k_top_p: bool = False, ) -> torch.Tensor: if not np.any(self.needs_logits_processing[idx_mapping_np]): @@ -223,39 +232,23 @@ def apply_sampling_params( # Copy logits to a new FP32 tensor. logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) - # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place. - self.logit_bias_state.apply_logit_bias( - logits, expanded_idx_mapping, idx_mapping_np, pos + ctx = LogitsContext( + expanded_idx_mapping=expanded_idx_mapping, + idx_mapping=idx_mapping, + idx_mapping_np=idx_mapping_np, + expanded_local_pos=expanded_local_pos, + input_ids=input_ids, + pos=pos, + seq_lens_upper_bound_np=seq_lens_upper_bound_np, ) - # Apply penalties in place. - self.penalties_state.apply_penalties( - logits, - expanded_idx_mapping, - idx_mapping_np, - input_ids, - expanded_local_pos, - ) - - # Apply bad words masking in place. - self.bad_words_state.apply_bad_words( - logits, - expanded_idx_mapping, - idx_mapping_np, - input_ids, - expanded_local_pos, - ) + # Apply logits processors (native + any custom). + for processor in self.logits_processors: + logits = processor.apply(logits, ctx) - # Force the reasoning end marker once a request's thinking budget is - # reached; applied before temperature so the forced token is always kept. - self.thinking_budget_state.apply( - logits, - expanded_idx_mapping, - idx_mapping, - idx_mapping_np, - input_ids, - expanded_local_pos, - ) + # Forcing runs last so no stage can overwrite the forced end marker + # or weaken it by scaling. + self.thinking_budget_state.apply(logits, ctx) # Apply temperature in place. self.sampling_states.apply_temperature( @@ -282,6 +275,7 @@ def sample( pos: torch.Tensor, input_ids: torch.Tensor, expanded_local_pos: torch.Tensor, + seq_lens_upper_bound_np: np.ndarray, return_logprobs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: processed_logits = self.apply_sampling_params( @@ -292,6 +286,7 @@ def sample( pos, input_ids, expanded_local_pos, + seq_lens_upper_bound_np, skip_top_k_top_p=True, ) top_k, top_p = self.sampling_states.get_top_k_top_p( diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index fe4dee6a6b10..47cf5de856e8 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -38,14 +38,18 @@ def __init__(self, max_num_reqs: int, vocab_size: int): # -1 means no logprobs are requested. self.num_logprobs.fill(NO_LOGPROBS) - def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: - self.temperature.np[req_idx] = sampling_params.temperature - self.top_p.np[req_idx] = sampling_params.top_p + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: + temperature = sampling_params.temperature top_k = sampling_params.top_k + top_p = sampling_params.top_p + min_p = sampling_params.min_p + + self.temperature.np[req_idx] = temperature + self.top_p.np[req_idx] = top_p if top_k <= 0 or top_k > self.vocab_size: top_k = self.vocab_size self.top_k.np[req_idx] = top_k - self.min_p.np[req_idx] = sampling_params.min_p + self.min_p.np[req_idx] = min_p seed = sampling_params.seed self.seeds_set[req_idx] = seed is not None @@ -60,6 +64,13 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: num_logprobs = self.vocab_size self.num_logprobs[req_idx] = num_logprobs + return temperature != 0.0 and ( + temperature != 1.0 + or min_p != 0.0 + or top_k != self.vocab_size + or top_p != 1.0 + ) + def apply_staged_writes(self) -> None: self.temperature.copy_to_uva() self.top_p.copy_to_uva() diff --git a/vllm/v1/worker/gpu/sample/thinking_budget.py b/vllm/v1/worker/gpu/sample/thinking_budget.py index eca792986027..48c960cbeb68 100644 --- a/vllm/v1/worker/gpu/sample/thinking_budget.py +++ b/vllm/v1/worker/gpu/sample/thinking_budget.py @@ -9,6 +9,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor +from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext from vllm.v1.worker.gpu.states import RequestState if TYPE_CHECKING: @@ -78,11 +79,12 @@ def __init__( end_ids, dtype=torch.int32, device=self.device ) - def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: if not self.enabled: - return + return False budget = sampling_params.thinking_token_budget - self.use_thinking_budget[req_idx] = budget is not None + use_thinking_budget = budget is not None + self.use_thinking_budget[req_idx] = use_thinking_budget if budget is None: budget = -1 else: @@ -91,6 +93,7 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: if self.thinking_token_budget.np[req_idx] != budget: self.thinking_token_budget.np[req_idx] = budget self._budget_dirty = True + return use_thinking_budget def apply_staged_writes(self) -> None: if not self.enabled: @@ -107,27 +110,20 @@ def apply_staged_writes(self) -> None: self.thinking_token_budget.copy_to_uva() self._budget_dirty = False - def apply( - self, - logits: torch.Tensor, - expanded_idx_mapping: torch.Tensor, - idx_mapping: torch.Tensor, - idx_mapping_np: np.ndarray, - input_ids: torch.Tensor, - expanded_local_pos: torch.Tensor, - ) -> None: + def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> None: + idx_mapping_np = ctx.idx_mapping_np if not self.enabled or not np.any(self.use_thinking_budget[idx_mapping_np]): return apply_thinking_budget( logits, - idx_mapping, - expanded_idx_mapping, + ctx.idx_mapping, + ctx.expanded_idx_mapping, self.thinking_token_budget.gpu, self.req_states.all_token_ids.gpu, self.req_states.total_len.gpu, - input_ids, - expanded_local_pos, + ctx.input_ids, + ctx.expanded_local_pos, self.cached_last_start, self.cached_last_end, self.cached_scan_pos, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 609924e87fd2..427c93b73cff 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -176,6 +176,7 @@ def _verify( idx_mapping_np: np.ndarray, expanded_idx_mapping: torch.Tensor, expanded_local_pos: torch.Tensor, + seq_lens_upper_bound_np: np.ndarray, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: processed_logits = self.sampler.apply_sampling_params( logits, @@ -185,6 +186,7 @@ def _verify( pos, draft_sampled, expanded_local_pos, + seq_lens_upper_bound_np, ) sampled, num_sampled = rejection_sample( processed_logits, @@ -251,6 +253,7 @@ def _verify_in_chunks( input_batch.idx_mapping_np[start:end], input_batch.expanded_idx_mapping[lo:hi], input_batch.expanded_local_pos[lo:hi], + input_batch.seq_lens_cpu_upper_bound.numpy()[start:end], ) chunk_logprobs = self._get_logprobs_tensors( sampled, From 133b71e0beec3c7bcaae5b0839a03e7e88d20bc7 Mon Sep 17 00:00:00 2001 From: Pavel Zakharov Date: Sat, 19 Sep 2026 13:14:29 -0700 Subject: [PATCH 517/774] [Bugfix] Take padded path for ragged decode batches in sparse_attn_indexer (#52500) Signed-off-by: pavelzak Co-authored-by: Claude Fable 5 Co-authored-by: Jee Jee Li --- .../layers/sparse_attn_indexer.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 7dde4db00fa8..ae45b77dcea8 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -642,6 +642,19 @@ def sparse_attn_indexer( assert decode_metadata is not None kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) decode_lens = decode_metadata.decode_lens + # requires_padding can be computed False for ragged warmup/mixed + # batches (seen on SM120 TP=2: 8 tokens over 6 seqs -> uniform + # reshape below would crash). The token-count/batch-size divisibility + # check is a necessary-but-not-sufficient proxy for "uniform batch" + # (e.g. decode_lens=[1, 2, 3] divides evenly by 3 but isn't uniform); + # it catches the specific case seen in practice without requiring a + # device sync to inspect decode_lens directly, which would break + # CUDA-graph-capture safety on this path. Shared by both the pack + # branch below and the matching unpack branch further down — they + # must agree on which path was taken. + needs_padded_path = decode_metadata.requires_padding or ( + num_decode_tokens % decode_lens.shape[0] != 0 + ) if num_decode_tokens == 0: padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:]) padded_q_scale = ( @@ -649,7 +662,7 @@ def sparse_attn_indexer( if q_scale is not None else None ) - elif decode_metadata.requires_padding: + elif needs_padded_path: # pad in edge case where we have short chunked prefill length < # decode_threshold since we unstrictly split # prefill and decode by decode_threshold @@ -773,7 +786,7 @@ def sparse_attn_indexer( cp_kv_cache_interleave_size, ) - if decode_metadata.requires_padding: + if needs_padded_path: # if padded, we need to unpack # the topk indices removing padded tokens topk_indices = unpack_seq_triton( From 36fa72d2d0d2f86c7c83e1e99c9012b7bd26463b Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:57:01 -0400 Subject: [PATCH 518/774] [GLM5.3 Perf] Size the GLM-5 sparse indexer decode workspace, 3072 MiB GPU memory saved (#57701) Signed-off-by: yewentao256 --- vllm/models/glm5next/amd/sparse_indexer.py | 16 ++++++++-------- vllm/models/glm5next/common/attention.py | 4 ++-- vllm/models/glm5next/nvidia/sparse_indexer.py | 14 +++++++------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index c9b68e973f84..761f79d95de7 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -100,7 +100,7 @@ def sparse_attn_indexer_kpool( scale_fmt: str | None, topk_tokens: int, head_dim: int, - max_model_len: int, + max_pool_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, @@ -136,7 +136,7 @@ def sparse_attn_indexer_kpool( ) # Reserve profiler-visible memory for the worst-case decode logits, - # whose shape is [B * next_n, max_model_len]. This profiling branch + # whose shape is [B * next_n, max_pool_len]. This profiling branch # returns before invoking the logits kernel itself. cfg = get_current_vllm_config_or_none() worst_decode_tokens = 0 @@ -152,7 +152,7 @@ def sparse_attn_indexer_kpool( sched.max_num_batched_tokens, ) # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes. - decode_logits_elems = worst_decode_tokens * max_model_len * 4 + decode_logits_elems = worst_decode_tokens * max_pool_len * 4 prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 max_logits_elems = max(decode_logits_elems, prefill_cap_elems) _ = torch.empty( @@ -566,7 +566,7 @@ def sparse_attn_indexer_kpool( seq_lens, decode_metadata.block_table, decode_metadata.schedule_metadata, - max_model_len=max_model_len, + max_model_len=max_pool_len, ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools, @@ -645,7 +645,7 @@ def __init__( scale_fmt: str, topk_tokens: int, head_dim: int, - max_model_len: int, + max_pool_len: int, max_total_seq_len: int, topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool = False, @@ -659,7 +659,7 @@ def __init__( self.scale_fmt = scale_fmt self.topk_tokens = topk_tokens self.head_dim = head_dim - self.max_model_len = max_model_len + self.max_pool_len = max_pool_len self.max_total_seq_len = max_total_seq_len self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert @@ -698,7 +698,7 @@ def forward_hip( self.scale_fmt, self.topk_tokens, self.head_dim, - self.max_model_len, + self.max_pool_len, self.max_total_seq_len, self.topk_indices_buffer, skip_k_cache_insert=self.skip_k_cache_insert, @@ -715,7 +715,7 @@ def forward_hip( self.scale_fmt, self.topk_tokens, self.head_dim, - self.max_model_len, + self.max_pool_len, self.max_total_seq_len, self.topk_indices_buffer, self.skip_k_cache_insert, diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index 21493174e8ba..4f776e8a8e76 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -296,7 +296,7 @@ def __init__( cache_config=cache_config, index_kpool=self.index_kpool, ) - self.max_model_len = vllm_config.model_config.max_model_len + self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool self.prefix = prefix from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size @@ -307,7 +307,7 @@ def __init__( self.scale_fmt, self.topk_tokens, self.head_dim, - self.max_model_len, + self.max_pool_len, self.max_total_seq_len, self.topk_indices_buffer, tail_cache=self.tail_cache, diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 9dab659d4cc0..f7204949fc5c 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -117,7 +117,7 @@ def sparse_attn_indexer_kpool( scale_fmt: str | None, topk_tokens: int, head_dim: int, - max_model_len: int, + max_pool_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, @@ -153,7 +153,7 @@ def sparse_attn_indexer_kpool( ) # Reserve profiler-visible memory for the worst-case decode logits, - # whose shape is [B * next_n, max_model_len]. This profiling branch + # whose shape is [B * next_n, max_pool_len]. This profiling branch # returns before invoking the logits kernel itself. cfg = get_current_vllm_config_or_none() worst_decode_tokens = 0 @@ -169,7 +169,7 @@ def sparse_attn_indexer_kpool( sched.max_num_batched_tokens, ) # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes. - decode_logits_elems = worst_decode_tokens * max_model_len * 4 + decode_logits_elems = worst_decode_tokens * max_pool_len * 4 prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 max_logits_elems = max(decode_logits_elems, prefill_cap_elems) _ = torch.empty( @@ -569,7 +569,7 @@ def sparse_attn_indexer_kpool( seq_lens, decode_metadata.block_table, decode_metadata.schedule_metadata, - max_model_len=max_model_len, + max_model_len=max_pool_len, clean_logits=False, ) num_rows = logits.shape[0] @@ -668,7 +668,7 @@ def __init__( scale_fmt: str, topk_tokens: int, head_dim: int, - max_model_len: int, + max_pool_len: int, max_total_seq_len: int, topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool = False, @@ -682,7 +682,7 @@ def __init__( self.scale_fmt = scale_fmt self.topk_tokens = topk_tokens self.head_dim = head_dim - self.max_model_len = max_model_len + self.max_pool_len = max_pool_len self.max_total_seq_len = max_total_seq_len self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert @@ -755,7 +755,7 @@ def forward_cuda( self.scale_fmt, self.topk_tokens, self.head_dim, - self.max_model_len, + self.max_pool_len, self.max_total_seq_len, self.topk_indices_buffer, self.skip_k_cache_insert, From 674b6d95d6d0bd63b44cae0abe2bf304328c73e9 Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Sat, 19 Sep 2026 14:42:21 -0700 Subject: [PATCH 519/774] [CI][Bugfix] Fix MoE reprocess test mock after #57405's kernel refactor (#57641) Signed-off-by: kimi-no-na-wa Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: kimi-no-na-wa Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- tests/kernels/moe/test_flashinfer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index d17416de99b8..9189c5c6415a 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing from dataclasses import dataclass +from unittest.mock import MagicMock import pytest import torch @@ -537,10 +538,11 @@ def _make_unquantized_flashinfer_test_layer( method.moe = moe_config method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM method.moe_kernel = None + mock_kernel = MagicMock() monkeypatch.setattr( method, "_init_moe_kernel", - lambda _: setattr(method, "moe_kernel", object()), + lambda _: setattr(method, "moe_kernel", mock_kernel), ) layer = torch.nn.Module() From a7fda4c88bfc421d31e33acc5e01e86ebe467ad8 Mon Sep 17 00:00:00 2001 From: Ariel Bereslavsky Date: Sun, 20 Sep 2026 03:15:07 +0300 Subject: [PATCH 520/774] [Bugfix] Fix SM100 fp8_ds_mla cache scales (#49435) Signed-off-by: Ariel Bereslavsky Signed-off-by: Ariel Bereslavsky Signed-off-by: Yongye Zhu Co-authored-by: Ariel Bereslavsky Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: Yongye Zhu Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- csrc/libtorch_stable/cache_kernels.cu | 6 +-- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 19 ++++----- tests/kernels/attention/test_cache.py | 10 ++--- .../kernels/attention/test_flashmla_sparse.py | 39 +++++++++++++++---- .../test_kimi_k3_mla_key_concat_kv_cache.py | 18 ++++++++- .../test_fused_deepseek_v32_norm_rope.py | 9 +++-- vllm/models/deepseek_v32/common/kernels.py | 7 ++-- .../ops/fused_mla_key_concat_kv_cache.py | 6 +-- 8 files changed, 76 insertions(+), 38 deletions(-) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index c5e35c52b82f..4b8f4c23a523 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -19,7 +19,6 @@ #include #include -#include #include #ifdef USE_ROCM @@ -589,8 +588,9 @@ __global__ void concat_and_cache_ds_mla_kernel( max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16)); } - // Compute the scale for the tile - float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + // Both SM90 and SM100 readers preserve power-of-two fp32 scales exactly. + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f); + tile_scale = exp2f(ceilf(log2f(tile_scale))); // The first lane of each half-warp writes the scale to kv_cache if ((lane_idx == 0) || (lane_idx == 16)) { diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 05c7bfc6d13e..3892790c4040 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -35,8 +35,9 @@ * * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): * - full key concat (bf16), and cache insert in DeepSeek's 656-byte - * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic - * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel. + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile + * power-of-two scales, RoPE bf16), bit-compatible with + * concat_and_cache_ds_mla_kernel. * * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the * producing GEMMs on sm_90+, and are structured after @@ -68,7 +69,6 @@ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #endif #include -#include #include #ifdef USE_ROCM @@ -111,9 +111,9 @@ constexpr float kFp8Max = 224.0f; #else constexpr float kFp8Max = 448.0f; #endif -// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu). -// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32 -// tile scales, [528,656) RoPE 64 bf16. +// Divisor for fp8_ds_mla per-tile power-of-two scales (matches +// cache_kernels.cu). fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), +// [512,528) 4 fp32 tile scales, [528,656) RoPE 64 bf16. constexpr float kFp8ScaleDivisor = kFp8Max; // Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a @@ -325,7 +325,7 @@ __device__ __forceinline__ void writeLatent576( } // Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512 -// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at +// as fp8 in 4 tiles of 128 (per-tile power-of-two scale, 4 fp32 scales at // [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with // concat_and_cache_ds_mla_kernel. template @@ -349,7 +349,8 @@ __device__ __forceinline__ void writeDsMlaCache( for (int offset = 4; offset > 0; offset /= 2) { max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8)); } - float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f); + tile_scale = exp2f(ceilf(log2f(tile_scale))); if ((laneId & 7) == 0) { reinterpret_cast(row)[kKvLoraRank / 4 + tile] = tile_scale; } @@ -598,7 +599,7 @@ __global__ void fusedKimiK3MLAKVConcatPackKernel( // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert // // Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel: -// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, power-of-two scaled // [512, 528) 4 fp32 per-tile scales // [528, 656) RoPE 64 vals as bf16 (unquantized) // The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16) diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 0d5066d2f112..41f0a1138ba0 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -960,14 +960,14 @@ def test_concat_and_cache_ds_mla( tile_end = (tile_idx + 1) * 128 tile_data[:] = kv_c_data[tile_start:tile_end] - # tile_scale = tile_data.amax().to(torch.float32) / 448. - # NOTE: Using torch's amax() gives different results, - # so this must be manually computed. + # Using torch's amax() gives different results, so this must be + # manually computed. tile_data_float = tile_data.to(torch.float32) manual_max = abs(tile_data_float[0]) for j in range(1, 128): manual_max = max(manual_max, abs(tile_data_float[j])) - tile_scale = manual_max / 448.0 + raw_scale = torch.clamp(manual_max / 448.0, min=1e-4) + tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale))) ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale @@ -1008,7 +1008,7 @@ def test_concat_and_cache_ds_mla( ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :] torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1) - torch.testing.assert_close(kv_scales, ref_scales, atol=0.001, rtol=0.1) + torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0) torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index fb93557a32c0..bfcd6ae10ab3 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -109,8 +109,9 @@ def test_sparse_flashmla_metadata_smoke(): assert num_splits is None -def test_sparse_flashmla_decode_smoke(): +def test_sparse_flashmla_decode_matches_cache_writer_scales(): import vllm.v1.attention.ops.flashmla as fm + from vllm import _custom_ops as ops ok, reason = fm.is_flashmla_sparse_supported() if not ok: @@ -130,7 +131,7 @@ def test_sparse_flashmla_decode_smoke(): # Metadata q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k # q_heads_per_hk = num_heads_q // num_heads_k - cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device) + cache_seqlens = torch.ones(batch_size, dtype=torch.int32, device=device) tile_md, num_splits = fm.get_mla_metadata( cache_seqlens, q_seq_per_hk, @@ -146,14 +147,33 @@ def test_sparse_flashmla_decode_smoke(): dtype=torch.bfloat16, device=device, ) - k_cache = torch.zeros( - (1, page_block_size, num_heads_k, bytes_per_token), - dtype=torch.uint8, - device=device, + cache_rows = torch.zeros( + (1, page_block_size, bytes_per_token), dtype=torch.uint8, device=device ) - indices = torch.zeros( - (batch_size, seqlen_q, topk), dtype=torch.int32, device=device + # 336 / 448 = 0.75 distinguishes arbitrary FP32 from power-of-two scales. + kv_c = torch.full((1, 512), 336.0, dtype=torch.bfloat16, device=device) + kv_c[:, -128:] = 0 + ops.concat_and_cache_mla( + kv_c, + torch.zeros((1, 64), dtype=torch.bfloat16, device=device), + cache_rows, + torch.zeros(1, dtype=torch.int64, device=device), + "fp8_ds_mla", + torch.ones(1, dtype=torch.float32, device=device), ) + scales = cache_rows[0, 0].view(torch.float32)[128:132] + expected_scales = torch.tensor([1.0, 1.0, 1.0, 2**-13], device=device) + torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0) + cached_nope = cache_rows[0, 0, :512].view(torch.float8_e4m3fn).float() + expected = (cached_nope * scales.repeat_interleave(128)).to(torch.bfloat16) + + k_cache = cache_rows.unsqueeze(2) + indices = torch.full( + (batch_size, seqlen_q, topk), -1, dtype=torch.int32, device=device + ) + indices[..., 0] = 0 + # SM90 sparse decode only supports a fixed top-k width. Padding with -1 + # exercises the same cache data on SM90 and SM100 without topk_length. block_table = torch.zeros((batch_size, 128), dtype=torch.int32, device=device) out, lse = fm.flash_mla_with_kvcache( @@ -170,6 +190,9 @@ def test_sparse_flashmla_decode_smoke(): assert out.shape[0] == batch_size assert out.shape[-1] == head_dim_v assert lse.shape[0] == batch_size + torch.testing.assert_close( + out, expected.view(1, 1, 1, 512).expand_as(out), rtol=0, atol=0 + ) @pytest.mark.parametrize("h_q", [64, 128]) diff --git a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py index e1582b7621a4..af41722605c6 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py +++ b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py @@ -222,7 +222,8 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache) -def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None: +@pytest.mark.parametrize("decode", [False, True]) +def test_ds_mla_cache_insert_bit_compatible_with_reference(decode: bool) -> None: """Fused ds_mla insertion must match the reference bit-for-bit.""" torch.manual_seed(0) num_tokens, num_blocks = 33, 16 @@ -242,5 +243,18 @@ def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None: got = ref.clone() scale = torch.ones(1, device="cuda", dtype=torch.float32) ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale) - fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots) + if decode: + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA_RANK, device="cuda", dtype=dt + ) + fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q[..., -ROPE_HEAD_DIM:], kv_c, k_pe, got, slots, ds_mla=True + ) + else: + fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots) assert torch.equal(ref, got) + rows = got.view(-1, DS_MLA_CACHE_ENTRY)[slots[slots >= 0]] + scales = rows[:, KV_LORA_RANK : KV_LORA_RANK + 16].view(torch.float32) + torch.testing.assert_close( + torch.log2(scales), torch.log2(scales).round(), rtol=0, atol=0 + ) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 7f74cde05d5b..61f977ddbdea 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -507,8 +507,8 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool): def test_fused_norm_rope_ds_mla(num_tokens: int): """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100). - Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32 - scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE. + Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales + stored as float32) | 4 float32 scales | 64 bf16 (unquantized) RoPE. """ torch.manual_seed(5) dev = "cuda" @@ -557,7 +557,8 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] tiles = kv_ref.view(num_tokens, 4, 128) - ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38) + ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4) + ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale))) ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8) cache = mla_cache[0, :num_tokens] # [N, 656] uint8 @@ -566,7 +567,7 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): rope_off = KV_LORA // 2 + 8 rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM] - torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6) + torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0) assert_fp8(nope, ref_nope, "ds_mla NoPE fp8") assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16") # No indexer on this call: top-k buffer must be untouched. diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index fd02a51076ed..1808d12f64af 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -315,7 +315,7 @@ def _fused_norm_rope_kernel( if MLA_CACHE_DS_MLA: # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element # tile of the NoPE is dynamically quantized to fp8 with its own - # float32 scale; the RoPE tail is stored unquantized in bf16. + # power-of-two scale stored as float32; the RoPE tail is bf16. # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE @@ -327,9 +327,8 @@ def _fused_norm_rope_kernel( ) kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM)) tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) - # scale = amax / 448 (fp8 e4m3 max), matching the reference - # concat_and_cache_ds_mla kernel; floored to FLT_MIN. - tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38) + tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1e-4) + tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale))) kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,)) tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8) tile_off = tl.arange(0, MLA_NUM_TILES) diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 24ae1a1a820c..79de3af0bc22 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -76,9 +76,9 @@ def fused_mla_key_concat_ds_mla_insert( """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout. The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of - 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no - scale argument. Returns the bf16 full key; optionally rotates ``q`` and - writes ``kv_cache`` in place. + 128 with per-tile power-of-two scales stored as float32, RoPE as bf16) -- + self-scaling, so no scale argument. Returns the bf16 full key; optionally + rotates ``q`` and writes ``kv_cache`` in place. """ k_pe = k_pe.reshape(k_pe.shape[0], -1) tp, num_heads, qk_nope_head_dim = k_nope.shape From 708c2495c0c6d0fc576224ecfc54efacc6414981 Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Sun, 20 Sep 2026 10:16:12 +0800 Subject: [PATCH 521/774] [Test][Core] Compute the expected hybrid prefix-cache hit instead of hard-coding it (#57467) Signed-off-by: Lin, Fanli --- .buildkite/intel_jobs/test-intel.yaml | 2 +- .../test_hybrid_prefix_cache_hit_rate.py | 62 ++++++++++++++++--- 2 files changed, 53 insertions(+), 11 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 458bef9518f2..ef2b4ba07396 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -98,7 +98,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pip install dpctl && - pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py --ignore=v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py && + pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && diff --git a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py index c75591b69724..ff62abab8c86 100644 --- a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py +++ b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py @@ -4,6 +4,7 @@ import asyncio import json +from math import lcm from uuid import uuid4 import httpx @@ -12,6 +13,7 @@ from prometheus_client.parser import text_string_to_metric_families from tests.utils import RemoteOpenAIServer +from vllm.utils.math_utils import round_down MODEL = "Qwen/Qwen3.5-0.8B" NUM_CONVERSATIONS = 8 @@ -34,6 +36,29 @@ def _user_turns(conversation: int) -> list[str]: ] +async def _scheduler_block_size(client: httpx.AsyncClient) -> int: + """Return the coarsest boundary an expected cache hit can land on. + + Mirrors ``scheduler_block_size`` in ``resolve_kv_cache_block_sizes``, which + the engine keeps internal: the LCM of the resolved group block sizes. A + finer ``prefix_match_unit`` only adds hits, so this stays conservative. + """ + response = await client.get("/metrics") + response.raise_for_status() + for family in text_string_to_metric_families(response.text): + for sample in family.samples: + if sample.name != "vllm:cache_config_info": + continue + sizes = [ + int(value) + for name in ("block_size", "mamba_block_size") + if (value := sample.labels.get(name, "None")) != "None" + ] + assert sizes, f"no block size in {sample.labels}" + return lcm(*sizes) + raise AssertionError("missing vllm:cache_config_info metric") + + async def _prefix_cache_counters( client: httpx.AsyncClient, ) -> tuple[float, float]: @@ -83,13 +108,13 @@ def server(request): @pytest.mark.asyncio @pytest.mark.parametrize( - "server, min_hit_rate", - # MTP re-prefills an extra block; both floors reject one more missed block. - [pytest.param(False, 0.90, id="base"), pytest.param(True, 0.75, id="mtp")], + "server, dropped_blocks", + # MTP's prefill lookahead pollutes the tail block, so EAGLE drops one. + [pytest.param(False, 0, id="base"), pytest.param(True, 1, id="mtp")], indirect=["server"], ) async def test_prefix_cache_hit_rate( - server: RemoteOpenAIServer, min_hit_rate: float + server: RemoteOpenAIServer, dropped_blocks: int ) -> None: # Isolate sessions so another conversation cannot hide a cache miss. salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)] @@ -102,6 +127,8 @@ async def test_prefix_cache_hit_rate( server.get_async_client() as client, httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics, ): + block_size = await _scheduler_block_size(metrics) + print(f"{block_size=}") async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]: history = histories[session] @@ -126,15 +153,30 @@ async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]: assert cached is not None, context assert 0 <= cached < prompt, context assert prompt >= MIN_PROMPT_TOKENS, context - if not replay and turn == 0: - assert cached == 0, context - else: - assert cached / prompt >= min_hit_rate, ( - f"{context}, {min_hit_rate=}" - ) if replay: assert prompt == prompt_lengths[session][turn], context + # The identical request computed this whole prompt, but at + # least one token is always recomputed. + computed = prompt - 1 + elif turn: + computed = prompt_lengths[session][turn - 1] else: + assert cached == 0, context + computed = 0 + # Every token below the last reachable boundary must be reused. + floor = max( + 0, round_down(computed, block_size) - dropped_blocks * block_size + ) + print( + f"{session=} {turn=} {replay=} {prompt=} {cached=} {floor=} " + f"hit_rate={cached / prompt:.4f}" + ) + assert cached >= floor, ( + f"{context}, {computed=}, {block_size=}, " + f"hit_rate={cached / prompt:.4f}, " + f"expected_hit_rate={floor / prompt:.4f}" + ) + if not replay: prompt_lengths[session].append(prompt) history.append( { From f4fc5aae7fea790e78bb2b5ed62f636dc20c71d8 Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Sun, 20 Sep 2026 10:52:50 +0800 Subject: [PATCH 522/774] [XPU][UT]Using the device_control_env_var to restrict visible device on different platform (#57581) Signed-off-by: yisheng Co-authored-by: Kunshang Ji --- tests/distributed/test_multiproc_executor.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index d0a417a7886d..cd9413865465 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -15,6 +15,7 @@ from tests.utils import multi_gpu_test from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.executor import multiproc_executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor @@ -378,8 +379,8 @@ def test_multiproc_executor_properties(): def test_multiproc_executor_multi_node(): """Test MultiprocExecutor with multi-node configuration. This simulates 2 nodes with TP=4: - - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2 - - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 + - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2 + - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2 Total world_size = 4, nnodes = 2 """ # Python 3.14+ changed default multiprocessing start method to 'forkserver' @@ -392,15 +393,20 @@ def test_multiproc_executor_multi_node(): # symm_mem does not work for simulating multi instance in single node os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" + # Use the platform-specific device control env var (e.g. + # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each + # node only sees its own devices on different platform. + device_control_env_var = current_platform.device_control_env_var + def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int): """Run a single node's executor.""" executor = None try: - # Set CUDA_VISIBLE_DEVICES for this node + # Restrict visible devices for this node if node_rank == 0: - os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" + os.environ[device_control_env_var] = "0,1" else: - os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" + os.environ[device_control_env_var] = "2,3" # Create config for this node vllm_config = create_vllm_config( From 1c3ef2ad4cabfbc2e5d74494b0b2bf8c3cb3f936 Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Sun, 20 Sep 2026 10:53:00 +0800 Subject: [PATCH 523/774] [XPU][UT]Fix test_mnnvl_alltoall device and distributed backend (#57591) Signed-off-by: yisustc Co-authored-by: Kunshang Ji --- tests/distributed/test_mnnvl_alltoall.py | 9 ++++++--- tests/utils.py | 5 +++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 1cd67524602f..362ffe4b1535 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -14,6 +14,7 @@ import torch.multiprocessing as mp from vllm.distributed import get_ep_group +from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, @@ -23,6 +24,8 @@ from ..utils import init_test_distributed_environment +DEVICE = current_platform.device_type + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -483,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_ep_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 64 tokens_per_rank = 16 @@ -624,7 +627,7 @@ def _two_sided_data_worker(rank, world_size): # Use DP group because MnnvlMoe workspace allocation calls get_dp_group() # internally and requires dp_size == ep_size. cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") num_gpus = torch.accelerator.device_count() hidden_size = 128 @@ -769,7 +772,7 @@ def _one_sided_data_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 256 tokens_per_rank = 32 diff --git a/tests/utils.py b/tests/utils.py index 81f82d98c8a7..7ee0e6680318 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1427,8 +1427,10 @@ def init_test_distributed_environment( get_current_vllm_config_or_none, set_current_vllm_config, ) + from vllm.platforms import current_platform distributed_init_method = f"tcp://localhost:{distributed_init_port}" + backend = current_platform.dist_backend if data_parallel_size > 1: # For DP we need to set a common DP master port @@ -1449,6 +1451,7 @@ def init_test_distributed_environment( rank=rank % tp_pp_world, distributed_init_method=distributed_init_method, local_rank=local_rank if local_rank >= 0 else rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) return @@ -1460,6 +1463,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) else: @@ -1470,6 +1474,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) From 03f83013872015f6addb66954e2752d2d9bccd96 Mon Sep 17 00:00:00 2001 From: Louie Tsai Date: Sat, 19 Sep 2026 20:24:09 -0700 Subject: [PATCH 524/774] Add one-step recipe serving (#57306) Signed-off-by: louie-tsai --- tools/recipes/README.md | 15 +++++++++ tools/recipes/serve_with_recipe.sh | 50 ++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100755 tools/recipes/serve_with_recipe.sh diff --git a/tools/recipes/README.md b/tools/recipes/README.md index 058470167a7b..253bd7a726ee 100644 --- a/tools/recipes/README.md +++ b/tools/recipes/README.md @@ -21,6 +21,21 @@ flowchart LR The recipe is the baseline. Hardware and workload information can optionally refine the initial configuration. Sweep tuning is an optional validation step. +## Getting Started + +Use `serve_with_recipe.sh` to generate the recipe configuration and start vLLM +in one step: + +```bash +tools/recipes/serve_with_recipe.sh \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 +``` + +For `xeon6`, the script enables hardware detection automatically. Generated +`config.yml` and `env.sh` files remain in the directory where the script was +invoked. + ## 1. vLLM Recipes Only Use the converter directly when the recipe already contains the deployment diff --git a/tools/recipes/serve_with_recipe.sh b/tools/recipes/serve_with_recipe.sh new file mode 100755 index 000000000000..cad1c06f3f3b --- /dev/null +++ b/tools/recipes/serve_with_recipe.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +CONVERTER="${SCRIPT_DIR}/recipe_json_to_vllm_config.py" +OUTPUT_DIR="$(pwd -P)" +CONFIG_PATH="${OUTPUT_DIR}/config.yml" +ENV_PATH="${OUTPUT_DIR}/env.sh" + +args=("$@") +hardware="" + +for ((i = 0; i < ${#args[@]}; i++)); do + case "${args[$i]}" in + --hardware) + if ((i + 1 < ${#args[@]})); then + hardware="${args[$((i + 1))]}" + fi + ;; + --hardware=*) + hardware="${args[$i]#*=}" + ;; + esac +done + +if [[ "${hardware,,}" == "xeon6" ]]; then + args+=(--detect-hardware) +fi + +echo "==> Generating vLLM configuration" +python3 "${CONVERTER}" "${args[@]}" + +echo +echo "==> Loading recipe environment" +# shellcheck disable=SC1090 +source "${ENV_PATH}" + +# Released vLLM images keep recipe-relative assets such as examples/*.jinja +# under /vllm-workspace. Launch from there when available while keeping the +# generated config and environment files in the caller's output directory. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +echo +echo "==> Starting vLLM" +exec vllm serve --config "${CONFIG_PATH}" From 1596fa5f8825d1b949463d51f3ed03650b0f539a Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sat, 19 Sep 2026 23:46:33 -0400 Subject: [PATCH 525/774] [ROCm][DSv4.1][Perf] Reuse the decode topk ragged metadata across layers (#57434) Signed-off-by: Fangzhou Ai Co-authored-by: Shanshan Shen <467638484@qq.com> --- vllm/models/deepseek_v41/amd/rocm.py | 79 +++++++++++++++++++++++----- 1 file changed, 66 insertions(+), 13 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 4d838e121e69..b7a67e85bd44 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -13,7 +13,10 @@ ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.attention import ( + DeepseekV4Attention, + _replace_layer_index, +) from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, @@ -391,6 +394,10 @@ def _copy_ragged_to_graph_buffers( return ragged_out, indptr_out +# (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode. +_TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + @dataclass class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None @@ -500,6 +507,23 @@ def __init__(self, *args, **kwargs): self._fused_compressor_weight: torch.Tensor | None self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None + # Decode ragged topk metadata is a pure function of the indices its + # index source published, so every consumer of a source rebuilds the + # same thing. The source owns one cache per compress ratio, refreshed + # when it runs; consumers below it read through. + self._topk_ragged_cache: dict[int, _TopkRagged] = {} + self._index_source_prefix: str | None = None + if self.compress_ratio > 0: + assert self.index_source_layer_id is not None + self._index_source_prefix = _replace_layer_index( + self.prefix, self.index_source_layer_id + ) + if self._index_source_prefix not in self._static_forward_context: + raise NotImplementedError( + f"Index source {self._index_source_prefix} not found on " + "this rank; PP splits inside a v4.1 index-sharing group " + "are not supported." + ) @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -733,6 +757,42 @@ def forward_mqa( output=output[:num_decode_tokens], ) + def _decode_topk_ragged( + self, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + num_decodes: int, + num_decode_tokens: int, + ) -> _TopkRagged: + """Ragged form of the topk indices this layer's index source published. + + The packing depends only on the shared ``topk_indices_buffer`` and + per-step metadata, apart from the layer's own compress ratio, so the + source memoizes one result per ratio for all the layers below it. + """ + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + assert self.topk_indices_buffer is not None + assert self._index_source_prefix is not None + + source = self._static_forward_context[self._index_source_prefix] + if source is self: + # Fresh indices as of this layer; drop what the last step cached. + source._topk_ragged_cache = {} + cached = source._topk_ragged_cache.get(self.compress_ratio) + if cached is not None: + return cached + + built = compute_global_topk_ragged_indices_and_indptr( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + attn_metadata.block_size // self.compress_ratio, + swa_metadata.is_valid_token[:num_decode_tokens], + ) + source._topk_ragged_cache[self.compress_ratio] = built + return built + def _forward_decode( self, q: torch.Tensor, @@ -749,22 +809,15 @@ def _forward_decode( topk_ragged_indices = None topk_ragged_indptr = None if not swa_only: - # Local indices filled by the index-source layer's indexer. - assert attn_metadata is not None - assert swa_metadata.is_valid_token is not None - assert self.topk_indices_buffer is not None - block_size = attn_metadata.block_size // self.compress_ratio - is_valid = swa_metadata.is_valid_token[:num_decode_tokens] ( topk_ragged_indices, topk_ragged_indptr, topk_lens, - ) = compute_global_topk_ragged_indices_and_indptr( - self.topk_indices_buffer[:num_decode_tokens], - swa_metadata.token_to_req_indices, - attn_metadata.block_table[:num_decodes], - block_size, - is_valid, + ) = self._decode_topk_ragged( + swa_metadata=swa_metadata, + attn_metadata=attn_metadata, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, ) rocm_sparse_attn_decode( From 70fc7bd695e6c9161ec7f3ba67b73e243189d571 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Sun, 20 Sep 2026 12:24:41 +0800 Subject: [PATCH 526/774] [Feature][EPD] Support dynamic EPD (EC-connector) proxy (#54176) Signed-off-by: Tianyu Guo Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex --- .../disaggregated_encoder/README.md | 78 ++- .../disagg_1e1p1d_example.sh | 25 +- .../disagg_1e1pd_example.sh | 24 +- .../disaggregated_encoder/disagg_epd_proxy.py | 631 ++++++++++++++++-- tests/v1/ec_connector/unit/test_epd_proxy.py | 279 ++++++++ .../ec_connector/unit/test_epd_proxy_retry.py | 47 +- .../unit/test_epd_proxy_round_robin.py | 2 + .../v1/ec_connector/unit/test_epd_registry.py | 156 +++++ 8 files changed, 1169 insertions(+), 73 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_epd_proxy.py create mode 100644 tests/v1/ec_connector/unit/test_epd_registry.py diff --git a/examples/disaggregated/disaggregated_encoder/README.md b/examples/disaggregated/disaggregated_encoder/README.md index 71a0d0e72ea8..ec10f23fc074 100644 --- a/examples/disaggregated/disaggregated_encoder/README.md +++ b/examples/disaggregated/disaggregated_encoder/README.md @@ -106,6 +106,82 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. | | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). | +### Dynamic registration + +Alternatively, let the external launcher register ready instances over HTTP. +No vLLM configuration changes or worker registration threads are required. +Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration +and removal. These are trusted control-plane APIs; do not expose them publicly. + +```bash +export ADMIN_API_KEY="your-admin-key" +python disagg_epd_proxy.py --port 8000 --dynamic-registration +``` + +Omit the static server URL flags and register every stage. The roles determine +the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` + +`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`, +even when D registers first or all P instances go offline; otherwise requests +return `503`. The proxy rejects mixing combined PD with standalone P/D, +including unhealthy instances still in the registry. Explicitly remove the +old topology's P/D or PD registrations before switching topologies. + +After each instance is ready, the launcher registers its reachable HTTP URL: + +```bash +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"encode","url":"http://e-host:8001"}' +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"prefill_decode","url":"http://pd-host:8002"}' +``` + +For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`. +Example and NIXL EC connectors need no additional registration fields. +For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode` +or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port +layout: supply the TP-rank-0 address for each DP replica, in DP-rank order; +replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers +the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`: + +```json +{ + "role": "prefill_decode", + "url": "http://pd-host:8002", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"] +} +``` + +The proxy selects one consumer replica and uses it for both the encoder push +and the consumer HTTP request. Standalone D does not need EC control addresses. +Port allocation and avoiding collisions remain the launcher's responsibility. + +Inspect or remove instances without restarting the proxy: + +```bash +curl http://proxy-host:8000/instances +curl --fail-with-body -X DELETE \ + 'http://proxy-host:8000/instances?url=http://e-host:8001' \ + -H "X-API-Key: $ADMIN_API_KEY" +``` + +Registration is idempotent. The proxy probes registered instances every +`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe +(default 2). After `--fail-threshold` consecutive failures (default 3), it stops +sending new requests to that instance. Healthy instances rejoin automatically; +unreachable ones are forgotten after `--evicted-ttl` seconds (default 900; +`0` retains them indefinitely). Removal stops new routing; already routed +requests retain their selected endpoints, so drain requests before stopping +the instance. The launcher must re-register instances after a proxy restart. + +The example launch scripts also support this flow: export `ADMIN_API_KEY` and +set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or +`disagg_1e1p1d_example.sh`. Their default remains static routing. + +### Static configuration + The proxy batches images from the same user request assigned to the same encoder. Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each @@ -137,6 +213,6 @@ For E + P + D setup: ```bash $ python disagg_encoder_proxy.py \ --encode-servers-urls "http://e1:8001,http://e2:8001" \ - --prefill-servers-urls "http://p1:8003,http://p2:8004" \ + --prefill-servers-urls "http://p1:8003,http://p2:8004" \ --decode-servers-urls "http://d1:8005,http://d2:8006" ``` diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh index cb51333e908d..57392a0408c8 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh @@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --prefill-servers-urls "http://localhost:$PREFILL_PORT" + --decode-servers-urls "http://localhost:$DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "http://localhost:$PREFILL_PORT" \ - --decode-servers-urls "http://localhost:$DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh index 4bf60ea66d61..44e43b8d7f6e 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh @@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "disable" \ - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 3317b091fb52..abae7eb65da1 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -1,12 +1,13 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""disagg_encoder_proxy.py +"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation. -Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two -clusters: - • encode (multimodal feature extraction) - • decode (language-model inference) +Use static server URLs, or --dynamic-registration for launcher-managed HTTP +registration of E, P/PD and D instances. Both modes route OpenAI-compatible +``/v1/chat/completions`` requests through encoder and inference clusters and +share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming +forwarding. For MM input we: 1. Extract *every* image/audio/video item. @@ -20,22 +21,27 @@ import argparse import asyncio +import enum import hashlib import itertools import json import logging import os import random +import secrets import time import uuid from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field from typing import Any import aiohttp import msgspec import uvicorn -from fastapi import FastAPI, HTTPException, Request +from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import AnyHttpUrl, BaseModel, Field, model_validator ############################################################################### # FastAPI app & global state @@ -44,6 +50,340 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") logger = logging.getLogger("proxy") +DEFAULT_PROBE_INTERVAL = 5.0 +DEFAULT_PROBE_TIMEOUT = 2.0 +DEFAULT_FAIL_THRESHOLD = 3 +# Stop probing an instance that has been down this long. 0 probes forever, +# which is what a cluster that restarts instances in place wants. +DEFAULT_EVICTED_TTL = 900.0 + + +class InstanceRole(str, enum.Enum): + ENCODE = "encode" + PREFILL = "prefill" + DECODE = "decode" + PREFILL_DECODE = "prefill_decode" + + +@dataclass +class InstanceRecord: + """One registered instance. + + Attributes: + role: Which stage this instance serves. + url: Base OpenAI-compatible URL, e.g. ``http://host:8000``. + ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica, + supplied by the launcher from the consumer's fixed port config. + dp_size: Data-parallel replicas behind `url`, so the proxy can pick + a replica and name the same one to both halves of a request. + + """ + + role: InstanceRole + url: str + ec_zmq_addrs: list[str] = field(default_factory=list) + dp_size: int = 1 + + +class InstanceRegistry: + def __init__( + self, + probe_interval: float = DEFAULT_PROBE_INTERVAL, + probe_timeout: float = DEFAULT_PROBE_TIMEOUT, + fail_threshold: int = DEFAULT_FAIL_THRESHOLD, + evicted_ttl: float = DEFAULT_EVICTED_TTL, + ): + self._probe_interval = probe_interval + self._probe_timeout = probe_timeout + self._fail_threshold = fail_threshold + self._evicted_ttl = evicted_ttl + + self._live: dict[str, InstanceRecord] = {} + self._evicted: dict[str, InstanceRecord] = {} + self._evicted_since: dict[str, float] = {} + self._fail_counts: dict[str, int] = {} + # One cursor per role, only ever incremented. Rebuilding it whenever + # the roster changes -- what an `itertools.cycle` over a mutable list + # forces -- restarts every fan-out at the first instance and hot-spots + # it after each registration. + self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole} + self._replica_cursors: dict[str, int] = {} + self._probe_task: asyncio.Task | None = None + + def register(self, record: InstanceRecord) -> bool: + """Add or refresh an instance without overriding its health status.""" + roles = { + other.role + for other in itertools.chain(self._live.values(), self._evicted.values()) + } + split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE} + if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or ( + record.role in split_roles and InstanceRole.PREFILL_DECODE in roles + ): + raise ValueError("Cannot mix prefill_decode with standalone prefill/decode") + key = record.url + previous = self._live.get(key) or self._evicted.get(key) + if previous is not None: + if previous == record: + return False + if previous.role is not record.role: + raise ValueError("Unregister the instance before changing its role") + target = self._live if key in self._live else self._evicted + target[key] = record + return False + self._live[key] = record + logger.info("Registered instance %s: %s", record.role.value, record.url) + return True + + def unregister(self, url: str) -> bool: + """Drop an instance for good, so a probe cannot bring it back.""" + key = url + if key not in self._live and key not in self._evicted: + return False + self._live.pop(key, None) + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + logger.info("Unregistered instance: %s", url) + return True + + def instances(self, role: InstanceRole) -> list[InstanceRecord]: + return [record for record in self._live.values() if record.role is role] + + def urls(self, role: InstanceRole) -> list[str]: + return [record.url for record in self.instances(role)] + + def pick(self, role: InstanceRole) -> InstanceRecord | None: + """Take the next instance of `role` in round-robin order.""" + picked = self.pick_many(role, 1) + return picked[0] if picked else None + + def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]: + """Take `count` instances, continuing the rotation across calls. + + A multimodal request fans out one encoder request per item, so the + assignment has to be contiguous with the previous request's rather + than restart at the first instance every time. + """ + alive = self.instances(role) + if not alive or count <= 0: + return [] + start = self._cursors[role] + self._cursors[role] = start + count + return [alive[(start + offset) % len(alive)] for offset in range(count)] + + def next_replica(self, record: InstanceRecord) -> int: + """Take the next data-parallel replica of `record`, round-robin. + + The encoder pushes to one replica's receive channel, so the request + has to run on that same replica; the caller names the rank to both + halves. + """ + if record.dp_size <= 1: + return 0 + cursor = self._replica_cursors.get(record.url, 0) + self._replica_cursors[record.url] = cursor + 1 + return cursor % record.dp_size + + def status(self) -> dict[str, Any]: + return { + role.value: { + "live": [record.url for record in self.instances(role)], + "evicted": [ + record.url + for record in self._evicted.values() + if record.role is role + ], + } + for role in InstanceRole + } + + def start_probing(self) -> None: + if self._probe_task is None and self._probe_interval > 0: + self._probe_task = asyncio.create_task(self._probe_loop()) + + async def stop_probing(self) -> None: + if self._probe_task is None: + return + self._probe_task.cancel() + with suppress(asyncio.CancelledError): + await self._probe_task + self._probe_task = None + + async def _probe_loop(self) -> None: + timeout = aiohttp.ClientTimeout(total=self._probe_timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + while True: + await asyncio.sleep(self._probe_interval) + try: + await self._probe_once(session) + except Exception: + logger.exception("EPD registry probe round failed") + + async def _probe_once(self, session: aiohttp.ClientSession) -> None: + targets = list(self._live.values()) + list(self._evicted.values()) + if not targets: + return + results = await asyncio.gather( + *(self._probe(session, record.url) for record in targets), + return_exceptions=True, + ) + now = time.monotonic() + for record, healthy in zip(targets, results): + current = self._live.get(record.url) or self._evicted.get(record.url) + if current is not record: + continue + if healthy is True: + self._on_probe_success(record) + else: + self._on_probe_failure(record, now) + self._drop_expired(now) + + async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool: + async with session.get(f"{url}/health") as resp: + return resp.status == 200 + + def _on_probe_success(self, record: InstanceRecord) -> None: + key = record.url + self._fail_counts.pop(key, None) + if key in self._evicted: + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._live[key] = record + logger.info( + "Instance %s (%s) is healthy again; routing resumed", + record.url, + record.role.value, + ) + + def _on_probe_failure(self, record: InstanceRecord, now: float) -> None: + key = record.url + if key not in self._live: + # Either already evicted, or unregistered while this probe was in + # flight. A round snapshots its targets and then awaits, so a + # removal inside that window would otherwise be undone by a result + # describing a registry that no longer exists. + return + failures = self._fail_counts.get(key, 0) + 1 + self._fail_counts[key] = failures + if failures < self._fail_threshold: + return + self._live.pop(key, None) + self._evicted[key] = record + self._evicted_since[key] = now + logger.warning( + "Instance %s (%s) failed %d consecutive probes; stopped routing " + "to it. It rejoins on its own once it responds again.", + record.url, + record.role.value, + failures, + ) + + def _drop_expired(self, now: float) -> None: + if self._evicted_ttl <= 0: + return + for key, since in list(self._evicted_since.items()): + if now - since < self._evicted_ttl: + continue + record = self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + if record is not None: + logger.warning("Instance %s stayed down; forgetting it", record.url) + + +class InstanceRegistration(BaseModel): + role: InstanceRole + url: AnyHttpUrl + ec_zmq_addrs: list[str] = Field(default_factory=list) + dp_size: int = Field(default=1, ge=1) + + @model_validator(mode="after") + def validate_consumer_addresses(self): + if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size: + raise ValueError("Provide one Mooncake control address per DP replica") + return self + + +def require_admin_key(x_api_key: str = Header(default="")) -> None: + expected = os.getenv("ADMIN_API_KEY", "") + if not expected or not secrets.compare_digest(x_api_key, expected): + raise HTTPException(403, "Invalid admin API key") + + +@dataclass +class EPDProxyConfig: + probe_interval: float = 5.0 + probe_timeout: float = 2.0 + fail_threshold: int = 3 + evicted_ttl: float = 900.0 + + +@dataclass +class _Route: + """The instances one request was assigned to. + + Attributes: + consumer: The stage that receives the embedding, if any stage does. + consumer_zmq: The receive address named to the encoders. + dp_rank: Which replica of `consumer` was named, so the request can be + pinned to it. + + """ + + encoder_urls: list[str] + prefill: InstanceRecord | None + decode: InstanceRecord + consumer: InstanceRecord | None = None + consumer_zmq: str | None = None + dp_rank: int | None = None + + +class EPDProxy: + def __init__(self, registry: InstanceRegistry): + self.registry = registry + + # ---------------------------------------------------------------- # + # Routing # + # ---------------------------------------------------------------- # + def route(self, num_items: int) -> _Route: + decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + prefill = ( + self.registry.pick(InstanceRole.PREFILL) + if decode.role is InstanceRole.DECODE + else None + ) + if decode.role is InstanceRole.DECODE and prefill is None: + raise HTTPException( + status_code=503, detail="No prefill instance is registered" + ) + encoders = self.registry.urls(InstanceRole.ENCODE) + if num_items and not encoders: + raise HTTPException( + status_code=503, detail="No encode instance is registered" + ) + route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode) + self._name_consumer(route) + return route + + def _name_consumer(self, route: _Route) -> None: + """Pin the EC consumer replica and, for push connectors, its endpoint.""" + candidate = route.prefill or route.decode + route.consumer = candidate + rank = self.registry.next_replica(candidate) + route.dp_rank = rank if candidate.dp_size > 1 else None + if candidate.ec_zmq_addrs: + route.consumer_zmq = candidate.ec_zmq_addrs[rank] + + app = FastAPI() encode_session: aiohttp.ClientSession | None = None prefill_session: aiohttp.ClientSession | None = None @@ -448,6 +788,7 @@ async def maybe_prefill( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """- Do prefill-only task if p_url exist; - Return a new body carrying kv transfer params (for nixl connector) @@ -459,7 +800,7 @@ async def maybe_prefill( if p_url: logger.info("[%s] Processing through prefill: %s", req_id, p_url) - prefill_response = await process_prefill_stage(req_data, p_url, req_id) + prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank) # for nixl connector to facilitate kv transfer... prefill_response_json = msgspec.json.decode(await prefill_response.read()) kv_transfer_params = prefill_response_json.get("kv_transfer_params", {}) @@ -473,6 +814,7 @@ async def process_prefill_stage( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """Process request through Prefill stage and return kv_transfer_params""" logger.info("[%s] Sending prefill request to: %s", req_id, p_url) @@ -494,6 +836,8 @@ async def process_prefill_stage( del prefill_request["stream_options"] headers = {"x-request-id": req_id, "Content-Type": "application/json"} + if dp_rank is not None: + headers["X-data-parallel-rank"] = str(dp_rank) try: prefill_response = await prefill_session.post( f"{p_url}/v1/chat/completions", @@ -594,9 +938,7 @@ async def on_startup() -> None: ), ) encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) - if app.state.p_urls: - # only setup if prefill instance(s) exist - prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) + prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) @@ -622,6 +964,7 @@ async def prepare_for_decode( e_urls: list[str], p_url: str, consumer_zmq: str | None, + prefill_dp_rank: int | None = None, ) -> tuple[dict, float, float]: """Encode, rewrite and prefill, returning the body to send to decode. @@ -642,7 +985,7 @@ async def prepare_for_decode( handles.update(ec_params) prepared = {**prepared, "ec_transfer_params": handles} _t2 = time.perf_counter() - prepared = await maybe_prefill(prepared, p_url, req_id) + prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank) return prepared, _t1 - _t0, _t2 - _t1 @@ -654,12 +997,13 @@ async def forward_non_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> Response: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -733,12 +1077,13 @@ async def forward_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> AsyncIterator[bytes]: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -976,10 +1321,147 @@ async def stop_profile(request: Request): return await _profile_cmd("stop", body, e_url, p_url, d_url) +def build_app(config: EPDProxyConfig | None = None) -> FastAPI: + config = config or EPDProxyConfig() + registry = InstanceRegistry( + probe_interval=config.probe_interval, + probe_timeout=config.probe_timeout, + fail_threshold=config.fail_threshold, + evicted_ttl=config.evicted_ttl, + ) + proxy = EPDProxy(registry) + + @asynccontextmanager + async def lifespan(app: FastAPI): + await on_startup() + try: + registry.start_probing() + yield + finally: + await registry.stop_probing() + await on_shutdown() + + app = FastAPI(lifespan=lifespan) + app.state.proxy = proxy + app.state.registry = registry + + @app.get("/instances") + async def list_instances(): + return registry.status() + + @app.post("/instances", dependencies=[Depends(require_admin_key)]) + async def register_instance(body: InstanceRegistration): + if body.ec_zmq_addrs and body.role not in ( + InstanceRole.PREFILL, + InstanceRole.PREFILL_DECODE, + ): + raise HTTPException(400, "Only EC consumers accept Mooncake addresses") + try: + created = registry.register( + InstanceRecord( + body.role, + str(body.url).rstrip("/"), + body.ec_zmq_addrs, + body.dp_size, + ) + ) + except ValueError as error: + raise HTTPException(409, str(error)) from error + return {"registered": created} + + @app.delete("/instances", dependencies=[Depends(require_admin_key)]) + async def unregister_instance(url: AnyHttpUrl): + return {"removed": registry.unregister(str(url).rstrip("/"))} + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request): + req_data = msgspec.json.decode(await request.body()) + req_id = request.headers.get("x-request-id", str(uuid.uuid4())) + route = proxy.route(len(extract_mm_items(req_data))) + args = ( + req_data, + req_id, + route.encoder_urls, + route.prefill.url if route.prefill else None, + route.decode.url, + route.consumer_zmq, + ) + ranks = { + "dp_rank": route.dp_rank if route.prefill is None else None, + "prefill_dp_rank": route.dp_rank if route.prefill is not None else None, + } + if req_data.get("stream", False): + return StreamingResponse( + forward_stream(*args, **ranks), media_type="text/event-stream" + ) + return await forward_non_stream(*args, **ranks) + + @app.get("/v1/models") + async def list_models(): + decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + async with decode_session.get(f"{decode.url}/v1/models") as resp: + resp.raise_for_status() + return await resp.json() + + @app.get("/health") + async def health(): + status = registry.status() + # An empty roster is not unhealthy: the proxy is meant to come up + # before anything registers with it. + return JSONResponse({"proxy": "healthy", "instances": status}) + + @app.post("/start_profile") + async def start_profile(request: Request): + return await _profile(registry, "start", await request.json()) + + @app.post("/stop_profile") + async def stop_profile(request: Request): + return await _profile(registry, "stop", await request.json()) + + return app + + +async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict: + headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} + targets = {role.value: registry.pick(role) for role in InstanceRole} + results = await asyncio.gather( + *( + _post_if_available( + decode_session, f"{record.url}/{cmd}_profile", payload, headers + ) + for record in targets.values() + if record is not None + ) + ) + reachable = [result for result in results if result is not None] + if not reachable: + raise HTTPException( + status_code=503, + detail="Profiling endpoints are disabled on every instance", + ) + live = [name for name, record in targets.items() if record is not None] + return dict(zip(live, results)) + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--dynamic-registration", + action="store_true", + help="Enable launcher-managed HTTP registration for E+PD or E+P+D.", + ) + parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL) + parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT) + parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD) + parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL) parser.add_argument( "--log-requests", action="store_true", @@ -995,12 +1477,12 @@ async def stop_profile(request: Request): ) parser.add_argument( "--encode-servers-urls", - required=True, + default="", help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")', ) parser.add_argument( "--prefill-servers-urls", - required=True, + default="none", help=( 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ' 'to enable E->P->D, set "disable" or "none" to enable E->PD' @@ -1008,7 +1490,7 @@ async def stop_profile(request: Request): ) parser.add_argument( "--decode-servers-urls", - required=True, + default="", help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")', ) parser.add_argument( @@ -1043,58 +1525,87 @@ async def stop_profile(request: Request): ) args = parser.parse_args() + if args.fail_threshold < 1: + parser.error("--fail-threshold must be at least 1") if args.log_requests: logging.getLogger().setLevel(logging.DEBUG) app.middleware("http")(log_requests) NO_REWRITE = args.no_rewrite DECODE_RETRIES = max(0, args.decode_retries) - app.state.e_urls = [ - u.strip() for u in args.encode_servers_urls.split(",") if u.strip() - ] - app.state.d_urls = [ - u.strip() for u in args.decode_servers_urls.split(",") if u.strip() - ] - app.state.d_ec_urls = [ - u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() - ] - if args.ec_consumer_dp_size < 1: - parser.error("--ec-consumer-dp-size must be at least 1") - app.state.ec_consumer_dp_size = args.ec_consumer_dp_size - app.state.replica_counter = itertools.count() - expected = len(app.state.d_urls) * args.ec_consumer_dp_size - if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: - parser.error( - "--ec-consumer-zmq-addrs must contain one address per consumer " - f"replica: expected {expected} " - f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), " - f"got {len(app.state.d_ec_urls)}" - ) - # handle prefill instances - if args.prefill_servers_urls.lower() in ("disable", "none", ""): - app.state.p_urls = [] - logger.info( - "Disaggregated prefill phase explicitly disabled by user. Running E + PD..." + if args.dynamic_registration: + if not os.getenv("ADMIN_API_KEY"): + parser.error("--dynamic-registration requires ADMIN_API_KEY") + if ( + args.encode_servers_urls + or args.prefill_servers_urls.lower() not in ("disable", "none", "") + or args.decode_servers_urls + or args.ec_consumer_zmq_addrs + ): + parser.error( + "With --dynamic-registration, instances register through /instances" + ) + app = build_app( + EPDProxyConfig( + probe_interval=args.probe_interval, + probe_timeout=args.probe_timeout, + fail_threshold=args.fail_threshold, + evicted_ttl=args.evicted_ttl, + ) ) + if args.log_requests: + app.middleware("http")(log_requests) else: - app.state.p_urls = [ - u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + if not args.encode_servers_urls or not args.decode_servers_urls: + parser.error( + "Static routing requires --encode-servers-urls " + "and --decode-servers-urls" + ) + app.state.e_urls = [ + u.strip() for u in args.encode_servers_urls.split(",") if u.strip() ] - logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") - try: - validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) - except ValueError as exc: - parser.error(str(exc)) - - logger.info("Proxy listening on %s:%s", args.host, args.port) - logger.info("Encode servers: %s", app.state.e_urls) - logger.info("Prefill instances %s", app.state.p_urls) - logger.info("Decode servers: %s", app.state.d_urls) - if app.state.ec_consumer_dp_size > 1: - logger.info( - "EC consumer replicas per server: %d (control addresses: %s)", - app.state.ec_consumer_dp_size, - app.state.d_ec_urls, - ) + app.state.d_urls = [ + u.strip() for u in args.decode_servers_urls.split(",") if u.strip() + ] + app.state.d_ec_urls = [ + u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() + ] + if args.ec_consumer_dp_size < 1: + parser.error("--ec-consumer-dp-size must be at least 1") + app.state.ec_consumer_dp_size = args.ec_consumer_dp_size + app.state.replica_counter = itertools.count() + expected = len(app.state.d_urls) * args.ec_consumer_dp_size + if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: + parser.error( + "--ec-consumer-zmq-addrs must contain one address per consumer " + f"replica: expected {expected} " + f"({len(app.state.d_urls)} servers x " + f"{args.ec_consumer_dp_size} replicas), " + f"got {len(app.state.d_ec_urls)}" + ) + # handle prefill instances + if args.prefill_servers_urls.lower() in ("disable", "none", ""): + app.state.p_urls = [] + logger.info("Disaggregated prefill disabled. Running E + PD...") + else: + app.state.p_urls = [ + u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + ] + logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") + try: + validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) + except ValueError as exc: + parser.error(str(exc)) + + logger.info("Proxy listening on %s:%s", args.host, args.port) + logger.info("Encode servers: %s", app.state.e_urls) + logger.info("Prefill instances %s", app.state.p_urls) + logger.info("Decode servers: %s", app.state.d_urls) + if app.state.ec_consumer_dp_size > 1: + logger.info( + "EC consumer replicas per server: %d (control addresses: %s)", + app.state.ec_consumer_dp_size, + app.state.d_ec_urls, + ) uvicorn.run( app, diff --git a/tests/v1/ec_connector/unit/test_epd_proxy.py b/tests/v1/ec_connector/unit/test_epd_proxy.py new file mode 100644 index 000000000000..13496a7ffb23 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_proxy.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routing and registration behaviour of the EPD proxy.""" + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, Mock + +import msgspec +import pytest +from fastapi import HTTPException +from fastapi.testclient import TestClient + +from examples.disaggregated.disaggregated_encoder import ( + disagg_epd_proxy as proxy_module, +) +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + EPDProxy, + EPDProxyConfig, + InstanceRecord, + InstanceRegistry, + InstanceRole, + build_app, + extract_mm_items, +) + +ENCODE = InstanceRole.ENCODE +PREFILL = InstanceRole.PREFILL +DECODE = InstanceRole.DECODE +PD = InstanceRole.PREFILL_DECODE + +IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}} + + +@pytest.fixture +def proxy(): + registry = InstanceRegistry(probe_interval=0) + return EPDProxy(registry) + + +class TestRouting: + def test_no_decode_instance_is_service_unavailable(self, proxy): + """The proxy comes up before anything registers, so this is normal.""" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + def test_media_without_an_encoder_is_service_unavailable(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).decode.url == "http://d0:8000" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=1) + assert excinfo.value.status_code == 503 + + def test_prefill_is_optional(self, proxy): + """An E+PD deployment registers no prefill instance at all.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).prefill is None + + def test_encoder_roster_tracks_new_instances(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + for index in range(2): + proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000")) + assert proxy.route(num_items=3).encoder_urls == [ + "http://e0:8000", + "http://e1:8000", + ] + proxy.registry.unregister("http://e0:8000") + assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"] + + @pytest.mark.parametrize("unregister", [False, True]) + def test_standalone_decode_always_requires_healthy_prefill(self, unregister): + app = build_app() + proxy = app.state.proxy + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000")) + route = proxy.route(num_items=0) + assert route.decode.url == "http://d0:8000" + assert route.consumer.url == "http://p0:8000" + if unregister: + proxy.registry.unregister("http://p0:8000") + else: + for _ in range(proxy.registry._fail_threshold): + proxy.registry._on_probe_failure(route.prefill, 0) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + +class TestConsumerAddress: + """Which stage receives the embedding depends on the topology.""" + + def test_shared_storage_connectors_name_no_target(self, proxy): + """Rotate replicas without requiring a push endpoint.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2)) + route = proxy.route(num_items=0) + assert route.consumer_zmq is None + assert route.dp_rank == 0 + assert proxy.route(num_items=0).dp_rank == 1 + + def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy): + proxy.registry.register( + InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"]) + ) + assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001" + + def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy): + """In E+P+D the prefill instance consumes the embedding, not decode. + + Reading the address out of whichever list it came from -- rather than + from the record that reported one -- sends the push to the wrong + instance in this topology. + """ + proxy.registry.register( + InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"]) + ) + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001" + + +class TestProxyApi: + @pytest.fixture + def client(self, monkeypatch): + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + config = EPDProxyConfig(probe_interval=0) + with TestClient(build_app(config)) as client: + client.headers["X-API-Key"] = "test-key" + yield client + + def test_proxy_serves_before_anything_registers(self, client): + assert client.get("/health").status_code == 200 + assert client.get("/instances").json()["encode"]["live"] == [] + + def test_requests_are_refused_until_a_decoder_registers(self, client): + response = client.post( + "/v1/chat/completions", json={"model": "m", "messages": []} + ) + assert response.status_code == 503 + assert client.get("/v1/models").status_code == 503 + + def test_http_registration_keeps_consumer_replica_and_address_together( + self, client + ): + payload: dict[str, Any] = { + "role": "prefill_decode", + "url": "http://d0:8000/", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"], + } + assert client.post("/instances", json=payload).json() == {"registered": True} + assert client.post("/instances", json=payload).json() == {"registered": False} + proxy = client.app.state.proxy + for rank in [0, 1, 0]: + route = proxy.route(0) + assert route.dp_rank == rank + assert route.consumer_zmq == payload["ec_zmq_addrs"][rank] + assert client.delete("/instances", params={"url": payload["url"]}).json() == { + "removed": True + } + assert proxy.registry.pick(PD) is None + + def test_registration_rejects_invalid_topology(self, client): + cases: list[tuple[dict[str, Any], int]] = [ + ({"dp_size": 0}, 422), + ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422), + ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ] + for fields, status in cases: + response = client.post( + "/instances", + json={"url": "http://d0:8000", "role": "prefill_decode", **fields}, + ) + assert response.status_code == status + + def test_registration_requires_admin_key(self, client): + client.headers.pop("X-API-Key") + assert ( + client.post( + "/instances", json={"role": "encode", "url": "http://e0:8000"} + ).status_code + == 403 + ) + assert ( + client.delete("/instances", params={"url": "http://e0:8000"}).status_code + == 403 + ) + + @pytest.mark.parametrize("split_role", ["prefill", "decode"]) + @pytest.mark.parametrize("combined_first", [False, True]) + def test_registration_rejects_mixed_topologies( + self, client, split_role, combined_first + ): + roles = ["prefill_decode", split_role] + if not combined_first: + roles.reverse() + first = {"role": roles[0], "url": "http://first:8000"} + second = {"role": roles[1], "url": "http://second:8000"} + assert client.post("/instances", json=first).status_code == 200 + assert client.post("/instances", json=second).status_code == 409 + # Health eviction must not silently permit a topology switch. + registry = client.app.state.registry + record = registry._live[first["url"]] + for _ in range(registry._fail_threshold): + registry._on_probe_failure(record, 0) + assert client.post("/instances", json=second).status_code == 409 + assert client.delete("/instances", params={"url": first["url"]}).json() == { + "removed": True + } + assert client.post("/instances", json=second).status_code == 200 + + +def test_extract_mm_items_finds_media_across_messages(): + req = { + "messages": [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]}, + {"role": "user", "content": [IMAGE_ITEM]}, + ] + } + assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM] + + +@pytest.mark.parametrize("push", [False, True]) +def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch): + """Keep main's hash-keyed handles and per-item push IDs after the proxy move.""" + proxy.registry.register(InstanceRecord(ENCODE, "http://encoder")) + proxy.registry.register( + InstanceRecord( + PD, + "http://decode", + ec_zmq_addrs=["tcp://decode:14579"] if push else [], + ) + ) + handle = { + "metadata": {"image_grid_thw": [[1, 2, 3]]}, + "item_indices": [0, 1], + "peer_port": 5601, + } + response = Mock( + status=200, + read=AsyncMock( + return_value=msgspec.json.encode( + {"ec_transfer_params": {"engine-hash": handle}} + ) + ), + ) + session = Mock(post=AsyncMock(return_value=response)) + monkeypatch.setattr(proxy_module, "encode_session", session) + original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]} + route = proxy.route(2) + prepared, _, _ = asyncio.run( + proxy_module.prepare_for_decode( + original, "request", route.encoder_urls, None, route.consumer_zmq + ) + ) + assert "ec_transfer_params" not in original + assert prepared["ec_transfer_params"]["engine-hash"] == handle + items = prepared["ec_transfer_params"]["ec_items"] + assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"] + assert items[0]["transfer_id"] != items[1]["transfer_id"] + assert session.post.call_count == 1 + sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get( + "ec_transfer_params" + ) + if push: + assert sent == { + "consumer_zmq": "tcp://decode:14579", + "ec_items": [{"transfer_id": item["transfer_id"]} for item in items], + } + else: + assert sent is None + for index in range(2): + assert prepared["messages"][0]["content"][index]["image_embeds"] == { + "image_grid_thw": [1, 2, 3] + } diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 0c48dc84d6f8..8db22fe1b976 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -8,6 +8,7 @@ import asyncio import importlib.util +import sys from pathlib import Path from types import SimpleNamespace from typing import Any @@ -27,6 +28,7 @@ def proxy(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module @@ -48,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch): """ served = [{"remote_block_ids": [1, 2]}, {}] - async def _stage(req_data, p_url, req_id): + async def _stage(req_data, p_url, req_id, dp_rank=None): assert "kv_transfer_params" not in req_data return _Response(served.pop(0)) @@ -222,7 +224,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch _EncoderSession([{"encoder-side-hash": handle}, {}]), ) - async def _no_prefill(req_data, p_url, req_id): + async def _no_prefill(req_data, p_url, req_id, dp_rank=None): return req_data monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill) @@ -280,6 +282,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "request_options", [ @@ -293,7 +296,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): ], ) async def test_http_roundtrip_preserves_payload_and_response_bytes( - proxy, monkeypatch, stream, prefill, request_options + proxy, monkeypatch, stream, prefill, dynamic, request_options ): """Preserve media and scheduling context across HTTP hops and retries.""" seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} @@ -307,6 +310,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( async def backend(request): stage = request.match_info["stage"] + if dynamic: + consumer = "prefill" if prefill else "decode" + expected_rank = "0" if stage == consumer else None + assert request.headers.get("X-data-parallel-rank") == expected_rank assert request.content_type == "application/json" body = await request.json() seen[stage].append(body) @@ -354,6 +361,10 @@ async def backend(request): monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False) monkeypatch.setattr(proxy, "DECODE_RETRIES", 1) monkeypatch.setattr(proxy, "NO_REWRITE", False) + target_app = proxy.app + if dynamic: + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0)) item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}} body = { "model": "test", @@ -377,9 +388,37 @@ async def backend(request): await proxy.on_startup() try: async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy" + transport=httpx.ASGITransport(app=target_app), base_url="http://proxy" ) as client: + if dynamic: + registrations = [("encode", "encode")] + registrations += ( + [("prefill", "prefill"), ("decode", "decode")] + if prefill + else [("prefill_decode", "decode")] + ) + for role, stage in registrations: + consumer = "prefill" if prefill else "decode" + registered = await client.post( + "/instances", + headers={"X-API-Key": "test-key"}, + json={ + "role": role, + "url": str(server.make_url(f"/{stage}")), + "dp_size": 2 if stage == consumer else 1, + }, + ) + assert registered.status_code == 200 response = await client.post("/v1/chat/completions", json=body) + if dynamic: + removed = await client.delete( + "/instances", + headers={"X-API-Key": "test-key"}, + params={"url": str(server.make_url("/decode"))}, + ) + assert removed.json() == {"removed": True} + unavailable = await client.post("/v1/chat/completions", json=body) + assert unavailable.status_code == 503 assert response.status_code == 200 assert response.content == expected assert response.headers["content-type"].startswith( diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py index 558a80f9cf82..8e1a730ca6b8 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -14,6 +14,7 @@ """ import importlib.util +import sys from collections import Counter from pathlib import Path @@ -27,6 +28,7 @@ def _load_proxy_module(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module diff --git a/tests/v1/ec_connector/unit/test_epd_registry.py b/tests/v1/ec_connector/unit/test_epd_registry.py new file mode 100644 index 000000000000..be5a8e9b9b3e --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_registry.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dynamic registration and liveness for the EPD proxy's instance registry.""" + +from unittest.mock import patch + +import pytest + +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + InstanceRecord, + InstanceRegistry, + InstanceRole, +) + +ENCODE = InstanceRole.ENCODE +DECODE = InstanceRole.DECODE + + +@pytest.fixture +def registry(): + return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0) + + +def _fill(registry, role, count, prefix="http://h"): + for index in range(count): + registry.register(InstanceRecord(role, f"{prefix}{index}:8000")) + + +async def _probe_round(registry, healthy: set[str], now: float = 0.0): + """Run one probe pass where only `healthy` URLs answer.""" + + async def fake_probe(self, session, url): + return url in healthy + + with ( + patch.object(InstanceRegistry, "_probe", fake_probe), + patch( + "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic", + return_value=now, + ), + ): + await registry._probe_once(session=None) + + +class TestRegistration: + def test_registering_is_idempotent_and_reports_novelty(self, registry): + record = InstanceRecord(ENCODE, "http://e0:8000") + assert registry.register(record) is True + assert registry.register(record) is False + assert registry.urls(ENCODE) == ["http://e0:8000"] + + def test_transfer_addresses_stay_with_the_registered_consumer(self, registry): + """Roster updates must not shift a consumer's transfer addresses.""" + registry.register(InstanceRecord(DECODE, "http://d0:8000")) + registry.register( + InstanceRecord( + DECODE, + "http://d0:8000", + ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"], + dp_size=2, + ) + ) + picked = registry.pick(DECODE) + assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"] + assert picked.dp_size == 2 + + def test_unregister_is_terminal(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + assert registry.unregister("http://e0:8000") is True + assert registry.unregister("http://e0:8000") is False + assert registry.pick(ENCODE) is None + + def test_round_robin_survives_registration(self, registry): + _fill(registry, ENCODE, 2) + assert [record.url for record in registry.pick_many(ENCODE, 2)] == [ + "http://h0:8000", + "http://h1:8000", + ] + registry.register(InstanceRecord(ENCODE, "http://h2:8000")) + assert registry.pick(ENCODE).url == "http://h2:8000" + + +class TestLiveness: + @pytest.mark.asyncio + async def test_identical_reregistration_preserves_inflight_probe_results( + self, registry + ): + """Repeated registrations must not suppress eviction or recovery.""" + url = "http://e0:8000" + registry.register(InstanceRecord(ENCODE, url)) + healthy = False + + async def probe(self, session, url): + registry.register(InstanceRecord(ENCODE, url)) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + for _ in range(registry._fail_threshold): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": [url]} + + healthy = True + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [url], "evicted": []} + + @pytest.mark.asyncio + async def test_consecutive_failures_stop_routing(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + registry.register(InstanceRecord(ENCODE, "http://e1:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy={"http://e1:8000"}) + assert registry.urls(ENCODE) == ["http://e1:8000"] + assert registry.status()["encode"]["evicted"] == ["http://e0:8000"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", [ENCODE, DECODE]) + async def test_a_recovered_instance_rejoins_without_re_registering( + self, registry, role + ): + """Restarting every encoder to recover from a blip is not acceptable.""" + registry.register(InstanceRecord(role, "http://e0:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + assert registry.urls(role) == [] + + await _probe_round(registry, healthy={"http://e0:8000"}, now=1) + assert registry.urls(role) == ["http://e0:8000"] + + @pytest.mark.asyncio + async def test_re_registration_does_not_override_failed_health_checks( + self, registry + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + registry.register(record) + assert registry.urls(ENCODE) == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("healthy", [False, True]) + async def test_inflight_probe_cannot_restore_unregistered_instance( + self, registry, healthy + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + + async def probe(self, session, url): + registry.unregister(url) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": []} From 92b40f5a144763c8b2ef6ebe848793de634ada05 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Sun, 20 Sep 2026 05:25:13 +0100 Subject: [PATCH 527/774] [CPU][Perf] refactor paged attention for Arm CPUs (#56045) Signed-off-by: Fadi Arafeh Co-authored-by: Codex Co-authored-by: Li, Jiang --- csrc/cpu/cpu_attn_impl.hpp | 49 +- csrc/cpu/cpu_attn_neon.hpp | 4 +- csrc/cpu/cpu_attn_neon_bfmmla.hpp | 847 +++++++++++-------------- vllm/v1/attention/backends/cpu_attn.py | 7 +- 4 files changed, 401 insertions(+), 506 deletions(-) diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index e9ddaacea573..4ade47eff373 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -956,6 +956,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col, std::printf("%s", ss.str().c_str()); } +// This allows us to store probabilities in a packed format +// for GEMM implementations that need packed A. +template +struct ProbabilityTokenStore { + using prob_buffer_t = typename attention_impl_t::prob_buffer_t; + static constexpr int32_t TokenStride = 1; + + FORCE_INLINE static void store_probabilities( + prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values, + const int32_t, const int64_t) { + using prob_buffer_vec_t = typename VecTypeTrait::vec_t; + prob_buffer_vec_t output_vec(values); + output_vec.save(probability); + } +}; + +template +struct ProbabilityTokenStore< + attention_impl_t, + std::void_t> + : attention_impl_t::ProbabilityTokenStore {}; + template class AttentionMainLoop { public: @@ -966,6 +988,7 @@ class AttentionMainLoop { using partial_output_buffer_t = typename attention_impl_t::partial_output_buffer_t; using prob_buffer_t = typename attention_impl_t::prob_buffer_t; + using probability_token_store_t = ProbabilityTokenStore; static constexpr int64_t max_q_head_num_per_iter = attention_impl_t::MaxQHeadNumPerIteration; @@ -974,6 +997,8 @@ class AttentionMainLoop { static constexpr int64_t headdim_alignment = attention_impl_t::HeadDimAlignment; static constexpr int64_t head_dim = attention_impl_t::HeadDim; + static constexpr int32_t probability_token_stride = + probability_token_store_t::TokenStride; static constexpr ISA ISAType = attention_impl_t::ISAType; static constexpr bool scale_on_logits = attention_impl_t::scale_on_logits; // apply scale on logits, otherwise @@ -1152,8 +1177,6 @@ class AttentionMainLoop { ? kv_tile_token_num : kv_tile_token_num * (sizeof(logits_buffer_t) / sizeof(prob_buffer_t)); - constexpr int64_t prob_buffer_elem_size = - prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t); partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer; bool accum_c = !is_first_iter; for (int32_t block_idx = start_block_idx; block_idx < end_block_idx; @@ -1187,9 +1210,7 @@ class AttentionMainLoop { remaining_group_num -= curr_group_num_in_block; curr_group_offset = 0; curr_group_num_in_block = token_group_num_per_block; - curr_prob_buffer = reinterpret_cast( - reinterpret_cast(curr_prob_buffer) + - curr_token_num * prob_buffer_elem_size); + curr_prob_buffer += curr_token_num * probability_token_stride; curr_partial_q_buffer = partial_q_buffer; accum_c = true; } @@ -1262,7 +1283,6 @@ class AttentionMainLoop { std::is_same_v; #endif - using prob_buffer_vec_t = typename VecTypeTrait::vec_t; constexpr bool prequantize_probabilities = []() { if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) { return tile_gemm_t::prequantize_probabilities; @@ -1331,12 +1351,6 @@ class AttentionMainLoop { vec = fast_exp(vec); } - if constexpr (prequantize_probabilities) { - vec.save(curr_logits_buffer_iter); - } else { - prob_buffer_vec_t output_vec(vec); - output_vec.save(curr_prob_buffer_iter); - } #else vec.save(curr_logits_buffer_iter); for (int32_t k = 0; k < 16; ++k) { @@ -1345,10 +1359,19 @@ class AttentionMainLoop { vec = vec_op::FP32Vec16(curr_logits_buffer_iter); #endif + if constexpr (prequantize_probabilities) { + vec.save(curr_logits_buffer_iter); + } else { + probability_token_store_t::store_probabilities( + curr_prob_buffer_iter, vec, i, + logits_buffer_stride * sizeof(logits_buffer_t) / + sizeof(prob_buffer_t)); + } + sum_vec = sum_vec + vec; curr_logits_buffer_iter += 16; - curr_prob_buffer_iter += 16; + curr_prob_buffer_iter += 16 * probability_token_stride; } if constexpr (prequantize_probabilities) { tile_gemm_t::quantize_probability_row( diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp index db4c5df2e88d..eed749944496 100644 --- a/csrc/cpu/cpu_attn_neon.hpp +++ b/csrc/cpu/cpu_attn_neon.hpp @@ -11,7 +11,7 @@ namespace cpu_attention { namespace { -#define BLOCK_SIZE_ALIGNMENT 32 +#define BLOCK_SIZE_ALIGNMENT 16 #define HEAD_SIZE_ALIGNMENT 32 #define MAX_Q_HEAD_NUM_PER_ITER 16 @@ -387,7 +387,7 @@ class AttentionImpl { }; #ifdef ARM_BF16_SUPPORT -// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment. +// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment. template class AttentionImpl : public AttentionImplNEONBFMMLA #include +#include +#include #include -#include namespace cpu_attention { -namespace { - -// BFMMLA tile dimensions -constexpr int32_t TILE_ROWS = 2; // M dimension -constexpr int32_t TILE_K = 4; // K reduction -constexpr int32_t TILE_COLS = 2; // N dimension (column-pair) - -// Derived constants -constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs -constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache -constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block -constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K; -constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS; -constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing - -// Matrix Packing and Accumulator -// Reshape two rows of Q into BFMMLA-friendly interleaved -// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3] -// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7] -// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk -FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0, - const c10::BFloat16* __restrict r1, - c10::BFloat16* __restrict dst, - int32_t K) { - const uint16_t* s0 = reinterpret_cast(r0); - const uint16_t* s1 = reinterpret_cast(r1); - uint16_t* d = reinterpret_cast(dst); - - // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output) - int32_t k = 0; - for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) { - vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k))); - } - - // Handle K tail: pack remaining elements with zero-padding - const int32_t tail = K - k; - if (tail > 0) { - // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1], - // pad] - for (int32_t t = 0; t < tail; ++t) { - d[t] = s0[k + t]; - d[t + TILE_K] = s1[k + t]; - } - // Zero-pad the rest - for (int32_t t = tail; t < TILE_K; ++t) { - d[t] = 0; - d[t + TILE_K] = 0; +class BfmmlaGemm { + public: + static constexpr int32_t KTile = 4; + static constexpr int32_t NTile = 8; + static constexpr int32_t MaxRows = 8; + + FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a, + const c10::BFloat16* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t n, const int32_t k, + const int64_t a_pair_stride, + const int64_t b_n_group_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const auto* a_ptr = reinterpret_cast(a); + const auto* b_ptr = reinterpret_cast(b); + + for (int32_t n_idx = 0; n_idx < n; n_idx += 16) { + const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride; + float* c_panel = c + n_idx; + + // Preserve this range so the inactive row pair is optimized away. + if (m <= 2) { + gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride, + b_n_group_stride, b_k_group_stride, ldc, accumulate); + } else if (m <= 4) { + gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride, + b_n_group_stride, b_k_group_stride, ldc, accumulate); + } else { + gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride, + ldc, accumulate); + gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k, + a_pair_stride, b_k_group_stride, ldc, accumulate); + } } } -} - -// 2x2 accumulator load/store with compile-time row count -template -FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) { - static_assert(m_rows == 1 || m_rows == 2); - float32x2_t row0 = vld1_f32(base + col_off); - float32x2_t row1 = - (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f); - return vcombine_f32(row0, row1); -} - -template -FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc, - int col_off) { - static_assert(m_rows == 1 || m_rows == 2); - vst1_f32(base + col_off, vget_low_f32(acc)); - if constexpr (m_rows == 2) { - vst1_f32(base + ldc + col_off, vget_high_f32(acc)); - } -} - -// Initialize 4 column-pair accumulators for 2 rows (8 columns total) -#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \ - do { \ - if (accum) { \ - if (m_rows == 2) { \ - a0 = load_acc_2x2<2>(Crow, ldc, 0); \ - a1 = load_acc_2x2<2>(Crow, ldc, 2); \ - a2 = load_acc_2x2<2>(Crow, ldc, 4); \ - a3 = load_acc_2x2<2>(Crow, ldc, 6); \ - } else { \ - a0 = load_acc_2x2<1>(Crow, ldc, 0); \ - a1 = load_acc_2x2<1>(Crow, ldc, 2); \ - a2 = load_acc_2x2<1>(Crow, ldc, 4); \ - a3 = load_acc_2x2<1>(Crow, ldc, 6); \ - } \ - } else { \ - a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \ - } \ - } while (0) - -// Store 4 column-pair accumulators back to C matrix -#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \ - do { \ - if (m_rows == 2) { \ - store_acc_2x2<2>(a0, Crow, ldc, 0); \ - store_acc_2x2<2>(a1, Crow, ldc, 2); \ - store_acc_2x2<2>(a2, Crow, ldc, 4); \ - store_acc_2x2<2>(a3, Crow, ldc, 6); \ - } else { \ - store_acc_2x2<1>(a0, Crow, ldc, 0); \ - store_acc_2x2<1>(a1, Crow, ldc, 2); \ - store_acc_2x2<1>(a2, Crow, ldc, 4); \ - store_acc_2x2<1>(a3, Crow, ldc, 6); \ - } \ - } while (0) - -// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs -#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \ - do { \ - r0 = vbfmmlaq_f32(r0, a, b0); \ - r1 = vbfmmlaq_f32(r1, a, b1); \ - r2 = vbfmmlaq_f32(r2, a, b2); \ - r3 = vbfmmlaq_f32(r3, a, b3); \ - } while (0) - -// Micro-kernel: updates a small fixed tile using BFMMLA. -// RP = number of row-pairs (1,2,4) -// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B. -// A_packed interleaves RP row-pairs; B layout is driven by the attention phase: -// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T) -// - AttentionGemmPhase::PV -> token-row layout (P @ V) -// K_static < 0 enables runtime K (PV only) -template -FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon( - const bfloat16_t* const* __restrict A_packed_rp, - const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk, - float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride, - int32_t K_runtime = 0) { - static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4"); - static_assert(K_static < 0 || K_static % TILE_K == 0, - "K must be divisible by TILE_K"); - static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, - "Runtime K only supported for PV"); - - constexpr bool runtime_k = (K_static < 0); - const int32_t K_iters = - runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K); - const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0; - - if (!runtime_k) { - // Help the compiler fold away unused K_runtime when K is compile-time - (void)K_runtime; - } - - auto* C_al = C; - const auto* B_al = B_blk; - // Setup A pointers - const bfloat16_t* a_ptr[4] = { - A_packed_rp[0], - (RP >= 2) ? A_packed_rp[1] : nullptr, - (RP >= 4) ? A_packed_rp[2] : nullptr, - (RP >= 4) ? A_packed_rp[3] : nullptr, - }; - - // Setup B pointers based on layout - const bfloat16_t* b_ptr[4]; - if constexpr (phase == AttentionGemmPhase::PV) { - b_ptr[0] = B_blk + 0 * b_stride; - b_ptr[1] = B_blk + 1 * b_stride; - b_ptr[2] = B_blk + 2 * b_stride; - b_ptr[3] = B_blk + 3 * b_stride; + private: + FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a, + const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); } - float32x4_t acc[4][4]; - -// Initialize accumulators -#define INIT_RP(rp) \ - if constexpr (RP > rp) { \ - INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ - C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \ + FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a, + const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); } - INIT_RP(0); - INIT_RP(1); - INIT_RP(2); - INIT_RP(3); -#undef INIT_RP - - // Main compute loop - for (int32_t ki = 0; ki < K_iters; ++ki) { - bfloat16x8_t b0, b1, b2, b3; - if constexpr (phase == AttentionGemmPhase::PV) { - b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE); - b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE); - b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE); - b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE); - } else { - const bfloat16_t* b_base = B_al + ki * b_stride; - b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE); - b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE); - b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE); - b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE); + + FORCE_INLINE static void init_accumulators( + float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45, + float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc, + const int32_t rows, const bool accumulate) { + if (!accumulate || rows == 0) { + acc01 = vdupq_n_f32(0.0f); + acc23 = vdupq_n_f32(0.0f); + acc45 = vdupq_n_f32(0.0f); + acc67 = vdupq_n_f32(0.0f); + return; } -#define COMPUTE_RP(rp) \ - if constexpr (RP > rp) { \ - bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \ - BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \ - b1, b2, b3); \ - } - COMPUTE_RP(0); - COMPUTE_RP(1); - COMPUTE_RP(2); - COMPUTE_RP(3); -#undef COMPUTE_RP + const float32x4_t row0_0123 = vld1q_f32(c); + const float32x4_t row0_4567 = vld1q_f32(c + 4); + const float32x4_t row1_0123 = + (rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f); + const float32x4_t row1_4567 = + (rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f); + + acc01 = zip_low_pairs(row0_0123, row1_0123); + acc23 = zip_high_pairs(row0_0123, row1_0123); + acc45 = zip_low_pairs(row0_4567, row1_4567); + acc67 = zip_high_pairs(row0_4567, row1_4567); } - // K tail for runtime PV: fallback path - if constexpr (runtime_k) { - if (K_tail > 0) { - const int32_t tail_offset = K_iters * V_INNER_STRIDE; - const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK; - for (int32_t kt = 0; kt < K_tail; ++kt) { - float32x4_t b_vecs[4]; - for (int32_t p = 0; p < 4; ++p) { - const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS; - const float b0 = vcvtah_f32_bf16(bp[0]); - const float b1 = vcvtah_f32_bf16(bp[1]); - const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1); - b_vecs[p] = vcombine_f32(b_pair, b_pair); - } + FORCE_INLINE static void store_accumulators( + const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45, + const float32x4_t acc67, float* __restrict__ c, const int64_t ldc, + const int32_t rows) { + if (rows == 0) { + return; + } -#define TAIL_RP(rp) \ - if constexpr (RP > rp) { \ - const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \ - float a_row0 = vcvtah_f32_bf16(ap[kt]); \ - float a_row1 = \ - (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \ - const float32x4_t a_vec = \ - vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \ - for (int32_t p = 0; p < 4; ++p) { \ - acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \ - } \ - } - TAIL_RP(0); - TAIL_RP(1); - TAIL_RP(2); - TAIL_RP(3); -#undef TAIL_RP - } + vst1q_f32(c, zip_low_pairs(acc01, acc23)); + vst1q_f32(c + 4, zip_low_pairs(acc45, acc67)); + if (rows == 2) { + vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23)); + vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67)); } } - // Store results -#define STORE_RP(rp) \ - if constexpr (RP > rp) { \ - STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ - C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \ - } - STORE_RP(0); - STORE_RP(1); - STORE_RP(2); - STORE_RP(3); -#undef STORE_RP -} - -// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the -// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables -// runtime K (PV only). -template -FORCE_INLINE void gemm_packA_compute_MB_xN( - const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, - float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc, - int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { - static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]"); - static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, - "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); - static_assert(K_static < 0 || K_static % TILE_K == 0, - "K must be divisible by TILE_K"); - static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, - "Runtime K only supported for PV"); - - constexpr bool runtime_k = (K_static < 0); - const int32_t K_val = runtime_k ? K_runtime : K_static; - - // Keep small packs on-stack to avoid heap churn - constexpr int32_t STACK_PACK_STRIDE = - (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; - - constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS; - const int32_t pack_stride = - runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK - : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; - - alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE]; - std::vector A_packed_heap; - c10::BFloat16* A_packed = - (pack_stride <= STACK_PACK_STRIDE) - ? A_packed_stack - : (A_packed_heap.resize(ROW_PAIRS * pack_stride), - A_packed_heap.data()); - - for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) { - const int32_t m = rp * TILE_ROWS; - const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1; - const c10::BFloat16* A0 = A + m * lda; - const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0; - reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val); + FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a, + const int32_t rows) { + if (rows == 0) { + return vdupq_n_bf16(bfloat16_t{}); + } + // Packed A reserves both rows for an M tail. + return vld1q_bf16(a); } - for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) { - const c10::BFloat16* B_blk_c10 = - (phase == AttentionGemmPhase::PV) - ? (B + (n / TILE_COLS) * b_layout_stride) - : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride); - const bfloat16_t* B_blk = reinterpret_cast(B_blk_c10); - - // Process row-pairs in groups of 4, 2, then 1 - int32_t row_pair_idx = 0; - -#define PROCESS_RP_GROUP(group_size) \ - for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \ - row_pair_idx += group_size) { \ - const bfloat16_t* Ap[group_size]; \ - int32_t mr[group_size]; \ - for (int32_t i = 0; i < group_size; ++i) { \ - Ap[i] = reinterpret_cast( \ - A_packed + (row_pair_idx + i) * pack_stride); \ - mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \ - } \ - float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \ - if constexpr (runtime_k) { \ - gemm_rowpairs_x8_bfmmla_neon( \ - Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \ - } else { \ - gemm_rowpairs_x8_bfmmla_neon( \ - Ap, mr, B_blk, C_blk, ldc, accumulate, \ - (phase == AttentionGemmPhase::PV) ? b_layout_stride \ - : b_reduction_stride); \ - } \ - } + FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a, + const bfloat16_t* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t k, + const int64_t a_pair_stride, + const int64_t b_n_group_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const int32_t rows01 = std::min(2, std::max(0, m)); + const int32_t rows23 = std::min(2, std::max(0, m - 2)); + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc0189, acc011011, acc011213, acc011415; + float32x4_t acc2389, acc231011, acc231213, acc231415; + init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01, + accumulate); + init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23, accumulate); + init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc, + rows01, accumulate); + init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8, + ldc, rows23, accumulate); + + const bfloat16_t* a01 = a; + const bfloat16_t* a23 = a + a_pair_stride; + const bfloat16_t* b0 = b; + const bfloat16_t* b1 = b + b_n_group_stride; + +#pragma GCC unroll 4 + for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) { + const bfloat16x8_t av01 = load_a_pair(a01, rows01); + const bfloat16x8_t av23 = load_a_pair(a23, rows23); + const bfloat16x8_t b01 = vld1q_bf16(b0); + const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile); + const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile); + const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile); + const bfloat16x8_t b89 = vld1q_bf16(b1); + const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile); + const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile); + const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile); + + acc0101 = vbfmmlaq_f32(acc0101, av01, b01); + acc2301 = vbfmmlaq_f32(acc2301, av23, b01); + acc0123 = vbfmmlaq_f32(acc0123, av01, b23); + acc2323 = vbfmmlaq_f32(acc2323, av23, b23); + acc0145 = vbfmmlaq_f32(acc0145, av01, b45); + acc2345 = vbfmmlaq_f32(acc2345, av23, b45); + acc0167 = vbfmmlaq_f32(acc0167, av01, b67); + acc2367 = vbfmmlaq_f32(acc2367, av23, b67); + acc0189 = vbfmmlaq_f32(acc0189, av01, b89); + acc2389 = vbfmmlaq_f32(acc2389, av23, b89); + acc011011 = vbfmmlaq_f32(acc011011, av01, b1011); + acc231011 = vbfmmlaq_f32(acc231011, av23, b1011); + acc011213 = vbfmmlaq_f32(acc011213, av01, b1213); + acc231213 = vbfmmlaq_f32(acc231213, av23, b1213); + acc011415 = vbfmmlaq_f32(acc011415, av01, b1415); + acc231415 = vbfmmlaq_f32(acc231415, av23, b1415); + + a01 += 2 * KTile; + a23 += 2 * KTile; + b0 += b_k_group_stride; + b1 += b_k_group_stride; + } - PROCESS_RP_GROUP(4); - PROCESS_RP_GROUP(2); - PROCESS_RP_GROUP(1); -#undef PROCESS_RP_GROUP + store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01); + store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23); + store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc, + rows01); + store_accumulators(acc2389, acc231011, acc231213, acc231415, + c + 2 * ldc + 8, ldc, rows23); } -} - -// Macro-kernel: iterates over M in MB={8,4,2,1} chunks. -// Supports compile-time K specialization when K >= 0; otherwise uses runtime K -// (runtime K path is only supported for PV). -template -FORCE_INLINE void gemm_macro_neon_bfmmla( - const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, - float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc, - int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { - static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, - "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); - - if constexpr (K >= 0) { - static_assert(K % TILE_K == 0, "K must be divisible by TILE_K"); - for (int32_t m = 0; m < M;) { - const int32_t rem = M - m; - const c10::BFloat16* A_blk = A + m * lda; - float* C_blk = C + m * ldc; - -#define DISPATCH_MB(mb) \ - gemm_packA_compute_MB_xN(A_blk, B, C_blk, 0, lda, ldc, \ - b_layout_stride, \ - b_reduction_stride, accumulate) - - if (rem >= 8) { - DISPATCH_MB(8); - m += 8; - } else if (rem >= 4) { - DISPATCH_MB(4); - m += 4; - } else if (rem >= 2) { - DISPATCH_MB(2); - m += 2; - } else { - DISPATCH_MB(1); - m += 1; - } -#undef DISPATCH_MB - } - } else { - static_assert(phase == AttentionGemmPhase::PV, - "Runtime K specialization only supported for PV."); - const int32_t K_val = K_runtime; - - for (int32_t m = 0; m < M;) { - const int32_t rem = M - m; - const c10::BFloat16* A_blk = A + m * lda; - float* C_blk = C + m * ldc; - -#define DISPATCH_MB_RUNTIME(mb) \ - gemm_packA_compute_MB_xN(A_blk, B, C_blk, K_val, lda, ldc, \ - b_layout_stride, \ - b_reduction_stride, accumulate) - - if (rem >= 8) { - DISPATCH_MB_RUNTIME(8); - m += 8; - } else if (rem >= 4) { - DISPATCH_MB_RUNTIME(4); - m += 4; - } else if (rem >= 2) { - DISPATCH_MB_RUNTIME(2); - m += 2; - } else { - DISPATCH_MB_RUNTIME(1); - m += 1; - } -#undef DISPATCH_MB_RUNTIME + + FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a, + const bfloat16_t* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t k, + const int64_t a_pair_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const int32_t rows01 = std::min(2, std::max(0, m)); + const int32_t rows23 = std::min(2, std::max(0, m - 2)); + const int32_t rows45 = std::min(2, std::max(0, m - 4)); + const int32_t rows67 = std::min(2, std::max(0, m - 6)); + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc4501, acc4523, acc4545, acc4567; + float32x4_t acc6701, acc6723, acc6745, acc6767; + init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01, + accumulate); + init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23, accumulate); + init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc, + rows45, accumulate); + init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc, + rows67, accumulate); + + const bfloat16_t* a01 = a; + const bfloat16_t* a23 = a + a_pair_stride; + const bfloat16_t* a45 = a + 2 * a_pair_stride; + const bfloat16_t* a67 = a + 3 * a_pair_stride; + const bfloat16_t* b_ptr = b; + +#pragma GCC unroll 4 + for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) { + const bfloat16x8_t av01 = load_a_pair(a01, rows01); + const bfloat16x8_t av23 = load_a_pair(a23, rows23); + const bfloat16x8_t av45 = load_a_pair(a45, rows45); + const bfloat16x8_t av67 = load_a_pair(a67, rows67); + const bfloat16x8_t b01 = vld1q_bf16(b_ptr); + const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile); + const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile); + const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile); + + acc0101 = vbfmmlaq_f32(acc0101, av01, b01); + acc2301 = vbfmmlaq_f32(acc2301, av23, b01); + acc4501 = vbfmmlaq_f32(acc4501, av45, b01); + acc6701 = vbfmmlaq_f32(acc6701, av67, b01); + acc0123 = vbfmmlaq_f32(acc0123, av01, b23); + acc2323 = vbfmmlaq_f32(acc2323, av23, b23); + acc4523 = vbfmmlaq_f32(acc4523, av45, b23); + acc6723 = vbfmmlaq_f32(acc6723, av67, b23); + acc0145 = vbfmmlaq_f32(acc0145, av01, b45); + acc2345 = vbfmmlaq_f32(acc2345, av23, b45); + acc4545 = vbfmmlaq_f32(acc4545, av45, b45); + acc6745 = vbfmmlaq_f32(acc6745, av67, b45); + acc0167 = vbfmmlaq_f32(acc0167, av01, b67); + acc2367 = vbfmmlaq_f32(acc2367, av23, b67); + acc4567 = vbfmmlaq_f32(acc4567, av45, b67); + acc6767 = vbfmmlaq_f32(acc6767, av67, b67); + + a01 += 2 * KTile; + a23 += 2 * KTile; + a45 += 2 * KTile; + a67 += 2 * KTile; + b_ptr += b_k_group_stride; } + + store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01); + store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23); + store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc, + rows45); + store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc, + rows67); } -} +}; + +namespace { -#undef INIT_ACC_ROWPAIR_4 -#undef STORE_ACC_ROWPAIR_4 -#undef BFMMLA_COMPUTE_4 +constexpr int32_t TILE_K = BfmmlaGemm::KTile; +constexpr int32_t TILE_COLS = 2; +constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile; +constexpr int32_t K_TOKENS_PER_GROUP = 8; +constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; +constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K; +constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS; } // namespace -// TileGemm Adapter for Attention - template class TileGemmNEONBFMMLA { public: @@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA { [[maybe_unused]] const int32_t block_size, [[maybe_unused]] const int32_t dynamic_k_size, const bool accum_c) { - static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0); - // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for - // API parity with other tile_gemm implementations. + static_assert(BlockTokens % 16 == 0); if constexpr (head_dim_ct >= 0) { - static_assert(head_dim_ct == HeadDim, - "BFMMLA expects head_dim_ct to match HeadDim; PV passes " - "-1 for API parity."); + static_assert(head_dim_ct == HeadDim); } + const auto* a = reinterpret_cast(a_tile); if constexpr (phase == AttentionGemmPhase::QK) { - const int64_t b_reduction_stride = K_INNER_STRIDE; - const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE; - - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride, - accum_c); + constexpr int64_t b_n_group_stride = + (HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE; + + for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) { + const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row); + BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m, + BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride, + K_CACHE_K_GROUP_STRIDE, ldc, accum_c); + } } else { - const int64_t b_pair_stride = - (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; - - // PV gemm with runtime K specialization - switch (dynamic_k_size) { - case 32: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 32, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 128: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 128, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 256: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 256, lda, ldc, b_pair_stride, 0, accum_c); - break; - default: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c); - break; + const int64_t b_n_group_stride = + (block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE; + + for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) { + const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row); + BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m, + HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride, + K_CACHE_K_GROUP_STRIDE, ldc, accum_c); } } } @@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA { static constexpr int64_t HeadDim = head_dim; static constexpr ISA ISAType = isa_type; static constexpr bool scale_on_logits = false; + static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK; + static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K; - static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0); - static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0); + static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0); + static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0); static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K"); public: @@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA { attention_iteration(CPU_ATTENTION_PARAMS); } + struct ProbabilityTokenStore { + static constexpr int32_t TokenStride = 2; + + FORCE_INLINE static void store_probabilities( + c10::BFloat16* __restrict__ probability, + const vec_op::FP32Vec16& values, const int32_t row, + const int64_t row_stride) { + const int32_t row_in_pair = row & 1; + auto* dst = reinterpret_cast( + probability + row_in_pair * (BfmmlaGemm::KTile - row_stride)); + vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0])); + vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1])); + vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2])); + vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3])); + } + }; + // Key cache stride per token group (TokenColumn layout; QK) static constexpr int64_t k_cache_token_group_stride( [[maybe_unused]] const int32_t block_size) { static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0); return (BlockSizeAlignment / K_TOKENS_PER_GROUP) * - ((head_dim / TILE_K) * K_INNER_STRIDE); + ((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE); } // Value cache stride per token group (TokenRow layout; PV) static constexpr int64_t v_cache_token_group_stride( [[maybe_unused]] const int32_t block_size) { static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0); - return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride; } // The stride to move to the "next" head_dim group @@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA { return head_dim * block_size; } - // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics + // Scale Q and write row pairs in BFMMLA reduction order. static void copy_q_heads_tile(c10::BFloat16* __restrict__ src, c10::BFloat16* __restrict__ q_buffer, const int32_t q_num, @@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA { const int64_t q_head_stride, float scale) { constexpr int32_t dim = static_cast(head_dim); const float32x4_t scale_vec = vdupq_n_f32(scale); - - for (int32_t qi = 0; qi < q_num; ++qi) { - for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) { - c10::BFloat16* __restrict__ curr_q = - src + qi * q_num_stride + hi * q_head_stride; - c10::BFloat16* __restrict__ dst = - q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim; - - for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) { - bfloat16x8_t in8 = - vld1q_bf16(reinterpret_cast(curr_q + i)); - float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec); - float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec); - - bfloat16x4_t lo_b = vcvt_bf16_f32(lo); - bfloat16x4_t hi_b = vcvt_bf16_f32(hi); - bfloat16x8_t out = vcombine_bf16(lo_b, hi_b); - vst1q_bf16(reinterpret_cast(dst + i), out); + const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{}); + const int32_t row_num = q_num * q_heads_per_kv; + + for (int32_t row = 0; row < row_num; row += 2) { + const int32_t q0 = row / q_heads_per_kv; + const int32_t h0 = row % q_heads_per_kv; + const auto* row0 = reinterpret_cast( + src + q0 * q_num_stride + h0 * q_head_stride); + const bool has_row1 = row + 1 < row_num; + const int32_t q1 = (row + 1) / q_heads_per_kv; + const int32_t h1 = (row + 1) % q_heads_per_kv; + const auto* row1 = has_row1 + ? reinterpret_cast( + src + q1 * q_num_stride + h1 * q_head_stride) + : nullptr; + auto* dst = reinterpret_cast(q_buffer + row * head_dim); + + for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) { + const bfloat16x8_t in0 = vld1q_bf16(row0 + k); + const bfloat16x4_t out0_lo = + vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec)); + const bfloat16x4_t out0_hi = + vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec)); + + bfloat16x4_t out1_lo = zero; + bfloat16x4_t out1_hi = zero; + if (has_row1) { + const bfloat16x8_t in1 = vld1q_bf16(row1 + k); + out1_lo = + vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec)); + out1_hi = + vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec)); } + + vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo)); + vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi)); } } } @@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA { public: // Reshape and cache K/V into BFMMLA-optimized layouts // K cache: - // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE] + // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K] + // [K_CACHE_K_GROUP_STRIDE] // - TokenColumn // V cache: - // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE] - // - TokenRows + // [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK] + // [VCacheKGroupStride] static void reshape_and_cache( const c10::BFloat16* __restrict__ key, const c10::BFloat16* __restrict__ value, @@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA { const int64_t block_size, [[maybe_unused]] const int64_t block_size_stride, const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) { - const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE; - const int64_t v_pair_stride = - (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE; + const int64_t v_n_group_stride = + (block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride; -#pragma omp parallel for - for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { - for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { +#pragma omp parallel for collapse(2) + for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { + for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { const int64_t pos = slot_mapping[token_idx]; if (pos < 0) continue; @@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA { key_base + block_in_block * k_block_stride; for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) { - uint16_t* dst_u16 = reinterpret_cast( - block_base + hd4 * K_INNER_STRIDE + - pair_in_block * V_INNER_STRIDE + lane_base); - const uint16_t* src_u16 = - reinterpret_cast(key_src + hd4 * TILE_K); - vst1_u16(dst_u16, vld1_u16(src_u16)); + bfloat16_t* dst = reinterpret_cast( + block_base + hd4 * K_CACHE_K_GROUP_STRIDE + + pair_in_block * B_COL_PAIR_STRIDE + lane_base); + const bfloat16_t* src = + reinterpret_cast(key_src + hd4 * TILE_K); + vst1_bf16(dst, vld1_bf16(src)); } } @@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA { value_cache + block_idx * num_blocks_stride + head_idx * cache_head_num_stride; - const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK; + const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK; const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1); - c10::BFloat16* __restrict row_block_base = - value_base + row_block * V_INNER_STRIDE; - - for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) { - c10::BFloat16* __restrict dst_val = - row_block_base + hd2 * v_pair_stride; - - const uint16_t* src_u16 = - reinterpret_cast(value_src); - uint16_t* dst_u16 = reinterpret_cast(dst_val); - dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0]; - dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] = - src_u16[hd2 * TILE_COLS + 1]; + const auto* src = reinterpret_cast(value_src); + auto* dst = reinterpret_cast(value_base); + for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) { + const bfloat16x8_t values = + vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK); + const bfloat16x4_t low = vget_low_bf16(values); + const bfloat16x4_t high = vget_high_bf16(values); + bfloat16_t* group = + dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride; + + vst1_lane_bf16(group + lane, low, 0); + vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1); + vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2); + vst1_lane_bf16( + group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low, + 3); + vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0); + vst1_lane_bf16( + group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, + high, 1); + vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2); + vst1_lane_bf16( + group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, + high, 3); } } } @@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA { } // namespace cpu_attention -#endif // CPU_ATTN_ASIMD_BFMMLA_HPP +#endif // CPU_ATTN_NEON_BFMMLA_HPP diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b6864897bf9b..24be5b8ec9c8 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -512,11 +512,10 @@ def _get_attn_isa( ) if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0: return "amx" + elif supports_arm: + return "neon" elif block_size % 32 == 0: - if supports_arm: - # support ARM NEON FMLA and BFMMLA (bf16) for block size 32 - return "neon" - elif supports_riscv and _riscv_supports_rvv(): + if supports_riscv and _riscv_supports_rvv(): return "rvv" elif supports_vxe: return "vxe" From 0c90804c0b8b139c252a86c79a3af26d181e11a0 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Sat, 19 Sep 2026 22:58:02 -0600 Subject: [PATCH 528/774] [Bugfix] DiffusionGemma: cast the self-conditioning soft embed to the buffer dtype (#57462) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 Co-authored-by: Lucas Wilkinson --- vllm/model_executor/models/diffusion_gemma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 0c1cc8c55e1a..10e0d8a2c224 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -543,7 +543,7 @@ def _compiled_sample_step( if tp_size > 1: soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name) soft_embeds = soft_embeds * normalizer - sc_embeds[decode_slots] = soft_embeds * sc_keep + sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype) # Overwrite canvas with argmax for newly converged denoise requests newly_converged = (converged & is_denoise).unsqueeze(1) From 601e1c7c51a10d6fd884d39f079f38bb942e436f Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Sat, 19 Sep 2026 22:58:30 -0600 Subject: [PATCH 529/774] [Perf] Give a prefill-only batch the model state's number of logit rows (#57416) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 --- vllm/v1/worker/gpu/model_runner.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e76c447ea5cc..8c736394cc52 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1291,7 +1291,7 @@ def prepare_inputs( # Get the number of draft tokens for each request. draft_tokens = scheduler_output.scheduled_spec_decode_tokens num_draft_tokens_per_req = None - if not draft_tokens: + if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1: # No draft token scheduled (common case). total_num_draft_tokens = 0 total_num_logits = num_reqs @@ -1322,9 +1322,9 @@ def prepare_inputs( np.cumsum(num_logits, out=cu_num_logits_np[1:]) cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) - adaptive_verification = ( - self.adaptive_verification if num_draft_tokens_per_req is not None else None - ) + # The general branch also serves a no-draft batch with k != 1, and + # compact_batch's budget is only primed when drafts are scheduled. + adaptive_verification = self.adaptive_verification if draft_tokens else None num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across @@ -1354,7 +1354,7 @@ def prepare_inputs( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens - if draft_tokens: + if num_draft_tokens_per_req is not None: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len ) @@ -1534,7 +1534,9 @@ def sample( sample_hidden_states = hidden_states[input_batch.logits_indices] logits = self.model.compute_logits(sample_hidden_states) - if grammar_output is not None: + # A diffusion prefill has no logit rows even when a bitmask row + # arrived for it. + if grammar_output is not None and logits.shape[0] > 0: # Apply grammar bitmask to the logits in-place. assert self.structured_outputs_worker is not None self.structured_outputs_worker.apply_grammar_bitmask( From d05da62e9ccdf8e342b15bf6785d83224cc165af Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Sun, 20 Sep 2026 14:00:31 +0800 Subject: [PATCH 530/774] [XPU] fix incorrect gdn kernel log for XPU path (#54871) Signed-off-by: Yan Ma --- vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index ccf1c60f0fe2..1ca07ad0cf3a 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -167,6 +167,8 @@ def _log_gdn_backend_decision( head_k_dim, ) return + elif current_platform.is_xpu(): + return chosen = { "flashinfer": "FlashInfer", @@ -524,6 +526,8 @@ def __init__( self.gdn_decode_kernel = "triton" elif current_platform.is_cpu(): self.gdn_decode_kernel = "CPU" + elif current_platform.is_xpu(): + self.gdn_decode_kernel = "XPU" self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) From e378275a8f20eac9e92212e4a1fee84ec5a31dc7 Mon Sep 17 00:00:00 2001 From: Mi Jiazhi <57009194+Mi-Jiazhi@users.noreply.github.com> Date: Sun, 20 Sep 2026 14:23:29 +0800 Subject: [PATCH 531/774] =?UTF-8?q?[Model]=20Pass=20intermediate=5Ftensors?= =?UTF-8?q?=20to=20the=20model=20when=20capturing=20CUDA=20gr=E2=80=A6=20(?= =?UTF-8?q?#57745)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: jiazhimi --- vllm/v1/worker/gpu/cudagraph_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4690c7eae70e..ddab7799d46c 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -662,6 +662,7 @@ def create_forward_fn( model_inputs = { "input_ids": input_buffers.input_ids[:num_tokens], "positions": input_buffers.positions[:num_tokens], + "intermediate_tensors": None, **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: From 27b7757f6364be5b0b714793ccdc3e3931443bca Mon Sep 17 00:00:00 2001 From: Dilber <44312867+dilberx@users.noreply.github.com> Date: Sun, 20 Sep 2026 14:22:27 +0530 Subject: [PATCH 532/774] [Bugfix][Frontend] Bound the prompt after multimodal expansion (#57076) Signed-off-by: Dilber P Shakir --- .../test_truncate_after_mm_expansion.py | 166 ++++++++++++++++++ vllm/renderers/base.py | 102 ++++++++++- vllm/renderers/hf.py | 52 +++--- 3 files changed, 285 insertions(+), 35 deletions(-) create mode 100644 tests/renderers/test_truncate_after_mm_expansion.py diff --git a/tests/renderers/test_truncate_after_mm_expansion.py b/tests/renderers/test_truncate_after_mm_expansion.py new file mode 100644 index 000000000000..96cdfb12c2d7 --- /dev/null +++ b/tests/renderers/test_truncate_after_mm_expansion.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Truncation must bound the prompt *after* multimodal placeholders expand. + +``TokenizeParams.apply_post_tokenization`` truncates the unexpanded prompt, so +one placeholder token can expand into many and push the result back over +``truncate_prompt_tokens`` (vllm-project/vllm#57038). +""" + +import numpy as np +import pytest +import torch + +from vllm.config import ModelConfig, VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.multimodal.inputs import PlaceholderRange +from vllm.renderers import renderer_from_config +from vllm.renderers.base import BaseRenderer +from vllm.renderers.hf import _ensure_prompt_embeds_placeholder_token +from vllm.renderers.params import TokenizeParams + + +def _renderer(): + """A BaseRenderer stub: only the tokenizer attribute is needed here.""" + + class _StubRenderer(BaseRenderer): + def __init__(self): + self.tokenizer = None + + def render_messages(self, *args, **kwargs): + raise NotImplementedError + + return _StubRenderer() + + +def _mm_input(num_tokens: int, offset: int, length: int, **extra): + return { + "type": "multimodal", + "prompt_token_ids": list(range(num_tokens)), + "mm_kwargs": {}, + "mm_hashes": {}, + "mm_placeholders": {"audio": [PlaceholderRange(offset=offset, length=length)]}, + **extra, + } + + +def _params(limit: int | None, side: str = "right"): + return TokenizeParams( + max_total_tokens=4096, + truncate_prompt_tokens=limit, + truncation_side=side, + ) + + +def test_expanded_prompt_is_bounded(): + """The reported case: audio fits, trailing text is cut to the limit.""" + engine_input = _mm_input(num_tokens=147, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(96)) + assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange( + offset=1, length=52 + ) + + +def test_prompt_within_limit_is_untouched(): + engine_input = _mm_input(num_tokens=64, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(64)) + + +def test_split_multimodal_span_is_rejected(): + """Cutting mid-span would desync features from token ids, so reject.""" + engine_input = _mm_input(num_tokens=147, offset=80, length=50) + + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + +def test_dropped_multimodal_span_is_rejected(): + """A span entirely past the cut would be discarded silently.""" + engine_input = _mm_input(num_tokens=147, offset=100, length=20) + + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + +def test_left_truncation_shifts_placeholder_offsets(): + engine_input = _mm_input(num_tokens=147, offset=100, length=20) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96, side="left")) + + assert engine_input["prompt_token_ids"] == list(range(51, 147)) + # The span started at 100 and 51 leading tokens were dropped. + assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange( + offset=49, length=20 + ) + + +def test_text_only_input_is_untouched(): + """Only multimodal prompts grow after tokenization, so nothing else runs.""" + engine_input = { + "type": "token", + "prompt_token_ids": list(range(147)), + "prompt_token_offsets": [(i, i + 1) for i in range(147)], + } + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(147)) + assert len(engine_input["prompt_token_offsets"]) == 147 + + +def test_no_truncation_requested_is_noop(): + engine_input = _mm_input(num_tokens=147, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(None)) + + assert engine_input["prompt_token_ids"] == list(range(147)) + + +def test_mixed_prompt_embeds_truncation_covers_embeds_spans(): + """Bound in _process_singleton, where prompt_embeds spans are already attached.""" + model_config = ModelConfig( + "Qwen/Qwen3-ASR-0.6B-hf", + revision="7f1569a48a89f3e3f4dc3a5c9d28bddd903bc76c", + dtype="float32", + max_model_len=4096, + mm_processor_cache_gb=0, + enable_prompt_embeds=True, + ) + renderer = renderer_from_config(VllmConfig(model_config=model_config)) + tok = renderer.tokenizer + pe_id = _ensure_prompt_embeds_placeholder_token(tok) + hidden = model_config.get_hidden_size() + audio = np.zeros(4 * 16000, dtype=np.float32) + + pe_len = 20 + tensors = [torch.zeros(pe_len, hidden, dtype=torch.float32)] + + def _build_prompt(): + text = "<|audio_start|><|audio_pad|><|audio_end|>" + ids = tok.encode(text) + [pe_id] + tok.encode(" tail" * 40) + return { + "prompt_token_ids": ids, + "multi_modal_data": {"audio": (audio, 16000)}, + "_prompt_embeds": (tensors, pe_id), + } + + # A limit past every span (audio ends at 52, prompt_embeds at 74): + # truncation should succeed and bound the prompt. + out = renderer._process_singleton( + _build_prompt(), + tok_params=_params(79), + ) + assert len(out["prompt_token_ids"]) == 79 + + # A limit inside the prompt_embeds span (54..74): must be rejected. + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + renderer._process_singleton( + _build_prompt(), + tok_params=_params(64), + ) diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index f4916faf6bb4..9b769686a7cb 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -7,11 +7,13 @@ from collections.abc import Mapping, Sequence from concurrent.futures import Future, ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import replace from functools import cached_property from typing import TYPE_CHECKING, Any, Generic, overload from typing_extensions import TypeVar +from vllm.exceptions import VLLMValidationError from vllm.inputs import ( EmbedsInput, EmbedsPrompt, @@ -962,27 +964,103 @@ async def _process_tokens_async( return engine_input + def _truncate_expanded_prompt( + self, + engine_input: TokensInput | MultiModalInput, + params: TokenizeParams | None, + ) -> None: + """Re-apply `truncate_prompt_tokens` once placeholders are expanded. + + `TokenizeParams.apply_post_tokenization` truncates the prompt before + multimodal expansion, so a single placeholder token can expand into + many and push the result back over the requested bound. + + Multimodal spans are never cut: a request whose spans would not all + survive intact is rejected, matching the post-expansion rejection in + `vllm.v1.engine.input_processor._validate_model_input`. + + Only multimodal inputs can grow after `apply_post_tokenization`, so + everything else returns immediately. + """ + if engine_input["type"] != "multimodal": + return + + if params is None or params.truncate_prompt_tokens is None: + return + + prompt_token_ids = engine_input.get("prompt_token_ids") + if prompt_token_ids is None: + return + + num_tokens = len(prompt_token_ids) + truncation = params._truncation_slice(self.tokenizer, num_tokens) + if truncation is None: + return + + start, stop, _ = truncation.indices(num_tokens) + + mm_placeholders = engine_input["mm_placeholders"] + for modality, positions in mm_placeholders.items(): + for position in positions: + end = position.offset + position.length + if position.offset >= start and end <= stop: + continue + raise VLLMValidationError( + f"truncate_prompt_tokens={params.truncate_prompt_tokens} " + f"would split or drop a(n) {modality} item, which spans " + f"tokens {position.offset} to {end} of the " + f"{num_tokens}-token prompt. Multimodal placeholders " + f"expand after truncation is applied, so the limit must " + f"leave every multimodal item intact. Please raise " + f"truncate_prompt_tokens or shorten the multimodal input.", + parameter="truncate_prompt_tokens", + value=params.truncate_prompt_tokens, + ) + + if start: + engine_input["mm_placeholders"] = { + modality: [ + replace(position, offset=position.offset - start) + for position in positions + ] + for modality, positions in mm_placeholders.items() + } + + engine_input["prompt_token_ids"] = prompt_token_ids[truncation] + def _process_singleton( self, prompt: SingletonTokPrompt, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> SingletonInput: if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] - return self._process_tokens(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] + engine_input = self._process_tokens( # type: ignore[arg-type] + prompt, # type: ignore[arg-type] + skip_mm_cache=skip_mm_cache, + ) + self._truncate_expanded_prompt(engine_input, tok_params) + return engine_input async def _process_singleton_async( self, prompt: SingletonTokPrompt, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> SingletonInput: if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] - return await self._process_tokens_async(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] + engine_input = await self._process_tokens_async( # type: ignore[arg-type] + prompt, # type: ignore[arg-type] + skip_mm_cache=skip_mm_cache, + ) + self._truncate_expanded_prompt(engine_input, tok_params) + return engine_input def _get_skip_decoder_start_token(self) -> bool: """Whether the multimodal processor supplies a complete decoder prefix.""" @@ -1048,12 +1126,15 @@ def process_for_engine( arrival_time: float, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: engine_input = self._process_enc_dec(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] else: - engine_input = self._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + engine_input = self._process_singleton( + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) engine_input["arrival_time"] = arrival_time @@ -1065,6 +1146,7 @@ async def process_for_engine_async( arrival_time: float, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: @@ -1074,7 +1156,7 @@ async def process_for_engine_async( ) else: engine_input = await self._process_singleton_async( - prompt, skip_mm_cache=skip_mm_cache + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) engine_input["arrival_time"] = arrival_time @@ -1101,7 +1183,9 @@ def render_cmpl( self._apply_prompt_extras(tok_prompts, prompt_extras) return [ - self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + self.process_for_engine( + prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) for prompt in tok_prompts ] @@ -1126,7 +1210,7 @@ async def render_cmpl_async( return await asyncio.gather( *( self.process_for_engine_async( - p, arrival_time, skip_mm_cache=skip_mm_cache + p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) for p in tok_prompts ) @@ -1164,7 +1248,9 @@ def render_chat( self._apply_prompt_extras(tok_prompts, prompt_extras) eng_prompts = [ - self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + self.process_for_engine( + prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) for prompt in tok_prompts ] @@ -1204,7 +1290,7 @@ async def render_chat_async( eng_prompts = await asyncio.gather( *( self.process_for_engine_async( - p, arrival_time, skip_mm_cache=skip_mm_cache + p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) for p in tok_prompts ) diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 1c35802d0d16..36605c62b914 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1194,22 +1194,20 @@ def _process_tokens( coordinate space, no offset shifting needed afterwards. """ prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) - if prompt_embeds_info is not None: - tensors, placeholder_token_id = prompt_embeds_info - mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) - cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( - list(prompt["prompt_token_ids"]), mm_updates - ) + if prompt_embeds_info is None: + return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) + tensors, placeholder_token_id = prompt_embeds_info + mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) + cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( + list(prompt["prompt_token_ids"]), mm_updates + ) engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) - - if prompt_embeds_info is not None: - tensors, _ = prompt_embeds_info - self._apply_prompt_embeds_to_engine_input( - cast(MultiModalInput, engine_input), - tensors, - mm_updates, - ) + self._apply_prompt_embeds_to_engine_input( + cast(MultiModalInput, engine_input), + tensors, + mm_updates, + ) return engine_input @@ -1222,24 +1220,24 @@ async def _process_tokens_async( ) -> TokensInput | MultiModalInput: """Async equivalent of `_process_tokens`.""" prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) - if prompt_embeds_info is not None: - tensors, placeholder_token_id = prompt_embeds_info - mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) - cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( - list(prompt["prompt_token_ids"]), mm_updates + if prompt_embeds_info is None: + return await super()._process_tokens_async( + prompt, skip_mm_cache=skip_mm_cache ) + tensors, placeholder_token_id = prompt_embeds_info + mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) + cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( + list(prompt["prompt_token_ids"]), mm_updates + ) engine_input = await super()._process_tokens_async( prompt, skip_mm_cache=skip_mm_cache ) - - if prompt_embeds_info is not None: - tensors, _ = prompt_embeds_info - self._apply_prompt_embeds_to_engine_input( - cast(MultiModalInput, engine_input), - tensors, - mm_updates, - ) + self._apply_prompt_embeds_to_engine_input( + cast(MultiModalInput, engine_input), + tensors, + mm_updates, + ) return engine_input From 0e110f696db450f5645ce1daf9945241212d6ef3 Mon Sep 17 00:00:00 2001 From: yinfengLiu Date: Sun, 20 Sep 2026 17:13:55 +0800 Subject: [PATCH 533/774] [ROCm][DSV4][Perf] Use FP8 WO_A output projection (#54894) Signed-off-by: Liuyinfeng01 Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/models/test_deepseek_v4_rocm_wo_a.py | 44 +++++ vllm/models/deepseek_v4/amd/rocm.py | 184 +++++++++++++++++++-- 2 files changed, 217 insertions(+), 11 deletions(-) create mode 100644 tests/models/test_deepseek_v4_rocm_wo_a.py diff --git a/tests/models/test_deepseek_v4_rocm_wo_a.py b/tests/models/test_deepseek_v4_rocm_wo_a.py new file mode 100644 index 000000000000..69c6ed95307c --- /dev/null +++ b/tests/models/test_deepseek_v4_rocm_wo_a.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0 + + +def test_wo_a_block_scale_to_e8m0_from_float(): + scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32) + + encoded = _wo_a_block_scale_to_e8m0(scale) + + assert encoded is not None + torch.testing.assert_close( + encoded, + torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8), + ) + assert encoded.is_contiguous() + + +def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales(): + raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8) + encoded = raw.view(torch.float8_e8m0fnu) + + converted = _wo_a_block_scale_to_e8m0(encoded) + + assert converted is not None + torch.testing.assert_close(converted, raw) + + +@pytest.mark.parametrize( + "scale", + [ + torch.tensor([[0.0, 1.0]]), + torch.tensor([[-1.0, 1.0]]), + torch.tensor([[0.75, 1.0]]), + torch.tensor([[float("inf"), 1.0]]), + torch.ones(1, dtype=torch.int32), + ], +) +def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor): + assert _wo_a_block_scale_to_e8m0(scale) is None diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3cf26f8767ce..a70a23e3ab79 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -7,6 +7,7 @@ import torch +from vllm import envs from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -42,6 +43,49 @@ logger = init_logger(__name__) +def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None: + """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes. + + E8M0 is an unsigned exponent-only scale format with bias 127. A finite + encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``; + ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used + by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP + Microscaling Formats (MX) Specification, section 5.4.1: + https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + + Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``, + or decode it to a floating-point power of two. Floating-point inputs are + accepted only when the original E8M0 byte can be recovered losslessly. + This function does not quantize or round arbitrary scales. + """ + if scale.dtype == torch.float8_e8m0fnu: + # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a + # numeric conversion instead of preserving the encoded exponent byte. + return scale.view(torch.uint8).contiguous() + if scale.dtype == torch.uint8: + # The checkpoint loader already exposed the E8M0 wire representation. + return scale.contiguous() + if not scale.dtype.is_floating_point: + return None + + scale_f32 = scale.detach().float() + if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()): + return None + + # With no mantissa, E8M0 can represent only exact powers of two. Rebuild + # the value before encoding so this adapter never silently quantizes a + # general floating-point checkpoint scale. + exponent = torch.round(torch.log2(scale_f32)) + if not torch.equal(torch.exp2(exponent), scale_f32): + return None + + encoded = exponent.to(torch.int32) + 127 + # 0xFF is NaN in OCP E8M0, not a finite exponent. + if int(encoded.min()) < 0 or int(encoded.max()) > 254: + return None + return encoded.to(torch.uint8).contiguous() + + def _trust_dsv4_extra_cache_nan_free( kv_cache_dtype: str, has_kv_transfer: bool, @@ -555,6 +599,10 @@ def __init__(self, *args, **kwargs): # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None + self._wo_a_fp8_weight: torch.Tensor | None = None + self._wo_a_e8m0_scale: torch.Tensor | None = None + self._wo_a_cos_cache: torch.Tensor | None = None + self._wo_a_sin_cache: torch.Tensor | None = None self._fused_compressor_weight: torch.Tensor | None self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None @@ -815,6 +863,93 @@ def _prep(linear) -> torch.Tensor | None: self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) self._wo_b_scale = _prep(self.wo_b) + if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM: + self._prepare_fp8_wo_a() + + def _prepare_fp8_wo_a(self) -> None: + try: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale as mxscale_op, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant as inverse_quant_op, + ) + except ImportError: + logger.warning_once( + "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; " + "falling back to BF16 WO_A." + ) + return + del mxscale_op, inverse_quant_op + + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_fp8_block_weight_scale, + is_fp8, + ) + + weight = getattr(self.wo_a, "weight", None) + scale = get_fp8_block_weight_scale(self.wo_a) + if scale is None: + # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the + # historical ``_inv`` suffix. + scale = getattr(self.wo_a, "weight_scale", None) + if weight is None or scale is None: + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; " + "the layer exposes no weight/weight scale. Falling back to " + "BF16 WO_A." + ) + return + if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype): + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D " + "block scale, got weight %s%s and scale %s. Falling back to " + "BF16 WO_A.", + weight.dtype, + tuple(weight.shape), + tuple(scale.shape), + ) + return + + groups = self.n_local_groups + out_per_group = self.o_lora_rank + out_features, in_features = weight.shape + if ( + out_features != groups * out_per_group + or out_per_group % 128 != 0 + or in_features % 128 != 0 + or scale.shape != (out_features // 128, in_features // 128) + ): + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of " + "%d outputs, got weight %s and scale %s. Falling back to BF16 " + "WO_A.", + groups, + out_per_group, + tuple(weight.shape), + tuple(scale.shape), + ) + return + + e8m0_scale = _wo_a_block_scale_to_e8m0(scale) + if e8m0_scale is None: + logger.warning_once( + "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a " + "block scale as OCP E8M0. Falling back to BF16 WO_A.", + scale.dtype, + ) + return + + self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features) + self._wo_a_e8m0_scale = e8m0_scale.view( + groups, out_per_group // 128, in_features // 128 + ) + cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None) + if cache is None: + cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16) + cos_cache, sin_cache = cache.chunk(2, dim=-1) + self._wo_a_cos_cache = cos_cache.contiguous() + self._wo_a_sin_cache = sin_cache.contiguous() def prepare_compressor_gemm_fusion(self) -> bool: if self._fused_compressor_weight is not None: @@ -973,17 +1108,44 @@ def _split_qkv_and_norm( ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. - z = rocm_inv_rope_einsum( - self.rotary_emb, - o, - positions, - self.rope_head_dim, - self.n_local_groups, - self.o_lora_rank, - self.wo_a, - ) - zf = z.flatten(1) + if self._wo_a_fp8_weight is not None: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant, + ) + + assert self._wo_a_cos_cache is not None + assert self._wo_a_sin_cache is not None + o_fp8, o_scale = inverse_rope_group_quant( + o.view(o.shape[0], self.n_local_heads, self.head_dim), + positions.to(torch.int64), + self._wo_a_cos_cache, + self._wo_a_sin_cache, + num_groups=self.n_local_groups, + quant_group_size=128, + ) + assert self._wo_a_e8m0_scale is not None + zf = batched_gemm_a8w8_mxscale( + o_fp8, + self._wo_a_fp8_weight, + o_scale, + self._wo_a_e8m0_scale, + dtype=o.dtype, + ).flatten(1) + else: + # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. + z = rocm_inv_rope_einsum( + self.rotary_emb, + o, + positions, + self.rope_head_dim, + self.n_local_groups, + self.o_lora_rank, + self.wo_a, + ) + zf = z.flatten(1) if self._wo_b_scale is not None and zf.dim() == 2: return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) return self.wo_b(zf) From e5fce7b56b07cdf992febd0862bfd599d6dca3fd Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Sun, 20 Sep 2026 05:28:33 -0400 Subject: [PATCH 534/774] [Core][Kernel] Share persistent workspaces for Marlin and Humming (#57421) Signed-off-by: mgoin Co-authored-by: ErinYin --- tests/kernels/moe/test_moe.py | 20 +- .../kernels/moe/test_moe_permute_unpermute.py | 175 ++++++++++++++++-- .../kernels/quantization/test_marlin_gemm.py | 58 ++++++ .../quantization/test_marlin_tile_padding.py | 13 +- .../model_loader/test_reload.py | 67 ++++--- tests/v1/worker/test_workspace.py | 63 +++++++ .../kernels/linear/mixed_precision/marlin.py | 8 +- .../kernels/linear/mxfp4/marlin.py | 2 +- .../kernels/linear/mxfp8/marlin.py | 2 +- .../kernels/linear/nvfp4/marlin.py | 2 +- .../kernels/linear/scaled_mm/marlin.py | 2 +- .../fused_moe/experts/fused_humming_moe.py | 34 ++-- .../layers/fused_moe/experts/marlin_moe.py | 4 +- .../layers/fused_moe/moe_permute_unpermute.py | 47 +++++ .../layers/quantization/auto_awq.py | 4 - .../layers/quantization/auto_gptq.py | 8 - .../compressed_tensors_moe_wna16.py | 14 -- .../layers/quantization/utils/marlin_utils.py | 115 +++++++++--- .../quantization/utils/marlin_utils_fp4.py | 26 +-- .../quantization/utils/marlin_utils_fp8.py | 49 +---- vllm/v1/worker/workspace.py | 96 +++++++--- 21 files changed, 598 insertions(+), 211 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8b4efaf91379..1aca8593ed7c 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1384,24 +1384,33 @@ def test_humming_global_valid_shape_m( assert result == expected -def test_humming_permute_scratch_is_keyed_by_runtime_topk( +def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): from types import SimpleNamespace from unittest.mock import Mock import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming + import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute + import vllm.v1.worker.workspace as workspace + + manager = workspace.WorkspaceManager(torch.device("cpu")) + monkeypatch.setattr(workspace, "_manager", manager) scratch_topk6 = Mock() scratch_topk1 = Mock() scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) - monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type) + monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) moe_config.moe_parallel_config.dp_size = 2 - experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config) + experts = SimpleNamespace(moe_config=moe_config) + other_layer = SimpleNamespace(moe_config=moe_config) assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + assert ( + humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 + ) assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 @@ -1412,6 +1421,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk( assert second_call.kwargs["max_num_tokens"] == 6144 assert second_call.kwargs["topk"] == 1 + manager.lock() + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + with pytest.raises(AssertionError, match="was not allocated during warmup"): + humming.HummingExpertsBase._get_permute_scratch(experts, 2) + def test_humming_delegates_to_instance_activation(): from types import SimpleNamespace diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76d1e57f58d0..52fc8963ece8 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -15,12 +15,14 @@ ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import current_workspace_manager NUM_EXPERTS = [16, 64, 256] TOP_KS = [2, 6, 8] @@ -120,6 +122,7 @@ def torch_unpermute( @pytest.mark.parametrize("topk", TOP_KS) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("ep_size", EP_SIZE) +@pytest.mark.parametrize("use_scratch", [False, True]) def test_moe_permute_unpermute( n_token: int, n_hidden: int, @@ -127,6 +130,8 @@ def test_moe_permute_unpermute( n_expert: int, ep_size: int, dtype: torch.dtype, + use_scratch: bool, + workspace_init, ): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") @@ -160,6 +165,18 @@ def test_moe_permute_unpermute( expert_map=expert_map, ) + scratch = None + if use_scratch: + scratch = get_moe_permute_scratch( + max_num_tokens=n_token, + topk=topk, + num_experts=n_expert, + num_local_experts=n_local_expert, + device=hidden_states.device, + hidden_size=n_hidden, + hidden_dtype=dtype, + ) + ( permuted_hidden_states, _, @@ -173,6 +190,7 @@ def test_moe_permute_unpermute( n_expert=n_expert, n_local_expert=n_local_expert, expert_map=expert_map, + scratch=scratch, ) # check expert_first_token_offset @@ -212,20 +230,22 @@ def test_moe_permute_unpermute( torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("n_token", [1, 33, 128]) +@pytest.mark.parametrize("topk", [1, 6, 8]) +def test_moe_permute_reuses_scratch_buffers( + dtype: torch.dtype, n_token: int, topk: int, workspace_init +): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") - n_token = 64 n_hidden = 2048 n_expert = 16 - topk = 4 hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype) gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype) _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) - scratch = MoEPermuteScratch( + scratch_config = dict( max_num_tokens=n_token, topk=topk, num_experts=n_expert, @@ -234,6 +254,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): hidden_size=n_hidden, hidden_dtype=hidden_states.dtype, ) + scratch = get_moe_permute_scratch(**scratch_config) first = moe_permute( hidden_states=hidden_states, @@ -242,13 +263,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): n_expert=n_expert, scratch=scratch, ) - second = moe_permute( - hidden_states=hidden_states, - a1q_scale=None, - topk_ids=topk_ids, - n_expert=n_expert, - scratch=scratch, - ) + current_workspace_manager().lock() + assert get_moe_permute_scratch(**scratch_config) is scratch + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + second = moe_permute( + hidden_states=hidden_states, + a1q_scale=None, + topk_ids=topk_ids, + n_expert=n_expert, + scratch=get_moe_permute_scratch(**scratch_config), + ) + + for _ in range(2): + hidden_states.add_(1) + topk_ids.copy_(topk_ids.roll(1, 0)) + graph.replay() + expected = moe_permute(hidden_states, None, topk_ids, n_expert) + for actual, reference in zip(second, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) ( permuted_hidden_states_1, @@ -288,6 +322,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): ) +def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: + """Switching graph sizes must not expose metadata left by a larger batch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + device = torch.device("cuda:0") + scratch = get_moe_permute_scratch( + max_num_tokens=128, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + current_workspace_manager().lock() + runs = [] + for n_token in (128, 1, 33): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((graph, hidden, topk_ids, output)) + + for graph, hidden, topk_ids, output in runs + runs[::-1]: + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None: + """Concurrent graph replays in different ubatches/lanes cannot share scratch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + device = torch.device("cuda:0") + manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2) + monkeypatch.setattr(workspace, "_manager", manager) + ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) + scratch_config = dict( + max_num_tokens=64, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + runs = [] + for i, n_token in enumerate((1, 7, 33, 64)): + ubatch = i // 2 + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + scratch = get_moe_permute_scratch(**scratch_config) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((stream, graph, hidden, topk_ids, output, scratch)) + torch.cuda.current_stream().wait_stream(stream) + + assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 + manager.lock() + for _ in range(3): + for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): + ubatch = i // 2 + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + assert get_moe_permute_scratch(**scratch_config) is scratch + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + for stream, _, hidden, topk_ids, output, _ in runs: + torch.cuda.current_stream().wait_stream(stream) + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_without_manager(monkeypatch) -> None: + """Standalone calls get independent scratch with initialized row indices.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + monkeypatch.setattr(workspace, "_manager", None) + config = dict( + max_num_tokens=4, + topk=2, + num_experts=4, + num_local_experts=4, + device=torch.device("cuda"), + ) + first = get_moe_permute_scratch(**config) + second = get_moe_permute_scratch(**config) + + assert ( + first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() + ) + torch.testing.assert_close( + first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda") + ) + + def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py index 4c290dada6b5..29910390bb9a 100644 --- a/tests/kernels/quantization/test_marlin_gemm.py +++ b/tests/kernels/quantization/test_marlin_gemm.py @@ -17,6 +17,7 @@ per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_permute_bias, @@ -43,6 +44,7 @@ ) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): pytest.skip( @@ -477,6 +479,62 @@ def test_marlin_gemm( assert max_diff < 0.04 +def test_marlin_persistent_locks_cuda_graph_streams(workspace_init): + """Compiled GEMMs select stream-local locks at runtime and survive replay.""" + size_m, size_k, size_n = 32, 1024, 2048 + inputs = [rand_data((size_m, size_k)) for _ in range(2)] + weight = rand_data((size_k, size_n)) + quant_type = scalar_types.uint4b8 + w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128) + zeros = marlin_make_empty(scales.device) + streams = [torch.cuda.Stream() for _ in inputs] + graphs = [torch.cuda.CUDAGraph() for _ in inputs] + locks = [] + outputs = [] + + @torch.compile(backend="eager", fullgraph=True) + def run(x): + return torch.ops.vllm.marlin_gemm( + x, + qweight, + None, + scales, + None, + None, + zeros, + None, + quant_type.id, + size_m, + size_n, + size_k, + use_atomic_add=False, + use_fp32_reduce=True, + ) + + for stream, graph, x in zip(streams, graphs, inputs): + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + locks.append(get_marlin_workspace(x.device)) + run(x) + with torch.cuda.graph(graph, stream=stream): + outputs.append(run(x)) + torch.cuda.current_stream().wait_stream(stream) + + assert locks[0].data_ptr() != locks[1].data_ptr() + current_workspace_manager().lock() + for _ in range(3): + for stream, graph, x, lock in zip(streams, graphs, inputs, locks): + with torch.cuda.stream(stream): + assert get_marlin_workspace(x.device) is lock + x.add_(0.1) + graph.replay() + for stream in streams: + torch.cuda.current_stream().wait_stream(stream) + for x, output, lock in zip(inputs, outputs, locks): + assert compute_max_diff(output, x @ w_ref) < 0.04 + assert torch.count_nonzero(lock) == 0 + + def test_marlin_gemm_subset_input(): quant_type = scalar_types.uint4b8 group_size = 128 diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c96..d224b0a2b257 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( GPTQ_MARLIN_TILE, apply_gptq_marlin_linear, + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_moe_padded_intermediate, @@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=layer.bias if use_bias else None, @@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale_inv, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=None, @@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -756,7 +757,7 @@ def q(w): # (out, in) -> fp8 weight, scale, dequant reference topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( @@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index fde8b86cdaef..155957160fb6 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -330,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer): ) -def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init): - """Marlin must reuse the workspace storage captured by CUDA graphs.""" +def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init): + """Weight reload must not create layer-owned Marlin lock storage.""" _stub_marlin_ops(monkeypatch) kernel = _make_marlin_kernel() @@ -339,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - workspace_ptr = kernel.workspace.data_ptr() + assert not hasattr(kernel, "workspace") _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - assert kernel.workspace.data_ptr() == workspace_ptr - assert torch.all(kernel.workspace == 0) + assert not hasattr(kernel, "workspace") @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) -def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): - """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold.""" +def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant): + """Weight preparation must not attach runtime workspace to model layers.""" from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -417,34 +415,57 @@ def load_checkpoint_format_weights(): load_checkpoint_format_weights() prepare(layer) - workspace_ptr = layer.workspace.data_ptr() + assert not hasattr(layer, "workspace") # Reload: fresh checkpoint-format tensors, prepare runs again load_checkpoint_format_weights() prepare(layer) - assert layer.workspace.data_ptr() == workspace_ptr - assert torch.all(layer.workspace == 0) + assert not hasattr(layer, "workspace") -def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): - """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that.""" +def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch): + """Calls reuse initialized locks; independent streams get separate storage.""" from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.utils import torch_utils + from vllm.v1.worker import workspace as workspace_module monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) device = torch.device("cpu") + manager = workspace_module.WorkspaceManager(device) + monkeypatch.setattr(workspace_module, "_manager", manager) + stream = "main" + monkeypatch.setattr(torch_utils, "current_stream", lambda: stream) - workspace = marlin_utils.marlin_make_workspace_new(device) - reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace) - assert reused is workspace + workspace = marlin_utils.get_marlin_workspace(device) + assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert workspace.dtype == torch.int32 + assert torch.count_nonzero(workspace) == 0 - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace) - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new( - device, existing=workspace.to(torch.int64) - ) + workspace.fill_(1) + assert marlin_utils.get_marlin_workspace(device) is workspace + assert torch.all(workspace == 1) + + stream = "aux" + aux_workspace = marlin_utils.get_marlin_workspace(device) + assert aux_workspace.data_ptr() != workspace.data_ptr() + assert torch.count_nonzero(aux_workspace) == 0 + + manager.lock() + assert marlin_utils.get_marlin_workspace(device) is aux_workspace + + +def test_marlin_workspace_without_manager(monkeypatch): + from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.v1.worker import workspace as workspace_module + + monkeypatch.setattr(workspace_module, "_manager", None) + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + first = marlin_utils.get_marlin_workspace(torch.device("cpu")) + second = marlin_utils.get_marlin_workspace(torch.device("cpu")) + assert first.data_ptr() != second.data_ptr() + assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert torch.count_nonzero(first) == 0 def test_model_cleanup(dist_init, default_vllm_config): diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 7824c8000d2b..3d03f25e0e76 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast +from unittest.mock import Mock import pytest import torch @@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None: with pytest.raises(ValueError, match="at least one"): workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) + + +def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None: + """Only the requesting slot allocates; each ubatch/lane owns its resource.""" + active_ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)}) + resources = [] + for ubatch in range(2): + active_ubatch = ubatch + for lane in range(2): + with workspace.use_workspace_lane(lane): + resource = manager.get_persistent_resource("scratch", factory) + resources.append(resource) + resource["buffer"].fill_(len(resources)) + assert manager.get_persistent_resource("scratch", factory) is resource + assert factory.call_count == len(resources) + + assert len({r["buffer"].data_ptr() for r in resources}) == 4 + for i, resource in enumerate(resources, 1): + assert torch.all(resource["buffer"] == i) + + +def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None: + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + factory = Mock(side_effect=object) + first = manager.get_persistent_resource("scratch", factory) + manager.lock() + assert manager.get_persistent_resource("scratch", factory) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("new", factory) + with workspace.use_workspace_lane(1): + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("scratch", factory) + assert factory.call_count == 1 + manager.unlock() + assert manager.get_persistent_resource("scratch", factory) is not first + assert manager.get_persistent_resource("scratch", factory) is first + + +def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None: + """Initialization happens once and transient scratch cannot overwrite locks.""" + manager = workspace.WorkspaceManager(torch.device("cpu")) + first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True) + assert torch.count_nonzero(first) == 0 + first.fill_(1) + (scratch,) = manager.get_simultaneous(((1024,), torch.uint8)) + scratch.zero_() + assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first + assert torch.all(first == 1) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (16,), torch.int32) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (8,), torch.int64) + manager.lock() + assert manager.get_persistent("locks", (8,), torch.int32) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent("new", (8,), torch.int32) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 387c33c11068..60687fc8602d 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -10,7 +10,6 @@ apply_gptq_marlin_linear, marlin_act_int8_process_scales, marlin_make_empty, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -91,11 +90,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: size_k, size_n = c.partition_weight_shape padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) - # Allocate marlin workspace, reusing existing storage on reload. - self.workspace = marlin_make_workspace_new( - device, existing=getattr(self, "workspace", None) - ) - # Default name since marlin requires empty parameter for zp, # TODO: remove this requirement from marlin (allow optional tensors) if self.w_zp_name is None: @@ -197,7 +191,7 @@ def apply_weights( weight=w_q, weight_scale=w_s, weight_zp=w_zp, # type: ignore - workspace=self.workspace, + workspace=None, wtype=c.weight_type, input_size_per_partition=c.partition_weight_shape[0], output_size_per_partition=c.partition_weight_shape[1], diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py index 30ff0d58d4c0..108f835e6435 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py @@ -58,7 +58,7 @@ def apply_weights( weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=None, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py index bec54cd942ed..6ebdbd647405 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -46,7 +46,7 @@ def apply_weights( input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/nvfp4/marlin.py b/vllm/model_executor/kernels/linear/nvfp4/marlin.py index a05d6823c881..1ab5336876bb 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/nvfp4/marlin.py @@ -50,7 +50,7 @@ def apply_weights( weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 3ca89d3cea1f..1df6ad55412d 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -93,7 +93,7 @@ def apply_weights( input=x, weight=layer.weight, weight_scale=weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, input_dtype=self.marlin_input_dtype, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 52b6d4e75334..18a2fb5cc6a8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -176,7 +177,6 @@ def __init__( max_num_tokens=max_num_tokens, num_dispatchers=num_dispatchers, ) - self._permute_scratch: dict[int, MoEPermuteScratch] = {} def init_humming_moe(self): from vllm.utils.humming import get_heuristics_config @@ -259,24 +259,20 @@ def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None - scratch = self._permute_scratch.get(topk) - if scratch is None: - max_expanded_rows = ( - self.moe_config.max_num_tokens - * self.moe_config.dp_size - * self.moe_config.experts_per_token - ) - scratch = MoEPermuteScratch( - max_num_tokens=math.ceil(max_expanded_rows / topk), - topk=topk, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - hidden_size=self.moe_config.hidden_dim, - hidden_dtype=self.moe_config.in_dtype, - ) - self._permute_scratch[topk] = scratch - return scratch + max_expanded_rows = ( + self.moe_config.max_num_tokens + * self.moe_config.dp_size + * self.moe_config.experts_per_token + ) + return get_moe_permute_scratch( + max_num_tokens=math.ceil(max_expanded_rows / topk), + topk=topk, + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + hidden_size=self.moe_config.hidden_dim, + hidden_dtype=self.moe_config.in_dtype, + ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): ctx = get_forward_context() diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index a6a70ca95264..cfd78de28dd1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -35,7 +35,7 @@ from vllm.model_executor.layers.fused_moe.utils import _resize_cache from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, - marlin_make_workspace_new, + get_marlin_workspace, marlin_moe_intermediate_size, marlin_quant_input, ) @@ -95,7 +95,7 @@ def _fused_marlin_moe( N = marlin_moe_intermediate_size(w1, w2) w13_num_shards = 2 if activation.is_gated else 1 if workspace is None: - workspace = marlin_make_workspace_new(hidden_states.device, 4) + workspace = get_marlin_workspace(hidden_states.device) if intermediate_cache13 is None: intermediate_cache13 = torch.empty( diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index 7a53a960e895..0aa9a3807194 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -5,6 +5,11 @@ import torch +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + @dataclass class MoEPermuteScratch: @@ -102,6 +107,48 @@ def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: return topk_ids_int32 +def get_moe_permute_scratch( + *, + max_num_tokens: int, + topk: int, + num_experts: int, + num_local_experts: int, + device: torch.device, + hidden_size: int | None = None, + hidden_dtype: torch.dtype | None = None, +) -> MoEPermuteScratch: + """Share scratch across sequential layers in the current ubatch and lane. + + Without a workspace manager, allocate new scratch for each call. + """ + + def create_scratch() -> MoEPermuteScratch: + return MoEPermuteScratch( + max_num_tokens=max_num_tokens, + topk=topk, + num_experts=num_experts, + num_local_experts=num_local_experts, + device=device, + hidden_size=hidden_size, + hidden_dtype=hidden_dtype, + ) + + if not is_workspace_manager_initialized(): + return create_scratch() + + key = ( + MoEPermuteScratch, + max_num_tokens, + topk, + num_experts, + num_local_experts, + device, + hidden_size, + hidden_dtype, + ) + return current_workspace_manager().get_persistent_resource(key, create_scratch) + + def moe_permute( hidden_states: torch.Tensor, a1q_scale: torch.Tensor | None, diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 71b858e91969..92132f3ffb97 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -48,7 +48,6 @@ check_marlin_supports_layer, check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, verify_marlin_supported, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -642,9 +641,6 @@ def create_weights( layer.register_parameter("w2_qzeros", w2_qzeros) set_weight_attrs(w2_qzeros, extra_weight_attrs) - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index e4478e897dae..ffbc6fba22bf 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -16,7 +16,6 @@ ) from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, - FusedMoEExpertsModular, FusedMoEMethodBase, FusedMoEQuantConfig, FusedMoeWeightScaleSupported, @@ -46,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_repeat_scales_on_all_ranks, verify_marlin_supported, ) @@ -601,12 +599,6 @@ def create_weights( layer.register_parameter("w2_bias", w2_bias) set_weight_attrs(w2_bias, extra_weight_attrs) - if self.experts_cls is not None and issubclass( - self.experts_cls, FusedMoEExpertsModular - ): - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: def replace_or_register(name: str, val: torch.Tensor | None): if val is None: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index bcbb687e3bb9..feddb77cffeb 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -11,7 +11,6 @@ from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoEExpertsModular, RoutedExperts, SharedExperts, ) @@ -36,7 +35,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_config, get_marlin_input_dtype, - marlin_make_workspace_new, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -466,18 +464,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - # Marlin workspace — only needed for Marlin-family backends, not emulation. - if ( - self.experts_cls is not None - and issubclass(self.experts_cls, FusedMoEExpertsModular) - and self.wna16_backend != WNA16MoEBackend.EMULATION - ): - layer.workspace = marlin_make_workspace_new( - layer.w13_weight_packed.device, - 4, - existing=getattr(layer, "workspace", None), - ) - # Alias packed weights to w13_weight/w2_weight for the modular kernel interface layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 82201dff3b29..1ee52a17802b 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -8,7 +8,6 @@ import torch import vllm.envs as envs -from vllm import _custom_ops as ops from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -23,6 +22,7 @@ from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.math_utils import round_up from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import direct_register_custom_op from .quant_utils import pack_cols, unpack_cols @@ -32,6 +32,7 @@ GPTQ_MARLIN_MIN_THREAD_N = 64 GPTQ_MARLIN_MIN_THREAD_K = 128 GPTQ_MARLIN_MAX_PARALLEL = 16 +MARLIN_MAX_BLOCKS_PER_SM = 4 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] @@ -407,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso def marlin_make_workspace_new( device: torch.device, max_blocks_per_sm: int = 1, - existing: torch.Tensor | None = None, ) -> torch.Tensor: # In the new marlin kernel, we use the num of threadblocks as workspace # size. The num of threadblocks is sms_count * max_blocks_per_sm. sms = num_compute_units(device.index) - size = sms * max_blocks_per_sm - # On weight reload, reuse the existing storage so the workspace address - # captured by CUDA graphs stays valid. - if existing is not None: - if ( - existing.device != device - or existing.dtype != torch.int - or existing.numel() != size - ): - raise ValueError( - f"Existing Marlin workspace is incompatible " - f"(device={existing.device}, dtype={existing.dtype}, " - f"numel={existing.numel()}; expected device={device}, " - f"dtype={torch.int}, numel={size}). Reload must reuse the " - f"workspace storage captured by CUDA graphs." - ) - return existing.zero_() - return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False) + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device) + + +def get_marlin_workspace(device: torch.device) -> torch.Tensor: + """Get zero-initialized locks for the current stream, ubatch, and lane. + + Marlin resets its locks after each call, so ordered calls can reuse them. + """ + from vllm.utils.torch_utils import current_stream + from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, + ) + + if not is_workspace_manager_initialized(): + return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM) + + return current_workspace_manager().get_persistent( + ("marlin_locks", current_stream()), + (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,), + torch.int, + zero_init=True, + ) + + +def _marlin_gemm( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + if workspace is None: + workspace = get_marlin_workspace(a.device) + return torch.ops._C.marlin_gemm( + a, + None, + b_q_weight, + b_bias, + b_scales, + a_scales, + global_scale, + b_zeros, + workspace, + b_q_type_id, + size_m, + size_n, + size_k, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ) + + +def _marlin_gemm_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype + return torch.empty((size_m, size_n), device=a.device, dtype=dtype) + + +direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake) def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool: @@ -686,7 +754,7 @@ def apply_gptq_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_zp: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, wtype: ScalarType, output_size_per_partition: int, input_size_per_partition: int, @@ -723,9 +791,8 @@ def apply_gptq_marlin_linear( reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( reshaped_x, - None, weight, bias, weight_scale, @@ -733,7 +800,7 @@ def apply_gptq_marlin_linear( None, weight_zp, workspace, - wtype, + wtype.id, size_m=reshaped_x.shape[0], size_n=padded_n, size_k=padded_k, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index a12d8b87e98c..7fb5621f2f9b 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -10,7 +10,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -161,7 +160,7 @@ def apply_fp4_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_global_scale: torch.Tensor | None, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -196,9 +195,8 @@ def apply_fp4_marlin_linear( inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -206,7 +204,7 @@ def apply_fp4_marlin_linear( global_scale=weight_global_scale, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float4_e2m1f, + b_q_type_id=scalar_types.float4_e2m1f.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -237,13 +235,6 @@ def prepare_fp4_layer_for_marlin( assert layer.weight.shape == (part_size_n, part_size_k // 2) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = layer.weight.view(torch.int32).T.contiguous() @@ -389,15 +380,9 @@ def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: return x return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) - device = w13.device param_dtype = layer.params_dtype is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: @@ -481,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin( k = layer.moe_config.hidden_dim n = layer.moe_config.intermediate_size_per_partition - # WORKSPACE - device = layer.w13_weight.device param_dtype = layer.params_dtype - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 # WEIGHT diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 4a5f5ee59a52..257c5bf44505 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -9,7 +9,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_moe_padded_intermediate, marlin_pad_dim, marlin_pad_qweight, @@ -50,7 +49,7 @@ def apply_fp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None, @@ -80,9 +79,8 @@ def apply_fp8_marlin_linear( # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) raise RuntimeError("Marlin W8A8 is not supported.") - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -90,7 +88,7 @@ def apply_fp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin( else: assert layer.weight.shape == (part_size_n, part_size_k) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first) @@ -241,12 +232,7 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight_scale: torch.Tensor, w2_weight_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Shuffle weights and scales into marlin format. - - Note that this function has the side effect of adding a `workspace` - attribute to the layer. This `workspace` does not need to be - registered as a Parameter as it is not used during weight reloading. - """ + """Shuffle weights and scales into marlin format.""" logger.warning_once( "Your GPU does not have native support for FP8 computation but " "FP8 quantization is being used. Weight-only FP8 compression will " @@ -272,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n) w2_weight = _moe_pad_last(w2_weight, n, padded_n) - # WORKSPACE - device = layer.w13_weight.device - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor: @@ -396,7 +376,7 @@ def apply_mxfp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -416,9 +396,8 @@ def apply_mxfp8_marlin_linear( dtype=input.dtype, ) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=reshaped_x, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -426,7 +405,7 @@ def apply_mxfp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -451,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: group_size = 32 # MX standard block size padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT - repack FP8 weights to Marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) qweight = qweight.T.contiguous() @@ -512,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin( """Repack MXFP8 MoE weights and scales into Marlin kernel format. Args: - layer: MoE layer (used to read params_dtype and attach workspace). + layer: MoE layer. w13: [E, 2*N, K] float8_e4m3fn weights. w2: [E, K, N] float8_e4m3fn weights. w13_scale: [E, 2*N, K//32] uint8 e8m0 scales. @@ -537,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin( n = padded_n w13_n = w13.shape[1] - device = w13.device param_dtype = torch.get_default_dtype() - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: if "w13" in name: size_n, size_k = w13_n, k diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index cddb760acd87..a86110f2f864 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -3,11 +3,12 @@ import inspect import os -from collections.abc import Iterator +from collections.abc import Callable, Hashable, Iterator from contextlib import contextmanager from contextvars import ContextVar from itertools import accumulate from math import prod +from typing import TypeVar, cast import torch @@ -17,6 +18,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) +T = TypeVar("T") def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: @@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]: class WorkspaceManager: """Manager for workspace allocation. - Manages one workspace buffer per active ``(ubatch, lane)`` slot. - Can be locked to prevent further growth during execution. + Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``. + Can be locked to prevent growth or new resources during execution. """ def __init__( @@ -66,6 +68,9 @@ def __init__( self._current_workspaces: list[torch.Tensor | None] = [None] * ( self._num_ubatches * self._num_lanes ) + self._persistent_resources: list[dict[Hashable, object]] = [ + {} for _ in self._current_workspaces + ] self._locked: bool = False @staticmethod @@ -76,10 +81,10 @@ def _workspace_size_bytes(workspace: torch.Tensor | None) -> int: return workspace.numel() * workspace.element_size() def lock(self) -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After locking, any attempt to allocate a larger workspace will raise - an assertion error. This ensures workspace size is fixed during execution. + Larger scratch requests and new persistent resources raise an assertion + error. Existing resources retain their storage throughout execution. """ self._locked = True if envs.VLLM_DEBUG_WORKSPACE: @@ -113,6 +118,63 @@ def is_locked(self) -> bool: """Check if workspace is locked.""" return self._locked + def _get_workspace_id(self) -> int: + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + return dbo_current_ubatch_id() * self._num_lanes + lane + + def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: + """Return the resource cached for ``key`` in the current ubatch and lane. + + The first request calls ``factory`` and stores its result for the lifetime + of this manager. Later requests return the same object without calling + the factory again. Each ubatch/lane pair has a separate cache. Once the + manager is locked, requesting a key missing from that cache raises an + AssertionError. + + Keys must identify the resource type and configuration. Callers must not + resize or replace tensor storage held by a cached resource. Uses of the + same resource must not overlap; include the CUDA stream in the key when + streams can execute concurrently. + + With torch.compile, call this inside a runtime custom op so resource + selection happens during execution. + """ + resources = self._persistent_resources[self._get_workspace_id()] + if key not in resources: + if self._locked: + raise AssertionError( + f"Workspace is locked but {key!r} was not allocated during " + "warmup for the current ubatch and lane." + ) + resources[key] = factory() + return cast(T, resources[key]) + + def get_persistent( + self, + key: Hashable, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + zero_init: bool = False, + ) -> torch.Tensor: + """Get a fixed-shape tensor, optionally zeroed on first allocation only.""" + factory = torch.zeros if zero_init else torch.empty + workspace = self.get_persistent_resource( + (torch.Tensor, key), + lambda: factory(shape, dtype=dtype, device=self._device), + ) + if workspace.shape != shape or workspace.dtype != dtype: + raise ValueError( + f"Persistent workspace {key!r} has shape={workspace.shape}, " + f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}." + ) + return workspace + def get_simultaneous( self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] ) -> list[torch.Tensor]: @@ -151,14 +213,7 @@ def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: The current workspace tensor. """ - ubatch_id = dbo_current_ubatch_id() - lane = _workspace_lane.get() - if lane >= self._num_lanes: - raise RuntimeError( - f"Workspace lane {lane} is not configured; manager has " - f"{self._num_lanes} lane(s)." - ) - workspace_id = ubatch_id * self._num_lanes + lane + workspace_id = self._get_workspace_id() current_workspace = self._current_workspaces[workspace_id] current_size = self._workspace_size_bytes(current_workspace) @@ -218,8 +273,8 @@ def get_caller_info() -> str: get_caller_info(), current_size / _MB, required_bytes / _MB, - ubatch_id, - lane, + workspace_id // self._num_lanes, + workspace_id % self._num_lanes, ) return current_workspace @@ -277,12 +332,11 @@ def init_workspace_manager( def lock_workspace() -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After calling this function, any attempt to allocate a workspace larger - than the current size will raise an AssertionError. This ensures that - workspace size is fixed during execution and prevents unexpected memory - allocations in the hot path. + Larger scratch requests and new persistent resources raise an AssertionError. + All required resources must be initialized in each execution slot during + warmup, so their storage remains fixed during execution. Example: # During initialization From 0748d3bd57cb5c4303e58c8465e21f069629480d Mon Sep 17 00:00:00 2001 From: mahird3 <142054011+mahird3@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:50:00 +0200 Subject: [PATCH 535/774] [Model][LoRA] Enable LoRA support for VoyageQwen3BidirectionalEmbedModel (#57708) Signed-off-by: Mahir Dursunoglu <142054011+mahird3@users.noreply.github.com> --- vllm/model_executor/models/voyage.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/voyage.py b/vllm/model_executor/models/voyage.py index d131ffff4b0b..9f4885a0f9db 100644 --- a/vllm/model_executor/models/voyage.py +++ b/vllm/model_executor/models/voyage.py @@ -8,6 +8,7 @@ import torch.nn as nn from vllm.config import VllmConfig +from vllm.model_executor.models.interfaces import SupportsLoRA from vllm.model_executor.models.qwen3 import Qwen3Model from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix from vllm.sequence import IntermediateTensors @@ -15,7 +16,7 @@ WeightItem = tuple[str, torch.Tensor] -class VoyageQwen3BidirectionalEmbedModel(nn.Module): +class VoyageQwen3BidirectionalEmbedModel(nn.Module, SupportsLoRA): """Qwen3Model + Voyage embedding head + bidirectional attention. Checkpoint conventions (HF): @@ -30,6 +31,11 @@ class VoyageQwen3BidirectionalEmbedModel(nn.Module): - No "model." prefix """ + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config From db1bfdd4fb0dd7b8226402ee00abc7a987561b7c Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 20 Sep 2026 18:15:04 +0800 Subject: [PATCH 536/774] [Bugfix][GLM-5.3-Flash] Address kpool tail blocks by the padded indexer stride in the NVIDIA prefill seed kernel (#57477) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/test_kpool_decode_update_batched.py | 13 ++++++++++--- vllm/models/glm5next/nvidia/ops/kpool_compress.py | 15 +++++++++++++-- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index afc0c1f836b7..37ba4b58c858 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot(): _assert_eq(r_ref, r_kern) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") -def test_amd_prefill_seed_honors_padded_tail_block_stride(): - """The tail shares a padded indexer allocation in production.""" +def test_prefill_seed_honors_padded_tail_block_stride(): + """The tail shares a padded indexer allocation in production. + + ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its + indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash + vs a dense 2048 B tail block), so a seed kernel that addresses blocks + densely writes into an unrelated indexer block and leaves the request's + tail block untouched. Runs on every platform; the NVIDIA kernel had this + bug while the AMD kernel did not. + """ kpool = 4 num_blocks = 6 logical_block_elems = 2 * kpool * HEAD_DIM diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 75ccfa8c2a37..416d3b5c48ea 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -383,6 +383,8 @@ def _kpool_tail_seed_kernel( tslot_ptr, tail_ptr, n_tokens, + TAIL_BLOCK_ELEMS: tl.constexpr, + KPOOL_HEAD: tl.constexpr, HEAD_DIM: tl.constexpr, KPOOL: tl.constexpr, BLOCK_D: tl.constexpr, @@ -393,6 +395,11 @@ def _kpool_tail_seed_kernel( ahead belongs to a different tail block (or is past the batch / padding, slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``. + + The tail cache aliases the indexer cache with the indexer's (padded) block + stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` / + ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense + ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array. """ i = tl.program_id(0) t = tl.load(tslot_ptr + i).to(tl.int64) @@ -409,11 +416,11 @@ def _kpool_tail_seed_kernel( return offs = tl.arange(0, BLOCK_D) m = offs < HEAD_DIM - base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM + base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m) s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m) tl.store(tail_ptr + base + offs, k, mask=m) - tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m) + tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m) def kpool_seed_tail_cache( @@ -426,6 +433,8 @@ def kpool_seed_tail_cache( ) -> None: """Seed the paged tail cache from a prefill batch (see the kernel).""" assert tail_kv_cache.dtype == torch.bfloat16 + assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2 + assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim assert key.dtype == torch.bfloat16 n = tslot.shape[0] if n == 0: @@ -436,6 +445,8 @@ def kpool_seed_tail_cache( tslot, tail_kv_cache, n, + TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0), + KPOOL_HEAD=tail_kv_cache.stride(1), HEAD_DIM=head_dim, KPOOL=kpool, BLOCK_D=triton.next_power_of_2(head_dim), From 1b9fa3eaa8ff23d1a583f2f77dd2b33c9f896ee8 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 20 Sep 2026 18:20:55 +0800 Subject: [PATCH 537/774] [Perf][GLM] Fuse the kpool tail slot mapping into one Triton kernel (#57534) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../attention/test_kpool_tail_slot_mapping.py | 63 +++++++++++++++++++ vllm/v1/attention/backends/mla/indexer.py | 59 ++++++++++++++++- 2 files changed, 120 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index cb39e5403719..ed0e250b1811 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -391,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular(): # The circular mapping keeps the rings isolated under interleaving. torch.testing.assert_close(circular, ground_truth) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize( + "per_req,num_actual,padded_len", + [ + ([list(range(10)), list(range(12))], 22, 22), + ([list(range(10)), list(range(12))], 22, 30), + ([[3, 4], [0], [7, 8, 9]], 6, 8), + ([[5]], 1, 1), + ], +) +def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): + """The CUDA (Triton) path must match the CPU torch reference, including + tokens between the last request boundary and num_actual_tokens (mapped to + the last request) and untouched padding beyond num_actual.""" + positions, qsl, slot_mapping, _, num_reqs = make_batch( + per_req, padded_len=padded_len + ) + # Replace the all--1 placeholder slots with sentinel values to check the + # padding range is copied through untouched. + slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 + bt = make_tail_block_table(list(range(5, 5 + num_reqs))) + + ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) + got = circular_tail_slots( + slot_mapping.cuda(), + bt.cuda(), + qsl.cuda().to(torch.int32), + positions.cuda(), + num_actual, + num_reqs, + ) + torch.testing.assert_close(got.cpu(), ref) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_triton_mapping_reads_strided_block_table(): + """The kernel must address the block table through its real stride(0), + not a dense assumption (the #57477 class of bug): a tail block table that + is a column view of a wider allocation must still read column 0.""" + per_req = [list(range(10)), list(range(12))] + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) + own_blocks = [11, 22] + wide = torch.zeros(num_reqs, 8, dtype=torch.int32) + bt = wide[:, 1:4] # non-contiguous view, stride(0) == 8 + bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) + + ref = circular_tail_slots( + slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs + ) + got = circular_tail_slots( + slot_mapping.cuda(), + bt.cuda(), + qsl.cuda().to(torch.int32), + positions.cuda(), + num_actual, + num_reqs, + ) + torch.testing.assert_close(got.cpu(), ref) + for req, blk in enumerate(own_blocks): + start, end = int(qsl[req]), int(qsl[req + 1]) + assert (got[start:end] // KPOOL == blk).all() diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 430bd10c990c..8af81c8d2900 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -634,6 +634,38 @@ class DeepseekV32IndexerMetadata: prefill: DeepseekV32IndexerPrefillMetadata | None = None +@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) +def _kpool_tail_slot_mapping_kernel( + slot_mapping_ptr, + block_table_ptr, + block_table_stride, + query_start_loc_ptr, + positions_ptr, + out_ptr, + num_reqs, + num_actual_tokens, + num_tokens, + kpool, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + if pid < num_reqs: + start = tl.load(query_start_loc_ptr + pid) + # Tokens past the last request's boundary (if any) also map to it. + end = tl.load(query_start_loc_ptr + pid + 1) + end = tl.where(pid == num_reqs - 1, num_actual_tokens, end) + own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64) + for i in range(start, end, BLOCK): + offs = i + tl.arange(0, BLOCK) + mask = offs < end + pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64) + tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask) + else: + offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK) + mask = offs < num_tokens + tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask) + + def compute_kpool_tail_slot_mapping( slot_mapping: torch.Tensor, block_table: torch.Tensor, @@ -646,10 +678,33 @@ def compute_kpool_tail_slot_mapping( ) -> torch.Tensor: """Map every token to its request's one circular tail block.""" if out is None: - out = slot_mapping.clone() + out = torch.empty_like(slot_mapping) else: assert out.shape == slot_mapping.shape - out.copy_(slot_mapping) + if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0: + block = 256 + num_tokens = slot_mapping.shape[0] + num_actual_tokens = min(num_actual_tokens, num_tokens) + grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),) + _kpool_tail_slot_mapping_kernel[grid]( + slot_mapping, + block_table, + block_table.stride(0), + query_start_loc, + positions, + out, + num_reqs, + num_actual_tokens, + num_tokens, + kpool, + BLOCK=block, + num_warps=4, + ) + return out + # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs. + # Production always takes the Triton path above — spec-decode tokens arrive + # flattened token-major, so slot_mapping is 1D there too. + out.copy_(slot_mapping) if num_actual_tokens == 0: return out tokens = torch.arange(num_actual_tokens, device=slot_mapping.device) From f648eed23dc48fcc8ba64be0c4182f8e775b5bfa Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 20 Sep 2026 18:58:55 +0800 Subject: [PATCH 538/774] [Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail and QSA) (#56810) Signed-off-by: Jared Wen Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: Zhewen Li --- tests/v1/simple_kv_offload/test_kv_events.py | 36 +++ tests/v1/simple_kv_offload/test_scheduler.py | 223 ++++++++++++++++++- vllm/v1/simple_kv_offload/manager.py | 10 +- 3 files changed, 267 insertions(+), 2 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index f21dae7e0bd8..68121f17e000 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -19,8 +19,10 @@ SchedulerFixture, _alloc_and_register, _allocate_cp_gpu_blocks, + _allocate_gpu_blocks, _make_cp_request, _make_kv_cache_config, + _make_scratch_kv_cache_config, _make_vllm_config, make_request, make_scheduler_output, @@ -757,3 +759,37 @@ def test_mamba_align_skips_positional_event_metadata() -> None: get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0)) ) assert ev.block_hashes == [expected_hash] + + +def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None: + """Event metadata is resolved only for prefix-cacheable groups. + + A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not + divide the hash block size, so viewing the request hashes at its block + size is impossible; the store must skip it and still emit BlockStored + for the attention blocks. + """ + fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16)) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + events = list(sched.take_events()) + assert len(events) == num_blocks + assert all(isinstance(e, BlockStored) for e in events) + assert {e.group_idx for e in events} == {0} + assert all(len(e.token_ids) == BLOCK_SIZE for e in events) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index ffc8b204787f..ed07718de60c 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -41,7 +41,9 @@ register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, + KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -125,6 +127,45 @@ def _make_kv_cache_config( ) +def _make_scratch_kv_cache_config( + num_blocks: int, scratch_block_size: int = 4 +) -> KVCacheConfig: + """FullAttention group plus a non-prefix-cacheable scratch group. + + Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request + scratch block of ``index_kpool`` tokens: it is excluded from prefix + caching, so its block size neither divides nor is divided by the + hash block size in general. + """ + register_all_kvcache_specs(vllm_config=None) + fa_config = _make_kv_cache_config(num_blocks, num_groups=1) + scratch_layers = ["layer_scratch"] + scratch_spec = KpoolTailSpec( + block_size=scratch_block_size, + num_kv_heads=2, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + sliding_window=scratch_block_size, + ) + assert not scratch_spec.prefix_cacheable + scratch_bytes = scratch_spec.page_size_bytes * num_blocks + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=fa_config.kv_cache_tensors + + [ + KVCacheTensor( + size=scratch_bytes, + layers=scratch_layers, + layer_stride=scratch_bytes, + block_stride=scratch_spec.page_size_bytes, + ) + ], + kv_cache_groups=fa_config.kv_cache_groups + + [KVCacheGroupSpec(scratch_layers, scratch_spec)], + ) + + def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig: """Minimal VllmConfig for scheduler tests (no GPU).""" model_config = ModelConfig( @@ -174,9 +215,13 @@ def make_scheduler( num_gpu_blocks: int = 16, num_groups: int = 1, lazy: bool = False, + kv_cache_config: KVCacheConfig | None = None, ) -> SchedulerFixture: """Build a SimpleCPUOffloadScheduler with small block pools.""" - kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + if kv_cache_config is None: + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + else: + num_groups = len(kv_cache_config.kv_cache_groups) vllm_config = _make_vllm_config() cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups @@ -1392,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None: ) +# --------------------------------------------------------------------------- +# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load +# --------------------------------------------------------------------------- +def test_scratch_group_excluded_from_store_and_load() -> None: + """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded. + + Its single per-request block carries no hash, so the store must offload + only the attention blocks and the load must pair only attention blocks. + """ + fix = make_scheduler( + num_cpu_blocks=8, + num_gpu_blocks=16, + kv_cache_config=_make_scratch_kv_cache_config(16), + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-scratch-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + scratch_block2 = gpu_pool.get_new_blocks(1) + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2)) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + sched_out2 = make_scheduler_output( + {req2.request_id: 1}, + new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, + ) + meta2 = sched.build_connector_meta(sched_out2) + assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2) + assert len(meta2.load_cpu_blocks) == num_blocks + + # --------------------------------------------------------------------------- # Test 9: Chunked prefill accumulates block_ids across steps # --------------------------------------------------------------------------- @@ -2844,3 +2946,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True + + +QSA_RING_BLOCK_SIZE = 4 + + +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: + register_all_kvcache_specs(vllm_config=None) + full_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring_spec = CircularBufferSpec( + block_size=QSA_RING_BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + ) + groups = [ + KVCacheGroupSpec(["layer_0"], full_spec), + KVCacheGroupSpec(["qsa_ring"], ring_spec), + ] + tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, + ) + for group in groups + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _make_qsa_scheduler( + lazy: bool = False, +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + sched = SimpleCPUOffloadScheduler( + vllm_config=_make_vllm_config(), + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + lazy_offload=lazy, + ) + gpu_pool = BlockPool( + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE + ) + sched.bind_gpu_block_pool(gpu_pool) + return sched, gpu_pool + + +def test_qsa_ring_group_is_never_stored_or_loaded() -> None: + sched, gpu_pool = _make_qsa_scheduler() + assert sched.prefix_cacheable_group_ids == (0,) + + num_blocks = 2 + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + ring_block = gpu_pool.get_new_blocks(1)[0] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} + assert ring_block.block_id not in meta.store_gpu_blocks + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-qsa-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + ring_block2 = gpu_pool.get_new_blocks(1)[0] + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} + assert ring_block2.block_id not in meta2.load_gpu_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + without_ring = KVCacheConfig( + num_blocks=16, + kv_cache_tensors=with_ring.kv_cache_tensors[:1], + kv_cache_groups=with_ring.kv_cache_groups[:1], + ) + max_batched = 64 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + with_ring, max_batched + ) + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + without_ring, max_batched + ) + assert target_with == target_without diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 85dde88f2d55..96bfb35a6760 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -165,6 +165,9 @@ def __init__( kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks + self.prefix_cacheable_group_ids = ( + self.cpu_kv_cache_config.prefix_cacheable_group_ids + ) self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU # Find the full attention kv group for prefix cache matching. self.fa_gidx = -1 @@ -297,7 +300,7 @@ def _estimate_lazy_target_blocks( """GPU blocks to keep available (free/offloaded) per step in lazy mode.""" WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks target = 0 - for g in kv_cache_config.kv_cache_groups: + for g in kv_cache_config.prefix_cacheable_groups: spec = g.kv_cache_spec # Only full attention is sharded across DCP ranks; replicated specs # (mamba, sliding window, chunked-local) keep their own block size. @@ -442,6 +445,9 @@ def update_state_after_alloc( # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): + if g not in self.prefix_cacheable_group_ids: + cpu_hit_blocks.append([]) + continue g_block_size = self.group_block_sizes[g] n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -828,6 +834,8 @@ def _select_eager_blocks_to_store( num_free = self.cpu_block_pool.get_num_free_blocks() for g, group_gpu_ids in enumerate(block_ids_by_group): + if g not in self.prefix_cacheable_group_ids: + continue if len(gpu_block_ids) >= num_free: break group_manager = self.cpu_coordinator.single_type_managers[g] From bf01fc4a313cc90863c470f37a26e41fd7b16fd1 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 20 Sep 2026 19:17:59 +0800 Subject: [PATCH 539/774] [GLM-5.3-Flash] Route kpool indexer top-k through the shared SparseIndexerTopk dispatcher (#57546) Signed-off-by: chaunceyjiang --- tests/kernels/test_top_k_per_row.py | 31 ++++++--- .../test_sparse_indexer_topk_dispatch.py | 54 ++++++++-------- vllm/model_executor/layers/indexer_topk.py | 28 ++++++-- vllm/models/glm5next/nvidia/sparse_indexer.py | 64 ++++++------------- 4 files changed, 94 insertions(+), 83 deletions(-) diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index db49b8091b40..b6c469181f09 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -1871,17 +1871,32 @@ def resolve( has_coop = _has_cooperative_topk() has_fi = _has_flashinfer_topk() - # "auto" is exactly the pre-existing chain: cooperative -> persistent -> - # per_row. It must never select the opt-in backends by itself. + # "auto" is the chain cooperative -> persistent -> per_row, with the + # cooperative/persistent switch at cooperative_topk's cluster-wave row + # limit. It must never select the opt-in backends by itself. if has_coop: - assert resolve("auto") == "cooperative" + from vllm.model_executor.layers.indexer_topk import ( + AUTO_COOPERATIVE_MAX_ROWS, + ) + + # cooperative for every batch within its cluster-wave row limit... + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative" + ) assert resolve("auto", k=512, num_rows=8) == "cooperative" - # Past cooperative's 64-row limit -> persistent, even where - # DeepSelect would be applicable. + # ...persistent past the row limit, even where DeepSelect would be + # applicable. + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1) + == "persistent" + ) assert resolve("auto", k=512, num_rows=128) == "persistent" - # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: - # still cooperative at <= 64 rows... - assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative" + # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the + # same row-limit logic applies. + assert ( + resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS) + == "cooperative" + ) # ...and persistent past the row limit. assert resolve("auto", unaligned_logits, num_rows=128) == "persistent" # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels) diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py index 83dfdc9f6ffd..9f94da459c15 100644 --- a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Focused unit tests for the GLM decode top-k dispatch.""" +"""Tests for the GLM kpool indexer's top-k backend wiring.""" import pytest import torch -from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk +from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -13,29 +13,31 @@ ) -@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"]) -@torch.inference_mode() -def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int): - logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda") - assert _use_cooperative_topk(logits, 512, num_rows) +def _require_deep_gemm() -> None: + from vllm.utils.deep_gemm import has_deep_gemm + if not has_deep_gemm(): + pytest.skip("kpool indexer requires DeepGEMM") -@pytest.mark.parametrize( - ("num_rows", "select_k", "stride_padding"), - [ - (65, 512, 0), - (1, 256, 0), - (1, 512, 1), - ], - ids=["too_many_rows", "unsupported_k", "unaligned_stride"], -) -@torch.inference_mode() -def test_cooperative_topk_falls_back_on_unsupported_batch( - num_rows: int, select_k: int, stride_padding: int -): - logits = torch.empty( - num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda" - ) - if stride_padding: - logits = logits[:, :8192] - assert not _use_cooperative_topk(logits, select_k, num_rows) + +@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: + """The kpool indexer must read kernel_config.sparse_indexer_topk_backend + and hand it to the shared SparseIndexerTopk dispatcher, rather than + hard-coding its own cooperative/persistent/per_row choice.""" + _require_deep_gemm() + from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool + + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + op = SparseAttnIndexerKpool( + k_cache=None, + quant_block_size=128, + scale_fmt="ue8m0", + topk_tokens=2048, + head_dim=128, + max_pool_len=4096, + max_total_seq_len=8192, + topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"), + ) + assert op.topk_backend == backend diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index 2ddf71070307..c3422445b2c5 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -14,6 +14,12 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 +# cooperative_topk's hard row limit: one cluster wave covers at most this many +# rows. "auto" uses cooperative_topk for every batch within the limit and falls +# back to persistent_topk past it; explicit ``cooperative`` requests are +# validated against the same bound. +AUTO_COOPERATIVE_MAX_ROWS = 64 + # --------------------------------------------------------------------------- # DeepSelect (vllm._deepselect_C) # --------------------------------------------------------------------------- @@ -209,8 +215,13 @@ def resolve_backend( def _resolve_auto( self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: - """The pre-existing priority chain: cooperative -> persistent -> - per_row. deep_select/flashinfer/torch are opt-in only.""" + """The priority chain: cooperative -> persistent -> per_row. + deep_select/flashinfer/torch are opt-in only. + + cooperative_topk is preferred whenever it is applicable, i.e. within + its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to + persistent_topk. + """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" if self._is_cuda and topk_tokens in (512, 1024, 2048): @@ -228,8 +239,10 @@ def _cooperative_constraints( failures.append( f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" ) - if num_rows > 64: - failures.append(f"num_rows must be <= 64, got {num_rows}") + if num_rows > AUTO_COOPERATIVE_MAX_ROWS: + failures.append( + f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}" + ) if logits.stride(0) % 4 != 0: failures.append( f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}" @@ -293,13 +306,18 @@ def forward( (topk_workspace,) = current_workspace_manager().get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) + # persistent_topk's last argument is the host-side "max seq_len + # across rows" bound (it gates the sampled_topk path and clamps + # per-row lengths), not the padded row width: logits is sized to + # max_model_len, so passing its width would incorrectly enable the + # long-row kernels on short-context batches. torch.ops._C.persistent_topk( logits, seq_lens, topk_indices, topk_workspace, topk_tokens, - logits.shape[1], + max_seq_len, ) elif backend == "flashinfer": # Deferred: importing flashinfer initializes CUDA at import time. diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index f7204949fc5c..c80e44201571 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -10,6 +10,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -39,21 +40,6 @@ logger = init_logger(__name__) -def _use_cooperative_topk( - logits: torch.Tensor, - select_k: int, - num_rows: int, -) -> bool: - return ( - current_platform.is_cuda() - and select_k in (512, 1024, 2048) - and num_rows <= 64 - and logits.stride(0) % 4 == 0 - and current_platform.has_device_capability(90) - and not current_platform.is_device_capability_family(120) - ) - - # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -134,6 +120,7 @@ def sparse_attn_indexer_kpool( # path and when the tail cache is disabled. tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -584,35 +571,18 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and select_k in (512, 1024, 2048): - workspace_manager = current_workspace_manager() - (topk_workspace,) = workspace_manager.get_simultaneous( - ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) - topk_op = ( - torch.ops._C.cooperative_topk - if _use_cooperative_topk(logits, select_k, num_rows) - else torch.ops._C.persistent_topk - ) - topk_op( - logits, - seq_lens, - topk_dst, - topk_workspace, - select_k, - attn_metadata_narrowed.max_seq_len, - ) - else: - torch.ops._C.top_k_per_row_decode( - logits, - next_n, - seq_lens, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) + # Shared dispatcher with the DSA sparse indexer, so the kpool select + # honors kernel_config.sparse_indexer_topk_backend and picks between + # cooperative/persistent by batch size instead of always taking the + # same kernel. + get_indexer_topk(topk_backend)( + logits, + seq_lens, + next_n, + topk_dst, + select_k, + attn_metadata_narrowed.max_seq_len, + ) # Resolve to token-level indices in the output buffer. if index_kpool > 1: @@ -692,6 +662,11 @@ def __init__( "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." ) _cfg = get_current_vllm_config_or_none() + self.topk_backend = ( + _cfg.kernel_config.sparse_indexer_topk_backend + if _cfg is not None + else "auto" + ) _parallel = _cfg.parallel_config if _cfg is not None else None if ( _parallel is not None @@ -766,4 +741,5 @@ def forward_cuda( positions, self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, + self.topk_backend, ) From 10e6a7f21094b76c65efb029b39b3f2227516418 Mon Sep 17 00:00:00 2001 From: Hongxin Xu <70438206+xhx1022@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:21:33 +0800 Subject: [PATCH 540/774] [Frontend] Expose per-request spec decode metrics in generate API (#43310) Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Codex --- .../engine-core-client/src/protocol/output.rs | 32 +++- rust/src/llm/src/output.rs | 12 +- rust/src/llm/tests/generate.rs | 1 + .../server/src/routes/inference/generate.rs | 141 ++++++++++++++++- .../src/routes/inference/generate/types.rs | 87 +++++++++++ .../scale_out/derender/test_derender.py | 45 ++++++ .../derender/test_derender_stream.py | 24 +++ .../test_generate_stream.py | 142 ++++++++++++++++++ .../entrypoints/scale_out/derender/serving.py | 4 + .../scale_out/token_in_token_out/protocol.py | 8 +- .../scale_out/token_in_token_out/serving.py | 26 +++- vllm/renderers/online_derenderer.py | 3 + 12 files changed, 516 insertions(+), 9 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index a45cbe37315e..0edd0ce788bd 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -130,10 +130,36 @@ pub struct EngineCoreOutput { #[serde(default)] pub new_sampling_mask: Option, /// Per-request speculative-decoding acceptance metrics, set on the final - /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here; - /// the Rust frontend does not yet surface it in responses. + /// output when `--per-request-spec-decode-metrics` is enabled. #[serde(default)] - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option, +} + +/// Raw per-sequence speculative-decoding accumulator. +/// +/// Python models this as a plain `@dataclass`, so it is serialized by msgspec +/// as a map (named fields). Derived values such as acceptance rates are +/// computed by the frontend, not carried on the wire. +/// +/// Original Python definition: `RequestSpecDecodeMetrics` in +/// `vllm/v1/metrics/stats.py`. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RequestSpecDecodeMetrics { + /// Configured `num_speculative_tokens`. + #[serde(default)] + pub num_spec_tokens: u64, + /// Verify-step counts indexed by accepted draft-token count. + #[serde(default)] + pub histogram: Vec, + /// Total proposed draft tokens. + #[serde(default)] + pub num_draft_tokens: u64, + /// Accepted draft count per verify step; empty unless `detailed`. + #[serde(default)] + pub per_step_accepted: Vec, + /// Proposed draft count per verify step; empty unless `detailed`. + #[serde(default)] + pub per_step_drafted: Vec, } impl EngineCoreOutput { diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index d0cab4a5076c..34eae2568238 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -11,7 +11,9 @@ use futures::stream::FusedStream; use futures::{Stream, StreamExt as _, pin_mut}; use serde::{Deserialize, Serialize}; use vllm_engine_core_client::protocol::logprobs::Logprobs; -use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason}; +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason, +}; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; @@ -47,6 +49,8 @@ pub struct CollectedGenerateOutput { pub ec_transfer_params: Option, /// Sampling support sets aligned one-to-one with generated token positions. pub sampling_mask: Option, + /// Per-request speculative-decoding metrics from the terminal output. + pub spec_decode_metrics: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -160,6 +164,8 @@ pub struct GenerateOutput { pub ec_transfer_params: Option, /// Sampling support sets aligned one-to-one with `token_ids`. pub sampling_mask: Option, + /// Per-request speculative-decoding metrics, present on terminal outputs. + pub spec_decode_metrics: Option, } impl GenerateOutput { @@ -208,6 +214,7 @@ impl GenerateOutput { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, } } } @@ -308,6 +315,7 @@ impl Stream for GenerateOutputStream { kv_transfer_params: raw.kv_transfer_params, ec_transfer_params: raw.ec_transfer_params, sampling_mask, + spec_decode_metrics: raw.spec_decode_metrics, }; Poll::Ready(Some(Ok(output))) @@ -396,6 +404,7 @@ impl> + Send> T { kv_transfer_params: None, ec_transfer_params: None, sampling_mask, + spec_decode_metrics: None, }); } @@ -409,6 +418,7 @@ impl> + Send> T { }; collected.kv_transfer_params = output.kv_transfer_params; collected.ec_transfer_params = output.ec_transfer_params; + collected.spec_decode_metrics = output.spec_decode_metrics; if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 09b519976772..c341737a3c68 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -413,6 +413,7 @@ async fn collect_output_rejects_partial_sampling_mask() { sampling_mask: Some(SamplingMask { rows: vec![vec![1, 33, 99]], }), + spec_decode_metrics: None, }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index f82059dffda4..eb3338a67f82 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -29,7 +29,8 @@ use vllm_llm::{ use self::convert::{ResponseOptions, prepare_generate_request}; use self::types::{ GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, - GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo, + GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo, + SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics, }; pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; @@ -160,10 +161,14 @@ async fn generate_chunk_stream( let mut prompt_tokens = None; let mut prompt_token_ids = None; let mut usage = TokenUsage::default(); + let mut spec_decode_metrics = None; while let Some(next) = stream.next().await { match next { Ok(output) => { + if let Some(metrics) = output.spec_decode_metrics { + spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics)); + } if prompt_tokens.is_none() && let Some(info) = output.prompt_info.as_ref() { @@ -223,6 +228,7 @@ async fn generate_chunk_stream( .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()), prompt_token_ids, + metrics: None, }) .await; } @@ -243,6 +249,11 @@ async fn generate_chunk_stream( usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)), prompt_token_ids: None, mm_placeholders: None, + metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics { + speculative_decoding: StreamingSpeculativeDecodingMetrics::from( + speculative_decoding, + ), + }), }) .await; } @@ -318,6 +329,9 @@ fn collect_generate( mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(), kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, + metrics: collected.spec_decode_metrics.map(|metrics| PerRequestMetrics { + speculative_decoding: SpeculativeDecodingMetrics::from(metrics), + }), }) } @@ -465,6 +479,7 @@ mod tests { use futures::{TryStreamExt as _, stream}; use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; + use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_llm::GeneratePromptInfo; use super::*; @@ -482,6 +497,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -496,6 +512,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, }), ]); @@ -578,6 +595,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, } } @@ -645,6 +663,80 @@ mod tests { assert!(chunks[1].mm_placeholders.is_none()); } + #[tokio::test] + async fn generate_chunk_stream_returns_terminal_spec_decode_metrics() { + let metrics = RequestSpecDecodeMetrics { + num_spec_tokens: 3, + histogram: vec![0, 1, 1, 0], + num_draft_tokens: 6, + ..Default::default() + }; + let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos())); + output.spec_decode_metrics = Some(metrics.clone()); + + for include_usage in [false, true] { + let chunks: Vec<_> = generate_chunk_stream( + stream::iter(vec![Ok(output.clone())]), + "raw-stream".to_string(), + ApiServerOptions::default(), + ResponseOptions { + include_usage, + ..Default::default() + }, + None, + ) + .try_collect() + .await + .expect("collect chunks"); + assert_eq!(chunks.len(), 1 + usize::from(include_usage)); + assert!(chunks[0].metrics.is_none()); + if include_usage { + assert!(chunks[1].choices.is_empty()); + assert_eq!( + chunks[1].metrics, + Some(PerRequestMetrics { + speculative_decoding: StreamingSpeculativeDecodingMetrics::from( + SpeculativeDecodingMetrics::from(metrics.clone()) + ) + }) + ); + let json = serde_json::to_value(&chunks[1]).expect("serialize chunk"); + let stats = &json["metrics"]["speculative_decoding"]; + assert_eq!(stats.as_object().expect("metrics object").len(), 7); + assert_eq!(stats["num_draft_tokens"], 6); + assert_eq!(stats["num_accepted_draft_tokens"], 3); + assert_eq!(stats["mean_acceptance_length"], 2.5); + assert_eq!( + stats["acceptance_histogram"], + serde_json::json!([0, 1, 1, 0]) + ); + assert!(stats.get("per_step_accepted").is_none()); + } + } + } + + #[test] + fn spec_decode_metrics_zero_steps_and_detailed() { + let empty = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics { + num_spec_tokens: 2, + histogram: vec![0, 0, 0], + ..Default::default() + }); + assert_eq!(empty.mean_acceptance_length, 1.0); + assert_eq!(empty.draft_acceptance_rate, 0.0); + assert_eq!(empty.per_step_accepted, None); + + let detailed = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics { + num_spec_tokens: 2, + histogram: vec![0, 1, 1], + num_draft_tokens: 4, + per_step_accepted: vec![1, 2], + per_step_drafted: vec![2, 2], + }); + assert_eq!(detailed.per_step_accepted, Some(vec![1, 2])); + assert_eq!(detailed.per_step_drafted, Some(vec![2, 2])); + } + #[test] fn collect_generate_omits_prompt_metadata_by_default() { let output = CollectedGenerateOutput { @@ -662,6 +754,7 @@ mod tests { ec_transfer_params: None, prompt_token_ids: vec![10, 20], sampling_mask: None, + spec_decode_metrics: None, }; let response = collect_generate( @@ -683,6 +776,50 @@ mod tests { assert!(response.mm_placeholders.is_none()); } + #[test] + fn collect_generate_returns_spec_decode_metrics() { + let metrics = RequestSpecDecodeMetrics { + num_spec_tokens: 3, + histogram: vec![0, 1, 1, 0], + num_draft_tokens: 6, + ..Default::default() + }; + let output = CollectedGenerateOutput { + request_id: "raw-metrics".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage::default(), + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10], + sampling_mask: None, + spec_decode_metrics: Some(metrics.clone()), + }; + + let response = collect_generate( + output, + "raw-metrics".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + None, + ) + .expect("response"); + + assert_eq!( + response.metrics, + Some(PerRequestMetrics { + speculative_decoding: SpeculativeDecodingMetrics::from(metrics) + }) + ); + let json = serde_json::to_value(&response).expect("serialize response"); + let stats = &json["metrics"]["speculative_decoding"]; + assert_eq!(stats.as_object().expect("metrics object").len(), 9); + assert!(stats["per_step_accepted"].is_null()); + assert!(stats["per_step_drafted"].is_null()); + } + #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); @@ -749,6 +886,7 @@ mod tests { ec_transfer_params: None, prompt_token_ids: vec![10, 20], sampling_mask: None, + spec_decode_metrics: None, }; let response = collect_generate( @@ -782,6 +920,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 406ec585692b..86d8ec00f4fc 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -7,6 +7,7 @@ use llm_multimodal::MediaContentPart; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; +use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage}; @@ -84,6 +85,7 @@ pub(super) struct GenerateStreamResponse { pub usage: Option, pub prompt_token_ids: Option>, pub mm_placeholders: Option, + pub metrics: Option>, } /// Mirrors the Python vLLM `GenerateResponse` class. @@ -96,6 +98,91 @@ pub(super) struct GenerateResponse { pub mm_placeholders: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, + pub metrics: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct PerRequestMetrics { + pub speculative_decoding: T, +} + +/// Mirrors the Python vLLM `SpeculativeDecodingMetrics` class. +/// +/// Derived from the raw engine accumulator the same way as Python +/// `RequestSpecDecodeMetrics.to_dict`. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct SpeculativeDecodingMetrics { + pub mean_acceptance_length: f64, + pub draft_acceptance_rate: f64, + pub acceptance_histogram: Vec, + pub num_spec_steps: u64, + pub num_accepted_draft_tokens: u64, + pub num_draft_tokens: u64, + pub num_spec_tokens: u64, + pub per_step_accepted: Option>, + pub per_step_drafted: Option>, +} + +impl From for SpeculativeDecodingMetrics { + fn from(raw: RequestSpecDecodeMetrics) -> Self { + let num_spec_steps: u64 = raw.histogram.iter().sum(); + let num_accepted_draft_tokens: u64 = + (0u64..).zip(&raw.histogram).map(|(accepted, count)| accepted * count).sum(); + let ratio = |num: u64, den: u64| { + if den == 0 { + 0.0 + } else { + num as f64 / den as f64 + } + }; + let detailed = !raw.per_step_accepted.is_empty(); + Self { + mean_acceptance_length: if num_spec_steps == 0 { + 1.0 + } else { + 1.0 + ratio(num_accepted_draft_tokens, num_spec_steps) + }, + draft_acceptance_rate: ratio(num_accepted_draft_tokens, raw.num_draft_tokens), + acceptance_histogram: raw.histogram, + num_spec_steps, + num_accepted_draft_tokens, + num_draft_tokens: raw.num_draft_tokens, + num_spec_tokens: raw.num_spec_tokens, + per_step_accepted: detailed.then_some(raw.per_step_accepted), + per_step_drafted: detailed.then_some(raw.per_step_drafted), + } + } +} + +/// Streaming form omits detailed fields when summary metrics are requested. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct StreamingSpeculativeDecodingMetrics { + pub mean_acceptance_length: f64, + pub draft_acceptance_rate: f64, + pub acceptance_histogram: Vec, + pub num_spec_steps: u64, + pub num_accepted_draft_tokens: u64, + pub num_draft_tokens: u64, + pub num_spec_tokens: u64, + pub per_step_accepted: Option>, + pub per_step_drafted: Option>, +} + +impl From for StreamingSpeculativeDecodingMetrics { + fn from(metrics: SpeculativeDecodingMetrics) -> Self { + Self { + mean_acceptance_length: metrics.mean_acceptance_length, + draft_acceptance_rate: metrics.draft_acceptance_rate, + acceptance_histogram: metrics.acceptance_histogram, + num_spec_steps: metrics.num_spec_steps, + num_accepted_draft_tokens: metrics.num_accepted_draft_tokens, + num_draft_tokens: metrics.num_draft_tokens, + num_spec_tokens: metrics.num_spec_tokens, + per_step_accepted: metrics.per_step_accepted, + per_step_drafted: metrics.per_step_drafted, + } + } } pub(super) type MultiModalPlaceholders = HashMap>; diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 9cc6124393d8..80ecc0f49cce 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -52,6 +52,7 @@ def _make_generate_response( logprobs: dict | None = None, prompt_logprobs: list | None = None, kv_transfer_params: dict | None = None, + metrics: dict | None = None, ) -> dict: choice: dict = { "index": 0, @@ -64,6 +65,7 @@ def _make_generate_response( "choices": [choice], "prompt_logprobs": prompt_logprobs, "kv_transfer_params": kv_transfer_params, + "metrics": metrics, } @@ -266,6 +268,49 @@ async def test_derender_chat_kv_transfer_params_passthrough(client): assert response.json()["kv_transfer_params"] == kv +@pytest.mark.asyncio +async def test_derender_chat_metrics_passthrough(client): + gen_req = await _render_chat(client) + metrics = { + "speculative_decoding": { + "mean_acceptance_length": 2.0, + "draft_acceptance_rate": 0.5, + "acceptance_histogram": [0, 1], + "num_spec_steps": 1, + "num_accepted_draft_tokens": 1, + "num_draft_tokens": 2, + "num_spec_tokens": 1, + } + } + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": _make_generate_response( + gen_req["token_ids"][:3], metrics=metrics + ), + }, + ) + assert response.status_code == 200 + actual = response.json()["metrics"] + assert set(actual) == { + "time_to_first_token_ms", + "generation_time_ms", + "queue_time_ms", + "mean_itl_ms", + "tokens_per_second", + "speculative_decoding", + } + assert all( + actual[name] is None for name in actual if name != "speculative_decoding" + ) + assert actual["speculative_decoding"] == { + **metrics["speculative_decoding"], + "per_step_accepted": None, + "per_step_drafted": None, + } + + @pytest.mark.asyncio async def test_derender_chat_empty_token_ids(client): """Empty token_ids list returns 400.""" diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 4032d9cbf308..e7a2e3b856cf 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -27,6 +27,8 @@ DeltaFunctionCall, DeltaMessage, DeltaToolCall, + PerRequestMetrics, + SpeculativeDecodingMetrics, ) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderStreamState, @@ -164,6 +166,7 @@ def _make_stream_chunk( finish_reason: str | None = None, request_id: str = "test-req", usage: dict | None = None, + metrics: PerRequestMetrics | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" return GenerateStreamResponse( @@ -176,6 +179,7 @@ def _make_stream_chunk( ) ], usage=UsageInfo(**usage) if usage else None, + metrics=metrics, ) @@ -512,6 +516,26 @@ async def test_role_on_first_chunk_only(self, derenderer, tokenizer): assert chunk1.choices[0].delta.role == "assistant" assert chunk2.choices[0].delta.role is None + @pytest.mark.asyncio + async def test_metrics_passthrough(self, derenderer): + metrics = PerRequestMetrics( + speculative_decoding=SpeculativeDecodingMetrics( + mean_acceptance_length=2.0, + draft_acceptance_rate=0.5, + acceptance_histogram=[0, 1], + num_spec_steps=1, + num_accepted_draft_tokens=1, + num_draft_tokens=2, + num_spec_tokens=1, + ) + ) + chunk, _ = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([], metrics=metrics), + state=DerenderStreamState(), + ) + assert chunk.metrics == metrics + @pytest.mark.asyncio async def test_chunked_equals_oneshot(self, derenderer, tokenizer): """Sum of streaming content deltas == one shot decode.""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index d9ae424cc85b..ba5839990c67 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -28,6 +28,7 @@ from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics MODEL_NAME = "openai-community/gpt2" BASE_MODEL_PATHS = [ @@ -130,6 +131,7 @@ def _make_request_output( logprobs: list[dict[int, Any] | None] | None = None, num_cached_tokens: int | None = None, index: int = 0, + spec_decode_metrics: RequestSpecDecodeMetrics | None = None, ) -> RequestOutput: return RequestOutput( request_id=request_id, @@ -144,6 +146,7 @@ def _make_request_output( cumulative_logprob=None, logprobs=logprobs, finish_reason=finish_reason, + spec_decode_metrics=spec_decode_metrics, ) ], finished=finished, @@ -321,6 +324,145 @@ async def mock_generate(*args, **kwargs): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_serve_tokens_returns_spec_decode_metrics(stream: bool): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=stream, + stream_options=StreamOptions(include_usage=True) if stream else None, + ) + + response = await serving.serve_tokens(request) + if stream: + chunks = [chunk async for chunk in response] + payload = next( + chunk + for chunk in _parse_sse_chunks(chunks) + if isinstance(chunk, dict) and chunk.get("metrics") + )["metrics"]["speculative_decoding"] + else: + assert isinstance(response, GenerateResponse) + assert response.metrics is not None + assert response.metrics.speculative_decoding is not None + payload = response.metrics.speculative_decoding.model_dump() + + assert payload["num_draft_tokens"] == 3 + assert payload["num_accepted_draft_tokens"] == 2 + assert payload["num_spec_steps"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output( + include_usage, +): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", + token_ids=[], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=include_usage), + ) + + chunks = [chunk async for chunk in await serving.serve_tokens(request)] + data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"] + + assert "metrics" not in data_chunks[0] + assert len(data_chunks) == 1 + int(include_usage) + if include_usage: + assert data_chunks[1]["choices"] == [] + assert ( + data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling( + stream: bool, +): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + result = _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + result.outputs.append( + CompletionOutput( + index=1, + text="", + token_ids=[20], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + spec_decode_metrics=metrics, + ) + ) + + async def mock_generate(*args, **kwargs): + yield result + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1, n=2), + model=MODEL_NAME, + stream=stream, + stream_options=StreamOptions(include_usage=True) if stream else None, + ) + + response = await serving.serve_tokens(request) + if stream: + chunks = [chunk async for chunk in response] + assert all( + "metrics" not in chunk + for chunk in _parse_sse_chunks(chunks) + if isinstance(chunk, dict) + ) + else: + assert isinstance(response, GenerateResponse) + assert response.metrics is None + + @pytest.mark.asyncio async def test_non_stream_returns_final_prompt_token_ids_when_requested(): engine = _mock_engine() diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index d3c5f3aac1cb..5da736640373 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -179,6 +179,7 @@ async def derender_chat_response( usage=usage, prompt_logprobs=gen.prompt_logprobs, kv_transfer_params=gen.kv_transfer_params, + metrics=gen.metrics, ) async def derender_completion_response( @@ -245,6 +246,9 @@ async def derender_completion_response( choices=choices, usage=usage, kv_transfer_params=kv_params, + # Metrics describe one prompt. Parallel samples are already + # suppressed by /generate; multi-prompt responses cannot be merged. + metrics=first.metrics if len(request.generate_responses) == 1 else None, ) async def derender_chat_stream_response( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index ae0fd62b6efb..088ec4d2931f 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -11,7 +11,11 @@ ) from vllm.config import ModelConfig -from vllm.entrypoints.generate.base.protocol import StreamOptions, validate_cache_salt +from vllm.entrypoints.generate.base.protocol import ( + PerRequestMetrics, + StreamOptions, + validate_cache_salt, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProbs, ChatCompletionRequest, @@ -360,6 +364,7 @@ class GenerateStreamResponse(BaseModel): usage: UsageInfo | None = Field(default=None) prompt_token_ids: list[int] | None = None mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None + metrics: PerRequestMetrics | None = None class GenerateResponse(BaseModel): @@ -378,6 +383,7 @@ class GenerateResponse(BaseModel): prompt_logprobs: list[dict[int, Logprob] | None] | None = None prompt_token_ids: list[int] | None = None mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None + metrics: PerRequestMetrics | None = None kv_transfer_params: dict[str, Any] | None = Field( default=None, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 7813c7facc60..59fda2dbf8cf 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -12,9 +12,13 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker -from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.generate.base.protocol import ( + PerRequestMetrics, + RequestResponseMetadata, +) from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, + build_spec_decoding_metrics, clamp_prompt_logprobs, ) from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -373,6 +377,11 @@ async def serve_tokens_full_generator( request_metadata.final_usage_info = usage + per_request_metrics = None + if request.sampling_params.n == 1: + spec_stats = build_spec_decoding_metrics(final_res) + if spec_stats is not None: + per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats) response = GenerateResponse( request_id=request_id, created=created_time, @@ -384,6 +393,7 @@ async def serve_tokens_full_generator( final_res.prompt_token_ids if request.return_token_ids else None ), mm_placeholders=request._response_mm_placeholders, + metrics=per_request_metrics, kv_transfer_params=final_res.kv_transfer_params, ec_transfer_params=final_res.ec_transfer_params, ) @@ -423,6 +433,7 @@ async def serve_tokens_stream_generator( prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params + last_res: RequestOutput | None = None include_usage, include_continuous_usage = should_include_usage( request.stream_options, False @@ -430,6 +441,7 @@ async def serve_tokens_stream_generator( try: async for res in result_generator: + last_res = res if first_iteration: if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) @@ -496,10 +508,10 @@ async def serve_tokens_stream_generator( total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent prompt metadata (Rust skips None fields). + # Omit fields that are absent from token-bearing chunks. exclude = { name - for name in ("prompt_token_ids", "mm_placeholders") + for name in ("prompt_token_ids", "mm_placeholders", "metrics") if getattr(chunk, name) is None } yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" @@ -517,10 +529,18 @@ async def serve_tokens_stream_generator( ) if include_usage: + per_request_metrics = None + if sampling_params.n == 1: + spec_stats = build_spec_decoding_metrics(last_res) + if spec_stats is not None: + per_request_metrics = PerRequestMetrics( + speculative_decoding=spec_stats + ) final_chunk = GenerateStreamResponse( request_id=request_id, choices=[], usage=final_usage_info, + metrics=per_request_metrics, ) yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index a89bc843b81e..9dfa0c412fb9 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -398,6 +398,7 @@ async def derender_chat_stream( model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state @@ -595,6 +596,7 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state @@ -741,6 +743,7 @@ async def derender_completion_stream( model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state From 9b2f34cad446f73b1699e8236ec0b611a65f48af Mon Sep 17 00:00:00 2001 From: Zyann <62597503+Zyann7@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:25:56 +0800 Subject: [PATCH 541/774] [Feature] support bf16 MoE router and mxfp4 MoE for MiMo V2 (#57784) Signed-off-by: Zyann7 <62597503+Zyann7@users.noreply.github.com> Co-authored-by: Abatom <182586866+Abatom@users.noreply.github.com> --- vllm/model_executor/models/mimo_v2.py | 13 ++++++++----- vllm/model_executor/models/mimo_v2_omni.py | 5 ++++- vllm/model_executor/models/qwen3_dflash.py | 9 +++++++++ 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index be665e7e064a..dfbed256551e 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -27,6 +27,7 @@ FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -156,14 +157,16 @@ def __init__( dtype = getattr(config, "moe_router_dtype", "float32") self.gate_dtype = str_dtype_to_torch_dtype(dtype) - self.gate = nn.Linear( + self.gate = GateLinear( config.hidden_size, config.n_routed_experts, bias=False, - dtype=self.gate_dtype, + params_dtype=self.gate_dtype, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", ) self.gate.e_score_correction_bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=self.gate_dtype) + torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype) ) self.experts = FusedMoEFactory( @@ -182,7 +185,7 @@ def __init__( num_expert_group=config.n_group, topk_group=config.topk_group, scoring_func="sigmoid", - router_logits_dtype=self.gate_dtype, + router_logits_dtype=self.gate.out_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -198,7 +201,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: gate_input = hidden_states.to(self.gate_dtype) else: gate_input = hidden_states - router_logits = self.gate(gate_input) + router_logits, _ = self.gate(gate_input) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 91a3177b7f50..f55cf44f3787 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -52,6 +52,7 @@ from .interfaces import ( MultiModalEmbeddings, + SupportsEagle3, SupportsMultiModal, SupportsPP, SupportsQuant, @@ -1209,7 +1210,9 @@ def get_dummy_mm_data( info=MiMoV2OmniProcessingInfo, dummy_inputs=MiMoV2OmniDummyInputsBuilder, ) -class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant): +class MiMoV2OmniForCausalLM( + nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3 +): # To ensure correct weight loading and mapping. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 5504304c733a..1636e3852c0f 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -168,6 +168,7 @@ def __init__( rms_norm_eps: float = 1e-06, attention_bias: bool = False, add_swa_attention_sink_bias: bool = False, + v_scale: float | None = None, sliding_window: int | None = None, causal: bool = False, is_neox_style: bool = True, @@ -193,6 +194,7 @@ def __init__( self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 + self.v_scale = v_scale self.qkv_proj = QKVParallelLinear( hidden_size, @@ -264,6 +266,8 @@ def forward( q, k = self.rotary_emb(positions, q, k) + if self.v_scale is not None: + v = v * self.v_scale attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -312,6 +316,7 @@ def __init__( rms_norm_eps=config.rms_norm_eps, attention_bias=getattr(config, "attention_bias", False), add_swa_attention_sink_bias=add_swa_attention_sink_bias, + v_scale=dflash_config.get("attention_value_scale"), sliding_window=sliding_window, causal=causal, is_neox_style=is_neox_style, @@ -633,6 +638,10 @@ def precompute_and_store_context_kv( if context_slot_mapping is None: return + v_scale = getattr(self.layers[0].self_attn, "v_scale", None) + if v_scale is not None: + all_v.mul_(v_scale) + # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) per_layer = isinstance(context_slot_mapping, (list, tuple)) From 27757dde020ecda4f9b0e2c5ca2df1c29badc82f Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 20 Sep 2026 21:51:13 +0800 Subject: [PATCH 542/774] [DSV4.1] Add encoder cuda graph support for deepseek-v4.1-flash (#56625) Signed-off-by: Isotr0py --- .../generation/test_vit_cudagraph.py | 23 ++ .../processing/test_tensor_schema.py | 6 + vllm/models/deepseek_v4/common/vision.py | 163 +++++++- vllm/models/deepseek_v41/amd/vl_model.py | 40 +- .../deepseek_v41/common/vl_cudagraph.py | 390 ++++++++++++++++++ vllm/models/deepseek_v41/nvidia/vl_model.py | 40 +- 6 files changed, 599 insertions(+), 63 deletions(-) create mode 100644 vllm/models/deepseek_v41/common/vl_cudagraph.py diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 1ed84ebc5275..9bda7ec85db9 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str: ) +def deepseek_v41_chat_template(content: str) -> str: + return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜>" + + def step3_vl_chat_template(content: str) -> str: return ( "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n " @@ -221,6 +225,25 @@ def minicpmv_chat_template(content: str) -> str: vllm_runner_kwargs={"gpu_memory_utilization": 0.80}, marks=[pytest.mark.core_model], ), + "deepseek_v41": VitCudagraphTestConfig( + model="deepseek-ai/DeepSeek-V4.1-Flash", + modalities=["image"], + image_prompt=deepseek_v41_chat_template( + "<｜deepseek_image｜>\n\nWhat is in this image?" + ), + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1024], + }, + vllm_runner_kwargs={ + "load_format": "dummy", + "attention_backend": "FLASHMLA_SPARSE_DSV41", + "hf_overrides": partial( + dummy_hf_overrides, + model_arch="DeepseekV41ForCausalLM", + exist_overrides={"vision_n_layers": 1}, + ), + }, + ), "step3_vl": VitCudagraphTestConfig( model="stepfun-ai/Step3-VL-10B", modalities=["image"], diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 403a01a9e2bb..1e1722537d4d 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -218,8 +218,14 @@ def test_model_tensor_schema(model_id: str): factories = model_cls._processor_factory + # Capture helpers return capture-buffer containers (e.g. + # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs. + capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"} + inputs_parse_methods = [] for attr_name in dir(model_cls): + if attr_name in capture_helpers: + continue attr = getattr(model_cls, attr_name) if hasattr(attr, "__annotations__"): return_type = attr.__annotations__.get("return", None) diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py index 4c633da6c549..bfc02768c06c 100644 --- a/vllm/models/deepseek_v4/common/vision.py +++ b/vllm/models/deepseek_v4/common/vision.py @@ -40,8 +40,7 @@ ) -@lru_cache(8) -def get_vision_cos_sin( +def _compute_vision_cos_sin( n_h: int, n_w: int, dim: int, theta: float ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) @@ -52,6 +51,13 @@ def get_vision_cos_sin( return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) +@lru_cache(8) +def get_vision_cos_sin( + n_h: int, n_w: int, dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor]: + return _compute_vision_cos_sin(n_h, n_w, dim, theta) + + def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: dtype = x.dtype x1, x2 = x.float().chunk(2, dim=-1) @@ -123,15 +129,23 @@ def __init__(self, config, prefix: str = ""): ) def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: n = x.size(0) qkv, _ = self.wqkv(x) q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) k = apply_rotary(k, cos, sin).unsqueeze(0) - # One image per call: a dense batch, no varlen packing metadata. - o = self.attn(q, k, v.unsqueeze(0)) + # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens + # set): packed multi-image batch for encoder CUDA graph replay. + o = self.attn( + q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen + ) out, _ = self.wo(o.reshape(n, -1)) return out @@ -173,9 +187,14 @@ def __init__(self, config, prefix: str = ""): self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), cos, sin) + x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen) return x + self.mlp(self.norm2(x)) @@ -206,6 +225,115 @@ def forward( x = block(x, cos, sin) return self.norm(x) + def forward_packed( + self, + patches: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + ) -> torch.Tensor: + """Varlen path for encoder CUDA graphs: multiple images packed along + rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed + on device. Pure-tensor (no host reads, no H2D), safe to capture.""" + x = self.patch_embed(patches) + for block in self.blocks: + x = block(x, cos, sin, cu_seqlens, max_seqlen) + return self.norm(x) + + +def build_packed_vit_metadata( + grids: list[list[int]] | list[tuple[int, int]], + *, + rope_dim: int, + rope_theta: float, + device: torch.device, + max_seqlen_override: int | None = None, + cached: bool = True, +) -> dict[str, torch.Tensor]: + """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``. + + Args: + grids: ``[n_vit_h, n_vit_w]`` per image, in packing order. + rope_dim: RoPE embedding dimension of the ViT. + rope_theta: RoPE base frequency of the ViT. + device: Device for the returned tensors (except ``max_seqlen``, + which stays on the host). + max_seqlen_override: Worst-case value baked in at CUDA graph capture + (the attention wrapper reads ``max_seqlen`` on the host, so the + capture-time value becomes a graph constant). + cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time + dummy grids pass ``False`` to avoid evicting real entries. + + """ + cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin + cos_list: list[torch.Tensor] = [] + sin_list: list[torch.Tensor] = [] + cu_seqlens = [0] + max_seqlen = 0 + for n_h, n_w in grids: + cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta) + cos_list.append(cos) + sin_list.append(sin) + n = n_h * n_w + cu_seqlens.append(cu_seqlens[-1] + n) + max_seqlen = max(max_seqlen, n) + if cos_list: + cos = torch.cat(cos_list).to(device) + sin = torch.cat(sin_list).to(device) + else: + cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + if max_seqlen_override is not None: + max_seqlen = max_seqlen_override + return { + "vit_cos": cos, + "vit_sin": sin, + "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device), + # Read on the host by the attention wrapper; keep on CPU. + "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32), + } + + +def build_packed_merge_metadata( + grids: list[list[int]] | list[tuple[int, int]], + downsample_ratio: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``. + + For each aligner output row, the ``r x r`` source patch rows in the + packed ViT output; positions past the image edge (the eager path pads + them with zeros) get mask 0 and a clamped in-range index. + """ + r = downsample_ratio + idx_list: list[torch.Tensor] = [] + mask_list: list[torch.Tensor] = [] + offset = 0 + for n_h, n_w in grids: + rows = -(-n_h // r) + cols = -(-n_w // r) + bi = torch.arange(rows).repeat_interleave(cols) + bj = torch.arange(cols).repeat(rows) + di = torch.arange(r).repeat_interleave(r) + dj = torch.arange(r).repeat(r) + hi = bi[:, None] * r + di[None, :] + wj = bj[:, None] * r + dj[None, :] + valid = (hi < n_h) & (wj < n_w) + idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1) + idx_list.append(idx) + mask_list.append(valid) + offset += n_h * n_w + if idx_list: + merge_idx = torch.cat(idx_list).to(device) + merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype) + else: + merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device) + merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device) + return {"merge_idx": merge_idx, "merge_mask": merge_mask} + class DeepseekV4Aligner(nn.Module): """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" @@ -240,6 +368,27 @@ def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: out, _ = self.w2(F.gelu(hidden)) return out + def forward_packed( + self, + x: torch.Tensor, + merge_idx: torch.Tensor, + merge_mask: torch.Tensor, + ) -> torch.Tensor: + """Packed multi-image path for encoder CUDA graphs. + + ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``; + gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge + exactly (unfold rows are channel-major, so the gathered patch rows + are transposed before flattening). + """ + r2 = self.downsample_ratio**2 + m = merge_idx.shape[0] + gathered = (x[merge_idx] * merge_mask).view(m, r2, -1) + gathered = gathered.permute(0, 2, 1).reshape(m, -1) + hidden, _ = self.w1(gathered) + out, _ = self.w2(F.gelu(hidden)) + return out + def run_dp_sharded_vision_tower( vision_model: DeepseekV4ViT, diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 3227b50a4e45..3c57df2ad258 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -26,6 +26,7 @@ from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ def _parse_and_validate_image_input( resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py new file mode 100644 index 000000000000..e2a718b7453c --- /dev/null +++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower. + +Mixin implementing the ``SupportsEncoderCudaGraph`` protocol +(``vllm/model_executor/models/interfaces.py``) on top of the shared +``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole +batch of images into one varlen run: ViT blocks attend per image via +``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask +against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/ +IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is +batched over all items of the packed batch. + +Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the +manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP +sharding and eager fallback. +""" + +from collections.abc import Hashable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.models.deepseek_v4.common.vision import ( + build_packed_merge_metadata, + build_packed_vit_metadata, +) +from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, +) + +from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.config.multimodal import MultiModalConfig + from vllm.models.deepseek_v4.common.vision import ( + DeepseekV4Aligner, + DeepseekV4ViT, + ) + from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config + + +class DeepseekV4VLEncoderCudaGraphMixin: + """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``. + + Expects ``self.vision``/``self.aligner``/``self.config``/ + ``self.multimodal_config`` and the ``image_start``/``image_newline``/ + ``image_end`` parameters from the host model. + """ + + # Provided by the host model; declared here so mixin methods type-check. + vision: "DeepseekV4ViT" + aligner: "DeepseekV4Aligner" + config: "DeepseekV41Config" + multimodal_config: "MultiModalConfig | None" + image_start: torch.nn.Parameter + image_end: torch.nn.Parameter + image_newline: torch.nn.Parameter + + # -- helpers shared with the eager multimodal path -- + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + ) -> torch.Tensor: + # Aligner rows in reading order, one per IMAGE slot. + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def _build_image_span( + self, image_embeds: torch.Tensor, types: torch.Tensor + ) -> torch.Tensor: + """Full image span: aligner rows at IMAGE slots, the learned + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" + types = types.to(image_embeds.device) + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) + dtype = image_embeds.dtype + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + span[types == IMAGE] = image_embeds + return span + + @staticmethod + def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: + grid = mm_kwargs[key] + if isinstance(grid, torch.Tensor): + grid = grid.tolist() + return [[int(v) for v in row] for row in grid] + + _encoder_cg_pad_totals: dict[int, int] = {} + + # -- SupportsEncoderCudaGraph protocol -- + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + mm_cfg = self.multimodal_config + if ( + mm_cfg is not None + and mm_cfg.mm_encoder_attn_dtype == "fp8" + and mm_cfg.mm_encoder_fp8_scale_path is None + ): + raise ValueError( + "cudagraph_mm_encoder is incompatible with dynamic FP8 " + "scaling for the vision encoder (the amax history is " + "host-side state that cannot be captured). Provide a static " + "scale file via --mm-encoder-fp8-scale-path or disable " + "cudagraph_mm_encoder." + ) + + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + backend = self.vision.blocks[0].attn.attn.attn_backend + # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls + # .tolist() on CUDA cu_seqlens; neither can be captured. + if backend in ( + AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.TORCH_SDPA, + ): + raise ValueError( + f"cudagraph_mm_encoder is not supported with the {backend.name} " + "ViT attention backend for DeepSeek-V4.1. Set " + "--mm-encoder-attn-backend FLASH_ATTN or disable " + "cudagraph_mm_encoder." + ) + + pad_totals = self._encoder_cg_pad_totals + + def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: + # Varlen attention requires cu_seqlens[-1] to equal the number of + # rows actually passed in. The captured buffers are sized for the + # full token budget, so a smaller real batch is completed with one + # trailing padding sequence; declaring fewer rows than the buffer + # holds is undefined behaviour and returns NaN on FlashAttn. + total = pad_totals.get(dst.data_ptr()) + if total is None: + raise RuntimeError( + "cu_seqlens replay buffer was not registered at capture " + "time; the manager must replay into the capture-time " + "buffers." + ) + n = min(src.shape[0], dst.shape[0]) + dst[:n].copy_(src[:n]) + dst[n:] = total + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "patches", + "vit_cos", + "vit_sin", + "cu_seqlens", + "max_seqlen", + "merge_idx", + "merge_mask", + ], + out_hidden_size=self.config.hidden_size, + padding_logics={"cu_seqlens": pad_cu_seqlens}, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: "VllmConfig", + ) -> tuple[int, int]: + # Min: comfortably above the smallest possible image span (4 tokens + # for a 1x1 grid); also sets max_batch_size = 64 via the manager's + # auto-inference (max_batch_size <= min_budget invariant). + min_budget = 64 + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ) -> list[EncoderItemSpec]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + return [ + EncoderItemSpec( + input_size=h * w, + output_tokens=lh * (lw + 1) + 2, + ) + for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True) + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + patches = mm_kwargs["patches"] + types = mm_kwargs["types"] + + if len(indices) == 0: + return { + "patches": patches[:0], + "vit_grid": torch.zeros((0, 2), dtype=torch.int64), + "llm_grid": torch.zeros((0, 2), dtype=torch.int64), + "types": types[:0], + } + + cum_patches = [0] + for h, w in vit_grid: + cum_patches.append(cum_patches[-1] + h * w) + cum_spans = [0] + for lh, lw in llm_grid: + cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2) + + return { + "patches": torch.cat( + [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ), + "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64), + "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64), + "types": torch.cat( + [types[cum_spans[i] : cum_spans[i + 1]] for i in indices] + ), + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, + ) -> EncoderCudaGraphCaptureInputs: + if ( + self.multimodal_config is not None + and self.multimodal_config.get_limit_per_prompt("image") == 0 + ): + raise RuntimeError( + "cudagraph_mm_encoder requires the vision tower, but it was " + "stubbed out because the image limit is 0. Disable " + "cudagraph_mm_encoder." + ) + + config = self.config + r = config.vision_downsample_ratio + # Span rows per dummy item (ceiling so the buffers tile the budget). + # Each dummy image is a single aligner row of per_item_out blocks: + # r x (r * per_item_out) patches. One real image using the whole + # budget needs at most r^2 * token_budget patches, which this fits. + per_item_out = -(-token_budget // max_batch_size) + grids = [[r, r * per_item_out]] * max_batch_size + total_patches = r * r * per_item_out * max_batch_size + + p = config.vision_patch_size + patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype) + + metadata = build_packed_vit_metadata( + grids, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=device, + max_seqlen_override=total_patches, + cached=False, + ) + # Spare cu_seqlens slots let replay append a padding sequence that + # covers the rows a smaller real batch does not fill. + real_cu = metadata.pop("cu_seqlens") + cu_seqlens = torch.full( + (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device + ) + cu_seqlens[: real_cu.numel()] = real_cu + + merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) + + self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches + + return EncoderCudaGraphCaptureInputs( + values={ + "patches": patches, + **metadata, + "cu_seqlens": cu_seqlens, + **merge, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ) -> EncoderCudaGraphReplayBuffers: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + patches = mm_kwargs["patches"] + dtype = self.aligner.w1.weight.dtype + if patches.dtype != dtype: + patches = patches.to(dtype) + + # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and + # pad_cu_seqlens appends the padding sequence covering the tail rows. + metadata = build_packed_vit_metadata( + vit_grid, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=patches.device, + ) + merge = build_packed_merge_metadata( + vit_grid, + self.config.vision_downsample_ratio, + device=patches.device, + dtype=dtype, + ) + return EncoderCudaGraphReplayBuffers( + values={"patches": patches, **metadata, **merge} + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + vit_out = self.vision.forward_packed( + inputs["patches"], + inputs["vit_cos"], + inputs["vit_sin"], + inputs["cu_seqlens"], + inputs["max_seqlen"], + ) + return self.aligner.forward_packed( + vit_out, inputs["merge_idx"], inputs["merge_mask"] + ) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype) + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + outs: list[torch.Tensor] = [] + offset = 0 + for n_vit_h, n_vit_w in vit_grid: + n_vit = n_vit_h * n_vit_w + outs.append( + self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w) + ) + offset += n_vit + if not outs: + return patches.new_zeros((0, self.config.hidden_size)) + return torch.cat(outs, dim=0) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + if not indices: + return + assert batch_mm_kwargs is not None + aligner_out = outputs["default"] + r = self.config.vision_downsample_ratio + vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid") + types = batch_mm_kwargs["types"].to(aligner_out.device) + + n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid) + span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid] + + # Batched span assembly: one masked fill per role across all items. + dtype = aligner_out.dtype + span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size) + span[types == IMAGE] = aligner_out[:n_rows] + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + + # Freshly allocated, so later replays cannot clobber the results. + offset = 0 + for idx, span_len in zip(indices, span_lens): + dest[idx] = span[offset : offset + span_len] + offset += span_len diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index bfbcaa89a1a1..901985bdd7fb 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -26,6 +26,7 @@ from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ def _parse_and_validate_image_input( resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs, From 7d99c2c4fd61cdeebd29cc7b60584bad10d1ef99 Mon Sep 17 00:00:00 2001 From: Jinzhen Lin Date: Sun, 20 Sep 2026 23:11:01 +0800 Subject: [PATCH 543/774] [Feature][Humming] Humming feature integration (#56685) Signed-off-by: jinzhen.ljz Signed-off-by: mgoin Co-authored-by: Codex Co-authored-by: mgoin --- .buildkite/test_areas/lm_eval.yaml | 86 ++- .../moe/moe_align_sum_kernels.cu | 66 +- csrc/libtorch_stable/moe/moe_ops.h | 12 +- .../moe/moe_permute_unpermute_op.cu | 93 ++- .../moe_permute_unpermute_kernel.cu | 54 +- .../moe_permute_unpermute_kernel.h | 13 +- csrc/libtorch_stable/moe/torch_bindings.cpp | 10 +- requirements/cuda.txt | 2 +- .../humming/Qwen3-30B-A3B-AWQ-humming.yaml | 2 + ...B-FP8-block-humming-deepep-ht-grouped.yaml | 14 + ...B-FP8-block-humming-deepep-ht-indexed.yaml | 14 + .../Qwen3-30B-A3B-MXFP4A16-humming.yaml | 2 + .../humming/Qwen3-30B-A3B-NVFP4-humming.yaml | 2 + ...B-A3B-int5wc-hadamard-humming-act-fp8.yaml | 13 - ...A3B-int6wc-hadamard-humming-act-int8.yaml} | 7 +- ...n3.5-35B-A3B-int6wc-hadamard-humming.yaml} | 7 +- ...5-9B-int6wc-hadamard-humming-act-int8.yaml | 13 + .../Qwen3.5-9B-int6wc-hadamard-humming.yaml | 12 + .../humming/config-a100-act-int8-shard-0.txt | 3 + .../humming/config-a100-act-int8-shard-1.txt | 4 + .../configs/humming/config-a100-shard-0.txt | 1 + .../gsm8k/configs/humming/config-act-int8.txt | 1 + .../gsm8k/configs/humming/config-deepep.txt | 2 + .../configs/humming/config-h100-shard-0.txt | 1 + .../humming/config-int5wc-hadamard.txt | 3 - .../evals/gsm8k/configs/humming/config-l4.txt | 9 + tests/kernels/core/test_activation.py | 151 +++-- tests/kernels/moe/test_moe.py | 101 ++- .../kernels/moe/test_moe_align_block_size.py | 40 ++ .../kernels/moe/test_moe_permute_unpermute.py | 27 +- .../test_mxfp4_kernel_selection.py | 3 +- .../test_humming_mxfp4_block_fp8.py | 20 +- tests/quantization/test_moe_wna16.py | 2 +- vllm/_custom_ops.py | 2 + vllm/envs.py | 3 +- .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/mixed_precision/humming.py | 5 +- .../kernels/linear/mxfp4/humming.py | 25 +- .../kernels/linear/mxfp6/base.py | 4 +- .../kernels/linear/mxfp6/humming.py | 84 +++ .../kernels/linear/mxfp8/humming.py | 10 +- .../kernels/linear/nvfp4/humming.py | 20 +- .../kernels/linear/scaled_mm/humming.py | 4 +- .../fused_moe/experts/fused_humming_moe.py | 430 ++++++------- .../layers/fused_moe/moe_align_block_size.py | 68 +- .../layers/fused_moe/moe_permute_unpermute.py | 31 + .../layers/fused_moe/oracle/fp8.py | 4 +- .../layers/fused_moe/oracle/int8.py | 4 +- .../layers/fused_moe/oracle/int_wna16.py | 2 +- .../layers/fused_moe/oracle/mxfp4.py | 6 +- .../layers/fused_moe/oracle/nvfp4.py | 4 +- .../layers/quantization/auto_awq.py | 2 +- .../layers/quantization/auto_gptq.py | 2 +- .../compressed_tensors_moe_wna16.py | 2 +- .../layers/quantization/humming.py | 158 +++-- .../layers/quantization/moe_wna16.py | 2 +- .../quantization/utils/humming/__init__.py | 51 ++ .../quantization/utils/humming/activation.py | 126 ++++ .../quantization/utils/humming/linear.py | 183 ++++++ .../{humming_utils.py => humming/moe.py} | 604 ++---------------- .../quantization/utils/humming/schema.py | 437 +++++++++++++ vllm/utils/humming.py | 55 +- 62 files changed, 2038 insertions(+), 1084 deletions(-) create mode 100644 tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml create mode 100644 tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml delete mode 100644 tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml rename tests/evals/gsm8k/configs/humming/{Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml => Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml} (62%) rename tests/evals/gsm8k/configs/humming/{Qwen3-30B-A3B-int5wc-hadamard-humming.yaml => Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml} (57%) create mode 100644 tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml create mode 100644 tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml create mode 100644 tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-deepep.txt delete mode 100644 tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-l4.txt create mode 100644 vllm/model_executor/kernels/linear/mxfp6/humming.py create mode 100644 vllm/model_executor/layers/quantization/utils/humming/__init__.py create mode 100644 vllm/model_executor/layers/quantization/utils/humming/activation.py create mode 100644 vllm/model_executor/layers/quantization/utils/humming/linear.py rename vllm/model_executor/layers/quantization/utils/{humming_utils.py => humming/moe.py} (51%) create mode 100644 vllm/model_executor/layers/quantization/utils/humming/schema.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 3eb7f0a44430..56aeca2a8eb3 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -347,47 +347,67 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N" - key: lm-eval-humming-f16-a100 +- label: ":nvidia: (L4) LM Eval Humming TEMPORARY" + key: lm-eval-humming-l4 + optional: true + timeout_in_minutes: 90 + device: l4 + num_devices: 1 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/configs/humming/ + - requirements/cuda.txt + - vllm/utils/humming.py + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming/ + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt + +- label: ":nvidia: (A100) LM Eval Humming Default TEMPORARY Shard %N" + key: lm-eval-humming-a100-default + optional: true timeout_in_minutes: 75 device: a100 - optional: true num_devices: 1 parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY" - key: lm-eval-humming-act-a100 +- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY Shard %N" + key: lm-eval-humming-a100-act + optional: true timeout_in_minutes: 45 device: a100 - optional: true num_devices: 1 + parallelism: 2 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N" - key: lm-eval-humming-f16-h100 +- label: ":nvidia: (H100) LM Eval Humming Default TEMPORARY Shard %N" + key: lm-eval-humming-h100-default + optional: true timeout_in_minutes: 70 device: h100 - optional: true num_devices: 1 parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -395,14 +415,14 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY" - key: lm-eval-humming-act-h100 + key: lm-eval-humming-h100-act + optional: true timeout_in_minutes: 70 device: h100 - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -410,15 +430,35 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt # - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY" - key: lm-eval-humming-f16-b200 +- label: ":nvidia: (H100) LM Eval Humming DeepEP TEMPORARY" + key: lm-eval-humming-h100-deepep + optional: true + timeout_in_minutes: 60 + device: h100 + num_devices: 2 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/configs/humming/ + - requirements/cuda.txt + - vllm/utils/humming.py + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/kernels/linear/ + - vllm/distributed/ + - csrc/moe/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-deepep.txt + +- label: ":nvidia: (B200) LM Eval Humming Default TEMPORARY" + key: lm-eval-humming-b200-default + optional: true timeout_in_minutes: 50 device: b200-k8s - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -426,14 +466,14 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt - label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY" - key: lm-eval-humming-act-b200 + key: lm-eval-humming-b200-act + optional: true timeout_in_minutes: 50 device: b200-k8s - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index d66b78fe838a..5c9a522feb7b 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -236,7 +236,7 @@ __device__ void _moe_align_block_size_small_batch_expert( int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num, - int32_t* token_mask, bool has_expert_map) { + int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) { // Compute input buffer offsets. Typically these will all be 0, except when // using Multi LoRA. int sorted_token_ids_offset = max_num_tokens_padded * model_offset; @@ -318,9 +318,14 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, - num_experts, has_expert_map); - expert_id != -1) { + int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, + has_expert_map); + bool valid = + expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); + if (scatter_idx != nullptr) { + scatter_idx[i] = valid ? static_cast(i) : -1; + } + if (expert_id != -1) { int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; @@ -338,14 +343,20 @@ __device__ void _count_and_sort_expert_tokens( int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask, - int32_t model_offset, int32_t topk_num, bool has_expert_map) { + int32_t model_offset, int32_t topk_num, bool has_expert_map, + int32_t* scatter_idx = nullptr) { const size_t tid = blockIdx.y * blockDim.x + threadIdx.x; const size_t stride = blockDim.x * gridDim.y; for (size_t i = tid; i < numel; i += stride) { - if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, - num_experts, has_expert_map); - expert_id != -1) { + int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, + has_expert_map); + bool valid = + expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); + if (scatter_idx != nullptr) { + scatter_idx[i] = valid ? static_cast(i) : -1; + } + if (expert_id != -1) { if (token_mask == nullptr || token_mask[i / topk_num]) { int32_t rank_post_pad = atomicAdd( &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); @@ -377,10 +388,11 @@ __global__ void count_and_sort_expert_tokens_kernel( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, - int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) { + int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map, + int32_t* scatter_idx) { _count_and_sort_expert_tokens( topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, - max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map); + max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx); } // Reduce the topk expert outputs per token (summed in fp32). The output is @@ -533,12 +545,12 @@ __global__ void moe_align_block_size_small_batch_expert_kernel( int32_t* __restrict__ total_tokens_post_pad, int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, size_t numel, int32_t max_num_tokens_padded, int32_t topk_num, - bool has_expert_map) { + bool has_expert_map, int32_t* scatter_idx) { _moe_align_block_size_small_batch_expert( topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, num_experts, block_size, numel, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr, - has_expert_map); + has_expert_map, scatter_idx); } template @@ -658,16 +670,31 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel( // taken from // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc -void moe_align_block_size( - torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, - torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, - torch::stable::Tensor num_tokens_post_pad, - std::optional maybe_expert_map) { +void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, + int64_t block_size, + torch::stable::Tensor sorted_token_ids, + torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map, + std::optional scatter_idx) { const torch::stable::accelerator::DeviceGuard device_guard( topk_ids.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(topk_ids.get_device_index()); + int32_t* scatter_ptr = nullptr; + if (scatter_idx.has_value()) { + STD_TORCH_CHECK( + scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int, + "scatter_idx must be int32"); + STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() && + scatter_idx->is_contiguous() && + scatter_idx->numel() == topk_ids.numel(), + "scatter_idx must be contiguous, on the same device, and " + "match topk_ids.numel()"); + scatter_ptr = reinterpret_cast(scatter_idx->mutable_data_ptr()); + } + int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; int experts_per_warp = WARP_SIZE; @@ -713,7 +740,8 @@ void moe_align_block_size( num_tokens_post_pad.mutable_data_ptr()), reinterpret_cast(expert_map.mutable_data_ptr()), num_experts, block_size, topk_ids.numel(), - sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map, + scatter_ptr); } else { torch::stable::Tensor cumsum_buffer = torch::stable::new_empty( topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); @@ -753,7 +781,7 @@ void moe_align_block_size( reinterpret_cast(cumsum_buffer.mutable_data_ptr()), reinterpret_cast(expert_map.mutable_data_ptr()), topk_ids.numel(), num_experts, sorted_token_ids.size(0), - topk_ids.size(1), has_expert_map); + topk_ids.size(1), has_expert_map, scatter_ptr); } }); } diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 03057c3b1421..549a19e98148 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, std::optional expert_map); -void moe_align_block_size( - torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, - torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, - torch::stable::Tensor num_tokens_post_pad, - std::optional maybe_expert_map); +void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, + int64_t block_size, + torch::stable::Tensor sorted_token_ids, + torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map, + std::optional scatter_idx); void batched_moe_align_block_size( int64_t max_tokens_per_batch, int64_t block_size, diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu index 251097456f9e..957cd60e49de 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu @@ -56,20 +56,18 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert)); } -void moe_permute_impl( - const torch::stable::Tensor& input, // [n_token, hidden] - const torch::stable::Tensor& topk_ids, // [n_token, topk] - const torch::stable::Tensor& token_expert_indices, // [n_token, topk] +torch::stable::Tensor moe_sort_routing( + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] const std::optional& expert_map, // [n_expert] int64_t n_expert, int64_t n_local_expert, int64_t topk, - torch::stable::Tensor& permuted_input, // [permuted_size, hidden] torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] - torch::stable::Tensor& permuted_idx, // [permute_size] const std::optional& maybe_sort_workspace, const std::optional& maybe_permuted_experts_id, const std::optional& maybe_sorted_row_idx, - const std::optional& maybe_topk_ids_for_sort) { + const std::optional& maybe_topk_ids_for_sort, + bool build_inverse) { STD_TORCH_CHECK(expert_first_token_offset.scalar_type() == torch::headeronly::ScalarType::Long, "expert_first_token_offset must be int64"); @@ -88,12 +86,11 @@ void moe_permute_impl( "token_expert_indices shape must be same as inv_permuted_idx"); const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - auto device = input.device(); - auto n_token = input.sizes()[0]; - auto n_hidden = input.sizes()[1]; + topk_ids.get_device_index()); + auto device = topk_ids.device(); + auto n_token = topk_ids.sizes()[0]; auto expanded_rows = n_token * topk; - auto stream = get_current_cuda_stream(input.get_device_index()); + auto stream = get_current_cuda_stream(topk_ids.get_device_index()); auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert); auto sort_workspace = maybe_allocate_tensor( @@ -107,13 +104,11 @@ void moe_permute_impl( torch::headeronly::ScalarType::Int, device, "sorted_row_idx"); CubKeyValueSorter sorter{}; - int64_t* valid_num_ptr = nullptr; torch::stable::Tensor topk_ids_for_sort = topk_ids; - if (expert_map.has_value()) { - const int* expert_map_ptr = get_ptr(expert_map.value()); - valid_num_ptr = - get_ptr(expert_first_token_offset) + n_local_expert; + if (expert_map.has_value() || build_inverse) { + const int* expert_map_ptr = + expert_map.has_value() ? get_ptr(expert_map.value()) : nullptr; topk_ids_for_sort = maybe_allocate_tensor( maybe_topk_ids_for_sort, topk_ids.sizes(), torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort"); @@ -126,7 +121,39 @@ void moe_permute_impl( get_ptr(topk_ids_for_sort), get_ptr(token_expert_indices), get_ptr(permuted_experts_id), get_ptr(sorted_row_idx), get_ptr(expert_first_token_offset), n_token, n_expert, - n_local_expert, topk, sorter, get_ptr(sort_workspace), stream); + n_local_expert, topk, sorter, get_ptr(sort_workspace), stream, + build_inverse ? get_ptr(inv_permuted_idx) : nullptr); + return sorted_row_idx; +} + +void moe_permute_impl( + const torch::stable::Tensor& input, // [n_token, hidden] + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] + const std::optional& expert_map, // [n_expert] + int64_t n_expert, int64_t n_local_expert, int64_t topk, + torch::stable::Tensor& permuted_input, // [permuted_size, hidden] + torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] + torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] + torch::stable::Tensor& permuted_idx, // [permute_size] + const std::optional& maybe_sort_workspace, + const std::optional& maybe_permuted_experts_id, + const std::optional& maybe_sorted_row_idx, + const std::optional& maybe_topk_ids_for_sort) { + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + auto sorted_row_idx = moe_sort_routing( + topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert, + topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace, + maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort, + false); + auto stream = get_current_cuda_stream(input.get_device_index()); + auto n_token = input.size(0); + auto n_hidden = input.size(1); + auto valid_num_ptr = + expert_map.has_value() + ? get_ptr(expert_first_token_offset) + n_local_expert + : nullptr; MOE_DISPATCH(input.scalar_type(), [&] { expandInputRowsKernelLauncher( @@ -171,6 +198,22 @@ void moe_permute_with_scratch( topk_ids_for_sort); } +void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, + int64_t n_expert, int64_t n_local_expert, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& scatter_idx, + torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& sorted_experts, + torch::stable::Tensor& sorted_rows, + torch::stable::Tensor& topk_ids_for_sort) { + moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert, + n_local_expert, topk_ids.size(1), expert_first_token_offset, + scatter_idx, sort_workspace, sorted_experts, sorted_rows, + topk_ids_for_sort, true); +} + void moe_unpermute( const torch::stable::Tensor& permuted_hidden_states, // [n_token * topk, hidden] @@ -315,6 +358,19 @@ void moe_unpermute( STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0"); } +void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, + int64_t n_expert, int64_t n_local_expert, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& scatter_idx, + torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& sorted_experts, + torch::stable::Tensor& sorted_rows, + torch::stable::Tensor& topk_ids_for_sort) { + STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0"); +} + #endif bool moe_permute_unpermute_supported() { @@ -328,5 +384,6 @@ bool moe_permute_unpermute_supported() { STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { m.impl("moe_permute", TORCH_BOX(&moe_permute)); m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch)); + m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter)); m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute)); } diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index ec7851da3d3d..8d8bdf8baaf6 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -82,34 +82,40 @@ __device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices, return target_location + 1; } -// Calculates the start offset of the tokens for a given expert. The last -// element is the total number of valid tokens -__global__ void computeExpertFirstTokenOffsetKernel( +// Computes expert offsets (ending with the valid row count) and optional +// inverse indices. +__global__ void computeExpertOffsetsAndInverseKernel( int const* sorted_experts, int64_t const sorted_experts_len, - int const num_experts, int64_t* expert_first_token_offset) { - // First, compute the global tid. We only need 1 thread per expert. - int const expert = blockIdx.x * blockDim.x + threadIdx.x; + int const num_experts, int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse) { + int const index = blockIdx.x * blockDim.x + threadIdx.x; // Note that expert goes [0, num_experts] (inclusive) because we want a count // for the total number of active tokens at the end of the scan. - if (expert >= num_experts + 1) { - return; + if (index < num_experts + 1) { + expert_first_token_offset[index] = + findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, index); + } + if (inverse != nullptr && index < sorted_experts_len) { + inverse[sorted_rows[index]] = index; } - expert_first_token_offset[expert] = - findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert); } -void computeExpertFirstTokenOffset(int const* sorted_indices, - int const total_indices, - int const num_experts, - int64_t* expert_first_token_offset, - cudaStream_t stream) { - int const num_entries = num_experts + 1; - int const threads = std::min(1024, num_entries); +void computeExpertOffsetsAndInverse(int const* sorted_indices, + int const total_indices, + int const num_experts, + int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse, + cudaStream_t stream) { + int const num_entries = inverse != nullptr + ? std::max(num_experts + 1, total_indices) + : num_experts + 1; + int const threads = inverse != nullptr ? 256 : std::min(1024, num_entries); int const blocks = (num_entries + threads - 1) / threads; - computeExpertFirstTokenOffsetKernel<<>>( - sorted_indices, total_indices, num_experts, expert_first_token_offset); + computeExpertOffsetsAndInverseKernel<<>>( + sorted_indices, total_indices, num_experts, expert_first_token_offset, + sorted_rows, inverse); } void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, @@ -117,7 +123,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, int64_t* expert_first_token_offset, int num_rows, int num_experts, int num_experts_per_node, int k, CubKeyValueSorter& sorter, void* sorter_ws, - cudaStream_t stream) { + cudaStream_t stream, int* inverse) { int64_t const expanded_num_rows = static_cast(k) * num_rows; // We need to use the full num_experts because that is the sentinel value used // by topk for disabled experts @@ -127,9 +133,9 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row, permuted_experts, source_rows, permuted_rows, expanded_num_rows, stream); - computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows, - num_experts_per_node, expert_first_token_offset, - stream); + computeExpertOffsetsAndInverse( + permuted_experts, expanded_num_rows, num_experts_per_node, + expert_first_token_offset, permuted_rows, inverse, stream); } __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size, @@ -142,7 +148,7 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size, extern __shared__ int smem_expert_map[]; // store expert_map in smem for (int i = tidx; i < num_experts; i += blockDim.x) { - smem_expert_map[i] = expert_map_ptr[i]; + smem_expert_map[i] = expert_map_ptr ? expert_map_ptr[i] : i; } __syncthreads(); diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index 89c278a4ed4b..53afad2f5d41 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -43,18 +43,19 @@ class CubKeyValueSorter { int num_bits_; }; -void computeExpertFirstTokenOffset(int const* sorted_indices, - int const total_indices, - int const num_experts, - int64_t* expert_first_token_offset, - cudaStream_t stream); +void computeExpertOffsetsAndInverse(int const* sorted_indices, + int const total_indices, + int const num_experts, + int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse, + cudaStream_t stream); void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, int* permuted_experts, int* permuted_rows, int64_t* expert_first_token_offset, int num_rows, int num_experts, int num_experts_per_node, int k, CubKeyValueSorter& sorter, void* sorter_ws, - cudaStream_t stream); + cudaStream_t stream, int* inverse = nullptr); template void expandInputRowsKernelLauncher( diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index c0e966b2be9c..2feec3838993 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -39,7 +39,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { " int block_size, Tensor! sorted_token_ids," " Tensor! experts_ids," " Tensor! num_tokens_post_pad," - " Tensor? maybe_expert_map) -> ()"); + " Tensor? maybe_expert_map," + " Tensor(a!)? scatter_idx=None) -> ()"); // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size, but for the batched case. @@ -105,6 +106,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, " "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()"); + m.def( + "moe_prepare_scatter(Tensor topk_ids, Tensor token_expert_indices, " + "Tensor? expert_map, int n_expert, int n_local_expert, " + "Tensor! expert_first_token_offset, Tensor! scatter_idx, " + "Tensor! sort_workspace, Tensor! sorted_experts, Tensor! sorted_rows, " + "Tensor! topk_ids_for_sort)->()"); + m.def( "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights," "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, " diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 00a56ac3bd6e..72609d8983b6 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -33,4 +33,4 @@ quack-kernels==0.6.5 # Required by MSA tokenspeed-mla==0.1.8; platform_system == "Linux" # Humming kernels for quantization gemm -humming-kernels[cu13]==0.1.12 +humming-kernels[cu13]==0.1.15 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml index 310255f6bf41..2ee157359a34 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml @@ -9,3 +9,5 @@ server_args: >- --moe-backend humming --linear-backend humming --dtype bfloat16 +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml new file mode 100644 index 000000000000..a6bf4f652528 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml @@ -0,0 +1,14 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --data-parallel-size 2 + --enable-expert-parallel + --all2all-backend deepep_high_throughput +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml new file mode 100644 index 000000000000..cf8e01d70f97 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml @@ -0,0 +1,14 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --data-parallel-size 2 + --enable-expert-parallel + --all2all-backend deepep_high_throughput +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "indexed" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml index 67725fce64ae..a8b289aaacd4 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml @@ -7,3 +7,5 @@ server_args: >- --max-model-len 8192 --quantization humming --linear-backend humming +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml index a20c14331036..57923c63b0e4 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml @@ -7,3 +7,5 @@ server_args: >- --max-model-len 8192 --moe-backend humming --linear-backend humming +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml deleted file mode 100644 index b58c0d710e4b..000000000000 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml +++ /dev/null @@ -1,13 +0,0 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 -num_questions: 1319 -num_fewshot: 5 -server_args: >- - --enforce-eager - --max-model-len 8192 - --tensor-parallel-size 1 - --quantization humming - --kernel-config.enable_flashinfer_autotune=False -env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' - VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml similarity index 62% rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml index c932091dbe35..d2f04348b9b3 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml @@ -1,7 +1,8 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 +gen_prefix: " \n\n\n" server_args: >- --enforce-eager --max-model-len 8192 @@ -9,5 +10,5 @@ server_args: >- --quantization humming --kernel-config.enable_flashinfer_autotune=False env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml similarity index 57% rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml index fa5095cc8824..08687aa6a307 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml @@ -1,7 +1,8 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 +gen_prefix: " \n\n\n" server_args: >- --enforce-eager --max-model-len 8192 @@ -9,4 +10,4 @@ server_args: >- --quantization humming --kernel-config.enable_flashinfer_autotune=False env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml new file mode 100644 index 000000000000..067b96924d25 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3.5-9B" +accuracy_threshold: 0.88 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml new file mode 100644 index 000000000000..24915abbc2b9 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-9B" +accuracy_threshold: 0.88 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt new file mode 100644 index 000000000000..0537ed974a3f --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt @@ -0,0 +1,3 @@ +gpt-oss-20b-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt new file mode 100644 index 000000000000..512cbe99ff34 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt @@ -0,0 +1,4 @@ +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-4B-mixed-quant-RTN-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt index 63b428ac0e30..bdfa9033df0c 100644 --- a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt @@ -1,4 +1,5 @@ Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml gpt-oss-20b-humming.yaml Qwen3-30B-A3B-GPTQ-Int4-humming.yaml Qwen3-30B-A3B-FP8-block-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-act-int8.txt b/tests/evals/gsm8k/configs/humming/config-act-int8.txt index a7efbac3d033..68f5cd1c15dc 100644 --- a/tests/evals/gsm8k/configs/humming/config-act-int8.txt +++ b/tests/evals/gsm8k/configs/humming/config-act-int8.txt @@ -3,4 +3,5 @@ Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml Qwen3-30B-A3B-AWQ-humming.yaml Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml Qwen3-4B-mixed-quant-RTN-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-deepep.txt b/tests/evals/gsm8k/configs/humming/config-deepep.txt new file mode 100644 index 000000000000..2c97531d75f1 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-deepep.txt @@ -0,0 +1,2 @@ +Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml +Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt index 60179f1159fe..737a92cf48ae 100644 --- a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt +++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt @@ -1,4 +1,5 @@ Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml Qwen3-30B-A3B-Fp8-v1-humming.yaml Qwen3-30B-A3B-NVFP4-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt deleted file mode 100644 index 2c10777a095a..000000000000 --- a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt +++ /dev/null @@ -1,3 +0,0 @@ -Qwen3-30B-A3B-int5wc-hadamard-humming.yaml -Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml -Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-l4.txt b/tests/evals/gsm8k/configs/humming/config-l4.txt new file mode 100644 index 000000000000..d1434b1cab2e --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-l4.txt @@ -0,0 +1,9 @@ +Qwen2-1.5B-Instruct-FP8W8-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml +Qwen3-0.6B-MXFP8-humming.yaml +Qwen3-0.6B-MXFP8-humming-act-fp8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml +gpt-oss-20b-humming.yaml +Qwen3.5-9B-int6wc-hadamard-humming.yaml +Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index b1a52a5d557d..c1ba96a79213 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -352,49 +352,46 @@ def test_masked_situ_and_mul( ) -@pytest.mark.parametrize( - ("activation", "activation_config"), - [ - (MoEActivation.SILU, ApplyMoEActivationConfig()), - ( - MoEActivation.SILU, - ApplyMoEActivationConfig(clamp_limit=3.0), +MOE_ACTIVATION_CASES = [ + pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"), + pytest.param( + MoEActivation.SILU, ApplyMoEActivationConfig(clamp_limit=3.0), id="silu_clamp" + ), + pytest.param(MoEActivation.GELU, ApplyMoEActivationConfig(), id="gelu"), + pytest.param(MoEActivation.GELU_TANH, ApplyMoEActivationConfig(), id="gelu_tanh"), + pytest.param( + MoEActivation.SITU, + ApplyMoEActivationConfig( + activation_situ_beta=1.5, + activation_situ_linear_beta=2.0, ), - (MoEActivation.GELU, ApplyMoEActivationConfig()), - (MoEActivation.GELU_TANH, ApplyMoEActivationConfig()), - ( - MoEActivation.SITU, - ApplyMoEActivationConfig( - activation_situ_beta=1.5, - activation_situ_linear_beta=2.0, - ), - ), - (MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig()), - ( - MoEActivation.SWIGLUOAI_UNINTERLEAVE, - ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5), - ), - (MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig()), - (MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.GELU_TANH_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig()), - ], - ids=[ - "silu", - "silu_clamp", - "gelu", - "gelu_tanh", - "situ", - "swigluoai", - "swigluoai_uninterleave", - "swiglustep", - "silu_no_mul", - "gelu_no_mul", - "gelu_tanh_no_mul", - "relu2_no_mul", - ], -) + id="situ", + ), + pytest.param(MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig(), id="swigluoai"), + pytest.param( + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5), + id="swigluoai_uninterleave", + ), + pytest.param(MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig(), id="swiglustep"), + pytest.param( + MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig(), id="silu_no_mul" + ), + pytest.param( + MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig(), id="gelu_no_mul" + ), + pytest.param( + MoEActivation.GELU_TANH_NO_MUL, + ApplyMoEActivationConfig(), + id="gelu_tanh_no_mul", + ), + pytest.param( + MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig(), id="relu2_no_mul" + ), +] + + +@pytest.mark.parametrize(("activation", "activation_config"), MOE_ACTIVATION_CASES) @torch.inference_mode() def test_masked_moe_activation_dispatch( default_vllm_config, @@ -497,3 +494,73 @@ def test_activation( out = torch.empty_like(x) opcheck(fn, (out, x)) + + +HUMMING_ACTIVATION_CASES = MOE_ACTIVATION_CASES + [ + pytest.param(MoEActivation.RELU2, ApplyMoEActivationConfig(), id="relu2"), +] +HUMMING_ACTIVATION_CASES += [ + pytest.param( + MoEActivation.SITU, + ApplyMoEActivationConfig( + activation_situ_beta=1.5, + activation_situ_linear_beta=linear_beta, + ), + id=f"situ-linear-beta-{linear_beta}", + ) + for linear_beta in (None, 0.0, -1.0) +] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Humming requires CUDA") +@pytest.mark.parametrize(("activation", "activation_config"), HUMMING_ACTIVATION_CASES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize(("num_tokens", "d"), [(1, 512), (7, 768), (83, 512)]) +@torch.inference_mode() +def test_humming_activation_matches_framework( + activation: MoEActivation, + activation_config: ApplyMoEActivationConfig, + dtype: torch.dtype, + num_tokens: int, + d: int, +) -> None: + """Compare activation math/layouts without quantization or Hadamard error.""" + pytest.importorskip("humming") + from humming.ops import process_input + + from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, + ) + + set_random_seed(0) + width = d * 2 if activation.is_gated else d + x = 4 * torch.randn(num_tokens, width, dtype=dtype, device=CUDA_DEVICES[0]) + # Exercise zero, saturation, and both sides of the clamp limits (3 and 7). + edges = x.new_tensor([0, 0.001, 1, 2.99, 3, 3.01, 6.99, 7, 7.01, 8, 16]) + edges = torch.cat((-edges[1:].flip(0), edges)) + x[0] = edges[torch.arange(width, device=x.device) % edges.numel()] + + actual, _, _ = process_input( + x, + quant_mode="none", + hadamard_block_size=0, + **get_humming_activation(activation, activation_config), + ) + expected = torch.empty(num_tokens, d, dtype=dtype, device=x.device) + if activation == MoEActivation.RELU2: + # apply_moe_activation has only the non-gated ReLU2 variant. + gate, up = x.float().chunk(2, dim=-1) + activated_gate = torch.empty_like(gate) + apply_moe_activation(MoEActivation.RELU2_NO_MUL, activated_gate, gate.clone()) + expected.copy_(activated_gate * up) + else: + # The framework's non-gated ReLU2 path modifies its input in place. + apply_moe_activation( + activation, expected, x.clone(), activation_config=activation_config + ) + torch.testing.assert_close( + actual, + expected, + atol=get_default_atol(expected), + rtol=get_default_rtol(expected), + ) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 1aca8593ed7c..8d6deac4f025 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1245,7 +1245,7 @@ def _make_humming_indexed_experts(activation: MoEActivation): from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( HummingIndexedExperts, ) - from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.model_executor.layers.quantization.utils import humming as humming_utils from vllm.utils import humming top_k, num_experts = 6, 12 @@ -1306,6 +1306,56 @@ def _make_humming_indexed_experts(activation: MoEActivation): return experts, layer +@pytest.mark.parametrize("valid_shape_m", [8192, 8193]) +@pytest.mark.parametrize("w2_block_size", [64, 96, 128]) +def test_humming_indexed_routing_matches_each_projection_tile( + valid_shape_m: int, w2_block_size: int +): + """Independently tuned projections preserve routed rows and expert identity.""" + from types import SimpleNamespace + + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingIndexedExperts, + ) + + topk_ids = torch.tensor([[0, 1], [1, 2], [2, 0]], device="cuda") + expert_map = torch.tensor([1, -1, 0], dtype=torch.int32, device="cuda") + experts = SimpleNamespace( + global_num_experts=3, + estimate_local_valid_shape_m=lambda _: valid_shape_m, + w13_tuning_config=[ + [0, 8192, {"block_shape": [64, 128, 64]}], + [8192, 16384, {"block_shape": [128, 128, 64]}], + ], + w2_tuning_config=[ + [0, 8192, {"block_shape": [64, 128, 64]}], + [8192, 16384, {"block_shape": [w2_block_size, 128, 64]}], + ], + compute_config_str="compute", + w13_tuning_config_str="w13", + w2_tuning_config_str="w2", + ) + first, second, _ = HummingIndexedExperts.prepare_humming_moe_kwargs( + experts, topk_ids, expert_map, None + ) + block_sizes = (64, 64) if valid_shape_m == 8192 else (128, w2_block_size) + expected_rows = torch.where(expert_map[topk_ids.flatten()] >= 0)[0] + for kwargs, block_size in zip((first, second), block_sizes): + padded = kwargs["num_tokens_padded"].item() + assert padded % block_size == 0 + rows = kwargs["sorted_ids"][:padded].long() + valid = rows < topk_ids.numel() + torch.testing.assert_close(rows[valid].sort().values, expected_rows) + routed_experts = kwargs["expert_ids"][: padded // block_size] + torch.testing.assert_close( + routed_experts.repeat_interleave(block_size)[valid], + expert_map[topk_ids.flatten()[rows[valid]]], + ) + assert (first["top_k"], second["top_k"]) == (2, 1) + for key in ("sorted_ids", "expert_ids", "num_tokens_padded"): + assert (first[key] is second[key]) == (block_sizes[0] == block_sizes[1]) + + @pytest.mark.parametrize("activation", list(MoEActivation)) def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation): from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( @@ -1464,54 +1514,46 @@ def test_humming_delegates_to_instance_activation(): def test_humming_grouped_apply_forwards_valid_prefix( monkeypatch: pytest.MonkeyPatch, ): - from types import MethodType, SimpleNamespace + from types import SimpleNamespace from unittest.mock import Mock import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( - HummingExpertsBase, HummingGroupedExperts, ) - activation_func = Mock() buffers = { "gate_up_output": torch.empty(6, 4), - "activation_output": torch.empty(6, 2), + "quanted_down_input": torch.empty(6, 2), "down_output": torch.empty(6, 4), } experts = SimpleNamespace( - activation=activation_func, num_experts=2, estimate_local_valid_shape_m=lambda _: 6, prepare_buffers=lambda *_: buffers, - _get_permute_scratch=lambda _: None, - quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: ( - inputs, - input_scale, + _get_permute_scratch=lambda _, *, indices_only=False: object(), + process_input=Mock( + side_effect=lambda _, **kwargs: ( + kwargs["inputs"], + kwargs.get("input_scale"), + None, + ) ), humming_forward=Mock(), compute_config_str="", w13_tuning_config_str="", w2_tuning_config_str="", ) - experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts) hidden_states = torch.empty(3, 4) output = torch.empty_like(hidden_states) topk_ids = torch.zeros((3, 2), dtype=torch.int64) topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) - expected_counts = torch.tensor([4], dtype=torch.int32) expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64) monkeypatch.setattr( humming, - "moe_permute", - lambda **_: ( - torch.empty(6, 4), - None, - expert_offsets, - torch.arange(6), - None, - ), + "moe_prepare_scatter", + lambda *_: (expert_offsets, torch.arange(6).view(3, 2)), ) monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None) @@ -1534,12 +1576,11 @@ def test_humming_grouped_apply_forwards_valid_prefix( apply_router_weight_on_input=False, ) - activation_func.assert_called_once() - call_kwargs = activation_func.call_args.kwargs + call_kwargs = experts.process_input.call_args.kwargs assert call_kwargs["activation"] == MoEActivation.SITU - assert call_kwargs["input"].shape == (6, 4) - assert call_kwargs["output"].shape == (6, 2) - torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts) + assert call_kwargs["inputs"].shape == (6, 4) + assert call_kwargs["quanted_input"].shape == (6, 2) + torch.testing.assert_close(call_kwargs["num_valid_tokens"], expert_offsets[-1:]) def test_batched_marlin_activation_uses_expert_token_counts( @@ -1653,7 +1694,8 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): activation=moe_config.activation, ) assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size - assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size + assert "activation_output" not in buffer_metas + assert buffer_metas["quanted_down_input"]["shape"][-1] == intermediate_size assert experts.moe_problem_size( a1=torch.empty(1, hidden_size), w1=torch.empty(num_experts, 1), @@ -1661,6 +1703,13 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): topk_ids=torch.empty(1, top_k, dtype=torch.long), ) == (num_experts, 1, intermediate_size, hidden_size, top_k) + # Quantized batched input storage must cover every dispatcher's token slots. + experts.is_batched = lambda: True + experts.max_num_tokens = 17 + experts.num_dispatchers = 2 + batched_metas, _ = experts.get_buffer_metas(17, top_k, activation) + assert batched_metas["quanted_gate_up_input"]["shape"][0] == num_experts * 17 * 2 + def test_humming_indexed_writes_supplied_output_buffer(): from vllm.forward_context import set_forward_context diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index a1599f5ad3da..ca339eeed8bd 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -22,6 +22,46 @@ set_random_seed(0) +@pytest.mark.parametrize("num_experts,num_tokens", [(32, 8), (256, 8), (256, 1024)]) +@pytest.mark.parametrize("use_expert_map", [False, True]) +def test_moe_align_scatter_masks_invalid_routes( + num_experts, num_tokens, use_expert_map +): + topk_ids = torch.arange(num_tokens * 4, device="cuda").reshape(-1, 4) + topk_ids %= num_experts + topk_ids[0] = torch.tensor([-1, num_experts, 0, 1], device="cuda") + expert_map = None + valid = (topk_ids >= 0) & (topk_ids < num_experts) + if use_expert_map: + expert_map = torch.full((num_experts,), -1, dtype=torch.int32, device="cuda") + expert_map[::2] = torch.arange( + num_experts // 2, device="cuda", dtype=torch.int32 + ) + valid &= topk_ids % 2 == 0 + + sorted_ids, expert_ids, num_padded, scatter_idx = moe_align_block_size( + topk_ids, + 32, + num_experts, + expert_map, + ignore_invalid_experts=True, + return_scatter_idx=True, + ) + expected = torch.arange(topk_ids.numel(), device="cuda", dtype=torch.int32) + expected[~valid.flatten()] = -1 + torch.testing.assert_close(scatter_idx, expected[:, None]) + active_sorted = sorted_ids[: num_padded.item()] + active_sorted = active_sorted[active_sorted < topk_ids.numel()] + torch.testing.assert_close(active_sorted.sort().values, expected[expected >= 0]) + + baseline = moe_align_block_size( + topk_ids, 32, num_experts, expert_map, ignore_invalid_experts=True + ) + assert len(baseline) == 3 + torch.testing.assert_close(num_padded, baseline[2]) + torch.testing.assert_close(expert_ids, baseline[1]) + + def _group_tokens_by_expert( sorted_ids: torch.Tensor, expert_ids: torch.Tensor, diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 52fc8963ece8..76f4b57cf07f 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -18,6 +18,7 @@ get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, + moe_prepare_scatter, moe_unpermute, ) from vllm.platforms import current_platform @@ -458,7 +459,7 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: hidden_dtype=hidden_states.dtype, ) - permuted, _, expert_offsets, _, _ = moe_permute( + permuted, _, expert_offsets, inverse, _ = moe_permute( hidden_states=hidden_states, a1q_scale=None, topk_ids=topk_ids, @@ -473,3 +474,27 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), ) torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]]) + assert torch.all(inverse[[1, 3, 4]] >= expert_offsets[-1]) + expected = torch.zeros_like(hidden_states) + expected[[0, 2]] = hidden_states[[0, 2]] + output = torch.empty_like(hidden_states) + moe_unpermute( + output, permuted, torch.ones(5, 1, device="cuda"), inverse, expert_offsets + ) + torch.testing.assert_close(output, expected) + + expected_inverse = inverse.clone() + expected_offsets = expert_offsets.clone() + expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) + torch.testing.assert_close(indices.flatten(), expected_inverse) + torch.testing.assert_close(expert_offsets, expected_offsets) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) + topk_ids.zero_() + graph.replay() + + expected_inverse = torch.arange(5, dtype=torch.int32, device="cuda") + expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") + torch.testing.assert_close(indices.flatten(), expected_inverse) + torch.testing.assert_close(expert_offsets, expected_offsets) diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 4adfffa7c87b..46bef106e644 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -14,7 +14,6 @@ AiterMxfp4LinearKernel, EmulationMxfp4LinearKernel, FlashInferMxFp4LinearKernel, - HummingMxFp4LinearKernel, MarlinMxFp4LinearKernel, MxFp4LinearKernel, MxFp4LinearLayerConfig, @@ -44,7 +43,7 @@ # Weight-only (A16) kernels: they never quantize activations. They still accept # MXFP4 activation keys as an intentional compatibility fallback. -_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel] +_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel] def test_can_implement_is_abstract(): diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py index dbec989df3f9..74c266e0bfe5 100644 --- a/tests/quantization/test_humming_mxfp4_block_fp8.py +++ b/tests/quantization/test_humming_mxfp4_block_fp8.py @@ -15,6 +15,7 @@ """ import pytest +import torch from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( HummingExpertsBase, @@ -56,8 +57,8 @@ def test_mxfp4_weight_with_per_token_fp8_activation_still_supported(): ], ) def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key): - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - _humming_input_schema_to_quant_key, + from vllm.model_executor.layers.quantization.utils.humming import ( + input_schema_to_quant_key, ) from vllm.utils.humming import HummingInputSchema from vllm.utils.humming import dtypes as humming_dtypes @@ -66,18 +67,21 @@ def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key): a_dtype=humming_dtypes.float8e4m3, input_scale_group_size=group_size, ) - assert _humming_input_schema_to_quant_key(schema) == expected_key + assert input_schema_to_quant_key(schema, param_dtype=torch.bfloat16) == expected_key @pytest.mark.skipif(not has_humming(), reason="humming is not installed") def test_humming_bf16_input_schema_is_unquantized(): - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - _humming_input_schema_to_quant_key, + from vllm.model_executor.layers.quantization.utils.humming import ( + input_schema_to_quant_key, ) from vllm.utils.humming import HummingInputSchema # No a_dtype -> unquantized (bf16/fp16) inputs -> None. - assert _humming_input_schema_to_quant_key(HummingInputSchema()) is None + assert ( + input_schema_to_quant_key(HummingInputSchema(), param_dtype=torch.bfloat16) + is None + ) def test_block_fp8_activation_quant_config_is_block_quantized(): @@ -86,7 +90,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): DeepEP prepare/finalize step quantize activations to block FP8 *before* the all-to-all dispatch instead of deferring to Humming -- and what makes HummingExpertsBase.expects_unquantized_inputs return False.""" - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( make_humming_moe_quant_config, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -108,7 +112,7 @@ def test_default_activation_quant_config_defers_to_humming(): """Without a block activation shape the config stays per-token (the deferred path): Humming quantizes internally, preserving pre-existing behavior for every non-block-FP8 scheme.""" - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( make_humming_moe_quant_config, ) from vllm.platforms import current_platform diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 0fd5a10120b3..989bf967fafd 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -229,7 +229,7 @@ def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): def test_moe_wna16_uses_humming_quant_config(monkeypatch): - from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.model_executor.layers.quantization.utils import humming as humming_utils method = object.__new__(MoeWNA16Method) method.wna16_backend = WNA16MoEBackend.HUMMING diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2a1654231fc2..1d38bcade013 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2272,6 +2272,7 @@ def moe_align_block_size( experts_ids: torch.Tensor, num_tokens_post_pad: torch.Tensor, expert_map: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.moe_align_block_size( topk_ids, @@ -2281,6 +2282,7 @@ def moe_align_block_size( experts_ids, num_tokens_post_pad, expert_map, + scatter_idx, ) diff --git a/vllm/envs.py b/vllm/envs.py index 26c00c6da623..259c08e7e7b3 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1524,7 +1524,8 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None) ), - # The activation dtype config for humming kernel + # The activation dtype config for humming kernel. Explicit schemas disable + # fallback unless the config includes "allow_fallback": true. "VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None) ), diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index fb7a170b82fb..764bb73e8825 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -101,6 +101,7 @@ from vllm.model_executor.kernels.linear.mxfp6.emulation import ( EmulationMxfp6LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp6.humming import HummingMxFp6LinearKernel from vllm.model_executor.kernels.linear.mxfp8 import ( Mxfp8LinearKernel, Mxfp8LinearLayerConfig, @@ -295,6 +296,7 @@ def _get_linear_backend(*, quantization: str) -> str: HummingLinearKernel, HummingMxfp8LinearKernel, HummingMxFp4LinearKernel, + HummingMxFp6LinearKernel, HummingNvFp4LinearKernel, }, "marlin": { @@ -574,6 +576,7 @@ def _resolve_backend_kernels( _POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = { PlatformEnum.CUDA: [ + HummingMxFp6LinearKernel, EmulationMxfp6LinearKernel, ], PlatformEnum.ROCM: [ @@ -1294,6 +1297,7 @@ def register_linear_kernel( "MxFp6LinearLayerConfig", "init_mxfp6_linear_kernel", "EmulationMxfp6LinearKernel", + "HummingMxFp6LinearKernel", "AiterMxfp4LinearKernel", "EmulationMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 3de8b9ca4151..b98dfee18b9d 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -24,7 +24,7 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, @@ -57,13 +57,14 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, ) return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index 660e31dcf4f2..0504fc5482ad 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -3,21 +3,18 @@ import torch -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig -logger = init_logger(__name__) - class HummingMxFp4LinearKernel(MxFp4LinearKernel): """Humming GEMM Kernel for MXFP4.""" @@ -39,14 +36,10 @@ def is_supported( @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: - if config.activation_quant_key not in (None, kMxfp4Dynamic): - return False, "only supports MXFP4 dynamic or unquantized activations" - if config.activation_quant_key is not None: - logger.warning_once( - "HummingMxFp4LinearKernel is a weight-only (A16) kernel; " - "the requested activation quantization (%s) is ignored.", - config.activation_quant_key, - ) + try: + quant_key_to_input_schema(config.activation_quant_key) + except ValueError as error: + return False, str(error) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -62,7 +55,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: } convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_schema = quant_key_to_input_schema(self.config.activation_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -75,6 +71,7 @@ def apply_weights( return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py index 320693698752..947ea60e3fdd 100644 --- a/vllm/model_executor/kernels/linear/mxfp6/base.py +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -13,8 +13,8 @@ class MxFp6LinearLayerConfig: """Configuration for an MXFP6 linear layer. - All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per - byte) and per-block weight scales (group size 32). + All MXFP6 layers share the same structure: packed uint8 weights (4 FP6 values + per 3 bytes) and per-block weight scales (group size 32). Attributes: weight_quant_key: Identifies the weight quantization format. Can be diff --git a/vllm/model_executor/kernels/linear/mxfp6/humming.py b/vllm/model_executor/kernels/linear/mxfp6/humming.py new file mode 100644 index 000000000000..fff7584d9332 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/humming.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming import ( + apply_humming_linear, + convert_linear_layer_to_humming_standard, + get_humming_linear_compute_config, + prepare_humming_linear_layer_config, + quant_key_to_input_schema, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp6E2M3Static +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming + +from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig + + +class HummingMxFp6LinearKernel(MxFp6LinearKernel): + """Humming GEMM for packed MXFP6 E2M3 and E3M2 weights.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not has_humming(): + return False, "Humming is not installed" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + try: + quant_key_to_input_schema(config.activation_quant_key) + except ValueError as error: + return False, str(error) + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_dtype = "float6e3m2" + if self.config.weight_quant_key == kMxfp6E2M3Static: + weight_dtype = "float6e2m3" + + quant_config = { + "quant_method": "humming", + "dtype": weight_dtype, + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + + convert_linear_layer_to_humming_standard( + layer=layer, + name_map={"weight": "weight", "weight_scale": "weight_scale"}, + ) + input_schema = quant_key_to_input_schema(self.config.activation_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) + self.compute_config = get_humming_linear_compute_config() + self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return apply_humming_linear( + layer, + x, + skip_bias_add=bias is None, + layer_config=self.layer_config, + compute_config=self.compute_config, + locks=self.locks, + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/humming.py b/vllm/model_executor/kernels/linear/mxfp8/humming.py index bde1565d2a3e..63e9514f4ea3 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp8/humming.py @@ -3,12 +3,14 @@ import torch -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming @@ -50,7 +52,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: } convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_schema = quant_key_to_input_schema(kMxfp8Dynamic) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -63,6 +68,7 @@ def apply_weights( return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py index 54a0cf6c1d8e..3829c39b7486 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/humming.py +++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py @@ -4,11 +4,13 @@ import torch from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming @@ -40,9 +42,7 @@ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Route through humming's compressed-tensors nvfp4 loader (same path as - # the MoE oracle); the native group_tensor schema mishandles a scalar - # global scale. + # Reuse the CT loader for packed FP4 weights and tensor scales. quant_config = { "quant_method": "compressed-tensors", "format": "nvfp4-pack-quantized", @@ -60,7 +60,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight_global_scale = torch.nn.Parameter( 1.0 / layer.weight_global_scale, requires_grad=False ) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_quant_key = None + if hasattr(layer, "input_global_scale"): + layer.input_scale_2 = torch.nn.Parameter( + layer.input_global_scale.reshape(-1), requires_grad=False + ) + input_quant_key = kNvfp4Dynamic + input_schema = quant_key_to_input_schema(input_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -73,6 +82,7 @@ def apply_weights( return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/humming.py b/vllm/model_executor/kernels/linear/scaled_mm/humming.py index b52e73b5a423..3381a3381f86 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/humming.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/humming.py @@ -4,7 +4,7 @@ import torch from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, @@ -92,6 +92,7 @@ def apply_weights( return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, @@ -158,6 +159,7 @@ def apply_weights( return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 18a2fb5cc6a8..a6d44c3141f7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -4,6 +4,7 @@ import json import math +from dataclasses import replace from typing import TYPE_CHECKING, Any, cast import torch @@ -30,6 +31,7 @@ get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, + moe_prepare_scatter, moe_unpermute, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( @@ -57,14 +59,16 @@ kMxfp8Static, kNvfp4Dynamic, kNvfp4Static, + kStaticTensorScale, ) from vllm.platforms import current_platform from vllm.scalar_type import ScalarType from vllm.utils.import_utils import has_humming +from vllm.utils.platform_utils import num_compute_units from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( HummingMoEQuantConfig, ) from vllm.utils.humming import ( @@ -118,38 +122,9 @@ def get_humming_moe_gemm_type(moe_config: FusedMoEConfig | None = None) -> str: return gemm_type -def _fixup_moe_tuning_config(tuning_config: list, max_k_block: int = 128) -> None: - """Fix up each MoE tile in place: cap the K-block and widen warp-N. - - - block_shape[2] (K-block) > ``max_k_block``: the driver rejects the TMA - descriptor at launch (CUDA_ERROR_MISALIGNED_ADDRESS). Cap at 128, which - Humming already uses for larger M. - - warp_shape[1] (warp-N) < 32: block-FP8 (group-128) activations route the - w13 (gate/up) GEMM to a tuning table that pins warp-N to 16, under-filling - the Hopper WGMMA N dimension and corrupting the GEMM output (gsm8k - 0.94 -> 0.89). Widen to 32 whenever block_n % 32 == 0 -- on every tile, - not just the K-capped ones. w2 (down) already uses warp_n=32. - """ - logger.info_once("Attempting to override humming GEMM config") - for entry in tuning_config: - config = entry[2] - block_shape = config.get("block_shape") - if not (block_shape and len(block_shape) == 3): - continue - block_m, block_n, block_k = block_shape - - logger.info_once(f"Overriding humming GEMM config. Previous config\n: {config}") - if block_k > max_k_block: - config["block_shape"] = [block_m, block_n, max_k_block] - - warp_shape = config.get("warp_shape") - if warp_shape and warp_shape[1] < 32 and block_n % 32 == 0: - config["warp_shape"] = [warp_shape[0], 32, warp_shape[2]] - - logger.info_once(f"Overridden humming GEMM config. Current config\n: {config}") - - class HummingExpertsBase(mk.FusedMoEExpertsModular): + quant_config: "HummingMoEQuantConfig" + def __init__( self, moe_config: FusedMoEConfig, @@ -165,7 +140,7 @@ def __init__( self.locks = torch.zeros(1024, dtype=torch.int32, device=moe_config.device) self.num_experts = moe_config.num_local_experts self.global_num_experts = moe_config.num_experts - self.quant_config = quant_config + self.quant_config = humming_quant_config self.init_humming_moe() if self.is_batched(): @@ -199,31 +174,61 @@ def init_humming_moe(self): gemm_type=self.humming_gemm_type(), ) self.compute_config_str = json.dumps(self.compute_config) - # Fix up the heuristic-chosen tiles before freezing to JSON -- see - # _fixup_moe_tuning_config. - _fixup_moe_tuning_config(self.w13_tuning_config) - _fixup_moe_tuning_config(self.w2_tuning_config) self.w13_tuning_config_str = json.dumps(self.w13_tuning_config) self.w2_tuning_config_str = json.dumps(self.w2_tuning_config) - def quantize_input( + def process_input( self, sublayer_name: str, inputs: torch.Tensor, quanted_input: torch.Tensor | None, input_scale: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - from vllm.utils.humming import may_quant_input + expert_tokens: torch.Tensor | None = None, + activation: MoEActivation | None = None, + scatter_idx: torch.Tensor | None = None, + num_valid_tokens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, + ) + from vllm.utils.humming import may_process_input - # input_scale is set for block-FP8 (group-128) activations that were - # quantized before the EP dispatch: may_quant_input then skips the - # redundant w13 quantization and forwards the pre-computed scale. - return may_quant_input( - self.humming_configs[sublayer_name], + quant_config = self.quant_config + config = self.humming_configs[sublayer_name] + if input_scale is not None: + assert scatter_idx is None + if activation is not None: + raise ValueError("Cannot apply activation to prequantized input") + return inputs, input_scale, None + + activation_kwargs = {} + if activation is not None: + activation_config = self.activation_config + activation_kwargs = get_humming_activation(activation, activation_config) + + prefix = "w1" if sublayer_name == "w13" else "w2" + scale = getattr(quant_config, f"{prefix}_input_scale") + scale2 = getattr(quant_config, f"{prefix}_input_scale_2") + mode = config.input_quant_mode + layout = "grouped_mask" if self.is_batched() else "normal" + if scatter_idx is not None: + layout = "scatter" + + inputs, group_scales, token_scales = may_process_input( + config, inputs=inputs, - input_scale=input_scale, - quanted_input=quanted_input, + outputs=quanted_input, + token_scales=scale2 if mode.has_secondary_scale else scale, + hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"), + layout=layout, + expert_tokens=expert_tokens if layout == "grouped_mask" else None, + scatter_idx=scatter_idx, + num_valid_tokens=num_valid_tokens, + **activation_kwargs, ) + input_scale = group_scales if mode.has_group_scale else token_scales + input_scale_2 = token_scales if mode.has_secondary_scale else None + return inputs, input_scale, input_scale_2 def humming_forward( self, @@ -232,30 +237,30 @@ def humming_forward( weight: torch.Tensor, input_scale: torch.Tensor | None, outputs: torch.Tensor, + input_scale_2: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: from vllm.utils.humming import humming_forward - is_w13 = sublayer_name == "w13" + index = 1 if sublayer_name == "w13" else 2 return humming_forward( self.humming_configs[sublayer_name], inputs=inputs, weight=weight, - weight_scale=( - self.quant_config.w1_scale if is_w13 else self.quant_config.w2_scale - ), - zero_point=(self.quant_config.w1_zp if is_w13 else self.quant_config.w2_zp), - bias=(self.quant_config.w1_bias if is_w13 else self.quant_config.w2_bias), - weight_scale_2=( - self.quant_config.g1_alphas if is_w13 else self.quant_config.g2_alphas - ), + weight_scale=getattr(self.quant_config, f"w{index}_scale"), + zero_point=getattr(self.quant_config, f"w{index}_zp"), + bias=getattr(self.quant_config, f"w{index}_bias"), + weight_scale_2=getattr(self.quant_config, f"g{index}_alphas"), input_scale=input_scale, + input_scale_2=input_scale_2, outputs=outputs, locks=self.locks, **kwargs, ) - def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None: + def _get_permute_scratch( + self, topk: int, indices_only: bool = False + ) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None @@ -270,8 +275,8 @@ def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None: num_experts=self.moe_config.num_experts, num_local_experts=self.moe_config.num_local_experts, device=torch.device(self.moe_config.device), - hidden_size=self.moe_config.hidden_dim, - hidden_dtype=self.moe_config.in_dtype, + hidden_size=None if indices_only else self.moe_config.hidden_dim, + hidden_dtype=None if indices_only else self.moe_config.in_dtype, ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): @@ -302,6 +307,16 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: + if ( + activation_key is not None + and activation_key.scale == kNvfp4Dynamic.scale + and activation_key.dtype == kNvfp4Dynamic.dtype + and activation_key.scale2 is not None + and not activation_key.scale2.static + and activation_key.scale2.group_shape == GroupShape.PER_TOKEN + and activation_key.scale2.dtype == torch.float32 + ): + activation_key = replace(activation_key, scale2=kStaticTensorScale) SUPPORTED_W_A = [ (kMxfp4Static, None), (kMxfp4Static, kMxfp4Dynamic), @@ -343,7 +358,7 @@ def _supports_quant_scheme( (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A or ( - activation_key in (None, kFp8DynamicTokenSym) + activation_key in (None, kFp8DynamicTokenSym, kInt8DynamicTokenSym) and _is_supported_wna16_weight_key(weight_key) ) @@ -355,7 +370,7 @@ def _prequantizes_dispatch_activation(self) -> bool: to FP8 before dispatch sends FP8 rather than BF16 over the interconnect, and Humming then consumes the pre-quantized FP8 + scale directly (see the apply() methods, which forward the dispatch scale into - HummingMethod.may_quant_input, and may_quant_input itself, which is a + HummingExpertsBase.process_input, which is a no-op when an input scale is already supplied). The scale layout produced by vLLM's block-FP8 quantization ([M, K // 128] float32, row-major) matches what the Humming WGMMA grouped GEMM expects. @@ -373,7 +388,7 @@ def expects_unquantized_inputs(self) -> bool: inputs). Humming normally quantizes inputs internally via - HummingMethod.may_quant_input() in apply(), so we defer quantization + HummingExpertsBase.process_input() in apply(), so we defer quantization (return True) to avoid quantizing twice -- once in prepare and once in Humming's apply(). @@ -398,8 +413,6 @@ def _supports_no_act_and_mul() -> bool: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - # Humming uses apply_moe_activation() callback for activation, - # so any activation supported there can be used here. return apply_moe_activation_supported(activation) @staticmethod @@ -456,10 +469,9 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): assert intermediate_dim == self.adjust_N_for_activation(gate_up_dim, activation) # hidden_states - # (-> quanted_gate_up_input) (if not BF16/FP16 activation) + # -> quanted_gate_up_input (Hadamard + optional quantization) # -> gate_up_output - # -> activation_output - # (-> quanted_down_input) (if not BF16/FP16 activation) + # -> quanted_down_input (activation + Hadamard + optional quantization) # -> down_output # (-> output) (if not is_batched) # Neighboring nodes are required to utilize distinct workspaces. @@ -471,8 +483,8 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): max_num_tokens = self.max_num_tokens num_dispatchers = self.num_dispatchers assert max_num_tokens is not None and num_dispatchers is not None - input_shape_m = num_experts * max_num_tokens real_shape_m = num_experts * max_num_tokens * num_dispatchers + input_shape_m = real_shape_m output_shape = (num_experts, max_num_tokens * num_dispatchers, K) else: input_shape_m = M @@ -483,15 +495,18 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): a_dtype = self.humming_configs["w13"].a_dtype c_dtype = self.humming_configs["w13"].c_dtype - num_bits = a_dtype.num_bits + down_a_dtype = w2_config.a_dtype torch_dtype_map = { dtypes.float16: torch.float16, dtypes.bfloat16: torch.bfloat16, dtypes.float32: torch.float32, + dtypes.float8e3m4: torch.uint8, dtypes.float8e4m3: torch.float8_e4m3fn, dtypes.float8e5m2: torch.float8_e5m2, dtypes.int8: torch.int8, dtypes.int4: torch.uint8, + dtypes.float4e0m3: torch.uint8, + dtypes.float4e2m1: torch.uint8, } buffer_metas = { @@ -503,13 +518,9 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): "shape": (real_shape_m, gate_up_dim), "dtype": torch_dtype_map[c_dtype], }, - "activation_output": { - "shape": (real_shape_m, intermediate_dim), - "dtype": torch_dtype_map[c_dtype], - }, "quanted_down_input": { "shape": (real_shape_m, intermediate_dim), - "dtype": torch_dtype_map[a_dtype], + "dtype": torch_dtype_map[down_a_dtype], }, "down_output": { "shape": output_shape if self.is_batched() else (real_shape_m, K), @@ -523,7 +534,8 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): for key in buffer_metas: meta = buffer_metas[key] - if "quanted" in key and a_dtype.num_bits == 4: + input_dtype = down_a_dtype if key == "quanted_down_input" else a_dtype + if "quanted" in key and input_dtype.num_bits == 4: last_dim = meta["shape"][-1] if last_dim % 2 != 0: raise ValueError( @@ -532,25 +544,12 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): ) meta["shape"] = meta["shape"][:-1] + (last_dim // 2,) - if num_bits == 16: - required_buffers = ["gate_up_output", "activation_output", "down_output"] - else: - required_buffers = [ - "quanted_gate_up_input", - "gate_up_output", - "activation_output", - "quanted_down_input", - "down_output", - ] - - # Note: The fused SITU+FP8 quant goes gate_up_output -> - # quanted_down_input directly, never materializing - # activation_output. Dropping activation_output from the chain - # flips the even/odd workspace 2-coloring below so - # gate_up_output and quanted_down_input land - # on DIFFERENT workspaces. - if self.fused_situ_quant_enabled(activation): - required_buffers.remove("activation_output") + required_buffers = [ + "quanted_gate_up_input", + "gate_up_output", + "quanted_down_input", + "down_output", + ] # batched moe use down_output as output if not self.is_batched(): @@ -673,81 +672,6 @@ def apply_activation( activation_kwargs["valid_token_counts"] = valid_token_counts self.activation(**activation_kwargs) - def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool: - """Whether the SITU activation + w2 quant can be fused into one kernel. - - Fused only for k-major block-FP8 group-128 e4m3 with float32 scales -- - the sole layout situ_and_mul_quant supports. A float32 as_dtype rules out - MXMMA (which uses e8m0 m-major scales), so a group-128 scale is - guaranteed k-major. Everything else (m-major MXFP8, other group sizes, - 16-bit passthrough, non-SITU) falls back to the separate - situ_and_mul + quantize_input path. - """ - if activation != MoEActivation.SITU: - return False - w2cfg = self.humming_configs["w2"] - # Report the specific blocker once so a silent fallback to the unfused - # situ_and_mul + quant pair is never a mystery in a trace. - reason: str | None = None - if not (w2cfg.a_dtype.num_bits == 8 and str(w2cfg.a_dtype) == "float8e4m3"): - reason = f"w2 a_dtype is {w2cfg.a_dtype} (need float8e4m3)" - elif w2cfg.input_scale_group_size != 128: - reason = ( - f"w2 input_scale_group_size is {w2cfg.input_scale_group_size} " - "(need 128)" - ) - elif not (w2cfg.as_dtype is not None and str(w2cfg.as_dtype) == "float32"): - reason = f"w2 as_dtype is {w2cfg.as_dtype} (need float32, k-major)" - if reason is not None: - logger.warning_once( - "Humming fused SITU+FP8 quant disabled, using unfused " - "situ_and_mul + quant: %s", - reason, - ) - return False - return True - - def fused_situ_quant( - self, - gate_up_output: torch.Tensor, - quanted_down_input: torch.Tensor, - num_valid_tokens: torch.Tensor | None, - topk: int = 1, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the fused SITU activation + FP8 quant for the w2 input. - - Returns ``(quanted_down_input, input_scale)`` in the same layout the - unfused ``apply_activation`` + ``quantize_input("w2")`` pair produces: an - fp8 [M, d] tensor plus a k-major block-FP8 float32 [M, d // 128] group - scale. Only reached when ``fused_situ_quant_enabled`` (group_size 128). - """ - from vllm.model_executor.layers.fused_moe.activation import ( - situ_and_mul_quant, - ) - - cfg = self.activation_config - assert cfg.activation_situ_beta is not None, ( - "SITU requires activation_situ_beta from FusedMoEConfig" - ) - group_size = self.humming_configs["w2"].input_scale_group_size - m, d = quanted_down_input.size(0), quanted_down_input.size(1) - input_scale = torch.empty( - (m, d // group_size), - dtype=torch.float32, - device=quanted_down_input.device, - ) - situ_and_mul_quant( - quanted_down_input, - input_scale, - gate_up_output, - beta=cfg.activation_situ_beta, - linear_beta=cfg.activation_situ_linear_beta, - group_size=group_size, - num_valid_tokens=num_valid_tokens, - topk=topk, - ) - return quanted_down_input, input_scale - class HummingIndexedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -768,7 +692,7 @@ def prepare_humming_moe_kwargs( topk_ids: torch.Tensor, expert_map: torch.Tensor | None, expert_tokens_meta: mk.ExpertTokensMetadata | None, - ) -> tuple[dict[str, Any], dict[str, Any]]: + ) -> tuple[dict[str, Any], dict[str, Any], torch.Tensor | None]: valid_shape_m = self.estimate_local_valid_shape_m(topk_ids) moe_block_size = None @@ -784,13 +708,21 @@ def prepare_humming_moe_kwargs( ) moe_block_size = 64 - sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size( + use_scatter = False + if expert_map is not None: + num_sms = num_compute_units(topk_ids.get_device()) + use_scatter = topk_ids.numel() > 2 * num_sms + + alignment = moe_align_block_size( topk_ids=topk_ids, block_size=moe_block_size, num_experts=self.global_num_experts, expert_map=expert_map, ignore_invalid_experts=True, + return_scatter_idx=use_scatter, ) + sorted_ids, expert_ids, num_tokens_padded = alignment[:3] + scatter_idx = alignment[3] if len(alignment) == 4 else None moe_common_kwargs = { "sorted_ids": sorted_ids, @@ -806,7 +738,26 @@ def prepare_humming_moe_kwargs( moe_kwargs1.update(moe_common_kwargs) moe_kwargs2.update(moe_common_kwargs) - return moe_kwargs1, moe_kwargs2 + w2_block_size = 64 + for lower, upper, config in self.w2_tuning_config: + if lower < valid_shape_m <= upper: + w2_block_size = config["block_shape"][0] + break + if w2_block_size != moe_block_size: + sorted_ids2, expert_ids2, num_tokens_padded2 = moe_align_block_size( + topk_ids=topk_ids, + block_size=w2_block_size, + num_experts=self.global_num_experts, + expert_map=expert_map, + ignore_invalid_experts=True, + ) + moe_kwargs2.update( + sorted_ids=sorted_ids2, + expert_ids=expert_ids2, + num_tokens_padded=num_tokens_padded2, + ) + + return moe_kwargs1, moe_kwargs2, scatter_idx def apply( self, @@ -832,7 +783,7 @@ def apply( object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to - may_quant_input so Humming skips the redundant w13 quantization. The + process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management. """ assert not apply_router_weight_on_input @@ -847,13 +798,13 @@ def apply( ) buffers["output"] = output - moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs( + moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs( topk_ids=topk_ids, expert_map=expert_map, expert_tokens_meta=expert_tokens_meta, ) - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, @@ -865,52 +816,25 @@ def apply( inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], **moe_kwargs1, ) - # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view. - # The fused path and mul_sum consume tokens directly. The unfused - # activation consumes the expanded row count (tokens * topk). - valid_tokens = None - topk = topk_ids.size(1) - if ( - expert_tokens_meta is not None - and expert_tokens_meta.psum_recv_per_rank is not None - ): - valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:] - - if self.fused_situ_quant_enabled(activation): - # Fused SITU + FP8 quant (per-token or block-FP8 group-128) straight - # into the w2 input, skipping the bf16 activation_output round-trip. - inputs, input_scale = self.fused_situ_quant( - gate_up_output=buffers["gate_up_output"], - quanted_down_input=buffers["quanted_down_input"], - num_valid_tokens=valid_tokens, - topk=topk, - ) - else: - valid_token_counts = ( - valid_tokens * topk if valid_tokens is not None else None - ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - valid_token_counts=valid_token_counts, - ) - - inputs, input_scale = self.quantize_input( - "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), - ) + inputs, input_scale, input_scale_2 = self.process_input( + "w2", + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + activation=activation, + scatter_idx=scatter_idx, + ) self.humming_forward( "w2", inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["down_output"].view(-1, hidden_states.size(-1)), **moe_kwargs2, ) @@ -918,6 +842,10 @@ def apply( # expert_map masks any non-local id; num_valid_tokens bounds the # persistent kernel to the real token rows [0, num_recv) so the padding # tail is never iterated (CUDA-graph-safe device scalar). + valid_tokens = None + if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None: + valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:] + moe_fused_mul_sum( inputs=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, @@ -983,21 +911,33 @@ def apply( ) buffers["output"] = output - hidden_states, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute( - hidden_states=hidden_states, - a1q_scale=a1q_scale, - topk_ids=topk_ids, - n_expert=global_num_experts, - n_local_expert=self.num_experts, - expert_map=expert_map, - scratch=self._get_permute_scratch(topk_ids.size(1)), - ) + if a1q_scale is None: + scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True) + assert scratch is not None + scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch) + expert_offsets, scatter_idx = scatter_metadata + inv_perm = scatter_idx.flatten() + num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None + else: + hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute( + hidden_states=hidden_states, + a1q_scale=a1q_scale, + topk_ids=topk_ids, + n_expert=global_num_experts, + n_local_expert=self.num_experts, + expert_map=expert_map, + scratch=self._get_permute_scratch(topk_ids.size(1)), + ) + scatter_idx = None + num_valid_tokens = None - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, quanted_input=buffers.get("quanted_gate_up_input", None), + scatter_idx=scatter_idx, + num_valid_tokens=num_valid_tokens, ) self.humming_forward( @@ -1005,28 +945,20 @@ def apply( inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], valid_shape_m=valid_shape_m, - expert_layout=expert_first_token_offset, + expert_layout=expert_offsets, compute_config=self.compute_config_str, tuning_config=self.w13_tuning_config_str, ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - valid_token_counts=( - expert_first_token_offset[-1:].to(torch.int32) - if expert_tokens_meta is None - else None - ), - ) - - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + activation=activation, + num_valid_tokens=expert_offsets[-1:], ) self.humming_forward( @@ -1034,9 +966,10 @@ def apply( inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["down_output"], valid_shape_m=valid_shape_m, - expert_layout=expert_first_token_offset, + expert_layout=expert_offsets, compute_config=self.compute_config_str, tuning_config=self.w2_tuning_config_str, ) @@ -1046,7 +979,7 @@ def apply( permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, inv_permuted_idx=inv_perm, - expert_first_token_offset=expert_first_token_offset, + expert_first_token_offset=expert_offsets, ) @@ -1088,7 +1021,7 @@ def apply( object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to - may_quant_input so Humming skips the redundant w13 quantization. The + process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management. """ assert not apply_router_weight_on_input @@ -1111,11 +1044,12 @@ def apply( ) buffers["down_output"] = output - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, quanted_input=buffers.get("quanted_gate_up_input", None), + expert_tokens=expert_num_tokens, ) self.humming_forward( @@ -1123,6 +1057,7 @@ def apply( inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], valid_shape_m=valid_shape_m, expert_layout=expert_num_tokens, @@ -1130,16 +1065,12 @@ def apply( tuning_config=self.w13_tuning_config_str, ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - ) - - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + expert_tokens=expert_num_tokens, + activation=activation, ) self.humming_forward( @@ -1147,6 +1078,7 @@ def apply( inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=output.view(-1, hidden_states.size(-1)), valid_shape_m=valid_shape_m, expert_layout=expert_num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py index c8587d3b719b..663475306cc3 100644 --- a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py +++ b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Literal, overload + import torch from vllm import _custom_ops as ops @@ -8,6 +10,7 @@ from vllm.utils.math_utils import round_up +@overload def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, @@ -15,7 +18,52 @@ def moe_align_block_size( expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ignore_invalid_experts: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + *, + return_scatter_idx: Literal[False] = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... + + +@overload +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + *, + return_scatter_idx: Literal[True], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ... + + +@overload +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + *, + return_scatter_idx: bool, +) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +): ... + + +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + return_scatter_idx: bool = False, +) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +): """Aligns the token distribution across experts to be compatible with block size for matrix multiplication. @@ -42,6 +90,8 @@ def moe_align_block_size( as -1. When True, all invalid expert_ids in topk_ids will be ignored and will not participate in counting or ranking, and there will be no -1 in expert_ids. + return_scatter_idx: Whether to additionally return a masked identity + mapping for the original routed rows. Returns: - sorted_token_ids: A tensor containing the sorted token indices according @@ -49,6 +99,10 @@ def moe_align_block_size( - expert_ids: A tensor indicating the assigned expert index for each block. - num_tokens_post_padded: The total number of tokens after padding, ensuring divisibility by block_size. + - scatter_idx: Only returned when return_scatter_idx=True. An int32 tensor + shaped [topk_ids.numel(), 1], containing the original routed row index + for valid routes and -1 otherwise. With expert_map, this requires + ignore_invalid_experts=True. This function pads the number of tokens that each expert needs to process so that it is divisible by block_size. @@ -71,6 +125,15 @@ def moe_align_block_size( by block_size for proper block matrix operations. """ + if return_scatter_idx and expert_map is not None and not ignore_invalid_experts: + raise ValueError("scatter_idx with expert_map requires ignore_invalid_experts") + scatter_idx = None + if return_scatter_idx: + scatter_idx = torch.empty( + (topk_ids.numel(), 1), + dtype=torch.int32, + device=topk_ids.device, + ) max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) if pad_sorted_ids: max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) @@ -95,11 +158,14 @@ def moe_align_block_size( expert_ids, num_tokens_post_pad, expert_map if ignore_invalid_experts else None, + scatter_idx, ) if expert_map is not None and not ignore_invalid_experts: expert_ids = expert_map[expert_ids] + if scatter_idx is not None: + return sorted_ids, expert_ids, num_tokens_post_pad, scatter_idx return sorted_ids, expert_ids, num_tokens_post_pad diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index 0aa9a3807194..d83152fc696a 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -107,6 +107,37 @@ def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: return topk_ids_int32 +def moe_prepare_scatter( + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, + scratch: MoEPermuteScratch, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate expert offsets and shared scatter/unpermute destination indices.""" + assert topk_ids.device == scratch.device + assert topk_ids.size(0) <= scratch.max_num_tokens + assert topk_ids.size(1) == scratch.topk + n_token, topk = topk_ids.shape + expanded_rows = topk_ids.numel() + inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) + if expanded_rows == 0: + scratch.expert_first_token_offset.zero_() + return scratch.expert_first_token_offset, inverse + torch.ops._moe_C.moe_prepare_scatter( + scratch.prepare_topk_ids(topk_ids), + scratch.token_expert_indices_view(n_token), + expert_map, + scratch.num_experts, + scratch.num_local_experts, + scratch.expert_first_token_offset, + inverse, + scratch.sort_workspace, + scratch.permuted_experts_id[:expanded_rows].view(n_token, topk), + scratch.sorted_row_idx[:expanded_rows].view(n_token, topk), + scratch.topk_ids_for_sort[:expanded_rows].view(n_token, topk), + ) + return scratch.expert_first_token_offset, inverse + + def get_moe_permute_scratch( *, max_num_tokens: int, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index abf089620408..5219fff617d3 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -603,7 +603,7 @@ def convert_to_fp8_moe_kernel_format( w13.is_shuffled = True w2.is_shuffled = True elif fp8_backend == Fp8MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -739,7 +739,7 @@ def make_fp8_moe_quant_config( ) elif fp8_backend == Fp8MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 1a23b0173663..29e9888098c1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -193,7 +193,7 @@ def make_int8_moe_quant_config( if int8_backend == Int8MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) @@ -254,7 +254,7 @@ def convert_to_int8_moe_kernel_format( ) -> tuple[torch.Tensor, torch.Tensor]: """Convert INT8 MoE weights to backend-specific kernel format.""" if int8_backend == Int8MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index c70bea9aa33a..ad591d9ea7e4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1456,7 +1456,7 @@ def convert_to_wna16_moe_kernel_format( """ if backend == WNA16MoEBackend.HUMMING: from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7fed4124f569..5450f73f8521 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -844,7 +844,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( sf_block_size = 32 # mxfp4 block size if mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -1428,7 +1428,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) if mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -1966,7 +1966,7 @@ def make_mxfp4_moe_quant_config( gemm1_clamp_limit=swiglu_limit, ) elif mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index ef6a9695f3c3..0d794584d74e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -409,7 +409,7 @@ def convert_to_nvfp4_moe_kernel_format( is_act_and_mul=is_act_and_mul, ) elif nvfp4_backend == NvFp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -523,7 +523,7 @@ def make_nvfp4_moe_quant_config( use_a16 = _use_a16(backend, use_a16) if backend == NvFp4MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 92132f3ffb97..b17a17d0836d 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -713,7 +713,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index ffbc6fba22bf..960ee95d60e4 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -689,7 +689,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index feddb77cffeb..6eb2455d0efd 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -474,7 +474,7 @@ def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: if self.wna16_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 128a5146f159..6bb0faa7c833 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from math import gcd from typing import TYPE_CHECKING, Any import regex as re @@ -28,12 +29,15 @@ QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( + check_and_fallback_input_schema, convert_to_humming_moe_kernel_format, get_humming_linear_compute_config, get_humming_moe_quant_config, + humming_update_schema_hadamard_block_size, input_schema_to_quant_key, make_humming_moe_kernel, + resolve_humming_layer_config, select_humming_moe_experts, weight_schema_to_quant_key, ) @@ -227,23 +231,13 @@ def get_layer_weight_schema(self, config: dict[str, Any], prefix: str): return None config = group_config - layer_config = config - layer_dynamic = config.get("dynamic", {}) - if not isinstance(layer_dynamic, dict): - layer_dynamic = {} - for regex, override_config in layer_dynamic.items(): - if regex[:1] != "+": - continue - if re.match(regex[2:], prefix): - layer_config = config.copy() - layer_config.update(override_config) - break + layer_config = resolve_humming_layer_config(config, prefix) if "quant_method" in layer_config: return _hm.BaseWeightSchema.from_config(layer_config) return None - def get_layer_input_schema(self, config: dict[str, Any], prefix: str): + def get_layer_input_config(self, config: dict[str, Any], prefix: str): if self.is_layer_skipped(config, prefix): return None if config["quant_method"] in ["compressed-tensors", "modelopt"]: @@ -252,10 +246,15 @@ def get_layer_input_schema(self, config: dict[str, Any], prefix: str): return None config = group_config + config = resolve_humming_layer_config(config, prefix) if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP: - return _hm.BaseInputSchema.from_config(config) + return config return None + def get_layer_input_schema(self, config: dict[str, Any], prefix: str): + config = self.get_layer_input_config(config, prefix) + return None if config is None else _hm.BaseInputSchema.from_config(config) + def get_quant_config_for_layer( self, prefix: str, layer_type: str ) -> "HummingLayerQuantizationConfig | None": @@ -279,14 +278,25 @@ def get_quant_config_for_layer( if weight_schema is not None: input_schema = None force_input_schema = None + allow_input_schema_fallback = True if self.full_config: - input_schema = self.get_layer_input_schema(self.full_config, prefix) + input_config = self.get_layer_input_config(self.full_config, prefix) + if input_config is not None: + allow_input_schema_fallback = input_config.pop( + "allow_fallback", True + ) + input_schema = _hm.BaseInputSchema.from_config(input_config) if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG: quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy() quant_config["quant_method"] = "humming" - force_input_schema = self.get_layer_input_schema(quant_config, prefix) + input_config = self.get_layer_input_config(quant_config, prefix) + if input_config is not None: + allow_input_schema_fallback = input_config.pop( + "allow_fallback", False + ) + force_input_schema = _hm.BaseInputSchema.from_config(input_config) if input_schema is None: input_schema = force_input_schema @@ -299,6 +309,7 @@ def get_quant_config_for_layer( force_weight_schema=force_weight_schema, force_input_schema=force_input_schema, is_online_quant=is_online_quant, + allow_input_schema_fallback=allow_input_schema_fallback, ) return None @@ -318,6 +329,11 @@ def get_quant_method( elif isinstance(layer, LinearBase): return UnquantizedLinearMethod() elif isinstance(layer, LinearBase): + input_size = layer.input_size + partition_size = getattr(layer, "input_size_per_partition", input_size) + is_online_quant = quant_config.is_online_quant + if is_online_quant and (input_size % 32 != 0 or partition_size % 32 != 0): + return UnquantizedLinearMethod() return HummingLinearMethod(quant_config) elif isinstance(layer, RoutedExperts): return HummingMoEMethod(quant_config, layer.moe_config) @@ -332,6 +348,7 @@ def __init__( force_weight_schema: "HummingWeightSchema | None" = None, force_input_schema: "HummingInputSchema | None" = None, is_online_quant: bool = False, + allow_input_schema_fallback: bool = True, ): self.weight_schema = weight_schema if input_schema is None: @@ -340,6 +357,7 @@ def __init__( self.force_weight_schema = force_weight_schema self.force_input_schema = force_input_schema self.is_online_quant = is_online_quant + self.allow_input_schema_fallback = allow_input_schema_fallback @classmethod def from_config(cls, config): @@ -373,21 +391,13 @@ def new_weight_loader( if is_unquantized and self.is_online_quant: # online quant (fp16/bf16 -> quant_type) assert isinstance(self.weight_schema, _hm.HummingWeightSchema) - f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype) - has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type) - tensor_list = _hm.quantize_weight( - weight=loaded_weight, - dtype=self.weight_schema.b_dtype, - scale_dtype=self.weight_schema.bs_dtype or f16_dtype, - group_size=self.weight_schema.weight_scale_group_size, - has_zero_point=self.weight_schema.has_zero_point, - has_global_scale=has_global_scale, - is_fp_zero_point=self.weight_schema.is_fp_zero_point, - pack=True, + tensors = self.weight_schema.quant_tensor( + loaded_weight.to(param.device), + self.weight_schema, + layer.param_dtype, ) - key_list = ["weight", "weight_scale", "zero_point", "global_scale"] - for key, tensor in zip(key_list, tensor_list): + for key, tensor in tensors.items(): if tensor is None or tensor.nelement() == 0: continue param = getattr(layer, key) @@ -459,6 +469,18 @@ def create_weights( layer.output_partition_sizes = output_partition_sizes layer.extra_weight_attrs = extra_weight_attrs.copy() + input_schema = self.force_input_schema or self.input_schema + input_schema = input_schema.to_humming_schema(params_dtype) + for name in ("weight_schema", "force_weight_schema"): + schema = getattr(self, name) + if getattr(schema, "hadamard_block_size", None) == -1: + schema = humming_update_schema_hadamard_block_size( + weight_schema=schema, + input_schema=input_schema, + shape_k=input_size_per_partition, + ) + setattr(self, name, schema) + weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) new_weight_loader = self.prepare_weight_loader(layer, weight_loader) extra_weight_attrs["weight_loader"] = new_weight_loader @@ -502,7 +524,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: return None # convert from checkpoint format to humming format - if not isinstance(self.weight_schema, _hm.HummingWeightSchema): + is_humming_weight = isinstance(self.weight_schema, _hm.HummingWeightSchema) + is_humming_input = isinstance(self.input_schema, _hm.HummingInputSchema) + + if not is_humming_weight or not is_humming_input: self.weight_schema, tensors = self.weight_schema.convert_humming( tensors=layer.state_dict(), shape_n_stacks=layer.output_partition_sizes, @@ -510,12 +535,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: param_dtype=layer.param_dtype, ) - self.input_schema, _ = self.input_schema.convert_humming( + self.input_schema, input_tensors = self.input_schema.convert_humming( tensors=layer.state_dict(), shape_n_stacks=layer.output_partition_sizes, shape_k_stacks=[layer.input_size_per_partition], param_dtype=layer.param_dtype, ) + tensors.update(input_tensors) for name, _ in list(layer.named_parameters()): delattr(layer, name) @@ -530,11 +556,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: assert isinstance(self.weight_schema, _hm.HummingWeightSchema) force_requant = self.force_weight_schema is not None if force_requant and self.weight_schema != self.force_weight_schema: + source_tensors = layer.state_dict() tensors = self.weight_schema.requant_tensors( - tensors=layer.state_dict(), + tensors=source_tensors, target_weight_schema=self.force_weight_schema, param_dtype=layer.param_dtype, ) + name = self.input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] self.weight_schema = self.force_weight_schema @@ -548,6 +578,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: del tensors + self.input_schema = check_and_fallback_input_schema( + weight_schema=self.weight_schema, + input_schema=self.input_schema, + param_dtype=layer.param_dtype, + allow_fallback=self.quant_config.allow_input_schema_fallback, + ) self.layer_config = _hm.prepare_layer_config( shape_n=layer.output_partition_sizes_sum, shape_k=layer.input_size_per_partition, @@ -557,10 +593,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: pad_k_to_multiple=128, has_bias=layer.has_bias, torch_dtype=layer.param_dtype, + device=layer.weight.device, ) - tensors = _hm.transform_humming_tensors( - self.layer_config, dict(layer.named_parameters()) - ) + source_tensors = dict(layer.named_parameters()) + tensors = _hm.transform_humming_tensors(self.layer_config, source_tensors) + name = self.input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] + if "bias" in tensors: + zero_bias = torch.zeros_like(tensors["bias"]) + layer.register_buffer("zero_bias", zero_bias) for name, _ in list(layer.named_parameters()): delattr(layer, name) for name, tensor in tensors.items(): @@ -583,8 +625,11 @@ def apply( weight=layer.weight, weight_scale=getattr(layer, "weight_scale", None), zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "bias", None), + bias=getattr(layer, "bias" if bias is not None else "zero_bias", None), weight_scale_2=getattr(layer, "weight_scale_2", None), + input_scale=getattr(layer, "input_scale", None), + input_scale_2=getattr(layer, "input_scale_2", None), + hadamard_block_size=self.weight_schema.hadamard_block_size, locks=self.locks, compute_config=self.compute_config, ) @@ -606,11 +651,15 @@ def __init__( # Derive QuantKeys from humming schemas. # Prefer force schemas (the final format after requant) over base. weight_key = weight_schema_to_quant_key( - self.force_weight_schema or self.weight_schema + self.force_weight_schema or self.weight_schema, moe.in_dtype ) - activation_key = input_schema_to_quant_key( - self.force_input_schema or self.input_schema + runtime_input_schema = check_and_fallback_input_schema( + weight_schema=self.force_weight_schema or self.weight_schema, + input_schema=self.force_input_schema or self.input_schema, + param_dtype=moe.in_dtype, + allow_fallback=quant_config.allow_input_schema_fallback, ) + activation_key = input_schema_to_quant_key(runtime_input_schema, moe.in_dtype) # Select Humming MoE experts self.experts_cls = select_humming_moe_experts( @@ -634,22 +683,14 @@ def new_weight_loader( # online quant (fp16/bf16 -> quant_type) if is_unquantized: assert isinstance(self.weight_schema, _hm.HummingWeightSchema) - f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype) - has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type) - tensor_list = _hm.quantize_weight( - weight=loaded_weight, - dtype=self.weight_schema.b_dtype, - scale_dtype=self.weight_schema.bs_dtype or f16_dtype, - group_size=self.weight_schema.weight_scale_group_size, - has_zero_point=self.weight_schema.has_zero_point, - has_global_scale=has_global_scale, - is_fp_zero_point=self.weight_schema.is_fp_zero_point, - pack=True, + tensors = self.weight_schema.quant_tensor( + loaded_weight.to(param.device), + self.weight_schema, + layer.param_dtype, ) - key_list = ["weight", "weight_scale", "zero_point", "global_scale"] success = True - for key, tensor in zip(key_list, tensor_list): + for key, tensor in tensors.items(): if tensor is None or tensor.nelement() == 0: continue sublayer_name = "w2" if shard_id == "w2" else "w13" @@ -695,6 +736,18 @@ def create_weights( layer.num_experts = num_experts layer.param_dtype = params_dtype layer.intermediate_size = intermediate_size_per_partition + input_schema = self.force_input_schema or self.input_schema + input_schema = input_schema.to_humming_schema(params_dtype) + shape_k = gcd(hidden_size, intermediate_size_per_partition) + for name in ("weight_schema", "force_weight_schema"): + schema = getattr(self, name) + if getattr(schema, "hadamard_block_size", None) == -1: + schema = humming_update_schema_hadamard_block_size( + weight_schema=schema, + input_schema=input_schema, + shape_k=shape_k, + ) + setattr(self, name, schema) weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) weight_loader = self.prepare_weight_loader(layer, weight_loader) extra_weight_attrs["weight_loader"] = weight_loader @@ -761,6 +814,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: weight_schema=self.weight_schema, input_schema=self.input_schema, force_weight_schema=self.force_weight_schema, + allow_input_schema_fallback=self.quant_config.allow_input_schema_fallback, ) # Build the MoE kernel diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index fef2c0c7eeae..6a27332b06c8 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -389,7 +389,7 @@ def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: if self.wna16_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/utils/humming/__init__.py b/vllm/model_executor/layers/quantization/utils/humming/__init__.py new file mode 100644 index 000000000000..4303f74221f0 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/__init__.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Humming quantization integration.""" + +from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, +) +from vllm.model_executor.layers.quantization.utils.humming.linear import ( + apply_humming_linear, + convert_linear_layer_to_humming_standard, + get_humming_linear_compute_config, + prepare_humming_linear_layer_config, +) +from vllm.model_executor.layers.quantization.utils.humming.moe import ( + HummingMoEQuantConfig, + convert_to_humming_moe_kernel_format, + get_humming_moe_quant_config, + make_humming_moe_kernel, + make_humming_moe_quant_config, + select_humming_moe_experts, +) +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + check_and_fallback_input_schema, + humming_is_layer_skipped, + humming_update_schema_hadamard_block_size, + input_schema_to_quant_key, + quant_key_to_input_schema, + resolve_humming_layer_config, + weight_schema_to_quant_key, +) + +__all__ = [ + "check_and_fallback_input_schema", + "get_humming_activation", + "HummingMoEQuantConfig", + "get_humming_moe_quant_config", + "humming_is_layer_skipped", + "humming_update_schema_hadamard_block_size", + "make_humming_moe_quant_config", + "apply_humming_linear", + "convert_linear_layer_to_humming_standard", + "get_humming_linear_compute_config", + "prepare_humming_linear_layer_config", + "make_humming_moe_kernel", + "select_humming_moe_experts", + "input_schema_to_quant_key", + "quant_key_to_input_schema", + "resolve_humming_layer_config", + "weight_schema_to_quant_key", + "convert_to_humming_moe_kernel_format", +] diff --git a/vllm/model_executor/layers/quantization/utils/humming/activation.py b/vllm/model_executor/layers/quantization/utils/humming/activation.py new file mode 100644 index 000000000000..3942944d2dd9 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/activation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE activation expressions for Humming input processing.""" + +from dataclasses import replace + +from vllm.model_executor.layers.fused_moe.activation import ( + ApplyMoEActivationConfig, + MoEActivation, +) + +_SILU_IMPL = "a / (1.f + expf(-a))" +_GELU_IMPL = "0.5f * a * (1.f + erff(a * 0.7071067811865476f))" +_RELU2_IMPL = "fmaxf(a, 0.f) * fmaxf(a, 0.f)" +_GELU_TANH_IMPL = """[](float a) { + const float cubic = 0.044715f * a * a * a; + const float inner = 0.7978845608028654f * (a + cubic); + return 0.5f * a * (1.f + tanhf(inner)); +}(a)""" +_CLAMPED_SILU_IMPL = """[](float a, float b) {{ + const float gate = fminf(a, {limit}); + const float up = fmaxf(fminf(b, {limit}), -{limit}); + const float activated_gate = gate / (1.f + expf(-{alpha} * gate)); + return activated_gate * (up + {beta}); +}}(a, b)""" + +# Parameterized expressions are resolved by get_humming_activation(). +_HUMMING_ACTIVATIONS: dict[MoEActivation, dict[str, str]] = { + MoEActivation.SILU: { + "activation_type": "binary_split", + "activation_impl": f"({_SILU_IMPL}) * b", + }, + MoEActivation.GELU: { + "activation_type": "binary_split", + "activation_impl": f"({_GELU_IMPL}) * b", + }, + MoEActivation.GELU_TANH: { + "activation_type": "binary_split", + "activation_impl": f"({_GELU_TANH_IMPL}) * b", + }, + MoEActivation.RELU2: { + "activation_type": "binary_split", + "activation_impl": f"({_RELU2_IMPL}) * b", + }, + MoEActivation.SWIGLUOAI: { + "activation_type": "binary_interleaved", + "activation_impl": _CLAMPED_SILU_IMPL, + }, + MoEActivation.SITU: { + "activation_type": "binary_split", + "activation_impl": """[](float a, float b) {{ + const float gate = {situ_beta} * tanhf(a / {situ_beta}); + const float activated_gate = gate / (1.f + expf(-a)); + const float up = {situ_up}; + return activated_gate * up; + }}(a, b)""", + }, + MoEActivation.SWIGLUOAI_UNINTERLEAVE: { + "activation_type": "binary_split", + "activation_impl": _CLAMPED_SILU_IMPL, + }, + MoEActivation.SWIGLUSTEP: { + "activation_type": "binary_split", + "activation_impl": """[](float a, float b) { + const float activated_gate = a / (1.f + expf(-a)); + const float gate = fminf(activated_gate, 7.f); + const float up = fmaxf(fminf(b, 7.f), -7.f); + return gate * up; + }(a, b)""", + }, + MoEActivation.SILU_NO_MUL: { + "activation_type": "unary", + "activation_impl": _SILU_IMPL, + }, + MoEActivation.GELU_NO_MUL: { + "activation_type": "unary", + "activation_impl": _GELU_IMPL, + }, + MoEActivation.GELU_TANH_NO_MUL: { + "activation_type": "unary", + "activation_impl": _GELU_TANH_IMPL, + }, + MoEActivation.RELU2_NO_MUL: { + "activation_type": "unary", + "activation_impl": _RELU2_IMPL, + }, +} + + +def get_humming_activation( + activation: MoEActivation, + config: ApplyMoEActivationConfig | None = None, +) -> dict[str, str]: + config = config or ApplyMoEActivationConfig() + use_clamped_silu = False + if activation == MoEActivation.SWIGLUOAI: + config = ApplyMoEActivationConfig(clamp_limit=7.0, alpha=1.702, beta=1.0) + use_clamped_silu = True + elif activation == MoEActivation.SILU and config.clamp_limit is not None: + config = replace(config, alpha=1.0, beta=0.0) + use_clamped_silu = True + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + if config.clamp_limit is None: + raise ValueError("SWIGLUOAI_UNINTERLEAVE requires clamp_limit") + use_clamped_silu = True + + definition = _HUMMING_ACTIVATIONS[activation].copy() + if activation == MoEActivation.SITU: + if config.activation_situ_beta is None: + raise ValueError("SITU requires activation_situ_beta") + linear_beta = config.activation_situ_linear_beta + situ_up = "b" + if linear_beta is not None and linear_beta > 0: + situ_up = f"({float(linear_beta)}f) * tanhf(b / ({float(linear_beta)}f))" + situ_beta = f"({float(config.activation_situ_beta)}f)" + activation_impl = definition["activation_impl"] + activation_impl = activation_impl.format(situ_beta=situ_beta, situ_up=situ_up) + definition["activation_impl"] = activation_impl + elif use_clamped_silu: + assert config.clamp_limit is not None + limit = f"({float(config.clamp_limit)}f)" + alpha = f"({float(config.alpha)}f)" + beta = f"({float(config.beta)}f)" + activation_impl = _CLAMPED_SILU_IMPL.format(limit=limit, alpha=alpha, beta=beta) + definition["activation_impl"] = activation_impl + return definition diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py new file mode 100644 index 000000000000..339ff640b03d --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepare and execute Humming linear layers.""" + +import json +from typing import TYPE_CHECKING + +import torch + +from vllm import envs +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + check_and_fallback_input_schema, +) + +if TYPE_CHECKING: + from vllm.utils.humming import BaseInputSchema, LayerConfig + + +def convert_linear_layer_to_humming_standard( + layer: LinearBase, name_map: dict[str, str] +): + """Rename/reshape a linear layer's quantized params (the canonical MPLinear + layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names + and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" + for name, checkpoint_name in name_map.items(): + tensor = getattr(layer, checkpoint_name) + delattr(layer, checkpoint_name) + + if name == "weight": + input_dim = getattr(tensor, "input_dim", 1) + output_dim = getattr(tensor, "output_dim", 0) + + if input_dim == 0 and output_dim == 1: + tensor = tensor.transpose(1, 0).contiguous() + else: + assert output_dim == 0 and input_dim == 1 + + tensor = tensor.view(tensor.size(0), -1).view(torch.int32) + elif name in ["weight_scale", "zero_point"]: + if getattr(tensor, "output_dim", 0) == 1: + tensor = tensor.transpose(0, 1).contiguous() + if tensor.ndim == 1: + tensor = tensor.unsqueeze(1) + + tensor = tensor.view(torch.int32) if name == "zero_point" else tensor + + if isinstance(tensor, torch.nn.Parameter): + param = tensor + else: + param = torch.nn.Parameter(tensor, requires_grad=False) + + setattr(layer, name, param) + + +def prepare_humming_linear_layer_config( + layer: LinearBase, + quant_config: dict, + input_quant_config: dict | None = None, + *, + input_schema: "BaseInputSchema | None" = None, +) -> "LayerConfig": + from vllm.utils.humming import ( + BaseInputSchema, + BaseWeightSchema, + HummingInputSchema, + prepare_layer_config, + transform_humming_tensors, + ) + + weight_schema = BaseWeightSchema.from_config(quant_config) + if input_schema is None: + if input_quant_config is not None: + input_schema = BaseInputSchema.from_config(input_quant_config) + else: + input_schema = HummingInputSchema() + + # ReplicatedLinear has no TP partitioning and so does not set + # input_size_per_partition; for it that is just input_size. Use hasattr + # rather than getattr's default arg, which is evaluated eagerly and would + # raise on layers lacking input_size (e.g. ParallelLMHead). + if hasattr(layer, "input_size_per_partition"): + input_size_per_partition = layer.input_size_per_partition + else: + input_size_per_partition = layer.input_size + shape_k_stacks = [input_size_per_partition] + shape_n_stacks = layer.output_partition_sizes + + # Step 1: convert weight and input schemas to humming standard format + source_tensors = dict(layer.named_parameters()) + weight_schema, tensors = weight_schema.convert_humming( + tensors=source_tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) + input_schema, input_tensors = input_schema.convert_humming( + tensors=source_tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) + + tensors.update(input_tensors) + layer.weight_schema = weight_schema + input_schema = check_and_fallback_input_schema( + weight_schema=weight_schema, + input_schema=input_schema, + param_dtype=layer.params_dtype, + ) + for name in ("input_scale", "input_scale_2"): + if name != input_schema.static_tensor_scale_name: + tensors.pop(name, None) + input_tensors.pop(name, None) + + # Step 2: transform weight (humming standard format) for forwarding. + config = prepare_layer_config( + shape_n=sum(layer.output_partition_sizes), + shape_k=input_size_per_partition, + weight_schema=weight_schema, + input_schema=input_schema, + pad_n_to_multiple=256, + pad_k_to_multiple=128, + has_bias=layer.has_bias, + torch_dtype=layer.params_dtype, + device=tensors["weight"].device, + ) + tensors = transform_humming_tensors(config, tensors) + tensors.update(input_tensors) + if "bias" in tensors: + zero_bias = torch.zeros_like(tensors["bias"]) + layer.register_buffer("zero_bias", zero_bias) + for name, _ in list(layer.named_parameters()): + delattr(layer, name) + for name, tensor in tensors.items(): + param = torch.nn.Parameter(tensor, requires_grad=False) + setattr(layer, name, param) + + return config + + +def get_humming_linear_compute_config() -> str: + return json.dumps( + { + "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, + "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, + "gemm_type": "dense", + } + ) + + +def apply_humming_linear( + layer: LinearBase, + x: torch.Tensor, + *, + skip_bias_add: bool = False, + layer_config: "LayerConfig", + compute_config: str, + locks: torch.Tensor, +) -> torch.Tensor: + from vllm.utils.humming import humming_forward + + mode = layer_config.input_quant_mode + input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None + input_scale_2 = ( + getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None + ) + flatten_inputs = x.reshape(-1, x.size(-1)) + output = humming_forward( + layer_config, + inputs=flatten_inputs, + weight=layer.weight, + weight_scale=getattr(layer, "weight_scale", None), + zero_point=getattr(layer, "zero_point", None), + bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None), + weight_scale_2=getattr(layer, "weight_scale_2", None), + input_scale=input_scale, + input_scale_2=input_scale_2, + hadamard_block_size=layer.weight_schema.hadamard_block_size, + locks=locks, + compute_config=compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming/moe.py similarity index 51% rename from vllm/model_executor/layers/quantization/utils/humming_utils.py rename to vllm/model_executor/layers/quantization/utils/humming/moe.py index f48c33d4cd31..37eb9c26f07d 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming/moe.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json +"""Configure, prepare weights for, and assemble Humming MoE kernels.""" + from dataclasses import dataclass from typing import TYPE_CHECKING, Any -import regex as re import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +import vllm.model_executor.layers.quantization.utils.humming.schema as humming_schema from vllm import envs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -23,34 +24,19 @@ HummingGroupedExperts, HummingIndexedExperts, ) -from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + humming_is_layer_skipped, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( - FP4_DTYPE, FP8_DTYPE, - INT4_DTYPE, - INT8_DTYPE, - MXFP_SCALE_DTYPE, GroupShape, QuantKey, - ScaleDesc, ) from vllm.utils.import_utils import has_humming if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from vllm.utils.humming import ( - AWQWeightSchema, - BaseInputSchema, - BaseWeightSchema, - CompressedTensorsInputSchema, - CompressedTensorsWeightSchema, - Fp8WeightSchema, - GPTQWeightSchema, - HummingInputSchema, - HummingWeightSchema, - LayerConfig, - ) - from vllm.utils.humming import dtypes as humming_dtypes + from vllm.utils.humming import LayerConfig logger = init_logger(__name__) @@ -59,504 +45,12 @@ class HummingMoEQuantConfig(FusedMoEQuantConfig): w1_humming_config: "LayerConfig" w2_humming_config: "LayerConfig" - - -if has_humming(): - from vllm.utils.humming import dtypes as humming_dtypes - - _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = { - humming_dtypes.float4e2m1: FP4_DTYPE, - humming_dtypes.float8e4m3: FP8_DTYPE, - humming_dtypes.float8e5m2: torch.float8_e5m2, - humming_dtypes.int8: torch.int8, - humming_dtypes.uint4: INT4_DTYPE, - humming_dtypes.uint8: INT8_DTYPE, - humming_dtypes.uint2: torch.uint8, - humming_dtypes.uint3: torch.uint8, - humming_dtypes.uint5: torch.uint8, - humming_dtypes.uint6: torch.uint8, - humming_dtypes.uint7: torch.uint8, - } - - _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { - humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE, - humming_dtypes.float8e4m3: FP8_DTYPE, - humming_dtypes.float16: torch.float16, - humming_dtypes.bfloat16: torch.bfloat16, - humming_dtypes.float32: torch.float32, - } - - -def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: - """Map humming group sizes to QuantKey GroupShape. - - group_size: elements per group along K (col); 0 means full dimension. - group_size_n: elements per group along N (row); 0 means 1 (per-row). - - GroupShape convention: row = N dim, col = K dim. - """ - if group_size == 0 and group_size_n == 0: - return GroupShape.PER_CHANNEL - - row = group_size_n if group_size_n > 0 else 1 - col = group_size if group_size > 0 else -1 - return GroupShape(row=row, col=col) - - -# ---- HummingWeightSchema (post-conversion) -------------------------------- - - -def _humming_weight_schema_to_quant_key( - schema: "HummingWeightSchema", -) -> QuantKey: - from vllm.utils.humming import WeightScaleType - - """Convert a HummingWeightSchema to a QuantKey.""" - dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype] - - if schema.bs_dtype is not None: - scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype] - else: - scale_dtype = torch.float32 - - group_shape = _group_shape( - schema.weight_scale_group_size, - schema.weight_scale_group_size_n, - ) - - scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) - - scale2 = None - if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR: - scale2 = ScaleDesc( - dtype=torch.float32, - static=True, - group_shape=GroupShape.PER_TENSOR, - ) - - return QuantKey( - dtype=dtype, - scale=scale, - scale2=scale2, - symmetric=not schema.has_zero_point, - ) - - -# ---- Checkpoint-format weight schemas (pre-conversion) -------------------- - - -def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey: - if schema.weight_block_size is not None: - gs_n, gs_k = schema.weight_block_size - group_shape = GroupShape(row=gs_n, col=gs_k) - else: - group_shape = GroupShape.PER_CHANNEL - - scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape) - return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True) - - -def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey: - group_shape = _group_shape(schema.group_size) - scale = ScaleDesc( - dtype=torch.float16, - static=True, - group_shape=group_shape, - ) - return QuantKey( - dtype=INT4_DTYPE, - scale=scale, - symmetric=not schema.zero_point, - ) - - -def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey: - group_shape = _group_shape(schema.group_size) - scale = ScaleDesc( - dtype=torch.float16, - static=True, - group_shape=group_shape, - ) - return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym) - - -def _compressed_tensors_weight_schema_to_quant_key( - schema: "CompressedTensorsWeightSchema", -) -> QuantKey: - # Determine dtype from format/type/num_bits - fmt = schema.format - if fmt in ("int-quantized", "float-quantized", "naive-quantized"): - dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE - elif "nvfp4" in fmt or "mxfp4" in fmt: - dtype = FP4_DTYPE - else: - dtype = _HUMMING_TO_QUANT_DTYPE[ - humming_dtypes.DataType.from_str(f"uint{schema.num_bits}") - ] - - # Determine group shape from strategy - if schema.strategy in ("group", "tensor_group"): - group_shape = _group_shape(schema.group_size or 0) - elif schema.strategy == "block" and schema.block_structure is not None: - group_shape = GroupShape( - row=schema.block_structure[0], - col=schema.block_structure[1], - ) - else: - group_shape = GroupShape.PER_CHANNEL - - # Determine scale dtype - if "mxfp" in fmt: - scale_dtype = MXFP_SCALE_DTYPE - elif "nvfp4" in fmt: - scale_dtype = FP8_DTYPE - else: - scale_dtype = torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) - - scale2 = None - if "nvfp4" in fmt or schema.strategy == "tensor_group": - scale2 = ScaleDesc( - dtype=torch.float32, - static=True, - group_shape=GroupShape.PER_TENSOR, - ) - - return QuantKey( - dtype=dtype, - scale=scale, - scale2=scale2, - symmetric=schema.symmetric, - ) - - -# ---- Dispatch for any BaseWeightSchema ------------------------------------ - - -def weight_schema_to_quant_key( - schema: "BaseWeightSchema", -) -> QuantKey: - from vllm.utils.humming import ( - AWQWeightSchema, - BitnetWeightSchema, - CompressedTensorsWeightSchema, - Fp8WeightSchema, - GptOssMxfp4WeightSchema, - GPTQWeightSchema, - HummingWeightSchema, - ModeloptMxfp8WeightSchema, - ModeloptNvfp4WeightSchema, - Mxfp4WeightSchema, - ) - - """Convert any BaseWeightSchema to a QuantKey.""" - if isinstance(schema, HummingWeightSchema): - return _humming_weight_schema_to_quant_key(schema) - - # Schemas with fixed QuantKeys - if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)): - return QuantKey( - dtype=FP4_DTYPE, - scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), - ) - if isinstance(schema, ModeloptMxfp8WeightSchema): - return QuantKey( - dtype=FP8_DTYPE, - scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), - ) - if isinstance(schema, ModeloptNvfp4WeightSchema): - return QuantKey( - dtype=FP4_DTYPE, - scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)), - scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR), - ) - if isinstance(schema, BitnetWeightSchema): - return QuantKey( - dtype=torch.uint8, - scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), - ) - - # Schemas requiring config inspection - if isinstance(schema, Fp8WeightSchema): - return _fp8_weight_schema_to_quant_key(schema) - if isinstance(schema, AWQWeightSchema): - return _awq_weight_schema_to_quant_key(schema) - if isinstance(schema, GPTQWeightSchema): - return _gptq_weight_schema_to_quant_key(schema) - if isinstance(schema, CompressedTensorsWeightSchema): - return _compressed_tensors_weight_schema_to_quant_key(schema) - - raise TypeError(f"Unsupported weight schema type: {type(schema)}") - - -# ---- HummingInputSchema (post-conversion) ---------------------------------- - - -def _humming_input_schema_to_quant_key( - schema: "HummingInputSchema", -) -> QuantKey | None: - """Convert a HummingInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs.""" - if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: - return None - - dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype] - - gs = schema.input_scale_group_size - group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN - - # Pick the scale dtype the Humming kernel actually consumes. An explicit - # input_scale_dtype always wins. Otherwise infer from the grouping: MX - # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8) - # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a - # float32 scale. Getting this right lets a grouped FP8 activation match - # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key. - if schema.input_scale_dtype is not None: - scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype] - elif gs == 32: - scale_dtype = MXFP_SCALE_DTYPE - else: - scale_dtype = torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) - - return QuantKey(dtype=dtype, scale=scale, symmetric=True) - - -# ---- Checkpoint-format input schemas (pre-conversion) ---------------------- - - -def _resolve_input_quant_key( - origin_a_dtype: "humming_dtypes.DataType", - group_size: int, -) -> QuantKey | None: - from vllm.utils.humming import HummingInputSchema - - """Resolve the actual activation QuantKey after platform fallback.""" - a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype) - if a_dtype is None or a_dtype.num_bits >= 16: - return None - - dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype] - gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0 - group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN - scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) - return QuantKey(dtype=dtype, scale=scale, symmetric=True) - - -def _compressed_tensors_input_schema_to_quant_key( - schema: "CompressedTensorsInputSchema", -) -> QuantKey | None: - type_bits_to_dtype = { - ("float", 8): humming_dtypes.float8e4m3, - ("float", 4): humming_dtypes.float4e2m1, - ("int", 8): humming_dtypes.int8, - ("int", 4): humming_dtypes.int4, - } - origin = type_bits_to_dtype.get((schema.type, schema.num_bits)) - if origin is None: - return None - return _resolve_input_quant_key(origin, schema.group_size) - - -# ---- Dispatch for any BaseInputSchema ------------------------------------- - - -def input_schema_to_quant_key( - schema: "BaseInputSchema", -) -> QuantKey | None: - from vllm.utils.humming import ( - CompressedTensorsInputSchema, - Fp8InputSchema, - HummingInputSchema, - ModeloptNvfp4InputSchema, - ) - - """Convert any BaseInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs.""" - if isinstance(schema, HummingInputSchema): - return _humming_input_schema_to_quant_key(schema) - - if isinstance(schema, Fp8InputSchema): - return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0) - - if isinstance(schema, ModeloptNvfp4InputSchema): - return _resolve_input_quant_key( - humming_dtypes.float8e4m3, - schema.group_size, - ) - - if isinstance(schema, CompressedTensorsInputSchema): - return _compressed_tensors_input_schema_to_quant_key(schema) - - raise TypeError(f"Unsupported input schema type: {type(schema)}") - - -def humming_is_layer_skipped(config: dict[str, Any], prefix: str): - if not config: - return True - - keys = ["ignored_layers", "ignore", "modules_to_not_convert"] - ignored_layers: list[str] = [] - for key in keys: - candidate = config.get(key, []) or [] - if candidate: - ignored_layers = candidate - break - - if any(module_name in prefix for module_name in ignored_layers): - return True - if "lm_head" in prefix: - return True - - for regex in config.get("dynamic", {}): - if regex[:1] != "-": - continue - if re.match(regex[2:], prefix): - return True - - return False - - -def convert_linear_layer_to_humming_standard( - layer: LinearBase, name_map: dict[str, str] -): - """Rename/reshape a linear layer's quantized params (the canonical MPLinear - layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names - and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" - for name, checkpoint_name in name_map.items(): - tensor = getattr(layer, checkpoint_name) - delattr(layer, checkpoint_name) - - if name == "weight": - input_dim = getattr(tensor, "input_dim", 1) - output_dim = getattr(tensor, "output_dim", 0) - - if input_dim == 0 and output_dim == 1: - tensor = tensor.transpose(1, 0).contiguous() - else: - assert output_dim == 0 and input_dim == 1 - - tensor = tensor.view(tensor.size(0), -1).view(torch.int32) - elif name in ["weight_scale", "zero_point"]: - if getattr(tensor, "output_dim", 0) == 1: - tensor = tensor.transpose(0, 1).contiguous() - if tensor.ndim == 1: - tensor = tensor.unsqueeze(1) - - tensor = tensor.view(torch.int32) if name == "zero_point" else tensor - - if isinstance(tensor, torch.nn.Parameter): - param = tensor - else: - param = torch.nn.Parameter(tensor, requires_grad=False) - - setattr(layer, name, param) - - -def prepare_humming_linear_layer_config( - layer: LinearBase, - quant_config: dict, - input_quant_config: dict | None = None, -) -> "LayerConfig": - from vllm.utils.humming import ( - BaseInputSchema, - BaseWeightSchema, - HummingInputSchema, - prepare_layer_config, - transform_humming_tensors, - ) - - weight_schema = BaseWeightSchema.from_config(quant_config) - if input_quant_config is not None: - input_schema = BaseInputSchema.from_config(input_quant_config) - else: - input_schema = HummingInputSchema() - - # ReplicatedLinear has no TP partitioning and so does not set - # input_size_per_partition; for it that is just input_size. Use hasattr - # rather than getattr's default arg, which is evaluated eagerly and would - # raise on layers lacking input_size (e.g. ParallelLMHead). - if hasattr(layer, "input_size_per_partition"): - input_size_per_partition = layer.input_size_per_partition - else: - input_size_per_partition = layer.input_size - shape_k_stacks = [input_size_per_partition] - shape_n_stacks = layer.output_partition_sizes - - # Step 1: convert weight and input schemas to humming standard format - weight_schema, tensors = weight_schema.convert_humming( - tensors=dict(layer.named_parameters()), - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.params_dtype, - ) - input_schema, _ = input_schema.convert_humming( - tensors={}, - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.params_dtype, - ) - - layer.weight_schema = weight_schema - - # Step 2: transform weight (humming standard format) for forwarding. - config = prepare_layer_config( - shape_n=sum(layer.output_partition_sizes), - shape_k=input_size_per_partition, - weight_schema=weight_schema, - input_schema=input_schema, - pad_n_to_multiple=256, - pad_k_to_multiple=128, - has_bias=layer.has_bias, - torch_dtype=layer.params_dtype, - ) - tensors = transform_humming_tensors(config, tensors) - for name, _ in list(layer.named_parameters()): - delattr(layer, name) - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - - return config - - -def get_humming_linear_compute_config() -> str: - return json.dumps( - { - "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, - "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, - "gemm_type": "dense", - } - ) - - -def apply_humming_linear( - layer: LinearBase, - x: torch.Tensor, - *, - layer_config: "LayerConfig", - compute_config: str, - locks: torch.Tensor, -) -> torch.Tensor: - from vllm.utils.humming import humming_forward - - flatten_inputs = x.reshape(-1, x.size(-1)) - output = humming_forward( - layer_config, - inputs=flatten_inputs, - weight=layer.weight, - weight_scale=getattr(layer, "weight_scale", None), - zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "bias", None), - weight_scale_2=getattr(layer, "weight_scale_2", None), - locks=locks, - compute_config=compute_config, - ) - return output.view(*x.shape[:-1], output.size(-1)) + w1_hadamard_block_size: int = 1 + w2_hadamard_block_size: int = 1 + w1_input_scale: torch.Tensor | None = None + w2_input_scale: torch.Tensor | None = None + w1_input_scale_2: torch.Tensor | None = None + w2_input_scale_2: torch.Tensor | None = None def make_humming_moe_quant_config( @@ -648,8 +142,12 @@ def get_humming_moe_quant_config( q_dtype is not None and q_dtype.startswith("float8") and input_scale_group_size == 128 + and weight_schema.hadamard_block_size <= 1 + and input_schema.input_quant_mode in (None, "dynamic_group") ): - q_dtype = _HUMMING_TO_QUANT_DTYPE.get(input_schema.a_dtype, FP8_DTYPE) + q_dtype = humming_schema._HUMMING_TO_QUANT_DTYPE.get( + input_schema.a_dtype, FP8_DTYPE + ) activation_group_shape = GroupShape(row=1, col=input_scale_group_size) weight_scale_group_size = weight_schema.weight_scale_group_size @@ -665,7 +163,7 @@ def get_humming_moe_quant_config( else: weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) - return make_humming_moe_quant_config( + config = make_humming_moe_quant_config( quant_dtype=q_dtype, weight_dtype=str(weight_schema.b_dtype), weight_group_shape=weight_group_shape, @@ -683,6 +181,17 @@ def get_humming_moe_quant_config( gemm1_clamp_limit=gemm1_clamp_limit, humming_configs=humming_configs, ) + for prefix, sublayer in (("w1", "w13"), ("w2", "w2")): + setattr( + config, + f"{prefix}_hadamard_block_size", + layer.weight_schemas[sublayer].hadamard_block_size, + ) + for name in ("input_scale", "input_scale_2"): + key = f"{prefix}_{name}" + value = getattr(layer, f"{sublayer}_{name}", None) + setattr(config, key, value) + return config def select_humming_moe_experts( @@ -846,12 +355,6 @@ def _convert_sublayer_to_humming( Tuple of (converted_weight_schema, converted_input_schema) """ - from vllm.utils.humming import HummingWeightSchema - - if isinstance(weight_schema, HummingWeightSchema): - # Already in Humming format - return weight_schema, input_schema - tensors = _extract_sublayer_tensors(layer, sublayer_name) shape_k_stacks = [shape_k] @@ -867,13 +370,14 @@ def _convert_sublayer_to_humming( num_experts=num_experts, ) - converted_input_schema, _ = input_schema.convert_humming( - tensors=converted_tensors, + converted_input_schema, input_tensors = input_schema.convert_humming( + tensors=tensors, shape_n_stacks=shape_n_stacks, shape_k_stacks=shape_k_stacks, param_dtype=param_dtype, num_experts=num_experts, ) + converted_tensors.update(input_tensors) _replace_layer_parameters(layer, sublayer_name, converted_tensors) @@ -907,11 +411,16 @@ def _prepare_and_transform_sublayer( has_bias=has_bias, num_experts=num_experts, torch_dtype=param_dtype, + device=getattr(layer, f"{sublayer_name}_weight").device, ) - tensors = transform_humming_tensors( - config, - _extract_sublayer_tensors(layer, sublayer_name), - ) + source_tensors = _extract_sublayer_tensors(layer, sublayer_name) + for name in ("input_scale", "input_scale_2"): + if name != input_schema.static_tensor_scale_name: + source_tensors.pop(name, None) + tensors = transform_humming_tensors(config, source_tensors) + for name in ("input_scale", "input_scale_2"): + if name in source_tensors: + tensors[name] = source_tensors[name] _replace_layer_parameters(layer, sublayer_name, tensors) return config @@ -927,6 +436,7 @@ def _process_single_sublayer( num_experts: int, param_dtype: torch.dtype, force_weight_schema: Any | None = None, + allow_input_schema_fallback: bool = True, ) -> tuple[Any, Any, "LayerConfig"]: """Process a single sublayer: convert, optionally requant, prepare, and transform. @@ -944,6 +454,7 @@ def _process_single_sublayer( num_experts: Number of experts param_dtype: Parameter data type force_weight_schema: Optional schema to force requantization to + allow_input_schema_fallback: Whether incompatible input schemas may be replaced. Returns: Tuple of the final weight schema, input schema, and Humming layer config. @@ -966,19 +477,28 @@ def _process_single_sublayer( # Step 2: Force requant if needed assert isinstance(current_weight_schema, HummingWeightSchema) if force_weight_schema is not None and current_weight_schema != force_weight_schema: - tensors = _extract_sublayer_tensors(layer, sublayer_name) + source_tensors = _extract_sublayer_tensors(layer, sublayer_name) tensors = current_weight_schema.requant_tensors( - tensors=tensors, + tensors=source_tensors, target_weight_schema=force_weight_schema, param_dtype=param_dtype, ) + name = current_input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] current_weight_schema = force_weight_schema _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True) del tensors # Step 3: Prepare layer metadata and transform weights + current_input_schema = humming_schema.check_and_fallback_input_schema( + weight_schema=current_weight_schema, + input_schema=current_input_schema, + param_dtype=param_dtype, + allow_fallback=allow_input_schema_fallback, + ) config = _prepare_and_transform_sublayer( layer=layer, sublayer_name=sublayer_name, @@ -1001,6 +521,7 @@ def convert_to_humming_moe_kernel_format( weight_schema: Any | None = None, input_schema: Any | None = None, force_weight_schema: Any | None = None, + allow_input_schema_fallback: bool = True, ) -> dict[str, "LayerConfig"]: """Convert MoE weights from checkpoint format to Humming kernel format. @@ -1023,6 +544,7 @@ def convert_to_humming_moe_kernel_format( input_schema: Optional initial input quantization schema. If None, built from quant_config or env vars. force_weight_schema: Optional schema to force requantization to + allow_input_schema_fallback: Whether incompatible input schemas may be replaced. Side effects: - Modifies layer parameters in place @@ -1041,20 +563,23 @@ def convert_to_humming_moe_kernel_format( "Must provide either weight_schema/input_schema or quant_config" ) - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - humming_is_layer_skipped, - ) from vllm.utils.humming import BaseWeightSchema, HummingInputSchema if weight_schema is None: weight_schema = BaseWeightSchema.from_config(quant_config) if input_schema is None: - input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {} + input_quant_config = (envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}).copy() if humming_is_layer_skipped(input_quant_config, layer.layer_name): input_schema = HummingInputSchema() else: # TODO: read input_quant_config from quant_config + input_quant_config = humming_schema.resolve_humming_layer_config( + input_quant_config, layer.layer_name + ) + allow_input_schema_fallback = input_quant_config.pop( + "allow_fallback", False + ) input_schema = HummingInputSchema.from_config(input_quant_config) # Build sublayer configs from layer properties if not provided @@ -1089,6 +614,7 @@ def convert_to_humming_moe_kernel_format( num_experts=num_experts, param_dtype=param_dtype, force_weight_schema=force_weight_schema, + allow_input_schema_fallback=allow_input_schema_fallback, ) ) diff --git a/vllm/model_executor/layers/quantization/utils/humming/schema.py b/vllm/model_executor/layers/quantization/utils/humming/schema.py new file mode 100644 index 000000000000..c84e8a370d22 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/schema.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Map Humming schemas and handle shared checkpoint quantization settings.""" + +import dataclasses +from typing import TYPE_CHECKING, Any + +import regex as re +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP4_DTYPE, + FP8_DTYPE, + INT4_DTYPE, + INT8_DTYPE, + MXFP_SCALE_DTYPE, + GroupShape, + QuantKey, + ScaleDesc, +) +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types +from vllm.utils.import_utils import has_humming + +if TYPE_CHECKING: + from vllm.utils.humming import ( + BaseInputSchema, + BaseWeightSchema, + HummingInputSchema, + HummingWeightSchema, + ) + from vllm.utils.humming import dtypes as humming_dtypes + + +def humming_update_schema_hadamard_block_size( + weight_schema: "HummingWeightSchema", + input_schema: "HummingInputSchema", + shape_k: int, +) -> "HummingWeightSchema": + assert shape_k > 0 + block_size = 32 + weight_group_size = weight_schema.weight_scale_group_size + input_group_size = input_schema.input_scale_group_size + for group_size in (weight_group_size, input_group_size): + if group_size is not None and group_size > 0: + block_size = min(block_size, group_size) + block_size = 1 << (block_size.bit_length() - 1) + while shape_k % block_size: + block_size //= 2 + return dataclasses.replace(weight_schema, hadamard_block_size=block_size) + + +if has_humming(): + from vllm.utils.humming import dtypes as humming_dtypes + + _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = { + humming_dtypes.float4e0m3: FP4_DTYPE, + humming_dtypes.float4e2m1: FP4_DTYPE, + humming_dtypes.float6e2m3: scalar_types.float6_e2m3f, + humming_dtypes.float6e3m2: scalar_types.float6_e3m2f, + humming_dtypes.float8e3m4: FP8_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float8e5m2: torch.float8_e5m2, + humming_dtypes.int8: torch.int8, + humming_dtypes.uint4: INT4_DTYPE, + humming_dtypes.uint8: INT8_DTYPE, + humming_dtypes.uint2: torch.uint8, + humming_dtypes.uint3: torch.uint8, + humming_dtypes.uint5: torch.uint8, + humming_dtypes.uint6: torch.uint8, + humming_dtypes.uint7: torch.uint8, + } + + _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { + humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float16: torch.float16, + humming_dtypes.bfloat16: torch.bfloat16, + humming_dtypes.float32: torch.float32, + } + + +logger = init_logger(__name__) + + +def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: + """Map humming group sizes to QuantKey GroupShape. + + group_size: elements per group along K (col); 0 means full dimension. + group_size_n: elements per group along N (row); 0 means 1 (per-row). + + GroupShape convention: row = N dim, col = K dim. + """ + if group_size == 0 and group_size_n == 0: + return GroupShape.PER_CHANNEL + + row = group_size_n if group_size_n > 0 else 1 + col = group_size if group_size > 0 else -1 + return GroupShape(row=row, col=col) + + +def weight_schema_to_quant_key( + schema: "BaseWeightSchema", + param_dtype: torch.dtype | None = None, +) -> QuantKey: + from vllm.utils.humming import WeightScale2Type, WeightScaleType + + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in [torch.float16, torch.bfloat16]: + param_dtype = torch.bfloat16 + if not current_platform.has_device_capability(80): + param_dtype = torch.float16 + + schema = schema.to_humming_schema(param_dtype) + dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype] + + if schema.bs_dtype is not None: + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype] + else: + scale_dtype = torch.float32 + + group_shape = _group_shape( + schema.weight_scale_group_size, + schema.weight_scale_group_size_n, + ) + if schema.weight_scale_type == WeightScaleType.TENSOR: + group_shape = GroupShape.PER_TENSOR + + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) + + scale2 = None + if schema.weight_scale_2_type == WeightScale2Type.TENSOR: + scale2 = ScaleDesc( + dtype=torch.float32, + static=True, + group_shape=GroupShape.PER_TENSOR, + ) + + return QuantKey( + dtype=dtype, + scale=scale, + scale2=scale2, + symmetric=not schema.has_zero_point, + ) + + +def input_schema_to_quant_key( + schema: "BaseInputSchema", + param_dtype: torch.dtype | None = None, +) -> QuantKey | None: + from vllm.utils.humming import InputQuantizationMode + + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in [torch.float16, torch.bfloat16]: + param_dtype = torch.bfloat16 + if not current_platform.has_device_capability(80): + param_dtype = torch.float16 + + schema = schema.to_humming_schema(param_dtype) + if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: + return None + + mode = schema.input_quant_mode + if mode == InputQuantizationMode.Disabled: + return None + dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype] + + gs = schema.input_scale_group_size + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN + + # Pick the scale dtype the Humming kernel actually consumes. An explicit + # input_scale_dtype always wins. Otherwise infer from the grouping: MX + # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8) + # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a + # float32 scale. Getting this right lets a grouped FP8 activation match + # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key. + if schema.input_scale_dtype is not None: + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype] + elif gs == 16: + scale_dtype = FP8_DTYPE + elif gs == 32: + scale_dtype = MXFP_SCALE_DTYPE + else: + scale_dtype = torch.float32 + + if mode == InputQuantizationMode.StaticTensor: + group_shape = GroupShape.PER_TENSOR + scale = ScaleDesc( + dtype=scale_dtype, + static=mode == InputQuantizationMode.StaticTensor, + group_shape=group_shape, + ) + scale2 = None + has_scale2 = ( + InputQuantizationMode.StaticTensorDynamicGroup, + InputQuantizationMode.DynamicGroupToken, + ) + if mode in has_scale2: + static = mode == InputQuantizationMode.StaticTensorDynamicGroup + scale2 = ScaleDesc( + torch.float32, + static, + GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN, + ) + + return QuantKey(dtype=dtype, scale=scale, scale2=scale2, symmetric=True) + + +def quant_key_to_input_schema(key: QuantKey | None) -> "HummingInputSchema": + from vllm.utils.humming import HummingInputSchema, InputQuantizationMode + + if key is None: + return HummingInputSchema(input_quant_mode=InputQuantizationMode.Disabled) + + if not key.symmetric: + raise ValueError("Humming input quantization must be symmetric") + + quant_dtypes = {value: dtype for dtype, value in _HUMMING_TO_QUANT_DTYPE.items()} + quant_dtypes[FP4_DTYPE] = humming_dtypes.float4e2m1 + quant_dtypes[FP8_DTYPE] = humming_dtypes.float8e4m3 + scale_dtypes = {value: dtype for dtype, value in _HUMMING_TO_SCALE_DTYPE.items()} + if key.dtype not in quant_dtypes or key.scale.dtype not in scale_dtypes: + raise ValueError(f"Unsupported Humming input or scale dtype: {key}") + + scale = key.scale + group_size = 0 + if scale.static and scale.group_shape.is_per_tensor() and key.scale2 is None: + mode = InputQuantizationMode.StaticTensor + elif not scale.static and scale.group_shape.is_per_token() and key.scale2 is None: + mode = InputQuantizationMode.DynamicToken + elif not scale.static and scale.group_shape.is_per_group(): + group_size = scale.group_shape.col + scale2 = key.scale2 + if scale2 is None: + mode = InputQuantizationMode.DynamicGroup + elif scale2.dtype != torch.float32: + raise ValueError("Humming secondary input scales must be float32") + elif scale2.static and scale2.group_shape.is_per_tensor(): + mode = InputQuantizationMode.StaticTensorDynamicGroup + elif not scale2.static and scale2.group_shape.is_per_token(): + mode = InputQuantizationMode.DynamicGroupToken + else: + raise ValueError(f"Unsupported Humming secondary input scale: {scale2}") + else: + raise ValueError(f"Unsupported Humming input scale: {key}") + + if group_size == 0 and scale.dtype != torch.float32: + raise ValueError("Humming tensor and token input scales must be float32") + + return HummingInputSchema( + a_dtype=quant_dtypes[key.dtype], + input_scale_group_size=group_size, + input_scale_dtype=scale_dtypes[scale.dtype], + input_quant_mode=mode, + ) + + +def check_and_fallback_input_schema( + weight_schema: "BaseWeightSchema", + input_schema: "BaseInputSchema", + param_dtype: torch.dtype | None = None, + allow_fallback: bool = True, +) -> "HummingInputSchema": + from vllm.utils.humming import HummingInputSchema, InputQuantizationMode + + capability = current_platform.get_device_capability() + assert capability is not None + sm_version = capability.to_int() + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in (torch.float16, torch.bfloat16): + param_dtype = torch.bfloat16 if sm_version >= 80 else torch.float16 + assert param_dtype in (torch.float16, torch.bfloat16) + + dtype_fallback_order_map = { + humming_dtypes.float8e5m2: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float8e4m3: ( + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float8e3m4: ( + humming_dtypes.float8e4m3, + humming_dtypes.int8, + ), + humming_dtypes.float4e2m1: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float4e0m3: ( + humming_dtypes.int4, + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.int8: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + ), + humming_dtypes.int4: ( + humming_dtypes.int8, + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + ), + } + + weight_schema = weight_schema.to_humming_schema(param_dtype) + input_schema = input_schema.to_humming_schema(param_dtype) + fp6_dtypes = (humming_dtypes.float6e2m3, humming_dtypes.float6e3m2) + if input_schema.a_dtype in fp6_dtypes: + input_schema.a_dtype = humming_dtypes.float8e4m3 + a_dtype = input_schema.a_dtype + input_bits = a_dtype.num_bits if a_dtype is not None else 16 + weight_group_size = weight_schema.weight_scale_group_size + input_group_size = input_schema.input_scale_group_size + input_scale_dtype = input_schema.input_scale_dtype + + def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: + if dtype is None or dtype.num_bits == 16: + return False + is_int4_deprecated = dtype == humming_dtypes.int4 and sm_version >= 90 + is_int8_deprecated = dtype == humming_dtypes.int8 and 103 <= sm_version < 110 + # Humming has not yet implemented high-performance FP8 computation + # on SM10X and SM11X. Therefore, it is temporarily considered deprecated. + # TODO: removed this after Humming has high-performance FP8 computation + is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120 + return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated + + if input_schema.is_compatible_with(weight_schema, param_dtype): + if not allow_fallback: + if is_deprecated(a_dtype): + msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now." + logger.warning_once(msg) + return input_schema + + if a_dtype is None or a_dtype.num_bits == 16: + return input_schema + + if not is_deprecated(a_dtype): + is_mxfp8 = ( + a_dtype.is_floating_point_type + and a_dtype.num_bits == 8 + and input_group_size == 32 + and input_scale_dtype == humming_dtypes.float8e8m0 + and sm_version >= 120 + ) + is_groupwise = input_group_size > 0 or weight_group_size > 0 + if input_bits == 8 and is_groupwise and not is_mxfp8: + # prefer wna16 for groupwise weight or groupwise input + return HummingInputSchema() + + return input_schema + + if allow_fallback: + fallback_dtypes = ( + a_dtype, + *dtype_fallback_order_map.get(a_dtype, ()), + humming_dtypes.DataType.from_any(param_dtype), + ) + for dtype in fallback_dtypes: + if dtype is None or is_deprecated(dtype): + continue + + group_size = 0 + scale_dtype = humming_dtypes.float32 + quant_mode = InputQuantizationMode.DynamicToken + if dtype.num_bits == 16: + scale_dtype = None + quant_mode = InputQuantizationMode.Disabled + elif dtype in (humming_dtypes.float8e4m3, humming_dtypes.float8e3m4): + is_channelwise_weight = weight_schema.weight_scale_group_size == 0 + weight_group_size = weight_schema.weight_scale_group_size + is_e8m0_scale = weight_schema.bs_dtype == humming_dtypes.float8e8m0 + is_mx_weight = weight_group_size == 32 and is_e8m0_scale + if sm_version >= 120 and (is_channelwise_weight or is_mx_weight): + group_size = 32 + scale_dtype = humming_dtypes.float8e8m0 + quant_mode = InputQuantizationMode.DynamicGroup + + candidate = HummingInputSchema( + a_dtype=dtype, + input_scale_group_size=group_size, + input_scale_dtype=scale_dtype, + input_quant_mode=quant_mode, + ) + if candidate.is_compatible_with(weight_schema, param_dtype): + return candidate + + raise ValueError( + f"No compatible Humming input schema for {input_schema} with " + f"weight schema {weight_schema} on SM{sm_version} " + f"({allow_fallback=})" + ) + + +def resolve_humming_layer_config(config: dict[str, Any], prefix: str) -> dict[str, Any]: + layer_config = config.copy() + dynamic = config.get("dynamic", {}) + if isinstance(dynamic, dict): + for pattern, overrides in dynamic.items(): + if pattern.startswith("+") and re.match(pattern[2:], prefix): + layer_config.update(overrides) + break + return layer_config + + +def humming_is_layer_skipped(config: dict[str, Any], prefix: str): + if not config: + return True + + keys = ["ignored_layers", "ignore", "modules_to_not_convert"] + ignored_layers: list[str] = [] + for key in keys: + candidate = config.get(key, []) or [] + if candidate: + ignored_layers = candidate + break + + if any(module_name in prefix for module_name in ignored_layers): + return True + if "lm_head" in prefix: + return True + + for regex in config.get("dynamic", {}): + if regex[:1] != "-": + continue + if re.match(regex[2:], prefix): + return True + + return False diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py index 014603ff39c3..574a00f16b8a 100644 --- a/vllm/utils/humming.py +++ b/vllm/utils/humming.py @@ -4,12 +4,56 @@ vLLM code should import humming symbols from here so that ``import humming`` (which has import-time side effects) is deferred until first use. Add new -symbols by appending one entry to ``_EXPORTS`` as ``"module.path:attr"``, -or ``"module.path"`` for a whole-module re-export. +symbols to both the ``TYPE_CHECKING`` imports and ``_EXPORTS``, using +``"module.path:attr"`` or ``"module.path"`` for a whole-module re-export. """ import importlib -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from humming import dtypes as dtypes + from humming.config import GemmType as GemmType + from humming.config import InputQuantizationMode as InputQuantizationMode + from humming.config import LayerConfig as LayerConfig + from humming.config import MmaType as MmaType + from humming.config import WeightScale2Type as WeightScale2Type + from humming.config import WeightScaleType as WeightScaleType + from humming.dtypes import DataType as DataType + from humming.forward import humming_forward as humming_forward + from humming.forward import may_process_input as may_process_input + from humming.forward import may_quant_input as may_quant_input + from humming.ops import process_input as process_input + from humming.schema import AWQWeightSchema as AWQWeightSchema + from humming.schema import BaseInputSchema as BaseInputSchema + from humming.schema import BaseWeightSchema as BaseWeightSchema + from humming.schema import BitnetWeightSchema as BitnetWeightSchema + from humming.schema import ( + CompressedTensorsInputSchema as CompressedTensorsInputSchema, + ) + from humming.schema import ( + CompressedTensorsWeightSchema as CompressedTensorsWeightSchema, + ) + from humming.schema import Fp8InputSchema as Fp8InputSchema + from humming.schema import GptOssMxfp4WeightSchema as GptOssMxfp4WeightSchema + from humming.schema import GPTQWeightSchema as GPTQWeightSchema + from humming.schema import HummingInputSchema as HummingInputSchema + from humming.schema import HummingWeightSchema as HummingWeightSchema + from humming.schema import Mxfp4WeightSchema as Mxfp4WeightSchema + from humming.schema.fp8 import Fp8WeightSchema as Fp8WeightSchema + from humming.schema.modelopt import ( + ModeloptMxfp8WeightSchema as ModeloptMxfp8WeightSchema, + ) + from humming.schema.modelopt import ( + ModeloptNvfp4InputSchema as ModeloptNvfp4InputSchema, + ) + from humming.schema.modelopt import ( + ModeloptNvfp4WeightSchema as ModeloptNvfp4WeightSchema, + ) + from humming.transform import prepare_layer_config as prepare_layer_config + from humming.transform import transform_humming_tensors as transform_humming_tensors + from humming.tune import get_heuristics_config as get_heuristics_config + from humming.utils.weight import quantize_weight as quantize_weight _EXPORTS: dict[str, str] = { "dtypes": "humming.dtypes", @@ -17,7 +61,12 @@ "GemmType": "humming.config:GemmType", "LayerConfig": "humming.config:LayerConfig", "WeightScaleType": "humming.config:WeightScaleType", + "WeightScale2Type": "humming.config:WeightScale2Type", + "InputQuantizationMode": "humming.config:InputQuantizationMode", + "MmaType": "humming.config:MmaType", "humming_forward": "humming.forward:humming_forward", + "may_process_input": "humming.forward:may_process_input", + "process_input": "humming.ops:process_input", "may_quant_input": "humming.forward:may_quant_input", "prepare_layer_config": "humming.transform:prepare_layer_config", "transform_humming_tensors": "humming.transform:transform_humming_tensors", From 49ee12d742b4a5524d2a3a3ee6c3fb77e4913073 Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:50:34 +0300 Subject: [PATCH 544/774] [Misc] Rename --enable-mamba-fine-grained-prefix-cache (#53945 follow-up) (#57382) Signed-off-by: Roi Koren Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/features/automatic_prefix_caching.md | 4 ++-- .../test_mamba_eagle_resume_checkpoint.py | 18 +++++++++--------- .../test_partial_prefix_cache_hits.py | 8 ++++---- tests/v1/core/test_mamba_align_chunk_split.py | 2 +- tests/v1/core/test_prefix_caching.py | 2 +- vllm/config/cache.py | 4 ++-- vllm/engine/arg_utils.py | 12 ++++++------ vllm/v1/core/kv_cache_manager.py | 10 +++++----- vllm/v1/core/sched/scheduler.py | 10 +++++----- vllm/v1/core/single_type_kv_cache_manager.py | 6 +++--- 10 files changed, 38 insertions(+), 38 deletions(-) diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md index 03cca3b38d6c..cfc8c19acb2b 100644 --- a/docs/features/automatic_prefix_caching.md +++ b/docs/features/automatic_prefix_caching.md @@ -15,7 +15,7 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl ## Hybrid Mamba models -Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-fine-grained-prefix-cache` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. +Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold: @@ -28,7 +28,7 @@ This helps when many requests share a long system prompt and then diverge. It is vllm serve \ --mamba-cache-mode align \ --prefix-match-unit 64 \ - --enable-mamba-fine-grained-prefix-cache + --enable-mamba-shared-prefix-checkpoint ``` `--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect. diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 059b4005f0db..4cb9afe87138 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -6,7 +6,7 @@ candidate and drops it. Mamba materializes state only on its own block grid, so the pruned position holds nothing and the hit floors back a whole Mamba block. -Gated by ``CacheConfig.enable_mamba_fine_grained_prefix_cache``. Every test here +Gated by ``CacheConfig.enable_mamba_shared_prefix_checkpoint``. Every test here drives the real ``Scheduler._mamba_block_aligned_split`` and the real ``KVCacheManager``; no chunk boundary is hard-coded. """ @@ -32,7 +32,7 @@ def _manager( block_size, hash_block_size, *, - fine_grained=True, + shared_prefix_checkpoint=True, num_blocks=8192, eagle_group=None, num_prefill_lookahead=0, @@ -52,7 +52,7 @@ def _manager( hash_block_size=hash_block_size, use_eagle=True, num_prefill_lookahead=num_prefill_lookahead, - enable_mamba_fine_grained_prefix_cache=fine_grained, + enable_mamba_shared_prefix_checkpoint=shared_prefix_checkpoint, ) @@ -76,8 +76,8 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True): hash_block_size=hash_block_size, mamba_has_prefill_checkpoint_blocks=False, # forced False under eagle mamba_partial_cache_hit=partial_hit, - mamba_fine_grained_prefix_cache=( - partial_hit and manager.mamba_fine_grained_prefix_cache + mamba_shared_prefix_checkpoint=( + partial_hit and manager.mamba_shared_prefix_checkpoint ), ) @@ -151,7 +151,7 @@ def test_sibling_resumes_from_the_observed_junction(): block_size, hash_block_size = 512, 32 manager = _manager(block_size, hash_block_size) stub = _stub(manager, block_size, hash_block_size) - assert stub.mamba_fine_grained_prefix_cache, "the feature must be armed" + assert stub.mamba_shared_prefix_checkpoint, "the feature must be armed" _orphaned_full_attention_tail(manager, stub, 2020) consumer = make_request( @@ -325,11 +325,11 @@ def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead): """ block_size, hash_block_size = 512, 32 - def sibling_hit(fine_grained): + def sibling_hit(shared_prefix_checkpoint): manager = _manager( block_size, hash_block_size, - fine_grained=fine_grained, + shared_prefix_checkpoint=shared_prefix_checkpoint, eagle_group=eagle_group, num_prefill_lookahead=num_prefill_lookahead, ) @@ -348,7 +348,7 @@ def sibling_hit(fine_grained): def test_disabled_by_default(): """Off, the junction is block-floored and no resume point is registered.""" block_size, hash_block_size = 512, 32 - manager = _manager(block_size, hash_block_size, fine_grained=False) + manager = _manager(block_size, hash_block_size, shared_prefix_checkpoint=False) stub = _stub(manager, block_size, hash_block_size) _orphaned_full_attention_tail(manager, stub, 2020) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 393ec484651e..f1301d928837 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -188,7 +188,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): dcp_world_size=dcp_world_size, scheduler_block_size=scheduler_block_size, mamba_partial_cache_hit=True, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, mamba_has_prefill_checkpoint_blocks=False, ) split = Scheduler._mamba_block_aligned_split @@ -236,7 +236,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget(): use_eagle_block_drop=False, hash_block_size=32, mamba_partial_cache_hit=False, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, mamba_has_prefill_checkpoint_blocks=False, ) req = make_request("0", [0] * prompt_length, 32, sha256) @@ -277,7 +277,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold(): use_eagle_block_drop=False, hash_block_size=32, mamba_partial_cache_hit=False, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, mamba_has_prefill_checkpoint_blocks=False, ) req = make_request("0", [0] * prompt_length, 32, sha256) @@ -1826,7 +1826,7 @@ def test_mamba_align_split_stops_below_eagle_proof_boundary(): use_eagle_block_drop=True, hash_block_size=hash_block_size, mamba_partial_cache_hit=True, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, mamba_has_prefill_checkpoint_blocks=False, ) split = Scheduler._mamba_block_aligned_split diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py index 1bd7f13b5935..b1575239158d 100644 --- a/tests/v1/core/test_mamba_align_chunk_split.py +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -144,7 +144,7 @@ def _split( scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), # `prefix_match_unit` finer than the block size (#46384). mamba_partial_cache_hit=partial_hit, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, hash_block_size=ATTN_BLOCK_SIZE, mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0), mamba_prefill_checkpoint_alignment=( diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index d6848a6a7f3d..0a34a6131dc3 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -2269,7 +2269,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): use_eagle_block_drop=False, hash_block_size=block_size, mamba_partial_cache_hit=False, - mamba_fine_grained_prefix_cache=False, + mamba_shared_prefix_checkpoint=False, mamba_has_prefill_checkpoint_blocks=False, ) req_2.shared_prefix_boundary = shared_prefix_boundary diff --git a/vllm/config/cache.py b/vllm/config/cache.py index c12433a367a8..6315a1e31763 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -184,7 +184,7 @@ class CacheConfig: when the token is at position i * block_size. This is the default when prefix caching is enabled. """ - enable_mamba_fine_grained_prefix_cache: bool = False + enable_mamba_shared_prefix_checkpoint: bool = False """Also register a Mamba "align" checkpoint at the shared-prefix junction -- where an EAGLE/MTP sibling was observed to resume -- instead of only at the prompt tail. Off by default; only takes effect with `mamba_cache_mode` @@ -274,7 +274,7 @@ def compute_hash(self) -> str: "prefix_cache_retention_interval", # Prefix-caching implementation detail (doesn't affect compiled graph). "prefix_match_unit", - "enable_mamba_fine_grained_prefix_cache", + "enable_mamba_shared_prefix_checkpoint", "mamba_page_size_padded", "skip_page_size_padded", "user_specified_block_size", diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 67d8a2f2a0ba..0d33caf625a2 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -753,8 +753,8 @@ class EngineArgs: mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode - enable_mamba_fine_grained_prefix_cache: bool = ( - CacheConfig.enable_mamba_fine_grained_prefix_cache + enable_mamba_shared_prefix_checkpoint: bool = ( + CacheConfig.enable_mamba_shared_prefix_checkpoint ) replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len use_replayssm: bool = CacheConfig.use_replayssm @@ -1330,8 +1330,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) cache_group.add_argument( - "--enable-mamba-fine-grained-prefix-cache", - **cache_kwargs["enable_mamba_fine_grained_prefix_cache"], + "--enable-mamba-shared-prefix-checkpoint", + **cache_kwargs["enable_mamba_shared_prefix_checkpoint"], ) cache_group.add_argument( "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"] @@ -2128,8 +2128,8 @@ def create_engine_config( mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, - enable_mamba_fine_grained_prefix_cache=( - self.enable_mamba_fine_grained_prefix_cache + enable_mamba_shared_prefix_checkpoint=( + self.enable_mamba_shared_prefix_checkpoint ), replayssm_buffer_len=self.replayssm_buffer_len, use_replayssm=self.use_replayssm, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 1adf386cab2e..f881fc981828 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -145,7 +145,7 @@ def __init__( pcp_world_size: int = 1, metrics_collector: KVCacheMetricsCollector | None = None, watermark: float = 0.0, - enable_mamba_fine_grained_prefix_cache: bool = False, + enable_mamba_shared_prefix_checkpoint: bool = False, ) -> None: self.max_model_len = max_model_len # When unset, fall back to `max_model_len` so the recycling-aware cap @@ -183,16 +183,16 @@ def __init__( # junction costs a forward pass and displaces the block-boundary stop. # Multi-module MTP is excluded because ``cache_blocks`` then hands the # manager ``num_computed - num_reprefillable`` rather than the chunk end. - self.mamba_fine_grained_prefix_cache = ( - enable_mamba_fine_grained_prefix_cache + self.mamba_shared_prefix_checkpoint = ( + enable_mamba_shared_prefix_checkpoint and bool(self.coordinator.eagle_group_ids) and self.coordinator.enable_partial_hash_hits and self.coordinator.num_reprefillable_tokens == 0 ) - if self.mamba_fine_grained_prefix_cache: + if self.mamba_shared_prefix_checkpoint: for manager in self.coordinator.single_type_managers: if isinstance(manager, MambaManager): - manager.fine_grained_prefix_cache = True + manager.shared_prefix_checkpoint = True self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.retained_hit_group_ids = tuple( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 19ac16284cd7..a2638942f880 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -320,8 +320,8 @@ def __init__( hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, watermark=self.scheduler_config.watermark, - enable_mamba_fine_grained_prefix_cache=( - self.cache_config.enable_mamba_fine_grained_prefix_cache + enable_mamba_shared_prefix_checkpoint=( + self.cache_config.enable_mamba_shared_prefix_checkpoint ), ) # Bind after construction so connectors can access the cache manager. @@ -376,9 +376,9 @@ def __init__( # manager decides whether it can check-point there (per-group eagle bit, # no MTP re-prefill tail); splitting for a stop it would refuse costs a # forward pass and displaces the block-boundary stop. - self.mamba_fine_grained_prefix_cache = ( + self.mamba_shared_prefix_checkpoint = ( self.mamba_partial_cache_hit - and self.kv_cache_manager.mamba_fine_grained_prefix_cache + and self.kv_cache_manager.mamba_shared_prefix_checkpoint ) # Counts of non-empty steps scheduled / processed. update_from_output @@ -509,7 +509,7 @@ def _mamba_block_aligned_split( # output tokens can still observe a junction there. junction_stop = ( junction - if self.mamba_fine_grained_prefix_cache + if self.mamba_shared_prefix_checkpoint and junction <= request.num_prompt_tokens else block_floored ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 9fe59dfdbe4d..70d56c6543bd 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -143,9 +143,9 @@ def __init__( # aligned segment (SWA). Initialized lazily by the coordinator after # determining the attention groups. self.use_eagle = False - # ``CacheConfig.enable_mamba_fine_grained_prefix_cache``, narrowed and set + # ``CacheConfig.enable_mamba_shared_prefix_checkpoint``, narrowed and set # by ``KVCacheManager``; only an EAGLE Mamba "align" group ever gets it. - self.fine_grained_prefix_cache = False + self.shared_prefix_checkpoint = False # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained # managers (full attention, mamba "align"); harmlessly empty elsewhere. self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {} @@ -2090,7 +2090,7 @@ def _cache_partial_tail_block( # running state block, mutated in place, which equals what its key # promises only after that step's forward. if num_tokens != latest_prompt_hash_boundary and not ( - self.fine_grained_prefix_cache + self.shared_prefix_checkpoint and num_tokens == request.shared_prefix_boundary and request.num_computed_tokens < num_tokens <= request.num_prompt_tokens ): From 4868312128172a424a7cf4c90f25c53b49186346 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sun, 20 Sep 2026 13:39:23 -0400 Subject: [PATCH 545/774] [Refactor] Remove dead kernel code (#57621) Signed-off-by: yewentao256 --- csrc/cpu/sgl-kernels/common.h | 16 -- csrc/libtorch_stable/cuda_vec_utils.cuh | 4 - csrc/libtorch_stable/moe/moe_wna16_utils.h | 8 - .../quantization/gptq/matrix_view.cuh | 33 --- .../quantization/marlin/marlin.cuh | 23 +- csrc/libtorch_stable/sampler.cu | 5 - .../quantization/w8a8/fp8/amd/quant_utils.cuh | 242 +----------------- .../w8a8/fp8/nvidia/quant_utils.cuh | 13 - csrc/rocm/attention.cu | 33 --- csrc/rocm/moe_q_gemm_rdna3.cu | 15 -- csrc/rocm/q_gemm_rdna3.cu | 60 +---- csrc/rocm/qdq_4_rdna3.cuh | 111 -------- 12 files changed, 4 insertions(+), 559 deletions(-) diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index de1260bb3321..dc0769450e36 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -280,26 +280,10 @@ inline int get_thread_num() { // balance payload across each thread template inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) { -#if 0 - // onednn partition pattern - T& n_my = n_end; - if (nth <= 1 || n == 0) { - n_start = 0; - n_my = n; - } else { - T n1 = div_up(n, nth); - T n2 = n1 - 1; - T T1 = n - n2 * nth; - n_my = ith < T1 ? n1 : n2; - n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2; - } - n_end += n_start; -#else // pytorch aten partition pattern T n_my = div_up(n, nth); n_start = ith * n_my; n_end = std::min(n_start + n_my, n); -#endif } template diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/libtorch_stable/cuda_vec_utils.cuh index ec6e60724e65..c962c99ec255 100644 --- a/csrc/libtorch_stable/cuda_vec_utils.cuh +++ b/csrc/libtorch_stable/cuda_vec_utils.cuh @@ -320,10 +320,6 @@ __host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) { return (reinterpret_cast(ptr) & 15) == 0; } -__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) { - return (reinterpret_cast(ptr) & 31) == 0; -} - // ============================================================ // Packed type conversion and arithmetic // ============================================================ diff --git a/csrc/libtorch_stable/moe/moe_wna16_utils.h b/csrc/libtorch_stable/moe/moe_wna16_utils.h index 8ef03f0e6052..6ce52ee5b1d8 100644 --- a/csrc/libtorch_stable/moe/moe_wna16_utils.h +++ b/csrc/libtorch_stable/moe/moe_wna16_utils.h @@ -34,10 +34,6 @@ class ScalarType { static __host__ __device__ float2 inline num22float2(const half2 x) { return __half22float2(x); } - - static __host__ __device__ half2 inline float22num2(const float2 x) { - return __float22half2_rn(x); - } }; template <> @@ -71,10 +67,6 @@ class ScalarType { static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) { return __bfloat1622float2(x); } - - static __host__ __device__ nv_bfloat162 inline float22num2(const float2 x) { - return __float22bfloat162_rn(x); - } #endif }; diff --git a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh b/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh index 3d82e4393166..abe3368cb16f 100644 --- a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh +++ b/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh @@ -74,18 +74,9 @@ class MatrixView_half_rw { const int width) : data(data), height(height), width(width) {} - __device__ __forceinline__ half item(int row, int column) const { - return data[row * width + column]; - } __device__ __forceinline__ half* item_ptr(int row, int column) { return &data[row * width + column]; } - __device__ __forceinline__ void set(int row, int column, half value) { - data[row * width + column] = value; - } - __device__ __forceinline__ void set_half2(int row, int column, half2 value) { - ((half2*)data)[(row * width + column) / 2] = value; - } __device__ __forceinline__ void set4(int row, int column, half v0, half v1, half v2, half v3) { @@ -113,14 +104,6 @@ class MatrixView_q4_row { return (data[row * width / 8 + column / 8] >> shift) & 0x0f; } - __device__ __forceinline__ void item2(int (&items)[2], int row, - int column) const { - int shift = (column & 0x07) * 4; - uint32_t d = data[row * width / 8 + column / 8] >> shift; - items[0] = d & 0x0f; - items[1] = (d >> 4) & 0x0f; - } - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { int shift = (column & 0x07) * 4; @@ -148,14 +131,6 @@ class MatrixView_q2_row { return (data[row * width / 16 + column / 16] >> shift) & 0x03; } - __device__ __forceinline__ void item2(int (&items)[2], int row, - int column) const { - int shift = (column & 0x0f) * 2; - uint32_t d = data[row * width / 16 + column / 16] >> shift; - items[0] = d & 0x03; - items[1] = (d >> 2) & 0x03; - } - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { int shift = (column & 0x0f) * 2; @@ -237,14 +212,6 @@ class MatrixView_q8_row { return (data[row * width / 4 + column / 4] >> shift) & 0xff; } - __device__ __forceinline__ void item2(int (&items)[2], int row, - int column) const { - int shift = (column & 0x03) * 8; - uint32_t d = data[row * width / 4 + column / 4] >> shift; - items[0] = d & 0xff; - items[1] = (d >> 8) & 0xff; - } - __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const { int shift = (column & 0x03) * 2; diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/csrc/libtorch_stable/quantization/marlin/marlin.cuh index bfb65e874b3d..af4d4730ba3a 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin.cuh +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cuh @@ -67,14 +67,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr, } } -__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = - reinterpret_cast(glob_ptr)[0]; - } -} - __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, bool pred = true) { if (pred) { @@ -121,19 +113,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr, "r"(smem), "l"(glob_ptr), "n"(BYTES)); } -__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, bool pred = true) { const int BYTES = 16; @@ -170,4 +149,4 @@ __device__ inline void cp_async_wait() { } // namespace MARLIN_NAMESPACE_NAME -#endif \ No newline at end of file +#endif diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index b8d64f7af130..703a748ce686 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -45,11 +45,6 @@ __global__ void apply_repetition_penalties_kernel( } } -__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; -} - template static inline __device__ uint32_t extractBinIdx(float x) { if constexpr (step == 0) { diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 58d19de73494..32925c0c3c95 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -67,104 +67,6 @@ using fp8_type = __hip_fp8_e4m3_fnuz; using fp8x2_type = __hip_fp8x2_e4m3_fnuz; #endif -// fp8 -> half -template <> -__inline__ __device__ uint16_t -vec_conversion(const uint8_t& a) { - return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x; -} - -// fp8x2 -> half2 -template <> -__inline__ __device__ uint32_t -vec_conversion(const uint16_t& a) { - union { - __half2_raw h2r; - uint32_t ui32; - } tmp; - tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret); - return tmp.ui32; -} - -// fp8x4 -> half2x2 -template <> -__inline__ __device__ uint2 vec_conversion(const uint32_t& a) { - union { - uint2 u32x2; - uint32_t u32[2]; - } tmp; - tmp.u32[0] = vec_conversion((uint16_t)a); - tmp.u32[1] = vec_conversion((uint16_t)(a >> 16U)); - return tmp.u32x2; -} - -// fp8x8 -> half2x4 -template <> -__inline__ __device__ uint4 vec_conversion(const uint2& a) { - union { - uint4 u64x2; - uint2 u64[2]; - } tmp; - tmp.u64[0] = vec_conversion(a.x); - tmp.u64[1] = vec_conversion(a.y); - return tmp.u64x2; -} - -using __nv_bfloat16 = __hip_bfloat16; - -// fp8 -> __nv_bfloat16 -template <> -__inline__ __device__ __nv_bfloat16 -vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) { - fp8_type f8; - f8.__x = a; - return __float2bfloat16(static_cast(f8)); -} - -using __nv_bfloat162 = __hip_bfloat162; - -// fp8x2 -> __nv_bfloat162 -template <> -__inline__ __device__ __nv_bfloat162 -vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) { - __nv_bfloat162 res; - res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a); - res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U)); - return res; -} - -// fp8x4 -> bf16_4_t -template <> -__inline__ __device__ bf16_4_t -vec_conversion(const uint32_t& a) { - bf16_4_t res; - res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a); - res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U)); - return res; -} - -// fp8x8 -> bf16_8_t -template <> -__inline__ __device__ bf16_8_t vec_conversion(const uint2& a) { - bf16_4_t tmp1, tmp2; - tmp1 = vec_conversion(a.x); - tmp2 = vec_conversion(a.y); - bf16_8_t res; - res.x = tmp1.x; - res.y = tmp1.y; - res.z = tmp2.x; - res.w = tmp2.y; - return res; -} - -// fp8 -> float -template <> -__inline__ __device__ float vec_conversion(const uint8_t& a) { - fp8_type f8; - f8.__x = a; - return static_cast(f8); -} - // fp8x2 -> float2 template <> __inline__ __device__ float2 @@ -193,138 +95,6 @@ vec_conversion(const uint32_t& a) { return res; } -// fp8x8 -> float8 -template <> -__inline__ __device__ Float8_ vec_conversion(const uint2& a) { - Float4_ tmp1, tmp2; - tmp1 = vec_conversion(a.x); - tmp2 = vec_conversion(a.y); - Float8_ res; - res.x = tmp1.x; - res.y = tmp1.y; - res.z = tmp2.x; - res.w = tmp2.y; - return res; -} - -// half -> fp8 -template <> -__inline__ __device__ uint8_t -vec_conversion(const uint16_t& a) { - __half_raw tmp; - tmp.x = a; - return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation, - fp8_type::__default_interpret); -} - -template <> -__inline__ __device__ uint16_t -vec_conversion(const uint32_t& a) { - union { - uint32_t ui32; - __half2_raw h2r; - } tmp; - tmp.ui32 = a; - return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation, - fp8_type::__default_interpret); -} - -// bf16 -> fp8 -template <> -__inline__ __device__ uint8_t -vec_conversion(const __nv_bfloat16& a) { - return __hip_cvt_float_to_fp8(__bfloat162float(a), - fp8_type::__default_saturation, - fp8_type::__default_interpret); -} - -// float -> fp8 -template <> -__inline__ __device__ uint8_t vec_conversion(const float& a) { - return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation, - fp8_type::__default_interpret); -} - -// float2 -> half2 -template <> -__inline__ __device__ uint32_t -vec_conversion(const float2& a) { - union { - half2 float16; - uint32_t uint32; - }; - - float16 = __float22half2_rn(a); - return uint32; -} - -// Float4 -> half2x2 -template <> -__inline__ __device__ uint2 vec_conversion(const Float4_& a) { - uint2 b; - float2 val; - val.x = a.x.x; - val.y = a.x.y; - b.x = vec_conversion(val); - - val.x = a.y.x; - val.y = a.y.y; - b.y = vec_conversion(val); - return b; -} - -// Float4 -> float4 -template <> -__inline__ __device__ float4 vec_conversion(const Float4_& a) { - float4 b; - b.x = a.x.x; - b.y = a.x.y; - b.z = a.y.x; - b.w = a.y.y; - return b; -} - -// Float8 -> half2x4 -template <> -__inline__ __device__ uint4 vec_conversion(const Float8_& a) { - uint4 b; - b.x = vec_conversion(a.x); - b.y = vec_conversion(a.y); - b.z = vec_conversion(a.z); - b.w = vec_conversion(a.w); - return b; -} - -// float2 -> bfloat162 -template <> -__inline__ __device__ __nv_bfloat162 -vec_conversion<__nv_bfloat162, float2>(const float2& a) { - __nv_bfloat162 b = __float22bfloat162_rn(a); - return b; -} - -// Float4 -> bfloat162x2 -template <> -__inline__ __device__ bf16_4_t -vec_conversion(const Float4_& a) { - bf16_4_t b; - b.x = __float22bfloat162_rn(a.x); - b.y = __float22bfloat162_rn(a.y); - return b; -} - -// Float8 -> bfloat162x4 -template <> -__inline__ __device__ bf16_8_t -vec_conversion(const Float8_& a) { - bf16_8_t b; - b.x = __float22bfloat162_rn(a.x); - b.y = __float22bfloat162_rn(a.y); - b.z = __float22bfloat162_rn(a.z); - b.w = __float22bfloat162_rn(a.w); - return b; -} - /* Scaled and vectorized conversions, for data exchange between high and low precision domains @@ -335,6 +105,7 @@ vec_conversion(const Float8_& a) { */ using __nv_bfloat16 = __hip_bfloat16; +using __nv_bfloat162 = __hip_bfloat162; // fp8 -> __nv_bfloat16 template <> @@ -613,17 +384,6 @@ scaled_vec_conversion(const float4& a, float scale) { } #endif // ENABLE_FP8 -template -__inline__ __device__ Tout convert(const Tin& x) { - #ifdef ENABLE_FP8 - if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { - return vec_conversion(x); - } - #endif - assert(false); - return {}; // Squash missing return statement warning -} - template __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { #ifdef ENABLE_FP8 diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 11a9b0818dc6..7fcf541fd64d 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -282,19 +282,6 @@ __inline__ __device__ float4 scaled_vec_conversion( } #endif // ENABLE_FP8 -template -__inline__ __device__ Tout convert(const Tin& x) { - #if 0 // Disable the following code to reduce the binary size. - if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { - return vec_conversion(x, __NV_E4M3); - } else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) { - return vec_conversion(x, __NV_E5M2); - } - #endif - assert(false); - __builtin_unreachable(); // Suppress missing return statement warning -} - template __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { #ifdef ENABLE_FP8 diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index b747b870e23f..2a11653a58d7 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -342,22 +342,6 @@ typedef union u64_cvt { int64_t i64; } _T8x8; -__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input, - _T8x8& Mtemp) { - _T8x8 Qtmp8x8; - - for (int i = 0; i < 2; i++) { - floatx4 q_out = {0, 0, 0, 0}; - q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i], - q_out); - Qtmp8x8.b16x4[i * 2] = - __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false); - Qtmp8x8.b16x4[i * 2 + 1] = - __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false); - } - return Qtmp8x8.b8x8; -} - __device__ float warpReduceMax(float val) { for (int offset = warpSize / 2; offset > 0; offset /= 2) { val = max( @@ -2583,23 +2567,6 @@ __device__ __forceinline__ float to_float(const T& inp) { } } -template -__device__ __forceinline__ float to_float_b16(const bit16_t& inp) { - union tmpcvt { - bit16_t u; - _Float16 f; - __hip_bfloat16 b; - } t16; - t16.u = inp; - if constexpr (std::is_same::value) { - return (float)t16.f; - } else if constexpr (std::is_same::value) { - return __bfloat162float(t16.b); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - template __device__ __forceinline__ T from_float(const float& inp) { if constexpr (std::is_same::value) { diff --git a/csrc/rocm/moe_q_gemm_rdna3.cu b/csrc/rocm/moe_q_gemm_rdna3.cu index 6c25ed7e4bc0..61ffbb9987e9 100644 --- a/csrc/rocm/moe_q_gemm_rdna3.cu +++ b/csrc/rocm/moe_q_gemm_rdna3.cu @@ -67,21 +67,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { return result; } -__forceinline__ __device__ float dot22_8_f(float (&dq)[8], - const bf16_t* a_ptr) { - float result = 0.0f; - #pragma unroll - for (int i = 0; i < 4; i++) { - uint32_t aw; - __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); - float a_x = __uint_as_float((aw & 0xFFFFu) << 16); - float a_y = __uint_as_float(aw & 0xFFFF0000u); - result = __fmaf_rn(dq[2 * i + 0], a_x, result); - result = __fmaf_rn(dq[2 * i + 1], a_y, result); - } - return result; -} - __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, half2 v23) { unsigned long long* addr_u = reinterpret_cast(addr); diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index 104af66d1cc5..bf2368269964 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -115,61 +115,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { return result; } -__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4], - const bf16_t* a_ptr) { - // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in - // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc - // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised - // fallback (single-element FMAs or fp32 round-trips), which empirically - // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16 - // decode path was paying that tax in full, scaling linearly with M (the - // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate - // and the kernel becomes memory-bound). - // - // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs) - // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA - // count, but each FMA is fast. Bonus: the accumulator is now fp32 - // throughout instead of bf16, which is also numerically more accurate - // (no compounding bf16-rounding inside the dot loop). - float result = 0.0f; - #pragma unroll - for (int i = 0; i < 4; i++) { - uint32_t aw, dw; - __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); - __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t)); - // bf16 in low 16 bits → fp32 by left-shifting into the upper half. - // bf16 in high 16 bits → already aligned with fp32's upper half. - float a_x = __uint_as_float((aw & 0xFFFFu) << 16); - float a_y = __uint_as_float(aw & 0xFFFF0000u); - float d_x = __uint_as_float((dw & 0xFFFFu) << 16); - float d_y = __uint_as_float(dw & 0xFFFF0000u); - result = __fmaf_rn(d_x, a_x, result); - result = __fmaf_rn(d_y, a_y, result); - } - return result; -} - -// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already -// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t -// overload above does for dq (still need to widen A from bf16). Wins more -// at high N: the bf162_t version's per-call widening cost scales with the -// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only -// for A widening (M_COUNT × 4 × 4 widens, half as many). -__forceinline__ __device__ float dot22_8_f(float (&dq)[8], - const bf16_t* a_ptr) { - float result = 0.0f; - #pragma unroll - for (int i = 0; i < 4; i++) { - uint32_t aw; - __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); - float a_x = __uint_as_float((aw & 0xFFFFu) << 16); - float a_y = __uint_as_float(aw & 0xFFFF0000u); - result = __fmaf_rn(dq[2 * i + 0], a_x, result); - result = __fmaf_rn(dq[2 * i + 1], a_y, result); - } - return result; -} - // --------------------------------------------------------------------------- // Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS). // RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those @@ -327,9 +272,8 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, // Per-column dequant constants. We hold one set of (z, y) pairs per column. // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper- - // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant - // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA - // bypass for the missing v_pk_fma_bf16 on gfx11. + // nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2 + // bias correction. half2 z1z16_h[4][2], y1y16_h[4][2]; float z_b_f[4], y_b_f[4]; diff --git a/csrc/rocm/qdq_4_rdna3.cuh b/csrc/rocm/qdq_4_rdna3.cuh index f4be668c1f08..43a7b33ba153 100644 --- a/csrc/rocm/qdq_4_rdna3.cuh +++ b/csrc/rocm/qdq_4_rdna3.cuh @@ -113,63 +113,6 @@ __forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4], // bf16 path // --------------------------------------------------------------------------- -// Bit-trick magic for bf16: -// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0). -// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because -// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So -// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi). -// -// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble -// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16] -// via a single 4/8/12-bit right-shift before the OR. That costs one extra -// shift per pair vs fp16, but keeps the FMA structure identical. -__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero, - bf16_t scale, - bf162_t& z_prep, - bf162_t& y_prep) { - // z = scale * -(128 + zero); y = scale. - float scale_f = __bfloat162float(scale); - float zf = -(128.0f + (float)zero) * scale_f; - bf16_t zb = __float2bfloat16(zf); - z_prep = __bfloat162bfloat162(zb); - y_prep = __bfloat162bfloat162(scale); -} - -__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa, - bf162_t (&dq)[4], - bf162_t z_prep, - bf162_t y_prep) { - const uint32_t c0 = 0x43004300; - - union { - uint32_t u; - bf162_t b2; - } q0, q1, q2, q3; - q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1]) - q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3]) - q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5]) - q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7]) - - // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale - dq[0] = __hfma2(q0.b2, y_prep, z_prep); - dq[1] = __hfma2(q1.b2, y_prep, z_prep); - dq[2] = __hfma2(q2.b2, y_prep, z_prep); - dq[3] = __hfma2(q3.b2, y_prep, z_prep); -} - -// --------------------------------------------------------------------------- -// bf16-input → fp32-output dequant (RDNA3 scalar path). -// -// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow -// fallback. Rather than computing dq in bf16 and widening at FMA time in -// the dot product, we widen to fp32 here once (a free left-shift by 16) and -// emit the (q - zero) * scale FMA directly in fp32. This: -// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32. -// * Eliminates 4× bf16→fp32 widens that the dot product would do. -// * Keeps the dot product accumulator in fp32 without a roundtrip. -// -// Output: fp32 dq[8], one element per K position (consumed by the -// fp32-overload of dot22_8_f in q_gemm_rdna3.cu). __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, bf16_t scale, float& z_prep, @@ -179,60 +122,6 @@ __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, y_prep = scale_f; } -// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32 -// values 128+nibble, without folding scale/zero. The caller folds scale/zb -// into the accumulator outside the inner loop using a precomputed sum_a, -// which saves ~27% of the FMA count vs the per-col-dequant approach above -// (only beneficial at M_COUNT=1; break-even at M_COUNT=2). -// -// Cost: 0 FMAs (pure bit-trick + as_float reinterprets). -__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa, - float (&q_f32)[8]) { - const uint32_t c0 = 0x43004300; - const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; - const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; - const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; - const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; - q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16); - q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u); - q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16); - q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u); - q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16); - q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u); - q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16); - q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u); -} - -__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa, - float (&dq)[8], - float z_prep, - float y_prep) { - const uint32_t c0 = 0x43004300; - const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; - const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; - const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; - const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; - // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16 - // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved). - const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16); - const float q0y = __uint_as_float(q0 & 0xFFFF0000u); - const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16); - const float q1y = __uint_as_float(q1 & 0xFFFF0000u); - const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16); - const float q2y = __uint_as_float(q2 & 0xFFFF0000u); - const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16); - const float q3y = __uint_as_float(q3 & 0xFFFF0000u); - // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale - dq[0] = __fmaf_rn(q0x, y_prep, z_prep); - dq[1] = __fmaf_rn(q0y, y_prep, z_prep); - dq[2] = __fmaf_rn(q1x, y_prep, z_prep); - dq[3] = __fmaf_rn(q1y, y_prep, z_prep); - dq[4] = __fmaf_rn(q2x, y_prep, z_prep); - dq[5] = __fmaf_rn(q2y, y_prep, z_prep); - dq[6] = __fmaf_rn(q3x, y_prep, z_prep); - dq[7] = __fmaf_rn(q3y, y_prep, z_prep); -} - } // namespace gptq_rdna3 } // namespace vllm From 01f1f58f10c616850a5f5ade129a7b58d62075fb Mon Sep 17 00:00:00 2001 From: Sheral Kumar Date: Sun, 20 Sep 2026 16:20:39 -0400 Subject: [PATCH 546/774] [ROCm][CI] Query HIP device memory for test GPU teardown waits. (#57450) Signed-off-by: Sheral Kumar Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Codex Co-authored-by: Andreas Karatzas --- tests/compile/passes/conftest.py | 22 +++++++++ tests/conftest.py | 7 +-- tests/kernels/moe/conftest.py | 2 + tests/utils.py | 78 ++++++++++++-------------------- 4 files changed, 56 insertions(+), 53 deletions(-) create mode 100644 tests/compile/passes/conftest.py diff --git a/tests/compile/passes/conftest.py b/tests/compile/passes/conftest.py new file mode 100644 index 000000000000..018747fd8cdc --- /dev/null +++ b/tests/compile/passes/conftest.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + + +@pytest.fixture(autouse=True) +def _workspace_manager_for_compile_passes(cleanup_fixture): + if not current_platform.is_rocm(): + yield + return + + # Release workspace tensors before cleanup_fixture flushes the allocator. + reset_workspace_manager() + if torch.accelerator.is_available(): + init_workspace_manager(torch.device(0)) + yield + reset_workspace_manager() diff --git a/tests/conftest.py b/tests/conftest.py index 0d70be62eb07..508e9fa3d50e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -901,13 +901,10 @@ def __enter__(self): # tests/basic_correctness/test_basic_correctness.py::test_models_distributed # where vllm worker processes are still alive and holding GPU # memory when hf_runner.__exit__ is called. - from tests.utils import ( - get_physical_device_indices, - record_gpu_memory_usage_stats, - ) + from tests.utils import record_gpu_memory_usage_stats if (device_count := current_platform.device_count()) > 0: - devices = get_physical_device_indices(devices=list(range(device_count))) + devices = list(range(device_count)) mem_usage_stats = record_gpu_memory_usage_stats(devices=devices) self.threshold_ratios = { device: 0.05 + mem_used / mem_tot diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py index 1405991c9ca5..fdb30880879a 100644 --- a/tests/kernels/moe/conftest.py +++ b/tests/kernels/moe/conftest.py @@ -3,6 +3,7 @@ import pytest from vllm.distributed import cleanup_dist_env_and_memory +from vllm.v1.worker.workspace import reset_workspace_manager NEEDS_CLEAN_ENTRY = frozenset( { @@ -58,4 +59,5 @@ def pytest_runtest_setup(item): # hook, so tests about to be skipped never reach this, while fixture setup # still runs afterwards. if item.path.name in NEEDS_CLEAN_ENTRY: + reset_workspace_manager() cleanup_dist_env_and_memory() diff --git a/tests/utils.py b/tests/utils.py index 7ee0e6680318..399a741df011 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -87,25 +87,7 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None: ) -if current_platform.is_rocm(): - from amdsmi import ( - amdsmi_get_gpu_vram_usage, - amdsmi_get_processor_handles, - amdsmi_init, - amdsmi_shut_down, - ) - - _amdsmi_lock = threading.Lock() - - @contextmanager - def _nvml(): - with _amdsmi_lock: - try: - amdsmi_init() - yield - finally: - amdsmi_shut_down() -elif current_platform.is_cuda(): +if current_platform.is_cuda(): from vllm.third_party.pynvml import ( nvmlDeviceGetHandleByIndex, nvmlDeviceGetHandleByUUID, @@ -595,27 +577,20 @@ def _find_pgid_members(pgid: int) -> list[int]: return members def _get_gpu_memory_used(self) -> float | None: - """Get total GPU memory used across all visible devices in bytes.""" + """Get device-wide usage across visible devices in bytes. + + On ROCm this initializes HIP in the caller; subsequent multiprocessing + GPU workers must use spawn, as the ROCm server helpers already do. + """ try: if current_platform.is_rocm(): - with _nvml(): - handles = amdsmi_get_processor_handles() - devices = get_physical_device_indices( - list(range(current_platform.device_count())) - ) - total_used_mib = 0 - for device in devices: - handle = handles[device] - vram_info = amdsmi_get_gpu_vram_usage(handle) - total_used_mib += vram_info["vram_used"] - # amdsmi reports VRAM in MiB; convert to bytes so this - # matches the CUDA/nvml branch (already bytes) and the - # byte-based target in _wait_for_gpu_memory_release. Without - # this, that wait compares MiB against a ~2e9-byte target, - # is always satisfied instantly, and returns "released to - # 0.00 GB" while the previous server's VRAM is still - # resident -- OOMing the next server's startup on ROCm. - return total_used_mib * 1024 * 1024 + # Use HIP logical devices. On MI355 DPX/NPS2, the primary + # AMD SMI device can expose whole-card memory counters. + total_used = 0 + for i in range(current_platform.device_count()): + free, total = torch.accelerator.get_memory_info(i) + total_used += total - free + return float(total_used) elif current_platform.is_cuda(): with _nvml(): total_used = 0 @@ -653,9 +628,12 @@ def _wait_for_gpu_memory_release( # Can't query GPU memory - nothing to do return - # Allow up to 2 GiB overhead above baseline for driver/context state - # that may persist between server instances. - headroom_bytes = 2 * 1024 * 1024 * 1024 + # Allow aggregate driver/context growth above the baseline. + headroom_bytes = ( + 4 * 1024 * 1024 * 1024 + if current_platform.is_rocm() + else 2 * 1024 * 1024 * 1024 + ) target = baseline + headroom_bytes start = time.time() @@ -1616,13 +1594,17 @@ def record_gpu_memory_usage_stats( *, devices: list[int], ) -> dict[int, tuple[float, float]]: + """Return device-wide used/total GiB; ROCm IDs are HIP logical devices. + + ROCm queries initialize HIP in the caller and include other processes' + allocations on the same memory partition. + """ output: dict[int, tuple[float, float]] = {} for device in devices: if current_platform.is_rocm(): - dev_handle = amdsmi_get_processor_handles()[device] - mem_info = amdsmi_get_gpu_vram_usage(dev_handle) - gb_used = mem_info["vram_used"] / 2**10 - gb_total = mem_info["vram_total"] / 2**10 + free, total = torch.accelerator.get_memory_info(device) + gb_used = (total - free) / 2**30 + gb_total = total / 2**30 elif current_platform.is_xpu(): # nvml/amdsmi are unavailable on XPU. Query device memory through # torch.accelerator.get_memory_info, which the XPU platform patches @@ -1650,7 +1632,9 @@ def wait_for_gpu_memory_to_clear( poll_interval_s: float = 5, ) -> None: assert threshold_bytes is not None or threshold_ratio is not None - devices = get_physical_device_indices(devices) + # HIP already applies device visibility; keep logical IDs and threshold keys. + if not current_platform.is_rocm(): + devices = get_physical_device_indices(devices) if isinstance(threshold_bytes, int): threshold_bytes = {device: threshold_bytes for device in devices} elif isinstance(threshold_bytes, dict): @@ -1672,8 +1656,6 @@ def wait_for_gpu_memory_to_clear( threshold_bytes.get(device, 0), min_threshold_b if ratio < 0.05 else 0 ) - # Use nvml instead of pytorch to reduce measurement error from torch cuda - # context. start_time = time.time() stable_since: float | None = None stable_used_bytes: dict[int, int] | None = None From 96791737887224c8bfec0a791517759f62e4d93b Mon Sep 17 00:00:00 2001 From: YukioZzz Date: Mon, 21 Sep 2026 05:30:34 +0800 Subject: [PATCH 547/774] [KVConnector][MoRIIO] Transfer hybrid mamba/KDA recurrent state in READ mode (#51052) Signed-off-by: Yichao Zhu --- .../moriio_toy_proxy_server.py | 2 + .../unit/test_moriio_connector.py | 11 + .../unit/test_moriio_hma_scheduler.py | 809 ++++++++++++ .../unit/test_moriio_kv_layout.py | 275 ++++ .../kv_connector/unit/test_moriio_routing.py | 1 + .../kv_connector/unit/test_moriio_tp_ack.py | 142 +- .../unit/test_moriio_transfer_completion.py | 129 ++ .../kv_connector/v1/moriio/moriio_common.py | 63 +- .../v1/moriio/moriio_connector.py | 1170 ++++++++++++++--- .../kv_connector/v1/moriio/moriio_engine.py | 38 +- .../kv_connector/v1/moriio/moriio_layout.py | 290 +++- 11 files changed, 2746 insertions(+), 184 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py create mode 100644 tests/v1/kv_connector/unit/test_moriio_transfer_completion.py diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index 8b2c46999efb..ce4d558fb919 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -367,6 +367,8 @@ async def handle_request(api: str, request: Request): "remote_block_ids" ] req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] + if prefill_kv.get("tp_size") is not None: + req_data["kv_transfer_params"]["tp_size"] = prefill_kv["tp_size"] req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[ "dp_size" diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index a30b5caf4bbd..f95e12c601bb 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -775,6 +775,17 @@ def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): assert scheduler.get_exchange_clipped_blocks(blocks) == blocks +def test_transfer_disabled_hybrid_group_is_ignored(): + config = _make_hybrid_kv_cache_config() + config.kv_cache_groups[1].enable_kv_transfer = False + scheduler = _read_scheduler(config) + + assert scheduler._is_hma_required is False + assert scheduler.get_exchange_clipped_blocks([[1, 2, 3], [20, 21, 22]]) == [ + [1, 2, 3] + ] + + def test_non_sliding_window_hybrid_is_rejected(): """A hybrid group that is not sliding-window (e.g. chunked-local attention) must fail closed rather than be silently mistransferred.""" diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py new file mode 100644 index 000000000000..13934f896afe --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -0,0 +1,809 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pure-CPU unit tests for the MoRIIO hybrid (attention + mamba/KDA) scheduler +and the KDA offset-template cache. + +These exercise the highest-risk, otherwise-untested scheduler logic without a +GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE +``N-1`` token accounting, P-side prompt truncation, and offset-template cache +wiring. + +Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is +running on ROCm with ``mori`` installed (importing the connector pulls in +``mori``). The authoritative run happens on the MIA recipe image. +""" + +import importlib +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +mori_available = importlib.util.find_spec("mori") is not None + +if not (current_platform.is_rocm() and mori_available): + pytest.skip( + "MoRIIOs are only available on ROCm with mori package installed", + allow_module_level=True, + ) + +moriio_connector = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector" +) +moriio_layout = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout" +) +ssm_conv_transfer_utils = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils" +) +MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo +MoRIIOMode = moriio_connector.MoRIIOMode +moriio_common = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common" +) + + +class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore[name-defined] + """Constructs a scheduler without the heavy real ``__init__``; only the + attributes read by the method under test are set.""" + + def __init__(self, **attrs): + self._mamba_group_ids: list[int] = [] + self._attn_group_ids: list[int] = [0] + self._ssm_state_slots_are_positional = False + self._is_hma_required = False + self.kv_cache_config = SimpleNamespace( + kv_cache_groups=[None, None], + select_transfer_block_ids=lambda block_ids: tuple(block_ids), + ) + for k, v in attrs.items(): + setattr(self, k, v) + + +class _FakeConnector(moriio_connector.MoRIIOConnector): # type: ignore[name-defined] + def __init__(self, connector_scheduler): + self.connector_scheduler = connector_scheduler + + +class _FakeWorker(moriio_connector.MoRIIOConnectorWorker): # type: ignore[name-defined] + def __init__(self, **attrs): + for k, v in attrs.items(): + setattr(self, k, v) + + +class _FakeBlocks: + def __init__(self, all_groups): + self._all_groups = all_groups + + def get_block_ids(self): + return self._all_groups + + +def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16): + inner = 1 + for d in per_slot_shape: + inner *= d + assert slot_stride >= inner + backing = torch.zeros(num_slots * slot_stride, dtype=dtype) + inner_strides: list[int] = [] + acc = 1 + for d in reversed(per_slot_shape): + inner_strides.insert(0, acc) + acc *= d + return backing.as_strided( + (num_slots, *per_slot_shape), (slot_stride, *inner_strides) + ) + + +def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2): + conv_dim = 2 * key_dim + value_dim + conv_state_bytes = conv_dim * conv_rows * dtype_size + ssm_state_bytes = 64 + return MambaConvSplitInfo( + conv_rows=conv_rows, + local_proj_dims=(key_dim, key_dim, value_dim), + conv_dtype_size=dtype_size, + ssm_sizes=(conv_state_bytes, ssm_state_bytes), + ) + + +def _mamba_spec(num_states: int = 2): + shapes = tuple((1, 1) for _ in range(num_states)) + return moriio_connector.MambaSpec( + block_size=16, + shapes=shapes, + dtypes=(torch.float32,) * num_states, + mamba_cache_mode="all", + ) + + +def _gate_vllm_config(*, read_mode: bool = True, speculative_config=None): + return SimpleNamespace( + kv_transfer_config=SimpleNamespace( + kv_connector_extra_config={"read_mode": read_mode} + ), + speculative_config=speculative_config, + ) + + +# -------------------------------------------------------------------------- +# split_block_groups +# -------------------------------------------------------------------------- +def test_split_block_groups_separates_attention_and_mamba(): + sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1]) + block_ids = ([1, 2], [99]) + attn, mamba = sched.split_block_groups(block_ids) + assert attn == [1, 2] + assert mamba == [[99]] + + +def test_split_block_groups_ignores_transfer_disabled_group(): + mamba_spec = moriio_connector.MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="all", + ) + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=mamba_spec), + ], + ) + config.transfer_groups = tuple( + group for group in config.kv_cache_groups if group.enable_kv_transfer + ) + config.select_transfer_block_ids = lambda block_ids: tuple( + block_ids[group_id] + for group_id, group in enumerate(config.kv_cache_groups) + if group.enable_kv_transfer + ) + attn_groups, mamba_groups = moriio_connector._split_kv_cache_group_kinds(config) + assert attn_groups == [0] + assert mamba_groups == [1] + + sched = _FakeScheduler( + _has_mamba=True, + _attn_group_ids=attn_groups, + _mamba_group_ids=mamba_groups, + kv_cache_config=config, + ) + assert sched.split_block_groups(([1, 2], [70, 71], [99])) == ( + [1, 2], + [[99]], + ) + + +def test_exchange_blocks_ignore_transfer_disabled_group(): + config = SimpleNamespace( + kv_cache_groups=[None, None, None], + select_transfer_block_ids=lambda block_ids: (block_ids[0], block_ids[2]), + ) + sched = _FakeScheduler( + kv_cache_config=config, + _is_hma_required=True, + blocks_per_sw=[0, 2], + ) + + assert sched.get_exchange_clipped_blocks(([1, 2], [70, 71], [90, 91, 92])) == [ + [1, 2], + [91, 92], + ] + + +def test_scheduler_rejects_multiple_attention_groups_with_mamba(): + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), + ], + ) + config.transfer_groups = tuple(config.kv_cache_groups) + + with pytest.raises(moriio_common.MoRIIOError, match="exactly one transferable"): + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) + + +def test_split_block_groups_preserves_multiple_mamba_groups(): + sched = _FakeScheduler( + _has_mamba=True, + _attn_group_ids=[0], + _mamba_group_ids=[1, 2], + ) + + assert sched.split_block_groups(([1, 2], [40], [70])) == ( + [1, 2], + [[40], [70]], + ) + + +def test_scheduler_rejects_incompatible_mamba_groups(): + incompatible = _mamba_spec() + object.__setattr__(incompatible, "page_size_padded", 4096) + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=incompatible), + ] + ) + config.transfer_groups = tuple(config.kv_cache_groups) + + with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"): + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) + + +def test_scheduler_rejects_mamba_group_without_two_states(): + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec(1)), + ], + ) + config.transfer_groups = tuple(config.kv_cache_groups) + + with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"): + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) + + +def test_scheduler_rejects_speculative_hybrid_read(): + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), + ], + ) + config.transfer_groups = tuple(config.kv_cache_groups) + + with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"): + moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(speculative_config=object()), "engine", config + ) + + +def test_scheduler_rejects_hybrid_write(): + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), + ], + ) + config.transfer_groups = tuple(config.kv_cache_groups) + + with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): + moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False), "engine", config + ) + + +def test_split_block_groups_rejects_non_hybrid_use(): + sched = _FakeScheduler(_has_mamba=False) + + with pytest.raises(moriio_common.MoRIIOError, match="non-hybrid"): + sched.split_block_groups(([5, 6, 7],)) + + +def test_split_block_groups_accepts_empty_abort_payload(): + sched = _FakeScheduler(_has_mamba=True) + + assert sched.split_block_groups(()) == ([], []) + + +def test_split_block_groups_keeps_positional_slots_in_all_mode(): + # mamba_cache_mode="all" keeps a state per block position. + sched = _FakeScheduler( + _has_mamba=True, + _attn_group_ids=[0], + _mamba_group_ids=[1], + _ssm_state_slots_are_positional=True, + ) + attn, mamba = sched.split_block_groups(([1], [40, 41, 42])) + assert attn == [1] + assert mamba == [[40, 41, 42]] + + +def test_split_block_groups_keeps_only_running_state_outside_all_mode(): + sched = _FakeScheduler( + _has_mamba=True, + _attn_group_ids=[0], + _mamba_group_ids=[1], + ) + assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]]) + + +# -------------------------------------------------------------------------- +# request_finished_all_groups +# -------------------------------------------------------------------------- +def test_request_finished_all_groups_carries_attn_and_mamba_in_one_field(): + seen = {} + + def _fake_request_finished(request, attn_block_ids, mamba_block_groups): + seen["attn"] = list(attn_block_ids) + seen["mamba"] = [list(group) for group in mamba_block_groups] + return True, { + "do_remote_prefill": True, + "remote_block_ids": [attn_block_ids, *mamba_block_groups], + } + + sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1]) + sched.request_finished = _fake_request_finished + conn = _FakeConnector(sched) + + request = SimpleNamespace(request_id="r0") + block_ids = ([10, 11], [42]) + delay_free, params = conn.request_finished_all_groups(request, block_ids) + + assert delay_free is True + # Attention ids drive request_finished; the mamba slot rides the SAME + # remote_block_ids channel as [attn, *mamba_groups], so the proxy/router + # needs no KDA-specific field. + assert seen["attn"] == [10, 11] + assert seen["mamba"] == [[42]] + assert params["remote_block_ids"] == [[10, 11], [42]] + assert "remote_mamba_block_ids" not in params + + +def test_request_finished_all_groups_pure_attention_stays_grouped(): + def _fake_request_finished(request, block_ids): + return True, {"remote_block_ids": block_ids} + + sched = _FakeScheduler(_has_mamba=False, _mamba_group_ids=[]) + sched.request_finished = _fake_request_finished + conn = _FakeConnector(sched) + + delay_free, params = conn.request_finished_all_groups( + SimpleNamespace(request_id="r1"), ([7, 8],) + ) + # Main's HMA contract keeps pure-attention block ids grouped. + assert params["remote_block_ids"] == ([7, 8],) + assert "remote_mamba_block_ids" not in params + + +def test_connector_supports_divergent_local_hybrid_hits(): + connector = _FakeConnector(None) + connector._transfers_mamba_state = True + connector.mode = MoRIIOMode.READ + assert connector.supports_divergent_local_hybrid_hits is True + connector.mode = MoRIIOMode.WRITE + assert connector.supports_divergent_local_hybrid_hits is False + + connector.mode = MoRIIOMode.READ + connector._transfers_mamba_state = False + assert connector.supports_divergent_local_hybrid_hits is False + + +def _make_read_scheduler(): + return _FakeScheduler( + is_producer=False, + mode=MoRIIOMode.READ, + _has_mamba=True, + _attn_group_ids=[0], + _mamba_group_ids=[1], + _ssm_state_slots_are_positional=True, + request_id_to_transfer_id={}, + transfer_id_to_request_id={}, + _reqs_need_recv={}, + _req_kv_params={}, + _max_decode_tail_blocks=1, + ) + + +def _make_read_request(remote_block_ids): + return SimpleNamespace( + request_id="req", + kv_transfer_params={ + "do_remote_prefill": True, + "transfer_id": "tx", + "remote_engine_id": "prefill", + "remote_block_ids": remote_block_ids, + }, + ) + + +def test_update_state_drops_decode_recompute_tail_block(): + sched = _make_read_scheduler() + request = _make_read_request([[10, 11], [90, 91]]) + blocks = _FakeBlocks( + all_groups=([100, 101, 102], [200, 201, 202]), + ) + + sched.update_state_after_alloc(request, blocks, num_external_tokens=256) + + assert sched._reqs_need_recv["req"][1] == [[100, 101], [200, 201]] + assert sched._req_kv_params["req"]["remote_block_ids"] == [ + [10, 11], + [90, 91], + ] + + +def test_update_state_pairs_shorter_local_blocks_with_remote_suffix(): + sched = _make_read_scheduler() + request = _make_read_request([[10, 11, 12], [92]]) + blocks = _FakeBlocks( + all_groups=([102], [202]), + ) + + sched.update_state_after_alloc(request, blocks, num_external_tokens=64) + + assert sched._reqs_need_recv["req"][1] == [[102], [202]] + assert sched._req_kv_params["req"]["remote_block_ids"] == [[12], [92]] + + +def test_update_state_full_attention_hit_still_carries_mamba_state(): + sched = _make_read_scheduler() + request = _make_read_request([[10], [90]]) + + sched.update_state_after_alloc( + request, + _FakeBlocks(all_groups=([100], [200])), + num_external_tokens=0, + ) + + assert sched._reqs_need_recv["req"][1] == [[], [200]] + + +def test_update_state_preserves_each_mamba_group(): + sched = _make_read_scheduler() + sched._mamba_group_ids = [1, 2] + request = _make_read_request([[10], [90], [80]]) + + sched.update_state_after_alloc( + request, + _FakeBlocks(all_groups=([100], [200], [300])), + num_external_tokens=64, + ) + + assert sched._reqs_need_recv["req"][1] == [[100], [200], [300]] + assert sched._req_kv_params["req"]["remote_block_ids"] == [ + [10], + [90], + [80], + ] + + +def test_align_read_blocks_allows_configured_lookahead_tail(): + assert moriio_connector.MoRIIOConnectorScheduler._align_read_blocks( + [100, 101, 102], [10], max_decode_tail_blocks=2 + ) == ([100], [10]) + with pytest.raises(ValueError, match="allowed local tail"): + moriio_connector.MoRIIOConnectorScheduler._align_read_blocks( + [100, 101, 102], [10], max_decode_tail_blocks=1 + ) + + +def test_update_state_aligns_attention_only_decode_tail(): + sched = _make_read_scheduler() + sched._has_mamba = False + sched._mamba_group_ids = [] + request = _make_read_request([[10, 11]]) + + sched.update_state_after_alloc( + request, + _FakeBlocks(all_groups=([100, 101, 102],)), + num_external_tokens=256, + ) + + assert sched._reqs_need_recv["req"][1] == [[100, 101]] + assert sched._req_kv_params["req"]["remote_block_ids"] == [[10, 11]] + + +def test_read_producer_does_not_queue_write_only_save_state(): + sched = _FakeScheduler( + is_producer=True, + mode=MoRIIOMode.READ, + _has_mamba=True, + _attn_group_ids=[0], + _mamba_group_ids=[1], + request_id_to_transfer_id={}, + transfer_id_to_request_id={}, + _reqs_need_save={}, + _req_kv_params={}, + ) + request = SimpleNamespace( + request_id="req", + kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"}, + ) + + sched.update_state_after_alloc( + request, + _FakeBlocks(all_groups=([100, 101], [200])), + num_external_tokens=0, + ) + + assert sched._reqs_need_save == {} + assert sched._req_kv_params == {} + + +def test_write_producer_queues_grouped_attention_blocks(): + sched = _FakeScheduler( + is_producer=True, + mode=MoRIIOMode.WRITE, + _has_mamba=False, + _attn_group_ids=[0], + _mamba_group_ids=[], + request_id_to_transfer_id={}, + transfer_id_to_request_id={}, + _reqs_need_save={}, + _req_kv_params={}, + ) + request = SimpleNamespace( + request_id="req", + kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"}, + ) + + sched.update_state_after_alloc( + request, + _FakeBlocks(all_groups=([100, 101],)), + num_external_tokens=0, + ) + + assert sched._reqs_need_save["req"][1] == ([100, 101],) + assert sched._req_kv_params["req"] == request.kv_transfer_params + + +def test_session_build_rejects_per_layer_region_count_mismatch(): + worker = _FakeWorker( + built_write_session={}, + layer_name_to_local_kv_cache_metadata={"kda.0": ["conv", "ssm"]}, + layer_name_to_remote_kv_cache_metadata={"prefill": {"kda.0": ["conv"]}}, + ) + + with pytest.raises(moriio_common.MoRIIOError, match="registered 1 region"): + worker._get_built_session("prefill") + + +def test_register_kv_caches_rejects_hybrid_write_before_registration(): + worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"}) + worker._is_mamba_layer = lambda _layer_name: True + + with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): + worker.register_kv_caches({"kda.0": object()}) + + +def test_register_kv_caches_rejects_no_transfer_enabled_layers(): + worker = _FakeWorker(mode=MoRIIOMode.READ, _transfer_layer_names=set()) + + with pytest.raises(moriio_common.MoRIIOError, match="no transfer-enabled"): + worker.register_kv_caches({"disabled.0": object()}) + + +def test_worker_rejects_incompatible_mamba_specs_before_registration(): + incompatible = _mamba_spec() + object.__setattr__(incompatible, "page_size_padded", 4096) + worker = _FakeWorker( + mode=MoRIIOMode.READ, + world_size=1, + _transfer_layer_names={"kda.0", "kda.1"}, + layer_to_spec={"kda.0": _mamba_spec(), "kda.1": incompatible}, + ) + + with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"): + worker.register_kv_caches({"kda.0": object(), "kda.1": object()}) + + +# -------------------------------------------------------------------------- +# _truncate_mamba_request_for_prefill +# -------------------------------------------------------------------------- +def _mk_request(prompt, max_tokens=64, params=None): + return SimpleNamespace( + kv_transfer_params={} if params is None else params, + num_prompt_tokens=len(prompt), + prompt_token_ids=list(prompt), + prompt_embeds=None, + _all_token_ids=list(prompt), + max_tokens=max_tokens, + ) + + +def test_truncate_mamba_request_pops_last_token_and_caps_tokens(): + sched = _FakeScheduler() + req = _mk_request([10, 11, 12, 13, 14]) + sched._truncate_mamba_request_for_prefill(req) + + assert req.prompt_token_ids == [10, 11, 12, 13] # last token popped + assert req._all_token_ids == [10, 11, 12, 13] + assert req.num_prompt_tokens == 4 # N-1 accounting + assert req.max_tokens == 1 + assert req.kv_transfer_params["_p_side_truncated"] is True + + +def test_truncate_mamba_request_is_idempotent_across_reschedule(): + sched = _FakeScheduler() + req = _mk_request([10, 11, 12]) + sched._truncate_mamba_request_for_prefill(req) + assert req.num_prompt_tokens == 2 + # A second call (e.g. after a preemption) must not truncate again. + sched._truncate_mamba_request_for_prefill(req) + assert req.num_prompt_tokens == 2 + assert req.prompt_token_ids == [10, 11] + + +def test_truncate_mamba_request_noop_for_single_token_prompt(): + sched = _FakeScheduler() + req = _mk_request([10]) + sched._truncate_mamba_request_for_prefill(req) + assert req.num_prompt_tokens == 1 + assert req.prompt_token_ids == [10] + assert "_p_side_truncated" not in req.kv_transfer_params + + +# -------------------------------------------------------------------------- +# get_num_new_matched_tokens (hybrid N-1 accounting) +# -------------------------------------------------------------------------- +def test_get_num_new_matched_tokens_read_recomputes_last_token(): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ, _has_mamba=True) + req = SimpleNamespace( + num_prompt_tokens=10, + prompt_token_ids=list(range(10)), + kv_transfer_params=None, + ) + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) + # READ always recomputes the final token locally: N-1 - computed. + assert n == 9 + assert is_async is False + + +def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens(): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False) + req = SimpleNamespace( + num_prompt_tokens=10, + prompt_token_ids=list(range(10)), + kv_transfer_params=None, + ) + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=2) + # Pure-attention WRITE: no N-1 drop; full length minus already-computed. + assert n == 8 + assert is_async is True + + +@pytest.mark.parametrize( + ("mode", "num_computed_tokens", "expected", "is_async"), + [ + (MoRIIOMode.READ, 0, 9, False), + (MoRIIOMode.READ, 10, 0, False), + (MoRIIOMode.WRITE, 2, 8, True), + ], +) +def test_get_num_new_matched_tokens_supports_embeds_only_prompts( + mode, num_computed_tokens, expected, is_async +): + sched = _FakeScheduler(is_producer=False, mode=mode, _has_mamba=True) + req = SimpleNamespace( + num_prompt_tokens=10, + prompt_token_ids=None, + prompt_embeds=object(), + kv_transfer_params=None, + ) + + assert sched.get_num_new_matched_tokens(req, num_computed_tokens) == ( + expected, + is_async, + ) + + +def test_get_num_new_matched_tokens_producer_returns_zero(): + sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(5)), params={"do_remote_decode": True}) + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) + assert (n, is_async) == (0, False) + # The producer never truncates here: the scheduler has already measured the + # prefix-cache hit against the full prompt by this point. + assert req.num_prompt_tokens == 5 + + +def test_on_new_request_truncates_producer_prompt(): + sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(5)), params={"do_remote_decode": True}) + sched.on_new_request(req) + # Producer stops at h(N-1): the last prompt token was dropped. + assert req.num_prompt_tokens == 4 + assert req.max_tokens == 1 + assert req.kv_transfer_params["_p_side_truncated"] is True + + +def test_on_new_request_is_noop_without_mamba_or_on_decode(): + plain = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=False) + req = _mk_request(list(range(5)), params={"do_remote_decode": True}) + plain.on_new_request(req) + assert req.num_prompt_tokens == 5 + + decode_side = _FakeScheduler( + is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True + ) + req2 = _mk_request(list(range(5)), params={"do_remote_prefill": True}) + decode_side.on_new_request(req2) + assert req2.num_prompt_tokens == 5 + + +# -------------------------------------------------------------------------- +# Offset-template cache +# -------------------------------------------------------------------------- +_SAMPLE_SLOT_SETS = [ + ([1, 2], [1, 2]), + ([0], [3]), + ([3, 0, 2], [3, 0, 2]), + ([2, 2], [5, 5]), + ([], []), +] + + +def test_worker_compute_mamba_offsets_caches_and_matches_template(): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv_shape = (conv_dim, split.conv_rows) + ssm_shape = (2, 4, 4) + page_bytes = 256 + pages = torch.zeros((8, 1, 1, page_bytes), dtype=torch.uint8) + spec = moriio_connector.MambaSpec( + block_size=1, + shapes=(conv_shape, ssm_shape), + dtypes=(torch.bfloat16, torch.bfloat16), + page_size_padded=page_bytes, + ) + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec) + + worker = _FakeWorker( + kv_caches={"kda.0": pages}, + layer_to_spec={"kda.0": spec}, + _conv_decomp=split, + _mamba_offset_templates={}, + tp_rank=0, + world_size=1, + ) + + cached_template = None + for local_slots, remote_slots in _SAMPLE_SLOT_SETS: + lo, ro, sz, n_conv = worker._compute_mamba_transfer_offsets( + "kda.0", local_slots, remote_slots + ) + template = moriio_layout.build_mamba_offset_template( + conv, ssm, split, tp_ratio=1 + ) + expected = moriio_layout.apply_mamba_offset_template( + template, local_slots, remote_slots + ) + assert (lo, ro, sz) == expected + assert n_conv == moriio_layout.compute_mamba_conv_split_count( + local_slots, split + ) + + current_template = worker._mamba_offset_templates["kda.0"] + if cached_template is None: + cached_template = current_template + else: + assert current_template is cached_template + + assert "kda.0" in worker._mamba_offset_templates + + +@pytest.mark.parametrize("remote_tp_size", [1, None]) +def test_mamba_tp_ratio_accepts_only_equal_tp(remote_tp_size): + worker = _FakeWorker(world_size=1) + assert worker._mamba_tp_ratio(remote_tp_size) == 1 + + +@pytest.mark.parametrize("remote_tp_size", [2, 3]) +def test_mamba_tp_ratio_rejects_heterogeneous_tp(remote_tp_size): + worker = _FakeWorker(world_size=1) + with pytest.raises(NotImplementedError, match="heterogeneous-TP"): + worker._mamba_tp_ratio(remote_tp_size) + + +def test_worker_selects_each_mamba_layers_own_block_group(): + worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 1, "kda.1": 2}) + groups = ([10], [40], [70]) + + assert worker._mamba_blocks_for_layer("kda.0", groups) == [40] + assert worker._mamba_blocks_for_layer("kda.1", groups) == [70] + + +def test_worker_rejects_missing_mamba_block_group(): + worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 2}) + + with pytest.raises(moriio_common.MoRIIOError, match="no block group"): + worker._mamba_blocks_for_layer("kda.0", ([10], [40])) diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index c9fa434ee9b9..1cb5e271bf4e 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.util +import math import threading from collections import OrderedDict, defaultdict from queue import Queue @@ -14,6 +15,7 @@ from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + MambaSpec, MLAAttentionSpec, ) @@ -36,6 +38,10 @@ "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout" ) msgpack = importlib.import_module("msgpack") +ssm_conv_transfer_utils = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils" +) +MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo ROLE = moriio_common.ROLE MoRIIOError = moriio_common.MoRIIOError @@ -353,6 +359,9 @@ class FakeWorker: kv_caches: dict[str, torch.Tensor] layer_name_to_local_kv_cache_metadata: dict[str, list[Any]] + def _region_session_indices(self, layer_name): + return [list(self.kv_caches).index(layer_name)] + def _compute_block_transfer_offsets( self, layer_name, @@ -673,6 +682,30 @@ def test_local_block_ids_longer_than_remote_raises_value_error(): ) +@pytest.mark.parametrize( + ("local_block_ids", "remote_block_ids", "match"), + [ + ([-1], [0], "local block id -1"), + ([0], [8], "remote block id 8"), + ], +) +def test_block_ids_must_fit_registered_regions( + local_block_ids, remote_block_ids, match +): + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + with pytest.raises(ValueError, match=match): + moriio_layout.compute_block_transfer_offsets( + "layer", + cache, + worker.layer_to_spec, + local_block_ids, + remote_block_ids, + _remote_meta(num_blocks=8).num_blocks, + ) + + def test_empty_local_block_ids_is_free_only_noop(): cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) worker = _worker({"layer": cache}, {"layer": _full_spec()}) @@ -745,3 +778,245 @@ def test_unsupported_shape_raises_value_error(): with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"): moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec) + + +# --------------------------------------------------------------------------- +# Mamba/KDA (conv + ssm) transfer geometry +# --------------------------------------------------------------------------- + + +def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16): + """Build a slot-strided view: shape (num_slots, *per_slot_shape) whose slot + stride (dim 0) exceeds the per-slot element count, mimicking the K3 KDA + conv/ssm cache tensors.""" + inner = 1 + for d in per_slot_shape: + inner *= d + assert slot_stride >= inner + backing = torch.zeros(num_slots * slot_stride, dtype=dtype) + inner_strides: list[int] = [] + acc = 1 + for d in reversed(per_slot_shape): + inner_strides.insert(0, acc) + acc *= d + return backing.as_strided( + (num_slots, *per_slot_shape), (slot_stride, *inner_strides) + ) + + +def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2): + """A GDN-style conv split: sub-projections (key, key, value).""" + conv_dim = 2 * key_dim + value_dim + conv_state_bytes = conv_dim * conv_rows * dtype_size + ssm_state_bytes = 64 # opaque here; ssm size is taken from the tensor + return MambaConvSplitInfo( + conv_rows=conv_rows, + local_proj_dims=(key_dim, key_dim, value_dim), + conv_dtype_size=dtype_size, + ssm_sizes=(conv_state_bytes, ssm_state_bytes), + ) + + +def test_mamba_transfer_geometry_is_slot_strided(): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(5, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(5, (2, 4, 4), slot_stride=64) + + geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm) + + assert geom.conv_slot_stride == 200 + assert geom.ssm_slot_stride == 64 + assert geom.ssm_slot_bytes == (2 * 4 * 4) * ssm.element_size() + # Region spans the full slot-strided extent (shape[0] * stride(0)), not just + # numel(), so the highest-slot transfers stay in-bounds. + assert geom.conv_region_len == conv.shape[0] * conv.stride(0) * conv.element_size() + assert geom.ssm_region_len == ssm.shape[0] * ssm.stride(0) * ssm.element_size() + + +def test_kda_conv_ssm_unpacks_padded_pages_as_aliasing_strided_views(): + num_blocks = 4 + page_bytes = 128 + conv_shape = (3, 4) + ssm_shape = (2, 2, 4) + conv_bytes = math.prod(conv_shape) * torch.bfloat16.itemsize + ssm_bytes = math.prod(ssm_shape) * torch.bfloat16.itemsize + assert conv_bytes + ssm_bytes < page_bytes + + pages = torch.zeros((num_blocks, 1, 1, page_bytes), dtype=torch.uint8) + spec = MambaSpec( + block_size=1, + shapes=(conv_shape, ssm_shape), + dtypes=(torch.bfloat16, torch.bfloat16), + page_size_padded=page_bytes, + ) + + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec) + + assert conv.shape == (num_blocks, *conv_shape) + assert ssm.shape == (num_blocks, *ssm_shape) + assert conv.stride() == (page_bytes // 2, 4, 1) + assert ssm.stride() == (page_bytes // 2, 8, 4, 1) + assert conv.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr() + assert ssm.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr() + assert conv.storage_offset() == 0 + assert ssm.storage_offset() * ssm.element_size() == conv_bytes + + conv[1].fill_(1) + ssm[2].fill_(2) + assert torch.all(pages[1, 0, 0, :conv_bytes].view(torch.bfloat16) == 1) + assert torch.all( + pages[2, 0, 0, conv_bytes : conv_bytes + ssm_bytes].view(torch.bfloat16) == 2 + ) + + +def test_kda_conv_ssm_accepts_strided_blocks_without_copying(): + backing = torch.zeros(4 * 256, dtype=torch.uint8) + pages = backing.as_strided((4, 1, 1, 128), (256, 128, 128, 1)) + spec = MambaSpec( + block_size=1, + shapes=((3, 4), (2, 2, 4)), + dtypes=(torch.bfloat16, torch.bfloat16), + page_size_padded=128, + ) + + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec) + + assert not pages.is_contiguous() + assert conv.stride(0) == 128 + assert ssm.stride(0) == 128 + assert conv.untyped_storage().data_ptr() == backing.data_ptr() + assert ssm.untyped_storage().data_ptr() == backing.data_ptr() + + +def test_kda_conv_ssm_rejects_noncontiguous_page_bytes(): + pages = torch.zeros((4, 1, 1, 256), dtype=torch.uint8)[..., ::2] + spec = MambaSpec( + block_size=1, + shapes=((3, 4), (2, 2, 4)), + dtypes=(torch.bfloat16, torch.bfloat16), + page_size_padded=128, + ) + + with pytest.raises(ValueError, match="contiguous byte dimension"): + moriio_layout.kda_conv_ssm(pages, spec) + + +def test_kda_conv_ssm_rejects_non_byte_packed_cache(): + pages = torch.zeros((4, 1, 1, 128), dtype=torch.bfloat16) + spec = MambaSpec( + block_size=1, + shapes=((3, 4), (2, 2, 4)), + dtypes=(torch.bfloat16, torch.bfloat16), + page_size_padded=128, + ) + + with pytest.raises(ValueError, match="byte-sized dtype"): + moriio_layout.kda_conv_ssm(pages, spec) + + +def test_mamba_transfer_geometry_requires_matching_slot_counts(): + conv = torch.empty((5, 4, 3), dtype=torch.bfloat16) + ssm = torch.empty((4, 2, 4, 4), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="conv/ssm slot count mismatch"): + moriio_layout.get_mamba_transfer_geometry(conv, ssm) + + +def test_mamba_transfer_geometry_rejects_empty_caches(): + conv = torch.empty((0, 4, 3), dtype=torch.bfloat16) + ssm = torch.empty((0, 2, 4, 4), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="at least one slot"): + moriio_layout.get_mamba_transfer_geometry(conv, ssm) + + +def test_mamba_conv_ssm_offsets_homogeneous_tp(): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64) + geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm) + + local_slots = [1, 2] + remote_slots = [1, 2] + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1) + local_offs, remote_offs, sizes = moriio_layout.apply_mamba_offset_template( + template, local_slots, remote_slots + ) + + n_conv = moriio_layout.compute_mamba_conv_split_count(local_slots, split) + assert n_conv == len(local_slots) * len(split.local_proj_dims) + assert len(local_offs) == n_conv + len(local_slots) + + # Homogeneous TP: remote conv slices equal local slices; since local==remote + # slots, conv local and remote offsets coincide. + loc_co = split.local_conv_offsets + rem_co = split.remote_conv_offsets(0, 1) + k = 0 + esz = conv.element_size() + for ls, rs in zip(local_slots, remote_slots): + lbase = ls * geom.conv_slot_stride * esz + rbase = rs * geom.conv_slot_stride * esz + for (loff, _), (roff, rsz) in zip(loc_co, rem_co): + assert local_offs[k] == lbase + loff + assert remote_offs[k] == rbase + roff + assert sizes[k] == rsz + k += 1 + + # SSM: whole per-slot block, no head sharding at tp_ratio == 1. + essz = ssm.element_size() + for j, (ls, rs) in enumerate(zip(local_slots, remote_slots)): + assert local_offs[n_conv + j] == ls * geom.ssm_slot_stride * essz + assert remote_offs[n_conv + j] == rs * geom.ssm_slot_stride * essz + assert sizes[n_conv + j] == geom.ssm_slot_bytes + + +def test_mamba_conv_ssm_offsets_heterogeneous_tp_is_gated(): + # Heterogeneous-TP mamba transfer is a documented follow-up (remote slot + # stride + multi-rank fan-in gather). Until then the offset helper must fail + # loudly rather than silently transfer wrong bytes, since the mamba path + # bypasses the attention KV-head validator. + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(4, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(4, (2, 4, 4), slot_stride=64) + with pytest.raises(NotImplementedError): + moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=2) + + +def test_local_mamba_slots_must_fit_registered_regions(): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64) + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1) + + with pytest.raises(ValueError, match="local mamba slot 6"): + moriio_layout.apply_mamba_offset_template(template, [6], [0]) + + +@pytest.mark.parametrize( + ("local_slots", "remote_slots"), + [([1, 2], [3]), ([1], [2, 3])], +) +def test_mamba_slot_counts_must_match(local_slots, remote_slots): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64) + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1) + + with pytest.raises(ValueError, match="slot counts must match"): + moriio_layout.apply_mamba_offset_template(template, local_slots, remote_slots) + + +def test_remote_mamba_slots_must_be_non_negative(): + split = _gdn_split_info() + conv_dim = sum(split.local_proj_dims) + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200) + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64) + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1) + + with pytest.raises(ValueError, match="must be non-negative"): + moriio_layout.apply_mamba_offset_template(template, [1], [-1]) diff --git a/tests/v1/kv_connector/unit/test_moriio_routing.py b/tests/v1/kv_connector/unit/test_moriio_routing.py index fc8b323218c2..b0ac6df3cba3 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing.py @@ -28,6 +28,7 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): s = SimpleNamespace( is_producer=True, _is_hma_required=False, + _has_mamba=False, blocks_per_sw=[0], _global_dp_rank=global_dp_rank, engine_id="engine-abc", diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index da412fcd9290..a35f0a1169b6 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +from types import SimpleNamespace import pytest from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( MoRIIOMode, MoRIIOTransferAck, + TransferError, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( MoRIIOConnector, @@ -17,6 +19,9 @@ resolve_moriio_transfer_ack, validate_moriio_heterogeneous_tp_kv_heads, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) def test_remote_tp_rank_same_tp_maps_to_self(): @@ -277,6 +282,9 @@ def __init__(self): self.lock = threading.Lock() self.sent = [] + # Exercise the real batch verdict rather than restating it here. + poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch + def send_notify( self, transfer_id, @@ -292,13 +300,18 @@ def shutdown(self): worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) worker.world_size = 8 + worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0) worker.moriio_wrapper = FakeWrapper() - worker._recving_transfers = {"req": {"layer0": DoneStatus()}} + # A layer maps to the list of reads posted for it (a KDA layer posts two). + worker._recving_transfers = {"req": {"layer0": [DoneStatus(), DoneStatus()]}} worker._recving_transfers_callback_addr = { "req": ("127.0.0.1", "7000", "tx-release") } # Transfer-timeout reaping state consulted by _pop_done_transfers. worker._recving_transfers_start = {} + # Load-error bookkeeping cleared alongside a completed transfer. + worker._recving_local_blocks = {} + worker._invalid_block_ids = set() assert worker._pop_done_transfers() == {"tx-release"} assert worker.moriio_wrapper.sent == [ @@ -314,6 +327,133 @@ def shutdown(self): assert worker._recving_transfers_callback_addr == {} +def test_aborted_read_without_local_blocks_only_releases_remote_blocks(): + class FakeWrapper: + def __init__(self): + self.sent = [] + + def send_notify( + self, + transfer_id, + host, + port, + message_type=None, + message_fields=None, + ): + self.sent.append((transfer_id, host, port, message_type, message_fields)) + + def shutdown(self): + pass + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.mode = MoRIIOMode.READ + worker.world_size = 8 + worker.moriio_wrapper = FakeWrapper() + + worker._read_blocks( + local_block_ids=[], + remote_block_ids=[[10, 11], [90]], + dst_engine_id="prefill", + request_id="req-aborted", + transfer_id="tx-aborted", + remote_host="127.0.0.1", + remote_notify_port=7000, + remote_tp_size=8, + remote_dp_rank=0, + chosen_tp=2, + ) + + assert worker.moriio_wrapper.sent == [ + ( + "tx-aborted", + "127.0.0.1", + "7002", + "release", + {"consumer_tp_size": 8}, + ) + ] + + +def test_read_completion_waits_for_every_posted_read(): + class Status: + def __init__(self, done: bool): + self.done = done + + def Succeeded(self): + return self.done + + def Failed(self): + return False + + class FakeWrapper: + def __init__(self): + self.lock = threading.Lock() + self.sent = [] + + poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch + + def send_notify(self, *args, **kwargs): + self.sent.append((args, kwargs)) + + def shutdown(self): + pass + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.world_size = 8 + worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0) + worker.moriio_wrapper = FakeWrapper() + worker._recving_transfers = { + "req": {"kda_layer": [Status(done=True), Status(done=False)]} + } + worker._recving_transfers_callback_addr = { + "req": ("127.0.0.1", "7000", "tx-pending") + } + worker._recving_transfers_start = {"req": float("inf")} + worker._recving_local_blocks = {"req": [1, 2]} + worker._invalid_block_ids = set() + + assert worker._pop_done_transfers() == set() + assert "req" in worker._recving_transfers + assert worker.moriio_wrapper.sent == [] + + +@pytest.mark.parametrize( + ("has_mamba", "expected_invalid"), + [(False, {1, 2}), (True, set())], +) +def test_failed_read_reports_blocks_only_without_hma(has_mamba, expected_invalid): + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = has_mamba + worker._recving_local_blocks = {"req": [1, 2]} + worker._invalid_block_ids = set() + + worker._record_failed_recv("req") + + assert worker.get_block_ids_with_load_errors() == expected_invalid + + +def test_hybrid_step_barrier_fails_closed(monkeypatch): + class FailingWrapper: + def waiting_for_transfer_complete(self, _statuses): + raise TransferError("failed") + + def shutdown(self): + pass + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = True + worker._reads_issued_this_step = [object()] + worker._mamba_reads_this_step = [object()] + worker.moriio_wrapper = FailingWrapper() + monkeypatch.setattr( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.get_forward_context", + lambda: SimpleNamespace(cudagraph_runtime_mode=None), + ) + + with pytest.raises(TransferError, match="failed"): + worker._await_reads_issued_this_step() + + def test_requested_cudagraph_mode_is_never_overridden(): # The configured cudagraph mode is always honored: the barrier fires when # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full diff --git a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py new file mode 100644 index 000000000000..2a3092ad5ae9 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Transfer-completion semantics of the MoRIIO wrapper. + +Covers the two behaviours the READ barrier depends on: + * poll_transfer_batch considers EVERY status, not just the newest one + * waiting_for_transfer_complete blocks until terminal and raises on failure. +""" + +import importlib +import importlib.util +import threading +import time + +import pytest + +from vllm.platforms import current_platform + +mori_available = importlib.util.find_spec("mori") is not None + +if not (current_platform.is_rocm() and mori_available): + pytest.skip( + "MoRIIOs are only available on ROCm with mori package installed", + allow_module_level=True, + ) + +moriio_common = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common" +) +moriio_engine = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine" +) + +MoRIIOWrapper = moriio_engine.MoRIIOWrapper +TransferBatchState = moriio_common.TransferBatchState +TransferError = moriio_common.TransferError + + +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan")]) +def test_timeout_must_be_positive_and_finite(value): + with pytest.raises(ValueError, match="must be finite and greater than zero"): + moriio_common._positive_finite_timeout("test_timeout", value) + + +def test_positive_finite_timeout_is_preserved(): + assert moriio_common._positive_finite_timeout("test_timeout", "1.5") == 1.5 + + +class FakeStatus: + """Mirrors the mori TransferStatus surface the wrapper uses.""" + + def __init__(self, state: str = "in_progress", message: str = ""): + self.state = state + self.message = message + + def Succeeded(self) -> bool: + return self.state == "success" + + def Failed(self) -> bool: + return self.state == "failed" + + def Message(self) -> str: + return self.message + + def Code(self) -> str: + return self.state + + +def _ok(): + return FakeStatus("success") + + +def _pending(): + return FakeStatus("in_progress") + + +def _bad(): + return FakeStatus("failed", "SQ full") + + +@pytest.mark.parametrize( + "statuses,expected", + [ + ([], TransferBatchState.DONE), + ([_ok(), _ok()], TransferBatchState.DONE), + # Each case below puts the interesting status BEFORE a success, which a + # statuses[-1] check would have missed entirely. + ([_bad(), _ok()], TransferBatchState.FAILED), + ([_pending(), _ok()], TransferBatchState.PENDING), + ([_pending(), _bad(), _ok()], TransferBatchState.FAILED), + ], +) +def test_poll_transfer_batch_considers_every_status(statuses, expected): + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05) + assert wrapper.poll_transfer_batch(statuses) is expected + + +@pytest.mark.parametrize( + "statuses_factory,expected_message", + [ + (lambda: [_ok(), _bad()], "SQ full"), + (lambda: [_ok(), _pending()], "timed out"), + ], + ids=["failed", "timeout"], +) +def test_wait_raises_with_detail(statuses_factory, expected_message): + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05) + with pytest.raises(TransferError, match=expected_message): + wrapper.waiting_for_transfer_complete(statuses_factory()) + + +def test_wait_returns_immediately_when_nothing_pending(): + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05) + wrapper.waiting_for_transfer_complete([]) + wrapper.waiting_for_transfer_complete([_ok(), _ok()]) + + +def test_wait_blocks_until_terminal(): + """The barrier is only worth anything if it actually waits.""" + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=5.0) + late = FakeStatus("in_progress") + timer = threading.Timer(0.2, lambda: setattr(late, "state", "success")) + timer.start() + try: + start = time.monotonic() + wrapper.waiting_for_transfer_complete([late]) + assert time.monotonic() - start >= 0.15 + finally: + timer.cancel() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index f25e53e3f0f9..5f1068562cf8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +import math import os import threading import time @@ -45,6 +46,13 @@ TransferOffsetsKey = tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype] +def _positive_finite_timeout(name: str, value: Any) -> float: + timeout = float(value) + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError(f"{name} must be finite and greater than zero") + return timeout + + class MoRIIOTransferAck(NamedTuple): transfer_id: TransferId consumer_tp_size: int = 1 @@ -160,6 +168,20 @@ class MoRIIOMode(Enum): WRITE = "write" +class TransferBatchState(Enum): + """Verdict for a group of transfer statuses that belong to one request. + + A request's KV transfer is spread over one status per layer (two for a KDA + layer), so a single status never decides the request: PENDING means at + least one is still in flight and none has failed, FAILED means at least + one failed, DONE means all succeeded. + """ + + DONE = "done" + FAILED = "failed" + PENDING = "pending" + + class MoRIIOError(Exception): """Base exception for MoRIIO operations.""" @@ -260,6 +282,7 @@ class MoRIIOConfig: tp_size: int transfer_timeout: float defer_timeout: float + recv_abort_timeout: float read_mode: bool = False qp_per_transfer: int = 1 post_batch_size: int = -1 @@ -284,6 +307,8 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": # raising TransferError (sec). # defer_timeout -> Timeout before a deferred send with no finished_sending # notification is reaped and its blocks force-freed (sec). + # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA + # completion never arrived is aborted (sec). # Knobs for RDMA transfers, ignored if on xgmi backend # qp_per_transfer -> Number of RDMA Queue Pairs per KV transfer. @@ -315,13 +340,29 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": "must be one of 'rdma' or 'xgmi'." ) - transfer_timeout = float( + transfer_timeout = _positive_finite_timeout( + "kv_connector_extra_config.transfer_timeout", extra_config.get( "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT - ) + ), ) - defer_timeout = float( - extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT) + defer_timeout = _positive_finite_timeout( + "kv_connector_extra_config.defer_timeout", + extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT), + ) + legacy_recv_abort_timeout = os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S") + if legacy_recv_abort_timeout is not None: + logger.warning_once( + "The environment variable VLLM_MORIIO_TRANSFER_TIMEOUT_S is " + "deprecated. Set 'recv_abort_timeout' inside " + "kv_transfer_config.kv_connector_extra_config instead." + ) + recv_abort_timeout = _positive_finite_timeout( + "kv_connector_extra_config.recv_abort_timeout", + extra_config.get( + "recv_abort_timeout", + legacy_recv_abort_timeout or MoRIIOConstants.DEFAULT_RECV_ABORT_TIMEOUT, + ), ) return cls( @@ -347,6 +388,7 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": backend=backend, transfer_timeout=transfer_timeout, defer_timeout=defer_timeout, + recv_abort_timeout=recv_abort_timeout, ) @@ -374,6 +416,10 @@ class MoRIIOConstants: # notification is reaped and its blocks force-freed. # Overridable via kv_connector_extra_config["defer_timeout"]. DEFAULT_DEFER_TIMEOUT = 60.0 + # Timeout (seconds) before an in-flight recv whose RDMA completion was lost + # is aborted, so the decode worker does not hang on it forever. + # Overridable via kv_connector_extra_config["recv_abort_timeout"]. + DEFAULT_RECV_ABORT_TIMEOUT = 120.0 # The router embeds both zmq_addresses in the request_id: @@ -567,9 +613,12 @@ def add_new_req( remote_handshake_port=int(remote_handshake_port), remote_notify_port=int(remote_notify_port), # Remote peer TP degree (used as remote_tp_size downstream). The - # proxy advertises it under "remote_tp_size"; #46332 read "tp_size" - # which is absent on WRITE producer requests -> defaulted to 1 -> - # rank collapse. Read the right key; 0 == unknown (== homogeneous). + # decode side must know the prefiller's TP degree to map producer + # ranks onto the correct decode rank (and vice versa). The old plain + # "tp_size" key (#46332) is absent on WRITE producer requests and + # defaulted to 1, collapsing every producer rank onto decode rank 0 + # and corrupting the transfer; prefer "remote_tp_size". 0 == unknown + # (treated as homogeneous downstream). tp_size=int( kv_transfer_params.get("remote_tp_size") or kv_transfer_params.get("tp_size") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index f15eee9bd57f..dce4a9806895 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2,14 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import math -import os import queue import threading import time from collections import defaultdict from collections.abc import Collection from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import msgpack import msgspec @@ -33,10 +32,13 @@ MoRIIOConfig, MoRIIOConnectorMetadata, MoRIIOConstants, + MoRIIOError, MoRIIOMode, MoRIIOTransferAck, ReqId, ReqMeta, + TransferBatchState, + TransferError, TransferId, WriteTask, fold_local_rank, @@ -56,11 +58,20 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import ( LayerTransferGeometry, + MambaOffsetTemplate, + apply_mamba_offset_template, build_layer_to_spec, + build_mamba_offset_template, compute_block_transfer_offsets, + compute_mamba_conv_split_count, get_layer_transfer_geometry, is_mla_cache_layer, iter_layer_registration_regions, + kda_conv_ssm, +) +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( + MambaConvSplitInfo, + derive_mamba_conv_split, ) from vllm.distributed.parallel_state import ( get_tensor_model_parallel_world_size, @@ -79,6 +90,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, + MambaSpec, SlidingWindowSpec, ) from vllm.v1.outputs import KVConnectorOutput @@ -91,6 +103,13 @@ logger = init_logger(__name__) +# Backoff bounds for re-posting an RDMA READ that the NIC rejected because the +# send queue was full. A separate CQ-poll thread drains completions, so the +# wait only needs to be long enough for it to make progress. +_SQ_FULL_BACKOFF_INITIAL_S = 0.001 +_SQ_FULL_BACKOFF_MAX_S = 0.05 +_MAX_LOCAL_DECODE_TAIL_BLOCKS = 1 + try: from mori.io import ( @@ -197,6 +216,12 @@ def resolve_moriio_transfer_ack( class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): + @property + def supports_divergent_local_hybrid_hits(self) -> bool: + # The READ path always transfers the recurrent-state slot, including + # when the attention prefix is already present locally. + return self.mode == MoRIIOMode.READ and self._transfers_mamba_state + def __init__( self, vllm_config: VllmConfig, @@ -217,19 +242,10 @@ def __init__( + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"]) ) self.mode = get_moriio_mode(self.kv_transfer_config) - if ( - self.mode == MoRIIOMode.READ - and self.kv_transfer_config.is_kv_consumer - and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs() - ): - # warn only; kv-read barrier requires PIECEWISE cudagraph mode - logger.warning_once( - "MoRIIO READ mode is running with %s CUDA graphs: per-layer " - "KV-read barrier can't fire inside full graph; accuracy may " - "degrade at high concurrency. Set cudagraph_mode=PIECEWISE " - "in --compilation-config.", - vllm_config.compilation_config.cudagraph_mode.name, - ) + self._transfers_mamba_state = any( + group.enable_kv_transfer and isinstance(group.kv_cache_spec, MambaSpec) + for group in kv_cache_config.kv_cache_groups + ) if role == KVConnectorRole.SCHEDULER: self.connector_scheduler: MoRIIOConnectorScheduler | None = ( MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) @@ -279,6 +295,11 @@ def update_state_after_alloc( request, blocks, num_external_tokens, self.connector_worker ) + def on_new_request(self, request: "Request") -> None: + """Fired when the scheduler admits a request, before any allocation.""" + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -300,6 +321,14 @@ def request_finished_all_groups( block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None + if self.connector_scheduler._has_mamba: + attn_block_ids, mamba_block_groups = ( + self.connector_scheduler.split_block_groups(block_ids) + ) + delay_free, params = self.connector_scheduler.request_finished( + request, attn_block_ids, mamba_block_groups + ) + return delay_free, params return self.connector_scheduler.request_finished(request, block_ids) def update_connector_output(self, connector_output: KVConnectorOutput) -> None: @@ -318,6 +347,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_block_ids_with_load_errors(self) -> set[int]: + """Blocks whose MoRIIO read failed, for the scheduler to recompute.""" + if self.connector_worker is None: + return set() + return self.connector_worker.get_block_ids_with_load_errors() + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.CONSUMER: @@ -381,6 +416,47 @@ def has_connector_metadata(self) -> bool: return False +def _split_kv_cache_group_kinds( + kv_cache_config: "KVCacheConfig", +) -> tuple[list[int], list[int]]: + """Classify kv_cache_groups into (attention_group_indices, + mamba_group_indices) by whether each group's spec is a MambaSpec. + + Groups that opt out of external transfer are absent from + ``KVCacheConfig.transfer_groups``. Returned indices address that filtered + sequence and therefore also the connector's transfer block payload. + """ + attn: list[int] = [] + mamba: list[int] = [] + for gi, group in enumerate(kv_cache_config.transfer_groups): + if isinstance(group.kv_cache_spec, MambaSpec): + mamba.append(gi) + else: + attn.append(gi) + return attn, mamba + + +def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None: + """Validate the common packed-state contract and return its representative.""" + specs = tuple(specs) + if not specs: + return None + for spec in specs: + if len(spec.shapes) != 2 or len(spec.dtypes) != 2: + raise MoRIIOError( + "MoRIIO hybrid READ supports exactly two Mamba states " + "(conv and SSM), got " + f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes" + ) + reference = specs[0] + if any(spec != reference for spec in specs[1:]): + raise MoRIIOError( + "MoRIIO hybrid READ requires all transferable Mamba groups " + "to use the same cache spec" + ) + return reference + + class MoRIIOConnectorScheduler: """Implementation of Scheduler side methods.""" @@ -391,35 +467,73 @@ def __init__( kv_cache_config: "KVCacheConfig", ): self.vllm_config = vllm_config - + self.kv_cache_config = kv_cache_config assert vllm_config.kv_transfer_config is not None, ( "kv_transfer_config must be set for MoRIIOConnector" ) self.kv_transfer_config = vllm_config.kv_transfer_config + self.mode = get_moriio_mode(self.kv_transfer_config) + self._attn_group_ids, self._mamba_group_ids = _split_kv_cache_group_kinds( + kv_cache_config + ) + self._has_mamba = bool(self._mamba_group_ids) + if self._has_mamba: + if len(self._attn_group_ids) != 1: + raise MoRIIOError( + "MoRIIO hybrid READ requires exactly one transferable " + "attention group, got " + f"{len(self._attn_group_ids)} attention groups" + ) + mamba_specs = [ + cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec) + for group_id in self._mamba_group_ids + ] + _validate_mamba_specs(mamba_specs) + if vllm_config.speculative_config is not None: + raise MoRIIOError( + "MoRIIO hybrid READ does not support speculative decoding" + ) + if self.mode != MoRIIOMode.READ: + raise MoRIIOError( + "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " + "mode only; set kv_connector_extra_config.read_mode=true" + ) + # Only "all" mode keeps a state per block position; the other modes keep + # a single running state in the last slot. + self._ssm_state_slots_are_positional = ( + vllm_config.cache_config.mamba_cache_mode == "all" + ) self.block_size = vllm_config.cache_config.block_size + self._max_decode_tail_blocks = ( + _MAX_LOCAL_DECODE_TAIL_BLOCKS + if self._has_mamba + else math.ceil((vllm_config.num_lookahead_tokens + 1) / self.block_size) + ) self.engine_id: EngineId = engine_id - self.mode = get_moriio_mode(self.kv_transfer_config) self._is_hma_required = ( not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager and any( not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ) ) if self._is_hma_required: - # TODO(simondanielsson): support non-sliding window hybrids + # MoRIIO handles sliding-window attention and mamba/KDA groups in + # READ mode. Other hybrid cache specs still fail closed. unsupported = [ type(g.kv_cache_spec).__name__ - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups if not isinstance( - g.kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec) + g.kv_cache_spec, + (FullAttentionSpec, SlidingWindowSpec, MambaSpec), ) ] if unsupported: raise NotImplementedError( "MoRIIO only supports sliding-window hybrid models (e.g. " - "Gemma) under the hybrid KV cache manager. Unsupported KV " + "Gemma) and mamba/KDA groups under the hybrid KV cache " + "manager. Unsupported KV " f"cache group spec(s): {sorted(set(unsupported))}. Pass " "--disable-hybrid-kv-cache-manager to run without HMA." ) @@ -435,7 +549,7 @@ def __init__( (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) if isinstance(g.kv_cache_spec, SlidingWindowSpec) else (0, self.block_size) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ] # add 1 to conservatively account for boundary overlap eg window isn't fully # aligned with blocks. @@ -451,8 +565,7 @@ def __init__( self._full_attn_group_idx = 0 self._full_attn_block_size = self.block_size for gi, group in enumerate(kv_cache_config.kv_cache_groups): - is_full_attn = self.blocks_per_sw[gi] == 0 - if is_full_attn: + if isinstance(group.kv_cache_spec, FullAttentionSpec): self._full_attn_group_idx = gi self._full_attn_block_size = getattr( group.kv_cache_spec, "block_size", self.block_size @@ -533,7 +646,11 @@ def get_exchange_clipped_blocks( - sliding window: keep only the last window size blocks - full attention: keep all blocks """ - if len(block_ids) == 0 or not self._is_hma_required: + if len(block_ids) == 0: + return [] + if len(block_ids) == len(self.kv_cache_config.kv_cache_groups): + block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids) + if not self._is_hma_required: return [list(blocks) for blocks in block_ids] return [ blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks @@ -595,21 +712,75 @@ def get_num_new_matched_tokens( """ if self.is_producer: + # The P-side prompt truncation for hybrid models happens in + # on_new_request, before the scheduler looks the request up in the + # prefix cache. return 0, False - token_ids = request.prompt_token_ids or [] + num_prompt_tokens = request.num_prompt_tokens + num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0) if self.mode == MoRIIOMode.WRITE: - # MoriiO in write mode, no remote prefill + # MoriiO in write mode, no remote prefill. Hybrid models never get + # here: register_kv_caches refuses WRITE mode for them, so there is + # no recurrent-state accounting to do. + return num_external_tokens, True - return len(token_ids) - num_computed_tokens, True + # READ mode always recomputes the last token locally on the decoder. + # + # The second element declares who waits for the KV, not whether it has + # arrived. False keeps the wait on this side: the forward runs in the + # same step, and the connector blocks in wait_for_layer_load (or in + # start_load_kv under full-graph capture) until the KV is really there. + return max(num_external_tokens - 1, 0), False + + def on_new_request(self, request: "Request") -> None: + """Scheduler hook fired when a request is admitted. + + The hybrid P-side truncation must land here rather than in + ``get_num_new_matched_tokens``: the scheduler runs the local prefix + cache lookup first and only then asks the connector, so truncating + there shortens a prompt whose cache hit was already measured against + the full length. On a full hit that leaves ``num_computed_tokens == + num_tokens`` and trips the scheduler's ``num_new_tokens > 0`` assert. + """ + if not (self.is_producer and self._has_mamba): + return + params = request.kv_transfer_params + if params is not None and params.get("do_remote_decode"): + self._truncate_mamba_request_for_prefill(request) - return len(token_ids) - 1 - num_computed_tokens, False + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the last prompt token so the prefiller computes + h(N-1) instead of h(N). + + The decoder recomputes the last token locally to derive h(N) correctly + (see the READ len-1 accounting in get_num_new_matched_tokens). Guarded + by ``_p_side_truncated`` to avoid repeated truncation across a + preemption/reschedule. Mirrors the NIXL scheduler. + """ + params = request.kv_transfer_params + if ( + params is not None + and not params.get("_p_side_truncated") + and request.num_prompt_tokens > 1 + ): + if request.prompt_token_ids is not None: + request.prompt_token_ids.pop() + elif request.prompt_embeds is not None: + request.prompt_embeds = request.prompt_embeds[:-1] + else: + return + + request._all_token_ids.pop() + request.num_prompt_tokens -= 1 + request.max_tokens = 1 + params["_p_side_truncated"] = True def send_notify_block( self, req_id: ReqId, transfer_id: TransferId, - block_notify_list: list[int], + block_notify_list: list, host=None, port=None, ): @@ -739,7 +910,7 @@ def update_state_after_alloc( params.setdefault("transfer_id", transfer_id) request_id = request.request_id self.map_request_id(request_id, transfer_id) - if params.get("do_remote_decode"): + if params.get("do_remote_decode") and self.mode == MoRIIOMode.WRITE: local_block_ids: BlockIds = blocks.get_block_ids() self._reqs_need_save[request.request_id] = (request, local_block_ids) # Snapshot params now so chunked-prefill build_connector_meta @@ -753,24 +924,88 @@ def update_state_after_alloc( # remote_engine_id is returned by the prefill's request_finished. # host/ports come from the request_id (parsed in add_new_req). if "remote_engine_id" in params: + adjusted_remote_block_ids = remote_block_ids if num_external_tokens > 0: - # Get unhashed blocks to pull from remote. - local_block_ids = self.get_exchange_clipped_blocks( - blocks.get_block_ids() - ) - # Partial cache hits can lead to fewer local blocks vs - # remote blocks, but never more - assert all( - len(local) <= len(remote) - for local, remote in zip( - local_block_ids, remote_block_ids, strict=True + if self._has_mamba: + # READ is synchronous, so allocate_slots may hash + # destination blocks before this hook runs. Pair + # the attention and recurrent-state suffixes + # independently. + remote_attn = list(remote_block_ids[0]) + remote_mamba_groups = [ + list(group) for group in remote_block_ids[1:] + ] + attn_block_ids, mamba_block_groups = ( + self.split_block_groups(blocks.get_block_ids()) ) - ) + if len(mamba_block_groups) != len(remote_mamba_groups): + raise MoRIIOError( + "MoRIIO hybrid READ Mamba group " + "count mismatch: " + f"local={len(mamba_block_groups)}, " + f"remote={len(remote_mamba_groups)}" + ) + local_attn = attn_block_ids + local_attn, remote_attn = self._align_read_blocks( + local_attn, + remote_attn, + self._max_decode_tail_blocks, + ) + aligned_local_mamba: list[list[int]] = [] + aligned_remote_mamba: list[list[int]] = [] + for local_group, remote_group in zip( + mamba_block_groups, + remote_mamba_groups, + strict=True, + ): + aligned_local, aligned_remote = ( + self._align_read_blocks( + local_group, remote_group, 1 + ) + ) + aligned_local_mamba.append(aligned_local) + aligned_remote_mamba.append(aligned_remote) + local_block_ids = [local_attn, *aligned_local_mamba] + adjusted_remote_block_ids = [ + remote_attn, + *aligned_remote_mamba, + ] + else: + # Preserve main's per-group HMA representation + # for attention-only hybrid models while + # accounting for the locally recomputed tail. + local_block_ids = self.get_exchange_clipped_blocks( + blocks.get_block_ids() + ) + adjusted_remote_block_ids = [] + aligned_local_block_ids: list[list[int]] = [] + for local_group, remote_group in zip( + local_block_ids, + remote_block_ids, + strict=True, + ): + aligned_local, aligned_remote = ( + self._align_read_blocks( + local_group, + remote_group, + self._max_decode_tail_blocks, + ) + ) + aligned_local_block_ids.append(aligned_local) + adjusted_remote_block_ids.append(aligned_remote) + local_block_ids = aligned_local_block_ids else: - # If remote_blocks and num_external_tokens = 0, we have - # a full prefix cache hit on the D worker. We need to call - # send_notify in _read_blocks to free the memory on the P. - local_block_ids = [[] for _ in remote_block_ids] + if self._has_mamba: + # Attention can be a complete local hit, but the + # recurrent state is never prefix-cacheable. + _, mamba_block_groups = self.split_block_groups( + blocks.get_block_ids() + ) + local_block_ids = [[], *mamba_block_groups] + else: + # Nothing to pull, but retain one empty list per + # cache group so the worker can notify P. + local_block_ids = [[] for _ in remote_block_ids] self._reqs_need_recv[request.request_id] = ( request, @@ -778,7 +1013,10 @@ def update_state_after_alloc( ) # Snapshot params for chunked prefill consumption # in build_connector_meta (see comment above). - self._req_kv_params[request.request_id] = dict(params) + saved_params = dict(params) + if num_external_tokens > 0: + saved_params["remote_block_ids"] = adjusted_remote_block_ids + self._req_kv_params[request.request_id] = saved_params else: logger.warning( "Got invalid KVTransferParams: %s. This " @@ -1003,13 +1241,78 @@ def shutdown(self): logger.warning("Error closing ZMQ socket for path %s: %s", path, e) self.paths.clear() + def split_block_groups( + self, block_ids: list[list[int]] | tuple[list[int], ...] + ) -> tuple[list[int], list[list[int]]]: + """Select transferable attention and Mamba block groups. + + The wire payload stores the attention group first, followed by every + Mamba group in transfer-group order. Outside ``mamba_cache_mode="all"``, + only each group's running-state slot is transferred. + """ + if not block_ids: + return [], [] + if not self._has_mamba: + raise MoRIIOError("Mamba block groups requested for a non-hybrid model") + transfer_block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids) + attn = list(transfer_block_ids[self._attn_group_ids[0]]) + mamba_groups = [ + self._clip_mamba_group(list(transfer_block_ids[group_id])) + for group_id in self._mamba_group_ids + ] + return attn, mamba_groups + + def _clip_mamba_group(self, blocks: list[int]) -> list[int]: + """Keep only the state-bearing slots of one mamba kv cache group.""" + if not blocks: + return blocks + if not self._ssm_state_slots_are_positional: + # Single running state: everything before it is a null placeholder + # or the previous step's superseded state. + blocks = blocks[-1:] + return blocks + + @staticmethod + def _align_read_blocks( + local_block_ids: list[int], + remote_block_ids: list[int], + max_decode_tail_blocks: int, + ) -> tuple[list[int], list[int]]: + """Pair a decode suffix with the corresponding prefill blocks. + + Prefix hits can make the decode's unhashed suffix shorter than the + prefill's full list, so read from the matching remote suffix. The + decoder can instead have one extra tail block for the locally + recomputed final token; that block has no remote source and is not + transferred. + """ + if len(local_block_ids) > len(remote_block_ids): + extra_decode_blocks = len(local_block_ids) - len(remote_block_ids) + if extra_decode_blocks > max_decode_tail_blocks: + raise ValueError( + "decode block ids exceed prefill block ids beyond the " + f"allowed local tail ({max_decode_tail_blocks} blocks): " + f"{len(local_block_ids)} > {len(remote_block_ids)}" + ) + local_block_ids = local_block_ids[: len(remote_block_ids)] + elif len(local_block_ids) < len(remote_block_ids): + remote_block_ids = ( + remote_block_ids[-len(local_block_ids) :] if local_block_ids else [] + ) + return local_block_ids, remote_block_ids + def request_finished( self, request: "Request", - block_ids: BlockIds, + block_ids: BlockIds | list[int], + mamba_block_groups: list[list[int]] | None = None, ) -> tuple[bool, dict[str, Any] | None]: """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. + + ``block_ids`` normally contains all cache groups. Hybrid callers from + ``request_finished_all_groups`` pass a flat attention list and the + recurrent-state groups separately in ``mamba_block_groups``. """ request_id = request.request_id params = request.kv_transfer_params @@ -1042,8 +1345,25 @@ def request_finished( if self.mode == MoRIIOMode.WRITE: self._release_write_prefill_blocks(request.request_id, params) else: - empty_groups: BlockIds = [[] for _ in self.blocks_per_sw] - self._reqs_need_recv[request.request_id] = (request, empty_groups) + # Carry the actually-allocated blocks (incl. the per-request KDA + # mamba slot) instead of []: a do_remote_prefill request finishing + # before update_state_after_alloc still has blocks, and dropping the + # mamba slot here would re-trip the KDA guard in _read_blocks. + if mamba_block_groups: + recv_blocks: BlockIds = [ + list(cast(list[int], block_ids)), + *[list(group) for group in mamba_block_groups], + ] + elif block_ids and isinstance(block_ids[0], int): + recv_blocks = [list(cast(list[int], block_ids))] + else: + grouped_block_ids = cast(BlockIds, block_ids) + recv_blocks = ( + [list(group) for group in grouped_block_ids] + if self._has_mamba + else [[] for _ in self.blocks_per_sw] + ) + self._reqs_need_recv[request.request_id] = (request, recv_blocks) params["do_remote_prefill"] = False return False, None @@ -1054,7 +1374,17 @@ def request_finished( return False, None # On producer - computed_block_ids = self.get_exchange_clipped_blocks(block_ids) + if mamba_block_groups is not None: + computed_block_ids: BlockIds = [ + list(cast(list[int], block_ids)), + *[list(group) for group in mamba_block_groups], + ] + elif self._has_mamba: + raise MoRIIOError("Hybrid request completion requires all KV cache groups") + else: + computed_block_ids = self.get_exchange_clipped_blocks( + cast(BlockIds, block_ids) + ) delay_free_blocks = any(len(group) > 0 for group in computed_block_ids) if delay_free_blocks: @@ -1209,10 +1539,14 @@ def __init__( self.is_producer = self.kv_transfer_config.is_kv_producer self.layer_to_spec = build_layer_to_spec(kv_cache_config) # we transfer by layer, but store by kv cache group - self.layer_to_group: dict[str, int] = { - layer_name: group_idx - for group_idx, group in enumerate(kv_cache_config.kv_cache_groups) - for layer_name in group.layer_names + self.layer_to_group = dict(kv_cache_config.transfer_group_index_by_layer) + self._transfer_layer_names = set(self.layer_to_group) + _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config) + self._num_mamba_transfer_groups = len(mamba_transfer_group_ids) + self._mamba_payload_index_by_layer = { + layer_name: payload_index + for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1) + for layer_name in kv_cache_config.transfer_groups[group_id].layer_names } if self.is_producer: @@ -1353,7 +1687,8 @@ def __init__( # Map of engine_id -> num_blocks. All ranks in the same deployment will # have the same number of blocks. self.dst_num_blocks: dict[EngineId, int] = {} - # In-progress READ transfers: req_id -> {layer_name: status}. + # In-progress READ transfers: req_id -> {layer_name: [statuses]}. + # A KDA layer posts two reads (conv + ssm), so a layer maps to a list. self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict) # Values are (remote_host, remote_notify_port, transfer_id). self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {} @@ -1361,6 +1696,18 @@ def __init__( # Used by _pop_done_transfers to abort transfers whose RDMA # completion is lost, instead of hanging forever. self._recving_transfers_start: dict[str, float] = {} + # Destination attention blocks per in-flight recv. Failed non-HMA reads + # report these to the scheduler's single-group recovery path. + self._recving_local_blocks: dict[ReqId, list[int]] = {} + # Drained by get_block_ids_with_load_errors. + self._invalid_block_ids: set[int] = set() + # This step's posted reads, awaited before the forward when the + # per-layer barrier cannot run. + self._reads_issued_this_step: list = [] + # This step's posted recurrent-state reads. Mamba layers are not + # wrapped by the attention KV-transfer decorator, so wait_for_layer_load + # never fires for them and these must be awaited in start_load_kv. + self._mamba_reads_this_step: list = [] # Track the expiration time of requests that are waiting to be sent. self._reqs_to_send: dict[ReqId, float] = {} @@ -1386,6 +1733,14 @@ def __init__( self.block_window_per_layer: list[int | None] = [] self.use_mla = self.model_config.use_mla + # Hybrid (mamba/KDA) recurrent-state transfer. Populated in + # register_kv_caches when the model has a MambaSpec kv_cache group. + self._has_mamba = False + self._conv_decomp: MambaConvSplitInfo | None = None + self._mamba_offset_templates: dict[str, MambaOffsetTemplate] = {} + # layer_name -> flat session indices (one per registered region: + # 1 for attention, 2 (conv, ssm) for a KDA layer). Built once lazily. + self._region_session_index: dict[str, list[int]] | None = None self.built_session = False self.built_write_session: defaultdict[str, list] = defaultdict(list) backend = get_attn_backend( @@ -1460,23 +1815,39 @@ def schedule_write_blocks( def _get_built_session(self, remote_engine_id): if remote_engine_id not in self.built_write_session: + remote_layer_map = self.layer_name_to_remote_kv_cache_metadata[ + remote_engine_id + ] cur_remote_engine_sessions = [] - for ln, local_meta in self.layer_name_to_local_kv_cache_metadata.items(): - unpacked_local_memory_meta = ( - self.moriio_wrapper.get_unpack_memory_metadata(local_meta[0]) - ) - unpacked_remote_memory_meta = ( - self.moriio_wrapper.get_unpack_memory_metadata( - self.layer_name_to_remote_kv_cache_metadata[remote_engine_id][ - ln - ][0] + for ln, local_metas in self.layer_name_to_local_kv_cache_metadata.items(): + remote_metas = remote_layer_map[ln] + if len(remote_metas) != len(local_metas): + # Sessions are addressed by region index, so a peer that + # registered a different number of regions for this layer + # would not just lose a region to zip(): every later index + # would point into the wrong layer's session. + raise MoRIIOError( + f"peer {remote_engine_id} registered " + f"{len(remote_metas)} region(s) for layer {ln}, this " + f"engine has {len(local_metas)}; the two sides must run " + "the same model and KV cache layout" ) - ) - cur_remote_engine_sessions.append( - self.moriio_wrapper.build_session( - unpacked_local_memory_meta, unpacked_remote_memory_meta + # One session per registered region, in registration order: + # attention layers have a single region; KDA layers have two + # (conv then ssm). This flat layout is indexed via + # _region_session_indices(layer_name). + for local_meta, remote_meta in zip(local_metas, remote_metas): + unpacked_local_memory_meta = ( + self.moriio_wrapper.get_unpack_memory_metadata(local_meta) + ) + unpacked_remote_memory_meta = ( + self.moriio_wrapper.get_unpack_memory_metadata(remote_meta) + ) + cur_remote_engine_sessions.append( + self.moriio_wrapper.build_session( + unpacked_local_memory_meta, unpacked_remote_memory_meta + ) ) - ) self.built_write_session[remote_engine_id] = cur_remote_engine_sessions return self.built_write_session[remote_engine_id], self.remote_moriio_metadata[ remote_engine_id @@ -1788,15 +2159,60 @@ def wait_all_dp(): def _is_mla_cache_layer(self, layer_name: str) -> bool: return is_mla_cache_layer(self.layer_to_spec, layer_name) + def _is_mamba_layer(self, layer_name: str) -> bool: + """True for a hybrid/KDA layer backed by a packed state page.""" + return isinstance(self.layer_to_spec.get(layer_name), MambaSpec) + + @staticmethod + def _contiguous_byte_alias(tensor: torch.Tensor) -> torch.Tensor: + """Return a contiguous uint8 view over ``tensor``'s full byte extent. + + KDA conv/ssm caches are slot-strided views (``stride(0)`` may exceed the + per-slot element count), so they are non-contiguous and cannot be passed + to ``register_torch_tensor``. This aliases the same storage (zero-copy, + identical ``data_ptr``) as a flat uint8 tensor spanning + ``shape[0] * stride(0)`` elements -- the full slot-strided extent that + transfer offsets (``slot * stride(0) * elem``) address. + """ + esz = tensor.element_size() + storage_nbytes = tensor.untyped_storage().nbytes() + offset_bytes = tensor.storage_offset() * esz + extent_bytes = tensor.shape[0] * tensor.stride(0) * esz + # Clamp to the bytes remaining from the view's storage offset. In a + # packed page, the ssm view starts after conv, so its slot-strided extent + # would otherwise run past the backing allocation. + extent_bytes = min(extent_bytes, storage_nbytes - offset_bytes) + assert extent_bytes > 0, ( + "KDA byte alias empty: storage_offset=" + f"{tensor.storage_offset()} * esz={esz} >= " + f"storage_nbytes={storage_nbytes}" + ) + alias = torch.empty(0, dtype=torch.uint8, device=tensor.device) + alias.set_( + tensor.untyped_storage(), + tensor.storage_offset() * esz, + (extent_bytes,), + (1,), + ) + return alias + def _get_layer_transfer_geometry( self, layer_name: str, remote_num_blocks: int | None = None ) -> LayerTransferGeometry: - return get_layer_transfer_geometry( + geometry = get_layer_transfer_geometry( layer_name, self.kv_caches[layer_name], self.layer_to_spec, remote_num_blocks, ) + # get_layer_transfer_geometry also serves mamba layers, whose geometry + # is a different NamedTuple with no block fields. Every caller here is + # on the attention path (mamba layers return before reaching it), so + # narrow the type instead of leaking the union. + assert isinstance(geometry, LayerTransferGeometry), ( + f"layer {layer_name} is a mamba layer; use get_mamba_transfer_geometry" + ) + return geometry def _iter_layer_registration_regions( self, layer_name: str @@ -1853,22 +2269,71 @@ def _remote_layer_mr_offset( def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in moriio.""" + kv_caches = { + layer_name: kv_cache + for layer_name, kv_cache in kv_caches.items() + if layer_name in self._transfer_layer_names + } + if not kv_caches: + raise MoRIIOError("MoRIIO has no transfer-enabled KV cache layers") self.kv_caches = kv_caches # layer name to kv cache + self._has_mamba = any( + self._is_mamba_layer(layer_name) for layer_name in kv_caches + ) + if self._has_mamba and self.mode == MoRIIOMode.WRITE: + # WRITE needs a distinct conv/ssm transfer plan. Reject before + # allocating or registering any regions rather than silently + # serving a zero recurrent state. + raise MoRIIOError( + "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " + "mode only; set kv_connector_extra_config.read_mode=true " + "for hybrid models" + ) + if self._has_mamba: + mamba_specs = [ + cast(MambaSpec, self.layer_to_spec[layer_name]) + for layer_name in kv_caches + if self._is_mamba_layer(layer_name) + ] + mamba_spec = _validate_mamba_specs(mamba_specs) + assert mamba_spec is not None + + from vllm.model_executor.layers.mamba.mamba_utils import ( + is_conv_state_dim_first, + ) + + if not is_conv_state_dim_first(): + raise MoRIIOError( + "MoRIIO KDA conv transfer requires the DS (dim-first) " + "state layout; set VLLM_SSM_CONV_STATE_LAYOUT=DS" + ) + self._conv_decomp = derive_mamba_conv_split(mamba_spec, self.world_size) + # KDA layers use packed pages with different geometry; only record + # attention tensor shapes here. self.kv_cache_shapes = { - layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items() + layer_name: kv_cache.shape + for layer_name, kv_cache in kv_caches.items() + if not self._is_mamba_layer(layer_name) } + # Geometry selection must pick an attention layer (never a KDA tuple): + # prefer a standard 5-D K/V layer, else the first attention layer. + attn_items = [ + (layer_name, kv_cache) + for layer_name, kv_cache in kv_caches.items() + if not self._is_mamba_layer(layer_name) + ] first_layer_name, first_kv_cache = next( ( (layer_name, kv_cache) - for layer_name, kv_cache in kv_caches.items() + for layer_name, kv_cache in attn_items if ( not self._is_mla_cache_layer(layer_name) and len(kv_cache.shape) == 5 and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2) ) ), - next(iter(kv_caches.items())), + attn_items[0] if attn_items else next(iter(kv_caches.items())), ) kv_elem_size = first_kv_cache.element_size() @@ -1915,6 +2380,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) for layer_name in kv_caches: self.layer_base_addr_index[layer_name] = base_addr_idx + if self._is_mamba_layer(layer_name): + # KDA layer: two whole-tensor regions (conv, ssm); no block-size + # geometry (recurrent state is per-slot, not paged). + for cache, region_len in self._iter_layer_registration_regions( + layer_name + ): + base_addr = cache.data_ptr() + caches_data.append((base_addr, region_len, cache.device.index, "")) + kv_caches_base_addr.append(base_addr) + base_addr_idx += 1 + continue geometry = self._get_layer_transfer_geometry(layer_name) if hma_enabled: if geometry.block_len != self.block_len: @@ -1941,31 +2417,87 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): kv_caches_base_addr.append(base_addr) base_addr_idx += 1 - shared_backing = len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1 + shared_backing = ( + all(isinstance(t, torch.Tensor) for t in kv_caches.values()) + and len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1 + ) shared_mr = None if shared_backing: reg_tensor, self.kv_layer_mr_offset = self._build_shared_kv_mr(kv_caches) shared_mr = self.moriio_wrapper.register_local_tensor(reg_tensor) + # Say it out loud when the groups disagree. This is expected on a hybrid + # model and used to be fatal, so a run that silently had one page size + # and a run that has three should not look identical in the log. + distinct_block_lens = sorted(set(self.block_lens.values())) + if len(distinct_block_lens) > 1: + logger.info( + "MoRIIO registered %d attention layers with %d distinct block " + "lengths (bytes): %s; transfer sizes and strides are per layer, " + "and %d is the page advertised to the peer.", + len(self.block_lens), + len(distinct_block_lens), + distinct_block_lens, + self.block_len, + ) + for layer_name, kv_cache in kv_caches.items(): if layer_name not in self.layer_name_to_local_kv_cache_metadata: self.layer_name_to_local_kv_cache_metadata[layer_name] = [] - moriio_mem_metadata = shared_mr or ( - self.moriio_wrapper.register_local_tensor(kv_cache) - ) - self.layer_name_to_local_kv_cache_metadata[layer_name].append( - moriio_mem_metadata - ) - - self.local_kv_cache_size.append( - kv_cache.nelement() * kv_cache.element_size() - ) + if self._is_mamba_layer(layer_name): + # Register conv and ssm as two whole-tensor regions, appending + # both metadata blobs (order: conv then ssm) so the session + # layout matches iter_layer_registration_regions. conv/ssm are + # slot-strided views and are not contiguous, which + # register_torch_tensor rejects; register a contiguous uint8 + # alias over each tensor's full byte extent instead (zero-copy: + # same storage and data_ptr). + conv, ssm = kda_conv_ssm( + kv_cache, cast(MambaSpec, self.layer_to_spec[layer_name]) + ) + logger.debug( + "MoRIIO KDA register %s: conv shape=%s stride=%s contig=%s " + "storage_nbytes=%s ; ssm shape=%s stride=%s contig=%s " + "storage_nbytes=%s", + layer_name, + tuple(conv.shape), + tuple(conv.stride()), + conv.is_contiguous(), + conv.untyped_storage().nbytes(), + tuple(ssm.shape), + tuple(ssm.stride()), + ssm.is_contiguous(), + ssm.untyped_storage().nbytes(), + ) + for tensor in (conv, ssm): + alias = self._contiguous_byte_alias(tensor) + meta = self.moriio_wrapper.register_local_tensor(alias) + self.layer_name_to_local_kv_cache_metadata[layer_name].append(meta) + self.local_kv_cache_size.append(alias.numel()) + else: + moriio_mem_metadata = shared_mr or ( + self.moriio_wrapper.register_local_tensor(kv_cache) + ) + self.layer_name_to_local_kv_cache_metadata[layer_name].append( + moriio_mem_metadata + ) + self.local_kv_cache_size.append( + kv_cache.nelement() * kv_cache.element_size() + ) self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr self.num_regions = len(caches_data) self.num_layers = len(self.kv_caches.keys()) + # Report the prevalidated KDA decomposition after registration. + if self._has_mamba: + assert self._conv_decomp is not None + logger.info( + "MoRIIO registered %d KDA (conv+ssm) layer(s); ssm sizes=%s", + sum(1 for ln in kv_caches if self._is_mamba_layer(ln)), + self._conv_decomp.ssm_sizes, + ) # Optimization for models with local attention (Llama 4) if self.vllm_config.model_config.hf_config.model_type == "llama4": from transformers import Llama4TextConfig @@ -2101,49 +2633,49 @@ def wait_for_layer_load(self, layer_name: str) -> None: if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: return - deadline = time.monotonic() + self.moriio_config.transfer_timeout - while True: - with self.moriio_wrapper.lock: - pending = [ - status_by_layer[layer_name] - for status_by_layer in self._recving_transfers.values() - if layer_name in status_by_layer - ] - - if not pending: - return - - still_running = False - for status in pending: - # A failed read is dropped in _pop_done_transfers. - if status.Succeeded() or status.Failed(): - continue - still_running = True + with self.moriio_wrapper.lock: + pending = [ + status + for status_by_layer in self._recving_transfers.values() + for status in status_by_layer.get(layer_name, ()) + ] - if not still_running: - return + if not pending: + return - if time.monotonic() > deadline: - logger.warning( - "MoRIIO READ barrier timed out for layer %s; proceeding " - "(request dropped via get_finished).", + try: + self.moriio_wrapper.waiting_for_transfer_complete(pending) + except TransferError: + if self._has_mamba: + logger.exception( + "MoRIIO hybrid READ failed before layer %s; aborting the " + "forward because HMA load-error recovery is unsupported.", layer_name, ) - return - - time.sleep(0.001) + raise + logger.warning( + "MoRIIO READ barrier did not complete for layer %s; proceeding " + "to transfer cleanup.", + layer_name, + exc_info=True, + ) def _pop_done_transfers(self) -> set[str]: done_req_ids: set[str] = set() - _xfer_timeout = int(os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S", "120")) + _xfer_timeout = self.moriio_config.recv_abort_timeout with self.moriio_wrapper.lock: to_remove = [] for req_id, status_by_layer in self._recving_transfers.items(): - statuses = list(status_by_layer.values()) - failed_status = next( - (status for status in statuses if status.Failed()), None - ) - if statuses and all(status.Succeeded() for status in statuses): + # Every layer's read must land before the request is complete, + # so the whole batch is polled rather than just the newest + # status. + statuses = [ + status + for layer_statuses in status_by_layer.values() + for status in layer_statuses + ] + state = self.moriio_wrapper.poll_transfer_batch(statuses) + if statuses and state is TransferBatchState.DONE: host, port, xfer_id = self._recving_transfers_callback_addr[req_id] done_req_ids.add(xfer_id) self.moriio_wrapper.send_notify( @@ -2154,15 +2686,19 @@ def _pop_done_transfers(self) -> set[str]: message_fields={"consumer_tp_size": self.world_size}, ) to_remove.append(req_id) - elif failed_status is not None: + elif state is TransferBatchState.FAILED: + failed_status = next( + (status for status in statuses if status.Failed()), None + ) logger.error( "RDMA transfer failed for request %s: %s (code=%s). " - "Notifying prefill to free blocks; request will be " - "aborted by timeout.", + "Notifying prefill to free blocks and recording the " + "load error where scheduler recovery is supported.", req_id, - failed_status.Message(), - failed_status.Code(), + failed_status.Message() if failed_status else "unknown", + failed_status.Code() if failed_status else "unknown", ) + self._record_failed_recv(req_id) host, port, xfer_id = self._recving_transfers_callback_addr[req_id] try: self.moriio_wrapper.send_notify( @@ -2178,8 +2714,9 @@ def _pop_done_transfers(self) -> set[str]: req_id, ) to_remove.append(req_id) - # Do NOT add to done_req_ids: decode KV cache is incomplete. - # The request will expire via the normal request timeout. + # Deliberately not in done_req_ids: the decode KV is + # incomplete, so this is a load error rather than a + # completion. elif req_id in self._recving_transfers_start: # Abort still-in-flight transfers that exceed the # configured deadline. Otherwise a lost RDMA @@ -2189,19 +2726,84 @@ def _pop_done_transfers(self) -> set[str]: if _age > _xfer_timeout: logger.error( "RDMA read TIMED OUT for req %s after %.1fs " - "(VLLM_MORIIO_TRANSFER_TIMEOUT_S=%d)", + "(kv_connector_extra_config.recv_abort_timeout=%.0f)", req_id, _age, _xfer_timeout, ) + self._record_failed_recv(req_id) to_remove.append(req_id) for req_id in to_remove: del self._recving_transfers[req_id] del self._recving_transfers_callback_addr[req_id] self._recving_transfers_start.pop(req_id, None) + self._recving_local_blocks.pop(req_id, None) return done_req_ids + def _await_reads_issued_this_step(self) -> None: + """Step-level READ barrier, for the reads no per-layer barrier covers. + + load_kv_async=False promises the KV is in place before the forward runs. + wait_for_layer_load keeps that promise per layer, letting early layers + overlap later transfers, but it only fires for layers that go through + the attention KV-transfer decorator. Two cases fall through it: + + * mamba/KDA layers have no such hook at all, so their conv+ssm reads + would otherwise be consumed while still in flight; + * under CUDAGraphMode.FULL a host-side blocking wait is illegal + anywhere in the graph, so every read is awaited here instead. + + start_load_kv runs outside the graph, so blocking here is safe. + Attention reads keep their per-layer overlap in the non-FULL case. + + Attention-only failures are reported via + get_block_ids_with_load_errors. Hybrid failures are re-raised because + the scheduler's recovery path does not support multiple cache groups; + this fails the step closed instead of running forward with incomplete + recurrent state. + """ + all_statuses = self._reads_issued_this_step + self._reads_issued_this_step = [] + mamba_statuses = self._mamba_reads_this_step + self._mamba_reads_this_step = [] + if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: + statuses = all_statuses + else: + statuses = mamba_statuses + if not statuses: + return + try: + self.moriio_wrapper.waiting_for_transfer_complete(statuses) + except TransferError: + if self._has_mamba: + logger.exception( + "MoRIIO hybrid READ failed before the forward; aborting " + "because HMA load-error recovery is unsupported." + ) + raise + logger.exception( + "MoRIIO reads did not complete before the forward; proceeding " + "to transfer cleanup." + ) + + def _record_failed_recv(self, req_id: ReqId) -> None: + """Hand a failed attention-only read to scheduler recovery. + + The scheduler's invalid-block recovery currently assumes exactly one + KV cache group. Reporting IDs for a hybrid request would crash it while + unpacking the request's groups, so match NIXL and keep HMA failures out + of this API until scheduler recovery carries group identity. + """ + if self._has_mamba: + return + self._invalid_block_ids.update(self._recving_local_blocks.get(req_id, ())) + + def get_block_ids_with_load_errors(self) -> set[int]: + """Drain the blocks whose read failed, for the scheduler to recompute.""" + invalid, self._invalid_block_ids = self._invalid_block_ids, set() + return invalid + def save_kv_layer( self, metadata: MoRIIOConnectorMetadata, @@ -2462,11 +3064,12 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): continue - # Handshake already completed, start async read xfer. + # Handshake already completed, start the read xfer. self._read_blocks_for_req(req_id, meta) # Start transfers for requests whose handshakes have now finished. if remote_engine_id is None and not wait_handshake_readd_req: + self._await_reads_issued_this_step() return _deadline = time.monotonic() + self.moriio_config.transfer_timeout while True: @@ -2493,6 +3096,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): else: break + self._await_reads_issued_this_step() self._reqs_to_send.update(metadata.reqs_to_send) def wait_for_save(self, metadata: MoRIIOConnectorMetadata): @@ -2717,6 +3321,163 @@ def _compute_block_transfer_offsets( remote = [x + remote_off for x in remote] return local, remote, sizes + def _region_session_indices(self, layer_name: str) -> list[int]: + """Flat session indices for a layer's registered regions. + + Sessions are built one per region in registration order (see + _get_built_session), so an attention layer maps to a single index and a + KDA layer to two indices ([conv, ssm]). Built once and cached. + """ + if self._region_session_index is None: + mapping: dict[str, list[int]] = {} + idx = 0 + for ln, metas in self.layer_name_to_local_kv_cache_metadata.items(): + mapping[ln] = list(range(idx, idx + len(metas))) + idx += len(metas) + self._region_session_index = mapping + return self._region_session_index[layer_name] + + def _mamba_tp_ratio(self, remote_tp_size: int | None) -> int: + """Validate the homogeneous-TP KDA contract.""" + remote = self.world_size if remote_tp_size is None else remote_tp_size + if remote != self.world_size: + raise NotImplementedError( + "heterogeneous-TP mamba/KDA transfer is not supported yet " + f"(local_tp={self.world_size}, remote_tp={remote}); use equal " + "prefill/decode TP for hybrid models" + ) + return 1 + + def _mamba_blocks_for_layer( + self, + layer_name: str, + block_groups: BlockIds, + ) -> list[int]: + payload_index = self._mamba_payload_index_by_layer.get(layer_name) + if payload_index is None or payload_index >= len(block_groups): + raise MoRIIOError( + f"KDA layer {layer_name} has no block group in the transfer payload" + ) + return list(block_groups[payload_index]) + + def _compute_mamba_transfer_offsets( + self, + layer_name: str, + local_slots: list[int], + remote_slots: list[int], + remote_tp_size: int | None = None, + ) -> tuple[list[int], list[int], list[int], int]: + """Compute (local, remote, sizes, n_conv) for a KDA layer. + + Entries [:n_conv] address the conv region/session; entries [n_conv:] + address the ssm region/session. + """ + assert self._conv_decomp is not None, "KDA decomposition not derived" + conv, ssm = kda_conv_ssm( + self.kv_caches[layer_name], + cast(MambaSpec, self.layer_to_spec[layer_name]), + ) + tp_ratio = self._mamba_tp_ratio(remote_tp_size) + # Cache the slot-independent geometry per layer; only the per-request + # slot bases vary, and heterogeneous TP is rejected above. + template = self._mamba_offset_templates.get(layer_name) + if template is None: + template = build_mamba_offset_template( + conv, + ssm, + self._conv_decomp, + tp_ratio, + ) + self._mamba_offset_templates[layer_name] = template + local_offs, remote_offs, sizes = apply_mamba_offset_template( + template, local_slots, remote_slots + ) + n_conv = compute_mamba_conv_split_count(local_slots, self._conv_decomp) + return local_offs, remote_offs, sizes, n_conv + + def _post_mamba_reads( + self, + layer_name: str, + sessions: list, + region_sessions: list[int], + local_slots: list[int], + remote_slots: list[int], + remote_tp_size: int, + request_id: str, + deadline: float, + ) -> list: + """Post one KDA layer's conv and ssm reads; return their statuses. + + The layer owns two registered regions, so the offset list is split at + the conv/ssm boundary and each half goes to its own session. A half + with nothing to move is skipped rather than posted empty. + """ + local, remote, sizes, n_conv = self._compute_mamba_transfer_offsets( + layer_name, + local_slots, + remote_slots, + remote_tp_size=remote_tp_size, + ) + statuses: list = [] + for sess_idx, sl in ( + (region_sessions[0], slice(0, n_conv)), + (region_sessions[1], slice(n_conv, None)), + ): + region_sizes = sizes[sl] + if not region_sizes: + continue + statuses.append( + self._post_read_with_backoff( + sessions[sess_idx], + region_sizes, + local[sl], + remote[sl], + request_id, + layer_name, + deadline, + ) + ) + return statuses + + def _post_read_with_backoff( + self, + session, + sizes: list[int], + local_offsets: list[int], + remote_offsets: list[int], + request_id: str, + layer_name: str, + deadline: float, + ): + """Post one RDMA READ, backing off on transient SQ-full rejection. + + read_remote_data posts synchronously, so a send-queue-full rejection is + a Failed() status on return; a separate CQ-poll thread drains + completions and frees SQ depth, so we back off and re-post until + transfer_timeout, then store the failed status (get_finished notifies + prefill and drops the request non-fatally). + """ + _backoff = _SQ_FULL_BACKOFF_INITIAL_S + while True: + transfer_status = self.moriio_wrapper.read_remote_data( + sizes, local_offsets, remote_offsets, session + ) + if not self._is_sq_full_status(transfer_status): + break + if time.monotonic() > deadline: + logger.warning( + "MoRIIO READ send queue stayed full past " + "transfer_timeout for req %s layer %s; storing failed " + "status (get_finished notifies prefill and drops the " + "request). Raise qp_per_transfer if frequent.", + request_id, + layer_name, + ) + break + time.sleep(_backoff) + _backoff = min(_backoff * 2, _SQ_FULL_BACKOFF_MAX_S) + return transfer_status + @staticmethod def _is_sq_full_status(status) -> bool: """True if a MoRIIO transfer status is a transient RDMA send-queue-full @@ -2766,6 +3527,44 @@ def _read_blocks( if chosen_tp is not None else self._remote_tp_rank(remote_tp_size) ) + if not any(local_block_ids): + # A request aborted before decode allocation still carries the + # producer's block table. There is no local destination or forward, + # so release the producer blocks without building sessions or + # posting RDMA reads. + self.moriio_wrapper.send_notify( + transfer_id, + remote_host, + str( + remote_notify_port + + get_port_offset( + int(remote_dp_rank), + eff_tp, + remote_tp_size, + ) + ), + message_type="release", + message_fields={"consumer_tp_size": self.world_size}, + ) + return + + if self._has_mamba: + expected_groups = 1 + self._num_mamba_transfer_groups + if ( + len(local_block_ids) != expected_groups + or len(remote_block_ids) != expected_groups + ): + raise MoRIIOError( + "MoRIIO hybrid READ block group count mismatch: " + f"expected={expected_groups}, local={len(local_block_ids)}, " + f"remote={len(remote_block_ids)}" + ) + local_attn = list(local_block_ids[0]) + remote_attn = list(remote_block_ids[0]) + self._mamba_tp_ratio(remote_tp_size) + else: + local_attn = remote_attn = [] + if flexible: remote_dp_engine_id = self.get_engine_name_with_dp_tp( dst_engine_id, int(remote_dp_rank), eff_tp @@ -2778,52 +3577,89 @@ def _read_blocks( # SQ-full backpressure deadline, shared across this request's layers. _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout + # TODO: apply multi-session batch-read when MoRIIO supports it. for layer_name in self.layer_name_to_local_kv_cache_metadata: - sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( - layer_name - ) - group_idx = self.layer_to_group[layer_name] - offs = self._compute_block_transfer_offsets( - layer_name, - local_block_ids[group_idx], - remote_block_ids[group_idx], - remote_moriio_meta, - remote_tp_size=remote_tp_size, - remote_engine_id=remote_dp_engine_id, - ) - # TODO : apply multi-session batch-read when moriio support it - # - # SQ-full backpressure: read_remote_data posts the RDMA READ - # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes - # back as a Failed() status right here. A SEPARATE CQ-poll thread - # drains completions and frees SQ depth, so back off and RE-POST - # rather than let a transient rejection abort the request. No - # self-deadlock (the drain is off-thread); the reserve is - # all-or-nothing (nothing posted on a rejected attempt). Bounded by - # transfer_timeout; on sustained overload store the failed status and - # let get_finished handle it non-fatally (notify prefill + drop). - _backoff = 0.001 - while True: - transfer_status = self.moriio_wrapper.read_remote_data( - offs[2], offs[0], offs[1], sessions[sess_idx] + region_sessions = self._region_session_indices(layer_name) + statuses = [] + if self._is_mamba_layer(layer_name): + local_mamba = self._mamba_blocks_for_layer(layer_name, local_block_ids) + remote_mamba = self._mamba_blocks_for_layer( + layer_name, remote_block_ids ) - if not self._is_sq_full_status(transfer_status): - break - if time.monotonic() > _sq_deadline: - logger.warning( - "MoRIIO READ send queue stayed full past " - "transfer_timeout for req %s layer %s; storing failed " - "status (get_finished notifies prefill and drops the " - "request). Raise qp_per_transfer if frequent.", + # KDA layer: pull conv (sub-projection offsets) and ssm at the + # request's recurrent-state slot. Two regions -> two sessions. + if not local_mamba or not remote_mamba: + if not local_attn and not remote_attn: + # Whole-request no-op (e.g. full prefix cache hit): + # nothing to transfer for any layer. + continue + # A KDA layer must have its recurrent-state slot in the + # block-id tuple; missing it (while attention blocks are + # present) means the mamba KV-cache group is absent -> bug. + raise MoRIIOError( + f"KDA layer {layer_name}: missing mamba recurrent-state " + f"block ids (local={local_mamba}, remote={remote_mamba}) " + f"with attention blocks present for request {request_id}; " + "the mamba KV-cache group is absent from the transfer " + "block-id tuple" + ) + statuses = self._post_mamba_reads( + layer_name, + sessions, + region_sessions, + local_mamba, + remote_mamba, + remote_tp_size, + request_id, + _sq_deadline, + ) + else: + if self._has_mamba: + local_layer_blocks = local_attn + remote_layer_blocks = remote_attn + else: + group_idx = self.layer_to_group[layer_name] + local_layer_blocks = local_block_ids[group_idx] + remote_layer_blocks = remote_block_ids[group_idx] + offs = self._compute_block_transfer_offsets( + layer_name, + local_layer_blocks, + remote_layer_blocks, + remote_moriio_meta, + remote_tp_size=remote_tp_size, + remote_engine_id=remote_dp_engine_id, + ) + statuses.append( + self._post_read_with_backoff( + sessions[region_sessions[0]], + offs[2], + offs[0], + offs[1], request_id, layer_name, + _sq_deadline, ) - break - time.sleep(_backoff) - _backoff = min(_backoff * 2, 0.05) + ) with self.moriio_wrapper.lock: - self._recving_transfers[request_id][layer_name] = transfer_status + self._recving_transfers[request_id][layer_name] = statuses self._recving_transfers_start.setdefault(request_id, time.monotonic()) + # Destination blocks, kept for _record_failed_recv. + if self._has_mamba: + failed_recv_blocks = local_attn + else: + failed_recv_blocks = [ + block_id for group in local_block_ids for block_id in group + ] + self._recving_local_blocks.setdefault( + request_id, list(failed_recv_blocks) + ) + # Awaited before the forward when the per-layer barrier cannot + # run (see _await_reads_issued_this_step). + self._reads_issued_this_step.extend(statuses) + if self._is_mamba_layer(layer_name): + # No per-layer barrier exists for mamba layers, so these are + # always awaited in start_load_kv. + self._mamba_reads_this_step.extend(statuses) self._recving_transfers_callback_addr[request_id] = ( remote_host, str( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 5c1330521dd6..a0c6ec1f6d7f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -29,6 +29,7 @@ MoRIIOError, MoRIIOTransferAck, RemoteAllocInfo, + TransferBatchState, TransferError, TransferId, WriteTask, @@ -58,7 +59,6 @@ except ImportError: logger.error("MoRIIO is not available") - """Write task execution logic for MoRIIO connector.""" @@ -395,9 +395,9 @@ def _prepare_transfer_plan( ) request_info.transfer_offsets[key] = offsets - # Get session index - layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys()) - sess_idx = layer_names.index(task.layer_name) + # One session per registered region; attention layers map to a single + # index (see MoRIIOConnectorWorker._region_session_indices). + sess_idx = self.worker._region_session_indices(task.layer_name)[0] local_off, remote_off, sizes = offsets @@ -662,6 +662,25 @@ def write_remote_data_single( ) return transfer_status + def poll_transfer_batch(self, transfer_statuses: list[Any]) -> TransferBatchState: + """Non-blocking verdict over every status of one request. + + Judging a request by its newest status alone calls it done while an + earlier layer is unfinished or already failed; failures do land on + layers other than the last, because _post_read_with_backoff returns a + failed status when the send queue never drains. + + This stays a Python scan so it does not drive the backend's progress + callback from the engine thread alongside MoRIIO's own poller. + """ + state = TransferBatchState.DONE + for status in transfer_statuses: + if status.Failed(): + return TransferBatchState.FAILED + if not status.Succeeded(): + state = TransferBatchState.PENDING + return state + def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None): if transfer_statuses is None: with self.lock: @@ -673,6 +692,15 @@ def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = No if not transfers_to_wait: return + self._poll_transfers_until_done(transfers_to_wait) + + def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None: + """Fallback for mori builds without the batched wait. + + Gives failure precedence over pending: once any status has failed the + request is lost, so waiting out the remaining ones would only delay the + barrier by up to transfer_timeout. + """ timeout = self._transfer_timeout deadline = time.monotonic() + timeout remaining = list(transfers_to_wait) @@ -698,6 +726,8 @@ def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = No ) else: still_waiting.append(status) + if errors: + break remaining = still_waiting if remaining: time.sleep(0.001) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 0ee2721139e5..b00692dc5188 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -2,14 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Mapping -from typing import NamedTuple +from typing import NamedTuple, cast import torch +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( + MambaConvSplitInfo, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, KVCacheSpec, + MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, UniformTypeKVCacheSpecs, @@ -29,6 +33,85 @@ class LayerTransferGeometry(NamedTuple): split_kv_regions: bool +class MambaTransferGeometry(NamedTuple): + """Geometry of a hybrid (mamba/KDA) layer's recurrent state. + + A KDA layer's packed page is exposed as ``conv_state`` and ``ssm_state`` + views. Each is slot-strided, so a slot's bytes live at + ``slot * slot_stride * element_size`` and span ``slot_bytes``. The two + tensors are registered as two separate MoRIIO regions per layer. + """ + + conv_slot_stride: int + ssm_slot_stride: int + ssm_slot_bytes: int + conv_element_size: int + ssm_element_size: int + conv_region_len: int + ssm_region_len: int + + +def get_mamba_transfer_geometry( + conv: torch.Tensor, + ssm: torch.Tensor, +) -> MambaTransferGeometry: + """Derive the slot-strided geometry of a KDA layer's conv + ssm state. + + ``conv``/``ssm`` are zero-copy views over the layer's packed page. Slot byte + offset is ``slot * slot_stride * element_size`` and each + slot spans ``subtensor[0].numel() * element_size`` bytes; the whole tensor + is registered as one region (``*_region_len`` bytes). + """ + if conv.ndim == 0 or ssm.ndim == 0 or conv.shape[0] == 0 or ssm.shape[0] == 0: + raise ValueError("KDA conv/ssm caches must contain at least one slot") + if conv.shape[0] != ssm.shape[0]: + raise ValueError( + f"KDA conv/ssm slot count mismatch: {conv.shape[0]} != {ssm.shape[0]}" + ) + conv_slot_bytes = conv[0].numel() * conv.element_size() + ssm_slot_bytes = ssm[0].numel() * ssm.element_size() + conv_available_bytes = ( + conv.untyped_storage().nbytes() - conv.storage_offset() * conv.element_size() + ) + ssm_available_bytes = ( + ssm.untyped_storage().nbytes() - ssm.storage_offset() * ssm.element_size() + ) + if conv.shape[0]: + conv_required_bytes = (conv.shape[0] - 1) * conv.stride( + 0 + ) * conv.element_size() + conv_slot_bytes + ssm_required_bytes = (ssm.shape[0] - 1) * ssm.stride( + 0 + ) * ssm.element_size() + ssm_slot_bytes + if conv_required_bytes > conv_available_bytes: + raise ValueError("KDA conv slots exceed the registered storage extent") + if ssm_required_bytes > ssm_available_bytes: + raise ValueError("KDA ssm slots exceed the registered storage extent") + + return MambaTransferGeometry( + conv_slot_stride=conv.stride(0), + ssm_slot_stride=ssm.stride(0), + ssm_slot_bytes=ssm_slot_bytes, + conv_element_size=conv.element_size(), + ssm_element_size=ssm.element_size(), + # Registered region must span the full slot-strided extent: stride(0) + # may exceed the per-slot element count, so numel() would under-size the + # region and truncate the highest-slot transfers. shape[0]*stride(0) + # covers every byte a ``slot * slot_stride * elem`` offset can address. + # Cap at the bytes remaining from each view's storage offset. A packed + # conv+ssm page places ssm at a non-zero intra-buffer offset, + # so do not register past the end of the shared backing allocation. + conv_region_len=min( + conv.shape[0] * conv.stride(0) * conv.element_size(), + conv_available_bytes, + ), + ssm_region_len=min( + ssm.shape[0] * ssm.stride(0) * ssm.element_size(), + ssm_available_bytes, + ), + ) + + def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]: layer_to_spec: dict[str, KVCacheSpec] = {} for group in kv_cache_config.kv_cache_groups: @@ -99,10 +182,10 @@ def get_layer_transfer_geometry( layer_to_spec: Mapping[str, KVCacheSpec], remote_num_blocks: int | None = None, ) -> LayerTransferGeometry: + spec = layer_to_spec[layer_name] shape = kv_cache.shape stride = kv_cache.stride() element_size = kv_cache.element_size() - spec = layer_to_spec[layer_name] is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name) if not is_mla_cache and len(shape) == 5 and shape[0] == 2: @@ -238,12 +321,60 @@ def get_layer_transfer_geometry( ) +def kda_conv_ssm( + kv_cache: torch.Tensor, + spec: MambaSpec, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return zero-copy conv and SSM views over a packed two-state page.""" + from math import prod + + from vllm.utils.torch_utils import get_dtype_size + + if len(spec.shapes) != 2 or len(spec.dtypes) != 2: + raise ValueError( + "MoRIIO KDA transfer requires exactly two states (conv and SSM)" + ) + if kv_cache.element_size() != 1: + raise ValueError("MoRIIO KDA packed cache must use a byte-sized dtype") + if kv_cache.ndim != 4 or kv_cache.shape[1:3] != (1, 1) or kv_cache.stride(3) != 1: + raise ValueError( + "MoRIIO KDA packed cache must have shape [B, 1, 1, C] with " + "a contiguous byte dimension" + ) + # Mirror MambaBase.bind_kv_cache: blocks may be strided within a shared + # multi-layer allocation, while every individual byte page is contiguous. + pages = kv_cache.squeeze(dim=(1, 2)) + states: list[torch.Tensor] = [] + offset = 0 + for shp, dt in zip(spec.shapes, spec.dtypes, strict=True): + nbytes = prod(shp) * get_dtype_size(dt) + if offset + nbytes > pages.shape[1]: + raise ValueError( + "MoRIIO KDA state layout exceeds the packed cache page size" + ) + state = pages[:, offset : offset + nbytes].view(dt) + states.append(state.view(-1, *shp)) + offset += nbytes + return states[0], states[1] + + def iter_layer_registration_regions( layer_name: str, kv_cache: torch.Tensor, layer_to_spec: Mapping[str, KVCacheSpec], ) -> list[tuple[torch.Tensor, int]]: - geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec) + spec = layer_to_spec[layer_name] + if isinstance(spec, MambaSpec): + # KDA layer: register the conv and ssm tensors as two whole-tensor + # regions (the conv 3-subprojection split is expressed as transfer + # offset triples, not extra registrations). + conv, ssm = kda_conv_ssm(kv_cache, spec) + geom = get_mamba_transfer_geometry(conv, ssm) + return [(conv, geom.conv_region_len), (ssm, geom.ssm_region_len)] + geometry = cast( + LayerTransferGeometry, + get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec), + ) region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len if geometry.regions_per_block == 1: # With padded or interleaved pages the block stride exceeds the @@ -311,9 +442,36 @@ def compute_block_transfer_offsets( "local_block_ids longer than remote_block_ids: " f"{len(local_block_ids)} > {len(remote_block_ids)}" ) - geometry = get_layer_transfer_geometry( - layer_name, kv_cache, layer_to_spec, remote_num_blocks + geometry = cast( + LayerTransferGeometry, + get_layer_transfer_geometry( + layer_name, kv_cache, layer_to_spec, remote_num_blocks + ), + ) + invalid_local = next( + ( + block_id + for block_id in local_block_ids + if not 0 <= block_id < geometry.num_blocks + ), + None, + ) + if invalid_local is not None: + raise ValueError( + f"local block id {invalid_local} is outside [0, {geometry.num_blocks})" + ) + invalid_remote = next( + ( + block_id + for block_id in remote_block_ids + if not 0 <= block_id < remote_num_blocks + ), + None, ) + if invalid_remote is not None: + raise ValueError( + f"remote block id {invalid_remote} is outside [0, {remote_num_blocks})" + ) element_size = kv_cache.element_size() transfer_size_byte = geometry.block_len per_block = geometry.transfers_per_block @@ -339,3 +497,125 @@ def compute_block_transfer_offsets( w += 1 return merge_fn(offset_local, offset_remote, sizes) + + +class MambaOffsetTemplate(NamedTuple): + """Slot-independent conv+ssm offset decomposition for a KDA layer. + + Homogeneous-TP geometry (conv/ssm slot strides, conv sub-projection + offsets, ssm per-slot size) is identical across requests and across the + homogeneous GDN layers, so it is computed once and reused; only the + per-request slot bases vary (see ``apply_mamba_offset_template``). + """ + + num_slots: int + conv_slot_stride_bytes: int + ssm_slot_stride_bytes: int + ssm_read_bytes: int + conv_subprojs: tuple[tuple[int, int, int], ...] + + +def build_mamba_offset_template( + conv: torch.Tensor, + ssm: torch.Tensor, + split_info: MambaConvSplitInfo, + tp_ratio: int, +) -> MambaOffsetTemplate: + """Precompute the slot-independent conv+ssm offset decomposition. + + Only homogeneous TP (``tp_ratio == 1``) is supported; heterogeneous TP is + gated with ``NotImplementedError`` (see the design doc). The returned + template is a pure function of the layer geometry and ``split_info``, so it + can be cached and reused across requests without recomputation. + """ + if tp_ratio != 1: + # Heterogeneous-TP recurrent-state transfer needs the remote page's + # slot stride (which differs from the local page) and, for + # P_TP > D_TP, a multi-rank gather of conv/ssm slices. Both are + # follow-ups. The mamba path bypasses the attention KV-head validator, + # so fail loudly here instead of silently transferring wrong bytes. + raise NotImplementedError( + "heterogeneous-TP mamba/KDA transfer is not supported yet " + f"(tp_ratio={tp_ratio}); use equal prefill/decode TP for " + "hybrid models" + ) + geom = get_mamba_transfer_geometry(conv, ssm) + conv_local_offsets = split_info.local_conv_offsets + conv_remote_offsets = split_info.remote_conv_offsets(0, 1) + ssm_read_bytes = geom.ssm_slot_bytes + conv_subprojs = tuple( + (loff, roff, rsz) + for (loff, _lsz), (roff, rsz) in zip(conv_local_offsets, conv_remote_offsets) + ) + return MambaOffsetTemplate( + num_slots=conv.shape[0], + conv_slot_stride_bytes=geom.conv_slot_stride * geom.conv_element_size, + ssm_slot_stride_bytes=geom.ssm_slot_stride * geom.ssm_element_size, + ssm_read_bytes=ssm_read_bytes, + conv_subprojs=conv_subprojs, + ) + + +def apply_mamba_offset_template( + template: MambaOffsetTemplate, + local_slots: list[int], + remote_slots: list[int], +) -> tuple[list[int], list[int], list[int]]: + """Apply per-request slot bases to a cached ``MambaOffsetTemplate``. + + Conv sub-projection entries (all slots) come first, followed by one ssm + entry per slot. + """ + if len(local_slots) != len(remote_slots): + raise ValueError( + "local and remote mamba slot counts must match: " + f"{len(local_slots)} != {len(remote_slots)}" + ) + invalid_local = next( + (slot for slot in local_slots if not 0 <= slot < template.num_slots), None + ) + if invalid_local is not None: + raise ValueError( + f"local mamba slot {invalid_local} is outside [0, {template.num_slots})" + ) + invalid_remote = next((slot for slot in remote_slots if slot < 0), None) + if invalid_remote is not None: + raise ValueError(f"remote mamba slot {invalid_remote} must be non-negative") + # The peer's KDA slot count is not part of MoRIIOAgentMetadata yet. Do not + # compare remote slot IDs against the local pool size: two otherwise + # compatible engines may reserve different numbers of slots. + local_offs: list[int] = [] + remote_offs: list[int] = [] + sizes: list[int] = [] + + # Conv sub-projections (region 0), all slots first. + conv_slot_stride_bytes = template.conv_slot_stride_bytes + for ls, rs in zip(local_slots, remote_slots): + lbase = ls * conv_slot_stride_bytes + rbase = rs * conv_slot_stride_bytes + for loff, roff, rsz in template.conv_subprojs: + local_offs.append(lbase + loff) + remote_offs.append(rbase + roff) + sizes.append(rsz) + + # SSM temporal state (region 1), one entry per slot. + ssm_slot_stride_bytes = template.ssm_slot_stride_bytes + ssm_read_bytes = template.ssm_read_bytes + for ls, rs in zip(local_slots, remote_slots): + local_offs.append(ls * ssm_slot_stride_bytes) + remote_offs.append(rs * ssm_slot_stride_bytes) + sizes.append(ssm_read_bytes) + + return local_offs, remote_offs, sizes + + +def compute_mamba_conv_split_count( + local_slots: list[int], + split_info: MambaConvSplitInfo, +) -> int: + """Number of leading conv entries in a KDA layer transfer plan. + + Entries ``[:count]`` are conv sub-projections (conv region/session); + entries ``[count:]`` are the per-slot ssm state (ssm region/session). + """ + return len(local_slots) * len(split_info.local_conv_offsets) From 17e50b9b761023d5f2499f062507df1ff49092a4 Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Mon, 21 Sep 2026 06:50:25 +0800 Subject: [PATCH 548/774] [XPU][UT]Bugfix when the process can't see all the world_size meet accuracy issue. (#57779) Signed-off-by: yisheng Co-authored-by: Kunshang Ji --- .../device_communicators/xpu_communicator.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index 58c64e79200b..e094f6fd50f3 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -99,12 +99,24 @@ def reduce_scatterv( output = torch.empty( output_shape, dtype=input_tensor.dtype, device=input_tensor.device ) - if sizes is not None and sizes.count(sizes[0]) != len(sizes): - # if inputs shape in different ranks is not the same using reduce_scatter + if sizes is None or sizes.count(sizes[0]) == len(sizes): + dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) + elif torch.xpu.device_count() < self.world_size: + # uneven reduce_scatter is not supported when the + # process cannot see every device in the group ( + # e.g. external-LBDP). + max_size = max(sizes) + padded = input_tensor.new_zeros( + (world_size * max_size,) + input_tensor.shape[1:] + ) + for rank, split in enumerate(input_tensor.split(sizes, dim=0)): + padded[rank * max_size : rank * max_size + sizes[rank]] = split + padded_output = input_tensor.new_empty((max_size,) + input_tensor.shape[1:]) + dist.reduce_scatter_tensor(padded_output, padded, group=self.device_group) + output = padded_output[:chunk_size] + else: input_splits = list(input_tensor.split(sizes, dim=0)) dist.reduce_scatter(output, input_splits, group=self.device_group) - else: - dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) # Reshape before returning return output.movedim(0, dim).contiguous() From 9b49f92344312c41ad61e05282c8e6a2d9bafb7f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 20 Sep 2026 17:33:17 -0700 Subject: [PATCH 549/774] [Perf][DSV4.1] Overlap mHC coefficients for small TP batches (#57603) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- tests/kernels/test_mhc_kernels.py | 141 ++++++++++++++++-- .../kernels/mhc/tilelang_kernels.py | 38 ++--- vllm/model_executor/kernels/mhc/warmup.py | 1 + vllm/models/deepseek_v41/nvidia/model.py | 76 ++++++++-- .../deepseek_v41/nvidia/ops/mega_mhc.py | 30 +++- vllm/models/deepseek_v41/nvidia/ops/mhc.py | 117 +++++++++++++++ 6 files changed, 357 insertions(+), 46 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/mhc.py diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 08b1a24c85e7..8f7df556023d 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -266,6 +266,77 @@ def test_deepseek_v41_mhc_pre_delayed( torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2) +@pytest.mark.skipif( + not HAS_TILELANG_MHC or not current_platform.is_device_capability_family(100), + reason="SM100 TileLang mHC required", +) +@pytest.mark.parametrize("num_tokens", [0, 1, 8, 9, 32, 128, 256]) +@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"]) +@pytest.mark.parametrize("hidden_size,hc_mult", [(4096, 2), (5120, 4), (7168, 4)]) +def test_deepseek_v41_mhc_overlap_preserves_outputs_on_replay( + num_tokens, entry, hidden_size, hc_mult +): + """The input is usable before joining; coefficients stay exact on graph replay.""" + from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if not is_deep_gemm_supported(): + pytest.skip("DeepGEMM required") + set_random_seed(42) + mix_size = hc_mult * (hc_mult + 2) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + fn = torch.randn(mix_size, hc_mult, hidden_size, device=DEVICE) * 0.02 + x = None + pre_mix = None + if entry == "broadcast": + x = residual[:, 0].contiguous() + residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + fn = fn.sum(1) + else: + fn = fn.flatten(1) + if entry == "carried": + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) + base = torch.randn(mix_size, device=DEVICE) + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE).uniform_( + 0.5, 1.5 + ) + args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20) + kwargs = dict(pre_mix=pre_mix, x=x, norm_weight=weight, norm_eps=1e-20) + side = torch.cuda.Stream() + + def run(): + outputs = mhc_pre_delayed_overlap(*args, **kwargs, stream=side) + consumed_input = outputs[2].clone() + torch.cuda.current_stream().wait_stream(side) + return (*outputs, consumed_input) + + def check(actual): + expected = mhc_pre_delayed_tilelang(*args, **kwargs) + for a, b in zip(actual[:4], expected, strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0) + torch.testing.assert_close(actual[4], expected[2], atol=0, rtol=0) + + check(run()) + if num_tokens == 0: + return + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + for _ in range(5): + if x is not None: + x.normal_() + residual.copy_(x.unsqueeze(1).expand_as(residual)) + else: + residual.normal_() + if pre_mix is not None: + pre_mix.uniform_() + graph.replay() + check(captured) + + def mhc_fused_post_pre_delayed_ref( x: torch.Tensor, residual: torch.Tensor, @@ -433,10 +504,16 @@ def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_ @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) +@pytest.mark.parametrize( + "mhc_mode", ["disabled", "overlap", "piecewise", "large_batch"] +) def test_deepseek_v41_decoder_mixes_match_torch( - entry, monkeypatch, default_vllm_config + entry, mhc_mode, monkeypatch, default_vllm_config ): """Preserve carried pre-mixes and Engram ordering at every decoder entry.""" + from vllm.models.deepseek_v41.nvidia.ops.mhc import MHC_OVERLAP_MAX_TOKENS + + num_tokens = MHC_OVERLAP_MAX_TOKENS + 1 if mhc_mode == "large_batch" else 3 set_random_seed(0) decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) nn.Module.__init__(decoder) @@ -446,6 +523,29 @@ def test_deepseek_v41_decoder_mixes_match_torch( decoder.rms_norm_eps = 1e-20 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.mhc_stream = None + if mhc_mode != "disabled": + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if ( + not current_platform.is_device_capability_family(100) + or not is_deep_gemm_supported() + ): + pytest.skip("SM100 DeepGEMM required for overlap") + decoder.mhc_stream = torch.cuda.Stream() + if mhc_mode in ("piecewise", "large_batch"): + + def unexpected_overlap(*args, **kwargs): + pytest.fail("unsupported execution must retain native mHC") + + monkeypatch.setattr( + "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_overlap", + unexpected_overlap, + ) + monkeypatch.setattr( + "vllm.models.deepseek_v41.nvidia.ops.mhc.mhc_pre_delayed_overlap", + unexpected_overlap, + ) decoder.engram = None from vllm.model_executor.layers.layernorm import RMSNorm @@ -461,18 +561,18 @@ def test_deepseek_v41_decoder_mixes_match_torch( decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) decoder.hc_attn_base = torch.randn(24) decoder.hc_ffn_base = torch.randn(24) - x = torch.randn(3, 5120, dtype=torch.bfloat16) - positions = torch.arange(3) + x = torch.randn(num_tokens, 5120, dtype=torch.bfloat16) + positions = torch.arange(num_tokens) kwargs = {} if entry != "broadcast": - kwargs["pre_mix"] = torch.rand(3, 4) + kwargs["pre_mix"] = torch.rand(num_tokens, 4) if entry == "pipeline": - x = torch.randn(3, 4, 5120, dtype=torch.bfloat16) + x = torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16) if entry in ("residual", "engram"): kwargs.update( - residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16), - post_mix=torch.rand(3, 4, 1), - res_mix=torch.rand(3, 4, 4), + residual=torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.rand(num_tokens, 4, 1), + res_mix=torch.rand(num_tokens, 4, 4), ) if entry == "engram": @@ -483,15 +583,22 @@ def forward(self, residual, hashes, mask): return residual + 0.125 decoder.engram = FakeEngram() - kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32) + kwargs["engram_hashes"] = torch.zeros(num_tokens, 1, 1, dtype=torch.int32) - actual = decoder(x, positions, None, **kwargs) + from vllm.config import CUDAGraphMode + from vllm.forward_context import set_forward_context + + mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE + with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode): + actual = decoder(x, positions, None, **kwargs) def reference(*args, norm_weight, norm_eps, **kwargs): post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) return post, res, decoder.attn_norm(collapsed), pre - def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, **kw): + def fused_reference( + x, residual, post_mix, res_mix, *args, capture_aux=False, stream=None, **kw + ): residual = mhc_post_torch(x, residual, post_mix, res_mix) aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0) return residual, *reference(residual, *args, **kw), aux @@ -500,6 +607,7 @@ def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, ** "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", reference, ) + decoder.mhc_stream = None monkeypatch.setattr( "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", fused_reference, @@ -530,6 +638,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf decoder.rms_norm_eps = 1e-6 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.mhc_stream = None decoder.engram = None from vllm.model_executor.layers.layernorm import RMSNorm @@ -772,7 +881,8 @@ def hc_head_ref( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) @pytest.mark.parametrize("hc_mult", [4]) -def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): +@pytest.mark.parametrize("fused_norm", [False, True]) +def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm): torch.set_default_device(DEVICE) set_random_seed(0) @@ -802,6 +912,12 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): hc_post_alpha, sinkhorn_repeat, ) + norm_kwargs = {} + if fused_norm: + # Non-shifted mHC must still use its freshly computed pre-mix. + weight = torch.ones(hidden_size, dtype=torch.bfloat16) + norm_kwargs = dict(norm_weight=weight, norm_eps=rms_eps) + ref = (*ref[:2], F.rms_norm(ref[2], (hidden_size,), weight, rms_eps)) out = torch.ops.vllm.mhc_pre_tilelang( residual, fn, @@ -812,6 +928,7 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): hc_sinkhorn_eps, hc_post_alpha, sinkhorn_repeat, + **norm_kwargs, ) for actual, expected in zip(out, ref, strict=True): diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index fa2d5e17004a..33de3b9ce57c 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -305,8 +305,12 @@ def mhc_pre_big_fuse_with_norm_tilelang( use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, + split_mode: str = "fused", write_aux: bool = False, ): + # Split modes require shifted mHC: input collapse uses a carried pre-mix. + assert split_mode in ("fused", "stats", "input") + assert split_mode == "fused" or save_pre_mix num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) if gemm_last_dim < 0: @@ -330,26 +334,26 @@ def mhc_pre_big_fuse_with_norm_tilelang( aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: - rms = T.alloc_fragment(1, T.float32) - mixes = T.alloc_fragment(hc_mult3, T.float32) - T.clear(mixes) - rms[0] = 0 - if ENABLE_PDL: T.pdl_sync() - - for i_split in T.serial(n_splits): - rms[0] += gemm_out_sqrsum[i_split, i] - rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) - for j in T.Parallel(hc_mult3): - mixes[j] = 0 - for i_split in T.serial(n_splits): - mixes[j] += gemm_out_mul[i_split, i, j] - mixes[j] *= rms[0] mixes_shared = T.alloc_shared(hc_mult3, T.float32) - T.copy(mixes, mixes_shared) + if split_mode != "input": + rms = T.alloc_fragment(1, T.float32) + mixes = T.alloc_fragment(hc_mult3, T.float32) + T.clear(mixes) + rms[0] = 0 - if T.get_thread_binding() < 32: + for i_split in T.serial(n_splits): + rms[0] += gemm_out_sqrsum[i_split, i] + rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) + for j in T.Parallel(hc_mult3): + mixes[j] = 0 + for i_split in T.serial(n_splits): + mixes[j] += gemm_out_mul[i_split, i, j] + mixes[j] *= rms[0] + T.copy(mixes, mixes_shared) + + if split_mode != "input" and T.get_thread_binding() < 32: cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) for j in T.Parallel(hc_mult): if save_pre_mix: @@ -395,7 +399,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( for j, k in T.Parallel(hc_mult, hc_mult): comb_mix[i, j * hc_mult + k] = cm[j, k] - else: + elif split_mode != "stats" and T.get_thread_binding() >= 32: pre_mix_shared = T.alloc_shared(hc_mult, T.float32) for j in T.Parallel(hc_mult): if use_pre_mix_in: diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 4e6a19781c82..deb1cf429a0c 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -36,6 +36,7 @@ class CompileKey: use_pre_mix_in: bool rms_numel: int save_pre_mix: bool = True + split_mode: str = "fused" write_aux: bool = False kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 6d1faace6510..b5caceab6fca 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing from collections.abc import Callable, Iterable +from functools import partial from itertools import islice import regex as re @@ -17,7 +18,11 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.forward_context import ( + get_forward_context, + in_piecewise_cudagraph, + is_forward_context_available, +) from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( mhc_post_tilelang, @@ -84,6 +89,11 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes from .ops.mega_mhc import mhc_shifted_post_pre +from .ops.mhc import ( + MHC_OVERLAP_MAX_TOKENS, + mhc_pre_delayed_overlap, + supports_mhc_overlap, +) if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -188,11 +198,13 @@ def __init__( aux_stream_list: list[torch.cuda.Stream] | None = None, candidate_block_buffer: torch.Tensor | None = None, engram_layout: EngramLayout | None = None, + mhc_stream: torch.cuda.Stream | None = None, ): super().__init__() config = vllm_config.model_config.hf_config self.hidden_size = config.hidden_size + self.mhc_stream = mhc_stream self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.engram: Engram | None = None @@ -269,19 +281,30 @@ def __init__( variants.append( {**hc_stream, "use_pre_mix_in": True, "write_aux": True} ) + # Shifted mHC splits the epilogue in two: the input collapse on the + # caller stream and coefficient generation on a side stream. + split_modes = ( + ("fused", "stats", "input") if mhc_stream is not None else ("fused",) + ) for variant in variants: - MHC_PRE_NORM_KERNEL.register_warmup( - max_tokens=max_tokens, - hidden_size=self.hidden_size, - rms_eps=self.rms_norm_eps, - hc_pre_eps=self.hc_eps, - hc_sinkhorn_eps=self.hc_eps, - hc_post_mult_value=self.hc_post_alpha, - sinkhorn_repeat=self.hc_sinkhorn_iters, - norm_eps=self.rms_norm_eps, - hc_mult=self.hc_mult, - **variant, - ) + for split_mode in split_modes: + MHC_PRE_NORM_KERNEL.register_warmup( + split_mode=split_mode, + max_tokens=( + max_tokens + if split_mode == "fused" + else min(max_tokens, MHC_OVERLAP_MAX_TOKENS) + ), + hidden_size=self.hidden_size, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + norm_eps=self.rms_norm_eps, + hc_mult=self.hc_mult, + **variant, + ) mix_hc = (2 + self.hc_mult) * self.hc_mult hc_dim = self.hc_mult * self.hidden_size self.hc_attn_fn = nn.Parameter( @@ -351,6 +374,18 @@ def forward( torch.Tensor | None, ]: previous_aux: torch.Tensor | None = None + mhc_stream = self.mhc_stream + if mhc_stream is not None and ( + in_piecewise_cudagraph() + or not 0 < positions.shape[0] <= MHC_OVERLAP_MAX_TOKENS + ): + # A side stream cannot remain unjoined across breakable graph segments. + mhc_stream = None + mhc_pre = ( + partial(mhc_pre_delayed_overlap, stream=mhc_stream) + if mhc_stream is not None + else mhc_pre_delayed_tilelang + ) # The reference collapses each sublayer's input with the *previous* # sublayer's pre-mix: attention uses the pre-mix carried in (identity # for the first layer), the FFN uses this layer's attention pre-mix. @@ -360,7 +395,7 @@ def forward( # copies and the identity pre-mix selects copy 0. assert self.hc_attn_fn_broadcast is not None residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous() - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang( + post_mix, res_mix, x, attn_pre = mhc_pre( residual, self.hc_attn_fn_broadcast, self.hc_attn_scale, @@ -376,7 +411,7 @@ def forward( ) else: residual = x - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang( + post_mix, res_mix, x, attn_pre = mhc_pre( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -403,7 +438,7 @@ def forward( engram_hashes[:, self.engram.layer_hash_index], engram_mask, ) - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang( + post_mix, res_mix, x, attn_pre = mhc_pre( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -437,6 +472,7 @@ def forward( norm_weight=self.attn_norm.weight, norm_eps=self.attn_norm.variance_epsilon, capture_aux=capture_previous_aux, + stream=mhc_stream, ) if capture_previous_aux: previous_aux = aux @@ -448,6 +484,8 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) + if mhc_stream is not None: + torch.cuda.current_stream().wait_stream(mhc_stream) residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( x, residual, @@ -464,8 +502,11 @@ def forward( pre_mix=attn_pre, norm_weight=self.ffn_norm.weight, norm_eps=self.ffn_norm.variance_epsilon, + stream=mhc_stream, ) x = self.ffn(x, input_ids, mega_gate_metadata) + if mhc_stream is not None: + torch.cuda.current_stream().wait_stream(mhc_stream) return x, residual, post_mix, res_mix, ffn_pre, previous_aux @@ -497,6 +538,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # (compressor kv_score, indexer.weights_proj). fused_wqa_wkv stays on # the default stream. aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + # Keep mHC independent of the streams used inside attention. + mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -541,6 +584,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): aux_stream_list=aux_stream_list, candidate_block_buffer=self.candidate_block_buffer, engram_layout=self.engram_layout, + mhc_stream=mhc_stream, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py index f2c74902616f..caae3d8d9b00 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -7,6 +7,7 @@ from vllm.model_executor.kernels.mhc.tilelang import ( mhc_fused_post_pre_delayed_tilelang, + mhc_post_tilelang, ) from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( @@ -106,10 +107,37 @@ def mhc_shifted_post_pre( norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, capture_aux: bool = False, + *, + stream: torch.cuda.Stream | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor ]: - """Use Mega-mHC for carried mixing, retaining TileLang for aux capture.""" + """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. + + When stream is supplied, join it before consuming the returned coefficients. + """ + if stream is not None: + from .mhc import mhc_pre_delayed_overlap + + residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix) + aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1]) + pre_outputs = mhc_pre_delayed_overlap( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + stream=stream, + ) + return residual, *pre_outputs, aux + if ( pre_mix is not None and norm_weight is not None diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py new file mode 100644 index 000000000000..29c9c3b95684 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Split shifted mHC pre so coefficient generation can overlap the sublayer.""" + +from functools import partial + +import torch + +from vllm.config import VllmConfig +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import is_deep_gemm_supported + +# GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress. +MHC_OVERLAP_MAX_TOKENS = 16 + + +def supports_mhc_overlap(vllm_config: VllmConfig) -> bool: + """Check kernel requirements and safety of sharing the coefficient stream.""" + config = vllm_config.model_config.hf_config + # DeepGEMM's prenorm kernel requires K % 64 == 0, N % 8 == 0, and N <= 32. + mix_size = config.hc_mult * (config.hc_mult + 2) + return ( + current_platform.is_device_capability_family(100) + and is_deep_gemm_supported() + and config.hidden_size % 64 == 0 + and 0 < mix_size <= 32 + and mix_size % 8 == 0 + and not vllm_config.parallel_config.use_ubatching + ) + + +def mhc_pre_delayed_overlap( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + x: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + *, + stream: torch.cuda.Stream, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare the input on the caller stream and coefficients on another stream. + + Returns post mix, residual mix, normalized input, and next pre-mix. Only the + input is ready on the caller stream; join stream before using coefficients. + The caller must retain the weights and coefficient outputs until the join. + """ + from vllm.model_executor.kernels.mhc.warmup import ( + MHC_PRE_NORM_KERNEL, + compute_mhc_pre_num_splits, + ) + from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm + + n, hc, hidden = residual.shape + assert residual.is_contiguous() and residual.dtype == torch.bfloat16 + assert norm_weight is not None + if x is None: + x = residual.view(n, hc * hidden) + assert x.is_contiguous() + post = torch.empty((n, hc), device=residual.device, dtype=torch.float32) + comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32) + next_pre = torch.empty_like(post) + layer_input = torch.empty((n, hidden), device=residual.device, dtype=torch.bfloat16) + outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre + if n == 0: + return outputs + splits = compute_mhc_pre_num_splits(x.shape[1], n) + mix = torch.empty( + (splits, n, hc * (hc + 2)), device=residual.device, dtype=torch.float32 + ) + sqr = torch.empty((splits, n), device=residual.device, dtype=torch.float32) + epilogue = partial( + MHC_PRE_NORM_KERNEL, + mix, + sqr, + hc_scale, + hc_base, + residual, + post, + comb, + layer_input, + norm_weight, + pre_mix if pre_mix is not None else post, + next_pre, + layer_input, # Unused aux output in split modes. + hidden_size=hidden, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=norm_eps, + hc_mult=hc, + use_pre_mix_in=pre_mix is not None, + save_pre_mix=True, + rms_numel=x.shape[1], + ) + main = torch.cuda.current_stream() + # Prioritize input readiness for small batches before releasing statistics. + if n <= 8: + epilogue(split_mode="input") + stream.wait_stream(main) + with torch.cuda.stream(stream): + tf32_hc_prenorm_gemm(x, fn, mix, sqr, splits) + epilogue(split_mode="stats") + for tensor in (x, mix, sqr): + tensor.record_stream(stream) + if n > 8: + epilogue(split_mode="input") + return outputs From d2983f2f163249acbe3053c073cb9830290984a0 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Sun, 20 Sep 2026 18:06:25 -0700 Subject: [PATCH 550/774] [Bugfix][Spec Decode] Stop dummy draft decode steps from writing KV through stale block-table rows (#56734) Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 --- tests/v1/worker/test_gpu_block_table.py | 39 +++++++++++++++++++ vllm/v1/worker/gpu/block_table.py | 7 +++- .../spec_decode/autoregressive/speculator.py | 1 + .../multi_module_mtp/speculator.py | 1 + vllm/v1/worker/gpu/spec_decode/speculator.py | 14 +++++-- 5 files changed, 56 insertions(+), 6 deletions(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index e77156eef901..5c245d1c4d9a 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,6 +5,7 @@ import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.block_table import BlockTables pytestmark = pytest.mark.skipif( @@ -273,3 +274,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() + + +def test_dummy_request_slot_mapping_is_pad(): + """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request. + + A dummy draft decode step must not resolve slots through the persistent + block-table row of a real request slot, which may be stale and point at + blocks that are now in the prefix cache. + """ + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[4], + max_num_reqs=2, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[4], + ) + block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) + block_tables.apply_staged_writes() + query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) + positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) + + real = block_tables.compute_slot_mappings( + torch.tensor([0], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3] + + dummy = block_tables.compute_slot_mappings( + torch.tensor([-1], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert dummy[0].tolist() == [PAD_SLOT_ID] * 3 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index c01895e63545..6563e30c342f 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -315,6 +315,9 @@ def _compute_slot_mappings_kernel( mapping_enabled = tl.load(slot_mapping_enabled + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) + # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns + # no blocks: never read its block-table row and emit PAD for its tokens. + is_real_req = req_state_idx >= 0 start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): @@ -342,12 +345,12 @@ def _compute_slot_mappings_kernel( block_offsets = local_positions % kernel_block_size block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local, + mask=is_local & is_real_req, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets if CP_SIZE != 1: slot_ids = tl.where(is_local, slot_ids, PAD_ID) - slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index f166d63e8efb..c787efa07fe0 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -254,6 +254,7 @@ def propose( input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) # Get the input ids and last token indices for the speculator. diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c6ae315cf46..fc0924e3a3f4 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -166,6 +166,7 @@ def propose( input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) num_tokens = input_batch.num_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 299cac641620..534b6ef8eb17 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -481,6 +481,7 @@ def _copy_request_inputs( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, + dummy_run: bool = False, ) -> None: # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. # NOTE(woosuk): For draft sampling, we only consider the temperature @@ -490,10 +491,15 @@ def _copy_request_inputs( # affect the output distribution after rejection sampling. self.temperature.copy_(temperature) self.seeds.copy_(seeds) - self.idx_mapping[:num_reqs].copy_(idx_mapping) - # idx_mapping for CG padded requests points to -1, which is ignored - # during sampling to prevent writing stale values to draft logits. - self.idx_mapping[num_reqs:].fill_(-1) + # idx_mapping == -1 marks a row the drafter must not act on: sampling + # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows + # always get it; a dummy batch gets it for every row, since its arange + # idx_mapping names request-state slots with stale block tables. + if dummy_run: + self.idx_mapping.fill_(-1) + else: + self.idx_mapping[:num_reqs].copy_(idx_mapping) + self.idx_mapping[num_reqs:].fill_(-1) def _build_uniform_batch_dp_sync( self, From 7ddd461a284d2c318410f29110f1604028b961ee Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Mon, 21 Sep 2026 09:15:02 +0800 Subject: [PATCH 551/774] [Docs] Add ECMooncakeConnector usage example for EPD (#57742) Signed-off-by: jiangkuaixue123 Co-authored-by: Codex --- docs/features/disagg_encoder.md | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index a4b2ece241a7..383e7d45ff75 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -32,8 +32,9 @@ Design doc: Date: Mon, 21 Sep 2026 09:25:02 +0800 Subject: [PATCH 552/774] [Bugfix][KV Connector] Propagate cache reset failure during sleep (#54581) Signed-off-by: CrazyCoder <18235787078@163.com> Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Codex --- vllm/v1/engine/core.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 5ee2c6f8c02c..a51278b7617c 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -867,10 +867,11 @@ def _reset_caches( # reset_connector=True so external connectors clear alongside # local caches, matching the pause_generation(clear_cache=True) # contract. No-op when no connector is configured. - self.reset_prefix_cache( + if not self.reset_prefix_cache( reset_running_requests=reset_running_requests, reset_connector=reset_connector, - ) + ): + raise RuntimeError("Failed to reset the KV connector cache.") self.reset_mm_cache() self.reset_encoder_cache() @@ -2001,8 +2002,12 @@ def pause_scheduler( raise ValueError(f"Invalid pause mode: {mode}") def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None: - engine._finish_pause(clear_cache) - future.set_result(None) + try: + engine._finish_pause(clear_cache) + except Exception as e: + future.set_exception(e) + else: + future.set_result(None) if mode == "abort": aborted_reqs = self.scheduler.finish_requests( From 3e67044972939772ef6cd06957bf3199a41b5139 Mon Sep 17 00:00:00 2001 From: Zhou Zihang Date: Mon, 21 Sep 2026 09:42:55 +0800 Subject: [PATCH 553/774] [Doc] Refresh Kthena integration guide (#57314) Signed-off-by: Zhou Zihang Co-authored-by: Jee Jee Li --- docs/deployment/integrations/kthena.md | 59 ++++++++++++++------------ 1 file changed, 32 insertions(+), 27 deletions(-) diff --git a/docs/deployment/integrations/kthena.md b/docs/deployment/integrations/kthena.md index 0285b733c6b8..73dcb3dafd72 100644 --- a/docs/deployment/integrations/kthena.md +++ b/docs/deployment/integrations/kthena.md @@ -2,7 +2,9 @@ [**Kthena**](https://github.com/volcano-sh/kthena) is a Kubernetes-native LLM inference platform that transforms how organizations deploy and manage Large Language Models in production. Built with declarative model lifecycle management and intelligent request routing, it provides high performance and enterprise-grade scalability for LLM inference workloads. -This guide shows how to deploy a production-grade, **multi-node vLLM** service on Kubernetes. +Kthena manages inference workloads through `ModelServing` and provides model-aware request routing through `ModelServer` and `ModelRoute`. + +This guide shows how to deploy a **multi-node vLLM** service on Kubernetes using `ModelServing` and expose it through a Kubernetes Service. We’ll: @@ -18,7 +20,7 @@ You need: - A Kubernetes cluster with **GPU nodes**. - `kubectl` access with cluster-admin or equivalent permissions. -- **Volcano** installed for gang scheduling. +- **Volcano 1.14 or later** for the role-level gang scheduling used in this example. - **Kthena** installed with the `ModelServing` CRD available. - A valid **Hugging Face token** if loading models from Hugging Face Hub. @@ -30,21 +32,22 @@ helm repo update helm install volcano volcano-sh/volcano -n volcano-system --create-namespace ``` -This provides the gang-scheduling and network topology features used by Kthena. +Volcano provides the gang scheduling used in this example. It is optional for Kthena deployments that do not use Volcano scheduling. ### 1.2 Install Kthena ```bash -helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v0.1.0 --namespace kthena-system --create-namespace +helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v1.0.0 --namespace kthena-system --create-namespace ``` - The `kthena-system` namespace is created. -- Kthena controllers and CRDs, including `ModelServing`, are installed and healthy. +- Kthena controllers and CRDs, including `ModelServing`, are installed. Validate: ```bash -kubectl get crd | grep modelserving +kubectl get pods -n kthena-system +kubectl get crd modelservings.workload.serving.volcano.sh ``` You should see: @@ -57,7 +60,7 @@ modelservings.workload.serving.volcano.sh ... ## 2. The Multi-Node vLLM `ModelServing` Example -Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Conceptually this is equivalent to the vLLM production stack Helm deployment, but expressed with `ModelServing`. +Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Each role replica contains an entry pod and its worker pods. A simplified version of the example (`llama-multinode`) looks like: @@ -116,10 +119,9 @@ Commands: - -c - > bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh - leader --ray_cluster_size=2; python3 -m - vllm.entrypoints.launchers.api_server.entry --port 8080 --model - meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 - --pipeline-parallel-size 2 + leader --ray_cluster_size=2; + vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 + --pipeline-parallel-size 2 --distributed-executor-backend ray ``` Worker: @@ -166,9 +168,9 @@ Save one of the following manifests to `modelserving.yaml`: restartGracePeriodSeconds: 60 gangPolicy: minRoleReplicas: - 405b: 1 + llama-405b: 1 roles: - - name: 405b + - name: llama-405b replicas: 2 entryTemplate: spec: @@ -256,9 +258,9 @@ Save one of the following manifests to `modelserving.yaml`: restartGracePeriodSeconds: 60 gangPolicy: minRoleReplicas: - 405b: 1 + llama-405b: 1 roles: - - name: 405b + - name: llama-405b replicas: 2 entryTemplate: spec: @@ -275,7 +277,7 @@ Save one of the following manifests to `modelserving.yaml`: - sh - -c - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2; - vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2" + vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray" resources: limits: nvidia.com/gpu: "8" @@ -383,24 +385,24 @@ kubectl get pod -owide -l modelserving.volcano.sh/name=llama-multinode Example output (from docs): ```text -NAMESPACE NAME READY STATUS RESTARTS AGE IP NODE ... -default llama-multinode-0-405b-0-0 1/1 Running 0 15m 10.244.0.56 192.168.5.12 ... -default llama-multinode-0-405b-0-1 1/1 Running 0 15m 10.244.0.58 192.168.5.43 ... -default llama-multinode-0-405b-1-0 1/1 Running 0 15m 10.244.0.57 192.168.5.58 ... -default llama-multinode-0-405b-1-1 1/1 Running 0 15m 10.244.0.53 192.168.5.36 ... +NAMESPACE NAME READY STATUS RESTARTS AGE IP NODE ... +default llama-multinode-0-llama-405b-0-0 1/1 Running 0 15m 10.244.0.56 192.168.5.12 ... +default llama-multinode-0-llama-405b-0-1 1/1 Running 0 15m 10.244.0.58 192.168.5.43 ... +default llama-multinode-0-llama-405b-1-0 1/1 Running 0 15m 10.244.0.57 192.168.5.58 ... +default llama-multinode-0-llama-405b-1-1 1/1 Running 0 15m 10.244.0.53 192.168.5.36 ... ``` Pod name pattern: - `llama-multinode----`. -The first number indicates `ServingGroup`. The second (`405b`) is the `Role`. The remaining indices identify the pod within the role. +The first index identifies the `ServingGroup`, followed by the role name (`llama-405b`), role replica index, and pod index. --- -## 6. Accessing the vLLM OpenAI-Compatible API +## 5. Accessing the vLLM OpenAI-Compatible API -Expose the entry via a Service: +Save the following Service as `service.yaml` to expose the entry pods: ```yaml apiVersion: v1 @@ -412,7 +414,6 @@ spec: selector: modelserving.volcano.sh/name: llama-multinode modelserving.volcano.sh/entry: "true" - # optionally further narrow to leader role if you label it ports: - name: http port: 80 @@ -420,9 +421,10 @@ spec: type: ClusterIP ``` -Port-forward from your local machine: +Apply the Service, then port-forward from your local machine: ```bash +kubectl apply -f service.yaml kubectl port-forward svc/llama-multinode-openai 30080:80 -n default ``` @@ -450,11 +452,12 @@ You should see an OpenAI-style response from vLLM. --- -## 7. Clean Up +## 6. Clean Up To remove the deployment and its resources: ```bash +kubectl delete -f service.yaml kubectl delete modelserving llama-multinode -n default ``` @@ -464,3 +467,5 @@ If you’re done with the entire stack: helm uninstall kthena -n kthena-system # or your Kthena release name helm uninstall volcano -n volcano-system ``` + +For model-aware routing and prefill-decode disaggregation, see the [Kthena documentation](https://kthena.volcano.sh/docs/user-guide/router-routing). From 6dfd87f59658eba0af0354418ca725c64ef1ed7d Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Sun, 20 Sep 2026 19:17:11 -0700 Subject: [PATCH 554/774] [ROCm][DSv4.1] Keep the Engram tables in host memory on ROCm (#57491) --- tests/kernels/test_engram.py | 8 ++++++-- tests/test_config.py | 12 +++++++----- vllm/config/engram.py | 2 +- vllm/config/vllm.py | 2 +- vllm/models/deepseek_v41/amd/model.py | 6 +++++- 5 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c2a7607cbac0..45cebcfe5711 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -599,7 +599,9 @@ def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" +) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch): @@ -701,7 +703,9 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): assert cache_sizes == [1] * len(cache_sizes) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" +) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 7, 256]) diff --git a/tests/test_config.py b/tests/test_config.py index 9da6eced9429..f31adfdb8d69 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1299,7 +1299,7 @@ def test_engram_dp_shared_memory_config_validation( monkeypatch, dp_size, load_format, multithread, error ): """Reject invalid shared configs before distributed init; allow threaded loads.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) config = cast( VllmConfig, SimpleNamespace( @@ -1324,7 +1324,7 @@ def test_engram_dp_shared_memory_config_validation( @pytest.mark.parametrize( - "architecture, ple_layers, cuda, supported", + "architecture, ple_layers, accelerator, supported", [ ("DeepseekV41ForCausalLM", [1], True, True), ("DeepseekV41ForCausalLM", [], True, False), @@ -1339,9 +1339,11 @@ def test_engram_dp_shared_memory_config_validation( (None, None, True, False), ], ) -def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, supported): +def test_engram_model_support( + monkeypatch, architecture, ple_layers, accelerator, supported +): """A similarly named HF field must not enable unsupported implementations.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator) model = ( cast( ModelConfig, @@ -1430,7 +1432,7 @@ def test_engram_explicit_config_requires_supported_model(): @pytest.mark.parametrize("explicit", [False, True]) def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit): """MTP may inherit cross-DP sharding without having its own PLE layers.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) target = SimpleNamespace( architecture="Qwen4ExpForCausalLM", hf_text_config=SimpleNamespace(ple_layer_ids=[1] if target_has_ple else []), diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 5abb2bd34a5e..773ca2091153 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -74,7 +74,7 @@ def verify_model_config(self, model_config: "ModelConfig | None") -> None: if ( model_config is None or field is None - or not current_platform.is_cuda() + or not current_platform.is_cuda_alike() or not getattr(model_config.hf_text_config, field, None) ): raise ValueError( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 78e30cf9d058..6974d1fbb3bd 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1260,7 +1260,7 @@ def _resolve_and_verify_engram_config(self) -> None: "Disable --enable-dbo and set --ubatch-size to 0." ) if self.engram_config is None: - if not current_platform.is_cuda() or not model_has_engram_layers( + if not current_platform.is_cuda_alike() or not model_has_engram_layers( model_config ): return diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index b3eddc202452..1383bac820cb 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -61,9 +61,13 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id -from ..common.engram import Engram, EngramLayout, NgramHashState +from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask +# Engram host offload and its prefetch stream are neither ROCm- nor +# NVIDIA-specific, so they are imported rather than duplicated. +from ..nvidia.engram import Engram + if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata From 3918f3c5a3e389fb17d9c0d629811f3745014db0 Mon Sep 17 00:00:00 2001 From: czhu-cohere Date: Sun, 20 Sep 2026 19:17:23 -0700 Subject: [PATCH 555/774] [Profiler] Unify platform-aware torch profiling (#57460) --- docs/contributing/profiling.md | 4 + tests/v1/engine/test_async_llm.py | 37 +++- tests/v1/worker/test_gpu_profiler.py | 241 +++++++++++++++++++++++++-- vllm/config/profiler.py | 26 ++- vllm/profiler/wrapper.py | 117 +++++++++++-- vllm/v1/engine/async_llm.py | 2 + vllm/v1/worker/cpu_worker.py | 22 --- vllm/v1/worker/gpu_worker.py | 41 ++--- vllm/v1/worker/xpu_worker.py | 28 ---- 9 files changed, 408 insertions(+), 110 deletions(-) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 880d065a2c5e..86bd68acef0c 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -22,6 +22,10 @@ To use the `torch.profiler` module, set the `profiler` entry to `'torch'` and `t - `torch_profiler_with_flops` to enable recording FLOPs, off by default - `torch_profiler_use_gzip` to control gzip-compressing profiling files, on by default - `torch_profiler_dump_cuda_time_total` to control dumping and printing the aggregated CUDA self time table, on by default +- `torch_profiler_activities` to select worker activities. Defaults are + platform-specific: `["CPU"]` on CPU, `["CPU", "CUDA"]` on NVIDIA GPUs, and + `["CPU", "XPU"]` on Intel GPUs. Selecting `["CUDA"]` omits CPU annotations + and the AsyncLLM CPU trace to reduce profiling overhead and trace size. When using `vllm bench serve`, you can enable profiling by passing the `--profile` flag. diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 53a68aa135d3..cac43d8204d0 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -60,13 +60,7 @@ } -def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): - vllm_config = MagicMock() - vllm_config.observability_config.otlp_traces_endpoint = None - vllm_config.scheduler_config.stream_interval = 1 - vllm_config.profiler_config.profiler = "cuda" - vllm_config.profiler_config.ignore_frontend = False - +def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch): renderer = MagicMock() engine_core = MagicMock() engine_core.profile_async = AsyncMock() @@ -90,6 +84,16 @@ def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatc "make_async_mp_client", MagicMock(return_value=engine_core), ) + return engine_core + + +def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): + vllm_config = MagicMock() + vllm_config.observability_config.otlp_traces_endpoint = None + vllm_config.scheduler_config.stream_interval = 1 + vllm_config.profiler_config.profiler = "cuda" + vllm_config.profiler_config.ignore_frontend = False + engine_core = _mock_async_llm_dependencies(monkeypatch) engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) @@ -103,6 +107,25 @@ async def profile(): engine_core.profile_async.assert_has_awaits([call(True, None), call(False)]) +def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( + monkeypatch: pytest.MonkeyPatch, +): + vllm_config = MagicMock() + vllm_config.observability_config.otlp_traces_endpoint = None + vllm_config.scheduler_config.stream_interval = 1 + vllm_config.profiler_config.profiler = "torch" + vllm_config.profiler_config.ignore_frontend = False + vllm_config.profiler_config.torch_profiler_activities = ["CUDA"] + _mock_async_llm_dependencies(monkeypatch) + profiler = MagicMock() + monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + assert engine.profiler is None + profiler.assert_not_called() + + async def generate( engine: AsyncLLM, request_id: str, diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index a283b6d932bc..cb3b177f41d4 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -7,6 +7,7 @@ from uuid import UUID import pytest +import torch from pydantic import ValidationError from vllm.config import ( @@ -17,10 +18,17 @@ ) from vllm.config.profiler import _is_uri_path from vllm.platforms import current_platform -from vllm.profiler.wrapper import ProtonProfilerWrapper, WorkerProfiler +from vllm.profiler.wrapper import ( + ProtonProfilerWrapper, + TorchProfilerWrapper, + WorkerProfiler, + create_worker_profiler, + validate_worker_profiler_config, +) from vllm.v1.core.sched.output import CachedRequestData from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker +from vllm.v1.worker.xpu_worker import XPUWorker class ConcreteWorkerProfiler(WorkerProfiler): @@ -51,6 +59,213 @@ def default_profiler_config(): ) +def test_torch_profiler_rebuilds_one_shot_profiler_each_round(tmp_path): + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_activities=["CUDA"], + torch_profiler_dump_cuda_time_total=False, + warmup_iterations=2, + ) + profilers = [MagicMock(), MagicMock()] + for profiler in profilers: + profiler.profiler = MagicMock() + + with patch( + "vllm.profiler.wrapper.torch.profiler.profile", side_effect=profilers + ) as profile: + wrapper = TorchProfilerWrapper( + config, worker_name="worker", local_rank=0, activities=["CUDA"] + ) + profile.assert_not_called() + + wrapper.start() + assert not wrapper.should_annotate + assert not wrapper._profiler_step() + wrapper.stop() + wrapper.start() + assert not wrapper._profiler_step() + wrapper.stop() + + assert profile.call_count == 2 + assert profile.call_args.kwargs["activities"] == [ + torch.profiler.ProfilerActivity.CUDA + ] + for profiler in profilers: + profiler.start.assert_called_once_with() + profiler.step.assert_called_once_with() + profiler.stop.assert_called_once_with() + + +def test_torch_profiler_records_each_profile_round(tmp_path): + traces: list[torch.profiler.profile] = [] + wrapper = TorchProfilerWrapper( + ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_dump_cuda_time_total=False, + ), + worker_name="worker", + local_rank=1, + activities=["CPU"], + on_trace_ready=traces.append, + ) + + for run in range(2): + wrapper.start() + with torch.profiler.record_function(f"run_{run}"): + pass + wrapper.stop() + + assert len(traces) == 2 + for run, trace in enumerate(traces): + assert { + event.name for event in trace.events() if event.name.startswith("run_") + } == {f"run_{run}"} + + +@pytest.mark.parametrize( + "activities", [["CPU", "CUDA"], ["CUDA"], ["CPU", "XPU"], ["XPU"]] +) +@pytest.mark.parametrize("dump_device_time", [True, False]) +def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device_time): + """Device summaries honor the dump option on both CUDA and XPU.""" + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_dump_cuda_time_total=dump_device_time, + ) + with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile: + profile.return_value.key_averages.return_value.table.return_value = ( + "device times" + ) + wrapper = TorchProfilerWrapper( + config, worker_name="worker", local_rank=0, activities=activities + ) + wrapper.start() + wrapper.stop() + + summary = tmp_path / "profiler_out_0.txt" + assert summary.exists() == dump_device_time + assert ("device times" in capsys.readouterr().out) == dump_device_time + if dump_device_time: + assert summary.read_text() == "device times\n" + + +@pytest.mark.parametrize( + "activities", + [[], ["CPU", "CPU"], ["INVALID"]], +) +def test_torch_profiler_activities_reject_invalid_values(activities): + with pytest.raises(ValueError, match="torch_profiler_activities"): + ProfilerConfig( + profiler="torch", + torch_profiler_dir="/tmp/mock", + torch_profiler_activities=activities, + ) + + +def test_torch_profiler_activities_require_torch_profiler(): + with pytest.raises(ValueError, match="only applicable"): + ProfilerConfig(torch_profiler_activities=["CPU"]) + + +@pytest.mark.parametrize( + ("device_type", "activities"), + [ + ("cuda", ["XPU"]), + ("xpu", ["CUDA"]), + ("cpu", ["CUDA"]), + ("cpu", ["XPU"]), + ], +) +def test_worker_rejects_unsupported_activities_at_startup(device_type, activities): + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir="/tmp/mock", + torch_profiler_activities=activities, + ) + + with ( + patch.object(current_platform, "device_type", device_type), + pytest.raises(ValueError, match="Unsupported torch profiler activities"), + ): + validate_worker_profiler_config(config) + + +@pytest.mark.parametrize("device_type", ["cpu", "xpu"]) +def test_worker_rejects_cuda_profiler_on_other_devices(device_type): + with ( + patch.object(current_platform, "device_type", device_type), + pytest.raises(ValueError, match="Unsupported profiler type"), + ): + validate_worker_profiler_config(ProfilerConfig(profiler="cuda")) + + +@pytest.mark.parametrize( + ("device_type", "activities", "expected"), + [ + ("cuda", None, ("CPU", "CUDA")), + ("xpu", None, ("CPU", "XPU")), + ("cpu", None, ("CPU",)), + ("cuda", ["CUDA"], ("CUDA",)), + ("xpu", ["XPU"], ("XPU",)), + ("cuda", ["CPU"], ("CPU",)), + ("xpu", ["CPU"], ("CPU",)), + ("cpu", ["CPU"], ("CPU",)), + ], +) +def test_worker_creates_platform_torch_profiler(device_type, activities, expected): + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir="/tmp/mock", + torch_profiler_activities=activities, + ) + + with ( + patch.object(current_platform, "device_type", device_type), + patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, + ): + validate_worker_profiler_config(config) + profiler = create_worker_profiler( + config, + worker_name="rank0", + local_rank=0, + ) + + assert profiler is wrapper.return_value + wrapper.assert_called_once_with( + config, + worker_name="rank0", + local_rank=0, + activities=expected, + ) + + +@pytest.mark.parametrize("worker_type", [Worker, XPUWorker]) +def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): + worker = object.__new__(worker_type) + worker.rank = 0 + worker.local_rank = 0 + worker.profiler = None + worker.profiler_config = ProfilerConfig( + profiler="torch", torch_profiler_dir="/tmp/mock" + ) + + with ( + patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"), + patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, + ): + worker.profile() + worker.profile(is_start=False) + worker.profile() + + assert worker.profiler is wrapper.return_value + wrapper.assert_called_once() + assert wrapper.return_value.start.call_count == 2 + wrapper.return_value.stop.assert_called_once_with() + + def test_immediate_start_stop(default_profiler_config): """Test standard start without delay.""" profiler = ConcreteWorkerProfiler(default_profiler_config) @@ -318,13 +533,17 @@ def test_detailed_format_mixed(self): "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)" ) - def test_skips_annotation_work_after_profiler_stops(self): + def test_skips_annotation_work_when_profiler_does_not_annotate(self): worker = MagicMock() - worker.profiler.is_running = False + worker.profiler.should_annotate = False - context = Worker.annotate_profile(worker, scheduler_output=None) + with patch( + "vllm.v1.worker.gpu_worker.compute_iteration_details" + ) as compute_iteration_details: + context = Worker.annotate_profile(worker, scheduler_output=None) worker.profiler.step.assert_called_once_with() + compute_iteration_details.assert_not_called() worker.profiler.annotate_context_manager.assert_not_called() assert isinstance(context, nullcontext) @@ -837,17 +1056,18 @@ def test_scope_annotations_delegate_to_proton(self, tmp_path): @_requires_cuda_for_proton def test_gpu_worker_creates_proton_profiler(): - worker = MagicMock() + worker = object.__new__(Worker) worker.rank = 1 + worker.local_rank = 1 worker.profiler = None - worker.profiler_config.profiler = "proton" + worker.profiler_config = MagicMock(profiler="proton") with ( patch( "vllm.distributed.utils.get_worker_rank_suffix", return_value="rank1", ), - patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, + patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper, ): Worker.profile(worker) @@ -857,17 +1077,18 @@ def test_gpu_worker_creates_proton_profiler(): @_requires_cuda_for_proton def test_gpu_worker_recreates_proton_profiler_for_each_run(): - worker = MagicMock() + worker = object.__new__(Worker) worker.rank = 1 + worker.local_rank = 1 worker.profiler = None - worker.profiler_config.profiler = "proton" + worker.profiler_config = MagicMock(profiler="proton") with ( patch( "vllm.distributed.utils.get_worker_rank_suffix", return_value="rank1", ), - patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, + patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper, ): wrapper.return_value.has_cuda_graph_session = False Worker.profile(worker, profile_prefix="first") diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index f3a69853243a..3ade4bceed87 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -14,6 +14,7 @@ logger = init_logger(__name__) ProfilerKind = Literal["torch", "cuda", "proton"] +TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"] ProtonBackend = Literal["cupti"] ProtonContext = Literal["shadow", "python"] ProtonData = Literal["tree", "trace"] @@ -51,6 +52,12 @@ class ProfilerConfig: worker's traces (CPU & GPU) will be saved under this directory. Note that it must be an absolute path.""" + torch_profiler_activities: list[TorchProfilerActivity] | None = Field( + default=None, min_length=1 + ) + """Activities recorded by workers using the torch profiler. When unset, + each worker uses its platform default: CPU; CPU and CUDA; or CPU and XPU.""" + proton_profiler_dir: str = "" """Directory to save Triton Proton profiles. Each worker writes a separate rank-qualified file.""" @@ -170,7 +177,24 @@ def compute_hash(self) -> str: @model_validator(mode="after") def _validate_profiler_config(self) -> Self: has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0 - if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend: + activities = self.torch_profiler_activities + if activities is not None: + if self.profiler != "torch": + raise ValueError( + "torch_profiler_activities is only applicable when profiler " + "is set to 'torch'" + ) + if len(activities) != len(set(activities)): + raise ValueError( + "torch_profiler_activities must not contain duplicates" + ) + records_cpu_activity = activities is None or "CPU" in activities + if ( + self.profiler == "torch" + and has_delay_or_limit + and not self.ignore_frontend + and records_cpu_activity + ): logger.warning_once( "Using 'torch' profiler with delay_iterations or max_iterations " "while ignore_frontend is False may result in high overhead." diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index bc679407c39e..bd9b6c231db7 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -6,10 +6,9 @@ import json import os from abc import ABC, abstractmethod -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress from glob import glob -from typing import Literal from uuid import uuid4 import torch @@ -18,7 +17,7 @@ import vllm.version from vllm.config import ProfilerConfig -from vllm.config.profiler import _is_uri_path +from vllm.config.profiler import ProfilerKind, TorchProfilerActivity, _is_uri_path from vllm.logger import init_logger from vllm.platforms import current_platform @@ -57,6 +56,11 @@ def is_running(self) -> bool: """Whether the underlying profiler is currently collecting data.""" return self._running + @property + def should_annotate(self) -> bool: + """Whether worker iterations should receive profiler annotations.""" + return self._running + @abstractmethod def _start(self) -> None: """Start the profiler.""" @@ -176,7 +180,6 @@ def annotate_context_manager(self, name: str): return nullcontext() -TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"] TorchProfilerActivityMap = { "CPU": torch.profiler.ProfilerActivity.CPU, "CUDA": torch.profiler.ProfilerActivity.CUDA, @@ -191,7 +194,7 @@ def __init__( profiler_config: ProfilerConfig, worker_name: str, local_rank: int, - activities: list[TorchProfilerActivity], + activities: Sequence[TorchProfilerActivity], on_trace_ready: Callable[[torch.profiler.profile], None] | None = None, ) -> None: super().__init__(profiler_config) @@ -224,7 +227,12 @@ def __init__( use_gzip=profiler_config.torch_profiler_use_gzip, ) - self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1 + self._records_cpu_activity = "CPU" in activities + self.dump_device_time_total = ( + any(activity != "CPU" for activity in activities) + and profiler_config.torch_profiler_dump_cuda_time_total + ) + self.dump_cpu_time_total = self._records_cpu_activity and len(activities) == 1 # Create profiler schedule if warmup or wait iterations are configured profiler_schedule = None @@ -244,7 +252,7 @@ def __init__( profiler_config.active_iterations, ) - self.profiler = torch.profiler.profile( + self._profiler_kwargs = dict( activities=[TorchProfilerActivityMap[activity] for activity in activities], schedule=profiler_schedule, record_shapes=profiler_config.torch_profiler_record_shapes, @@ -253,6 +261,7 @@ def __init__( with_flops=profiler_config.torch_profiler_with_flops, on_trace_ready=trace_handler, ) + self.profiler: torch.profiler.profile # Track if we're using a schedule (need to call step()) self._uses_schedule = profiler_schedule is not None @@ -260,10 +269,11 @@ def __init__( # Subtract 1 because profiler.start() already consumes step 0 # (WAIT or WARMUP), so only wait + warmup - 1 non-active steps # remain to be advanced through via profiler.step() calls. - self._warmup_steps_remaining = max( + self._initial_warmup_steps_remaining = max( profiler_config.wait_iterations + profiler_config.warmup_iterations - 1, 0, ) + self._warmup_steps_remaining = self._initial_warmup_steps_remaining self._version_metadata_added = False def _build_profiler_table( @@ -319,6 +329,9 @@ def _maybe_add_version_metadata(self) -> None: @override def _start(self) -> None: + self.profiler = torch.profiler.profile(**self._profiler_kwargs) + self._warmup_steps_remaining = self._initial_warmup_steps_remaining + self._version_metadata_added = False self.profiler.start() # No-schedule case: Kineto is live immediately. With a schedule this # no-ops and _profiler_step stamps it once WAIT ends. @@ -328,10 +341,9 @@ def _start(self) -> None: def _stop(self) -> None: self.profiler.stop() - profiler_config = self.profiler_config rank = self.local_rank - if profiler_config.torch_profiler_dump_cuda_time_total: - table = self._build_profiler_table(sort_key="self_cuda_time_total") + if self.dump_device_time_total: + table = self._build_profiler_table(sort_key="self_device_time_total") self._write_profiler_table(rank, table) # only print profiler results on rank 0 @@ -367,8 +379,15 @@ def _profiler_step(self) -> bool: return False return True + @property + @override + def should_annotate(self) -> bool: + return self._running and self._records_cpu_activity + @override def annotate_context_manager(self, name: str): + if not self.should_annotate: + return nullcontext() return torch.profiler.record_function(name) @@ -606,6 +625,82 @@ def annotate_context_manager(self, name: str): return torch.cuda.nvtx.range(name) +_DEFAULT_TORCH_PROFILER_ACTIVITIES: dict[str, tuple[TorchProfilerActivity, ...]] = { + "cpu": ("CPU",), + "cuda": ("CPU", "CUDA"), + "xpu": ("CPU", "XPU"), +} +_SUPPORTED_TORCH_PROFILER_ACTIVITIES = { + device: frozenset(activities) + for device, activities in _DEFAULT_TORCH_PROFILER_ACTIVITIES.items() +} +_SUPPORTED_PROFILER_KINDS: dict[str, frozenset[ProfilerKind]] = { + "cpu": frozenset(("torch",)), + "cuda": frozenset(("torch", "cuda", "proton")), + "xpu": frozenset(("torch",)), +} + + +def validate_worker_profiler_config(profiler_config: ProfilerConfig) -> None: + """Validate profiler selections against the worker's capabilities.""" + profiler_type = profiler_config.profiler + if profiler_type is None: + return + device_type = current_platform.device_type + if device_type not in _SUPPORTED_PROFILER_KINDS: + raise ValueError(f"Unsupported profiler device type: {device_type}") + supported_kinds = _SUPPORTED_PROFILER_KINDS[device_type] + if profiler_type not in supported_kinds: + supported_names = ", ".join(sorted(supported_kinds)) + raise ValueError( + f"Unsupported profiler type for {device_type}: " + f"{profiler_type}. Supported profiler types: {supported_names}." + ) + if profiler_type == "torch": + default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[device_type] + supported_activities = _SUPPORTED_TORCH_PROFILER_ACTIVITIES[device_type] + configured = profiler_config.torch_profiler_activities + activities = default_activities if configured is None else configured + unsupported = set(activities) - supported_activities + if unsupported: + unsupported_names = ", ".join(sorted(unsupported)) + supported_names = ", ".join(sorted(supported_activities)) + raise ValueError( + f"Unsupported torch profiler activities for " + f"{device_type}: {unsupported_names}. " + f"Supported activities: {supported_names}." + ) + + +def create_worker_profiler( + profiler_config: ProfilerConfig, + *, + worker_name: str, + local_rank: int, +) -> WorkerProfiler: + """Create a profiler using a validated config and platform defaults.""" + profiler_type = profiler_config.profiler + if profiler_type == "torch": + default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[ + current_platform.device_type + ] + configured = profiler_config.torch_profiler_activities + logger.debug("Starting torch profiler with trace name: %s", worker_name) + return TorchProfilerWrapper( + profiler_config, + worker_name=worker_name, + local_rank=local_rank, + activities=default_activities if configured is None else tuple(configured), + ) + if profiler_type == "cuda": + logger.debug("Starting CUDA profiler") + return CudaProfilerWrapper(profiler_config) + + assert profiler_type == "proton", f"Unknown profiler type: {profiler_type}" + logger.debug("Starting Proton profiler with trace name: %s", worker_name) + return ProtonProfilerWrapper(profiler_config, worker_name=worker_name) + + def create_graph_capture_profiler( profiler_config: ProfilerConfig, global_rank: int ) -> WorkerProfiler | None: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 67b8ac560d52..ee913b052f62 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -208,9 +208,11 @@ def __init__( pass self.profiler = profiler + configured_activities = vllm_config.profiler_config.torch_profiler_activities if ( vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend + and (configured_activities is None or "CPU" in configured_activities) ): profiler_dir = vllm_config.profiler_config.torch_profiler_dir logger.info( diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index 1badcfe6e82e..e36b15300b25 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -7,7 +7,6 @@ import math import os import sys -from typing import Any import psutil import torch @@ -16,7 +15,6 @@ from vllm.config import CompilationMode, VllmConfig from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform -from vllm.profiler.wrapper import TorchProfilerWrapper from vllm.utils.cpu_resource_utils import ( get_allowed_cpu_list, get_memory_node_info, @@ -106,18 +104,6 @@ def __init__( self.parallel_config.disable_custom_all_reduce = True - # Torch profiler. Enabled and configured through profiler_config. - self.profiler: Any | None = None - profiler_config = vllm_config.profiler_config - if profiler_config.profiler == "torch": - worker_name = f"{vllm_config.instance_id}-rank-{self.rank}" - self.profiler = TorchProfilerWrapper( - profiler_config, - worker_name=worker_name, - local_rank=self.local_rank, - activities=["CPU"], - ) - def init_device(self): self.device = torch.device("cpu") @@ -287,11 +273,3 @@ def compile_or_warm_up_model(self) -> CompilationTimes: language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, ) - - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - if self.profiler is None: - raise RuntimeError("Profiler is not enabled.") - if is_start: - self.profiler.start() - else: - self.profiler.stop() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 8ced3bdd4307..d914bd2deab3 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -61,10 +61,9 @@ from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory from vllm.platforms import current_platform from vllm.profiler.wrapper import ( - CudaProfilerWrapper, - ProtonProfilerWrapper, - TorchProfilerWrapper, create_graph_capture_profiler, + create_worker_profiler, + validate_worker_profiler_config, ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask @@ -222,6 +221,7 @@ def __init__( # so we have all the information needed for proper trace naming. self.profiler: Any | None = None self.profiler_config = vllm_config.profiler_config + validate_worker_profiler_config(self.profiler_config) self.use_v2_model_runner = vllm_config.use_v2_model_runner @@ -1038,7 +1038,7 @@ def annotate_profile(self, scheduler_output): return nullcontext() self.profiler.step() - if not self.profiler.is_running: + if not self.profiler.should_annotate: return nullcontext() iteration_details = compute_iteration_details(scheduler_output) @@ -1262,7 +1262,6 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): ) if is_start: - profiler_type = self.profiler_config.profiler # Generate the trace name by combining prefix with comprehensive rank suffix from vllm.distributed.utils import get_worker_rank_suffix @@ -1274,36 +1273,16 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): else: trace_name = rank_suffix - if profiler_type == "proton" and self.profiler is not None: + if self.profiler_config.profiler == "proton" and self.profiler is not None: self.profiler.set_output_name(trace_name) # Create the profiler wrapper only on the first start call if self.profiler is None: - if profiler_type == "torch": - self.profiler = TorchProfilerWrapper( - self.profiler_config, - worker_name=trace_name, - local_rank=self.local_rank, - activities=["CPU", "CUDA"], - ) - logger.debug( - "Starting torch profiler with trace name: %s", trace_name - ) - elif profiler_type == "cuda": - self.profiler = CudaProfilerWrapper(self.profiler_config) - logger.debug("Starting CUDA profiler") - elif profiler_type == "proton": - self.profiler = ProtonProfilerWrapper( - self.profiler_config, worker_name=trace_name - ) - logger.debug( - "Starting Proton profiler with trace name: %s", trace_name - ) - else: - # Config validation should prevent this code being reached - raise ValueError( - f"Invalid profiler value of {self.profiler_config.profiler}" - ) + self.profiler = create_worker_profiler( + self.profiler_config, + worker_name=trace_name, + local_rank=self.local_rank, + ) self.profiler.start() else: diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 5c5d125f4bfa..02cc2c9e87d5 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -8,7 +8,6 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.profiler.wrapper import TorchProfilerWrapper from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats @@ -225,33 +224,6 @@ def init_device(self): # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - if self.profiler_config is None or self.profiler_config.profiler is None: - raise RuntimeError( - "Profiling is not enabled. Please set --profiler-config to enable " - "profiling. Example: " - "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir" - "=YOUR_DIR_PATH_TO_DUMP_TRACE'" - ) - - if is_start and self.profiler is None: - from vllm.distributed.utils import get_worker_rank_suffix - - rank_suffix = get_worker_rank_suffix(global_rank=self.rank) - trace_name = ( - f"{profile_prefix}_{rank_suffix}" if profile_prefix else rank_suffix - ) - - self.profiler = TorchProfilerWrapper( - self.profiler_config, - worker_name=trace_name, - local_rank=self.local_rank, - activities=["CPU", "XPU"], - ) - logger.debug("Starting torch profiler with trace name: %s", trace_name) - - super().profile(is_start=is_start, profile_prefix=profile_prefix) - def shutdown(self) -> None: logger.info( "XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)", From 86ce4d10e290e6179e6f1d5ee392f582814d5747 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Mon, 21 Sep 2026 10:59:11 +0800 Subject: [PATCH 556/774] [BUGFIX][HY4] Record indexer completion event for full CUDA graph capture (#57811) Signed-off-by: Jee Jee Li --- vllm/models/hy_v4/nvidia/attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/models/hy_v4/nvidia/attention.py b/vllm/models/hy_v4/nvidia/attention.py index 5e0a33f18373..7f963cc707ee 100644 --- a/vllm/models/hy_v4/nvidia/attention.py +++ b/vllm/models/hy_v4/nvidia/attention.py @@ -758,6 +758,8 @@ def _indexer_and_attn( """ if self.indexer is not None and self.is_sparse and not self.skip_topk: self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) + if self.is_sparse: + self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined] out.copy_( self.mla_attn( q, From eb87980585c10f7e5ee8f498fd3830898101d208 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 20 Sep 2026 20:14:00 -0700 Subject: [PATCH 557/774] [Build] Fix DeepGEMM CUDA 12.9 release builds (#57554) Signed-off-by: khluu Co-authored-by: OpenAI Codex Co-authored-by: Jee Jee Li --- cmake/external_projects/deepgemm.cmake | 7 ++++--- tools/install_deepgemm.sh | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 9ac0d05414c1..a17031cc5e0f 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -29,10 +29,11 @@ if(DEEPGEMM_SRC_DIR) else() # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") - # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 + # Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120 # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU - # alpha/beta, and SiTU for FP8/FP4 Mega MoE. - set(_DEEPGEMM_UPSTREAM_TAG "a6bbb8000161c0dc3a85a0300a905f76898a7913") + # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x + # layout header fix from vllm-project/DeepGEMM#12. + set(_DEEPGEMM_UPSTREAM_TAG "e1f418c2a4f20818221f6b0e578b4c2f634d4c3f") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index a3ec840cd897..949aa28da533 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -7,10 +7,10 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" -# NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 +# NOTE: This targets the vLLM fork's dev branch, which carries the sm120 # and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE -# activations -DEEPGEMM_GIT_REF="a6bbb8000161c0dc3a85a0300a905f76898a7913" +# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. +DEEPGEMM_GIT_REF="e1f418c2a4f20818221f6b0e578b4c2f634d4c3f" WHEEL_DIR="" # Parse command line arguments From 9a70c233cd5d3ef3b14d79c582a6f9ddfa03cf80 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 20 Sep 2026 20:28:06 -0700 Subject: [PATCH 558/774] [Perf][DSV4.1] Fuse TP all-reduce with mHC input preparation (#57643) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- CMakeLists.txt | 2 + csrc/libtorch_stable/all_reduce_mhc.cu | 212 ++++++++++++++++++ tests/distributed/test_custom_all_reduce.py | 111 +++++++++ tests/distributed/test_engram_dp_shard.py | 1 + tests/kernels/test_mhc_kernels.py | 46 +++- vllm/models/deepseek_v4/nvidia/model.py | 3 + vllm/models/deepseek_v41/nvidia/model.py | 21 +- .../deepseek_v41/nvidia/ops/mega_mhc.py | 98 -------- vllm/models/deepseek_v41/nvidia/ops/mhc.py | 160 ++++++++++++- 9 files changed, 549 insertions(+), 105 deletions(-) create mode 100644 csrc/libtorch_stable/all_reduce_mhc.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index fbacfd9091ce..dbf4ca90e2c3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -460,6 +460,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) + list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/all_reduce_mhc.cu") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") diff --git a/csrc/libtorch_stable/all_reduce_mhc.cu b/csrc/libtorch_stable/all_reduce_mhc.cu new file mode 100644 index 000000000000..dc58371156a2 --- /dev/null +++ b/csrc/libtorch_stable/all_reduce_mhc.cu @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "torch_utils.h" +#include "custom_all_gather_reduce_scatter.cuh" +#include +#include + +namespace { +using Tensor = torch::stable::Tensor; +using Pack = vllm::packed_t::P; +constexpr int kHidden = 5120; +constexpr int kRanks = 4; +constexpr int kThreads = 128; +constexpr int kCluster = 5; +constexpr int kPacks = kHidden / Pack::size; + +__global__ __launch_bounds__(kThreads) + __cluster_dims__(1, kCluster, 1) void all_reduce_mhc_kernel( + const Pack* input, const Pack* residual, const float* post, + const float* comb, const float* pre, const Pack* weight, Pack* output, + Pack* normalized, Pack* local, Pack* multicast, uint32_t* epochs, + int rank, int stage_size, float eps) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + const int token = blockIdx.x; + const int column = blockIdx.y * kThreads + threadIdx.x; + const int idx = token * kPacks + column; + const int size = gridDim.x * kPacks; + const int stage = epochs[0] % 3; + const int dirty = vllm::mnnvl_lamport_dirty_stage(stage); + const int dirty_size = epochs[2 + dirty]; + Pack value = input[idx]; + vllm::store_multimem_lamport_payload( + multicast + stage * stage_size + rank * size + idx, + vllm::sanitize_lamport_payload(value)); + // Multi-token forwards benefit from launching dependents before mHC finishes. + if (gridDim.x > 1) cudaTriggerProgrammaticLaunchCompletion(); + vllm::lamport_cta_arrive(epochs + 1); + for (int i = idx; i < dirty_size; i += size) { + local[dirty * stage_size + i] = vllm::lamport_sentinel(); + } + Pack peers[kRanks]; + bool pending; + do { + vllm::LamportPack packets[kRanks]; + #pragma unroll + for (int peer = 0; peer < kRanks; ++peer) { + auto ptr = local + stage * stage_size + idx + peer * size; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(packets[peer].words[0]), "=r"(packets[peer].words[1]), + "=r"(packets[peer].words[2]), "=r"(packets[peer].words[3]) + : "l"(ptr) + : "memory"); + } + pending = false; + #pragma unroll + for (int peer = 0; peer < kRanks; ++peer) { + pending |= vllm::is_lamport_dirty(packets[peer]); + peers[peer] = packets[peer].packed; + } + } while (pending); + auto sum = vllm::upcast(peers[0]); + #pragma unroll + for (int r = 1; r < kRanks; ++r) { + vllm::packed_assign_add(sum, vllm::upcast(peers[r])); + } + // Keep the BF16 boundary between the collective and mHC post. + auto reduced = vllm::upcast(vllm::downcast(sum)); + decltype(sum) streams[4]; + #pragma unroll + for (int r = 0; r < 4; ++r) { + streams[r] = vllm::upcast(residual[(token * 4 + r) * kPacks + column]); + } + decltype(sum) collapsed{}; + #pragma unroll + for (int c = 0; c < 4; ++c) { + decltype(sum) mixed; + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + // Match post*x followed by the four residual FMAs. + float acc = __fmul_rn(post[token * 4 + c], reduced.data[e]); + #pragma unroll + for (int r = 0; r < 4; ++r) { + acc = fmaf(comb[token * 16 + r * 4 + c], streams[r].data[e], acc); + } + mixed.data[e] = acc; + } + Pack rounded = vllm::downcast(mixed); + output[(token * 4 + c) * kPacks + column] = rounded; + auto as_float = vllm::upcast(rounded); + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + collapsed.data[e] = + fmaf(pre[token * 4 + c], as_float.data[e], collapsed.data[e]); + } + } + // RMSNorm consumes the rounded BF16 collapse in the unfused path. + collapsed = vllm::upcast(vllm::downcast(collapsed)); + float sqr = 0; + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + sqr = fmaf(collapsed.data[e], collapsed.data[e], sqr); + } + #pragma unroll + for (int delta = 16; delta; delta /= 2) { + sqr += __shfl_xor_sync(0xffffffff, sqr, delta); + } + __shared__ float warp_sums[kCluster][kThreads / 32]; + auto cluster = cooperative_groups::this_cluster(); + const int lane = threadIdx.x % 32; + if (lane < kCluster) { + *cluster.map_shared_rank(&warp_sums[blockIdx.y][threadIdx.x / 32], lane) = + sqr; + } + cluster.sync(); + float total = 0; + #pragma unroll + for (int c = 0; c < kCluster; ++c) { + #pragma unroll + for (int w = 0; w < kThreads / 32; ++w) total += warp_sums[c][w]; + } + const float scale = rsqrtf(__fadd_rn(total / kHidden, eps)); + auto weights = vllm::upcast(weight[column]); + #pragma unroll + for (int e = 0; e < Pack::size; ++e) { + collapsed.data[e] = collapsed.data[e] * scale * weights.data[e]; + } + normalized[idx] = vllm::downcast(collapsed); + if (idx == 0) { + while (*reinterpret_cast(epochs + 1) < + gridDim.x * gridDim.y) { + } + epochs[2 + stage] = kRanks * size; + epochs[0] = vllm::mnnvl_lamport_next_stage(stage); + epochs[1] = 0; + } + if (gridDim.x == 1) cudaTriggerProgrammaticLaunchCompletion(); +#else + asm volatile("trap;"); +#endif +} + +void all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, + Tensor pre, Tensor weight, Tensor output, Tensor normalized, + int64_t local, int64_t multicast, Tensor epochs, + int64_t rank, int64_t stage_bytes, double eps) { + using torch::headeronly::ScalarType; + const torch::stable::accelerator::DeviceGuard guard(input.get_device_index()); + const int64_t n = input.size(0); + STD_TORCH_CHECK(input.dim() == 2 && input.size(1) == kHidden && n > 0 && + n <= 16); + for (const auto& t : {input, residual, weight, output, normalized}) { + STD_TORCH_CHECK(t.is_contiguous() && + t.scalar_type() == ScalarType::BFloat16); + STD_TORCH_CHECK(t.get_device_index() == input.get_device_index()); + STD_TORCH_CHECK(reinterpret_cast(t.const_data_ptr()) % 16 == 0); + } + for (const auto& t : {post, comb, pre}) { + STD_TORCH_CHECK(t.is_contiguous() && t.scalar_type() == ScalarType::Float); + STD_TORCH_CHECK(t.get_device_index() == input.get_device_index()); + } + STD_TORCH_CHECK(residual.numel() == n * 4 * kHidden && + output.numel() == residual.numel() && + normalized.numel() == input.numel() && + weight.numel() == kHidden); + STD_TORCH_CHECK(post.numel() == n * 4 && pre.numel() == n * 4 && + comb.numel() == n * 16); + STD_TORCH_CHECK(epochs.is_contiguous() && + epochs.scalar_type() == ScalarType::Int && + epochs.numel() >= 5 && + epochs.get_device_index() == input.get_device_index()); + STD_TORCH_CHECK(local && multicast && rank >= 0 && rank < kRanks && + stage_bytes % 16 == 0 && + stage_bytes >= n * kHidden * 2 * kRanks); + STD_TORCH_CHECK(get_device_prop()->major == 10); + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{ + .gridDim = dim3(n, kCluster), + .blockDim = dim3(kThreads), + .dynamicSmemBytes = 0, + .stream = get_current_cuda_stream(input.get_device_index()), + .attrs = &attr, + .numAttrs = 1}; + STD_CUDA_CHECK(cudaLaunchKernelEx( + &config, all_reduce_mhc_kernel, + static_cast(input.const_data_ptr()), + static_cast(residual.const_data_ptr()), + static_cast(post.const_data_ptr()), + static_cast(comb.const_data_ptr()), + static_cast(pre.const_data_ptr()), + static_cast(weight.const_data_ptr()), + static_cast(output.mutable_data_ptr()), + static_cast(normalized.mutable_data_ptr()), + reinterpret_cast(local), reinterpret_cast(multicast), + static_cast(epochs.mutable_data_ptr()), static_cast(rank), + static_cast(stage_bytes / sizeof(Pack)), static_cast(eps))); +} +} // namespace + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, mhc) { + mhc.def( + "all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, " + "Tensor pre, Tensor weight, Tensor! output, Tensor! normalized, " + "int local, int multicast, Tensor! epochs, int rank, int stage_bytes, " + "float eps) -> ()"); +} +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, mhc) { + mhc.impl("all_reduce_mhc", TORCH_BOX(&all_reduce_mhc)); +} diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 7f91f11d7c33..7119e596a097 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -11,6 +11,7 @@ from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa from vllm.distributed.device_communicators import custom_all_reduce as car from vllm.distributed.parallel_state import get_tp_group, graph_capture +from vllm.platforms import current_platform from ..utils import ( ensure_model_parallel_initialized, @@ -24,6 +25,116 @@ test_sizes[i] -= v % 8 +@ray.remote(num_gpus=1, max_calls=1) +def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_post_tilelang, + mhc_pre_delayed_tilelang, + ) + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + comm = get_tp_group().device_communicator.ca_comm + assert comm is not None and comm.mnnvl_lamport_ag_multicast_ptr + fn = torch.zeros(24, 20480, device=device) + scale = torch.ones(3, device=device) + base = torch.zeros(24, device=device) + + # Changing shapes and interleaving all-gather exercise the shared + # Lamport stages, including cleanup of a larger previous payload. + def run(n): + torch.manual_seed(42 + rank) + x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) + # Packed +0/-0 pairs collide with the Lamport sentinel. + x[:, :16] = 0 + x[:, 9:16:2] = -0.0 + torch.manual_seed(123) + residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16) + post = torch.rand(n, 4, device=device) + comb = torch.randn(n, 4, 4, device=device) * 0.1 + pre = torch.rand(n, 4, device=device) + weight = torch.randn(5120, device=device, dtype=torch.bfloat16) + output = torch.empty_like(residual) + normalized = torch.empty_like(x) + + def fused(): + torch.ops._C_custom_ar.all_reduce_mhc( + x, + residual, + post, + comb, + pre, + weight, + output, + normalized, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_epochs[0], + rank, + comm.mnnvl_buffer_size, + 1e-6, + ) + + def check(): + gathered = comm.custom_all_gather(x) + assert gathered is not None + peers = gathered.view(tp_size, n, 5120).float() + reduced = peers[0].clone() + for peer in peers[1:]: + reduced.add_(peer) + expected = mhc_post_tilelang( + reduced.bfloat16(), residual, post.unsqueeze(-1), comb + ) + expected_norm = mhc_pre_delayed_tilelang( + expected, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 2.0, + 20, + pre_mix=pre, + norm_weight=weight, + )[2] + torch.testing.assert_close(output, expected, rtol=0.008, atol=1e-6) + torch.testing.assert_close( + normalized, expected_norm, rtol=0.008, atol=0.008 + ) + + fused() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(6): + fused() + for _ in range(20): + graph.replay() + check() + x.mul_(0.5) + residual.mul_(2) + graph.replay() + check() + + # Cover fixed Q6 and shrinking/growing adaptive verification batches. + for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1): + run(n) + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), reason="Requires SM100" +) +def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch): + if torch.accelerator.device_count() < 4: + pytest.skip("Requires four GPUs with NVLink multicast") + multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) + + @pytest.mark.parametrize( ("dtype", "expected"), [ diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index c2fa0f857fc8..dd7105a23ba9 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -512,6 +512,7 @@ def __call__(self, hidden, positions, input_ids, *args, **kwargs): model = SimpleNamespace( use_mega_moe=False, use_sequence_parallel=False, + fuse_mhc_all_reduce=False, engram_hash=state, engram_swa_prefix="swa", engram_dp_shared_memory=dp_shared_memory, diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 8f7df556023d..63648ebf2d91 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -50,6 +50,40 @@ DEVICE = current_platform.device_type +@pytest.mark.parametrize( + "tp,ep,hidden,hc,multicast,expected", + [ + (4, False, 5120, 4, 1, True), + (2, False, 5120, 4, 1, False), + (4, True, 5120, 4, 1, False), + (4, False, 4096, 4, 1, False), + (4, False, 5120, 2, 1, False), + (4, False, 5120, 4, 0, False), + ], +) +def test_deepseek_v41_all_reduce_fusion_requires_kernel_support( + monkeypatch, tp, ep, hidden, hc, multicast, expected +): + """Eligibility depends on kernel inputs, not serving or attention settings.""" + from vllm.models.deepseek_v41.nvidia.ops import mhc + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=tp, enable_expert_parallel=ep + ), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(hidden_size=hidden, hc_mult=hc) + ), + ) + group = SimpleNamespace( + device_communicator=SimpleNamespace( + ca_comm=SimpleNamespace(mnnvl_lamport_ag_multicast_ptr=multicast) + ) + ) + monkeypatch.setattr(mhc, "get_tp_group", lambda: group) + assert mhc.supports_mhc_all_reduce(config) is expected + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize( "num_tokens,hc_mult,hidden_size", @@ -523,6 +557,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( decoder.rms_norm_eps = 1e-20 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.fuse_mhc_all_reduce = False decoder.mhc_stream = None if mhc_mode != "disabled": from vllm.utils.deep_gemm import is_deep_gemm_supported @@ -597,7 +632,15 @@ def reference(*args, norm_weight, norm_eps, **kwargs): return post, res, decoder.attn_norm(collapsed), pre def fused_reference( - x, residual, post_mix, res_mix, *args, capture_aux=False, stream=None, **kw + x, + residual, + post_mix, + res_mix, + *args, + capture_aux=False, + stream=None, + reduce_results=False, + **kw, ): residual = mhc_post_torch(x, residual, post_mix, res_mix) aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0) @@ -638,6 +681,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf decoder.rms_norm_eps = 1e-6 decoder.hc_post_alpha = 2.0 decoder.use_sequence_parallel = False + decoder.fuse_mhc_all_reduce = False decoder.mhc_stream = None decoder.engram = None from vllm.model_executor.layers.layernorm import RMSNorm diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 39d3c87bc64f..7edff8ae0202 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -793,6 +793,7 @@ def __init__( use_sequence_parallel: bool = False, *, num_hash_layers: int, + reduce_results: bool = True, n_routed_experts: int | None = None, n_activated_experts: int | None = None, image_sentinel_lo: int = IMAGE_SENTINEL_BASE_ID, @@ -804,6 +805,7 @@ def __init__( quant_config = vllm_config.quant_config self.prefix = prefix self.use_sequence_parallel = use_sequence_parallel + self.reduce_results = reduce_results moe_backend = vllm_config.kernel_config.moe_backend validate_fi_moe_ep_config(vllm_config) self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS @@ -1010,6 +1012,7 @@ def _init_fused_moe_experts( self.physical_expert_end = self.experts_end_idx self.experts = FusedMoEFactory( + reduce_results=self.reduce_results, shared_experts=self.shared_experts, gate=self.gate, num_experts=self.n_routed_experts, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b5caceab6fca..d7c9ddc1284f 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -17,6 +17,7 @@ get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( get_forward_context, @@ -88,10 +89,11 @@ from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes -from .ops.mega_mhc import mhc_shifted_post_pre from .ops.mhc import ( MHC_OVERLAP_MAX_TOKENS, mhc_pre_delayed_overlap, + mhc_shifted_post_pre, + supports_mhc_all_reduce, supports_mhc_overlap, ) @@ -107,6 +109,7 @@ def __init__( vllm_config: VllmConfig, prefix: str = "", use_sequence_parallel: bool = False, + reduce_results: bool = True, ): config = vllm_config.model_config.hf_config n_routed_experts = config.n_routed_experts @@ -125,6 +128,7 @@ def __init__( n_routed_experts=n_routed_experts, n_activated_experts=n_activated_experts, num_hash_layers=0, + reduce_results=reduce_results, image_sentinel_lo=IMAGE_SENTINEL_BASE_ID, ) @@ -199,12 +203,14 @@ def __init__( candidate_block_buffer: torch.Tensor | None = None, engram_layout: EngramLayout | None = None, mhc_stream: torch.cuda.Stream | None = None, + fuse_mhc_all_reduce: bool = False, ): super().__init__() config = vllm_config.model_config.hf_config self.hidden_size = config.hidden_size self.mhc_stream = mhc_stream + self.fuse_mhc_all_reduce = fuse_mhc_all_reduce self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.engram: Engram | None = None @@ -228,12 +234,13 @@ def __init__( aux_stream_list=aux_stream_list, candidate_block_buffer=candidate_block_buffer, ) - if self.use_sequence_parallel: + if self.use_sequence_parallel or fuse_mhc_all_reduce: self.attn.wo_b.reduce_results = False self.ffn = DeepseekV4MoE( vllm_config, prefix=f"{prefix}.ffn", use_sequence_parallel=self.use_sequence_parallel, + reduce_results=not fuse_mhc_all_reduce, ) self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) @@ -430,6 +437,8 @@ def forward( # and this block's pre, on the full hc stream, so the mix # coefficients see the injected stream. The injection also keeps # the post out of the pre-norm GEMM's fused prologue. + if self.fuse_mhc_all_reduce: + x = tensor_model_parallel_all_reduce(x) previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix) if capture_previous_aux: previous_aux = previous_post.mean(dim=1) @@ -473,6 +482,7 @@ def forward( norm_eps=self.attn_norm.variance_epsilon, capture_aux=capture_previous_aux, stream=mhc_stream, + reduce_results=self.fuse_mhc_all_reduce, ) if capture_previous_aux: previous_aux = aux @@ -503,6 +513,7 @@ def forward( norm_weight=self.ffn_norm.weight, norm_eps=self.ffn_norm.variance_epsilon, stream=mhc_stream, + reduce_results=self.fuse_mhc_all_reduce, ) x = self.ffn(x, input_ids, mega_gate_metadata) if mhc_stream is not None: @@ -540,6 +551,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): aux_stream_list = [torch.cuda.Stream() for _ in range(3)] # Keep mHC independent of the streams used inside attention. mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None + self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce( + vllm_config + ) # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -585,6 +599,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): candidate_block_buffer=self.candidate_block_buffer, engram_layout=self.engram_layout, mhc_stream=mhc_stream, + fuse_mhc_all_reduce=self.fuse_mhc_all_reduce, ), prefix=f"{prefix}.layers", ) @@ -794,6 +809,8 @@ def forward( aux_hidden_by_layer[idx] = previous_aux if layer is not None: # The last layer has no successor to fold its post into. + if self.fuse_mhc_all_reduce: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) hidden_states = mhc_post_tilelang( hidden_states, residual, post_mix, res_mix ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py index caae3d8d9b00..ef56b5a82393 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -5,10 +5,6 @@ import torch -from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, - mhc_post_tilelang, -) from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( _import_deep_gemm, @@ -88,97 +84,3 @@ def mhc_shifted_post_pre_deep_gemm( y_bf16, new_prev_mix.squeeze(-1), ) - - -def mhc_shifted_post_pre( - x: torch.Tensor, - residual: torch.Tensor, - post_layer_mix: torch.Tensor, - comb_res_mix: torch.Tensor, - fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, - rms_eps: float, - hc_pre_eps: float, - hc_sinkhorn_eps: float, - hc_post_mult_value: float, - sinkhorn_repeat: int, - pre_mix: torch.Tensor | None = None, - norm_weight: torch.Tensor | None = None, - norm_eps: float = 1e-6, - capture_aux: bool = False, - *, - stream: torch.cuda.Stream | None = None, -) -> tuple[ - torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor -]: - """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. - - When stream is supplied, join it before consuming the returned coefficients. - """ - if stream is not None: - from .mhc import mhc_pre_delayed_overlap - - residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix) - aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1]) - pre_outputs = mhc_pre_delayed_overlap( - residual, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - pre_mix=pre_mix, - norm_weight=norm_weight, - norm_eps=norm_eps, - stream=stream, - ) - return residual, *pre_outputs, aux - - if ( - pre_mix is not None - and norm_weight is not None - and not capture_aux - and x.shape[0] <= 1 << 20 - and is_mega_mhc_supported(x.shape[1], residual.shape[1]) - ): - outputs = mhc_shifted_post_pre_deep_gemm( - x, - residual, - pre_mix, - post_layer_mix, - comb_res_mix, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_post_mult_value, - hc_sinkhorn_eps, - sinkhorn_repeat, - norm_weight, - norm_eps, - ) - return *outputs, x.new_empty(0, x.shape[1]) - - return mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_layer_mix, - comb_res_mix, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - pre_mix=pre_mix, - norm_weight=norm_weight, - norm_eps=norm_eps, - capture_aux=capture_aux, - ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index 29c9c3b95684..084a102436ed 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -1,15 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Split shifted mHC pre so coefficient generation can overlap the sublayer.""" +"""Dispatch DSV4.1 mHC operations and overlap coefficient generation.""" from functools import partial +from typing import TYPE_CHECKING, cast import torch from vllm.config import VllmConfig +from vllm.distributed import get_tp_group +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, + mhc_post_tilelang, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported +from .mega_mhc import is_mega_mhc_supported, mhc_shifted_post_pre_deep_gemm + +if TYPE_CHECKING: + from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator + # GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress. MHC_OVERLAP_MAX_TOKENS = 16 @@ -29,6 +40,20 @@ def supports_mhc_overlap(vllm_config: VllmConfig) -> bool: ) +def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: + parallel = vllm_config.parallel_config + config = vllm_config.model_config.hf_config + if ( + parallel.tensor_parallel_size != 4 + or parallel.enable_expert_parallel + or config.hidden_size != 5120 + or config.hc_mult != 4 + ): + return False + comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm + return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr) + + def mhc_pre_delayed_overlap( residual: torch.Tensor, fn: torch.Tensor, @@ -45,6 +70,7 @@ def mhc_pre_delayed_overlap( norm_eps: float = 1e-6, *, stream: torch.cuda.Stream, + layer_input: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Prepare the input on the caller stream and coefficients on another stream. @@ -67,7 +93,11 @@ def mhc_pre_delayed_overlap( post = torch.empty((n, hc), device=residual.device, dtype=torch.float32) comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32) next_pre = torch.empty_like(post) - layer_input = torch.empty((n, hidden), device=residual.device, dtype=torch.bfloat16) + compute_input = layer_input is None + if layer_input is None: + layer_input = torch.empty( + (n, hidden), device=residual.device, dtype=torch.bfloat16 + ) outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre if n == 0: return outputs @@ -104,7 +134,7 @@ def mhc_pre_delayed_overlap( ) main = torch.cuda.current_stream() # Prioritize input readiness for small batches before releasing statistics. - if n <= 8: + if compute_input and n <= 8: epilogue(split_mode="input") stream.wait_stream(main) with torch.cuda.stream(stream): @@ -112,6 +142,128 @@ def mhc_pre_delayed_overlap( epilogue(split_mode="stats") for tensor in (x, mix, sqr): tensor.record_stream(stream) - if n > 8: + if compute_input and n > 8: epilogue(split_mode="input") return outputs + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, + *, + stream: torch.cuda.Stream | None = None, + reduce_results: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang. + + When stream is supplied, join it before consuming the returned coefficients. + """ + layer_input = None + if reduce_results: + tp = get_tp_group() + if stream is not None and 0 < x.shape[0] <= MHC_OVERLAP_MAX_TOKENS: + comm = cast("CudaCommunicator", tp.device_communicator).ca_comm + assert comm is not None and comm.mnnvl_lamport_epochs is not None + output = torch.empty_like(residual) + layer_input = torch.empty_like(x) + torch.ops._C_custom_ar.all_reduce_mhc( + x, + residual, + post_layer_mix, + comb_res_mix, + pre_mix, + norm_weight, + output, + layer_input, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_epochs[0], + comm.rank, + comm.mnnvl_buffer_size, + norm_eps, + ) + residual = output + else: + x = tp.all_reduce(x) + if stream is not None: + if layer_input is None: + residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix) + aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1]) + pre_outputs = mhc_pre_delayed_overlap( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + stream=stream, + layer_input=layer_input, + ) + return residual, *pre_outputs, aux + + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( + x, + residual, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + return *outputs, x.new_empty(0, x.shape[1]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + capture_aux=capture_aux, + ) From ff0c8854946b75dc0fe160756e8adff10a8da053 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Sun, 20 Sep 2026 22:05:24 -0600 Subject: [PATCH 559/774] [Bugfix] DiffusionGemma: fix multimodal support (#57589) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/models/diffusion_gemma.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 10e0d8a2c224..aef9d816258b 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -197,6 +197,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): text_config.attention_k_eq_v = True # ---- Vision tower ---- + # Gemma4's image path, borrowed below, reads this flag. + lora_config = vllm_config.lora_config + self._enable_mm_lora = bool( + lora_config is not None and lora_config.enable_tower_connector_lora + ) vision_config = getattr(config, "vision_config", None) self.embed_vision: Gemma4MultimodalEmbedder | None if vision_config is not None: From 3df4ae153eb385e27b52f26c81f8edb9e20b9984 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Mon, 21 Sep 2026 12:46:11 +0800 Subject: [PATCH 560/774] [Core][KDA] Generalize Mamba prefill checkpoint builder and exporter (#57783) Signed-off-by: chaunceyjiang Signed-off-by: zjy0516 Co-authored-by: zjy0516 --- tests/models/kimi_k3/test_kda.py | 42 ++---- .../model_executor/layers/mamba/checkpoint.py | 137 ++++++++++++++++++ .../layers/mamba/kda_checkpoint.py | 130 +++++++++++++++++ vllm/models/kimi_k3/nvidia/kda.py | 119 +++------------ vllm/models/kimi_k3/nvidia/kda_metadata.py | 106 +++----------- 5 files changed, 314 insertions(+), 220 deletions(-) create mode 100644 vllm/model_executor/layers/mamba/checkpoint.py create mode 100644 vllm/model_executor/layers/mamba/kda_checkpoint.py diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 84a5e7eed409..f03d2cdcae3f 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -14,6 +14,10 @@ from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.model_executor.layers.mamba.checkpoint import MambaPrefillCheckpointMetadata +from vllm.model_executor.layers.mamba.kda_checkpoint import ( + FlashKDAPrefillCheckpointExporter, +) from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, @@ -25,7 +29,6 @@ from vllm.models.kimi_k3.nvidia.kda import ( _flashinfer_kda_prefill, _flashkda_prefill, - _store_cache_checkpoints_kernel, is_flashinfer_fused_kda_decode_supported, is_flashinfer_recurrent_kda_prefill_supported, is_flashkda_supported, @@ -1498,36 +1501,13 @@ def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: fl checkpoint_state_indices = torch.tensor( [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE ) - state_len = conv_state.shape[-1] - width = H * D - recurrent_row_size = checkpoint_state[0].numel() - block_size = 256 - _store_cache_checkpoints_kernel[ - ( - checkpoint_state_indices.numel(), - (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size, - ) - ]( - conv_input, - conv_state, - checkpoint_state, - recurrent_state, - cu_seqlens, - checkpoint_offsets, - checkpoint_state_indices, - conv_input.stride(0), - conv_input.stride(1), - conv_state.stride(0), - conv_state.stride(1), - conv_state.stride(2), - checkpoint_state.stride(0), - recurrent_state.stride(0), - checkpoint_offsets.stride(0), - state_len, - width, - recurrent_row_size, - NULL_BLOCK_ID, - block_size, + FlashKDAPrefillCheckpointExporter().export( + MambaPrefillCheckpointMetadata(checkpoint_offsets, checkpoint_state_indices), + raw_qkv=conv_input, + conv_state=conv_state, + recurrent_checkpoint=checkpoint_state, + recurrent_state=recurrent_state, + cu_seqlens=cu_seqlens, ) torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1)) torch.testing.assert_close(recurrent_state[1], checkpoint_state[0]) diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py new file mode 100644 index 000000000000..3ba4062079cc --- /dev/null +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + +from vllm.config import VllmConfig +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import ( + MambaSpec, + get_mamba_prefill_checkpoint_position, + is_mamba_prefill_checkpoint_valid, +) + + +def compute_mamba_prefill_checkpoints( + seq_lens: list[int], + query_lens: list[int], + hash_block_size: int, + mamba_block_size: int, + checkpoint_alignment: int | None, + drop_eagle_block: bool, +) -> tuple[list[int], list[int]]: + """Per-row internal prefill checkpoint offsets and cache block columns. + + Backends call this instead of re-deriving the rules, so they decline in + lockstep with the scheduler and ``MambaManager``: allocating a checkpoint + block without writing it leaves the prefix cache serving uninitialized + state. + + Returns: + ``(offsets, cols)``: the checkpoint's offset into each row's query and + its block-table column. ``0`` and ``-1`` mean the row has none. + + """ + offsets: list[int] = [] + cols: list[int] = [] + for seq_len, query_len in zip(seq_lens, query_lens): + query_start = seq_len - query_len + position = get_mamba_prefill_checkpoint_position( + seq_len, hash_block_size, drop_eagle_block=drop_eagle_block + ) + valid = is_mamba_prefill_checkpoint_valid( + query_start=query_start, + query_end=seq_len, + checkpoint_position=position, + hash_block_size=hash_block_size, + mamba_block_size=mamba_block_size, + checkpoint_alignment=checkpoint_alignment, + ) + offsets.append(position - query_start if valid else 0) + cols.append(cdiv(seq_len, mamba_block_size) - 2 if valid else -1) + return offsets, cols + + +@dataclass +class MambaPrefillCheckpointMetadata: + checkpoint_offsets: torch.Tensor + state_indices: torch.Tensor + + +class MambaPrefillCheckpointBuilder: + """Build backend-neutral prefill checkpoint locations.""" + + def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None: + self.vllm_config = vllm_config + self.kv_cache_spec = kv_cache_spec + + def build( + self, + m: CommonAttentionMetadata, + request_rows: list[int], + ) -> MambaPrefillCheckpointMetadata | None: + if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0: + return None + assert m.seq_lens_cpu_upper_bound is not None + all_query_lens = m.query_start_loc_cpu.diff().tolist() + query_lens = [all_query_lens[row] for row in request_rows] + seq_lens = m.seq_lens_cpu_upper_bound.tolist() + block_size = self.kv_cache_spec.block_size + hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size + speculative_config = self.vllm_config.speculative_config + drop_eagle_block = ( + speculative_config is not None and speculative_config.use_eagle_block_drop() + ) + checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints( + [seq_lens[row] for row in request_rows], + query_lens, + hash_block_size=hash_block_size, + mamba_block_size=block_size, + checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment, + drop_eagle_block=drop_eagle_block, + ) + if not any(checkpoint_offsets): + return None + checkpoint_offsets_tensor = async_tensor_h2d( + checkpoint_offsets, + dtype=torch.int32, + device=m.query_start_loc.device, + ) + request_rows_tensor = async_tensor_h2d( + request_rows, dtype=torch.int64, device=m.query_start_loc.device + ) + checkpoint_cols_tensor = async_tensor_h2d( + checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device + ) + checkpoint_state_indices = m.block_table_tensor[ + request_rows_tensor, checkpoint_cols_tensor + ] + checkpoint_state_indices = torch.where( + checkpoint_cols_tensor >= 0, + checkpoint_state_indices, + NULL_BLOCK_ID, + ) + return MambaPrefillCheckpointMetadata( + checkpoint_offsets_tensor, + checkpoint_state_indices, + ) + + +class MambaPrefillCheckpointExporter(ABC): + """Export a mid-prefill checkpoint into backend-specific paged states.""" + + @abstractmethod + def export( + self, + checkpoint: MambaPrefillCheckpointMetadata, + *args, + **kwargs, + ) -> None: + """Write checkpoint state into the paged cache.""" + raise NotImplementedError diff --git a/vllm/model_executor/layers/mamba/kda_checkpoint.py b/vllm/model_executor/layers/mamba/kda_checkpoint.py new file mode 100644 index 000000000000..20741623d105 --- /dev/null +++ b/vllm/model_executor/layers/mamba/kda_checkpoint.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointExporter, + MambaPrefillCheckpointMetadata, +) +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def kda_prefill_checkpoint_alignment(backend: str) -> int | None: + return 16 if backend == "flashkda" else None + + +@dataclass(frozen=True) +class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter): + """Store FlashKDA recurrent and convolution checkpoint states.""" + + state_len: int | None = None + + def export( + self, + checkpoint: MambaPrefillCheckpointMetadata, + *, + raw_qkv: torch.Tensor, + conv_state: torch.Tensor, + recurrent_checkpoint: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ) -> None: + state_len = ( + self.state_len if self.state_len is not None else conv_state.shape[-1] + ) + width = raw_qkv.shape[-1] + recurrent_row_size = recurrent_checkpoint[0].numel() + block_size = 256 + store_cache_checkpoints_kernel[ + ( + checkpoint.checkpoint_offsets.numel(), + triton.cdiv(max(width * state_len, recurrent_row_size), block_size), + ) + ]( + raw_qkv, + conv_state, + recurrent_checkpoint, + recurrent_state, + cu_seqlens, + checkpoint.checkpoint_offsets, + checkpoint.state_indices, + raw_qkv.stride(0), + raw_qkv.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + recurrent_checkpoint.stride(0), + recurrent_state.stride(0), + checkpoint.checkpoint_offsets.stride(0), + state_len, + width, + recurrent_row_size, + NULL_BLOCK_ID, + block_size, + ) + + +@triton.jit +def store_cache_checkpoints_kernel( + x_ptr, + conv_state_ptr, + recurrent_checkpoint_ptr, + recurrent_state_ptr, + query_start_loc_ptr, + checkpoint_offsets_ptr, + checkpoint_state_indices_ptr, + x_stride_0: tl.constexpr, + x_stride_1: tl.constexpr, + state_stride_0: tl.constexpr, + state_stride_1: tl.constexpr, + state_stride_2: tl.constexpr, + checkpoint_stride_0: tl.constexpr, + recurrent_state_stride_0: tl.constexpr, + checkpoint_offset_stride: tl.constexpr, + STATE_LEN: tl.constexpr, + WIDTH: tl.constexpr, + RECURRENT_ROW_SIZE: tl.constexpr, + NULL_STATE_IDX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + seq_idx = tl.program_id(0) + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64) + checkpoint_offset = tl.load( + checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride + ) + valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) + valid_conv = ( + (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) + ) + width_idx = cols // STATE_LEN + history_idx = cols % STATE_LEN + checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset + token_idx = checkpoint_end - STATE_LEN + history_idx + values = tl.load( + x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1, + mask=valid_conv, + ) + tl.store( + conv_state_ptr + + state_idx * state_stride_0 + + width_idx * state_stride_1 + + history_idx * state_stride_2, + values, + mask=valid_conv, + ) + + valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint + recurrent = tl.load( + recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols, + mask=valid_recurrent, + ) + tl.store( + recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols, + recurrent, + mask=valid_recurrent, + ) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 92326944f592..c920cdbbdfa4 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -21,6 +21,10 @@ RowParallelLinear, ) from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.kda_checkpoint import ( + FlashKDAPrefillCheckpointExporter, + kda_prefill_checkpoint_alignment, +) from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, @@ -49,7 +53,6 @@ from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig -from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import ( flashinfer_fused_kda_decode, flashinfer_recurrent_kda, @@ -57,7 +60,6 @@ has_flashinfer_recurrent_kda, ) from vllm.v1.attention.backend import AttentionBackend -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager @@ -389,69 +391,6 @@ def _flashinfer_kda_prefill( return output, initial_state -@triton.jit -def _store_cache_checkpoints_kernel( - x_ptr, - conv_state_ptr, - recurrent_checkpoint_ptr, - recurrent_state_ptr, - query_start_loc_ptr, - checkpoint_offsets_ptr, - checkpoint_state_indices_ptr, - x_stride_0: tl.constexpr, - x_stride_1: tl.constexpr, - state_stride_0: tl.constexpr, - state_stride_1: tl.constexpr, - state_stride_2: tl.constexpr, - checkpoint_stride_0: tl.constexpr, - recurrent_state_stride_0: tl.constexpr, - checkpoint_offset_stride: tl.constexpr, - STATE_LEN: tl.constexpr, - WIDTH: tl.constexpr, - RECURRENT_ROW_SIZE: tl.constexpr, - NULL_STATE_IDX: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - # store checkpoints to cache - seq_idx = tl.program_id(0) - cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64) - checkpoint_offset = tl.load( - checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride - ) - valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) - valid_conv = ( - (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) - ) - width_idx = cols // STATE_LEN - history_idx = cols % STATE_LEN - checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset - token_idx = checkpoint_end - STATE_LEN + history_idx - values = tl.load( - x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1, - mask=valid_conv, - ) - tl.store( - conv_state_ptr - + state_idx * state_stride_0 - + width_idx * state_stride_1 - + history_idx * state_stride_2, - values, - mask=valid_conv, - ) - - valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint - recurrent = tl.load( - recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols, - mask=valid_recurrent, - ) - tl.store( - recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols, - recurrent, - mask=valid_recurrent, - ) - - def resolve_kda_prefill_backend( backend: str, head_dim: int, @@ -735,6 +674,7 @@ def __init__( self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = ( None ) + self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None if self.kda_prefill_backend == "flashkda": T = vllm_config.scheduler_config.max_num_batched_tokens N = vllm_config.scheduler_config.max_num_seqs @@ -748,6 +688,7 @@ def __init__( ((N, H, D, D), self.get_state_dtype()[1]), ((workspace_size,), torch.uint8), ) + self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter() elif self.kda_prefill_backend == "flashinfer": T = vllm_config.scheduler_config.max_num_batched_tokens H, D = self.local_num_heads, self.head_dim @@ -800,12 +741,11 @@ def __init__( def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: spec = super().get_kv_cache_spec(vllm_config) assert isinstance(spec, MambaSpec) + alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend) return replace( spec, - num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"), - prefill_checkpoint_alignment=( - 16 if self.kda_prefill_backend == "flashkda" else None - ), + num_prefill_checkpoint_blocks=int(alignment is not None), + prefill_checkpoint_alignment=alignment, ) def forward( @@ -1173,39 +1113,14 @@ def _prefill_conv( ) core_attn_out_non_spec = flashkda_out last_recurrent_state = final_state - state_len = conv_state.shape[-1] - width = mixed_qkv_ns.shape[-1] - recurrent_row_size = checkpoint_state[0].numel() - block_size = 256 - _store_cache_checkpoints_kernel[ - ( - checkpoint_offsets.numel(), - triton.cdiv( - max(width * state_len, recurrent_row_size), - block_size, - ), - ) - ]( - mixed_qkv_ns, - conv_state, - checkpoint_state, - recurrent_state, - non_spec_query_start_loc, - checkpoint_offsets, - checkpoint.state_indices, - mixed_qkv_ns.stride(0), - mixed_qkv_ns.stride(1), - conv_state.stride(0), - conv_state.stride(1), - conv_state.stride(2), - checkpoint_state.stride(0), - recurrent_state.stride(0), - checkpoint_offsets.stride(0), - state_len, - width, - recurrent_row_size, - NULL_BLOCK_ID, - block_size, + assert self._checkpoint_exporter is not None + self._checkpoint_exporter.export( + checkpoint, + raw_qkv=mixed_qkv_ns, + conv_state=conv_state, + recurrent_checkpoint=checkpoint_state, + recurrent_state=recurrent_state, + cu_seqlens=non_spec_query_start_loc, ) else: ( diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index d342adb6a688..d1c2b319215c 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -18,9 +18,12 @@ import torch from vllm.config import VllmConfig +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointBuilder, + MambaPrefillCheckpointMetadata, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.gdn_attn import ( @@ -37,11 +40,7 @@ compute_causal_conv1d_metadata, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import ( - MambaSpec, - get_mamba_prefill_checkpoint_position, - is_mamba_prefill_checkpoint_valid, -) +from vllm.v1.kv_cache_interface import MambaSpec if TYPE_CHECKING: from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( @@ -262,12 +261,6 @@ class KDARecoverSSMCommitMetadata: align: KDARecoverSSMAlignMetadata | None -@dataclass -class KDACheckpointMetadata: - checkpoint_offsets: torch.Tensor - state_indices: torch.Tensor - - @dataclass class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): spec_token_start: int | None = None @@ -278,7 +271,7 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): recoverssm_context: "KDARecoverSSMCommitContext | None" = field( default=None, repr=False, compare=False ) - checkpoint: KDACheckpointMetadata | None = None + checkpoint: MambaPrefillCheckpointMetadata | None = None def commit_recoverssm_state( self, num_accepted_tokens: torch.Tensor @@ -322,6 +315,9 @@ def __init__( device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.checkpoint_builder = MambaPrefillCheckpointBuilder( + vllm_config, kv_cache_spec + ) additional_config = vllm_config.additional_config self.use_flashinfer_prefill = ( isinstance(additional_config, dict) @@ -400,8 +396,10 @@ def build( # type: ignore[override] if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: spec_sequence_masks_cpu = None num_spec_decodes = 0 + active_non_spec_mask_cpu = None else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + active_non_spec_mask_cpu = None if self.use_recoverssm: assert m.is_prefilling is not None assert m.is_prefilling.device.type == "cpu" @@ -418,6 +416,7 @@ def build( # type: ignore[override] ): spec_sequence_masks_cpu = None num_spec_decodes = 0 + active_non_spec_mask_cpu = None else: num_spec_decodes = spec_sequence_masks_cpu.sum().item() if num_spec_decodes == 0: @@ -624,80 +623,6 @@ def build( # type: ignore[override] else: has_initial_state = None - checkpoint = None - checkpoint_enabled = ( - self.vllm_config.cache_config.mamba_cache_mode == "align" - and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 - ) - if num_prefills > 0 and checkpoint_enabled: - # prepare checkpoint metadata - assert m.seq_lens_cpu_upper_bound is not None - request_rows = list(range(m.num_reqs)) - if spec_sequence_masks_cpu is not None: - # get non spec request rows - request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() - all_query_lens = query_start_loc_cpu.diff().tolist() - query_lens = [all_query_lens[row] for row in request_rows] - seq_lens = m.seq_lens_cpu_upper_bound.tolist() - block_size = self.kv_cache_spec.block_size - hash_block_size = ( - self.vllm_config.cache_config.prefix_match_unit or block_size - ) - speculative_config = self.vllm_config.speculative_config - drop_eagle_block = ( - speculative_config is not None - and speculative_config.use_eagle_block_drop() - ) - checkpoint_splits = [] - checkpoint_cols = [] - for row, query_len in zip(request_rows, query_lens): - seq_len = seq_lens[row] - query_start = seq_len - query_len - checkpoint_position = get_mamba_prefill_checkpoint_position( - seq_len, - hash_block_size, - drop_eagle_block=drop_eagle_block, - ) - offset = checkpoint_position - query_start - valid = is_mamba_prefill_checkpoint_valid( - query_start=query_start, - query_end=seq_len, - checkpoint_position=checkpoint_position, - hash_block_size=hash_block_size, - mamba_block_size=block_size, - checkpoint_alignment=( - self.kv_cache_spec.prefill_checkpoint_alignment - ), - ) - offset = offset if valid else 0 - first_len = offset or query_len - checkpoint_splits.append((first_len, query_len - first_len)) - checkpoint_cols.append(cdiv(seq_len, block_size) - 2 if valid else -1) - if any(tail for _, tail in checkpoint_splits): - checkpoint_offsets_tensor = async_tensor_h2d( - [first if tail else 0 for first, tail in checkpoint_splits], - dtype=torch.int32, - device=query_start_loc.device, - ) - request_rows_tensor = async_tensor_h2d( - request_rows, dtype=torch.int64, device=query_start_loc.device - ) - checkpoint_cols_tensor = async_tensor_h2d( - checkpoint_cols, dtype=torch.int64, device=query_start_loc.device - ) - checkpoint_state_indices = m.block_table_tensor[ - request_rows_tensor, checkpoint_cols_tensor - ] - checkpoint_state_indices = torch.where( - checkpoint_cols_tensor >= 0, - checkpoint_state_indices, - NULL_BLOCK_ID, - ) - checkpoint = KDACheckpointMetadata( - checkpoint_offsets_tensor, - checkpoint_state_indices, - ) - # Prepare per-request tensors for cudagraph replay. num_actual_tokens # may be token-padded, while state/query/acceptance metadata is indexed # by request. @@ -776,6 +701,13 @@ def build( # type: ignore[override] flashinfer_prefill_query_start_loc.diff(), descending=True ).to(torch.int32) + checkpoint = None + if num_prefills > 0: + request_rows = list(range(m.num_reqs)) + if active_non_spec_mask_cpu is not None: + request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() + checkpoint = self.checkpoint_builder.build(m, request_rows) + return KimiK3KDAMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, From a15dbbd63daaba913b3302a858ad7e20e6dcc219 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 21 Sep 2026 13:15:29 +0800 Subject: [PATCH 561/774] [XPU][CI] fallback to 7-args of moe_align_block_size (#57855) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- vllm/_custom_ops.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 1d38bcade013..dcecce71b3eb 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2274,6 +2274,19 @@ def moe_align_block_size( expert_map: torch.Tensor | None = None, scatter_idx: torch.Tensor | None = None, ) -> None: + if current_platform.is_xpu(): + if scatter_idx is not None: + raise NotImplementedError("scatter_idx is not supported on XPU") + torch.ops._moe_C.moe_align_block_size( + topk_ids, + num_experts, + block_size, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + expert_map, + ) + return torch.ops._moe_C.moe_align_block_size( topk_ids, num_experts, From 76ee20c1ec5a9d6e094ecaecd46d71a5d591d05c Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 00:21:35 -0500 Subject: [PATCH 562/774] [ROCm][Tests] Avoid pooling cleanup waits against live shared engines (#57864) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/models/language/pooling/conftest.py | 2 ++ tests/models/language/pooling_mteb_test/conftest.py | 3 +++ tests/models/multimodal/pooling/conftest.py | 3 +++ 3 files changed, 8 insertions(+) diff --git a/tests/models/language/pooling/conftest.py b/tests/models/language/pooling/conftest.py index 2da4c65133f5..ceb21b5a409e 100644 --- a/tests/models/language/pooling/conftest.py +++ b/tests/models/language/pooling/conftest.py @@ -33,6 +33,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: diff --git a/tests/models/language/pooling_mteb_test/conftest.py b/tests/models/language/pooling_mteb_test/conftest.py index 5a31a5f41580..ce34d82563a8 100644 --- a/tests/models/language/pooling_mteb_test/conftest.py +++ b/tests/models/language/pooling_mteb_test/conftest.py @@ -9,6 +9,7 @@ from tests.utils import wait_for_memory_to_settle from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -26,6 +27,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: diff --git a/tests/models/multimodal/pooling/conftest.py b/tests/models/multimodal/pooling/conftest.py index 32f598b30a8b..fdd58e2fba81 100644 --- a/tests/models/multimodal/pooling/conftest.py +++ b/tests/models/multimodal/pooling/conftest.py @@ -9,6 +9,7 @@ from tests.utils import wait_for_memory_to_settle from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -26,6 +27,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: From 6654bcbfca86a8c17c8a377f6813c7a86e64cd6d Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 21 Sep 2026 13:24:41 +0800 Subject: [PATCH 563/774] [MRV2][XPU] use xpu sample kernel in mrv2 sampler (#57277) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- vllm/v1/sample/ops/topk_topp_sampler.py | 98 ++++++++++++++++--------- vllm/v1/watermarking/gpu_sampler.py | 4 +- vllm/v1/worker/gpu/sample/sampler.py | 30 ++++++-- 3 files changed, 89 insertions(+), 43 deletions(-) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 477e99c2fabf..dac855456814 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import torch import torch.nn as nn @@ -127,6 +126,69 @@ def flashinfer_sampler_supported() -> bool: return False +def xpu_sampler_supported() -> bool: + """Decide whether the fused XPU top-k/top-p sampler kernel can be used. + + Returns False (with appropriate logging) when the platform isn't XPU, when + ``VLLM_XPU_USE_SAMPLER_KERNEL`` is 0. + + Note: callers must additionally ensure no request needs a per-request seed + or greedy sampling, since the kernel draws from the device's default + generator and always samples randomly. + """ + if not current_platform.is_xpu(): + return False + if not envs.VLLM_XPU_USE_SAMPLER_KERNEL: + logger.info_once( + "Fused XPU top-p/top-k sampling disabled via VLLM_XPU_USE_SAMPLER_KERNEL=0." + ) + return False + + logger.info_once("Using the fused XPU kernel for top-p & top-k sampling.") + return True + + +def xpu_sample( + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + logprobs_mode: LogprobsMode = "raw_logprobs", +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Sample from the logits using the fused XPU top-k/top-p kernel. + + Statistically equivalent to `random_sample`, but avoids sorting the vocab + and never materializes the probability tensor. The noise comes from the + device's default generator, so per-request generators aren't supported. + + Returns the sampled token ids and, for processed logprobs modes, the + post-top-k/top-p logits (or logprobs); otherwise None. + """ + logits = logits.to(dtype=torch.float32).contiguous() + sampled = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device) + logits_to_return = ( + torch.empty_like(logits) if logprobs_mode in PROCESSED_LOGPROBS_MODES else None + ) + + generator = torch.xpu.default_generators[logits.device.index] + state = generator.get_state() + seed, offset = state.view(torch.int64).tolist() + seeds = torch.tensor([seed, offset], dtype=torch.int64, device="cpu") + # The XPU kernel expects k as int64 (Long), but the input batch + # stores top_k as int32. Cast here to avoid dtype mismatch. + if k is not None: + k = k.to(torch.int64) + torch.ops.vllm.xpu_topk_topp_sampler( + sampled, logits_to_return, logits, k, p, logprobs_mode, seeds + ) + # The custom XPU sampler kernel consumes RNG values internally, so advance + # the default generator's offset to keep future draws deterministic. + # pytorch: offset must be multiple of 4 + offset = (offset + logits.numel() + 3) // 4 * 4 + state.view(torch.int64)[1] = offset + generator.set_state(state) + return sampled, logits_to_return + + class TopKTopPSampler(nn.Module): """Module that performs optional top-k and top-p filtering followed by weighted random sampling of logits. @@ -162,7 +224,7 @@ def __init__( else: self.forward = self.forward_cpu elif current_platform.is_xpu(): - if envs.VLLM_XPU_USE_SAMPLER_KERNEL: + if xpu_sampler_supported(): self.forward = self.forward_xpu else: self.forward = self.forward_native @@ -356,37 +418,7 @@ def forward_xpu( "PyTorch-native implementation." ) return self.forward_native(logits, generators, k, p) - random_sampled = torch.empty( - logits.shape[0], dtype=torch.int64, device=logits.device - ) - logits_to_return = None - if self.logprobs_mode in PROCESSED_LOGPROBS_MODES: - logits_to_return = torch.empty_like(logits) - - assert len(generators) != logits.shape[0], ( - "xpu kernel topk_topp_sampler does not support batch-wise generators." - ) - generator = torch.xpu.default_generators[logits.device.index] - - state = generator.get_state() - seed, offset = state.view(torch.int64) - seeds = torch.tensor( - [seed, offset], dtype=torch.int64, device=torch.device("cpu") - ) - # The XPU kernel expects k as int64 (Long), but the input batch - # stores top_k as int32. Cast here to avoid dtype mismatch. - if k is not None: - k = k.to(torch.int64) - torch.ops.vllm.xpu_topk_topp_sampler( - random_sampled, logits_to_return, logits, k, p, self.logprobs_mode, seeds - ) - # The custom XPU sampler kernel consumes RNG values internally, so advance - # the default generator's offset to keep future draws deterministic. - # pytorch: offset must be multiple of 4 - offset = (offset + logits.numel() + 3) // 4 * 4 - state.view(torch.int64)[1] = offset - generator.set_state(state) - return random_sampled, logits_to_return + return xpu_sample(logits, k, p, self.logprobs_mode) # Note: this is a workaround for diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index c9624b2b25ce..a4e6e2dd5bb0 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -58,7 +58,7 @@ def _sample_random( pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, - use_flashinfer: bool, + use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: enabled = self.watermarking.np[idx_mapping_np] & ( self.sampling_states.temperature.np[idx_mapping_np] != 0 @@ -71,7 +71,7 @@ def _sample_random( pos, top_k, top_p, - use_flashinfer, + use_fused_sampler, ) processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 92e9ecee4d66..802adb6f642f 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -14,6 +14,8 @@ apply_top_k_top_p, flashinfer_sample, flashinfer_sampler_supported, + xpu_sample, + xpu_sampler_supported, ) from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected from vllm.v1.worker.gpu.metrics.logits import get_num_nans @@ -86,6 +88,10 @@ def __init__( self.use_flashinfer = ( not return_sampling_mask and flashinfer_sampler_supported() ) + # The XPU kernel draws fp32 exponential noise, so it can't honor fp64. + self.use_xpu_sampler = ( + not return_sampling_mask and not use_fp64_gumbel and xpu_sampler_supported() + ) def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: needs_processing = self.sampling_states.add_request(req_idx, sampling_params) @@ -292,15 +298,18 @@ def sample( top_k, top_p = self.sampling_states.get_top_k_top_p( expanded_idx_mapping, idx_mapping_np ) - use_flashinfer = self.use_flashinfer and not ( - # Don't use FI sampler if no requests use top_k/top_p, if there are - # any greedy requests or per-request seeds, or if post-processed - # logprobs need to be returned for any requests. + # Don't use a fused sampler if no requests use top_k/top_p, if there are + # any greedy requests or per-request seeds, or if post-processed + # logprobs need to be returned for any requests. + fused_sampler_eligible = not ( (top_k is None and top_p is None) or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES) or self.sampling_states.any_greedy(idx_mapping_np) or self.sampling_states.any_explicit_seed(idx_mapping_np) ) + use_fused_sampler = ( + self.use_flashinfer or self.use_xpu_sampler + ) and fused_sampler_eligible return self._sample_random( processed_logits, @@ -309,7 +318,7 @@ def sample( pos, top_k, top_p, - use_flashinfer, + use_fused_sampler, ) def _sample_random( @@ -320,10 +329,15 @@ def _sample_random( pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, - use_flashinfer: bool, + use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if use_flashinfer: - sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64) + if use_fused_sampler: + if self.use_flashinfer: + sampled = flashinfer_sample(processed_logits, top_k, top_p).to( + torch.int64 + ) + else: # Use XPU sampler + sampled, _ = xpu_sample(processed_logits, top_k, top_p) else: processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) sampled = gumbel_sample( From 8ec00ec2bc2d294b742f7814041d3a3043e0a90e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 00:27:29 -0500 Subject: [PATCH 564/774] [Tests] Release HF embedding weights after prompt extraction (#57868) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/models/language/generation/test_common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index f70d5a0c6892..450e8badcebe 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -158,7 +158,9 @@ def test_models( hf_model.model.device ) if prompt_embeds is not None: - embed = hf_model.model.get_input_embeddings()(token_ids) + # Retained prompt values must not keep reference weights alive. + with torch.no_grad(): + embed = hf_model.model.get_input_embeddings()(token_ids) # MiniCPM models apply scale_emb to embeddings internally. # vLLM expects pre-scaled embeddings when using inputs_embeds. From 59f15272802077724c4ed483ea8ad393725554e4 Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:44:40 +0800 Subject: [PATCH 565/774] [Bugfix] Fix Whisper engine crash on audio clips longer than 30s (#57769) Signed-off-by: Tianyao Wu --- .../multimodal/processing/test_whisper.py | 37 +++++++++++++++++++ vllm/model_executor/models/whisper.py | 1 + 2 files changed, 38 insertions(+) create mode 100644 tests/models/multimodal/processing/test_whisper.py diff --git a/tests/models/multimodal/processing/test_whisper.py b/tests/models/multimodal/processing/test_whisper.py new file mode 100644 index 000000000000..1627b48914c3 --- /dev/null +++ b/tests/models/multimodal/processing/test_whisper.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for Whisper's multimodal preprocessing.""" + +import numpy as np +import pytest + +from vllm.multimodal import MULTIMODAL_REGISTRY + +from ...utils import build_model_context + + +@pytest.mark.parametrize("model_id", ["openai/whisper-large-v3-turbo"]) +@pytest.mark.parametrize("audio_duration_s", [5, 30, 35]) +def test_audio_features_fit_encoder_window( + model_id: str, + audio_duration_s: int, +) -> None: + """The encoder only has positions for one 30 s window, so a longer clip + has to be cut to that window like the HF feature extractor does by default. + """ + ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": 1}) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + feature_extractor = processor.info.get_feature_extractor() + + sampling_rate = feature_extractor.sampling_rate + rng = np.random.RandomState(0) + audio = rng.rand(sampling_rate * audio_duration_s).astype(np.float32) + + processed_inputs = processor( + "<|startoftranscript|>", + mm_items=processor.info.parse_mm_data({"audio": [(audio, sampling_rate)]}), + hf_processor_mm_kwargs={}, + ) + + input_features = processed_inputs["mm_kwargs"].get_data()["input_features"] + assert input_features.shape[-1] == feature_extractor.nb_max_frames diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index 59268619b0a4..f67e3f0cb9fd 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -767,6 +767,7 @@ def _get_hf_mm_inputs( hf_kwargs=dict( hf_inputs.hf_kwargs, sampling_rate=feature_extractor.sampling_rate, + truncation=True, ) ) From 9598487bc1829504c8af81d6007ab6b9deeaa15a Mon Sep 17 00:00:00 2001 From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:36:59 +0100 Subject: [PATCH 566/774] [Bugfix][Model] Mistral3: align placeholder grid with public processor (#53758) Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../multimodal/processing/test_mistral3.py | 76 ++++++++++++------- vllm/model_executor/models/mistral3.py | 15 ++-- 2 files changed, 54 insertions(+), 37 deletions(-) diff --git a/tests/models/multimodal/processing/test_mistral3.py b/tests/models/multimodal/processing/test_mistral3.py index fcb783f22f6e..4d900e627fdf 100644 --- a/tests/models/multimodal/processing/test_mistral3.py +++ b/tests/models/multimodal/processing/test_mistral3.py @@ -5,7 +5,7 @@ import pytest import torch from PIL import Image -from transformers import BatchFeature +from transformers import AutoProcessor, BatchFeature from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo @@ -22,25 +22,41 @@ _LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025" -def _processed_pixel_values( +def _process_images_with_hf( hf_processor, images: list[Image.Image], mm_processor_kwargs: dict[str, object], -) -> list[torch.Tensor]: - """Resize via the HF image processor and un-pad to per-image H×W.""" - image_processor = hf_processor.image_processor - hf_out = image_processor(images=images, return_tensors="pt", **mm_processor_kwargs) +) -> tuple[list[torch.Tensor], list[int]]: + """Process images and placeholders through the public HF processor.""" + hf_out = hf_processor( + text=hf_processor.image_token * len(images), + images=images, + return_tensors="pt", + **mm_processor_kwargs, + ) pixel_values = hf_out["pixel_values"] image_sizes = hf_out["image_sizes"] - return [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)] - - -def _placeholder_count_from_prompt_updates( + unpadded = [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)] + + special_token_ids = { + hf_processor.image_token_id, + hf_processor.image_break_token_id, + hf_processor.image_end_token_id, + } + placeholder_tokens = [ + token_id + for token_id in hf_out["input_ids"][0].tolist() + if token_id in special_token_ids + ] + return unpadded, placeholder_tokens + + +def _placeholder_tokens_from_prompt_updates( processor, images: list[Image.Image], pixel_values: list[torch.Tensor], mm_processor_kwargs: dict[str, object], -) -> int: +) -> list[int]: hf_inputs = BatchFeature({"pixel_values": pixel_values}) fields_config = processor._get_mm_fields_config(hf_inputs, mm_processor_kwargs) out_mm_kwargs = MultiModalKwargsItems.from_hf_inputs(hf_inputs, fields_config) @@ -49,13 +65,11 @@ def _placeholder_count_from_prompt_updates( updates = processor._get_prompt_updates( mm_items, mm_processor_kwargs, out_mm_kwargs ) - image_token_id = processor.info.get_hf_config().image_token_index - - total = 0 + placeholder_tokens: list[int] = [] for item_idx in range(len(images)): details = updates[0].resolve(item_idx).content - total += details.full.count(image_token_id) - return total + placeholder_tokens.extend(details.full) + return placeholder_tokens def _expected_placeholder_tokens_per_image( @@ -72,12 +86,12 @@ def _expected_placeholder_tokens_per_image( patch_h = patch_w = int(patch_size) spatial_merge_size = getattr(hf_processor, "spatial_merge_size", 1) - assert image_h % patch_h == 0 - assert image_w % patch_w == 0 + merged_patch_h = patch_h * spatial_merge_size + merged_patch_w = patch_w * spatial_merge_size + assert image_h % merged_patch_h == 0 + assert image_w % merged_patch_w == 0 - return (image_h // (patch_h * spatial_merge_size)) * ( - image_w // (patch_w * spatial_merge_size) - ) + return (image_h // merged_patch_h) * (image_w // merged_patch_w) @pytest.mark.parametrize("model_id", [_MODEL_ID]) @@ -87,8 +101,8 @@ def _expected_placeholder_tokens_per_image( ({}, (448, 448), 256), ({"size": {"longest_edge": 1008}}, (1540, 1540), 1296), ({"size": {"longest_edge": 1288}}, (1536, 1187), 1656), - ({"size": {"longest_edge": 1008}}, (29, 29), 1), - ({"size": {"longest_edge": 1000}}, (1540, 1700), 1152), + ({"size": {"longest_edge": 1008}}, (29, 29), 4), + ({"size": {"longest_edge": 1000}}, (1540, 1700), 1188), ], ) @pytest.mark.parametrize("num_imgs", [1, 2]) @@ -109,21 +123,29 @@ def test_processor_size_override( ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs - hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs) + hf_processor = AutoProcessor.from_pretrained( + model_id, + fix_mistral_regex=True, + ) dummy_image = Image.new("RGB", image_size, color=(127, 127, 127)) images = [dummy_image] * num_imgs merged_mm_kwargs = processor.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs) - pixel_values = _processed_pixel_values(hf_processor, images, merged_mm_kwargs) + pixel_values, hf_placeholder_tokens = _process_images_with_hf( + hf_processor, images, merged_mm_kwargs + ) - image_token_count = _placeholder_count_from_prompt_updates( + prompt_update_tokens = _placeholder_tokens_from_prompt_updates( processor, images, pixel_values, hf_processor_mm_kwargs ) expected_from_pixel_values = _expected_placeholder_tokens_per_image( hf_processor, pixel_values[0] ) assert expected_from_pixel_values == expected_toks_per_img - assert image_token_count == expected_from_pixel_values * num_imgs + assert hf_placeholder_tokens.count(hf_processor.image_token_id) == ( + expected_from_pixel_values * num_imgs + ) + assert prompt_update_tokens == hf_placeholder_tokens def test_lightonocr_keeps_vision_config_image_size(): diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 41352dc42766..a70d8e1de531 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -192,18 +192,13 @@ def get_patch_grid_size( image_width = math.floor(image_width / ratio) image_height = math.floor(image_height / ratio) - patch_size = self.vision_config.patch_size assert isinstance(self.hf_config, Mistral3Config) - spatial_merge_size = self.hf_config.spatial_merge_size - - # The HF processor rounds each dimension up to the vision patch size - # before the projector drops incomplete spatial-merge groups. This is - # not equivalent to rounding directly to the merged patch size. - num_width_patches = (image_width - 1) // patch_size + 1 - num_height_patches = (image_height - 1) // patch_size + 1 + merged_patch_size = ( + self.vision_config.patch_size * self.hf_config.spatial_merge_size + ) - ncols = num_width_patches // spatial_merge_size - nrows = num_height_patches // spatial_merge_size + ncols = (image_width - 1) // merged_patch_size + 1 + nrows = (image_height - 1) // merged_patch_size + 1 return ncols, nrows From 3c84ad13ed3f9fcd816bf1d15acb83c28cfa02f9 Mon Sep 17 00:00:00 2001 From: CaiJohn Date: Mon, 21 Sep 2026 14:59:20 +0800 Subject: [PATCH 567/774] [Bugfix] Fix xgrammar feature gate bypass when JSON Schema type is a list (#48416) Signed-off-by: john <6712432+CaiJohn@users.noreply.github.com> Co-authored-by: john <6712432+CaiJohn@users.noreply.github.com> Co-authored-by: Codex --- tests/v1/structured_output/test_utils.py | 46 +++++++++++++++++++ tests/v1/structured_output/test_validation.py | 15 ++++++ vllm/v1/structured_output/backend_xgrammar.py | 28 ++++++++--- 3 files changed, 82 insertions(+), 7 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index 3e4a12f7e5c1..ff2b58dfe194 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -224,6 +224,52 @@ def test_supported_json_features(schema_type, request): ) +class TestHasXGrammarUnsupportedJsonFeatures: + class TestPR48416Regressions: + @pytest.mark.parametrize( + "schema", + [ + {"type": ["number"], "multipleOf": 3}, + { + "type": ["string", "null"], + "pattern": "^a+$", + "maxLength": 2, + }, + { + "type": ["array", "null"], + "items": {"type": "integer"}, + "uniqueItems": True, + }, + { + "type": ["object", "null"], + "properties": {"Bad": {"type": "integer"}}, + "propertyNames": {"pattern": "^[a-z]+$"}, + }, + ], + ) + def test_list_type_does_not_bypass_unsupported_feature_check(self, schema): + assert has_xgrammar_unsupported_json_features(schema) + + @pytest.mark.parametrize( + "schema", + [ + {"type": ["string", "null"]}, + {"type": ["integer", "string"]}, + {"type": ["array", "null"], "items": {"type": "string"}}, + { + "type": ["object", "null"], + "propertyNames": {"pattern": "^[a-z_]+$"}, + }, + { + "type": ["object", "null"], + "patternProperties": {"^S": {"type": "string"}}, + }, + ], + ) + def test_supported_list_type_json_features(self, schema): + assert not has_xgrammar_unsupported_json_features(schema) + + def test_property_names_constrains_keys(property_names_schema): assert grammar_accepts(property_names_schema, '{"score": 5}') assert grammar_accepts(property_names_schema, '{"score": "seven"}') diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index a23189438475..a82b3e51177b 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -145,6 +145,21 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): }, "guidance", ), + ( + { + "type": "object", + "properties": {"n": {"type": ["number", "null"], "multipleOf": 3}}, + }, + "guidance", + ), + ( + { + "type": ["string", "null"], + "pattern": "^a+$", + "maxLength": 2, + }, + "guidance", + ), # patternProperties + properties is also unsupported by guidance. ( { diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 1ba7e2734d9e..3c48ff64a0bd 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -242,6 +242,18 @@ def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool: ) +# FIXME(arpera): The approach used here needs to be redesigned because of +# existing bugs: https://github.com/vllm-project/vllm/issues/57550 +def _schema_types(schema: dict[str, Any]) -> set[str]: + """Normalize a scalar or list-valued JSON Schema type.""" + schema_type = schema.get("type") + if isinstance(schema_type, str): + return {schema_type} + if isinstance(schema_type, list): + return {item for item in schema_type if isinstance(item, str)} + return set() + + def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: """Check if JSON schema contains features unsupported by xgrammar.""" @@ -249,12 +261,14 @@ def check_object(obj: dict[str, Any]) -> bool: if not isinstance(obj, dict): return False + schema_types = _schema_types(obj) + # Check for numeric ranges - if obj.get("type") in ("integer", "number") and ("multipleOf" in obj): + if (schema_types & {"integer", "number"}) and ("multipleOf" in obj): return True # Check for array unsupported keywords - if obj.get("type") == "array" and any( + if "array" in schema_types and any( key in obj for key in ("uniqueItems", "contains", "minContains", "maxContains") ): @@ -262,7 +276,7 @@ def check_object(obj: dict[str, Any]) -> bool: # Unsupported keywords for strings if ( - obj.get("type") == "string" + "string" in schema_types and "format" in obj and obj["format"] not in STRING_SUPPORTED_FORMATS ): @@ -275,13 +289,13 @@ def check_object(obj: dict[str, Any]) -> bool: # the compiled EBNF: pattern/format grammars come out byte-identical # with and without the length keywords, while maxLength alone lowers # to {0, N} correctly. - if obj.get("type") == "string" and _has_pattern_and_length_bounds(obj): + if "string" in schema_types and _has_pattern_and_length_bounds(obj): return True # propertyNames validates names, so it is a string schema even when it # omits "type", which is the form that escapes the check above. if ( - obj.get("type") == "object" + "object" in schema_types and isinstance(obj.get("propertyNames"), dict) and _has_pattern_and_length_bounds(obj["propertyNames"]) ): @@ -291,7 +305,7 @@ def check_object(obj: dict[str, Any]) -> bool: # additionalProperties/unevaluatedProperties under xgrammar. # https://github.com/mlc-ai/xgrammar/issues/826 if ( - obj.get("type") == "object" + "object" in schema_types and "propertyNames" in obj and ( "properties" in obj @@ -305,7 +319,7 @@ def check_object(obj: dict[str, Any]) -> bool: # FIXME: multiple patternProperties, or patternProperties alongside # properties, conflict under xgrammar. if ( - obj.get("type") == "object" + "object" in schema_types and isinstance(obj.get("patternProperties"), dict) and ("properties" in obj or len(obj["patternProperties"]) > 1) ): From 97dc6b19d2fe92b794b49e684677a2e1a8b2c540 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 21 Sep 2026 00:08:43 -0700 Subject: [PATCH 568/774] [Perf][Attention] Avoid redundant sparse attention metadata operations (#57885) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- vllm/v1/attention/backends/mla/indexer.py | 5 +++-- vllm/v1/attention/backends/mla/sparse_swa.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 8af81c8d2900..00d14ceeb7eb 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1261,7 +1261,6 @@ def build( assert num_decode_tokens + num_prefill_tokens == num_tokens compressed_slot_mapping = slot_mapping - compressed_seq_lens = seq_lens indexer_block_table = block_table if self.compress_ratio > 1: kernel_block_size = self.kernel_block_size @@ -1297,10 +1296,12 @@ def build( self.compressed_slot_mapping_buffer[:padded_num_tokens], dim=0, ) - compressed_seq_lens = seq_lens // self.compress_ratio prefill_metadata = None if num_prefills > 0: + compressed_seq_lens = ( + seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens + ) # This CPU value is an upper bound for async-spec extend rows. It # is safe for chunking/allocation because CUDA metadata below is # built from exact device seq_lens and gather ignores the tail. diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 6fe20e0e6934..7b07dae71c9e 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -613,7 +613,7 @@ def build( ) is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] - is_valid_token.copy_(slot_mapping >= 0) + torch.ge(slot_mapping, 0, out=is_valid_token) if replay_start is None: # Graph captures build without it; their dummy batches replay nothing. replay_start = self.no_replay_start From 67513c8b67e2439934ce7c0ba5a18c2e859f5c53 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 21 Sep 2026 01:16:59 -0700 Subject: [PATCH 569/774] [Bugfix][DSV4.1] Restrict mHC overlap to full CUDA graphs (#57874) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- tests/kernels/test_mhc_kernels.py | 23 +++++++++++++++++-- vllm/models/deepseek_v41/nvidia/model.py | 3 ++- .../deepseek_v41/nvidia/ops/mega_mhc.py | 16 +++++++++++++ vllm/models/deepseek_v41/nvidia/ops/mhc.py | 14 +++++------ 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 63648ebf2d91..4735fb96d18d 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -539,7 +539,7 @@ def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_ @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) @pytest.mark.parametrize( - "mhc_mode", ["disabled", "overlap", "piecewise", "large_batch"] + "mhc_mode", ["disabled", "eager", "overlap", "piecewise", "large_batch"] ) def test_deepseek_v41_decoder_mixes_match_torch( entry, mhc_mode, monkeypatch, default_vllm_config @@ -568,7 +568,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( ): pytest.skip("SM100 DeepGEMM required for overlap") decoder.mhc_stream = torch.cuda.Stream() - if mhc_mode in ("piecewise", "large_batch"): + if mhc_mode in ("eager", "piecewise", "large_batch"): def unexpected_overlap(*args, **kwargs): pytest.fail("unsupported execution must retain native mHC") @@ -624,8 +624,27 @@ def forward(self, residual, hashes, mask): from vllm.forward_context import set_forward_context mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE + overlap_calls = [] + if mhc_mode == "overlap": + from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap + + def checked_overlap(*args, **kwargs): + overlap_calls.append(torch.cuda.is_current_stream_capturing()) + return mhc_pre_delayed_overlap(*args, **kwargs) + + for module in ("model", "ops.mhc"): + monkeypatch.setattr( + f"vllm.models.deepseek_v41.nvidia.{module}.mhc_pre_delayed_overlap", + checked_overlap, + ) with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode): actual = decoder(x, positions, None, **kwargs) + if mhc_mode == "overlap": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = decoder(x, positions, None, **kwargs) + graph.replay() + assert overlap_calls and all(overlap_calls) def reference(*args, norm_weight, norm_eps, **kwargs): post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index d7c9ddc1284f..b91ad1d7ba4a 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -385,8 +385,9 @@ def forward( if mhc_stream is not None and ( in_piecewise_cudagraph() or not 0 < positions.shape[0] <= MHC_OVERLAP_MAX_TOKENS + or not torch.cuda.is_current_stream_capturing() ): - # A side stream cannot remain unjoined across breakable graph segments. + # Use overlap only in FULL graphs; eager warmup initializes Mega mHC. mhc_stream = None mhc_pre = ( partial(mhc_pre_delayed_overlap, stream=mhc_stream) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py index ef56b5a82393..bdb5b95bbeb9 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -28,6 +28,22 @@ def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) +def can_use_mega_mhc( + x: torch.Tensor, + residual: torch.Tensor, + pre_mix: torch.Tensor | None, + norm_weight: torch.Tensor | None, + capture_aux: bool, +) -> bool: + return ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ) + + def mhc_shifted_post_pre_deep_gemm( x: torch.Tensor, residual: torch.Tensor, diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index 084a102436ed..88288ead4d9d 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -16,7 +16,10 @@ from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported -from .mega_mhc import is_mega_mhc_supported, mhc_shifted_post_pre_deep_gemm +from .mega_mhc import ( + can_use_mega_mhc, + mhc_shifted_post_pre_deep_gemm, +) if TYPE_CHECKING: from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator @@ -223,13 +226,8 @@ def mhc_shifted_post_pre( ) return residual, *pre_outputs, aux - if ( - pre_mix is not None - and norm_weight is not None - and not capture_aux - and x.shape[0] <= 1 << 20 - and is_mega_mhc_supported(x.shape[1], residual.shape[1]) - ): + if can_use_mega_mhc(x, residual, pre_mix, norm_weight, capture_aux): + assert pre_mix is not None and norm_weight is not None outputs = mhc_shifted_post_pre_deep_gemm( x, residual, From 0ff0477ff955154f37db48791260cdf27fa35e60 Mon Sep 17 00:00:00 2001 From: akii96 Date: Mon, 21 Sep 2026 11:20:43 +0300 Subject: [PATCH 570/774] [Perf][ROCm] Add a ROCm path for Hy4 and compile the backbone (#57526) Signed-off-by: Aakif Nawaz --- tests/models/test_hyv4_rocm.py | 48 ++ vllm/models/hy_v4/__init__.py | 21 +- vllm/models/hy_v4/amd/__init__.py | 2 + vllm/models/hy_v4/amd/model.py | 718 ++++++++++++++++++++++++++++++ 4 files changed, 782 insertions(+), 7 deletions(-) create mode 100644 vllm/models/hy_v4/amd/__init__.py create mode 100644 vllm/models/hy_v4/amd/model.py diff --git a/tests/models/test_hyv4_rocm.py b/tests/models/test_hyv4_rocm.py index b6dc8d70423f..ea740c5cb099 100644 --- a/tests/models/test_hyv4_rocm.py +++ b/tests/models/test_hyv4_rocm.py @@ -51,6 +51,18 @@ def _num_captured_graphs(_worker): return compilation_counter.num_cudagraph_captured +def _piecewise_compile_state(worker): + from vllm.compilation.counter import compilation_counter + + # `cudagraph_mode` is returned by name: the enum does not survive the + # collective_rpc round trip, it arrives as its bare value. + return ( + compilation_counter.num_models_seen, + compilation_counter.num_cudagraph_captured, + worker.vllm_config.compilation_config.cudagraph_mode.name, + ) + + @pytest.mark.parametrize( "tp,enforce_eager,mtp,dtype,kv_cache_dtype", [ @@ -135,3 +147,39 @@ def test_hyv4_dummy_generation( and metric.value > 0 for metric in model.llm.get_metrics() ), "MTP drafter did not run" + + +def test_hyv4_backbone_reaches_the_compiler(vllm_runner, enable_pickle) -> None: + """Guard against the backbone losing `@support_torch_compile`. + + An undecorated backbone is not a hard failure: vLLM downgrades + `cudagraph_mode` to NONE and serves the whole model eager, which on TP8 + gfx942 cost ~4x decode latency. + """ + with vllm_runner( + "tencent/Hy4-preview", + trust_remote_code=True, + skip_tokenizer_init=True, + load_format="dummy", + dtype="bfloat16", + hf_overrides=_small_hyv4_config, + max_model_len=256, + max_num_batched_tokens=128, + max_num_seqs=8, + enable_chunked_prefill=True, + enable_prefix_caching=False, + kv_cache_memory_bytes=128 * 1024 * 1024, + enforce_eager=False, + seed=17, + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + cudagraph_mode=CUDAGraphMode.PIECEWISE, + cudagraph_capture_sizes=[1, 2, 4], + ), + ) as model: + for models_seen, captured, mode in model.collective_rpc( + _piecewise_compile_state + ): + assert models_seen > 0, "backbone was never handed to the compiler" + assert mode == CUDAGraphMode.PIECEWISE.name + assert captured > 0 diff --git a/vllm/models/hy_v4/__init__.py b/vllm/models/hy_v4/__init__.py index f09f25bfb21f..f6ffafb1c516 100644 --- a/vllm/models/hy_v4/__init__.py +++ b/vllm/models/hy_v4/__init__.py @@ -14,20 +14,27 @@ only public entry point and dispatches on the current platform, so registry entries never reach into a platform subpackage. -NVIDIA and ROCm are supported. The port also drops the reference -implementation's HPC/TPCP fusion paths, which depend on infrastructure that -does not exist in this tree. +NVIDIA and ROCm are supported, the latter through `amd/`. The port also drops +the reference implementation's HPC/TPCP fusion paths, which depend on +infrastructure that does not exist in this tree. """ +from typing import TYPE_CHECKING + from vllm.platforms import current_platform if current_platform.is_xpu(): raise NotImplementedError("hy_v4 does not yet support XPU.") -else: - # The implementation is shared by CUDA and ROCm; platform-specific - # attention behavior is selected by the normal attention backend router. + +# The NVIDIA branch is the static default that type-checkers see; the ROCm +# branch overrides it at runtime (kept type-compatible via type: ignore). +if TYPE_CHECKING or not current_platform.is_rocm(): from .nvidia.model import HYV4ForCausalLM - from .nvidia.mtp import HYV4MTP +else: + from .amd.model import HYV4ForCausalLM # type: ignore[assignment] + +# The draft head has no ROCm-specific implementation yet. +from .nvidia.mtp import HYV4MTP # noqa: E402 __all__ = [ "HYV4ForCausalLM", diff --git a/vllm/models/hy_v4/amd/__init__.py b/vllm/models/hy_v4/amd/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/hy_v4/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/hy_v4/amd/model.py b/vllm/models/hy_v4/amd/model.py new file mode 100644 index 000000000000..753ac92a3878 --- /dev/null +++ b/vllm/models/hy_v4/amd/model.py @@ -0,0 +1,718 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Adapted from +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only HY V4 model compatible with HuggingFace weights (ROCm).""" + +import typing +from collections.abc import Callable, Iterable, MutableSequence, Sequence +from itertools import islice + +import regex as re +import torch +from torch import nn +from transformers import PretrainedConfig + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + get_pp_missing_layer_names, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from ..nvidia.attention import ( + HYV4MLAAttention, + compute_skip_topk_layers, + is_skip_topk_indexer_weight, +) +from ..nvidia.hc import HYV4HCHeadLayer, HYV4HCLayer +from ..nvidia.moe import HYV4FeedForward, HYV4MoEFused + +logger = init_logger(__name__) + + +def _normalize_hyv4_config(config: PretrainedConfig) -> PretrainedConfig: + """Populate the aliases consumed by the shared MoE implementation.""" + config.router_scaling_factor = config.routed_scaling_factor + config.num_experts = config.n_routed_experts + config.expert_hidden_dim = config.moe_intermediate_size + config.num_shared_experts = config.n_shared_experts + config.route_norm = config.norm_topk_prob + return config + + +class HYV4DecoderLayer(nn.Module): + """One HY V4 decoder layer: MLA attention plus a dense or MoE MLP. + + When``config.enable_ihc`` is set the layer runs on ``hc_mult`` residual + channels and each sub-block is wrapped by an `HYV4HCLayer` boundary; + otherwise it uses the standard single-stream residual. + """ + + def __init__( + self, + config: PretrainedConfig, + vllm_config: VllmConfig, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.enable_ihc = getattr(config, "enable_ihc", False) + layer_idx = int(prefix.split(".")[-1]) + self.layer_idx = layer_idx + + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = HYV4MLAAttention( + vllm_config=vllm_config, + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + max_position_embeddings=max_position_embeddings, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + layer_idx=layer_idx, + topk_indices_buffer=topk_indices_buffer, + ) + self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) + if config.mlp_layer_types[layer_idx] == "dense": + self.mlp = HYV4FeedForward( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.block_type = "feedforward" + else: + self.mlp = HYV4MoEFused( + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + self.block_type = "moe" + self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.hc_attn_layer = HYV4HCLayer( + config, layer_idx, prefix=f"{prefix}.hc_attn_layer" + ) + self.hc_mlp_layer = HYV4HCLayer( + config, layer_idx, prefix=f"{prefix}.hc_mlp_layer" + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if self.enable_ihc: + return self._forward_ihc(positions, hidden_states) + return self._forward_normal(positions, hidden_states, residual) + + def _forward_normal( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Standard single-residual-stream forward (iHC disabled).""" + if residual is not None: + hidden_states = hidden_states + residual + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + hidden_states = hidden_states + residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual + + def _forward_ihc( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, None]: + """IHC forward: each sub-block reduces and re-scatters the channels.""" + hidden_states = self.hc_attn_layer.prepare_input(hidden_states) + hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states) + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states = self.hc_attn_layer.post(hidden_states, residual, post_gates) + + hidden_states = self.hc_mlp_layer.prepare_input(hidden_states) + hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post_gates) + + # Under iHC the residual is carried inside hidden_states. + return hidden_states, None + + +@support_torch_compile +class HYV4Model(nn.Module): + """HY V4 backbone.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = _normalize_hyv4_config(vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + eplb_config = parallel_config.eplb_config + self.num_redundant_experts = eplb_config.num_redundant_experts + self.device = current_platform.device_type + self.vocab_size = config.vocab_size + self.is_sparse = hasattr(config, "index_topk") + if self.is_sparse: + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + else: + self.topk_indices_buffer = None + self.config = config + self.quant_config = quant_config + self.enable_ihc = getattr(config, "enable_ihc", False) + + if get_pp_group().is_first_rank or ( + self.config.tie_word_embeddings and get_pp_group().is_last_rank + ): + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + else: + self.embed_tokens = PPMissingLayer() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: HYV4DecoderLayer( + config=config, + vllm_config=vllm_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + if self.enable_ihc: + # iHC head layer: merge the residual channels into one stream. + if get_pp_group().is_last_rank: + self.hc_head = HYV4HCHeadLayer( + config, + hidden_size=config.hidden_size, + hc_mult=config.hc_mult, + hc_eps=config.hc_eps, + prefix=f"{prefix}.hc_head", + ) + else: + self.hc_head = PPMissingLayer() + self.make_empty_intermediate_tensors = ( + make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hc_mult * config.hidden_size + ) + ) + else: + self.make_empty_intermediate_tensors = ( + make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + ) + + # MoE hyperparameters (consumed by EPLB). + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = [] + self.num_expert_groups = 1 + self.moe_layers: list[nn.Module] = [] + example_layer: HYV4MoEFused | None = None + for layer in self.layers: + if isinstance(layer, PPMissingLayer): + continue + + assert isinstance(layer, HYV4DecoderLayer) + if layer.block_type == "moe": + assert isinstance(layer.mlp, HYV4MoEFused) + example_layer = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_layer is None: + self.num_moe_layers = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_redundant_experts = 0 + return + + self.num_moe_layers = len(self.moe_layers) + self.num_logical_experts = example_layer.n_logical_experts + self.num_physical_experts = example_layer.n_physical_experts + self.num_local_physical_experts = example_layer.n_local_physical_experts + self.num_routed_experts = example_layer.n_routed_experts + self.num_redundant_experts = example_layer.n_redundant_experts + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, HYV4MoEFused): + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # (param_name, weight_name, expert_id, shard_id) for weights, fp8 + # weight scales and fp8 activation scales. + if not hasattr(self, "_cached_expert_params_mapping"): + self._cached_expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + return self._cached_expert_params_mapping + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + # In iHC mode the flattened [num_tokens, hc*h] tensor from the + # previous PP stage is reshaped back to 3D by the first layer's + # prepare_input, and residual is unused. + residual = None if self.enable_ihc else intermediate_tensors["residual"] + + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual = layer(positions, hidden_states, residual) + + if not get_pp_group().is_last_rank: + if self.enable_ihc: + # hidden_states is [num_tokens, hc, h]; flatten the channel dim + # for PP transfer (matches the 2D receive buffer). + return IntermediateTensors({"hidden_states": hidden_states.flatten(1)}) + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + if self.enable_ihc: + hidden_states = self.hc_head(hidden_states) + else: + hidden_states = hidden_states + residual + + return self.norm(hidden_states) + + def load_fused_expert_weights( + self, + name: str, + params_dict: dict, + loaded_weight: torch.Tensor, + shard_id: str, + num_experts: int, + ) -> bool: + param = params_dict[name] + weight_loader = typing.cast(Callable[..., bool], param.weight_loader) + loaded_local_expert = False + for expert_id in range(num_experts): + curr_expert_weight = loaded_weight[expert_id] + success = weight_loader( + param, + curr_expert_weight, + name, + shard_id, + expert_id, + return_success=True, + ) + if success: + loaded_local_expert = True + + return loaded_local_expert + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # FP8 indexer wk dequant buffer (weight and scale arrive separately). + pending_wk_fp8: dict[str, dict[str, torch.Tensor]] = {} + pp_missing_layer_names = get_pp_missing_layer_names(self) + skip_topk_layers = compute_skip_topk_layers(self.config) + + # Must not be cached on `self`: `process_weights_after_loading` swaps in + # kernel-specific expert weights via `replace_parameter`, and a cache + # outliving this call would pin the pre-shuffle storage (OOM on large + # MoE) and make a later reload target orphaned tensors. + params_dict = dict(self.named_parameters()) + # Split per-expert mapping (V3 style): experts.0.gate_proj.weight + split_expert_params_mapping = self.get_expert_mapping() + loaded_params: set[str] = set() + + # Sink weights are sharded like the q/k/v linears. + sink_tp_size = get_tensor_model_parallel_world_size() + sink_tp_rank = get_tensor_model_parallel_rank() + n_local_head = self.config.num_attention_heads // sink_tp_size + head_rank_start = n_local_head * sink_tp_rank + head_rank_end = n_local_head * (sink_tp_rank + 1) + base_layer = ( + "base_layer." if any(".base_layer." in name for name in params_dict) else "" + ) + # Fused expert mapping: experts.gate_up_proj (all experts in one tensor). + # The packed weights are owned by the RoutedExperts submodule, so the + # targets are experts.routed_experts.[base_layer.]w{13,2}_weight. + fused_expert_prefix = f".experts.routed_experts.{base_layer}" + fused_expert_params_mapping = [ + (f"{fused_expert_prefix}w13_weight", ".experts.gate_up_proj", 0, "w1"), + (f"{fused_expert_prefix}w2_weight", ".experts.down_proj", 0, "w2"), + ] + num_experts = getattr(self.config, "num_experts", 0) + + def _should_skip_missing_param(param_name: str) -> bool: + # Sparse checkpoints may contain indexer weights for layers that + # fall back to dense MLA on unsupported GPUs. + return ".indexer." in param_name + + def _is_split_expert_weight(weight_name: str) -> bool: + """Whether this weight is in split (per-expert) format.""" + # Split format: mlp.experts..gate_proj.weight + return bool(re.search(r"\.experts\.\d+\.", weight_name)) + + def _is_fused_expert_weight(weight_name: str) -> bool: + """Whether this weight is in fused (all-experts-packed) format.""" + return ".experts.gate_up_proj" in weight_name or ( + ".experts.down_proj" in weight_name + and not _is_split_expert_weight(weight_name) + ) + + for name, loaded_weight in weights: + if is_skip_topk_indexer_weight(name, skip_topk_layers): + continue + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, + ): + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + # Names and stacked shard ids are normalized upstream by + # AutoWeightsLoader, through quant_config.get_cache_scale_mapper() + # and hf_to_vllm_mapper. + stacked_shard_id = getattr(loaded_weight, "shard_id", None) + if ".indexer.wk." in name: + name = name.replace(".indexer.wk.", ".indexer.wk_weights_proj.") + stacked_shard_id = 0 + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + # Determine per-weight whether this is fused or split format. + is_fused_expert = _is_fused_expert_weight(name) + expert_params_mapping = ( + fused_expert_params_mapping + if is_fused_expert + else split_expert_params_mapping + ) + + is_expert_weight = False + loaded_expert_param_names: set[str] = set() + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + + # This is an expert weight and must not be attempted as any + # other kind of weight later on. + is_expert_weight = True + + # Do not modify `name`: the loop may continue past this point. + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + if is_fused_expert: + if "experts.gate_up_proj" in name: + chunks = loaded_weight.chunk(2, dim=-2) + success_w1 = self.load_fused_expert_weights( + name_mapped, params_dict, chunks[0], "w1", num_experts + ) + success_w3 = self.load_fused_expert_weights( + name_mapped, params_dict, chunks[1], "w3", num_experts + ) + success = success_w1 and success_w3 + else: + success = self.load_fused_expert_weights( + name_mapped, + params_dict, + loaded_weight, + shard_id, + num_experts, + ) + if success: + name = name_mapped + break + else: + # Split per-expert format (V3 style). + if name_mapped not in params_dict: + if _should_skip_missing_param(name_mapped): + continue + logger.warning_once( + "Skipping unknown checkpoint weight: %s", + name_mapped, + ) + continue + param = params_dict[name_mapped] + # Ask the weight loader whether it succeeded, otherwise we + # may skip experts that have other available replicas. + weight_loader: Callable[..., typing.Any] = param.weight_loader + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fused_expert: + loaded_expert_param_names.add(name_mapped) + continue + name = name_mapped + break + else: + if loaded_expert_param_names: + loaded_params.update(loaded_expert_param_names) + continue + if "learnable_sink_param" in name: + if is_pp_missing_parameter(name, self): + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + with torch.no_grad(): + params_dict[name][:n].copy_(narrow_weight) + else: + if is_expert_weight: + # An expert weight that is not mapped to this rank. + continue + if "gate.e_score_correction_bias" in name: + name = name.replace( + "gate.e_score_correction_bias", "expert_bias" + ) + if "router.gate." in name: + name = name.replace("router.", "") + if "hc_fn" in name: + name = name.replace("hc_fn", "hc_fn.weight") + if "hc_head_fn" in name: + name = name.replace("hc_head_fn", "hc_head_fn.weight") + + if is_pp_missing_parameter(name, self): + continue + + if name not in params_dict: + if _should_skip_missing_param(name): + continue + logger.warning_once( + "Skipping unknown checkpoint weight: %s", + name, + ) + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + if stacked_shard_id is None: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, stacked_shard_id) + loaded_params.add(name) + + return loaded_params + + +class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_a_proj": (".fused_qkv_a_proj", 0), + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0), + ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1), + ".indexer.weights_proj.": (".indexer.wk_weights_proj.", 1), + } + ) + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + # MLA runs both latent down-projections as one GEMM. + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + # The indexer fuses wk and weights_proj into one GEMM. + "wk_weights_proj": ["wk", "weights_proj"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + if quant_config is not None: + quant_config.packed_modules_mapping = self.packed_modules_mapping + + parallel_config = vllm_config.parallel_config + eplb_config = parallel_config.eplb_config + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.model = HYV4Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + # The head stays in the model dtype; ``enable_lm_head_fp32`` is + # surfaced as ``head_dtype`` on the config (see HYV4Config), which + # makes LogitsProcessor accumulate the projection straight into + # fp32 instead of materializing an fp32 copy of the weight. + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + # With tie_word_embeddings, embed_tokens is kept on the last rank + # (see HYV4Model.__init__) so the weight can be shared here. + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + + if getattr(self.config, "soft_logits_capping", False): + soft_cap = self.config.soft_logits_capping_logits + logits = soft_cap * torch.nn.functional.tanh(logits / soft_cap) + + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # MTP layers are loaded by the MTP head, not by the target model. They + # appear as model.mtp_layers..* and as the layer ids past the backbone. + mtp_start = self.config.num_hidden_layers + drop_prefixes: dict[str, str | None] = {"model.mtp_layers.": None} + for i in range(getattr(self.config, "num_nextn_predict_layers", 0)): + drop_prefixes[f"model.layers.{mtp_start + i}."] = None + if self.config.tie_word_embeddings: + drop_prefixes["lm_head."] = None + drop_mapper = WeightsMapper(orig_to_new_prefix=drop_prefixes) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper | drop_mapper) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() From 76ffb0d37488988ec158cab609d5569e9f9ca34c Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Mon, 21 Sep 2026 01:26:56 -0700 Subject: [PATCH 571/774] [Model][Engram] Share host tables across co-located DP replicas by default (#57651) Signed-off-by: Juntian Liu Signed-off-by: zjy0516 Co-authored-by: Claude Fable 5.1 Co-authored-by: zjy0516 --- tests/distributed/test_engram_dp_shard.py | 7 ++++++ tests/engine/test_arg_utils.py | 2 +- tests/test_config.py | 18 ++++++++++++--- vllm/config/engram.py | 28 +++++++++++------------ vllm/config/vllm.py | 2 +- vllm/models/deepseek_v41/nvidia/engram.py | 24 ++++++++++++++++++- vllm/models/deepseek_v41/nvidia/model.py | 9 +++++++- 7 files changed, 68 insertions(+), 22 deletions(-) diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index dd7105a23ba9..ba03387aac63 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -685,6 +685,13 @@ def test_engram_dp_shared_memory_runtime_requirements( ) +def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch): + """A /dev/shm smaller than the tables must fall back instead of failing startup.""" + monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: 2) + layout = SimpleNamespace(num_embeddings=(1 << 50,), head_dim=DIM) + assert not engram_ops.can_share_engram_tables(layout) + + @pytest.mark.parametrize( "node_ids,dp_size,replica_size,expected", [ diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 68fe25067713..bded621fad6c 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -114,7 +114,7 @@ def test_engram_config_cli(options): "options,provided,dp_shared_memory", [ ([], False, False), - (["--engram-config", "{}"], True, False), + (["--engram-config", "{}"], True, None), ( ["--engram-config", '{"dp_shared_memory": true}'], True, diff --git a/tests/test_config.py b/tests/test_config.py index f31adfdb8d69..cad5ffd627fd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1284,15 +1284,27 @@ def test_engram_dp_shared_memory_requires_cpu_offload(): EngramConfig(cpu_offload=False, dp_shared_memory=True) +@pytest.mark.parametrize( + "cpu_offload,dp_size,elastic_ep,expected", + [(True, 2, False, True), (False, 2, False, False), (True, 1, False, False)], +) +def test_engram_dp_shared_memory_defaults_when_supported( + cpu_offload, dp_size, elastic_ep, expected +): + """Unset dp_shared_memory enables sharing only for offloaded, non-elastic DP.""" + parallel = ParallelConfig(data_parallel_size=dp_size) + parallel.enable_elastic_ep = elastic_ep + config = EngramConfig(cpu_offload=cpu_offload) + config.resolve_dp_shared_memory(parallel) + assert config.dp_shared_memory is expected + + @pytest.mark.parametrize( "dp_size,load_format,multithread,error", [ (1, "auto", False, "requires data_parallel_size > 1"), - (2, "dummy", False, "requires load_format"), - (2, "sharded_state", False, "requires load_format"), (2, "auto", False, None), (2, "safetensors", True, None), - (2, "pt", True, None), ], ) def test_engram_dp_shared_memory_config_validation( diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 773ca2091153..c9f05c93828f 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -10,7 +10,6 @@ from vllm.config.utils import config, get_hash_factors, hash_factors if TYPE_CHECKING: - from vllm.config.load import LoadConfig from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig @@ -50,11 +49,13 @@ class EngramConfig: """Shard embeddings across TP and all DP ranks when enabled. Otherwise, each DP rank has a separate TP-sharded embedding replica.""" - dp_shared_memory: bool = False + dp_shared_memory: bool | None = None """Share CPU-offloaded embedding weights between co-located DP replicas. Each node stores one copy of every TP shard, reducing host memory without per-step Engram DP collectives. Requires sufficient - /dev/shm capacity and a shared IPC namespace.""" + /dev/shm capacity and a shared IPC namespace. Defaults to enabled whenever + the other settings allow it, falling back to per-replica tables when DP + replicas are not co-located on one node or /dev/shm cannot hold them.""" @model_validator(mode="after") def _validate_shared_memory(self) -> Self: @@ -82,6 +83,15 @@ def verify_model_config(self, model_config: "ModelConfig | None") -> None: "embeddings, non-empty n-gram layer ids, and CUDA." ) + def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None: + """Share host tables by default wherever the configuration permits.""" + if self.dp_shared_memory is None: + self.dp_shared_memory = ( + self.cpu_offload + and parallel_config.data_parallel_size > 1 + and not parallel_config.enable_elastic_ep + ) + def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: """Reject unsupported embedding parallel topologies.""" if self.dp_shared_memory: @@ -98,18 +108,6 @@ def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: "Engram embedding_across_dp is not supported with elastic EP yet." ) - def verify_load_config(self, load_config: "LoadConfig") -> None: - """Shared tables require a loader that invokes parameter weight callbacks.""" - if self.dp_shared_memory and load_config.load_format not in ( - "auto", - "safetensors", - "pt", - ): - raise ValueError( - "dp_shared_memory requires load_format 'auto', " - f"'safetensors' or 'pt'; got {load_config.load_format!r}." - ) - def get_parallel_size(self, parallel_config: "ParallelConfig") -> int: """Derive the embedding group size from the parallel configuration.""" size = parallel_config.tensor_parallel_size diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6974d1fbb3bd..1ff675a8f010 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1266,8 +1266,8 @@ def _resolve_and_verify_engram_config(self) -> None: return self.engram_config = EngramConfig() self.engram_config.verify_model_config(model_config) + self.engram_config.resolve_dp_shared_memory(self.parallel_config) self.engram_config.verify_parallel_config(self.parallel_config) - self.engram_config.verify_load_config(self.load_config) logger.info_once("Resolved Engram configuration: %s", str(self.engram_config)) def __post_init__(self): diff --git a/vllm/models/deepseek_v41/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py index 624a36c214fc..7dd820fee46d 100644 --- a/vllm/models/deepseek_v41/nvidia/engram.py +++ b/vllm/models/deepseek_v41/nvidia/engram.py @@ -3,6 +3,7 @@ """NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch.""" import mmap +import shutil import tempfile import weakref from contextlib import ExitStack @@ -89,6 +90,27 @@ def gather_engram_hashes( return dp_group.all_gather(hash_ids, dim=0) +def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool: + """Whether co-located DP replicas exist and /dev/shm can hold the full tables.""" + if get_engram_dp_size() == 1: + logger.warning_once( + "Engram DP replicas are not co-located on one node; " + "storing the offloaded tables per rank instead of sharing them." + ) + return False + num_bytes = sum(layout.num_embeddings) * ( + layout.head_dim + layout.head_dim // block_size + ) + if shutil.disk_usage("/dev/shm").total < num_bytes: + logger.warning_once( + "Sharing Engram tables across DP replicas needs %.1f GiB of /dev/shm " + "(--shm-size or --ipc=host); sharding them across replicas instead.", + num_bytes / 1024**3, + ) + return False + return True + + class DPSharedEngramStorage: """Registered host weights shared by a node-local DP group with one writer.""" @@ -362,7 +384,7 @@ def _create_embedding( layout.head_dim, tuple(size for order in layout.primes[layer_hash_index] for size in order), cpu_offload=engram_config.cpu_offload, - dp_shared_memory=engram_config.dp_shared_memory, + dp_shared_memory=bool(engram_config.dp_shared_memory), ) def _init_staging(self, max_tokens: int, head_dim: int) -> None: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b91ad1d7ba4a..d525aff3c707 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -88,7 +88,7 @@ from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask -from .engram import Engram, gather_engram_hashes +from .engram import Engram, can_share_engram_tables, gather_engram_hashes from .ops.mhc import ( MHC_OVERLAP_MAX_TOKENS, mhc_pre_delayed_overlap, @@ -589,6 +589,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.embed_tokens = PPMissingLayer() self.engram_layout = EngramLayout.from_config(config) + engram_config = vllm_config.engram_config + if ( + self.engram_layout is not None + and engram_config is not None + and engram_config.dp_shared_memory + ): + engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, From f05b88751ec74b66ee27eb38f65fdfdda6161239 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 03:29:02 -0500 Subject: [PATCH 572/774] [CI][ROCm] Make Python-only installation failures blocking (#57870) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/misc.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 1bff8e18500b..ef096e9bd338 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -416,7 +416,6 @@ steps: dind: false device: mi300_1 timeout_in_minutes: 55 - soft_fail: true depends_on: - image-build-amd source_file_dependencies: From b8cf2753825d05d7f41eb4747480b78934953aac Mon Sep 17 00:00:00 2001 From: weitliao Date: Mon, 21 Sep 2026 01:38:30 -0700 Subject: [PATCH 573/774] [AMD][Minimax-M3][perf] Enable packed LBHNC AITER QK-norm fusion for MiniMax-M3 on ROCm (#54535) Signed-off-by: Wei-Ting Liao Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/kernels/attention/test_minimax_m3.py | 371 ++++++++++++++++++++- vllm/_aiter_ops.py | 112 +++++++ vllm/models/minimax_m3/amd/model.py | 244 +++++++++----- 3 files changed, 636 insertions(+), 91 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index a7e6d499c176..9d3b1500e1d2 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -19,7 +19,10 @@ minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode, ) -from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl +from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseTritonImpl, + minimax_m3_rebase_slots_to_page16, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.utils import record_kv_cache_layout @@ -28,6 +31,7 @@ KVCacheLayout, MLAAttentionSpec, compute_layer_kv_cache_shape_bytes, + get_kv_quant_mode, ) if not (current_platform.is_cuda() or current_platform.is_rocm()): @@ -1672,8 +1676,9 @@ def test_main_cache_layout_contract(): assert set(order) == set(range(len(order))) -def test_aiter_sparse_pa_layout_contract(monkeypatch): - """The shuffle-only AITER path retains separately contiguous K/V storage.""" +@pytest.mark.parametrize("layout", [KVCacheLayout.LHBNC, KVCacheLayout.LBHNC]) +def test_aiter_sparse_pa_layout_contract(monkeypatch, layout): + """Build the exact separate and packed page-16 views used by the model.""" import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True) @@ -1685,21 +1690,43 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch): backend = sparse_attn_mod.MiniMaxM3SparseBackend assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts() + assert KVCacheLayout.LBHNC in backend.supported_kv_cache_layouts() - nb, h = 7, 1 + nb, h = 3, 1 + cache_dtype = current_platform.fp8_dtype() spec = FullAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=h, head_size=HEAD_DIM, - dtype=DTYPE, + dtype=cache_dtype, + kv_quant_mode=get_kv_quant_mode("fp8"), num_head_slots=2, - state_content_bytes=h * HEAD_DIM * DTYPE.itemsize, + state_content_bytes=h * HEAD_DIM * cache_dtype.itemsize, ) raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) - view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0] - key_cache, value_cache = view.unbind(1) - assert key_cache.is_contiguous() - assert value_cache.is_contiguous() + kv_cache = dense_kv_cache_views(raw, spec, nb, 1, layout)[0].view(cache_dtype) + key_cache, value_cache = kv_cache.unbind(1) + x = 16 // cache_dtype.itemsize + pages_per_side = BLOCK_SIZE // 16 + if key_cache.is_contiguous() and value_cache.is_contiguous(): + k_src, v_src = key_cache, value_cache + block_pages, v_page_offset = pages_per_side, 0 + else: + assert kv_cache.is_contiguous() + k_src = v_src = kv_cache + block_pages, v_page_offset = 2 * pages_per_side, pages_per_side + + num_phys16 = nb * block_pages + k_view = k_src.view(num_phys16, h, HEAD_DIM // x, 16, x) + v_view = v_src.view(num_phys16, h, 16 // x, HEAD_DIM, x)[v_page_offset:] + assert k_view.is_contiguous() + assert v_view.is_contiguous() + if layout is KVCacheLayout.LHBNC: + assert k_view.shape == (24, 1, 8, 16, 16) + assert v_view.shape == (24, 1, 1, 128, 16) + else: + assert k_view.shape == (48, 1, 8, 16, 16) + assert v_view.shape == (40, 1, 1, 128, 16) def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch): @@ -1769,6 +1796,330 @@ def test_aiter_indexer_requires_the_aiter_attend(monkeypatch): assert reason is None or "AITER sparse PA attend" not in reason +def test_aiter_consolidated_qknorm_enabled(monkeypatch): + from vllm._aiter_ops import rocm_aiter_ops + + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True) + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) + + scale = torch.tensor([0.25]) + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto") + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, scale) + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e4m3", scale, scale) + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8") + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, None) + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale) + + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", False) + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto") + + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: False) + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto") + + +def _fake_aiter_qknorm_args(kv_cache_dtype: str = "auto") -> dict: + tensor = torch.empty(1) + return { + "qkv": tensor, + "q_norm_weight": tensor, + "k_norm_weight": tensor, + "cos_sin_cache": tensor, + "positions": tensor, + "num_heads": 16, + "num_kv_heads": 1, + "rotary_dim": 64, + "eps": 1e-6, + "slot_mapping": tensor, + "kv_cache_k": tensor, + "kv_cache_v": tensor, + "q_out": tensor, + "kv_cache_dtype": kv_cache_dtype, + } + + +def _install_fake_aiter_qknorm(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + """Stub ``aiter.fused_qknorm_idxrqknorm`` so mapping tests run without AITER.""" + import sys + import types + from typing import Any + + calls: list[dict] = [] + + def fake_op(*args, **kwargs): + calls.append(kwargs) + + existing = sys.modules.get("aiter") + if existing is not None: + monkeypatch.setattr(existing, "fused_qknorm_idxrqknorm", fake_op, raising=False) + else: + fake_mod: Any = types.ModuleType("aiter") + fake_mod.fused_qknorm_idxrqknorm = fake_op + monkeypatch.setitem(sys.modules, "aiter", fake_mod) + return calls + + +def test_aiter_consolidated_qknorm_packed_shuffle_unequal_spans(monkeypatch): + from vllm._aiter_ops import rocm_aiter_ops + + calls = _install_fake_aiter_qknorm(monkeypatch) + call_args = _fake_aiter_qknorm_args() + call_args["kv_cache_k"] = torch.empty(2, 1) + call_args["kv_cache_v"] = torch.empty(1, 1) + + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args) + assert len(calls) == 1 + + +@pytest.mark.parametrize( + ("kv_cache_dtype", "aiter_dtype", "passes_scales"), + [ + ("auto", "auto", False), + ("fp8", "fp8_e4m3_static", True), + ("fp8_e4m3", "fp8_e4m3_static", True), + ], +) +def test_aiter_consolidated_qknorm_dtype_and_layout_mapping( + monkeypatch, kv_cache_dtype, aiter_dtype, passes_scales +): + from vllm._aiter_ops import rocm_aiter_ops + + calls = _install_fake_aiter_qknorm(monkeypatch) + call_args = _fake_aiter_qknorm_args(kv_cache_dtype) + k_scale = torch.tensor([0.25]) + v_scale = torch.tensor([0.5]) + call_args.update(k_scale=k_scale, v_scale=v_scale) + + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args) + assert len(calls) == 1 + kwargs = calls[0] + assert kwargs["kv_cache_dtype"] == aiter_dtype + assert kwargs["index_cache_dtype"] == "auto" + assert kwargs["block_size"] == 16 + assert kwargs["asm_layout"] is True + if passes_scales: + assert kwargs["k_scale"] is k_scale + assert kwargs["v_scale"] is v_scale + else: + assert kwargs["k_scale"] is None + assert kwargs["v_scale"] is None + + +def test_aiter_consolidated_qknorm_fp8_index_dtype_mapping(monkeypatch): + from vllm._aiter_ops import rocm_aiter_ops + + calls = _install_fake_aiter_qknorm(monkeypatch) + fp8 = current_platform.fp8_dtype() + call_args = _fake_aiter_qknorm_args("fp8") + call_args.update( + k_scale=torch.tensor([0.25]), + v_scale=torch.tensor([0.5]), + index_cache=torch.empty(1, dtype=fp8), + index_q_out=torch.empty(1, dtype=fp8), + num_index_heads=1, + ) + + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args) + assert calls[0]["index_cache_dtype"] == "fp8" + assert calls[0]["kv_cache_dtype"] == "fp8_e4m3_static" + + +@pytest.mark.parametrize("skip_index_branch", [False, True]) +def test_aiter_consolidated_qknorm_full_and_skip_index_args( + monkeypatch, skip_index_branch +): + from vllm._aiter_ops import rocm_aiter_ops + + calls = _install_fake_aiter_qknorm(monkeypatch) + call_args = _fake_aiter_qknorm_args() + if skip_index_branch: + call_args.update(skip_index_branch=True) + else: + tensor = torch.empty(1) + call_args.update( + index_q_norm_weight=tensor, + index_k_norm_weight=tensor, + num_index_heads=1, + index_cache=tensor, + index_q_out=tensor, + index_slot_mapping=tensor, + ) + + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args) + kwargs = calls[0] + assert kwargs["skip_index_branch"] is skip_index_branch + if skip_index_branch: + assert kwargs["index_q_norm_weight"] is None + assert kwargs["index_k_norm_weight"] is None + assert kwargs["index_cache"] is None + assert kwargs["index_q_out"] is None + assert kwargs["index_slot_mapping"] is None + else: + assert kwargs["index_q_norm_weight"] is not None + assert kwargs["index_k_norm_weight"] is not None + assert kwargs["index_cache"] is not None + assert kwargs["index_q_out"] is not None + assert kwargs["index_slot_mapping"] is not None + + +def test_aiter_consolidated_qknorm_rejects_fp8_e5m2(monkeypatch): + from vllm._aiter_ops import rocm_aiter_ops + + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True) + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) + scale = torch.tensor([0.25]) + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale) + + +def test_aiter_sparse_pa_slot_rebase_preserves_page_offsets(): + slots = torch.tensor([0, 15, 16, 127, 128, 129, -1]) + actual = minimax_m3_rebase_slots_to_page16(slots, block_size=128) + expected = torch.tensor([0, 15, 16, 127, 256, 257, -1]) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="The consolidated packed-SHUFFLE AITER op is ROCm-only", +) +@pytest.mark.parametrize("skip_index_branch", [False, True]) +@pytest.mark.parametrize("fp8_index", [False, True]) +def test_aiter_consolidated_qknorm_real_packed_shuffle(skip_index_branch, fp8_index): + from vllm._aiter_ops import rocm_aiter_ops + + torch.manual_seed(20260831) + num_tokens = 4 + num_heads = 16 + num_kv_heads = 1 + num_index_heads = 1 + rotary_dim = 64 + head_dim = 128 + logical_block_size = 128 + num_blocks = 3 + fp8_dtype = current_platform.fp8_dtype() + + qkv = torch.randn( + num_tokens, + (num_heads + 2 * num_kv_heads + num_index_heads + 1) * head_dim, + dtype=torch.bfloat16, + device="cuda", + ) + weights = [ + torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") for _ in range(4) + ] + cos_sin_cache = torch.randn(16, rotary_dim, dtype=torch.bfloat16, device="cuda") + positions = torch.arange(num_tokens, dtype=torch.int64, device="cuda") + logical_slots = torch.tensor([0, 127, 128, -1], dtype=torch.int64, device="cuda") + slot_mapping = minimax_m3_rebase_slots_to_page16(logical_slots, logical_block_size) + index_slot_mapping = logical_slots.clone() + + packed = torch.zeros( + num_blocks, + 2, + logical_block_size, + head_dim, + dtype=fp8_dtype, + device="cuda", + ) + x = 16 // packed.element_size() + pages_per_side = logical_block_size // 16 + num_phys16 = num_blocks * 2 * pages_per_side + kv_cache_k = packed.view(num_phys16, num_kv_heads, head_dim // x, 16, x) + kv_cache_v = packed.view(num_phys16, num_kv_heads, 16 // x, head_dim, x)[ + pages_per_side: + ] + assert kv_cache_k.shape == (48, 1, 8, 16, 16) + assert kv_cache_v.shape == (40, 1, 1, 128, 16) + + q_out = torch.zeros( + num_tokens, num_heads * head_dim, dtype=torch.bfloat16, device="cuda" + ) + index_dtype = fp8_dtype if fp8_index else torch.bfloat16 + if fp8_index: + index_q_out = torch.empty( + (num_tokens, num_index_heads * head_dim), + dtype=index_dtype, + device="cuda", + ) + index_cache = torch.empty( + (num_blocks, logical_block_size, head_dim), + dtype=index_dtype, + device="cuda", + ) + index_q_out.view(torch.uint8).fill_(7) + index_cache.view(torch.uint8).fill_(7) + else: + index_q_out = torch.full( + (num_tokens, num_index_heads * head_dim), + 7, + dtype=torch.bfloat16, + device="cuda", + ) + index_cache = torch.full( + (num_blocks, logical_block_size, head_dim), + 7, + dtype=torch.bfloat16, + device="cuda", + ) + scale = torch.tensor([0.02], dtype=torch.float32, device="cuda") + + kwargs = {} + if not skip_index_branch: + kwargs = { + "index_q_norm_weight": weights[2], + "index_k_norm_weight": weights[3], + "index_cache": index_cache, + "index_q_out": index_q_out, + "index_slot_mapping": index_slot_mapping, + } + rocm_aiter_ops.fused_qknorm_idxrqknorm( + qkv, + weights[0], + weights[1], + cos_sin_cache, + positions, + num_heads, + num_kv_heads, + rotary_dim, + 1e-6, + slot_mapping, + kv_cache_k, + kv_cache_v, + q_out, + "fp8", + scale, + scale, + num_index_heads=num_index_heads, + skip_index_branch=skip_index_branch, + **kwargs, + ) + torch.accelerator.synchronize() + + assert torch.isfinite(q_out).all() + assert torch.count_nonzero(q_out) > 0 + assert torch.count_nonzero(packed) > 0 + if fp8_index: + index_q_untouched = torch.all(index_q_out.view(torch.uint8) == 7) + index_cache_untouched = torch.all(index_cache.view(torch.uint8) == 7) + index_cache_t0 = torch.all(index_cache[0, 0].view(torch.uint8) == 7) + index_cache_t127 = torch.all(index_cache[0, 127].view(torch.uint8) == 7) + index_cache_t128 = torch.all(index_cache[1, 0].view(torch.uint8) == 7) + else: + index_q_untouched = torch.all(index_q_out == 7) + index_cache_untouched = torch.all(index_cache == 7) + index_cache_t0 = torch.all(index_cache[0, 0] == 7) + index_cache_t127 = torch.all(index_cache[0, 127] == 7) + index_cache_t128 = torch.all(index_cache[1, 0] == 7) + if skip_index_branch: + assert index_q_untouched + assert index_cache_untouched + else: + assert not index_q_untouched + assert not index_cache_t0 + assert not index_cache_t127 + assert not index_cache_t128 + + def test_indexer_cache_squeezes_to_contiguous_3d(): """The indexer side cache is standardized 4D with H=1: under both layouts the allocator's logical view stays contiguous and squeezes (as diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index c514b7c0bfef..1829d9f5e247 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -35,10 +35,22 @@ # which is a host op, so we cache it once here. FP8_DTYPE = current_platform.fp8_dtype() _HIPB_MM_INITIALIZED_DEVICES: set[int] = set() +_FP8_E4M3_DTYPES = { + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e4m3fnuz", None), + ) + if dtype is not None +} KB = 1024 MB = 1024 * KB +def _is_fp8_e4m3_tensor(t: torch.Tensor | None) -> bool: + return t is not None and t.dtype in _FP8_E4M3_DTYPES + + def _get_or_create_aiter_qr_rmsnorm_comm( device_comm: Any, ) -> tuple[int, int, bool] | None: @@ -2200,6 +2212,106 @@ def is_custom_all_reduce_enabled(cls) -> bool: def is_shuffle_kv_cache_enabled(cls) -> bool: return cls._SHUFFLE_KV_CACHE_ENABLED + @classmethod + @if_aiter_supported + def fused_qknorm_idxrqknorm_enabled( + cls, + kv_cache_dtype: str, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + ) -> bool: + """Whether MiniMax-M3 can use AITER's consolidated fused QK-norm. + + Requires AITER to be installed and enabled. bf16 (``auto``) always + qualifies; fp8 e4m3 also needs K/V scales. Other cache dtypes fall + back to vLLM's fused kernel. + """ + if not cls._AITER_ENABLED: + return False + if kv_cache_dtype == "auto": + return True + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + return k_scale is not None and v_scale is not None + return False + + @classmethod + def fused_qknorm_idxrqknorm( + cls, + qkv: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads: int, + num_kv_heads: int, + rotary_dim: int, + eps: float, + slot_mapping: torch.Tensor, + kv_cache_k: torch.Tensor, + kv_cache_v: torch.Tensor, + q_out: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + index_q_norm_weight: torch.Tensor | None = None, + index_k_norm_weight: torch.Tensor | None = None, + num_index_heads: int = 0, + index_cache: torch.Tensor | None = None, + index_q_out: torch.Tensor | None = None, + index_slot_mapping: torch.Tensor | None = None, + skip_index_branch: bool = False, + ) -> None: + """Run consolidated MiniMax-M3 QK-norm fusion. + + Callers must check ``fused_qknorm_idxrqknorm_enabled`` first. Runtime + failures from the AITER op are deliberately propagated. + """ + if kv_cache_dtype == "auto": + aiter_kv_cache_dtype = "auto" + aiter_k_scale = None + aiter_v_scale = None + elif kv_cache_dtype in ("fp8", "fp8_e4m3"): + aiter_kv_cache_dtype = "fp8_e4m3_static" + aiter_k_scale = k_scale + aiter_v_scale = v_scale + else: + raise ValueError( + "AITER fused QK-norm requires kv_cache_dtype 'auto', 'fp8', " + f"or 'fp8_e4m3', got {kv_cache_dtype!r}" + ) + aiter_index_cache_dtype = "fp8" if _is_fp8_e4m3_tensor(index_cache) else "auto" + + from aiter import fused_qknorm_idxrqknorm + + fused_qknorm_idxrqknorm( + qkv, + q_norm_weight, + k_norm_weight, + cos_sin_cache, + positions, + num_heads, + num_kv_heads, + rotary_dim, + eps, + index_q_norm_weight=index_q_norm_weight, + index_k_norm_weight=index_k_norm_weight, + num_index_heads=num_index_heads, + slot_mapping=slot_mapping, + kv_cache_k=kv_cache_k, + kv_cache_v=kv_cache_v, + index_cache=index_cache, + block_size=16, + q_out=q_out, + index_q_out=index_q_out, + index_slot_mapping=index_slot_mapping, + kv_cache_dtype=aiter_kv_cache_dtype, + index_cache_dtype=aiter_index_cache_dtype, + k_scale=aiter_k_scale, + v_scale=aiter_v_scale, + asm_layout=True, + skip_index_branch=skip_index_branch, + ) + @classmethod @if_aiter_supported def is_triton_unified_attn_enabled(cls) -> bool: diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index a0a58253d1a8..88651e944a82 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -890,6 +890,30 @@ def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]: self._ensure_aiter_sparse_pa_kv_cache() return self.kv_cache_k, self.kv_cache_v + def _get_aiter_sparse_pa_slot_mapping( + self, + slot_mapping: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + ) -> torch.Tensor: + if value_cache.shape[0] == key_cache.shape[0]: + return slot_mapping + + attn_metadata = get_forward_context().attn_metadata + page16_slot_mapping = None + if isinstance(attn_metadata, dict): + main_md = attn_metadata[self.layer_name] + assert isinstance(main_md, MiniMaxM3SparseMetadata) + page16_slot_mapping = main_md.page16_slot_mapping + if ( + page16_slot_mapping is None + or page16_slot_mapping.shape != slot_mapping.shape + ): + page16_slot_mapping = minimax_m3_rebase_slots_to_page16( + slot_mapping, self.kv_cache.shape[2] + ) + return page16_slot_mapping + def _insert_aiter_sparse_pa_kv( self, k: torch.Tensor, @@ -897,6 +921,8 @@ def _insert_aiter_sparse_pa_kv( index_k: torch.Tensor | None, slot_mapping: torch.Tensor, index_slot_mapping: torch.Tensor | None, + key_cache: torch.Tensor, + value_cache: torch.Tensor, ) -> None: if self.kv_cache.numel() == 0: return @@ -906,22 +932,6 @@ def _insert_aiter_sparse_pa_kv( minimax_m3_insert_index_cache, ) - key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache() - if value_cache.shape[0] != key_cache.shape[0]: - attn_metadata = get_forward_context().attn_metadata - page16_slot_mapping = None - if isinstance(attn_metadata, dict): - main_md = attn_metadata[self.layer_name] - assert isinstance(main_md, MiniMaxM3SparseMetadata) - page16_slot_mapping = main_md.page16_slot_mapping - if ( - page16_slot_mapping is None - or page16_slot_mapping.shape != slot_mapping.shape - ): - page16_slot_mapping = minimax_m3_rebase_slots_to_page16( - slot_mapping, self.kv_cache.shape[2] - ) - slot_mapping = page16_slot_mapping kv_cache_dtype = ( self.kv_cache_dtype if is_quantized_kv_cache(self.kv_cache_dtype) @@ -976,38 +986,78 @@ def forward( main_slot_mapping = fwd_slot_mapping[self.layer_name] q = qkv.new_empty((num_tokens, self.q_size)) + key_cache = value_cache = page16_slot_mapping = None + k_scale = v_scale = None + use_fused_qknorm = False + if self.use_aiter_sparse_pa: + key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache() + page16_slot_mapping = self._get_aiter_sparse_pa_slot_mapping( + main_slot_mapping, key_cache, value_cache + ) + k_scale = getattr(self, "_k_scale", None) + v_scale = getattr(self, "_v_scale", None) + use_fused_qknorm = rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled( + self.kv_cache_dtype, k_scale, v_scale + ) if self.skip_index_topk: index_q = None if self.use_aiter_sparse_pa: - ops.fused_minimax_m3_qknorm_rope_kv_insert( - qkv, - self.q_norm.weight, - self.k_norm.weight, - cos_sin_cache, - positions, - self.num_heads, - self.num_kv_heads, - rotary_dim, - eps, - num_index_heads=self.num_idx_heads, - q_out=q, - skip_index_branch=True, - ) - k_start = self.q_size - v_start = k_start + self.kv_size - k = qkv[:, k_start:v_start].view( - num_tokens, self.num_kv_heads, self.head_dim - ) - v = qkv[:, v_start : v_start + self.kv_size].view( - num_tokens, self.num_kv_heads, self.head_dim - ) - self._insert_aiter_sparse_pa_kv( - k, - v, - None, - main_slot_mapping, - None, - ) + assert key_cache is not None + assert value_cache is not None + assert page16_slot_mapping is not None + if use_fused_qknorm: + rocm_aiter_ops.fused_qknorm_idxrqknorm( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + page16_slot_mapping, + key_cache, + value_cache, + q, + self.kv_cache_dtype, + k_scale, + v_scale, + num_index_heads=self.num_idx_heads, + skip_index_branch=True, + ) + else: + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + num_index_heads=self.num_idx_heads, + q_out=q, + skip_index_branch=True, + ) + k_start = self.q_size + v_start = k_start + self.kv_size + k = qkv[:, k_start:v_start].view( + num_tokens, self.num_kv_heads, self.head_dim + ) + v = qkv[:, v_start : v_start + self.kv_size].view( + num_tokens, self.num_kv_heads, self.head_dim + ) + self._insert_aiter_sparse_pa_kv( + k, + v, + None, + page16_slot_mapping, + None, + key_cache, + value_cache, + ) else: ops.fused_minimax_m3_qknorm_rope_kv_insert( qkv, @@ -1037,42 +1087,74 @@ def forward( dtype=self.indexer.index_cache.dtype, ) if self.use_aiter_sparse_pa: - ops.fused_minimax_m3_qknorm_rope_kv_insert( - qkv, - self.q_norm.weight, - self.k_norm.weight, - cos_sin_cache, - positions, - self.num_heads, - self.num_kv_heads, - rotary_dim, - eps, - self.index_q_norm.weight, - self.index_k_norm.weight, - self.num_idx_heads, - q_out=q, - index_q_out=index_q, - kv_cache_dtype=self.kv_cache_dtype, - ) - k_start = self.q_size - v_start = k_start + self.kv_size - index_k_start = v_start + self.kv_size + self.index_q_size - k = qkv[:, k_start:v_start].view( - num_tokens, self.num_kv_heads, self.head_dim - ) - v = qkv[:, v_start : v_start + self.kv_size].view( - num_tokens, self.num_kv_heads, self.head_dim - ) - index_k = qkv[ - :, index_k_start : index_k_start + self.idx_head_dim - ].view(num_tokens, self.idx_head_dim) - self._insert_aiter_sparse_pa_kv( - k, - v, - index_k, - main_slot_mapping, - index_slot_mapping, - ) + assert key_cache is not None + assert value_cache is not None + assert page16_slot_mapping is not None + index_cache = self.indexer.index_cache.kv_cache + if use_fused_qknorm: + rocm_aiter_ops.fused_qknorm_idxrqknorm( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + page16_slot_mapping, + key_cache, + value_cache, + q, + self.kv_cache_dtype, + k_scale, + v_scale, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.num_idx_heads, + index_cache, + index_q, + index_slot_mapping, + ) + else: + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.num_idx_heads, + q_out=q, + index_q_out=index_q, + kv_cache_dtype=self.kv_cache_dtype, + ) + k_start = self.q_size + v_start = k_start + self.kv_size + index_k_start = v_start + self.kv_size + self.index_q_size + k = qkv[:, k_start:v_start].view( + num_tokens, self.num_kv_heads, self.head_dim + ) + v = qkv[:, v_start : v_start + self.kv_size].view( + num_tokens, self.num_kv_heads, self.head_dim + ) + index_k = qkv[ + :, index_k_start : index_k_start + self.idx_head_dim + ].view(num_tokens, self.idx_head_dim) + self._insert_aiter_sparse_pa_kv( + k, + v, + index_k, + page16_slot_mapping, + index_slot_mapping, + key_cache, + value_cache, + ) else: ops.fused_minimax_m3_qknorm_rope_kv_insert( qkv, From 8902dbb1770b7b1780bf3d978f16cd208702472a Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 04:16:26 -0500 Subject: [PATCH 574/774] [CI][ROCm] Mirror remaining portable suites without duplicate coverage (#57877) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test-amd.yaml | 174 ++++++++++++++---- .buildkite/test_areas/engine.yaml | 19 +- .buildkite/test_areas/kernels.yaml | 72 +++++--- .buildkite/test_areas/models_basic.yaml | 11 ++ .buildkite/test_areas/models_language.yaml | 23 ++- .buildkite/test_areas/pytorch.yaml | 14 ++ .buildkite/test_areas/samplers.yaml | 23 ++- tests/compile/h100/test_startup.py | 24 ++- .../test_triton_unified_attention.py | 3 +- .../test_triton_unified_attention_diffkv.py | 81 ++++---- tests/models/glm5next/test_kda_recurrent.py | 17 +- 11 files changed, 352 insertions(+), 109 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index fe2cfd18c89e..e0c081286608 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -436,7 +436,7 @@ steps: - tests/v1/e2e/general/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" - label: ":amd: (MI250) E2E Scheduling" # TBD timeout_in_minutes: 55 @@ -1737,29 +1737,6 @@ steps: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir -- label: ":amd: (MI300) MLA Kernels" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - - vllm/platforms/rocm.py - - vllm/_aiter_ops.py - - tests/kernels/attention/test_rocm_aiter_mla_decode.py - - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - optional: true - commands: - - amd-smi - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - - label: ":amd: (MI300) Attention Kernels Shard %N" # TBD timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1778,7 +1755,7 @@ steps: - vllm/platforms/rocm.py - vllm/utils/flashinfer.py commands: - - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD timeout_in_minutes: 45 @@ -2566,7 +2543,7 @@ steps: - vllm/platforms/rocm.py - vllm/entrypoints/generate/beam_search commands: - - pytest -v -s samplers + - pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data" #------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# @@ -3364,6 +3341,48 @@ steps: commands: - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" +#---------------------------------------------------------- mi355 · compile ----------------------------------------------------------# + +- label: ":amd: (MI355) PyTorch Compilation H100 Cases" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/__init__.py + - vllm/_aiter_ops.py + - vllm/_custom_ops.py + - vllm/compilation/ + - vllm/config/ + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/env_override.py + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/ir/ + - vllm/kernels/ + - vllm/logger.py + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/plugins/ + - vllm/sampling_params.py + - vllm/sequence.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/compile/h100/ + - tests/utils.py + - csrc/ + commands: + - pytest -v -s compile/h100/test_startup.py + #-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# - label: ":amd: (MI355) Distributed DP Basic" # TBD @@ -3790,6 +3809,23 @@ steps: #---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------# +- label: ":amd: (MI355) Attention DiffKV Kernels" + timeout_in_minutes: 35 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py + - vllm/v1/attention/backends/triton_attn_diffkv.py + - tests/kernels/attention/test_triton_unified_attention.py + - tests/kernels/attention/test_triton_unified_attention_diffkv.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py + - label: ":amd: (MI355) Core Operation Kernels Shard %N" # TBD timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3859,19 +3895,19 @@ steps: - vllm/_aiter_ops.py - examples/basic/offline_inference/chat.py - tests/kernels/attention/test_attention_selector.py - - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py optional: true commands: - amd-smi - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN - pytest -v -s tests/kernels/attention/test_attention_selector.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py -- label: ":amd: (MI355) MLA Kernels" # TBD - timeout_in_minutes: 25 +- label: ":amd: (MI355) MLA Kernels" + timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 + num_gpus: 1 + optional: true working_dir: "/vllm-workspace/" source_file_dependencies: - vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -3882,7 +3918,6 @@ steps: - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - optional: true commands: - amd-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -3908,7 +3943,7 @@ steps: - vllm/utils/flashinfer.py optional: true commands: - - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355) MoE Kernels Shard %N" # TBD timeout_in_minutes: 55 @@ -4050,6 +4085,21 @@ steps: #------------------------------------------------------ mi355 · models / basic -------------------------------------------------------# +- label: ":amd: (MI355) GLM5Next" + timeout_in_minutes: 35 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/models/glm5next/ + - tests/models/glm5next/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s models/glm5next + - label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N" # TBD timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4087,7 +4137,27 @@ steps: - export GIT_CONFIG_VALUE_0=HTTP/1.1 - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +- label: ":amd: (MI355) Granite Language Model Compatibility" + timeout_in_minutes: 80 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - "!vllm/distributed/kv_transfer/" + - tests/models/language/generation + commands: + - export GIT_CONFIG_COUNT=1 + - export GIT_CONFIG_KEY_0=http.version + - export GIT_CONFIG_VALUE_0=HTTP/1.1 + - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' - label: ":amd: (MI355) Language Models (Extended Generation)" # TBD timeout_in_minutes: 95 @@ -4308,8 +4378,46 @@ steps: # commands: # - pytest -s -v tests/quantization/test_gfx950_moe.py +#--------------------------------------------------------- mi355 · samplers ----------------------------------------------------------# + +- label: ":amd: (MI355) Samplers Multimodal Beam Search" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - vllm/v1/sample/ + - vllm/entrypoints/generate/beam_search/ + - tests/samplers + - tests/conftest.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data + #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# +- label: ":amd: (MI355) E2E Core Large Memory" + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py + - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + - label: ":amd: (MI355) V1 Attention Shard %N" # TBD timeout_in_minutes: 125 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 8e47a84e18cd..e38cd925732d 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -94,7 +94,7 @@ steps: label: ":amd: (MI250) E2E Core" dind: false commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" device: mi250_1 timeout_in_minutes: 50 depends_on: @@ -115,6 +115,23 @@ steps: commands: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + mirror: + amd: + label: ":amd: (MI355) E2E Core Large Memory" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 65 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py + - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async - label: ":nvidia: (L4) V1 E2E" key: v1-e2e-2-gpus diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index aab6e52b8eb9..2134d6dd040b 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -196,6 +196,8 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (H100) Attention DiffKV Kernels" key: kernels-attention-diffkv-test-h100 @@ -209,6 +211,24 @@ steps: - tests/kernels/attention/test_triton_unified_attention_diffkv.py commands: - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py + mirror: + amd: + label: ":amd: (MI355) Attention DiffKV Kernels" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 35 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py + - vllm/v1/attention/backends/triton_attn_diffkv.py + - tests/kernels/attention/test_triton_unified_attention.py + - tests/kernels/attention/test_triton_unified_attention_diffkv.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py - label: ":nvidia: (H100) MLA Kernel Test" key: kernels-flashmla-test-h100 @@ -231,6 +251,31 @@ steps: - pytest -v -s kernels/attention/test_flashmla_sparse.py - pytest -v -s kernels/attention/test_flashinfer_mla_decode.py - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py + mirror: + amd: + label: ":amd: (MI355) MLA Kernels" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + - tests/kernels/attention/test_rocm_aiter_mla_decode.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py + - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py + commands: + - amd-smi + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: ":nvidia: (L4) Quantization Kernels Shard %N" key: kernels-quantization-test @@ -640,31 +685,6 @@ steps: - pytest -v -s kernels/moe/test_moe_layer.py - pytest -v -s kernels/moe/test_deepep_v2_moe.py -- label: ":amd: (MI355) MLA Kernels" - key: amd-mla-kernels-mi355 - dind: false - device: mi355_1 - num_devices: 1 - timeout_in_minutes: 25 - working_dir: /vllm-workspace/ - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - - vllm/platforms/rocm.py - - vllm/_aiter_ops.py - - tests/kernels/attention/test_rocm_aiter_mla_decode.py - - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - commands: - - amd-smi - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - - label: ":amd: (MI355) FP8 MoE Kernels" key: amd-fp8-moe-kernels-mi355 dind: false @@ -720,9 +740,7 @@ steps: - vllm/_aiter_ops.py - examples/basic/offline_inference/chat.py - tests/kernels/attention/test_attention_selector.py - - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py commands: - amd-smi - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN - pytest -v -s tests/kernels/attention/test_attention_selector.py - - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 6702eb6dd3fe..930c2dd5902e 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -119,6 +119,17 @@ steps: - tests/models/glm5next/ commands: - pytest -v -s models/glm5next + mirror: + amd: + label: ":amd: (MI355) GLM5Next" + dind: false + device: mi355_1 + timeout_in_minutes: 35 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (H200) Qwen4 Exp" key: qwen4-exp-unit-tests diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 3f80dfdfa3c4..c8473e7f0b60 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -91,7 +91,7 @@ steps: commands: - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB # Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD # autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one @@ -111,6 +111,27 @@ steps: - rm -rf /tmp/mamba-src && git clone --depth 1 --branch v2.3.0 https://github.com/state-spaces/mamba /tmp/mamba-src && sed -i 's/-std=c++17/-std=c++20/g' /tmp/mamba-src/setup.py && MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation /tmp/mamba-src - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' + mirror: + amd: + label: ":amd: (MI355) Granite Language Model Compatibility" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 80 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/ + - "!vllm/distributed/kv_transfer/" + - tests/models/language/generation + commands: + - export GIT_CONFIG_COUNT=1 + - export GIT_CONFIG_KEY_0=http.version + - export GIT_CONFIG_VALUE_0=HTTP/1.1 + - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' - label: ":nvidia: (H200 MIG 35GB) Language Models (Extended Generation)" device: h200_35gb diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 459f7faeb4a1..8204fb9670a2 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -87,6 +87,20 @@ steps: - tests/compile/h100/ commands: - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" + mirror: + amd: + label: ":amd: (MI355) PyTorch Compilation H100 Cases" + dind: false + device: mi355_1 + timeout_in_minutes: 45 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - tests/utils.py + - csrc/ + commands: + - pytest -v -s compile/h100/test_startup.py - label: ":nvidia: (L4) PyTorch Compilation Passes" key: pytorch-compilation-passes-unit-tests diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index e1bc4c08d8ad..084583c8a732 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -35,7 +35,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s samplers + - pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data" # Qwen2-Audio-7B leaves no room for KV cache on a 16 GiB slice. - label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search" @@ -51,3 +51,24 @@ steps: commands: - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data + mirror: + amd: + label: ":amd: (MI355) Samplers Multimodal Beam Search" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - vllm/v1/sample/ + - vllm/entrypoints/generate/beam_search/ + - tests/samplers + - tests/conftest.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data diff --git a/tests/compile/h100/test_startup.py b/tests/compile/h100/test_startup.py index e57cf85b4a77..a56cf15e5c67 100644 --- a/tests/compile/h100/test_startup.py +++ b/tests/compile/h100/test_startup.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Cold start and warm start tests for vLLM-compile. -Cold start runs in a forked child (must fork before CUDA init) which +Cold start runs in a separate child process which populates on-disk caches and asserts cold-start counters. Warm start then runs in the parent with clean in-memory state but populated caches. """ @@ -18,7 +18,7 @@ from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig from vllm.utils.torch_utils import is_torch_equal_or_newer -from ...utils import fork_new_process_for_each_test +from ...utils import create_new_process_for_each_test, requires_spawn_multiprocessing MODEL = "microsoft/Phi-tiny-MoE-instruct" @@ -54,17 +54,21 @@ def _cold_start(vllm_runner): assert counters["aot_autograd"]["autograd_cache_hit"] == 0 -@fork_new_process_for_each_test +@create_new_process_for_each_test() @pytest.mark.parametrize("mega_aot_artifact", ["0", "1"]) def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact): monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact) monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") - # Cold start in a forked child (must fork before CUDA init). + # ROCm and XPU cannot fork safely after device discovery. # This model has 32 identical transformer layers which produce # 33 subgraphs after splitting on attention — only 3 are unique. - ctx = mp.get_context("fork") + ctx = ( + mp.get_context("spawn") + if requires_spawn_multiprocessing() + else mp.get_context("fork") + ) p = ctx.Process(target=_cold_start, args=(vllm_runner,)) p.start() p.join() @@ -223,13 +227,17 @@ def _cold_start_model(vllm_runner, spec: ModelStartupSpec): @pytest.mark.parametrize("spec", MODEL_SPECS) -@fork_new_process_for_each_test +@create_new_process_for_each_test() def test_model_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec): monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") - # Cold start in a forked child (must fork before CUDA init). - ctx = mp.get_context("fork") + # ROCm and XPU cannot fork safely after device discovery. + ctx = ( + mp.get_context("spawn") + if requires_spawn_multiprocessing() + else mp.get_context("fork") + ) p = ctx.Process(target=_cold_start_model, args=(vllm_runner, spec)) p.start() p.join() diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 772fee137aa6..c3a2919fac88 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -135,6 +135,7 @@ def ref_paged_attn( num_seqs = len(query_lens) block_tables = block_tables.cpu().numpy() _, block_size, num_kv_heads, head_size = key_cache.shape + head_size_v = value_cache.shape[-1] outputs: list[torch.Tensor] = [] start_idx = 0 @@ -149,7 +150,7 @@ def ref_paged_attn( k = key_cache[block_indices].view(-1, num_kv_heads, head_size) k = k[:kv_len] - v = value_cache[block_indices].view(-1, num_kv_heads, head_size) + v = value_cache[block_indices].view(-1, num_kv_heads, head_size_v) v = v[:kv_len] if q.shape[1] != k.shape[1]: diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 04b7755853eb..ea90b60a7514 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -5,6 +5,7 @@ import pytest import torch +from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn from vllm.platforms import current_platform from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import ( @@ -81,7 +82,7 @@ def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES) @torch.inference_mode() -def test_triton_unified_attn_diffkv_vs_fa( +def test_triton_unified_attn_diffkv_vs_reference( seq_lens: list[tuple[int, int]], num_heads: tuple[int, int], head_sizes: tuple[int, int], @@ -93,12 +94,13 @@ def test_triton_unified_attn_diffkv_vs_fa( ) -> None: head_size_qk, head_size_v = head_sizes - # DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference. - fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v) - if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4): - pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).") - - from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func + # Keep the FA3/FA4 comparison on NVIDIA; ROCm uses the PyTorch oracle. + if not current_platform.is_rocm(): + fa_version = get_flash_attn_version( + head_size=head_size_qk, head_size_v=head_size_v + ) + if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4): + pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).") torch.set_default_device(DEVICE_TYPE) set_random_seed(0) @@ -135,28 +137,44 @@ def test_triton_unified_attn_diffkv_vs_fa( 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 ) - # ---- FlashAttention DiffKV (ground truth) --------------------------- - # Mirror the backend: fix degenerate strides on size-1 dims so FA's - # TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1). - fa_k = canonicalize_singleton_dim_strides(key_cache) - fa_v = canonicalize_singleton_dim_strides(value_cache) - fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype) - flash_attn_varlen_func( - q=query, - k=fa_k, - v=fa_v, - out=fa_out, - cu_seqlens_q=cu_query_lens, - max_seqlen_q=max_query_len, - seqused_k=kv_lens_t, - max_seqlen_k=max_kv_len, - softmax_scale=scale, - causal=True, - window_size=list(window_size), - block_table=block_tables, - softcap=soft_cap if soft_cap is not None else 0, - fa_version=fa_version, - ) + if current_platform.is_rocm(): + # FP32 also keeps the helper's in-place scaling off the kernel input. + ref_out = ref_paged_attn( + query.float(), + key_cache.float(), + value_cache.float(), + query_lens, + kv_lens, + block_tables, + scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + ).to(dtype) + else: + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func + + # FA's TMA path needs aligned singleton strides (num_kv_heads == 1). + fa_k = canonicalize_singleton_dim_strides(key_cache) + fa_v = canonicalize_singleton_dim_strides(value_cache) + ref_out = torch.empty( + sum(query_lens), num_query_heads, head_size_v, dtype=dtype + ) + flash_attn_varlen_func( + q=query, + k=fa_k, + v=fa_v, + out=ref_out, + cu_seqlens_q=cu_query_lens, + max_seqlen_q=max_query_len, + seqused_k=kv_lens_t, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=list(window_size), + block_table=block_tables, + softcap=soft_cap if soft_cap is not None else 0, + fa_version=fa_version, + ) # ---- Triton DiffKV -------------------------------------------------- segm_output, segm_max, segm_expsum = _alloc_segm_buffers( @@ -183,7 +201,4 @@ def test_triton_unified_attn_diffkv_vs_fa( softmax_segm_expsum=segm_expsum, ) - ( - torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2), - f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}", - ) + torch.testing.assert_close(triton_out, ref_out, atol=2e-2, rtol=2e-2) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 68cf493b443a..234a63fb73e8 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -4,18 +4,23 @@ The decode path hands the kernel column slices of the merged ``q|k|v`` conv output and of the fused ``qkvbfg_a`` projection (beta), so q/k/v/beta are -token-strided rather than contiguous. The kernel must read them in place, -match a pure-PyTorch recurrence, and reject layouts it cannot address. +token-strided rather than contiguous. Both backends must match a pure-PyTorch +recurrence. CUDA reads these slices in place and rejects unsupported layouts; +the ROCm wrapper makes contiguous copies before launching its kernel. """ import pytest import torch -from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda from vllm.platforms import current_platform +if current_platform.is_rocm(): + from vllm.models.glm5next.amd.ops.third_party.kda import fused_recurrent_kda +else: + from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda + pytestmark = pytest.mark.skipif( - not current_platform.is_cuda(), reason="CUDA-only Triton kernel" + not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm" ) H, D = 16, 128 @@ -160,6 +165,10 @@ def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous( torch.testing.assert_close(state, state_ref, rtol=0, atol=0) +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Only the CUDA implementation requires zero-copy token-strided layouts", +) @torch.inference_mode() def test_fused_recurrent_kda_rejects_unaddressable_layouts(): """Layouts the token-stride addressing cannot express must fail loudly From 0aee727ff6131a1b647941186ceef8f4777bdac2 Mon Sep 17 00:00:00 2001 From: djramic Date: Mon, 21 Sep 2026 11:18:59 +0200 Subject: [PATCH 575/774] [CI][ROCm] Increase timeout for MI300 Distributed DP Extended (#57903) Signed-off-by: Djordje Ramic --- .buildkite/test-amd.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e0c081286608..3211840c9b06 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3207,7 +3207,7 @@ steps: - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) Distributed DP Extended" # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_4 From 0b7f11a1eece3dc9c689f1b5ab90d60c03b10793 Mon Sep 17 00:00:00 2001 From: Hongxin Xu <70438206+xhx1022@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:25:05 +0800 Subject: [PATCH 576/774] [AuxOutput] Add block-keyed storage for routed-expert outputs (#45635) Signed-off-by: aoshen02 Signed-off-by: Nick Hill Co-authored-by: aoshen02 Co-authored-by: OpenAI Codex Co-authored-by: Roger Wang Co-authored-by: Nick Hill Co-authored-by: Codex --- tests/config/test_aux_output_config.py | 119 ++ .../aux_output_connector/test_store.py | 1854 +++++++++++++++++ .../openai/test_return_routed_experts.py | 19 +- .../test_return_routed_experts.py | 16 +- .../test_routed_experts_capture.py | 412 ++-- tests/v1/core/test_async_scheduler.py | 3 +- tests/v1/core/test_scheduler.py | 94 +- .../v1/e2e/general/test_mamba_prefix_cache.py | 1 + tests/v1/engine/test_output_processor.py | 48 +- tests/v1/executor/test_ray_utils.py | 29 +- vllm/config/__init__.py | 3 + vllm/config/aux_output.py | 31 + vllm/config/model.py | 2 - vllm/config/vllm.py | 86 +- .../aux_output_connector/__init__.py | 2 + .../aux_output_connector/connector.py | 157 ++ .../aux_output_connector/routed_experts.py | 219 ++ .../distributed/aux_output_connector/store.py | 234 +++ .../aux_output_connector/worker.py | 412 ++++ vllm/engine/arg_utils.py | 18 +- .../fused_moe/routed_experts_capturer.py | 274 +-- vllm/outputs.py | 4 + vllm/v1/core/sched/output.py | 7 + vllm/v1/core/sched/scheduler.py | 131 +- vllm/v1/executor/ray_utils.py | 12 +- vllm/v1/outputs.py | 80 +- vllm/v1/worker/gpu/async_utils.py | 22 +- vllm/v1/worker/gpu/model_runner.py | 48 +- vllm/v1/worker/gpu/sample/batch_shard.py | 2 - vllm/v1/worker/gpu/sample/output.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 127 -- vllm/v1/worker/gpu_worker.py | 3 - 32 files changed, 3600 insertions(+), 873 deletions(-) create mode 100644 tests/config/test_aux_output_config.py create mode 100644 tests/distributed/aux_output_connector/test_store.py create mode 100644 vllm/config/aux_output.py create mode 100644 vllm/distributed/aux_output_connector/__init__.py create mode 100644 vllm/distributed/aux_output_connector/connector.py create mode 100644 vllm/distributed/aux_output_connector/routed_experts.py create mode 100644 vllm/distributed/aux_output_connector/store.py create mode 100644 vllm/distributed/aux_output_connector/worker.py diff --git a/tests/config/test_aux_output_config.py b/tests/config/test_aux_output_config.py new file mode 100644 index 000000000000..bf428ea2dd9e --- /dev/null +++ b/tests/config/test_aux_output_config.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm.config import AuxOutputConfig, VllmConfig +from vllm.engine.arg_utils import EngineArgs + +pytestmark = pytest.mark.cpu_test + + +def _config( + *, + use_v2: bool = True, + pp: int = 1, + dcp: int = 1, + pcp: int = 1, + connector: str | None = None, + runner_type: str = "generate", + is_moe: bool = True, + sliding_window: int | None = None, + attention_chunk_size: int | None = None, + enable_prefix_caching: bool = True, + adaptive_verification: bool = False, +): + return SimpleNamespace( + model_config=SimpleNamespace( + runner_type=runner_type, + is_moe=is_moe, + get_sliding_window=lambda: sliding_window, + attention_chunk_size=attention_chunk_size, + ), + use_v2_model_runner=use_v2, + parallel_config=SimpleNamespace( + pipeline_parallel_size=pp, + decode_context_parallel_size=dcp, + prefill_context_parallel_size=pcp, + ), + aux_output_config=AuxOutputConfig(enable_return_routed_experts=True), + cache_config=SimpleNamespace( + enable_prefix_caching=enable_prefix_caching, + ), + speculative_config=( + SimpleNamespace(enable_adaptive_verification=True) + if adaptive_verification + else None + ), + kv_transfer_config=( + None + if connector is None + else SimpleNamespace( + is_kv_transfer_instance=True, + kv_connector=connector, + ) + ), + ) + + +def test_aux_output_config_defaults(): + config = AuxOutputConfig() + + assert not config.enabled + assert not config.enable_return_routed_experts + assert config.max_bytes is None + + +def test_aux_output_capture_changes_compilation_hash(): + disabled = AuxOutputConfig() + enabled = AuxOutputConfig(enable_return_routed_experts=True) + + assert disabled.compute_hash() != enabled.compute_hash() + + +def test_legacy_routed_experts_flag_updates_aux_output_config(): + args = EngineArgs(enable_return_routed_experts=True) + + assert args.aux_output_config.enabled + assert args.aux_output_config.enable_return_routed_experts + + +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"use_v2": False}, "requires Model Runner V2"), + ({"runner_type": "pooling"}, "only supports generate runners"), + ({"is_moe": False}, "only supports MoE models"), + ({"enable_prefix_caching": False}, "requires prefix caching"), + ( + {"adaptive_verification": True}, + "adaptive speculative verification", + ), + ({"pp": 2}, "pipeline parallelism"), + ({"dcp": 2}, "context parallelism"), + ({"pcp": 2}, "context parallelism"), + ({"connector": "MooncakeConnector"}, "incompatible with KV connectors"), + ], +) +def test_aux_output_connector_rejects_unsupported_configuration(kwargs, error): + with pytest.raises(ValueError, match=error): + VllmConfig._verify_aux_output_compatibility(_config(**kwargs)) + + +@pytest.mark.parametrize( + "kwargs", + [{"sliding_window": 4096}, {"attention_chunk_size": 4096}], +) +def test_aux_output_config_defers_attention_layout_to_kv_config(kwargs): + VllmConfig._verify_aux_output_compatibility(_config(**kwargs)) + + +def test_aux_output_guards_are_inactive_when_capture_is_disabled(): + config: Any = VllmConfig.__new__(VllmConfig) + config.model_config = SimpleNamespace() + config.aux_output_config = AuxOutputConfig() + + config._verify_aux_output_compatibility() diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py new file mode 100644 index 000000000000..39277f3e7e0e --- /dev/null +++ b/tests/distributed/aux_output_connector/test_store.py @@ -0,0 +1,1854 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from contextlib import nullcontext +from dataclasses import dataclass, field +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from vllm.distributed.aux_output_connector.connector import ( + AuxOutputConnectorMetadata, + AuxOutputSchedulerConnector, + AuxRequestOutput, + PackedBlockHashes, +) +from vllm.distributed.aux_output_connector.routed_experts import ( + RoutedExpertsBuffer, + materialize_routed_experts, + publish_routed_experts, + routed_experts_keys, +) +from vllm.distributed.aux_output_connector.store import ( + BackgroundBlockObjectStore, + BlockObject, + BlockObjectStore, + BlockObjectStoreError, +) +from vllm.distributed.aux_output_connector.worker import ( + AuxOutputWorkerConnector, +) +from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.worker.gpu import async_utils +from vllm.v1.worker.gpu.sample.output import SamplerOutput + +pytestmark = pytest.mark.cpu_test + +_SHAPE = (3, 2) +_DTYPE = np.dtype("uint8") +_BLOCK_SIZE = 4 + + +def test_background_store_publishes_without_blocking_caller(): + started = threading.Event() + release = threading.Event() + underlying = Mock() + + def put(*_args, **_kwargs): + started.set() + release.wait() + + underlying.put.side_effect = put + store = BackgroundBlockObjectStore(underlying, max_pending_batches=1) + + store.put([BlockObject("key", b"value")]) + assert started.wait(timeout=1) + release.set() + store.close() + + underlying.close.assert_called_once_with() + + +def _make_connector(): + return AuxOutputSchedulerConnector() + + +@dataclass +class _Execution: + request_id: str + token_start: int + num_tokens: int + emit_start: int + block_hashes: list[bytes] | PackedBlockHashes + + +@dataclass +class _WorkerStep: + metadata: AuxOutputConnectorMetadata + executions: list[_Execution] + + +@dataclass +class _SchedulerRequest: + request_id: str + block_hashes: list[bytes] + num_tokens: int + num_output_tokens: int = 1 + num_computed_tokens: int = 0 + num_in_flight_tokens: int = 0 + finished: bool = False + sampling_params: SimpleNamespace = field( + default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0) + ) + + def is_finished(self) -> bool: + return self.finished + + +def _request_metadata( + request_id: str, + token_start: int, + num_tokens: int, + emit_start: int, + block_hashes, +) -> _Execution: + return _Execution(request_id, token_start, num_tokens, emit_start, block_hashes) + + +def _metadata( + generation: int, + requests: list[_Execution], + finished_requests: dict[str, list[bytes] | PackedBlockHashes], + hash_updates: dict[str, list[bytes] | PackedBlockHashes] | None = None, +) -> _WorkerStep: + block_hashes = { + request.request_id: request.block_hashes + for request in requests + if request.block_hashes + } + if hash_updates: + block_hashes.update( + (request_id, hashes) + for request_id, hashes in hash_updates.items() + if hashes + ) + block_hashes.update( + (request_id, hashes) + for request_id, hashes in finished_requests.items() + if hashes + ) + return _WorkerStep( + AuxOutputConnectorMetadata( + generation, + {request.request_id: request.emit_start for request in requests}, + block_hashes, + tuple(finished_requests), + ), + requests, + ) + + +def _execution_ranges(metadata, request_ids): + by_request = {request.request_id: request for request in metadata.executions} + token_starts = tuple( + by_request[request_id].token_start for request_id in request_ids + ) + num_tokens = tuple(by_request[request_id].num_tokens for request_id in request_ids) + return token_starts, num_tokens + + +def _begin_step(worker: AuxOutputWorkerConnector, step: _WorkerStep) -> None: + worker.begin_step(step.metadata) + + +def _make_worker( + max_num_seqs: int, + max_num_batched_tokens: int | None = None, + max_concurrent_batches: int = 2, + max_store_blocks: int | None = None, +) -> AuxOutputWorkerConnector: + object_nbytes = _BLOCK_SIZE * int(np.prod(_SHAPE)) + store = _make_store( + max_bytes=( + 1 << 20 if max_store_blocks is None else max_store_blocks * object_nbytes + ), + object_nbytes=object_nbytes, + ) + worker = object.__new__(AuxOutputWorkerConnector) + worker._store = store + worker._buffer = RoutedExpertsBuffer( + _DTYPE, + _SHAPE, + _BLOCK_SIZE, + max_num_seqs, + max_num_batched_tokens or max_num_seqs * _BLOCK_SIZE, + max_concurrent_batches, + ) + worker._requests = {} + worker._generation = 0 + worker._step_metadata = None + worker._pending_outputs = [] + worker._lock = threading.Lock() + worker._max_concurrent_batches = max_concurrent_batches + return worker + + +def test_worker_rejects_metadata_generation_rollback(): + worker = _make_worker(1) + _begin_step(worker, _metadata(1, [], {})) + + with pytest.raises(AssertionError, match="generation moved backwards"): + _begin_step(worker, _metadata(0, [], {})) + + worker.close() + + +def test_worker_rejects_run_and_finish_in_one_step(): + worker = _make_worker(1) + request = _request_metadata("request", 0, 1, 0, []) + metadata = _metadata( + 0, + [request], + {"request": []}, + ) + + with pytest.raises(AssertionError, match="cannot run and finish"): + _begin_step(worker, metadata) + + worker.close() + + +def _input_batch(request_ids, token_starts, query_start_loc): + return SimpleNamespace( + req_ids=request_ids, + num_computed_tokens_np=token_starts, + query_start_loc_np=query_start_loc, + ) + + +def test_non_output_rank_skips_capture_snapshot(): + worker = object.__new__(AuxOutputWorkerConnector) + worker._store = None + worker._buffer = None + worker._capturer = Mock() + worker._step_metadata = Mock() + assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None + worker._capturer.snapshot_routing_data.assert_not_called() + + +def test_worker_skips_aux_outputs_for_internal_warmup_step(): + worker = _make_worker(1) + worker._capturer = Mock() + + worker.begin_step(None) + + assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None + worker._capturer.snapshot_routing_data.assert_not_called() + worker.close() + + +def test_next_step_does_not_consume_pending_output(monkeypatch): + """begin_step must not wait for or consume an unconsumed step output.""" + event = Mock() + monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event) + monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext()) + monkeypatch.setattr( + async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy() + ) + worker = _make_worker(1) + worker.begin_step( + _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata + ) + rows = torch.ones((1, *_SHAPE), dtype=torch.uint8) + worker._capturer = Mock() + worker._capturer.snapshot_routing_data.return_value = rows + process_output = Mock(wraps=worker.process_output) + monkeypatch.setattr(worker, "process_output", process_output) + pending = worker.prepare_output( + _input_batch(["request"], np.array([0]), np.array([0, 1])) + ) + assert pending is not None + output = async_utils.AsyncOutput( + ModelRunnerOutput(["request"], {"request": 0}), + SamplerOutput( + torch.tensor([[7]]), None, None, None, num_rejected=torch.tensor([0]) + ), + torch.tensor([1]), + Mock(), + Mock(), + check_ep_fault=False, + pending_aux_output=pending, + ) + event.record.assert_called_once() + assert worker._pending_outputs == [pending] + + worker.begin_step(_metadata(0, [], {}).metadata) + process_output.assert_not_called() + assert worker._pending_outputs == [pending] + + result = output.get_output() + + assert result.sampled_token_ids == [[7]] + np.testing.assert_array_equal( + result.aux_output_connector_output["request"].rows, rows.numpy() + ) + process_output.assert_called_once() + assert worker._pending_outputs == [] + worker.close() + + +@pytest.mark.parametrize("num_pending", [1, 2]) +def test_finished_request_teardown_waits_for_pending_output(num_pending): + """Terminal cleanup waits for every outstanding output, not just the first.""" + worker = _make_worker(1) + rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, *_SHAPE) + worker._capturer = Mock() + worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows) + pending_outputs = [] + for i in range(num_pending): + request = _request_metadata("request", 4 * i, 4, 4 * i, [bytes([i]) * 4]) + worker.begin_step(_metadata(0, [request], {}).metadata) + pending = worker.prepare_output( + _input_batch(["request"], np.array([4 * i]), np.array([0, 4])) + ) + assert pending is not None + pending.enqueue_cpu_copy(np.array([1]), np.array([0])) + pending_outputs.append(pending) + + # The next step finishes the request while its output is unconsumed. + worker.begin_step(_metadata(0, [], {"request": []}).metadata) + assert "request" in worker._requests + + for i, pending in enumerate(pending_outputs): + output = worker.process_output(pending) + np.testing.assert_array_equal(output["request"].rows, rows) + if i + 1 < num_pending: + assert "request" in worker._requests + assert worker._store._references + assert "request" not in worker._requests + # The finished request's store keys must be released after commit. + assert not worker._store._references + worker.close() + + +def test_pending_outputs_bounded_by_max_concurrent_batches(): + worker = _make_worker(1, max_concurrent_batches=2) + worker._capturer = Mock() + worker._capturer.snapshot_routing_data.side_effect = lambda num_rows: torch.zeros( + (num_rows, *_SHAPE), dtype=torch.uint8 + ) + metadata = _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata + batch = _input_batch(["request"], np.array([0]), np.array([0, 1])) + for _ in range(2): + worker.begin_step(metadata) + assert worker.prepare_output(batch) is not None + + worker.begin_step(metadata) + with pytest.raises(AssertionError, match="not consumed in order"): + worker.prepare_output(batch) + + worker.close() + + +def test_worker_rejects_invalid_rejected_token_count(): + worker = _make_worker(1) + metadata = _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [])], + {}, + ) + with pytest.raises(AssertionError, match="rejected-token count is invalid"): + _process_output( + worker, + metadata, + np.zeros((1, *_SHAPE), dtype=_DTYPE), + ["request"], + np.array([-1]), + token_starts=(0,), + num_tokens=(1,), + ) + assert worker._pending_outputs == [] + + worker.close() + + +def _process_output( + worker, + step, + rows, + request_ids, + num_rejected, + num_sampled=None, + *, + token_starts=None, + num_tokens=None, + query_start_loc=None, +): + if num_sampled is None: + num_sampled = np.ones(len(request_ids), dtype=np.int32) + if token_starts is None or num_tokens is None: + assert isinstance(step, _WorkerStep) + token_starts, num_tokens = _execution_ranges(step, request_ids) + worker._capturer = Mock() + worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows) + metadata = step.metadata if isinstance(step, _WorkerStep) else step + worker.begin_step(metadata) + if query_start_loc is None: + query_start_loc = np.concatenate( + (np.zeros(1, dtype=np.int32), np.cumsum(num_tokens, dtype=np.int32)) + ) + pending = worker.prepare_output( + _input_batch(request_ids, np.asarray(token_starts), query_start_loc) + ) + assert pending is not None + pending.enqueue_cpu_copy(num_sampled, num_rejected) + return worker.process_output(pending) + + +def test_worker_ignores_cudagraph_query_padding(): + worker = _make_worker(2) + logical = np.arange(2 * 3 * 2, dtype=np.uint8).reshape(2, 3, 2) + metadata = _metadata( + 0, + [ + _request_metadata("first", 0, 1, 0, []), + _request_metadata("second", 0, 1, 0, []), + ], + {}, + ) + + output = _process_output( + worker, + metadata, + logical, + ["first", "second"], + np.zeros(2, dtype=np.int32), + query_start_loc=np.array([0, 1, 2, 2], dtype=np.int32), + ) + + np.testing.assert_array_equal(output["first"].rows, logical[:1]) + np.testing.assert_array_equal(output["second"].rows, logical[1:]) + worker.close() + + +def test_worker_rejects_scheduler_emit_cursor_ahead(): + worker = _make_worker(1) + _begin_step( + worker, + _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [])], + {}, + ), + ) + + with pytest.raises(AssertionError, match="Scheduler emit cursor moved ahead"): + _begin_step( + worker, + _metadata( + 0, + [_request_metadata("request", 1, 1, 1, [])], + {}, + ), + ) + + worker.close() + + +@pytest.mark.parametrize( + "routing", + [ + np.zeros((1, _SHAPE[0], _SHAPE[1] + 1), dtype=_DTYPE), + np.zeros((1, *_SHAPE), dtype=np.int32), + ], + ids=["shape", "dtype"], +) +def test_worker_rejects_mismatched_capture_profile(routing): + worker = _make_worker(1) + metadata = _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [])], + {}, + ) + + with pytest.raises(AssertionError, match="capture profile changed"): + _process_output( + worker, + metadata, + routing, + ["request"], + np.array([0]), + ) + worker.close() + + +@pytest.mark.parametrize( + ("initial_start", "initial_rows", "invalid_start"), + [ + (4, [4, 5, 6], 6), + (0, [0], 2), + (None, [], 3), + ], + ids=["overlap", "gap", "unaligned-initial-capture"], +) +def test_logical_buffer_rejects_noncontiguous_capture( + initial_start, initial_rows, invalid_start +): + buffer = RoutedExpertsBuffer(np.dtype("uint8"), (1,), 4, 1, 4, 1) + if initial_start is not None: + rows = np.asarray(initial_rows, dtype=np.uint8).reshape(-1, 1) + assert buffer.capture("request", initial_start, rows) == [] + + with pytest.raises(AssertionError, match="not contiguous"): + buffer.capture( + "request", invalid_start, np.array([[invalid_start]], dtype=np.uint8) + ) + + +def test_logical_buffer_captures_one_row_per_decode_step(): + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2) + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, 3, 2) + + completed = [] + for step in range(_BLOCK_SIZE): + completed += buffer.capture("request", step, logical[step : step + 1]) + + assert len(completed) == 1 + np.testing.assert_array_equal(completed[0][1], logical) + + +def test_logical_buffer_copies_borrowed_block_when_retained(): + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 4, 2) + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, *_SHAPE) + + expected = logical.copy() + completed = buffer.capture("request", 0, logical) + retained = buffer.retain_block(completed[0][1]) + logical[:] = 0 + + assert not np.shares_memory(retained, logical) + np.testing.assert_array_equal(retained, expected) + buffer.release_block(retained) + + +def _make_store( + *, + max_bytes: int = 1 << 20, + object_nbytes: int = 4, +): + return BlockObjectStore( + max_bytes=max_bytes, + object_nbytes=object_nbytes, + ) + + +def test_publish_routed_experts_publishes_full_blocks(): + store = _make_store(object_nbytes=_BLOCK_SIZE * int(np.prod(_SHAPE))) + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2) + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2) + hashes = [b"a" * 32, b"b" * 32] + keys = routed_experts_keys(hashes, "0") + blocks = buffer.capture("request", 0, logical) + publish_routed_experts( + store, + batches=[(keys, blocks)], + block_size=_BLOCK_SIZE, + ) + expected = logical.copy() + logical[:] = 0 + + np.testing.assert_array_equal( + materialize_routed_experts( + store, + keys, + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + expected, + ) + store.close() + + +def test_worker_data_plane_publishes_blocks_and_reuses_prefix(): + worker = _make_worker(2) + + hashes = [b"a" * 32, b"b" * 32, b"c" * 32] + logical = np.arange(10 * 3 * 2, dtype=np.uint8).reshape(10, 3, 2) + first = _metadata( + generation=0, + requests=[_request_metadata("first", 0, 8, 0, hashes)], + finished_requests={}, + ) + output = _process_output(worker, first, logical[:8], ["first"], np.array([0])) + assert output is not None + np.testing.assert_array_equal(output["first"].rows, logical[:8]) + + second = _metadata( + generation=0, + requests=[_request_metadata("second", 8, 2, 0, hashes)], + finished_requests={"first": []}, + ) + output = _process_output(worker, second, logical[8:], ["second"], np.array([0])) + assert output is not None + np.testing.assert_array_equal(output["second"].rows, logical) + worker.close() + + +def test_worker_matches_kv_eviction_order_after_requests_finish(): + worker = _make_worker(2, max_store_blocks=3) + hashes = [bytes([value]) * 32 for value in range(5)] + blocks = [ + np.full((_BLOCK_SIZE, *_SHAPE), value, dtype=_DTYPE) for value in range(5) + ] + + def publish_and_finish(request_id, block_indexes): + rows = np.concatenate([blocks[index] for index in block_indexes]) + step = _metadata( + 0, + [ + _request_metadata( + request_id, + 0, + len(rows), + 0, + [hashes[index] for index in block_indexes], + ) + ], + {}, + ) + _process_output(worker, step, rows, [request_id], np.array([0])) + _begin_step(worker, _metadata(0, [], {request_id: []})) + + publish_and_finish("s", [0]) + publish_and_finish("ab", [1, 2]) + publish_and_finish("c", [3]) + publish_and_finish("d", [4]) + + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([hashes[1], hashes[3], hashes[4]], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + np.concatenate([blocks[1], blocks[3], blocks[4]]), + ) + with pytest.raises(BlockObjectStoreError, match="does not exist"): + materialize_routed_experts( + worker._store, + routed_experts_keys([hashes[2]], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ) + worker.close() + + +def test_worker_rejects_unbacked_capture_gap(): + worker = _make_worker(1) + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + first = _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [])], + {}, + ) + _process_output( + worker, + first, + logical[:1], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + second = _metadata( + 0, + [_request_metadata("request", 2, 1, 0, [])], + {}, + ) + + with pytest.raises(AssertionError, match="unbacked token gap"): + _process_output( + worker, + second, + logical[2:3], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + worker.close() + + +def test_worker_replays_rejected_speculative_gap(): + worker = _make_worker(1) + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2) + first = _metadata( + 0, + [_request_metadata("request", 0, 5, 0, [])], + {}, + ) + first_output = _process_output( + worker, first, logical[:5], ["request"], np.array([4]) + ) + + queued = _metadata( + 0, + [_request_metadata("request", 5, 4, 0, [])], + {}, + ) + queued_output = _process_output( + worker, + queued, + logical[1:], + ["request"], + np.array([3]), + ) + + rolled_back = _metadata( + 0, + [_request_metadata("request", 6, 3, 0, [b"a" * 32])], + {}, + ) + + rolled_back_output = _process_output( + worker, + rolled_back, + logical[2:], + ["request"], + np.array([0]), + ) + + np.testing.assert_array_equal(first_output["request"].rows, logical[:1]) + assert queued_output["request"].token_start == 1 + np.testing.assert_array_equal( + queued_output["request"].rows, + logical[1:2], + ) + np.testing.assert_array_equal( + rolled_back_output["request"].rows, + logical[2:], + ) + assert worker._requests["request"].capture_cursor == len(logical) + worker.close() + + +def test_worker_rejects_overlapping_capture_rows(): + worker = _make_worker(1) + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + first = _metadata( + 0, + [_request_metadata("request", 0, 3, 0, [])], + {}, + ) + _process_output( + worker, + first, + logical[:3], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + stale_overlap = logical[2:].copy() + stale_overlap[0] += 100 + second = _metadata( + 0, + [_request_metadata("request", 2, 2, 0, [b"a" * 32])], + {}, + ) + with pytest.raises(AssertionError, match="capture moved backwards"): + _process_output( + worker, + second, + stale_overlap, + ["request"], + np.array([0]), + ) + worker.close() + + +def test_worker_publishes_entire_batch_before_materializing_prefix(): + worker = _make_worker(2) + hashes = [b"a" * 32, b"b" * 32] + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2) + metadata = _metadata( + generation=0, + requests=[ + _request_metadata("consumer", 4, 4, 0, hashes), + _request_metadata("producer", 0, 4, 0, hashes), + ], + finished_requests={}, + ) + + output = _process_output( + worker, + metadata, + np.concatenate((logical[4:], logical[:4])), + ["consumer", "producer"], + np.array([0, 0]), + ) + + assert output is not None + np.testing.assert_array_equal(output["consumer"].rows, logical) + np.testing.assert_array_equal(output["producer"].rows, logical[:4]) + worker.close() + + +def test_worker_defers_full_block_until_kv_hash_arrives(): + worker = _make_worker(1) + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2) + first = _metadata( + 0, + [_request_metadata("request", 0, 4, 0, [])], + {}, + ) + + output = _process_output( + worker, + first, + logical[:4], + ["request"], + np.array([0]), + ) + + assert output is not None + np.testing.assert_array_equal(output["request"].rows, logical[:4]) + assert worker._requests["request"].pending_blocks[0][0] == 0 + + block_hash = b"a" * 32 + second = _metadata( + 0, + [_request_metadata("request", 4, 1, 4, [block_hash])], + {}, + ) + output = _process_output(worker, second, logical[4:], ["request"], np.array([0])) + + assert output is not None + np.testing.assert_array_equal(output["request"].rows, logical[4:]) + assert not worker._requests["request"].pending_blocks + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([block_hash], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical[:4], + ) + worker.close() + + +def test_worker_releases_newly_keyed_blocks_before_capture(): + worker = _make_worker(1, max_num_batched_tokens=16) + connector = _make_connector() + block_hashes = [b"a" * 32] + request = _scheduler_request("request", block_hashes, num_tokens=100) + logical = np.arange(52 * 3 * 2, dtype=np.uint8).reshape(52, 3, 2) + + for token_start, num_tokens in [(0, 4), (4, 16), (20, 16)]: + request.num_computed_tokens = token_start + metadata = connector.build_connector_meta( + _step_output([request.request_id], [token_start], [num_tokens]), + {request.request_id: request}, + ) + _process_output( + worker, + metadata, + logical[token_start : token_start + num_tokens], + [request.request_id], + np.array([0]), + token_starts=(token_start,), + num_tokens=(num_tokens,), + ) + + block_hashes.extend(bytes([value]) * 32 for value in range(1, 5)) + request.num_computed_tokens = 36 + metadata = connector.build_connector_meta( + _step_output([request.request_id], [36], [16]), + {request.request_id: request}, + ) + _process_output( + worker, + metadata, + logical[36:52], + [request.request_id], + np.array([0]), + token_starts=(36,), + num_tokens=(16,), + ) + + state = worker._requests[request.request_id] + assert [start for start, _ in state.pending_blocks] == [ + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + ] + worker.close() + + +def test_worker_publishes_pending_block_without_forward(): + worker = _make_worker(1) + _begin_step( + worker, + _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [])], + {}, + ), + ) + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2) + assert worker._buffer is not None + worker._requests["request"].pending_blocks = [ + (0, worker._buffer.retain_block(logical)) + ] + block_hash = b"a" * 32 + + _begin_step(worker, _metadata(0, [], {}, {"request": [block_hash]})) + + assert not worker._requests["request"].pending_blocks + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([block_hash], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical, + ) + worker.close() + + +def test_worker_releases_unscheduled_request_blocks_before_capture(): + worker = _make_worker(2, 20, max_concurrent_batches=1) + logical = np.arange(48 * 3 * 2, dtype=np.uint8).reshape(48, 3, 2) + prompt_hashes = [b"a" * 32, b"b" * 32] + prompt = _metadata( + 0, + [ + _request_metadata("first", 0, 4, 0, prompt_hashes[:1]), + _request_metadata("second", 0, 4, 0, prompt_hashes[1:]), + ], + {}, + ) + _process_output( + worker, + prompt, + logical[:8], + ["first", "second"], + np.zeros(2), + np.zeros(2, dtype=np.int32), + ) + + first_decode = _metadata( + 0, + [_request_metadata("first", 4, 20, 0, [])], + {}, + ) + _process_output( + worker, + first_decode, + logical[8:28], + ["first"], + np.zeros(1), + np.zeros(1, dtype=np.int32), + ) + assert len(worker._requests["first"].pending_blocks) == 5 + + first_hashes = [bytes([value]) * 32 for value in range(2, 7)] + second_decode = _metadata( + 0, + [_request_metadata("second", 4, 20, 0, [])], + {}, + {"first": first_hashes}, + ) + _process_output( + worker, + second_decode, + logical[28:48], + ["second"], + np.zeros(1), + np.zeros(1, dtype=np.int32), + ) + + assert not worker._requests["first"].pending_blocks + assert len(worker._requests["second"].pending_blocks) == 5 + worker.close() + + +def test_worker_does_not_rematerialize_emitted_rows(): + worker = _make_worker(1) + hashes = [b"a" * 32, b"b" * 32] + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2) + + first = _metadata( + 0, + [_request_metadata("request", 0, 4, 0, hashes)], + {}, + ) + output = _process_output(worker, first, logical[:4], ["request"], np.array([0])) + assert output is not None + + assert worker._store is not None + worker._store.get_concatenated = Mock(wraps=worker._store.get_concatenated) + second = _metadata( + 0, + # Async scheduling may build this before the scheduler consumes first. + [_request_metadata("request", 4, 4, 0, [])], + {}, + ) + output = _process_output(worker, second, logical[4:], ["request"], np.array([0])) + + assert output is not None + np.testing.assert_array_equal(output["request"].rows, logical[4:]) + worker._store.get_concatenated.assert_not_called() + worker.close() + + +def test_worker_resumes_from_scheduler_emit_boundary(): + worker = _make_worker(1) + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2) + resumed = _metadata( + 0, + [_request_metadata("request", 4, 1, 4, [b"a" * 32])], + {}, + ) + + output = _process_output(worker, resumed, logical[4:], ["request"], np.array([0])) + + assert output["request"].token_start == 4 + np.testing.assert_array_equal(output["request"].rows, logical[4:]) + assert worker._requests["request"].emit_cursor == 5 + worker.close() + + +def test_worker_emits_mid_block_chunked_prefill(): + worker = _make_worker(1) + logical = np.arange(3 * 3 * 2, dtype=np.uint8).reshape(3, 3, 2) + + first = _metadata( + 0, + [_request_metadata("request", 0, 2, 0, [])], + {}, + ) + output = _process_output( + worker, + first, + logical[:2], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + assert output is not None and not output + + second = _metadata( + 0, + [_request_metadata("request", 2, 1, 0, [])], + {}, + ) + output = _process_output(worker, second, logical[2:], ["request"], np.array([0])) + + np.testing.assert_array_equal(output["request"].rows, logical) + worker.close() + + +def test_worker_emits_published_block_and_mid_block_tail(): + worker = _make_worker(1) + logical = np.arange(7 * 3 * 2, dtype=np.uint8).reshape(7, 3, 2) + + first = _metadata( + 0, + [_request_metadata("request", 0, 6, 0, [b"a" * 32])], + {}, + ) + output = _process_output( + worker, + first, + logical[:6], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + assert output is not None and not output + + second = _metadata( + 0, + [_request_metadata("request", 6, 1, 0, [])], + {}, + ) + output = _process_output(worker, second, logical[6:], ["request"], np.array([0])) + + np.testing.assert_array_equal(output["request"].rows, logical) + worker.close() + + +def test_worker_output_does_not_alias_released_tail_buffer(): + worker = _make_worker(1) + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + first = _metadata( + 0, + [_request_metadata("request", 0, 2, 0, [b"a" * 32])], + {}, + ) + _process_output( + worker, + first, + logical[:2], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + second = _metadata( + 0, + [_request_metadata("request", 2, 2, 0, [])], + {}, + ) + + output = _process_output(worker, second, logical[2:], ["request"], np.array([0])) + emitted = output["request"].rows + expected = logical.copy() + + replacement = np.full_like(logical, 99) + reuse = _metadata( + 0, + [_request_metadata("reuse", 0, 4, 0, [b"b" * 32])], + {}, + ) + _process_output( + worker, + reuse, + replacement, + ["reuse"], + np.array([0]), + np.zeros(1, dtype=np.int32), + ) + + np.testing.assert_array_equal(emitted, expected) + worker.close() + + +def test_worker_merges_block_hash_deltas(): + worker = _make_worker(1) + _begin_step( + worker, + _metadata( + 0, + [_request_metadata("request", 0, 1, 0, [b"a" * 32, b"b" * 32])], + {}, + ), + ) + _begin_step( + worker, + _metadata( + 0, + [ + _request_metadata( + "request", + 1, + 1, + 0, + [b"c" * 32], + ) + ], + {}, + ), + ) + + assert worker._requests["request"].aux_output_keys == [ + "vllm-artifact/0/" + (b"a" * 32).hex(), + "vllm-artifact/0/" + (b"b" * 32).hex(), + "vllm-artifact/0/" + (b"c" * 32).hex(), + ] + worker.close() + + +def test_worker_finish_publishes_keyed_block_and_discards_request_state(): + worker = _make_worker(1) + logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape( + 2 * _BLOCK_SIZE, 3, 2 + ) + running = _metadata( + 0, + [_request_metadata("request", 0, len(logical), 0, [])], + {}, + ) + _process_output( + worker, + running, + logical, + ["request"], + np.array([0]), + np.array([0]), + ) + block_hash = b"a" * 32 + + _begin_step( + worker, + _metadata( + 0, + [], + {"request": [block_hash]}, + ), + ) + + assert "request" not in worker._requests + assert worker._buffer is not None + assert "request" not in worker._buffer._requests + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([block_hash], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical[:_BLOCK_SIZE], + ) + worker.close() + + +def test_worker_fails_when_cached_aux_output_is_missing(): + worker = _make_worker(1) + worker._generation = 0 + metadata = _metadata( + 0, + [ + _request_metadata( + "request", + 8, + 1, + 0, + [b"a" * 32, b"b" * 32], + ) + ], + {}, + ) + + with pytest.raises(BlockObjectStoreError, match="does not exist"): + _process_output( + worker, + metadata, + np.zeros((1, *_SHAPE), dtype=_DTYPE), + ["request"], + np.array([0]), + ) + worker.close() + + +def test_store_rejects_oversized_batch_without_partial_write(): + store = _make_store(max_bytes=6, object_nbytes=3) + store.put([BlockObject("retained", b"rrr")]) + + with pytest.raises(BlockObjectStoreError, match="cannot retain"): + store.put( + [ + BlockObject("first", b"111"), + BlockObject("second", b"222"), + BlockObject("third", b"333"), + ] + ) + + assert store.get_concatenated(["retained"]) == b"rrr" + with pytest.raises(BlockObjectStoreError, match="does not exist"): + store.get_concatenated(["first"]) + store.close() + + +def test_store_lru_and_immutable_put(): + store = _make_store(max_bytes=8) + store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")]) + assert store.get_concatenated(["first"]) == b"1111" + + store.put([BlockObject("first", b"1111"), BlockObject("third", b"3333")]) + + assert store.get_concatenated(["first", "third"]) == b"11113333" + with pytest.raises(BlockObjectStoreError, match="Increase aux_output_config"): + store.get_concatenated(["second"]) + store.close() + + +def test_store_does_not_evict_referenced_aux_outputs(): + store = _make_store(max_bytes=8) + store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")]) + store.put([], retain_keys=["first"]) + store.put([], retain_keys=["first"], release_keys=["first"]) + + store.put([BlockObject("third", b"3333")]) + + assert store.get_concatenated(["first", "third"]) == b"11113333" + with pytest.raises(BlockObjectStoreError, match="does not exist"): + store.get_concatenated(["second"]) + store.put([], release_keys=["first"]) + store.close() + + +def test_store_terminal_put_preserves_tail_first_eviction_order(): + store = _make_store(max_bytes=8) + keys = ["first", "second", "third"] + store.put( + [BlockObject("first", b"1111"), BlockObject("second", b"2222")], + retain_keys=keys, + ) + + store.put( + [BlockObject("third", b"3333")], + release_keys=reversed(keys), + ) + + assert store.get_concatenated(keys[:2]) == b"11112222" + with pytest.raises(BlockObjectStoreError, match="does not exist"): + store.get_concatenated(["third"]) + store.close() + + +def test_store_terminal_put_orders_survivors_for_later_eviction(): + store = _make_store(max_bytes=12) + keys = ["first", "second", "third"] + store.put( + [BlockObject("first", b"1111"), BlockObject("second", b"2222")], + retain_keys=keys, + ) + + store.put( + [BlockObject("third", b"3333")], + release_keys=reversed(keys), + ) + store.put([BlockObject("fourth", b"4444")]) + + assert store.get_concatenated(["first", "second", "fourth"]) == b"111122224444" + with pytest.raises(BlockObjectStoreError, match="does not exist"): + store.get_concatenated(["third"]) + store.close() + + +def test_store_terminal_put_keeps_a_shared_incoming_key(): + store = _make_store(max_bytes=8) + keys = ["first", "second", "third"] + store.put( + [BlockObject("first", b"1111"), BlockObject("second", b"2222")], + retain_keys=(*keys, "third"), + ) + + store.put( + [BlockObject("third", b"3333")], + release_keys=reversed(keys), + ) + + assert store.get_concatenated(["first", "third"]) == b"11113333" + with pytest.raises(BlockObjectStoreError, match="does not exist"): + store.get_concatenated(["second"]) + store.put([], release_keys=["third"]) + store.close() + + +def test_store_rejects_access_after_close(): + store = _make_store(max_bytes=8) + store.close() + + with pytest.raises(RuntimeError, match="closed"): + store.put([BlockObject("first", b"1111")]) + with pytest.raises(RuntimeError, match="closed"): + store.get_concatenated(["first"]) + + +def test_store_reuses_evicted_slot_without_moving_live_objects(): + store = _make_store(max_bytes=12) + store.put( + [ + BlockObject("first", b"1111"), + BlockObject("second", b"2222"), + BlockObject("third", b"3333"), + ] + ) + assert store.get_concatenated(["second"]) == b"2222" + store.put([BlockObject("fourth", b"4444")]) + + assert store.get_concatenated(["second", "third", "fourth"]) == b"222233334444" + store.close() + + +def _scheduler_request( + request_id: str, + block_hashes: list[bytes], + *, + num_tokens: int = 10, + num_output_tokens: int = 1, + prompt_start: int = 0, +): + return _SchedulerRequest( + request_id=request_id, + block_hashes=block_hashes, + num_tokens=num_tokens, + num_output_tokens=num_output_tokens, + num_computed_tokens=num_tokens, + sampling_params=SimpleNamespace(routed_experts_prompt_start=prompt_start), + ) + + +def _step_output( + request_ids: list[str], token_starts: list[int], token_counts: list[int] +) -> SchedulerOutput: + output = SchedulerOutput.make_empty() + output.scheduled_cached_reqs = CachedRequestData( + req_ids=request_ids, + resumed_req_ids=set(), + new_token_ids=[[] for _ in request_ids], + all_token_ids={}, + new_block_ids=[None for _ in request_ids], + num_computed_tokens=token_starts, + num_output_tokens=[0 for _ in request_ids], + ) + output.num_scheduled_tokens = dict(zip(request_ids, token_counts, strict=True)) + output.total_num_scheduled_tokens = sum(token_counts) + return output + + +def test_scheduler_connector_builds_worker_metadata_and_forwards_output(): + connector = _make_connector() + request = _scheduler_request("request", [b"a" * 32], num_tokens=5) + request.num_computed_tokens = 0 + scheduler_output = _step_output([request.request_id], [0], [4]) + + metadata = connector.build_connector_meta( + scheduler_output, {request.request_id: request} + ) + + assert metadata.generation == 0 + assert metadata.requests == {request.request_id: 4} + assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32] + + routing = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + output = {"request": AuxRequestOutput(0, routing)} + request.num_computed_tokens = 4 + np.testing.assert_array_equal(connector.take_output(request, output), routing) + + +def test_scheduler_starts_worker_output_at_requested_prompt_token(): + connector = _make_connector() + request = _scheduler_request( + "request", + [b"a" * 32], + num_tokens=4, + num_output_tokens=0, + prompt_start=2, + ) + + metadata = connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + + assert metadata.requests == {request.request_id: 2} + + +def test_scheduler_connector_preserves_request_finish_order(): + connector = _make_connector() + first = _scheduler_request("first", [b"a" * 32]) + second = _scheduler_request("second", [b"b" * 32]) + connector.build_connector_meta( + _step_output([first.request_id, second.request_id], [0, 0], [0, 0]), + {first.request_id: first, second.request_id: second}, + ) + connector.request_finished(second) + connector.request_finished(first) + + metadata = connector.build_connector_meta(_step_output([], [], []), {}) + + assert metadata.finished_requests == (second.request_id, first.request_id) + assert not metadata.block_hashes + + +def test_worker_preemption_terminal_keeps_inflight_hashes_for_publish(): + worker = _make_worker(1) + block_hashes = [b"a" * 32, b"b" * 32] + logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape( + 2 * _BLOCK_SIZE, 3, 2 + ) + first = _metadata( + 0, + [ + _request_metadata( + "request", + 0, + _BLOCK_SIZE, + 0, + block_hashes[:1], + ) + ], + {}, + ) + _process_output( + worker, + first, + logical[:_BLOCK_SIZE], + ["request"], + np.array([0]), + ) + inflight = _metadata( + 0, + [ + _request_metadata( + "request", + _BLOCK_SIZE, + _BLOCK_SIZE, + 0, + block_hashes[1:], + ) + ], + {}, + ) + _process_output( + worker, + inflight, + logical[_BLOCK_SIZE:], + ["request"], + np.array([0]), + np.zeros(1, dtype=np.int32), + token_starts=(_BLOCK_SIZE,), + num_tokens=(_BLOCK_SIZE,), + ) + _begin_step( + worker, + _metadata( + 0, + [], + {"request": []}, + ), + ) + + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys(block_hashes[:1], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical[:_BLOCK_SIZE], + ) + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys(block_hashes[1:], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical[_BLOCK_SIZE:], + ) + worker.close() + + +def test_worker_releases_terminal_pins_before_same_step_publish(): + worker = _make_worker(2, max_store_blocks=1) + victim_hash = b"a" * 32 + producer_hash = b"b" * 32 + victim_rows = np.zeros((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE) + producer_rows = np.ones((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE) + running = _metadata( + 0, + [ + _request_metadata("victim", 0, _BLOCK_SIZE, 0, [victim_hash]), + _request_metadata("producer", 0, _BLOCK_SIZE, 0, []), + ], + {}, + ) + _process_output( + worker, + running, + np.concatenate((victim_rows, producer_rows)), + ["victim", "producer"], + np.zeros(2, dtype=np.int32), + np.zeros(2, dtype=np.int32), + ) + + _begin_step( + worker, + _metadata( + 0, + [], + {"victim": []}, + {"producer": [producer_hash]}, + ), + ) + + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([producer_hash], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + producer_rows, + ) + with pytest.raises(BlockObjectStoreError, match="does not exist"): + worker._store.get_concatenated(routed_experts_keys([victim_hash], "0")) + worker.close() + + +def test_terminal_request_publishes_hash_discovered_after_last_schedule(): + connector = _make_connector() + worker = _make_worker(1) + block_hashes: list[bytes] = [] + request = _scheduler_request("request", block_hashes, num_tokens=5) + request.num_computed_tokens = 0 + metadata = connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2) + + _process_output( + worker, + metadata, + logical, + [request.request_id], + np.array([0]), + token_starts=(0,), + num_tokens=(_BLOCK_SIZE,), + ) + assert worker._requests[request.request_id].pending_blocks + + block_hash = b"a" * 32 + block_hashes.append(block_hash) + request.num_computed_tokens = _BLOCK_SIZE + connector.request_finished(request) + cleanup = connector.build_connector_meta(_step_output([], [], []), {}) + worker.begin_step(cleanup) + + assert worker._store is not None + np.testing.assert_array_equal( + materialize_routed_experts( + worker._store, + routed_experts_keys([block_hash], "0"), + shape_per_token=_SHAPE, + dtype=_DTYPE, + ), + logical, + ) + worker.close() + + +def test_scheduler_connector_recreates_preempted_request_state(): + connector = _make_connector() + request = _scheduler_request("request", [b"a" * 32], num_tokens=5) + connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + connector.request_finished(request) + cleanup = connector.build_connector_meta( + _step_output([], [], []), + {request.request_id: request}, + ) + assert request.request_id in cleanup.finished_requests + assert request.request_id not in cleanup.block_hashes + + resumed = connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + assert list(resumed.block_hashes[request.request_id]) == [b"a" * 32] + + +def test_scheduler_consumes_ordered_stale_aux_outputs(): + connector = _make_connector() + request = _scheduler_request("request", [], num_tokens=6) + connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + first_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + second_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + 40 + first = {"request": AuxRequestOutput(0, first_rows)} + second = {"request": AuxRequestOutput(4, second_rows)} + connector.request_finished(request) + + request.num_tokens = 5 + np.testing.assert_array_equal(connector.take_output(request, first), first_rows) + request.num_tokens = 6 + np.testing.assert_array_equal( + connector.take_output(request, second), second_rows[:1] + ) + + +def test_scheduler_rejects_stale_output_without_aux_outputs(): + connector = _make_connector() + request = _scheduler_request("request", [], num_tokens=1) + connector.build_connector_meta( + _step_output([request.request_id], [0], [1]), + {request.request_id: request}, + ) + + with pytest.raises( + AssertionError, match="auxiliary output worker output is missing" + ): + connector.take_output( + request, + {}, + ) + + +def test_scheduler_rejects_missing_accepted_aux_output_rows(): + connector = _make_connector() + request = _scheduler_request("request", [], num_tokens=2) + connector.build_connector_meta( + _step_output([request.request_id], [0], [1]), + {request.request_id: request}, + ) + output = { + "request": AuxRequestOutput( + 0, + np.empty((0, *_SHAPE), dtype=_DTYPE), + ) + } + + with pytest.raises(AssertionError, match="invalid token range"): + connector.take_output(request, output) + + +def test_scheduler_rejects_aux_output_past_finished_request(): + connector = _make_connector() + request = _scheduler_request("request", [], num_tokens=1) + request.finished = True + connector.build_connector_meta( + _step_output([request.request_id], [0], [1]), + {request.request_id: request}, + ) + output = { + "request": AuxRequestOutput( + 1, + np.empty((0, *_SHAPE), dtype=_DTYPE), + ) + } + + with pytest.raises( + AssertionError, match="finished auxiliary output output has no accepted" + ): + connector.take_output(request, output) + + +@pytest.mark.parametrize("chunk_size", [2, 4]) +@pytest.mark.parametrize("max_tokens", [1, 2]) +def test_prompt_end_offset_through_worker_and_scheduler(chunk_size, max_tokens): + """Omit prompt R3 without losing the empty first output or later decode rows.""" + connector = _make_connector() + worker = _make_worker(1) + request = _scheduler_request( + "request", [b"a" * 32], num_tokens=4, num_output_tokens=0, prompt_start=4 + ) + rows = np.arange(5 * 3 * 2, dtype=_DTYPE).reshape(5, *_SHAPE) + try: + for start in range(0, 4, chunk_size): + metadata = connector.build_connector_meta( + _step_output([request.request_id], [start], [chunk_size]), + {request.request_id: request}, + ) + output = _process_output( + worker, + metadata, + rows[start : start + chunk_size], + [request.request_id], + np.array([0]), + num_sampled=np.array([int(start + chunk_size == 4)]), + token_starts=(start,), + num_tokens=(chunk_size,), + ) + if start + chunk_size < 4: + assert output == {} + request.num_tokens = 5 + request.num_output_tokens = 1 + request.finished = max_tokens == 1 + result = connector.take_output(request, output) + np.testing.assert_array_equal(result, rows[:0]) + + if max_tokens == 2: + metadata = connector.build_connector_meta( + _step_output([request.request_id], [4], [1]), + {request.request_id: request}, + ) + output = _process_output( + worker, + metadata, + rows[4:], + [request.request_id], + np.array([0]), + token_starts=(4,), + num_tokens=(1,), + ) + request.num_tokens = 6 + request.finished = True + np.testing.assert_array_equal( + connector.take_output(request, output), rows[4:] + ) + finally: + worker.close() + + +def test_scheduler_connector_sends_each_block_hash_once(): + connector = _make_connector() + request = _scheduler_request("request", [b"a" * 32, b"b" * 32], num_tokens=12) + scheduler_output = _step_output([request.request_id], [0], [4]) + + first = connector.build_connector_meta( + scheduler_output, {request.request_id: request} + ) + request.block_hashes.append(b"c" * 32) + second = connector.build_connector_meta( + _step_output([request.request_id], [4], [4]), + {request.request_id: request}, + ) + third = connector.build_connector_meta( + _step_output([request.request_id], [8], [4]), + {request.request_id: request}, + ) + + assert list(first.block_hashes[request.request_id]) == [b"a" * 32, b"b" * 32] + assert list(second.block_hashes[request.request_id]) == [b"c" * 32] + assert request.request_id not in third.block_hashes + + +def test_scheduler_connector_sends_unscheduled_hash_update(): + connector = _make_connector() + request = _scheduler_request("request", [b"a" * 32], num_tokens=4) + connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + + request.block_hashes.append(b"b" * 32) + request.num_tokens = 9 + request.num_computed_tokens = 8 + metadata = connector.build_connector_meta( + _step_output([], [], []), + {request.request_id: request}, + ) + + assert not metadata.requests + assert list(metadata.block_hashes[request.request_id]) == [b"b" * 32] + + +def test_scheduler_connector_reset_derives_emit_start_from_request(): + connector = _make_connector() + request = _scheduler_request("request", [b"a" * 32], num_tokens=5) + request.num_computed_tokens = 0 + before_reset = connector.build_connector_meta( + _step_output([request.request_id], [0], [4]), + {request.request_id: request}, + ) + assert list(before_reset.block_hashes[request.request_id]) == [b"a" * 32] + + connector.reset() + empty = connector.build_connector_meta( + _step_output([], [], []), + {request.request_id: request}, + ) + assert not empty.block_hashes + request.num_computed_tokens = 4 + scheduler_output = _step_output([request.request_id], [4], [1]) + metadata = connector.build_connector_meta( + scheduler_output, {request.request_id: request} + ) + + assert metadata.generation == 1 + assert metadata.requests[request.request_id] == 4 + assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32] diff --git a/tests/entrypoints/openai/test_return_routed_experts.py b/tests/entrypoints/openai/test_return_routed_experts.py index a64e80c39868..46eeed179531 100644 --- a/tests/entrypoints/openai/test_return_routed_experts.py +++ b/tests/entrypoints/openai/test_return_routed_experts.py @@ -12,15 +12,13 @@ MODEL_NAME = "TitanML/tiny-mixtral" # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers. -# The published config has sliding_window=4096, which produces -# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a -# FullAttentionSpec group, so we override sliding_window=null below. +# Its published sliding window is incompatible with AuxOutput retention. NUM_LOCAL_EXPERTS = 8 NUM_EXPERTS_PER_TOK = 2 NUM_HIDDEN_LAYERS = 2 -def assert_valid_routed_experts(encoded: str | None) -> None: +def decode_routed_experts(encoded: str | None) -> np.ndarray: assert encoded is not None routed_experts = np.load(io.BytesIO(base64.b64decode(encoded))) assert routed_experts.ndim == 3 @@ -30,6 +28,7 @@ def assert_valid_routed_experts(encoded: str | None) -> None: assert topk == NUM_EXPERTS_PER_TOK assert (routed_experts >= 0).all() assert (routed_experts < NUM_LOCAL_EXPERTS).all() + return routed_experts @pytest.fixture(scope="module") @@ -44,7 +43,11 @@ def server(): "--hf-overrides", '{"sliding_window": null}', ] - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + with RemoteOpenAIServer( + MODEL_NAME, + args, + env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"}, + ) as remote_server: yield remote_server @@ -63,4 +66,8 @@ async def test_routed_experts(server): choice = result.model_dump()["choices"][0] assert choice["token_ids"] is not None - assert_valid_routed_experts(choice["routed_experts"]) + assert result.usage is not None + routed_experts = decode_routed_experts(choice["routed_experts"]) + assert len(routed_experts) == ( + result.usage.prompt_tokens + len(choice["token_ids"]) - 1 + ) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py index e7527a071b75..9359006562bb 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py @@ -15,9 +15,7 @@ GEN_ENDPOINT = "/inference/v1/generate" # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers. -# The published config has sliding_window=4096, which produces -# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a -# FullAttentionSpec group, so we override sliding_window=null below. +# Its published sliding window is incompatible with AuxOutput retention. NUM_LOCAL_EXPERTS = 8 NUM_EXPERTS_PER_TOK = 2 NUM_HIDDEN_LAYERS = 2 @@ -36,7 +34,13 @@ def server(): '{"sliding_window": null}', "--enable-scale-out", ] - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + with RemoteOpenAIServer( + MODEL_NAME, + args, + env_dict={ + "VLLM_USE_V2_MODEL_RUNNER": "1", + }, + ) as remote_server: yield remote_server @@ -56,9 +60,10 @@ async def client(server: RemoteOpenAIServer): @pytest.mark.asyncio async def test_generate_routed_experts(client): """Test that /inference/v1/generate returns routed_experts when enabled.""" + prompt_token_ids = [1, 2, 3] payload = { "model": MODEL_NAME, - "token_ids": [1, 2, 3], + "token_ids": prompt_token_ids, "sampling_params": {"max_tokens": 10, "temperature": 0.0}, "stream": False, } @@ -80,3 +85,4 @@ async def test_generate_routed_experts(client): assert topk == NUM_EXPERTS_PER_TOK assert (routed_experts >= 0).all() assert (routed_experts < NUM_LOCAL_EXPERTS).all() + assert num_tokens == len(prompt_token_ids) + len(choice["token_ids"]) - 1 diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index b146a235c1c9..ff683ae44c67 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -7,26 +7,19 @@ import pytest import torch -from vllm.config import ModelConfig, VllmConfig +from vllm.config import VllmConfig from vllm.config.compilation import CompilationMode from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, - RoutedExpertsManager, bind_routed_experts_capturer, - get_routed_experts_attn_gid, ) from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter -from vllm.transformers_utils.model_arch_config_convertor import ( - ModelArchConfigConvertorBase, -) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, - KVCacheConfig, KVCacheGroupSpec, MLAAttentionSpec, - SlidingWindowMLASpec, UniformTypeKVCacheSpecs, ) @@ -42,6 +35,7 @@ def _capturer_with_buffer( num_experts_per_tok: int = 2, dp_rank: int = 0, tp_size: int = 1, + dtype: torch.dtype = torch.int32, ) -> RoutedExpertsCapturer: # Bypass __init__ so the test can use a CPU buffer and skip the # VllmConfig dependency. The CUDA device-tensor allocation in the @@ -49,10 +43,11 @@ def _capturer_with_buffer( c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) c.dp_rank = dp_rank c.tp_size = tp_size + c.output_dtype = torch.uint8 c.device_buffer = torch.full( (max_tokens, num_layers, num_experts_per_tok), -1, - dtype=torch.int32, + dtype=dtype, ) return c @@ -88,74 +83,36 @@ def _make_modular_routed_experts(): ) -def _make_model_config(hf_config): - num_experts_per_token = ModelArchConfigConvertorBase( - hf_config, hf_config - ).get_num_experts_per_token() - model_config = SimpleNamespace( - hf_text_config=hf_config, - model_arch_config=SimpleNamespace( - num_experts_per_token=num_experts_per_token, - ), - ) - model_config.get_num_experts = lambda: hf_config.num_experts - model_config.get_num_experts_per_tok = lambda: ( - ModelConfig.get_num_experts_per_tok(model_config) - ) - model_config.get_total_num_hidden_layers = lambda: hf_config.num_hidden_layers - return model_config - - -def test_routed_experts_manager_uses_gemma4_top_k_experts(): - hf_config = SimpleNamespace( - num_experts=8, - top_k_experts=2, - num_hidden_layers=3, - ) - vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config)) - kv_cache_spec = FullAttentionSpec( - block_size=4, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ) - kv_cache_config = KVCacheConfig( - num_blocks=2, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], - ) - - manager = RoutedExpertsManager(vllm_config, kv_cache_config) - - assert manager.routed_experts_by_slot.shape == (8, 3, 2) - - -def test_routed_experts_manager_uses_kimi_k3_experts_per_token(): - hf_config = SimpleNamespace( - num_experts=8, - num_experts_per_token=2, - num_hidden_layers=3, - ) - vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config)) - kv_cache_spec = FullAttentionSpec( - block_size=4, +def _full_attention_kv_group( + spec_type: type[FullAttentionSpec] = FullAttentionSpec, +) -> KVCacheGroupSpec: + full_attention = spec_type( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32, ) - kv_cache_config = KVCacheConfig( - num_blocks=2, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + return KVCacheGroupSpec( + ["layer"], + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer": full_attention}, + ), ) - manager = RoutedExpertsManager(vllm_config, kv_cache_config) - - assert manager.routed_experts_by_slot.shape == (8, 3, 2) +@pytest.mark.parametrize("eplb_enabled", [False, True]) +def test_base_router_capture_pre_eplb_mapping(eplb_enabled): + eplb_state = None + if eplb_enabled: + eplb_state = EplbLayerState() + eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64) + eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1) + eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64) + eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool) + eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)] + router = _make_router(eplb_state) -def test_base_router_capture_pre_eplb_mapping(): - router = _make_router() captured = [] def capture_fn(ids): @@ -173,34 +130,7 @@ def capture_fn(ids): assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]])) -def test_base_router_capture_with_eplb_enabled(): - eplb_state = EplbLayerState() - eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64) - eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1) - eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64) - eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool) - eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)] - router = _make_router(eplb_state=eplb_state) - - captured = [] - - def capture_fn(ids): - captured.append(ids.clone()) - - router.set_capture_fn(capture_fn) - _, topk_ids = router.select_experts( - hidden_states=torch.empty(1), - router_logits=torch.empty(1), - ) - - assert len(captured) == 1 - # Capture should see logical ids pre-EPLB mapping. - assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]])) - # Our DummyRouter mapping adds +10. - assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]])) - - -def test_public_binding_only_visits_target_model(monkeypatch): +def test_public_binding_binds_target_model_router(monkeypatch): class DummyFusedMoE: def __init__(self, layer_id): self.layer_id = layer_id @@ -208,7 +138,6 @@ def __init__(self, layer_id): self._quant_method = _make_modular_routed_experts().quant_method target_module = DummyFusedMoE(layer_id=7) - draft_module = DummyFusedMoE(layer_id=0) import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer @@ -221,12 +150,30 @@ def __init__(self, layer_id): ) assert target_module.router.capture_fn is not None - assert draft_module.router.capture_fn is None topk_ids = torch.tensor([[5, 6]]) target_module.router.capture_fn(topk_ids) assert calls == [(7, topk_ids)] +def test_public_binding_supports_direct_capture_source(): + class DirectCaptureSource(torch.nn.Module): + def __init__(self): + super().__init__() + self.layer_id = 3 + self.capture_fn = None + + source = DirectCaptureSource() + calls = [] + capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args)) + + bind_routed_experts_capturer(source, capturer) + + assert source.capture_fn is not None + topk_ids = torch.tensor([[1, 2]]) + source.capture_fn(topk_ids) + assert calls == [(3, topk_ids)] + + def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch): class DummyFusedMoE: def __init__(self): @@ -277,6 +224,51 @@ def test_routed_experts_capturer_single_dp_no_metadata(): assert capturer.device_buffer[3, 0, 0].item() == -1 +@pytest.mark.parametrize( + ("num_experts", "dtype_name", "torch_dtype"), + [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)], +) +def test_routed_experts_capturer_exposes_output_profile( + monkeypatch, num_experts, dtype_name, torch_dtype +): + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as module + + monkeypatch.setattr(module, "current_platform", SimpleNamespace(device_type="cpu")) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_total_num_hidden_layers=lambda: 3, + get_num_experts=lambda: num_experts, + get_num_experts_per_tok=lambda: 2, + hf_text_config=SimpleNamespace(model_type="test"), + ), + parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1), + ) + + capturer = RoutedExpertsCapturer(8, config) + + assert capturer.shape_per_token == (3, 2) + assert capturer.output_dtype_name == dtype_name + assert capturer.output_dtype == torch_dtype + assert capturer.device_buffer.shape == (8, 3, 2) + + +@pytest.mark.parametrize("output_dtype", [torch.uint8, torch.uint16]) +def test_routed_experts_capturer_narrows_snapshot(output_dtype): + capturer = _capturer_with_buffer(dtype=torch.int32) + capturer.output_dtype = output_dtype + topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64) + ctx = SimpleNamespace(dp_metadata=None) + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): + capturer.capture(layer_id=0, topk_ids=topk) + + output = capturer.snapshot_routing_data(2) + assert capturer.device_buffer.dtype == torch.int32 + assert capturer.device_buffer[:2, 0, :].tolist() == topk.tolist() + assert output.dtype == output_dtype + assert output[:, 0, :].tolist() == topk.tolist() + assert output.data_ptr() != capturer.device_buffer.data_ptr() + + def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" capturer = _capturer_with_buffer(dp_rank=1) @@ -307,6 +299,47 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): assert torch.equal(capturer.device_buffer[:3, 0, :], topk) +def test_routed_experts_capturer_dp_ep_gathered_shards(): + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32) + ctx = SimpleNamespace( + dp_metadata=SimpleNamespace( + num_tokens_across_dp_cpu=num_tokens_dp, + local_sizes=[2, 2, 1, 1], + ) + ) + topk = torch.tensor( + [[0, 1], [2, 3], [4, 5], [-1, -1], [10, 11], [12, 13]], + dtype=torch.int32, + ) + + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): + capturer.capture(layer_id=0, topk_ids=topk) + + assert torch.equal(capturer.device_buffer[:2, 0], topk[4:]) + + +def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) + ctx = SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ) + local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) + gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32) + tp_group = Mock() + tp_group.all_gather.return_value = gathered + + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), + ): + capturer.capture(layer_id=0, topk_ids=local_shard) + + tp_group.all_gather.assert_called_once_with(local_shard, dim=0) + assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3]) + + def test_routed_experts_capturer_dp_unexpected_batch_raises(): """Mismatch between topk batch dim and DP layout: fail fast.""" capturer = _capturer_with_buffer(dp_rank=0) @@ -324,128 +357,119 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises(): assert capturer.device_buffer[0, 0, 0].item() == -1 -def test_routed_experts_attention_group_is_shared_and_fail_closed(): - """Both sides key routing data by this gid, so it must skip non-full-attention - groups rather than defaulting to 0, and fail closed when none exists.""" - common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) - config = SimpleNamespace( - kv_cache_groups=[ - KVCacheGroupSpec( - ["swa_layer"], - SlidingWindowMLASpec(block_size=4, sliding_window=8, **common), - ), - KVCacheGroupSpec(["full_layer"], FullAttentionSpec(block_size=4, **common)), - ] - ) - assert get_routed_experts_attn_gid(config) == 1 - - with pytest.raises(ValueError, match="requires a full-attention KV cache group"): - get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[])) - - -def test_routed_experts_attention_group_unwraps_uniform_type_specs(): - """DeepSeek-V4-shaped groups wrap their specs in ``UniformTypeKVCacheSpecs``. +def test_get_aux_output_connector_passes_config(monkeypatch): + import vllm.distributed.aux_output_connector.worker as aux_output_worker - The wrapper is not a ``FullAttentionSpec``, so a bare isinstance check finds - no group and fails closed on every worker. Unwrap semantics themselves are - covered by ``test_is_full_attention_spec_*`` in tests/v1/core. - """ - common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32) - swa_spec = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common) - mla_spec = MLAAttentionSpec(block_size=4, **common) + connector = Mock() + constructor = Mock(return_value=connector) + monkeypatch.setattr(aux_output_worker, "AuxOutputWorkerConnector", constructor) config = SimpleNamespace( - kv_cache_groups=[ - KVCacheGroupSpec( - ["swa_layer"], - UniformTypeKVCacheSpecs( - block_size=4, kv_cache_specs={"swa_layer": swa_spec} - ), - ), - KVCacheGroupSpec( - ["mla_layer"], - UniformTypeKVCacheSpecs( - block_size=4, kv_cache_specs={"mla_layer": mla_spec} - ), - ), - ] + scheduler_config=SimpleNamespace(max_num_batched_tokens=32) ) + model = Mock() + kv_cache_config = Mock() - assert get_routed_experts_attn_gid(config) == 1 + result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config) - swa_only = SimpleNamespace(kv_cache_groups=[config.kv_cache_groups[0]]) - with pytest.raises(ValueError, match="requires a full-attention KV cache group"): - get_routed_experts_attn_gid(swa_only) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_mrv2_async_output_returns_existing_routed_experts_field(): - from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors - from vllm.v1.worker.gpu.async_utils import AsyncOutput - from vllm.v1.worker.gpu.sample.output import SamplerOutput - - routed_experts = RoutedExpertsTensors( - routing_data=torch.arange(6, dtype=torch.int32, device="cuda").reshape(3, 1, 2), - slot_mapping=torch.tensor([11, 12, 13], device="cuda"), - ) - num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - sampler_output = SamplerOutput( - sampled_token_ids=torch.tensor([[1]], device="cuda"), - logprobs_tensors=None, - num_nans=None, - num_sampled=num_sampled, - num_rejected=torch.tensor([0], dtype=torch.int32, device="cuda"), + constructor.assert_called_once_with( + model=model, + kv_cache_config=kv_cache_config, + vllm_config=config, ) - output = AsyncOutput( - model_runner_output=ModelRunnerOutput(req_ids=["req"], req_id_to_index={}), - sampler_output=sampler_output, - num_sampled_tokens=num_sampled, - main_stream=torch.cuda.current_stream(), - copy_stream=torch.cuda.Stream(), - routed_experts=routed_experts, - ).get_output() - - assert output.routed_experts is not None - assert output.routed_experts.routing_data[:, 0, 0].tolist() == [0, 2, 4] - assert output.routed_experts.slot_mapping.tolist() == [11, 12, 13] + assert result is connector -@pytest.mark.parametrize("rank", [0, 1]) -def test_all_tp_ranks_initialize_capture(monkeypatch, rank): - pytest.importorskip("vllm.vllm_flash_attn", exc_type=ImportError) - import vllm.v1.worker.gpu.model_runner as model_runner +def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatch): + import vllm.distributed.aux_output_connector.worker as aux_output_worker capturer = Mock() constructor = Mock(return_value=capturer) bind = Mock() - monkeypatch.setattr(model_runner, "RoutedExpertsCapturer", constructor) - monkeypatch.setattr(model_runner, "bind_routed_experts_capturer", bind) - - runner = model_runner.GPUModelRunner.__new__(model_runner.GPUModelRunner) - runner.max_num_tokens = 32 - runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(rank=rank)) - runner.kv_cache_config = SimpleNamespace() - runner.model = Mock() + monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", constructor) + monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind) + monkeypatch.setattr( + aux_output_worker, + "get_tp_group", + lambda: SimpleNamespace(is_first_rank=False, world_size=1), + ) - runner.init_routed_experts_capturer() + config = SimpleNamespace( + aux_output_config=SimpleNamespace(enable_return_routed_experts=True), + kv_transfer_config=None, + max_concurrent_batches=2, + scheduler_config=SimpleNamespace(max_num_batched_tokens=32), + ) + model = Mock() + connector = aux_output_worker.AuxOutputWorkerConnector( + vllm_config=config, + model=model, + kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]), + ) constructor.assert_called_once_with( max_num_batched_tokens=32, - vllm_config=runner.vllm_config, - kv_cache_config=runner.kv_cache_config, + vllm_config=config, + ) + bind.assert_called_once_with(model, capturer) + connector.begin_step(Mock()) + assert connector.prepare_output(Mock()) is None + capturer.snapshot_routing_data.assert_not_called() + + +def test_aux_output_worker_connector_default_capacity(monkeypatch): + import vllm.distributed.aux_output_connector.worker as aux_output_worker + + tp_group = SimpleNamespace(is_first_rank=True, world_size=1) + store_constructor = Mock() + background_store_constructor = Mock(side_effect=lambda store, **_: store) + capturer = SimpleNamespace(shape_per_token=(2,), output_dtype_name="int32") + monkeypatch.setattr(aux_output_worker, "get_tp_group", lambda: tp_group) + monkeypatch.setattr( + aux_output_worker, "RoutedExpertsCapturer", Mock(return_value=capturer) + ) + monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", Mock()) + monkeypatch.setattr( + aux_output_worker, + "resolve_kv_cache_block_sizes", + lambda *_: (32, 16), + ) + monkeypatch.setattr(aux_output_worker, "BlockObjectStore", store_constructor) + monkeypatch.setattr( + aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor ) - bind.assert_called_once_with(runner.model, capturer) - assert runner.routed_experts_capturer is capturer + monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock()) + + config = SimpleNamespace( + aux_output_config=SimpleNamespace(max_bytes=None), + kv_transfer_config=None, + cache_config=SimpleNamespace(enable_prefix_caching=True), + scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32), + max_concurrent_batches=2, + ) + kwargs = dict( + vllm_config=config, + model=Mock(), + kv_cache_config=SimpleNamespace( + num_blocks=10, + kv_cache_groups=[_full_attention_kv_group(MLAAttentionSpec)], + ), + ) + + aux_output_worker.AuxOutputWorkerConnector(**kwargs) + assert store_constructor.call_args.kwargs["max_bytes"] == 2560 + assert store_constructor.call_args.kwargs["object_nbytes"] == 128 + assert background_store_constructor.call_args.kwargs["max_pending_batches"] == 16 def test_v2_model_runner_accepts_routed_experts(monkeypatch): monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ()) config = SimpleNamespace( model_config=SimpleNamespace( - enable_return_routed_experts=True, use_mla=False, logits_processors=None, enable_prompt_embeds=False, ), + aux_output_config=SimpleNamespace(enable_return_routed_experts=True), speculative_config=None, parallel_config=SimpleNamespace( prefill_context_parallel_size=1, diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7e27a4517e03..27107aa6b00f 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -354,8 +354,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.finished_req_ids_dict = None scheduler.grammar_compile_error_reqs = set() scheduler.vllm_config = Mock() - scheduler.vllm_config.model_config.enable_return_routed_experts = False - scheduler.enable_return_routed_experts = False + scheduler.aux_output_connector = None scheduler.return_sampling_mask = False scheduler.recompute_kv_load_failures = False scheduler.defer_block_free = False diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 4951014d1ba1..a58bcab7d5ea 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -19,6 +19,10 @@ SpeculativeConfig, VllmConfig, ) +from vllm.distributed.aux_output_connector.connector import ( + AuxOutputSchedulerConnector, + AuxRequestOutput, +) from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( HiSparseConnector, @@ -55,7 +59,6 @@ ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, - RoutedExpertsLists, SamplingMaskLists, make_empty_encoder_model_runner_output, ) @@ -136,21 +139,18 @@ def test_add_requests(): assert len(scheduler.waiting) == i + 1 -def test_routed_experts_prompt_start_at_prompt_end(): +@pytest.mark.parametrize("max_tokens", [1, 2]) +def test_routed_experts_prompt_start_at_prompt_end(max_tokens): """A prompt-end offset should return empty routing data, not crash.""" scheduler = create_scheduler() - (request,) = create_requests(num_requests=1) + (request,) = create_requests(num_requests=1, max_tokens=max_tokens) request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens + scheduler.aux_output_connector = AuxOutputSchedulerConnector() scheduler.add_request(request) scheduler_output = scheduler.schedule() - num_tokens = scheduler_output.num_scheduled_tokens[request.request_id] - - manager = Mock() - manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8) - manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8) - scheduler.enable_return_routed_experts = True - scheduler.routed_experts_mgr = manager - scheduler._re_block_ids = {request.request_id: []} + assert scheduler_output.aux_output_connector_metadata.requests == { + request.request_id: request.num_prompt_tokens + } outputs = scheduler.update_from_output( scheduler_output, @@ -161,21 +161,23 @@ def test_routed_experts_prompt_start_at_prompt_end(): logprobs=None, prompt_logprobs_dict={}, pooler_output=[], - routed_experts=RoutedExpertsLists( - routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8), - slot_mapping=np.arange(num_tokens), - ), + aux_output_connector_output={ + request.request_id: AuxRequestOutput( + request.num_prompt_tokens, + np.empty((0, 1, 1), dtype=np.uint8), + ) + }, ), ) - manager.get.assert_called_once_with( - [], request.num_prompt_tokens, token_start=request.num_prompt_tokens - ) - assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1) + output = outputs[request.client_index].outputs[0] + assert output.new_token_ids == [0] + assert output.routed_experts.shape == (0, 1, 1) def test_finish_request(): scheduler = create_scheduler() + scheduler.aux_output_connector = Mock() requests = create_requests(num_requests=10) for request in requests: scheduler.add_request(request) @@ -184,6 +186,7 @@ def test_finish_request(): scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) assert request.request_id not in scheduler.requests assert len(scheduler.waiting) == 9 - i + scheduler.aux_output_connector.request_finished.assert_called_with(request) def test_get_num_unfinished_requests(): @@ -1373,6 +1376,31 @@ def test_preemption_re_records_prefix_cache_query(): assert stats.preempted_requests == 1 +def test_preemption_processes_stale_aux_output(): + scheduler = create_scheduler(enable_prefix_caching=True) + request = create_requests(num_requests=1)[0] + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + scheduler.running.remove(request) + scheduler.aux_output_connector = Mock() + scheduler._preempt_request(request, 0.0) + scheduler.aux_output_connector.request_finished.assert_called_once_with(request) + + scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[1000]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + scheduler.aux_output_connector.take_output.assert_called_once_with(request, None) + + def test_prefix_cache_stats_not_recorded_when_caching_disabled(): """With prefix caching off there is no local lookup, so admitting a request records no phantom miss.""" @@ -1461,6 +1489,7 @@ def spy_allocate_slots(*args, **kwargs): def test_scheduler_reset_prefix_cache(): scheduler = create_scheduler(enable_prefix_caching=True) + scheduler.aux_output_connector = Mock() requests = create_requests(num_requests=10) for request in requests: scheduler.add_request(request) @@ -1477,10 +1506,20 @@ def test_scheduler_reset_prefix_cache(): # Reset prefix cache should fail since there are still running requests # and they are taking KV cache assert not scheduler.reset_prefix_cache() + scheduler.aux_output_connector.reset.assert_not_called() + + with pytest.raises(RuntimeError, match=r"pause\(mode='keep'\)"): + scheduler.reset_prefix_cache(reset_running_requests=True) + + # pause(mode="keep") also waits for scheduled model outputs to drain. + scheduler.set_pause_state(PauseState.PAUSED_ALL) + with pytest.raises(RuntimeError, match="model output is in flight"): + scheduler.reset_prefix_cache(reset_running_requests=True) + for request in requests: + request.num_in_flight_tokens = 0 - # Reset prefix cache with reset_running_requests=True. All running requests - # Should be pushed back to the waiting queue and kv cache should be freed assert scheduler.reset_prefix_cache(reset_running_requests=True) + scheduler.aux_output_connector.reset.assert_called_once_with() # Verify requests moved from running to waiting assert len(scheduler.waiting) == len(requests) @@ -1490,6 +1529,16 @@ def test_scheduler_reset_prefix_cache(): assert scheduler.waiting[i] == request +@pytest.mark.parametrize("reset_successful", [False, True]) +def test_aux_output_reset_follows_kv_reset_result(reset_successful: bool): + scheduler = create_scheduler(enable_prefix_caching=True) + scheduler.aux_output_connector = Mock() + scheduler.kv_cache_manager.reset_prefix_cache = Mock(return_value=reset_successful) + + assert scheduler.reset_prefix_cache() is reset_successful + assert scheduler.aux_output_connector.reset.call_count == int(reset_successful) + + def test_kv_cache_release_after_keep_pause_preserves_requests(): """A completed keep pause permits release and recomputation of live requests.""" scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True) @@ -3791,10 +3840,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.kv_event_publisher = Mock() scheduler.finished_req_ids = set() scheduler.finished_req_ids_dict = None + scheduler.aux_output_connector = None scheduler.grammar_compile_error_reqs = set() scheduler.vllm_config = Mock() - scheduler.vllm_config.model_config.enable_return_routed_experts = False - scheduler.enable_return_routed_experts = False scheduler.return_sampling_mask = False scheduler.recompute_kv_load_failures = False scheduler.defer_block_free = False diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index e388b73468d3..5b6826a4b3fc 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -1122,6 +1122,7 @@ def fake_sample( logprobs_tensors=None, num_nans=None, num_sampled=num_sampled, + num_rejected=num_rejected, ) return sampler_output, num_sampled, num_rejected diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 69128f55b704..485187f22211 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -3,8 +3,9 @@ import math import time -from unittest.mock import MagicMock +from unittest.mock import MagicMock, Mock +import numpy as np import pytest from tests.v1.engine.utils import ( @@ -1307,6 +1308,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors): async def test_request_output_collector(): NUM_REQS = 3 TEXT = "a" + routed_experts = np.arange(12, dtype=np.uint8).reshape(2, 3, 2) def make_outputs() -> list[RequestOutput]: return [ @@ -1322,6 +1324,9 @@ def make_outputs() -> list[RequestOutput]: token_ids=[idx], cumulative_logprob=(idx + 1 * 1.0), logprobs=[{"a": idx, "b": idx}], + routed_experts=( + routed_experts if idx == NUM_REQS - 1 else None + ), finish_reason="length" if (idx == NUM_REQS - 1) else None, ) ], @@ -1374,6 +1379,7 @@ def make_outputs() -> list[RequestOutput]: assert output.finished assert output.outputs[0].finish_reason == "length" + np.testing.assert_array_equal(output.outputs[0].routed_experts, routed_experts) # Text, token_ids, and logprobs should get merged. assert output.outputs[0].text == TEXT * num_to_put for tok_0, tok_1 in zip(output.outputs[0].token_ids, list(range(num_to_put))): @@ -1385,6 +1391,46 @@ def make_outputs() -> list[RequestOutput]: assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected +def test_routed_experts_are_accumulated_until_finish(): + state = RequestState( + request_id="request-int", + external_req_id="request", + parent_req=None, + request_index=0, + lora_request=None, + output_kind=RequestOutputKind.DELTA, + prompt="prompt", + prompt_token_ids=[1, 2, 3, 4], + prompt_embeds=None, + logprobs_processor=Mock( + logprobs=None, + cumulative_logprob=0.0, + pop_prompt_logprobs=Mock(return_value=None), + ), + detokenizer=Mock(get_next_output_text=Mock(return_value="")), + max_tokens_param=2, + arrival_time=0.0, + queue=None, + log_stats=False, + stream_interval=1, + ) + prompt_chunk = np.arange(24, dtype=np.uint8).reshape(4, 3, 2) + decode_chunk = np.arange(12, dtype=np.uint8).reshape(2, 3, 2) + 24 + state.routed_experts_chunks.append(prompt_chunk) + + partial = state.make_request_output([5], None, None, None) + assert partial is not None + assert partial.outputs[0].routed_experts is None + + state.routed_experts_chunks.append(decode_chunk) + finished = state.make_request_output([6], None, FinishReason.LENGTH, None) + assert finished is not None + np.testing.assert_array_equal( + finished.outputs[0].routed_experts, + np.concatenate((prompt_chunk, decode_chunk)), + ) + + @pytest.mark.asyncio async def test_cumulative_output_collector_n(): """Test collector correctly handles multiple outputs by index.""" diff --git a/tests/v1/executor/test_ray_utils.py b/tests/v1/executor/test_ray_utils.py index 0fa2c129f1b3..2f945c5a3f12 100644 --- a/tests/v1/executor/test_ray_utils.py +++ b/tests/v1/executor/test_ray_utils.py @@ -3,12 +3,14 @@ import numpy as np +from vllm.distributed.aux_output_connector.connector import ( + AuxRequestOutput, +) from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output from vllm.v1.outputs import ( LogprobsLists, LogprobsTensors, ModelRunnerOutput, - RoutedExpertsLists, ) @@ -59,26 +61,19 @@ def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views(): assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs -def test_detach_zero_copy_routed_experts_without_logprobs(): +def test_detach_zero_copy_aux_output_without_logprobs(): + rows = _make_readonly(np.arange(12, dtype=np.uint8).reshape(2, 3, 2)) output = ModelRunnerOutput( req_ids=["req-0"], req_id_to_index={"req-0": 0}, - routed_experts=RoutedExpertsLists( - routing_data=_make_readonly(np.arange(12, dtype=np.int32).reshape(2, 3, 2)), - slot_mapping=_make_readonly(np.array([7, 8], dtype=np.int64)), - ), + aux_output_connector_output={"req-0": AuxRequestOutput(0, rows)}, ) - original = output.routed_experts - assert output.logprobs is None detach_zero_copy_from_model_runner_output(output) - detached = output.routed_experts - assert detached is not None - assert detached is not original - assert detached.routing_data is not original.routing_data - assert detached.slot_mapping is not original.slot_mapping - assert detached.routing_data.flags.writeable - assert detached.slot_mapping.flags.writeable - np.testing.assert_array_equal(detached.routing_data, original.routing_data) - np.testing.assert_array_equal(detached.slot_mapping, original.slot_mapping) + aux_output = output.aux_output_connector_output + assert aux_output is not None + detached = aux_output["req-0"].rows + assert detached is not rows + assert detached.flags.writeable + np.testing.assert_array_equal(detached, rows) diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index ff2ba0493d8d..940aea75521f 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config.attention import AttentionConfig, HiSparseConfig +from vllm.config.aux_output import AuxOutputConfig from vllm.config.cache import CacheConfig from vllm.config.compilation import ( CompilationConfig, @@ -66,6 +67,8 @@ # __all__ should only contain classes and functions. # Types and globals should be imported from their respective modules. __all__ = [ + # From vllm.config.aux_output + "AuxOutputConfig", # From vllm.config.attention "AttentionConfig", "HiSparseConfig", diff --git a/vllm/config/aux_output.py b/vllm/config/aux_output.py new file mode 100644 index 000000000000..f70f57e70f4e --- /dev/null +++ b/vllm/config/aux_output.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for execution auxiliary outputs.""" + +from pydantic import Field + +from vllm.config.utils import config + + +@config +class AuxOutputConfig: + """Configuration for auxiliary-output delivery.""" + + enable_return_routed_experts: bool = False + """Capture and return routed-experts auxiliary outputs.""" + + max_bytes: int | None = Field(default=None, gt=0) + """LRU capacity, or ``None`` to derive it from the KV cache capacity.""" + + @property + def enabled(self) -> bool: + """Whether any execution auxiliary output is enabled.""" + return self.enable_return_routed_experts + + def compute_hash(self) -> str: + """Hash AuxOutput settings that alter the model forward graph.""" + from vllm.config.utils import hash_factors + + return hash_factors( + {"enable_return_routed_experts": self.enable_return_routed_experts} + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index 9ee35069dab4..145814a26866 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -247,8 +247,6 @@ class ModelConfig: NOTE: This disables both `torch.compile` and CUDA graphs, and is equivalent to setting `-cc.mode=none -cc.cudagraph_mode=none`.""" - enable_return_routed_experts: bool = False - """Whether to return routed experts.""" return_sampling_mask: bool = False """Whether to return the post-processing token support for each sample.""" max_logprobs: int = Field(default=20, ge=-1) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 1ff675a8f010..bf11255e8e0a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -28,6 +28,7 @@ from vllm.utils.hashing import safe_hash from .attention import AttentionConfig, HiSparseConfig +from .aux_output import AuxOutputConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig @@ -376,6 +377,8 @@ class VllmConfig: """Model weight offloading configuration.""" attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" + aux_output_config: AuxOutputConfig = Field(default_factory=AuxOutputConfig) + """Execution auxiliary output configuration.""" engram_config: EngramConfig | None = None """N-gram embedding storage and sharding settings.""" mamba_config: MambaConfig = Field(default_factory=MambaConfig) @@ -554,6 +557,7 @@ def compute_hash(self, include_version: bool = True) -> str: vllm_factors.append(self.ec_transfer_config.compute_hash()) else: vllm_factors.append("None") + vllm_factors.append(self.aux_output_config.compute_hash()) if self.additional_config: if isinstance(additional_config := self.additional_config, dict): additional_config_hash = safe_hash( @@ -1084,6 +1088,53 @@ def _post_init_kv_transfer_config(self) -> None: # This is the same for all backends self.kv_transfer_config.kv_role = "kv_both" + def _verify_aux_output_compatibility(self) -> None: + """Reject configurations unsupported by enabled auxiliary outputs.""" + if not self.aux_output_config.enabled: + return + if not self.use_v2_model_runner: + raise ValueError( + "AuxOutput Connector requires Model Runner V2; set " + "VLLM_USE_V2_MODEL_RUNNER=1." + ) + if self.model_config.runner_type != "generate": + raise ValueError("AuxOutput Connector only supports generate runners.") + if not self.model_config.is_moe: + raise ValueError("AuxOutput Connector only supports MoE models.") + if not self.cache_config.enable_prefix_caching: + raise ValueError("AuxOutput Connector requires prefix caching.") + if ( + self.speculative_config is not None + and self.speculative_config.enable_adaptive_verification + ): + raise ValueError( + "--enable-return-routed-experts is incompatible with " + "adaptive speculative verification." + ) + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError( + "--enable-return-routed-experts is incompatible with " + "pipeline parallelism (PP > 1)." + ) + if ( + self.parallel_config.decode_context_parallel_size > 1 + or self.parallel_config.prefill_context_parallel_size > 1 + ): + raise ValueError( + "--enable-return-routed-experts is incompatible with " + "context parallelism (DCP/PCP > 1)." + ) + + kv_transfer_config = self.kv_transfer_config + if ( + kv_transfer_config is not None + and kv_transfer_config.is_kv_transfer_instance + ): + raise ValueError( + "--enable-return-routed-experts is incompatible with KV " + "connectors (PD disaggregation and KV cache offload)." + ) + def _verify_kv_transfer_compat(self) -> None: """Reject configurations that silently corrupt KV transfers.""" if ( @@ -1294,38 +1345,6 @@ def __post_init__(self): self.parallel_config.is_moe_model = self.model_config.is_moe - if ( - self.model_config is not None - and self.model_config.enable_return_routed_experts - ): - if self.parallel_config.pipeline_parallel_size > 1: - raise ValueError( - "--enable-return-routed-experts is incompatible with " - "pipeline parallelism (PP > 1)." - ) - if ( - self.parallel_config.decode_context_parallel_size > 1 - or self.parallel_config.prefill_context_parallel_size > 1 - ): - raise ValueError( - "--enable-return-routed-experts is incompatible with context " - "parallelism (DCP > 1 or PCP > 1)." - ) - - # Incompatible with any KV connector — covers both PD disaggregation - # (kv_producer/kv_consumer: routing captured on P can't reach D) and - # single-instance KV offload/sharing (kv_both: slot_mapping semantics - # change when KV blocks live outside local GPU memory, breaking the - # slot-indexed routed_experts buffer). - if ( - self.kv_transfer_config is not None - and self.kv_transfer_config.is_kv_transfer_instance - ): - raise ValueError( - "--enable-return-routed-experts is incompatible with KV " - "connectors (PD disaggregation, KV cache offload)." - ) - if ( self.model_config is not None and self.model_config.multimodal_config is not None @@ -2032,6 +2051,7 @@ def has_blocked_weights(): # Resolve kv_offloading-derived connector name into kv_transfer_config # before the HMA check below, which inspects the connector class. self._post_init_kv_transfer_config() + self._verify_aux_output_compatibility() if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching: # Such an instance publishes encoder embeddings and runs no language @@ -2689,7 +2709,7 @@ def __str__(self): f"quantization={self.model_config.quantization}, " f"quantization_config={self.model_config.quantization_config}, " # noqa f"enforce_eager={self.model_config.enforce_eager}, " - f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, " # noqa + f"aux_output_config={self.aux_output_config!r}, " f"kv_cache_dtype={self.cache_config.cache_dtype}, " f"device_config={self.device_config.device}, " f"structured_outputs_config={self.structured_outputs_config!r}, " diff --git a/vllm/distributed/aux_output_connector/__init__.py b/vllm/distributed/aux_output_connector/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/distributed/aux_output_connector/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/distributed/aux_output_connector/connector.py b/vllm/distributed/aux_output_connector/connector.py new file mode 100644 index 000000000000..bd8cbf4d183d --- /dev/null +++ b/vllm/distributed/aux_output_connector/connector.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Scheduler-side control plane for execution auxiliary outputs.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.request import Request + + +@dataclass +class PackedBlockHashes: + """Contiguous, self-contained block hashes for scheduler-worker IPC.""" + + data: bytes + item_size: int + + def __iter__(self) -> Iterator[bytes]: + for start in range(0, len(self.data), self.item_size): + yield self.data[start : start + self.item_size] + + +@dataclass +class AuxOutputConnectorMetadata: + generation: int + requests: dict[str, int] + block_hashes: dict[str, PackedBlockHashes] + finished_requests: tuple[str, ...] + + +@dataclass +class AuxRequestOutput: + token_start: int + rows: np.ndarray + + +class AuxOutputSchedulerConnector: + """Build worker metadata without owning auxiliary output payloads or stores.""" + + def __init__(self) -> None: + # Number of hashes already sent to the worker for each active request. + self._sent_hash_counts: dict[str, int] = {} + # Terminal events are delivered with the next connector metadata. + self._finished_requests: dict[str, PackedBlockHashes | None] = {} + self._generation = 0 + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + requests: dict[str, Request], + ) -> AuxOutputConnectorMetadata: + """Build one step's incremental worker metadata.""" + scheduled_requests: dict[str, int] = {} + block_hashes_by_request: dict[str, PackedBlockHashes] = {} + for request_id in scheduler_output.num_scheduled_tokens: + num_sent = self._sent_hash_counts.setdefault(request_id, 0) + request = requests[request_id] + packed = self._pack_new_hashes(request.block_hashes, num_sent) + if packed is not None: + block_hashes_by_request[request_id] = packed + self._sent_hash_counts[request_id] = len(request.block_hashes) + assert request.sampling_params is not None + scheduled_requests[request_id] = max( + request.sampling_params.routed_experts_prompt_start, + 0 if request.num_output_tokens == 0 else request.num_tokens - 1, + ) + # A settled token can complete a hash block after the next async schedule + # was built. Send a hash-only update if the request was not rescheduled. + if len(self._sent_hash_counts) != len(scheduled_requests): + for request_id, num_sent in self._sent_hash_counts.items(): + if request_id in scheduled_requests: + continue + request = requests[request_id] + packed = self._pack_new_hashes(request.block_hashes, num_sent) + if packed is None: + continue + block_hashes_by_request[request_id] = packed + self._sent_hash_counts[request_id] = len(request.block_hashes) + # Sending transfers ownership of these one-shot events. + finished_requests = tuple(self._finished_requests) + block_hashes_by_request.update( + (request_id, block_hashes) + for request_id, block_hashes in self._finished_requests.items() + if block_hashes is not None + ) + self._finished_requests = {} + return AuxOutputConnectorMetadata( + self._generation, + scheduled_requests, + block_hashes_by_request, + finished_requests, + ) + + def take_output( + self, + request: Request, + output: dict[str, AuxRequestOutput] | None, + ) -> np.ndarray | None: + """Return the accepted R3 rows for one scheduled request.""" + request_id = request.request_id + assert output is not None and request_id in output, ( + f"auxiliary output worker output is missing {request_id}" + ) + request_output = output[request_id] + token_end = request.num_tokens - 1 + local_end = token_end - request_output.token_start + if local_end < 0: + assert not request.is_finished(), ( + "finished auxiliary output output has no accepted token range: " + f"request={request_id}, token_end={token_end}, " + f"output_start={request_output.token_start}, " + "output_end=" + f"{request_output.token_start + len(request_output.rows)}" + ) + return None + assert local_end <= len(request_output.rows), ( + "auxiliary output worker output has an invalid token range: " + f"request={request_id}, token_end={token_end}, " + f"output_start={request_output.token_start}, " + f"output_end={request_output.token_start + len(request_output.rows)}" + ) + return request_output.rows[:local_end] + + def request_finished(self, request: Request) -> None: + """Queue a request's terminal event and final block hashes.""" + request_id = request.request_id + num_sent = self._sent_hash_counts.pop(request_id, None) + if num_sent is None: + return + # The next metadata delivers this terminal event to the worker. + self._finished_requests[request_id] = self._pack_new_hashes( + request.block_hashes, num_sent + ) + + @staticmethod + def _pack_new_hashes( + block_hashes: Sequence[bytes], num_sent: int + ) -> PackedBlockHashes | None: + assert num_sent <= len(block_hashes), "KV block-hash history shrank" + new_hashes = block_hashes[num_sent:] + if not new_hashes: + return None + return PackedBlockHashes(b"".join(new_hashes), len(new_hashes[0])) + + def reset(self) -> None: + """Start a new auxiliary output namespace after a prefix-cache reset.""" + # The worker drops temporary state on generation changes; resend hashes. + self._sent_hash_counts.clear() + self._finished_requests.clear() + self._generation += 1 diff --git a/vllm/distributed/aux_output_connector/routed_experts.py b/vllm/distributed/aux_output_connector/routed_experts.py new file mode 100644 index 000000000000..1c89b97215a6 --- /dev/null +++ b/vllm/distributed/aux_output_connector/routed_experts.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routed-experts auxiliary output keys, publication, and materialization.""" + +from __future__ import annotations + +from collections.abc import Hashable, Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from vllm.distributed.aux_output_connector.store import ( + BackgroundBlockObjectStore, + BlockObject, + BlockObjectStore, +) + + +@dataclass +class _RequestTail: + slot: int + block_start: int + length: int = 0 + + +class RoutedExpertsBuffer: + """Own incomplete and unkeyed full routed-experts blocks.""" + + def __init__( + self, + dtype: np.dtype[Any], + shape_per_token: tuple[int, ...], + block_size: int, + max_num_seqs: int, + max_num_batched_tokens: int, + max_concurrent_batches: int, + ) -> None: + self.dtype = dtype + self.shape_per_token = shape_per_token + self.block_size = block_size + max_step_blocks = ( + max_num_batched_tokens + block_size - 1 + ) // block_size + max_num_seqs + # In-flight full blocks plus one incomplete tail per active request. + max_blocks = max_concurrent_batches * max_step_blocks + max_num_seqs + self._rows: np.ndarray = np.empty( + (max_blocks, block_size, *shape_per_token), dtype=dtype + ) + self._free_slots = list(range(len(self._rows) - 1, -1, -1)) + self._owned_slots: dict[int, int] = {} + self._requests: dict[Hashable, _RequestTail] = {} + + def _tail(self, request_id: Hashable, block_start: int) -> _RequestTail: + tail = self._requests.get(request_id) + if tail is not None: + assert tail.block_start == block_start, ( + "auxiliary output capture skipped an incomplete block: " + f"request={request_id}, expected={tail.block_start}, " + f"actual={block_start}" + ) + return tail + assert self._free_slots, "auxiliary output block pool is exhausted" + tail = _RequestTail(self._free_slots.pop(), block_start) + self._requests[request_id] = tail + return tail + + def capture( + self, request_id: Hashable, token_start: int, rows: np.ndarray + ) -> list[tuple[int, np.ndarray]]: + """Stage rows and return completed blocks without retaining them.""" + rows = np.asarray(rows) + assert rows.shape[1:] == self.shape_per_token and rows.dtype == self.dtype, ( + "routed-experts capture profile changed" + ) + if token_start < 0: + raise ValueError("auxiliary output token start must be non-negative") + + completed: list[tuple[int, np.ndarray]] = [] + offset = 0 + while offset < len(rows): + position = token_start + offset + block_start = position // self.block_size * self.block_size + local_start = position - block_start + + # Full aligned input blocks do not need tail staging. + if ( + request_id not in self._requests + and local_start == 0 + and len(rows) - offset >= self.block_size + ): + completed.append((block_start, rows[offset : offset + self.block_size])) + offset += self.block_size + continue + + tail = self._tail(request_id, block_start) + assert local_start == tail.length, ( + "auxiliary output capture is not contiguous: " + f"request={request_id}, expected={block_start + tail.length}, " + f"actual={position}" + ) + count = min(self.block_size - local_start, len(rows) - offset) + if count == 1: + self._rows[tail.slot, local_start] = rows[offset] + else: + self._rows[tail.slot, local_start : local_start + count] = rows[ + offset : offset + count + ] + tail.length += count + offset += count + if tail.length == self.block_size: + block = self._rows[tail.slot] + del self._requests[request_id] + self._owned_slots[id(block)] = tail.slot + completed.append((block_start, block)) + + return completed + + def read( + self, request_id: Hashable, token_start: int, token_end: int + ) -> np.ndarray: + tail = self._requests.get(request_id) + assert tail is not None, ( + f"auxiliary output buffer is missing request {request_id}" + ) + local_start = token_start - tail.block_start + local_end = token_end - tail.block_start + assert 0 <= local_start < local_end <= tail.length, ( + "auxiliary output range is unavailable: " + f"request={request_id}, range=[{token_start}, {token_end}), " + f"available=[{tail.block_start}, {tail.block_start + tail.length})" + ) + return self._rows[tail.slot, local_start:local_end].copy() + + def retain_block(self, rows: np.ndarray) -> np.ndarray: + """Retain one unkeyed block after the current capture call.""" + if id(rows) in self._owned_slots: + return rows + assert self._free_slots, "auxiliary output block pool is exhausted" + slot = self._free_slots.pop() + retained = self._rows[slot] + retained[...] = rows + self._owned_slots[id(retained)] = slot + return retained + + def release_block(self, rows: np.ndarray) -> None: + slot = self._owned_slots.pop(id(rows), None) + if slot is not None: + self._free_slots.append(slot) + + def discard(self, request_id: Hashable) -> None: + tail = self._requests.pop(request_id, None) + if tail is not None: + self._free_slots.append(tail.slot) + + def reset(self) -> None: + self._requests.clear() + self._owned_slots.clear() + self._free_slots[:] = range(len(self._rows) - 1, -1, -1) + + +def routed_experts_keys( + block_hashes: Iterable[bytes], aux_output_namespace: str +) -> list[str]: + prefix = f"vllm-artifact/{aux_output_namespace}/" + return [prefix + block_hash.hex() for block_hash in block_hashes] + + +def materialize_routed_experts( + store: BackgroundBlockObjectStore | BlockObjectStore, + aux_output_keys: list[str], + *, + shape_per_token: tuple[int, ...], + dtype: np.dtype[Any], +) -> np.ndarray: + if not aux_output_keys: + raise ValueError("routed-experts auxiliary output key list must not be empty") + payload = store.get_concatenated(aux_output_keys) + return np.frombuffer(payload, dtype=dtype).reshape((-1, *shape_per_token)) + + +def publish_routed_experts( + store: BackgroundBlockObjectStore | BlockObjectStore, + *, + batches: Sequence[tuple[Sequence[str], list[tuple[int, np.ndarray]]]], + block_size: int, + retain_keys: Sequence[str] = (), + release_keys: Sequence[str] = (), +) -> None: + """Publish immutable full R3 blocks.""" + objects = [] + for aux_output_keys, blocks in batches: + for block_start, array in blocks: + if block_start < 0 or block_start % block_size: + raise ValueError( + "auxiliary output block start is not hash-block aligned" + ) + block_index = block_start // block_size + if block_index >= len(aux_output_keys): + raise ValueError( + "auxiliary output block has no corresponding key: " + f"start={block_start}, index={block_index}, " + f"keys={len(aux_output_keys)}, block_size={block_size}" + ) + if len(array) != block_size: + raise ValueError( + "auxiliary output block length does not match hash block size" + ) + objects.append( + BlockObject( + key=aux_output_keys[block_index], + payload=array.tobytes(order="C"), + ) + ) + store.put( + objects, + retain_keys=retain_keys, + release_keys=release_keys, + ) diff --git a/vllm/distributed/aux_output_connector/store.py b/vllm/distributed/aux_output_connector/store.py new file mode 100644 index 000000000000..2b364288a588 --- /dev/null +++ b/vllm/distributed/aux_output_connector/store.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""In-process auxiliary-output storage.""" + +from __future__ import annotations + +import mmap +import queue +import threading +from collections import OrderedDict +from collections.abc import Iterable +from dataclasses import dataclass + + +@dataclass(frozen=True) +class BlockObject: + """One immutable auxiliary output object.""" + + key: str + payload: bytes + + +class BlockObjectStoreError(RuntimeError): + """AuxOutput storage or retrieval failed.""" + + +class BackgroundBlockObjectStore: + """Serialize store mutations on a background thread.""" + + def __init__(self, store: BlockObjectStore, *, max_pending_batches: int) -> None: + self._store = store + self._queue: queue.Queue[ + tuple[list[BlockObject], tuple[str, ...], tuple[str, ...]] | None + ] = queue.Queue(maxsize=max_pending_batches) + self._error: BaseException | None = None + self._closed = False + self._thread = threading.Thread( + target=self._run, + daemon=True, + name="vllm-aux-output-writer", + ) + self._thread.start() + + def _run(self) -> None: + while True: + update = self._queue.get() + try: + if update is None: + return + if self._error is None: + objects, retain_keys, release_keys = update + self._store.put( + objects, + retain_keys=retain_keys, + release_keys=release_keys, + ) + except BaseException as error: + self._error = error + finally: + self._queue.task_done() + + def _raise_if_failed(self) -> None: + if self._error is not None: + raise BlockObjectStoreError( + "auxiliary output publication failed" + ) from self._error + + def put( + self, + objects: list[BlockObject], + *, + retain_keys: Iterable[str] = (), + release_keys: Iterable[str] = (), + ) -> None: + retains = tuple(retain_keys) + releases = tuple(release_keys) + if not objects and not retains and not releases: + return + if self._closed: + raise RuntimeError("auxiliary output store is closed") + self._queue.put((objects, retains, releases)) + self._raise_if_failed() + + def get_concatenated(self, keys: list[str]) -> bytes: + self._queue.join() + self._raise_if_failed() + return self._store.get_concatenated(keys) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._queue.join() + self._queue.put(None) + self._thread.join() + self._store.close() + self._raise_if_failed() + + +class BlockObjectStore: + """Single-owner bounded store that fails closed after an eviction.""" + + _UNALLOCATED_SLOT = -1 + + def __init__( + self, + *, + max_bytes: int, + object_nbytes: int, + ) -> None: + if object_nbytes <= 0: + raise ValueError("auxiliary output object size must be positive") + if max_bytes < object_nbytes: + raise ValueError("auxiliary output store must fit at least one object") + self.object_nbytes = object_nbytes + self.num_slots = max_bytes // object_nbytes + self._lru: OrderedDict[str, int] = OrderedDict() + self._references: dict[str, int] = {} + self._free_slots: list[int] = [] + self._next_slot = 0 + self._arena: mmap.mmap | None = mmap.mmap( + -1, self.num_slots * self.object_nbytes + ) + + def _evict_to_fit(self, protected: set[str]) -> None: + excess = len(self._lru) - self.num_slots + if excess <= 0: + return + victims = [] + for key in self._lru: + if key not in self._references and key not in protected: + victims.append(key) + if len(victims) == excess: + break + if len(victims) != excess: + raise BlockObjectStoreError( + "auxiliary output store cannot retain the requested batch: " + f"limit={self.num_slots} objects" + ) + for victim in victims: + slot = self._lru.pop(victim) + if slot != self._UNALLOCATED_SLOT: + self._free_slots.append(slot) + + def _allocate_slot(self) -> int: + if self._free_slots: + return self._free_slots.pop() + slot = self._next_slot + self._next_slot += 1 + return slot + + def put( + self, + objects: list[BlockObject], + *, + retain_keys: Iterable[str] = (), + release_keys: Iterable[str] = (), + ) -> None: + arena = self._arena + if arena is None: + raise RuntimeError("auxiliary output store is closed") + unique = {obj.key: obj for obj in objects} + self._retain(retain_keys) + terminal_order = self._release(release_keys) + for key in unique: + self._lru.setdefault(key, self._UNALLOCATED_SLOT) + self._lru.move_to_end(key) + for key in terminal_order: + if key in self._lru: + self._lru.move_to_end(key) + try: + self._evict_to_fit(set(unique) - set(terminal_order)) + except BlockObjectStoreError: + for key in unique: + if self._lru.get(key) == self._UNALLOCATED_SLOT: + del self._lru[key] + raise + for object_id, obj in unique.items(): + slot = self._lru.get(object_id) + if slot is None: + continue + if slot != self._UNALLOCATED_SLOT: + continue + slot = self._allocate_slot() + offset = slot * self.object_nbytes + arena[offset : offset + self.object_nbytes] = obj.payload + self._lru[object_id] = slot + + def _retain(self, keys: Iterable[str]) -> None: + for key in keys: + references = self._references.get(key, 0) + self._references[key] = references + 1 + + def _release(self, keys: Iterable[str]) -> list[str]: + terminal_order = [] + for key in keys: + references = self._references[key] - 1 + if references: + self._references[key] = references + continue + del self._references[key] + terminal_order.append(key) + return terminal_order + + def get_concatenated(self, keys: list[str]) -> bytes: + arena_obj = self._arena + if arena_obj is None: + raise RuntimeError("auxiliary output store is closed") + try: + entries = [self._lru[key] for key in keys] + except KeyError as error: + raise BlockObjectStoreError( + "auxiliary output object does not exist; the object may have been " + f"evicted (used={len(self._lru)}, " + f"limit={self.num_slots} objects). Increase " + "aux_output_config.max_bytes when a KV cache hit requires it." + ) from error + arena = memoryview(arena_obj) + try: + payload = b"".join( + arena[slot * self.object_nbytes : (slot + 1) * self.object_nbytes] + for slot in entries + ) + finally: + arena.release() + for key in keys: + self._lru.move_to_end(key) + return payload + + def close(self) -> None: + arena = self._arena + if arena is not None: + self._arena = None + arena.close() diff --git a/vllm/distributed/aux_output_connector/worker.py b/vllm/distributed/aux_output_connector/worker.py new file mode 100644 index 000000000000..7959e15ea884 --- /dev/null +++ b/vllm/distributed/aux_output_connector/worker.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker-side auxiliary-output data plane.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from threading import Lock +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +from vllm.config import VllmConfig +from vllm.distributed.aux_output_connector.connector import ( + AuxOutputConnectorMetadata, + AuxRequestOutput, +) +from vllm.distributed.aux_output_connector.routed_experts import ( + RoutedExpertsBuffer, + materialize_routed_experts, + publish_routed_experts, + routed_experts_keys, +) +from vllm.distributed.aux_output_connector.store import ( + BackgroundBlockObjectStore, + BlockObjectStore, +) +from vllm.distributed.parallel_state import get_tp_group +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, + bind_routed_experts_capturer, +) +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes + +if TYPE_CHECKING: + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.worker.gpu.input_batch import InputBatch + + +@dataclass +class _WorkerRequestState: + aux_output_keys: list[str] = field(default_factory=list) + pending_blocks: list[tuple[int, np.ndarray]] = field(default_factory=list) + capture_cursor: int | None = None + scheduled_cursor: int = 0 + emit_cursor: int = 0 + pending_outputs: int = 0 + # Terminal event received; teardown waits for in-flight step outputs. + finished: bool = False + + +@dataclass +class PendingAuxOutput: + """Own one step's R3 snapshot until its asynchronous copy is consumed.""" + + connector: AuxOutputWorkerConnector + request_ids: list[str] + token_starts: np.ndarray + query_start_loc: np.ndarray + # GPU snapshot taken on the main stream. + routed_experts_gpu: torch.Tensor + # Filled by enqueue_cpu_copy on the output copy stream. + routed_experts: np.ndarray | None = None + num_sampled: np.ndarray | None = None + num_rejected: np.ndarray | None = None + + def enqueue_cpu_copy( + self, num_sampled: np.ndarray, num_rejected: np.ndarray + ) -> None: + """Enqueue asynchronous D2H copies; call on the output copy stream.""" + self.num_sampled = num_sampled + self.num_rejected = num_rejected + self.routed_experts = self.routed_experts_gpu.to( + "cpu", non_blocking=True + ).numpy() + + def process_output(self) -> dict[str, AuxRequestOutput]: + return self.connector.process_output(self) + + +class AuxOutputWorkerConnector: + """Own capture, request tails, and backend resources on the output worker.""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + model: torch.nn.Module, + kv_cache_config: KVCacheConfig, + ) -> None: + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + capturer = RoutedExpertsCapturer( + max_num_batched_tokens=max_num_batched_tokens, vllm_config=vllm_config + ) + bind_routed_experts_capturer(model, capturer) + self._capturer = capturer + self._store: BackgroundBlockObjectStore | None = None + self._buffer: RoutedExpertsBuffer | None = None + self._requests: dict[str, _WorkerRequestState] = {} + self._generation = 0 + self._step_metadata: AuxOutputConnectorMetadata | None = None + # Steps whose asynchronous copy has not been consumed yet. The engine + # consumes step outputs in order, so at most max_concurrent_batches + # are outstanding; guarded by _lock against the output thread. + self._pending_outputs: list[PendingAuxOutput] = [] + self._lock = Lock() + self._max_concurrent_batches = vllm_config.max_concurrent_batches + # Every TP rank participates in capture collectives, but only the + # executor output rank owns the auxiliary output data plane. + if not get_tp_group().is_first_rank: + return + + shape_per_token = self._capturer.shape_per_token + dtype: np.dtype[Any] = np.dtype(self._capturer.output_dtype_name) + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) + hashes_per_kv_block = scheduler_block_size // hash_block_size + block_nbytes = hash_block_size * int(np.prod(shape_per_token)) * dtype.itemsize + max_bytes = vllm_config.aux_output_config.max_bytes + if max_bytes is None: + max_bytes = kv_cache_config.num_blocks * hashes_per_kv_block * block_nbytes + self._store = BackgroundBlockObjectStore( + BlockObjectStore(max_bytes=max_bytes, object_nbytes=block_nbytes), + max_pending_batches=2 * vllm_config.scheduler_config.max_num_seqs, + ) + self._buffer = RoutedExpertsBuffer( + dtype, + shape_per_token, + hash_block_size, + vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens, + vllm_config.max_concurrent_batches, + ) + + def prepare_output(self, input_batch: InputBatch) -> PendingAuxOutput | None: + """Snapshot one step's R3 tensor for asynchronous CPU transfer.""" + if self._buffer is None or self._step_metadata is None: + return None + + request_ids = list(input_batch.req_ids) + query_start_loc = input_batch.query_start_loc_np[: len(request_ids) + 1] + num_rows = int(query_start_loc[-1]) + pending_output = PendingAuxOutput( + connector=self, + request_ids=request_ids, + token_starts=input_batch.num_computed_tokens_np, + query_start_loc=query_start_loc, + routed_experts_gpu=self._capturer.snapshot_routing_data(num_rows), + ) + with self._lock: + assert len(self._pending_outputs) < self._max_concurrent_batches, ( + "auxiliary output step outputs are not consumed in order" + ) + self._pending_outputs.append(pending_output) + for request_id in pending_output.request_ids: + self._requests[request_id].pending_outputs += 1 + return pending_output + + def process_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]: + """Commit one consumed R3 snapshot and build request outputs. + + Request teardown deferred by begin_step is completed here once no + in-flight step covers the request. + """ + with self._lock: + teardown = [] + release_keys: list[str] = [] + try: + outputs = self._commit_output(pending) + finally: + self._pending_outputs.remove(pending) + for request_id in pending.request_ids: + state = self._requests[request_id] + state.pending_outputs -= 1 + if state.pending_outputs == 0 and state.finished: + teardown.append(request_id) + release_keys.extend(reversed(state.aux_output_keys)) + if teardown: + self._publish_blocks([], release_keys=release_keys) + self._teardown(teardown) + return outputs + + def _commit_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]: + buffer = self._buffer + store = self._store + assert buffer is not None and store is not None + block_size = buffer.block_size + + routed_experts = pending.routed_experts + num_sampled = pending.num_sampled + num_rejected = pending.num_rejected + assert ( + routed_experts is not None + and num_sampled is not None + and num_rejected is not None + ), "auxiliary output CPU copy was not enqueued" + + # Publish the whole batch before materializing any consumer output. + materialize_outputs: list[tuple[str, int, int]] = [] + block_batches = [] + outputs: dict[str, AuxRequestOutput] = {} + + # Use the ModelRunner's actual batch boundaries rather than rebuilding them. + for request_id, token_start, start, end, sampled, rejected in zip( + pending.request_ids, + pending.token_starts, + pending.query_start_loc[:-1], + pending.query_start_loc[1:], + num_sampled, + num_rejected, + strict=True, + ): + request_num_tokens = end - start + assert request_num_tokens > 0, ( + "auxiliary output request token count must be positive" + ) + state = self._requests[request_id] + + # Capture precedes speculative acceptance, so discard the rejected + # suffix. Batch boundaries still span the full executed range. + rejected = int(rejected) + assert 0 <= rejected <= request_num_tokens, ( + "auxiliary output rejected-token count is invalid" + ) + rows = routed_experts[start : end - rejected] + + capture_start = token_start + capture_cursor = state.capture_cursor + if capture_cursor is None: + capture_cursor = capture_start + + assert capture_start >= capture_cursor, ( + "auxiliary output capture moved backwards" + ) + if capture_start > capture_cursor: + # Reattach after an optimistically scheduled suffix was rejected. + assert capture_cursor < state.scheduled_cursor, ( + "auxiliary output capture has an unbacked token gap" + ) + capture_start = capture_cursor + + emit_start = state.emit_cursor + # Complete blocks without keys remain pending until a hash update. + completed = buffer.capture(request_id, capture_start, rows) + token_end = capture_start + len(rows) + state.capture_cursor = token_end + state.scheduled_cursor = token_start + request_num_tokens + block_batches.append((state, completed)) + + if sampled > 0 and emit_start <= token_end: + if emit_start >= capture_start: + outputs[request_id] = AuxRequestOutput( + emit_start, rows[emit_start - capture_start :] + ) + state.emit_cursor = token_end + else: + materialize_outputs.append((request_id, emit_start, token_end)) + + # A consumer may reuse a block produced earlier in the same batch. + self._publish_blocks(block_batches) + + for request_id, emit_start, token_end in materialize_outputs: + state = self._requests[request_id] + stored_end = ( + min(token_end // block_size, len(state.aux_output_keys)) * block_size + ) + if emit_start < stored_end: + first_block = emit_start // block_size + stored = materialize_routed_experts( + store, + state.aux_output_keys[first_block : stored_end // block_size], + shape_per_token=buffer.shape_per_token, + dtype=buffer.dtype, + ) + local_start = emit_start % block_size + rows = stored[local_start : local_start + stored_end - emit_start] + if stored_end < token_end: + rows = np.concatenate( + (rows, buffer.read(request_id, stored_end, token_end)) + ) + else: + rows = buffer.read(request_id, emit_start, token_end) + outputs[request_id] = AuxRequestOutput(emit_start, rows) + state.emit_cursor = token_end + return outputs + + def _publish_blocks( + self, + batches: list[tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]], + retain_keys: Sequence[str] = (), + release_keys: Sequence[str] = (), + ) -> None: + store = self._store + buffer = self._buffer + assert store is not None and buffer is not None + ready_batches = [] + for state, completed in batches: + blocks = state.pending_blocks + completed + keyed_end = len(state.aux_output_keys) * buffer.block_size + ready = [(start, rows) for start, rows in blocks if start < keyed_end] + state.pending_blocks = [ + (start, buffer.retain_block(rows)) + for start, rows in blocks + if start >= keyed_end + ] + if ready: + ready_batches.append((state.aux_output_keys, ready)) + if ready_batches or retain_keys or release_keys: + publish_routed_experts( + store, + batches=ready_batches, + block_size=buffer.block_size, + retain_keys=retain_keys, + release_keys=release_keys, + ) + for _, blocks in ready_batches: + for _, rows in blocks: + buffer.release_block(rows) + + def _teardown(self, request_ids: Iterable[str]) -> None: + """Drop finished requests' temporary state after publishing key releases.""" + buffer = self._buffer + assert buffer is not None + for request_id in request_ids: + state = self._requests.pop(request_id) + for _, rows in state.pending_blocks: + buffer.release_block(rows) + buffer.discard(request_id) + + def begin_step(self, metadata: AuxOutputConnectorMetadata | None) -> None: + """Apply one scheduler step's request and block-hash updates. + + While an unconsumed step covers a finished request, its teardown is + deferred to that step's process_output. + """ + self._step_metadata = metadata + if self._buffer is None or metadata is None: + return + assert not metadata.requests.keys() & metadata.finished_requests, ( + "auxiliary output request cannot run and finish in one step" + ) + assert metadata.generation >= self._generation, ( + "auxiliary output metadata generation moved backwards" + ) + with self._lock: + release_keys: list[str] = [] + if metadata.generation > self._generation: + # A generation change follows a prefix-cache reset, which the + # scheduler only performs once all model output is consumed. + assert not self._pending_outputs, ( + "auxiliary output generation changed with output in flight" + ) + release_keys.extend( + key + for state in self._requests.values() + for key in reversed(state.aux_output_keys) + ) + self._buffer.reset() + self._requests.clear() + self._generation = metadata.generation + for request_id, emit_start in metadata.requests.items(): + state = self._requests.setdefault( + request_id, _WorkerRequestState(emit_cursor=emit_start) + ) + if state.pending_outputs == 0: + # In-flight steps leave the worker cursor behind the + # scheduler's optimistic view; only check settled requests. + assert emit_start <= state.emit_cursor, ( + "auxiliary output Scheduler emit cursor moved ahead" + ) + block_batches: list[ + tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]] + ] = [] + retained_keys: list[str] = [] + for request_id, block_hashes in metadata.block_hashes.items(): + state = self._requests[request_id] + keys = routed_experts_keys(block_hashes, str(self._generation)) + state.aux_output_keys.extend(keys) + retained_keys.extend(keys) + block_batches.append((state, [])) + finished_now: list[str] = [] + for request_id in metadata.finished_requests: + state = self._requests[request_id] + state.finished = True + if state.pending_outputs: + continue + finished_now.append(request_id) + release_keys.extend(reversed(state.aux_output_keys)) + self._publish_blocks(block_batches, retained_keys, release_keys) + self._teardown(finished_now) + + def close(self) -> None: + with self._lock: + self._pending_outputs.clear() + if self._store is not None: + try: + self._store.close() + finally: + self._store = None + + +def get_aux_output_connector( + model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig +) -> AuxOutputWorkerConnector: + return AuxOutputWorkerConnector( + model=model, + kv_cache_config=kv_cache_config, + vllm_config=vllm_config, + ) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0d33caf625a2..9fad681d7dd9 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -35,6 +35,7 @@ import vllm.envs as envs from vllm.config import ( AttentionConfig, + AuxOutputConfig, CacheConfig, CompilationConfig, ConfigType, @@ -447,7 +448,9 @@ class EngineArgs: """Arguments for vLLM engine.""" model: str = ModelConfig.model - enable_return_routed_experts: bool = ModelConfig.enable_return_routed_experts + # Public compatibility argument. Canonical runtime state lives in + # AuxOutputConfig. + enable_return_routed_experts: bool = False return_sampling_mask: bool = ModelConfig.return_sampling_mask model_weights: str = ModelConfig.model_weights served_model_name: str | list[str] | None = ModelConfig.served_model_name @@ -674,6 +677,7 @@ class EngineArgs: structured_outputs_config: StructuredOutputsConfig = get_field( VllmConfig, "structured_outputs_config" ) + aux_output_config: AuxOutputConfig = get_field(VllmConfig, "aux_output_config") reasoning_parser: str = StructuredOutputsConfig.reasoning_parser reasoning_parser_plugin: str | None = None @@ -816,6 +820,10 @@ def __post_init__(self): self.compilation_config = CompilationConfig(**self.compilation_config) if isinstance(self.attention_config, dict): self.attention_config = AttentionConfig(**self.attention_config) + if isinstance(self.aux_output_config, dict): + self.aux_output_config = AuxOutputConfig(**self.aux_output_config) + if self.enable_return_routed_experts: + self.aux_output_config.enable_return_routed_experts = True if isinstance(self.engram_config, dict): self.engram_config = EngramConfig(**self.engram_config) if isinstance(self.mamba_config, dict): @@ -884,6 +892,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: """Shared CLI arguments for vLLM engine.""" # Model arguments model_kwargs = get_kwargs(ModelConfig) + aux_output_kwargs = get_kwargs(AuxOutputConfig) model_group = parser.add_argument_group( title="ModelConfig", description=ModelConfig.__doc__, @@ -923,7 +932,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: model_group.add_argument("--enforce-eager", **model_kwargs["enforce_eager"]) model_group.add_argument( "--enable-return-routed-experts", - **model_kwargs["enable_return_routed_experts"], + **aux_output_kwargs["enable_return_routed_experts"], ) model_group.add_argument( "--return-sampling-mask", @@ -1749,6 +1758,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: vllm_group.add_argument( "--structured-outputs-config", **vllm_kwargs["structured_outputs_config"] ) + vllm_group.add_argument( + "--aux-output-config", **vllm_kwargs["aux_output_config"] + ) vllm_group.add_argument("--profiler-config", **vllm_kwargs["profiler_config"]) vllm_group.add_argument( "--optimization-level", **vllm_kwargs["optimization_level"] @@ -1856,7 +1868,6 @@ def create_model_config(self) -> ModelConfig: quantization_config=self.quantization_config, allow_deprecated_quantization=self.allow_deprecated_quantization, enforce_eager=self.enforce_eager, - enable_return_routed_experts=self.enable_return_routed_experts, return_sampling_mask=self.return_sampling_mask, max_logprobs=self.max_logprobs, logprobs_mode=self.logprobs_mode, @@ -2669,6 +2680,7 @@ def create_engine_config( load_config=load_config, offload_config=offload_config, attention_config=attention_config, + aux_output_config=self.aux_output_config, engram_config=self.engram_config, mamba_config=mamba_config, kernel_config=kernel_config, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index 235808960263..c49d614549d8 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -10,15 +10,12 @@ from functools import partial from typing import Protocol, runtime_checkable -import numpy as np import torch from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_tp_group from vllm.forward_context import get_forward_context from vllm.platforms import current_platform -from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec -from vllm.v1.outputs import RoutedExpertsTensors logger = logging.getLogger(__name__) @@ -46,37 +43,29 @@ def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]: class RoutedExpertsCapturer: """Worker-side capturer for routed experts, lives on GPU. - Layer-level hooks call :meth:`capture` from inside the forward pass - with the per-layer ``topk_ids`` tensor. The tensor is sliced to the - tokens owned by this DP rank and written into a preallocated device - buffer. At the end of the step, :class:`GPUModelRunner` reads the - device buffer, issues a D2H copy into a pinned CPU buffer, and hands - the result to the scheduler via :class:`RoutedExpertsLists`. + Layer-level hooks call :meth:`capture` inside the forward pass. Routing + rows owned by this DP rank are written into a preallocated device buffer. - The device / pinned-CPU transit buffers use ``torch.int32`` (not a - narrow ``uint8``/``uint16`` sized by ``num_experts``). This keeps the - SP all-gather path free of dtype casts, matches the router's native - ``topk_ids`` indices dtype more closely, and costs only a few MB per - worker (``max_num_batched_tokens * num_layers * top_k * 4`` bytes). - The scheduler-side slot buffer - (``RoutedExpertsManager.routed_experts_by_slot``) still uses the - narrow dtype -- numpy fancy-index assignment in ``store_batch`` - narrows the data on the way in. + The device buffer uses ``int32``. Stable snapshots use the narrowest dtype + that can represent every logical expert ID. Invariants: - One instance per worker; shape is fixed at init and covers the worst-case step (``max_num_batched_tokens`` tokens). - Every routed layer overwrites the current step's token rows. - - ``device_buffer.dtype`` is ``torch.int32``. """ def __init__( self, max_num_batched_tokens: int, vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig, ) -> None: - num_layers, _, num_experts_per_tok = _get_routed_experts_shape(vllm_config) + num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape( + vllm_config + ) + self.shape_per_token = (num_layers, num_experts_per_tok) + self.output_dtype_name = "uint8" if num_experts <= 256 else "uint16" + self.output_dtype = getattr(torch, self.output_dtype_name) logger.info( "RoutedExpertsCapturer: allocating buffer with " "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d " @@ -87,23 +76,12 @@ def __init__( vllm_config.model_config.hf_text_config.model_type, ) self.device_buffer = torch.zeros( - ( - max_num_batched_tokens, - num_layers, - num_experts_per_tok, - ), - # Use int32 for the device / host transit buffers: it - # matches the router's native topk_ids dtype, is universally - # supported by NCCL (uint8/uint16 are version-dependent), - # and the extra bytes are small (few MB per worker). The - # big scheduler-side slot buffer stays narrow. + (max_num_batched_tokens, *self.shape_per_token), dtype=torch.int32, device=current_platform.device_type, ) self.dp_rank = vllm_config.parallel_config.data_parallel_rank self.tp_size = vllm_config.parallel_config.tensor_parallel_size - # KV cache group whose slot layout the routing data is keyed by. - self.attn_gid = get_routed_experts_attn_gid(kv_cache_config) def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: """Capture expert routing decisions for a specific layer. @@ -139,31 +117,21 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: """ ctx = get_forward_context() - if ctx.dp_metadata is None: # single dp - start_loc = 0 - end_loc = topk_ids.shape[0] - token_num_per_dp = topk_ids.shape[0] - else: # multi dp + if ctx.dp_metadata is None: + local_topk_ids = topk_ids + else: num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu - token_num_per_dp = int(num_tokens_dp[self.dp_rank].item()) + num_local_tokens = int(num_tokens_dp[self.dp_rank].item()) total = int(num_tokens_dp.sum().item()) n = topk_ids.shape[0] shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None) gathered_size = sum(shard_sizes) if shard_sizes is not None else -1 if n == total: - # Naive dispatch: all DP ranks' tokens concatenated - # before routing. This rank owns tokens - # [end_loc - token_num_per_dp, end_loc). - cumsum = torch.cumsum(num_tokens_dp, dim=0) - end_loc = int(cumsum[self.dp_rank].item()) - start_loc = end_loc - token_num_per_dp - elif n == token_num_per_dp: - # Modular-kernel path: DP combine happens inside - # quant_method.apply; select_experts only sees this - # rank's tokens, take the whole tensor. - start_loc = 0 - end_loc = token_num_per_dp + end = int(num_tokens_dp[: self.dp_rank + 1].sum().item()) + local_topk_ids = topk_ids[end - num_local_tokens : end] + elif n == num_local_tokens: + local_topk_ids = topk_ids elif shard_sizes is not None and n == gathered_size: # Naive DP+EP dispatch gathers the sequence-parallel shards in # flattened DP-then-TP order. ``local_sizes`` is the exact @@ -173,12 +141,11 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: assert len(shard_sizes) % num_dp_ranks == 0 shards_per_dp_rank = len(shard_sizes) // num_dp_ranks first_shard = self.dp_rank * shards_per_dp_rank - start_loc = sum(shard_sizes[:first_shard]) - end_loc = start_loc + token_num_per_dp + start = sum(shard_sizes[:first_shard]) + local_topk_ids = topk_ids[start : start + num_local_tokens] elif ( self.tp_size > 1 - and n != token_num_per_dp - and n == (token_num_per_dp + self.tp_size - 1) // self.tp_size + and n == (num_local_tokens + self.tp_size - 1) // self.tp_size ): # SP + modular-kernel path. All-gather across the TP # group along dim=0 to reconstruct the full per-DP-rank @@ -193,60 +160,30 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: # # ``topk_ids`` is already whatever the router produced # (typically int32/int64, both supported by NCCL); the - # downstream ``device_buffer[...] = topk_ids[...]`` - # setitem narrows into int32 automatically. - topk_ids = get_tp_group().all_gather(topk_ids, dim=0) - start_loc = 0 - end_loc = token_num_per_dp + # downstream buffer assignment narrows to the capture dtype. + local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[ + :num_local_tokens + ] else: - sp_expected = ( - (token_num_per_dp + self.tp_size - 1) // self.tp_size - if self.tp_size > 0 - else -1 - ) + sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size raise AssertionError( "RoutedExpertsCapturer: unexpected topk_ids batch " - f"dim {n} (expected {total}, {token_num_per_dp}, " + f"dim {n} (expected {total}, {num_local_tokens}, " f"{gathered_size}, or {sp_expected} for " f"dp_rank={self.dp_rank}, tp_size={self.tp_size})" ) - if layer_id >= self.device_buffer.shape[1]: + if not 0 <= layer_id < self.device_buffer.shape[1]: raise IndexError( f"routed-experts layer {layer_id} exceeds capture buffer " f"layer count {self.device_buffer.shape[1]}" ) - self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[ - start_loc:end_loc, : - ] - - def get_device_buffer(self) -> torch.Tensor: - """Return the underlying device buffer so the model runner can - issue the D2H copy. The tensor is shared; callers must either - clone or fully drain it before the next forward pass overwrites it. - """ - return self.device_buffer - - def get_routed_experts( - self, slot_mappings: torch.Tensor, num_tokens: int - ) -> RoutedExpertsTensors: - """Snapshot this step's routing data and its attention slot mapping. - - Both tensors are cloned since the capture buffer and the shared - ``slot_mappings`` are overwritten by the next step, which may race - with an in-flight D2H copy. + self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids - Args: - slot_mappings: Per-KV-cache-group slot mappings for this step, - shape ``(num_kv_cache_groups, max_num_batched_tokens)``. - num_tokens: Total number of tokens scheduled in this step. - - """ - return RoutedExpertsTensors( - routing_data=self.device_buffer[:num_tokens].clone(), - slot_mapping=slot_mappings[self.attn_gid, :num_tokens].clone(), - ) + def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor: + """Return a stable snapshot of the current routing data.""" + return self.device_buffer[:num_tokens].to(self.output_dtype) def bind_routed_experts_capturer( @@ -268,20 +205,12 @@ def bind_routed_experts_capturer( continue if not isinstance(module, MoERunner): continue - layer_id = module.layer_id - - def capture_fn( - topk_ids: torch.Tensor, - layer_id: int = layer_id, - capturer: RoutedExpertsCapturer = capturer, - ) -> None: - capturer.capture(layer_id, topk_ids) - + capture_fn = partial(capturer.capture, module.layer_id) quant_method = module._quant_method - moe_kernel = getattr(quant_method, "moe_kernel", None) - impl = getattr(moe_kernel, "impl", None) - fused_experts = getattr(impl, "fused_experts", None) if quant_method.is_monolithic: + moe_kernel = getattr(quant_method, "moe_kernel", None) + impl = getattr(moe_kernel, "impl", None) + fused_experts = getattr(impl, "fused_experts", None) if not ( isinstance(fused_experts, FusedMoEExpertsMonolithic) and fused_experts.supports_routing_replay_capture() @@ -303,132 +232,3 @@ def capture_fn( if num_bound == 0: raise ValueError("No supported MoE router found for routed-experts capture.") - - -def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int: - """Return the full-attention KV cache group used for routed experts.""" - for gid, group in enumerate(kv_cache_config.kv_cache_groups): - if is_full_attention_spec(group.kv_cache_spec): - return gid - raise ValueError("Routed-experts capture requires a full-attention KV cache group.") - - -class RoutedExpertsManager: - """Scheduler-side slot-indexed buffer for routed experts. - - Lives on CPU in the scheduler process. Each slot corresponds to - ``block_id * block_size + offset_in_block`` where ``block_id`` is - drawn from the physical KV-cache block pool, so routing data is - tied to physical blocks and naturally survives preemption for - prefix-cached blocks (prefix hits re-expose the same slots). - - Data flow per step: - 1. Worker D2Hs its device capture buffer into - :class:`RoutedExpertsLists` and returns it via - :class:`ModelRunnerOutput`. - 2. Scheduler calls :meth:`store_batch` with that step's - ``(routing_data, slot_mapping)`` — a single CPU->CPU - fancy-index assign, ~few MB per step. - 3. On request completion / abort / preemption, the scheduler - calls :meth:`get` with the request's block IDs to recover - the full per-token routing. - - Memory: ``routed_experts_by_slot`` is sized for the whole block - pool (``num_blocks * block_size`` slots). For large block pools - this can reach multiple GB; see the init log for the exact size. - """ - - def __init__( - self, - vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig, - ) -> None: - # Hybrid models also have KV groups whose slot layout differs. - self.attn_gid = get_routed_experts_attn_gid(kv_cache_config) - attn_group = kv_cache_config.kv_cache_groups[self.attn_gid] - self.block_size = attn_group.kv_cache_spec.block_size - - # All kv_cache_groups share the same physical block pool, so - # block IDs span [0, num_blocks) regardless of how many groups - # exist. Sizing to the full pool avoids index-out-of-range - # when different groups happen to land on the same block. - num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape( - vllm_config - ) - max_num_slots = kv_cache_config.num_blocks * self.block_size - # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct - # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping - # this narrow matters because the slot buffer is sized for the - # whole block pool and can reach multiple GB. - expert_id_dtype = np.uint8 if num_experts <= 256 else np.uint16 - self.routed_experts_by_slot = np.zeros( - ( - max_num_slots, - num_layers, - num_experts_per_tok, - ), - dtype=expert_id_dtype, - ) - logger.info( - "RoutedExpertsManager CPU buffer: %.2f GB " - "(slots=%d, layers=%d, top_k=%d, dtype=%s)", - self.routed_experts_by_slot.nbytes / 1e9, - max_num_slots, - num_layers, - num_experts_per_tok, - self.routed_experts_by_slot.dtype.name, - ) - - def store_batch(self, data: np.ndarray, slot_mapping: np.ndarray) -> None: - """Persist one step's routed experts into the slot buffer. - - Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy - indexing handles repeated / out-of-order indices. Called once - per scheduler step in ``update_from_output``. - """ - self.routed_experts_by_slot[slot_mapping] = data - - def get( - self, - block_ids: list[int], - num_tokens: int, - token_start: int = 0, - ) -> np.ndarray: - """Read routed experts data for a completed / preempted request. - - Reconstructs a per-token slot_mapping from the request's block - IDs and returns the routing slice. Because numpy fancy indexing - returns a **copy** (not a view), the returned ndarray is safe - to hold across subsequent :meth:`store_batch` calls — do not - replace the fancy index with a slice without re-verifying. - - Args: - block_ids: Block IDs from the attention KV-cache group. - num_tokens: Number of tokens that have gone through a forward - pass and therefore have routing data written to their - slots (typically ``request.num_tokens - 1``; the last - sampled token has not been forwarded yet). Slots beyond - ``request.num_computed_tokens`` are zero-initialized. - token_start: Skip the first ``token_start`` tokens from the - result. The slot_mapping is sliced before the fancy-index - read, so only the requested slots are fetched — no large - intermediate array is allocated. Clamped to - ``[0, num_tokens]`` automatically. - - Returns: - Array of shape (num_tokens - token_start, num_layers, - num_experts_per_tok). - - """ - bs = self.block_size - block_ids_array = np.array(block_ids, dtype=np.int32) - block_offsets = np.arange(bs) - # slot = block_id * block_size + offset_in_block; flatten the - # (num_blocks, block_size) grid and trim to num_tokens, then - # skip the first token_start entries so only the requested - # range is fetched in a single fancy-index read. - slot_mapping = ( - block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1) - ).flatten()[:num_tokens] - slot_mapping = slot_mapping[token_start:] - return self.routed_experts_by_slot[slot_mapping] diff --git a/vllm/outputs.py b/vllm/outputs.py index 9d46f327d787..e9a08f7c4628 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -194,6 +194,10 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None: completion.cumulative_logprob = ( next_completion.cumulative_logprob ) + # R3 is returned on the terminal output and must survive + # aggregation with earlier chunks that have no R3. + if next_completion.routed_experts is not None: + completion.routed_experts = next_completion.routed_experts completion.finish_reason = next_completion.finish_reason completion.stop_reason = next_completion.stop_reason else: diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 7501c51b2ccd..0da64cf27e0f 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -14,6 +14,9 @@ import numpy.typing as npt import torch + from vllm.distributed.aux_output_connector.connector import ( + AuxOutputConnectorMetadata, + ) from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.lora.request import LoRARequest @@ -23,6 +26,7 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.request import Request else: + AuxOutputConnectorMetadata = object ECConnectorMetadata = object KVConnectorMetadata = object KVCacheBlockCopy = object @@ -285,6 +289,9 @@ class SchedulerOutput: # synchronously during this step (load_async=False). has_sync_kv_loads: bool = False + # Execution-auxiliary output control metadata consumed by the worker connector. + aux_output_connector_metadata: AuxOutputConnectorMetadata | None = None + # EC Cache Connector metadata ec_connector_metadata: ECConnectorMetadata | None = None # EC Cache Manager metadata diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a2638942f880..7d8aee55dc6e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -9,6 +9,7 @@ from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import KVEventsConfig, VllmConfig +from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorBase, ECConnectorMetadata, @@ -26,9 +27,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsManager, -) from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.utils import get_mm_features_in_window @@ -393,26 +391,13 @@ def __init__( if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: self.perf_metrics = ModelMetrics(vllm_config) - self.enable_return_routed_experts = ( - vllm_config.model_config.enable_return_routed_experts + self.aux_output_connector = ( + AuxOutputSchedulerConnector() + if vllm_config.aux_output_config.enabled + else None ) self.return_sampling_mask = vllm_config.model_config.return_sampling_mask - if self.enable_return_routed_experts: - assert self.dcp_world_size == 1 and self.pcp_world_size == 1, ( - "enable_return_routed_experts does not support context parallelism " - "(dcp_world_size > 1 or pcp_world_size > 1)" - ) - - self.routed_experts_mgr = RoutedExpertsManager( - vllm_config=vllm_config, - kv_cache_config=kv_cache_config, - ) - # Block-ID snapshot taken at schedule time (before forward), - # so update_from_output can read slot data even if a later - # schedule() frees the blocks (async scheduling race). - self._re_block_ids: dict[str, list[int]] = {} - self._pause_state: PauseState = PauseState.UNPAUSED # In-flight requests still prefilling (prefill chunks + in-progress @@ -1487,6 +1472,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta + if self.aux_output_connector is not None: + scheduler_output.aux_output_connector_metadata = ( + self.aux_output_connector.build_connector_meta( + scheduler_output, self.requests + ) + ) + # Build the connector meta for ECConnector if self.ec_connector is not None: ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( @@ -1541,6 +1533,8 @@ def _preempt_request( assert request.status == RequestStatus.RUNNING, ( "Only running requests can be preempted" ) + if self.aux_output_connector is not None: + self.aux_output_connector.request_finished(request) self._free_request_blocks(request) self.encoder_cache_manager.free(request) self._inflight_prefills.discard(request) @@ -1596,22 +1590,6 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: if not request.is_prefill_chunk: self._inflight_prefills.discard(request) - # Snapshot block IDs for routed experts before forward starts. - # A concurrent schedule() may preempt requests and free blocks - # before update_from_output runs; the snapshot survives that. - # Use update() to preserve entries from the previous step that - # have not yet been consumed by update_from_output (async - # scheduling may call _update_after_schedule again before the - # prior update_from_output runs). - if self.enable_return_routed_experts: - gid = self.routed_experts_mgr.attn_gid - self._re_block_ids.update( - { - rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid] - for rid in num_scheduled_tokens - } - ) - # Clear the finished and preempted request IDs. # NOTE: We shouldn't just clear() here because it will also affect # the scheduler output. @@ -1962,28 +1940,6 @@ def update_from_output( num_scheduled_tokens, ) ) - - # Persist per-step routed experts into the scheduler-side slot - # buffer (CPU->CPU fancy-index assign; ~few MB per step). - # MUST precede the per-request routing reads below: stopped - # requests may terminate on tokens generated in this very step, - # whose routing was just D2H'd into model_runner_output. - routing_data = None - routing_offsets: dict[str, int] = {} - if model_runner_output.routed_experts is not None: - re = model_runner_output.routed_experts - self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping) - routing_data = re.routing_data.astype( - self.routed_experts_mgr.routed_experts_by_slot.dtype, - copy=False, - ) - # Build offset map using model runner's request order - # (input_batch ordering), NOT scheduler dict order. - offset = 0 - for rid in model_runner_output.req_ids: - routing_offsets[rid] = offset - offset += num_scheduled_tokens[rid] - # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, # the below loop can be a performance bottleneck. We should do our best # to avoid expensive operations inside the loop. @@ -2076,7 +2032,6 @@ def update_from_output( ec_transfer_params = None prefill_stats = None status_before_stop = request.status - num_output_tokens_before = len(request._output_token_ids) # Check for stop and update request status. if new_token_ids: @@ -2113,48 +2068,13 @@ def update_from_output( stopped = True routed_experts = None - if ( - self.enable_return_routed_experts - and routing_data is not None - and new_token_ids - ): - req_offset = routing_offsets[req_id] - end = req_offset + num_tokens_scheduled - block_ids = self._re_block_ids.pop(req_id, []) - if num_output_tokens_before == 0: - # Prefill completed: read full prompt routing from - # slot buffer using the block-ID snapshot taken at - # schedule time (immune to async preemption). - if ( - request.sampling_params is not None - and request.sampling_params.routed_experts_prompt_start - is not None - ): - prompt_start = ( - request.sampling_params.routed_experts_prompt_start - ) - assert prompt_start <= request.num_prompt_tokens - else: - prompt_start = 0 - routed_experts = self.routed_experts_mgr.get( - block_ids, - request.num_prompt_tokens, - token_start=prompt_start, - ) - else: - if scheduled_spec_token_ids: - # Spec decode: accepted tokens at the START of - # the scheduled range, rejected at the end. - routed_experts = routing_data[ - req_offset : req_offset + len(new_token_ids) - ] - else: - # Normal decode / re-prefill: token(s) at the END. - routed_experts = routing_data[end - len(new_token_ids) : end] - should_emit_output = bool( new_token_ids or pooler_output is not None or stopped ) + if self.aux_output_connector is not None and should_emit_output: + routed_experts = self.aux_output_connector.take_output( + request, model_runner_output.aux_output_connector_output + ) if should_emit_output: prefill_stats = request.take_prefill_stats() if prefill_stats is not None: @@ -2641,6 +2561,8 @@ def _free_request( ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: assert request.is_finished() + if self.aux_output_connector is not None: + self.aux_output_connector.request_finished(request) self._inflight_prefills.discard(request) connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) @@ -2773,6 +2695,16 @@ def reset_prefix_cache( Otherwise, this method will only reset the KV prefix cache when there is no running requests taking KV cache. """ + if reset_running_requests and self.aux_output_connector is not None: + if self._pause_state != PauseState.PAUSED_ALL: + raise RuntimeError( + "AuxOutput Connector only supports resetting running requests " + "after pause(mode='keep')." + ) + if any(request.num_in_flight_tokens for request in self.requests.values()): + raise RuntimeError( + "AuxOutput Connector cannot reset while model output is in flight." + ) if reset_running_requests: # For logging. timestamp = time.monotonic() @@ -2803,6 +2735,9 @@ def reset_prefix_cache( if reset_connector: reset_successful = self.reset_connector_cache() and reset_successful + if reset_successful and self.aux_output_connector is not None: + self.aux_output_connector.reset() + return reset_successful def reset_connector_cache(self) -> bool: diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index e609fb6e7e3d..a5edb67dee85 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -233,14 +233,10 @@ def _copy_if_readonly(arr): token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens ) - if output.routed_experts is not None: - routing_data, slot_mapping = output.routed_experts - routing_data_c = _copy_if_readonly(routing_data) - slot_mapping_c = _copy_if_readonly(slot_mapping) - if routing_data_c is not routing_data or slot_mapping_c is not slot_mapping: - output.routed_experts = type(output.routed_experts)( - routing_data_c, slot_mapping_c - ) + aux_output = output.aux_output_connector_output + if aux_output is not None: + for request_output in aux_output.values(): + request_output.rows = _copy_if_readonly(request_output.rows) class FutureWrapper(Future): diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index c4535f47f227..c38eafa4a480 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -15,6 +15,7 @@ from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: + from vllm.distributed.aux_output_connector.connector import AuxRequestOutput from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats from vllm.distributed.kv_events import KVConnectorKVEvents @@ -23,6 +24,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats else: + AuxRequestOutput = object KVConnectorStats = object KVConnectorWorkerMetadata = object KVConnectorKVEvents = object @@ -188,73 +190,6 @@ def empty_cpu( ) -class RoutedExpertsTensors(NamedTuple): - """Device-side snapshot of routed experts data, pending async D2H. - - Produced by :class:`GPUModelRunner` at the end of each async-scheduled - step. The copy stream waits on the default stream, then issues - non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU - buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes - the copy before the scheduler reads it. - - Sliced to ``total_num_scheduled_tokens`` (step-level, across all - requests — NOT per-request). Both ``routing_data`` and - ``slot_mapping`` must be private clones when sourced from shared - capturer / prepare-input buffers, so the next forward pass / - ``_prepare_inputs`` on the default stream does not race with a - D2H still pending on the copy stream. - """ - - # (num_scheduled_tokens, num_layers, num_experts_per_tok) - routing_data: torch.Tensor - # (num_scheduled_tokens,) - slot_mapping: torch.Tensor - - def to_cpu_nonblocking(self) -> "RoutedExpertsTensors": - """Issue non-blocking D2H on the current stream. - - NOTE: ``non_blocking=True`` only delivers true overlap when the - CPU target is pinned. The current fallback here allocates a - new pageable CPU tensor per call, which silently degrades to a - synchronous copy; acceptable because the sync happens on the - dedicated copy stream, not the default stream. - """ - if self.routing_data.device.type == "cpu": - return self - return RoutedExpertsTensors( - self.routing_data.to("cpu", non_blocking=True), - self.slot_mapping.to("cpu", non_blocking=True), - ) - - def tolists(self) -> "RoutedExpertsLists": - """Convert to the numpy-backed form consumed by the scheduler. - - ``.cpu()`` is a no-op when the tensor is already on CPU, so this - is cheap for the post-D2H case; for raw device tensors it will - synchronously block, which is only reached in tests. - """ - return RoutedExpertsLists( - self.routing_data.cpu().numpy(), - self.slot_mapping.cpu().numpy(), - ) - - -class RoutedExpertsLists(NamedTuple): - """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch` - consumes. - - Batched per scheduler step: the leading dim is the number of tokens - scheduled across all requests in this step (``total_num_scheduled_tokens``), - not per-request tokens. ``slot_mapping[i]`` tells the scheduler which - physical KV-cache slot row ``i`` of ``routing_data`` belongs to. - """ - - # (num_scheduled_tokens, num_layers, num_experts_per_tok) - routing_data: np.ndarray - # (num_scheduled_tokens,) - slot_mapping: np.ndarray - - # [num_reqs, ] # The shape of each element depends on the pooler used PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None] @@ -361,16 +296,7 @@ class ModelRunnerOutput: # information related to cudagraph execution cudagraph_stats: CUDAGraphStat | None = None - # Per-step routed experts data captured by the worker. - # ``routing_data`` shape: (num_scheduled_tokens, num_layers, - # num_experts_per_tok); expert IDs as uint8/uint16. - # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache - # slot for each row of routing_data. - # ``num_scheduled_tokens`` is step-level (total across all requests - # in this step), not per-request. The scheduler persists this into - # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``. - # ``None`` when ``enable_return_routed_experts`` is off. - routed_experts: RoutedExpertsLists | None = None + aux_output_connector_output: dict[str, AuxRequestOutput] | None = None # ``None`` when ``return_sampling_mask`` is off. sampling_masks: SamplingMaskLists | None = None diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 0ab0ab3f3cd8..02592980a856 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -15,12 +15,12 @@ LogprobsTensors, ModelRunnerOutput, PoolerOutput, - RoutedExpertsTensors, ) from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors from vllm.v1.worker.utils import raise_if_nan_logits if TYPE_CHECKING: + from vllm.distributed.aux_output_connector.worker import PendingAuxOutput from vllm.v1.worker.gpu.input_batch import InputBatch @@ -120,8 +120,8 @@ def __init__( num_sampled_tokens: torch.Tensor, main_stream: torch.cuda.Stream, copy_stream: torch.cuda.Stream, - check_ep_fault: bool = False, - routed_experts: RoutedExpertsTensors | None = None, + check_ep_fault: bool, + pending_aux_output: "PendingAuxOutput | None", ): # NOTE(woosuk): We must retain references to the GPU tensors, # as the copy operations are performed on a different CUDA stream than @@ -129,7 +129,7 @@ def __init__( self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.routed_experts = routed_experts + self.pending_aux_output = pending_aux_output # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. self.copy_event = torch.cuda.Event(blocking=True) self._has_fault: torch.Tensor | None = None @@ -152,13 +152,15 @@ def __init__( self.sampling_mask_tensors = ( sampler_output.sampling_mask_tensors.to_cpu_nonblocking() ) - self.routed_experts_cpu: RoutedExpertsTensors | None = None - if routed_experts is not None: - self.routed_experts_cpu = routed_experts.to_cpu_nonblocking() self.prompt_logprobs_dict = { k: v.to_cpu_nonblocking() if v is not None else None for k, v in self.model_runner_output.prompt_logprobs_dict.items() } + if self.pending_aux_output is not None: + self.pending_aux_output.enqueue_cpu_copy( + num_sampled=self.num_sampled_tokens_np, + num_rejected=async_copy_to_np(sampler_output.num_rejected), + ) if check_ep_fault: has_fault = get_ep_all2all_manager().query_fault() self._has_fault = has_fault.to("cpu", non_blocking=True) @@ -192,8 +194,10 @@ def get_output(self) -> ModelRunnerOutput: if self.logprobs_tensors is not None: self.model_runner_output.logprobs = self.logprobs_tensors.tolists() self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict - if self.routed_experts_cpu is not None: - self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists() + if self.pending_aux_output is not None: + self.model_runner_output.aux_output_connector_output = ( + self.pending_aux_output.process_output() + ) if self._has_fault is not None and self._has_fault.item(): mask = get_ep_all2all_manager().query_active_mask() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8c736394cc52..1ba40c5fdae5 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -32,6 +32,10 @@ from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.aux_output_connector.worker import ( + AuxOutputWorkerConnector, + get_aux_output_connector, +) from vllm.distributed.parallel_state import ( get_dcp_group, get_pp_group, @@ -39,10 +43,6 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsCapturer, - bind_routed_experts_capturer, -) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, ) @@ -73,7 +73,6 @@ DraftTokenIds, ECConnectorOutput, ModelRunnerOutput, - RoutedExpertsTensors, ) from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler @@ -348,7 +347,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Expert parallelism load balancer. self.eplb = EPLBController(self.parallel_config, self.device) - self.routed_experts_capturer: RoutedExpertsCapturer | None = None + # The AuxOutput Connector owns R3 capture, copying, and storage. + self.aux_output_connector: AuxOutputWorkerConnector | None = None set_offloader(create_offloader(self.vllm_config.offload_config)) @@ -356,15 +356,6 @@ def update_max_model_len(self, max_model_len: int) -> None: self.max_model_len = max_model_len self.req_states.max_model_len = max_model_len - def init_routed_experts_capturer(self) -> None: - """Initialize target-model capture on every participating worker.""" - self.routed_experts_capturer = RoutedExpertsCapturer( - max_num_batched_tokens=self.max_num_tokens, - vllm_config=self.vllm_config, - kv_cache_config=self.kv_cache_config, - ) - bind_routed_experts_capturer(self.model, self.routed_experts_capturer) - def get_supported_tasks(self) -> tuple[SupportedTask, ...]: tasks: list[SupportedTask] = [] if self.model_config.runner_type == "generate": @@ -756,6 +747,12 @@ def initialize_kv_cache( else: self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) + # AuxOutput connector requires resolved kv_cache_config. + if self.vllm_config.aux_output_config.enabled: + self.aux_output_connector = get_aux_output_connector( + self.model, self.vllm_config, kv_cache_config + ) + def _init_kv_zero_meta(self) -> None: """Build KV-block zeroing metadata; invoked from gpu_worker.""" self.kv_block_zeroer = KVBlockZeroer( @@ -1648,6 +1645,11 @@ def execute_model( self.add_requests(scheduler_output) self.update_requests(scheduler_output) self.block_tables.apply_staged_writes() + if self.aux_output_connector is not None: + # Register this step before the GPU forward. + self.aux_output_connector.begin_step( + scheduler_output.aux_output_connector_metadata + ) if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) @@ -1969,11 +1971,6 @@ def execute_model( aux_hidden_states = None output_intermediate_tensors = model_output - routed_experts = None - if not dummy_run and (capturer := self.routed_experts_capturer) is not None: - assert slot_mappings is not None - routed_experts = capturer.get_routed_experts(slot_mappings, num_toks) - finished_req_ids = scheduler_output.finished_req_ids self.execute_model_state = ExecuteModelState( input_batch=input_batch, @@ -1984,7 +1981,6 @@ def execute_model( dp_sync=dp_sync, finished_req_ids=finished_req_ids, ec_connector_output=ec_connector_output, - routed_experts=routed_experts, cudagraph_stats=cudagraph_stats, ) @@ -2014,7 +2010,6 @@ def sample_tokens( dp_sync = self.execute_model_state.dp_sync finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output - routed_experts = self.execute_model_state.routed_experts cudagraph_stats = self.execute_model_state.cudagraph_stats self.execute_model_state = None @@ -2083,6 +2078,10 @@ def sample_tokens( prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type] cudagraph_stats=cudagraph_stats, ) + pending_aux_output = None + if self.aux_output_connector is not None: + pending_aux_output = self.aux_output_connector.prepare_output(input_batch) + # Start async output copy here so that it can overlap with speculator proposal. async_output = AsyncOutput( model_runner_output=model_runner_output, @@ -2091,7 +2090,7 @@ def sample_tokens( main_stream=self.main_stream, copy_stream=self.output_copy_stream, check_ep_fault=self.check_ep_fault, - routed_experts=routed_experts, + pending_aux_output=pending_aux_output, ) mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None @@ -2244,6 +2243,8 @@ def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() + if self.aux_output_connector is not None: + self.aux_output_connector.close() self.cudagraph_manager = None self.fast_prefill = None if hasattr(self, "kv_caches"): @@ -2316,7 +2317,6 @@ class ExecuteModelState(NamedTuple): dp_sync: DPSyncState | None finished_req_ids: set[str] ec_connector_output: ECConnectorOutput | None - routed_experts: RoutedExpertsTensors | None cudagraph_stats: CUDAGraphStat | None diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index b2803de9daeb..3d41c49720f2 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -617,8 +617,6 @@ def gather_sampler_output( device=device, ) if local_output is not None: - assert local_output.num_sampled is not None - assert local_output.num_rejected is not None assert not gather_num_nans or local_output.num_nans is not None num_src_cols = min( local_output.sampled_token_ids.shape[1], max_num_logits_per_req diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py index ff3675db130f..aedff5ce1873 100644 --- a/vllm/v1/worker/gpu/sample/output.py +++ b/vllm/v1/worker/gpu/sample/output.py @@ -17,8 +17,8 @@ class SamplerOutput: sampled_token_ids: torch.Tensor logprobs_tensors: LogprobsTensors | None num_nans: torch.Tensor | None - num_sampled: torch.Tensor | None - num_rejected: torch.Tensor | None = None + num_sampled: torch.Tensor + num_rejected: torch.Tensor sampling_mask_tensors: SamplingMaskTensors | None = None diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 7d933917bc53..bdabe398f676 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -61,10 +61,6 @@ from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsCapturer, - bind_routed_experts_capturer, -) from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, @@ -176,8 +172,6 @@ LogprobsTensors, ModelRunnerOutput, PoolerOutput, - RoutedExpertsLists, - RoutedExpertsTensors, SamplerOutput, make_empty_encoder_model_runner_output, ) @@ -293,7 +287,6 @@ def __init__( invalid_req_indices: list[int], async_output_copy_stream: torch.cuda.Stream, vocab_size: int, - routed_experts: RoutedExpertsTensors | None = None, check_ep_fault: bool = False, num_nans: torch.Tensor | None = None, ): @@ -309,7 +302,6 @@ def __init__( self._sampled_token_ids = sampled_token_ids self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors - self._routed_experts = routed_experts self._num_nans = num_nans self._has_fault: torch.Tensor | None = None @@ -325,11 +317,6 @@ def __init__( if self._logprobs_tensors else None ) - self._routed_experts_cpu = ( - self._routed_experts.to_cpu_nonblocking() - if self._routed_experts is not None - else None - ) self._num_nans_cpu = ( self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None @@ -370,10 +357,6 @@ def get_output(self) -> ModelRunnerOutput: output.sampled_token_ids = valid_sampled_token_ids output.logprobs = logprobs_lists - if self._routed_experts_cpu is not None: - output.routed_experts = self._routed_experts_cpu.tolists() - del self._routed_experts - if self._num_nans_cpu is not None: output.num_nans_in_logits = nans_to_dict( self._num_nans_cpu.tolist(), output.req_id_to_index @@ -537,9 +520,6 @@ def __init__( # These will be overridden in load_model() self.is_multimodal_pruning_enabled = False self.requires_sequential_video_encoding = False - # Set to True after init_routed_experts_capturer() completes. - # Prevents routed experts code from running during profiling/dummy run. - self.routed_experts_initialized = False self.max_model_len = model_config.max_model_len # Always set to false after the first forward pass @@ -2360,18 +2340,6 @@ def _get_block_table(kv_cache_gid: int): block_table_gid_0 = _get_block_table(0) slot_mapping_gid_0 = slot_mappings[0] - if self.routed_experts_initialized: - # Copy this step's attention slot_mapping into our private - # device buffer. The shared ``slot_mappings[attn_gid]`` is - # owned by the attention block table and will be overwritten - # by the next ``_prepare_inputs``; we need a stable snapshot - # because the async D2H may still be in flight on the copy - # stream when the next step runs. - slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid] - self.routed_experts_slot_mapping_device[:num_tokens].copy_( - slot_mapping_attn[:num_tokens] - ) - num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[ :num_reqs_padded ] @@ -3745,20 +3713,6 @@ def _bookkeeping_sync( invalid_req_indices = [] logprobs_lists = None if not self.use_async_scheduling: - # Sync scheduling: issue routed experts D2H into the pinned - # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does - # ``event.synchronize()`` on the async copy stream which - # waits for every D2H queued on the default stream since - # the last sync, so this enqueue is naturally covered - # without requiring its own synchronize. - if self.routed_experts_initialized: - buf = self.routed_experts_capturer.get_device_buffer() - total = scheduler_output.total_num_scheduled_tokens - self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) - self.routed_experts_slot_mapping_cpu[:total].copy_( - self.routed_experts_slot_mapping_device[:total], - non_blocking=True, - ) with gpu_sync_allowed(): # Get the valid generated tokens. max_gen_len = sampled_token_ids.shape[-1] @@ -4802,39 +4756,14 @@ def propose_draft_token_ids(sampled_token_ids): else None, num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, - routed_experts=None, ) if not self.use_async_scheduling: - if self.routed_experts_initialized: - # Sync path: D2H was issued in ``_bookkeeping_sync`` and - # synchronized by ``_to_list``'s event.synchronize(), so - # the pinned buffers are ready to be wrapped as numpy. - total = scheduler_output.total_num_scheduled_tokens - output.routed_experts = RoutedExpertsLists( - routing_data=self.routed_experts_cpu[:total].numpy(), - slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), - ) return output with record_function_or_nullcontext( "gpu_model_runner: AsyncGPUModelRunnerOutput" ): - # Async path: produce a device-side snapshot that the async - # copy stream can D2H later. Both tensors must be private - # clones because: - # - ``routing_data`` source is the shared capturer buffer, - # which the next forward overwrites on the default stream. - # - ``slot_mapping`` source is our own - # ``routed_experts_slot_mapping_device``, which the - # next ``_prepare_inputs`` overwrites on the default - # stream while the D2H is still pending on the copy - # stream. - # Without clones, the copy stream would read torn data. - routed_experts_snapshot = self.get_routed_experts( - scheduler_output.total_num_scheduled_tokens - ) - async_output = AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=sampler_output.sampled_token_ids, @@ -4842,7 +4771,6 @@ def propose_draft_token_ids(sampled_token_ids): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self._get_or_create_async_output_copy_stream(), vocab_size=self.input_batch.vocab_size, - routed_experts=routed_experts_snapshot, check_ep_fault=self.check_ep_fault, num_nans=num_nans_device, ) @@ -7468,61 +7396,6 @@ def initialize_kv_cache( kv_transfer_group.register_kv_caches(kv_caches) kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks) - def get_routed_experts( - self, - num_tokens: int, - ) -> RoutedExpertsTensors | None: - if not self.routed_experts_initialized: - return None - - device_buffer = self.routed_experts_capturer.get_device_buffer() - return RoutedExpertsTensors( - routing_data=device_buffer[:num_tokens].clone(), - slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(), - ) - - def init_routed_experts_capturer(self): - logger.info( - "Initializing routed experts capturer, enable_return_routed_experts: %s", - self.model_config.enable_return_routed_experts, - ) - self.routed_experts_capturer = RoutedExpertsCapturer( - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, - vllm_config=self.vllm_config, - kv_cache_config=self.kv_cache_config, - ) - bind_routed_experts_capturer(self.model, self.routed_experts_capturer) - - # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on - # the sync scheduling path. Shape / dtype mirror the device - # capturer exactly so ``copy_`` is a straight memcpy. - self.routed_experts_cpu = torch.empty( - self.routed_experts_capturer.device_buffer.shape, - dtype=self.routed_experts_capturer.device_buffer.dtype, - device="cpu", - pin_memory=PIN_MEMORY, - ) - # ``slot_mapping`` dtype is fixed to int64 by - # ``block_table.slot_mapping``; we mirror that here. - max_tokens = self.scheduler_config.max_num_batched_tokens - self.routed_experts_slot_mapping_cpu = torch.empty( - (max_tokens,), - dtype=torch.int64, - device="cpu", - pin_memory=PIN_MEMORY, - ) - # Private device buffer so the shared ``block_table.slot_mapping`` - # can be overwritten by the next ``_prepare_inputs`` while the - # D2H is still pending on the copy stream. Written in - # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path) - # or cloned into a snapshot (async path). - self.routed_experts_slot_mapping_device = torch.empty( - (max_tokens,), - dtype=torch.int64, - device=self.device, - ) - self.routed_experts_initialized = True - def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: """Add encoder-only layers to the KV cache config.""" block_size = self.vllm_config.cache_config.block_size diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d914bd2deab3..b6bc122a28b7 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -761,9 +761,6 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: ), ) - if self.model_config.enable_return_routed_experts: - self.model_runner.init_routed_experts_capturer() - # Build KV-zero metadata outside the CuMem pool so the bookkeeping # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch # allocator and are not discarded during sleep/wake cycles. From 04c1f4a407962a37d01b573f3a2aba91dd2ed3a9 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:29:14 -0400 Subject: [PATCH 577/774] [Bugfix][ROCm][DSv4.1] Disable SWA bounded replay on ROCm (#57906) Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- vllm/models/deepseek_v41/attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index bf645a8f156b..1e25704dd8b7 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -531,6 +531,15 @@ def __init__( "instead." ) swa_bounded_replay = False + if swa_bounded_replay and current_platform.is_rocm(): + logger.warning_once( + "SWA bounded replay is off on ROCm (the sparse SWA metadata " + "builders forward replay_start, but the window clamp it relies " + "on lives in the FlashInfer and FlashMLA prefill kernels, so " + "the padded slots fault); the sliding-window cache takes part " + "in prefix caching instead." + ) + swa_bounded_replay = False self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, window_size=self.window_size, From 15859bb3a1d81a709b64eff2a1d1f38a958a362b Mon Sep 17 00:00:00 2001 From: CaiJohn Date: Mon, 21 Sep 2026 18:56:28 +0800 Subject: [PATCH 578/774] [Test] Organize structured output utility tests (#57545) Signed-off-by: john <6712432+CaiJohn@users.noreply.github.com> Co-authored-by: john <6712432+CaiJohn@users.noreply.github.com> --- tests/v1/structured_output/test_utils.py | 89 ++++++++++++------------ 1 file changed, 44 insertions(+), 45 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index ff2b58dfe194..d3187c428cc0 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -190,41 +190,39 @@ def pattern_properties_schema(): } -@pytest.mark.parametrize( - "schema_type", - [ - "unsupported_string_schemas", - "unsupported_integer_schemas", - "unsupported_number_schemas", - "unsupported_array_schemas", - "unsupported_property_names_combinations", - "unsupported_pattern_properties_combinations", - ], -) -def test_unsupported_json_features_by_type(schema_type, request): - schemas = request.getfixturevalue(schema_type) - for schema in schemas: - assert has_xgrammar_unsupported_json_features(schema), ( - f"Schema should be unsupported: {schema}" - ) - - -@pytest.mark.parametrize( - "schema_type", - [ - "supported_frankenstein_schema", - "pattern_properties_schema", - "property_names_schema", - ], -) -def test_supported_json_features(schema_type, request): - schema = request.getfixturevalue(schema_type) - assert not has_xgrammar_unsupported_json_features(schema), ( - f"Schema should be supported: {schema}" +class TestHasXGrammarUnsupportedJsonFeatures: + @pytest.mark.parametrize( + "schema_type", + [ + "unsupported_string_schemas", + "unsupported_integer_schemas", + "unsupported_number_schemas", + "unsupported_array_schemas", + "unsupported_property_names_combinations", + "unsupported_pattern_properties_combinations", + ], ) + def test_unsupported_json_features_by_type(self, schema_type, request): + schemas = request.getfixturevalue(schema_type) + for schema in schemas: + assert has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be unsupported: {schema}" + ) + @pytest.mark.parametrize( + "schema_type", + [ + "supported_frankenstein_schema", + "pattern_properties_schema", + "property_names_schema", + ], + ) + def test_supported_json_features(self, schema_type, request): + schema = request.getfixturevalue(schema_type) + assert not has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be supported: {schema}" + ) -class TestHasXGrammarUnsupportedJsonFeatures: class TestPR48416Regressions: @pytest.mark.parametrize( "schema", @@ -270,20 +268,21 @@ def test_supported_list_type_json_features(self, schema): assert not has_xgrammar_unsupported_json_features(schema) -def test_property_names_constrains_keys(property_names_schema): - assert grammar_accepts(property_names_schema, '{"score": 5}') - assert grammar_accepts(property_names_schema, '{"score": "seven"}') - assert not grammar_accepts(property_names_schema, '{"Score": 5}') - assert not grammar_accepts(property_names_schema, '{"score_1": 5}') - - -def test_pattern_properties_constrains_keys_and_values(pattern_properties_schema): - assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') - assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') - assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') - - class TestIsGrammarAcceptString: + class TestPR42904Support: + def test_property_names_constrains_keys(self, property_names_schema): + assert grammar_accepts(property_names_schema, '{"score": 5}') + assert grammar_accepts(property_names_schema, '{"score": "seven"}') + assert not grammar_accepts(property_names_schema, '{"Score": 5}') + assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + + def test_pattern_properties_constrains_keys_and_values( + self, pattern_properties_schema + ): + assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') + class TestPR48115Regressions: @pytest.mark.parametrize("codepoint", [*range(0x20), 0x7F]) def test_choice_as_grammar_preserves_control_characters(self, codepoint): From 82daf9f5756e1868be0aa751afaec4726beca12a Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 06:08:06 -0500 Subject: [PATCH 579/774] [CI][ROCm] Add ten AMD parity groups (#57876) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test-amd.yaml | 239 ++++++++++++++++++ .buildkite/test_areas/disaggregated.yaml | 21 ++ .buildkite/test_areas/disaggregated_ec.yaml | 19 ++ .../test_areas/disaggregated_mooncake.yaml | 42 +++ .buildkite/test_areas/distributed.yaml | 23 ++ .buildkite/test_areas/kernels.yaml | 17 ++ .buildkite/test_areas/lm_eval.yaml | 94 ++++++- .buildkite/test_areas/misc.yaml | 21 ++ .../Qwen3-30B-A3B-BF16-triton-rocm.yaml | 5 + .../moe-refactor/config-rocm-shard-0.txt | 1 + .../moe-refactor/config-rocm-shard-1.txt | 1 + .../moe-refactor/config-rocm-shard-2.txt | 1 + .../moe-refactor/config-rocm-shard-3.txt | 1 + .../configs/Qwen3.8-Flash-Next-FP8-rocm.yaml | 26 ++ tests/evals/qwen4_exp/configs/models-rocm.txt | 1 + ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 16 +- 16 files changed, 525 insertions(+), 3 deletions(-) create mode 100644 tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml create mode 100644 tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt create mode 100644 tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt create mode 100644 tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt create mode 100644 tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt create mode 100644 tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml create mode 100644 tests/evals/qwen4_exp/configs/models-rocm.txt diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 3211840c9b06..9a2116a16d0a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1361,6 +1361,23 @@ steps: - bash ../tools/install_torchcodec_rocm.sh || exit 1 - pytest -s entrypoints/openai/correctness/ +- label: ":amd: (MI300) Scale-out EC Connector E2E" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ec_transfer/ + - vllm/entrypoints/scale_out/ + - tests/entrypoints/scale_out/ec_integration/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh + #----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# - label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD @@ -3190,6 +3207,58 @@ steps: v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism +- label: ":amd: (MI300) Mooncake EC TCP E2E" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/distributed/ec_transfer/ + - vllm/config/ec_transfer.py + - vllm/config/multimodal.py + - vllm/config/vllm.py + - vllm/multimodal/ + - vllm/v1/core/encoder_cache_manager.py + - vllm/v1/core/sched/ + - vllm/v1/engine/ + - vllm/v1/worker/ec_connector_model_runner_mixin.py + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/worker/gpu_worker.py + - vllm/v1/worker/gpu/ec_connector.py + - vllm/v1/worker/gpu/model_runner.py + - vllm/v1/worker/mm_encoder_model_runner.py + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py + - vllm/v1/attention/ops/vit_attn_wrappers.py + - vllm/v1/worker/gpu/model_states/ + - vllm/v1/worker/gpu/mm/ + - tests/v1/cudagraph/test_encoder_cudagraph.py + - tests/models/multimodal/generation/test_vit_cudagraph.py + - examples/disaggregated/disaggregated_encoder/ + - tests/v1/ec_connector/ + - requirements/kv_connectors.txt + - vllm/platforms/rocm.py + - requirements/rocm.txt + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTHON_BIN='python' + - export PYTHONUNBUFFERED='1' + - export VLLM_HOST_IP='127.0.0.1' + - export VLLM_USE_V2_MODEL_RUNNER='1' + - export E_CUDAGRAPH_MM_ENCODER='1' + - export E_MM_ENCODER_ATTN_BACKEND='TRITON_ATTN' + - export MOONCAKE_EC_PROTOCOL='tcp' + - export USE_MM_PROMPTS='1' + - export SKIP_BASELINE='0' + - export CONCURRENCY='3' + - export REPEAT='2' + - export LOG_PATH='/tmp/mooncake-ec-e2e' + - export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json' + - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh + - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3385,6 +3454,27 @@ steps: #-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# +- label: ":amd: (MI355) Sharded RDT Weight Transfer" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/weight_transfer/ + - tests/distributed/test_sharded_rdt_plan.py + - tests/distributed/test_sharded_rdt_producer.py + - tests/distributed/test_sharded_rdt_trainer.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - uv pip install --system --no-deps 'ray==2.56.1' + - pytest -v -s distributed/test_sharded_rdt_plan.py + - pytest -v -s distributed/test_sharded_rdt_producer.py + - pytest -v -s distributed/test_sharded_rdt_trainer.py + - label: ":amd: (MI355) Distributed DP Basic" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3464,6 +3554,35 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" + timeout_in_minutes: 105 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/platforms/rocm.py + - vllm/models/kimi_k3/ + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - requirements/kv_connectors_rocm.txt + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export ENABLE_HMA_FLAG='1' + - export DP_EP='1' + - export GPU_MEMORY_UTILIZATION='0.9' + - export VLLM_SSM_CONV_STATE_LAYOUT='DS' + - export PREFILLER_TP_SIZE='2' + - export DECODER_TP_SIZE='2' + - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' + - export PREFILL_BLOCK_SIZE='2048' + - export DECODE_BLOCK_SIZE='2048' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768' + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh + - label: ":amd: (MI355) Distributed Features" # TBD timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3661,6 +3780,51 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt +- label: ":amd: (MI355) LM Eval Watermarking" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt + +- label: ":amd: (MI355) LM Eval Watermark Feature Combination" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/speculative.py + - vllm/config/vllm.py + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/warmup/ + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt + - label: ":amd: (MI355) LM Eval Qwen3-5 Models" # TBD timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3684,6 +3848,27 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt +- label: ":amd: (MI355) Qwen3-8-Flash-Next-FP8 Accuracy Eval" + timeout_in_minutes: 105 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/libtorch_stable/gdn/ + - tests/evals/qwen4_exp/ + - vllm/model_executor/layers/mamba/ + - vllm/models/qwen4_exp/ + - vllm/transformers_utils/configs/qwen4_exp.py + - vllm/v1/attention/backends/short_conv_attn.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - uv pip install --system 'evalscope==1.10.0' + - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt + - label: ":amd: (MI355) LM Eval Small Models Distributed" timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3705,6 +3890,25 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt +- label: ":amd: (MI355) MoE Refactor Integration Shard %N" + timeout_in_minutes: 105 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + parallelism: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/platforms/rocm.py + - tests/evals/gsm8k/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/_aiter_ops.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt + - label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD timeout_in_minutes: 25 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3925,6 +4129,22 @@ steps: - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py +- label: ":amd: (MI355) DeepSeek V4 Fused Kernels" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - vllm/models/deepseek_v4/common/ops/ + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - label: ":amd: (MI355) Attention Kernels Shard %N" # TBD timeout_in_minutes: 90 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4418,6 +4638,25 @@ steps: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async +- label: ":amd: (MI355) Extract Hidden States Integration Single GPU" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/v1/spec_decode/extract_hidden_states.py + - vllm/model_executor/models/extract_hidden_states.py + - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py + - tests/v1/kv_connector/extract_hidden_states_integration + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration + - label: ":amd: (MI355) V1 Attention Shard %N" # TBD timeout_in_minutes: 125 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b187169e7675..811301f58d91 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -318,3 +318,24 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh + mirror: + amd: + label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" + dind: false + device: mi355_4 + num_devices: 4 + timeout_in_minutes: 105 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py + - vllm/models/kimi_k3/ + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - requirements/kv_connectors_rocm.txt + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated_ec.yaml b/.buildkite/test_areas/disaggregated_ec.yaml index 9100ad0a431e..67bfebaf8cac 100644 --- a/.buildkite/test_areas/disaggregated_ec.yaml +++ b/.buildkite/test_areas/disaggregated_ec.yaml @@ -14,3 +14,22 @@ steps: - tests/entrypoints/scale_out/ec_integration/ commands: - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh + mirror: + amd: + label: ":amd: (MI300) Scale-out EC Connector E2E" + dind: false + device: mi300_2 + num_devices: 2 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ec_transfer/ + - vllm/entrypoints/scale_out/ + - tests/entrypoints/scale_out/ec_integration/ + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index f6e944d70581..d89996c9a215 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.buildkite/test_areas/disaggregated_mooncake.yaml @@ -83,3 +83,45 @@ steps: PY )" - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh + mirror: + amd: + label: ":amd: (MI300) Mooncake EC TCP E2E" + dind: false + device: mi300_2 + num_devices: 2 + timeout_in_minutes: 60 + working_dir: /vllm-workspace + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ec_transfer/ + - vllm/config/ec_transfer.py + - vllm/config/multimodal.py + - vllm/config/vllm.py + - vllm/multimodal/ + - vllm/v1/core/encoder_cache_manager.py + - vllm/v1/core/sched/ + - vllm/v1/engine/ + - vllm/v1/worker/ec_connector_model_runner_mixin.py + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/worker/gpu_worker.py + - vllm/v1/worker/gpu/ec_connector.py + - vllm/v1/worker/gpu/model_runner.py + - vllm/v1/worker/mm_encoder_model_runner.py + - vllm/v1/worker/encoder_cudagraph.py + - vllm/v1/worker/encoder_cudagraph_defs.py + - vllm/v1/attention/ops/vit_attn_wrappers.py + - vllm/v1/worker/gpu/model_states/ + - vllm/v1/worker/gpu/mm/ + - tests/v1/cudagraph/test_encoder_cudagraph.py + - tests/models/multimodal/generation/test_vit_cudagraph.py + - examples/disaggregated/disaggregated_encoder/ + - tests/v1/ec_connector/ + - requirements/kv_connectors.txt + - vllm/platforms/rocm.py + - requirements/rocm.txt + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + PYTHON_BIN: "python" + commands: + - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index a10824656395..953d9a9f8bc8 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -33,6 +33,29 @@ steps: - pytest -v -s distributed/test_sharded_rdt_plan.py - pytest -v -s distributed/test_sharded_rdt_producer.py - pytest -v -s distributed/test_sharded_rdt_trainer.py + mirror: + amd: + label: ":amd: (MI355) Sharded RDT Weight Transfer" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/weight_transfer/ + - tests/distributed/test_sharded_rdt_plan.py + - tests/distributed/test_sharded_rdt_producer.py + - tests/distributed/test_sharded_rdt_trainer.py + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - uv pip install --system --no-deps 'ray==2.56.1' + - pytest -v -s distributed/test_sharded_rdt_plan.py + - pytest -v -s distributed/test_sharded_rdt_producer.py + - pytest -v -s distributed/test_sharded_rdt_trainer.py - label: ":nvidia: (L4) Distributed DP Basic" key: distributed-dp-tests-2-gpus diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 2134d6dd040b..78d57ca4559f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -111,6 +111,23 @@ steps: commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py - pytest -v -s kernels/test_top_k_per_row.py + mirror: + amd: + label: ":amd: (MI355) DeepSeek V4 Fused Kernels" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - vllm/models/deepseek_v4/common/ops/ + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - label: ":nvidia: (B200) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-b200 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 56aeca2a8eb3..5a284d172ab3 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -66,6 +66,28 @@ steps: autorun_on_main: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt + mirror: + amd: + label: ":amd: (MI355) LM Eval Watermarking" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt # - label: LM Eval Large Models (4xA100) # key: lm-eval-large-models-4xa100 @@ -203,6 +225,33 @@ steps: autorun_on_main: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt + mirror: + amd: + label: ":amd: (MI355) LM Eval Watermark Feature Combination" + dind: false + device: mi355_2 + num_devices: 2 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/speculative.py + - vllm/config/vllm.py + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/warmup/ + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt - label: ":nvidia: (B200) LM Eval Spec Decode" key: lm-eval-spec-decode-4xb200 @@ -329,7 +378,27 @@ steps: # build 83851): shards ~12.4/13.6/12.4/14.5m of eval time. New configs must # be added to exactly one config-h100-shard-*.txt. - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt - + mirror: + amd: + label: ":amd: (MI355) MoE Refactor Integration Shard %N" + dind: false + device: mi355_2 + num_devices: 2 + timeout_in_minutes: 105 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py + - tests/evals/gsm8k/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/_aiter_ops.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt + - label: ":nvidia: (B200) MoE Refactor Integration TEMPORARY Shard %N" key: moe-refactor-integration-test-b200-temporary device: b200-k8s @@ -578,6 +647,29 @@ steps: commands: - uv pip install --system 'evalscope==1.10.0' - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt + mirror: + amd: + label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" + dind: false + device: mi355_4 + num_devices: 4 + timeout_in_minutes: 105 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/libtorch_stable/gdn/ + - tests/evals/qwen4_exp/ + - vllm/model_executor/layers/mamba/ + - vllm/models/qwen4_exp/ + - vllm/transformers_utils/configs/qwen4_exp.py + - vllm/v1/attention/backends/short_conv_attn.py + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - uv pip install --system 'evalscope==1.10.0' + - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt - label: ":nvidia: (DGX) Spark GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-spark diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index ef096e9bd338..759f3900c5ea 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -237,6 +237,27 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration + mirror: + amd: + label: ":amd: (MI355) Extract Hidden States Integration Single GPU" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/extract_hidden_states.py + - vllm/model_executor/models/extract_hidden_states.py + - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py + - tests/v1/kv_connector/extract_hidden_states_integration + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration - label: ":nvidia: (L4) Extract Hidden States Integration" key: extract-hidden-states-integration-2-gpus diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml new file mode 100644 index 000000000000..c6da3748b8b8 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.88 +num_questions: 1319 +num_fewshot: 5 +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton" diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt new file mode 100644 index 000000000000..4ca4943fc78f --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt @@ -0,0 +1 @@ +Qwen3-30B-A3B-BF16-triton-rocm.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt new file mode 100644 index 000000000000..c46cb2c243fa --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt @@ -0,0 +1 @@ +Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt new file mode 100644 index 000000000000..928653a58039 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt @@ -0,0 +1 @@ +Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt new file mode 100644 index 000000000000..8b175f5b0ef1 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt @@ -0,0 +1 @@ +Mixtral-8x7B-Fp8-AutoFp8-triton.yaml diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml new file mode 100644 index 000000000000..db43869eb4ae --- /dev/null +++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "Qwen/Qwen3.8-Flash-Next-FP8" +datasets: + gsm8k: + metric_threshold: 0.98 + tolerance: 0.02 + aime25: + metric_threshold: 0.90 + tolerance: 0.05 +eval_batch_size: 32 +startup_max_wait_seconds: 1800 +generation_config: + seed: 1236 + do_sample: true + temperature: 0.7 + top_p: 0.8 + top_k: 20 + repetition_penalty: 1.0 + presence_penalty: 1.5 + max_tokens: 32768 + stream: true + extra_body: + chat_template_kwargs: + enable_thinking: true +server_args: "--tensor-parallel-size 4 --mamba-cache-mode align" diff --git a/tests/evals/qwen4_exp/configs/models-rocm.txt b/tests/evals/qwen4_exp/configs/models-rocm.txt new file mode 100644 index 000000000000..9999574a4d43 --- /dev/null +++ b/tests/evals/qwen4_exp/configs/models-rocm.txt @@ -0,0 +1 @@ +Qwen3.8-Flash-Next-FP8-rocm.yaml diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index ced96bc40199..c41af1e4b9ca 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -96,7 +96,7 @@ def apply_rope_gptj_last_k( # `e*c - o*s` fused form. This keeps the reference close to the kernel, but # the fp32 reference and the fp32 GPU kernel can still round to bf16 on # opposite sides of a round-to-nearest tie for a tiny number of elements at - # high positions, so callers compare the RoPE region within 1 bf16 ULP. + # high positions, so callers allow for bf16 rounding in the RoPE region. new_even = torch.addcmul(-odd * sin, even, cos) new_odd = torch.addcmul(odd * cos, even, sin) rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape) @@ -314,7 +314,19 @@ def test_q_path_without_qnorm_matches_rope_only_reference(num_tokens: int): apply_q_norm=False, ) - torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=0, atol=0) + torch.testing.assert_close( + q_out[:, :n_heads, :NOPE_DIM], q_ref[..., :NOPE_DIM], rtol=0, atol=0 + ) + # HIP and PyTorch can contract different FP32 products into an FMA. + # Allow bf16 rounding and small cancellation errors only in the RoPE region + # on ROCm; the unrotated dimensions and padded heads must remain exact. + rope_rtol, rope_atol = (1e-2, 1e-6) if current_platform.is_rocm() else (0, 0) + torch.testing.assert_close( + q_out[:, :n_heads, NOPE_DIM:], + q_ref[..., NOPE_DIM:], + rtol=rope_rtol, + atol=rope_atol, + ) assert q_out[:, n_heads:].count_nonzero().item() == 0 From 7268f6e38ebbb3d61700f16c638c6eba0c88bb0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:50:48 +0200 Subject: [PATCH 580/774] [Security] Prefer fresh multimodal payloads over a stale receiver cache (#57833) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Cursor --- tests/multimodal/test_cache.py | 45 ++++++++++++++++++++++++++++++++++ vllm/multimodal/cache.py | 24 ++++++++++++------ 2 files changed, 61 insertions(+), 8 deletions(-) diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index 02bd5119700b..37ec5b0cb8c1 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -330,6 +330,51 @@ def test_mm_cache_miss_raises_and_recovers(): assert p1.get_and_update_item(None, mm_hash) == item +def test_receiver_cache_replaces_stale_item_when_payload_resent(): + """A P0 miss can resend a different item under the same identity. + + Independent LRU eviction leaves P1 holding the old tensor. Substituting it + for the new payload pairs the request's placeholders with the wrong item + and kills EngineCore. Prefer the fresh payload and keep it for later hits. + """ + model_config = _StubModelConfig(mm_processor_cache_gb=1) + p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + small = MultiModalKwargsItem.dummy(nbytes=64) + large = MultiModalKwargsItem.dummy(nbytes=256) + mm_hash = "shared-id" + + assert p1.get_and_update_item(small, mm_hash) is small + assert p1.get_and_update_item(large, mm_hash) is large + assert p1.get_and_update_item(None, mm_hash) is large + + +def test_receiver_cache_features_keep_resent_payload(): + """EngineCore updates features through get_and_update_features.""" + model_config = _StubModelConfig(mm_processor_cache_gb=1) + p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + small = MultiModalKwargsItem.dummy(nbytes=64) + large = MultiModalKwargsItem.dummy(nbytes=256) + mm_hash = "shared-id" + + def _feature( + data: MultiModalKwargsItem | None, length: int + ) -> MultiModalFeatureSpec: + return MultiModalFeatureSpec( + data=data, + modality="image", + identifier=mm_hash, + mm_position=PlaceholderRange(offset=0, length=length), + mm_hash=mm_hash, + ) + + seeded = p1.get_and_update_features([_feature(small, length=4)]) + assert seeded[0].data is small + + updated = p1.get_and_update_features([_feature(large, length=2048)]) + assert updated[0].data is large + assert updated[0].mm_position.length == 2048 + + def test_mm_cache_miss_batches_all_drifted_hashes(): """All hashes drifted within one request must surface in a single error. diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache.py index b4fabfc21ffb..15d433112cf2 100644 --- a/vllm/multimodal/cache.py +++ b/vllm/multimodal/cache.py @@ -710,9 +710,12 @@ class MultiModalReceiverCache(BaseMultiModalReceiverCache): How to update each item: - - If the item is in the cache, replace the input with the cached item. - - If the item is not in the cache, store that item (which includes tensor - data) into the cache, and return the input. + - If the caller sent tensor data, store it (replacing any cached item + under the same key) and return that data. P0 can miss after independent + LRU eviction and resend a different item for the same identity. + - If the caller sent no data and the item is cached, return the cached item. + - If the caller sent no data and the item is not cached, raise + ``MultiModalCacheMissError``. """ def __init__(self, model_config: "ModelConfig") -> None: @@ -731,17 +734,22 @@ def get_and_update_item( mm_item: MultiModalKwargsItem | None, mm_hash: str, ) -> MultiModalKwargsItem: + if mm_item is not None: + # P0 sent a payload. Never keep a stale cached tensor under this + # identity: a later P0 hit would pair new placeholders with the old + # item and crash the engine. Drop first so a too-large replacement + # cannot leave the previous entry in place. + self._cache.pop(mm_hash, None) + self.cache_if_fits(self._cache, mm_hash, mm_item) + return mm_item + if (cached_item := self._cache.get(mm_hash)) is not None: return cached_item # No data and not cached here: P0 sent data=None trusting its shadow, but # the P0/P1 caches have drifted. Raise a retryable error (not assert) so the # engine can have P0 drop the stale entry and the client resend the data. - if mm_item is None: - raise MultiModalCacheMissError([mm_hash]) - - self.cache_if_fits(self._cache, mm_hash, mm_item) - return mm_item + raise MultiModalCacheMissError([mm_hash]) @override def touch_receiver_cache_item( From db7f1f671419535b85fa73e29ff061e287e39459 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Mon, 21 Sep 2026 20:54:58 +0800 Subject: [PATCH 581/774] [CI][XPU] Deselect Ray UT in XPU V1 test Job (#57889) Signed-off-by: Chaojun Zhang Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index ef2b4ba07396..d5ae40ec6e02 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -106,7 +106,7 @@ steps: pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py' + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --deselect "v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]"' parallelism: 2 - label: "XPU server test" key: xpu-server-test From 606d124b377e6462fe0f9076df4d50c88e4dcf48 Mon Sep 17 00:00:00 2001 From: Hotragn Pettugani <103170876+Hotragn@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:27:40 -0400 Subject: [PATCH 582/774] [Bugfix][Multimodal] Tolerate malformed EXIF in Molmo 2 image preprocessing (#57234) Signed-off-by: hotragn Co-authored-by: Claude Opus 5 --- vllm/model_executor/models/molmo2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 881be146e790..b44419a42ab6 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -11,7 +11,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from PIL import ImageOps from PIL.Image import Image from transformers import ( BaseImageProcessor, @@ -54,6 +53,7 @@ ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.image import normalize_image from vllm.multimodal.inputs import ( MultiModalFieldConfig, MultiModalKwargsItems, @@ -1299,7 +1299,7 @@ def exif_transpose( exif_transpose(img) if isinstance(img, Image) else img for img in images ] elif images is not None and isinstance(images, Image): - images = ImageOps.exif_transpose(images) + images = normalize_image(images) return images From 8b98b7d0b465920740cfa1acf5dde1b6a57d8d52 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 21 Sep 2026 15:13:52 +0100 Subject: [PATCH 583/774] [Frontend] Add streaming parity tests and docs for derender (#57922) Signed-off-by: Martin Hickey --- docs/serving/online_serving/derenderer.md | 111 +++++++++- .../derender/test_derender_parity.py | 189 +++++++++++++++++- .../derender/test_derender_stream.py | 104 ++-------- tests/entrypoints/scale_out/derender/utils.py | 103 ++++++++++ .../scale_out/token_in_token_out/protocol.py | 4 + 5 files changed, 410 insertions(+), 101 deletions(-) create mode 100644 tests/entrypoints/scale_out/derender/utils.py diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index 773131d03379..0e865796f81d 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -6,7 +6,7 @@ This closes the loop for a token-in / token-out engine in disaggregated serving: - **GPU less post processing**: Detokenization, reasoning parsing, and tool call parsing run on the same GPU less frontend that hosts `/render` - **Parser parity**: The derenderer reuses vLLM's tool and reasoning parsers, so a disaggregated deployment produces the same `content`/`reasoning`/ `tool_calls` split as a standard `vllm serve` server -- **Streaming and one-shot parsing**: Non-streaming calls send a complete `GenerateResponse` with all token IDs present and get one-shot parsing. Both endpoints also accept `stream: true`, taking one `GenerateStreamResponse` delta plus a client carried `stream_state` and returning `{chunk, stream_state}`. The chat endpoint's streaming path supports reasoning and tool call parsing while emitting the same `reasoning`/`content`/`tool_calls` deltas the generate streaming path would for the same token IDs +- **Streaming and one-shot parsing**: Non-streaming calls send a complete `GenerateResponse` with all token IDs present and get one-shot parsing. Both endpoints also accept `stream: true`, taking one `GenerateStreamResponse` delta plus a client carried `stream_state` and returning `{chunk, stream_state}`. The chat endpoint's streaming path supports reasoning and tool call parsing while emitting the same `reasoning`/`content`/`tool_calls` deltas the generate streaming path would for the same token IDs (see [Streaming](#streaming)) Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render` endpoints. @@ -27,8 +27,10 @@ The derender step needs more than the engine's `token_ids`. It also consumes the - Chat Completions Derender API (`/v1/chat/completions/derender`) - Post process a single `GenerateResponse` into a `ChatCompletionResponse` + - With `stream: true`, post process one `GenerateStreamResponse` chunk into a `ChatCompletionStreamResponse` chunk - Completions Derender API (`/v1/completions/derender`) - Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse` + - With `stream: true`, post process one `GenerateStreamResponse` chunk into a `CompletionStreamResponse` chunk ## Request format @@ -50,8 +52,54 @@ Each request wraps the engine's `GenerateResponse`(s) together with the caller m --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request" ``` +Streaming requests set `stream: true` and carry one generate chunk (`generate_chunk`) plus `stream_state` instead of a complete response. + +`/v1/chat/completions/derender` with `stream: true`: + +??? code + + ```python + --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-stream-request" + ``` + +`/v1/completions/derender` with `stream: true`: + +??? code + + ```python + --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-stream-request" + ``` + +Both return `{"chunk": ..., "stream_state": ...}`. `chunk` is a `ChatCompletionStreamResponse` or `CompletionStreamResponse` and `stream_state` goes with the next call. + Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs. +## Parser configuration + +The derenderer builds its tool and reasoning parsers from its own server flags plus the `chat_request` sent with each call. It can't see how the request was rendered or served anywhere else, so for its output to match a `vllm serve` server: + +- Start the render server with the same `--tool-call-parser`, `--reasoning-parser`, `--enable-auto-tool-choice`, `--chat-template` and `--default-chat-template-kwargs` you'd give `vllm serve` for this model. If `/render` and `/derender` run on different servers, give both the same values. +- Send the full `chat_request` that went to `/render`, not just `messages` and `tools`. Fields like `chat_template_kwargs`, `reasoning_effort`, `tool_choice` and `include_reasoning` change how the output is parsed. + +A mismatch doesn't fail. The parser just splits `reasoning`, `content` and `tool_calls` differently from what `vllm serve` would return. + +`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. It only reads `skip_special_tokens` from `completion_request`. + +## Streaming + +Streaming derender mirrors `/inference/v1/generate`. Set `stream: true` in the body on the same path and the endpoint takes one generate stream chunk instead of a complete response. It answers with JSON, not SSE. Each call turns one generate chunk into one derendered chunk and the client forwards that to its own caller. + +The server keeps no state between calls. Everything the next call needs is in `stream_state` which the client carries: + +- **Echo `stream_state` back.** Leave it out (or send `null`) on the first call, then send the `stream_state` from each response with the next request, unchanged. State that doesn't add up (e.g. `output_chunk_lens` that don't sum to the number of output tokens) is rejected with a `400`. +- **One choice per chunk.** `/inference/v1/generate` sends one choice per SSE event and each derender call accepts at most one. More than one is rejected with a `400`. +- **`n > 1` is N streams.** Generate interleaves chunks for different choices and tags each with its `index`. Keep one `stream_state` per index and send each chunk with the state for its index. Each index gets its own `role` delta on its first chunk. +- **Send every chunk through, including the last ones.** The finish reason usually arrives with the final tokens. A chunk with a `finish_reason` and no tokens is accepted too and flushes any buffered tool call arguments. With `stream_options: {"include_usage": true}`, generate ends with a usage only chunk (`choices: []`). Send it through derender like any other chunk to get the usage chunk for the response. `usage.prompt_tokens` is the request's `prompt_tokens` if set, otherwise the generate chunk's. +- **Send the same context on every call.** `chat_request` and `prompt_token_ids` aren't kept between calls, so they go with every chunk including the usage chunk. Both are required when a tool or reasoning parser is configured. `prompt_token_ids` is the `token_ids` of the `GenerateRequest` returned by `/render`. +- **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk. + +Streaming chunks don't carry logprobs yet, so `logprobs` on a generate chunk are dropped. The non-streaming endpoints do resolve them, including `token_id:N` placeholders. + ## Streaming cost Streaming derender threads a client carried `stream_state` across per-chunk calls instead of keeping session state on the server. @@ -118,3 +166,64 @@ print(response["choices"][0]["message"]["content"]) ``` Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. `chat_request` can only be omitted for a model with no tool or reasoning parser configured. A parser configured model rejects a missing `chat_request` with a 400 rather than silently falling back to plain detokenization. + +### Streaming example + +The same round trip, streamed. `/render` keeps `stream` and `stream_options` on the `GenerateRequest` it returns, so the generate call streams too. Each generate chunk goes through derender with the `stream_state` from the previous call (see [Streaming](#streaming)). + +```python +import json + +import httpx + +MODEL = "meta-llama/Llama-3.2-1B-Instruct" +RENDER = "http://localhost:8100" # vllm launch render ... +ENGINE = "http://localhost:8200" # token-in / token-out engine + +chat_request = { + "model": MODEL, + "messages": [{"role": "user", "content": "What is 2+2?"}], + "max_tokens": 32, + "stream": True, + "stream_options": {"include_usage": True}, +} + +with httpx.Client(timeout=60.0) as client: + # 1. Render: request -> token IDs (GPU less) + generate_request = client.post( + f"{RENDER}/v1/chat/completions/render", json=chat_request + ).json() + prompt_token_ids = generate_request["token_ids"] + + # 2. Generate: stream token IDs (token-in / token-out engine) + stream_state = None # one per choice index when n > 1 + with client.stream( + "POST", f"{ENGINE}/inference/v1/generate", json=generate_request + ) as generate_stream: + for line in generate_stream.iter_lines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + + # 3. Derender: one generate chunk -> one chat completion chunk + derendered = client.post( + f"{RENDER}/v1/chat/completions/derender", + json={ + "stream": True, + "model": MODEL, + "generate_chunk": json.loads(line[len("data: ") :]), + "stream_state": stream_state, + "prompt_tokens": len(prompt_token_ids), + "prompt_token_ids": prompt_token_ids, + "chat_request": chat_request, + }, + ).json() + stream_state = derendered["stream_state"] + + chunk = derendered["chunk"] + for choice in chunk["choices"]: + print(choice["delta"].get("content") or "", end="", flush=True) + if chunk.get("usage"): + print(f"\n{chunk['usage']}") +``` + +With a tool or reasoning parser configured, `delta` also carries `reasoning` and `tool_calls`, the same deltas `/v1/chat/completions` streams for those tokens. diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index adece09c0a1a..a89f8d68db53 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -3,20 +3,25 @@ """Round trip parity CI: render -> generate -> derender -> render == render. -Pins the coupled parsing path (``/v1/chat/completions`` on a normal GPU +Pins the coupled parsing path (`/v1/chat/completions` on a normal GPU server which parses generated tokens incrementally as they stream out) -against the disaggregated path (``OnlineDerenderer.derender_chat`` via -``/v1/chat/completions/derender`` which parses the same tokens all at once, -out of process). A standard ``vllm serve`` GPU server mounts both, so one +against the disaggregated path (`OnlineDerenderer.derender_chat` via +`/v1/chat/completions/derender` which parses the same tokens all at once, +out of process). A standard `vllm serve` GPU server mounts both, so one real generation lets both parsers run on identical input. Both paths consume the same generated token IDs (extracted from the coupled -response's ``token_ids`` via ``return_token_ids=True``), so generation +response's `token_ids` via `return_token_ids=True`), so generation nondeterminism is irrelevant. The only variable under test is whether the two parsing code paths agree. Parity is asserted unconditionally as only the stronger per case assertions (e.g. "a tool call was produced") are gated behind the marker actually having been emitted since a 1.5B model is not -guaranteed to emit ```` / ````. +guaranteed to emit `` / ``. + +The streaming cases do the same for `stream=true`. Each coupled chunk's +`token_ids` becomes one `/inference/v1/generate` stream chunk, fed +through the chunked derender path with `stream_state` threaded across calls +the way a client would, so both parsers see the same chunk boundaries. """ import json @@ -24,7 +29,12 @@ import httpx import pytest import pytest_asyncio +import regex as re +from tests.entrypoints.scale_out.derender.utils import ( + assemble_stream, + stream_chat_derender, +) from tests.utils import RemoteOpenAIServer MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" @@ -228,8 +238,8 @@ async def test_parity_logprobs(client): """token_id:N resolution parity vs. the coupled server's real strings. A real disaggregated worker only has token IDs so it emits logprobs - with ``token_id:N`` placeholders (``return_tokens_as_token_ids=True`` - reproduces that shape here). ``/derender`` must resolve those + with `token_id:N` placeholders (`return_tokens_as_token_ids=True` + reproduces that shape here). `/derender` must resolve those placeholders to the same token strings/bytes the coupled server resolves them to directly. """ @@ -267,3 +277,166 @@ async def test_parity_logprobs(client): for d_entry, r_entry in zip(d_content, r_content): assert d_entry["token"] == r_entry["token"] assert d_entry["bytes"] == r_entry["bytes"] + + +# --------------------------------------------------------------------------- +# Streaming helpers +# --------------------------------------------------------------------------- + + +async def _coupled_stream( + client: httpx.AsyncClient, messages: list[dict], **extra +) -> tuple[list[dict], list[int]]: + """Stream `/v1/chat/completions` and return the choices that carry + `token_ids` along with the prompt token IDs. + + The leading role only chunk has no `token_ids` and no counterpart on + `/inference/v1/generate`, so it is dropped. + """ + choices: list[dict] = [] + prompt_token_ids: list[int] | None = None + async with client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": MODEL, + "messages": messages, + "temperature": 0, + "max_tokens": 128, + "return_token_ids": True, + "stream": True, + **extra, + }, + ) as resp: + assert resp.status_code == 200, await resp.aread() + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[len("data: ") :] + if data == "[DONE]": + break + chunk = json.loads(data) + assert "error" not in chunk, chunk + if chunk.get("prompt_token_ids") is not None: + prompt_token_ids = chunk["prompt_token_ids"] + choices.extend( + ch for ch in chunk["choices"] if ch.get("token_ids") is not None + ) + assert prompt_token_ids is not None + return choices, prompt_token_ids + + +def _tool_call_id_shape(tool_call_id: str | None) -> str | None: + """Drop the random suffix so IDs compare by format, e.g. `chatcmpl-tool-`. + Deterministic IDs such as kimi_k2's `functions.:` are kept + whole.""" + if tool_call_id is None: + return None + return re.sub(r"[0-9a-f]{16}$", "", tool_call_id) + + +async def _run_stream_parity_case( + client: httpx.AsyncClient, + messages: list[dict], + finish_only_tail: bool = False, + **extra, +) -> tuple[dict, list[dict]]: + """Stream the coupled request, replay its chunks through the streaming + derender endpoint and assert the assembled messages agree. Returns the + assembled coupled message and the coupled choices. + + Derender gets the engine's finish reason, since `"tool_calls"` is a chat + level rewrite of `"stop"` that derender has to reapply itself. With + `finish_only_tail` it goes on a trailing chunk with no tokens. + """ + coupled_choices, prompt_token_ids = await _coupled_stream(client, messages, **extra) + finish_reason = coupled_choices[-1]["finish_reason"] + disagg = await stream_chat_derender( + client, + [tid for ch in coupled_choices for tid in ch["token_ids"]], + [len(ch["token_ids"]) for ch in coupled_choices] + + ([0] if finish_only_tail else []), + {"model": MODEL, "messages": messages, **extra}, + len(prompt_token_ids), + prompt_token_ids, + finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, + ) + coupled = assemble_stream(coupled_choices) + for message in (coupled, disagg): + for tc in message["tool_calls"]: + tc["id"] = _tool_call_id_shape(tc["id"]) + assert disagg == coupled + return coupled, coupled_choices + + +# --------------------------------------------------------------------------- +# Streaming parity cases +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_stream_parity_plain(client): + """Plain streamed detokenization parity.""" + messages = [ + {"role": "user", "content": "What is 2+2? Answer in one short sentence."} + ] + await _run_stream_parity_case(client, messages) + + +@pytest.mark.asyncio +async def test_stream_parity_reasoning(client): + """Streamed reasoning/content split parity.""" + messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] + coupled, _ = await _run_stream_parity_case( + client, messages, include_reasoning=True, max_tokens=256 + ) + + if not coupled["reasoning"]: + pytest.skip("Model did not emit a block") + + +@pytest.mark.asyncio +async def test_stream_parity_tool_call(client): + """Streamed tool call parity under `tool_choice="auto"`, where a tool + call also rewrites `finish_reason` to `"tool_calls"`.""" + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + coupled, _ = await _run_stream_parity_case( + client, messages, tools=TOOLS, tool_choice="auto", max_tokens=512 + ) + + if not coupled["tool_calls"]: + pytest.skip("Model did not emit a tool call") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunking", ["engine_steps", "multi_token", "finish_only_tail"] +) +async def test_stream_parity_reasoning_and_tool_call(client, chunking): + """Streamed reasoning + forced tool call parity across chunk shapes. + + `multi_token` batches tokens per chunk with `stream_interval` the way + speculative decoding does, so replay has to reuse the original chunk + boundaries. `finish_only_tail` sends the finish reason on a chunk with no + tokens, which must still flush buffered tool call arguments. + """ + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + extra: dict = { + "tools": TOOLS, + "tool_choice": FORCE_WEATHER_TOOL, + "include_reasoning": True, + "max_tokens": 512, + } + if chunking == "multi_token": + extra["stream_interval"] = 4 + coupled, coupled_choices = await _run_stream_parity_case( + client, + messages, + finish_only_tail=chunking == "finish_only_tail", + **extra, + ) + + if chunking == "multi_token": + assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices) + if not (coupled["reasoning"] and coupled["tool_calls"]): + pytest.skip("Model did not emit both a block and a tool call") diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index e7a2e3b856cf..4e325db32c16 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -5,24 +5,24 @@ Tests are split into two layers: -1. Unit tests (no server): covers ``_detokenize_delta`` correctness - (chunked == one-shot) and ``derender_completion_stream`` / - ``derender_chat_stream`` logic via a real tokenizer on a tiny model. +1. Unit tests (no server): covers `_detokenize_delta` correctness + (chunked == one-shot) and `derender_completion_stream` / + `derender_chat_stream` logic via a real tokenizer on a tiny model. The parser path is covered both with a deterministic stub parser and - with the real ``HarmonyParser`` (skipped without ``openai_harmony``). + with the real `HarmonyParser` (skipped without `openai_harmony`). 2. Integration tests (require a running render server): covers the full HTTP round-trip through the streaming endpoint. Marked with - ``@pytest.mark.asyncio`` and gated by the ``server`` / ``client`` - fixtures from the sibling ``test_derender.py``. + `@pytest.mark.asyncio` and gated by the `server` / `client` + fixtures from the sibling `test_derender.py`. """ import json -from collections.abc import Callable import pytest import pytest_asyncio +from tests.entrypoints.scale_out.derender.utils import stream_chat_derender from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -356,7 +356,7 @@ def test_state_window_stays_bounded(self, derenderer, tokenizer): """prev_tokens must not grow with the number of chunks (bounded transport). Guards that the carried decode window is a small constant - tail, so cumulative ``stream_state`` transport is O(n) and not O(n^2). + tail, so cumulative `stream_state` transport is O(n) and not O(n^2). """ token_ids = tokenizer.encode( "a reasonably long ascii stream of tokens used to exercise the " @@ -1616,86 +1616,6 @@ async def _render_parser_chat(client, messages: list[dict]) -> dict: return resp.json() -async def _stream_chat_derender( - client, - output_ids: list[int], - chunk_sizes: list[int], - chat_request: dict, - prompt_tokens: int, - prompt_token_ids: list[int], - on_chunk: Callable[[list[dict]], None] | None = None, -) -> dict: - """Feed `output_ids` through the streaming chat derender endpoint in - the given `chunk_sizes`, threading `stream_state` across calls and - return the assembled message. - - If `on_chunk` is given, it is called after every chunk with a - snapshot (deep copy) of the `tool_calls` accumulator so far, letting - callers assert properties of the intermediate deltas (e.g. monotonic - argument growth) rather than only the final assembled result. - """ - state = None - content = "" - reasoning = "" - tool_calls: list[dict] = [] - finish_reason = None - - pos = 0 - for i, size in enumerate(chunk_sizes): - tids = output_ids[pos : pos + size] - pos += size - is_last = i == len(chunk_sizes) - 1 - resp = await client.post( - "/v1/chat/completions/derender", - json={ - "stream": True, - "model": PARSER_MODEL, - "generate_chunk": { - "request_id": "stream-test", - "choices": [ - { - "index": 0, - "token_ids": tids, - "finish_reason": "stop" if is_last else None, - } - ], - }, - "stream_state": state, - "prompt_tokens": prompt_tokens, - "prompt_token_ids": prompt_token_ids, - "chat_request": chat_request, - }, - ) - assert resp.status_code == 200, resp.text - data = resp.json() - state = data["stream_state"] - delta = data["chunk"]["choices"][0]["delta"] - content += delta.get("content") or "" - reasoning += delta.get("reasoning") or "" - for tc in delta.get("tool_calls") or []: - idx = tc["index"] - while len(tool_calls) <= idx: - tool_calls.append({"id": None, "name": None, "arguments": ""}) - if tc.get("id"): - tool_calls[idx]["id"] = tc["id"] - fn = tc.get("function") or {} - if fn.get("name"): - tool_calls[idx]["name"] = fn["name"] - if fn.get("arguments"): - tool_calls[idx]["arguments"] += fn["arguments"] - if is_last: - finish_reason = data["chunk"]["choices"][0]["finish_reason"] - if on_chunk is not None: - on_chunk([dict(tc) for tc in tool_calls]) - - return { - "content": content or None, - "reasoning": reasoning or None, - "tool_calls": tool_calls, - "finish_reason": finish_reason, - } - - @pytest.mark.asyncio @pytest.mark.parametrize( "chunking", @@ -1755,7 +1675,7 @@ async def test_stream_parsed_matches_batch_reasoning( "triples": [3] * (n // 3) + [n % 3], }[chunking] chunk_sizes = [c for c in chunk_sizes if c > 0] - streamed = await _stream_chat_derender( + streamed = await stream_chat_derender( parser_client, output_ids, chunk_sizes, @@ -1823,7 +1743,7 @@ def _record( if tool_calls and tool_calls[0]["arguments"]: _snapshots.append(tool_calls[0]["arguments"]) - streamed = await _stream_chat_derender( + streamed = await stream_chat_derender( parser_client, output_ids, chunk_sizes, @@ -1868,7 +1788,7 @@ async def test_stream_parsed_cjk_across_chunk_boundaries( "messages": messages, "include_reasoning": True, } - streamed = await _stream_chat_derender( + streamed = await stream_chat_derender( parser_client, output_ids, [1] * len(output_ids), @@ -1896,7 +1816,7 @@ async def test_stream_parsed_include_reasoning_false(parser_client, parser_token "messages": messages, "include_reasoning": False, } - streamed = await _stream_chat_derender( + streamed = await stream_chat_derender( parser_client, output_ids, [1] * len(output_ids), diff --git a/tests/entrypoints/scale_out/derender/utils.py b/tests/entrypoints/scale_out/derender/utils.py new file mode 100644 index 000000000000..d8128353bf3c --- /dev/null +++ b/tests/entrypoints/scale_out/derender/utils.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Helpers shared by the streaming chat derender tests.""" + +from collections.abc import Callable + +import httpx + + +def assemble_stream(choices: list[dict]) -> dict: + """Fold streamed chat choices into one message. + + A tool call index must keep the first ID it was given. On the derender + side that is what pinning the IDs in `stream_state` guarantees. + """ + content = "" + reasoning = "" + tool_calls: list[dict] = [] + finish_reason = None + for choice in choices: + delta = choice["delta"] + content += delta.get("content") or "" + reasoning += delta.get("reasoning") or "" + for tc in delta.get("tool_calls") or []: + idx = tc["index"] + while len(tool_calls) <= idx: + tool_calls.append({"id": None, "name": None, "arguments": ""}) + if tc.get("id"): + assert tool_calls[idx]["id"] in (None, tc["id"]), ( + f"tool call {idx} changed ID mid stream" + ) + tool_calls[idx]["id"] = tc["id"] + fn = tc.get("function") or {} + if fn.get("name"): + tool_calls[idx]["name"] = fn["name"] + tool_calls[idx]["arguments"] += fn.get("arguments") or "" + finish_reason = choice.get("finish_reason") or finish_reason + return { + "content": content or None, + "reasoning": reasoning or None, + "tool_calls": tool_calls, + "finish_reason": finish_reason, + } + + +async def stream_chat_derender( + client: httpx.AsyncClient, + output_ids: list[int], + chunk_sizes: list[int], + chat_request: dict, + prompt_tokens: int, + prompt_token_ids: list[int], + finish_reason: str = "stop", + on_chunk: Callable[[list[dict]], None] | None = None, +) -> dict: + """Feed `output_ids` through the streaming chat derender endpoint in + the given `chunk_sizes`, threading `stream_state` across calls and + return the assembled message. + + `finish_reason` goes on the last chunk, so a trailing size of 0 sends it + on a chunk with no tokens. + + If `on_chunk` is given, it is called after every chunk with the + `tool_calls` assembled so far, letting callers assert properties of the + intermediate deltas (e.g. monotonic argument growth) rather than only + the final assembled result. + """ + state = None + choices: list[dict] = [] + pos = 0 + for i, size in enumerate(chunk_sizes): + tids = output_ids[pos : pos + size] + pos += size + is_last = i == len(chunk_sizes) - 1 + resp = await client.post( + "/v1/chat/completions/derender", + json={ + "stream": True, + "model": chat_request["model"], + "generate_chunk": { + "request_id": "stream-test", + "choices": [ + { + "index": 0, + "token_ids": tids, + "finish_reason": finish_reason if is_last else None, + } + ], + }, + "stream_state": state, + "prompt_tokens": prompt_tokens, + "prompt_token_ids": prompt_token_ids, + "chat_request": chat_request, + }, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + state = data["stream_state"] + choices.extend(data["chunk"]["choices"]) + if on_chunk is not None: + on_chunk(assemble_stream(choices)["tool_calls"]) + + return assemble_stream(choices) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 088ec4d2931f..fb2b93872388 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -602,6 +602,7 @@ class DerenderChatStreamRequest(BaseModel): the client carried ``stream_state``. """ + # --8<-- [start:derender-chat-stream-request] stream: Literal[True] model: str | None = None @@ -629,6 +630,7 @@ class DerenderChatStreamRequest(BaseModel): chat_request: ChatCompletionRequest | None = None """The original (post adjust_request) ChatCompletionRequest from /render.""" + # --8<-- [end:derender-chat-stream-request] class DerenderCompletionStreamRequest(BaseModel): @@ -639,6 +641,7 @@ class DerenderCompletionStreamRequest(BaseModel): returns the derendered chunk plus updated state. """ + # --8<-- [start:derender-completion-stream-request] stream: Literal[True] model: str | None = None @@ -653,6 +656,7 @@ class DerenderCompletionStreamRequest(BaseModel): completion_request: CompletionRequest | None = None """The original (post adjust_request) CompletionRequest from /render.""" + # --8<-- [end:derender-completion-stream-request] class DerenderChatStreamResponse(BaseModel): From 7488b5780f585fc0a00a6ab28b6a2477696574d1 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 09:42:47 -0500 Subject: [PATCH 584/774] [ROCm][Tests] Reduce host memory when downcasting FP32 HF references (#57869) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/conftest.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 508e9fa3d50e..cd1d11917b98 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -504,6 +504,36 @@ def _init( if model_cls is not None: _fix_v4_tied_weights_keys(model_cls) + from transformers.integrations import is_deepspeed_zero3_enabled + + # On ROCm, avoid a converted CPU copy when downcasting FP32 references. + # Same-dtype CPU loads can stay lazily mapped and use less host RAM. + # Leave customized, remote, and quantized loading unchanged. + if ( + current_platform.is_rocm() + and self.device == "cuda" + and auto_cls is AutoModelForCausalLM + and revision is None + and model_kwargs.keys() == {"dtype"} + and self.config.dtype == torch.float32 + and model_kwargs["dtype"] in (torch.float16, torch.bfloat16) + and not (trust_remote_code and hasattr(self.config, "auto_map")) + and getattr(self.config, "quantization_config", None) is None + and getattr( + self.config.get_text_config(decoder=True), + "quantization_config", + None, + ) + is None + and not is_deepspeed_zero3_enabled() + ): + model_kwargs = { + **model_kwargs, + "device_map": torch.device( + "cuda", torch.accelerator.current_device_index() + ), + } + model = cast( nn.Module, auto_cls.from_pretrained( From c8824a1e42bbd674751f808d657d58024c68fb0f Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Mon, 21 Sep 2026 22:43:07 +0800 Subject: [PATCH 585/774] [MM] Move `supports_multimodal_inputs` and cache out of registry (#57913) Signed-off-by: DarkLight1337 --- tests/config/test_multimodal_config.py | 100 +++-- .../test_generative_scoring.py | 5 +- .../openai/chat_completion/test_chat_error.py | 3 +- .../chat_completion/test_serving_chat.py | 1 + .../completion/test_completion_error.py | 3 +- .../openai/completion/test_lora_resolvers.py | 1 + .../entrypoints/openai/test_render_parity.py | 3 +- .../test_generate_stream.py | 3 +- .../tokenize/test_serving_tokenization.py | 3 +- .../processing/test_audio_in_video.py | 4 +- tests/multimodal/test_cache.py | 49 +-- tests/multimodal/test_registry.py | 22 -- tests/renderers/test_cohere.py | 1 + tests/renderers/test_completions.py | 1 + tests/renderers/test_kimi_k3.py | 1 + tests/renderers/test_mistral.py | 1 + .../unit/test_handshake_pp_aggregation.py | 11 +- .../test_scheduler_streaming.py | 1 + .../test_gpu_autoregressive_speculator.py | 50 +-- vllm/config/model.py | 88 ++++- vllm/multimodal/cache/__init__.py | 38 ++ vllm/multimodal/{cache.py => cache/base.py} | 367 +----------------- vllm/multimodal/cache/factories.py | 107 +++++ vllm/multimodal/cache/lru.py | 142 +++++++ vllm/multimodal/cache/shm.py | 261 +++++++++++++ vllm/multimodal/encoder_budget.py | 9 +- vllm/multimodal/registry.py | 140 +------ vllm/renderers/base.py | 14 +- vllm/v1/core/sched/scheduler.py | 4 +- vllm/v1/engine/core.py | 11 +- vllm/v1/engine/input_processor.py | 2 +- vllm/v1/spec_decode/llm_base_proposer.py | 6 +- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 5 +- vllm/v1/worker/gpu_model_runner.py | 4 +- vllm/v1/worker/worker_base.py | 26 +- 36 files changed, 784 insertions(+), 707 deletions(-) create mode 100644 vllm/multimodal/cache/__init__.py rename vllm/multimodal/{cache.py => cache/base.py} (54%) create mode 100644 vllm/multimodal/cache/factories.py create mode 100644 vllm/multimodal/cache/lru.py create mode 100644 vllm/multimodal/cache/shm.py diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 84cf092c2159..225f41bd014e 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy from unittest.mock import MagicMock, patch import pytest import torch from transformers import PretrainedConfig +from tests.models.utils import build_model_context from vllm.config.ec_transfer import ECRole, ECTransferConfig from vllm.config.model import ModelConfig from vllm.config.multimodal import MultiModalConfig @@ -17,6 +19,26 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum +@pytest.mark.parametrize( + "model_id,limit_mm_per_prompt,expected", + [ + ("Qwen/Qwen2-0.5B-Instruct", {}, False), + ("Qwen/Qwen2.5-VL-3B-Instruct", {}, True), + ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0, "video": 0}, False), + ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0}, True), + ], +) +@pytest.mark.core_model +def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected): + """Test supports_multimodal_inputs returns correct boolean for various + configs.""" + ctx = build_model_context( + model_id, + limit_mm_per_prompt=limit_mm_per_prompt, + ) + assert ctx.model_config.supports_multimodal_inputs is expected + + def test_mm_encoder_attn_backend_str_conversion(): config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN") # type: ignore[arg-type] assert config.mm_encoder_attn_backend == AttentionBackendEnum.FLASH_ATTN @@ -84,59 +106,33 @@ def test_mm_encoder_attn_dtype_hash_updates(tmp_path): assert fp8_hash != fp8_static_hash +_MULTIMODAL_MODEL = "llava-hf/llava-1.5-7b-hf" +_TEXT_ONLY_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" + + def _make_mm_prefix_model_config( + model: str = _MULTIMODAL_MODEL, *, language_model_only: bool = False, ) -> ModelConfig: - model_config = MagicMock(spec=ModelConfig) - model_config.multimodal_config = MultiModalConfig( - language_model_only=language_model_only - ) - # Bind real helper methods onto the mock. - model_config._supports_multimodal_for_mm_prefix = ( - ModelConfig._supports_multimodal_for_mm_prefix.__get__( - model_config, ModelConfig - ) - ) - return model_config - - -@pytest.mark.parametrize("supports_mm", [True, False]) -def test_supports_multimodal_for_mm_prefix_uses_registry(supports_mm: bool): - model_config = _make_mm_prefix_model_config() - - with patch( - "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", - return_value=supports_mm, - ) as mocked: - assert model_config._supports_multimodal_for_mm_prefix() is supports_mm - mocked.assert_called_once_with(model_config) - - # Sticky cache — registry must not be consulted again. - with patch( - "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", - side_effect=AssertionError("should use cache"), - ): - assert model_config._supports_multimodal_for_mm_prefix() is supports_mm + return ModelConfig(model, language_model_only=language_model_only) def test_supports_multimodal_for_mm_prefix_before_multimodal_config(): - model_config = _make_mm_prefix_model_config() - model_config.multimodal_config = None + """A text-only model never builds a multimodal config, so the early + return must neither clear mm_prefix nor write the sticky cache.""" + model_config = _make_mm_prefix_model_config(_TEXT_ONLY_MODEL) + assert model_config.multimodal_config is None assert model_config._supports_multimodal_for_mm_prefix() is True - assert not hasattr(model_config, "_supports_multimodal_inputs_cached") + assert not getattr(model_config, "_supports_multimodal_inputs_cache", None) def test_language_model_only_disables_via_supports_multimodal_inputs(): """language_model_only zeros all limits, so registry reports text-only.""" model_config = _make_mm_prefix_model_config(language_model_only=True) - with patch( - "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", - return_value=False, - ): - assert model_config._supports_multimodal_for_mm_prefix() is False + assert model_config._supports_multimodal_for_mm_prefix() is False def test_convertor_clears_mm_prefix_when_multimodal_disabled(): @@ -159,23 +155,19 @@ def test_convertor_clears_mm_prefix_when_multimodal_disabled(): def test_sticky_cache_survives_text_subconfig_regeneration(): """with_hf_config deepcopies the cached decision onto text submodules.""" model_config = _make_mm_prefix_model_config() - with patch( - "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", - return_value=False, - ): - assert model_config._supports_multimodal_for_mm_prefix() is False + assert model_config._supports_multimodal_for_mm_prefix() is True - # Simulate deepcopy onto a Gemma4ForCausalLM-like config that would - # otherwise fail registry lookup / return False incorrectly. - text_config = _make_mm_prefix_model_config() - text_config._supports_multimodal_inputs_cached = ( - model_config._supports_multimodal_inputs_cached - ) - with patch( - "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", - side_effect=AssertionError("must not re-query registry"), - ): - assert text_config._supports_multimodal_for_mm_prefix() is False + # `with_hf_config` deep-copies this config and swaps `hf_config` for a + # text-only submodule (e.g. Gemma4ForCausalLM). + text_config = copy.deepcopy(model_config) + text_config.hf_config = model_config.hf_text_config + assert text_config._supports_multimodal_for_mm_prefix() is True + + # Without the copied cache the submodule architecture has no registered + # multimodal processor, so re-querying the registry wrongly settles on + # text-only and would clear mm_prefix. + del text_config._supports_multimodal_inputs_cache + assert text_config._supports_multimodal_for_mm_prefix() is False @pytest.mark.parametrize( diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index b67efcc22931..f413869c9009 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -60,7 +60,10 @@ class MockModelConfig: encoder_config = None generation_config: str = "auto" media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False vocab_size = 151936 def get_diff_sampling_param(self): diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index b83341cc5014..a1665a71de6d 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -56,9 +56,10 @@ class MockModelConfig: encoder_config = None generation_config: str = "auto" media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 def get_diff_sampling_param(self): diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 27573203623e..9f92e1ce3594 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -556,6 +556,7 @@ class MockModelConfig: skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 enable_prompt_embeds: bool = False diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 46f4fe452cb9..28a499e1296b 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -61,9 +61,10 @@ class MockModelConfig: encoder_config = None generation_config: str = "auto" media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 def get_diff_sampling_param(self): diff --git a/tests/entrypoints/openai/completion/test_lora_resolvers.py b/tests/entrypoints/openai/completion/test_lora_resolvers.py index 4d8c3a1fe2a0..8cd874996cd8 100644 --- a/tests/entrypoints/openai/completion/test_lora_resolvers.py +++ b/tests/entrypoints/openai/completion/test_lora_resolvers.py @@ -54,6 +54,7 @@ class MockModelConfig: skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 def get_diff_sampling_param(self): diff --git a/tests/entrypoints/openai/test_render_parity.py b/tests/entrypoints/openai/test_render_parity.py index 92b7fe71bd80..4fc44c5a7172 100644 --- a/tests/entrypoints/openai/test_render_parity.py +++ b/tests/entrypoints/openai/test_render_parity.py @@ -86,9 +86,10 @@ class MockModelConfig: generation_config: str = "auto" override_generation_config: dict[str, Any] = field(default_factory=dict) media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 enable_prompt_embeds: bool = False diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index ba5839990c67..f7c6f05e7fcd 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -61,9 +61,10 @@ class MockModelConfig: encoder_config = None generation_config: str = "auto" media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 def get_diff_sampling_param(self): diff --git a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py index a74962ab0ddd..057d7d1e8a74 100644 --- a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py +++ b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py @@ -56,9 +56,10 @@ class MockModelConfig: encoder_config = None generation_config: str = "auto" media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) - skip_tokenizer_init = False + skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 def get_diff_sampling_param(self): diff --git a/tests/models/multimodal/processing/test_audio_in_video.py b/tests/models/multimodal/processing/test_audio_in_video.py index acb102f60fb0..bde7bd51ab06 100644 --- a/tests/models/multimodal/processing/test_audio_in_video.py +++ b/tests/models/multimodal/processing/test_audio_in_video.py @@ -22,7 +22,7 @@ import pytest from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.cache import MultiModalProcessorSenderCache +from vllm.multimodal.cache import LruKeyReplicatedSenderCache from ....multimodal.utils import random_audio, random_video from ...utils import build_model_context @@ -68,7 +68,7 @@ def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> Non # Sender cache: on a cache hit returns (None, prompt_updates) for each # item, setting mm_kwargs["video"] = [None] – the exact condition that # triggered the original bug. - sender_cache = MultiModalProcessorSenderCache(ctx.model_config) + sender_cache = LruKeyReplicatedSenderCache(ctx.model_config) video_token_id = processor.info.get_hf_config().video_token_id diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index 37ec5b0cb8c1..ffd7d61cf25b 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -10,20 +10,23 @@ from vllm.config import ModelConfig, ParallelConfig, VllmConfig from vllm.config.multimodal import MultiModalConfig -from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import ( BaseMultiModalProcessorCache, BaseMultiModalReceiverCache, - MultiModalCache, + LruKeyReplicatedReceiverCache, + LruKeyReplicatedSenderCache, MultiModalCacheMissError, - MultiModalProcessorCacheInItem, - MultiModalProcessorCacheItem, - MultiModalProcessorCacheItemMetadata, MultiModalProcessorOnlyCache, - MultiModalProcessorSenderCache, - MultiModalReceiverCache, ShmObjectStoreReceiverCache, ShmObjectStoreSenderCache, + engine_receiver_cache_from_config, + processor_cache_from_config, +) +from vllm.multimodal.cache.base import ( + MultiModalCache, + MultiModalProcessorCacheInItem, + MultiModalProcessorCacheItem, + MultiModalProcessorCacheItemMetadata, ) from vllm.multimodal.hasher import MultiModalHasher from vllm.multimodal.inputs import ( @@ -137,10 +140,10 @@ def _compare_caches( n_iter: int = 100, seed: int = 0, ): - cache_0_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_0) - cache_0_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_0) - cache_1_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_1) - cache_1_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_1) + cache_0_p0 = processor_cache_from_config(config_0) + cache_0_p1 = engine_receiver_cache_from_config(config_0) + cache_1_p0 = processor_cache_from_config(config_1) + cache_1_p1 = engine_receiver_cache_from_config(config_1) cache_size_gb = max( config_0.model_config.multimodal_config.mm_processor_cache_gb, @@ -275,8 +278,8 @@ def test_oversized_item_is_served_uncached(): assert p0_only.get_and_update_item((small, []), "small")[0] is small assert p0_only.is_cached_item("small") - p0 = MultiModalProcessorSenderCache(model_config) # type: ignore[arg-type] - p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + p0 = LruKeyReplicatedSenderCache(model_config) # type: ignore[arg-type] + p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type] assert p0.get_and_update_item((item, []), "big")[0] is item assert not p0.is_cached_item("big") assert p1.get_and_update_item(item, "big") is item @@ -296,8 +299,8 @@ def test_mm_cache_miss_raises_and_recovers(): model or network. """ model_config = _StubModelConfig(mm_processor_cache_gb=1) - p0 = MultiModalProcessorSenderCache(model_config) # type: ignore[arg-type] - p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + p0 = LruKeyReplicatedSenderCache(model_config) # type: ignore[arg-type] + p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type] mm_hash = "image_A" item = MultiModalKwargsItem.dummy(nbytes=64) @@ -338,7 +341,7 @@ def test_receiver_cache_replaces_stale_item_when_payload_resent(): and kills EngineCore. Prefer the fresh payload and keep it for later hits. """ model_config = _StubModelConfig(mm_processor_cache_gb=1) - p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type] small = MultiModalKwargsItem.dummy(nbytes=64) large = MultiModalKwargsItem.dummy(nbytes=256) mm_hash = "shared-id" @@ -351,7 +354,7 @@ def test_receiver_cache_replaces_stale_item_when_payload_resent(): def test_receiver_cache_features_keep_resent_payload(): """EngineCore updates features through get_and_update_features.""" model_config = _StubModelConfig(mm_processor_cache_gb=1) - p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type] small = MultiModalKwargsItem.dummy(nbytes=64) large = MultiModalKwargsItem.dummy(nbytes=256) mm_hash = "shared-id" @@ -384,7 +387,7 @@ def test_mm_cache_miss_batches_all_drifted_hashes(): non-drifted items in the same request are still ingested. """ model_config = _StubModelConfig(mm_processor_cache_gb=1) - p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type] item = MultiModalKwargsItem.dummy(nbytes=64) def _feature( @@ -582,8 +585,8 @@ def test_cache_eviction_lru_cache(): model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf", mm_processor_cache_gb=6 / GiB_bytes, ) - sender_cache = MultiModalProcessorSenderCache(model_config) - receiver_cache = MultiModalReceiverCache(model_config) + sender_cache = LruKeyReplicatedSenderCache(model_config) + receiver_cache = LruKeyReplicatedReceiverCache(model_config) _run_test_cache_eviction_lru(sender_cache, receiver_cache, base_item_size=1) @@ -775,7 +778,7 @@ def test_processor_cache_shared_across_loras(): model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf", mm_processor_cache_gb=1, ) - receiver_cache = MultiModalReceiverCache(model_config) + receiver_cache = LruKeyReplicatedReceiverCache(model_config) base_mm_hash = "image_hash_abc123" lora_a_identifier = f"12345:{base_mm_hash}" @@ -827,8 +830,8 @@ async def test_release_kv_cache_resends_mm_payload(use_async, release_error): ) model_config = ctx.model_config - sender = MultiModalProcessorSenderCache(model_config) - receiver = MultiModalReceiverCache(model_config) + sender = LruKeyReplicatedSenderCache(model_config) + receiver = LruKeyReplicatedReceiverCache(model_config) item = _dummy_item({"pixel_values": 16}) mm_hash = "image_A" payload, _ = sender.get_and_update_item((item, []), mm_hash) diff --git a/tests/multimodal/test_registry.py b/tests/multimodal/test_registry.py index d7947845afe3..f395d3693334 100644 --- a/tests/multimodal/test_registry.py +++ b/tests/multimodal/test_registry.py @@ -10,31 +10,9 @@ from vllm.multimodal import MULTIMODAL_REGISTRY -from ..models.utils import build_model_context - pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize( - "model_id,limit_mm_per_prompt,expected", - [ - ("Qwen/Qwen2-0.5B-Instruct", {}, False), - ("Qwen/Qwen2.5-VL-3B-Instruct", {}, True), - ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0, "video": 0}, False), - ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0}, True), - ], -) -@pytest.mark.core_model -def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected): - """Test supports_multimodal_inputs returns correct boolean for various - configs.""" - ctx = build_model_context( - model_id, - limit_mm_per_prompt=limit_mm_per_prompt, - ) - assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected - - def test_create_processor_error_uses_served_model_name(): model_config = SimpleNamespace( is_multimodal_model=False, diff --git a/tests/renderers/test_cohere.py b/tests/renderers/test_cohere.py index 76b38734c212..ae2dfe6c7a88 100644 --- a/tests/renderers/test_cohere.py +++ b/tests/renderers/test_cohere.py @@ -997,6 +997,7 @@ class _MockModelConfig: skip_tokenizer_init: bool = True is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 3094eb6b1178..99862219b97e 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -39,6 +39,7 @@ class MockModelConfig: skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 hidden_size: int = 768 dtype: torch.dtype = torch.float32 diff --git a/tests/renderers/test_kimi_k3.py b/tests/renderers/test_kimi_k3.py index af8723f801a2..de43c0f6bbcc 100644 --- a/tests/renderers/test_kimi_k3.py +++ b/tests/renderers/test_kimi_k3.py @@ -40,6 +40,7 @@ class MockHFConfig: class MockModelConfig: runner_type: str = "generate" is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False multimodal_config: Any = None hf_config: MockHFConfig = field(default_factory=MockHFConfig) allowed_local_media_path: str = "" diff --git a/tests/renderers/test_mistral.py b/tests/renderers/test_mistral.py index 3a8b49a4db29..6ad55f2fb27f 100644 --- a/tests/renderers/test_mistral.py +++ b/tests/renderers/test_mistral.py @@ -37,6 +37,7 @@ class MockModelConfig: skip_tokenizer_init: bool = False is_encoder_decoder: bool = False is_multimodal_model: bool = False + supports_multimodal_inputs: bool = False renderer_num_workers: int = 1 diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index e5cbd977b6d8..509d3d3234b9 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -79,11 +79,6 @@ def get_kv_connector(self) -> KVConnectorBase_V1: "resolve_kv_cache_block_sizes", lambda kv_cache_config, vllm_config: (16, 16), ) - monkeypatch.setattr( - engine_core_module, - "MULTIMODAL_REGISTRY", - SimpleNamespace(engine_receiver_cache_from_config=lambda vllm_config: None), - ) monkeypatch.setattr(engine_core_module, "freeze_gc_heap", lambda: None) monkeypatch.setattr( engine_core_module, "maybe_attach_gc_debug_callback", lambda: None @@ -105,7 +100,11 @@ def get_kv_connector(self) -> KVConnectorBase_V1: speculative_config=None, ec_transfer_config=None, max_concurrent_batches=1, - model_config=SimpleNamespace(runner_type="generate", is_diffusion=False), + model_config=SimpleNamespace( + runner_type="generate", + is_diffusion=False, + supports_multimodal_inputs=False, + ), cache_config=SimpleNamespace( enable_prefix_caching=False, prefix_caching_hash_algo="builtin", diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index af858eaa8c04..0d259c1c6948 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -53,6 +53,7 @@ def create_scheduler() -> Scheduler: vllm_config.model_config = MagicMock() vllm_config.model_config.skip_tokenizer_init = True vllm_config.model_config.is_multimodal_model = False + vllm_config.model_config.supports_multimodal_inputs = False vllm_config.model_config.is_encoder_decoder = False vllm_config.model_config.is_diffusion = False vllm_config.model_config.max_model_len = 1024 diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index b9b3570b8a7b..f73023cee359 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -148,8 +148,16 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected) def test_mm_support_configured_after_model_load(monkeypatch): - target_model_config = object() - draft_model_config = object() + checked_configs = [] + + class _ModelConfig: + @property + def supports_multimodal_inputs(self): + checked_configs.append(self) + return True + + target_model_config = _ModelConfig() + draft_model_config = _ModelConfig() vllm_config = SimpleNamespace(model_config=target_model_config) draft_model = _MultimodalDraftModel() @@ -164,19 +172,8 @@ def init_base(speculator, vllm_config, device): speculator.supports_mm_inputs = False speculator.use_acceptance_estimator = False - checked_configs = [] - - def supports_multimodal_inputs(model_config): - checked_configs.append(model_config) - return True - monkeypatch.setattr(DraftModelSpeculator, "__init__", init_base) _mock_base_model_load(monkeypatch) - monkeypatch.setattr( - base_spec_module.MULTIMODAL_REGISTRY, - "supports_multimodal_inputs", - supports_multimodal_inputs, - ) speculator = _TestSpeculator(vllm_config, torch.device("cpu")) @@ -198,7 +195,9 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch): speculator.supports_mm_inputs = False speculator.inputs_embeds = None speculator.use_acceptance_estimator = False - speculator.vllm_config = SimpleNamespace(model_config=object()) + speculator.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(supports_multimodal_inputs=True) + ) speculator.max_num_tokens = 4 speculator.hidden_size = 3 speculator.dtype = torch.float32 @@ -206,11 +205,6 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch): draft_model = _MultimodalDraftModel() speculator.test_draft_model = draft_model _mock_base_model_load(monkeypatch) - monkeypatch.setattr( - base_spec_module.MULTIMODAL_REGISTRY, - "supports_multimodal_inputs", - lambda model_config: True, - ) speculator.load_model(torch.nn.Module()) @@ -223,16 +217,13 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch): speculator.supports_mm_inputs = False speculator.inputs_embeds = None speculator.use_acceptance_estimator = False - speculator.vllm_config = SimpleNamespace(model_config=object()) + speculator.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(supports_multimodal_inputs=True) + ) draft_model = _TextOnlyDraftModel() speculator.test_draft_model = draft_model warning_messages = [] _mock_base_model_load(monkeypatch) - monkeypatch.setattr( - base_spec_module.MULTIMODAL_REGISTRY, - "supports_multimodal_inputs", - lambda model_config: True, - ) monkeypatch.setattr( base_spec_module.logger, "warning_once", @@ -257,7 +248,9 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch): speculator.inputs_embeds = None speculator.use_acceptance_estimator = False speculator.cached_draft_input_embeds = None - speculator.vllm_config = SimpleNamespace(model_config=object()) + speculator.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(supports_multimodal_inputs=True) + ) speculator.max_num_tokens = 4 speculator.max_num_reqs = 2 speculator.num_speculative_steps = 3 @@ -271,11 +264,6 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch): "load_draft_model", lambda self, target_model, target_attn_layer_names: draft_model, ) - monkeypatch.setattr( - base_spec_module.MULTIMODAL_REGISTRY, - "supports_multimodal_inputs", - lambda model_config: True, - ) speculator.load_model(torch.nn.Module()) diff --git a/vllm/config/model.py b/vllm/config/model.py index 145814a26866..6b36afd91068 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -865,6 +865,80 @@ def __post_init__( self._verify_quantization() self._verify_cuda_graph() + def _supports_multimodal_inputs(self) -> bool: + """Checks if the model supports multimodal inputs. + Returns True if the model is multimodal with any non-zero supported + modalities, otherwise returns False, effectively running in + text-only mode. + """ + if not self.is_multimodal_model: + return False + + from vllm.multimodal import MULTIMODAL_REGISTRY + + mm_config = self.get_multimodal_config() + try: + info = MULTIMODAL_REGISTRY.get_processing_info(self) + except ValueError: + # Speculative drafters for multimodal targets (e.g. Qwen3_5MTP, + # Exaone4_5_MTP, MiMoV2OmniMTP) declare `SupportsMultiModal` so + # that they can consume the embeddings merged by the target model, + # but they never run a multi-modal processor of their own. Running + # in text-only mode is the expected outcome for them, not a + # misconfiguration worth warning about. + if self.runner_type != "draft": + logger.warning_once( + "Model %s is treated as multimodal but has no registered " + "multimodal processor; running in text-only mode.", + self.model, + ) + return False + + # Check if all supported modalities have limit == 0 + if all( + mm_config.get_limit_per_prompt(modality) == 0 + for modality in info.supported_mm_limits + ): + # If enable_mm_embeds is True, we still need MM infrastructure + # to process pre-computed embeddings even though encoder won't run + if mm_config.enable_mm_embeds: + return True + + logger.info_once( + "All limits of multimodal modalities supported by the model " + "are set to 0, running in text-only mode." + ) + return False + + return True + + def _cached_supports_multimodal_inputs(self) -> bool: + cache = getattr(self, "_supports_multimodal_inputs_cache", None) + if cache is None: + self._supports_multimodal_inputs_cache = cache = {} + + mm_config = self.multimodal_config + if mm_config is None: + mm_cache_key = None + else: + limits_per_prompt = { + modality: mm_config.get_limit_per_prompt(modality) + for modality in mm_config.limit_per_prompt + } + mm_cache_key = ( + mm_config.language_model_only, + tuple(limits_per_prompt.items()), + mm_config.enable_mm_embeds, + ) + + cache_key = (self.is_multimodal_model, self.runner_type, mm_cache_key) + if cache_key in cache: + return cache[cache_key] + + supports_mm = self._supports_multimodal_inputs() + cache[cache_key] = supports_mm + return supports_mm + def _supports_multimodal_for_mm_prefix(self) -> bool: """Whether multimodal inputs can still appear for this deployment. @@ -881,24 +955,18 @@ def _supports_multimodal_for_mm_prefix(self) -> bool: vision modality is still enabled (e.g. ``image=0`` but video allowed). The deep-copied cache preserves the top-level decision instead. """ - cached = getattr(self, "_supports_multimodal_inputs_cached", None) - if cached is not None: - return cached - if self.multimodal_config is None: # Early call before multimodal init — do not clear mm_prefix yet. return True - from vllm.multimodal import MULTIMODAL_REGISTRY - - supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(self) - self._supports_multimodal_inputs_cached = supports_mm + supports_mm = self._cached_supports_multimodal_inputs() if not supports_mm: logger.info_once( "Disabled mm_prefix attention mode because multimodal inputs " "are configuration-disabled. Attention backends without " "mm_prefix support may now be selected." ) + return supports_mm def get_model_arch_config(self) -> ModelArchitectureConfig: @@ -1859,6 +1927,10 @@ def is_multimodal_model(self) -> bool: def is_multimodal_raw_input_only_model(self) -> bool: return self._model_info.supports_multimodal_raw_input_only + @property + def supports_multimodal_inputs(self) -> bool: + return self._cached_supports_multimodal_inputs() + @property def requires_raw_input_tokens(self) -> bool: return self._model_info.requires_raw_input_tokens diff --git a/vllm/multimodal/cache/__init__.py b/vllm/multimodal/cache/__init__.py new file mode 100644 index 000000000000..29411c185922 --- /dev/null +++ b/vllm/multimodal/cache/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .base import ( + BaseMultiModalProcessorCache, + BaseMultiModalReceiverCache, + MultiModalCacheMissError, + MultiModalProcessorOnlyCache, +) +from .factories import ( + engine_receiver_cache_from_config, + processor_cache_from_config, + processor_only_cache_from_config, + worker_receiver_cache_from_config, +) +from .lru import ( + LruKeyReplicatedReceiverCache, + LruKeyReplicatedSenderCache, +) +from .shm import ( + ShmObjectStoreReceiverCache, + ShmObjectStoreSenderCache, +) + +__all__ = [ + "MultiModalCacheMissError", + "BaseMultiModalProcessorCache", + "BaseMultiModalReceiverCache", + "MultiModalProcessorOnlyCache", + "LruKeyReplicatedSenderCache", + "LruKeyReplicatedReceiverCache", + "ShmObjectStoreReceiverCache", + "ShmObjectStoreSenderCache", + "processor_cache_from_config", + "processor_only_cache_from_config", + "engine_receiver_cache_from_config", + "worker_receiver_cache_from_config", +] diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache/base.py similarity index 54% rename from vllm/multimodal/cache.py rename to vllm/multimodal/cache/base.py index 15d433112cf2..1b9b757917a3 100644 --- a/vllm/multimodal/cache.py +++ b/vllm/multimodal/cache/base.py @@ -4,26 +4,19 @@ import sys from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence -from multiprocessing.synchronize import Lock as LockType -from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar, cast +from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar import torch from typing_extensions import override -import vllm.envs as envs -from vllm.distributed.device_communicators.shm_object_storage import ( - MsgpackSerde, - SingleWriterShmObjectStorage, - SingleWriterShmRingBuffer, -) +from vllm.config import ModelConfig from vllm.logger import init_logger from vllm.utils.cache import CacheInfo, LRUCache from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves -from vllm.utils.mem_constants import GiB_bytes, MiB_bytes +from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import format_gib -from .inputs import ( - MultiModalBatchedField, +from ..inputs import ( MultiModalFeatureSpec, MultiModalFieldElem, MultiModalKwargsItem, @@ -32,9 +25,7 @@ ) if TYPE_CHECKING: - from vllm.config import ModelConfig, VllmConfig - - from .processing.processor import ResolvedPromptUpdate + from ..processing.processor import ResolvedPromptUpdate logger = init_logger(__name__) @@ -399,7 +390,7 @@ class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache): tensor data and metadata) into the cache, and return the input. """ - def __init__(self, model_config: "ModelConfig") -> None: + def __init__(self, model_config: ModelConfig) -> None: super().__init__() mm_config = model_config.get_multimodal_config() @@ -440,216 +431,6 @@ def make_stats(self, *, delta: bool = False) -> CacheInfo: return self._cache.stat(delta=delta) -class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache): - """The cache which is used on P0 when IPC caching is enabled. - - How to update each item: - - - If the item is already in the cache, clear the input to avoid - unnecessary IPC. - - - If the item is not in the cache, store the metadata of that item so - that the eviction policy remains the same as the cache on P1, - and return the input. - By only storing the metadata, we avoid keeping the data itself in - memory inside P0. - """ - - def __init__(self, model_config: "ModelConfig") -> None: - super().__init__() - - mm_config = model_config.get_multimodal_config() - - self._cache = MultiModalCache.get_lru_cache( - mm_config.mm_processor_cache_gb, - MultiModalProcessorCacheItemMetadata, - ) - - @override - def is_cached_item(self, mm_hash: str) -> bool: - return mm_hash in self._cache - - @override - def get_and_update_item( - self, - mm_item: MultiModalProcessorCacheInItem, - mm_hash: str, - ) -> MultiModalProcessorCacheOutItem: - if (cached_item := self._cache.get(mm_hash)) is not None: - return None, cached_item.prompt_updates - - assert mm_item is not None, f"Expected a cached item for {mm_hash=}" - - self.cache_if_fits( - self._cache, mm_hash, MultiModalProcessorCacheItemMetadata(*mm_item) - ) - return mm_item - - @override - def touch_sender_cache_item(self, mm_hash: str) -> None: - self._cache.touch(mm_hash) - - @override - def clear_cache(self) -> None: - self._cache.clear() - - @override - def make_stats(self, *, delta: bool = False) -> CacheInfo: - return self._cache.stat(delta=delta) - - @override - def invalidate(self, mm_hash: str) -> None: - # Drop our stale shadow entry so the next request for this hash re-sends the - # data and repopulates P1 (see MultiModalCacheMissError). - self._cache.pop(mm_hash, None) - - -class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): - """The cache which is used on P0 when IPC caching is enabled. - - How to update each item: - - - If the item is already in the cache, clear the input to avoid - unnecessary IPC. - - - If the item is not in the cache, store the data in shared memory. - """ - - def __init__(self, vllm_config: "VllmConfig") -> None: - super().__init__() - - self.world_size = vllm_config.parallel_config.world_size - mm_config = vllm_config.model_config.get_multimodal_config() - - ring_buffer = SingleWriterShmRingBuffer( - data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), - name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, - create=True, # sender is the writer - ) - self._shm_cache = SingleWriterShmObjectStorage( - max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, - n_readers=self.world_size, - ring_buffer=ring_buffer, - serde_class=MsgpackSerde, - ) - # cache prompt_updates for P0 only - self._p0_cache: dict[str, Sequence[ResolvedPromptUpdate]] = {} - - self._hits = 0 - self._total = 0 - self._last_info = CacheInfo(hits=0, total=0) - - def _stat(self, *, delta: bool = False) -> CacheInfo: - info = CacheInfo(hits=self._hits, total=self._total) - - if delta: - info_delta = info - self._last_info - self._last_info = info - info = info_delta - - return info - - @override - def is_cached_item(self, mm_hash: str) -> bool: - return self._shm_cache.is_cached(mm_hash) - - @override - def get_and_update_item( - self, - mm_item: MultiModalProcessorCacheInItem, - mm_hash: str, - ) -> MultiModalProcessorCacheOutItem: - if self._shm_cache.is_cached(mm_hash): - self._hits += 1 - self._total += 1 - - address, monotonic_id = self._shm_cache.get_cached(mm_hash) - prompt_updates = self._p0_cache[mm_hash] - return self.address_as_item(address, monotonic_id), prompt_updates - - assert mm_item is not None, f"Expected a cached item for {mm_hash=}" - item, prompt_updates = mm_item - - self._total += 1 - - try: - address, monotonic_id = self._shm_cache.put(mm_hash, item) - # Try to remove dangling items if p0 cache is too large. - if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index): - self.remove_dangling_items() - - self._p0_cache[mm_hash] = prompt_updates - return self.address_as_item(address, monotonic_id), prompt_updates - except ValueError as e: - # `put` raises ValueError either for an oversize item or for a - # duplicate key (concurrent insert); the latter is benign so we - # only warn on the oversize case. Subsequent UUID-only requests - # for an oversize item will fail with a cache miss. - if "already exists" not in str(e): - logger.warning_once( - "mm_input %s too large to cache; " - "raise --mm-shm-cache-max-object-size-mb. (%s)", - mm_hash, - str(e), - ) - return mm_item - except MemoryError as e: - # Cache full and protected items prevent eviction. - logger.debug( - "mm_input %s not cached; shm cache full, " - "consider raising --mm-processor-cache-gb. (%s)", - mm_hash, - str(e), - ) - return mm_item - - @override - def touch_sender_cache_item(self, mm_hash: str) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments writer_flag on sender side.""" - self._shm_cache.touch(mm_hash) - - @override - def clear_cache(self) -> None: - self._shm_cache.clear() - self._p0_cache.clear() - - self._hits = 0 - self._total = 0 - self._last_info = CacheInfo(hits=0, total=0) - - @override - def make_stats(self, *, delta: bool = False) -> CacheInfo: - return self._stat(delta=delta) - - @override - def close(self) -> None: - self._shm_cache.close() - - def remove_dangling_items(self) -> None: - """Remove items that are no longer in the shared memory cache.""" - cached_hashes = self._shm_cache.key_index.keys() - dangling_hashes = set(self._p0_cache.keys()) - cached_hashes - for mm_hash in dangling_hashes: - del self._p0_cache[mm_hash] - - def address_as_item( - self, - address: int, - monotonic_id: int, - ) -> MultiModalKwargsItem: - addr_elem = MultiModalFieldElem( - data=address, - field=MultiModalBatchedField(), - ) - id_elem = MultiModalFieldElem( - data=monotonic_id, - field=MultiModalBatchedField(), - ) - - return MultiModalKwargsItem({"address": addr_elem, "monotonic_id": id_elem}) - - class BaseMultiModalReceiverCache( BaseMultiModalCache[MultiModalKwargsItem | None, MultiModalKwargsItem] ): @@ -703,139 +484,3 @@ def touch_receiver_cache_item( """ raise NotImplementedError - - -class MultiModalReceiverCache(BaseMultiModalReceiverCache): - """The cache which is used on P1 when IPC caching is enabled. - - How to update each item: - - - If the caller sent tensor data, store it (replacing any cached item - under the same key) and return that data. P0 can miss after independent - LRU eviction and resend a different item for the same identity. - - If the caller sent no data and the item is cached, return the cached item. - - If the caller sent no data and the item is not cached, raise - ``MultiModalCacheMissError``. - """ - - def __init__(self, model_config: "ModelConfig") -> None: - super().__init__() - - mm_config = model_config.get_multimodal_config() - - self._cache = MultiModalCache.get_lru_cache( - mm_config.mm_processor_cache_gb, - MultiModalKwargsItem, - ) - - @override - def get_and_update_item( - self, - mm_item: MultiModalKwargsItem | None, - mm_hash: str, - ) -> MultiModalKwargsItem: - if mm_item is not None: - # P0 sent a payload. Never keep a stale cached tensor under this - # identity: a later P0 hit would pair new placeholders with the old - # item and crash the engine. Drop first so a too-large replacement - # cannot leave the previous entry in place. - self._cache.pop(mm_hash, None) - self.cache_if_fits(self._cache, mm_hash, mm_item) - return mm_item - - if (cached_item := self._cache.get(mm_hash)) is not None: - return cached_item - - # No data and not cached here: P0 sent data=None trusting its shadow, but - # the P0/P1 caches have drifted. Raise a retryable error (not assert) so the - # engine can have P0 drop the stale entry and the client resend the data. - raise MultiModalCacheMissError([mm_hash]) - - @override - def touch_receiver_cache_item( - self, - mm_hash: str, - mm_item: MultiModalKwargsItem | None = None, - ) -> None: - self._cache.touch(mm_hash) - - @override - def clear_cache(self) -> None: - self._cache.clear() - - -class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): - """The cache which is used on P1 Worker Process when IPC caching is enabled. - - How to update each item: - - - If the item has an address, replace the input with the cached item. - - If not, return the input. - """ - - def __init__( - self, - vllm_config: "VllmConfig", - shared_worker_lock: LockType, - ) -> None: - super().__init__() - - self.world_size = vllm_config.parallel_config.world_size - mm_config = vllm_config.model_config.get_multimodal_config() - - ring_buffer = SingleWriterShmRingBuffer( - data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), - name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, - create=False, # Server is a reader - ) - self._shm_cache = SingleWriterShmObjectStorage( - max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, - n_readers=self.world_size, - ring_buffer=ring_buffer, - serde_class=MsgpackSerde, - reader_lock=shared_worker_lock, - ) - - @override - def get_and_update_features( - self, mm_features: list["MultiModalFeatureSpec"] - ) -> list["MultiModalFeatureSpec"]: - # strip_covered_mm_data preserves address items, so None here represents - # a stripped uncached payload and has no SHM reference to acknowledge. - features_with_data = [ - feature for feature in mm_features if feature.data is not None - ] - super().get_and_update_features(features_with_data) - return mm_features - - @override - def get_and_update_item( - self, - mm_item: MultiModalKwargsItem | None, - mm_hash: str, - ) -> MultiModalKwargsItem: - assert mm_item is not None, f"Expected an address item for {mm_hash=}" - if "address" in mm_item: - address = cast(int, mm_item["address"].data) - monotonic_id = cast(int, mm_item["monotonic_id"].data) - return self._shm_cache.get(address, monotonic_id) - - return mm_item - - @override - def touch_receiver_cache_item( - self, - mm_hash: str, - mm_item: MultiModalKwargsItem | None = None, - ) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments reader_count on receiver side.""" - assert mm_item is not None - if "address" in mm_item: - address = cast(int, mm_item["address"].data) - monotonic_id = cast(int, mm_item["monotonic_id"].data) - self._shm_cache.touch(mm_hash, address=address, monotonic_id=monotonic_id) - - @override - def clear_cache(self) -> None: - self._shm_cache.clear() diff --git a/vllm/multimodal/cache/factories.py b/vllm/multimodal/cache/factories.py new file mode 100644 index 000000000000..aa4f5a4d2cb5 --- /dev/null +++ b/vllm/multimodal/cache/factories.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from multiprocessing.synchronize import Lock as LockType +from typing import Literal + +from vllm.config import VllmConfig + +from .base import ( + BaseMultiModalProcessorCache, + BaseMultiModalReceiverCache, + MultiModalProcessorOnlyCache, +) +from .lru import ( + LruKeyReplicatedReceiverCache, + LruKeyReplicatedSenderCache, +) +from .shm import ( + ShmObjectStoreReceiverCache, + ShmObjectStoreSenderCache, +) + + +def _get_cache_type( + vllm_config: VllmConfig, +) -> Literal[None, "processor_only", "lru", "shm"]: + model_config = vllm_config.model_config + if not model_config.supports_multimodal_inputs: + return None + + # Check if the cache is disabled. + mm_config = model_config.get_multimodal_config() + if mm_config.mm_processor_cache_gb <= 0: + return None + + # Check if IPC caching is supported. + parallel_config = vllm_config.parallel_config + is_ipc_supported = parallel_config._api_process_count == 1 and ( + parallel_config.data_parallel_size == 1 + or parallel_config.data_parallel_external_lb + ) + + if not is_ipc_supported: + return "processor_only" + + mm_config = model_config.get_multimodal_config() + return mm_config.mm_processor_cache_type + + +def processor_cache_from_config( + vllm_config: VllmConfig, +) -> BaseMultiModalProcessorCache | None: + """Return a `BaseMultiModalProcessorCache`, if enabled.""" + cache_type = _get_cache_type(vllm_config) + if cache_type is None: + return None + elif cache_type == "processor_only": + return MultiModalProcessorOnlyCache(vllm_config.model_config) + elif cache_type == "lru": + return LruKeyReplicatedSenderCache(vllm_config.model_config) + elif cache_type == "shm": + return ShmObjectStoreSenderCache(vllm_config) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") + + +def processor_only_cache_from_config( + vllm_config: VllmConfig, +) -> MultiModalProcessorOnlyCache | None: + """Return a `MultiModalProcessorOnlyCache`, if enabled.""" + cache_type = _get_cache_type(vllm_config) + if cache_type is None: + return None + + return MultiModalProcessorOnlyCache(vllm_config.model_config) + + +def engine_receiver_cache_from_config( + vllm_config: VllmConfig, +) -> BaseMultiModalReceiverCache | None: + """Return a `BaseMultiModalReceiverCache` for the engine process.""" + cache_type = _get_cache_type(vllm_config) + if cache_type in (None, "processor_only", "shm"): + return None + elif cache_type == "lru": + return LruKeyReplicatedReceiverCache(vllm_config.model_config) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") + + +def worker_receiver_cache_from_config( + vllm_config: VllmConfig, + shared_worker_lock: LockType | None, +) -> BaseMultiModalReceiverCache | None: + """Return a `BaseMultiModalReceiverCache` for the worker process.""" + cache_type = _get_cache_type(vllm_config) + if cache_type in (None, "processor_only", "lru"): + return None + elif cache_type == "shm": + if shared_worker_lock is None: + raise ValueError( + "Missing `shared_worker_lock` argument from executor. " + "This argument is needed for mm_processor_cache_type='shm'." + ) + + return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") diff --git a/vllm/multimodal/cache/lru.py b/vllm/multimodal/cache/lru.py new file mode 100644 index 000000000000..6b5ec3516a57 --- /dev/null +++ b/vllm/multimodal/cache/lru.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Implementation of Key-Replicated Cache (see docs/configuration/optimization.md).""" + +from typing_extensions import override + +from vllm.config import ModelConfig +from vllm.utils.cache import CacheInfo + +from ..inputs import MultiModalKwargsItem +from .base import ( + BaseMultiModalProcessorCache, + BaseMultiModalReceiverCache, + MultiModalCache, + MultiModalCacheMissError, + MultiModalProcessorCacheInItem, + MultiModalProcessorCacheItemMetadata, + MultiModalProcessorCacheOutItem, +) + + +class LruKeyReplicatedSenderCache(BaseMultiModalProcessorCache): + """The cache which is used on P0 when LRU caching is enabled. + + How to update each item: + + - If the item is already in the cache, clear the input to avoid + unnecessary IPC. + + - If the item is not in the cache, store the metadata of that item so + that the eviction policy remains the same as the cache on P1, + and return the input. + By only storing the metadata, we avoid keeping the data itself in + memory inside P0. + """ + + def __init__(self, model_config: ModelConfig) -> None: + super().__init__() + + mm_config = model_config.get_multimodal_config() + + self._cache = MultiModalCache.get_lru_cache( + mm_config.mm_processor_cache_gb, + MultiModalProcessorCacheItemMetadata, + ) + + @override + def is_cached_item(self, mm_hash: str) -> bool: + return mm_hash in self._cache + + @override + def get_and_update_item( + self, + mm_item: MultiModalProcessorCacheInItem, + mm_hash: str, + ) -> MultiModalProcessorCacheOutItem: + if (cached_item := self._cache.get(mm_hash)) is not None: + return None, cached_item.prompt_updates + + assert mm_item is not None, f"Expected a cached item for {mm_hash=}" + + self.cache_if_fits( + self._cache, mm_hash, MultiModalProcessorCacheItemMetadata(*mm_item) + ) + return mm_item + + @override + def touch_sender_cache_item(self, mm_hash: str) -> None: + self._cache.touch(mm_hash) + + @override + def clear_cache(self) -> None: + self._cache.clear() + + @override + def make_stats(self, *, delta: bool = False) -> CacheInfo: + return self._cache.stat(delta=delta) + + @override + def invalidate(self, mm_hash: str) -> None: + # Drop our stale shadow entry so the next request for this hash re-sends the + # data and repopulates P1 (see MultiModalCacheMissError). + self._cache.pop(mm_hash, None) + + +class LruKeyReplicatedReceiverCache(BaseMultiModalReceiverCache): + """The cache which is used on P1 when LRU caching is enabled. + + How to update each item: + + - If the caller sent tensor data, store it (replacing any cached item + under the same key) and return that data. P0 can miss after independent + LRU eviction and resend a different item for the same identity. + - If the caller sent no data and the item is cached, return the cached item. + - If the caller sent no data and the item is not cached, raise + `MultiModalCacheMissError`. + """ + + def __init__(self, model_config: ModelConfig) -> None: + super().__init__() + + mm_config = model_config.get_multimodal_config() + + self._cache = MultiModalCache.get_lru_cache( + mm_config.mm_processor_cache_gb, + MultiModalKwargsItem, + ) + + @override + def get_and_update_item( + self, + mm_item: MultiModalKwargsItem | None, + mm_hash: str, + ) -> MultiModalKwargsItem: + if mm_item is not None: + # P0 sent a payload. Never keep a stale cached tensor under this + # identity: a later P0 hit would pair new placeholders with the old + # item and crash the engine. Drop first so a too-large replacement + # cannot leave the previous entry in place. + self._cache.pop(mm_hash, None) + self.cache_if_fits(self._cache, mm_hash, mm_item) + return mm_item + + if (cached_item := self._cache.get(mm_hash)) is not None: + return cached_item + + # No data and not cached here: P0 sent data=None trusting its shadow, but + # the P0/P1 caches have drifted. Raise a retryable error (not assert) so the + # engine can have P0 drop the stale entry and the client resend the data. + raise MultiModalCacheMissError([mm_hash]) + + @override + def touch_receiver_cache_item( + self, + mm_hash: str, + mm_item: MultiModalKwargsItem | None = None, + ) -> None: + self._cache.touch(mm_hash) + + @override + def clear_cache(self) -> None: + self._cache.clear() diff --git a/vllm/multimodal/cache/shm.py b/vllm/multimodal/cache/shm.py new file mode 100644 index 000000000000..dddf3494cf5d --- /dev/null +++ b/vllm/multimodal/cache/shm.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Implementation of Shared Memory Cache (see docs/configuration/optimization.md).""" + +from collections.abc import Sequence +from multiprocessing.synchronize import Lock as LockType +from typing import TYPE_CHECKING, cast + +from typing_extensions import override + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed.device_communicators.shm_object_storage import ( + MsgpackSerde, + SingleWriterShmObjectStorage, + SingleWriterShmRingBuffer, +) +from vllm.logger import init_logger +from vllm.utils.cache import CacheInfo +from vllm.utils.mem_constants import GiB_bytes, MiB_bytes + +from ..inputs import ( + MultiModalBatchedField, + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, +) +from .base import ( + BaseMultiModalProcessorCache, + BaseMultiModalReceiverCache, + MultiModalProcessorCacheInItem, + MultiModalProcessorCacheOutItem, +) + +if TYPE_CHECKING: + from ..processing.processor import ResolvedPromptUpdate + +logger = init_logger(__name__) + + +class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): + """The cache which is used on P0 when SHM caching is enabled. + + How to update each item: + + - If the item is already in the cache, clear the input to avoid + unnecessary IPC. + + - If the item is not in the cache, store the data in shared memory. + """ + + def __init__(self, vllm_config: VllmConfig) -> None: + super().__init__() + + self.world_size = vllm_config.parallel_config.world_size + mm_config = vllm_config.model_config.get_multimodal_config() + + ring_buffer = SingleWriterShmRingBuffer( + data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), + name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, + create=True, # sender is the writer + ) + self._shm_cache = SingleWriterShmObjectStorage( + max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, + n_readers=self.world_size, + ring_buffer=ring_buffer, + serde_class=MsgpackSerde, + ) + # cache prompt_updates for P0 only + self._p0_cache: dict[str, Sequence[ResolvedPromptUpdate]] = {} + + self._hits = 0 + self._total = 0 + self._last_info = CacheInfo(hits=0, total=0) + + def _stat(self, *, delta: bool = False) -> CacheInfo: + info = CacheInfo(hits=self._hits, total=self._total) + + if delta: + info_delta = info - self._last_info + self._last_info = info + info = info_delta + + return info + + @override + def is_cached_item(self, mm_hash: str) -> bool: + return self._shm_cache.is_cached(mm_hash) + + @override + def get_and_update_item( + self, + mm_item: MultiModalProcessorCacheInItem, + mm_hash: str, + ) -> MultiModalProcessorCacheOutItem: + if self._shm_cache.is_cached(mm_hash): + self._hits += 1 + self._total += 1 + + address, monotonic_id = self._shm_cache.get_cached(mm_hash) + prompt_updates = self._p0_cache[mm_hash] + return self.address_as_item(address, monotonic_id), prompt_updates + + assert mm_item is not None, f"Expected a cached item for {mm_hash=}" + item, prompt_updates = mm_item + + self._total += 1 + + try: + address, monotonic_id = self._shm_cache.put(mm_hash, item) + # Try to remove dangling items if p0 cache is too large. + if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index): + self.remove_dangling_items() + + self._p0_cache[mm_hash] = prompt_updates + return self.address_as_item(address, monotonic_id), prompt_updates + except ValueError as e: + # `put` raises ValueError either for an oversize item or for a + # duplicate key (concurrent insert); the latter is benign so we + # only warn on the oversize case. Subsequent UUID-only requests + # for an oversize item will fail with a cache miss. + if "already exists" not in str(e): + logger.warning_once( + "mm_input %s too large to cache; " + "raise --mm-shm-cache-max-object-size-mb. (%s)", + mm_hash, + str(e), + ) + return mm_item + except MemoryError as e: + # Cache full and protected items prevent eviction. + logger.debug( + "mm_input %s not cached; shm cache full, " + "consider raising --mm-processor-cache-gb. (%s)", + mm_hash, + str(e), + ) + return mm_item + + @override + def touch_sender_cache_item(self, mm_hash: str) -> None: + """Touch the item in shared memory cache to prevent eviction. + Increments writer_flag on sender side.""" + self._shm_cache.touch(mm_hash) + + @override + def clear_cache(self) -> None: + self._shm_cache.clear() + self._p0_cache.clear() + + self._hits = 0 + self._total = 0 + self._last_info = CacheInfo(hits=0, total=0) + + @override + def make_stats(self, *, delta: bool = False) -> CacheInfo: + return self._stat(delta=delta) + + @override + def close(self) -> None: + self._shm_cache.close() + + def remove_dangling_items(self) -> None: + """Remove items that are no longer in the shared memory cache.""" + cached_hashes = self._shm_cache.key_index.keys() + dangling_hashes = set(self._p0_cache.keys()) - cached_hashes + for mm_hash in dangling_hashes: + del self._p0_cache[mm_hash] + + def address_as_item( + self, + address: int, + monotonic_id: int, + ) -> MultiModalKwargsItem: + addr_elem = MultiModalFieldElem( + data=address, + field=MultiModalBatchedField(), + ) + id_elem = MultiModalFieldElem( + data=monotonic_id, + field=MultiModalBatchedField(), + ) + + return MultiModalKwargsItem({"address": addr_elem, "monotonic_id": id_elem}) + + +class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): + """The cache which is used on P1 Worker Process when SHM caching is enabled. + + How to update each item: + + - If the item has an address, replace the input with the cached item. + - If not, return the input. + """ + + def __init__( + self, + vllm_config: VllmConfig, + shared_worker_lock: LockType, + ) -> None: + super().__init__() + + self.world_size = vllm_config.parallel_config.world_size + mm_config = vllm_config.model_config.get_multimodal_config() + + ring_buffer = SingleWriterShmRingBuffer( + data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), + name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, + create=False, # Server is a reader + ) + self._shm_cache = SingleWriterShmObjectStorage( + max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, + n_readers=self.world_size, + ring_buffer=ring_buffer, + serde_class=MsgpackSerde, + reader_lock=shared_worker_lock, + ) + + @override + def get_and_update_features( + self, mm_features: list["MultiModalFeatureSpec"] + ) -> list["MultiModalFeatureSpec"]: + # strip_covered_mm_data preserves address items, so None here represents + # a stripped uncached payload and has no SHM reference to acknowledge. + features_with_data = [ + feature for feature in mm_features if feature.data is not None + ] + super().get_and_update_features(features_with_data) + return mm_features + + @override + def get_and_update_item( + self, + mm_item: MultiModalKwargsItem | None, + mm_hash: str, + ) -> MultiModalKwargsItem: + assert mm_item is not None, f"Expected an address item for {mm_hash=}" + if "address" in mm_item: + address = cast(int, mm_item["address"].data) + monotonic_id = cast(int, mm_item["monotonic_id"].data) + return self._shm_cache.get(address, monotonic_id) + + return mm_item + + @override + def touch_receiver_cache_item( + self, + mm_hash: str, + mm_item: MultiModalKwargsItem | None = None, + ) -> None: + """Touch the item in shared memory cache to prevent eviction. + Increments reader_count on receiver side.""" + assert mm_item is not None + if "address" in mm_item: + address = cast(int, mm_item["address"].data) + monotonic_id = cast(int, mm_item["monotonic_id"].data) + self._shm_cache.touch(mm_hash, address=address, monotonic_id=monotonic_id) + + @override + def clear_cache(self) -> None: + self._shm_cache.clear() diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index 13448b9096dc..32db904bfa14 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/vllm/multimodal/encoder_budget.py @@ -4,7 +4,10 @@ from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.multimodal.cache import BaseMultiModalProcessorCache +from vllm.multimodal.cache import ( + BaseMultiModalProcessorCache, + processor_only_cache_from_config, +) from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.multimodal.registry import MultiModalRegistry @@ -57,9 +60,7 @@ def __init__( with set_default_torch_num_threads(): # Avoid hang during startup cache = ( - mm_registry.processor_only_cache_from_config(vllm_config) - if enable_cache - else None + processor_only_cache_from_config(vllm_config) if enable_cache else None ) processor = mm_registry.create_processor(model_config) diff --git a/vllm/multimodal/registry.py b/vllm/multimodal/registry.py index 24a7a696b4dd..28383979f917 100644 --- a/vllm/multimodal/registry.py +++ b/vllm/multimodal/registry.py @@ -3,21 +3,11 @@ import threading from collections import defaultdict from dataclasses import dataclass -from multiprocessing.synchronize import Lock as LockType -from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config -from .cache import ( - BaseMultiModalProcessorCache, - BaseMultiModalReceiverCache, - MultiModalProcessorOnlyCache, - MultiModalProcessorSenderCache, - MultiModalReceiverCache, - ShmObjectStoreReceiverCache, - ShmObjectStoreSenderCache, -) from .processing import ( BaseDummyInputsBuilder, BaseMultiModalProcessor, @@ -27,7 +17,7 @@ ) if TYPE_CHECKING: - from vllm.config import ModelConfig, ObservabilityConfig, VllmConfig + from vllm.config import ModelConfig, ObservabilityConfig from vllm.model_executor.models.interfaces import SupportsMultiModal logger = init_logger(__name__) @@ -89,51 +79,6 @@ def build_processor( class MultiModalRegistry: """A registry that dispatches data processing according to the model.""" - def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool: - """Checks if the model supports multimodal inputs. - Returns True if the model is multimodal with any non-zero supported - modalities, otherwise returns False, effectively running in - text-only mode. - """ - if not model_config.is_multimodal_model: - return False - - mm_config = model_config.get_multimodal_config() - try: - info = self._create_processing_info(model_config, tokenizer=None) - except ValueError: - # Speculative drafters for multimodal targets (e.g. Qwen3_5MTP, - # Exaone4_5_MTP, MiMoV2OmniMTP) declare `SupportsMultiModal` so - # that they can consume the embeddings merged by the target model, - # but they never run a multi-modal processor of their own. Running - # in text-only mode is the expected outcome for them, not a - # misconfiguration worth warning about. - if model_config.runner_type != "draft": - logger.warning_once( - "Model %s is treated as multimodal but has no registered " - "multimodal processor; running in text-only mode.", - model_config.model, - ) - return False - - # Check if all supported modalities have limit == 0 - if all( - mm_config.get_limit_per_prompt(modality) == 0 - for modality in info.supported_mm_limits - ): - # If enable_mm_embeds is True, we still need MM infrastructure - # to process pre-computed embeddings even though encoder won't run - if mm_config.enable_mm_embeds: - return True - - logger.info_once( - "All limits of multimodal modalities supported by the model " - "are set to 0, running in text-only mode." - ) - return False - - return True - def register_processor( self, processor: MultiModalProcessorFactory[_I], @@ -220,87 +165,6 @@ def create_processor( return factories.build_processor(ctx) - def _get_cache_type( - self, - vllm_config: "VllmConfig", - ) -> Literal[None, "processor_only", "lru", "shm"]: - model_config = vllm_config.model_config - if not self.supports_multimodal_inputs(model_config): - return None - - # Check if the cache is disabled. - mm_config = model_config.get_multimodal_config() - if mm_config.mm_processor_cache_gb <= 0: - return None - - # Check if IPC caching is supported. - parallel_config = vllm_config.parallel_config - is_ipc_supported = parallel_config._api_process_count == 1 and ( - parallel_config.data_parallel_size == 1 - or parallel_config.data_parallel_external_lb - ) - - if not is_ipc_supported: - return "processor_only" - - mm_config = model_config.get_multimodal_config() - return mm_config.mm_processor_cache_type - - def processor_cache_from_config( - self, - vllm_config: "VllmConfig", - ) -> BaseMultiModalProcessorCache | None: - """Return a `BaseMultiModalProcessorCache`, if enabled.""" - cache_type = self._get_cache_type(vllm_config) - if cache_type is None: - return None - elif cache_type == "processor_only": - return MultiModalProcessorOnlyCache(vllm_config.model_config) - elif cache_type == "lru": - return MultiModalProcessorSenderCache(vllm_config.model_config) - elif cache_type == "shm": - return ShmObjectStoreSenderCache(vllm_config) - else: - raise ValueError(f"Unknown cache type: {cache_type!r}") - - def processor_only_cache_from_config( - self, - vllm_config: "VllmConfig", - ) -> MultiModalProcessorOnlyCache | None: - """Return a `MultiModalProcessorOnlyCache`, if enabled.""" - cache_type = self._get_cache_type(vllm_config) - if cache_type is None: - return None - - return MultiModalProcessorOnlyCache(vllm_config.model_config) - - def engine_receiver_cache_from_config( - self, - vllm_config: "VllmConfig", - ) -> BaseMultiModalReceiverCache | None: - """Return a `BaseMultiModalReceiverCache` for the engine process.""" - cache_type = self._get_cache_type(vllm_config) - if cache_type in (None, "processor_only", "shm"): - return None - elif cache_type == "lru": - return MultiModalReceiverCache(vllm_config.model_config) - else: - raise ValueError(f"Unknown cache type: {cache_type!r}") - - def worker_receiver_cache_from_config( - self, - vllm_config: "VllmConfig", - shared_worker_lock: LockType, - ) -> BaseMultiModalReceiverCache | None: - """Return a `BaseMultiModalReceiverCache` for the worker process.""" - cache_type = self._get_cache_type(vllm_config) - if cache_type in (None, "processor_only", "lru"): - return None - elif cache_type == "shm": - return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock) - else: - raise ValueError(f"Unknown cache type: {cache_type!r}") - class MultiModalTimingRegistry: def __init__(self, observability_config: "ObservabilityConfig | None") -> None: diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 9b769686a7cb..53e3f50b4e8e 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -32,7 +32,11 @@ ) from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY as mm_registry -from vllm.multimodal.cache import BaseMultiModalProcessorCache +from vllm.multimodal.cache import ( + BaseMultiModalProcessorCache, + processor_cache_from_config, + processor_only_cache_from_config, +) from vllm.multimodal.gpu_ipc_memory import maybe_init_mm_gpu_ipc_pool from vllm.multimodal.parse import ( MultiModalDataItems, @@ -149,7 +153,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: self._mm_cache_stats: MultiModalCacheStats | None = None - if mm_registry.supports_multimodal_inputs(config.model_config): + if config.model_config.supports_multimodal_inputs: # Install the process-global GPU memory pool used to gate # frontend GPU-side multimodal decoding (no-op when the budget # is 0). Lives in the API-server process only. @@ -166,10 +170,8 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None: tokenizer=self.tokenizer, ) - self._mm_processor_cache = mm_registry.processor_cache_from_config(config) - self._mm_processor_only_cache = ( - mm_registry.processor_only_cache_from_config(config) - ) + self._mm_processor_cache = processor_cache_from_config(config) + self._mm_processor_only_cache = processor_only_cache_from_config(config) # This is used to generate internal request ID for MM processing # It has no relation to the request ID for engine core diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7d8aee55dc6e..36a09f9adcc4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -234,9 +234,7 @@ def __init__( # Encoder-related. # Calculate encoder cache size if applicable - supports_mm_inputs = mm_registry.supports_multimodal_inputs( - vllm_config.model_config - ) + supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs mm_budget = ( MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index a51278b7617c..9b27f50edc03 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -31,8 +31,10 @@ from vllm.logger import init_logger from vllm.logging_utils.dump_input import dump_engine_exception from vllm.lora.request import LoRARequest -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.cache import MultiModalCacheMissError +from vllm.multimodal.cache import ( + MultiModalCacheMissError, + engine_receiver_cache_from_config, +) from vllm.tasks import POOLING_TASKS, SupportedTask from vllm.tracing import instrument, maybe_init_worker_tracer from vllm.transformers_utils.config import maybe_register_config_serialize_by_value @@ -182,10 +184,7 @@ def __init__( if self.scheduler.ec_connector is not None: # type: ignore self.model_executor.init_ec_output_aggregator() - mm_registry = MULTIMODAL_REGISTRY - self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config( - vllm_config - ) + self.mm_receiver_cache = engine_receiver_cache_from_config(vllm_config) # If a KV connector is initialized for scheduler, we want to collect # handshake metadata from all workers so the connector in the scheduler diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 06b2b32efd2a..d0ccb1dacf4c 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -62,7 +62,7 @@ def __init__( self.renderer = renderer or renderer_from_config(vllm_config) - self.supports_mm_inputs = mm_registry.supports_multimodal_inputs(model_config) + self.supports_mm_inputs = model_config.supports_multimodal_inputs self.mm_encoder_cache_size = 0 self.skip_prompt_length_check = False if self.supports_mm_inputs: diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f7ad68a88a6..2affce2fdd55 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -35,7 +35,6 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM -from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backend import CommonAttentionMetadata @@ -165,10 +164,7 @@ def __init__( self.max_positions = self.max_num_tokens # Multi-modal data support - self.mm_registry = MULTIMODAL_REGISTRY - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( - vllm_config.model_config - ) + self.supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs self.draft_attn_groups: list[AttentionGroup] = [] self.kv_cache_gid: int = -1 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1ba40c5fdae5..92a37008cebe 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -257,9 +257,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Multimodal self.mm_registry = MULTIMODAL_REGISTRY - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( - self.model_config - ) + self.supports_mm_inputs = self.model_config.supports_multimodal_inputs self.uses_inputs_embeds = ( self.supports_mm_inputs or self.model_config.enable_prompt_embeds ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 534b6ef8eb17..244fa8bc18fb 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -15,7 +15,6 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.models import supports_multimodal_embeddings -from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.spec_decode import ( @@ -226,9 +225,7 @@ def load_model(self, target_model: nn.Module) -> None: ) self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names - target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs( - self.vllm_config.model_config - ) + target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs draft_supports_mm = supports_multimodal_embeddings(self.model) self.supports_mm_inputs = target_supports_mm and draft_supports_mm if target_supports_mm and not draft_supports_mm: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bdabe398f676..00e33b47c609 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -550,9 +550,7 @@ def __init__( self.mm_registry = MULTIMODAL_REGISTRY self.uses_mrope = model_config.uses_mrope self.mrope_num_dims = model_config.mrope_num_dims - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( - model_config - ) + self.supports_mm_inputs = model_config.supports_multimodal_inputs if self.model_config.is_encoder_decoder: # Maximum length of the encoder input, only for encoder-decoder diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index cf869edb26ff..b0839bc4cae2 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -12,7 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import worker_receiver_cache_from_config from vllm.tracing import instrument from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.system_utils import update_environment_variables @@ -315,26 +315,10 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: ) shared_worker_lock = kwargs.pop("shared_worker_lock", None) - if shared_worker_lock is None: - msg = ( - "Missing `shared_worker_lock` argument from executor. " - "This argument is needed for mm_processor_cache_type='shm'." - ) - - mm_config = vllm_config.model_config.multimodal_config - if mm_config and mm_config.mm_processor_cache_type == "shm": - raise ValueError(msg) - else: - logger.warning_once(msg) - - self.mm_receiver_cache = None - else: - self.mm_receiver_cache = ( - MULTIMODAL_REGISTRY.worker_receiver_cache_from_config( - vllm_config, - shared_worker_lock, - ) - ) + self.mm_receiver_cache = worker_receiver_cache_from_config( + vllm_config, + shared_worker_lock, + ) with set_current_vllm_config(self.vllm_config): # To make vLLM config available during worker initialization From 887cf91e30545373f90eb967c981b456c9cbbf8d Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 21 Sep 2026 09:46:39 -0500 Subject: [PATCH 586/774] [Pooling] Fix normalization of chunked long-text embeddings (#57498) Signed-off-by: Taneem Ibrahim --- .../entrypoints/pooling/embed/io_processor.py | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 57728c88377c..3c5b75cee62c 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -57,6 +57,8 @@ class _ChunkedPromptAggregator: weighted_sum: torch.Tensor | None = None total_weight: int = 0 + num_chunks: int = 0 + all_chunks_l2_normalized: bool = True class EmbedIOProcessor(PoolingIOProcessor): @@ -204,7 +206,17 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None: weight = len(result.prompt_token_ids) embedding_data = result.outputs.data - weighted_embedding = embedding_data.to(dtype=torch.float32) * weight + float_embedding = embedding_data.to(dtype=torch.float32) + if ctx.pooling_params.use_activation: + # Custom activations need not L2-normalize; allow BF16 rounding. + embedding_norm = torch.linalg.vector_norm(float_embedding) + aggregator.all_chunks_l2_normalized &= torch.allclose( + embedding_norm, + torch.ones_like(embedding_norm), + rtol=1e-2, + atol=1e-5, + ) + weighted_embedding = float_embedding * weight if aggregator.weighted_sum is None: # First chunk @@ -214,6 +226,7 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None: aggregator.weighted_sum += weighted_embedding aggregator.total_weight += weight + aggregator.num_chunks += 1 if ctx.original_engine_inputs is None: raise ValueError("Original engine inputs not available") @@ -239,6 +252,14 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None: ): # Compute final mean embedding final_embedding = weighted_sum / total_weight + if ( + ctx.pooling_params.use_activation + and aggregator.num_chunks > 1 + and aggregator.all_chunks_l2_normalized + ): + final_embedding = torch.nn.functional.normalize( + final_embedding, dim=-1 + ) # Create a PoolingRequestOutput # for the aggregated result From 21aa17c1289f9c376c941b5a5e3d140590d4cae1 Mon Sep 17 00:00:00 2001 From: sfigs Date: Mon, 21 Sep 2026 11:55:38 -0300 Subject: [PATCH 587/774] [Bugfix][Frontend] Accept diarized transcription responses in run-batch (#57948) Signed-off-by: Sergio Figueras --- vllm/entrypoints/launchers/run_batch.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py index 8c05a1201f9c..b5732ba65a16 100644 --- a/vllm/entrypoints/launchers/run_batch.py +++ b/vllm/entrypoints/launchers/run_batch.py @@ -53,8 +53,7 @@ ) from vllm.entrypoints.speech_to_text.transcription.protocol import ( TranscriptionRequest, - TranscriptionResponse, - TranscriptionResponseVerbose, + TranscriptionResponseVariant, ) from vllm.entrypoints.speech_to_text.translation.protocol import ( TranslationRequest, @@ -203,8 +202,7 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo): | EmbeddingResponse | ScoreResponse | RerankResponse - | TranscriptionResponse - | TranscriptionResponseVerbose + | TranscriptionResponseVariant | TranslationResponse | TranslationResponseVerbose ) @@ -660,8 +658,7 @@ def wrapper(handler_fn: Callable): async def transcription_wrapper( batch_request_body: (BatchTranscriptionRequest | BatchTranslationRequest), ) -> ( - TranscriptionResponse - | TranscriptionResponseVerbose + TranscriptionResponseVariant | TranslationResponse | TranslationResponseVerbose | ErrorResponse From e34685dfc0c9f46ae58650fe7b08cc89b615a5fe Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Mon, 21 Sep 2026 18:03:41 +0300 Subject: [PATCH 588/774] [ROCm][DSv4] Enable DSpark adaptive verification (#52362) Signed-off-by: Tuukka Sarvi Signed-off-by: larryli2-amd Co-authored-by: larryli2-amd --- .../attention/test_rocm_triton_attn_dsv4.py | 260 +++++++++++++++++- tests/models/test_deepseek_v4_dspark_rocm.py | 130 +++++++++ .../test_deepseek_v4_rocm_adaptive.py | 249 +++++++++++++++++ vllm/models/deepseek_v4/amd/dspark.py | 33 ++- vllm/models/deepseek_v4/amd/rocm.py | 31 +++ vllm/v1/attention/backends/mla/indexer.py | 29 +- 6 files changed, 725 insertions(+), 7 deletions(-) create mode 100644 tests/models/test_deepseek_v4_dspark_rocm.py create mode 100644 tests/v1/attention/test_deepseek_v4_rocm_adaptive.py diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 256b50d49a3b..5b68732f7a5c 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -188,6 +188,7 @@ def _ref_sparse_decode_ragged( extra_rows: list[list[int]] | None = None, main_use_fnuz: bool = False, extra_use_fnuz: bool = False, + extra_block_size: int | None = None, ) -> torch.Tensor: q_f32 = q.float() out = torch.empty_like(q_f32) @@ -209,7 +210,10 @@ def _ref_sparse_decode_ragged( ) row_kv.append( _read_fp8_ds_mla_cache_rows( - extra_cache, extra_slots, block_size, extra_use_fnuz + extra_cache, + extra_slots, + extra_block_size or block_size, + extra_use_fnuz, ) ) @@ -244,6 +248,12 @@ def _ragged_from_rows( ) +def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]: + ends = indptr.cpu().tolist() + values = indices[: ends[-1]].cpu().tolist() + return [values[start:end] for start, end in zip(ends, ends[1:])] + + def _launch_sparse_decode_reduce( part_m: torch.Tensor, part_l: torch.Tensor, @@ -1192,6 +1202,254 @@ def run_decode() -> torch.Tensor: torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: + """Replay the production ROCm decode path after device-only reallocation.""" + from tests.v1.attention.utils import create_vllm_config + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + ) + from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec + + device = torch.device("cuda") + num_reqs = 3 + upper_query_len = 8 + graph_tokens = num_reqs * upper_query_len + block_size = 256 + compressed_block_size = block_size // 128 + window_size = 32 + + vllm_config = create_vllm_config( + model_name="facebook/opt-125m", + max_model_len=1024, + block_size=block_size, + max_num_seqs=num_reqs, + max_num_batched_tokens=graph_tokens, + hf_config_override={ + "compress_ratios": [128], + "index_topk": 2048, + "sliding_window": window_size, + }, + ) + vllm_config.speculative_config = SimpleNamespace( + num_speculative_tokens=upper_query_len - 1, + parallel_drafting=False, + enable_adaptive_verification=True, + use_dspark=lambda: True, + ) + mla_spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=HEAD_DIM, + dtype=torch.bfloat16, + tokens_per_state=128, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + swa_spec = SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + head_size=HEAD_DIM, + dtype=torch.bfloat16, + sliding_window=window_size, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder( + mla_spec, ["c128a"], vllm_config, device + ) + swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder( + swa_spec, ["c128a"], vllm_config, device + ) + assert ( + mla_builder.get_cudagraph_support(vllm_config, mla_spec) + == AttentionCGSupport.ALWAYS + ) + assert ( + swa_builder.get_cudagraph_support(vllm_config, swa_spec) + == AttentionCGSupport.ALWAYS + ) + + seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device) + seq_lens_cpu = seq_lens.cpu() + cpu_query_start_loc = torch.arange( + 0, + graph_tokens + 1, + upper_query_len, + dtype=torch.int32, + ) + max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size + block_table = torch.arange( + num_reqs * max_blocks, dtype=torch.int32, device=device + ).view(num_reqs, max_blocks) + + def build_metadata(query_lens: list[int]): + query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device) + query_start_loc = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + query_lens_tensor.cumsum(0), + ] + ) + active_tokens = sum(query_lens) + positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device) + position_rows = [ + torch.arange( + seq_len - query_len, + seq_len, + dtype=torch.int64, + device=device, + ) + for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens) + ] + positions[:active_tokens] = torch.cat(position_rows) + slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device) + slot_mapping[:active_tokens] = torch.arange( + active_tokens, dtype=torch.int64, device=device + ) + common = CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=cpu_query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=seq_lens_cpu, + num_reqs=num_reqs, + num_actual_tokens=graph_tokens, + max_query_len=upper_query_len, + max_seq_len=int(seq_lens_cpu.max()), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + positions=positions, + causal=True, + ) + return ( + mla_builder.build_for_cudagraph_capture(common), + swa_builder.build_for_cudagraph_capture(common), + ) + + mla_metadata, swa_metadata = build_metadata([8, 8, 8]) + assert mla_metadata.for_cudagraph_capture + metadata_ptrs = ( + mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(), + mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(), + swa_metadata.decode_swa_ragged_indices.data_ptr(), + swa_metadata.decode_swa_ragged_indptr.data_ptr(), + ) + + torch.manual_seed(19) + num_heads = 16 + q = ( + torch.randn( + graph_tokens, + num_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + num_cache_blocks = num_reqs * max_blocks + swa_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_cache_blocks * block_size, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + block_size, + use_fnuz=False, + ) + compressed_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_cache_blocks * compressed_block_size, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + compressed_block_size, + use_fnuz=False, + ) + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + out = torch.empty_like(q) + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1) + + def run_decode(mla_md, swa_md) -> None: + mod.rocm_sparse_attn_decode( + q=q, + kv_cache=compressed_cache, + swa_k_cache=swa_cache, + swa_only=False, + topk_indices=mla_md.c128a_global_decode_topk_indices, + topk_lens=mla_md.c128a_decode_topk_lens, + swa_indices=swa_md.decode_swa_indices, + swa_lens=swa_md.decode_swa_lens, + swa_ragged_indices=swa_md.decode_swa_ragged_indices, + swa_ragged_indptr=swa_md.decode_swa_ragged_indptr, + topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices, + topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr, + attn_sink=attn_sink, + scale=HEAD_DIM**-0.5, + head_dim=HEAD_DIM, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + output=out, + extra_cache_nan_free=True, + adaptive_splits=True, + ) + + def reference(mla_md, swa_md) -> torch.Tensor: + return _ref_sparse_decode_ragged( + q=q, + main_cache=swa_cache, + main_rows=_rows_from_ragged( + swa_md.decode_swa_ragged_indices, + swa_md.decode_swa_ragged_indptr, + ), + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=compressed_cache, + extra_rows=_rows_from_ragged( + mla_md.c128a_decode_topk_ragged_indices, + mla_md.c128a_decode_topk_ragged_indptr, + ), + extra_block_size=compressed_block_size, + ) + + run_decode(mla_metadata, swa_metadata) + torch.accelerator.synchronize() + expected_full = reference(mla_metadata, swa_metadata) + torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_decode(mla_metadata, swa_metadata) + torch.accelerator.synchronize() + captured_full = out.clone() + + reallocated_mla, reallocated_swa = build_metadata([3, 8, 5]) + assert metadata_ptrs == ( + reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(), + reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(), + reallocated_swa.decode_swa_ragged_indices.data_ptr(), + reallocated_swa.decode_swa_ragged_indptr.data_ptr(), + ) + expected_reallocated = reference(reallocated_mla, reallocated_swa) + run_decode(reallocated_mla, reallocated_swa) + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2) + assert not torch.equal(captured_full, expected_reallocated) + + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py new file mode 100644 index 000000000000..efa97e10d7f9 --- /dev/null +++ b/tests/models/test_deepseek_v4_dspark_rocm.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.deepseek_v4.amd import dspark as dspark_module +from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM + + +def _make_uninitialized_model(confidence_head): + model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM) + object.__setattr__( + model, + "model", + SimpleNamespace(confidence_head=confidence_head), + ) + return model + + +def _prepare_loader_model(model, named_parameters): + object.__setattr__( + model, + "config", + SimpleNamespace( + n_routed_experts=1, + expert_dtype="fp4", + num_attention_heads=1, + ), + ) + object.__setattr__(model, "named_parameters", lambda: named_parameters) + + +def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + dspark_module, + "fused_moe_make_expert_params_mapping", + lambda *args, **kwargs: [], + ) + monkeypatch.setattr( + dspark_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head(): + model = _make_uninitialized_model(object()) + + assert ( + model._remap_dspark_name("mtp.2.confidence_head.proj.weight") + == "model.confidence_head.proj.weight" + ) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head(): + model = _make_uninitialized_model(None) + + assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head( + monkeypatch: pytest.MonkeyPatch, +): + model = _make_uninitialized_model(object()) + _prepare_loader_model(model, []) + _disable_distributed_loader_paths(monkeypatch) + + assert model.load_weights([]) == set() + assert model.model.confidence_head is None + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head( + monkeypatch: pytest.MonkeyPatch, +): + class FakeParameter: + loaded_weight = None + + def weight_loader(self, param, loaded_weight): + assert param is self + self.loaded_weight = loaded_weight + + confidence_head = object() + parameter = FakeParameter() + model = _make_uninitialized_model(confidence_head) + _prepare_loader_model( + model, + [("model.confidence_head.proj.weight", parameter)], + ) + _disable_distributed_loader_paths(monkeypatch) + loaded_weight = torch.tensor([[1.0, 2.0]]) + + loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)]) + + assert loaded == {"model.confidence_head.proj.weight"} + assert parameter.loaded_weight is loaded_weight + assert model.model.confidence_head is confidence_head + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_confidence_is_probability(): + class ConfidenceHead: + def __call__(self, head_hidden, markov_embed): + return (head_hidden[:, 0] + markov_embed[:, 0]).float() + + model = _make_uninitialized_model(ConfidenceHead()) + head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16) + markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16) + + confidence = model.compute_confidence(head_hidden, markov_embed) + + torch.testing.assert_close( + confidence, + torch.sigmoid(torch.tensor([0.0, -1.0])), + ) + assert torch.all((confidence >= 0) & (confidence <= 1)) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_confidence_requires_a_head(): + model = _make_uninitialized_model(None) + empty = torch.zeros((1, 1)) + + with pytest.raises(RuntimeError, match="confidence_head"): + model.compute_confidence(empty, empty) diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py new file mode 100644 index 000000000000..914cd249fea2 --- /dev/null +++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, +) +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.mla import indexer +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, +) + + +def _make_indexer_config(*, architecture: str, adaptive: bool): + return SimpleNamespace( + model_config=SimpleNamespace(architectures=[architecture]), + num_speculative_tokens=1, + speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive), + ) + + +def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None: + # Platform predicates are mutually exclusive in production. Override both + # so these ROCm policy tests do not inherit CUDA capabilities from the CI + # host running them. + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False) + monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True) + + +def _make_indexer_builder(*, adaptive: bool, capacity: int = 12): + builder = DeepseekV32IndexerMetadataBuilder.__new__( + DeepseekV32IndexerMetadataBuilder + ) + builder.vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive) + ) + builder.supports_varlen = False + builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32) + builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32) + builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32) + builder.arange_buffer = torch.arange(capacity, dtype=torch.int32) + return builder + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs(): + adaptive_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=True) + ) + fixed_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=False) + ) + + for builder_cls in ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + ): + assert ( + builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace()) + == AttentionCGSupport.ALWAYS + ) + assert ( + builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace()) + == AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch): + _mock_rocm_platform(monkeypatch) + adaptive_config = _make_indexer_config( + architecture="DeepseekV4ForCausalLM", adaptive=True + ) + fixed_config = _make_indexer_config( + architecture="DeepseekV4ForCausalLM", adaptive=False + ) + + assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch() + assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config) + assert ( + DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder + ) + assert indexer._use_flattening(adaptive_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + adaptive_config, SimpleNamespace() + ) + == AttentionCGSupport.ALWAYS + ) + assert not indexer._use_flattening(fixed_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + fixed_config, SimpleNamespace() + ) + == AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "architecture", + ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"], +) +def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4( + monkeypatch: pytest.MonkeyPatch, + architecture: str, +): + _mock_rocm_platform(monkeypatch) + adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True) + + assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config) + assert not indexer._use_flattening(adaptive_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + adaptive_config, SimpleNamespace() + ) + == AttentionCGSupport.UNIFORM_BATCH + ) + + if architecture == "DeepseekV41ForCausalLM": + assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch() + + +@pytest.mark.cpu_test +def test_rocm_adaptive_indexer_preserves_single_request_uniform_path( + monkeypatch: pytest.MonkeyPatch, +): + builder = _make_indexer_builder(adaptive=True, capacity=8) + + class FakeUniformKernel: + called = False + + def __call__( + self, + seq_lens, + decode_seq_lens, + block_table, + expanded_block_table, + decode_lens, + num_decode_tokens, + max_decode_len, + ): + self.called = True + assert num_decode_tokens == 4 + assert block_table.stride(0) == expanded_block_table.stride(0) == 2 + decode_seq_lens[:max_decode_len] = torch.arange( + seq_lens[0] - max_decode_len + 1, + seq_lens[0] + 1, + dtype=torch.int32, + ) + expanded_block_table[:max_decode_len] = block_table[0] + decode_lens[:max_decode_len] = 1 + + fake_kernel = FakeUniformKernel() + monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel) + + seq_lens, block_table, decode_lens, batch_size, requires_padding = ( + builder._prepare_decode_tensors( + seq_lens=torch.tensor([10], dtype=torch.int32), + block_table=torch.tensor([[1, 2]], dtype=torch.int32), + decode_lens=torch.tensor([4], dtype=torch.int32), + decode_lens_cpu=torch.tensor([4], dtype=torch.int32), + query_start_loc=torch.tensor([0], dtype=torch.int32), + num_decodes=1, + num_decode_tokens=4, + use_native=False, + next_n=8, + max_decode_len=4, + ) + ) + + assert fake_kernel.called + torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32)) + torch.testing.assert_close( + block_table, + torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32), + ) + torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32)) + assert batch_size == 4 + assert not requires_padding + + +@pytest.mark.cpu_test +def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers(): + builder = _make_indexer_builder(adaptive=True) + buffer_ptrs = ( + builder.decode_seq_lens_buffer.data_ptr(), + builder.expanded_block_table_buffer.data_ptr(), + builder.decode_lens_buffer.data_ptr(), + ) + block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) + decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32) + + first = builder._prepare_decode_tensors( + seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), + block_table=block_table, + decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32), + decode_lens_cpu=decode_lens_cpu, + query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32), + num_decodes=3, + num_decode_tokens=10, + use_native=False, + next_n=8, + max_decode_len=4, + ) + torch.testing.assert_close( + first[0], + torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close( + first[1][:, 0], + torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32), + ) + + second = builder._prepare_decode_tensors( + seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), + block_table=block_table, + decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32), + decode_lens_cpu=decode_lens_cpu, + query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32), + num_decodes=3, + num_decode_tokens=10, + use_native=False, + next_n=8, + max_decode_len=4, + ) + torch.testing.assert_close( + second[0], + torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close( + second[1][:, 0], + torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32)) + assert second[3:] == (10, False) + assert buffer_ptrs == ( + builder.decode_seq_lens_buffer.data_ptr(), + builder.expanded_block_table_buffer.data_ptr(), + builder.decode_lens_buffer.data_ptr(), + ) diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index edca59032aa3..c3f249fa80d5 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -44,6 +44,7 @@ ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.qwen3_dspark import ( + DSparkConfidenceHead, DSparkMarkovHead, ) from vllm.model_executor.models.utils import maybe_prefix @@ -102,7 +103,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ] ) - # Heads: final norm + hc_head, and the Markov head + # Heads: final norm + hc_head, and the Markov + confidence heads # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) hc_dim = self.hc_mult * config.hidden_size @@ -125,6 +126,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config.dspark_markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(config, "enable_confidence_head", True): + self.confidence_head = DSparkConfidenceHead( + config.hidden_size + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "confidence_head"), + ) # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch), # replacing the direct nvidia tilelang kernel call. @@ -358,6 +365,17 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) + def compute_confidence( + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor + ) -> torch.Tensor: + """Per-position acceptance probability for each drafted token.""" + if self.model.confidence_head is None: + raise RuntimeError( + "compute_confidence() requires a confidence head, but the " + "checkpoint did not provide confidence_head weights." + ) + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed)) + # --- Weight loading ---------------------------------------------------- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -391,6 +409,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + loaded_confidence_head = False tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() @@ -403,6 +422,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if mapped is None: continue name = mapped + if "confidence_head." in name: + loaded_confidence_head = True # ``.scale`` -> per-method scale suffix. if name.endswith(".scale"): @@ -439,8 +460,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue # Stacked rules only apply to decoder-layer weights. Head-stack params - # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g. - # "markov_w1" would collide with the "w1" shard rule. + # (main_proj/norm/hc_head/markov_head/confidence_head) load directly — + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule. is_layer_param = name.startswith("model.layers.") for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if not is_layer_param or weight_name not in name: @@ -469,6 +490,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) + if self.model.confidence_head is not None and not loaded_confidence_head: + self.model.confidence_head = None logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -482,8 +505,7 @@ def _remap_dspark_name(self, name: str) -> str | None: return None stage = int(m.group(1)) rest = m.group(2) - # The confidence head is not wired into inference yet; drop its weights. - if rest.startswith("confidence_head."): + if rest.startswith("confidence_head.") and self.model.confidence_head is None: return None # Head-stack params live at model level (mtp.last), context combiner at # model level (mtp.0); everything else is a per-layer decoder block. @@ -493,6 +515,7 @@ def _remap_dspark_name(self, name: str) -> str | None: "hc_head_base", "hc_head_scale", "markov_head.", + "confidence_head.", ) if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( head_prefixes diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index a70a23e3ab79..b9478b508f23 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -8,6 +8,7 @@ import torch from vllm import envs +from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -26,6 +27,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( + AttentionCGSupport, CommonAttentionMetadata, ) from vllm.v1.attention.backends.mla.sparse_swa import ( @@ -38,6 +40,7 @@ rocm_sparse_attn_decode, rocm_sparse_attn_prefill, ) +from vllm.v1.kv_cache_interface import KVCacheSpec from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) @@ -445,6 +448,20 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.enable_adaptive_verification: + # All per-token metadata is built from device query boundaries into + # persistent buffers, so adaptive verification can replay varlen + # FULL decode graphs after reallocating drafts across requests. + return AttentionCGSupport.ALWAYS + return super().get_cudagraph_support(vllm_config, kv_cache_spec) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None @@ -512,6 +529,20 @@ def build_for_cudagraph_capture( class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.enable_adaptive_verification: + # SWA indices, lengths, and token-to-request mappings are built from + # device boundaries into persistent buffers, so adaptive verification + # can replay varlen FULL decode graphs safely. + return AttentionCGSupport.ALWAYS + return super().get_cudagraph_support(vllm_config, kv_cache_spec) + # Keep fused multi-step decode disabled until update_draft_decode_metadata() # also refreshes the ROCm-specific ragged SWA indices and indptrs. supports_draft_decode_metadata_update = False diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 00d14ceeb7eb..1ede0442bc72 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -237,6 +237,15 @@ def get_builder_cls() -> type["KpoolTailMetadataBuilder"]: # type: ignore[overr class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # ROCm runs adaptive verification through the per-token flattened + # indexer path, which derives row ownership from device decode lengths. + return ( + current_platform.is_rocm() + or super().supports_device_cpu_query_lens_mismatch() + ) + @staticmethod def get_name() -> str: return "DEEPSEEK_V4_INDEXER" @@ -254,6 +263,12 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # The ROCm flattened-query support above is validated for the + # DeepSeek-V4 adaptive DSpark path only. + return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch() + @staticmethod def get_name() -> str: return "DEEPSEEK_V41_INDEXER" @@ -802,6 +817,15 @@ def _supports_flattened_device_query_lens() -> bool: ) +def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool: + model_config = vllm_config.model_config + return ( + current_platform.is_rocm() + and model_config is not None + and "DeepseekV4ForCausalLM" in model_config.architectures + ) + + def _supports_native_decode(next_n: int) -> bool: """Whether decode can pass `next_n` Q rows per request to the kernel instead of flattening to one single-token row per query, which re-reads @@ -822,7 +846,10 @@ def _use_flattening(vllm_config: VllmConfig) -> bool: return not _supports_native_decode(next_n) or ( speculative_config is not None and speculative_config.enable_adaptive_verification - and _supports_flattened_device_query_lens() + and ( + _supports_flattened_device_query_lens() + or _rocm_supports_flattened_device_query_lens(vllm_config) + ) ) From dc479a2d839d5c474b36ac521dc5a72294bcb186 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Mon, 21 Sep 2026 23:19:14 +0800 Subject: [PATCH 589/774] [Bugfix] Fix external LB DP rank handling when replicas share nodes (#53743) Signed-off-by: Tony Lin --- docs/serving/data_parallel_deployment.md | 8 +++ tests/test_config.py | 35 +++++++++++ tests/v1/engine/test_engine_args.py | 75 ++++++++++++++++++++++++ vllm/config/parallel.py | 20 ++++++- vllm/engine/arg_utils.py | 22 +++++-- 5 files changed, 153 insertions(+), 7 deletions(-) diff --git a/docs/serving/data_parallel_deployment.md b/docs/serving/data_parallel_deployment.md index 3b92fa06a5ee..d1f9b4a4624d 100644 --- a/docs/serving/data_parallel_deployment.md +++ b/docs/serving/data_parallel_deployment.md @@ -126,6 +126,14 @@ vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \ --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345 ``` +When using `--data-parallel-external-lb` with `--nnodes`, the data parallel rank +can be inferred from `--node-rank` only when `--nnodes` is evenly divisible by +`--data-parallel-size`. If it is not evenly divisible, automatic rank inference +is rejected and `--data-parallel-rank` must be set explicitly. When multiple +external-LB ranks are co-located on a node, assign each process a disjoint +device set with `CUDA_VISIBLE_DEVICES` or `--device-ids`; devices are not +partitioned automatically between processes. + The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
diff --git a/tests/test_config.py b/tests/test_config.py index cad5ffd627fd..e97e2cebaad2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1518,6 +1518,41 @@ def test_reconfigure_for_independent_dp_rank_on_multinode_dense_model(): assert parallel_config.world_size == 8 +@pytest.mark.parametrize("data_parallel_size", [4, 8]) +def test_nnodes_within_dp_when_replicas_outnumber_nodes(data_parallel_size): + """A replica that fits on one node spans one node, never zero. + + External LB pins ``data_parallel_size_local`` to 1, so the ratio rounds + down to 0 as soon as there are more DP replicas than nodes. Anything + dividing by ``nnodes_within_dp`` then raises ZeroDivisionError. + """ + parallel_config = ParallelConfig( + data_parallel_size=data_parallel_size, + data_parallel_size_local=1, + data_parallel_external_lb=True, + distributed_executor_backend="mp", + nnodes=2, + node_rank=1, + ) + + assert parallel_config.nnodes_within_dp == 1 + assert parallel_config.node_rank_within_dp == 0 + assert parallel_config.local_world_size == parallel_config.world_size + + +@pytest.mark.parametrize("nnodes", [4, 9]) +def test_nnodes_within_dp_rejects_uneven_internal_lb(nnodes): + parallel_config = ParallelConfig( + data_parallel_size=8, + data_parallel_size_local=1, + distributed_executor_backend="mp", + nnodes=nnodes, + ) + + with pytest.raises(ValueError, match="Invalid data parallel configuration"): + _ = parallel_config.nnodes_within_dp + + def test_draft_model_enables_async_scheduling_by_default(): parallel_config = ParallelConfig(distributed_executor_backend="uni") model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048) diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index b1ba1adcfd0a..3f9f617bc7ed 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from argparse import ArgumentError +from unittest.mock import patch import pytest +from vllm.config import ModelConfig from vllm.engine.arg_utils import EngineArgs from vllm.usage.usage_lib import UsageContext from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -121,3 +123,76 @@ def test_data_parallel_start_rank_zero_infers_hybrid_lb(): assert vllm_config.parallel_config.data_parallel_hybrid_lb is True assert vllm_config.parallel_config.data_parallel_rank == 0 + + +def test_external_lb_preserves_explicit_rank_when_dp_exceeds_nodes(): + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=4, + data_parallel_rank=3, + data_parallel_external_lb=True, + nnodes=2, + node_rank=1, + ) + + with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)): + vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + assert vllm_config.parallel_config.data_parallel_rank == 3 + + +@pytest.mark.parametrize( + ("data_parallel_size", "nnodes", "tensor_parallel_size"), + [(4, 2, 1), (2, 3, 3)], +) +def test_external_lb_requires_explicit_rank_when_nodes_are_not_evenly_partitioned( + data_parallel_size, nnodes, tensor_parallel_size +): + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=data_parallel_size, + data_parallel_external_lb=True, + tensor_parallel_size=tensor_parallel_size, + nnodes=nnodes, + node_rank=1, + ) + + with ( + patch.object(ModelConfig, "is_moe", new=property(lambda self: True)), + pytest.raises( + ValueError, + match="Set a unique `--data-parallel-rank`", + ), + ): + engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + +def test_external_lb_infers_rank_when_dp_does_not_exceed_nodes(): + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=2, + data_parallel_external_lb=True, + nnodes=2, + node_rank=1, + ) + + with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)): + vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + assert vllm_config.parallel_config.data_parallel_rank == 1 + + +def test_external_lb_infers_rank_for_multinode_replicas(): + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=2, + data_parallel_external_lb=True, + tensor_parallel_size=12, + nnodes=4, + node_rank=1, + ) + + with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)): + vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + assert vllm_config.parallel_config.data_parallel_rank == 0 diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 666f2441ae82..3929d3be1266 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -760,12 +760,30 @@ def node_rank_within_dp(self) -> int: @property def nnodes_within_dp(self) -> int: + """Number of nodes one DP replica spans. + + External LB pins ``data_parallel_size_local`` to 1, so the ratio + rounds down to 0 once DP replicas outnumber nodes. A replica that + does not span nodes still occupies exactly one. + """ if self.nnodes == 1: return 1 data_parallel_node_size = ( self.data_parallel_size // self.data_parallel_size_local ) - return self.nnodes // data_parallel_node_size + nnodes_within_dp = self.nnodes // data_parallel_node_size + if self.data_parallel_external_lb: + return max(nnodes_within_dp, 1) + if self.nnodes % data_parallel_node_size != 0: + raise ValueError( + "Invalid data parallel configuration: " + f"nnodes ({self.nnodes}) must be divisible by the number of " + "data parallel node groups " + f"({data_parallel_node_size} = data_parallel_size " + f"{self.data_parallel_size} / data_parallel_size_local " + f"{self.data_parallel_size_local})" + ) + return nnodes_within_dp @property def local_world_size(self) -> int: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 9fad681d7dd9..555f1aebb67f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2232,12 +2232,22 @@ def create_engine_config( self.node_rank * local_world_size ) // world_size_within_dp if self.data_parallel_size > 1 and self.data_parallel_external_lb: - self.data_parallel_rank = inferred_data_parallel_rank - logger.info( - "Inferred data_parallel_rank %d from node_rank %d for external lb", - self.data_parallel_rank, - self.node_rank, - ) + if self.data_parallel_rank is None: + if self.nnodes % self.data_parallel_size != 0: + raise ValueError( + "Invalid data-parallel launch options: " + "`--node-rank` cannot unambiguously identify external " + "data-parallel ranks when `--nnodes` is not divisible " + "by `--data-parallel-size`. Set a unique " + "`--data-parallel-rank` for each external-LB process." + ) + self.data_parallel_rank = inferred_data_parallel_rank + logger.info( + "Inferred data_parallel_rank %d from node_rank %d " + "for external lb", + self.data_parallel_rank, + self.node_rank, + ) elif self.data_parallel_size_local is None: # Infer data parallel size local for internal dplb: self.data_parallel_size_local = max( From 6b858751f6e99c4f33913657dc0c2afd34ff13e6 Mon Sep 17 00:00:00 2001 From: Lucain Date: Mon, 21 Sep 2026 17:32:27 +0200 Subject: [PATCH 590/774] [docs] Fix legacy hf CLI references (vllm) (#57958) Signed-off-by: Wauplin Co-authored-by: Claude Opus 5 (1M context) --- docker/Dockerfile.rock | 2 +- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.rocm_gfx1250 | 2 +- docs/getting_started/installation/gpu.rocm.inc.md | 2 +- docs/models/supported_models.md | 2 +- docs/usage/troubleshooting.md | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index 6e52199fa371..16e01ad723b9 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -914,7 +914,7 @@ RUN case "$(which python3)" in \ *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub[cli] + uv pip install --system --upgrade huggingface-hub # Install vLLM using uv (inherited from base stage) # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 483aa9bb97e3..0bcc8a6ca6f1 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -938,7 +938,7 @@ RUN case "$(which python3)" in \ *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub[cli] + uv pip install --system --upgrade huggingface-hub # Install vLLM using uv (inherited from base stage) # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index c68fdfebfe54..ecb618d99070 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -671,7 +671,7 @@ RUN case "$(which python3)" in \ *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub[cli] + uv pip install --system --upgrade huggingface-hub # Install vLLM using uv (inherited from base stage) # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index ad5af7994054..0169d0154608 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -324,7 +324,7 @@ You can find more information about vLLM's wheels in # Install dependencies pip install --upgrade numba \ scipy \ - huggingface-hub[cli] \ + huggingface-hub \ setuptools_scm pip install -r requirements/rocm.txt diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index ddc50ad41cdb..c6e1786ee080 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -240,7 +240,7 @@ Alternatively, you can [open an issue on GitHub](https://github.com/vllm-project #### Download a model -If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-download) or specific files from a model repository: +If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#hf-download) or specific files from a model repository: ```bash # Download a model diff --git a/docs/usage/troubleshooting.md b/docs/usage/troubleshooting.md index dc1cd89f8209..bbd4109a4196 100644 --- a/docs/usage/troubleshooting.md +++ b/docs/usage/troubleshooting.md @@ -8,7 +8,7 @@ This document outlines some troubleshooting strategies you can consider. If you ## Hangs downloading a model If the model isn't already downloaded to disk, vLLM will download it from the internet which can take time and depend on your internet connection. -It's recommended to download the model first using the [huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue. +It's recommended to download the model first using the [`hf` CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue. ## Hangs loading a model from disk From 8c0825090df8b7939a92ff22b02bc7d03cb3605a Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 21 Sep 2026 11:57:49 -0400 Subject: [PATCH 591/774] [Bugfix][NIXL] Fix DCP pulls across MLA cache regions (#57389) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../unit/test_nixl_connector_hma.py | 77 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 44 +++++++++-- .../kv_connector/v1/nixl/pull_worker.py | 40 +++++----- 3 files changed, 133 insertions(+), 28 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index a4dfc168f31f..3be1730ec3e0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -58,6 +58,7 @@ def region_pull_worker(): worker._recving_transfers = {} worker.use_mla, worker._has_mamba = True, False worker.dcp_size = 1 + worker.dcp_rank = 0 spec = MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 ) @@ -165,6 +166,82 @@ def test_region_pull_ignores_allocation_padding( ] +@pytest.mark.cpu_test +@pytest.mark.parametrize("num_pages", [0, 3, 19]) +@pytest.mark.parametrize("region_groups", [(0, 1), (1, 0, 1)]) +def test_dcp_region_pull(region_pull_worker, num_pages, region_groups): + """Read each uncached page from its DCP rank; notify ranks with no pages.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + RemoteMeta, + ReqMeta, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping + + worker = region_pull_worker + worker.region_group_ids = list(region_groups) + worker.num_regions = len(region_groups) + worker.dst_region_group_ids = {"P": [0] * worker.num_regions} + worker.dst_region_num_blocks = { + engine: [100] * worker.num_regions for engine in ("P", "D") + } + worker.block_len_per_layer = [1024] * worker.num_regions + worker.dcp_rank = 0 + remote = worker.transfer_topo.get_engine_info.return_value + remote.remote_tp_size = remote.remote_dcp_size = 8 + remote.remote_physical_blocks_per_logical = 1 + ranks = tuple(range(8)) + worker.tp_mappings["P"] = TPMapping((ranks, ranks), ranks, {r: r for r in ranks}, 8) + worker.dst_xfer_side_handles = {"P": {r: 1000 + r for r in ranks}} + worker._remote_agents = {"P": {(0, r): f"P-rank{r}" for r in ranks}} + # Two regions with different prefix hits, plus one padding page each. + local = ( + [list(range(31, 31 + num_pages)), list(range(42, 41 + num_pages))] + if num_pages + else [] + ) + meta = ReqMeta( + local, + local, + tp_size=8, + local_num_computed_blocks=(0, 1, 2), + remote=RemoteMeta( + [[10, 11, 12]], + "localhost", + 1, + "P", + "request-P", + num_tokens=max(0, num_pages * 64 - 7), + ), + ) + worker._read_blocks_for_req("request", meta) + reads = { + call.kwargs["remote_xfer_side_handle"] - 1000: call.kwargs + for call in worker._read_blocks_mixed.call_args_list + } + notified = {call.args[0] for call in worker.nixl_wrapper.send_notif.call_args_list} + for rank in ranks: + expected = [ + (region * 100 + 30 + group * 10 + page, region * 100 + 10 + page // 8) + for region, group in enumerate(region_groups) + for page in range(group + 1, num_pages) + if page % 8 == rank + ] + if expected: + read = reads[rank] + assert ( + list(zip(read["local_block_descs_ids"], read["remote_block_descs_ids"])) + == expected + ) + assert f"P-rank{rank}" not in notified + else: + assert rank not in reads and f"P-rank{rank}" in notified + assert meta.region_blocks_to_zero == ( + [[30 + group * 10 + num_pages] for group in region_groups] + if num_pages + else None + ) + + @pytest.mark.cpu_test def test_region_pull_completion_zeros_only_own_padding(region_pull_worker): """Completion must preserve neighboring layers and already imported pages.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index c2b20ffc597e..0f13a265dfd7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -75,6 +75,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.kv_cache_interface import ( @@ -2430,6 +2431,18 @@ def _validate_remote_agent_handshake( remote_physical_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block ) + if ( + (self.dcp_size > 1 or remote_dcp_size > 1) + and self.region_group_ids + and self.region_group_ids != self.dst_region_group_ids[remote_engine_id] + and ( + self._physical_blocks_per_logical_kv_block != 1 + or remote_physical_per_logical != 1 + ) + ): + raise NotImplementedError( + "DCP region pulls require matching logical and physical block sizes" + ) if ( self._has_mamba and remote_physical_per_logical @@ -3234,11 +3247,14 @@ def _apply_dcp_prefix_caching( local_dcp_rank: int, remote_dcp_size: int, local_num_computed_blocks: int, + num_remote_blocks: int | None = None, ) -> tuple[list[int], list[int]]: """Match local and remote block IDs by global DCP position. ``local_ids`` excludes blocks already satisfied by the local prefix cache, while ``remote_ids`` contains the full transferable remote list. + ``num_remote_blocks``, when provided, is the global block count before + DCP partitioning and excludes allocation padding. Example: local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6] @@ -3254,6 +3270,9 @@ def _apply_dcp_prefix_caching( """ local_size, remote_size = local_dcp_size, remote_dcp_size + if num_remote_blocks is not None: + num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size) + local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)] if local_size == remote_size: local_slice = local_ids @@ -3274,6 +3293,8 @@ def _apply_dcp_prefix_caching( remote_slice = remote_ids[start_remote::k] matched_blocks = min(len(local_slice), len(remote_slice)) + if num_remote_blocks is not None and matched_blocks < len(local_slice): + raise ValueError("Remote KV pages do not cover the requested range") return local_slice[:matched_blocks], remote_slice[:matched_blocks] @staticmethod @@ -3290,13 +3311,15 @@ def _block_ids_by_region( block_ids_by_region.append(list(block_ids[group_id])) return block_ids_by_region - @staticmethod def _apply_prefix_caching_by_region( + self, decode_block_ids: BlockIds, prefill_block_ids: BlockIds, *, num_computed_blocks: list[int] | None = None, num_remote_blocks: int | None = None, + remote_rank: int = 0, + remote_dcp_size: int = 1, ) -> tuple[BlockIds, BlockIds]: """Pair an uncached decode suffix with the same prefill regions.""" assert len(decode_block_ids) == len(prefill_block_ids) @@ -3310,13 +3333,18 @@ def _apply_prefix_caching_by_region( for decode_region, prefill_region, start in zip( decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True ): - count = min(len(decode_region), max(num_remote_blocks - start, 0)) - if start + count > len(prefill_region): - raise ValueError("Remote KV pages do not cover the requested range") - # An extra producer page can be allocation padding, not a - # cached prefix. Select by token position instead of list length. - matched_decode.append(list(decode_region[:count])) - matched_prefill.append(list(prefill_region[start : start + count])) + local, remote = self._apply_dcp_prefix_caching( + decode_region, + prefill_region, + remote_rank, + self.dcp_size, + self.dcp_rank, + remote_dcp_size, + start, + num_remote_blocks=num_remote_blocks, + ) + matched_decode.append(local) + matched_prefill.append(remote) return matched_decode, matched_prefill trimmed_prefill: list[list[int]] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 43fbbe0c3c25..6539b7e57ed8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -191,7 +191,6 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): "Different NIXL cache-group layouts are only supported for " "pure MLA models" ) - assert len(plan.all_source_ranks) == 1 if self.block_size != remote_info.remote_block_size: raise NotImplementedError( "Region-mapped NIXL transfers require matching physical block sizes" @@ -213,7 +212,6 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): and meta.local_num_computed_blocks and all(group >= 0 for group in local_region_groups) and all(group >= 0 for group in remote_region_groups) - and not dcp_active ): transfer_groups = self.kv_cache_config.transfer_group_ids num_computed_blocks = [ @@ -224,32 +222,34 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): num_remote_blocks = cdiv( meta.remote.num_tokens, remote_info.remote_block_size ) - elif ( - remote_info.remote_physical_blocks_per_logical + elif any(local_by_region) and ( + dcp_active + or remote_info.remote_physical_blocks_per_logical != self._physical_blocks_per_logical_kv_block ): raise NotImplementedError( - "Region-mapped pulls with different logical block sizes require " - "remote_num_tokens, per-group prefix counts, unshared regions " - "and DCP=1" + "Region-mapped pulls with DCP or different logical block sizes " + "require remote_num_tokens, per-group prefix counts " + "and unshared regions" ) - matched_local, matched_remote = self._apply_prefix_caching_by_region( - local_by_region, - remote_by_region, - num_computed_blocks=num_computed_blocks, - num_remote_blocks=num_remote_blocks, - ) - meta.region_blocks_to_zero = [ - list(blocks[len(matched) :]) - for blocks, matched in zip(local_by_region, matched_local, strict=True) - ] read_specs = [ ReadSpec( - remote_rank=plan.all_source_ranks[0], - local_block_ids=matched_local, - remote_block_ids=matched_remote, + rank, + *self._apply_prefix_caching_by_region( + local_by_region, + remote_by_region, + num_computed_blocks=num_computed_blocks, + num_remote_blocks=num_remote_blocks, + remote_rank=rank, + remote_dcp_size=remote_info.remote_dcp_size, + ), block_ids_by_region=True, ) + for rank in plan.all_source_ranks + ] + meta.region_blocks_to_zero = [ + list(blocks[sum(len(spec.local_block_ids[r]) for spec in read_specs) :]) + for r, blocks in enumerate(local_by_region) ] else: remote_logical_block_ids = meta.remote.block_ids From 874a6b5a06f19c9d565d14ebe64131d6aaab91cd Mon Sep 17 00:00:00 2001 From: afriedri Date: Mon, 21 Sep 2026 11:30:16 -0500 Subject: [PATCH 592/774] [ROCm] Refactor tuned gemms (#55001) Signed-off-by: Andy Friedrich Signed-off-by: afriedri Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/quantization/test_rocm_mxfp4.py | 14 +- vllm/_aiter_ops.py | 137 +++++++----------- 2 files changed, 52 insertions(+), 99 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 8d1a4bb0a734..d4e16629afac 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -50,16 +50,7 @@ } PRESHUFFLED_SHAPES = [ (64, 4096, 8192), - # aiter 0.1.20 Triton preshuffled fp4 GEMM OOBs (GPU memory access fault -> - # SIGABRT, crashing the process) for (M=32, N=8192, K=8192) on gfx950, despite - # advertising the (N, K) as tuned. A hard GPU fault can't be xfail'd, so skip it - # until the aiter kernel is fixed (ROCm/aiter#4867). - pytest.param( - (32, 8192, 8192), - marks=pytest.mark.skip( - reason="aiter 0.1.20 preshuffled fp4 GEMM OOB (M32,N8192,K8192, gfx950)" - ), - ), + (32, 8192, 16384), ] @@ -430,7 +421,8 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): M, K, N = shape assert M <= 64 - assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K) + # The predicate takes K as packed bytes, matching `weight.shape[1]`. + assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2) A = torch.randn(M, K, dtype=torch.bfloat16) B = torch.randn(N, K, dtype=torch.bfloat16) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 1829d9f5e247..c0b90e0be6b6 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -13,7 +13,6 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.utils.import_utils import PlaceholderModule from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_aiter_sparse_attn_indexer, @@ -25,11 +24,6 @@ logger = init_logger(__name__) -try: - import pandas as pd -except ImportError: - pd = PlaceholderModule("pandas") - # fp8_dtype is not cached. # on ROCm the fp8_dtype always calls is_fp8_fnuz # which is a host op, so we cache it once here. @@ -243,27 +237,42 @@ def is_aiter_found_and_supported_on_rdna4() -> bool: return False -@functools.cache -def _load_gemm_tuned_configs( - q_dtype_w: torch.dtype, csv_path: str -) -> set[tuple[int, int, int]]: - try: - df = pd.read_csv(csv_path).drop_duplicates() - df = df[df["q_dtype_w"] == str(q_dtype_w)] - return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int))) - except Exception: - return set() +def _triton_gemm_config_is_tuned(config_name: str, N: int, K: int) -> bool: + from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config + + # any M shape under 1024 should be tuned + M_TUNE_PROBE = 128 + return get_gemm_config(config_name, M_TUNE_PROBE, N, K)[1] -def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool: - configs = _load_gemm_tuned_configs(q_dtype_w, csv_path) + +def _ck_gemm_shape_is_tuned( + N: int, K: int, q_dtype_w: torch.dtype, csv_attr: str +) -> bool: l_m = ( [1, 2, 4] + list(range(8, 513, 8)) + [1024, 1536] + [2**i for i in range(11, 19)] ) - return any((N, K, M) in configs for M in l_m) + try: + from aiter.ops.gemm_op_a8w8 import ( + AITER_CONFIGS, + get_GEMM_config_with_quant_type, + ) + + csv_path = getattr(AITER_CONFIGS, csv_attr) + return any( + get_GEMM_config_with_quant_type(M, N, K, q_dtype_w, csv_path) is not None + for M in l_m + ) + except (AttributeError, ImportError, OSError): + logger.warning_once( + "Could not read aiter CK GEMM configs from AITER_CONFIGS.%s; " + "treating all shapes as untuned.", + csv_attr, + ) + return False def if_aiter_supported(func: Callable) -> Callable: @@ -3458,89 +3467,41 @@ def group_fp8_quant( ) @staticmethod + @functools.cache def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool: if not current_platform.is_rocm(): return False - from vllm.platforms.rocm import on_gfx950, on_rdna4 - - gfx950_tuned = { - (1024, 8192), - (2112, 7168), - (3072, 1536), - (32768, 8192), - (4096, 7168), - (4608, 7168), - (512, 7168), - (7168, 2048), - (7168, 256), - (8192, 1024), - (8192, 32768), - } - rdna4_tuned = gfx950_tuned | { - (2048, 2048), - (2624, 6144), - (3072, 6144), - (3584, 512), - (4096, 512), - (6144, 1536), - (6144, 2048), - (7168, 2304), - (7168, 16384), - (7168, 18432), - (8192, 8192), - (16384, 1536), - (24576, 1536), - (32768, 512), - (36864, 7168), - } - if on_rdna4(): - return (n, k) in rdna4_tuned - if on_gfx950(): - return (n, k) in gfx950_tuned - return False + try: + return _triton_gemm_config_is_tuned("GEMM-A8W8_BLOCKSCALE", n, k) + except (AssertionError, ImportError): + return False @staticmethod - def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool: - return (n, k) in [ - (8192, 4096), - (1280, 8192), - (16384, 53248), - (106496, 16384), - (57344, 8192), - (8192, 2048), - (2560, 8192), - (10240, 8192), - (16384, 16384), - (8192, 28672), - (28672, 8192), - (18432, 16384), - (8192, 1024), - (7168, 8192), - (5120, 8192), - (8192, 8192), - (8192, 7168), - (14336, 8192), - (8192, 14336), - (8192, 3584), - ] + @functools.cache + def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool: + if not current_platform.is_rocm(): + return False + # Input = weight.shape[1] (bytes); *4 for fp4 and aiter config 2*k + try: + return _triton_gemm_config_is_tuned( + "GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes + ) + except (AssertionError, ImportError): + return False @staticmethod + @functools.cache def is_shuffled_per_token_w8a8_gemm_tuned( N: int, K: int, q_dtype_w: torch.dtype ) -> bool: - import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops - - csv_path = ( - aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE + return _ck_gemm_shape_is_tuned( + N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE" ) - return _check_kernel_tuned(N, K, q_dtype_w, csv_path) @staticmethod + @functools.cache def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool: - import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops - - csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE - return _check_kernel_tuned(N, K, q_dtype_w, csv_path) + return _ck_gemm_shape_is_tuned(N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_FILE") @staticmethod def shuffle_weight( From 4e61ad8ee126bf01dbc4d28979014a30628309ed Mon Sep 17 00:00:00 2001 From: Sahil Patel <91423311+Sip4818@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:38:25 +0100 Subject: [PATCH 593/774] [Bugfix] unskip InternViT test for transformers v5 compatibility (#55767) Signed-off-by: sahil Co-authored-by: Cyrus Leung --- .../multimodal/pooling/test_intern_vit.py | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index 72d2a7018bf9..4d5e8728589c 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -3,7 +3,7 @@ import pytest import torch import torch.nn as nn -from transformers import AutoConfig, AutoModel, CLIPImageProcessor +from transformers import AutoConfig, AutoModel, CLIPImageProcessor, PreTrainedModel from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform @@ -12,11 +12,6 @@ from ....conftest import ImageTestAssets -pytestmark = pytest.mark.skip( - reason="InternVisionModel's custom code is incompatible with " - "transformers v5 (missing all_tied_weights_keys)" -) - # we use snapshot_download to prevent conflicts between # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] @@ -45,9 +40,23 @@ def run_intern_vit_test( if not getattr(config, "norm_type", None): config.norm_type = "rms_norm" - hf_model = AutoModel.from_pretrained( - model, dtype=torch_dtype, trust_remote_code=True - ).to(DEVICE_TYPE) + # Monkey-patch for transformers v5 compatibility: InternVisionModel's + # custom code doesn't call post_init(), so all_tied_weights_keys is + # never set. Remove once https://github.com/OpenGVLab/InternVL fixes it. + _orig_init = PreTrainedModel.__init__ + + def _patched_init(self, *args, **kwargs): + _orig_init(self, *args, **kwargs) + self.all_tied_weights_keys = {} + + PreTrainedModel.__init__ = _patched_init + + try: + hf_model = AutoModel.from_pretrained( + model, dtype=torch_dtype, trust_remote_code=True + ).to(DEVICE_TYPE) + finally: + PreTrainedModel.__init__ = _orig_init hf_outputs_per_image = [ hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state for pixel_value in pixel_values From 7853700d4cdc2b84c728c840114ed42cda7454a7 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Tue, 22 Sep 2026 00:39:20 +0800 Subject: [PATCH 594/774] [MM] Move get_dummy_processor_inputs into MM processor (#57967) Signed-off-by: DarkLight1337 --- .../multimodal/processing/test_common.py | 5 +- .../processing/test_tensor_schema.py | 3 +- tests/renderers/test_warmup.py | 14 ++--- .../models/hyperclovax_vision_v2.py | 56 ++++++++----------- vllm/model_executor/models/kimi_audio.py | 47 ++++++++-------- .../models/moss_transcribe_diarize.py | 39 +++++++------ vllm/model_executor/models/pixtral.py | 12 ++-- vllm/model_executor/models/voxtral.py | 12 ++-- vllm/multimodal/processing/dummy_inputs.py | 37 ------------ vllm/multimodal/processing/processor.py | 39 ++++++++++++- 10 files changed, 130 insertions(+), 134 deletions(-) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index db6f37dc3d5f..fa239560171d 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -159,7 +159,6 @@ def get_token_prompt( processor: BaseMultiModalProcessor, mm_data: MultiModalDataDict, ) -> list[int]: - dummy_inputs = processor.dummy_inputs tokenizer: TokenizerLike = processor.info.get_tokenizer() model_config = processor.info.ctx.model_config @@ -174,7 +173,7 @@ def get_token_prompt( mm_counts = {k: len(vs) for k, vs in parsed_data.items()} if is_mistral_tokenizer(tokenizer): - inputs = dummy_inputs.get_dummy_processor_inputs( + inputs = processor.get_dummy_inputs( model_config.max_model_len, mm_counts, mm_options=MultiModalDummyOptions(), @@ -182,7 +181,7 @@ def get_token_prompt( mm_data=mm_data, # type: ignore[call-arg] ) else: - inputs = dummy_inputs.get_dummy_processor_inputs( + inputs = processor.get_dummy_inputs( model_config.max_model_len, mm_counts, mm_options=MultiModalDummyOptions(), diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 1e1722537d4d..73be3ab76c0b 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -90,13 +90,12 @@ def create_batched_mm_kwargs( size_factors: tuple[float, ...] = (1.0, 0.5, 0.25), ) -> Iterable[tuple[str, int, BatchedTensorInputs]]: processing_info = processor.info - dummy_inputs = processor.dummy_inputs supported_mm_limits = processing_info.get_supported_mm_limits() mm_counts = { modality: 3 if limit is None else limit for modality, limit in supported_mm_limits.items() } - processor_inputs = dummy_inputs.get_dummy_processor_inputs( + processor_inputs = processor.get_dummy_inputs( seq_len=model_config.max_model_len, mm_counts=mm_counts, mm_options=MultiModalDummyOptions(), diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index 2627cdd98501..b0e3dcb54796 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -4,7 +4,7 @@ These tests exercise: - Zero-limit modalities are filtered from mm_counts passed to - get_dummy_processor_inputs (e.g. --limit-mm-per-prompt image=0 ...) + get_dummy_inputs (e.g. --limit-mm-per-prompt image=0 ...) - The warmup runs through the processor-only cache and clears it afterwards, so its dummy inputs never land in the sender cache used by the serving path - MM warmup is skipped entirely when mm_processor is None @@ -92,7 +92,7 @@ def test_zero_limit_modality_excluded_from_mm_counts(self): BaseRenderer.warmup(renderer, ChatParams()) - get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs + get_inputs = renderer.mm_processor.get_dummy_inputs get_inputs.assert_called_once() _, kwargs = get_inputs.call_args assert "video" not in kwargs["mm_counts"] @@ -104,7 +104,7 @@ def test_all_zero_limits_passes_empty_mm_counts(self): BaseRenderer.warmup(renderer, ChatParams()) - get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs + get_inputs = renderer.mm_processor.get_dummy_inputs get_inputs.assert_called_once() _, kwargs = get_inputs.call_args assert kwargs["mm_counts"] == {} @@ -115,7 +115,7 @@ def test_positive_limits_all_included_in_mm_counts(self): BaseRenderer.warmup(renderer, ChatParams()) - get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs + get_inputs = renderer.mm_processor.get_dummy_inputs get_inputs.assert_called_once() _, kwargs = get_inputs.call_args assert kwargs["mm_counts"] == {"image": 1, "video": 1} @@ -134,7 +134,7 @@ def test_warmup_sequence_length_is_bounded(self, max_model_len, expected_seq_len BaseRenderer.warmup(renderer, ChatParams()) - get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs + get_inputs = renderer.mm_processor.get_dummy_inputs assert get_inputs.call_args.kwargs["seq_len"] == expected_seq_len def test_processor_apply_called(self): @@ -207,9 +207,7 @@ def test_chat_failure_does_not_abort_mm_warmup(self): def test_mm_failure_is_swallowed_and_cache_cleared(self): renderer = _make_renderer_mock({"image": 1}) # the warmup blows up before apply() - renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs.side_effect = ( - RuntimeError("mm boom") - ) + renderer.mm_processor.get_dummy_inputs.side_effect = RuntimeError("mm boom") BaseRenderer.warmup(renderer, ChatParams()) # must not raise diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py index e269e2a05993..ee23d2c4df02 100644 --- a/vllm/model_executor/models/hyperclovax_vision_v2.py +++ b/vllm/model_executor/models/hyperclovax_vision_v2.py @@ -189,38 +189,6 @@ def get_dummy_text( num_videos = mm_counts.get("video", 0) return V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos - def get_dummy_processor_inputs( - self, - seq_len: int, - mm_counts: Mapping[str, int], - mm_options: MultiModalDummyOptions | None = None, - mm_processor_kwargs: Mapping[str, object] | None = None, - ) -> ProcessorInputs: - """Build dummy processor inputs for memory profiling.""" - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - prompt_text = V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos - - dummy_mm_data = self.get_dummy_mm_data( - seq_len, - mm_counts, - mm_options, - mm_processor_kwargs=mm_processor_kwargs, - ) - dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) - - tokenizer = self.info.get_tokenizer() - prompt = tokenizer.encode( - prompt_text, - **self.info.default_tok_params.get_encode_kwargs(), - ) - - return ProcessorInputs( - prompt=prompt, - mm_data_items=dummy_mm_items, - hf_processor_mm_kwargs=mm_processor_kwargs or {}, - ) - def get_dummy_mm_data( self, seq_len: int, @@ -255,6 +223,28 @@ class HCXVisionV2MultiModalProcessor( ): """Multimodal processor for HyperCLOVAX V2 (32B Think model).""" + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> ProcessorInputs: + num_images = mm_counts.get("image", 0) + num_videos = mm_counts.get("video", 0) + prompt_text = V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos + + builder = self.dummy_inputs + dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) + + tokenizer = self.info.get_tokenizer() + prompt = tokenizer.encode( + prompt_text, + **self.info.default_tok_params.get_encode_kwargs(), + ) + + return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items) + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) @@ -267,7 +257,7 @@ def _get_prompt_updates( hf_config = self.info.get_hf_config() # Use token IDs directly from config. - # This matches what get_dummy_processor_inputs uses, ensuring consistency. + # This matches what get_dummy_inputs uses, ensuring consistency. placeholder: dict[str, int] = { "image": hf_config.image_token_id, # 128060 for <|IMAGE_PAD|> "video": hf_config.video_token_id, # 128061 for <|VIDEO_PAD|> diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index 3541fcf371aa..a2357836f1aa 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -176,29 +176,6 @@ def get_dummy_mm_data( ), } - def get_dummy_processor_inputs( - self, - seq_len: int, - mm_counts: Mapping[str, int], - mm_options: MultiModalDummyOptions, - ) -> ProcessorInputs: - dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) - dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) - - num_audios = mm_counts.get("audio", 0) - dummy_tokens = ( - [198] - if num_audios == 0 - else [ - KimiAudioProcessor.KIMIA_MEDIA_BEGIN, - KimiAudioProcessor.KIMIA_TEXT_BLANK, - KimiAudioProcessor.KIMIA_MEDIA_END, - ] - * num_audios - ) - - return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) - # Field config for Kimi-Audio multimodal data _KIMIAUDIO_FIELD_CONFIG = { @@ -228,6 +205,30 @@ def _parse_audio_data( class KimiAudioMultiModalProcessor(BaseMultiModalProcessor[KimiAudioProcessingInfo]): """vLLM multi-modal processor wrapper for Kimi-Audio.""" + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> ProcessorInputs: + builder = self.dummy_inputs + dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) + + num_audios = mm_counts.get("audio", 0) + dummy_tokens = ( + [198] + if num_audios == 0 + else [ + KimiAudioProcessor.KIMIA_MEDIA_BEGIN, + KimiAudioProcessor.KIMIA_TEXT_BLANK, + KimiAudioProcessor.KIMIA_MEDIA_END, + ] + * num_audios + ) + + return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) + def _get_hf_mm_inputs( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 208d5d4bc1ab..e3fdab76ff04 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -438,31 +438,38 @@ def get_dummy_mm_data( ) } - def get_dummy_processor_inputs( + +class MossTranscribeDiarizeMultiModalProcessor( + BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo] +): + def get_dummy_inputs( self, seq_len: int, mm_counts: Mapping[str, int], mm_options: MultiModalDummyOptions, ) -> ProcessorInputs: - dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + builder = self.dummy_inputs + tokenizer = self.info.get_tokenizer() + + dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) + num_audios = mm_counts.get("audio", 0) - tokenizer = self.info.get_tokenizer() - prompt = cached_encode( - tokenizer, - AUDIO_PLACEHOLDER * num_audios, - add_special_tokens=False, - ) or cached_encode( - tokenizer, - "\n", - add_special_tokens=False, - ) - return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items) + if num_audios > 0: + prompt = cached_encode( + tokenizer, + AUDIO_PLACEHOLDER * num_audios, + add_special_tokens=False, + ) + else: + prompt = cached_encode( + tokenizer, + "\n", + add_special_tokens=False, + ) + return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items) -class MossTranscribeDiarizeMultiModalProcessor( - BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo] -): def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index f5c167aec3da..91ac38e11087 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -221,18 +221,22 @@ def get_dummy_mm_data( ) } - def get_dummy_processor_inputs( + +class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]): + def get_dummy_inputs( self, seq_len: int, mm_counts: Mapping[str, int], mm_options: MultiModalDummyOptions, + # For test_common.py only mm_data: MultiModalDataDict | None = None, ) -> ProcessorInputs: + builder = self.dummy_inputs tokenizer = self.info.get_tokenizer() - dummy_text = self.get_dummy_text(mm_counts) + dummy_text = builder.get_dummy_text(mm_counts) dummy_mm_data = ( - self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) if mm_data is None else mm_data ) @@ -256,8 +260,6 @@ def get_dummy_processor_inputs( return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) - -class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]): # The tokens are already inserted by the chat template, # so we just double check that they exist def _maybe_apply_prompt_updates( diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index 0b8d889bbc69..21b291344d90 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -149,19 +149,23 @@ def get_dummy_mm_data( ) } - def get_dummy_processor_inputs( + +class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]): + def get_dummy_inputs( self, seq_len: int, mm_counts: Mapping[str, int], mm_options: MultiModalDummyOptions, + # For test_common.py only mm_data: MultiModalDataDict | None = None, ) -> ProcessorInputs: + builder = self.dummy_inputs tokenizer = self.info.get_tokenizer() feature_extractor = self.info.get_feature_extractor() - dummy_text = self.get_dummy_text(mm_counts) + dummy_text = builder.get_dummy_text(mm_counts) dummy_mm_data = ( - self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) if mm_data is None else mm_data ) @@ -197,8 +201,6 @@ def get_dummy_processor_inputs( return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) - -class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]): # The tokens are already inserted by the chat template, # so we just double check that they exist def _maybe_apply_prompt_updates( diff --git a/vllm/multimodal/processing/dummy_inputs.py b/vllm/multimodal/processing/dummy_inputs.py index 53e9f4f5b69d..dd49b6c4fabc 100644 --- a/vllm/multimodal/processing/dummy_inputs.py +++ b/vllm/multimodal/processing/dummy_inputs.py @@ -18,7 +18,6 @@ from vllm.logger import init_logger from .context import BaseProcessingInfo -from .inputs import ProcessorInputs _I = TypeVar("_I", bound=BaseProcessingInfo) @@ -61,42 +60,6 @@ def get_dummy_mm_data( """ raise NotImplementedError - def get_dummy_processor_inputs( - self, - seq_len: int, - mm_counts: Mapping[str, int], - mm_options: MultiModalDummyOptions, - ) -> ProcessorInputs: - """Build the input which, after processing, results in - the maximum possible number of placeholder tokens. - - Args: - seq_len: Sequence length - mm_counts: Count of items per modality - mm_options: Configurable options per modality (optional) - - """ - dummy_text = self.get_dummy_text(mm_counts) - dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) - dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) - - tokenizer = self.info.ctx.tokenizer - dummy_prompt: list[int] - if tokenizer is None: - # Tokenizer-less models (e.g. `skip_tokenizer_init=True`) only - # accept embeddings and have an empty dummy text, so there are no - # prompt tokens. - dummy_prompt = [] - else: - from .processor import cached_encode - - dummy_prompt = cached_encode(tokenizer, dummy_text, truncation=False) - - return ProcessorInputs( - prompt=dummy_prompt, - mm_data_items=dummy_mm_items, - ) - def _get_dummy_audios( self, *, diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index b8e51b205301..eb53023f0f35 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -39,7 +39,7 @@ ) from ..parse import MultiModalDataItems, MultiModalUUIDItems from .context import BaseProcessingInfo, TimingContext -from .dummy_inputs import BaseDummyInputsBuilder +from .dummy_inputs import BaseDummyInputsBuilder, MultiModalDummyOptions from .inputs import ProcessorInputs if TYPE_CHECKING: @@ -1055,6 +1055,41 @@ def __init__( self.data_parser = self.info.get_data_parser() + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> ProcessorInputs: + """Build the input which, after processing, results in + the maximum possible number of placeholder tokens. + + Args: + seq_len: Sequence length + mm_counts: Count of items per modality + mm_options: Configurable options per modality (optional) + + """ + builder = self.dummy_inputs + dummy_text = builder.get_dummy_text(mm_counts) + dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) + + tokenizer = self.info.ctx.tokenizer + dummy_prompt: list[int] + if tokenizer is None: + # Tokenizer-less models (e.g. `skip_tokenizer_init=True`) only + # accept embeddings and have an empty dummy text, so there are no + # prompt tokens. + dummy_prompt = [] + else: + dummy_prompt = cached_encode(tokenizer, dummy_text, truncation=False) + + return ProcessorInputs( + prompt=dummy_prompt, + mm_data_items=dummy_mm_items, + ) + def get_dummy_mm_inputs( self, mm_counts: Mapping[str, int], @@ -1070,7 +1105,7 @@ def get_dummy_mm_inputs( mm_config = model_config.get_multimodal_config() - processor_inputs = self.dummy_inputs.get_dummy_processor_inputs( + processor_inputs = self.get_dummy_inputs( seq_len=seq_len, mm_counts=mm_counts, mm_options=mm_config.limit_per_prompt, From 4f1451679088e5832bce0965a254295c854aaa09 Mon Sep 17 00:00:00 2001 From: djramic Date: Mon, 21 Sep 2026 18:39:31 +0200 Subject: [PATCH 595/774] [ROCm][CI] Use ROCm backend for DeepSeek V4.1 ViT test (#57931) Signed-off-by: Djordje Ramic Co-authored-by: Andreas Karatzas --- tests/models/multimodal/generation/test_vit_cudagraph.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 9bda7ec85db9..544623319f52 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -236,7 +236,11 @@ def minicpmv_chat_template(content: str) -> str: }, vllm_runner_kwargs={ "load_format": "dummy", - "attention_backend": "FLASHMLA_SPARSE_DSV41", + "attention_backend": ( + "ROCM_FLASHMLA_SPARSE_DSV4" + if current_platform.is_rocm() + else "FLASHMLA_SPARSE_DSV41" + ), "hf_overrides": partial( dummy_hf_overrides, model_arch="DeepseekV41ForCausalLM", From 986e217158a9ed54eca8307c6bb9abac7c5b2278 Mon Sep 17 00:00:00 2001 From: Karen Chung Date: Mon, 21 Sep 2026 09:56:11 -0700 Subject: [PATCH 596/774] [Feature] Add first-class KV hints request envelope for programmatic KV management (#53423) Signed-off-by: Karen Chung --- rust/proto/inference.proto | 15 +++++ rust/src/chat/src/lib.rs | 1 + .../src/protocol/kv_hints.rs | 21 +++++++ .../engine-core-client/src/protocol/mod.rs | 1 + .../src/protocol/request.rs | 17 ++++- .../engine-core-client/src/tests/client.rs | 11 ++++ .../src/tests/python_compat.py | 26 ++++++++ .../src/llm/examples/external_engine_smoke.rs | 1 + rust/src/llm/src/request.rs | 20 ++++++ rust/src/llm/tests/generate.rs | 1 + rust/src/server/src/grpc/convert.rs | 62 ++++++++++++++++++- .../src/routes/inference/generate/convert.rs | 1 + .../src/routes/openai/completions/convert.rs | 1 + rust/src/server/src/routes/tokenize/types.rs | 1 + rust/src/text/src/lower.rs | 31 ++++++++++ rust/src/text/src/request.rs | 5 ++ .../offloading_connector/test_scheduler.py | 18 ++++++ tests/v1/kv_offload/tiering/test_kvcr_tier.py | 12 ++-- tests/v1/test_request.py | 34 ++++++++++ .../kv_connector/v1/offloading/scheduler.py | 1 + vllm/engine/protocol.py | 2 + vllm/v1/engine/__init__.py | 2 + vllm/v1/engine/async_llm.py | 9 +++ vllm/v1/engine/input_processor.py | 3 + vllm/v1/engine/llm_engine.py | 3 + vllm/v1/kv_hints/__init__.py | 6 ++ vllm/v1/kv_hints/protocol.py | 26 ++++++++ vllm/v1/kv_offload/base.py | 6 +- vllm/v1/kv_offload/tiering/kvcr/manager.py | 17 +++-- vllm/v1/request.py | 4 ++ 30 files changed, 340 insertions(+), 18 deletions(-) create mode 100644 rust/src/engine-core-client/src/protocol/kv_hints.rs create mode 100644 vllm/v1/kv_hints/__init__.py create mode 100644 vllm/v1/kv_hints/protocol.py diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 4eaf6f92f4b4..2bc35bcb64a2 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -56,6 +56,21 @@ message GenerateRequest { // Apply the engine's configured watermark. Defaults to true when omitted. optional bool watermarking = 16; + + KVHintsEnvelope kv_hints = 17; +} + +message KVHintsEnvelope { + string protocol_version = 1; + string message_id = 2; + repeated KVHintAction actions = 3; +} + +message KVHintAction { + string action_id = 1; + string action_type = 2; + string action_version = 3; + google.protobuf.Struct payload = 4; } message RandomSampling { diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 0273396ebd85..d1a34455d2a0 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -194,6 +194,7 @@ impl ChatRequestProcessor { add_special_tokens: request.add_special_tokens, data_parallel_rank: request.data_parallel_rank, session_id: request.session_id, + kv_hints: None, reasoning_parser_kwargs, lora_request: request.lora_request, arrival_time: Some(arrival_time), diff --git a/rust/src/engine-core-client/src/protocol/kv_hints.rs b/rust/src/engine-core-client/src/protocol/kv_hints.rs new file mode 100644 index 000000000000..418ca2fcedb6 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/kv_hints.rs @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct KvHintAction { + pub action_id: String, + pub action_type: String, + pub action_version: String, + pub payload: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct KvHintsEnvelope { + pub protocol_version: String, + pub message_id: String, + pub actions: Vec, +} diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 742fed59bdeb..0586dde135ce 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -24,6 +24,7 @@ fn is_false(v: &bool) -> bool { pub mod dtype; pub mod handshake; +pub mod kv_hints; pub mod logprobs; pub mod lora; pub mod multimodal; diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs index 973ada71c6c1..58c6efb641de 100644 --- a/rust/src/engine-core-client/src/protocol/request.rs +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use crate::protocol::kv_hints::KvHintsEnvelope; use crate::protocol::multimodal::MmFeatures; use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::{OpaqueValue, lora}; @@ -127,6 +128,8 @@ pub struct EngineCoreRequest { /// Stable session identity shared by related requests. #[serde(default)] pub session_id: Option, + #[serde(default)] + pub kv_hints: Option, } impl EngineCoreRequest { @@ -159,6 +162,7 @@ mod tests { use super::*; use crate::protocol::dtype::TensorDtype; + use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use crate::protocol::multimodal::{ MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, MmModality, PlaceholderRange, @@ -181,6 +185,16 @@ mod tests { arrival_time: 1234.5, client_index: 7, session_id: Some("session-1".to_string()), + kv_hints: Some(KvHintsEnvelope { + protocol_version: "0.1".to_string(), + message_id: "msg-1".to_string(), + actions: vec![KvHintAction { + action_id: "action-1".to_string(), + action_type: "example.action".to_string(), + action_version: "1.0".to_string(), + payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]), + }], + }), ..EngineCoreRequest::default() }; @@ -191,13 +205,14 @@ mod tests { other => panic!("expected array, got {other:?}"), }; - assert_eq!(array.len(), 21); + assert_eq!(array.len(), 22); assert_eq!(array[0], Value::from("req-1")); assert_eq!(array[2], Value::Nil); assert_eq!(array[4], Value::Nil); assert_eq!(array[10], Value::Nil); assert_eq!(array[11], Value::from(7)); assert_eq!(array[20], Value::from("session-1")); + assert!(matches!(&array[21], Value::Map(_))); } #[test] diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index eb8e713fb402..bd3847c483d8 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -20,6 +20,7 @@ use zeromq::util::PeerIdentity; use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage}; use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; +use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use crate::protocol::logprobs::MaybeWireLogprobs; use crate::protocol::multimodal::{ MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmModality, MmSlice, @@ -165,6 +166,16 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest { }), arrival_time: 42.5, session_id: Some("session-1".to_string()), + kv_hints: Some(KvHintsEnvelope { + protocol_version: "0.1".to_string(), + message_id: "msg-1".to_string(), + actions: vec![KvHintAction { + action_id: "action-1".to_string(), + action_type: "example.action".to_string(), + action_version: "1.0".to_string(), + payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]), + }], + }), ..EngineCoreRequest::default() } } diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index a19dc9fdaa45..76978328bb1b 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -52,6 +52,19 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True): output_kind: RequestOutputKind = RequestOutputKind.DELTA +class KvHintAction(msgspec.Struct, frozen=True): + action_id: str + action_type: str + action_version: str + payload: dict[str, object] + + +class KvHintsEnvelope(msgspec.Struct, frozen=True): + protocol_version: str + message_id: str + actions: list[KvHintAction] + + class EngineCoreRequest( msgspec.Struct, array_like=True, @@ -78,6 +91,7 @@ class EngineCoreRequest( reasoning_parser_kwargs: dict[str, object] | None = None abort_immediately: bool = False session_id: str | None = None + kv_hints: KvHintsEnvelope | None = None class EngineCoreOutput( @@ -153,6 +167,18 @@ class EngineCoreOutputs( arrival_time=42.5, client_index=0, session_id="session-1", + kv_hints=KvHintsEnvelope( + protocol_version="0.1", + message_id="msg-1", + actions=[ + KvHintAction( + action_id="action-1", + action_type="example.action", + action_version="1.0", + payload={"key": "value"}, + ) + ], + ), ) # All defaults -> empty map. Regression guard for the sparse-map decode. diff --git a/rust/src/llm/examples/external_engine_smoke.rs b/rust/src/llm/examples/external_engine_smoke.rs index 9d318285ae81..0a580a222156 100644 --- a/rust/src/llm/examples/external_engine_smoke.rs +++ b/rust/src/llm/examples/external_engine_smoke.rs @@ -60,6 +60,7 @@ fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest { priority: 0, data_parallel_rank: None, session_id: None, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: None, } diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs index c647f54d8df4..1f9d38ae0e13 100644 --- a/rust/src/llm/src/request.rs +++ b/rust/src/llm/src/request.rs @@ -4,6 +4,7 @@ use std::collections::BTreeMap; use uuid::Uuid; +use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs}; @@ -48,6 +49,8 @@ pub struct GenerateRequest { pub data_parallel_rank: Option, /// Stable session identity shared by related requests. pub session_id: Option, + /// Optional orchestrator-originated KV hints. + pub kv_hints: Option, /// Optional reasoning-parser kwargs forwarded to engine-side structured /// output logic. pub reasoning_parser_kwargs: Option, @@ -79,6 +82,7 @@ impl GenerateRequest { priority, data_parallel_rank, session_id, + kv_hints, reasoning_parser_kwargs, lora_request, } = self; @@ -109,6 +113,7 @@ impl GenerateRequest { trace_headers, resumable: false, session_id, + kv_hints, external_req_id: Some(external_request_id), // Rust parser doesn't expose this information, leave it unset and let the // reasoning logic in engine-sided structured output manager handle it. @@ -134,6 +139,7 @@ impl PreparedGenerateRequest { mod tests { use std::collections::BTreeMap; + use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; @@ -155,6 +161,16 @@ mod tests { priority: 3, data_parallel_rank: Some(2), session_id: Some("session-1".to_string()), + kv_hints: Some(KvHintsEnvelope { + protocol_version: "0.1".to_string(), + message_id: "msg-1".to_string(), + actions: vec![KvHintAction { + action_id: "action-1".to_string(), + action_type: "example.action".to_string(), + action_version: "1.0".to_string(), + payload: BTreeMap::new(), + }], + }), reasoning_parser_kwargs: Some(ReasoningParserKwargs { chat_template_kwargs: [( "chat_template_kwargs".to_string(), @@ -183,6 +199,10 @@ mod tests { assert_eq!(request.cache_salt.as_deref(), Some("salt")); assert_eq!(request.data_parallel_rank, Some(2)); assert_eq!(request.session_id.as_deref(), Some("session-1")); + assert_eq!( + request.kv_hints.as_ref().map(|hints| hints.message_id.as_str()), + Some("msg-1") + ); assert_eq!( request.trace_headers, Some(BTreeMap::from([( diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index c341737a3c68..c9849cd57c19 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -167,6 +167,7 @@ fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest priority: 0, data_parallel_rank: None, session_id: None, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: None, } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 0f28e147ef07..538167b41f0d 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -7,6 +7,7 @@ use thiserror_ext::AsReport as _; use tonic::Status; use uuid::Uuid; +use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_text::{ @@ -16,6 +17,32 @@ use vllm_text::{ use super::pb; +fn kv_hints_from_proto(hints: pb::KvHintsEnvelope) -> KvHintsEnvelope { + KvHintsEnvelope { + protocol_version: hints.protocol_version, + message_id: hints.message_id, + actions: hints + .actions + .into_iter() + .map(|action| KvHintAction { + action_id: action.action_id, + action_type: action.action_type, + action_version: action.action_version, + payload: action + .payload + .map(|payload| { + payload + .fields + .into_iter() + .map(|(key, value)| (key, proto_value_to_json(&value))) + .collect() + }) + .unwrap_or_default(), + }) + .collect(), + } +} + // ======================================================================================== // Request conversion // ======================================================================================== @@ -56,6 +83,7 @@ pub fn to_text_request( req.request_id }; let session_id = req.session_id.filter(|s| !s.is_empty()); + let kv_hints = req.kv_hints.map(kv_hints_from_proto); let sampling = req.sampling.as_ref(); let decoding = req.decoding.as_ref(); @@ -108,6 +136,7 @@ pub fn to_text_request( add_special_tokens: true, data_parallel_rank: None, session_id, + kv_hints, reasoning_parser_kwargs: None, lora_request: None, arrival_time: None, @@ -537,7 +566,9 @@ mod tests { use vllm_tokenizer::test_utils::TestTokenizer; use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; - use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request}; + use super::{ + ResponseOpts, json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request, + }; fn base_request() -> pb::GenerateRequest { pb::GenerateRequest { @@ -578,6 +609,35 @@ mod tests { assert_eq!(text.sampling_params.temperature, Some(0.7)); } + #[test] + fn kv_hints_propagate_from_top_level_request_field() { + let req = pb::GenerateRequest { + kv_hints: Some(pb::KvHintsEnvelope { + protocol_version: "0.1".to_string(), + message_id: "msg-1".to_string(), + actions: vec![pb::KvHintAction { + action_id: "action-1".to_string(), + action_type: "example.action".to_string(), + action_version: "1.0".to_string(), + payload: json_to_proto_struct(&serde_json::json!({"source": "g2"})), + }], + }), + ..base_request() + }; + + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok"); + let hints = text.kv_hints.expect("kv hints"); + assert_eq!(hints.protocol_version, "0.1"); + assert_eq!(hints.message_id, "msg-1"); + assert_eq!(hints.actions[0].action_id, "action-1"); + assert_eq!(hints.actions[0].action_type, "example.action"); + assert_eq!(hints.actions[0].action_version, "1.0"); + assert_eq!( + hints.actions[0].payload.get("source"), + Some(&serde_json::json!("g2")) + ); + } + #[test] fn unset_temperature_defaults_to_greedy() { let text = to_text_request(base_request(), false, &["test-model".to_string()]) diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 8093a34ca29d..c3ec49dd23f8 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -82,6 +82,7 @@ pub(super) fn prepare_generate_request( add_special_tokens: false, data_parallel_rank: ctx.data_parallel_rank, session_id: ctx.session_id, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), arrival_time: None, diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index a4e8b6d3e423..ddace3970e40 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -155,6 +155,7 @@ pub(super) fn prepare_completion_request( add_special_tokens: request.add_special_tokens, data_parallel_rank: ctx.data_parallel_rank, session_id, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), arrival_time: None, diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 35a027c1c0f7..ba27211ebc6d 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -54,6 +54,7 @@ impl TokenizeCompletionRequest { add_special_tokens: self.add_special_tokens, data_parallel_rank: None, session_id: None, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: None, arrival_time: None, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 466b9d5b31ad..d8c639977322 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -60,6 +60,7 @@ pub fn lower_text_request( priority: request.priority, data_parallel_rank: request.data_parallel_rank, session_id: request.session_id.clone(), + kv_hints: request.kv_hints.clone(), reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(), lora_request: request.lora_request.clone(), arrival_time: request.arrival_time, @@ -320,6 +321,7 @@ mod tests { use std::collections::{BTreeSet, HashMap}; use serial_test::file_serial; + use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::multimodal::{ MmFeatureSpec, MmModality, PlaceholderRange, }; @@ -1301,6 +1303,35 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_passes_kv_hints_through() { + let hints = KvHintsEnvelope { + protocol_version: "0.1".to_string(), + message_id: "msg-1".to_string(), + actions: vec![KvHintAction { + action_id: "action-1".to_string(), + action_type: "example.action".to_string(), + action_version: "1.0".to_string(), + payload: Default::default(), + }], + }; + let request = TextRequest { + kv_hints: Some(hints.clone()), + ..sample_request() + }; + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.kv_hints, Some(hints)); + } + #[test] fn lower_text_request_passes_arrival_time_through() { let request = TextRequest { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index f0dd2fb1f88c..7890977ae4f5 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; use serde_json::Value; +use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; @@ -211,6 +212,9 @@ pub struct TextRequest { /// Stable session identity shared by related requests. #[serde(default)] pub session_id: Option, + /// Optional orchestrator-originated KV hints. + #[serde(default)] + pub kv_hints: Option, /// Optional reasoning-parser kwargs forwarded to engine-side structured /// output logic. #[serde(default)] @@ -242,6 +246,7 @@ impl TextRequest { add_special_tokens: false, data_parallel_rank: None, session_id: None, + kv_hints: None, reasoning_parser_kwargs: None, lora_request: None, arrival_time: None, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 035fb6b0054f..a69129561620 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -33,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, + _create_req_context, get_sliding_window_size_in_chunks, ) from vllm.v1.core.block_pool import BlockPool @@ -60,6 +61,7 @@ SlidingWindowSpec, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, LookupResult, @@ -214,6 +216,7 @@ def _make_partial_tail_request( request = MagicMock() request.request_id = "req" request.kv_transfer_params = kv_transfer_params + request.kv_hints = None request.num_prompt_tokens = 30 request.num_tokens = 30 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)] @@ -225,6 +228,21 @@ def _make_partial_tail_request( return request +def test_create_req_context_preserves_kv_hints(): + kv_hints = KvHintsEnvelope( + protocol_version="0.1", + message_id="message-1", + actions=[], + ) + request = MagicMock( + request_id="req", + kv_transfer_params=None, + kv_hints=kv_hints, + ) + + assert _create_req_context(request).kv_hints is kv_hints + + def _reduce_kv_connector_stats(runner): reduced: dict[str, int | float] = {} for payload in runner.kv_connector_stats: diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index 38a2ae044ce2..5cbc85230ffb 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -5,12 +5,13 @@ from types import SimpleNamespace from typing import Any +import msgspec import numpy as np import pytest pytest.importorskip("kvcr") -from kvcr import ROUTER_HINT_KEY, KVCRBindings +from kvcr import KVCRBindings from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy from kvcr.types import ( @@ -27,6 +28,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.kv_offload.base import ( LookupResult, Medium, @@ -271,10 +273,10 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch): """Check hint forwarding, key conversion, load descriptors, and cleanup.""" kvcr = RecordingKVCR() tier = _make_tier(monkeypatch, kvcr) - router_hint = {"opaque": True} + kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[]) ctx = ReqContext( req_id="req", - kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()}, + kv_hints=kv_hint, ) key = make_offload_key((123).to_bytes(8, "big"), 0) same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) @@ -282,7 +284,9 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch): tier.on_new_request(ctx) - assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})] + assert kvcr.submit_hint_calls == [ + ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)}) + ] bindings = kvcr.constructor_bindings assert bindings is not None diff --git a/tests/v1/test_request.py b/tests/v1/test_request.py index be417b9b2ff5..6c45e1a6d056 100644 --- a/tests/v1/test_request.py +++ b/tests/v1/test_request.py @@ -3,6 +3,7 @@ from vllm import SamplingParams from vllm.v1.engine import EngineCoreRequest +from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope from vllm.v1.request import Request, RequestStatus @@ -24,6 +25,7 @@ def test_request_status_fmt_str(): def test_request_copies_session_id_from_engine_core_request(): + """Test that Request preserves session ID from EngineCoreRequest.""" engine_request = EngineCoreRequest( request_id="request-1", prompt_token_ids=[1, 2, 3], @@ -40,3 +42,35 @@ def test_request_copies_session_id_from_engine_core_request(): request = Request.from_engine_core_request(engine_request, block_hasher=None) assert request.session_id == "session-1" + + +def test_request_copies_kv_hints_from_engine_core_request(): + """Test that Request preserves KV hints from EngineCoreRequest.""" + kv_hints = KvHintsEnvelope( + protocol_version="0.1", + message_id="msg-1", + actions=[ + KvHintAction( + action_id="action-1", + action_type="example.action", + action_version="1.0", + payload={}, + ) + ], + ) + engine_request = EngineCoreRequest( + request_id="request-1", + prompt_token_ids=[1, 2, 3], + mm_features=None, + sampling_params=SamplingParams(max_tokens=1), + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + kv_hints=kv_hints, + ) + + request = Request.from_engine_core_request(engine_request, block_hasher=None) + + assert request.kv_hints == kv_hints diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index d6c210a6f9bf..47aa8654feb8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -553,6 +553,7 @@ def _create_req_context(req: Request) -> ReqContext: return ReqContext( req_id=req.request_id, kv_transfer_params=params, + kv_hints=req.kv_hints, load_tier_filter=load_filter, ) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 1e6011a2300f..74849a85cef9 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -22,6 +22,7 @@ from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult +from vllm.v1.kv_hints import KvHintsEnvelope if TYPE_CHECKING: from vllm.v1.engine import PauseMode @@ -102,6 +103,7 @@ def generate( priority: int = 0, data_parallel_rank: int | None = None, session_id: str | None = None, + kv_hints: KvHintsEnvelope | None = None, reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, ) -> AsyncGenerator[RequestOutput, None]: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 7c424ea0df0c..e5bbb96a9cf7 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -16,6 +16,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.metrics.stats import ( PrefillStats, RequestSpecDecodeMetrics, @@ -157,6 +158,7 @@ class EngineCoreRequest( abort_immediately: bool = False session_id: str | None = None + kv_hints: KvHintsEnvelope | None = None @property def params(self) -> SamplingParams | PoolingParams: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index ee913b052f62..7d71761990b0 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -53,6 +53,7 @@ from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.executor import Executor from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.metrics.loggers import ( StatLoggerFactory, StatLoggerManager, @@ -386,6 +387,7 @@ async def add_request( prompt_text: str | None = None, reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, + kv_hints: KvHintsEnvelope | None = None, ) -> RequestOutputCollector: """Add new request to the AsyncLLM.""" if self.errored: @@ -420,6 +422,7 @@ async def add_request( priority, data_parallel_rank, session_id, + kv_hints, ) # Convert Input --> Request. @@ -453,6 +456,7 @@ async def add_request( priority=priority, data_parallel_rank=data_parallel_rank, session_id=session_id, + kv_hints=kv_hints, ) else: # Raw prompts require tokenization and possibly multimodal @@ -469,6 +473,7 @@ async def add_request( priority=priority, data_parallel_rank=data_parallel_rank, session_id=session_id, + kv_hints=kv_hints, ) prompt_text, _, _ = extract_prompt_components(self.model_config, prompt) @@ -557,6 +562,7 @@ async def _add_streaming_input_request( priority: int = 0, data_parallel_rank: int | None = None, session_id: str | None = None, + kv_hints: KvHintsEnvelope | None = None, ) -> RequestOutputCollector: self._validate_streaming_input_sampling_params(sampling_params) @@ -569,6 +575,7 @@ async def _add_streaming_input_request( priority=priority, data_parallel_rank=data_parallel_rank, session_id=session_id, + kv_hints=kv_hints, ) if not sampling_params.skip_clone: @@ -670,6 +677,7 @@ async def generate( priority: int = 0, data_parallel_rank: int | None = None, session_id: str | None = None, + kv_hints: KvHintsEnvelope | None = None, reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, ) -> AsyncGenerator[RequestOutput, None]: @@ -710,6 +718,7 @@ async def generate( priority=priority, data_parallel_rank=data_parallel_rank, session_id=session_id, + kv_hints=kv_hints, prompt_text=prompt_text, reasoning_ended=reasoning_ended, reasoning_parser_kwargs=reasoning_parser_kwargs, diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index d0ccb1dacf4c..582e1abac685 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -32,6 +32,7 @@ from vllm.utils.async_utils import make_async from vllm.utils.jsontree import json_iter_leaves from vllm.v1.engine import EngineCoreRequest +from vllm.v1.kv_hints import KvHintsEnvelope logger = init_logger(__name__) @@ -318,6 +319,7 @@ def process_inputs( data_parallel_rank: int | None = None, resumable: bool = False, session_id: str | None = None, + kv_hints: KvHintsEnvelope | None = None, ) -> EngineCoreRequest: self._validate_params(params, supported_tasks) self._validate_lora(lora_request) @@ -461,6 +463,7 @@ def process_inputs( trace_headers=trace_headers, resumable=resumable, session_id=session_id, + kv_hints=kv_hints, ) def _validate_prompt_len( diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index a2d97b5544fb..c925ba90e8e0 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -34,6 +34,7 @@ from vllm.v1.engine.output_processor import OutputProcessor from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.executor import Executor +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager from vllm.v1.metrics.reader import Metric, get_metrics_snapshot from vllm.v1.metrics.stats import IterationStats @@ -230,6 +231,7 @@ def add_request( priority: int = 0, session_id: str | None = None, prompt_text: str | None = None, + kv_hints: KvHintsEnvelope | None = None, ) -> str: # Validate the request_id type. if not isinstance(request_id, str): @@ -263,6 +265,7 @@ def add_request( trace_headers=trace_headers, priority=priority, session_id=session_id, + kv_hints=kv_hints, ) prompt_text, _, _ = extract_prompt_components(self.model_config, prompt) diff --git a/vllm/v1/kv_hints/__init__.py b/vllm/v1/kv_hints/__init__.py new file mode 100644 index 000000000000..d0807572c7a8 --- /dev/null +++ b/vllm/v1/kv_hints/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.v1.kv_hints.protocol import KvHintAction, KvHintsEnvelope + +__all__ = ["KvHintAction", "KvHintsEnvelope"] diff --git a/vllm/v1/kv_hints/protocol.py b/vllm/v1/kv_hints/protocol.py new file mode 100644 index 000000000000..a3f93ab0bf23 --- /dev/null +++ b/vllm/v1/kv_hints/protocol.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import msgspec + + +class KvHintAction(msgspec.Struct, frozen=True): # type: ignore[call-arg] + """A versioned orchestrator-provided KV hint action.""" + + action_id: str + action_type: str + action_version: str + # `payload` must contain JSON-compatible values. Action-specific schemas are + # validated by the component implementing the action. + payload: dict[str, Any] + + +class KvHintsEnvelope(msgspec.Struct, frozen=True): # type: ignore[call-arg] + """Orchestrator-provided KV hints optionally attached to one inference request. + A KV hint envelope is a collection of KV hint action(s).""" + + protocol_version: str + message_id: str + actions: list[KvHintAction] diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 1144477628d2..f8e92abed604 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -16,6 +16,7 @@ OffloadingConnectorStats, ) +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.kv_offload.config import OffloadingConfig # `OffloadKey` identifies an offloaded block. It combines a block hash with @@ -89,10 +90,11 @@ def allows(self, medium: Medium | None, locality: Locality | None) -> bool: class ReqContext: req_id: str kv_transfer_params: dict[str, Any] | None = None + kv_hints: KvHintsEnvelope | None = None load_tier_filter: TierFilter = TierFilter.ALL # Per-request scratch space keyed by value type, so a tier can parse - # kv_transfer_params once (in on_new_request) and read the result back - # on later calls for the same request. + # kv_transfer_params and kv_hints once (in on_new_request) and read the + # result back on later calls for the same request. _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False) # End-token position for each key in this request. The scheduler records # these positions so managers can recover prefix order even when store diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 00186279a4bc..18da2a6c2764 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -7,15 +7,15 @@ import socket import time import uuid -from collections.abc import Collection, Iterable, Mapping +from collections.abc import Collection, Iterable from pathlib import Path from typing import TYPE_CHECKING, Any, cast +import msgspec from kvcr import ( DURATION_METRIC, KVCR, ROUTER_HINT_CAPABILITIES, - ROUTER_HINT_KEY, STATE_METRIC, TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC, @@ -575,14 +575,11 @@ def take_events(self) -> Iterable[OffloadingEvent]: @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - params = getattr(req_context, "kv_transfer_params", None) - if isinstance(params, Mapping): - hint = params.get(ROUTER_HINT_KEY) - if hint is not None: - self._kvcr.submit_hint( - request_id=req_context.req_id, - hints=hint, - ) + if (hint := req_context.kv_hints) is not None: + self._kvcr.submit_hint( + request_id=req_context.req_id, + hints=msgspec.to_builtins(hint), + ) return RequestOffloadingContext() @override diff --git a/vllm/v1/request.py b/vllm/v1/request.py index dcf679093b66..4b1b9ecc2eed 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -20,6 +20,7 @@ EngineCoreRequest, FinishReason, ) +from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics from vllm.v1.structured_output.request import StructuredOutputRequest from vllm.v1.utils import ConstantList @@ -78,6 +79,7 @@ def __init__( reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, abort_immediately: bool = False, + kv_hints: KvHintsEnvelope | None = None, ) -> None: self.request_id = request_id self.client_index = client_index @@ -193,6 +195,7 @@ def __init__( # trace_headers self.trace_headers = trace_headers self.session_id = session_id + self.kv_hints = kv_hints # True if this request is scheduled as a non-final prefill chunk. self.is_prefill_chunk = False @@ -260,6 +263,7 @@ def from_engine_core_request( block_hasher=block_hasher, resumable=request.resumable, session_id=request.session_id, + kv_hints=request.kv_hints, reasoning_ended=request.reasoning_ended, reasoning_parser_kwargs=request.reasoning_parser_kwargs, abort_immediately=request.abort_immediately, From fba33e39cb74c7c31cc648ae530a85e0c6636f92 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 21 Sep 2026 19:15:28 +0200 Subject: [PATCH 597/774] [Docs] Add an Engram feature page explaining Engram usage in vLLM (#57910) Signed-off-by: NickLucche --- docs/features/engram.md | 148 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 docs/features/engram.md diff --git a/docs/features/engram.md b/docs/features/engram.md new file mode 100644 index 000000000000..73dc05e9f144 --- /dev/null +++ b/docs/features/engram.md @@ -0,0 +1,148 @@ +# Engram: conditional memory via n-gram lookups + +Engram is a *conditional memory* module attached to specific layers of a +transformer backbone. At each position, the preceding few tokens are hashed +into large static embedding tables, and the retrieved rows are fused into the +hidden state, gated by how well they match it. Where Mixture-of-Experts scales +capacity through conditional computation, Engram scales it through conditional +memory: parameters that are never computed over, only looked up, so capacity +grows at O(1) per-token cost. The mechanism and its design rationale are +described in [Conditional Memory via Scalable Lookup: A New Axis of Sparsity +for Large Language Models](https://arxiv.org/abs/2601.07372). + +vLLM enables Engram automatically whenever the checkpoint declares Engram +layers, so basic serving needs no extra flags. This page explains what Engram +does and how to control its memory footprint and multi-GPU topology. + +## Supported models + +Engram layers are declared by the checkpoint, not chosen at serve time: + +| Architecture | HF config field | Notes | +| --- | --- | --- | +| `DeepseekV41ForCausalLM` | `engram_layer_ids` | e.g. `deepseek-ai/DeepSeek-V4.1-Flash` | +| `Qwen4ExpForCausalLM` / `Qwen4ExpForConditionalGeneration` | `ple_layer_ids` | same mechanism, called PLE | + +Passing `--engram-config` for a model without n-gram embedding layers fails at +startup, as does running on a non-CUDA-alike platform. + +## How it works + +Engram processes each position in two phases: *retrieval* and *fusion*. + +**Retrieval** maps the local token context to static table rows: + +1. **Tokenizer compression.** At startup, vLLM builds a many-to-one map from + the tokenizer that collapses tokens which normalize alike (NFKC, accent + stripping, lowercasing, whitespace collapsing), so `" The"`, `"the"` and + `"THE"` hash identically. This concentrates semantic density into a much + smaller id space; its size must match `engram_compressed_vocab_size` from + the checkpoint, or startup fails — every hash multiplier derives from it. +2. **Multi-head hashing.** For each n-gram order `n` (up to + `engram_max_ngram_size - 1`), several hash heads each apply an independent + multiplicative-XOR hash over the compressed suffix n-gram. Every (n-gram + order, head) pair indexes its own prime-sized bucket range in the layer's + table — disjoint across pairs, so collisions cannot collide twice — and + the retrieved rows are concatenated into one memory vector per position. + +**Fusion** turns the retrieved rows into a context-aware contribution to the +hidden state. The rows are context-independent priors and may be noisy (hash +collisions, polysemy), so they are gated: a `wkv` projection produces one key +per hyper-connection branch plus a shared value, and each branch computes a +sigmoid gate over the normalized dot product between its hidden state — +which has already aggregated global context through preceding attention — +and its key. If the retrieved memory contradicts the current context, the +gate tends toward zero and suppresses it. The gated value is added +residually into the branch's stream (`hidden + gate * value`); positions that +take no part in an n-gram (e.g. image spans) are masked so they pass through +untouched. + +Engram sits only on the layers listed in the checkpoint. Placement is a +deliberate hardware-algorithm trade-off: early enough that the backbone is +relieved from reconstructing static patterns in its first layers, deep enough +that vLLM can hide table lookups behind the compute of the preceding layers +(see below). + +In vLLM, hashing, lookup, and gating run as fused Triton kernels over FP8 +rows with ue8m0 block scales. The layer layout — table sizes, n-gram size, +head count — is fixed by the checkpoint and is not user-configurable. Only the +storage and sharding described below are configurable, via +`--engram-config`, which takes a JSON object and can also be set field by +field as `--engram-config.` +([CLI reference](../cli/serve.md#-engram-config)). + +## CPU offload + +Engram tables dwarf typical embedding weights: [DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) +carries two Engram layers of roughly 384M table rows each, with 256-byte FP8 +entries plus their ue8m0 block scales — about 200 GB of table weights alone. +Because lookups are deterministic — the indices are known from the token ids +alone, before the forward pass — the tables need not sit in GPU memory at +all. By default vLLM stores them in **pinned host memory** and serves lookups +directly from the Triton kernel through unified virtual addressing (UVA), +prefetching on a side CUDA stream so transfers overlap with the compute of +preceding layers. The freed GPU memory goes to the KV cache. + +`cpu_offload` defaults to `true` (on): it follows the +`VLLM_PLE_CPU_OFFLOAD` environment variable, which also defaults to `true`. +An explicit `--engram-config` value takes precedence over the environment +variable. + +To keep the tables resident on GPU instead (e.g. when host memory is +scarce): + +```bash +vllm serve deepseek-ai/DeepSeek-V4.1-Flash \ + --engram-config.cpu_offload false +``` + +or equivalently `VLLM_PLE_CPU_OFFLOAD=0`. CPU offload requires a GPU with UVA +support; vLLM fails fast when it is unavailable. + +## Data-parallel topologies (DeepSeek V4.1) + +By default (`embedding_across_dp: false`), each DP replica keeps its own +TP-sharded copy of the tables — no cross-replica communication, but one host +copy per replica. + +With `embedding_across_dp: true`, the hash heads are sharded across all +TP × DP ranks into a single table copy. Each step gathers the hash ids of +every co-located DP replica, each rank looks up the heads it owns, and the +rows are exchanged back. This trades per-step DP collectives for a much +smaller table footprint: + +```bash +vllm serve deepseek-ai/DeepSeek-V4.1-Flash \ + --tensor-parallel-size 2 --data-parallel-size 4 \ + --engram-config '{"embedding_across_dp": true}' +``` + +`embedding_across_dp` is not supported with elastic expert parallelism yet. +Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only. + +## Sharing host tables across DP replicas + +When the tables are CPU-offloaded and DP replicas are co-located on one node, +`dp_shared_memory` stores **one copy of each TP shard in `/dev/shm`**, shared +by all co-located replicas through `cudaHostRegister`. Each replica then +prefetches only its own tokens — no per-step Engram DP collectives — and host +memory drops from one table copy per replica to one per node. + +`dp_shared_memory` defaults to enabled whenever the other settings allow it, +and falls back to per-replica (or DP-sharded) tables with a warning when: + +- the DP replicas are not co-located on a single node, +- `/dev/shm` cannot hold the full tables — in containers, raise the limit + with `--shm-size` or `--ipc=host`, +- expert parallelism is elastic, which is unsupported. + +Requirements: `cpu_offload` enabled, `--data-parallel-size > 1`, and a shared +IPC namespace across the co-located replicas. + +## Limitations + +- CUDA-alike platforms only. +- DeepSeek V4.1 Engram does not support DBO or microbatching; disable + `--enable-dbo` and set `--ubatch-size 0`. +- `dp_shared_memory` requires `cpu_offload`, `data_parallel_size > 1`, and is + unsupported with elastic expert parallelism. From 17322375a2d7815ed3d7c9c3ffd718196564314e Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Mon, 21 Sep 2026 20:17:55 +0300 Subject: [PATCH 598/774] [ROCm][Perf] Route the fused shared-expert gate GEMM through the platform dispatcher (#54185) Signed-off-by: Mikko Tukiainen Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/layers/fused_moe/runner/moe_runner.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 8d3bc7942872..870d04c0c856 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -45,6 +45,7 @@ SharedExperts, SharedExpertsOrder, ) +from vllm.model_executor.layers.utils import dispatch_unquantized_gemm from vllm.platforms import current_platform from vllm.utils.torch_utils import ( _USE_LAYERNAME, @@ -273,7 +274,7 @@ def __init__( # When both gates are present and FSE is enabled, fuse their # weight matrices into [num_experts + num_shared, hidden] so one - # F.linear produces combined logits. The topk kernel can then + # GEMM produces combined logits. The topk kernel can then # apply routing softmax and shared expert activation (sigmoid) # in a single launch. self._fse_fuse_gate = gate is not None and shared_expert_gate is not None @@ -898,7 +899,9 @@ def _forward_impl( if self.gate is not None: if self._fse_fuse_gate: self._maybe_fuse_gate_weights() - router_logits = F.linear(hidden_states, self._combined_gate_weight) + router_logits = dispatch_unquantized_gemm()( + self, hidden_states, self._combined_gate_weight, None + ) else: router_logits, _ = self.gate(hidden_states) From 54020c3c3ec9de219929a5a956010a0251b87520 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 21 Sep 2026 20:16:02 +0200 Subject: [PATCH 599/774] [Engram] Drop redundant VLLM_PLE_CPU_OFFLOAD env var (#57937) Signed-off-by: NickLucche --- tests/test_config.py | 12 ++---------- vllm/config/engram.py | 13 +++---------- vllm/envs.py | 3 --- 3 files changed, 5 insertions(+), 23 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index e97e2cebaad2..7a6ea1d4c435 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1392,16 +1392,8 @@ def test_engram_model_support( assert resolved.engram_config.cpu_offload is True -@pytest.mark.parametrize( - ("value", "expected"), [(None, True), ("0", False), ("1", True)] -) -def test_engram_cpu_offload_environment_default(monkeypatch, value, expected): - monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) - if value is not None: - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", value) - assert EngramConfig().cpu_offload is expected - assert EngramConfig(cpu_offload=False).cpu_offload is False - assert EngramConfig(cpu_offload=True).cpu_offload is True +def test_engram_cpu_offload_default(): + assert EngramConfig().cpu_offload is True def test_engram_config_defaults_to_none(): diff --git a/vllm/config/engram.py b/vllm/config/engram.py index c9f05c93828f..0f5ec31b335f 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -3,10 +3,9 @@ from typing import TYPE_CHECKING -from pydantic import Field, model_validator +from pydantic import model_validator from typing_extensions import Self -import vllm.envs as envs from vllm.config.utils import config, get_hash_factors, hash_factors if TYPE_CHECKING: @@ -22,10 +21,6 @@ } -def _default_cpu_offload() -> bool: - return envs.VLLM_PLE_CPU_OFFLOAD - - def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: """Whether the model carries n-gram embedding layers.""" if model_config is None: @@ -40,10 +35,8 @@ def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: class EngramConfig: """Configuration for Engram embedding storage and sharding.""" - cpu_offload: bool = Field(default_factory=_default_cpu_offload) - """Store embedding weights in pinned CPU memory for UVA lookup. - Defaults to VLLM_PLE_CPU_OFFLOAD, which is enabled by default. An explicit - value takes precedence over the environment variable.""" + cpu_offload: bool = True + """Store embedding weights in pinned CPU memory for UVA lookup.""" embedding_across_dp: bool = False """Shard embeddings across TP and all DP ranks when enabled. diff --git a/vllm/envs.py b/vllm/envs.py index 259c08e7e7b3..fa98db1cf12a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -156,7 +156,6 @@ VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False VLLM_ENABLE_V1_MULTIPROCESSING: bool = True VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 - VLLM_PLE_CPU_OFFLOAD: bool = True VLLM_DISABLE_COMPILE_CACHE: bool = False VLLM_REPLICATE_EMBED: bool = False VLLM_USE_LAYERNAME: bool = True @@ -2066,8 +2065,6 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_LOG_MODEL_INSPECTION": lambda: bool( int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0")) ), - # Store n-gram embedding tables in pinned CPU memory for UVA lookup. - "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "1"))), # Debug logging for --enable-mfu-metrics "VLLM_DEBUG_MFU_METRICS": lambda: bool( int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0")) From 9179bcd4166bd03f4903aafed39d2c284eeb8c91 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 15:28:35 -0500 Subject: [PATCH 600/774] [Bugfix][MoE] Reject hash routing for unsupported monolithic backends (#57867) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/kernels/moe/conftest.py | 2 +- tests/kernels/moe/test_topk_softplus_sqrt.py | 59 +++++++++++++++++++ .../model_executor/layers/fused_moe/config.py | 1 + vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/modular_kernel.py | 2 + 5 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py index fdb30880879a..af99dc756b96 100644 --- a/tests/kernels/moe/conftest.py +++ b/tests/kernels/moe/conftest.py @@ -58,6 +58,6 @@ def pytest_runtest_setup(item): # Not tryfirst: the skipping plugin evaluates skip marks in a tryfirst # hook, so tests about to be skipped never reach this, while fixture setup # still runs afterwards. - if item.path.name in NEEDS_CLEAN_ENTRY: + if item.path.name in NEEDS_CLEAN_ENTRY or "dist_init" in item.fixturenames: reset_workspace_manager() cleanup_dist_env_and_memory() diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index ab9072576d92..bc438c66931b 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -81,6 +81,65 @@ def _torch_topk_softplus_sqrt( return topk_weights.to(torch.float32), topk_ids.to(torch.int32) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") +@pytest.mark.parametrize("has_hash_routing", [False, True]) +@pytest.mark.parametrize( + "backend,expected_experts", + [ + ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), + ("aiter", "AiterExperts"), + ("triton_unfused", "UnfusedOAITritonExperts"), + ], +) +def test_hash_routing_backend_selection( + dist_init, + default_vllm_config, + monkeypatch, + has_hash_routing, + backend, + expected_experts, +): + """Hash tables must reach backend selection; modular AITER remains eligible.""" + from dataclasses import replace + + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + select_deepseek_v4_mxfp4_moe_backend, + ) + from vllm.platforms import rocm + + # Exercise both native MX and monolithic AITER selection on any ROCm device. + monkeypatch.setattr(rocm, "on_gfx950", lambda: True) + monkeypatch.setattr(rocm, "on_gfx1250", lambda: False) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: True) + default_vllm_config.kernel_config.moe_backend = "triton" + layer = FusedMoEFactory( + num_experts=8, + top_k=2, + hidden_size=256, + intermediate_size=256, + params_dtype=torch.bfloat16, + scoring_func="sqrtsoftplus", + e_score_correction_bias=torch.zeros(8), + hash_indices_table=( + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + if has_hash_routing + else None + ), + prefix="hash_routing_selection", + ) + config = replace(layer.routed_experts.moe_config, moe_backend=backend) + + if has_hash_routing and backend == "aiter_triton_mxfp4_bf16": + with pytest.raises(ValueError, match="hash routing"): + select_deepseek_v4_mxfp4_moe_backend(config) + else: + _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) + assert experts_cls.__name__ == expected_experts + + def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id(): gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]]) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 62ab2bf88ae6..0c82d28096dc 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1282,6 +1282,7 @@ class FusedMoEConfig: elastic_ep_max_dp_size: int | None = None has_bias: bool = False is_lora_enabled: bool = False + has_hash_routing: bool = False # When True, the MoE skips its final cross-rank all-reduce (and the separate # shared-expert reduce), returning the partial per-rank sum. The caller is diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 49557bb3f9a0..ff620622c1f1 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -363,6 +363,7 @@ def FusedMoEFactory( activation=moe_activation, device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal + has_hash_routing=hash_indices_table is not None, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d91743dd951f..fba16f67b0f6 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -544,6 +544,8 @@ def _make_reason(reason: str) -> str: return False, _make_reason( f"parallel config {moe_config.moe_parallel_config}" ) + elif moe_config.has_hash_routing and cls.is_monolithic(): + return False, _make_reason("hash routing") elif not cls._supports_routing_method( moe_config.routing_method, weight_key, activation_key ): From 12972ae40bece6f9b9898fadd0c50d6a317df6fd Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 15:30:21 -0500 Subject: [PATCH 601/774] [Bugfix][ROCm] Reject unsupported EP for monolithic AITER MXFP4 MoE (#57866) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/kernels/moe/test_ocp_mx_moe.py | 55 +++++++++++++++++++ .../experts/aiter_mxfp4_w4a16_moe.py | 2 +- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 2 +- 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index d513a14f775a..1cd204d4b00b 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1868,6 +1868,61 @@ def is_supported_config( mxfp4_oracle.select_mxfp4_moe_backend(moe_config) +@pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"]) +@pytest.mark.parametrize("use_ep", [False, True]) +def test_aiter_mxfp4_monolithic_rejects_expert_parallel( + monkeypatch: pytest.MonkeyPatch, backend_name: str, use_ep: bool +): + """Global routing cannot index local expert weights, even without all-to-all.""" + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import RoutingMethodType + from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEActivationFormat, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + backend_to_kernel_cls, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kMxfp4Static, + ) + + experts_cls = backend_to_kernel_cls(Mxfp4MoeBackend[backend_name])[0] + monkeypatch.setattr(experts_cls, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + num_experts=256, + num_local_experts=64 if use_ep else 256, + experts_per_token=6, + hidden_dim=4096, + intermediate_size=2048, + activation=MoEActivation.SWIGLUOAI, + ) + config = replace( + config, + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=replace( + config.moe_parallel_config, + tp_size=1 if use_ep else 4, + ep_size=4 if use_ep else 1, + use_ep=use_ep, + ), + ) + assert not config.moe_parallel_config.use_all2all_kernels + supported, reason = experts_cls.is_supported_config( + experts_cls, + config, + kMxfp4Static, + kFp8StaticTensorSym if backend_name == "AITER_MXFP4_FP8" else None, + FusedMoEActivationFormat.Standard, + ) + + assert supported is not use_ep + if use_ep: + assert "parallel config" in reason + + # Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that # `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns # the activation untouched, so a name it does not know (e.g. "mxfp6" instead of diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index f6afd6fbfa34..fb3256e1a4c1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -321,7 +321,7 @@ def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, ) -> bool: return ( - not moe_parallel_config.use_all2all_kernels + not moe_parallel_config.use_ep and not moe_parallel_config.enable_eplb and moe_parallel_config.dp_size <= 1 ) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index b327cbeb594c..407dc097b16e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -246,7 +246,7 @@ def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, ) -> bool: return ( - not moe_parallel_config.use_all2all_kernels + not moe_parallel_config.use_ep and not moe_parallel_config.enable_eplb and moe_parallel_config.dp_size <= 1 ) From 9fe1cbd463055314863a6de210d5f898200a0a9a Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Mon, 21 Sep 2026 13:40:15 -0700 Subject: [PATCH 602/774] [Docker] Use zstd for CI images and offer a Docker Hub variant (#55608) Signed-off-by: Nils Matteson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/image_build/image_build.sh | 8 ++-- .buildkite/image_build/zstd.hcl | 3 ++ .buildkite/scripts/publish-release-images.sh | 7 +++ .buildkite/scripts/publish-zstd-image.sh | 45 ++++++++++++++++++++ 4 files changed, 60 insertions(+), 3 deletions(-) create mode 100644 .buildkite/image_build/zstd.hcl create mode 100755 .buildkite/scripts/publish-zstd-image.sh diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh index b81cd4ec0cbe..a1134e9a4d9a 100755 --- a/.buildkite/image_build/image_build.sh +++ b/.buildkite/image_build/image_build.sh @@ -154,7 +154,7 @@ print_bake_config() { local bake_tmp bake_tmp="$(mktemp -d)" BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json" - docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true + docker buildx bake "${BAKE_FILES[@]}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true echo "Saved bake config to ${BAKE_CONFIG_FILE}" echo "--- :arrow_down: Uploading bake config to Buildkite" (cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")") @@ -230,8 +230,10 @@ TARGET="test-ci" VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}" BUILDER_NAME="${BUILDER_NAME:-vllm-builder}" CI_HCL_URL="${CI_HCL_URL:-https://raw.githubusercontent.com/vllm-project/ci-infra/main/docker/ci.hcl}" -CI_HCL_PATH="/tmp/ci.hcl" +CI_HCL_PATH="${CI_HCL_PATH:-/tmp/ci.hcl}" +ZSTD_HCL_PATH="${ZSTD_HCL_PATH:-.buildkite/image_build/zstd.hcl}" BUILDKIT_SOCKET="/run/buildkit/buildkitd.sock" +BAKE_FILES=(-f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" -f "${ZSTD_HCL_PATH}") prepare_cache_tags ecr_login @@ -309,7 +311,7 @@ BUILD_TMP_DIR="$(mktemp -d)" trap 'rm -rf -- "${BUILD_TMP_DIR}"' EXIT BUILD_METADATA_FILE="${BUILD_TMP_DIR}/build-metadata.json" BUILD_STATUS=0 -docker --debug buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" \ +docker --debug buildx bake "${BAKE_FILES[@]}" \ --progress plain --metadata-file "${BUILD_METADATA_FILE}" "${TARGET}" || BUILD_STATUS=$? record_buildkit_trace "${BUILD_METADATA_FILE}" || true diff --git a/.buildkite/image_build/zstd.hcl b/.buildkite/image_build/zstd.hcl new file mode 100644 index 000000000000..1f795c5f9150 --- /dev/null +++ b/.buildkite/image_build/zstd.hcl @@ -0,0 +1,3 @@ +target "test-ci" { + output = ["type=registry,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true"] +} diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index 2fa334241fce..dec366390b00 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -61,6 +61,13 @@ if target_enabled cuda-13-0; then docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" docker manifest push vllm/vllm-openai:latest docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}" + + ZSTD_DIGEST=$(docker buildx imagetools inspect \ + "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" \ + --format '{{json .Manifest.Digest}}' | tr -d '"') + .buildkite/scripts/publish-zstd-image.sh \ + "public.ecr.aws/q9t5s3a7/vllm-release-repo@${ZSTD_DIGEST}" \ + "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-zstd" fi # ---- CUDA 12.9 ---- diff --git a/.buildkite/scripts/publish-zstd-image.sh b/.buildkite/scripts/publish-zstd-image.sh new file mode 100755 index 000000000000..bd3d2f99d65c --- /dev/null +++ b/.buildkite/scripts/publish-zstd-image.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +if [ "$#" -ne 2 ]; then + echo "Usage: $0 SOURCE@sha256:DIGEST DESTINATION-zstd" >&2 + exit 2 +fi + +SOURCE=$1 +DESTINATION=$2 +if [[ ! "$SOURCE" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*@sha256:[a-f0-9]{64}$ ]] || + [[ ! "$DESTINATION" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*:[a-zA-Z0-9_][a-zA-Z0-9_.-]*-zstd$ ]]; then + echo "ERROR: source must use an immutable sha256 digest and destination must end in -zstd" >&2 + exit 2 +fi + +BUILDER= +cleanup() { + status=$? + trap - EXIT INT TERM + if [ -n "$BUILDER" ] && ! docker buildx rm --force "$BUILDER"; then + echo "ERROR: failed to remove Buildx builder $BUILDER" >&2 + [ "$status" -ne 0 ] || status=1 + fi + exit "$status" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +BUILDER=$(docker buildx create \ + --driver docker-container \ + --driver-opt network=host \ + --driver-opt 'image=moby/buildkit:v0.32.2@sha256:28a898719c18a33f4e8000685287fa36fd0dd9560c6440227d3a732d79bb41d8') + +printf '# check=error=true\n\nFROM %s\n' "$SOURCE" | docker buildx build \ + --builder "$BUILDER" \ + --platform linux/amd64 \ + --progress plain \ + --provenance=mode=min \ + --output "type=image,name=${DESTINATION},push=true,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true" \ + - From 382970ee6ca490aeaaaf4e32c53695b581ff61ba Mon Sep 17 00:00:00 2001 From: yinfengLiu Date: Tue, 22 Sep 2026 04:44:41 +0800 Subject: [PATCH 603/774] [Kimi-K3][AMD] Return KDA and MLA projection outputs directly (#50592) Signed-off-by: Liuyinfeng01 Co-authored-by: Liuyinfeng01 <199041580+LiuYinfeng01@users.noreply.github.com> Co-authored-by: Andreas Karatzas --- .../kimi_k3/test_amd_kda_direct_return.py | 72 +++++++++++++++++++ .../kimi_k3/test_amd_mla_direct_return.py | 38 ++++++++++ .../kimi_k3/test_nvidia_kda_direct_return.py | 31 ++++++++ .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 5 +- vllm/models/kimi_k3/amd/kda.py | 5 +- vllm/models/kimi_k3/amd/linear.py | 41 ++++++----- vllm/models/kimi_k3/nvidia/model.py | 11 --- 7 files changed, 170 insertions(+), 33 deletions(-) create mode 100644 tests/models/kimi_k3/test_amd_kda_direct_return.py create mode 100644 tests/models/kimi_k3/test_amd_mla_direct_return.py create mode 100644 tests/models/kimi_k3/test_nvidia_kda_direct_return.py diff --git a/tests/models/kimi_k3/test_amd_kda_direct_return.py b/tests/models/kimi_k3/test_amd_kda_direct_return.py new file mode 100644 index 000000000000..e46efe7681a1 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_direct_return.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention +from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer + + +class _ReturningAttention(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + return self.result + + +class _TupleProjection(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward(self, _: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.result, None + + +def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer: + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = self_attn + return layer + + +def test_kda_self_attention_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = _make_layer(_ReturningAttention(projected)) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() + + +def test_k3_kda_forward_returns_output_projection_storage() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + + def fill_core_output(**kwargs: torch.Tensor) -> None: + kwargs["core_attn_out"].fill_(1) + + attention = SimpleNamespace( + local_projection_size=4, + head_dim=2, + local_num_heads=2, + in_proj_padding=0, + in_proj_qkvgfab=_TupleProjection(torch.randn(4, 20)), + f_b_proj=_TupleProjection(torch.randn(4, 4)), + o_proj=_TupleProjection(projected), + _forward=fill_core_output, + ) + + output = KimiK3DeltaAttention.forward(attention, hidden_states, torch.arange(4)) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py new file mode 100644 index 000000000000..a1c120737df8 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer, KimiMLAAttention + + +class _DirectMLA(KimiMLAAttention): + def __init__(self, result: torch.Tensor) -> None: + nn.Module.__init__(self) + self.result = result + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + return self.result + + +def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer: + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = self_attn + return layer + + +def test_mla_self_attention_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = _make_layer(_DirectMLA(projected)) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/tests/models/kimi_k3/test_nvidia_kda_direct_return.py b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py new file mode 100644 index 000000000000..5b40d7ef6fa1 --- /dev/null +++ b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.kimi_k3.nvidia.model import KimiDecoderLayer + + +class _ReturningSharedKDA(nn.Module): + def __init__(self, result: torch.Tensor) -> None: + super().__init__() + self.result = result + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + return self.result + + +def test_low_rank_kda_caller_returns_projection_storage_directly() -> None: + hidden_states = torch.randn(4, 8) + projected = torch.randn_like(hidden_states) + layer = object.__new__(KimiDecoderLayer) + nn.Module.__init__(layer) + layer.self_attn = _ReturningSharedKDA(projected) + + output = layer._run_self_attn(torch.arange(4), hidden_states) + + assert output is projected + assert output.data_ptr() == projected.data_ptr() diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index bf81cdd75f49..3892d1a27aa3 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -329,8 +329,7 @@ def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, - ) -> None: + ) -> torch.Tensor: num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] if self.use_full_rank_gate: @@ -375,7 +374,7 @@ def forward( core_attn_out=core_attn_out, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index a02bf90d38cd..ed631aea7e85 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -271,8 +271,7 @@ def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, - ) -> None: + ) -> torch.Tensor: num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] @@ -306,7 +305,7 @@ def forward( core_attn_out=core_attn_out, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") - output[:] = self.o_proj(core_attn_out)[0] + return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 2a013ed0fc08..65dad661ab12 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -456,9 +456,8 @@ def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, - output: torch.Tensor, - ) -> None: - output[:] = self.mla_attn(positions, hidden_states) + ) -> torch.Tensor: + return self.mla_attn(positions, hidden_states) class KimiDecoderLayer(nn.Module): @@ -582,25 +581,29 @@ def _run_self_attn( positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: - attn_output = torch.empty_like(hidden_states) - self.self_attn( + return self.self_attn( hidden_states=hidden_states, positions=positions, - output=attn_output, ) - return attn_output def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, + prefix_delta: torch.Tensor | None = None, **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> ( + tuple[torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ): if self.use_attn_residuals: assert residual is not None - return self.forward_attn_residual(positions, hidden_states, residual) + return self.forward_attn_residual( + positions, hidden_states, residual, prefix_delta + ) + assert prefix_delta is None # Self Attention if residual is None: residual = hidden_states @@ -620,7 +623,8 @@ def forward_attn_residual( positions: torch.Tensor, hidden_states: torch.Tensor, block_residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + prefix_delta: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: prefix_sum = hidden_states hidden_states = _apply_attn_res( prefix_sum, @@ -628,6 +632,7 @@ def forward_attn_residual( self.self_attention_res_proj, self.self_attention_res_norm, self.prev_valid_blocks, + delta=prefix_delta, output_norm=self.input_layernorm, block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), ) @@ -657,8 +662,7 @@ def forward_attn_residual( ) hidden_states = self.mlp(hidden_states) - prefix_sum = prefix_sum + hidden_states - return prefix_sum, block_residual + return prefix_sum, block_residual, hidden_states class KimiLinearModel(nn.Module, EagleModelMixin): @@ -813,24 +817,28 @@ def forward( if residual is not None: block_residual[:, : residual.size(1), :].copy_(residual) residual = block_residual + prefix_delta = None for layer_idx, layer in enumerate( self.layers[self.start_layer : self.end_layer], start=self.start_layer, ): - hidden_states, residual = layer( + hidden_states, residual, prefix_delta = layer( positions=positions, hidden_states=hidden_states, residual=residual, + prefix_delta=prefix_delta, ) if (layer_idx + 1) in self.aux_hidden_state_layers: - # AMD attn-res layer already returns prefix_sum + MLP delta as - # hidden_states; the override drops the block bank in residual. self._maybe_add_hidden_state( - aux_hidden_states, layer_idx + 1, hidden_states, residual + aux_hidden_states, + layer_idx + 1, + hidden_states + prefix_delta, + residual, ) if not get_pp_group().is_last_rank: + hidden_states = hidden_states + prefix_delta return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) @@ -841,6 +849,7 @@ def forward( self.output_attn_res_proj, self.output_attn_res_norm, attn_res_block_num, + delta=prefix_delta, ) # NOTE: the final norm is applied in compute_logits instead of here, so # the MTP draft model receives the pre-norm hidden states. diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 737184b9d836..5dedf44b7c3d 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -922,14 +922,12 @@ def __init__( aux_stream=aux_stream, run_gemm_rs_ar=run_gemm_rs_ar, ) - self._self_attn_writes_output = False else: self.self_attn = KimiLinearGatedDeltaNetAttention( config, vllm_config, prefix=f"{prefix}.self_attn", ) - self._self_attn_writes_output = True else: qk_nope_head_dim = config.qk_nope_head_dim qk_rope_head_dim = config.qk_rope_head_dim @@ -959,7 +957,6 @@ def __init__( aux_stream=aux_stream, run_gemm_rs_ar=run_gemm_rs_ar, ) - self._self_attn_writes_output = False if self.use_sequence_parallel: self.self_attn.o_proj.reduce_results = False @@ -1025,14 +1022,6 @@ def _run_self_attn( positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: - if self._self_attn_writes_output: - output = torch.empty_like(hidden_states) - self.self_attn( - hidden_states=hidden_states, - positions=positions, - output=output, - ) - return output return self.self_attn( hidden_states=hidden_states, positions=positions, From 8644d2af2fb7c4588190a8d7f6c26a3cbef78a0f Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Mon, 21 Sep 2026 15:30:14 -0700 Subject: [PATCH 604/774] [CI][Build] Harden triton-cpu sleef submodule fetch in CPU image build (#57871) Signed-off-by: kimi-no-na-wa Co-authored-by: kimi-no-na-wa --- docker/Dockerfile.cpu | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index f81ae880966c..7898e57cf39a 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -140,6 +140,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \ cd triton-cpu; \ git checkout "270e696d"; \ + if [ ! -f third_party/sleef/CMakeLists.txt ]; then \ + echo "sleef submodule empty after shallow fetch; retrying unshallow"; \ + git submodule foreach --recursive 'git fetch --unshallow origin || true'; \ + git submodule update --init --recursive --force; \ + fi; \ + test -f third_party/sleef/CMakeLists.txt || { echo "FATAL: sleef submodule unavailable"; exit 1; }; \ if [ "$USE_SCCACHE" = "1" ]; then \ if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \ unset SCCACHE_BUCKET SCCACHE_REGION; \ From 8bf051cd51adff79e1bc94e910adaea543aaa2e1 Mon Sep 17 00:00:00 2001 From: Garrett Goon <44747910+garrett361@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:38:38 -0400 Subject: [PATCH 605/774] [Bugfix][Kernel] Skip the fused silu-mul block-quant fast path when a swiglu clamp is set (#57984) Signed-off-by: Garrett Goon Co-authored-by: Claude Opus 5 (1M context) --- tests/kernels/moe/test_block_fp8.py | 55 +++++++++++++++++++ .../layers/fused_moe/experts/triton_moe.py | 3 + 2 files changed, 58 insertions(+) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 1cbbbb45fb09..63f64c735ab8 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a8_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.experts import triton_moe from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( _valid_deep_gemm_shape, ) @@ -246,6 +247,60 @@ def test_w8a8_block_fp8_fused_moe( torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol) +@pytest.mark.parametrize("clamp_limit", [None, 10.0]) +@torch.inference_mode() +def test_w8a8_block_fp8_clamp_skips_fused_silu_quant( + clamp_limit, monkeypatch, workspace_init, disable_deepgemm_ue8m0 +): + """A SwiGLU clamp must bypass the fused SiLU+quant kernel, which cannot clamp.""" + torch.manual_seed(0) + M, N, K, E, topk = 16, 256, 512, 8, 2 + block_size = [128, 128] + + a = torch.randn((M, K), dtype=torch.bfloat16) / 10 + score = torch.randn((M, E), dtype=torch.bfloat16) + (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( + E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size + ) + quant_config = fp8_w8a8_moe_quant_config( + w1_scale=w1_s, + w2_scale=w2_s, + block_shape=block_size, + gemm1_clamp_limit=clamp_limit, + ) + m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) + + fused_calls = 0 + fused_op = triton_moe.ops.silu_and_mul_per_block_quant + + def counting_fused_op(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return fused_op(*args, **kwargs) + + monkeypatch.setattr( + triton_moe.ops, "silu_and_mul_per_block_quant", counting_fused_op + ) + + with set_current_vllm_config(vllm_config): + m_fused_moe.apply( + a, + w1, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, + global_num_experts=E, + ) + + assert (fused_calls > 0) == (clamp_limit is None), ( + f"fused fast path ran {fused_calls} times with clamp_limit={clamp_limit}" + ) + + @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG) @pytest.mark.parametrize("E", E) @pytest.mark.parametrize("topk", TOP_KS) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 700d113921e8..1e21a4fe2c15 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -479,8 +479,11 @@ def _lora_w13_fn(): # Fuse SiLU+Mul + FP8 block quantize into a single kernel # when conditions permit (gated SiLU, fp8 block quant with # group_size=128, no LoRA requiring the BF16 intermediate). + # The fused kernel has no clamp parameter, so a configured + # SwiGLU clamp limit falls through to the unfused path. if ( activation == MoEActivation.SILU + and self.activation_config.clamp_limit is None and self.quant_config.use_fp8_w8a8 and self.block_shape == [128, 128] and lora_context is None From 13058faa49f5e4a4d50e1ca562ba0c2fdd73c7bc Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Mon, 21 Sep 2026 16:18:49 -0700 Subject: [PATCH 606/774] [Frontend] Add reusable TP1 initialized-engine snapshots (#51360) Signed-off-by: Nils Matteson Signed-off-by: Taneem Ibrahim Signed-off-by: Andreas Karatzas Co-authored-by: Taneem Ibrahim Co-authored-by: elehayym <52448798+Yuzu23@users.noreply.github.com> Co-authored-by: Andreas Karatzas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../scripts/initialized-snapshot-e2e.sh | 311 ++++ .buildkite/test_areas/entrypoints.yaml | 20 + docker/Dockerfile | 16 + docs/features/initialized_snapshots.md | 152 ++ .../entrypoints/launchers/test_launch_cli.py | 1043 +++++++++++++- .../quantization/test_mimo_v2_qkv_shard.py | 8 + tools/install_snapshot_runtime.sh | 109 ++ vllm/entrypoints/cli/main.py | 12 +- vllm/entrypoints/cli/snapshot.py | 143 ++ vllm/snapshot/__init__.py | 4 + vllm/snapshot/controller.py | 132 ++ vllm/snapshot/manifest.py | 277 ++++ vllm/snapshot/runtime.py | 1246 +++++++++++++++++ vllm/snapshot/server.py | 252 ++++ vllm/snapshot/types.py | 37 + vllm/utils/argparse_utils.py | 21 +- 16 files changed, 3773 insertions(+), 10 deletions(-) create mode 100755 .buildkite/scripts/initialized-snapshot-e2e.sh create mode 100644 docs/features/initialized_snapshots.md create mode 100755 tools/install_snapshot_runtime.sh create mode 100644 vllm/entrypoints/cli/snapshot.py create mode 100644 vllm/snapshot/__init__.py create mode 100644 vllm/snapshot/controller.py create mode 100644 vllm/snapshot/manifest.py create mode 100644 vllm/snapshot/runtime.py create mode 100644 vllm/snapshot/server.py create mode 100644 vllm/snapshot/types.py diff --git a/.buildkite/scripts/initialized-snapshot-e2e.sh b/.buildkite/scripts/initialized-snapshot-e2e.sh new file mode 100755 index 000000000000..cd0511ffdd9a --- /dev/null +++ b/.buildkite/scripts/initialized-snapshot-e2e.sh @@ -0,0 +1,311 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -Eeuo pipefail +readonly MODEL="Qwen/Qwen3-0.6B" +readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca" +readonly EXPECTED_TOKEN_ID="12095" +readonly EXPECTED_TEXT=" Paris" +readonly CLEANUP_RESERVE_S=300 +RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}" +[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \ + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; } +(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \ + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; } +readonly RUN_LIMIT_S +readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}" +HF_CACHE_MOUNT=() +if [[ -n "$HOST_HF_CACHE" ]]; then + [[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \ + || { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; } + HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly") + PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e' +else + PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf' +fi +readonly PREPARE_ARTIFACT_ROOT +if (( $# != 2 )); then + echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2 + exit 2 +fi +readonly IMAGE_TAG="$1" +readonly EXPECTED_COMMIT="$2" +readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))" +RUN_ID="" CONTAINER_NAME="" CONTAINER_ID="" +GPU_UUID="" PORT_ONE="" PORT_TWO="" +LINK_REMAP_BASELINE="" +SNAPSHOT_PIDS=() +die() { + echo "initialized snapshot E2E: $*" >&2 + exit 1 +} + +run() { + local label="$1" maximum_s="$2" available limit + shift 2 + available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))" + (( available > 0 )) || die "no time remains before the cleanup reserve" + (( available < maximum_s )) && limit="$available" || limit="$maximum_s" + echo "--- ${label} (timeout ${limit}s)" >&2 + timeout --signal=TERM --kill-after=30s "${limit}s" "$@" +} + +link_remaps() { + find /dev/shm -maxdepth 1 -name 'link_remap.*' \ + -printf '%p|%D|%i|%m|%s\n' | sort +} + +gpu_pids() { + nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \ + --format=csv,noheader,nounits | sed '/^[[:space:]]*$/d' +} + +port_listeners() { + ss -H -ltn "sport = :$1" +} + +state_clean() { + local pid port output + for pid in "${SNAPSHOT_PIDS[@]}"; do + [[ ! -e "/proc/$pid" ]] || return 1 + done + output="$(gpu_pids)" || return 1 + [[ -z "$output" ]] || return 1 + output="$(link_remaps)" || return 1 + [[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1 + for port in "$PORT_ONE" "$PORT_TWO"; do + [[ -n "$port" ]] || continue + output="$(port_listeners "$port")" || return 1 + [[ -z "$output" ]] || return 1 + done +} + +wait_clean() { + local _attempt + for ((_attempt = 0; _attempt < 60; _attempt++)); do + state_clean && return 0 + sleep 1 + done + return 1 +} + +container_names() { + timeout 15s docker container ls --all --format '{{.Names}}' +} + +cleanup() { + local failed=0 names="" actual="" expected="" + if [[ -n "$CONTAINER_NAME" ]]; then + names="$(container_names 2>/dev/null)" || return 1 + if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then + actual="$(docker container inspect \ + --format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \ + "$CONTAINER_NAME" 2>/dev/null)" || return 1 + expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID" + if [[ "$actual" != "$expected" ]]; then + echo "refusing to remove a container with mismatched identity" >&2 + failed=1 + else + timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \ + || failed=1 + fi + fi + names="$(container_names 2>/dev/null)" || return 1 + grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1 + fi + if [[ -n "$PORT_TWO" ]] && ! wait_clean; then + echo "process, GPU, port, or link-remap residue remained" >&2 + failed=1 + fi + return "$failed" +} + +on_exit() { + local status=$? + trap - EXIT INT TERM + set +e + cleanup || status=1 + exit "$status" +} + +trap on_exit EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +parse_inspect() { + timeout 30s python3 -c ' +import json, math, sys +data = json.load(sys.stdin) +model, revision, gpu_uuid, token_id, text = sys.argv[1:] +expected = { + "boundary": "post-engine-init-reloadable-state-released", + "model": model, "model_revision": revision, + "gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)], + "oracle_text": text, +} +bad = next((field for field, value in expected.items() if data.get(field) != value), None) +if bad: + raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}") +logprob = data.get("oracle_sampled_token_logprob") +if type(logprob) is not float or not math.isfinite(logprob): + raise SystemExit("inspect logprob is not finite") +pids = data.get("process_tree") +if not pids: + raise SystemExit("inspect process_tree is empty") +for pid in pids: + print(f"pid={pid}") +print(f"oracle_logprob={logprob}") +' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \ + "$EXPECTED_TEXT" +} + +assert_oracle() { + local port="$1" ordinal="$2" response + response="$(run "restore ${ordinal} public completion" 120 curl \ + --fail --silent --show-error --max-time 120 \ + --header 'Content-Type: application/json' \ + --data "{\"model\":\"${MODEL}\",\"prompt\":\"The capital of France is\",\"temperature\":0,\"seed\":0,\"min_tokens\":1,\"max_tokens\":1,\"logprobs\":0,\"return_token_ids\":true}" \ + "http://127.0.0.1:${port}/v1/completions")" + printf '%s' "$response" | timeout 30s python3 -c ' +import json, math, sys +choice = json.load(sys.stdin)["choices"][0] +token_id, text, expected, ordinal = sys.argv[1:] +values = choice["logprobs"]["token_logprobs"] +actual = float(values[0]) if len(values) == 1 else math.nan +if (choice["token_ids"] != [int(token_id)] or choice["text"] != text + or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3): + raise SystemExit(f"public oracle mismatch: {choice!r}") +print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}") +' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal" +} + +[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \ + || die "host must be Linux x86_64" +[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \ + || die "BUILDKITE_COMMIT must be a full lowercase commit hash" +[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \ + || die "kernel.io_uring_disabled must already equal 2" +run "Docker preflight" 60 docker info >/dev/null + +VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}" +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then + VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}" +fi +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then + CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}" + if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then + VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}" + else + VISIBLE_GPU="" + fi +fi +[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \ + && "$VISIBLE_GPU" != *","* ]] \ + || die "the runner must expose exactly one selected GPU" +GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \ + --query-gpu=name,uuid,mig.mode.current,memory.total \ + --format=csv,noheader,nounits)" +IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW" +for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do + printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}" +done +[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \ + || die "selected device is not a full non-MIG H200: $GPU_ROW" +(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle" + +RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM" +CONTAINER_NAME="vllm-$RUN_ID" +read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))') +LINK_REMAP_BASELINE="$(link_remaps)" + +run "pull exact candidate image" 900 docker pull --platform linux/amd64 \ + "$IMAGE_TAG" >/dev/null +IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <( + docker image inspect \ + --format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \ + "$IMAGE_TAG" +) +[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \ + || die "candidate image identity is invalid" +[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \ + || die "candidate image commit labels do not match" +CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \ + --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \ + --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \ + --network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \ + --env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \ + ${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \ + --entrypoint sleep "$IMAGE_ID" infinity)" +[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID" +run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \ + "$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'" +run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \ + 'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so' +run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c ' +import sys, torch +if torch.cuda.device_count() != 1: + raise SystemExit("expected exactly one visible CUDA device") +actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-") +if actual != sys.argv[1].removeprefix("GPU-"): + raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}") +print(f"selected_cuda_uuid=GPU-{actual}") +' "$GPU_UUID" +OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \ + --env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \ + "$CONTAINER_NAME") +if [[ -n "$HOST_HF_CACHE" ]]; then + run "validate pinned model in read-only host cache" 120 \ + "${OFFLINE_EXEC[@]}" python3 -c \ + 'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2], local_files_only=True); assert Path(p).name == sys.argv[2]' \ + "$MODEL" "$MODEL_REVISION" +else + run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \ + "$CONTAINER_NAME" python3 -c \ + 'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2]); assert Path(p).name == sys.argv[2]' \ + "$MODEL" "$MODEL_REVISION" +fi + +run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \ + vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \ + --tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \ + --dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50 +wait_clean || die "create left process, GPU, port, or link-remap residue" + +INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \ + "${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)" +INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)" +mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT") +INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")" +(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs" +state_clean || die "create left recorded state alive" +echo "$INSPECT_RECEIPT" + +run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \ + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE" +PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port" +[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU" +assert_oracle "$PORT_ONE" 1 +run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \ + || die "docker stop replaced the container" +wait_clean || die "restore 1 teardown left residue" + +run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \ + || die "docker start replaced the container" +run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \ + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO" +PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port" +[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU" +assert_oracle "$PORT_TWO" 2 +run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null +wait_clean || die "restore 2 teardown left residue" +cleanup || die "exact cleanup failed" +echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO" +echo "result=PASS" diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 60ac5d932c3a..9e0a8399a36d 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -2,6 +2,26 @@ group: Entrypoints depends_on: - image-build steps: +- label: ":nvidia: (H200) Initialized Snapshot E2E" + key: initialized-snapshot-e2e + device: h200 + num_devices: 1 + no_plugin: true + optional: true + timeout_in_minutes: 60 + working_dir: "." + env: + SNAPSHOT_E2E_RUN_LIMIT_S: "1200" + source_file_dependencies: + - .buildkite/scripts/initialized-snapshot-e2e.sh + - .buildkite/test_areas/entrypoints.yaml + - docker/Dockerfile + - tools/install_snapshot_runtime.sh + - vllm/entrypoints/cli/ + - vllm/snapshot/ + commands: + - bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT" + - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit" device: h200_35gb key: entrypoints-unit-tests diff --git a/docker/Dockerfile b/docker/Dockerfile index ad010285b181..5fd25f538855 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -735,6 +735,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \ FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base ARG INSTALL_RUBIN_PRERELEASE +ARG TARGETPLATFORM +ARG max_jobs=2 ARG CUDA_VERSION ARG PYTHON_VERSION ARG NCCL_VERSION @@ -854,6 +856,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ && ldconfig; \ fi +# Install the initialized-snapshot runtime on the architecture supported by +# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned +# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's +# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of +# redistributing NVIDIA's proprietary cuda-checkpoint binary. +COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh +RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \ + MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \ + else \ + echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \ + fi + +ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu + # Install uv for faster pip installs RUN python3 -m pip install uv diff --git a/docs/features/initialized_snapshots.md b/docs/features/initialized_snapshots.md new file mode 100644 index 000000000000..538695d0a8be --- /dev/null +++ b/docs/features/initialized_snapshots.md @@ -0,0 +1,152 @@ +# Initialized engine snapshots + +Initialized engine snapshots are an experimental way to trade local disk space +and host privileges for a faster vLLM activation. Snapshot creation initializes +the engine and records deterministic generation output before CRIU captures the +process tree and CUDA state. Restore validates the saved environment, restores +the engine, binds the HTTP server, and reproduces the recorded token and sampled +token log probability before returning. + +This path is intended for repeatedly activating the same model and engine +configuration on the same machine. It is not a portable model artifact. + +## Requirements + +Snapshots currently require: + +- Linux on x86-64 with one NVIDIA GPU. +- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes, + TLS, middleware, Unix sockets, and speculative decoding are unsupported. +- Because snapshot mode requires an unauthenticated plaintext HTTP server, use + a trusted host or network boundary, or external controls. See the vLLM + [Security guide](../usage/security.md). +- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a + `cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`. +- Root or passwordless `sudo` for CRIU. +- `io_uring` disabled before launch because CRIU cannot dump it. Use + `kernel.io_uring_disabled=1` for an unprivileged process. A process running + as root, including the `docker exec` flow below, bypasses `=1`, so use `=2` + host-wide. +- No established TCP connection to a peer outside the captured process tree. + Current Hugging Face hub clients hold their connections for the process + lifetime, so download the model in a separate step and run create with + `HF_HUB_OFFLINE=1`, as the quickstart below does. +- A remote model ID and an immutable 40-character `--revision`. Local model + directories and mutable revisions are not supported. +- Enough disk for the artifact, with the same installed vLLM package, model + files, container filesystem, and generated-cache paths available at restore. + Generated-cache files are not copied into the artifact. The manifest + fingerprints the ones the captured tree holds open, so restore fails early + and names the file when one is removed or replaced. Only open descriptors + are recorded. A library the tree mapped and then closed is not fingerprinted, + and CRIU still reopens it by path, so replacing that file permanently + invalidates the artifact without an early error. + +The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot +runtime. It still requires a compatible host driver, kernel, and privileges. +Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the +directory containing `cuda_plugin.so`. + +Run snapshot commands with `docker exec` inside a long-lived container. Restore +hands the API server off as a detached process, so a one-shot container would +stop that server when its PID 1 exits. Snapshot preflight requires every +component of the artifact path to be owned by the invoking user or root, with +the directory itself at mode 0700, and the commands below run as root inside +the official image, so the bind-mounted host directory is created with `sudo` +and root ownership. The model downloads in its own step so the captured tree +holds no hub connection, and create runs offline. This example also keeps the +container filesystem and `/dev/shm` namespace stable for the lifetime of the +artifact: + +```bash +sudo sysctl kernel.io_uring_disabled=2 + +snapshot_root="$(pwd)/vllm-snapshots" +sudo install -d -m 0700 -o root -g root "${snapshot_root}" + +docker run --detach --name vllm-snapshot \ + --gpus all \ + --privileged \ + --pid=host \ + --ipc=host \ + --network=host \ + --mount "type=bind,source=${snapshot_root},target=/snapshots" \ + --entrypoint sleep \ + vllm/vllm-openai:latest infinity + +docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \ + --revision c1899de289a04d12100db370d81485cdf75e47ca + +docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \ + --snapshot-dir /snapshots/qwen3-0.6b \ + --revision c1899de289a04d12100db370d81485cdf75e47ca \ + --dtype float16 \ + --max-model-len 512 + +docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b + +docker exec vllm-snapshot vllm snapshot restore \ + /snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000 +``` + +Keep that container and its mounts available while the snapshot is in use. +Stop and remove it only after the restored API server is no longer needed. + +Create initializes the engine, records a one-token canary, releases and reloads +weights and KV cache to rehearse restore, then releases them again for capture. +The manifest is published only after CRIU completes and the source tree stops. +Creation is offline preparation and is not part of restore latency. + +The private `0700` artifact contains process memory, CUDA state, engine +arguments, compatibility identity, and canary output; its manifest is `0600`. +Literal API keys and Hugging Face tokens are redacted from the manifest's engine +arguments. The manifest records selected environment names plus deterministic +name-bound fingerprints, not their values. The protected CRIU artifact can +still contain process secrets, so treat it as sensitive data. Restore reloads +model files and KV cache before binding HTTP, then reproduces the canary or +tears down the restored tree. The inspect command prints that identity and +canary without executing the saved process. + +## Restore behavior + +Restore fails before CRIU runs if the saved identity does not match the current +host. It does not silently fall back to ordinary startup. After CRIU restores +the process tree, vLLM releases the saved engine to bind the requested HTTP +address and checks the first generated token and sampled-token log probability +against the snapshot canary. The command returns only after that check passes. +The restored API server continues to run as a detached process. + +The pre-release port probe is best-effort only. It neither reserves the port +nor authenticates the listener that appears afterward. + +Rollback terminates and waits for the restored tree after process identity +verification succeeds. If verification fails, vLLM writes an abort marker for +the snapshot server instead of signaling unverified PIDs. If that server has +already exited, surviving engine processes can require operator cleanup. +Verify that the previous tree has stopped before retrying; do not kill a process +solely because its PID appears in the manifest, since PIDs can be reused. + +The artifact is reusable after its previous restored tree stops. Only one +snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time. + +## Tradeoffs and limitations + +- Creation has its own latency and briefly requires the full engine. Artifact + size can approach the captured process and GPU memory. +- Artifacts are not guaranteed to survive a power loss. The manifest is + fsynced, but publication does not explicitly flush all captured payload files + and directories. Recreate the artifact after an unclean host shutdown. +- Restore currently requires the same host, GPU, driver, kernel, Python, + PyTorch, installed vLLM version, model revision, engine arguments, + selected environment variables, and CRIU plugin binaries. +- Only dense float16 TP1 has been validated. Other model formats depend on their + existing sleep level 2 reload support; NCCL state is not restored. +- CRIU support varies by kernel and driver. Preserve package, library, model, + and generated-cache paths for the artifact lifetime. +- A snapshot can include application secrets or request state present in the + process. Build it before traffic and protect it like process memory. +- Any feature that opens an external connection must close it before capture. + +For a lower-complexity option that retains a live process, see +[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain +different amounts of state and have different idle resource costs. diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py index 443dd82fdd05..fc8301069858 100644 --- a/tests/entrypoints/launchers/test_launch_cli.py +++ b/tests/entrypoints/launchers/test_launch_cli.py @@ -3,15 +3,48 @@ """Unit tests for the `vllm launch` CLI subcommand.""" import argparse -from unittest.mock import patch +import hashlib +import json +import os +import signal +import subprocess +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any +from unittest.mock import create_autospec, patch import pytest +import vllm.snapshot.runtime as snapshot_runtime +import vllm.snapshot.server as snapshot_server +from vllm.entrypoints.cli import snapshot as snapshot_cli from vllm.entrypoints.cli.launch import ( LaunchSubcommand, RenderSubcommand, cmd_init, ) +from vllm.snapshot.controller import ( + LocalSnapshotTools, + SnapshotCreateError, + SnapshotRestoreError, + _child_engine_argv, + create_snapshot, + restore_snapshot, +) +from vllm.snapshot.manifest import ( + SnapshotCompatibilityError, + SnapshotManifest, + SnapshotRuntimeIdentity, + SnapshotSecurityError, + read_manifest, + validate_artifact_root, + validate_identity, + write_manifest_atomic, +) +from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord +from vllm.snapshot.server import SnapshotCanaryError +from vllm.snapshot.types import Oracle, oracles_match from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -109,3 +142,1011 @@ def test_launch_registered_in_main(): assert hasattr(launch_module, "cmd_init") subcmds = launch_module.cmd_init() assert any(s.name == "launch" for s in subcmds) + + +_MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca" + + +def _run_cli(argv: list[str]) -> None: + from vllm.entrypoints.cli import main as cli_main + + with patch.object(sys, "argv", ["vllm", *argv]): + cli_main.main() + + +def test_snapshot_environment_contract( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +): + from vllm.entrypoints.serve.utils import api_utils + + monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail) + secret = "snapshot-secret" + monkeypatch.setenv("VLLM_API_KEY", secret) + monkeypatch.setenv("VLLM_USER_SETTING", secret) + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret) + environment = dict(LocalSnapshotTools()._environment_identity()) + assert ( + environment["VLLM_USER_SETTING"] + == hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest() + ) + assert ( + environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"] + ) + assert secret not in str(environment) + assert "VLLM_API_KEY" not in environment + + monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD") + implicit_worker_default = dict(LocalSnapshotTools()._environment_identity()) + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default + + monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed") + changed_environment = dict(LocalSnapshotTools()._environment_identity()) + assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"] + + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items()))) + before = os.environ.copy() + + _run_cli(["snapshot", "inspect", str(artifact)]) + + assert os.environ == before + output = capsys.readouterr().out + assert secret not in output + inspected = json.loads(output) + inspected_environment = dict(inspected["environment"]) + assert inspected_environment == environment + + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) + manifest = tools.make_manifest( + argparse.Namespace( + hf_token=secret, + model_tag="Qwen/Qwen3-0.6B", + revision=_MODEL_REVISION, + served_model_name=None, + tokenizer_revision=None, + ), + ( + "Qwen/Qwen3-0.6B", + "--api-key", + secret, + f"{secret}-alternate", + f"--api-key={secret}", + "--hf-token", + secret, + f"--hf-token={secret}", + ), + ProcessInventory(100, (100, 101), (101,), "GPU-abc"), + _oracle(), + tmp_path, + ) + + assert manifest.engine_argv == ( + "Qwen/Qwen3-0.6B", + "--api-key", + "***", + "***", + "--api-key=***", + "--hf-token", + "***", + "--hf-token=***", + ) + + +@pytest.mark.parametrize( + "case", + [ + (("--hf_token", "SECRET"), ("--hf_token", "***")), + (("--hf_token=SECRET",), ("--hf_token=***",)), + (("--api_key", "SECRET"), ("--api_key", "***")), + (("--api_key=SECRET",), ("--api_key=***",)), + (("--hf-tok", "SECRET"), ("--hf-tok", "***")), + (("--hf-tok=SECRET",), ("--hf-tok=***",)), + (("--hf_tok", "SECRET"), ("--hf_tok", "***")), + (("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")), + (("--api-k=SECRET",), ("--api-k=***",)), + (("--hf-overrides", "{}"), ("--hf-overrides", "{}")), + ], +) +def test_snapshot_manifest_redacts_accepted_secret_options(case): + engine_argv, expected = case + tools = create_autospec(LocalSnapshotTools, instance=True) + tools.current_identity.return_value = _runtime_identity() + tools._artifact_bytes.return_value = 0 + manifest = LocalSnapshotTools.make_manifest( + tools, + snapshot_server.parse_vllm_args(["model", *engine_argv]), + engine_argv, + ProcessInventory(100, (100, 101), (101,), "GPU-abc"), + _oracle(), + Path(), + ) + + assert manifest.engine_argv == expected + if "SECRET" in " ".join(engine_argv): + assert "***" in " ".join(manifest.engine_argv) + assert "SECRET" not in " ".join(manifest.engine_argv) + + +def parse_snapshot(*argv: str): + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers(dest="subparser", required=True) + snapshot_cli.SnapshotSubcommand( + create_requested=argv[:1] == ("create",) + ).subparser_init(subparsers) + return parser.parse_args(["snapshot", *argv]) + + +@pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"]) +@pytest.mark.parametrize("joined", [False, True]) +def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch): + directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"] + argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory] + args = parse_snapshot(*argv) + monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) + + child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) + + assert child.model == "Qwen/Qwen3-0.6B" + assert child.revision == _MODEL_REVISION + assert child.enable_sleep_mode + + +@pytest.mark.parametrize("model_in_cli", [False, True]) +def test_snapshot_create_config_preserves_cli_precedence( + tmp_path, monkeypatch, model_in_cli +): + config = tmp_path / "engine.yaml" + config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n") + model = ["Qwen/Qwen3-0.6B"] if model_in_cli else [] + argv = [ + "create", + *model, + "--snapshot-dir", + str(tmp_path / "snapshot"), + "--config", + str(config), + "--revision", + _MODEL_REVISION, + "--max-model-len", + "1024", + ] + args = parse_snapshot(*argv) + monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) + child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) + + assert ( + (args.model_tag or args.model) + == child.model + == ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model") + ) + assert args.max_model_len == child.max_model_len == 1024 + assert args.dtype == child.dtype == "float16" + + +@pytest.mark.parametrize( + ("probe", "status", "output"), + [ + (None, 0, ""), + ("nvidia-smi", 1, ""), + ("find", 1, ""), + ("ss", 1, ""), + ("nvidia-smi", 0, "12345"), + ("find", 0, "link_remap.fixture"), + ("ss", 0, "LISTEN"), + ], +) +def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output): + script = ( + Path(__file__).resolve().parents[3] + / ".buildkite/scripts/initialized-snapshot-e2e.sh" + ) + if not script.exists(): + pytest.skip("snapshot E2E script is not packaged in this test image") + source = script.read_text() + # Exercise the real probes without entering Docker/CRIU setup. + functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")] + shell = ( + """set -Eeuo pipefail +GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE='' +SNAPSHOT_PIDS=(999999999) +nvidia-smi() { :; } +find() { :; } +ss() { :; } +""" + + functions + ) + if probe: + shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n' + shell += "\nif state_clean; then exit 0; else exit 1; fi\n" + + result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10) + + assert result.returncode == (0 if probe is None else 1), result.stderr + + +def test_snapshot_create_cli_accepts_only_pinned_compact_mode( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux") + monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64") + base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot") + compact = parse_snapshot(*base, "--revision", _MODEL_REVISION) + + snapshot_cli.validate_create_args(compact) + assert compact.model_tag == "Qwen/Qwen3-0.6B" + + dispatched: list[argparse.Namespace] = [] + monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append) + _run_cli(["snapshot", *base, "--revision", _MODEL_REVISION]) + assert dispatched[0].revision == _MODEL_REVISION + + +@pytest.mark.parametrize( + ("attribute", "value", "message"), + [ + ("prefill_context_parallel_size", 2, "prefill context parallel size 1"), + ("model_tag", None, "immutable remote model"), + ("logprobs_mode", "raw_logits", "log-probability mode"), + ("speculative_config", {}, "does not support speculative decoding"), + ("spec_method", "ngram", "does not support speculative decoding"), + ("spec_model", "draft-model", "does not support speculative decoding"), + ("spec_tokens", 3, "does not support speculative decoding"), + ], +) +def test_snapshot_create_rejects_invalid_config( + attribute: str, + value: object, + message: str, + tmp_path: Path, +): + args = parse_snapshot( + "create", + "Qwen/Qwen3-0.6B", + "--snapshot-dir=/tmp/snapshot", + "--revision", + _MODEL_REVISION, + ) + if attribute == "model_tag": + local_model = tmp_path / "model" + local_model.mkdir() + value = str(local_model) + setattr(args, attribute, value) + with pytest.raises(ValueError, match=message): + snapshot_cli.validate_create_args(args) + + +def _oracle( + token_ids: tuple[int, ...] = (12095,), + text: str = " Paris", + logprob: float = -0.125, +) -> Oracle: + return Oracle(token_ids, text, logprob) + + +def _manifest(**changes: object) -> SnapshotManifest: + manifest = SnapshotManifest( + schema_version=1, + boundary="post-engine-init-reloadable-state-released", + created_at="2026-08-06T00:00:00Z", + artifact_bytes=8, + vllm_version="0.1.dev1+snapshot", + python_version="3.12.3", + torch_version="2.9.0", + cuda_runtime="12.9", + driver_version="575.57.08", + criu_version="4.1", + cuda_checkpoint_sha256="a" * 64, + kernel_release="6.8.0", + host_id="host-a", + gpu_name="NVIDIA A10", + gpu_uuid="GPU-abc", + model="Qwen/Qwen3-0.6B", + served_model_name="Qwen/Qwen3-0.6B", + model_revision="model-sha", + tokenizer_revision="tokenizer-sha", + engine_argv=("Qwen/Qwen3-0.6B",), + environment=(("VLLM_USE_V1", "1"),), + process_tree=(100, 101), + cuda_holders=(101,), + oracle_token_ids=(12095,), + oracle_text=" Paris", + oracle_sampled_token_logprob=-0.125, + ) + return SnapshotManifest.model_validate( + {**manifest.model_dump(mode="python"), **changes}, strict=True + ) + + +def _runtime_identity(**changes: object) -> Any: + fields = SnapshotRuntimeIdentity.model_fields.keys() + payload = _manifest().model_dump(mode="python", include=fields) + return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True) + + +def _fake_snapshot_tools() -> Any: + tools = create_autospec(LocalSnapshotTools, instance=True) + tools.launch_child.return_value = 100 + tools.wait_ready.return_value = _oracle() + tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc") + + tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity() + tools.restore.return_value = 100 + tools.request_oracle.return_value = _oracle() + return tools + + +@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"]) +def test_snapshot_timeout_must_be_positive_and_finite( + value: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value) + with pytest.raises(ValueError, match="positive finite number"): + LocalSnapshotTools() + + +def test_snapshot_create_abort_reaps_after_process_group_cleanup( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + tools = LocalSnapshotTools() + process = subprocess.Popen( + [sys.executable, "-c", "raise SystemExit(1)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + tools._children[process.pid] = process + (tmp_path / "child.log").write_text("root failed") + killed: list[tuple[int, signal.Signals]] = [] + monkeypatch.setattr( + os, + "killpg", + lambda pid, sig: killed.append((pid, sig)), + ) + monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False) + + try: + with pytest.raises(SnapshotCreateError, match="root failed"): + tools.wait_ready(tmp_path, process.pid) + assert process.returncode is None + tools.abort_create(process.pid) + finally: + if process.pid in tools._children: + tools.abort_create(process.pid) + + assert killed == [(process.pid, signal.SIGKILL)] + assert process.pid not in tools._children + + +def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None: + """`wait_ready` binds Linux-only `os.waitid` before it reads the marker.""" + monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False) + + +def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + _supply_waitid(monkeypatch) + snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle()) + + assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle() + + +@pytest.mark.parametrize( + "payload", + [ + {"token_ids": [12095], "sampled_token_logprob": -0.125}, + {"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125}, + ], +) +def test_snapshot_wait_ready_rejects_a_malformed_marker( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object] +): + _supply_waitid(monkeypatch) + (tmp_path / "ready.json").write_text(json.dumps(payload)) + + with pytest.raises(SnapshotCreateError, match="ready marker invalid"): + LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) + + +def test_snapshot_create_abort_rejects_surviving_process_group( + monkeypatch: pytest.MonkeyPatch, +): + tools = LocalSnapshotTools() + process = create_autospec(subprocess.Popen, instance=True) + process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10) + tools._children[100] = process + monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None) + + with pytest.raises(SnapshotCreateError, match="survived SIGKILL"): + tools.abort_create(100) + + +@pytest.mark.asyncio +async def test_snapshot_child_rejects_mismatch_before_publication( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from vllm.entrypoints.launchers.api_server import entry as api_server + + canaries = iter((_oracle(), _oracle((42,), " Lyon"))) + + @asynccontextmanager + async def engine_context(): + yield object() + + async def next_canary(_engine: object) -> Oracle: + return next(canaries) + + async def return_none(_engine: object) -> None: + return None + + monkeypatch.setattr( + api_server, "build_async_engine_client", lambda _args: engine_context() + ) + monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary) + monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none) + monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none) + monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None) + monkeypatch.setattr( + snapshot_server, + "write_ready_atomic", + lambda *_args: pytest.fail("mismatched snapshot was published"), + ) + control = snapshot_server.ControlArgs( + ready_file=tmp_path / "ready.json", + release_file=tmp_path / "release.json", + release_timeout_s=1, + ) + with pytest.raises(SnapshotCanaryError, match="rehearsal changed"): + await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace()) + + +def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path): + marker = tmp_path / "release.json" + marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1})) + + listener = snapshot_server.read_release_marker(marker) + + assert listener == snapshot_server.ListenerConfig(host=None, port=8000) + + +@pytest.mark.parametrize( + "payload", + [ + ["release"], + {"port": 8000}, + {"release": "true", "port": 8000}, + {"release": 1, "port": 8000}, + {"release": True, "port": 0}, + {"release": True, "port": 70000}, + {"release": True, "port": True}, + {"release": True, "port": "8000"}, + {"release": True, "port": 8000, "host": 5}, + ], +) +def test_snapshot_release_marker_rejects_an_invalid_payload( + tmp_path: Path, payload: object +): + marker = tmp_path / "release.json" + marker.write_text(json.dumps(payload)) + + with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"): + snapshot_server.read_release_marker(marker) + + +def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path): + target = tmp_path / "snapshot" + tools = _fake_snapshot_tools() + tools.dump.side_effect = RuntimeError("dump failed") + with pytest.raises(RuntimeError, match="dump failed"): + create_snapshot( + argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"), + tools=tools, + ) + tools.abort_create.assert_called_once_with(100) + assert not target.exists() + + +def _restore_fixture(tmp_path: Path) -> argparse.Namespace: + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + write_manifest_atomic(artifact, _manifest()) + return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000) + + +def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path): + args = _restore_fixture(tmp_path) + + identity_mismatch = _fake_snapshot_tools() + identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: ( + _runtime_identity(vllm_version="different") + ) + with pytest.raises(SnapshotCompatibilityError, match="vllm_version"): + restore_snapshot(args, tools=identity_mismatch) + identity_mismatch.restore.assert_not_called() + + +def test_snapshot_identity_mismatch_names_every_differing_field(): + with pytest.raises(SnapshotCompatibilityError) as mismatch: + validate_identity( + _manifest(), + _runtime_identity(vllm_version="different", gpu_uuid="GPU-other"), + ) + + assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid" + + +def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path): + args = _restore_fixture(tmp_path) + oracle_mismatch = _fake_snapshot_tools() + oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5) + with pytest.raises(SnapshotRestoreError, match="oracle mismatch"): + restore_snapshot(args, tools=oracle_mismatch) + oracle_mismatch.cleanup.assert_called_once() + oracle_mismatch.complete_restore.assert_not_called() + assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005)) + assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002)) + + +@pytest.mark.parametrize( + ("invalid_update", "diagnostic"), + [ + ({"unexpected": True}, "unexpected"), + ({"schema_version": True}, "schema_version"), + ({"boundary": "unknown"}, "boundary"), + ({"process_tree": [100, 100]}, "root"), + ({"cuda_holders": [999]}, "root"), + ({"oracle_token_ids": [1, 2]}, "oracle_token_ids"), + ({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"), + ( + {"oracle_sampled_token_logprob": float("nan")}, + "oracle_sampled_token_logprob", + ), + ], +) +def test_snapshot_manifest_validation( + tmp_path: Path, + invalid_update: dict[str, object], + diagnostic: str, +): + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + write_manifest_atomic(artifact, _manifest()) + invalid = _manifest().model_dump(mode="json") + invalid.update(invalid_update) + (artifact / "manifest.json").write_text(json.dumps(invalid)) + with pytest.raises( + SnapshotCompatibilityError, + match=f"invalid snapshot manifest: {diagnostic}", + ): + read_manifest(artifact) + + +def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path: + artifact = tmp_path / name + artifact.mkdir(mode=0o700) + (artifact / "child.log").write_bytes(b"startup log") + (artifact / "child.log.snapshot-size").write_text("11\n") + return artifact + + +def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]: + probed_pids: list[int] = [] + + def missing(pid: int, _signal: int) -> None: + probed_pids.append(pid) + raise ProcessLookupError + + monkeypatch.setattr(snapshot_runtime.os, "kill", missing) + return probed_pids + + +@pytest.mark.parametrize( + ("probe_error", "diagnostic"), + [ + (None, "already occupied: 100"), + (PermissionError(), "already occupied: 100"), + (OSError("probe failed"), "availability probe failed: 100"), + ], +) +def test_snapshot_restore_rejects_unavailable_pid_before_mutation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + probe_error: OSError | None, + diagnostic: str, +): + args = _restore_fixture(tmp_path) + artifact = Path(args.snapshot_dir) + (artifact / "child.log").write_bytes(b"startup log") + (artifact / "release.json").write_bytes(b"release sentinel") + before = {path.name: path.read_bytes() for path in artifact.iterdir()} + tools = LocalSnapshotTools() + + def fail_probe(_pid: int, _signal: int) -> None: + if probe_error is not None: + raise probe_error + + monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe) + monkeypatch.setattr(tools, "preflight", lambda *_args: None) + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + + with pytest.raises(SnapshotRestoreError, match=diagnostic): + restore_snapshot(args, tools=tools) + assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before + + +@pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"]) +def test_snapshot_restore_terminates_only_a_pinned_tree( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mode: str, +): + _mark_snapshot_pids_free(monkeypatch) + artifact = _local_restore_artifact(tmp_path) + manifest = _manifest() + root_pid = manifest.process_tree[0] + tools = LocalSnapshotTools() + terminated: list[int] = [] + + def terminate_and_wait(pid: int) -> None: + terminated.append(pid) + if mode == "cleanup-fails": + raise SnapshotRestoreError("restored process cleanup is incomplete") + tools._restored_processes.pop(pid) + + def pin( + _artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...] + ) -> None: + if mode == "unpinned": + raise SnapshotRestoreError( + "restored session does not match the captured process tree" + ) + tools._restored_processes[process_tree[0]] = () + + monkeypatch.setattr(tools, "cleanup", terminate_and_wait) + monkeypatch.setattr(tools, "_pin_restored_tree", pin) + monkeypatch.setattr(tools, "_criu", lambda *_args: None) + monkeypatch.setattr( + tools, + "_run", + lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""), + ) + + with pytest.raises(SnapshotRestoreError, match="does not match") as failure: + tools.restore(artifact, manifest) + + # A pinned tree is ours: it must be terminated and waited for, not left to + # honour the abort marker. An unpinned PID is never signalled. + assert terminated == ([] if mode == "unpinned" else [root_pid]) + if mode == "cleanup-fails": + # The failure to terminate is reported beside the primary error, and a + # tree that may still be running stays pinned. + assert "restore cleanup failed" in str(failure.value) + assert tools._restored_processes == {root_pid: ()} + else: + assert tools._restored_processes == {} + assert json.loads((artifact / "release.json").read_text())["release"] is False + + +def test_snapshot_private_path_and_link_remap_security( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + _mark_snapshot_pids_free(monkeypatch) + real = tmp_path / "real" + real.mkdir(mode=0o700) + linked = tmp_path / "linked" + linked.symlink_to(real, target_is_directory=True) + with pytest.raises(SnapshotSecurityError, match="symlink"): + validate_artifact_root(linked, creating=False) + + unsafe = tmp_path / "unsafe" + unsafe.mkdir() + unsafe.chmod(0o777) + with pytest.raises(SnapshotSecurityError, match="world-writable"): + validate_artifact_root(unsafe / "snapshot", creating=True) + + artifact = _local_restore_artifact(tmp_path) + saved_remaps = artifact / "link-remaps" + saved_remaps.mkdir() + (saved_remaps / "link_remap.270").write_bytes(b"semaphore state") + shm_dir = tmp_path / "shm" + shm_dir.mkdir() + target = shm_dir / "link_remap.270" + replacement = tmp_path / "replacement" + replacement.write_bytes(b"semaphore state") + target.symlink_to(replacement) + tools = LocalSnapshotTools() + tools.shm_dir = shm_dir + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"): + tools.restore(artifact, _manifest()) + + +@pytest.mark.parametrize( + ("blocked_state", "message"), + [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")], +) +def test_snapshot_rejects_unsafe_process_state_before_criu( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str +): + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101)) + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) + monkeypatch.setattr( + tools, + "_descriptor_inventory", + lambda _tree: ( + (101,) if blocked_state == "io_uring" else (), + {41: 101}, + ), + ) + tcp_record = _TcpSocketRecord( + family="AF_INET", + local_raw="3C00000A:D6F6", + remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000", + inode=41, + ) + monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,)) + with pytest.raises(SnapshotCreateError, match=message) as excinfo: + tools.inventory(100, tmp_path) + if blocked_state == "tcp": + assert "pid 101" in str(excinfo.value) + assert "1.1.1.1:443" in str(excinfo.value) + + +def test_decode_endpoint_families(): + assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030" + assert ( + _decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB") + == "[::1]:443" + ) + + +def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch): + tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6" + contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n" + tcp.write_text(contents) + tables = (("AF_INET", tcp), ("AF_INET6", tcp6)) + monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables) + + tools = LocalSnapshotTools() + assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",) + + tcp6.write_text(contents) + families = tuple(record.family for record in tools._tcp_records()) + assert families == ("AF_INET", "AF_INET6") + + +def test_snapshot_manifest_records_external_cache_files(tmp_path: Path): + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + recorded = ( + (str(tmp_path / "kernel.so"), "b" * 64), + (str(tmp_path / "launcher.lock"), "locked"), + ) + write_manifest_atomic(artifact, _manifest(external_cache_files=recorded)) + + assert read_manifest(artifact).external_cache_files == recorded + + # A manifest written before this field existed still loads. + legacy = _manifest().model_dump(mode="json") + del legacy["external_cache_files"] + (artifact / "manifest.json").write_text(json.dumps(legacy)) + + assert read_manifest(artifact).external_cache_files == () + + +def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Keep the resolver off whatever cache roots the host running this has.""" + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache")) + monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default")) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache")) + + +def _unreadable(_path: Path) -> str: + raise OSError("cache file is unreadable") + + +def test_snapshot_inventory_records_open_generated_cache_files( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + cache_root = tmp_path / "triton_cache" + cache_root.mkdir() + kernel = cache_root / "kernel.so" + kernel.write_bytes(b"compiled kernel") + lock = cache_root / "launcher.lock" + lock.write_bytes(b"held by pid 100") + weights = tmp_path / "model.safetensors" + weights.write_bytes(b"weights") + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + (artifact / "child.log.lock").write_bytes(b"") + _pin_generated_cache_roots(monkeypatch, tmp_path) + monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) + + targets = { + 100: (str(kernel), str(weights), "socket:[41]"), + 101: (str(lock), str(cache_root), str(artifact / "child.log.lock")), + 102: (f"{cache_root / 'evicted.so'} (deleted)",), + } + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102)) + monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid]) + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) + monkeypatch.setattr(tools, "_tcp_records", lambda: ()) + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) + + manifest = tools.make_manifest( + argparse.Namespace( + model_tag="Qwen/Qwen3-0.6B", + revision=_MODEL_REVISION, + served_model_name=None, + tokenizer_revision=None, + ), + ("Qwen/Qwen3-0.6B",), + tools.inventory(100, artifact), + _oracle(), + artifact, + ) + + # Weights, sockets, the cache directory itself, files inside the artifact, + # and deleted targets that CRIU remaps are all left out. + assert manifest.external_cache_files == ( + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), + (str(lock), "locked"), + ) + + +def test_snapshot_inventory_marks_a_cache_file_it_cannot_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + cache_root = tmp_path / "triton_cache" + cache_root.mkdir() + kernel = cache_root / "kernel.so" + kernel.write_bytes(b"compiled kernel") + lock = cache_root / "launcher.lock" + lock.write_bytes(b"held by pid 100") + artifact = tmp_path / "snapshot" + artifact.mkdir(mode=0o700) + _pin_generated_cache_roots(monkeypatch, tmp_path) + monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) + + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) + monkeypatch.setattr( + tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock)) + ) + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) + monkeypatch.setattr(tools, "_tcp_records", lambda: ()) + monkeypatch.setattr(tools, "_sha256", _unreadable) + + # A file that cannot be digested keeps its own marker, so the manifest says + # which entries restore can check for existence only. + assert tools.inventory(100, artifact).external_cache_files == ( + (str(kernel), "unreadable"), + (str(lock), "locked"), + ) + + +def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + # Descriptor targets come back kernel-canonical, so the fixture must be too. + base = Path(os.path.realpath(tmp_path)) + real_cache = base / "mnt_cache" + real_cache.mkdir() + kernel = real_cache / "kernel.so" + kernel.write_bytes(b"compiled kernel") + linked_cache = base / "linked_cache" + linked_cache.symlink_to(real_cache, target_is_directory=True) + sibling = base / "vllm_cache-backup" + sibling.mkdir() + stale = sibling / "stale.so" + stale.write_bytes(b"stale kernel") + artifact = base / "snapshot" + artifact.mkdir(mode=0o700) + + _pin_generated_cache_roots(monkeypatch, base) + monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache)) + + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) + monkeypatch.setattr( + tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale)) + ) + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) + monkeypatch.setattr(tools, "_tcp_records", lambda: ()) + + # A cache root reached through a symlink still covers the file the kernel + # reports by its real path, and a directory that merely shares the root's + # name prefix is not part of the cache at all. + assert tools.inventory(100, artifact).external_cache_files == ( + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), + ) + + +def test_snapshot_restore_rejects_rotated_cache_file_before_mutation( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + _mark_snapshot_pids_free(monkeypatch) + kernel = tmp_path / "kernel.so" + kernel.write_bytes(b"compiled kernel") + lock = tmp_path / "launcher.lock" + lock.write_bytes(b"held by pid 100") + manifest = _manifest( + external_cache_files=( + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), + (str(lock), "locked"), + ) + ) + artifact = _local_restore_artifact(tmp_path) + before = {path.name: path.read_bytes() for path in artifact.iterdir()} + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + + kernel.write_bytes(b"recompiled kernel") + with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed: + tools.restore(artifact, manifest) + assert str(kernel) in str(changed.value) + + # Restoring the digest gets past the first entry, so the lock file is only + # checked for existence: the second failure proves the first one passed. + kernel.write_bytes(b"compiled kernel") + lock.unlink() + with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone: + tools.restore(artifact, manifest) + assert str(lock) in str(gone.value) + + assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before + + +def test_snapshot_restore_rejects_a_cache_file_it_cannot_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + _mark_snapshot_pids_free(monkeypatch) + kernel = tmp_path / "kernel.so" + kernel.write_bytes(b"compiled kernel") + manifest = _manifest( + external_cache_files=( + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), + ) + ) + artifact = _local_restore_artifact(tmp_path) + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + monkeypatch.setattr(tools, "_sha256", _unreadable) + + # A recorded digest that cannot be read again is a failed check, not a file + # to wave through on existence. + with pytest.raises( + SnapshotRestoreError, match="unreadable since capture" + ) as blocked: + tools.restore(artifact, manifest) + assert str(kernel) in str(blocked.value) + + +def test_snapshot_runtime_installer_logs_the_failing_exit_code(): + installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh" + if not installer.is_file(): + pytest.skip("installer script is not shipped in this test image") + timing = [ + line + for line in installer.read_text().splitlines() + if line.startswith(("_snapshot_install_started=", "trap ")) + ] + assert len(timing) == 2 + + result = subprocess.run( + ["bash", "-c", "\n".join([*timing, "exit 7"])], + capture_output=True, + text=True, + ) + + assert result.returncode == 7 + assert "(exit 7)" in result.stdout diff --git a/tests/models/quantization/test_mimo_v2_qkv_shard.py b/tests/models/quantization/test_mimo_v2_qkv_shard.py index 91eb036bc7c5..68afe7fdecfc 100644 --- a/tests/models/quantization/test_mimo_v2_qkv_shard.py +++ b/tests/models/quantization/test_mimo_v2_qkv_shard.py @@ -19,6 +19,7 @@ from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import set_default_torch_num_threads pytestmark = pytest.mark.cpu_test @@ -39,6 +40,13 @@ } +@pytest.fixture(scope="module", autouse=True) +def single_threaded_cpu_ops(): + # Do not prime the parent's thread pool before later engine forks. + with set_default_torch_num_threads(1): + yield + + def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): return ( num_heads // ckpt_tp * head_dim diff --git a/tools/install_snapshot_runtime.sh b/tools/install_snapshot_runtime.sh new file mode 100755 index 000000000000..40d6ad944fee --- /dev/null +++ b/tools/install_snapshot_runtime.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euxo pipefail + +# Surface the isolated snapshot-runtime install cost in the CI build log. The +# trap captures the status first: the date substitution below would reset it. +_snapshot_install_started=$(date +%s) +code=0 +trap 'code=$?; echo "[TIMING] snapshot runtime install: $(( $(date +%s) - _snapshot_install_started )) s (exit $code)"' EXIT + +readonly CRIU_VERSION="4.2.1" +readonly CRIU_SOURCE_SHA256="feffdf4638125ebb12d2434754f80a1d7bbba85a3e6bee98c216f88fb99a5d96" +readonly DYNAMO_CHECKPOINT_HELPER_COMMIT="fdf25efad60f696a73393caf341c527a93d84190" +readonly DYNAMO_CHECKPOINT_HELPER_SHA256="7b8ea21baecb1da729202c32e0fecbbbafb789d65e5c1c0649ea253ad09b8f82" +readonly DYNAMO_LICENSE_SHA256="aeeca3d74a13b91d17ab9211c407b8d94a77bf8d9e5750c467818d3066b8adcf" + +# Keep the compiler and NL headers required by runtime JIT and RDMA. +snapshot_shared_build_deps=( + build-essential + libnl-3-dev + libnl-route-3-dev +) +snapshot_build_deps=( + pkg-config + libbsd-dev + libcap-dev + libnet1-dev + libprotobuf-dev + libprotobuf-c-dev + protobuf-c-compiler + protobuf-compiler + python3-protobuf + libgnutls28-dev + libnftables-dev + uuid-dev +) +snapshot_runtime_deps=( + libbsd0 + libcap2 + libnet1 + libnl-3-200 + libnl-route-3-200 + libprotobuf-c1 + libgnutls30t64 + libnftables1 + iproute2 + iptables + procps + uuid-runtime + util-linux +) + +apt-get update -y +apt-get install -y --no-install-recommends \ + "${snapshot_shared_build_deps[@]}" \ + "${snapshot_build_deps[@]}" \ + "${snapshot_runtime_deps[@]}" + +snapshot_build_dir="$(mktemp -d)" +curl -fsSL \ + "https://github.com/checkpoint-restore/criu/archive/refs/tags/v${CRIU_VERSION}.tar.gz" \ + -o "${snapshot_build_dir}/criu.tar.gz" +echo "${CRIU_SOURCE_SHA256} ${snapshot_build_dir}/criu.tar.gz" | sha256sum -c - +mkdir "${snapshot_build_dir}/criu" +tar -xzf "${snapshot_build_dir}/criu.tar.gz" \ + --strip-components=1 -C "${snapshot_build_dir}/criu" +make -C "${snapshot_build_dir}/criu" -j"${MAX_JOBS:-$(nproc)}" criu cuda_plugin +install -D -m 0755 "${snapshot_build_dir}/criu/criu/criu" \ + /usr/local/sbin/criu +install -D -m 0755 \ + "${snapshot_build_dir}/criu/plugins/cuda/cuda_plugin.so" \ + /usr/local/lib/criu/cuda_plugin.so + +helper_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/deploy/snapshot/cmd/cuda-checkpoint-helper/main.c" +curl -fsSL "${helper_url}" -o "${snapshot_build_dir}/cuda-checkpoint.c" +echo "${DYNAMO_CHECKPOINT_HELPER_SHA256} ${snapshot_build_dir}/cuda-checkpoint.c" \ + | sha256sum -c - +gcc -O2 -Wall -Wextra \ + -I/usr/local/cuda/include \ + -L/usr/local/cuda/lib64/stubs \ + "${snapshot_build_dir}/cuda-checkpoint.c" -lcuda \ + -o /usr/local/sbin/cuda-checkpoint + +license_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/LICENSE" +curl -fsSL "${license_url}" -o "${snapshot_build_dir}/DYNAMO-LICENSE" +echo "${DYNAMO_LICENSE_SHA256} ${snapshot_build_dir}/DYNAMO-LICENSE" \ + | sha256sum -c - + +install -D -m 0644 "${snapshot_build_dir}/criu.tar.gz" \ + /usr/share/doc/criu/source/criu.tar.gz +install -D -m 0644 "${snapshot_build_dir}/criu/COPYING" \ + /usr/share/doc/criu/COPYING +install -D -m 0644 "${snapshot_build_dir}/cuda-checkpoint.c" \ + /usr/share/doc/cuda-checkpoint-helper/main.c +install -D -m 0644 "${snapshot_build_dir}/DYNAMO-LICENSE" \ + /usr/share/doc/cuda-checkpoint-helper/LICENSE + +criu --version +ln -s /usr/local/cuda/lib64/stubs/libcuda.so \ + "${snapshot_build_dir}/libcuda.so.1" +LD_LIBRARY_PATH="${snapshot_build_dir}" \ + cuda-checkpoint --help | grep -q -- "--action" +test -x /usr/local/lib/criu/cuda_plugin.so + +apt-get purge -y --auto-remove "${snapshot_build_deps[@]}" +rm -rf "${snapshot_build_dir}" +rm -rf /var/lib/apt/lists/* diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 4ac960d8cb93..ab9afcb008d9 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -46,6 +46,7 @@ def main(): import vllm.entrypoints.cli.openai import vllm.entrypoints.cli.run_batch import vllm.entrypoints.cli.serve + import vllm.entrypoints.cli.snapshot from vllm.entrypoints.serve.utils.api_utils import ( VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup, @@ -59,9 +60,11 @@ def main(): vllm.entrypoints.cli.benchmark.main, vllm.entrypoints.cli.collect_env, vllm.entrypoints.cli.run_batch, + vllm.entrypoints.cli.snapshot, ] - cli_env_setup() + if sys.argv[1:2] != ["snapshot"]: + cli_env_setup() vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench() @@ -94,7 +97,12 @@ def main(): subparsers = parser.add_subparsers(required=False, dest="subparser") cmds = {} for cmd_module in CMD_MODULES: - new_cmds = cmd_module.cmd_init() + if cmd_module is vllm.entrypoints.cli.snapshot: + new_cmds = cmd_module.cmd_init( + create_requested=sys.argv[1:3] == ["snapshot", "create"] + ) + else: + new_cmds = cmd_module.cmd_init() for cmd in new_cmds: cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) cmds[cmd.name] = cmd diff --git a/vllm/entrypoints/cli/snapshot.py b/vllm/entrypoints/cli/snapshot.py new file mode 100644 index 000000000000..f5201675d40c --- /dev/null +++ b/vllm/entrypoints/cli/snapshot.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import os +import platform +import typing +from pathlib import Path + +from vllm.entrypoints.cli.types import CLISubcommand + +if typing.TYPE_CHECKING: + from vllm.utils.argparse_utils import FlexibleArgumentParser +else: + FlexibleArgumentParser = argparse.ArgumentParser + + +def validate_create_args(args: argparse.Namespace) -> None: + model = str(getattr(args, "model_tag", None) or args.model) + revision = str(getattr(args, "revision", None) or "") + if ( + Path(model).exists() + or len(revision) != 40 + or any(character not in "0123456789abcdefABCDEF" for character in revision) + ): + raise ValueError( + "snapshot create requires an immutable remote model and " + "40-character --revision" + ) + parallel_sizes = ( + ("tensor parallel", args.tensor_parallel_size), + ("pipeline parallel", args.pipeline_parallel_size), + ("data parallel", args.data_parallel_size), + ("prefill context parallel", args.prefill_context_parallel_size), + ) + for name, size in parallel_sizes: + if size != 1: + raise ValueError(f"snapshot create currently requires {name} size 1") + unsupported_frontend = ( + args.api_server_count not in (None, 1) + or args.grpc + or args.headless + or args.api_key + or os.environ.get("VLLM_API_KEY") + or args.uds + or args.middleware + ) + if unsupported_frontend or any( + ( + args.ssl_keyfile, + args.ssl_certfile, + args.ssl_ca_certs, + args.enable_ssl_refresh, + args.ssl_ciphers, + ) + ): + raise ValueError( + "snapshot create requires one unauthenticated plaintext HTTP server" + ) + if args.logprobs_mode in {"raw_logits", "processed_logits"}: + raise ValueError("snapshot canary requires a log-probability mode") + if any( + getattr(args, name, None) is not None + for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens") + ): + raise ValueError("snapshot create does not support speculative decoding") + if platform.system() != "Linux" or platform.machine() != "x86_64": + raise ValueError("snapshot create requires Linux x86_64") + + +def run_create(args: argparse.Namespace) -> None: + validate_create_args(args) + from vllm.snapshot.controller import create_snapshot + + create_snapshot(args) + + +def run_inspect(args: argparse.Namespace) -> None: + from vllm.snapshot.manifest import inspect_snapshot + + print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2)) + + +def run_restore(args: argparse.Namespace) -> None: + from vllm.snapshot.controller import restore_snapshot + + restore_snapshot(args) + + +class SnapshotSubcommand(CLISubcommand): + """The `snapshot` subcommand for the vLLM CLI.""" + + name = "snapshot" + + def __init__(self, *, create_requested: bool = False) -> None: + self.create_requested = create_requested + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + args.snapshot_dispatch(args) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + self.name, + help="Create, inspect, or restore an initialized vLLM snapshot.", + usage="vllm snapshot [options]", + ) + actions = parser.add_subparsers(required=True, dest="snapshot_action") + + create_parser = actions.add_parser( + "create", help="Create a snapshot from an initialized TP1 engine." + ) + if self.create_requested: + from vllm.entrypoints.launchers.cli_args import make_arg_parser + + create_parser = make_arg_parser(create_parser) + else: + create_parser.add_argument("model_tag", nargs="?") + create_parser.add_argument("--snapshot-dir", required=True) + create_parser.set_defaults(snapshot_dispatch=run_create) + + inspect_parser = actions.add_parser( + "inspect", help="Inspect a snapshot without restoring it." + ) + inspect_parser.add_argument("snapshot_dir") + inspect_parser.set_defaults(snapshot_dispatch=run_inspect) + + restore_parser = actions.add_parser( + "restore", help="Restore a same-host TP1 snapshot." + ) + restore_parser.add_argument("snapshot_dir") + restore_parser.add_argument("--host", default=None) + restore_parser.add_argument("--port", type=int, default=8000) + restore_parser.set_defaults(snapshot_dispatch=run_restore) + + return parser + + +def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]: + return [SnapshotSubcommand(create_requested=create_requested)] diff --git a/vllm/snapshot/__init__.py b/vllm/snapshot/__init__.py new file mode 100644 index 000000000000..8c4fde128286 --- /dev/null +++ b/vllm/snapshot/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Create and restore explicitly compatible vLLM engine snapshots.""" diff --git a/vllm/snapshot/controller.py b/vllm/snapshot/controller.py new file mode 100644 index 000000000000..d5564fb8632d --- /dev/null +++ b/vllm/snapshot/controller.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Create and restore transaction orchestration for local snapshots.""" + +import argparse +import shutil +import sys +from pathlib import Path + +from vllm.snapshot.manifest import ( + read_manifest, + validate_artifact_root, + validate_identity, +) +from vllm.snapshot.runtime import ( + LocalSnapshotTools, + SnapshotCreateError, + SnapshotRestoreError, + _error_detail, +) +from vllm.snapshot.types import Oracle, oracles_match + + +def _child_engine_argv( + args: argparse.Namespace, engine_argv: tuple[str, ...] | None +) -> tuple[str, ...]: + """Build the child argv from the first argv source that is present. + + In order: an explicit ``engine_argv`` from the caller, the process argv + tail when invoked as ``vllm snapshot create ...``, then the bare model from + the parsed args. Whichever source wins, snapshot options are stripped and + sleep mode is appended. + """ + if not engine_argv: + process_argv = tuple(sys.argv[1:]) + if process_argv[:2] == ("snapshot", "create"): + engine_argv = process_argv[2:] + else: + model = getattr(args, "model_tag", None) or getattr(args, "model", None) + if not model: + raise SnapshotCreateError("snapshot create requires a model argument") + engine_argv = (str(model),) + child_argv: list[str] = [] + iterator = iter(engine_argv) + for item in iterator: + option, separator, _ = item.partition("=") + option = option.replace("_", "-") + if len(option) > 2 and "--snapshot-dir".startswith(option): + if not separator: + next(iterator, None) + else: + child_argv.append(item) + if "--enable-sleep-mode" not in child_argv: + child_argv.append("--enable-sleep-mode") + return tuple(child_argv) + + +def create_snapshot( + args: argparse.Namespace, + *, + engine_argv: tuple[str, ...] | None = None, + tools: LocalSnapshotTools | None = None, +) -> None: + target = Path(args.snapshot_dir).absolute() + toolset = tools or LocalSnapshotTools() + toolset.preflight("create", target) + if target.exists() or target.is_symlink(): + raise SnapshotCreateError(f"snapshot target already exists: {target}") + validate_artifact_root(target, creating=True) + target.parent.mkdir(mode=0o700, parents=False, exist_ok=True) + validate_artifact_root(target.parent, creating=False) + target.mkdir(mode=0o700) + validate_artifact_root(target, creating=False) + published = False + root_pid: int | None = None + try: + child_argv = _child_engine_argv(args, engine_argv) + root_pid = toolset.launch_child(target, child_argv) + oracle = toolset.wait_ready(target, root_pid) + inventory = toolset.inventory(root_pid, target) + toolset.dump(target, inventory) + toolset.verify_dead(inventory) + manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target) + toolset.publish(target, manifest) + published = True + except BaseException: + if root_pid is not None: + toolset.abort_create(root_pid) + raise + finally: + # Runs even when abort_create raised, so no partial artifact is left. + if not published and target.exists(): + shutil.rmtree(target) + + +def restore_snapshot( + args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None +) -> None: + artifact = Path(args.snapshot_dir).absolute() + toolset = tools or LocalSnapshotTools() + toolset.preflight("restore", artifact) + manifest = read_manifest(artifact) + validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid)) + + root_pid: int | None = None + try: + root_pid = toolset.restore(artifact, manifest) + toolset.release(artifact, args.host, args.port) + toolset.wait_listener(root_pid, args.host, args.port) + actual = toolset.request_oracle(args.host, args.port, manifest) + expected = Oracle( + token_ids=manifest.oracle_token_ids, + text=manifest.oracle_text, + sampled_token_logprob=manifest.oracle_sampled_token_logprob, + ) + if not oracles_match(expected, actual): + raise SnapshotRestoreError( + f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}" + ) + toolset.complete_restore(root_pid) + except BaseException as error: + if root_pid is not None: + try: + toolset.cleanup(root_pid) + except BaseException as cleanup_error: + raise SnapshotRestoreError( + f"snapshot restore failed: {_error_detail(error)}; " + f"cleanup failed: {_error_detail(cleanup_error)}" + ) from error + if isinstance(error, SnapshotRestoreError): + raise + raise SnapshotRestoreError(str(error)) from error diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py new file mode 100644 index 000000000000..a029c544c52e --- /dev/null +++ b/vllm/snapshot/manifest.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import math +import os +import stat +from pathlib import Path +from typing import Annotated, Any, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + StrictFloat, + StrictInt, + StrictStr, + ValidationError, + field_validator, + model_validator, +) +from typing_extensions import Self + + +class SnapshotCompatibilityError(RuntimeError): + """The snapshot identity does not match the requested runtime.""" + + +class SnapshotSecurityError(RuntimeError): + """The snapshot artifact does not meet the private-path contract.""" + + +_MAX_PID = 2**31 - 1 +_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)] +_OracleTokenId = Annotated[int, Field(ge=0)] + + +class SnapshotRuntimeIdentity(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + vllm_version: str + python_version: str + torch_version: str + cuda_runtime: str + driver_version: str + criu_version: str + cuda_checkpoint_sha256: str + kernel_release: str + host_id: str + gpu_name: str + gpu_uuid: str + environment: tuple[tuple[str, str], ...] + + +class SnapshotManifest(SnapshotRuntimeIdentity): + schema_version: Literal[1] + boundary: Literal["post-engine-init-reloadable-state-released"] + created_at: str + artifact_bytes: Annotated[int, Field(ge=0)] + model: str + served_model_name: str + model_revision: str + tokenizer_revision: str + engine_argv: tuple[str, ...] + process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] + cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] + external_cache_files: tuple[tuple[str, str], ...] = () + oracle_token_ids: Annotated[ + tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1) + ] + oracle_text: str + oracle_sampled_token_logprob: float + + @field_validator("schema_version", mode="before") + @classmethod + def _require_schema_version_one(cls, value: Any) -> Any: + if type(value) is not int or value != 1: + raise ValueError("schema version must be integer 1") + return value + + @field_validator("oracle_sampled_token_logprob", mode="before") + @classmethod + def _require_finite_float(cls, value: Any) -> Any: + if type(value) is not float or not math.isfinite(value): + raise ValueError("sampled token log probability must be a finite float") + return value + + @model_validator(mode="after") + def _validate_process_relationships(self) -> Self: + if len(set(self.process_tree)) != len(self.process_tree): + raise ValueError("process_tree contains duplicate PIDs") + if len(set(self.cuda_holders)) != len(self.cuda_holders): + raise ValueError("cuda_holders contains duplicate PIDs") + if not set(self.cuda_holders).issubset(self.process_tree): + raise ValueError("cuda_holders must be a subset of process_tree") + return self + + +class ReleaseMarker(BaseModel): + """Release marker the controller writes; unknown keys are ignored.""" + + release: Literal[True] + host: StrictStr | None = None + port: Annotated[StrictInt, Field(ge=1, le=65535)] + + @field_validator("release", mode="before") + @classmethod + def _require_json_true(cls, value: Any) -> Any: + # A literal accepts 1 as True; the barrier requires JSON true itself. + if value is not True: + raise ValueError("release must be JSON true") + return value + + +class ReadyMarker(BaseModel): + """Ready marker the child writes; unknown keys are ignored.""" + + token_ids: tuple[StrictInt, ...] + text: StrictStr + sampled_token_logprob: StrictFloat | StrictInt + + +def validate_identity( + expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity +) -> None: + """Require an exact match for every field that can affect compatibility.""" + # ``expected`` may be a manifest, so compare the identity fields only. + fields = set(SnapshotRuntimeIdentity.model_fields) + expected_values = expected.model_dump(include=fields) + actual_values = actual.model_dump(include=fields) + differing = [ + name + for name in SnapshotRuntimeIdentity.model_fields + if expected_values[name] != actual_values[name] + ] + if differing: + raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}") + + +def _existing_path_chain(path: Path): + current = path + while True: + if current.exists() or current.is_symlink(): + yield current + if current.parent == current: + break + current = current.parent + + +def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool: + return ( + path_stat.st_uid == 0 + and stat.S_ISDIR(path_stat.st_mode) + and bool(path_stat.st_mode & stat.S_ISVTX) + ) + + +def validate_artifact_root(path: Path, *, creating: bool) -> None: + """Validate that an artifact path cannot be replaced by another user.""" + path = Path(os.path.abspath(path)) + if path.is_symlink(): + raise SnapshotSecurityError(f"snapshot path is a symlink: {path}") + if not creating and not path.exists(): + raise SnapshotSecurityError(f"snapshot path does not exist: {path}") + + effective_uid = os.geteuid() + for component in _existing_path_chain(path): + component_stat = component.lstat() + if stat.S_ISLNK(component_stat.st_mode): + raise SnapshotSecurityError( + f"snapshot path contains a symlink: {component}" + ) + if component_stat.st_uid not in (0, effective_uid): + raise SnapshotSecurityError( + f"snapshot path has an untrusted owner: {component}" + ) + if component_stat.st_mode & ( + stat.S_IWGRP | stat.S_IWOTH + ) and not _is_trusted_sticky_directory(component_stat): + raise SnapshotSecurityError( + f"snapshot path has a group- or world-writable ancestor: {component}" + ) + + if path.exists(): + path_stat = path.lstat() + if not stat.S_ISDIR(path_stat.st_mode): + raise SnapshotSecurityError(f"snapshot path is not a directory: {path}") + if path_stat.st_uid != effective_uid: + raise SnapshotSecurityError( + f"snapshot directory is not owned by the current user: {path}" + ) + if stat.S_IMODE(path_stat.st_mode) & 0o077: + raise SnapshotSecurityError( + f"snapshot directory must have mode 0700: {path}" + ) + + +def _validation_path(error: ValidationError) -> str: + detail = error.errors(include_url=False)[0] + location = detail["loc"] + if not location: + return "JSON" if detail["type"] == "json_invalid" else "root" + return ".".join(map(str, location)) + + +def _fsync_directory(path: Path) -> None: + directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None: + temporary = path.with_name(f"{path.name}.tmp") + if overwrite: + temporary.unlink(missing_ok=True) + elif path.exists() or path.is_symlink(): + raise FileExistsError(path) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(temporary, flags, 0o600) + try: + with os.fdopen(descriptor, "wb") as output: + output.write( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + _fsync_directory(path.parent) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + +def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None: + """Write a new private manifest and make its directory entry durable.""" + path = Path(path) + validate_artifact_root(path, creating=False) + destination = path / "manifest.json" + try: + _write_json_atomic(destination, manifest.model_dump(mode="json")) + except FileExistsError as error: + raise SnapshotSecurityError( + f"manifest already exists: {destination}" + ) from error + + +def read_manifest(path: Path) -> SnapshotManifest: + """Read a snapshot published by atomically installing its manifest.""" + path = Path(path) + validate_artifact_root(path, creating=False) + manifest_path = path / "manifest.json" + if manifest_path.is_symlink(): + raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}") + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + file_descriptor = os.open(manifest_path, flags) + except FileNotFoundError as error: + raise SnapshotSecurityError( + f"manifest does not exist: {manifest_path}" + ) from error + try: + with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file: + payload = manifest_file.read() + except UnicodeDecodeError as error: + raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error + try: + return SnapshotManifest.model_validate_json(payload, strict=True) + except ValidationError as error: + raise SnapshotCompatibilityError( + f"invalid snapshot manifest: {_validation_path(error)}" + ) from error + + +def inspect_snapshot(path: Path) -> dict[str, object]: + return read_manifest(path).model_dump(mode="json") diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py new file mode 100644 index 000000000000..582bc9a7d9cb --- /dev/null +++ b/vllm/snapshot/runtime.py @@ -0,0 +1,1246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host runtime operations for local snapshots.""" + +import argparse +import ast +import hashlib +import importlib.metadata +import importlib.util +import json +import os +import platform +import select +import shutil +import signal +import socket +import stat +import struct +import subprocess +import sys +import tempfile +import time +import urllib.request +from collections.abc import Iterable +from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import NoReturn + +import regex as re +from pydantic import ValidationError + +from vllm.snapshot.manifest import ( + ReadyMarker, + SnapshotManifest, + SnapshotRuntimeIdentity, + _fsync_directory, + _validation_path, + _write_json_atomic, + validate_artifact_root, + write_manifest_atomic, +) +from vllm.snapshot.types import Oracle + + +class SnapshotCreateError(RuntimeError): + """Snapshot creation did not produce a complete artifact.""" + + +class SnapshotRestoreError(RuntimeError): + """Snapshot restore failed.""" + + +_COMMON_CRIU_OPTIONS = ( + "--shell-job", + "--ext-unix-sk", + "--tcp-established", + "--link-remap", + "--file-locks", +) +_TCP_TABLES = ( + ("AF_INET", Path("/proc/net/tcp")), + ("AF_INET6", Path("/proc/net/tcp6")), +) + +_CACHE_REMEDY = "recreate the snapshot or restore with the same cache state" +# Recorded instead of a digest for files whose content is volatile by design. +_CACHE_LOCKED = "locked" +# Recorded instead of a digest for a file capture found but could not read. +_CACHE_UNREADABLE = "unreadable" + + +def _generated_cache_prefixes() -> tuple[str, ...]: + """Path prefixes of the generated caches an engine process keeps open. + + vLLM redirects Inductor and Triton under ``VLLM_CACHE_ROOT`` when it + compiles, but a process can still hold their upstream defaults open. + A cache root is resolved and separator-terminated so that it still matches + the kernel-canonical descriptor targets under a symlinked root, and so that + a sibling directory sharing its name does not. Inductor's default location + is a bare prefix instead, because its directory carries a per-user suffix. + """ + import vllm.envs as envs + + inductor_root = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + roots = tuple( + f"{os.path.realpath(os.path.expanduser(root)).rstrip(os.sep)}{os.sep}" + for root in ( + envs.VLLM_CACHE_ROOT, + os.environ.get("TRITON_CACHE_DIR") or "~/.triton/cache", + "~/.cache/flashinfer", + *((inductor_root,) if inductor_root else ()), + ) + ) + if inductor_root: + return roots + return roots + ( + os.path.join(os.path.realpath(tempfile.gettempdir()), "torchinductor_"), + ) + + +def _error_detail(error: BaseException) -> str: + message = str(error) + return f"{type(error).__name__}: {message}" if message else type(error).__name__ + + +def _close_pidfds(handles: Iterable[tuple[int, int]]) -> None: + for _pid, pidfd in handles: + with suppress(OSError): + os.close(pidfd) + + +@dataclass(frozen=True) +class ProcessInventory: + root_pid: int + process_tree: tuple[int, ...] + cuda_holders: tuple[int, ...] + gpu_uuid: str + external_cache_files: tuple[tuple[str, str], ...] = () + + +@dataclass(frozen=True) +class _TcpSocketRecord: + family: str + local_raw: str + remote_raw: str + inode: int + + +def _process_comm(pid: int) -> str: + try: + return (Path("/proc") / str(pid) / "comm").read_text().strip() + except OSError: + return "unknown" + + +def _decode_endpoint(family: str, raw: str) -> str: + """Render a ``/proc/net/tcp`` ``HEXADDR:HEXPORT`` field as ``ip:port``. + + The kernel prints the address as little-endian 32-bit words: one for IPv4, + four in order for IPv6. + """ + hex_address, _, hex_port = raw.partition(":") + port = int(hex_port, 16) + if family == "AF_INET": + packed = struct.pack(" None: + owned = tuple(record for record in records if record.inode in socket_owners) + endpoints = { + (record.family, record.local_raw, record.remote_raw) for record in owned + } + offenders = tuple( + record + for record in owned + if (record.family, record.remote_raw, record.local_raw) not in endpoints + ) + if not offenders: + return + named = "; ".join( + f"pid {socket_owners[record.inode]} " + f"({_process_comm(socket_owners[record.inode])}) " + f"{_decode_endpoint(record.family, record.local_raw)} -> " + f"{_decode_endpoint(record.family, record.remote_raw)}" + for record in offenders[:5] + ) + remaining = f" (+{len(offenders) - 5} more)" if len(offenders) > 5 else "" + raise SnapshotCreateError( + f"snapshot tree has an external established TCP connection: {named}{remaining}" + ) + + +class LocalSnapshotTools: + """Host-owned CRIU and cuda-checkpoint implementation.""" + + def __init__(self) -> None: + self.criu = shutil.which("criu") or "criu" + self.cuda_checkpoint = shutil.which("cuda-checkpoint") or "cuda-checkpoint" + self.nvidia_smi = shutil.which("nvidia-smi") or "nvidia-smi" + plugin_dir = os.environ.get("CRIU_CUDA_PLUGIN_DIR", "") + self.plugin_dir = Path(plugin_dir) if plugin_dir else None + self.shm_dir = Path("/dev/shm") + self.timeout_s = float(os.environ.get("VLLM_SNAPSHOT_TIMEOUT_S", "900")) + if not 0 < self.timeout_s < float("inf"): + raise ValueError("VLLM_SNAPSHOT_TIMEOUT_S must be a positive finite number") + self._children: dict[int, subprocess.Popen[bytes]] = {} + self._restored_processes: dict[int, tuple[tuple[int, int], ...]] = {} + + def _privileged(self) -> list[str]: + return [] if os.geteuid() == 0 else ["sudo", "-n"] + + def _run( + self, + command: list[str], + *, + timeout: float | None = None, + ) -> subprocess.CompletedProcess[str]: + return subprocess.run( + command, + check=True, + capture_output=True, + text=True, + timeout=timeout or self.timeout_s, + ) + + def preflight(self, action: str, artifact: Path) -> None: + if platform.system() != "Linux" or platform.machine() != "x86_64": + raise RuntimeError("snapshot requires Linux x86_64") + if action == "restore" and ( + not callable(getattr(os, "pidfd_open", None)) + or not callable(getattr(signal, "pidfd_send_signal", None)) + ): + raise RuntimeError("snapshot requires Linux pidfd support") + for command in (self.criu, self.cuda_checkpoint, self.nvidia_smi): + if shutil.which(command) is None: + raise RuntimeError(f"snapshot dependency not found: {command}") + if self.plugin_dir is None: + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") + if not (self.plugin_dir / "cuda_plugin.so").is_file(): + raise RuntimeError(f"CRIU CUDA plugin not found in {self.plugin_dir}") + if os.geteuid() != 0: + self._run(["sudo", "-n", "true"], timeout=5) + if action == "restore": + validate_artifact_root(artifact, creating=False) + else: + validate_artifact_root(artifact, creating=True) + + def launch_child( + self, + workdir: Path, + engine_argv: tuple[str, ...], + ) -> int: + command = [ + sys.executable, + "-m", + "vllm.snapshot.server", + "--ready-file", + str(workdir / "ready.json"), + "--release-file", + str(workdir / "release.json"), + "--release-timeout-s", + str(self.timeout_s), + "--", + *engine_argv, + ] + child_environment = os.environ.copy() + child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + with (workdir / "child.log").open("wb") as log_file: + process = subprocess.Popen( + command, + env=child_environment, + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + self._children[process.pid] = process + return process.pid + + def wait_ready(self, workdir: Path, root_pid: int) -> Oracle: + ready_file = workdir / "ready.json" + deadline = time.monotonic() + self.timeout_s + waitid = os.waitid # type: ignore[attr-defined] + while time.monotonic() < deadline: + if ready_file.is_file() and ready_file.stat().st_size: + try: + marker = ReadyMarker.model_validate_json(ready_file.read_text()) + except ValidationError as error: + raise SnapshotCreateError( + f"snapshot ready marker invalid: {_validation_path(error)}" + ) from error + return Oracle( + token_ids=marker.token_ids, + text=marker.text, + sampled_token_logprob=marker.sampled_token_logprob, + ) + process = self._children.get(root_pid) + if process is not None and waitid( + os.P_PID, + root_pid, + os.WEXITED | os.WNOHANG | os.WNOWAIT, + ): + log = (workdir / "child.log").read_text(errors="replace") + raise SnapshotCreateError(f"snapshot child exited before ready:\n{log}") + try: + os.kill(root_pid, 0) + except ProcessLookupError as error: + log = (workdir / "child.log").read_text(errors="replace") + raise SnapshotCreateError( + f"snapshot child exited before ready:\n{log}" + ) from error + time.sleep(0.05) + raise SnapshotCreateError("snapshot child did not become ready before timeout") + + def _tree_pids(self, root_pid: int) -> tuple[int, ...]: + rows = self._run(["ps", "-eo", "pid=,ppid="]).stdout.splitlines() + children: dict[int, list[int]] = {} + for row in rows: + pid, parent_pid = map(int, row.split()) + children.setdefault(parent_pid, []).append(pid) + stack = [root_pid] + found: list[int] = [] + while stack: + pid = stack.pop() + if pid in found: + continue + found.append(pid) + stack.extend(children.get(pid, [])) + return tuple(found) + + def _cuda_process_rows(self) -> tuple[str, ...]: + output = self._run( + [ + self.nvidia_smi, + "--query-compute-apps=pid,gpu_uuid", + "--format=csv,noheader,nounits", + ] + ).stdout + return tuple(line for line in output.splitlines() if line.strip()) + + def _cuda_pids(self, rows: tuple[str, ...]) -> tuple[int, ...]: + return tuple(sorted({int(line.split(",", 1)[0].strip()) for line in rows})) + + def _descriptor_targets(self, pid: int) -> tuple[str, ...]: + descriptor_dir = Path("/proc") / str(pid) / "fd" + try: + descriptors = tuple(descriptor_dir.iterdir()) + except FileNotFoundError as error: + raise SnapshotCreateError( + f"snapshot process exited during descriptor inventory: {pid}" + ) from error + targets: list[str] = [] + for descriptor in descriptors: + try: + targets.append(os.readlink(descriptor)) + except FileNotFoundError: + continue + return tuple(targets) + + def _descriptor_inventory( + self, process_tree: tuple[int, ...] + ) -> tuple[tuple[int, ...], dict[int, int]]: + """Collect io_uring holders and each socket inode's owning pid.""" + io_uring_pids: list[int] = [] + socket_owners: dict[int, int] = {} + for pid in process_tree: + targets = self._descriptor_targets(pid) + if "anon_inode:[io_uring]" in targets: + io_uring_pids.append(pid) + for target in targets: + if target.startswith("socket:[") and target.endswith("]"): + # Shared descriptors name one holder; any of them locates + # the connection for the operator. + socket_owners.setdefault(int(target[len("socket:[") : -1]), pid) + return tuple(io_uring_pids), socket_owners + + def _tcp_records(self) -> tuple[_TcpSocketRecord, ...]: + records: list[_TcpSocketRecord] = [] + for family, table in _TCP_TABLES: + if not table.exists(): + continue + for row in table.read_text().splitlines()[1:]: + fields = row.split() + if fields[3] != "01": + continue + records.append( + _TcpSocketRecord( + family=family, + local_raw=fields[1], + remote_raw=fields[2], + inode=int(fields[9]), + ) + ) + return tuple(records) + + def _cache_fingerprint(self, path: Path) -> str | None: + """Digest a cache file, or mark the ones whose content is volatile.""" + try: + if not stat.S_ISREG(path.stat().st_mode): + return None + except OSError: + return None + if path.name.endswith(".lock"): + return _CACHE_LOCKED + try: + return self._sha256(path) + except OSError: + # Restore can only check that this one still exists, so record the + # downgrade rather than passing it off as a volatile file. + return _CACHE_UNREADABLE + + def _external_cache_files( + self, process_tree: tuple[int, ...], artifact: Path + ) -> tuple[tuple[str, str], ...]: + """Fingerprint the generated-cache files the captured tree holds open. + + CRIU reopens these by path on restore, so a JIT cache that rotates + between capture and restore fails the restore permanently. Recording + them lets restore fail early and name the file that changed. Files + inside the artifact are excluded: they belong to the snapshot itself. + Deleted targets are excluded too, because CRIU remaps those links. + """ + prefixes = _generated_cache_prefixes() + artifact_prefix = f"{os.path.abspath(artifact)}{os.sep}" + candidates: set[str] = set() + for pid in process_tree: + for target in self._descriptor_targets(pid): + if target.endswith(" (deleted)") or target.startswith(artifact_prefix): + continue + if target.endswith(".lock") or target.startswith(prefixes): + candidates.add(target) + recorded = ( + (path, self._cache_fingerprint(Path(path))) for path in sorted(candidates) + ) + return tuple( + (path, fingerprint) + for path, fingerprint in recorded + if fingerprint is not None + ) + + def inventory(self, root_pid: int, artifact: Path) -> ProcessInventory: + process_tree = self._tree_pids(root_pid) + cuda_rows = self._cuda_process_rows() + cuda_holders = tuple( + pid for pid in self._cuda_pids(cuda_rows) if pid in process_tree + ) + if not cuda_holders: + raise SnapshotCreateError("snapshot tree has no CUDA-holding process") + io_uring_pids, socket_owners = self._descriptor_inventory(process_tree) + if io_uring_pids: + raise SnapshotCreateError( + "snapshot tree owns io_uring state that CRIU cannot dump; set " + "kernel.io_uring_disabled=1 before starting an unprivileged " + "vLLM process, or use =2 to disable it host-wide " + f"(pids: {', '.join(map(str, io_uring_pids))})" + ) + _validate_tcp_connections(self._tcp_records(), socket_owners) + gpu_uuid = self._gpu_uuid_for_pids(cuda_holders, cuda_rows) + return ProcessInventory( + root_pid=root_pid, + process_tree=process_tree, + cuda_holders=cuda_holders, + gpu_uuid=gpu_uuid, + external_cache_files=self._external_cache_files(process_tree, artifact), + ) + + def _gpu_uuid_for_pids(self, pids: tuple[int, ...], rows: tuple[str, ...]) -> str: + wanted = set(pids) + matches: set[str] = set() + try: + for line in rows: + pid, gpu_uuid = (item.strip() for item in line.split(",", 1)) + if int(pid) in wanted: + matches.add(gpu_uuid) + except (TypeError, ValueError) as error: + raise SnapshotCreateError( + "could not identify the GPU used by the snapshot process tree" + ) from error + if len(matches) != 1: + raise SnapshotCreateError( + "snapshot process tree must use exactly one identifiable GPU" + ) + return matches.pop() + + def _criu( + self, + action: str, + artifact: Path, + arguments: list[str], + ) -> None: + if self.plugin_dir is None: + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") + env = os.environ.copy() + env["PATH"] = f"{Path(self.cuda_checkpoint).parent}:{env['PATH']}" + env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] = str(artifact / "ignored-nvidia-fds.tsv") + env["RIG_CRIU_IGNORED_FD_MAP"] = env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] + command = [ + *self._privileged(), + "env", + f"PATH={env['PATH']}", + f"VLLM_SNAPSHOT_IGNORED_FD_MAP={env['VLLM_SNAPSHOT_IGNORED_FD_MAP']}", + f"RIG_CRIU_IGNORED_FD_MAP={env['RIG_CRIU_IGNORED_FD_MAP']}", + self.criu, + action, + "-L", + str(self.plugin_dir), + *arguments, + ] + self._run(command) + + def _record_child_log_size(self, artifact: Path) -> None: + child_log = artifact / "child.log" + size_path = artifact / "child.log.snapshot-size" + payload = f"{child_log.stat().st_size}\n".encode() + descriptor = os.open( + size_path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + with os.fdopen(descriptor, "wb") as output: + output.write(payload) + output.flush() + os.fsync(output.fileno()) + + def _link_remap_names(self) -> set[str]: + return { + path.name + for path in self.shm_dir.glob("link_remap.*") + if path.name.removeprefix("link_remap.").isdigit() + } + + def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None: + if not names: + return + destination_dir = artifact / "link-remaps" + destination_dir.mkdir(mode=0o700) + for name in sorted(names): + source = self.shm_dir / name + source_fd = os.open( + source, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + metadata = os.fstat(source_fd) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid(): + raise SnapshotCreateError(f"invalid CRIU link remap: {source}") + destination_fd = os.open( + destination_dir / name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + with os.fdopen(destination_fd, "wb") as destination: + with os.fdopen(os.dup(source_fd), "rb") as contents: + shutil.copyfileobj(contents, destination) + destination.flush() + os.fsync(destination.fileno()) + finally: + os.close(source_fd) + source.unlink() + + def _rollback_link_remaps(self, created: dict[str, tuple[int, int]]) -> None: + for name, identity in created.items(): + target = self.shm_dir / name + try: + metadata = target.lstat() + except FileNotFoundError: + continue + if (metadata.st_dev, metadata.st_ino) == identity: + target.unlink() + + def _restore_link_remaps(self, artifact: Path) -> dict[str, tuple[int, int]]: + source_dir = artifact / "link-remaps" + if not source_dir.exists(): + return {} + created: dict[str, tuple[int, int]] = {} + try: + for source in sorted(source_dir.iterdir()): + if ( + not source.name.removeprefix("link_remap.").isdigit() + or source.is_symlink() + or not source.is_file() + ): + raise SnapshotRestoreError(f"invalid saved link remap: {source}") + payload = source.read_bytes() + target = self.shm_dir / source.name + try: + descriptor = os.open( + target, + os.O_WRONLY + | os.O_CREAT + | os.O_EXCL + | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + except FileExistsError as error: + try: + descriptor = os.open( + target, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + except OSError: + raise SnapshotRestoreError( + f"conflicting CRIU link remap exists: {target}" + ) from error + with os.fdopen(descriptor, "rb") as existing: + metadata = os.fstat(existing.fileno()) + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_uid != os.getuid() + or existing.read() != payload + ): + raise SnapshotRestoreError( + f"conflicting CRIU link remap exists: {target}" + ) from error + continue + metadata = os.fstat(descriptor) + created[source.name] = (metadata.st_dev, metadata.st_ino) + with os.fdopen(descriptor, "wb") as output: + output.write(payload) + output.flush() + os.fsync(output.fileno()) + except BaseException: + self._rollback_link_remaps(created) + raise + return created + + def _reset_child_log(self, artifact: Path) -> None: + size_path = artifact / "child.log.snapshot-size" + descriptor = os.open( + size_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + with os.fdopen(descriptor, "r") as source: + size = int(source.read().strip()) + if size < 0: + raise SnapshotRestoreError("captured child log size is invalid") + + child_log = artifact / "child.log" + descriptor = os.open( + child_log, + os.O_WRONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size < size: + raise SnapshotRestoreError("captured child log is invalid") + os.ftruncate(descriptor, size) + os.fsync(descriptor) + finally: + os.close(descriptor) + + def dump(self, workdir: Path, inventory: ProcessInventory) -> None: + images = workdir / "images" + images.mkdir(mode=0o700) + fd_map = workdir / "ignored-nvidia-fds.tsv" + fd_map.touch(mode=0o600) + remaps_before = self._link_remap_names() + try: + self._criu( + "dump", + workdir, + [ + "--tree", + str(inventory.root_pid), + "--images-dir", + str(images), + "--log-file", + "dump.log", + *_COMMON_CRIU_OPTIONS, + ], + ) + self._capture_link_remaps( + workdir, + self._link_remap_names() - remaps_before, + ) + except BaseException: + for name in self._link_remap_names() - remaps_before: + (self.shm_dir / name).unlink(missing_ok=True) + raise + self._record_child_log_size(workdir) + + def verify_dead(self, inventory: ProcessInventory) -> None: + process = self._children.get(inventory.root_pid) + if process is not None: + process.wait(timeout=10) + for pid in inventory.process_tree: + try: + os.kill(pid, 0) + except ProcessLookupError: + continue + raise SnapshotCreateError(f"process survived CRIU dump: {pid}") + self._children.pop(inventory.root_pid, None) + + def _version(self, command: list[str]) -> str: + try: + result = self._run(command, timeout=10) + except subprocess.CalledProcessError as error: + return (error.stdout or error.stderr).splitlines()[0] + return result.stdout.splitlines()[0] if result.stdout else "unknown" + + def _sha256(self, path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def _plugin_identity(self) -> tuple[tuple[str, str], ...]: + if self.plugin_dir is None: + return () + return tuple( + (f"snapshot_plugin_sha256:{path.name}", self._sha256(path)) + for path in sorted(self.plugin_dir.glob("*.so")) + ) + + def _torch_identity(self) -> tuple[str, str]: + torch_version = importlib.metadata.version("torch") + torch_spec = importlib.util.find_spec("torch") + locations = torch_spec and torch_spec.submodule_search_locations + if not locations: + raise RuntimeError("installed torch package could not be located") + version_path = Path(next(iter(locations))) / "version.py" + # Parse the version file rather than importing torch: snapshot create's + # CLI path stays import-light (vllm/entrypoints/cli/snapshot.py). + syntax = ast.parse(version_path.read_text(), filename=str(version_path)) + cuda_runtime = "" + for statement in syntax.body: + target = None + value = None + if isinstance(statement, ast.Assign) and len(statement.targets) == 1: + target = statement.targets[0] + value = statement.value + elif isinstance(statement, ast.AnnAssign): + target = statement.target + value = statement.value + if isinstance(target, ast.Name) and target.id == "cuda" and value: + parsed = ast.literal_eval(value) + cuda_runtime = str(parsed or "") + break + return torch_version, cuda_runtime + + def _gpu_identity(self, gpu_uuid: str) -> tuple[str, str, str]: + output = self._run( + [ + self.nvidia_smi, + "--query-gpu=name,uuid,driver_version", + "--format=csv,noheader,nounits", + f"--id={gpu_uuid}", + ] + ).stdout.strip() + name, uuid, driver = (item.strip() for item in output.split(",", 2)) + return name, uuid, driver + + def _artifact_bytes(self, workdir: Path) -> int: + return sum(path.stat().st_size for path in workdir.rglob("*") if path.is_file()) + + def _environment_identity(self) -> tuple[tuple[str, str], ...]: + prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_") + environment = os.environ.copy() + environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + selected = tuple( + sorted( + (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest()) + for key, value in environment.items() + if key.startswith(prefixes) + and key not in {"VLLM_API_KEY", "VLLM_SNAPSHOT_TIMEOUT_S"} + ) + ) + return selected + self._plugin_identity() + + def make_manifest( + self, + args: argparse.Namespace, + engine_argv: tuple[str, ...], + inventory: ProcessInventory, + oracle: Oracle, + workdir: Path, + ) -> SnapshotManifest: + identity = self.current_identity(inventory.gpu_uuid) + redacted_argv = list(engine_argv) + redact_api_keys = False + for index, argument in enumerate(redacted_argv): + option, separator, _ = argument.partition("=") + normalized_option = option.replace("_", "-") + # Parsed argv can contain unambiguous long-option abbreviations. + if len(normalized_option) > 2 and "--api-key".startswith(normalized_option): + redact_api_keys = not separator + if separator: + redacted_argv[index] = f"{option}=***" + elif len(normalized_option) > 2 and "--hf-token".startswith( + normalized_option + ): + if separator: + redacted_argv[index] = f"{option}=***" + elif index + 1 < len(redacted_argv): + redacted_argv[index + 1] = "***" + redact_api_keys = False + elif argument.startswith("-"): + redact_api_keys = False + elif redact_api_keys: + redacted_argv[index] = "***" + engine_argv = tuple(redacted_argv) + revision = str(getattr(args, "revision", None) or "") + tokenizer_revision = str(getattr(args, "tokenizer_revision", None) or revision) + source_model = str(getattr(args, "model_tag", None) or args.model) + from vllm.config.model import get_served_model_name + + served_model_name = get_served_model_name( + source_model, getattr(args, "served_model_name", None) + ) + return SnapshotManifest( + schema_version=1, + boundary="post-engine-init-reloadable-state-released", + created_at=datetime.now(timezone.utc).isoformat(), + artifact_bytes=self._artifact_bytes(workdir), + **identity.model_dump(), + model=source_model, + served_model_name=str(served_model_name), + model_revision=revision, + tokenizer_revision=tokenizer_revision, + engine_argv=engine_argv, + process_tree=inventory.process_tree, + cuda_holders=inventory.cuda_holders, + external_cache_files=inventory.external_cache_files, + oracle_token_ids=oracle.token_ids, + oracle_text=oracle.text, + oracle_sampled_token_logprob=oracle.sampled_token_logprob, + ) + + def publish(self, workdir: Path, manifest: SnapshotManifest) -> None: + write_manifest_atomic(workdir, manifest) + _fsync_directory(workdir.parent) + + def current_identity(self, gpu_uuid: str) -> SnapshotRuntimeIdentity: + gpu_name, gpu_uuid, driver_version = self._gpu_identity(gpu_uuid) + torch_version, cuda_runtime = self._torch_identity() + host_id = Path("/etc/machine-id").read_text().strip() + return SnapshotRuntimeIdentity( + vllm_version=importlib.metadata.version("vllm"), + python_version=platform.python_version(), + torch_version=torch_version, + cuda_runtime=cuda_runtime, + driver_version=driver_version, + criu_version=self._version([self.criu, "--version"]), + cuda_checkpoint_sha256=self._sha256(Path(self.cuda_checkpoint)), + kernel_release=platform.release(), + host_id=host_id, + gpu_name=gpu_name, + gpu_uuid=gpu_uuid, + environment=self._environment_identity(), + ) + + def _read_restored_pid(self, pidfile: Path) -> int: + try: + payload = self._run([*self._privileged(), "cat", str(pidfile)]).stdout + except Exception as error: + raise SnapshotRestoreError("restored PID file is missing") from error + match = re.fullmatch(r"([1-9][0-9]{0,9})\n?", payload) + if match is None: + raise SnapshotRestoreError("restored PID file is invalid") + restored_pid = int(match.group(1)) + if restored_pid > 2**31 - 1: + raise SnapshotRestoreError("restored PID file is invalid") + return restored_pid + + def _process_state(self, pid: int) -> tuple[int, int, int, int]: + payload = (Path("/proc") / str(pid) / "stat").read_text() + command_end = payload.rfind(")") + fields = payload[command_end + 2 :].split() + if command_end < 0 or len(fields) < 20: + raise ValueError(f"invalid process state for PID {pid}") + return (int(fields[1]), int(fields[2]), int(fields[3]), int(fields[19])) + + def _process_states(self) -> dict[int, tuple[int, int, int, int]]: + states: dict[int, tuple[int, int, int, int]] = {} + for entry in Path("/proc").iterdir(): + if not entry.name.isdigit(): + continue + pid = int(entry.name) + try: + states[pid] = self._process_state(pid) + except (OSError, ValueError): + continue + return states + + def _process_command(self, pid: int) -> tuple[str, ...]: + payload = (Path("/proc") / str(pid) / "cmdline").read_bytes() + return tuple( + os.fsdecode(argument) for argument in payload.split(b"\0") if argument + ) + + def _pidfd_open(self, pid: int) -> int: + pidfd_open = getattr(os, "pidfd_open", None) + if not callable(pidfd_open): + raise SnapshotRestoreError("snapshot requires Linux pidfd support") + return pidfd_open(pid) + + def _pidfd_exited(self, pidfd: int) -> bool: + poller = select.poll() + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) + return any(self._pidfd_event_is_exit(event) for _fd, event in poller.poll(0)) + + def _pidfd_event_is_exit(self, event: int) -> bool: + if event & (select.POLLERR | select.POLLNVAL): + raise SnapshotRestoreError("restored process pidfd poll failed") + return bool(event & (select.POLLIN | select.POLLHUP)) + + def _pin_restored_tree( + self, + artifact: Path, + process_tree: tuple[int, ...], + cuda_holders: tuple[int, ...], + ) -> None: + root_pid = process_tree[0] + expected_pids = set(process_tree) + expected_cuda_holders = set(cuda_holders) + if not expected_cuda_holders or not expected_cuda_holders.issubset( + expected_pids + ): + raise SnapshotRestoreError( + "captured CUDA holders are outside the captured process tree" + ) + if root_pid in self._restored_processes: + raise SnapshotRestoreError("restored PID already has a pinned transaction") + handles: list[tuple[int, int]] = [] + try: + for pid in process_tree: + handles.append((pid, self._pidfd_open(pid))) + handles_by_pid = dict(handles) + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): + raise SnapshotRestoreError( + "restored process exited during verification" + ) + + root_state = self._process_state(root_pid) + if root_state[1:3] != (root_pid, root_pid): + raise SnapshotRestoreError( + "restored root is not the session and process-group leader" + ) + command = self._process_command(root_pid) + release_file = str(artifact / "release.json") + if not any( + command[index : index + 2] == ("-m", "vllm.snapshot.server") + for index in range(len(command) - 1) + ) or not any( + command[index : index + 2] == ("--release-file", release_file) + for index in range(len(command) - 1) + ): + raise SnapshotRestoreError( + "restored root command does not match the snapshot barrier" + ) + if self._process_state(root_pid) != root_state: + raise SnapshotRestoreError("restored PID changed during verification") + if self._pidfd_exited(handles[0][1]): + raise SnapshotRestoreError("restored root exited during verification") + + first_states = self._process_states() + live_session = { + pid for pid, state in first_states.items() if state[2] == root_pid + } + if not expected_cuda_holders.issubset(live_session): + raise SnapshotRestoreError( + "restored session is missing captured CUDA holders" + ) + if live_session != expected_pids: + raise SnapshotRestoreError( + "restored session does not match the captured process tree" + ) + if first_states[root_pid] != root_state: + raise SnapshotRestoreError("restored PID changed during verification") + for pid in process_tree[1:]: + state = first_states[pid] + pidfd = handles_by_pid[pid] + if self._process_state(pid) != state or self._pidfd_exited(pidfd): + raise SnapshotRestoreError( + "restored process tree changed during verification" + ) + + second_states = self._process_states() + second_session = { + pid for pid, state in second_states.items() if state[2] == root_pid + } + if second_session != expected_pids or any( + second_states[pid] != first_states[pid] for pid in process_tree + ): + raise SnapshotRestoreError( + "restored process tree changed during verification" + ) + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): + raise SnapshotRestoreError( + "restored process exited during verification" + ) + except BaseException as error: + _close_pidfds(handles) + if isinstance(error, SnapshotRestoreError): + raise + if not isinstance(error, Exception): + raise + raise SnapshotRestoreError( + "restored process tree could not be pinned" + ) from error + self._restored_processes[root_pid] = tuple(handles) + + def _abort_failed_restore( + self, + artifact: Path, + root_pid: int, + staged_remaps: dict[str, tuple[int, int]], + primary: BaseException, + ) -> NoReturn: + secondary: list[str] = [] + try: + _write_json_atomic( + artifact / "release.json", {"release": False}, overwrite=True + ) + except BaseException as error: + secondary.append(f"abort marker failed: {_error_detail(error)}") + # The marker above is only a request. Terminate and wait for the tree we + # own, before releasing the shared state it may still be using. A tree + # that could not be pinned is not ours to signal. + if root_pid in self._restored_processes: + try: + self.cleanup(root_pid) + except BaseException as error: + secondary.append(f"restore cleanup failed: {_error_detail(error)}") + try: + self._rollback_link_remaps(staged_remaps) + except BaseException as error: + secondary.append(f"remap rollback failed: {_error_detail(error)}") + if secondary: + raise SnapshotRestoreError( + f"snapshot restore failed: {_error_detail(primary)}; " + + "; ".join(secondary) + ) from primary + raise primary + + def _verify_external_cache_files(self, manifest: SnapshotManifest) -> None: + """Reject a rotated generated cache before CRIU reopens it by path. + + Like the PID probe above this reports a stale host early rather than + guarding it: CRIU stays authoritative for a cache replaced afterwards. + """ + for path, fingerprint in manifest.external_cache_files: + target = Path(path) + if not target.is_file(): + raise SnapshotRestoreError( + f"external cache file vanished since capture: {path}; " + f"{_CACHE_REMEDY}" + ) + if fingerprint in (_CACHE_LOCKED, _CACHE_UNREADABLE): + continue + try: + current = self._sha256(target) + except OSError as error: + raise SnapshotRestoreError( + f"external cache file is unreadable since capture: {path}; " + f"{_CACHE_REMEDY}" + ) from error + if current != fingerprint: + raise SnapshotRestoreError( + f"external cache file changed since capture: {path}; " + f"{_CACHE_REMEDY}" + ) + + def restore(self, artifact: Path, manifest: SnapshotManifest) -> int: + # This is necessarily a best-effort TOCTOU check. CRIU restore and the + # rollback below remain authoritative if a PID is claimed after it. + for pid in manifest.process_tree: + try: + os.kill(pid, 0) + except ProcessLookupError: + continue + except PermissionError: + pass + except Exception as error: + raise SnapshotRestoreError( + f"captured PID availability probe failed: {pid}" + ) from error + raise SnapshotRestoreError(f"captured PID is already occupied: {pid}") + self._verify_external_cache_files(manifest) + expected_root_pid = manifest.process_tree[0] + (artifact / "release.json").unlink(missing_ok=True) + pidfile = artifact / "restored.pid" + pidfile.unlink(missing_ok=True) + self._reset_child_log(artifact) + staged_remaps = self._restore_link_remaps(artifact) + try: + self._criu( + "restore", + artifact, + [ + "--images-dir", + str(artifact / "images"), + "--log-file", + "restore.log", + *_COMMON_CRIU_OPTIONS, + "--restore-detached", + "--pidfile", + str(pidfile), + ], + ) + # Pin before reading the PID file so that a failure between restore + # and the read cannot orphan the tree. + self._pin_restored_tree( + artifact, manifest.process_tree, manifest.cuda_holders + ) + restored_pid = self._read_restored_pid(pidfile) + if restored_pid != expected_root_pid: + raise SnapshotRestoreError( + "restored PID does not match the captured process tree" + ) + except BaseException as error: + self._abort_failed_restore( + artifact, expected_root_pid, staged_remaps, error + ) + return restored_pid + + def release(self, artifact: Path, host: str | None, port: int) -> None: + if not 1 <= port <= 65535: + raise SnapshotRestoreError(f"snapshot restore port is invalid: {port}") + family = socket.AF_INET6 if host and ":" in host else socket.AF_INET + with socket.socket(family=family, type=socket.SOCK_STREAM) as listener_probe: + listener_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + listener_probe.bind((host or "", port)) + except OSError as error: + address = f"{host or '0.0.0.0'}:{port}" + raise SnapshotRestoreError( + f"snapshot restore address is already in use: {address}" + ) from error + + _write_json_atomic( + artifact / "release.json", + {"release": True, "host": host, "port": port}, + overwrite=True, + ) + + def _connect_host(self, host: str | None) -> str: + if not host or host == "0.0.0.0": + return "127.0.0.1" + if host == "::": + return "::1" + return host + + def wait_listener(self, root_pid: int, host: str | None, port: int) -> None: + deadline = time.monotonic() + self.timeout_s + connect_host = self._connect_host(host) + while time.monotonic() < deadline: + try: + with socket.create_connection((connect_host, port), timeout=1): + return + except OSError as error: + handles = self._restored_processes.get(root_pid) + if handles is None or self._pidfd_exited(handles[0][1]): + raise SnapshotRestoreError( + "restored process exited before binding HTTP" + ) from error + time.sleep(0.1) + raise SnapshotRestoreError("restored HTTP listener did not become ready") + + def request_oracle( + self, host: str | None, port: int, manifest: SnapshotManifest + ) -> Oracle: + connect_host = self._connect_host(host) + url_host = f"[{connect_host}]" if ":" in connect_host else connect_host + payload = json.dumps( + { + "model": manifest.served_model_name, + "prompt": "The capital of France is", + "min_tokens": 1, + "max_tokens": 1, + "temperature": 0, + "seed": 0, + "logprobs": 0, + "return_token_ids": True, + } + ).encode() + request = urllib.request.Request( + f"http://{url_host}:{port}/v1/completions", + data=payload, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + body = json.load(response) + choice = body["choices"][0] + try: + return Oracle( + token_ids=tuple(choice["token_ids"]), + text=choice["text"], + sampled_token_logprob=choice["logprobs"]["token_logprobs"][0], + ) + except ValueError as error: + raise SnapshotRestoreError( + "snapshot HTTP canary must return exactly one finite token logprob" + ) from error + except (IndexError, KeyError, TypeError) as error: + raise SnapshotRestoreError( + "snapshot HTTP canary response is missing sampled token logprob" + ) from error + + def cleanup(self, root_pid: int) -> None: + handles = self._restored_processes.get(root_pid) + if handles is None: + raise SnapshotRestoreError("no pinned restore transaction for cleanup") + pidfd_send_signal = getattr(signal, "pidfd_send_signal", None) + if not callable(pidfd_send_signal): + raise SnapshotRestoreError("restored process cleanup is incomplete") + try: + for _pid, pidfd in handles: + with suppress(ProcessLookupError): + pidfd_send_signal(pidfd, signal.SIGTERM, None, 0) + time.sleep(min(0.5, self.timeout_s)) + for _pid, pidfd in handles: + with suppress(ProcessLookupError): + pidfd_send_signal(pidfd, signal.SIGKILL, None, 0) + except BaseException as error: + raise SnapshotRestoreError( + f"restored process cleanup is incomplete: {_error_detail(error)}" + ) from error + + poller = select.poll() + for _pid, pidfd in handles: + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) + pending = {pidfd for _pid, pidfd in handles} + deadline = time.monotonic() + self.timeout_s + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise SnapshotRestoreError("restored process cleanup is incomplete") + for pidfd, event in poller.poll(max(1, int(remaining * 1000))): + if not self._pidfd_event_is_exit(event): + continue + pending.discard(pidfd) + with suppress(KeyError, OSError): + poller.unregister(pidfd) + self._restored_processes.pop(root_pid) + _close_pidfds(handles) + + def complete_restore(self, root_pid: int) -> None: + handles = self._restored_processes.pop(root_pid, None) + if handles is None: + raise SnapshotRestoreError("no pinned restore transaction to complete") + _close_pidfds(handles) + + def abort_create(self, root_pid: int) -> None: + process = self._children.get(root_pid) + if process is None: + return + with suppress(ProcessLookupError): + os.killpg(root_pid, signal.SIGKILL) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired as error: + raise SnapshotCreateError( + f"snapshot child process group {root_pid} survived SIGKILL" + ) from error + self._children.pop(root_pid, None) diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py new file mode 100644 index 000000000000..2dc0f3384aeb --- /dev/null +++ b/vllm/snapshot/server.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import asyncio +import json +import os +import stat +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from pydantic import ValidationError + +from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic +from vllm.snapshot.types import Oracle, oracles_match + +_CANARY_PROMPT = "The capital of France is" + + +class SnapshotCanaryError(RuntimeError): + """The initialized engine did not produce a valid snapshot oracle.""" + + +class SnapshotBarrierError(RuntimeError): + """The controller supplied an invalid snapshot release marker.""" + + +@dataclass(frozen=True) +class ListenerConfig: + host: str | None + port: int + + +@dataclass(frozen=True) +class ControlArgs: + ready_file: Path + release_file: Path + release_timeout_s: float + + +def write_ready_atomic(path: Path, oracle: Oracle) -> None: + _write_json_atomic( + path, + { + "sampled_token_logprob": oracle.sampled_token_logprob, + "token_ids": oracle.token_ids, + "text": oracle.text, + }, + ) + + +def read_release_marker(path: Path) -> ListenerConfig: + path = Path(path) + path_stat = path.lstat() + if stat.S_ISLNK(path_stat.st_mode): + raise SnapshotBarrierError("release marker must not be a symlink") + try: + payload = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise SnapshotBarrierError("release marker is not valid JSON") from error + try: + marker = ReleaseMarker.model_validate(payload) + except ValidationError as error: + raise SnapshotBarrierError( + f"release marker is invalid: {_validation_path(error)}" + ) from error + return ListenerConfig(host=marker.host, port=marker.port) + + +def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--ready-file", type=Path, required=True) + parser.add_argument("--release-file", type=Path, required=True) + parser.add_argument("--release-timeout-s", type=float, default=900.0) + control_args, remaining = parser.parse_known_args(argv) + if remaining and remaining[0] == "--": + remaining = remaining[1:] + if control_args.release_timeout_s <= 0: + raise ValueError("release timeout must be positive") + return ( + ControlArgs( + ready_file=control_args.ready_file, + release_file=control_args.release_file, + release_timeout_s=control_args.release_timeout_s, + ), + remaining, + ) + + +async def _release_reloadable_state(engine: Any) -> None: + """Discard model and KV state before the process image is captured.""" + await engine.sleep(level=2) + + +async def _restore_reloadable_state(engine: Any) -> None: + """Rebuild state discarded by ``_release_reloadable_state``.""" + await engine.wake_up(tags=["weights"]) + await engine.collective_rpc("reload_weights") + await engine.wake_up(tags=["kv_cache"]) + + +def oracle_from_request_output(request_output: Any) -> Oracle: + try: + candidate = request_output.outputs[0] + token_ids = tuple(candidate.token_ids) + (sampled_token_id,) = token_ids + return Oracle( + token_ids=token_ids, + text=candidate.text, + sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, + ) + except ValueError as error: + raise SnapshotCanaryError( + "snapshot canary must produce exactly one finite token logprob" + ) from error + except (AttributeError, IndexError, KeyError, TypeError) as error: + raise SnapshotCanaryError( + "snapshot canary did not return sampled token logprob" + ) from error + + +async def run_engine_canary(engine: Any) -> Oracle: + from vllm import SamplingParams + + final_output = None + sampling_params = SamplingParams( + temperature=0, + min_tokens=1, + max_tokens=1, + seed=0, + logprobs=0, + ) + async for output in engine.generate( + _CANARY_PROMPT, + sampling_params, + request_id="vllm-snapshot-canary", + ): + final_output = output + if final_output is None: + raise SnapshotCanaryError("canary generation returned no output") + return oracle_from_request_output(final_output) + + +def detach_snapshot_streams() -> None: + """Remove launch-log descriptors from the reusable process image.""" + sys.stdout.flush() + sys.stderr.flush() + sink = os.open(os.devnull, os.O_WRONLY) + try: + os.dup2(sink, sys.stdout.fileno()) + os.dup2(sink, sys.stderr.fileno()) + finally: + os.close(sink) + + +async def wait_for_release_marker( + path: Path, + *, + timeout_s: float, + poll_interval_s: float = 0.05, +) -> ListenerConfig: + if timeout_s <= 0: + raise ValueError("release timeout must be positive") + if poll_interval_s <= 0: + raise ValueError("release poll interval must be positive") + remaining_s = timeout_s + while not path.exists(): + if remaining_s <= 0: + raise TimeoutError(f"release marker not found before timeout: {path}") + delay_s = min(poll_interval_s, remaining_s) + await asyncio.sleep(delay_s) + remaining_s -= delay_s + return read_release_marker(path) + + +def parse_vllm_args(argv: list[str]) -> Any: + from vllm.entrypoints.launchers.cli_args import ( + make_arg_parser, + validate_parsed_serve_args, + ) + from vllm.utils.argparse_utils import FlexibleArgumentParser + + parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server") + args = make_arg_parser(parser).parse_args(argv) + if args.model_tag is not None: + args.model = args.model_tag + if args.grpc or args.headless: + raise ValueError("snapshot server requires the HTTP frontend") + if args.api_server_count not in (None, 1): + raise ValueError("snapshot server supports one API server") + args.api_server_count = None + validate_parsed_serve_args(args) + return args + + +async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None: + from vllm.entrypoints.launchers.api_server.entry import ( + build_and_serve, + build_async_engine_client, + ) + from vllm.entrypoints.launchers.launcher import setup_server + + args.enable_sleep_mode = True + + async with build_async_engine_client(args) as engine: + oracle = await run_engine_canary(engine) + try: + await _release_reloadable_state(engine) + await _restore_reloadable_state(engine) + rehearsal_oracle = await run_engine_canary(engine) + if not oracles_match(oracle, rehearsal_oracle): + raise SnapshotCanaryError("snapshot rehearsal changed canary output") + await _release_reloadable_state(engine) + except SnapshotCanaryError: + raise + except Exception as error: + raise SnapshotCanaryError("snapshot rehearsal failed") from error + + detach_snapshot_streams() + write_ready_atomic(control.ready_file, oracle) + listener = await wait_for_release_marker( + control.release_file, + timeout_s=control.release_timeout_s, + ) + args.host = listener.host + args.port = listener.port + await _restore_reloadable_state(engine) + listen_address, sock = setup_server(args, reuse_port=False) + try: + shutdown_task = await build_and_serve( + engine, + listen_address, + sock, + args, + ) + await shutdown_task + finally: + sock.close() + + +def main(argv: list[str] | None = None) -> None: + control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv) + args = parse_vllm_args(remaining) + import uvloop + + uvloop.run(run_vllm_snapshot_child(control, args)) + + +if __name__ == "__main__": + main() diff --git a/vllm/snapshot/types.py b/vllm/snapshot/types.py new file mode 100644 index 000000000000..236ba124365a --- /dev/null +++ b/vllm/snapshot/types.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Oracle: + token_ids: tuple[int, ...] + text: str + sampled_token_logprob: float + + def __post_init__(self) -> None: + if len(self.token_ids) != 1: + raise ValueError("snapshot oracle must contain exactly one token") + value = self.sampled_token_logprob + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise ValueError("sampled token logprob must be finite") + object.__setattr__(self, "sampled_token_logprob", float(value)) + + +def oracles_match(expected: Oracle, actual: Oracle) -> bool: + return ( + expected.token_ids == actual.token_ids + and expected.text == actual.text + and math.isclose( + expected.sampled_token_logprob, + actual.sampled_token_logprob, + rel_tol=0.0, + abs_tol=1e-3, + ) + ) diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index d214da3402cd..b91b1f26fd13 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -519,7 +519,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]: config_args = self.load_config_file(file_path) # 0th index might be the sub command {serve,chat,complete,...} - # optionally followed by model_tag (only for serve) + # optionally followed by model_tag (serve or snapshot create) # followed by config args # followed by rest of cli args. # maintaining this order will enforce the precedence @@ -527,8 +527,11 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]: if args[0].startswith("-"): # No sub command (e.g., api_server entry point) args = config_args + args[0:index] + args[index + 2 :] - elif args[0] == "serve": - model_in_cli = len(args) > 1 and not args[1].startswith("-") + elif args[0] == "serve" or args[:2] == ["snapshot", "create"]: + model_index = 1 if args[0] == "serve" else 2 + model_in_cli = len(args) > model_index and not args[model_index].startswith( + "-" + ) model_in_config = any(arg == "--model" for arg in config_args) if not model_in_cli and not model_in_config: @@ -540,15 +543,19 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]: if model_in_cli: # Model specified as positional arg, keep CLI version args = ( - [args[0]] - + [args[1]] + args[: model_index + 1] + config_args - + args[2:index] + + args[model_index + 1 : index] + args[index + 2 :] ) else: # No model in CLI, use config if available - args = [args[0]] + config_args + args[1:index] + args[index + 2 :] + args = ( + args[:model_index] + + config_args + + args[model_index:index] + + args[index + 2 :] + ) else: args = [args[0]] + config_args + args[1:index] + args[index + 2 :] From 7d06dd1ffaab13bb6b9a34e86a8a62e4d715bd1e Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 21 Sep 2026 18:54:57 -0500 Subject: [PATCH 607/774] [Pooling] MRV2 pooling shutdown model ref (#57737) Signed-off-by: Taneem Ibrahim --- vllm/v1/worker/gpu/model_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 92a37008cebe..8be523d3de15 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2245,6 +2245,7 @@ def shutdown(self) -> None: self.aux_output_connector.close() self.cudagraph_manager = None self.fast_prefill = None + self.pooling_runner = None if hasattr(self, "kv_caches"): self.kv_caches.clear() if hasattr(self, "attn_groups"): From 78d78365254dfa20d3a265b6da695ab8151ec8b1 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:19:17 -0400 Subject: [PATCH 608/774] [CI] Split (H200) LM Eval Large Models into per-model jobs (#57965) Signed-off-by: Thang Nguyen Co-authored-by: Kimi Code Co-authored-by: Claude Opus 5 (1M context) --- .buildkite/test_areas/lm_eval.yaml | 75 ++++++++++++++----- .../configs/models-h200-deepseek-r1-dp.txt | 1 + .../configs/models-h200-deepseek-r1-tp.txt | 1 + .../configs/models-h200-deepseek-v32-dp.txt | 1 + .../configs/models-h200-deepseek-v32-tp.txt | 1 + .../configs/models-h200-nemotron-3-super.txt | 2 + tests/evals/gsm8k/configs/models-h200.txt | 6 -- 7 files changed, 62 insertions(+), 25 deletions(-) create mode 100644 tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt create mode 100644 tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt create mode 100644 tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt create mode 100644 tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt create mode 100644 tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt delete mode 100644 tests/evals/gsm8k/configs/models-h200.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 5a284d172ab3..3450b70c2802 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -343,29 +343,66 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt -- label: ":nvidia: (H200) LM Eval Large Models" - key: lm-eval-large-models-8xh200 - timeout_in_minutes: 120 +- label: ":nvidia: (H200) LM Eval DeepSeek R1 TP" + key: lm-eval-deepseek-r1-tp-8xh200 + timeout_in_minutes: 35 device: h200 optional: true num_devices: 8 commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt - mirror: - amd: - label: ":amd: (MI300) LM Eval Large Models" - dind: false - device: mi300_8 - timeout_in_minutes: 65 - depends_on: - - image-build-amd - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - source_file_dependencies: - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-tp.txt + +- label: ":nvidia: (H200) LM Eval DeepSeek R1 DP" + key: lm-eval-deepseek-r1-dp-8xh200 + timeout_in_minutes: 35 + device: h200 + optional: true + num_devices: 8 + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-dp.txt + +- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 TP" + key: lm-eval-deepseek-v32-tp-8xh200 + timeout_in_minutes: 35 + device: h200 + optional: true + num_devices: 8 + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-tp.txt + +- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 DP" + key: lm-eval-deepseek-v32-dp-8xh200 + timeout_in_minutes: 35 + device: h200 + optional: true + num_devices: 8 + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-dp.txt + +- label: ":nvidia: (H200) LM Eval Nemotron 3 Super" + key: lm-eval-nemotron-3-super-8xh200 + timeout_in_minutes: 35 + device: h200 + optional: true + num_devices: 8 + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-nemotron-3-super.txt + +- label: ":amd: (MI300) LM Eval Large Models" + key: amd-lm-eval-large-models + dind: false + device: mi300_8 + optional: true + timeout_in_minutes: 65 + depends_on: + - image-build-amd + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt + source_file_dependencies: + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) MoE Refactor Integration TEMPORARY" key: moe-refactor-integration-test-h100-temporary diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt new file mode 100644 index 000000000000..9e293daf16f4 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt @@ -0,0 +1 @@ +DeepSeek-R1-DP.yaml diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt new file mode 100644 index 000000000000..30826a6f7698 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt @@ -0,0 +1 @@ +DeepSeek-R1-TP.yaml diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt new file mode 100644 index 000000000000..57faba312ffc --- /dev/null +++ b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt @@ -0,0 +1 @@ +DeepSeek-V3.2-DP.yaml diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt new file mode 100644 index 000000000000..027f168fd58e --- /dev/null +++ b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt @@ -0,0 +1 @@ +DeepSeek-V3.2-TP.yaml diff --git a/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt b/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt new file mode 100644 index 000000000000..ec6fc214cdb8 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt @@ -0,0 +1,2 @@ +Nemotron-3-Super-120B-A12B-BF16.yaml +Nemotron-3-Super-120B-A12B-FP8.yaml diff --git a/tests/evals/gsm8k/configs/models-h200.txt b/tests/evals/gsm8k/configs/models-h200.txt deleted file mode 100644 index 255e93b67443..000000000000 --- a/tests/evals/gsm8k/configs/models-h200.txt +++ /dev/null @@ -1,6 +0,0 @@ -DeepSeek-R1-TP.yaml -DeepSeek-R1-DP.yaml -DeepSeek-V3.2-TP.yaml -DeepSeek-V3.2-DP.yaml -Nemotron-3-Super-120B-A12B-BF16.yaml -Nemotron-3-Super-120B-A12B-FP8.yaml From c64b15cde583c7dd3cbb906b30164131a030d6b4 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:19:42 -0400 Subject: [PATCH 609/774] [Scheduler] Soften Long Prefill Tokens Threshhold (#57951) Signed-off-by: Robert Shaw Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw --- tests/v1/core/test_deferred_block_free.py | 2 +- tests/v1/core/test_scheduler.py | 30 +++++++++++++++++++ .../general/test_pooling_chunked_prefill.py | 8 ++--- vllm/config/scheduler.py | 5 +++- vllm/v1/core/sched/scheduler.py | 18 +++++++---- 5 files changed, 52 insertions(+), 11 deletions(-) diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py index ebbf492e208d..44cdf567f7c9 100644 --- a/tests/v1/core/test_deferred_block_free.py +++ b/tests/v1/core/test_deferred_block_free.py @@ -417,7 +417,7 @@ def test_abort_mid_prefill_defers_free(): chunk is still in flight must withhold its blocks. """ scheduler = create_scheduler( - model=MODEL, async_scheduling=True, long_prefill_token_threshold=16 + model=MODEL, async_scheduling=True, max_num_batched_tokens=16 ) scheduler.defer_block_free = True pool = scheduler.kv_cache_manager.block_pool diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index a58bcab7d5ea..e31163923b29 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -762,6 +762,36 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool): assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224 +def test_long_prefill_threshold_ignored_when_alone(): + """A lone long prefill is not capped by the threshold: it has no other + request to starve.""" + scheduler = create_scheduler( + max_num_batched_tokens=1024, + long_prefill_token_threshold=400, + ) + request = create_requests(num_requests=1, num_tokens=2000)[0] + scheduler.add_request(request) + + output = scheduler.schedule() + assert output.num_scheduled_tokens[request.request_id] == 1024 + + +def test_long_prefill_threshold_applies_with_other_requests(): + """The threshold caps the prefill as soon as another request is queued.""" + scheduler = create_scheduler( + max_num_batched_tokens=1024, + long_prefill_token_threshold=400, + ) + long_req = create_requests(num_requests=1, num_tokens=2000)[0] + short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0] + for request in [long_req, short_req]: + scheduler.add_request(request) + + output = scheduler.schedule() + assert output.num_scheduled_tokens[long_req.request_id] == 400 + assert output.num_scheduled_tokens[short_req.request_id] == 10 + + def test_update_from_output_routes_sampling_masks_by_request(): """Each request receives the sampler row at its own batch index.""" scheduler = create_scheduler() diff --git a/tests/v1/e2e/general/test_pooling_chunked_prefill.py b/tests/v1/e2e/general/test_pooling_chunked_prefill.py index e6b63eaf61aa..ab36e9427d98 100644 --- a/tests/v1/e2e/general/test_pooling_chunked_prefill.py +++ b/tests/v1/e2e/general/test_pooling_chunked_prefill.py @@ -77,13 +77,13 @@ def test_pooling_chunked_prefill(vllm_runner, monkeypatch): chunk_size = 10 - # Set chunking parameters to force chunked prefill - # Note: Chunked prefill is automatically handled by vLLM - # internally based on the model size and prompt + # Cap the token budget to force chunked prefill; the + # long_prefill_token_threshold is not applied to a lone request. with vllm_runner( model_id, runner="pooling", - long_prefill_token_threshold=chunk_size, + max_num_batched_tokens=chunk_size, + max_num_seqs=1, tensor_parallel_size=1, enforce_eager=True, enable_chunked_prefill=True, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 321c76fe3058..9185c26c4d85 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -79,7 +79,10 @@ class SchedulerConfig: long_prefill_token_threshold: int = Field(default=0, ge=0) """For chunked prefill, a request is considered long if the prompt is - longer than this number of tokens. 0 disables the cap (default).""" + longer than this number of tokens. 0 disables the cap (default). + + The cap is not applied when the request is the only one in the batch, + since there is no other request for it to starve.""" max_num_queued_reqs: int | None = Field(default=None, ge=0) """Maximum number of requests that can be in-flight (waiting or running) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 36a09f9adcc4..18bbc102830b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -599,6 +599,15 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: throttle_prefills and not self.prefill_capacity_bound ) and any(not r.is_prefill_chunk for r in self.running) + # `long_prefill_token_threshold` exists to stop a long prefill from + # starving other requests of the token budget. When it is the only + # request there is nobody to starve, so let it use the whole budget. + long_prefill_token_threshold = ( + self.scheduler_config.long_prefill_token_threshold + if len(self.running) + len(self.waiting) + len(self.skipped_waiting) > 1 + else 0 + ) + # First, schedule the RUNNING requests. req_index = 0 while req_index < len(self.running) and token_budget > 0: @@ -650,8 +659,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + request.num_output_placeholders - request.num_computed_tokens ) - if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold + if 0 < long_prefill_token_threshold < num_new_tokens: + num_new_tokens = long_prefill_token_threshold num_new_tokens = min( num_new_tokens, token_budget, input_budget - draft_slots ) @@ -1090,9 +1099,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_new_tokens = padded_num_tokens pad_spec_decode = True - threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold + if 0 < long_prefill_token_threshold < num_new_tokens: + num_new_tokens = long_prefill_token_threshold # chunked prefill has to be enabled explicitly to allow # pooling requests to be chunked From a369a7becc20c6c4dda601aedef51b8cc8387971 Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Mon, 21 Sep 2026 17:24:25 -0700 Subject: [PATCH 610/774] [Perf][Attention] Reduce GLM sparse MLA preparation overhead (#57458) Signed-off-by: Summer Yang Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- .../attention/test_flashinfer_mla_decode.py | 152 ++++++++++++------ tests/kernels/test_concat_mla_q.py | 29 ++++ .../v1/attention/test_indexer_dcp_localize.py | 7 +- .../v1/attention/test_sparse_mla_backends.py | 25 ++- .../layers/attention/mla_attention.py | 5 +- .../backends/mla/flashinfer_mla_sparse.py | 34 ++-- .../v1/attention/backends/mla/sparse_utils.py | 90 ++++++++--- 7 files changed, 248 insertions(+), 94 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index 0d3791964398..c3f7b4c96f78 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -130,69 +130,123 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int): torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("mode", ["tp", "dcp", "hisparse"]) @requires_sm10x -def test_flashinfer_trtllm_sparse_mla_decode_without_rope(): - """The native sparse MLA path supports a zero-width rotary tail.""" +def test_flashinfer_trtllm_sparse_mla_decode_without_rope(dtype, mode, monkeypatch): + """TP skips preparation; empty DCP contributions and HiSparse outputs vanish.""" + import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as backend + from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup + torch.set_default_device("cuda") torch.manual_seed(42) - - batch_size = 2 - block_size = 64 - num_blocks = 4 - sparse_topk = 128 - valid_lens = torch.tensor([17, 73], dtype=torch.int32) - - query = torch.randn( - batch_size, - 1, - NUM_HEADS, - KV_LORA_RANK, - dtype=torch.bfloat16, - ) - kv_cache = torch.randn( - num_blocks, - block_size, - KV_LORA_RANK, - dtype=torch.bfloat16, - ) - + batch_size, num_heads, block_size, num_blocks, sparse_topk = 3, 8, 64, 4, 2176 + valid_lens = [1 if mode == "tp" else 0, 17, 73] + counts = torch.tensor(valid_lens, dtype=torch.int32) + query = torch.randn(batch_size, 1, num_heads, KV_LORA_RANK).to(dtype) + kv_cache = torch.randn(num_blocks, block_size, KV_LORA_RANK).to(dtype) num_slots = num_blocks * block_size - slot_tables = torch.stack( - [torch.randperm(num_slots)[:sparse_topk] for _ in range(batch_size)] - ).to(torch.int32) - for row, valid_len in zip(slot_tables, valid_lens.tolist()): - row[valid_len:] = -1 - - workspace_buffer = torch.empty( - FLASHINFER_WORKSPACE_BUFFER_SIZE, - dtype=torch.int8, + indices = torch.full((batch_size, sparse_topk), -1, dtype=torch.int32) + for token, length in enumerate(valid_lens): + indices[token, :length] = torch.randperm(num_slots)[:length].int() + empty = counts == 0 + reference_indices = indices.clone() + reference_indices[:, 0].masked_fill_(empty, 0) + workspace = torch.empty(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.int8) + + impl = object.__new__(backend.FlashInferMLASparseImpl) + impl.qk_nope_head_dim = QK_NOPE_HEAD_DIM + impl.kv_lora_rank = KV_LORA_RANK + impl.qk_rope_head_dim = 0 + impl.index_group = ( + object.__new__(HiSparseMLAIndexGroup) if mode == "hisparse" else None ) - out = trtllm_batch_decode_with_kv_cache_mla( + impl._workspace_buffer = workspace + impl.bmm1_scale = QK_NOPE_HEAD_DIM**-0.5 + impl.bmm2_scale = 1.0 + impl.is_nope_mla = True + impl.need_to_return_lse_for_decode = mode == "dcp" + if mode == "tp": + + def unexpected_helper(*args, **kwargs): + pytest.fail("TP must not launch empty-query preparation") + + monkeypatch.setattr( + backend, "prepare_sparse_mla_safe_lengths", unexpected_helper + ) + + out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts) + reference = trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache.unsqueeze(1), - workspace_buffer=workspace_buffer, + workspace_buffer=workspace, qk_nope_head_dim=QK_NOPE_HEAD_DIM, kv_lora_rank=KV_LORA_RANK, qk_rope_head_dim=0, - block_tables=slot_tables.unsqueeze(1), - seq_lens=valid_lens, + block_tables=reference_indices.unsqueeze(1), + seq_lens=counts, max_seq_len=sparse_topk, sparse_mla_top_k=sparse_topk, - sparse_mla_top_k_lens=valid_lens, + sparse_mla_top_k_lens=counts.clamp(min=1), bmm1_scale=QK_NOPE_HEAD_DIM**-0.5, bmm2_scale=1.0, - ).squeeze(1) - - flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float() - refs = [] - for batch_idx, valid_len in enumerate(valid_lens.tolist()): - selected_kv = flat_cache[slot_tables[batch_idx, :valid_len].long()] - scores = torch.einsum("hd,kd->hk", query[batch_idx, 0].float(), selected_kv) - probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1) - refs.append(torch.einsum("hk,kd->hd", probs, selected_kv)) - ref = torch.stack(refs).to(torch.bfloat16) - - torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + return_lse=mode == "dcp", + ) + if mode == "dcp": + reference, reference_lse = reference + reference_lse = reference_lse.reshape(batch_size, num_heads) + reference_lse.masked_fill_(empty[:, None], float("-inf")) + torch.testing.assert_close(lse, reference_lse, atol=0, rtol=0) + + from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out + + # The peer owns only our empty selections. Check our contribution after + # the production merge, without zeroing the local attention output. + peer_lse = torch.full_like(lse, float("-inf")).masked_fill(empty[:, None], 0) + out, _ = correct_attn_out( + out, torch.stack((lse, peer_lse)), cp_rank=0, ctx=CPTritonContext() + ) + else: + assert lse is None + reference = reference.squeeze(1).masked_fill(empty[:, None, None], 0) + torch.testing.assert_close(out, reference, atol=0, rtol=0) + + if dtype == torch.bfloat16: + flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float() + refs = [] + for token, length in enumerate(valid_lens): + selected_kv = flat_cache[reference_indices[token, :length].long()] + scores = torch.einsum("hd,kd->hk", query[token, 0].float(), selected_kv) + probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1) + refs.append(torch.einsum("hk,kd->hd", probs, selected_kv)) + ref = torch.stack(refs).to(torch.bfloat16) + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + + +@requires_sm10x +def test_sparse_mla_safe_lengths(): + from vllm.v1.attention.backends.mla.sparse_utils import ( + prepare_sparse_mla_safe_lengths, + ) + + device = "cuda" + num_tokens = 259 + # Include a partial tile and expose accidental writes into adjacent storage. + counts_storage = torch.full((num_tokens, 2), 17, dtype=torch.int32, device=device) + counts = counts_storage[:, 0] + counts[::3] = 0 + original_counts = counts_storage.clone() + empty = counts == 0 + indices_storage = torch.full( + (num_tokens, 256), 42, dtype=torch.int32, device=device + ) + indices = indices_storage[:, ::2] + expected_indices = indices_storage.clone() + expected_indices[:, 0].masked_fill_(empty, 0) + safe = prepare_sparse_mla_safe_lengths(indices, counts) + torch.testing.assert_close(safe, counts.clamp(min=1), rtol=0, atol=0) + torch.testing.assert_close(indices_storage, expected_indices, rtol=0, atol=0) + torch.testing.assert_close(counts_storage, original_counts, rtol=0, atol=0) @requires_sm90 diff --git a/tests/kernels/test_concat_mla_q.py b/tests/kernels/test_concat_mla_q.py index 6ace34e503e4..39c2742e4f8a 100644 --- a/tests/kernels/test_concat_mla_q.py +++ b/tests/kernels/test_concat_mla_q.py @@ -148,3 +148,32 @@ def test_concat_mla_q_values_preserved(num_tokens): assert torch.equal(out_bits[..., :nope_dim], ql_nope_bits) assert torch.equal(out_bits[..., nope_dim:], q_pe_bits) + + +@pytest.mark.parametrize("rope_dim,transposed", [(0, False), (0, True), (64, False)]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled): + """FP8 query packing preserves values with empty RoPE and strided BMM output.""" + from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.attention.mla_attention import _DecodeConcatQuantFP8 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + torch.manual_seed(42) + shape = (8, 17, 512) if transposed else (17, 8, 512) + ql = torch.randn(shape, dtype=torch.bfloat16, device="cuda") * 128 + if transposed: + ql = ql.transpose(0, 1) + qp = torch.randn(17, 8, rope_dim, dtype=torch.bfloat16, device="cuda") * 128 + scale = torch.tensor(0.3, device="cuda", dtype=torch.float32) + joined = torch.cat((ql, qp), dim=-1) + expected = (joined.float() * scale.reciprocal()).clamp(-448, 448) + expected = expected.to(torch.float8_e4m3fn) + custom_ops = ["none"] + (["+mla_decode_concat_quant_fp8"] if enabled else []) + config = VllmConfig( + compilation_config=CompilationConfig(mode=0, custom_ops=custom_ops) + ) + with set_current_vllm_config(config): + quant = _DecodeConcatQuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + actual = quant(ql, qp, scale) + torch.testing.assert_close(actual.float(), expected.float(), atol=0, rtol=0) + assert actual.is_contiguous() diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 8bc409adb704..b099159bc766 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -700,9 +700,8 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel( @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("interleave", [1, 2]) @pytest.mark.parametrize("dcp_rank", [0, 1]) -# 384 is not a power of two, so it exercises the multi-tile atomic allocator -# rather than the single-tile path 1024 takes. -@pytest.mark.parametrize("num_topk", [1024, 384]) +# Include padded single-tile compaction and the multi-tile atomic fallback. +@pytest.mark.parametrize("num_topk", [1024, 2176, 4224]) def test_dcp_filter_compaction_matches_reference( interleave: int, dcp_rank: int, num_topk: int ): @@ -729,7 +728,7 @@ def test_dcp_filter_compaction_matches_reference( for r in range(num_rows): # Ids are distinct, so a row holds at most `seq` valid entries. hi = min(num_topk * 3 // 4, seq) - n_valid = int(torch.randint(num_topk // 4, hi, (1,)).item()) + n_valid = int(torch.randint(hi // 3, hi, (1,)).item()) perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32) token_indices[r, :n_valid] = perm diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 7704600d09d4..1d7fc389ef69 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1663,10 +1663,12 @@ def test_split_indexer_prefill_chunks_single_request_overflow(): assert out == expected -# 384 is not a power of two, so it counts via the tiled atomic accumulation -# rather than the single-tile path 128 takes. -@pytest.mark.parametrize("num_topk_tokens", [128, 384]) -def test_triton_convert_returns_valid_counts(num_topk_tokens: int): +# Power-of-two, GLM's padded tile, and atomic fallback, with reused buffers. +@pytest.mark.parametrize( + "num_topk_tokens,reuse_buffers", + [(128, False), (2176, False), (2176, True), (4224, False), (4224, True)], +) +def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers: bool): """Test that return_valid_counts correctly counts non-negative indices.""" device = torch.device(DEVICE_TYPE) num_tokens = 8 @@ -1679,8 +1681,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): num_requests * max_blocks_per_req, dtype=torch.int32, device=device ).view(num_requests, max_blocks_per_req) - # Create token indices with varying numbers of valid entries: half the row, - # a quarter of it, the whole row, then a single valid entry -- twice over. + # Include partly valid, fully valid, and empty selections for each request. token_indices = torch.full( (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device ) @@ -1688,7 +1689,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): num_topk_tokens // 2, num_topk_tokens // 4, num_topk_tokens, - 1, + 0, ] * 2 expected_valid = [] for i in range(num_tokens): @@ -1702,6 +1703,15 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): expected_valid, dtype=torch.int32, device=device ) + buffers = {} + if reuse_buffers: + buffers = dict( + out=torch.full_like(token_indices, -99), + valid_counts_out=torch.full( + (num_tokens,), 99, dtype=torch.int32, device=device + ), + ) + # Test with return_valid_counts=True result, valid_counts = triton_convert_req_index_to_global_index( req_id, @@ -1710,6 +1720,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): BLOCK_SIZE=block_size, NUM_TOPK_TOKENS=num_topk_tokens, return_valid_counts=True, + **buffers, ) torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 3104ac25c6ef..c37801e88716 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1553,7 +1553,10 @@ def forward( scale: torch.Tensor, scale_ub: torch.Tensor | None = None, ) -> torch.Tensor: - decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1) + if decode_q_pe.shape[-1] == 0 and decode_ql_nope.is_contiguous(): + decode_q0 = decode_ql_nope + else: + decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1) decode_q_flat = decode_q0.reshape(decode_q0.shape[0], -1) decode_q, _ = quant_fn(self, decode_q_flat, scale, scale_ub) return decode_q.view(decode_q0.shape) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index ba10a2e91cea..142a63ec7234 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -30,6 +30,7 @@ from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, + prepare_sparse_mla_safe_lengths, triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) @@ -697,18 +698,16 @@ def _run_mqa_kernel( sparse_topk_capacity = topk_indices.shape[1] extra_kwargs: dict[str, torch.Tensor] = {} - empty_rows: torch.Tensor | None = None + needs_empty_query_guard = self.need_to_return_lse_for_decode or isinstance( + self.index_group, HiSparseMLAIndexGroup + ) if self.is_nope_mla: - # The native no-rope kernel takes the active top-k length per query - # token (``seq_lens`` here is already the compacted per-token valid - # count, int32) and rejects zero-length rows. Point empty rows at a - # single valid dummy slot with length 1 and zero their output after - # the launch. ``triton_convert_req_index_to_global_index`` packs the - # valid indices into a contiguous prefix, which is what the kernel - # requires of the page table. - empty_rows = seq_lens == 0 - topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0) - extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1) + # Resident TP queries have nonempty selections. Preserve empty-query + # handling for DCP's local selections and host-backed HiSparse. + topk_lens = seq_lens + if needs_empty_query_guard: + topk_lens = prepare_sparse_mla_safe_lengths(topk_indices, seq_lens) + extra_kwargs["sparse_mla_top_k_lens"] = topk_lens kernel_out = trtllm_batch_decode_with_kv_cache_mla( query=query, @@ -737,12 +736,17 @@ def _run_mqa_kernel( out = o.view(-1, o.shape[-2], o.shape[-1]) if lse is not None: lse = self._normalize_lse(lse, out.shape[0], out.shape[1]) - if empty_rows is None and lse is not None: + if self.is_nope_mla and needs_empty_query_guard: + empty_queries = seq_lens == 0 + if lse is not None: + # DCP combine already suppresses outputs with zero LSE weight. + lse.masked_fill_(empty_queries[:, None], float("-inf")) + else: + out.masked_fill_(empty_queries[:, None, None], 0.0) + elif lse is not None: empty_rows = (topk_indices == -1).all(dim=-1) - if empty_rows is not None: out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) - if lse is not None: - lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) return out, lse @staticmethod diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 0ac131cee5df..2ec10218d99f 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -86,7 +86,7 @@ class CompileKey: @staticmethod @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"]) - def kernel( + def sparse_mla_index_remap_kernel( req_id_ptr, # int32 [num_tokens] block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] @@ -100,15 +100,16 @@ def kernel( BLOCK_SIZE: tl.constexpr, BLOCK_STRIDE_ROWS: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns + NUM_TOPK_TOKENS: tl.constexpr, HAS_PREFILL: tl.constexpr, COUNT_VALID: tl.constexpr, # whether to count valid indices - # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count + # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count # is an in-register reduction and needs no atomic. SINGLE_TILE: tl.constexpr, # When set, scatter valid slots to a contiguous prefix [0, valid_count) using # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries). - # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the + # Requires COUNT_VALID; multi-tile output must be pre-filled with -1. The # prefix is unspecified (only the selected set matters). COMPACT_TO_FRONT: tl.constexpr, # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op @@ -136,7 +137,7 @@ def kernel( # Load token indices for this tile ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 - tok = tl.load(ti_ptr) # int32 + tok = tl.load(ti_ptr, mask=indice_id < NUM_TOPK_TOKENS, other=-1) # Only token == -1 should propagate as -1 is_invalid_tok = tok < 0 @@ -197,10 +198,19 @@ def kernel( dest = base + local_offset out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1 tl.store(out_ptr_dest, out_val, mask=is_valid == 1) + if SINGLE_TILE: + # This program owns both the compacted prefix and its padding. + # The stores are disjoint, so no separate initialization is needed. + tl.store( + out_ptr + token_id * out_stride0 + indice_id * out_stride1, + -1, + mask=(indice_id >= tile_valid_count) + & (indice_id < NUM_TOPK_TOKENS), + ) else: # Store results in place (input column == output column). out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 - tl.store(out_ptr_ij, out_val) + tl.store(out_ptr_ij, out_val, mask=indice_id < NUM_TOPK_TOKENS) # Accumulate the tile's valid count into the row total; a single tile's # reduction *is* the total. @@ -211,6 +221,9 @@ def kernel( else: tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + # Keep the warmup wrapper interface while exposing a descriptive CUDA symbol. + kernel = sparse_mla_index_remap_kernel + def dispatch( # type: ignore[override] self, *, @@ -421,18 +434,19 @@ def _remap_tiling( Counting the valid slots per row is the only reason the column tiles have to talk to each other, so when counting give one program the whole row: the count becomes an in-register reduction plus a plain store, needing neither - atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this - needs a power-of-two width; other top-k sizes stay tiled and atomic. + atomics nor a zero-initialized counter. Pad modest non-power-of-two widths + (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic. Returns: (single_tile, block_n, tiles_per_row, num_warps) """ - single_tile = ( - count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS + padded_width = triton.next_power_of_2(NUM_TOPK_TOKENS) + single_tile = count_valid and ( + padded_width == NUM_TOPK_TOKENS or padded_width <= 4096 ) if single_tile: - return True, NUM_TOPK_TOKENS, 1, 8 + return True, padded_width, 1, 8 return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4 @@ -504,14 +518,11 @@ def triton_convert_req_index_to_global_index( req_id_c = req_id.contiguous() block_table_c = block_table.contiguous() token_indices_c = token_indices.contiguous() - # When return_valid_counts, the kernel scatters valid entries to a - # contiguous prefix [0, valid_count) and leaves the tail unwritten, so - # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to - # [:topk_length] == exactly the valid set (no dropped tokens). + # Only multi-tile compaction needs separate padding initialization. if out is None: out = ( torch.full_like(token_indices_c, -1) - if return_valid_counts + if return_valid_counts and not single_tile else torch.empty_like(token_indices_c) ) else: @@ -519,7 +530,7 @@ def triton_convert_req_index_to_global_index( assert out.device == token_indices_c.device assert out.shape == token_indices_c.shape assert out.is_contiguous() - if return_valid_counts: + if return_valid_counts and not single_tile: out.fill_(-1) valid_counts: torch.Tensor | None = None @@ -636,13 +647,13 @@ def triton_filter_and_convert_dcp_index( token_indices_c = token_indices.contiguous() # The compaction uses the valid-count buffer as a slot allocator, so it - # requires counting. Pre-fill out with -1 so the unwritten tail stays -1. + # requires counting. Only the multi-tile path needs pre-filled padding. count_valid = return_valid_counts or compact_valid_to_front # The compaction builds on the counting, so it shares the tiling. single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid) - if compact_valid_to_front: + if compact_valid_to_front and not single_tile: out = torch.full_like(token_indices_c, -1) else: out = torch.empty_like(token_indices_c) @@ -685,3 +696,46 @@ def triton_filter_and_convert_dcp_index( _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel() + + +@triton.jit(do_not_specialize=["num_tokens"]) +def sparse_mla_prepare_safe_lengths_kernel( + indices_ptr, + counts_ptr, + safe_lengths_ptr, + num_tokens, + index_stride0: tl.constexpr, + count_stride: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + count = tl.load(counts_ptr + token * count_stride, token < num_tokens, other=1) + tl.store(safe_lengths_ptr + token, tl.maximum(count, 1), token < num_tokens) + # TRTLLM's NoPE sparse kernel requires at least one valid KV-cache entry. + tl.store( + indices_ptr + token * index_stride0, 0, (token < num_tokens) & (count == 0) + ) + + +def prepare_sparse_mla_safe_lengths( + physical_indices: torch.Tensor, valid_counts: torch.Tensor +) -> torch.Tensor: + """Install dummy slots for empty queries and return nonzero kernel lengths. + + Preserve the raw counts so empty outputs and LSE can be neutralized later. + """ + num_tokens = valid_counts.numel() + safe_lengths = torch.empty( + (num_tokens,), dtype=valid_counts.dtype, device=valid_counts.device + ) + if num_tokens: + sparse_mla_prepare_safe_lengths_kernel[(triton.cdiv(num_tokens, 256),)]( + physical_indices, + valid_counts, + safe_lengths, + num_tokens, + physical_indices.stride(0), + valid_counts.stride(0), + BLOCK=256, + ) + return safe_lengths From 0bce411a073e0da2c9bded9bb6ced1cbee51c564 Mon Sep 17 00:00:00 2001 From: Navjot Singh Date: Mon, 21 Sep 2026 20:31:26 -0400 Subject: [PATCH 611/774] [Bugfix] Annotate MTP draft KV cache groups positionally on the hybrid grouping path (#55390) Signed-off-by: Navjot Singh Co-authored-by: Codex --- tests/v1/core/test_kv_cache_utils.py | 96 ++++++++++++++++++++++++++-- vllm/v1/core/kv_cache_utils.py | 58 ++++++++++++----- 2 files changed, 133 insertions(+), 21 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 07c222214586..bc2a761d86ec 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -4275,12 +4275,16 @@ def _deepseek_v4_specs(model_version="deepseek_v4"): } -def test_deepseek_v4_draft_group_annotated_on_packed_path(): +@pytest.mark.parametrize( + ("method", "model_type"), + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], +) +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): # DeepseekV4's MTP block reuses the target's decoder layer, so its spec # carries no draft marker and only the positional rule can find it. This # pins the pre-existing behaviour that the unified annotator must preserve. groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="deepseek_v4"), + _spec_decode_grouping_config(method=method, model_type=model_type), _deepseek_v4_specs(model_version=None), ) @@ -4289,13 +4293,93 @@ def test_deepseek_v4_draft_group_annotated_on_packed_path(): assert "model.layers.3.self_attn.attn" in flagged[0].layer_names -def test_deepseek_v4_annotation_requires_model_type(): - # The positional rule is only sound for DeepseekV4, where the draft layer - # is known to be registered last. Without that model gate nothing may be - # flagged, however the grouping happens to fall out. +def test_trailing_layer_fallback_applies_to_any_mtp_model(): + # The positional rule is sound for every MTP drafter, not just DeepseekV4: + # MTP blocks reuse the target's decoder layer (no spec marker) and always + # register after every target layer. The model_type must not gate it. groups = get_kv_cache_groups( _spec_decode_grouping_config(method="mtp", model_type="other"), _deepseek_v4_specs(), ) + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def _qwen3_5_hybrid_specs(with_mtp_layer: bool): + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with + the MTP drafter's full-attn layer (spec-identical to the target's) + registered last.""" + specs = {} + idx = 0 + for _ in range(2): + for _ in range(3): + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( + mamba_cache_mode="align" + ) + idx += 1 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() + idx += 1 + if with_mtp_layer: + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() + return specs + + +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): + # A hybrid mamba + full-attention model with an MTP drafter that is + # spec-indistinguishable from the target reaches the general multi-group + # path. The trailing-layer rule must locate the draft group there so the + # Mamba groups are not swept up by the flag-all consumer fallback. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names + for group in groups: + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + assert not group.is_eagle_group + assert "could be identified as the draft model's" not in caplog_vllm.text + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): + # Other drafters are not covered by the trailing-layer rule, so an + # unidentifiable hybrid draft must still warn. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" in caplog_vllm.text + + +def test_trailing_layer_fallback_requires_exact_partition(): + # If the groups do not partition the layers exactly (e.g. a caller that + # dropped or duplicated layers), the positional rule is meaningless and + # must not fire. + from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups + + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") + groups = get_kv_cache_groups(config, specs) + for g in groups: + g.is_eagle_group = False + # Remove one layer from its group: no longer an exact partition. + trimmed = [ + KVCacheGroupSpec( + [n for n in g.layer_names if n != "model.layers.0.linear_attn"], + g.kv_cache_spec, + ) + for g in groups + ] + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) + + assert not any(g.is_eagle_group for g in trimmed) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 656765c89604..af1112c54aaf 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2117,16 +2117,24 @@ def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: vllm_config, kv_cache_spec, groups, - use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config), + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), ) _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups -def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: + """Whether the drafter's KV layers can be located positionally. + + MTP drafters register their KV layers after the target layers, without a + distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. + The annotator separately checks that the groups partition the layers exactly. + """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle(): return False + if spec_config.method == "mtp": + return True model_config = vllm_config.model_config return model_config is not None and model_config.hf_config.model_type in ( "deepseek_v4", @@ -2134,11 +2142,24 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: ) +def _groups_partition_layers_exactly( + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. + + The trailing-layer draft fallback is only meaningful when the last + registered layer is guaranteed to appear in exactly one group. + """ + names = [name for group in kv_cache_groups for name in group.layer_names] + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) + + def _annotate_eagle_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], kv_cache_groups: list[KVCacheGroupSpec], - use_deepseek_v4_fallback: bool = False, + use_trailing_layer_fallback: bool = False, ) -> None: """Flag the KV cache groups that hold drafter attention layers. @@ -2151,14 +2172,14 @@ def _annotate_eagle_groups( spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block - reuses the target's own decoder layer and so carries no spec marker. Its - draft attention layer is always the last registered layer, so flag whichever - group holds it. This rule is only valid where the groups partition - exactly the layers of ``kv_cache_spec``, which is true on the packed - grouping path and not in general; other callers must leave - ``use_deepseek_v4_fallback`` False. The caller gates this fallback on - the configured model type. + 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose + MTP block reuses the target's own decoder layer and so carries no spec + marker. Their draft attention layers always register after every + target layer, so flag whichever group holds the last registered layer. + This rule is only valid where the groups partition exactly the layers + of ``kv_cache_spec``, which is re-checked here before applying it. + When a drafter's trailing caches span several groups, this rule flags + only the group holding the very last layer and must be generalized. FIXME(yifan): avoid/generalize this hacky check. Args: @@ -2166,8 +2187,8 @@ def _annotate_eagle_groups( kv_cache_spec: The kv cache spec of each attention layer, in layer registration order. Only read by rule 2. kv_cache_groups: Groups to annotate in place. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed - group. + use_trailing_layer_fallback: Enable rule 2. Callers gate this on + ``_uses_trailing_mtp_layers``. """ spec_config = vllm_config.speculative_config @@ -2181,7 +2202,9 @@ def _annotate_eagle_groups( ): group.is_eagle_group = True - if not use_deepseek_v4_fallback: + if not use_trailing_layer_fallback: + return + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): return last_layer = next(reversed(kv_cache_spec)) for group in kv_cache_groups: @@ -2352,7 +2375,12 @@ def get_kv_cache_groups( aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) groups.append(KVCacheGroupSpec([name], aligned)) - _annotate_eagle_groups(vllm_config, kv_cache_spec, groups) + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups From 79468c20ef23227d4e051f807e10fd54fb24e24f Mon Sep 17 00:00:00 2001 From: cherry77-cloud <1615405@qq.com> Date: Tue, 22 Sep 2026 09:11:49 +0800 Subject: [PATCH 612/774] [Bugfix][GDN] Fix stateless first-chunk classification (#51565) Signed-off-by: taking-lying-flat <1615405@qq.com> Signed-off-by: zjy0516 Co-authored-by: zjy0516 --- .../v1/attention/test_gdn_metadata_builder.py | 83 +++++++++++++++++++ vllm/v1/attention/backends/gdn_attn.py | 27 +++++- 2 files changed, 109 insertions(+), 1 deletion(-) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 221f933d8942..b4d3acf07cbb 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -221,3 +221,86 @@ def test_full_cudagraph_spec_metadata_uses_request_count(): assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,) assert meta.num_accepted_tokens is not None assert meta.num_accepted_tokens.shape == (batch.batch_size,) + + +def _build_non_spec( + batch: BatchSpec, + is_prefilling: list[bool] | None, + full_cuda_graph: bool = False, +): + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace( + is_prefilling=None + if is_prefilling is None + else torch.tensor(is_prefilling, dtype=torch.bool) + ) + builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph) + return builder, common_attn_metadata, builder.build(0, common_attn_metadata) + + +@pytest.mark.parametrize( + ("seq_len", "query_len", "is_prefilling", "num_prefills"), + [ + pytest.param(1, 1, True, 1, id="first-chunk"), + pytest.param(65, 1, True, 0, id="resumed-chunk"), + pytest.param(0, 0, True, 0, id="padding"), + pytest.param(1, 1, None, 0, id="missing-prefill-flag"), + ], +) +def test_one_token_chunk_classification( + seq_len: int, + query_len: int, + is_prefilling: bool | None, + num_prefills: int, +): + """Only a real first chunk with a prefill flag needs state initialization.""" + _, _, meta = _build_non_spec( + BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]), + is_prefilling=None if is_prefilling is None else [False, is_prefilling], + ) + + assert meta.num_prefills == num_prefills + assert meta.num_decodes == 2 - num_prefills + assert meta.num_prefill_tokens == num_prefills + assert meta.num_decode_tokens == 1 + query_len - num_prefills + if num_prefills: + assert meta.has_initial_state is not None + assert meta.has_initial_state.tolist() == [True, False] + else: + assert meta.has_initial_state is None + + +def test_one_token_first_chunk_excludes_padding(): + """Neither padding requests nor padding tokens count as prefill work.""" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]), + BLOCK_SIZE, + DEVICE, + ).replace( + is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool), + num_actual_tokens=4, + ) + meta = _create_gdn_builder().build(0, common) + + assert meta.num_decodes == 1 + assert meta.num_prefills == 1 + assert meta.num_decode_tokens == 1 + assert meta.num_prefill_tokens == 1 + + +def test_cudagraph_capture_batch_stays_decode_only(): + """Capture rows have no history, but must still select decode kernels.""" + batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4) + builder, common_attn_metadata, _ = _build_non_spec( + batch, [False] * 4, full_cuda_graph=True + ) + meta = builder.build_for_cudagraph_capture(common_attn_metadata) + + assert meta.num_prefills == 0 + assert meta.num_decodes == 4 + assert meta.has_initial_state is None + staged = meta.non_spec_state_indices_tensor + assert staged is not None + assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr() + torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0]) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 7a839ff73b4a..c2cfd003d90d 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -247,9 +247,34 @@ def build( # type: ignore[override] ) if spec_sequence_masks is None: + # V2 already excludes prefills from full decode graphs via + # has_prefill. Classify first chunks as prefills to mask recycled + # state; resumed one-token chunks can still use the decode kernels. + assert m.seq_lens_cpu_upper_bound is not None + query_lens_cpu = query_start_loc_cpu.diff() + no_prior_state = (query_lens_cpu > 0) & ( + m.seq_lens_cpu_upper_bound <= query_lens_cpu + ) + # Capture batches also have seq_len == query_len, but are not + # prefills. + if m.is_prefilling is not None: + no_prior_state &= m.is_prefilling + else: + no_prior_state = torch.zeros_like(no_prior_state) num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills(m, decode_threshold=1) + split_decodes_and_prefills( + m.replace(is_prefilling=no_prior_state), + decode_threshold=1, + treat_short_extends_as_decodes=False, + ) ) + # Exclude trailing padding from both prefill counts. + if num_prefills: + num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum()) + num_prefill_tokens = ( + int(query_start_loc_cpu[num_decodes + num_prefills]) + - num_decode_tokens + ) num_spec_decode_tokens = 0 spec_token_indx = None non_spec_token_indx = None From e9f169d16b9408bb9ae44f75072b91a5521d733c Mon Sep 17 00:00:00 2001 From: Turner Jabbour Date: Mon, 21 Sep 2026 19:36:45 -0600 Subject: [PATCH 613/774] [CI] Add pre-commit check that new tests are tethered to Buildkite jobs (#54867) Signed-off-by: Turner Jabbour Signed-off-by: Turner Co-authored-by: Claude Sonnet 5 Co-authored-by: Kevin H. Luu Co-authored-by: Nick Hill --- .buildkite/test_areas/misc.yaml | 4 + .pre-commit-config.yaml | 8 + docker/Dockerfile.cpu | 4 + tests/tools/test_check_test_tethering.py | 879 +++++++++++++++++ tools/__init__.py | 0 tools/pre_commit/__init__.py | 0 tools/pre_commit/check_test_tethering.py | 925 ++++++++++++++++++ tools/pre_commit/test_tethering_allowlist.txt | 100 ++ 8 files changed, 1920 insertions(+) create mode 100644 tests/tools/test_check_test_tethering.py create mode 100644 tools/__init__.py create mode 100644 tools/pre_commit/__init__.py create mode 100644 tools/pre_commit/check_test_tethering.py create mode 100644 tools/pre_commit/test_tethering_allowlist.txt diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 759f3900c5ea..6709dbb56585 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -525,6 +525,9 @@ steps: - tests/renderers - tests/standalone_tests/lazy_imports.py - tests/reasoning + - tools/pre_commit/check_test_tethering.py + - tools/pre_commit/test_tethering_allowlist.txt + - tests/tools/test_check_test_tethering.py - tests/tool_parsers - tests/tokenizers_ - tests/parser @@ -552,6 +555,7 @@ steps: commands: - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s config + - pytest -v -s tools/test_check_test_tethering.py - label: ":computer: (CPU) Reasoning + Renderers" key: cpu-reasoning-renderers diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b12c8ce508ea..4823f788dc3a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -184,6 +184,14 @@ repos: files: ^\.buildkite/.*\.ya?ml$ pass_filenames: false additional_dependencies: [PyYAML] + - id: check-test-tethering + name: Check new tests are wired into Buildkite CI + language: python + entry: python tools/pre_commit/check_test_tethering.py + files: ^(tests/.*\.py|\.buildkite/(test_areas|hardware_tests|intel_jobs)/.*\.yaml|\.buildkite/(test-amd|ci_config[^/]*)\.yaml|tools/pre_commit/test_tethering_allowlist\.txt)$ + additional_dependencies: [PyYAML, regex] + verbose: true # full-scan `note:` lines are advisory, so the hook passes + require_serial: true # one full scan, not one per pre-commit file batch - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward name: Run mypy for Python 3.10 entry: python tools/pre_commit/mypy.py "3.10" diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 7898e57cf39a..d61e186dac4b 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -471,6 +471,10 @@ ADD ./benchmarks/ ./benchmarks/ ADD ./vllm/collect_env.py . ADD ./docker/ ./docker/ ADD ./.buildkite/ ./.buildkite/ +# tools/ carries repo-tooling code that a few tests exercise directly +# (e.g. tests/tools/test_check_test_tethering.py imports +# tools.pre_commit.check_test_tethering). +ADD ./tools/ ./tools/ # install development dependencies (for testing) RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/tests/tools/test_check_test_tethering.py b/tests/tools/test_check_test_tethering.py new file mode 100644 index 000000000000..34f4fd179c8e --- /dev/null +++ b/tests/tools/test_check_test_tethering.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the test-tethering pre-commit check. + +The check answers one question - "does any Buildkite job (in +``.buildkite/test_areas/`` or the legacy ``.buildkite/test-amd.yaml``) collect +this test file?" - by parsing job ``commands`` as shell. The risk that +matters is a *false tether*: a misparse that invents coverage makes the gate +pass while the test still never runs, which is the exact failure the check +exists to catch. So the parser cases below pin both directions, and +:func:`test_no_selection_matches_everything` guards the catastrophic version of +it. +""" + +import textwrap +from pathlib import Path + +import pytest + +import tools.pre_commit.check_test_tethering as checker +from tools.pre_commit.check_test_tethering import ( + FindSelection, + PytestSelection, + _parse_command, + _to_repo_relative, + all_test_modules, + allowlist_entries_missing_reason, + is_test_module, + is_tethered, + load_allowlist, + load_selections, + main, + normalize_test_path, + run_changed_files_check, + run_full_scan, +) + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def runs(command: str, test_file: str) -> bool: + """True if `command` collects `test_file` (path relative to ``tests/``).""" + return any(sel.runs(test_file) for sel in _parse_command(command)) + + +# --------------------------------------------------------------------------- # +# pytest path arguments +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("command", "test_file", "expected"), + [ + # Directory sweeps cover everything nested underneath. + ("pytest -v -s kernels", "kernels/test_attention.py", True), + ("pytest -v -s kernels/", "kernels/moe/test_moe.py", True), + ("pytest -v -s kernels", "lora/test_lora.py", False), + # An exact file argument covers that file and nothing beside it. + ("pytest -v -s kernels/test_a.py", "kernels/test_a.py", True), + ("pytest -v -s kernels/test_a.py", "kernels/test_b.py", False), + # Commands run from the repo root and from tests/ alike. + ("pytest -v -s tests/kernels/test_a.py", "kernels/test_a.py", True), + ("pytest -v -s ./kernels/test_a.py", "kernels/test_a.py", True), + ( + "pytest -v -s /vllm-workspace/tests/kernels/test_a.py", + "kernels/test_a.py", + True, + ), + # A ::nodeid selector still collects the file. + ("pytest -v -s kernels/test_a.py::test_one", "kernels/test_a.py", True), + ("pytest -v -s kernels/test_a.py::Cls::test_one", "kernels/test_a.py", True), + # Several positional paths, including a bare root-level file. + ("pytest -v -s test_envs.py test_outputs.py", "test_outputs.py", True), + ("pytest -v -s test_envs.py test_outputs.py", "test_regression.py", False), + # A bare word is a path only when it is really a directory under tests/. + ("pytest -v -s samplers", "samplers/test_sampler.py", True), + # Glob path args do not cross a directory separator, matching the shell. + ("pytest -v -s kernels/test_*.py", "kernels/test_a.py", True), + ("pytest -v -s kernels/test_*.py", "kernels/moe/test_a.py", False), + ("pytest -v -s kernels/test_*.py", "kernels/helper.py", False), + ], +) +def test_pytest_path_arguments(command, test_file, expected): + assert runs(command, test_file) is expected + + +# --------------------------------------------------------------------------- # +# --ignore / --deselect +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("command", "test_file", "expected"), + [ + # --ignore removes a file from an otherwise covering sweep, in both the + # `=` and space-separated spellings. + ("pytest -v -s kernels --ignore=kernels/test_a.py", "kernels/test_a.py", False), + ("pytest -v -s kernels --ignore kernels/test_a.py", "kernels/test_a.py", False), + # ...and only that file. + ("pytest -v -s kernels --ignore=kernels/test_a.py", "kernels/test_b.py", True), + # --ignore of a directory removes the whole subtree. + ("pytest -v -s kernels --ignore=kernels/moe", "kernels/moe/test_moe.py", False), + ( + "pytest -v -s lora --ignore-glob=lora/*_tp.py", + "lora/test_llama_tp.py", + False, + ), + # --ignore-glob is a whole-path fnmatch (pytest semantics): its `*` + # crosses `/`, so a bare `*_tp.py` drops a nested match too. A positional + # path glob would not - that is the distinction the parser has to keep. + ("pytest -v -s kernels --ignore-glob=*_tp.py", "kernels/test_a_tp.py", False), + ( + "pytest -v -s kernels --ignore-glob=*_tp.py", + "kernels/nested/test_a_tp.py", + False, + ), + ("pytest -v -s kernels --ignore-glob '*_tp.py'", "kernels/test_a.py", True), + # A node-id deselect leaves the file collected; a bare-path one does not. + ( + "pytest -v -s kernels --deselect kernels/test_a.py::test_one", + "kernels/test_a.py", + True, + ), + ( + "pytest -v -s kernels --deselect kernels/test_a.py", + "kernels/test_a.py", + False, + ), + ], +) +def test_ignore_and_deselect(command, test_file, expected): + assert runs(command, test_file) is expected + + +# --------------------------------------------------------------------------- # +# Options, env prefixes and other token noise +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("command", "test_file", "expected"), + [ + # A -k / -m filter narrows which tests run but does not un-run the file. + ("pytest -v -s kernels -k 'not slow'", "kernels/test_a.py", True), + ("pytest -v -s kernels -m core_model", "kernels/test_a.py", True), + # A value-taking option's value is never mistaken for a path. + ("pytest -m core_model", "core_model", False), + ("pytest -k kernels", "kernels/test_a.py", False), + ("pytest -n 4 kernels/test_a.py", "kernels/test_a.py", True), + # Buildkite shard flags carry a $$VAR and must not swallow the path. + ( + ( + "pytest -v -s kernels --shard-id=$$BUILDKITE_PARALLEL_JOB " + "--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT" + ), + "kernels/test_a.py", + True, + ), + # A leading env-var assignment is not a path. + ( + "VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s kernels", + "kernels/test_a.py", + True, + ), + # `python -m pytest` is the same invocation. + ("python3 -m pytest -v -s kernels/test_a.py", "kernels/test_a.py", True), + # Quoted paths, wrapper commands and redirects all survive tokenizing. + ("pytest -v -s 'kernels/test_a.py'", "kernels/test_a.py", True), + ("timeout 600 pytest kernels/test_a.py", "kernels/test_a.py", True), + ("pytest kernels/test_a.py > out.log 2>&1", "kernels/test_a.py", True), + # Each sub-command of a compound line is classified on its own. + ("pytest kernels/test_a.py && pytest lora/test_b.py", "lora/test_b.py", True), + ("pytest kernels/test_a.py ; pytest lora/test_b.py", "lora/test_b.py", True), + # An unresolvable $VAR path yields no coverage rather than a guess. + ("pytest -v -s $TEST_TARGET", "kernels/test_a.py", False), + ], +) +def test_option_and_token_handling(command, test_file, expected): + assert runs(command, test_file) is expected + + +# --------------------------------------------------------------------------- # +# find | xargs pytest pipelines +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("command", "test_file", "expected"), + [ + ("find kernels -name 'test_*.py' | xargs pytest -v", "kernels/test_a.py", True), + ( + "find kernels -name 'test_*.py' | xargs pytest -v", + "kernels/moe/test_b.py", + True, + ), + # A non-matching filename is not collected by the find expression. + ( + "find kernels -name 'test_*.py' | xargs pytest -v", + "kernels/helper.py", + False, + ), + # -maxdepth 1 keeps direct children and drops nested ones. + ( + "find kernels -maxdepth 1 -name 'test_*.py' | xargs pytest", + "kernels/test_a.py", + True, + ), + ( + "find kernels -maxdepth 1 -name 'test_*.py' | xargs pytest", + "kernels/moe/test_b.py", + False, + ), + # -not -name excludes its target and leaves siblings alone. + ( + "find kernels -name 'test_*.py' -not -name 'test_skip.py' | xargs pytest", + "kernels/test_skip.py", + False, + ), + ( + "find kernels -name 'test_*.py' -not -name 'test_skip.py' | xargs pytest", + "kernels/test_a.py", + True, + ), + # The root confines the sweep. + ("find kernels -name 'test_*.py' | xargs pytest", "lora/test_a.py", False), + # -maxdepth is honored even when the find root is the tests/ root itself + # ("." or "tests/"), not just a subdir. + ( + "find . -maxdepth 1 -name 'test_*.py' | xargs pytest", + "test_regression.py", + True, + ), + ( + "find . -maxdepth 1 -name 'test_*.py' | xargs pytest", + "v1/test_scheduler.py", + False, + ), + ( + "find tests/ -maxdepth 1 -name 'test_*.py' | xargs pytest", + "v1/test_scheduler.py", + False, + ), + # -exec pytest wires the sweep just as `| xargs pytest` does. + ( + "find compile -name 'test_*.py' -exec pytest -s -v {} \\;", + "compile/test_a.py", + True, + ), + # A pytest in a *separate* command on the line (not fed by the find) does + # not arm the sweep - the find here is piped to `wc`. + ( + "find kernels -name 'test_*.py' | wc -l ; pytest test_regression.py", + "kernels/test_a.py", + False, + ), + ( + "find kernels -name 'test_*.py' -delete && pytest test_regression.py", + "kernels/test_a.py", + False, + ), + # ...but only a *non-filtering* intermediate stage. `sort` reorders, + # every match still reaches pytest. + ( + "find kernels -name 'test_*.py' | sort | xargs pytest", + "kernels/test_a.py", + True, + ), + # `head` drops all but one match - the sweep is not fully tethered. + ( + "find kernels -name 'test_*.py' | head -n 1 | xargs pytest", + "kernels/test_a.py", + False, + ), + ( + "find kernels -name 'test_*.py' | grep attention | xargs pytest", + "kernels/test_a.py", + False, + ), + # `xargs` runs `echo` here, not pytest - the trailing `pytest` is echo's + # argument. + ( + "find kernels -name 'test_*.py' | xargs echo pytest", + "kernels/test_a.py", + False, + ), + # An `-exec` with no `\;` / `+` terminator is a `find` syntax error - it + # never runs pytest. + ( + "find kernels -name 'test_*.py' -exec pytest {}", + "kernels/test_a.py", + False, + ), + # An `-exec pytest` with no `{}` runs only the named file for every + # match - the matched files never reach pytest. + ( + "find kernels -name 'test_*.py' -exec pytest test_regression.py \\;", + "kernels/test_a.py", + False, + ), + # `{}` as an option *value* (`--ignore {}`) doesn't reach pytest as a + # test path either - only `test_regression.py` runs. + ( + ( + "find kernels -name 'test_*.py' -exec pytest test_regression.py " + "--ignore {} \\;" + ), + "kernels/test_a.py", + False, + ), + # xargs value-flags with a space don't swallow the command. + ( + "find kernels -name 'test_*.py' | xargs -0 -I {} pytest {}", + "kernels/test_a.py", + True, + ), + # `xargs env FOO=1 pytest` - xargs runs `env`, which runs pytest. + ( + "find kernels -name 'test_*.py' | xargs env FOO=1 pytest", + "kernels/test_a.py", + True, + ), + ], +) +def test_find_pipelines(command, test_file, expected): + assert runs(command, test_file) is expected + + +def test_find_pipeline_parses_as_find_selection(): + (selection,) = [ + s + for s in _parse_command("find kernels -name 'test_*.py' | xargs pytest -v") + if isinstance(s, FindSelection) + ] + assert selection.root == "kernels" + assert selection.name_globs == ["test_*.py"] + + +# --------------------------------------------------------------------------- # +# Direct file runners +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("command", "test_file", "expected"), + [ + ( + "python3 standalone_tests/lazy_imports.py", + "standalone_tests/lazy_imports.py", + True, + ), + ( + "python standalone_tests/lazy_imports.py", + "standalone_tests/lazy_imports.py", + True, + ), + # `timeout -k ` wrapping a runner: the grace value must not + # be consumed as the duration (which would hide `python`). + ( + "timeout -k 5s 30s python standalone_tests/lazy_imports.py", + "standalone_tests/lazy_imports.py", + True, + ), + ( + "torchrun --nproc-per-node 2 distributed/test_a.py", + "distributed/test_a.py", + True, + ), + # The option's value is a count, not a path. + ("torchrun --nproc-per-node 2 distributed/test_a.py", "2", False), + ( + "VLLM_X=1 python3 standalone_tests/lazy_imports.py | grep -q ok", + "standalone_tests/lazy_imports.py", + True, + ), + # Only the executed script counts; the tokens after it are that script's + # own argv. `--suffix v1` must not tether the whole tests/v1 subtree. + ( + ( + "python3 features/tensorize_vllm_model.py --model facebook/opt-125m " + "serialize --serialized-directory /tmp/ --suffix v1" + ), + "v1/test_scheduler.py", + False, + ), + ( + ( + "python3 features/tensorize_vllm_model.py --model facebook/opt-125m " + "serialize --serialized-directory /tmp/ --suffix v1" + ), + "features/tensorize_vllm_model.py", + True, + ), + # A script named by an absolute `$VAR` path still resolves. + ( + 'python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/e2e/test_x.py"', + "v1/e2e/test_x.py", + True, + ), + # `python -m ` runs the module, not a script - only `-m pytest` + # collects tests. `compileall kernels` byte-compiles, it does not run. + ("python -m compileall kernels", "kernels/test_a.py", False), + ("python -m mypy kernels/test_a.py", "kernels/test_a.py", False), + # `coverage run -m pytest` still resolves through to pytest. + ("coverage run -m pytest kernels/test_a.py", "kernels/test_a.py", True), + # `python -c '' file.py` runs the code string; `file.py` is argv. + ("python -c 'pass' kernels/test_a.py", "kernels/test_a.py", False), + ("python -c 'import sys' kernels/test_a.py", "kernels/test_a.py", False), + # a `-m pytest` *after* the `-c` code operand is that program's argv too. + ("python -c 'pass' -m pytest kernels/test_a.py", "kernels/test_a.py", False), + # ...but `-c` as the *script's* own arg still tethers the script. + ("python kernels/test_a.py -c config.py", "kernels/test_a.py", True), + # ...and pytest's own `-c ` (after `-m pytest`) is fine. + ("python -m pytest kernels/test_a.py -c setup.cfg", "kernels/test_a.py", True), + ], +) +def test_direct_runners(command, test_file, expected): + assert runs(command, test_file) is expected + + +# --------------------------------------------------------------------------- # +# Commands that must contribute no coverage +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "command", + [ + "export VLLM_WORKER_MULTIPROC_METHOD=spawn", + "pip install pytest-timeout pytest-forked", + "uv pip install -r requirements/test/cuda.in", + "bash tests/does_not_exist.sh", + "echo 'pytest kernels/test_a.py'", + # `pytest` as an argument to another command - `echo` runs, not pytest. + "echo pytest kernels/test_a.py", + 'pytest -v -s "unbalanced', + # A `.sh` that is only an *argument* (Buildkite printing / copying a + # filename) is not a script invocation. + "echo tests/foo.sh", + "cp scripts/run_tests.sh /tmp/", + # A `find` whose output never reaches pytest is not a test selection. + "find kernels -name 'test_*.py'", + "find kernels -name 'test_*.py' | wc -l", + "find kernels -name 'test_*.py' -delete", + "find kernels -name 'test_*.py' | xargs echo pytest", + "find kernels -name 'test_*.py' -exec pytest {}", + ], +) +def test_non_test_commands_contribute_nothing(command): + assert _parse_command(command) == [] or not any( + sel.runs("kernels/test_a.py") for sel in _parse_command(command) + ) + + +# --------------------------------------------------------------------------- # +# Job-level yaml parsing +# --------------------------------------------------------------------------- # + + +def _write_area(tmp_path, monkeypatch, body: str): + mod = checker + area_dir = tmp_path / "test_areas" + area_dir.mkdir(exist_ok=True) + (area_dir / "synthetic.yaml").write_text(textwrap.dedent(body)) + monkeypatch.setattr(mod, "TEST_AREAS_DIR", area_dir) + # Keep the real test-amd.yaml out of these isolated parsing tests. + monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "no-such-test-amd.yaml") + # The yaml-error path reports a repo-relative name, so REPO_ROOT has to + # contain the synthetic area too. + monkeypatch.setattr(mod, "REPO_ROOT", tmp_path) + return load_selections() + + +def test_plain_step_commands_are_collected(tmp_path, monkeypatch): + selections = _write_area( + tmp_path, + monkeypatch, + """ + steps: + - label: ":nvidia: (H200) Kernels" + commands: + - pytest -v -s kernels + """, + ) + assert is_tethered("kernels/test_a.py", selections) + + +def test_singular_command_key_is_collected(tmp_path, monkeypatch): + selections = _write_area( + tmp_path, + monkeypatch, + """ + steps: + - label: "Singular" + command: pytest -v -s kernels/test_a.py + """, + ) + assert is_tethered("kernels/test_a.py", selections) + + +def test_test_amd_yaml_coverage_counts_as_tethered(tmp_path, monkeypatch): + """A test wired only into the legacy test-amd.yaml is still run by CI.""" + mod = checker + (tmp_path / "test_areas").mkdir() + amd_yaml = tmp_path / "test-amd.yaml" + amd_yaml.write_text( + textwrap.dedent( + """ + steps: + - label: ":amd: (MI300) Quantization" + commands: + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ + - pytest -v -s rocm/test_moe_weight_replay.py + """ + ) + ) + monkeypatch.setattr(mod, "TEST_AREAS_DIR", tmp_path / "test_areas") + monkeypatch.setattr(mod, "TEST_AMD_YAML", amd_yaml) + monkeypatch.setattr(mod, "REPO_ROOT", tmp_path) + selections = load_selections() + assert is_tethered("rocm/test_moe_weight_replay.py", selections) + assert is_tethered("quantization/test_foo.py", selections) + + +def test_missing_test_amd_yaml_is_not_fatal(tmp_path, monkeypatch): + """When the migration finally deletes test-amd.yaml, the checker still runs.""" + mod = checker + (tmp_path / "test_areas").mkdir() + monkeypatch.setattr(mod, "TEST_AREAS_DIR", tmp_path / "test_areas") + monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "gone.yaml") + monkeypatch.setattr(mod, "REPO_ROOT", tmp_path) + assert load_selections() == [] + + +def test_job_dirs_come_from_ci_config(tmp_path, monkeypatch): + """The pipeline directories are whatever `.buildkite/ci_config*.yaml` lists + under `job_dirs:` - so a test wired only into e.g. `hardware_tests/` counts, + and an unreadable config falls back to `test_areas/` alone.""" + mod = checker + bk = tmp_path / ".buildkite" + (bk / "test_areas").mkdir(parents=True) + (bk / "hardware_tests").mkdir() + (bk / "ci_config.yaml").write_text("job_dirs:\n - .buildkite/hardware_tests\n") + (bk / "ci_config_broken.yaml").write_text("job_dirs: [oops\n") # ignored + (bk / "hardware_tests" / "cpu.yaml").write_text( + textwrap.dedent( + """ + steps: + - label: "CPU" + commands: + - pytest -v -s v1/e2e/test_cpu_thing.py + """ + ) + ) + monkeypatch.setattr(mod, "REPO_ROOT", tmp_path) + monkeypatch.setattr(mod, "TEST_AREAS_DIR", bk / "test_areas") + monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "no-amd.yaml") + + assert (bk / "hardware_tests") in checker._pipeline_job_dirs() + assert is_tethered("v1/e2e/test_cpu_thing.py", load_selections()) + + +def test_step_level_ignore_is_honored(tmp_path, monkeypatch): + selections = _write_area( + tmp_path, + monkeypatch, + """ + steps: + - label: "Kernels" + commands: + - pytest -v -s kernels --ignore=kernels/test_a.py + """, + ) + assert not is_tethered("kernels/test_a.py", selections) + assert is_tethered("kernels/test_b.py", selections) + + +def test_amd_mirror_only_command_is_collected(tmp_path, monkeypatch): + """A mirror that overrides `commands` runs tests the base step never does.""" + selections = _write_area( + tmp_path, + monkeypatch, + """ + steps: + - label: ":nvidia: (H200) Base" + commands: + - pytest -v -s kernels/test_base.py + mirror: + amd: + label: ":amd: (MI300) Base" + commands: + - pytest -v -s kernels/test_amd_only.py + """, + ) + assert is_tethered("kernels/test_base.py", selections) + assert is_tethered("kernels/test_amd_only.py", selections) + + +def test_shell_script_continuation_vars_and_nesting(tmp_path, monkeypatch): + """The v1 P/D integration scripts name their test through a `\\`-continued + line, a literal shell var, and a nested `bash "$SCRIPT"`. Each indirection + has to be followed or the check falsely reports the test untethered.""" + scripts = tmp_path / "tests" / "v1" / "kv_connector" / "nixl_integration" + scripts.mkdir(parents=True) + (scripts / "inner.sh").write_text( + "python3 -m pytest -s -v \\\n" + ' "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_inner.py"\n' + ) + (scripts / "sweep.sh").write_text( + 'SCRIPT="v1/kv_connector/nixl_integration/inner.sh"\n' + 'IMPORT_CANARY="v1/kv_connector/nixl_integration/test_canary.py"\n' + 'LITERAL="v1/kv_connector/nixl_integration/test_not_run.py"\n' + 'python3 -m pytest -s -x "${IMPORT_CANARY}"\n' + # single-quoted: the shell passes a literal $LITERAL, collecting nothing + "python3 -m pytest -s -x '${LITERAL}'\n" + 'if ! env FOO=1 bash "${SCRIPT}"; then exit 1; fi\n' + ) + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "TESTS_DIR", tmp_path / "tests") + + selections = _parse_command("bash v1/kv_connector/nixl_integration/sweep.sh") + + def tethered(path): + return any(s.runs(path) for s in selections) + + # `\`-continuation + $VAR-prefixed absolute path, resolved through a nested + # `env FOO=1 bash "$SCRIPT"`: + assert tethered("v1/kv_connector/nixl_integration/test_inner.py") + # literal `VAR=path` + double-quoted `"$VAR"` (the import-canary pattern): + assert tethered("v1/kv_connector/nixl_integration/test_canary.py") + # single-quoted `'$VAR'` is NOT expanded by the shell, so not tethered: + assert not tethered("v1/kv_connector/nixl_integration/test_not_run.py") + + +def test_bash_dash_c_does_not_follow_a_script_operand(tmp_path, monkeypatch): + """`bash -c '' foo.sh` runs the `-c` string; `foo.sh` is `$0` and is + never executed, so its contents must not count as coverage.""" + (tmp_path / "tests").mkdir() + (tmp_path / "tests" / "decoy.sh").write_text("pytest -v -s lora/test_decoy.py\n") + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "TESTS_DIR", tmp_path / "tests") + + with_c = _parse_command("bash -c 'echo hi' tests/decoy.sh") + assert not any(s.runs("lora/test_decoy.py") for s in with_c) + # ...but a plain `bash tests/decoy.sh` still follows it. + plain = _parse_command("bash tests/decoy.sh") + assert any(s.runs("lora/test_decoy.py") for s in plain) + # ...and a `-c` that is the *script's* own flag (not bash's) is not `bash -c`. + script_flag = _parse_command("bash tests/decoy.sh -c models.txt") + assert any(s.runs("lora/test_decoy.py") for s in script_flag) + # A `bash -c` trailing arg that looks like a command string is `$0`, not run. + cmd_string_arg = _parse_command("bash -c 'echo ok' 'pytest lora/test_x.py'") + assert not any(s.runs("lora/test_x.py") for s in cmd_string_arg) + # ...but the `-c` code string *itself* is parsed - it is what runs. + code_runs = _parse_command("bash -c 'pytest lora/test_x.py'") + assert any(s.runs("lora/test_x.py") for s in code_runs) + # a combined short-option group resolves the `-c` operand with `c` in any + # position (`-ec`, `-ce`), and its trailing args are still `$0`, not run. + for grp in ("-ec", "-ce", "-euxc"): + assert any( + s.runs("lora/test_x.py") + for s in _parse_command(f"bash {grp} 'pytest lora/test_x.py'") + ) + assert not any( + s.runs("lora/test_x.py") + for s in _parse_command("bash -ce 'echo ok' 'pytest lora/test_x.py'") + ) + + +def test_unparsable_yaml_is_fatal(tmp_path, monkeypatch): + """Silently skipping a bad yaml would drop its coverage and report false + untethered files, so it must fail loudly instead.""" + with pytest.raises(SystemExit): + _write_area(tmp_path, monkeypatch, "steps: [oops\n") + + +def test_non_mapping_yaml_is_fatal(tmp_path, monkeypatch): + """A yaml that parses but isn't a pipeline mapping (a bare list or scalar) + must raise the actionable error, not crash on ``.get``.""" + with pytest.raises(SystemExit): + _write_area(tmp_path, monkeypatch, "- just\n- a\n- list\n") + + +# --------------------------------------------------------------------------- # +# is_test_module / path normalization +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("path", "expected"), + [ + ("tests/kernels/test_a.py", True), + ("tests/kernels/a_test.py", True), + ("tests/kernels/conftest.py", False), + ("tests/kernels/__init__.py", False), + ("tests/kernels/utils.py", False), + ("tests/kernels/test_a.txt", False), + ("vllm/kernels/test_a.py", False), + ], +) +def test_is_test_module(path, expected): + assert is_test_module(path) is expected + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("tests/kernels/test_a.py", "kernels/test_a.py"), + ("./kernels/test_a.py", "kernels/test_a.py"), + ("/vllm-workspace/tests/kernels/test_a.py", "kernels/test_a.py"), + ("'kernels/test_a.py'", "kernels/test_a.py"), + ("kernels/test_a.py::test_one", "kernels/test_a.py"), + ("kernels/", "kernels"), + ("${GIT_ROOT}/tests/kernels/test_a.py", "kernels/test_a.py"), + ("$GIT_ROOT/tests/kernels/test_a.py", "kernels/test_a.py"), + ], +) +def test_normalize_test_path(raw, expected): + assert normalize_test_path(raw) == expected + + +# --------------------------------------------------------------------------- # +# Whole-tree invariants +# --------------------------------------------------------------------------- # + + +def test_real_tree_parses_into_selections(): + assert len(load_selections()) > 200 + assert len(all_test_modules()) > 1000 + + +def test_no_selection_matches_everything(): + """A path arg that normalizes to '' or '.' would tether every test in the + repo and silently disarm the gate - the worst possible misparse.""" + for selection in load_selections(): + if isinstance(selection, PytestSelection): + offenders = [ + p + for p in selection.included_paths + if normalize_test_path(p) in ("", ".") + ] + assert not offenders, f"match-everything pytest arg: {offenders}" + else: + assert normalize_test_path(selection.root) not in ("", ".") + + +def test_tethering_is_deterministic(): + selections = load_selections() + modules = all_test_modules()[:200] + first = [is_tethered(m, selections) for m in modules] + assert first == [is_tethered(m, load_selections()) for m in modules] + + +def test_allowlist_entries_are_well_formed(): + allowlist = load_allowlist() + assert allowlist, "allowlist should not be empty while gaps remain" + assert all(is_test_module(path) for path in allowlist) + assert all(not path.startswith("/") for path in allowlist) + + +def test_real_allowlist_entries_all_have_a_reason(): + assert allowlist_entries_missing_reason() == [] + + +def test_bare_allowlist_entry_is_flagged(tmp_path, monkeypatch): + allowlist = tmp_path / "allowlist.txt" + allowlist.write_text( + "tests/a/test_ok.py # torchrun multi-GPU, no job\n" + "tests/b/test_bare.py\n" + " # a standalone comment line is fine\n" + ) + monkeypatch.setattr(checker, "ALLOWLIST_PATH", allowlist) + assert allowlist_entries_missing_reason() == ["tests/b/test_bare.py"] + + +def test_real_tree_has_no_unallowlisted_gaps(): + """Every test module is either collected by a job or allowlisted. + + Deliberately non-strict: a stale allowlist entry is advisory by design, so + asserting `strict=True` here would make unrelated allowlist cleanup block + CI - the exact behaviour the two-mode design avoids. + """ + assert run_full_scan(load_selections(), load_allowlist(), strict=False) == 0 + + +# --------------------------------------------------------------------------- # +# CLI behaviour +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def fake_repo(tmp_path, monkeypatch): + """Point the module at a throwaway tree so the changed-files path can be + exercised without writing test files into the real repo.""" + mod = checker + (tmp_path / "tests" / "kernels").mkdir(parents=True) + (tmp_path / "tests" / "kernels" / "test_new.py").write_text("def test_x(): pass\n") + monkeypatch.setattr(mod, "REPO_ROOT", tmp_path) + monkeypatch.setattr(mod, "TESTS_DIR", tmp_path / "tests") + monkeypatch.setattr(mod, "ALLOWLIST_PATH", tmp_path / "allowlist.txt") + # The throwaway tree is not a git checkout, so the tracked-file inventory + # has to be stubbed rather than shelled out. + monkeypatch.setattr(mod, "all_test_modules", lambda: ["tests/kernels/test_new.py"]) + return tmp_path + + +def test_untethered_changed_file_fails(fake_repo, capsys): + rc = run_changed_files_check(["tests/kernels/test_new.py"], [], set()) + assert rc == 1 + assert "is not run by any Buildkite job" in capsys.readouterr().out + + +def test_tethered_changed_file_passes(fake_repo): + selections = _parse_command("pytest -v -s kernels") + assert ( + run_changed_files_check(["tests/kernels/test_new.py"], selections, set()) == 0 + ) + + +def test_allowlisted_changed_file_passes(fake_repo): + allowlist = {"tests/kernels/test_new.py"} + assert run_changed_files_check(["tests/kernels/test_new.py"], [], allowlist) == 0 + + +def test_non_test_and_missing_paths_are_ignored(fake_repo): + paths = ["tests/kernels/conftest.py", "vllm/config.py", "tests/kernels/gone.py"] + assert run_changed_files_check(paths, [], set()) == 0 + + +def test_stale_allowlist_entries_are_advisory_but_reported(fake_repo, capsys): + """Dangling entries must not block an unrelated PR, but they have to be + printed - the hook sets `verbose: true` so a passing run still shows them.""" + allowlist = {"tests/kernels/gone.py"} + rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=False) + out = capsys.readouterr().out + assert rc == 0 + assert "no longer exists" in out + + +def test_stale_allowlist_entries_are_fatal_under_all(fake_repo, capsys): + allowlist = {"tests/kernels/gone.py"} + rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=True) + assert rc == 1 + assert "no longer exists" in capsys.readouterr().out + + +def test_now_tethered_allowlist_entry_is_reported(fake_repo, capsys): + allowlist = {"tests/kernels/test_new.py"} + rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=True) + assert rc == 1 + assert "is now tethered" in capsys.readouterr().out + + +def test_main_passes_a_tethered_file(monkeypatch): + """The CLI's normal pre-commit path: a tethered file exits 0. + + `--all` is not asserted here on purpose - it implies `strict=True`, so a + stale allowlist entry would turn unrelated cleanup into a CI failure. + """ + monkeypatch.setattr( + "sys.argv", + ["check_test_tethering.py", "tests/tools/test_check_test_tethering.py"], + ) + assert main() == 0 + + +def test_main_flags_an_untethered_file(monkeypatch): + # Must be named like a test module, or the check correctly skips it. + orphan = REPO_ROOT / "tests" / "test_tethering_probe_orphan.py" + orphan.write_text("def test_x(): pass\n") + monkeypatch.setattr("sys.argv", ["check_test_tethering.py", str(orphan)]) + try: + assert main() == 1 + finally: + orphan.unlink() + + +def test_to_repo_relative_accepts_absolute_paths(): + absolute = str(REPO_ROOT / "tests" / "kernels" / "test_a.py") + assert _to_repo_relative(absolute) == "tests/kernels/test_a.py" diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tools/pre_commit/__init__.py b/tools/pre_commit/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tools/pre_commit/check_test_tethering.py b/tools/pre_commit/check_test_tethering.py new file mode 100644 index 000000000000..ee080a675734 --- /dev/null +++ b/tools/pre_commit/check_test_tethering.py @@ -0,0 +1,925 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check that test files are wired into ("tethered to") the Buildkite CI. + +A test file is *tethered* when at least one Buildkite job actually collects it - +some job's ``commands`` run ``pytest`` (or ``python`` / ``torchrun``, or a +``find ... | xargs pytest`` pipeline) over a path that includes the file. A test +that is collected by no job silently never runs: it passes review, merges, and +rots. + +How it works +------------ +1. Parse every job in every Buildkite pipeline yaml (and each ``mirror.amd`` + sub-step), and turn each shell command into a :class:`Selection` describing + which test files it runs. The pipeline directories are read from + ``.buildkite/ci_config*.yaml`` (``job_dirs:``) - currently ``test_areas/``, + ``hardware_tests/`` and ``intel_jobs/`` across the CUDA, ROCm and Intel + pipelines - plus the legacy ``.buildkite/test-amd.yaml``. +2. Build the inventory of test files under ``tests/``. +3. A file is a problem if no selection runs it and it is not in the allowlist + (``tools/pre_commit/test_tethering_allowlist.txt`` - the set of pre-existing + gaps, which is only ever meant to shrink). Every allowlist entry must carry a + trailing ``# ``; a bare path line is itself an error. + +Modes +----- +* **changed-files** (default): pre-commit passes the changed test files as + arguments; each one must be tethered or allowlisted. +* **full scan**: ``--all``, or automatically whenever the change set touches a + pipeline yaml or the allowlist itself - because editing a job command can + orphan a test that is not otherwise in the diff. The full scan also reports + allowlist entries that have since become tethered, or whose file is gone, so + the list stays honest. + +Coverage from *any* of those pipelines counts - the question is "does some +Buildkite job run this test", not "which pipeline". ``test-amd.yaml`` is being +folded into ``test_areas/`` a group at a time; when it is gone the checker stops +parsing it automatically. + +Usage:: + + python tools/pre_commit/check_test_tethering.py [FILES...] + python tools/pre_commit/check_test_tethering.py --all +""" + +import argparse +import fnmatch +import shlex +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path + +import regex as re +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +BUILDKITE_DIR = REPO_ROOT / ".buildkite" +# Fallback pipeline dir, always parsed even if ``ci_config*.yaml`` can't be read. +TEST_AREAS_DIR = BUILDKITE_DIR / "test_areas" +# The legacy hand-maintained AMD pipeline. It is being migrated into +# ``test_areas/`` piecemeal; until that finishes, a test wired only into this +# file is still genuinely run by CI, so it counts as tethered. When the file is +# removed, ``_pipeline_yaml_paths()`` stops appending it automatically. +TEST_AMD_YAML = BUILDKITE_DIR / "test-amd.yaml" +TESTS_DIR = REPO_ROOT / "tests" +ALLOWLIST_PATH = Path(__file__).resolve().parent / "test_tethering_allowlist.txt" + +# `pytest` options that take a value in the *next* token, so that token must not +# be mistaken for a test path (e.g. `-k expr`, `--shard-id 3`). Options that only +# ever take their value with `=` (`--cov=...`) don't belong here - a bare +# `--cov` takes no value and would wrongly swallow the following test path. +PYTEST_OPTIONS_WITH_VALUE = { + "-k", "-m", "-p", "-n", "-c", "-o", "-W", "-r", + "--shard-id", "--num-shards", "--dist", "--durations", "--maxfail", + "--timeout", "--rootdir", "--junitxml", "--tb", +} # fmt: skip + +# `pytest` options whose value is a path that is *removed* from collection at +# file granularity. ``--ignore`` takes a path prefix; ``--ignore-glob`` takes a +# glob that pytest matches against the whole path with ``fnmatch`` (so its ``*`` +# *does* cross a ``/``), which is why the two are modelled separately. +PYTEST_IGNORE_OPTIONS = {"--ignore", "--ignore-glob"} + +# `--deselect` takes a value too, but it removes individual node IDs - a file +# with one parametrization deselected is still collected and still "run". +PYTEST_DESELECT_OPTION = "--deselect" + +PYTEST_COMMANDS = {"pytest", "py.test"} + +# A leading ``$VAR/`` or ``${VAR}/`` path prefix - CI shell scripts routinely +# name a test by an absolute path rooted at a shell var +# (``"${GIT_ROOT}/tests/v1/..."``). The prefix is dropped before the path is +# interpreted; an unresolved ``$`` anywhere else still makes the token unusable. +_LEADING_VAR_PREFIX = re.compile(r"^\$\{?\w+\}?/") + +# Commands that execute a Python file directly. Only the *first* path argument is +# the script that runs; everything after it is that script's own argv, not more +# test files (``python tensorize.py --suffix v1`` runs one file, not ``tests/v1``). +FILE_RUNNER_COMMANDS = {"python", "python3", "torchrun", "coverage"} + + +# --------------------------------------------------------------------------- # +# Shell-token helpers +# --------------------------------------------------------------------------- # + + +def _is_env_assignment(token: str) -> bool: + """True for a leading ``FOO=bar`` shell env-var assignment - the kind some + commands put before ``pytest`` (``VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest``). + """ + if "=" not in token: + return False + name = token.split("=", 1)[0] + return bool(name) and name.replace("_", "").isupper() + + +def normalize_test_path(raw: str) -> str: + """Reduce a path as it appears in a command to one comparable form: + relative to ``tests/``, no surrounding quotes, no ``::nodeid`` selector, no + trailing slash. + + Commands run from either the repo root or from ``tests/``, so both a + ``tests/`` and a ``/vllm-workspace/tests/`` prefix are stripped, as is a + leading ``$VAR/`` (``"${GIT_ROOT}/tests/..."``). + """ + path = raw.strip().strip("'\"") + path = path.split("::", 1)[0] + path = _LEADING_VAR_PREFIX.sub("", path) + for prefix in ("./", "tests/", "/vllm-workspace/tests/"): + if path.startswith(prefix): + path = path[len(prefix) :] + return path.rstrip("/") + + +def _token_is_test_path(token: str) -> bool: + """True if a shell token is a path argument rather than a flag, an env-var + assignment, a Buildkite ``$VAR``, or a ``-m`` / ``-k`` expression value. + + A token with a ``/`` or a ``.py`` ending is always a path. A bare word + (``samplers``) only counts if it is really a directory under ``tests/``. A + leading ``$VAR/`` is dropped first, but a ``$`` anywhere else means the token + can't be resolved to a known path. + """ + if not token or token.startswith("-"): + return False + if _is_env_assignment(token): + return False + token = _LEADING_VAR_PREFIX.sub("", token) + if "$" in token: + return False + if "/" in token or token.endswith(".py"): + return True + return (TESTS_DIR / normalize_test_path(token)).is_dir() + + +def _path_is_within(parent: str, path: str) -> bool: + """True if ``path`` is ``parent`` itself or nested underneath it. A ``''`` / + ``'.'`` parent (a command run with no path arg) contains everything.""" + if parent in ("", "."): + return True + return path == parent or path.startswith(parent + "/") + + +def _glob_match_path(pattern: str, path: str) -> bool: + """Like :func:`fnmatch.fnmatch`, but ``*`` and ``?`` do not match across a + ``/`` - the way the shell and ``pytest`` treat an unexpanded glob path + argument (``kernels/test_foo_*.py`` selects one directory, not the tree). + A ``**`` component keeps the cross-separator meaning.""" + if not any(ch in pattern for ch in "*?["): + return pattern == path + if "**" in pattern: + return fnmatch.fnmatch(path, pattern) + pattern_parts = pattern.split("/") + path_parts = path.split("/") + return len(pattern_parts) == len(path_parts) and all( + fnmatch.fnmatch(part, glob) for part, glob in zip(path_parts, pattern_parts) + ) + + +# --------------------------------------------------------------------------- # +# Selections - "which test files does this command run?" +# --------------------------------------------------------------------------- # + + +@dataclass +class PytestSelection: + """A ``pytest``-style command: run everything under ``included_paths`` + except what an ``--ignore`` (path prefix) or ``--ignore-glob`` (whole-path + ``fnmatch``) removes.""" + + included_paths: list[str] = field(default_factory=list) + ignored_paths: list[str] = field(default_factory=list) + ignored_globs: list[str] = field(default_factory=list) + + def runs(self, test_file: str) -> bool: + included = any(self._arg_covers(arg, test_file) for arg in self.included_paths) + if not included: + return False + if any(self._arg_covers(arg, test_file) for arg in self.ignored_paths): + return False + # `--ignore-glob` is a plain `fnmatch` over the whole path - `*_tp.py` + # drops `kernels/nested/test_a_tp.py`, unlike a positional glob arg. + return not any( + fnmatch.fnmatch(test_file, normalize_test_path(glob)) + for glob in self.ignored_globs + ) + + @staticmethod + def _arg_covers(path_arg: str, test_file: str) -> bool: + """True if a single pytest path argument selects ``test_file`` - by + exact match, by being a parent directory, or as a glob.""" + path_arg = normalize_test_path(path_arg) + if _path_is_within(path_arg, test_file): + return True + return _glob_match_path(path_arg, test_file) + + +@dataclass +class FindSelection: + """A ``find -name ... [-not -name ...] [-maxdepth N] | xargs pytest`` + pipeline. ``pytest`` runs one file per ``find`` hit, so the ``find`` + expression alone defines the selected set.""" + + root: str + name_globs: list[str] = field(default_factory=list) + exclude_name_globs: list[str] = field(default_factory=list) + max_depth: int | None = None + + def runs(self, test_file: str) -> bool: + root = normalize_test_path(self.root) + if not _path_is_within(root, test_file): + return False + + # `find ` puts a direct child at depth 1, whether is a real + # subdir ("compile") or the tests/ root itself ("" / "."). + below = test_file if root in ("", ".") else test_file[len(root) :] + depth_below_root = below.lstrip("/").count("/") + 1 + if self.max_depth is not None and depth_below_root > self.max_depth: + return False + + filename = test_file.rsplit("/", 1)[-1] + if self.name_globs and not _glob_any(filename, self.name_globs): + return False + return not _glob_any(filename, self.exclude_name_globs) + + +Selection = PytestSelection | FindSelection + + +def _glob_any(name: str, globs: list[str]) -> bool: + return any(fnmatch.fnmatch(name, glob) for glob in globs) + + +# --------------------------------------------------------------------------- # +# Parsing CI commands into selections +# --------------------------------------------------------------------------- # + + +def _parse_pytest_command(tokens: list[str]) -> PytestSelection: + """Read a ``pytest ...`` invocation's tokens into a :class:`PytestSelection`. + + Positional path args are the included paths; ``--ignore`` values (and a + bare-path ``--deselect``) are ignored paths, ``--ignore-glob`` values are + ignored globs. Value-taking options and their values, and every other flag, + are skipped. ``-k`` / ``-m`` filtering is intentionally *not* modelled - a + file that a marker expression narrows is still "run". + """ + selection = PytestSelection() + + # Advance past the pytest token itself (and any env-var / wrapper prefix). + tokens = iter(tokens) + for token in tokens: + if token in PYTEST_COMMANDS: + break + + for token in tokens: + opt, eq, inline = token.partition("=") + if opt in PYTEST_IGNORE_OPTIONS: + value = inline if eq else next(tokens, "") + if opt == "--ignore-glob": + selection.ignored_globs.append(value) + else: + selection.ignored_paths.append(value) + elif opt == PYTEST_DESELECT_OPTION: + value = inline if eq else next(tokens, "") + # A `path::node_id` deselect leaves the file collected; only a + # bare-path deselect removes anything at file granularity. + if value and "::" not in value: + selection.ignored_paths.append(value) + elif token in PYTEST_OPTIONS_WITH_VALUE: + next(tokens, None) # consume and discard the option's value + elif token.startswith("-"): + continue + elif _token_is_test_path(token): + selection.included_paths.append(token) + + return selection + + +def _parse_find_command(tokens: list[str]) -> FindSelection | None: + """Read a ``find -name ...`` command's tokens into a + :class:`FindSelection`, or return ``None`` if there is no usable ``find``.""" + if "find" not in tokens: + return None + + root: str | None = None + name_globs: list[str] = [] + exclude_name_globs: list[str] = [] + max_depth: int | None = None + + tokens = iter(tokens[tokens.index("find") + 1 :]) + negated = False + for token in tokens: + if token in ("-not", "!"): + negated = True + continue + if token in ("-name", "-iname"): + glob = next(tokens, "").strip("'\"") + (exclude_name_globs if negated else name_globs).append(glob) + elif token == "-maxdepth": + depth = next(tokens, "") + max_depth = int(depth) if depth.isdigit() else None + elif ( + root is None and not token.startswith("-") and token not in PYTEST_COMMANDS + ): + root = token + negated = False + + if root is None: + return None + return FindSelection(root, name_globs, exclude_name_globs, max_depth) + + +def _dash_c_code_operand(tokens: list[str]) -> str | None: + """For ``bash -c '' [$0] [args]`` / ``sh -c ...`` / a combined short + group with ``c`` anywhere in it (``-ec``, ``-ce``, ``-euxc``), the ```` + string the shell runs. ``None`` when this is not a ``-c`` invocation (plain + ``bash script.sh`` / ``bash -x script.sh``). Everything after the code + operand is a positional parameter (``$0``, ``$1``, ...), not a script or + command that executes.""" + if not tokens or tokens[0].rsplit("/", 1)[-1] not in ("bash", "sh"): + return None + for i, token in enumerate(tokens[1:], start=1): + if not token.startswith("-"): + return None # script operand reached; a later `-c` is *its* flag + # a short-option group (single dash) with `c` among its flags. + if not token.startswith("--") and "c" in token[1:]: + return tokens[i + 1] if i + 1 < len(tokens) else None + return None + + +def _parse_shell_script(tokens: list[str], visited: set[str]) -> list[Selection]: + """Follow a ``bash tests/foo.sh`` command: read the script and parse the + pytest / runner lines inside it. ``visited`` stops a script that (directly or + otherwise) refers back to itself; nesting is one level deep in practice. + """ + if _dash_c_code_operand(tokens) is not None: + return [] # the `-c` code string is parsed by the caller, not a file + script_arg = next((t for t in tokens if t.endswith(".sh")), None) + if script_arg is None: + return [] + + rel = script_arg.strip().strip("'\"").replace("/vllm-workspace/", "") + rel = rel.replace("vllm-workspace/", "").removeprefix("./").removeprefix("/") + candidates = [REPO_ROOT / rel, TESTS_DIR / normalize_test_path(rel)] + script_path = next((p for p in candidates if p.is_file()), None) + if script_path is None or str(script_path) in visited: + return [] + visited.add(str(script_path)) + + # Join `\`-continued lines (a `pytest ... \` with its path on the next line + # is common), then track literal `VAR=value` assignments so a later + # `pytest "$VAR"` - the import-canary pattern - resolves to a real path. + text = script_path.read_text().replace("\\\n", " ") + var_values: dict[str, str] = {} + selections: list[Selection] = [] + for line in text.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + assignment = re.match(r"^(\w+)=(\S+)$", line) + if assignment: + var_values[assignment.group(1)] = assignment.group(2).strip("'\"") + continue + if var_values: + line = _expand_vars(line, var_values) + selections.extend(_parse_command(line, visited)) + return selections + + +def _expand_vars(line: str, var_values: dict[str, str]) -> str: + """Substitute known ``$VAR`` / ``${VAR}`` in a script line - but not inside + single quotes, where the shell does not expand (``pytest '$X'`` passes a + literal ``$X``, collecting nothing).""" + + def sub(segment: str) -> str: + return re.sub( + r"\$\{(\w+)\}|\$(\w+)", + lambda m: var_values.get(m.group(1) or m.group(2), m.group(0)), + segment, + ) + + # Splitting on `'`: even-indexed segments are outside single quotes, odd are + # inside. (A `'` within double quotes splits wrongly, but that only ever + # drops a substitution - the safe direction.) + parts = line.split("'") + return "'".join(sub(p) if i % 2 == 0 else p for i, p in enumerate(parts)) + + +# Shell tokens (or runs of them) that separate one sub-command from the next. +_SHELL_SEP_CHARS = set("|&;()<>") +_RUNNER_KEYWORDS = PYTEST_COMMANDS | FILE_RUNNER_COMMANDS + +# Tokens that can precede the real command word without being it. +_COMMAND_WRAPPERS = {"if", "!", "then", "time", "exec", "nice", "command", "builtin"} + + +def _effective_command_index(tokens: list[str]) -> int: + """Index of the token that names the command actually executed, past leading + ``FOO=bar`` assignments and simple wrappers - ``if ! env FOO=1 bash x.sh`` + resolves to the ``bash``, so ``echo x.sh`` stays ``echo`` and is not + followed. Returns ``len(tokens)`` when nothing executable is left.""" + i = 0 + while i < len(tokens): + token = tokens[i] + if _is_env_assignment(token) or token in _COMMAND_WRAPPERS: + i += 1 + elif token == "env": + i += 1 + while i < len(tokens) and ( + _is_env_assignment(tokens[i]) or "$" in tokens[i] or tokens[i] == "-" + ): + i += 1 + elif token == "timeout" and i + 1 < len(tokens): + i += 1 + while i < len(tokens) and tokens[i].startswith("-"): + # `-k` / `--kill-after` / `-s` / `--signal` take a value. + if tokens[i] in ("-k", "--kill-after", "-s", "--signal"): + i += 1 + i += 1 + i += 1 # the DURATION operand + else: + break + return i + + +def _split_shell_stages(command: str) -> list[tuple[str, list[str]]]: + """Tokenize a shell line - honouring quotes - and split it at the shell + operators, tagging each stage with the operator token that *precedes* it + (``""`` for the first stage, then ``"|"`` / ``"&&"`` / ``"||"`` / ``";"`` / + ``"&"`` / a subshell paren). Knowing which operator joins two stages is what + lets a ``find`` sweep count only when it is genuinely piped into ``pytest``, + not merely on the same line as one. A line that cannot be tokenized + (unbalanced quote, ...) yields nothing: better to miss a selection than to + misparse one. Splitting *after* tokenizing means an operator inside a quoted + argument stays part of that argument.""" + lexer = shlex.shlex(command, posix=True, punctuation_chars=True) + lexer.whitespace_split = True + try: + tokens = list(lexer) + except ValueError: + return [] + + stages: list[tuple[str, list[str]]] = [("", [])] + in_exec = False # inside a `find -exec ... \;` action - its `;` terminates + for token in tokens: + if in_exec and token in (";", "+"): + stages[-1][1].append(token) + in_exec = False + elif token and set(token) <= _SHELL_SEP_CHARS: + stages.append((token, [])) + in_exec = False + else: + stages[-1][1].append(token) + if token in ("-exec", "-execdir"): + in_exec = True + return [(sep, sub) for sep, sub in stages if sub] + + +def _looks_like_command_string(token: str) -> bool: + """True for a single token that is itself a command line - e.g. a quoted + ``"pytest a.py && torchrun b.py"`` passed as an argument to a runner + script.""" + return " " in token and any(word in _RUNNER_KEYWORDS for word in token.split()) + + +def _classify_subcommand(tokens: list[str], visited: set[str]) -> list[Selection]: + """Map one already-tokenized sub-command to the selections it implies: + find / pytest / ``bash